[RC] Support configuring the core RC service with JWT auth (#32430)

Co-authored-by: mellon85 <[email protected]>

pkg/config/remote/api/http.go

@@ -13,6 +13,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"sync"
 	"time"
 
 	"google.golang.org/protobuf/proto"
@@ -46,11 +47,13 @@ type API interface {
 	Fetch(context.Context, *pbgo.LatestConfigsRequest) (*pbgo.LatestConfigsResponse, error)
 	FetchOrgData(context.Context) (*pbgo.OrgDataResponse, error)
 	FetchOrgStatus(context.Context) (*pbgo.OrgStatusResponse, error)
+	UpdatePARJWT(string)
 }
 
 // Auth defines the possible Authentication data to access the RC backend
 type Auth struct {
 	APIKey    string
+	PARJWT    string
 	AppKey    string
 	UseAppKey bool
 }
@@ -59,18 +62,26 @@ type Auth struct {
 type HTTPClient struct {
 	baseURL string
 	client  *http.Client
-	header  http.Header
+
+	headerLock sync.RWMutex
+	header     http.Header
 }
 
 // NewHTTPClient returns a new HTTP configuration client
 func NewHTTPClient(auth Auth, cfg model.Reader, baseURL *url.URL) (*HTTPClient, error) {
 	header := http.Header{
 		"Content-Type": []string{"application/x-protobuf"},
-		"DD-Api-Key":   []string{auth.APIKey},
+	}
+	if auth.PARJWT != "" {
+		header["DD-PAR-JWT"] = []string{auth.PARJWT}
+	}
+	if auth.APIKey != "" {
+		header["DD-Api-Key"] = []string{auth.APIKey}
 	}
 	if auth.UseAppKey {
 		header["DD-Application-Key"] = []string{auth.AppKey}
 	}
+
 	transport := httputils.CreateHTTPTransport(cfg)
 	// Set the keep-alive timeout to 30s instead of the default 90s, so the http RC client is not closed by the backend
 	transport.IdleConnTimeout = 30 * time.Second
@@ -104,7 +115,8 @@ func (c *HTTPClient) Fetch(ctx context.Context, request *pbgo.LatestConfigsReque
 	if err != nil {
 		return nil, fmt.Errorf("failed to create org data request: %w", err)
 	}
-	req.Header = c.header
+
+	c.addHeaders(req)
 
 	resp, err := c.client.Do(req)
 	if err != nil {
@@ -150,7 +162,8 @@ func (c *HTTPClient) FetchOrgData(ctx context.Context) (*pbgo.OrgDataResponse, e
 	if err != nil {
 		return nil, fmt.Errorf("failed to create org data request: %w", err)
 	}
-	req.Header = c.header
+
+	c.addHeaders(req)
 
 	resp, err := c.client.Do(req)
 	if err != nil {
@@ -187,7 +200,8 @@ func (c *HTTPClient) FetchOrgStatus(ctx context.Context) (*pbgo.OrgStatusRespons
 	if err != nil {
 		return nil, fmt.Errorf("failed to create org data request: %w", err)
 	}
-	req.Header = c.header
+
+	c.addHeaders(req)
 
 	resp, err := c.client.Do(req)
 	if err != nil {
@@ -216,6 +230,22 @@ func (c *HTTPClient) FetchOrgStatus(ctx context.Context) (*pbgo.OrgStatusRespons
 	return response, err
 }
 
+// UpdatePARJWT allows for dynamic setting of a Private Action Runners JWT
+// Token for authentication to the RC backend.
+func (c *HTTPClient) UpdatePARJWT(jwt string) {
+	c.headerLock.Lock()
+	c.header.Set("DD-PAR-JWT", jwt)
+	c.headerLock.Unlock()
+}
+
+func (c *HTTPClient) addHeaders(req *http.Request) {
+	c.headerLock.RLock()
+	for k, v := range c.header {
+		req.Header[k] = v
+	}
+	c.headerLock.RUnlock()
+}
+
 func checkStatusCode(resp *http.Response) error {
 	// Specific case: authentication method is wrong
 	// we want to be descriptive about what can be done

pkg/config/remote/service/service.go

@@ -220,6 +220,7 @@ type options struct {
 	site                           string
 	rcKey                          string
 	apiKey                         string
+	parJWT                         string
 	traceAgentEnv                  string
 	databaseFileName               string
 	databaseFilePath               string
@@ -236,6 +237,7 @@ type options struct {
 var defaultOptions = options{
 	rcKey:                          "",
 	apiKey:                         "",
+	parJWT:                         "",
 	traceAgentEnv:                  "",
 	databaseFileName:               "remote-config.db",
 	databaseFilePath:               "",
@@ -327,6 +329,11 @@ func WithAPIKey(apiKey string) func(s *options) {
 	return func(s *options) { s.apiKey = apiKey }
 }
 
+// WithPARJWT sets the JWT for the private action runner
+func WithPARJWT(jwt string) func(s *options) {
+	return func(s *options) { s.parJWT = jwt }
+}
+
 // WithClientCacheBypassLimit validates and sets the service client cache bypass limit
 func WithClientCacheBypassLimit(limit int, cfgPath string) func(s *options) {
 	if limit < minCacheBypassLimit || limit > maxCacheBypassLimit {
@@ -387,7 +394,7 @@ func NewService(cfg model.Reader, rcType, baseRawURL, hostname string, tagsGette
 	backoffPolicy := backoff.NewExpBackoffPolicy(minBackoffFactor, baseBackoffTime,
 		options.maxBackoff.Seconds(), recoveryInterval, recoveryReset)
 
-	authKeys, err := getRemoteConfigAuthKeys(options.apiKey, options.rcKey)
+	authKeys, err := getRemoteConfigAuthKeys(options.apiKey, options.rcKey, options.parJWT)
 	if err != nil {
 		return nil, err
 	}
@@ -504,6 +511,12 @@ func (s *CoreAgentService) Start() {
 	}()
 }
 
+// UpdatePARJWT updates the stored JWT for Private Action Runners
+// for authentication to the remote config backend.
+func (s *CoreAgentService) UpdatePARJWT(jwt string) {
+	s.api.UpdatePARJWT(jwt)
+}
+
 func startWithAgentPollLoop(s *CoreAgentService) {
 	err := s.refresh()
 	if err != nil {

pkg/config/remote/service/service_test.go

@@ -69,6 +69,10 @@ func (m *mockAPI) FetchOrgStatus(ctx context.Context) (*pbgo.OrgStatusResponse,
 	return args.Get(0).(*pbgo.OrgStatusResponse), args.Error(1)
 }
 
+func (m *mockAPI) UpdatePARJWT(jwt string) {
+	m.Called(jwt)
+}
+
 type mockUptane struct {
 	mock.Mock
 }

pkg/config/remote/service/util.go

@@ -139,13 +139,16 @@ func openCacheDB(path string, agentVersion string, apiKey string) (*bbolt.DB, er
 type remoteConfigAuthKeys struct {
 	apiKey string
 
+	parJWT string
+
 	rcKeySet bool
 	rcKey    *msgpgo.RemoteConfigKey
 }
 
 func (k *remoteConfigAuthKeys) apiAuth() api.Auth {
 	auth := api.Auth{
 		APIKey: k.apiKey,
+		PARJWT: k.parJWT,
 	}
 	if k.rcKeySet {
 		auth.UseAppKey = true
@@ -154,12 +157,15 @@ func (k *remoteConfigAuthKeys) apiAuth() api.Auth {
 	return auth
 }
 
-func getRemoteConfigAuthKeys(apiKey string, rcKey string) (remoteConfigAuthKeys, error) {
+func getRemoteConfigAuthKeys(apiKey string, rcKey string, parJWT string) (remoteConfigAuthKeys, error) {
 	if rcKey == "" {
 		return remoteConfigAuthKeys{
 			apiKey: apiKey,
+			parJWT: parJWT,
 		}, nil
 	}
+
+	// Legacy auth with RC specific keys
 	rcKey = strings.TrimPrefix(rcKey, "DDRCM_")
 	encoding := base32.StdEncoding.WithPadding(base32.NoPadding)
 	rawKey, err := encoding.DecodeString(rcKey)
@@ -176,6 +182,7 @@ func getRemoteConfigAuthKeys(apiKey string, rcKey string) (remoteConfigAuthKeys,
 	}
 	return remoteConfigAuthKeys{
 		apiKey:   apiKey,
+		parJWT:   parJWT,
 		rcKeySet: true,
 		rcKey:    &key,
 	}, nil

pkg/config/remote/service/util_test.go

@@ -27,6 +27,7 @@ func TestAuthKeys(t *testing.T) {
 	tests := []struct {
 		rcKey  string
 		apiKey string
+		parJWT string
 		err    bool
 		output remoteConfigAuthKeys
 	}{
@@ -42,10 +43,17 @@ func TestAuthKeys(t *testing.T) {
 		{rcKey: generateKey(t, 2, "datadoghq.com", ""), err: true},
 		{rcKey: generateKey(t, 2, "", "app_Key"), err: true},
 		{rcKey: generateKey(t, 0, "datadoghq.com", "app_Key"), err: true},
+		{parJWT: "myJWT", err: false, output: remoteConfigAuthKeys{
+			parJWT: "myJWT",
+		}},
+		{parJWT: "myJWT", apiKey: "myAPIKey", err: false, output: remoteConfigAuthKeys{
+			parJWT: "myJWT",
+			apiKey: "myAPIKey",
+		}},
 	}
 	for _, test := range tests {
 		t.Run(fmt.Sprintf("%s|%s", test.apiKey, test.rcKey), func(tt *testing.T) {
-			output, err := getRemoteConfigAuthKeys(test.apiKey, test.rcKey)
+			output, err := getRemoteConfigAuthKeys(test.apiKey, test.rcKey, test.parJWT)
 			if test.err {
 				assert.Error(tt, err)
 			} else {