From 860d6d00bb1488fbb04c314f5f34718f74dbaff9 Mon Sep 17 00:00:00 2001 From: Marvin Drees Date: Mon, 22 Jun 2026 14:56:38 +0200 Subject: [PATCH] Add atomic id guard to prevent TOCTOU (#131) * add atomic id guard to prevent TOCTOU Signed-off-by: Marvin Drees * implement review feedback to tcp conn guard Signed-off-by: Marvin Drees --------- Signed-off-by: Marvin Drees --- internet/conn_closerace_test.go | 214 ++++++++++++++++++++++++++++++++ tcp/conn.go | 53 ++++++-- 2 files changed, 257 insertions(+), 10 deletions(-) create mode 100644 internet/conn_closerace_test.go diff --git a/internet/conn_closerace_test.go b/internet/conn_closerace_test.go new file mode 100644 index 0000000..935f0ea --- /dev/null +++ b/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) + } +} diff --git a/tcp/conn.go b/tcp/conn.go index 5aac264..cf3d5b8 100644 --- a/tcp/conn.go +++ b/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() - err := conn.checkPipeOpen() - if err != nil { - return 0, err + connID = conn.h.connid + conn.mu.Unlock() + var backoffs uint + 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