add Hijacker-like functionality

This commit is contained in:
Patricio Whittingslow
2026-07-25 21:11:25 -03:00
parent 5b83b18da0
commit 962f6864be
3 changed files with 161 additions and 30 deletions
+51
View File
@@ -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)
}
}
}
+77 -30
View File
@@ -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
+33
View File
@@ -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()
}