Suggested/conflicts 2 (#2)

* merge conflicts

* conflict

* conflicts

* lint

* skip empty keys

* wee changes and fix

Co-authored-by: Michal Pristas <michal.pristas@gmail.com>

dev-tools/integration/.env

@@ -1,4 +1,4 @@
-ELASTICSEARCH_VERSION=8.5.0-c7913db3-SNAPSHOT
+ELASTICSEARCH_VERSION=8.5.0-56d2c52d-SNAPSHOT
 ELASTICSEARCH_USERNAME=elastic
 ELASTICSEARCH_PASSWORD=changeme
 TEST_ELASTICSEARCH_HOSTS=localhost:9200

internal/pkg/api/handleAck.go

@@ -26,6 +26,7 @@ import (
 	"github.com/elastic/fleet-server/v7/internal/pkg/logger"
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
 	"github.com/elastic/fleet-server/v7/internal/pkg/policy"
+	"github.com/elastic/fleet-server/v7/internal/pkg/smap"
 
 	"github.com/julienschmidt/httprouter"
 	"github.com/rs/zerolog"
@@ -351,38 +352,120 @@ func (ack *AckT) handlePolicyChange(ctx context.Context, zlog zerolog.Logger, ag
 		return nil
 	}
 
-	ack.invalidateAPIKeys(ctx, agent)
+	for _, output := range agent.Outputs {
+		if output.Type != policy.OutputTypeElasticsearch {
+			continue
+		}
+
+		err := ack.updateAPIKey(ctx,
+			zlog,
+			agent.Id,
+			currRev, currCoord,
+			agent.PolicyID,
+			output.APIKeyID, output.PermissionsHash, output.ToRetireAPIKeyIds)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+
+}
+
+func (ack *AckT) updateAPIKey(ctx context.Context,
+	zlog zerolog.Logger,
+	agentID string,
+	currRev, currCoord int64,
+	policyID, apiKeyID, permissionHash string,
+	toRetireAPIKeyIDs []model.ToRetireAPIKeyIdsItems) error {
+
+	if apiKeyID != "" {
+		res, err := ack.bulk.APIKeyRead(ctx, apiKeyID, true)
+		if err != nil {
+			zlog.Error().
+				Err(err).
+				Str(LogAPIKeyID, apiKeyID).
+				Msg("Failed to read API Key roles")
+		} else {
+			clean, removedRolesCount, err := cleanRoles(res.RoleDescriptors)
+			if err != nil {
+				zlog.Error().
+					Err(err).
+					RawJSON("roles", res.RoleDescriptors).
+					Str(LogAPIKeyID, apiKeyID).
+					Msg("Failed to cleanup roles")
+			} else if removedRolesCount > 0 {
+				if err := ack.bulk.APIKeyUpdate(ctx, apiKeyID, permissionHash, clean); err != nil {
+					zlog.Error().Err(err).RawJSON("roles", clean).Str(LogAPIKeyID, apiKeyID).Msg("Failed to update API Key")
+				} else {
+					zlog.Debug().
+						Str("hash.sha256", permissionHash).
+						Str(LogAPIKeyID, apiKeyID).
+						RawJSON("roles", clean).
+						Int("removedRoles", removedRolesCount).
+						Msg("Updating agent record to pick up reduced roles.")
+				}
+			}
+		}
+		ack.invalidateAPIKeys(ctx, toRetireAPIKeyIDs, apiKeyID)
+	}
 
 	body := makeUpdatePolicyBody(
-		agent.PolicyID,
+		policyID,
 		currRev,
 		currCoord,
 	)
 
 	err := ack.bulk.Update(
 		ctx,
 		dl.FleetAgents,
-		agent.Id,
+		agentID,
 		body,
 		bulk.WithRefresh(),
 		bulk.WithRetryOnConflict(3),
 	)
 
-	zlog.Info().Err(err).
-		Str(LogPolicyID, agent.PolicyID).
+	zlog.Err(err).
+		Str(LogPolicyID, policyID).
 		Int64("policyRevision", currRev).
 		Int64("policyCoordinator", currCoord).
 		Msg("ack policy")
 
 	return errors.Wrap(err, "handlePolicyChange update")
 }
 
-func (ack *AckT) invalidateAPIKeys(ctx context.Context, agent *model.Agent) {
-	var ids []string
-	for _, out := range agent.Outputs {
-		for _, k := range out.ToRetireAPIKeyIds {
-			ids = append(ids, k.ID)
+func cleanRoles(roles json.RawMessage) (json.RawMessage, int, error) {
+	rr := smap.Map{}
+	if err := json.Unmarshal(roles, &rr); err != nil {
+		return nil, 0, errors.Wrap(err, "failed to unmarshal provided roles")
+	}
+
+	keys := make([]string, 0, len(rr))
+	for k := range rr {
+		if strings.HasSuffix(k, "-rdstale") {
+			keys = append(keys, k)
+		}
+	}
+
+	if len(keys) == 0 {
+		return roles, 0, nil
+	}
+
+	for _, k := range keys {
+		delete(rr, k)
+	}
+
+	r, err := json.Marshal(rr)
+	return r, len(keys), errors.Wrap(err, "failed to marshal resulting role definition")
+}
+
+func (ack *AckT) invalidateAPIKeys(ctx context.Context, toRetireAPIKeyIDs []model.ToRetireAPIKeyIdsItems, skip string) {
+	ids := make([]string, 0, len(toRetireAPIKeyIDs))
+	for _, k := range toRetireAPIKeyIDs {
+		if k.ID == skip || k.ID == "" {
+			continue
 		}
+		ids = append(ids, k.ID)
 	}
 
 	if len(ids) > 0 {

internal/pkg/api/handleAck_test.go

@@ -452,7 +452,16 @@ func TestInvalidateAPIKeys(t *testing.T) {
 	}}
 	var toRetire3 []model.ToRetireAPIKeyIdsItems
 
-	want := []string{"toRetire1", "toRetire2_0", "toRetire2_1"}
+	skips := map[string]string{
+		"1": "toRetire1",
+		"2": "toRetire2_0",
+		"3": "",
+	}
+	wants := map[string][]string{
+		"1": {},
+		"2": {"toRetire2_1"},
+		"3": {},
+	}
 
 	agent := model.Agent{
 		Outputs: map[string]*model.PolicyOutput{
@@ -462,17 +471,24 @@ func TestInvalidateAPIKeys(t *testing.T) {
 		},
 	}
 
-	bulker := ftesting.NewMockBulk()
-	bulker.On("APIKeyInvalidate",
-		context.Background(), mock.MatchedBy(func(ids []string) bool {
-			// if A contains B and B contains A => A = B
-			return assert.Subset(t, ids, want) &&
-				assert.Subset(t, want, ids)
-		})).
-		Return(nil)
+	for i, out := range agent.Outputs {
+		skip := skips[i]
+		want := wants[i]
+
+		bulker := ftesting.NewMockBulk()
+		if len(want) > 0 {
+			bulker.On("APIKeyInvalidate",
+				context.Background(), mock.MatchedBy(func(ids []string) bool {
+					// if A contains B and B contains A => A = B
+					return assert.Subset(t, ids, want) &&
+						assert.Subset(t, want, ids)
+				})).
+				Return(nil)
+		}
 
-	ack := &AckT{bulk: bulker}
-	ack.invalidateAPIKeys(context.Background(), &agent)
+		ack := &AckT{bulk: bulker}
+		ack.invalidateAPIKeys(context.Background(), out.ToRetireAPIKeyIds, skip)
 
-	bulker.AssertExpectations(t)
+		bulker.AssertExpectations(t)
+	}
 }

internal/pkg/api/handleEnroll.go

@@ -298,7 +298,7 @@ func invalidateAPIKey(ctx context.Context, zlog zerolog.Logger, bulker bulk.Bulk
 LOOP:
 	for {
 
-		_, err := bulker.APIKeyRead(ctx, apikeyID)
+		_, err := bulker.APIKeyRead(ctx, apikeyID, false)
 
 		switch {
 		case err == nil:

internal/pkg/apikey/apikey.go

@@ -6,18 +6,17 @@
 package apikey
 
 import (
-	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"net/http"
 	"strings"
 	"unicode/utf8"
 
 	"github.com/elastic/go-elasticsearch/v7"
 	"github.com/elastic/go-elasticsearch/v7/esapi"
+	"github.com/pkg/errors"
 )
 
 const (
@@ -36,20 +35,26 @@ var AuthKey = http.CanonicalHeaderKey("Authorization")
 
 // APIKeyMetadata tracks Metadata associated with an APIKey.
 type APIKeyMetadata struct {
-	ID       string
-	Metadata Metadata
+	ID              string
+	Metadata        Metadata
+	RoleDescriptors json.RawMessage
 }
 
 // Read gathers APIKeyMetadata from Elasticsearch using the given client.
-func Read(ctx context.Context, client *elasticsearch.Client, id string) (*APIKeyMetadata, error) {
+func Read(ctx context.Context, client *elasticsearch.Client, id string, withOwner bool) (*APIKeyMetadata, error) {
+
 	opts := []func(*esapi.SecurityGetAPIKeyRequest){
 		client.Security.GetAPIKey.WithContext(ctx),
 		client.Security.GetAPIKey.WithID(id),
 	}
+	if withOwner {
+		opts = append(opts, client.Security.GetAPIKey.WithOwner(true))
+	}
 
 	res, err := client.Security.GetAPIKey(
 		opts...,
 	)
+
 	if err != nil {
 		return nil, fmt.Errorf("request to elasticsearch failed: %w", err)
 	}
@@ -60,32 +65,31 @@ func Read(ctx context.Context, client *elasticsearch.Client, id string) (*APIKey
 	}
 
 	type APIKeyResponse struct {
-		ID       string   `json:"id"`
-		Metadata Metadata `json:"metadata"`
+		ID              string          `json:"id"`
+		Metadata        Metadata        `json:"metadata"`
+		RoleDescriptors json.RawMessage `json:"role_descriptors"`
 	}
 	type GetAPIKeyResponse struct {
 		APIKeys []APIKeyResponse `json:"api_keys"`
 	}
 
-	var buff bytes.Buffer
-	if _, err := buff.ReadFrom(res.Body); err != nil {
-		return nil, fmt.Errorf("could not read from response body: %w", err)
-	}
-
 	var resp GetAPIKeyResponse
-	if err = json.Unmarshal(buff.Bytes(), &resp); err != nil {
+	d := json.NewDecoder(res.Body)
+	if err = d.Decode(&resp); err != nil {
 		return nil, fmt.Errorf(
-			"could not Unmarshal elasticsearch GetAPIKeyResponse: %w", err)
+			"could not decode elasticsearch GetAPIKeyResponse: %w", err)
 	}
 
 	if len(resp.APIKeys) == 0 {
 		return nil, ErrAPIKeyNotFound
 	}
 
 	first := resp.APIKeys[0]
+
 	return &APIKeyMetadata{
-		ID:       first.ID,
-		Metadata: first.Metadata,
+		ID:              first.ID,
+		Metadata:        first.Metadata,
+		RoleDescriptors: first.RoleDescriptors,
 	}, nil
 }
 

internal/pkg/apikey/apikey_integration_test.go

@@ -30,7 +30,58 @@ const testFleetRoles = `
 }
 `
 
-func TestRead(t *testing.T) {
+func TestRead_existingKey(t *testing.T) {
+	ctx, cn := context.WithCancel(context.Background())
+	defer cn()
+
+	cfg := elasticsearch.Config{
+		Username: "elastic",
+		Password: "changeme",
+	}
+
+	es, err := elasticsearch.NewClient(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create the key
+	agentID := uuid.Must(uuid.NewV4()).String()
+	name := uuid.Must(uuid.NewV4()).String()
+	akey, err := Create(ctx, es, name, "", "true", []byte(testFleetRoles),
+		NewMetadata(agentID, "", TypeAccess))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Get the key and verify that metadata was saved correctly
+	aKeyMeta, err := Read(ctx, es, akey.ID, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	diff := cmp.Diff(ManagedByFleetServer, aKeyMeta.Metadata.ManagedBy)
+	if diff != "" {
+		t.Error(diff)
+	}
+
+	diff = cmp.Diff(true, aKeyMeta.Metadata.Managed)
+	if diff != "" {
+		t.Error(diff)
+	}
+
+	diff = cmp.Diff(agentID, aKeyMeta.Metadata.AgentID)
+	if diff != "" {
+		t.Error(diff)
+	}
+
+	diff = cmp.Diff(TypeAccess.String(), aKeyMeta.Metadata.Type)
+	if diff != "" {
+		t.Error(diff)
+	}
+
+}
+
+func TestRead_noKey(t *testing.T) {
 	ctx, cn := context.WithCancel(context.Background())
 	defer cn()
 
@@ -45,12 +96,12 @@ func TestRead(t *testing.T) {
 	}
 
 	// Try to get the key that doesn't exist, expect ErrApiKeyNotFound
-	_, err = Read(ctx, es, "0000000000000")
+	_, err = Read(ctx, es, "0000000000000", false)
 	if !errors.Is(err, ErrAPIKeyNotFound) {
-		t.Errorf("Unexpected error type: %v", err)
+		t.Errorf("Unexpected error: %v", err)
 	}
-
 }
+
 func TestCreateAPIKeyWithMetadata(t *testing.T) {
 	tts := []struct {
 		name       string
@@ -92,7 +143,7 @@ func TestCreateAPIKeyWithMetadata(t *testing.T) {
 			}
 
 			// Get the API key and verify that the metadata was saved correctly
-			aKeyMeta, err := Read(ctx, es, apiKey.ID)
+			aKeyMeta, err := Read(ctx, es, apiKey.ID, false)
 			if err != nil {
 				t.Fatal(err)
 			}