bastion 0.56.0

client-go/lane_bounds_test.go

@@ -0,0 +1,481 @@
+//go:build !windows
+
+package main
+
+// What a paired host sends down the harness, session, root-exec and mirror
+// lanes, and what a coder run sends its egress proxy, is read at the size this
+// machine chose (lane_bounds.go). Every test here drives only code that main
+// already has, so each one also runs against main, where it fails.
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// hostSendsAJob is a host that writes one job frame while it reads the box's
+// replies, and reports whether its write completed and what came back.
+type hostResult struct {
+	werr   error
+	frames map[byte][]byte
+}
+
+func hostSendsAJob(frame []byte, out chan<- hostResult) func(c *websocket.Conn) {
+	return hostSendsAJobAs(func(c *websocket.Conn) error {
+		return c.WriteMessage(websocket.BinaryMessage, frame)
+	}, out)
+}
+
+// writeInPieces writes msg as ONE message carried by many small frames: the
+// writer flushes a continuation frame each time its 4 KiB buffer fills
+// (gorilla's default write buffer), so a bound that counted a frame instead
+// of the message would let all of it through.
+func writeInPieces(msg []byte) func(c *websocket.Conn) error {
+	return func(c *websocket.Conn) error {
+		w, err := c.NextWriter(websocket.BinaryMessage)
+		if err != nil {
+			return err
+		}
+		for off := 0; off < len(msg); off += 1 << 10 {
+			if _, err := w.Write(msg[off:min(off+1<<10, len(msg))]); err != nil {
+				return err
+			}
+		}
+		return w.Close()
+	}
+}
+
+func hostSendsAJobAs(write func(c *websocket.Conn) error, out chan<- hostResult) func(c *websocket.Conn) {
+	return func(c *websocket.Conn) {
+		got := map[byte][]byte{}
+		readDone := make(chan struct{})
+		go func() {
+			defer close(readDone)
+			_ = c.SetReadDeadline(time.Now().Add(20 * time.Second))
+			for {
+				_, d, err := c.ReadMessage()
+				if err != nil || len(d) == 0 {
+					return
+				}
+				got[d[0]] = append(got[d[0]], d[1:]...)
+				if d[0] == 'X' {
+					return
+				}
+			}
+		}()
+		werr := write(c)
+		<-readDone
+		out <- hostResult{werr, got}
+	}
+}
+
+// jsonFrame is tag + a JSON object whose field `pad` carries n bytes of 'a'.
+func jsonFrame(tag byte, head string, n int) []byte {
+	b := make([]byte, 0, len(head)+n+16)
+	b = append(b, tag)
+	b = append(b, head...)
+	b = append(b, `,"pad":"`...)
+	b = append(b, bytes.Repeat([]byte{'a'}, n)...)
+	b = append(b, `"}`...)
+	return b
+}
+
+// A harness spec far past the bound is not read to its end: the box answers
+// that it is too large, names the bound, and closes, so the host's write of
+// the rest fails. On main the box read all of it into memory, then ran it.
+func TestHarnessWorkerStopsReadingASpecPastItsBound(t *testing.T) {
+	const frame = 96 << 20
+	out := make(chan hostResult, 1)
+	box := boxLaneHost(t, hostSendsAJob(jsonFrame('H', `{"harness_id":"big","mode":"echo","prompt":"ping"`, frame), out))
+	handleHarnessConn(box, "", "")
+	_ = box.Close()
+	r := <-out
+	if r.werr == nil {
+		t.Fatalf("the box read the whole %d MiB spec (it replied %.200q); it must stop at its bound", frame>>20, r.frames['J'])
+	}
+	if !strings.Contains(string(r.frames['J']), "more than the 32 MiB") || string(r.frames['X']) != "1" {
+		t.Fatalf("want a refusal naming the 32 MiB bound and exit 1: J=%.200q X=%q", r.frames['J'], r.frames['X'])
+	}
+}
+
+// The bound counts the whole message, not one frame of it: a spec past the
+// bound, sent as one message in thousands of 4 KiB frames, is refused the same
+// way and not read to its end. On main the box read all of it.
+func TestHarnessWorkerBoundCountsEveryFrameOfTheSpec(t *testing.T) {
+	const frame = 40 << 20
+	out := make(chan hostResult, 1)
+	spec := jsonFrame('H', `{"harness_id":"pieces","mode":"echo","prompt":"ping"`, frame)
+	box := boxLaneHost(t, hostSendsAJobAs(writeInPieces(spec), out))
+	handleHarnessConn(box, "", "")
+	_ = box.Close()
+	r := <-out
+	if r.werr == nil {
+		t.Fatalf("the box read the whole %d MiB spec sent in small frames (it replied %.200q); it must stop at its bound", frame>>20, r.frames['J'])
+	}
+	if !strings.Contains(string(r.frames['J']), "more than the 32 MiB") || string(r.frames['X']) != "1" {
+		t.Fatalf("want a refusal naming the 32 MiB bound and exit 1: J=%.200q X=%q", r.frames['J'], r.frames['X'])
+	}
+}
+
+// The same for a session spec sent in small frames.
+func TestSessionWorkerBoundCountsEveryFrameOfTheSpec(t *testing.T) {
+	env, _, _ := sessionTestEnv(t, []string{sessionOpOpen}, nil)
+	withSink(t)
+	const frame = 64 << 20
+	out := make(chan hostResult, 1)
+	spec := jsonFrame('S', `{"op":"session:open","session_id":"sess-pieces"`, frame)
+	box := boxLaneHost(t, hostSendsAJobAs(writeInPieces(spec), out))
+	handleSessionWorkerConn(box, "", env)
+	_ = box.Close()
+	r := <-out
+	if r.werr == nil {
+		t.Fatalf("the box read the whole %d MiB session spec sent in small frames (E=%q X=%q); it must stop at its bound",
+			frame>>20, r.frames['E'], r.frames['X'])
+	}
+	if !strings.Contains(string(r.frames['E']), "more than the 1024 KiB") || string(r.frames['X']) != "1" {
+		t.Fatalf("want a refusal naming the bound and exit 1: E=%q X=%q", r.frames['E'], r.frames['X'])
+	}
+}
+
+// A spec the lane can carry still runs: a tree as large as the most this lane
+// hands back (12 MiB in base64), with room on top, past the 16 MiB a host
+// takes in.
+func TestHarnessWorkerTakesALargeSpecTheLaneCarries(t *testing.T) {
+	out := make(chan hostResult, 1)
+	box := boxLaneHost(t, hostSendsAJob(jsonFrame('H', `{"harness_id":"tree","mode":"echo","prompt":"ping","worktree_b64":"dGFy"`, 20<<20), out))
+	handleHarnessConn(box, "", "")
+	r := <-out
+	if r.werr != nil {
+		t.Fatalf("the host could not send a 20 MiB spec: %v", r.werr)
+	}
+	var res map[string]any
+	if err := json.Unmarshal(r.frames['J'], &res); err != nil || res["ok"] != true || res["echo"] != "ping" {
+		t.Fatalf("a 20 MiB spec did not run: J=%.300q X=%q", r.frames['J'], r.frames['X'])
+	}
+}
+
+// The same on a session worker: an S frame far past its 1 MiB bound is
+// refused by name and not read to its end.
+func TestSessionWorkerStopsReadingASpecPastItsBound(t *testing.T) {
+	env, _, _ := sessionTestEnv(t, []string{sessionOpOpen}, nil)
+	withSink(t) // a box that wrongly took the spec opens no real shell
+	const frame = 64 << 20
+	out := make(chan hostResult, 1)
+	box := boxLaneHost(t, hostSendsAJob(jsonFrame('S', `{"op":"session:open","session_id":"sess-big"`, frame), out))
+	handleSessionWorkerConn(box, "", env)
+	_ = box.Close()
+	r := <-out
+	if r.werr == nil {
+		t.Fatalf("the box read the whole %d MiB session spec (E=%q X=%q); it must stop at its bound",
+			frame>>20, r.frames['E'], r.frames['X'])
+	}
+	if !strings.Contains(string(r.frames['E']), "more than the 1024 KiB") || string(r.frames['X']) != "1" {
+		t.Fatalf("want a refusal naming the bound and exit 1: E=%q X=%q", r.frames['E'], r.frames['X'])
+	}
+}
+
+// A session spec near its bound (a grant, and room the size of a root command
+// escaped six times over) still opens a terminal.
+func TestSessionWorkerOpensASpecNearItsBound(t *testing.T) {
+	env, terms, sig := sessionTestEnv(t, []string{sessionOpOpen}, nil)
+	withFake(t)
+	// Here the test is the host: it takes the server end, and the box runs on
+	// the dialled end.
+	hostEnd := make(chan *websocket.Conn, 1)
+	hold := make(chan struct{})
+	t.Cleanup(func() { close(hold) })
+	boxEnd := boxLaneHost(t, func(c *websocket.Conn) { hostEnd <- c; <-hold })
+	b := <-hostEnd
+	head := fmt.Sprintf(`{"op":"session:open","session_id":"sess-near","grant_terms_b64":%q,"grant_sig_b64":%q`, terms, sig)
+	done := make(chan struct{})
+	go func() { defer close(done); handleSessionWorkerConn(boxEnd, "", env) }()
+	if err := b.WriteMessage(websocket.BinaryMessage, jsonFrame('S', head, 900<<10)); err != nil {
+		t.Fatal(err)
+	}
+	if err := b.WriteMessage(websocket.BinaryMessage, []byte("Ihello\n")); err != nil {
+		t.Fatal(err)
+	}
+	got := readUntil(t, b, 'O', 5*time.Second)
+	if string(got['O']) != "hello\n" {
+		t.Fatalf("a 900 KiB spec did not open a terminal: %q", got)
+	}
+	_ = b.Close()
+	<-done
+}
+
+// sinkProc is a terminal that counts what is typed into it and prints nothing.
+type sinkProc struct {
+	n      atomic.Int64
+	killed chan struct{}
+	once   sync.Once
+}
+
+func (p *sinkProc) Read(b []byte) (int, error)  { <-p.killed; return 0, io.EOF }
+func (p *sinkProc) Write(b []byte) (int, error) { p.n.Add(int64(len(b))); return len(b), nil }
+func (p *sinkProc) Resize(c, r int) error       { return nil }
+func (p *sinkProc) Wait() int                   { <-p.killed; return 0 }
+func (p *sinkProc) Kill()                       { p.once.Do(func() { close(p.killed) }) }
+func (p *sinkProc) User() string                { return "sink" }
+func (p *sinkProc) Shell() string               { return "/bin/sink" }
+
+func withSink(t *testing.T) *sinkProc {
+	t.Helper()
+	sp := &sinkProc{killed: make(chan struct{})}
+	old := sessionSpawn
+	sessionSpawn = func(op, shell string, cols, rows int) (sessionProc, error) { return sp, nil }
+	t.Cleanup(func() { sessionSpawn = old; sp.Kill() })
+	return sp
+}
+
+func waitFor(d time.Duration, cond func() bool) bool {
+	end := time.Now().Add(d)
+	for time.Now().Before(end) {
+		if cond() {
+			return true
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+	return cond()
+}
+
+// Terminal input is streamed into the pty: a 16 MiB paste, the most a host
+// relays in one frame, all arrives, and the box never holds it whole.
+func TestSessionInputIsStreamedIntoTheTerminal(t *testing.T) {
+	env, terms, sig := sessionTestEnv(t, []string{sessionOpOpen}, nil)
+	sp := withSink(t)
+	ws := sessionDial(t, env)
+	sendSpec(t, ws, sessionSpec{Op: sessionOpOpen, GrantB64: terms, SigB64: sig, SessionID: "sess-paste"})
+	const paste = 16 << 20
+	frame := append([]byte{'I'}, bytes.Repeat([]byte{'x'}, paste)...)
+	runtime.GC()
+	var before, after runtime.MemStats
+	runtime.ReadMemStats(&before)
+	if err := ws.WriteMessage(websocket.BinaryMessage, frame); err != nil {
+		t.Fatal(err)
+	}
+	if !waitFor(20*time.Second, func() bool { return sp.n.Load() >= paste }) {
+		t.Fatalf("the terminal got %d of %d bytes", sp.n.Load(), paste)
+	}
+	runtime.ReadMemStats(&after)
+	if grew := after.TotalAlloc - before.TotalAlloc; grew > 8<<20 {
+		t.Fatalf("typing a %d MiB paste allocated %d MiB: the frame was held whole", paste>>20, grew>>20)
+	}
+}
+
+// One frame of input past the most a host relays ends the session, and the
+// terminal never receives all of it.
+func TestSessionInputFramePastItsBoundEndsTheSession(t *testing.T) {
+	env, terms, sig := sessionTestEnv(t, []string{sessionOpOpen}, nil)
+	sp := withSink(t)
+	ws := sessionDial(t, env)
+	sendSpec(t, ws, sessionSpec{Op: sessionOpOpen, GrantB64: terms, SigB64: sig, SessionID: "sess-huge"})
+	const frame = 17 << 20
+	go func() {
+		_ = ws.WriteMessage(websocket.BinaryMessage, append([]byte{'I'}, bytes.Repeat([]byte{'x'}, frame)...))
+	}()
+	ended := func() bool {
+		select {
+		case <-sp.killed:
+			return true
+		default:
+			return false
+		}
+	}
+	if !waitFor(20*time.Second, ended) {
+		t.Fatalf("a %d MiB input frame did not end the session; the terminal got %d bytes", frame>>20, sp.n.Load())
+	}
+	if n := sp.n.Load(); n >= frame {
+		t.Fatalf("the terminal received all %d bytes of a frame past the bound", n)
+	}
+}
+
+// The mirror forwarder streams a host frame into the replica's socket: reading
+// 4 KiB of a 64 MiB frame does not need the frame held whole.
+func TestMirrorForwarderStreamsAHostFrame(t *testing.T) {
+	const frame = 64 << 20
+	payload := bytes.Repeat([]byte{'m'}, frame)
+	start := make(chan struct{})
+	up := websocket.Upgrader{}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c, err := up.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer c.Close()
+		_, _, _ = c.ReadMessage() // the BEARER hello
+		<-start
+		_ = c.WriteMessage(websocket.BinaryMessage, payload)
+		time.Sleep(3 * time.Second)
+	}))
+	defer srv.Close()
+	addr := srv.Listener.Addr().String()
+	dialer := &websocket.Dialer{NetDialContext: func(ctx context.Context, n, a string) (net.Conn, error) {
+		return (&net.Dialer{}).DialContext(ctx, "tcp", addr)
+	}}
+	boxEnd, replica := net.Pipe()
+	defer replica.Close()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go pipeMirrorConn(ctx, boxEnd, dialer, "ws://"+addr+"/", "b")
+	buf := make([]byte, 4096)
+	runtime.GC()
+	var before, after runtime.MemStats
+	runtime.ReadMemStats(&before)
+	close(start)
+	_ = replica.SetReadDeadline(time.Now().Add(20 * time.Second))
+	if _, err := io.ReadFull(replica, buf); err != nil || buf[0] != 'm' {
+		t.Fatalf("the replica read: %v", err)
+	}
+	runtime.ReadMemStats(&after)
+	if grew := after.TotalAlloc - before.TotalAlloc; grew > 8<<20 {
+		t.Fatalf("passing 4 KiB of a %d MiB host frame to the replica allocated %d MiB: the frame was held whole",
+			frame>>20, grew>>20)
+	}
+}
+
+// The forwarder carries at most eight connections at once. It listens on a
+// fixed port (5433), so the test needs that port free: run it in its own
+// network namespace on a machine where something already holds it.
+func TestMirrorForwarderCapsConnectionsInFlight(t *testing.T) {
+	probe, err := net.Listen("tcp", "127.0.0.1:5433")
+	if err != nil {
+		t.Skipf("127.0.0.1:5433 is in use here (%v); run this test in its own network namespace", err)
+	}
+	_ = probe.Close()
+	var live, most atomic.Int64
+	release := make(chan struct{})
+	up := websocket.Upgrader{}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c, err := up.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer c.Close()
+		n := live.Add(1)
+		for {
+			m := most.Load()
+			if n <= m || most.CompareAndSwap(m, n) {
+				break
+			}
+		}
+		<-release
+		live.Add(-1)
+	}))
+	defer srv.Close()
+	defer close(release)
+	addr := srv.Listener.Addr().String()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go serveMirrorForward(ctx, "127.0.0.1", "127.0.0.0/8", "ws://"+addr+"/", nil,
+		func(ctx context.Context) (net.Conn, error) { return (&net.Dialer{}).DialContext(ctx, "tcp", addr) }, "b")
+	var conns []net.Conn
+	defer func() {
+		for _, c := range conns {
+			_ = c.Close()
+		}
+	}()
+	if !waitFor(5*time.Second, func() bool {
+		c, err := net.Dial("tcp", "127.0.0.1:5433")
+		if err == nil {
+			conns = append(conns, c)
+		}
+		return err == nil
+	}) {
+		t.Fatal("the forwarder never listened")
+	}
+	for i := 1; i < 12; i++ {
+		c, err := net.Dial("tcp", "127.0.0.1:5433")
+		if err != nil {
+			t.Fatal(err)
+		}
+		conns = append(conns, c)
+	}
+	waitFor(3*time.Second, func() bool { return most.Load() >= 12 })
+	if m := most.Load(); m != 8 {
+		t.Fatalf("12 local connections made %d WebSockets to the host at once; want 8", m)
+	}
+}
+
+// A coder run's egress proxy reads a request head under a bound: a request
+// line that never ends is not read to its end.
+func TestCoderEgressHeadIsBounded(t *testing.T) {
+	boxEnd, runEnd := net.Pipe()
+	var refused int64
+	done := make(chan struct{})
+	go func() { defer close(done); handleCoderEgressConn(boxEnd, nil, &refused) }()
+	chunk := bytes.Repeat([]byte{'a'}, 64<<10)
+	_, _ = runEnd.Write([]byte("GET http://example.invalid/"))
+	sent := 0
+	var werr error
+	for sent < 8<<20 {
+		_ = runEnd.SetWriteDeadline(time.Now().Add(5 * time.Second))
+		var n int
+		n, werr = runEnd.Write(chunk)
+		sent += n
+		if werr != nil {
+			break
+		}
+	}
+	_ = runEnd.Close()
+	<-done
+	if werr == nil {
+		t.Fatalf("the proxy read %d MiB of a request line with no end; it must stop at a bound", sent>>20)
+	}
+}
+
+// A coder run's proxy carries at most 1024 connections at once; the next is
+// answered 503 at once, never left waiting.
+func TestCoderEgressCapsConnectionsInFlight(t *testing.T) {
+	dir := t.TempDir()
+	run := "lane-bounds-test"
+	if err := coderListenEgress(run, dir, os.Getuid()); err != nil {
+		t.Fatal(err)
+	}
+	defer coderCloseEgress(run)
+	sock := filepath.Join(dir, coderEgressSock)
+	var conns []net.Conn
+	defer func() {
+		for _, c := range conns {
+			_ = c.Close()
+		}
+	}()
+	for i := 0; i < 1024; i++ {
+		c, err := net.Dial("unix", sock)
+		if err != nil {
+			if errors.Is(err, syscall.EMFILE) || errors.Is(err, syscall.ENFILE) {
+				t.Skipf("this process may not open %d sockets: %v", 2*1024, err)
+			}
+			t.Fatal(err)
+		}
+		conns = append(conns, c)
+	}
+	// Every one of them is idle in its head, so the proxy holds them all.
+	time.Sleep(200 * time.Millisecond)
+	extra, err := net.Dial("unix", sock)
+	if err != nil {
+		t.Fatal(err)
+	}
+	conns = append(conns, extra)
+	_ = extra.SetReadDeadline(time.Now().Add(3 * time.Second))
+	line, _ := io.ReadAll(io.LimitReader(extra, 256))
+	if !strings.HasPrefix(string(line), "HTTP/1.1 503") {
+		t.Fatalf("connection 1025 of one run got %q; want a 503 at once", line)
+	}
+}