From 43c07cf27efd9f6302fb035f38173291f356b12d Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 14:35:33 -0700 Subject: [PATCH 01/16] Add bounded manual relay capacity runner --- internal/capacity/README.md | 227 ++++++++++ internal/capacity/admission_test.go | 30 ++ internal/capacity/canary.go | 82 ++++ internal/capacity/churn_test.go | 48 +++ internal/capacity/client.go | 304 +++++++++++++ internal/capacity/client_test.go | 34 ++ internal/capacity/cmd/main.go | 65 +++ internal/capacity/config.go | 271 ++++++++++++ internal/capacity/config_test.go | 97 +++++ internal/capacity/guard.go | 288 +++++++++++++ internal/capacity/guard_test.go | 110 +++++ internal/capacity/report.go | 72 ++++ internal/capacity/runner.go | 643 ++++++++++++++++++++++++++++ internal/capacity/runner_test.go | 317 ++++++++++++++ internal/capacity/safety_test.go | 98 +++++ internal/capacity/transfer.go | 313 ++++++++++++++ internal/capacity/workload.go | 175 ++++++++ internal/capacity/workload_test.go | 50 +++ 18 files changed, 3224 insertions(+) create mode 100644 internal/capacity/README.md create mode 100644 internal/capacity/admission_test.go create mode 100644 internal/capacity/canary.go create mode 100644 internal/capacity/churn_test.go create mode 100644 internal/capacity/client.go create mode 100644 internal/capacity/client_test.go create mode 100644 internal/capacity/cmd/main.go create mode 100644 internal/capacity/config.go create mode 100644 internal/capacity/config_test.go create mode 100644 internal/capacity/guard.go create mode 100644 internal/capacity/guard_test.go create mode 100644 internal/capacity/report.go create mode 100644 internal/capacity/runner.go create mode 100644 internal/capacity/runner_test.go create mode 100644 internal/capacity/safety_test.go create mode 100644 internal/capacity/transfer.go create mode 100644 internal/capacity/workload.go create mode 100644 internal/capacity/workload_test.go diff --git a/internal/capacity/README.md b/internal/capacity/README.md new file mode 100644 index 000000000..16701cc01 --- /dev/null +++ b/internal/capacity/README.md @@ -0,0 +1,227 @@ +# Manual relay capacity experiment + +This internal tool creates real headless `host.Host` sessions, authenticated SSH +or WebSocket guests, PTY shell channels and optional SFTP traffic. It is opt-in; +no test or normal invocation contacts Fly. It changes no production client or +server behavior and supplies no infrastructure credentials. + +Build and invoke explicitly: + +```sh +go build -o /tmp/upterm-capacity ./internal/capacity/cmd +/tmp/upterm-capacity -manifest /secure/run/manifest.json -phase /secure/run/phase.json -validate +/tmp/upterm-capacity -manifest /secure/run/manifest.json -phase /secure/run/phase.json > /secure/run/events.jsonl +``` + +Both files are required. JSON decoding rejects unknown fields. Validation is +pure: it checks syntax, limits and authorization without DNS queries or dials. +A load run also resolves every destination immediately before connecting and +rejects the entire DNS answer if any address is outside the target's IP list. +It then dials a literal approved IP, preventing DNS check/dial races. Invalid +configuration and preexisting watchdog stops refuse admission. Workload, +readiness, telemetry, budget or watchdog failures produce a nonzero exit. + +All examples below use placeholders, not operational destinations or keys. +Replace `EXPIRES_UTC` with a UTC RFC3339 timestamp within the next 24 hours. + +Manifest: + +```json +{ + "run_id": "RUN_ID", + "expires_at": "EXPIRES_UTC", + "local_only": false, + "targets": [{ + "app_id": "ISOLATED_APP_ID", + "machine_id": "RELAY_MACHINE_ID", + "ssh_url": "ssh://RELAY_MACHINE_ID.vm.ISOLATED_APP_ID.internal:2222", + "ws_url": "ws://RELAY_MACHINE_ID.vm.ISOLATED_APP_ID.internal:8080", + "public_ssh_url": "ssh://ISOLATED_APP_ID.fly.dev:22", + "public_ws_url": "wss://ISOLATED_APP_ID.fly.dev:443", + "node_addresses": ["RELAY_MACHINE_ID.vm.ISOLATED_APP_ID.internal:2222"], + "allowed_ips": ["RELAY_6PN_IP", "AUTHORIZED_PUBLIC_IPV4", "AUTHORIZED_PUBLIC_IPV6"], + "host_keys": ["ssh-ed25519 BASE64_RUN_RELAY_PUBLIC_KEY"], + "telemetry_url": "http://RELAY_MACHINE_ID.vm.ISOLATED_APP_ID.internal:9091/metrics" + }] +} +``` + +There may be one or two target records, with exact machine IDs. `host_target` +and `entry_target` reference these IDs. For crossed-node experiments, list both +nodes and choose different IDs. A guest session node must match a manifest node +exactly, resolve within that node target's IP set, and agree with the node +encoded in an embedded-mode SSH username. Consul-mode usernames must equal the +session ID. This check runs at initial creation, before every guest connection, +and after host restoration. Hosts always redial their original guarded URL. + +`host_keys` pins raw relay keys or certificate signing keys. Host certificates +must have the expected hostname principal, valid signature/time and a pinned +signing key. Include the exact private machine hostname and authorized public +hostname in the relay certificate for a run using both. Known production relay +names/apps and the published production relay key are refused. The manifest +is an authorization record supplied by the operator; it cannot independently +prove which Fly resource owns an arbitrary IP. + +Phase: + +```json +{ + "hosts": 10, + "guests_per_host": 2, + "host_target": "RELAY_MACHINE_ID", + "entry_target": "RELAY_MACHINE_ID", + "host_public": false, + "entry_public": false, + "host_ws_percent": 50, + "guest_ws_percent": 50, + "channels_per_guest": 1, + "duration_seconds": 300, + "admissions_per_second": 10, + "payload_budget_bytes": 5368709120, + "active": true, + "transfer_concurrency": 0, + "churn_percent": 0, + "host_identity_file": "/secure/run/host-identity", + "watchdog_file": "/secure/run/STOP", + "watchdog_heartbeat_file": "/secure/run/PERMIT" +} +``` + +`host_public`/`entry_public` select the optional public URLs; otherwise the +private URLs are used. Both public URLs must be present when selected. Transport +percentages assign the first `floor(role_count * percentage / 100)` clients to +WS/WSS, the rest to SSH. JSONL records each role, actual node and transport. +Idle phases with zero guests per host reserve one extra external connection for +a transient one-channel guest canary at readiness and every 30 seconds. Canary +hosts are sampled round-robin; canary counts, maximum ready latency and RTT +samples are separate from the nominal zero steady guests. Each canary closes +before the next, within this same runner. Channels per steady guest must +be one or four. Only the first channel generates steady input; all channels +drain fan-out and each proves readiness independently. The optional identity +file authenticates host registration only; the runner generates distinct +session host keys and isolated temporary directories. Its private contents and +path are never logged. Without it, each host gets an ephemeral identity. + +Bounds apply before any host starts: at most two relay targets, 1,000 external +connections including one telemetry connection per target and the idle host-only canary +slot, 20 new external +connections/second, 24 hours of authorization and 50 GiB of application payload. +Up to twenty admission workers overlap readiness; every external connection +and retry passes the same global limiter, including `Host`'s own reconnect +attempts. Admission workers are joined before steady observation or cleanup. +The process rejects nonempty uppercase or lowercase `HTTP_PROXY`, `HTTPS_PROXY` +and `ALL_PROXY`. Its private loopback CONNECT doors enforce destination/IP +allowlists and own the host sockets used for cuts. `local_only: true` is a +separate mode requiring literal loopback endpoints and only loopback allowed +IPs; it does not authorize arbitrary external destinations. + +The steady duration starts after all requested hosts, guests and channels are +ready. Setup has a separate `(H + G) / rate + 30 seconds` deadline; manifest +expiry and parent cancellation bound both. Summary `setup_seconds`, +`steady_seconds` and total `duration_seconds` preserve the distinction. A +five-minute hold therefore remains five minutes after admission. Host keepalive +uses the released client's 50-second interval and normal reconnect timings. + +## Offered data and measurement + +The hosted command is this binary's own `helper` subcommand. It sets its PTY to +raw mode, answers a distinct 128-byte readiness challenge, echoes deterministic +data and never interprets +input as commands. Active guests send exactly 128 bytes once per second, +beginning at steady observation start. Active helpers emit a 1 KiB frame every +ten seconds from helper startup. Idle helpers emit only their startup marker +and readiness challenge/probe responses. A request tagged `R` receives a +helper-generated `A` response, so a PTY kernel echo cannot prove readiness. +This also works after the original startup marker leaves terminal history. +All receivers drain output, track sequence gaps +and reordering, and fail if a requested response remains outstanding for more +than ten seconds. Initial terminal replay is excluded until a channel's fresh +probe response establishes the live measurement boundary. Planned reconnects +reset tracking only for the affected source generation; survivor generations +retain gap/reordering checks. + +Each 128-byte input/echo and 1,024-byte host frame contains a 44-byte ASCII +header (`CAP`, role, eight-digit source and sequence, nineteen-digit Unix +nanoseconds), padding, and one final LF. The 45 framing bytes are included in +those fixed sizes. Readiness challenge and response frames are also 128 bytes; `handshake_frames` +records their requests separately from `input_frames`. The helper startup +marker is ten bytes. SSH, TLS, +WebSocket, TCP, terminal reconstruction and relay forwarding overhead are not +measured as application framing. RTT uses the originating guest's echoed tag; +summary quantiles use a uniform reservoir of at most 100,000 samples and also +report the complete sample count. Every receiver observes all host fan-out. `peak_external_connections` includes +transient canaries, telemetry and admitted in-flight dials; it conservatively +counts a dial until it fails or its connection closes. + +The payload budget uses conservative pre-generation reservations: host output +through the maximum bounded lifetime is reserved before hosts start; each +input reserves its echo to all configured fan-out channels before writing; +each transfer reserves upload and download before creating its file. +`reserved_payload_bytes` is the safety bound; `payload_bytes` records observed +input, live received frames and actual SFTP bytes. Unused fan-out/output +reservations are intentionally retained. Preflight rejects a budget too small +for the intended duration/workload. Reconnect probes reserve extra bytes at +runtime. These byte figures are per phase; the operator must also enforce the +approved aggregate budget across attempts. + +`transfer_concurrency` can be zero, four or eight. Each transfer concurrently +uploads one deterministic 16 MiB file into its host's temporary directory, +then downloads it and verifies both length and SHA-256. Source/download data +are streamed with bounded memory; at most eight temporary 16 MiB files exist. +Files and host directories are removed during teardown; `cleanup_complete` +requires the owned directory to be absent before return. JSONL contains actual +upload/download bytes, digest and verification outcome. Transfers incomplete +at phase expiry are a failure. + +`churn_percent` can be zero, ten or fifty. It requires at least two hosts, two +guests per host and 25 steady seconds. Counts are floored; use enough hosts for +the requested percentage to select a host. Selected host and guest transports +are cut over a ten-second schedule starting one second into steady load. +Host restoration uses the original `Host.Run`, host key and session ID; +guests reuse their identities and transport kinds. Reconnect events and +latency are separate from first admissions. A survivor on an unaffected host +keeps its transport and must continue responding. Intentional input pauses are +reported as `paused_active_frames`. Idle churn explicitly probes its survivor. +Churn incomplete at phase expiry is a failure. Combining transfers and churn +may fail a transfer cut in progress; run distinct phases to measure each. + +## Telemetry and external stop control + +Private relay Prometheus text is collected initially and every five seconds. +Redirects, non-200 responses and responses over 1 MiB fail collection. A +separate one-second watchdog stops the phase once any target has had more than +30 seconds without successful telemetry. JSONL also records generator PID, +Go version, goroutines, heap/runtime memory and, on Linux, `/proc/self/status` +and `/proc/self/stat` for RSS/CPU analysis. Sampling does not diagnose pressure; +the coordinator evaluates Fly metrics/OOM and generator pressure externally. + +Before startup and once per second, an existing `watchdog_file` requests a +graceful safety stop. If `watchdog_heartbeat_file` is configured, it must exist +at startup and have an mtime no older than 30 seconds; stale, missing or +future-dated heartbeats stop the run. A coordinator can `touch` the heartbeat +at least every five seconds while continued authorization holds, or create +`STOP` when an OOM, pressure, cost or other limit is reached. SIGINT/SIGTERM +also cancel gracefully. All transports are closed, owned goroutines and host +processes are joined, and temporary directories are removed before the final +summary. The generator requires no Fly API credential. + +Library callers use `Run(ctx, Config, io.Writer) (Summary, error)`. Supply the +parsed manifest in `Config.Manifest`. `HelperExecutable` defaults to the running +executable; an embedding program must either dispatch the `helper` subcommand +or supply the path of a built capacity binary. Tests build a dedicated helper +and only use embedded loopback relays. + +## Focused checks + +```sh +go test -race -count=1 ./internal/capacity/... +go vet ./internal/capacity/... +GOOS=windows go vet ./internal/capacity/... +CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o /tmp/upterm-capacity-linux ./internal/capacity/cmd +``` + +The real integration suite covers SSH, WS and both crossed transport mixes, +four shell channels, deterministic concurrent SFTP, actual host/guest cuts, +cancellation, cleanup, confinement and safety accounting. It runs no external +load and requires no local Consul. The Windows check establishes compilation +and vet compatibility; the manual benchmark targets Linux. diff --git a/internal/capacity/admission_test.go b/internal/capacity/admission_test.go new file mode 100644 index 000000000..841050fac --- /dev/null +++ b/internal/capacity/admission_test.go @@ -0,0 +1,30 @@ +package capacity + +import ( + "context" + "testing" + "time" +) + +func TestAdmissionWorkersOverlapReadinessWithinBound(t *testing.T) { + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(context.Canceled) + r := &runState{ctx: ctx, cancel: cancel} + started := make(chan struct{}, 40) + release := make(chan struct{}) + finished := make(chan error, 1) + go func() { finished <- r.admissions(40, func(int) error { started <- struct{}{}; <-release; return nil }) }() + defer func() { close(release); <-finished }() + for i := 0; i < 20; i++ { + select { + case <-started: + case <-time.After(time.Second): + t.Fatalf("only %d readiness checks overlapped", i) + } + } + select { + case <-started: + t.Fatal("more than twenty admission workers started") + case <-time.After(20 * time.Millisecond): + } +} diff --git a/internal/capacity/canary.go b/internal/capacity/canary.go new file mode 100644 index 000000000..9f4591de6 --- /dev/null +++ b/internal/capacity/canary.go @@ -0,0 +1,82 @@ +package capacity + +import ( + "time" +) + +// Idle host-only phases keep zero steady guests. One transient guest at +// readiness and every thirty seconds proves a sampled host still serves a +// shell, without a separate generator or unbounded extra connections. +func (r *runState) hostCanary(sample int) error { + et, _ := r.cfg.target(r.cfg.EntryTarget) + if r.canary == nil { + identity, err := makeSigner() + if err != nil { + return err + } + q := &guest{id: 0, target: et, identity: identity, canary: true, expected: true} + r.clientsMu.Lock() + r.guests = append(r.guests, q) + r.clientsMu.Unlock() + r.canary = q + } + q := r.canary + q.host = r.hosts[sample%len(r.hosts)] + u, err := transport(et, r.cfg.GuestWSPercent, (sample*37)%100, 100, r.cfg.EntryPublic) + if err != nil { + return err + } + q.endpoint = u + proxy, err := newTunnelProxy(r.ctx, r.guard, et) + if err != nil { + return err + } + defer proxy.close() + q.proxy = proxy + q.mu.Lock() + q.generation++ + q.mu.Unlock() + r.rec.mu.Lock() + r.rec.summary.CanaryOffered++ + r.rec.mu.Unlock() + began := time.Now() + err = q.connect(r.ctx, r) + readyElapsed := time.Since(began) + q.setExpected(true) + q.close() + if err != nil { + return err + } + ms := float64(readyElapsed) / float64(time.Millisecond) + r.rec.mu.Lock() + r.rec.summary.CanaryEstablished++ + if ms > r.rec.summary.CanaryReadyMillis { + r.rec.summary.CanaryReadyMillis = ms + } + r.rec.mu.Unlock() + r.rec.event("host_only_canary", map[string]any{"sample": sample, "host": q.host.id, "transport": u.Scheme, "ready_ms": ms, "transient_guest": true, "channels": 1}) + return nil +} +func (r *runState) canaries() { + r.workers.Add(1) + go func() { + defer r.workers.Done() + tick := time.NewTicker(30 * time.Second) + defer tick.Stop() + sample := 1 + for { + select { + case <-r.ctx.Done(): + return + case <-tick.C: + if err := r.hostCanary(sample); err != nil { + if r.ctx.Err() == nil { + r.fail(err) + } + return + } + sample++ + } + } + }() +} diff --git a/internal/capacity/churn_test.go b/internal/capacity/churn_test.go new file mode 100644 index 000000000..8d58b9377 --- /dev/null +++ b/internal/capacity/churn_test.go @@ -0,0 +1,48 @@ +package capacity + +import ( + "bytes" + "testing" + "time" +) + +func TestChurnSchedulesCutsAcrossTenSeconds(t *testing.T) { + actions := churnSchedule(20, 4, 10) + if len(actions) != 10 { + t.Fatalf("want two host and eight guest cuts, got %d", len(actions)) + } + if actions[0].offset != 0 || actions[len(actions)-1].offset != 10*time.Second { + t.Fatalf("cuts do not span ten seconds: %+v", actions) + } + for i := 1; i < len(actions); i++ { + if actions[i].offset <= actions[i-1].offset { + t.Fatal("cuts not staggered") + } + } +} +func TestRealTransportChurnRestoresOriginalSessions(t *testing.T) { + c := localRelay(t) + c.Hosts = 2 + c.GuestsPerHost = 2 + c.DurationSeconds = 25 + c.ChurnPercent = 50 + c.PayloadBudgetBytes = 4 << 20 + c.HostWSPercent = 50 + c.GuestWSPercent = 50 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if err != nil { + t.Fatalf("churn: %v\n%s", err, out.String()) + } + if s.Reconnects < 3 { + t.Fatalf("real host/guest restoration absent: %+v", s) + } + for _, event := range []string{`"event":"transport_cut"`, `"event":"host_restored"`, `"event":"guest_restored"`, `"event":"survivor_canary"`} { + if !bytes.Contains(out.Bytes(), []byte(event)) { + t.Fatalf("missing %s", event) + } + } + if s.EstablishedHosts != 2 || s.EstablishedGuests != 4 { + t.Fatalf("reconnect was counted as new admission: %+v", s) + } +} diff --git a/internal/capacity/client.go b/internal/capacity/client.go new file mode 100644 index 000000000..e09d6032a --- /dev/null +++ b/internal/capacity/client.go @@ -0,0 +1,304 @@ +package capacity + +import ( + "context" + "crypto/ed25519" + "crypto/rand" + "encoding/base64" + "errors" + "fmt" + "io" + mathrand "math/rand/v2" + "net" + "net/url" + "sync" + "time" + + "github.com/owenthereal/upterm/host/api" + "github.com/owenthereal/upterm/upterm" + "github.com/owenthereal/upterm/ws" + "golang.org/x/crypto/ssh" +) + +func makeSigner() (ssh.Signer, error) { + _, key, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + return nil, err + } + return ssh.NewSignerFromKey(key) +} +func guestHandshake(ctx context.Context, raw net.Conn, addr string, cfg *ssh.ClientConfig) (*ssh.Client, error) { + stop := context.AfterFunc(ctx, func() { _ = raw.Close() }) + defer stop() + if deadline, ok := ctx.Deadline(); ok { + _ = raw.SetDeadline(deadline) + } + conn, chans, reqs, err := ssh.NewClientConn(raw, addr, cfg) + if err != nil { + _ = raw.Close() + if ctx.Err() != nil { + return nil, ctx.Err() + } + return nil, err + } + if !stop() { + _ = conn.Close() + return nil, ctx.Err() + } + _ = raw.SetDeadline(time.Time{}) + return ssh.NewClient(conn, chans, reqs), nil +} +func openGuest(ctx context.Context, g *guard, t Target, u *url.URL, s *api.GetSessionResponse, identity ssh.Signer, proxy *tunnelProxy) (*ssh.Client, net.Conn, error) { + if err := g.session(ctx, s); err != nil { + return nil, nil, err + } + var raw net.Conn + var err error + if u.Scheme == "ssh" { + raw, err = g.connect(ctx, t, u.Host) + } else { + copy := *u + copy.User = url.UserPassword(s.SessionId, base64.URLEncoding.EncodeToString([]byte(s.NodeAddr))) + raw, err = ws.NewWSConnContext(ctx, ©, true, proxy.url) + } + if err != nil { + return nil, nil, err + } + cfg := &ssh.ClientConfig{User: s.SshUser, Auth: []ssh.AuthMethod{ssh.PublicKeys(identity)}, ClientVersion: upterm.ClientSSHClientVersion, HostKeyCallback: pinnedKey(t)} + client, err := guestHandshake(ctx, raw, u.Host, cfg) + return client, raw, err +} + +type guest struct { + canary bool + probeSeq [4]int + pending map[string]time.Time + lifetimeStop func() bool + acks int64 + ops sync.Mutex + id int + host *hostRun + target Target + endpoint *url.URL + identity ssh.Signer + proxy *tunnelProxy + client *ssh.Client + raw net.Conn + channels []*ssh.Session + writers []io.WriteCloser + done []chan struct{} + seq int + generation int + expected bool + mu sync.Mutex +} + +func (q *guest) close() { + q.ops.Lock() + defer q.ops.Unlock() + if q.lifetimeStop != nil { + q.lifetimeStop() + q.lifetimeStop = nil + } + if q.client != nil { + _ = q.client.Close() + } + if q.raw != nil { + _ = q.raw.Close() + } + for _, d := range q.done { + <-d + } + q.done = nil + q.channels = nil + q.writers = nil +} +func (q *guest) connect(ctx context.Context, r *runState) error { + q.ops.Lock() + defer q.ops.Unlock() + s, err := q.host.session(ctx) + if err != nil { + return err + } + if err = r.guard.session(ctx, s); err != nil { + return err + } + attempt, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + q.client, q.raw, err = openGuest(attempt, r.guard, q.target, q.endpoint, s, q.identity, q.proxy) + if err != nil { + return err + } + client := q.client + q.lifetimeStop = context.AfterFunc(ctx, func() { _ = client.Close() }) + stop := context.AfterFunc(attempt, func() { _ = client.Close() }) + defer stop() + channels := r.cfg.ChannelsPerGuest + if q.canary { + channels = 1 + } + for ch := 0; ch < channels; ch++ { + session, err := q.client.NewSession() + if err != nil { + return err + } + q.channels = append(q.channels, session) + if err = session.RequestPty("dumb", 100, 2048, ssh.TerminalModes{ssh.ECHO: 0}); err != nil { + return err + } + in, err := session.StdinPipe() + if err != nil { + return err + } + out, err := session.StdoutPipe() + if err != nil { + return err + } + session.Stderr = io.Discard + if err = session.Shell(); err != nil { + return err + } + q.writers = append(q.writers, in) + stream := q.id*4 + ch + probeSeq := q.probeSeq[ch] + q.probeSeq[ch]++ + if ch == 0 { + probeSeq = q.seq + q.seq++ + } + ready := make(chan struct{}, 1) + helperReady := make(chan struct{}, 1) + done := make(chan struct{}) + q.done = append(q.done, done) + go func() { + defer close(done) + tracker := sequenceTracker{} + epochs := map[string]int{} + live := false + err := scanFrames(out, func(o observation) error { + if !live { + if o.kind != "I" || o.id != stream || o.seq != probeSeq { + return nil + } + live = true + } + if err := r.observe(tracker, epochs, o); err != nil { + return err + } + r.actual(int64(frameLength(o.kind))) + r.rec.mu.Lock() + r.rec.summary.ReceivedFrames++ + if o.kind == "I" && o.id == stream { + q.mu.Lock() + q.acks++ + delete(q.pending, fmt.Sprintf("%d:%d", o.id, o.seq)) + q.mu.Unlock() + ms := float64(time.Since(o.sent)) / float64(time.Millisecond) + if q.canary { + r.rec.summary.CanaryRTTCount++ + r.rec.summary.CanaryRTTMillis = append(r.rec.summary.CanaryRTTMillis, ms) + } else { + r.rec.summary.RTTCount++ + if len(r.rec.summary.RTTMillis) < 100000 { + r.rec.summary.RTTMillis = append(r.rec.summary.RTTMillis, ms) + } else if i := mathrand.IntN(r.rec.summary.RTTCount); i < 100000 { + r.rec.summary.RTTMillis[i] = ms + } + } + } + r.rec.mu.Unlock() + if o.kind == "I" && o.id == stream { + select { + case ready <- struct{}{}: + default: + } + } + return nil + }, func(ack observation) { + r.actual(128) + if ack.id != stream || ack.seq != probeSeq { + return + } + select { + case helperReady <- struct{}{}: + default: + } + }) + q.mu.Lock() + expected := q.expected + q.mu.Unlock() + if err != nil && !expected && ctx.Err() == nil { + r.fail(fmt.Errorf("guest %d channel %d: %w", q.id, ch, err)) + } + }() + if err = r.sendKind(q, in, stream, probeSeq, "R"); err != nil { + return err + } + select { + case <-helperReady: + case <-attempt.Done(): + return fmt.Errorf("helper readiness: %w", attempt.Err()) + case <-ctx.Done(): + return ctx.Err() + } + if err = r.send(q, in, stream, probeSeq); err != nil { + return err + } + select { + case <-ready: + case <-attempt.Done(): + return fmt.Errorf("terminal readiness: %w", attempt.Err()) + case <-ctx.Done(): + return ctx.Err() + } + } + if !stop() { + return errors.New("guest readiness cancelled") + } + r.rec.event("guest_ready", map[string]any{"guest": q.id, "host": q.host.id, "transport": q.endpoint.Scheme, "node": s.NodeAddr, "session_id": s.SessionId, "channels": len(q.channels), "generation": q.generation, "canary": q.canary}) + return nil +} +func frameLength(kind string) int { + if kind == "H" { + return 1024 + } + return 128 +} +func (q *guest) active(r *runState) error { + q.ops.Lock() + defer q.ops.Unlock() + q.mu.Lock() + expected := q.expected + q.mu.Unlock() + if expected { + r.rec.mu.Lock() + r.rec.summary.PausedActiveFrames++ + r.rec.mu.Unlock() + return nil + } + if len(q.writers) == 0 { + return errors.New("guest has no terminal") + } + err := r.send(q, q.writers[0], q.id*4, q.seq) + q.seq++ + if err == nil { + r.rec.mu.Lock() + r.rec.summary.ActiveFrames++ + r.rec.mu.Unlock() + } + return err +} +func (q *guest) setExpected(b bool) { q.mu.Lock(); q.expected = b; q.mu.Unlock() } +func transport(t Target, percent, index, total int, public bool) (*url.URL, error) { + raw := t.SSHURL + if public { + raw = t.PublicSSHURL + } + if index < total*percent/100 { + raw = t.WSURL + if public { + raw = t.PublicWSURL + } + } + return endpoint(raw, "ssh", "ws", "wss") +} diff --git a/internal/capacity/client_test.go b/internal/capacity/client_test.go new file mode 100644 index 000000000..9a01b6677 --- /dev/null +++ b/internal/capacity/client_test.go @@ -0,0 +1,34 @@ +package capacity + +import ( + "context" + "net" + "testing" + "time" + + "golang.org/x/crypto/ssh" +) + +func TestSilentHandshakeCancellationClosesTransport(t *testing.T) { + client, server := net.Pipe() + defer server.Close() + ctx, cancel := context.WithCancel(t.Context()) + finished := make(chan error, 1) + go func() { + _, err := guestHandshake(ctx, client, "127.0.0.1:22", &ssh.ClientConfig{User: "session", HostKeyCallback: ssh.InsecureIgnoreHostKey()}) + finished <- err + }() + cancel() + select { + case err := <-finished: + if err == nil { + t.Fatal("silent handshake reported success") + } + case <-time.After(time.Second): + t.Fatal("handshake ignored cancellation") + } + _ = server.SetReadDeadline(time.Now().Add(time.Second)) + if _, err := server.Read(make([]byte, 1024)); err == nil { + t.Fatal("cancelled transport remained open") + } +} diff --git a/internal/capacity/cmd/main.go b/internal/capacity/cmd/main.go new file mode 100644 index 000000000..3bfa678f0 --- /dev/null +++ b/internal/capacity/cmd/main.go @@ -0,0 +1,65 @@ +package main + +import ( + "context" + "flag" + "fmt" + "os" + "os/signal" + "strconv" + "syscall" + + "github.com/owenthereal/upterm/internal/capacity" +) + +func main() { + if err := run(); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} +func run() error { + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + if len(os.Args) > 1 && os.Args[1] == "helper" { + if len(os.Args) != 4 { + return fmt.Errorf("helper requires host ID and active flag") + } + id, err := strconv.Atoi(os.Args[2]) + if err != nil { + return err + } + active, err := strconv.ParseBool(os.Args[3]) + if err != nil { + return err + } + return capacity.Helper(ctx, os.Stdin, os.Stdout, id, active) + } + manifest := flag.String("manifest", "", "explicit bounded-run manifest JSON") + phase := flag.String("phase", "", "explicit phase JSON") + validate := flag.Bool("validate", false, "validate only; never dial") + flag.Parse() + if *manifest == "" || *phase == "" { + return fmt.Errorf("both -manifest and -phase are required") + } + m, err := os.Open(*manifest) + if err != nil { + return err + } + defer m.Close() + p, err := os.Open(*phase) + if err != nil { + return err + } + defer p.Close() + cfg, err := capacity.Decode(m, p) + if err != nil { + return err + } + if *validate { + fmt.Fprintln(os.Stdout, `{"event":"validated","dialed":false}`) + return nil + } + _, err = capacity.Run(ctx, cfg, os.Stdout) + return err +} diff --git a/internal/capacity/config.go b/internal/capacity/config.go new file mode 100644 index 000000000..2ed489d47 --- /dev/null +++ b/internal/capacity/config.go @@ -0,0 +1,271 @@ +// Package capacity is the opt-in, bounded relay experiment runner. +package capacity + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/netip" + "net/url" + "os" + "strings" + "time" + + "golang.org/x/crypto/ssh" +) + +type Manifest struct { + RunID string `json:"run_id"` + ExpiresAt time.Time `json:"expires_at"` + LocalOnly bool `json:"local_only"` + Targets []Target `json:"targets"` +} + +type Target struct { + PublicSSHURL string `json:"public_ssh_url,omitempty"` + PublicWSURL string `json:"public_ws_url,omitempty"` + AppID string `json:"app_id"` + MachineID string `json:"machine_id"` + SSHURL string `json:"ssh_url"` + WSURL string `json:"ws_url"` + NodeAddresses []string `json:"node_addresses"` + AllowedIPs []string `json:"allowed_ips"` + HostKeys []string `json:"host_keys"` + TelemetryURL string `json:"telemetry_url"` +} + +type Config struct { + HostPublic bool `json:"host_public"` + EntryPublic bool `json:"entry_public"` + HostIdentityFile string `json:"host_identity_file"` + WatchdogHeartbeatFile string `json:"watchdog_heartbeat_file"` + // HelperExecutable must be this runner binary (defaults to os.Executable). + HelperExecutable string `json:"-"` + Manifest Manifest `json:"-"` + Hosts int `json:"hosts"` + GuestsPerHost int `json:"guests_per_host"` + HostTarget string `json:"host_target"` + EntryTarget string `json:"entry_target"` + HostWSPercent int `json:"host_ws_percent"` + GuestWSPercent int `json:"guest_ws_percent"` + ChannelsPerGuest int `json:"channels_per_guest"` + DurationSeconds int `json:"duration_seconds"` + AdmissionsPerSecond int `json:"admissions_per_second"` + PayloadBudgetBytes int64 `json:"payload_budget_bytes"` + Active bool `json:"active"` + TransferConcurrency int `json:"transfer_concurrency"` + ChurnPercent int `json:"churn_percent"` + WatchdogFile string `json:"watchdog_file"` +} + +const MaxPayload = int64(50 << 30) +const TransferSize = int64(16 << 20) + +func decodeOne(r io.Reader, v any) error { + if r == nil { + return errors.New("explicit manifest and phase inputs required") + } + d := json.NewDecoder(io.LimitReader(r, 1<<20)) + d.DisallowUnknownFields() + if err := d.Decode(v); err != nil { + return err + } + var extra any + if err := d.Decode(&extra); err != io.EOF { + return errors.New("configuration must contain one JSON object") + } + return nil +} + +func Decode(manifest, phase io.Reader) (Config, error) { + var c Config + if err := decodeOne(manifest, &c.Manifest); err != nil { + return c, fmt.Errorf("manifest: %w", err) + } + if err := decodeOne(phase, &c); err != nil { + return c, fmt.Errorf("phase: %w", err) + } + return c, c.Validate(time.Now()) +} + +func production(s string) bool { + s = strings.ToLower(strings.TrimSuffix(s, ".")) + return s == "upterm" || s == "uptermd" || s == "upterm.fly.dev" || s == "upterm.internal" || strings.HasSuffix(s, ".upterm.internal") || s == "upterm.dev" || strings.HasSuffix(s, ".upterm.dev") || s == "upterm.io" || strings.HasSuffix(s, ".upterm.io") +} + +func endpoint(raw string, schemes ...string) (*url.URL, error) { + u, err := url.Parse(raw) + if err != nil { + return nil, errors.New("invalid endpoint URL") + } + ok := false + for _, s := range schemes { + ok = ok || u.Scheme == s + } + if !ok || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || u.Opaque != "" || production(u.Hostname()) { + return nil, errors.New("unauthorized endpoint") + } + if u.Path != "" && u.Path != "/" && !(u.Scheme == "http" && u.Path == "/metrics") { + return nil, errors.New("unexpected endpoint path") + } + if u.Port() == "" { + ports := map[string]string{"ssh": "22", "ws": "80", "wss": "443", "http": "80"} + u.Host = net.JoinHostPort(u.Hostname(), ports[u.Scheme]) + } + if _, err := net.LookupPort("tcp", u.Port()); err != nil { + return nil, errors.New("invalid endpoint port") + } + return u, nil +} + +func (c Config) target(id string) (Target, bool) { + for _, t := range c.Manifest.Targets { + if t.MachineID == id { + return t, true + } + } + return Target{}, false +} + +func (c Config) Validate(now time.Time) error { + m := c.Manifest + if m.RunID == "" || !m.ExpiresAt.After(now) || m.ExpiresAt.Sub(now) > 24*time.Hour || len(m.Targets) < 1 || len(m.Targets) > 2 { + return errors.New("manifest missing, expired, or beyond 24-hour/two-target limit") + } + for _, k := range []string{"HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"} { + if os.Getenv(k) != "" { + return fmt.Errorf("implicit proxy environment %s refused", k) + } + } + if c.Hosts < 1 || c.Hosts > 998 || c.GuestsPerHost < 0 || c.GuestsPerHost > 998 || c.Hosts*(1+c.GuestsPerHost)+len(m.Targets)+c.canaryConnection() > 1000 { + return errors.New("host/guest external connection limit exceeded") + } + if c.AdmissionsPerSecond < 1 || c.AdmissionsPerSecond > 20 || c.DurationSeconds < 1 || c.DurationSeconds > 86400 || time.Duration(c.DurationSeconds)*time.Second > m.ExpiresAt.Sub(now) { + return errors.New("rate/duration outside approved bounds") + } + if c.PayloadBudgetBytes < 1 || c.PayloadBudgetBytes > MaxPayload || (c.ChannelsPerGuest != 1 && c.ChannelsPerGuest != 4) { + return errors.New("byte/channel limit outside approved bounds") + } + if c.HostWSPercent < 0 || c.HostWSPercent > 100 || c.GuestWSPercent < 0 || c.GuestWSPercent > 100 { + return errors.New("transport percentage outside 0..100") + } + if c.TransferConcurrency != 0 && c.TransferConcurrency != 4 && c.TransferConcurrency != 8 { + return errors.New("transfer concurrency must be zero, four or eight") + } + if c.TransferConcurrency > 0 && (c.GuestsPerHost == 0 || int64(c.TransferConcurrency)*2*TransferSize > c.PayloadBudgetBytes) { + return errors.New("transfer requires guests and sufficient byte budget") + } + if c.ChurnPercent != 0 && c.ChurnPercent != 10 && c.ChurnPercent != 50 { + return errors.New("churn must be zero, ten or fifty percent") + } + if c.ChurnPercent > 0 && (c.Hosts < 2 || c.GuestsPerHost < 2 || c.DurationSeconds < 25 || c.Hosts*c.ChurnPercent/100 < 1) { + return errors.New("churn needs survivor hosts/guests and at least 25 seconds") + } + fanout := int64(1 + c.GuestsPerHost*c.ChannelsPerGuest) + probes := int64(c.Hosts*c.GuestsPerHost*c.ChannelsPerGuest) * 256 * fanout + if c.idleHostOnly() { + probes += int64(1+(c.DurationSeconds-1)/30) * 512 + } + active := int64(0) + if c.Active { + active = int64(c.Hosts*c.GuestsPerHost*c.DurationSeconds) * 128 * fanout + } + if c.hostOutputReservation()+probes+active+int64(c.TransferConcurrency)*2*TransferSize > c.PayloadBudgetBytes { + return errors.New("payload budget cannot cover offered workload") + } + seen := map[string]bool{} + for _, t := range m.Targets { + if t.AppID == "" || t.MachineID == "" || seen[t.MachineID] || production(t.AppID) || len(t.AllowedIPs) == 0 || len(t.HostKeys) == 0 || len(t.NodeAddresses) == 0 { + return errors.New("manifest requires isolated, unique machine IDs, IPs, nodes and keys") + } + seen[t.MachineID] = true + for _, ip := range t.AllowedIPs { + a, err := netip.ParseAddr(ip) + if err != nil || a.IsUnspecified() || a.IsMulticast() || a.IsLoopback() != m.LocalOnly { + return errors.New("manifest IP violates explicit local-only isolation") + } + } + for _, k := range t.HostKeys { + key, _, _, rest, err := ssh.ParseAuthorizedKey([]byte(k)) + if err != nil || len(rest) != 0 || ssh.FingerprintSHA256(key) == "SHA256:9ajV8JqMe6jJE/s3TYjb/9xw7T0pfJ2+gADiBIJWDPE" { + return errors.New("invalid or production host key") + } + } + endpoints := map[string][]string{t.SSHURL: {"ssh"}, t.WSURL: {"ws", "wss"}, t.TelemetryURL: {"http"}} + if t.PublicSSHURL != "" { + endpoints[t.PublicSSHURL] = []string{"ssh"} + } + if t.PublicWSURL != "" { + endpoints[t.PublicWSURL] = []string{"ws", "wss"} + } + for raw, schemes := range endpoints { + u, err := endpoint(raw, schemes...) + if err != nil { + return err + } + if a, err := netip.ParseAddr(u.Hostname()); err == nil && !contains(t.AllowedIPs, a.String()) { + return errors.New("endpoint IP absent from manifest") + } + if m.LocalOnly && net.ParseIP(u.Hostname()) == nil { + return errors.New("local-only endpoints must use literal loopback IPs") + } + } + for _, node := range t.NodeAddresses { + h, p, err := net.SplitHostPort(node) + if err != nil || h == "" || p == "" || production(h) { + return errors.New("invalid/production session node") + } + if a, err := netip.ParseAddr(h); err == nil && !contains(t.AllowedIPs, a.String()) { + return errors.New("session node IP absent from manifest") + } + } + } + if _, ok := c.target(c.HostTarget); !ok { + return errors.New("unknown host target") + } + if _, ok := c.target(c.EntryTarget); !ok { + return errors.New("unknown entry target") + } + for _, selection := range []struct { + id string + public bool + }{{c.HostTarget, c.HostPublic}, {c.EntryTarget, c.EntryPublic}} { + if selection.public { + t, _ := c.target(selection.id) + if t.PublicSSHURL == "" || t.PublicWSURL == "" { + return errors.New("public phase requires explicit public endpoints") + } + } + } + return nil +} + +func (c Config) hostOutputReservation() int64 { + fanout := int64(1 + c.GuestsPerHost*c.ChannelsPerGuest) + // Startup has an independent, enforced ramp+30s deadline. Teardown can + // produce one final tick, reserved here before any helper is started. + seconds := int64(c.DurationSeconds + 42 + c.Hosts*(1+c.GuestsPerHost)/c.AdmissionsPerSecond) + output := int64(10) // helper-ready marker + if c.Active { + output += (seconds/10 + 1) * 1024 + } + return int64(c.Hosts) * output * fanout +} +func contains(ss []string, s string) bool { + for _, v := range ss { + if s == v { + return true + } + } + return false +} + +func (c Config) idleHostOnly() bool { return c.GuestsPerHost == 0 && !c.Active } +func (c Config) canaryConnection() int { + if c.idleHostOnly() { + return 1 + } + return 0 +} diff --git a/internal/capacity/config_test.go b/internal/capacity/config_test.go new file mode 100644 index 000000000..910f7834e --- /dev/null +++ b/internal/capacity/config_test.go @@ -0,0 +1,97 @@ +package capacity + +import ( + "context" + "crypto/ed25519" + "crypto/rand" + "io" + "strings" + "testing" + "time" + + "golang.org/x/crypto/ssh" +) + +func signer(t *testing.T) ssh.Signer { + t.Helper() + _, key, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatal(err) + } + s, err := ssh.NewSignerFromKey(key) + if err != nil { + t.Fatal(err) + } + return s +} + +func validConfig(t *testing.T) Config { + t.Helper() + return Config{Manifest: Manifest{RunID: "test-run", ExpiresAt: time.Now().Add(time.Hour), LocalOnly: true, + Targets: []Target{{AppID: "test-app", MachineID: "test-machine", SSHURL: "ssh://127.0.0.1:2222", WSURL: "ws://127.0.0.1:8080", NodeAddresses: []string{"127.0.0.1:2222"}, AllowedIPs: []string{"127.0.0.1"}, HostKeys: []string{string(ssh.MarshalAuthorizedKey(signer(t).PublicKey()))}, TelemetryURL: "http://127.0.0.1:9091/metrics"}}}, + Hosts: 1, GuestsPerHost: 1, HostTarget: "test-machine", EntryTarget: "test-machine", ChannelsPerGuest: 1, DurationSeconds: 2, AdmissionsPerSecond: 20, PayloadBudgetBytes: 1 << 20} +} + +func clearProxies(t *testing.T) { + t.Helper() + for _, k := range []string{"HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"} { + t.Setenv(k, "") + } +} + +// Removing any bound or manifest check must allow one of these forbidden runs. +func TestInvalidConfigurationRefusesRun(t *testing.T) { + clearProxies(t) + tests := []struct { + name string + change func(*Config) + }{ + {"absent manifest", func(c *Config) { c.Manifest = Manifest{} }}, + {"expired", func(c *Config) { c.Manifest.ExpiresAt = time.Now().Add(-time.Second) }}, + {"production", func(c *Config) { c.Manifest.Targets[0].SSHURL = "ssh://uptermd.upterm.dev:22" }}, + {"production app", func(c *Config) { c.Manifest.Targets[0].AppID = "upterm" }}, + {"local escape", func(c *Config) { c.Manifest.Targets[0].AllowedIPs = []string{"8.8.8.8"} }}, + {"zero hosts", func(c *Config) { c.Hosts = 0 }}, + {"negative guests", func(c *Config) { c.GuestsPerHost = -1 }}, + {"overflow counts", func(c *Config) { c.Hosts = int(^uint(0) >> 1); c.GuestsPerHost = 100 }}, + {"connections", func(c *Config) { c.Hosts = 500; c.GuestsPerHost = 1 }}, + {"rate", func(c *Config) { c.AdmissionsPerSecond = 21 }}, + {"zero rate", func(c *Config) { c.AdmissionsPerSecond = 0 }}, + {"duration", func(c *Config) { c.DurationSeconds = 86401 }}, + {"bytes", func(c *Config) { c.PayloadBudgetBytes = 50*(1<<30) + 1 }}, + {"channels", func(c *Config) { c.ChannelsPerGuest = 5 }}, + {"transport", func(c *Config) { c.GuestWSPercent = 101 }}, + {"unknown target", func(c *Config) { c.EntryTarget = "other" }}, + {"transfer", func(c *Config) { c.TransferConcurrency = 5 }}, + {"churn", func(c *Config) { c.ChurnPercent = 99 }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := validConfig(t) + tt.change(&cfg) + if _, err := Run(context.Background(), cfg, io.Discard); err == nil { + t.Fatal("unsafe configuration started successfully") + } + }) + } +} + +func TestProxyEnvironmentRefused(t *testing.T) { + clearProxies(t) + for _, k := range []string{"HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"} { + t.Run(k, func(t *testing.T) { + t.Setenv(k, "socks5://127.0.0.1:9") + if err := validConfig(t).Validate(time.Now()); err == nil { + t.Fatal("implicit proxy accepted") + } + }) + } +} + +func TestDecodeRequiresBothExplicitInputs(t *testing.T) { + for _, input := range []struct{ m, p io.Reader }{{nil, strings.NewReader("{}")}, {strings.NewReader("{}"), nil}, {strings.NewReader("{}"), strings.NewReader(`{"unrecognized":true}`)}} { + if _, err := Decode(input.m, input.p); err == nil { + t.Fatal("missing or unknown configuration accepted") + } + } +} diff --git a/internal/capacity/guard.go b/internal/capacity/guard.go new file mode 100644 index 000000000..ac0cfb923 --- /dev/null +++ b/internal/capacity/guard.go @@ -0,0 +1,288 @@ +package capacity + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "net/http" + "net/url" + "strings" + "sync" + "time" + + "github.com/owenthereal/upterm/host/api" + "github.com/owenthereal/upterm/routing" + "golang.org/x/crypto/ssh" +) + +type guard struct { + cfg Config + lookup func(context.Context, string) ([]net.IPAddr, error) + dial func(context.Context, string, string) (net.Conn, error) + mu sync.Mutex + next time.Time + active int + peak int +} + +func newGuard(c Config) *guard { + return &guard{cfg: c, lookup: net.DefaultResolver.LookupIPAddr, dial: (&net.Dialer{Timeout: 10 * time.Second}).DialContext} +} +func (g *guard) resolved(ctx context.Context, t Target, addr string) (string, error) { + h, p, err := net.SplitHostPort(addr) + if err != nil || production(h) { + return "", errors.New("unauthorized destination") + } + var ips []net.IPAddr + if ip := net.ParseIP(h); ip != nil { + ips = []net.IPAddr{{IP: ip}} + } else { + ips, err = g.lookup(ctx, h) + if err != nil { + return "", err + } + } + if len(ips) == 0 { + return "", errors.New("empty DNS response") + } + for _, ip := range ips { + if !contains(t.AllowedIPs, ip.IP.String()) { + return "", errors.New("DNS address outside manifest") + } + } + return net.JoinHostPort(ips[0].IP.String(), p), nil +} +func (g *guard) admit(ctx context.Context) error { + g.mu.Lock() + when := g.next + if when.Before(time.Now()) { + when = time.Now() + } + g.next = when.Add(time.Second / time.Duration(g.cfg.AdmissionsPerSecond)) + g.mu.Unlock() + timer := time.NewTimer(time.Until(when)) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} +func (g *guard) connect(ctx context.Context, t Target, addr string) (net.Conn, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + allowed := false + for _, raw := range []string{t.SSHURL, t.WSURL, t.PublicSSHURL, t.PublicWSURL, t.TelemetryURL} { + u, err := endpoint(raw, "ssh", "ws", "wss", "http") + if err == nil && u.Host == addr { + allowed = true + } + } + if !allowed { + return nil, errors.New("destination absent from manifest endpoints") + } + pinned, err := g.resolved(ctx, t, addr) + if err != nil { + return nil, err + } + if err = g.admit(ctx); err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } + if !time.Now().Before(g.cfg.Manifest.ExpiresAt) { + return nil, errors.New("manifest expired") + } + g.mu.Lock() + if g.active >= 1000 { + g.mu.Unlock() + return nil, errors.New("external connection safety stop") + } + g.active++ + if g.active > g.peak { + g.peak = g.active + } + g.mu.Unlock() + conn, err := g.dial(ctx, "tcp", pinned) + if err != nil { + g.mu.Lock() + g.active-- + g.mu.Unlock() + return nil, err + } + return &countedConn{Conn: conn, release: func() { g.mu.Lock(); g.active--; g.mu.Unlock() }}, nil +} + +type countedConn struct { + net.Conn + once sync.Once + release func() +} + +func (c *countedConn) Close() error { err := c.Conn.Close(); c.once.Do(c.release); return err } +func (g *guard) session(ctx context.Context, s *api.GetSessionResponse) error { + if s == nil || s.SessionId == "" { + return errors.New("missing session") + } + for _, t := range g.cfg.Manifest.Targets { + if contains(t.NodeAddresses, s.NodeAddr) { + if _, err := g.resolved(ctx, t, s.NodeAddr); err != nil { + return err + } + if strings.Contains(s.SshUser, ":") { + id, node, err := (&routing.EmbeddedEncodeDecoder{}).Decode(s.SshUser) + if err != nil || id != s.SessionId || node != s.NodeAddr { + return errors.New("SSH route disagrees with allowed session node") + } + } else if s.SshUser != s.SessionId { + return errors.New("SSH route disagrees with session ID") + } + return nil + } + } + return errors.New("session node outside manifest") +} +func pinnedKey(t Target) ssh.HostKeyCallback { + var pins []ssh.PublicKey + for _, s := range t.HostKeys { + k, _, _, _, _ := ssh.ParseAuthorizedKey([]byte(s)) + pins = append(pins, k) + } + matches := func(k ssh.PublicKey) bool { + for _, p := range pins { + if p != nil && bytes.Equal(p.Marshal(), k.Marshal()) { + return true + } + } + return false + } + return func(addr string, _ net.Addr, k ssh.PublicKey) error { + if cert, ok := k.(*ssh.Certificate); ok { + h, _, err := net.SplitHostPort(addr) + if err != nil { + return err + } + if cert.CertType != ssh.HostCert || !matches(cert.SignatureKey) { + return errors.New("untrusted relay certificate") + } + return (&ssh.CertChecker{}).CheckCert(h, cert) + } + if !matches(k) { + return errors.New("relay key differs from manifest") + } + return nil + } +} + +// A private CONNECT door keeps Host's own reconnect machinery behind the same +// DNS pinning and admission limiter as guests. It also owns the actual sockets +// severed during host churn; no host or session is replaced. +type tunnelProxy struct { + g *guard + t Target + ctx context.Context + server *http.Server + listener net.Listener + url *url.URL + mu sync.Mutex + conns map[net.Conn]bool + closing bool + workers sync.WaitGroup + done chan struct{} +} + +func newTunnelProxy(ctx context.Context, g *guard, t Target) (*tunnelProxy, error) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, err + } + p := &tunnelProxy{g: g, t: t, ctx: ctx, listener: ln, conns: map[net.Conn]bool{}, done: make(chan struct{})} + p.url = &url.URL{Scheme: "http", Host: ln.Addr().String()} + p.server = &http.Server{Handler: p, ReadHeaderTimeout: 5 * time.Second} + go func() { defer close(p.done); _ = p.server.Serve(ln) }() + return p, nil +} +func (p *tunnelProxy) track(c net.Conn) bool { + p.mu.Lock() + defer p.mu.Unlock() + if p.closing { + _ = c.Close() + return false + } + p.conns[c] = true + return true +} +func (p *tunnelProxy) untrack(c net.Conn) { + _ = c.Close() + p.mu.Lock() + delete(p.conns, c) + p.mu.Unlock() +} +func (p *tunnelProxy) cut() { + p.mu.Lock() + defer p.mu.Unlock() + for c := range p.conns { + _ = c.Close() + } +} +func (p *tunnelProxy) close() { + p.mu.Lock() + p.closing = true + for c := range p.conns { + _ = c.Close() + } + p.mu.Unlock() + _ = p.server.Close() + <-p.done + p.workers.Wait() +} +func (p *tunnelProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { + p.mu.Lock() + if p.closing { + p.mu.Unlock() + return + } + p.workers.Add(1) + p.mu.Unlock() + defer p.workers.Done() + if r.Method != http.MethodConnect { + http.Error(w, "CONNECT required", 405) + return + } + ctx, cancel := context.WithTimeout(p.ctx, 10*time.Second) + defer cancel() + upstream, err := p.g.connect(ctx, p.t, r.Host) + if err != nil { + http.Error(w, "destination refused", 403) + return + } + if !p.track(upstream) { + return + } + defer p.untrack(upstream) + downstream, rw, err := w.(http.Hijacker).Hijack() + if err != nil { + return + } + if !p.track(downstream) { + return + } + defer p.untrack(downstream) + if _, err = rw.WriteString("HTTP/1.1 200 Connection Established\r\n\r\n"); err != nil { + return + } + if err = rw.Flush(); err != nil { + return + } + copied := make(chan struct{}) + go func() { defer close(copied); _, _ = io.Copy(upstream, rw) }() + _, _ = io.Copy(downstream, upstream) + _ = downstream.Close() + _ = upstream.Close() + <-copied +} diff --git a/internal/capacity/guard_test.go b/internal/capacity/guard_test.go new file mode 100644 index 000000000..3bb6a5717 --- /dev/null +++ b/internal/capacity/guard_test.go @@ -0,0 +1,110 @@ +package capacity + +import ( + "context" + "encoding/base64" + "errors" + "net" + "testing" + + "github.com/owenthereal/upterm/host/api" + "github.com/owenthereal/upterm/server" + "golang.org/x/crypto/ssh" +) + +func TestGuardInvalidDestinationNeverDials(t *testing.T) { + c := validConfig(t) + g := newGuard(c) + calls := 0 + g.dial = func(context.Context, string, string) (net.Conn, error) { + calls++ + return nil, errors.New("dial attempted") + } + if _, err := g.connect(t.Context(), c.Manifest.Targets[0], "uptermd.upterm.dev:22"); err == nil { + t.Fatal("escaped endpoint accepted") + } + if calls != 0 { + t.Fatalf("invalid endpoint caused %d dials", calls) + } +} + +func TestGuardDNSRebindingNeverDials(t *testing.T) { + c := validConfig(t) + c.Manifest.LocalOnly = false + c.Manifest.Targets[0].SSHURL = "ssh://isolated.internal:2222" + c.Manifest.Targets[0].AllowedIPs = []string{"fdaa::1"} + g := newGuard(c) + calls := 0 + g.lookup = func(context.Context, string) ([]net.IPAddr, error) { + return []net.IPAddr{{IP: net.ParseIP("fdaa::1")}, {IP: net.ParseIP("8.8.8.8")}}, nil + } + g.dial = func(context.Context, string, string) (net.Conn, error) { + calls++ + return nil, errors.New("dial attempted") + } + if _, err := g.connect(t.Context(), c.Manifest.Targets[0], "isolated.internal:2222"); err == nil { + t.Fatal("unapproved DNS response accepted") + } + if calls != 0 { + t.Fatalf("DNS escape caused %d dials", calls) + } +} + +func TestGuardSessionNodeEscape(t *testing.T) { + g := newGuard(validConfig(t)) + for _, s := range []*api.GetSessionResponse{ + {SessionId: "session", SshUser: "session", NodeAddr: "uptermd.upterm.dev:22"}, + {SessionId: "session", SshUser: "session:" + base64.URLEncoding.EncodeToString([]byte("8.8.8.8:22")), NodeAddr: "127.0.0.1:2222"}, + {SessionId: "session", SshUser: "different", NodeAddr: "127.0.0.1:2222"}, + } { + if err := g.session(t.Context(), s); err == nil { + t.Fatal("session route escape accepted") + } + } +} +func TestPinnedKeysRejectForeignCertificatesAndPrincipals(t *testing.T) { + c := validConfig(t) + key := signer(t) + c.Manifest.Targets[0].HostKeys = []string{string(ssh.MarshalAuthorizedKey(key.PublicKey()))} + check := pinnedKey(c.Manifest.Targets[0]) + addr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 22} + if err := check("127.0.0.1:22", addr, key.PublicKey()); err != nil { + t.Fatal(err) + } + foreign := signer(t) + if err := check("127.0.0.1:22", addr, foreign.PublicKey()); err == nil { + t.Fatal("foreign relay key accepted") + } + for _, principal := range []string{"127.0.0.1", "another.internal"} { + cert, err := (&server.HostCertSigner{Hostnames: []string{principal}}).SignCert(key) + if err != nil { + t.Fatal(err) + } + err = check("127.0.0.1:22", addr, cert.PublicKey()) + if (err == nil) != (principal == "127.0.0.1") { + t.Fatalf("certificate principal %s: %v", principal, err) + } + } +} +func TestExternalPeakIncludesTransientConnections(t *testing.T) { + c := validConfig(t) + g := newGuard(c) + g.dial = func(context.Context, string, string) (net.Conn, error) { + client, server := net.Pipe() + t.Cleanup(func() { _ = server.Close() }) + return client, nil + } + first, err := g.connect(t.Context(), c.Manifest.Targets[0], "127.0.0.1:2222") + if err != nil { + t.Fatal(err) + } + second, err := g.connect(t.Context(), c.Manifest.Targets[0], "127.0.0.1:2222") + if err != nil { + t.Fatal(err) + } + _ = first.Close() + _ = second.Close() + if g.active != 0 || g.peak != 2 { + t.Fatalf("transient external connections missing from peak: active=%d peak=%d", g.active, g.peak) + } +} diff --git a/internal/capacity/report.go b/internal/capacity/report.go new file mode 100644 index 000000000..c13ae6531 --- /dev/null +++ b/internal/capacity/report.go @@ -0,0 +1,72 @@ +package capacity + +import ( + "encoding/json" + "io" + "sync" + "time" +) + +type Summary struct { + PeakExternalConnections int `json:"peak_external_connections"` + HandshakeFrames int64 `json:"handshake_frames"` + CanaryRTTCount int `json:"canary_rtt_count"` + CanaryRTTMillis []float64 `json:"canary_rtt_millis"` + CanaryOffered int `json:"canary_offered"` + CanaryEstablished int `json:"canary_established"` + CanaryReadyMillis float64 `json:"canary_ready_millis"` + CleanupComplete bool `json:"cleanup_complete"` + PausedActiveFrames int64 `json:"paused_active_frames"` + InputFrames int64 `json:"input_frames"` + ActiveFrames int64 `json:"active_frames"` + ChurnCompleted bool `json:"churn_completed"` + RTTSampling string `json:"rtt_sampling"` + ReservedPayloadBytes int64 `json:"reserved_payload_bytes"` + SetupSeconds float64 `json:"setup_seconds"` + SteadySeconds float64 `json:"steady_seconds"` + RunID string `json:"run_id"` + OfferedHosts int `json:"offered_hosts"` + OfferedGuests int `json:"offered_guests"` + OfferedChannels int `json:"offered_channels"` + EstablishedHosts int `json:"established_hosts"` + EstablishedGuests int `json:"established_guests"` + EstablishedChannels int `json:"established_channels"` + AdmissionFailures int `json:"admission_failures"` + PayloadBytes int64 `json:"payload_bytes"` + PayloadBudgetBytes int64 `json:"payload_budget_bytes"` + ReceivedFrames int64 `json:"received_frames"` + RTTCount int `json:"rtt_count"` + RTTMillis []float64 `json:"rtt_millis"` + RTTP50Millis float64 `json:"rtt_p50_millis"` + RTTP95Millis float64 `json:"rtt_p95_millis"` + RTTP99Millis float64 `json:"rtt_p99_millis"` + Reconnects int `json:"reconnects"` + Transfers int `json:"transfers"` + TransferBytes int64 `json:"transfer_bytes"` + DurationSeconds float64 `json:"duration_seconds"` + StopReason string `json:"stop_reason"` +} + +type recorder struct { + mu sync.Mutex + encoder *json.Encoder + err error + summary Summary +} + +func (r *recorder) event(kind string, fields map[string]any) { + r.mu.Lock() + defer r.mu.Unlock() + if fields == nil { + fields = map[string]any{} + } + fields["event"] = kind + fields["at"] = time.Now().UTC() + fields["run_id"] = r.summary.RunID + if r.err == nil { + r.err = r.encoder.Encode(fields) + } +} +func newRecorder(out io.Writer, c Config) *recorder { + return &recorder{encoder: json.NewEncoder(out), summary: Summary{RunID: c.Manifest.RunID, RTTSampling: "uniform_reservoir_100000", OfferedHosts: c.Hosts, OfferedGuests: c.Hosts * c.GuestsPerHost, OfferedChannels: c.Hosts * c.GuestsPerHost * c.ChannelsPerGuest, PayloadBudgetBytes: c.PayloadBudgetBytes}} +} diff --git a/internal/capacity/runner.go b/internal/capacity/runner.go new file mode 100644 index 000000000..197011a94 --- /dev/null +++ b/internal/capacity/runner.go @@ -0,0 +1,643 @@ +package capacity + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "os" + "path/filepath" + "runtime" + "sort" + "strconv" + "sync" + "time" + + "github.com/owenthereal/upterm/host" + "github.com/owenthereal/upterm/host/api" + "golang.org/x/crypto/ssh" +) + +var errDuration = errors.New("duration") + +type runState struct { + canary *guest + clientsMu sync.RWMutex + telemetryMu sync.Mutex + lastTelemetry map[string]time.Time + cfg Config + guard *guard + rec *recorder + ctx context.Context + cancel context.CancelCauseFunc + hosts []*hostRun + guests []*guest + proxies []*tunnelProxy + workers sync.WaitGroup + helper string + dir string +} +type hostRun struct { + repair sync.Mutex + id int + socket string + proxy *tunnelProxy + h *host.Host + done chan error + mu sync.Mutex + initial *api.GetSessionResponse + tunnel chan host.TunnelEvent +} + +func (h *hostRun) session(ctx context.Context) (*api.GetSessionResponse, error) { + conn, err := host.AdminConn(h.socket) + if err != nil { + return nil, err + } + defer conn.Close() + s, err := api.NewAdminServiceClient(conn).GetSession(ctx, &api.GetSessionRequest{}) + if err != nil { + return nil, err + } + h.mu.Lock() + initial := h.initial + h.mu.Unlock() + if initial == nil || s.SessionId != initial.SessionId { + return nil, errors.New("host session identity changed") + } + return s, nil +} +func (r *runState) fail(err error) { r.cancel(err) } +func (r *runState) reserve(n int64) error { + r.rec.mu.Lock() + defer r.rec.mu.Unlock() + if n > r.cfg.PayloadBudgetBytes-r.rec.summary.ReservedPayloadBytes { + return errors.New("payload_budget") + } + r.rec.summary.ReservedPayloadBytes += n + return nil +} +func (r *runState) actual(n int64) { + r.rec.mu.Lock() + r.rec.summary.PayloadBytes += n + r.rec.mu.Unlock() +} +func (r *runState) send(q *guest, w io.Writer, stream, seq int) error { + return r.sendKind(q, w, stream, seq, "I") +} +func (r *runState) sendKind(q *guest, w io.Writer, stream, seq int, kind string) error { + fanout := r.cfg.GuestsPerHost * r.cfg.ChannelsPerGuest + if q.canary { + fanout = 1 + } + if err := r.reserve(128 * int64(1+fanout)); err != nil { + return err + } + now := time.Now() + key := fmt.Sprintf("%d:%d", stream, seq) + if kind == "I" { + q.mu.Lock() + if q.pending == nil { + q.pending = map[string]time.Time{} + } + q.pending[key] = now + q.mu.Unlock() + } + n, err := w.Write(frame(kind, stream, seq, now, 128)) + if n == 128 && err == nil { + r.rec.mu.Lock() + if kind == "I" { + r.rec.summary.InputFrames++ + } else { + r.rec.summary.HandshakeFrames++ + } + r.rec.mu.Unlock() + } + r.actual(int64(n)) + return err +} +func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { + u, err := transport(t, r.cfg.HostWSPercent, index, r.cfg.Hosts, r.cfg.HostPublic) + if err != nil { + return err + } + proxy, err := newTunnelProxy(r.ctx, r.guard, t) + if err != nil { + return err + } + r.clientsMu.Lock() + r.proxies = append(r.proxies, proxy) + r.clientsMu.Unlock() + dir := filepath.Join(r.dir, "h"+strconv.Itoa(index)) + if err = os.Mkdir(dir, 0700); err != nil { + return err + } + ready := make(chan string, 1) + h := &hostRun{id: index, socket: filepath.Join(dir, "admin.sock"), proxy: proxy, done: make(chan error, 1), tunnel: make(chan host.TunnelEvent, 16)} + if identity == nil { + identity, err = makeSigner() + if err != nil { + return err + } + } + key, err := makeSigner() + if err != nil { + return err + } + h.h = &host.Host{KeepAliveDuration: 50 * time.Second, Host: u.String(), Command: []string{r.helper, "helper", strconv.Itoa(index), strconv.FormatBool(r.cfg.Active)}, CommandDir: dir, Signers: []ssh.Signer{identity}, IdentitiesOnly: true, HostKey: key, HostKeyCallback: pinnedKey(t), ProxyURL: proxy.url, AdminSocketFile: h.socket, AutoAccept: true, SFTPAutoAccept: true, ForwardingDisabled: true, ExecDisabled: true, Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), StopGrace: time.Second, + SessionCreatedCallback: func(ctx context.Context, s *api.GetSessionResponse) error { + if err := r.guard.session(ctx, s); err != nil { + return err + } + h.mu.Lock() + h.initial = s + h.mu.Unlock() + return nil + }, + SessionReadyCallback: func(status string) { + select { + case ready <- status: + default: + } + }, + TunnelCallback: func(e host.TunnelEvent) { + r.rec.event("host_tunnel", map[string]any{"host": index, "up": e.Up, "can_reconnect": e.CanReconnect, "route_changed": e.RouteChanged}) + select { + case h.tunnel <- e: + default: + r.fail(errors.New("host tunnel event overflow")) + } + }, + } + r.hosts[index] = h + began := time.Now() + go func() { + err := h.h.Run(r.ctx) + h.done <- err + if r.ctx.Err() == nil { + r.fail(fmt.Errorf("host %d stopped: %v", index, err)) + } + }() + timer := time.NewTimer(12 * time.Second) + defer timer.Stop() + select { + case status := <-ready: + if status != "ready" { + return fmt.Errorf("host readiness is %s", status) + } + case err := <-h.done: + h.done <- err + return fmt.Errorf("host startup: %w", err) + case <-timer.C: + return errors.New("host readiness deadline") + case <-r.ctx.Done(): + return context.Cause(r.ctx) + } + r.rec.mu.Lock() + r.rec.summary.EstablishedHosts++ + r.rec.mu.Unlock() + h.mu.Lock() + s := h.initial + h.mu.Unlock() + r.rec.event("host_ready", map[string]any{"host": index, "transport": u.Scheme, "node": s.NodeAddr, "session_id": s.SessionId, "ready_ms": float64(time.Since(began)) / float64(time.Millisecond)}) + return nil +} +func (r *runState) admit(identity ssh.Signer) error { + ht, _ := r.cfg.target(r.cfg.HostTarget) + et, _ := r.cfg.target(r.cfg.EntryTarget) + r.hosts = make([]*hostRun, r.cfg.Hosts) + if err := r.admissions(r.cfg.Hosts, func(i int) error { return r.startHost(i, ht, identity) }); err != nil { + return err + } + for _, h := range r.hosts { + for j := 0; j < r.cfg.GuestsPerHost; j++ { + id := len(r.guests) + u, err := transport(et, r.cfg.GuestWSPercent, id, r.cfg.Hosts*r.cfg.GuestsPerHost, r.cfg.EntryPublic) + if err != nil { + return err + } + key, err := makeSigner() + if err != nil { + return err + } + proxy, err := newTunnelProxy(r.ctx, r.guard, et) + if err != nil { + return err + } + r.proxies = append(r.proxies, proxy) + q := &guest{id: id, host: h, target: et, endpoint: u, identity: key, proxy: proxy} + r.clientsMu.Lock() + r.guests = append(r.guests, q) + r.clientsMu.Unlock() + } + } + if err := r.admissions(len(r.guests), func(i int) error { + q := r.guests[i] + began := time.Now() + if err := q.connect(r.ctx, r); err != nil { + return err + } + r.rec.mu.Lock() + r.rec.summary.EstablishedGuests++ + r.rec.summary.EstablishedChannels += len(q.channels) + r.rec.mu.Unlock() + r.rec.event("admission", map[string]any{"role": "guest", "guest": q.id, "ready_ms": float64(time.Since(began)) / float64(time.Millisecond)}) + return nil + }); err != nil { + return err + } + if r.cfg.idleHostOnly() { + return r.hostCanary(0) + } + return nil +} +func (r *runState) cleanup() error { + r.cancel(context.Canceled) + r.workers.Wait() + for _, q := range r.guests { + q.setExpected(true) + q.close() + } + for _, p := range r.proxies { + p.close() + } + for _, h := range r.hosts { + if h != nil { + <-h.done + } + } + if err := os.RemoveAll(r.dir); err != nil { + return err + } + _, err := os.Stat(r.dir) + complete := os.IsNotExist(err) + r.rec.mu.Lock() + r.rec.summary.CleanupComplete = complete + r.rec.mu.Unlock() + if !complete { + return errors.New("temporary directory still exists after cleanup") + } + return nil +} + +func watchFiles(c Config, now time.Time) error { + if c.WatchdogFile != "" { + if _, err := os.Stat(c.WatchdogFile); err == nil { + return errors.New("external_watchdog_stop") + } else if !os.IsNotExist(err) { + return err + } + } + if c.WatchdogHeartbeatFile != "" { + info, err := os.Stat(c.WatchdogHeartbeatFile) + if err != nil || now.Sub(info.ModTime()) > 30*time.Second || info.ModTime().After(now.Add(time.Second)) { + return errors.New("external_watchdog_heartbeat_lost") + } + } + return nil +} +func (r *runState) telemetry() { + r.workers.Add(1) + go func() { + defer r.workers.Done() + clients := map[string]*http.Client{} + + for _, t := range r.cfg.Manifest.Targets { + target := t + tr := &http.Transport{Proxy: nil, DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + return r.guard.connect(ctx, target, addr) + }} + defer tr.CloseIdleConnections() + clients[t.MachineID] = &http.Client{Transport: tr, Timeout: 3 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return errors.New("telemetry redirect refused") }} + r.telemetryMu.Lock() + r.lastTelemetry[t.MachineID] = time.Now() + r.telemetryMu.Unlock() + } + tick := time.NewTicker(5 * time.Second) + defer tick.Stop() + collect := func() { + for _, t := range r.cfg.Manifest.Targets { + req, err := http.NewRequestWithContext(r.ctx, http.MethodGet, t.TelemetryURL, nil) + if err != nil { + r.fail(err) + return + } + resp, err := clients[t.MachineID].Do(req) + if err == nil { + body, e := io.ReadAll(io.LimitReader(resp.Body, (1<<20)+1)) + _ = resp.Body.Close() + if e == nil && resp.StatusCode == 200 && len(body) <= 1<<20 { + r.telemetryMu.Lock() + r.lastTelemetry[t.MachineID] = time.Now() + r.telemetryMu.Unlock() + r.rec.event("relay_metrics", map[string]any{"machine_id": t.MachineID, "prometheus": string(body)}) + } else { + err = errors.New("invalid telemetry response") + } + } + if err != nil { + r.rec.event("telemetry_error", map[string]any{"machine_id": t.MachineID, "error": err.Error()}) + } + + } + var m runtime.MemStats + runtime.ReadMemStats(&m) + fields := map[string]any{"pid": os.Getpid(), "goroutines": runtime.NumGoroutine(), "heap_alloc": m.HeapAlloc, "heap_sys": m.HeapSys, "go_sys": m.Sys, "gc_count": m.NumGC, "go_version": runtime.Version()} + if status, err := os.ReadFile("/proc/self/status"); err == nil { + fields["proc_status"] = string(status) + } + if stat, err := os.ReadFile("/proc/self/stat"); err == nil { + fields["proc_stat"] = string(stat) + } + r.rec.event("generator", fields) + } + collect() + for { + select { + case <-r.ctx.Done(): + return + case <-tick.C: + collect() + } + } + }() +} + +func Run(ctx context.Context, c Config, out io.Writer) (result Summary, err error) { + if err = c.Validate(time.Now()); err != nil { + return result, err + } + if err = watchFiles(c, time.Now()); err != nil { + return result, err + } + rec := newRecorder(out, c) + started := time.Now() + deadline := c.Manifest.ExpiresAt + bounded, stop := context.WithDeadline(ctx, deadline) + defer stop() + runctx, cancel := context.WithCancelCause(bounded) + r := &runState{lastTelemetry: map[string]time.Time{}, cfg: c, guard: newGuard(c), rec: rec, ctx: runctx, cancel: cancel} + r.helper = c.HelperExecutable + if r.helper == "" { + r.helper, err = os.Executable() + if err != nil { + return result, err + } + } + r.dir, err = os.MkdirTemp("", "upterm-cap-") + if err != nil { + return result, err + } + defer func() { + if e := r.cleanup(); e != nil { + err = errors.Join(err, e) + rec.mu.Lock() + rec.summary.StopReason = "cleanup_failure" + rec.mu.Unlock() + } + rec.mu.Lock() + r.guard.mu.Lock() + rec.summary.PeakExternalConnections = r.guard.peak + r.guard.mu.Unlock() + rec.summary.DurationSeconds = time.Since(started).Seconds() + sort.Float64s(rec.summary.RTTMillis) + samples := rec.summary.RTTMillis + if len(samples) > 0 { + rec.summary.RTTP50Millis = samples[(len(samples)-1)*50/100] + rec.summary.RTTP95Millis = samples[(len(samples)-1)*95/100] + rec.summary.RTTP99Millis = samples[(len(samples)-1)*99/100] + } + result = rec.summary + rec.mu.Unlock() + rec.event("summary", map[string]any{"summary": result}) + if rec.err != nil && err == nil { + err = rec.err + } + }() + rec.event("phase_start", map[string]any{"config": safeConfig(c), "input_frame_bytes": 128, "host_output_frame_bytes": 1024, "framing_header_bytes": 44, "framing_newline_bytes": 1, "wire_overhead": "not measured", "duration_includes_admission": false}) + var identity ssh.Signer + if c.HostIdentityFile != "" { + key, e := os.ReadFile(c.HostIdentityFile) + if e != nil { + err = errors.New("cannot read host identity") + } else { + identity, e = ssh.ParsePrivateKey(key) + if e != nil { + err = errors.New("invalid host identity") + } + } + if err != nil { + rec.summary.StopReason = "identity_error" + return result, err + } + } + if err = r.reserve(c.hostOutputReservation()); err != nil { + rec.summary.StopReason = "payload_budget" + return result, err + } + r.telemetry() + r.watch() + admissionTimer := time.AfterFunc(time.Duration(c.Hosts*(1+c.GuestsPerHost))*time.Second/time.Duration(c.AdmissionsPerSecond)+30*time.Second, func() { r.fail(errors.New("admission_deadline")) }) + defer admissionTimer.Stop() + if err = r.admit(identity); err != nil { + rec.mu.Lock() + rec.summary.AdmissionFailures++ + rec.summary.StopReason = "admission_failure" + rec.mu.Unlock() + return result, err + } + admissionTimer.Stop() + steadyStarted := time.Now() + rec.mu.Lock() + rec.summary.SetupSeconds = steadyStarted.Sub(started).Seconds() + rec.mu.Unlock() + steady := time.AfterFunc(time.Duration(c.DurationSeconds)*time.Second, func() { r.fail(errDuration) }) + defer steady.Stop() + defer func() { + rec.mu.Lock() + rec.summary.SteadySeconds = time.Since(steadyStarted).Seconds() + rec.mu.Unlock() + }() + rec.event("steady_start", map[string]any{"duration_seconds": c.DurationSeconds}) + if c.idleHostOnly() { + r.canaries() + } + if c.Active { + for _, q := range r.guests { + if e := q.active(r); e != nil { + r.fail(e) + break + } + } + } + if c.TransferConcurrency > 0 { + r.workers.Add(1) + go func() { + defer r.workers.Done() + if e := r.transfers(); e != nil { + r.fail(e) + } + }() + } + if c.ChurnPercent > 0 { + r.workers.Add(1) + go func() { + defer r.workers.Done() + if e := r.churn(); e != nil { + r.fail(e) + } + }() + } + tick := time.NewTicker(time.Second) + defer tick.Stop() + for { + select { + case <-runctx.Done(): + cause := context.Cause(runctx) + rec.mu.Lock() + rec.summary.StopReason = cause.Error() + err = cause + if errors.Is(cause, errDuration) { + err = nil + if c.TransferConcurrency > 0 && rec.summary.Transfers != c.TransferConcurrency { + err = errors.New("incomplete_transfers") + } + if c.ChurnPercent > 0 && !rec.summary.ChurnCompleted { + err = errors.New("incomplete_churn") + } + if c.Active && rec.summary.ActiveFrames+rec.summary.PausedActiveFrames != int64(c.Hosts*c.GuestsPerHost*c.DurationSeconds) { + err = errors.New("offered_workload_incomplete") + } + if err != nil { + rec.summary.StopReason = err.Error() + } + } else if errors.Is(cause, context.DeadlineExceeded) && ctx.Err() == nil { + rec.summary.StopReason = "manifest_expired" + } + rec.mu.Unlock() + return result, err + case <-tick.C: + if c.Active && time.Since(steadyStarted) < time.Duration(c.DurationSeconds)*time.Second { + for _, q := range r.guests { + if e := q.active(r); e != nil { + r.fail(e) + break + } + } + } + } + } +} +func safeConfig(c Config) any { + copy := c + copy.Manifest = Manifest{} + copy.HostIdentityFile = "" + copy.HelperExecutable = "" + return copy +} +func telemetryLost(last, now time.Time) bool { return now.Sub(last) > 30*time.Second } +func (r *runState) watch() { + r.workers.Add(1) + go func() { + defer r.workers.Done() + tick := time.NewTicker(time.Second) + defer tick.Stop() + for { + select { + case <-r.ctx.Done(): + return + case now := <-tick.C: + if err := watchFiles(r.cfg, now); err != nil { + r.fail(err) + return + } + r.telemetryMu.Lock() + lost := false + for _, last := range r.lastTelemetry { + lost = lost || telemetryLost(last, now) + } + r.telemetryMu.Unlock() + if lost { + r.fail(errors.New("telemetry_lost")) + return + } + r.clientsMu.RLock() + guests := append([]*guest(nil), r.guests...) + r.clientsMu.RUnlock() + for _, q := range guests { + q.mu.Lock() + stale := false + if !q.expected { + for _, sent := range q.pending { + stale = stale || now.Sub(sent) > 10*time.Second + } + } + q.mu.Unlock() + if stale { + r.fail(errors.New("terminal_response_timeout")) + return + } + } + r.rec.mu.Lock() + err := r.rec.err + r.rec.mu.Unlock() + if err != nil { + r.fail(err) + return + } + } + } + }() +} +func (r *runState) admissions(count int, work func(int) error) error { + if count == 0 { + return r.ctx.Err() + } + workers := min(count, 20) + jobs := make(chan int) + failed := make(chan error, 1) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-r.ctx.Done(): + return + case index, ok := <-jobs: + if !ok { + return + } + if err := work(index); err != nil { + select { + case failed <- err: + default: + } + r.fail(err) + return + } + } + } + }() + } +offer: + for i := 0; i < count; i++ { + select { + case jobs <- i: + case <-r.ctx.Done(): + break offer + } + } + close(jobs) + wg.Wait() + select { + case err := <-failed: + return err + default: + return context.Cause(r.ctx) + } +} diff --git a/internal/capacity/runner_test.go b/internal/capacity/runner_test.go new file mode 100644 index 000000000..8544e5f7c --- /dev/null +++ b/internal/capacity/runner_test.go @@ -0,0 +1,317 @@ +package capacity + +import ( + "bytes" + "context" + "crypto/ed25519" + "crypto/rand" + "encoding/pem" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/go-kit/kit/metrics/provider" + "github.com/owenthereal/upterm/routing" + "github.com/owenthereal/upterm/server" + "golang.org/x/crypto/ssh" +) + +var helperDir string + +func TestMain(m *testing.M) { + var err error + helperDir, err = os.MkdirTemp("", "capacity-test-") + if err != nil { + panic(err) + } + code := m.Run() + _ = os.RemoveAll(helperDir) + os.Exit(code) +} + +var helperOnce sync.Once +var helperBinary string +var helperError error + +func buildHelper(t *testing.T) string { + t.Helper() + helperOnce.Do(func() { + helperBinary = filepath.Join(helperDir, "capacity") + cmd := exec.Command("go", "build", "-o", helperBinary, "./cmd") + if out, err := cmd.CombinedOutput(); err != nil { + helperError = fmt.Errorf("helper build: %w: %s", err, out) + } + }) + if helperError != nil { + t.Fatal(helperError) + } + return helperBinary +} +func localRelay(t *testing.T, options ...func(*server.Server)) Config { + t.Helper() + clearProxies(t) + sshln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + wsln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + key := signer(t) + cert, err := (&server.HostCertSigner{Hostnames: []string{"127.0.0.1"}}).SignCert(key) + if err != nil { + t.Fatal(err) + } + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + network := &server.MemoryProvider{} + if err = network.SetOpts(nil); err != nil { + t.Fatal(err) + } + sm, err := server.NewSessionManager(routing.ModeEmbedded, server.WithSessionManagerLogger(logger)) + if err != nil { + t.Fatal(err) + } + relay := &server.Server{NodeAddr: sshln.Addr().String(), HostSigners: []ssh.Signer{cert}, Signers: []ssh.Signer{key}, NetworkProvider: network, SessionManager: sm, MetricsProvider: provider.NewDiscardProvider(), Logger: logger} + for _, option := range options { + option(relay) + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- relay.ServeListeners(ctx, server.Listeners{SSH: sshln, WS: wsln}) }() + t.Cleanup(func() { + cancel() + _ = relay.Shutdown() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Error("relay leaked") + } + }) + metrics := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, "test_metric 1\n") })) + t.Cleanup(metrics.Close) + c := validConfig(t) + c.HelperExecutable = buildHelper(t) + c.Manifest.Targets[0].SSHURL = "ssh://" + sshln.Addr().String() + c.Manifest.Targets[0].WSURL = "ws://" + wsln.Addr().String() + c.Manifest.Targets[0].NodeAddresses = []string{sshln.Addr().String()} + c.Manifest.Targets[0].HostKeys = []string{string(ssh.MarshalAuthorizedKey(key.PublicKey()))} + c.Manifest.Targets[0].TelemetryURL = metrics.URL + "/metrics" + c.Active = true + return c +} +func TestRealSessionsAcrossTransports(t *testing.T) { + for _, mix := range []struct { + name string + h, g int + }{{"ssh", 0, 0}, {"ws", 100, 100}, {"ssh-to-ws", 0, 100}, {"ws-to-ssh", 100, 0}} { + t.Run(mix.name, func(t *testing.T) { + c := localRelay(t) + c.HostWSPercent = mix.h + c.GuestWSPercent = mix.g + c.ChannelsPerGuest = 4 + c.DurationSeconds = 3 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if err != nil { + t.Fatalf("Run: %v\n%s", err, out.String()) + } + if s.EstablishedHosts != 1 || s.EstablishedGuests != 1 || s.EstablishedChannels != 4 { + t.Fatalf("real readiness not counted: %+v", s) + } + if s.RTTCount < 4 || s.ReceivedFrames < 4 || s.PayloadBytes == 0 { + t.Fatalf("usable terminal traffic missing: %+v", s) + } + if s.SteadySeconds < 3 || s.InputFrames != 7 { + t.Fatalf("offered workload or hold was shortened: %+v", s) + } + if s.StopReason != "duration" { + t.Fatalf("unexpected stop: %+v", s) + } + }) + } +} +func TestFailedReadinessNeverCountsEstablished(t *testing.T) { + c := localRelay(t) + c.HelperExecutable = "/does/not/exist" + s, err := Run(t.Context(), c, io.Discard) + if err == nil { + t.Fatal("failed helper readiness reported success") + } + if s.EstablishedHosts != 0 || s.EstablishedGuests != 0 { + t.Fatalf("failed readiness counted: %+v", s) + } +} + +func TestRealSFTPTransfersHaveVerifiedBytes(t *testing.T) { + c := localRelay(t) + c.Active = false + c.DurationSeconds = 6 + c.TransferConcurrency = 4 + c.PayloadBudgetBytes = 256 << 20 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if err != nil { + t.Fatalf("Run: %v\n%s", err, out.String()) + } + if s.Transfers != 4 || s.TransferBytes != 134217728 { + t.Fatalf("four verified upload/download pairs missing: %+v", s) + } + if !bytes.Contains(out.Bytes(), []byte(`"sha256"`)) { + t.Fatal("transfer hashes absent from evidence") + } +} +func TestRunCancellationJoinsOwnedWork(t *testing.T) { + c := localRelay(t) + c.DurationSeconds = 120 + ctx, cancel := context.WithCancel(t.Context()) + time.AfterFunc(500*time.Millisecond, cancel) + start := time.Now() + s, err := Run(ctx, c, io.Discard) + if err == nil || s.StopReason == "duration" { + t.Fatalf("cancel was reported successful: %+v %v", s, err) + } + if !s.CleanupComplete { + t.Fatal("cancellation did not verify temporary directory cleanup") + } + if time.Since(start) > 8*time.Second { + t.Fatal("owned work did not join promptly") + } +} +func TestInvalidLimitsNeverStartHostTransport(t *testing.T) { + clearProxies(t) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + calls := make(chan struct{}, 1) + done := make(chan struct{}) + go func() { + defer close(done) + for { + c, e := ln.Accept() + if e != nil { + return + } + select { + case calls <- struct{}{}: + default: + } + _ = c.Close() + } + }() + c := validConfig(t) + c.Manifest.Targets[0].SSHURL = "ssh://" + ln.Addr().String() + c.AdmissionsPerSecond = 21 + if _, err := Run(t.Context(), c, io.Discard); err == nil { + t.Fatal("invalid rate admitted host") + } + _ = ln.Close() + <-done + select { + case <-calls: + t.Fatal("invalid limits started a host transport") + default: + } +} + +func TestSuppliedHostIdentityAuthenticatesRestrictedRelay(t *testing.T) { + _, private, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatal(err) + } + identity, err := ssh.NewSignerFromKey(private) + if err != nil { + t.Fatal(err) + } + block, err := ssh.MarshalPrivateKey(private, "capacity-test") + if err != nil { + t.Fatal(err) + } + dir := t.TempDir() + keyFile := filepath.Join(dir, "identity") + authorized := filepath.Join(dir, "authorized_keys") + if err = os.WriteFile(keyFile, pem.EncodeToMemory(block), 0600); err != nil { + t.Fatal(err) + } + if err = os.WriteFile(authorized, ssh.MarshalAuthorizedKey(identity.PublicKey()), 0600); err != nil { + t.Fatal(err) + } + c := localRelay(t, func(s *server.Server) { s.AuthorizedKeysFiles = []string{authorized} }) + c.HostIdentityFile = keyFile + c.DurationSeconds = 1 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if err != nil { + t.Fatalf("restricted relay: %v", err) + } + if s.EstablishedHosts != 1 || !s.CleanupComplete { + t.Fatalf("identity/cleanup missing: %+v", s) + } + if bytes.Contains(out.Bytes(), []byte(keyFile)) || bytes.Contains(out.Bytes(), []byte("BEGIN OPENSSH PRIVATE KEY")) { + t.Fatal("private identity leaked into events") + } +} +func TestIdleHostOnlyPhaseProvesSparseGuestCanary(t *testing.T) { + c := localRelay(t) + c.Hosts = 2 + c.GuestsPerHost = 0 + c.Active = false + c.DurationSeconds = 1 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if err != nil { + t.Fatalf("host-only: %v\n%s", err, out.String()) + } + if s.OfferedGuests != 0 || s.EstablishedGuests != 0 { + t.Fatalf("transient canary changed nominal guests: %+v", s) + } + if s.CanaryOffered != 1 || s.CanaryEstablished != 1 || s.CanaryReadyMillis <= 0 { + t.Fatalf("host-only usability was not proved: %+v", s) + } + if !s.CleanupComplete { + t.Fatal("canary transport cleanup failed") + } +} +func TestIdleCanariesReuseOneSlotAndRotateHosts(t *testing.T) { + c := localRelay(t) + c.Hosts = 2 + c.GuestsPerHost = 0 + c.Active = false + ctx, cancel := context.WithCancelCause(t.Context()) + dir, err := os.MkdirTemp("", "capacity-canary-") + if err != nil { + t.Fatal(err) + } + r := &runState{cfg: c, ctx: ctx, cancel: cancel, guard: newGuard(c), rec: newRecorder(io.Discard, c), helper: c.HelperExecutable, dir: dir} + defer func() { + if err := r.cleanup(); err != nil { + t.Error(err) + } + }() + if err = r.admit(nil); err != nil { + t.Fatal(err) + } + if err = r.hostCanary(1); err != nil { + t.Fatal(err) + } + if len(r.guests) != 1 || r.canary.host.id != 1 || r.rec.summary.CanaryEstablished != 2 || r.rec.summary.CanaryRTTCount != 2 { + t.Fatalf("canaries did not rotate within one slot: guests=%d summary=%+v", len(r.guests), r.rec.summary) + } + r.guard.mu.Lock() + active := r.guard.active + r.guard.mu.Unlock() + if active != 2 { + t.Fatalf("transient canary transport leaked: active=%d want two hosts", active) + } +} diff --git a/internal/capacity/safety_test.go b/internal/capacity/safety_test.go new file mode 100644 index 000000000..2331fb009 --- /dev/null +++ b/internal/capacity/safety_test.go @@ -0,0 +1,98 @@ +package capacity + +import ( + "context" + "errors" + "io" + "net" + "os" + "path/filepath" + "testing" + "time" +) + +func TestTelemetryLossUsesThirtySecondBound(t *testing.T) { + now := time.Unix(100, 0) + if telemetryLost(now, now.Add(30*time.Second)) { + t.Fatal("telemetry stopped before grace expired") + } + if !telemetryLost(now, now.Add(30*time.Second+time.Nanosecond)) { + t.Fatal("telemetry loss did not stop after 30 seconds") + } +} +func TestWatchdogStopAndStaleHeartbeat(t *testing.T) { + c := validConfig(t) + dir := t.TempDir() + c.WatchdogFile = filepath.Join(dir, "stop") + c.WatchdogHeartbeatFile = filepath.Join(dir, "heartbeat") + now := time.Now() + if err := os.WriteFile(c.WatchdogHeartbeatFile, nil, 0600); err != nil { + t.Fatal(err) + } + if err := watchFiles(c, now); err != nil { + t.Fatal(err) + } + if err := watchFiles(c, now.Add(31*time.Second)); err == nil { + t.Fatal("stale watchdog heartbeat accepted") + } + if err := os.WriteFile(c.WatchdogFile, nil, 0600); err != nil { + t.Fatal(err) + } + if _, err := Run(t.Context(), c, io.Discard); err == nil { + t.Fatal("watchdog stop admitted workload") + } +} +func TestAdmissionRateIncludesRetriesAndCancellation(t *testing.T) { + c := validConfig(t) + g := newGuard(c) + var starts []time.Time + g.dial = func(context.Context, string, string) (net.Conn, error) { + starts = append(starts, time.Now()) + return nil, errors.New("failed dial") + } + for i := 0; i < 3; i++ { + _, _ = g.connect(t.Context(), c.Manifest.Targets[0], "127.0.0.1:2222") + } + if len(starts) != 3 || starts[2].Sub(starts[0]) < 95*time.Millisecond { + t.Fatalf("retry handshakes escaped limiter: %v", starts) + } + ctx, cancel := context.WithCancel(t.Context()) + cancel() + _, _ = g.connect(ctx, c.Manifest.Targets[0], "127.0.0.1:2222") + if len(starts) != 3 { + t.Fatal("cancelled admission dialed") + } +} +func TestPayloadStopReservesBeforeWriting(t *testing.T) { + c := validConfig(t) + c.PayloadBudgetBytes = 255 + r := &runState{cfg: c, rec: newRecorder(io.Discard, c)} + q := &guest{} + sink := &countWriter{} + if err := r.send(q, sink, 0, 0); err == nil { + t.Fatal("over-budget input was sent") + } + if sink.n != 0 { + t.Fatalf("over-budget input emitted %d bytes", sink.n) + } +} + +type countWriter struct{ n int } + +func (w *countWriter) Write(p []byte) (int, error) { w.n += len(p); return len(p), nil } +func TestTelemetryLossCancelsRunningPhase(t *testing.T) { + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(context.Canceled) + c := validConfig(t) + r := &runState{cfg: c, ctx: ctx, cancel: cancel, rec: newRecorder(io.Discard, c), lastTelemetry: map[string]time.Time{"test-machine": time.Now().Add(-31 * time.Second)}} + r.watch() + defer func() { cancel(context.Canceled); r.workers.Wait() }() + select { + case <-ctx.Done(): + if context.Cause(ctx).Error() != "telemetry_lost" { + t.Fatalf("wrong stop cause %v", context.Cause(ctx)) + } + case <-time.After(2 * time.Second): + t.Fatal("telemetry loss did not cancel the phase") + } +} diff --git a/internal/capacity/transfer.go b/internal/capacity/transfer.go new file mode 100644 index 000000000..d04789998 --- /dev/null +++ b/internal/capacity/transfer.go @@ -0,0 +1,313 @@ +package capacity + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "path/filepath" + "sync" + "time" + + "github.com/pkg/sftp" +) + +type deterministicReader struct{ offset int64 } + +func (d *deterministicReader) Read(p []byte) (int, error) { + for i := range p { + p[i] = byte((d.offset + int64(i)) % 251) + } + d.offset += int64(len(p)) + return len(p), nil +} +func deterministicHash() string { + h := sha256.New() + _, _ = io.CopyN(h, &deterministicReader{}, TransferSize) + return hex.EncodeToString(h.Sum(nil)) +} +func (r *runState) transfers() error { + expected := deterministicHash() + var wg sync.WaitGroup + results := make(chan error, r.cfg.TransferConcurrency) + for i := 0; i < r.cfg.TransferConcurrency; i++ { + wg.Add(1) + go func(index int) { defer wg.Done(); results <- r.transfer(index, expected) }(i) + } + wg.Wait() + close(results) + for err := range results { + if err != nil { + return err + } + } + return nil +} +func (r *runState) transfer(index int, expected string) (err error) { + q := r.guests[index%len(r.guests)] + q.ops.Lock() + client := q.client + q.ops.Unlock() + c, err := sftp.NewClient(client) + if err != nil { + return err + } + defer c.Close() + path := filepath.ToSlash(filepath.Join(filepath.Dir(q.host.socket), fmt.Sprintf("transfer-%d.bin", index))) + defer func() { _ = c.Remove(path) }() + began := time.Now() + var uploaded, downloaded int64 + defer func() { + fields := map[string]any{"transfer": index, "guest": q.id, "upload_bytes": uploaded, "download_bytes": downloaded, "sha256": expected, "verified": err == nil, "duration_seconds": time.Since(began).Seconds()} + if err != nil { + fields["error"] = err.Error() + } + r.rec.event("transfer", fields) + }() + if err = r.reserve(2 * TransferSize); err != nil { + return err + } + f, err := c.Create(path) + if err != nil { + return err + } + uploaded, err = io.CopyN(f, &deterministicReader{}, TransferSize) + closeErr := f.Close() + r.actual(uploaded) + if err != nil { + return err + } + if closeErr != nil { + return closeErr + } + if uploaded != TransferSize { + return errors.New("upload length mismatch") + } + f, err = c.Open(path) + if err != nil { + return err + } + h := sha256.New() + downloaded, err = io.Copy(h, io.LimitReader(f, TransferSize+1)) + closeErr = f.Close() + r.actual(downloaded) + if err != nil { + return err + } + if closeErr != nil { + return closeErr + } + got := hex.EncodeToString(h.Sum(nil)) + if downloaded != TransferSize || got != expected { + return errors.New("download hash/length mismatch") + } + r.rec.mu.Lock() + r.rec.summary.Transfers++ + r.rec.summary.TransferBytes += uploaded + downloaded + r.rec.mu.Unlock() + return nil +} + +// Churn cuts existing sockets. The same Host.Run (and session ID), guest key, +// role and transport survive; fresh guest TCP/SSH transports are reported as +// reconnects, separately from first admission. +func (r *runState) churn() error { + actions := churnSchedule(r.cfg.Hosts, r.cfg.GuestsPerHost, r.cfg.ChurnPercent) + if len(actions) == 0 { + return errors.New("churn percentage selects no transports") + } + affectedHosts := map[int]bool{} + affectedGuests := map[int]bool{} + for _, a := range actions { + if a.host { + affectedHosts[a.index] = true + } else { + affectedGuests[a.index] = true + } + } + var canary *guest + for _, q := range r.guests { + if !affectedHosts[q.host.id] && !affectedGuests[q.id] { + canary = q + break + } + } + if canary == nil { + return errors.New("churn has no survivor canary") + } + canary.mu.Lock() + before := canary.acks + canary.mu.Unlock() + start := time.Now().Add(time.Second) + var wg sync.WaitGroup + results := make(chan error, len(actions)) + for _, a := range actions { + if err := pause(r.ctx, time.Until(start.Add(a.offset))); err != nil { + wg.Wait() + return err + } + wg.Add(1) + go func(action churnAction) { + defer wg.Done() + var err error + if action.host { + err = r.cutHost(r.hosts[action.index]) + } else { + err = r.cutGuest(r.guests[action.index]) + } + if err != nil { + r.fail(err) + } + results <- err + }(a) + } + wg.Wait() + close(results) + for err := range results { + if err != nil { + return err + } + } + if !r.cfg.Active { + if err := canary.active(r); err != nil { + return err + } + if err := pause(r.ctx, time.Second); err != nil { + return err + } + } + canary.mu.Lock() + after := canary.acks + canary.mu.Unlock() + if after <= before { + return errors.New("survivor canary stopped responding") + } + r.rec.mu.Lock() + r.rec.summary.ChurnCompleted = true + r.rec.mu.Unlock() + r.rec.event("survivor_canary", map[string]any{"guest": canary.id, "host": canary.host.id, "responses_before": before, "responses_after": after, "transport_replaced": false}) + return nil +} +func (r *runState) restoreGuest(q *guest, began time.Time) error { + q.close() + q.mu.Lock() + q.pending = map[string]time.Time{} + q.generation++ + q.mu.Unlock() + var last error + for tries := 0; tries < 3; tries++ { + last = q.connect(r.ctx, r) + if last == nil { + q.setExpected(false) + r.rec.mu.Lock() + r.rec.summary.Reconnects++ + r.rec.mu.Unlock() + r.rec.event("guest_restored", map[string]any{"guest": q.id, "host": q.host.id, "generation": q.generation, "restoration_ms": float64(time.Since(began)) / float64(time.Millisecond), "same_identity": true}) + return nil + } + q.close() + if err := pause(r.ctx, time.Second); err != nil { + return err + } + } + return fmt.Errorf("guest restoration failed: %w", last) +} +func (r *runState) cutGuest(q *guest) error { + q.host.repair.Lock() + defer q.host.repair.Unlock() + q.setExpected(true) + began := time.Now() + r.rec.event("transport_cut", map[string]any{"role": "guest", "guest": q.id, "host": q.host.id}) + q.ops.Lock() + if q.raw != nil { + _ = q.raw.Close() + } + q.ops.Unlock() + return r.restoreGuest(q, began) +} +func (r *runState) cutHost(h *hostRun) error { + h.repair.Lock() + defer h.repair.Unlock() + var guests []*guest + for _, q := range r.guests { + if q.host == h { + q.setExpected(true) + guests = append(guests, q) + } + } + began := time.Now() + r.rec.event("transport_cut", map[string]any{"role": "host", "host": h.id}) + h.proxy.cut() + ctx, cancel := context.WithTimeout(r.ctx, 15*time.Second) + defer cancel() + down := false + for { + select { + case e := <-h.tunnel: + if !e.Up { + down = true + if !e.CanReconnect { + return errors.New("relay cannot restore host") + } + } + if e.Up && down { + s, err := h.session(ctx) + if err != nil { + return err + } + if err = r.guard.session(ctx, s); err != nil { + return err + } + r.rec.mu.Lock() + r.rec.summary.Reconnects++ + r.rec.mu.Unlock() + r.rec.event("host_restored", map[string]any{"host": h.id, "session_id": s.SessionId, "restoration_ms": float64(time.Since(began)) / float64(time.Millisecond), "same_host_run": true}) + for _, q := range guests { + if err = r.restoreGuest(q, began); err != nil { + return err + } + } + return nil + } + case <-ctx.Done(): + return fmt.Errorf("host restoration: %w", ctx.Err()) + } + } +} + +func pause(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +type churnAction struct { + host bool + index int + offset time.Duration +} + +func churnSchedule(hosts, guests, percent int) []churnAction { + hn := hosts * percent / 100 + gn := hosts * guests * percent / 100 + var actions []churnAction + for i := 0; i < hn; i++ { + actions = append(actions, churnAction{host: true, index: i * hosts / hn}) + } + for i := 0; i < gn; i++ { + actions = append(actions, churnAction{index: i * hosts * guests / gn}) + } + for i := range actions { + if len(actions) > 1 { + actions[i].offset = time.Duration(i) * 10 * time.Second / time.Duration(len(actions)-1) + } + } + return actions +} diff --git a/internal/capacity/workload.go b/internal/capacity/workload.go new file mode 100644 index 000000000..661445a4f --- /dev/null +++ b/internal/capacity/workload.go @@ -0,0 +1,175 @@ +package capacity + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "os" + "strconv" + "strings" + "time" + + "golang.org/x/term" +) + +// A frame's entire fixed length is application payload: its header is 44 +// bytes (including the following space), its final byte is LF, the rest 'x'. +func frame(kind string, id, seq int, at time.Time, size int) []byte { + header := fmt.Sprintf("CAP %s %08d %08d %019d ", kind, id, seq, at.UnixNano()) + return []byte(header + strings.Repeat("x", size-len(header)-1) + "\n") +} + +type observation struct { + kind string + id, seq int + sent time.Time +} + +func parseFrame(line string) (observation, bool) { + i := strings.Index(line, "CAP ") + if i < 0 { + return observation{}, false + } + fields := strings.Fields(line[i:]) + if len(fields) != 6 || (fields[1] != "I" && fields[1] != "H" && fields[1] != "R" && fields[1] != "A") { + return observation{}, false + } + id, e1 := strconv.Atoi(fields[2]) + seq, e2 := strconv.Atoi(fields[3]) + ns, e3 := strconv.ParseInt(fields[4], 10, 64) + if e1 != nil || e2 != nil || e3 != nil { + return observation{}, false + } + return observation{fields[1], id, seq, time.Unix(0, ns)}, true +} + +// Helper is the runner's harmless hosted command. It echoes fixed 128-byte +// input frames and, in active mode, emits exactly 1 KiB every ten seconds. +// It never starts a shell or interprets received text as a command. +func Helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active bool) error { + if f, ok := in.(*os.File); ok && term.IsTerminal(int(f.Fd())) { + state, err := term.MakeRaw(int(f.Fd())) + if err != nil { + return err + } + defer term.Restore(int(f.Fd()), state) + } + if _, err := io.WriteString(out, "CAP READY\n"); err != nil { + return err + } + ctx, cancel := context.WithCancel(ctx) + defer cancel() + received := make(chan []byte) + failed := make(chan error, 1) + done := make(chan struct{}) + go func() { + defer close(done) + for { + b := make([]byte, 128) + _, err := io.ReadFull(in, b) + if err != nil { + failed <- err + return + } + select { + case received <- b: + case <-ctx.Done(): + return + } + } + }() + defer func() { + cancel() + if closer, ok := in.(io.Closer); ok { + _ = closer.Close() + } + <-done + }() + tick := time.NewTicker(10 * time.Second) + defer tick.Stop() + seq := 0 + for { + select { + case <-ctx.Done(): + return ctx.Err() + case err := <-failed: + if errors.Is(err, io.EOF) { + return nil + } + return err + case b := <-received: + if len(b) > 4 && b[4] == 'R' { + b[4] = 'A' + } + if _, err := out.Write(b); err != nil { + return err + } + case now := <-tick.C: + if active { + if _, err := out.Write(frame("H", hostID, seq, now, 1024)); err != nil { + return err + } + seq++ + } + } + } +} + +type sequenceTracker map[string]int + +func (s sequenceTracker) observe(o observation) error { + key := o.kind + ":" + strconv.Itoa(o.id) + if last, ok := s[key]; ok && o.seq != last+1 { + return fmt.Errorf("sequence gap/reordering %s: previous=%d current=%d", key, last, o.seq) + } + s[key] = o.seq + return nil +} +func scanFrames(in io.Reader, consume func(observation) error, helperReady func(observation)) error { + scanner := bufio.NewScanner(in) + scanner.Buffer(make([]byte, 4096), 1<<20) + for scanner.Scan() { + + if o, ok := parseFrame(scanner.Text()); ok { + if o.kind == "R" { + continue + } + if o.kind == "A" { + helperReady(o) + continue + } + if err := consume(o); err != nil { + return err + } + } + } + if err := scanner.Err(); err != nil { + return err + } + return io.EOF +} +func (r *runState) observe(tracker sequenceTracker, epochs map[string]int, o observation) error { + if o.kind == "I" { + r.clientsMu.RLock() + if o.id < 0 || o.id/4 >= len(r.guests) { + r.clientsMu.RUnlock() + return errors.New("unknown input source") + } + peer := r.guests[o.id/4] + r.clientsMu.RUnlock() + peer.mu.Lock() + generation, expected := peer.generation, peer.expected + peer.mu.Unlock() + key := "I:" + strconv.Itoa(o.id) + if old, ok := epochs[key]; expected || !ok || old != generation { + delete(tracker, key) + } + epochs[key] = generation + if expected { + return nil + } + } + return tracker.observe(o) +} diff --git a/internal/capacity/workload_test.go b/internal/capacity/workload_test.go new file mode 100644 index 000000000..cc92799b0 --- /dev/null +++ b/internal/capacity/workload_test.go @@ -0,0 +1,50 @@ +package capacity + +import ( + "bytes" + "context" + "strings" + "testing" + "time" +) + +func TestFixedFramesAndSequenceGaps(t *testing.T) { + b := frame("I", 7, 1, time.Unix(1, 2), 128) + if len(b) != 128 || !strings.HasPrefix(string(b), "CAP I 00000007 00000001 0000000001000000002 ") { + t.Fatalf("incorrect offered frame %q", b) + } + o, ok := parseFrame(string(b)) + if !ok || o.id != 7 || o.seq != 1 { + t.Fatal("frame lost its tag") + } + tracker := sequenceTracker{} + if err := tracker.observe(o); err != nil { + t.Fatal(err) + } + o.seq = 3 + if tracker.observe(o) == nil { + t.Fatal("missing frame went undetected") + } +} +func TestIntentionalReconnectResetsOnlyAffectedSource(t *testing.T) { + peer := &guest{generation: 1} + r := &runState{guests: []*guest{peer}} + tracker := sequenceTracker{"I:0": 1} + epochs := map[string]int{"I:0": 0} + if err := r.observe(tracker, epochs, observation{kind: "I", id: 0, seq: 4}); err != nil { + t.Fatalf("intentional reconnect gap was treated as survivor loss: %v", err) + } + if err := r.observe(tracker, epochs, observation{kind: "I", id: 0, seq: 6}); err == nil { + t.Fatal("unplanned gap within a generation went undetected") + } +} +func TestHelperReadinessChallengeCannotBeKernelEcho(t *testing.T) { + request := frame("R", 2, 0, time.Unix(1, 0), 128) + var out bytes.Buffer + if err := Helper(context.Background(), bytes.NewReader(request), &out, 0, false); err != nil { + t.Fatal(err) + } + if !bytes.Contains(out.Bytes(), []byte("CAP A 00000002")) { + t.Fatalf("helper did not return distinct readiness proof: %q", out.String()) + } +} From 4814f3f5e36192150af92ae1e3a45e1b3d916faf Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 16:19:19 -0700 Subject: [PATCH 02/16] Fix capacity delivery, pacing and churn measurement --- internal/capacity/README.md | 71 ++- internal/capacity/canary.go | 4 +- internal/capacity/churn_test.go | 13 + internal/capacity/client.go | 235 ++++++-- internal/capacity/completion.go | 166 ++++++ internal/capacity/config.go | 6 +- internal/capacity/guard.go | 192 +++++-- internal/capacity/report.go | 80 +-- internal/capacity/review_regression_test.go | 588 ++++++++++++++++++++ internal/capacity/runner.go | 187 ++++--- internal/capacity/transfer.go | 207 +++++-- internal/capacity/workload.go | 68 ++- 12 files changed, 1569 insertions(+), 248 deletions(-) create mode 100644 internal/capacity/completion.go create mode 100644 internal/capacity/review_regression_test.go diff --git a/internal/capacity/README.md b/internal/capacity/README.md index 16701cc01..8f442b546 100644 --- a/internal/capacity/README.md +++ b/internal/capacity/README.md @@ -108,7 +108,13 @@ slot, 20 new external connections/second, 24 hours of authorization and 50 GiB of application payload. Up to twenty admission workers overlap readiness; every external connection and retry passes the same global limiter, including `Host`'s own reconnect -attempts. Admission workers are joined before steady observation or cleanup. +attempts. Initial hosts reserve a dial slot before `Host.Run`; guests reserve +before their ten-second service attempt, including WS CONNECT/upgrade. Queue +waiting is generator pacing, excluded from `ready_ms`, and reported separately +as `queue_ms` on admission/readiness events. A slow chosen ramp cannot consume +a client's service deadline. A permit remains exclusive until the actual dial +starts, so delayed reservations and cancelled waiters cannot bunch dials. +Admission workers are joined before steady observation or cleanup. The process rejects nonempty uppercase or lowercase `HTTP_PROXY`, `HTTPS_PROXY` and `ALL_PROXY`. Its private loopback CONNECT doors enforce destination/IP allowlists and own the host sockets used for cuts. `local_only: true` is a @@ -120,7 +126,13 @@ ready. Setup has a separate `(H + G) / rate + 30 seconds` deadline; manifest expiry and parent cancellation bound both. Summary `setup_seconds`, `steady_seconds` and total `duration_seconds` preserve the distinction. A five-minute hold therefore remains five minutes after admission. Host keepalive -uses the released client's 50-second interval and normal reconnect timings. +uses the released client's 50-second interval and normal reconnect backoff. +Released host redial budgets include CONNECT waiting: their normal 30-second +service budget gets an explicit allowance of at most five seconds for generator +pacing. The CONNECT door enforces the original 30-second service bound after +dispatch. The released WS dial/upgrade path remains in use. A retry queue that +cannot dispatch within five seconds stops as `generator_admission_backlog`; +this is invalid generator-limited evidence, never a relay capacity ceiling. ## Offered data and measurement @@ -133,18 +145,52 @@ ten seconds from helper startup. Idle helpers emit only their startup marker and readiness challenge/probe responses. A request tagged `R` receives a helper-generated `A` response, so a PTY kernel echo cannot prove readiness. This also works after the original startup marker leaves terminal history. +The authenticated `R` challenge may repeat every 100 ms within the same +original ten-second service attempt until a matching live `A` is seen. +`readiness_challenge_retry` events and summary `readiness_challenge_retries` +expose every additional attempt; all requests reserve payload before writing. +Ready/recovery latency includes this work. These retries establish the live +boundary and never retry or waive a steady `I`/`H` delivery failure. Challenges, +barriers and own input replies match their original timestamps, so replayed +responses cannot acknowledge or remeasure a request. The helper resynchronizes +its bounded line frames if a cut truncates a preceding input. All receivers drain output, track sequence gaps and reordering, and fail if a requested response remains outstanding for more -than ten seconds. Initial terminal replay is excluded until a channel's fresh -probe response establishes the live measurement boundary. Planned reconnects +than ten seconds. Initial terminal replay is excluded by a fresh helper `Q`/`B` barrier before +its live probe. The barrier reports the helper's emitted host-output sequence +count and establishes each channel's live measurement boundary. Every new +transport clears old barriers so replay cannot activate its live reader. Planned reconnects reset tracking only for the affected source generation; survivor generations retain gap/reordering checks. +Steady input has an independent schedule per guest. Expected restoration +pauses only that guest; it does not hold up healthy survivors. A missed +one-second interval fails as `generator_cadence_late`; it is never caught up by +sending active input after the offered interval. At steady end, +new input stops and a separate drain of at most ten seconds retains transports, +telemetry, watchdog and manifest bounds. It first requires every outstanding +own reply, then sends `F`/`B` barriers on every channel. `F` quiesces helper +output; `B` reports its final emitted output count. Every receiver must have +the final active input sequence from every live source and the final host +output sequence. Missing suffixes fail even when no later frame exposes a gap. +Late completed replies retain their original RTT timestamps and enter the RTT +reservoir. After successful completion and cleanup, `input_frames` equals +`rtt_count + canary_rtt_count + interrupted_replies`: each complete input write +is either acknowledged once or explicitly interrupted by a planned cut. +Readiness `I` probes and accepted inputs during churn are included. `R`, `Q`, +and `F` requests count in `handshake_frames`; their `A`/`B` replies do not create +input RTT samples. Idle host-only canaries contribute one `input_frames` and +one `canary_rtt_count` per successful canary, with zero normal `rtt_count`. A successful summary requires `delivery_complete: true`, keeps +`stop_reason: "duration"`, and adds `drain_seconds`; `steady_seconds` continues +to report only the configured offered interval. Inputs interrupted by an +intentional cut are counted explicitly in `interrupted_replies`; the restored +generation must prove delivery again. + Each 128-byte input/echo and 1,024-byte host frame contains a 44-byte ASCII header (`CAP`, role, eight-digit source and sequence, nineteen-digit Unix nanoseconds), padding, and one final LF. The 45 framing bytes are included in those fixed sizes. Readiness challenge and response frames are also 128 bytes; `handshake_frames` -records their requests separately from `input_frames`. The helper startup +records readiness and barrier requests separately from `input_frames`. The helper startup marker is ten bytes. SSH, TLS, WebSocket, TCP, terminal reconstruction and relay forwarding overhead are not measured as application framing. RTT uses the originating guest's echoed tag; @@ -177,9 +223,22 @@ at phase expiry are a failure. guests per host and 25 steady seconds. Counts are floored; use enough hosts for the requested percentage to select a host. Selected host and guest transports are cut over a ten-second schedule starting one second into steady load. +Selected guests are cut before selected hosts so an overlapping selection +always severs a real guest transport. Actual cutting takes no repair lock; +repairs may serialize independently. `transport_cut` records `planned_at`, +`actual_at`, `schedule_delay_ms`, `actual_transports`, and for host cuts +`already_disrupted_guests`. Already-disrupted guest transports cause failure, +not an extra claimed cut. Summary `actual_host_cuts` and `actual_guest_cuts` +must equal the requested floored counts. Host restoration uses the original `Host.Run`, host key and session ID; guests reuse their identities and transport kinds. Reconnect events and -latency are separate from first admissions. A survivor on an unaffected host +latency are separate from first admissions. Restoration `restoration_ms` is +actual cut-to-usable wall time, including generator pacing and repair waits; +the approved recovery objective continues to use this total. Additive +`queue_ms` separates admission pacing and `service_ms` reports the remaining +elapsed recovery path, including normal retry/backoff and repair waiting. +An overlapping later host cut supersedes guest repair without resetting that +guest's original outage start. A survivor on an unaffected host keeps its transport and must continue responding. Intentional input pauses are reported as `paused_active_frames`. Idle churn explicitly probes its survivor. Churn incomplete at phase expiry is a failure. Combining transfers and churn diff --git a/internal/capacity/canary.go b/internal/capacity/canary.go index 9f4591de6..3ede6391b 100644 --- a/internal/capacity/canary.go +++ b/internal/capacity/canary.go @@ -41,7 +41,7 @@ func (r *runState) hostCanary(sample int) error { r.rec.mu.Unlock() began := time.Now() err = q.connect(r.ctx, r) - readyElapsed := time.Since(began) + readyElapsed := time.Since(began) - q.queueElapsed q.setExpected(true) q.close() if err != nil { @@ -54,7 +54,7 @@ func (r *runState) hostCanary(sample int) error { r.rec.summary.CanaryReadyMillis = ms } r.rec.mu.Unlock() - r.rec.event("host_only_canary", map[string]any{"sample": sample, "host": q.host.id, "transport": u.Scheme, "ready_ms": ms, "transient_guest": true, "channels": 1}) + r.rec.event("host_only_canary", map[string]any{"sample": sample, "host": q.host.id, "transport": u.Scheme, "ready_ms": ms, "queue_ms": float64(q.queueElapsed) / float64(time.Millisecond), "transient_guest": true, "channels": 1}) return nil } func (r *runState) canaries() { diff --git a/internal/capacity/churn_test.go b/internal/capacity/churn_test.go index 8d58b9377..73087b23f 100644 --- a/internal/capacity/churn_test.go +++ b/internal/capacity/churn_test.go @@ -2,6 +2,8 @@ package capacity import ( "bytes" + "os" + "path/filepath" "testing" "time" ) @@ -45,4 +47,15 @@ func TestRealTransportChurnRestoresOriginalSessions(t *testing.T) { if s.EstablishedHosts != 2 || s.EstablishedGuests != 4 { t.Fatalf("reconnect was counted as new admission: %+v", s) } + if !s.DeliveryComplete || s.ActualHostCuts != 1 || s.ActualGuestCuts != 2 { + t.Fatalf("actual churn delivery/cuts incomplete: %+v", s) + } + if s.InputFrames != int64(s.RTTCount+s.CanaryRTTCount+s.InterruptedReplies) { + t.Fatalf("churn completion accounting lost or duplicated replies: %+v", s) + } + if dir := os.Getenv("CAPACITY_EVIDENCE_DIR"); dir != "" { + if err := os.WriteFile(filepath.Join(dir, "churn-events.jsonl"), out.Bytes(), 0600); err != nil { + t.Fatal(err) + } + } } diff --git a/internal/capacity/client.go b/internal/capacity/client.go index e09d6032a..292eeb468 100644 --- a/internal/capacity/client.go +++ b/internal/capacity/client.go @@ -55,7 +55,7 @@ func openGuest(ctx context.Context, g *guard, t Target, u *url.URL, s *api.GetSe var raw net.Conn var err error if u.Scheme == "ssh" { - raw, err = g.connect(ctx, t, u.Host) + raw, err = g.connectPermit(ctx, t, u.Host, proxy.takePermit()) } else { copy := *u copy.User = url.UserPassword(s.SessionId, base64.URLEncoding.EncodeToString([]byte(s.NodeAddr))) @@ -70,12 +70,24 @@ func openGuest(ctx context.Context, g *guard, t Target, u *url.URL, s *api.GetSe } type guest struct { + queueElapsed time.Duration + barriers map[int]*frameBarrier + transportLive bool + disruption int + disruptedAt time.Time + lastActiveSeq int + lastActiveAt time.Time + activeGeneration int + activeOffered bool + canary bool probeSeq [4]int pending map[string]time.Time + challenges map[string]time.Time lifetimeStop func() bool acks int64 ops sync.Mutex + inputOps sync.Mutex id int host *hostRun target Target @@ -106,16 +118,27 @@ func (q *guest) close() { if q.raw != nil { _ = q.raw.Close() } + q.inputOps.Lock() + defer q.inputOps.Unlock() for _, d := range q.done { <-d } q.done = nil + q.mu.Lock() q.channels = nil q.writers = nil + q.transportLive = false + q.mu.Unlock() } func (q *guest) connect(ctx context.Context, r *runState) error { q.ops.Lock() defer q.ops.Unlock() + q.mu.Lock() + q.challenges = map[string]time.Time{} + // A baseline belongs to one transport. Reconstructed old B frames must + // not activate the new reader before its own live Q/B exchange. + q.barriers = nil + q.mu.Unlock() s, err := q.host.session(ctx) if err != nil { return err @@ -123,13 +146,48 @@ func (q *guest) connect(ctx context.Context, r *runState) error { if err = r.guard.session(ctx, s); err != nil { return err } - attempt, cancel := context.WithTimeout(ctx, 10*time.Second) + timeout := r.guestReadyTimeout + if timeout == 0 { + timeout = 10 * time.Second + } + q.queueElapsed = 0 + queuedAt := time.Now() + queuectx := ctx + cancelQueue := func() {} + q.mu.Lock() + restoring := q.expected && q.generation > 0 + q.mu.Unlock() + if restoring && !q.canary { + queuectx, cancelQueue = context.WithTimeout(ctx, retryQueueLimit) + } + permit, err := r.guard.permit(queuectx) + q.queueElapsed = time.Since(queuedAt) + cancelQueue() + if err != nil { + if ctx.Err() == nil { + r.rec.event("admission_queue", map[string]any{"role": "guest", "guest": q.id, "retry": true, "queue_ms": float64(q.queueElapsed) / float64(time.Millisecond), "generator_limited": true}) + err = errors.New("generator_admission_backlog") + r.fail(err) + } + return err + } + defer permit.release() + q.proxy.mu.Lock() + q.proxy.permit = permit + q.proxy.mu.Unlock() + r.rec.event("admission_queue", map[string]any{"role": "guest", "guest": q.id, "queue_ms": float64(q.queueElapsed) / float64(time.Millisecond), "generation": q.generation}) + serviceBegan := time.Now() + attempt, cancel := context.WithTimeout(ctx, timeout) defer cancel() - q.client, q.raw, err = openGuest(attempt, r.guard, q.target, q.endpoint, s, q.identity, q.proxy) + client, raw, err := openGuest(attempt, r.guard, q.target, q.endpoint, s, q.identity, q.proxy) + q.mu.Lock() + q.client = client + q.raw = raw + q.transportLive = err == nil + q.mu.Unlock() if err != nil { return err } - client := q.client q.lifetimeStop = context.AfterFunc(ctx, func() { _ = client.Close() }) stop := context.AfterFunc(attempt, func() { _ = client.Close() }) defer stop() @@ -175,12 +233,35 @@ func (q *guest) connect(ctx context.Context, r *runState) error { tracker := sequenceTracker{} epochs := map[string]int{} live := false + baseline := time.Time{} + hostKey := fmt.Sprintf("H:%d", q.host.id) err := scanFrames(out, func(o observation) error { - if !live { - if o.kind != "I" || o.id != stream || o.seq != probeSeq { + if o.kind == "B" { + r.actual(128) + q.mu.Lock() + barrier := q.barriers[stream] + matches := o.id == stream && barrier != nil && o.sent.Equal(barrier.sent) + q.mu.Unlock() + if !matches { return nil } - live = true + if !live { + tracker[hostKey] = o.seq - 1 + baseline = barrier.sent + live = true + } else if last := tracker[hostKey]; last != o.seq-1 { + return fmt.Errorf("host output suffix missing: host=%d received_through=%d emitted=%d", q.host.id, last, o.seq) + } + if barrier.final { + if err := r.verifyFinalInputs(tracker, q.host.id, baseline); err != nil { + return err + } + } + barrier.once.Do(func() { close(barrier.done) }) + return nil + } + if !live { + return nil } if err := r.observe(tracker, epochs, o); err != nil { return err @@ -188,26 +269,27 @@ func (q *guest) connect(ctx context.Context, r *runState) error { r.actual(int64(frameLength(o.kind))) r.rec.mu.Lock() r.rec.summary.ReceivedFrames++ + acknowledged := false if o.kind == "I" && o.id == stream { - q.mu.Lock() - q.acks++ - delete(q.pending, fmt.Sprintf("%d:%d", o.id, o.seq)) - q.mu.Unlock() - ms := float64(time.Since(o.sent)) / float64(time.Millisecond) - if q.canary { - r.rec.summary.CanaryRTTCount++ - r.rec.summary.CanaryRTTMillis = append(r.rec.summary.CanaryRTTMillis, ms) - } else { - r.rec.summary.RTTCount++ - if len(r.rec.summary.RTTMillis) < 100000 { - r.rec.summary.RTTMillis = append(r.rec.summary.RTTMillis, ms) - } else if i := mathrand.IntN(r.rec.summary.RTTCount); i < 100000 { - r.rec.summary.RTTMillis[i] = ms + sent, ok := q.acknowledge(o) + acknowledged = ok + if ok { + ms := float64(time.Since(sent)) / float64(time.Millisecond) + if q.canary { + r.rec.summary.CanaryRTTCount++ + r.rec.summary.CanaryRTTMillis = append(r.rec.summary.CanaryRTTMillis, ms) + } else { + r.rec.summary.RTTCount++ + if len(r.rec.summary.RTTMillis) < 100000 { + r.rec.summary.RTTMillis = append(r.rec.summary.RTTMillis, ms) + } else if i := mathrand.IntN(r.rec.summary.RTTCount); i < 100000 { + r.rec.summary.RTTMillis[i] = ms + } } } } r.rec.mu.Unlock() - if o.kind == "I" && o.id == stream { + if acknowledged && o.seq == probeSeq { select { case ready <- struct{}{}: default: @@ -219,6 +301,17 @@ func (q *guest) connect(ctx context.Context, r *runState) error { if ack.id != stream || ack.seq != probeSeq { return } + q.mu.Lock() + key := fmt.Sprintf("%d:%d:%d", ack.id, ack.seq, ack.sent.UnixNano()) + sent, exists := q.challenges[key] + valid := exists && sent.Equal(ack.sent) + if valid { + delete(q.challenges, key) + } + q.mu.Unlock() + if !valid { + return + } select { case helperReady <- struct{}{}: default: @@ -231,15 +324,11 @@ func (q *guest) connect(ctx context.Context, r *runState) error { r.fail(fmt.Errorf("guest %d channel %d: %w", q.id, ch, err)) } }() - if err = r.sendKind(q, in, stream, probeSeq, "R"); err != nil { + if err = q.helperChallenge(attempt, r, in, stream, probeSeq, ch, helperReady); err != nil { return err } - select { - case <-helperReady: - case <-attempt.Done(): - return fmt.Errorf("helper readiness: %w", attempt.Err()) - case <-ctx.Done(): - return ctx.Err() + if err = q.barrier(attempt, r, in, stream, "Q"); err != nil { + return err } if err = r.send(q, in, stream, probeSeq); err != nil { return err @@ -255,7 +344,7 @@ func (q *guest) connect(ctx context.Context, r *runState) error { if !stop() { return errors.New("guest readiness cancelled") } - r.rec.event("guest_ready", map[string]any{"guest": q.id, "host": q.host.id, "transport": q.endpoint.Scheme, "node": s.NodeAddr, "session_id": s.SessionId, "channels": len(q.channels), "generation": q.generation, "canary": q.canary}) + r.rec.event("guest_ready", map[string]any{"guest": q.id, "host": q.host.id, "transport": q.endpoint.Scheme, "node": s.NodeAddr, "session_id": s.SessionId, "channels": len(q.channels), "generation": q.generation, "canary": q.canary, "ready_ms": float64(time.Since(serviceBegan)) / float64(time.Millisecond), "queue_ms": float64(q.queueElapsed) / float64(time.Millisecond)}) return nil } func frameLength(kind string) int { @@ -265,29 +354,52 @@ func frameLength(kind string) int { return 128 } func (q *guest) active(r *runState) error { - q.ops.Lock() - defer q.ops.Unlock() q.mu.Lock() expected := q.expected q.mu.Unlock() if expected { - r.rec.mu.Lock() - r.rec.summary.PausedActiveFrames++ - r.rec.mu.Unlock() + r.paused() + return nil + } + q.inputOps.Lock() + defer q.inputOps.Unlock() + q.mu.Lock() + if q.expected { + q.mu.Unlock() + r.paused() return nil } if len(q.writers) == 0 { + q.mu.Unlock() return errors.New("guest has no terminal") } - err := r.send(q, q.writers[0], q.id*4, q.seq) + w, seq := q.writers[0], q.seq q.seq++ - if err == nil { - r.rec.mu.Lock() - r.rec.summary.ActiveFrames++ - r.rec.mu.Unlock() + q.mu.Unlock() + sent := time.Now() + err := r.send(q, w, q.id*4, seq) + if err != nil { + q.mu.Lock() + expected := q.expected + q.mu.Unlock() + if expected { + r.paused() + return nil + } + return err } - return err + q.mu.Lock() + q.lastActiveSeq = seq + q.lastActiveAt = sent + q.activeGeneration = q.generation + q.activeOffered = true + q.mu.Unlock() + r.rec.mu.Lock() + r.rec.summary.ActiveFrames++ + r.rec.mu.Unlock() + return nil } +func (r *runState) paused() { r.rec.mu.Lock(); r.rec.summary.PausedActiveFrames++; r.rec.mu.Unlock() } func (q *guest) setExpected(b bool) { q.mu.Lock(); q.expected = b; q.mu.Unlock() } func transport(t Target, percent, index, total int, public bool) (*url.URL, error) { raw := t.SSHURL @@ -302,3 +414,44 @@ func transport(t Target, percent, index, total int, public bool) (*url.URL, erro } return endpoint(raw, "ssh", "ws", "wss") } + +// Replay and duplicates are observations, never a second completed request. +func (q *guest) acknowledge(o observation) (time.Time, bool) { + q.mu.Lock() + defer q.mu.Unlock() + key := fmt.Sprintf("%d:%d", o.id, o.seq) + sent, ok := q.pending[key] + if !ok || !sent.Equal(o.sent) { + return time.Time{}, false + } + delete(q.pending, key) + q.acks++ + return sent, true +} + +func (q *guest) helperChallenge(ctx context.Context, r *runState, w io.Writer, stream, seq, ch int, ready <-chan struct{}) error { + for attempt := 1; ; attempt++ { + if err := ctx.Err(); err != nil { + return fmt.Errorf("helper readiness: %w", err) + } + if attempt > 1 { + r.rec.mu.Lock() + r.rec.summary.ReadinessChallengeRetries++ + r.rec.mu.Unlock() + r.rec.event("readiness_challenge_retry", map[string]any{"guest": q.id, "host": q.host.id, "channel": ch, "attempt": attempt, "canary": q.canary}) + } + if err := r.sendKind(q, w, stream, seq, "R"); err != nil { + return err + } + timer := time.NewTimer(100 * time.Millisecond) + select { + case <-ready: + timer.Stop() + return nil + case <-ctx.Done(): + timer.Stop() + return fmt.Errorf("helper readiness: %w", ctx.Err()) + case <-timer.C: + } + } +} diff --git a/internal/capacity/completion.go b/internal/capacity/completion.go new file mode 100644 index 000000000..011d060b5 --- /dev/null +++ b/internal/capacity/completion.go @@ -0,0 +1,166 @@ +package capacity + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + "time" +) + +type frameBarrier struct { + sent time.Time + done chan struct{} + once sync.Once + final bool +} + +// B reports the helper's emitted H count after processing all earlier input. +// A baseline Q excludes replay; a final F quiesces periodic helper output. +func (q *guest) barrier(ctx context.Context, r *runState, w io.Writer, stream int, kind string) error { + b := &frameBarrier{sent: time.Now(), done: make(chan struct{}), final: kind == "F"} + q.mu.Lock() + if q.barriers == nil { + q.barriers = map[int]*frameBarrier{} + } + q.barriers[stream] = b + q.mu.Unlock() + if err := r.sendBarrier(q, w, stream, kind, b.sent); err != nil { + return err + } + select { + case <-b.done: + return nil + case <-ctx.Done(): + return fmt.Errorf("terminal delivery barrier: %w", ctx.Err()) + } +} +func (r *runState) sendBarrier(q *guest, w io.Writer, stream int, kind string, at time.Time) error { + fanout := r.cfg.GuestsPerHost * r.cfg.ChannelsPerGuest + if q.canary { + fanout = 1 + } + if err := r.reserve(128 * int64(1+fanout)); err != nil { + return err + } + n, err := w.Write(frame(kind, stream, 0, at, 128)) + r.actual(int64(n)) + if n != 128 && err == nil { + err = io.ErrShortWrite + } + if n == 128 { + r.rec.mu.Lock() + r.rec.summary.HandshakeFrames++ + r.rec.mu.Unlock() + } + return err +} +func (r *runState) completeDelivery() error { + // Own acknowledgments prove every offered input reached the helper before + // final barriers are queued. Each receiver then verifies every active + // source's final sequence, covering a lost fan-out suffix as well. + for { + pending := 0 + for _, q := range r.guests { + q.mu.Lock() + pending += len(q.pending) + q.mu.Unlock() + } + if pending == 0 { + break + } + if err := pause(r.ctx, 10*time.Millisecond); err != nil { + return err + } + } + var wg sync.WaitGroup + failures := make(chan error, len(r.guests)) + for _, q := range r.guests { + if q.canary { + continue + } + wg.Add(1) + go func() { + defer wg.Done() + q.mu.Lock() + expected := q.expected + q.mu.Unlock() + if expected { + failures <- errors.New("terminal_restoration_incomplete") + return + } + q.ops.Lock() + defer q.ops.Unlock() + for ch, w := range q.writers { + if err := q.barrier(r.ctx, r, w, q.id*4+ch, "F"); err != nil { + failures <- err + return + } + } + q.mu.Lock() + pending := len(q.pending) + q.mu.Unlock() + if pending != 0 { + failures <- fmt.Errorf("terminal delivery missing %d replies", pending) + } + }() + } + wg.Wait() + close(failures) + for err := range failures { + if err != nil { + return err + } + } + return nil +} + +// Each guest owns its own absolute one-second schedule. A slow writer or an +// expected restoration cannot hold up any other guest's offered cadence. +func (r *runState) activeWorkload(start time.Time, wg *sync.WaitGroup) { + if !r.cfg.Active { + return + } + for _, q := range r.guests { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < r.cfg.DurationSeconds; i++ { + if err := pause(r.ctx, time.Until(start.Add(time.Duration(i)*time.Second))); err != nil { + return + } + if now := time.Now(); !now.Before(start.Add(time.Duration(r.cfg.DurationSeconds)*time.Second)) || now.Sub(start.Add(time.Duration(i)*time.Second)) >= time.Second { + r.fail(errors.New("generator_cadence_late")) + return + } + if err := q.active(r); err != nil { + r.fail(err) + return + } + } + }() + } +} + +func (r *runState) verifyFinalInputs(tracker sequenceTracker, hostID int, baseline time.Time) error { + r.clientsMu.RLock() + defer r.clientsMu.RUnlock() + for _, source := range r.guests { + if source.host.id != hostID { + continue + } + source.mu.Lock() + expected := source.lastActiveSeq + required := source.activeOffered && source.activeGeneration == source.generation && !source.lastActiveAt.Before(baseline) + id := source.id + source.mu.Unlock() + if required { + got, ok := tracker[fmt.Sprintf("I:%d", id*4)] + if !ok || got != expected { + return fmt.Errorf("input suffix missing: guest=%d received_through=%d expected=%d", id, got, expected) + } + } + } + return nil +} diff --git a/internal/capacity/config.go b/internal/capacity/config.go index 2ed489d47..5ff0de664 100644 --- a/internal/capacity/config.go +++ b/internal/capacity/config.go @@ -164,9 +164,9 @@ func (c Config) Validate(now time.Time) error { return errors.New("churn needs survivor hosts/guests and at least 25 seconds") } fanout := int64(1 + c.GuestsPerHost*c.ChannelsPerGuest) - probes := int64(c.Hosts*c.GuestsPerHost*c.ChannelsPerGuest) * 256 * fanout + probes := int64(c.Hosts*c.GuestsPerHost*c.ChannelsPerGuest) * 512 * fanout if c.idleHostOnly() { - probes += int64(1+(c.DurationSeconds-1)/30) * 512 + probes += int64(1+(c.DurationSeconds-1)/30) * 768 } active := int64(0) if c.Active { @@ -246,7 +246,7 @@ func (c Config) hostOutputReservation() int64 { fanout := int64(1 + c.GuestsPerHost*c.ChannelsPerGuest) // Startup has an independent, enforced ramp+30s deadline. Teardown can // produce one final tick, reserved here before any helper is started. - seconds := int64(c.DurationSeconds + 42 + c.Hosts*(1+c.GuestsPerHost)/c.AdmissionsPerSecond) + seconds := int64(c.DurationSeconds + 52 + c.Hosts*(1+c.GuestsPerHost)/c.AdmissionsPerSecond) output := int64(10) // helper-ready marker if c.Active { output += (seconds/10 + 1) * 1024 diff --git a/internal/capacity/guard.go b/internal/capacity/guard.go index ac0cfb923..d3edff2e8 100644 --- a/internal/capacity/guard.go +++ b/internal/capacity/guard.go @@ -23,12 +23,13 @@ type guard struct { dial func(context.Context, string, string) (net.Conn, error) mu sync.Mutex next time.Time + gate chan struct{} active int peak int } func newGuard(c Config) *guard { - return &guard{cfg: c, lookup: net.DefaultResolver.LookupIPAddr, dial: (&net.Dialer{Timeout: 10 * time.Second}).DialContext} + return &guard{gate: make(chan struct{}, 1), cfg: c, lookup: net.DefaultResolver.LookupIPAddr, dial: (&net.Dialer{Timeout: 10 * time.Second}).DialContext} } func (g *guard) resolved(ctx context.Context, t Target, addr string) (string, error) { h, p, err := net.SplitHostPort(addr) @@ -54,24 +55,43 @@ func (g *guard) resolved(ctx context.Context, t Target, addr string) (string, er } return net.JoinHostPort(ips[0].IP.String(), p), nil } -func (g *guard) admit(ctx context.Context) error { - g.mu.Lock() - when := g.next - if when.Before(time.Now()) { - when = time.Now() + +// A permit holds the one dial gate until an actual external dial starts. +// Waiting for it is generator pacing, outside initial service deadlines. +type dialPermit struct { + g *guard + once sync.Once + queued time.Duration +} + +func (p *dialPermit) release() { + if p != nil { + p.once.Do(func() { <-p.g.gate }) } - g.next = when.Add(time.Second / time.Duration(g.cfg.AdmissionsPerSecond)) - g.mu.Unlock() - timer := time.NewTimer(time.Until(when)) - defer timer.Stop() +} +func (g *guard) permit(ctx context.Context) (*dialPermit, error) { + began := time.Now() select { + case g.gate <- struct{}{}: case <-ctx.Done(): - return ctx.Err() - case <-timer.C: - return nil + return nil, ctx.Err() + } + p := &dialPermit{g: g} + g.mu.Lock() + when := g.next + g.mu.Unlock() + if err := pause(ctx, time.Until(when)); err != nil { + p.release() + return nil, err } + p.queued = time.Since(began) + return p, nil } func (g *guard) connect(ctx context.Context, t Target, addr string) (net.Conn, error) { + return g.connectPermit(ctx, t, addr, nil) +} +func (g *guard) connectPermit(ctx context.Context, t Target, addr string, permit *dialPermit) (net.Conn, error) { + defer func() { permit.release() }() if err := ctx.Err(); err != nil { return nil, err } @@ -89,8 +109,11 @@ func (g *guard) connect(ctx context.Context, t Target, addr string) (net.Conn, e if err != nil { return nil, err } - if err = g.admit(ctx); err != nil { - return nil, err + if permit == nil { + permit, err = g.permit(ctx) + if err != nil { + return nil, err + } } if err := ctx.Err(); err != nil { return nil, err @@ -108,7 +131,15 @@ func (g *guard) connect(ctx context.Context, t Target, addr string) (net.Conn, e g.peak = g.active } g.mu.Unlock() - conn, err := g.dial(ctx, "tcp", pinned) + // Update pacing at the actual dial, so delayed/cancelled reservations + // cannot bunch later external handshakes together. + g.mu.Lock() + g.next = time.Now().Add(time.Second / time.Duration(g.cfg.AdmissionsPerSecond)) + g.mu.Unlock() + permit.release() + dialctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + conn, err := g.dial(dialctx, "tcp", pinned) if err != nil { g.mu.Lock() g.active-- @@ -183,17 +214,25 @@ func pinnedKey(t Target) ssh.HostKeyCallback { // DNS pinning and admission limiter as guests. It also owns the actual sockets // severed during host churn; no host or session is replaced. type tunnelProxy struct { - g *guard - t Target - ctx context.Context - server *http.Server - listener net.Listener - url *url.URL - mu sync.Mutex - conns map[net.Conn]bool - closing bool - workers sync.WaitGroup - done chan struct{} + g *guard + t Target + ctx context.Context + server *http.Server + listener net.Listener + url *url.URL + mu sync.Mutex + conns map[net.Conn]bool + closing bool + workers sync.WaitGroup + done chan struct{} + permit *dialPermit + dialStarted chan time.Time + readyTimer *time.Timer + serviceTimeout time.Duration + queueLimit time.Duration + queueTime time.Duration + failure func(error) + queueEvent func(time.Duration, bool) } func newTunnelProxy(ctx context.Context, g *guard, t Target) (*tunnelProxy, error) { @@ -201,20 +240,20 @@ func newTunnelProxy(ctx context.Context, g *guard, t Target) (*tunnelProxy, erro if err != nil { return nil, err } - p := &tunnelProxy{g: g, t: t, ctx: ctx, listener: ln, conns: map[net.Conn]bool{}, done: make(chan struct{})} + p := &tunnelProxy{g: g, t: t, ctx: ctx, listener: ln, conns: map[net.Conn]bool{}, done: make(chan struct{}), dialStarted: make(chan time.Time, 1), queueLimit: 5 * time.Second} p.url = &url.URL{Scheme: "http", Host: ln.Addr().String()} p.server = &http.Server{Handler: p, ReadHeaderTimeout: 5 * time.Second} go func() { defer close(p.done); _ = p.server.Serve(ln) }() return p, nil } -func (p *tunnelProxy) track(c net.Conn) bool { +func (p *tunnelProxy) track(c net.Conn, external bool) bool { p.mu.Lock() defer p.mu.Unlock() if p.closing { _ = c.Close() return false } - p.conns[c] = true + p.conns[c] = external return true } func (p *tunnelProxy) untrack(c net.Conn) { @@ -223,16 +262,49 @@ func (p *tunnelProxy) untrack(c net.Conn) { delete(p.conns, c) p.mu.Unlock() } -func (p *tunnelProxy) cut() { +func (p *tunnelProxy) cut() (int, time.Time) { p.mu.Lock() defer p.mu.Unlock() + n := 0 + var actual time.Time + // Sever the external socket first. Closing the loopback side first could + // race its copier and turn the external close into an uncounted no-op. + for c, external := range p.conns { + if !external { + continue + } + if err := c.Close(); err == nil { + n++ + actual = time.Now() + } + delete(p.conns, c) + } for c := range p.conns { _ = c.Close() + delete(p.conns, c) + } + return n, actual +} +func (p *tunnelProxy) ready() { + p.mu.Lock() + defer p.mu.Unlock() + if p.readyTimer != nil { + p.readyTimer.Stop() + p.readyTimer = nil + } + if p.serviceTimeout > 0 { + p.serviceTimeout = 30 * time.Second } } +func (p *tunnelProxy) queued() time.Duration { p.mu.Lock(); defer p.mu.Unlock(); return p.queueTime } func (p *tunnelProxy) close() { p.mu.Lock() p.closing = true + if p.readyTimer != nil { + p.readyTimer.Stop() + } + p.permit.release() + p.permit = nil for c := range p.conns { _ = c.Close() } @@ -254,14 +326,54 @@ func (p *tunnelProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { http.Error(w, "CONNECT required", 405) return } - ctx, cancel := context.WithTimeout(p.ctx, 10*time.Second) - defer cancel() - upstream, err := p.g.connect(ctx, p.t, r.Host) + p.mu.Lock() + permit := p.permit + p.permit = nil + p.mu.Unlock() + if permit == nil { + // Released retry/WS dial clocks include CONNECT. Keep that queue bounded, + // classify saturation as generator-limited, and give retry service its + // normal 30 seconds after this at-most-five-second pacing allowance. + queuedAt := time.Now() + queuectx, cancel := context.WithTimeout(p.ctx, p.queueLimit) + var err error + permit, err = p.g.permit(queuectx) + cancel() + if p.queueEvent != nil { + p.queueEvent(time.Since(queuedAt), err != nil) + } + if err != nil { + if p.ctx.Err() == nil && p.failure != nil { + p.failure(errors.New("generator_admission_backlog")) + } + http.Error(w, "generator admission backlog", 503) + return + } + } + defer permit.release() + p.mu.Lock() + p.queueTime += permit.queued + p.mu.Unlock() + began := time.Now() + select { + case p.dialStarted <- began: + default: + } + upstream, err := p.g.connectPermit(p.ctx, p.t, r.Host, permit) if err != nil { http.Error(w, "destination refused", 403) return } - if !p.track(upstream) { + p.mu.Lock() + if p.serviceTimeout > 0 { + p.readyTimer = time.AfterFunc(p.serviceTimeout-time.Since(began), func() { _ = upstream.Close() }) + } + serviceTimer := p.readyTimer + p.mu.Unlock() + if serviceTimer != nil { + defer serviceTimer.Stop() + } + if !p.track(upstream, true) { return } defer p.untrack(upstream) @@ -269,7 +381,7 @@ func (p *tunnelProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { if err != nil { return } - if !p.track(downstream) { + if !p.track(downstream, false) { return } defer p.untrack(downstream) @@ -286,3 +398,11 @@ func (p *tunnelProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { _ = upstream.Close() <-copied } + +func (p *tunnelProxy) takePermit() *dialPermit { + p.mu.Lock() + defer p.mu.Unlock() + permit := p.permit + p.permit = nil + return permit +} diff --git a/internal/capacity/report.go b/internal/capacity/report.go index c13ae6531..ca9ca86fa 100644 --- a/internal/capacity/report.go +++ b/internal/capacity/report.go @@ -8,43 +8,49 @@ import ( ) type Summary struct { - PeakExternalConnections int `json:"peak_external_connections"` - HandshakeFrames int64 `json:"handshake_frames"` - CanaryRTTCount int `json:"canary_rtt_count"` - CanaryRTTMillis []float64 `json:"canary_rtt_millis"` - CanaryOffered int `json:"canary_offered"` - CanaryEstablished int `json:"canary_established"` - CanaryReadyMillis float64 `json:"canary_ready_millis"` - CleanupComplete bool `json:"cleanup_complete"` - PausedActiveFrames int64 `json:"paused_active_frames"` - InputFrames int64 `json:"input_frames"` - ActiveFrames int64 `json:"active_frames"` - ChurnCompleted bool `json:"churn_completed"` - RTTSampling string `json:"rtt_sampling"` - ReservedPayloadBytes int64 `json:"reserved_payload_bytes"` - SetupSeconds float64 `json:"setup_seconds"` - SteadySeconds float64 `json:"steady_seconds"` - RunID string `json:"run_id"` - OfferedHosts int `json:"offered_hosts"` - OfferedGuests int `json:"offered_guests"` - OfferedChannels int `json:"offered_channels"` - EstablishedHosts int `json:"established_hosts"` - EstablishedGuests int `json:"established_guests"` - EstablishedChannels int `json:"established_channels"` - AdmissionFailures int `json:"admission_failures"` - PayloadBytes int64 `json:"payload_bytes"` - PayloadBudgetBytes int64 `json:"payload_budget_bytes"` - ReceivedFrames int64 `json:"received_frames"` - RTTCount int `json:"rtt_count"` - RTTMillis []float64 `json:"rtt_millis"` - RTTP50Millis float64 `json:"rtt_p50_millis"` - RTTP95Millis float64 `json:"rtt_p95_millis"` - RTTP99Millis float64 `json:"rtt_p99_millis"` - Reconnects int `json:"reconnects"` - Transfers int `json:"transfers"` - TransferBytes int64 `json:"transfer_bytes"` - DurationSeconds float64 `json:"duration_seconds"` - StopReason string `json:"stop_reason"` + ReadinessChallengeRetries int `json:"readiness_challenge_retries"` + DrainSeconds float64 `json:"drain_seconds"` + DeliveryComplete bool `json:"delivery_complete"` + ActualHostCuts int `json:"actual_host_cuts"` + ActualGuestCuts int `json:"actual_guest_cuts"` + InterruptedReplies int `json:"interrupted_replies"` + PeakExternalConnections int `json:"peak_external_connections"` + HandshakeFrames int64 `json:"handshake_frames"` + CanaryRTTCount int `json:"canary_rtt_count"` + CanaryRTTMillis []float64 `json:"canary_rtt_millis"` + CanaryOffered int `json:"canary_offered"` + CanaryEstablished int `json:"canary_established"` + CanaryReadyMillis float64 `json:"canary_ready_millis"` + CleanupComplete bool `json:"cleanup_complete"` + PausedActiveFrames int64 `json:"paused_active_frames"` + InputFrames int64 `json:"input_frames"` + ActiveFrames int64 `json:"active_frames"` + ChurnCompleted bool `json:"churn_completed"` + RTTSampling string `json:"rtt_sampling"` + ReservedPayloadBytes int64 `json:"reserved_payload_bytes"` + SetupSeconds float64 `json:"setup_seconds"` + SteadySeconds float64 `json:"steady_seconds"` + RunID string `json:"run_id"` + OfferedHosts int `json:"offered_hosts"` + OfferedGuests int `json:"offered_guests"` + OfferedChannels int `json:"offered_channels"` + EstablishedHosts int `json:"established_hosts"` + EstablishedGuests int `json:"established_guests"` + EstablishedChannels int `json:"established_channels"` + AdmissionFailures int `json:"admission_failures"` + PayloadBytes int64 `json:"payload_bytes"` + PayloadBudgetBytes int64 `json:"payload_budget_bytes"` + ReceivedFrames int64 `json:"received_frames"` + RTTCount int `json:"rtt_count"` + RTTMillis []float64 `json:"rtt_millis"` + RTTP50Millis float64 `json:"rtt_p50_millis"` + RTTP95Millis float64 `json:"rtt_p95_millis"` + RTTP99Millis float64 `json:"rtt_p99_millis"` + Reconnects int `json:"reconnects"` + Transfers int `json:"transfers"` + TransferBytes int64 `json:"transfer_bytes"` + DurationSeconds float64 `json:"duration_seconds"` + StopReason string `json:"stop_reason"` } type recorder struct { diff --git a/internal/capacity/review_regression_test.go b/internal/capacity/review_regression_test.go new file mode 100644 index 000000000..545a2bf13 --- /dev/null +++ b/internal/capacity/review_regression_test.go @@ -0,0 +1,588 @@ +package capacity + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +func TestSilentAfterReadinessFailsDelivery(t *testing.T) { + c := localRelay(t) + helper := filepath.Join(t.TempDir(), "silent-helper") + script := "#!/bin/sh\nstty raw -echo\ndd bs=128 count=1 2>/dev/null | tr R A\ndd bs=128 count=1 2>/dev/null | tr Q B\ndd bs=128 count=1 2>/dev/null\ncat >/dev/null\n" + if err := os.WriteFile(helper, []byte(script), 0700); err != nil { + t.Fatal(err) + } + c.HelperExecutable = helper + c.DurationSeconds = 3 + s, err := Run(t.Context(), c, io.Discard) + t.Logf("delivery result: err=%v active=%d rtt=%d stop=%s", err, s.ActiveFrames, s.RTTCount, s.StopReason) + if s.EstablishedGuests != 1 || s.ActiveFrames != 3 { + t.Fatalf("fixture did not reach steady workload: %+v", s) + } + if err == nil { + t.Fatalf("silent post-readiness workload succeeded: active=%d rtt=%d stop=%s", s.ActiveFrames, s.RTTCount, s.StopReason) + } + if !s.CleanupComplete { + t.Fatal("failed delivery did not clean up") + } +} + +func TestExpectedRestoreDoesNotBlockHealthyCadence(t *testing.T) { + c := validConfig(t) + r := &runState{cfg: c, rec: newRecorder(io.Discard, c)} + disrupted := &guest{expected: true} + healthy := &guest{writers: []io.WriteCloser{discardWriteCloser{}}} + disrupted.ops.Lock() // connect holds this throughout a slow restoration. + defer disrupted.ops.Unlock() + finished := make(chan error, 1) + go func() { + for i := 0; i < 3; i++ { + if err := disrupted.active(r); err != nil { + finished <- err + return + } + if err := healthy.active(r); err != nil { + finished <- err + return + } + } + finished <- nil + }() + select { + case err := <-finished: + if err != nil { + t.Fatal(err) + } + case <-time.After(100 * time.Millisecond): + t.Fatal("restoring guest blocked healthy offered cadence") + } + if r.rec.summary.ActiveFrames != 3 || r.rec.summary.PausedActiveFrames != 3 { + t.Fatalf("wrong independent cadence accounting: %+v", r.rec.summary) + } +} + +type discardWriteCloser struct{} + +func (discardWriteCloser) Write(b []byte) (int, error) { return len(b), nil } +func (discardWriteCloser) Close() error { return nil } + +func TestGuestTransportCutDoesNotWaitForHostRepair(t *testing.T) { + c := validConfig(t) + ctx, cancel := context.WithCancelCause(t.Context()) + r := &runState{ctx: ctx, cancel: cancel, cfg: c, rec: newRecorder(io.Discard, c)} + raw, peer := net.Pipe() + defer peer.Close() + h := &hostRun{} + q := &guest{host: h, raw: raw, transportLive: true} + r.guests = []*guest{q} + h.repair.Lock() + finished := make(chan error, 1) + go func() { finished <- r.cutGuest(q) }() + defer func() { cancel(context.Canceled); h.repair.Unlock(); <-finished }() + _ = peer.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) + if _, err := peer.Read(make([]byte, 1)); err == nil { + t.Fatal("transport was not cut") + } else if e, ok := err.(net.Error); ok && e.Timeout() { + t.Fatal("actual transport cut waited behind slow host repair") + } +} + +func TestAdmissionQueueIsOutsideServiceReadinessDeadline(t *testing.T) { + for _, mix := range []struct { + name string + ws int + }{{"ssh", 0}, {"ws", 100}} { + t.Run(mix.name, func(t *testing.T) { pacedAdmissionFixture(t, mix.ws) }) + } +} +func pacedAdmissionFixture(t *testing.T, ws int) { + c := localRelay(t) + c.HostWSPercent = ws + ctx, cancel := context.WithCancelCause(t.Context()) + dir, err := os.MkdirTemp("", "cap-admit-") + if err != nil { + t.Fatal(err) + } + var events bytes.Buffer + r := &runState{ctx: ctx, cancel: cancel, cfg: c, rec: newRecorder(&events, c), guard: newGuard(c), dir: dir, helper: c.HelperExecutable, hostReadyTimeout: 200 * time.Millisecond, guestReadyTimeout: 200 * time.Millisecond} + defer func() { + if err := r.cleanup(); err != nil { + t.Error(err) + } + }() + r.hosts = make([]*hostRun, 1) + r.guard.next = time.Now().Add(450 * time.Millisecond) + target := c.Manifest.Targets[0] + if err := r.startHost(0, target, nil); err != nil { + t.Fatalf("generator queue consumed host service deadline: %v", err) + } + u, err := transport(target, ws, 0, 1, false) + if err != nil { + t.Fatal(err) + } + key, err := makeSigner() + if err != nil { + t.Fatal(err) + } + proxy, err := newTunnelProxy(ctx, r.guard, target) + if err != nil { + t.Fatal(err) + } + r.proxies = append(r.proxies, proxy) + q := &guest{host: r.hosts[0], target: target, endpoint: u, identity: key, proxy: proxy} + r.guests = append(r.guests, q) + r.guard.mu.Lock() + r.guard.next = time.Now().Add(450 * time.Millisecond) + r.guard.mu.Unlock() + if err = q.connect(ctx, r); err != nil { + t.Fatalf("generator queue consumed guest service deadline: %v", err) + } + decoder := json.NewDecoder(&events) + ready := 0 + for decoder.More() { + var event map[string]any + if err := decoder.Decode(&event); err != nil { + t.Fatal(err) + } + if event["event"] == "host_ready" || event["event"] == "guest_ready" { + ready++ + if event["queue_ms"].(float64) < 400 || event["ready_ms"].(float64) >= 200 { + t.Fatalf("pacing polluted service latency: %+v", event) + } + } + } + if ready != 2 { + t.Fatalf("missing service readiness timing events: %d", ready) + } + +} + +func scriptedHelper(t *testing.T, body string) Config { + t.Helper() + c := localRelay(t) + path := filepath.Join(t.TempDir(), "helper") + if err := os.WriteFile(path, []byte("#!/bin/sh\nstty raw -echo\n"+body), 0700); err != nil { + t.Fatal(err) + } + c.HelperExecutable = path + return c +} + +func TestLateFinalReplyIsDrainedAndMeasured(t *testing.T) { + c := scriptedHelper(t, "dd bs=128 count=1 2>/dev/null | tr R A\ndd bs=128 count=1 2>/dev/null | tr Q B\ndd bs=128 count=1 2>/dev/null\ndd bs=128 count=1 2>/dev/null | (sleep 1.2; cat)\ndd bs=128 count=1 2>/dev/null | tr F B\ncat >/dev/null\n") + c.DurationSeconds = 1 + s, err := Run(t.Context(), c, io.Discard) + if err != nil { + t.Fatalf("late final reply failed: %v %+v", err, s) + } + late := false + for _, ms := range s.RTTMillis { + late = late || ms >= 1000 + } + if !s.DeliveryComplete || s.RTTCount != 2 || !late || s.DrainSeconds < .1 || s.SteadySeconds < 1 || s.SteadySeconds > 1.1 { + t.Fatalf("late reply was not measured separately from offered interval: %+v", s) + } +} + +func TestMissingHostOutputSuffixFailsDelivery(t *testing.T) { + // The barrier advertises one emitted H frame, but that entire suffix is + // dropped. Every I frame is echoed, so own-response and gap-only checks pass. + c := scriptedHelper(t, "while IFS= read -r line; do\ncase \"$line\" in\n'CAP R '*) printf '%s\\n' \"$line\" | tr R A;;\n'CAP Q '*) printf '%s\\n' \"$line\" | tr Q B;;\n'CAP F '*) printf '%s\\n' \"$line\" | awk '{$2=\"B\";$4=\"00000001\";print}';;\n*) printf '%s\\n' \"$line\";;\nesac\ndone\n") + c.DurationSeconds = 1 + s, err := Run(t.Context(), c, io.Discard) + if err == nil { + t.Fatalf("missing H suffix succeeded: %+v", s) + } + if !strings.Contains(err.Error(), "host output suffix missing") { + t.Fatalf("fixture did not detect missing H suffix: %v", err) + } +} + +func TestFinalFanoutSuffixMustReachEveryReceiver(t *testing.T) { + h := &hostRun{id: 0} + r := &runState{guests: []*guest{{id: 0, host: h, activeOffered: true, lastActiveSeq: 4}, {id: 1, host: h, activeOffered: true, lastActiveSeq: 7}}} + if err := r.verifyFinalInputs(sequenceTracker{"I:0": 4, "I:4": 6}, 0, time.Time{}); err == nil { + t.Fatal("missing final peer input was accepted") + } + if err := r.verifyFinalInputs(sequenceTracker{"I:0": 4, "I:4": 7}, 0, time.Time{}); err != nil { + t.Fatal(err) + } +} + +func TestActualCutsKeepScheduleDuringSlowOverlappingRepairs(t *testing.T) { + c := validConfig(t) + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(context.Canceled) + r := &runState{cfg: c, ctx: ctx, cancel: cancel, rec: newRecorder(io.Discard, c)} + h := &hostRun{id: 0, proxy: &tunnelProxy{conns: map[net.Conn]bool{}}} + hostRaw, hostPeer := net.Pipe() + defer hostPeer.Close() + h.proxy.conns[hostRaw] = true + r.hosts = []*hostRun{h} + var peers []net.Conn + for i := 0; i < 2; i++ { + raw, peer := net.Pipe() + peers = append(peers, peer) + defer peer.Close() + r.guests = append(r.guests, &guest{id: i, host: h, raw: raw, transportLive: true}) + } + actions := []churnAction{{index: 0}, {index: 1, offset: 25 * time.Millisecond}, {host: true, index: 0, offset: 50 * time.Millisecond}} + release := make(chan struct{}) + actual := make(chan time.Time, len(actions)) + finished := make(chan error, 1) + start := time.Now().Add(10 * time.Millisecond) + go func() { + finished <- executeCuts(ctx, actions, start, func(a churnAction, planned time.Time) (func() error, error) { + _, err := r.cutAction(a, planned) + actual <- time.Now() + return func() error { <-release; return nil }, err + }) + }() + for i, a := range actions { + select { + case at := <-actual: + if delay := at.Sub(start.Add(a.offset)); delay < 0 || delay > 40*time.Millisecond { + t.Errorf("actual cut %d missed planned time by %s", i, delay) + } + case <-time.After(150 * time.Millisecond): + t.Error("actual cut waited for repair") + } + } + close(release) + if err := <-finished; err != nil { + t.Fatal(err) + } + if r.rec.summary.ActualHostCuts != 1 || r.rec.summary.ActualGuestCuts != 2 { + t.Fatalf("cuts were silently reduced: %+v", r.rec.summary) + } + for _, peer := range append(peers, hostPeer) { + _ = peer.SetReadDeadline(time.Now().Add(time.Second)) + if _, err := peer.Read(make([]byte, 1)); err != io.EOF { + t.Fatalf("scheduled transport did not actually close: %v", err) + } + } + if _, err := r.cutAction(actions[0], time.Now()); err == nil { + t.Fatal("already-disrupted transport was counted as another cut") + } +} + +func TestIndependentGuestCadenceDuringSlowRestore(t *testing.T) { + c := validConfig(t) + c.DurationSeconds = 3 + c.Active = true + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(context.Canceled) + r := &runState{cfg: c, ctx: ctx, cancel: cancel, rec: newRecorder(io.Discard, c)} + disrupted := &guest{expected: true} + healthy := &guest{id: 1, writers: []io.WriteCloser{timedWriter{calls: make(chan time.Time, 3)}}} + calls := healthy.writers[0].(timedWriter).calls + r.guests = []*guest{disrupted, healthy} + disrupted.ops.Lock() + defer disrupted.ops.Unlock() + began := time.Now() + r.activeWorkload(began, &r.activeWorkers) + for i := 0; i < 3; i++ { + select { + case at := <-calls: + if delay := at.Sub(began.Add(time.Duration(i) * time.Second)); delay < 0 || delay > 150*time.Millisecond { + t.Errorf("healthy guest cadence %d delayed %s", i, delay) + } + case <-time.After(1200 * time.Millisecond): + t.Errorf("healthy cadence %d blocked by restoration", i) + } + } + r.activeWorkers.Wait() + if r.rec.summary.ActiveFrames != 3 || r.rec.summary.PausedActiveFrames != 3 { + t.Fatalf("healthy work was paused: %+v", r.rec.summary) + } +} + +type timedWriter struct{ calls chan time.Time } + +func (w timedWriter) Write(b []byte) (int, error) { w.calls <- time.Now(); return len(b), nil } +func (w timedWriter) Close() error { return nil } + +func TestActualDialsRemainPacedAcrossReservationsAndRetries(t *testing.T) { + c := validConfig(t) + c.AdmissionsPerSecond = 20 + g := newGuard(c) + var mu sync.Mutex + var starts []time.Time + g.dial = func(context.Context, string, string) (net.Conn, error) { + mu.Lock() + starts = append(starts, time.Now()) + mu.Unlock() + a, b := net.Pipe() + _ = b.Close() + return a, nil + } + var wg sync.WaitGroup + for i := 0; i < 5; i++ { + wg.Add(1) + go func() { + defer wg.Done() + conn, err := g.connect(t.Context(), c.Manifest.Targets[0], "127.0.0.1:2222") + if err != nil { + t.Error(err) + return + } + _ = conn.Close() + }() + } + wg.Wait() + for i := 1; i < len(starts); i++ { + if starts[i].Sub(starts[i-1]) < 49*time.Millisecond { + t.Fatalf("external dials bunched: %v", starts) + } + } + if len(starts) != 5 { + t.Fatalf("missing real dial attempts: %d", len(starts)) + } +} + +func TestRetryBacklogIsGeneratorLimitedBeforeServiceStarts(t *testing.T) { + c := validConfig(t) + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(context.Canceled) + g := newGuard(c) + proxy, err := newTunnelProxy(ctx, g, c.Manifest.Targets[0]) + if err != nil { + t.Fatal(err) + } + defer proxy.close() + proxy.queueLimit = 30 * time.Millisecond + proxy.failure = cancel + g.next = time.Now().Add(time.Second) + req := httptest.NewRequest(http.MethodConnect, "http://127.0.0.1:2222", nil) + req.Host = "127.0.0.1:2222" + proxy.ServeHTTP(httptest.NewRecorder(), req) + if cause := context.Cause(ctx); cause == nil || cause.Error() != "generator_admission_backlog" { + t.Fatalf("retry backlog looked like a relay service failure: %v", cause) + } + if g.peak != 0 { + t.Fatal("timed-out generator queue started an external dial") + } +} + +func TestRealDeliveryEvidence(t *testing.T) { + c := localRelay(t) + c.DurationSeconds = 1 + c.GuestWSPercent = 100 + c.GuestsPerHost = 2 + c.ChannelsPerGuest = 4 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if err != nil { + t.Fatalf("real crossed delivery: %v\n%s", err, out.String()) + } + if !s.DeliveryComplete || s.RTTCount != 10 || s.InputFrames != 10 || s.InputFrames != int64(s.RTTCount+s.CanaryRTTCount+s.InterruptedReplies) || !s.CleanupComplete { + t.Fatalf("delivery evidence incomplete: %+v", s) + } + if dir := os.Getenv("CAPACITY_EVIDENCE_DIR"); dir != "" { + manifest, e := json.MarshalIndent(c.Manifest, "", " ") + if e != nil { + t.Fatal(e) + } + phase, e := json.MarshalIndent(c, "", " ") + if e != nil { + t.Fatal(e) + } + for name, data := range map[string][]byte{"manifest.json": manifest, "phase.json": phase, "events.jsonl": out.Bytes()} { + if err := os.WriteFile(filepath.Join(dir, name), data, 0600); err != nil { + t.Fatal(err) + } + } + t.Logf("retained successful local fixture evidence in %s (Darwin has no /proc telemetry)", dir) + } +} + +func TestHelperFinalBarrierFlushesAndQuiescesHostOutput(t *testing.T) { + in, input := io.Pipe() + output, out := io.Pipe() + defer input.Close() + defer output.Close() + finished := make(chan error, 1) + go func() { + err := helper(t.Context(), in, out, 0, true, 20*time.Millisecond) + _ = out.Close() + finished <- err + }() + observations := make(chan observation, 16) + scanned := make(chan error, 1) + go func() { + scanned <- scanFrames(output, func(o observation) error { observations <- o; return nil }, func(observation) {}) + }() + seen := 0 + for seen < 2 { + select { + case o := <-observations: + if o.kind != "H" || o.seq != seen { + t.Fatalf("host cadence/sequence missing: %+v", o) + } + seen++ + case <-time.After(100 * time.Millisecond): + t.Fatal("helper skipped a due host-output interval") + } + } + if _, err := input.Write(frame("F", 0, 0, time.Now(), 128)); err != nil { + t.Fatal(err) + } + select { + case o := <-observations: + if o.kind != "B" || o.seq != 2 { + t.Fatalf("final host-output count missing: %+v", o) + } + case <-time.After(100 * time.Millisecond): + t.Fatal("helper final barrier blocked") + } + select { + case o := <-observations: + t.Fatalf("helper emitted output after final barrier: %+v", o) + case <-time.After(60 * time.Millisecond): + } + _ = input.Close() + if err := <-finished; err != nil { + t.Fatal(err) + } + if err := <-scanned; err != io.EOF { + t.Fatal(err) + } +} + +func TestReplayCannotAcknowledgeOrRemeasureCurrentRequest(t *testing.T) { + now := time.Now() + q := &guest{pending: map[string]time.Time{"0:4": now}} + old := observation{kind: "I", id: 0, seq: 4, sent: now.Add(-time.Minute)} + if _, ok := q.acknowledge(old); ok { + t.Fatal("replayed timestamp acknowledged current input") + } + if len(q.pending) != 1 { + t.Fatal("replay removed outstanding input") + } + old.sent = now + if sent, ok := q.acknowledge(old); !ok || !sent.Equal(now) { + t.Fatal("matching reply was not acknowledged") + } + if _, ok := q.acknowledge(old); ok || q.acks != 1 { + t.Fatal("duplicate reply inflated RTT/count") + } +} + +func TestFreshChallengeSurvivesUnterminatedTerminalReplay(t *testing.T) { + // Terminal reconstruction may finish with a partial screen row rather than + // LF. Its old text and the first fresh helper frame then share a scan line. + replay := strings.TrimSuffix(string(frame("I", 4, 3, time.Unix(1, 0), 128)), "\n") + "\x1b[0m" + fresh := frame("A", 4, 11, time.Unix(2, 0), 128) + o, ok := parseFrame(replay + string(fresh)) + if !ok || o.kind != "A" || o.seq != 11 || !o.sent.Equal(time.Unix(2, 0)) { + t.Fatalf("fresh challenge swallowed by terminal replay suffix: %+v valid=%v", o, ok) + } +} + +func TestReconnectCannotReuseReplayedBaseline(t *testing.T) { + old := time.Unix(1, 0) + // The old transport's baseline can be replayed with an incomplete screen + // history. None of it belongs to this connection's delivery interval. + replay := string(frame("B", 0, 0, old, 128)) + string(frame("H", 0, 0, old, 1024)) + string(frame("H", 0, 2, old, 1024)) + c := scriptedHelper(t, "printf '%s' '"+replay+"'\nwhile IFS= read -r line; do\ncase \"$line\" in\n'CAP R '*) printf '%s\\n' \"$line\" | tr R A;;\n'CAP Q '*) printf '%s\\n' \"$line\" | tr Q B;;\n*) printf '%s\\n' \"$line\";;\nesac\ndone\n") + ctx, cancel := context.WithCancelCause(t.Context()) + dir, err := os.MkdirTemp("", "cap-replay-") + if err != nil { + t.Fatal(err) + } + r := &runState{ctx: ctx, cancel: cancel, cfg: c, rec: newRecorder(io.Discard, c), guard: newGuard(c), dir: dir, helper: c.HelperExecutable, guestReadyTimeout: 2 * time.Second, hosts: make([]*hostRun, 1)} + defer func() { + if err := r.cleanup(); err != nil { + t.Error(err) + } + }() + target := c.Manifest.Targets[0] + if err := r.startHost(0, target, nil); err != nil { + t.Fatal(err) + } + u, err := transport(target, 0, 0, 1, false) + if err != nil { + t.Fatal(err) + } + proxy, err := newTunnelProxy(ctx, r.guard, target) + if err != nil { + t.Fatal(err) + } + r.proxies = append(r.proxies, proxy) + q := &guest{host: r.hosts[0], target: target, endpoint: u, identity: signer(t), proxy: proxy, expected: true, generation: 1, barriers: map[int]*frameBarrier{0: {sent: old, done: make(chan struct{})}}} + r.guests = append(r.guests, q) + if err := q.connect(ctx, r); err != nil { + t.Fatalf("replayed old baseline prevented fresh readiness: %v", err) + } + q.mu.Lock() + acks, barrier := q.acks, q.barriers[0] + q.mu.Unlock() + if acks != 1 || barrier == nil || !barrier.sent.After(old) { + t.Fatalf("fresh connection did not establish its own live boundary: acks=%d barrier=%+v", acks, barrier) + } +} + +func TestShortWriteCannotBecomeOfferedOrOutstandingFrame(t *testing.T) { + c := validConfig(t) + r := &runState{cfg: c, rec: newRecorder(io.Discard, c)} + q := &guest{writers: []io.WriteCloser{shortFrameWriter{}}} + if err := q.active(r); err != io.ErrShortWrite { + t.Fatalf("short write became a successful active frame: err=%v active=%d", err, r.rec.summary.ActiveFrames) + } + if len(q.pending) != 0 || r.rec.summary.InputFrames != 0 || r.rec.summary.ActiveFrames != 0 { + t.Fatal("partial input inflated completion accounting") + } +} + +type shortFrameWriter struct{} + +func (shortFrameWriter) Write(b []byte) (int, error) { return len(b) / 2, nil } +func (shortFrameWriter) Close() error { return nil } + +func TestHelperResynchronizesAfterCutTruncatesInput(t *testing.T) { + // A transport cut may leave a partial old input before the next guest's + // complete challenge on the shared PTY. Fixed-size reads stay misaligned. + broken := frame("I", 0, 5, time.Unix(1, 0), 128)[:63] + challenge := frame("R", 4, 11, time.Unix(2, 0), 128) + var out bytes.Buffer + err := Helper(t.Context(), bytes.NewReader(append(broken, challenge...)), &out, 0, false) + if err != nil || !bytes.Contains(out.Bytes(), frame("A", 4, 11, time.Unix(2, 0), 128)) { + t.Fatalf("partial disrupted input prevented a fresh challenge: err=%v output=%q", err, out.Bytes()) + } +} + +func TestReadinessChallengeCanRetryWithinSameServiceBudget(t *testing.T) { + // A reconstruction can consume/transform the first challenge response. + // A repeated authenticated challenge must prove the same live channel. + c := scriptedHelper(t, "dd bs=128 count=1 of=/dev/null 2>/dev/null\ndd bs=128 count=1 2>/dev/null | tr R A\ndd bs=128 count=1 2>/dev/null | tr Q B\nwhile IFS= read -r line; do\ncase \"$line\" in\n'CAP F '*) printf '%s\\n' \"$line\" | tr F B;;\n*) printf '%s\\n' \"$line\";;\nesac\ndone\n") + c.DurationSeconds = 1 + s, err := Run(t.Context(), c, io.Discard) + if err != nil || !s.DeliveryComplete || s.EstablishedGuests != 1 || s.InputFrames != 2 || s.RTTCount != 2 || s.HandshakeFrames < 4 { + t.Fatalf("bounded repeated challenge did not prove live delivery: err=%v summary=%+v", err, s) + } +} + +func TestActiveWorkloadNeverOffersAfterSteadyEnd(t *testing.T) { + c := validConfig(t) + c.Active = true + c.DurationSeconds = 2 + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(context.Canceled) + r := &runState{cfg: c, ctx: ctx, cancel: cancel, rec: newRecorder(io.Discard, c), guests: []*guest{{writers: []io.WriteCloser{discardWriteCloser{}}}}} + r.activeWorkload(time.Now().Add(-2*time.Second), &r.activeWorkers) + r.activeWorkers.Wait() + if r.rec.summary.ActiveFrames != 0 || context.Cause(ctx) == nil { + t.Fatalf("missed offered intervals were sent in final drain: active=%d cause=%v", r.rec.summary.ActiveFrames, context.Cause(ctx)) + } +} diff --git a/internal/capacity/runner.go b/internal/capacity/runner.go index 197011a94..704918191 100644 --- a/internal/capacity/runner.go +++ b/internal/capacity/runner.go @@ -21,24 +21,32 @@ import ( "golang.org/x/crypto/ssh" ) -var errDuration = errors.New("duration") +const finalDrainLimit = 10 * time.Second +const retryQueueLimit = 5 * time.Second type runState struct { - canary *guest - clientsMu sync.RWMutex - telemetryMu sync.Mutex - lastTelemetry map[string]time.Time - cfg Config - guard *guard - rec *recorder - ctx context.Context - cancel context.CancelCauseFunc - hosts []*hostRun - guests []*guest - proxies []*tunnelProxy - workers sync.WaitGroup - helper string - dir string + hostReadyTimeout time.Duration + guestReadyTimeout time.Duration + canary *guest + clientsMu sync.RWMutex + telemetryMu sync.Mutex + lastTelemetry map[string]time.Time + cfg Config + guard *guard + rec *recorder + ctx context.Context + cancel context.CancelCauseFunc + hosts []*hostRun + guests []*guest + proxies []*tunnelProxy + workers sync.WaitGroup + activeWorkers sync.WaitGroup + helper string + dir string +} +type tunnelObservation struct { + event host.TunnelEvent + at time.Time } type hostRun struct { repair sync.Mutex @@ -49,7 +57,7 @@ type hostRun struct { done chan error mu sync.Mutex initial *api.GetSessionResponse - tunnel chan host.TunnelEvent + tunnel chan tunnelObservation } func (h *hostRun) session(ctx context.Context) (*api.GetSessionResponse, error) { @@ -105,9 +113,16 @@ func (r *runState) sendKind(q *guest, w io.Writer, stream, seq int, kind string) } q.pending[key] = now q.mu.Unlock() + } else if kind == "R" { + q.mu.Lock() + if q.challenges == nil { + q.challenges = map[string]time.Time{} + } + q.challenges[fmt.Sprintf("%d:%d:%d", stream, seq, now.UnixNano())] = now + q.mu.Unlock() } n, err := w.Write(frame(kind, stream, seq, now, 128)) - if n == 128 && err == nil { + if n == 128 { r.rec.mu.Lock() if kind == "I" { r.rec.summary.InputFrames++ @@ -117,6 +132,14 @@ func (r *runState) sendKind(q *guest, w io.Writer, stream, seq int, kind string) r.rec.mu.Unlock() } r.actual(int64(n)) + if n != 128 && err == nil { + err = io.ErrShortWrite + } + if kind == "I" && n != 128 { + q.mu.Lock() + delete(q.pending, key) + q.mu.Unlock() + } return err } func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { @@ -136,7 +159,7 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { return err } ready := make(chan string, 1) - h := &hostRun{id: index, socket: filepath.Join(dir, "admin.sock"), proxy: proxy, done: make(chan error, 1), tunnel: make(chan host.TunnelEvent, 16)} + h := &hostRun{id: index, socket: filepath.Join(dir, "admin.sock"), proxy: proxy, done: make(chan error, 1), tunnel: make(chan tunnelObservation, 16)} if identity == nil { identity, err = makeSigner() if err != nil { @@ -147,7 +170,16 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { if err != nil { return err } - h.h = &host.Host{KeepAliveDuration: 50 * time.Second, Host: u.String(), Command: []string{r.helper, "helper", strconv.Itoa(index), strconv.FormatBool(r.cfg.Active)}, CommandDir: dir, Signers: []ssh.Signer{identity}, IdentitiesOnly: true, HostKey: key, HostKeyCallback: pinnedKey(t), ProxyURL: proxy.url, AdminSocketFile: h.socket, AutoAccept: true, SFTPAutoAccept: true, ForwardingDisabled: true, ExecDisabled: true, Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), StopGrace: time.Second, + timeout := r.hostReadyTimeout + if timeout == 0 { + timeout = 12 * time.Second + } + proxy.serviceTimeout = timeout + proxy.failure = r.fail + proxy.queueEvent = func(wait time.Duration, limited bool) { + r.rec.event("admission_queue", map[string]any{"role": "host", "host": index, "retry": true, "queue_ms": float64(wait) / float64(time.Millisecond), "generator_limited": limited}) + } + h.h = &host.Host{Reconnect: host.ReconnectTiming{AttemptDeadline: 30*time.Second + retryQueueLimit}, KeepAliveDuration: 50 * time.Second, Host: u.String(), Command: []string{r.helper, "helper", strconv.Itoa(index), strconv.FormatBool(r.cfg.Active)}, CommandDir: dir, Signers: []ssh.Signer{identity}, IdentitiesOnly: true, HostKey: key, HostKeyCallback: pinnedKey(t), ProxyURL: proxy.url, AdminSocketFile: h.socket, AutoAccept: true, SFTPAutoAccept: true, ForwardingDisabled: true, ExecDisabled: true, Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), StopGrace: time.Second, SessionCreatedCallback: func(ctx context.Context, s *api.GetSessionResponse) error { if err := r.guard.session(ctx, s); err != nil { return err @@ -158,21 +190,35 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { return nil }, SessionReadyCallback: func(status string) { + proxy.ready() select { case ready <- status: default: } }, TunnelCallback: func(e host.TunnelEvent) { + if e.Up { + proxy.ready() + } r.rec.event("host_tunnel", map[string]any{"host": index, "up": e.Up, "can_reconnect": e.CanReconnect, "route_changed": e.RouteChanged}) select { - case h.tunnel <- e: + case h.tunnel <- tunnelObservation{event: e, at: time.Now()}: default: r.fail(errors.New("host tunnel event overflow")) } }, } r.hosts[index] = h + queuedAt := time.Now() + permit, err := r.guard.permit(r.ctx) + if err != nil { + h.done <- err + return err + } + proxy.mu.Lock() + proxy.permit = permit + proxy.mu.Unlock() + r.rec.event("admission_queue", map[string]any{"role": "host", "host": index, "queue_ms": float64(time.Since(queuedAt)) / float64(time.Millisecond)}) began := time.Now() go func() { err := h.h.Run(r.ctx) @@ -181,7 +227,17 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { r.fail(fmt.Errorf("host %d stopped: %v", index, err)) } }() - timer := time.NewTimer(12 * time.Second) + // The dial permit is consumed before the dependency starts waiting on + // remote registration. Key generation and intentional pacing are excluded. + select { + case began = <-proxy.dialStarted: + case err := <-h.done: + h.done <- err + return fmt.Errorf("host startup: %w", err) + case <-r.ctx.Done(): + return context.Cause(r.ctx) + } + timer := time.NewTimer(timeout) defer timer.Stop() select { case status := <-ready: @@ -202,7 +258,7 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { h.mu.Lock() s := h.initial h.mu.Unlock() - r.rec.event("host_ready", map[string]any{"host": index, "transport": u.Scheme, "node": s.NodeAddr, "session_id": s.SessionId, "ready_ms": float64(time.Since(began)) / float64(time.Millisecond)}) + r.rec.event("host_ready", map[string]any{"host": index, "transport": u.Scheme, "node": s.NodeAddr, "session_id": s.SessionId, "ready_ms": float64(time.Since(began)) / float64(time.Millisecond), "queue_ms": float64(permit.queued) / float64(time.Millisecond)}) return nil } func (r *runState) admit(identity ssh.Signer) error { @@ -244,7 +300,7 @@ func (r *runState) admit(identity ssh.Signer) error { r.rec.summary.EstablishedGuests++ r.rec.summary.EstablishedChannels += len(q.channels) r.rec.mu.Unlock() - r.rec.event("admission", map[string]any{"role": "guest", "guest": q.id, "ready_ms": float64(time.Since(began)) / float64(time.Millisecond)}) + r.rec.event("admission", map[string]any{"role": "guest", "guest": q.id, "ready_ms": float64(time.Since(began)-q.queueElapsed) / float64(time.Millisecond), "queue_ms": float64(q.queueElapsed) / float64(time.Millisecond)}) return nil }); err != nil { return err @@ -256,6 +312,7 @@ func (r *runState) admit(identity ssh.Signer) error { } func (r *runState) cleanup() error { r.cancel(context.Canceled) + r.activeWorkers.Wait() r.workers.Wait() for _, q := range r.guests { q.setExpected(true) @@ -311,7 +368,8 @@ func (r *runState) telemetry() { return r.guard.connect(ctx, target, addr) }} defer tr.CloseIdleConnections() - clients[t.MachineID] = &http.Client{Transport: tr, Timeout: 3 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return errors.New("telemetry redirect refused") }} + tr.ResponseHeaderTimeout = 3 * time.Second + clients[t.MachineID] = &http.Client{Transport: tr, CheckRedirect: func(*http.Request, []*http.Request) error { return errors.New("telemetry redirect refused") }} r.telemetryMu.Lock() r.lastTelemetry[t.MachineID] = time.Now() r.telemetryMu.Unlock() @@ -327,8 +385,10 @@ func (r *runState) telemetry() { } resp, err := clients[t.MachineID].Do(req) if err == nil { + bodyTimer := time.AfterFunc(3*time.Second, func() { _ = resp.Body.Close() }) body, e := io.ReadAll(io.LimitReader(resp.Body, (1<<20)+1)) _ = resp.Body.Close() + bodyTimer.Stop() if e == nil && resp.StatusCode == 200 && len(body) <= 1<<20 { r.telemetryMu.Lock() r.lastTelemetry[t.MachineID] = time.Now() @@ -454,25 +514,18 @@ func Run(ctx context.Context, c Config, out io.Writer) (result Summary, err erro rec.mu.Lock() rec.summary.SetupSeconds = steadyStarted.Sub(started).Seconds() rec.mu.Unlock() - steady := time.AfterFunc(time.Duration(c.DurationSeconds)*time.Second, func() { r.fail(errDuration) }) + steady := time.NewTimer(time.Duration(c.DurationSeconds) * time.Second) defer steady.Stop() defer func() { rec.mu.Lock() - rec.summary.SteadySeconds = time.Since(steadyStarted).Seconds() + rec.summary.SteadySeconds = min(time.Since(steadyStarted).Seconds(), float64(c.DurationSeconds)) rec.mu.Unlock() }() rec.event("steady_start", map[string]any{"duration_seconds": c.DurationSeconds}) if c.idleHostOnly() { r.canaries() } - if c.Active { - for _, q := range r.guests { - if e := q.active(r); e != nil { - r.fail(e) - break - } - } - } + r.activeWorkload(steadyStarted, &r.activeWorkers) if c.TransferConcurrency > 0 { r.workers.Add(1) go func() { @@ -491,46 +544,54 @@ func Run(ctx context.Context, c Config, out io.Writer) (result Summary, err erro } }() } - tick := time.NewTicker(time.Second) - defer tick.Stop() for { select { case <-runctx.Done(): - cause := context.Cause(runctx) + err = context.Cause(runctx) rec.mu.Lock() - rec.summary.StopReason = cause.Error() - err = cause - if errors.Is(cause, errDuration) { - err = nil - if c.TransferConcurrency > 0 && rec.summary.Transfers != c.TransferConcurrency { - err = errors.New("incomplete_transfers") - } - if c.ChurnPercent > 0 && !rec.summary.ChurnCompleted { - err = errors.New("incomplete_churn") - } - if c.Active && rec.summary.ActiveFrames+rec.summary.PausedActiveFrames != int64(c.Hosts*c.GuestsPerHost*c.DurationSeconds) { - err = errors.New("offered_workload_incomplete") - } - if err != nil { - rec.summary.StopReason = err.Error() - } - } else if errors.Is(cause, context.DeadlineExceeded) && ctx.Err() == nil { + rec.summary.StopReason = err.Error() + if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil { rec.summary.StopReason = "manifest_expired" } rec.mu.Unlock() return result, err - case <-tick.C: - if c.Active && time.Since(steadyStarted) < time.Duration(c.DurationSeconds)*time.Second { - for _, q := range r.guests { - if e := q.active(r); e != nil { - r.fail(e) - break - } - } + case <-steady.C: + rec.event("steady_end", map[string]any{"duration_seconds": c.DurationSeconds}) + drainStarted := time.Now() + // This independent bound also closes transports beneath a blocked final + // input/barrier write. Manifest expiry and all safety watches stay active. + drain := time.AfterFunc(finalDrainLimit, func() { r.fail(errors.New("terminal_delivery_timeout")) }) + r.activeWorkers.Wait() + if runctx.Err() == nil { + err = r.completeDelivery() } + drain.Stop() + if runctx.Err() != nil { + err = context.Cause(runctx) + } + rec.mu.Lock() + rec.summary.DrainSeconds = time.Since(drainStarted).Seconds() + if err == nil && c.TransferConcurrency > 0 && rec.summary.Transfers != c.TransferConcurrency { + err = errors.New("incomplete_transfers") + } + if err == nil && c.ChurnPercent > 0 && !rec.summary.ChurnCompleted { + err = errors.New("incomplete_churn") + } + if err == nil && c.Active && rec.summary.ActiveFrames+rec.summary.PausedActiveFrames != int64(c.Hosts*c.GuestsPerHost*c.DurationSeconds) { + err = errors.New("offered_workload_incomplete") + } + rec.summary.StopReason = "duration" + rec.summary.DeliveryComplete = err == nil + if err != nil { + rec.summary.StopReason = err.Error() + } + rec.mu.Unlock() + rec.event("delivery_complete", map[string]any{"complete": err == nil, "drain_seconds": time.Since(drainStarted).Seconds()}) + return result, err } } } + func safeConfig(c Config) any { copy := c copy.Manifest = Manifest{} diff --git a/internal/capacity/transfer.go b/internal/capacity/transfer.go index d04789998..2e0d2413c 100644 --- a/internal/capacity/transfer.go +++ b/internal/capacity/transfer.go @@ -141,34 +141,21 @@ func (r *runState) churn() error { before := canary.acks canary.mu.Unlock() start := time.Now().Add(time.Second) - var wg sync.WaitGroup - results := make(chan error, len(actions)) - for _, a := range actions { - if err := pause(r.ctx, time.Until(start.Add(a.offset))); err != nil { - wg.Wait() - return err + if err := executeCuts(r.ctx, actions, start, func(a churnAction, planned time.Time) (func() error, error) { + repair, err := r.cutAction(a, planned) + if err != nil { + r.fail(err) + return nil, err } - wg.Add(1) - go func(action churnAction) { - defer wg.Done() - var err error - if action.host { - err = r.cutHost(r.hosts[action.index]) - } else { - err = r.cutGuest(r.guests[action.index]) - } + return func() error { + err := repair() if err != nil { r.fail(err) } - results <- err - }(a) - } - wg.Wait() - close(results) - for err := range results { - if err != nil { return err - } + }, nil + }); err != nil { + return err } if !r.cfg.Active { if err := canary.active(r); err != nil { @@ -193,59 +180,151 @@ func (r *runState) churn() error { func (r *runState) restoreGuest(q *guest, began time.Time) error { q.close() q.mu.Lock() + interrupted := len(q.pending) q.pending = map[string]time.Time{} q.generation++ + disruption := q.disruption + q.activeOffered = false + if !q.disruptedAt.IsZero() { + began = q.disruptedAt + } q.mu.Unlock() + r.rec.mu.Lock() + r.rec.summary.InterruptedReplies += interrupted + r.rec.mu.Unlock() + var queue time.Duration var last error for tries := 0; tries < 3; tries++ { last = q.connect(r.ctx, r) + queue += q.queueElapsed + q.mu.Lock() + superseded := disruption != q.disruption + q.mu.Unlock() + if superseded { + q.close() + return nil + } // a later actual host cut owns repair. if last == nil { q.setExpected(false) r.rec.mu.Lock() r.rec.summary.Reconnects++ r.rec.mu.Unlock() - r.rec.event("guest_restored", map[string]any{"guest": q.id, "host": q.host.id, "generation": q.generation, "restoration_ms": float64(time.Since(began)) / float64(time.Millisecond), "same_identity": true}) + r.rec.event("guest_restored", map[string]any{"guest": q.id, "host": q.host.id, "generation": q.generation, "restoration_ms": float64(time.Since(began)) / float64(time.Millisecond), "same_identity": true, "queue_ms": float64(queue) / float64(time.Millisecond), "service_ms": float64(time.Since(began)-queue) / float64(time.Millisecond)}) return nil } + r.rec.event("guest_restore_attempt_failed", map[string]any{"guest": q.id, "host": q.host.id, "attempt": tries + 1, "error": last.Error()}) q.close() + q.mu.Lock() + interrupted := len(q.pending) + q.pending = map[string]time.Time{} + q.challenges = map[string]time.Time{} + q.mu.Unlock() + r.rec.mu.Lock() + r.rec.summary.InterruptedReplies += interrupted + r.rec.mu.Unlock() if err := pause(r.ctx, time.Second); err != nil { return err } } return fmt.Errorf("guest restoration failed: %w", last) } -func (r *runState) cutGuest(q *guest) error { - q.host.repair.Lock() - defer q.host.repair.Unlock() - q.setExpected(true) + +// Cutting is synchronous and never takes the repair/connection-operation +// locks. Repairs may serialize, while every actual transport cut keeps its +// independently planned time. Guest selections run before host selections: +// an overlap therefore cuts a real guest transport before its host is cut. +func (r *runState) cutAction(a churnAction, planned time.Time) (func() error, error) { began := time.Now() - r.rec.event("transport_cut", map[string]any{"role": "guest", "guest": q.id, "host": q.host.id}) - q.ops.Lock() - if q.raw != nil { - _ = q.raw.Close() + if a.host { + h := r.hosts[a.index] + overlaps := 0 + var newlyDisrupted []*guest + for _, q := range r.guests { + if q.host != h { + continue + } + q.mu.Lock() + if q.expected { + overlaps++ + } + if !q.expected { + newlyDisrupted = append(newlyDisrupted, q) + } + q.expected = true + q.disruption++ + q.mu.Unlock() + } + n, actual := h.proxy.cut() + began = actual + if n != 1 { + return nil, fmt.Errorf("host cut requires one live transport, got %d", n) + } + for _, q := range newlyDisrupted { + q.mu.Lock() + q.disruptedAt = actual + q.mu.Unlock() + } + for _, q := range r.guests { + if q.host != h { + continue + } + q.mu.Lock() + if q.raw != nil { + _ = q.raw.Close() + } + q.transportLive = false + q.mu.Unlock() + } + r.rec.mu.Lock() + r.rec.summary.ActualHostCuts++ + r.rec.mu.Unlock() + r.rec.event("transport_cut", map[string]any{"role": "host", "host": h.id, "planned_at": planned.UTC(), "actual_at": began.UTC(), "schedule_delay_ms": float64(began.Sub(planned)) / float64(time.Millisecond), "actual_transports": n, "already_disrupted_guests": overlaps}) + queued := h.proxy.queued() + return func() error { h.repair.Lock(); defer h.repair.Unlock(); return r.repairHost(h, began, queued) }, nil } - q.ops.Unlock() - return r.restoreGuest(q, began) + q := r.guests[a.index] + q.mu.Lock() + if q.expected || !q.transportLive || q.raw == nil { + q.mu.Unlock() + return nil, fmt.Errorf("guest %d already disrupted before its planned cut", q.id) + } + q.expected = true + q.disruption++ + q.transportLive = false + err := q.raw.Close() + began = time.Now() + q.disruptedAt = began + q.mu.Unlock() + if err != nil { + return nil, fmt.Errorf("guest %d transport was already closed: %w", q.id, err) + } + r.rec.mu.Lock() + r.rec.summary.ActualGuestCuts++ + r.rec.mu.Unlock() + r.rec.event("transport_cut", map[string]any{"role": "guest", "guest": q.id, "host": q.host.id, "planned_at": planned.UTC(), "actual_at": began.UTC(), "schedule_delay_ms": float64(began.Sub(planned)) / float64(time.Millisecond), "actual_transports": 1}) + return func() error { q.host.repair.Lock(); defer q.host.repair.Unlock(); return r.restoreGuest(q, began) }, nil } -func (r *runState) cutHost(h *hostRun) error { - h.repair.Lock() - defer h.repair.Unlock() +func (r *runState) cutGuest(q *guest) error { + repair, err := r.cutAction(churnAction{index: q.id}, time.Now()) + if err != nil { + return err + } + return repair() +} +func (r *runState) repairHost(h *hostRun, began time.Time, queuedAtCut time.Duration) error { var guests []*guest for _, q := range r.guests { if q.host == h { - q.setExpected(true) guests = append(guests, q) } } - began := time.Now() - r.rec.event("transport_cut", map[string]any{"role": "host", "host": h.id}) - h.proxy.cut() ctx, cancel := context.WithTimeout(r.ctx, 15*time.Second) defer cancel() down := false for { select { - case e := <-h.tunnel: + case observation := <-h.tunnel: + e := observation.event if !e.Up { down = true if !e.CanReconnect { @@ -263,7 +342,7 @@ func (r *runState) cutHost(h *hostRun) error { r.rec.mu.Lock() r.rec.summary.Reconnects++ r.rec.mu.Unlock() - r.rec.event("host_restored", map[string]any{"host": h.id, "session_id": s.SessionId, "restoration_ms": float64(time.Since(began)) / float64(time.Millisecond), "same_host_run": true}) + r.rec.event("host_restored", map[string]any{"host": h.id, "session_id": s.SessionId, "restoration_ms": float64(observation.at.Sub(began)) / float64(time.Millisecond), "usable_at": observation.at.UTC(), "same_host_run": true, "queue_ms": float64(h.proxy.queued()-queuedAtCut) / float64(time.Millisecond), "service_ms": float64(observation.at.Sub(began)-(h.proxy.queued()-queuedAtCut)) / float64(time.Millisecond)}) for _, q := range guests { if err = r.restoreGuest(q, began); err != nil { return err @@ -278,6 +357,9 @@ func (r *runState) cutHost(h *hostRun) error { } func pause(ctx context.Context, d time.Duration) error { + if err := ctx.Err(); err != nil { + return err + } timer := time.NewTimer(d) defer timer.Stop() select { @@ -298,12 +380,12 @@ func churnSchedule(hosts, guests, percent int) []churnAction { hn := hosts * percent / 100 gn := hosts * guests * percent / 100 var actions []churnAction - for i := 0; i < hn; i++ { - actions = append(actions, churnAction{host: true, index: i * hosts / hn}) - } for i := 0; i < gn; i++ { actions = append(actions, churnAction{index: i * hosts * guests / gn}) } + for i := 0; i < hn; i++ { + actions = append(actions, churnAction{host: true, index: i * hosts / hn}) + } for i := range actions { if len(actions) > 1 { actions[i].offset = time.Duration(i) * 10 * time.Second / time.Duration(len(actions)-1) @@ -311,3 +393,36 @@ func churnSchedule(hosts, guests, percent int) []churnAction { } return actions } + +// The cut callback must only sever transports and return deferred repair work. +// This seam lets tests deliberately stall repairs while timing real cuts. +func executeCuts(ctx context.Context, actions []churnAction, start time.Time, cut func(churnAction, time.Time) (func() error, error)) error { + var wg sync.WaitGroup + results := make(chan error, len(actions)) + var cutErr error + for _, a := range actions { + planned := start.Add(a.offset) + if err := pause(ctx, time.Until(planned)); err != nil { + cutErr = err + break + } + repair, err := cut(a, planned) + if err != nil { + cutErr = err + break + } + wg.Add(1) + go func() { defer wg.Done(); results <- repair() }() + } + wg.Wait() + close(results) + if cutErr != nil { + return cutErr + } + for err := range results { + if err != nil { + return err + } + } + return nil +} diff --git a/internal/capacity/workload.go b/internal/capacity/workload.go index 661445a4f..425af206d 100644 --- a/internal/capacity/workload.go +++ b/internal/capacity/workload.go @@ -28,12 +28,15 @@ type observation struct { } func parseFrame(line string) (observation, bool) { - i := strings.Index(line, "CAP ") + // Reconstruction can end mid-row, placing old screen text before the + // first fresh frame on this scan line. Only the final complete tag can be + // that live frame; timestamp-bound challenges/barriers reject stale replay. + i := strings.LastIndex(line, "CAP ") if i < 0 { return observation{}, false } fields := strings.Fields(line[i:]) - if len(fields) != 6 || (fields[1] != "I" && fields[1] != "H" && fields[1] != "R" && fields[1] != "A") { + if len(fields) != 6 || (fields[1] != "I" && fields[1] != "H" && fields[1] != "R" && fields[1] != "A" && fields[1] != "Q" && fields[1] != "F" && fields[1] != "B") { return observation{}, false } id, e1 := strconv.Atoi(fields[2]) @@ -49,6 +52,9 @@ func parseFrame(line string) (observation, bool) { // input frames and, in active mode, emits exactly 1 KiB every ten seconds. // It never starts a shell or interprets received text as a command. func Helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active bool) error { + return helper(ctx, in, out, hostID, active, 10*time.Second) +} +func helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active bool, outputInterval time.Duration) error { if f, ok := in.(*os.File); ok && term.IsTerminal(int(f.Fd())) { state, err := term.MakeRaw(int(f.Fd())) if err != nil { @@ -66,19 +72,30 @@ func Helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active done := make(chan struct{}) go func() { defer close(done) - for { - b := make([]byte, 128) - _, err := io.ReadFull(in, b) - if err != nil { - failed <- err - return + scanner := bufio.NewScanner(in) + scanner.Buffer(make([]byte, 4096), 1<<20) + for scanner.Scan() { + line := scanner.Text() + i := strings.LastIndex(line, "CAP ") + if i < 0 || len(line)-i != 127 { + continue } + o, ok := parseFrame(line[i:]) + if !ok || (o.kind != "I" && o.kind != "R" && o.kind != "Q" && o.kind != "F") { + continue + } + b := []byte(line[i:] + "\n") select { case received <- b: case <-ctx.Done(): return } } + err := scanner.Err() + if err == nil { + err = io.EOF + } + failed <- err }() defer func() { cancel() @@ -87,9 +104,23 @@ func Helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active } <-done }() - tick := time.NewTicker(10 * time.Second) + started := time.Now() + tick := time.NewTicker(outputInterval) defer tick.Stop() seq := 0 + quiesced := false + emitDue := func(now time.Time) error { + if active && !quiesced { + due := int(now.Sub(started) / outputInterval) + for seq < due { + if _, err := out.Write(frame("H", hostID, seq, now, 1024)); err != nil { + return err + } + seq++ + } + } + return nil + } for { select { case <-ctx.Done(): @@ -100,6 +131,18 @@ func Helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active } return err case b := <-received: + if o, ok := parseFrame(string(b)); ok && (o.kind == "Q" || o.kind == "F") { + if err := emitDue(time.Now()); err != nil { + return err + } + if o.kind == "F" { + quiesced = true + } + if _, err := out.Write(frame("B", o.id, seq, o.sent, 128)); err != nil { + return err + } + continue + } if len(b) > 4 && b[4] == 'R' { b[4] = 'A' } @@ -107,11 +150,8 @@ func Helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active return err } case now := <-tick.C: - if active { - if _, err := out.Write(frame("H", hostID, seq, now, 1024)); err != nil { - return err - } - seq++ + if err := emitDue(now); err != nil { + return err } } } From 1dca18e58c3cdd6ed5f34299d4fae967c0e606fa Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 16:41:49 -0700 Subject: [PATCH 03/16] fix(capacity): cross mixed host and guest transports --- internal/capacity/README.md | 9 +++- internal/capacity/mixed_test.go | 96 +++++++++++++++++++++++++++++++++ internal/capacity/runner.go | 4 +- 3 files changed, 106 insertions(+), 3 deletions(-) create mode 100644 internal/capacity/mixed_test.go diff --git a/internal/capacity/README.md b/internal/capacity/README.md index 8f442b546..ba49ffcf7 100644 --- a/internal/capacity/README.md +++ b/internal/capacity/README.md @@ -89,8 +89,13 @@ Phase: `host_public`/`entry_public` select the optional public URLs; otherwise the private URLs are used. Both public URLs must be present when selected. Transport -percentages assign the first `floor(role_count * percentage / 100)` clients to -WS/WSS, the rest to SSH. JSONL records each role, actual node and transport. +percentages assign exactly `floor(role_count * percentage / 100)` clients of +each role to WS/WSS, the rest to SSH. Hosts select WS/WSS from the beginning of +host order; guests select from the end of guest order (grouped by host), without +changing IDs or host association. Thus a 50/50 workload with an even host count +includes SSH hosts with WS guests and WS hosts with SSH guests. All-SSH and +all-WS phases cover same-transport pairs. JSONL records each role, actual node +and transport. Idle phases with zero guests per host reserve one extra external connection for a transient one-channel guest canary at readiness and every 30 seconds. Canary hosts are sampled round-robin; canary counts, maximum ready latency and RTT diff --git a/internal/capacity/mixed_test.go b/internal/capacity/mixed_test.go new file mode 100644 index 000000000..4cd1eee29 --- /dev/null +++ b/internal/capacity/mixed_test.go @@ -0,0 +1,96 @@ +package capacity + +import ( + "bytes" + "encoding/json" + "io" + "maps" + "os" + "path/filepath" + "testing" +) + +// Ascending transport selection for both grouped roles loses crossed pairs. +func TestRealMixedTransportPairing(t *testing.T) { + for _, tc := range []struct { + name string + hosts, guests int + wantHost, wantGuest map[string]int + wantPairs map[string]int + }{ + {"even", 4, 1, map[string]int{"ssh": 2, "ws": 2}, map[string]int{"ssh": 2, "ws": 2}, map[string]int{"ssh->ws": 2, "ws->ssh": 2}}, + {"rounded-host-count", 3, 2, map[string]int{"ssh": 2, "ws": 1}, map[string]int{"ssh": 3, "ws": 3}, map[string]int{"ssh->ws": 3, "ws->ssh": 2, "ssh->ssh": 1}}, + } { + t.Run(tc.name, func(t *testing.T) { + c := localRelay(t) + c.Hosts, c.GuestsPerHost = tc.hosts, tc.guests + c.HostWSPercent, c.GuestWSPercent = 50, 50 + c.DurationSeconds = 1 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if err != nil { + t.Fatalf("mixed workload: %v\n%s", err, out.String()) + } + if s.EstablishedHosts != tc.hosts || s.EstablishedGuests != tc.hosts*tc.guests || s.EstablishedChannels != tc.hosts*tc.guests { + t.Fatalf("admission totals: %+v", s) + } + if !s.DeliveryComplete || !s.CleanupComplete || s.StopReason != "duration" || s.ActiveFrames == 0 || s.ReceivedFrames == 0 || s.InputFrames != int64(s.RTTCount) { + t.Fatalf("mixed delivery or cleanup incomplete: %+v", s) + } + hosts := map[int]string{} + hostCounts, guestCounts, pairs := map[string]int{}, map[string]int{}, map[string]int{} + guests := map[int]bool{} + decoder := json.NewDecoder(bytes.NewReader(out.Bytes())) + for { + var event struct { + Event string `json:"event"` + Host int `json:"host"` + Guest int `json:"guest"` + Transport string `json:"transport"` + } + if err := decoder.Decode(&event); err == io.EOF { + break + } else if err != nil { + t.Fatal(err) + } + switch event.Event { + case "host_ready": + if _, exists := hosts[event.Host]; exists || event.Host < 0 || event.Host >= tc.hosts { + t.Fatalf("duplicate or invalid host identity: %+v", event) + } + hosts[event.Host] = event.Transport + hostCounts[event.Transport]++ + case "guest_ready": + if guests[event.Guest] || event.Guest < 0 || event.Guest >= tc.hosts*tc.guests || event.Host != event.Guest/tc.guests || hosts[event.Host] == "" { + t.Fatalf("duplicate, invalid or reassigned guest identity: %+v", event) + } + guests[event.Guest] = true + guestCounts[event.Transport]++ + pairs[hosts[event.Host]+"->"+event.Transport]++ + } + } + if !maps.Equal(hostCounts, tc.wantHost) || !maps.Equal(guestCounts, tc.wantGuest) { + t.Errorf("role transport counts: hosts=%v guests=%v; want hosts=%v guests=%v", hostCounts, guestCounts, tc.wantHost, tc.wantGuest) + } + if !maps.Equal(pairs, tc.wantPairs) { + t.Fatalf("mixed host->guest transport pairs = %v; want %v", pairs, tc.wantPairs) + } + t.Logf("actual admissions: hosts=%v guests=%v host->guest=%v; delivery and cleanup complete", hostCounts, guestCounts, pairs) + if dir := os.Getenv("CAPACITY_EVIDENCE_DIR"); dir != "" { + manifest, err := json.MarshalIndent(c.Manifest, "", " ") + if err != nil { + t.Fatal(err) + } + phase, err := json.MarshalIndent(c, "", " ") + if err != nil { + t.Fatal(err) + } + for name, data := range map[string][]byte{"manifest.json": manifest, "phase.json": phase, "events.jsonl": out.Bytes()} { + if err := os.WriteFile(filepath.Join(dir, "mixed-"+tc.name+"-"+name), data, 0600); err != nil { + t.Fatal(err) + } + } + } + }) + } +} diff --git a/internal/capacity/runner.go b/internal/capacity/runner.go index 704918191..8f2d22139 100644 --- a/internal/capacity/runner.go +++ b/internal/capacity/runner.go @@ -271,7 +271,9 @@ func (r *runState) admit(identity ssh.Signer) error { for _, h := range r.hosts { for j := 0; j < r.cfg.GuestsPerHost; j++ { id := len(r.guests) - u, err := transport(et, r.cfg.GuestWSPercent, id, r.cfg.Hosts*r.cfg.GuestsPerHost, r.cfg.EntryPublic) + total := r.cfg.Hosts * r.cfg.GuestsPerHost + // Select guests from the opposite end so a 50/50 mix crosses host transports. + u, err := transport(et, r.cfg.GuestWSPercent, total-1-id, total, r.cfg.EntryPublic) if err != nil { return err } From 8b97ae5aad5a7d49ce089f327a0eae03d3c5148d Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 17:01:36 -0700 Subject: [PATCH 04/16] fix(capacity): authorize guests and reject unexpected host loss --- internal/capacity/canary.go | 6 +- internal/capacity/review_regression_test.go | 53 +++- internal/capacity/runner.go | 71 ++++-- internal/capacity/transfer.go | 8 + internal/capacity/whole_review_test.go | 266 ++++++++++++++++++++ 5 files changed, 367 insertions(+), 37 deletions(-) create mode 100644 internal/capacity/whole_review_test.go diff --git a/internal/capacity/canary.go b/internal/capacity/canary.go index 3ede6391b..5c4d8f0d3 100644 --- a/internal/capacity/canary.go +++ b/internal/capacity/canary.go @@ -10,11 +10,7 @@ import ( func (r *runState) hostCanary(sample int) error { et, _ := r.cfg.target(r.cfg.EntryTarget) if r.canary == nil { - identity, err := makeSigner() - if err != nil { - return err - } - q := &guest{id: 0, target: et, identity: identity, canary: true, expected: true} + q := &guest{id: 0, target: et, identity: r.guestIdentity, canary: true, expected: true} r.clientsMu.Lock() r.guests = append(r.guests, q) r.clientsMu.Unlock() diff --git a/internal/capacity/review_regression_test.go b/internal/capacity/review_regression_test.go index 545a2bf13..2452e9956 100644 --- a/internal/capacity/review_regression_test.go +++ b/internal/capacity/review_regression_test.go @@ -16,10 +16,47 @@ import ( "time" ) +const silentAfterReadinessScript = `inputs=0 +while IFS= read -r line; do +case "$line" in +'CAP R '*) printf 'CAP A %s\n' "${line#CAP R }";; +'CAP Q '*) printf 'CAP B %s\n' "${line#CAP Q }";; +'CAP I '*) + if [ "$inputs" -eq 0 ]; then printf '%s\n' "$line"; fi + inputs=$((inputs + 1));; +esac +done +` +const lateFinalReplyScript = `inputs=0 +while IFS= read -r line; do +case "$line" in +'CAP R '*) printf 'CAP A %s\n' "${line#CAP R }";; +'CAP Q '*) printf 'CAP B %s\n' "${line#CAP Q }";; +'CAP I '*) + if [ "$inputs" -eq 1 ]; then sleep 1.2; fi + if [ "$inputs" -le 1 ]; then printf '%s\n' "$line"; fi + inputs=$((inputs + 1));; +'CAP F '*) printf 'CAP B %s\n' "${line#CAP F }";; +esac +done +` +const retryReadinessScript = `challenges=0 +while IFS= read -r line; do +case "$line" in +'CAP R '*) + if [ "$challenges" -gt 0 ]; then printf 'CAP A %s\n' "${line#CAP R }"; fi + challenges=$((challenges + 1));; +'CAP Q '*) printf 'CAP B %s\n' "${line#CAP Q }";; +'CAP F '*) printf 'CAP B %s\n' "${line#CAP F }";; +'CAP I '*) printf '%s\n' "$line";; +esac +done +` + func TestSilentAfterReadinessFailsDelivery(t *testing.T) { c := localRelay(t) helper := filepath.Join(t.TempDir(), "silent-helper") - script := "#!/bin/sh\nstty raw -echo\ndd bs=128 count=1 2>/dev/null | tr R A\ndd bs=128 count=1 2>/dev/null | tr Q B\ndd bs=128 count=1 2>/dev/null\ncat >/dev/null\n" + script := "#!/bin/sh\nstty raw -echo\n" + silentAfterReadinessScript if err := os.WriteFile(helper, []byte(script), 0700); err != nil { t.Fatal(err) } @@ -122,6 +159,7 @@ func pacedAdmissionFixture(t *testing.T, ws int) { } }() r.hosts = make([]*hostRun, 1) + r.guestIdentity = signer(t) r.guard.next = time.Now().Add(450 * time.Millisecond) target := c.Manifest.Targets[0] if err := r.startHost(0, target, nil); err != nil { @@ -131,16 +169,12 @@ func pacedAdmissionFixture(t *testing.T, ws int) { if err != nil { t.Fatal(err) } - key, err := makeSigner() - if err != nil { - t.Fatal(err) - } proxy, err := newTunnelProxy(ctx, r.guard, target) if err != nil { t.Fatal(err) } r.proxies = append(r.proxies, proxy) - q := &guest{host: r.hosts[0], target: target, endpoint: u, identity: key, proxy: proxy} + q := &guest{host: r.hosts[0], target: target, endpoint: u, identity: r.guestIdentity, proxy: proxy} r.guests = append(r.guests, q) r.guard.mu.Lock() r.guard.next = time.Now().Add(450 * time.Millisecond) @@ -180,7 +214,7 @@ func scriptedHelper(t *testing.T, body string) Config { } func TestLateFinalReplyIsDrainedAndMeasured(t *testing.T) { - c := scriptedHelper(t, "dd bs=128 count=1 2>/dev/null | tr R A\ndd bs=128 count=1 2>/dev/null | tr Q B\ndd bs=128 count=1 2>/dev/null\ndd bs=128 count=1 2>/dev/null | (sleep 1.2; cat)\ndd bs=128 count=1 2>/dev/null | tr F B\ncat >/dev/null\n") + c := scriptedHelper(t, lateFinalReplyScript) c.DurationSeconds = 1 s, err := Run(t.Context(), c, io.Discard) if err != nil { @@ -508,6 +542,7 @@ func TestReconnectCannotReuseReplayedBaseline(t *testing.T) { } }() target := c.Manifest.Targets[0] + r.guestIdentity = signer(t) if err := r.startHost(0, target, nil); err != nil { t.Fatal(err) } @@ -520,7 +555,7 @@ func TestReconnectCannotReuseReplayedBaseline(t *testing.T) { t.Fatal(err) } r.proxies = append(r.proxies, proxy) - q := &guest{host: r.hosts[0], target: target, endpoint: u, identity: signer(t), proxy: proxy, expected: true, generation: 1, barriers: map[int]*frameBarrier{0: {sent: old, done: make(chan struct{})}}} + q := &guest{host: r.hosts[0], target: target, endpoint: u, identity: r.guestIdentity, proxy: proxy, expected: true, generation: 1, barriers: map[int]*frameBarrier{0: {sent: old, done: make(chan struct{})}}} r.guests = append(r.guests, q) if err := q.connect(ctx, r); err != nil { t.Fatalf("replayed old baseline prevented fresh readiness: %v", err) @@ -565,7 +600,7 @@ func TestHelperResynchronizesAfterCutTruncatesInput(t *testing.T) { func TestReadinessChallengeCanRetryWithinSameServiceBudget(t *testing.T) { // A reconstruction can consume/transform the first challenge response. // A repeated authenticated challenge must prove the same live channel. - c := scriptedHelper(t, "dd bs=128 count=1 of=/dev/null 2>/dev/null\ndd bs=128 count=1 2>/dev/null | tr R A\ndd bs=128 count=1 2>/dev/null | tr Q B\nwhile IFS= read -r line; do\ncase \"$line\" in\n'CAP F '*) printf '%s\\n' \"$line\" | tr F B;;\n*) printf '%s\\n' \"$line\";;\nesac\ndone\n") + c := scriptedHelper(t, retryReadinessScript) c.DurationSeconds = 1 s, err := Run(t.Context(), c, io.Discard) if err != nil || !s.DeliveryComplete || s.EstablishedGuests != 1 || s.InputFrames != 2 || s.RTTCount != 2 || s.HandshakeFrames < 4 { diff --git a/internal/capacity/runner.go b/internal/capacity/runner.go index 8f2d22139..2703dbbb2 100644 --- a/internal/capacity/runner.go +++ b/internal/capacity/runner.go @@ -28,6 +28,7 @@ type runState struct { hostReadyTimeout time.Duration guestReadyTimeout time.Duration canary *guest + guestIdentity ssh.Signer clientsMu sync.RWMutex telemetryMu sync.Mutex lastTelemetry map[string]time.Time @@ -49,15 +50,40 @@ type tunnelObservation struct { at time.Time } type hostRun struct { - repair sync.Mutex - id int - socket string - proxy *tunnelProxy - h *host.Host - done chan error - mu sync.Mutex - initial *api.GetSessionResponse - tunnel chan tunnelObservation + repair sync.Mutex + id int + socket string + proxy *tunnelProxy + h *host.Host + done chan error + mu sync.Mutex + initial *api.GetSessionResponse + established bool + expectedCut bool + tunnel chan tunnelObservation +} + +func (r *runState) hostTunnel(h *hostRun, e host.TunnelEvent) { + h.mu.Lock() + unexpected := !e.Up && h.established && !h.expectedCut + if e.Up { + // Up means the host's guest door is live again. A later loss is a + // separate outage, even while planned guest repairs are still running. + h.expectedCut = false + } + h.mu.Unlock() + if unexpected && r.ctx.Err() == nil { + r.fail(fmt.Errorf("unexpected host transport loss: host=%d", h.id)) + } + if e.Up { + h.proxy.ready() + } + r.rec.event("host_tunnel", map[string]any{"host": h.id, "up": e.Up, "can_reconnect": e.CanReconnect, "route_changed": e.RouteChanged}) + select { + case h.tunnel <- tunnelObservation{event: e, at: time.Now()}: + default: + r.fail(errors.New("host tunnel event overflow")) + } } func (h *hostRun) session(ctx context.Context) (*api.GetSessionResponse, error) { @@ -180,6 +206,7 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { r.rec.event("admission_queue", map[string]any{"role": "host", "host": index, "retry": true, "queue_ms": float64(wait) / float64(time.Millisecond), "generator_limited": limited}) } h.h = &host.Host{Reconnect: host.ReconnectTiming{AttemptDeadline: 30*time.Second + retryQueueLimit}, KeepAliveDuration: 50 * time.Second, Host: u.String(), Command: []string{r.helper, "helper", strconv.Itoa(index), strconv.FormatBool(r.cfg.Active)}, CommandDir: dir, Signers: []ssh.Signer{identity}, IdentitiesOnly: true, HostKey: key, HostKeyCallback: pinnedKey(t), ProxyURL: proxy.url, AdminSocketFile: h.socket, AutoAccept: true, SFTPAutoAccept: true, ForwardingDisabled: true, ExecDisabled: true, Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), StopGrace: time.Second, + AuthorizedKeys: []*host.AuthorizedKey{{PublicKeys: []ssh.PublicKey{r.guestIdentity.PublicKey()}}}, SessionCreatedCallback: func(ctx context.Context, s *api.GetSessionResponse) error { if err := r.guard.session(ctx, s); err != nil { return err @@ -190,6 +217,9 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { return nil }, SessionReadyCallback: func(status string) { + h.mu.Lock() + h.established = status == "ready" + h.mu.Unlock() proxy.ready() select { case ready <- status: @@ -197,15 +227,7 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { } }, TunnelCallback: func(e host.TunnelEvent) { - if e.Up { - proxy.ready() - } - r.rec.event("host_tunnel", map[string]any{"host": index, "up": e.Up, "can_reconnect": e.CanReconnect, "route_changed": e.RouteChanged}) - select { - case h.tunnel <- tunnelObservation{event: e, at: time.Now()}: - default: - r.fail(errors.New("host tunnel event overflow")) - } + r.hostTunnel(h, e) }, } r.hosts[index] = h @@ -262,6 +284,13 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { return nil } func (r *runState) admit(identity ssh.Signer) error { + // One fresh in-memory guest identity belongs to this run, independent of + // host registration, host session and relay keys. Reconnects retain it. + var err error + r.guestIdentity, err = makeSigner() + if err != nil { + return err + } ht, _ := r.cfg.target(r.cfg.HostTarget) et, _ := r.cfg.target(r.cfg.EntryTarget) r.hosts = make([]*hostRun, r.cfg.Hosts) @@ -277,16 +306,12 @@ func (r *runState) admit(identity ssh.Signer) error { if err != nil { return err } - key, err := makeSigner() - if err != nil { - return err - } proxy, err := newTunnelProxy(r.ctx, r.guard, et) if err != nil { return err } r.proxies = append(r.proxies, proxy) - q := &guest{id: id, host: h, target: et, endpoint: u, identity: key, proxy: proxy} + q := &guest{id: id, host: h, target: et, endpoint: u, identity: r.guestIdentity, proxy: proxy} r.clientsMu.Lock() r.guests = append(r.guests, q) r.clientsMu.Unlock() diff --git a/internal/capacity/transfer.go b/internal/capacity/transfer.go index 2e0d2413c..33181c02d 100644 --- a/internal/capacity/transfer.go +++ b/internal/capacity/transfer.go @@ -254,7 +254,15 @@ func (r *runState) cutAction(a churnAction, planned time.Time) (func() error, er q.disruption++ q.mu.Unlock() } + // Publish the expected outage before closing any socket. Keep the + // callback classification synchronized with the actual-cut result. + h.mu.Lock() + h.expectedCut = true n, actual := h.proxy.cut() + if n != 1 { + h.expectedCut = false + } + h.mu.Unlock() began = actual if n != 1 { return nil, fmt.Errorf("host cut requires one live transport, got %d", n) diff --git a/internal/capacity/whole_review_test.go b/internal/capacity/whole_review_test.go new file mode 100644 index 000000000..5812a450f --- /dev/null +++ b/internal/capacity/whole_review_test.go @@ -0,0 +1,266 @@ +package capacity + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/owenthereal/upterm/host" + "github.com/pkg/sftp" + "golang.org/x/crypto/ssh" +) + +func TestHostTunnelLossClassification(t *testing.T) { + for _, tc := range []struct { + name string + established, expected bool + events []host.TunnelEvent + wantFailure bool + }{ + {"startup", false, false, []host.TunnelEvent{{CanReconnect: true}, {Up: true}}, false}, + {"unexpected", true, false, []host.TunnelEvent{{CanReconnect: true}, {Up: true}}, true}, + {"permanent loss", true, false, []host.TunnelEvent{{}}, true}, + {"planned recovery", true, true, []host.TunnelEvent{{CanReconnect: true}, {Up: true}}, false}, + {"loss after planned recovery", true, true, []host.TunnelEvent{{CanReconnect: true}, {Up: true}, {CanReconnect: true}}, true}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(context.Canceled) + r := &runState{ctx: ctx, cancel: cancel, rec: newRecorder(io.Discard, Config{})} + h := &hostRun{established: tc.established, expectedCut: tc.expected, proxy: &tunnelProxy{}, tunnel: make(chan tunnelObservation, 8)} + for _, e := range tc.events { + r.hostTunnel(h, e) + } + if failed := context.Cause(ctx) != nil; failed != tc.wantFailure { + t.Fatalf("host loss classification: failed=%v want=%v cause=%v", failed, tc.wantFailure, context.Cause(ctx)) + } + }) + } + t.Run("normal teardown", func(t *testing.T) { + ctx, cancel := context.WithCancelCause(t.Context()) + cancel(context.Canceled) + r := &runState{ctx: ctx, cancel: func(err error) { t.Errorf("teardown caused run failure: %v", err) }, rec: newRecorder(io.Discard, Config{})} + h := &hostRun{established: true, tunnel: make(chan tunnelObservation, 1)} + r.hostTunnel(h, host.TunnelEvent{}) + if !errors.Is(context.Cause(ctx), context.Canceled) { + t.Fatal("teardown changed cancellation cause") + } + }) +} + +func TestUnexpectedHostDownFailsDespiteImmediateRecovery(t *testing.T) { + c := localRelay(t) + c.Active = false + c.GuestsPerHost = 0 + r := admittedRun(t, c) + h := r.hosts[0] + h.h.TunnelCallback(host.TunnelEvent{CanReconnect: true}) + h.h.TunnelCallback(host.TunnelEvent{Up: true}) + if cause := context.Cause(r.ctx); cause == nil || !strings.Contains(cause.Error(), "unexpected host transport loss") { + t.Fatalf("established idle host loss escaped failure after recovery: %v", cause) + } +} + +func TestRealIdleHostTransportLossFailsBeforeNextCanary(t *testing.T) { + c := localRelay(t) + c.Active = false + c.GuestsPerHost = 0 + r := admittedRun(t, c) + r.canaries() + if n, _ := r.hosts[0].proxy.cut(); n != 1 { + t.Fatalf("test did not sever a real host transport: %d", n) + } + timer := time.NewTimer(3 * time.Second) + defer timer.Stop() + recovered := false + for { + select { + case event := <-r.hosts[0].tunnel: + recovered = recovered || event.event.Up + case <-r.ctx.Done(): + if cause := context.Cause(r.ctx); !strings.Contains(cause.Error(), "unexpected host transport loss") { + t.Fatalf("host loss failed for the wrong reason: %v", cause) + } + r.workers.Wait() + if r.rec.summary.CanaryEstablished != 1 { + t.Fatal("unexpected loss was only detected by another canary") + } + return + case <-timer.C: + t.Fatalf("real idle host loss did not fail before the next canary (recovered=%v)", recovered) + } + } +} + +func TestFaultHelpersDispatchCompleteRepeatedReadinessFrames(t *testing.T) { + for _, tc := range []struct{ name, script, input, output string }{ + {"silent", silentAfterReadinessScript, "RRQII", "AABI-"}, + {"late", lateFinalReplyScript, "RRQIIF", "AABIIB"}, + {"retry", retryReadinessScript, "RRRQIIF", "-AABIIB"}, + } { + for _, fragmented := range []bool{false, true} { + name := tc.name + "/whole" + if fragmented { + name = tc.name + "/fragmented" + } + t.Run(name, func(t *testing.T) { + var input, want []byte + for i, kind := range tc.input { + at := time.Unix(1700000000, int64(i)) + input = append(input, frame(string(kind), 0, i, at, 128)...) + if tc.output[i] != '-' { + want = append(want, frame(string(tc.output[i]), 0, i, at, 128)...) + } + } + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, "/bin/sh", "-c", tc.script) + var output bytes.Buffer + cmd.Stdout = &output + cmd.Stderr = &output + stdin, err := cmd.StdinPipe() + if err != nil { + t.Fatal(err) + } + if err = cmd.Start(); err != nil { + t.Fatal(err) + } + for len(input) > 0 { + n := len(input) + if fragmented { + n = min(n, 11) + } + if _, err = stdin.Write(input[:n]); err != nil { + t.Error(err) + break + } + input = input[n:] + if fragmented { + time.Sleep(time.Millisecond) + } + } + _ = stdin.Close() + if err = cmd.Wait(); err != nil { + t.Fatal(err) + } + if !bytes.Equal(output.Bytes(), want) { + t.Fatalf("complete repeated-R admission/fault responses mismatch: got %q, want %q", output.Bytes(), want) + } + }) + } + } +} + +func admittedRun(t *testing.T, c Config) *runState { + t.Helper() + ctx, cancel := context.WithCancelCause(t.Context()) + dir, err := os.MkdirTemp("", "capacity-review-") + if err != nil { + t.Fatal(err) + } + r := &runState{cfg: c, ctx: ctx, cancel: cancel, guard: newGuard(c), rec: newRecorder(io.Discard, c), helper: c.HelperExecutable, dir: dir} + t.Cleanup(func() { + if err := r.cleanup(); err != nil { + t.Error(err) + } + }) + if err := r.admit(nil); err != nil { + t.Fatal(err) + } + return r +} + +func TestRunGuestAuthorizationAcrossTransports(t *testing.T) { + var previousRunKey []byte + for _, ws := range []int{0, 100} { + name := "ssh" + if ws == 100 { + name = "ws" + } + t.Run(name, func(t *testing.T) { + c := localRelay(t) + c.Active = false + c.Hosts = 2 + c.GuestWSPercent = ws + r := admittedRun(t, c) + runKey := r.guests[0].identity.PublicKey().Marshal() + if bytes.Equal(runKey, previousRunKey) { + t.Fatal("separate runs reused a guest authorization key") + } + previousRunKey = runKey + // Admission has proved a usable authenticated shell on each host. + for _, q := range r.guests { + client, err := sftp.NewClient(q.client) + if err != nil { + t.Fatalf("authorized guest SFTP: %v", err) + } + path := filepath.Join(q.host.h.CommandDir, "authorized-transfer") + f, err := client.Create(path) + if err != nil { + t.Fatal(err) + } + payload := []byte("run-owned transfer\n") + if _, err = f.Write(payload); err != nil { + t.Fatal(err) + } + if err = f.Close(); err != nil { + t.Fatal(err) + } + f, err = client.Open(path) + if err != nil { + t.Fatal(err) + } + got, err := io.ReadAll(f) + _ = f.Close() + _ = client.Close() + if err != nil || !bytes.Equal(got, payload) { + t.Fatalf("authorized transfer mismatch: %q, %v", got, err) + } + } + q := r.guests[0] + s, err := q.host.session(r.ctx) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(r.ctx, 3*time.Second) + defer cancel() + foreign, raw, err := openGuest(ctx, r.guard, q.target, q.endpoint, s, signer(t), q.proxy) + if raw != nil { + defer raw.Close() + } + if err == nil { + defer foreign.Close() + files, sftpErr := sftp.NewClient(foreign) + if sftpErr == nil { + _ = files.Close() + t.Error("fresh foreign key opened SFTP on a benchmark host") + } + t.Fatal("fresh foreign key joined a benchmark session") + } + if ctx.Err() != nil { + t.Fatalf("foreign key was not rejected within the handshake budget: %v", err) + } + for _, q := range r.guests { + if !bytes.Equal(q.identity.PublicKey().Marshal(), r.guests[0].identity.PublicKey().Marshal()) { + t.Fatal("synthetic guests do not share the run-owned authorization key") + } + for _, other := range []ssh.Signer{q.host.h.Signers[0], q.host.h.HostKey} { + if bytes.Equal(q.identity.PublicKey().Marshal(), other.PublicKey().Marshal()) { + t.Fatal("guest identity reused a host registration or session key") + } + } + key, _, _, _, err := ssh.ParseAuthorizedKey([]byte(q.target.HostKeys[0])) + if err != nil || bytes.Equal(q.identity.PublicKey().Marshal(), key.Marshal()) { + t.Fatal("guest identity reused the relay key") + } + } + }) + } +} From 90415dd42700b7fa05c32910fbeae9395ab010d7 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 17:09:40 -0700 Subject: [PATCH 05/16] fix(capacity): allow balanced host guest churn --- internal/capacity/README.md | 4 +- internal/capacity/churn_test.go | 156 +++++++++++++++++++++++++++++++ internal/capacity/config.go | 2 +- internal/capacity/config_test.go | 45 +++++++++ 4 files changed, 204 insertions(+), 3 deletions(-) diff --git a/internal/capacity/README.md b/internal/capacity/README.md index ba49ffcf7..d64706271 100644 --- a/internal/capacity/README.md +++ b/internal/capacity/README.md @@ -224,8 +224,8 @@ requires the owned directory to be absent before return. JSONL contains actual upload/download bytes, digest and verification outcome. Transfers incomplete at phase expiry are a failure. -`churn_percent` can be zero, ten or fifty. It requires at least two hosts, two -guests per host and 25 steady seconds. Counts are floored; use enough hosts for +`churn_percent` can be zero, ten or fifty. It requires at least two hosts, one +guest per host and 25 steady seconds. Counts are floored; use enough hosts for the requested percentage to select a host. Selected host and guest transports are cut over a ten-second schedule starting one second into steady load. Selected guests are cut before selected hosts so an overlapping selection diff --git a/internal/capacity/churn_test.go b/internal/capacity/churn_test.go index 73087b23f..83546fd96 100644 --- a/internal/capacity/churn_test.go +++ b/internal/capacity/churn_test.go @@ -2,8 +2,11 @@ package capacity import ( "bytes" + "encoding/json" + "io" "os" "path/filepath" + "reflect" "testing" "time" ) @@ -59,3 +62,156 @@ func TestRealTransportChurnRestoresOriginalSessions(t *testing.T) { } } } + +// A 1:1 selection cuts matching pairs and leaves all other pairs untouched. +func TestBalancedChurnSchedulesPreserveSurvivors(t *testing.T) { + for _, tc := range []struct { + name string + hosts, percent int + want []churnAction + survivors []int + }{ + {"ten-percent", 10, 10, []churnAction{{index: 0}, {host: true, index: 0, offset: 10 * time.Second}}, []int{1, 2, 3, 4, 5, 6, 7, 8, 9}}, + {"fifty-percent", 2, 50, []churnAction{{index: 0}, {host: true, index: 0, offset: 10 * time.Second}}, []int{1}}, + {"rounded", 5, 50, []churnAction{{index: 0}, {index: 2, offset: 10 * time.Second / 3}, {host: true, index: 0, offset: 20 * time.Second / 3}, {host: true, index: 2, offset: 10 * time.Second}}, []int{1, 3, 4}}, + } { + t.Run(tc.name, func(t *testing.T) { + actions := churnSchedule(tc.hosts, 1, tc.percent) + if !reflect.DeepEqual(actions, tc.want) { + t.Fatalf("cut schedule = %+v; want %+v", actions, tc.want) + } + hosts, guests := map[int]bool{}, map[int]bool{} + for _, a := range actions { + if a.host { + hosts[a.index] = true + } else { + guests[a.index] = true + } + } + if !reflect.DeepEqual(hosts, guests) { + t.Fatalf("balanced affected sets differ: hosts=%v guests=%v", hosts, guests) + } + var survivors []int + for i := 0; i < tc.hosts; i++ { + if !hosts[i] && !guests[i] { + survivors = append(survivors, i) + } + } + if !reflect.DeepEqual(survivors, tc.survivors) { + t.Fatalf("untouched pairs = %v; want %v", survivors, tc.survivors) + } + }) + } +} + +func TestRealBalancedMixedChurnRestoresOriginalSessions(t *testing.T) { + c := localRelay(t) + c.Hosts, c.GuestsPerHost, c.ChurnPercent, c.DurationSeconds = 2, 1, 50, 25 + c.HostWSPercent, c.GuestWSPercent = 50, 50 + c.PayloadBudgetBytes = 4 << 20 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + // Preserve the actual run even if an assertion fails; no extra evidence run. + if dir := os.Getenv("CAPACITY_EVIDENCE_DIR"); dir != "" { + manifest, e := json.MarshalIndent(c.Manifest, "", " ") + if e != nil { + t.Fatal(e) + } + phase, e := json.MarshalIndent(c, "", " ") + if e != nil { + t.Fatal(e) + } + for name, data := range map[string][]byte{"manifest.json": manifest, "phase.json": phase, "events.jsonl": out.Bytes()} { + if e := os.WriteFile(filepath.Join(dir, "balanced-churn-"+name), data, 0600); e != nil { + t.Fatal(e) + } + } + } + if err != nil { + t.Fatalf("balanced churn: %v\n%s", err, out.String()) + } + if s.EstablishedHosts != 2 || s.EstablishedGuests != 2 || s.EstablishedChannels != 2 || s.AdmissionFailures != 0 { + t.Fatalf("balanced original admissions changed: %+v", s) + } + if s.ActualHostCuts != 1 || s.ActualGuestCuts != 1 || s.Reconnects < 2 || !s.ChurnCompleted { + t.Fatalf("balanced real cuts/restoration missing: %+v", s) + } + if !s.DeliveryComplete || !s.CleanupComplete || s.StopReason != "duration" || s.SteadySeconds < 25 { + t.Fatalf("balanced hold/delivery/cleanup incomplete: %+v", s) + } + if s.ActiveFrames == 0 || s.ActiveFrames+s.PausedActiveFrames != 50 || s.InputFrames != int64(s.RTTCount+s.CanaryRTTCount+s.InterruptedReplies) || s.ReceivedFrames <= int64(s.RTTCount) { + t.Fatalf("balanced input/reply/host-output accounting incomplete: %+v", s) + } + type event struct { + Event string `json:"event"` + Role string `json:"role"` + Host int `json:"host"` + Guest int `json:"guest"` + Session string `json:"session_id"` + Transport string `json:"transport"` + Generation int `json:"generation"` + SameIdentity bool `json:"same_identity"` + SameHostRun bool `json:"same_host_run"` + Replaced bool `json:"transport_replaced"` + Before int `json:"responses_before"` + After int `json:"responses_after"` + Actual int `json:"actual_transports"` + Planned time.Time `json:"planned_at"` + } + sessions, transports := map[int]string{}, map[int]string{} + ready := map[int]int{} + hostRestores, guestRestores, canaries := 0, 0, 0 + var cuts []event + decoder := json.NewDecoder(bytes.NewReader(out.Bytes())) + for { + var e event + if err := decoder.Decode(&e); err == io.EOF { + break + } else if err != nil { + t.Fatal(err) + } + switch e.Event { + case "host_ready": + if e.Host < 0 || e.Host > 1 || sessions[e.Host] != "" || e.Session == "" { + t.Fatalf("invalid original host: %+v", e) + } + sessions[e.Host], transports[e.Host] = e.Session, e.Transport + case "guest_ready": + if e.Guest < 0 || e.Guest > 1 || e.Host != e.Guest || e.Session != sessions[e.Host] { + t.Fatalf("guest moved sessions: %+v", e) + } + want := []string{"ssh", "ws"}[e.Guest] + if e.Transport != want || (e.Guest == 1 && (e.Generation != 0 || ready[1] != 0)) { + t.Fatalf("guest transport/identity changed: %+v", e) + } + ready[e.Guest]++ + case "transport_cut": + if e.Host != 0 || e.Actual != 1 || (e.Role == "guest" && e.Guest != 0) { + t.Fatalf("survivor or absent transport cut: %+v", e) + } + cuts = append(cuts, e) + case "host_restored": + if e.Host != 0 || !e.SameHostRun || e.Session != sessions[0] { + t.Fatalf("host/session replaced: %+v", e) + } + hostRestores++ + case "guest_restored": + if e.Guest != 0 || e.Host != 0 || !e.SameIdentity || e.Generation < 1 { + t.Fatalf("guest identity replaced: %+v", e) + } + guestRestores++ + case "survivor_canary": + if e.Host != 1 || e.Guest != 1 || e.Replaced || e.After <= e.Before { + t.Fatalf("original survivor stopped responding: %+v", e) + } + canaries++ + } + } + if transports[0] != "ws" || transports[1] != "ssh" || len(sessions) != 2 || ready[0] < 2 || ready[1] != 1 || hostRestores != 1 || guestRestores < 1 || canaries != 1 { + t.Fatalf("missing crossed transports, same sessions or untouched survivor: hosts=%v ready=%v restores=%d/%d canaries=%d", transports, ready, hostRestores, guestRestores, canaries) + } + if len(cuts) != 2 || cuts[0].Role != "guest" || cuts[1].Role != "host" || cuts[1].Planned.Sub(cuts[0].Planned) != 10*time.Second { + t.Fatalf("actual balanced cut schedule changed: %+v", cuts) + } + t.Logf("balanced 2 hosts/2 guests: one host and one guest cut, original identities/sessions, unchanged generation-0 survivor, 50 offered active slots, complete delivery/cleanup; summary=%+v", s) +} diff --git a/internal/capacity/config.go b/internal/capacity/config.go index 5ff0de664..347d739d5 100644 --- a/internal/capacity/config.go +++ b/internal/capacity/config.go @@ -160,7 +160,7 @@ func (c Config) Validate(now time.Time) error { if c.ChurnPercent != 0 && c.ChurnPercent != 10 && c.ChurnPercent != 50 { return errors.New("churn must be zero, ten or fifty percent") } - if c.ChurnPercent > 0 && (c.Hosts < 2 || c.GuestsPerHost < 2 || c.DurationSeconds < 25 || c.Hosts*c.ChurnPercent/100 < 1) { + if c.ChurnPercent > 0 && (c.Hosts < 2 || c.GuestsPerHost < 1 || c.DurationSeconds < 25 || c.Hosts*c.ChurnPercent/100 < 1) { return errors.New("churn needs survivor hosts/guests and at least 25 seconds") } fanout := int64(1 + c.GuestsPerHost*c.ChannelsPerGuest) diff --git a/internal/capacity/config_test.go b/internal/capacity/config_test.go index 910f7834e..309a19593 100644 --- a/internal/capacity/config_test.go +++ b/internal/capacity/config_test.go @@ -95,3 +95,48 @@ func TestDecodeRequiresBothExplicitInputs(t *testing.T) { } } } + +// Balanced reconnect phases must use the candidate's original 1:1 load. +func TestBalancedChurnValidation(t *testing.T) { + clearProxies(t) + for _, tc := range []struct { + name string + hosts, percent int + }{ + {"ten-percent", 10, 10}, {"fifty-percent", 2, 50}, + } { + t.Run(tc.name, func(t *testing.T) { + c := validConfig(t) + c.Hosts, c.GuestsPerHost, c.ChurnPercent = tc.hosts, 1, tc.percent + c.DurationSeconds = 25 + if err := c.Validate(time.Now()); err != nil { + t.Fatalf("balanced churn refused: %v", err) + } + }) + } + for _, tc := range []struct { + name string + change func(*Config) + }{ + {"zero guests", func(c *Config) { c.GuestsPerHost = 0 }}, + {"no survivor host", func(c *Config) { c.Hosts = 1 }}, + {"zero selected hosts", func(c *Config) { c.Hosts = 9; c.ChurnPercent = 10 }}, + {"short duration", func(c *Config) { c.DurationSeconds = 24 }}, + {"excess duration", func(c *Config) { c.DurationSeconds = 86401 }}, + {"manifest deadline", func(c *Config) { c.Manifest.ExpiresAt = time.Now().Add(24 * time.Second) }}, + {"zero rate", func(c *Config) { c.AdmissionsPerSecond = 0 }}, + {"excess rate", func(c *Config) { c.AdmissionsPerSecond = 21 }}, + {"external connections", func(c *Config) { c.Hosts = 500 }}, + {"insufficient payload", func(c *Config) { c.PayloadBudgetBytes = 1 }}, + {"excess payload", func(c *Config) { c.PayloadBudgetBytes = MaxPayload + 1 }}, + } { + t.Run(tc.name, func(t *testing.T) { + c := validConfig(t) + c.Hosts, c.GuestsPerHost, c.ChurnPercent, c.DurationSeconds = 2, 1, 50, 25 + tc.change(&c) + if err := c.Validate(time.Now()); err == nil { + t.Fatal("out-of-bounds balanced churn accepted") + } + }) + } +} From 944f592ebb71dae28069193fde3892081e05ff91 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 17:53:47 -0700 Subject: [PATCH 06/16] fix(capacity): bind forwarded WebSocket certificates to origin --- internal/capacity/client.go | 6 +- internal/capacity/forwarded_cert_test.go | 352 +++++++++++++++++++++++ internal/capacity/guard.go | 23 ++ 3 files changed, 378 insertions(+), 3 deletions(-) create mode 100644 internal/capacity/forwarded_cert_test.go diff --git a/internal/capacity/client.go b/internal/capacity/client.go index 292eeb468..d90b842d2 100644 --- a/internal/capacity/client.go +++ b/internal/capacity/client.go @@ -49,11 +49,11 @@ func guestHandshake(ctx context.Context, raw net.Conn, addr string, cfg *ssh.Cli return ssh.NewClient(conn, chans, reqs), nil } func openGuest(ctx context.Context, g *guard, t Target, u *url.URL, s *api.GetSessionResponse, identity ssh.Signer, proxy *tunnelProxy) (*ssh.Client, net.Conn, error) { - if err := g.session(ctx, s); err != nil { + check, err := g.guestKey(ctx, t, u.Scheme, s) + if err != nil { return nil, nil, err } var raw net.Conn - var err error if u.Scheme == "ssh" { raw, err = g.connectPermit(ctx, t, u.Host, proxy.takePermit()) } else { @@ -64,7 +64,7 @@ func openGuest(ctx context.Context, g *guard, t Target, u *url.URL, s *api.GetSe if err != nil { return nil, nil, err } - cfg := &ssh.ClientConfig{User: s.SshUser, Auth: []ssh.AuthMethod{ssh.PublicKeys(identity)}, ClientVersion: upterm.ClientSSHClientVersion, HostKeyCallback: pinnedKey(t)} + cfg := &ssh.ClientConfig{User: s.SshUser, Auth: []ssh.AuthMethod{ssh.PublicKeys(identity)}, ClientVersion: upterm.ClientSSHClientVersion, HostKeyCallback: check} client, err := guestHandshake(ctx, raw, u.Host, cfg) return client, raw, err } diff --git a/internal/capacity/forwarded_cert_test.go b/internal/capacity/forwarded_cert_test.go new file mode 100644 index 000000000..71ae52800 --- /dev/null +++ b/internal/capacity/forwarded_cert_test.go @@ -0,0 +1,352 @@ +package capacity + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "maps" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/go-kit/kit/metrics/provider" + consul "github.com/hashicorp/consul/api" + "github.com/owenthereal/upterm/host/api" + "github.com/owenthereal/upterm/routing" + "github.com/owenthereal/upterm/server" + "golang.org/x/crypto/ssh" +) + +// Trusting the entry/all manifest keys for a forwarded certificate, or using +// the origin principal for an SSH/embedded/direct guest, breaks these checks. +func TestGuestKeyBindsExactRouteAndPreservesEntryChecks(t *testing.T) { + c := validConfig(t) + originKey, entryKey := signer(t), signer(t) + origin := c.Manifest.Targets[0] + origin.HostKeys = []string{string(ssh.MarshalAuthorizedKey(originKey.PublicKey()))} + entry := origin + entry.MachineID = "entry" + entry.NodeAddresses, entry.AllowedIPs = []string{"[::1]:2222"}, []string{"::1"} + entry.HostKeys = []string{string(ssh.MarshalAuthorizedKey(entryKey.PublicKey()))} + c.Manifest.Targets = []Target{origin, entry} + g := newGuard(c) + s := &api.GetSessionResponse{SessionId: "session", SshUser: "session", NodeAddr: "127.0.0.1:2222"} + newCert := func(key ssh.Signer, principal string, change func(*ssh.Certificate)) *ssh.Certificate { + cert := &ssh.Certificate{Key: key.PublicKey(), CertType: ssh.HostCert, ValidPrincipals: []string{principal}, ValidBefore: ssh.CertTimeInfinity} + if change != nil { + change(cert) + } + if err := cert.SignCert(rand.Reader, key); err != nil { + t.Fatal(err) + } + return cert + } + check, err := g.guestKey(t.Context(), entry, "ws", s) + if err != nil { + t.Fatal(err) + } + valid := newCert(originKey, "127.0.0.1", nil) + if err := check("[::1]:8080", nil, valid); err != nil { + t.Fatalf("exact origin certificate refused: %v", err) + } + if err := check("[::1]:8080", nil, originKey.PublicKey()); err != nil { + t.Fatalf("exact origin raw pin refused: %v", err) + } + corrupt := newCert(originKey, "127.0.0.1", nil) + corrupt.Signature.Blob[0] ^= 1 + for name, key := range map[string]ssh.PublicKey{ + "wrong origin principal": newCert(originKey, "::1", nil), + "another manifest authority": newCert(entryKey, "127.0.0.1", nil), + "foreign authority": newCert(signer(t), "127.0.0.1", nil), + "another manifest raw key": entryKey.PublicKey(), + "user certificate": newCert(originKey, "127.0.0.1", func(c *ssh.Certificate) { c.CertType = ssh.UserCert }), + "expired": newCert(originKey, "127.0.0.1", func(c *ssh.Certificate) { c.ValidBefore = uint64(time.Now().Add(-time.Hour).Unix()) }), + "future": newCert(originKey, "127.0.0.1", func(c *ssh.Certificate) { c.ValidAfter = uint64(time.Now().Add(time.Hour).Unix()) }), + "invalid signature": corrupt, + } { + t.Run(name, func(t *testing.T) { + if err := check("[::1]:8080", nil, key); err == nil { + t.Fatal("invalid forwarded certificate/key accepted") + } + }) + } + for _, tc := range []struct { + name, transport, addr, principal string + entry Target + route string + key ssh.Signer + }{ + {"ssh entry", "ssh", "[::1]:2222", "::1", entry, "session", entryKey}, + {"embedded ws entry", "ws", "[::1]:8080", "::1", entry, (&routing.EmbeddedEncodeDecoder{}).Encode("session", s.NodeAddr), entryKey}, + {"direct private", "ws", "127.0.0.1:8080", "127.0.0.1", origin, "session", originKey}, + {"direct public", "wss", "isolated.example:443", "isolated.example", origin, "session", originKey}, + } { + t.Run(tc.name, func(t *testing.T) { + session := &api.GetSessionResponse{SessionId: "session", SshUser: tc.route, NodeAddr: "127.0.0.1:2222"} + check, err := g.guestKey(t.Context(), tc.entry, tc.transport, session) + if err != nil { + t.Fatal(err) + } + if err := check(tc.addr, nil, newCert(tc.key, tc.principal, nil)); err != nil { + t.Fatalf("entry endpoint certificate refused: %v", err) + } + if err := check(tc.addr, nil, newCert(tc.key, "wrong.example", nil)); err == nil { + t.Fatal("entry endpoint principal check bypassed") + } + }) + } +} + +func TestGuestInvalidRouteFailsBeforeAdmission(t *testing.T) { + c := validConfig(t) + g := newGuard(c) + calls := 0 + g.dial = func(context.Context, string, string) (net.Conn, error) { + calls++ + return nil, errors.New("unexpected dial") + } + u, err := url.Parse(c.Manifest.Targets[0].SSHURL) + if err != nil { + t.Fatal(err) + } + for _, session := range []*api.GetSessionResponse{ + nil, + {SessionId: "session", SshUser: "session", NodeAddr: "[::1]:2222"}, + {SessionId: "session", SshUser: "different", NodeAddr: "127.0.0.1:2222"}, + {SessionId: "session", SshUser: (&routing.EmbeddedEncodeDecoder{}).Encode("session", "[::1]:2222"), NodeAddr: "127.0.0.1:2222"}, + } { + if _, _, err := openGuest(t.Context(), g, c.Manifest.Targets[0], u, session, signer(t), &tunnelProxy{}); err == nil { + t.Fatal("unvalidated guest route accepted") + } + } + if calls != 0 || g.peak != 0 || !g.next.IsZero() { + t.Fatalf("invalid route entered admission: dials=%d peak=%d next=%v", calls, g.peak, g.next) + } +} + +// Distinct loopback principals reproduce a real forwarded certificate: the +// guest dials B, while routing presents A's certificate to the guest. +func TestRealForwardedGuestCertificate(t *testing.T) { + for _, mode := range []routing.Mode{routing.ModeEmbedded, routing.ModeConsul} { + t.Run(string(mode), func(t *testing.T) { + var registry *url.URL + if mode == routing.ModeConsul { + registry = forwardedSessionRegistry(t) + } + for _, tc := range []struct { + name string + h, g int + }{{"ssh", 0, 0}, {"ws", 100, 100}, {"ssh-to-ws", 0, 100}, {"ws-to-ssh", 100, 0}} { + t.Run(tc.name, func(t *testing.T) { + clearProxies(t) + key := signer(t) + origin := forwardedRelay(t, "origin", "127.0.0.1", key, mode, registry) + entry := forwardedRelay(t, "entry", "::1", key, mode, registry) + c := validConfig(t) + c.Manifest.Targets = []Target{origin, entry} + c.HostTarget, c.EntryTarget = "origin", "entry" + c.HelperExecutable = buildHelper(t) + c.HostWSPercent, c.GuestWSPercent = tc.h, tc.g + c.Active, c.DurationSeconds, c.ChannelsPerGuest = true, 1, 4 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if dir := os.Getenv("CAPACITY_EVIDENCE_DIR"); dir != "" { + manifest, e := json.MarshalIndent(c.Manifest, "", " ") + if e != nil { + t.Fatal(e) + } + phase, e := json.MarshalIndent(c, "", " ") + if e != nil { + t.Fatal(e) + } + for name, data := range map[string][]byte{"manifest.json": manifest, "phase.json": phase, "events.jsonl": out.Bytes()} { + if e := os.WriteFile(filepath.Join(dir, "forwarded-"+string(mode)+"-"+tc.name+"-"+name), data, 0600); e != nil { + t.Fatal(e) + } + } + } + if err != nil { + t.Fatalf("forwarded shell workload: %v\n%s", err, out.String()) + } + if s.EstablishedHosts != 1 || s.EstablishedGuests != 1 || s.EstablishedChannels != 4 || s.AdmissionFailures != 0 { + t.Fatalf("forwarded admissions incomplete: %+v", s) + } + if !s.DeliveryComplete || !s.CleanupComplete || s.StopReason != "duration" || s.SteadySeconds < 1 || s.ActiveFrames != 1 || s.InputFrames != 5 || s.RTTCount != 5 || s.PayloadBytes == 0 { + t.Fatalf("forwarded framed delivery/cleanup incomplete: %+v", s) + } + nodes := map[string]int{} + dec := json.NewDecoder(bytes.NewReader(out.Bytes())) + for { + var event struct { + Event string `json:"event"` + Node string `json:"node"` + } + if err := dec.Decode(&event); err == io.EOF { + break + } else if err != nil { + t.Fatal(err) + } + if event.Event == "guest_ready" { + nodes[event.Node]++ + } + } + if !maps.Equal(nodes, map[string]int{origin.NodeAddresses[0]: 1}) { + t.Fatalf("guest readiness named wrong origin: %v", nodes) + } + }) + } + }) + } +} + +func forwardedRelay(t *testing.T, id, ip string, key ssh.Signer, mode routing.Mode, registry *url.URL) Target { + t.Helper() + sshln, err := net.Listen("tcp", net.JoinHostPort(ip, "0")) + if err != nil { + t.Fatalf("required distinct-principal loopback %s unavailable: %v", ip, err) + } + t.Cleanup(func() { _ = sshln.Close() }) + wsln, err := net.Listen("tcp", net.JoinHostPort(ip, "0")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = wsln.Close() }) + cert, err := (&server.HostCertSigner{Hostnames: []string{ip}}).SignCert(key) + if err != nil { + t.Fatal(err) + } + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + network := &server.MemoryProvider{} + if err := network.SetOpts(nil); err != nil { + t.Fatal(err) + } + sm, err := server.NewSessionManager(mode, server.WithSessionManagerLogger(logger), server.WithSessionManagerConsulURL(registry)) + if err != nil { + t.Fatal(err) + } + relay := &server.Server{NodeAddr: sshln.Addr().String(), HostSigners: []ssh.Signer{cert}, Signers: []ssh.Signer{key}, NetworkProvider: network, SessionManager: sm, MetricsProvider: provider.NewDiscardProvider(), Logger: logger} + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- relay.ServeListeners(ctx, server.Listeners{SSH: sshln, WS: wsln}) }() + t.Cleanup(func() { + cancel() + _ = relay.Shutdown() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Error("forwarded relay leaked") + } + }) + metrics := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, "test_metric 1\n") })) + _ = metrics.Listener.Close() + metrics.Listener, err = net.Listen("tcp", net.JoinHostPort(ip, "0")) + if err != nil { + t.Fatal(err) + } + metrics.Start() + t.Cleanup(metrics.Close) + return Target{AppID: "test-app", MachineID: id, SSHURL: "ssh://" + sshln.Addr().String(), WSURL: "ws://" + wsln.Addr().String(), NodeAddresses: []string{sshln.Addr().String()}, AllowedIPs: []string{ip}, HostKeys: []string{string(ssh.MarshalAuthorizedKey(key.PublicKey()))}, TelemetryURL: metrics.URL + "/metrics"} +} + +// The registry supplies only the local storage protocol needed for a single +// non-churning session. Certificate exchange, route lookup, shell framing and +// cleanup all run through the real relay and host implementations. +func forwardedSessionRegistry(t *testing.T) *url.URL { + t.Helper() + var mu sync.Mutex + index, leases := uint64(1), 0 + pairs := map[string]*consul.KVPair{} + changed := make(chan struct{}) + poke := func() { index++; close(changed); changed = make(chan struct{}) } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + defer mu.Unlock() + w.Header().Set("Content-Type", "application/json") + switch { + case r.URL.Path == "/v1/catalog/register": + _, _ = io.WriteString(w, "true") + case r.URL.Path == "/v1/session/create": + leases++ + _ = json.NewEncoder(w).Encode(map[string]string{"ID": fmt.Sprintf("lease-%d", leases)}) + case strings.HasPrefix(r.URL.Path, "/v1/session/destroy/"): + lease := strings.TrimPrefix(r.URL.Path, "/v1/session/destroy/") + for key, pair := range pairs { + if pair.Session == lease { + delete(pairs, key) + poke() + } + } + _, _ = io.WriteString(w, "true") + case r.URL.Path == "/v1/txn": + var ops consul.TxnOps + if err := json.NewDecoder(r.Body).Decode(&ops); err != nil { + t.Error(err) + w.WriteHeader(http.StatusBadRequest) + return + } + var results consul.TxnResults + for _, op := range ops { + kv := op.KV + switch kv.Verb { + case consul.KVCheckNotExists, consul.KVCheckIndex, consul.KVCheckSession: + case consul.KVLock: + poke() + pair := &consul.KVPair{Key: kv.Key, Value: kv.Value, Session: kv.Session, ModifyIndex: index} + pairs[kv.Key] = pair + results = append(results, &consul.TxnResult{KV: pair}) + case consul.KVDelete: + delete(pairs, kv.Key) + poke() + default: + t.Errorf("unexpected registry transaction: %s", kv.Verb) + } + } + _ = json.NewEncoder(w).Encode(consul.TxnResponse{Results: results}) + case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/v1/kv/"): + if wait, _ := strconv.ParseUint(r.URL.Query().Get("index"), 10, 64); wait == index { + wake := changed + mu.Unlock() + select { + case <-wake: + case <-r.Context().Done(): + } + mu.Lock() + } + w.Header().Set("X-Consul-Index", strconv.FormatUint(index, 10)) + key := strings.TrimPrefix(r.URL.Path, "/v1/kv/") + var found consul.KVPairs + for k, pair := range pairs { + if k == key || (r.URL.Query().Has("recurse") && strings.HasPrefix(k, key)) { + found = append(found, pair) + } + } + if len(found) == 0 { + w.WriteHeader(http.StatusNotFound) + return + } + _ = json.NewEncoder(w).Encode(found) + default: + t.Errorf("unexpected registry request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusInternalServerError) + } + })) + t.Cleanup(srv.Close) + u, err := url.Parse(srv.URL) + if err != nil { + t.Fatal(err) + } + return u +} diff --git a/internal/capacity/guard.go b/internal/capacity/guard.go index d3edff2e8..b340c6b24 100644 --- a/internal/capacity/guard.go +++ b/internal/capacity/guard.go @@ -178,6 +178,29 @@ func (g *guard) session(ctx context.Context, s *api.GetSessionResponse) error { } return errors.New("session node outside manifest") } + +func (g *guard) guestKey(ctx context.Context, entry Target, transport string, s *api.GetSessionResponse) (ssh.HostKeyCallback, error) { + if err := g.session(ctx, s); err != nil { + return nil, err + } + // SSH terminates at the entry. Embedded-mode WebSocket routing also goes + // through the entry's SSH proxy. Consul-mode WebSocket routing instead + // pipes SSH to the session's node, whose certificate the guest receives. + // Its plain SSH user distinguishes it from the validated embedded route. + if transport != "ssh" && s.SshUser == s.SessionId { + for _, origin := range g.cfg.Manifest.Targets { + if origin.MachineID != entry.MachineID && contains(origin.NodeAddresses, s.NodeAddr) { + check := pinnedKey(origin) + node := s.NodeAddr + return func(_ string, remote net.Addr, key ssh.PublicKey) error { + return check(node, remote, key) + }, nil + } + } + } + return pinnedKey(entry), nil +} + func pinnedKey(t Target) ssh.HostKeyCallback { var pins []ssh.PublicKey for _, s := range t.HostKeys { From 9c96079892b97a15f83dfc0438d4b240d015ace6 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 19:45:57 -0700 Subject: [PATCH 07/16] fix(capacity): validate payloads and release closed tunnels --- internal/capacity/admission_test.go | 7 +- internal/capacity/client_test.go | 2 +- internal/capacity/cmd/main.go | 10 +- internal/capacity/config.go | 2 +- internal/capacity/fault_helper_test.go | 156 +++++++++++++++++ internal/capacity/forwarded_cert_test.go | 3 +- internal/capacity/guard.go | 15 +- internal/capacity/review_regression_test.go | 81 ++------- internal/capacity/round1_regression_test.go | 179 ++++++++++++++++++++ internal/capacity/runner.go | 12 +- internal/capacity/runner_test.go | 15 +- internal/capacity/transfer.go | 4 +- internal/capacity/whole_review_test.go | 14 +- internal/capacity/workload.go | 16 +- 14 files changed, 418 insertions(+), 98 deletions(-) create mode 100644 internal/capacity/fault_helper_test.go create mode 100644 internal/capacity/round1_regression_test.go diff --git a/internal/capacity/admission_test.go b/internal/capacity/admission_test.go index 841050fac..a47ddf1d7 100644 --- a/internal/capacity/admission_test.go +++ b/internal/capacity/admission_test.go @@ -14,7 +14,12 @@ func TestAdmissionWorkersOverlapReadinessWithinBound(t *testing.T) { release := make(chan struct{}) finished := make(chan error, 1) go func() { finished <- r.admissions(40, func(int) error { started <- struct{}{}; <-release; return nil }) }() - defer func() { close(release); <-finished }() + defer func() { + close(release) + if err := <-finished; err != nil { + t.Error(err) + } + }() for i := 0; i < 20; i++ { select { case <-started: diff --git a/internal/capacity/client_test.go b/internal/capacity/client_test.go index 9a01b6677..b46e691ea 100644 --- a/internal/capacity/client_test.go +++ b/internal/capacity/client_test.go @@ -11,7 +11,7 @@ import ( func TestSilentHandshakeCancellationClosesTransport(t *testing.T) { client, server := net.Pipe() - defer server.Close() + defer func() { _ = server.Close() }() ctx, cancel := context.WithCancel(t.Context()) finished := make(chan error, 1) go func() { diff --git a/internal/capacity/cmd/main.go b/internal/capacity/cmd/main.go index 3bfa678f0..d9b592be9 100644 --- a/internal/capacity/cmd/main.go +++ b/internal/capacity/cmd/main.go @@ -14,7 +14,7 @@ import ( func main() { if err := run(); err != nil { - fmt.Fprintln(os.Stderr, err) + _, _ = fmt.Fprintln(os.Stderr, err) os.Exit(1) } } @@ -46,19 +46,19 @@ func run() error { if err != nil { return err } - defer m.Close() + defer func() { _ = m.Close() }() p, err := os.Open(*phase) if err != nil { return err } - defer p.Close() + defer func() { _ = p.Close() }() cfg, err := capacity.Decode(m, p) if err != nil { return err } if *validate { - fmt.Fprintln(os.Stdout, `{"event":"validated","dialed":false}`) - return nil + _, err = fmt.Fprintln(os.Stdout, `{"event":"validated","dialed":false}`) + return err } _, err = capacity.Run(ctx, cfg, os.Stdout) return err diff --git a/internal/capacity/config.go b/internal/capacity/config.go index 347d739d5..d238b8ac4 100644 --- a/internal/capacity/config.go +++ b/internal/capacity/config.go @@ -107,7 +107,7 @@ func endpoint(raw string, schemes ...string) (*url.URL, error) { if !ok || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || u.Opaque != "" || production(u.Hostname()) { return nil, errors.New("unauthorized endpoint") } - if u.Path != "" && u.Path != "/" && !(u.Scheme == "http" && u.Path == "/metrics") { + if u.Path != "" && u.Path != "/" && (u.Scheme != "http" || u.Path != "/metrics") { return nil, errors.New("unexpected endpoint path") } if u.Port() == "" { diff --git a/internal/capacity/fault_helper_test.go b/internal/capacity/fault_helper_test.go new file mode 100644 index 000000000..7d5b1a6b4 --- /dev/null +++ b/internal/capacity/fault_helper_test.go @@ -0,0 +1,156 @@ +package capacity + +import ( + "bufio" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "golang.org/x/term" +) + +// Fault behavior belongs to the test executable, never the deployed helper. +// Each copy gets a sidecar so Host can invoke its normal helper arguments on +// every platform without a shell, inherited environment, or production mode. +type faultHelperConfig struct { + Mode string + Replay string +} + +func faultHelperExecutable(t *testing.T, mode, replay string) string { + t.Helper() + executable, err := os.Executable() + if err != nil { + t.Fatal(err) + } + source, err := os.Open(executable) + if err != nil { + t.Fatal(err) + } + defer func() { _ = source.Close() }() + name := "fault-helper" + if runtime.GOOS == "windows" { + name += ".exe" + } + path := filepath.Join(t.TempDir(), name) + destination, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0700) + if err != nil { + t.Fatal(err) + } + _, copyErr := io.Copy(destination, source) + closeErr := destination.Close() + if copyErr != nil || closeErr != nil { + t.Fatalf("copy fault helper: copy=%v close=%v", copyErr, closeErr) + } + config, err := json.Marshal(faultHelperConfig{Mode: mode, Replay: replay}) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path+".json", config, 0600); err != nil { + t.Fatal(err) + } + return path +} + +func faultHelper(t *testing.T, mode, replay string) Config { + t.Helper() + c := localRelay(t) + c.HelperExecutable = faultHelperExecutable(t, mode, replay) + return c +} + +func runFaultHelper(config faultHelperConfig, in *os.File, out io.Writer) error { + if term.IsTerminal(int(in.Fd())) { + state, err := term.MakeRaw(int(in.Fd())) + if err != nil { + return err + } + defer func() { _ = term.Restore(int(in.Fd()), state) }() + } + if _, err := io.WriteString(out, config.Replay); err != nil { + return err + } + inputs, challenges := 0, 0 + scanner := bufio.NewScanner(in) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "CAP ") || len(line) != 127 { + continue + } + b := []byte(line + "\n") + switch b[4] { + case 'R': + challenges++ + if config.Mode == "retry" && challenges == 1 { + continue + } + b[4] = 'A' + case 'Q': + b[4] = 'B' + case 'F': + if config.Mode == "silent" { + continue + } + if config.Mode != "replay" { + b[4] = 'B' + } + if config.Mode == "missing-host-suffix" { + copy(b[15:23], "00000001") + } + case 'I': + inputs++ + switch config.Mode { + case "silent": + if inputs > 1 { + continue + } + case "late": + if inputs == 2 { + time.Sleep(1200 * time.Millisecond) + } + if inputs > 2 { + continue + } + case "corrupt": + if inputs > 1 { + b[126] = 'y' + } + } + default: + if config.Mode != "missing-host-suffix" && config.Mode != "replay" { + continue + } + } + if _, err := out.Write(b); err != nil { + return err + } + } + return scanner.Err() +} + +func runFaultHelperProcess() int { + configPath, err := os.Executable() + if err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + return 1 + } + data, err := os.ReadFile(configPath + ".json") + var config faultHelperConfig + if err == nil { + err = json.Unmarshal(data, &config) + } + if err == nil { + err = runFaultHelper(config, os.Stdin, os.Stdout) + } + if err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + return 1 + } + return 0 +} diff --git a/internal/capacity/forwarded_cert_test.go b/internal/capacity/forwarded_cert_test.go index 71ae52800..17d7f57b0 100644 --- a/internal/capacity/forwarded_cert_test.go +++ b/internal/capacity/forwarded_cert_test.go @@ -245,7 +245,8 @@ func forwardedRelay(t *testing.T, id, ip string, key ssh.Signer, mode routing.Mo cancel() _ = relay.Shutdown() select { - case <-done: + case err := <-done: + _ = err // Shutdown owns the expected serving error. case <-time.After(10 * time.Second): t.Error("forwarded relay leaked") } diff --git a/internal/capacity/guard.go b/internal/capacity/guard.go index b340c6b24..6a41e9fd3 100644 --- a/internal/capacity/guard.go +++ b/internal/capacity/guard.go @@ -346,7 +346,7 @@ func (p *tunnelProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { p.mu.Unlock() defer p.workers.Done() if r.Method != http.MethodConnect { - http.Error(w, "CONNECT required", 405) + http.Error(w, "CONNECT required", http.StatusMethodNotAllowed) return } p.mu.Lock() @@ -369,7 +369,7 @@ func (p *tunnelProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { if p.ctx.Err() == nil && p.failure != nil { p.failure(errors.New("generator_admission_backlog")) } - http.Error(w, "generator admission backlog", 503) + http.Error(w, "generator admission backlog", http.StatusServiceUnavailable) return } } @@ -384,7 +384,7 @@ func (p *tunnelProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { } upstream, err := p.g.connectPermit(p.ctx, p.t, r.Host, permit) if err != nil { - http.Error(w, "destination refused", 403) + http.Error(w, "destination refused", http.StatusForbidden) return } p.mu.Lock() @@ -415,7 +415,14 @@ func (p *tunnelProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } copied := make(chan struct{}) - go func() { defer close(copied); _, _ = io.Copy(upstream, rw) }() + go func() { + defer close(copied) + _, _ = io.Copy(upstream, rw) + // Either peer can disappear while the other remains silent. Release + // both sockets to unblock the opposite copier before joining it. + _ = upstream.Close() + _ = downstream.Close() + }() _, _ = io.Copy(downstream, upstream) _ = downstream.Close() _ = upstream.Close() diff --git a/internal/capacity/review_regression_test.go b/internal/capacity/review_regression_test.go index 2452e9956..66fb941e9 100644 --- a/internal/capacity/review_regression_test.go +++ b/internal/capacity/review_regression_test.go @@ -16,51 +16,8 @@ import ( "time" ) -const silentAfterReadinessScript = `inputs=0 -while IFS= read -r line; do -case "$line" in -'CAP R '*) printf 'CAP A %s\n' "${line#CAP R }";; -'CAP Q '*) printf 'CAP B %s\n' "${line#CAP Q }";; -'CAP I '*) - if [ "$inputs" -eq 0 ]; then printf '%s\n' "$line"; fi - inputs=$((inputs + 1));; -esac -done -` -const lateFinalReplyScript = `inputs=0 -while IFS= read -r line; do -case "$line" in -'CAP R '*) printf 'CAP A %s\n' "${line#CAP R }";; -'CAP Q '*) printf 'CAP B %s\n' "${line#CAP Q }";; -'CAP I '*) - if [ "$inputs" -eq 1 ]; then sleep 1.2; fi - if [ "$inputs" -le 1 ]; then printf '%s\n' "$line"; fi - inputs=$((inputs + 1));; -'CAP F '*) printf 'CAP B %s\n' "${line#CAP F }";; -esac -done -` -const retryReadinessScript = `challenges=0 -while IFS= read -r line; do -case "$line" in -'CAP R '*) - if [ "$challenges" -gt 0 ]; then printf 'CAP A %s\n' "${line#CAP R }"; fi - challenges=$((challenges + 1));; -'CAP Q '*) printf 'CAP B %s\n' "${line#CAP Q }";; -'CAP F '*) printf 'CAP B %s\n' "${line#CAP F }";; -'CAP I '*) printf '%s\n' "$line";; -esac -done -` - func TestSilentAfterReadinessFailsDelivery(t *testing.T) { - c := localRelay(t) - helper := filepath.Join(t.TempDir(), "silent-helper") - script := "#!/bin/sh\nstty raw -echo\n" + silentAfterReadinessScript - if err := os.WriteFile(helper, []byte(script), 0700); err != nil { - t.Fatal(err) - } - c.HelperExecutable = helper + c := faultHelper(t, "silent", "") c.DurationSeconds = 3 s, err := Run(t.Context(), c, io.Discard) t.Logf("delivery result: err=%v active=%d rtt=%d stop=%s", err, s.ActiveFrames, s.RTTCount, s.StopReason) @@ -119,14 +76,19 @@ func TestGuestTransportCutDoesNotWaitForHostRepair(t *testing.T) { ctx, cancel := context.WithCancelCause(t.Context()) r := &runState{ctx: ctx, cancel: cancel, cfg: c, rec: newRecorder(io.Discard, c)} raw, peer := net.Pipe() - defer peer.Close() + defer func() { _ = peer.Close() }() h := &hostRun{} q := &guest{host: h, raw: raw, transportLive: true} r.guests = []*guest{q} h.repair.Lock() finished := make(chan error, 1) go func() { finished <- r.cutGuest(q) }() - defer func() { cancel(context.Canceled); h.repair.Unlock(); <-finished }() + defer func() { + cancel(context.Canceled) + h.repair.Unlock() + err := <-finished + _ = err // Cancellation intentionally stops the queued repair. + }() _ = peer.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) if _, err := peer.Read(make([]byte, 1)); err == nil { t.Fatal("transport was not cut") @@ -202,19 +164,8 @@ func pacedAdmissionFixture(t *testing.T, ws int) { } -func scriptedHelper(t *testing.T, body string) Config { - t.Helper() - c := localRelay(t) - path := filepath.Join(t.TempDir(), "helper") - if err := os.WriteFile(path, []byte("#!/bin/sh\nstty raw -echo\n"+body), 0700); err != nil { - t.Fatal(err) - } - c.HelperExecutable = path - return c -} - func TestLateFinalReplyIsDrainedAndMeasured(t *testing.T) { - c := scriptedHelper(t, lateFinalReplyScript) + c := faultHelper(t, "late", "") c.DurationSeconds = 1 s, err := Run(t.Context(), c, io.Discard) if err != nil { @@ -232,7 +183,7 @@ func TestLateFinalReplyIsDrainedAndMeasured(t *testing.T) { func TestMissingHostOutputSuffixFailsDelivery(t *testing.T) { // The barrier advertises one emitted H frame, but that entire suffix is // dropped. Every I frame is echoed, so own-response and gap-only checks pass. - c := scriptedHelper(t, "while IFS= read -r line; do\ncase \"$line\" in\n'CAP R '*) printf '%s\\n' \"$line\" | tr R A;;\n'CAP Q '*) printf '%s\\n' \"$line\" | tr Q B;;\n'CAP F '*) printf '%s\\n' \"$line\" | awk '{$2=\"B\";$4=\"00000001\";print}';;\n*) printf '%s\\n' \"$line\";;\nesac\ndone\n") + c := faultHelper(t, "missing-host-suffix", "") c.DurationSeconds = 1 s, err := Run(t.Context(), c, io.Discard) if err == nil { @@ -261,14 +212,14 @@ func TestActualCutsKeepScheduleDuringSlowOverlappingRepairs(t *testing.T) { r := &runState{cfg: c, ctx: ctx, cancel: cancel, rec: newRecorder(io.Discard, c)} h := &hostRun{id: 0, proxy: &tunnelProxy{conns: map[net.Conn]bool{}}} hostRaw, hostPeer := net.Pipe() - defer hostPeer.Close() + defer func() { _ = hostPeer.Close() }() h.proxy.conns[hostRaw] = true r.hosts = []*hostRun{h} var peers []net.Conn for i := 0; i < 2; i++ { raw, peer := net.Pipe() peers = append(peers, peer) - defer peer.Close() + defer func() { _ = peer.Close() }() r.guests = append(r.guests, &guest{id: i, host: h, raw: raw, transportLive: true}) } actions := []churnAction{{index: 0}, {index: 1, offset: 25 * time.Millisecond}, {host: true, index: 0, offset: 50 * time.Millisecond}} @@ -444,8 +395,8 @@ func TestRealDeliveryEvidence(t *testing.T) { func TestHelperFinalBarrierFlushesAndQuiescesHostOutput(t *testing.T) { in, input := io.Pipe() output, out := io.Pipe() - defer input.Close() - defer output.Close() + defer func() { _ = input.Close() }() + defer func() { _ = output.Close() }() finished := make(chan error, 1) go func() { err := helper(t.Context(), in, out, 0, true, 20*time.Millisecond) @@ -529,7 +480,7 @@ func TestReconnectCannotReuseReplayedBaseline(t *testing.T) { // The old transport's baseline can be replayed with an incomplete screen // history. None of it belongs to this connection's delivery interval. replay := string(frame("B", 0, 0, old, 128)) + string(frame("H", 0, 0, old, 1024)) + string(frame("H", 0, 2, old, 1024)) - c := scriptedHelper(t, "printf '%s' '"+replay+"'\nwhile IFS= read -r line; do\ncase \"$line\" in\n'CAP R '*) printf '%s\\n' \"$line\" | tr R A;;\n'CAP Q '*) printf '%s\\n' \"$line\" | tr Q B;;\n*) printf '%s\\n' \"$line\";;\nesac\ndone\n") + c := faultHelper(t, "replay", replay) ctx, cancel := context.WithCancelCause(t.Context()) dir, err := os.MkdirTemp("", "cap-replay-") if err != nil { @@ -600,7 +551,7 @@ func TestHelperResynchronizesAfterCutTruncatesInput(t *testing.T) { func TestReadinessChallengeCanRetryWithinSameServiceBudget(t *testing.T) { // A reconstruction can consume/transform the first challenge response. // A repeated authenticated challenge must prove the same live channel. - c := scriptedHelper(t, retryReadinessScript) + c := faultHelper(t, "retry", "") c.DurationSeconds = 1 s, err := Run(t.Context(), c, io.Discard) if err != nil || !s.DeliveryComplete || s.EstablishedGuests != 1 || s.InputFrames != 2 || s.RTTCount != 2 || s.HandshakeFrames < 4 { diff --git a/internal/capacity/round1_regression_test.go b/internal/capacity/round1_regression_test.go new file mode 100644 index 000000000..fa8b2f2c8 --- /dev/null +++ b/internal/capacity/round1_regression_test.go @@ -0,0 +1,179 @@ +package capacity + +import ( + "bufio" + "context" + "fmt" + "io" + "net" + "net/http" + "strings" + "testing" + "time" +) + +func TestTunnelProxyJoinsBothCopiesWhenEitherEndpointCloses(t *testing.T) { + for _, side := range []string{"downstream", "upstream", "shutdown"} { + t.Run(side, func(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ln.Close() }) + accepted := make(chan net.Conn, 1) + go func() { + conn, err := ln.Accept() + if err == nil { + accepted <- conn + } + }() + c := validConfig(t) + c.Manifest.Targets[0].SSHURL = "ssh://" + ln.Addr().String() + g := newGuard(c) + p, err := newTunnelProxy(t.Context(), g, c.Manifest.Targets[0]) + if err != nil { + t.Fatal(err) + } + t.Cleanup(p.close) + client, err := net.Dial("tcp", p.listener.Addr().String()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Close() }) + if err := client.SetDeadline(time.Now().Add(3 * time.Second)); err != nil { + t.Fatal(err) + } + if _, err := fmt.Fprintf(client, "CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", ln.Addr(), ln.Addr()); err != nil { + t.Fatal(err) + } + response, err := http.ReadResponse(bufio.NewReader(client), &http.Request{Method: http.MethodConnect}) + if err != nil || response.StatusCode != http.StatusOK { + t.Fatalf("CONNECT failed: response=%v err=%v", response, err) + } + var upstream net.Conn + select { + case upstream = <-accepted: + case <-time.After(3 * time.Second): + t.Fatal("upstream was never connected") + } + t.Cleanup(func() { _ = upstream.Close() }) + joined := make(chan struct{}) + go func() { p.workers.Wait(); close(joined) }() + switch side { + case "downstream": + _ = client.Close() + case "upstream": + _ = upstream.Close() + case "shutdown": + p.close() + } + select { + case <-joined: + case <-time.After(time.Second): + t.Fatal("CONNECT handler remained blocked after one endpoint closed") + } + peer := upstream + if side == "upstream" { + peer = client + } + if err := peer.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatal(err) + } + if _, err := peer.Read(make([]byte, 1)); err != io.EOF { + t.Fatalf("opposite silent endpoint was not released: %v", err) + } + g.mu.Lock() + active, peak := g.active, g.peak + g.mu.Unlock() + p.mu.Lock() + tracked := len(p.conns) + p.mu.Unlock() + if active != 0 || peak != 1 || tracked != 0 || len(g.gate) != 0 { + t.Fatalf("tunnel leaked accounting/tracking/permit: active=%d peak=%d tracked=%d gate=%d", active, peak, tracked, len(g.gate)) + } + }) + } +} + +func TestParseFrameRequiresCompleteDeterministicPayload(t *testing.T) { + at := time.Unix(1700000000, 123) + for _, kind := range []string{"I", "H", "R", "A", "Q", "F", "B"} { + t.Run(kind, func(t *testing.T) { + size := 128 + if kind == "H" { + size = 1024 + } + line := strings.TrimSuffix(string(frame(kind, 7, 3, at, size)), "\n") + for _, prefix := range []string{"", "old row\x1b[0m", "CAP I old partial row "} { + for _, ending := range []string{"", "\n", "\r\n"} { + o, ok := parseFrame(prefix + line + ending) + if !ok || o.kind != kind || o.id != 7 || o.seq != 3 || !o.sent.Equal(at) { + t.Fatalf("valid %d-byte frame rejected with prefix=%q ending=%q: %+v valid=%v", size, prefix, ending, o, ok) + } + } + } + for _, bad := range []struct{ name, line string }{ + {"truncated padding", line[:len(line)-1]}, + {"expanded padding", line + "x"}, + {"corrupt padding", line[:len(line)-1] + "y"}, + {"extra header space", strings.Replace(line, " 00000007", " 00000007", 1)}, + {"tab in header", strings.Replace(line, " 00000007", "\t00000007", 1)}, + {"short numeric field", strings.Replace(line, "00000007", "7", 1)}, + {"space in padding", line[:len(line)-1] + " "}, + } { + t.Run(bad.name, func(t *testing.T) { + if _, ok := parseFrame("old screen row" + bad.line); ok { + t.Fatal("damaged payload was accepted as a complete frame") + } + }) + } + }) + } +} + +func TestCorruptOutputCannotAcknowledgeDelivery(t *testing.T) { + // A valid header with one damaged padding byte used to acknowledge the + // outstanding input and let completeDelivery return successfully. + c := validConfig(t) + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + at := time.Unix(1700000000, 0) + q := &guest{pending: map[string]time.Time{"0:3": at}, canary: true} + r := &runState{cfg: c, ctx: ctx, rec: newRecorder(io.Discard, c), guests: []*guest{q}} + damaged := frame("I", 0, 3, at, 128) + damaged[126] = 'y' + err := scanFrames(strings.NewReader(string(damaged)), func(o observation) error { + if _, ok := q.acknowledge(o); ok { + r.rec.summary.ReceivedFrames++ + r.actual(int64(frameLength(o.kind))) + } + return nil + }, func(observation) { t.Error("corrupt input became readiness proof") }) + if err != io.EOF { + t.Fatal(err) + } + if err := r.completeDelivery(); err == nil { + t.Error("corrupt payload satisfied terminal delivery completion") + } + if len(q.pending) != 1 || r.rec.summary.ReceivedFrames != 0 || r.rec.summary.PayloadBytes != 0 { + t.Fatalf("corrupt output inflated acknowledgments or nominal accounting: pending=%d received=%d bytes=%d", len(q.pending), r.rec.summary.ReceivedFrames, r.rec.summary.PayloadBytes) + } +} + +func TestCorruptHostedOutputCannotCompleteDelivery(t *testing.T) { + c := faultHelper(t, "corrupt", "") + c.DurationSeconds = 1 + s, err := Run(t.Context(), c, io.Discard) + if s.EstablishedGuests != 1 || s.ActiveFrames != 1 { + t.Fatalf("corrupt helper did not reach the steady workload: %+v err=%v", s, err) + } + if err == nil || s.DeliveryComplete { + t.Fatalf("corrupt hosted output satisfied delivery completion: %+v err=%v", s, err) + } + if s.RTTCount != 1 || s.ReceivedFrames != 1 { + t.Fatalf("corrupt output inflated delivered frames beyond the valid readiness input: %+v", s) + } + if !s.CleanupComplete { + t.Fatal("corrupt delivery did not clean up") + } +} diff --git a/internal/capacity/runner.go b/internal/capacity/runner.go index 2703dbbb2..366fd2132 100644 --- a/internal/capacity/runner.go +++ b/internal/capacity/runner.go @@ -91,7 +91,7 @@ func (h *hostRun) session(ctx context.Context) (*api.GetSessionResponse, error) if err != nil { return nil, err } - defer conn.Close() + defer func() { _ = conn.Close() }() s, err := api.NewAdminServiceClient(conn).GetSession(ctx, &api.GetSessionRequest{}) if err != nil { return nil, err @@ -132,14 +132,15 @@ func (r *runState) sendKind(q *guest, w io.Writer, stream, seq int, kind string) } now := time.Now() key := fmt.Sprintf("%d:%d", stream, seq) - if kind == "I" { + switch kind { + case "I": q.mu.Lock() if q.pending == nil { q.pending = map[string]time.Time{} } q.pending[key] = now q.mu.Unlock() - } else if kind == "R" { + case "R": q.mu.Lock() if q.challenges == nil { q.challenges = map[string]time.Time{} @@ -241,7 +242,7 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { proxy.permit = permit proxy.mu.Unlock() r.rec.event("admission_queue", map[string]any{"role": "host", "host": index, "queue_ms": float64(time.Since(queuedAt)) / float64(time.Millisecond)}) - began := time.Now() + var began time.Time go func() { err := h.h.Run(r.ctx) h.done <- err @@ -350,7 +351,8 @@ func (r *runState) cleanup() error { } for _, h := range r.hosts { if h != nil { - <-h.done + hostErr := <-h.done + _ = hostErr // Exit errors were classified before this cancellation/join. } } if err := os.RemoveAll(r.dir); err != nil { diff --git a/internal/capacity/runner_test.go b/internal/capacity/runner_test.go index 8544e5f7c..3a060ddc3 100644 --- a/internal/capacity/runner_test.go +++ b/internal/capacity/runner_test.go @@ -15,6 +15,7 @@ import ( "os" "os/exec" "path/filepath" + "runtime" "sync" "testing" "time" @@ -28,6 +29,9 @@ import ( var helperDir string func TestMain(m *testing.M) { + if len(os.Args) > 1 && os.Args[1] == "helper" { + os.Exit(runFaultHelperProcess()) + } var err error helperDir, err = os.MkdirTemp("", "capacity-test-") if err != nil { @@ -45,7 +49,11 @@ var helperError error func buildHelper(t *testing.T) string { t.Helper() helperOnce.Do(func() { - helperBinary = filepath.Join(helperDir, "capacity") + name := "capacity" + if runtime.GOOS == "windows" { + name += ".exe" + } + helperBinary = filepath.Join(helperDir, name) cmd := exec.Command("go", "build", "-o", helperBinary, "./cmd") if out, err := cmd.CombinedOutput(); err != nil { helperError = fmt.Errorf("helper build: %w: %s", err, out) @@ -92,7 +100,8 @@ func localRelay(t *testing.T, options ...func(*server.Server)) Config { cancel() _ = relay.Shutdown() select { - case <-done: + case err := <-done: + _ = err // Shutdown owns the expected serving error. case <-time.After(10 * time.Second): t.Error("relay leaked") } @@ -193,7 +202,7 @@ func TestInvalidLimitsNeverStartHostTransport(t *testing.T) { if err != nil { t.Fatal(err) } - defer ln.Close() + defer func() { _ = ln.Close() }() calls := make(chan struct{}, 1) done := make(chan struct{}) go func() { diff --git a/internal/capacity/transfer.go b/internal/capacity/transfer.go index 33181c02d..1ba3bcd2f 100644 --- a/internal/capacity/transfer.go +++ b/internal/capacity/transfer.go @@ -54,7 +54,7 @@ func (r *runState) transfer(index int, expected string) (err error) { if err != nil { return err } - defer c.Close() + defer func() { _ = c.Close() }() path := filepath.ToSlash(filepath.Join(filepath.Dir(q.host.socket), fmt.Sprintf("transfer-%d.bin", index))) defer func() { _ = c.Remove(path) }() began := time.Now() @@ -234,7 +234,7 @@ func (r *runState) restoreGuest(q *guest, began time.Time) error { // independently planned time. Guest selections run before host selections: // an overlap therefore cuts a real guest transport before its host is cut. func (r *runState) cutAction(a churnAction, planned time.Time) (func() error, error) { - began := time.Now() + var began time.Time if a.host { h := r.hosts[a.index] overlaps := 0 diff --git a/internal/capacity/whole_review_test.go b/internal/capacity/whole_review_test.go index 5812a450f..4d97f2061 100644 --- a/internal/capacity/whole_review_test.go +++ b/internal/capacity/whole_review_test.go @@ -100,10 +100,10 @@ func TestRealIdleHostTransportLossFailsBeforeNextCanary(t *testing.T) { } func TestFaultHelpersDispatchCompleteRepeatedReadinessFrames(t *testing.T) { - for _, tc := range []struct{ name, script, input, output string }{ - {"silent", silentAfterReadinessScript, "RRQII", "AABI-"}, - {"late", lateFinalReplyScript, "RRQIIF", "AABIIB"}, - {"retry", retryReadinessScript, "RRRQIIF", "-AABIIB"}, + for _, tc := range []struct{ name, input, output string }{ + {"silent", "RRQII", "AABI-"}, + {"late", "RRQIIF", "AABIIB"}, + {"retry", "RRRQIIF", "-AABIIB"}, } { for _, fragmented := range []bool{false, true} { name := tc.name + "/whole" @@ -121,7 +121,7 @@ func TestFaultHelpersDispatchCompleteRepeatedReadinessFrames(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) defer cancel() - cmd := exec.CommandContext(ctx, "/bin/sh", "-c", tc.script) + cmd := exec.CommandContext(ctx, faultHelperExecutable(t, tc.name, ""), "helper", "0", "false") var output bytes.Buffer cmd.Stdout = &output cmd.Stderr = &output @@ -233,10 +233,10 @@ func TestRunGuestAuthorizationAcrossTransports(t *testing.T) { defer cancel() foreign, raw, err := openGuest(ctx, r.guard, q.target, q.endpoint, s, signer(t), q.proxy) if raw != nil { - defer raw.Close() + defer func() { _ = raw.Close() }() } if err == nil { - defer foreign.Close() + defer func() { _ = foreign.Close() }() files, sftpErr := sftp.NewClient(foreign) if sftpErr == nil { _ = files.Close() diff --git a/internal/capacity/workload.go b/internal/capacity/workload.go index 425af206d..718a7aff8 100644 --- a/internal/capacity/workload.go +++ b/internal/capacity/workload.go @@ -35,7 +35,10 @@ func parseFrame(line string) (observation, bool) { if i < 0 { return observation{}, false } - fields := strings.Fields(line[i:]) + // Scanner removes LF (and an optional preceding CR). Accept the same + // endings for direct callers without trimming any payload whitespace. + payload := strings.TrimSuffix(strings.TrimSuffix(line[i:], "\n"), "\r") + fields := strings.Fields(payload) if len(fields) != 6 || (fields[1] != "I" && fields[1] != "H" && fields[1] != "R" && fields[1] != "A" && fields[1] != "Q" && fields[1] != "F" && fields[1] != "B") { return observation{}, false } @@ -45,7 +48,14 @@ func parseFrame(line string) (observation, bool) { if e1 != nil || e2 != nil || e3 != nil { return observation{}, false } - return observation{fields[1], id, seq, time.Unix(0, ns)}, true + sent := time.Unix(0, ns) + // Counting nominal frame bytes is valid only for the entire canonical + // header and deterministic padding, not merely six parseable fields. + want := frame(fields[1], id, seq, sent, frameLength(fields[1])) + if payload != string(want[:len(want)-1]) { + return observation{}, false + } + return observation{fields[1], id, seq, sent}, true } // Helper is the runner's harmless hosted command. It echoes fixed 128-byte @@ -60,7 +70,7 @@ func helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active if err != nil { return err } - defer term.Restore(int(f.Fd()), state) + defer func() { _ = term.Restore(int(f.Fd()), state) }() } if _, err := io.WriteString(out, "CAP READY\n"); err != nil { return err From b86d705e96a7625ad47bc990213c71ae91b24432 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 20:20:29 -0700 Subject: [PATCH 08/16] test(capacity): capture raw concurrent churn terminal traffic --- internal/capacity/churn_test.go | 1 + internal/capacity/client.go | 3 + internal/capacity/config.go | 2 + internal/capacity/terminal_capture_test.go | 71 ++++++++++++++++++++++ 4 files changed, 77 insertions(+) create mode 100644 internal/capacity/terminal_capture_test.go diff --git a/internal/capacity/churn_test.go b/internal/capacity/churn_test.go index 83546fd96..660736c4d 100644 --- a/internal/capacity/churn_test.go +++ b/internal/capacity/churn_test.go @@ -34,6 +34,7 @@ func TestRealTransportChurnRestoresOriginalSessions(t *testing.T) { c.PayloadBudgetBytes = 4 << 20 c.HostWSPercent = 50 c.GuestWSPercent = 50 + captureChurnTerminals(t, &c) var out bytes.Buffer s, err := Run(t.Context(), c, &out) if err != nil { diff --git a/internal/capacity/client.go b/internal/capacity/client.go index d90b842d2..2be5cbe96 100644 --- a/internal/capacity/client.go +++ b/internal/capacity/client.go @@ -212,6 +212,9 @@ func (q *guest) connect(ctx context.Context, r *runState) error { if err != nil { return err } + if r.cfg.terminalIO != nil { + in, out = r.cfg.terminalIO(q.id, ch, q.generation, in, out) + } session.Stderr = io.Discard if err = session.Shell(); err != nil { return err diff --git a/internal/capacity/config.go b/internal/capacity/config.go index d238b8ac4..6853b1a96 100644 --- a/internal/capacity/config.go +++ b/internal/capacity/config.go @@ -58,6 +58,8 @@ type Config struct { TransferConcurrency int `json:"transfer_concurrency"` ChurnPercent int `json:"churn_percent"` WatchdogFile string `json:"watchdog_file"` + // terminalIO observes a local test's raw PTY traffic before frame parsing. + terminalIO func(guest, channel, generation int, in io.WriteCloser, out io.Reader) (io.WriteCloser, io.Reader) } const MaxPayload = int64(50 << 30) diff --git a/internal/capacity/terminal_capture_test.go b/internal/capacity/terminal_capture_test.go new file mode 100644 index 000000000..6c26bf441 --- /dev/null +++ b/internal/capacity/terminal_capture_test.go @@ -0,0 +1,71 @@ +package capacity + +import ( + "fmt" + "io" + "sync" + "testing" + "time" +) + +// Capture the same concurrent admissions as the churn test, before parsing or +// terminal normalization. Keep diagnostics bounded and print only on failure. +func captureChurnTerminals(t *testing.T, c *Config) { + t.Helper() + var mu sync.Mutex + var captures []*terminalCapture + c.terminalIO = func(guest, channel, generation int, in io.WriteCloser, out io.Reader) (io.WriteCloser, io.Reader) { + capture := &terminalCapture{name: fmt.Sprintf("guest=%d channel=%d generation=%d", guest, channel, generation), started: time.Now()} + mu.Lock() + captures = append(captures, capture) + mu.Unlock() + return terminalCaptureInput{in, capture}, io.TeeReader(out, terminalCaptureOutput{capture}) + } + t.Cleanup(func() { + if !t.Failed() { + return + } + mu.Lock() + defer mu.Unlock() + for _, capture := range captures { + capture.mu.Lock() + t.Logf("terminal %s (last 16384 diagnostic bytes):\n%s", capture.name, capture.tail) + capture.mu.Unlock() + } + }) +} + +type terminalCapture struct { + mu sync.Mutex + name string + started time.Time + tail []byte +} + +func (c *terminalCapture) record(direction string, p []byte) { + line := fmt.Appendf(nil, "%s %s %q\n", time.Since(c.started), direction, p) + c.mu.Lock() + c.tail = append(c.tail, line...) + if len(c.tail) > 16384 { + c.tail = c.tail[len(c.tail)-16384:] + } + c.mu.Unlock() +} + +type terminalCaptureInput struct { + io.WriteCloser + capture *terminalCapture +} + +func (w terminalCaptureInput) Write(p []byte) (int, error) { + n, err := w.WriteCloser.Write(p) + w.capture.record("sent", p[:n]) + return n, err +} + +type terminalCaptureOutput struct{ capture *terminalCapture } + +func (w terminalCaptureOutput) Write(p []byte) (int, error) { + w.capture.record("received", p) + return len(p), nil +} From f1275b8b8edd1d476bf8a621c56509ebb07040cd Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 20:32:09 -0700 Subject: [PATCH 09/16] fix(capacity): authenticate Consul WebSocket entry fallback --- internal/capacity/fallback_cert_test.go | 248 +++++++++++++++++++++++ internal/capacity/forwarded_cert_test.go | 7 +- internal/capacity/guard.go | 13 +- 3 files changed, 263 insertions(+), 5 deletions(-) create mode 100644 internal/capacity/fallback_cert_test.go diff --git a/internal/capacity/fallback_cert_test.go b/internal/capacity/fallback_cert_test.go new file mode 100644 index 000000000..651d3e67c --- /dev/null +++ b/internal/capacity/fallback_cert_test.go @@ -0,0 +1,248 @@ +package capacity + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/json" + "io" + "log/slog" + "net" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/owenthereal/upterm/host/api" + "github.com/owenthereal/upterm/routing" + "github.com/owenthereal/upterm/server" + "golang.org/x/crypto/ssh" +) + +// Falling back to all manifest pins/principals, or permitting raw entry keys, +// would admit one of the rejected cases below. Each authority stays paired +// with its own principal even when the origin and entry share an authority. +func TestGuestKeyConsulWebSocketFallback(t *testing.T) { + for _, shared := range []bool{false, true} { + name := "distinct authorities" + if shared { + name = "shared authority" + } + t.Run(name, func(t *testing.T) { + c := validConfig(t) + originKey, entryKey, thirdKey := signer(t), signer(t), signer(t) + if shared { + entryKey = originKey + } + origin := c.Manifest.Targets[0] + origin.HostKeys = []string{string(ssh.MarshalAuthorizedKey(originKey.PublicKey()))} + entry := origin + entry.MachineID = "entry" + entry.SSHURL, entry.WSURL = "ssh://[::1]:2222", "ws://[::1]:8080" + entry.PublicWSURL = "wss://entry.example:443" + entry.NodeAddresses, entry.AllowedIPs = []string{"[::1]:2222"}, []string{"::1"} + entry.HostKeys = []string{string(ssh.MarshalAuthorizedKey(entryKey.PublicKey()))} + third := origin + third.MachineID = "third" + third.NodeAddresses, third.AllowedIPs = []string{"127.0.0.2:2222"}, []string{"127.0.0.2"} + third.HostKeys = []string{string(ssh.MarshalAuthorizedKey(thirdKey.PublicKey()))} + c.Manifest.Targets = []Target{origin, entry, third} + g := newGuard(c) + s := &api.GetSessionResponse{SessionId: "session", SshUser: "session", NodeAddr: "127.0.0.1:2222"} + cert := func(authority ssh.Signer, principal string, change func(*ssh.Certificate)) *ssh.Certificate { + c := &ssh.Certificate{Key: authority.PublicKey(), CertType: ssh.HostCert, ValidPrincipals: []string{principal}, ValidBefore: ssh.CertTimeInfinity} + if change != nil { + change(c) + } + if err := c.SignCert(rand.Reader, authority); err != nil { + t.Fatal(err) + } + return c + } + for _, endpoint := range []struct{ transport, addr, principal string }{ + {"ws", "[::1]:8080", "::1"}, + {"wss", "entry.example:443", "entry.example"}, + } { + t.Run(endpoint.transport, func(t *testing.T) { + check, err := g.guestKey(t.Context(), entry, endpoint.transport, s) + if err != nil { + t.Fatal(err) + } + for name, key := range map[string]ssh.PublicKey{ + "origin certificate": cert(originKey, "127.0.0.1", nil), + "origin raw pin": originKey.PublicKey(), + "entry certificate": cert(entryKey, endpoint.principal, nil), + } { + t.Run(name, func(t *testing.T) { + if err := check(endpoint.addr, nil, key); err != nil { + t.Fatalf("authenticated exact route refused: %v", err) + } + }) + } + corrupt := cert(entryKey, endpoint.principal, nil) + corrupt.Signature.Blob[0] ^= 1 + bad := map[string]ssh.PublicKey{ + "entry authority third principal": cert(entryKey, "127.0.0.2", nil), + "origin authority third principal": cert(originKey, "127.0.0.2", nil), + "third authority entry principal": cert(thirdKey, endpoint.principal, nil), + "third authority origin principal": cert(thirdKey, "127.0.0.1", nil), + "third certificate": cert(thirdKey, "127.0.0.2", nil), + "third raw key": thirdKey.PublicKey(), + "foreign authority": cert(signer(t), endpoint.principal, nil), + "user certificate": cert(entryKey, endpoint.principal, func(c *ssh.Certificate) { c.CertType = ssh.UserCert }), + "expired": cert(entryKey, endpoint.principal, func(c *ssh.Certificate) { c.ValidBefore = uint64(time.Now().Add(-time.Hour).Unix()) }), + "future": cert(entryKey, endpoint.principal, func(c *ssh.Certificate) { c.ValidAfter = uint64(time.Now().Add(time.Hour).Unix()) }), + "invalid signature": corrupt, + } + if !shared { + bad["entry authority origin principal"] = cert(entryKey, "127.0.0.1", nil) + bad["origin authority entry principal"] = cert(originKey, endpoint.principal, nil) + bad["entry raw key"] = entryKey.PublicKey() + } + for name, key := range bad { + t.Run(name, func(t *testing.T) { + if err := check(endpoint.addr, nil, key); err == nil { + t.Fatal("invalid fallback certificate/key accepted") + } + }) + } + }) + } + }) + } +} + +// An initially unavailable advertised node forces the real Consul WebSocket +// handler to hand off to the entry's SSH proxy. Restore that TCP path when the +// handler reports the refused dial, before the entry authenticates the guest. +// Removing entry certificate fallback fails the handshake before any shell. +func TestRealConsulWebSocketFallbackCertificate(t *testing.T) { + clearProxies(t) + registry := forwardedSessionRegistry(t) + reserved, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + node := reserved.Addr().String() + if err := reserved.Close(); err != nil { + t.Fatal(err) + } + originKey, entryKey := signer(t), signer(t) + origin := forwardedRelay(t, "origin", "127.0.0.1", originKey, routing.ModeConsul, registry, func(s *server.Server) { + s.NodeAddr = node + s.HopTrusted = []ssh.PublicKey{entryKey.PublicKey()} + }) + var fallback atomic.Int32 + var once sync.Once + entry := forwardedRelay(t, "entry", "::1", entryKey, routing.ModeConsul, registry, func(s *server.Server) { + s.NodeTrusted = []ssh.PublicKey{originKey.PublicKey()} + s.Logger = slog.New(&fallbackLogHandler{Handler: slog.NewTextHandler(io.Discard, nil), record: func(r slog.Record) { + if r.Message == "error dialing the session's node" { + fallback.Add(1) + once.Do(func() { restoreFallbackNode(t, node, strings.TrimPrefix(origin.SSHURL, "ssh://")) }) + } + }}) + }) + c := validConfig(t) + c.Manifest.Targets = []Target{origin, entry} + c.HostTarget, c.EntryTarget = "origin", "entry" + c.HelperExecutable = buildHelper(t) + c.GuestWSPercent = 100 + c.Active, c.DurationSeconds, c.ChannelsPerGuest = true, 1, 4 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if dir := os.Getenv("CAPACITY_EVIDENCE_DIR"); dir != "" { + manifest, e := json.MarshalIndent(c.Manifest, "", " ") + if e != nil { + t.Fatal(e) + } + phase, e := json.MarshalIndent(c, "", " ") + if e != nil { + t.Fatal(e) + } + for name, data := range map[string][]byte{"manifest.json": manifest, "phase.json": phase, "events.jsonl": out.Bytes()} { + if e := os.WriteFile(filepath.Join(dir, "fallback-"+name), data, 0600); e != nil { + t.Fatal(e) + } + } + } + if fallback.Load() != 1 { + t.Fatalf("expected exactly one real TCP dial failure and handoff, got %d", fallback.Load()) + } + if err != nil { + t.Fatalf("fallback shell workload: %v\n%s", err, out.String()) + } + if s.EstablishedHosts != 1 || s.EstablishedGuests != 1 || s.EstablishedChannels != 4 || s.AdmissionFailures != 0 { + t.Fatalf("fallback admissions incomplete: %+v", s) + } + if !s.DeliveryComplete || !s.CleanupComplete || s.StopReason != "duration" || s.SteadySeconds < 1 || s.ActiveFrames != 1 || s.InputFrames != 5 || s.RTTCount != 5 || s.PayloadBytes == 0 { + t.Fatalf("fallback framed delivery/cleanup incomplete: %+v", s) + } +} + +type fallbackLogHandler struct { + slog.Handler + record func(slog.Record) +} + +func (h *fallbackLogHandler) Handle(ctx context.Context, r slog.Record) error { + h.record(r) + return h.Handler.Handle(ctx, r) +} +func (h *fallbackLogHandler) WithAttrs(attrs []slog.Attr) slog.Handler { + return &fallbackLogHandler{Handler: h.Handler.WithAttrs(attrs), record: h.record} +} +func (h *fallbackLogHandler) WithGroup(name string) slog.Handler { + return &fallbackLogHandler{Handler: h.Handler.WithGroup(name), record: h.record} +} + +func restoreFallbackNode(t *testing.T, node, upstream string) { + t.Helper() + ln, err := net.Listen("tcp", node) + if err != nil { + t.Errorf("restore origin TCP listener: %v", err) + return + } + ctx, cancel := context.WithCancel(context.Background()) + var workers sync.WaitGroup + workers.Add(1) + go func() { + defer workers.Done() + for { + down, err := ln.Accept() + if err != nil { + return + } + workers.Go(func() { + defer down.Close() + stopDown := context.AfterFunc(ctx, func() { _ = down.Close() }) + defer stopDown() + up, err := (&net.Dialer{Timeout: time.Second}).DialContext(ctx, "tcp", upstream) + if err != nil { + t.Errorf("restored origin dial: %v", err) + return + } + defer up.Close() + stopUp := context.AfterFunc(ctx, func() { _ = up.Close() }) + defer stopUp() + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = io.Copy(up, down) + _ = up.Close() + }() + _, _ = io.Copy(down, up) + _ = down.Close() + <-done + }) + } + }() + t.Cleanup(func() { + cancel() + _ = ln.Close() + workers.Wait() + }) +} diff --git a/internal/capacity/forwarded_cert_test.go b/internal/capacity/forwarded_cert_test.go index 17d7f57b0..b212c3cca 100644 --- a/internal/capacity/forwarded_cert_test.go +++ b/internal/capacity/forwarded_cert_test.go @@ -212,7 +212,7 @@ func TestRealForwardedGuestCertificate(t *testing.T) { } } -func forwardedRelay(t *testing.T, id, ip string, key ssh.Signer, mode routing.Mode, registry *url.URL) Target { +func forwardedRelay(t *testing.T, id, ip string, key ssh.Signer, mode routing.Mode, registry *url.URL, options ...func(*server.Server)) Target { t.Helper() sshln, err := net.Listen("tcp", net.JoinHostPort(ip, "0")) if err != nil { @@ -238,6 +238,9 @@ func forwardedRelay(t *testing.T, id, ip string, key ssh.Signer, mode routing.Mo t.Fatal(err) } relay := &server.Server{NodeAddr: sshln.Addr().String(), HostSigners: []ssh.Signer{cert}, Signers: []ssh.Signer{key}, NetworkProvider: network, SessionManager: sm, MetricsProvider: provider.NewDiscardProvider(), Logger: logger} + for _, option := range options { + option(relay) + } ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- relay.ServeListeners(ctx, server.Listeners{SSH: sshln, WS: wsln}) }() @@ -259,7 +262,7 @@ func forwardedRelay(t *testing.T, id, ip string, key ssh.Signer, mode routing.Mo } metrics.Start() t.Cleanup(metrics.Close) - return Target{AppID: "test-app", MachineID: id, SSHURL: "ssh://" + sshln.Addr().String(), WSURL: "ws://" + wsln.Addr().String(), NodeAddresses: []string{sshln.Addr().String()}, AllowedIPs: []string{ip}, HostKeys: []string{string(ssh.MarshalAuthorizedKey(key.PublicKey()))}, TelemetryURL: metrics.URL + "/metrics"} + return Target{AppID: "test-app", MachineID: id, SSHURL: "ssh://" + sshln.Addr().String(), WSURL: "ws://" + wsln.Addr().String(), NodeAddresses: []string{relay.NodeAddr}, AllowedIPs: []string{ip}, HostKeys: []string{string(ssh.MarshalAuthorizedKey(key.PublicKey()))}, TelemetryURL: metrics.URL + "/metrics"} } // The registry supplies only the local storage protocol needed for a single diff --git a/internal/capacity/guard.go b/internal/capacity/guard.go index 6a41e9fd3..f0376f49c 100644 --- a/internal/capacity/guard.go +++ b/internal/capacity/guard.go @@ -185,15 +185,22 @@ func (g *guard) guestKey(ctx context.Context, entry Target, transport string, s } // SSH terminates at the entry. Embedded-mode WebSocket routing also goes // through the entry's SSH proxy. Consul-mode WebSocket routing instead - // pipes SSH to the session's node, whose certificate the guest receives. + // pipes SSH to the session's node, or hands off to the entry's SSH proxy + // if that TCP dial fails. Each certificate must match its own authority + // and principal; a raw key cannot establish the entry fallback identity. // Its plain SSH user distinguishes it from the validated embedded route. if transport != "ssh" && s.SshUser == s.SessionId { for _, origin := range g.cfg.Manifest.Targets { if origin.MachineID != entry.MachineID && contains(origin.NodeAddresses, s.NodeAddr) { check := pinnedKey(origin) + fallback := pinnedKey(entry) node := s.NodeAddr - return func(_ string, remote net.Addr, key ssh.PublicKey) error { - return check(node, remote, key) + return func(addr string, remote net.Addr, key ssh.PublicKey) error { + err := check(node, remote, key) + if _, cert := key.(*ssh.Certificate); err == nil || !cert { + return err + } + return fallback(addr, remote, key) }, nil } } From b4346add55122dc748fe30cc2c2201917e061ab0 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 20:41:58 -0700 Subject: [PATCH 10/16] fix(capacity): require explicit websocket relay principals --- internal/capacity/fallback_cert_test.go | 4 ++++ internal/capacity/guard.go | 8 +++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/internal/capacity/fallback_cert_test.go b/internal/capacity/fallback_cert_test.go index 651d3e67c..c37bf1971 100644 --- a/internal/capacity/fallback_cert_test.go +++ b/internal/capacity/fallback_cert_test.go @@ -85,6 +85,10 @@ func TestGuestKeyConsulWebSocketFallback(t *testing.T) { corrupt := cert(entryKey, endpoint.principal, nil) corrupt.Signature.Blob[0] ^= 1 bad := map[string]ssh.PublicKey{ + "missing entry principal": cert(entryKey, endpoint.principal, func(c *ssh.Certificate) { c.ValidPrincipals = nil }), + "missing origin principal": cert(originKey, "127.0.0.1", func(c *ssh.Certificate) { c.ValidPrincipals = nil }), + "empty entry principal list": cert(entryKey, endpoint.principal, func(c *ssh.Certificate) { c.ValidPrincipals = []string{} }), + "empty origin principal list": cert(originKey, "127.0.0.1", func(c *ssh.Certificate) { c.ValidPrincipals = []string{} }), "entry authority third principal": cert(entryKey, "127.0.0.2", nil), "origin authority third principal": cert(originKey, "127.0.0.2", nil), "third authority entry principal": cert(thirdKey, endpoint.principal, nil), diff --git a/internal/capacity/guard.go b/internal/capacity/guard.go index f0376f49c..c191cf994 100644 --- a/internal/capacity/guard.go +++ b/internal/capacity/guard.go @@ -196,8 +196,14 @@ func (g *guard) guestKey(ctx context.Context, entry Target, transport string, s fallback := pinnedKey(entry) node := s.NodeAddr return func(addr string, remote net.Addr, key ssh.PublicKey) error { + // CheckCert treats no principals as valid for every host; + // this route requires an explicit origin or entry principal. + cert, isCert := key.(*ssh.Certificate) + if isCert && len(cert.ValidPrincipals) == 0 { + return errors.New("relay certificate has no explicit principal") + } err := check(node, remote, key) - if _, cert := key.(*ssh.Certificate); err == nil || !cert { + if err == nil || !isCert { return err } return fallback(addr, remote, key) From 9c712d38cdff8cb8744e5d10fbd26aa87edaa532 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 20:54:49 -0700 Subject: [PATCH 11/16] test(capacity): make fallback cleanup error discards explicit --- internal/capacity/fallback_cert_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/capacity/fallback_cert_test.go b/internal/capacity/fallback_cert_test.go index c37bf1971..0013c7386 100644 --- a/internal/capacity/fallback_cert_test.go +++ b/internal/capacity/fallback_cert_test.go @@ -221,7 +221,7 @@ func restoreFallbackNode(t *testing.T, node, upstream string) { return } workers.Go(func() { - defer down.Close() + defer func() { _ = down.Close() }() stopDown := context.AfterFunc(ctx, func() { _ = down.Close() }) defer stopDown() up, err := (&net.Dialer{Timeout: time.Second}).DialContext(ctx, "tcp", upstream) @@ -229,7 +229,7 @@ func restoreFallbackNode(t *testing.T, node, upstream string) { t.Errorf("restored origin dial: %v", err) return } - defer up.Close() + defer func() { _ = up.Close() }() stopUp := context.AfterFunc(ctx, func() { _ = up.Close() }) defer stopUp() done := make(chan struct{}) From c6690c07f885e43e58553571c89fa13d6a94f5d0 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 21:33:51 -0700 Subject: [PATCH 12/16] fix(capacity): enforce complete input epochs and bounded evidence --- internal/capacity/README.md | 16 + internal/capacity/client.go | 6 +- internal/capacity/config.go | 4 + internal/capacity/mixed_test.go | 1 + internal/capacity/report.go | 34 ++- internal/capacity/round2_regression_test.go | 316 ++++++++++++++++++++ internal/capacity/runner.go | 9 +- internal/capacity/transfer.go | 9 +- internal/capacity/workload.go | 50 +++- internal/capacity/workload_test.go | 4 +- 10 files changed, 437 insertions(+), 12 deletions(-) create mode 100644 internal/capacity/round2_regression_test.go diff --git a/internal/capacity/README.md b/internal/capacity/README.md index d64706271..cb8326f73 100644 --- a/internal/capacity/README.md +++ b/internal/capacity/README.md @@ -167,6 +167,12 @@ count and establishes each channel's live measurement boundary. Every new transport clears old barriers so replay cannot activate its live reader. Planned reconnects reset tracking only for the affected source generation; survivor generations retain gap/reordering checks. +At the live boundary, each receiver snapshots every peer stream's input +allocation counter. Earlier frames may arrive late or be replayed; every +sequence allocated afterward, including the first, must arrive in order. +After a source reconnects, its successful restoration publishes the next input +sequence atomically with clearing the expected outage. Receivers then require +the complete new range without treating disrupted input or probes as gaps. Steady input has an independent schedule per guest. Expected restoration pauses only that guest; it does not hold up healthy survivors. A missed @@ -259,6 +265,16 @@ Go version, goroutines, heap/runtime memory and, on Linux, `/proc/self/status` and `/proc/self/stat` for RSS/CPU analysis. Sampling does not diagnose pressure; the coordinator evaluates Fly metrics/OOM and generator pressure externally. +Each run can write at most 512 MiB of encoded JSONL, including escaping and +newlines. A complete event that would exceed the remaining budget is rejected +before writing; the run stops with `output_budget`, closes its resources, and +reports incomplete delivery. Earlier complete events remain available; a final +summary may not fit. This limit is separate from application payload accounting +and can end a phase before its configured duration. At observed normal sample +sizes, two-target metrics and generator samples use about 32 MiB per two hours, leaving headroom +for admission, churn and summary events. Each metrics response is still bounded +to 1 MiB before JSON encoding. + Before startup and once per second, an existing `watchdog_file` requests a graceful safety stop. If `watchdog_heartbeat_file` is configured, it must exist at startup and have an mtime no older than 30 seconds; stale, missing or diff --git a/internal/capacity/client.go b/internal/capacity/client.go index 2be5cbe96..920313d95 100644 --- a/internal/capacity/client.go +++ b/internal/capacity/client.go @@ -82,6 +82,7 @@ type guest struct { canary bool probeSeq [4]int + inputStart [4]int pending map[string]time.Time challenges map[string]time.Time lifetimeStop func() bool @@ -221,12 +222,14 @@ func (q *guest) connect(ctx context.Context, r *runState) error { } q.writers = append(q.writers, in) stream := q.id*4 + ch + q.mu.Lock() probeSeq := q.probeSeq[ch] q.probeSeq[ch]++ if ch == 0 { probeSeq = q.seq q.seq++ } + q.mu.Unlock() ready := make(chan struct{}, 1) helperReady := make(chan struct{}, 1) done := make(chan struct{}) @@ -234,7 +237,7 @@ func (q *guest) connect(ctx context.Context, r *runState) error { go func() { defer close(done) tracker := sequenceTracker{} - epochs := map[string]int{} + epochs := map[string]inputEpoch{} live := false baseline := time.Time{} hostKey := fmt.Sprintf("H:%d", q.host.id) @@ -250,6 +253,7 @@ func (q *guest) connect(ctx context.Context, r *runState) error { } if !live { tracker[hostKey] = o.seq - 1 + r.baselineInputs(tracker, epochs, q.host.id) baseline = barrier.sent live = true } else if last := tracker[hostKey]; last != o.seq-1 { diff --git a/internal/capacity/config.go b/internal/capacity/config.go index 6853b1a96..8d3ddf46a 100644 --- a/internal/capacity/config.go +++ b/internal/capacity/config.go @@ -60,6 +60,10 @@ type Config struct { WatchdogFile string `json:"watchdog_file"` // terminalIO observes a local test's raw PTY traffic before frame parsing. terminalIO func(guest, channel, generation int, in io.WriteCloser, out io.Reader) (io.WriteCloser, io.Reader) + // beforeRestoreCommit lets local tests interleave a later actual cut. + beforeRestoreCommit func() + // outputLimit lets local tests exercise the same quota with small logs. + outputLimit int64 } const MaxPayload = int64(50 << 30) diff --git a/internal/capacity/mixed_test.go b/internal/capacity/mixed_test.go index 4cd1eee29..f2241e9b3 100644 --- a/internal/capacity/mixed_test.go +++ b/internal/capacity/mixed_test.go @@ -26,6 +26,7 @@ func TestRealMixedTransportPairing(t *testing.T) { c.Hosts, c.GuestsPerHost = tc.hosts, tc.guests c.HostWSPercent, c.GuestWSPercent = 50, 50 c.DurationSeconds = 1 + captureChurnTerminals(t, &c) var out bytes.Buffer s, err := Run(t.Context(), c, &out) if err != nil { diff --git a/internal/capacity/report.go b/internal/capacity/report.go index ca9ca86fa..79fe94ef7 100644 --- a/internal/capacity/report.go +++ b/internal/capacity/report.go @@ -2,6 +2,7 @@ package capacity import ( "encoding/json" + "errors" "io" "sync" "time" @@ -57,9 +58,33 @@ type recorder struct { mu sync.Mutex encoder *json.Encoder err error + fail func(error) summary Summary } +const maxOutputBytes = int64(512 << 20) + +var errOutputBudget = errors.New("output_budget") + +// json.Encoder writes one complete encoded event at a time. Check its actual +// encoded size, including escaping and LF, before forwarding any of it. +type outputBudgetWriter struct { + out io.Writer + remaining int64 +} + +func (w *outputBudgetWriter) Write(p []byte) (int, error) { + if int64(len(p)) > w.remaining { + return 0, errOutputBudget + } + n, err := w.out.Write(p) + w.remaining -= int64(n) + if n != len(p) && err == nil { + err = io.ErrShortWrite + } + return n, err +} + func (r *recorder) event(kind string, fields map[string]any) { r.mu.Lock() defer r.mu.Unlock() @@ -71,8 +96,15 @@ func (r *recorder) event(kind string, fields map[string]any) { fields["run_id"] = r.summary.RunID if r.err == nil { r.err = r.encoder.Encode(fields) + if r.err != nil && r.fail != nil { + r.fail(r.err) + } } } func newRecorder(out io.Writer, c Config) *recorder { - return &recorder{encoder: json.NewEncoder(out), summary: Summary{RunID: c.Manifest.RunID, RTTSampling: "uniform_reservoir_100000", OfferedHosts: c.Hosts, OfferedGuests: c.Hosts * c.GuestsPerHost, OfferedChannels: c.Hosts * c.GuestsPerHost * c.ChannelsPerGuest, PayloadBudgetBytes: c.PayloadBudgetBytes}} + limit := maxOutputBytes + if c.outputLimit > 0 && c.outputLimit < limit { + limit = c.outputLimit + } + return &recorder{encoder: json.NewEncoder(&outputBudgetWriter{out: out, remaining: limit}), summary: Summary{RunID: c.Manifest.RunID, RTTSampling: "uniform_reservoir_100000", OfferedHosts: c.Hosts, OfferedGuests: c.Hosts * c.GuestsPerHost, OfferedChannels: c.Hosts * c.GuestsPerHost * c.ChannelsPerGuest, PayloadBudgetBytes: c.PayloadBudgetBytes}} } diff --git a/internal/capacity/round2_regression_test.go b/internal/capacity/round2_regression_test.go new file mode 100644 index 000000000..ddaeeefb2 --- /dev/null +++ b/internal/capacity/round2_regression_test.go @@ -0,0 +1,316 @@ +package capacity + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +// The first source finishes its probe before the second terminal subscribes. +// Losing only the first steady fan-out frame must not turn the rest of that +// source's sequence into a new, apparently complete range. +func TestFirstSteadyPeerFrameMustReachEveryReceiver(t *testing.T) { + for _, damage := range []string{"lost", "corrupt"} { + t.Run(damage, func(t *testing.T) { + c := localRelay(t) + c.GuestsPerHost = 2 + c.DurationSeconds = 2 + firstReady := make(chan struct{}) + var readyOnce sync.Once + c.terminalIO = func(guest, channel, generation int, in io.WriteCloser, out io.Reader) (io.WriteCloser, io.Reader) { + if guest == 1 { + select { + case <-firstReady: + case <-t.Context().Done(): + case <-time.After(10 * time.Second): + t.Error("first source never completed its probe") + } + } + return in, &frameTransformReader{reader: bufio.NewReader(out), transform: func(line []byte) []byte { + o, ok := parseFrame(string(line)) + if ok && o.kind == "I" && o.id == 0 { + if guest == 0 && o.seq == 0 { + readyOnce.Do(func() { close(firstReady) }) + } + if guest == 1 && o.seq == 1 { + if damage == "lost" { + return nil + } + line[len(line)-3] = 'y' + } + } + return line + }} + } + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if s.EstablishedGuests != 2 { + t.Fatalf("fixture did not admit both receivers: %v\n%s", err, out.String()) + } + if err == nil || s.DeliveryComplete { + t.Fatalf("%s first peer frame was accepted as complete delivery: %+v", damage, s) + } + if !strings.Contains(err.Error(), "sequence gap/reordering") { + t.Fatalf("missing first peer frame was not detected at the next sequence: %v\n%s", err, out.String()) + } + if !s.CleanupComplete { + t.Fatal("failed delivery did not clean up") + } + }) + } +} + +type frameTransformReader struct { + reader *bufio.Reader + transform func([]byte) []byte + pending []byte + err error +} + +func (r *frameTransformReader) Read(p []byte) (int, error) { + for len(r.pending) == 0 && r.err == nil { + line, err := r.reader.ReadBytes('\n') + r.pending, r.err = r.transform(line), err + } + n := copy(p, r.pending) + r.pending = r.pending[n:] + if n != 0 { + return n, nil + } + return 0, r.err +} + +func TestLaterHostCutOwnsGuestRestorationCommit(t *testing.T) { + c := localRelay(t) + c.Hosts = 2 + ctx, cancel := context.WithCancelCause(t.Context()) + dir, err := os.MkdirTemp("", "cap-commit-") + if err != nil { + t.Fatal(err) + } + r := &runState{ctx: ctx, cancel: cancel, cfg: c, guard: newGuard(c), rec: newRecorder(io.Discard, c), helper: c.HelperExecutable, dir: dir} + t.Cleanup(func() { + if err := r.cleanup(); err != nil { + t.Error(err) + } + }) + if err := r.admit(nil); err != nil { + t.Fatal(err) + } + var repairHost func() error + r.cfg.beforeRestoreCommit = func() { + // q.connect has completed, but the older repair has not committed its + // successful transition. This cut must keep ownership of the outage. + repairHost, err = r.cutAction(churnAction{host: true, index: 0}, time.Now()) + if err != nil { + t.Fatal(err) + } + } + if err := r.cutGuest(r.guests[0]); err != nil { + t.Fatal(err) + } + r.cfg.beforeRestoreCommit = nil + q := r.guests[0] + q.mu.Lock() + expected, disruption := q.expected, q.disruption + q.mu.Unlock() + if !expected || disruption != 2 || r.rec.summary.Reconnects != 0 { + t.Fatalf("stale repair cleared or counted the later host outage: expected=%v disruption=%d reconnects=%d", expected, disruption, r.rec.summary.Reconnects) + } + if err := q.active(r); err != nil { + t.Fatalf("new outage did not keep input paused: %v", err) + } + if err := r.guests[1].active(r); err != nil { + t.Fatalf("unaffected survivor lost its terminal: %v", err) + } + if r.rec.summary.PausedActiveFrames != 1 || r.rec.summary.ActiveFrames != 1 { + t.Fatal("overlap changed survivor or interrupted input accounting") + } + if repairHost == nil { + t.Fatal("fixture did not perform the later host cut") + } + if err := repairHost(); err != nil { + t.Fatal(err) + } + q.mu.Lock() + expected, live := q.expected, q.transportLive + q.mu.Unlock() + if expected || !live || r.rec.summary.Reconnects != 2 || r.rec.summary.EstablishedGuests != 2 || r.rec.summary.ActualGuestCuts != 1 || r.rec.summary.ActualHostCuts != 1 { + t.Fatalf("later owner did not restore the same guest exactly once: expected=%v live=%v summary=%+v", expected, live, r.rec.summary) + } +} + +func TestInputRangeExcludesOnlyPreBoundaryFrames(t *testing.T) { + for _, ch := range []int{0, 3} { + for _, missing := range []bool{false, true} { + h := &hostRun{id: 0} + q := &guest{host: h, seq: 4, probeSeq: [4]int{0, 0, 0, 4}} + r := &runState{cfg: Config{ChannelsPerGuest: 4}, guests: []*guest{q}} + tracker, epochs := sequenceTracker{}, map[string]inputEpoch{} + r.baselineInputs(tracker, epochs, 0) + // Replayed and in-flight earlier input may arrive after the B frame. + for _, seq := range []int{0, 3, 2} { + if err := r.observe(tracker, epochs, observation{kind: "I", id: ch, seq: seq}); err != nil { + t.Fatalf("pre-boundary input became a live gap: %v", err) + } + } + if !missing { + if err := r.observe(tracker, epochs, observation{kind: "I", id: ch, seq: 4}); err != nil { + t.Fatal(err) + } + } + err := r.observe(tracker, epochs, observation{kind: "I", id: ch, seq: 5}) + if (err != nil) != missing { + t.Fatalf("channel %d complete allocated range: missing=%v err=%v", ch, missing, err) + } + } + } +} + +func TestFirstPeerFrameAfterSourceReconnectIsRequired(t *testing.T) { + q := &guest{generation: 1, inputStart: [4]int{7}, seq: 9} + r := &runState{guests: []*guest{q}} + tracker := sequenceTracker{"I:0": 3} + epochs := map[string]inputEpoch{"I:0": {generation: 0, floor: 0}} + // The reconnect probe and interrupted old epoch are excluded. A delayed + // receiver still requires sequence 7, even when 8 has already been offered. + if err := r.observe(tracker, epochs, observation{kind: "I", id: 0, seq: 6}); err != nil { + t.Fatal(err) + } + if err := r.observe(tracker, epochs, observation{kind: "I", id: 0, seq: 8}); err == nil { + t.Fatal("lost first input after peer restoration became a new baseline") + } + if err := r.observe(tracker, epochs, observation{kind: "I", id: 0, seq: 7}); err != nil { + t.Fatal(err) + } + if err := r.observe(tracker, epochs, observation{kind: "I", id: 0, seq: 8}); err != nil { + t.Fatal(err) + } +} + +func TestRecorderBoundsCumulativeEncodedOutput(t *testing.T) { + // JSON escapes every NUL to six bytes. Retain only byte counts, so the + // regression exercises the real cap without allocating a huge log. + out := &countOutput{} + r := newRecorder(out, Config{}) + body := strings.Repeat("\x00", 1<<20) + for i := 0; i < 100 && r.err == nil; i++ { + r.event("relay_metrics", map[string]any{"prometheus": body}) + } + if r.err == nil || r.err.Error() != "output_budget" { + t.Fatalf("unbounded encoded metrics were retained: bytes=%d err=%v", out.bytes, r.err) + } + if out.bytes > 512<<20 { + t.Fatalf("encoded output exceeded the cap before rejection: %d", out.bytes) + } + before := out.bytes + r.event("summary", nil) + if out.bytes != before { + t.Fatal("failed recorder kept writing after exhausting its output budget") + } +} + +type countOutput struct{ bytes int64 } + +func (w *countOutput) Write(p []byte) (int, error) { + w.bytes += int64(len(p)) + return len(p), nil +} + +func TestEncodedOutputBudgetChecksExactLimitBeforeWrite(t *testing.T) { + for _, limit := range []int64{8, 9, 10} { + var out bytes.Buffer + w := &outputBudgetWriter{out: &out, remaining: limit} + encoder := json.NewEncoder(w) + err := encoder.Encode("\x00") // "\u0000" plus LF is exactly nine bytes. + if limit == 8 { + if !errors.Is(err, errOutputBudget) || out.Len() != 0 || w.remaining != 8 { + t.Fatalf("oversized encoded event wrote a prefix: limit=%d err=%v output=%q remaining=%d", limit, err, out.String(), w.remaining) + } + continue + } + if err != nil || out.String() != "\"\\u0000\"\n" || w.remaining != limit-9 { + t.Fatalf("complete event at the limit was rejected/miscounted: limit=%d err=%v output=%q remaining=%d", limit, err, out.String(), w.remaining) + } + if err := encoder.Encode("x"); !errors.Is(err, errOutputBudget) || out.Len() != 9 || w.remaining != limit-9 { + t.Fatalf("cumulative overflow wrote a partial second event: limit=%d err=%v output=%q remaining=%d", limit, err, out.String(), w.remaining) + } + } +} + +func TestOutputBudgetStopsLiveWorkAndCleansUp(t *testing.T) { + c := localRelay(t) + c.DurationSeconds = 20 + c.outputLimit = 64 << 10 + var samples atomic.Int32 + metrics := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if samples.Add(1) == 1 { + _, _ = io.WriteString(w, "test_metric 1\n") + return + } + _, _ = io.WriteString(w, strings.Repeat("\x00", 32<<10)) + })) + t.Cleanup(metrics.Close) + c.Manifest.Targets[0].TelemetryURL = metrics.URL + "/metrics" + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if !errors.Is(err, errOutputBudget) || s.StopReason != "output_budget" || s.DeliveryComplete || !s.CleanupComplete { + t.Fatalf("output exhaustion did not stop and invalidate the phase: err=%v summary=%+v", err, s) + } + if s.EstablishedGuests != 1 || s.ActiveFrames == 0 || s.DurationSeconds >= 10 { + t.Fatalf("fixture did not cancel a live workload promptly: %+v", s) + } + if int64(out.Len()) > c.outputLimit { + t.Fatalf("output exceeded its budget: %d", out.Len()) + } + decoder := json.NewDecoder(&out) + for decoder.More() { + var event map[string]any + if err := decoder.Decode(&event); err != nil { + t.Fatalf("quota retained a partial JSON event: %v", err) + } + } +} + +func TestOutputBudgetDuringStartupStillCleansUp(t *testing.T) { + c := localRelay(t) + c.outputLimit = 1 + var out bytes.Buffer + s, err := Run(t.Context(), c, &out) + if !errors.Is(err, errOutputBudget) || s.StopReason != "output_budget" || s.DeliveryComplete || !s.CleanupComplete || out.Len() != 0 { + t.Fatalf("startup budget failure lost its cause or cleanup: err=%v summary=%+v output=%d", err, s, out.Len()) + } +} + +func TestFinalOutputFailureInvalidatesDelivery(t *testing.T) { + c := localRelay(t) + c.DurationSeconds = 1 + s, err := Run(t.Context(), c, rejectSummaryOutput{}) + if !errors.Is(err, errOutputBudget) || s.StopReason != "output_budget" || s.DeliveryComplete || !s.CleanupComplete || s.EstablishedGuests != 1 { + t.Fatalf("final output failure reported successful delivery: err=%v summary=%+v", err, s) + } +} + +type rejectSummaryOutput struct{} + +func (rejectSummaryOutput) Write(p []byte) (int, error) { + var event struct{ Event string } + if err := json.Unmarshal(p, &event); err != nil { + return 0, err + } + if event.Event == "summary" { + return 0, errOutputBudget + } + return len(p), nil +} diff --git a/internal/capacity/runner.go b/internal/capacity/runner.go index 366fd2132..3792e8545 100644 --- a/internal/capacity/runner.go +++ b/internal/capacity/runner.go @@ -468,6 +468,7 @@ func Run(ctx context.Context, c Config, out io.Writer) (result Summary, err erro bounded, stop := context.WithDeadline(ctx, deadline) defer stop() runctx, cancel := context.WithCancelCause(bounded) + rec.fail = cancel r := &runState{lastTelemetry: map[string]time.Time{}, cfg: c, guard: newGuard(c), rec: rec, ctx: runctx, cancel: cancel} r.helper = c.HelperExecutable if r.helper == "" { @@ -502,8 +503,12 @@ func Run(ctx context.Context, c Config, out io.Writer) (result Summary, err erro result = rec.summary rec.mu.Unlock() rec.event("summary", map[string]any{"summary": result}) - if rec.err != nil && err == nil { - err = rec.err + if rec.err != nil { + err = errors.Join(err, rec.err) + result.DeliveryComplete = false + if result.StopReason != "cleanup_failure" { + result.StopReason = rec.err.Error() + } } }() rec.event("phase_start", map[string]any{"config": safeConfig(c), "input_frame_bytes": 128, "host_output_frame_bytes": 1024, "framing_header_bytes": 44, "framing_newline_bytes": 1, "wire_overhead": "not measured", "duration_includes_admission": false}) diff --git a/internal/capacity/transfer.go b/internal/capacity/transfer.go index 1ba3bcd2f..3851f880b 100644 --- a/internal/capacity/transfer.go +++ b/internal/capacity/transfer.go @@ -197,15 +197,22 @@ func (r *runState) restoreGuest(q *guest, began time.Time) error { for tries := 0; tries < 3; tries++ { last = q.connect(r.ctx, r) queue += q.queueElapsed + if last == nil && r.cfg.beforeRestoreCommit != nil { + r.cfg.beforeRestoreCommit() + } q.mu.Lock() superseded := disruption != q.disruption + if !superseded && last == nil { + q.inputStart = q.probeSeq + q.inputStart[0] = q.seq + q.expected = false + } q.mu.Unlock() if superseded { q.close() return nil } // a later actual host cut owns repair. if last == nil { - q.setExpected(false) r.rec.mu.Lock() r.rec.summary.Reconnects++ r.rec.mu.Unlock() diff --git a/internal/capacity/workload.go b/internal/capacity/workload.go index 718a7aff8..cbb4d4dfb 100644 --- a/internal/capacity/workload.go +++ b/internal/capacity/workload.go @@ -200,7 +200,40 @@ func scanFrames(in io.Reader, consume func(observation) error, helperReady func( } return io.EOF } -func (r *runState) observe(tracker sequenceTracker, epochs map[string]int, o observation) error { + +type inputEpoch struct { + generation int + floor int +} + +// A live receiver requires every input allocated after this snapshot. Earlier +// input can still be in flight or replayed, and is outside its complete range. +// Snapshot allocation counters, not timestamps from independently queued SSH +// channels. The steady workload starts only after every initial live barrier. +func (r *runState) baselineInputs(tracker sequenceTracker, epochs map[string]inputEpoch, hostID int) { + r.clientsMu.RLock() + defer r.clientsMu.RUnlock() + for _, source := range r.guests { + if source.host.id != hostID { + continue + } + source.mu.Lock() + if !source.expected { + for ch := 0; ch < r.cfg.ChannelsPerGuest; ch++ { + next := source.probeSeq[ch] + if ch == 0 { + next = source.seq + } + key := "I:" + strconv.Itoa(source.id*4+ch) + tracker[key] = next - 1 + epochs[key] = inputEpoch{source.generation, next - 1} + } + } + source.mu.Unlock() + } +} + +func (r *runState) observe(tracker sequenceTracker, epochs map[string]inputEpoch, o observation) error { if o.kind == "I" { r.clientsMu.RLock() if o.id < 0 || o.id/4 >= len(r.guests) { @@ -210,14 +243,21 @@ func (r *runState) observe(tracker sequenceTracker, epochs map[string]int, o obs peer := r.guests[o.id/4] r.clientsMu.RUnlock() peer.mu.Lock() - generation, expected := peer.generation, peer.expected + generation, expected, first := peer.generation, peer.expected, peer.inputStart[o.id%4] peer.mu.Unlock() key := "I:" + strconv.Itoa(o.id) - if old, ok := epochs[key]; expected || !ok || old != generation { + if expected { delete(tracker, key) + delete(epochs, key) + return nil } - epochs[key] = generation - if expected { + epoch, ok := epochs[key] + if !ok || epoch.generation != generation { + epoch = inputEpoch{generation, first - 1} + epochs[key] = epoch + tracker[key] = epoch.floor + } + if o.seq <= epoch.floor { return nil } } diff --git a/internal/capacity/workload_test.go b/internal/capacity/workload_test.go index cc92799b0..c84e2fdfc 100644 --- a/internal/capacity/workload_test.go +++ b/internal/capacity/workload_test.go @@ -27,10 +27,10 @@ func TestFixedFramesAndSequenceGaps(t *testing.T) { } } func TestIntentionalReconnectResetsOnlyAffectedSource(t *testing.T) { - peer := &guest{generation: 1} + peer := &guest{generation: 1, inputStart: [4]int{4}} r := &runState{guests: []*guest{peer}} tracker := sequenceTracker{"I:0": 1} - epochs := map[string]int{"I:0": 0} + epochs := map[string]inputEpoch{"I:0": {generation: 0, floor: 0}} if err := r.observe(tracker, epochs, observation{kind: "I", id: 0, seq: 4}); err != nil { t.Fatalf("intentional reconnect gap was treated as survivor loss: %v", err) } From 0efeb680049ba23510b03f0d7b1cc5a136a47b9c Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 22:37:48 -0700 Subject: [PATCH 13/16] Confine capacity SFTP and pin helper terminal geometry --- host/host.go | 65 +++++- host/internal/server.go | 25 +- host/internal/sftp.go | 42 +++- host/internal/sftp_root.go | 219 ++++++++++++++++++ host/internal/sftp_root_test.go | 77 ++++++ host/internal/sftp_root_unix_test.go | 117 ++++++++++ host/sftp_root_test.go | 34 +++ internal/capacity/README.md | 23 +- internal/capacity/fault_helper_test.go | 3 + internal/capacity/geometry_test.go | 199 ++++++++++++++++ internal/capacity/review_regression_test.go | 1 + internal/capacity/runner.go | 35 ++- internal/capacity/sftp_confinement_test.go | 200 ++++++++++++++++ internal/capacity/sftp_namespace_unix_test.go | 122 ++++++++++ internal/capacity/sftp_symlink_test.go | 68 ++++++ internal/capacity/transfer.go | 2 +- internal/capacity/whole_review_test.go | 2 + 17 files changed, 1217 insertions(+), 17 deletions(-) create mode 100644 host/internal/sftp_root.go create mode 100644 host/internal/sftp_root_test.go create mode 100644 host/internal/sftp_root_unix_test.go create mode 100644 host/sftp_root_test.go create mode 100644 internal/capacity/geometry_test.go create mode 100644 internal/capacity/sftp_confinement_test.go create mode 100644 internal/capacity/sftp_namespace_unix_test.go create mode 100644 internal/capacity/sftp_symlink_test.go diff --git a/host/host.go b/host/host.go index 4f27f3170..a861eaf35 100644 --- a/host/host.go +++ b/host/host.go @@ -325,6 +325,14 @@ type Host struct { // SFTP configuration SFTPDisabled bool // Disable SFTP subsystem entirely (--no-sftp) + // SFTPRoot confines SFTP content and metadata to this directory handle. + // Nil preserves the ordinary filesystem/consent behavior. The caller owns + // the root and must keep it open until Run has returned AND WaitSFTP succeeds. + // Rooted hosts are single-use. Configure this before Run; do not change it. + SFTPRoot *os.Root + sftpMu sync.Mutex + sftpRunDone chan struct{} + sftpServer *internal.Server // PermissionChecker supplies custom consent policy. Nil denies ungranted operations. PermissionChecker consent.PermissionChecker // TerminalConsent selects consent on the primary local terminal; mutually exclusive with PermissionChecker. @@ -600,13 +608,56 @@ func (c *Host) claimTimeout() time.Duration { return ClaimTimeout } +// WaitSFTP certifies that a rooted host has returned from Run, closed SFTP +// admission, and joined every rooted request worker and open handle. Concurrent +// waits are safe. Before Run it returns an error; failed startup with no server +// is already quiescent. Nil-root hosts have nothing to wait for. +// +// On timeout the caller must retain the root and its directory and may wait +// again. Cancellation cannot force a custom permission checker or blocked OS +// filesystem call to return, so Run returning alone is not this guarantee. +func (c *Host) WaitSFTP(ctx context.Context) error { + if c.SFTPRoot == nil { + return nil + } + c.sftpMu.Lock() + done := c.sftpRunDone + c.sftpMu.Unlock() + if done == nil { + return errors.New("rooted host has not run") + } + select { + case <-done: + case <-ctx.Done(): + return ctx.Err() + } + c.sftpMu.Lock() + server := c.sftpServer + c.sftpMu.Unlock() + if server == nil { + return nil + } + return server.WaitSFTP(ctx) +} + // Run hosts one session and returns when it ends. // -// It may be called again afterwards: everything Run claims, it gives back -// before returning, including the two fields it fills in on the Host itself. -// A caller that supplied AdminSocketFile keeps it across runs, since managing -// the path is what supplying it means. +// Without SFTPRoot it may be called again afterwards: everything Run claims, +// it gives back before returning, including the two fields it fills in on the +// Host itself. A caller that supplied AdminSocketFile keeps it across runs, +// since managing the path is what supplying it means. Rooted hosts are single- +// use and require WaitSFTP before the caller closes their root. func (c *Host) Run(ctx context.Context) (runErr error) { + if c.SFTPRoot != nil { + c.sftpMu.Lock() + if c.sftpRunDone != nil { + c.sftpMu.Unlock() + return errors.New("a host with SFTPRoot cannot be reused") + } + c.sftpRunDone = make(chan struct{}) + c.sftpMu.Unlock() + defer close(c.sftpRunDone) + } // An empty command is no command, as the guest door reads one. if len(c.LoginCommand) > 0 && len(c.ForceCommand) > 0 { return errors.New("LoginCommand and ForceCommand are mutually exclusive") @@ -1246,6 +1297,7 @@ func (c *Host) Run(ctx context.Context) (runErr error) { InitialClientTimeout: c.InitialClientTimeout, StopGrace: c.StopGrace, SFTPDisabled: c.SFTPDisabled, + SFTPRoot: c.SFTPRoot, PermissionChecker: c.PermissionChecker, TerminalConsent: c.TerminalConsent, SFTPAutoAccept: c.AutoAccept || c.SFTPAutoAccept, @@ -1263,6 +1315,11 @@ func (c *Host) Run(ctx context.Context) (runErr error) { logger.Warn("reverse tunnel stopped serving guests; command continues", "error", err) }, } + if c.SFTPRoot != nil { + c.sftpMu.Lock() + c.sftpServer = &sshServer + c.sftpMu.Unlock() + } g.Add(func() error { return sshServer.ServeWithContext(ctx, guestDoor, attachLn) }, func(err error) { diff --git a/host/internal/server.go b/host/internal/server.go index a24d520c2..08afefb8f 100644 --- a/host/internal/server.go +++ b/host/internal/server.go @@ -7,6 +7,7 @@ import ( "io" "log/slog" "net" + "os" "path/filepath" "slices" "strings" @@ -131,6 +132,9 @@ type Server struct { // SFTP configuration SFTPDisabled bool // Disable SFTP subsystem entirely + // SFTPRoot is caller-owned; WaitSFTP must succeed before it is closed. + SFTPRoot *os.Root + sftpWorkers rootedSFTPWorkers // PermissionChecker supplies custom consent policy. Nil denies ungranted operations. PermissionChecker consent.PermissionChecker // TerminalConsent selects consent on the primary local terminal; mutually exclusive with PermissionChecker. @@ -219,6 +223,12 @@ func logReplayFallback(logger *slog.Logger) func(error) { // guest, and — when host is not nil — the host door on host, for clients that // are already on this machine. func (s *Server) ServeWithContext(ctx context.Context, guest, host net.Listener) error { + if s.SFTPRoot != nil { + defer s.sftpWorkers.stop() + if !rootedSFTPSupported() { + return errors.New("rooted SFTP is unsupported on this platform") + } + } if s.TerminalConsent && s.PermissionChecker != nil { return errors.New("TerminalConsent and PermissionChecker are mutually exclusive") } @@ -375,9 +385,13 @@ func (s *Server) ServeWithContext(ctx context.Context, guest, host net.Listener) kind: kindGuest, cmdDone: cmdDone, commandResult: cmd.Result, + sftpRoot: s.SFTPRoot, onPrivateCommandExit: s.onPrivateCommandExit, } + if s.SFTPRoot != nil { + sh.sftpWorkers = &s.sftpWorkers + } ss := []gssh.Signer{s.HostKey} @@ -390,9 +404,9 @@ func (s *Server) ServeWithContext(ctx context.Context, guest, host net.Listener) // Set up subsystem handlers (SFTP). The door admits a subsystem with // the session's end, as it does a shell (see admitGuest), and SFTP - // hands that admission back as it starts: an SFTP session in progress - // is not waited for, and is cut off when the session ends, as it - // always was. + // hands that admission back as it starts. Rooted sessions separately + // register their workers for WaitSFTP; nil-root sessions retain their + // existing shutdown behavior. var subsystemHandlers map[string]gssh.SubsystemHandler if !s.SFTPDisabled { subsystemHandlers = map[string]gssh.SubsystemHandler{ @@ -485,6 +499,9 @@ func (s *Server) ServeWithContext(ctx context.Context, guest, host net.Listener) // door's listener's close. A command-led exit has already closed // cmdDone by the time this runs. releaseSessions(cmdDone, timing.outputDrain+guestFlushTimeout, cancel) + if s.SFTPRoot != nil { + s.sftpWorkers.stop() + } // That cancel hangs up every forced command, and each one's guest is // sent what it says on the way out and a reset once it has gone. @@ -1087,6 +1104,8 @@ type sessionHandler struct { // SFTP configuration permissionChecker consent.PermissionChecker sftpAutoAccept bool + sftpRoot *os.Root + sftpWorkers *rootedSFTPWorkers // subsystems are the guest door's subsystem handlers, by name, so that // admitGuest refuses one the door doesn't serve. Nil in a test's handler. diff --git a/host/internal/sftp.go b/host/internal/sftp.go index e9f184be2..7b794924d 100644 --- a/host/internal/sftp.go +++ b/host/internal/sftp.go @@ -39,10 +39,19 @@ type SFTPSession struct { logger *slog.Logger ctx context.Context autoAccept bool + root *os.Root } // HandleSFTP handles SFTP subsystem requests func (h *sessionHandler) HandleSFTP(sess gssh.Session) { + if h.sftpRoot != nil { + finished, ok := h.sftpWorkers.admit(sess) + if !ok { + _ = sess.Exit(1) + return + } + defer finished() + } sessionID := sess.Context().Value(gssh.ContextKeySessionID).(string) // Get client info for terminal approval. The fingerprint is the guest's @@ -80,6 +89,7 @@ func (h *sessionHandler) HandleSFTP(sess gssh.Session) { permissionChecker: h.permissionChecker, clientInfo: clientInfo, logger: h.logger, + root: h.sftpRoot, } handlers := sftp.Handlers{ @@ -89,10 +99,18 @@ func (h *sessionHandler) HandleSFTP(sess gssh.Session) { FileList: &sftpFileLister{session: sftpSession, logger: h.logger}, } - // Get user's home directory for SFTP start directory - userHome, err := os.UserHomeDir() + // Rooted sessions resolve relative requests under their declared root; + // ordinary sessions retain the user's home as their starting directory. + var userHome string + var err error + if h.sftpRoot != nil { + userHome, err = filepath.Abs(h.sftpRoot.Name()) + handlers.FileList = &rootedSFTPFileLister{sftpFileLister{session: sftpSession, logger: h.logger}} + } else { + userHome, err = os.UserHomeDir() + } if err != nil { - h.logger.Error("failed to get user home directory", "error", err) + h.logger.Error("failed to get SFTP start directory", "error", err) _ = sess.Exit(1) return } @@ -249,6 +267,9 @@ func sftpOperationError(ctx context.Context, r *sftp.Request, s *SFTPSession) er // paths already resolve to the home directory, so "notes.txt" covers what // "~/notes.txt" was reaching for. func (s *SFTPSession) resolvePath(reqPath string) (string, error) { + if s.root != nil { + return rootedSFTPPath(s.root, reqPath) + } if runtime.GOOS == "windows" { reqPath = undoubleDriveLetter(reqPath) } @@ -343,6 +364,9 @@ func (h *sftpFileReader) Fileread(r *sftp.Request) (io.ReaderAt, error) { if err := sftpOperationError(ctx, r, h.session); err != nil { return nil, err } + if h.session.root != nil { + return h.session.root.Open(path) + } return os.Open(path) } @@ -401,6 +425,9 @@ func (h *sftpFileWriter) Filewrite(r *sftp.Request) (io.WriterAt, error) { if err := sftpOperationError(ctx, r, h.session); err != nil { return nil, err } + if h.session.root != nil { + return h.session.root.OpenFile(path, osFlags, 0600) + } return os.OpenFile(path, osFlags, 0644) } @@ -411,6 +438,9 @@ type sftpFileCmder struct { } func (h *sftpFileCmder) Filecmd(r *sftp.Request) error { + if h.session.root != nil && r.Method != "Remove" { + return sftp.ErrSSHFxOpUnsupported + } arrived := time.Now() requestID := nextSFTPRequestID() ctx, cancel := h.session.operationContext(r.Context()) @@ -467,6 +497,9 @@ func (h *sftpFileCmder) Filecmd(r *sftp.Request) error { // Execute the operation switch r.Method { case "Remove": + if h.session.root != nil { + return h.session.root.Remove(path) + } return os.Remove(path) case "Mkdir": return os.Mkdir(path, 0755) @@ -559,6 +592,9 @@ type sftpFileLister struct { } func (h *sftpFileLister) Filelist(r *sftp.Request) (sftp.ListerAt, error) { + if h.session.root != nil { + return h.rootedList(r) + } path, err := h.session.resolvePath(r.Filepath) if err != nil { return nil, sftp.ErrSSHFxPermissionDenied diff --git a/host/internal/sftp_root.go b/host/internal/sftp_root.go new file mode 100644 index 000000000..6ffece803 --- /dev/null +++ b/host/internal/sftp_root.go @@ -0,0 +1,219 @@ +package internal + +import ( + "context" + "os" + "path" + "path/filepath" + "runtime" + "strings" + "sync" + + gssh "charm.land/ssh" + "github.com/pkg/sftp" + "golang.org/x/crypto/ssh" +) + +// Only descriptor-backed Root implementations are supported. In particular, +// js and plan9 cannot keep this boundary across concurrent namespace changes. +func rootedSFTPSupported() bool { + switch runtime.GOOS { + case "aix", "android", "darwin", "dragonfly", "freebsd", "illumos", "ios", "linux", "netbsd", "openbsd", "solaris", "windows": + return true + default: + return false + } +} + +// This conversion selects a name, never authorizes an ambient filesystem +// operation. Every subsequent access is through Root, even after a rename or +// symlink replacement. Absolute names must lie under the root's declared name. +func rootedSFTPPath(root *os.Root, name string) (string, error) { + if !rootedSFTPSupported() { + return "", sftp.ErrSSHFxOpUnsupported + } + base, err := filepath.Abs(root.Name()) + if err != nil { + return "", err + } + base = filepath.ToSlash(base) + if runtime.GOOS == "windows" { + // pkg/sftp resolves a native drive path as relative to its POSIX + // start directory. Restore that spelling before checking our root. + name = undoubleDriveLetter(name) + base = "/" + base + } + if strings.ContainsAny(name, "\\\x00") { + return "", sftp.ErrSSHFxPermissionDenied + } + base = path.Clean(base) + prefix := strings.TrimSuffix(base, "/") + "/" + name = path.Clean(name) + if path.IsAbs(name) { + switch { + case name == base: + name = "." + case strings.HasPrefix(name, prefix): + name = strings.TrimPrefix(name, prefix) + default: + return "", sftp.ErrSSHFxPermissionDenied + } + } + if !filepath.IsLocal(name) || strings.Contains(name, ":") { + return "", sftp.ErrSSHFxPermissionDenied + } + for _, component := range strings.Split(name, "/") { + // Reserve Windows device spellings on every platform; a runner's + // remotely supplied paths must never become devices when moved to Windows. + stem, _, _ := strings.Cut(strings.ToUpper(strings.TrimRight(component, " .")), ".") + switch stem { + case "CON", "PRN", "AUX", "NUL", "CONIN$", "CONOUT$": + return "", sftp.ErrSSHFxPermissionDenied + } + if strings.HasPrefix(stem, "COM") || strings.HasPrefix(stem, "LPT") { + suffix := []rune(stem[3:]) + if len(suffix) == 1 && strings.ContainsRune("123456789¹²³", suffix[0]) { + return "", sftp.ErrSSHFxPermissionDenied + } + } + } + return filepath.FromSlash(name), nil +} + +func (h *sftpFileLister) rootedList(r *sftp.Request) (sftp.ListerAt, error) { + if err := sftpOperationError(r.Context(), r, h.session); err != nil { + return nil, err + } + name, err := h.session.resolvePath(r.Filepath) + if err != nil { + return nil, err + } + switch r.Method { + case "List": + dir, err := h.session.root.Open(name) + if err != nil { + return nil, err + } + defer func() { _ = dir.Close() }() + // Root-opened directory descriptors keep entry metadata descriptor- + // relative too; never use ambient ReadDir or lazy pathname Info calls. + infos, err := dir.Readdir(-1) + return listerat(infos), err + case "Stat": + info, err := h.session.root.Stat(name) + if err != nil { + return nil, err + } + return listerat{info}, nil + case "Lstat": + info, err := h.session.root.Lstat(name) + if err != nil { + return nil, err + } + return listerat{info}, nil + default: + return nil, sftp.ErrSSHFxOpUnsupported + } +} + +type rootedSFTPFileLister struct{ sftpFileLister } + +func (h *rootedSFTPFileLister) Lstat(r *sftp.Request) (sftp.ListerAt, error) { + return h.Filelist(r) +} + +// RealPath only returns the canonical protocol spelling. It does not resolve +// symlinks, and cannot disclose an outside path or bypass the access handlers. +func (h *rootedSFTPFileLister) RealPath(name string) (string, error) { + rel, err := h.session.resolvePath(name) + if err != nil { + return "", err + } + base, err := filepath.Abs(h.session.root.Name()) + if err != nil { + return "", err + } + base = filepath.ToSlash(base) + if runtime.GOOS == "windows" { + base = "/" + base + } + return path.Join(base, filepath.ToSlash(rel)), nil +} + +// Closing a channel alone can wait for an unresponsive peer. At host shutdown, +// close the rooted sessions' SSH transports to release both reads and writes. +// A closed admission gate and the final Serve return are both needed: Serve +// joins pkg/sftp workers and closes their handles before the registry releases. +type rootedSFTPWorkers struct { + mu sync.Mutex + stopped bool + done chan struct{} + active map[chan struct{}]ssh.Conn +} + +func (w *rootedSFTPWorkers) admit(sess gssh.Session) (func(), bool) { + w.mu.Lock() + defer w.mu.Unlock() + conn, ok := sess.Context().Value(gssh.ContextKeyConn).(ssh.Conn) + if w.stopped || !ok { + return nil, false + } + if w.done == nil { + w.done = make(chan struct{}) + } + if w.active == nil { + w.active = make(map[chan struct{}]ssh.Conn) + } + id := make(chan struct{}) + w.active[id] = conn + return func() { + w.mu.Lock() + defer w.mu.Unlock() + delete(w.active, id) + if w.stopped && len(w.active) == 0 { + close(w.done) + } + }, true +} + +func (w *rootedSFTPWorkers) stop() { + w.mu.Lock() + if w.stopped { + w.mu.Unlock() + return + } + w.stopped = true + if w.done == nil { + w.done = make(chan struct{}) + } + connections := make([]ssh.Conn, 0, len(w.active)) + for _, conn := range w.active { + connections = append(connections, conn) + } + if len(w.active) == 0 { + close(w.done) + } + w.mu.Unlock() + for _, conn := range connections { + _ = conn.Close() + } +} + +func (s *Server) WaitSFTP(ctx context.Context) error { + if s.SFTPRoot == nil { + return nil + } + w := &s.sftpWorkers + w.mu.Lock() + if w.done == nil { + w.done = make(chan struct{}) + } + done := w.done + w.mu.Unlock() + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} diff --git a/host/internal/sftp_root_test.go b/host/internal/sftp_root_test.go new file mode 100644 index 000000000..0d3096040 --- /dev/null +++ b/host/internal/sftp_root_test.go @@ -0,0 +1,77 @@ +package internal + +import ( + "context" + "io" + "os" + "path/filepath" + "testing" + + "github.com/owenthereal/upterm/host/consent" + "github.com/pkg/sftp" + "github.com/stretchr/testify/require" +) + +func TestRootedSFTPRetainsConsentAndReadOnly(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + require.NoError(t, root.WriteFile("file", []byte("inside"), 0600)) + for _, policy := range []string{"no checker", "denied", "readonly", "auto", "allowed"} { + t.Run(policy, func(t *testing.T) { + s := &SFTPSession{root: root, ctx: t.Context(), logger: discardLogger()} + called := false + switch policy { + case "readonly": + s.readOnly, s.autoAccept = true, true + case "auto": + s.autoAccept = true + } + if policy != "no checker" { + s.permissionChecker = sftpConsentChecker{check: func(_ context.Context, _ consent.Operation, _ consent.ClientInfo, paths ...string) (consent.Decision, error) { + called = true + require.Equal(t, []string{"file"}, paths) + if policy == "allowed" { + return consent.AllowConnection, nil + } + return consent.Deny, nil + }} + } + reader, err := (&sftpFileReader{session: s, logger: s.logger}).Fileread(sftp.NewRequest("Get", filepath.Join(root.Name(), "file"))) + if policy == "no checker" || policy == "denied" { + require.Error(t, err) + } else { + require.NoError(t, err) + data := make([]byte, 6) + _, err = reader.ReadAt(data, 0) + require.NoError(t, err) + require.Equal(t, "inside", string(data)) + require.NoError(t, reader.(io.Closer).Close()) + } + if policy == "readonly" { + _, err = (&sftpFileWriter{session: s, logger: s.logger}).Filewrite(sftp.NewRequest("Put", "file")) + require.Error(t, err) + require.Error(t, (&sftpFileCmder{session: s, logger: s.logger}).Filecmd(sftp.NewRequest("Remove", "file"))) + } + require.Equal(t, policy == "denied" || policy == "allowed", called) + }) + } + closedRoot, err := os.OpenRoot(root.Name()) + require.NoError(t, err) + require.NoError(t, closedRoot.Close()) + s := &SFTPSession{root: closedRoot, ctx: t.Context(), autoAccept: true, logger: discardLogger()} + _, err = (&sftpFileReader{session: s, logger: s.logger}).Fileread(sftp.NewRequest("Get", filepath.Join(root.Name(), "file"))) + require.Error(t, err, "a closed root must never fall back to ambient access") +} + +func TestRootedSFTPRejectsForeignVolumesAndDevices(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + for _, name := range []string{"..", "../escape", "/outside", "C:/outside", "C:relative", "//server/share/file", `\\?\C:\file`, "file:stream", "NUL", "NUL.txt", "COM1", "COM¹", "LPT9", "CONOUT$", "dir/PRN", "bad\x00name"} { + t.Run(name, func(t *testing.T) { + _, err := rootedSFTPPath(root, name) + require.Error(t, err) + }) + } +} diff --git a/host/internal/sftp_root_unix_test.go b/host/internal/sftp_root_unix_test.go new file mode 100644 index 000000000..0c4279f22 --- /dev/null +++ b/host/internal/sftp_root_unix_test.go @@ -0,0 +1,117 @@ +//go:build !windows + +package internal + +import ( + "context" + "crypto/ed25519" + "crypto/rand" + "os" + "path/filepath" + "testing" + "time" + + "github.com/owenthereal/upterm/host/consent" + "github.com/owenthereal/upterm/server" + "github.com/pkg/sftp" + "github.com/stretchr/testify/require" + "golang.org/x/crypto/ssh" +) + +func TestRootedSFTPWaitIncludesBlockedWorkers(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + entered, canceled, release := make(chan struct{}), make(chan struct{}), make(chan struct{}) + defer close(release) + srv := &Server{Command: readsALine("READY", 0), SFTPRoot: root, PermissionChecker: sftpConsentChecker{check: func(ctx context.Context, _ consent.Operation, _ consent.ClientInfo, _ ...string) (consent.Decision, error) { + close(entered) + <-ctx.Done() + close(canceled) + <-release // deliberately prove a checker that ignores cancellation is not joined + return consent.AllowConnection, nil + }}} + h := startHost(t, srv) + client, err := sftp.NewClient(h.guestClient(t)) + require.NoError(t, err) + defer func() { _ = client.Close() }() + opened := make(chan error, 1) + go func() { + f, err := client.Create("pending") + if f != nil { + _ = f.Close() + } + opened <- err + }() + select { + case <-entered: + case <-time.After(harnessTimeout): + t.Fatal("permission worker never started") + } + h.stop(t) + select { + case <-canceled: + case <-time.After(harnessTimeout): + t.Fatal("rooted request was not canceled") + } + ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond) + defer cancel() + require.ErrorIs(t, srv.WaitSFTP(ctx), context.DeadlineExceeded, "returning from Serve is not a worker join") + _, err = root.Stat(".") + require.NoError(t, err, "server closed the caller-owned root") + // A separate sender releases the checker without closing release twice. + release <- struct{}{} + joined, cancelJoin := context.WithTimeout(t.Context(), harnessTimeout) + defer cancelJoin() + require.NoError(t, srv.WaitSFTP(joined)) + require.Error(t, <-opened) + _, err = root.Stat("pending") + require.True(t, os.IsNotExist(err), "a late permission grant created a file after cancellation") +} + +func TestRootedSFTPTrustsRelayIdentityButConfinesItsAccess(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + public, _, err := ed25519.GenerateKey(rand.Reader) + require.NoError(t, err) + guestKey, err := ssh.NewPublicKey(public) + require.NoError(t, err) + h := startHost(t, &Server{Command: readsALine("READY", 0), SFTPRoot: root, SFTPAutoAccept: true, AuthorizedKeys: []ssh.PublicKey{guestKey}}) + // The relay can attest an allowed guest key without possessing its private + // key. This is the threat the filesystem boundary must still contain. + h.guestSigner, err = (&server.UserCertSigner{User: "guest", SessionID: t.Name(), AuthRequest: &server.AuthRequest{AuthorizedKey: ssh.MarshalAuthorizedKey(guestKey)}}).SignCert(h.relaySigner) + require.NoError(t, err) + client, err := sftp.NewClient(h.guestClient(t)) + require.NoError(t, err) + defer func() { _ = client.Close() }() + outside := filepath.Join(t.TempDir(), "private") + require.NoError(t, os.WriteFile(outside, []byte("private"), 0600)) + _, err = client.Open(outside) + require.Error(t, err) + _, err = client.Stat(outside) + require.Error(t, err) + h.stop(t) + require.NoError(t, h.srv.WaitSFTP(t.Context())) +} + +func TestNilRootSFTPKeepsAbsoluteAccessAndConsent(t *testing.T) { + file := filepath.Join(t.TempDir(), "existing") + require.NoError(t, os.WriteFile(file, []byte("ordinary"), 0600)) + for _, accepted := range []bool{false, true} { + h := startHost(t, &Server{Command: readsALine("READY", 0), SFTPAutoAccept: accepted}) + client, err := sftp.NewClient(h.guestClient(t)) + require.NoError(t, err) + f, err := client.Open(file) + if accepted { + require.NoError(t, err) + require.NoError(t, f.Close()) + } else { + require.Error(t, err) + } + _, err = client.Stat(file) + require.NoError(t, err, "nil root preserves existing metadata semantics") + require.NoError(t, client.Close()) + h.stop(t) + } +} diff --git a/host/sftp_root_test.go b/host/sftp_root_test.go new file mode 100644 index 000000000..8ee90d0da --- /dev/null +++ b/host/sftp_root_test.go @@ -0,0 +1,34 @@ +package host + +import ( + "context" + "os" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestRootedHostWaitContract(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + h := &Host{SFTPRoot: root, LoginCommand: []string{"a"}, ForceCommand: []string{"b"}} + require.Error(t, h.WaitSFTP(t.Context()), "pre-Run wait cannot certify a future server") + require.Error(t, h.Run(t.Context()), "invalid startup must fail before serving") + var wg sync.WaitGroup + results := make(chan error, 8) + for range 8 { + wg.Add(1) + go func() { defer wg.Done(); results <- h.WaitSFTP(t.Context()) }() + } + wg.Wait() + close(results) + for err := range results { + require.NoError(t, err) + } + _, err = root.Stat(".") + require.NoError(t, err, "host closed the caller's root") + require.ErrorContains(t, h.Run(t.Context()), "cannot be reused") + require.NoError(t, (&Host{}).WaitSFTP(context.Background())) +} diff --git a/internal/capacity/README.md b/internal/capacity/README.md index cb8326f73..1a6c895c9 100644 --- a/internal/capacity/README.md +++ b/internal/capacity/README.md @@ -141,8 +141,10 @@ this is invalid generator-limited evidence, never a relay capacity ceiling. ## Offered data and measurement -The hosted command is this binary's own `helper` subcommand. It sets its PTY to -raw mode, answers a distinct 128-byte readiness challenge, echoes deterministic +The hosted command is this binary's own `helper` subcommand. The host creates +and pins its PTY at 2048 columns by 100 rows, matching guest requests and +preventing later joins/resizes from repainting old Windows terminal output. +The helper sets raw mode, answers a distinct 128-byte readiness challenge, echoes deterministic data and never interprets input as commands. Active guests send exactly 128 bytes once per second, beginning at steady observation start. Active helpers emit a 1 KiB frame every @@ -222,9 +224,17 @@ runtime. These byte figures are per phase; the operator must also enforce the approved aggregate budget across attempts. `transfer_concurrency` can be zero, four or eight. Each transfer concurrently -uploads one deterministic 16 MiB file into its host's temporary directory, +uploads one deterministic 16 MiB file into its host's private transfer directory, then downloads it and verifies both length and SHA-256. Source/download data are streamed with bounded memory; at most eight temporary 16 MiB files exist. +Zero-transfer phases disable SFTP entirely. Transfer-enabled hosts confine all +content and metadata access through a caller-owned `os.Root`, separate from +admin sockets and other control files. Only open/create, remove, stat and list +are supported; links, rename, mkdir/rmdir, readlink, attribute changes and +other operations are refused. Remove can also remove an empty in-root directory. +This protects against path traversal and symlink replacement, not privileged +mount/device changes or hostile preexisting hardlinks. Transfer roots start +empty and the helper executes no remote commands. Files and host directories are removed during teardown; `cleanup_complete` requires the owned directory to be absent before return. JSONL contains actual upload/download bytes, digest and verification outcome. Transfers incomplete @@ -282,8 +292,11 @@ future-dated heartbeats stop the run. A coordinator can `touch` the heartbeat at least every five seconds while continued authorization holds, or create `STOP` when an OOM, pressure, cost or other limit is reached. SIGINT/SIGTERM also cancel gracefully. All transports are closed, owned goroutines and host -processes are joined, and temporary directories are removed before the final -summary. The generator requires no Fly API credential. +processes are joined, and temporary directories are removed before a successful +cleanup report. Rooted SFTP admission closes and request workers/handles must +join before their root closes. If the bounded join fails, cleanup reports an +error and retains the root and directories; it never deletes beneath an active +worker. The generator requires no Fly API credential. Library callers use `Run(ctx, Config, io.Writer) (Summary, error)`. Supply the parsed manifest in `Config.Manifest`. `HelperExecutable` defaults to the running diff --git a/internal/capacity/fault_helper_test.go b/internal/capacity/fault_helper_test.go index 7d5b1a6b4..e44f8c126 100644 --- a/internal/capacity/fault_helper_test.go +++ b/internal/capacity/fault_helper_test.go @@ -73,6 +73,9 @@ func runFaultHelper(config faultHelperConfig, in *os.File, out io.Writer) error } defer func() { _ = term.Restore(int(in.Fd()), state) }() } + if config.Mode == "geometry" { + return runGeometryHelper(config.Replay, in, out) + } if _, err := io.WriteString(out, config.Replay); err != nil { return err } diff --git a/internal/capacity/geometry_test.go b/internal/capacity/geometry_test.go new file mode 100644 index 000000000..66ee25080 --- /dev/null +++ b/internal/capacity/geometry_test.go @@ -0,0 +1,199 @@ +package capacity + +import ( + "bufio" + "bytes" + "context" + "fmt" + "io" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "golang.org/x/crypto/ssh" + "golang.org/x/term" +) + +// Record the actual child terminal before the first guest can resize it. +func runGeometryHelper(report string, in *os.File, out io.Writer) error { + writeSize := func(suffix string) error { + cols, rows, err := term.GetSize(int(os.Stdout.Fd())) + if err != nil { + return err + } + return os.WriteFile(report+suffix, fmt.Appendf(nil, "%dx%d", cols, rows), 0600) + } + if err := writeSize("-startup"); err != nil { + return err + } + scanner := bufio.NewScanner(in) + for scanner.Scan() { + line := scanner.Text() + if line == "geometry" { + if err := writeSize("-after"); err != nil { + return err + } + continue + } + if observation, ok := parseFrame(line); ok { + kind := observation.kind + if kind == "Q" { + kind = "B" + } + if _, err := out.Write(frame(kind, observation.id, observation.seq, observation.sent, 128)); err != nil { + return err + } + } + } + return scanner.Err() +} + +type geometryOutput struct { + mu sync.Mutex + data bytes.Buffer + changed chan struct{} +} + +func (o *geometryOutput) Write(p []byte) (int, error) { + o.mu.Lock() + defer o.mu.Unlock() + n, err := o.data.Write(p) + select { + case o.changed <- struct{}{}: + default: + } + return n, err +} + +func (o *geometryOutput) wait(t *testing.T, marker []byte) { + t.Helper() + timer := time.NewTimer(5 * time.Second) + defer timer.Stop() + for { + o.mu.Lock() + rows := bytes.Split(o.data.Bytes(), []byte{'\n'}) + found := false + for _, row := range rows[:len(rows)-1] { + if bytes.Equal(bytes.TrimSuffix(row, []byte{'\r'}), marker) { + found = true + } + } + o.mu.Unlock() + if found { + return + } + select { + case <-o.changed: + case <-timer.C: + t.Fatal("helper response missing") + } + } +} + +func waitGeometryFile(t *testing.T, path string) string { + t.Helper() + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + for { + if data, err := os.ReadFile(path); err == nil && len(data) > 0 { + return string(data) + } + if err := pause(ctx, 10*time.Millisecond); err != nil { + t.Fatal("helper geometry report missing") + } + } +} + +func TestCapacityHelperStartsPinnedAndDoesNotRepaintOnJoin(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + report := filepath.Join(t.TempDir(), "geometry") + c := faultHelper(t, "geometry", report) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + ctx, cancel := context.WithCancelCause(t.Context()) + dir, err := os.MkdirTemp("", "capgeo-") + if err != nil { + t.Fatal(err) + } + r := &runState{cfg: c, ctx: ctx, cancel: cancel, guard: newGuard(c), rec: newRecorder(io.Discard, c), helper: c.HelperExecutable, dir: dir, hosts: make([]*hostRun, 1), guestIdentity: signer(t)} + t.Cleanup(func() { + if err := r.cleanup(); err != nil { + t.Error(err) + } + }) + target := c.Manifest.Targets[0] + if err := r.startHost(0, target, nil); err != nil { + t.Fatal(err) + } + if got := waitGeometryFile(t, report+"-startup"); got != "2048x100" { + t.Errorf("startup geometry = %s, want 2048x100", got) + } + endpoint, err := transport(target, ws, 0, 1, false) + if err != nil { + t.Fatal(err) + } + proxy, err := newTunnelProxy(ctx, r.guard, target) + if err != nil { + t.Fatal(err) + } + r.proxies = append(r.proxies, proxy) + connect := func(cols, rows int) (*ssh.Session, io.WriteCloser, *geometryOutput) { + attempt, end := context.WithTimeout(ctx, 5*time.Second) + defer end() + client, raw, err := openGuest(attempt, r.guard, target, endpoint, r.hosts[0].initial, r.guestIdentity, proxy) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Close(); _ = raw.Close() }) + sess, err := client.NewSession() + if err != nil { + t.Fatal(err) + } + if err := sess.RequestPty("dumb", rows, cols, ssh.TerminalModes{ssh.ECHO: 0}); err != nil { + t.Fatal(err) + } + in, err := sess.StdinPipe() + if err != nil { + t.Fatal(err) + } + out := &geometryOutput{changed: make(chan struct{}, 1)} + sess.Stdout, sess.Stderr = out, io.Discard + if err := sess.Shell(); err != nil { + t.Fatal(err) + } + return sess, in, out + } + _, first, output := connect(2048, 100) + stamp := time.Now() + marker := frame("I", 314, 7, stamp, 128) + if _, err := first.Write(marker); err != nil { + t.Fatal(err) + } + output.wait(t, marker[:127]) + second, _, _ := connect(80, 24) + if err := second.WindowChange(31, 90); err != nil { + t.Fatal(err) + } + if _, err := io.WriteString(first, "geometry\n"); err != nil { + t.Fatal(err) + } + if got := waitGeometryFile(t, report+"-after"); got != "2048x100" { + t.Errorf("joined/resized geometry = %s, want 2048x100", got) + } + barrier := frame("Q", 314, 8, stamp, 128) + if _, err := first.Write(barrier); err != nil { + t.Fatal(err) + } + output.wait(t, frame("B", 314, 8, stamp, 128)[:127]) + output.mu.Lock() + count := bytes.Count(output.data.Bytes(), marker[:127]) + captured := output.data.String() + output.mu.Unlock() + if count != 1 { + t.Fatalf("one-time marker rendered %d times: %q", count, captured) + } + }) + } +} diff --git a/internal/capacity/review_regression_test.go b/internal/capacity/review_regression_test.go index 66fb941e9..b58d7759d 100644 --- a/internal/capacity/review_regression_test.go +++ b/internal/capacity/review_regression_test.go @@ -362,6 +362,7 @@ func TestRetryBacklogIsGeneratorLimitedBeforeServiceStarts(t *testing.T) { func TestRealDeliveryEvidence(t *testing.T) { c := localRelay(t) + captureChurnTerminals(t, &c) c.DurationSeconds = 1 c.GuestWSPercent = 100 c.GuestsPerHost = 2 diff --git a/internal/capacity/runner.go b/internal/capacity/runner.go index 3792e8545..feaf93be7 100644 --- a/internal/capacity/runner.go +++ b/internal/capacity/runner.go @@ -18,6 +18,7 @@ import ( "github.com/owenthereal/upterm/host" "github.com/owenthereal/upterm/host/api" + "github.com/owenthereal/upterm/internal/termsize" "golang.org/x/crypto/ssh" ) @@ -61,6 +62,7 @@ type hostRun struct { established bool expectedCut bool tunnel chan tunnelObservation + started bool } func (r *runState) hostTunnel(h *hostRun, e host.TunnelEvent) { @@ -206,7 +208,19 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { proxy.queueEvent = func(wait time.Duration, limited bool) { r.rec.event("admission_queue", map[string]any{"role": "host", "host": index, "retry": true, "queue_ms": float64(wait) / float64(time.Millisecond), "generator_limited": limited}) } - h.h = &host.Host{Reconnect: host.ReconnectTiming{AttemptDeadline: 30*time.Second + retryQueueLimit}, KeepAliveDuration: 50 * time.Second, Host: u.String(), Command: []string{r.helper, "helper", strconv.Itoa(index), strconv.FormatBool(r.cfg.Active)}, CommandDir: dir, Signers: []ssh.Signer{identity}, IdentitiesOnly: true, HostKey: key, HostKeyCallback: pinnedKey(t), ProxyURL: proxy.url, AdminSocketFile: h.socket, AutoAccept: true, SFTPAutoAccept: true, ForwardingDisabled: true, ExecDisabled: true, Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), StopGrace: time.Second, + commandDir := dir + var root *os.Root + if r.cfg.TransferConcurrency > 0 { + commandDir = filepath.Join(dir, "transfers") + if err := os.Mkdir(commandDir, 0700); err != nil { + return err + } + root, err = os.OpenRoot(commandDir) + if err != nil { + return err + } + } + h.h = &host.Host{Reconnect: host.ReconnectTiming{AttemptDeadline: 30*time.Second + retryQueueLimit}, KeepAliveDuration: 50 * time.Second, Host: u.String(), Command: []string{r.helper, "helper", strconv.Itoa(index), strconv.FormatBool(r.cfg.Active)}, CommandDir: commandDir, Signers: []ssh.Signer{identity}, IdentitiesOnly: true, HostKey: key, HostKeyCallback: pinnedKey(t), ProxyURL: proxy.url, AdminSocketFile: h.socket, AutoAccept: true, SFTPAutoAccept: true, SFTPDisabled: root == nil, SFTPRoot: root, ForwardingDisabled: true, ExecDisabled: true, Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), StopGrace: time.Second, PtySize: termsize.Size{Cols: 2048, Rows: 100}, PinPtySize: true, AuthorizedKeys: []*host.AuthorizedKey{{PublicKeys: []ssh.PublicKey{r.guestIdentity.PublicKey()}}}, SessionCreatedCallback: func(ctx context.Context, s *api.GetSessionResponse) error { if err := r.guard.session(ctx, s); err != nil { @@ -243,6 +257,7 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { proxy.mu.Unlock() r.rec.event("admission_queue", map[string]any{"role": "host", "host": index, "queue_ms": float64(time.Since(queuedAt)) / float64(time.Millisecond)}) var began time.Time + h.started = true go func() { err := h.h.Run(r.ctx) h.done <- err @@ -349,12 +364,30 @@ func (r *runState) cleanup() error { for _, p := range r.proxies { p.close() } + var rootErr error for _, h := range r.hosts { if h != nil { hostErr := <-h.done _ = hostErr // Exit errors were classified before this cancellation/join. + if h.h.SFTPRoot != nil { + if h.started { + ctx, cancel := context.WithTimeout(context.Background(), finalDrainLimit) + err := h.h.WaitSFTP(ctx) + cancel() + if err != nil { + rootErr = errors.Join(rootErr, fmt.Errorf("host %d SFTP workers did not join; retained transfer root and directory: %w", h.id, err)) + continue + } + } + rootErr = errors.Join(rootErr, h.h.SFTPRoot.Close()) + } } } + // A timed-out filesystem/permission worker may still reference the root. + // Never let the parent RemoveAll bypass that failed lifetime proof. + if rootErr != nil { + return rootErr + } if err := os.RemoveAll(r.dir); err != nil { return err } diff --git a/internal/capacity/sftp_confinement_test.go b/internal/capacity/sftp_confinement_test.go new file mode 100644 index 000000000..33b8bed2b --- /dev/null +++ b/internal/capacity/sftp_confinement_test.go @@ -0,0 +1,200 @@ +package capacity + +import ( + "bytes" + "context" + "fmt" + "io" + "os" + "path/filepath" + "testing" + "time" + + "github.com/owenthereal/upterm/host" + "github.com/pkg/sftp" + "github.com/stretchr/testify/require" +) + +func TestCapacityZeroTransferRejectsSFTP(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + c := localRelay(t) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + r := admittedRun(t, c) + client, err := sftp.NewClient(r.guests[0].client) + if err == nil { + _ = client.Close() + t.Fatal("zero-transfer host admitted SFTP") + } + }) + } +} + +func TestCapacitySFTPCancellationClosesWorkersBeforeRoot(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + c := localRelay(t) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + c.TransferConcurrency, c.PayloadBudgetBytes = 4, 256<<20 + var root *os.Root + t.Cleanup(func() { + if root != nil { + _, err := root.Stat(".") + require.Error(t, err, "cleanup did not close the joined root") + } + }) + r := admittedRun(t, c) + h := r.hosts[0] + root = h.h.SFTPRoot + client, err := sftp.NewClient(r.guests[0].client) + require.NoError(t, err) + defer func() { _ = client.Close() }() + file, err := client.Create("open-during-cancel") + require.NoError(t, err) + defer func() { _ = file.Close() }() + _, err = file.Write([]byte("pending transfer")) + require.NoError(t, err) + r.cancel(context.Canceled) + select { + case err := <-h.done: + h.done <- err + case <-time.After(8 * time.Second): + t.Fatal("host cancellation did not finish promptly") + } + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + require.NoError(t, h.h.WaitSFTP(ctx)) + _, err = root.Stat(".") + require.NoError(t, err, "host took ownership of caller's root") + _, err = file.Write([]byte("late")) + require.Error(t, err, "open handle survived joined cancellation") + }) + } +} + +func TestCapacityRetainsRootWhenQuiescenceCannotBeProved(t *testing.T) { + dir := t.TempDir() + root, err := os.OpenRoot(dir) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + ctx, cancel := context.WithCancelCause(t.Context()) + h := &hostRun{h: &host.Host{SFTPRoot: root}, started: true, done: make(chan error, 1)} + // A completion notification without the matching Host lifecycle is not + // proof that its rooted workers have stopped. WaitSFTP must reject it. + h.done <- nil + r := &runState{ctx: ctx, cancel: cancel, hosts: []*hostRun{h}, dir: dir, rec: newRecorder(io.Discard, Config{})} + require.ErrorContains(t, r.cleanup(), "retained transfer root and directory") + _, err = root.Stat(".") + require.NoError(t, err) + _, err = os.Stat(dir) + require.NoError(t, err) + require.False(t, r.rec.summary.CleanupComplete) +} + +func TestCapacitySFTPConfinesContentAndMetadata(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + c := localRelay(t) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + c.TransferConcurrency, c.PayloadBudgetBytes = 4, 256<<20 + r := admittedRun(t, c) + q := r.guests[0] + client, err := sftp.NewClient(q.client) + if err != nil { + t.Fatal(err) + } + defer func() { _ = client.Close() }() + outsideDir := t.TempDir() + outside := filepath.Join(outsideDir, "sentinel") + payload := []byte("outside private contents") + if err := os.WriteFile(outside, payload, 0600); err != nil { + t.Fatal(err) + } + control := filepath.Join(filepath.Dir(q.host.socket), "control-secret") + if err := os.WriteFile(control, payload, 0600); err != nil { + t.Fatal(err) + } + inside := filepath.Join(q.host.h.CommandDir, "allowed") + f, err := client.Create(inside) + if err != nil { + t.Fatal(err) + } + if _, err := f.Write([]byte("inside")); err != nil { + t.Fatal(err) + } + if _, err := f.Stat(); err != nil { + t.Fatal(err) + } + if err := f.Chmod(0600); err == nil { + t.Error("handle setstat was allowed") + } + if err := f.Close(); err != nil { + t.Fatal(err) + } + f, err = client.Open("allowed") + if err != nil { + t.Error(err) + } else { + got, readErr := io.ReadAll(f) + _ = f.Close() + if readErr != nil || string(got) != "inside" { + t.Errorf("relative in-root contents: %q, %v", got, readErr) + } + } + for _, path := range []string{outside, control, q.host.socket, "../control-secret", "../../h0/control-secret", inside + ":stream", filepath.Join(q.host.h.CommandDir, "NUL")} { + t.Run(path, func(t *testing.T) { + if f, err := client.Open(path); err == nil { + _ = f.Close() + t.Error("outside/special open succeeded") + } + if f, err := client.Create(path); err == nil { + _ = f.Close() + t.Error("outside/special create succeeded") + } + if _, err := client.Stat(path); err == nil { + t.Error("outside/special stat succeeded") + } + if _, err := client.Lstat(path); err == nil { + t.Error("outside/special lstat succeeded") + } + }) + } + if _, err := client.ReadDir(outsideDir); err == nil { + t.Error("outside directory enumerated") + } + if err := client.Remove(outside); err == nil { + t.Error("outside remove succeeded") + } + emptyDir := filepath.Join(q.host.h.CommandDir, "empty") + require.NoError(t, os.Mkdir(emptyDir, 0700)) + require.Error(t, client.RemoveDirectory(emptyDir), "rmdir is unsupported") + require.NoError(t, client.Remove(emptyDir), "Root.Remove also permits an empty in-root directory") + _, err = client.RealPath(outside) + require.Error(t, err) + _, err = client.RealPath(".") + require.NoError(t, err) + for name, operation := range map[string]func() error{ + "mkdir": func() error { return client.Mkdir(filepath.Join(q.host.h.CommandDir, "newdir")) }, + "rename": func() error { return client.Rename(inside, inside+"-moved") }, + "posix rename": func() error { return client.PosixRename(inside, inside+"-moved") }, + "symlink": func() error { return client.Symlink(outside, filepath.Join(q.host.h.CommandDir, "link")) }, + "hardlink": func() error { return client.Link(outside, filepath.Join(q.host.h.CommandDir, "hardlink")) }, + "setstat": func() error { return client.Truncate(inside, 0) }, + "readlink": func() error { _, err := client.ReadLink("link"); return err }, + "statvfs": func() error { _, err := client.StatVFS(inside); return err }, + } { + t.Run(name, func(t *testing.T) { + if err := operation(); err == nil { + t.Error("unsupported operation was allowed") + } + }) + } + if data, err := os.ReadFile(outside); err != nil || !bytes.Equal(data, payload) { + t.Errorf("outside changed: %q, %v", data, err) + } + if data, err := os.ReadFile(control); err != nil || !bytes.Equal(data, payload) { + t.Errorf("control file changed: %q, %v", data, err) + } + }) + } +} diff --git a/internal/capacity/sftp_namespace_unix_test.go b/internal/capacity/sftp_namespace_unix_test.go new file mode 100644 index 000000000..a7a05149c --- /dev/null +++ b/internal/capacity/sftp_namespace_unix_test.go @@ -0,0 +1,122 @@ +//go:build !windows + +package capacity + +import ( + "context" + "fmt" + "io" + "os" + "path/filepath" + "testing" + + "github.com/pkg/sftp" + "github.com/stretchr/testify/require" +) + +// Unix permits replacing/renaming directories while they are open. All of the +// reads, writes and entry metadata must stay on the directory descriptors. +func TestCapacitySFTPNamespaceReplacement(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + c := localRelay(t) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + c.TransferConcurrency, c.PayloadBudgetBytes = 4, 256<<20 + r := admittedRun(t, c) + client, err := sftp.NewClient(r.guests[0].client) + require.NoError(t, err) + defer func() { _ = client.Close() }() + root := r.hosts[0].h.CommandDir + outside := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(outside, "sentinel"), []byte("outside secret"), 0600)) + race, parked := filepath.Join(root, "race"), filepath.Join(root, "parked") + require.NoError(t, os.Mkdir(race, 0700)) + require.NoError(t, os.WriteFile(filepath.Join(race, "sentinel"), []byte("inside"), 0600)) + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { + for ctx.Err() == nil { + if err := os.Rename(race, parked); err != nil { + done <- err + return + } + if err := os.Symlink(outside, race); err != nil { + done <- err + return + } + if err := os.Remove(race); err != nil { + done <- err + return + } + if err := os.Rename(parked, race); err != nil { + done <- err + return + } + } + done <- nil + }() + defer func() { cancel(); require.NoError(t, <-done) }() + for range 200 { + if f, err := client.Open("race/sentinel"); err == nil { + data, readErr := io.ReadAll(f) + _ = f.Close() + require.NoError(t, readErr) + require.Equal(t, "inside", string(data)) + } + if info, err := client.Stat("race/sentinel"); err == nil { + require.EqualValues(t, 6, info.Size()) + } + if infos, err := client.ReadDir("race"); err == nil { + for _, info := range infos { + if info.Name() == "sentinel" { + require.EqualValues(t, 6, info.Size(), "entry metadata escaped through a replaced parent") + } + } + } + if f, err := client.Create("race/created"); err == nil { + _, _ = f.Write([]byte("inside")) + _ = f.Close() + } + _ = client.Remove("race/created") + } + _, err = os.Stat(filepath.Join(outside, "created")) + require.True(t, os.IsNotExist(err), "a raced write escaped") + data, err := os.ReadFile(filepath.Join(outside, "sentinel")) + require.NoError(t, err) + require.Equal(t, "outside secret", string(data)) + }) + } +} + +func TestCapacitySFTPMetadataSurvivesRootRename(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + c := localRelay(t) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + c.TransferConcurrency, c.PayloadBudgetBytes = 4, 256<<20 + r := admittedRun(t, c) + client, err := sftp.NewClient(r.guests[0].client) + require.NoError(t, err) + defer func() { _ = client.Close() }() + root := r.hosts[0].h.CommandDir + require.NoError(t, os.WriteFile(filepath.Join(root, "file"), []byte("inside"), 0600)) + moved := root + "-moved" + require.NoError(t, os.Rename(root, moved)) + require.NoError(t, os.Mkdir(root, 0700)) + require.NoError(t, os.WriteFile(filepath.Join(root, "file"), []byte("outside substitute"), 0600)) + infos, err := client.ReadDir(root) + require.NoError(t, err) + require.Len(t, infos, 1) + require.EqualValues(t, 6, infos[0].Size(), "listing reopened the old path for entry metadata") + info, err := client.Stat(filepath.Join(root, "file")) + require.NoError(t, err) + require.EqualValues(t, 6, info.Size()) + f, err := client.Open(filepath.Join(root, "file")) + require.NoError(t, err) + data, err := io.ReadAll(f) + _ = f.Close() + require.NoError(t, err) + require.Equal(t, "inside", string(data)) + }) + } +} diff --git a/internal/capacity/sftp_symlink_test.go b/internal/capacity/sftp_symlink_test.go new file mode 100644 index 000000000..d7e553be4 --- /dev/null +++ b/internal/capacity/sftp_symlink_test.go @@ -0,0 +1,68 @@ +package capacity + +import ( + "fmt" + "io" + "os" + "path/filepath" + "testing" + + "github.com/pkg/sftp" + "github.com/stretchr/testify/require" +) + +func TestCapacitySFTPRejectsOutsideSymlinks(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + c := localRelay(t) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + c.TransferConcurrency, c.PayloadBudgetBytes = 4, 256<<20 + r := admittedRun(t, c) + client, err := sftp.NewClient(r.guests[0].client) + require.NoError(t, err) + defer func() { _ = client.Close() }() + root := r.hosts[0].h.CommandDir + outside := t.TempDir() + secret := filepath.Join(outside, "secret") + require.NoError(t, os.WriteFile(secret, []byte("outside secret"), 0600)) + require.NoError(t, os.Symlink(outside, filepath.Join(root, "escape"))) + require.NoError(t, os.Symlink(secret, filepath.Join(root, "link"))) + for _, path := range []string{"link", "escape/secret", filepath.Join(root, "escape", "secret")} { + f, err := client.Open(path) + if f != nil { + _ = f.Close() + } + require.Error(t, err, "outside content via %s", path) + f, err = client.Create(path) + if f != nil { + _ = f.Close() + } + require.Error(t, err, "outside write via %s", path) + _, err = client.Stat(path) + require.Error(t, err, "outside stat via %s", path) + } + _, err = client.ReadDir("escape") + require.Error(t, err) + _, err = client.ReadLink("link") + require.Error(t, err) + info, err := client.Lstat("link") + require.NoError(t, err) + require.NotZero(t, info.Mode()&os.ModeSymlink, "lstat returns in-root link metadata without following it") + require.NoError(t, os.WriteFile(filepath.Join(root, "file"), []byte("inside"), 0600)) + f, err := client.Open("file") + require.NoError(t, err) + defer func() { _ = f.Close() }() + require.NoError(t, os.Rename(filepath.Join(root, "file"), filepath.Join(root, "saved"))) + require.NoError(t, os.Symlink(secret, filepath.Join(root, "file"))) + _, err = f.Stat() + require.Error(t, err, "path-based protocol fstat must not follow a replaced path outside") + require.Error(t, f.Truncate(0), "handle setstat remains denied") + data, err := io.ReadAll(f) + require.NoError(t, err) + require.Equal(t, "inside", string(data), "open content handle stayed on its original file") + data, err = os.ReadFile(secret) + require.NoError(t, err) + require.Equal(t, "outside secret", string(data)) + }) + } +} diff --git a/internal/capacity/transfer.go b/internal/capacity/transfer.go index 3851f880b..1085975c8 100644 --- a/internal/capacity/transfer.go +++ b/internal/capacity/transfer.go @@ -55,7 +55,7 @@ func (r *runState) transfer(index int, expected string) (err error) { return err } defer func() { _ = c.Close() }() - path := filepath.ToSlash(filepath.Join(filepath.Dir(q.host.socket), fmt.Sprintf("transfer-%d.bin", index))) + path := filepath.ToSlash(filepath.Join(q.host.h.CommandDir, fmt.Sprintf("transfer-%d.bin", index))) defer func() { _ = c.Remove(path) }() began := time.Now() var uploaded, downloaded int64 diff --git a/internal/capacity/whole_review_test.go b/internal/capacity/whole_review_test.go index 4d97f2061..f73b60533 100644 --- a/internal/capacity/whole_review_test.go +++ b/internal/capacity/whole_review_test.go @@ -188,6 +188,8 @@ func TestRunGuestAuthorizationAcrossTransports(t *testing.T) { c := localRelay(t) c.Active = false c.Hosts = 2 + c.TransferConcurrency = 4 + c.PayloadBudgetBytes = 256 << 20 c.GuestWSPercent = ws r := admittedRun(t, c) runKey := r.guests[0].identity.PublicKey().Marshal() From d3f397803655d074859f14e9f693d3ac0a2b99f9 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Thu, 5 Nov 2026 23:21:47 -0700 Subject: [PATCH 14/16] test(capacity): log bounded output on geometry timeout --- internal/capacity/geometry_test.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/internal/capacity/geometry_test.go b/internal/capacity/geometry_test.go index 66ee25080..fc1c6a04b 100644 --- a/internal/capacity/geometry_test.go +++ b/internal/capacity/geometry_test.go @@ -87,7 +87,13 @@ func (o *geometryOutput) wait(t *testing.T, marker []byte) { select { case <-o.changed: case <-timer.C: - t.Fatal("helper response missing") + o.mu.Lock() + data := o.data.Bytes() + start := max(0, len(data)-4096) + diagnostic := fmt.Sprintf("helper response missing: marker=%q output_bytes=%d marker_offset=%d output_tail[%d:]=%q", + marker, len(data), bytes.Index(data, marker), start, data[start:]) + o.mu.Unlock() + t.Fatal(diagnostic) } } } From 24e9c35a5db5312a0b4d35eb5a5b8caf0d442721 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Fri, 6 Nov 2026 06:18:36 -0700 Subject: [PATCH 15/16] test(capacity): accept exact geometry frames after terminal prefixes --- internal/capacity/geometry_test.go | 47 +++++++++++++++++++++++++----- 1 file changed, 40 insertions(+), 7 deletions(-) diff --git a/internal/capacity/geometry_test.go b/internal/capacity/geometry_test.go index fc1c6a04b..411502513 100644 --- a/internal/capacity/geometry_test.go +++ b/internal/capacity/geometry_test.go @@ -67,19 +67,52 @@ func (o *geometryOutput) Write(p []byte) (int, error) { return n, err } +func geometryHasMarker(data, marker []byte) bool { + rows := bytes.Split(data, []byte{'\n'}) + for _, row := range rows[:len(rows)-1] { + row = bytes.TrimSuffix(row, []byte{'\r'}) + if start := bytes.LastIndex(row, []byte("CAP ")); start >= 0 && bytes.Equal(row[start:], marker) { + return true + } + } + return false +} + +func TestGeometryOutputMarker(t *testing.T) { + // Captured from Windows CI run 42174271728, job 127748995046. + const prefix = "\x1b[?9001h\x1b[?1004h\x1b[?25l\x1b[2J\x1b[m\x1b[H" + const marker = "CAP I 00000314 00000007 1793946543794605400 xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" + for _, tc := range []struct { + name string + output string + want bool + }{ + {"captured Windows prefix", prefix + marker + "\r\n", true}, + {"plain LF", marker + "\n", true}, + {"plain CRLF", marker + "\r\n", true}, + {"corrupt padding", prefix + marker[:len(marker)-1] + "y\r\n", false}, + {"corrupt header", prefix + "CAP B" + marker[5:] + "\r\n", false}, + {"wrong timestamp", prefix + marker[:24] + "1793946543794605401" + marker[43:] + "\r\n", false}, + {"trailing data", prefix + marker + "extra\r\n", false}, + {"trailing erase", prefix + marker + "\x1b[K\r\n", false}, + {"incomplete LF framing", prefix + marker + "\r", false}, + {"later mismatched CAP", prefix + marker + "CAP I 00000314 00000008 1793946543794605400\r\n", false}, + } { + t.Run(tc.name, func(t *testing.T) { + if got := geometryHasMarker([]byte(tc.output), []byte(marker)); got != tc.want { + t.Fatalf("geometryHasMarker(%q) = %t, want %t", tc.output, got, tc.want) + } + }) + } +} + func (o *geometryOutput) wait(t *testing.T, marker []byte) { t.Helper() timer := time.NewTimer(5 * time.Second) defer timer.Stop() for { o.mu.Lock() - rows := bytes.Split(o.data.Bytes(), []byte{'\n'}) - found := false - for _, row := range rows[:len(rows)-1] { - if bytes.Equal(bytes.TrimSuffix(row, []byte{'\r'}), marker) { - found = true - } - } + found := geometryHasMarker(o.data.Bytes(), marker) o.mu.Unlock() if found { return From b870950580f5a060893a1c4b012da63a59045e34 Mon Sep 17 00:00:00 2001 From: Owen Ou Date: Fri, 6 Nov 2026 09:28:38 -0700 Subject: [PATCH 16/16] fix(capacity): bound terminal buffers and rooted SFTP writes --- host/host.go | 21 +- host/internal/server.go | 11 +- host/internal/sftp.go | 8 + host/internal/sftp_limits.go | 111 ++++++++++ host/internal/sftp_limits_test.go | 234 ++++++++++++++++++++ host/internal/sftp_limits_unix_test.go | 140 ++++++++++++ host/sftp_limits_test.go | 20 ++ internal/capacity/bounded_frames_test.go | 245 +++++++++++++++++++++ internal/capacity/runner.go | 12 +- internal/capacity/sftp_confinement_test.go | 6 +- internal/capacity/sftp_limits_test.go | 201 +++++++++++++++++ internal/capacity/whole_review_test.go | 3 +- internal/capacity/workload.go | 98 ++++++++- 13 files changed, 1094 insertions(+), 16 deletions(-) create mode 100644 host/internal/sftp_limits.go create mode 100644 host/internal/sftp_limits_test.go create mode 100644 host/internal/sftp_limits_unix_test.go create mode 100644 host/sftp_limits_test.go create mode 100644 internal/capacity/bounded_frames_test.go create mode 100644 internal/capacity/sftp_limits_test.go diff --git a/host/host.go b/host/host.go index a861eaf35..4b534b8b0 100644 --- a/host/host.go +++ b/host/host.go @@ -96,6 +96,12 @@ type RelayWarning struct { Host string } +// SFTPWriteLimits bounds optional rooted uploads by expected filename, +// per-file extent and lifetime write bytes, and host-wide lifetime write bytes. +// Each filename has one non-reusable writable instance. Nil keeps the existing +// rooted behavior; non-nil limits require SFTPRoot. +type SFTPWriteLimits = internal.SFTPWriteLimits + type Host struct { Host string // KeepAliveDuration paces the guest sshd's keepalive to each guest. The @@ -329,10 +335,12 @@ type Host struct { // Nil preserves the ordinary filesystem/consent behavior. The caller owns // the root and must keep it open until Run has returned AND WaitSFTP succeeds. // Rooted hosts are single-use. Configure this before Run; do not change it. - SFTPRoot *os.Root - sftpMu sync.Mutex - sftpRunDone chan struct{} - sftpServer *internal.Server + SFTPRoot *os.Root + // SFTPWriteLimits is copied at Run entry; configure it before Run. + SFTPWriteLimits *SFTPWriteLimits + sftpMu sync.Mutex + sftpRunDone chan struct{} + sftpServer *internal.Server // PermissionChecker supplies custom consent policy. Nil denies ungranted operations. PermissionChecker consent.PermissionChecker // TerminalConsent selects consent on the primary local terminal; mutually exclusive with PermissionChecker. @@ -658,6 +666,10 @@ func (c *Host) Run(ctx context.Context) (runErr error) { c.sftpMu.Unlock() defer close(c.sftpRunDone) } + sftpWriteLimits, err := internal.CopySFTPWriteLimits(c.SFTPRoot, c.SFTPWriteLimits) + if err != nil { + return err + } // An empty command is no command, as the guest door reads one. if len(c.LoginCommand) > 0 && len(c.ForceCommand) > 0 { return errors.New("LoginCommand and ForceCommand are mutually exclusive") @@ -1298,6 +1310,7 @@ func (c *Host) Run(ctx context.Context) (runErr error) { StopGrace: c.StopGrace, SFTPDisabled: c.SFTPDisabled, SFTPRoot: c.SFTPRoot, + SFTPWriteLimits: sftpWriteLimits, PermissionChecker: c.PermissionChecker, TerminalConsent: c.TerminalConsent, SFTPAutoAccept: c.AutoAccept || c.SFTPAutoAccept, diff --git a/host/internal/server.go b/host/internal/server.go index 08afefb8f..a357683d9 100644 --- a/host/internal/server.go +++ b/host/internal/server.go @@ -133,8 +133,9 @@ type Server struct { // SFTP configuration SFTPDisabled bool // Disable SFTP subsystem entirely // SFTPRoot is caller-owned; WaitSFTP must succeed before it is closed. - SFTPRoot *os.Root - sftpWorkers rootedSFTPWorkers + SFTPRoot *os.Root + SFTPWriteLimits *SFTPWriteLimits + sftpWorkers rootedSFTPWorkers // PermissionChecker supplies custom consent policy. Nil denies ungranted operations. PermissionChecker consent.PermissionChecker // TerminalConsent selects consent on the primary local terminal; mutually exclusive with PermissionChecker. @@ -229,6 +230,10 @@ func (s *Server) ServeWithContext(ctx context.Context, guest, host net.Listener) return errors.New("rooted SFTP is unsupported on this platform") } } + sftpWriteLimits, err := newRootedSFTPWriteLimits(s.SFTPRoot, s.SFTPWriteLimits) + if err != nil { + return err + } if s.TerminalConsent && s.PermissionChecker != nil { return errors.New("TerminalConsent and PermissionChecker are mutually exclusive") } @@ -386,6 +391,7 @@ func (s *Server) ServeWithContext(ctx context.Context, guest, host net.Listener) cmdDone: cmdDone, commandResult: cmd.Result, sftpRoot: s.SFTPRoot, + sftpWriteLimits: sftpWriteLimits, onPrivateCommandExit: s.onPrivateCommandExit, } @@ -1105,6 +1111,7 @@ type sessionHandler struct { permissionChecker consent.PermissionChecker sftpAutoAccept bool sftpRoot *os.Root + sftpWriteLimits *rootedSFTPWriteLimits sftpWorkers *rootedSFTPWorkers // subsystems are the guest door's subsystem handlers, by name, so that diff --git a/host/internal/sftp.go b/host/internal/sftp.go index 7b794924d..3fd8cbd8b 100644 --- a/host/internal/sftp.go +++ b/host/internal/sftp.go @@ -40,6 +40,7 @@ type SFTPSession struct { ctx context.Context autoAccept bool root *os.Root + writeLimits *rootedSFTPWriteLimits } // HandleSFTP handles SFTP subsystem requests @@ -90,6 +91,7 @@ func (h *sessionHandler) HandleSFTP(sess gssh.Session) { clientInfo: clientInfo, logger: h.logger, root: h.sftpRoot, + writeLimits: h.sftpWriteLimits, } handlers := sftp.Handlers{ @@ -426,6 +428,12 @@ func (h *sftpFileWriter) Filewrite(r *sftp.Request) (io.WriterAt, error) { return nil, err } if h.session.root != nil { + if h.session.writeLimits != nil { + if !pflags.Write { + return nil, sftp.ErrSSHFxPermissionDenied + } + return h.session.writeLimits.open(h.session.root, path, osFlags) + } return h.session.root.OpenFile(path, osFlags, 0600) } return os.OpenFile(path, osFlags, 0644) diff --git a/host/internal/sftp_limits.go b/host/internal/sftp_limits.go new file mode 100644 index 000000000..f6cfe530e --- /dev/null +++ b/host/internal/sftp_limits.go @@ -0,0 +1,111 @@ +package internal + +import ( + "fmt" + "io" + "os" + "path/filepath" + "sync" + + "github.com/pkg/sftp" +) + +// SFTPWriteLimits optionally restricts rooted uploads for one host lifetime. +// Files maps canonical relative names to both their maximum logical extent +// and cumulative accepted write bytes. TotalBytes bounds accepted writes +// across all names. Each name permits only one writable OPEN, including a +// failed OPEN; close, remove, truncation and reconnect never reset allowances. +// Configure before Run/Serve. The server copies the declaration before use. +type SFTPWriteLimits struct { + Files map[string]int64 + TotalBytes int64 +} + +// CopySFTPWriteLimits validates and snapshots the caller's declaration before +// host startup can serve or publish anything. Nil preserves ordinary behavior. +func CopySFTPWriteLimits(root *os.Root, limits *SFTPWriteLimits) (*SFTPWriteLimits, error) { + if limits == nil { + return nil, nil + } + if root == nil { + return nil, fmt.Errorf("SFTPWriteLimits requires SFTPRoot") + } + if limits.TotalBytes < 0 { + return nil, fmt.Errorf("SFTPWriteLimits: negative total byte limit") + } + copy := &SFTPWriteLimits{Files: make(map[string]int64, len(limits.Files)), TotalBytes: limits.TotalBytes} + for name, size := range limits.Files { + rel, err := rootedSFTPPath(root, name) + if err != nil || name == "." || filepath.ToSlash(rel) != name || size < 0 { + return nil, fmt.Errorf("SFTPWriteLimits: invalid relative file or byte limit %q", name) + } + copy.Files[name] = size + } + return copy, nil +} + +// All sessions and transports of a Server share this state. In particular it +// must not be constructed in HandleSFTP, where reconnect would replenish it. +type rootedSFTPWriteLimits struct { + mu sync.Mutex + files map[string]*sftpFileAllowance + remaining int64 +} + +type sftpFileAllowance struct { + limit, remaining int64 + opened bool +} + +func newRootedSFTPWriteLimits(root *os.Root, limits *SFTPWriteLimits) (*rootedSFTPWriteLimits, error) { + declaration, err := CopySFTPWriteLimits(root, limits) + if err != nil || declaration == nil { + return nil, err + } + state := &rootedSFTPWriteLimits{files: make(map[string]*sftpFileAllowance, len(declaration.Files)), remaining: declaration.TotalBytes} + for name, size := range declaration.Files { + state.files[filepath.FromSlash(name)] = &sftpFileAllowance{limit: size, remaining: size} + } + return state, nil +} + +func (s *rootedSFTPWriteLimits) open(root *os.Root, name string, flags int) (io.WriterAt, error) { + s.mu.Lock() + allowance, allowed := s.files[name] + if !allowed || allowance.opened || flags&os.O_CREATE == 0 || flags&os.O_APPEND != 0 || flags&(os.O_TRUNC|os.O_EXCL) == 0 { + s.mu.Unlock() + return nil, sftp.ErrSSHFxPermissionDenied + } + // Claim before filesystem access, never recycle even after an open error. + allowance.opened = true + s.mu.Unlock() + f, err := root.OpenFile(name, flags, 0600) + if err != nil { + return nil, err + } + return &limitedSFTPWriter{file: f, limits: s, allowance: allowance}, nil +} + +type limitedSFTPWriter struct { + file *os.File + limits *rootedSFTPWriteLimits + allowance *sftpFileAllowance +} + +func (w *limitedSFTPWriter) WriteAt(p []byte, offset int64) (int, error) { + size := int64(len(p)) + w.limits.mu.Lock() + if offset < 0 || offset > w.allowance.limit || size > w.allowance.limit-offset || + size > w.allowance.remaining || size > w.limits.remaining { + w.limits.mu.Unlock() + return 0, sftp.ErrSSHFxPermissionDenied + } + w.allowance.remaining -= size + w.limits.remaining -= size + w.limits.mu.Unlock() + // Reserve before I/O, without refund on a short or failed write. This + // remains safe for concurrent writes and deleted-but-still-open handles. + return w.file.WriteAt(p, offset) +} + +func (w *limitedSFTPWriter) Close() error { return w.file.Close() } diff --git a/host/internal/sftp_limits_test.go b/host/internal/sftp_limits_test.go new file mode 100644 index 000000000..6f8b5c9fc --- /dev/null +++ b/host/internal/sftp_limits_test.go @@ -0,0 +1,234 @@ +package internal + +import ( + "io" + "math" + "os" + "path/filepath" + "sync" + "testing" + + "github.com/pkg/sftp" + "github.com/stretchr/testify/require" +) + +func limitedSession(t *testing.T, root *os.Root, state *rootedSFTPWriteLimits) *SFTPSession { + t.Helper() + return &SFTPSession{root: root, writeLimits: state, autoAccept: true, ctx: t.Context(), logger: discardLogger()} +} + +func openLimitedFile(s *SFTPSession, name string, flags uint32) (io.WriterAt, error) { + r := sftp.NewRequest("Put", filepath.Join(s.root.Name(), name)) + r.Flags = flags + return (&sftpFileWriter{session: s, logger: s.logger}).Filewrite(r) +} + +// SFTP wire flags: WRITE | CREAT | TRUNC. Client.Create also sets READ; +// pkg/sftp still routes it to Filewrite for a write-only handler. +const limitedCreateFlags = 2 | 8 | 16 + +func TestRootedSFTPLimitsValidateBeforeServing(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + for _, tc := range []struct { + name string + root *os.Root + limits *SFTPWriteLimits + }{ + {"no root", nil, &SFTPWriteLimits{}}, + {"negative total", root, &SFTPWriteLimits{TotalBytes: -1}}, + {"negative file", root, &SFTPWriteLimits{Files: map[string]int64{"file": -1}}}, + } { + t.Run(tc.name, func(t *testing.T) { + s := &Server{SFTPRoot: tc.root, SFTPWriteLimits: tc.limits} + require.ErrorContains(t, s.ServeWithContext(t.Context(), nil, nil), "SFTPWriteLimits") + }) + } + for _, name := range []string{"", ".", "..", "../outside", "/absolute", "a/../file", "a//file", "file/", `a\file`, "file:stream", "NUL", "dir/COM1"} { + _, err := CopySFTPWriteLimits(root, &SFTPWriteLimits{Files: map[string]int64{name: 1}, TotalBytes: 1}) + require.Error(t, err, "invalid policy name %q", name) + } + for _, limits := range []*SFTPWriteLimits{nil, {}, {Files: map[string]int64{"empty": 0}}, {Files: map[string]int64{"dir/file": 8}, TotalBytes: 8}} { + copy, err := CopySFTPWriteLimits(root, limits) + require.NoError(t, err) + _, err = CopySFTPWriteLimits(root, copy) + require.NoError(t, err, "a valid declaration remains canonical when copied again") + } +} + +func TestRootedSFTPLimitsShareExtentFileAndTotalBudgets(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + declaration := &SFTPWriteLimits{Files: map[string]int64{"a": 8, "b": 8}, TotalBytes: 10} + state, err := newRootedSFTPWriteLimits(root, declaration) + require.NoError(t, err) + // Caller changes cannot enlarge a running host's immutable allowance. + declaration.Files["a"], declaration.Files["unexpected"], declaration.TotalBytes = 100, 100, 100 + one, two := limitedSession(t, root, state), limitedSession(t, root, state) + _, err = openLimitedFile(one, "unexpected", limitedCreateFlags) + require.Error(t, err) + a, err := openLimitedFile(one, "a", limitedCreateFlags) + require.NoError(t, err) + defer func() { _ = a.(io.Closer).Close() }() + for _, offset := range []int64{-1, 8, math.MaxInt64} { + n, err := a.WriteAt([]byte("x"), offset) + require.Error(t, err) + require.Zero(t, n) + } + n, err := a.WriteAt([]byte("efgh"), 4) + require.NoError(t, err) + require.Equal(t, 4, n) + n, err = a.WriteAt([]byte("abcd"), 0) + require.NoError(t, err) + require.Equal(t, 4, n) + n, err = a.WriteAt([]byte("x"), 0) + require.Error(t, err, "overwriting does not replenish the per-file allowance") + require.Zero(t, n) + b, err := openLimitedFile(two, "b", limitedCreateFlags) + require.NoError(t, err) + defer func() { _ = b.(io.Closer).Close() }() + n, err = b.WriteAt([]byte("xy"), 0) + require.NoError(t, err) + require.Equal(t, 2, n) + n, err = b.WriteAt([]byte("z"), 2) + require.Error(t, err, "a second session cannot exceed the host aggregate") + require.Zero(t, n) + data, err := root.ReadFile("a") + require.NoError(t, err) + require.Equal(t, "abcdefgh", string(data)) + data, err = root.ReadFile("b") + require.NoError(t, err) + require.Equal(t, "xy", string(data), "denied aggregate write reached disk") +} + +func TestRootedSFTPLimitsReserveBeforeConcurrentIO(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + state, err := newRootedSFTPWriteLimits(root, &SFTPWriteLimits{Files: map[string]int64{"a": 8, "b": 8, "c": 8}, TotalBytes: 10}) + require.NoError(t, err) + var wg sync.WaitGroup + accepted := make(chan int, 30) + for _, name := range []string{"a", "b", "c"} { + f, err := openLimitedFile(limitedSession(t, root, state), name, limitedCreateFlags) + require.NoError(t, err) + defer func() { _ = f.(io.Closer).Close() }() + for range 10 { + wg.Go(func() { n, _ := f.WriteAt([]byte("x"), 0); accepted <- n }) + } + } + wg.Wait() + close(accepted) + total := 0 + for n := range accepted { + total += n + } + require.Equal(t, 10, total, "concurrent sessions exceeded the shared aggregate") +} + +func TestRootedSFTPLimitsNeverRefundFailedWritesOrOpens(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + state, err := newRootedSFTPWriteLimits(root, &SFTPWriteLimits{Files: map[string]int64{"a": 2, "b": 2, "dir/file": 2}, TotalBytes: 2}) + require.NoError(t, err) + s := limitedSession(t, root, state) + _, err = openLimitedFile(s, "dir/file", limitedCreateFlags) + require.Error(t, err, "missing parent should fail OPEN") + require.NoError(t, root.Mkdir("dir", 0700)) + _, err = openLimitedFile(s, "dir/file", limitedCreateFlags) + require.ErrorIs(t, err, sftp.ErrSSHFxPermissionDenied, "failed OPEN recycled a creation slot") + f, err := openLimitedFile(s, "a", limitedCreateFlags) + require.NoError(t, err) + require.NoError(t, f.(io.Closer).Close()) + n, err := f.WriteAt([]byte("xx"), 0) + require.Error(t, err, "underlying handle is closed") + require.Zero(t, n) + b, err := openLimitedFile(s, "b", limitedCreateFlags) + require.NoError(t, err) + defer func() { _ = b.(io.Closer).Close() }() + n, err = b.WriteAt([]byte("x"), 0) + require.ErrorIs(t, err, sftp.ErrSSHFxPermissionDenied, "failed I/O refunded the host budget") + require.Zero(t, n) +} + +func TestRootedSFTPLimitsKeepConsentAndRejectReopenModes(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + state, err := newRootedSFTPWriteLimits(root, &SFTPWriteLimits{Files: map[string]int64{"a": 8}, TotalBytes: 8}) + require.NoError(t, err) + s := limitedSession(t, root, state) + s.autoAccept = false + _, err = openLimitedFile(s, "a", limitedCreateFlags) + require.Error(t, err, "limits must not bypass consent") + s.autoAccept, s.readOnly = true, true + _, err = openLimitedFile(s, "a", limitedCreateFlags) + require.Error(t, err, "limits must not bypass read-only") + s.readOnly = false + for _, flags := range []uint32{2, 2 | 16, limitedCreateFlags | 4, 8 | 16} { + _, err = openLimitedFile(s, "a", flags) + require.Error(t, err, "unsupported writable mode %d", flags) + } + f, err := openLimitedFile(s, "a", limitedCreateFlags|1) + require.NoError(t, err, "normal Client.Create flags remain supported") + require.NoError(t, f.(io.Closer).Close()) + _, err = openLimitedFile(s, "a", limitedCreateFlags) + require.Error(t, err, "close recycled the creation slot") + _, err = newRootedSFTPWriteLimits(root, nil) + require.NoError(t, err) + // A rooted host without limits retains ordinary writable reopen behavior. + s = limitedSession(t, root, nil) + f, err = openLimitedFile(s, "a", limitedCreateFlags) + require.NoError(t, err) + _, err = f.WriteAt([]byte("unrestricted"), 100) + require.NoError(t, err) + require.NoError(t, f.(io.Closer).Close()) + _, err = root.Stat("a") + require.NoError(t, err) +} + +func TestRootedSFTPLimitsRejectNontruncatingCreateOnExistingFile(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + for _, name := range []string{"existing", "truncate", "ordinary"} { + require.NoError(t, root.WriteFile(name, []byte("123456789"), 0600)) + } + state, err := newRootedSFTPWriteLimits(root, &SFTPWriteLimits{Files: map[string]int64{"existing": 8, "truncate": 8, "exclusive": 8}, TotalBytes: 24}) + require.NoError(t, err) + s := limitedSession(t, root, state) + f, openErr := openLimitedFile(s, "existing", 2|8) // WRITE | CREATE, neither TRUNC nor EXCL + if f != nil { + require.NoError(t, f.(io.Closer).Close()) + } + data, err := root.ReadFile("existing") + require.NoError(t, err) + require.Equal(t, "123456789", string(data), "rejected mode must not perform I/O") + // Both supported creation modes still provide a fresh bounded instance. + f, err = openLimitedFile(s, "truncate", limitedCreateFlags|1) + require.NoError(t, err, "Client.Create remains supported") + info, err := root.Stat("truncate") + require.NoError(t, err) + require.Zero(t, info.Size()) + _, err = f.WriteAt([]byte("12345678"), 0) + require.NoError(t, err) + require.NoError(t, f.(io.Closer).Close()) + f, err = openLimitedFile(s, "exclusive", 2|8|32) // WRITE | CREATE | EXCL + require.NoError(t, err) + _, err = f.WriteAt([]byte("12345678"), 0) + require.NoError(t, err) + require.NoError(t, f.(io.Closer).Close()) + // Optional limits do not change the existing rooted/nil-policy modes. + f, err = openLimitedFile(limitedSession(t, root, nil), "ordinary", 2|8) + require.NoError(t, err) + _, err = f.WriteAt([]byte("x"), 9) + require.NoError(t, err) + require.NoError(t, f.(io.Closer).Close()) + info, err = root.Stat("ordinary") + require.NoError(t, err) + require.EqualValues(t, 10, info.Size()) + require.ErrorIs(t, openErr, sftp.ErrSSHFxPermissionDenied, "nontruncating create preserved an extent above its limit") +} diff --git a/host/internal/sftp_limits_unix_test.go b/host/internal/sftp_limits_unix_test.go new file mode 100644 index 000000000..1dd0cd880 --- /dev/null +++ b/host/internal/sftp_limits_unix_test.go @@ -0,0 +1,140 @@ +//go:build !windows + +package internal + +import ( + "context" + "io" + "os" + "path/filepath" + "testing" + + "github.com/pkg/sftp" + "github.com/stretchr/testify/require" +) + +func TestRootedSFTPLimitsSurviveSessionAndTransportReplacement(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + limits := &SFTPWriteLimits{Files: map[string]int64{"a": 8, "b": 8, "c": 8}, TotalBytes: 10} + h := startHost(t, &Server{Command: readsALine("READY", 0), SFTPRoot: root, SFTPWriteLimits: limits, SFTPAutoAccept: true}) + connection := h.guestClient(t) + first, err := sftp.NewClient(connection) + require.NoError(t, err) + f, err := first.Create("a") + require.NoError(t, err) + // Changing the caller's configuration after serving cannot add allowance. + limits.Files["extra"], limits.TotalBytes = 100, 100 + _, err = f.Write([]byte("abcdefgh")) + require.NoError(t, err) + require.NoError(t, f.Close()) + require.NoError(t, first.Close()) + second, err := sftp.NewClient(connection) + require.NoError(t, err) + _, err = second.Create("extra") + require.Error(t, err) + _, err = second.Create("a") + require.Error(t, err, "new channel reset file-instance state") + f, err = second.Create("b") + require.NoError(t, err) + _, err = f.Write([]byte("xy")) + require.NoError(t, err) + _, err = f.Write([]byte("z")) + require.Error(t, err, "new channel reset aggregate accounting") + require.NoError(t, f.Close()) + require.NoError(t, second.Close()) + require.NoError(t, connection.Close()) + third, err := sftp.NewClient(h.guestClient(t)) + require.NoError(t, err) + defer func() { _ = third.Close() }() + f, err = third.Create("c") + require.NoError(t, err) + _, err = f.Write([]byte("x")) + require.Error(t, err, "new SSH transport reset aggregate accounting") + require.NoError(t, f.Close()) + h.stop(t) + require.NoError(t, h.srv.WaitSFTP(t.Context())) +} + +// The capacity allowlist deliberately rejects arbitrary namespace names. +// Exercise the descriptor-rooted write path without that optional policy so +// its race protection cannot pass solely because of an allowlist denial. +func TestRootedSFTPNilLimitsConfinesWritesDuringNamespaceReplacement(t *testing.T) { + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + h := startHost(t, &Server{Command: readsALine("READY", 0), SFTPRoot: root, SFTPAutoAccept: true}) + client, err := sftp.NewClient(h.guestClient(t)) + require.NoError(t, err) + defer func() { _ = client.Close() }() + outside := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(outside, "sentinel"), []byte("outside secret"), 0600)) + race, parked := filepath.Join(root.Name(), "race"), filepath.Join(root.Name(), "parked") + require.NoError(t, os.Mkdir(race, 0700)) + require.NoError(t, os.WriteFile(filepath.Join(race, "sentinel"), []byte("inside"), 0600)) + // A real successful writable control proves this reaches filesystem I/O. + f, err := client.Create("race/created") + require.NoError(t, err) + _, err = f.Write([]byte("inside")) + require.NoError(t, err) + require.NoError(t, f.Close()) + control, err := root.ReadFile("race/created") + require.NoError(t, err) + require.Equal(t, "inside", string(control)) + require.NoError(t, client.Remove("race/created")) + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { + for ctx.Err() == nil { + if err := os.Rename(race, parked); err != nil { + done <- err + return + } + if err := os.Symlink(outside, race); err != nil { + done <- err + return + } + if err := os.Remove(race); err != nil { + done <- err + return + } + if err := os.Rename(parked, race); err != nil { + done <- err + return + } + } + done <- nil + }() + defer func() { cancel(); require.NoError(t, <-done) }() + for range 200 { + if f, err := client.Open("race/sentinel"); err == nil { + data, readErr := io.ReadAll(f) + _ = f.Close() + require.NoError(t, readErr) + require.Equal(t, "inside", string(data)) + } + if info, err := client.Stat("race/sentinel"); err == nil { + require.EqualValues(t, 6, info.Size()) + } + if infos, err := client.ReadDir("race"); err == nil { + for _, info := range infos { + if info.Name() == "sentinel" { + require.EqualValues(t, 6, info.Size()) + } + } + } + if f, err := client.Create("race/created"); err == nil { + _, _ = f.Write([]byte("inside")) + _ = f.Close() + } + _ = client.Remove("race/created") + } + _, err = os.Stat(filepath.Join(outside, "created")) + require.True(t, os.IsNotExist(err), "a raced write escaped") + data, err := os.ReadFile(filepath.Join(outside, "sentinel")) + require.NoError(t, err) + require.Equal(t, "outside secret", string(data)) + h.stop(t) + require.NoError(t, h.srv.WaitSFTP(t.Context())) +} diff --git a/host/sftp_limits_test.go b/host/sftp_limits_test.go new file mode 100644 index 000000000..7a0ea1775 --- /dev/null +++ b/host/sftp_limits_test.go @@ -0,0 +1,20 @@ +package host + +import ( + "os" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHostRejectsInvalidSFTPWriteLimitsBeforeStartup(t *testing.T) { + require.ErrorContains(t, (&Host{SFTPWriteLimits: &SFTPWriteLimits{}}).Run(t.Context()), "requires SFTPRoot") + root, err := os.OpenRoot(t.TempDir()) + require.NoError(t, err) + defer func() { require.NoError(t, root.Close()) }() + h := &Host{SFTPRoot: root, SFTPWriteLimits: &SFTPWriteLimits{Files: map[string]int64{"../outside": 1}, TotalBytes: 1}} + require.ErrorContains(t, h.Run(t.Context()), "SFTPWriteLimits") + require.NoError(t, h.WaitSFTP(t.Context())) + _, err = root.Stat(".") + require.NoError(t, err, "invalid policy must not take ownership of the root") +} diff --git a/internal/capacity/bounded_frames_test.go b/internal/capacity/bounded_frames_test.go new file mode 100644 index 000000000..fa45c5cf1 --- /dev/null +++ b/internal/capacity/bounded_frames_test.go @@ -0,0 +1,245 @@ +package capacity + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// Generate a large unfinished screen row without allocating that row. Keep +// readers alive at its end so the bound applies to resident, usable channels. +type heldFramePrefix struct { + remaining int + tail *strings.Reader + reached chan struct{} + release chan struct{} + once sync.Once + released sync.Once + maxRead atomic.Int64 +} + +func (r *heldFramePrefix) Read(p []byte) (int, error) { + for old := r.maxRead.Load(); int64(len(p)) > old; old = r.maxRead.Load() { + if r.maxRead.CompareAndSwap(old, int64(len(p))) { + break + } + } + if r.remaining > 0 { + n := min(len(p), r.remaining) + for i := range n { + p[i] = 'x' + } + r.remaining -= n + return n, nil + } + r.once.Do(func() { close(r.reached) }) + <-r.release + return r.tail.Read(p) +} + +func (r *heldFramePrefix) Close() error { + r.released.Do(func() { close(r.release) }) + return nil +} + +func TestTerminalReadersKeepFixedBuffers(t *testing.T) { + for _, helperInput := range []bool{false, true} { + t.Run(fmt.Sprintf("helper=%t", helperInput), func(t *testing.T) { + const channels = 4 + var observations atomic.Int64 + var wg sync.WaitGroup + errors := make(chan error, channels) + readers := make([]*heldFramePrefix, channels) + t.Cleanup(func() { + for _, r := range readers { + if r != nil { + _ = r.Close() + } + } + wg.Wait() + }) + for i := range readers { + fresh := string(frame("I", i, 9, time.Unix(123, 456), 128)) + r := &heldFramePrefix{remaining: 600 << 10, tail: strings.NewReader(fresh), reached: make(chan struct{}), release: make(chan struct{})} + readers[i] = r + wg.Go(func() { + if helperInput { + var out bytes.Buffer + err := helper(context.Background(), r, &out, 0, false, time.Hour) + if err == nil && out.String() != "CAP READY\n"+fresh { + err = fmt.Errorf("helper lost fresh frame: %q", out.String()) + } + errors <- err + return + } + err := scanFrames(r, func(o observation) error { + observations.Add(1) + if o.kind != "I" || o.seq != 9 { + return fmt.Errorf("wrong fresh frame: %+v", o) + } + return nil + }, func(observation) { observations.Add(1) }) + if err == io.EOF { + err = nil + } + errors <- err + }) + } + for _, r := range readers { + <-r.reached + require.Zero(t, observations.Load(), "unfinished junk became a frame") + if r.maxRead.Load() > 4096 { + t.Errorf("terminal reader grew its read buffer to %d bytes", r.maxRead.Load()) + } + } + for _, r := range readers { + _ = r.Close() + } + wg.Wait() + close(errors) + for err := range errors { + require.NoError(t, err) + } + if !helperInput { + require.EqualValues(t, channels, observations.Load()) + } + }) + } +} + +// A Reader may deliver its final bytes and error together or separately. +type terminalEndingReader struct { + data string + end error + step int + withData bool +} + +func (r *terminalEndingReader) Read(p []byte) (int, error) { + if len(r.data) == 0 { + return 0, r.end + } + if r.step > 0 && len(p) > r.step { + p = p[:r.step] + } + n := copy(p, r.data) + r.data = r.data[n:] + if len(r.data) == 0 && r.withData { + return n, r.end + } + return n, nil +} + +func TestTerminalFrameRawBoundaries(t *testing.T) { + broken := errors.New("terminal read failed") + for _, kind := range []string{"I", "H", "R", "A", "Q", "F", "B"} { + for _, ending := range []string{"\n", "\r\n", "\r\r\n", "\r\r\r\n", "", "\r", "\r\r"} { + for _, delta := range []int{-1, 0, 1} { + for _, end := range []error{io.EOF, broken} { + for _, together := range []bool{false, true} { + name := fmt.Sprintf("%s/ending=%q/delta=%d/%v/together=%t", kind, ending, delta, end, together) + t.Run(name, func(t *testing.T) { + canonical := string(frame(kind, 3, 7, time.Unix(123, 456), frameLength(kind))) + body := strings.TrimSuffix(canonical, "\n") + ending + data := strings.Repeat("x", (1<<20)+delta-len(body)) + body + reader := &terminalEndingReader{data: data, end: end, withData: together} + var delivered, ready int + err := scanFrames(reader, func(o observation) error { delivered++; return nil }, func(observation) { ready++ }) + terminated := strings.HasSuffix(ending, "\n") + tooLong := delta > 0 || (delta == 0 && !terminated && !together) + valid := !tooLong && ending != "\r\r\r\n" + wantDelivered, wantReady := 0, 0 + if valid && kind == "A" { + wantReady = 1 + } + if valid && kind != "A" && kind != "R" { + wantDelivered = 1 + } + require.Equal(t, wantDelivered, delivered) + require.Equal(t, wantReady, ready) + if tooLong { + require.ErrorContains(t, err, "token too long") + } else { + require.ErrorIs(t, err, end) + } + }) + } + } + } + } + } +} + +func TestTerminalFrameSplitTagsAndFinalCanonicalPayload(t *testing.T) { + for _, kind := range []string{"I", "H", "R", "A", "Q", "F", "B"} { + for _, step := range []int{1, 2, 3, 127, 1024} { + canonical := strings.TrimSuffix(string(frame(kind, 3, 7, time.Unix(123, 456), frameLength(kind))), "\n") + for _, suffix := range []string{canonical, canonical + "x", canonical[:len(canonical)-1], canonical + "CAP bad", canonical[:len(canonical)-1] + "y"} { + data := "CAP stale replay " + suffix + "\r\n" + delivered, ready := 0, 0 + err := scanFrames(&terminalEndingReader{data: data, end: io.EOF, step: step}, func(observation) error { delivered++; return nil }, func(observation) { ready++ }) + require.ErrorIs(t, err, io.EOF) + wantDelivered, wantReady := 0, 0 + if suffix == canonical && kind == "A" { + wantReady = 1 + } + if suffix == canonical && kind != "A" && kind != "R" { + wantDelivered = 1 + } + require.Equal(t, wantDelivered, delivered, "kind=%s step=%d", kind, step) + require.Equal(t, wantReady, ready, "kind=%s step=%d", kind, step) + } + } + } +} + +func TestHelperFrameRawBoundaries(t *testing.T) { + broken := errors.New("helper read failed") + for _, kind := range []string{"I", "H", "R", "A", "Q", "F", "B"} { + for _, ending := range []string{"\n", "\r\n", "\r\r\n", "", "\r", "\r\r"} { + for _, delta := range []int{-1, 0, 1} { + for _, end := range []error{io.EOF, broken} { + for _, together := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/%q/%d/%v/%t", kind, ending, delta, end, together), func(t *testing.T) { + stamp := time.Unix(123, 456) + body := strings.TrimSuffix(string(frame(kind, 3, 7, stamp, frameLength(kind))), "\n") + ending + data := strings.Repeat("x", (1<<20)+delta-len(body)) + body + var out bytes.Buffer + err := helper(context.Background(), &terminalEndingReader{data: data, end: end, withData: together}, &out, 0, false, time.Hour) + tooLong := delta > 0 || (delta == 0 && !strings.HasSuffix(ending, "\n") && !together) + valid := !tooLong && ending != "\r\r\n" && ending != "\r\r" + want := "CAP READY\n" + if valid { + switch kind { + case "I": + want += string(frame("I", 3, 7, stamp, 128)) + case "R": + want += string(frame("A", 3, 7, stamp, 128)) + case "Q", "F": + want += string(frame("B", 3, 0, stamp, 128)) + } + } + require.Equal(t, want, out.String()) + if tooLong { + require.ErrorContains(t, err, "token too long") + } else if end == io.EOF { + require.NoError(t, err) + } else { + require.ErrorIs(t, err, end) + } + }) + } + } + } + } + } +} diff --git a/internal/capacity/runner.go b/internal/capacity/runner.go index feaf93be7..253b3b5e7 100644 --- a/internal/capacity/runner.go +++ b/internal/capacity/runner.go @@ -210,7 +210,16 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { } commandDir := dir var root *os.Root + var writeLimits *host.SFTPWriteLimits if r.cfg.TransferConcurrency > 0 { + writeLimits = &host.SFTPWriteLimits{Files: make(map[string]int64)} + for transfer := 0; transfer < r.cfg.TransferConcurrency; transfer++ { + guest := transfer % (r.cfg.Hosts * r.cfg.GuestsPerHost) + if guest/r.cfg.GuestsPerHost == index { + writeLimits.Files[fmt.Sprintf("transfer-%d.bin", transfer)] = TransferSize + writeLimits.TotalBytes += TransferSize + } + } commandDir = filepath.Join(dir, "transfers") if err := os.Mkdir(commandDir, 0700); err != nil { return err @@ -221,7 +230,8 @@ func (r *runState) startHost(index int, t Target, identity ssh.Signer) error { } } h.h = &host.Host{Reconnect: host.ReconnectTiming{AttemptDeadline: 30*time.Second + retryQueueLimit}, KeepAliveDuration: 50 * time.Second, Host: u.String(), Command: []string{r.helper, "helper", strconv.Itoa(index), strconv.FormatBool(r.cfg.Active)}, CommandDir: commandDir, Signers: []ssh.Signer{identity}, IdentitiesOnly: true, HostKey: key, HostKeyCallback: pinnedKey(t), ProxyURL: proxy.url, AdminSocketFile: h.socket, AutoAccept: true, SFTPAutoAccept: true, SFTPDisabled: root == nil, SFTPRoot: root, ForwardingDisabled: true, ExecDisabled: true, Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), StopGrace: time.Second, PtySize: termsize.Size{Cols: 2048, Rows: 100}, PinPtySize: true, - AuthorizedKeys: []*host.AuthorizedKey{{PublicKeys: []ssh.PublicKey{r.guestIdentity.PublicKey()}}}, + SFTPWriteLimits: writeLimits, + AuthorizedKeys: []*host.AuthorizedKey{{PublicKeys: []ssh.PublicKey{r.guestIdentity.PublicKey()}}}, SessionCreatedCallback: func(ctx context.Context, s *api.GetSessionResponse) error { if err := r.guard.session(ctx, s); err != nil { return err diff --git a/internal/capacity/sftp_confinement_test.go b/internal/capacity/sftp_confinement_test.go index 33b8bed2b..09435e21c 100644 --- a/internal/capacity/sftp_confinement_test.go +++ b/internal/capacity/sftp_confinement_test.go @@ -49,7 +49,7 @@ func TestCapacitySFTPCancellationClosesWorkersBeforeRoot(t *testing.T) { client, err := sftp.NewClient(r.guests[0].client) require.NoError(t, err) defer func() { _ = client.Close() }() - file, err := client.Create("open-during-cancel") + file, err := client.Create("transfer-0.bin") require.NoError(t, err) defer func() { _ = file.Close() }() _, err = file.Write([]byte("pending transfer")) @@ -114,7 +114,7 @@ func TestCapacitySFTPConfinesContentAndMetadata(t *testing.T) { if err := os.WriteFile(control, payload, 0600); err != nil { t.Fatal(err) } - inside := filepath.Join(q.host.h.CommandDir, "allowed") + inside := filepath.Join(q.host.h.CommandDir, "transfer-0.bin") f, err := client.Create(inside) if err != nil { t.Fatal(err) @@ -131,7 +131,7 @@ func TestCapacitySFTPConfinesContentAndMetadata(t *testing.T) { if err := f.Close(); err != nil { t.Fatal(err) } - f, err = client.Open("allowed") + f, err = client.Open("transfer-0.bin") if err != nil { t.Error(err) } else { diff --git a/internal/capacity/sftp_limits_test.go b/internal/capacity/sftp_limits_test.go new file mode 100644 index 000000000..601dd8916 --- /dev/null +++ b/internal/capacity/sftp_limits_test.go @@ -0,0 +1,201 @@ +package capacity + +import ( + "bytes" + "context" + "fmt" + "io" + "math" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/pkg/sftp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCapacitySFTPWriteLimitsAcrossTransports(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + c := localRelay(t) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + c.TransferConcurrency, c.PayloadBudgetBytes = 4, 256<<20 + r := admittedRun(t, c) + q := r.guests[0] + client, err := sftp.NewClient(q.client) + require.NoError(t, err) + defer func() { _ = client.Close() }() + outside := filepath.Join(t.TempDir(), "outside") + control := filepath.Join(filepath.Dir(q.host.socket), "control-secret") + for _, name := range []string{outside, control} { + require.NoError(t, os.WriteFile(name, []byte("unchanged"), 0600)) + } + t.Run("unassigned name", func(t *testing.T) { + f, err := client.Create("unexpected") + if f != nil { + _ = f.Close() + } + require.Error(t, err, "unassigned upload name was created") + }) + t.Run("extent before IO", func(t *testing.T) { + f, err := client.Create("transfer-0.bin") + require.NoError(t, err) + defer func() { _ = f.Close() }() + for _, offset := range []int64{TransferSize, -1, math.MaxInt64} { + n, err := f.WriteAt([]byte("x"), offset) + assert.Error(t, err, "write outside extent accepted at %d", offset) + assert.Zero(t, n) + } + info, err := f.Stat() + require.NoError(t, err) + assert.Zero(t, info.Size(), "rejected write reached the filesystem") + _, err = f.WriteAt([]byte("z"), TransferSize-1) + require.NoError(t, err, "last legal byte remains writable") + }) + t.Run("one lifetime instance including deleted handles", func(t *testing.T) { + f, err := client.Create("transfer-1.bin") + require.NoError(t, err) + defer func() { _ = f.Close() }() + _, err = f.WriteAt([]byte("old"), 0) + require.NoError(t, err) + other, err := sftp.NewClient(q.client) + require.NoError(t, err) + defer func() { _ = other.Close() }() + duplicate, err := other.Create("transfer-1.bin") + if duplicate != nil { + _ = duplicate.Close() + } + assert.Error(t, err, "second session received a new writable instance") + require.NoError(t, client.Remove("transfer-1.bin")) + duplicate, err = other.Create("transfer-1.bin") + if duplicate != nil { + _ = duplicate.Close() + } + assert.Error(t, err, "remove recycled an outstanding inode's creation slot") + _, err = f.WriteAt([]byte("still bounded"), 3) + require.NoError(t, err) + require.NoError(t, f.Close()) + s, err := q.host.session(r.ctx) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(r.ctx, 3*time.Second) + defer cancel() + connection, raw, err := openGuest(ctx, r.guard, q.target, q.endpoint, s, q.identity, q.proxy) + require.NoError(t, err) + defer func() { _ = raw.Close(); _ = connection.Close() }() + reconnected, err := sftp.NewClient(connection) + require.NoError(t, err) + defer func() { _ = reconnected.Close() }() + duplicate, err = reconnected.Create("transfer-1.bin") + if duplicate != nil { + _ = duplicate.Close() + } + require.Error(t, err, "reconnect replenished the creation slot") + }) + t.Run("concurrent cumulative overwrites", func(t *testing.T) { + f, err := client.Create("transfer-2.bin") + require.NoError(t, err) + defer func() { _ = f.Close() }() + payload := bytes.Repeat([]byte("x"), 1<<20) + var wg sync.WaitGroup + written := make(chan int, 17) + for range 17 { + wg.Go(func() { n, _ := f.WriteAt(payload, 0); written <- n }) + } + wg.Wait() + close(written) + total := 0 + for n := range written { + total += n + } + require.EqualValues(t, TransferSize, total, "overwrites bypassed lifetime accepted-write accounting") + n, err := f.WriteAt([]byte("x"), 0) + require.Error(t, err) + require.Zero(t, n) + }) + t.Run("out of order read stat remove", func(t *testing.T) { + f, err := client.Create("transfer-3.bin") + require.NoError(t, err) + _, err = f.WriteAt([]byte("world"), 6) + require.NoError(t, err) + _, err = f.WriteAt([]byte("hello "), 0) + require.NoError(t, err) + require.NoError(t, f.Close()) + f, err = client.Open("transfer-3.bin") + require.NoError(t, err) + data, err := io.ReadAll(f) + require.NoError(t, err) + require.Equal(t, "hello world", string(data)) + info, err := f.Stat() + require.NoError(t, err) + require.EqualValues(t, 11, info.Size()) + require.NoError(t, f.Close()) + require.NoError(t, client.Remove("transfer-3.bin")) + }) + for _, name := range []string{outside, control} { + data, err := os.ReadFile(name) + require.NoError(t, err) + require.Equal(t, "unchanged", string(data)) + } + }) + } +} + +func TestCapacitySFTPAllowedNameStillUsesRoot(t *testing.T) { + for _, ws := range []int{0, 100} { + t.Run(fmt.Sprintf("ws=%d", ws), func(t *testing.T) { + c := localRelay(t) + c.Active, c.HostWSPercent, c.GuestWSPercent = false, ws, ws + c.TransferConcurrency, c.PayloadBudgetBytes = 4, 256<<20 + r := admittedRun(t, c) + client, err := sftp.NewClient(r.guests[0].client) + require.NoError(t, err) + defer func() { _ = client.Close() }() + outside := filepath.Join(t.TempDir(), "sentinel") + require.NoError(t, os.WriteFile(outside, []byte("outside"), 0600)) + require.NoError(t, os.Symlink(outside, filepath.Join(r.hosts[0].h.CommandDir, "transfer-0.bin"))) + f, err := client.Create("transfer-0.bin") + if f != nil { + _ = f.Close() + } + require.Error(t, err, "an allowed filename followed a symlink outside Root") + f, err = client.Create("transfer-1.bin") + require.NoError(t, err, "allowed in-root writable control") + _, err = f.Write([]byte("inside")) + require.NoError(t, err) + require.NoError(t, f.Close()) + data, err := os.ReadFile(outside) + require.NoError(t, err) + require.Equal(t, "outside", string(data)) + }) + } +} + +func TestCapacitySFTPAssignsOnlyEachHostsTransfers(t *testing.T) { + for _, guestsPerHost := range []int{1, 2} { + t.Run(fmt.Sprintf("guests=%d", guestsPerHost), func(t *testing.T) { + c := localRelay(t) + c.Active, c.Hosts, c.GuestsPerHost = false, 5, guestsPerHost + c.TransferConcurrency, c.PayloadBudgetBytes = 4, 256<<20 + r := admittedRun(t, c) + for hostID := range 5 { + client, err := sftp.NewClient(r.guests[hostID*guestsPerHost].client) + require.NoError(t, err) + for transfer := range 4 { + f, err := client.Create(fmt.Sprintf("transfer-%d.bin", transfer)) + if f != nil { + _ = f.Close() + } + if transfer/guestsPerHost == hostID { + require.NoError(t, err, "assigned transfer %d on host %d", transfer, hostID) + } else { + require.Error(t, err, "unassigned transfer %d on host %d", transfer, hostID) + } + } + require.NoError(t, client.Close()) + } + }) + } +} diff --git a/internal/capacity/whole_review_test.go b/internal/capacity/whole_review_test.go index f73b60533..783ef5771 100644 --- a/internal/capacity/whole_review_test.go +++ b/internal/capacity/whole_review_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "errors" + "fmt" "io" "os" "os/exec" @@ -203,7 +204,7 @@ func TestRunGuestAuthorizationAcrossTransports(t *testing.T) { if err != nil { t.Fatalf("authorized guest SFTP: %v", err) } - path := filepath.Join(q.host.h.CommandDir, "authorized-transfer") + path := filepath.Join(q.host.h.CommandDir, fmt.Sprintf("transfer-%d.bin", q.id)) f, err := client.Create(path) if err != nil { t.Fatal(err) diff --git a/internal/capacity/workload.go b/internal/capacity/workload.go index cbb4d4dfb..cb78ca2ce 100644 --- a/internal/capacity/workload.go +++ b/internal/capacity/workload.go @@ -2,6 +2,7 @@ package capacity import ( "bufio" + "bytes" "context" "errors" "fmt" @@ -35,7 +36,7 @@ func parseFrame(line string) (observation, bool) { if i < 0 { return observation{}, false } - // Scanner removes LF (and an optional preceding CR). Accept the same + // The line reader removes LF (and an optional preceding CR). Accept the same // endings for direct callers without trimming any payload whitespace. payload := strings.TrimSuffix(strings.TrimSuffix(line[i:], "\n"), "\r") fields := strings.Fields(payload) @@ -58,6 +59,95 @@ func parseFrame(line string) (observation, bool) { return observation{fields[1], id, seq, sent}, true } +// frameLines preserves ScanLines' raw 1 MiB limit and error/CR semantics, +// while retaining only the largest canonical payload plus two trailing CRs. +// Replay prefixes can be arbitrarily long within that raw limit; parseFrame +// still selects the final CAP tag and verifies every canonical payload byte. +type frameLines struct { + in io.Reader + read [4096]byte + tail [1025]byte + start, end int + length int + err error + done bool +} + +func (s *frameLines) Scan() bool { + if s.done { + return false + } + s.length = 0 + raw := 0 + for { + if s.start < s.end { + data := s.read[s.start:s.end] + newline := bytes.IndexByte(data, '\n') + if newline >= 0 { + data = data[:newline] + } + raw += len(data) + s.start += len(data) + if len(data) >= len(s.tail) { + s.length = copy(s.tail[:], data[len(data)-len(s.tail):]) + } else { + if excess := s.length + len(data) - len(s.tail); excess > 0 { + s.length = copy(s.tail[:], s.tail[excess:s.length]) + } + s.length += copy(s.tail[s.length:], data) + } + if newline >= 0 { + s.start++ + return s.line() + } + } + // ScanLines offers the final unterminated token even when its reader + // returned data together with a non-EOF error. Report that error next. + if s.err != nil { + s.done = true + return raw > 0 && s.line() + } + if raw == 1<<20 { + s.err, s.done = bufio.ErrTooLong, true + return false + } + // Never read beyond the old Scanner's remaining raw-line allowance: + // exactly-full unterminated input is accepted only with a same-read + // error, whereas a newline may occupy the last allowed byte. + p := s.read[:min(len(s.read), (1<<20)-raw)] + for empty := 0; ; empty++ { + n, err := s.in.Read(p) + if n < 0 || n > len(p) { + s.err = bufio.ErrBadReadCount + break + } + s.start, s.end, s.err = 0, n, err + if n > 0 || err != nil { + break + } + if empty == 100 { + s.err = io.ErrNoProgress + break + } + } + } +} + +func (s *frameLines) line() bool { + if s.length > 0 && s.tail[s.length-1] == '\r' { + s.length-- + } + return true +} + +func (s *frameLines) Text() string { return string(s.tail[:s.length]) } +func (s *frameLines) Err() error { + if s.err == io.EOF { + return nil + } + return s.err +} + // Helper is the runner's harmless hosted command. It echoes fixed 128-byte // input frames and, in active mode, emits exactly 1 KiB every ten seconds. // It never starts a shell or interprets received text as a command. @@ -82,8 +172,7 @@ func helper(ctx context.Context, in io.Reader, out io.Writer, hostID int, active done := make(chan struct{}) go func() { defer close(done) - scanner := bufio.NewScanner(in) - scanner.Buffer(make([]byte, 4096), 1<<20) + scanner := &frameLines{in: in} for scanner.Scan() { line := scanner.Text() i := strings.LastIndex(line, "CAP ") @@ -178,8 +267,7 @@ func (s sequenceTracker) observe(o observation) error { return nil } func scanFrames(in io.Reader, consume func(observation) error, helperReady func(observation)) error { - scanner := bufio.NewScanner(in) - scanner.Buffer(make([]byte, 4096), 1<<20) + scanner := &frameLines{in: in} for scanner.Scan() { if o, ok := parseFrame(scanner.Text()); ok {