feat(apiclient): add token save functionality (#3639)

Author: sabban

cmd/crowdsec-cli/clicapi/capi.go

@@ -198,7 +198,7 @@ func queryCAPIStatus(ctx context.Context, db *database.Client, hub *cwhub.Hub, c
 		return false, false, err
 	}
 
-	if err := db.SaveAPICToken(ctx, authResp.Token); err != nil {
+	if err := db.SaveAPICToken(ctx, apiclient.TokenDBField, authResp.Token); err != nil {
 		return false, false, err
 	}
 

pkg/apiclient/auth_jwt.go

@@ -33,6 +33,7 @@ type JWTTransport struct {
 	Transport         http.RoundTripper
 	UpdateScenario    func(context.Context) ([]string, error)
 	refreshTokenMutex sync.Mutex
+	TokenSave         TokenSave
 }
 
 func (t *JWTTransport) refreshJwtToken() error {
@@ -134,6 +135,12 @@ func (t *JWTTransport) refreshJwtToken() error {
 
 	t.Token = response.Token
 
+	if t.TokenSave != nil {
+		err = t.TokenSave(ctx, TokenDBField, t.Token)
+		if err != nil {
+			log.Errorf("unable to save token: %s", err)
+		}
+	}
 	log.Debugf("token %s will expire on %s", t.Token, t.Expiration.String())
 
 	return nil

pkg/apiclient/client.go

@@ -26,6 +26,8 @@ var (
 	lapiClient         *ApiClient
 )
 
+type TokenSave func(ctx context.Context, tokenKey string, token string) error
+
 type ApiClient struct {
 	/*The http client used to make requests*/
 	client *http.Client
@@ -147,6 +149,7 @@ func NewClient(config *Config) (*ApiClient, error) {
 			WithStatusCodeConfig(http.StatusServiceUnavailable, 5, true, false),
 			WithStatusCodeConfig(http.StatusGatewayTimeout, 5, true, false),
 		),
+		TokenSave: config.TokenSave,
 	}
 
 	transport, baseURL := createTransport(config.URL)

pkg/apiclient/config.go

@@ -7,6 +7,8 @@ import (
 	"github.com/go-openapi/strfmt"
 )
 
+const TokenDBField = "apic_token"
+
 type Config struct {
 	MachineID         string
 	Password          strfmt.Password
@@ -16,4 +18,5 @@ type Config struct {
 	UserAgent         string
 	RegistrationToken string
 	UpdateScenario    func(context.Context) ([]string, error)
+	TokenSave         func(context.Context, string, string) error
 }

pkg/apiserver/apic.go

@@ -74,6 +74,8 @@ type apic struct {
 	pullBlocklists bool
 	pullCommunity  bool
 	shareSignals   bool
+
+	TokenSave apiclient.TokenSave
 }
 
 // randomDuration returns a duration value between d-delta and d+delta
@@ -228,6 +230,9 @@ func NewAPIC(ctx context.Context, config *csconfig.OnlineApiClientCfg, dbClient
 		PapiURL:        papiURL,
 		VersionPrefix:  "v3",
 		UpdateScenario: ret.FetchScenariosListFromDB,
+		TokenSave: func(ctx context.Context, tokenKey string, token string) error {
+			return dbClient.SaveAPICToken(ctx, tokenKey, token)
+		},
 	})
 	if err != nil {
 		return nil, fmt.Errorf("while creating api client: %w", err)
@@ -277,7 +282,7 @@ func (a *apic) Authenticate(ctx context.Context, config *csconfig.OnlineApiClien
 
 	a.apiClient.GetClient().Transport.(*apiclient.JWTTransport).Token = authResp.Token
 
-	return a.dbClient.SaveAPICToken(ctx, authResp.Token)
+	return a.dbClient.SaveAPICToken(ctx, apiclient.TokenDBField, authResp.Token)
 }
 
 // keep track of all alerts in cache and push it to CAPI every PushInterval.

pkg/database/config.go

@@ -9,12 +9,11 @@ import (
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 
+	"github.com/crowdsecurity/crowdsec/pkg/apiclient"
 	"github.com/crowdsecurity/crowdsec/pkg/database/ent"
 	"github.com/crowdsecurity/crowdsec/pkg/database/ent/configitem"
 )
 
-const apicTokenKey = "apic_token"
-
 func (c *Client) GetConfigItem(ctx context.Context, key string) (string, error) {
 	result, err := c.Ent.ConfigItem.Query().Where(configitem.NameEQ(key)).First(ctx)
 
@@ -53,7 +52,7 @@ func (c *Client) SetConfigItem(ctx context.Context, key string, value string) er
 //   - it is a properly formatted JWT with an "exp" claim,
 //   - it is not expired or near expiry.
 func (c *Client) LoadAPICToken(ctx context.Context, logger logrus.FieldLogger) (string, time.Time, bool) {
-	token, err := c.GetConfigItem(ctx, apicTokenKey)
+	token, err := c.GetConfigItem(ctx, apiclient.TokenDBField) // TokenKey is a constant string representing the key for the token in the database
 	if err != nil {
 		logger.Debugf("error fetching token from DB: %s", err)
 		return "", time.Time{}, false
@@ -78,6 +77,18 @@ func (c *Client) LoadAPICToken(ctx context.Context, logger logrus.FieldLogger) (
 		return "", time.Time{}, false
 	}
 
+	iatFloat, ok := claims["iat"].(float64)
+	if !ok {
+		logger.Debug("token missing 'iat' claim")
+		return "", time.Time{}, false
+	}
+
+	iat := time.Unix(int64(iatFloat), 0)
+	if time.Now().UTC().After(iat.Add(1 * time.Minute)) {
+		logger.Debug("token is more than 1 minute old, not using it")
+		return "", time.Time{}, false
+	}
+
 	expFloat, ok := claims["exp"].(float64)
 	if !ok {
 		logger.Debug("token missing 'exp' claim")
@@ -94,8 +105,8 @@ func (c *Client) LoadAPICToken(ctx context.Context, logger logrus.FieldLogger) (
 }
 
 // SaveAPICToken stores the given JWT token in the local database under the appropriate config item.
-func (c *Client) SaveAPICToken(ctx context.Context, token string) error {
-	if err := c.SetConfigItem(ctx, apicTokenKey, token); err != nil {
+func (c *Client) SaveAPICToken(ctx context.Context, tokenKey string, token string) error {
+	if err := c.SetConfigItem(ctx, tokenKey, token); err != nil {
 		return fmt.Errorf("saving token to db: %w", err)
 	}