handler/openid: Populate at_hash in explicit/refresh flows (#315)

Signed-off-by: Wenhao Ni <niwenhao@gmail.com>

handler/openid/flow_explicit_auth_test.go

@@ -34,9 +34,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// expose key to verify id_token
+var key = internal.MustRSAKey()
 var j = &DefaultStrategy{
 	JWTStrategy: &jwt.RS256JWTStrategy{
-		PrivateKey: internal.MustRSAKey(),
+		PrivateKey: key,
 	},
 }
 

handler/openid/flow_explicit_token.go

@@ -52,6 +52,18 @@ func (c *OpenIDConnectExplicitHandler) PopulateTokenEndpointResponse(ctx context
 		return errors.WithStack(fosite.ErrInvalidGrant.WithHint("The OAuth 2.0 Client is not allowed to use the authorization grant \"authorization_code\"."))
 	}
 
+	sess, ok := requester.GetSession().(Session)
+	if !ok {
+		return errors.WithStack(fosite.ErrServerError.WithDebug("Failed to generate id token because session must be of type fosite/handler/openid.Session."))
+	}
+
+	claims := sess.IDTokenClaims()
+	if claims.Subject == "" {
+		return errors.WithStack(fosite.ErrServerError.WithDebug("Failed to generate id token because subject is an empty string."))
+	}
+
+	claims.AccessTokenHash = c.GetAccessTokenHash(ctx, requester, responder)
+
 	// The response type `id_token` is only required when performing the implicit or hybrid flow, see:
 	// https://openid.net/specs/openid-connect-registration-1_0.html
 	//

handler/openid/flow_explicit_token_test.go

@@ -33,6 +33,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	jwtgo "github.com/dgrijalva/jwt-go"
 )
 
 func TestHandleTokenEndpointRequest(t *testing.T) {
@@ -68,6 +69,7 @@ func TestExplicit_PopulateTokenEndpointResponse(t *testing.T) {
 		description string
 		setup       func()
 		expectErr   error
+		check       func(t *testing.T, aresp *fosite.AccessResponse)
 	}{
 		{
 			description: "should fail because invalid response type",
@@ -112,6 +114,29 @@ func TestExplicit_PopulateTokenEndpointResponse(t *testing.T) {
 				r.Form.Set("nonce", "1111111111111111")
 				store.EXPECT().GetOpenIDConnectSession(nil, gomock.Any(), areq).AnyTimes().Return(r, nil)
 			},
+			check: func(t *testing.T, aresp *fosite.AccessResponse) {
+				assert.NotEmpty(t, aresp.GetExtra("id_token"))
+				idToken, _ := aresp.GetExtra("id_token").(string)
+				decodedIdToken, _ := jwtgo.Parse(idToken, func(token *jwtgo.Token)(interface{}, error) {
+					return key.PublicKey, nil
+				})
+				claims, _ := decodedIdToken.Claims.(jwtgo.MapClaims)
+				assert.NotEmpty(t, claims["at_hash"])			
+			},
+		},
+		{
+			description: "should fail because missing subject claim",
+			setup: func() {
+				session.Claims.Subject = ""
+			},
+			expectErr: fosite.ErrServerError,
+		},
+		{
+			description: "should fail because missing session",
+			setup: func() {
+				areq.Session = nil
+			},
+			expectErr: fosite.ErrServerError,
 		},
 	} {
 		t.Run(fmt.Sprintf("case=%d", k), func(t *testing.T) {
@@ -123,6 +148,9 @@ func TestExplicit_PopulateTokenEndpointResponse(t *testing.T) {
 			} else {
 				require.NoError(t, err)
 			}
+			if c.check != nil {
+				c.check(t, aresp)
+			}
 		})
 	}
 }

handler/openid/flow_refresh_token.go

@@ -82,5 +82,19 @@ func (c *OpenIDConnectRefreshHandler) PopulateTokenEndpointResponse(ctx context.
 	// 	 return errors.WithStack(fosite.ErrUnknownRequest.WithDebug("The client is not allowed to use response type id_token"))
 	// }
 
+	sess, ok := requester.GetSession().(Session)
+	if !ok {
+		return errors.WithStack(fosite.ErrServerError.WithDebug("Failed to generate id token because session must be of type fosite/handler/openid.Session."))
+	}
+
+	claims := sess.IDTokenClaims()
+	if claims.Subject == "" {
+		return errors.WithStack(fosite.ErrServerError.WithDebug("Failed to generate id token because subject is an empty string."))
+	}
+
+	hash := c.GetAccessTokenHash(ctx, requester, responder)
+
+	claims.AccessTokenHash = hash
+
 	return c.IssueExplicitIDToken(ctx, requester, responder)
 }

handler/openid/flow_refresh_token_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	jwtgo "github.com/dgrijalva/jwt-go"
 )
 
 func TestOpenIDConnectRefreshHandler_HandleTokenEndpointRequest(t *testing.T) {
@@ -155,8 +156,45 @@ func TestOpenIDConnectRefreshHandler_PopulateTokenEndpointResponse(t *testing.T)
 			},
 			check: func(t *testing.T, aresp *fosite.AccessResponse) {
 				assert.NotEmpty(t, aresp.GetExtra("id_token"))
+				idToken, _ := aresp.GetExtra("id_token").(string)
+				decodedIdToken, _ := jwtgo.Parse(idToken, func(token *jwtgo.Token)(interface{}, error) {
+					return key.PublicKey, nil
+				})
+				claims, _ := decodedIdToken.Claims.(jwtgo.MapClaims)
+				assert.NotEmpty(t, claims["at_hash"])			
 			},
 		},
+		{
+			description: "should fail because missing subject claim",
+			areq: &fosite.AccessRequest{
+				GrantTypes: []string{"refresh_token"},
+				Request: fosite.Request{
+					GrantedScopes: []string{"openid"},
+					Client: &fosite.DefaultClient{
+						GrantTypes: []string{"refresh_token"},
+						//ResponseTypes: []string{"id_token"},
+					},
+					Session: &DefaultSession{
+						Subject: "foo",
+						Claims:  &jwt.IDTokenClaims{},
+					},
+				},
+			},
+			expectedErr: fosite.ErrServerError,
+		},
+		{
+			description: "should fail because missing session",
+			areq: &fosite.AccessRequest{
+				GrantTypes: []string{"refresh_token"},
+				Request: fosite.Request{
+					GrantedScopes: []string{"openid"},
+					Client: &fosite.DefaultClient{
+						GrantTypes: []string{"refresh_token"},
+					},
+				},
+			},
+			expectedErr: fosite.ErrServerError,
+		},
 	} {
 		t.Run("case="+c.description, func(t *testing.T) {
 			aresp := fosite.NewAccessResponse()

handler/openid/helper.go

@@ -22,7 +22,10 @@
 package openid
 
 import (
+	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/base64"
 
 	"github.com/ory/fosite"
 )
@@ -31,6 +34,18 @@ type IDTokenHandleHelper struct {
 	IDTokenStrategy OpenIDConnectTokenStrategy
 }
 
+func (i *IDTokenHandleHelper) GetAccessTokenHash(ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder) string {
+	token := responder.GetAccessToken()
+
+	buffer := bytes.NewBufferString(token)
+	hash := sha256.New()
+	hash.Write(buffer.Bytes())
+	hashBuf := bytes.NewBuffer(hash.Sum([]byte{}))
+	len := hashBuf.Len()
+
+	return base64.RawURLEncoding.EncodeToString(hashBuf.Bytes()[:len/2])
+}
+
 func (i *IDTokenHandleHelper) generateIDToken(ctx context.Context, fosr fosite.Requester) (token string, err error) {
 	token, err = i.IDTokenStrategy.GenerateIDToken(ctx, fosr)
 	if err != nil {

handler/openid/helper_test.go

@@ -120,3 +120,18 @@ func TestIssueImplicitToken(t *testing.T) {
 	err := h.IssueImplicitIDToken(nil, ar, resp)
 	assert.NoError(t, err)
 }
+
+func TestGetAccessTokenHash(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	req := internal.NewMockAccessRequester(ctrl)
+	resp := internal.NewMockAccessResponder(ctrl)
+
+	defer ctrl.Finish()
+
+	resp.EXPECT().GetAccessToken().Return("7a35f818-9164-48cb-8c8f-e1217f44228431c41102-d410-4ed5-9276-07ba53dfdcd8")
+
+	h := &IDTokenHandleHelper{IDTokenStrategy: strat}
+
+	hash := h.GetAccessTokenHash(nil, req, resp)
+	assert.Equal(t, "Zfn_XBitThuDJiETU3OALQ", hash)
+}