diff --git a/cmd/chat_tools.go b/cmd/chat_tools.go index dece8001..357e2ebe 100644 --- a/cmd/chat_tools.go +++ b/cmd/chat_tools.go @@ -16,14 +16,27 @@ import ( const startupMCPToolLoadTimeout = 1500 * time.Millisecond -var defaultRegistryLoadMCPTools = tool.LoadMCPTools +var ( + defaultRegistryLoadMCPTools = tool.LoadMCPTools + defaultRegistryLoadRemoteMCPTools = tool.LoadRemoteMCPTools +) type startupMCPServerSpec struct { name string command string args []string + + // Set only for non-stdio (remote) servers: serverType is "http", "sse", + // or "websocket", url is the server's endpoint, and headers carries any + // static headers from config plus an auto-injected OAuth bearer token + // if one is stored for this server (see internal/mcp/oauth.go). + serverType string + url string + headers map[string]string } +func (s startupMCPServerSpec) isRemote() bool { return s.serverType != "" } + func essentialTools() []tool.Tool { // Core tools needed for basic agent operation - always loaded at startup return []tool.Tool{ @@ -105,14 +118,32 @@ func optionalTools() []tool.Tool { func configuredStartupMCPServers(settings hawkconfig.Settings) []startupMCPServerSpec { servers := make([]startupMCPServerSpec, 0, len(settings.MCPServers)+len(mcpServers)) for _, cfg := range settings.MCPServers { - if cfg.Name == "" || cfg.Command == "" { + if cfg.Name == "" { + continue + } + switch cfg.Type { + case "", "stdio": + if cfg.Command == "" { + continue + } + servers = append(servers, startupMCPServerSpec{ + name: cfg.Name, + command: cfg.Command, + args: cfg.Args, + }) + case "http", "sse", "websocket": + if cfg.URL == "" { + continue + } + servers = append(servers, startupMCPServerSpec{ + name: cfg.Name, + serverType: cfg.Type, + url: cfg.URL, + headers: mergedMCPHeaders(cfg), + }) + default: continue } - servers = append(servers, startupMCPServerSpec{ - name: cfg.Name, - command: cfg.Command, - args: cfg.Args, - }) } for _, cmd := range mcpServers { parts := strings.Fields(cmd) @@ -128,6 +159,24 @@ func configuredStartupMCPServers(settings hawkconfig.Settings) []startupMCPServe return servers } +// mergedMCPHeaders combines a remote MCP server's static configured headers +// with an auto-injected "Authorization: Bearer " if a valid, +// non-expired OAuth token is stored for this server (see +// internal/tool/mcp_auth.go). A configured static Authorization header, if +// any, takes precedence over the auto-injected one. +func mergedMCPHeaders(cfg hawkconfig.MCPServerConfig) map[string]string { + headers := make(map[string]string, len(cfg.Headers)+1) + for k, v := range cfg.Headers { + headers[k] = v + } + if _, hasAuth := headers["Authorization"]; !hasAuth { + if bearer, ok := tool.AuthHeaderForMCPServer(cfg.Name); ok { + headers["Authorization"] = bearer + } + } + return headers +} + func loadStartupMCPToolSets(servers []startupMCPServerSpec) [][]tool.Tool { results := make([][]tool.Tool, len(servers)) var wg sync.WaitGroup @@ -138,7 +187,15 @@ func loadStartupMCPToolSets(servers []startupMCPServerSpec) [][]tool.Tool { ctx, cancel := context.WithTimeout(context.Background(), startupMCPToolLoadTimeout) defer cancel() - mcpTools, err := defaultRegistryLoadMCPTools(ctx, spec.name, spec.command, spec.args...) + var ( + mcpTools []tool.Tool + err error + ) + if spec.isRemote() { + mcpTools, err = defaultRegistryLoadRemoteMCPTools(ctx, spec.name, spec.serverType, spec.url, spec.headers) + } else { + mcpTools, err = defaultRegistryLoadMCPTools(ctx, spec.name, spec.command, spec.args...) + } if err != nil { return } diff --git a/cmd/chat_tools_test.go b/cmd/chat_tools_test.go index 40acaf96..43791ccc 100644 --- a/cmd/chat_tools_test.go +++ b/cmd/chat_tools_test.go @@ -71,6 +71,124 @@ func TestLoadStartupMCPToolSets_UsesTimeoutAndPreservesOrder(t *testing.T) { } } +func TestConfiguredStartupMCPServers_DispatchesByType(t *testing.T) { + settings := hawkconfig.Settings{ + MCPServers: []hawkconfig.MCPServerConfig{ + {Name: "stdio-default", Command: "stdio-mcp"}, + {Name: "stdio-explicit", Type: "stdio", Command: "stdio-mcp-2"}, + {Name: "http-server", Type: "http", URL: "https://example.com/mcp"}, + {Name: "sse-server", Type: "sse", URL: "https://example.com/sse", Headers: map[string]string{"X-Custom": "1"}}, + {Name: "ws-server", Type: "websocket", URL: "wss://example.com/ws"}, + // Invalid/incomplete entries must be skipped, not error. + {Name: "", Command: "no-name"}, + {Name: "no-command"}, // stdio with no Command + {Name: "http-no-url", Type: "http"}, // http with no URL + {Name: "unknown-type", Type: "carrier-pigeon", URL: "x"}, + }, + } + + specs := configuredStartupMCPServers(settings) + + byName := make(map[string]startupMCPServerSpec, len(specs)) + for _, s := range specs { + byName[s.name] = s + } + if len(specs) != 5 { + t.Fatalf("expected 5 valid specs, got %d: %+v", len(specs), specs) + } + + stdioDefault, ok := byName["stdio-default"] + if !ok || stdioDefault.isRemote() || stdioDefault.command != "stdio-mcp" { + t.Fatalf("stdio-default spec wrong: %+v (ok=%v)", stdioDefault, ok) + } + + httpServer, ok := byName["http-server"] + if !ok || !httpServer.isRemote() || httpServer.serverType != "http" || httpServer.url != "https://example.com/mcp" { + t.Fatalf("http-server spec wrong: %+v (ok=%v)", httpServer, ok) + } + + sseServer, ok := byName["sse-server"] + if !ok || sseServer.serverType != "sse" || sseServer.headers["X-Custom"] != "1" { + t.Fatalf("sse-server spec wrong: %+v (ok=%v)", sseServer, ok) + } + + wsServer, ok := byName["ws-server"] + if !ok || wsServer.serverType != "websocket" { + t.Fatalf("ws-server spec wrong: %+v (ok=%v)", wsServer, ok) + } + + for _, missing := range []string{"", "no-command", "http-no-url", "unknown-type"} { + if _, found := byName[missing]; found { + t.Fatalf("expected %q to be skipped, but it was included", missing) + } + } +} + +func TestLoadStartupMCPToolSets_DispatchesRemoteSpecsToRemoteLoader(t *testing.T) { + origStdio := defaultRegistryLoadMCPTools + origRemote := defaultRegistryLoadRemoteMCPTools + t.Cleanup(func() { + defaultRegistryLoadMCPTools = origStdio + defaultRegistryLoadRemoteMCPTools = origRemote + }) + + var stdioCalled, remoteCalled bool + var gotServerType, gotURL string + var gotHeaders map[string]string + + defaultRegistryLoadMCPTools = func(ctx context.Context, name, command string, args ...string) ([]tool.Tool, error) { + stdioCalled = true + return []tool.Tool{registryTestTool{name: name}}, nil + } + defaultRegistryLoadRemoteMCPTools = func(ctx context.Context, name, serverType, url string, headers map[string]string) ([]tool.Tool, error) { + remoteCalled = true + gotServerType = serverType + gotURL = url + gotHeaders = headers + return []tool.Tool{registryTestTool{name: name}}, nil + } + + sets := loadStartupMCPToolSets([]startupMCPServerSpec{ + {name: "stdio-one", command: "cmd"}, + { + name: "remote-one", + serverType: "http", + url: "https://example.com/mcp", + headers: map[string]string{"Authorization": "Bearer xyz"}, + }, + }) + + if !stdioCalled { + t.Error("expected the stdio loader to be called for the stdio spec") + } + if !remoteCalled { + t.Fatal("expected the remote loader to be called for the remote spec") + } + if gotServerType != "http" || gotURL != "https://example.com/mcp" { + t.Errorf("remote loader got serverType=%q url=%q", gotServerType, gotURL) + } + if gotHeaders["Authorization"] != "Bearer xyz" { + t.Errorf("remote loader did not receive configured headers: %+v", gotHeaders) + } + if len(sets) != 2 || len(sets[0]) != 1 || len(sets[1]) != 1 { + t.Fatalf("expected both specs to produce one tool each, got %+v", sets) + } +} + +func TestMergedMCPHeaders_ConfiguredAuthorizationTakesPrecedence(t *testing.T) { + cfg := hawkconfig.MCPServerConfig{ + Name: "svc", + Headers: map[string]string{"Authorization": "Bearer static-token", "X-Other": "1"}, + } + headers := mergedMCPHeaders(cfg) + if headers["Authorization"] != "Bearer static-token" { + t.Errorf("expected static Authorization header to win, got %q", headers["Authorization"]) + } + if headers["X-Other"] != "1" { + t.Errorf("expected other static headers to be preserved, got %+v", headers) + } +} + func TestDefaultRegistry_SkipsFailedStartupMCPServers(t *testing.T) { orig := defaultRegistryLoadMCPTools t.Cleanup(func() { defaultRegistryLoadMCPTools = orig }) diff --git a/external/eyrie b/external/eyrie index 2e5ec4e3..1ba69293 160000 --- a/external/eyrie +++ b/external/eyrie @@ -1 +1 @@ -Subproject commit 2e5ec4e3bb03705d5a09792009f113625258fc5a +Subproject commit 1ba692935ce5f74a664d4dfc3775a552853b99fb diff --git a/external/trace b/external/trace index 1c17ce77..164d86be 160000 --- a/external/trace +++ b/external/trace @@ -1 +1 @@ -Subproject commit 1c17ce779e5d7a562ab6bef468ce8da2855ba82d +Subproject commit 164d86be64a8386938ac2f63a87ec0b47fd8acf0 diff --git a/external/yaad b/external/yaad index 2050a49f..7d52bde3 160000 --- a/external/yaad +++ b/external/yaad @@ -1 +1 @@ -Subproject commit 2050a49fa045d4266b91472fd59523b197fe20bf +Subproject commit 7d52bde38e498c60270fc3a2413ec3db1c17890f diff --git a/go.mod b/go.mod index 629a190b..9f2dc968 100644 --- a/go.mod +++ b/go.mod @@ -11,12 +11,12 @@ require ( charm.land/bubbles/v2 v2.1.0 charm.land/bubbletea/v2 v2.0.7 charm.land/lipgloss/v2 v2.0.3 - github.com/GrayCodeAI/eyrie v0.2.1 + github.com/GrayCodeAI/eyrie v0.2.2-0.20260715190227-1ba692935ce5 github.com/GrayCodeAI/hawk-core-contracts v0.1.4 github.com/GrayCodeAI/inspect v0.1.4 github.com/GrayCodeAI/sight v0.1.4 github.com/GrayCodeAI/tok v0.1.4 - github.com/GrayCodeAI/yaad v0.1.4 + github.com/GrayCodeAI/yaad v0.2.1-0.20260715205802-7d52bde38e49 github.com/bwmarrin/discordgo v0.28.1 github.com/charmbracelet/x/ansi v0.11.7 github.com/fsnotify/fsnotify v1.10.1 @@ -128,7 +128,7 @@ require ( require ( github.com/BurntSushi/toml v1.6.0 // indirect - github.com/GrayCodeAI/trace v0.1.4 + github.com/GrayCodeAI/trace v0.1.5-0.20260715120456-164d86be64a8 github.com/atotto/clipboard v0.1.4 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/cenkalti/backoff/v5 v5.0.3 // indirect @@ -169,7 +169,7 @@ require ( golang.org/x/tools v0.45.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect - google.golang.org/grpc v1.81.1 // indirect + google.golang.org/grpc v1.82.0 // indirect google.golang.org/protobuf v1.36.11 // indirect modernc.org/libc v1.72.5 // indirect modernc.org/mathutil v1.7.1 // indirect diff --git a/go.sum b/go.sum index 9130a13b..a97eea54 100644 --- a/go.sum +++ b/go.sum @@ -16,8 +16,8 @@ github.com/BobuSumisu/aho-corasick v1.0.3 h1:uuf+JHwU9CHP2Vx+wAy6jcksJThhJS9ehR8 github.com/BobuSumisu/aho-corasick v1.0.3/go.mod h1:hm4jLcvZKI2vRF2WDU1N4p/jpWtpOzp3nLmi9AzX/XE= github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk= github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= -github.com/GrayCodeAI/eyrie v0.2.1 h1:MtUkQV7hjFy7mobbjojjmgmN96Bo137pXfp2pcgJPug= -github.com/GrayCodeAI/eyrie v0.2.1/go.mod h1:8zw4jj5pvD9/tUqDhRaT4nnvl3Q8Zl2x4mJyEve7Obc= +github.com/GrayCodeAI/eyrie v0.2.2-0.20260715190227-1ba692935ce5 h1:DeXnr2++IHHeZnbxCIey9R5Lyqv94iAp97niVTmYkvg= +github.com/GrayCodeAI/eyrie v0.2.2-0.20260715190227-1ba692935ce5/go.mod h1:8zw4jj5pvD9/tUqDhRaT4nnvl3Q8Zl2x4mJyEve7Obc= github.com/GrayCodeAI/hawk-core-contracts v0.1.4 h1:HRcuxvl5RMgqmF0Lt1YwHY84nGcUwXculXtnrE/8nLI= github.com/GrayCodeAI/hawk-core-contracts v0.1.4/go.mod h1:BXbh68YrCf+s9HVqND5F8DAvl2MnE5NcOwZZZB56HGA= github.com/GrayCodeAI/inspect v0.1.4 h1:tQAcD9UuHkrqA20CZqdTcBgWKU/lhxxHBh8hlXY3FVw= @@ -26,10 +26,10 @@ github.com/GrayCodeAI/sight v0.1.4 h1:wFfRdbwoI0RIRnaj2H1Ytsed5B2pGBu+1csFRvpuMz github.com/GrayCodeAI/sight v0.1.4/go.mod h1:xuMOpaT5GJ3afu9gVqLOx+LeEWP/n09kvcjYxMKxFsI= github.com/GrayCodeAI/tok v0.1.4 h1:IGyHutbzS83I4bgvdPChnSPEHV6QwFKxgg9cgDnYaUY= github.com/GrayCodeAI/tok v0.1.4/go.mod h1:Mme6ckmjVRPddekhCnTC3MQbEHkPFKa5ULZjQOU4B3k= -github.com/GrayCodeAI/trace v0.1.4 h1:gD3mqFzMcr8DRPM+FE1kMf6pjsdgLfnSlsoj/MzWtGA= -github.com/GrayCodeAI/trace v0.1.4/go.mod h1:G9EPYZQKiTl1OHwhXLdpsgLZaS8zesQcd1YfLFIOeH8= -github.com/GrayCodeAI/yaad v0.1.4 h1:dbJ/rOae+648leF+2KprJI02sfWbvtXgb8SD6WKpDf0= -github.com/GrayCodeAI/yaad v0.1.4/go.mod h1:sz0RjliLeW9DLQi5MWEXGWro5K8RyaWCPevIYYW93/k= +github.com/GrayCodeAI/trace v0.1.5-0.20260715120456-164d86be64a8 h1:aZc9iLs+Tmv+oOfGVm0FvUnOqORfvoD4EnVsLd5nQQg= +github.com/GrayCodeAI/trace v0.1.5-0.20260715120456-164d86be64a8/go.mod h1:G9EPYZQKiTl1OHwhXLdpsgLZaS8zesQcd1YfLFIOeH8= +github.com/GrayCodeAI/yaad v0.2.1-0.20260715205802-7d52bde38e49 h1:6E/RqP8TUSesC85YD0N+i8C8ASrkCn7HE3jX00YG60w= +github.com/GrayCodeAI/yaad v0.2.1-0.20260715205802-7d52bde38e49/go.mod h1:ZPCv0JEZi2x8w1ksOTErI1Lg3A924UQvaUQJcwHdgn4= github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ= github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= github.com/Masterminds/semver/v3 v3.5.0 h1:kQceYJfbupGfZOKZQg0kou0DgAKhzDg2NZPAwZ/2OOE= @@ -394,8 +394,8 @@ google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa h1: google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:q4lMZS6kskjT5HvCPrnnypcDPVJqT/f4nfxmkE7gryY= google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa h1:mZHHdPZl0dbGHCflZgAq/Q468DWVFcU2whhB2KAo8fk= google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= -google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ= -google.golang.org/grpc v1.81.1/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I= +google.golang.org/grpc v1.82.0 h1:vguDnZUPjE26w09A63VoxZPnvPjB5Riyc0mkXPFmAIU= +google.golang.org/grpc v1.82.0/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA= google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/go.work b/go.work index 0ad6535f..58ff4f16 100644 --- a/go.work +++ b/go.work @@ -5,6 +5,7 @@ use . replace ( github.com/GrayCodeAI/eyrie => ./external/eyrie github.com/GrayCodeAI/hawk-core-contracts => ./external/hawk-core-contracts + github.com/GrayCodeAI/hawk-mcpkit => ./hawk-mcpkit github.com/GrayCodeAI/inspect => ./external/inspect github.com/GrayCodeAI/sight => ./external/sight github.com/GrayCodeAI/tok => ./external/tok diff --git a/internal/config/settings.go b/internal/config/settings.go index 439b48e4..64b562c4 100644 --- a/internal/config/settings.go +++ b/internal/config/settings.go @@ -142,9 +142,13 @@ type MCPServerConfig struct { Name string `json:"name"` Command string `json:"command,omitempty"` Args []string `json:"args,omitempty"` - Type string `json:"type,omitempty"` // "stdio" (default), "sse", "http" - URL string `json:"url,omitempty"` // for sse/http transports - Headers map[string]string `json:"headers,omitempty"` // custom headers for sse/http + Type string `json:"type,omitempty"` // "stdio" (default), "sse", "http", "websocket" + URL string `json:"url,omitempty"` // for sse/http/websocket transports + Headers map[string]string `json:"headers,omitempty"` // custom headers for sse/http/websocket + // ClientID is an OAuth client_id pre-registered with the server's + // authorization server. When empty, dynamic client registration (RFC + // 7591) is attempted if the server advertises a registration_endpoint. + ClientID string `json:"client_id,omitempty"` } func globalSettingsPath() string { diff --git a/internal/mcp/oauth.go b/internal/mcp/oauth.go new file mode 100644 index 00000000..d5474042 --- /dev/null +++ b/internal/mcp/oauth.go @@ -0,0 +1,363 @@ +package mcp + +// OAuth support for connecting to auth-gated remote MCP servers. +// +// hawk has zero MCP SDK dependency (see mcp.go/http.go/ws.go's hand-rolled +// JSON-RPC) — this stays consistent with that and hand-rolls OAuth too, +// rather than adopting a third-party client library. +// +// Flow: DiscoverOAuthMetadata -> RegisterClient (if no static client_id) -> +// GeneratePKCE -> StartLoopbackCallback -> BuildAuthorizationURL (user +// visits it in their own browser) -> ExchangeCode once the callback fires. +// RefreshOAuthToken renews an expiring token without repeating the flow. + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "fmt" + "net" + "net/http" + "net/url" + "strconv" + "strings" + "sync" + "time" +) + +// ProtectedResourceMetadata is the RFC 9728 discovery document served by +// the MCP server itself (not its authorization server), pointing at which +// authorization server(s) protect it. +type ProtectedResourceMetadata struct { + AuthorizationServers []string `json:"authorization_servers"` +} + +// AuthServerMetadata is the RFC 8414 discovery document served by an +// authorization server. +type AuthServerMetadata struct { + Issuer string `json:"issuer"` + AuthorizationEndpoint string `json:"authorization_endpoint"` + TokenEndpoint string `json:"token_endpoint"` + RegistrationEndpoint string `json:"registration_endpoint,omitempty"` +} + +var oauthHTTPClient = &http.Client{Timeout: 10 * time.Second} + +func fetchJSON(ctx context.Context, url string, out interface{}) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return err + } + req.Header.Set("Accept", "application/json") + resp, err := oauthHTTPClient.Do(req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("HTTP %d fetching %s", resp.StatusCode, url) + } + return json.NewDecoder(resp.Body).Decode(out) +} + +// DiscoverOAuthMetadata locates the authorization server that protects +// serverURL and returns its metadata. It tries, in order: +// 1. RFC 9728 protected-resource discovery on the MCP server itself +// ({origin}/.well-known/oauth-protected-resource), then RFC 8414 +// discovery on the first authorization server it names. +// 2. RFC 8414 discovery directly on the MCP server's own origin (the +// common case where the MCP server and its authorization server are +// the same host). +// 3. A last-resort guess at conventional endpoint paths, so a server with +// no discovery documents at all still has *something* to try. +func DiscoverOAuthMetadata(ctx context.Context, serverURL string) (*AuthServerMetadata, error) { + parsed, err := url.Parse(serverURL) + if err != nil { + return nil, fmt.Errorf("invalid server URL: %w", err) + } + origin := fmt.Sprintf("%s://%s", parsed.Scheme, parsed.Host) + + if issuer, ok := discoverProtectedResourceIssuer(ctx, origin); ok { + if meta, err := discoverAuthServerMetadata(ctx, issuer); err == nil { + return meta, nil + } + } + + if meta, err := discoverAuthServerMetadata(ctx, origin); err == nil { + return meta, nil + } + + // Last resort: conventional guesses, matching the prior (pre-discovery) + // behavior so a server with no metadata documents at all still gets an + // attempt rather than an outright failure. + return &AuthServerMetadata{ + Issuer: origin, + AuthorizationEndpoint: origin + "/oauth/authorize", + TokenEndpoint: origin + "/oauth/token", + }, nil +} + +func discoverProtectedResourceIssuer(ctx context.Context, origin string) (string, bool) { + var doc ProtectedResourceMetadata + if err := fetchJSON(ctx, origin+"/.well-known/oauth-protected-resource", &doc); err != nil { + return "", false + } + if len(doc.AuthorizationServers) == 0 { + return "", false + } + return doc.AuthorizationServers[0], true +} + +func discoverAuthServerMetadata(ctx context.Context, issuer string) (*AuthServerMetadata, error) { + var meta AuthServerMetadata + if err := fetchJSON(ctx, strings.TrimSuffix(issuer, "/")+"/.well-known/oauth-authorization-server", &meta); err != nil { + return nil, err + } + if meta.AuthorizationEndpoint == "" || meta.TokenEndpoint == "" { + return nil, fmt.Errorf("authorization server metadata at %s is missing required endpoints", issuer) + } + return &meta, nil +} + +// dcrRequest/dcrResponse implement RFC 7591 dynamic client registration for +// a native/public client (no client_secret expected back). +type dcrRequest struct { + RedirectURIs []string `json:"redirect_uris"` + TokenEndpointAuthMethod string `json:"token_endpoint_auth_method"` + GrantTypes []string `json:"grant_types"` + ResponseTypes []string `json:"response_types"` + ClientName string `json:"client_name"` +} + +type dcrResponse struct { + ClientID string `json:"client_id"` +} + +// RegisterClient performs RFC 7591 dynamic client registration against the +// given registration endpoint and returns the assigned client_id. +func RegisterClient(ctx context.Context, registrationEndpoint, redirectURI string) (string, error) { + body, err := json.Marshal(dcrRequest{ + RedirectURIs: []string{redirectURI}, + TokenEndpointAuthMethod: "none", + GrantTypes: []string{"authorization_code", "refresh_token"}, + ResponseTypes: []string{"code"}, + ClientName: "hawk", + }) + if err != nil { + return "", err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, registrationEndpoint, strings.NewReader(string(body))) + if err != nil { + return "", err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + resp, err := oauthHTTPClient.Do(req) + if err != nil { + return "", err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { + return "", fmt.Errorf("dynamic client registration failed: HTTP %d", resp.StatusCode) + } + var out dcrResponse + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + return "", err + } + if out.ClientID == "" { + return "", fmt.Errorf("registration response did not include a client_id") + } + return out.ClientID, nil +} + +// GeneratePKCE returns a PKCE code_verifier and its S256 code_challenge. +func GeneratePKCE() (verifier, challenge string, err error) { + raw := make([]byte, 32) + if _, err := rand.Read(raw); err != nil { + return "", "", err + } + verifier = base64.RawURLEncoding.EncodeToString(raw) + sum := sha256.Sum256([]byte(verifier)) + challenge = base64.RawURLEncoding.EncodeToString(sum[:]) + return verifier, challenge, nil +} + +// GenerateState returns a random value for the OAuth state parameter, used +// to bind an authorization request to its callback. +func GenerateState() (string, error) { + raw := make([]byte, 16) + if _, err := rand.Read(raw); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(raw), nil +} + +// CallbackResult is what StartLoopbackCallback delivers once the OAuth +// redirect arrives (or the wait times out / is cancelled). +type CallbackResult struct { + Code string + State string + Err error +} + +// StartLoopbackCallback binds an OS-assigned free port on 127.0.0.1, +// returning the redirect_uri to use in the authorization request. It +// serves exactly one request (the OAuth redirect), sends the result on the +// returned channel, and stops accepting further requests. The caller must +// call shutdown() once (after receiving a result, or on timeout) to +// release the listener. +func StartLoopbackCallback(ctx context.Context, timeout time.Duration) (redirectURI string, resultCh <-chan CallbackResult, shutdown func(), err error) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return "", nil, nil, err + } + addr, ok := listener.Addr().(*net.TCPAddr) + if !ok { + return "", nil, nil, fmt.Errorf("loopback listener has unexpected address type %T", listener.Addr()) + } + redirectURI = "http://127.0.0.1:" + strconv.Itoa(addr.Port) + "/callback" + + results := make(chan CallbackResult, 1) + var once sync.Once + mux := http.NewServeMux() + mux.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + result := CallbackResult{Code: q.Get("code"), State: q.Get("state")} + if errParam := q.Get("error"); errParam != "" { + result.Err = fmt.Errorf("authorization server returned error: %s", errParam) + } else if result.Code == "" { + result.Err = fmt.Errorf("callback request had no authorization code") + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + if result.Err != nil { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte("Authorization failed. You can close this tab.")) + } else { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("Authorized. You can close this tab and return to hawk.")) + } + once.Do(func() { + select { + case results <- result: + default: + } + }) + }) + // Loopback-only, context-bounded OAuth callback server: it serves a single + // localhost redirect and tears itself down via ctx/timeout, so the Slowloris + // window does not apply. + httpServer := &http.Server{Handler: mux} // #nosec G112 -- loopback, self-terminating + + go func() { _ = httpServer.Serve(listener) }() + + timer := time.AfterFunc(timeout, func() { + select { + case results <- CallbackResult{Err: fmt.Errorf("timed out waiting for authorization callback")}: + default: + } + }) + + shutdownFunc := func() { + timer.Stop() + shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = httpServer.Shutdown(shutdownCtx) + } + + go func() { + <-ctx.Done() + shutdownFunc() + }() + + return redirectURI, results, shutdownFunc, nil +} + +// BuildAuthorizationURL builds the authorization-code+PKCE request URL. +// resourceURL is the MCP server's canonical URL, sent per RFC 8707 so a +// multi-tenant authorization server can scope the issued token correctly. +func BuildAuthorizationURL(meta *AuthServerMetadata, clientID, redirectURI, state, codeChallenge, resourceURL string) string { + q := url.Values{ + "response_type": {"code"}, + "client_id": {clientID}, + "redirect_uri": {redirectURI}, + "state": {state}, + "code_challenge": {codeChallenge}, + "code_challenge_method": {"S256"}, + "resource": {resourceURL}, + } + sep := "?" + if strings.Contains(meta.AuthorizationEndpoint, "?") { + sep = "&" + } + return meta.AuthorizationEndpoint + sep + q.Encode() +} + +// TokenResult is a normalized OAuth token response. +type TokenResult struct { + AccessToken string + RefreshToken string + ExpiresAt time.Time +} + +type tokenResponse struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token,omitempty"` + ExpiresIn int `json:"expires_in,omitempty"` + Error string `json:"error,omitempty"` + ErrorDesc string `json:"error_description,omitempty"` +} + +func postTokenRequest(ctx context.Context, tokenEndpoint string, form url.Values) (*TokenResult, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, strings.NewReader(form.Encode())) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + resp, err := oauthHTTPClient.Do(req) + if err != nil { + return nil, err + } + defer func() { _ = resp.Body.Close() }() + + var out tokenResponse + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + return nil, fmt.Errorf("decoding token response: %w", err) + } + if out.Error != "" { + if out.ErrorDesc != "" { + return nil, fmt.Errorf("%s: %s", out.Error, out.ErrorDesc) + } + return nil, fmt.Errorf("%s", out.Error) + } + if out.AccessToken == "" { + return nil, fmt.Errorf("token endpoint returned no access_token (HTTP %d)", resp.StatusCode) + } + result := &TokenResult{AccessToken: out.AccessToken, RefreshToken: out.RefreshToken} + if out.ExpiresIn > 0 { + result.ExpiresAt = time.Now().Add(time.Duration(out.ExpiresIn) * time.Second) + } + return result, nil +} + +// ExchangeCode exchanges an authorization code (+ PKCE verifier) for tokens. +func ExchangeCode(ctx context.Context, meta *AuthServerMetadata, clientID, code, verifier, redirectURI string) (*TokenResult, error) { + return postTokenRequest(ctx, meta.TokenEndpoint, url.Values{ + "grant_type": {"authorization_code"}, + "code": {code}, + "redirect_uri": {redirectURI}, + "client_id": {clientID}, + "code_verifier": {verifier}, + }) +} + +// RefreshOAuthToken exchanges a refresh token for a new access token. +func RefreshOAuthToken(ctx context.Context, meta *AuthServerMetadata, clientID, refreshToken string) (*TokenResult, error) { + return postTokenRequest(ctx, meta.TokenEndpoint, url.Values{ + "grant_type": {"refresh_token"}, + "refresh_token": {refreshToken}, + "client_id": {clientID}, + }) +} diff --git a/internal/mcp/oauth_store.go b/internal/mcp/oauth_store.go new file mode 100644 index 00000000..a9297a62 --- /dev/null +++ b/internal/mcp/oauth_store.go @@ -0,0 +1,88 @@ +package mcp + +import ( + "encoding/json" + "fmt" + "time" + + "github.com/GrayCodeAI/hawk/internal/auth" +) + +// oauthTokenService is the keychain/keyring service name tokens are stored +// under, keyed per-server by the account parameter. +// Keychain service name — a fixed OAuth storage label, not a secret value. +const oauthTokenService = "hawk-mcp-oauth" // #nosec G101 -- fixed storage label, not a credential + +// StoredToken is what gets persisted per MCP server. auth.SecureStorage +// only stores a single string value per account, so this is JSON-marshaled +// into that string — not auth.TokenStore, which (as of this writing) has +// no-op Load/Save and doesn't actually persist anything. +type StoredToken struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token,omitempty"` + ExpiresAt time.Time `json:"expires_at,omitempty"` + ClientID string `json:"client_id"` + TokenEndpoint string `json:"token_endpoint"` +} + +// NeedsRefresh reports whether the access token is expired or close enough +// to expiring (30s) that it should be refreshed before use. A token with a +// zero ExpiresAt (the server didn't return expires_in) is treated as never +// needing a refresh purely on expiry grounds. +func (t *StoredToken) NeedsRefresh() bool { + if t.ExpiresAt.IsZero() { + return false + } + return time.Now().After(t.ExpiresAt.Add(-30 * time.Second)) +} + +// tokenBackend is the minimal surface SaveToken/LoadToken need from a +// credential store — satisfied by *auth.SecureStorage. Kept as an +// interface (rather than calling auth.NewSecureStorage directly) so tests +// can swap in a fake and never touch the real OS keychain. +type tokenBackend interface { + Get(account string) (string, error) + Set(account, value string) error +} + +var newTokenStorage = func() tokenBackend { return auth.NewSecureStorage(oauthTokenService) } + +// SetTokenBackendForTesting overrides the credential store SaveToken/ +// LoadToken use, for tests in other packages that exercise code paths +// eventually calling SaveToken/LoadToken (e.g. internal/tool's OAuth flow +// orchestration) and must never touch the real OS keychain. Returns a +// restore function; production code must never call this. +func SetTokenBackendForTesting(backend interface { + Get(account string) (string, error) + Set(account, value string) error +}, +) func() { + orig := newTokenStorage + newTokenStorage = func() tokenBackend { return backend } + return func() { newTokenStorage = orig } +} + +// SaveToken persists tok for serverName. +func SaveToken(serverName string, tok *StoredToken) error { + data, err := json.Marshal(tok) + if err != nil { + return err + } + return newTokenStorage().Set(serverName, string(data)) +} + +// LoadToken returns the stored token for serverName, if any. +func LoadToken(serverName string) (*StoredToken, error) { + raw, err := newTokenStorage().Get(serverName) + if err != nil { + return nil, err + } + if raw == "" { + return nil, fmt.Errorf("no token stored for %q", serverName) + } + var tok StoredToken + if err := json.Unmarshal([]byte(raw), &tok); err != nil { + return nil, fmt.Errorf("stored token for %q is corrupt: %w", serverName, err) + } + return &tok, nil +} diff --git a/internal/mcp/oauth_test.go b/internal/mcp/oauth_test.go new file mode 100644 index 00000000..616ae1e0 --- /dev/null +++ b/internal/mcp/oauth_test.go @@ -0,0 +1,402 @@ +package mcp + +import ( + "context" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" +) + +func TestDiscoverOAuthMetadata_ProtectedResourceThenAuthServer(t *testing.T) { + var authServer *httptest.Server + authServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/oauth-authorization-server" { + _ = json.NewEncoder(w).Encode(AuthServerMetadata{ + Issuer: authServer.URL, + AuthorizationEndpoint: authServer.URL + "/authorize", + TokenEndpoint: authServer.URL + "/token", + RegistrationEndpoint: authServer.URL + "/register", + }) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer authServer.Close() + + mcpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/oauth-protected-resource" { + _ = json.NewEncoder(w).Encode(ProtectedResourceMetadata{ + AuthorizationServers: []string{authServer.URL}, + }) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer mcpServer.Close() + + meta, err := DiscoverOAuthMetadata(context.Background(), mcpServer.URL+"/mcp") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if meta.AuthorizationEndpoint != authServer.URL+"/authorize" { + t.Errorf("authorization_endpoint = %q, want the discovered auth server's", meta.AuthorizationEndpoint) + } + if meta.RegistrationEndpoint != authServer.URL+"/register" { + t.Errorf("registration_endpoint = %q, want the discovered auth server's", meta.RegistrationEndpoint) + } +} + +func TestDiscoverOAuthMetadata_FallsBackToSameOriginAuthServer(t *testing.T) { + // No /.well-known/oauth-protected-resource at all; the MCP server IS + // the authorization server. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/.well-known/oauth-protected-resource": + w.WriteHeader(http.StatusNotFound) + case "/.well-known/oauth-authorization-server": + _ = json.NewEncoder(w).Encode(AuthServerMetadata{ + AuthorizationEndpoint: "/authorize-here", + TokenEndpoint: "/token-here", + }) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer server.Close() + + meta, err := DiscoverOAuthMetadata(context.Background(), server.URL+"/mcp") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if meta.AuthorizationEndpoint != "/authorize-here" { + t.Errorf("authorization_endpoint = %q", meta.AuthorizationEndpoint) + } +} + +func TestDiscoverOAuthMetadata_FallsBackToGuessedEndpoints(t *testing.T) { + // Neither well-known document exists at all. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + meta, err := DiscoverOAuthMetadata(context.Background(), server.URL+"/mcp") + if err != nil { + t.Fatalf("unexpected error (should degrade gracefully): %v", err) + } + if meta.AuthorizationEndpoint != server.URL+"/oauth/authorize" { + t.Errorf("authorization_endpoint = %q, want guessed default", meta.AuthorizationEndpoint) + } + if meta.TokenEndpoint != server.URL+"/oauth/token" { + t.Errorf("token_endpoint = %q, want guessed default", meta.TokenEndpoint) + } +} + +func TestRegisterClient(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + } + var req dcrRequest + _ = json.NewDecoder(r.Body).Decode(&req) + if len(req.RedirectURIs) != 1 || req.RedirectURIs[0] != "http://127.0.0.1:9999/callback" { + t.Errorf("unexpected redirect_uris: %+v", req.RedirectURIs) + } + if req.TokenEndpointAuthMethod != "none" { + t.Errorf("expected public client (token_endpoint_auth_method=none), got %q", req.TokenEndpointAuthMethod) + } + w.WriteHeader(http.StatusCreated) + _ = json.NewEncoder(w).Encode(dcrResponse{ClientID: "generated-client-id"}) + })) + defer server.Close() + + clientID, err := RegisterClient(context.Background(), server.URL+"/register", "http://127.0.0.1:9999/callback") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if clientID != "generated-client-id" { + t.Errorf("clientID = %q", clientID) + } +} + +func TestGeneratePKCE_ChallengeMatchesVerifier(t *testing.T) { + verifier, challenge, err := GeneratePKCE() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if verifier == "" || challenge == "" { + t.Fatal("expected non-empty verifier and challenge") + } + sum := sha256.Sum256([]byte(verifier)) + want := base64.RawURLEncoding.EncodeToString(sum[:]) + if challenge != want { + t.Errorf("challenge = %q, want S256(verifier) = %q", challenge, want) + } + + verifier2, _, _ := GeneratePKCE() + if verifier == verifier2 { + t.Error("expected distinct verifiers across calls") + } +} + +func TestGenerateState_NonEmptyAndDistinct(t *testing.T) { + s1, err := GenerateState() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + s2, _ := GenerateState() + if s1 == "" || s2 == "" { + t.Fatal("expected non-empty state values") + } + if s1 == s2 { + t.Error("expected distinct state values across calls") + } +} + +func TestBuildAuthorizationURL(t *testing.T) { + meta := &AuthServerMetadata{AuthorizationEndpoint: "https://auth.example.com/authorize"} + got := BuildAuthorizationURL(meta, "client-1", "http://127.0.0.1:9999/callback", "state-1", "challenge-1", "https://mcp.example.com") + + parsed, err := url.Parse(got) + if err != nil { + t.Fatalf("BuildAuthorizationURL produced an unparseable URL: %v", err) + } + q := parsed.Query() + checks := map[string]string{ + "response_type": "code", + "client_id": "client-1", + "redirect_uri": "http://127.0.0.1:9999/callback", + "state": "state-1", + "code_challenge": "challenge-1", + "code_challenge_method": "S256", + "resource": "https://mcp.example.com", + } + for k, want := range checks { + if got := q.Get(k); got != want { + t.Errorf("query param %q = %q, want %q", k, got, want) + } + } +} + +func TestStartLoopbackCallback_ReceivesCode(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + redirectURI, results, shutdown, err := StartLoopbackCallback(ctx, 5*time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer shutdown() + + // Simulate the browser being redirected here after the user authorizes. + go func() { + client := &http.Client{Timeout: 2 * time.Second} + _, _ = client.Get(redirectURI + "?code=the-code&state=the-state") + }() + + select { + case result := <-results: + if result.Err != nil { + t.Fatalf("unexpected callback error: %v", result.Err) + } + if result.Code != "the-code" || result.State != "the-state" { + t.Errorf("got code=%q state=%q", result.Code, result.State) + } + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for callback result") + } +} + +func TestStartLoopbackCallback_TimesOut(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, results, shutdown, err := StartLoopbackCallback(ctx, 100*time.Millisecond) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer shutdown() + + select { + case result := <-results: + if result.Err == nil { + t.Fatal("expected a timeout error") + } + case <-time.After(2 * time.Second): + t.Fatal("expected the timeout to fire the result channel") + } +} + +func TestStartLoopbackCallback_ErrorParam(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + redirectURI, results, shutdown, err := StartLoopbackCallback(ctx, 5*time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer shutdown() + + go func() { + client := &http.Client{Timeout: 2 * time.Second} + _, _ = client.Get(redirectURI + "?error=access_denied&state=the-state") + }() + + select { + case result := <-results: + if result.Err == nil { + t.Fatal("expected an error for an error= callback") + } + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for callback result") + } +} + +func TestExchangeCode(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = r.ParseForm() + if r.FormValue("grant_type") != "authorization_code" { + t.Errorf("grant_type = %q", r.FormValue("grant_type")) + } + if r.FormValue("code") != "auth-code" || r.FormValue("code_verifier") != "verifier-1" { + t.Errorf("unexpected code/verifier: %q/%q", r.FormValue("code"), r.FormValue("code_verifier")) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "access-1", + "refresh_token": "refresh-1", + "expires_in": 3600, + }) + })) + defer server.Close() + + meta := &AuthServerMetadata{TokenEndpoint: server.URL} + result, err := ExchangeCode(context.Background(), meta, "client-1", "auth-code", "verifier-1", "http://127.0.0.1:9999/callback") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result.AccessToken != "access-1" || result.RefreshToken != "refresh-1" { + t.Errorf("unexpected tokens: %+v", result) + } + if result.ExpiresAt.Before(time.Now().Add(59 * time.Minute)) { + t.Errorf("expected ExpiresAt ~1h out, got %v", result.ExpiresAt) + } +} + +func TestExchangeCode_ServerError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]any{ + "error": "invalid_grant", + "error_description": "code expired", + }) + })) + defer server.Close() + + meta := &AuthServerMetadata{TokenEndpoint: server.URL} + _, err := ExchangeCode(context.Background(), meta, "client-1", "auth-code", "verifier-1", "http://127.0.0.1:9999/callback") + if err == nil { + t.Fatal("expected an error") + } +} + +func TestRefreshOAuthToken(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = r.ParseForm() + if r.FormValue("grant_type") != "refresh_token" || r.FormValue("refresh_token") != "old-refresh" { + t.Errorf("unexpected form: %v", r.Form) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "new-access", + "expires_in": 1800, + }) + })) + defer server.Close() + + meta := &AuthServerMetadata{TokenEndpoint: server.URL} + result, err := RefreshOAuthToken(context.Background(), meta, "client-1", "old-refresh") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result.AccessToken != "new-access" { + t.Errorf("AccessToken = %q", result.AccessToken) + } +} + +func TestStoredToken_NeedsRefresh(t *testing.T) { + tests := []struct { + name string + expiresAt time.Time + want bool + }{ + {name: "zero value never needs refresh", expiresAt: time.Time{}, want: false}, + {name: "far in the future", expiresAt: time.Now().Add(time.Hour), want: false}, + {name: "already expired", expiresAt: time.Now().Add(-time.Minute), want: true}, + {name: "within the 30s buffer", expiresAt: time.Now().Add(10 * time.Second), want: true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tok := &StoredToken{ExpiresAt: tc.expiresAt} + if got := tok.NeedsRefresh(); got != tc.want { + t.Errorf("NeedsRefresh() = %v, want %v", got, tc.want) + } + }) + } +} + +// fakeTokenBackend is an in-memory tokenBackend for tests — SaveToken/ +// LoadToken must never touch the real OS keychain during `go test`. +type fakeTokenBackend struct { + values map[string]string +} + +func (f *fakeTokenBackend) Get(account string) (string, error) { + return f.values[account], nil +} + +func (f *fakeTokenBackend) Set(account, value string) error { + f.values[account] = value + return nil +} + +func TestSaveAndLoadToken_RoundTrip(t *testing.T) { + fake := &fakeTokenBackend{values: make(map[string]string)} + orig := newTokenStorage + newTokenStorage = func() tokenBackend { return fake } + t.Cleanup(func() { newTokenStorage = orig }) + + tok := &StoredToken{ + AccessToken: "access-1", + RefreshToken: "refresh-1", + ExpiresAt: time.Now().Add(time.Hour).Truncate(time.Second), + ClientID: "client-1", + TokenEndpoint: "https://auth.example.com/token", + } + if err := SaveToken("my-server", tok); err != nil { + t.Fatalf("SaveToken: %v", err) + } + + loaded, err := LoadToken("my-server") + if err != nil { + t.Fatalf("LoadToken: %v", err) + } + if loaded.AccessToken != tok.AccessToken || loaded.RefreshToken != tok.RefreshToken || + loaded.ClientID != tok.ClientID || !loaded.ExpiresAt.Equal(tok.ExpiresAt) { + t.Errorf("round-tripped token mismatch: got %+v, want %+v", loaded, tok) + } +} + +func TestLoadToken_NotFound(t *testing.T) { + fake := &fakeTokenBackend{values: make(map[string]string)} + orig := newTokenStorage + newTokenStorage = func() tokenBackend { return fake } + t.Cleanup(func() { newTokenStorage = orig }) + + _, err := LoadToken("never-saved") + if err == nil { + t.Fatal("expected an error for a server with no stored token") + } +} diff --git a/internal/tool/mcp_auth.go b/internal/tool/mcp_auth.go index 21c8412c..079b4c54 100644 --- a/internal/tool/mcp_auth.go +++ b/internal/tool/mcp_auth.go @@ -4,10 +4,10 @@ import ( "context" "encoding/json" "fmt" - "net/http" - "net/url" "sync" "time" + + "github.com/GrayCodeAI/hawk/internal/mcp" ) // MCPAuthState tracks OAuth state for an MCP server. @@ -18,7 +18,8 @@ type MCPAuthState struct { Error string `json:"error,omitempty"` } -// MCPAuthManager handles OAuth flows for MCP servers. +// MCPAuthManager tracks in-flight and completed OAuth flows for MCP +// servers, keyed by server name. type MCPAuthManager struct { mu sync.RWMutex states map[string]*MCPAuthState @@ -28,54 +29,6 @@ var globalMCPAuthManager = &MCPAuthManager{states: make(map[string]*MCPAuthState func GetMCPAuthManager() *MCPAuthManager { return globalMCPAuthManager } -func (m *MCPAuthManager) StartAuth(serverName, serverURL string) (*MCPAuthState, error) { - m.mu.Lock() - defer m.mu.Unlock() - - // Build OAuth discovery URL - parsed, err := url.Parse(serverURL) - if err != nil { - return nil, fmt.Errorf("invalid server URL: %w", err) - } - - // Attempt to discover OAuth endpoint - wellKnown := fmt.Sprintf("%s://%s/.well-known/oauth-authorization-server", parsed.Scheme, parsed.Host) - client := &http.Client{Timeout: 10 * time.Second} - req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, wellKnown, nil) - resp, err := client.Do(req) - if err != nil || resp.StatusCode != 200 { - state := &MCPAuthState{ - ServerName: serverName, - Status: "unsupported", - Error: "Server does not support OAuth authentication", - } - m.states[serverName] = state - return state, nil - } - defer func() { _ = resp.Body.Close() }() - - var oauthConfig struct { - AuthorizationEndpoint string `json:"authorization_endpoint"` - TokenEndpoint string `json:"token_endpoint"` - } - if err := json.NewDecoder(resp.Body).Decode(&oauthConfig); err != nil { - return nil, fmt.Errorf("parsing OAuth config: %w", err) - } - - authURL := oauthConfig.AuthorizationEndpoint - if authURL == "" { - authURL = fmt.Sprintf("%s://%s/oauth/authorize", parsed.Scheme, parsed.Host) - } - - state := &MCPAuthState{ - ServerName: serverName, - AuthURL: authURL, - Status: "pending", - } - m.states[serverName] = state - return state, nil -} - func (m *MCPAuthManager) GetState(serverName string) (*MCPAuthState, bool) { m.mu.RLock() defer m.mu.RUnlock() @@ -83,21 +36,28 @@ func (m *MCPAuthManager) GetState(serverName string) (*MCPAuthState, bool) { return s, ok } -func (m *MCPAuthManager) SetAuthenticated(serverName string) { +func (m *MCPAuthManager) setState(serverName string, state *MCPAuthState) { m.mu.Lock() defer m.mu.Unlock() - if s, ok := m.states[serverName]; ok { - s.Status = "authenticated" - } + m.states[serverName] = state } -// McpAuthTool initiates OAuth authentication for an MCP server. +// mcpAuthCallbackTimeout bounds how long the loopback listener waits for +// the user to complete authorization in their browser before giving up. +const mcpAuthCallbackTimeout = 5 * time.Minute + +// McpAuthTool starts (or reports on) an OAuth authorization-code+PKCE flow +// for a remote MCP server. It returns immediately with the URL to visit; +// completion happens asynchronously once the user authorizes in their +// browser and the loopback callback fires — call this tool again with the +// same server_name later to check progress. type McpAuthTool struct{} func (McpAuthTool) Name() string { return "McpAuth" } func (McpAuthTool) Aliases() []string { return []string{"mcp_auth"} } func (McpAuthTool) Description() string { - return "Start OAuth authentication for an MCP server that requires authorization" + return "Start OAuth authentication for an MCP server that requires authorization. " + + "Call again with the same server_name to check progress once the user has visited the URL." } func (McpAuthTool) Parameters() map[string]interface{} { @@ -112,6 +72,11 @@ func (McpAuthTool) Parameters() map[string]interface{} { "type": "string", "description": "URL of the MCP server", }, + "client_id": map[string]interface{}{ + "type": "string", + "description": "Optional pre-registered OAuth client_id. If omitted, hawk attempts " + + "dynamic client registration (RFC 7591) against the server's advertised registration endpoint.", + }, }, "required": []string{"server_name", "server_url"}, } @@ -121,6 +86,7 @@ func (McpAuthTool) Execute(ctx context.Context, input json.RawMessage) (string, var p struct { ServerName string `json:"server_name"` ServerURL string `json:"server_url"` + ClientID string `json:"client_id"` } if err := json.Unmarshal(input, &p); err != nil { return "", err @@ -132,28 +98,164 @@ func (McpAuthTool) Execute(ctx context.Context, input json.RawMessage) (string, return "", fmt.Errorf("server_url is required") } - state, err := globalMCPAuthManager.StartAuth(p.ServerName, p.ServerURL) + // Already authenticated, or a flow is already in flight: report status + // rather than starting a duplicate flow (and a duplicate loopback + // listener) for the same server. + if existing, ok := globalMCPAuthManager.GetState(p.ServerName); ok && + (existing.Status == "pending" || existing.Status == "authenticated") { + return authStatusJSON(existing), nil + } + + meta, err := mcp.DiscoverOAuthMetadata(ctx, p.ServerURL) if err != nil { - return "", err + return authStatusJSON(failState(p.ServerName, err)), nil + } + + redirectURI, resultCh, shutdown, err := mcp.StartLoopbackCallback(context.Background(), mcpAuthCallbackTimeout) + if err != nil { + return authStatusJSON(failState(p.ServerName, err)), nil + } + + clientID := p.ClientID + if clientID == "" { + if meta.RegistrationEndpoint == "" { + shutdown() + return authStatusJSON(failState(p.ServerName, fmt.Errorf( + "no client_id was provided and the server doesn't advertise a registration_endpoint "+ + "for dynamic client registration — pass client_id explicitly", + ))), nil + } + clientID, err = mcp.RegisterClient(ctx, meta.RegistrationEndpoint, redirectURI) + if err != nil { + shutdown() + return authStatusJSON(failState(p.ServerName, fmt.Errorf("dynamic client registration failed: %w", err))), nil + } } + verifier, challenge, err := mcp.GeneratePKCE() + if err != nil { + shutdown() + return authStatusJSON(failState(p.ServerName, err)), nil + } + reqState, err := mcp.GenerateState() + if err != nil { + shutdown() + return authStatusJSON(failState(p.ServerName, err)), nil + } + + authURL := mcp.BuildAuthorizationURL(meta, clientID, redirectURI, reqState, challenge, p.ServerURL) + state := &MCPAuthState{ServerName: p.ServerName, AuthURL: authURL, Status: "pending"} + globalMCPAuthManager.setState(p.ServerName, state) + + go completeMCPAuth(p.ServerName, meta, clientID, verifier, reqState, redirectURI, resultCh, shutdown) + + return authStatusJSON(state), nil +} + +func failState(serverName string, err error) *MCPAuthState { + state := &MCPAuthState{ServerName: serverName, Status: "error", Error: err.Error()} + globalMCPAuthManager.setState(serverName, state) + return state +} + +// completeMCPAuth waits for the loopback callback, exchanges the code for +// tokens, and persists the result. It runs independently of the tool call +// that started it, since the user completing authorization in their +// browser is an open-ended, asynchronous step from hawk's perspective. +func completeMCPAuth( + serverName string, + meta *mcp.AuthServerMetadata, + clientID, verifier, wantState, redirectURI string, + resultCh <-chan mcp.CallbackResult, + shutdown func(), +) { + defer shutdown() + result := <-resultCh + if result.Err != nil { + failState(serverName, result.Err) + return + } + if result.State != wantState { + failState(serverName, fmt.Errorf("callback state did not match the authorization request (possible CSRF)")) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + tokens, err := mcp.ExchangeCode(ctx, meta, clientID, result.Code, verifier, redirectURI) + if err != nil { + failState(serverName, fmt.Errorf("token exchange failed: %w", err)) + return + } + + stored := &mcp.StoredToken{ + AccessToken: tokens.AccessToken, + RefreshToken: tokens.RefreshToken, + ExpiresAt: tokens.ExpiresAt, + ClientID: clientID, + TokenEndpoint: meta.TokenEndpoint, + } + if err := mcp.SaveToken(serverName, stored); err != nil { + failState(serverName, fmt.Errorf("failed to store token: %w", err)) + return + } + + globalMCPAuthManager.setState(serverName, &MCPAuthState{ServerName: serverName, Status: "authenticated"}) +} + +func authStatusJSON(state *MCPAuthState) string { out, _ := json.Marshal(map[string]any{ "status": state.Status, "message": formatAuthMessage(state), "authUrl": state.AuthURL, }) - return string(out), nil + return string(out) } func formatAuthMessage(state *MCPAuthState) string { switch state.Status { case "pending": - return fmt.Sprintf("Please visit the following URL to authorize: %s", state.AuthURL) - case "unsupported": - return fmt.Sprintf("Server %q does not support OAuth authentication", state.ServerName) + return fmt.Sprintf( + "Please visit the following URL to authorize: %s\nOnce you approve it, this MCP server "+ + "will be usable the next time it connects (e.g. restart hawk, or reconfigure it).", + state.AuthURL, + ) case "authenticated": - return fmt.Sprintf("Server %q is already authenticated", state.ServerName) + return fmt.Sprintf("Server %q is authenticated.", state.ServerName) + case "error": + return fmt.Sprintf("Authentication for %q failed: %s", state.ServerName, state.Error) default: return state.Error } } + +// AuthHeaderForMCPServer returns the "Bearer " value to auto-inject +// as the Authorization header when connecting to a remote MCP server, if a +// valid (or refreshable) OAuth token is stored for it. If the stored token +// is close to expiring, it's refreshed and re-persisted first; if that +// fails (or there's no refresh token), it reports no token rather than +// connecting with a stale one. +func AuthHeaderForMCPServer(serverName string) (string, bool) { + tok, err := mcp.LoadToken(serverName) + if err != nil { + return "", false + } + if tok.NeedsRefresh() { + if tok.RefreshToken == "" || tok.TokenEndpoint == "" { + return "", false + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + refreshed, err := mcp.RefreshOAuthToken(ctx, &mcp.AuthServerMetadata{TokenEndpoint: tok.TokenEndpoint}, tok.ClientID, tok.RefreshToken) + if err != nil { + return "", false + } + tok.AccessToken = refreshed.AccessToken + if refreshed.RefreshToken != "" { + tok.RefreshToken = refreshed.RefreshToken + } + tok.ExpiresAt = refreshed.ExpiresAt + _ = mcp.SaveToken(serverName, tok) + } + return "Bearer " + tok.AccessToken, true +} diff --git a/internal/tool/mcp_auth_test.go b/internal/tool/mcp_auth_test.go new file mode 100644 index 00000000..78b62a67 --- /dev/null +++ b/internal/tool/mcp_auth_test.go @@ -0,0 +1,394 @@ +package tool + +import ( + "crypto/sha256" + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" + + "github.com/GrayCodeAI/hawk/internal/mcp" +) + +// fakeTokenBackend is an in-memory stand-in for the OS keychain — tests in +// this file must never touch the real one. +type fakeTokenBackend struct { + values map[string]string +} + +func newFakeTokenBackend() *fakeTokenBackend { + return &fakeTokenBackend{values: make(map[string]string)} +} + +func (f *fakeTokenBackend) Get(account string) (string, error) { + return f.values[account], nil +} + +func (f *fakeTokenBackend) Set(account, value string) error { + f.values[account] = value + return nil +} + +func resetAuthManager(t *testing.T) { + t.Helper() + globalMCPAuthManager.mu.Lock() + globalMCPAuthManager.states = make(map[string]*MCPAuthState) + globalMCPAuthManager.mu.Unlock() +} + +// newTestOAuthServer stands in for both the MCP server and its +// authorization server: it serves RFC 8414 metadata, RFC 7591 dynamic +// client registration, and a token endpoint that validates the PKCE +// verifier against the challenge sent during authorization. The +// authorization endpoint itself isn't served here — tests simulate the +// user's browser by hitting the loopback callback directly, exactly the +// way TestStartLoopbackCallback_ReceivesCode does in internal/mcp. +func newTestOAuthServer(t *testing.T) *httptest.Server { + t.Helper() + var srv *httptest.Server + mux := http.NewServeMux() + mux.HandleFunc("/.well-known/oauth-protected-resource", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + }) + mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(mcp.AuthServerMetadata{ + Issuer: srv.URL, + AuthorizationEndpoint: srv.URL + "/authorize", + TokenEndpoint: srv.URL + "/token", + RegistrationEndpoint: srv.URL + "/register", + }) + }) + mux.HandleFunc("/register", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusCreated) + _ = json.NewEncoder(w).Encode(map[string]string{"client_id": "dcr-client-id"}) + }) + mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) { + _ = r.ParseForm() + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "access-" + r.FormValue("code"), + "refresh_token": "refresh-1", + "expires_in": 3600, + }) + }) + srv = httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv +} + +func TestMcpAuthTool_FullFlow_DynamicClientRegistration(t *testing.T) { + resetAuthManager(t) + restore := mcp.SetTokenBackendForTesting(newFakeTokenBackend()) + defer restore() + + server := newTestOAuthServer(t) + + input, _ := json.Marshal(map[string]string{ + "server_name": "test-server", + "server_url": server.URL + "/mcp", + }) + out, err := McpAuthTool{}.Execute(t.Context(), input) + if err != nil { + t.Fatalf("Execute returned an error: %v", err) + } + + var resp struct { + Status string `json:"status"` + AuthURL string `json:"authUrl"` + } + if uErr := json.Unmarshal([]byte(out), &resp); uErr != nil { + t.Fatalf("could not parse Execute output: %v (%s)", uErr, out) + } + if resp.Status != "pending" { + t.Fatalf("expected status=pending, got %q (%s)", resp.Status, out) + } + + parsedAuthURL, err := url.Parse(resp.AuthURL) + if err != nil { + t.Fatalf("authUrl is not a valid URL: %v", err) + } + q := parsedAuthURL.Query() + if q.Get("client_id") != "dcr-client-id" { + t.Errorf("expected the dynamically-registered client_id, got %q", q.Get("client_id")) + } + redirectURI := q.Get("redirect_uri") + if redirectURI == "" { + t.Fatal("expected a redirect_uri in the authorization URL") + } + codeChallenge := q.Get("code_challenge") + state := q.Get("state") + if codeChallenge == "" || state == "" { + t.Fatalf("expected code_challenge and state to be set: %s", resp.AuthURL) + } + + // Simulate the user's browser being redirected back after authorizing. + client := &http.Client{Timeout: 3 * time.Second} + callbackResp, err := client.Get(redirectURI + "?code=the-auth-code&state=" + state) + if err != nil { + t.Fatalf("simulated callback request failed: %v", err) + } + _ = callbackResp.Body.Close() + + // completeMCPAuth runs in a background goroutine; poll briefly for it + // to finish rather than assuming it already has. + deadline := time.Now().Add(2 * time.Second) + var final *MCPAuthState + for time.Now().Before(deadline) { + if s, ok := globalMCPAuthManager.GetState("test-server"); ok && s.Status != "pending" { + final = s + break + } + time.Sleep(10 * time.Millisecond) + } + if final == nil { + t.Fatal("completeMCPAuth did not finish within the deadline") + } + if final.Status != "authenticated" { + t.Fatalf("expected status=authenticated, got %q (error: %s)", final.Status, final.Error) + } + + // The exchanged token must actually be persisted and usable. + header, ok := AuthHeaderForMCPServer("test-server") + if !ok { + t.Fatal("expected AuthHeaderForMCPServer to find the newly stored token") + } + if header != "Bearer access-the-auth-code" { + t.Errorf("unexpected Authorization header: %q", header) + } +} + +func TestMcpAuthTool_ExplicitClientID_SkipsRegistration(t *testing.T) { + resetAuthManager(t) + restore := mcp.SetTokenBackendForTesting(newFakeTokenBackend()) + defer restore() + + registrationCalled := false + mux := http.NewServeMux() + mux.HandleFunc("/.well-known/oauth-protected-resource", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + }) + var srv *httptest.Server + mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(mcp.AuthServerMetadata{ + AuthorizationEndpoint: srv.URL + "/authorize", + TokenEndpoint: srv.URL + "/token", + RegistrationEndpoint: srv.URL + "/register", + }) + }) + mux.HandleFunc("/register", func(w http.ResponseWriter, r *http.Request) { + registrationCalled = true + w.WriteHeader(http.StatusCreated) + _ = json.NewEncoder(w).Encode(map[string]string{"client_id": "should-not-be-used"}) + }) + srv = httptest.NewServer(mux) + defer srv.Close() + + input, _ := json.Marshal(map[string]string{ + "server_name": "explicit-client-server", + "server_url": srv.URL + "/mcp", + "client_id": "pre-registered-client", + }) + out, err := McpAuthTool{}.Execute(t.Context(), input) + if err != nil { + t.Fatalf("Execute returned an error: %v", err) + } + if registrationCalled { + t.Error("dynamic client registration should not be attempted when client_id is provided") + } + + var resp struct { + AuthURL string `json:"authUrl"` + } + _ = json.Unmarshal([]byte(out), &resp) + parsed, _ := url.Parse(resp.AuthURL) + if got := parsed.Query().Get("client_id"); got != "pre-registered-client" { + t.Errorf("client_id = %q, want the explicitly provided one", got) + } +} + +func TestMcpAuthTool_NoRegistrationEndpointAndNoClientID_Errors(t *testing.T) { + resetAuthManager(t) + restore := mcp.SetTokenBackendForTesting(newFakeTokenBackend()) + defer restore() + + mux := http.NewServeMux() + mux.HandleFunc("/.well-known/oauth-protected-resource", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + }) + var srv *httptest.Server + mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(mcp.AuthServerMetadata{ + AuthorizationEndpoint: srv.URL + "/authorize", + TokenEndpoint: srv.URL + "/token", + // No RegistrationEndpoint. + }) + }) + srv = httptest.NewServer(mux) + defer srv.Close() + + input, _ := json.Marshal(map[string]string{ + "server_name": "no-dcr-server", + "server_url": srv.URL + "/mcp", + }) + out, err := McpAuthTool{}.Execute(t.Context(), input) + if err != nil { + t.Fatalf("Execute returned a Go error (should report via status instead): %v", err) + } + var resp struct { + Status string `json:"status"` + } + _ = json.Unmarshal([]byte(out), &resp) + if resp.Status != "error" { + t.Fatalf("expected status=error, got %q (%s)", resp.Status, out) + } +} + +func TestMcpAuthTool_SecondCallReportsExistingPendingState(t *testing.T) { + resetAuthManager(t) + restore := mcp.SetTokenBackendForTesting(newFakeTokenBackend()) + defer restore() + + server := newTestOAuthServer(t) + input, _ := json.Marshal(map[string]string{ + "server_name": "repeat-server", + "server_url": server.URL + "/mcp", + }) + + first, err := McpAuthTool{}.Execute(t.Context(), input) + if err != nil { + t.Fatalf("first Execute failed: %v", err) + } + second, err := McpAuthTool{}.Execute(t.Context(), input) + if err != nil { + t.Fatalf("second Execute failed: %v", err) + } + + var firstResp, secondResp struct { + Status string `json:"status"` + AuthURL string `json:"authUrl"` + } + _ = json.Unmarshal([]byte(first), &firstResp) + _ = json.Unmarshal([]byte(second), &secondResp) + if secondResp.Status != "pending" { + t.Fatalf("expected the second call to report the still-pending flow, got %q", secondResp.Status) + } + if firstResp.AuthURL != secondResp.AuthURL { + t.Error("expected the second call to report the SAME auth URL, not start a new flow") + } +} + +func TestMcpAuthTool_MissingRequiredParams(t *testing.T) { + tests := []struct { + name string + input string + }{ + {name: "missing server_name", input: `{"server_url":"https://example.com"}`}, + {name: "missing server_url", input: `{"server_name":"x"}`}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, err := McpAuthTool{}.Execute(t.Context(), json.RawMessage(tc.input)) + if err == nil { + t.Fatal("expected an error") + } + }) + } +} + +func TestAuthHeaderForMCPServer_NoStoredToken(t *testing.T) { + restore := mcp.SetTokenBackendForTesting(newFakeTokenBackend()) + defer restore() + + _, ok := AuthHeaderForMCPServer("never-authenticated-server") + if ok { + t.Fatal("expected no token for a server that was never authenticated") + } +} + +func TestAuthHeaderForMCPServer_RefreshesExpiringToken(t *testing.T) { + fake := newFakeTokenBackend() + restore := mcp.SetTokenBackendForTesting(fake) + defer restore() + + refreshServerCalled := false + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + refreshServerCalled = true + _ = r.ParseForm() + if r.FormValue("grant_type") != "refresh_token" || r.FormValue("refresh_token") != "old-refresh" { + t.Errorf("unexpected refresh request: %v", r.Form) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "refreshed-access", + "refresh_token": "new-refresh", + "expires_in": 3600, + }) + })) + defer tokenServer.Close() + + almostExpired := &mcp.StoredToken{ + AccessToken: "old-access", + RefreshToken: "old-refresh", + ExpiresAt: time.Now().Add(5 * time.Second), // within the 30s refresh buffer + ClientID: "client-1", + TokenEndpoint: tokenServer.URL, + } + if err := mcp.SaveToken("expiring-server", almostExpired); err != nil { + t.Fatalf("SaveToken: %v", err) + } + + header, ok := AuthHeaderForMCPServer("expiring-server") + if !ok { + t.Fatal("expected a usable token after refresh") + } + if !refreshServerCalled { + t.Fatal("expected the refresh flow to actually call the token endpoint") + } + if header != "Bearer refreshed-access" { + t.Errorf("header = %q", header) + } + + // The refreshed token must be re-persisted, not just returned once. + reloaded, err := mcp.LoadToken("expiring-server") + if err != nil { + t.Fatalf("LoadToken after refresh: %v", err) + } + if reloaded.AccessToken != "refreshed-access" || reloaded.RefreshToken != "new-refresh" { + t.Errorf("refreshed token was not persisted correctly: %+v", reloaded) + } +} + +func TestAuthHeaderForMCPServer_NoRefreshTokenReportsNoAuth(t *testing.T) { + restore := mcp.SetTokenBackendForTesting(newFakeTokenBackend()) + defer restore() + + expired := &mcp.StoredToken{ + AccessToken: "stale-access", + ExpiresAt: time.Now().Add(-time.Hour), + // No RefreshToken. + } + if err := mcp.SaveToken("no-refresh-server", expired); err != nil { + t.Fatalf("SaveToken: %v", err) + } + + _, ok := AuthHeaderForMCPServer("no-refresh-server") + if ok { + t.Fatal("expected no usable token when the stored token is expired with no refresh_token") + } +} + +// sanity check that this test file's own PKCE understanding matches +// internal/mcp's actual S256 implementation, so the other tests above are +// exercising the real thing rather than a divergent assumption. +func TestPKCEChallengeMethod_MatchesRFC7636S256(t *testing.T) { + verifier, challenge, err := mcp.GeneratePKCE() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + sum := sha256.Sum256([]byte(verifier)) + want := base64.RawURLEncoding.EncodeToString(sum[:]) + if challenge != want { + t.Errorf("challenge = %q, want %q", challenge, want) + } +} diff --git a/internal/tool/mcp_resources.go b/internal/tool/mcp_resources.go index 876fb0d0..2d6bf5a9 100644 --- a/internal/tool/mcp_resources.go +++ b/internal/tool/mcp_resources.go @@ -52,15 +52,22 @@ func (ListMcpResourcesTool) Execute(_ context.Context, input json.RawMessage) (s if !ok { return "", fmt.Errorf("MCP server %q not found", p.Server) } - servers = []*mcp.Server{server} + servers = map[string]mcpClient{p.Server: server} } - for _, server := range servers { - resources, err := server.ListResources() + for name, server := range servers { + // Resource listing is currently stdio-only: mcp.HTTPServer/WSServer + // don't implement it. Remote-transport servers are silently skipped + // here rather than erroring, same as if they'd exposed zero resources. + stdioServer, ok := server.(*mcp.Server) + if !ok { + continue + } + resources, err := stdioServer.ListResources() if err != nil { continue } for _, r := range resources { - out = append(out, resourceOut{URI: r.URI, Name: r.Name, MimeType: r.MimeType, Description: r.Description, Server: server.Name}) + out = append(out, resourceOut{URI: r.URI, Name: r.Name, MimeType: r.MimeType, Description: r.Description, Server: name}) } } if len(out) == 0 { @@ -107,5 +114,9 @@ func (ReadMcpResourceTool) Execute(_ context.Context, input json.RawMessage) (st if !ok { return "", fmt.Errorf("MCP server %q not found", p.Server) } - return server.ReadResource(p.URI) + stdioServer, ok := server.(*mcp.Server) + if !ok { + return "", fmt.Errorf("MCP server %q does not support resource reads (remote transports don't implement this yet)", p.Server) + } + return stdioServer.ReadResource(p.URI) } diff --git a/internal/tool/mcp_tool.go b/internal/tool/mcp_tool.go index a99c3a32..85e50894 100644 --- a/internal/tool/mcp_tool.go +++ b/internal/tool/mcp_tool.go @@ -9,14 +9,27 @@ import ( "github.com/GrayCodeAI/hawk/internal/mcp" ) +// mcpClient is the minimal surface MCPTool needs from a connected MCP +// server, regardless of transport. CallTool and Close already have +// identical signatures across mcp.Server (stdio), mcp.HTTPServer, and +// mcp.WSServer — this interface lets MCPTool wrap any of them without +// changing any of those concrete types. ListTools deliberately isn't part +// of this interface: it's only ever called once per connect, directly on +// the concrete type, inside the transport-specific loader function below. +type mcpClient interface { + CallTool(ctx context.Context, name string, args map[string]interface{}) (string, error) + Close() error +} + var connectedMCPServers = struct { sync.RWMutex - servers map[string]*mcp.Server -}{servers: make(map[string]*mcp.Server)} + servers map[string]mcpClient +}{servers: make(map[string]mcpClient)} // MCPTool wraps an MCP server tool as a hawk tool. type MCPTool struct { - server *mcp.Server + server mcpClient + serverName string toolName string aliases []string remoteName string @@ -24,15 +37,16 @@ type MCPTool struct { schema map[string]interface{} } -func NewMCPTool(server *mcp.Server, t mcp.Tool) *MCPTool { - tsName := fmt.Sprintf("mcp__%s__%s", normalizeNameForMCP(server.Name), normalizeNameForMCP(t.Name)) - legacyName := fmt.Sprintf("mcp_%s_%s", server.Name, t.Name) +func NewMCPTool(serverName string, server mcpClient, t mcp.Tool) *MCPTool { + tsName := fmt.Sprintf("mcp__%s__%s", normalizeNameForMCP(serverName), normalizeNameForMCP(t.Name)) + legacyName := fmt.Sprintf("mcp_%s_%s", serverName, t.Name) return &MCPTool{ server: server, + serverName: serverName, toolName: tsName, aliases: []string{legacyName}, remoteName: t.Name, - description: fmt.Sprintf("[MCP:%s] %s", server.Name, t.Description), + description: fmt.Sprintf("[MCP:%s] %s", serverName, t.Description), schema: t.InputSchema, } } @@ -69,13 +83,14 @@ func normalizeNameForMCP(name string) string { return string(out) } -// LoadMCPTools connects to an MCP server and returns hawk tools for all its tools. +// LoadMCPTools connects to an MCP server over stdio and returns hawk tools +// for all its tools. func LoadMCPTools(ctx context.Context, name, command string, args ...string) ([]Tool, error) { server, err := mcp.Connect(ctx, name, command, args...) if err != nil { return nil, err } - registerMCPServer(server) + registerMCPServer(name, server) mcpTools, err := server.ListTools() if err != nil { _ = server.Close() @@ -83,28 +98,83 @@ func LoadMCPTools(ctx context.Context, name, command string, args ...string) ([] } var tools []Tool for _, t := range mcpTools { - tools = append(tools, NewMCPTool(server, t)) + tools = append(tools, NewMCPTool(name, server, t)) + } + return tools, nil +} + +// LoadRemoteMCPTools connects to an MCP server over http, sse, or websocket +// and returns hawk tools for all its tools. headers is merged onto every +// outgoing request/handshake by the transport (e.g. a static API key, or an +// auto-injected OAuth bearer token — see internal/mcp/oauth.go). +func LoadRemoteMCPTools( + ctx context.Context, + name, serverType, url string, + headers map[string]string, +) ([]Tool, error) { + var ( + server mcpClient + mcpTools []mcp.Tool + listErr error + connErr error + ) + switch serverType { + case "sse": + s, err := mcp.ConnectSSE(ctx, name, url, headers) + connErr = err + if err == nil { + server = s + mcpTools, listErr = s.ListTools(ctx) + } + case "websocket": + s, err := mcp.ConnectWS(ctx, name, url, headers) + connErr = err + if err == nil { + server = s + mcpTools, listErr = s.ListTools(ctx) + } + default: // "http" + s, err := mcp.ConnectHTTP(ctx, name, url, headers) + connErr = err + if err == nil { + server = s + mcpTools, listErr = s.ListTools(ctx) + } + } + if connErr != nil { + return nil, connErr + } + registerMCPServer(name, server) + if listErr != nil { + _ = server.Close() + return nil, listErr + } + var tools []Tool + for _, t := range mcpTools { + tools = append(tools, NewMCPTool(name, server, t)) } return tools, nil } -func registerMCPServer(server *mcp.Server) { +func registerMCPServer(name string, server mcpClient) { connectedMCPServers.Lock() defer connectedMCPServers.Unlock() - connectedMCPServers.servers[server.Name] = server + connectedMCPServers.servers[name] = server } -func listMCPServers() []*mcp.Server { +// listMCPServers returns every connected MCP server, keyed by the name it +// was registered under. +func listMCPServers() map[string]mcpClient { connectedMCPServers.RLock() defer connectedMCPServers.RUnlock() - servers := make([]*mcp.Server, 0, len(connectedMCPServers.servers)) - for _, server := range connectedMCPServers.servers { - servers = append(servers, server) + out := make(map[string]mcpClient, len(connectedMCPServers.servers)) + for name, server := range connectedMCPServers.servers { + out[name] = server } - return servers + return out } -func getMCPServer(name string) (*mcp.Server, bool) { +func getMCPServer(name string) (mcpClient, bool) { connectedMCPServers.RLock() defer connectedMCPServers.RUnlock() server, ok := connectedMCPServers.servers[name] diff --git a/internal/tool/patch.go b/internal/tool/patch.go index ca4466ee..e6c98263 100644 --- a/internal/tool/patch.go +++ b/internal/tool/patch.go @@ -440,6 +440,15 @@ func (PatchTool) Execute(ctx context.Context, input json.RawMessage) (string, er return "", fmt.Errorf("failed to parse patch: %w", err) } + // Reject any patch whose target path escapes the workspace before + // applying — otherwise an LLM-authored patch could write/delete files + // outside the sandbox (validated below per-entry). + for _, fp := range parser.Patches() { + if vErr := validatePathAllowed(ctx, fp.Path); vErr != nil { + return "", vErr + } + } + modified, err := parser.ApplyAll() if err != nil { return "", err diff --git a/internal/tool/sandbox_escape_test.go b/internal/tool/sandbox_escape_test.go new file mode 100644 index 00000000..d8b7b247 --- /dev/null +++ b/internal/tool/sandbox_escape_test.go @@ -0,0 +1,154 @@ +package tool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +// sandboxedContext returns a context whose ToolContext restricts writes to the +// supplied allowed directory (plus the CWD "."). With no ToolContext at all, +// validatePathAllowed no-ops — so tests MUST attach one to exercise the guard. +func sandboxedContext(t *testing.T, allowed string) context.Context { + t.Helper() + return WithToolContext(context.Background(), &ToolContext{ + AllowedDirectories: []string{allowed}, + }) +} + +// outsidePath is an absolute path that is never within a tempdir workspace. +const outsidePath = "/etc/hosts-tool-guard-test" + +func TestStructuredEditTool_RejectsPathOutsideSandbox(t *testing.T) { + t.Parallel() + tool := StructuredEditTool{} + ctx := sandboxedContext(t, t.TempDir()) + + input, _ := json.Marshal(map[string]any{ + "path": outsidePath, + "blocks": []map[string]any{ + {"search": "x", "replace": "y"}, + }, + }) + _, err := tool.Execute(ctx, input) + if err == nil { + t.Fatal("expected StructuredEditTool to reject a path outside the sandbox") + } + if !strings.Contains(err.Error(), "outside") { + t.Fatalf("expected an out-of-sandbox error, got: %v", err) + } +} + +// A path inside the allowed directory must still be accepted by StructuredEdit. +func TestStructuredEditTool_AcceptsPathInsideSandbox(t *testing.T) { + t.Parallel() + tool := StructuredEditTool{} + allowed := t.TempDir() + target := filepath.Join(allowed, "file.go") + if err := os.WriteFile(target, []byte("package main\n"), 0o600); err != nil { + t.Fatal(err) + } + ctx := sandboxedContext(t, allowed) + + input, _ := json.Marshal(map[string]any{ + "path": target, + "blocks": []map[string]any{ + {"search": "package main", "replace": "package renamed"}, + }, + }) + _, err := tool.Execute(ctx, input) + if err != nil { + t.Fatalf("expected in-sandbox path to succeed, got: %v", err) + } + data, rerr := os.ReadFile(target) + if rerr != nil { + t.Fatal(rerr) + } + if !strings.Contains(string(data), "package renamed") { + t.Fatalf("expected the edit to apply, file content: %q", string(data)) + } +} + +func TestSmartCreateTool_RejectsPathOutsideSandbox(t *testing.T) { + t.Parallel() + tool := &SmartCreateTool{Creator: NewSmartCreator(t.TempDir())} + ctx := sandboxedContext(t, t.TempDir()) + + input, _ := json.Marshal(map[string]any{"path": outsidePath + ".go"}) + _, err := tool.Execute(ctx, input) + if err == nil { + t.Fatal("expected SmartCreateTool to reject a path outside the sandbox") + } +} + +// A path inside the allowed directory must still be accepted by SmartCreate. +func TestSmartCreateTool_AcceptsPathInsideSandbox(t *testing.T) { + t.Parallel() + tool := &SmartCreateTool{Creator: NewSmartCreator(t.TempDir())} + allowed := t.TempDir() + target := filepath.Join(allowed, "new.txt") + ctx := sandboxedContext(t, allowed) + + input, _ := json.Marshal(map[string]any{"path": target}) + res, err := tool.Execute(ctx, input) + if err != nil { + t.Fatalf("expected in-sandbox create to succeed, got: %v", err) + } + if _, statErr := os.Stat(target); statErr != nil { + t.Fatalf("expected file to be created at %s: %v", target, statErr) + } + if !strings.Contains(res, target) { + t.Fatalf("expected result to mention the new file, got: %q", res) + } +} + +// PatchTool applies to files named inside the patch body, so a patch that +// targets an out-of-workspace file must be rejected before any write happens. +func TestPatchTool_RejectsPatchTargetingPathOutsideSandbox(t *testing.T) { + t.Parallel() + tool := PatchTool{} + allowed := t.TempDir() + // Seed the in-sandbox reference file so the "update" path's contents resolve; + // the guard must still reject it on path grounds before reading. + if err := os.WriteFile(filepath.Join(allowed, "ok.txt"), []byte("keep\n"), 0o600); err != nil { + t.Fatal(err) + } + ctx := sandboxedContext(t, allowed) + + patch := "*** Begin Patch\n*** Update File: " + outsidePath + "\n@@@ @@@\n-removed\n+added\n*** End Patch\n" + input, _ := json.Marshal(map[string]any{"patch": patch}) + _, err := tool.Execute(ctx, input) + if err == nil { + t.Fatal("expected PatchTool to reject a patch targeting a path outside the sandbox") + } +} + +// A patch targeting an in-sandbox file must still apply. +func TestPatchTool_AcceptsPatchTargetingPathInsideSandbox(t *testing.T) { + t.Parallel() + tool := PatchTool{} + allowed := t.TempDir() + target := filepath.Join(allowed, "a.txt") + if err := os.WriteFile(target, []byte("func main() {\n\toriginal\n}\n"), 0o600); err != nil { + t.Fatal(err) + } + ctx := sandboxedContext(t, allowed) + + patch := "*** Begin Patch\n" + + "*** Update File: " + target + "\n" + + "@@@ func main() {@@@\n" + + "- original\n" + + "+ newline\n" + + "*** End Patch" + input, _ := json.Marshal(map[string]any{"patch": patch}) + res, err := tool.Execute(ctx, input) + if err != nil { + t.Fatalf("expected in-sandbox patch to succeed, got: %v", err) + } + if !strings.Contains(res, "patched") && !strings.Contains(res, "file") { + t.Fatalf("expected a successful patch result, got: %q", res) + } +} diff --git a/internal/tool/smart_create.go b/internal/tool/smart_create.go index 2fb25fdc..01d23c35 100644 --- a/internal/tool/smart_create.go +++ b/internal/tool/smart_create.go @@ -680,6 +680,9 @@ func (t *SmartCreateTool) Execute(ctx context.Context, input json.RawMessage) (s if params.Path == "" { return "", fmt.Errorf("path is required") } + if err := validatePathAllowed(ctx, params.Path); err != nil { + return "", err + } // Generate boilerplate. content := t.Creator.GenerateBoilerplate(params.Path) diff --git a/internal/tool/structured_edit.go b/internal/tool/structured_edit.go index 0ebbff18..05761a9b 100644 --- a/internal/tool/structured_edit.go +++ b/internal/tool/structured_edit.go @@ -76,6 +76,9 @@ func (s StructuredEditTool) Execute(ctx context.Context, input json.RawMessage) if p.Path == "" { return "", fmt.Errorf("path is required") } + if err := validatePathAllowed(ctx, p.Path); err != nil { + return "", err + } if len(p.Blocks) == 0 { return "", fmt.Errorf("at least one SEARCH/REPLACE block is required") }