From fef7e653cf8b3b4003ccba8760504ed6a32c0bb4 Mon Sep 17 00:00:00 2001 From: jiashuoz Date: Wed, 30 Sep 2026 21:47:19 +0800 Subject: [PATCH] Authorize guest reconnect through one runner control connection --- .../design/2026-09-30-guest-reconnect-host.md | 55 +++ docs/design/2026-09-30-guest-reconnect-rpc.md | 3 + internal/driver/reconnect.go | 23 ++ internal/runnerd/agent.go | 8 +- internal/runnerd/microvm.go | 12 +- internal/runnerd/microvm_test.go | 10 + internal/runnerd/reconnect.go | 166 +++++++++ internal/runnerd/reconnect_process_test.go | 113 ++++++ internal/runnerd/reconnect_test.go | 346 ++++++++++++++++++ internal/runnerd/runnerd.go | 6 + .../runnerd/testdata/reconnect-host/main.go | 50 +++ 11 files changed, 786 insertions(+), 6 deletions(-) create mode 100644 docs/design/2026-09-30-guest-reconnect-host.md create mode 100644 internal/driver/reconnect.go create mode 100644 internal/runnerd/reconnect.go create mode 100644 internal/runnerd/reconnect_process_test.go create mode 100644 internal/runnerd/reconnect_test.go create mode 100644 internal/runnerd/testdata/reconnect-host/main.go diff --git a/docs/design/2026-09-30-guest-reconnect-host.md b/docs/design/2026-09-30-guest-reconnect-host.md new file mode 100644 index 0000000..51387ab --- /dev/null +++ b/docs/design/2026-09-30-guest-reconnect-host.md @@ -0,0 +1,55 @@ +# Runner reconnect authorization callback + +This dependent slice follows the shared [RPC contract](2026-09-30-guest-reconnect-rpc.md). +It implements the optional `driver.GuestReconnectHost` port on runnerd. It does +not advertise a capability or enable the guest listener. The next driver handshake +can call `host.AuthorizeGuestReconnect(ctx, sessionID, prove)`; `prove` exchanges +only the validated challenge for a signature on its bounded guest connection. + +One call captures the accepted agent connection, local placement generation and +registry boot identity. Both RPCs use that connection's existing outbound queue +and runner-origin pending table. Redial never moves an old proof to the new queue. +Before the guest callback, before accept, and before returning authority, check the +original connection and live local placement again. The control plane remains the +source of durable placement, expiry, membership, policy and single-use decisions. +The driver supplies session identity; guest data cannot choose workload scope. +Unknown/zero placement and unaccepted connections refuse. Current cold-resume +placement bookkeeping is insufficient for this guard and must be reconciled before +enabling recovery; no permissive fallback is introduced. + +The whole call has a five-second context budget, at most one in-flight operation +per session and 64 across a runner. Admission fails immediately when full. Slots +and pending RPC entries are released on return. Guest proof callbacks must honor +the context and bound peer/frame resources; the callback runs inline, so a broken +implementation that ignores cancellation retains its slot rather than spawning +unbounded detached work. It must not send configuration or replace a relay. + +Strict shared decoders validate challenge, proof and acceptance. The challenge +must match the expected session, placement and original connection incarnation. +Refusals return a zero acceptance and only invalid/expired/fenced/unavailable; +upstream and guest error text is discarded. A failure after committed acceptance +can strand that attempt: start a fresh challenge, never replay or cache authority. +Successful return still requires fresh configuration and old-relay fencing by the +caller. No token, proof or challenge is logged. + +The generic guest relay now explicitly refuses begin/accept by method name, for +both Docker and microVM guests, in addition to its existing high-bit request and +response fence. Enrollment remains guest-originated. Legacy bootstrap and other +session RPCs preserve their behavior. + +A pair of separately callable begin/accept methods was rejected because callers +could accidentally migrate an attempt across agent reconnection, omit validation +or omit cleanup. A new public RPC/HTTP route was rejected: the existing agent +session-RPC queue and driver host callback already own this boundary. + +Verification covers real WebSocket round trips, wrong scope, malformed replies, +fixed refusals, cancellation, deadline, admission, local replacement and control +redial. A built execution probe calls the port through real runnerd/RunAgent with +a synthetic driver and proof provider. No shipping CLI or listener invokes this +port yet, so this is executable port coverage, not VM or guest continuity evidence. + +Depends on OSS PR113. Before enablement: guest key custody/enrollment, negotiated +capability, strict 4 KiB/five-second guest handshake, bounded peer acceptance, +current configuration delivery, cold lifecycle ordering, relay fencing and real +PID/PTY/agent qualification. Hosted deployment additionally depends on Cloud +PR146/147/148. No RAM persistence is introduced. diff --git a/docs/design/2026-09-30-guest-reconnect-rpc.md b/docs/design/2026-09-30-guest-reconnect-rpc.md index 5d4544c..3673dda 100644 --- a/docs/design/2026-09-30-guest-reconnect-rpc.md +++ b/docs/design/2026-09-30-guest-reconnect-rpc.md @@ -64,3 +64,6 @@ initial and cold-resume enrollment ordering, bounded guest peer handling, configuration refresh, and relay takeover. It must retain the five-second challenge/handshake budget and prove original-socket fencing, replay refusal, tenant isolation and real PID/PTY/agent continuity before advertising support. + +The follow-on [runner host callback](2026-09-30-guest-reconnect-host.md) now supplies +connection-bound begin/proof/accept orchestration without enabling the listener. diff --git a/internal/driver/reconnect.go b/internal/driver/reconnect.go new file mode 100644 index 0000000..b6c4c14 --- /dev/null +++ b/internal/driver/reconnect.go @@ -0,0 +1,23 @@ +package driver + +import ( + "context" + "github.com/tokencanopy/rainier/protocol/runner" +) + +// GuestReconnectProof exchanges a validated challenge with the guest and returns +// its canonical signature. It must honor ctx, bound unauthenticated frames and +// never send configuration or replace the active relay. Errors are not exposed. +type GuestReconnectProof func(context.Context, runner.GuestReconnectChallenge) (string, error) + +// GuestReconnectHost is an optional authorization port for a future reconnect +// listener. Implementing it does not advertise or enable reconnect capability. +type GuestReconnectHost interface { + // AuthorizeGuestReconnect performs challenge and acceptance on one accepted + // control connection, within a single five-second budget. sessionID is the + // driver's assignment, never guest input. Proof is called only after challenge + // scope is checked. Failures return a zero response and one fixed error code: + // invalid, expired, fenced or unavailable. Success is committed authority, not + // permission to skip fresh configuration or old-relay fencing. + AuthorizeGuestReconnect(context.Context, string, GuestReconnectProof) (runner.GuestReconnectAcceptResponse, error) +} diff --git a/internal/runnerd/agent.go b/internal/runnerd/agent.go index b28e558..77302c5 100644 --- a/internal/runnerd/agent.go +++ b/internal/runnerd/agent.go @@ -258,6 +258,10 @@ func (s *Server) agentSession(ctx context.Context, cfg AgentConfig) (established ag := &agentSessionState{} out := make(chan runner.FromRunner, 64) + rc := &reconnectControl{ctx: connCtx, state: ag, out: out} + s.reconnectControl.Store(rc) + defer s.reconnectControl.CompareAndSwap(rc, nil) + send := func(m runner.FromRunner) { cctx, ccancel := context.WithTimeout(ctx, 5*time.Second) m.Used, m.Total, _ = s.drv.Capacity(cctx) // best-effort; piggybacked on every message @@ -607,8 +611,8 @@ func (s *Server) forwardSessionRPC(m runner.ToRunner, send func(runner.FromRunne return } // A response to a request this RUNNER originated stops here: the runner - // is a pure forwarder for everything a sandbox asked, and the one thing - // it asks for itself (a bootstrap token on a cold resume) is answered to + // is a pure forwarder for everything a sandbox asked; cold bootstrap and + // reconnect authorization requested by the runner are answered to // a caller inside this process, not to the guest. The id spaces are // disjoint so the two can share one connection — see // runnerOriginatedIDBase. diff --git a/internal/runnerd/microvm.go b/internal/runnerd/microvm.go index 72efbc7..161c7bb 100644 --- a/internal/runnerd/microvm.go +++ b/internal/runnerd/microvm.go @@ -82,7 +82,9 @@ func isRunnerOriginated(id uint64) bool { return id&runnerOriginatedIDBase != 0 // rather than answering it there. `method` is the envelope's method, which // is "resp" for a response, so both arms consult the same fence. // -// Two things are refused, and only this hop can refuse either of them. +// Runner-only methods and runner-owned IDs are refused at this hop. +// Reconnect begin/accept are also runner-only: a guest may enroll its fresh +// boot but may not originate the host's later authorization flow. // // A sandbox may not mint its own bootstrap token. The design's whole point // is that the token is SINGLE-USE and lives 120 seconds: a guest that could @@ -103,6 +105,8 @@ func isRunnerOriginated(id uint64) bool { return id&runnerOriginatedIDBase != 0 // waiting on that number. func refuseSandboxOrigin(method string, id uint64) string { switch { + case method == runner.MethodBeginGuestReconnect || method == runner.MethodAcceptGuestReconnect: + return "fenced" case method == runner.MethodMintSessionBootstrap: return "a sandbox may not mint its own bootstrap token; a fresh one is minted by the runner on a cold resume" case isRunnerOriginated(id): @@ -114,8 +118,8 @@ func refuseSandboxOrigin(method string, id uint64) string { // runnerRPCTable is the pending table for requests this runner originated. // // It is deliberately tiny and deliberately separate from the sandbox's: the -// runner asks the control plane exactly one thing (a bootstrap token, on a -// cold resume), and a table shared with the forwarding path would make the +// runner asks for cold bootstrap and guest reconnect authorization, and a +// table shared with the forwarding path would make the // forwarder a participant in the conversation it is supposed to be a pipe // for. type runnerRPCTable struct { @@ -191,7 +195,7 @@ const mintBootstrapTimeout = 10 * time.Second // MintSessionBootstrap asks the control plane for a fresh single-use token // for sessionID, on the runner's own control connection. // -// It is the one request this runner originates on the session RPC. The +// This cold-resume request originates inside the runner. The // message carries no session id of its own — FromRunner.Session is the id, // and the control plane answers from the row its placement guard read — which // is what stops a runner holding session A from minting for session B. diff --git a/internal/runnerd/microvm_test.go b/internal/runnerd/microvm_test.go index be214a8..16b3e05 100644 --- a/internal/runnerd/microvm_test.go +++ b/internal/runnerd/microvm_test.go @@ -264,6 +264,16 @@ func TestASandboxMayNotMintItsOwnBootstrapToken(t *testing.T) { ev relay.ControlEvent wantReason string }{ + { + name: "guest begin reconnect", + ev: relay.ControlEvent{Kind: "req:" + runner.MethodBeginGuestReconnect, ID: 5, Payload: []byte(`{"protocol":1}`)}, + wantReason: "fenced", + }, + { + name: "guest accept reconnect", + ev: relay.ControlEvent{Kind: "req:" + runner.MethodAcceptGuestReconnect, ID: 6, Payload: []byte(`{"protocol":1}`)}, + wantReason: "fenced", + }, { name: "the mint method itself", ev: relay.ControlEvent{Kind: "req:" + runner.MethodMintSessionBootstrap, ID: 7, Payload: []byte(`{"protocol":1}`)}, diff --git a/internal/runnerd/reconnect.go b/internal/runnerd/reconnect.go new file mode 100644 index 0000000..1598f2c --- /dev/null +++ b/internal/runnerd/reconnect.go @@ -0,0 +1,166 @@ +package runnerd + +import ( + "context" + "encoding/json" + "errors" + "strconv" + "time" + + "github.com/tokencanopy/rainier/internal/driver" + "github.com/tokencanopy/rainier/protocol/runner" +) + +const reconnectBudget = 5 * time.Second +const maxGuestReconnects = 64 + +var ( + errReconnectInvalid = errors.New("invalid") + errReconnectExpired = errors.New("expired") + errReconnectFenced = errors.New("fenced") + errReconnectUnavailable = errors.New("unavailable") +) + +// A snapshot of one agent connection. out belongs only to this connection; +// a late request can never land on the replacement's queue. +type reconnectControl struct { + ctx context.Context + state *agentSessionState + out chan<- runner.FromRunner +} + +var _ driver.GuestReconnectHost = (*Server)(nil) + +// AuthorizeGuestReconnect implements driver.GuestReconnectHost. It does not +// deliver configuration, attach a guest, or advertise reconnect capability. +func (s *Server) AuthorizeGuestReconnect(ctx context.Context, id string, prove driver.GuestReconnectProof) (runner.GuestReconnectAcceptResponse, error) { + var zero runner.GuestReconnectAcceptResponse + if prove == nil { + return zero, errReconnectInvalid + } + row, ok := s.reg.snapshot(id) + if !ok || row.placementGen == 0 || (row.state != "running" && row.state != "starting") { + return zero, errReconnectFenced + } + rc := s.reconnectControl.Load() + if rc == nil || rc.ctx.Err() != nil || rc.state.generation.Load() == 0 { + return zero, errReconnectUnavailable + } + generation := rc.state.generation.Load() + if !s.claimReconnect(id) { + return zero, errReconnectUnavailable + } + defer s.releaseReconnect(id) + ctx, cancel := context.WithTimeout(ctx, reconnectBudget) + defer cancel() + stop := context.AfterFunc(rc.ctx, cancel) + defer stop() + valid := func() bool { + current, exists := s.reg.snapshot(id) + return ctx.Err() == nil && rc.ctx.Err() == nil && s.reconnectControl.Load() == rc && rc.state.generation.Load() == generation && exists && current.placementGen == row.placementGen && current.boot == row.boot && current.handle == row.handle && (current.state == "running" || current.state == "starting") + } + if !valid() { + return zero, errReconnectFenced + } + payload, err := s.reconnectCall(ctx, rc, id, runner.MethodBeginGuestReconnect, runner.GuestReconnectBeginRequest{Protocol: runner.GuestReconnectProtocol}) + if err != nil { + return zero, err + } + if !valid() { + return zero, errReconnectFenced + } + challenge, err := runner.DecodeGuestReconnectChallenge(payload) + if err != nil || challenge.SessionID != id || challenge.PlacementGeneration != row.placementGen || challenge.HostIncarnation != strconv.FormatUint(generation, 10) { + return zero, errReconnectInvalid + } + signature, err := prove(ctx, challenge) + if err != nil { + return zero, errReconnectUnavailable + } + if !valid() { + return zero, errReconnectFenced + } + request := runner.GuestReconnectAcceptRequest{Protocol: runner.GuestReconnectProtocol, AttemptID: challenge.AttemptID, Signature: signature} + body, _ := json.Marshal(request) + if _, err := runner.DecodeGuestReconnectAcceptRequest(body); err != nil { + return zero, errReconnectInvalid + } + payload, err = s.reconnectCall(ctx, rc, id, runner.MethodAcceptGuestReconnect, request) + if err != nil { + return zero, err + } + if !valid() { + return zero, errReconnectFenced + } + accepted, err := runner.DecodeGuestReconnectAcceptResponse(payload) + if err != nil { + return zero, errReconnectInvalid + } + return accepted, nil +} + +func (s *Server) claimReconnect(id string) bool { + s.reconnectMu.Lock() + defer s.reconnectMu.Unlock() + if _, ok := s.reconnecting[id]; ok || len(s.reconnecting) >= maxGuestReconnects { + return false + } + if s.reconnecting == nil { + s.reconnecting = make(map[string]struct{}) + } + s.reconnecting[id] = struct{}{} + return true +} +func (s *Server) releaseReconnect(id string) { + s.reconnectMu.Lock() + defer s.reconnectMu.Unlock() + delete(s.reconnecting, id) +} + +func (s *Server) reconnectCall(ctx context.Context, rc *reconnectControl, session, method string, request any) ([]byte, error) { + if ctx.Err() != nil || rc.ctx.Err() != nil { + return nil, errReconnectUnavailable + } + body, err := json.Marshal(request) + if err != nil { + return nil, errReconnectInvalid + } + id, ch := s.runnerRPC.begin(session) + defer s.runnerRPC.end(id) + msg := runner.FromRunner{Type: "session_req", Session: session, RPC: &runner.RPCEnvelope{ID: id, Method: method, Payload: body}} + msg.Used, msg.Total, _ = s.drv.Capacity(ctx) + msg.Active, msg.IdleExited = s.reg.counts() + if ctx.Err() != nil || rc.ctx.Err() != nil { + return nil, errReconnectUnavailable + } + select { + case rc.out <- msg: + default: + return nil, errReconnectUnavailable + } + select { + case <-ctx.Done(): + return nil, errReconnectUnavailable + case answer := <-ch: + if ctx.Err() != nil || rc.ctx.Err() != nil { + return nil, errReconnectUnavailable + } + if answer.OK { + return answer.Payload, nil + } + refusal, err := runner.DecodeGuestReconnectErrorResponse(answer.Payload) + if err != nil { + return nil, errReconnectUnavailable + } + switch refusal.Error { + case "invalid": + return nil, errReconnectInvalid + case "expired": + return nil, errReconnectExpired + case "fenced": + return nil, errReconnectFenced + default: + return nil, errReconnectUnavailable + } + } +} diff --git a/internal/runnerd/reconnect_process_test.go b/internal/runnerd/reconnect_process_test.go new file mode 100644 index 0000000..9c48c46 --- /dev/null +++ b/internal/runnerd/reconnect_process_test.go @@ -0,0 +1,113 @@ +package runnerd + +import ( + "bufio" + "bytes" + "context" + "crypto/ed25519" + "encoding/base64" + "github.com/tokencanopy/rainier/protocol/runner" + "io" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" +) + +// No shipping command invokes the optional callback yet. Build its public call +// site and exercise that artifact across a real agent WebSocket; do not imply +// that a live VM, guest key owner or relay takeover was exercised. +func TestReconnectHostBuiltExecution(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + t.Cleanup(cancel) + binary := filepath.Join(t.TempDir(), "reconnect-host") + build := exec.CommandContext(ctx, "go", "build", "-o", binary, "./testdata/reconnect-host") + if output, err := build.CombinedOutput(); err != nil { + t.Fatalf("build probe: %v %s", err, output) + } + fc := newFakeControld(t, testToken) + cmd := exec.CommandContext(ctx, binary, fc.wsURL()) + input, err := cmd.StdinPipe() + if err != nil { + t.Fatal(err) + } + output, err := cmd.StdoutPipe() + if err != nil { + t.Fatal(err) + } + var logs bytes.Buffer + cmd.Stderr = &logs + if err := cmd.Start(); err != nil { + t.Fatal(err) + } + answers := make(chan string, 8) + go func() { + scan := bufio.NewScanner(output) + for scan.Scan() { + answers <- scan.Text() + } + close(answers) + }() + waited := make(chan error, 1) + go func() { waited <- cmd.Wait() }() + t.Cleanup(func() { + input.Close() + select { + case err := <-waited: + if err != nil { + t.Error("probe failed") + } + case <-time.After(3 * time.Second): + cancel() + <-waited + t.Error("probe did not stop") + } + if strings.Contains(logs.String(), "panic:") || strings.Contains(logs.String(), strings.Repeat("A", 43)) { + t.Error("probe log failure or cryptographic disclosure") + } + }) + conn := fc.nextConn(t) + conn.readAnnounce(t) + acceptReconnectPeer(t, conn, 7) + conn.send(t, runner.ToRunner{Type: "create", ReqID: 1, Session: "session_test", PlacementGeneration: 3, Spec: &runner.Spec{Image: "synthetic.invalid/session:test"}}) + if result := conn.readMsg(t); result.Type != "result" || !result.OK { + t.Fatal("probe placement failed") + } + for _, tc := range []struct{ name, want string }{{"valid", "accepted"}, {"wrong_scope", "invalid"}, {"revoked", "fenced"}} { + t.Run(tc.name, func(t *testing.T) { + if _, err := io.WriteString(input, "authorize\n"); err != nil { + t.Fatal(err) + } + req := conn.readMsg(t) + if req.RPC == nil || req.RPC.Method != runner.MethodBeginGuestReconnect || !isRunnerOriginated(req.RPC.ID) { + t.Fatal("invalid begin") + } + switch tc.name { + case "revoked": + replyReconnect(t, conn, req, false, []byte(`{"error":"fenced"}`)) + case "wrong_scope": + replyReconnect(t, conn, req, true, []byte(strings.Replace(string(reconnectChallengeJSON()), "session_test", "wrong_test", 1))) + default: + replyReconnect(t, conn, req, true, reconnectChallengeJSON()) + req = conn.readMsg(t) + proof, err := runner.DecodeGuestReconnectAcceptRequest(req.RPC.Payload) + challenge, _ := runner.DecodeGuestReconnectChallenge(reconnectChallengeJSON()) + public := ed25519.NewKeyFromSeed(make([]byte, 32)).Public().(ed25519.PublicKey) + if err != nil || req.RPC.Method != runner.MethodAcceptGuestReconnect || challenge.VerifyProof(base64.RawURLEncoding.EncodeToString(public), proof.Signature) != nil { + t.Fatal("invalid built-process proof") + } + replyReconnect(t, conn, req, true, []byte(`{"epoch":2,"token":"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA","expires_in_sec":120}`)) + } + select { + case got := <-answers: + if got != tc.want { + t.Fatalf("result=%s want %s", got, tc.want) + } + case <-time.After(6 * time.Second): + t.Fatal("no probe result") + } + }) + } + t.Log("built host callback: accepted proof, wrong-scope and revoked authority paths passed; synthetic driver/guest, real agent WebSocket") +} diff --git a/internal/runnerd/reconnect_test.go b/internal/runnerd/reconnect_test.go new file mode 100644 index 0000000..a6dffd9 --- /dev/null +++ b/internal/runnerd/reconnect_test.go @@ -0,0 +1,346 @@ +package runnerd + +import ( + "context" + "crypto/ed25519" + "crypto/rand" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/tokencanopy/rainier/protocol/runner" +) + +func TestReconnectHostRoundTrip(t *testing.T) { + s, _ := testMicrovmServer(t) + s.reg.put("session_test", &sessionEntry{id: "session_test", state: "running", placementGen: 3}) + fc := newFakeControld(t, testToken) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go s.RunAgent(ctx, AgentConfig{ControldURL: fc.wsURL(), Token: testToken, RunnerName: "runner_test"}) + conn := fc.nextConn(t) + conn.readAnnounce(t) + conn.send(t, runner.ToRunner{Type: "accept", Generation: 7}) + // A command/result barrier proves the accept has run before the host call. + conn.send(t, runner.ToRunner{Type: "destroy", ReqID: 99, Session: "absent_test"}) + conn.readMsg(t) + public, private, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatal(err) + } + done := make(chan error, 1) + go func() { + accepted, err := s.AuthorizeGuestReconnect(ctx, "session_test", func(ctx context.Context, c runner.GuestReconnectChallenge) (string, error) { + msg, err := c.SigningMessage() + if err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(ed25519.Sign(private, msg)), nil + }) + if err == nil && (accepted.Epoch != 2 || accepted.ExpiresInSec != 120) { + err = context.Canceled + } + done <- err + }() + req := conn.readMsg(t) + if req.Type != "session_req" || req.Session != "session_test" || req.RPC == nil || req.RPC.Method != runner.MethodBeginGuestReconnect || !isRunnerOriginated(req.RPC.ID) { + t.Fatal("wrong begin request") + } + c := runner.GuestReconnectChallenge{Protocol: 1, SessionID: "session_test", BootEpoch: "boot_test", HostIncarnation: "7", AttemptID: "attempt_test", PlacementGeneration: 3, Challenge: base64.RawURLEncoding.EncodeToString(make([]byte, 32))} + body, _ := json.Marshal(c) + conn.send(t, runner.ToRunner{Type: "session_rpc", Session: req.Session, RPC: &runner.RPCEnvelope{ID: req.RPC.ID, Method: "resp", OK: true, Payload: body}}) + req = conn.readMsg(t) + if req.RPC == nil || req.RPC.Method != runner.MethodAcceptGuestReconnect || !isRunnerOriginated(req.RPC.ID) { + t.Fatal("wrong accept request") + } + proof, err := runner.DecodeGuestReconnectAcceptRequest(req.RPC.Payload) + if err != nil || proof.AttemptID != c.AttemptID || c.VerifyProof(base64.RawURLEncoding.EncodeToString(public), proof.Signature) != nil { + t.Fatal("proof was not preserved") + } + body = []byte(`{"epoch":2,"token":"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA","expires_in_sec":120}`) + conn.send(t, runner.ToRunner{Type: "session_rpc", Session: req.Session, RPC: &runner.RPCEnvelope{ID: req.RPC.ID, Method: "resp", OK: true, Payload: body}}) + select { + case err := <-done: + if err != nil { + t.Fatalf("authorization failed: %v", err) + } + case <-time.After(time.Second): + t.Fatal("no authorization result") + } +} + +func reconnectPeer(t *testing.T) (*Server, *fakeControld, *fakeConn) { + t.Helper() + s, _ := testMicrovmServer(t) + s.reg.put("session_test", &sessionEntry{id: "session_test", handle: "vm_test", state: "running", placementGen: 3}) + fc := newFakeControld(t, testToken) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + _ = s.RunAgent(ctx, AgentConfig{ControldURL: fc.wsURL(), Token: testToken, RunnerName: "runner_test", RedialBackoffMin: time.Millisecond}) + }() + t.Cleanup(func() { + cancel() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Error("agent did not stop") + } + }) + conn := fc.nextConn(t) + conn.readAnnounce(t) + acceptReconnectPeer(t, conn, 7) + return s, fc, conn +} +func acceptReconnectPeer(t *testing.T, c *fakeConn, generation uint64) { + t.Helper() + c.send(t, runner.ToRunner{Type: "accept", Generation: generation}) + c.send(t, runner.ToRunner{Type: "destroy", ReqID: 99, Session: "absent_test"}) + if c.readMsg(t).Type != "result" { + t.Fatal("accept barrier failed") + } +} +func reconnectChallengeJSON() []byte { + return []byte(`{"protocol":1,"session_id":"session_test","boot_epoch":"boot_test","host_incarnation":"7","attempt_id":"attempt_test","placement_generation":3,"challenge":"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"}`) +} +func replyReconnect(t *testing.T, c *fakeConn, r runner.FromRunner, ok bool, body []byte) { + t.Helper() + if r.RPC == nil { + t.Fatal("missing request") + } + c.send(t, runner.ToRunner{Type: "session_rpc", Session: r.Session, RPC: &runner.RPCEnvelope{ID: r.RPC.ID, Method: "resp", OK: ok, Payload: body}}) +} +func waitReconnectError(t *testing.T, done <-chan error, want string) { + t.Helper() + select { + case err := <-done: + if err == nil || err.Error() != want { + t.Fatalf("refusal=%v want %s", err, want) + } + case <-time.After(2 * time.Second): + t.Fatal("authorization did not refuse") + } +} +func TestReconnectHostRejectsWrongChallengeBeforeProof(t *testing.T) { + for _, tc := range []struct{ name, from, to string }{ + {"session", "session_test", "other_test"}, {"host", `"7"`, `"8"`}, {"placement", `"placement_generation":3`, `"placement_generation":4`}, + {"unknown", "{", `{"extra":true,`}, {"duplicate", "{", `{"protocol":1,`}, {"oversize", "boot_test", strings.Repeat("b", 4096)}, + } { + t.Run(tc.name, func(t *testing.T) { + s, _, c := reconnectPeer(t) + var called atomic.Bool + done := make(chan error, 1) + go func() { + out, err := s.AuthorizeGuestReconnect(context.Background(), "session_test", func(context.Context, runner.GuestReconnectChallenge) (string, error) { + called.Store(true) + return strings.Repeat("A", 86), nil + }) + if out != (runner.GuestReconnectAcceptResponse{}) { + t.Error("failure returned authority") + } + done <- err + }() + req := c.readMsg(t) + replyReconnect(t, c, req, true, []byte(strings.Replace(string(reconnectChallengeJSON()), tc.from, tc.to, 1))) + waitReconnectError(t, done, "invalid") + if called.Load() { + t.Fatal("unvalidated challenge reached guest") + } + }) + } +} +func TestReconnectHostChecksProofAndAcceptance(t *testing.T) { + for _, tc := range []struct { + name, signature, body, want string + ok bool + }{ + {"bad_signature", "secret.test", "", "invalid", false}, + {"raw_error", strings.Repeat("A", 86), `{"error":"secret.test"}`, "unavailable", false}, + {"fenced", strings.Repeat("A", 86), `{"error":"fenced"}`, "fenced", false}, + {"expired", strings.Repeat("A", 86), `{"error":"expired"}`, "expired", false}, + {"invalid", strings.Repeat("A", 86), `{"error":"invalid"}`, "invalid", false}, + {"bad_success", strings.Repeat("A", 86), `{"epoch":1,"token":"secret.test","expires_in_sec":120}`, "invalid", true}, + } { + t.Run(tc.name, func(t *testing.T) { + s, _, c := reconnectPeer(t) + done := make(chan error, 1) + go func() { + out, err := s.AuthorizeGuestReconnect(context.Background(), "session_test", func(context.Context, runner.GuestReconnectChallenge) (string, error) { return tc.signature, nil }) + if out != (runner.GuestReconnectAcceptResponse{}) { + t.Error("failure returned authority") + } + done <- err + }() + replyReconnect(t, c, c.readMsg(t), true, reconnectChallengeJSON()) + if tc.body != "" { + replyReconnect(t, c, c.readMsg(t), tc.ok, []byte(tc.body)) + } + waitReconnectError(t, done, tc.want) + }) + } +} +func TestReconnectHostDoesNotMigrateProofToReplacementSocket(t *testing.T) { + s, fc, c := reconnectPeer(t) + entered := make(chan struct{}) + release := make(chan struct{}) + defer close(release) + done := make(chan error, 1) + go func() { + out, err := s.AuthorizeGuestReconnect(context.Background(), "session_test", func(ctx context.Context, _ runner.GuestReconnectChallenge) (string, error) { + close(entered) + select { + case <-ctx.Done(): + case <-release: + } + return strings.Repeat("A", 86), nil + }) + if out != (runner.GuestReconnectAcceptResponse{}) { + t.Error("stale proof returned authority") + } + done <- err + }() + replyReconnect(t, c, c.readMsg(t), true, reconnectChallengeJSON()) + <-entered + c.c.CloseNow() + replacement := fc.nextConn(t) + replacement.readAnnounce(t) + acceptReconnectPeer(t, replacement, 8) + waitReconnectError(t, done, "fenced") + // A current call still works; its first request must be begin, never the old proof. + retry := make(chan error, 1) + go func() { + _, err := s.AuthorizeGuestReconnect(context.Background(), "session_test", func(context.Context, runner.GuestReconnectChallenge) (string, error) { + return "", errors.New("guest_secret.test") + }) + retry <- err + }() + req := replacement.readMsg(t) + if req.RPC == nil || req.RPC.Method != runner.MethodBeginGuestReconnect { + t.Fatal("old proof moved to new socket") + } + body := []byte(strings.Replace(string(reconnectChallengeJSON()), `"7"`, `"8"`, 1)) + replyReconnect(t, replacement, req, true, body) + waitReconnectError(t, retry, "unavailable") +} +func TestReconnectHostBoundsAttemptsAndCleansUp(t *testing.T) { + s, _, c := reconnectPeer(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + prove := func(context.Context, runner.GuestReconnectChallenge) (string, error) { + return "", errors.New("guest_error.test") + } + go func() { _, err := s.AuthorizeGuestReconnect(ctx, "session_test", prove); done <- err }() + c.readMsg(t) + if out, err := s.AuthorizeGuestReconnect(ctx, "session_test", prove); err == nil || err.Error() != "unavailable" || out != (runner.GuestReconnectAcceptResponse{}) { + t.Fatal("parallel attempt admitted") + } + cancel() + waitReconnectError(t, done, "unavailable") + retry := make(chan error, 1) + go func() { _, err := s.AuthorizeGuestReconnect(context.Background(), "session_test", prove); retry <- err }() + replyReconnect(t, c, c.readMsg(t), false, []byte(`{"error":"expired"}`)) + waitReconnectError(t, retry, "expired") + s.runnerRPC.mu.Lock() + pending := len(s.runnerRPC.waiting) + s.runnerRPC.mu.Unlock() + if pending != 0 { + t.Fatal("pending RPC leaked") + } +} +func TestReconnectHostRefusesReplacedInstance(t *testing.T) { + s, _, c := reconnectPeer(t) + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + done := make(chan error, 1) + go func() { + _, err := s.AuthorizeGuestReconnect(ctx, "session_test", func(context.Context, runner.GuestReconnectChallenge) (string, error) { + s.reg.put("session_test", &sessionEntry{id: "session_test", handle: "replacement_test", state: "running", placementGen: 3}) + return strings.Repeat("A", 86), nil + }) + done <- err + }() + replyReconnect(t, c, c.readMsg(t), true, reconnectChallengeJSON()) + waitReconnectError(t, done, "fenced") +} + +func TestReconnectHostGlobalAdmissionBound(t *testing.T) { + s, _, c := reconnectPeer(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 64) + prove := func(context.Context, runner.GuestReconnectChallenge) (string, error) { return "", nil } + // Pin the documented production limit independently of its constant. + for i := 0; i < 64; i++ { + id := fmt.Sprintf("session_%d_test", i) + s.reg.put(id, &sessionEntry{id: id, state: "running", placementGen: 3}) + go func() { _, err := s.AuthorizeGuestReconnect(ctx, id, prove); done <- err }() + req := c.readMsg(t) + if req.RPC == nil || req.RPC.Method != runner.MethodBeginGuestReconnect { + t.Fatal("expected pending begin") + } + } + if out, err := s.AuthorizeGuestReconnect(ctx, "session_test", prove); err == nil || err.Error() != "unavailable" || out != (runner.GuestReconnectAcceptResponse{}) { + t.Fatal("global admission limit exceeded") + } + cancel() + for i := 0; i < 64; i++ { + waitReconnectError(t, done, "unavailable") + } + s.reconnectMu.Lock() + active := len(s.reconnecting) + s.reconnectMu.Unlock() + if active != 0 { + t.Fatal("admission slots leaked") + } +} +func TestReconnectHostDeadlineAndLateProof(t *testing.T) { + s, _, c := reconnectPeer(t) + done := make(chan error, 1) + go func() { + out, err := s.AuthorizeGuestReconnect(context.Background(), "session_test", func(ctx context.Context, _ runner.GuestReconnectChallenge) (string, error) { + deadline, ok := ctx.Deadline() + if !ok || time.Until(deadline) > 5*time.Second { + t.Error("unbounded proof context") + } + <-ctx.Done() + return strings.Repeat("A", 86), nil + }) + if out != (runner.GuestReconnectAcceptResponse{}) { + t.Error("late proof returned authority") + } + done <- err + }() + replyReconnect(t, c, c.readMsg(t), true, reconnectChallengeJSON()) + select { + case err := <-done: + if err == nil || err.Error() != "fenced" { + t.Fatalf("late proof error=%v", err) + } + case <-time.After(6 * time.Second): + t.Fatal("handshake exceeded budget") + } +} +func TestReconnectHostRequiresPlacementAndAcceptedConnection(t *testing.T) { + s, _ := testMicrovmServer(t) + prove := func(context.Context, runner.GuestReconnectChallenge) (string, error) { + t.Error("unauthorized proof callback") + return "", nil + } + if out, err := s.AuthorizeGuestReconnect(context.Background(), "absent_test", prove); err == nil || err.Error() != "fenced" || out != (runner.GuestReconnectAcceptResponse{}) { + t.Fatal("missing session admitted") + } + s.reg.put("session_test", &sessionEntry{id: "session_test", state: "running", placementGen: 3}) + if _, err := s.AuthorizeGuestReconnect(context.Background(), "session_test", prove); err == nil || err.Error() != "unavailable" { + t.Fatal("missing connection admitted") + } + if _, err := s.AuthorizeGuestReconnect(context.Background(), "session_test", nil); err == nil || err.Error() != "invalid" { + t.Fatal("nil proof accepted") + } +} diff --git a/internal/runnerd/runnerd.go b/internal/runnerd/runnerd.go index 4220fd4..ef0ed32 100644 --- a/internal/runnerd/runnerd.go +++ b/internal/runnerd/runnerd.go @@ -26,6 +26,12 @@ import ( ) type Server struct { + // reconnectControl is a connection-scoped view of the same agent RPC queue. + // A handshake captures it once; reconnecting never migrates a pending proof. + reconnectControl atomic.Pointer[reconnectControl] + reconnectMu sync.Mutex + reconnecting map[string]struct{} + drv driver.Driver reg *registry dialBase string // e.g. ws://runnerd:8080 — what sessiond dials to register diff --git a/internal/runnerd/testdata/reconnect-host/main.go b/internal/runnerd/testdata/reconnect-host/main.go new file mode 100644 index 0000000..caa5711 --- /dev/null +++ b/internal/runnerd/testdata/reconnect-host/main.go @@ -0,0 +1,50 @@ +// A built execution probe for the optional host callback, not a shipping command. +// The driver and guest proof provider are synthetic; RunAgent and the host callback +// are real. No capability or local RPC endpoint is enabled by this probe. +package main + +import ( + "bufio" + "context" + "crypto/ed25519" + "encoding/base64" + "fmt" + "github.com/tokencanopy/rainier/internal/driver" + "github.com/tokencanopy/rainier/internal/runnerd" + "github.com/tokencanopy/rainier/protocol/runner" + "os" +) + +func main() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + s := runnerd.New(driver.NewFake(4), "", "", "") + done := make(chan struct{}) + go func() { + defer close(done) + _ = s.RunAgent(ctx, runnerd.AgentConfig{ControldURL: os.Args[1], Token: "testtoken", RunnerName: "runner_process_test"}) + }() + private := ed25519.NewKeyFromSeed(make([]byte, 32)) + input := bufio.NewScanner(os.Stdin) + for input.Scan() { + if input.Text() != "authorize" { + break + } + out, err := s.AuthorizeGuestReconnect(ctx, "session_test", func(ctx context.Context, c runner.GuestReconnectChallenge) (string, error) { + message, err := c.SigningMessage() + if err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(ed25519.Sign(private, message)), nil + }) + if err != nil { + fmt.Println(err.Error()) + } else if out.Epoch == 2 && out.ExpiresInSec == 120 { + fmt.Println("accepted") + } else { + fmt.Println("unexpected") + } + } + cancel() + <-done +}