contrib/internal/httptrace: add Block() method to stop following call…

…s to Write and WriteHeader

Signed-off-by: Eliott Bouhana <[email protected]>

contrib/internal/httptrace/response_writer.go

@@ -7,28 +7,58 @@ 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
 }
 
+// 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 +68,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")
+	})
 }