From 62a15962f1a331a998b6fe6b724b814490be5fe4 Mon Sep 17 00:00:00 2001
From: Florian Forster <ff@octo.it>
Date: Thu, 11 Jan 2018 12:56:30 +0100
Subject: [PATCH] Package app: Wrap oauth2.TokenSource to ensure datastore is
 always updated.

---
 app/user.go      | 44 +++++++++++++++++++++++++++++++++++++++++++-
 fitbit/fitbit.go | 25 ++-----------------------
 2 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/app/user.go b/app/user.go
index 8465c2c..e47a5cc 100644
--- a/app/user.go
+++ b/app/user.go
@@ -3,11 +3,13 @@ package app
 import (
 	"context"
 	"fmt"
+	"net/http"
 
 	"github.com/google/uuid"
 	legacy_context "golang.org/x/net/context"
 	"golang.org/x/oauth2"
 	"google.golang.org/appengine/datastore"
+	"google.golang.org/appengine/log"
 )
 
 type User struct {
@@ -78,7 +80,47 @@ func (u *User) Token(ctx context.Context, svc string) (*oauth2.Token, error) {
 }
 
 func (u *User) SetToken(ctx context.Context, svc string, tok *oauth2.Token) error {
-	key := datastore.NewKey(ctx, "Token", "Fitbit", 0, u.key)
+	key := datastore.NewKey(ctx, "Token", svc, 0, u.key)
 	_, err := datastore.Put(ctx, key, tok)
 	return err
 }
+
+func (u *User) OAuthClient(ctx context.Context, svc string, cfg *oauth2.Config) (*http.Client, error) {
+	key := datastore.NewKey(ctx, "Token", svc, 0, u.key)
+
+	var tok oauth2.Token
+	if err := datastore.Get(ctx, key, &tok); err != nil {
+		return nil, err
+	}
+
+	src := cfg.TokenSource(ctx, &tok)
+	return oauth2.NewClient(ctx, &persistingTokenSource{
+		ctx: ctx,
+		t:   &tok,
+		src: src,
+		key: key,
+	}), nil
+}
+
+type persistingTokenSource struct {
+	ctx context.Context
+	t   *oauth2.Token
+	src oauth2.TokenSource
+	key *datastore.Key
+}
+
+func (s *persistingTokenSource) Token() (*oauth2.Token, error) {
+	tok, err := s.src.Token()
+	if err != nil {
+		return nil, err
+	}
+
+	if s.t.RefreshToken != tok.RefreshToken {
+		if _, err := datastore.Put(s.ctx, s.key, tok); err != nil {
+			log.Errorf(s.ctx, "persisting OAuth token in datastore failed: %v", err)
+		}
+	}
+
+	s.t = tok
+	return tok, nil
+}
diff --git a/fitbit/fitbit.go b/fitbit/fitbit.go
index bdfc3dd..81d4ebc 100644
--- a/fitbit/fitbit.go
+++ b/fitbit/fitbit.go
@@ -127,36 +127,15 @@ func NewClient(ctx context.Context, fitbitUserID string, u *app.User) (*Client,
 		fitbitUserID = "-"
 	}
 
-	storedToken, err := u.Token(ctx, "Fitbit")
+	c, err := u.OAuthClient(ctx, "Fitbit", oauth2Config)
 	if err != nil {
 		return nil, err
 	}
 
-	// The oauth2 package will refresh the token when it is valid for less
-	// than 10 seconds. To avoid a race with later calls (which would
-	// refresh the token but the new RefreshToken wouldn't make it back
-	// into datastore), we refresh earlier than that. The Fitbit tokens are
-	// quite long-lived (six hours?); the additional load this puts on the
-	// backends is negligible.
-	if storedToken.Expiry.Round(0).Add(-5 * time.Minute).Before(time.Now()) {
-		storedToken.Expiry = time.Now()
-	}
-
-	refreshedToken, err := oauth2Config.TokenSource(ctx, storedToken).Token()
-	if err != nil {
-		return nil, err
-	}
-
-	if refreshedToken.RefreshToken != storedToken.RefreshToken {
-		if err := u.SetToken(ctx, "Fitbit", refreshedToken); err != nil {
-			return nil, err
-		}
-	}
-
 	return &Client{
 		fitbitUserID: fitbitUserID,
 		appUser:      u,
-		client:       oauth2Config.Client(ctx, refreshedToken),
+		client:       c,
 	}, nil
 }
 
-- 
2.11.0