DA-4596: dto changes for caching auth0 jwks

Signed-off-by: Ajinkya Nahar <[email protected]>

Author: Ajinkya Nahar

auth0/dto.go

@@ -9,6 +9,13 @@ type AuthToken struct {
 	CreatedAt time.Time `json:"created_at"`
 }
 
+// AuthJwks Struct
+type AuthJwks struct {
+	Name      string    `json:"name"`
+	Jwks      string    `json:"jwks"`
+	CreatedAt time.Time `json:"created_at"`
+}
+
 // Resp struct
 type Resp struct {
 	AccessToken string `json:"access_token"`
@@ -46,6 +53,36 @@ type ESTokenSchema struct {
 	} `json:"hits"`
 }
 
+// ESJwksSchema ...
+type ESJwksSchema struct {
+	Took     int  `json:"took"`
+	TimedOut bool `json:"timed_out"`
+	Shards   struct {
+		Total      int `json:"total"`
+		Successful int `json:"successful"`
+		Skipped    int `json:"skipped"`
+		Failed     int `json:"failed"`
+	} `json:"_shards"`
+	Hits struct {
+		Total struct {
+			Value    int    `json:"value"`
+			Relation string `json:"relation"`
+		} `json:"total"`
+		MaxScore float64 `json:"max_score"`
+		Hits     []struct {
+			Index  string  `json:"_index"`
+			Type   string  `json:"_type"`
+			ID     string  `json:"_id"`
+			Score  float64 `json:"_score"`
+			Source struct {
+				Name      string    `json:"name"`
+				Jwks      string    `json:"jwks"`
+				CreatedAt time.Time `json:"created_at"`
+			} `json:"_source"`
+		} `json:"hits"`
+	} `json:"hits"`
+}
+
 // LastActionSchema ...
 type LastActionSchema struct {
 	Took     int  `json:"took"`
@@ -79,12 +116,14 @@ const (
 	lastAuth0TokenRequest = "last-auth0-token-request-"
 	auth0TokenCache       = "auth0-token-cache-"
 	tokenDoc              = "token"
+	auth0JwksCache        = "auth0-jwks-cache-"
+	jwksDoc               = "jwks"
 )
 
 // RefreshResult ...
 type RefreshResult string
 
-const(
+const (
 	// RefreshError ...
 	RefreshError RefreshResult = "error refreshing auth0 token"
 	// RefreshSuccessful ...

auth0/jwks.go

@@ -0,0 +1,126 @@
+package auth0
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"strings"
+	"time"
+
+	"github.com/dgrijalva/jwt-go"
+)
+
+// Jwks result from auth0 well know keys
+type Jwks struct {
+	Keys []JSONWebKeys `json:"keys"`
+}
+
+// JSONWebKeys auth0 token key
+type JSONWebKeys struct {
+	Kty string   `json:"kty"`
+	Kid string   `json:"kid"`
+	Use string   `json:"use"`
+	N   string   `json:"n"`
+	E   string   `json:"e"`
+	X5c []string `json:"x5c"`
+}
+
+func (a *ClientProvider) createAuthJwks(cert string) error {
+	log.Println("creating new auth jwks cert string")
+	at := AuthJwks{
+		Name:      "AuthJwks",
+		Jwks:      cert,
+		CreatedAt: time.Now().UTC(),
+	}
+	_, err := a.esClient.UpdateDocument(fmt.Sprintf("%s%s", auth0JwksCache, a.Environment), jwksDoc, at)
+	if err != nil {
+		log.Println("could not write the data", err)
+		return err
+	}
+
+	return nil
+}
+
+func (a *ClientProvider) getPemCert(token *jwt.Token, refreshJwks bool) (string, error) {
+	cert := ""
+	cert, expired, err := a.getCachedJwks()
+	if err != nil {
+		return cert, err
+	}
+
+	// check if the cache expired as well is not invoked via refresh token cron
+	if !expired && !refreshJwks {
+		return cert, nil
+	}
+
+	_, resp, err := a.httpClient.Request(fmt.Sprintf("%s/oauth/.well-known/jwks.json", a.AuthURL), "GET", nil, nil, nil)
+	if err != nil {
+		return cert, err
+	}
+
+	var jwks = Jwks{}
+	if err := json.Unmarshal(resp, &jwks); err != nil {
+		return cert, err
+	}
+
+	for _, k := range jwks.Keys {
+		if token.Header["kid"] == k.Kid {
+			cert = "-----BEGIN CERTIFICATE-----\n" + k.X5c[0] + "\n-----END CERTIFICATE-----"
+		}
+	}
+
+	if cert == "" {
+		err := errors.New("unable to find appropriate key")
+		return cert, err
+	}
+
+	err = a.createAuthJwks(cert)
+	if err != nil {
+		return "", err
+	}
+
+	return cert, nil
+}
+
+func (a *ClientProvider) getCachedJwks() (string, bool, error) {
+	expired := true
+	res, err := a.esClient.Search(strings.TrimSpace(auth0JwksCache+a.Environment), searchJwksQuery)
+	if err != nil {
+		go func() {
+			errMsg := fmt.Sprintf("%s-%s: error cached jwks not found\n %s", a.appName, a.Environment, err)
+			err := a.slackClient.SendText(errMsg)
+			fmt.Println("Err: send to slack: ", err)
+		}()
+
+		return "", expired, err
+	}
+
+	var e ESJwksSchema
+	err = json.Unmarshal(res, &e)
+	if err != nil {
+		log.Println("repository: GetOauthJwks: could not unmarshal the data", err)
+		return "", expired, err
+	}
+
+	if len(e.Hits.Hits) > 0 {
+		data := e.Hits.Hits[0]
+		// compare current time v/s existing cached time + 30 mins
+		if data.Source.CreatedAt.Add(30*time.Minute).Unix() <= time.Now().UTC().Unix() {
+			expired = false
+		}
+
+		return data.Source.Jwks, expired, nil
+	}
+
+	return "", expired, errors.New("GetJwks: could not find the associated jwks")
+}
+
+var searchJwksQuery = map[string]interface{}{
+	"size": 1,
+	"query": map[string]interface{}{
+		"term": map[string]interface{}{
+			"_id": jwksDoc,
+		},
+	},
+}

auth0/token.go

@@ -85,9 +85,10 @@ func (a *ClientProvider) GetToken(input bool) (string, error) {
 	if authToken == "" {
 		return authToken, errors.New("cached token is empty")
 	}
+
 	if input {
 		// check token validity
-		ok, _, err := a.isValid(authToken)
+		ok, _, err := a.isValid(authToken, false)
 		if err != nil {
 			log.Println(err)
 			return "", err
@@ -99,6 +100,7 @@ func (a *ClientProvider) GetToken(input bool) (string, error) {
 
 		return authToken, errors.New("cached token is not valid")
 	}
+
 	return authToken, nil
 }
 
@@ -137,6 +139,7 @@ func (a *ClientProvider) generateToken() (string, error) {
 		}()
 		log.Println("Err: GenerateToken ", err)
 	}
+
 	go func() {
 		err = a.createLastActionDate()
 		log.Println(err)
@@ -149,7 +152,8 @@ func (a *ClientProvider) generateToken() (string, error) {
 	if result.AccessToken != "" {
 		log.Println("GenerateToken: Token generated successfully.")
 	}
-	ok, _, err := a.isValid(result.AccessToken)
+
+	ok, _, err := a.isValid(result.AccessToken, true)
 	if !ok || err != nil {
 		go func() {
 			errMsg := fmt.Sprintf("%s-%s: error validating the newly created token\n %s", a.appName, a.Environment, err)
@@ -222,13 +226,13 @@ var searchCacheQuery = map[string]interface{}{
 	},
 }
 
-func (a *ClientProvider) isValid(token string) (bool, jwt.MapClaims, error) {
+func (a *ClientProvider) isValid(token string, refreshJwks bool) (bool, jwt.MapClaims, error) {
 	p, err := jwt.Parse(token, func(t *jwt.Token) (interface{}, error) {
 		if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
 			return nil, errors.New("unexpected signing method")
 		}
 
-		cert, err := a.getPemCert(t)
+		cert, err := a.getPemCert(t, refreshJwks)
 		if err != nil {
 			return nil, err
 		}
@@ -240,6 +244,7 @@ func (a *ClientProvider) isValid(token string) (bool, jwt.MapClaims, error) {
 
 		return key, nil
 	})
+
 	if err != nil {
 		return false, nil, err
 	}
@@ -252,47 +257,6 @@ func (a *ClientProvider) isValid(token string) (bool, jwt.MapClaims, error) {
 	return p.Valid, claims, err
 }
 
-// Jwks result from auth0 well know keys
-type Jwks struct {
-	Keys []JSONWebKeys `json:"keys"`
-}
-
-// JSONWebKeys auth0 token key
-type JSONWebKeys struct {
-	Kty string   `json:"kty"`
-	Kid string   `json:"kid"`
-	Use string   `json:"use"`
-	N   string   `json:"n"`
-	E   string   `json:"e"`
-	X5c []string `json:"x5c"`
-}
-
-func (a *ClientProvider) getPemCert(token *jwt.Token) (string, error) {
-	cert := ""
-	_, resp, err := a.httpClient.Request(fmt.Sprintf("%s/oauth/.well-known/jwks.json", a.AuthURL), "GET", nil, nil, nil)
-	if err != nil {
-		return cert, err
-	}
-
-	var jwks = Jwks{}
-	if err := json.Unmarshal(resp, &jwks); err != nil {
-		return cert, err
-	}
-
-	for _, k := range jwks.Keys {
-		if token.Header["kid"] == k.Kid {
-			cert = "-----BEGIN CERTIFICATE-----\n" + k.X5c[0] + "\n-----END CERTIFICATE-----"
-		}
-	}
-
-	if cert == "" {
-		err := errors.New("unable to find appropriate key")
-		return cert, err
-	}
-
-	return cert, nil
-}
-
 func (a *ClientProvider) createLastActionDate() error {
 	s := struct {
 		Date time.Time `json:"date"`
@@ -357,22 +321,28 @@ func (a *ClientProvider) RefreshToken() (RefreshResult, error) {
 	}
 
 	if authToken == "" || err != nil {
-		authToken, err = a.refreshCachedToken()
+		_, err = a.refreshCachedToken()
 		if err != nil {
 			return RefreshError, err
 		}
+
 		return RefreshSuccessful, nil
 	}
 
-	ok, claims, err := a.isValid(authToken)
+	ok, claims, err := a.isValid(authToken, false)
 	if ok && err == nil {
-		if claims.VerifyExpiresAt(time.Now().Add(60*time.Minute).Unix(), false) == false {
+		if !claims.VerifyExpiresAt(time.Now().Add(60*time.Minute).Unix(), false) {
+			if _, err := a.refreshCachedToken(); err != nil {
+				log.Printf("Error refresh auth0 token %s\n", err.Error())
+				return RefreshError, err
+			}
 			if _, err := a.refreshCachedToken(); err != nil {
 				log.Printf("Error refresh auth0 token %s\n", err.Error())
 				return RefreshError, err
 			}
 			return RefreshSuccessful, nil
 		}
+
 		return NotExpireSoon, nil
 	}