From 06e557a36b6ec81eead6a1ad4569ca78ef873acd Mon Sep 17 00:00:00 2001
From: Xuran <37136584+Duslia@users.noreply.github.com>
Date: Fri, 30 Jun 2023 14:15:07 +0800
Subject: [PATCH] fix: 304 status code return error (#24)

---
 factory/http2_test.go | 31 +++++++++++++++++++++++++++++++
 response_writer.go    |  7 +------
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/factory/http2_test.go b/factory/http2_test.go
index 8ebe279..0e1158f 100644
--- a/factory/http2_test.go
+++ b/factory/http2_test.go
@@ -339,3 +339,34 @@ func TestTrailer(t *testing.T) {
 		}
 	}
 }
+
+func TestBodyNotAllowedStatus(t *testing.T) {
+	var acceptCount int32 = 0
+	h := server.New(
+		server.WithHostPorts(":8894"),
+		server.WithH2C(true),
+		server.WithOnAccept(func(conn net.Conn) context.Context {
+			atomic.AddInt32(&acceptCount, 1)
+			return context.Background()
+		}))
+
+	// register http2 server factory
+	h.AddProtocol("h2", NewServerFactory())
+
+	h.POST("/", func(c context.Context, ctx *app.RequestContext) {
+		ctx.Data(304, "application/json", []byte("test data"))
+	})
+	go h.Spin()
+	time.Sleep(time.Second)
+
+	c, _ := client.NewClient()
+	c.SetClientFactory(NewClientFactory(config.WithAllowHTTP(true)))
+	req, rsp := protocol.AcquireRequest(), protocol.AcquireResponse()
+	req.SetMethod("POST")
+	req.SetRequestURI("http://127.0.0.1:8894")
+
+	err := c.Do(context.Background(), req, rsp)
+	assert.Nil(t, err)
+	assert.DeepEqual(t, rsp.StatusCode(), 304)
+	assert.DeepEqual(t, len(rsp.Body()), 0)
+}
diff --git a/response_writer.go b/response_writer.go
index 360dbbb..3a9a7a2 100644
--- a/response_writer.go
+++ b/response_writer.go
@@ -299,12 +299,7 @@ func (w *responseWriter) write(lenData int, dataB []byte, dataS string) (n int,
 	if rws == nil {
 		panic("Write called after Handler finished")
 	}
-	if !rws.wroteHeader {
-		w.WriteHeader(200)
-	}
-	if !bodyAllowedForStatus(rws.status) {
-		return 0, http.ErrBodyNotAllowed
-	}
+
 	rws.wroteBytes += int64(len(dataB)) + int64(len(dataS)) // only one can be set
 	if rws.sentContentLen != 0 && rws.wroteBytes > rws.sentContentLen {
 		// TODO: send a RST_STREAM