From afbef23a5137cfd26a5986b990536f213fb55d0a Mon Sep 17 00:00:00 2001
From: Matt Moyer <moyerm@vmware.com>
Date: Tue, 8 Dec 2020 10:17:03 -0800
Subject: [PATCH] WIP implementing TokenExchangeHandler methods

Signed-off-by: Margo Crawford <margaretc@vmware.com>
---
 internal/oidc/oidc.go           |  3 ++-
 internal/oidc/token_exchange.go | 41 ++++++++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/internal/oidc/oidc.go b/internal/oidc/oidc.go
index 81c7d54c0..607abff37 100644
--- a/internal/oidc/oidc.go
+++ b/internal/oidc/oidc.go
@@ -83,7 +83,7 @@ func PinnipedCLIOIDCClient() *fosite.DefaultOpenIDConnectClient {
 			Public:        true,
 			RedirectURIs:  []string{"http://127.0.0.1/callback"},
 			ResponseTypes: []string{"code"},
-			GrantTypes:    []string{"authorization_code", "token_exchange"},
+			GrantTypes:    []string{"authorization_code", "urn:ietf:params:oauth:grant-type:token-exchange"},
 			Scopes:        []string{"openid", "profile", "email"},
 		},
 		TokenEndpointAuthMethod: "none",
@@ -124,6 +124,7 @@ func FositeOauth2Helper(
 			OpenIDConnectTokenStrategy: newDynamicOpenIDConnectECDSAStrategy(oauthConfig, jwksProvider),
 		},
 		nil, // hasher, defaults to using BCrypt when nil. Used for hashing client secrets.
+		TokenExchangeFactory,
 		compose.OAuth2AuthorizeExplicitFactory,
 		// compose.OAuth2RefreshTokenGrantFactory,
 		compose.OpenIDConnectExplicitFactory,
diff --git a/internal/oidc/token_exchange.go b/internal/oidc/token_exchange.go
index 2f3f1ee15..c4fbc7088 100644
--- a/internal/oidc/token_exchange.go
+++ b/internal/oidc/token_exchange.go
@@ -3,21 +3,60 @@ package oidc
 import (
 	"context"
 
+	"github.com/ory/fosite/handler/oauth2"
+
+	"github.com/ory/fosite/handler/openid"
+
+	"github.com/pkg/errors"
+
 	"github.com/ory/fosite"
 	"github.com/ory/fosite/compose"
 )
 
 func TokenExchangeFactory(config *compose.Config, storage interface{}, strategy interface{}) interface{} {
-	return &TokenExchangeHandler{}
+	return &TokenExchangeHandler{
+		strategy.(openid.OpenIDConnectTokenStrategy),
+		strategy.(oauth2.AccessTokenStrategy),
+		storage.(oauth2.AccessTokenStorage),
+	}
 }
 
 type TokenExchangeHandler struct {
+	idTokenStrategy     openid.OpenIDConnectTokenStrategy
+	accessTokenStrategy oauth2.AccessTokenStrategy
+	accessTokenStorage  oauth2.AccessTokenStorage
 }
 
 func (t *TokenExchangeHandler) HandleTokenEndpointRequest(ctx context.Context, requester fosite.AccessRequester) error {
+	if !(requester.GetGrantTypes().ExactOne("urn:ietf:params:oauth:grant-type:token-exchange")) {
+		return errors.WithStack(fosite.ErrUnknownRequest)
+	}
 	return nil
 }
 
 func (t *TokenExchangeHandler) PopulateTokenEndpointResponse(ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder) error {
+	params := requester.GetRequestForm()
+	accessToken := params.Get("subject_token")
+	if err := t.accessTokenStrategy.ValidateAccessToken(ctx, requester, accessToken); err != nil {
+		return errors.WithStack(err)
+	}
+	signature := t.accessTokenStrategy.AccessTokenSignature(accessToken)
+	accessTokenSession, err := t.accessTokenStorage.GetAccessTokenSession(ctx, signature, requester.GetSession())
+	if err != nil {
+		return errors.WithStack(err)
+	}
+	if !accessTokenSession.GetGrantedScopes().Has("pinniped.sts.unrestricted") {
+		return errors.WithStack(fosite.ErrScopeNotGranted)
+	}
+	// TODO check the other requester fields
+	scopedDownRequester := fosite.NewAccessRequest(requester.GetSession())
+	scopedDownRequester.GrantedAudience = []string{params.Get("audience")}
+	newToken, err := t.idTokenStrategy.GenerateIDToken(ctx, scopedDownRequester)
+	if err != nil {
+		return errors.WithStack(err)
+	}
+	responder.SetAccessToken(newToken)
+	responder.SetTokenType("N_A")
+	responder.SetExtra("issued_token_type", "urn:ietf:params:oauth:token-type:jwt")
 	return nil
 }