diff --git a/connectctl/main.go b/connectctl/main.go index 7261e47..b7d2c6e 100644 --- a/connectctl/main.go +++ b/connectctl/main.go @@ -699,18 +699,18 @@ func sink(opts docopt.Opts) { } type Receive struct { - source connect.TransferPath - frames []*protocol.Frame - provideMode protocol.ProvideMode + source connect.TransferPath + frames []*protocol.Frame + peer connect.Peer } receives := make(chan *Receive) - client.AddReceiveCallback(func(source connect.TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { + client.AddReceiveCallback(func(source connect.TransferPath, frames []*protocol.Frame, peer connect.Peer) { receives <- &Receive{ - source: source, - frames: frames, - provideMode: provideMode, + source: source, + frames: frames, + peer: peer, } }) @@ -718,7 +718,7 @@ func sink(opts docopt.Opts) { for i := 0; messageCount < 0 || i < messageCount; i += 1 { select { case receive := <-receives: - fmt.Printf("[%s %s] %s\n", receive.source, receive.provideMode, receive.frames) + fmt.Printf("[%s %s] %s\n", receive.source, receive.peer.ProvideMode, receive.frames) } } } diff --git a/egress.go b/egress.go new file mode 100644 index 0000000..f634bed --- /dev/null +++ b/egress.go @@ -0,0 +1,83 @@ +package connect + +import ( + "net" + "sync/atomic" + "syscall" +) + +// Socket self-exclusion (R1 for the Windows VPN service). +// +// When this process itself provides a VPN tunnel (a DeviceLocal writing a tun +// adapter), its own outbound sockets must not route back into that tunnel, or +// the platform/provider connections loop. macOS network extensions and Android +// VpnService.protect solve this at the OS layer; Windows has no per-application +// exclusion, so connect binds its outbound sockets to the physical interface +// via IP_UNICAST_IF / IPV6_UNICAST_IF (the wireguard-windows pattern). A socket +// with a forced egress interface ignores the tun's default route. +// +// The binding is a no-op unless an egress index is set, and a no-op on every +// non-Windows platform (see egress_other.go), so this is inert for the mobile +// and desktop app builds. + +var egressInterfaceIndex4 atomic.Uint32 +var egressInterfaceIndex6 atomic.Uint32 + +// SetEgressInterfaceIndex binds connect's outbound sockets to the given +// physical interface indices (IPv4 and IPv6). Pass 0 for a family to leave that +// family unbound. Safe to call repeatedly, e.g. from a network-change handler. +// Only takes effect on Windows. +func SetEgressInterfaceIndex(index4 uint32, index6 uint32) { + egressInterfaceIndex4.Store(index4) + egressInterfaceIndex6.Store(index6) +} + +// EgressInterfaceIndex returns the current (index4, index6). Zero means unset. +func EgressInterfaceIndex() (uint32, uint32) { + return egressInterfaceIndex4.Load(), egressInterfaceIndex6.Load() +} + +// egressControl is a net.Dialer / net.ListenConfig Control callback that applies +// the egress interface binding to a socket at creation time. +func egressControl(_ string, _ string, c syscall.RawConn) error { + index4, index6 := EgressInterfaceIndex() + if index4 == 0 && index6 == 0 { + return nil + } + var innerErr error + err := c.Control(func(fd uintptr) { + innerErr = applyEgressInterface(fd, index4, index6) + }) + if err != nil { + return err + } + return innerErr +} + +// egressDialer installs the egress control on d, chaining any existing Control. +func egressDialer(d *net.Dialer) *net.Dialer { + prev := d.Control + d.Control = func(network string, address string, c syscall.RawConn) error { + if prev != nil { + if err := prev(network, address, c); err != nil { + return err + } + } + return egressControl(network, address, c) + } + return d +} + +// applyEgress binds an already-created socket-backed conn (e.g. *net.UDPConn) to +// the egress interface. A no-op when no egress index is set or off Windows. +func applyEgress(conn syscall.Conn) error { + index4, index6 := EgressInterfaceIndex() + if index4 == 0 && index6 == 0 { + return nil + } + raw, err := conn.SyscallConn() + if err != nil { + return err + } + return egressControl("", "", raw) +} diff --git a/egress_net.go b/egress_net.go new file mode 100644 index 0000000..0d780f5 --- /dev/null +++ b/egress_net.go @@ -0,0 +1,50 @@ +//go:build !js + +package connect + +import ( + "net" + "syscall" + + "github.com/pion/transport/v4" + "github.com/pion/transport/v4/stdnet" +) + +// egressNet is a pion transport.Net that binds the UDP sockets pion creates for +// ICE to the egress interface, so p2p (webrtc) traffic does not loop back into +// the tunnel this process provides (R1). It embeds the standard net and only +// overrides socket creation; a no-op unless an egress index is set / off +// Windows. +type egressNet struct { + transport.Net +} + +func newEgressNet() (transport.Net, error) { + base, err := stdnet.NewNet() + if err != nil { + return nil, err + } + return &egressNet{Net: base}, nil +} + +func (self *egressNet) ListenUDP(network string, locAddr *net.UDPAddr) (transport.UDPConn, error) { + conn, err := self.Net.ListenUDP(network, locAddr) + if err != nil { + return nil, err + } + if sc, ok := conn.(syscall.Conn); ok { + _ = applyEgress(sc) + } + return conn, nil +} + +func (self *egressNet) ListenPacket(network string, address string) (net.PacketConn, error) { + conn, err := self.Net.ListenPacket(network, address) + if err != nil { + return nil, err + } + if sc, ok := conn.(syscall.Conn); ok { + _ = applyEgress(sc) + } + return conn, nil +} diff --git a/egress_other.go b/egress_other.go new file mode 100644 index 0000000..9a06e4c --- /dev/null +++ b/egress_other.go @@ -0,0 +1,11 @@ +//go:build !windows + +package connect + +// On non-Windows platforms self-exclusion is handled at the OS layer: macOS +// network extensions bypass their own tunnel automatically, and Android uses +// VpnService.protect / addDisallowedApplication. The forced-egress binding is a +// no-op so the egress control is inert in the mobile and desktop app builds. +func applyEgressInterface(_ uintptr, _ uint32, _ uint32) error { + return nil +} diff --git a/egress_windows.go b/egress_windows.go new file mode 100644 index 0000000..b51453d --- /dev/null +++ b/egress_windows.go @@ -0,0 +1,37 @@ +//go:build windows + +package connect + +import ( + "golang.org/x/sys/windows" +) + +// IP_UNICAST_IF (IPPROTO_IP, option 31) forces the outgoing interface for +// unicast, overriding the route table. The interface index must be passed in +// NETWORK byte order. IPV6_UNICAST_IF (IPPROTO_IPV6, option 31) does the same +// for IPv6 but takes the index in HOST byte order. This asymmetry matches +// wireguard-windows' bindSocketRoute. +const ( + ipUnicastIf = 31 + ipv6UnicastIf = 31 +) + +func htonl(x uint32) uint32 { + return (x&0x000000ff)<<24 | (x&0x0000ff00)<<8 | (x&0x00ff0000)>>8 | (x&0xff000000)>>24 +} + +// applyEgressInterface sets the forced egress interface on the socket. Both +// families are set best-effort: the option for the wrong family on a +// single-family socket fails harmlessly, so per-option errors are ignored (as +// in wireguard-windows). This guarantees the call never breaks connection +// setup even when only one family applies. +func applyEgressInterface(fd uintptr, index4 uint32, index6 uint32) error { + h := windows.Handle(fd) + if index4 != 0 { + _ = windows.SetsockoptInt(h, windows.IPPROTO_IP, ipUnicastIf, int(htonl(index4))) + } + if index6 != 0 { + _ = windows.SetsockoptInt(h, windows.IPPROTO_IPV6, ipv6UnicastIf, int(index6)) + } + return nil +} diff --git a/ip.go b/ip.go index 597465f..c14036f 100644 --- a/ip.go +++ b/ip.go @@ -19,14 +19,10 @@ import ( "sync/atomic" "time" - "github.com/google/gopacket" - "github.com/google/gopacket/layers" - "golang.org/x/exp/maps" // "google.golang.org/protobuf/proto" - "github.com/urnetwork/connect/protocol" ) @@ -50,6 +46,228 @@ const TcpHeaderSizeWithoutExtensions = 20 const debugVerifyHeaders = false +type IPProtocol uint8 + +const ( + IP_PROTOCOL_UDP IPProtocol = 17 + IP_PROTOCOL_TCP IPProtocol = 6 +) + +type UDPPort uint16 +type TCPPort uint16 + +type parsedUdp struct { + sourceIp net.IP + destinationIp net.IP + sourcePort UDPPort + destinationPort UDPPort + payload []byte +} + +type parsedTcp struct { + sourceIp net.IP + destinationIp net.IP + sourcePort TCPPort + destinationPort TCPPort + fin bool + syn bool + rst bool + psh bool + ack bool + seq uint32 + ackNumber uint32 + windowSize uint16 + options []byte + payload []byte +} + +func (self *parsedTcp) flagsString() string { + flags := []string{} + if self.fin { + flags = append(flags, "FIN") + } + if self.syn { + flags = append(flags, "SYN") + } + if self.rst { + flags = append(flags, "RST") + } + if self.psh { + flags = append(flags, "PSH") + } + if self.ack { + flags = append(flags, "ACK") + } + return strings.Join(flags, ", ") +} + +func parseIpv4(ipPacket []byte) (ipProtocol IPProtocol, sourceIp net.IP, destinationIp net.IP, transport []byte, ok bool) { + if len(ipPacket) < Ipv4HeaderSizeWithoutExtensions { + return + } + headerByteCount := int(ipPacket[0]&0xf) * 4 + totalByteCount := int(binary.BigEndian.Uint16(ipPacket[2:4])) + if headerByteCount < Ipv4HeaderSizeWithoutExtensions || totalByteCount < headerByteCount || len(ipPacket) < totalByteCount { + return + } + ipProtocol = IPProtocol(ipPacket[9]) + sourceIp = net.IP(ipPacket[12:16]) + destinationIp = net.IP(ipPacket[16:20]) + transport = ipPacket[headerByteCount:totalByteCount] + ok = true + return +} + +func parseIpv6(ipPacket []byte) (ipProtocol IPProtocol, sourceIp net.IP, destinationIp net.IP, transport []byte, ok bool) { + if len(ipPacket) < Ipv6HeaderSize { + return + } + payloadByteCount := int(binary.BigEndian.Uint16(ipPacket[4:6])) + if len(ipPacket) < Ipv6HeaderSize+payloadByteCount { + return + } + ipProtocol = IPProtocol(ipPacket[6]) + sourceIp = net.IP(ipPacket[8:24]) + destinationIp = net.IP(ipPacket[24:40]) + transport = ipPacket[Ipv6HeaderSize : Ipv6HeaderSize+payloadByteCount] + ok = true + return +} + +func parseUdpPacket(sourceIp net.IP, destinationIp net.IP, transport []byte, udp *parsedUdp) bool { + if len(transport) < UdpHeaderSize { + return false + } + udpByteCount := int(binary.BigEndian.Uint16(transport[4:6])) + if udpByteCount < UdpHeaderSize || len(transport) < udpByteCount { + return false + } + udp.sourceIp = sourceIp + udp.destinationIp = destinationIp + udp.sourcePort = UDPPort(binary.BigEndian.Uint16(transport[0:2])) + udp.destinationPort = UDPPort(binary.BigEndian.Uint16(transport[2:4])) + udp.payload = transport[UdpHeaderSize:udpByteCount] + return true +} + +func parseTcpPacket(sourceIp net.IP, destinationIp net.IP, transport []byte, tcp *parsedTcp) bool { + if len(transport) < TcpHeaderSizeWithoutExtensions { + return false + } + headerByteCount := int(transport[12]>>4) * 4 + if headerByteCount < TcpHeaderSizeWithoutExtensions || len(transport) < headerByteCount { + return false + } + flags := transport[13] + tcp.sourceIp = sourceIp + tcp.destinationIp = destinationIp + tcp.sourcePort = TCPPort(binary.BigEndian.Uint16(transport[0:2])) + tcp.destinationPort = TCPPort(binary.BigEndian.Uint16(transport[2:4])) + tcp.seq = binary.BigEndian.Uint32(transport[4:8]) + tcp.ackNumber = binary.BigEndian.Uint32(transport[8:12]) + tcp.fin = (flags & 0x01) != 0 + tcp.syn = (flags & 0x02) != 0 + tcp.rst = (flags & 0x04) != 0 + tcp.psh = (flags & 0x08) != 0 + tcp.ack = (flags & 0x10) != 0 + tcp.windowSize = binary.BigEndian.Uint16(transport[14:16]) + tcp.options = transport[TcpHeaderSizeWithoutExtensions:headerByteCount] + tcp.payload = transport[headerByteCount:] + return true +} + +func ParseTcpWindowScaleOpts(opts []byte) (bool, uint32) { + for i := 0; i < len(opts); { + kind := opts[i] + if kind == 0 { + break + } + if kind == 1 { + i += 1 + continue + } + if i+1 >= len(opts) { + break + } + length := opts[i+1] + if length < 2 || i+int(length) > len(opts) { + break + } + if kind == 3 && length >= 3 { + shift := uint32(opts[i+2]) + if 14 < shift { + shift = 14 + } + return true, shift + } + i += int(length) + } + return false, 0 +} + +const ( + tcpFlagFin = byte(0x01) + tcpFlagSyn = byte(0x02) + tcpFlagRst = byte(0x04) + tcpFlagPsh = byte(0x08) + tcpFlagAck = byte(0x10) +) + +func checksumAdd(sum uint32, b []byte) uint32 { + i := 0 + for ; i+1 < len(b); i += 2 { + sum += uint32(binary.BigEndian.Uint16(b[i : i+2])) + } + if i < len(b) { + sum += uint32(b[i]) << 8 + } + return sum +} + +func checksumFinish(sum uint32) uint16 { + for 0xffff < sum { + sum = (sum & 0xffff) + (sum >> 16) + } + return ^uint16(sum) +} + +func transportChecksum(protocol IPProtocol, packetSourceIp net.IP, packetDestinationIp net.IP, transport []byte) uint16 { + sum := checksumAdd(0, packetSourceIp) + sum = checksumAdd(sum, packetDestinationIp) + sum += uint32(protocol) + sum += uint32(len(transport)) + return checksumFinish(checksumAdd(sum, transport)) +} + +func writeIpv4Header(packet []byte, ipProtocol IPProtocol, packetSourceIp net.IP, packetDestinationIp net.IP) { + packet[0] = 0x45 + packet[1] = 0 + binary.BigEndian.PutUint16(packet[2:4], uint16(len(packet))) + packet[4] = 0 + packet[5] = 0 + packet[6] = 0 + packet[7] = 0 + packet[8] = 64 + packet[9] = byte(ipProtocol) + packet[10] = 0 + packet[11] = 0 + copy(packet[12:16], packetSourceIp) + copy(packet[16:20], packetDestinationIp) + binary.BigEndian.PutUint16(packet[10:12], checksumFinish(checksumAdd(0, packet[0:Ipv4HeaderSizeWithoutExtensions]))) +} + +func writeIpv6Header(packet []byte, ipProtocol IPProtocol, packetSourceIp net.IP, packetDestinationIp net.IP) { + packet[0] = 0x60 + packet[1] = 0 + packet[2] = 0 + packet[3] = 0 + binary.BigEndian.PutUint16(packet[4:6], uint16(len(packet)-Ipv6HeaderSize)) + packet[6] = byte(ipProtocol) + packet[7] = 64 + copy(packet[8:24], packetSourceIp) + copy(packet[24:40], packetDestinationIp) +} + // send from a raw socket // note `ipProtocol` is not supplied. The implementation must do a packet inspection to determine protocol // `provideMode` is the relationship between the source and this device @@ -342,22 +560,32 @@ func (self *LocalUserNat) runShard(shard int) { return case sendPacket := <-shardCh: ipPacket := sendPacket.packet + if len(ipPacket) == 0 { + MessagePoolReturn(ipPacket) + continue + } + var udpPacket parsedUdp + var tcpPacket parsedTcp ipVersion := uint8(ipPacket[0]) >> 4 switch ipVersion { case 4: - ipv4 := layers.IPv4{} - ipv4.DecodeFromBytes(ipPacket, gopacket.NilDecodeFeedback) - switch ipv4.Protocol { - case layers.IPProtocolUDP: - udp := layers.UDP{} - udp.DecodeFromBytes(ipv4.Payload, gopacket.NilDecodeFeedback) + ipProtocol, sourceIp, destinationIp, transport, ok := parseIpv4(ipPacket) + if !ok { + MessagePoolReturn(ipPacket) + continue + } + switch ipProtocol { + case IP_PROTOCOL_UDP: + if !parseUdpPacket(sourceIp, destinationIp, transport, &udpPacket) { + MessagePoolReturn(ipPacket) + continue + } c := func() bool { success, err := udp4Buffer.send( sendPacket.source, sendPacket.provideMode, - &ipv4, - &udp, + &udpPacket, self.settings.UdpBufferSettings.WriteTimeout, ipPacket, ) @@ -371,16 +599,17 @@ func (self *LocalUserNat) runShard(shard int) { } else { c() } - case layers.IPProtocolTCP: - tcp := layers.TCP{} - tcp.DecodeFromBytes(ipv4.Payload, gopacket.NilDecodeFeedback) + case IP_PROTOCOL_TCP: + if !parseTcpPacket(sourceIp, destinationIp, transport, &tcpPacket) { + MessagePoolReturn(ipPacket) + continue + } c := func() bool { success, err := tcp4Buffer.send( sendPacket.source, sendPacket.provideMode, - &ipv4, - &tcp, + &tcpPacket, self.settings.TcpBufferSettings.WriteTimeout, ipPacket, ) @@ -399,19 +628,23 @@ func (self *LocalUserNat) runShard(shard int) { MessagePoolReturn(ipPacket) } case 6: - ipv6 := layers.IPv6{} - ipv6.DecodeFromBytes(ipPacket, gopacket.NilDecodeFeedback) - switch ipv6.NextHeader { - case layers.IPProtocolUDP: - udp := layers.UDP{} - udp.DecodeFromBytes(ipv6.Payload, gopacket.NilDecodeFeedback) + ipProtocol, sourceIp, destinationIp, transport, ok := parseIpv6(ipPacket) + if !ok { + MessagePoolReturn(ipPacket) + continue + } + switch ipProtocol { + case IP_PROTOCOL_UDP: + if !parseUdpPacket(sourceIp, destinationIp, transport, &udpPacket) { + MessagePoolReturn(ipPacket) + continue + } c := func() bool { success, err := udp6Buffer.send( sendPacket.source, sendPacket.provideMode, - &ipv6, - &udp, + &udpPacket, self.settings.UdpBufferSettings.WriteTimeout, ipPacket, ) @@ -425,16 +658,17 @@ func (self *LocalUserNat) runShard(shard int) { } else { c() } - case layers.IPProtocolTCP: - tcp := layers.TCP{} - tcp.DecodeFromBytes(ipv6.Payload, gopacket.NilDecodeFeedback) + case IP_PROTOCOL_TCP: + if !parseTcpPacket(sourceIp, destinationIp, transport, &tcpPacket) { + MessagePoolReturn(ipPacket) + continue + } c := func() bool { success, err := tcp6Buffer.send( sendPacket.source, sendPacket.provideMode, - &ipv6, - &tcp, + &tcpPacket, self.settings.TcpBufferSettings.WriteTimeout, ipPacket, ) @@ -452,6 +686,9 @@ func (self *LocalUserNat) runShard(shard int) { // no support for this protocol, drop MessagePoolReturn(ipPacket) } + default: + // no support for this ip version, drop + MessagePoolReturn(ipPacket) } } } @@ -534,17 +771,17 @@ func NewUdp4Buffer(ctx context.Context, receiveCallback ReceivePacketFunction, } func (self *Udp4Buffer) send(source TransferPath, provideMode protocol.ProvideMode, - ipv4 *layers.IPv4, udp *layers.UDP, timeout time.Duration, ipPacket []byte) (bool, error) { + udp *parsedUdp, timeout time.Duration, ipPacket []byte) (bool, error) { bufferId := NewBufferId4( source, - ipv4.SrcIP, int(udp.SrcPort), - ipv4.DstIP, int(udp.DstPort), + udp.sourceIp, int(udp.sourcePort), + udp.destinationIp, int(udp.destinationPort), ) return self.udpSend( bufferId, - ipv4.SrcIP, - ipv4.DstIP, + udp.sourceIp, + udp.destinationIp, source, provideMode, 4, @@ -566,17 +803,17 @@ func NewUdp6Buffer(ctx context.Context, receiveCallback ReceivePacketFunction, } func (self *Udp6Buffer) send(source TransferPath, provideMode protocol.ProvideMode, - ipv6 *layers.IPv6, udp *layers.UDP, timeout time.Duration, ipPacket []byte) (bool, error) { + udp *parsedUdp, timeout time.Duration, ipPacket []byte) (bool, error) { bufferId := NewBufferId6( source, - ipv6.SrcIP, int(udp.SrcPort), - ipv6.DstIP, int(udp.DstPort), + udp.sourceIp, int(udp.sourcePort), + udp.destinationIp, int(udp.destinationPort), ) return self.udpSend( bufferId, - ipv6.SrcIP, - ipv6.DstIP, + udp.sourceIp, + udp.destinationIp, source, provideMode, 6, @@ -623,7 +860,7 @@ func (self *UdpBuffer[BufferId]) udpSend( source TransferPath, provideMode protocol.ProvideMode, ipVersion int, - udp *layers.UDP, + udp *parsedUdp, timeout time.Duration, ipPacket []byte, ) (bool, error) { @@ -679,9 +916,9 @@ func (self *UdpBuffer[BufferId]) udpSend( provideMode, ipVersion, sourceIpCopy, - udp.SrcPort, + udp.sourcePort, destinationIpCopy, - udp.DstPort, + udp.destinationPort, self.udpBufferSettings, ) self.sequences[bufferId] = sequence @@ -751,8 +988,8 @@ func NewUdpSequence(ctx context.Context, receiveCallback ReceivePacketFunction, source TransferPath, provideMode protocol.ProvideMode, ipVersion int, - sourceIp net.IP, sourcePort layers.UDPPort, - destinationIp net.IP, destinationPort layers.UDPPort, + sourceIp net.IP, sourcePort UDPPort, + destinationIp net.IP, destinationPort UDPPort, udpBufferSettings *UdpBufferSettings) *UdpSequence { cancelCtx, cancel := context.WithCancel(ctx) // e2e-pqe merge: upstream now inlines StreamState in the return below; keep @@ -774,7 +1011,7 @@ func NewUdpSequence(ctx context.Context, receiveCallback ReceivePacketFunction, sourcePort: sourcePort, destinationIp: destinationIp, destinationPort: destinationPort, - buffer: gopacket.NewSerializeBufferExpectedSize(128, 2048), + userLimited: userLimited{ lastActivityTime: time.Now(), }, @@ -1087,7 +1324,7 @@ func (self *UdpSequence) Run() { if !ok { return } - payload := sendItem.udp.Payload + payload := sendItem.udp.payload if 0 < len(payload) { writePayload := writePayload{ @@ -1134,7 +1371,7 @@ func (self *UdpSequence) Close() { type UdpSendItem struct { source TransferPath provideMode protocol.ProvideMode - udp *layers.UDP + udp *parsedUdp ipPacket []byte } @@ -1143,10 +1380,10 @@ type StreamState struct { provideMode protocol.ProvideMode ipVersion int sourceIp net.IP - sourcePort layers.UDPPort + sourcePort UDPPort destinationIp net.IP - destinationPort layers.UDPPort - buffer gopacket.SerializeBuffer + destinationPort UDPPort + userLimited // cached immutable ip path for this stream (see IpPath). The stream @@ -1180,92 +1417,79 @@ func (self *StreamState) IpPath() *IpPath { // this must only be called from one goroutine // this is called from the writer only and does not need to syncrhronize with the reader state func (self *StreamState) DataPackets(payload []byte, n int, mtu int) ([][]byte, error) { - headerSize := 0 - var ip gopacket.NetworkLayer + var headerByteCount int switch self.ipVersion { case 4: - ip = &layers.IPv4{ - Version: 4, - TTL: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - Protocol: layers.IPProtocolUDP, - } - headerSize += Ipv4HeaderSizeWithoutExtensions + headerByteCount = Ipv4HeaderSizeWithoutExtensions + UdpHeaderSize case 6: - ip = &layers.IPv6{ - Version: 6, - HopLimit: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - NextHeader: layers.IPProtocolUDP, - } - headerSize += Ipv6HeaderSize + headerByteCount = Ipv6HeaderSize + UdpHeaderSize } - udp := layers.UDP{ - SrcPort: self.destinationPort, - DstPort: self.sourcePort, + if mtu <= headerByteCount { + return nil, fmt.Errorf("mtu %d is too small for IP+UDP headers (%d bytes)", mtu, headerByteCount) } - udp.SetNetworkLayerForChecksum(ip) - headerSize += UdpHeaderSize - - options := gopacket.SerializeOptions{ - ComputeChecksums: true, - FixLengths: true, + packetByteCount := mtu - headerByteCount + if n <= packetByteCount { + self.singleDataPacket[0] = self.udpPacket(payload[0:n]) + return self.singleDataPacket[:], nil } - - if debugVerifyHeaders { - buffer := gopacket.NewSerializeBufferExpectedSize(headerSize, 0) - err := gopacket.SerializeLayers(buffer, options, - ip.(gopacket.SerializableLayer), - &udp, - ) - if err != nil { - return nil, err - } - packetHeaders := buffer.Bytes() - if headerSize != len(packetHeaders) { - return nil, errors.New(fmt.Sprintf("Header check failed %d <> %d", headerSize, len(packetHeaders))) - } + packets := make([][]byte, 0, (n+packetByteCount-1)/packetByteCount) + for i := 0; i < n; { + j := min(i+packetByteCount, n) + packets = append(packets, self.udpPacket(payload[i:j])) + i = j } + return packets, nil +} - if headerSize+n <= mtu { - self.buffer.Clear() - err := gopacket.SerializeLayers(self.buffer, options, - ip.(gopacket.SerializableLayer), - &udp, - gopacket.Payload(payload[0:n]), - ) - if err != nil { - return nil, err - } - packet := MessagePoolCopy(self.buffer.Bytes()) - // reuse the single-packet backing for the common unfragmented case - // (see singleDataPacket); the result is consumed before the next call. - self.singleDataPacket[0] = packet - return self.singleDataPacket[:], nil - } else { - // fragment - packetSize := mtu - headerSize - packets := make([][]byte, 0, (n+packetSize)/packetSize) - for i := 0; i < n; { - self.buffer.Clear() - j := min(i+packetSize, n) - err := gopacket.SerializeLayers(self.buffer, options, - ip.(gopacket.SerializableLayer), - &udp, - gopacket.Payload(payload[i:j]), - ) - if err != nil { - return nil, err - } - packet := MessagePoolCopy(self.buffer.Bytes()) - packets = append(packets, packet) - i = j - } - return packets, nil +// builds a udp packet from the stream destination to the stream source +// into a single pool buffer +func (self *StreamState) udpPacket(payload []byte) []byte { + var ipHeaderByteCount int + switch self.ipVersion { + case 4: + ipHeaderByteCount = Ipv4HeaderSizeWithoutExtensions + case 6: + ipHeaderByteCount = Ipv6HeaderSize } + + packet := MessagePoolGet(ipHeaderByteCount + UdpHeaderSize + len(payload)) + switch self.ipVersion { + case 4: + writeIpv4Header(packet, IP_PROTOCOL_UDP, self.destinationIp, self.sourceIp) + case 6: + writeIpv6Header(packet, IP_PROTOCOL_UDP, self.destinationIp, self.sourceIp) + } + + udp := packet[ipHeaderByteCount:] + binary.BigEndian.PutUint16(udp[0:2], uint16(self.destinationPort)) + binary.BigEndian.PutUint16(udp[2:4], uint16(self.sourcePort)) + binary.BigEndian.PutUint16(udp[4:6], uint16(UdpHeaderSize+len(payload))) + udp[6] = 0 + udp[7] = 0 + copy(udp[UdpHeaderSize:], payload) + checksum := transportChecksum(IP_PROTOCOL_UDP, self.destinationIp, self.sourceIp, udp) + if checksum == 0 { + checksum = 0xffff + } + binary.BigEndian.PutUint16(udp[6:8], checksum) + return packet +} + +func writeTcpHeader(tcp []byte, sourcePort, destinationPort uint16, seq, ackNumber uint32, flags uint8, windowSize uint16, options []byte) { + binary.BigEndian.PutUint16(tcp[0:2], sourcePort) + binary.BigEndian.PutUint16(tcp[2:4], destinationPort) + binary.BigEndian.PutUint32(tcp[4:8], seq) + binary.BigEndian.PutUint32(tcp[8:12], ackNumber) + headerWordCount := 5 + (len(options)+3)/4 + tcp[12] = uint8(headerWordCount << 4) + tcp[13] = flags + binary.BigEndian.PutUint16(tcp[14:16], windowSize) + tcp[16] = 0 + tcp[17] = 0 + tcp[18] = 0 + tcp[19] = 0 + copy(tcp[20:], options) } type TcpBufferSettings struct { @@ -1309,17 +1533,17 @@ func NewTcp4Buffer(ctx context.Context, receiveCallback ReceivePacketFunction, } func (self *Tcp4Buffer) send(source TransferPath, provideMode protocol.ProvideMode, - ipv4 *layers.IPv4, tcp *layers.TCP, timeout time.Duration, ipPacket []byte) (bool, error) { + tcp *parsedTcp, timeout time.Duration, ipPacket []byte) (bool, error) { bufferId := NewBufferId4( source, - ipv4.SrcIP, int(tcp.SrcPort), - ipv4.DstIP, int(tcp.DstPort), + tcp.sourceIp, int(tcp.sourcePort), + tcp.destinationIp, int(tcp.destinationPort), ) return self.tcpSend( bufferId, - ipv4.SrcIP, - ipv4.DstIP, + tcp.sourceIp, + tcp.destinationIp, source, provideMode, 4, @@ -1341,17 +1565,17 @@ func NewTcp6Buffer(ctx context.Context, receiveCallback ReceivePacketFunction, } func (self *Tcp6Buffer) send(source TransferPath, provideMode protocol.ProvideMode, - ipv6 *layers.IPv6, tcp *layers.TCP, timeout time.Duration, ipPacket []byte) (bool, error) { + tcp *parsedTcp, timeout time.Duration, ipPacket []byte) (bool, error) { bufferId := NewBufferId6( source, - ipv6.SrcIP, int(tcp.SrcPort), - ipv6.DstIP, int(tcp.DstPort), + tcp.sourceIp, int(tcp.sourcePort), + tcp.destinationIp, int(tcp.destinationPort), ) return self.tcpSend( bufferId, - ipv6.SrcIP, - ipv6.DstIP, + tcp.sourceIp, + tcp.destinationIp, source, provideMode, 6, @@ -1398,7 +1622,7 @@ func (self *TcpBuffer[BufferId]) tcpSend( source TransferPath, provideMode protocol.ProvideMode, ipVersion int, - tcp *layers.TCP, + tcp *parsedTcp, timeout time.Duration, ipPacket []byte, ) (bool, error) { @@ -1407,7 +1631,7 @@ func (self *TcpBuffer[BufferId]) tcpSend( defer self.mutex.Unlock() if sequence, ok := self.sequences[bufferId]; ok { - if tcp.RST { + if tcp.rst { sequence.Cancel() delete(self.sequences, bufferId) sourceSequences := self.sourceSequences[sequence.source] @@ -1419,14 +1643,15 @@ func (self *TcpBuffer[BufferId]) tcpSend( self.bw.RemoveSession(sequence.source) } } + MessagePoolReturn(ipPacket) return nil } return sequence } - if !tcp.SYN { + if !tcp.syn { MessagePoolReturn(ipPacket) - self.log.V(2).Infof("[lnr]tcp drop no syn (%s)\n", tcpFlagsString(tcp)) + self.log.V(2).Infof("[lnr]tcp drop no syn (%s)\n", tcp.flagsString()) return nil } @@ -1459,9 +1684,9 @@ func (self *TcpBuffer[BufferId]) tcpSend( provideMode, ipVersion, sourceIpCopy, - tcp.SrcPort, + tcp.sourcePort, destinationIpCopy, - tcp.DstPort, + tcp.destinationPort, self.tcpBufferSettings, ) self.sequences[bufferId] = sequence @@ -1539,8 +1764,8 @@ func NewTcpSequence(ctx context.Context, receiveCallback ReceivePacketFunction, source TransferPath, provideMode protocol.ProvideMode, ipVersion int, - sourceIp net.IP, sourcePort layers.TCPPort, - destinationIp net.IP, destinationPort layers.TCPPort, + sourceIp net.IP, sourcePort TCPPort, + destinationIp net.IP, destinationPort TCPPort, tcpBufferSettings *TcpBufferSettings) *TcpSequence { cancelCtx, cancel := context.WithCancel(ctx) @@ -1569,7 +1794,7 @@ func NewTcpSequence(ctx context.Context, receiveCallback ReceivePacketFunction, // FIXME initial window size should be ~4k, set max window size as a 2^amount multiplier of initial size windowSize: tcpBufferSettings.MinWindowSize, windowScale: 0, - buffer: gopacket.NewSerializeBufferExpectedSize(128, 2048), + userLimited: userLimited{ lastActivityTime: time.Now(), }, @@ -1696,9 +1921,9 @@ func (self *TcpSequence) Run() { case <-self.ctx.Done(): return case sendItem := <-self.sendItems: - self.log.V(2).Infof("[init]send(%d)\n", len(sendItem.tcp.Payload)) + self.log.V(2).Infof("[init]send(%d)\n", len(sendItem.tcp.payload)) // the first packet must be a syn - if sendItem.tcp.SYN { + if sendItem.tcp.syn { self.log.V(2).Infof("[init]SYN\n") func() { @@ -1707,31 +1932,14 @@ func (self *TcpSequence) Run() { // sendSeq is the next expected sequence number // SYN and FIN consume one - self.sendSeq = sendItem.tcp.Seq + 1 + self.sendSeq = sendItem.tcp.seq + 1 // start the send seq at send seq // this is arbitrary, and since there is no transport security risk back to sender is fine - self.receiveSeq = sendItem.tcp.Seq - self.receiveSeqAck = sendItem.tcp.Seq - - parseWindowScaleOpts := func() (bool, uint32) { - for _, opt := range sendItem.tcp.Options { - if opt.OptionType == layers.TCPOptionKindWindowScale { - // fmt.Printf("[init]OPTION DATA = %v (%d,%d)\n", opt.OptionData, len(opt.OptionData), opt.OptionLength) - windowScaleBytes := make([]byte, 4) - copy(windowScaleBytes[4-len(opt.OptionData):4], opt.OptionData) - windowScale := min( - binary.BigEndian.Uint32(windowScaleBytes[0:4]), - // see 2.3 Using the Window Scale Option - 14, - ) - return true, windowScale - } - } - return false, 0 - } + self.receiveSeq = sendItem.tcp.seq + self.receiveSeqAck = sendItem.tcp.seq - self.enableWindowScale, self.receiveWindowScale = parseWindowScaleOpts() - self.receiveWindowSize = uint32(sendItem.tcp.Window) << self.receiveWindowScale + self.enableWindowScale, self.receiveWindowScale = ParseTcpWindowScaleOpts(sendItem.tcp.options) + self.receiveWindowSize = uint32(sendItem.tcp.windowSize) << self.receiveWindowScale if self.enableWindowScale { // compute the window scale to fit the window size in uint16 bits := math.Log2(float64(self.tcpBufferSettings.MaxWindowSize) / float64(math.MaxUint16)) @@ -2146,11 +2354,11 @@ func (self *TcpSequence) Run() { lastActivityTime = time.Now() if self.log.V(2).Enabled() { if "ACK" != tcpFlagsString(sendItem.tcp) { - self.log.Infof("[r%d]receive(%d %s)\n", sendIter, len(sendItem.tcp.Payload), tcpFlagsString(sendItem.tcp)) + self.log.Infof("[r%d]receive(%d %s)\n", sendIter, len(sendItem.tcp.payload), tcpFlagsString(sendItem.tcp)) } } - if sendItem.tcp.RST { + if sendItem.tcp.rst { // a RST typically appears for a bad TCP segment self.log.V(2).Infof("[r%d]RST\n", sendIter) MessagePoolReturn(sendItem.ipPacket) @@ -2163,19 +2371,19 @@ func (self *TcpSequence) Run() { self.mutex.Lock() defer self.mutex.Unlock() - if self.sendSeq != sendItem.tcp.Seq || sendItem.tcp.SYN { + if self.sendSeq != sendItem.tcp.seq || sendItem.tcp.syn { // a retransmit // since the transfer from local to remote is lossless and preserves order, // the packet is already pending. Ignore. drop = true - } else if sendItem.tcp.ACK { + } else if sendItem.tcp.ack { // acks are reliably delivered (see above) // we do not need to resend receive packets on missing acks // note the window size can be be adjusted at any time for the same receive seq number, // e.g. ->0 then ->full on receiver full - if self.receiveSeqAck <= sendItem.tcp.Ack { - self.receiveWindowSize = uint32(sendItem.tcp.Window) << self.receiveWindowScale - self.receiveSeqAck = sendItem.tcp.Ack + if self.receiveSeqAck <= sendItem.tcp.ackNumber { + self.receiveWindowSize = uint32(sendItem.tcp.windowSize) << self.receiveWindowScale + self.receiveSeqAck = sendItem.tcp.ackNumber receiveAckCond.Broadcast() } } @@ -2186,7 +2394,7 @@ func (self *TcpSequence) Run() { continue } - if sendItem.tcp.FIN { + if sendItem.tcp.fin { self.log.V(2).Infof("[r%d]FIN\n", sendIter) func() { self.mutex.Lock() @@ -2197,7 +2405,7 @@ func (self *TcpSequence) Run() { }() } - payload := sendItem.tcp.Payload + payload := sendItem.tcp.payload if 0 < len(payload) { writePayload := writePayload{ payload: payload, @@ -2245,7 +2453,7 @@ func (self *TcpSequence) Run() { MessagePoolReturn(sendItem.ipPacket) } - if sendItem.tcp.FIN { + if sendItem.tcp.fin { // flush the write channel to propage the FIN and close the sequence close(writePayloads) fin = true @@ -2299,7 +2507,7 @@ func (self *TcpSequence) Close() { type TcpSendItem struct { provideMode protocol.ProvideMode - tcp *layers.TCP + tcp *parsedTcp ipPacket []byte } @@ -2308,9 +2516,9 @@ type ConnectionState struct { provideMode protocol.ProvideMode ipVersion int sourceIp net.IP - sourcePort layers.TCPPort + sourcePort TCPPort destinationIp net.IP - destinationPort layers.TCPPort + destinationPort TCPPort mutex sync.Mutex @@ -2324,8 +2532,6 @@ type ConnectionState struct { windowScale uint32 // encodedWindowSize uint16 - buffer gopacket.SerializeBuffer - userLimited } @@ -2348,357 +2554,184 @@ func (self *ConnectionState) encodedWindowSize() uint16 { } func (self *ConnectionState) SynAck() ([]byte, error) { - headerSize := 0 - var ip gopacket.NetworkLayer + var ipHeaderByteCount int switch self.ipVersion { case 4: - ip = &layers.IPv4{ - Version: 4, - TTL: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - Protocol: layers.IPProtocolTCP, - } - headerSize += Ipv4HeaderSizeWithoutExtensions + ipHeaderByteCount = Ipv4HeaderSizeWithoutExtensions case 6: - ip = &layers.IPv6{ - Version: 6, - HopLimit: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - NextHeader: layers.IPProtocolTCP, - } - headerSize += Ipv6HeaderSize + ipHeaderByteCount = Ipv6HeaderSize } - opts := []layers.TCPOption{} - + var optsBytes []byte if self.enableWindowScale { windowScaleBytes := make([]byte, 4) binary.BigEndian.PutUint32(windowScaleBytes[0:4], self.windowScale) - - windowScaleOpt := layers.TCPOption{ - OptionType: layers.TCPOptionKindWindowScale, - OptionLength: 3, - // one byte - OptionData: windowScaleBytes[3:4], - } - opts = append(opts, windowScaleOpt) + // options must be padded to a 4-byte boundary to match writeTcpHeader's + // data-offset (headerWordCount) computation; make() zero-fills the pad, + // which reads as an EOL terminator (kind 0) + const optionsByteCount = 3 + paddedOptionsByteCount := (optionsByteCount + 3) &^ 3 + optsBytes = make([]byte, paddedOptionsByteCount) + optsBytes[0] = 3 + optsBytes[1] = 3 + optsBytes[2] = windowScaleBytes[3] + } + + tcpHeaderByteCount := TcpHeaderSizeWithoutExtensions + len(optsBytes) + packet := MessagePoolGet(ipHeaderByteCount + tcpHeaderByteCount) + switch self.ipVersion { + case 4: + writeIpv4Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) + case 6: + writeIpv6Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) } - tcp := layers.TCP{ - SrcPort: self.destinationPort, - DstPort: self.sourcePort, - Seq: self.receiveSeq, - Ack: self.sendSeq, - ACK: true, - SYN: true, - Window: self.encodedWindowSize(), - Options: opts, - } - tcp.SetNetworkLayerForChecksum(ip) - headerSize += TcpHeaderSizeWithoutExtensions + flags := tcpFlagSyn | tcpFlagAck + writeTcpHeader(packet[ipHeaderByteCount:], uint16(self.destinationPort), uint16(self.sourcePort), self.receiveSeq, self.sendSeq, flags, self.encodedWindowSize(), optsBytes) - options := gopacket.SerializeOptions{ - ComputeChecksums: true, - FixLengths: true, - } + // checksum covers the full segment (header + options), not just the fixed header + tcpBytes := packet[ipHeaderByteCount:] + checksum := transportChecksum(IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp, tcpBytes) + binary.BigEndian.PutUint16(tcpBytes[16:18], checksum) - self.buffer.Clear() - err := gopacket.SerializeLayers(self.buffer, options, - ip.(gopacket.SerializableLayer), - &tcp, - ) - - if err != nil { - return nil, err - } - packet := MessagePoolCopy(self.buffer.Bytes()) return packet, nil } - func (self *ConnectionState) PureAck() ([]byte, error) { - headerSize := 0 - var ip gopacket.NetworkLayer + var ipHeaderByteCount int switch self.ipVersion { case 4: - ip = &layers.IPv4{ - Version: 4, - TTL: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - Protocol: layers.IPProtocolTCP, - } - headerSize += Ipv4HeaderSizeWithoutExtensions + ipHeaderByteCount = Ipv4HeaderSizeWithoutExtensions case 6: - ip = &layers.IPv6{ - Version: 6, - HopLimit: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - NextHeader: layers.IPProtocolTCP, - } - headerSize += Ipv6HeaderSize + ipHeaderByteCount = Ipv6HeaderSize } - tcp := layers.TCP{ - SrcPort: self.destinationPort, - DstPort: self.sourcePort, - Seq: self.receiveSeq, - Ack: self.sendSeq, - ACK: true, - Window: self.encodedWindowSize(), + packet := MessagePoolGet(ipHeaderByteCount + TcpHeaderSizeWithoutExtensions) + switch self.ipVersion { + case 4: + writeIpv4Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) + case 6: + writeIpv6Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) } - tcp.SetNetworkLayerForChecksum(ip) - headerSize += TcpHeaderSizeWithoutExtensions - options := gopacket.SerializeOptions{ - ComputeChecksums: true, - FixLengths: true, - } + writeTcpHeader(packet[ipHeaderByteCount:], uint16(self.destinationPort), uint16(self.sourcePort), self.receiveSeq, self.sendSeq, tcpFlagAck, self.encodedWindowSize(), nil) - self.buffer.Clear() - err := gopacket.SerializeLayers(self.buffer, options, - ip.(gopacket.SerializableLayer), - &tcp, - ) + tcpBytes := packet[ipHeaderByteCount : ipHeaderByteCount+TcpHeaderSizeWithoutExtensions] + checksum := transportChecksum(IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp, tcpBytes) + binary.BigEndian.PutUint16(tcpBytes[16:18], checksum) - if err != nil { - return nil, err - } - packet := MessagePoolCopy(self.buffer.Bytes()) return packet, nil } func (self *ConnectionState) FinAck() ([]byte, error) { - headerSize := 0 - var ip gopacket.NetworkLayer + var ipHeaderByteCount int switch self.ipVersion { case 4: - ip = &layers.IPv4{ - Version: 4, - TTL: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - Protocol: layers.IPProtocolTCP, - } - headerSize += Ipv4HeaderSizeWithoutExtensions + ipHeaderByteCount = Ipv4HeaderSizeWithoutExtensions case 6: - ip = &layers.IPv6{ - Version: 6, - HopLimit: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - NextHeader: layers.IPProtocolTCP, - } - headerSize += Ipv6HeaderSize + ipHeaderByteCount = Ipv6HeaderSize } - tcp := layers.TCP{ - SrcPort: self.destinationPort, - DstPort: self.sourcePort, - Seq: self.receiveSeq, - Ack: self.sendSeq, - ACK: true, - FIN: true, - Window: self.encodedWindowSize(), + packet := MessagePoolGet(ipHeaderByteCount + TcpHeaderSizeWithoutExtensions) + switch self.ipVersion { + case 4: + writeIpv4Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) + case 6: + writeIpv6Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) } - tcp.SetNetworkLayerForChecksum(ip) - headerSize += TcpHeaderSizeWithoutExtensions - options := gopacket.SerializeOptions{ - ComputeChecksums: true, - FixLengths: true, - } + flags := tcpFlagFin | tcpFlagAck + writeTcpHeader(packet[ipHeaderByteCount:], uint16(self.destinationPort), uint16(self.sourcePort), self.receiveSeq, self.sendSeq, flags, self.encodedWindowSize(), nil) - self.buffer.Clear() - err := gopacket.SerializeLayers(self.buffer, options, - ip.(gopacket.SerializableLayer), - &tcp, - ) + tcpBytes := packet[ipHeaderByteCount : ipHeaderByteCount+TcpHeaderSizeWithoutExtensions] + checksum := transportChecksum(IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp, tcpBytes) + binary.BigEndian.PutUint16(tcpBytes[16:18], checksum) - if err != nil { - return nil, err - } - packet := MessagePoolCopy(self.buffer.Bytes()) return packet, nil } - func (self *ConnectionState) RstAck() ([]byte, error) { - headerSize := 0 - var ip gopacket.NetworkLayer + var ipHeaderByteCount int switch self.ipVersion { case 4: - ip = &layers.IPv4{ - Version: 4, - TTL: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - Protocol: layers.IPProtocolTCP, - } - headerSize += Ipv4HeaderSizeWithoutExtensions + ipHeaderByteCount = Ipv4HeaderSizeWithoutExtensions case 6: - ip = &layers.IPv6{ - Version: 6, - HopLimit: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - NextHeader: layers.IPProtocolTCP, - } - headerSize += Ipv6HeaderSize + ipHeaderByteCount = Ipv6HeaderSize } - tcp := layers.TCP{ - SrcPort: self.destinationPort, - DstPort: self.sourcePort, - Seq: self.receiveSeq, - Ack: self.sendSeq, - ACK: true, - RST: true, - Window: self.encodedWindowSize(), + packet := MessagePoolGet(ipHeaderByteCount + TcpHeaderSizeWithoutExtensions) + switch self.ipVersion { + case 4: + writeIpv4Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) + case 6: + writeIpv6Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) } - tcp.SetNetworkLayerForChecksum(ip) - headerSize += TcpHeaderSizeWithoutExtensions - options := gopacket.SerializeOptions{ - ComputeChecksums: true, - FixLengths: true, - } + flags := tcpFlagRst | tcpFlagAck + writeTcpHeader(packet[ipHeaderByteCount:], uint16(self.destinationPort), uint16(self.sourcePort), self.receiveSeq, self.sendSeq, flags, self.encodedWindowSize(), nil) - self.buffer.Clear() - err := gopacket.SerializeLayers(self.buffer, options, - ip.(gopacket.SerializableLayer), - &tcp, - ) + tcpBytes := packet[ipHeaderByteCount : ipHeaderByteCount+TcpHeaderSizeWithoutExtensions] + checksum := transportChecksum(IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp, tcpBytes) + binary.BigEndian.PutUint16(tcpBytes[16:18], checksum) - if err != nil { - return nil, err - } - packet := MessagePoolCopy(self.buffer.Bytes()) return packet, nil } - func (self *ConnectionState) DataPackets(payload []byte, n int, mtu int) ([][]byte, error) { - headerSize := 0 - var ip gopacket.NetworkLayer + var ipHeaderByteCount int switch self.ipVersion { case 4: - ip = &layers.IPv4{ - Version: 4, - TTL: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - Protocol: layers.IPProtocolTCP, - } - headerSize += Ipv4HeaderSizeWithoutExtensions + ipHeaderByteCount = Ipv4HeaderSizeWithoutExtensions case 6: - ip = &layers.IPv6{ - Version: 6, - HopLimit: 64, - SrcIP: self.destinationIp, - DstIP: self.sourceIp, - NextHeader: layers.IPProtocolTCP, - } - headerSize += Ipv6HeaderSize + ipHeaderByteCount = Ipv6HeaderSize } - tcp := layers.TCP{ - SrcPort: self.destinationPort, - DstPort: self.sourcePort, - Seq: self.receiveSeq, - Ack: self.sendSeq, - ACK: true, - Window: self.encodedWindowSize(), + headerByteCount := ipHeaderByteCount + TcpHeaderSizeWithoutExtensions + if mtu <= headerByteCount { + return nil, fmt.Errorf("mtu %d is too small for IP+TCP headers (%d bytes)", mtu, headerByteCount) } - tcp.SetNetworkLayerForChecksum(ip) - headerSize += TcpHeaderSizeWithoutExtensions - - options := gopacket.SerializeOptions{ - ComputeChecksums: true, - FixLengths: true, + packetByteCount := mtu - headerByteCount + if n <= packetByteCount { + pkt := self.tcpPacket(payload[0:n], self.receiveSeq) + return [][]byte{pkt}, nil } - - if debugVerifyHeaders { - buffer := gopacket.NewSerializeBufferExpectedSize(headerSize, 0) - err := gopacket.SerializeLayers(buffer, options, - ip.(gopacket.SerializableLayer), - &tcp, - ) - if err != nil { - return nil, err - } - packetHeaders := buffer.Bytes() - if headerSize != len(packetHeaders) { - return nil, errors.New(fmt.Sprintf("Header check failed %d <> %d", headerSize, len(packetHeaders))) - } - } - - if headerSize+n <= mtu { - self.buffer.Clear() - err := gopacket.SerializeLayers(self.buffer, options, - ip.(gopacket.SerializableLayer), - &tcp, - gopacket.Payload(payload[0:n]), - ) - if err != nil { - return nil, err - } - packet := MessagePoolCopy(self.buffer.Bytes()) - return [][]byte{packet}, nil - } else { - // fragment - packetSize := mtu - headerSize - packets := [][]byte{} - for i := 0; i < n; { - self.buffer.Clear() - j := min(i+packetSize, n) - tcp.Seq = self.receiveSeq + uint32(i) - err := gopacket.SerializeLayers(self.buffer, options, - ip.(gopacket.SerializableLayer), - &tcp, - gopacket.Payload(payload[i:j]), - ) - if err != nil { - return nil, err - } - packet := MessagePoolCopy(self.buffer.Bytes()) - packets = append(packets, packet) - i = j - } - return packets, nil + packets := make([][]byte, 0, (n+packetByteCount-1)/packetByteCount) + for i := 0; i < n; { + j := min(i+packetByteCount, n) + packets = append(packets, self.tcpPacket(payload[i:j], self.receiveSeq+uint32(i))) + i = j } + return packets, nil } -func tcpFlagsString(tcp *layers.TCP) string { - flags := []string{} - if tcp.FIN { - flags = append(flags, "FIN") - } - if tcp.SYN { - flags = append(flags, "SYN") - } - if tcp.RST { - flags = append(flags, "RST") - } - if tcp.PSH { - flags = append(flags, "PSH") - } - if tcp.ACK { - flags = append(flags, "ACK") - } - if tcp.URG { - flags = append(flags, "URG") - } - if tcp.ECE { - flags = append(flags, "ECE") - } - if tcp.CWR { - flags = append(flags, "CWR") +func (self *ConnectionState) tcpPacket(payload []byte, seq uint32) []byte { + var ipHeaderByteCount int + switch self.ipVersion { + case 4: + ipHeaderByteCount = Ipv4HeaderSizeWithoutExtensions + case 6: + ipHeaderByteCount = Ipv6HeaderSize } - if tcp.NS { - flags = append(flags, "NS") + + tcpHeaderByteCount := TcpHeaderSizeWithoutExtensions + totalLen := ipHeaderByteCount + tcpHeaderByteCount + len(payload) + packet := MessagePoolGet(totalLen) + switch self.ipVersion { + case 4: + writeIpv4Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) + case 6: + writeIpv6Header(packet, IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp) } - return strings.Join(flags, ", ") + + writeTcpHeader(packet[ipHeaderByteCount:], uint16(self.destinationPort), uint16(self.sourcePort), seq, self.sendSeq, tcpFlagAck, self.encodedWindowSize(), nil) + copy(packet[ipHeaderByteCount+tcpHeaderByteCount:], payload) + + // checksum covers the full segment (header + payload), not just the header + tcpBytes := packet[ipHeaderByteCount:] + checksum := transportChecksum(IP_PROTOCOL_TCP, self.destinationIp, self.sourceIp, tcpBytes) + binary.BigEndian.PutUint16(tcpBytes[16:18], checksum) + + return packet +} +func tcpFlagsString(tcp *parsedTcp) string { + return tcp.flagsString() } func DefaultRemoteUserNatProviderSettings() *RemoteUserNatProviderSettings { @@ -2823,11 +2856,11 @@ func (self *RemoteUserNatProvider) Receive( } // `connect.ReceiveFunction` -func (self *RemoteUserNatProvider) ClientReceive(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { +func (self *RemoteUserNatProvider) ClientReceive(source TransferPath, frames []*protocol.Frame, peer Peer) { for _, frame := range frames { switch frame.MessageType { case protocol.MessageType_IpIpPing: - self.client.log.V(1).Infof("[ip]provider ping <- %s(%d)\n", source, provideMode) + self.client.log.V(1).Infof("[ip]provider ping <- %s(%d)\n", source, peer.ProvideMode) // echo back over a companion contract, like the provider's other // return traffic; the source only provides ProvideMode_Stream, so a // forward contract here would be rejected (no permission). @@ -2850,7 +2883,7 @@ func (self *RemoteUserNatProvider) ClientReceive(source TransferPath, frames []* ipPath, payload, err := ParseIpPathWithPayload(ipPacketToProvider.IpPacket.PacketBytes) if err == nil { - r, err := self.securityPolicy.Inspect(provideMode, ipPath, payload) + r, err := self.securityPolicy.Inspect(peer.ProvideMode, ipPath, payload) if err == nil { switch r { case SecurityPolicyResultAllow: @@ -2864,7 +2897,7 @@ func (self *RemoteUserNatProvider) ClientReceive(source TransferPath, frames []* // self.client.log.Infof("[trace]provider send packet from %s\n", source.SourceId) success := self.localUserNat.SendPacketWithTimeout( source, - provideMode, + peer.ProvideMode, packet, self.settings.WriteTimeout, ) @@ -3021,7 +3054,7 @@ func (self *RemoteUserNatClient) SendPacket(source TransferPath, provideMode pro } // `connect.ReceiveFunction` -func (self *RemoteUserNatClient) ClientReceive(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { +func (self *RemoteUserNatClient) ClientReceive(source TransferPath, frames []*protocol.Frame, peer Peer) { // only process frames from the destinations // if allow := self.sourceFilter[source]; !allow { // return @@ -3044,7 +3077,7 @@ func (self *RemoteUserNatClient) ClientReceive(source TransferPath, frames []*pr HandleError(func() { self.receivePacketCallback( source, - provideMode, + peer.ProvideMode, ipPath, packet, ) @@ -3196,99 +3229,114 @@ func ParseIpPath(ipPacket []byte) (*IpPath, error) { } func ParseIpPathWithPayload(ipPacket []byte) (*IpPath, []byte, error) { + if len(ipPacket) == 0 { + return nil, nil, fmt.Errorf("Empty packet.") + } ipVersion := uint8(ipPacket[0]) >> 4 switch ipVersion { case 4: - ipv4 := layers.IPv4{} - ipv4.DecodeFromBytes(ipPacket, gopacket.NilDecodeFeedback) + ipProtocol4, sourceIp4, destinationIp4, transport, ok := parseIpv4(ipPacket) + if !ok { + return nil, nil, fmt.Errorf("No support for protocol") + } // copy the ips so the ip path can be retained independently of the // shared packet buffer (which is recycled after the handoff call). both // copies share one backing allocation instead of one per address. - ipBacking := make(net.IP, len(ipv4.SrcIP)+len(ipv4.DstIP)) - sn := copy(ipBacking, ipv4.SrcIP) - copy(ipBacking[sn:], ipv4.DstIP) + ipBacking := make(net.IP, len(sourceIp4)+len(destinationIp4)) + sn := copy(ipBacking, sourceIp4) + copy(ipBacking[sn:], destinationIp4) sourceIpCopy := ipBacking[:sn:sn] destinationIpCopy := ipBacking[sn:] - switch ipv4.Protocol { - case layers.IPProtocolUDP: - udp := layers.UDP{} - udp.DecodeFromBytes(ipv4.Payload, gopacket.NilDecodeFeedback) + switch ipProtocol4 { + case IP_PROTOCOL_UDP: + var udpPacket parsedUdp + if !parseUdpPacket(sourceIp4, destinationIp4, transport, &udpPacket) { + return nil, nil, fmt.Errorf("No support for protocol") + } return &IpPath{ Version: int(ipVersion), Protocol: IpProtocolUdp, SourceIp: sourceIpCopy, - SourcePort: int(udp.SrcPort), + SourcePort: int(udpPacket.sourcePort), DestinationIp: destinationIpCopy, - DestinationPort: int(udp.DstPort), - }, udp.Payload, nil - case layers.IPProtocolTCP: - tcp := layers.TCP{} - tcp.DecodeFromBytes(ipv4.Payload, gopacket.NilDecodeFeedback) + DestinationPort: int(udpPacket.destinationPort), + }, udpPacket.payload, nil + case IP_PROTOCOL_TCP: + var tcpPacket parsedTcp + if !parseTcpPacket(sourceIp4, destinationIp4, transport, &tcpPacket) { + return nil, nil, fmt.Errorf("No support for protocol") + } return &IpPath{ Version: int(ipVersion), Protocol: IpProtocolTcp, SourceIp: sourceIpCopy, - SourcePort: int(tcp.SrcPort), + SourcePort: int(tcpPacket.sourcePort), DestinationIp: destinationIpCopy, - DestinationPort: int(tcp.DstPort), - SequenceNumber: tcp.Seq, - Syn: tcp.SYN, - Rst: tcp.RST, - Ack: tcp.ACK, - }, tcp.Payload, nil + DestinationPort: int(tcpPacket.destinationPort), + SequenceNumber: tcpPacket.seq, + Syn: tcpPacket.syn, + Rst: tcpPacket.rst, + Ack: tcpPacket.ack, + }, tcpPacket.payload, nil default: // no support for this protocol - return nil, nil, fmt.Errorf("No support for protocol %d", ipv4.Protocol) + return nil, nil, fmt.Errorf("No support for protocol %d", ipProtocol4) } case 6: - ipv6 := layers.IPv6{} - ipv6.DecodeFromBytes(ipPacket, gopacket.NilDecodeFeedback) + ipProtocol6, sourceIp6, destinationIp6, transport, ok := parseIpv6(ipPacket) + if !ok { + return nil, nil, fmt.Errorf("No support for protocol") + } // copy the ips so the ip path can be retained independently of the // shared packet buffer (which is recycled after the handoff call). both // copies share one backing allocation instead of one per address. - ipBacking := make(net.IP, len(ipv6.SrcIP)+len(ipv6.DstIP)) - sn := copy(ipBacking, ipv6.SrcIP) - copy(ipBacking[sn:], ipv6.DstIP) + ipBacking := make(net.IP, len(sourceIp6)+len(destinationIp6)) + sn := copy(ipBacking, sourceIp6) + copy(ipBacking[sn:], destinationIp6) sourceIpCopy := ipBacking[:sn:sn] destinationIpCopy := ipBacking[sn:] - switch ipv6.NextHeader { - case layers.IPProtocolUDP: - udp := layers.UDP{} - udp.DecodeFromBytes(ipv6.Payload, gopacket.NilDecodeFeedback) + switch ipProtocol6 { + case IP_PROTOCOL_UDP: + var udpPacket parsedUdp + if !parseUdpPacket(sourceIp6, destinationIp6, transport, &udpPacket) { + return nil, nil, fmt.Errorf("No support for protocol") + } return &IpPath{ Version: int(ipVersion), Protocol: IpProtocolUdp, SourceIp: sourceIpCopy, - SourcePort: int(udp.SrcPort), + SourcePort: int(udpPacket.sourcePort), DestinationIp: destinationIpCopy, - DestinationPort: int(udp.DstPort), - }, udp.Payload, nil - case layers.IPProtocolTCP: - tcp := layers.TCP{} - tcp.DecodeFromBytes(ipv6.Payload, gopacket.NilDecodeFeedback) + DestinationPort: int(udpPacket.destinationPort), + }, udpPacket.payload, nil + case IP_PROTOCOL_TCP: + var tcpPacket parsedTcp + if !parseTcpPacket(sourceIp6, destinationIp6, transport, &tcpPacket) { + return nil, nil, fmt.Errorf("No support for protocol") + } return &IpPath{ Version: int(ipVersion), Protocol: IpProtocolTcp, SourceIp: sourceIpCopy, - SourcePort: int(tcp.SrcPort), + SourcePort: int(tcpPacket.sourcePort), DestinationIp: destinationIpCopy, - DestinationPort: int(tcp.DstPort), - SequenceNumber: tcp.Seq, - Syn: tcp.SYN, - Rst: tcp.RST, - Ack: tcp.ACK, - }, tcp.Payload, nil + DestinationPort: int(tcpPacket.destinationPort), + SequenceNumber: tcpPacket.seq, + Syn: tcpPacket.syn, + Rst: tcpPacket.rst, + Ack: tcpPacket.ack, + }, tcpPacket.payload, nil default: // no support for this protocol - return nil, nil, fmt.Errorf("No support for protocol %d", ipv6.NextHeader) + return nil, nil, fmt.Errorf("No support for protocol %d", ipProtocol6) } default: // no support for this version diff --git a/ip_assoc.go b/ip_assoc.go new file mode 100644 index 0000000..3a7d418 --- /dev/null +++ b/ip_assoc.go @@ -0,0 +1,1135 @@ +package connect + +import ( + "context" + "net/netip" + "slices" + "sync" + "sync/atomic" + "time" + + "golang.org/x/net/publicsuffix" +) + +// IpAssoc is a time-blocked co-association matrix that tracks clusters +// of IpPaths by ip address and server base names. +// The core association is packet activity that occurs within `PacketAssociationTime` +// of each other on each path. The association probability in a time block is +// `P(a|b) = (co-occuring packets of a with b)/min(total packets of a, total packets of b)` +// The clustering uses an agglomative clustering with a mean inter-entity co-association threshold, +// which means each member has a mean of at least the inter-entity co-association threshold. +// The algorithm is greedy to build the largest clusters that satisfy the mean inter-entity co-association threshold, +// starting with all paths in one cluster and recursively splitting off paths to keep the largest amount that satisfy the mean inter-entity co-association threshold. +// the association clustering should use data from all available time blocks, not just the latest +// +// implementation notes: +// - an entity is a remote endpoint ip (egress destination, ingress source). +// entities that share a server base name (eTLD+1, resolved via the optional +// `ServerNameLookup`) are merged before clustering, since they are the same site. +// - packet activity is coalesced per `ActivityGranule` so the per-packet hot path is +// one lock + map probe in the steady state. co-occurrence therefore counts active +// granules, not raw packets, which bounds the pair-update cost under load and is +// robust to packet-rate skew between paths. +// - each activity pairs with at most `MaxPairPartnersPerActivity` co-active entities +// (a uniform sample via randomized map iteration), so a burst of n concurrent +// endpoints cannot mint n^2 pairs; persistent co-activity still accumulates the +// highest counts over a block. +// - a block stores entities on a dense uint16 index (addresses interned once per +// block) and pairs on a packed uint32 index-pair key, so the matrix costs a few +// bytes per pair instead of two addresses per pair. +// - clustering runs on the `ClusterEpoch` only when new activity was recorded, over +// reusable scratch (the aggregate is refilled, not reallocated), and only over +// connected components of the co-association graph with at most +// `MaxComponentNodeCount` nodes — a larger component is a degenerate burst clique +// with quadratic split cost and little affinity value, and is skipped. +// - a clustering pass checks the assoc ctx between phases, so teardown aborts an +// in-flight pass instead of letting it run out over a large matrix. the run loop +// drops the matrix and unregisters the memory shedder on exit (Close or parent +// ctx cancel), and `ShedMemory` (registered with `AddMemoryShedder`) drops it +// under host memory pressure; both publish empty clusters so consumers do not act +// on stale members. + +type IpAssocSettings struct { + AssociationBlockDuration time.Duration + AssociationBlockCount int + + // activity on two paths within this time counts as co-occurring + PacketAssociationTime time.Duration + // per-entity activity is coalesced to at most one count per granule + ActivityGranule time.Duration + + MinMeanAssociation float64 + + ClusterEpoch time.Duration + + // bounds per block. entities/pairs beyond these are not counted. + // MaxEntityCount is clamped to 65535 (entities are block-indexed on uint16). + MaxEntityCount int + MaxAssociationCount int + + // MaxPairPartnersPerActivity caps how many co-active entities one activity pairs + // with — a uniform sample of the active set (randomized map iteration order) — so + // a burst of concurrent endpoints grows pairs linearly, not quadratically. + // 0 uses a default. + MaxPairPartnersPerActivity int + + // MaxComponentNodeCount skips clustering a connected component larger than this + // many nodes: a degenerate burst clique has quadratic split cost and little + // affinity value. 0 uses a default. + MaxComponentNodeCount int +} + +func DefaultIpAssocSettings() *IpAssocSettings { + return &IpAssocSettings{ + AssociationBlockDuration: 300 * time.Second, + AssociationBlockCount: 8, + PacketAssociationTime: 5 * time.Second, + ActivityGranule: 1 * time.Second, + MinMeanAssociation: 0.5, + ClusterEpoch: 1 * time.Second, + MaxEntityCount: 2048, + MaxAssociationCount: 16384, + MaxPairPartnersPerActivity: defaultMaxPairPartnersPerActivity, + MaxComponentNodeCount: defaultMaxComponentNodeCount, + } +} + +const ( + // defaultMaxPairPartnersPerActivity is the MaxPairPartnersPerActivity default. + defaultMaxPairPartnersPerActivity = 16 + // defaultMaxComponentNodeCount is the MaxComponentNodeCount default. + defaultMaxComponentNodeCount = 256 + // maxIpAssocEntityCount is the hard entity-per-block bound (dense uint16 index). + maxIpAssocEntityCount = 65535 + // ipAssocPruneScanMin is the minimum lastActive entries an activity's pair scan + // visits before it may stop early, so pruning of expired activity stays amortized + // even when the pair-partner cap is reached immediately. + ipAssocPruneScanMin = 64 + // ipAssocScratchShrinkMin is the scratch map fill above which a large-then-idle + // pass replaces the map (clear keeps buckets, which would pin a burst's footprint). + ipAssocScratchShrinkMin = 4096 +) + +// canonical pair key with a < b. used by the map-input convenience form of +// clustering (see clusterIpAssoc); blocks store pairs on packed dense indexes. +type ipAssocPair struct { + a netip.Addr + b netip.Addr +} + +func newIpAssocPair(a netip.Addr, b netip.Addr) ipAssocPair { + if b.Less(a) { + a, b = b, a + } + return ipAssocPair{a: a, b: b} +} + +// ipAssocPackIndexPair packs two block entity indexes into the canonical (low, high) +// pair key. +func ipAssocPackIndexPair(i uint16, j uint16) uint32 { + if j < i { + i, j = j, i + } + return uint32(i)<<16 | uint32(j) +} + +// ipAssocPackNodePair packs two dense ids into the canonical (low, high) pair key. +func ipAssocPackNodePair(a uint32, b uint32) uint64 { + if b < a { + a, b = b, a + } + return uint64(a)<<32 | uint64(b) +} + +// ipAssocBlock stores one time block of the matrix on a dense entity index: +// an address is interned once (indexes/addrs), activity is a plain slice, and a +// pair is a packed index pair — a few bytes per pair instead of two addresses. +// An entity holds an index (and counts toward `MaxEntityCount`) once it records +// activity or appears as a pair partner. +type ipAssocBlock struct { + endTime time.Time + // entity -> dense index into addrs/counts + indexes map[netip.Addr]uint16 + // index -> entity + addrs []netip.Addr + // index -> active granule count (0 for a pure pair partner) + counts []uint32 + // packed index pair -> co-occurring granule count + coCounts map[uint32]uint32 +} + +func newIpAssocBlock(endTime time.Time) *ipAssocBlock { + return &ipAssocBlock{ + endTime: endTime, + indexes: map[netip.Addr]uint16{}, + coCounts: map[uint32]uint32{}, + } +} + +// index returns the dense index for addr, interning it if there is entity capacity. +func (self *ipAssocBlock) index(addr netip.Addr, maxEntityCount int) (uint16, bool) { + if i, ok := self.indexes[addr]; ok { + return i, true + } + if maxEntityCount <= len(self.addrs) { + return 0, false + } + i := uint16(len(self.addrs)) + self.indexes[addr] = i + self.addrs = append(self.addrs, addr) + self.counts = append(self.counts, 0) + return i, true +} + +// countFor returns the entity's active granule count in this block (0 if absent). +func (self *ipAssocBlock) countFor(addr netip.Addr) uint32 { + if i, ok := self.indexes[addr]; ok { + return self.counts[i] + } + return 0 +} + +// immutable clustering snapshot, read lock-free on the packet path +type ipAssocClusters struct { + version uint64 + // entity -> cluster member entities (including the entity), for clusters with 2+ members + members map[netip.Addr][]netip.Addr +} + +type IpAssoc struct { + ctx context.Context + cancel context.CancelFunc + + settings *IpAssocSettings + + serverNameLookup atomic.Pointer[serverNameLookupHolder] + + stateLock sync.Mutex + // newest last + blocks []*ipAssocBlock + // entities active within `PacketAssociationTime`, entity -> last activity + lastActive map[netip.Addr]time.Time + // entity -> server base names (eTLD+1), bounded small + baseNames map[netip.Addr][]string + dirty bool + // generation increments when the matrix is dropped (shed or teardown), so an + // in-flight clustering pass does not publish a stale result over the freshly + // emptied clusters + generation uint64 + + unregisterShed func() + + clusters atomic.Pointer[ipAssocClusters] + + // scratch is the reusable aggregation/clustering workspace, owned exclusively by + // the clustering pass (the run goroutine; tests with a long ClusterEpoch drive + // updateClusters directly instead) + scratch ipAssocScratch +} + +// indirection so a nil interface can be stored/cleared atomically +type serverNameLookupHolder struct { + lookup ServerNameLookup +} + +func NewIpAssoc(ctx context.Context, settings *IpAssocSettings) *IpAssoc { + cancelCtx, cancel := context.WithCancel(ctx) + + ipAssoc := &IpAssoc{ + ctx: cancelCtx, + cancel: cancel, + settings: settings, + lastActive: map[netip.Addr]time.Time{}, + baseNames: map[netip.Addr][]string{}, + } + ipAssoc.clusters.Store(&ipAssocClusters{ + version: 0, + members: map[netip.Addr][]netip.Addr{}, + }) + // drop the matrix under host memory pressure; the signal rebuilds from live traffic + ipAssoc.unregisterShed = AddMemoryShedder(ipAssoc.ShedMemory) + + go HandleError(ipAssoc.run, cancel) + + return ipAssoc +} + +// maxEntityCount is the per-block entity bound, clamped to the dense index range. +func (self *IpAssoc) maxEntityCount() int { + return min(self.settings.MaxEntityCount, maxIpAssocEntityCount) +} + +// maxPairPartners is the per-activity pair fan-out bound (0 settings value = default). +func (self *IpAssoc) maxPairPartners() int { + if 0 < self.settings.MaxPairPartnersPerActivity { + return self.settings.MaxPairPartnersPerActivity + } + return defaultMaxPairPartnersPerActivity +} + +// maxComponentNodes is the clustering component bound (0 settings value = default). +func (self *IpAssoc) maxComponentNodes() int { + if 0 < self.settings.MaxComponentNodeCount { + return self.settings.MaxComponentNodeCount + } + return defaultMaxComponentNodeCount +} + +// SetServerNameLookup installs (or clears, with nil) the lookup used to resolve +// entity ips to server base names. Safe to call at runtime. +func (self *IpAssoc) SetServerNameLookup(serverNameLookup ServerNameLookup) { + self.serverNameLookup.Store(&serverNameLookupHolder{ + lookup: serverNameLookup, + }) +} + +func (self *IpAssoc) AddEgressPacket(path *IpPath) { + if addr, ok := ipAssocAddr(path.DestinationIp); ok { + self.addActivity(addr) + } +} + +// the ingress path is the return packet as received, before reverse +// (the remote endpoint is the source) +func (self *IpAssoc) AddIngressPacket(path *IpPath) { + if addr, ok := ipAssocAddr(path.SourceIp); ok { + self.addActivity(addr) + } +} + +func ipAssocAddr(ip []byte) (netip.Addr, bool) { + addr, ok := netip.AddrFromSlice(ip) + if !ok { + return netip.Addr{}, false + } + return addr.Unmap(), true +} + +func (self *IpAssoc) addActivity(addr netip.Addr) { + now := time.Now() + + needBaseNames := false + func() { + self.stateLock.Lock() + defer self.stateLock.Unlock() + + if last, ok := self.lastActive[addr]; ok && now.Sub(last) < self.settings.ActivityGranule { + // coalesced into the current granule. steady-state fast path + return + } + + block := self.blockWithLock(now) + + idx, ok := block.index(addr, self.maxEntityCount()) + if !ok { + // the block is at entity capacity + return + } + block.counts[idx] += 1 + + // pair with a bounded sample of the entities active within the association + // time, pruning expired entries during the same scan. map iteration starts + // at a random position, so the sample is uniform and the amortized pruning + // (at least ipAssocPruneScanMin entries per scan) reaches every entry. + maxPartners := self.maxPairPartners() + partners := 0 + visited := 0 + for other, otherLast := range self.lastActive { + visited += 1 + if other != addr { + if self.settings.PacketAssociationTime < now.Sub(otherLast) { + delete(self.lastActive, other) + } else if partners < maxPartners { + if otherIdx, ok := block.index(other, self.maxEntityCount()); ok { + key := ipAssocPackIndexPair(idx, otherIdx) + if _, ok := block.coCounts[key]; ok || len(block.coCounts) < self.settings.MaxAssociationCount { + block.coCounts[key] += 1 + } + partners += 1 + } + } + } + if maxPartners <= partners && ipAssocPruneScanMin <= visited { + break + } + } + self.lastActive[addr] = now + + if _, ok := self.baseNames[addr]; !ok { + // mark looked-up so nameless entities do not repeat the lookup. + // names observed later refresh after the entity ages out of all blocks + self.baseNames[addr] = nil + needBaseNames = true + } + + self.dirty = true + }() + + if needBaseNames { + // the lookup calls an external object, so it must not hold the state lock + if baseNames := self.lookupBaseNames(addr); baseNames != nil { + self.stateLock.Lock() + defer self.stateLock.Unlock() + self.baseNames[addr] = baseNames + } + } +} + +func (self *IpAssoc) lookupBaseNames(addr netip.Addr) []string { + holder := self.serverNameLookup.Load() + if holder == nil || holder.lookup == nil { + return nil + } + serverNames := holder.lookup.ServerNames(addr.String()) + if len(serverNames) == 0 { + return nil + } + var baseNames []string + seen := map[string]bool{} + for _, serverName := range serverNames { + baseName := serverName + if rootDomain, err := publicsuffix.EffectiveTLDPlusOne(serverName); err == nil { + baseName = rootDomain + } + if !seen[baseName] { + seen[baseName] = true + baseNames = append(baseNames, baseName) + } + } + return baseNames +} + +// called with stateLock +func (self *IpAssoc) blockWithLock(now time.Time) *ipAssocBlock { + if n := len(self.blocks); 0 < n && now.Before(self.blocks[n-1].endTime) { + return self.blocks[n-1] + } + block := newIpAssocBlock(now.Add(self.settings.AssociationBlockDuration)) + self.blocks = append(self.blocks, block) + if self.settings.AssociationBlockCount < len(self.blocks) { + self.blocks = self.blocks[len(self.blocks)-self.settings.AssociationBlockCount:] + // prune base names for entities no longer in any block + liveAddrs := map[netip.Addr]bool{} + for _, block := range self.blocks { + for addr := range block.indexes { + liveAddrs[addr] = true + } + } + for addr := range self.baseNames { + if !liveAddrs[addr] { + delete(self.baseNames, addr) + } + } + } + return block +} + +func (self *IpAssoc) run() { + defer self.cancel() + // teardown (Close or parent ctx cancel): a stopped assoc must not hold the + // matrix or stay registered for memory shedding. the scratch is owned by this + // goroutine, so it is released here too. + defer func() { + self.unregisterShed() + self.ShedMemory() + self.scratch = ipAssocScratch{} + }() + + for { + select { + case <-self.ctx.Done(): + return + case <-time.After(self.settings.ClusterEpoch): + } + + self.updateClusters() + } +} + +// ShedMemory drops the association matrix — blocks, activity, the name cache, and the +// published clusters — under host memory pressure (or at teardown). The signal rebuilds +// from live traffic; consumers see empty clusters rather than stale members. +func (self *IpAssoc) ShedMemory() { + self.stateLock.Lock() + defer self.stateLock.Unlock() + + self.blocks = nil + clear(self.lastActive) + clear(self.baseNames) + self.dirty = false + self.generation += 1 + previous := self.clusters.Load() + if 0 < len(previous.members) { + self.clusters.Store(&ipAssocClusters{ + version: previous.version + 1, + members: map[netip.Addr][]netip.Addr{}, + }) + } +} + +func (self *IpAssoc) updateClusters() { + scratch := &self.scratch + var generation uint64 + dirty := func() bool { + self.stateLock.Lock() + defer self.stateLock.Unlock() + + if !self.dirty { + return false + } + self.dirty = false + generation = self.generation + + // aggregate all available time blocks into the reusable scratch: the matrix + // summed on dense aggregate indexes, base names carried by reference. pairs + // collect as a flat slice here (cheap under the lock, and several times + // smaller than a map); the clustering pass sorts and coalesces them outside + // the lock. + scratch.resetAggregate() + for _, block := range self.blocks { + remap := scratch.remap[:0] + for i, addr := range block.addrs { + aggIdx := scratch.index(addr) + scratch.counts[aggIdx] += uint64(block.counts[i]) + scratch.baseNames[aggIdx] = self.baseNames[addr] + remap = append(remap, aggIdx) + } + scratch.remap = remap + for packed, coCount := range block.coCounts { + i := remap[uint16(packed>>16)] + j := remap[uint16(packed)] + scratch.pairs = append(scratch.pairs, ipAssocAggPair{ + key: ipAssocPackNodePair(i, j), + count: coCount, + }) + } + } + return true + }() + if !dirty { + return + } + + members := clusterIpAssocScratch(self.ctx, scratch, self.settings.MinMeanAssociation, self.maxComponentNodes()) + scratch.finishPass() + if members == nil { + // aborted by teardown + return + } + + previous := self.clusters.Load() + if ipAssocClustersEqual(previous.members, members) { + return + } + // do not publish over a shed/teardown that raced this pass + self.stateLock.Lock() + defer self.stateLock.Unlock() + if generation != self.generation { + return + } + self.clusters.Store(&ipAssocClusters{ + version: previous.version + 1, + members: members, + }) +} + +// ipAssocAggPair is one aggregated pair observation: a packed pair key and its +// summed co-occurrence count. The aggregate keeps pairs as a flat slice, sorted and +// coalesced outside the state lock — the clustering only ever iterates pairs, and a +// slice is several times smaller than a map at the same fill. +type ipAssocAggPair struct { + key uint64 + count uint32 +} + +// ipAssocScratch is the reusable aggregation and clustering workspace: the map is +// cleared and refilled (buckets are kept), slices are truncated and regrown, so a +// steady-state clustering pass allocates almost nothing. finishPass shrinks storage +// whose fill collapsed well below its high water, so one burst does not pin its +// peak footprint. +type ipAssocScratch struct { + // aggregate matrix: entity -> dense aggregate index + indexes map[netip.Addr]uint32 + // aggregate index -> entity / summed count / base names (shared, read-only) + addrs []netip.Addr + counts []uint64 + baseNames [][]string + // pair observations on packed aggregate index pairs; sorted and coalesced by the + // clustering pass, and remapped in place onto node pairs after the name merge + pairs []ipAssocAggPair + // block index -> aggregate index, reused per block + remap []uint32 + + // clustering workspace, sized to the aggregate per pass + parent []uint32 + nodeIds []int32 + nodeCounts []uint64 + nodeMemberInit []int32 + nodeMembers []uint32 + memberFill []int32 + compParent []uint32 + compSizes []int32 + compIds []int32 + compInit []int32 + compNodes []uint32 + degrees []int32 + adjInit []int32 + adjNodes []uint32 + adjPs []float64 + + entityHighWater int +} + +func (self *ipAssocScratch) resetAggregate() { + if self.indexes == nil { + self.indexes = map[netip.Addr]uint32{} + } + clear(self.indexes) + self.addrs = self.addrs[:0] + self.counts = self.counts[:0] + self.baseNames = self.baseNames[:0] + self.pairs = self.pairs[:0] +} + +// index returns the dense aggregate index for addr, interning it on first use. +func (self *ipAssocScratch) index(addr netip.Addr) uint32 { + if i, ok := self.indexes[addr]; ok { + return i + } + i := uint32(len(self.addrs)) + self.indexes[addr] = i + self.addrs = append(self.addrs, addr) + self.counts = append(self.counts, 0) + self.baseNames = append(self.baseNames, nil) + return i +} + +// finishPass tracks fill high water and shrinks storage whose fill collapsed well +// below it (a cleared map keeps its buckets and a truncated slice its backing array, +// which would otherwise pin a burst's peak footprint). +func (self *ipAssocScratch) finishPass() { + if self.entityHighWater < len(self.indexes) { + self.entityHighWater = len(self.indexes) + } + if ipAssocScratchShrinkMin <= self.entityHighWater && len(self.indexes)*4 < self.entityHighWater { + fresh := make(map[netip.Addr]uint32, len(self.indexes)) + for addr, i := range self.indexes { + fresh[addr] = i + } + self.indexes = fresh + self.entityHighWater = len(self.indexes) + } + if ipAssocScratchShrinkMin <= cap(self.pairs) && len(self.pairs)*4 < cap(self.pairs) { + self.pairs = make([]ipAssocAggPair, 0, len(self.pairs)) + } +} + +// coalesceIpAssocPairs sorts the pair observations by key and sums duplicates in +// place, returning the compacted slice. +func coalesceIpAssocPairs(pairs []ipAssocAggPair) []ipAssocAggPair { + slices.SortFunc(pairs, func(x ipAssocAggPair, y ipAssocAggPair) int { + switch { + case x.key < y.key: + return -1 + case y.key < x.key: + return 1 + default: + return 0 + } + }) + coalesced := pairs[:0] + for _, pair := range pairs { + if n := len(coalesced); 0 < n && coalesced[n-1].key == pair.key { + coalesced[n-1].count += pair.count + } else { + coalesced = append(coalesced, pair) + } + } + return coalesced +} + +// find is path-compressing union-find over the aggregate indexes. +func (self *ipAssocScratch) find(i uint32) uint32 { + for self.parent[i] != i { + self.parent[i] = self.parent[self.parent[i]] + i = self.parent[i] + } + return i +} + +func (self *ipAssocScratch) union(a uint32, b uint32) { + ra := self.find(a) + rb := self.find(b) + if ra != rb { + self.parent[rb] = ra + } +} + +// growUint32 / growInt32 / growUint64 / growBool size a reusable slice to n, reusing +// the backing array when it is large enough. +func growUint32(s []uint32, n int) []uint32 { + if cap(s) < n { + return make([]uint32, n) + } + return s[:n] +} + +func growInt32(s []int32, n int) []int32 { + if cap(s) < n { + return make([]int32, n) + } + return s[:n] +} + +func growUint64(s []uint64, n int) []uint64 { + if cap(s) < n { + return make([]uint64, n) + } + return s[:n] +} + +func growFloat64(s []float64, n int) []float64 { + if cap(s) < n { + return make([]float64, n) + } + return s[:n] +} + +// clusterIpAssoc merges entities by shared base name into nodes, then greedily +// clusters the nodes: starting from each connected component of the co-association +// graph, the node with the lowest mean association is split off until every member's +// mean association with the rest of the cluster is at least `minMeanAssociation`. +// split-off nodes re-cluster among themselves the same way. +// returns entity -> cluster member entities, for clusters with 2+ entities. +// +// this map-input form is a convenience over the dense scratch pass that +// updateClusters uses directly (it builds a one-off scratch from the maps). +func clusterIpAssoc( + counts map[netip.Addr]uint32, + coCounts map[ipAssocPair]uint32, + baseNames map[netip.Addr][]string, + minMeanAssociation float64, +) map[netip.Addr][]netip.Addr { + scratch := &ipAssocScratch{} + scratch.resetAggregate() + for addr, count := range counts { + i := scratch.index(addr) + scratch.counts[i] += uint64(count) + scratch.baseNames[i] = baseNames[addr] + } + for pair, coCount := range coCounts { + // pairs for entities without counts are dropped (they aged out of all blocks) + i, iOk := scratch.indexes[pair.a] + j, jOk := scratch.indexes[pair.b] + if !iOk || !jOk { + continue + } + scratch.pairs = append(scratch.pairs, ipAssocAggPair{ + key: ipAssocPackNodePair(i, j), + count: coCount, + }) + } + return clusterIpAssocScratch(context.Background(), scratch, minMeanAssociation, defaultMaxComponentNodeCount) +} + +// clusterIpAssocScratch runs the clustering pass over the filled aggregate scratch. +// It checks ctx between phases and returns nil when aborted (teardown), so a pass +// over a large matrix stops promptly instead of running out. +func clusterIpAssocScratch( + ctx context.Context, + scratch *ipAssocScratch, + minMeanAssociation float64, + maxComponentNodes int, +) map[netip.Addr][]netip.Addr { + entityCount := len(scratch.addrs) + members := map[netip.Addr][]netip.Addr{} + if entityCount == 0 { + return members + } + + // aggregate duplicate pair observations (the same pair seen in several blocks) + scratch.pairs = coalesceIpAssocPairs(scratch.pairs) + if ctx.Err() != nil { + return nil + } + + // merge entities that share a base name (union-find; transitive across names) + scratch.parent = growUint32(scratch.parent, entityCount) + for i := range entityCount { + scratch.parent[i] = uint32(i) + } + merged := false + nameFirst := map[string]uint32{} + for i := range entityCount { + for _, baseName := range scratch.baseNames[i] { + if first, ok := nameFirst[baseName]; ok { + scratch.union(first, uint32(i)) + merged = true + } else { + nameFirst[baseName] = uint32(i) + } + } + } + if ctx.Err() != nil { + return nil + } + + // dense node ids per union root, node counts, and node -> member entities (csr) + scratch.nodeIds = growInt32(scratch.nodeIds, entityCount) + for i := range entityCount { + scratch.nodeIds[i] = -1 + } + nodeCount := 0 + for i := range entityCount { + root := scratch.find(uint32(i)) + if scratch.nodeIds[root] < 0 { + scratch.nodeIds[root] = int32(nodeCount) + nodeCount += 1 + } + scratch.nodeIds[i] = scratch.nodeIds[root] + } + scratch.nodeCounts = growUint64(scratch.nodeCounts, nodeCount) + clear(scratch.nodeCounts) + // nodeMemberInit doubles as the running fill cursor while nodeMembers is built + scratch.nodeMemberInit = growInt32(scratch.nodeMemberInit, nodeCount+1) + clear(scratch.nodeMemberInit) + for i := range entityCount { + node := scratch.nodeIds[i] + scratch.nodeCounts[node] += scratch.counts[i] + scratch.nodeMemberInit[node+1] += 1 + } + for node := range nodeCount { + scratch.nodeMemberInit[node+1] += scratch.nodeMemberInit[node] + } + scratch.nodeMembers = growUint32(scratch.nodeMembers, entityCount) + scratch.memberFill = growInt32(scratch.memberFill, nodeCount) + copy(scratch.memberFill, scratch.nodeMemberInit[:nodeCount]) + for i := range entityCount { + node := scratch.nodeIds[i] + scratch.nodeMembers[scratch.memberFill[node]] = uint32(i) + scratch.memberFill[node] += 1 + } + if ctx.Err() != nil { + return nil + } + + // remap the pairs onto node pairs. without a name merge the node ids are the + // identity, so the coalesced entity pairs already are the node pairs + if merged { + remapped := scratch.pairs[:0] + for _, pair := range scratch.pairs { + a := scratch.nodeIds[uint32(pair.key>>32)] + b := scratch.nodeIds[uint32(pair.key)] + if a == b { + continue + } + remapped = append(remapped, ipAssocAggPair{ + key: ipAssocPackNodePair(uint32(a), uint32(b)), + count: pair.count, + }) + } + scratch.pairs = coalesceIpAssocPairs(remapped) + } + if ctx.Err() != nil { + return nil + } + + // connected components via union-find over the node pairs, discovered before any + // adjacency is built: a degenerate component (a burst clique beyond + // maxComponentNodes) is then skipped without paying for its edges at all + scratch.compParent = growUint32(scratch.compParent, nodeCount) + for node := range nodeCount { + scratch.compParent[node] = uint32(node) + } + compFind := func(i uint32) uint32 { + for scratch.compParent[i] != i { + scratch.compParent[i] = scratch.compParent[scratch.compParent[i]] + i = scratch.compParent[i] + } + return i + } + pairNodes := func(pair ipAssocAggPair) (uint32, uint32, bool) { + a := uint32(pair.key >> 32) + b := uint32(pair.key) + return a, b, 0 < pair.count && 0 < min(scratch.nodeCounts[a], scratch.nodeCounts[b]) + } + for _, pair := range scratch.pairs { + if a, b, ok := pairNodes(pair); ok { + ra := compFind(a) + rb := compFind(b) + if ra != rb { + scratch.compParent[rb] = ra + } + } + } + scratch.compSizes = growInt32(scratch.compSizes, nodeCount) + clear(scratch.compSizes) + for node := range nodeCount { + scratch.compSizes[compFind(uint32(node))] += 1 + } + keepNode := func(node uint32) bool { + size := int(scratch.compSizes[compFind(node)]) + return 2 <= size && size <= maxComponentNodes + } + if ctx.Err() != nil { + return nil + } + + // association edges as csr adjacency (node -> (neighbor, p)), kept components only + scratch.degrees = growInt32(scratch.degrees, nodeCount) + clear(scratch.degrees) + edgeCount := 0 + for _, pair := range scratch.pairs { + if a, b, ok := pairNodes(pair); ok && keepNode(a) { + scratch.degrees[a] += 1 + scratch.degrees[b] += 1 + edgeCount += 1 + } + } + scratch.adjInit = growInt32(scratch.adjInit, nodeCount+1) + scratch.adjInit[0] = 0 + for node := range nodeCount { + scratch.adjInit[node+1] = scratch.adjInit[node] + scratch.degrees[node] + } + scratch.adjNodes = growUint32(scratch.adjNodes, 2*edgeCount) + scratch.adjPs = growFloat64(scratch.adjPs, 2*edgeCount) + adjFill := scratch.degrees + clear(adjFill) + for _, pair := range scratch.pairs { + a, b, ok := pairNodes(pair) + if !ok || !keepNode(a) { + continue + } + // coalesced counts can overlap; cap at 1 + p := min(float64(pair.count)/float64(min(scratch.nodeCounts[a], scratch.nodeCounts[b])), 1.0) + ai := scratch.adjInit[a] + adjFill[a] + scratch.adjNodes[ai] = b + scratch.adjPs[ai] = p + adjFill[a] += 1 + bi := scratch.adjInit[b] + adjFill[b] + scratch.adjNodes[bi] = a + scratch.adjPs[bi] = p + adjFill[b] += 1 + } + if ctx.Err() != nil { + return nil + } + + emit := func(clusterNodes []uint32) { + clusterAddrs := []netip.Addr{} + for _, node := range clusterNodes { + for m := scratch.nodeMemberInit[node]; m < scratch.nodeMemberInit[node+1]; m += 1 { + clusterAddrs = append(clusterAddrs, scratch.addrs[scratch.nodeMembers[m]]) + } + } + for _, addr := range clusterAddrs { + members[addr] = clusterAddrs + } + } + + // gather the kept components (csr by dense component id) + scratch.compIds = growInt32(scratch.compIds, nodeCount) + for node := range nodeCount { + scratch.compIds[node] = -1 + } + compCount := 0 + for node := range nodeCount { + if !keepNode(uint32(node)) { + continue + } + root := compFind(uint32(node)) + if scratch.compIds[root] < 0 { + scratch.compIds[root] = int32(compCount) + compCount += 1 + } + scratch.compIds[node] = scratch.compIds[root] + } + scratch.compInit = growInt32(scratch.compInit, compCount+1) + clear(scratch.compInit) + for node := range nodeCount { + if 0 <= scratch.compIds[node] { + scratch.compInit[scratch.compIds[node]+1] += 1 + } + } + for comp := range compCount { + scratch.compInit[comp+1] += scratch.compInit[comp] + } + scratch.compNodes = growUint32(scratch.compNodes, int(scratch.compInit[compCount])) + scratch.memberFill = growInt32(scratch.memberFill, compCount) + copy(scratch.memberFill, scratch.compInit[:compCount]) + for node := range nodeCount { + if comp := scratch.compIds[node]; 0 <= comp { + scratch.compNodes[scratch.memberFill[comp]] = uint32(node) + scratch.memberFill[comp] += 1 + } + } + + // greedy split within each component, re-clustering the split-off pool + for comp := range compCount { + if ctx.Err() != nil { + return nil + } + pool := scratch.compNodes[scratch.compInit[comp]:scratch.compInit[comp+1]] + for 2 <= len(pool) { + cluster, rest := splitIpAssocCluster(scratch, pool, minMeanAssociation) + if 2 <= len(cluster) { + emit(cluster) + } + if len(rest) == len(pool) { + // no progress + break + } + pool = rest + } + } + + // base-name merged nodes are a cluster on their own if not already clustered + for node := range nodeCount { + first := scratch.nodeMemberInit[node] + end := scratch.nodeMemberInit[node+1] + if 2 <= end-first { + if _, ok := members[scratch.addrs[scratch.nodeMembers[first]]]; !ok { + emit([]uint32{uint32(node)}) + } + } + } + + return members +} + +// splitIpAssocCluster removes the member with the lowest mean association until all +// members have mean association >= minMeanAssociation, returning the surviving +// cluster and the removed members. component sizes are bounded by +// MaxComponentNodeCount, so the per-component maps stay small. +func splitIpAssocCluster( + scratch *ipAssocScratch, + pool []uint32, + minMeanAssociation float64, +) (cluster []uint32, rest []uint32) { + cluster = append([]uint32{}, pool...) + inCluster := map[uint32]bool{} + for _, node := range cluster { + inCluster[node] = true + } + // sum of associations to other members + sums := map[uint32]float64{} + for _, node := range cluster { + for ai := scratch.adjInit[node]; ai < scratch.adjInit[node+1]; ai += 1 { + if inCluster[scratch.adjNodes[ai]] { + sums[node] += scratch.adjPs[ai] + } + } + } + + for 2 <= len(cluster) { + minNode := uint32(0) + minMean := 0.0 + found := false + for _, node := range cluster { + mean := sums[node] / float64(len(cluster)-1) + if !found || mean < minMean { + minNode = node + minMean = mean + found = true + } + } + if minMeanAssociation <= minMean { + // all members satisfy the threshold + return + } + // split off the weakest member + delete(inCluster, minNode) + for ai := scratch.adjInit[minNode]; ai < scratch.adjInit[minNode+1]; ai += 1 { + if inCluster[scratch.adjNodes[ai]] { + sums[scratch.adjNodes[ai]] -= scratch.adjPs[ai] + } + } + next := make([]uint32, 0, len(cluster)-1) + for _, node := range cluster { + if node != minNode { + next = append(next, node) + } + } + cluster = next + rest = append(rest, minNode) + } + rest = append(rest, cluster...) + cluster = nil + return +} + +// ipAssocClustersEqual compares two cluster snapshots. cluster member slices are +// shared among their members, so each distinct cluster pairing is verified once +// (set compare) and every further member key checks by backing-array identity — +// O(keys + total members) instead of a set per key. +func ipAssocClustersEqual(a map[netip.Addr][]netip.Addr, b map[netip.Addr][]netip.Addr) bool { + if len(a) != len(b) { + return false + } + matched := map[*netip.Addr]*netip.Addr{} + for addr, aMembers := range a { + bMembers, ok := b[addr] + if !ok || len(aMembers) != len(bMembers) { + return false + } + aFirst := &aMembers[0] + bFirst := &bMembers[0] + if prev, ok := matched[aFirst]; ok { + if prev != bFirst { + return false + } + continue + } + aSet := make(map[netip.Addr]bool, len(aMembers)) + for _, member := range aMembers { + aSet[member] = true + } + for _, member := range bMembers { + if !aSet[member] { + return false + } + } + matched[aFirst] = bFirst + } + return true +} + +// ClusterVersion changes whenever the clustering output changes, +// so callers can cache cluster-derived decisions +func (self *IpAssoc) ClusterVersion() uint64 { + return self.clusters.Load().version +} + +// this is consulted when applying a routing rule that needs to match against an entire cluster not just a single path +func (self *IpAssoc) GetCluster(path *IpPath) []*IpPath { + addr, ok := ipAssocAddr(path.DestinationIp) + if !ok { + return nil + } + members := self.clusters.Load().members[addr] + if len(members) == 0 { + return nil + } + memberPaths := make([]*IpPath, 0, len(members)) + for _, member := range members { + version := 4 + if member.Is6() { + version = 6 + } + memberPaths = append(memberPaths, &IpPath{ + Version: version, + DestinationIp: member.AsSlice(), + }) + } + return memberPaths +} + +// GetClusterAddrs returns the cluster member ips for an entity ip, +// or nil if the entity is not in a multi-member cluster +func (self *IpAssoc) GetClusterAddrs(addr netip.Addr) []netip.Addr { + return self.clusters.Load().members[addr.Unmap()] +} + +func (self *IpAssoc) Close() { + self.cancel() +} diff --git a/ip_datapackets_test.go b/ip_datapackets_test.go new file mode 100644 index 0000000..de458fa --- /dev/null +++ b/ip_datapackets_test.go @@ -0,0 +1,116 @@ +package connect + +import ( + "net" + "testing" + + "github.com/go-playground/assert/v2" + + "github.com/urnetwork/connect/protocol" +) + +// Regression tests for findings 2.1/2.2 from the 2026-07-13 audit of PR #263 +// (ip.go), verifying the fix in commit c4065b4: `StreamState.DataPackets` +// (UDP) and `ConnectionState.DataPackets` (TCP) both used to compute +// `packetByteCount := mtu - headerByteCount` unconditionally. With an MTU +// configured at or below the IP+transport header size, this went to zero or +// negative, causing a divide-by-zero panic in the chunk-count formula +// `(n+packetByteCount-1)/packetByteCount` or a negative slice-capacity panic +// in `make`. Both functions now return an error instead of panicking. + +func newTestStreamState(ipVersion int) *StreamState { + return &StreamState{ + ipVersion: ipVersion, + sourceIp: net.ParseIP("10.0.0.1"), + sourcePort: 1234, + destinationIp: net.ParseIP("10.0.0.2"), + destinationPort: 5678, + provideMode: protocol.ProvideMode_Network, + } +} + +func newTestConnectionState(ipVersion int) *ConnectionState { + return &ConnectionState{ + ipVersion: ipVersion, + sourceIp: net.ParseIP("10.0.0.1"), + sourcePort: 1234, + destinationIp: net.ParseIP("10.0.0.2"), + destinationPort: 5678, + provideMode: protocol.ProvideMode_Network, + receiveSeq: 1, + } +} + +func TestStreamStateDataPacketsRejectsUndersizedMtu(t *testing.T) { + payload := make([]byte, 100) + + // IPv4+UDP header is 28 bytes; mtu == header and mtu < header must both + // error rather than panic (divide-by-zero at mtu==header, negative + // capacity/slice bounds at mtu>4) * 4 + if dataOffset != 24 { + t.Fatalf("SynAck data offset: want 24, got %d", dataOffset) + } +} + +func TestConnectionStatePureAck(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1), + destinationPort: 443, + sendSeq: 2000, + receiveSeq: 1500, + windowSize: 65535, + } + + packet, err := cs.PureAck() + if err != nil { + t.Fatalf("PureAck failed: %v", err) + } + defer MessagePoolReturn(packet) + + ipHeaderLen := int(packet[0]&0x0f) * 4 + tcpBytes := packet[ipHeaderLen:] + + if tcpBytes[13]&tcpFlagAck == 0 { + t.Fatal("ACK flag not set") + } + if tcpBytes[13]&tcpFlagSyn != 0 { + t.Fatal("SYN flag should not be set on PureAck") + } + + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatalf("PureAck TCP checksum verification failed: stored=0x%x", + binary.BigEndian.Uint16(tcpBytes[16:18])) + } +} + +func TestConnectionStateRstAck(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1), + destinationPort: 443, + sendSeq: 2000, + receiveSeq: 1500, + windowSize: 65535, + } + + packet, err := cs.RstAck() + if err != nil { + t.Fatalf("RstAck failed: %v", err) + } + defer MessagePoolReturn(packet) + + ipHeaderLen := int(packet[0]&0x0f) * 4 + tcpBytes := packet[ipHeaderLen:] + + if tcpBytes[13]&tcpFlagRst == 0 { + t.Fatal("RST flag not set") + } + if tcpBytes[13]&tcpFlagAck == 0 { + t.Fatal("ACK flag not set") + } + + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatalf("RstAck TCP checksum verification failed: stored=0x%x", + binary.BigEndian.Uint16(tcpBytes[16:18])) + } +} + +func TestParseIpv4RoundTrip(t *testing.T) { + payload := []byte{0x01, 0x02, 0x03, 0x04} + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + UdpHeaderSize + len(payload)) + writeIpv4Header(packet, IP_PROTOCOL_UDP, net.IPv4(10, 0, 0, 1).To4(), net.IPv4(192, 168, 1, 1).To4()) + udp := packet[Ipv4HeaderSizeWithoutExtensions:] + binary.BigEndian.PutUint16(udp[0:2], 12345) + binary.BigEndian.PutUint16(udp[2:4], 80) + binary.BigEndian.PutUint16(udp[4:6], uint16(UdpHeaderSize+len(payload))) + binary.BigEndian.PutUint16(udp[6:8], 0) + copy(udp[UdpHeaderSize:], payload) + cs := checksumFinish(checksumAdd(0, udp)) + binary.BigEndian.PutUint16(udp[6:8], cs) + defer MessagePoolReturn(packet) + + ipProtocol, srcIP, dstIP, transport, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed") + } + if ipProtocol != IP_PROTOCOL_UDP { + t.Fatalf("protocol: want UDP(17), got %d", ipProtocol) + } + if !srcIP.Equal(net.IPv4(10, 0, 0, 1)) { + t.Fatalf("source IP mismatch: %v", srcIP) + } + if !dstIP.Equal(net.IPv4(192, 168, 1, 1)) { + t.Fatalf("dest IP mismatch: %v", dstIP) + } + var udpPkt parsedUdp + if !parseUdpPacket(srcIP, dstIP, transport, &udpPkt) { + t.Fatal("parseUdpPacket failed") + } + if udpPkt.sourcePort != 12345 { + t.Fatalf("UDP src port: want 12345, got %d", udpPkt.sourcePort) + } + if udpPkt.destinationPort != 80 { + t.Fatalf("UDP dst port: want 80, got %d", udpPkt.destinationPort) + } +} + +func TestParseIpv6RoundTrip(t *testing.T) { + payload := []byte{0x05, 0x06} + src := net.ParseIP("2001:db8::1") + dst := net.ParseIP("2001:db8::2") + packet := MessagePoolGet(Ipv6HeaderSize + UdpHeaderSize + len(payload)) + writeIpv6Header(packet, IP_PROTOCOL_UDP, src, dst) + udp := packet[Ipv6HeaderSize:] + binary.BigEndian.PutUint16(udp[0:2], 33445) + binary.BigEndian.PutUint16(udp[2:4], 53) + binary.BigEndian.PutUint16(udp[4:6], uint16(UdpHeaderSize+len(payload))) + binary.BigEndian.PutUint16(udp[6:8], 0) + copy(udp[UdpHeaderSize:], payload) + cs := checksumFinish(checksumAdd(0, udp)) + binary.BigEndian.PutUint16(udp[6:8], cs) + defer MessagePoolReturn(packet) + + ipProtocol, srcIP, dstIP, transport, ok := parseIpv6(packet) + if !ok { + t.Fatal("parseIpv6 failed") + } + if ipProtocol != IP_PROTOCOL_UDP { + t.Fatalf("protocol: want UDP(17), got %d", ipProtocol) + } + if !srcIP.Equal(src) { + t.Fatalf("source IP mismatch: %v", srcIP) + } + if !dstIP.Equal(dst) { + t.Fatalf("dest IP mismatch: %v", dstIP) + } + var udpPkt parsedUdp + if !parseUdpPacket(srcIP, dstIP, transport, &udpPkt) { + t.Fatal("parseUdpPacket failed") + } + if udpPkt.sourcePort != 33445 { + t.Fatalf("UDP src port: want 33445, got %d", udpPkt.sourcePort) + } + if udpPkt.destinationPort != 53 { + t.Fatalf("UDP dst port: want 53, got %d", udpPkt.destinationPort) + } +} + +func TestParseUdpRoundTrip(t *testing.T) { + payload := []byte("hello udp") + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + UdpHeaderSize + len(payload)) + writeIpv4Header(packet, IP_PROTOCOL_UDP, src.To4(), dst.To4()) + udp := packet[Ipv4HeaderSizeWithoutExtensions:] + binary.BigEndian.PutUint16(udp[0:2], 40000) + binary.BigEndian.PutUint16(udp[2:4], 53) + binary.BigEndian.PutUint16(udp[4:6], uint16(UdpHeaderSize+len(payload))) + copy(udp[UdpHeaderSize:], payload) + cs := transportChecksum(IP_PROTOCOL_UDP, src, dst, udp) + binary.BigEndian.PutUint16(udp[6:8], cs) + defer MessagePoolReturn(packet) + + var udpPkt parsedUdp + ipProtocol, _, _, transport, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed") + } + if ipProtocol != IP_PROTOCOL_UDP { + t.Fatalf("protocol: want UDP, got %d", ipProtocol) + } + if !parseUdpPacket(nil, nil, transport, &udpPkt) { + t.Fatal("parseUdpPacket failed") + } + if udpPkt.sourcePort != 40000 { + t.Fatalf("UDP src port: want 40000, got %d", udpPkt.sourcePort) + } + if udpPkt.destinationPort != 53 { + t.Fatalf("UDP dst port: want 53, got %d", udpPkt.destinationPort) + } +} + +func TestParseTcpRoundTrip(t *testing.T) { + payload := []byte("hello tcp") + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + TcpHeaderSizeWithoutExtensions + len(payload)) + writeIpv4Header(packet, IP_PROTOCOL_TCP, src.To4(), dst.To4()) + tcp := packet[Ipv4HeaderSizeWithoutExtensions:] + writeTcpHeader(tcp, 40000, 80, 100, 200, tcpFlagSyn|tcpFlagAck, 65535, nil) + copy(tcp[TcpHeaderSizeWithoutExtensions:], payload) + cs := transportChecksum(IP_PROTOCOL_TCP, src, dst, tcp) + binary.BigEndian.PutUint16(tcp[16:18], cs) + defer MessagePoolReturn(packet) + + var tcpPkt parsedTcp + ipProtocol, _, _, transport, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed") + } + if ipProtocol != IP_PROTOCOL_TCP { + t.Fatalf("protocol: want TCP(6), got %d", ipProtocol) + } + if !parseTcpPacket(nil, nil, transport, &tcpPkt) { + t.Fatal("parseTcpPacket failed") + } + if tcpPkt.sourcePort != 40000 { + t.Fatalf("TCP src port: want 40000, got %d", tcpPkt.sourcePort) + } + if tcpPkt.destinationPort != 80 { + t.Fatalf("TCP dst port: want 80, got %d", tcpPkt.destinationPort) + } + if tcpPkt.seq != 100 { + t.Fatalf("TCP seq: want 100, got %d", tcpPkt.seq) + } + if !tcpPkt.ack { + t.Fatal("ACK flag not set") + } + if tcpPkt.ackNumber != 200 { + t.Fatalf("TCP ack: want 200, got %d", tcpPkt.ackNumber) + } + if !tcpPkt.syn || !tcpPkt.ack { + t.Fatal("SYN+ACK flags not set") + } +} + +func TestParseTcpSynWithWindowScale(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + optsBytes := make([]byte, 8) + optsBytes[0] = 3 + optsBytes[1] = 3 + optsBytes[2] = 7 + tcpOptionLen := len(optsBytes) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + TcpHeaderSizeWithoutExtensions + tcpOptionLen) + writeIpv4Header(packet, IP_PROTOCOL_TCP, src.To4(), dst.To4()) + tcp := packet[Ipv4HeaderSizeWithoutExtensions:] + writeTcpHeader(tcp, 40000, 80, 100, 200, tcpFlagSyn, 65535, optsBytes) + fullTcp := packet[Ipv4HeaderSizeWithoutExtensions:] + cs := transportChecksum(IP_PROTOCOL_TCP, src, dst, fullTcp) + binary.BigEndian.PutUint16(tcp[16:18], cs) + defer MessagePoolReturn(packet) + + var tcpPkt parsedTcp + _, _, _, transport, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed") + } + if !parseTcpPacket(nil, nil, transport, &tcpPkt) { + t.Fatal("parseTcpPacket failed") + } + if !tcpPkt.syn { + t.Fatal("SYN flag not set") + } +} + +func TestParseTcpWindowScaleOpts_Nil(t *testing.T) { + found, shift := ParseTcpWindowScaleOpts(nil) + if found { + t.Fatal("expected no window scale found for nil options") + } + if shift != 0 { + t.Fatalf("shift: want 0, got %d", shift) + } +} + +func TestParseTcpWindowScaleOpts_Empty(t *testing.T) { + found, shift := ParseTcpWindowScaleOpts([]byte{}) + if found { + t.Fatal("expected no window scale for empty options") + } + if shift != 0 { + t.Fatalf("shift: want 0, got %d", shift) + } +} + +func TestParseTcpWindowScaleOpts_WithWindowScale(t *testing.T) { + opts := []byte{3, 3, 7} + found, shift := ParseTcpWindowScaleOpts(opts) + if !found { + t.Fatal("expected to find window scale") + } + if shift != 7 { + t.Fatalf("window scale: want 7, got %d", shift) + } +} + +func TestParseTcpWindowScaleOpts_WindowScaleCappedAt14(t *testing.T) { + opts := []byte{3, 3, 20} + found, shift := ParseTcpWindowScaleOpts(opts) + if !found { + t.Fatal("expected to find window scale") + } + if shift != 14 { + t.Fatalf("window scale: want 14 (capped), got %d", shift) + } +} + +func TestParseTcpWindowScaleOpts_NopBeforeWindowScale(t *testing.T) { + opts := []byte{1, 3, 3, 7} + found, shift := ParseTcpWindowScaleOpts(opts) + if !found { + t.Fatal("expected to find window scale after NOP") + } + if shift != 7 { + t.Fatalf("window scale: want 7, got %d", shift) + } +} + +func TestParseTcpWindowScaleOpts_Eol(t *testing.T) { + opts := []byte{0, 3, 3, 7} + found, shift := ParseTcpWindowScaleOpts(opts) + if found { + t.Fatal("expected no window scale (EOL stops parsing)") + } + if shift != 0 { + t.Fatalf("shift: want 0, got %d", shift) + } +} + +func TestParseTcpWindowScaleOpts_EolAfterOptions(t *testing.T) { + opts := []byte{3, 3, 5, 0} + found, shift := ParseTcpWindowScaleOpts(opts) + if !found { + t.Fatal("expected to find window scale before EOL") + } + if shift != 5 { + t.Fatalf("window scale: want 5, got %d", shift) + } +} + +func TestParseTcpWindowScaleOpts_UnknownOption(t *testing.T) { + opts := []byte{5, 3, 1, 3, 3, 7} + found, shift := ParseTcpWindowScaleOpts(opts) + if !found { + t.Fatal("expected to find window scale after unknown option") + } + if shift != 7 { + t.Fatalf("window scale: want 7, got %d", shift) + } +} + +func TestParseTcpWindowScaleOpts_MultipleNops(t *testing.T) { + opts := []byte{1, 1, 1, 3, 3, 3} + found, shift := ParseTcpWindowScaleOpts(opts) + if !found { + t.Fatal("expected to find window scale after multiple NOPs") + } + if shift != 3 { + t.Fatalf("window scale: want 3, got %d", shift) + } +} + +func TestWriteIpv4RoundTrip(t *testing.T) { + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + UdpHeaderSize) + writeIpv4Header(packet, IP_PROTOCOL_UDP, net.IPv4(10, 0, 0, 1).To4(), net.IPv4(192, 168, 1, 1).To4()) + defer MessagePoolReturn(packet) + + _, srcIP, dstIP, _, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed on self-constructed header") + } + if !srcIP.Equal(net.IPv4(10, 0, 0, 1)) { + t.Fatalf("source IP: want 10.0.0.1, got %v", srcIP) + } + if !dstIP.Equal(net.IPv4(192, 168, 1, 1)) { + t.Fatalf("dest IP: want 192.168.1.1, got %v", dstIP) + } + // the IPv4 header checksum covers only the 20-byte header, not whatever + // follows in the allocated buffer (pool buffers aren't guaranteed zeroed, + // so verifying over the full buffer is flaky depending on pool reuse) + if checksumFinish(checksumAdd(0, packet[:Ipv4HeaderSizeWithoutExtensions])) != 0 { + t.Fatalf("IPv4 checksum verification failed") + } +} + +func TestWriteIpv6RoundTrip(t *testing.T) { + src := net.ParseIP("2001:db8::1") + dst := net.ParseIP("2001:db8::2") + packet := MessagePoolGet(Ipv6HeaderSize) + writeIpv6Header(packet, IP_PROTOCOL_UDP, src, dst) + defer MessagePoolReturn(packet) + + _, srcIP, dstIP, _, ok := parseIpv6(packet) + if !ok { + t.Fatal("parseIpv6 failed on self-constructed header") + } + if !srcIP.Equal(src) { + t.Fatalf("source IP: want 2001:db8::1, got %v", srcIP) + } + if !dstIP.Equal(dst) { + t.Fatalf("dest IP: want 2001:db8::2, got %v", dstIP) + } +} + +func TestWriteTcpRoundTrip(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + opts := []byte{3, 3, 7, 0} + tcpOptionLen := len(opts) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + TcpHeaderSizeWithoutExtensions + tcpOptionLen) + writeIpv4Header(packet, IP_PROTOCOL_TCP, src.To4(), dst.To4()) + tcp := packet[Ipv4HeaderSizeWithoutExtensions:] + writeTcpHeader(tcp, 40000, 80, 100, 200, tcpFlagSyn|tcpFlagAck, 65535, opts) + fullTcp := packet[Ipv4HeaderSizeWithoutExtensions:] + cs := transportChecksum(IP_PROTOCOL_TCP, src, dst, fullTcp) + binary.BigEndian.PutUint16(tcp[16:18], cs) + defer MessagePoolReturn(packet) + + var tcpPkt parsedTcp + _, _, _, transport, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed") + } + if !parseTcpPacket(nil, nil, transport, &tcpPkt) { + t.Fatal("parseTcpPacket failed") + } + if tcpPkt.sourcePort != 40000 { + t.Fatalf("TCP src port: want 40000, got %d", tcpPkt.sourcePort) + } + if tcpPkt.destinationPort != 80 { + t.Fatalf("TCP dst port: want 80, got %d", tcpPkt.destinationPort) + } + if !tcpPkt.syn || !tcpPkt.ack { + t.Fatal("SYN+ACK flags not set") + } +} + +func TestParseUdpTooShort(t *testing.T) { + var udp parsedUdp + if parseUdpPacket(nil, nil, []byte{0x00}, &udp) { + t.Fatal("expected parseUdpPacket to fail on short buffer") + } +} + +func TestParseTcpTooShort(t *testing.T) { + var tcp parsedTcp + if parseTcpPacket(nil, nil, []byte{0x00}, &tcp) { + t.Fatal("expected parseTcpPacket to fail on short buffer") + } +} + +func TestParseTcpMalformedOptions(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + TcpHeaderSizeWithoutExtensions) + writeIpv4Header(packet, IP_PROTOCOL_TCP, src.To4(), dst.To4()) + tcp := packet[Ipv4HeaderSizeWithoutExtensions:] + tcp[12] = 0xFF + tcp[13] = tcpFlagAck + cs := transportChecksum(IP_PROTOCOL_TCP, src, dst, tcp) + binary.BigEndian.PutUint16(tcp[16:18], cs) + defer MessagePoolReturn(packet) + + var tcpPkt parsedTcp + _, _, _, transport, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed") + } + if parseTcpPacket(nil, nil, transport, &tcpPkt) { + t.Fatal("expected parseTcpPacket to fail with malformed data offset") + } +} + +func TestEmptyPacketGuard(t *testing.T) { + _, _, _, _, ok := parseIpv4(nil) + if ok { + t.Fatal("expected parseIpv4 to fail on nil") + } + _, _, _, _, ok = parseIpv4([]byte{}) + if ok { + t.Fatal("expected parseIpv4 to fail on empty") + } + _, _, _, _, ok = parseIpv6(nil) + if ok { + t.Fatal("expected parseIpv6 to fail on nil") + } + _, _, _, _, ok = parseIpv6([]byte{}) + if ok { + t.Fatal("expected parseIpv6 to fail on empty") + } +} + +func TestFlowHashIPv4Consistent(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + TcpHeaderSizeWithoutExtensions) + writeIpv4Header(packet, IP_PROTOCOL_TCP, src.To4(), dst.To4()) + tcp := packet[Ipv4HeaderSizeWithoutExtensions:] + writeTcpHeader(tcp, 40000, 80, 100, 200, tcpFlagSyn, 65535, nil) + cs := transportChecksum(IP_PROTOCOL_TCP, src, dst, tcp) + binary.BigEndian.PutUint16(tcp[16:18], cs) + defer MessagePoolReturn(packet) + + h1 := flowHash(packet) + h2 := flowHash(packet) + if h1 != h2 { + t.Fatalf("flowHash not consistent: %d vs %d", h1, h2) + } +} + +func TestFlowHashIPv4DifferentPort(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + packet1 := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + TcpHeaderSizeWithoutExtensions) + writeIpv4Header(packet1, IP_PROTOCOL_TCP, src.To4(), dst.To4()) + tcp1 := packet1[Ipv4HeaderSizeWithoutExtensions:] + writeTcpHeader(tcp1, 40000, 80, 100, 200, tcpFlagSyn, 65535, nil) + cs1 := transportChecksum(IP_PROTOCOL_TCP, src, dst, tcp1) + binary.BigEndian.PutUint16(tcp1[16:18], cs1) + defer MessagePoolReturn(packet1) + + packet2 := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + TcpHeaderSizeWithoutExtensions) + writeIpv4Header(packet2, IP_PROTOCOL_TCP, src.To4(), dst.To4()) + tcp2 := packet2[Ipv4HeaderSizeWithoutExtensions:] + writeTcpHeader(tcp2, 40001, 80, 100, 200, tcpFlagSyn, 65535, nil) + cs2 := transportChecksum(IP_PROTOCOL_TCP, src, dst, tcp2) + binary.BigEndian.PutUint16(tcp2[16:18], cs2) + defer MessagePoolReturn(packet2) + + flowHash(packet1) + flowHash(packet2) +} + +func TestFlowHashIPv6Consistent(t *testing.T) { + src := net.ParseIP("2001:db8::1") + dst := net.ParseIP("2001:db8::2") + packet := MessagePoolGet(Ipv6HeaderSize + TcpHeaderSizeWithoutExtensions) + writeIpv6Header(packet, IP_PROTOCOL_TCP, src, dst) + tcp := packet[Ipv6HeaderSize:] + writeTcpHeader(tcp, 40000, 80, 100, 200, tcpFlagSyn, 65535, nil) + defer MessagePoolReturn(packet) + + h1 := flowHash(packet) + h2 := flowHash(packet) + if h1 != h2 { + t.Fatalf("flowHash not consistent for IPv6: %d vs %d", h1, h2) + } +} + +func TestFlowHashShortPacket(t *testing.T) { + if h := flowHash([]byte{0x45}); h != 0 { + t.Fatalf("expected 0 for short packet, got %d", h) + } + if h := flowHash([]byte{}); h != 0 { + t.Fatalf("expected 0 for empty packet, got %d", h) + } +} + +func TestSynAckDataOffsetMatchesOptions(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1).To4(), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1).To4(), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: true, + windowScale: 7, + } + + packet, err := cs.SynAck() + if err != nil { + t.Fatalf("SynAck failed: %v", err) + } + defer MessagePoolReturn(packet) + + ipHeaderLen := int(packet[0]&0x0f) * 4 + tcpBytes := packet[ipHeaderLen:] + dataOffsetBytes := int(tcpBytes[12]>>4) * 4 + actualLen := len(tcpBytes) + if dataOffsetBytes != actualLen { + t.Fatalf("SynAck: header claims %d bytes but segment is %d bytes", dataOffsetBytes, actualLen) + } +} + +func TestUdpChecksumNonZero(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + UdpHeaderSize) + writeIpv4Header(packet, IP_PROTOCOL_UDP, src.To4(), dst.To4()) + udp := packet[Ipv4HeaderSizeWithoutExtensions:] + binary.BigEndian.PutUint16(udp[0:2], 40000) + binary.BigEndian.PutUint16(udp[2:4], 53) + binary.BigEndian.PutUint16(udp[4:6], UdpHeaderSize) + binary.BigEndian.PutUint16(udp[6:8], 0) + cs := transportChecksum(IP_PROTOCOL_UDP, src, dst, udp) + binary.BigEndian.PutUint16(udp[6:8], cs) + defer MessagePoolReturn(packet) + + if cs == 0 { + t.Fatal("UDP checksum should be non-zero for this packet") + } + + if transportChecksum(IP_PROTOCOL_UDP, src, dst, udp) != 0 { + t.Fatal("UDP checksum verification failed") + } +} + +func TestUdpChecksumZero(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + UdpHeaderSize) + writeIpv4Header(packet, IP_PROTOCOL_UDP, src.To4(), dst.To4()) + udp := packet[Ipv4HeaderSizeWithoutExtensions:] + binary.BigEndian.PutUint16(udp[0:2], 40000) + binary.BigEndian.PutUint16(udp[2:4], 53) + binary.BigEndian.PutUint16(udp[4:6], UdpHeaderSize) + binary.BigEndian.PutUint16(udp[6:8], 0) + defer MessagePoolReturn(packet) + + var udpPkt parsedUdp + _, _, _, transport, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed") + } + if !parseUdpPacket(nil, nil, transport, &udpPkt) { + t.Fatal("parseUdpPacket failed with zero checksum (RFC 768 allows)") + } +} + +func TestConnectionStateFinAck(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1), + destinationPort: 443, + sendSeq: 2000, + receiveSeq: 1500, + windowSize: 65535, + } + + packet, err := cs.FinAck() + if err != nil { + t.Fatalf("FinAck failed: %v", err) + } + defer MessagePoolReturn(packet) + + ipHeaderLen := int(packet[0]&0x0f) * 4 + tcpBytes := packet[ipHeaderLen:] + + if tcpBytes[13]&tcpFlagFin == 0 { + t.Fatal("FIN flag not set") + } + if tcpBytes[13]&tcpFlagAck == 0 { + t.Fatal("ACK flag not set") + } + + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("FinAck TCP checksum verification failed") + } +} + +func TestSynAckIPv6(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 6, + sourceIp: net.ParseIP("2001:db8::1"), + sourcePort: 40000, + destinationIp: net.ParseIP("2001:db8::2"), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: true, + windowScale: 7, + } + + packet, err := cs.SynAck() + if err != nil { + t.Fatalf("SynAck IPv6 failed: %v", err) + } + defer MessagePoolReturn(packet) + + _, srcIP, dstIP, _, ok := parseIpv6(packet) + if !ok { + t.Fatal("parseIpv6 failed") + } + if !srcIP.Equal(cs.destinationIp) { + t.Fatalf("IPv6 source IP mismatch: %v", srcIP) + } + if !dstIP.Equal(cs.sourceIp) { + t.Fatalf("IPv6 dest IP mismatch: %v", dstIP) + } + + if packet[0]>>4 != 6 { + t.Fatal("not an IPv6 packet") + } + + tcpBytes := packet[Ipv6HeaderSize:] + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("IPv6 SynAck TCP checksum verification failed") + } + + dataOffset := int(tcpBytes[12]>>4) * 4 + actualLen := len(tcpBytes) + if dataOffset != actualLen { + t.Fatalf("IPv6 SynAck: header claims %d bytes but segment is %d bytes", dataOffset, actualLen) + } +} + +func TestConnectionStatePureAckIPv6(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 6, + sourceIp: net.ParseIP("2001:db8::1"), + sourcePort: 40000, + destinationIp: net.ParseIP("2001:db8::2"), + destinationPort: 443, + sendSeq: 2000, + receiveSeq: 1500, + windowSize: 65535, + } + + packet, err := cs.PureAck() + if err != nil { + t.Fatalf("PureAck IPv6 failed: %v", err) + } + defer MessagePoolReturn(packet) + + if packet[0]>>4 != 6 { + t.Fatal("not an IPv6 packet") + } + + tcpBytes := packet[Ipv6HeaderSize:] + if tcpBytes[13]&tcpFlagAck == 0 { + t.Fatal("ACK flag not set") + } + if tcpBytes[13]&tcpFlagSyn != 0 { + t.Fatal("SYN flag should not be set on PureAck") + } + + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("IPv6 PureAck TCP checksum verification failed") + } +} + +func TestConnectionStateRstAckIPv6(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 6, + sourceIp: net.ParseIP("2001:db8::1"), + sourcePort: 40000, + destinationIp: net.ParseIP("2001:db8::2"), + destinationPort: 443, + sendSeq: 2000, + receiveSeq: 1500, + windowSize: 65535, + } + + packet, err := cs.RstAck() + if err != nil { + t.Fatalf("RstAck IPv6 failed: %v", err) + } + defer MessagePoolReturn(packet) + + if packet[0]>>4 != 6 { + t.Fatal("not an IPv6 packet") + } + + tcpBytes := packet[Ipv6HeaderSize:] + if tcpBytes[13]&tcpFlagRst == 0 { + t.Fatal("RST flag not set") + } + if tcpBytes[13]&tcpFlagAck == 0 { + t.Fatal("ACK flag not set") + } + + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("IPv6 RstAck TCP checksum verification failed") + } +} + +func TestConnectionStateFinAckIPv6(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 6, + sourceIp: net.ParseIP("2001:db8::1"), + sourcePort: 40000, + destinationIp: net.ParseIP("2001:db8::2"), + destinationPort: 443, + sendSeq: 2000, + receiveSeq: 1500, + windowSize: 65535, + } + + packet, err := cs.FinAck() + if err != nil { + t.Fatalf("FinAck IPv6 failed: %v", err) + } + defer MessagePoolReturn(packet) + + if packet[0]>>4 != 6 { + t.Fatal("not an IPv6 packet") + } + + tcpBytes := packet[Ipv6HeaderSize:] + if tcpBytes[13]&tcpFlagFin == 0 { + t.Fatal("FIN flag not set") + } + if tcpBytes[13]&tcpFlagAck == 0 { + t.Fatal("ACK flag not set") + } + + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("IPv6 FinAck TCP checksum verification failed") + } +} + +func TestTcpPacketIPv6WithPayload(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 6, + sourceIp: net.ParseIP("2001:db8::1"), + sourcePort: 40000, + destinationIp: net.ParseIP("2001:db8::2"), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: false, + } + + payload := []byte("hello ipv6") + packet := cs.tcpPacket(payload, 100) + defer MessagePoolReturn(packet) + + if packet[0]>>4 != 6 { + t.Fatal("not an IPv6 packet") + } + + _, srcIP, dstIP, _, ok := parseIpv6(packet) + if !ok { + t.Fatal("parseIpv6 failed") + } + if !srcIP.Equal(cs.destinationIp) { + t.Fatalf("IPv6 source IP mismatch: %v", srcIP) + } + if !dstIP.Equal(cs.sourceIp) { + t.Fatalf("IPv6 dest IP mismatch: %v", dstIP) + } + + tcpBytes := packet[Ipv6HeaderSize:] + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("IPv6 tcpPacket checksum verification failed") + } +} + +func TestSynAckNoWindowScale(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1).To4(), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1).To4(), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: false, + } + + packet, err := cs.SynAck() + if err != nil { + t.Fatalf("SynAck (no WS) failed: %v", err) + } + defer MessagePoolReturn(packet) + + ipHeaderLen := int(packet[0]&0x0f) * 4 + tcpBytes := packet[ipHeaderLen:] + + if tcpBytes[13]&(tcpFlagSyn|tcpFlagAck) != (tcpFlagSyn | tcpFlagAck) { + t.Fatal("SYN+ACK flags not set") + } + + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("SynAck (no WS) TCP checksum verification failed") + } + + dataOffset := int(tcpBytes[12]>>4) * 4 + if dataOffset != TcpHeaderSizeWithoutExtensions { + t.Fatalf("SynAck (no WS) data offset: want %d, got %d", TcpHeaderSizeWithoutExtensions, dataOffset) + } +} + +func TestDataPacketsSinglePacket(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1).To4(), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1).To4(), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: false, + } + + payload := []byte("small") + packets, err := cs.DataPackets(payload, len(payload), 1500) + if err != nil { + t.Fatalf("DataPackets failed: %v", err) + } + if len(packets) != 1 { + t.Fatalf("expected 1 packet, got %d", len(packets)) + } + defer MessagePoolReturn(packets[0]) + + ipHeaderLen := int(packets[0][0]&0x0f) * 4 + tcpBytes := packets[0][ipHeaderLen:] + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("DataPackets single packet checksum failed") + } +} + +func TestDataPacketsMultiMtu(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1).To4(), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1).To4(), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: false, + } + + mtu := 256 + maxPayloadPerPacket := mtu - Ipv4HeaderSizeWithoutExtensions - TcpHeaderSizeWithoutExtensions + payload := make([]byte, maxPayloadPerPacket*3+1) + for i := range payload { + payload[i] = byte(i) + } + + packets, err := cs.DataPackets(payload, len(payload), mtu) + if err != nil { + t.Fatalf("DataPackets failed: %v", err) + } + + expectedCount := (len(payload) + maxPayloadPerPacket - 1) / maxPayloadPerPacket + if len(packets) != expectedCount { + t.Fatalf("expected %d packets for %d bytes at mtu %d, got %d", + expectedCount, len(payload), mtu, len(packets)) + } + for i, pkt := range packets { + defer MessagePoolReturn(pkt) + ipHeaderLen := int(pkt[0]&0x0f) * 4 + tcpBytes := pkt[ipHeaderLen:] + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatalf("DataPackets packet %d checksum failed", i) + } + } +} + +func TestDataPacketsExactFit(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1).To4(), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1).To4(), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: false, + } + + mtu := 256 + maxPayloadPerPacket := mtu - Ipv4HeaderSizeWithoutExtensions - TcpHeaderSizeWithoutExtensions + payload := make([]byte, maxPayloadPerPacket) + + packets, err := cs.DataPackets(payload, len(payload), mtu) + if err != nil { + t.Fatalf("DataPackets failed: %v", err) + } + if len(packets) != 1 { + t.Fatalf("expected 1 packet for exact MTU fit, got %d", len(packets)) + } + defer MessagePoolReturn(packets[0]) + + ipHeaderLen := int(packets[0][0]&0x0f) * 4 + tcpBytes := packets[0][ipHeaderLen:] + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("DataPackets exact-fit checksum failed") + } +} + +func TestDataPacketsIPv6(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 6, + sourceIp: net.ParseIP("2001:db8::1"), + sourcePort: 40000, + destinationIp: net.ParseIP("2001:db8::2"), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: false, + } + + payload := make([]byte, 100) + packets, err := cs.DataPackets(payload, len(payload), 1500) + if err != nil { + t.Fatalf("DataPackets IPv6 failed: %v", err) + } + if len(packets) != 1 { + t.Fatalf("expected 1 packet for IPv6, got %d", len(packets)) + } + defer MessagePoolReturn(packets[0]) + + if packets[0][0]>>4 != 6 { + t.Fatal("not an IPv6 packet") + } + + tcpBytes := packets[0][Ipv6HeaderSize:] + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("DataPackets IPv6 checksum failed") + } +} + +func TestTcpPacketEmptyPayload(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1).To4(), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1).To4(), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: false, + } + + packet := cs.tcpPacket(nil, 100) + defer MessagePoolReturn(packet) + + ipHeaderLen := int(packet[0]&0x0f) * 4 + tcpBytes := packet[ipHeaderLen:] + if len(tcpBytes) != TcpHeaderSizeWithoutExtensions { + t.Fatalf("empty payload tcpPacket: want %d bytes, got %d", + TcpHeaderSizeWithoutExtensions, len(tcpBytes)) + } + + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("empty payload tcpPacket checksum failed") + } +} + +func TestFlowHashUdp(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + packet := MessagePoolGet(Ipv4HeaderSizeWithoutExtensions + UdpHeaderSize) + writeIpv4Header(packet, IP_PROTOCOL_UDP, src.To4(), dst.To4()) + udp := packet[Ipv4HeaderSizeWithoutExtensions:] + binary.BigEndian.PutUint16(udp[0:2], 40000) + binary.BigEndian.PutUint16(udp[2:4], 53) + binary.BigEndian.PutUint16(udp[4:6], UdpHeaderSize) + binary.BigEndian.PutUint16(udp[6:8], 0) + defer MessagePoolReturn(packet) + + h := flowHash(packet) + if h == 0 { + t.Fatal("flowHash returned 0 for valid UDP packet") + } + h2 := flowHash(packet) + if h != h2 { + t.Fatal("flowHash not consistent for UDP") + } +} + +func TestTransportChecksumIPv6(t *testing.T) { + src := net.ParseIP("2001:db8::1") + dst := net.ParseIP("2001:db8::2") + transport := make([]byte, TcpHeaderSizeWithoutExtensions) + + cs := transportChecksum(IP_PROTOCOL_TCP, src, dst, transport) + if cs == 0 { + t.Fatal("expected non-zero transport checksum for IPv6 with zeroed transport") + } + + cs2 := transportChecksum(IP_PROTOCOL_TCP, src, dst, transport) + if cs != cs2 { + t.Fatal("transportChecksum IPv6 not consistent") + } +} + +func TestParseIpv4TooShort(t *testing.T) { + _, _, _, _, ok := parseIpv4(make([]byte, 19)) + if ok { + t.Fatal("expected parseIpv4 to fail on 19-byte buffer") + } + + _, _, _, _, ok = parseIpv4(make([]byte, 18)) + if ok { + t.Fatal("expected parseIpv4 to fail on 18-byte buffer") + } + + _, _, _, _, ok = parseIpv4(make([]byte, 21)) + if ok { + t.Fatal("expected parseIpv4 to fail on 21-byte with header=24 (bogus ihl)") + } +} + +func TestParseIpv4InvalidLength(t *testing.T) { + buf := make([]byte, Ipv4HeaderSizeWithoutExtensions+UdpHeaderSize) + buf[0] = 0x45 + binary.BigEndian.PutUint16(buf[2:4], uint16(Ipv4HeaderSizeWithoutExtensions-1)) + _, _, _, _, ok := parseIpv4(buf) + if ok { + t.Fatal("expected parseIpv4 to fail when total length < header length") + } + + binary.BigEndian.PutUint16(buf[2:4], uint16(Ipv4HeaderSizeWithoutExtensions+UdpHeaderSize+1)) + _, _, _, _, ok = parseIpv4(buf) + if ok { + t.Fatal("expected parseIpv4 to fail when total length > buffer length") + } +} + +func TestParseIpv6TooShort(t *testing.T) { + _, _, _, _, ok := parseIpv6(make([]byte, 39)) + if ok { + t.Fatal("expected parseIpv6 to fail on 39-byte buffer") + } + + _, _, _, _, ok = parseIpv6(make([]byte, 0)) + if ok { + t.Fatal("expected parseIpv6 to fail on empty buffer") + } +} + +func TestWriteIpv4HeaderTo4Required(t *testing.T) { + src := net.IPv4(10, 0, 0, 1) + dst := net.IPv4(192, 168, 1, 1) + + packet := make([]byte, Ipv4HeaderSizeWithoutExtensions+UdpHeaderSize) + writeIpv4Header(packet, IP_PROTOCOL_UDP, src.To4(), dst.To4()) + + if checksumFinish(checksumAdd(0, packet)) != 0 { + t.Fatal("IPv4 checksum verification failed with To4() IPs") + } + + _, srcIP, dstIP, _, ok := parseIpv4(packet) + if !ok { + t.Fatal("parseIpv4 failed on header built with To4() IPs") + } + if !srcIP.Equal(src) { + t.Fatalf("source IP mismatch: %v", srcIP) + } + if !dstIP.Equal(dst) { + t.Fatalf("dest IP mismatch: %v", dstIP) + } +} + +func TestParseTcpWindowScaleCombined(t *testing.T) { + tests := []struct { + name string + opts []byte + wantFound bool + wantShift uint32 + }{ + {"MSS then WS", []byte{2, 4, 0x05, 0xdc, 3, 3, 7}, true, 7}, + {"SACK permitted then WS", []byte{4, 2, 3, 3, 5}, true, 5}, + {"NOP NOP MSS WS", []byte{1, 1, 2, 4, 0x05, 0xdc, 3, 3, 7}, true, 7}, + {"WS at shift=0", []byte{3, 3, 0}, true, 0}, + {"WS at max=14", []byte{3, 3, 14}, true, 14}, + {"only unknown opts", []byte{5, 3, 1, 7, 3, 1}, false, 0}, + {"short WS (len=2)", []byte{3, 2, 7}, false, 0}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + found, shift := ParseTcpWindowScaleOpts(tc.opts) + if found != tc.wantFound { + t.Fatalf("found: want %v, got %v", tc.wantFound, found) + } + if shift != tc.wantShift { + t.Fatalf("shift: want %d, got %d", tc.wantShift, shift) + } + }) + } +} + +func TestFinAckIPv6ChecksumCoverage(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 6, + sourceIp: net.ParseIP("2001:db8::1"), + sourcePort: 40000, + destinationIp: net.ParseIP("2001:db8::2"), + destinationPort: 443, + sendSeq: 2000, + receiveSeq: 1500, + windowSize: 65535, + } + + packet, err := cs.FinAck() + if err != nil { + t.Fatalf("FinAck IPv6 failed: %v", err) + } + defer MessagePoolReturn(packet) + + tcpBytes := packet[Ipv6HeaderSize:] + if transportChecksum(IP_PROTOCOL_TCP, cs.destinationIp, cs.sourceIp, tcpBytes) != 0 { + t.Fatal("FinAck IPv6 checksum verification failed") + } +} + +func TestSeqNumIncrementAcrossDataPackets(t *testing.T) { + cs := &ConnectionState{ + ipVersion: 4, + sourceIp: net.IPv4(10, 0, 0, 1).To4(), + sourcePort: 40000, + destinationIp: net.IPv4(192, 168, 1, 1).To4(), + destinationPort: 443, + sendSeq: 100, + receiveSeq: 200, + windowSize: 65535, + enableWindowScale: false, + } + + mtu := 256 + maxPayloadPerPacket := mtu - Ipv4HeaderSizeWithoutExtensions - TcpHeaderSizeWithoutExtensions + payload := make([]byte, maxPayloadPerPacket*2) + for i := range payload { + payload[i] = byte(i) + } + + packets, err := cs.DataPackets(payload, len(payload), mtu) + if err != nil { + t.Fatalf("DataPackets failed: %v", err) + } + if len(packets) < 2 { + t.Fatalf("expected at least 2 packets, got %d", len(packets)) + } + for _, pkt := range packets { + defer MessagePoolReturn(pkt) + } + + p0tcp := packets[0][Ipv4HeaderSizeWithoutExtensions:] + p1tcp := packets[1][Ipv4HeaderSizeWithoutExtensions:] + + seq0 := binary.BigEndian.Uint32(p0tcp[4:8]) + seq1 := binary.BigEndian.Uint32(p1tcp[4:8]) + + payloadLen0 := len(p0tcp) - TcpHeaderSizeWithoutExtensions + expectedSeq1 := seq0 + uint32(payloadLen0) + if seq1 != expectedSeq1 { + t.Fatalf("seq not incremented by payload length: seq0=%d + payload=%d, got seq1=%d", + seq0, payloadLen0, seq1) + } +} diff --git a/ip_remote_multi_client.go b/ip_remote_multi_client.go index bd8405e..e39cf72 100644 --- a/ip_remote_multi_client.go +++ b/ip_remote_multi_client.go @@ -190,6 +190,8 @@ func DefaultMultiClientSettings() *MultiClientSettings { IngressSecurityPolicyGenerator: DefaultIngressSecurityPolicyWithStats, EgressSecurityPolicyGenerator: DefaultEgressSecurityPolicyWithStats, + EventEpoch: 1 * time.Second, + RemoteUserNatMultiClientMonitorSettings: *DefaultRemoteUserNatMultiClientMonitorSettings(), } } @@ -282,6 +284,8 @@ type MultiClientSettings struct { // nil resolves to DefaultLogger(). Log Logger + EventEpoch time.Duration + RemoteUserNatMultiClientMonitorSettings } @@ -360,6 +364,16 @@ type receivePacket struct { Pooled bool } +type ServerNameLookup interface { + ServerNames(ip string) []string +} + +type multiClientConfig struct { + performanceProfile *PerformanceProfile + localSecurityBypass bool + serverNameLookup ServerNameLookup +} + type RemoteUserNatMultiClient struct { ctx context.Context cancel context.CancelFunc @@ -389,8 +403,11 @@ type RemoteUserNatMultiClient struct { affinityIp6Paths map[Ip6Path]map[Ip6Path]time.Time clientUpdates map[*multiClientChannel]map[*multiClientChannelUpdate]bool - performanceProfile *PerformanceProfile - localSecurityBypass bool + // config is an immutable snapshot of the rarely-changed routing config + // (performance profile + local security bypass). it is rebuilt under + // stateLock by the setters and read lock-free by selectWindowTypes, the + // affinity selection, and the SendPacket drop path. + config atomic.Pointer[multiClientConfig] localUserNat *LocalUserNat localUserNatUnsub func() @@ -448,9 +465,13 @@ func NewRemoteUserNatMultiClient( affinityIp4Paths: map[Ip4Path]map[Ip4Path]time.Time{}, affinityIp6Paths: map[Ip6Path]map[Ip6Path]time.Time{}, clientUpdates: map[*multiClientChannel]map[*multiClientChannelUpdate]bool{}, - performanceProfile: settings.DefaultPerformanceProfile, localUserNat: localUserNat, } + multiClient.config.Store(&multiClientConfig{ + performanceProfile: settings.DefaultPerformanceProfile, + localSecurityBypass: false, + serverNameLookup: nil, + }) multiClient.windows[WindowTypeQuality] = newMultiClientWindow( cancelCtx, @@ -521,7 +542,12 @@ func (self *RemoteUserNatMultiClient) SetPerformanceProfile(performanceProfile * self.stateLock.Lock() defer self.stateLock.Unlock() - self.performanceProfile = performanceProfile + prev := self.config.Load() + self.config.Store(&multiClientConfig{ + performanceProfile: performanceProfile, + localSecurityBypass: prev.localSecurityBypass, + serverNameLookup: prev.serverNameLookup, + }) }() for _, window := range self.windows { window.SetPerformanceProfile(performanceProfile) @@ -533,13 +559,29 @@ func (self *RemoteUserNatMultiClient) SetPerformanceProfile(performanceProfile * func (self *RemoteUserNatMultiClient) SetLocalSecurityBypass(localSecurityBypass bool) { self.stateLock.Lock() defer self.stateLock.Unlock() - self.localSecurityBypass = localSecurityBypass + prev := self.config.Load() + self.config.Store(&multiClientConfig{ + performanceProfile: prev.performanceProfile, + localSecurityBypass: localSecurityBypass, + serverNameLookup: prev.serverNameLookup, + }) } func (self *RemoteUserNatMultiClient) LocalSecurityBypass() bool { + return self.config.Load().localSecurityBypass +} + +// SetServerNameLookup installs (or clears, with nil) the ServerNameLookup used for +// ServerName-based path affinity. Safe to call at runtime. +func (self *RemoteUserNatMultiClient) SetServerNameLookup(serverNameLookup ServerNameLookup) { self.stateLock.Lock() defer self.stateLock.Unlock() - return self.localSecurityBypass + prev := self.config.Load() + self.config.Store(&multiClientConfig{ + performanceProfile: prev.performanceProfile, + localSecurityBypass: prev.localSecurityBypass, + serverNameLookup: serverNameLookup, + }) } // ordered by choice descending @@ -552,13 +594,9 @@ func (self *RemoteUserNatMultiClient) selectWindowTypes(sendPacket *parsedPacket } var fixedWindowType *WindowType - func() { - self.stateLock.Lock() - defer self.stateLock.Unlock() - if self.performanceProfile != nil { - fixedWindowType = &self.performanceProfile.WindowType - } - }() + if pp := self.config.Load().performanceProfile; pp != nil { + fixedWindowType = &pp.WindowType + } if fixedWindowType != nil { return []WindowType{*fixedWindowType} @@ -573,10 +611,11 @@ func (self *RemoteUserNatMultiClient) selectWindowTypes(sendPacket *parsedPacket // called with stateLock func (self *RemoteUserNatMultiClient) affinityIpPathsWithLock(ipPath *IpPath) (affinityPaths []*IpPath, attach bool) { attach = true + config := self.config.Load() singleIp := false - if self.performanceProfile != nil { - singleIp = (self.performanceProfile.WindowSize.FixedWindowSize == 1) + if pp := config.performanceProfile; pp != nil { + singleIp = (pp.WindowSize.FixedWindowSize == 1) } if singleIp { @@ -585,7 +624,7 @@ func (self *RemoteUserNatMultiClient) affinityIpPathsWithLock(ipPath *IpPath) (a } affinityPaths = append(affinityPaths, singlePath) } else if ipPath.DestinationPort == 80 || ipPath.DestinationPort == 53 || ipPath.DestinationPort == 443 { - // for these ports, cycle the path per destination ip, regardless of protocol + // for these ports, cycle the path per destination ip/port, regardless of protocol destinationPath := &IpPath{ Version: ipPath.Version, DestinationIp: ipPath.DestinationIp, @@ -3487,7 +3526,7 @@ func (self *multiClientChannel) windowStatsWithCoalesce(coalesce bool) (*clientW } // `connect.ReceiveFunction` -func (self *multiClientChannel) clientReceive(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { +func (self *multiClientChannel) clientReceive(source TransferPath, frames []*protocol.Frame, peer Peer) { select { case <-self.ctx.Done(): return @@ -3514,7 +3553,7 @@ func (self *multiClientChannel) clientReceive(source TransferPath, frames []*pro if ipPath.Syn { self.addReceiveSyn(1) } - self.clientReceivePacketCallback(self, source, provideMode, ipPath, packet) + self.clientReceivePacketCallback(self, source, peer.ProvideMode, ipPath, packet) } // else not an ip packet, drop } else { diff --git a/ip_test.go b/ip_test.go index 647ca4f..8b089d3 100644 --- a/ip_test.go +++ b/ip_test.go @@ -329,7 +329,7 @@ func testClient[P comparable]( natClient, err := userNatClientGenerator(ctx, providerClient, receivePacketCallback) assert.Equal(t, err, nil) - providerClient.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { + providerClient.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, peer Peer) { // cMutex.Lock() // cReceiveCount += 1 // // fmt.Printf("C Receive %d/%d (%.2f%%)\n", cReceiveCount, totalCount, 100.0 * float32(cReceiveCount) / float32(totalCount)) diff --git a/ip_udp_write_pipeline_leak_test.go b/ip_udp_write_pipeline_leak_test.go index 9aa1a4f..67eadef 100644 --- a/ip_udp_write_pipeline_leak_test.go +++ b/ip_udp_write_pipeline_leak_test.go @@ -6,8 +6,6 @@ import ( "testing" "time" - "github.com/google/gopacket/layers" - "github.com/urnetwork/connect/protocol" ) @@ -63,8 +61,8 @@ func TestUdpSequenceWritePipelineDrainsBuffersOnCancel(t *testing.T) { DestinationId(NewId()), protocol.ProvideMode_Network, 4, - net.ParseIP("127.0.0.1"), layers.UDPPort(40000), - destAddr.IP, layers.UDPPort(destAddr.Port), + net.ParseIP("127.0.0.1"), UDPPort(40000), + destAddr.IP, UDPPort(destAddr.Port), udpBufferSettings, ) @@ -78,8 +76,8 @@ func TestUdpSequenceWritePipelineDrainsBuffersOnCancel(t *testing.T) { // drain them, then cancel immediately - this is the exact window the // fix's drain-on-exit defer covers. for j := 0; j < sendsPerIteration; j += 1 { - udp := &layers.UDP{} - udp.Payload = []byte("x") + udp := &parsedUdp{} + udp.payload = []byte("x") ipPacket := MessagePoolGet(2048) sendItem := &UdpSendItem{ udp: udp, diff --git a/memory_budget.go b/memory_budget.go new file mode 100644 index 0000000..f0663e0 --- /dev/null +++ b/memory_budget.go @@ -0,0 +1,56 @@ +package connect + +import ( + "sync/atomic" +) + +// a process-wide memory budget that scales the default settings whose values +// dominate the per-connection and per-peer memory ceilings (queue caps, +// receive windows, socket buffers, cache bounds). The budget is advisory +// sizing state, separate from the go runtime soft memory limit; hosts set +// both together through the sdk (see sdk.SetMemoryLimit). +// +// A zero budget (the default) leaves every setting at its unscaled default. +// Settings sample the budget when a Default*Settings constructor runs, so the +// budget must be set before constructing the objects it should size. The +// mobile hosts set it at process start, before any device exists. + +// budgets at or above the reference use the unscaled defaults; smaller +// budgets scale the memory-dominant settings proportionally, down to +// per-setting floors +var referenceMemoryBudgetByteCount = mib(64) + +var memoryBudgetByteCount atomic.Int64 + +// SetMemoryBudget sets the process-wide memory budget that scales the +// memory-dominant default settings. 0 (the default) disables scaling. +func SetMemoryBudget(budgetByteCount ByteCount) { + memoryBudgetByteCount.Store(budgetByteCount) +} + +func MemoryBudget() ByteCount { + return memoryBudgetByteCount.Load() +} + +// memoryScale returns the budget scale in (0, 1] +func memoryScale() float64 { + budgetByteCount := memoryBudgetByteCount.Load() + if budgetByteCount <= 0 || referenceMemoryBudgetByteCount <= budgetByteCount { + return 1 + } + return float64(budgetByteCount) / float64(referenceMemoryBudgetByteCount) +} + +// MemoryScaledByteCount scales a default byte count by the memory budget, +// with a floor that preserves a working minimum +func MemoryScaledByteCount(unscaledByteCount ByteCount, floorByteCount ByteCount) ByteCount { + scaledByteCount := ByteCount(memoryScale() * float64(unscaledByteCount)) + return max(floorByteCount, scaledByteCount) +} + +// MemoryScaledCount scales a default count by the memory budget, +// with a floor that preserves a working minimum +func MemoryScaledCount(unscaledCount int, floorCount int) int { + scaledCount := int(memoryScale() * float64(unscaledCount)) + return max(floorCount, scaledCount) +} diff --git a/message_pool.go b/message_pool.go index e113e01..3dd25f7 100644 --- a/message_pool.go +++ b/message_pool.go @@ -58,7 +58,7 @@ const debugTags = false const MessagePoolMetaByteCount = 13 const MessagePoolFlagShared = uint8(0x01) -var InitialMessagePoolByteCount = mib(1) +var InitialMessagePoolByteCount = mib(2) type poolShard struct { mutex sync.Mutex @@ -240,11 +240,24 @@ func poolStats(pools []*messagePool) { } func ResizeMessagePools(maxByteCount ByteCount) { - for _, pool := range orderedMessagePools() { + pools := orderedMessagePools() + poolSizeCount := ByteCount(len(pools)) + for _, pool := range pools { + pool.Resize(int(maxByteCount / poolSizeCount / ByteCount(pool.size))) + } +} + +// ResizeMessagePoolsPerClass applies the given maximum byte count directly as a ceiling +// for each message pool size class, instead of dividing it across the classes. +// This preserves the legacy buffer allocation semantics required by the provider. +func ResizeMessagePoolsPerClass(maxByteCount ByteCount) { + pools := orderedMessagePools() + for _, pool := range pools { pool.Resize(int(maxByteCount / ByteCount(pool.size))) } } + func ClearMessagePools() { for _, pool := range orderedMessagePools() { pool.Clear() diff --git a/protocol/frame.pb.go b/protocol/frame.pb.go index 8ba5d00..7339d36 100644 --- a/protocol/frame.pb.go +++ b/protocol/frame.pb.go @@ -1,7 +1,7 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: // protoc-gen-go v1.36.11 -// protoc v7.35.0 +// protoc v7.35.1 // source: frame.proto package protocol @@ -71,6 +71,15 @@ const ( // pipeline that the platform itself authors. The corresponding // private key is held only by the client process. MessageType_TransferClientKey MessageType = 25 + // control message that marks the start of a complete network peers + // list for the receiving client. `Frame.message_bytes` is a marshaled + // `NetworkPeersReset`. The receiver clears its peer state; the + // complete list follows as `TransferNetworkPeersUpdate` messages. + MessageType_TransferNetworkPeersReset MessageType = 26 + // control message that upserts network peers (and disconnect markers) + // into the receiving client's peer state. `Frame.message_bytes` is a + // marshaled `NetworkPeersUpdate`. + MessageType_TransferNetworkPeersUpdate MessageType = 27 ) // Enum value maps for MessageType. @@ -102,6 +111,8 @@ var ( 23: "TransferEncryptedControl", 24: "TransferEncryptedKey", 25: "TransferClientKey", + 26: "TransferNetworkPeersReset", + 27: "TransferNetworkPeersUpdate", } MessageType_value = map[string]int32{ "TransferPack": 0, @@ -130,6 +141,8 @@ var ( "TransferEncryptedControl": 23, "TransferEncryptedKey": 24, "TransferClientKey": 25, + "TransferNetworkPeersReset": 26, + "TransferNetworkPeersUpdate": 27, } ) @@ -230,7 +243,7 @@ const file_frame_proto_rawDesc = "" + "\x05Frame\x129\n" + "\fmessage_type\x18\x01 \x01(\x0e2\x16.bringyour.MessageTypeR\vmessageType\x12#\n" + "\rmessage_bytes\x18\x02 \x01(\fR\fmessageBytes\x12\x10\n" + - "\x03raw\x18\x03 \x01(\bR\x03raw*\x8f\x05\n" + + "\x03raw\x18\x03 \x01(\bR\x03raw*\xce\x05\n" + "\vMessageType\x12\x10\n" + "\fTransferPack\x10\x00\x12\x0f\n" + "\vTransferAck\x10\x01\x12\x14\n" + @@ -258,7 +271,9 @@ const file_frame_proto_rawDesc = "" + "\x13TransferStreamReset\x10\x16\x12\x1c\n" + "\x18TransferEncryptedControl\x10\x17\x12\x18\n" + "\x14TransferEncryptedKey\x10\x18\x12\x15\n" + - "\x11TransferClientKey\x10\x19B'Z%github.com/urnetwork/connect/protocolb\x06proto3" + "\x11TransferClientKey\x10\x19\x12\x1d\n" + + "\x19TransferNetworkPeersReset\x10\x1a\x12\x1e\n" + + "\x1aTransferNetworkPeersUpdate\x10\x1bB'Z%github.com/urnetwork/connect/protocolb\x06proto3" var ( file_frame_proto_rawDescOnce sync.Once diff --git a/protocol/transfer.pb.go b/protocol/transfer.pb.go index 9cfb7a6..adae1ad 100644 --- a/protocol/transfer.pb.go +++ b/protocol/transfer.pb.go @@ -1,7 +1,7 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: // protoc-gen-go v1.36.11 -// protoc v7.35.0 +// protoc v7.35.1 // source: transfer.proto package protocol @@ -1149,6 +1149,192 @@ func (x *StreamReset) GetStreams() []*StreamOpen { return nil } +// a network peer: a connected top-level client in the same network and its +// identity metadata. When `disconnect_time` is set, this entry is a +// disconnect marker for a recently disconnected peer. +type NetworkPeer struct { + state protoimpl.MessageState `protogen:"open.v1"` + // ulid + ClientId []byte `protobuf:"bytes,1,opt,name=client_id,json=clientId,proto3" json:"client_id,omitempty"` + // the peer's enabled provide modes + ProvideModes []ProvideMode `protobuf:"varint,2,rep,packed,name=provide_modes,json=provideModes,proto3,enum=bringyour.ProvideMode" json:"provide_modes,omitempty"` + // identity principal assigned at client creation + Principal string `protobuf:"bytes,3,opt,name=principal,proto3" json:"principal,omitempty"` + // roles assigned at client creation + Roles []string `protobuf:"bytes,4,rep,name=roles,proto3" json:"roles,omitempty"` + DeviceName string `protobuf:"bytes,5,opt,name=device_name,json=deviceName,proto3" json:"device_name,omitempty"` + DeviceSpec string `protobuf:"bytes,6,opt,name=device_spec,json=deviceSpec,proto3" json:"device_spec,omitempty"` + // unix milliseconds. Set when the peer recently disconnected; + // the receiver ages the marker out locally + DisconnectTime *uint64 `protobuf:"varint,7,opt,name=disconnect_time,json=disconnectTime,proto3,oneof" json:"disconnect_time,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *NetworkPeer) Reset() { + *x = NetworkPeer{} + mi := &file_transfer_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *NetworkPeer) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*NetworkPeer) ProtoMessage() {} + +func (x *NetworkPeer) ProtoReflect() protoreflect.Message { + mi := &file_transfer_proto_msgTypes[13] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use NetworkPeer.ProtoReflect.Descriptor instead. +func (*NetworkPeer) Descriptor() ([]byte, []int) { + return file_transfer_proto_rawDescGZIP(), []int{13} +} + +func (x *NetworkPeer) GetClientId() []byte { + if x != nil { + return x.ClientId + } + return nil +} + +func (x *NetworkPeer) GetProvideModes() []ProvideMode { + if x != nil { + return x.ProvideModes + } + return nil +} + +func (x *NetworkPeer) GetPrincipal() string { + if x != nil { + return x.Principal + } + return "" +} + +func (x *NetworkPeer) GetRoles() []string { + if x != nil { + return x.Roles + } + return nil +} + +func (x *NetworkPeer) GetDeviceName() string { + if x != nil { + return x.DeviceName + } + return "" +} + +func (x *NetworkPeer) GetDeviceSpec() string { + if x != nil { + return x.DeviceSpec + } + return "" +} + +func (x *NetworkPeer) GetDisconnectTime() uint64 { + if x != nil && x.DisconnectTime != nil { + return *x.DisconnectTime + } + return 0 +} + +// control message. Marks the start of a complete network peers list; +// the receiver clears its peer state. The complete list follows as +// `NetworkPeersUpdate` messages. +type NetworkPeersReset struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *NetworkPeersReset) Reset() { + *x = NetworkPeersReset{} + mi := &file_transfer_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *NetworkPeersReset) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*NetworkPeersReset) ProtoMessage() {} + +func (x *NetworkPeersReset) ProtoReflect() protoreflect.Message { + mi := &file_transfer_proto_msgTypes[14] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use NetworkPeersReset.ProtoReflect.Descriptor instead. +func (*NetworkPeersReset) Descriptor() ([]byte, []int) { + return file_transfer_proto_rawDescGZIP(), []int{14} +} + +// control message. Upserts peers (and disconnect markers) into the +// receiver's peer state. +type NetworkPeersUpdate struct { + state protoimpl.MessageState `protogen:"open.v1"` + Peers []*NetworkPeer `protobuf:"bytes,1,rep,name=peers,proto3" json:"peers,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *NetworkPeersUpdate) Reset() { + *x = NetworkPeersUpdate{} + mi := &file_transfer_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *NetworkPeersUpdate) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*NetworkPeersUpdate) ProtoMessage() {} + +func (x *NetworkPeersUpdate) ProtoReflect() protoreflect.Message { + mi := &file_transfer_proto_msgTypes[15] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use NetworkPeersUpdate.ProtoReflect.Descriptor instead. +func (*NetworkPeersUpdate) Descriptor() ([]byte, []int) { + return file_transfer_proto_rawDescGZIP(), []int{15} +} + +func (x *NetworkPeersUpdate) GetPeers() []*NetworkPeer { + if x != nil { + return x.Peers + } + return nil +} + type ExchangeSignals struct { state protoimpl.MessageState `protogen:"open.v1"` // ulid @@ -1161,7 +1347,7 @@ type ExchangeSignals struct { func (x *ExchangeSignals) Reset() { *x = ExchangeSignals{} - mi := &file_transfer_proto_msgTypes[13] + mi := &file_transfer_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1173,7 +1359,7 @@ func (x *ExchangeSignals) String() string { func (*ExchangeSignals) ProtoMessage() {} func (x *ExchangeSignals) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[13] + mi := &file_transfer_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1186,7 +1372,7 @@ func (x *ExchangeSignals) ProtoReflect() protoreflect.Message { // Deprecated: Use ExchangeSignals.ProtoReflect.Descriptor instead. func (*ExchangeSignals) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{13} + return file_transfer_proto_rawDescGZIP(), []int{16} } func (x *ExchangeSignals) GetStreamId() []byte { @@ -1224,7 +1410,7 @@ type ExchangeSignal struct { func (x *ExchangeSignal) Reset() { *x = ExchangeSignal{} - mi := &file_transfer_proto_msgTypes[14] + mi := &file_transfer_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1236,7 +1422,7 @@ func (x *ExchangeSignal) String() string { func (*ExchangeSignal) ProtoMessage() {} func (x *ExchangeSignal) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[14] + mi := &file_transfer_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1249,7 +1435,7 @@ func (x *ExchangeSignal) ProtoReflect() protoreflect.Message { // Deprecated: Use ExchangeSignal.ProtoReflect.Descriptor instead. func (*ExchangeSignal) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{14} + return file_transfer_proto_rawDescGZIP(), []int{17} } func (x *ExchangeSignal) GetSignalType() SignalType { @@ -1297,7 +1483,7 @@ type CreateContract struct { func (x *CreateContract) Reset() { *x = CreateContract{} - mi := &file_transfer_proto_msgTypes[15] + mi := &file_transfer_proto_msgTypes[18] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1309,7 +1495,7 @@ func (x *CreateContract) String() string { func (*CreateContract) ProtoMessage() {} func (x *CreateContract) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[15] + mi := &file_transfer_proto_msgTypes[18] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1322,7 +1508,7 @@ func (x *CreateContract) ProtoReflect() protoreflect.Message { // Deprecated: Use CreateContract.ProtoReflect.Descriptor instead. func (*CreateContract) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{15} + return file_transfer_proto_rawDescGZIP(), []int{18} } func (x *CreateContract) GetDestinationId() []byte { @@ -1388,7 +1574,7 @@ type CreateContractResult struct { func (x *CreateContractResult) Reset() { *x = CreateContractResult{} - mi := &file_transfer_proto_msgTypes[16] + mi := &file_transfer_proto_msgTypes[19] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1400,7 +1586,7 @@ func (x *CreateContractResult) String() string { func (*CreateContractResult) ProtoMessage() {} func (x *CreateContractResult) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[16] + mi := &file_transfer_proto_msgTypes[19] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1413,7 +1599,7 @@ func (x *CreateContractResult) ProtoReflect() protoreflect.Message { // Deprecated: Use CreateContractResult.ProtoReflect.Descriptor instead. func (*CreateContractResult) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{16} + return file_transfer_proto_rawDescGZIP(), []int{19} } func (x *CreateContractResult) GetError() ContractError { @@ -1473,7 +1659,7 @@ type Contract struct { func (x *Contract) Reset() { *x = Contract{} - mi := &file_transfer_proto_msgTypes[17] + mi := &file_transfer_proto_msgTypes[20] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1485,7 +1671,7 @@ func (x *Contract) String() string { func (*Contract) ProtoMessage() {} func (x *Contract) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[17] + mi := &file_transfer_proto_msgTypes[20] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1498,7 +1684,7 @@ func (x *Contract) ProtoReflect() protoreflect.Message { // Deprecated: Use Contract.ProtoReflect.Descriptor instead. func (*Contract) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{17} + return file_transfer_proto_rawDescGZIP(), []int{20} } func (x *Contract) GetStoredContractBytes() []byte { @@ -1578,13 +1764,21 @@ type StoredContract struct { // whether the cert chain in this contract is admitted into the // session's trusted-peer-cert set. DestinationClientKeySignedTlsCertificate []byte `protobuf:"bytes,9,opt,name=destination_client_key_signed_tls_certificate,json=destinationClientKeySignedTlsCertificate,proto3" json:"destination_client_key_signed_tls_certificate,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // roles assigned to the source client at creation. Sealed into the + // platform-signed contract bytes. Set only when `provide_mode` is + // `Network`; empty for all other provide modes. + Roles []string `protobuf:"bytes,10,rep,name=roles,proto3" json:"roles,omitempty"` + // identity principal assigned to the source client at creation. + // Sealed into the platform-signed contract bytes. Set only when + // `provide_mode` is `Network`; empty for all other provide modes. + Principal string `protobuf:"bytes,11,opt,name=principal,proto3" json:"principal,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *StoredContract) Reset() { *x = StoredContract{} - mi := &file_transfer_proto_msgTypes[18] + mi := &file_transfer_proto_msgTypes[21] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1596,7 +1790,7 @@ func (x *StoredContract) String() string { func (*StoredContract) ProtoMessage() {} func (x *StoredContract) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[18] + mi := &file_transfer_proto_msgTypes[21] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1609,7 +1803,7 @@ func (x *StoredContract) ProtoReflect() protoreflect.Message { // Deprecated: Use StoredContract.ProtoReflect.Descriptor instead. func (*StoredContract) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{18} + return file_transfer_proto_rawDescGZIP(), []int{21} } func (x *StoredContract) GetContractId() []byte { @@ -1675,6 +1869,20 @@ func (x *StoredContract) GetDestinationClientKeySignedTlsCertificate() []byte { return nil } +func (x *StoredContract) GetRoles() []string { + if x != nil { + return x.Roles + } + return nil +} + +func (x *StoredContract) GetPrincipal() string { + if x != nil { + return x.Principal + } + return "" +} + // control message type CloseContract struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -1688,7 +1896,7 @@ type CloseContract struct { func (x *CloseContract) Reset() { *x = CloseContract{} - mi := &file_transfer_proto_msgTypes[19] + mi := &file_transfer_proto_msgTypes[22] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1700,7 +1908,7 @@ func (x *CloseContract) String() string { func (*CloseContract) ProtoMessage() {} func (x *CloseContract) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[19] + mi := &file_transfer_proto_msgTypes[22] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1713,7 +1921,7 @@ func (x *CloseContract) ProtoReflect() protoreflect.Message { // Deprecated: Use CloseContract.ProtoReflect.Descriptor instead. func (*CloseContract) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{19} + return file_transfer_proto_rawDescGZIP(), []int{22} } func (x *CloseContract) GetContractId() []byte { @@ -1770,7 +1978,7 @@ type PeerAudit struct { func (x *PeerAudit) Reset() { *x = PeerAudit{} - mi := &file_transfer_proto_msgTypes[20] + mi := &file_transfer_proto_msgTypes[23] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1782,7 +1990,7 @@ func (x *PeerAudit) String() string { func (*PeerAudit) ProtoMessage() {} func (x *PeerAudit) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[20] + mi := &file_transfer_proto_msgTypes[23] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1795,7 +2003,7 @@ func (x *PeerAudit) ProtoReflect() protoreflect.Message { // Deprecated: Use PeerAudit.ProtoReflect.Descriptor instead. func (*PeerAudit) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{20} + return file_transfer_proto_rawDescGZIP(), []int{23} } func (x *PeerAudit) GetPeerId() []byte { @@ -1897,7 +2105,7 @@ type ControlPing struct { func (x *ControlPing) Reset() { *x = ControlPing{} - mi := &file_transfer_proto_msgTypes[21] + mi := &file_transfer_proto_msgTypes[24] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1909,7 +2117,7 @@ func (x *ControlPing) String() string { func (*ControlPing) ProtoMessage() {} func (x *ControlPing) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[21] + mi := &file_transfer_proto_msgTypes[24] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1922,7 +2130,7 @@ func (x *ControlPing) ProtoReflect() protoreflect.Message { // Deprecated: Use ControlPing.ProtoReflect.Descriptor instead. func (*ControlPing) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{21} + return file_transfer_proto_rawDescGZIP(), []int{24} } type ProvidePing struct { @@ -1933,7 +2141,7 @@ type ProvidePing struct { func (x *ProvidePing) Reset() { *x = ProvidePing{} - mi := &file_transfer_proto_msgTypes[22] + mi := &file_transfer_proto_msgTypes[25] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1945,7 +2153,7 @@ func (x *ProvidePing) String() string { func (*ProvidePing) ProtoMessage() {} func (x *ProvidePing) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[22] + mi := &file_transfer_proto_msgTypes[25] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1958,7 +2166,7 @@ func (x *ProvidePing) ProtoReflect() protoreflect.Message { // Deprecated: Use ProvidePing.ProtoReflect.Descriptor instead. func (*ProvidePing) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{22} + return file_transfer_proto_rawDescGZIP(), []int{25} } // EncryptedKey publishes this client's public TLS certificate to the @@ -1988,7 +2196,7 @@ type EncryptedKey struct { func (x *EncryptedKey) Reset() { *x = EncryptedKey{} - mi := &file_transfer_proto_msgTypes[23] + mi := &file_transfer_proto_msgTypes[26] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2000,7 +2208,7 @@ func (x *EncryptedKey) String() string { func (*EncryptedKey) ProtoMessage() {} func (x *EncryptedKey) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[23] + mi := &file_transfer_proto_msgTypes[26] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2013,7 +2221,7 @@ func (x *EncryptedKey) ProtoReflect() protoreflect.Message { // Deprecated: Use EncryptedKey.ProtoReflect.Descriptor instead. func (*EncryptedKey) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{23} + return file_transfer_proto_rawDescGZIP(), []int{26} } func (x *EncryptedKey) GetProvideTlsCertificate() [][]byte { @@ -2049,7 +2257,7 @@ type ClientKey struct { func (x *ClientKey) Reset() { *x = ClientKey{} - mi := &file_transfer_proto_msgTypes[24] + mi := &file_transfer_proto_msgTypes[27] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2061,7 +2269,7 @@ func (x *ClientKey) String() string { func (*ClientKey) ProtoMessage() {} func (x *ClientKey) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[24] + mi := &file_transfer_proto_msgTypes[27] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2074,7 +2282,7 @@ func (x *ClientKey) ProtoReflect() protoreflect.Message { // Deprecated: Use ClientKey.ProtoReflect.Descriptor instead. func (*ClientKey) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{24} + return file_transfer_proto_rawDescGZIP(), []int{27} } func (x *ClientKey) GetPublicKey() []byte { @@ -2124,7 +2332,7 @@ type EncryptedControl struct { func (x *EncryptedControl) Reset() { *x = EncryptedControl{} - mi := &file_transfer_proto_msgTypes[25] + mi := &file_transfer_proto_msgTypes[28] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2136,7 +2344,7 @@ func (x *EncryptedControl) String() string { func (*EncryptedControl) ProtoMessage() {} func (x *EncryptedControl) ProtoReflect() protoreflect.Message { - mi := &file_transfer_proto_msgTypes[25] + mi := &file_transfer_proto_msgTypes[28] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2149,7 +2357,7 @@ func (x *EncryptedControl) ProtoReflect() protoreflect.Message { // Deprecated: Use EncryptedControl.ProtoReflect.Descriptor instead. func (*EncryptedControl) Descriptor() ([]byte, []int) { - return file_transfer_proto_rawDescGZIP(), []int{25} + return file_transfer_proto_rawDescGZIP(), []int{28} } func (x *EncryptedControl) GetControlType() EncryptedControlType { @@ -2264,7 +2472,21 @@ const file_transfer_proto_rawDesc = "" + "\vStreamClose\x12\x1b\n" + "\tstream_id\x18\x03 \x01(\fR\bstreamId\">\n" + "\vStreamReset\x12/\n" + - "\astreams\x18\x01 \x03(\v2\x15.bringyour.StreamOpenR\astreams\"\x88\x01\n" + + "\astreams\x18\x01 \x03(\v2\x15.bringyour.StreamOpenR\astreams\"\x9f\x02\n" + + "\vNetworkPeer\x12\x1b\n" + + "\tclient_id\x18\x01 \x01(\fR\bclientId\x12;\n" + + "\rprovide_modes\x18\x02 \x03(\x0e2\x16.bringyour.ProvideModeR\fprovideModes\x12\x1c\n" + + "\tprincipal\x18\x03 \x01(\tR\tprincipal\x12\x14\n" + + "\x05roles\x18\x04 \x03(\tR\x05roles\x12\x1f\n" + + "\vdevice_name\x18\x05 \x01(\tR\n" + + "deviceName\x12\x1f\n" + + "\vdevice_spec\x18\x06 \x01(\tR\n" + + "deviceSpec\x12,\n" + + "\x0fdisconnect_time\x18\a \x01(\x04H\x00R\x0edisconnectTime\x88\x01\x01B\x12\n" + + "\x10_disconnect_time\"\x13\n" + + "\x11NetworkPeersReset\"B\n" + + "\x12NetworkPeersUpdate\x12,\n" + + "\x05peers\x18\x01 \x03(\v2\x16.bringyour.NetworkPeerR\x05peers\"\x88\x01\n" + "\x0fExchangeSignals\x12\x1b\n" + "\tstream_id\x18\x01 \x01(\fR\bstreamId\x12#\n" + "\rreset_signals\x18\x02 \x01(\bR\fresetSignals\x123\n" + @@ -2300,7 +2522,7 @@ const file_transfer_proto_rawDesc = "" + "\fprovide_mode\x18\x03 \x01(\x0e2\x16.bringyour.ProvideModeR\vprovideMode\x126\n" + "\x17provide_tls_certificate\x18\x04 \x03(\fR\x15provideTlsCertificate\x12A\n" + "\x1ddestination_client_public_key\x18\x05 \x01(\fR\x1adestinationClientPublicKey\x12_\n" + - "-destination_client_key_signed_tls_certificate\x18\x06 \x01(\fR(destinationClientKeySignedTlsCertificate\"\x8a\x04\n" + + "-destination_client_key_signed_tls_certificate\x18\x06 \x01(\fR(destinationClientKeySignedTlsCertificate\"\xbe\x04\n" + "\x0eStoredContract\x12\x1f\n" + "\vcontract_id\x18\x01 \x01(\fR\n" + "contractId\x12.\n" + @@ -2311,7 +2533,10 @@ const file_transfer_proto_rawDesc = "" + "\bpriority\x18\x06 \x01(\rH\x03R\bpriority\x88\x01\x01\x126\n" + "\x17provide_tls_certificate\x18\a \x03(\fR\x15provideTlsCertificate\x12A\n" + "\x1ddestination_client_public_key\x18\b \x01(\fR\x1adestinationClientPublicKey\x12_\n" + - "-destination_client_key_signed_tls_certificate\x18\t \x01(\fR(destinationClientKeySignedTlsCertificateB\f\n" + + "-destination_client_key_signed_tls_certificate\x18\t \x01(\fR(destinationClientKeySignedTlsCertificate\x12\x14\n" + + "\x05roles\x18\n" + + " \x03(\tR\x05roles\x12\x1c\n" + + "\tprincipal\x18\v \x01(\tR\tprincipalB\f\n" + "\n" + "_source_idB\x11\n" + "\x0f_destination_idB\f\n" + @@ -2399,7 +2624,7 @@ func file_transfer_proto_rawDescGZIP() []byte { } var file_transfer_proto_enumTypes = make([]protoimpl.EnumInfo, 5) -var file_transfer_proto_msgTypes = make([]protoimpl.MessageInfo, 26) +var file_transfer_proto_msgTypes = make([]protoimpl.MessageInfo, 29) var file_transfer_proto_goTypes = []any{ (SequenceRole)(0), // 0: bringyour.SequenceRole (ProvideMode)(0), // 1: bringyour.ProvideMode @@ -2419,51 +2644,56 @@ var file_transfer_proto_goTypes = []any{ (*StreamOpen)(nil), // 15: bringyour.StreamOpen (*StreamClose)(nil), // 16: bringyour.StreamClose (*StreamReset)(nil), // 17: bringyour.StreamReset - (*ExchangeSignals)(nil), // 18: bringyour.ExchangeSignals - (*ExchangeSignal)(nil), // 19: bringyour.ExchangeSignal - (*CreateContract)(nil), // 20: bringyour.CreateContract - (*CreateContractResult)(nil), // 21: bringyour.CreateContractResult - (*Contract)(nil), // 22: bringyour.Contract - (*StoredContract)(nil), // 23: bringyour.StoredContract - (*CloseContract)(nil), // 24: bringyour.CloseContract - (*PeerAudit)(nil), // 25: bringyour.PeerAudit - (*ControlPing)(nil), // 26: bringyour.ControlPing - (*ProvidePing)(nil), // 27: bringyour.ProvidePing - (*EncryptedKey)(nil), // 28: bringyour.EncryptedKey - (*ClientKey)(nil), // 29: bringyour.ClientKey - (*EncryptedControl)(nil), // 30: bringyour.EncryptedControl - (*Frame)(nil), // 31: bringyour.Frame - (MessageType)(0), // 32: bringyour.MessageType + (*NetworkPeer)(nil), // 18: bringyour.NetworkPeer + (*NetworkPeersReset)(nil), // 19: bringyour.NetworkPeersReset + (*NetworkPeersUpdate)(nil), // 20: bringyour.NetworkPeersUpdate + (*ExchangeSignals)(nil), // 21: bringyour.ExchangeSignals + (*ExchangeSignal)(nil), // 22: bringyour.ExchangeSignal + (*CreateContract)(nil), // 23: bringyour.CreateContract + (*CreateContractResult)(nil), // 24: bringyour.CreateContractResult + (*Contract)(nil), // 25: bringyour.Contract + (*StoredContract)(nil), // 26: bringyour.StoredContract + (*CloseContract)(nil), // 27: bringyour.CloseContract + (*PeerAudit)(nil), // 28: bringyour.PeerAudit + (*ControlPing)(nil), // 29: bringyour.ControlPing + (*ProvidePing)(nil), // 30: bringyour.ProvidePing + (*EncryptedKey)(nil), // 31: bringyour.EncryptedKey + (*ClientKey)(nil), // 32: bringyour.ClientKey + (*EncryptedControl)(nil), // 33: bringyour.EncryptedControl + (*Frame)(nil), // 34: bringyour.Frame + (MessageType)(0), // 35: bringyour.MessageType } var file_transfer_proto_depIdxs = []int32{ 5, // 0: bringyour.TransferFrame.transfer_path:type_name -> bringyour.TransferPath - 31, // 1: bringyour.TransferFrame.frame:type_name -> bringyour.Frame - 32, // 2: bringyour.TransferFrame.message_type:type_name -> bringyour.MessageType + 34, // 1: bringyour.TransferFrame.frame:type_name -> bringyour.Frame + 35, // 2: bringyour.TransferFrame.message_type:type_name -> bringyour.MessageType 8, // 3: bringyour.TransferFrame.pack:type_name -> bringyour.Pack 10, // 4: bringyour.TransferFrame.ack:type_name -> bringyour.Ack 0, // 5: bringyour.TransferFrame.session_role:type_name -> bringyour.SequenceRole 5, // 6: bringyour.FilteredTransferFrame.transfer_path:type_name -> bringyour.TransferPath - 31, // 7: bringyour.Pack.frames:type_name -> bringyour.Frame - 31, // 8: bringyour.Pack.contract_frame:type_name -> bringyour.Frame + 34, // 7: bringyour.Pack.frames:type_name -> bringyour.Frame + 34, // 8: bringyour.Pack.contract_frame:type_name -> bringyour.Frame 11, // 9: bringyour.Pack.tag:type_name -> bringyour.Tag - 31, // 10: bringyour.FilteredPack.contract_frame:type_name -> bringyour.Frame + 34, // 10: bringyour.FilteredPack.contract_frame:type_name -> bringyour.Frame 11, // 11: bringyour.Ack.tag:type_name -> bringyour.Tag 14, // 12: bringyour.Provide.keys:type_name -> bringyour.ProvideKey 1, // 13: bringyour.ProvideKey.mode:type_name -> bringyour.ProvideMode 15, // 14: bringyour.StreamReset.streams:type_name -> bringyour.StreamOpen - 19, // 15: bringyour.ExchangeSignals.signals:type_name -> bringyour.ExchangeSignal - 2, // 16: bringyour.ExchangeSignal.signal_type:type_name -> bringyour.SignalType - 3, // 17: bringyour.CreateContractResult.error:type_name -> bringyour.ContractError - 22, // 18: bringyour.CreateContractResult.contract:type_name -> bringyour.Contract - 20, // 19: bringyour.CreateContractResult.create_contract:type_name -> bringyour.CreateContract - 1, // 20: bringyour.Contract.provide_mode:type_name -> bringyour.ProvideMode - 4, // 21: bringyour.EncryptedControl.control_type:type_name -> bringyour.EncryptedControlType - 0, // 22: bringyour.EncryptedControl.session_role:type_name -> bringyour.SequenceRole - 23, // [23:23] is the sub-list for method output_type - 23, // [23:23] is the sub-list for method input_type - 23, // [23:23] is the sub-list for extension type_name - 23, // [23:23] is the sub-list for extension extendee - 0, // [0:23] is the sub-list for field type_name + 1, // 15: bringyour.NetworkPeer.provide_modes:type_name -> bringyour.ProvideMode + 18, // 16: bringyour.NetworkPeersUpdate.peers:type_name -> bringyour.NetworkPeer + 22, // 17: bringyour.ExchangeSignals.signals:type_name -> bringyour.ExchangeSignal + 2, // 18: bringyour.ExchangeSignal.signal_type:type_name -> bringyour.SignalType + 3, // 19: bringyour.CreateContractResult.error:type_name -> bringyour.ContractError + 25, // 20: bringyour.CreateContractResult.contract:type_name -> bringyour.Contract + 23, // 21: bringyour.CreateContractResult.create_contract:type_name -> bringyour.CreateContract + 1, // 22: bringyour.Contract.provide_mode:type_name -> bringyour.ProvideMode + 4, // 23: bringyour.EncryptedControl.control_type:type_name -> bringyour.EncryptedControlType + 0, // 24: bringyour.EncryptedControl.session_role:type_name -> bringyour.SequenceRole + 25, // [25:25] is the sub-list for method output_type + 25, // [25:25] is the sub-list for method input_type + 25, // [25:25] is the sub-list for extension type_name + 25, // [25:25] is the sub-list for extension extendee + 0, // [0:25] is the sub-list for field type_name } func init() { file_transfer_proto_init() } @@ -2478,17 +2708,18 @@ func file_transfer_proto_init() { file_transfer_proto_msgTypes[4].OneofWrappers = []any{} file_transfer_proto_msgTypes[5].OneofWrappers = []any{} file_transfer_proto_msgTypes[10].OneofWrappers = []any{} - file_transfer_proto_msgTypes[14].OneofWrappers = []any{} - file_transfer_proto_msgTypes[15].OneofWrappers = []any{} - file_transfer_proto_msgTypes[16].OneofWrappers = []any{} + file_transfer_proto_msgTypes[13].OneofWrappers = []any{} + file_transfer_proto_msgTypes[17].OneofWrappers = []any{} file_transfer_proto_msgTypes[18].OneofWrappers = []any{} + file_transfer_proto_msgTypes[19].OneofWrappers = []any{} + file_transfer_proto_msgTypes[21].OneofWrappers = []any{} type x struct{} out := protoimpl.TypeBuilder{ File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_transfer_proto_rawDesc), len(file_transfer_proto_rawDesc)), NumEnums: 5, - NumMessages: 26, + NumMessages: 29, NumExtensions: 0, NumServices: 0, }, diff --git a/provider/main.go b/provider/main.go index 937d0c7..a9e2fae 100644 --- a/provider/main.go +++ b/provider/main.go @@ -303,7 +303,7 @@ func applyPoolAutoSize(maxMemory connect.ByteCount) { if poolBytes > ceiling { poolBytes = ceiling } - connect.ResizeMessagePools(poolBytes) + connect.ResizeMessagePoolsPerClass(poolBytes) tlog("📦 [pool] message pool %dMiB (RAM=%dMiB)\n", poolBytes/1024/1024, connect.ByteCount(ram)/1024/1024) } @@ -831,7 +831,7 @@ func auth(opts docopt.Opts) { } } if 0 < maxMemory { - connect.ResizeMessagePools(maxMemory / 8) + connect.ResizeMessagePoolsPerClass(maxMemory / 8) debug.SetMemoryLimit(maxMemory) } @@ -1874,7 +1874,7 @@ func provide(opts docopt.Opts) { } } if 0 < maxMemory { - connect.ResizeMessagePools(maxMemory / 8) + connect.ResizeMessagePoolsPerClass(maxMemory / 8) debug.SetMemoryLimit(maxMemory) } applyPoolAutoSize(maxMemory) diff --git a/race_off_test.go b/race_off_test.go new file mode 100644 index 0000000..c40d33e --- /dev/null +++ b/race_off_test.go @@ -0,0 +1,7 @@ +//go:build !race + +package connect + +// raceEnabled reports whether the test binary was built with -race. See the +// -race build variant in race_on_test.go. +const raceEnabled = false diff --git a/race_on_test.go b/race_on_test.go new file mode 100644 index 0000000..f2e96f2 --- /dev/null +++ b/race_on_test.go @@ -0,0 +1,9 @@ +//go:build race + +package connect + +// raceEnabled reports whether the test binary was built with -race. Under -race the +// gvisor stack runs orders of magnitude slower, so throughput/timing assertions that +// are meaningful at full speed (e.g. TestTunTCPThroughput's byte count and floor) are +// scaled down to still exercise the path within the per-connection deadlines. +const raceEnabled = true diff --git a/transfer.go b/transfer.go index 788b2f2..af8b8b7 100644 --- a/transfer.go +++ b/transfer.go @@ -116,7 +116,22 @@ func safeAck(cb func(error), err error) { // provideMode is the mode of where these frames are from: network, friends and family, public // provideMode nil means no contract -type ReceiveFunction = func(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) +type ReceiveFunction = func(source TransferPath, frames []*protocol.Frame, peer Peer) + +type Peer struct { + ProvideMode protocol.ProvideMode +} + +type NetworkPeer struct { + ClientId Id + ProvideModes []protocol.ProvideMode + ProvideEnabled bool + Principal string + Roles []string + DeviceName string + DeviceSpec string +} + type ForwardFunction = func(path TransferPath, transferFrameBytes []byte) func DefaultClientSettings() *ClientSettings { @@ -176,11 +191,12 @@ func DefaultSendBufferSettings() *SendBufferSettings { AckBufferSize: DefaultTransferBufferSize, MinMessageByteCount: ByteCount(1), // this includes transport reconnections - WriteTimeout: 15 * time.Second, - ResendQueueMaxByteCount: mib(4), - ContractFillFraction: 0.7, - ProtocolVersion: DefaultProtocolVersion, - MaxResendCount: 16, + WriteTimeout: 15 * time.Second, + ResendQueueMaxByteCount: MemoryScaledByteCount(mib(4), kib(256)), + ResendQueueMinByteCount: kib(256), + ContractFillFraction: 0.7, + ProtocolVersion: DefaultProtocolVersion, + MaxResendCount: 16, } } @@ -197,11 +213,12 @@ func DefaultReceiveBufferSettings() *ReceiveBufferSettings { // ResendAbuseMultiple: 0.5, MaxPeerAuditDuration: 60 * time.Second, // this includes transport reconnections - WriteTimeout: 15 * time.Second, - ReceiveQueueMaxByteCount: mib(2) + kib(512), - AllowLegacyNack: true, - MaxOpenReceiveContract: 4, - ProtocolVersion: DefaultProtocolVersion, + WriteTimeout: 15 * time.Second, + ReceiveQueueMaxByteCount: MemoryScaledByteCount(mib(2)+kib(512), kib(320)), + ReceiveQueueMinByteCount: kib(320), + AllowLegacyNack: true, + MaxOpenReceiveContract: 4, + ProtocolVersion: DefaultProtocolVersion, } } @@ -865,11 +882,11 @@ func (self *Client) SendMultiHop(frame *protocol.Frame, destination MultiHopId, } // ReceiveFunction -func (self *Client) receive(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { +func (self *Client) receive(source TransferPath, frames []*protocol.Frame, peer Peer) { for _, receiveCallback := range self.receiveCallbacks.Get() { c := func() any { return HandleError(func() { - receiveCallback(source, frames, provideMode) + receiveCallback(source, frames, peer) }) } if self.log.V(2).Enabled() { @@ -993,7 +1010,7 @@ func (self *Client) run() { self.receive( SourceId(self.clientId), []*protocol.Frame{sendPack.Frame}, - protocol.ProvideMode_Network, + Peer{ProvideMode: protocol.ProvideMode_Network}, ) safeAck(sendPack.AckCallback, nil) MessagePoolReturn(sendPack.Frame.MessageBytes) @@ -1431,6 +1448,8 @@ type SendBufferSettings struct { WriteTimeout time.Duration ResendQueueMaxByteCount ByteCount + ResendQueueMinByteCount ByteCount + ResendQueueBudget *TransferMemoryBudget // as this ->1, there is more risk that noack messages will get dropped due to out of sync contracts ContractFillFraction float32 @@ -1846,7 +1865,7 @@ func NewSendSequence( openSendContracts: map[Id]*sequenceContract{}, packs: make(chan *SendPack, sendBufferSettings.SequenceBufferSize), acks: make(chan *protocol.Ack, sendBufferSettings.AckBufferSize), - resendQueue: newResendQueue(), + resendQueue: newResendQueue(sendBufferSettings.ResendQueueBudget, sendBufferSettings.ResendQueueMinByteCount), sendItems: []*sendItem{}, nextSequenceNumber: 0, idleCondition: NewIdleCondition(), @@ -2016,7 +2035,7 @@ func (self *SendSequence) Run() { } // drain the buffer - for _, item := range self.resendQueue.orderedItems { + for _, item := range self.resendQueue.Clear() { safeAck(item.ackCallback, errors.New("Send sequence closed.")) item.messagePoolReturn() } @@ -2201,12 +2220,7 @@ func (self *SendSequence) Run() { // approximate since this cannot consider the next message byte size canQueue := func() bool { - // always allow at least one item in the resend queue - queueSize, queueByteCount := self.resendQueue.QueueSize() - if 0 == queueSize { - return true - } - return queueByteCount < self.sendBufferSettings.ResendQueueMaxByteCount + return self.resendQueue.CanAdd(0, self.sendBufferSettings.ResendQueueMaxByteCount) } if !canQueue() { // wait for acks @@ -2476,6 +2490,13 @@ func (self *SendSequence) setContract(nextSendContract *sequenceContract) { self.openSendContracts[nextSendContract.contractId] = nextSendContract self.sendContract = nextSendContract self.sendContractAcked = false + nextSendContract.contractStatsEntry = self.client.ContractManager().registerContractStats( + nextSendContract.contractId, + false, + self.companionContract, + nextSendContract.path, + nextSendContract.transferByteCount, + ) // The contract carries the destination's `ProvideTlsCertificate` // commitment (possibly empty). Fold the chain into the session's // trusted-peer-cert set so the peer's TLS-handshake cert can be matched @@ -2852,6 +2873,9 @@ func (self *SendSequence) ackItem(item *sendItem) { if item.contractId != nil { if itemSendContract, ok := self.openSendContracts[*item.contractId]; ok { itemSendContract.ack(item.messageByteCount) + if itemSendContract.contractStatsEntry != nil { + itemSendContract.contractStatsEntry.updateUsedByteCount(itemSendContract.ackedByteCount) + } // not current and closed if self.sendContract != itemSendContract && itemSendContract.unackedByteCount == 0 { self.client.ContractManager().CloseContract( @@ -3101,8 +3125,8 @@ func (self *sendItem) messagePoolReturn() { // - ack timeouts type resendQueue = transferQueue[*sendItem] -func newResendQueue() *resendQueue { - return newTransferQueue[*sendItem](func(a *sendItem, b *sendItem) int { +func newResendQueue(budget *TransferMemoryBudget, minByteCount ByteCount) *resendQueue { + q := newTransferQueue[*sendItem](func(a *sendItem, b *sendItem) int { if a.resendTime.Before(b.resendTime) { return -1 } else if b.resendTime.Before(a.resendTime) { @@ -3111,6 +3135,8 @@ func newResendQueue() *resendQueue { return 0 } }) + q.setBudget(budget, minByteCount) + return q } type ReceiveBufferSettings struct { @@ -3134,6 +3160,8 @@ type ReceiveBufferSettings struct { WriteTimeout time.Duration ReceiveQueueMaxByteCount ByteCount + ReceiveQueueMinByteCount ByteCount + ReceiveQueueBudget *TransferMemoryBudget // whether to allow nacks without a contract_id AllowLegacyNack bool @@ -3403,6 +3431,7 @@ type ReceiveSequence struct { // it never initiates one. Holding it keeps the session (and its cipher) // alive for the stream's lifetime; released when the sequence terminates. // Nil when encryption is disabled or this is control-plane traffic. + session *peerEncryptionSession } @@ -3428,7 +3457,7 @@ func NewReceiveSequence( openReceiveContracts: map[Id]*sequenceContract{}, receiveContract: nil, packs: make(chan *ReceivePack, receiveBufferSettings.SequenceBufferSize), - receiveQueue: newReceiveQueue(), + receiveQueue: newReceiveQueue(receiveBufferSettings.ReceiveQueueBudget, receiveBufferSettings.ReceiveQueueMinByteCount), nextSequenceNumber: 0, idleCondition: NewIdleCondition(), ackWindow: newSequenceAckWindow(), @@ -3546,7 +3575,7 @@ func (self *ReceiveSequence) Run() { } // drain the buffer - for _, item := range self.receiveQueue.orderedItems { + for _, item := range self.receiveQueue.Clear() { self.peerAudit.Update(func(a *PeerAudit) { a.discard(item.messageByteCount) }) @@ -3951,12 +3980,7 @@ func (self *ReceiveSequence) receive(receivePack *ReceivePack) (bool, error) { } else { // store only up to a max size in the receive queue canQueue := func(byteCount ByteCount) bool { - // always allow at least one item in the receive queue - queueSize, queueByteCount := self.receiveQueue.QueueSize() - if 0 == queueSize { - return true - } - return queueByteCount+byteCount < self.receiveBufferSettings.ReceiveQueueMaxByteCount + return self.receiveQueue.CanAdd(byteCount, self.receiveBufferSettings.ReceiveQueueMaxByteCount) } // remove later items to fit @@ -4066,9 +4090,13 @@ func (self *ReceiveSequence) receiveHead(item *receiveItem) { var provideMode protocol.ProvideMode if item.contractId != nil { - receiveContract := self.openReceiveContracts[*item.contractId] - receiveContract.ack(item.messageByteCount) - provideMode = receiveContract.provideMode + if receiveContract, ok := self.openReceiveContracts[*item.contractId]; ok { + receiveContract.ack(item.messageByteCount) + self.updateContractStats(receiveContract) + provideMode = receiveContract.provideMode + } else { + provideMode = protocol.ProvideMode_Network + } } else { // no contract peers are considered in network provideMode = protocol.ProvideMode_Network @@ -4085,7 +4113,7 @@ func (self *ReceiveSequence) receiveHead(item *receiveItem) { item.receiveCallback( self.source, appFrames, - provideMode, + Peer{ProvideMode: provideMode}, ) } if item.ack { @@ -4201,6 +4229,13 @@ func (self *ReceiveSequence) setContract(nextReceiveContract *sequenceContract) self.openReceiveContracts[nextReceiveContract.contractId] = nextReceiveContract self.receiveContract = nextReceiveContract + nextReceiveContract.contractStatsEntry = self.client.ContractManager().registerContractStats( + nextReceiveContract.contractId, + true, + false, + nextReceiveContract.path, + nextReceiveContract.transferByteCount, + ) if d := len(self.openReceiveContracts) - self.receiveBufferSettings.MaxOpenReceiveContract; 0 < d { // remove the least recently added @@ -4243,6 +4278,12 @@ func (self *ReceiveSequence) updateContract(item *receiveItem) bool { return false } +func (self *ReceiveSequence) updateContractStats(receiveContract *sequenceContract) { + if receiveContract.contractStatsEntry != nil { + receiveContract.contractStatsEntry.updateUsedByteCount(receiveContract.ackedByteCount) + } +} + func (self *ReceiveSequence) Close() { self.cancel() @@ -4317,8 +4358,8 @@ func (self *receiveItem) messagePoolReturn() { // ordered by sequenceNumber type receiveQueue = transferQueue[*receiveItem] -func newReceiveQueue() *receiveQueue { - return newTransferQueue[*receiveItem](func(a *receiveItem, b *receiveItem) int { +func newReceiveQueue(budget *TransferMemoryBudget, minByteCount ByteCount) *receiveQueue { + q := newTransferQueue[*receiveItem](func(a *receiveItem, b *receiveItem) int { if a.sequenceNumber < b.sequenceNumber { return -1 } else if b.sequenceNumber < a.sequenceNumber { @@ -4327,6 +4368,8 @@ func newReceiveQueue() *receiveQueue { return 0 } }) + q.setBudget(budget, minByteCount) + return q } type sequenceAck struct { @@ -4471,6 +4514,13 @@ type sequenceContract struct { path TransferPath + // per-contract stats entry. Stored here (not on the owning sequence) + // because a sequence can have multiple open contracts in flight during + // contract rollover; keeping the entry on the sequence would misattribute + // acks/updates for a non-current contract to whichever contract is + // currently active. + contractStatsEntry *contractStatsEntry + ackedByteCount ByteCount unackedByteCount ByteCount diff --git a/transfer_contract_manager.go b/transfer_contract_manager.go index 971be4a..bbcb3eb 100644 --- a/transfer_contract_manager.go +++ b/transfer_contract_manager.go @@ -156,6 +156,25 @@ func (self *ContractManagerStats) ContractOpenByteCount() ByteCount { return netContractOpenByteCount } +func SignStoredContract(settings *ContractManagerSettings, provideSecretKey []byte, storedContractBytes []byte) []byte { + mac := hmac.New(sha256.New, provideSecretKey) + if time.Now().Before(settings.NetworkEventTimeChangeHmac) { + return mac.Sum(storedContractBytes) + } + mac.Write(storedContractBytes) + return mac.Sum(nil) +} + +func VerifyStoredContract(settings *ContractManagerSettings, provideSecretKey []byte, storedContractBytes []byte, storedContractHmac []byte) bool { + legacyMac := hmac.New(sha256.New, provideSecretKey) + if hmac.Equal(storedContractHmac, legacyMac.Sum(storedContractBytes)) { + return true + } + standardMac := hmac.New(sha256.New, provideSecretKey) + standardMac.Write(storedContractBytes) + return hmac.Equal(storedContractHmac, standardMac.Sum(nil)) +} + func DefaultContractManagerSettings() *ContractManagerSettings { return DefaultContractManagerSettingsWithBufferSize(defaultTransferBufferSize) } @@ -167,6 +186,10 @@ func DefaultContractManagerSettingsWithBufferSize(bufferSize int) *ContractManag if err != nil { panic(err) } + networkEventTimeChangeHmac, err := time.Parse(time.RFC3339, "2026-09-01T00:00:00Z") + if err != nil { + panic(err) + } return &ContractManagerSettings{ SequenceBufferSize: bufferSize, InitialContractTransferByteCount: mib(2), @@ -174,11 +197,14 @@ func DefaultContractManagerSettingsWithBufferSize(bufferSize int) *ContractManag ContractTransferByteSeqScale: 3, NetworkEventTimeEnableContracts: networkEventTimeEnableContracts, + NetworkEventTimeChangeHmac: networkEventTimeChangeHmac, ProvidePingTimeout: 0, OriginContractLinger: 300 * time.Second, + ContractStatsEpoch: 60 * time.Second, + ContractQueueExpireTimeout: 120 * time.Second, ProtocolVersion: DefaultProtocolVersion, @@ -188,6 +214,7 @@ func DefaultContractManagerSettingsWithBufferSize(bufferSize int) *ContractManag func DefaultContractManagerSettingsNoNetworkEvents() *ContractManagerSettings { settings := DefaultContractManagerSettings() settings.NetworkEventTimeEnableContracts = time.Time{} + settings.NetworkEventTimeChangeHmac = time.Time{} return settings } @@ -204,6 +231,8 @@ type ContractManagerSettings struct { // this can be removed after wide adoption NetworkEventTimeEnableContracts time.Time + NetworkEventTimeChangeHmac time.Time + // an active ping to the control fast-tracks any timeouts ProvidePingTimeout time.Duration @@ -218,6 +247,8 @@ type ContractManagerSettings struct { ContractQueueExpireTimeout time.Duration ProtocolVersion int + + ContractStatsEpoch time.Duration } func (self *ContractManagerSettings) ContractsEnabled() bool { @@ -250,6 +281,11 @@ type ContractManager struct { localStats *ContractManagerStats + contractStatsLock sync.Mutex + contractStatsEntries map[contractStatsKey]*contractStatsEntry + contractStatsCallbacks *CallbackList[ContractStatsFunction] + contractStatsStarted bool + controlSyncProvide *ControlSync controlSyncProvideOob *ControlSyncOob } @@ -263,6 +299,9 @@ func NewContractManager( client *Client, settings *ContractManagerSettings, ) *ContractManager { + if settings.ContractStatsEpoch <= 0 { + settings.ContractStatsEpoch = 10 * time.Second + } // at a minimum // - messages to/from the platform (ControlId) do not need a contract // this is because the platform is needed to create contracts @@ -289,6 +328,8 @@ func NewContractManager( sendNoContractClientIds: sendNoContractClientIds, contractStatusCallbacks: NewCallbackList[*contractStatusCallbackWorker](), localStats: NewContractManagerStats(), + contractStatsEntries: map[contractStatsKey]*contractStatsEntry{}, + contractStatsCallbacks: NewCallbackList[ContractStatsFunction](), controlSyncProvide: NewControlSync(ctx, client, "provide"), controlSyncProvideOob: NewControlSyncOob(ctx, client, "provide-oob"), } @@ -681,10 +722,12 @@ func (self *ContractManager) provideFrame() (*protocol.Frame, error) { var provide *protocol.Provide if self.providePaused { - // keep only ProvideMode_Stream to allow return traffic, if set + // pause stops providing to public/ff only. + // keep ProvideMode_Stream to allow return traffic and + // ProvideMode_Network so network peers never fall back to stream, if set provideKeys := []*protocol.ProvideKey{} for provideMode, allow := range self.provideModes { - if allow && provideMode == protocol.ProvideMode_Stream { + if allow && (provideMode == protocol.ProvideMode_Stream || provideMode == protocol.ProvideMode_Network) { provideSecretKey, ok := self.provideSecretKeys[provideMode] if ok { provideKeys = append(provideKeys, &protocol.ProvideKey{ @@ -825,8 +868,9 @@ func (self *ContractManager) Verify(storedContractHmac []byte, storedContractByt self.mutex.Lock() defer self.mutex.Unlock() - // only allow ProvideMode_Stream when paused, for return traffic - if self.providePaused && provideMode != protocol.ProvideMode_Stream { + // when paused, only allow ProvideMode_Stream for return traffic and + // ProvideMode_Network for network peers (pause stops public/ff only) + if self.providePaused && provideMode != protocol.ProvideMode_Stream && provideMode != protocol.ProvideMode_Network { return false } @@ -1135,6 +1179,16 @@ func (self *ContractManager) CloseContractWithCheckpoint( } } + // Unlike the localStats.ContractOpenByteCounts/ContractOpenKeys bookkeeping + // above (which legitimately stays untouched on checkpoint — the wire-level + // contract may still be reused by a new sequence), the per-contract stats + // entry's lifecycle is tied to this sequence's Run() ending, not to the + // wire contract's lifecycle. It must always be released here, checkpoint + // or not, or it leaks for every contract whose sequence ends via + // checkpoint instead of a hard close — which per ReceiveSequence.Run's + // defer is the common case, not the exception. + self.closeContractStats(contractId) + // Reliable delivery via a per-contract `ControlSync`. The // previous implementation called `ClientOob().SendControl(...)` // once and dropped the result on the floor — a single transient diff --git a/transfer_contract_stats.go b/transfer_contract_stats.go new file mode 100644 index 0000000..c7e9b6a --- /dev/null +++ b/transfer_contract_stats.go @@ -0,0 +1,164 @@ +package connect + +import ( + "sync/atomic" + "time" +) + +// contract stats surface per-contract usage from the send and receive sequences +// without per-packet dispatch. open contracts register an entry whose used byte +// count the owning sequence updates with an atomic store on debit/ack. an epoch +// worker (started on the first callback) snapshots the registry every +// `ContractStatsEpoch` and emits events only for contracts that changed. + +type ContractStatsEvent struct { + ContractId Id + // true for a receive (ingress) contract, false for a send (egress) contract + Receive bool + // true when the send sequence opened this as a companion contract. + // always false on the receive side (the wire contract does not carry it). + // pair contracts to their companions with the peer client id on `Path` + Companion bool + // source -> destination of the contract + Path TransferPath + TransferByteCount ByteCount + UsedByteCount ByteCount + // used bytes since the previous event for this contract + UsedByteCountDelta ByteCount + // false when the contract closed. the final event for a contract + Open bool +} + +type ContractStatsFunction func(contractStatsEvents []*ContractStatsEvent) + +type contractStatsKey struct { + contractId Id + receive bool +} + +type contractStatsEntry struct { + contractId Id + receive bool + companion bool + path TransferPath + transferByteCount ByteCount + + // stored by the owning sequence on debit/ack + usedByteCount atomic.Int64 + closed atomic.Bool + + // owned by the epoch worker + emittedUsedByteCount ByteCount + emitted bool +} + +func (self *contractStatsEntry) updateUsedByteCount(usedByteCount ByteCount) { + self.usedByteCount.Store(int64(usedByteCount)) +} + +// register adds an entry for an open contract and returns it. +// the owning sequence stores the used byte count on the entry +func (self *ContractManager) registerContractStats( + contractId Id, + receive bool, + companion bool, + path TransferPath, + transferByteCount ByteCount, +) *contractStatsEntry { + entry := &contractStatsEntry{ + contractId: contractId, + receive: receive, + companion: companion, + path: path, + transferByteCount: transferByteCount, + } + + self.contractStatsLock.Lock() + defer self.contractStatsLock.Unlock() + self.contractStatsEntries[contractStatsKey{contractId: contractId, receive: receive}] = entry + return entry +} + +// called from the contract close paths (`CloseContract`, `CheckpointContract`). +// the entry emits a final closed event and is removed by the epoch worker, +// or removed immediately when the worker never started +func (self *ContractManager) closeContractStats(contractId Id) { + self.contractStatsLock.Lock() + defer self.contractStatsLock.Unlock() + for _, receive := range []bool{false, true} { + key := contractStatsKey{contractId: contractId, receive: receive} + if entry, ok := self.contractStatsEntries[key]; ok { + if self.contractStatsStarted { + entry.closed.Store(true) + } else { + delete(self.contractStatsEntries, key) + } + } + } +} + +// AddContractStatsCallback registers a listener for the epoch contract stats +// events. the epoch worker starts on the first callback +func (self *ContractManager) AddContractStatsCallback(contractStatsCallback ContractStatsFunction) func() { + func() { + self.contractStatsLock.Lock() + defer self.contractStatsLock.Unlock() + if !self.contractStatsStarted { + self.contractStatsStarted = true + go HandleError(self.runContractStats, self.client.Cancel) + } + }() + callbackId := self.contractStatsCallbacks.Add(contractStatsCallback) + return func() { + self.contractStatsCallbacks.Remove(callbackId) + } +} + +func (self *ContractManager) runContractStats() { + for { + select { + case <-self.ctx.Done(): + return + case <-time.After(self.settings.ContractStatsEpoch): + } + + callbacks := self.contractStatsCallbacks.Get() + + var events []*ContractStatsEvent + func() { + self.contractStatsLock.Lock() + defer self.contractStatsLock.Unlock() + for key, entry := range self.contractStatsEntries { + usedByteCount := ByteCount(entry.usedByteCount.Load()) + closed := entry.closed.Load() + if !entry.emitted || usedByteCount != entry.emittedUsedByteCount || closed { + if 0 < len(callbacks) { + events = append(events, &ContractStatsEvent{ + ContractId: entry.contractId, + Receive: entry.receive, + Companion: entry.companion, + Path: entry.path, + TransferByteCount: entry.transferByteCount, + UsedByteCount: usedByteCount, + UsedByteCountDelta: usedByteCount - entry.emittedUsedByteCount, + Open: !closed, + }) + } + entry.emitted = true + entry.emittedUsedByteCount = usedByteCount + } + if closed { + delete(self.contractStatsEntries, key) + } + } + }() + + if 0 < len(events) { + for _, callback := range callbacks { + HandleError(func() { + callback(events) + }) + } + } + } +} diff --git a/transfer_contract_stats_test.go b/transfer_contract_stats_test.go new file mode 100644 index 0000000..06bfc1a --- /dev/null +++ b/transfer_contract_stats_test.go @@ -0,0 +1,301 @@ +package connect + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/go-playground/assert/v2" + + "github.com/urnetwork/connect/protocol" +) + +// Regression tests for the three contract-stats bugs found in the 2026-07-13 +// audit of PR #264 (transfer_contract_manager.go / transfer.go / +// transfer_contract_stats.go), verifying the fixes landed in commit 7440e8b: +// - checkpoint-path leak of `contractStatsEntries` (closeContractStats was +// skipped whenever a contract was checkpointed rather than closed) +// - receive-side stats updated one item early, from `updateContract`'s +// reservation step instead of `receiveHead`'s ack step, permanently +// losing the final processed item's bytes on close +// - `runContractStats` hot-spinning at 100% CPU when `ContractStatsEpoch` +// is left at its zero value + +// TestContractStatsEpochDefaultsWhenInvalid pins finding 3.3: a +// ContractManagerSettings with an unset/invalid ContractStatsEpoch must not +// leave `runContractStats` ticking on `time.After(0)`. +func TestContractStatsEpochDefaultsWhenInvalid(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + client := NewClient(ctx, NewId(), NewNoContractClientOob(), DefaultClientSettings()) + defer client.Cancel() + + settings := DefaultContractManagerSettings() + settings.ContractStatsEpoch = 0 + contractManager := NewContractManager(ctx, client, settings) + + assert.Equal(t, true, 0 < contractManager.settings.ContractStatsEpoch) + + // Confirm the worker actually respects a sane interval rather than + // spinning: with the epoch normalized, a callback should not fire faster + // than the (small but nonzero) floor the constructor applies. + var fireCount int + var mu sync.Mutex + remove := contractManager.AddContractStatsCallback(func(events []*ContractStatsEvent) { + mu.Lock() + fireCount += 1 + mu.Unlock() + }) + defer remove() + + time.Sleep(200 * time.Millisecond) + + mu.Lock() + defer mu.Unlock() + // A hot-spin (time.After(0)) would fire thousands of times in 200ms; + // a normalized multi-second-scale epoch fires at most once or twice. + assert.Equal(t, true, fireCount < 50) +} + +// TestContractStatsClosedOnCheckpoint pins finding 3.1: checkpointing a +// contract (the path taken when a sequence terminates but the wire contract +// stays open server-side) must still release the client-side stats entry, +// not just a hard CloseContract. +func TestContractStatsClosedOnCheckpoint(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + clientSettings := DefaultClientSettings() + clientSettings.ContractManagerSettings.ContractStatsEpoch = 20 * time.Millisecond + client := NewClient(ctx, NewId(), NewNoContractClientOob(), clientSettings) + defer client.Cancel() + contractManager := client.ContractManager() + + contractId := NewId() + path := TransferPath{DestinationId: NewId()} + + t.Run("worker never started: checkpoint deletes the entry immediately", func(t *testing.T) { + entry := contractManager.registerContractStats(contractId, true, false, path, Mib(2)) + assert.Equal(t, false, entry.closed.Load()) + + contractManager.CheckpointContract(contractId, Kib(10), 0) + + contractManager.contractStatsLock.Lock() + _, stillPresent := contractManager.contractStatsEntries[contractStatsKey{contractId: contractId, receive: true}] + contractManager.contractStatsLock.Unlock() + + // Before the fix, `CloseContractWithCheckpoint` skipped + // `closeContractStats` entirely on the checkpoint path, so this entry + // would remain in the map forever (the leak). + assert.Equal(t, false, stillPresent) + }) + + t.Run("worker started: checkpoint marks the entry closed for the epoch worker to reap", func(t *testing.T) { + contractId2 := NewId() + entry := contractManager.registerContractStats(contractId2, true, false, path, Mib(2)) + + var events []*ContractStatsEvent + var mu sync.Mutex + remove := contractManager.AddContractStatsCallback(func(e []*ContractStatsEvent) { + mu.Lock() + events = append(events, e...) + mu.Unlock() + }) + defer remove() + + entry.updateUsedByteCount(Kib(5)) + contractManager.CheckpointContract(contractId2, Kib(5), 0) + + deadline := time.After(5 * time.Second) + for { + mu.Lock() + found := false + for _, e := range events { + if e.ContractId == contractId2 && !e.Open { + found = true + } + } + n := len(events) + mu.Unlock() + if found { + break + } + select { + case <-deadline: + t.Fatalf("no closed event observed for checkpointed contract after 5s (saw %d events)", n) + case <-time.After(10 * time.Millisecond): + } + } + + contractManager.contractStatsLock.Lock() + _, stillPresent := contractManager.contractStatsEntries[contractStatsKey{contractId: contractId2, receive: true}] + contractManager.contractStatsLock.Unlock() + assert.Equal(t, false, stillPresent) + }) +} + +// TestContractStatsReceiveSideNotLagging pins finding 3.2 end-to-end: the +// final contract-stats event for a receive-side contract must account for +// every acked byte, including the last message processed before the sequence +// closes. Before the fix, `updateContractStats` was called from +// `updateContract` (the reservation step, before `ackedByteCount` is +// incremented) instead of `receiveHead` (the ack step), so the reported +// UsedByteCount always lagged by one message and the last message's bytes +// were lost entirely once the sequence closed. +func TestContractStatsReceiveSideNotLagging(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + aClientId := NewId() + bClientId := NewId() + + aSend := make(chan []byte) + bSend := make(chan []byte) + + _, bReceive := newConditioner(ctx, aSend) + _, aReceive := newConditioner(ctx, bSend) + + aSendTransport := NewSendGatewayTransport() + aReceiveTransport := NewReceiveGatewayTransport() + bSendTransport := NewSendGatewayTransport() + bReceiveTransport := NewReceiveGatewayTransport() + + provideModes := map[protocol.ProvideMode]bool{protocol.ProvideMode_Network: true} + + clientSettingsA := DefaultClientSettings() + clientSettingsA.SendBufferSettings.SequenceBufferSize = 0 + clientSettingsA.SendBufferSettings.AckBufferSize = 0 + clientSettingsA.SendBufferSettings.AckTimeout = 60 * time.Second + clientSettingsA.SendBufferSettings.IdleTimeout = 60 * time.Second + clientSettingsA.ReceiveBufferSettings.SequenceBufferSize = 0 + clientSettingsA.ReceiveBufferSettings.GapTimeout = 60 * time.Second + clientSettingsA.ReceiveBufferSettings.IdleTimeout = 60 * time.Second + clientSettingsA.ForwardBufferSettings.SequenceBufferSize = 0 + clientSettingsA.ForwardBufferSettings.IdleTimeout = 60 * time.Second + a := NewClient(ctx, aClientId, NewNoContractClientOob(), clientSettingsA) + defer a.Cancel() + a.RouteManager().UpdateTransport(aSendTransport, []Route{aSend}) + a.RouteManager().UpdateTransport(aReceiveTransport, []Route{aReceive}) + a.ContractManager().SetProvideModes(provideModes) + + clientSettingsB := DefaultClientSettings() + clientSettingsB.SendBufferSettings.SequenceBufferSize = 0 + clientSettingsB.SendBufferSettings.AckBufferSize = 0 + clientSettingsB.SendBufferSettings.AckTimeout = 60 * time.Second + clientSettingsB.SendBufferSettings.IdleTimeout = 60 * time.Second + clientSettingsB.ReceiveBufferSettings.SequenceBufferSize = 0 + clientSettingsB.ReceiveBufferSettings.GapTimeout = 60 * time.Second + clientSettingsB.ReceiveBufferSettings.IdleTimeout = 60 * time.Second + clientSettingsB.ForwardBufferSettings.SequenceBufferSize = 0 + clientSettingsB.ForwardBufferSettings.IdleTimeout = 60 * time.Second + // Short epoch so the closing event flushes quickly once the sequence + // checkpoints. + clientSettingsB.ContractManagerSettings.ContractStatsEpoch = 20 * time.Millisecond + b := NewClient(ctx, bClientId, NewNoContractClientOob(), clientSettingsB) + defer b.Cancel() + b.RouteManager().UpdateTransport(bSendTransport, []Route{bSend}) + b.RouteManager().UpdateTransport(bReceiveTransport, []Route{bReceive}) + b.ContractManager().SetProvideModes(provideModes) + + var eventsMu sync.Mutex + var lastClosedUsedByteCount ByteCount + var sawClosed bool + removeCallback := b.ContractManager().AddContractStatsCallback(func(events []*ContractStatsEvent) { + eventsMu.Lock() + defer eventsMu.Unlock() + for _, e := range events { + if e.Receive && !e.Open { + sawClosed = true + lastClosedUsedByteCount = e.UsedByteCount + } + } + }) + defer removeCallback() + + // Grant a single receive-side contract, generously sized so all messages + // fit in it (no rollover — this test isolates the lagging-stats bug, not + // the cross-contract rollover case covered by 0823eca). + err := a.ContractManager().HandleControlFrame( + ContractKey{Destination: DestinationId(bClientId)}, + requireContractResult( + protocol.ProvideMode_Network, + b.ContractManager().RequireProvideSecretKey(protocol.ProvideMode_Network), + aClientId, + bClientId, + ), + ) + assert.Equal(t, err, nil) + + n := 20 + acks := make(chan error, n) + receives := make(chan *protocol.SimpleMessage, n) + b.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, peer Peer) { + for _, frame := range frames { + m, err := FromFrame(frame) + if err != nil { + continue + } + if sm, ok := m.(*protocol.SimpleMessage); ok { + receives <- sm + } + } + }) + + for i := 0; i < n; i += 1 { + message := &protocol.SimpleMessage{Content: fmt.Sprintf("msg-%d", i)} + frame, err := ToFrame(message, DefaultProtocolVersion) + if err != nil { + t.Fatal(err) + } + success := a.Send(frame, DestinationId(bClientId), func(err error) { + acks <- err + }) + assert.Equal(t, success, true) + } + + ackCount, receiveCount := 0, 0 + deadline := time.After(30 * time.Second) + for ackCount < n || receiveCount < n { + select { + case err := <-acks: + assert.Equal(t, err, nil) + ackCount += 1 + case <-receives: + receiveCount += 1 + case <-deadline: + t.Fatalf("timed out waiting for all messages (acked=%d received=%d of %d)", ackCount, receiveCount, n) + } + } + + // All n messages are acked and received; ground truth for "bytes actually + // processed by the receive sequence" is `b`'s own contract accounting, + // captured by closing the receive sequence (checkpoint) and reading the + // accumulated ReceiveContractCloseByteCount off LocalStats. This is + // exactly the `receiveContract.ackedByteCount` value at close time, + // independent of the stats-event machinery under test. + b.Cancel() + + deadline2 := time.After(5 * time.Second) + for { + groundTruth := b.ContractManager().LocalStats().ReceiveContractCloseByteCount + eventsMu.Lock() + closed, used := sawClosed, lastClosedUsedByteCount + eventsMu.Unlock() + if closed && 0 < groundTruth { + // Before the fix this would fail: the closing stats event's + // UsedByteCount was captured one ack behind, so it never + // included the final message's bytes and would be strictly + // less than groundTruth. + assert.Equal(t, groundTruth, used) + return + } + select { + case <-deadline2: + t.Fatalf("no closed contract-stats event observed after close (sawClosed=%t)", closed) + case <-time.After(10 * time.Millisecond): + } + } +} diff --git a/transfer_control_test.go b/transfer_control_test.go index ef2e91d..b79e2d5 100644 --- a/transfer_control_test.go +++ b/transfer_control_test.go @@ -44,7 +44,7 @@ func TestControlSync(t *testing.T) { receive := make(chan *protocol.SimpleMessage) // on receive a test message, set the channel - controlClientA.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { + controlClientA.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, peer Peer) { for _, frame := range frames { m, err := FromFrame(frame) assert.Equal(t, err, nil) diff --git a/transfer_encrypt_contract_test.go b/transfer_encrypt_contract_test.go index d7015dd..b934f71 100644 --- a/transfer_encrypt_contract_test.go +++ b/transfer_encrypt_contract_test.go @@ -254,7 +254,7 @@ func TestSendReceiveEncryptedWithContracts(t *testing.T) { receivesA := make(chan string, 1024) receivesB := make(chan string, 1024) - a.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, _ protocol.ProvideMode) { + a.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, _ Peer) { for _, frame := range frames { if m, err := FromFrame(frame); err == nil { if sm, ok := m.(*protocol.SimpleMessage); ok { @@ -263,7 +263,7 @@ func TestSendReceiveEncryptedWithContracts(t *testing.T) { } } }) - b.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, _ protocol.ProvideMode) { + b.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, _ Peer) { for _, frame := range frames { if m, err := FromFrame(frame); err == nil { if sm, ok := m.(*protocol.SimpleMessage); ok { diff --git a/transfer_memory_budget.go b/transfer_memory_budget.go new file mode 100644 index 0000000..b2cd750 --- /dev/null +++ b/transfer_memory_budget.go @@ -0,0 +1,74 @@ +package connect + +import ( + "sync/atomic" +) + +// A byte budget shared by transfer queues across sequences — typically all +// clients of one device (the control client plus every window client) — so +// the aggregate queue memory stays flat as the number of peers grows, while +// a single fast peer can still borrow a deep queue. +// +// Queues reserve the bytes they hold above their guaranteed floor +// (`ResendQueueMinByteCount`/`ReceiveQueueMinByteCount`) and release them as +// items leave (see `transferQueue`). Admission gates on `Available` before a +// queue grows above its floor, so `Reserve` may transiently overdraft the +// total by up to one message per sequence past an admission that saw +// headroom. The floor keeps every sequence progressing when the pool is +// empty, which makes cross-sequence deadlock impossible. +// +// All methods are safe for concurrent use. +type TransferMemoryBudget struct { + totalByteCount ByteCount + usedByteCount atomic.Int64 + // cumulative counters, so tests can assert reserve/release balance after + // a build/load/teardown cycle (the message pool counts pattern) + reservedByteCount atomic.Int64 + releasedByteCount atomic.Int64 +} + +func NewTransferMemoryBudget(totalByteCount ByteCount) *TransferMemoryBudget { + return &TransferMemoryBudget{ + totalByteCount: totalByteCount, + } +} + +func (self *TransferMemoryBudget) TotalByteCount() ByteCount { + return self.totalByteCount +} + +// Available is the unreserved remainder of the budget +func (self *TransferMemoryBudget) Available() ByteCount { + return max(0, self.totalByteCount-self.usedByteCount.Load()) +} + +func (self *TransferMemoryBudget) UsedByteCount() ByteCount { + return self.usedByteCount.Load() +} + +// Reserve takes bytes from the budget. It always succeeds (see the overdraft +// note in the type doc); admission gates on `Available`. +func (self *TransferMemoryBudget) Reserve(byteCount ByteCount) { + self.usedByteCount.Add(byteCount) + self.reservedByteCount.Add(byteCount) +} + +// Release returns bytes to the budget +func (self *TransferMemoryBudget) Release(byteCount ByteCount) { + used := self.usedByteCount.Add(-byteCount) + self.releasedByteCount.Add(byteCount) + if used < 0 { + // accounting bug: more released than reserved. + // log unconditionally so production sees it (tests see it as a + // negative used count breaking the balance assertions) + DefaultLogger().Errorf("[tmb]release below zero (%d)", used) + } +} + +// Counts returns the cumulative reserved/released byte counts. +// reserved-released equals the currently used bytes: it returns to zero when +// every borrowing queue is drained or cleared, so growth across a +// build/load/teardown cycle attributes a lost release. +func (self *TransferMemoryBudget) Counts() (reservedByteCount ByteCount, releasedByteCount ByteCount) { + return self.reservedByteCount.Load(), self.releasedByteCount.Load() +} diff --git a/transfer_peer_manager.go b/transfer_peer_manager.go new file mode 100644 index 0000000..55c6d90 --- /dev/null +++ b/transfer_peer_manager.go @@ -0,0 +1,166 @@ +package connect + +import ( + "context" + "slices" + "sync" + "time" + + "github.com/urnetwork/connect/protocol" +) + +func DefaultPeerManagerSettings() *PeerManagerSettings { + return &PeerManagerSettings{ + // matches the platform's recently disconnected peer window + DisconnectedPeerWindow: 5 * time.Minute, + } +} + +type PeerManagerSettings struct { + // how long a disconnect marker counts toward the disconnected peer count + DisconnectedPeerWindow time.Duration +} + +// PeerManager tracks the network peers of this client: +// the connected top-level clients in the same network and their identity metadata. +// The platform announces the complete list on each connect +// (`NetworkPeersReset` followed by `NetworkPeersUpdate` messages) +// and diffs after. Disconnected peers arrive as disconnect markers +// and are aged out locally by the disconnected peer window. +type PeerManager struct { + ctx context.Context + + client *Client + + peerManagerSettings *PeerManagerSettings + + peersMonitor *Monitor + + stateLock sync.Mutex + connectedPeers map[Id]*NetworkPeer + disconnectTimes map[Id]time.Time +} + +func NewPeerManager(ctx context.Context, client *Client, peerManagerSettings *PeerManagerSettings) *PeerManager { + return &PeerManager{ + ctx: ctx, + client: client, + peerManagerSettings: peerManagerSettings, + peersMonitor: NewMonitor(), + connectedPeers: map[Id]*NetworkPeer{}, + disconnectTimes: map[Id]time.Time{}, + } +} + +func (self *PeerManager) Client() *Client { + return self.client +} + +// PeersMonitor notifies on any change to the peer state +func (self *PeerManager) PeersMonitor() *Monitor { + return self.peersMonitor +} + +// ReceiveFunction +func (self *PeerManager) Receive(source TransferPath, frames []*protocol.Frame, peer Peer) { + if source.IsControlSource() { + for _, frame := range frames { + // ignore error + self.handleControlFrame(frame) + } + } +} + +func (self *PeerManager) handleControlFrame(frame *protocol.Frame) error { + switch frame.MessageType { + case protocol.MessageType_TransferNetworkPeersReset, protocol.MessageType_TransferNetworkPeersUpdate: + message, err := FromFrame(frame) + if err != nil { + return err + } + + changed := false + switch v := message.(type) { + case *protocol.NetworkPeersReset: + changed = self.resetPeers() + case *protocol.NetworkPeersUpdate: + changed, err = self.updatePeers(v) + if err != nil { + return err + } + } + if changed { + self.peersMonitor.NotifyAll() + } + } + return nil +} + +func (self *PeerManager) resetPeers() bool { + self.stateLock.Lock() + defer self.stateLock.Unlock() + + changed := 0 < len(self.connectedPeers) || 0 < len(self.disconnectTimes) + clear(self.connectedPeers) + clear(self.disconnectTimes) + return changed +} + +func (self *PeerManager) updatePeers(update *protocol.NetworkPeersUpdate) (bool, error) { + self.stateLock.Lock() + defer self.stateLock.Unlock() + + changed := false + for _, p := range update.Peers { + clientId, err := IdFromBytes(p.ClientId) + if err != nil { + return changed, err + } + if p.DisconnectTime != nil { + if _, ok := self.connectedPeers[clientId]; ok { + delete(self.connectedPeers, clientId) + changed = true + } + disconnectTime := time.UnixMilli(int64(*p.DisconnectTime)) + if t, ok := self.disconnectTimes[clientId]; !ok || !t.Equal(disconnectTime) { + self.disconnectTimes[clientId] = disconnectTime + changed = true + } + } else { + delete(self.disconnectTimes, clientId) + self.connectedPeers[clientId] = &NetworkPeer{ + ClientId: clientId, + ProvideModes: p.ProvideModes, + ProvideEnabled: slices.Contains(p.ProvideModes, protocol.ProvideMode_Network), + Principal: p.Principal, + Roles: p.Roles, + DeviceName: p.DeviceName, + DeviceSpec: p.DeviceSpec, + } + changed = true + } + } + return changed, nil +} + +// NetworkPeers enumerates the connected peers and the count of +// peers disconnected within the disconnected peer window +func (self *PeerManager) NetworkPeers() (connected []*NetworkPeer, disconnectedCount int) { + self.stateLock.Lock() + defer self.stateLock.Unlock() + + // age out disconnect markers older than the window + windowStart := time.Now().Add(-self.peerManagerSettings.DisconnectedPeerWindow) + for clientId, disconnectTime := range self.disconnectTimes { + if disconnectTime.Before(windowStart) { + delete(self.disconnectTimes, clientId) + } + } + + connected = make([]*NetworkPeer, 0, len(self.connectedPeers)) + for _, networkPeer := range self.connectedPeers { + connected = append(connected, networkPeer) + } + disconnectedCount = len(self.disconnectTimes) + return +} diff --git a/transfer_queue.go b/transfer_queue.go index e010fd1..df4c913 100644 --- a/transfer_queue.go +++ b/transfer_queue.go @@ -68,6 +68,15 @@ type transferQueue[T transferQueueItem] struct { byteCount ByteCount stateLock sync.Mutex + // shared budget accounting (see `TransferMemoryBudget`): the bytes held + // above `minByteCount` are borrowed from `budget`. every byte count + // change funnels through `updateByteCountWithLock`, so every add and + // removal path reserves/releases by construction. nil budget disables + // borrowing (independent per-queue caps). + budget *TransferMemoryBudget + minByteCount ByteCount + borrowedByteCount ByteCount + cmp TransferQueueCmpFunction[T] } @@ -84,6 +93,56 @@ func newTransferQueue[T transferQueueItem](cmp TransferQueueCmpFunction[T]) *tra return transferQueue } +func (self *transferQueue[T]) setBudget(budget *TransferMemoryBudget, minByteCount ByteCount) { + self.budget = budget + self.minByteCount = minByteCount +} + +func (self *transferQueue[T]) updateByteCountWithLock(deltaByteCount ByteCount) { + self.byteCount += deltaByteCount + if self.budget != nil { + borrowTargetByteCount := max(0, self.byteCount-self.minByteCount) + if borrowTargetByteCount < self.borrowedByteCount { + self.budget.Release(self.borrowedByteCount - borrowTargetByteCount) + self.borrowedByteCount = borrowTargetByteCount + } else if self.borrowedByteCount < borrowTargetByteCount { + self.budget.Reserve(borrowTargetByteCount - self.borrowedByteCount) + self.borrowedByteCount = borrowTargetByteCount + } + } +} + +func (self *transferQueue[T]) CanAdd(byteCount ByteCount, maxByteCount ByteCount) bool { + self.stateLock.Lock() + defer self.stateLock.Unlock() + + if len(self.orderedItems) == 0 { + return true + } + if maxByteCount <= self.byteCount+byteCount { + return false + } + if self.budget == nil { + return true + } + probeByteCount := max(self.byteCount+byteCount, self.byteCount+1) + borrowTargetByteCount := max(0, probeByteCount-self.minByteCount) + return borrowTargetByteCount-self.borrowedByteCount <= self.budget.Available() +} + +func (self *transferQueue[T]) Clear() []T { + self.stateLock.Lock() + defer self.stateLock.Unlock() + + items := self.orderedItems + self.orderedItems = []T{} + self.maxHeap = newTransferQueueMaxHeap[T](self.cmp) + clear(self.messageIdItems) + clear(self.sequenceNumberItems) + self.updateByteCountWithLock(-self.byteCount) + return items +} + func (self *transferQueue[T]) QueueSize() (int, ByteCount) { self.stateLock.Lock() defer self.stateLock.Unlock() @@ -112,7 +171,7 @@ func (self *transferQueue[T]) Add(item T) { self.sequenceNumberItems[item.SequenceNumber()] = item heap.Push(self, item) heap.Push(self.maxHeap, item) - self.byteCount += item.MessageByteCount() + self.updateByteCountWithLock(item.MessageByteCount()) } func (self *transferQueue[T]) ContainsMessageId(messageId Id) (sequenceNumber uint64, ok bool) { @@ -182,7 +241,7 @@ func (self *transferQueue[T]) remove(item T) T { panic("Heap invariant broken.") } heap.Remove(self.maxHeap, item.MaxHeapIndex()) - self.byteCount -= item.MessageByteCount() + self.updateByteCountWithLock(-item.MessageByteCount()) return item } @@ -199,7 +258,7 @@ func (self *transferQueue[T]) RemoveFirst() T { heap.Remove(self.maxHeap, item.MaxHeapIndex()) delete(self.messageIdItems, item.MessageId()) delete(self.sequenceNumberItems, item.SequenceNumber()) - self.byteCount -= item.MessageByteCount() + self.updateByteCountWithLock(-item.MessageByteCount()) return item } diff --git a/transfer_stream_manager.go b/transfer_stream_manager.go index 5b115d7..26ba40d 100644 --- a/transfer_stream_manager.go +++ b/transfer_stream_manager.go @@ -69,7 +69,7 @@ func (self *StreamManager) WebRtcManager() *WebRtcManager { } // ReceiveFunction -func (self *StreamManager) Receive(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { +func (self *StreamManager) Receive(source TransferPath, frames []*protocol.Frame, peer Peer) { if source.IsControlSource() { for _, frame := range frames { // ignore error diff --git a/transfer_test.go b/transfer_test.go index 216bfac..2cca51e 100644 --- a/transfer_test.go +++ b/transfer_test.go @@ -200,7 +200,7 @@ func runSendReceiveSenderReset(t *testing.T, encMode encryptionMode) { acks := make(chan error) receives := make(chan *protocol.SimpleMessage) - b.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { + b.AddReceiveCallback(func(source TransferPath, frames []*protocol.Frame, peer Peer) { for _, frame := range frames { m, err := FromFrame(frame) if err != nil { diff --git a/transport_p2p_webrtc.go b/transport_p2p_webrtc.go index a83cbed..b248f84 100644 --- a/transport_p2p_webrtc.go +++ b/transport_p2p_webrtc.go @@ -127,7 +127,7 @@ func ReceiveSignalsFromClient(client *Client, receiver SignalReceiver) func() { } // ReceiveFunction -func (self *clientSignalReceiver) Receive(source TransferPath, frames []*protocol.Frame, provideMode protocol.ProvideMode) { +func (self *clientSignalReceiver) Receive(source TransferPath, frames []*protocol.Frame, peer Peer) { for _, frame := range frames { self.handleControlFrame(source, frame) } diff --git a/util.go b/util.go index 227f5e6..a50cec6 100644 --- a/util.go +++ b/util.go @@ -59,6 +59,21 @@ func DetectEffectiveRAMLimitBytes() int64 { return 850 * 1024 * 1024 } +var memoryShedders = NewCallbackList[func()]() + +func AddMemoryShedder(shed func()) func() { + callbackId := memoryShedders.Add(shed) + return func() { + memoryShedders.Remove(callbackId) + } +} + +func ShedMemory() { + for _, shed := range memoryShedders.Get() { + HandleError(shed) + } +} + type Monitor struct { mutex sync.Mutex notify chan struct{}