shark-expose/e2e/e2e_test.go

// Package e2e runs the shark CLI against a real Hammerhead, Blacktip and Bamboo in one process.
package e2e

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"git.veillax.com/Veillax/blacktip"
	"git.veillax.com/Veillax/hammerhead"
	"git.veillax.com/Veillax/shark-expose/internal/cli"
)

type stack struct {
	url, key string
	admin    *http.Client
	org      string
}

func freeUDPPort(t *testing.T) int {
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer pc.Close()
	return pc.LocalAddr().(*net.UDPAddr).Port
}

func freePort(t *testing.T) int {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port
}

func (s *stack) call(t *testing.T, method, path string, in any) []byte {
	t.Helper()
	var body io.Reader
	if in != nil {
		b, _ := json.Marshal(in)
		body = bytes.NewReader(b)
	}
	req, _ := http.NewRequest(method, s.url+path, body)
	req.Header.Set("Content-Type", "application/json")
	resp, err := s.admin.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		t.Fatalf("%s %s -> %d %s", method, path, resp.StatusCode, raw)
	}
	return raw
}

// newStack starts Hammerhead and one Blacktip edge with a domain, makes an API key, and points
// shark's settings at it.
func newStack(t *testing.T) *stack {
	t.Helper()
	s := &stack{}
	var handler http.Handler
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { handler.ServeHTTP(w, r) }))
	t.Cleanup(ts.Close)
	hh, err := hammerhead.New(t.TempDir(), hammerhead.Options{AdminEmail: "admin@test.local", AdminPassword: "password123", DashboardURL: ts.URL})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(hh.Close)
	handler = hh.Handler()
	s.url = ts.URL
	jar, _ := cookiejar.New(nil)
	s.admin = &http.Client{Jar: jar, Timeout: 15 * time.Second}
	s.call(t, "POST", "/api/login", map[string]string{"email": "admin@test.local", "password": "password123"})
	var me struct{ Orgs []struct{ ID string } }
	json.Unmarshal(s.call(t, "GET", "/api/me", nil), &me)
	s.org = me.Orgs[0].ID

	tunnel := freeUDPPort(t)
	var edge struct{ ID, Token string }
	json.Unmarshal(s.call(t, "POST", "/api/edges", map[string]any{"name": "edge", "tunnelAddr": "127.0.0.1:" + strconv.Itoa(tunnel)}), &edge)
	s.call(t, "POST", "/api/orgs/"+s.org+"/domains", map[string]any{"name": "dev.test", "edgeId": edge.ID})
	var key struct{ Key string }
	json.Unmarshal(s.call(t, "POST", "/api/orgs/"+s.org+"/apikeys", map[string]string{"name": "shark"}), &key)
	s.key = key.Key

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	bt, err := blacktip.New(blacktip.Config{Hammerhead: ts.URL, Token: edge.Token, TunnelAddr: "127.0.0.1:" + strconv.Itoa(tunnel), DataDir: t.TempDir(),
		GetCertificate: func(*tls.ClientHelloInfo) (*tls.Certificate, error) { return nil, fmt.Errorf("no tls in this test") }})
	if err != nil {
		t.Fatal(err)
	}
	go bt.Run(ctx)
	t.Setenv("SHARK_CONFIG", t.TempDir()+"/config.json")
	t.Setenv("SHARK_HAMMERHEAD", s.url)
	t.Setenv("SHARK_TOKEN", s.key)
	return s
}

// syncBuf is a buffer that can be read while shark is still writing to it.
type syncBuf struct {
	mu sync.Mutex
	b  bytes.Buffer
}

func (s *syncBuf) Write(p []byte) (int, error) { s.mu.Lock(); defer s.mu.Unlock(); return s.b.Write(p) }
func (s *syncBuf) String() string              { s.mu.Lock(); defer s.mu.Unlock(); return s.b.String() }

var addrLine = regexp.MustCompile(`shark (\S+)\s+->`)

// shark runs `shark <args>` and waits for it to print its public address. stop is Ctrl+C: it
// waits for shark to exit and returns the exit code.
func shark(t *testing.T, args ...string) (addr string, stop func() int) {
	t.Helper()
	out, errb := &syncBuf{}, &syncBuf{}
	ctx, cancel := context.WithCancel(context.Background())
	code := make(chan int, 1)
	go func() { code <- cli.Run(ctx, args, out, errb) }()
	var once sync.Once
	var exit int
	stop = func() int {
		once.Do(func() {
			cancel()
			select {
			case exit = <-code:
			case <-time.After(20 * time.Second):
				t.Error("shark did not exit after Ctrl+C")
			}
		})
		return exit
	}
	t.Cleanup(func() { stop() })
	for i := 0; i < 200; i++ {
		if m := addrLine.FindStringSubmatch(out.String()); m != nil {
			return m[1], stop
		}
		select {
		case c := <-code:
			t.Fatalf("shark exited early (%d): %s%s", c, out.String(), errb.String())
		case <-time.After(100 * time.Millisecond):
		}
	}
	t.Fatalf("no address printed: %s%s", out.String(), errb.String())
	return
}

func (s *stack) exposures(t *testing.T) int {
	t.Helper()
	req, _ := http.NewRequest("GET", s.url+"/api/orgs/-/expose", nil)
	req.Header.Set("Authorization", "Bearer "+s.key)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	var list []map[string]any
	json.NewDecoder(resp.Body).Decode(&list)
	return len(list)
}

func TestSharkTCP(t *testing.T) {
	s := newStack(t)
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()
	go func() {
		for {
			c, err := l.Accept()
			if err != nil {
				return
			}
			go func() {
				defer c.Close()
				b := make([]byte, 64)
				n, _ := c.Read(b)
				fmt.Fprintf(c, "echo:%s", b[:n])
			}()
		}
	}()
	addr, stop := shark(t, "expose", "tcp", strconv.Itoa(l.Addr().(*net.TCPAddr).Port))
	hostport := strings.TrimPrefix(addr, "tcp://")

	var got string
	for i := 0; i < 60 && got != "echo:hi"; i++ {
		if c, err := net.DialTimeout("tcp", hostport, time.Second); err == nil {
			c.SetDeadline(time.Now().Add(2 * time.Second))
			fmt.Fprint(c, "hi")
			b, _ := io.ReadAll(c)
			c.Close()
			got = string(b)
		}
		time.Sleep(150 * time.Millisecond)
	}
	if got != "echo:hi" {
		t.Fatalf("through shark: %q", got)
	}
	if n := s.exposures(t); n != 1 {
		t.Fatalf("want 1 exposure while running, have %d", n)
	}
	if code := stop(); code != 0 {
		t.Errorf("exit code %d", code)
	}
	if n := s.exposures(t); n != 0 {
		t.Errorf("exposure left behind after exit: %d", n)
	}
	closed := false
	for i := 0; i < 30 && !closed; i++ {
		c, err := net.DialTimeout("tcp", hostport, 300*time.Millisecond)
		if err != nil {
			closed = true
			break
		}
		c.Close()
		time.Sleep(200 * time.Millisecond)
	}
	if !closed {
		t.Error("public port still open after exit")
	}
}

func TestSharkHTTPWithPassword(t *testing.T) {
	s := newStack(t)
	app := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, "hello from the app") }))
	defer app.Close()
	port := freePort(t)
	addr, stop := shark(t, "expose", app.Listener.Addr().String(), "--name", "demo", "--port", strconv.Itoa(port), "--password", "hunter22")
	host := strings.TrimPrefix(addr, "http://")
	host = host[:strings.LastIndex(host, ":")]
	if !strings.HasPrefix(host, "demo-") || !strings.HasSuffix(host, ".dev.test") {
		t.Fatalf("public name %q", host)
	}
	// the edge is on loopback, so dial it directly while presenting the public host name
	hc := &http.Client{Timeout: 5 * time.Second,
		Transport: &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, "tcp", "127.0.0.1:"+strconv.Itoa(port))
		}},
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	var status int
	var loc string
	for i := 0; i < 60 && status != 302; i++ {
		if resp, err := hc.Get(addr + "/"); err == nil {
			resp.Body.Close()
			status, loc = resp.StatusCode, resp.Header.Get("Location")
		}
		time.Sleep(150 * time.Millisecond)
	}
	if status != 302 || !strings.HasPrefix(loc, s.url+"/auth/r/") {
		t.Fatalf("a protected site should redirect to the login, got %d %q", status, loc)
	}
	stop()
	if n := s.exposures(t); n != 0 {
		t.Errorf("exposure left behind: %d", n)
	}
}