Skip to content

Commit acaf6b4

Browse files
dgellowdgellow
committed
Replace TokenSetup with UserAuthentication config
Support both manual tokens and OAuth flows. Add Secret type for sensitive fields.
1 parent 00a2309 commit acaf6b4

16 files changed

+434
-171
lines changed

‎internal/client/client.go

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,7 @@ func (c *Client) addToolsToServer(
162162
tokenStore storage.UserTokenStore,
163163
serverName string,
164164
setupBaseURL string,
165-
tokenSetup *config.TokenSetupConfig,
165+
userAuth *config.UserAuthentication,
166166
session server.ClientSession,
167167
) error {
168168
toolsRequest := mcp.ListToolsRequest{}
@@ -251,7 +251,7 @@ func (c *Client) addToolsToServer(
251251
userEmail,
252252
serverName,
253253
setupBaseURL,
254-
tokenSetup,
254+
userAuth,
255255
)
256256
} else {
257257
handler = c.client.CallTool
@@ -412,7 +412,7 @@ func (c *Client) wrapToolHandler(
412412
userEmail string,
413413
serverName string,
414414
setupBaseURL string,
415-
tokenSetup *config.TokenSetupConfig,
415+
userAuth *config.UserAuthentication,
416416
) func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
417417
return func(toolCtx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
418418
// Log tool invocation
@@ -435,7 +435,7 @@ func (c *Client) wrapToolHandler(
435435
errorData := createTokenRequiredError(
436436
serverName,
437437
setupBaseURL,
438-
tokenSetup,
438+
userAuth,
439439
"configuration error: this service requires user tokens but OAuth is not properly configured.",
440440
)
441441

@@ -449,14 +449,14 @@ func (c *Client) wrapToolHandler(
449449
tokenSetupURL := fmt.Sprintf("%s/my/tokens", setupBaseURL)
450450

451451
var errorMessage string
452-
if tokenSetup != nil {
452+
if userAuth != nil {
453453
errorMessage = fmt.Sprintf(
454454
"token required: %s requires a user token to access the API. "+
455455
"please visit %s to set up your %s token. %s",
456-
tokenSetup.DisplayName,
456+
userAuth.DisplayName,
457457
tokenSetupURL,
458-
tokenSetup.DisplayName,
459-
tokenSetup.Instructions,
458+
userAuth.DisplayName,
459+
userAuth.Instructions,
460460
)
461461
} else {
462462
errorMessage = fmt.Sprintf(
@@ -469,7 +469,7 @@ func (c *Client) wrapToolHandler(
469469
errorData := createTokenRequiredError(
470470
serverName,
471471
setupBaseURL,
472-
tokenSetup,
472+
userAuth,
473473
errorMessage,
474474
)
475475

@@ -509,7 +509,7 @@ func (c *Client) Close() error {
509509
}
510510

511511
// createTokenRequiredError creates the structured error for missing user tokens
512-
func createTokenRequiredError(serverName, setupBaseURL string, tokenSetup *config.TokenSetupConfig, message string) map[string]interface{} {
512+
func createTokenRequiredError(serverName, setupBaseURL string, userAuth *config.UserAuthentication, message string) map[string]interface{} {
513513
tokenSetupURL := fmt.Sprintf("%s/my/tokens", setupBaseURL)
514514

515515
return map[string]interface{}{

‎internal/client/session_manager.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,7 @@ func (s *StdioSession) DiscoverAndRegisterCapabilities(
251251
tokenStore storage.UserTokenStore,
252252
serverName string,
253253
setupBaseURL string,
254-
tokenSetup *config.TokenSetupConfig,
254+
userAuth *config.UserAuthentication,
255255
session server.ClientSession,
256256
) error {
257257
// Initialize the client
@@ -289,11 +289,11 @@ func (s *StdioSession) DiscoverAndRegisterCapabilities(
289289
"sessionID": session.SessionID(),
290290
"userEmail": userEmail,
291291
"requiresUserToken": requiresToken,
292-
"hasTokenSetup": tokenSetup != nil,
292+
"hasTokenSetup": userAuth != nil,
293293
})
294294

295295
// Discover and register tools
296-
if err := s.client.addToolsToServer(ctx, mcpServer, userEmail, requiresToken, tokenStore, serverName, setupBaseURL, tokenSetup, session); err != nil {
296+
if err := s.client.addToolsToServer(ctx, mcpServer, userEmail, requiresToken, tokenStore, serverName, setupBaseURL, userAuth, session); err != nil {
297297
return err
298298
}
299299

‎internal/config/load.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -204,9 +204,9 @@ func validateMCPServer(name string, server *MCPClientConfig) error {
204204
return fmt.Errorf("server %s has invalid transportType: %s", name, server.TransportType)
205205
}
206206

207-
// Validate token setup if required
208-
if server.RequiresUserToken && server.TokenSetup == nil {
209-
return fmt.Errorf("server %s requires user token but has no tokenSetup", name)
207+
// Validate user authentication if required
208+
if server.RequiresUserToken && server.UserAuthentication == nil {
209+
return fmt.Errorf("server %s requires user token but has no userAuthentication", name)
210210
}
211211

212212
return nil

‎internal/config/load_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,8 @@ func TestValidateConfig_UserTokensRequireOAuth(t *testing.T) {
2525
TransportType: MCPClientTypeSSE,
2626
URL: "https://notion.example.com",
2727
RequiresUserToken: true,
28-
TokenSetup: &TokenSetupConfig{
28+
UserAuthentication: &UserAuthentication{
29+
Type: UserAuthTypeManual,
2930
DisplayName: "Notion",
3031
},
3132
},
@@ -56,7 +57,8 @@ func TestValidateConfig_UserTokensRequireOAuth(t *testing.T) {
5657
TransportType: MCPClientTypeSSE,
5758
URL: "https://notion.example.com",
5859
RequiresUserToken: true,
59-
TokenSetup: &TokenSetupConfig{
60+
UserAuthentication: &UserAuthentication{
61+
Type: UserAuthTypeManual,
6062
DisplayName: "Notion",
6163
},
6264
},

‎internal/config/secret_test.go

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
package config
2+
3+
import (
4+
"encoding/json"
5+
"fmt"
6+
"strings"
7+
"testing"
8+
)
9+
10+
func TestSecretRedaction(t *testing.T) {
11+
tests := []struct {
12+
name string
13+
secret Secret
14+
want string
15+
}{
16+
{
17+
name: "non-empty secret",
18+
secret: Secret("super-secret-password"),
19+
want: "***",
20+
},
21+
{
22+
name: "empty secret",
23+
secret: Secret(""),
24+
want: "",
25+
},
26+
}
27+
28+
for _, tt := range tests {
29+
t.Run(tt.name, func(t *testing.T) {
30+
// Test String() method
31+
if got := tt.secret.String(); got != tt.want {
32+
t.Errorf("Secret.String() = %v, want %v", got, tt.want)
33+
}
34+
35+
// Test fmt.Sprintf behavior
36+
formatted := fmt.Sprintf("value: %s", tt.secret)
37+
expectedFormatted := "value: " + tt.want
38+
if formatted != expectedFormatted {
39+
t.Errorf("fmt.Sprintf = %v, want %v", formatted, expectedFormatted)
40+
}
41+
42+
// Test fmt.Printf (capture output)
43+
output := fmt.Sprintf("password: %v", tt.secret)
44+
if tt.secret != "" && strings.Contains(output, string(tt.secret)) {
45+
t.Errorf("fmt.Printf leaked secret: %v", output)
46+
}
47+
})
48+
}
49+
}
50+
51+
func TestSecretJSONMarshal(t *testing.T) {
52+
type configWithSecrets struct {
53+
Username string `json:"username"`
54+
Password Secret `json:"password"`
55+
APIKey Secret `json:"apiKey"`
56+
}
57+
58+
cfg := configWithSecrets{
59+
Username: "admin",
60+
Password: Secret("super-secret-password"),
61+
APIKey: Secret("sk-1234567890abcdef"),
62+
}
63+
64+
data, err := json.Marshal(cfg)
65+
if err != nil {
66+
t.Fatalf("json.Marshal failed: %v", err)
67+
}
68+
69+
jsonStr := string(data)
70+
71+
// Check that secrets are redacted
72+
if strings.Contains(jsonStr, "super-secret-password") {
73+
t.Errorf("JSON contains unredacted password: %s", jsonStr)
74+
}
75+
if strings.Contains(jsonStr, "sk-1234567890abcdef") {
76+
t.Errorf("JSON contains unredacted API key: %s", jsonStr)
77+
}
78+
79+
// Check that username is not redacted
80+
if !strings.Contains(jsonStr, "admin") {
81+
t.Errorf("JSON doesn't contain username: %s", jsonStr)
82+
}
83+
84+
// Check expected JSON structure
85+
expected := `{"username":"admin","password":"***","apiKey":"***"}`
86+
if jsonStr != expected {
87+
t.Errorf("JSON = %s, want %s", jsonStr, expected)
88+
}
89+
}
90+
91+
func TestSecretInStruct(t *testing.T) {
92+
auth := ServiceAuth{
93+
Type: ServiceAuthTypeBasic,
94+
Username: "testuser",
95+
HashedPassword: Secret("$2a$10$abcdef..."),
96+
UserToken: Secret("token-12345"),
97+
}
98+
99+
// Test struct string representation
100+
str := fmt.Sprintf("%+v", auth)
101+
if strings.Contains(str, "$2a$10$abcdef") {
102+
t.Errorf("Struct representation leaked hashed password: %s", str)
103+
}
104+
if strings.Contains(str, "token-12345") {
105+
t.Errorf("Struct representation leaked token: %s", str)
106+
}
107+
108+
// Individual field access should still redact
109+
if auth.HashedPassword.String() != "***" {
110+
t.Errorf("HashedPassword.String() = %v, want ***", auth.HashedPassword.String())
111+
}
112+
}

‎internal/config/types.go

Lines changed: 71 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,25 @@ import (
88
"time"
99
)
1010

11+
// Secret is a string type that redacts itself when printed
12+
type Secret string
13+
14+
// String implements fmt.Stringer to redact the secret
15+
func (s Secret) String() string {
16+
if s == "" {
17+
return ""
18+
}
19+
return "***"
20+
}
21+
22+
// MarshalJSON implements json.Marshaler to prevent secrets in JSON logs
23+
func (s Secret) MarshalJSON() ([]byte, error) {
24+
if s == "" {
25+
return json.Marshal("")
26+
}
27+
return json.Marshal("***")
28+
}
29+
1130
// MCPClientType represents the transport type for MCP clients
1231
type MCPClientType string
1332

@@ -46,15 +65,6 @@ type Options struct {
4665
ToolFilter *ToolFilterConfig `json:"toolFilter,omitempty"`
4766
}
4867

49-
// TokenSetupConfig provides information for users to set up their tokens
50-
type TokenSetupConfig struct {
51-
DisplayName string `json:"displayName"`
52-
Instructions string `json:"instructions"`
53-
HelpURL string `json:"helpUrl,omitempty"`
54-
TokenFormat string `json:"tokenFormat,omitempty"`
55-
CompiledRegex *regexp.Regexp `json:"-"`
56-
}
57-
5868
// ServiceAuthType represents the type of service authentication
5969
type ServiceAuthType string
6070

@@ -63,23 +73,65 @@ const (
6373
ServiceAuthTypeBasic ServiceAuthType = "basic"
6474
)
6575

76+
// UserAuthType represents the type of user authentication
77+
type UserAuthType string
78+
79+
const (
80+
// UserAuthTypeManual indicates that users manually provide API tokens/keys
81+
// through the web UI. These tokens are stored encrypted and injected into
82+
// MCP servers as configured.
83+
UserAuthTypeManual UserAuthType = "manual"
84+
85+
// UserAuthTypeOAuth indicates OAuth 2.0 authorization code flow is used.
86+
// Users click "Connect with X" and are redirected to the service's OAuth
87+
// consent page. The resulting access tokens are stored, automatically
88+
// refreshed, and injected into MCP servers.
89+
UserAuthTypeOAuth UserAuthType = "oauth"
90+
)
91+
6692
// ServiceAuth represents authentication method for service-to-service communication
6793
type ServiceAuth struct {
6894
Type ServiceAuthType `json:"type"`
6995

7096
// For basic auth
71-
Username string `json:"username,omitempty"`
72-
Password json.RawMessage `json:"password,omitempty"`
97+
Username string `json:"username,omitempty"`
98+
PasswordRaw json.RawMessage `json:"password,omitempty"`
7399

74100
// For bearer auth
75101
Tokens []string `json:"tokens,omitempty"`
76102

77103
// User token to inject when requiresUserToken is true
78-
UserToken json.RawMessage `json:"userToken,omitempty"`
104+
UserTokenRaw json.RawMessage `json:"userToken,omitempty"`
105+
106+
// Computed fields
107+
HashedPassword Secret `json:"-"` // bcrypt hash for basic auth
108+
UserToken Secret `json:"-"` // parsed user token
109+
}
110+
111+
// UserAuthentication represents authentication configuration for end users
112+
type UserAuthentication struct {
113+
Type UserAuthType `json:"type"`
114+
DisplayName string `json:"displayName"`
115+
116+
// For OAuth
117+
ClientIDRaw json.RawMessage `json:"clientId,omitempty"`
118+
ClientSecretRaw json.RawMessage `json:"clientSecret,omitempty"`
119+
AuthorizationURL string `json:"authorizationUrl,omitempty"`
120+
TokenURL string `json:"tokenUrl,omitempty"`
121+
Scopes []string `json:"scopes,omitempty"`
122+
123+
// For Manual
124+
Instructions string `json:"instructions,omitempty"`
125+
HelpURL string `json:"helpUrl,omitempty"`
126+
Validation string `json:"validation,omitempty"`
127+
128+
// Common
129+
TokenFormat string `json:"tokenFormat,omitempty"`
79130

80131
// Computed fields
81-
HashedPassword string `json:"-"` // bcrypt hash for basic auth
82-
ResolvedUserToken string `json:"-"` // resolved user token
132+
ClientID Secret `json:"-"`
133+
ClientSecret Secret `json:"-"`
134+
ValidationRegex *regexp.Regexp `json:"-"`
83135
}
84136

85137
// MCPClientConfig represents the configuration for an MCP client after parsing.
@@ -125,8 +177,8 @@ type MCPClientConfig struct {
125177
Options *Options `json:"options,omitempty"`
126178

127179
// User token requirements
128-
RequiresUserToken bool `json:"requiresUserToken,omitempty"`
129-
TokenSetup *TokenSetupConfig `json:"tokenSetup,omitempty"`
180+
RequiresUserToken bool `json:"requiresUserToken,omitempty"`
181+
UserAuthentication *UserAuthentication `json:"userAuthentication,omitempty"`
130182

131183
// Service-to-service authentication
132184
ServiceAuths []ServiceAuth `json:"serviceAuths,omitempty"`
@@ -160,10 +212,10 @@ type OAuthAuthConfig struct {
160212
FirestoreDatabase string `json:"firestoreDatabase,omitempty"` // Optional: Firestore database name
161213
FirestoreCollection string `json:"firestoreCollection,omitempty"` // Optional: Firestore collection name
162214
GoogleClientID string `json:"googleClientId"`
163-
GoogleClientSecret string `json:"googleClientSecret"`
215+
GoogleClientSecret Secret `json:"googleClientSecret"`
164216
GoogleRedirectURI string `json:"googleRedirectUri"`
165-
JWTSecret string `json:"jwtSecret"`
166-
EncryptionKey string `json:"encryptionKey"`
217+
JWTSecret Secret `json:"jwtSecret"`
218+
EncryptionKey Secret `json:"encryptionKey"`
167219
}
168220

169221
// ProxyConfig represents the proxy configuration with resolved values

0 commit comments

Comments
 (0)