Improve frontend-backend connection resiliency (#86)

* Improve frontend backend reconnect handling

* Reconnect after missed stream heartbeats

* Harden frontend stream failure handling

* Harden backend connection lifecycle

* Harden socket reconnect edge cases

* Fix stale frontend connection races

* Add server websocket liveness deadlines

* Harden frontend socket cleanup paths

* Fix socket cleanup review findings

* Address AI review connection feedback

* Finalize already-closing streams during clear

Avoid calling close a second time when Streams.clear sees a stream that is already shutting down. This lets the existing shutdown path finalize with completed() instead of treating the normal already-closing state as a cleanup failure.

Add a regression test that starts local stream shutdown before clearing the transport and verifies the command is completed without reporting a clear error.

* Complete streams after clear close failures

During final stream teardown, preserve command close errors for the cleared callback but still run stream completion callbacks. This releases command subscribers and clears stream slots even when the transport is already broken.

Update the stream clear regression test to assert both error propagation and terminal completion after a command close failure.

* Clear stale socket handler after callback failure

If the connected callback throws after a stream handler is assigned, clear the cached handler and close the dialed connection before reporting failure. This prevents later get() calls from reusing a non-serving handler.

Add a socket regression test that throws from connected(), verifies cleanup, and confirms the next get() opens a fresh stream.

Author: Snuffy2

application/command/handler.go

@@ -21,6 +21,8 @@ import (
 // length exceeds the read buffer.
 // ErrHandlerInvalidControlMessage is returned when a control frame carries zero
 // bytes of payload.
+// ErrHandlerUnknownControlSignal is returned when a control frame carries an
+// unsupported control signal byte.
 var (
 	ErrHandlerUnknownHeaderType = errors.New(
 		"unknown command header type")
@@ -30,6 +32,9 @@ var (
 
 	ErrHandlerInvalidControlMessage = errors.New(
 		"invalid control message")
+
+	ErrHandlerUnknownControlSignal = errors.New(
+		"unknown control signal")
 )
 
 // HandlerCancelSignal is a channel that, when closed or written to, signals
@@ -247,6 +252,9 @@ func (e *Handler) handleControl(d byte, l log.Logger) error {
 		} else {
 			l.Debug("Repeated Resume Stream command, ignore")
 		}
+
+	default:
+		return ErrHandlerUnknownControlSignal
 	}
 
 	return nil

application/command/handler_echo_test.go

@@ -6,6 +6,7 @@ package command
 
 import (
 	"bytes"
+	"errors"
 	"io"
 	"sync"
 	"testing"
@@ -93,3 +94,33 @@ func TestHandlerHandleEcho(t *testing.T) {
 		return
 	}
 }
+
+func TestHandlerRejectsUnknownControlSignal(t *testing.T) {
+	w := dummyWriter{
+		written: make([]byte, 0, 64),
+	}
+	s := []byte{
+		byte(HeaderControl | 1),
+		0xff,
+	}
+	lock := sync.Mutex{}
+	bufferPool := NewBufferPool(4096)
+	handler := newHandler(
+		Configuration{},
+		nil,
+		rw.NewFetchReader(testDummyFetchGen(s)),
+		&w,
+		&lock,
+		0,
+		0,
+		log.NewDitch(),
+		NewHooks(configuration.HookSettings{}),
+		&bufferPool,
+	)
+
+	hErr := handler.Handle()
+
+	if !errors.Is(hErr, ErrHandlerUnknownControlSignal) {
+		t.Fatalf("expected unknown control signal error, got %v", hErr)
+	}
+}

application/commands/ssh.go

@@ -801,6 +801,7 @@ func (d *sshClient) local(
 			if wErr != nil {
 				remote.closer()
 				d.l.Debug("Failed to write data to remote: %s", wErr)
+				return wErr
 			}
 		}
 
@@ -910,12 +911,13 @@ func (d *sshClient) Close() error {
 		d.fingerprintVerifyResultReceiveClosed = true
 	}
 
+	d.baseCtxCancel()
+
 	remote, remoteErr := d.getRemote()
 	if remoteErr == nil {
 		remote.closer()
 	}
 
-	d.baseCtxCancel()
 	d.remoteCloseWait.Wait()
 
 	return nil

application/commands/ssh_lifecycle_test.go

@@ -5,11 +5,15 @@
 package commands
 
 import (
+	"context"
+	"errors"
 	"testing"
+	"time"
 
 	"github.com/Snuffy2/shellport/application/command"
 	"github.com/Snuffy2/shellport/application/configuration"
 	"github.com/Snuffy2/shellport/application/log"
+	"golang.org/x/crypto/ssh"
 )
 
 // TestSSHCommandKeepsBufferPoolScopedToSession verifies that SSH clients retain
@@ -38,3 +42,77 @@ func TestSSHCommandKeepsBufferPoolScopedToSession(t *testing.T) {
 		)
 	}
 }
+
+// TestSSHCloseCancelsBeforeWaitingForRemote verifies Close can unblock remote
+// startup paths that only exit after the base context is cancelled.
+func TestSSHCloseCancelsBeforeWaitingForRemote(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	client := &sshClient{
+		baseCtx:                              ctx,
+		baseCtxCancel:                        cancel,
+		credentialReceive:                    make(chan []byte),
+		fingerprintVerifyResultReceive:       make(chan bool),
+		remoteConnReceive:                    make(chan sshRemoteConn),
+		credentialReceiveClosed:              false,
+		fingerprintVerifyResultReceiveClosed: false,
+	}
+	client.remoteCloseWait.Add(1)
+
+	go func() {
+		<-ctx.Done()
+		close(client.remoteConnReceive)
+		client.remoteCloseWait.Done()
+	}()
+
+	done := make(chan struct{})
+	go func() {
+		_ = client.Close()
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("Close did not cancel base context before waiting for remote")
+	}
+
+	select {
+	case <-done:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("Close did not return after remote shutdown")
+	}
+}
+
+type failingSSHWriter struct {
+	err error
+}
+
+func (w failingSSHWriter) Write(_ []byte) (int, error) {
+	return 0, w.err
+}
+
+// TestSSHLocalReturnsRemoteWriteErrors verifies stdin write failures surface to
+// the stream handler instead of leaving the UI in a misleading connected state.
+func TestSSHLocalReturnsRemoteWriteErrors(t *testing.T) {
+	writeErr := errors.New("remote write failed")
+	closed := false
+	client := &sshClient{
+		l: log.NewDitch(),
+		remoteConn: sshRemoteConn{
+			writer:  failingSSHWriter{err: writeErr},
+			closer:  func() error { closed = true; return nil },
+			session: &ssh.Session{},
+		},
+	}
+	header := command.StreamHeader{}
+	header.Set(SSHClientStdIn, 5)
+
+	err := client.local(nil, newLimitedReader([]byte("hello")), header, make([]byte, 16))
+
+	if !errors.Is(err, writeErr) {
+		t.Fatalf("expected remote write error, got %v", err)
+	}
+	if !closed {
+		t.Fatal("expected remote closer to run after write failure")
+	}
+}

application/commands/telnet.go

@@ -267,19 +267,21 @@ func (d *telnetClient) client(
 		if wErr != nil {
 			remoteConn.Close()
 			d.l.Debug("Failed to write data to remote: %s", wErr)
+			return wErr
 		}
 	}
 
 	return nil
 }
 
 func (d *telnetClient) Close() error {
+	d.baseCtxCancel()
+
 	remoteConn, remoteConnErr := d.getRemote()
 	if remoteConnErr == nil {
 		remoteConn.Close()
 	}
 
-	d.baseCtxCancel()
 	d.closeWait.Wait()
 	return nil
 }

application/commands/telnet_lifecycle_test.go

@@ -5,8 +5,13 @@
 package commands
 
 import (
+	"context"
+	"errors"
+	"io"
+	"net"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/Snuffy2/shellport/application/command"
 	"github.com/Snuffy2/shellport/application/configuration"
@@ -40,3 +45,100 @@ func TestTelnetCommandKeepsBufferPoolScopedToSession(t *testing.T) {
 		)
 	}
 }
+
+// TestTelnetCloseCancelsBeforeWaitingForRemote verifies Close can unblock
+// remote startup paths that only exit after the base context is cancelled.
+func TestTelnetCloseCancelsBeforeWaitingForRemote(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	client := &telnetClient{
+		baseCtx:       ctx,
+		baseCtxCancel: cancel,
+		remoteChan:    make(chan net.Conn),
+	}
+	client.closeWait.Add(1)
+
+	go func() {
+		<-ctx.Done()
+		close(client.remoteChan)
+		client.closeWait.Done()
+	}()
+
+	done := make(chan struct{})
+	go func() {
+		_ = client.Close()
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("Close did not cancel base context before waiting for remote")
+	}
+
+	select {
+	case <-done:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("Close did not return after remote shutdown")
+	}
+}
+
+type failingTelnetConn struct {
+	net.Conn
+	writeErr error
+	closed   bool
+}
+
+func (c *failingTelnetConn) Read(_ []byte) (int, error) {
+	return 0, io.EOF
+}
+
+func (c *failingTelnetConn) Write(_ []byte) (int, error) {
+	return 0, c.writeErr
+}
+
+func (c *failingTelnetConn) Close() error {
+	c.closed = true
+	return nil
+}
+
+func (c *failingTelnetConn) LocalAddr() net.Addr {
+	return nil
+}
+
+func (c *failingTelnetConn) RemoteAddr() net.Addr {
+	return nil
+}
+
+func (c *failingTelnetConn) SetDeadline(_ time.Time) error {
+	return nil
+}
+
+func (c *failingTelnetConn) SetReadDeadline(_ time.Time) error {
+	return nil
+}
+
+func (c *failingTelnetConn) SetWriteDeadline(_ time.Time) error {
+	return nil
+}
+
+// TestTelnetClientReturnsRemoteWriteErrors verifies stdin write failures
+// surface to the stream handler instead of leaving the UI connected to dead IO.
+func TestTelnetClientReturnsRemoteWriteErrors(t *testing.T) {
+	writeErr := errors.New("remote write failed")
+	remote := &failingTelnetConn{writeErr: writeErr}
+	client := &telnetClient{l: log.NewDitch(), remoteConn: remote}
+
+	err := client.client(
+		nil,
+		newLimitedReader([]byte("hello")),
+		command.StreamHeader{},
+		make([]byte, 16),
+	)
+
+	if !errors.Is(err, writeErr) {
+		t.Fatalf("expected remote write error, got %v", err)
+	}
+	if !remote.closed {
+		t.Fatal("expected remote connection to close after write failure")
+	}
+}