Hammerhead/e2e/terminal_relay_test.go

package e2e

import (
	"bytes"
	"context"
	"encoding/json"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git.veillax.com/Veillax/bamboo"
	"git.veillax.com/Veillax/bamboo/sshd"
	"git.veillax.com/Veillax/blacktip"
	"git.veillax.com/Veillax/reef/model"
	"golang.org/x/crypto/ssh"
	"golang.org/x/net/websocket"
)

// With Bamboo reporting its SSH server, the terminal works with no SSH port anywhere: Hammerhead
// asks the agent to call back and signs in over that connection.
func TestWebTerminalWithoutAnyOpenPort(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	siteCfgHook = func(c *bamboo.Config) { c.SSH = ln.Addr().String() } // what bamboo -ssh tells the agent
	w := setup(t, blacktip.TraefikConfig{})
	siteCfgHook = nil
	base := "/api/orgs/" + w.org
	u, _ := url.Parse(w.ts.URL)

	var key struct{ PublicKey string }
	w.admin.must("GET", base+"/terminal", nil, &key)
	dir := t.TempDir()
	keys := filepath.Join(dir, "authorized_keys")
	if err := os.WriteFile(keys, []byte(key.PublicKey+"\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	srv, err := sshd.New(sshd.Config{AuthorizedKeys: keys, HostKey: filepath.Join(dir, "ssh_host_ed25519_key")})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go srv.Serve(ctx, ln)

	// the agent's next report tells Hammerhead where its SSH server is
	var st struct{ Sites []model.Site }
	for i := 0; i < 60; i++ {
		w.admin.must("GET", base+"/state", nil, &st)
		for _, s := range st.Sites {
			if s.ID == w.site && s.SSH != "" {
				i = 60
			}
		}
		if i < 60 {
			time.Sleep(200 * time.Millisecond)
		}
	}
	var site model.Site
	for _, s := range st.Sites {
		if s.ID == w.site {
			site = s
		}
	}
	if site.SSH != ln.Addr().String() || !site.SSHRelay {
		t.Fatalf("the site did not report its SSH server: %+v", site)
	}

	wsURL := "ws" + strings.TrimPrefix(w.ts.URL, "http") + base + "/sites/" + w.site + "/terminal?cols=100&rows=30"
	dial := func() *websocket.Conn {
		t.Helper()
		cfg, err := websocket.NewConfig(wsURL, w.ts.URL)
		if err != nil {
			t.Fatal(err)
		}
		var ck []string
		for _, c := range w.admin.c.Jar.Cookies(u) {
			ck = append(ck, c.Name+"="+c.Value)
		}
		cfg.Header.Set("Cookie", strings.Join(ck, "; "))
		ws, err := websocket.DialConfig(cfg)
		if err != nil {
			t.Fatal(err)
		}
		return ws
	}
	first := func() (m termMsg) {
		t.Helper()
		ws := dial()
		defer ws.Close()
		for {
			ws.SetReadDeadline(time.Now().Add(20 * time.Second))
			var f []byte
			if err := websocket.Message.Receive(ws, &f); err != nil {
				t.Fatalf("waiting for the terminal: %v", err)
			}
			if len(f) > 0 && f[0] == '{' {
				return parseTermMsg(f)
			}
		}
	}

	m := first()
	if sshd.Supported {
		if m.Type != "ready" {
			t.Fatalf("opening the terminal with no port: %+v", m)
		}
		ws := dial()
		defer ws.Close()
		// $((40+2)) is expanded by the shell, so "relay-42" can only come from the command running
		if err := websocket.Message.Send(ws, []byte("echo relay-$((40+2))\n")); err != nil {
			t.Fatal(err)
		}
		var out bytes.Buffer
		for !strings.Contains(out.String(), "relay-42\r\n") {
			ws.SetReadDeadline(time.Now().Add(15 * time.Second))
			var f []byte
			if err := websocket.Message.Receive(ws, &f); err != nil {
				t.Fatalf("waiting for the command's output (got %q): %v", out.String(), err)
			}
			out.Write(f)
		}
	} else if m.Type != "error" || !strings.Contains(m.Message, "terminal") {
		// this platform can't run shells, but asking for one proves the call-back and Hammerhead's key worked
		t.Fatalf("with no port, on a platform without shells: %+v", m)
	}

	// it really did not need a port
	var maps struct{ Mappings []model.Mapping }
	w.admin.must("GET", base+"/state", nil, &maps)
	if len(maps.Mappings) != 0 {
		t.Errorf("the test opened a port after all: %+v", maps.Mappings)
	}

	// a machine with no shell (the published container image has none) says why, instead of the
	// terminal just disconnecting: the shell's error output reaches the browser, then the exit code
	if sshd.Supported {
		t.Setenv("SHELL", "/nonexistent/shell")
		ws := dial()
		defer ws.Close()
		var out bytes.Buffer
		var exit termMsg
		for exit.Type != "exit" {
			ws.SetReadDeadline(time.Now().Add(15 * time.Second))
			var f []byte
			if err := websocket.Message.Receive(ws, &f); err != nil {
				t.Fatalf("waiting for the shell to end (got %q): %v", out.String(), err)
			}
			if len(f) > 0 && f[0] == '{' {
				if m := parseTermMsg(f); m.Type == "exit" {
					exit = m
				}
				continue
			}
			out.Write(f)
		}
		if exit.Code != 127 || !strings.Contains(out.String(), "no shell to run") {
			t.Errorf("a missing shell: exit %d, output %q", exit.Code, out.String())
		}
	}

	// the SSH server still decides who gets in
	other := newSSHSigner(t)
	if err := os.WriteFile(keys, ssh.MarshalAuthorizedKey(other.PublicKey()), 0o600); err != nil {
		t.Fatal(err)
	}
	if m := first(); m.Type != "error" || !strings.Contains(m.Message, "authorized_keys") {
		t.Errorf("after removing Hammerhead's key: %+v", m)
	}

	// and a stopped SSH server is reported at once rather than after a timeout
	cancel()
	ln.Close()
	time.Sleep(300 * time.Millisecond)
	start := time.Now()
	if m := first(); m.Type != "error" || !strings.Contains(m.Message, "SSH server") {
		t.Errorf("with the SSH server stopped: %+v", m)
	}
	if time.Since(start) > 8*time.Second {
		t.Errorf("a stopped SSH server took %s to be reported", time.Since(start))
	}
}

func parseTermMsg(b []byte) termMsg {
	var m termMsg
	json.Unmarshal(b, &m)
	return m
}