From f2841a599857eaedfa1163ea2c6dcc2f814cab9b Mon Sep 17 00:00:00 2001
From: Blake Rouse <blake.rouse@elastic.co>
Date: Fri, 5 Mar 2021 08:52:51 -0500
Subject: [PATCH] Add API key invalidation on unenroll ACK. (#124) (#125)

* Add API key invalidation on unenroll ACK.

* Fix import location.

* Fix test.

(cherry picked from commit dc762b5a39500afb68d969d43e6de83656ddbff7)
---
 cmd/fleet/handleAck.go               | 19 +++++++++++
 cmd/fleet/server_integration_test.go |  2 +-
 internal/pkg/apikey/auth.go          |  2 +-
 internal/pkg/apikey/create.go        |  2 +-
 internal/pkg/apikey/invalidate.go    | 50 ++++++++++++++++++++++++++++
 5 files changed, 72 insertions(+), 3 deletions(-)
 create mode 100644 internal/pkg/apikey/invalidate.go

diff --git a/cmd/fleet/handleAck.go b/cmd/fleet/handleAck.go
index bdf1d7ade..8ce7e4b7e 100644
--- a/cmd/fleet/handleAck.go
+++ b/cmd/fleet/handleAck.go
@@ -13,6 +13,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/elastic/fleet-server/v7/internal/pkg/apikey"
 	"github.com/elastic/fleet-server/v7/internal/pkg/bulk"
 	"github.com/elastic/fleet-server/v7/internal/pkg/cache"
 	"github.com/elastic/fleet-server/v7/internal/pkg/dl"
@@ -191,6 +192,13 @@ func _handlePolicyChange(ctx context.Context, bulker bulk.Bulk, agent *model.Age
 }
 
 func _handleUnenroll(ctx context.Context, bulker bulk.Bulk, agent *model.Agent) error {
+	apiKeys := _getAPIKeyIDs(agent)
+	if len(apiKeys) > 0 {
+		if err := apikey.Invalidate(ctx, bulker.Client(), apiKeys...); err != nil {
+			return err
+		}
+	}
+
 	updates := make([]bulk.BulkOp, 0, 1)
 	now := time.Now().UTC().Format(time.RFC3339)
 	fields := map[string]interface{}{
@@ -214,3 +222,14 @@ func _handleUnenroll(ctx context.Context, bulker bulk.Bulk, agent *model.Agent)
 
 	return bulker.MUpdate(ctx, updates, bulk.WithRefresh())
 }
+
+func _getAPIKeyIDs(agent *model.Agent) []string {
+	keys := make([]string, 0, 1)
+	if agent.AccessApiKeyId != "" {
+		keys = append(keys, agent.AccessApiKeyId)
+	}
+	if agent.DefaultApiKeyId != "" {
+		keys = append(keys, agent.DefaultApiKeyId)
+	}
+	return keys
+}
diff --git a/cmd/fleet/server_integration_test.go b/cmd/fleet/server_integration_test.go
index 5f9d726c9..e36527910 100644
--- a/cmd/fleet/server_integration_test.go
+++ b/cmd/fleet/server_integration_test.go
@@ -181,7 +181,7 @@ func TestServerUnauthorized(t *testing.T) {
 
 	// Unauthorized, expecting error from /_security/_authenticate
 	t.Run("unauthorized", func(t *testing.T) {
-		const expectedErrResponsePrefix = `Fail Auth: [401 Unauthorized]`
+		const expectedErrResponsePrefix = `fail Auth: [401 Unauthorized]`
 		for _, u := range agenturls {
 			req, err := http.NewRequest("POST", u, bytes.NewBuffer([]byte("{}")))
 			require.NoError(t, err)
diff --git a/internal/pkg/apikey/auth.go b/internal/pkg/apikey/auth.go
index e097bf185..0a7675770 100644
--- a/internal/pkg/apikey/auth.go
+++ b/internal/pkg/apikey/auth.go
@@ -48,7 +48,7 @@ func (k ApiKey) Authenticate(ctx context.Context, es *elasticsearch.Client) (*Se
 	}
 
 	if res.IsError() {
-		return nil, fmt.Errorf("Fail Auth: %s", res.String())
+		return nil, fmt.Errorf("fail Auth: %s", res.String())
 	}
 
 	var info SecurityInfo
diff --git a/internal/pkg/apikey/create.go b/internal/pkg/apikey/create.go
index dc244871d..35d4b66b2 100644
--- a/internal/pkg/apikey/create.go
+++ b/internal/pkg/apikey/create.go
@@ -48,7 +48,7 @@ func Create(ctx context.Context, client *elasticsearch.Client, name, ttl string,
 	defer res.Body.Close()
 
 	if res.IsError() {
-		return nil, fmt.Errorf("Fail CreateAPIKey: %s", res.String())
+		return nil, fmt.Errorf("fail CreateAPIKey: %s", res.String())
 	}
 
 	type APIKeyResponse struct {
diff --git a/internal/pkg/apikey/invalidate.go b/internal/pkg/apikey/invalidate.go
new file mode 100644
index 000000000..938e2bd52
--- /dev/null
+++ b/internal/pkg/apikey/invalidate.go
@@ -0,0 +1,50 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License;
+// you may not use this file except in compliance with the Elastic License.
+
+package apikey
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/elastic/go-elasticsearch/v8"
+	"github.com/elastic/go-elasticsearch/v8/esapi"
+)
+
+// Invalidate invalidates the provides API keys by ID.
+func Invalidate(ctx context.Context, client *elasticsearch.Client, ids ...string) error {
+
+	payload := struct {
+		IDs []string `json:"ids,omitempty"`
+	}{
+		ids,
+	}
+
+	body, err := json.Marshal(&payload)
+	if err != nil {
+		return err
+	}
+
+	opts := []func(*esapi.SecurityInvalidateAPIKeyRequest){
+		client.Security.InvalidateAPIKey.WithContext(ctx),
+	}
+
+	res, err := client.Security.InvalidateAPIKey(
+		bytes.NewReader(body),
+		opts...,
+	)
+
+	if err != nil {
+		return err
+	}
+
+	defer res.Body.Close()
+
+	if res.IsError() {
+		return fmt.Errorf("fail InvalidateAPIKey: %s", res.String())
+	}
+	return nil
+}