internal/server/api/console_test.go
Ref: Size: 6.6 KiB History
package api
import (
"context"
"encoding/json"
"errors"
"io"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/coder/websocket"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// fakeConsole hands back an in-memory duplex pipe and records the ask.
// Tests that touch serverEnd/hostID/vmID after a successful dial must set
// opened and receive from it first: OpenConsole runs on the handler's
// goroutine AFTER the WS handshake completes, so the dial returning does not
// order the test's reads after the fake's writes — the channel does.
type fakeConsole struct {
hostID, vmID string
err error
serverEnd io.ReadWriteCloser
opened chan struct{} // optional; signalled once per OpenConsole call
}
type rwc struct {
io.Reader
io.Writer
closeFn func() error
}
func (c rwc) Close() error { return c.closeFn() }
func duplexPipe() (a, b io.ReadWriteCloser) {
ar, bw := io.Pipe()
br, aw := io.Pipe()
return rwc{ar, aw, func() error { aw.Close(); return ar.Close() }},
rwc{br, bw, func() error { bw.Close(); return br.Close() }}
}
func (f *fakeConsole) OpenConsole(_ context.Context, hostID, vmID string) (io.ReadWriteCloser, error) {
f.hostID, f.vmID = hostID, vmID
if f.err != nil {
return nil, f.err
}
var clientEnd io.ReadWriteCloser
clientEnd, f.serverEnd = duplexPipe()
if f.opened != nil {
f.opened <- struct{}{}
}
return clientEnd, nil
}
// consoleFixture is the seed data console WS tests need: the running server
// plus the enrolled host + created VM pair the endpoint routes to.
type consoleFixture struct {
ts *httptest.Server
hostID string
vmID string
}
// newConsoleAPI builds a test server (shared fixture) and seeds one host with
// one VM over the HTTP API, returning the *API for SetConsoleDialer wiring.
func newConsoleAPI(t *testing.T) (*API, consoleFixture) {
t.Helper()
ts, _, _, _, a := newServer(t)
out := enroll(t, ts)
resp := do(t, "POST", ts.URL+"/api/v1/vms", testPAT,
map[string]any{"host_id": out["host_id"], "name": "console-vm"})
require.Equal(t, 201, resp.StatusCode)
var created map[string]string
require.NoError(t, json.NewDecoder(resp.Body).Decode(&created))
require.NotEmpty(t, created["id"])
return a, consoleFixture{ts: ts, hostID: out["host_id"], vmID: created["id"]}
}
// consoleWSURL builds the ws:// URL for the fixture VM carrying the ticket.
func consoleWSURL(fix consoleFixture, ticket string) string {
return strings.Replace(fix.ts.URL, "http://", "ws://", 1) +
"/api/v1/vms/" + fix.vmID + "/console/ws?ticket=" + ticket
}
func TestConsoleWSRequiresTicket(t *testing.T) {
_, fix := newConsoleAPI(t)
// No ticket: plain GET (the handler rejects before upgrading).
resp, err := fix.ts.Client().Get(fix.ts.URL + "/api/v1/vms/whatever/console/ws")
require.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 401, resp.StatusCode)
}
func TestConsoleWSRejectsDeletedVM(t *testing.T) {
a, fix := newConsoleAPI(t)
a.SetConsoleDialer(&fakeConsole{})
// Tombstone the VM (GetVM does not filter tombstones).
resp := do(t, "DELETE", fix.ts.URL+"/api/v1/vms/"+fix.vmID, testPAT, nil)
require.Equal(t, 204, resp.StatusCode)
ticket := mintTicket(t, fix.ts.URL, testPAT)
r, err := fix.ts.Client().Get(fix.ts.URL + "/api/v1/vms/" + fix.vmID + "/console/ws?ticket=" + ticket)
require.NoError(t, err)
defer r.Body.Close()
assert.Equal(t, 404, r.StatusCode, "console to a tombstoned VM must be rejected")
}
func TestConsoleWSBridgesBytes(t *testing.T) {
a, fix := newConsoleAPI(t)
fc := &fakeConsole{opened: make(chan struct{}, 1)}
a.SetConsoleDialer(fc)
ticket := mintTicket(t, fix.ts.URL, testPAT)
c, _, err := websocket.Dial(context.Background(), consoleWSURL(fix, ticket), nil)
require.NoError(t, err)
defer c.CloseNow()
<-fc.opened // handler has dialed the fake; serverEnd is set
// guest → browser
_, err = fc.serverEnd.Write([]byte("login:"))
require.NoError(t, err)
typ, data, err := c.Read(context.Background())
require.NoError(t, err)
assert.Equal(t, websocket.MessageBinary, typ)
assert.Equal(t, "login:", string(data))
// browser → guest
require.NoError(t, c.Write(context.Background(), websocket.MessageBinary, []byte("root\r")))
buf := make([]byte, 5)
_, err = io.ReadFull(fc.serverEnd, buf)
require.NoError(t, err)
assert.Equal(t, "root\r", string(buf))
assert.Equal(t, fix.vmID, fc.vmID)
assert.Equal(t, fix.hostID, fc.hostID)
}
// TestConsoleWSCloseTearsDownAgentStream pins the no-leaked-agent-stream
// invariant: when the browser side goes away, the handler's pumps must close
// the agent-leg stream — otherwise every abandoned tab would hold a serial
// pump hostage forever.
func TestConsoleWSCloseTearsDownAgentStream(t *testing.T) {
a, fix := newConsoleAPI(t)
fc := &fakeConsole{opened: make(chan struct{}, 1)}
a.SetConsoleDialer(fc)
ticket := mintTicket(t, fix.ts.URL, testPAT)
c, _, err := websocket.Dial(context.Background(), consoleWSURL(fix, ticket), nil)
require.NoError(t, err)
<-fc.opened // handler has dialed the fake; serverEnd is set
// Browser disconnects.
require.NoError(t, c.Close(websocket.StatusNormalClosure, ""))
// The agent end must observe teardown (EOF/closed-pipe), not block.
readErr := make(chan error, 1)
go func() {
_, err := fc.serverEnd.Read(make([]byte, 1))
readErr <- err
}()
select {
case err := <-readErr:
require.Error(t, err, "agent stream must be closed, not left readable")
case <-time.After(5 * time.Second):
t.Fatal("agent stream still open after browser WS close — leaked stream")
}
}
func TestConsoleWSHostOfflineClosesWithReason(t *testing.T) {
a, fix := newConsoleAPI(t)
a.SetConsoleDialer(&fakeConsole{err: errors.New("agent not connected")})
ticket := mintTicket(t, fix.ts.URL, testPAT)
c, _, err := websocket.Dial(context.Background(), consoleWSURL(fix, ticket), nil)
require.NoError(t, err, "the upgrade succeeds; failure arrives as a close frame")
defer c.CloseNow()
_, _, err = c.Read(context.Background())
var ce websocket.CloseError
require.ErrorAs(t, err, &ce)
assert.Equal(t, websocket.StatusInternalError, ce.Code)
assert.Contains(t, ce.Reason, "agent not connected", "the reason must reach the browser")
}
func TestConsoleWSTicketIsSingleUse(t *testing.T) {
a, fix := newConsoleAPI(t)
a.SetConsoleDialer(&fakeConsole{})
ticket := mintTicket(t, fix.ts.URL, testPAT)
url := fix.ts.URL + "/api/v1/vms/" + fix.vmID + "/console/ws?ticket=" + ticket
// Consume it once (plain GET is fine — the ticket is consumed before upgrade).
_, _ = fix.ts.Client().Get(url)
resp, err := fix.ts.Client().Get(url)
require.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 401, resp.StatusCode, "one-time ticket must not replay")
}