byre/internal/server/terminal_test.go

package server_test

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"

	"g.fc.run/misc/byre/internal/api"
)

type termConn struct {
	t    *testing.T
	conn *websocket.Conn
	seen strings.Builder
}

func (h *harness) openTerminal(taskID string) *termConn {
	h.t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	header := http.Header{}
	header.Set("Tailscale-User-Login", login)
	header.Set("Tailscale-User-Name", "Alice")
	url := "ws" + strings.TrimPrefix(h.url, "http") + "/api/tasks/" + taskID + "/terminal?cols=80&rows=24"
	conn, _, err := websocket.Dial(ctx, url, &websocket.DialOptions{HTTPHeader: header})
	if err != nil {
		h.t.Fatalf("dial terminal: %v", err)
	}
	h.t.Cleanup(func() { conn.CloseNow() })
	return &termConn{t: h.t, conn: conn}
}

func (c *termConn) send(text string) {
	c.t.Helper()
	err := c.conn.Write(context.Background(), websocket.MessageBinary, []byte(text))
	if err != nil {
		c.t.Fatalf("write terminal: %v", err)
	}
}

func (c *termConn) waitFor(want string) {
	c.t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	for !strings.Contains(c.seen.String(), want) {
		_, b, err := c.conn.Read(ctx)
		if err != nil {
			c.t.Fatalf("waiting for %q: %v; seen %q", want, err, c.seen.String())
		}
		c.seen.Write(b)
	}
}

func (c *termConn) waitClosed() websocket.CloseError {
	c.t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	for {
		_, b, err := c.conn.Read(ctx)
		if err == nil {
			c.seen.Write(b)
			continue
		}
		ce := websocket.CloseError{}
		if !errors.As(err, &ce) {
			c.t.Fatalf("want a close frame, got %v", err)
		}
		return ce
	}
}

func TestTaskTerminal(t *testing.T) {
	t.Setenv("SHELL", "/bin/sh")
	h := newHarness(t)
	p := h.createProject("demo")
	task := api.Task{}
	h.do("POST", "/api/projects/"+p.Id+"/tasks", api.TaskCreate{Title: "terminal"}, &task)

	wt, err := filepath.EvalSymlinks(task.WorktreePath)
	if err != nil {
		t.Fatalf("resolve worktree: %v", err)
	}
	header := http.Header{}
	header.Set("Tailscale-User-Login", login)
	header.Set("Tailscale-User-Name", "Alice")
	header.Set("Origin", "https://elsewhere.example")
	_, resp, err := websocket.Dial(context.Background(), "ws"+strings.TrimPrefix(h.url, "http")+"/api/tasks/"+task.Id+"/terminal?cols=80&rows=24", &websocket.DialOptions{HTTPHeader: header})
	if err == nil || resp == nil || resp.StatusCode != http.StatusForbidden {
		t.Fatalf("cross-origin terminal: want 403, got %v %v", resp, err)
	}

	a := h.openTerminal(task.Id)
	a.send("echo \"$PORT:$(pwd -P)\"\n")
	a.waitFor(fmt.Sprintf("%d:%s", task.PortBase, wt))
	a.send("echo byre-$((40+2))\n")
	a.waitFor("byre-42")

	size, err := json.Marshal(api.TerminalSize{Cols: 100, Rows: 30})
	if err != nil {
		t.Fatalf("marshal size: %v", err)
	}
	err = a.conn.Write(context.Background(), websocket.MessageText, size)
	if err != nil {
		t.Fatalf("resize: %v", err)
	}
	a.send("stty size\n")
	a.waitFor("30 100")

	b := h.openTerminal(task.Id)
	b.waitFor("byre-42")
	b.send("exit\n")
	for _, c := range []*termConn{a, b} {
		ce := c.waitClosed()
		if ce.Code != 4000 || ce.Reason != "exit status 0" {
			t.Fatalf("close after exit: %d %q", ce.Code, ce.Reason)
		}
	}

	c := h.openTerminal(task.Id)
	c.send("echo fresh\n")
	c.waitFor("fresh")
	if strings.Contains(c.seen.String(), "byre-42") {
		t.Fatalf("new shell replayed the old one's output: %q", c.seen.String())
	}

	code := h.do("POST", "/api/tasks/"+task.Id+"/discard", nil, nil)
	if code != http.StatusOK {
		t.Fatalf("discard: %d", code)
	}
	ce := c.waitClosed()
	if ce.Code != 4000 {
		t.Fatalf("close after discard: %d %q", ce.Code, ce.Reason)
	}

	code = h.do("GET", "/api/tasks/"+task.Id+"/terminal?cols=80&rows=24", nil, nil)
	if code != http.StatusConflict {
		t.Fatalf("terminal of a discarded task: want 409, got %d", code)
	}
}