From 1e7f32936487c0d8052d63fd32d84f4c1121986e Mon Sep 17 00:00:00 2001
From: Roland Shoemaker <roland@golang.org>
Date: Mon, 27 Mar 2023 14:21:45 -0700
Subject: [PATCH] oauth2: add ReuseTokenSourceWithExpiry

Add a constructor which allows for the configuration of the expiryDelta
buffer. Due to the construction of reuseTokenSource and Token we need
to store the new delta in both places, so the behavior of Valid is
consistent regardless of where it is called from.

Fixes #623

Change-Id: I89f9c206a9cc16bb473b8c619605c8410a82fff0
Reviewed-on: https://go-review.googlesource.com/c/oauth2/+/479676
Run-TryBot: Roland Shoemaker <roland@golang.org>
Reviewed-by: Cody Oss <codyoss@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
---
 oauth2.go     | 31 +++++++++++++++++++++++++++++++
 token.go      | 14 ++++++++++++--
 token_test.go |  8 ++++++--
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/oauth2.go b/oauth2.go
index de52304..9085fab 100644
--- a/oauth2.go
+++ b/oauth2.go
@@ -16,6 +16,7 @@ import (
 	"net/url"
 	"strings"
 	"sync"
+	"time"
 
 	"golang.org/x/oauth2/internal"
 )
@@ -290,6 +291,8 @@ type reuseTokenSource struct {
 
 	mu sync.Mutex // guards t
 	t  *Token
+
+	expiryDelta time.Duration
 }
 
 // Token returns the current token if it's still valid, else will
@@ -305,6 +308,7 @@ func (s *reuseTokenSource) Token() (*Token, error) {
 	if err != nil {
 		return nil, err
 	}
+	t.expiryDelta = s.expiryDelta
 	s.t = t
 	return t, nil
 }
@@ -379,3 +383,30 @@ func ReuseTokenSource(t *Token, src TokenSource) TokenSource {
 		new: src,
 	}
 }
+
+// ReuseTokenSource returns a TokenSource that acts in the same manner as the
+// TokenSource returned by ReuseTokenSource, except the expiry buffer is
+// configurable. The expiration time of a token is calculated as
+// t.Expiry.Add(-earlyExpiry).
+func ReuseTokenSourceWithExpiry(t *Token, src TokenSource, earlyExpiry time.Duration) TokenSource {
+	// Don't wrap a reuseTokenSource in itself. That would work,
+	// but cause an unnecessary number of mutex operations.
+	// Just build the equivalent one.
+	if rt, ok := src.(*reuseTokenSource); ok {
+		if t == nil {
+			// Just use it directly, but set the expiryDelta to earlyExpiry,
+			// so the behavior matches what the user expects.
+			rt.expiryDelta = earlyExpiry
+			return rt
+		}
+		src = rt.new
+	}
+	if t != nil {
+		t.expiryDelta = earlyExpiry
+	}
+	return &reuseTokenSource{
+		t:           t,
+		new:         src,
+		expiryDelta: earlyExpiry,
+	}
+}
diff --git a/token.go b/token.go
index 8227203..7c64006 100644
--- a/token.go
+++ b/token.go
@@ -16,10 +16,10 @@ import (
 	"golang.org/x/oauth2/internal"
 )
 
-// expiryDelta determines how earlier a token should be considered
+// defaultExpiryDelta determines how earlier a token should be considered
 // expired than its actual expiration time. It is used to avoid late
 // expirations due to client-server time mismatches.
-const expiryDelta = 10 * time.Second
+const defaultExpiryDelta = 10 * time.Second
 
 // Token represents the credentials used to authorize
 // the requests to access protected resources on the OAuth 2.0
@@ -52,6 +52,11 @@ type Token struct {
 	// raw optionally contains extra metadata from the server
 	// when updating a token.
 	raw interface{}
+
+	// expiryDelta is used to calculate when a token is considered
+	// expired, by subtracting from Expiry. If zero, defaultExpiryDelta
+	// is used.
+	expiryDelta time.Duration
 }
 
 // Type returns t.TokenType if non-empty, else "Bearer".
@@ -127,6 +132,11 @@ func (t *Token) expired() bool {
 	if t.Expiry.IsZero() {
 		return false
 	}
+
+	expiryDelta := defaultExpiryDelta
+	if t.expiryDelta != 0 {
+		expiryDelta = t.expiryDelta
+	}
 	return t.Expiry.Round(0).Add(-expiryDelta).Before(timeNow())
 }
 
diff --git a/token_test.go b/token_test.go
index ee97b4f..0d8c7df 100644
--- a/token_test.go
+++ b/token_test.go
@@ -43,9 +43,13 @@ func TestTokenExpiry(t *testing.T) {
 		want bool
 	}{
 		{name: "12 seconds", tok: &Token{Expiry: now.Add(12 * time.Second)}, want: false},
-		{name: "10 seconds", tok: &Token{Expiry: now.Add(expiryDelta)}, want: false},
-		{name: "10 seconds-1ns", tok: &Token{Expiry: now.Add(expiryDelta - 1*time.Nanosecond)}, want: true},
+		{name: "10 seconds", tok: &Token{Expiry: now.Add(defaultExpiryDelta)}, want: false},
+		{name: "10 seconds-1ns", tok: &Token{Expiry: now.Add(defaultExpiryDelta - 1*time.Nanosecond)}, want: true},
 		{name: "-1 hour", tok: &Token{Expiry: now.Add(-1 * time.Hour)}, want: true},
+		{name: "12 seconds, custom expiryDelta", tok: &Token{Expiry: now.Add(12 * time.Second), expiryDelta: time.Second * 5}, want: false},
+		{name: "5 seconds, custom expiryDelta", tok: &Token{Expiry: now.Add(time.Second * 5), expiryDelta: time.Second * 5}, want: false},
+		{name: "5 seconds-1ns, custom expiryDelta", tok: &Token{Expiry: now.Add(time.Second*5 - 1*time.Nanosecond), expiryDelta: time.Second * 5}, want: true},
+		{name: "-1 hour, custom expiryDelta", tok: &Token{Expiry: now.Add(-1 * time.Hour), expiryDelta: time.Second * 5}, want: true},
 	}
 	for _, tc := range cases {
 		if got, want := tc.tok.expired(), tc.want; got != want {