wip

Signed-off-by: Eliott Bouhana <eliott.bouhana@datadoghq.com>

contrib/internal/httptrace/response_writer.go

@@ -39,6 +39,11 @@ func (w *responseWriter) Status() int {
 	return w.status
 }
 
+// Blocked returns whether the response has been blocked.
+func (w *responseWriter) Blocked() bool {
+	return w.blocked
+}
+
 // Block is supposed only once, after a response (one made by appsec code) as been sent. If it not the case, the function will do nothing.
 // All subsequent calls to Write and WriteHeader will be trigger a log warning users that the response has been blocked.
 func (w *responseWriter) Block() {

contrib/labstack/echo.v4/appsec.go

@@ -21,7 +21,10 @@ func withAppSec(next echo.HandlerFunc, span tracer.Span) echo.HandlerFunc {
 		for _, n := range c.ParamNames() {
 			params[n] = c.Param(n)
 		}
-		var err error
+		var (
+			err error
+			writer = &statusResponseWriter{Response: c.Response()}
+		)
 		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			c.SetRequest(r)
 			err = next(c)
@@ -32,7 +35,7 @@ func withAppSec(next echo.HandlerFunc, span tracer.Span) echo.HandlerFunc {
 			}
 		})
 		// Wrap the echo response to allow monitoring of the response status code in httpsec.WrapHandler()
-		httpsec.WrapHandler(handler, span, params, nil).ServeHTTP(&statusResponseWriter{Response: c.Response()}, c.Request())
+		httpsec.WrapHandler(handler, span, params, nil).ServeHTTP(, c.Request())
 		// If an error occurred, wrap it under an echo.HTTPError. We need to do this so that APM doesn't override
 		// the response code tag with 500 in case it doesn't recognize the error type.
 		if _, ok := err.(*echo.HTTPError); !ok && err != nil {

internal/appsec/emitter/waf/actions/block.go

@@ -141,13 +141,12 @@ func newHTTPBlockRequestAction(statusCode int, template blockingTemplateType) *B
 			template = blockingTemplateTypeFromHeaders(request.Header)
 		}
 
-		if UnwrapGetStatusCode(writer) != 0 {
+		if code, found := UnwrapGetStatusCode(writer); found && code != 0 {
 			// The status code has already been set, so we can't change it, do nothing
 			return
 		}
 
-		blocker, found := UnwrapBlocker(writer)
-		if found {
+		if blocker, found := UnwrapBlocker(writer); found {
 			// We found our custom response writer, so we can block futur calls to Write and WriteHeader
 			defer blocker()
 		}
@@ -203,15 +202,15 @@ func UnwrapBlocker(writer http.ResponseWriter) (func(), bool) {
 
 // UnwrapGetStatusCode unwraps the right struct method from contrib/internal/httptrace.responseWriter
 // and calls it to know if a call to WriteHeader has been made and returns the status code.
-func UnwrapGetStatusCode(writer http.ResponseWriter) int {
+func UnwrapGetStatusCode(writer http.ResponseWriter) (int, bool) {
 	// this is part of the contrib/internal/httptrace.responseWriter interface
 	wrapped, ok := writer.(interface {
 		Status() int
 	})
 	if !ok {
 		// Somehow we can't access the wrapped response writer, so we can't get the status code
-		return 0
+		return 0, false
 	}
 
-	return wrapped.Status()
+	return wrapped.Status(), true
 }