Wish Ssh Code Review

Wish SSH Code Review

Quick Reference

Review gates

Run these in order when producing a written review. Do not claim a defect in a later step until the Pass when for the current step is satisfied for the code under review.

1. Locate Wish entry points — Pass when: you have at least one repo path per server surface that calls wish.NewServer, wish.WithMiddleware, registers bubbletea.Middleware, or defines the top-level ssh.Handler chain (list the paths explicitly). 2. Capture server-setup evidence — Pass when: for each path from step 1, you have the actual wish.WithHostKey* / host-key configuration and the full middleware list in source order as written (not recalled from memory). If graceful shutdown exists, note the file(s) where ListenAndServe and Shutdown run. 3. Capture session / TUI evidence — Pass when: for each teaHandler (or equivalent), you have noted from source whether s.Pty() is checked before using window size, and whether per-session renderers (bubbletea.MakeRenderer) are used where Lipgloss styles apply. 4. Write findings — Pass when: each finding uses [FILE:LINE] ISSUE_TITLE (line range allowed where needed) and points to the relevant row in Quick Reference (or the matching section in references/).

Review Checklist

Use alongside Review gates; for a written review, complete the gates first so each item below can be tied to cited source.

• [ ] Host keys are loaded from file or generated securely
• [ ] Middleware order is correct (logging first, auth early)
• [ ] Session context is used for per-connection state
• [ ] Graceful shutdown handles active sessions
• [ ] PTY requests are handled for terminal apps
• [ ] Connection limits prevent resource exhaustion
• [ ] Timeout middleware prevents hung connections
• [ ] BubbleTea middleware correctly configured

Critical Patterns

Server Setup

// GOOD - complete server setup
s, err := wish.NewServer(
    wish.WithAddress(fmt.Sprintf("%s:%d", host, port)),
    wish.WithHostKeyPath(".ssh/id_ed25519"),
    wish.WithMiddleware(
        logging.Middleware(),       // first: log all connections
        activeterm.Middleware(),    // handle terminal sizing
        bubbletea.Middleware(teaHandler),
    ),
)
if err != nil {
    return fmt.Errorf("creating server: %w", err)
}

Graceful Shutdown

// BAD - abrupt shutdown
log.Fatal(s.ListenAndServe())

// GOOD - graceful shutdown
done := make(chan os.Signal, 1)
signal.Notify(done, os.Interrupt, syscall.SIGTERM)

go func() {
    if err := s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
        log.Error("server error", "error", err)
    }
}()

<-done
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
if err := s.Shutdown(ctx); err != nil {
    log.Error("shutdown error", "error", err)
}

BubbleTea Handler

func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
    pty, _, _ := s.Pty()

    model := NewModel(pty.Window.Width, pty.Window.Height)

    return model, []tea.ProgramOption{
        tea.WithAltScreen(),
        tea.WithMouseCellMotion(),
    }
}

When to Load References

• Reviewing server initialization → server.md
• Reviewing authentication, session state → sessions.md

Review Questions

1. Are host keys handled securely? 2. Is middleware order correct? 3. Is graceful shutdown implemented? 4. Are PTY window sizes passed to the TUI? 5. Are connection timeouts configured?

Server Setup

Host Key Management

1. Use Persistent Keys

// BAD - generates new key each start (fingerprint changes)
s, err := wish.NewServer(
    wish.WithAddress(":22"),
    // no host key specified - generates random
)

// GOOD - load from file
s, err := wish.NewServer(
    wish.WithAddress(":22"),
    wish.WithHostKeyPath("/data/ssh_host_ed25519_key"),
)

// GOOD - generate if missing, persist for reuse
func ensureHostKey(path string) error {
    if _, err := os.Stat(path); os.IsNotExist(err) {
        _, priv, err := ed25519.GenerateKey(rand.Reader)
        if err != nil {
            return err
        }
        // save to file...
    }
    return nil
}

2. Support Multiple Key Types

s, err := wish.NewServer(
    wish.WithAddress(":22"),
    wish.WithHostKeyPath("/data/ssh_host_ed25519_key"),
    wish.WithHostKeyPEM(rsaKeyBytes),  // additional key type
)

Middleware Configuration

1. Correct Middleware Order

// Middleware executes in order - first added runs first
wish.WithMiddleware(
    // 1. Logging - see all connections
    logging.Middleware(),

    // 2. Timeout - prevent hung connections
    wish.WithIdleTimeout(10*time.Minute),
    wish.WithMaxTimeout(30*time.Minute),

    // 3. Active terminal - handle PTY/window sizing
    activeterm.Middleware(),

    // 4. Your app handler - BubbleTea or custom
    bubbletea.Middleware(teaHandler),
)

2. Custom Middleware

func customMiddleware() wish.Middleware {
    return func(next ssh.Handler) ssh.Handler {
        return func(s ssh.Session) {
            // Before handling
            log.Info("connection", "user", s.User(), "remote", s.RemoteAddr())

            // Call next handler
            next(s)

            // After handling (session ended)
            log.Info("disconnected", "user", s.User())
        }
    }
}

3. Metrics Middleware

func metricsMiddleware(metrics *Metrics) wish.Middleware {
    return func(next ssh.Handler) ssh.Handler {
        return func(s ssh.Session) {
            metrics.ActiveConnections.Inc()
            start := time.Now()

            defer func() {
                metrics.ActiveConnections.Dec()
                metrics.SessionDuration.Observe(time.Since(start).Seconds())
            }()

            next(s)
        }
    }
}

Server Lifecycle

1. Graceful Shutdown

func run() error {
    s, err := wish.NewServer(...)
    if err != nil {
        return err
    }

    // Start server in goroutine
    errCh := make(chan error, 1)
    go func() {
        errCh <- s.ListenAndServe()
    }()

    // Wait for shutdown signal
    quit := make(chan os.Signal, 1)
    signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

    select {
    case err := <-errCh:
        return err
    case <-quit:
    }

    // Graceful shutdown with timeout
    ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
    defer cancel()

    return s.Shutdown(ctx)
}

2. Health Checks

// Run HTTP health endpoint alongside SSH
http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusOK)
    w.Write([]byte("ok"))
})

go http.ListenAndServe(":8080", nil)

Connection Handling

1. Connection Limits

// Limit concurrent connections
var connLimiter = make(chan struct{}, 100)

func connectionLimitMiddleware() wish.Middleware {
    return func(next ssh.Handler) ssh.Handler {
        return func(s ssh.Session) {
            select {
            case connLimiter <- struct{}{}:
                defer func() { <-connLimiter }()
                next(s)
            default:
                s.Exit(1)
            }
        }
    }
}

2. Rate Limiting

import "golang.org/x/time/rate"

var limiter = rate.NewLimiter(rate.Every(time.Second), 10)  // 10/sec

func rateLimitMiddleware() wish.Middleware {
    return func(next ssh.Handler) ssh.Handler {
        return func(s ssh.Session) {
            if !limiter.Allow() {
                io.WriteString(s, "Too many connections, try again later\n")
                s.Exit(1)
                return
            }
            next(s)
        }
    }
}

Anti-Patterns

1. No Error Handling on ListenAndServe

// BAD
go s.ListenAndServe()

// GOOD
go func() {
    if err := s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
        log.Fatal("server error", "error", err)
    }
}()

2. Ignoring Context in Shutdown

// BAD - no timeout
s.Shutdown(context.Background())  // could hang forever

// GOOD
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
s.Shutdown(ctx)

Review Questions

1. Are host keys persisted (not regenerated on restart)? 2. Is middleware order correct (logging first)? 3. Is graceful shutdown implemented with timeout? 4. Are connection/rate limits in place? 5. Is there a health check endpoint?

Sessions & Security

Session Handling

1. Access Session Info

func handler(s ssh.Session) {
    // User info
    user := s.User()
    remoteAddr := s.RemoteAddr()

    // Public key (if key auth)
    key := s.PublicKey()

    // Environment variables
    env := s.Environ()

    // Command (if not interactive)
    cmd := s.Command()

    // PTY info (if allocated)
    pty, winCh, isPty := s.Pty()
    if isPty {
        width := pty.Window.Width
        height := pty.Window.Height
    }
}

2. Handle Window Resize

func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
    pty, winCh, _ := s.Pty()

    model := NewModel(pty.Window.Width, pty.Window.Height)

    // Window change channel is passed via activeterm middleware
    // BubbleTea handles this automatically when using bubbletea.Middleware

    return model, []tea.ProgramOption{tea.WithAltScreen()}
}

3. Session Context for State

// Store per-session state using context
type contextKey string

const sessionDataKey contextKey = "sessionData"

type SessionData struct {
    User      string
    ConnectAt time.Time
    PageViews int
}

func sessionMiddleware() wish.Middleware {
    return func(next ssh.Handler) ssh.Handler {
        return func(s ssh.Session) {
            data := &SessionData{
                User:      s.User(),
                ConnectAt: time.Now(),
            }
            ctx := context.WithValue(s.Context(), sessionDataKey, data)
            // Note: wish.Session doesn't expose SetContext
            // Store in sync.Map keyed by session ID instead
            next(s)
        }
    }
}

Security

1. Authentication

// Public key authentication
wish.WithPublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
    // Check against authorized keys
    authorized := loadAuthorizedKeys()
    for _, authKey := range authorized {
        if ssh.KeysEqual(key, authKey) {
            return true
        }
    }
    return false
}),

// Password authentication (not recommended for production)
wish.WithPasswordAuth(func(ctx ssh.Context, password string) bool {
    // Never do this - use public key auth
    return password == os.Getenv("SSH_PASSWORD")
}),

2. Authorization

func authorizationMiddleware(allowedUsers map[string]bool) wish.Middleware {
    return func(next ssh.Handler) ssh.Handler {
        return func(s ssh.Session) {
            if !allowedUsers[s.User()] {
                io.WriteString(s, "Access denied\n")
                s.Exit(1)
                return
            }
            next(s)
        }
    }
}

3. Secure Defaults

s, err := wish.NewServer(
    wish.WithAddress(":22"),
    wish.WithHostKeyPath("./host_key"),

    // Timeouts prevent hung connections
    wish.WithIdleTimeout(10*time.Minute),
    wish.WithMaxTimeout(60*time.Minute),

    // Require public key auth
    wish.WithPublicKeyAuth(authHandler),

    wish.WithMiddleware(
        logging.Middleware(),  // audit trail
        activeterm.Middleware(),
        bubbletea.Middleware(teaHandler),
    ),
)

BubbleTea Integration

1. Basic Handler

func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
    pty, _, _ := s.Pty()

    renderer := bubbletea.MakeRenderer(s)
    model := NewModel(renderer, pty.Window.Width, pty.Window.Height)

    return model, []tea.ProgramOption{
        tea.WithAltScreen(),
    }
}

2. Passing Session to Model

type Model struct {
    renderer *lipgloss.Renderer
    user     string
    width    int
    height   int
}

func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
    pty, _, _ := s.Pty()
    renderer := bubbletea.MakeRenderer(s)

    model := Model{
        renderer: renderer,
        user:     s.User(),
        width:    pty.Window.Width,
        height:   pty.Window.Height,
    }

    return model, []tea.ProgramOption{tea.WithAltScreen()}
}

3. Per-Session Styles

// Each session needs its own renderer for correct color detection
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
    renderer := bubbletea.MakeRenderer(s)

    // Create styles with session's renderer
    styles := NewStyles(renderer)

    model := Model{
        styles: styles,
    }

    return model, nil
}

type Styles struct {
    Title lipgloss.Style
    Item  lipgloss.Style
}

func NewStyles(r *lipgloss.Renderer) Styles {
    return Styles{
        Title: r.NewStyle().Bold(true).Foreground(lipgloss.Color("205")),
        Item:  r.NewStyle().PaddingLeft(2),
    }
}

Anti-Patterns

1. Ignoring PTY

// BAD - assumes PTY always exists
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
    pty, _, _ := s.Pty()  // may be nil!
    model := NewModel(pty.Window.Width, pty.Window.Height)  // panic!
}

// GOOD - handle non-PTY connections
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
    pty, _, hasPty := s.Pty()

    width, height := 80, 24  // sensible defaults
    if hasPty {
        width = pty.Window.Width
        height = pty.Window.Height
    }

    model := NewModel(width, height)
    return model, nil
}

2. Global Lipgloss Styles

// BAD - global styles don't detect terminal capabilities per-session
var titleStyle = lipgloss.NewStyle().Bold(true)

// GOOD - per-session renderer
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
    renderer := bubbletea.MakeRenderer(s)
    titleStyle := renderer.NewStyle().Bold(true)
    // ...
}

Review Questions

1. Is PTY presence checked before accessing window size? 2. Are per-session renderers used for Lipgloss? 3. Is authentication configured (public key preferred)? 4. Are session timeouts set? 5. Is logging middleware capturing connection info?