From ceadbebcb24bccb1a9a2a5c7a30c17e0ef59cb3e Mon Sep 17 00:00:00 2001 From: Alex Florenti Date: Fri, 16 Oct 2026 17:12:54 +0200 Subject: [PATCH] Bound shell transport close and disconnect bridge on stdin EOF --- .../app/shell/lifecycle_regression_test.go | 121 ++++++++++++++++++ internal/app/shell/shell.go | 17 ++- internal/term/eof_regression_test.go | 80 ++++++++++++ internal/term/term.go | 16 ++- 4 files changed, 224 insertions(+), 10 deletions(-) create mode 100644 internal/app/shell/lifecycle_regression_test.go create mode 100644 internal/term/eof_regression_test.go diff --git a/internal/app/shell/lifecycle_regression_test.go b/internal/app/shell/lifecycle_regression_test.go new file mode 100644 index 0000000..b4a5acb --- /dev/null +++ b/internal/app/shell/lifecycle_regression_test.go @@ -0,0 +1,121 @@ +package shell + +import ( + "encoding/json" + "errors" + "sync" + "testing" + "time" +) + +// This isolated transport blocks writes until explicitly released or closed. +// No account, remote agent, network socket, or saved credentials are used. +type diagnosticBlockedConn struct { + entered chan struct{} + release chan struct{} + closed chan struct{} + enterOnce sync.Once + closeOnce sync.Once +} + +func (f *diagnosticBlockedConn) SendText([]byte) error { + f.enterOnce.Do(func() { close(f.entered) }) + select { + case <-f.release: + return nil + case <-f.closed: + return errors.New("test transport closed") + } +} +func (f *diagnosticBlockedConn) Read() ([]byte, error) { + <-f.closed + return nil, errors.New("test transport closed") +} +func (f *diagnosticBlockedConn) Close() error { + f.closeOnce.Do(func() { close(f.closed) }) + return nil +} + +func TestCloseAbortsBlockedInput(t *testing.T) { + f := &diagnosticBlockedConn{entered: make(chan struct{}), release: make(chan struct{}), closed: make(chan struct{})} + s := &Shell{sock: f, termID: 1, out: make(chan []byte, 64), done: make(chan struct{})} + inputDone, closeDone := make(chan struct{}), make(chan struct{}) + go func() { _ = s.Input("pwd\r"); close(inputDone) }() + <-f.entered + go func() { _ = s.Close(); close(closeDone) }() + select { + case <-f.closed: + t.Log("Close interrupted the blocked input transport") + case <-time.After(closeGracePeriod + time.Second): + close(f.release) + t.Fatal("Close did not interrupt the blocked input within its grace period") + } + close(f.release) + select { + case <-inputDone: + case <-time.After(time.Second): + t.Fatal("input cleanup failed") + } + select { + case <-closeDone: + case <-time.After(time.Second): + t.Fatal("close cleanup failed") + } +} + +func TestCloseAbortsStalledClosingFrame(t *testing.T) { + f := &diagnosticBlockedConn{entered: make(chan struct{}), release: make(chan struct{}), closed: make(chan struct{})} + s := &Shell{sock: f, termID: 1, out: make(chan []byte, 64), done: make(chan struct{})} + done := make(chan struct{}) + go func() { _ = s.Close(); close(done) }() + <-f.entered + select { + case <-f.closed: + t.Log("Close interrupted its own stalled closing frame") + case <-time.After(closeGracePeriod + time.Second): + close(f.release) + t.Fatal("Close did not abort its stalled closing frame within its grace period") + } + close(f.release) + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("close cleanup failed") + } +} + +func TestNormalCloseRequestsTerminalAndAppCleanup(t *testing.T) { + s, f := newTestShell() + if err := s.Close(); err != nil { + t.Fatal(err) + } + frames := drain(f) + if len(frames) != 2 { + t.Fatalf("expected terminal close plus app termination, got %v", frames) + } + var terminal struct { + ID int `json:"id"` + Type int `json:"type"` + } + if err := json.Unmarshal([]byte(frames[0]), &terminal); err != nil { + t.Fatal(err) + } + if terminal.ID != s.termID || terminal.Type != reqClose { + t.Fatalf("unexpected terminal close: %s", frames[0]) + } + var app struct { + Type string `json:"type"` + } + if err := json.Unmarshal([]byte(frames[1]), &app); err != nil { + t.Fatal(err) + } + if app.Type != "term" { + t.Fatalf("unexpected app termination: %s", frames[1]) + } + if err := s.Close(); err != nil { + t.Fatal(err) + } + if extra := drain(f); len(extra) != 0 { + t.Fatalf("repeated Close sent extra frames: %v", extra) + } +} diff --git a/internal/app/shell/shell.go b/internal/app/shell/shell.go index bd393b1..e582b47 100644 --- a/internal/app/shell/shell.go +++ b/internal/app/shell/shell.go @@ -17,11 +17,12 @@ import ( // message type codes (client → server). const ( - reqOpen = 0 - reqClose = 1 - reqInput = 2 - reqChangeSize = 3 - keepAlivePeriod = 30 * time.Second + reqOpen = 0 + reqClose = 1 + reqInput = 2 + reqChangeSize = 3 + keepAlivePeriod = 30 * time.Second + closeGracePeriod = 2 * time.Second ) // maxMessageBytes caps the size of a single client → server message. @@ -193,9 +194,13 @@ func (s *Shell) Resize(cols, rows int) error { func (s *Shell) Close() error { s.once.Do(func() { close(s.done) + // Closing frames are best effort. A pending or stalled socket write must + // not prevent transport closure: closing the connection unblocks it. + abort := time.AfterFunc(closeGracePeriod, func() { _ = s.sock.Close() }) + defer abort.Stop() + defer s.sock.Close() _ = s.send(map[string]any{"id": s.termID, "type": reqClose}) _ = s.send(map[string]any{"type": "term"}) - s.sock.Close() }) return nil } diff --git a/internal/term/eof_regression_test.go b/internal/term/eof_regression_test.go new file mode 100644 index 0000000..606a38f --- /dev/null +++ b/internal/term/eof_regression_test.go @@ -0,0 +1,80 @@ +package term + +import ( + "context" + "errors" + "os" + "testing" + "time" +) + +type diagnosticIdleRemote struct{ out chan []byte } + +func (r *diagnosticIdleRemote) Input(string) error { return nil } +func (r *diagnosticIdleRemote) Resize(int, int) error { return nil } +func (r *diagnosticIdleRemote) Output() <-chan []byte { return r.out } + +func TestStdinEOFDisconnects(t *testing.T) { + reader, writer, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + _ = writer.Close() // immediate real EOF on a non-interactive stdin + original := os.Stdin + os.Stdin = reader + defer func() { os.Stdin = original; reader.Close() }() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + err = Bridge(ctx, &diagnosticIdleRemote{out: make(chan []byte)}) + if err != nil { + t.Fatalf("stdin EOF should disconnect cleanly before context deadline, got %v", err) + } + t.Log("Real stdin EOF disconnected cleanly") +} + +type failingInputRemote struct { + diagnosticIdleRemote + err error +} + +func (r *failingInputRemote) Input(string) error { return r.err } + +func TestStdinInputFailureDisconnects(t *testing.T) { + reader, writer, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer reader.Close() + defer writer.Close() + if _, err := writer.Write([]byte("pwd\n")); err != nil { + t.Fatal(err) + } + original := os.Stdin + os.Stdin = reader + defer func() { os.Stdin = original }() + want := errors.New("input transport failed") + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + err = Bridge(ctx, &failingInputRemote{diagnosticIdleRemote{out: make(chan []byte)}, want}) + if !errors.Is(err, want) { + t.Fatalf("expected input failure to leave Bridge, got %v", err) + } +} + +func TestStdinReadFailureDisconnects(t *testing.T) { + reader, writer, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + _ = reader.Close() + _ = writer.Close() + original := os.Stdin + os.Stdin = reader + defer func() { os.Stdin = original }() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + err = Bridge(ctx, &diagnosticIdleRemote{out: make(chan []byte)}) + if !errors.Is(err, os.ErrClosed) { + t.Fatalf("expected closed stdin error to leave Bridge, got %v", err) + } +} diff --git a/internal/term/term.go b/internal/term/term.go index 3a0f37b..e1b33c9 100644 --- a/internal/term/term.go +++ b/internal/term/term.go @@ -80,6 +80,8 @@ const resizePoll = 250 * time.Millisecond // stdin reaches EOF. When stdin is a TTY it is switched to raw mode and restored // on return. func Bridge(ctx context.Context, r Remote) error { + ctx, cancel := context.WithCancel(ctx) + defer cancel() // Stop resize polling on EOF, escape, and input errors too. inFd := int(os.Stdin.Fd()) if xterm.IsTerminal(inFd) { old, err := xterm.MakeRaw(inFd) @@ -89,7 +91,7 @@ func Bridge(ctx context.Context, r Remote) error { } // stdin → remote input, with the SSH-style "~." disconnect escape. - quit := make(chan struct{}) + quit := make(chan error, 1) go func() { buf := make([]byte, 4096) esc := newEscapeScanner() @@ -99,15 +101,21 @@ func Bridge(ctx context.Context, r Remote) error { out, disc := esc.feed(buf[:n]) if len(out) > 0 { if werr := r.Input(string(out)); werr != nil { + quit <- werr return } } if disc { - close(quit) + quit <- nil return } } if err != nil { + if errors.Is(err, io.EOF) { + quit <- nil + } else { + quit <- err + } return } } @@ -138,8 +146,8 @@ func Bridge(ctx context.Context, r Remote) error { select { case <-ctx.Done(): return ctx.Err() - case <-quit: // user typed the "~." disconnect escape - return nil + case err := <-quit: // escape, stdin EOF, or input failure + return err case chunk, ok := <-out: if !ok { return nil