mcp: handle and pass upstream oauth2 tokens (#5595)

internal/mcp/handler_authorization.go

@@ -8,7 +8,9 @@ import (
 	"net/url"
 	"time"
 
+	"github.com/bufbuild/protovalidate-go"
 	"github.com/go-jose/go-jose/v3/jwt"
+	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -27,19 +29,38 @@ func (srv *Handler) Authorize(w http.ResponseWriter, r *http.Request) {
 
 	ctx := r.Context()
 
-	sessionID, err := getSessionFromRequest(r)
+	claims, err := getClaimsFromRequest(r)
 	if err != nil {
+		log.Ctx(ctx).Error().Err(err).Msg("failed to get claims from request")
+		http.Error(w, "invalid request", http.StatusBadRequest)
+		return
+	}
+	sessionID, ok := getSessionIDFromClaims(claims)
+	if !ok {
 		log.Ctx(ctx).Error().Err(err).Msg("session is not present, this is a misconfigured request")
 		http.Error(w, "internal server error", http.StatusInternalServerError)
 		return
 	}
+	userID, ok := getUserIDFromClaims(claims)
+	if !ok {
+		log.Ctx(ctx).Error().Err(err).Msg("user id is not present, this is a misconfigured request")
+		http.Error(w, "internal server error", http.StatusInternalServerError)
+		return
+	}
 
-	v, err := oauth21.ParseCodeGrantAuthorizeRequest(r, sessionID)
+	v, err := oauth21.ParseCodeGrantAuthorizeRequest(r)
 	if err != nil {
 		log.Ctx(ctx).Error().Err(err).Msg("failed to parse authorization request")
 		oauth21.ErrorResponse(w, http.StatusBadRequest, oauth21.InvalidRequest)
 		return
 	}
+	v.UserId = userID
+	v.SessionId = sessionID
+	if err := protovalidate.Validate(v); err != nil {
+		log.Ctx(ctx).Error().Err(err).Msg("failed to validate authorization request")
+		oauth21.ErrorResponse(w, http.StatusBadRequest, oauth21.InvalidRequest)
+		return
+	}
 
 	client, err := srv.storage.GetClient(ctx, v.ClientId)
 	if err != nil && status.Code(err) == codes.NotFound {
@@ -61,19 +82,51 @@ func (srv *Handler) Authorize(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	id, err := srv.storage.CreateAuthorizationRequest(ctx, v)
-	if err != nil {
-		log.Ctx(ctx).Error().Err(err).Msg("failed to store authorization request")
-		http.Error(w, "cannot create authorization request", http.StatusInternalServerError)
+	requiresUpstreamOAuth2Token := srv.relyingParties.HasConfigForHost(r.Host)
+	var authReqID string
+	var hasUpstreamOAuth2Token bool
+	{
+		eg, ctx := errgroup.WithContext(ctx)
+		eg.Go(func() error {
+			var err error
+			authReqID, err = srv.storage.CreateAuthorizationRequest(ctx, v)
+			if err != nil {
+				return fmt.Errorf("failed to create authorization request: %w", err)
+			}
+			return nil
+		})
+		eg.Go(func() error {
+			if !requiresUpstreamOAuth2Token {
+				return nil
+			}
+
+			var err error
+			token, err := srv.GetUpstreamOAuth2Token(ctx, r.Host, userID)
+			if err != nil && status.Code(err) != codes.NotFound {
+				return fmt.Errorf("failed to get upstream oauth2 token: %w", err)
+			}
+			hasUpstreamOAuth2Token = token != ""
+			return nil
+		})
+
+		err := eg.Wait()
+		if err != nil {
+			log.Ctx(ctx).Error().Err(err).Msg("prepare for authorization redirect")
+			http.Error(w, "internal error", http.StatusInternalServerError)
+			return
+		}
+	}
+
+	if !requiresUpstreamOAuth2Token || hasUpstreamOAuth2Token {
+		srv.AuthorizationResponse(ctx, w, r, authReqID, v)
 		return
 	}
 
-	loginURL, ok := srv.relyingParties.GetLoginURLForHost(r.Host, id)
+	loginURL, ok := srv.relyingParties.GetLoginURLForHost(r.Host, authReqID)
 	if ok {
 		http.Redirect(w, r, loginURL, http.StatusFound)
-	} else {
-		srv.AuthorizationResponse(ctx, w, r, id, v)
 	}
+	log.Ctx(ctx).Error().Msg("authorize: must have login URL, this is a bug")
 }
 
 // AuthorizationResponse generates the successful authorization response
@@ -111,22 +164,31 @@ func (srv *Handler) AuthorizationResponse(
 	http.Redirect(w, r, to.String(), http.StatusFound)
 }
 
-func getSessionFromRequest(r *http.Request) (string, error) {
+func getClaimsFromRequest(r *http.Request) (map[string]any, error) {
 	h := r.Header.Get(httputil.HeaderPomeriumJWTAssertion)
 	if h == "" {
-		return "", fmt.Errorf("missing %s header", httputil.HeaderPomeriumJWTAssertion)
+		return nil, fmt.Errorf("missing %s header", httputil.HeaderPomeriumJWTAssertion)
 	}
 
 	token, err := jwt.ParseSigned(h)
 	if err != nil {
-		return "", fmt.Errorf("failed to parse JWT: %w", err)
+		return nil, fmt.Errorf("failed to parse JWT: %w", err)
 	}
 	var m map[string]any
-	_ = token.UnsafeClaimsWithoutVerification(&m)
-	sessionID, ok := m["sid"].(string)
-	if !ok {
-		return "", fmt.Errorf("missing session ID in JWT")
+	err = token.UnsafeClaimsWithoutVerification(&m)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse JWT claims: %w", err)
 	}
 
-	return sessionID, nil
+	return m, nil
+}
+
+func getSessionIDFromClaims(claims map[string]any) (string, bool) {
+	sessionID, ok := claims["sid"].(string)
+	return sessionID, ok
+}
+
+func getUserIDFromClaims(claims map[string]any) (string, bool) {
+	userID, ok := claims["sub"].(string)
+	return userID, ok
 }