Bamboo/terminal_test.go

package bamboo

import (
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"golang.org/x/net/websocket"
)

func TestSSHTarget(t *testing.T) {
	for in, want := range map[string]string{
		"":               "",
		"127.0.0.1:2222": "127.0.0.1:2222",
		":2222":          "127.0.0.1:2222",
		"0.0.0.0:2200":   "127.0.0.1:2200",
		"[::]:2222":      "127.0.0.1:2222",
		"10.0.0.5:2222":  "10.0.0.5:2222",
	} {
		if got := sshTarget(in); got != want {
			t.Errorf("sshTarget(%q) = %q, want %q", in, got, want)
		}
	}
}

// Hammerhead lists a terminal session; the agent calls back with its token and joins the connection
// to its SSH server. Here the "SSH server" is a local echo, and Hammerhead is a test server.
func TestTerminalRelayJoinsHammerheadToTheLocalSSHServer(t *testing.T) {
	echo, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer echo.Close()
	go func() {
		for {
			c, err := echo.Accept()
			if err != nil {
				return
			}
			go func() { defer c.Close(); io.Copy(c, c) }()
		}
	}()

	got := make(chan *websocket.Conn, 4)
	auth := make(chan string, 4)
	paths := make(chan string, 4)
	hh := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		auth <- r.Header.Get("Authorization")
		paths <- r.URL.Path
		release := make(chan struct{})
		websocket.Server{Handshake: func(*websocket.Config, *http.Request) error { return nil }, Handler: func(ws *websocket.Conn) {
			ws.PayloadType = websocket.BinaryFrame
			got <- ws
			<-release
		}}.ServeHTTP(w, r)
		close(release)
	}))
	defer hh.Close()

	b, err := New(Config{Hammerhead: hh.URL, Token: "site-token", DataDir: t.TempDir(), SSH: echo.Addr().String()})
	if err != nil {
		t.Fatal(err)
	}
	b.openTerminals([]string{"abc"})
	b.openTerminals([]string{"abc"}) // fetching the config again must not answer the same session twice

	var ws *websocket.Conn
	select {
	case ws = <-got:
	case <-time.After(5 * time.Second):
		t.Fatal("the agent never called back")
	}
	if a := <-auth; a != "Bearer site-token" {
		t.Errorf("call-back authorization: %q", a)
	}
	if p := <-paths; p != "/site/terminal/abc" {
		t.Errorf("call-back path: %q", p)
	}
	if _, err := ws.Write([]byte("hello through the relay")); err != nil {
		t.Fatal(err)
	}
	buf := make([]byte, 64)
	ws.SetReadDeadline(time.Now().Add(5 * time.Second))
	n, err := ws.Read(buf)
	if err != nil || string(buf[:n]) != "hello through the relay" {
		t.Fatalf("echo through the relay: %q, %v", buf[:n], err)
	}
	select {
	case <-got:
		t.Error("the same session was answered twice")
	case <-time.After(300 * time.Millisecond):
	}
	ws.Close()
}

// With no SSH server the agent does not call back at all, rather than connect and hang.
func TestTerminalRelayDoesNothingWithoutAnSSHServer(t *testing.T) {
	called := make(chan struct{}, 1)
	hh := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { called <- struct{}{} }))
	defer hh.Close()
	b, err := New(Config{Hammerhead: hh.URL, Token: "t", DataDir: t.TempDir()})
	if err != nil {
		t.Fatal(err)
	}
	b.openTerminals([]string{"x"})
	select {
	case <-called:
		t.Error("the agent called Hammerhead although it has no SSH server")
	case <-time.After(400 * time.Millisecond):
	}
}