refactor(go-forwarder): decouple key resolution, cache clients and share AWS config (#1151)

Author: ndakkoune

aws/logs_monitoring_go/cmd/forwarder/main.go

@@ -11,7 +11,9 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"log/slog"
 
+	"github.com/DataDog/datadog-serverless-functions/aws/logs_monitoring_go/internal/apikey"
 	"github.com/DataDog/datadog-serverless-functions/aws/logs_monitoring_go/internal/config"
 	"github.com/DataDog/datadog-serverless-functions/aws/logs_monitoring_go/internal/filtering"
 	"github.com/DataDog/datadog-serverless-functions/aws/logs_monitoring_go/internal/forwarding"
@@ -37,16 +39,11 @@ func main() {
 	}
 	httpclient.Init(tlsOpts...)
 
-	// Will refactor this in the future to not stop the forwarder if the api key resolution or validation fails.
-	// We may want to store the events in the storage retry mechanism even in case of API key resolution/expiration
-	// so we let the customer some time to configure it properly and not lose any of the events from then.
-	err = cfg.ResolveAPIKey(context.Background())
-	if err != nil {
-		log.Fatal(err)
-	}
-	err = cfg.ValidateAPIKey(context.Background())
-	if err != nil {
-		log.Fatal(err)
+	if err = apikey.Validate(context.Background(), httpclient.Client, cfg.APIURL, cfg.APIKey); err != nil {
+		if !cfg.StoreOnFail {
+			log.Fatal(fmt.Errorf("no failed event storage set, validate API key: %w", err))
+		}
+		slog.Error("API key validation", slog.Any("error", err))
 	}
 
 	lambda.Start(handleRequest(cfg))

aws/logs_monitoring_go/internal/apikey/apikey.go

@@ -0,0 +1,60 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2026-Present Datadog, Inc.
+
+package apikey
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+
+	"github.com/DataDog/datadog-serverless-functions/aws/logs_monitoring_go/internal/httpclient"
+)
+
+const (
+	apiKeyHeader      = "DD-API-KEY"
+	validationTimeout = 3 * time.Second
+)
+
+func Validate(ctx context.Context, client *http.Client, url, key string) error {
+	if !validFormat(key) {
+		return invalidAPIKeyError{"bad format"}
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, validationTimeout)
+	defer cancel()
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return fmt.Errorf("new request: %w", err)
+	}
+
+	req.Header.Set(apiKeyHeader, key)
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return fmt.Errorf("client do: %w", err)
+	}
+	defer httpclient.DrainClose(resp)
+
+	if resp.StatusCode == http.StatusForbidden {
+		return invalidAPIKeyError{"denied from validation endpoint"}
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		if body, err := io.ReadAll(resp.Body); err == nil {
+			return fmt.Errorf("unexpected HTTP/%d response: %s", resp.StatusCode, body)
+		}
+		return fmt.Errorf("unexpected HTTP/%d response", resp.StatusCode)
+	}
+
+	return nil
+}
+
+func validFormat(key string) bool {
+	return len(key) == 32
+}

aws/logs_monitoring_go/internal/apikey/apikey_test.go

@@ -0,0 +1,71 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2026-Present Datadog, Inc.
+
+package apikey
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestValidate(t *testing.T) {
+	t.Parallel()
+
+	tests := map[string]struct {
+		key        string
+		statusCode int
+		wantErr    bool
+		err        error
+	}{
+		"valid": {
+			key:        "0123456789abcdefghij0123456789ab",
+			statusCode: http.StatusOK,
+		},
+		"wrong format": {
+			key:     "not32characters",
+			wantErr: true,
+			err:     &invalidAPIKeyError{},
+		},
+		"invalid": {
+			key:        "myapikeyisexpiredorinvalid012345",
+			statusCode: http.StatusForbidden,
+			wantErr:    true,
+			err:        &invalidAPIKeyError{},
+		},
+		"unexpected error": {
+			key:        "0123456789abcdefghij0123456789ab",
+			statusCode: http.StatusInternalServerError,
+			wantErr:    true,
+		},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+				assert.Equal(t, tc.key, req.Header.Get("DD-API-KEY"))
+				w.WriteHeader(tc.statusCode)
+			}))
+			t.Cleanup(server.Close)
+
+			err := Validate(t.Context(), server.Client(), server.URL, tc.key)
+
+			if tc.wantErr {
+				if tc.err != nil {
+					require.ErrorAs(t, err, tc.err)
+				}
+				require.Error(t, err)
+				return
+			}
+
+			require.NoError(t, err)
+		})
+	}
+}

aws/logs_monitoring_go/internal/apikey/errors.go

@@ -0,0 +1,14 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2026-Present Datadog, Inc.
+
+package apikey
+
+type invalidAPIKeyError struct {
+	message string
+}
+
+func (e invalidAPIKeyError) Error() string {
+	return "invalid Datadog API key: " + e.message + ". See https://docs.datadoghq.com/serverless/forwarder/ for more information."
+}

aws/logs_monitoring_go/internal/config/apikey.go

@@ -7,23 +7,24 @@ package config
 
 import (
 	"compress/gzip"
+	"context"
 	"errors"
+	"fmt"
 	"log/slog"
 	"os"
 	"regexp"
 
 	"github.com/DataDog/datadog-serverless-functions/aws/logs_monitoring_go/internal/model"
+	"github.com/DataDog/datadog-serverless-functions/aws/logs_monitoring_go/internal/sdkclient"
 )
 
 const (
-	DefaultSite     = "datadoghq.com"
-	DefaultPort     = "443"
-	DefaultProtocol = "https"
-	DefaultLogLevel = "INFO"
-
+	DefaultSite                 = "datadoghq.com"
+	DefaultPort                 = "443"
+	DefaultProtocol             = "https"
+	DefaultLogLevel             = "INFO"
 	EnvAPIKey                   = "DD_API_KEY"
 	EnvSite                     = "DD_SITE"
-	EnvURL                      = "DD_URL"
 	EnvAPIURL                   = "DD_API_URL"
 	EnvLogLevel                 = "DD_LOG_LEVEL"
 	EnvPort                     = "DD_PORT"
@@ -45,6 +46,17 @@ const (
 	ForwarderVersion            = "6.0"
 )
 
+type apiKeyResolver struct {
+	env     string
+	resolve func(ctx context.Context, value string) (string, error)
+}
+
+var apiKeyResolvers = []apiKeyResolver{
+	{"DD_API_KEY_SECRET_ARN", sdkclient.ResolveFromSecretsManager},
+	{"DD_API_KEY_SSM_NAME", sdkclient.ResolveFromSSM},
+	{"DD_KMS_API_KEY", sdkclient.ResolveFromKMS},
+}
+
 type Config struct {
 	APIKey                string
 	APIURL                string
@@ -92,6 +104,10 @@ func Load() (*Config, error) {
 		}
 	}
 
+	var resolutionErr error
+	cfg.APIKey, resolutionErr = resolveAPIKey(context.Background(), apiKeyResolvers)
+	errs = append(errs, resolutionErr)
+
 	return &cfg, errors.Join(errs...)
 }
 
@@ -113,8 +129,7 @@ func (c *Config) loadEnv() {
 	}
 	c.CompressionLevel = compressionLevel
 
-	c.IntakeURL = envOrDefault(EnvURL, protocol+"://http-intake.logs."+site+":"+port+"/api/v2/logs")
-	c.APIURL = envOrDefault(EnvAPIURL, protocol+"://api."+site)
+	c.IntakeURL, c.APIURL = buildURLs(protocol, site, port)
 
 	c.Source = envOrDefault(EnvSource, "")
 	c.Host = envOrDefault(EnvHost, "")
@@ -126,3 +141,27 @@ func (c *Config) loadEnv() {
 	c.ScrubIP = envOrDefaultBool(EnvRedactIP, false)
 	c.ScrubEmail = envOrDefaultBool(EnvRedactEmail, false)
 }
+
+func buildURLs(protocol, site, port string) (intakeURL string, apiURL string) {
+	intakeURL = protocol + "://http-intake.logs." + site + ":" + port + "/api/v2/logs"
+	apiURL = protocol + "://api." + site + "/api/v1/validate"
+	return
+}
+
+func resolveAPIKey(ctx context.Context, resolvers []apiKeyResolver) (string, error) {
+	for _, resolver := range resolvers {
+		v, ok := os.LookupEnv(resolver.env)
+		if !ok {
+			continue
+		}
+
+		key, err := resolver.resolve(ctx, v)
+		if err != nil {
+			return "", fmt.Errorf("resolve: %w", err)
+		}
+
+		return key, nil
+	}
+
+	return "", errors.New("no Datadog API key configured")
+}