From a283521938e080178b8d4c536b0b004250c8aa74 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Edu=20G=C3=B3mez=20Escandell?=
 <edu.gomez.escandell@canonical.com>
Date: Mon, 30 Oct 2023 17:13:08 +0100
Subject: [PATCH] Avoid redundant talk with the contracts server

We don't need to querry the contract server if whe have a non-expired,
store-backed Ubuntu Pro token.
---
 windows-agent/internal/config/config.go       | 28 +++++++++-
 windows-agent/internal/contracts/contracts.go | 51 ++++++++++---------
 2 files changed, 53 insertions(+), 26 deletions(-)

diff --git a/windows-agent/internal/config/config.go b/windows-agent/internal/config/config.go
index 902297700..6222cba62 100644
--- a/windows-agent/internal/config/config.go
+++ b/windows-agent/internal/config/config.go
@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"sync"
 
+	"github.com/canonical/ubuntu-pro-for-windows/common"
 	"github.com/canonical/ubuntu-pro-for-windows/windows-agent/internal/config/registry"
 	"github.com/canonical/ubuntu-pro-for-windows/windows-agent/internal/contracts"
 	"github.com/canonical/ubuntu-pro-for-windows/windows-agent/internal/distros/database"
@@ -219,11 +220,36 @@ func (c *Config) FetchMicrosoftStoreSubscription(ctx context.Context) (err error
 		return fmt.Errorf("subscription cannot be user-managed")
 	}
 
-	proToken, err := contracts.ProToken(ctx)
+	_, src, err := c.Subscription(ctx)
+	if err != nil {
+		return fmt.Errorf("could not get current subscription status: %v", err)
+	}
+
+	// Shortcut to avoid spamming the contract server
+	// We don't need to request a new token if we have a non-expired one
+	if src == SourceMicrosoftStore {
+		expired, err := contracts.Expired()
+		if err != nil {
+			return fmt.Errorf("could not obtain current subscription expiration: %v", err)
+		}
+
+		if !expired {
+			log.Debug(ctx, "Microsoft Store subscription is active")
+			return nil
+		}
+
+		log.Debug(ctx, "Microsoft Store subscription is expired")
+	}
+
+	proToken, err := contracts.NewProToken(ctx)
 	if err != nil {
 		return fmt.Errorf("could not get ProToken from Microsoft Store: %v", err)
 	}
 
+	if proToken != "" {
+		log.Debugf(ctx, "Obtained Store-backed pro token: %q", common.Obfuscate(proToken))
+	}
+
 	if err := c.setStoreSubscription(ctx, proToken); err != nil {
 		return err
 	}
diff --git a/windows-agent/internal/contracts/contracts.go b/windows-agent/internal/contracts/contracts.go
index aa7794bd3..fb9016ebb 100644
--- a/windows-agent/internal/contracts/contracts.go
+++ b/windows-agent/internal/contracts/contracts.go
@@ -10,7 +10,6 @@ import (
 
 	"github.com/canonical/ubuntu-pro-for-windows/storeapi/go-wrapper/microsoftstore"
 	"github.com/canonical/ubuntu-pro-for-windows/windows-agent/internal/contracts/contractclient"
-	"github.com/ubuntu/decorate"
 )
 
 type options struct {
@@ -52,10 +51,31 @@ func (msftStoreDLL) GetSubscriptionExpirationDate() (tm time.Time, err error) {
 	return microsoftstore.GetSubscriptionExpirationDate()
 }
 
-// ProToken directs the dance between the Microsoft Store and the Ubuntu Pro contract server to
+func Expired(args ...Option) (bool, error) {
+	opts := options{
+		microsoftStore: msftStoreDLL{},
+	}
+
+	for _, f := range args {
+		f(&opts)
+	}
+
+	expiration, err := opts.microsoftStore.GetSubscriptionExpirationDate()
+	if err != nil {
+		return false, fmt.Errorf("could not get subscription expiration date: %v", err)
+	}
+
+	if expiration.After(time.Now()) {
+		return false, nil
+	}
+
+	return true, nil
+}
+
+// NewProToken directs the dance between the Microsoft Store and the Ubuntu Pro contract server to
 // validate a store entitlement and obtain its associated pro token. If there is no entitlement,
 // the token is returned as an empty string.
-func ProToken(ctx context.Context, args ...Option) (token string, err error) {
+func NewProToken(ctx context.Context, args ...Option) (token string, err error) {
 	opts := options{
 		microsoftStore: msftStoreDLL{},
 	}
@@ -73,28 +93,9 @@ func ProToken(ctx context.Context, args ...Option) (token string, err error) {
 	}
 
 	contractClient := contractclient.New(opts.proURL, &http.Client{Timeout: 30 * time.Second})
+	msftStore := opts.microsoftStore
 
-	token, err = proToken(ctx, contractClient, opts.microsoftStore)
-	if err != nil {
-		return "", err
-	}
-
-	return token, nil
-}
-
-func proToken(ctx context.Context, serverClient *contractclient.Client, msftStore MicrosoftStore) (proToken string, err error) {
-	defer decorate.OnError(&err, "could not obtain pro token")
-
-	expiration, err := msftStore.GetSubscriptionExpirationDate()
-	if err != nil {
-		return "", fmt.Errorf("could not get subscription expiration date: %v", err)
-	}
-
-	if expiration.Before(time.Now()) {
-		return "", fmt.Errorf("the subscription has been expired since %s", expiration)
-	}
-
-	adToken, err := serverClient.GetServerAccessToken(ctx)
+	adToken, err := contractClient.GetServerAccessToken(ctx)
 	if err != nil {
 		return "", err
 	}
@@ -104,7 +105,7 @@ func proToken(ctx context.Context, serverClient *contractclient.Client, msftStor
 		return "", err
 	}
 
-	proToken, err = serverClient.GetProToken(ctx, storeToken)
+	proToken, err := contractClient.GetProToken(ctx, storeToken)
 	if err != nil {
 		return "", err
 	}