diff --git a/internal/cmd/delegation.go b/internal/cmd/delegation.go new file mode 100644 index 00000000..a39745af --- /dev/null +++ b/internal/cmd/delegation.go @@ -0,0 +1,62 @@ +package cmd + +import ( + "context" + "fmt" + "net" + "net/http" + + "github.com/github/gh-aw-mcpg/internal/delegation" + "github.com/github/gh-aw-mcpg/internal/logger" + "github.com/github/gh-aw-mcpg/internal/server" +) + +func startUnifiedDelegationControl( + ctx context.Context, + cancel context.CancelFunc, + unifiedServer *server.UnifiedServer, + delegationConfig *delegation.RuntimeConfig, +) (*http.Server, <-chan error, error) { + controlListenerErrCh := make(chan error, 1) + if delegationConfig == nil { + return nil, controlListenerErrCh, nil + } + controlListener, err := net.Listen("tcp", delegationConfig.ControlListenAddr) + if err != nil { + return nil, controlListenerErrCh, fmt.Errorf("failed to listen on private delegation control channel %s: %w", delegationConfig.ControlListenAddr, err) + } + controlHTTPServer := &http.Server{ + Handler: unifiedServer.ControlHandler(), + BaseContext: func(_ net.Listener) context.Context { + return ctx + }, + } + go func() { + if err := controlHTTPServer.Serve(controlListener); err != nil && err != http.ErrServerClosed { + logger.LogError("delegation", "Private delegation control channel exited unexpectedly, shutting down: %v", err) + controlListenerErrCh <- err + cancel() + } + }() + logger.LogInfo("startup", "Private delegation control channel listening on %s", controlListener.Addr()) + return controlHTTPServer, controlListenerErrCh, nil +} + +func persistUnifiedDelegationState(delegationConfig *delegation.RuntimeConfig, delegationStatePath string) error { + if delegationConfig == nil { + return nil + } + if err := delegationConfig.Store.SaveState(delegationStatePath); err != nil { + return fmt.Errorf("failed to persist delegation state: %w", err) + } + return nil +} + +func selectDelegationControlError(err error, controlListenerErrCh <-chan error) error { + select { + case controlErr := <-controlListenerErrCh: + return fmt.Errorf("private delegation control channel failed: %w", controlErr) + default: + return err + } +} diff --git a/internal/cmd/delegation_test.go b/internal/cmd/delegation_test.go new file mode 100644 index 00000000..4acb191b --- /dev/null +++ b/internal/cmd/delegation_test.go @@ -0,0 +1,111 @@ +package cmd + +import ( + "bytes" + "context" + "encoding/json" + "net" + "net/http" + "testing" + "time" + + "github.com/github/gh-aw-mcpg/internal/config" + "github.com/github/gh-aw-mcpg/internal/delegation" + "github.com/github/gh-aw-mcpg/internal/sanitize" + "github.com/github/gh-aw-mcpg/internal/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestStartUnifiedDelegationControlServesStatus(t *testing.T) { + previousRedaction := sanitize.PrivateSelectorRedactionEnabled() + t.Cleanup(func() { sanitize.SetPrivateSelectorRedaction(previousRedaction) }) + + const capabilityKey = "control-capability-key-32-bytes!!" + delegationConfig, err := testRuntimeDelegationConfig(t, availableLoopbackAddr(t), capabilityKey) + require.NoError(t, err) + cfg := &config.Config{ + Servers: map[string]*config.ServerConfig{}, + Delegation: delegationConfig, + } + us, err := server.NewUnified(context.Background(), cfg) + require.NoError(t, err) + defer us.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + controlServer, controlErrCh, err := startUnifiedDelegationControl(ctx, cancel, us, delegationConfig) + require.NoError(t, err) + require.NotNil(t, controlServer) + defer func() { _ = controlServer.Shutdown(context.Background()) }() + + body := bytes.NewBufferString(`{"run_id":"run-1","enclave_entry_id":"entry-1"}`) + req, err := http.NewRequest(http.MethodPost, "http://"+delegationConfig.ControlListenAddr+delegation.ControlPathPrefix+"status", body) + require.NoError(t, err) + req.Header.Set("Authorization", capabilityKey) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + var payload map[string]any + require.NoError(t, json.NewDecoder(resp.Body).Decode(&payload)) + assert.InEpsilon(t, 1.0, payload["generation"], 0) + assert.NoError(t, selectDelegationControlError(nil, controlErrCh)) +} + +func TestStartUnifiedDelegationControlFailsWhenListenAddrOccupied(t *testing.T) { + previousRedaction := sanitize.PrivateSelectorRedactionEnabled() + t.Cleanup(func() { sanitize.SetPrivateSelectorRedaction(previousRedaction) }) + + const capabilityKey = "control-capability-key-32-bytes!!" + occupiedAddr := availableLoopbackAddr(t) + occupyingListener, err := net.Listen("tcp", occupiedAddr) + require.NoError(t, err) + defer occupyingListener.Close() + + delegationConfig, err := testRuntimeDelegationConfig(t, occupiedAddr, capabilityKey) + require.NoError(t, err) + cfg := &config.Config{ + Servers: map[string]*config.ServerConfig{}, + Delegation: delegationConfig, + } + us, err := server.NewUnified(context.Background(), cfg) + require.NoError(t, err) + defer us.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + controlServer, _, err := startUnifiedDelegationControl(ctx, cancel, us, delegationConfig) + require.Error(t, err) + assert.Nil(t, controlServer) + assert.Contains(t, err.Error(), "failed to listen on private delegation control channel") +} + +func testRuntimeDelegationConfig(t *testing.T, listenAddr, capabilityKey string) (*delegation.RuntimeConfig, error) { + t.Helper() + envelope := &delegation.Envelope{ + RunID: "run-1", + EnclaveBackend: "awf-enclave", + AllowedRepositories: []string{"github/gh-aw"}, + ToolPolicy: delegation.ToolPolicyGitHubRepositoryReadV1, + AllowedSchemaHashes: []string{"sha256:test"}, + MaxIdentityTTL: 120 * time.Second, + ExpiresAt: time.Now().Add(time.Hour), + } + store, err := delegation.NewStore(envelope, 1) + if err != nil { + return nil, err + } + capability, err := delegation.NewControlCapability(capabilityKey) + if err != nil { + return nil, err + } + return &delegation.RuntimeConfig{ + Store: store, + Capability: capability, + StatePath: t.TempDir() + "/state.json", + ControlListenAddr: listenAddr, + }, nil +} diff --git a/internal/cmd/proxy.go b/internal/cmd/proxy.go index f836e002..11bef932 100644 --- a/internal/cmd/proxy.go +++ b/internal/cmd/proxy.go @@ -58,7 +58,7 @@ func init() { rootCmd.AddCommand(newProxyCmd()) } -func resolveDelegationProxyConfig() (*proxy.DelegationConfig, string, error) { +func resolveDelegationProxyConfig() (*delegation.RuntimeConfig, string, error) { envelopeJSON := os.Getenv("MCP_GATEWAY_DELEGATION_ENVELOPE") capabilityKey := os.Getenv(delegation.EnvControlCapabilityKey) statePath := os.Getenv("MCP_GATEWAY_DELEGATION_STATE_PATH") @@ -95,7 +95,7 @@ func resolveDelegationProxyConfig() (*proxy.DelegationConfig, string, error) { if err != nil { return nil, "", err } - return &proxy.DelegationConfig{Store: store, Capability: capability, StatePath: statePath, ControlListenAddr: controlListenAddr}, statePath, nil + return &delegation.RuntimeConfig{Store: store, Capability: capability, StatePath: statePath, ControlListenAddr: controlListenAddr}, statePath, nil } func newProxyCmd() *cobra.Command { diff --git a/internal/cmd/root.go b/internal/cmd/root.go index 8cfb58e5..66675034 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -1,6 +1,7 @@ package cmd import ( + "context" "fmt" "log" "os" @@ -284,6 +285,12 @@ func run(cmd *cobra.Command, args []string) error { debugLog.Printf("Server mode: %s, guards mode: %s", mode, cfg.DIFCMode) + delegationConfig, delegationStatePath, err := resolveDelegationProxyConfig() + if err != nil { + return err + } + cfg.Delegation = delegationConfig + // Apply tracing flags: CLI flags and env var overrides take precedence over config values. applyTracingOverrides(cmd, cfg) @@ -340,6 +347,16 @@ func run(cmd *cobra.Command, args []string) error { debugLog.Printf("Unified MCP server created successfully") defer unifiedServer.Close() + delegationControlServer, delegationControlErrCh, err := startUnifiedDelegationControl(ctx, cancel, unifiedServer, delegationConfig) + if err != nil { + return err + } + if delegationControlServer != nil { + defer func() { + _ = delegationControlServer.Shutdown(context.Background()) + }() + } + // Handle graceful shutdown via context cancellation go func() { <-ctx.Done() @@ -370,7 +387,7 @@ func run(cmd *cobra.Command, args []string) error { log.Printf("Warning: failed to write gateway configuration to stdout: %v", err) } - if err := serveAndWait( + err = serveAndWait( ctx, cancel, httpServer, @@ -381,10 +398,15 @@ func run(cmd *cobra.Command, args []string) error { func() error { return httpServer.Serve(listener) }, - ); err != nil { + ) + err = selectDelegationControlError(err, delegationControlErrCh) + if err != nil { debugLog.Printf("Server exited with error: %v", err) return err } + if err := persistUnifiedDelegationState(delegationConfig, delegationStatePath); err != nil { + return err + } return nil } diff --git a/internal/config/config_core.go b/internal/config/config_core.go index 096be162..595bad43 100644 --- a/internal/config/config_core.go +++ b/internal/config/config_core.go @@ -33,6 +33,7 @@ import ( "github.com/BurntSushi/toml" + "github.com/github/gh-aw-mcpg/internal/delegation" "github.com/github/gh-aw-mcpg/internal/logger" ) @@ -94,6 +95,11 @@ type Config struct { // GuardPolicySource describes where GuardPolicy was resolved from (cli|env|config|legacy). GuardPolicySource string `toml:"-" json:"-"` + + // Delegation optionally enables github-repository-delegation-v1 runtime + // control/data-plane authorization. It is resolved from environment + // activation inputs, never from user configuration. + Delegation *delegation.RuntimeConfig `toml:"-" json:"-"` } // GatewayConfig holds global gateway settings. diff --git a/internal/delegation/config.go b/internal/delegation/config.go new file mode 100644 index 00000000..1f85ae13 --- /dev/null +++ b/internal/delegation/config.go @@ -0,0 +1,14 @@ +package delegation + +// ControlPathPrefix is the private AWF control-plane URL prefix for +// github-repository-delegation-v1 operations. +const ControlPathPrefix = "/internal/awf-enclave-mcp-control/" + +// RuntimeConfig enables runtime repository-read delegation and its +// AWF-authenticated private control channel. +type RuntimeConfig struct { + Store *Store + Capability *ControlCapability + StatePath string + ControlListenAddr string +} diff --git a/internal/delegation/store.go b/internal/delegation/store.go index 43b51570..07d0ae07 100644 --- a/internal/delegation/store.go +++ b/internal/delegation/store.go @@ -373,6 +373,25 @@ func (s *Store) AuthorizeExecutor(executorBearer, repository, tool string) (stri return identity.Handle, nil } +// HasLiveExecutorBearer reports whether executorBearer currently identifies a +// live delegated executor identity. It intentionally does not authorize any +// repository or tool; callers must still use AuthorizeExecutor at the actual +// data-plane operation. +func (s *Store) HasLiveExecutorBearer(executorBearer string) bool { + s.mu.Lock() + defer s.mu.Unlock() + now := time.Now() + s.cleanupExpiredLocked(now) + if s.recoveryIncomplete || !now.Before(s.envelope.ExpiresAt) { + return false + } + if bearer, ok := strings.CutPrefix(executorBearer, "Bearer "); ok { + executorBearer = bearer + } + _, ok := s.byBearer[sha256.Sum256([]byte(executorBearer))] + return ok +} + func (s *Store) authorize(executorBearer, repository, tool string, bindingMatches func(*Identity) bool) (*Identity, error) { s.mu.Lock() defer s.mu.Unlock() diff --git a/internal/delegation/store_test.go b/internal/delegation/store_test.go index 26b312c7..7ab9b94a 100644 --- a/internal/delegation/store_test.go +++ b/internal/delegation/store_test.go @@ -247,6 +247,20 @@ func TestAuthorizeExecutor_ReturnsIdentityHandleForIsolation(t *testing.T) { assert.Error(t, err) } +func TestHasLiveExecutorBearerTracksRevocation(t *testing.T) { + store, _ := newTestStore(t) + req := validRequest() + created, err := store.CreateOrConfirm(req) + require.NoError(t, err) + + assert.True(t, store.HasLiveExecutorBearer(created.ExecutorBearer)) + assert.True(t, store.HasLiveExecutorBearer("Bearer "+created.ExecutorBearer)) + assert.False(t, store.HasLiveExecutorBearer(created.Handle), "control handles must not be accepted as executor bearers") + + require.NoError(t, store.Revoke(created.Handle)) + assert.False(t, store.HasLiveExecutorBearer(created.ExecutorBearer)) +} + func TestExpiry_AutomaticAndExplicit(t *testing.T) { store, _ := newTestStore(t) req := validRequest() diff --git a/internal/proxy/delegation.go b/internal/proxy/delegation.go index 13ed9b34..5d77694c 100644 --- a/internal/proxy/delegation.go +++ b/internal/proxy/delegation.go @@ -18,7 +18,7 @@ import ( var logDelegation = logger.ForFile() -const delegationControlPath = "/internal/awf-enclave-mcp-control/" +const delegationControlPath = delegation.ControlPathPrefix type delegationState struct { store *delegation.Store @@ -28,12 +28,7 @@ type delegationState struct { // DelegationConfig enables runtime repository-read delegation and its // AWF-authenticated private control channel. -type DelegationConfig struct { - Store *delegation.Store - Capability *delegation.ControlCapability - StatePath string - ControlListenAddr string -} +type DelegationConfig = delegation.RuntimeConfig func newDelegationState(cfg *DelegationConfig) (*delegationState, error) { if cfg == nil { diff --git a/internal/server/agent_policy_enforce.go b/internal/server/agent_policy_enforce.go index 30e9f4b5..ecc5316a 100644 --- a/internal/server/agent_policy_enforce.go +++ b/internal/server/agent_policy_enforce.go @@ -4,6 +4,7 @@ import ( "context" "strings" + "github.com/github/gh-aw-mcpg/internal/delegation" "github.com/github/gh-aw-mcpg/internal/logger" "github.com/github/gh-aw-mcpg/internal/util" sdk "github.com/modelcontextprotocol/go-sdk/mcp" @@ -102,3 +103,67 @@ func createAgentFilteredUnifiedServer(us *UnifiedServer, agentID string) *sdk.Se util.HashIdentifierForLog(agentID), registered) return server } + +func createDelegationFilteredUnifiedServer(us *UnifiedServer) *sdk.Server { + logUnified.Print("createDelegationFilteredUnifiedServer: building delegated github-repository-read-v1 tool view") + server := newSDKServer("awmg-unified-delegation", logTransport) + + us.toolsMu.RLock() + tools := make([]ToolInfo, 0, len(us.tools)) + for _, t := range us.tools { + tools = append(tools, *t) + } + us.toolsMu.RUnlock() + + registered := registerFilteredTools( + server, + tools, + "delegation", + func(toolInfo ToolInfo) (string, string) { + return toolInfo.BackendID, strings.TrimPrefix(toolInfo.Name, toolInfo.BackendID+"___") + }, + func(_ string, serverID, toolName string) bool { + return serverID == "github" && delegation.IsDelegatedTool(toolName) + }, + func(toolInfo ToolInfo) func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error) { + return toolInfo.Handler + }, + ) + + logger.LogInfo("client", "Built delegated unified tool view: tools=%d", registered) + return server +} + +func createDelegationFilteredServer(unifiedServer *UnifiedServer, backendID string) *sdk.Server { + logRouted.Printf("Creating delegated filtered server: backend=%s", backendID) + server := newSDKServer("awmg-"+backendID+"-delegation", logRouted) + if backendID != "github" { + return server + } + + tools := unifiedServer.GetToolsForBackend(backendID) + registerFilteredTools( + server, + tools, + "delegation", + func(toolInfo ToolInfo) (string, string) { + return backendID, toolInfo.Name + }, + func(_ string, _ string, toolName string) bool { + return delegation.IsDelegatedTool(toolName) + }, + func(toolInfo ToolInfo) func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error) { + handler := unifiedServer.GetToolHandler(backendID, toolInfo.Name) + if handler == nil { + logRouted.Printf("WARNING: No handler found for %s___%s", backendID, toolInfo.Name) + return nil + } + return func(ctx context.Context, req *sdk.CallToolRequest, _ interface{}) (*sdk.CallToolResult, interface{}, error) { + logRouted.Printf("[ROUTED] Calling delegated unified handler for: %s", toolInfo.Name) + return handler(ctx, req, nil) + } + }, + ) + + return server +} diff --git a/internal/server/agent_policy_visibility_test.go b/internal/server/agent_policy_visibility_test.go index b73ef504..234dfbfd 100644 --- a/internal/server/agent_policy_visibility_test.go +++ b/internal/server/agent_policy_visibility_test.go @@ -142,6 +142,14 @@ func TestCreateAgentFilteredUnifiedServer_PolicyIsolation(t *testing.T) { assert.Empty(t, ghostTools, "an agent without a policy sees no tools") } +func TestCreateDelegationFilteredUnifiedServer_ClosedToolSurface(t *testing.T) { + us := agentVisibilityServer(t) + + tools, err := listToolsViaInMemory(createDelegationFilteredUnifiedServer(us)) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"github___issue_read"}, tools, "delegated sessions see only the closed github-repository-read-v1 tools registered by github") +} + // TestCreateAgentFilteredUnifiedServer_ConcurrentIsolation exercises the per-agent // filtered-server construction concurrently to catch data races (run with -race). func TestCreateAgentFilteredUnifiedServer_ConcurrentIsolation(t *testing.T) { diff --git a/internal/server/backend_call.go b/internal/server/backend_call.go index 2f13c9d6..e78786a0 100644 --- a/internal/server/backend_call.go +++ b/internal/server/backend_call.go @@ -209,6 +209,12 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName }() sessionID := us.getSessionID(ctx) + ctx, err := us.authorizeDelegatedToolCall(ctx, serverID, toolName, args) + if err != nil { + httpStatusCode = 403 + tracing.RecordSpanError(toolSpan, err, "delegation denied") + return mcp.NewErrorCallToolResult(err) + } // Propagate a redacted, stable session attribution to the tool call span so it // is queryable on child spans without exposing the raw authenticated identity. if toolSpan.IsRecording() { @@ -219,7 +225,7 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName // authenticated agent is not permitted to use. This is defense-in-depth // alongside the per-agent tool-visibility filtering applied at session // establishment (createAgentFilteredServer / createAgentFilteredUnifiedServer). - if us.agentPoliciesEnforced() { + if us.agentPoliciesEnforced() && !delegatedToolAuthorized(ctx, serverID, toolName) { agentIdentity := guard.GetAgentIDFromContext(ctx) if !us.agentCanUseTool(agentIdentity, serverID, toolName) { logger.LogWarn("client", "tools/call denied by per-agent policy: agent=%s tool=%q server=%s", diff --git a/internal/server/delegation.go b/internal/server/delegation.go new file mode 100644 index 00000000..b8d888e3 --- /dev/null +++ b/internal/server/delegation.go @@ -0,0 +1,294 @@ +package server + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + + "github.com/github/gh-aw-mcpg/internal/delegation" + "github.com/github/gh-aw-mcpg/internal/guard" + "github.com/github/gh-aw-mcpg/internal/httputil" + "github.com/github/gh-aw-mcpg/internal/logger" + "github.com/github/gh-aw-mcpg/internal/util" +) + +var logServerDelegation = logger.ForFile() + +type delegatedToolAuthorizationKey struct{} + +type delegatedToolAuthorization struct { + serverID string + toolName string +} + +func (us *UnifiedServer) delegationEnabled() bool { + return us != nil && us.delegation != nil +} + +func (us *UnifiedServer) isDelegatedExecutorAuth(authorizationHeader string) bool { + if !us.delegationEnabled() { + return false + } + return us.delegation.Store.HasLiveExecutorBearer(authorizationHeader) +} + +func (us *UnifiedServer) isDelegatedExecutorSession(sessionID string) bool { + if !us.delegationEnabled() { + return false + } + return us.delegation.Store.HasLiveExecutorBearer(sessionID) +} + +func (us *UnifiedServer) authorizeDelegatedToolCall(ctx context.Context, serverID, toolName string, args interface{}) (context.Context, error) { + if !us.delegationEnabled() { + return ctx, nil + } + sessionID := us.getSessionID(ctx) + if !us.delegation.Store.HasLiveExecutorBearer(sessionID) { + return ctx, nil + } + if serverID != "github" { + return ctx, fmt.Errorf("delegated identity is not authorized for server %q", serverID) + } + repository, ok := delegatedToolRepository(toolName, args) + if !ok { + return ctx, fmt.Errorf("delegated identity requires canonical owner/repo arguments for tool %q", toolName) + } + handle, err := us.delegation.Store.AuthorizeExecutor(sessionID, repository, toolName) + if err != nil { + logServerDelegation.Printf("Delegated tool call denied: tool=%s repo_hash=%s", toolName, util.HashForLog(repository, 16, "")) + return ctx, err + } + ctx = guard.SetAgentIDInContext(ctx, "delegation:"+handle) + ctx = context.WithValue(ctx, delegatedToolAuthorizationKey{}, delegatedToolAuthorization{ + serverID: serverID, + toolName: toolName, + }) + return ctx, nil +} + +func delegatedToolRepository(toolName string, args interface{}) (string, bool) { + if !delegation.IsDelegatedTool(toolName) { + return "", false + } + argsMap, ok := args.(map[string]interface{}) + if !ok { + return "", false + } + owner, ownerOK := argsMap["owner"].(string) + repo, repoOK := argsMap["repo"].(string) + if !ownerOK || !repoOK { + return "", false + } + repository := owner + "/" + repo + return repository, delegation.IsCanonicalRepositorySelector(repository) +} + +func delegatedToolAuthorized(ctx context.Context, serverID, toolName string) bool { + authorization, ok := ctx.Value(delegatedToolAuthorizationKey{}).(delegatedToolAuthorization) + return ok && authorization.serverID == serverID && authorization.toolName == toolName +} + +// delegatedAllowedMethods are the only JSON-RPC methods a delegated executor +// bearer may invoke on the /mcp data plane. Everything else, including +// prompts/resources capabilities and unrecognized methods, is denied. +var delegatedAllowedMethods = map[string]bool{ + "initialize": true, + "notifications/initialized": true, + "ping": true, + "tools/list": true, + "tools/call": true, +} + +func (us *UnifiedServer) rejectDelegatedNonToolMethods(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !us.isDelegatedExecutorAuth(r.Header.Get("Authorization")) || r.Method != http.MethodPost { + next.ServeHTTP(w, r) + return + } + body, err := readAndRestoreRequestBody(r) + if err != nil { + logServerDelegation.Printf("Delegated MCP request denied: failed to read body: %v", err) + httputil.WriteErrorResponse(w, http.StatusForbidden, "delegation_method_denied", "delegated identity is not authorized for this MCP request") + return + } + methods, ok := parseDelegatedRequestMethods(body) + if !ok || len(methods) == 0 { + logServerDelegation.Printf("Delegated MCP request denied: unparsable or empty request envelope") + httputil.WriteErrorResponse(w, http.StatusForbidden, "delegation_method_denied", "delegated identity is not authorized for this MCP request") + return + } + for _, method := range methods { + if !delegatedAllowedMethods[method] { + logServerDelegation.Printf("Delegated MCP method denied: method=%s", method) + httputil.WriteErrorResponse(w, http.StatusForbidden, "delegation_method_denied", "delegated identity is not authorized for this MCP method") + return + } + } + next.ServeHTTP(w, r) + }) +} + +// parseDelegatedRequestMethods extracts every JSON-RPC "method" value from a +// request body, which may be either a single request object or a batch +// array of request objects (accepted by the SDK for protocol versions +// before 2025-06-18). It fails closed: any parse failure, or any batch +// element missing a non-empty method, returns ok=false so the caller denies +// the request instead of passing it through unauthorized. +func parseDelegatedRequestMethods(body []byte) ([]string, bool) { + trimmed := bytes.TrimLeft(body, " \t\r\n") + if len(trimmed) == 0 { + return nil, false + } + var envelopes []json.RawMessage + if trimmed[0] == '[' { + if err := json.Unmarshal(trimmed, &envelopes); err != nil { + return nil, false + } + if len(envelopes) == 0 { + return nil, false + } + } else { + envelopes = []json.RawMessage{json.RawMessage(trimmed)} + } + methods := make([]string, 0, len(envelopes)) + for _, envelope := range envelopes { + var request struct { + Method string `json:"method"` + } + if err := json.Unmarshal(envelope, &request); err != nil || request.Method == "" { + return nil, false + } + methods = append(methods, request.Method) + } + return methods, true +} + +// ControlHandler returns the private delegation control-plane handler. It is +// intentionally separate from the MCP data plane so executor bearers cannot +// reach control operations through /mcp. +func (us *UnifiedServer) ControlHandler() http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !us.delegationEnabled() || !strings.HasPrefix(r.URL.Path, delegation.ControlPathPrefix) { + http.NotFound(w, r) + return + } + us.handleDelegationControl(w, r) + }) +} + +func (us *UnifiedServer) handleDelegationControl(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost || us.delegation.Capability.Authenticate(r.Header.Get("Authorization")) != nil { + logServerDelegation.Printf("Delegation control access denied: method=%s path=%s", r.Method, r.URL.Path) + httputil.WriteErrorResponse(w, http.StatusForbidden, "delegation_access_denied", "delegation control access denied") + return + } + + switch r.URL.Path { + case delegation.ControlPathPrefix + "create-or-confirm": + var requestWire delegation.CreateOrConfirmRequestWire + if !decodeServerDelegationJSON(w, r, &requestWire) { + return + } + request, err := requestWire.ToRequest() + if err != nil { + httputil.WriteErrorResponse(w, http.StatusBadRequest, "invalid_delegation_request", "invalid delegation request") + return + } + result, err := us.delegation.Store.CreateOrConfirm(request) + if err != nil { + if !us.persistDelegationState(w) { + return + } + httputil.WriteErrorResponse(w, http.StatusForbidden, "delegation_request_denied", "delegation request denied") + return + } + if !us.persistDelegationState(w) { + return + } + httputil.WriteJSONResponse(w, http.StatusOK, result) + case delegation.ControlPathPrefix + "revoke": + var request struct { + Handle string `json:"handle"` + } + if !decodeServerDelegationJSON(w, r, &request) { + return + } + if err := us.delegation.Store.Revoke(request.Handle); err != nil { + httputil.WriteErrorResponse(w, http.StatusInternalServerError, "delegation_revoke_failed", "delegation revoke failed") + return + } + if !us.persistDelegationState(w) { + return + } + httputil.WriteJSONResponse(w, http.StatusOK, map[string]bool{"revoked": true}) + case delegation.ControlPathPrefix + "revoke-by-labels": + var request struct { + RunID string `json:"run_id"` + EnclaveEntryID string `json:"enclave_entry_id"` + } + if !decodeServerDelegationJSON(w, r, &request) { + return + } + revoked := us.delegation.Store.RevokeByLabels(request.RunID, request.EnclaveEntryID) + if !us.persistDelegationState(w) { + return + } + httputil.WriteJSONResponse(w, http.StatusOK, map[string]int{"revoked": revoked}) + case delegation.ControlPathPrefix + "status": + var request struct { + RunID string `json:"run_id"` + EnclaveEntryID string `json:"enclave_entry_id"` + } + if !decodeServerDelegationJSON(w, r, &request) { + return + } + if request.RunID == "" || request.EnclaveEntryID == "" { + httputil.WriteErrorResponse(w, http.StatusBadRequest, "delegation_status_invalid_request", "run_id and enclave_entry_id are required") + return + } + status := us.delegation.Store.Status() + httputil.WriteJSONResponse(w, http.StatusOK, map[string]any{ + "recovery_incomplete": status.RecoveryIncomplete, + "generation": status.Generation, + "live_identity_count": status.LiveIdentityCount, + "labelled_handles": us.delegation.Store.LabelHandles(request.RunID, request.EnclaveEntryID), + }) + case delegation.ControlPathPrefix + "reconcile": + var request struct{} + if !decodeServerDelegationJSON(w, r, &request) { + return + } + if err := us.delegation.Store.MarkReconciledAndSaveState(us.delegation.StatePath); err != nil { + httputil.WriteErrorResponse(w, http.StatusInternalServerError, "delegation_state_persist_failed", "delegation state persistence failed") + return + } + httputil.WriteJSONResponse(w, http.StatusOK, map[string]bool{"reconciled": true}) + default: + http.NotFound(w, r) + } +} + +func (us *UnifiedServer) persistDelegationState(w http.ResponseWriter) bool { + if err := us.delegation.Store.SaveState(us.delegation.StatePath); err != nil { + httputil.WriteErrorResponse(w, http.StatusInternalServerError, "delegation_state_persist_failed", "delegation state persistence failed") + return false + } + return true +} + +func decodeServerDelegationJSON(w http.ResponseWriter, r *http.Request, value any) bool { + r.Body = http.MaxBytesReader(w, r.Body, 64*1024) + defer r.Body.Close() + decoder := json.NewDecoder(r.Body) + decoder.DisallowUnknownFields() + if decoder.Decode(value) != nil || decoder.Decode(&struct{}{}) != io.EOF { + httputil.WriteErrorResponse(w, http.StatusBadRequest, "invalid_delegation_request", "invalid delegation request") + return false + } + return true +} diff --git a/internal/server/delegation_test.go b/internal/server/delegation_test.go new file mode 100644 index 00000000..52153d4b --- /dev/null +++ b/internal/server/delegation_test.go @@ -0,0 +1,271 @@ +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/github/gh-aw-mcpg/internal/config" + "github.com/github/gh-aw-mcpg/internal/delegation" + "github.com/github/gh-aw-mcpg/internal/guard" + "github.com/github/gh-aw-mcpg/internal/sanitize" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newUnifiedDelegationConfig(t *testing.T) (*delegation.RuntimeConfig, delegation.CreateOrConfirmRequest) { + t.Helper() + envelope := &delegation.Envelope{ + RunID: "run-1", + EnclaveBackend: "awf-enclave", + AllowedRepositories: []string{"github/gh-aw"}, + ToolPolicy: delegation.ToolPolicyGitHubRepositoryReadV1, + AllowedSchemaHashes: []string{"sha256:test"}, + MaxIdentityTTL: 120 * time.Second, + ExpiresAt: time.Now().Add(time.Hour), + } + store, err := delegation.NewStore(envelope, 1) + require.NoError(t, err) + capability, err := delegation.NewControlCapability("control-capability-key-32-bytes!!") + require.NoError(t, err) + return &delegation.RuntimeConfig{ + Store: store, + Capability: capability, + StatePath: t.TempDir() + "/state.json", + ControlListenAddr: "127.0.0.1:0", + }, delegation.CreateOrConfirmRequest{ + RunID: "run-1", + EnclaveBackend: "awf-enclave", + EnclaveEntryID: "entry-1", + InvocationID: "inv-1", + Repository: "github/gh-aw", + ToolPolicy: delegation.ToolPolicyGitHubRepositoryReadV1, + SchemaHash: "sha256:test", + RequestedTTL: time.Minute, + IdempotencyKey: "key-1", + } +} + +func TestDelegatedAuthAdmitsOnlyLiveExecutorBearer(t *testing.T) { + delegationConfig, createReq := newUnifiedDelegationConfig(t) + created, err := delegationConfig.Store.CreateOrConfirm(createReq) + require.NoError(t, err) + + us := &UnifiedServer{delegation: delegationConfig} + called := false + handler := applyAuthIfConfiguredWithDelegation([]string{"gateway-key"}, us.isDelegatedExecutorAuth, func(w http.ResponseWriter, _ *http.Request) { + called = true + w.WriteHeader(http.StatusNoContent) + }) + + req, err := http.NewRequest(http.MethodPost, "/mcp", nil) + require.NoError(t, err) + req.Header.Set("Authorization", created.ExecutorBearer) + rec := httptest.NewRecorder() + handler(rec, req) + assert.Equal(t, http.StatusNoContent, rec.Code) + assert.True(t, called) + + require.NoError(t, delegationConfig.Store.Revoke(created.Handle)) + called = false + rec = httptest.NewRecorder() + handler(rec, req) + assert.Equal(t, http.StatusUnauthorized, rec.Code) + assert.False(t, called) +} + +func TestDelegatedAuthWithoutGatewayKeyRequiresLiveBearer(t *testing.T) { + delegationConfig, createReq := newUnifiedDelegationConfig(t) + created, err := delegationConfig.Store.CreateOrConfirm(createReq) + require.NoError(t, err) + + us := &UnifiedServer{delegation: delegationConfig} + called := false + handler := applyAuthIfConfiguredWithDelegation(nil, us.isDelegatedExecutorAuth, func(w http.ResponseWriter, _ *http.Request) { + called = true + w.WriteHeader(http.StatusNoContent) + }) + + req, err := http.NewRequest(http.MethodPost, "/mcp", nil) + require.NoError(t, err) + req.Header.Set("Authorization", created.ExecutorBearer) + req.Header.Set("X-Agent-ID", "attacker-agent") + rec := httptest.NewRecorder() + handler(rec, req) + assert.Equal(t, http.StatusNoContent, rec.Code) + assert.True(t, called) + + require.NoError(t, delegationConfig.Store.Revoke(created.Handle)) + called = false + rec = httptest.NewRecorder() + handler(rec, req) + assert.Equal(t, http.StatusUnauthorized, rec.Code) + assert.False(t, called) +} + +func TestRejectDelegatedNonToolMethodsDeniesUnauthorizedMethod(t *testing.T) { + delegationConfig, createReq := newUnifiedDelegationConfig(t) + created, err := delegationConfig.Store.CreateOrConfirm(createReq) + require.NoError(t, err) + + us := &UnifiedServer{delegation: delegationConfig} + called := false + handler := us.rejectDelegatedNonToolMethods(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + called = true + w.WriteHeader(http.StatusNoContent) + })) + + body := `{"jsonrpc":"2.0","id":1,"method":"prompts/get"}` + req, err := http.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Authorization", created.ExecutorBearer) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + assert.Equal(t, http.StatusForbidden, rec.Code) + assert.False(t, called) +} + +func TestRejectDelegatedNonToolMethodsDeniesBatchedUnauthorizedMethod(t *testing.T) { + delegationConfig, createReq := newUnifiedDelegationConfig(t) + created, err := delegationConfig.Store.CreateOrConfirm(createReq) + require.NoError(t, err) + + us := &UnifiedServer{delegation: delegationConfig} + + testCases := []struct { + name string + body string + }{ + { + name: "batched prompts/list", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/call"},{"jsonrpc":"2.0","id":2,"method":"prompts/list"}]`, + }, + { + name: "batched prompts/get", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","id":2,"method":"prompts/get"}]`, + }, + { + name: "malformed envelope", + body: `{"jsonrpc":"2.0","id":1,`, + }, + { + name: "batch element missing method", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","id":2}]`, + }, + { + name: "empty batch", + body: `[]`, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + called := false + handler := us.rejectDelegatedNonToolMethods(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + called = true + w.WriteHeader(http.StatusNoContent) + })) + req, err := http.NewRequest(http.MethodPost, "/mcp", strings.NewReader(tc.body)) + require.NoError(t, err) + req.Header.Set("Authorization", created.ExecutorBearer) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + assert.Equal(t, http.StatusForbidden, rec.Code) + assert.False(t, called) + }) + } +} + +func TestRejectDelegatedNonToolMethodsAllowsBatchedToolMethods(t *testing.T) { + delegationConfig, createReq := newUnifiedDelegationConfig(t) + created, err := delegationConfig.Store.CreateOrConfirm(createReq) + require.NoError(t, err) + + us := &UnifiedServer{delegation: delegationConfig} + called := false + handler := us.rejectDelegatedNonToolMethods(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + called = true + w.WriteHeader(http.StatusNoContent) + })) + + body := `[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","id":2,"method":"tools/call"}]` + req, err := http.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Authorization", created.ExecutorBearer) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + assert.Equal(t, http.StatusNoContent, rec.Code) + assert.True(t, called) +} + +func TestUnifiedDelegationAuthorizesExactGitHubRepositoryTool(t *testing.T) { + previousRedaction := sanitize.PrivateSelectorRedactionEnabled() + t.Cleanup(func() { sanitize.SetPrivateSelectorRedaction(previousRedaction) }) + + delegationConfig, createReq := newUnifiedDelegationConfig(t) + created, err := delegationConfig.Store.CreateOrConfirm(createReq) + require.NoError(t, err) + + backend := newBackendWithToolResponse(t, "list_issues", defaultToolResponse) + defer backend.Close() + + guardName := "unified-delegation-guard" + guard.RegisterGuardType(guardName, func() (guard.Guard, error) { return &difcTestGuard{}, nil }) + cfg := &config.Config{ + DIFCMode: "filter", + Gateway: &config.GatewayConfig{ + AgentID: "enclave-agent", + AgentPolicies: map[string]*config.AgentPolicy{ + "enclave-agent": {Servers: []string{"github"}, Tools: map[string][]string{"github": {"list_issues"}}}, + }, + }, + Servers: map[string]*config.ServerConfig{ + "github": { + Type: "http", + URL: backend.URL, + Guard: guardName, + GuardPolicies: map[string]interface{}{ + "allow-only": map[string]interface{}{"repos": "public", "min-integrity": "none"}, + }, + }, + }, + Guards: map[string]*config.GuardConfig{ + guardName: {Type: guardName}, + }, + GuardPolicy: &config.GuardPolicy{ + AllowOnly: &config.AllowOnlyPolicy{Repos: "public", MinIntegrity: config.IntegrityNone}, + }, + GuardPolicySource: "cli", + Delegation: delegationConfig, + } + us, err := NewUnified(context.Background(), cfg) + require.NoError(t, err) + defer us.Close() + + result, _, err := us.callBackendTool(callCtx(created.ExecutorBearer), "github", "list_issues", map[string]interface{}{ + "owner": "github", + "repo": "gh-aw", + }) + require.NoError(t, err) + require.NotNil(t, result) + assert.False(t, result.IsError) + + result, _, err = us.callBackendTool(callCtx(created.ExecutorBearer), "github", "list_issues", map[string]interface{}{ + "owner": "github", + "repo": "gh-aw-firewall", + }) + require.Error(t, err) + require.NotNil(t, result) + assert.True(t, result.IsError) + + result, _, err = us.callBackendTool(callCtx(created.ExecutorBearer), "github", "search_issues", map[string]interface{}{ + "owner": "github", + "repo": "gh-aw", + }) + require.Error(t, err) + require.NotNil(t, result) + assert.True(t, result.IsError) +} diff --git a/internal/server/http_server.go b/internal/server/http_server.go index 6b34103b..db16ee52 100644 --- a/internal/server/http_server.go +++ b/internal/server/http_server.go @@ -63,12 +63,13 @@ func buildMCPHTTPServer( // signature (ASI-07); common endpoints (e.g. /health, /close) are not HMAC-protected. func CreateHTTPServerForMCP(addr string, unifiedServer *UnifiedServer, apiKeys []string, hmacSecret string) *http.Server { logTransport.Printf("Creating HTTP server for MCP: addr=%s, auth_enabled=%v, hmac_enabled=%v", addr, len(apiKeys) > 0, hmacSecret != "") - authEnabled := len(apiKeys) > 0 + authEnabled := len(apiKeys) > 0 || unifiedServer.delegationEnabled() return buildMCPHTTPServer(addr, unifiedServer, apiKeys, hmacSecret, func(mux *http.ServeMux, sessionTimeout time.Duration) { logTransport.Print("Registering streamable HTTP handler for MCP protocol") // Per-agent unified servers expose only the tools an agent may see. Built // lazily and cached per identity; only used when per-agent policies are set. agentServerCache := syncutil.NewTTLCache[string, *sdk.Server](sessionTimeout, filteredServerCacheMaxSize) + delegationServerCache := syncutil.NewTTLCache[string, *sdk.Server](sessionTimeout, filteredServerCacheMaxSize) // Create the standard MCP handler stack (StreamableHTTP + session auto-init + middleware). // This is what Codex uses with transport = "streamablehttp" finalHandler := buildMCPHandler(func(r *http.Request) *sdk.Server { @@ -83,6 +84,12 @@ func CreateHTTPServerForMCP(addr string, unifiedServer *UnifiedServer, apiKeys [ return nil } + if unifiedServer.isDelegatedExecutorSession(sessionID) { + return delegationServerCache.GetOrCreate(sessionID, func() *sdk.Server { + return createDelegationFilteredUnifiedServer(unifiedServer) + }) + } + // When per-agent policies are configured, expose only the tools this // authenticated agent may use (tool visibility enforcement). if unifiedServer.agentPoliciesEnforced() { @@ -117,11 +124,12 @@ func CreateHTTPServerForRoutedMode(addr string, unifiedServer *UnifiedServer, ap allBackends := unifiedServer.GetServerIDs() logRouted.Printf("Registering routes for %d backends: %v", len(allBackends), allBackends) - authEnabled := len(apiKeys) > 0 + authEnabled := len(apiKeys) > 0 || unifiedServer.delegationEnabled() return buildMCPHTTPServer(addr, unifiedServer, apiKeys, hmacSecret, func(mux *http.ServeMux, sessionTimeout time.Duration) { logRouted.Printf("[CACHE] Creating filtered server cache: ttl=%s, maxSize=%d", sessionTimeout, filteredServerCacheMaxSize) serverCache := syncutil.NewTTLCache[string, *sdk.Server](sessionTimeout, filteredServerCacheMaxSize) + delegationServerCache := syncutil.NewTTLCache[string, *sdk.Server](sessionTimeout, filteredServerCacheMaxSize) for _, serverID := range allBackends { backendID := serverID @@ -143,6 +151,12 @@ func CreateHTTPServerForRoutedMode(addr string, unifiedServer *UnifiedServer, ap return nil } + if unifiedServer.isDelegatedExecutorSession(sessionID) { + return delegationServerCache.GetOrCreate(backendID+"|"+sessionID, func() *sdk.Server { + return createDelegationFilteredServer(unifiedServer, backendID) + }) + } + // Per-agent server-access enforcement at session establishment: // reject routed sessions for backends this agent's policy does not permit. if !unifiedServer.agentCanAccessServer(sessionID, backendID) { diff --git a/internal/server/middleware.go b/internal/server/middleware.go index 304b183b..02698039 100644 --- a/internal/server/middleware.go +++ b/internal/server/middleware.go @@ -85,9 +85,11 @@ func wrapWithMiddleware(handler http.Handler, logTag string, unifiedServer *Unif // Wrap SDK handler with detailed logging for JSON-RPC translation debugging loggedHandler := WithSDKLogging(handler, logTag) + delegationMethodHandler := unifiedServer.rejectDelegatedNonToolMethods(loggedHandler) + // Apply shutdown check middleware (spec 5.1.3) // This must come before auth to ensure shutdown takes precedence - shutdownHandler := rejectIfShutdown(unifiedServer, loggedHandler, "server:"+logTag) + shutdownHandler := rejectIfShutdown(unifiedServer, delegationMethodHandler, "server:"+logTag) // Apply HMAC signature verification if secret is configured (ASI-07). // HMAC wraps the shutdown handler so only post-auth requests pay the body-read cost. @@ -96,7 +98,11 @@ func wrapWithMiddleware(handler http.Handler, logTag string, unifiedServer *Unif // Apply auth middleware if API key is configured (spec 7.1). // Auth is the outermost application-level check so unauthenticated requests are // rejected before HMAC validation (and its body-read overhead) runs. - authedHandler := applyAuthIfConfigured(apiKeys, hmacHandler) + var delegatedAuthenticator func(string) bool + if unifiedServer.delegationEnabled() { + delegatedAuthenticator = unifiedServer.isDelegatedExecutorAuth + } + authedHandler := applyAuthIfConfiguredWithDelegation(apiKeys, delegatedAuthenticator, hmacHandler) // Wrap with OTEL tracing span (outermost, so it covers auth + HMAC + shutdown + logging) tracingHandler := WithOTELTracing(authedHandler, logTag) diff --git a/internal/server/middleware_auth.go b/internal/server/middleware_auth.go index 4f8c3d69..de5d0262 100644 --- a/internal/server/middleware_auth.go +++ b/internal/server/middleware_auth.go @@ -60,6 +60,10 @@ func matchesAnyKey(authHeader string, keys []string) bool { // allows multiple concurrent identities, e.g. primary/enclave, to each // authenticate with their own identifier). func authMiddleware(apiKeys []string, next http.HandlerFunc) http.HandlerFunc { + return authMiddlewareWithDelegation(apiKeys, nil, next) +} + +func authMiddlewareWithDelegation(apiKeys []string, delegatedAuthenticator func(string) bool, next http.HandlerFunc) http.HandlerFunc { logAuth.Printf("Initialized auth middleware") return func(w http.ResponseWriter, r *http.Request) { logAuth.Printf("Authenticating request: method=%s, path=%s, remote=%s", r.Method, r.URL.Path, r.RemoteAddr) @@ -83,7 +87,7 @@ func authMiddleware(apiKeys []string, next http.HandlerFunc) http.HandlerFunc { } // Spec 7.1: Authorization header must contain one of the configured API keys directly. - if !matchesAnyKey(authHeader, apiKeys) { + if !matchesAnyKey(authHeader, apiKeys) && (delegatedAuthenticator == nil || !delegatedAuthenticator(authHeader)) { logAuth.Printf("Rejecting auth request: status=%d, code=%s, detail=%s, path=%s, remote=%s", http.StatusUnauthorized, "unauthorized", "invalid_api_key", r.URL.Path, r.RemoteAddr) rejectRequest(w, r, http.StatusUnauthorized, "unauthorized", "invalid API key", "auth", "authentication_failed", "invalid_api_key") return @@ -98,9 +102,13 @@ func authMiddleware(apiKeys []string, next http.HandlerFunc) http.HandlerFunc { // applyAuthIfConfigured applies authentication middleware if at least one API key is provided. // Returns the handler unchanged if apiKeys is empty. func applyAuthIfConfigured(apiKeys []string, handler http.HandlerFunc) http.HandlerFunc { - if len(apiKeys) > 0 { + return applyAuthIfConfiguredWithDelegation(apiKeys, nil, handler) +} + +func applyAuthIfConfiguredWithDelegation(apiKeys []string, delegatedAuthenticator func(string) bool, handler http.HandlerFunc) http.HandlerFunc { + if len(apiKeys) > 0 || delegatedAuthenticator != nil { logAuth.Print("Auth key configured, applying middleware") - return authMiddleware(apiKeys, handler) + return authMiddlewareWithDelegation(apiKeys, delegatedAuthenticator, handler) } logAuth.Print("No auth key configured, skipping middleware") return handler diff --git a/internal/server/unified.go b/internal/server/unified.go index 7da1b1e2..5b062f86 100644 --- a/internal/server/unified.go +++ b/internal/server/unified.go @@ -8,12 +8,14 @@ import ( "time" "github.com/github/gh-aw-mcpg/internal/config" + "github.com/github/gh-aw-mcpg/internal/delegation" "github.com/github/gh-aw-mcpg/internal/difc" "github.com/github/gh-aw-mcpg/internal/githubhttp" "github.com/github/gh-aw-mcpg/internal/guard" "github.com/github/gh-aw-mcpg/internal/launcher" "github.com/github/gh-aw-mcpg/internal/logger" "github.com/github/gh-aw-mcpg/internal/mcp" + "github.com/github/gh-aw-mcpg/internal/sanitize" "github.com/github/gh-aw-mcpg/internal/tracing" "github.com/github/gh-aw-mcpg/internal/util" sdk "github.com/modelcontextprotocol/go-sdk/mcp" @@ -112,6 +114,8 @@ type UnifiedServer struct { // Configuration reference for guard loading cfg *config.Config + delegation *delegation.RuntimeConfig + // Shutdown state tracking isShutdown bool shutdownMu sync.RWMutex @@ -144,6 +148,12 @@ type UnifiedServer struct { // NewUnified creates a new unified MCP server func NewUnified(ctx context.Context, cfg *config.Config) (*UnifiedServer, error) { logUnified.Printf("Creating new unified server: sequentialLaunch=%v, servers=%d", cfg.SequentialLaunch, len(cfg.Servers)) + if cfg.Delegation != nil { + if cfg.Delegation.Store == nil || cfg.Delegation.Capability == nil || cfg.Delegation.StatePath == "" { + return nil, fmt.Errorf("delegation store, control capability, and state path are required") + } + sanitize.EnablePrivateSelectorRedaction() + } l := launcher.New(ctx, cfg) @@ -181,6 +191,7 @@ func NewUnified(ctx context.Context, cfg *config.Config) (*UnifiedServer, error) guardRegistry: guard.NewRegistry(), DIFCComponents: difcComponents, cfg: cfg, // Store config for guard loading + delegation: cfg.Delegation, // Cache tracer at construction to avoid calling otel.Tracer on every request. CachedTracer: tracing.CachedTracer{Tracer: tracing.Tracer()}, diff --git a/test/integration/delegation_unified_test.go b/test/integration/delegation_unified_test.go new file mode 100644 index 00000000..601abfd0 --- /dev/null +++ b/test/integration/delegation_unified_test.go @@ -0,0 +1,502 @@ +package integration + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strconv" + "testing" + "time" + + "github.com/github/gh-aw-mcpg/internal/delegation" + sdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestUnifiedDelegationRootCommandWithAgentPolicies(t *testing.T) { + if testing.Short() { + t.Skip("Skipping binary integration test in short mode") + } + + binaryPath := findBinary(t) + backend := createDelegationMockMCPBackend(t) + defer backend.Close() + + gatewayAddr := freeLoopbackAddr(t) + controlAddr := freeLoopbackAddr(t) + statePath := filepath.Join(t.TempDir(), "delegation-state.json") + capabilityKey := "control-capability-key-32-bytes!!" + envelopeJSON := delegationEnvelopeJSON(t) + + cmd, stdout, stderr := startDelegationGateway(t, binaryPath, gatewayAddr, controlAddr, statePath, capabilityKey, envelopeJSON, backend.URL+"/mcp") + defer stopCommand(t, cmd) + + serverURL := "http://" + gatewayAddr + if !waitForServer(t, serverURL+"/health", 10*time.Second) { + t.Logf("STDOUT: %s", stdout.String()) + t.Logf("STDERR: %s", stderr.String()) + t.Fatal("Server did not start in time") + } + + status := postDelegationControl(t, controlAddr, capabilityKey, "status", map[string]any{ + "run_id": "run-1", + "enclave_entry_id": "entry-1", + }) + assert.Empty(t, status["labelled_handles"]) + + reconcile := postDelegationControl(t, controlAddr, capabilityKey, "reconcile", map[string]any{}) + reconciled, ok := reconcile["reconciled"].(bool) + require.True(t, ok) + assert.True(t, reconciled) + + created := postDelegationControl(t, controlAddr, capabilityKey, "create-or-confirm", map[string]any{ + "run_id": "run-1", + "enclave_backend": "awf-enclave", + "enclave_entry_id": "entry-1", + "invocation_id": "inv-1", + "repository": "github/gh-aw", + "tool_policy": delegation.ToolPolicyGitHubRepositoryReadV1, + "schema_hash": "sha256:test", + "requested_ttl": 30, + "idempotency_key": "key-1", + }) + bearer, ok := created["executor_bearer"].(string) + require.True(t, ok) + require.NotEmpty(t, bearer) + handle, ok := created["handle"].(string) + require.True(t, ok) + require.NotEmpty(t, handle) + + // Negative credential cross-use: an executor bearer must not be usable + // as a control-plane capability key, and the control capability key + // must not be admitted as a data-plane executor bearer. + crossUseReq, err := http.NewRequest(http.MethodPost, "http://"+controlAddr+delegation.ControlPathPrefix+"status", bytes.NewReader([]byte(`{}`))) + require.NoError(t, err) + crossUseReq.Header.Set("Authorization", bearer) + crossUseReq.Header.Set("Content-Type", "application/json") + crossUseResp, err := http.DefaultClient.Do(crossUseReq) + require.NoError(t, err) + crossUseResp.Body.Close() + assert.Equal(t, http.StatusForbidden, crossUseResp.StatusCode, "executor bearer must not authenticate to the control listener") + + capabilityAsBearer := delegatedMCPRequest(t, serverURL+"/mcp", capabilityKey, "", map[string]any{ + "jsonrpc": "2.0", + "id": 100, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2024-11-05", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "cross-use-test", "version": "1.0.0"}, + }, + }) + assert.Contains(t, capabilityAsBearer, "error", "control capability key must not be admitted as a data-plane executor bearer") + + mcpSessionID := initializeDelegatedMCP(t, serverURL+"/mcp", bearer) + tools := delegatedMCPRequest(t, serverURL+"/mcp", bearer, mcpSessionID, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/list", + }) + result, ok := tools["result"].(map[string]any) + require.True(t, ok, "tools/list response: %#v", tools) + assert.ElementsMatch(t, []string{"github___issue_read", "github___list_issues"}, toolNamesFromResult(t, result)) + + allowed := delegatedMCPRequest(t, serverURL+"/mcp", bearer, mcpSessionID, toolCallPayload(2, "github___list_issues", map[string]any{ + "owner": "github", + "repo": "gh-aw", + })) + assertDelegatedToolSucceeded(t, allowed) + + wrongRepo := delegatedMCPRequest(t, serverURL+"/mcp", bearer, mcpSessionID, toolCallPayload(3, "github___list_issues", map[string]any{ + "owner": "github", + "repo": "sibling", + })) + assertDelegatedToolDenied(t, wrongRepo) + + disallowedTool := delegatedMCPRequest(t, serverURL+"/mcp", bearer, mcpSessionID, toolCallPayload(4, "github___search_issues", map[string]any{ + "owner": "github", + "repo": "gh-aw", + })) + assertDelegatedToolDenied(t, disallowedTool) + + expiring := postDelegationControl(t, controlAddr, capabilityKey, "create-or-confirm", map[string]any{ + "run_id": "run-1", + "enclave_backend": "awf-enclave", + "enclave_entry_id": "entry-expiring", + "invocation_id": "inv-expiring", + "repository": "github/gh-aw", + "tool_policy": delegation.ToolPolicyGitHubRepositoryReadV1, + "schema_hash": "sha256:test", + "requested_ttl": 1, + "idempotency_key": "key-expiring", + }) + expiringBearer := expiring["executor_bearer"].(string) + expiringSessionID := initializeDelegatedMCP(t, serverURL+"/mcp", expiringBearer) + time.Sleep(1100 * time.Millisecond) + expiredReplay := delegatedMCPRequest(t, serverURL+"/mcp", expiringBearer, expiringSessionID, toolCallPayload(5, "github___list_issues", map[string]any{ + "owner": "github", + "repo": "gh-aw", + })) + assertDelegatedToolDenied(t, expiredReplay) + + prompts := delegatedMCPRequest(t, serverURL+"/mcp", bearer, mcpSessionID, map[string]any{ + "jsonrpc": "2.0", + "id": 6, + "method": "prompts/list", + }) + assert.Contains(t, prompts, "error", "delegated sessions must not expose non-tool MCP capabilities") + + postDelegationControl(t, controlAddr, capabilityKey, "revoke", map[string]any{"handle": handle}) + replayed := delegatedMCPRequest(t, serverURL+"/mcp", bearer, mcpSessionID, toolCallPayload(7, "github___list_issues", map[string]any{ + "owner": "github", + "repo": "gh-aw", + })) + assertDelegatedToolDenied(t, replayed) + + createdForLabels := postDelegationControl(t, controlAddr, capabilityKey, "create-or-confirm", map[string]any{ + "run_id": "run-1", + "enclave_backend": "awf-enclave", + "enclave_entry_id": "entry-1", + "invocation_id": "inv-2", + "repository": "github/gh-aw", + "tool_policy": delegation.ToolPolicyGitHubRepositoryReadV1, + "schema_hash": "sha256:test", + "requested_ttl": 30, + "idempotency_key": "key-2", + }) + labelBearer := createdForLabels["executor_bearer"].(string) + labelSessionID := initializeDelegatedMCP(t, serverURL+"/mcp", labelBearer) + revokedByLabels := postDelegationControl(t, controlAddr, capabilityKey, "revoke-by-labels", map[string]any{ + "run_id": "run-1", + "enclave_entry_id": "entry-1", + }) + revoked, ok := revokedByLabels["revoked"].(float64) + require.True(t, ok) + assert.Equal(t, 1, int(revoked)) + labelReplay := delegatedMCPRequest(t, serverURL+"/mcp", labelBearer, labelSessionID, toolCallPayload(8, "github___issue_read", map[string]any{ + "owner": "github", + "repo": "gh-aw", + })) + assertDelegatedToolDenied(t, labelReplay) + + // A live (non-revoked, non-expired) identity must survive a gateway + // restart and continue to authorize calls against the persisted store. + persistent := postDelegationControl(t, controlAddr, capabilityKey, "create-or-confirm", map[string]any{ + "run_id": "run-1", + "enclave_backend": "awf-enclave", + "enclave_entry_id": "entry-persist", + "invocation_id": "inv-persist", + "repository": "github/gh-aw", + "tool_policy": delegation.ToolPolicyGitHubRepositoryReadV1, + "schema_hash": "sha256:test", + "requested_ttl": 50, + "idempotency_key": "key-persist", + }) + persistentBearer := persistent["executor_bearer"].(string) + + stopCommand(t, cmd) + cmd, stdout, stderr = startDelegationGateway(t, binaryPath, gatewayAddr, controlAddr, statePath, capabilityKey, envelopeJSON, backend.URL+"/mcp") + defer stopCommand(t, cmd) + if !waitForServer(t, serverURL+"/health", 10*time.Second) { + t.Logf("STDOUT: %s", stdout.String()) + t.Logf("STDERR: %s", stderr.String()) + t.Fatal("Restarted server did not start in time") + } + recoveredStatus := postDelegationControl(t, controlAddr, capabilityKey, "status", map[string]any{ + "run_id": "run-1", + "enclave_entry_id": "entry-1", + }) + assert.Empty(t, recoveredStatus["labelled_handles"], "revoked delegations must remain revoked after restart") + + persistentStatus := postDelegationControl(t, controlAddr, capabilityKey, "status", map[string]any{ + "run_id": "run-1", + "enclave_entry_id": "entry-persist", + }) + assert.NotEmpty(t, persistentStatus["labelled_handles"], "a live identity must remain reconciled after restart") + + persistentSessionID := initializeDelegatedMCP(t, serverURL+"/mcp", persistentBearer) + persistentCall := delegatedMCPRequest(t, serverURL+"/mcp", persistentBearer, persistentSessionID, toolCallPayload(9, "github___list_issues", map[string]any{ + "owner": "github", + "repo": "gh-aw", + })) + assertDelegatedToolSucceeded(t, persistentCall) +} + +func TestUnifiedDelegationRootCommandRejectsPartialActivation(t *testing.T) { + if testing.Short() { + t.Skip("Skipping binary integration test in short mode") + } + + binaryPath := findBinary(t) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, binaryPath, + "--config-stdin", + "--listen", freeLoopbackAddr(t), + "--unified", + ) + cmd.Env = append(os.Environ(), + "MCP_GATEWAY_DELEGATION_ENVELOPE="+delegationEnvelopeJSON(t), + delegation.EnvControlCapabilityKey+"=control-capability-key-32-bytes!!", + delegation.EnvControlListenAddr+"="+freeLoopbackAddr(t), + "MCP_GATEWAY_DELEGATION_GENERATION=1", + ) + configJSON := map[string]any{ + "mcpServers": map[string]any{ + "github": map[string]any{ + "type": "http", + "url": "http://127.0.0.1:1/mcp", + }, + }, + "gateway": map[string]any{ + "port": portFromAddr(t, freeLoopbackAddr(t)), + "domain": "localhost", + "agentId": "primary-agent", + }, + } + configBytes, err := json.Marshal(configJSON) + require.NoError(t, err) + cmd.Stdin = bytes.NewReader(configBytes) + output, err := cmd.CombinedOutput() + require.Error(t, err) + assert.Contains(t, string(output), "must be configured together") +} + +func createDelegationMockMCPBackend(t *testing.T) *httptest.Server { + t.Helper() + impl := &sdk.Implementation{Name: "delegation-mock-backend", Version: "1.0.0"} + mcpServer := sdk.NewServer(impl, nil) + for _, toolName := range []string{"issue_read", "list_issues", "search_issues"} { + name := toolName + mcpServer.AddTool(&sdk.Tool{ + Name: name, + Description: "mock " + name, + InputSchema: map[string]any{"type": "object"}, + }, func(_ context.Context, _ *sdk.CallToolRequest) (*sdk.CallToolResult, error) { + return &sdk.CallToolResult{ + Content: []sdk.Content{&sdk.TextContent{Text: name + " response"}}, + }, nil + }) + } + handler := sdk.NewStreamableHTTPHandler(func(_ *http.Request) *sdk.Server { + return mcpServer + }, &sdk.StreamableHTTPOptions{Stateless: false}) + mux := http.NewServeMux() + mux.Handle("/mcp", handler) + mux.Handle("/mcp/", handler) + return httptest.NewServer(mux) +} + +func delegationEnvelopeJSON(t *testing.T) string { + t.Helper() + envelope := delegation.EnvelopeWire{ + RunID: "run-1", + EnclaveBackend: "awf-enclave", + AllowedRepositories: []string{"github/gh-aw"}, + ToolPolicy: delegation.ToolPolicyGitHubRepositoryReadV1, + AllowedSchemaHashes: []string{"sha256:test"}, + MaxIdentityTTLSeconds: 60, + ExpiresAt: time.Now().Add(time.Minute), + } + raw, err := json.Marshal(envelope) + require.NoError(t, err) + return string(raw) +} + +func startDelegationGateway(t *testing.T, binaryPath, gatewayAddr, controlAddr, statePath, capabilityKey, envelopeJSON, backendURL string) (*exec.Cmd, *bytes.Buffer, *bytes.Buffer) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) + t.Cleanup(cancel) + cmd := exec.CommandContext(ctx, binaryPath, + "--config-stdin", + "--listen", gatewayAddr, + "--unified", + ) + cmd.Env = append(os.Environ(), + "MCP_GATEWAY_DELEGATION_ENVELOPE="+envelopeJSON, + delegation.EnvControlCapabilityKey+"="+capabilityKey, + "MCP_GATEWAY_DELEGATION_STATE_PATH="+statePath, + "MCP_GATEWAY_DELEGATION_GENERATION=1", + delegation.EnvControlListenAddr+"="+controlAddr, + ) + configJSON := map[string]any{ + "mcpServers": map[string]any{ + "github": map[string]any{ + "type": "http", + "url": backendURL, + }, + }, + "gateway": map[string]any{ + "port": portFromAddr(t, gatewayAddr), + "domain": "localhost", + "agentIds": []string{"primary-agent", "enclave-agent"}, + "agentPolicies": map[string]any{ + "primary-agent": map[string]any{"servers": []string{"github"}}, + "enclave-agent": map[string]any{ + "servers": []string{"github"}, + "tools": map[string]any{"github": []string{"list_issues", "issue_read"}}, + }, + }, + }, + } + configBytes, err := json.Marshal(configJSON) + require.NoError(t, err) + stdout := &bytes.Buffer{} + stderr := &bytes.Buffer{} + cmd.Stdin = bytes.NewReader(configBytes) + cmd.Stdout = stdout + cmd.Stderr = stderr + require.NoError(t, cmd.Start()) + return cmd, stdout, stderr +} + +func portFromAddr(t *testing.T, addr string) int { + t.Helper() + _, port, err := net.SplitHostPort(addr) + require.NoError(t, err) + parsed, err := strconv.Atoi(port) + require.NoError(t, err) + return parsed +} + +func stopCommand(t *testing.T, cmd *exec.Cmd) { + t.Helper() + if cmd == nil || cmd.Process == nil { + return + } + _ = cmd.Process.Kill() + _, _ = cmd.Process.Wait() +} + +func freeLoopbackAddr(t *testing.T) string { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer listener.Close() + return listener.Addr().String() +} + +func postDelegationControl(t *testing.T, controlAddr, capabilityKey, operation string, payload map[string]any) map[string]any { + t.Helper() + raw, err := json.Marshal(payload) + require.NoError(t, err) + req, err := http.NewRequest(http.MethodPost, "http://"+controlAddr+delegation.ControlPathPrefix+operation, bytes.NewReader(raw)) + require.NoError(t, err) + req.Header.Set("Authorization", capabilityKey) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, string(body)) + var result map[string]any + require.NoError(t, json.Unmarshal(body, &result), string(body)) + return result +} + +func initializeDelegatedMCP(t *testing.T, url, bearer string) string { + t.Helper() + _, sessionID := delegatedMCPRequestWithSession(t, url, bearer, "", map[string]any{ + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2024-11-05", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "delegation-integration-test", "version": "1.0.0"}, + }, + }) + require.NotEmpty(t, sessionID) + return sessionID +} + +func delegatedMCPRequest(t *testing.T, url, bearer, mcpSessionID string, payload map[string]any) map[string]any { + t.Helper() + result, _ := delegatedMCPRequestWithSession(t, url, bearer, mcpSessionID, payload) + return result +} + +func delegatedMCPRequestWithSession(t *testing.T, url, bearer, mcpSessionID string, payload map[string]any) (map[string]any, string) { + t.Helper() + raw, err := json.Marshal(payload) + require.NoError(t, err) + req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(raw)) + require.NoError(t, err) + req.Header.Set("Authorization", bearer) + req.Header.Set("X-Agent-ID", "spoofed-agent") + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + if mcpSessionID != "" { + req.Header.Set("Mcp-Session-Id", mcpSessionID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + nextSessionID := resp.Header.Get("Mcp-Session-Id") + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + if resp.StatusCode != http.StatusOK { + return map[string]any{"error": map[string]any{"code": resp.StatusCode, "message": string(body)}}, nextSessionID + } + contentType := resp.Header.Get("Content-Type") + if contentType == "text/event-stream" { + return parseSSEResponse(t, string(body)), nextSessionID + } + var result map[string]any + require.NoError(t, json.Unmarshal(body, &result), string(body)) + return result, nextSessionID +} + +func toolCallPayload(id int, name string, arguments map[string]any) map[string]any { + return map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": "tools/call", + "params": map[string]any{ + "name": name, + "arguments": arguments, + }, + } +} + +func toolNamesFromResult(t *testing.T, result map[string]any) []string { + t.Helper() + tools, ok := result["tools"].([]any) + require.True(t, ok) + names := make([]string, 0, len(tools)) + for _, tool := range tools { + toolMap, ok := tool.(map[string]any) + require.True(t, ok) + name, ok := toolMap["name"].(string) + require.True(t, ok) + names = append(names, name) + } + return names +} + +func assertDelegatedToolSucceeded(t *testing.T, response map[string]any) { + t.Helper() + require.NotContains(t, response, "error") + result, ok := response["result"].(map[string]any) + require.True(t, ok) + assert.NotEqual(t, true, result["isError"]) +} + +func assertDelegatedToolDenied(t *testing.T, response map[string]any) { + t.Helper() + if _, ok := response["error"]; ok { + return + } + result, ok := response["result"].(map[string]any) + require.True(t, ok) + assert.Equal(t, true, result["isError"]) +}