udp_packetconn.go

     1  // Copyright 2024 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package quic
     6  
     7  import (
     8  	"net"
     9  	"net/netip"
    10  )
    11  
    12  // netPacketConn is a packetConn implementation wrapping a net.PacketConn.
    13  //
    14  // This is mostly useful for tests, since PacketConn doesn't provide access to
    15  // important features such as identifying the local address packets were received on.
    16  type netPacketConn struct {
    17  	c         net.PacketConn
    18  	localAddr netip.AddrPort
    19  }
    20  
    21  func newNetPacketConn(pc net.PacketConn) (*netPacketConn, error) {
    22  	addr, err := addrPortFromAddr(pc.LocalAddr())
    23  	if err != nil {
    24  		return nil, err
    25  	}
    26  	return &netPacketConn{
    27  		c:         pc,
    28  		localAddr: addr,
    29  	}, nil
    30  }
    31  
    32  func (c *netPacketConn) Close() error {
    33  	return c.c.Close()
    34  }
    35  
    36  func (c *netPacketConn) LocalAddr() netip.AddrPort {
    37  	return c.localAddr
    38  }
    39  
    40  func (c *netPacketConn) Read(f func(*datagram)) {
    41  	for {
    42  		dgram := newDatagram()
    43  		n, peerAddr, err := c.c.ReadFrom(dgram.b)
    44  		if err != nil {
    45  			return
    46  		}
    47  		dgram.peerAddr, err = addrPortFromAddr(peerAddr)
    48  		if err != nil {
    49  			continue
    50  		}
    51  		dgram.b = dgram.b[:n]
    52  		f(dgram)
    53  	}
    54  }
    55  
    56  func (c *netPacketConn) Write(dgram datagram) error {
    57  	_, err := c.c.WriteTo(dgram.b, net.UDPAddrFromAddrPort(dgram.peerAddr))
    58  	return err
    59  }
    60  
    61  func addrPortFromAddr(addr net.Addr) (netip.AddrPort, error) {
    62  	switch a := addr.(type) {
    63  	case *net.UDPAddr:
    64  		return a.AddrPort(), nil
    65  	}
    66  	return netip.ParseAddrPort(addr.String())
    67  }
    68
View as plain text