suggestions from @RomainMuller

Signed-off-by: Eliott Bouhana <[email protected]>
DataDog · Jul 10, 2024 · fefe28b · fefe28b
1 parent 6e1344c
commit fefe28b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/internal/appsec/listener/sharedsec/shared.go b/internal/appsec/listener/sharedsec/shared.go
@@ -8,6 +8,7 @@ package sharedsec
 import (
 	"encoding/json"
 	"errors"
+
 	"gopkg.in/DataDog/dd-trace-go.v1/appsec/events"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/sharedsec"
@@ -37,7 +38,12 @@ func RunWAF(wafCtx *waf.Context, values waf.RunAddressData) waf.Result {
 	return result
 }
 
-func MakeWAFRunListener[O dyngo.Operation, T dyngo.ArgOf[O]](events *trace.SecurityEventsHolder, wafCtx *waf.Context, limiter limiter.Limiter, toRunAddressData func(T) waf.RunAddressData) func(O, T) {
+func MakeWAFRunListener[O dyngo.Operation, T dyngo.ArgOf[O]](
+	events *trace.SecurityEventsHolder,
+	wafCtx *waf.Context,
+	limiter limiter.Limiter,
+	toRunAddressData func(T) waf.RunAddressData,
+) func(O, T) {
 	return func(op O, args T) {
 		wafResult := RunWAF(wafCtx, toRunAddressData(args))
 		if !wafResult.HasEvents() {

diff --git a/internal/appsec/listener/sqlsec/sql.go b/internal/appsec/listener/sqlsec/sql.go
@@ -8,6 +8,7 @@ package sqlsec
 import (
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/dyngo"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/sqlsec/types"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/listener"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/listener/sharedsec"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/trace"
 
@@ -26,7 +27,7 @@ func RegisterSQLListener(op dyngo.Operation, events *trace.SecurityEventsHolder,
 	}))
 }
 
-func SQLAddressesPresent(addresses map[string]struct{}) bool {
+func SQLAddressesPresent(addresses listener.AddressSet) bool {
 	_, queryAddr := addresses[ServerDBStatementAddr]
 	_, driverAddr := addresses[ServerDBTypeAddr]