github.com/anacrolix/torrent@v1.61.0/peer_protocol/ut-holepunch/ut-holepunch.go (about)

     1  package utHolepunch
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"fmt"
     7  	"net/netip"
     8  )
     9  
    10  const ExtensionName = "ut_holepunch"
    11  
    12  type (
    13  	Msg struct {
    14  		MsgType  MsgType
    15  		AddrPort netip.AddrPort
    16  		ErrCode  ErrCode
    17  	}
    18  	MsgType  byte
    19  	AddrType byte
    20  )
    21  
    22  const (
    23  	Rendezvous MsgType = iota
    24  	Connect
    25  	Error
    26  )
    27  
    28  func (me MsgType) String() string {
    29  	switch me {
    30  	case Rendezvous:
    31  		return "rendezvous"
    32  	case Connect:
    33  		return "connect"
    34  	case Error:
    35  		return "error"
    36  	default:
    37  		return fmt.Sprintf("unknown %d", me)
    38  	}
    39  }
    40  
    41  const (
    42  	Ipv4 AddrType = iota
    43  	Ipv6 AddrType = iota
    44  )
    45  
    46  func (m *Msg) UnmarshalBinary(b []byte) error {
    47  	if len(b) < 12 {
    48  		return fmt.Errorf("buffer too small to be valid")
    49  	}
    50  	m.MsgType = MsgType(b[0])
    51  	b = b[1:]
    52  	addrType := AddrType(b[0])
    53  	b = b[1:]
    54  	var addr netip.Addr
    55  	switch addrType {
    56  	case Ipv4:
    57  		addr = netip.AddrFrom4(*(*[4]byte)(b[:4]))
    58  		b = b[4:]
    59  	case Ipv6:
    60  		if len(b) < 22 {
    61  			return fmt.Errorf("not enough bytes")
    62  		}
    63  		addr = netip.AddrFrom16(*(*[16]byte)(b[:16]))
    64  		b = b[16:]
    65  	default:
    66  		return fmt.Errorf("unhandled addr type value %v", addrType)
    67  	}
    68  	port := binary.BigEndian.Uint16(b[:])
    69  	b = b[2:]
    70  	m.AddrPort = netip.AddrPortFrom(addr, port)
    71  	m.ErrCode = ErrCode(binary.BigEndian.Uint32(b[:]))
    72  	b = b[4:]
    73  	if len(b) != 0 {
    74  		return fmt.Errorf("%v trailing unused bytes", len(b))
    75  	}
    76  	return nil
    77  }
    78  
    79  func (m *Msg) MarshalBinary() (_ []byte, err error) {
    80  	var buf bytes.Buffer
    81  	buf.Grow(24)
    82  	buf.WriteByte(byte(m.MsgType))
    83  	addr := m.AddrPort.Addr()
    84  	switch {
    85  	case addr.Is4():
    86  		buf.WriteByte(byte(Ipv4))
    87  	case addr.Is6():
    88  		buf.WriteByte(byte(Ipv6))
    89  	default:
    90  		err = fmt.Errorf("unhandled addr type: %v", addr)
    91  		return
    92  	}
    93  	buf.Write(addr.AsSlice())
    94  	binary.Write(&buf, binary.BigEndian, m.AddrPort.Port())
    95  	binary.Write(&buf, binary.BigEndian, m.ErrCode)
    96  	return buf.Bytes(), nil
    97  }