Files
privatae-bastion/client-go/lane_bounds_helpers_test.go
T
Privatae 4012da0b75
ci / personal-data (push) Successful in 10s
ci / test (push) Successful in 5m30s
ci / fuzz (push) Successful in 3m16s
bastion 0.56.0
2026-10-28 02:24:07 +00:00

127 lines
3.6 KiB
Go

package main
// The helpers lane_bounds.go and harness.go added for the reads that were
// inline before (the heartbeat reply in run.go, the agent's /run reply in
// driveHarnessAgent, the drains in harness.go and rootexec.go). They have no
// counterpart on main to run against.
import (
"bytes"
"io"
"net/http"
"net/http/httptest"
"runtime"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/gorilla/websocket"
)
// boxLaneHost stands up the HOST side of one worker socket (host runs on the
// upgraded server end) and returns the box side, dialled.
func boxLaneHost(t *testing.T, host func(c *websocket.Conn)) *websocket.Conn {
t.Helper()
up := websocket.Upgrader{}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
c, err := up.Upgrade(w, r, nil)
if err != nil {
return
}
defer c.Close()
host(c)
}))
t.Cleanup(srv.Close)
ws, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(srv.URL, "http"), nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = ws.Close() })
return ws
}
// endless is a body that never ends; served counts what was taken from it.
type endless struct{ served *atomic.Int64 }
func (e endless) Read(p []byte) (int, error) {
for i := range p {
p[i] = 'z'
}
e.served.Add(int64(len(p)))
return len(p), nil
}
func TestHeartbeatReplyIsReadOnlyToItsBound(t *testing.T) {
var served atomic.Int64
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(200)
_, _ = io.Copy(w, io.LimitReader(endless{&served}, 64<<20))
}))
defer srv.Close()
resp, err := http.Get(srv.URL)
if err != nil {
t.Fatal(err)
}
body, ok := readHeartbeatReply(resp)
resp.Body.Close()
if ok || body != nil {
t.Fatalf("a 64 MiB heartbeat reply was taken (%d bytes)", len(body))
}
small := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`{"ok":true,"dcn_grants":{"stores":[]}}`))
}))
defer small.Close()
resp, err = http.Get(small.URL)
if err != nil {
t.Fatal(err)
}
body, ok = readHeartbeatReply(resp)
resp.Body.Close()
if !ok || !strings.Contains(string(body), "dcn_grants") {
t.Fatalf("an ordinary reply was not taken: ok=%v %q", ok, body)
}
}
func TestHarnessRunReplyIsReadOnlyToItsBound(t *testing.T) {
var served atomic.Int64
if _, err := readHarnessRunReply(endless{&served}); err == nil || !strings.Contains(err.Error(), "more than the 16 MiB") {
t.Fatalf("an endless /run reply: %v", err)
}
if n := served.Load(); n > 17<<20 {
t.Fatalf("an endless /run reply was read to %d MiB", n>>20)
}
big := append([]byte(`{"ok":true,"pad":"`), bytes.Repeat([]byte{'a'}, 15<<20)...)
big = append(big, `"}`...)
if rb, err := readHarnessRunReply(bytes.NewReader(big)); err != nil || len(rb) != len(big) {
t.Fatalf("a 15 MiB reply the lane carries: %d bytes, %v", len(rb), err)
}
}
// A drained lane holds no frame whole, and ends when the host goes away.
func TestDrainLaneHoldsNoFrame(t *testing.T) {
const frame = 64 << 20
payload := bytes.Repeat([]byte{'d'}, frame)
start := make(chan struct{})
box := boxLaneHost(t, func(c *websocket.Conn) {
<-start
_ = c.WriteMessage(websocket.BinaryMessage, payload)
})
runtime.GC()
var before, after runtime.MemStats
runtime.ReadMemStats(&before)
close(start)
done := make(chan error, 1)
go func() { done <- drainLane(box) }()
select {
case <-done:
case <-time.After(20 * time.Second):
t.Fatal("drainLane did not return after the host went away")
}
runtime.ReadMemStats(&after)
if grew := after.TotalAlloc - before.TotalAlloc; grew > 8<<20 {
t.Fatalf("draining a %d MiB frame allocated %d MiB", frame>>20, grew>>20)
}
}