internal/server/api/client/console_test.go

Ref: main v0.0.5   Size: 5.4 KiB   History

package client_test

import (
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"

	"github.com/a73x/eitri/internal/server/api/client"
)

func TestMintStreamTicket(t *testing.T) {
	var cap capture
	srv := serve(t, &cap, http.StatusCreated, `{"ticket":"tkt-123"}`)
	c := &client.Client{BaseURL: srv.URL, Token: "tok"}

	ticket, err := c.MintStreamTicket(context.Background())
	if err != nil {
		t.Fatalf("MintStreamTicket: %v", err)
	}
	if ticket != "tkt-123" {
		t.Errorf("ticket = %q, want tkt-123", ticket)
	}
	if cap.method != http.MethodPost || cap.path != "/api/v1/stream-tickets" {
		t.Errorf("request = %s %s, want POST /api/v1/stream-tickets", cap.method, cap.path)
	}
	if cap.auth != "Bearer tok" {
		t.Errorf("auth = %q, want the PAT in the header", cap.auth)
	}
}

// TestMintStreamTicketRejectsAnEmptyTicket: an empty ticket dials a console
// that answers 401, so the failure belongs at the mint, naming the mint.
func TestMintStreamTicketRejectsAnEmptyTicket(t *testing.T) {
	var cap capture
	srv := serve(t, &cap, http.StatusCreated, `{"ticket":""}`)
	c := &client.Client{BaseURL: srv.URL, Token: "tok"}

	if _, err := c.MintStreamTicket(context.Background()); err == nil {
		t.Fatal("MintStreamTicket: want error for an empty ticket, got nil")
	} else if !strings.Contains(err.Error(), "stream-tickets") {
		t.Errorf("error = %q, want it to name the endpoint", err.Error())
	}
}

func TestMintStreamTicketPropagatesTheAPIError(t *testing.T) {
	var cap capture
	srv := serve(t, &cap, http.StatusUnauthorized, `unauthorized`)
	c := &client.Client{BaseURL: srv.URL, Token: "stale"}

	_, err := c.MintStreamTicket(context.Background())
	var apiErr *client.Error
	if !errors.As(err, &apiErr) {
		t.Fatalf("error = %v, want a *client.Error", err)
	}
	if apiErr.Status != http.StatusUnauthorized {
		t.Errorf("status = %d, want 401", apiErr.Status)
	}
}

// consoleServer stands in for the control plane's console endpoint: it mints a
// ticket on POST /api/v1/stream-tickets and accepts the WebSocket only when
// that same ticket comes back in the query, writing text to whoever attaches.
func consoleServer(t *testing.T, text string) (*httptest.Server, *string) {
	t.Helper()
	dialedPath := new(string)
	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/v1/stream-tickets", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusCreated)
		io.WriteString(w, `{"ticket":"tkt-abc"}`)
	})
	mux.HandleFunc("GET /api/v1/vms/{id}/console/ws", func(w http.ResponseWriter, r *http.Request) {
		*dialedPath = r.URL.EscapedPath() + "?" + r.URL.RawQuery
		if r.URL.Query().Get("ticket") != "tkt-abc" {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		c, err := websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true})
		if err != nil {
			return
		}
		defer c.CloseNow()
		nc := websocket.NetConn(r.Context(), c, websocket.MessageBinary)
		io.WriteString(nc, text)
		time.Sleep(50 * time.Millisecond)
		_ = c.Close(websocket.StatusNormalClosure, "")
	})
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	return srv, dialedPath
}

// TestDialConsoleMintsThenDials pins the two-step: the console is reached with
// a freshly minted ticket in the URL, and the bytes come back raw.
func TestDialConsoleMintsThenDials(t *testing.T) {
	srv, dialedPath := consoleServer(t, "ubuntu-vm login: ")
	c := &client.Client{BaseURL: srv.URL, Token: "tok"}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	stream, err := c.DialConsole(ctx, "vm-1")
	if err != nil {
		t.Fatalf("DialConsole: %v", err)
	}
	defer stream.Close()

	got, err := io.ReadAll(stream)
	if err != nil && err != io.EOF {
		t.Fatalf("read console: %v", err)
	}
	if !strings.Contains(string(got), "login:") {
		t.Errorf("console bytes = %q, want the guest's output", string(got))
	}
	if *dialedPath != "/api/v1/vms/vm-1/console/ws?ticket=tkt-abc" {
		t.Errorf("dialed %q, want the VM's console path carrying the minted ticket", *dialedPath)
	}
}

// TestDialConsoleFailsWhenTheTicketIsRefused: a console that will not open is
// an error at dial, not a silent stream that never says anything.
func TestDialConsoleFailsWhenTheTicketIsRefused(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/v1/stream-tickets", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusCreated)
		io.WriteString(w, `{"ticket":"expired"}`)
	})
	mux.HandleFunc("GET /api/v1/vms/{id}/console/ws", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()

	c := &client.Client{BaseURL: srv.URL, Token: "tok"}
	if _, err := c.DialConsole(context.Background(), "vm-1"); err == nil {
		t.Fatal("DialConsole: want error when the console refuses the ticket, got nil")
	} else if !strings.Contains(err.Error(), "vm-1") {
		t.Errorf("error = %q, want it to name the VM", err.Error())
	}
}

// TestDialConsoleEscapesTheVMID pins that an id with URL-significant characters
// addresses one VM rather than reshaping the path.
func TestDialConsoleEscapesTheVMID(t *testing.T) {
	srv, dialedPath := consoleServer(t, "x")
	c := &client.Client{BaseURL: srv.URL, Token: "tok"}

	stream, err := c.DialConsole(context.Background(), "vm/1")
	if err == nil {
		stream.Close()
	}
	if !strings.Contains(*dialedPath, "vm%2F1") {
		t.Errorf("dialed %q, want the id escaped into one path segment", *dialedPath)
	}
}