Skip to content

Commit bd3806b

Browse files
committed
Enhance: Support cookie flow for mcp-ui
Signed-off-by: Daishan Peng <[email protected]> wip Signed-off-by: Daishan Peng <[email protected]>
1 parent 95fc784 commit bd3806b

File tree

22 files changed

+1461
-187
lines changed

22 files changed

+1461
-187
lines changed

Makefile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ test-coverage:
3636

3737
# Linting
3838
lint:
39-
golangci-lint run
39+
golangci-lint run --tests=false
4040

4141
# Build targets
4242
build:

cmd/root.go

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
package cmd
2+
3+
import (
4+
"fmt"
5+
"log"
6+
"net/http"
7+
"net/url"
8+
9+
"github.com/gptscript-ai/cmd"
10+
"github.com/obot-platform/mcp-oauth-proxy/pkg/proxy"
11+
"github.com/obot-platform/mcp-oauth-proxy/pkg/types"
12+
"github.com/spf13/cobra"
13+
)
14+
15+
var (
16+
version = "dev"
17+
buildTime = "unknown"
18+
)
19+
20+
// RootCmd represents the base command when called without any subcommands
21+
type RootCmd struct {
22+
// Database configuration
23+
DatabaseDSN string `name:"database-dsn" env:"DATABASE_DSN" usage:"Database connection string (PostgreSQL or SQLite file path). If empty, uses SQLite at data/oauth_proxy.db"`
24+
25+
// OAuth Provider configuration
26+
OAuthClientID string `name:"oauth-client-id" env:"OAUTH_CLIENT_ID" usage:"OAuth client ID from your OAuth provider" required:"true"`
27+
OAuthClientSecret string `name:"oauth-client-secret" env:"OAUTH_CLIENT_SECRET" usage:"OAuth client secret from your OAuth provider" required:"true"`
28+
OAuthAuthorizeURL string `name:"oauth-authorize-url" env:"OAUTH_AUTHORIZE_URL" usage:"Authorization endpoint URL from your OAuth provider (e.g., https://accounts.google.com)" required:"true"`
29+
30+
// Scopes and MCP configuration
31+
ScopesSupported string `name:"scopes-supported" env:"SCOPES_SUPPORTED" usage:"Comma-separated list of supported OAuth scopes (e.g., 'openid,profile,email')" required:"true"`
32+
MCPServerURL string `name:"mcp-server-url" env:"MCP_SERVER_URL" usage:"URL of the MCP server to proxy requests to" required:"true"`
33+
34+
// Security configuration
35+
EncryptionKey string `name:"encryption-key" env:"ENCRYPTION_KEY" usage:"Base64-encoded 32-byte AES-256 key for encrypting sensitive data (optional)"`
36+
37+
// Server configuration
38+
Port string `name:"port" env:"PORT" usage:"Port to run the server on" default:"8080"`
39+
Host string `name:"host" env:"HOST" usage:"Host to bind the server to" default:"localhost"`
40+
41+
// Logging
42+
Verbose bool `name:"verbose,v" usage:"Enable verbose logging"`
43+
Version bool `name:"version" usage:"Show version information"`
44+
45+
Mode string `name:"mode" env:"MODE" usage:"Mode to run the server in" default:"proxy"`
46+
}
47+
48+
func (c *RootCmd) Run(cobraCmd *cobra.Command, args []string) error {
49+
if c.Version {
50+
fmt.Printf("MCP OAuth Proxy\n")
51+
fmt.Printf("Version: %s\n", version)
52+
fmt.Printf("Built: %s\n", buildTime)
53+
return nil
54+
}
55+
56+
// Configure logging
57+
if c.Verbose {
58+
log.SetFlags(log.LstdFlags | log.Lshortfile)
59+
log.Println("Verbose logging enabled")
60+
}
61+
62+
// Convert CLI config to internal config format
63+
config := &types.Config{
64+
DatabaseDSN: c.DatabaseDSN,
65+
OAuthClientID: c.OAuthClientID,
66+
OAuthClientSecret: c.OAuthClientSecret,
67+
OAuthAuthorizeURL: c.OAuthAuthorizeURL,
68+
ScopesSupported: c.ScopesSupported,
69+
MCPServerURL: c.MCPServerURL,
70+
EncryptionKey: c.EncryptionKey,
71+
Mode: c.Mode,
72+
}
73+
74+
// Validate configuration
75+
if err := c.validateConfig(); err != nil {
76+
return fmt.Errorf("configuration validation failed: %w", err)
77+
}
78+
79+
// Create OAuth proxy
80+
oauthProxy, err := proxy.NewOAuthProxy(config)
81+
if err != nil {
82+
return fmt.Errorf("failed to create OAuth proxy: %w", err)
83+
}
84+
defer func() {
85+
if err := oauthProxy.Close(); err != nil {
86+
log.Printf("Error closing database: %v", err)
87+
}
88+
}()
89+
90+
// Get HTTP handler
91+
handler := oauthProxy.GetHandler()
92+
93+
// Start server
94+
address := fmt.Sprintf("%s:%s", c.Host, c.Port)
95+
log.Printf("Starting OAuth proxy server on %s", address)
96+
log.Printf("OAuth Provider: %s", c.OAuthAuthorizeURL)
97+
log.Printf("MCP Server: %s", c.MCPServerURL)
98+
log.Printf("Database: %s", c.getDatabaseType())
99+
100+
return http.ListenAndServe(address, handler)
101+
}
102+
103+
func (c *RootCmd) validateConfig() error {
104+
if c.OAuthClientID == "" {
105+
return fmt.Errorf("oauth-client-id is required")
106+
}
107+
if c.OAuthClientSecret == "" {
108+
return fmt.Errorf("oauth-client-secret is required")
109+
}
110+
if c.OAuthAuthorizeURL == "" {
111+
return fmt.Errorf("oauth-authorize-url is required")
112+
}
113+
if c.ScopesSupported == "" {
114+
return fmt.Errorf("scopes-supported is required")
115+
}
116+
if c.MCPServerURL == "" {
117+
return fmt.Errorf("mcp-server-url is required")
118+
}
119+
if c.Mode == proxy.ModeProxy {
120+
if u, err := url.Parse(c.MCPServerURL); err != nil || u.Scheme != "http" && u.Scheme != "https" {
121+
return fmt.Errorf("invalid MCP server URL: %w", err)
122+
} else if u.Path != "" && u.Path != "/" || u.RawQuery != "" || u.Fragment != "" {
123+
return fmt.Errorf("MCP server URL must not contain a path, query, or fragment")
124+
}
125+
}
126+
return nil
127+
}
128+
129+
func (c *RootCmd) getDatabaseType() string {
130+
if c.DatabaseDSN == "" {
131+
return "SQLite (data/oauth_proxy.db)"
132+
}
133+
if len(c.DatabaseDSN) > 10 && (c.DatabaseDSN[:11] == "postgres://" || c.DatabaseDSN[:14] == "postgresql://") {
134+
return "PostgreSQL"
135+
}
136+
return fmt.Sprintf("SQLite (%s)", c.DatabaseDSN)
137+
}
138+
139+
// Customizer interface implementation for additional command customization
140+
func (c *RootCmd) Customize(cobraCmd *cobra.Command) {
141+
cobraCmd.Use = "mcp-oauth-proxy"
142+
cobraCmd.Short = "OAuth 2.1 proxy server for MCP (Model Context Protocol)"
143+
cobraCmd.Long = `MCP OAuth Proxy is a comprehensive OAuth 2.1 proxy server that provides
144+
OAuth authorization server functionality with PostgreSQL/SQLite storage.
145+
146+
This proxy supports multiple OAuth providers (Google, Microsoft, GitHub) and
147+
proxies requests to MCP servers with user context headers.
148+
149+
Examples:
150+
# Start with environment variables
151+
export OAUTH_CLIENT_ID="your-google-client-id"
152+
export OAUTH_CLIENT_SECRET="your-secret"
153+
export OAUTH_AUTHORIZE_URL="https://accounts.google.com"
154+
export SCOPES_SUPPORTED="openid,profile,email"
155+
export MCP_SERVER_URL="http://localhost:3000"
156+
mcp-oauth-proxy
157+
158+
# Start with CLI flags
159+
mcp-oauth-proxy \
160+
--oauth-client-id="your-google-client-id" \
161+
--oauth-client-secret="your-secret" \
162+
--oauth-authorize-url="https://accounts.google.com" \
163+
--scopes-supported="openid,profile,email" \
164+
--mcp-server-url="http://localhost:3000"
165+
166+
# Use PostgreSQL database
167+
mcp-oauth-proxy \
168+
--database-dsn="postgres://user:pass@localhost:5432/oauth_db?sslmode=disable" \
169+
--oauth-client-id="your-client-id" \
170+
# ... other required flags
171+
172+
Configuration:
173+
Configuration values are loaded in this order (later values override earlier ones):
174+
1. Default values
175+
2. Environment variables
176+
3. Command line flags
177+
178+
Database Support:
179+
- PostgreSQL: Full ACID compliance, recommended for production
180+
- SQLite: Zero configuration, perfect for development and small deployments`
181+
182+
cobraCmd.Version = version
183+
}
184+
185+
// Execute is the main entry point for the CLI
186+
func Execute() error {
187+
rootCmd := &RootCmd{}
188+
cobraCmd := cmd.Command(rootCmd)
189+
return cobraCmd.Execute()
190+
}

go.mod

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,12 @@ module github.com/obot-platform/mcp-oauth-proxy
33
go 1.23.0
44

55
require (
6+
github.com/golang-jwt/jwt/v5 v5.3.0
67
github.com/gorilla/handlers v1.5.2
8+
github.com/gptscript-ai/cmd v0.0.0-20250530150401-bc71fddf8070
9+
github.com/spf13/cobra v1.7.0
710
github.com/stretchr/testify v1.10.0
11+
golang.org/x/oauth2 v0.30.0
812
gorm.io/driver/postgres v1.6.0
913
gorm.io/driver/sqlite v1.6.0
1014
gorm.io/gorm v1.30.1
@@ -13,6 +17,7 @@ require (
1317
require (
1418
github.com/davecgh/go-spew v1.1.1 // indirect
1519
github.com/felixge/httpsnoop v1.0.3 // indirect
20+
github.com/inconshreveable/mousetrap v1.1.0 // indirect
1621
github.com/jackc/pgpassfile v1.0.0 // indirect
1722
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
1823
github.com/jackc/pgx/v5 v5.7.5 // indirect
@@ -23,8 +28,8 @@ require (
2328
github.com/mattn/go-sqlite3 v1.14.32 // indirect
2429
github.com/pmezard/go-difflib v1.0.0 // indirect
2530
github.com/rogpeppe/go-internal v1.8.0 // indirect
31+
github.com/spf13/pflag v1.0.5 // indirect
2632
golang.org/x/crypto v0.41.0 // indirect
27-
golang.org/x/oauth2 v0.30.0 // indirect
2833
golang.org/x/sync v0.16.0 // indirect
2934
golang.org/x/text v0.28.0 // indirect
3035
gopkg.in/yaml.v3 v3.0.1 // indirect

go.sum

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,17 @@
1+
github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
12
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
23
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
34
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
45
github.com/felixge/httpsnoop v1.0.3 h1:s/nj+GCswXYzN5v2DpNMuMQYe+0DDwt5WVCU6CWBdXk=
56
github.com/felixge/httpsnoop v1.0.3/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
7+
github.com/golang-jwt/jwt/v5 v5.3.0 h1:pv4AsKCKKZuqlgs5sUmn4x8UlGa0kEVt/puTpKx9vvo=
8+
github.com/golang-jwt/jwt/v5 v5.3.0/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
69
github.com/gorilla/handlers v1.5.2 h1:cLTUSsNkgcwhgRqvCNmdbRWG0A3N4F+M2nWKdScwyEE=
710
github.com/gorilla/handlers v1.5.2/go.mod h1:dX+xVpaxdSw+q0Qek8SSsl3dfMk3jNddUkMzo0GtH0w=
11+
github.com/gptscript-ai/cmd v0.0.0-20250530150401-bc71fddf8070 h1:xm5ZZFraWFwxyE7TBEncCXArubCDZTwG6s5bpMzqhSY=
12+
github.com/gptscript-ai/cmd v0.0.0-20250530150401-bc71fddf8070/go.mod h1:DJAo1xTht1LDkNYFNydVjTHd576TC7MlpsVRl3oloVw=
13+
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
14+
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
815
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
916
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
1017
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
@@ -30,6 +37,11 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
3037
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
3138
github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8=
3239
github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=
40+
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
41+
github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
42+
github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
43+
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
44+
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
3345
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
3446
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
3547
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=

main.go

Lines changed: 4 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,13 @@
11
package main
22

33
import (
4-
"log"
5-
"net/http"
4+
"os"
65

7-
"github.com/obot-platform/mcp-oauth-proxy/pkg/proxy"
6+
"github.com/obot-platform/mcp-oauth-proxy/cmd"
87
)
98

109
func main() {
11-
// Load configuration from environment variables
12-
config, err := proxy.LoadConfigFromEnv()
13-
if err != nil {
14-
log.Fatalf("Failed to load configuration: %v", err)
10+
if err := cmd.Execute(); err != nil {
11+
os.Exit(1)
1512
}
16-
17-
proxy, err := proxy.NewOAuthProxy(config)
18-
if err != nil {
19-
log.Fatalf("Failed to create OAuth proxy: %v", err)
20-
}
21-
defer func() {
22-
if err := proxy.Close(); err != nil {
23-
log.Printf("Error closing database: %v", err)
24-
}
25-
}()
26-
27-
// Get HTTP handler
28-
handler := proxy.GetHandler()
29-
30-
// Start server
31-
log.Printf("Starting OAuth proxy server on localhost:" + config.Port)
32-
log.Fatal(http.ListenAndServe(":"+config.Port, handler))
3313
}

main_test.go

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"time"
1111

1212
"github.com/obot-platform/mcp-oauth-proxy/pkg/proxy"
13+
"github.com/obot-platform/mcp-oauth-proxy/pkg/types"
1314
"github.com/stretchr/testify/assert"
1415
"github.com/stretchr/testify/require"
1516
)
@@ -62,7 +63,10 @@ func TestIntegrationFlow(t *testing.T) {
6263
}()
6364

6465
// Create OAuth proxy
65-
config, err := proxy.LoadConfigFromEnv()
66+
config := &types.Config{
67+
Mode: proxy.ModeProxy,
68+
}
69+
_, err := proxy.NewOAuthProxy(config)
6670
if err != nil {
6771
log.Fatalf("Failed to load configuration: %v", err)
6872
}
@@ -180,7 +184,10 @@ func TestOAuthProxyCreation(t *testing.T) {
180184
}()
181185

182186
// Create OAuth proxy
183-
config, err := proxy.LoadConfigFromEnv()
187+
config := &types.Config{
188+
Mode: proxy.ModeProxy,
189+
}
190+
_, err := proxy.NewOAuthProxy(config)
184191
if err != nil {
185192
log.Fatalf("Failed to load configuration: %v", err)
186193
}
@@ -233,7 +240,10 @@ func TestOAuthProxyStart(t *testing.T) {
233240
}()
234241

235242
// Create OAuth proxy
236-
config, err := proxy.LoadConfigFromEnv()
243+
config := &types.Config{
244+
Mode: proxy.ModeProxy,
245+
}
246+
_, err := proxy.NewOAuthProxy(config)
237247
if err != nil {
238248
log.Fatalf("Failed to load configuration: %v", err)
239249
}
@@ -296,7 +306,7 @@ func TestForwardAuthIntegrationFlow(t *testing.T) {
296306
"OAUTH_AUTHORIZE_URL": "https://accounts.google.com",
297307
"SCOPES_SUPPORTED": "openid,profile,email",
298308
"PROXY_MODE": "forward_auth",
299-
"PORT": "8082", // Different port to avoid conflicts
309+
"PORT": "8082", // Different port to avoid conflicts
300310
"DATABASE_DSN": os.Getenv("TEST_DATABASE_DSN"), // Use test database if available
301311
}
302312

@@ -320,7 +330,10 @@ func TestForwardAuthIntegrationFlow(t *testing.T) {
320330
}()
321331

322332
// Create OAuth proxy in forward auth mode
323-
config, err := proxy.LoadConfigFromEnv()
333+
config := &types.Config{
334+
Mode: proxy.ModeForwardAuth,
335+
}
336+
_, err := proxy.NewOAuthProxy(config)
324337
if err != nil {
325338
log.Fatalf("Failed to load configuration: %v", err)
326339
}
@@ -375,7 +388,7 @@ func TestForwardAuthIntegrationFlow(t *testing.T) {
375388
// Test that forward auth mode requires authorization for protected endpoints
376389
t.Run("ForwardAuthRequiresAuth", func(t *testing.T) {
377390
testPaths := []string{"/api", "/data", "/protected", "/mcp", "/test"}
378-
391+
379392
for _, path := range testPaths {
380393
t.Run("Path_"+path, func(t *testing.T) {
381394
w := httptest.NewRecorder()
@@ -399,7 +412,7 @@ func TestForwardAuthIntegrationFlow(t *testing.T) {
399412
// Should get unauthorized (no proxying attempt)
400413
assert.Equal(t, http.StatusUnauthorized, w.Code)
401414
assert.Contains(t, w.Header().Get("WWW-Authenticate"), "Bearer")
402-
415+
403416
// Should not have any proxy-related error messages
404417
assert.NotContains(t, w.Body.String(), "proxy")
405418
assert.NotContains(t, w.Body.String(), "502")

pkg/db/db.go

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -295,14 +295,6 @@ func (d *Store) RevokeToken(token string) error {
295295
return result.Error
296296
}
297297

298-
// UpdateTokenRefreshToken updates the refresh token for an existing token
299-
func (d *Store) UpdateTokenRefreshToken(accessToken, newRefreshToken string) error {
300-
hashedAccessToken := hashToken(accessToken)
301-
hashedNewRefreshToken := hashToken(newRefreshToken)
302-
303-
return d.db.Model(&types.TokenData{}).Where("access_token = ?", hashedAccessToken).Update("refresh_token", hashedNewRefreshToken).Error
304-
}
305-
306298
// CleanupExpiredTokens removes expired tokens and authorization codes
307299
func (d *Store) CleanupExpiredTokens() error {
308300
now := time.Now()

0 commit comments

Comments
 (0)