diff --git a/host/host.go b/host/host.go index 4f27f3170..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 @@ -325,6 +331,16 @@ 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 + // 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. @@ -600,13 +616,60 @@ 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) + } + 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") @@ -1246,6 +1309,8 @@ func (c *Host) Run(ctx context.Context) (runErr error) { InitialClientTimeout: c.InitialClientTimeout, StopGrace: c.StopGrace, SFTPDisabled: c.SFTPDisabled, + SFTPRoot: c.SFTPRoot, + SFTPWriteLimits: sftpWriteLimits, PermissionChecker: c.PermissionChecker, TerminalConsent: c.TerminalConsent, SFTPAutoAccept: c.AutoAccept || c.SFTPAutoAccept, @@ -1263,6 +1328,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..a357683d9 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,10 @@ 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 + 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. @@ -219,6 +224,16 @@ 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") + } + } + 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") } @@ -375,9 +390,14 @@ func (s *Server) ServeWithContext(ctx context.Context, guest, host net.Listener) kind: kindGuest, cmdDone: cmdDone, commandResult: cmd.Result, + sftpRoot: s.SFTPRoot, + sftpWriteLimits: sftpWriteLimits, onPrivateCommandExit: s.onPrivateCommandExit, } + if s.SFTPRoot != nil { + sh.sftpWorkers = &s.sftpWorkers + } ss := []gssh.Signer{s.HostKey} @@ -390,9 +410,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 +505,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 +1110,9 @@ type sessionHandler struct { // SFTP configuration permissionChecker consent.PermissionChecker sftpAutoAccept bool + sftpRoot *os.Root + sftpWriteLimits *rootedSFTPWriteLimits + 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..3fd8cbd8b 100644 --- a/host/internal/sftp.go +++ b/host/internal/sftp.go @@ -39,10 +39,20 @@ type SFTPSession struct { logger *slog.Logger ctx context.Context autoAccept bool + root *os.Root + writeLimits *rootedSFTPWriteLimits } // 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 +90,8 @@ func (h *sessionHandler) HandleSFTP(sess gssh.Session) { permissionChecker: h.permissionChecker, clientInfo: clientInfo, logger: h.logger, + root: h.sftpRoot, + writeLimits: h.sftpWriteLimits, } handlers := sftp.Handlers{ @@ -89,10 +101,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 +269,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 +366,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 +427,15 @@ 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 { + 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) } @@ -411,6 +446,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 +505,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 +600,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_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/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_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/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 new file mode 100644 index 000000000..1a6c895c9 --- /dev/null +++ b/internal/capacity/README.md @@ -0,0 +1,320 @@ +# 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 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 +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. 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 +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 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 + +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 +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 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. +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 +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 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; +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 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 +at phase expiry are a failure. + +`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 +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. 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 +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. + +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 +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 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 +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..a47ddf1d7 --- /dev/null +++ b/internal/capacity/admission_test.go @@ -0,0 +1,35 @@ +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) + if err := <-finished; err != nil { + t.Error(err) + } + }() + 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/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/canary.go b/internal/capacity/canary.go new file mode 100644 index 000000000..5c4d8f0d3 --- /dev/null +++ b/internal/capacity/canary.go @@ -0,0 +1,78 @@ +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 { + 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() + 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.queueElapsed + 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, "queue_ms": float64(q.queueElapsed) / float64(time.Millisecond), "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..660736c4d --- /dev/null +++ b/internal/capacity/churn_test.go @@ -0,0 +1,218 @@ +package capacity + +import ( + "bytes" + "encoding/json" + "io" + "os" + "path/filepath" + "reflect" + "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 + captureChurnTerminals(t, &c) + 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) + } + 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) + } + } +} + +// 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/client.go b/internal/capacity/client.go new file mode 100644 index 000000000..920313d95 --- /dev/null +++ b/internal/capacity/client.go @@ -0,0 +1,464 @@ +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) { + check, err := g.guestKey(ctx, t, u.Scheme, s) + if err != nil { + return nil, nil, err + } + var raw net.Conn + if u.Scheme == "ssh" { + 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))) + 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: check} + client, err := guestHandshake(ctx, raw, u.Host, cfg) + return client, raw, err +} + +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 + inputStart [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 + 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() + } + 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 + } + if err = r.guard.session(ctx, s); err != nil { + return err + } + 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() + 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 + } + 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 + } + 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 + } + 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{}) + q.done = append(q.done, done) + go func() { + defer close(done) + tracker := sequenceTracker{} + epochs := map[string]inputEpoch{} + live := false + baseline := time.Time{} + hostKey := fmt.Sprintf("H:%d", q.host.id) + err := scanFrames(out, func(o observation) error { + 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 + } + 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 { + 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 + } + r.actual(int64(frameLength(o.kind))) + r.rec.mu.Lock() + r.rec.summary.ReceivedFrames++ + acknowledged := false + if o.kind == "I" && o.id == stream { + 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 acknowledged && o.seq == probeSeq { + select { + case ready <- struct{}{}: + default: + } + } + return nil + }, func(ack observation) { + r.actual(128) + 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: + } + }) + 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 = q.helperChallenge(attempt, r, in, stream, probeSeq, ch, helperReady); err != nil { + return 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 + } + 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, "ready_ms": float64(time.Since(serviceBegan)) / float64(time.Millisecond), "queue_ms": float64(q.queueElapsed) / float64(time.Millisecond)}) + return nil +} +func frameLength(kind string) int { + if kind == "H" { + return 1024 + } + return 128 +} +func (q *guest) active(r *runState) error { + q.mu.Lock() + expected := q.expected + q.mu.Unlock() + if expected { + 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") + } + w, seq := q.writers[0], q.seq + q.seq++ + 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 + } + 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 + if public { + raw = t.PublicSSHURL + } + if index < total*percent/100 { + raw = t.WSURL + if public { + raw = t.PublicWSURL + } + } + 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/client_test.go b/internal/capacity/client_test.go new file mode 100644 index 000000000..b46e691ea --- /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 func() { _ = 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..d9b592be9 --- /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 func() { _ = m.Close() }() + p, err := os.Open(*phase) + if err != nil { + return err + } + defer func() { _ = p.Close() }() + cfg, err := capacity.Decode(m, p) + if err != nil { + return err + } + if *validate { + _, 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/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 new file mode 100644 index 000000000..8d3ddf46a --- /dev/null +++ b/internal/capacity/config.go @@ -0,0 +1,277 @@ +// 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"` + // 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) +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 < 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) + probes := int64(c.Hosts*c.GuestsPerHost*c.ChannelsPerGuest) * 512 * fanout + if c.idleHostOnly() { + probes += int64(1+(c.DurationSeconds-1)/30) * 768 + } + 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 + 52 + 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..309a19593 --- /dev/null +++ b/internal/capacity/config_test.go @@ -0,0 +1,142 @@ +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") + } + } +} + +// 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") + } + }) + } +} diff --git a/internal/capacity/fallback_cert_test.go b/internal/capacity/fallback_cert_test.go new file mode 100644 index 000000000..0013c7386 --- /dev/null +++ b/internal/capacity/fallback_cert_test.go @@ -0,0 +1,252 @@ +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{ + "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), + "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 func() { _ = 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 func() { _ = 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/fault_helper_test.go b/internal/capacity/fault_helper_test.go new file mode 100644 index 000000000..e44f8c126 --- /dev/null +++ b/internal/capacity/fault_helper_test.go @@ -0,0 +1,159 @@ +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 config.Mode == "geometry" { + return runGeometryHelper(config.Replay, in, out) + } + 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 new file mode 100644 index 000000000..b212c3cca --- /dev/null +++ b/internal/capacity/forwarded_cert_test.go @@ -0,0 +1,356 @@ +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, options ...func(*server.Server)) 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} + 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 err := <-done: + _ = err // Shutdown owns the expected serving error. + 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{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 +// 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/geometry_test.go b/internal/capacity/geometry_test.go new file mode 100644 index 000000000..411502513 --- /dev/null +++ b/internal/capacity/geometry_test.go @@ -0,0 +1,238 @@ +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 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() + found := geometryHasMarker(o.data.Bytes(), marker) + o.mu.Unlock() + if found { + return + } + select { + case <-o.changed: + case <-timer.C: + 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) + } + } +} + +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/guard.go b/internal/capacity/guard.go new file mode 100644 index 000000000..c191cf994 --- /dev/null +++ b/internal/capacity/guard.go @@ -0,0 +1,451 @@ +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 + gate chan struct{} + active int + peak int +} + +func newGuard(c Config) *guard { + 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) + 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 +} + +// 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 }) + } +} +func (g *guard) permit(ctx context.Context) (*dialPermit, error) { + began := time.Now() + select { + case g.gate <- struct{}{}: + case <-ctx.Done(): + 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 + } + 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 permit == nil { + permit, err = g.permit(ctx) + if 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() + // 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-- + 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 (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, 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(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 err == nil || !isCert { + return err + } + return fallback(addr, remote, key) + }, nil + } + } + } + return pinnedKey(entry), nil +} + +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{} + 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) { + 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{}), 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, external bool) bool { + p.mu.Lock() + defer p.mu.Unlock() + if p.closing { + _ = c.Close() + return false + } + p.conns[c] = external + 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() (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() + } + 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", http.StatusMethodNotAllowed) + return + } + 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", http.StatusServiceUnavailable) + 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", http.StatusForbidden) + return + } + 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) + downstream, rw, err := w.(http.Hijacker).Hijack() + if err != nil { + return + } + if !p.track(downstream, false) { + 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) + // 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() + <-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/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/mixed_test.go b/internal/capacity/mixed_test.go new file mode 100644 index 000000000..f2241e9b3 --- /dev/null +++ b/internal/capacity/mixed_test.go @@ -0,0 +1,97 @@ +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 + captureChurnTerminals(t, &c) + 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/report.go b/internal/capacity/report.go new file mode 100644 index 000000000..79fe94ef7 --- /dev/null +++ b/internal/capacity/report.go @@ -0,0 +1,110 @@ +package capacity + +import ( + "encoding/json" + "errors" + "io" + "sync" + "time" +) + +type Summary struct { + 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 { + 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() + 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) + if r.err != nil && r.fail != nil { + r.fail(r.err) + } + } +} +func newRecorder(out io.Writer, c Config) *recorder { + 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/review_regression_test.go b/internal/capacity/review_regression_test.go new file mode 100644 index 000000000..b58d7759d --- /dev/null +++ b/internal/capacity/review_regression_test.go @@ -0,0 +1,575 @@ +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 := 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) + 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 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() + 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") + } 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.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 { + t.Fatalf("generator queue consumed host service deadline: %v", err) + } + u, 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) + 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) + 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 TestLateFinalReplyIsDrainedAndMeasured(t *testing.T) { + c := faultHelper(t, "late", "") + 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 := faultHelper(t, "missing-host-suffix", "") + 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 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 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}} + 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) + captureChurnTerminals(t, &c) + 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 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) + _ = 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 := faultHelper(t, "replay", replay) + 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] + r.guestIdentity = signer(t) + 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: 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) + } + 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 := 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 { + 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/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/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 new file mode 100644 index 000000000..253b3b5e7 --- /dev/null +++ b/internal/capacity/runner.go @@ -0,0 +1,781 @@ +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" + "github.com/owenthereal/upterm/internal/termsize" + "golang.org/x/crypto/ssh" +) + +const finalDrainLimit = 10 * time.Second +const retryQueueLimit = 5 * time.Second + +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 + 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 + 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 + started bool +} + +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) { + conn, err := host.AdminConn(h.socket) + if err != nil { + return nil, err + } + defer func() { _ = 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) + switch kind { + case "I": + q.mu.Lock() + if q.pending == nil { + q.pending = map[string]time.Time{} + } + q.pending[key] = now + q.mu.Unlock() + case "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 { + r.rec.mu.Lock() + if kind == "I" { + r.rec.summary.InputFrames++ + } else { + r.rec.summary.HandshakeFrames++ + } + 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 { + 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 tunnelObservation, 16)} + if identity == nil { + identity, err = makeSigner() + if err != nil { + return err + } + } + key, err := makeSigner() + if err != nil { + return err + } + 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}) + } + 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 + } + 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, + 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 + } + h.mu.Lock() + h.initial = s + h.mu.Unlock() + return nil + }, + SessionReadyCallback: func(status string) { + h.mu.Lock() + h.established = status == "ready" + h.mu.Unlock() + proxy.ready() + select { + case ready <- status: + default: + } + }, + TunnelCallback: func(e host.TunnelEvent) { + r.hostTunnel(h, e) + }, + } + 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)}) + var began time.Time + h.started = true + 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)) + } + }() + // 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: + 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), "queue_ms": float64(permit.queued) / float64(time.Millisecond)}) + 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) + 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) + 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 + } + 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: r.guestIdentity, 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)-q.queueElapsed) / float64(time.Millisecond), "queue_ms": float64(q.queueElapsed) / 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.activeWorkers.Wait() + r.workers.Wait() + for _, q := range r.guests { + q.setExpected(true) + q.close() + } + 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 + } + _, 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() + 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() + } + 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 { + 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() + 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) + 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 == "" { + 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 = 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}) + 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.NewTimer(time.Duration(c.DurationSeconds) * time.Second) + defer steady.Stop() + defer func() { + rec.mu.Lock() + 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() + } + r.activeWorkload(steadyStarted, &r.activeWorkers) + 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) + } + }() + } + for { + select { + case <-runctx.Done(): + err = context.Cause(runctx) + rec.mu.Lock() + 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 <-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{} + 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..3a060ddc3 --- /dev/null +++ b/internal/capacity/runner_test.go @@ -0,0 +1,326 @@ +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" + "runtime" + "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) { + if len(os.Args) > 1 && os.Args[1] == "helper" { + os.Exit(runFaultHelperProcess()) + } + 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() { + 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) + } + }) + 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 err := <-done: + _ = err // Shutdown owns the expected serving error. + 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 func() { _ = 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/sftp_confinement_test.go b/internal/capacity/sftp_confinement_test.go new file mode 100644 index 000000000..09435e21c --- /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("transfer-0.bin") + 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, "transfer-0.bin") + 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("transfer-0.bin") + 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_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/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/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 +} diff --git a/internal/capacity/transfer.go b/internal/capacity/transfer.go new file mode 100644 index 000000000..1085975c8 --- /dev/null +++ b/internal/capacity/transfer.go @@ -0,0 +1,443 @@ +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 func() { _ = c.Close() }() + 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 + 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) + 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 + } + return func() error { + err := repair() + if err != nil { + r.fail(err) + } + return err + }, nil + }); 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() + 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 + 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 { + 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, "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) +} + +// 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) { + var began time.Time + 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() + } + // 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) + } + 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 := 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) 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 { + guests = append(guests, q) + } + } + ctx, cancel := context.WithTimeout(r.ctx, 15*time.Second) + defer cancel() + down := false + for { + select { + case observation := <-h.tunnel: + e := observation.event + 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(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 + } + } + return nil + } + case <-ctx.Done(): + return fmt.Errorf("host restoration: %w", ctx.Err()) + } + } +} + +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 { + 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 < 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) + } + } + 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/whole_review_test.go b/internal/capacity/whole_review_test.go new file mode 100644 index 000000000..783ef5771 --- /dev/null +++ b/internal/capacity/whole_review_test.go @@ -0,0 +1,269 @@ +package capacity + +import ( + "bytes" + "context" + "errors" + "fmt" + "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, input, output string }{ + {"silent", "RRQII", "AABI-"}, + {"late", "RRQIIF", "AABIIB"}, + {"retry", "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, faultHelperExecutable(t, tc.name, ""), "helper", "0", "false") + 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.TransferConcurrency = 4 + c.PayloadBudgetBytes = 256 << 20 + 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, fmt.Sprintf("transfer-%d.bin", q.id)) + 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 func() { _ = raw.Close() }() + } + if err == nil { + defer func() { _ = 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") + } + } + }) + } +} diff --git a/internal/capacity/workload.go b/internal/capacity/workload.go new file mode 100644 index 000000000..cb78ca2ce --- /dev/null +++ b/internal/capacity/workload.go @@ -0,0 +1,353 @@ +package capacity + +import ( + "bufio" + "bytes" + "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) { + // 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 + } + // 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) + 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]) + seq, e2 := strconv.Atoi(fields[3]) + ns, e3 := strconv.ParseInt(fields[4], 10, 64) + if e1 != nil || e2 != nil || e3 != nil { + return observation{}, false + } + 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 +} + +// 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. +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 { + return err + } + defer func() { _ = 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) + scanner := &frameLines{in: in} + 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() + if closer, ok := in.(io.Closer); ok { + _ = closer.Close() + } + <-done + }() + 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(): + return ctx.Err() + case err := <-failed: + if errors.Is(err, io.EOF) { + return nil + } + 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' + } + if _, err := out.Write(b); err != nil { + return err + } + case now := <-tick.C: + if err := emitDue(now); err != nil { + return err + } + } + } +} + +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 := &frameLines{in: in} + 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 +} + +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) { + r.clientsMu.RUnlock() + return errors.New("unknown input source") + } + peer := r.guests[o.id/4] + r.clientsMu.RUnlock() + peer.mu.Lock() + generation, expected, first := peer.generation, peer.expected, peer.inputStart[o.id%4] + peer.mu.Unlock() + key := "I:" + strconv.Itoa(o.id) + if expected { + delete(tracker, key) + delete(epochs, key) + return nil + } + 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 + } + } + return tracker.observe(o) +} diff --git a/internal/capacity/workload_test.go b/internal/capacity/workload_test.go new file mode 100644 index 000000000..c84e2fdfc --- /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, inputStart: [4]int{4}} + r := &runState{guests: []*guest{peer}} + tracker := sequenceTracker{"I:0": 1} + 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) + } + 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()) + } +}