Fix fasthttpadaptor Content-Type detection

Make this in line with how net/http handles Content-Type.
valyala · Jun 21, 2020 · ef51a7e · dink10 · Jun 21, 2020 · Quons
1 parent ac4cc17
commit ef51a7e
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/fasthttpadaptor/adaptor.go b/fasthttpadaptor/adaptor.go
@@ -85,11 +85,26 @@ func NewFastHTTPHandler(h http.Handler) fasthttp.RequestHandler {
 		h.ServeHTTP(&w, r.WithContext(ctx))
 
 		ctx.SetStatusCode(w.StatusCode())
+		haveContentType := false
 		for k, vv := range w.Header() {
+			if k == fasthttp.HeaderContentType {
+				haveContentType = true
+			}
+
 			for _, v := range vv {
 				ctx.Response.Header.Set(k, v)
 			}
 		}
+		if !haveContentType {
+			// From net/http.ResponseWriter.Write:
+			// If the Header does not contain a Content-Type line, Write adds a Content-Type set
+			// to the result of passing the initial 512 bytes of written data to DetectContentType.
+			l := 512
+			if len(w.body) < 512 {
+				l = len(w.body)
+			}
+			ctx.Response.Header.Set(fasthttp.HeaderContentType, http.DetectContentType(w.body[:l]))
+		}
 		ctx.Write(w.body) //nolint:errcheck
 	}
 }

diff --git a/fasthttpadaptor/adaptor_test.go b/fasthttpadaptor/adaptor_test.go
@@ -141,3 +141,30 @@ func setContextValueMiddleware(next fasthttp.RequestHandler, key string, value i
 		next(ctx)
 	}
 }
+
+func TestContentType(t *testing.T) {
+	nethttpH := func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("<!doctype html><html>"))
+	}
+	fasthttpH := NewFastHTTPHandler(http.HandlerFunc(nethttpH))
+
+	var ctx fasthttp.RequestCtx
+	var req fasthttp.Request
+
+	req.SetRequestURI("http://example.com")
+
+	remoteAddr, err := net.ResolveTCPAddr("tcp", "1.2.3.4:80")
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	ctx.Init(&req, remoteAddr, nil)
+
+	fasthttpH(&ctx)
+
+	resp := &ctx.Response
+	got := string(resp.Header.Peek("Content-Type"))
+	expected := "text/html; charset=utf-8"
+	if got != expected {
+		t.Errorf("expected %q got %q", expected, got)
+	}
+}