Add RestoreHijackPrevention and RestoreSession interface stub (#4405)

* Add RestoreHijackPrevention and RestoreSession interface stub (RC-15, #4216)

Add the infrastructure needed to reconstruct hijack-prevention state from
Redis-persisted metadata, enabling cross-pod token validation in horizontal
scaling scenarios.

- security.go: Add RestoreHijackPrevention(), the restore counterpart to
  PreventSessionHijacking(). Rebuilds a hijackPreventionDecorator from
  persisted tokenHash + tokenSaltHex + hmacSecret without re-hashing a
  live token. Returns errors for nil session, missing salt on authenticated
  sessions, and invalid hex salt.

- factory.go: Add RestoreSession() to the MultiSessionFactory interface with
  full doc comment (backend ID parsing, session hint lookup, routing-table
  rebuild, hijack-prevention re-application). Add a stub implementation on
  defaultMultiSessionFactory (returns "not yet implemented"); full
  reconnection logic is deferred. Document the cross-replica HMAC secret
  consistency requirement on defaultHMACSecret.

- decorating_factory.go: Forward RestoreSession() to the base factory;
  decorators are not re-applied during restore.

- mocks/mock_factory.go: Regenerate mock to include RestoreSession().

- restore_test.go: Unit tests covering nil session, missing salt, invalid
  hex, anonymous session round-trip, authenticated store→restore→validate
  round-trip, and cross-replica secret mismatch.

Closes #4216.

* changes from rebase

---------

Co-authored-by: taskbot <taskbot@users.noreply.github.com>

Author: yrobla

pkg/vmcp/session/connector_integration_test.go

@@ -21,6 +21,7 @@ import (
 	vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth"
 	"github.com/stacklok/toolhive/pkg/vmcp/auth/strategies"
 	authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
+	"github.com/stacklok/toolhive/pkg/vmcp/session/internal/security"
 	sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types"
 )
 
@@ -347,6 +348,96 @@ func TestTokenBinding_DifferentSecretsProduceDifferentHashes(t *testing.T) {
 		"different HMAC secrets must produce different token hashes for the same input token")
 }
 
+// TestRestoreHijackPrevention_Integration_RoundTrip verifies the full
+// store-then-restore flow across a real factory-created session:
+//
+//  1. Create a session via the factory (writes tokenHash + tokenSalt to metadata).
+//  2. Extract the persisted values.
+//  3. Wrap a fresh base session with RestoreHijackPrevention using those values.
+//  4. Confirm the restored decorator accepts the original token and rejects others.
+func TestRestoreHijackPrevention_Integration_RoundTrip(t *testing.T) {
+	t.Parallel()
+
+	const rawToken = "integration-token"
+	hmacSecret := []byte("test-hmac-secret-exactly-32bytes")
+	identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: "alice"}, Token: rawToken}
+
+	factory := newSessionFactoryWithConnector(nilBackendConnector(), WithHMACSecret(hmacSecret))
+	sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), identity, false, nil)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = sess.Close() })
+
+	// Extract persisted values — these simulate what would be read back from Redis.
+	meta := sess.GetMetadata()
+	persistedHash := meta[MetadataKeyTokenHash]
+	persistedSalt := meta[sessiontypes.MetadataKeyTokenSalt]
+	require.NotEmpty(t, persistedHash, "factory must write tokenHash to metadata")
+	require.NotEmpty(t, persistedSalt, "factory must write tokenSalt to metadata")
+
+	// Simulate "Pod B": restore the decorator from persisted metadata.
+	// We use a nil-connector session as the inner session (no real backend needed
+	// to test auth path).
+	innerSess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), identity, false, nil)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = innerSess.Close() })
+
+	restored, err := security.RestoreHijackPrevention(innerSess, persistedHash, persistedSalt, hmacSecret)
+	require.NoError(t, err)
+
+	ctx := context.Background()
+
+	// Original caller is accepted.
+	_, err = restored.CallTool(ctx, identity, "any-tool", nil, nil)
+	// ErrToolNotFound is expected (no backends), not an auth error.
+	require.NotErrorIs(t, err, sessiontypes.ErrUnauthorizedCaller)
+	require.NotErrorIs(t, err, sessiontypes.ErrNilCaller)
+
+	// A different caller is rejected at the auth layer — before any backend routing.
+	wrongCaller := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: "eve"}, Token: "eve-token"}
+	_, err = restored.CallTool(ctx, wrongCaller, "any-tool", nil, nil)
+	require.ErrorIs(t, err, sessiontypes.ErrUnauthorizedCaller)
+
+	// Nil caller is rejected at the auth layer.
+	_, err = restored.CallTool(ctx, nil, "any-tool", nil, nil)
+	require.ErrorIs(t, err, sessiontypes.ErrNilCaller)
+}
+
+// TestRestoreHijackPrevention_Integration_CrossReplicaSecretMismatch verifies
+// that a session restored on a replica with a different HMAC secret rejects
+// the original caller's token, documenting the operational requirement that
+// all replicas must share the same secret.
+func TestRestoreHijackPrevention_Integration_CrossReplicaSecretMismatch(t *testing.T) {
+	t.Parallel()
+
+	secretA := []byte("secret-A-exactly-32-bytes-long!!")
+	secretB := []byte("secret-B-exactly-32-bytes-long!!")
+
+	identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: "alice"}, Token: "alice-token"}
+
+	// Pod A creates the session with secretA, persisting the hash.
+	factoryA := newSessionFactoryWithConnector(nilBackendConnector(), WithHMACSecret(secretA))
+	sessA, err := factoryA.MakeSessionWithID(context.Background(), uuid.New().String(), identity, false, nil)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = sessA.Close() })
+
+	persistedHash := sessA.GetMetadata()[MetadataKeyTokenHash]
+	persistedSalt := sessA.GetMetadata()[sessiontypes.MetadataKeyTokenSalt]
+
+	// Pod B restores with secretB — the persisted hash was computed with secretA,
+	// so validation will produce a different HMAC and reject the caller.
+	factoryB := newSessionFactoryWithConnector(nilBackendConnector(), WithHMACSecret(secretB))
+	innerSess, err := factoryB.MakeSessionWithID(context.Background(), uuid.New().String(), identity, false, nil)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = innerSess.Close() })
+
+	restored, err := security.RestoreHijackPrevention(innerSess, persistedHash, persistedSalt, secretB)
+	require.NoError(t, err)
+
+	_, err = restored.CallTool(context.Background(), identity, "any-tool", nil, nil)
+	require.ErrorIs(t, err, sessiontypes.ErrUnauthorizedCaller,
+		"cross-replica secret mismatch must reject the original caller")
+}
+
 // TestTokenBinding_MetadataEncoding verifies that the token hash and salt stored
 // in session metadata are valid hex strings of the expected lengths:
 //   - token hash: 64 hex chars (32-byte HMAC-SHA256)

pkg/vmcp/session/decorating_factory.go

@@ -44,24 +44,7 @@ func (f *decoratingMultiSessionFactory) RestoreSession(
 	if err != nil {
 		return nil, err
 	}
-	for _, dec := range f.decorators {
-		var decorated MultiSession
-		decorated, err = dec(ctx, sess)
-		if err != nil {
-			if closeErr := sess.Close(); closeErr != nil {
-				slog.Warn("failed to close session after decorator error", "error", closeErr)
-			}
-			return nil, err
-		}
-		if decorated == nil {
-			if closeErr := sess.Close(); closeErr != nil {
-				slog.Warn("failed to close session after decorator returned nil", "error", closeErr)
-			}
-			return nil, fmt.Errorf("decorator returned nil session without error")
-		}
-		sess = decorated
-	}
-	return sess, nil
+	return f.applyDecorators(ctx, sess)
 }
 
 func (f *decoratingMultiSessionFactory) MakeSessionWithID(
@@ -75,6 +58,13 @@ func (f *decoratingMultiSessionFactory) MakeSessionWithID(
 	if err != nil {
 		return nil, err
 	}
+	return f.applyDecorators(ctx, sess)
+}
+
+// applyDecorators runs the decorator chain over sess in order, closing sess on
+// any error and returning the fully-decorated session on success.
+func (f *decoratingMultiSessionFactory) applyDecorators(ctx context.Context, sess MultiSession) (MultiSession, error) {
+	var err error
 	for _, dec := range f.decorators {
 		var decorated MultiSession
 		decorated, err = dec(ctx, sess)

pkg/vmcp/session/decorating_factory_test.go

@@ -160,3 +160,158 @@ func TestNewDecoratingFactory_HappyPath_ReturnsFinalSession(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, finalSess, got)
 }
+
+// ---------------------------------------------------------------------------
+// RestoreSession — mirrors the MakeSessionWithID tests above
+// ---------------------------------------------------------------------------
+
+func TestRestoreSession_BaseError_Propagated(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl)
+	baseErr := errors.New("restore failed")
+	base.EXPECT().
+		RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(nil, baseErr)
+
+	factory := session.NewDecoratingFactory(base,
+		func(_ context.Context, s session.MultiSession) (session.MultiSession, error) { return s, nil },
+	)
+
+	_, err := factory.RestoreSession(context.Background(), "id", nil, nil)
+	require.ErrorIs(t, err, baseErr)
+}
+
+func TestRestoreSession_DecoratorsAppliedInOrder(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	sess := sessionmocks.NewMockMultiSession(ctrl)
+	base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl)
+	base.EXPECT().
+		RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(sess, nil)
+
+	var order []int
+	dec1 := func(_ context.Context, s session.MultiSession) (session.MultiSession, error) {
+		order = append(order, 1)
+		return s, nil
+	}
+	dec2 := func(_ context.Context, s session.MultiSession) (session.MultiSession, error) {
+		order = append(order, 2)
+		return s, nil
+	}
+
+	factory := session.NewDecoratingFactory(base, dec1, dec2)
+	_, err := factory.RestoreSession(context.Background(), "id", nil, nil)
+	require.NoError(t, err)
+	assert.Equal(t, []int{1, 2}, order)
+}
+
+func TestRestoreSession_DecoratorError_ClosesSession(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	sess := sessionmocks.NewMockMultiSession(ctrl)
+	sess.EXPECT().Close().Return(nil)
+
+	base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl)
+	base.EXPECT().
+		RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(sess, nil)
+
+	decErr := errors.New("decorator boom")
+	factory := session.NewDecoratingFactory(base, func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) {
+		return nil, decErr
+	})
+
+	_, err := factory.RestoreSession(context.Background(), "id", nil, nil)
+	require.ErrorIs(t, err, decErr)
+}
+
+func TestRestoreSession_SecondDecoratorError_ClosesCurrentSession(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	sess := sessionmocks.NewMockMultiSession(ctrl)
+	wrappedSess := sessionmocks.NewMockMultiSession(ctrl)
+	// Only the session that is current at the time of failure should be closed.
+	wrappedSess.EXPECT().Close().Return(nil)
+
+	base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl)
+	base.EXPECT().
+		RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(sess, nil)
+
+	decErr := errors.New("second decorator boom")
+	dec1 := func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) { return wrappedSess, nil }
+	dec2 := func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) { return nil, decErr }
+
+	factory := session.NewDecoratingFactory(base, dec1, dec2)
+	_, err := factory.RestoreSession(context.Background(), "id", nil, nil)
+	require.ErrorIs(t, err, decErr)
+}
+
+func TestRestoreSession_NilReturnWithNoError_ClosesSession(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	sess := sessionmocks.NewMockMultiSession(ctrl)
+	sess.EXPECT().Close().Return(nil)
+
+	base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl)
+	base.EXPECT().
+		RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(sess, nil)
+
+	factory := session.NewDecoratingFactory(base, func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) {
+		return nil, nil // buggy decorator
+	})
+
+	_, err := factory.RestoreSession(context.Background(), "id", nil, nil)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "nil session")
+}
+
+func TestRestoreSession_CloseErrorDoesNotSuppressOriginalError(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	sess := sessionmocks.NewMockMultiSession(ctrl)
+	sess.EXPECT().Close().Return(errors.New("close failed"))
+
+	base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl)
+	base.EXPECT().
+		RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(sess, nil)
+
+	decErr := errors.New("decorator error")
+	factory := session.NewDecoratingFactory(base, func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) {
+		return nil, decErr
+	})
+
+	_, err := factory.RestoreSession(context.Background(), "id", nil, nil)
+	require.ErrorIs(t, err, decErr)
+}
+
+func TestRestoreSession_HappyPath_ReturnsFinalSession(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	sess := sessionmocks.NewMockMultiSession(ctrl)
+	finalSess := sessionmocks.NewMockMultiSession(ctrl)
+
+	base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl)
+	base.EXPECT().
+		RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(sess, nil)
+
+	factory := session.NewDecoratingFactory(base,
+		func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) { return finalSess, nil },
+	)
+
+	got, err := factory.RestoreSession(context.Background(), "id", nil, nil)
+	require.NoError(t, err)
+	assert.Equal(t, finalSess, got)
+}

pkg/vmcp/session/factory.go

@@ -50,6 +50,11 @@ var (
 	// defaultHMACSecret is the fallback HMAC secret used when WithHMACSecret is not provided.
 	// WARNING: This is INSECURE and should ONLY be used for testing/development.
 	// Production deployments MUST provide a secure secret via WithHMACSecret option.
+	//
+	// NOTE: In multi-replica deployments, all replicas must use the same HMAC secret,
+	// injected via the VMCP_SESSION_HMAC_SECRET environment variable. If replicas use
+	// different secrets, cross-pod token validation will silently reject legitimate
+	// callers. The default insecure secret must NOT be used in production.
 	defaultHMACSecret = []byte("insecure-default-for-testing-only-change-in-production")
 )
 
@@ -582,7 +587,7 @@ func (f *defaultMultiSessionFactory) RestoreSession(
 		return nil, fmt.Errorf("RestoreSession: token hash metadata key absent (corrupted session metadata)")
 	}
 	storedSalt := storedMetadata[sessiontypes.MetadataKeyTokenSalt]
-	restored, err := security.RestoreHijackPrevention(baseSession, f.hmacSecret, storedHash, storedSalt)
+	restored, err := security.RestoreHijackPrevention(baseSession, storedHash, storedSalt, f.hmacSecret)
 	if err != nil {
 		_ = baseSession.Close()
 		return nil, fmt.Errorf("RestoreSession: failed to restore hijack prevention: %w", err)

pkg/vmcp/session/internal/security/hijack_prevention_test.go

@@ -227,7 +227,7 @@ func TestRestoreHijackPrevention(t *testing.T) {
 		t.Parallel()
 
 		base := newMockSession("s1")
-		restored, err := RestoreHijackPrevention(base, testSecret, "", "")
+		restored, err := RestoreHijackPrevention(base, "", "", testSecret)
 		require.NoError(t, err)
 		require.NotNil(t, restored)
 	})
@@ -236,7 +236,7 @@ func TestRestoreHijackPrevention(t *testing.T) {
 		t.Parallel()
 
 		base := newMockSession("s2")
-		_, err := RestoreHijackPrevention(base, testSecret, "somehash", "")
+		_, err := RestoreHijackPrevention(base, "somehash", "", testSecret)
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "salt is missing")
 	})
@@ -245,15 +245,15 @@ func TestRestoreHijackPrevention(t *testing.T) {
 		t.Parallel()
 
 		base := newMockSession("s3")
-		_, err := RestoreHijackPrevention(base, testSecret, "", "deadbeef")
+		_, err := RestoreHijackPrevention(base, "", "deadbeef", testSecret)
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "hash is missing")
 	})
 
 	t.Run("nil session is rejected", func(t *testing.T) {
 		t.Parallel()
 
-		_, err := RestoreHijackPrevention(nil, testSecret, "", "")
+		_, err := RestoreHijackPrevention(nil, "", "", testSecret)
 		require.Error(t, err)
 	})
 }