Skip to content

Commit e9cbf6d

Browse files
sarg3ntclaude
andauthored
feat(rotation): Phase 5 drift detection in agent.Client (#72) (#132)
When the dashboard signs an outbound request with kid X but the agent matches kid Y instead (because rotation propagated on one side but not the other), drift detection logs the disagreement so an operator can resync. Closes the observability loop on the install→use→remove cycle: Phase 2 added the request header, the agent's auth middleware already echoes the matched kid, and this commit wires the dashboard-side comparison. What's here ----------- `agent.HeaderKID = "X-Gearbox-Kid"` is the shared constant the dashboard sends and the agent echoes back. The agent's middleware sets it on every authenticated response (see Phase 1). `agent.Client` - `SetDriftHandler(DriftHandler)` installs an optional callback invoked when `resp.Header.Get(HeaderKID)` differs from the kid the client was built with. Reads c.onDrift at RoundTrip time, so the handler can be installed AFTER construction (typical for long- lived per-box clients held by the dashboard). - Transport wrap: `kidObservingTransport` sits between the http client and the underlying TLS transport, calling `c.checkDrift` on every successful response. One central point of inspection — no invasive edits to every `doRequest*` method. - `LogDriftHandler(logger, boxID)` builds a ready-to-use DriftHandler that emits a structured warn-level log. Lowest-friction wiring for the long-lived clients held in the WebSocketManager. Test fix -------- `TestClientTimeout` was a flaky pre-existing test whose substring check was case-sensitive — Go's net/http error message capitalises "Timeout" sometimes and emits "context deadline exceeded" other times. Fixed by lower-casing the error message before substring matching. Verified stable across 5 runs. Tests ----- `client_drift_test.go` exercises the four corners of the matrix: - Drift handler fires when kid mismatches. - Doesn't fire when kid matches. - Doesn't fire when the agent omits the header (older agents). - Doesn't fire when the dashboard client has no kid. Not yet wired into production code ---------------------------------- Adding `agent.LogDriftHandler` is the small API surface; deciding *where* to call SetDriftHandler is a separate design choice (WebSocketManager? capability poller? every short-lived handler.agentClient()?). Deferring that integration so this PR stays focused on the observability primitive itself. A follow-up can install LogDriftHandler at every site that constructs a kid- bearing client. Refs: Phase 5 of the implementation plan posted to #72. Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 7645c9b commit e9cbf6d

3 files changed

Lines changed: 182 additions & 7 deletions

File tree

gearbox/internal/framework/agent/client.go

Lines changed: 79 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"fmt"
99
"io"
1010
"log"
11+
"log/slog"
1112
"net"
1213
"net/http"
1314
"net/url"
@@ -33,17 +34,43 @@ type Client struct {
3334
baseURL string
3435
apiKey string
3536
kid string // optional; when set, sent as X-Gearbox-Kid header
37+
onDrift DriftHandler
3638
httpClient *http.Client
3739
}
3840

3941
// HeaderKID is the request/response header name carrying the keyring
4042
// entry id. The agent's middleware echoes the matched kid on every
4143
// authenticated response (see middleware.ResponseHeaderKID); the
4244
// dashboard sends this kid on outbound requests so the agent + audit
43-
// log can correlate keys. Phase 5 (drift detection) compares the
45+
// log can correlate keys. Drift detection (Phase 5) compares the
4446
// request-time kid with the echoed response kid.
4547
const HeaderKID = "X-Gearbox-Kid"
4648

49+
// DriftHandler is invoked when the agent's echoed kid differs from
50+
// the kid the client sent. Implementations typically log + surface a
51+
// "rotation drift" banner; the default is to do nothing (set via
52+
// SetDriftHandler).
53+
//
54+
// expected = what the client sent on the request
55+
// actual = what the agent echoed back on the response
56+
//
57+
// The handler is called on the request goroutine; cheap operations
58+
// only (logging is fine, network calls are not).
59+
type DriftHandler func(expected, actual string)
60+
61+
// LogDriftHandler returns a DriftHandler that emits a structured
62+
// warn-level log line — the lowest-friction wiring for a long-lived
63+
// agent.Client. The "box_id" tag lets the operator filter the journal
64+
// to a single box.
65+
func LogDriftHandler(logger *slog.Logger, boxID int64) DriftHandler {
66+
return func(expected, actual string) {
67+
logger.Warn("rotation drift: agent matched a different kid than expected",
68+
"box_id", boxID,
69+
"expected_kid", expected,
70+
"actual_kid", actual)
71+
}
72+
}
73+
4774
// NewClient creates a new HAProxy Agent API client.
4875
func NewClient(baseURL, apiKey string) *Client {
4976
return NewClientWithTimeout(baseURL, apiKey, DefaultTimeout)
@@ -74,6 +101,29 @@ func (c *Client) WithKID(kid string) *Client {
74101
// paths.
75102
func (c *Client) KID() string { return c.kid }
76103

104+
// SetDriftHandler installs a callback invoked when the agent's
105+
// echoed X-Gearbox-Kid response header differs from the kid the
106+
// client sent. Nil disables drift detection (the default).
107+
//
108+
// Use this on long-lived Client instances cached per box; the rotator
109+
// builds short-lived clients per call and wouldn't benefit from
110+
// installing a handler.
111+
func (c *Client) SetDriftHandler(h DriftHandler) { c.onDrift = h }
112+
113+
// checkDrift inspects resp's X-Gearbox-Kid header against the kid the
114+
// client sent. Called from each doRequest* path on success. No-op when
115+
// the client has no kid or no handler is installed.
116+
func (c *Client) checkDrift(resp *http.Response) {
117+
if resp == nil || c.kid == "" || c.onDrift == nil {
118+
return
119+
}
120+
actual := resp.Header.Get(HeaderKID)
121+
if actual == "" || actual == c.kid {
122+
return
123+
}
124+
c.onDrift(c.kid, actual)
125+
}
126+
77127
// setAuthHeaders applies the standard auth + Accept headers and, when
78128
// the client was built with a kid, the X-Gearbox-Kid request header.
79129
// Callers must call this BEFORE setting Content-Type so their override
@@ -135,14 +185,37 @@ func NewClientWithTimeout(baseURL, apiKey string, timeout time.Duration) *Client
135185
},
136186
}
137187

138-
return &Client{
188+
c := &Client{
139189
baseURL: baseURL,
140190
apiKey: apiKey,
141-
httpClient: &http.Client{
142-
Timeout: timeout,
143-
Transport: transport,
144-
},
145191
}
192+
c.httpClient = &http.Client{
193+
Timeout: timeout,
194+
// Wrap the base transport so every response is inspected for the
195+
// X-Gearbox-Kid header against c.kid. The transport reads c.kid
196+
// and c.onDrift at RoundTrip time, so SetDriftHandler is reflected
197+
// immediately without rebuilding the client.
198+
Transport: &kidObservingTransport{base: transport, c: c},
199+
}
200+
return c
201+
}
202+
203+
// kidObservingTransport wraps an http.RoundTripper and invokes the
204+
// owning Client's checkDrift on every successful response. Implements
205+
// the drift-detection half of Phase 5 (issue #72) — the dashboard
206+
// learns immediately when an agent's matched kid disagrees with the
207+
// kid the dashboard sent, signalling a partial rotation.
208+
type kidObservingTransport struct {
209+
base http.RoundTripper
210+
c *Client
211+
}
212+
213+
func (t *kidObservingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
214+
resp, err := t.base.RoundTrip(req)
215+
if err == nil {
216+
t.c.checkDrift(resp)
217+
}
218+
return resp, err
146219
}
147220

148221
// BuildTLSConfig is the exported alias for createTLSConfig — used by
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
package agent
2+
3+
import (
4+
"net/http"
5+
"net/http/httptest"
6+
"net/url"
7+
"sync/atomic"
8+
"testing"
9+
)
10+
11+
func TestDriftHandler_FiresOnKIDMismatch(t *testing.T) {
12+
// Mock agent that always echoes a hard-coded kid in the response.
13+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
14+
w.Header().Set(HeaderKID, "actual-kid")
15+
w.WriteHeader(http.StatusOK)
16+
_, _ = w.Write([]byte("{}"))
17+
}))
18+
defer srv.Close()
19+
20+
c := NewClientWithKID(srv.URL, "ignored", "expected-kid")
21+
var fired atomic.Bool
22+
var seenExpected, seenActual string
23+
c.SetDriftHandler(func(expected, actual string) {
24+
fired.Store(true)
25+
seenExpected, seenActual = expected, actual
26+
})
27+
28+
if _, err := c.doRequest("GET", "/anything", url.Values{}); err != nil {
29+
t.Fatalf("doRequest: %v", err)
30+
}
31+
if !fired.Load() {
32+
t.Errorf("drift handler not invoked")
33+
}
34+
if seenExpected != "expected-kid" || seenActual != "actual-kid" {
35+
t.Errorf("drift handler args: expected=%q actual=%q", seenExpected, seenActual)
36+
}
37+
}
38+
39+
func TestDriftHandler_DoesNotFireOnMatch(t *testing.T) {
40+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
41+
w.Header().Set(HeaderKID, "matched")
42+
w.WriteHeader(http.StatusOK)
43+
_, _ = w.Write([]byte("{}"))
44+
}))
45+
defer srv.Close()
46+
47+
c := NewClientWithKID(srv.URL, "ignored", "matched")
48+
var fired atomic.Bool
49+
c.SetDriftHandler(func(_, _ string) { fired.Store(true) })
50+
51+
if _, err := c.doRequest("GET", "/anything", url.Values{}); err != nil {
52+
t.Fatalf("doRequest: %v", err)
53+
}
54+
if fired.Load() {
55+
t.Errorf("drift handler incorrectly invoked on matching kid")
56+
}
57+
}
58+
59+
func TestDriftHandler_DoesNotFireWhenAgentDoesNotEchoKID(t *testing.T) {
60+
// Older agents won't set X-Gearbox-Kid. The drift handler should
61+
// stay silent rather than flagging every response as "drift".
62+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
63+
w.WriteHeader(http.StatusOK)
64+
_, _ = w.Write([]byte("{}"))
65+
}))
66+
defer srv.Close()
67+
68+
c := NewClientWithKID(srv.URL, "ignored", "expected-kid")
69+
var fired atomic.Bool
70+
c.SetDriftHandler(func(_, _ string) { fired.Store(true) })
71+
72+
if _, err := c.doRequest("GET", "/anything", url.Values{}); err != nil {
73+
t.Fatalf("doRequest: %v", err)
74+
}
75+
if fired.Load() {
76+
t.Errorf("drift handler incorrectly invoked when agent sent no kid header")
77+
}
78+
}
79+
80+
func TestDriftHandler_DoesNotFireWhenClientHasNoKID(t *testing.T) {
81+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
82+
w.Header().Set(HeaderKID, "something")
83+
w.WriteHeader(http.StatusOK)
84+
_, _ = w.Write([]byte("{}"))
85+
}))
86+
defer srv.Close()
87+
88+
c := NewClient(srv.URL, "ignored") // no kid
89+
var fired atomic.Bool
90+
c.SetDriftHandler(func(_, _ string) { fired.Store(true) })
91+
92+
if _, err := c.doRequest("GET", "/anything", url.Values{}); err != nil {
93+
t.Fatalf("doRequest: %v", err)
94+
}
95+
if fired.Load() {
96+
t.Errorf("drift handler invoked when client has no kid; should no-op")
97+
}
98+
}

gearbox/internal/framework/agent/client_test.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -579,7 +579,11 @@ func TestClientTimeout(t *testing.T) {
579579
t.Fatal("expected timeout error")
580580
}
581581

582-
if !strings.Contains(err.Error(), "context deadline exceeded") && !strings.Contains(err.Error(), "timeout") {
582+
// The error message format varies by Go version + timing: sometimes
583+
// "context deadline exceeded", sometimes "Client.Timeout exceeded
584+
// while awaiting headers" (capital T). Match either by lower-casing.
585+
msg := strings.ToLower(err.Error())
586+
if !strings.Contains(msg, "context deadline exceeded") && !strings.Contains(msg, "timeout") {
583587
t.Errorf("expected timeout error, got %v", err)
584588
}
585589
}

0 commit comments

Comments
 (0)