termd/internal/server/attach_test.go

package server_test

import (
	"context"
	"encoding/json"
	"io"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/coder/websocket"

	"termd/internal/attach"
)

// syncBuf is a goroutine-safe writer the attach client can stream into.
type syncBuf struct {
	mu sync.Mutex
	b  strings.Builder
}

func (s *syncBuf) Write(p []byte) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.b.Write(p)
}

func (s *syncBuf) String() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.b.String()
}

// TestAttachFlow drives the real attach client (minus raw tty mode) against
// the real websocket endpoint: snapshot on join, human keystrokes, API input
// interleaving, and detach.
func TestAttachFlow(t *testing.T) {
	e := newEnv(t)
	id := e.newSession("sh")
	e.send(id, "echo before-attach", "Enter")
	e.waitFor("pre-attach output", func() bool {
		return strings.Contains(e.screenText(id), "before-attach")
	})

	stdinR, stdinW := io.Pipe()
	var stdout syncBuf
	done := make(chan error, 1)
	go func() {
		done <- attach.Run(context.Background(), e.c, id, attach.Options{
			Write:  true,
			Stdin:  stdinR,
			Stdout: &stdout,
		})
	}()

	// The join snapshot repaints existing content.
	e.waitFor("snapshot on join", func() bool {
		return strings.Contains(stdout.String(), "before-attach")
	})

	// Human keystrokes go straight to the PTY.
	if _, err := stdinW.Write([]byte("echo via-human\r")); err != nil {
		t.Fatal(err)
	}
	e.waitFor("human input reaches the session", func() bool {
		return strings.Contains(e.screenText(id), "via-human")
	})

	// API input shows up live on the attached terminal.
	e.send(id, "echo via-api", "Enter")
	e.waitFor("api input streams to attach client", func() bool {
		out := stdout.String()
		return strings.Contains(out, "via-api") && strings.Contains(out, "via-human")
	})

	// C-] detaches cleanly; the session keeps running.
	if _, err := stdinW.Write([]byte{0x1d}); err != nil {
		t.Fatal(err)
	}
	if err := <-done; err != nil {
		t.Fatalf("attach.Run after detach: %v", err)
	}
	e.send(id, "echo after-detach", "Enter")
	e.waitFor("session alive after detach", func() bool {
		return strings.Contains(e.screenText(id), "after-detach")
	})
}

// TestAttachSessionExit: when the process dies while attached, the client is
// told the exit code and returns cleanly.
func TestAttachSessionExit(t *testing.T) {
	e := newEnv(t)
	id := e.newSession("sh")

	stdinR, stdinW := io.Pipe()
	defer stdinW.Close()
	var stdout syncBuf
	done := make(chan error, 1)
	go func() {
		done <- attach.Run(context.Background(), e.c, id, attach.Options{
			Write:  true,
			Stdin:  stdinR,
			Stdout: &stdout,
		})
	}()
	e.waitFor("attached", func() bool { return strings.Contains(stdout.String(), "$") })

	if _, err := stdinW.Write([]byte("exit 7\r")); err != nil {
		t.Fatal(err)
	}
	if err := <-done; err != nil {
		t.Fatalf("attach.Run after session exit: %v", err)
	}
	if out := stdout.String(); !strings.Contains(out, "exited with code 7") {
		t.Errorf("no exit notice in attach output:\n%q", out)
	}
}

// TestAttachReadOnly: without Write, the viewer streams everything but its
// keystrokes never reach the session — and the detach byte still works.
func TestAttachReadOnly(t *testing.T) {
	e := newEnv(t)
	id := e.newSession("sh")

	stdinR, stdinW := io.Pipe()
	var stdout syncBuf
	done := make(chan error, 1)
	go func() {
		done <- attach.Run(context.Background(), e.c, id, attach.Options{
			Stdin:  stdinR, // Write deliberately unset: read-only is the zero value
			Stdout: &stdout,
		})
	}()
	e.waitFor("attached", func() bool { return strings.Contains(stdout.String(), "$") })

	// Typing goes nowhere...
	if _, err := stdinW.Write([]byte("echo should-not-appear\r")); err != nil {
		t.Fatal(err)
	}
	// ...but API-driven output still streams to the viewer.
	e.send(id, "echo marker", "Enter")
	e.waitFor("marker streams to read-only viewer", func() bool {
		return strings.Contains(stdout.String(), "marker")
	})
	if strings.Contains(e.screenText(id), "should-not-appear") {
		t.Errorf("read-only viewer's keystrokes reached the session:\n%s", e.screenText(id))
	}

	// Detach byte still works read-only.
	if _, err := stdinW.Write([]byte{0x1d}); err != nil {
		t.Fatal(err)
	}
	if err := <-done; err != nil {
		t.Fatalf("read-only detach: %v", err)
	}
}

// TestAttachReadOnlyServerEnforced: the server drops binary frames on a
// ?readonly=1 connection even from a client that misbehaves and sends them.
func TestAttachReadOnlyServerEnforced(t *testing.T) {
	e := newEnv(t)
	id := e.newSession("sh")

	conn, _, err := websocket.Dial(ctx(), e.c.AttachURL(id)+"?readonly=1", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.CloseNow()
	conn.SetReadLimit(16 << 20)

	if err := conn.Write(ctx(), websocket.MessageBinary, []byte("echo smuggled\r")); err != nil {
		t.Fatal(err)
	}
	// Prove ordering with a write that must land after the dropped frame.
	e.send(id, "echo marker", "Enter")
	e.waitFor("marker", func() bool { return strings.Contains(e.screenText(id), "marker") })
	if strings.Contains(e.screenText(id), "smuggled") {
		t.Errorf("server let a binary frame through on a readonly connection:\n%s", e.screenText(id))
	}
}

// TestAttachRenderMode: with ?view=render the server streams full emulator
// frames — absolute positioning, per-line clear, cursor placement — instead
// of raw bytes, so a viewer larger than the session shows a clean session box
// padded with blanks.
func TestAttachRenderMode(t *testing.T) {
	e := newEnv(t)
	id := e.newSession("-cols", "40", "-rows", "6", "sh")

	conn, _, err := websocket.Dial(ctx(), e.c.AttachURL(id)+"?view=render", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.CloseNow()
	conn.SetReadLimit(16 << 20)

	readFrame := func() string {
		t.Helper()
		rctx, rcancel := context.WithTimeout(ctx(), 10*time.Second)
		defer rcancel()
		for {
			typ, data, err := conn.Read(rctx)
			if err != nil {
				t.Fatalf("reading frame: %v", err)
			}
			if typ == websocket.MessageBinary {
				return string(data)
			}
		}
	}

	first := readFrame()
	if !strings.Contains(first, "\x1b[2J") || !strings.Contains(first, "\x1b[H") {
		t.Errorf("first frame is not a clearing repaint: %q", first)
	}

	// Every frame is the full screen, so once the echo output exists, any
	// subsequent frame carries it along with the padding sequences.
	e.send(id, "echo render-me", "Enter")
	e.waitFor("frame with new output", func() bool {
		f := readFrame()
		return strings.Contains(f, "render-me") &&
			strings.Contains(f, "\x1b[0m\x1b[K") && // per-line reset+clear (the padding)
			strings.Contains(f, "\x1b[J") // clear below the session box
	})
}

// TestAttachRenderThroughClient drives the real attach client with Render on.
func TestAttachRenderThroughClient(t *testing.T) {
	e := newEnv(t)
	id := e.newSession("-cols", "40", "-rows", "6", "sh")

	stdinR, stdinW := io.Pipe()
	var stdout syncBuf
	done := make(chan error, 1)
	go func() {
		done <- attach.Run(context.Background(), e.c, id, attach.Options{
			Render: true,
			Stdin:  stdinR,
			Stdout: &stdout,
		})
	}()

	e.send(id, "echo boxed", "Enter")
	e.waitFor("rendered frame reaches client", func() bool {
		out := stdout.String()
		return strings.Contains(out, "boxed") && strings.Contains(out, "\x1b[0m\x1b[K")
	})

	if _, err := stdinW.Write([]byte{0x1d}); err != nil {
		t.Fatal(err)
	}
	if err := <-done; err != nil {
		t.Fatalf("detach from render view: %v", err)
	}
}

// TestAttachResizeControl exercises the server side of the resize control
// frame with a bare websocket client (the real client only sends it when
// stdout is a terminal).
func TestAttachResizeControl(t *testing.T) {
	e := newEnv(t)
	id := e.newSession("sh")

	conn, _, err := websocket.Dial(ctx(), e.c.AttachURL(id), nil)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.CloseNow()
	conn.SetReadLimit(16 << 20)

	frame, err := json.Marshal(map[string]any{"resize": map[string]int{"cols": 90, "rows": 30}})
	if err != nil {
		t.Fatal(err)
	}
	if err := conn.Write(ctx(), websocket.MessageText, frame); err != nil {
		t.Fatal(err)
	}
	e.waitFor("resize applied", func() bool {
		scr := e.screen(id)
		return scr.Cols == 90 && scr.Rows == 30
	})
}