mirror of
https://github.com/netbirdio/netbird.git
synced 2026-04-21 17:56:39 +00:00
233
client/iface/bind/udp_muxed_conn.go
Normal file
233
client/iface/bind/udp_muxed_conn.go
Normal file
@@ -0,0 +1,233 @@
|
||||
package bind
|
||||
|
||||
/*
|
||||
Most of this code was copied from https://github.com/pion/ice and modified to fulfill NetBird's requirements
|
||||
*/
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/pion/logging"
|
||||
"github.com/pion/transport/v3/packetio"
|
||||
)
|
||||
|
||||
type udpMuxedConnParams struct {
|
||||
Mux *UDPMuxDefault
|
||||
AddrPool *sync.Pool
|
||||
Key string
|
||||
LocalAddr net.Addr
|
||||
Logger logging.LeveledLogger
|
||||
}
|
||||
|
||||
// udpMuxedConn represents a logical packet conn for a single remote as identified by ufrag
|
||||
type udpMuxedConn struct {
|
||||
params *udpMuxedConnParams
|
||||
// remote addresses that we have sent to on this conn
|
||||
addresses []string
|
||||
|
||||
// channel holding incoming packets
|
||||
buf *packetio.Buffer
|
||||
closedChan chan struct{}
|
||||
closeOnce sync.Once
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func newUDPMuxedConn(params *udpMuxedConnParams) *udpMuxedConn {
|
||||
p := &udpMuxedConn{
|
||||
params: params,
|
||||
buf: packetio.NewBuffer(),
|
||||
closedChan: make(chan struct{}),
|
||||
}
|
||||
|
||||
return p
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) ReadFrom(b []byte) (n int, rAddr net.Addr, err error) {
|
||||
buf := c.params.AddrPool.Get().(*bufferHolder) //nolint:forcetypeassert
|
||||
defer c.params.AddrPool.Put(buf)
|
||||
|
||||
// read address
|
||||
total, err := c.buf.Read(buf.buf)
|
||||
if err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
|
||||
dataLen := int(binary.LittleEndian.Uint16(buf.buf[:2]))
|
||||
if dataLen > total || dataLen > len(b) {
|
||||
return 0, nil, io.ErrShortBuffer
|
||||
}
|
||||
|
||||
// read data and then address
|
||||
offset := 2
|
||||
copy(b, buf.buf[offset:offset+dataLen])
|
||||
offset += dataLen
|
||||
|
||||
// read address len & decode address
|
||||
addrLen := int(binary.LittleEndian.Uint16(buf.buf[offset : offset+2]))
|
||||
offset += 2
|
||||
|
||||
if rAddr, err = decodeUDPAddr(buf.buf[offset : offset+addrLen]); err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
|
||||
return dataLen, rAddr, nil
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) WriteTo(buf []byte, rAddr net.Addr) (n int, err error) {
|
||||
if c.isClosed() {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
// each time we write to a new address, we'll register it with the mux
|
||||
addr := rAddr.String()
|
||||
if !c.containsAddress(addr) {
|
||||
c.addAddress(addr)
|
||||
}
|
||||
|
||||
return c.params.Mux.writeTo(buf, rAddr)
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) LocalAddr() net.Addr {
|
||||
return c.params.LocalAddr
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) SetDeadline(tm time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) SetReadDeadline(tm time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) SetWriteDeadline(tm time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) CloseChannel() <-chan struct{} {
|
||||
return c.closedChan
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) Close() error {
|
||||
var err error
|
||||
c.closeOnce.Do(func() {
|
||||
err = c.buf.Close()
|
||||
close(c.closedChan)
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) isClosed() bool {
|
||||
select {
|
||||
case <-c.closedChan:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) getAddresses() []string {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
addresses := make([]string, len(c.addresses))
|
||||
copy(addresses, c.addresses)
|
||||
return addresses
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) addAddress(addr string) {
|
||||
c.mu.Lock()
|
||||
c.addresses = append(c.addresses, addr)
|
||||
c.mu.Unlock()
|
||||
|
||||
// map it on mux
|
||||
c.params.Mux.registerConnForAddress(c, addr)
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) containsAddress(addr string) bool {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
for _, a := range c.addresses {
|
||||
if addr == a {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (c *udpMuxedConn) writePacket(data []byte, addr *net.UDPAddr) error {
|
||||
// write two packets, address and data
|
||||
buf := c.params.AddrPool.Get().(*bufferHolder) //nolint:forcetypeassert
|
||||
defer c.params.AddrPool.Put(buf)
|
||||
|
||||
// format of buffer | data len | data bytes | addr len | addr bytes |
|
||||
if len(buf.buf) < len(data)+maxAddrSize {
|
||||
return io.ErrShortBuffer
|
||||
}
|
||||
// data len
|
||||
binary.LittleEndian.PutUint16(buf.buf, uint16(len(data)))
|
||||
offset := 2
|
||||
|
||||
// data
|
||||
copy(buf.buf[offset:], data)
|
||||
offset += len(data)
|
||||
|
||||
// write address first, leaving room for its length
|
||||
n, err := encodeUDPAddr(addr, buf.buf[offset+2:])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
total := offset + n + 2
|
||||
|
||||
// address len
|
||||
binary.LittleEndian.PutUint16(buf.buf[offset:], uint16(n))
|
||||
|
||||
if _, err := c.buf.Write(buf.buf[:total]); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func encodeUDPAddr(addr *net.UDPAddr, buf []byte) (int, error) {
|
||||
ipData, err := addr.IP.MarshalText()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
total := 2 + len(ipData) + 2 + len(addr.Zone)
|
||||
if total > len(buf) {
|
||||
return 0, io.ErrShortBuffer
|
||||
}
|
||||
|
||||
binary.LittleEndian.PutUint16(buf, uint16(len(ipData)))
|
||||
offset := 2
|
||||
n := copy(buf[offset:], ipData)
|
||||
offset += n
|
||||
binary.LittleEndian.PutUint16(buf[offset:], uint16(addr.Port))
|
||||
offset += 2
|
||||
copy(buf[offset:], addr.Zone)
|
||||
return total, nil
|
||||
}
|
||||
|
||||
func decodeUDPAddr(buf []byte) (*net.UDPAddr, error) {
|
||||
addr := net.UDPAddr{}
|
||||
|
||||
offset := 0
|
||||
ipLen := int(binary.LittleEndian.Uint16(buf[:2]))
|
||||
offset += 2
|
||||
// basic bounds checking
|
||||
if ipLen+offset > len(buf) {
|
||||
return nil, io.ErrShortBuffer
|
||||
}
|
||||
if err := addr.IP.UnmarshalText(buf[offset : offset+ipLen]); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
offset += ipLen
|
||||
addr.Port = int(binary.LittleEndian.Uint16(buf[offset : offset+2]))
|
||||
offset += 2
|
||||
zone := make([]byte, len(buf[offset:]))
|
||||
copy(zone, buf[offset:])
|
||||
addr.Zone = string(zone)
|
||||
|
||||
return &addr, nil
|
||||
}
|
||||
Reference in New Issue
Block a user