puma-plus-core 0.1.0

data/go/internal/frontend/handler_test.go

@@ -0,0 +1,486 @@
+package frontend
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/evanphx/puma-plus/internal/queue"
+	"github.com/evanphx/puma-plus/internal/wire"
+)
+
+// fakeWorker speaks the worker side of the protocol in Go, so the frontend can
+// be tested end to end without spawning Ruby. It deliberately implements the
+// same obligations the real worker has, including draining the request body
+// before RESP_END.
+type fakeWorker struct {
+	t  *testing.T
+	c  net.Conn
+	r  *wire.Reader
+	w  *wire.Writer
+	fn func(env map[string]string, body []byte) (int, [][2]string, [][]byte)
+}
+
+func newFakeWorker(t *testing.T, reg *queue.Registry,
+	fn func(env map[string]string, body []byte) (int, [][2]string, [][]byte)) *fakeWorker {
+	t.Helper()
+
+	// socketpair via net.Pipe would not exercise real framing across a kernel
+	// buffer, and the concurrent body pump depends on that buffering to make
+	// progress. Use a real unix socketpair.
+	fds, err := makeSocketPair()
+	if err != nil {
+		t.Fatalf("socketpair: %v", err)
+	}
+	goSide, workerSide := fds[0], fds[1]
+
+	fw := &fakeWorker{t: t, c: workerSide, r: wire.NewReader(workerSide), w: wire.NewWriter(workerSide), fn: fn}
+	conn := queue.NewConn(goSide, 1234, 0, 0)
+	reg.Add(conn)
+
+	go fw.serve()
+	t.Cleanup(func() { workerSide.Close() })
+	return fw
+}
+
+func (fw *fakeWorker) serve() {
+	for {
+		typ, payload, err := fw.r.ReadFrame()
+		if err != nil {
+			return
+		}
+		if typ != wire.TypeRequest {
+			return
+		}
+
+		meta, rest, err := wire.DecodeRequestMeta(payload)
+		if err != nil {
+			return
+		}
+		pairs, rest, err := wire.DecodeKV(rest)
+		if err != nil {
+			return
+		}
+		env := make(map[string]string, len(pairs))
+		for _, p := range pairs {
+			env[string(p.Key)] = string(p.Value)
+		}
+
+		var body []byte
+		switch meta.BodyMode {
+		case wire.BodyInline:
+			body = append(body, rest...)
+		case wire.BodyStream:
+			// Drain to BODY_END, as the protocol requires.
+			for {
+				bt, bp, err := fw.r.ReadFrame()
+				if err != nil {
+					return
+				}
+				if bt == wire.TypeBodyEnd {
+					break
+				}
+				if bt == wire.TypeBodyChunk {
+					body = append(body, bp...)
+				}
+			}
+		}
+
+		status, headers, chunks := fw.fn(env, body)
+
+		resp := []byte{byte(status >> 8), byte(status)}
+		hp := make([]wire.Pair, 0, len(headers))
+		for _, h := range headers {
+			hp = append(hp, wire.Pair{Key: []byte(h[0]), Value: []byte(h[1])})
+		}
+		if err := fw.w.WriteFrame(wire.TypeResponse, wire.AppendKV(resp, hp)); err != nil {
+			return
+		}
+		for _, c := range chunks {
+			if err := fw.w.WriteFrame(wire.TypeRespChunk, c); err != nil {
+				return
+			}
+		}
+		end := wire.AppendRespEndMeta(nil, wire.RespEndMeta{ServiceNs: 1000})
+		if err := fw.w.WriteFrame(wire.TypeRespEnd, wire.AppendKV(end, nil)); err != nil {
+			return
+		}
+	}
+}
+
+func testHandler(t *testing.T, fn func(map[string]string, []byte) (int, [][2]string, [][]byte)) *Handler {
+	t.Helper()
+	reg := queue.NewRegistry(8)
+	newFakeWorker(t, reg, fn)
+	// Give the registration a moment; Add is synchronous but the goroutine that
+	// serves is not.
+	time.Sleep(10 * time.Millisecond)
+	return NewHandler(reg, Options{
+		Logger:         log.New(io.Discard, "", 0),
+		AcquireTimeout: 5 * time.Second,
+	})
+}
+
+// echoBody returns the SHA-256 of whatever body arrived, plus its length. Any
+// byte lost in transit changes the digest.
+func echoBody(_ map[string]string, body []byte) (int, [][2]string, [][]byte) {
+	sum := sha256.Sum256(body)
+	return 200, [][2]string{{"content-type", "text/plain"}},
+		[][]byte{[]byte(fmt.Sprintf("%d %s", len(body), hex.EncodeToString(sum[:])))}
+}
+
+// TestRequestBodyIntegrity is the regression test for a silent one-byte data
+// loss: readBody probed one byte past BodyBufferMax to detect overflow, then
+// returned only the first BodyBufferMax bytes, discarding a byte it had already
+// consumed from the request. Every streamed body arrived exactly one byte short
+// and the only symptom was a wrong digest.
+//
+// The sizes bracket both boundaries that matter: wire.InlineBodyMax (INLINE vs
+// STREAM carriage) and BodyBufferMax (buffered vs live-streamed).
+func TestRequestBodyIntegrity(t *testing.T) {
+	h := testHandler(t, echoBody)
+
+	sizes := []int{
+		0, 1, 1024,
+		wire.InlineBodyMax - 1, wire.InlineBodyMax, wire.InlineBodyMax + 1,
+		wire.ChunkTarget, wire.ChunkTarget + 1,
+		1<<20 - 1, 1 << 20, 1<<20 + 1,
+		3 << 20,
+	}
+
+	for _, size := range sizes {
+		t.Run(fmt.Sprintf("size=%d", size), func(t *testing.T) {
+			payload := make([]byte, size)
+			for i := range payload {
+				payload[i] = byte(i * 31)
+			}
+			want := sha256.Sum256(payload)
+
+			req := httptest.NewRequest(http.MethodPost, "/upload", bytes.NewReader(payload))
+			rec := httptest.NewRecorder()
+			h.ServeHTTP(rec, req)
+
+			if rec.Code != 200 {
+				t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
+			}
+			gotLen, gotSum := 0, ""
+			if _, err := fmt.Sscanf(rec.Body.String(), "%d %s", &gotLen, &gotSum); err != nil {
+				t.Fatalf("parsing worker reply %q: %v", rec.Body.String(), err)
+			}
+			if gotLen != size {
+				t.Errorf("worker received %d bytes, sent %d", gotLen, size)
+			}
+			if gotSum != hex.EncodeToString(want[:]) {
+				t.Errorf("digest mismatch: body was corrupted in transit")
+			}
+		})
+	}
+}
+
+// TestChunkedRequestBody covers the no-Content-Length path, which always
+// streams because the size is unknown until EOF.
+func TestChunkedRequestBody(t *testing.T) {
+	h := testHandler(t, echoBody)
+
+	payload := bytes.Repeat([]byte("chunky"), 50_000) // 300 KB
+	want := sha256.Sum256(payload)
+
+	req := httptest.NewRequest(http.MethodPost, "/upload", bytes.NewReader(payload))
+	req.ContentLength = -1 // as net/http presents a chunked request
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+
+	gotLen, gotSum := 0, ""
+	fmt.Sscanf(rec.Body.String(), "%d %s", &gotLen, &gotSum)
+	if gotLen != len(payload) || gotSum != hex.EncodeToString(want[:]) {
+		t.Errorf("chunked body corrupted: got %d bytes, digest %s", gotLen, gotSum)
+	}
+}
+
+// TestBodyModeSelection pins the carriage decision at both boundaries.
+func TestBodyModeSelection(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		size int
+		more bool
+		want uint8
+	}{
+		{"empty", 0, false, wire.BodyNone},
+		{"one byte", 1, false, wire.BodyInline},
+		{"at inline max", wire.InlineBodyMax, false, wire.BodyInline},
+		{"past inline max", wire.InlineBodyMax + 1, false, wire.BodyStream},
+		{"more to come", 16, true, wire.BodyStream},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := bodyMode(make([]byte, tc.size), tc.more); got != tc.want {
+				t.Errorf("bodyMode(%d bytes, more=%v) = %d, want %d",
+					tc.size, tc.more, got, tc.want)
+			}
+		})
+	}
+}
+
+// TestWorkerMustDrainBody checks the protocol invariant: a worker that ignores
+// the request body but still drains it leaves the connection reusable, so the
+// next request on that connection succeeds. Without the drain the connection
+// desynchronises and the following request reads a body chunk where it expects
+// a REQUEST frame.
+func TestWorkerMustDrainBody(t *testing.T) {
+	h := testHandler(t, func(_ map[string]string, body []byte) (int, [][2]string, [][]byte) {
+		// The fake worker drains in serve() before calling us; we ignore the
+		// body entirely, as a 413 handler would.
+		return 413, [][2]string{{"content-type", "text/plain"}}, [][]byte{[]byte("too large")}
+	})
+
+	big := bytes.NewReader(make([]byte, 3<<20))
+	req := httptest.NewRequest(http.MethodPost, "/upload", big)
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+	if rec.Code != 413 {
+		t.Fatalf("first request: status %d", rec.Code)
+	}
+
+	// Same worker connection, reused. This is what breaks if the drain is missing.
+	for i := 0; i < 3; i++ {
+		req = httptest.NewRequest(http.MethodGet, "/", nil)
+		rec = httptest.NewRecorder()
+		h.ServeHTTP(rec, req)
+		if rec.Code != 413 {
+			t.Fatalf("follow-up request %d on the reused connection: status %d (%s)",
+				i, rec.Code, rec.Body.String())
+		}
+	}
+}
+
+// TestResponseHeadersAndStatus covers repeated header keys, which is how Rack 3
+// array-valued headers and multiple Set-Cookie lines survive the round trip.
+func TestResponseHeadersAndStatus(t *testing.T) {
+	h := testHandler(t, func(_ map[string]string, _ []byte) (int, [][2]string, [][]byte) {
+		return 201, [][2]string{
+			{"content-type", "text/plain"},
+			{"set-cookie", "a=1"},
+			{"set-cookie", "b=2"},
+		}, [][]byte{[]byte("created")}
+	})
+
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	if rec.Code != 201 {
+		t.Errorf("status = %d, want 201", rec.Code)
+	}
+	cookies := rec.Result().Header.Values("Set-Cookie")
+	if len(cookies) != 2 || cookies[0] != "a=1" || cookies[1] != "b=2" {
+		t.Errorf("Set-Cookie = %v, want two separate values in order", cookies)
+	}
+	if rec.Body.String() != "created" {
+		t.Errorf("body = %q", rec.Body.String())
+	}
+}
+
+// TestEnvConstruction checks the CGI-ification Go performs on Ruby's behalf.
+func TestEnvConstruction(t *testing.T) {
+	var got map[string]string
+	h := testHandler(t, func(env map[string]string, _ []byte) (int, [][2]string, [][]byte) {
+		got = env
+		return 200, nil, nil
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/path/here?x=1&y=2", nil)
+	req.Host = "example.com:8080"
+	req.Header.Set("User-Agent", "test-agent")
+	req.Header.Set("X-Custom-Thing", "value")
+	req.Header.Add("X-Multi", "one")
+	req.Header.Add("X-Multi", "two")
+	req.Header.Set("Content-Type", "application/json")
+	req.RemoteAddr = "10.0.0.5:54321"
+
+	h.ServeHTTP(httptest.NewRecorder(), req)
+
+	for _, tc := range []struct{ key, want string }{
+		{"REQUEST_METHOD", "GET"},
+		{"PATH_INFO", "/path/here"},
+		{"QUERY_STRING", "x=1&y=2"},
+		{"SCRIPT_NAME", ""},
+		{"SERVER_NAME", "example.com"},
+		{"SERVER_PORT", "8080"},
+		{"SERVER_PROTOCOL", "HTTP/1.1"},
+		{"REMOTE_ADDR", "10.0.0.5"},
+		{"rack.url_scheme", "http"},
+		{"HTTP_USER_AGENT", "test-agent"},
+		{"HTTP_X_CUSTOM_THING", "value"},
+		{"CONTENT_TYPE", "application/json"},
+		// Repeated headers join with ", ", matching puma.
+		{"HTTP_X_MULTI", "one, two"},
+	} {
+		if v, ok := got[tc.key]; !ok || v != tc.want {
+			t.Errorf("env[%q] = %q (present=%v), want %q", tc.key, v, ok, tc.want)
+		}
+	}
+
+	// Content-Type is a CGI variable, not an HTTP_ one.
+	if _, ok := got["HTTP_CONTENT_TYPE"]; ok {
+		t.Error("Content-Type must not also appear as HTTP_CONTENT_TYPE")
+	}
+}
+
+// TestUnderscoreHeadersDropped documents where we deliberately diverge from
+// puma. puma's parser maps '_' to ',' then un-mangles later; Go sees pristine
+// names and simply drops them, because after CGI-ification an underscore header
+// is indistinguishable from its dashed sibling and that ambiguity is a request
+// smuggling vector.
+func TestUnderscoreHeadersDropped(t *testing.T) {
+	var got map[string]string
+	h := testHandler(t, func(env map[string]string, _ []byte) (int, [][2]string, [][]byte) {
+		got = env
+		return 200, nil, nil
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header["X_Spoofed_Thing"] = []string{"evil"}
+	req.Header.Set("X-Spoofed-Thing", "legitimate")
+	h.ServeHTTP(httptest.NewRecorder(), req)
+
+	if got["HTTP_X_SPOOFED_THING"] != "legitimate" {
+		t.Errorf("dashed header should win: got %q", got["HTTP_X_SPOOFED_THING"])
+	}
+}
+
+// TestSplitHostPort pins SERVER_NAME and SERVER_PORT against what puma produces
+// for the same Host header.
+//
+// The IPv6 cases changed: this used to strip the brackets, giving SERVER_NAME
+// "::1". Puma keeps them (puma/lib/puma/client_env.rb:39 takes host[0, colon+1],
+// which includes the ']'), and keeping them is also the only form that can be
+// reassembled -- "::1" with SERVER_PORT 443 joins to "::1:443", which is
+// ambiguous, while "[::1]:443" is not. Rack::Request#host strips the brackets
+// itself when an app wants the bare address.
+func TestSplitHostPort(t *testing.T) {
+	for _, tc := range []struct{ in, host, port string }{
+		{"example.com:8080", "example.com", "8080"},
+		{"example.com", "example.com", "80"},
+		{"[::1]:9292", "[::1]", "9292"},
+		{"[::1]", "[::1]", "80"},
+		{"[2001:db8::1]:443", "[2001:db8::1]", "443"},
+		{"", "localhost", "80"},
+	} {
+		host, port := splitHostPort(tc.in, "80")
+		if host != tc.host || port != tc.port {
+			t.Errorf("splitHostPort(%q) = (%q, %q), want (%q, %q)",
+				tc.in, host, port, tc.host, tc.port)
+		}
+	}
+}
+
+// TestShedWhenNoCapacity checks that a request with no worker available is shed
+// with a 503 and a Retry-After rather than hanging.
+func TestShedWhenNoCapacity(t *testing.T) {
+	reg := queue.NewRegistry(4) // no workers ever connect
+	h := NewHandler(reg, Options{
+		Logger:         log.New(io.Discard, "", 0),
+		AcquireTimeout: 50 * time.Millisecond,
+	})
+
+	rec := httptest.NewRecorder()
+	start := time.Now()
+	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/", nil))
+	elapsed := time.Since(start)
+
+	if rec.Code != http.StatusServiceUnavailable {
+		t.Errorf("status = %d, want 503", rec.Code)
+	}
+	if got := rec.Header().Get("Retry-After"); got == "" {
+		t.Error("503 should carry Retry-After")
+	}
+	if elapsed > time.Second {
+		t.Errorf("shed took %v; the acquire timeout should have bounded it", elapsed)
+	}
+	if !strings.Contains(rec.Body.String(), "busy") {
+		t.Errorf("body = %q", rec.Body.String())
+	}
+}
+
+// TestRetriesWhenWorkerDiesBeforeDelivery is the regression test for the 502s
+// found by the first real benchmark run.
+//
+// A worker whose process exits between being acquired and being written to
+// produces a broken pipe. Nothing was delivered, so the app never ran and the
+// request is safe to retry -- returning 502 there was throwing away requests we
+// could have served. The benchmark saw 100 such 502s.
+func TestRetriesWhenWorkerDiesBeforeDelivery(t *testing.T) {
+	reg := queue.NewRegistry(8)
+
+	// A worker that is already dead: its socket is closed, so any write fails.
+	fds, err := makeSocketPair()
+	if err != nil {
+		t.Fatal(err)
+	}
+	fds[1].Close() // kill the worker side
+	reg.Add(queue.NewConn(fds[0], 999, 0, 0))
+
+	// A healthy worker that can actually serve.
+	newFakeWorker(t, reg, func(_ map[string]string, _ []byte) (int, [][2]string, [][]byte) {
+		return 200, [][2]string{{"content-type", "text/plain"}}, [][]byte{[]byte("served")}
+	})
+	time.Sleep(20 * time.Millisecond)
+
+	h := NewHandler(reg, Options{
+		Logger:         log.New(io.Discard, "", 0),
+		AcquireTimeout: 2 * time.Second,
+	})
+
+	// Several requests, so the dead connection is hit and retried past.
+	for i := 0; i < 5; i++ {
+		rec := httptest.NewRecorder()
+		h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/", nil))
+		if rec.Code != 200 {
+			t.Fatalf("request %d: status %d (%s); a worker dying before delivery "+
+				"must be retried, not surfaced to the client",
+				i, rec.Code, rec.Body.String())
+		}
+		if rec.Body.String() != "served" {
+			t.Errorf("request %d: body %q", i, rec.Body.String())
+		}
+	}
+}
+
+// TestNoRetryAfterDelivery: once the request has been handed to a worker we
+// cannot know whether the app ran, so a death there must NOT be retried. A
+// non-idempotent request executing twice is far worse than a 502.
+func TestNoRetryAfterDelivery(t *testing.T) {
+	reg := queue.NewRegistry(8)
+
+	// A worker that reads the request and then dies without responding.
+	fds, err := makeSocketPair()
+	if err != nil {
+		t.Fatal(err)
+	}
+	reg.Add(queue.NewConn(fds[0], 998, 0, 0))
+	go func() {
+		r := wire.NewReader(fds[1])
+		_, _, _ = r.ReadFrame() // take delivery, then vanish
+		fds[1].Close()
+	}()
+	time.Sleep(20 * time.Millisecond)
+
+	h := NewHandler(reg, Options{
+		Logger:         log.New(io.Discard, "", 0),
+		AcquireTimeout: 500 * time.Millisecond,
+	})
+
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/charge", nil))
+	if rec.Code != http.StatusBadGateway {
+		t.Errorf("status = %d, want 502; a request that may already have executed "+
+			"must not be silently retried", rec.Code)
+	}
+}

data/go/internal/frontend/hijack.go

@@ -0,0 +1,108 @@
+package frontend
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"sync"
+
+	"github.com/evanphx/puma-plus/internal/queue"
+)
+
+// ErrHijackUnsupported is returned when the transport cannot be hijacked.
+var ErrHijackUnsupported = errors.New("frontend: this transport cannot be hijacked")
+
+// spliceHijack takes over the client connection and joins it to the worker
+// connection as a raw bidirectional byte tunnel.
+//
+// After the worker sends HIJACK, the protocol is over: no more frames, no
+// response, no return to the pool. Ruby owns the conversation, including writing
+// the HTTP/1.1 101 Switching Protocols response itself. This is what Rack full
+// hijack means -- the application takes the socket -- and it is how WebSockets
+// work here without puma-plus needing to know anything about WebSockets.
+//
+// The worker thread is retired rather than released, because it will spend the
+// rest of its life on this tunnel. That is the honest accounting: a hijacked
+// connection is capacity permanently removed from the pool, and the controller
+// must see it disappear rather than believe it is idle.
+func (h *Handler) spliceHijack(w http.ResponseWriter, conn *queue.Conn) error {
+	hj, ok := w.(http.Hijacker)
+	if !ok {
+		// h2 and h3 multiplex many streams over one connection, so there is no
+		// single socket to hand over -- taking the connection would steal every
+		// other in-flight stream on it. Rack full hijack is an HTTP/1.1 concept
+		// and we refuse it honestly rather than corrupting the connection.
+		return fmt.Errorf("%w: full hijack requires HTTP/1.1 (this is %s)",
+			ErrHijackUnsupported, currentProto(w))
+	}
+
+	clientConn, buf, err := hj.Hijack()
+	if err != nil {
+		return fmt.Errorf("hijacking the client connection: %w", err)
+	}
+
+	// The worker leaves the pool permanently.
+	h.reg.Retire(conn)
+
+	workerConn := conn.Net()
+
+	// Anything the HTTP server already read from the client but has not consumed
+	// -- pipelined bytes, the start of a WebSocket frame -- is sitting in this
+	// buffer. Dropping it would silently corrupt the tunnel's first message.
+	if n := buf.Reader.Buffered(); n > 0 {
+		pending := make([]byte, n)
+		if _, err := io.ReadFull(buf.Reader, pending); err == nil {
+			if _, err := workerConn.Write(pending); err != nil {
+				clientConn.Close()
+				workerConn.Close()
+				return fmt.Errorf("forwarding buffered client bytes: %w", err)
+			}
+		}
+	}
+
+	go spliceBoth(clientConn, workerConn)
+	return nil
+}
+
+// spliceBoth copies in both directions until either side finishes, then closes
+// both.
+//
+// Half-close is deliberately not honored: a WebSocket peer that closes its write
+// side has ended the conversation, and keeping the other direction open would
+// leak a goroutine and a worker thread for as long as the other side stays
+// silent. Closing both on the first EOF bounds that.
+func spliceBoth(a, b net.Conn) {
+	var once sync.Once
+	closeBoth := func() {
+		once.Do(func() {
+			a.Close()
+			b.Close()
+		})
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		_, _ = io.Copy(a, b)
+		closeBoth()
+	}()
+	go func() {
+		defer wg.Done()
+		_, _ = io.Copy(b, a)
+		closeBoth()
+	}()
+	wg.Wait()
+	closeBoth()
+}
+
+func currentProto(w http.ResponseWriter) string {
+	// http.ResponseWriter does not expose the protocol, but only h1 implements
+	// Hijacker, so reaching here means h2 or h3.
+	if _, ok := w.(http.Flusher); ok {
+		return "HTTP/2 or HTTP/3"
+	}
+	return "an unknown transport"
+}