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 const (
  29         Ipv4 AddrType = iota
  30         Ipv6 AddrType = iota
  31 )
  32
  33 func (m *Msg) UnmarshalBinary(b []byte) error {
  34         if len(b) < 12 {
  35                 return fmt.Errorf("buffer too small to be valid")
  36         }
  37         m.MsgType = MsgType(b[0])
  38         b = b[1:]
  39         addrType := AddrType(b[0])
  40         b = b[1:]
  41         var addr netip.Addr
  42         switch addrType {
  43         case Ipv4:
  44                 addr = netip.AddrFrom4([4]byte(b[:4]))
  45                 b = b[4:]
  46         case Ipv6:
  47                 if len(b) < 22 {
  48                         return fmt.Errorf("not enough bytes")
  49                 }
  50                 addr = netip.AddrFrom16([16]byte(b[:16]))
  51                 b = b[16:]
  52         default:
  53                 return fmt.Errorf("unhandled addr type value %v", addrType)
  54         }
  55         port := binary.BigEndian.Uint16(b[:])
  56         b = b[2:]
  57         m.AddrPort = netip.AddrPortFrom(addr, port)
  58         m.ErrCode = ErrCode(binary.BigEndian.Uint32(b[:]))
  59         b = b[4:]
  60         if len(b) != 0 {
  61                 return fmt.Errorf("%v trailing unused bytes", len(b))
  62         }
  63         return nil
  64 }
  65
  66 func (m *Msg) MarshalBinary() (_ []byte, err error) {
  67         var buf bytes.Buffer
  68         buf.Grow(24)
  69         buf.WriteByte(byte(m.MsgType))
  70         addr := m.AddrPort.Addr()
  71         switch {
  72         case addr.Is4():
  73                 buf.WriteByte(byte(Ipv4))
  74         case addr.Is6():
  75                 buf.WriteByte(byte(Ipv6))
  76         default:
  77                 err = fmt.Errorf("unhandled addr type: %v", addr)
  78                 return
  79         }
  80         buf.Write(addr.AsSlice())
  81         binary.Write(&buf, binary.BigEndian, m.AddrPort.Port())
  82         binary.Write(&buf, binary.BigEndian, m.ErrCode)
  83         return buf.Bytes(), nil
  84 }