appsec: early blocking response

contrib/internal/httptrace/response_writer.go

@@ -7,28 +7,63 @@ package httptrace
 
 //go:generate sh -c "go run make_responsewriter.go | gofmt > trace_gen.go"
 
-import "net/http"
+import (
+	"net/http"
+	"sync"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+)
+
+var warnLogOnce sync.Once
+
+const warnLogMsg = `appsec: http.ResponseWriter was used after a security blocking decision was enacted.
+Please check for gopkg.in/DataDog/dd-trace-go.v1/appsec/events.BlockingSecurityEvent in the error result value of instrumented functions.`
+
+// TODO(eliott.bouhana): add a link to the appsec SDK documentation ^^^ here ^^^
 
 // responseWriter is a small wrapper around an http response writer that will
 // intercept and store the status of a request.
 type responseWriter struct {
 	http.ResponseWriter
-	status int
+	status  int
+	blocked bool
 }
 
 func newResponseWriter(w http.ResponseWriter) *responseWriter {
-	return &responseWriter{w, 0}
+	return &responseWriter{w, 0, false}
 }
 
 // Status returns the status code that was monitored.
 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() {
+	if !appsec.Enabled() || w.status == 0 {
+		return
+	}
+
+	w.blocked = true
+}
+
 // Write writes the data to the connection as part of an HTTP reply.
 // We explicitly call WriteHeader with the 200 status code
 // in order to get it reported into the span.
 func (w *responseWriter) Write(b []byte) (int, error) {
+	if w.blocked {
+		warnLogOnce.Do(func() {
+			log.Warn(warnLogMsg)
+		})
+		return len(b), nil
+	}
 	if w.status == 0 {
 		w.WriteHeader(http.StatusOK)
 	}
@@ -38,11 +73,16 @@ func (w *responseWriter) Write(b []byte) (int, error) {
 // WriteHeader sends an HTTP response header with status code.
 // It also sets the status code to the span.
 func (w *responseWriter) WriteHeader(status int) {
-	if w.status != 0 {
+	if w.blocked {
+		warnLogOnce.Do(func() {
+			log.Warn(warnLogMsg)
+		})
 		return
 	}
+	if w.status == 0 {
+		w.status = status
+	}
 	w.ResponseWriter.WriteHeader(status)
-	w.status = status
 }
 
 // Unwrap returns the underlying wrapped http.ResponseWriter.

contrib/internal/httptrace/response_writer_test.go

@@ -7,9 +7,12 @@ package httptrace
 
 import (
 	"net/http"
+	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec"
 )
 
 func Test_wrapResponseWriter(t *testing.T) {
@@ -32,5 +35,40 @@ func Test_wrapResponseWriter(t *testing.T) {
 		_, ok = w.(http.Pusher)
 		assert.True(t, ok)
 	})
+}
+
+func TestBlock(t *testing.T) {
+	appsec.Start()
+	defer appsec.Stop()
+
+	if !appsec.Enabled() {
+		t.Skip("appsec is not enabled")
+	}
+
+	t.Run("block-before-first-write", func(t *testing.T) {
+		recorder := httptest.NewRecorder()
+		rw := newResponseWriter(recorder)
+		rw.Block()
+		assert.False(t, rw.blocked)
+
+		rw.WriteHeader(http.StatusForbidden)
 
+		rw.Block()
+		assert.True(t, rw.blocked)
+
+		assert.Equal(t, http.StatusForbidden, recorder.Code)
+	})
+
+	t.Run("write-after-block", func(t *testing.T) {
+		recorder := httptest.NewRecorder()
+		rw := newResponseWriter(recorder)
+		rw.WriteHeader(http.StatusForbidden)
+		rw.Write([]byte("foo"))
+		rw.Block()
+		rw.WriteHeader(http.StatusOK)
+		rw.Write([]byte("bar"))
+
+		assert.Equal(t, http.StatusForbidden, recorder.Code)
+		assert.Equal(t, recorder.Body.String(), "foo")
+	})
 }

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/httpsec/http.go

@@ -56,20 +56,31 @@ type (
 func (HandlerOperationArgs) IsArgOf(*HandlerOperation)   {}
 func (HandlerOperationRes) IsResultOf(*HandlerOperation) {}
 
-func StartOperation(ctx context.Context, args HandlerOperationArgs) (*HandlerOperation, *atomic.Pointer[actions.BlockHTTP], context.Context) {
-	wafOp, ctx := waf.StartContextOperation(ctx)
+func StartOperation(w http.ResponseWriter, r *http.Request, pathParams map[string]string, opts *Config) (*HandlerOperation, *atomic.Bool, context.Context) {
+	wafOp, ctx := waf.StartContextOperation(r.Context())
 	op := &HandlerOperation{
 		Operation:        dyngo.NewOperation(wafOp),
 		ContextOperation: wafOp,
 	}
 
-	// We need to use an atomic pointer to store the action because the action may be created asynchronously in the future
-	var action atomic.Pointer[actions.BlockHTTP]
+	var blocked atomic.Bool
 	dyngo.OnData(op, func(a *actions.BlockHTTP) {
-		action.Store(a)
+		a.Handler.ServeHTTP(w, r)
+		for _, f := range opts.OnBlock {
+			f()
+		}
 	})
 
-	return op, &action, dyngo.StartAndRegisterOperation(ctx, op, args)
+	return op, &blocked, dyngo.StartAndRegisterOperation(ctx, op, HandlerOperationArgs{
+		Method:      r.Method,
+		RequestURI:  r.RequestURI,
+		Host:        r.Host,
+		RemoteAddr:  r.RemoteAddr,
+		Headers:     r.Header,
+		Cookies:     makeCookies(r.Cookies()),
+		QueryParams: r.URL.Query(),
+		PathParams:  pathParams,
+	})
 }
 
 // Finish the HTTP handler operation and its children operations and write everything to the service entry span.
@@ -125,18 +136,7 @@ func BeforeHandle(
 		opts.ResponseHeaderCopier = defaultWrapHandlerConfig.ResponseHeaderCopier
 	}
 
-	op, blockAtomic, ctx := StartOperation(r.Context(), HandlerOperationArgs{
-		Method:      r.Method,
-		RequestURI:  r.RequestURI,
-		Host:        r.Host,
-		RemoteAddr:  r.RemoteAddr,
-		Headers:     r.Header,
-		Cookies:     makeCookies(r.Cookies()),
-		QueryParams: r.URL.Query(),
-		PathParams:  pathParams,
-	})
-	tr := r.WithContext(ctx)
-	var blocked atomic.Bool
+	op, blocked, ctx := StartOperation(w, r, pathParams, opts)
 
 	afterHandle := func() {
 		var statusCode int
@@ -147,28 +147,9 @@ func BeforeHandle(
 			Headers:    opts.ResponseHeaderCopier(w),
 			StatusCode: statusCode,
 		}, span)
-
-		if blockPtr := blockAtomic.Swap(nil); blockPtr != nil {
-			blockPtr.Handler.ServeHTTP(w, tr)
-			blocked.Store(true)
-		}
-
-		// Execute the onBlock functions to make sure blocking works properly
-		// in case we are instrumenting the Gin framework
-		if blocked.Load() {
-			for _, f := range opts.OnBlock {
-				f()
-			}
-		}
-	}
-
-	if blockPtr := blockAtomic.Swap(nil); blockPtr != nil {
-		// handler is replaced
-		blockPtr.Handler.ServeHTTP(w, tr)
-		blocked.Store(true)
 	}
 
-	return w, tr, afterHandle, blocked.Load()
+	return w, r.WithContext(ctx), afterHandle, blocked.Load()
 }
 
 // WrapHandler wraps the given HTTP handler with the abstract HTTP operation defined by HandlerOperationArgs and
@@ -177,7 +158,6 @@ func BeforeHandle(
 // It is a specific patch meant for Gin, for which we must abort the
 // context since it uses a queue of handlers and it's the only way to make
 // sure other queued handlers don't get executed.
-// TODO: this patch must be removed/improved when we rework our actions/operations system
 func WrapHandler(handler http.Handler, span ddtrace.Span, pathParams map[string]string, opts *Config) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		tw, tr, afterHandle, handled := BeforeHandle(w, r, span, pathParams, opts)

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

@@ -25,7 +25,7 @@ var actionHandlers = map[string]actionHandler{}
 
 func registerActionHandler(aType string, handler actionHandler) {
 	if _, ok := actionHandlers[aType]; ok {
-		log.Warn("appsec: action type `%s` already registered", aType)
+		log.Debug("appsec: action type `%s` already registered", aType)
 		return
 	}
 	actionHandlers[aType] = handler

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

@@ -17,9 +17,9 @@ import (
 func TestNewHTTPBlockRequestAction(t *testing.T) {
 	mux := http.NewServeMux()
 	srv := httptest.NewServer(mux)
-	mux.HandleFunc("/json", newHTTPBlockRequestAction(403, "json").ServeHTTP)
-	mux.HandleFunc("/html", newHTTPBlockRequestAction(403, "html").ServeHTTP)
-	mux.HandleFunc("/auto", newHTTPBlockRequestAction(403, "auto").ServeHTTP)
+	mux.HandleFunc("/json", newHTTPBlockRequestAction(403, BlockingTemplateJSON).ServeHTTP)
+	mux.HandleFunc("/html", newHTTPBlockRequestAction(403, BlockingTemplateHTML).ServeHTTP)
+	mux.HandleFunc("/auto", newHTTPBlockRequestAction(403, BlockingTemplateAuto).ServeHTTP)
 	defer srv.Close()
 
 	t.Run("json", func(t *testing.T) {