Skip to content

Commit d66b79b

Browse files
joelverhagenjoelverhagen
committed
Address comments
1 parent cfb4717 commit d66b79b

File tree

3 files changed

+72
-64
lines changed

3 files changed

+72
-64
lines changed

‎internal/api/handlers/v0/auth/common.go‎

Lines changed: 50 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,16 @@ import (
1414
"github.com/modelcontextprotocol/registry/internal/config"
1515
)
1616

17-
// CoreTokenExchangeInput represents the common input structure for token exchange
18-
type CoreTokenExchangeInput struct {
17+
// SignatureTokenExchangeInput represents the common input structure for token exchange
18+
type SignatureTokenExchangeInput struct {
1919
Domain string `json:"domain" doc:"Domain name" example:"example.com" required:"true"`
2020
Timestamp string `json:"timestamp" doc:"RFC3339 timestamp" example:"2023-01-01T00:00:00Z" required:"true"`
2121
SignedTimestamp string `json:"signed_timestamp" doc:"Hex-encoded Ed25519 signature of timestamp" example:"abcdef1234567890" required:"true"`
2222
}
2323

24+
// KeyFetcher defines a function type for fetching keys from external sources
25+
type KeyFetcher func(ctx context.Context, domain string) ([]string, error)
26+
2427
// CoreAuthHandler represents the common handler structure
2528
type CoreAuthHandler struct {
2629
config *config.Config
@@ -90,9 +93,6 @@ func BuildPermissions(domain string, includeSubdomains bool) []auth.Permission {
9093
}
9194

9295
if includeSubdomains {
93-
// DNS implies a hierarchy where subdomains are treated as part of the parent domain,
94-
// therefore we grant permissions for all subdomains (e.g., com.example.*)
95-
// This is in line with other DNS-based authentication methods e.g. ACME DNS-01 challenges
9696
permissions = append(permissions, auth.Permission{
9797
Action: auth.PermissionActionPublish,
9898
ResourcePattern: fmt.Sprintf("%s.*", reverseDomain),
@@ -120,6 +120,51 @@ func (h *CoreAuthHandler) CreateJWTClaimsAndToken(ctx context.Context, authMetho
120120
return tokenResponse, nil
121121
}
122122

123+
// ExchangeToken is a shared method for token exchange that takes a key fetcher function,
124+
// subdomain inclusion flag, and auth method
125+
func (h *CoreAuthHandler) ExchangeToken(
126+
ctx context.Context,
127+
domain, timestamp, signedTimestamp string,
128+
keyFetcher KeyFetcher,
129+
includeSubdomains bool,
130+
authMethod auth.Method) (*auth.TokenResponse, error) {
131+
_, err := ValidateDomainAndTimestamp(domain, timestamp)
132+
if err != nil {
133+
return nil, err
134+
}
135+
136+
signature, err := DecodeAndValidateSignature(signedTimestamp)
137+
if err != nil {
138+
return nil, err
139+
}
140+
141+
keyStrings, err := keyFetcher(ctx, domain)
142+
if err != nil {
143+
return nil, fmt.Errorf("failed to fetch keys: %w", err)
144+
}
145+
146+
publicKeys := ParseMCPKeysFromStrings(keyStrings)
147+
if len(publicKeys) == 0 {
148+
switch authMethod {
149+
case auth.MethodHTTP:
150+
return nil, fmt.Errorf("failed to parse public key")
151+
case auth.MethodDNS:
152+
return nil, fmt.Errorf("no valid MCP public keys found in DNS TXT records")
153+
default:
154+
return nil, fmt.Errorf("no valid MCP public keys found using %s authentication", authMethod)
155+
}
156+
}
157+
158+
messageBytes := []byte(timestamp)
159+
if !VerifySignatureWithKeys(publicKeys, messageBytes, signature) {
160+
return nil, fmt.Errorf("signature verification failed")
161+
}
162+
163+
permissions := BuildPermissions(domain, includeSubdomains)
164+
165+
return h.CreateJWTClaimsAndToken(ctx, authMethod, domain, permissions)
166+
}
167+
123168
func ParseMCPKeysFromStrings(inputs []string) []ed25519.PublicKey {
124169
var publicKeys []ed25519.PublicKey
125170
mcpPattern := regexp.MustCompile(`v=MCPv1;\s*k=ed25519;\s*p=([A-Za-z0-9+/=]+)`)

‎internal/api/handlers/v0/auth/dns.go‎

Lines changed: 13 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import (
1414

1515
// DNSTokenExchangeInput represents the input for DNS-based authentication
1616
type DNSTokenExchangeInput struct {
17-
Body CoreTokenExchangeInput
17+
Body SignatureTokenExchangeInput
1818
}
1919

2020
// DNSResolver defines the interface for DNS resolution
@@ -75,35 +75,18 @@ func RegisterDNSEndpoint(api huma.API, cfg *config.Config) {
7575

7676
// ExchangeToken exchanges DNS signature for a Registry JWT token
7777
func (h *DNSAuthHandler) ExchangeToken(ctx context.Context, domain, timestamp, signedTimestamp string) (*auth.TokenResponse, error) {
78-
_, err := ValidateDomainAndTimestamp(domain, timestamp)
79-
if err != nil {
80-
return nil, err
81-
}
82-
83-
signature, err := DecodeAndValidateSignature(signedTimestamp)
84-
if err != nil {
85-
return nil, err
86-
}
87-
88-
// Lookup DNS TXT records
89-
txtRecords, err := h.resolver.LookupTXT(ctx, domain)
90-
if err != nil {
91-
return nil, fmt.Errorf("failed to lookup DNS TXT records: %w", err)
92-
}
93-
94-
publicKeys := ParseMCPKeysFromStrings(txtRecords)
95-
96-
if len(publicKeys) == 0 {
97-
return nil, fmt.Errorf("no valid MCP public keys found in DNS TXT records")
98-
}
99-
100-
messageBytes := []byte(timestamp)
101-
if !VerifySignatureWithKeys(publicKeys, messageBytes, signature) {
102-
return nil, fmt.Errorf("signature verification failed")
78+
keyFetcher := func(ctx context.Context, domain string) ([]string, error) {
79+
// Lookup DNS TXT records
80+
// DNS implies a hierarchy where subdomains are treated as part of the parent domain,
81+
// therefore we grant permissions for all subdomains (e.g., com.example.*)
82+
// This is in line with other DNS-based authentication methods e.g. ACME DNS-01 challenges
83+
txtRecords, err := h.resolver.LookupTXT(ctx, domain)
84+
if err != nil {
85+
return nil, fmt.Errorf("failed to lookup DNS TXT records: %w", err)
86+
}
87+
return txtRecords, nil
10388
}
10489

105-
// Build permissions for domain (DNS includes subdomains)
106-
permissions := BuildPermissions(domain, true)
107-
108-
return h.CreateJWTClaimsAndToken(ctx, auth.MethodDNS, domain, permissions)
90+
allowSubdomains := true
91+
return h.CoreAuthHandler.ExchangeToken(ctx, domain, timestamp, signedTimestamp, keyFetcher, allowSubdomains, auth.MethodDNS)
10992
}

‎internal/api/handlers/v0/auth/http.go‎

Lines changed: 9 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ const MaxKeyResponseSize = 4096
1919

2020
// HTTPTokenExchangeInput represents the input for HTTP-based authentication
2121
type HTTPTokenExchangeInput struct {
22-
Body CoreTokenExchangeInput
22+
Body SignatureTokenExchangeInput
2323
}
2424

2525
// HTTPKeyFetcher defines the interface for fetching HTTP keys
@@ -133,34 +133,14 @@ func RegisterHTTPEndpoint(api huma.API, cfg *config.Config) {
133133

134134
// ExchangeToken exchanges HTTP signature for a Registry JWT token
135135
func (h *HTTPAuthHandler) ExchangeToken(ctx context.Context, domain, timestamp, signedTimestamp string) (*auth.TokenResponse, error) {
136-
_, err := ValidateDomainAndTimestamp(domain, timestamp)
137-
if err != nil {
138-
return nil, err
139-
}
140-
141-
signature, err := DecodeAndValidateSignature(signedTimestamp)
142-
if err != nil {
143-
return nil, err
144-
}
145-
146-
// Fetch the key from well known HTTP endpoint
147-
keyResponse, err := h.fetcher.FetchKey(ctx, domain)
148-
if err != nil {
149-
return nil, fmt.Errorf("failed to fetch public key: %w", err)
150-
}
151-
152-
publicKeys := ParseMCPKeysFromStrings([]string{keyResponse})
153-
if len(publicKeys) == 0 {
154-
return nil, fmt.Errorf("failed to parse public key")
155-
}
156-
157-
messageBytes := []byte(timestamp)
158-
if !VerifySignatureWithKeys(publicKeys, messageBytes, signature) {
159-
return nil, fmt.Errorf("signature verification failed")
136+
keyFetcher := func(ctx context.Context, domain string) ([]string, error) {
137+
keyResponse, err := h.fetcher.FetchKey(ctx, domain)
138+
if err != nil {
139+
return nil, fmt.Errorf("failed to fetch public key: %w", err)
140+
}
141+
return []string{keyResponse}, nil
160142
}
161143

162-
// Build permissions for domain (HTTP does not include subdomains)
163-
permissions := BuildPermissions(domain, false)
164-
165-
return h.CreateJWTClaimsAndToken(ctx, auth.MethodHTTP, domain, permissions)
144+
allowSubdomains := false
145+
return h.CoreAuthHandler.ExchangeToken(ctx, domain, timestamp, signedTimestamp, keyFetcher, allowSubdomains, auth.MethodHTTP)
166146
}

0 commit comments

Comments
 (0)