PAPI: auto enable on upgrade (#3659)

Author: blotus

cmd/crowdsec-cli/clicapi/capi.go

@@ -2,7 +2,6 @@ package clicapi
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"io"
 	"net/url"
@@ -155,34 +154,36 @@ func (cli *cliCapi) newRegisterCmd() *cobra.Command {
 	return cmd
 }
 
+type capiStatus struct {
+	authenticated    bool
+	enrolled         bool
+	subscriptionType string
+}
+
 // queryCAPIStatus checks if the Central API is reachable, and if the credentials are correct. It then checks if the instance is enrolled in the console.
-func queryCAPIStatus(ctx context.Context, db *database.Client, hub *cwhub.Hub, credURL string, login string, password string) (bool, bool, error) {
+func queryCAPIStatus(ctx context.Context, db *database.Client, hub *cwhub.Hub, credURL string, login string, password string) (capiStatus, error) {
 	apiURL, err := url.Parse(credURL)
 	if err != nil {
-		return false, false, err
+		return capiStatus{}, err
 	}
 
 	itemsForAPI := hub.GetInstalledListForAPI()
 
-	if len(itemsForAPI) == 0 {
-		return false, false, errors.New("no scenarios or appsec-rules installed, abort")
-	}
-
 	passwd := strfmt.Password(password)
 
 	client, err := apiclient.NewClient(&apiclient.Config{
 		MachineID: login,
 		Password:  passwd,
 		URL:       apiURL,
-		// I don't believe papi is neede to check enrollement
+		// I don't believe papi is needed to check enrollement
 		// PapiURL:       papiURL,
 		VersionPrefix: "v3",
 		UpdateScenario: func(_ context.Context) ([]string, error) {
 			return itemsForAPI, nil
 		},
 	})
 	if err != nil {
-		return false, false, err
+		return capiStatus{}, err
 	}
 
 	pw := strfmt.Password(password)
@@ -195,20 +196,20 @@ func queryCAPIStatus(ctx context.Context, db *database.Client, hub *cwhub.Hub, c
 
 	authResp, _, err := client.Auth.AuthenticateWatcher(ctx, t)
 	if err != nil {
-		return false, false, err
+		return capiStatus{}, err
 	}
 
 	if err := db.SaveAPICToken(ctx, apiclient.TokenDBField, authResp.Token); err != nil {
-		return false, false, err
+		return capiStatus{}, err
 	}
 
 	client.GetClient().Transport.(*apiclient.JWTTransport).Token = authResp.Token
 
 	if client.IsEnrolled() {
-		return true, true, nil
+		return capiStatus{true, true, client.GetSubscriptionType()}, nil
 	}
 
-	return true, false, nil
+	return capiStatus{true, false, ""}, nil
 }
 
 func (cli *cliCapi) Status(ctx context.Context, db *database.Client, out io.Writer, hub *cwhub.Hub) error {
@@ -223,17 +224,18 @@ func (cli *cliCapi) Status(ctx context.Context, db *database.Client, out io.Writ
 	fmt.Fprintf(out, "Loaded credentials from %s\n", cfg.API.Server.OnlineClient.CredentialsFilePath)
 	fmt.Fprintf(out, "Trying to authenticate with username %s on %s\n", cred.Login, cred.URL)
 
-	auth, enrolled, err := queryCAPIStatus(ctx, db, hub, cred.URL, cred.Login, cred.Password)
+	status, err := queryCAPIStatus(ctx, db, hub, cred.URL, cred.Login, cred.Password)
 	if err != nil {
 		return fmt.Errorf("failed to authenticate to Central API (CAPI): %w", err)
 	}
 
-	if auth {
+	if status.authenticated {
 		fmt.Fprint(out, "You can successfully interact with Central API (CAPI)\n")
 	}
 
-	if enrolled {
+	if status.enrolled {
 		fmt.Fprint(out, "Your instance is enrolled in the console\n")
+		fmt.Fprintf(out, "Subscription type: %s\n", status.subscriptionType)
 	}
 
 	switch *cfg.API.Server.OnlineClient.Sharing {

pkg/apiclient/auth_jwt.go

@@ -30,15 +30,19 @@ type JWTTransport struct {
 	RetryConfig   *RetryConfig
 	// Transport is the underlying HTTP transport to use when making requests.
 	// It will default to http.DefaultTransport if nil.
-	Transport         http.RoundTripper
-	UpdateScenario    func(context.Context) ([]string, error)
+	Transport        http.RoundTripper
+	UpdateScenario   func(context.Context) ([]string, error)
+	TokenRefreshChan chan struct{} // will write to this channel when the token is refreshed
+
 	refreshTokenMutex sync.Mutex
 	TokenSave         TokenSave
 }
 
 func (t *JWTTransport) refreshJwtToken(ctx context.Context) error {
 	var err error
 
+	log.Debugf("refreshing jwt token for '%s'", *t.MachineID)
+
 	if t.UpdateScenario != nil {
 		t.Scenarios, err = t.UpdateScenario(ctx)
 		if err != nil {
@@ -142,6 +146,12 @@ func (t *JWTTransport) refreshJwtToken(ctx context.Context) error {
 
 	log.Debugf("token %s will expire on %s", t.Token, t.Expiration.String())
 
+	select {
+	case t.TokenRefreshChan <- struct{}{}:
+	default:
+		// Do not block if no one is waiting for the token refresh (ie, PAPI fully disabled)
+	}
+
 	return nil
 }
 

pkg/apiclient/client.go

@@ -69,6 +69,29 @@ func (c *ApiClient) IsEnrolled() bool {
 	return ok
 }
 
+func (c *ApiClient) GetSubscriptionType() string {
+	jwtTransport := c.client.Transport.(*JWTTransport)
+	tokenStr := jwtTransport.Token
+
+	token, _ := jwt.Parse(tokenStr, nil)
+	if token == nil {
+		return ""
+	}
+
+	claims := token.Claims.(jwt.MapClaims)
+	subscriptionType, ok := claims["subscription_type"].(string)
+	if ok {
+		return subscriptionType
+	}
+
+	return ""
+}
+
+func (c *ApiClient) GetTokenRefreshChan() chan struct{} {
+	jwtTransport := c.client.Transport.(*JWTTransport)
+	return jwtTransport.TokenRefreshChan
+}
+
 type service struct {
 	client *ApiClient
 }
@@ -149,7 +172,8 @@ func NewClient(config *Config) (*ApiClient, error) {
 			WithStatusCodeConfig(http.StatusServiceUnavailable, 5, true, false),
 			WithStatusCodeConfig(http.StatusGatewayTimeout, 5, true, false),
 		),
-		TokenSave: config.TokenSave,
+		TokenSave:        config.TokenSave,
+		TokenRefreshChan: make(chan struct{}),
 	}
 
 	transport, baseURL := createTransport(config.URL)

pkg/apiclient/subscription_types.go

@@ -0,0 +1,7 @@
+package apiclient
+
+const (
+	SubscriptionTypeEnterprise = "ENTERPRISE"
+	SubscriptionTypeSecOps     = "SECOPS"
+	SubscriptionTypeCommunity  = "COMMUNITY"
+)

pkg/apiserver/apiserver.go

@@ -260,19 +260,17 @@ func NewServer(ctx context.Context, config *csconfig.LocalApiServerCfg) (*APISer
 
 		controller.AlertsAddChan = apiClient.AlertsAddChan
 
-		if config.ConsoleConfig.IsPAPIEnabled() && config.OnlineClient.Credentials.PapiURL != "" {
-			if apiClient.apiClient.IsEnrolled() {
-				log.Info("Machine is enrolled in the console, Loading PAPI Client")
+		if apiClient.apiClient.IsEnrolled() {
+			log.Info("Machine is enrolled in the console, Loading PAPI Client")
 
-				papiClient, err = NewPAPI(apiClient, dbClient, config.ConsoleConfig, *config.PapiLogLevel)
-				if err != nil {
-					return nil, err
-				}
-
-				controller.DecisionDeleteChan = papiClient.Channels.DeleteDecisionChannel
-			} else {
-				log.Error("Machine is not enrolled in the console, can't synchronize with the console")
+			papiClient, err = NewPAPI(apiClient, dbClient, config.ConsoleConfig, *config.PapiLogLevel)
+			if err != nil {
+				return nil, err
 			}
+
+			controller.DecisionDeleteChan = papiClient.Channels.DeleteDecisionChannel
+		} else {
+			log.Error("Machine is not enrolled in the console, can't synchronize with the console")
 		}
 	}
 
@@ -343,9 +341,8 @@ func (s *APIServer) initAPIC(ctx context.Context) {
 	s.apic.pushTomb.Go(func() error { return s.apicPush(ctx) })
 	s.apic.pullTomb.Go(func() error { return s.apicPull(ctx) })
 
-	// csConfig.API.Server.ConsoleConfig.ShareCustomScenarios
 	if s.apic.apiClient.IsEnrolled() {
-		if s.consoleConfig.IsPAPIEnabled() && s.papi != nil {
+		if s.papi != nil {
 			if s.papi.URL != "" {
 				log.Info("Starting PAPI decision receiver")
 				s.papi.pullTomb.Go(func() error { return s.papiPull(ctx) })

pkg/apiserver/papi.go

@@ -258,32 +258,66 @@ func (p *Papi) Pull(ctx context.Context) error {
 		}
 	}
 
-	p.Logger.Infof("Starting PAPI pull (since:%s)", lastTimestamp)
+	tokenRefreshChan := p.apiClient.GetTokenRefreshChan()
+	var papiChan chan longpollclient.Event // Chan is nil by default to block until PAPI actually establishes the connection
+	papiCtx, cancel := context.WithCancel(ctx)
 
-	for event := range p.Client.Start(ctx, lastTimestamp) {
-		logger := p.Logger.WithField("request-id", event.RequestId)
-		// update last timestamp in database
-		newTime := time.Now().UTC()
+	currentSubscriptionType := p.apiClient.GetSubscriptionType()
 
-		binTime, err := newTime.MarshalText()
-		if err != nil {
-			return fmt.Errorf("failed to serialize last timestamp: %w", err)
-		}
+	p.Logger.Debugf("current subscription type is %s", currentSubscriptionType)
 
-		err = p.handleEvent(event, false)
-		if err != nil {
-			logger.Errorf("failed to handle event: %s", err)
-			continue
-		}
+	if currentSubscriptionType == apiclient.SubscriptionTypeEnterprise || currentSubscriptionType == apiclient.SubscriptionTypeSecOps {
+		// If allowed to use PAPI, start it
+		// Otherwise it will be started when the token is refreshed with an ent subscription
+		p.Logger.Infof("Starting PAPI pull (since:%s)", lastTimestamp)
+		papiChan = p.Client.Start(papiCtx, lastTimestamp)
+	}
 
-		if err := p.DBClient.SetConfigItem(ctx, PapiPullKey, string(binTime)); err != nil {
-			return fmt.Errorf("failed to update last timestamp: %w", err)
-		}
+	for {
+		select {
+		case <-tokenRefreshChan:
+			subType := p.apiClient.GetSubscriptionType()
+			if subType == currentSubscriptionType {
+				continue
+			}
+			currentSubscriptionType = subType
+			p.Logger.Infof("Subscription type changed to %s", subType)
+			switch subType {
+			case apiclient.SubscriptionTypeEnterprise, apiclient.SubscriptionTypeSecOps:
+				p.Logger.Infof("Starting PAPI pull (since:%s)", lastTimestamp)
+				papiChan = p.Client.Start(papiCtx, lastTimestamp)
+			default:
+				// PAPI got started but the user downgraded (or removed the engine from the console)
+				p.Logger.Info("Stopping PAPI because of plan downgrade or engine removal")
+				cancel() // This will stop any ongoing PAPI pull
+				p.Client.Stop()
+				papiCtx, cancel = context.WithCancel(ctx) //nolint:fatcontext // Recreate the context if the pull is restarted
+				papiChan = nil
+				p.Logger.Debug("done stopping PAPI pull")
+			}
+		case event := <-papiChan:
+			logger := p.Logger.WithField("request-id", event.RequestId)
+			// update last timestamp in database
+			newTime := time.Now().UTC()
 
-		logger.Debugf("set last timestamp to %s", newTime)
-	}
+			binTime, _ := newTime.MarshalText() // No need to check the error, time.Now().UTC() always returns a valid time
 
-	return nil
+			lastTimestamp = newTime
+
+			err = p.handleEvent(event, false)
+			if err != nil {
+				logger.Errorf("failed to handle event: %s", err)
+				continue
+			}
+
+			if err := p.DBClient.SetConfigItem(ctx, PapiPullKey, string(binTime)); err != nil {
+				// Killing PAPI is overkill if we cannot update the last timestamp
+				logger.Errorf("failed to update last timestamp in database: %s", err)
+			}
+
+			logger.Debugf("set last timestamp to %s", newTime)
+		}
+	}
 }
 
 func (p *Papi) SyncDecisions(ctx context.Context) error {

pkg/csconfig/api.go

@@ -20,6 +20,7 @@ import (
 	"github.com/crowdsecurity/go-cs-lib/yamlpatch"
 
 	"github.com/crowdsecurity/crowdsec/pkg/apiclient"
+	"github.com/crowdsecurity/crowdsec/pkg/types"
 )
 
 type APICfg struct {
@@ -123,6 +124,10 @@ func (o *OnlineApiClientCfg) Load() error {
 		o.Credentials = nil
 	}
 
+	if o.Credentials != nil && o.Credentials.PapiURL == "" {
+		o.Credentials.PapiURL = types.PAPIBaseURL
+	}
+
 	return nil
 }
 

pkg/csconfig/api_test.go

@@ -13,6 +13,8 @@ import (
 
 	"github.com/crowdsecurity/go-cs-lib/cstest"
 	"github.com/crowdsecurity/go-cs-lib/ptr"
+
+	"github.com/crowdsecurity/crowdsec/pkg/types"
 )
 
 func TestLoadLocalApiClientCfg(t *testing.T) {
@@ -93,6 +95,7 @@ func TestLoadOnlineApiClientCfg(t *testing.T) {
 				URL:      "http://crowdsec.api",
 				Login:    "test",
 				Password: "testpassword",
+				PapiURL:  types.PAPIBaseURL,
 			},
 		},
 		{
@@ -211,6 +214,7 @@ func TestLoadAPIServer(t *testing.T) {
 						URL:      "http://crowdsec.api",
 						Login:    "test",
 						Password: "testpassword",
+						PapiURL:  types.PAPIBaseURL,
 					},
 					Sharing: ptr.Of(true),
 					PullConfig: CapiPullConfig{