Add atomic id guard to prevent TOCTOU (#131)

* add atomic id guard to prevent TOCTOU

Signed-off-by: Marvin Drees <marvin.drees@9elements.com>

* implement review feedback to tcp conn guard

Signed-off-by: Marvin Drees <marvin.drees@9elements.com>

---------

Signed-off-by: Marvin Drees <marvin.drees@9elements.com>

internet/conn_closerace_test.go

@@ -0,0 +1,214 @@
+package internet
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"math/rand"
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/soypat/lneto/tcp"
+)
+
+// TestConn_ConcurrentCloseDoesNotDropWrite is a regression test for issue #82:
+// when one goroutine has a Write in progress (blocked because the TX buffer is
+// full) and another goroutine calls Close, the in-flight write data must not be
+// silently dropped. With the atomic write-lock, Close waits for the in-progress
+// write to finish queueing all of its data before tearing the connection down.
+//
+// The scenario is made deterministic by filling the TX buffer before any packets
+// are exchanged: the writer blocks holding the write lock, Close is then issued
+// (and must wait), and only afterwards is the packet pump started so the buffer
+// can drain and the writer can run to completion.
+func TestConn_ConcurrentCloseDoesNotDropWrite(t *testing.T) {
+	rng := rand.New(rand.NewSource(1))
+	var sbCl, sbSv StackIPv4
+	var connCl, connSv tcp.Conn
+	setupClientServerEstablished(t, rng, &sbCl, &sbSv, &connCl, &connSv)
+
+	// Payload several times larger than the 2048-byte TX buffer so the writer
+	// must block waiting for the buffer to drain across many segments.
+	payload := make([]byte, 4*2048)
+	for i := range payload {
+		payload[i] = byte(i)
+	}
+
+	// Watchdog: break any deadlock so a regression fails fast instead of hanging.
+	watchdog := time.AfterFunc(10*time.Second, func() {
+		t.Error("test timed out: Close likely blocked waiting on a stuck write")
+		connCl.Abort()
+		connSv.Abort()
+	})
+	defer watchdog.Stop()
+
+	// Writer (G1): writes the whole payload. Blocks once the TX buffer fills.
+	type writeResult struct {
+		n   int
+		err error
+	}
+	writeDone := make(chan writeResult, 1)
+	go func() {
+		n, err := connCl.Write(payload)
+		writeDone <- writeResult{n, err}
+	}()
+
+	// Wait until the writer has filled the TX buffer and is blocked. At this
+	// point it owns the write lock, reproducing "Write in progress".
+	for connCl.FreeOutput() != 0 {
+		runtime.Gosched()
+	}
+
+	// Close (G2): issued while the write is in progress. Must wait for the
+	// writer to drain instead of dropping its data.
+	closeDone := make(chan error, 1)
+	go func() {
+		closeDone <- connCl.Close()
+	}()
+
+	// Reader: drains the server side until EOF, accumulating everything received.
+	readDone := make(chan []byte, 1)
+	go func() {
+		got := make([]byte, 0, len(payload))
+		rbuf := make([]byte, 1024)
+		for {
+			n, err := connSv.Read(rbuf)
+			got = append(got, rbuf[:n]...)
+			if err != nil {
+				break
+			}
+		}
+		readDone <- got
+	}()
+
+	// Packet pump: only now do we move packets between the two stacks, letting
+	// the buffer drain so the blocked writer can finish.
+	var stop atomic.Bool
+	pumpStopped := make(chan struct{})
+	go func() {
+		defer close(pumpStopped)
+		var buf [2048]byte
+		for !stop.Load() {
+			progressed := false
+			if n, err := sbCl.Encapsulate(buf[:], 0, 0); err == nil && n > 0 {
+				_ = sbSv.Demux(buf[:n], 0)
+				progressed = true
+			}
+			if n, err := sbSv.Encapsulate(buf[:], 0, 0); err == nil && n > 0 {
+				_ = sbCl.Demux(buf[:n], 0)
+				progressed = true
+			}
+			if !progressed {
+				runtime.Gosched()
+			}
+		}
+	}()
+
+	wr := <-writeDone
+	if wr.err != nil {
+		t.Errorf("issue #82: Write returned error %v; concurrent Close dropped in-flight data", wr.err)
+	}
+	if wr.n != len(payload) {
+		t.Errorf("issue #82: Write wrote %d of %d bytes; concurrent Close truncated the write", wr.n, len(payload))
+	}
+	if err := <-closeDone; err != nil {
+		t.Errorf("Close returned error: %v", err)
+	}
+	got := <-readDone
+	stop.Store(true)
+	<-pumpStopped
+
+	if len(got) != len(payload) {
+		t.Fatalf("server received %d of %d bytes; data was dropped by concurrent Close", len(got), len(payload))
+	}
+	if !bytes.Equal(got, payload) {
+		t.Fatal("server received corrupted data")
+	}
+}
+
+// TestConn_ConcurrentWritesSerialize verifies that the atomic write-lock
+// serializes concurrent Write calls on the same connection so their payloads are
+// not interleaved on the wire, and that all bytes from both writers are delivered.
+func TestConn_ConcurrentWritesSerialize(t *testing.T) {
+	rng := rand.New(rand.NewSource(2))
+	var sbCl, sbSv StackIPv4
+	var connCl, connSv tcp.Conn
+	setupClientServerEstablished(t, rng, &sbCl, &sbSv, &connCl, &connSv)
+
+	const chunk = 1500
+	a := bytes.Repeat([]byte{0xAA}, chunk)
+	b := bytes.Repeat([]byte{0xBB}, chunk)
+
+	watchdog := time.AfterFunc(10*time.Second, func() {
+		t.Error("test timed out")
+		connCl.Abort()
+		connSv.Abort()
+	})
+	defer watchdog.Stop()
+
+	var stop atomic.Bool
+	pumpStopped := make(chan struct{})
+	go func() {
+		defer close(pumpStopped)
+		var buf [2048]byte
+		for !stop.Load() {
+			progressed := false
+			if n, err := sbCl.Encapsulate(buf[:], 0, 0); err == nil && n > 0 {
+				_ = sbSv.Demux(buf[:n], 0)
+				progressed = true
+			}
+			if n, err := sbSv.Encapsulate(buf[:], 0, 0); err == nil && n > 0 {
+				_ = sbCl.Demux(buf[:n], 0)
+				progressed = true
+			}
+			if !progressed {
+				runtime.Gosched()
+			}
+		}
+	}()
+
+	writeErr := make(chan error, 2)
+	writer := func(b []byte) {
+		n, err := connCl.Write(b)
+		if err == nil && n != len(b) {
+			err = io.ErrShortWrite
+		}
+		writeErr <- err
+	}
+	go writer(a)
+	go writer(b)
+
+	got := make([]byte, 0, 2*chunk)
+	rbuf := make([]byte, 1024)
+	for len(got) < 2*chunk {
+		n, err := connSv.Read(rbuf)
+		got = append(got, rbuf[:n]...)
+		if err != nil && !errors.Is(err, io.EOF) {
+			break
+		}
+	}
+	for range 2 {
+		if err := <-writeErr; err != nil {
+			t.Errorf("concurrent write failed: %v", err)
+		}
+	}
+	stop.Store(true)
+	<-pumpStopped
+
+	if len(got) != 2*chunk {
+		t.Fatalf("received %d bytes, want %d", len(got), 2*chunk)
+	}
+	// Serialized writes mean one chunk's bytes appear fully before the other's;
+	// the boundary is a single transition, never interleaved.
+	transitions := 0
+	for i := 1; i < len(got); i++ {
+		if got[i] != got[i-1] {
+			transitions++
+		}
+	}
+	if transitions != 1 {
+		t.Fatalf("expected exactly one A/B boundary (serialized writes), got %d transitions", transitions)
+	}
+}

tcp/conn.go

@@ -8,6 +8,7 @@ import (
 	"net/netip"
 	"os"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/soypat/lneto"
@@ -29,6 +30,13 @@ type Conn struct {
 	h          Handler
 	remoteAddr []byte
 
+	// writeLock holds the connection ID ([Handler.connid]) of the goroutine
+	// that currently owns the write side via [Conn.Write] or [Conn.Flush].
+	// A zero value means no write is in progress. It serializes concurrent
+	// writers and lets [Conn.Close] acquire the write side before closing so it
+	// does not drop in-flight data (see issue #82).
+	writeLock atomic.Uint64
+
 	_backoff lneto.BackoffStrategy
 	rdead    time.Time
 	wdead    time.Time
@@ -49,6 +57,7 @@ func (conn *Conn) reset(h Handler) {
 	conn.wdead = time.Time{}
 	conn.abortErr = nil
 	conn.ipID = 0
+	conn.writeLock.Store(0)
 }
 
 // ConnConfig provides configuration parameters for [Conn].
@@ -213,6 +222,11 @@ func (conn *Conn) CloseRead() error {
 
 // Close will initiate TCP close sequence. After Close is called future [Conn.Write] calls will fail with [net.ErrClosed].
 func (conn *Conn) Close() error {
+	connid, err := conn.acquireWriteLock(nil)
+	if err != nil {
+		return err
+	}
+	defer conn.releaseWriteLock(connid)
 	conn.mu.Lock()
 	defer conn.mu.Unlock()
 	conn.trace("TCPConn.Close", slog.Uint64("lport", uint64(conn.h.localPort)), slog.Uint64("rport", uint64(conn.h.remotePort)))
@@ -237,10 +251,11 @@ func (conn *Conn) InternalHandler() *Handler {
 
 // Write writes argument data to the TCPConns's output buffer which is queued to be sent.
 func (conn *Conn) Write(b []byte) (int, error) {
-	connid, err := conn.lockPipeConnID()
+	connid, err := conn.acquireWriteLock(&conn.wdead)
 	if err != nil {
 		return 0, err
 	}
+	defer conn.releaseWriteLock(connid)
 	rport := conn.RemotePort()
 	plen := len(b)
 	lport := conn.LocalPort()
@@ -286,10 +301,11 @@ func (conn *Conn) Write(b []byte) (int, error) {
 
 // Flush blocks until all buffered TCP data has been sent.
 func (conn *Conn) Flush() error {
-	connid, err := conn.lockPipeConnID()
+	connid, err := conn.acquireWriteLock(&conn.wdead)
 	if err != nil {
 		return err
 	}
+	defer conn.releaseWriteLock(connid)
 	var backoffs uint
 	for conn.BufferedUnsent() != 0 {
 		if err := conn.checkPipe(connid, &conn.wdead); err != nil {
@@ -350,14 +366,31 @@ func (conn *Conn) Read(b []byte) (int, error) {
 	return n, nil
 }
 
-func (conn *Conn) lockPipeConnID() (uint64, error) {
+// acquireWriteLock validates the connection is open and acquires the write lock
+// for the current connection, returning its connection ID. It serializes
+// concurrent writers: while another goroutine owns the write side it blocks with
+// backoff until that writer releases, the connection closes, or deadline elapses.
+// The returned connID must be released with [Conn.releaseWriteLock].
+func (conn *Conn) acquireWriteLock(deadline *time.Time) (connID uint64, err error) {
 	conn.mu.Lock()
-	defer conn.mu.Unlock()
+	connID = conn.h.connid
-	err := conn.checkPipeOpen()
+	conn.mu.Unlock()
-	if err != nil {
+	var backoffs uint
-		return 0, err
+	for {
+		if err := conn.checkPipe(connID, deadline); err != nil {
+			return 0, err
+		}
+		if conn.writeLock.CompareAndSwap(0, connID) {
+			return connID, nil
+		}
+		conn.backoff(backoffs)
+		backoffs++
 	}
-	return conn.h.connid, nil
+}
+
+// releaseWriteLock releases a write lock previously acquired by [Conn.acquireWriteLock].
+func (conn *Conn) releaseWriteLock(connID uint64) {
+	conn.writeLock.CompareAndSwap(connID, 0)
 }
 
 func (conn *Conn) checkPipe(connID uint64, deadline *time.Time) (err error) {
@@ -365,9 +398,9 @@ func (conn *Conn) checkPipe(connID uint64, deadline *time.Time) (err error) {
 	defer conn.mu.Unlock()
 	if conn.abortErr != nil {
 		err = conn.abortErr
-	} else if connID != conn.h.connid {
+	} else if connID != conn.h.connid || conn.h.State().IsClosed() {
 		err = net.ErrClosed
-	} else if !deadline.IsZero() && time.Since(*deadline) > 0 {
+	} else if deadline != nil && !deadline.IsZero() && time.Since(*deadline) > 0 {
 		err = errDeadlineExceeded
 	}
 	return err