ceremonyclient/go-libp2p/p2p/transport/webrtc/udpmux/mux.go

299 lines
8.2 KiB
Go
Raw Normal View History

2024-06-08 11:32:45 +00:00
package udpmux
import (
"bytes"
"context"
"fmt"
"io"
"net"
"strings"
"sync"
logging "github.com/ipfs/go-log/v2"
pool "github.com/libp2p/go-buffer-pool"
"github.com/pion/ice/v2"
"github.com/pion/stun"
)
var log = logging.Logger("webrtc-udpmux")
// ReceiveBufSize is the size of the buffer used to receive packets from the PacketConn.
// It is fine for this number to be higher than the actual path MTU as this value is not
// used to decide the packet size on the write path.
const ReceiveBufSize = 1500
type Candidate struct {
Ufrag string
Addr *net.UDPAddr
}
// UDPMux multiplexes multiple ICE connections over a single net.PacketConn,
// generally a UDP socket.
//
// The connections are indexed by (ufrag, IP address family) and by remote
// address from which the connection has received valid STUN/RTC packets.
//
// When a new packet is received on the underlying net.PacketConn, we
// first check the address map to see if there is a connection associated with the
// remote address:
// If found, we pass the packet to that connection.
// Otherwise, we check to see if the packet is a STUN packet.
// If it is, we read the ufrag from the STUN packet and use it to check if there
// is a connection associated with the (ufrag, IP address family) pair.
// If found we add the association to the address map.
type UDPMux struct {
socket net.PacketConn
queue chan Candidate
mx sync.Mutex
// ufragMap allows us to multiplex incoming STUN packets based on ufrag
ufragMap map[ufragConnKey]*muxedConnection
// addrMap allows us to correctly direct incoming packets after the connection
// is established and ufrag isn't available on all packets
addrMap map[string]*muxedConnection
// ufragAddrMap allows cleaning up all addresses from the addrMap once the connection is closed
// During the ICE connectivity checks, the same ufrag might be used on multiple addresses.
ufragAddrMap map[ufragConnKey][]net.Addr
// the context controls the lifecycle of the mux
wg sync.WaitGroup
ctx context.Context
cancel context.CancelFunc
}
var _ ice.UDPMux = &UDPMux{}
func NewUDPMux(socket net.PacketConn) *UDPMux {
ctx, cancel := context.WithCancel(context.Background())
mux := &UDPMux{
ctx: ctx,
cancel: cancel,
socket: socket,
ufragMap: make(map[ufragConnKey]*muxedConnection),
addrMap: make(map[string]*muxedConnection),
ufragAddrMap: make(map[ufragConnKey][]net.Addr),
queue: make(chan Candidate, 32),
}
return mux
}
func (mux *UDPMux) Start() {
mux.wg.Add(1)
go func() {
defer mux.wg.Done()
mux.readLoop()
}()
}
// GetListenAddresses implements ice.UDPMux
func (mux *UDPMux) GetListenAddresses() []net.Addr {
return []net.Addr{mux.socket.LocalAddr()}
}
// GetConn implements ice.UDPMux
// It creates a net.PacketConn for a given ufrag if an existing one cannot be found.
// We differentiate IPv4 and IPv6 addresses, since a remote is can be reachable at multiple different
// UDP addresses of the same IP address family (eg. server-reflexive addresses and peer-reflexive addresses).
func (mux *UDPMux) GetConn(ufrag string, addr net.Addr) (net.PacketConn, error) {
a, ok := addr.(*net.UDPAddr)
if !ok {
return nil, fmt.Errorf("unexpected address type: %T", addr)
}
select {
case <-mux.ctx.Done():
return nil, io.ErrClosedPipe
default:
isIPv6 := ok && a.IP.To4() == nil
_, conn := mux.getOrCreateConn(ufrag, isIPv6, mux, addr)
return conn, nil
}
}
// Close implements ice.UDPMux
func (mux *UDPMux) Close() error {
select {
case <-mux.ctx.Done():
return nil
default:
}
mux.cancel()
mux.socket.Close()
mux.wg.Wait()
return nil
}
// writeTo writes a packet to the underlying net.PacketConn
func (mux *UDPMux) writeTo(buf []byte, addr net.Addr) (int, error) {
return mux.socket.WriteTo(buf, addr)
}
func (mux *UDPMux) readLoop() {
for {
select {
case <-mux.ctx.Done():
return
default:
}
buf := pool.Get(ReceiveBufSize)
n, addr, err := mux.socket.ReadFrom(buf)
if err != nil {
if strings.Contains(err.Error(), "use of closed network connection") {
log.Debugf("readLoop exiting: socket %s closed", mux.socket.LocalAddr())
} else {
log.Errorf("error reading from socket %s: %v", mux.socket.LocalAddr(), err)
}
pool.Put(buf)
return
}
buf = buf[:n]
if processed := mux.processPacket(buf, addr); !processed {
pool.Put(buf)
}
}
}
func (mux *UDPMux) processPacket(buf []byte, addr net.Addr) (processed bool) {
udpAddr, ok := addr.(*net.UDPAddr)
if !ok {
log.Errorf("received a non-UDP address: %s", addr)
return false
}
isIPv6 := udpAddr.IP.To4() == nil
// Connections are indexed by remote address. We first
// check if the remote address has a connection associated
// with it. If yes, we push the received packet to the connection
mux.mx.Lock()
conn, ok := mux.addrMap[addr.String()]
mux.mx.Unlock()
if ok {
if err := conn.Push(buf, addr); err != nil {
log.Debugf("could not push packet: %v", err)
return false
}
return true
}
if !stun.IsMessage(buf) {
log.Debug("incoming message is not a STUN message")
return false
}
msg := &stun.Message{Raw: buf}
if err := msg.Decode(); err != nil {
log.Debugf("failed to decode STUN message: %s", err)
return false
}
if msg.Type != stun.BindingRequest {
log.Debugf("incoming message should be a STUN binding request, got %s", msg.Type)
return false
}
ufrag, err := ufragFromSTUNMessage(msg)
if err != nil {
log.Debugf("could not find STUN username: %s", err)
return false
}
connCreated, conn := mux.getOrCreateConn(ufrag, isIPv6, mux, udpAddr)
if connCreated {
select {
case mux.queue <- Candidate{Addr: udpAddr, Ufrag: ufrag}:
default:
log.Debugw("queue full, dropping incoming candidate", "ufrag", ufrag, "addr", udpAddr)
conn.Close()
return false
}
}
if err := conn.Push(buf, addr); err != nil {
log.Debugf("could not push packet: %v", err)
return false
}
return true
}
func (mux *UDPMux) Accept(ctx context.Context) (Candidate, error) {
select {
case c := <-mux.queue:
return c, nil
case <-ctx.Done():
return Candidate{}, ctx.Err()
case <-mux.ctx.Done():
return Candidate{}, mux.ctx.Err()
}
}
type ufragConnKey struct {
ufrag string
isIPv6 bool
}
// ufragFromSTUNMessage returns the local or ufrag
// from the STUN username attribute. Local ufrag is the ufrag of the
// peer which initiated the connectivity check, e.g in a connectivity
// check from A to B, the username attribute will be B_ufrag:A_ufrag
// with the local ufrag value being A_ufrag. In case of ice-lite, the
// localUfrag value will always be the remote peer's ufrag since ICE-lite
// implementations do not generate connectivity checks. In our specific
// case, since the local and remote ufrag is equal, we can return
// either value.
func ufragFromSTUNMessage(msg *stun.Message) (string, error) {
attr, err := msg.Get(stun.AttrUsername)
if err != nil {
return "", err
}
index := bytes.Index(attr, []byte{':'})
if index == -1 {
return "", fmt.Errorf("invalid STUN username attribute")
}
return string(attr[index+1:]), nil
}
// RemoveConnByUfrag removes the connection associated with the ufrag and all the
// addresses associated with that connection. This method is called by pion when
// a peerconnection is closed.
func (mux *UDPMux) RemoveConnByUfrag(ufrag string) {
if ufrag == "" {
return
}
mux.mx.Lock()
defer mux.mx.Unlock()
for _, isIPv6 := range [...]bool{true, false} {
key := ufragConnKey{ufrag: ufrag, isIPv6: isIPv6}
if _, ok := mux.ufragMap[key]; ok {
delete(mux.ufragMap, key)
for _, addr := range mux.ufragAddrMap[key] {
delete(mux.addrMap, addr.String())
}
delete(mux.ufragAddrMap, key)
}
}
}
func (mux *UDPMux) getOrCreateConn(ufrag string, isIPv6 bool, _ *UDPMux, addr net.Addr) (created bool, _ *muxedConnection) {
key := ufragConnKey{ufrag: ufrag, isIPv6: isIPv6}
mux.mx.Lock()
defer mux.mx.Unlock()
if conn, ok := mux.ufragMap[key]; ok {
mux.addrMap[addr.String()] = conn
mux.ufragAddrMap[key] = append(mux.ufragAddrMap[key], addr)
return false, conn
}
conn := newMuxedConnection(mux, func() { mux.RemoveConnByUfrag(ufrag) })
mux.ufragMap[key] = conn
mux.addrMap[addr.String()] = conn
mux.ufragAddrMap[key] = append(mux.ufragAddrMap[key], addr)
return true, conn
}