peer_protocol/ut-holepunch/ut-holepunch.go

   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 }