From 51187a78a6451c0ee49d68202e6327a94fbabb1e Mon Sep 17 00:00:00 2001
From: Eric Chiang <eric.chiang.m@gmail.com>
Date: Wed, 31 Aug 2022 20:41:12 -0700
Subject: [PATCH] oidc: add verifier benchmark and fix up errors.As in test

---
 oidc/jwks_test.go   | 48 +++++++++++++++++++++++++++++++++++++++++++--
 oidc/verify_test.go | 21 +++++++++++---------
 2 files changed, 58 insertions(+), 11 deletions(-)

diff --git a/oidc/jwks_test.go b/oidc/jwks_test.go
index 69b55605..5c5ef4da 100644
--- a/oidc/jwks_test.go
+++ b/oidc/jwks_test.go
@@ -8,6 +8,7 @@ import (
 	"crypto/rand"
 	"crypto/rsa"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strconv"
@@ -39,7 +40,7 @@ type signingKey struct {
 }
 
 // sign creates a JWS using the private key from the provided payload.
-func (s *signingKey) sign(t *testing.T, payload []byte) string {
+func (s *signingKey) sign(t testing.TB, payload []byte) string {
 	privKey := &jose.JSONWebKey{Key: s.priv, Algorithm: string(s.alg), KeyID: s.keyID}
 
 	signer, err := jose.NewSigner(jose.SigningKey{Algorithm: s.alg, Key: privKey}, nil)
@@ -62,7 +63,7 @@ func (s *signingKey) jwk() jose.JSONWebKey {
 	return jose.JSONWebKey{Key: s.pub, Use: "sig", Algorithm: string(s.alg), KeyID: s.keyID}
 }
 
-func newRSAKey(t *testing.T) *signingKey {
+func newRSAKey(t testing.TB) *signingKey {
 	priv, err := rsa.GenerateKey(rand.Reader, 1028)
 	if err != nil {
 		t.Fatal(err)
@@ -236,3 +237,46 @@ func TestRotation(t *testing.T) {
 		t.Errorf("failed to verify valid signature: %v", err)
 	}
 }
+
+func BenchmarkVerify(b *testing.B) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	key := newRSAKey(b)
+
+	now := time.Date(2022, 01, 29, 0, 0, 0, 0, time.UTC)
+	exp := now.Add(time.Hour)
+	payload := []byte(fmt.Sprintf(`{
+		"iss": "https://example.com",
+		"sub": "test_user",
+		"aud": "test_client_id",
+		"exp": %d
+	}`, exp.Unix()))
+
+	idToken := key.sign(b, payload)
+	server := &keyServer{
+		keys: jose.JSONWebKeySet{
+			Keys: []jose.JSONWebKey{key.jwk()},
+		},
+	}
+	s := httptest.NewServer(server)
+	defer s.Close()
+
+	rks := NewRemoteKeySet(ctx, s.URL)
+	verifier := NewVerifier("https://example.com", rks, &Config{
+		ClientID: "test_client_id",
+		Now:      func() time.Time { return now },
+	})
+
+	// Trigger the remote key set to query the public keys and cache them.
+	if _, err := verifier.Verify(ctx, idToken); err != nil {
+		b.Fatalf("verifying id token: %v", err)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := verifier.Verify(ctx, idToken); err != nil {
+			b.Fatalf("verifying id token: %v", err)
+		}
+	}
+}
diff --git a/oidc/verify_test.go b/oidc/verify_test.go
index 3a93f2af..2ee59541 100644
--- a/oidc/verify_test.go
+++ b/oidc/verify_test.go
@@ -73,8 +73,8 @@ func TestVerify(t *testing.T) {
 			config: Config{
 				SkipClientIDCheck: true,
 			},
-			signKey:   newRSAKey(t),
-			wantErrAs: &TokenExpiredError{},
+			signKey:       newRSAKey(t),
+			wantErrExpiry: true,
 		},
 		{
 			name:    "unexpired token",
@@ -531,9 +531,9 @@ type verificationTest struct {
 	// testing invalid signatures.
 	verificationKey *signingKey
 
-	config    Config
-	wantErr   bool
-	wantErrAs error
+	config        Config
+	wantErr       bool
+	wantErrExpiry bool
 }
 
 func (v verificationTest) runGetToken(t *testing.T) (*IDToken, error) {
@@ -559,13 +559,16 @@ func (v verificationTest) runGetToken(t *testing.T) (*IDToken, error) {
 
 func (v verificationTest) run(t *testing.T) {
 	_, err := v.runGetToken(t)
-	if err != nil && !v.wantErr && v.wantErrAs == nil {
+	if err != nil && !v.wantErr && !v.wantErrExpiry {
 		t.Errorf("%v", err)
 	}
-	if err == nil && (v.wantErr || v.wantErrAs != nil) {
+	if err == nil && (v.wantErr || v.wantErrExpiry) {
 		t.Errorf("expected error")
 	}
-	if v.wantErrAs != nil && !errors.As(err, &v.wantErrAs) {
-		t.Errorf("expected error %q but got %q", v.wantErrAs, err)
+	if v.wantErrExpiry {
+		var errExp *TokenExpiredError
+		if !errors.As(err, &errExp) {
+			t.Errorf("expected *TokenExpiryError but got %q", err)
+		}
 	}
 }