chore: oauth accepts APIHostResolver as argument

Author: atharva1051

pkg/http/oauth/oauth.go

@@ -3,11 +3,13 @@
 package oauth
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"strings"
 
 	"github.com/github/github-mcp-server/pkg/http/headers"
+	"github.com/github/github-mcp-server/pkg/utils"
 	"github.com/go-chi/chi/v5"
 	"github.com/modelcontextprotocol/go-sdk/auth"
 	"github.com/modelcontextprotocol/go-sdk/oauthex"
@@ -43,8 +45,13 @@ type Config struct {
 	// This is used to construct the OAuth resource URL.
 	BaseURL string
 
+	// APIHost is the GitHub API host resolver that provides OAuth URL.
+	// If set, this takes precedence over AuthorizationServer.
+	APIHost utils.APIHostResolver
+
 	// AuthorizationServer is the OAuth authorization server URL.
 	// Defaults to GitHub's OAuth server if not specified.
+	// This field is ignored if APIHost is set.
 	AuthorizationServer string
 
 	// ResourcePath is the externally visible base path for the MCP server (e.g., "/mcp").
@@ -64,8 +71,15 @@ func NewAuthHandler(cfg *Config) (*AuthHandler, error) {
 		cfg = &Config{}
 	}
 
-	// Default authorization server to GitHub
-	if cfg.AuthorizationServer == "" {
+	// Resolve authorization server from APIHost if provided
+	if cfg.APIHost != nil {
+		oauthURL, err := cfg.APIHost.OAuthURL(context.Background())
+		if err != nil {
+			return nil, fmt.Errorf("failed to get OAuth URL from API host: %w", err)
+		}
+		cfg.AuthorizationServer = oauthURL.String()
+	} else if cfg.AuthorizationServer == "" {
+		// Default authorization server to GitHub if not provided
 		cfg.AuthorizationServer = DefaultAuthorizationServer
 	}
 

pkg/http/oauth/oauth_test.go

@@ -1,18 +1,49 @@
 package oauth
 
 import (
+	"context"
 	"crypto/tls"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"testing"
 
 	"github.com/github/github-mcp-server/pkg/http/headers"
+	"github.com/github/github-mcp-server/pkg/utils"
 	"github.com/go-chi/chi/v5"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+// mockAPIHostResolver is a test implementation of utils.APIHostResolver
+type mockAPIHostResolver struct {
+	oauthURL string
+}
+
+func (m mockAPIHostResolver) BaseRESTURL(_ context.Context) (*url.URL, error) {
+	return nil, nil
+}
+
+func (m mockAPIHostResolver) GraphqlURL(_ context.Context) (*url.URL, error) {
+	return nil, nil
+}
+
+func (m mockAPIHostResolver) UploadURL(_ context.Context) (*url.URL, error) {
+	return nil, nil
+}
+
+func (m mockAPIHostResolver) RawURL(_ context.Context) (*url.URL, error) {
+	return nil, nil
+}
+
+func (m mockAPIHostResolver) OAuthURL(_ context.Context) (*url.URL, error) {
+	return url.Parse(m.oauthURL)
+}
+
+// Ensure mockAPIHostResolver implements utils.APIHostResolver
+var _ utils.APIHostResolver = mockAPIHostResolver{}
+
 func TestNewAuthHandler(t *testing.T) {
 	t.Parallel()
 
@@ -51,6 +82,31 @@ func TestNewAuthHandler(t *testing.T) {
 			expectedAuthServer:   DefaultAuthorizationServer,
 			expectedResourcePath: "/mcp",
 		},
+		{
+			name: "APIHost with HTTPS GHES",
+			cfg: &Config{
+				APIHost: mockAPIHostResolver{oauthURL: "https://ghes.example.com/login/oauth"},
+			},
+			expectedAuthServer:   "https://ghes.example.com/login/oauth",
+			expectedResourcePath: "",
+		},
+		{
+			name: "APIHost with HTTP GHES",
+			cfg: &Config{
+				APIHost: mockAPIHostResolver{oauthURL: "http://ghes.local/login/oauth"},
+			},
+			expectedAuthServer:   "http://ghes.local/login/oauth",
+			expectedResourcePath: "",
+		},
+		{
+			name: "APIHost takes precedence over AuthorizationServer",
+			cfg: &Config{
+				APIHost:             mockAPIHostResolver{oauthURL: "https://ghes.example.com/login/oauth"},
+				AuthorizationServer: "https://should-be-ignored.example.com/oauth",
+			},
+			expectedAuthServer:   "https://ghes.example.com/login/oauth",
+			expectedResourcePath: "",
+		},
 	}
 
 	for _, tc := range tests {

pkg/http/server.go

@@ -126,13 +126,7 @@ func RunHTTPServer(cfg ServerConfig) error {
 	oauthCfg := &oauth.Config{
 		BaseURL:      cfg.BaseURL,
 		ResourcePath: cfg.ResourcePath,
-	}
-	if cfg.Host != "" {
-		oauthURL, err := apiHost.OAuthURL(ctx)
-		if err != nil {
-			return fmt.Errorf("failed to get OAuth URL: %w", err)
-		}
-		oauthCfg.AuthorizationServer = oauthURL.String()
+		APIHost:      apiHost,
 	}
 
 	serverOptions := []HandlerOption{}