feat: persist client session fencing

2026-09-02 05:59:06 +00:00
parent 3e25edd337
commit 2764024d36
3 changed files with 448 additions and 0 deletions
@@ -883,6 +883,127 @@ func TestIrreparableIncidentRequiresExplicitAcknowledgement_BH_STORE_09(t *testi
 	}
 }

+func TestSessionRegistrationFencingAndExactTakeover_HP_SES_02(t *testing.T) {
+	t.Parallel()
+
+	opened := openStore(t, filepath.Join(t.TempDir(), "state"))
+	now := time.Unix(7_000_000, 0).UTC()
+	instanceA := uuidBytes(1)
+	first := registerClient(t, opened, "client-a", instanceA, uuidBytes(20), now)
+	if first.Generation != 1 {
+		t.Fatalf("first generation = %d", first.Generation)
+	}
+	if clientID, err := opened.ValidateLiveSession(context.Background(), first.SessionID, first.Generation); err != nil || clientID != "client-a" {
+		t.Fatalf("first session validation = (%q, %v)", clientID, err)
+	}
+	replacement := registerClient(t, opened, "client-a", instanceA, uuidBytes(40), now.Add(time.Second))
+	if replacement.Generation != 2 {
+		t.Fatalf("replacement generation = %d", replacement.Generation)
+	}
+	if _, err := opened.ValidateLiveSession(context.Background(), first.SessionID, first.Generation); !errors.Is(err, store.ErrSessionUnknown) {
+		t.Fatalf("fenced first session error = %v", err)
+	}
+
+	instanceB := uuidBytes(60)
+	_, err := opened.RegisterClientSession(context.Background(), store.ClientRegistration{
+		ClientID: "client-a", Platform: 3, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\work`,
+		SupportedShells: []byte{1}, ClientInstanceID: instanceB, SessionID: uuidBytes(80), ConnectedAt: now.Add(2 * time.Second),
+	})
+	if !errors.Is(err, store.ErrTakeoverRequired) {
+		t.Fatalf("different live instance error = %v", err)
+	}
+	var pending []byte
+	if err := opened.DB().QueryRow(`SELECT pending_instance_id FROM clients WHERE client_id = 'client-a'`).Scan(&pending); err != nil || !bytes.Equal(pending, instanceB[:]) {
+		t.Fatalf("pending instance = %x, err = %v", pending, err)
+	}
+	grantRequest := uuidBytes(100)
+	expires := time.Now().UTC().Add(5 * time.Minute)
+	granted, replay, err := opened.AuthorizeClientTakeover(context.Background(), store.TakeoverAuthorization{
+		ClientID: "client-a", ClientInstanceID: instanceB, RequestID: grantRequest, ExpiresAt: expires,
+	})
+	if err != nil || replay || !granted.Equal(expires) {
+		t.Fatalf("takeover grant = (%s, replay=%v, err=%v)", granted, replay, err)
+	}
+	granted, replay, err = opened.AuthorizeClientTakeover(context.Background(), store.TakeoverAuthorization{
+		ClientID: "client-a", ClientInstanceID: instanceB, RequestID: grantRequest, ExpiresAt: expires,
+	})
+	if err != nil || !replay || !granted.Equal(expires) {
+		t.Fatalf("takeover replay = (%s, replay=%v, err=%v)", granted, replay, err)
+	}
+	accepted := registerClient(t, opened, "client-a", instanceB, uuidBytes(120), now.Add(3*time.Second))
+	if accepted.Generation != 3 {
+		t.Fatalf("takeover generation = %d", accepted.Generation)
+	}
+	if _, err := opened.ValidateLiveSession(context.Background(), replacement.SessionID, replacement.Generation); !errors.Is(err, store.ErrSessionUnknown) {
+		t.Fatalf("takeover did not fence prior session: %v", err)
+	}
+	if clientID, err := opened.ValidateLiveSession(context.Background(), accepted.SessionID, accepted.Generation); err != nil || clientID != "client-a" {
+		t.Fatalf("accepted session validation = (%q, %v)", clientID, err)
+	}
+	if err := opened.Close(); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestTakeoverBoundariesAndClosedSession_BH_SES_02(t *testing.T) {
+	t.Parallel()
+
+	opened := openStore(t, filepath.Join(t.TempDir(), "state"))
+	now := time.Unix(8_000_000, 0).UTC()
+	instanceA := uuidBytes(140)
+	active := registerClient(t, opened, "client-b", instanceA, uuidBytes(160), now)
+	instanceB := uuidBytes(180)
+	_, err := opened.RegisterClientSession(context.Background(), store.ClientRegistration{
+		ClientID: "client-b", Platform: 3, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\work`,
+		SupportedShells: []byte{1}, ClientInstanceID: instanceB, SessionID: uuidBytes(200), ConnectedAt: now.Add(time.Second),
+	})
+	if !errors.Is(err, store.ErrTakeoverRequired) {
+		t.Fatalf("takeover-needed error = %v", err)
+	}
+	if _, _, err := opened.AuthorizeClientTakeover(context.Background(), store.TakeoverAuthorization{
+		ClientID: "client-b", ClientInstanceID: uuidBytes(220), RequestID: uuidBytes(10), ExpiresAt: time.Now().UTC().Add(time.Minute),
+	}); !errors.Is(err, store.ErrTakeoverMismatch) {
+		t.Fatalf("mismatched grant error = %v", err)
+	}
+	request := uuidBytes(30)
+	if _, _, err := opened.AuthorizeClientTakeover(context.Background(), store.TakeoverAuthorization{
+		ClientID: "client-b", ClientInstanceID: instanceB, RequestID: request, ExpiresAt: time.Now().UTC().Add(time.Minute),
+	}); err != nil {
+		t.Fatal(err)
+	}
+	if _, _, err := opened.AuthorizeClientTakeover(context.Background(), store.TakeoverAuthorization{
+		ClientID: "client-b", ClientInstanceID: instanceB, RequestID: uuidBytes(50), ExpiresAt: time.Now().UTC().Add(time.Minute),
+	}); !errors.Is(err, store.ErrTakeoverAlreadyGranted) {
+		t.Fatalf("second active grant error = %v", err)
+	}
+	if err := opened.CloseLiveSession(context.Background(), active.SessionID, active.Generation, "test close", now.Add(2*time.Second)); err != nil {
+		t.Fatal(err)
+	}
+	if err := opened.CloseLiveSession(context.Background(), active.SessionID, active.Generation, "again", now.Add(3*time.Second)); !errors.Is(err, store.ErrSessionUnknown) {
+		t.Fatalf("duplicate close error = %v", err)
+	}
+	// With no live session, a new instance is accepted without consuming the prior grant.
+	accepted := registerClient(t, opened, "client-b", instanceB, uuidBytes(70), now.Add(4*time.Second))
+	if accepted.Generation != 2 {
+		t.Fatalf("closed-session registration generation = %d", accepted.Generation)
+	}
+	instanceC := uuidBytes(90)
+	if _, err := opened.RegisterClientSession(context.Background(), store.ClientRegistration{
+		ClientID: "client-b", Platform: 3, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\work`,
+		SupportedShells: []byte{1}, ClientInstanceID: instanceC, SessionID: uuidBytes(110), ConnectedAt: now.Add(5 * time.Second),
+	}); !errors.Is(err, store.ErrTakeoverRequired) {
+		t.Fatalf("post-cleanup takeover-needed error = %v", err)
+	}
+	if _, _, err := opened.AuthorizeClientTakeover(context.Background(), store.TakeoverAuthorization{
+		ClientID: "client-b", ClientInstanceID: instanceC, RequestID: uuidBytes(130), ExpiresAt: time.Now().UTC().Add(time.Minute),
+	}); err != nil {
+		t.Fatalf("stale grant blocked fresh pending instance: %v", err)
+	}
+	if err := opened.Close(); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func openStore(t *testing.T, dataDir string) *store.Store {
 	t.Helper()
 	return openStoreWithOptions(t, store.Options{DataDir: dataDir, BusyTimeout: busyTimeout})
@@ -999,3 +1120,15 @@ func markTerminal(t *testing.T, database *sql.DB, issue [16]byte, issueTime, ter
 		t.Fatal(err)
 	}
 }
+
+func registerClient(t *testing.T, opened *store.Store, clientID string, instance, session [16]byte, at time.Time) store.SessionRegistration {
+	t.Helper()
+	registration, err := opened.RegisterClientSession(context.Background(), store.ClientRegistration{
+		ClientID: clientID, Platform: 3, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\work`,
+		SupportedShells: []byte{1}, ClientInstanceID: instance, SessionID: session, ConnectedAt: at,
+	})
+	if err != nil {
+		t.Fatalf("RegisterClientSession: %v", err)
+	}
+	return registration
+}