fix(dcutr): Fix end to end tests and add legacy behavior flag (default=true) (#3044)

Author: MarcoPolo

p2p/net/simconn/router.go

@@ -0,0 +1,179 @@
+package simconn
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"sync"
+	"time"
+)
+
+type PacketReciever interface {
+	RecvPacket(p Packet)
+}
+
+// PerfectRouter is a router that has no latency or jitter and can route to
+// every node
+type PerfectRouter struct {
+	mu    sync.Mutex
+	nodes map[net.Addr]PacketReciever
+}
+
+// SendPacket implements Router.
+func (r *PerfectRouter) SendPacket(p Packet) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	conn, ok := r.nodes[p.To]
+	if !ok {
+		return errors.New("unknown destination")
+	}
+
+	conn.RecvPacket(p)
+	return nil
+}
+
+func (r *PerfectRouter) AddNode(addr net.Addr, conn PacketReciever) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.nodes == nil {
+		r.nodes = make(map[net.Addr]PacketReciever)
+	}
+	r.nodes[addr] = conn
+}
+
+func (r *PerfectRouter) RemoveNode(addr net.Addr) {
+	delete(r.nodes, addr)
+}
+
+var _ Router = &PerfectRouter{}
+
+type DelayedPacketReciever struct {
+	inner PacketReciever
+	delay time.Duration
+}
+
+func (r *DelayedPacketReciever) RecvPacket(p Packet) {
+	time.AfterFunc(r.delay, func() { r.inner.RecvPacket(p) })
+}
+
+type FixedLatencyRouter struct {
+	PerfectRouter
+	latency time.Duration
+}
+
+func (r *FixedLatencyRouter) SendPacket(p Packet) error {
+	return r.PerfectRouter.SendPacket(p)
+}
+
+func (r *FixedLatencyRouter) AddNode(addr net.Addr, conn PacketReciever) {
+	r.PerfectRouter.AddNode(addr, &DelayedPacketReciever{
+		inner: conn,
+		delay: r.latency,
+	})
+}
+
+var _ Router = &FixedLatencyRouter{}
+
+type simpleNodeFirewall struct {
+	mu                sync.Mutex
+	publiclyReachable bool
+	packetsOutTo      map[string]struct{}
+	node              *SimConn
+}
+
+func (f *simpleNodeFirewall) MarkPacketSentOut(p Packet) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if f.packetsOutTo == nil {
+		f.packetsOutTo = make(map[string]struct{})
+	}
+	f.packetsOutTo[p.To.String()] = struct{}{}
+}
+
+func (f *simpleNodeFirewall) IsPacketInAllowed(p Packet) bool {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if f.publiclyReachable {
+		return true
+	}
+
+	_, ok := f.packetsOutTo[p.From.String()]
+	return ok
+}
+
+func (f *simpleNodeFirewall) String() string {
+	return fmt.Sprintf("public: %v, packetsOutTo: %v", f.publiclyReachable, f.packetsOutTo)
+}
+
+type SimpleFirewallRouter struct {
+	mu    sync.Mutex
+	nodes map[string]*simpleNodeFirewall
+}
+
+func (r *SimpleFirewallRouter) String() string {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	nodes := make([]string, 0, len(r.nodes))
+	for _, node := range r.nodes {
+		nodes = append(nodes, node.String())
+	}
+	return fmt.Sprintf("%v", nodes)
+}
+
+func (r *SimpleFirewallRouter) SendPacket(p Packet) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	toNode, exists := r.nodes[p.To.String()]
+	if !exists {
+		return errors.New("unknown destination")
+	}
+
+	// Record that this node is sending a packet to the destination
+	fromNode, exists := r.nodes[p.From.String()]
+	if !exists {
+		return errors.New("unknown source")
+	}
+	fromNode.MarkPacketSentOut(p)
+
+	if !toNode.IsPacketInAllowed(p) {
+		return nil // Silently drop blocked packets
+	}
+
+	toNode.node.RecvPacket(p)
+	return nil
+}
+
+func (r *SimpleFirewallRouter) AddNode(addr net.Addr, conn *SimConn) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.nodes == nil {
+		r.nodes = make(map[string]*simpleNodeFirewall)
+	}
+	r.nodes[addr.String()] = &simpleNodeFirewall{
+		packetsOutTo: make(map[string]struct{}),
+		node:         conn,
+	}
+}
+
+func (r *SimpleFirewallRouter) AddPubliclyReachableNode(addr net.Addr, conn *SimConn) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.nodes == nil {
+		r.nodes = make(map[string]*simpleNodeFirewall)
+	}
+	r.nodes[addr.String()] = &simpleNodeFirewall{
+		publiclyReachable: true,
+		node:              conn,
+	}
+}
+
+func (r *SimpleFirewallRouter) RemoveNode(addr net.Addr) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.nodes == nil {
+		return
+	}
+	delete(r.nodes, addr.String())
+}
+
+var _ Router = &SimpleFirewallRouter{}

p2p/net/simconn/simconn.go

@@ -0,0 +1,218 @@
+package simconn
+
+import (
+	"errors"
+	"net"
+	"slices"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+var ErrDeadlineExceeded = errors.New("deadline exceeded")
+
+type Router interface {
+	SendPacket(p Packet) error
+}
+
+type Packet struct {
+	To   net.Addr
+	From net.Addr
+	buf  []byte
+}
+
+type SimConn struct {
+	mu         sync.Mutex
+	closed     bool
+	closedChan chan struct{}
+
+	packetsSent atomic.Uint64
+	packetsRcvd atomic.Uint64
+	bytesSent   atomic.Int64
+	bytesRcvd   atomic.Int64
+
+	router Router
+
+	myAddr        *net.UDPAddr
+	myLocalAddr   net.Addr
+	packetsToRead chan Packet
+
+	readDeadline  time.Time
+	writeDeadline time.Time
+}
+
+// NewSimConn creates a new simulated connection with the specified parameters
+func NewSimConn(addr *net.UDPAddr, rtr Router) *SimConn {
+	return &SimConn{
+		router:        rtr,
+		myAddr:        addr,
+		packetsToRead: make(chan Packet, 512), // buffered channel to prevent blocking
+		closedChan:    make(chan struct{}),
+	}
+}
+
+type ConnStats struct {
+	BytesSent   int
+	BytesRcvd   int
+	PacketsSent int
+	PacketsRcvd int
+}
+
+func (c *SimConn) Stats() ConnStats {
+	return ConnStats{
+		BytesSent:   int(c.bytesSent.Load()),
+		BytesRcvd:   int(c.bytesRcvd.Load()),
+		PacketsSent: int(c.packetsSent.Load()),
+		PacketsRcvd: int(c.packetsRcvd.Load()),
+	}
+}
+
+// SetLocalAddr only changes what `.LocalAddr()` returns.
+// Packets will still come From the initially configured addr.
+func (c *SimConn) SetLocalAddr(addr net.Addr) {
+	c.myLocalAddr = addr
+}
+
+func (c *SimConn) RecvPacket(p Packet) {
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return
+	}
+	c.mu.Unlock()
+	c.packetsRcvd.Add(1)
+	c.bytesRcvd.Add(int64(len(p.buf)))
+
+	select {
+	case c.packetsToRead <- p:
+	default:
+		// drop the packet if the channel is full
+	}
+}
+
+var _ net.PacketConn = &SimConn{}
+
+// Close implements net.PacketConn
+func (c *SimConn) Close() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed {
+		return nil
+	}
+	c.closed = true
+	close(c.closedChan)
+	return nil
+}
+
+// ReadFrom implements net.PacketConn
+func (c *SimConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return 0, nil, net.ErrClosed
+	}
+	deadline := c.readDeadline
+	c.mu.Unlock()
+
+	if !deadline.IsZero() && time.Now().After(deadline) {
+		return 0, nil, ErrDeadlineExceeded
+	}
+
+	var pkt Packet
+	if !deadline.IsZero() {
+		select {
+		case pkt = <-c.packetsToRead:
+		case <-time.After(time.Until(deadline)):
+			return 0, nil, ErrDeadlineExceeded
+		}
+	} else {
+		pkt = <-c.packetsToRead
+	}
+
+	n = copy(p, pkt.buf)
+	// if the provided buffer is not enough to read the whole packet, we drop
+	// the rest of the data. this is similar to what `recvfrom` does on Linux
+	// and macOS.
+	return n, pkt.From, nil
+}
+
+// WriteTo implements net.PacketConn
+func (c *SimConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return 0, net.ErrClosed
+	}
+	deadline := c.writeDeadline
+	c.mu.Unlock()
+
+	if !deadline.IsZero() && time.Now().After(deadline) {
+		return 0, ErrDeadlineExceeded
+	}
+
+	c.packetsSent.Add(1)
+	c.bytesSent.Add(int64(len(p)))
+
+	pkt := Packet{
+		From: c.myAddr,
+		To:   addr,
+		buf:  slices.Clone(p),
+	}
+	return len(p), c.router.SendPacket(pkt)
+}
+
+func (c *SimConn) UnicastAddr() net.Addr {
+	return c.myAddr
+}
+
+// LocalAddr implements net.PacketConn
+func (c *SimConn) LocalAddr() net.Addr {
+	if c.myLocalAddr != nil {
+		return c.myLocalAddr
+	}
+	return c.myAddr
+}
+
+// SetDeadline implements net.PacketConn
+func (c *SimConn) SetDeadline(t time.Time) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.readDeadline = t
+	c.writeDeadline = t
+	return nil
+}
+
+// SetReadDeadline implements net.PacketConn
+func (c *SimConn) SetReadDeadline(t time.Time) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.readDeadline = t
+	return nil
+}
+
+// SetWriteDeadline implements net.PacketConn
+func (c *SimConn) SetWriteDeadline(t time.Time) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.writeDeadline = t
+	return nil
+}
+
+func IntToPublicIPv4(n int) net.IP {
+	n += 1
+	// Avoid private IP ranges
+	b := make([]byte, 4)
+	b[0] = byte((n>>24)&0xFF | 1)
+	b[1] = byte((n >> 16) & 0xFF)
+	b[2] = byte((n >> 8) & 0xFF)
+	b[3] = byte(n & 0xFF)
+
+	ip := net.IPv4(b[0], b[1], b[2], b[3])
+
+	// Check and modify if it's in private ranges
+	if ip.IsPrivate() {
+		b[0] = 1 // Use 1.x.x.x as public range
+	}
+
+	return ip
+}