Files
privatae-bastion/client-go/lanehold_test.go
T
Privatae 48899b41a9
ci / personal-data (push) Successful in 4s
ci / test (push) Successful in 4m58s
ci / fuzz (push) Successful in 3m33s
bastion 0.63.1
2026-11-03 22:56:18 +00:00

252 lines
8.2 KiB
Go

package main
import (
"context"
"crypto/tls"
"crypto/x509"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/gorilla/websocket"
)
// heldHost stands in for a Host's lane port: it counts each worker's hello,
// then either closes the socket with the code a Host gives a machine it holds,
// or parks it.
type heldHost struct {
hellos atomic.Int32
hold atomic.Bool
// together: a held socket is closed only once this many hellos came in
// (or after barrierWait), so every worker of a pool has dialled before
// the first refusal starts the hold. 0 closes each at once.
together atomic.Int32
tls *tls.Config
dial func(ctx context.Context) (net.Conn, error)
url string
}
func startHeldHost(t *testing.T) *heldHost {
t.Helper()
h := &heldHost{}
h.hold.Store(true)
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 {
return
}
h.hellos.Add(1)
if h.hold.Load() {
for end := time.Now().Add(barrierWait); h.hellos.Load() < h.together.Load() && time.Now().Before(end); {
time.Sleep(5 * time.Millisecond)
}
_ = c.WriteControl(websocket.CloseMessage,
websocket.FormatCloseMessage(laneHeldCloseCode, "this bastion is quarantined"), time.Now().Add(time.Second))
return
}
for {
if _, _, err := c.ReadMessage(); err != nil {
return
}
}
}))
srv.StartTLS()
t.Cleanup(srv.Close)
pool := x509.NewCertPool()
pool.AddCert(srv.Certificate())
h.tls = &tls.Config{RootCAs: pool, ServerName: "example.com"}
addr := srv.Listener.Addr().String()
h.dial = func(ctx context.Context) (net.Conn, error) {
var d net.Dialer
return d.DialContext(ctx, "tcp", addr)
}
h.url = "wss://example.com/lane"
return h
}
// barrierWait bounds how long a held socket waits for the rest of its pool.
const barrierWait = 15 * time.Second
// shortHolds makes the holds short enough for a test. Set before any worker
// runs, and never from a parallel subtest.
func shortHolds(t *testing.T, first, max time.Duration) {
t.Helper()
f, m := laneHoldFirst, laneHoldMax
laneHoldFirst, laneHoldMax = first, max
t.Cleanup(func() { laneHoldFirst, laneHoldMax = f, m })
}
// ownHold gives a test its own pairing's hold, and forgets it after.
func ownHold(t *testing.T, host string) {
t.Helper()
laneHolds.Delete(host)
t.Cleanup(func() { laneHolds.Delete(host) })
}
// Every lane's pool, so a lane added later without the hold is caught.
func lanePools(dir string) map[string]func(ctx context.Context, host string, h *heldHost, n int) {
return map[string]func(ctx context.Context, host string, h *heldHost, n int){
"egress": func(ctx context.Context, host string, h *heldHost, n int) {
var d net.Dialer
runEgressWorkers(ctx, host, h.url, h.tls, h.dial, "b", d.DialContext, n)
},
"serve": func(ctx context.Context, host string, h *heldHost, n int) {
runServeWorkers(ctx, host, dir, h.url, h.tls, h.dial, "b", n)
},
"mgmt": func(ctx context.Context, host string, h *heldHost, n int) {
runMgmtWorkers(ctx, host, h.url, h.tls, h.dial, "b", dir, dir, dir, n)
},
"harness": func(ctx context.Context, host string, h *heldHost, n int) {
runHarnessWorkers(ctx, host, dir, h.url, h.tls, h.dial, "b", n)
},
"session": func(ctx context.Context, host string, h *heldHost, n int) {
runSessionWorkers(ctx, host, h.url, h.tls, h.dial, "b", n,
sessionEnv{CtxRoot: dir, DcnDir: dir, HostName: host, Server: "test"})
},
}
}
// A Host that holds this machine closes every worker socket with 4403. No
// worker of any lane dials it again before the hold is over: one hello per
// worker, where a pool used to dial again at once or a second later, 30 to 80
// times a second from one held machine.
//
// The Host closes the sockets only once every worker has said hello
// (together). Closing the first at once made the test race the pool's own
// start: a worker not yet past its first waitClear when another was refused
// rightly waited out the hold before dialling, and the count read one of two
// (a full gate, 2026-11-03). The hold is timed from the refusals, which come
// after the last hello, so the 2 s look is inside the 3 s hold however slowly
// the workers start.
//
// MUTATION: drop hold.refused from a lane's worker (mgmt.go). RESULT: red,
// that lane dials again within the hold.
func TestEveryLanePoolWaitsOutAHostThatHoldsThisMachine(t *testing.T) {
shortHolds(t, 3*time.Second, time.Minute)
for lane, run := range lanePools(t.TempDir()) {
t.Run(lane, func(t *testing.T) {
t.Parallel()
host := "held-" + lane + ".example.test"
ownHold(t, host)
h := startHeldHost(t)
const n = 2
h.together.Store(n)
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
go run(ctx, host, h, n)
for end := time.Now().Add(barrierWait); h.hellos.Load() < n && time.Now().Before(end); {
time.Sleep(5 * time.Millisecond)
}
if got := h.hellos.Load(); got < n {
t.Fatalf("only %d of %d workers dialled the Host in %s", got, n, barrierWait)
}
time.Sleep(2 * time.Second)
if got := h.hellos.Load(); got != n {
t.Fatalf("%d dials within the hold from %d workers; want one each", got, n)
}
})
}
}
// A beat whose reply says quarantined stops the lanes before any of them
// dials, and the first beat that no longer says it brings them back at once.
func TestABeatThatSaysQuarantinedStopsTheLanesUntilOneDoesNot(t *testing.T) {
host := "quarantined.example.test"
shortHolds(t, time.Hour, time.Hour)
ownHold(t, host)
h := startHeldHost(t)
h.hold.Store(false)
laneHoldFor(host).fromBeat(host, beatSaysQuarantined([]byte(`{"ok":true,"quarantined":true}`)))
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
const n = 2
go lanePools(t.TempDir())["mgmt"](ctx, host, h, n)
time.Sleep(700 * time.Millisecond)
if got := h.hellos.Load(); got != 0 {
t.Fatalf("%d dials while the Host says it holds this machine", got)
}
laneHoldFor(host).fromBeat(host, beatSaysQuarantined([]byte(`{"ok":true}`)))
deadline := time.Now().Add(3 * time.Second)
for h.hellos.Load() < n && time.Now().Before(deadline) {
time.Sleep(20 * time.Millisecond)
}
if got := h.hellos.Load(); got != n {
t.Fatalf("after the release, %d of %d workers parked again", got, n)
}
}
// A pool's workers are refused together, so a refusal inside a hold is that
// hold's; each hold after one runs out is twice as long, up to the most; a
// refusal long after the last starts again from the first; a release ends a
// hold at once.
func TestAHoldDoublesOncePerHoldUpToItsMost(t *testing.T) {
host := "doubling.example.test"
shortHolds(t, time.Second, 4*time.Second)
ownHold(t, host)
h := laneHoldFor(host)
expire := func() {
h.mu.Lock()
h.until = time.Now().Add(-time.Millisecond)
h.mu.Unlock()
}
wait := func() time.Duration {
h.mu.Lock()
defer h.mu.Unlock()
return h.wait
}
h.refused(host)
h.refused(host)
h.refused(host)
if got := wait(); got != time.Second {
t.Fatalf("three refusals in one hold: %s, want the first hold", got)
}
var got []string
for range 3 {
expire()
h.refused(host)
got = append(got, wait().String())
}
if strings.Join(got, " ") != "2s 4s 4s" {
t.Fatalf("holds after the first: %v, want 2s 4s 4s", got)
}
expire()
h.mu.Lock()
h.last = time.Now().Add(-9 * time.Second)
h.mu.Unlock()
h.refused(host)
if got := wait(); got != time.Second {
t.Fatalf("a refusal after a long quiet: %s, want the first hold again", got)
}
h.fromBeat(host, true)
h.fromBeat(host, false)
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
if !h.waitClear(ctx) {
t.Fatal("a release must end the hold at once")
}
}
func TestBeatSaysQuarantinedReadsOnlyTheHostsTrue(t *testing.T) {
for body, want := range map[string]bool{
`{"ok":true,"quarantined":true}`: true,
`{"ok":true,"quarantined":false}`: false,
`{"ok":true}`: false,
`{"quarantined":"yes"}`: false,
`not json`: false,
``: false,
} {
if got := beatSaysQuarantined([]byte(body)); got != want {
t.Errorf("%q: %v, want %v", body, got, want)
}
}
}