bastion 0.56.0

client-go/egress_relay_test.go

@@ -0,0 +1,166 @@
+//go:build !windows
+
+// The egress worker pool (egress_relay.go: runEgressWorkers) against a
+// stand-in Host's /bastion-egress: how fast workers park again after a job
+// they took and after one they refused. Uses localEgressFixture
+// (egress_local_test.go), so it carries the same build tag.
+
+package main
+
+import (
+	"context"
+	"io"
+	"log"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// egressHost is a Host's /bastion-egress as the pool meets it: every worker
+// that parks is counted and handed the job job(n) (n is its place in line),
+// and what the bastion answered is read to its end.
+type egressHost struct {
+	parked atomic.Int32
+	ok     atomic.Int32 // answers that began "HTTP/1.1 200"
+}
+
+func startEgressHost(t *testing.T, job func(n int32) string) (*egressHost, string, func(context.Context) (net.Conn, error)) {
+	t.Helper()
+	h := &egressHost{}
+	up := websocket.Upgrader{}
+	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c, err := up.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer c.Close()
+		if _, _, err := c.ReadMessage(); err != nil { // BEARER
+			return
+		}
+		n := h.parked.Add(1)
+		if err := c.WriteMessage(websocket.BinaryMessage, append([]byte{'D'}, job(n)...)); err != nil {
+			return
+		}
+		_ = c.SetReadDeadline(time.Now().Add(5 * time.Second))
+		var answer strings.Builder
+		for {
+			_, m, err := c.ReadMessage()
+			if err != nil || len(m) == 0 || m[0] == 'E' {
+				break
+			}
+			answer.Write(m[1:])
+		}
+		if strings.HasPrefix(answer.String(), "HTTP/1.1 200") {
+			h.ok.Add(1)
+		}
+	}))
+	srv.Config.ErrorLog = log.New(io.Discard, "", 0)
+	srv.Start()
+	t.Cleanup(srv.Close)
+	addr := srv.Listener.Addr().String()
+	dial := func(ctx context.Context) (net.Conn, error) {
+		return (&net.Dialer{}).DialContext(ctx, "tcp", addr)
+	}
+	return h, "ws://" + addr + "/bastion-egress", dial
+}
+
+// runEgressPool runs egressWorkerCount workers for the test's life, and at its
+// end stops them and waits until each has exited: a worker still vetting a job
+// would otherwise read the fixture's stubs while the test's cleanup puts them
+// back.
+func runEgressPool(t *testing.T, wsURL string, dialHost func(context.Context) (net.Conn, error), dialOut egressDialer) {
+	t.Helper()
+	ctx, cancel := context.WithCancel(context.Background())
+	wait := runEgressWorkers(ctx, "home.host.example", wsURL, nil, dialHost, "test-bearer", dialOut, egressWorkerCount)
+	t.Cleanup(func() {
+		cancel()
+		done := make(chan struct{})
+		go func() { wait(); close(done) }()
+		select {
+		case <-done:
+		case <-time.After(10 * time.Second):
+			t.Error("the egress workers did not exit 10 s after the pool stopped")
+		}
+	})
+}
+
+// SECURITY (this machine's resources): a Host that answers every parked egress
+// worker with a job this machine refuses (a request for this machine, a line
+// that is no request) cannot make it dial, handshake and refuse in a loop as
+// fast as the Host answers. Each worker that refused waits before it parks
+// again, so the pool costs at most egressWorkerCount dials a second.
+//
+// MUTATION: in runEgressWorkers, park again at once whatever handleEgressConn
+// reports (as the pool did).
+// RESULT: red — hundreds of workers park in the 2 s window.
+func TestARefusingHostCannotSpinTheEgressWorkers(t *testing.T) {
+	localEgressFixture(t)
+	var dialled atomic.Int32
+	dialOut := func(ctx context.Context, network, addr string) (net.Conn, error) {
+		dialled.Add(1)
+		return nil, io.EOF
+	}
+	h, wsURL, dialHost := startEgressHost(t, func(n int32) string {
+		if n%2 == 0 {
+			return "NOT-A-REQUEST\r\n"
+		}
+		return "CONNECT localhost:443 HTTP/1.1\r\nHost: localhost:443\r\n\r\n"
+	})
+	runEgressPool(t, wsURL, dialHost, dialOut)
+	time.Sleep(2 * time.Second)
+	// egressWorkerCount park at the start, then each worker once a second:
+	// about 12 in 2 s. A pool that parks again at once makes hundreds.
+	n := h.parked.Load()
+	t.Logf("%d workers parked in 2 s", n)
+	if limit := int32(egressWorkerCount * 4); n > limit {
+		t.Fatalf("%d egress workers parked in 2 s on a Host that refuses every job (at most %d): the pool re-dials in a loop", n, limit)
+	}
+	if n < egressWorkerCount {
+		t.Fatalf("only %d workers parked: the pool did not run, so the count proves nothing", n)
+	}
+	if d := dialled.Load(); d != 0 {
+		t.Fatalf("a refused job was dialled out %d times", d)
+	}
+}
+
+// The person's own traffic is not slowed: a worker that took a job parks again
+// at once, so 40 jobs through 4 workers finish in a fraction of the time one
+// wait per job would take (40 / 4 x egressRefusedWait = 10 s).
+//
+// MUTATION: in runEgressWorkers, wait egressRefusedWait after every job.
+// RESULT: red — 40 jobs need about 10 s.
+func TestEgressWorkersParkAgainAtOnceAfterAJobTheyTook(t *testing.T) {
+	localEgressFixture(t)
+	dialOut := func(ctx context.Context, network, addr string) (net.Conn, error) {
+		if addr != "198.51.100.9:443" {
+			t.Errorf("dialled %q, not the address the request named", addr)
+		}
+		ours, theirs := net.Pipe()
+		_ = theirs.Close() // the destination accepts and ends at once
+		return ours, nil
+	}
+	h, wsURL, dialHost := startEgressHost(t, func(int32) string {
+		return "CONNECT 198.51.100.9:443 HTTP/1.1\r\nHost: 198.51.100.9:443\r\n\r\n"
+	})
+	start := time.Now()
+	runEgressPool(t, wsURL, dialHost, dialOut)
+	const jobs = 40
+	deadline := start.Add(10 * time.Second)
+	for h.ok.Load() < jobs && time.Now().Before(deadline) {
+		time.Sleep(10 * time.Millisecond)
+	}
+	took := time.Since(start)
+	t.Logf("%d accepted jobs answered in %v", h.ok.Load(), took)
+	if got := h.ok.Load(); got < jobs {
+		t.Fatalf("only %d of %d accepted jobs were answered 200 in %v", got, jobs, took)
+	}
+	if limit := 3 * egressRefusedWait; took > limit {
+		t.Fatalf("%d accepted jobs took %v (over %v): a worker that took a job waited before it parked again", jobs, took, limit)
+	}
+}