feat: distinguishing request and response

cloudwego · Sep 9, 2023 · 935bb29 · 935bb29
1 parent 404144e
commit 935bb29
Show file tree

Hide file tree

Showing 7 changed files with 98 additions and 41 deletions.
diff --git a/pkg/app/server/hertz_test.go b/pkg/app/server/hertz_test.go
@@ -784,3 +784,40 @@ func TestSilentMode(t *testing.T) {
 		t.Fatalf("unexpected error in log: %s", b.String())
 	}
 }
+
+func TestHertzDisableHeaderNamesNormalizing(t *testing.T) {
+	h := New(
+		WithHostPorts("localhost:9212"),
+		WithDisableResponseHeaderNamesNormalizing(true),
+		WithDisableRequestHeaderNamesNormalizing(true),
+	)
+	headerName := "CASE-senSITive-HEAder-NAME"
+	headerValue := "foobar baz"
+	succeed := false
+	h.GET("/test", func(c context.Context, ctx *app.RequestContext) {
+		ctx.VisitAllHeaders(func(key, value []byte) {
+			if string(key) == headerName && string(value) == headerValue {
+				succeed = true
+				return
+			}
+		})
+		if !succeed {
+			t.Fatalf("DisableHeaderNamesNormalizing failed")
+		} else {
+			ctx.Header(headerName, headerValue)
+		}
+	})
+
+	go h.Spin()
+	time.Sleep(100 * time.Millisecond)
+
+	cli, _ := c.NewClient(c.WithDisableHeaderNamesNormalizing(true))
+
+	r := protocol.NewRequest("GET", "http://localhost:9212/test", nil)
+	r.Header.DisableNormalizing()
+	r.Header.Set(headerName, headerValue)
+	res := protocol.AcquireResponse()
+	err := cli.Do(context.Background(), r, res)
+	assert.Nil(t, err)
+	assert.DeepEqual(t, headerValue, res.Header.Get(headerName))
+}
diff --git a/pkg/app/server/option.go b/pkg/app/server/option.go
@@ -347,9 +347,16 @@ func WithOnConnect(fn func(ctx context.Context, conn network.Conn) context.Conte
 	}}
 }
 
-// WithDisableHeaderNamesNormalizing is used to set whether disable header names normalizing.
-func WithDisableHeaderNamesNormalizing(disable bool) config.Option {
+// WithDisableRequestHeaderNamesNormalizing is used to set whether disable header names normalizing.
+func WithDisableRequestHeaderNamesNormalizing(disable bool) config.Option {
 	return config.Option{F: func(o *config.Options) {
-		o.DisableHeaderNamesNormalizing = disable
+		o.DisableRequestHeaderNamesNormalizing = disable
+	}}
+}
+
+// WithDisableResponseHeaderNamesNormalizing is used to set whether disable header names normalizing.
+func WithDisableResponseHeaderNamesNormalizing(disable bool) config.Option {
+	return config.Option{F: func(o *config.Options) {
+		o.DisableResponseHeaderNamesNormalizing = disable
 	}}
 }
diff --git a/pkg/app/server/option_test.go b/pkg/app/server/option_test.go
@@ -75,6 +75,8 @@ func TestOptions(t *testing.T) {
 		WithAutoReloadRender(true, 5*time.Second),
 		WithListenConfig(cfg),
 		WithAltTransport(transporter),
+		WithDisableRequestHeaderNamesNormalizing(true),
+		WithDisableResponseHeaderNamesNormalizing(true),
 	})
 	assert.DeepEqual(t, opt.ReadTimeout, time.Second)
 	assert.DeepEqual(t, opt.WriteTimeout, time.Second)
@@ -107,6 +109,8 @@ func TestOptions(t *testing.T) {
 	assert.DeepEqual(t, opt.AutoReloadInterval, 5*time.Second)
 	assert.DeepEqual(t, opt.ListenConfig, cfg)
 	assert.Assert(t, reflect.TypeOf(opt.AltTransporterNewer) == reflect.TypeOf(transporter))
+	assert.DeepEqual(t, opt.DisableRequestHeaderNamesNormalizing, true)
+	assert.DeepEqual(t, opt.DisableResponseHeaderNamesNormalizing, true)
 }
 
 func TestDefaultOptions(t *testing.T) {
@@ -139,6 +143,8 @@ func TestDefaultOptions(t *testing.T) {
 	assert.Assert(t, opt.RegistryInfo == nil)
 	assert.DeepEqual(t, opt.AutoReloadRender, false)
 	assert.DeepEqual(t, opt.AutoReloadInterval, time.Duration(0))
+	assert.DeepEqual(t, opt.DisableRequestHeaderNamesNormalizing, false)
+	assert.DeepEqual(t, opt.DisableResponseHeaderNamesNormalizing, false)
 }
 
 type mockTransporter struct{}

diff --git a/pkg/common/config/option.go b/pkg/common/config/option.go
@@ -104,15 +104,17 @@ type Options struct {
 	// Disabled header names' normalization may be useful only for proxying
 	// responses to other clients expecting case-sensitive header names.
 	//
-	// By default request and response header names are normalized, i.e.
+	// By default, request and response header names are normalized, i.e.
 	// The first letter and the first letters following dashes
 	// are uppercased, while all the other letters are lowercased.
 	// Examples:
 	//
 	//     * HOST -> Host
 	//     * content-type -> Content-Type
 	//     * cONTENT-lenGTH -> Content-Length
-	DisableHeaderNamesNormalizing bool
+	DisableRequestHeaderNamesNormalizing bool
+
+	DisableResponseHeaderNamesNormalizing bool
 }
 
 func (o *Options) Apply(opts []Option) {
@@ -243,7 +245,9 @@ func NewOptions(opts []Option) *Options {
 		Registry: registry.NoopRegistry,
 
 		// Disabled header names' normalization, default false
-		DisableHeaderNamesNormalizing: false,
+		DisableRequestHeaderNamesNormalizing: false,
+
+		DisableResponseHeaderNamesNormalizing: false,
 	}
 	options.Apply(opts)
 	return options

diff --git a/pkg/common/config/option_test.go b/pkg/common/config/option_test.go
@@ -53,6 +53,8 @@ func TestDefaultOptions(t *testing.T) {
 	assert.DeepEqual(t, []interface{}{}, options.Tracers)
 	assert.DeepEqual(t, new(interface{}), options.TraceLevel)
 	assert.DeepEqual(t, registry.NoopRegistry, options.Registry)
+	assert.DeepEqual(t, false, options.DisableRequestHeaderNamesNormalizing)
+	assert.DeepEqual(t, false, options.DisableResponseHeaderNamesNormalizing)
 }
 
 // TestApplyCustomOptions test apply options with custom values after init

diff --git a/pkg/protocol/http1/server.go b/pkg/protocol/http1/server.go
@@ -54,21 +54,20 @@ var (
 )
 
 type Option struct {
-	StreamRequestBody             bool
-	GetOnly                       bool
-	DisablePreParseMultipartForm  bool
-	DisableKeepalive              bool
-	NoDefaultServerHeader         bool
-	MaxRequestBodySize            int
-	IdleTimeout                   time.Duration
-	ReadTimeout                   time.Duration
-	ServerName                    []byte
-	TLS                           *tls.Config
-	HTMLRender                    render.HTMLRender
-	EnableTrace                   bool
-	ContinueHandler               func(header *protocol.RequestHeader) bool
-	HijackConnHandle              func(c network.Conn, h app.HijackHandler)
-	DisableHeaderNamesNormalizing bool
+	StreamRequestBody            bool
+	GetOnly                      bool
+	DisablePreParseMultipartForm bool
+	DisableKeepalive             bool
+	NoDefaultServerHeader        bool
+	MaxRequestBodySize           int
+	IdleTimeout                  time.Duration
+	ReadTimeout                  time.Duration
+	ServerName                   []byte
+	TLS                          *tls.Config
+	HTMLRender                   render.HTMLRender
+	EnableTrace                  bool
+	ContinueHandler              func(header *protocol.RequestHeader) bool
+	HijackConnHandle             func(c network.Conn, h app.HijackHandler)
 }
 
 type Server struct {
@@ -181,11 +180,6 @@ func (s Server) Serve(c context.Context, conn network.Conn) (err error) {
 			})
 		}
 
-		if s.DisableHeaderNamesNormalizing {
-			ctx.Request.Header.DisableNormalizing()
-			ctx.Response.Header.DisableNormalizing()
-		}
-
 		// Read Headers
 		if err = req.ReadHeader(&ctx.Request.Header, zr); err == nil {
 			if s.EnableTrace {

diff --git a/pkg/route/engine.go b/pkg/route/engine.go
@@ -737,6 +737,14 @@ func (engine *Engine) allocateContext() *app.RequestContext {
 	ctx.Response.SetMaxKeepBodySize(engine.options.MaxKeepBodySize)
 	ctx.SetClientIPFunc(engine.clientIPFunc)
 	ctx.SetFormValueFunc(engine.formValueFunc)
+	if engine.options.DisableRequestHeaderNamesNormalizing {
+		ctx.Request.Header.DisableNormalizing()
+		ctx.Request.Header.Trailer().DisableNormalizing()
+	}
+	if engine.options.DisableResponseHeaderNamesNormalizing {
+		ctx.Response.Header.DisableNormalizing()
+		ctx.Response.Header.Trailer().DisableNormalizing()
+	}
 	return ctx
 }
 
@@ -990,21 +998,20 @@ func iterate(method string, routes RoutesInfo, root *node) RoutesInfo {
 // for built-in http1 impl only.
 func newHttp1OptionFromEngine(engine *Engine) *http1.Option {
 	opt := &http1.Option{
-		StreamRequestBody:             engine.options.StreamRequestBody,
-		GetOnly:                       engine.options.GetOnly,
-		DisablePreParseMultipartForm:  engine.options.DisablePreParseMultipartForm,
-		DisableKeepalive:              engine.options.DisableKeepalive,
-		NoDefaultServerHeader:         engine.options.NoDefaultServerHeader,
-		MaxRequestBodySize:            engine.options.MaxRequestBodySize,
-		IdleTimeout:                   engine.options.IdleTimeout,
-		ReadTimeout:                   engine.options.ReadTimeout,
-		ServerName:                    engine.GetServerName(),
-		ContinueHandler:               engine.ContinueHandler,
-		TLS:                           engine.options.TLS,
-		HTMLRender:                    engine.htmlRender,
-		EnableTrace:                   engine.IsTraceEnable(),
-		HijackConnHandle:              engine.HijackConnHandle,
-		DisableHeaderNamesNormalizing: engine.options.DisableHeaderNamesNormalizing,
+		StreamRequestBody:            engine.options.StreamRequestBody,
+		GetOnly:                      engine.options.GetOnly,
+		DisablePreParseMultipartForm: engine.options.DisablePreParseMultipartForm,
+		DisableKeepalive:             engine.options.DisableKeepalive,
+		NoDefaultServerHeader:        engine.options.NoDefaultServerHeader,
+		MaxRequestBodySize:           engine.options.MaxRequestBodySize,
+		IdleTimeout:                  engine.options.IdleTimeout,
+		ReadTimeout:                  engine.options.ReadTimeout,
+		ServerName:                   engine.GetServerName(),
+		ContinueHandler:              engine.ContinueHandler,
+		TLS:                          engine.options.TLS,
+		HTMLRender:                   engine.htmlRender,
+		EnableTrace:                  engine.IsTraceEnable(),
+		HijackConnHandle:             engine.HijackConnHandle,
 	}
 	// Idle timeout of standard network must not be zero. Set it to -1 seconds if it is zero.
 	// Due to the different triggering ways of the network library, see the actual use of this value for the detailed reasons.