Fix MSI token refresh

Signed-off-by: Philip Laine <philip.laine@xenit.se>

pkg/objstore/azure/helpers.go

@@ -38,20 +38,31 @@ func getAzureStorageCredentials(conf Config) (blob.Credential, error) {
 	if conf.MSIResource != "" {
 		msiConfig := auth.NewMSIConfig()
 		msiConfig.Resource = conf.MSIResource
-
 		azureServicePrincipalToken, err := msiConfig.ServicePrincipalToken()
 		if err != nil {
 			return nil, err
 		}
 
-		// Get a new token.
+		var tokenRefresher blob.TokenRefresher = func(credential blob.TokenCredential) time.Duration {
+			err := azureServicePrincipalToken.Refresh()
+			if err != nil {
+				// Retry later as the error can be related to API throttling
+				return 30 * time.Second
+			}
+			token := azureServicePrincipalToken.Token()
+			credential.SetToken(token.AccessToken)
+			// Refresh token two minutes before actual expiration
+			exp := token.Expires().Sub(time.Now().Add(2 * time.Minute))
+			return exp
+		}
+
 		err = azureServicePrincipalToken.Refresh()
 		if err != nil {
 			return nil, err
 		}
-		token := azureServicePrincipalToken.Token()
+		initialToken := azureServicePrincipalToken.Token()
 
-		return blob.NewTokenCredential(token.AccessToken, nil), nil
+		return blob.NewTokenCredential(initialToken.AccessToken, tokenRefresher), nil
 	}
 
 	credential, err := blob.NewSharedKeyCredential(conf.StorageAccountName, conf.StorageAccountKey)