Detect loop and return 404

imranismail · Nov 25, 2020 · e9aa6c9 · e9aa6c9
1 parent 979a9c1
commit e9aa6c9
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/proxy/error_boundary.go b/proxy/error_boundary.go
@@ -77,6 +77,18 @@ func (eb *ErrorBoundary) ModifyRequest(req *http.Request) error {
 func (eb *ErrorBoundary) ModifyResponse(res *http.Response) error {
 	defer eb.resv.ResetResponseVerifications()
 
+	ctx := martian.NewContext(res.Request)
+
+	if err, ok := ctx.Get("via.LoopDetection"); err != nil && ok {
+		res.Body.Close()
+		res.ContentLength = 0
+		res.Body = ioutil.NopCloser(bytes.NewReader([]byte{}))
+		res.StatusCode = http.StatusNotFound
+		res.Status = http.StatusText(res.StatusCode)
+
+		return nil
+	}
+
 	merr := martian.NewMultiError()
 
 	if eb.resmod != nil {

diff --git a/proxy/proxy.go b/proxy/proxy.go
@@ -25,6 +25,7 @@ import (
 	_ "github.com/google/martian/v3/body"
 	_ "github.com/google/martian/v3/cookie"
 	_ "github.com/google/martian/v3/failure"
+	_ "github.com/google/martian/v3/fifo"
 	_ "github.com/google/martian/v3/header"
 	_ "github.com/google/martian/v3/martianurl"
 	_ "github.com/google/martian/v3/method"