serverless/appsec: send enhanced_metric to track supported invocations

DataDog · Jan 9, 2024 · 3d70921 · 3d70921
1 parent 2167824
commit 3d70921
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 6 deletions.
diff --git a/cmd/serverless/main.go b/cmd/serverless/main.go
@@ -257,7 +257,7 @@ func runAgent() {
 	go func() {
 		defer wg.Done()
 		var err error
-		appsecProxyProcessor, err = appsec.New()
+		appsecProxyProcessor, err = appsec.New(serverlessDaemon.MetricAgent.Demux)
 		if err != nil {
 			log.Error("appsec: could not start: ", err)
 		}

diff --git a/pkg/serverless/appsec/appsec.go b/pkg/serverless/appsec/appsec.go
@@ -15,14 +15,15 @@ import (
 	waf "github.com/DataDog/go-libddwaf/v2"
 	json "github.com/json-iterator/go"
 
+	"github.com/DataDog/datadog-agent/pkg/aggregator"
 	"github.com/DataDog/datadog-agent/pkg/serverless/appsec/config"
 	"github.com/DataDog/datadog-agent/pkg/serverless/appsec/httpsec"
 	"github.com/DataDog/datadog-agent/pkg/serverless/proxy"
 	"github.com/DataDog/datadog-agent/pkg/util/log"
 )
 
 //nolint:revive // TODO(ASM) Fix revive linter
-func New() (*httpsec.ProxyLifecycleProcessor, error) {
+func New(demux aggregator.Demultiplexer) (*httpsec.ProxyLifecycleProcessor, error) {
 	appsecInstance, err := newAppSec() // note that the assigned variable is in the parent scope
 	if err != nil {
 		return nil, err
@@ -32,7 +33,7 @@ func New() (*httpsec.ProxyLifecycleProcessor, error) {
 	}
 
 	// AppSec monitors the invocations by acting as a proxy of the AWS Lambda Runtime API.
-	lp := httpsec.NewProxyLifecycleProcessor(appsecInstance)
+	lp := httpsec.NewProxyLifecycleProcessor(appsecInstance, demux)
 	proxy.Start(
 		"127.0.0.1:9000",
 		"127.0.0.1:9001",

diff --git a/pkg/serverless/appsec/appsec_test.go b/pkg/serverless/appsec/appsec_test.go
@@ -26,7 +26,7 @@ func TestNew(t *testing.T) {
 		appsecEnabledStr := strconv.FormatBool(appsecEnabled)
 		t.Run(fmt.Sprintf("DD_SERVERLESS_APPSEC_ENABLED=%s", appsecEnabledStr), func(t *testing.T) {
 			t.Setenv("DD_SERVERLESS_APPSEC_ENABLED", appsecEnabledStr)
-			lp, err := New()
+			lp, err := New(nil)
 			if err := wafHealth(); err != nil {
 				if ok, _ := waf.SupportsTarget(); ok {
 					// host should be supported by appsec, error is unexpected

diff --git a/pkg/serverless/appsec/httpsec/proxy.go b/pkg/serverless/appsec/httpsec/proxy.go
@@ -8,9 +8,11 @@ package httpsec
 import (
 	"bytes"
 
+	"github.com/DataDog/datadog-agent/pkg/aggregator"
 	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace"
 	"github.com/DataDog/datadog-agent/pkg/serverless/appsec/config"
 	"github.com/DataDog/datadog-agent/pkg/serverless/invocationlifecycle"
+	serverlessMetrics "github.com/DataDog/datadog-agent/pkg/serverless/metrics"
 	"github.com/DataDog/datadog-agent/pkg/serverless/trigger"
 	"github.com/DataDog/datadog-agent/pkg/trace/sampler"
 	"github.com/DataDog/datadog-agent/pkg/util/log"
@@ -31,13 +33,16 @@ type ProxyLifecycleProcessor struct {
 
 	// Parsed invocation event value
 	invocationEvent interface{}
+
+	Demux aggregator.Demultiplexer
 }
 
 // NewProxyLifecycleProcessor returns a new httpsec proxy processor monitored with the
 // given Monitorer.
-func NewProxyLifecycleProcessor(appsec Monitorer) *ProxyLifecycleProcessor {
+func NewProxyLifecycleProcessor(appsec Monitorer, demux aggregator.Demultiplexer) *ProxyLifecycleProcessor {
 	return &ProxyLifecycleProcessor{
 		appsec: appsec,
+		Demux:  demux,
 	}
 }
 
@@ -83,6 +88,7 @@ func (lp *ProxyLifecycleProcessor) OnInvokeStart(startDetails *invocationlifecyc
 	case trigger.LambdaFunctionURLEvent:
 		event = &events.LambdaFunctionURLRequest{}
 	}
+	serverlessMetrics.SendASMInvocationEnhancedMetric(nil, lp.Demux)
 
 	if err := json.Unmarshal(payloadBytes, event); err != nil {
 		log.Errorf("appsec: proxy-lifecycle: unexpected lambda event parsing error: %v", err)

diff --git a/pkg/serverless/appsec/httpsec/proxy_test.go b/pkg/serverless/appsec/httpsec/proxy_test.go
@@ -26,7 +26,7 @@ func init() {
 
 func TestProxyLifecycleProcessor(t *testing.T) {
 	t.Setenv("DD_SERVERLESS_APPSEC_ENABLED", "true")
-	lp, err := appsec.New()
+	lp, err := appsec.New(nil)
 	if err != nil {
 		t.Skipf("appsec disabled: %v", err)
 	}