Added support for secrets into action data (#5482)

Author: michalpristas

internal/pkg/action/dispatcher.go

@@ -10,10 +10,12 @@ import (
 	"sync"
 	"time"
 
+	"github.com/elastic/fleet-server/v7/internal/pkg/bulk"
 	"github.com/elastic/fleet-server/v7/internal/pkg/es"
 	"github.com/elastic/fleet-server/v7/internal/pkg/logger/ecs"
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
 	"github.com/elastic/fleet-server/v7/internal/pkg/monitor"
+	"github.com/elastic/fleet-server/v7/internal/pkg/secret"
 	"github.com/elastic/fleet-server/v7/internal/pkg/sqn"
 
 	"github.com/rs/zerolog"
@@ -34,23 +36,25 @@ func (s Sub) Ch() chan []model.Action {
 
 // Dispatcher tracks agent subscriptions and emits actions to the subscriptions.
 type Dispatcher struct {
-	am    monitor.SimpleMonitor
-	limit *rate.Limiter
+	am     monitor.SimpleMonitor
+	limit  *rate.Limiter
+	bulker bulk.Bulk
 
 	mx   sync.RWMutex
 	subs map[string]Sub
 }
 
 // NewDispatcher creates a Dispatcher using the provided monitor.
-func NewDispatcher(am monitor.SimpleMonitor, throttle time.Duration, i int) *Dispatcher {
+func NewDispatcher(am monitor.SimpleMonitor, throttle time.Duration, i int, bulker bulk.Bulk) *Dispatcher {
 	r := rate.Inf
 	if throttle > 0 {
 		r = rate.Every(throttle)
 	}
 	return &Dispatcher{
-		am:    am,
-		limit: rate.NewLimiter(r, i),
-		subs:  make(map[string]Sub),
+		am:     am,
+		limit:  rate.NewLimiter(r, i),
+		subs:   make(map[string]Sub),
+		bulker: bulker,
 	}
 }
 
@@ -117,6 +121,16 @@ func (d *Dispatcher) process(ctx context.Context, hits []es.HitT) {
 			zerolog.Ctx(ctx).Error().Err(err).Msg("Failed to unmarshal action document")
 			break
 		}
+
+		secretData, err := secret.GetActionDataWithSecrets(ctx, action.Data, action.SecretReferences, d.bulker)
+		if err != nil {
+			zerolog.Ctx(ctx).Error().Err(err).Msg("Failed to replace secrets in action document")
+			break
+		}
+
+		action.Data = secretData
+		action.SecretReferences = nil
+
 		numAgents := len(action.Agents)
 		for i, agentID := range action.Agents {
 			arr := agentActions[agentID]

internal/pkg/action/dispatcher_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/elastic/fleet-server/v7/internal/pkg/es"
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
 	"github.com/elastic/fleet-server/v7/internal/pkg/sqn"
+	ftesting "github.com/elastic/fleet-server/v7/internal/pkg/testing"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"golang.org/x/time/rate"
@@ -24,7 +25,7 @@ type mockMonitor struct {
 
 func (m *mockMonitor) Output() <-chan []es.HitT {
 	args := m.Called()
-	return args.Get(0).(<-chan []es.HitT)
+	return args.Get(0).(<-chan []es.HitT) //nolint:errcheck // we don't need to check here
 }
 
 func (m *mockMonitor) Run(ctx context.Context) error {
@@ -34,15 +35,17 @@ func (m *mockMonitor) Run(ctx context.Context) error {
 
 func (m *mockMonitor) GetCheckpoint() sqn.SeqNo {
 	args := m.Called()
-	return args.Get(0).(sqn.SeqNo)
+	return args.Get(0).(sqn.SeqNo) //nolint:errcheck // we don't need to check here
 }
 
 func TestNewDispatcher(t *testing.T) {
 	m := &mockMonitor{}
-	d := NewDispatcher(m, 0, 0)
+	bulker := ftesting.NewMockBulk()
+	d := NewDispatcher(m, 0, 0, bulker)
 
 	assert.NotNil(t, d.am)
 	assert.NotNil(t, d.subs)
+	assert.NotNil(t, d.bulker)
 }
 
 func compareActions(t *testing.T, expects, results []model.Action) {
@@ -259,7 +262,7 @@ func Test_Dispatcher_Run(t *testing.T) {
 			}
 
 			now := time.Now()
-			ctx, cancel := context.WithCancel(context.Background())
+			ctx, cancel := context.WithCancel(t.Context())
 			defer cancel()
 			go func() {
 				err := d.Run(ctx)
@@ -366,7 +369,7 @@ func Test_offsetStartTime(t *testing.T) {
 	}}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			r := offsetStartTime(context.Background(), tt.start, tt.dur, tt.i, tt.total)
+			r := offsetStartTime(t.Context(), tt.start, tt.dur, tt.i, tt.total)
 			assert.Equal(t, tt.result, r)
 		})
 	}

internal/pkg/api/handleCheckin.go

@@ -30,6 +30,7 @@ import (
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
 	"github.com/elastic/fleet-server/v7/internal/pkg/monitor"
 	"github.com/elastic/fleet-server/v7/internal/pkg/policy"
+	"github.com/elastic/fleet-server/v7/internal/pkg/secret"
 	"github.com/elastic/fleet-server/v7/internal/pkg/sqn"
 
 	"github.com/hashicorp/go-version"
@@ -861,7 +862,7 @@ func processPolicy(ctx context.Context, zlog zerolog.Logger, bulker bulk.Bulk, a
 	data := model.ClonePolicyData(pp.Policy.Data)
 	for policyName, policyOutput := range data.Outputs {
 		// NOTE: Not sure if output secret keys collected here include new entries, but they are collected for completeness
-		ks, err := policy.ProcessOutputSecret(ctx, policyOutput, bulker) // makes a bulk request to get secret values
+		ks, err := secret.ProcessOutputSecret(ctx, policyOutput, bulker) // makes a bulk request to get secret values
 		if err != nil {
 			return nil, fmt.Errorf("failed to process output secrets %q: %w",
 				policyName, err)

internal/pkg/model/schema.go

@@ -14,13 +14,13 @@ import (
 
 	"github.com/elastic/fleet-server/v7/internal/pkg/bulk"
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
+	"github.com/elastic/fleet-server/v7/internal/pkg/secret"
 	"github.com/elastic/fleet-server/v7/internal/pkg/smap"
 )
 
 const (
 	FieldOutputs            = "outputs"
 	FieldOutputType         = "type"
-	FieldOutputSecrets      = "secrets"
 	FieldOutputFleetServer  = "fleet_server"
 	FieldOutputServiceToken = "service_token"
 	FieldOutputPermissions  = "output_permissions"
@@ -68,7 +68,7 @@ func NewParsedPolicy(ctx context.Context, bulker bulk.Bulk, p model.Policy) (*Pa
 		return nil, err
 	}
 	for name, policyOutput := range p.Data.Outputs {
-		ks, err := ProcessOutputSecret(ctx, policyOutput, bulker)
+		ks, err := secret.ProcessOutputSecret(ctx, policyOutput, bulker)
 		if err != nil {
 			return nil, err
 		}
@@ -80,7 +80,7 @@ func NewParsedPolicy(ctx context.Context, bulker bulk.Bulk, p model.Policy) (*Pa
 	if err != nil {
 		return nil, err
 	}
-	policyInputs, keys, err := getPolicyInputsWithSecrets(ctx, p.Data, bulker)
+	policyInputs, keys, err := secret.GetPolicyInputsWithSecrets(ctx, p.Data, bulker)
 	if err != nil {
 		return nil, err
 	}

internal/pkg/policy/parsed_policy.go

@@ -2,10 +2,11 @@
 // or more contributor license agreements. Licensed under the Elastic License;
 // you may not use this file except in compliance with the Elastic License.
 
-package policy
+package secret
 
 import (
 	"context"
+	"encoding/json"
 	"regexp"
 	"strconv"
 	"strings"
@@ -15,12 +16,16 @@ import (
 	"github.com/elastic/fleet-server/v7/internal/pkg/smap"
 )
 
+const (
+	FieldOutputSecrets = "secrets"
+)
+
 var (
 	secretRegex = regexp.MustCompile(`\$co\.elastic\.secret{([^}]*)}`)
 )
 
 // read secret values that belong to the agent policy's secret references, returns secrets as id:value map
-func getSecretValues(ctx context.Context, secretRefs []model.SecretReferencesItems, bulker bulk.Bulk) (map[string]string, error) {
+func GetSecretValues(ctx context.Context, secretRefs []model.SecretReferencesItems, bulker bulk.Bulk) (map[string]string, error) {
 	if len(secretRefs) == 0 {
 		return nil, nil
 	}
@@ -40,7 +45,7 @@ func getSecretValues(ctx context.Context, secretRefs []model.SecretReferencesIte
 
 // read inputs and secret_references from agent policy
 // replace values of secret refs in inputs and input streams properties
-func getPolicyInputsWithSecrets(ctx context.Context, data *model.PolicyData, bulker bulk.Bulk) ([]map[string]interface{}, []string, error) {
+func GetPolicyInputsWithSecrets(ctx context.Context, data *model.PolicyData, bulker bulk.Bulk) ([]map[string]interface{}, []string, error) {
 	if len(data.Inputs) == 0 {
 		return nil, nil, nil
 	}
@@ -49,7 +54,7 @@ func getPolicyInputsWithSecrets(ctx context.Context, data *model.PolicyData, bul
 		return data.Inputs, nil, nil
 	}
 
-	secretValues, err := getSecretValues(ctx, data.SecretReferences, bulker)
+	secretValues, err := GetSecretValues(ctx, data.SecretReferences, bulker)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -67,6 +72,31 @@ func getPolicyInputsWithSecrets(ctx context.Context, data *model.PolicyData, bul
 	return result, keys, nil
 }
 
+func GetActionDataWithSecrets(ctx context.Context, data json.RawMessage, refs []model.SecretReferencesItems, bulker bulk.Bulk) (json.RawMessage, error) {
+	if len(refs) == 0 {
+		return data, nil
+	}
+
+	secretValues, err := GetSecretValues(ctx, refs, bulker)
+	if err != nil {
+		return data, err
+	}
+
+	parsedData, err := smap.Parse(data)
+	if err != nil {
+		return data, err
+	}
+
+	result, _ := replaceMapRef(parsedData, secretValues)
+
+	b, err := json.Marshal(result)
+	if err != nil {
+		return data, err
+	}
+
+	return b, nil
+}
+
 // replaceMapRef replaces all nested secret values in the passed input and returns the resulting input along with a list of keys where inputs have been replaced.
 func replaceMapRef(input map[string]any, secrets map[string]string) (map[string]any, []string) {
 	keys := make([]string, 0)
@@ -207,7 +237,7 @@ func ProcessOutputSecret(ctx context.Context, output smap.Map, bulker bulk.Bulk)
 	if len(secretReferences) == 0 {
 		return nil, nil
 	}
-	secretValues, err := getSecretValues(ctx, secretReferences, bulker)
+	secretValues, err := GetSecretValues(ctx, secretReferences, bulker)
 	if err != nil {
 		return nil, err
 	}

internal/pkg/policy/secret.go renamed to internal/pkg/secret/secret.go

@@ -4,10 +4,11 @@
 
 //go:build !integration
 
-package policy
+package secret
 
 import (
 	"context"
+	"encoding/json"
 	"testing"
 
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
@@ -86,7 +87,7 @@ func TestGetSecretValues(t *testing.T) {
 	refs := []model.SecretReferencesItems{{ID: "ref1"}, {ID: "ref2"}}
 	bulker := ftesting.NewMockBulk()
 
-	secretRefs, _ := getSecretValues(context.TODO(), refs, bulker)
+	secretRefs, _ := GetSecretValues(context.TODO(), refs, bulker)
 
 	expectedRefs := map[string]string{
 		"ref1": "ref1_value",
@@ -95,6 +96,40 @@ func TestGetSecretValues(t *testing.T) {
 	assert.Equal(t, expectedRefs, secretRefs)
 }
 
+func TestGetActionDataWithSecrets(t *testing.T) {
+	refs := []model.SecretReferencesItems{
+		{ID: "ref1"},
+		{ID: "ref2"},
+	}
+	// Input JSON with secret references
+	input := map[string]interface{}{
+		"username": "user1",
+		"password": "$co.elastic.secret{ref1}",
+		"nested": map[string]interface{}{
+			"token": "$co.elastic.secret{ref2}",
+		},
+	}
+	b, err := json.Marshal(input)
+	require.NoError(t, err)
+
+	bulker := ftesting.NewMockBulk()
+	result, err := GetActionDataWithSecrets(t.Context(), b, refs, bulker)
+	require.NoError(t, err)
+
+	var out map[string]interface{}
+	err = json.Unmarshal(result, &out)
+	require.NoError(t, err)
+
+	assert.Equal(t, "user1", out["username"])
+	assert.Equal(t, "ref1_value", out["password"])
+
+	nestedMap, ok := out["nested"].(map[string]interface{})
+	assert.True(t, ok)
+
+	require.NoError(t, err)
+	assert.Equal(t, "ref2_value", nestedMap["token"])
+}
+
 func TestGetPolicyInputsWithSecretsAndStreams(t *testing.T) {
 	refs := []model.SecretReferencesItems{{ID: "ref1"}, {ID: "ref2"}, {ID: "ref3"}}
 	inputs := []map[string]interface{}{
@@ -126,7 +161,7 @@ func TestGetPolicyInputsWithSecretsAndStreams(t *testing.T) {
 		{"id": "input2", "streams": []interface{}{expectedStream}},
 	}
 
-	result, keys, _ := getPolicyInputsWithSecrets(context.TODO(), &pData, bulker)
+	result, keys, _ := GetPolicyInputsWithSecrets(context.TODO(), &pData, bulker)
 
 	assert.Equal(t, expectedResult, result)
 	assert.ElementsMatch(t, []string{"inputs.0.package_var_secret", "inputs.0.input_var_secret", "inputs.1.streams.0.package_var_secret", "inputs.1.streams.0.input_var_secret", "inputs.1.streams.0.stream_var_secret"}, keys)
@@ -173,7 +208,7 @@ func TestPolicyInputSteamsEmbedded(t *testing.T) {
 		}},
 	}
 
-	result, keys, err := getPolicyInputsWithSecrets(context.TODO(), &pData, bulker)
+	result, keys, err := GetPolicyInputsWithSecrets(context.TODO(), &pData, bulker)
 	require.NoError(t, err)
 
 	assert.Equal(t, expected, result)
@@ -201,7 +236,7 @@ func TestGetPolicyInputsNoopWhenNoSecrets(t *testing.T) {
 		{"id": "input2", "streams": []interface{}{expectedStream}},
 	}
 
-	result, keys, _ := getPolicyInputsWithSecrets(context.TODO(), &pData, bulker)
+	result, keys, _ := GetPolicyInputsWithSecrets(context.TODO(), &pData, bulker)
 
 	assert.Equal(t, expectedResult, result)
 	assert.Empty(t, keys)

internal/pkg/policy/secret_test.go renamed to internal/pkg/secret/secret_test.go

@@ -518,7 +518,7 @@ func (f *Fleet) runSubsystems(ctx context.Context, cfg *config.Config, g *errgro
 	}
 	g.Go(loggedRunFunc(ctx, "Action monitor", am.Run))
 
-	ad := action.NewDispatcher(am, cfg.Inputs[0].Server.Limits.ActionLimit.Interval, cfg.Inputs[0].Server.Limits.ActionLimit.Burst)
+	ad := action.NewDispatcher(am, cfg.Inputs[0].Server.Limits.ActionLimit.Interval, cfg.Inputs[0].Server.Limits.ActionLimit.Burst, bulker)
 	g.Go(loggedRunFunc(ctx, "Action dispatcher", ad.Run))
 
 	bc := checkin.NewBulk(bulker)

internal/pkg/server/fleet.go

@@ -61,6 +61,19 @@
           "description": "The optional action timeout in seconds",
           "type": "integer"
         },
+        "secret_references": {
+          "description": "A list of all secrets fleet-server needs to inject into the actions data before passing it to the agent. This attribute is removed when action data is send to an agent.",
+          "type": "array",
+          "items": {
+            "type": "object",
+            "properties": {
+              "id": {
+                "type": "string"
+              }
+            },
+            "required": ["id"]
+          }
+        },
         "user_id": {
           "description": "The ID of the user who created the action.",
           "type": "string"