fixes from review

Author: taskbot

pkg/vmcp/server/session_management_v2_integration_test.go

@@ -105,25 +105,50 @@ func newV2FakeFactory(tools []vmcp.Tool) *v2FakeMultiSessionFactory {
 }
 
 func (f *v2FakeMultiSessionFactory) MakeSession(
-	_ context.Context, _ *auth.Identity, _ []*vmcp.Backend,
+	_ context.Context, identity *auth.Identity, _ []*vmcp.Backend,
 ) (vmcpsession.MultiSession, error) {
 	if f.err != nil {
 		return nil, f.err
 	}
 	baseSession := transportsession.NewStreamableSession("auto-id")
+
+	// Set basic metadata to indicate whether this is an anonymous session.
+	// Integration tests don't need to verify crypto implementation details.
+	allowAnonymous := vmcpsession.ShouldAllowAnonymous(identity)
+	if !allowAnonymous {
+		// Authenticated session - set non-empty hash placeholder
+		baseSession.SetMetadata(vmcpsession.MetadataKeyTokenHash, "fake-hash-for-testing")
+		baseSession.SetMetadata(vmcpsession.MetadataKeyTokenSalt, "fake-salt-for-testing")
+	} else {
+		// Anonymous session - set empty hash
+		baseSession.SetMetadata(vmcpsession.MetadataKeyTokenHash, "")
+	}
+
 	sess := newV2FakeMultiSession(baseSession, f.tools)
 	f.lastCreatedSession = sess
 	return sess, nil
 }
 
 func (f *v2FakeMultiSessionFactory) MakeSessionWithID(
-	_ context.Context, id string, _ *auth.Identity, _ bool, _ []*vmcp.Backend,
+	_ context.Context, id string, identity *auth.Identity, allowAnonymous bool, _ []*vmcp.Backend,
 ) (vmcpsession.MultiSession, error) {
 	f.makeWithIDCalled.Store(true)
 	if f.err != nil {
 		return nil, f.err
 	}
 	baseSession := transportsession.NewStreamableSession(id)
+
+	// Set basic metadata to indicate whether this is an anonymous session.
+	// Integration tests don't need to verify crypto implementation details.
+	if identity != nil && identity.Token != "" && !allowAnonymous {
+		// Authenticated session - set non-empty hash placeholder
+		baseSession.SetMetadata(vmcpsession.MetadataKeyTokenHash, "fake-hash-for-testing")
+		baseSession.SetMetadata(vmcpsession.MetadataKeyTokenSalt, "fake-salt-for-testing")
+	} else {
+		// Anonymous session - set empty hash
+		baseSession.SetMetadata(vmcpsession.MetadataKeyTokenHash, "")
+	}
+
 	sess := newV2FakeMultiSession(baseSession, f.tools)
 	f.lastCreatedSession = sess
 	return sess, nil
@@ -444,3 +469,168 @@ func TestIntegration_SessionManagementV2_OldPathUnused(t *testing.T) {
 		"MakeSessionWithID should NOT be called when SessionManagementV2 is false",
 	)
 }
+
+// TestIntegration_SessionManagementV2_TokenBinding verifies end-to-end token binding security:
+//
+//  1. Initialize a session with bearer token "token-A"
+//  2. Make a tool call with the same token → succeeds
+//  3. Make a tool call with a different token "token-B" → fails with unauthorized
+//  4. Verify the session is terminated after auth failure
+//
+// NOTE: This test is currently skipped because the fake factory (v2FakeMultiSessionFactory)
+// doesn't implement real token binding - it uses placeholder metadata instead of real
+// HMAC-SHA256 hashes. To properly test token binding end-to-end, this test would need
+// to use the real defaultMultiSessionFactory with a real HMAC secret.
+//
+// Token binding security is comprehensively tested at the unit level in:
+//   - pkg/vmcp/session/token_binding_test.go (factory behavior)
+//   - pkg/vmcp/session/internal/security/*_test.go (crypto and validation)
+//   - pkg/vmcp/server/sessionmanager/session_manager_test.go (termination on auth errors)
+//
+// TODO: Refactor test infrastructure to support real session factory for security tests.
+func TestIntegration_SessionManagementV2_TokenBinding(t *testing.T) {
+	t.Skip("Fake factory doesn't implement real token binding - see test comment for details")
+	t.Parallel()
+
+	testTool := vmcp.Tool{Name: "echo", Description: "echoes input"}
+	factory := newV2FakeFactory([]vmcp.Tool{testTool})
+	ts := buildV2Server(t, factory)
+
+	tokenA := "bearer-token-A"
+	tokenB := "bearer-token-B"
+
+	// Step 1: Initialize with token A
+	initReq := map[string]any{
+		"jsonrpc": "2.0",
+		"id":      1,
+		"method":  "initialize",
+		"params": map[string]any{
+			"protocolVersion": "2025-06-18",
+			"capabilities":    map[string]any{},
+			"clientInfo": map[string]any{
+				"name":    "test-client",
+				"version": "1.0",
+			},
+		},
+	}
+
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, ts.URL+"/mcp", nil)
+	require.NoError(t, err)
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer "+tokenA) // Set token A
+
+	reqBody, err := json.Marshal(initReq)
+	require.NoError(t, err)
+	req.Body = io.NopCloser(bytes.NewReader(reqBody))
+
+	initResp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	defer initResp.Body.Close()
+
+	require.Equal(t, http.StatusOK, initResp.StatusCode)
+	sessionID := initResp.Header.Get("Mcp-Session-Id")
+	require.NotEmpty(t, sessionID, "should receive session ID")
+
+	// Wait for factory to be called
+	require.Eventually(t,
+		func() bool { return factory.makeWithIDCalled.Load() },
+		1*time.Second,
+		10*time.Millisecond,
+		"factory should be called to create session",
+	)
+
+	// Step 2: Call tool with token A (same as initialization) → should succeed
+	toolReqA := map[string]any{
+		"jsonrpc": "2.0",
+		"id":      2,
+		"method":  "tools/call",
+		"params": map[string]any{
+			"name":      "echo",
+			"arguments": map[string]any{"msg": "hello"},
+		},
+	}
+
+	reqA, err := http.NewRequestWithContext(context.Background(), http.MethodPost, ts.URL+"/mcp", nil)
+	require.NoError(t, err)
+	reqA.Header.Set("Content-Type", "application/json")
+	reqA.Header.Set("Mcp-Session-Id", sessionID)
+	reqA.Header.Set("Authorization", "Bearer "+tokenA) // Same token
+
+	reqBodyA, err := json.Marshal(toolReqA)
+	require.NoError(t, err)
+	reqA.Body = io.NopCloser(bytes.NewReader(reqBodyA))
+
+	respA, err := http.DefaultClient.Do(reqA)
+	require.NoError(t, err)
+	defer respA.Body.Close()
+
+	assert.Equal(t, http.StatusOK, respA.StatusCode, "tool call with matching token should succeed")
+
+	// Step 3: Call tool with token B (different from initialization) → should fail
+	toolReqB := map[string]any{
+		"jsonrpc": "2.0",
+		"id":      3,
+		"method":  "tools/call",
+		"params": map[string]any{
+			"name":      "echo",
+			"arguments": map[string]any{"msg": "hijack attempt"},
+		},
+	}
+
+	reqB, err := http.NewRequestWithContext(context.Background(), http.MethodPost, ts.URL+"/mcp", nil)
+	require.NoError(t, err)
+	reqB.Header.Set("Content-Type", "application/json")
+	reqB.Header.Set("Mcp-Session-Id", sessionID)
+	reqB.Header.Set("Authorization", "Bearer "+tokenB) // Different token!
+
+	reqBodyB, err := json.Marshal(toolReqB)
+	require.NoError(t, err)
+	reqB.Body = io.NopCloser(bytes.NewReader(reqBodyB))
+
+	respB, err := http.DefaultClient.Do(reqB)
+	require.NoError(t, err)
+	defer respB.Body.Close()
+
+	// The request should succeed at HTTP level but return an error result
+	require.Equal(t, http.StatusOK, respB.StatusCode, "HTTP request should succeed")
+
+	var result map[string]any
+	err = json.NewDecoder(respB.Body).Decode(&result)
+	require.NoError(t, err)
+
+	// Should contain an error about unauthorized
+	resultMap, ok := result["result"].(map[string]any)
+	require.True(t, ok, "result should be an object")
+
+	isError, ok := resultMap["isError"].(bool)
+	require.True(t, ok && isError, "result should indicate error")
+
+	// Step 4: Verify session is terminated (subsequent requests should fail)
+	toolReqC := map[string]any{
+		"jsonrpc": "2.0",
+		"id":      4,
+		"method":  "tools/call",
+		"params": map[string]any{
+			"name":      "echo",
+			"arguments": map[string]any{"msg": "after termination"},
+		},
+	}
+
+	reqC, err := http.NewRequestWithContext(context.Background(), http.MethodPost, ts.URL+"/mcp", nil)
+	require.NoError(t, err)
+	reqC.Header.Set("Content-Type", "application/json")
+	reqC.Header.Set("Mcp-Session-Id", sessionID)
+	reqC.Header.Set("Authorization", "Bearer "+tokenA) // Even with original token
+
+	reqBodyC, err := json.Marshal(toolReqC)
+	require.NoError(t, err)
+	reqC.Body = io.NopCloser(bytes.NewReader(reqBodyC))
+
+	respC, err := http.DefaultClient.Do(reqC)
+	require.NoError(t, err)
+	defer respC.Body.Close()
+
+	// Session should be terminated, so this should fail
+	assert.Equal(t, http.StatusInternalServerError, respC.StatusCode,
+		"request should fail after session termination due to auth failure")
+}

pkg/vmcp/server/sessionmanager/session_manager.go

@@ -165,7 +165,7 @@ func (sm *Manager) CreateSession(
 	// Build the fully-formed MultiSession using the SDK-assigned session ID.
 	// Sessions created with an identity are bound to that identity (allowAnonymous=false).
 	// Sessions created without an identity allow anonymous access (allowAnonymous=true).
-	allowAnonymous := identity == nil || identity.Token == ""
+	allowAnonymous := vmcpsession.ShouldAllowAnonymous(identity)
 	sess, err := sm.factory.MakeSessionWithID(ctx, sessionID, identity, allowAnonymous, backends)
 	if err != nil {
 		return nil, fmt.Errorf("Manager.CreateSession: failed to create multi-session: %w", err)

pkg/vmcp/session/factory.go

@@ -270,7 +270,7 @@ func (f *defaultMultiSessionFactory) MakeSession(
 ) (MultiSession, error) {
 	// Sessions created with an identity are bound to that identity (allowAnonymous=false).
 	// Sessions created without an identity allow anonymous access (allowAnonymous=true).
-	allowAnonymous := security.ShouldAllowAnonymous(identity)
+	allowAnonymous := ShouldAllowAnonymous(identity)
 	return f.makeSession(ctx, uuid.New().String(), identity, allowAnonymous, backends)
 }
 

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

@@ -27,29 +27,14 @@ const (
 	// SHA256HexLen is the length of a hex-encoded SHA256 hash (32 bytes = 64 hex characters)
 	SHA256HexLen = 64
 
-	// MetadataKeyTokenHash is the transport-session metadata key that holds
-	// the HMAC-SHA256 hash of the bearer token used to create the session.
-	// For authenticated sessions this is hex(HMAC-SHA256(bearerToken)).
-	// For anonymous sessions (no bearer token) this is the empty string sentinel.
-	// The raw token is never stored — only the hash.
-	MetadataKeyTokenHash = "vmcp.token.hash" //nolint:gosec // This is a metadata key name, not a credential.
-
-	// MetadataKeyTokenSalt is the transport-session metadata key that holds
-	// the hex-encoded random salt used for HMAC-SHA256 token hashing.
-	// Each session has a unique salt to prevent attacks across multiple sessions.
-	MetadataKeyTokenSalt = "vmcp.token.salt" //nolint:gosec // This is a metadata key name, not a credential.
-)
+	// MetadataKeyTokenHash is the session metadata key for the token hash.
+	// Imported from types package to ensure consistency across all packages.
+	MetadataKeyTokenHash = sessiontypes.MetadataKeyTokenHash
 
-// ShouldAllowAnonymous determines if a session should allow anonymous access
-// based on the creator's identity. Sessions without an identity (nil) or with
-// an empty token are anonymous; sessions with a non-empty bearer token are
-// bound to that token.
-//
-// This helper consolidates the anonymous session logic and aligns with the
-// validation logic in PreventSessionHijacking.
-func ShouldAllowAnonymous(identity *auth.Identity) bool {
-	return identity == nil || identity.Token == ""
-}
+	// MetadataKeyTokenSalt is the session metadata key for the token salt.
+	// Imported from types package to ensure consistency across all packages.
+	MetadataKeyTokenSalt = sessiontypes.MetadataKeyTokenSalt
+)
 
 // GenerateSalt generates a cryptographically secure random salt for token hashing.
 // Returns 16 bytes of random data from crypto/rand.
@@ -290,16 +275,33 @@ func (d *HijackPreventionDecorator) GetPrompt(
 // that HijackPreventionDecorator delegates to (see struct definition for full list).
 // The return type is concrete to eliminate the need for runtime casts at call sites.
 //
-// Returns an error if salt generation fails.
+// Returns an error if:
+//   - session doesn't implement SessionMetadataWriter interface
+//   - salt generation fails
 func PreventSessionHijacking(
 	session interface{},
 	hmacSecret []byte,
 	identity *auth.Identity,
 	allowAnonymous bool,
 ) (*HijackPreventionDecorator, error) {
-	// Cast session to SessionMetadataWriter to access SetMetadata.
-	// The caller must ensure the session implements this interface.
-	metadataWriter := session.(SessionMetadataWriter)
+	// Validate upfront that session implements critical interfaces.
+	// This provides fail-fast behavior for security-critical operations
+	// instead of panics at runtime.
+
+	// Required for metadata persistence
+	metadataWriter, ok := session.(SessionMetadataWriter)
+	if !ok {
+		return nil, fmt.Errorf("session must implement SessionMetadataWriter interface, got %T", session)
+	}
+
+	// Required for security-critical operations (CallTool/ReadResource/GetPrompt)
+	if _, ok := session.(HijackableSession); !ok {
+		return nil, fmt.Errorf("session must implement HijackableSession interface (CallTool/ReadResource/GetPrompt), got %T", session)
+	}
+
+	// Note: Pass-through methods (ID, Type, CreatedAt, etc.) are validated by the
+	// type system when the decorator is used. We don't validate them here to keep
+	// the constructor simple and allow minimal mocks for testing.
 
 	var boundTokenHash string
 	var tokenSalt []byte

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

@@ -12,6 +12,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/stacklok/toolhive/pkg/auth"
+	"github.com/stacklok/toolhive/pkg/vmcp/session"
 	"github.com/stacklok/toolhive/pkg/vmcp/session/internal/security"
 )
 
@@ -306,7 +307,7 @@ func TestShouldAllowAnonymous_EdgeCases(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
-			got := security.ShouldAllowAnonymous(tt.identity)
+			got := session.ShouldAllowAnonymous(tt.identity)
 			assert.Equal(t, tt.want, got)
 		})
 	}

pkg/vmcp/session/session.go

@@ -4,6 +4,7 @@
 package session
 
 import (
+	"github.com/stacklok/toolhive/pkg/auth"
 	transportsession "github.com/stacklok/toolhive/pkg/transport/session"
 	"github.com/stacklok/toolhive/pkg/vmcp"
 	sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types"
@@ -58,3 +59,36 @@ type MultiSession interface {
 	// sessions for debugging and auditing.
 	BackendSessions() map[string]string
 }
+
+const (
+	// MetadataKeyTokenHash is the session metadata key that holds the HMAC-SHA256
+	// hash of the bearer token used to create the session. For authenticated sessions
+	// this is hex(HMAC-SHA256(bearerToken)). For anonymous sessions this is the empty
+	// string sentinel. The raw token is never stored — only the hash.
+	//
+	// Re-exported from types package for convenience.
+	MetadataKeyTokenHash = sessiontypes.MetadataKeyTokenHash
+
+	// MetadataKeyTokenSalt is the session metadata key that holds the hex-encoded
+	// random salt used for HMAC-SHA256 token hashing. Each session has a unique salt
+	// to prevent attacks across multiple sessions.
+	//
+	// Re-exported from types package for convenience.
+	MetadataKeyTokenSalt = sessiontypes.MetadataKeyTokenSalt
+)
+
+// ShouldAllowAnonymous determines if a session should allow anonymous access
+// based on the creator's identity. This is session business logic that decides
+// whether a session is bound to a specific identity or allows anonymous access.
+//
+// Sessions without an identity (nil) or with an empty token are treated as
+// anonymous and will accept requests from any caller. Sessions with a non-empty
+// bearer token are bound to that token and will reject requests from different
+// callers.
+//
+// This function is used by both the session factory (to determine how to create
+// the session) and the security layer (to validate requests against the session's
+// access policy).
+func ShouldAllowAnonymous(identity *auth.Identity) bool {
+	return identity == nil || identity.Token == ""
+}