From d5ff5ab876fbb2a8c14cc5470db12f6b60fb90ef Mon Sep 17 00:00:00 2001
From: robnorman <rob.norman@infinitycloud.com>
Date: Wed, 16 Sep 2015 16:33:09 +0000
Subject: [PATCH] jwt: allow setting a custom expiry time for JWT tokens

The current implementation of JWS/JWT in this package uses a fixed
1 hour expiry time for JWT tokens.

Some services do not accept such a long expiry time, e.g. Salesforce,
which defaults to a 5 minute expiry.
https://help.salesforce.com/HTViewHelpDoc?id=remoteaccess_oauth_jwt_flow.htm

This change adds an Expires time.Duration property to the jwt.Config
struct that, if set, will be used to calculate the jws.ClaimSet Exp property.
It allows a custom expiry to be set on a JWT token.

This change is backward compatible and will revert to previous behaviour if
the Expires property is not set.

Fixes golang/oauth2#151

Change-Id: I3159ac2a5711ef10389d83c0e290bfc7a9f54015
Reviewed-on: https://go-review.googlesource.com/14681
Reviewed-by: Burcu Dogan <jbd@google.com>
---
 jws/jws.go | 25 ++++++++++++-------------
 jwt/jwt.go |  6 ++++++
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/jws/jws.go b/jws/jws.go
index 37d8651..6e03198 100644
--- a/jws/jws.go
+++ b/jws/jws.go
@@ -41,23 +41,22 @@ type ClaimSet struct {
 	// See http://tools.ietf.org/html/draft-jones-json-web-token-10#section-4.3
 	// This array is marshalled using custom code (see (c *ClaimSet) encode()).
 	PrivateClaims map[string]interface{} `json:"-"`
-
-	exp time.Time
-	iat time.Time
 }
 
 func (c *ClaimSet) encode() (string, error) {
-	if c.exp.IsZero() || c.iat.IsZero() {
-		// Reverting time back for machines whose time is not perfectly in sync.
-		// If client machine's time is in the future according
-		// to Google servers, an access token will not be issued.
-		now := time.Now().Add(-10 * time.Second)
-		c.iat = now
-		c.exp = now.Add(time.Hour)
+	// Reverting time back for machines whose time is not perfectly in sync.
+	// If client machine's time is in the future according
+	// to Google servers, an access token will not be issued.
+	now := time.Now().Add(-10 * time.Second)
+	if c.Iat == 0 {
+		c.Iat = now.Unix()
+	}
+	if c.Exp == 0 {
+		c.Exp = now.Add(time.Hour).Unix()
+	}
+	if c.Exp < c.Iat {
+		return "", fmt.Errorf("jws: invalid Exp must be later than Iat", c.Exp)
 	}
-
-	c.Exp = c.exp.Unix()
-	c.Iat = c.iat.Unix()
 
 	b, err := json.Marshal(c)
 	if err != nil {
diff --git a/jwt/jwt.go b/jwt/jwt.go
index 205d23e..11a2687 100644
--- a/jwt/jwt.go
+++ b/jwt/jwt.go
@@ -54,6 +54,9 @@ type Config struct {
 
 	// TokenURL is the endpoint required to complete the 2-legged JWT flow.
 	TokenURL string
+
+	// Expires optionally specifies how long the token is valid for.
+	Expires time.Duration
 }
 
 // TokenSource returns a JWT TokenSource using the configuration
@@ -95,6 +98,9 @@ func (js jwtSource) Token() (*oauth2.Token, error) {
 		// to be compatible with legacy OAuth 2.0 providers.
 		claimSet.Prn = subject
 	}
+	if t := js.conf.Expires; t > 0 {
+		claimSet.Exp = time.Now().Add(t).Unix()
+	}
 	payload, err := jws.Encode(defaultHeader, claimSet, pk)
 	if err != nil {
 		return nil, err