Unable to customize http.ResponseWriter because of wrong implementation of Kratos' http.wrapper.Result()

[xbchaos]

What happened:

i implemented ResponseWriter of net/http and use it in a custom Filter, but in my custom ResponseEncoder, the type assertion code below will fail:

customResponseWriter := w.(*CustomResponseWriter)

finally, i found that, in Kratos' http.wrapper.Result(), wrapper.res (which is right actually my custom ResponseWriter) is not used, instead, wrapper.w is used:

func (c *wrapper) Result(code int, v interface{}) error {
	c.w.WriteHeader(code)
	return c.router.srv.enc(&c.w, c.req, v)
}

and what's more, we can see any other method of the wrapper, like JSON() or XML(), they just use wrapper.res:

func (c *wrapper) JSON(code int, v interface{}) error {
	c.res.Header().Set("Content-Type", "application/json")
	c.res.WriteHeader(code)
	return json.NewEncoder(c.res).Encode(v)
}

func (c *wrapper) XML(code int, v interface{}) error {
	c.res.Header().Set("Content-Type", "application/xml")
	c.res.WriteHeader(code)
	return xml.NewEncoder(c.res).Encode(v)
}

so, i'm wondering, the Result() method of wrapper should use wrapper.res too, just like other methods.

What you expected to happen:

http.wrapper.Result() changed to:

func (c *wrapper) Result(code int, v interface{}) error {
	c.res.WriteHeader(code)
	return c.router.srv.enc(c.res, c.req, v)
}

How to reproduce it (as minimally and precisely as possible):

Anything else we need to know?:

Environment:

• Kratos version (use kratos -v): 2.1
• Go version (use go version): 1.20
• OS (e.g: cat /etc/os-release): macOS
• Others:

[dosubot]

🤖

It's great that you've identified a potential improvement in the http.wrapper.Result() method. Your understanding of the code and the proposed change seem well-reasoned. If you believe this change would benefit the project, I encourage you to create a pull request with your proposed modification. Your contribution could help improve the project for everyone. Thank you for considering to contribute!

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[kvii]

我看完源码后猜测可能是因为这个原因：

1. 仔细看 responseWriter 类型的方法，它实现了一个类似于“缓存响应码”的功能。
   

   kratos/transport/http/context.go

   Lines 55 to 59 in a556a2b

   	func (w *responseWriter) WriteHeader(statusCode int) { w.code = statusCode }
   	func (w *responseWriter) Write(data []byte) (int, error) {
   	w.w.WriteHeader(w.code)
   	return w.w.Write(data)
   	}

2. 这是 Result 方法，可以看到和其他方法相比，如 "JSON"，它缺少了设置 "Content-Type" 响应头的步骤，直接设置了响应码。
   

   kratos/transport/http/context.go

   Lines 110 to 119 in a556a2b

   	func (c *wrapper) Result(code int, v interface{}) error {
   	c.w.WriteHeader(code)
   	return c.router.srv.enc(&c.w, c.req, v)
   	}
   	func (c *wrapper) JSON(code int, v interface{}) error {
   	c.res.Header().Set("Content-Type", "application/json")
   	c.res.WriteHeader(code)
   	return json.NewEncoder(c.res).Encode(v)
   	}

3. 在默认的 enc 函数中操作了 header，根据数据类型添加了对应响应头。
   

   kratos/transport/http/codec.go

   Line 98 in a556a2b

   w.Header().Set("Content-Type", httputil.ContentType(codec.Name()))

4. go 源码中 ResponseWriter 的 Header 的注释中说了，设置了响应码之后，Header 的更改就无效了。
   https://github.com/golang/go/blob/f719d5cffdb8298eff7a5ef533fe95290e8c869c/src/net/http/server.go#L99-L101

所以我推测，是因为必须在 Result 等方法中先设置响应码。而先设置响应码会导致响应头的修改无效。所以才有了 responseWriter。不过不论如何，这里肯定有问题就是了，起码 c.res 和 c.w 不应该同时存在。

[kratos-ci-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

After reading the source code, I guessed it might be because of this:

1. Look carefully at the method of the responseWriter type. It implements a function similar to "caching response codes".
   

   kratos/transport/http/context.go

   Lines 55 to 59 in a556a2b

   	func (w *responseWriter) WriteHeader(statusCode int) { w.code = statusCode }
   	func (w *responseWriter) Write(data []byte) (int, error) {
   	w.w.WriteHeader(w.code)
   	return w.w.Write(data)
   	}

2. This is the Result method. You can see that compared with other methods, such as "JSON", it lacks the step of setting the "Content-Type" response header and directly sets the response code.
   

   kratos/transport/http/context.go

   Lines 110 to 113 in a556a2b

   	func (c *wrapper) Result(code int, v interface{}) error {
   	c.w.WriteHeader(code)
   	return c.router.srv.enc(&c.w, c.req, v)
   	}

3. The header is operated in the default enc function and the corresponding response header is added according to the data type.
   

   kratos/transport/http/codec.go

   Line 98 in a556a2b

   w.Header().Set("Content-Type", httputil.ContentType(codec.Name()))

4. The comments on the Header of ResponseWriter in the go source code say that after the response code is set, changes to the Header will be invalid.
   https://github.com/golang/go/blob/f719d5cffdb8298eff7a5ef533fe95290e8c869c/src/net/http/server.go#L99-L101

So I speculate that it is because the response code must be set first in methods such as Result. Setting the response code first will cause the modification of the response header to be invalid. That's why there is responseWriter. But no matter what, there must be a problem here. At least c.res and c.w should not exist at the same time.

[kvii]

@xbchaos Use can use the code below when kratos version >= v2.7.3 to get the underlying custom writer that you've set before:
If you think this code can solve your problem, remember to close this issue.
PR reference: #3265

package main

import (
	"context"
	"fmt"
	"io"
	nt "net/http"
	"net/http/httptest"

	"github.com/go-kratos/kratos/v2/middleware"
	"github.com/go-kratos/kratos/v2/transport/http"
)

type customWriter struct {
	http.ResponseWriter
}

func main() {
	s := http.NewServer(
		http.Middleware(func(h middleware.Handler) middleware.Handler {
			return func(ctx context.Context, req any) (any, error) {
				hc := ctx.(http.Context)
				hc.Reset(customWriter{hc.Response()}, hc.Request()) // reset response writer
				return h(ctx, req)
			}
		}),
		http.ResponseEncoder(func(w http.ResponseWriter, r *http.Request, v any) error {
			// get underlying response writer
			var ok bool
			under := w
		out:
			for {
				switch cur := under.(type) {
				case customWriter:
					under = cur
					ok = true
					break out
				case interface{ Unwrap() http.ResponseWriter }:
					under = cur.Unwrap()
				default:
					break out
				}
			}
			return http.DefaultResponseEncoder(w, r, ok)
		}),
	)

	// simulate what happened in xx_http.pb.go
	s.Route("/").GET("/a", func(ctx http.Context) error {
		h := ctx.Middleware(func(ctx context.Context, req any) (any, error) {
			return nil, nil
		})
		out, _ := h(ctx, nil)
		return ctx.Result(200, out)
	})

	w := httptest.NewRecorder()
	r := httptest.NewRequest(nt.MethodGet, "/a", nil)
	s.ServeHTTP(w, r)
	res := w.Result()
	defer res.Body.Close()

	bs, _ := io.ReadAll(res.Body)
	fmt.Println(string(bs)) // true
}