[Fleet] Replace findByApiKeyID with a get call when possible (#3611)

Author: nchaulet

internal/pkg/api/auth.go

@@ -116,7 +116,13 @@ func authAgent(r *http.Request, id *string, bulker bulk.Bulk, c cache.Cache) (*m
 			Msg("authApiKey slow")
 	}
 
-	agent, err := findAgentByAPIKeyID(ctx, bulker, key.ID)
+	var agent *model.Agent
+	// If we have the agentID retrieve the agent document with a get (more performant) instead of triggering a search
+	if id != nil {
+		agent, err = getAgentAndVerifyAPIKeyID(ctx, bulker, *id, key.ID)
+	} else {
+		agent, err = findAgentByAPIKeyID(ctx, bulker, key.ID)
+	}
 	if err != nil {
 		return nil, err
 	}

internal/pkg/api/handleCheckin.go

@@ -855,6 +855,25 @@ func processPolicy(ctx context.Context, zlog zerolog.Logger, bulker bulk.Bulk, a
 	return &resp, nil
 }
 
+func getAgentAndVerifyAPIKeyID(ctx context.Context, bulker bulk.Bulk, agentID string, apiKeyID string) (*model.Agent, error) {
+	span, ctx := apm.StartSpan(ctx, "getAgentAndVerifyAPIKeyID", "read")
+	defer span.End()
+	agent, err := dl.GetAgent(ctx, bulker, agentID)
+	if err != nil {
+		if errors.Is(err, dl.ErrNotFound) {
+			err = ErrAgentNotFound
+		} else {
+			err = fmt.Errorf("GetAgent: %w", err)
+		}
+	}
+
+	if agent.AccessAPIKeyID != apiKeyID {
+		err = fmt.Errorf("invalid API Key ID %w", ErrAgentIdentity)
+	}
+
+	return &agent, err
+}
+
 func findAgentByAPIKeyID(ctx context.Context, bulker bulk.Bulk, id string) (*model.Agent, error) {
 	span, ctx := apm.StartSpan(ctx, "findAgentByID", "search")
 	defer span.End()

internal/pkg/bulk/engine.go

@@ -47,6 +47,7 @@ type Bulk interface {
 	// Synchronous operations run in the bulk engine
 	Create(ctx context.Context, index, id string, body []byte, opts ...Opt) (string, error)
 	Read(ctx context.Context, index, id string, opts ...Opt) ([]byte, error)
+	ReadRaw(ctx context.Context, index, id string, opts ...Opt) (*MgetResponseItem, error)
 	Update(ctx context.Context, index, id string, body []byte, opts ...Opt) error
 	Delete(ctx context.Context, index, id string, opts ...Opt) error
 	Index(ctx context.Context, index, id string, body []byte, opts ...Opt) (string, error)

internal/pkg/bulk/opRead.go

@@ -21,8 +21,8 @@ const (
 	rSuffix = "]}"
 )
 
-func (b *Bulker) Read(ctx context.Context, index, id string, opts ...Opt) ([]byte, error) {
-	span, ctx := apm.StartSpan(ctx, "Bulker: read", "bulker")
+func (b *Bulker) ReadRaw(ctx context.Context, index, id string, opts ...Opt) (*MgetResponseItem, error) {
+	span, ctx := apm.StartSpan(ctx, "Bulker: readRaw", "bulker")
 	defer span.End()
 	opt := b.parseOpts(append(opts, withAPMLinkedContext(ctx))...)
 	blk := b.newBlk(ActionRead, opt)
@@ -49,6 +49,17 @@ func (b *Bulker) Read(ctx context.Context, index, id string, opts ...Opt) ([]byt
 	if !ok {
 		return nil, fmt.Errorf("unable to cast response to *MgetResponseItem, detected type: %T", resp.data)
 	}
+	return r, nil
+}
+
+func (b *Bulker) Read(ctx context.Context, index, id string, opts ...Opt) ([]byte, error) {
+	span, ctx := apm.StartSpan(ctx, "Bulker: read", "bulker")
+	defer span.End()
+	r, err := b.ReadRaw(ctx, index, id, opts...)
+	if err != nil {
+		return nil, err
+	}
+
 	return r.Source, nil
 }
 

internal/pkg/bulk/schema.go

@@ -80,9 +80,9 @@ type MgetResponse struct {
 type MgetResponseItem struct {
 	//	Index      string          `json:"_index"`
 	//	Type       string          `json:"_type"`
-	//	DocumentID string          `json:"_id"`
-	//	Version    int64           `json:"_version"`
-	//	SeqNo      int64           `json:"_seq_no"`
+	DocumentID string `json:"_id"`
+	Version    int64  `json:"_version"`
+	SeqNo      int64  `json:"_seq_no"`
 	//	PrimTerm   int64           `json:"_primary_term"`
 	Found bool `json:"found"`
 	//	Routing    string          `json:"_routing"`

internal/pkg/dl/agent.go

@@ -6,10 +6,13 @@ package dl
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
 	"fmt"
 
 	"github.com/elastic/fleet-server/v7/internal/pkg/bulk"
 	"github.com/elastic/fleet-server/v7/internal/pkg/dsl"
+	"github.com/elastic/fleet-server/v7/internal/pkg/es"
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
 )
 
@@ -39,6 +42,29 @@ func prepareAgentFindByField(field string) *dsl.Tmpl {
 	return prepareFindByField(field, map[string]interface{}{"version": true})
 }
 
+func GetAgent(ctx context.Context, bulker bulk.Bulk, agentID string, opt ...Option) (model.Agent, error) {
+	o := newOption(FleetAgents, opt...)
+	var agent model.Agent
+	data, err := bulker.ReadRaw(ctx, o.indexName, agentID)
+	if err != nil {
+		if errors.Is(err, es.ErrElasticNotFound) {
+			return model.Agent{}, ErrNotFound
+		} else {
+			return model.Agent{}, err
+		}
+	}
+	err = json.Unmarshal(data.Source, &agent)
+	if err != nil {
+		return model.Agent{}, err
+	}
+
+	agent.Id = agentID
+	agent.SeqNo = data.SeqNo
+	agent.Version = data.Version
+
+	return agent, err
+}
+
 func FindAgent(ctx context.Context, bulker bulk.Bulk, tmpl *dsl.Tmpl, name string, v interface{}, opt ...Option) (model.Agent, error) {
 	o := newOption(FleetAgents, opt...)
 	res, err := SearchWithOneParam(ctx, bulker, tmpl, o.indexName, name, v)

internal/pkg/server/fleet_integration_test.go

@@ -897,7 +897,7 @@ func Test_Agent_Auth_errors(t *testing.T) {
 		res, err := cli.Do(req)
 		require.NoError(t, err)
 		res.Body.Close()
-		require.Equal(t, http.StatusNotFound, res.StatusCode) // NOTE this is a 404 and not a 400
+		require.Equal(t, http.StatusForbidden, res.StatusCode)
 	})
 	t.Run("wrong agent ID", func(t *testing.T) {
 		ctx := testlog.SetLogger(t).WithContext(ctx)

internal/pkg/testing/bulk.go

@@ -45,6 +45,11 @@ func (m *MockBulk) Read(ctx context.Context, index, id string, opts ...bulk.Opt)
 	return args.Get(0).([]byte), args.Error(1)
 }
 
+func (m *MockBulk) ReadRaw(ctx context.Context, index, id string, opts ...bulk.Opt) (*bulk.MgetResponseItem, error) {
+	args := m.Called(ctx, index, id, opts)
+	return args.Get(0).(*bulk.MgetResponseItem), args.Error(1)
+}
+
 func (m *MockBulk) Delete(ctx context.Context, index, id string, opts ...bulk.Opt) error {
 	args := m.Called(ctx, index, id, opts)
 	return args.Error(0)