From e5d53ad4acfd1bb9094172c718cdb9885e633bc8 Mon Sep 17 00:00:00 2001
From: Josh Carp <jm.carp@gmail.com>
Date: Tue, 30 Nov 2021 10:22:11 -0500
Subject: [PATCH] Refresh public keys on lookup failure.

[Fixes #10]
---
 config.go     | 25 ++-----------------------
 jwt/config.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 jwt/token.go  | 11 ++---------
 3 files changed, 49 insertions(+), 32 deletions(-)

diff --git a/config.go b/config.go
index 77831c6..5234f7c 100644
--- a/config.go
+++ b/config.go
@@ -42,7 +42,8 @@ func initConfig() error {
 		return err
 	}
 	initDomains(*domains)
-	if err := initPublicKeys(*publicKeysPath); err != nil {
+	cfg.PublicKeyPath = *publicKeysPath
+	if err := cfg.RefreshPublicKeys(); err != nil {
 		return err
 	}
 	return nil
@@ -116,25 +117,3 @@ func initDomains(domains string) {
 		}
 	}
 }
-
-func initPublicKeys(filePath string) error {
-	var err error
-	if len(filePath) != 0 {
-		cfg.PublicKeys, err = loadPublicKeysFromFile(filePath)
-	} else {
-		cfg.PublicKeys, err = jwt.FetchPublicKeys()
-	}
-	if err != nil {
-		return err
-	}
-	return cfg.Validate()
-}
-
-func loadPublicKeysFromFile(filePath string) (map[string]jwt.PublicKey, error) {
-	f, err := os.Open(filePath)
-	if err != nil {
-		return nil, err
-	}
-	defer f.Close()
-	return jwt.DecodePublicKeys(f)
-}
diff --git a/jwt/config.go b/jwt/config.go
index c50efc2..a49cd9c 100644
--- a/jwt/config.go
+++ b/jwt/config.go
@@ -2,17 +2,53 @@ package jwt
 
 import (
 	"errors"
+	"fmt"
+	"os"
 	"regexp"
+
+	jwt "github.com/dgrijalva/jwt-go"
 )
 
 // Config specifies the parameters for which to perform validation of JWT
 // tokens in requests against.
 type Config struct {
+	PublicKeyPath  string
 	PublicKeys     map[string]PublicKey
 	MatchAudiences *regexp.Regexp
 	MatchDomains   map[string]bool
 }
 
+func (cfg *Config) RefreshPublicKeys() error {
+	var err error
+	if len(cfg.PublicKeyPath) != 0 {
+		cfg.PublicKeys, err = loadPublicKeysFromFile(cfg.PublicKeyPath)
+	} else {
+		cfg.PublicKeys, err = FetchPublicKeys()
+	}
+	return err
+}
+
+func (cfg *Config) GetPublicKey(keyID string) (interface{}, error) {
+	key, ok := cfg.PublicKeys[keyID]
+	// Refresh public keys on lookup failure. For details, see
+	// https://github.com/imkira/gcp-iap-auth/issues/10 and
+	// https://stackoverflow.com/questions/44828856/google-iap-public-keys-expiry.
+	if !ok {
+		if err := cfg.RefreshPublicKeys(); err != nil {
+			return nil, err
+		}
+		key, ok = cfg.PublicKeys[keyID]
+		if !ok {
+			return nil, fmt.Errorf("No public key for %q", keyID)
+		}
+	}
+	parsedKey, err := jwt.ParseECPublicKeyFromPEM(key)
+	if err != nil {
+		return nil, fmt.Errorf("Failed to parse key: %v", err)
+	}
+	return parsedKey, nil
+}
+
 // Validate validates the Configuration.
 func (cfg *Config) Validate() error {
 	if cfg.MatchAudiences == nil {
@@ -31,3 +67,12 @@ func (cfg *Config) matchesAudience(aud *Audience) bool {
 func (cfg *Config) matchesDomain(hd string) bool {
 	return len(cfg.MatchDomains) == 0 || cfg.MatchDomains[hd]
 }
+
+func loadPublicKeysFromFile(filePath string) (map[string]PublicKey, error) {
+	f, err := os.Open(filePath)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	return DecodePublicKeys(f)
+}
diff --git a/jwt/token.go b/jwt/token.go
index 150b2cd..fd0576e 100644
--- a/jwt/token.go
+++ b/jwt/token.go
@@ -19,15 +19,8 @@ func tokenKey(token *jwt.Token) (interface{}, error) {
 		return nil, fmt.Errorf("Invalid algorithm: %v", token.Header[algorithmClaim])
 	}
 	keyID, _ := token.Header[keyIDClaim].(string)
-	key := token.Claims.(*Claims).cfg.PublicKeys[keyID]
-	if len(key) == 0 {
-		return nil, fmt.Errorf("No public key for %q", keyID)
-	}
-	parsedKey, err := jwt.ParseECPublicKeyFromPEM(key)
-	if err != nil {
-		return nil, fmt.Errorf("Failed to parse key: %v", err)
-	}
-	return parsedKey, nil
+	cfg := token.Claims.(*Claims).cfg
+	return cfg.GetPublicKey(keyID)
 }
 
 func tokenMethod(token *jwt.Token) (jwt.SigningMethod, bool) {