Token renewal in a interceptor

Author: majst01

generate/go_client.tpl

@@ -56,15 +56,16 @@ func New(config *DialConfig) (Client, error) {
 		authInterceptor := &authInterceptor{config: config}
 		c.interceptors = append(c.interceptors, authInterceptor)
 	}
+	if config.TokenRenewal != nil {
+		tokenRenewingInterceptor := &tokenRenewingInterceptor{config: config, client: c}
+		c.interceptors = append(c.interceptors, tokenRenewingInterceptor)
+	}
 	if config.Log != nil {
 		loggingInterceptor := &loggingInterceptor{config: config}
 		c.interceptors = append(c.interceptors, loggingInterceptor)
 	}
 	c.interceptors = append(c.interceptors, config.Interceptors...)
 
-	// TODO convert to interceptor
-	go c.startTokenRenewal()
-
 	return c, nil
 }
 

go/client/client-interceptors.go

@@ -2,8 +2,14 @@ package client
 
 import (
 	"context"
+	"fmt"
+	"log/slog"
+	"sync"
+	"sync/atomic"
+	"time"
 
 	"connectrpc.com/connect"
+	apiv2models "github.com/metal-stack/api/go/metalstack/api/v2"
 )
 
 // authinterceptor adds the required auth headers
@@ -65,3 +71,79 @@ func (i *loggingInterceptor) WrapStreamingClient(next connect.StreamingClientFun
 func (i *loggingInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
 	return next
 }
+
+type tokenRenewingInterceptor struct {
+	config *DialConfig
+	client *client
+
+	renewing atomic.Bool
+
+	sync.Mutex
+}
+
+func (i *tokenRenewingInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
+	return connect.UnaryFunc(func(ctx context.Context, request connect.AnyRequest) (connect.AnyResponse, error) {
+		err := i.renewTokenIfNeeded()
+		if err != nil {
+			return nil, err
+		}
+		return next(ctx, request)
+	})
+}
+
+func (i *tokenRenewingInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
+	return next
+}
+
+func (i *tokenRenewingInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
+	return next
+}
+
+func (i *tokenRenewingInterceptor) renewTokenIfNeeded() error {
+	if i.config.expiresAt.IsZero() {
+		return nil
+	}
+	if i.renewing.Load() {
+		return nil
+	}
+	if i.config.Log == nil {
+		i.config.Log = slog.Default()
+	}
+
+	replaceBefore := i.config.expiresAt.Sub(i.config.issuedAt) / tokenRenewChecksDuringLifetime
+
+	if time.Until(i.config.expiresAt) > replaceBefore {
+		return nil
+	}
+
+	i.renewing.Store(true)
+	defer i.renewing.Store(false)
+
+	i.config.Log.Info("call token refresh, current token expires soon", "expires", i.config.expiresAt.String())
+
+	i.Lock()
+	defer i.Unlock()
+
+	resp, err := i.client.Apiv2().Token().Refresh(context.Background(), &apiv2models.TokenServiceRefreshRequest{})
+	if err != nil {
+		return fmt.Errorf("unable to refresh token %w", err)
+	}
+
+	i.config.Token = resp.Secret
+	err = i.config.parse()
+	if err != nil {
+		return fmt.Errorf("unable to parse token %w", err)
+	}
+
+	if i.config.TokenRenewal.PersistTokenFn == nil {
+		return nil
+	}
+
+	err = i.config.TokenRenewal.PersistTokenFn(i.config.Token)
+	if err != nil {
+		return fmt.Errorf("unable to persist token %w", err)
+	}
+
+	i.config.Log.Info("token refreshed, new token expires in", "expires", i.config.expiresAt.String())
+	return nil
+}

go/client/client.go

@@ -1,7 +1,6 @@
 package client
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -10,7 +9,6 @@ import (
 
 	"connectrpc.com/connect"
 	"github.com/golang-jwt/jwt/v5"
-	api "github.com/metal-stack/api/go/metalstack/api/v2"
 )
 
 const tokenRenewChecksDuringLifetime = 4
@@ -84,70 +82,3 @@ func (dc *DialConfig) parse() error {
 	}
 	return nil
 }
-
-func (c *client) startTokenRenewal() {
-	if c.config.TokenRenewal == nil {
-		return
-	}
-	if c.config.expiresAt.IsZero() {
-		return
-	}
-	if c.config.Log == nil {
-		c.config.Log = slog.Default()
-	}
-
-	replaceBefore := c.config.expiresAt.Sub(c.config.issuedAt) / tokenRenewChecksDuringLifetime
-
-	err := c.renewTokenIfNeeded(replaceBefore)
-	if err != nil {
-		c.config.Log.Error("unable to renew token", "error", err)
-	}
-
-	ticker := time.NewTicker(replaceBefore)
-	defer ticker.Stop()
-	done := make(chan bool)
-	for {
-		select {
-		case <-done:
-			return
-		case <-ticker.C:
-			err := c.renewTokenIfNeeded(replaceBefore)
-			if err != nil {
-				c.config.Log.Error("unable to renew token", "error", err)
-			}
-		}
-	}
-}
-
-func (c *client) renewTokenIfNeeded(replaceBefore time.Duration) error {
-	if time.Until(c.config.expiresAt) > replaceBefore {
-		return nil
-	}
-	c.config.Log.Info("call token refresh, current token expires soon", "expires", c.config.expiresAt.String())
-
-	c.Lock()
-	defer c.Unlock()
-
-	resp, err := c.Apiv2().Token().Refresh(context.Background(), &api.TokenServiceRefreshRequest{})
-	if err != nil {
-		return fmt.Errorf("unable to refresh token %w", err)
-	}
-
-	c.config.Token = resp.Secret
-	err = c.config.parse()
-	if err != nil {
-		return fmt.Errorf("unable to parse token %w", err)
-	}
-
-	if c.config.TokenRenewal.PersistTokenFn == nil {
-		return nil
-	}
-
-	err = c.config.TokenRenewal.PersistTokenFn(c.config.Token)
-	if err != nil {
-		return fmt.Errorf("unable to persist token %w", err)
-	}
-
-	c.config.Log.Info("token refreshed, new token expires in", "expires", c.config.expiresAt.String())
-	return nil
-}