WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/mcp/oauth_manager.go

@@ -0,0 +1,280 @@
+package mcp
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	secutils "github.com/Tencent/WeKnora/internal/utils"
+	"github.com/mark3labs/mcp-go/client/transport"
+	"github.com/redis/go-redis/v9"
+)
+
+// clientRegistrationName is sent as client_name during dynamic client
+// registration (RFC 7591).
+const clientRegistrationName = "WeKnora"
+
+// oauthCallbackTimeout bounds token exchange after the browser lands on the
+// public callback route. The Gin request context is canceled once the client
+// receives the redirect, so CompleteAuthorization must detach from it.
+const oauthCallbackTimeout = 60 * time.Second
+
+// OAuthManager orchestrates the MCP OAuth2 authorization-code flow:
+// discovery, dynamic client registration, the authorize redirect, and the
+// callback code exchange. Tokens are persisted per (tenant, principal, service);
+// the registered client is persisted per (tenant, service) and reused.
+type OAuthManager struct {
+	repo        interfaces.MCPOAuthRepository
+	serviceRepo interfaces.MCPServiceRepository
+	states      *oauthStateStore
+}
+
+// NewOAuthManager constructs the OAuth manager. rdb may be nil (Lite mode),
+// in which case in-flight authorization states are kept in memory.
+func NewOAuthManager(
+	repo interfaces.MCPOAuthRepository,
+	serviceRepo interfaces.MCPServiceRepository,
+	rdb *redis.Client,
+) *OAuthManager {
+	return &OAuthManager{
+		repo:        repo,
+		serviceRepo: serviceRepo,
+		states:      newOAuthStateStore(rdb),
+	}
+}
+
+// newHandler builds an OAuth handler bound to a service + per-principal token store.
+func (m *OAuthManager) newHandler(
+	ctx context.Context, service *types.MCPService, tenantID uint64, principal types.Principal, redirectURI string,
+) (*transport.OAuthHandler, error) {
+	if service.URL == nil || *service.URL == "" {
+		return nil, fmt.Errorf("MCP service URL is required for OAuth")
+	}
+	if err := ValidateServiceOutboundURLs(service); err != nil {
+		return nil, err
+	}
+	httpCfg := secutils.DefaultSSRFSafeHTTPClientConfig()
+	httpCfg.Timeout = 30 * time.Second
+	cfg := transport.OAuthConfig{
+		RedirectURI:           redirectURI,
+		Scopes:                service.AuthConfig.Scopes,
+		TokenStore:            newDBTokenStore(m.repo, tenantID, principal, service.ID),
+		PKCEEnabled:           true,
+		AuthServerMetadataURL: service.AuthConfig.AuthServerMetadataURL,
+		HTTPClient:            secutils.NewSSRFSafeHTTPClient(httpCfg),
+	}
+	if existing, err := m.repo.GetClient(ctx, tenantID, service.ID); err == nil && existing != nil {
+		cfg.ClientID = existing.ClientID
+		cfg.ClientSecret = existing.ClientSecret
+	}
+	h := transport.NewOAuthHandler(cfg)
+	h.SetBaseURL(*service.URL)
+	return h, nil
+}
+
+// StartAuthorization performs discovery + (one-time) dynamic client
+// registration, then returns the authorization URL and an opaque attempt ID.
+// redirectURI is the backend callback URL registered with the auth server;
+// frontendRedirect is where the callback bounces the browser when finished.
+func (m *OAuthManager) StartAuthorization(
+	ctx context.Context,
+	service *types.MCPService,
+	tenantID uint64,
+	principal types.Principal,
+	redirectURI, frontendRedirect string,
+) (authorizationURL, attemptID string, err error) {
+	if !service.AuthConfig.IsOAuth() {
+		return "", "", fmt.Errorf("MCP service %s does not use OAuth", service.ID)
+	}
+	principal = principal.Normalize()
+	if !principal.Valid() {
+		return "", "", fmt.Errorf("principal context is required to authorize OAuth MCP service %s", service.ID)
+	}
+
+	h, err := m.newHandler(ctx, service, tenantID, principal, redirectURI)
+	if err != nil {
+		return "", "", err
+	}
+
+	// Register a client dynamically if we don't have one yet for this service.
+	existing, _ := m.repo.GetClient(ctx, tenantID, service.ID)
+	if existing == nil {
+		if err := h.RegisterClient(ctx, clientRegistrationName); err != nil {
+			return "", "", fmt.Errorf("dynamic client registration failed: %w", err)
+		}
+		clientID := h.GetClientID()
+		if clientID == "" {
+			return "", "", fmt.Errorf("dynamic client registration returned an empty client_id")
+		}
+		if err := m.repo.SaveClient(ctx, &types.MCPOAuthClient{
+			TenantID:    tenantID,
+			ServiceID:   service.ID,
+			ClientID:    clientID,
+			RedirectURI: redirectURI,
+		}); err != nil {
+			logger.GetLogger(ctx).Warnf("failed to persist MCP oauth client: %v", err)
+		}
+	}
+
+	verifier, err := transport.GenerateCodeVerifier()
+	if err != nil {
+		return "", "", fmt.Errorf("failed to generate PKCE verifier: %w", err)
+	}
+	challenge := transport.GenerateCodeChallenge(verifier)
+	state, err := transport.GenerateState()
+	if err != nil {
+		return "", "", fmt.Errorf("failed to generate state: %w", err)
+	}
+
+	authURL, err := h.GetAuthorizationURL(ctx, state, challenge)
+	if err != nil {
+		return "", "", fmt.Errorf("failed to build authorization URL: %w", err)
+	}
+
+	if err := m.states.Put(ctx, state, OAuthState{
+		TenantID:         tenantID,
+		UserID:           principal.StorageID(),
+		Principal:        principal,
+		ServiceID:        service.ID,
+		CodeVerifier:     verifier,
+		ClientID:         h.GetClientID(),
+		RedirectURI:      redirectURI,
+		FrontendRedirect: frontendRedirect,
+	}); err != nil {
+		return "", "", fmt.Errorf("failed to persist authorization state: %w", err)
+	}
+
+	return authURL, state, nil
+}
+
+// StartAuthorizationForService loads the MCP service by ID and starts the
+// authorization-code flow, returning the URL the user must open. It is a
+// convenience for callers (e.g. IM channels) that only hold a service ID and
+// cannot reach the MCP service lookup directly.
+func (m *OAuthManager) StartAuthorizationForService(
+	ctx context.Context,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, redirectURI, frontendRedirect string,
+) (string, error) {
+	service, err := m.serviceRepo.GetByID(ctx, tenantID, serviceID)
+	if err != nil {
+		return "", fmt.Errorf("failed to load MCP service: %w", err)
+	}
+	if service == nil {
+		return "", fmt.Errorf("MCP service not found")
+	}
+	authURL, _, err := m.StartAuthorization(ctx, service, tenantID, principal, redirectURI, frontendRedirect)
+	return authURL, err
+}
+
+// CompleteAuthorization handles the provider callback: it validates state,
+// exchanges the code for tokens (PKCE), and persists the per-user token.
+// Returns the frontend redirect URL and service ID recorded at
+// StartAuthorization time so the caller can recycle any cached transport that
+// still carries the previous OAuth client registration.
+func (m *OAuthManager) CompleteAuthorization(
+	ctx context.Context, state, code string,
+) (frontendRedirect, serviceID string, err error) {
+	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), oauthCallbackTimeout)
+	defer cancel()
+
+	st, err := m.states.Take(ctx, state)
+	if err != nil {
+		return "", "", err
+	}
+	frontendRedirect = st.FrontendRedirect
+	serviceID = st.ServiceID
+	principal := st.Principal.Normalize()
+	if !principal.Valid() && st.UserID != "" {
+		principal = types.Principal{Type: types.PrincipalWebUser, ID: st.UserID}.Normalize()
+	}
+	if !principal.Valid() {
+		return frontendRedirect, serviceID, fmt.Errorf("principal context is missing from OAuth state")
+	}
+
+	service, err := m.serviceRepo.GetByID(ctx, st.TenantID, st.ServiceID)
+	if err != nil {
+		return frontendRedirect, serviceID, fmt.Errorf("failed to load MCP service: %w", err)
+	}
+	if service == nil {
+		return frontendRedirect, serviceID, fmt.Errorf("MCP service not found")
+	}
+
+	h, err := m.newHandler(ctx, service, st.TenantID, principal, st.RedirectURI)
+	if err != nil {
+		return frontendRedirect, serviceID, err
+	}
+	// Re-prime the expected state so the library's CSRF check passes after
+	// reconstructing the handler in this separate request.
+	h.SetExpectedState(state)
+
+	if err := h.ProcessAuthorizationResponse(ctx, code, state, st.CodeVerifier); err != nil {
+		return frontendRedirect, serviceID, fmt.Errorf("token exchange failed: %w", err)
+	}
+	if err := m.states.CompleteAttempt(ctx, state); err != nil {
+		return frontendRedirect, serviceID, fmt.Errorf("failed to record authorization completion: %w", err)
+	}
+	// ProcessAuthorizationResponse persists the token via the TokenStore.
+	logger.GetLogger(ctx).Infof(
+		"MCP OAuth authorized: service=%s principal=%s", st.ServiceID, principal.StorageID(),
+	)
+	return frontendRedirect, serviceID, nil
+}
+
+// IsAuthorizationAttemptComplete reports whether this exact authorization
+// attempt completed for the requested principal and service. A pre-existing
+// token must never satisfy a newly opened OAuth popup.
+func (m *OAuthManager) IsAuthorizationAttemptComplete(
+	ctx context.Context,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, attemptID string,
+) (bool, error) {
+	attempt, err := m.states.Attempt(ctx, attemptID)
+	if err != nil {
+		return false, err
+	}
+	principal = principal.Normalize()
+	attemptPrincipal := attempt.Principal.Normalize()
+	if attempt.TenantID != tenantID || attempt.ServiceID != serviceID ||
+		attemptPrincipal.Type != principal.Type || attemptPrincipal.ID != principal.ID {
+		return false, fmt.Errorf("oauth authorization attempt does not match the current principal or service")
+	}
+	return attempt.Completed, nil
+}
+
+// AuthorizationStatus reports whether the stored access token is usable now,
+// or is expired but still has a refresh token that runtime use can rotate.
+func (m *OAuthManager) AuthorizationStatus(
+	ctx context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) (OAuthAuthorizationStatus, error) {
+	tok, err := m.repo.GetTokenForPrincipal(ctx, tenantID, principal, serviceID)
+	if err != nil {
+		return OAuthAuthorizationStatus{}, err
+	}
+	return tokenStatus(tok, time.Now()), nil
+}
+
+// IsAuthorized reports whether the given principal has an access token that is
+// usable now. An expired row is not authorization success merely because its
+// encrypted token columns remain non-empty.
+func (m *OAuthManager) IsAuthorized(
+	ctx context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) (bool, error) {
+	status, err := m.AuthorizationStatus(ctx, tenantID, principal, serviceID)
+	if err != nil {
+		return false, err
+	}
+	return status.Authorized, nil
+}
+
+// Revoke removes the principal's stored token for the service.
+func (m *OAuthManager) Revoke(
+	ctx context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) error {
+	return m.repo.DeleteTokenForPrincipal(ctx, tenantID, principal, serviceID)
+}