diff --git a/h2transport.go b/h2transport.go index b482656..908fe07 100644 --- a/h2transport.go +++ b/h2transport.go @@ -21,15 +21,15 @@ import ( type h2Transport struct { tlsConn net.Conn h2Roundtrip http.RoundTripper - serverAddr string - localAddr net.Addr - remoteAddr net.Addr - shaper *Shaper + serverAddr string + localAddr net.Addr + remoteAddr net.Addr + shaper *Shaper - mu sync.Mutex + mu sync.Mutex activeStreams atomic.Int32 maxStreams int - closed bool + closed bool } // newH2Transport creates an HTTP/2 transport over an existing TLS connection. @@ -50,7 +50,7 @@ func newH2Transport(tlsConn net.Conn, serverAddr string, maxStreams int, shaper serverAddr: serverAddr, localAddr: tlsConn.LocalAddr(), remoteAddr: tlsConn.RemoteAddr(), - maxStreams: maxStreams, + maxStreams: maxStreams, shaper: shaper, } @@ -205,32 +205,33 @@ func (s *h2StreamRWC) Write(b []byte) (int, error) { return s.writer.Write(b) } +// abort force-closes the response body to unblock a read stalled on a stream +// the server never ends. This sends RST_STREAM and is a last resort; the +// graceful path is Close, which drains to EOF first. +func (s *h2StreamRWC) abort() { + s.reader.Close() +} + +// Close finalizes the stream. It drains any remaining response body to EOF +// before closing the reader: closing resp.Body while data is still pending +// makes net/http2 abort the stream with RST_STREAM, an abrupt reset a censor +// could fingerprint, whereas draining first lets the stream end with a normal +// END_STREAM. A server that never ends the stream is force-closed after +// drainForceTimeout. streamConn's read pump has normally already drained (or +// aborted) the body before Close runs, so the copy returns at once; the drain +// keeps Close self-safe for any direct caller. func (s *h2StreamRWC) Close() error { var closeErr error s.once.Do(func() { closeErr = s.closeWriter() - - // Do NOT close s.reader (resp.Body) synchronously — that calls - // abortStream which sends RST_STREAM, disrupting other multiplexed - // streams. Instead, drain resp.Body to EOF in the background so the - // H2 stream completes cleanly (server sends END_STREAM, writeRequest - // goroutine finishes via forgetStreamID). The drain also ensures the - // response body is closed and resources are released promptly even if - // the caller didn't read to EOF. - go func() { - // Use a timeout to prevent permanent goroutine leak if the - // server never sends END_STREAM (crash, network partition). - timer := time.AfterFunc(5*time.Second, func() { - s.reader.Close() - }) - io.Copy(io.Discard, s.reader) - timer.Stop() - s.reader.Close() - if s.tunnelCancel != nil { - s.tunnelCancel() - } - s.transport.activeStreams.Add(-1) - }() + timer := time.AfterFunc(drainForceTimeout, func() { s.reader.Close() }) + io.Copy(io.Discard, s.reader) + timer.Stop() + s.reader.Close() + if s.tunnelCancel != nil { + s.tunnelCancel() + } + s.transport.activeStreams.Add(-1) }) return closeErr } diff --git a/samizdat_test.go b/samizdat_test.go index 16dbf17..f3512a3 100644 --- a/samizdat_test.go +++ b/samizdat_test.go @@ -199,7 +199,6 @@ func TestShaperNoOp(t *testing.T) { } } - func TestRecordFragmenter(t *testing.T) { rf := NewRecordFragmenter(true) @@ -311,6 +310,224 @@ func TestStreamConnDeadline(t *testing.T) { } } +func TestStreamConnDeadlineInterruptsBlockedRead(t *testing.T) { + // The server side is never written to, so the underlying read blocks + // indefinitely — the exact failure mode the deadline must interrupt. + server, client := net.Pipe() + defer server.Close() + + sc := newStreamConn( + client, + &streamAddr{"tcp", "local"}, + &streamAddr{"tcp", "remote"}, + "remote", + nil, + ) + defer sc.Close() + + sc.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) + buf := make([]byte, 16) + start := time.Now() + _, err := sc.Read(buf) + if netErr, ok := err.(net.Error); !ok || !netErr.Timeout() { + t.Fatalf("expected timeout error, got %v", err) + } + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("read blocked %v; deadline did not interrupt it", elapsed) + } + + // The timeout is recoverable: extending the deadline lets a subsequent + // read succeed on the same conn. + go func() { + server.Write([]byte("late")) + }() + sc.SetReadDeadline(time.Now().Add(2 * time.Second)) + n, err := sc.Read(buf) + if err != nil { + t.Fatalf("read after extending deadline: %v", err) + } + if string(buf[:n]) != "late" { + t.Errorf("got %q, want %q", buf[:n], "late") + } +} + +func TestStreamConnReadDeliversDataSplitAcrossBuffers(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + + sc := newStreamConn( + client, + &streamAddr{"tcp", "local"}, + &streamAddr{"tcp", "remote"}, + "remote", + nil, + ) + defer sc.Close() + + go func() { + server.Write([]byte("abcdef")) + }() + + // Read with a buffer smaller than the delivered chunk; the remainder must + // survive to the next read rather than being dropped. + buf := make([]byte, 4) + n, err := sc.Read(buf) + if err != nil || string(buf[:n]) != "abcd" { + t.Fatalf("first read: got %q err %v", buf[:n], err) + } + n, err = sc.Read(buf) + if err != nil || string(buf[:n]) != "ef" { + t.Fatalf("second read: got %q err %v", buf[:n], err) + } +} + +func TestStreamConnReadDeliversFinalDataThenError(t *testing.T) { + // A reader that returns data and io.EOF together on its final read; the + // pump must not block on a second send, and reads past EOF must keep + // returning the error rather than hanging. + sc := newStreamConn( + &dataThenEOF{data: []byte("bye")}, + &streamAddr{"tcp", "local"}, + &streamAddr{"tcp", "remote"}, + "remote", + nil, + ) + defer sc.Close() + + buf := make([]byte, 16) + n, err := sc.Read(buf) + if err != nil || string(buf[:n]) != "bye" { + t.Fatalf("first read: got %q err %v, want \"bye\" nil", buf[:n], err) + } + + // The pump has exited; these must return EOF promptly, not block. + for i := 0; i < 3; i++ { + done := make(chan struct{}) + go func() { + _, err = sc.Read(buf) + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("read past EOF blocked") + } + if err != io.EOF { + t.Fatalf("read past EOF: got %v, want io.EOF", err) + } + } +} + +// dataThenEOF returns its data and io.EOF on the first read, then io.EOF. +type dataThenEOF struct { + data []byte + done bool +} + +func (d *dataThenEOF) Read(p []byte) (int, error) { + if d.done { + return 0, io.EOF + } + d.done = true + return copy(p, d.data), io.EOF +} +func (d *dataThenEOF) Write(p []byte) (int, error) { return len(p), nil } +func (d *dataThenEOF) Close() error { return nil } + +// blockingReadCloser has no CloseWrite; its Read blocks until Close. +type blockingReadCloser struct { + closeCh chan struct{} + closeOnce sync.Once + closed bool +} + +func (b *blockingReadCloser) Read(p []byte) (int, error) { + <-b.closeCh + return 0, io.EOF +} +func (b *blockingReadCloser) Write(p []byte) (int, error) { return len(p), nil } +func (b *blockingReadCloser) Close() error { + b.closeOnce.Do(func() { + b.closed = true + close(b.closeCh) + }) + return nil +} + +func TestStreamConnZeroLengthReadReturnsImmediately(t *testing.T) { + // A silent stream: nothing is ever written. A zero-length read must return + // (0, nil) at once rather than blocking on the pump. + _, client := net.Pipe() + sc := newStreamConn(client, &streamAddr{"tcp", "l"}, &streamAddr{"tcp", "r"}, "r", nil) + defer sc.Close() + + done := make(chan struct{}) + var n int + var err error + go func() { + n, err = sc.Read(nil) + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("zero-length Read blocked") + } + if n != 0 || err != nil { + t.Fatalf("zero-length Read = (%d, %v), want (0, nil)", n, err) + } +} + +func TestStreamConnReadAfterCloseFails(t *testing.T) { + sc := newStreamConn( + &blockingReadCloser{closeCh: make(chan struct{})}, + &streamAddr{"tcp", "l"}, &streamAddr{"tcp", "r"}, "r", nil, + ) + sc.Close() + + done := make(chan struct{}) + var err error + go func() { + _, err = sc.Read(make([]byte, 8)) + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Read after Close blocked") + } + if err == nil { + t.Fatal("Read after Close returned nil error, want failure") + } +} + +func TestStreamConnWriteAfterCloseFails(t *testing.T) { + sc := newStreamConn( + &blockingReadCloser{closeCh: make(chan struct{})}, + &streamAddr{"tcp", "l"}, &streamAddr{"tcp", "r"}, "r", nil, + ) + sc.Close() + if _, err := sc.Write([]byte("data")); err == nil { + t.Fatal("Write after Close returned nil error, want failure") + } +} + +func TestStreamConnCloseWithoutHalfCloseClosesImmediately(t *testing.T) { + // An rwc with no CloseWrite can't elicit a clean remote EOF, so Close must + // close it right away instead of waiting out drainForceTimeout. + b := &blockingReadCloser{closeCh: make(chan struct{})} + sc := newStreamConn(b, &streamAddr{"tcp", "l"}, &streamAddr{"tcp", "r"}, "r", nil) + + start := time.Now() + sc.Close() + if !b.closed { + t.Fatal("Close did not close rwc when half-close is unavailable") + } + if elapsed := time.Since(start); elapsed >= drainForceTimeout { + t.Fatalf("Close took %v; expected an immediate close", elapsed) + } +} + func TestStreamConnCloseWrite(t *testing.T) { // Use TCP connections instead of net.Pipe() because net.Pipe doesn't // support half-close (CloseWrite). TCP connections do. @@ -414,6 +631,50 @@ func TestStreamConnCloseWriteNoSupport(t *testing.T) { } } +// trackReadCloser records whether it was drained to EOF and closed. +type trackReadCloser struct { + r *bytes.Reader + readAll bool + closed bool +} + +func (t *trackReadCloser) Read(p []byte) (int, error) { + n, err := t.r.Read(p) + if err == io.EOF { + t.readAll = true + } + return n, err +} +func (t *trackReadCloser) Close() error { t.closed = true; return nil } + +type nopWriteCloser struct{} + +func (nopWriteCloser) Write(p []byte) (int, error) { return len(p), nil } +func (nopWriteCloser) Close() error { return nil } + +func TestH2StreamRWCCloseDrainsBeforeClosing(t *testing.T) { + // Close must drain the response body to EOF before closing it; closing + // while data is pending is what makes net/http2 send a fingerprintable RST. + trc := &trackReadCloser{r: bytes.NewReader([]byte("leftover data"))} + rwc := &h2StreamRWC{ + reader: trc, + writer: nopWriteCloser{}, + transport: &h2Transport{maxStreams: 100}, + tunnelCancel: func() {}, + } + rwc.transport.activeStreams.Add(1) + + if err := rwc.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if !trc.readAll { + t.Error("Close did not drain the reader to EOF before closing") + } + if !trc.closed { + t.Error("Close did not close the reader") + } +} + func TestH2StreamRWCCloseWriteThenClose(t *testing.T) { pr, pw := io.Pipe() rwc := &h2StreamRWC{ @@ -453,7 +714,7 @@ func TestConnPoolBasic(t *testing.T) { return &h2Transport{ tlsConn: client, serverAddr: "test:443", - maxStreams: 100, + maxStreams: 100, localAddr: &streamAddr{"tcp", "local"}, remoteAddr: &streamAddr{"tcp", "remote"}, }, nil @@ -581,7 +842,7 @@ func TestMasqueradeDefaults(t *testing.T) { type panicResponseWriter struct{} func (pw *panicResponseWriter) Header() http.Header { return http.Header{} } -func (pw *panicResponseWriter) WriteHeader(statusCode int) {} +func (pw *panicResponseWriter) WriteHeader(statusCode int) {} func (pw *panicResponseWriter) Write(b []byte) (int, error) { panic("Write called after Handler finished") } @@ -636,7 +897,7 @@ type slowResponseWriter struct { once sync.Once } -func (sw *slowResponseWriter) Header() http.Header { return http.Header{} } +func (sw *slowResponseWriter) Header() http.Header { return http.Header{} } func (sw *slowResponseWriter) WriteHeader(statusCode int) {} func (sw *slowResponseWriter) Write(b []byte) (int, error) { sw.once.Do(func() { close(sw.started) }) diff --git a/streamconn.go b/streamconn.go index 6ae87d2..4326fda 100644 --- a/streamconn.go +++ b/streamconn.go @@ -7,9 +7,24 @@ import ( "time" ) +const ( + // readChunkSize bounds how far ahead the read pump reads from the + // underlying stream. The pump blocks on delivery until the caller consumes + // a chunk, so at most one chunk is buffered. + readChunkSize = 32 * 1024 + + // drainForceTimeout bounds how long Close waits for the server to end the + // stream cleanly before force-closing the reader. Draining to EOF rather + // than sending RST_STREAM avoids an abrupt-reset fingerprint. + drainForceTimeout = 5 * time.Second +) + // streamConn wraps an io.ReadWriteCloser (typically an HTTP/2 stream body) -// as a net.Conn. It implements all required net.Conn methods, delegating -// Read/Write to the underlying stream and supporting deadline-based timeouts. +// as a net.Conn. Because the underlying stream body reads block in a way that +// a socket deadline cannot interrupt (the body is fed by a background HTTP/2 +// frame-reader goroutine), a dedicated read pump owns all reads from rwc and +// Read selects the pump's output against the read deadline. This makes +// SetReadDeadline actually interrupt an in-flight blocked read. type streamConn struct { rwc io.ReadWriteCloser localAddr net.Addr @@ -17,36 +32,154 @@ type streamConn struct { shaper *Shaper destination string - readDeadline *deadlineTimer - writeDeadline *deadlineTimer + readDeadline *pipeDeadline + writeDeadline *pipeDeadline - mu sync.Mutex - closed bool + readResults chan readResult + done chan struct{} // closed by Close; stops Read and drains the pump + pumpDone chan struct{} // closed when the pump goroutine exits + + bufPool sync.Pool // reusable readChunkSize buffers for the pump + + // readMu serializes Read so concurrent callers (which net.Conn permits) + // don't race on readBuf/readErr. readBuf holds the tail of a chunk that + // didn't fit in the caller's buffer; readErr is the sticky terminal error + // from the pump, returned by every subsequent Read (matching io.Reader) so + // a read past EOF returns the error rather than blocking on the exited pump. + readMu sync.Mutex + readBuf []byte + readErr error + + closeOnce sync.Once + rwcCloseOnce sync.Once + mu sync.Mutex + closed bool +} + +// readResult carries a pooled buffer plus the length read and any terminal +// error. The receiver copies the data out and returns buf to the pool. buf is a +// *[]byte because sync.Pool must hold pointer-like values to avoid an +// allocation on every Put (staticcheck SA6002). +type readResult struct { + buf *[]byte + n int + err error } // newStreamConn creates a net.Conn backed by the given ReadWriteCloser. func newStreamConn(rwc io.ReadWriteCloser, localAddr, remoteAddr net.Addr, destination string, shaper *Shaper) *streamConn { - return &streamConn{ + sc := &streamConn{ rwc: rwc, localAddr: localAddr, remoteAddr: remoteAddr, destination: destination, shaper: shaper, - readDeadline: newDeadlineTimer(), - writeDeadline: newDeadlineTimer(), + readDeadline: newPipeDeadline(), + writeDeadline: newPipeDeadline(), + readResults: make(chan readResult), + done: make(chan struct{}), + pumpDone: make(chan struct{}), + } + sc.bufPool.New = func() any { b := make([]byte, readChunkSize); return &b } + go sc.readLoop() + return sc +} + +// readLoop is the sole reader of rwc. It delivers chunks to Read over an +// unbuffered channel (providing backpressure) and, after Close, keeps reading +// to drain the stream to EOF. +func (sc *streamConn) readLoop() { + defer close(sc.pumpDone) + for { + bufp := sc.bufPool.Get().(*[]byte) + n, err := sc.rwc.Read(*bufp) + if n == 0 && err == nil { + sc.bufPool.Put(bufp) + continue + } + // Deliver data and a terminal error in one result so the pump can exit + // after a single send. A separate error send could block forever if the + // caller consumes the final data but never reads again or closes. + select { + case sc.readResults <- readResult{buf: bufp, n: n, err: err}: + // Read owns bufp now and returns it to the pool. + case <-sc.done: // draining after Close: discard + sc.bufPool.Put(bufp) + } + if err != nil { + return + } } } func (sc *streamConn) Read(b []byte) (int, error) { - if err := sc.readDeadline.wait(); err != nil { - return 0, err + // Per io.Reader/net.Conn convention a zero-length read returns immediately + // and must not block on the pump, deadline, or done channels. + if len(b) == 0 { + return 0, nil + } + + sc.readMu.Lock() + defer sc.readMu.Unlock() + + // Once closed, reads stop promptly with an error rather than draining any + // buffered tail or sticky error (net.Conn semantics). + select { + case <-sc.done: + return 0, net.ErrClosed + default: + } + + // An already-expired deadline takes precedence over buffered or new data. + select { + case <-sc.readDeadline.wait(): + return 0, &timeoutError{} + default: + } + + if len(sc.readBuf) > 0 { + n := copy(b, sc.readBuf) + sc.readBuf = sc.readBuf[n:] + return n, nil + } + if sc.readErr != nil { + return 0, sc.readErr + } + + select { + case res := <-sc.readResults: + if res.err != nil { + sc.readErr = res.err + } + if res.n == 0 { + sc.bufPool.Put(res.buf) + return 0, sc.readErr + } + data := (*res.buf)[:res.n] + n := copy(b, data) + if n < res.n { + // Copy the tail into our own buffer so the pooled buffer can be + // reused immediately; aliasing it would race with the pump's next + // read. Any terminal error stays in readErr and surfaces once this + // buffered tail is drained, so the caller sees all bytes first. + sc.readBuf = append([]byte(nil), data[n:]...) + } + sc.bufPool.Put(res.buf) + return n, nil + case <-sc.readDeadline.wait(): + return 0, &timeoutError{} + case <-sc.done: + return 0, net.ErrClosed } - return sc.rwc.Read(b) } func (sc *streamConn) Write(b []byte) (int, error) { - if err := sc.writeDeadline.wait(); err != nil { - return 0, err + select { + case <-sc.done: + return 0, net.ErrClosed + case <-sc.writeDeadline.wait(): + return 0, &timeoutError{} + default: } if sc.shaper != nil { return sc.shaper.Write(sc.rwc, b) @@ -55,15 +188,59 @@ func (sc *streamConn) Write(b []byte) (int, error) { } func (sc *streamConn) Close() error { - sc.mu.Lock() - defer sc.mu.Unlock() - if sc.closed { - return nil + sc.closeOnce.Do(func() { + sc.mu.Lock() + sc.closed = true + sc.mu.Unlock() + + sc.readDeadline.set(time.Time{}) + sc.writeDeadline.set(time.Time{}) + + close(sc.done) + + cw, canHalfClose := sc.rwc.(interface{ CloseWrite() error }) + if !canHalfClose { + // Without half-close there's no way to elicit a clean remote EOF, so + // draining would just block until the force timer. Close rwc now to + // unblock the pump and release resources immediately. + sc.closeRWC() + return + } + + // Half-close so the server ends the stream; the pump then drains the + // response to EOF and the goroutine runs the graceful close once the + // pump exits. The timer only aborts the read to unblock a stalled pump; + // it must not run closeRWC's drain while the pump is still reading rwc. + cw.CloseWrite() + forceTimer := time.AfterFunc(drainForceTimeout, sc.abortRWC) + go func() { + <-sc.pumpDone + forceTimer.Stop() + sc.closeRWC() + }() + }) + return nil +} + +// readAborter can force its read side closed to unblock a stalled read, +// bypassing the drain that Close performs. +type readAborter interface { + abort() +} + +// abortRWC unblocks a pump stalled in rwc.Read. A readAborter (the H2 stream) +// closes just its reader; a plain conn is closed outright, since that is its +// only way to interrupt a blocked read. +func (sc *streamConn) abortRWC() { + if a, ok := sc.rwc.(readAborter); ok { + a.abort() + return } - sc.closed = true - sc.readDeadline.stop() - sc.writeDeadline.stop() - return sc.rwc.Close() + sc.closeRWC() +} + +func (sc *streamConn) closeRWC() { + sc.rwcCloseOnce.Do(func() { sc.rwc.Close() }) } func (sc *streamConn) LocalAddr() net.Addr { return sc.localAddr } @@ -91,8 +268,9 @@ func (sc *streamConn) SetWriteDeadline(t time.Time) error { // data after the client signals it's done writing. func (sc *streamConn) CloseWrite() error { sc.mu.Lock() - defer sc.mu.Unlock() - if sc.closed { + closed := sc.closed + sc.mu.Unlock() + if closed { return nil } if cw, ok := sc.rwc.(interface{ CloseWrite() error }); ok { @@ -101,62 +279,65 @@ func (sc *streamConn) CloseWrite() error { return nil } -// deadlineTimer supports net.Conn deadline semantics. -type deadlineTimer struct { - mu sync.Mutex - timer *time.Timer - expired bool +// pipeDeadline is a net.Conn deadline whose expiry is observable on a channel, +// so a blocked Read can select against it. It follows the implementation of +// the standard library's net.Pipe deadline. +type pipeDeadline struct { + mu sync.Mutex + timer *time.Timer + cancel chan struct{} // closed when the deadline is reached; never nil } -func newDeadlineTimer() *deadlineTimer { - return &deadlineTimer{} +func newPipeDeadline() *pipeDeadline { + return &pipeDeadline{cancel: make(chan struct{})} } -// set configures the deadline. A zero time clears the deadline. -func (dt *deadlineTimer) set(t time.Time) { - dt.mu.Lock() - defer dt.mu.Unlock() +// set arms the deadline. A zero time clears it. +func (d *pipeDeadline) set(t time.Time) { + d.mu.Lock() + defer d.mu.Unlock() - dt.expired = false - if dt.timer != nil { - dt.timer.Stop() - dt.timer = nil + if d.timer != nil && !d.timer.Stop() { + <-d.cancel // a fired timer already closed cancel; drain it } + d.timer = nil + closed := isClosedChan(d.cancel) if t.IsZero() { + if closed { + d.cancel = make(chan struct{}) + } return } - d := time.Until(t) - if d <= 0 { - dt.expired = true + if dur := time.Until(t); dur > 0 { + if closed { + d.cancel = make(chan struct{}) + } + d.timer = time.AfterFunc(dur, func() { close(d.cancel) }) return } - dt.timer = time.AfterFunc(d, func() { - dt.mu.Lock() - dt.expired = true - dt.mu.Unlock() - }) + // Deadline in the past: expire immediately. + if !closed { + close(d.cancel) + } } -func (dt *deadlineTimer) stop() { - dt.mu.Lock() - defer dt.mu.Unlock() - if dt.timer != nil { - dt.timer.Stop() - } +// wait returns a channel that is closed when the deadline is reached. +func (d *pipeDeadline) wait() <-chan struct{} { + d.mu.Lock() + defer d.mu.Unlock() + return d.cancel } -// wait returns a timeout error if the deadline has expired, nil otherwise. -func (dt *deadlineTimer) wait() error { - dt.mu.Lock() - expired := dt.expired - dt.mu.Unlock() - if expired { - return &timeoutError{} +func isClosedChan(c <-chan struct{}) bool { + select { + case <-c: + return true + default: + return false } - return nil } // timeoutError implements the net.Error interface for deadline timeouts.