From d53a7db37c60914564f2596a799268e77f61433a Mon Sep 17 00:00:00 2001 From: Patricio Whittingslow Date: Sat, 25 Jul 2026 21:11:25 -0300 Subject: [PATCH] add Hijacker-like functionality --- http/httphi/exchange_test.go | 51 +++++++++++++++++ http/httphi/httphi.go | 107 +++++++++++++++++++++++++---------- http/httphi/router_test.go | 33 +++++++++++ 3 files changed, 161 insertions(+), 30 deletions(-) diff --git a/http/httphi/exchange_test.go b/http/httphi/exchange_test.go index 2368e25..58afac5 100644 --- a/http/httphi/exchange_test.go +++ b/http/httphi/exchange_test.go @@ -279,3 +279,54 @@ func TestExchangeReadBody(t *testing.T) { t.Errorf("want body %q, got %q", body, got) } } + +// SetHeader must budget every byte it writes: colon, CRLF, and the CRLF that +// FlushHeader appends after the last field. Buffers that fit all but the last +// byte must be refused, never overrun. +func TestExchangeSetHeaderExactFit(t *testing.T) { + const key, value = "K", "V" + const field = len(key) + len(value) + len(":\r\n") + for _, bufLen := range []int{field + 2, field + 1, field} { + conn := newConn("") + exch := new(Exchange) + exch.Configure(make([]byte, bufLen), bufLen, false) + if !exch.Acquire(conn) { + t.Fatal("fresh exchange failed to acquire connection") + } + set := exch.SetHeader(key, value) + exch.WriteHeader(200) + + want := "HTTP/1.1 200 OK\r\n" + if set { + want += key + ":" + value + "\r\n" + } + want += "\r\n" + if got := conn.ViewWritten(); got != want { + t.Errorf("buffer %d: want %q, got %q", bufLen, want, got) + } + if wantSet := bufLen >= field+2; set != wantSet { + t.Errorf("buffer %d: want SetHeader=%v, got %v", bufLen, wantSet, set) + } + } +} + +// Handle never closes the connection, on any outcome: the caller owns it so +// that error policy and connection reuse stay the caller's decision. +func TestHandleLeavesConnOpen(t *testing.T) { + var sm sliceMux + sm.Handle("GET /", staticPage(t, "ok")) + for _, request := range []string{ + "GET / HTTP/1.1\r\nHost: h\r\n\r\n", // Served. + "GET /nowhere HTTP/1.1\r\nHost: h\r\n\r\n", // 404. + "GET /\r\nHost: h\r\n\r\n", // Rejected: no HTTP version. + "GET / HTTP/1.1\r\nBadFieldNoColon\r\n\r\n", // Rejected: parse error. + } { + conn := newConn(request) + conn.Hangup() + exch := newExchange(t, conn, 1024, false) + Handle(exch, &sm, nopBackoff) + if conn.IsClosed() { + t.Errorf("Handle closed the connection for %q", request) + } + } +} diff --git a/http/httphi/httphi.go b/http/httphi/httphi.go index dd24305..7a01104 100644 --- a/http/httphi/httphi.go +++ b/http/httphi/httphi.go @@ -134,6 +134,7 @@ func (r *Router) Configure(cfg RouterConfig) error { // Handle is a extremely low-level HTTP handling method used internally in [Router]. // Requires exchange to be acquired and configured. Will panic if any argument is nil. +// Handle does not close the connection on any outcome: the caller owns it. func Handle(exch *Exchange, mux Mux, backoff lneto.BackoffStrategy) error { reqhdr := &exch.reqHdr reqhdr.Reset(nil) @@ -167,7 +168,6 @@ func Handle(exch *Exchange, mux Mux, backoff lneto.BackoffStrategy) error { // Request line with no HTTP version is a HTTP/0.9 simple-request, which // httpraw tolerates. It is not a valid HTTP/1.1 request-line, RFC 9112 3. exch.WriteHeader(int(StatusBadRequest)) - exch.rw.Close() return errNoRequestProto } // Mux URI. @@ -181,22 +181,26 @@ func Handle(exch *Exchange, mux Mux, backoff lneto.BackoffStrategy) error { exch.WriteHeader(404) } // TODO write response from exchange here. - exch.rw.Close() return nil } func (r *Router) Handle(conn conn) error { - exch := r.getExch(conn) + // Exchange acquisition and the configuration it is served with must be read + // under the same lock: [Router.Configure] may run concurrently. + r.mu.Lock() + exch := r.getExchLocked(conn) + numGoro, backoff, mux, queue := r.numGoro, r.backoff, r.mux, r.pendingConns + r.mu.Unlock() if exch == nil { return lneto.ErrExhausted } - if r.numGoro == 0 { - go r.goroHandle(exch, r.backoff, r.mux) + if numGoro == 0 { + go r.goroHandle(exch, backoff, mux) return nil } select { - case r.pendingConns <- job{exch: exch}: + case queue <- job{exch: exch}: return nil default: // pendingConns cannot store another Conn, we drop and return error. @@ -247,9 +251,8 @@ func (r *Router) freeExch(exch *Exchange) { r.mu.Unlock() } -func (r *Router) getExch(conn conn) (exch *Exchange) { - r.mu.Lock() - defer r.mu.Unlock() +// getExchLocked returns an exchange acquired on conn. Requires r.mu held. +func (r *Router) getExchLocked(conn conn) (exch *Exchange) { if r.freeList != nil { if r.freeList.Acquire(conn) { exch = r.freeList @@ -295,7 +298,9 @@ type Exchange struct { respHeaderOff uint16 respHeaderLen uint16 reqHdr httpraw.Header - rw conn + + hijacked bool + rw conn respRemains int headerWritten bool @@ -304,6 +309,35 @@ type Exchange struct { readErr error } +// HijackRaw is a low-level implementation of http.Hijacker interface. +// A Hijack method is not exposed due to heap allocation implications and correctness concerns. +// Below is what an actual implementation may look like: +// +// func (exch *Exchange) Hijack() (net.Conn, *bufio.ReadWriter, error) { +// conn, ok := exch.rw.(net.Conn) +// if !ok { +// return nil, nil, errors.New("net.Conn not implemented") +// } +// _, data, err := exch.HijackRaw(nil) +// if err != nil { +// return nil, nil, err +// } +// var rd *bufio.ReadWriter +// if len(data) > 0 { +// rd = &bufio.ReadWriter{Reader: bufio.NewReader(bytes.NewReader(data))} +// } +// return conn, rd, nil +// } +func (exch *Exchange) HijackRaw(dstBody []byte) (conn, []byte, error) { + data, err := exch.remainingSurplusBody() + if err != nil { + return nil, nil, err + } + exch.hijacked = true + dstBody = append(dstBody, data...) + return exch.rw, dstBody, nil +} + func (exch *Exchange) Configure(rawbuf []byte, requestLim int, normalizeKeys bool) { respSize := len(rawbuf) - requestLim if respSize < 0 { @@ -331,6 +365,9 @@ func (exch *Exchange) Acquire(conn conn) bool { } func (exch *Exchange) Release() { + if !exch.hijacked { + exch.rw.Close() + } exch.rw = nil exch.used.Store(false) } @@ -338,7 +375,9 @@ func (exch *Exchange) Release() { func (exch *Exchange) SetHeader(key, value string) (enoughMemory bool) { off := int(exch.respHeaderOff) + int(exch.respHeaderLen) free := len(exch.rawbuf) - off - if len(key)+len(value)+2 > free { + // Field costs key+':'+value+CRLF, plus the CRLF [Exchange.FlushHeader] + // appends past the last field to close the header block. + if len(key)+len(value)+len(":\r\n")+len("\r\n") > free { return false } n := copy(exch.rawbuf[off:], key) @@ -410,6 +449,33 @@ func (exch *Exchange) Write(buf []byte) (int, error) { return exch.rw.Write(buf) } +func (exch *Exchange) ReadBody(dst []byte) (n int, _ error) { + if exch.respRemains > 0 { + toRead, err := exch.remainingSurplusBody() + if err != nil { + return 0, err + } + n = copy(dst, toRead) + exch.respRemains -= n + if len(dst) == n { + return n, nil + } + dst = dst[n:] + } + nr, err := exch.rw.Read(dst) + return nr + n, err +} + +func (exch *Exchange) remainingSurplusBody() ([]byte, error) { + _, err := exch.reqHdr.Body() + if err != nil { + return nil, err // Returns mangled buffer error if request header has been misused. + } + surplus := exch.rawbuf[exch.reqHdr.BufferParsed():exch.reqHdr.BufferReceived()] + toRead := surplus[len(surplus)-exch.respRemains:] + return toRead, nil +} + func (exch *Exchange) RequestHeaderRaw() *httpraw.Header { return &exch.reqHdr } @@ -436,25 +502,6 @@ func (exch *Exchange) RequestConnectionClose() bool { return exch.RequestHeaderRaw().ConnectionClose() } -func (exch *Exchange) ReadBody(dst []byte) (n int, _ error) { - if exch.respRemains > 0 { - _, err := exch.reqHdr.Body() - if err != nil { - return 0, err // Returns mangled buffer error if request header has been misused. - } - surplus := exch.rawbuf[exch.reqHdr.BufferParsed():exch.reqHdr.BufferReceived()] - toRead := surplus[len(surplus)-exch.respRemains:] - n = copy(dst, toRead) - exch.respRemains -= n - if len(dst) == n { - return n, nil - } - dst = dst[n:] - } - nr, err := exch.rw.Read(dst) - return nr + n, err -} - type HandlerFunc func(ex *Exchange) type Method uint8 diff --git a/http/httphi/router_test.go b/http/httphi/router_test.go index 6a5ca99..7b9940c 100644 --- a/http/httphi/router_test.go +++ b/http/httphi/router_test.go @@ -101,6 +101,13 @@ func (r *rwconn) AddReadable(b []byte) { defer r.mu.Unlock() r.readable.Write(b) } +// IsClosed reports whether the connection was closed by its handler. +func (r *rwconn) IsClosed() bool { + r.mu.Lock() + defer r.mu.Unlock() + return r.closed +} + func (r *rwconn) ViewWritten() string { r.mu.Lock() defer r.mu.Unlock() @@ -296,3 +303,29 @@ func staticPage(t *testing.T, page string) HandlerFunc { } } } + +// Configure writes the fields Handle reads; concurrent use must not race. +func TestRouterConfigureHandleRace(t *testing.T) { + const bufferSize = 1024 + var ( + sm sliceMux + router Router + ) + sm.Handle("GET /", staticPage(t, "ok")) + configSynchronousRouter(t, &router, bufferSize, &sm) + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + configSynchronousRouter(t, &router, bufferSize, &sm) + }() + go func() { + defer wg.Done() + conn := newConn("GET / HTTP/1.1\r\nHost: h\r\n\r\n") + if err := router.Handle(conn); err != nil { + t.Error(err) + } + }() + wg.Wait() +}