at-container-registry/pkg/auth/oauth/discovery.go

package oauth

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
)

// ProtectedResourceMetadata represents the OAuth protected resource metadata
// as defined in ATProto OAuth spec
type ProtectedResourceMetadata struct {
	Resource             string   `json:"resource"`
	AuthorizationServers []string `json:"authorization_servers"`
}

// AuthServerMetadata represents the OAuth authorization server metadata
// as defined in RFC 8414
type AuthServerMetadata struct {
	Issuer                                     string   `json:"issuer"`
	AuthorizationEndpoint                      string   `json:"authorization_endpoint"`
	TokenEndpoint                              string   `json:"token_endpoint"`
	PushedAuthorizationRequestEndpoint         string   `json:"pushed_authorization_request_endpoint,omitempty"`
	RegistrationEndpoint                       string   `json:"registration_endpoint,omitempty"`
	JWKsURI                                    string   `json:"jwks_uri,omitempty"`
	ScopesSupported                            []string `json:"scopes_supported,omitempty"`
	ResponseTypesSupported                     []string `json:"response_types_supported,omitempty"`
	GrantTypesSupported                        []string `json:"grant_types_supported,omitempty"`
	TokenEndpointAuthMethodsSupported          []string `json:"token_endpoint_auth_methods_supported,omitempty"`
	DPoPSigningAlgValuesSupported              []string `json:"dpop_signing_alg_values_supported,omitempty"`
	CodeChallengeMethodsSupported              []string `json:"code_challenge_methods_supported,omitempty"`
	AuthorizationResponseIssParameterSupported bool     `json:"authorization_response_iss_parameter_supported,omitempty"`
}

// DiscoverProtectedResource discovers the protected resource metadata
// from the PDS endpoint to find the authorization servers
func DiscoverProtectedResource(ctx context.Context, pdsEndpoint string) (*ProtectedResourceMetadata, error) {
	// Construct the well-known URL per ATProto OAuth spec
	discoveryURL := fmt.Sprintf("%s/.well-known/oauth-protected-resource", pdsEndpoint)

	req, err := http.NewRequestWithContext(ctx, "GET", discoveryURL, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create protected resource discovery request: %w", err)
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch protected resource metadata: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("protected resource discovery failed with status %d", resp.StatusCode)
	}

	var metadata ProtectedResourceMetadata
	if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
		return nil, fmt.Errorf("failed to decode protected resource metadata: %w", err)
	}

	// Validate required fields
	if len(metadata.AuthorizationServers) == 0 {
		return nil, fmt.Errorf("protected resource metadata missing authorization_servers")
	}

	return &metadata, nil
}

// DiscoverAuthServer discovers the OAuth authorization server metadata
// using the ATProto two-step discovery process:
// 1. Fetch protected resource metadata from PDS to get authorization server URL
// 2. Fetch authorization server metadata from that URL
func DiscoverAuthServer(ctx context.Context, pdsEndpoint string) (*AuthServerMetadata, error) {
	// Step 1: Discover the authorization server URL from the protected resource
	protectedResource, err := DiscoverProtectedResource(ctx, pdsEndpoint)
	if err != nil {
		return nil, fmt.Errorf("step 1 failed - discover protected resource from PDS %s: %w", pdsEndpoint, err)
	}

	// Use the first authorization server (ATProto spec allows multiple, but typically one)
	authServerURL := protectedResource.AuthorizationServers[0]

	// Step 2: Fetch authorization server metadata
	discoveryURL := fmt.Sprintf("%s/.well-known/oauth-authorization-server", authServerURL)

	req, err := http.NewRequestWithContext(ctx, "GET", discoveryURL, nil)
	if err != nil {
		return nil, fmt.Errorf("step 2 failed - create request: %w", err)
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("step 2 failed - fetch auth server metadata from %s: %w", discoveryURL, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("step 2 failed - auth server discovery at %s returned status %d", discoveryURL, resp.StatusCode)
	}

	var metadata AuthServerMetadata
	if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
		return nil, fmt.Errorf("failed to decode authorization server metadata: %w", err)
	}

	// Validate required fields
	if metadata.Issuer == "" {
		return nil, fmt.Errorf("authorization server metadata missing issuer")
	}
	if metadata.AuthorizationEndpoint == "" {
		return nil, fmt.Errorf("authorization server metadata missing authorization_endpoint")
	}
	if metadata.TokenEndpoint == "" {
		return nil, fmt.Errorf("authorization server metadata missing token_endpoint")
	}

	return &metadata, nil
}