fix: Ensure `SecureConnection` `Read` and `Write` methods are safe during concurrent `Close` operations by adding a closed state and mutex, verified by a new test.

lemon-mint committed Nov 19, 2025 at 09:38 UTC 4dc3c5128ace5fd396683563f730bddb5896498f
2 files changed +85
portal/core/cryptoops/handshaker.go
+31
@@ -9,6 +9,7 @@ import (
9 "errors"
10 "fmt"
11 "io"
12 + "net"
13 "slices"
14 "sync"
15 "time"
@@ -111,6 +112,8 @@ type SecureConnection struct {
112 readBuffer *bytebufferpool.ByteBuffer
113
114 // Ensure Close is safe and idempotent
115 + mu sync.RWMutex
116 + closed bool
117 closeOnce sync.Once
118 closeErr error
119 }
@@ -146,6 +149,13 @@ func (sc *SecureConnection) RemoteID() string {
149
150 // Write encrypts and writes data to the underlying connection
151 func (sc *SecureConnection) Write(p []byte) (int, error) {
152 + sc.mu.RLock()
153 + if sc.closed {
154 + sc.mu.RUnlock()
155 + return 0, net.ErrClosed
156 + }
157 + sc.mu.RUnlock()
158 +
159 const fragSize = maxRawPacketSize / 2
160 if len(p) > fragSize {
161 for i := 0; i < (len(p)+fragSize-1)/fragSize; i++ {
@@ -190,12 +200,20 @@ func (sc *SecureConnection) writeFragmentation(p []byte) (int, error) {
200
201 // Read reads and decrypts data from the underlying connection
202 func (sc *SecureConnection) Read(p []byte) (int, error) {
203 + sc.mu.RLock()
204 + if sc.closed {
205 + sc.mu.RUnlock()
206 + return 0, net.ErrClosed
207 + }
208 +
209 if sc.readBuffer != nil && len(sc.readBuffer.B) > 0 {
210 n := copy(p, sc.readBuffer.B)
211 copy(sc.readBuffer.B[:len(sc.readBuffer.B)-n], sc.readBuffer.B[n:])
212 sc.readBuffer.B = sc.readBuffer.B[:len(sc.readBuffer.B)-n]
213 + sc.mu.RUnlock()
214 return n, nil
215 }
216 + sc.mu.RUnlock()
217
218 // Read length prefix first (4 bytes)
219 lengthBuf := _lengthBufferPool.Get().(*[4]byte)
@@ -235,9 +253,19 @@ func (sc *SecureConnection) Read(p []byte) (int, error) {
253 return 0, ErrDecryptionFailed
254 }
255
256 + sc.mu.Lock()
257 + defer sc.mu.Unlock()
258 +
259 + if sc.closed {
260 + return 0, net.ErrClosed
261 + }
262 +
263 // Copy decrypted data to the provided buffer
264 n := copy(p, decrypted)
265 if n < len(decrypted) {
266 + if sc.readBuffer == nil {
267 + sc.readBuffer = acquireBuffer(len(decrypted) - n)
268 + }
269 sc.readBuffer.B = append(sc.readBuffer.B, decrypted[n:]...)
270 }
271
@@ -247,10 +275,13 @@ func (sc *SecureConnection) Read(p []byte) (int, error) {
275 // Close closes the underlying connection and releases resources
276 func (sc *SecureConnection) Close() error {
277 sc.closeOnce.Do(func() {
278 + sc.mu.Lock()
279 + sc.closed = true
280 if sc.readBuffer != nil {
281 releaseBuffer(sc.readBuffer)
282 sc.readBuffer = nil
283 }
284 + sc.mu.Unlock()
285 sc.closeErr = sc.conn.Close()
286 })
287 return sc.closeErr
portal/core/cryptoops/handshaker_test.go
+54
@@ -821,3 +821,57 @@ func BenchmarkEncryption(b *testing.B) {
821 clientSecure.Close()
822 serverSecure.Close()
823 }
824 +
825 +// TestConcurrentReadClose tests that closing the connection while reading is safe
826 +func TestConcurrentReadClose(t *testing.T) {
827 + clientCred, _ := NewCredential()
828 + serverCred, _ := NewCredential()
829 +
830 + clientConn, serverConn := pipeConn()
831 +
832 + clientHandshaker := NewHandshaker(clientCred)
833 + serverHandshaker := NewHandshaker(serverCred)
834 +
835 + var clientSecure, serverSecure *SecureConnection
836 + var wg sync.WaitGroup
837 + wg.Add(2)
838 +
839 + go func() {
840 + defer wg.Done()
841 + clientSecure, _ = clientHandshaker.ClientHandshake(clientConn, "test-alpn")
842 + }()
843 +
844 + go func() {
845 + defer wg.Done()
846 + serverSecure, _ = serverHandshaker.ServerHandshake(serverConn, []string{"test-alpn"})
847 + }()
848 +
849 + wg.Wait()
850 +
851 + // Start a goroutine that reads continuously
852 + readErrCh := make(chan error, 1)
853 + go func() {
854 + buf := make([]byte, 1024)
855 + _, err := clientSecure.Read(buf)
856 + readErrCh <- err
857 + }()
858 +
859 + // Give the reader a moment to start and block
860 + time.Sleep(10 * time.Millisecond)
861 +
862 + // Close the connection
863 + clientSecure.Close()
864 +
865 + // Check the read error
866 + select {
867 + case err := <-readErrCh:
868 + if err == nil {
869 + t.Error("Expected error from Read after Close, got nil")
870 + }
871 + // We expect either net.ErrClosed or an IO error depending on timing
872 + case <-time.After(1 * time.Second):
873 + t.Error("Read did not return after Close")
874 + }
875 +
876 + serverSecure.Close()
877 +}