From 12e65319b3190ba5d8e6eebd7ecf0e3bab6b6df2 Mon Sep 17 00:00:00 2001
From: Yotam loewenbach <yotam.loewenbach@logz.io>
Date: Mon, 14 Oct 2024 13:17:08 +0300
Subject: [PATCH] modify `initLogger`

---
 handler/handler.go | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/handler/handler.go b/handler/handler.go
index ebc2039..c69c81a 100644
--- a/handler/handler.go
+++ b/handler/handler.go
@@ -44,27 +44,27 @@ type handlerConfig struct {
 	EnvId       string
 	url         string
 	RequestId   string
+	FirehoseArn string
 }
 
-func initLogger(ctx context.Context, request events.APIGatewayProxyRequest, token string) *zap.Logger {
+func initLogger(ctx context.Context, handlerCfg handlerConfig) *zap.Logger {
 	awsRequestId, account, logzioIdentifier := "", "", ""
 	if lambdaContext, ok := lambdacontext.FromContext(ctx); ok {
 		awsRequestId = lambdaContext.AwsRequestID
 	}
-	if arnStr := strings.Split(request.Headers["X-Amz-Firehose-Source-Arn"], ":"); len(arnStr) > 4 {
+	if arnStr := strings.Split(handlerCfg.FirehoseArn, ":"); len(arnStr) > 4 {
 		account = arnStr[4]
 	}
-	if len(token) >= 5 {
-		logzioIdentifier = token[len(token)-5:]
+	if len(handlerCfg.LogzioToken) >= 5 {
+		logzioIdentifier = handlerCfg.LogzioToken[len(handlerCfg.LogzioToken)-5:]
 	}
-	firehoseRequestId := request.Headers["X-Amz-Firehose-Request-Id"]
 	config := zap.NewProductionConfig()
 	config.EncoderConfig.StacktraceKey = ""
 	config.OutputPaths = []string{"stdout"}
 	config.InitialFields = map[string]interface{}{
 		"aws_account":               account,
 		"lambda_invocation_id":      awsRequestId,
-		"firehose_request_id":       firehoseRequestId,
+		"firehose_request_id":       handlerCfg.RequestId,
 		"logzio_account_identifier": logzioIdentifier,
 	}
 	logger, configErr := config.Build()
@@ -129,6 +129,11 @@ func extractHeaders(request events.APIGatewayProxyRequest) handlerConfig {
 		logzioToken = request.Headers["x-amz-firehose-access-key"]
 	}
 	config.LogzioToken = logzioToken
+	firehoseArn := request.Headers["X-Amz-Firehose-Source-Arn"]
+	if firehoseArn == "" {
+		firehoseArn = request.Headers["x-amz-firehose-source-arn"]
+	}
+	config.FirehoseArn = firehoseArn
 	commonAttributes := request.Headers["X-Amz-Firehose-Common-Attributes"]
 	if commonAttributes == "" {
 		commonAttributes = request.Headers["x-amz-firehose-common-attributes"]
@@ -301,7 +306,7 @@ func HandleRequest(ctx context.Context, request events.APIGatewayProxyRequest) (
 	metricCount, dataPointCount := 0, 0
 	shippingErrors := new(internal.ErrorCollector)
 	handlerCfg := extractHeaders(request)
-	log := initLogger(ctx, request, handlerCfg.LogzioToken)
+	log := initLogger(ctx, handlerCfg)
 	firehoseResponseClient := internal.NewResponseClient(handlerCfg.RequestId, log)
 	defer log.Sync()
 	if handlerCfg.LogzioToken == "" {