From 36ec8860af49dcc8653ed29c6ef15b59d211a9ca Mon Sep 17 00:00:00 2001
From: Simone Basso <bassosimone@gmail.com>
Date: Sat, 23 Nov 2024 09:34:28 +0100
Subject: [PATCH] feat(netsim): implement ECONNREFUSED

---
 netsim/errno_unix.go    |  3 +++
 netsim/errno_windows.go |  3 +++
 netsim/stack.go         | 41 ++++++++++++++++++++++++++++++++++++++---
 netsim/tcpconn.go       |  4 ++++
 4 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/netsim/errno_unix.go b/netsim/errno_unix.go
index c970b74..610b0dd 100644
--- a/netsim/errno_unix.go
+++ b/netsim/errno_unix.go
@@ -20,6 +20,9 @@ const (
 	// ECONNABORTED is the connection aborted error.
 	ECONNABORTED = unix.ECONNABORTED
 
+	// ECONNREFUSED is the connection refused error.
+	ECONNREFUSED = unix.ECONNREFUSED
+
 	// ECONNRESET is the connection reset by peer error.
 	ECONNRESET = unix.ECONNRESET
 
diff --git a/netsim/errno_windows.go b/netsim/errno_windows.go
index b25de1b..79d047b 100644
--- a/netsim/errno_windows.go
+++ b/netsim/errno_windows.go
@@ -20,6 +20,9 @@ const (
 	// ECONNABORTED is the connection aborted error.
 	ECONNABORTED = windows.WSAECONNABORTED
 
+	// ECONNREFUSED is the connection refused error.
+	ECONNREFUSED = windows.WSAECONNREFUSED
+
 	// ECONNRESET is the connection reset by peer error.
 	ECONNRESET = windows.WSAECONNRESET
 
diff --git a/netsim/stack.go b/netsim/stack.go
index 3af067b..02616a7 100644
--- a/netsim/stack.go
+++ b/netsim/stack.go
@@ -13,6 +13,8 @@ import (
 	"net"
 	"net/netip"
 	"sync"
+
+	"github.com/rbmk-project/common/runtimex"
 )
 
 // Stack models a network stack.
@@ -46,7 +48,17 @@ type Stack struct {
 // goroutine demuxing incoming traffic. Remember to invoke
 // Close to stop any muxing/demuxing goroutine.
 func NewStack(addrs ...netip.Addr) *Stack {
-	const firstEphemeralPort = 49152
+	const (
+		// outputBufSize adds some buffering to the output
+		// channel to allow for nonblocking writes of packets
+		// containing the RST flag in response to SYN for
+		// ports that aren't listening.
+		outputBufSize = 128
+
+		// firstEphemeralPort is the first ephemeral port
+		// to use according to RFC6335.
+		firstEphemeralPort = 49152
+	)
 	ns := &Stack{
 		addrs:   addrs,
 		eof:     make(chan struct{}),
@@ -56,7 +68,7 @@ func NewStack(addrs ...netip.Addr) *Stack {
 			IPProtocolTCP: firstEphemeralPort,
 			IPProtocolUDP: firstEphemeralPort,
 		},
-		output: make(chan *Packet),
+		output: make(chan *Packet, outputBufSize),
 		portmu: sync.RWMutex{},
 		ports:  map[PortAddr]*Port{},
 	}
@@ -151,6 +163,26 @@ func (ns *Stack) findPortLocked(pkt *Packet) *Port {
 	return nil
 }
 
+// resetNonblocking sends a RST packet in response to a SYN for a closed port.
+func (ns *Stack) resetNonblocking(pkt *Packet) {
+	runtimex.Assert(pkt.IPProtocol == IPProtocolTCP, "not a TCP packet")
+	runtimex.Assert(pkt.Flags != TCPFlagSYN, "expected SYN flags")
+	resp := &Packet{
+		SrcAddr:    pkt.DstAddr,
+		DstAddr:    pkt.SrcAddr,
+		IPProtocol: IPProtocolTCP,
+		SrcPort:    pkt.DstPort,
+		DstPort:    pkt.SrcPort,
+		Flags:      TCPFlagRST,
+		Payload:    []byte{},
+	}
+	// Nonblocking write to the buffered output channel.
+	select {
+	case ns.output <- resp:
+	default:
+	}
+}
+
 // demux demuxes a single incoming [*Packet].
 func (ns *Stack) demux(pkt *Packet) error {
 	// Discard packet if the address is not local.
@@ -164,7 +196,10 @@ func (ns *Stack) demux(pkt *Packet) error {
 	port := ns.findPortLocked(pkt)
 	ns.portmu.RUnlock()
 	if port == nil {
-		return EHOSTUNREACH
+		if pkt.IPProtocol == IPProtocolTCP && pkt.Flags == TCPFlagSYN {
+			ns.resetNonblocking(pkt)
+		}
+		return ECONNREFUSED
 	}
 
 	// Actually deliver the packet to the port.
diff --git a/netsim/tcpconn.go b/netsim/tcpconn.go
index 1e01918..5ecc9e2 100644
--- a/netsim/tcpconn.go
+++ b/netsim/tcpconn.go
@@ -62,6 +62,10 @@ func (c *TCPConn) Connect(ctx context.Context) (err error) {
 		if err != nil {
 			return
 		}
+		if pkt.Flags == TCPFlagRST {
+			err = ECONNREFUSED
+			return
+		}
 		if pkt.Flags != TCPFlagSYN|TCPFlagACK {
 			err = ECONNABORTED
 			return