Repository: cloudwego/hertz
Pull request #685 (2023-03-23 12:32:50)

feat(client): support last middleware for client
#### What type of PR is this?
feat
#### Check the PR title.
- [x] This PR title match the format: \<type\>(optional scope): \<description\>
- [x] The description of this PR title is user-oriented and clear enough for others to understand.


#### (Optional) Translate the PR title into Chinese.
Client 支持设置一个last中间件，主要用于对于最终发起请求前的数据最终处理操作

#### (Optional) More detail description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review. If it is a perf type PR, perf data is suggested to give.
-->
en: 上下文中间件会将一些 ctx 转换为 http 的 header 信息，但如果 ctx 中间件位置过于靠前（往往没法选择），将导致用户后来新增的中间件中针对 ctx 元数据的修改等操作无法生效到最终的header中来，因此增加一个配置最终中间件的方式，让这类操作有机会在用户的中间件执行之后最终收尾。
zh(optional): The context middleware will convert some ctx to http header information, but if the ctx middleware is located too high (often unselectable), it will cause the modification of ctx metadata in the middleware added later by the user to fail to take effect Come to the final header, so add a way to configure the final middleware, so that this type of operation has the opportunity to be finalized after the user's middleware is executed.

#### Which issue(s) this PR fixes:
None


welkeyever: Possibly add the .GetLastMiddleware() which will let users have the chance to get the last Middleware and chain it explicitly.
welkeyever: > Possibly add the .GetLastMiddleware() which will let users have the chance to get the last Middleware and chain it explicitly.

Add `TakeOutLastMiddleware` to let user decide the final order of the middleware.

Patch:
diff --git a/pkg/app/client/client.go b/pkg/app/client/client.go
index 1f98cbe..33f9fb5 100644
--- a/pkg/app/client/client.go
+++ b/pkg/app/client/client.go
@@ -66,7 +66,10 @@ import (
 	"github.com/cloudwego/hertz/pkg/protocol/suite"
 )
 
-var errorInvalidURI = errors.NewPublic("invalid uri")
+var (
+	errorInvalidURI          = errors.NewPublic("invalid uri")
+	errorLastMiddlewareExist = errors.NewPublic("last middleware already set")
+)
 
 // Do performs the given http request and fills the given http response.
 //
@@ -269,10 +272,11 @@ type Client struct {
 
 	clientFactory suite.ClientFactory
 
-	mLock sync.Mutex
-	m     map[string]client.HostClient
-	ms    map[string]client.HostClient
-	mws   Middleware
+	mLock          sync.Mutex
+	m              map[string]client.HostClient
+	ms             map[string]client.HostClient
+	mws            Middleware
+	lastMiddleware Middleware
 }
 
 func (c *Client) GetOptions() *config.ClientOptions {
@@ -454,6 +458,9 @@ func (c *Client) Do(ctx context.Context, req *protocol.Request, resp *protocol.R
 	if c.mws == nil {
 		return c.do(ctx, req, resp)
 	}
+	if c.lastMiddleware != nil {
+		return c.mws(c.lastMiddleware(c.do))(ctx, req, resp)
+	}
 	return c.mws(c.do)(ctx, req, resp)
 }
 
@@ -616,6 +623,28 @@ func (c *Client) Use(mws ...Middleware) {
 	c.mws = chain(middlewares...)
 }
 
+// UseAsLast is used to add middleware to the end of the middleware chain.
+//
+// Will return an error if last middleware has been set before, to ensure all middleware has the change to work,
+// Please use `TakeOutLastMiddleware` to take out the already set middleware.
+// Chain the middleware after or before is both Okay - but remember to put it back.
+func (c *Client) UseAsLast(mw Middleware) error {
+	if c.lastMiddleware != nil {
+		return errorLastMiddlewareExist
+	}
+	c.lastMiddleware = mw
+	return nil
+}
+
+// TakeOutLastMiddleware will return the set middleware and remove it from client.
+//
+// Remember to set it back after chain it with other middleware.
+func (c *Client) TakeOutLastMiddleware() Middleware {
+	last := c.lastMiddleware
+	c.lastMiddleware = nil
+	return last
+}
+
 func newHttp1OptionFromClient(c *Client) *http1.ClientOptions {
 	return &http1.ClientOptions{
 		Name:                          c.options.Name,
diff --git a/pkg/network/netpoll/dial.go b/pkg/network/netpoll/dial.go
index 86a371e..ee1feb0 100644
--- a/pkg/network/netpoll/dial.go
+++ b/pkg/network/netpoll/dial.go
@@ -45,6 +45,9 @@ func (d dialer) DialConnection(n, address string, timeout time.Duration, tlsConf
 }
 
 func (d dialer) DialTimeout(network, address string, timeout time.Duration, tlsConfig *tls.Config) (conn net.Conn, err error) {
+	if tlsConfig != nil {
+		return nil, errNotSupportTLS
+	}
 	conn, err = d.Dialer.DialTimeout(network, address, timeout)
 	return
 }


Test patch:
diff --git a/pkg/app/client/client_test.go b/pkg/app/client/client_test.go
index 1754044..b6eb263 100644
--- a/pkg/app/client/client_test.go
+++ b/pkg/app/client/client_test.go
@@ -55,6 +55,7 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"reflect"
 	"regexp"
 	"strings"
 	"sync"
@@ -1803,13 +1804,82 @@ func TestClientMiddleware(t *testing.T) {
 	client.Use(mw1)
 	client.Use(mw2)
 
-	req, resp := protocol.AcquireRequest(), protocol.AcquireResponse()
-	err := client.Do(context.Background(), req, resp)
+	request, response := protocol.AcquireRequest(), protocol.AcquireResponse()
+	defer func() {
+		protocol.ReleaseRequest(request)
+		protocol.ReleaseResponse(response)
+	}()
+	err := client.Do(context.Background(), request, response)
 	if err != nil {
 		t.Errorf("unexpected error: %s", err.Error())
 	}
 }
 
+func TestClientLastMiddleware(t *testing.T) {
+	client, _ := NewClient()
+	mw0 := func(next Endpoint) Endpoint {
+		return func(ctx context.Context, req *protocol.Request, resp *protocol.Response) (err error) {
+			finalValue0 := ctx.Value("final0")
+			assert.DeepEqual(t, "final3", finalValue0)
+			finalValue1 := ctx.Value("final1")
+			assert.DeepEqual(t, "final1", finalValue1)
+			finalValue2 := ctx.Value("final2")
+			assert.DeepEqual(t, "final2", finalValue2)
+			return nil
+		}
+	}
+	mw1 := func(next Endpoint) Endpoint {
+		return func(ctx context.Context, req *protocol.Request, resp *protocol.Response) (err error) {
+			ctx = context.WithValue(ctx, "final0", "final0")
+			return next(ctx, req, resp)
+		}
+	}
+	mw2 := func(next Endpoint) Endpoint {
+		return func(ctx context.Context, req *protocol.Request, resp *protocol.Response) (err error) {
+			ctx = context.WithValue(ctx, "final1", "final1")
+			return next(ctx, req, resp)
+		}
+	}
+	mw3 := func(next Endpoint) Endpoint {
+		return func(ctx context.Context, req *protocol.Request, resp *protocol.Response) (err error) {
+			ctx = context.WithValue(ctx, "final2", "final2")
+			return next(ctx, req, resp)
+		}
+	}
+	mw4 := func(next Endpoint) Endpoint {
+		return func(ctx context.Context, req *protocol.Request, resp *protocol.Response) (err error) {
+			ctx = context.WithValue(ctx, "final0", "final3")
+			return next(ctx, req, resp)
+		}
+	}
+	err := client.UseAsLast(mw0)
+	assert.Nil(t, err)
+	err = client.UseAsLast(func(endpoint Endpoint) Endpoint {
+		return nil
+	})
+	assert.DeepEqual(t, errorLastMiddlewareExist, err)
+	client.Use(mw1)
+	client.Use(mw2)
+	client.Use(mw3)
+	client.Use(mw4)
+
+	request, response := protocol.AcquireRequest(), protocol.AcquireResponse()
+	defer func() {
+		protocol.ReleaseRequest(request)
+		protocol.ReleaseResponse(response)
+	}()
+	err = client.Do(context.Background(), request, response)
+	if err != nil {
+		t.Errorf("unexpected error: %s", err.Error())
+	}
+
+	last := client.TakeOutLastMiddleware()
+
+	assert.DeepEqual(t, reflect.ValueOf(last).Pointer(), reflect.ValueOf(mw0).Pointer())
+	last = client.TakeOutLastMiddleware()
+	assert.Nil(t, last)
+}
+
 func TestClientReadResponseBodyStreamWithDoubleRequest(t *testing.T) {
 	part1 := ""
 	for i := 0; i < 8192; i++ {
diff --git a/pkg/common/test/assert/assert.go b/pkg/common/test/assert/assert.go
index 4d52201..9ed544f 100644
--- a/pkg/common/test/assert/assert.go
+++ b/pkg/common/test/assert/assert.go
@@ -51,6 +51,7 @@ func DeepEqual(t testing.TB, expected, actual interface{}) {
 }
 
 func Nil(t testing.TB, data interface{}) {
+	t.Helper()
 	if data == nil {
 		return
 	}
@@ -60,6 +61,7 @@ func Nil(t testing.TB, data interface{}) {
 }
 
 func NotNil(t testing.TB, data interface{}) {
+	t.Helper()
 	if data == nil {
 		return
 	}
diff --git a/pkg/network/netpoll/dial_test.go b/pkg/network/netpoll/dial_test.go
new file mode 100644
index 0000000..1830d1e
--- /dev/null
+++ b/pkg/network/netpoll/dial_test.go
@@ -0,0 +1,71 @@
+// Copyright 2023 CloudWeGo Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+//go:build !windows
+// +build !windows
+
+package netpoll
+
+import (
+	"context"
+	"crypto/tls"
+	"testing"
+	"time"
+
+	"github.com/cloudwego/hertz/pkg/common/config"
+	"github.com/cloudwego/hertz/pkg/common/test/assert"
+	"github.com/cloudwego/hertz/pkg/common/test/mock"
+)
+
+func TestDial(t *testing.T) {
+	t.Run("NetpollDial", func(t *testing.T) {
+		const nw = "tcp"
+		const addr = "localhost:10100"
+		transporter := NewTransporter(&config.Options{
+			Addr:    addr,
+			Network: nw,
+		})
+		go transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
+			return nil
+		})
+		defer transporter.Close()
+		time.Sleep(100 * time.Millisecond)
+
+		dial := NewDialer()
+		// DialConnection
+		_, err := dial.DialConnection("tcp", "localhost:10101", time.Second, nil) // wrong addr
+		assert.NotNil(t, err)
+
+		nwConn, err := dial.DialConnection(nw, addr, time.Second, nil)
+		assert.Nil(t, err)
+		defer nwConn.Close()
+		_, err = nwConn.Write([]byte("abcdef"))
+		assert.Nil(t, err)
+		// DialTimeout
+		nConn, err := dial.DialTimeout(nw, addr, time.Second, nil)
+		assert.Nil(t, err)
+		defer nConn.Close()
+	})
+
+	t.Run("NotSupportTLS", func(t *testing.T) {
+		dial := NewDialer()
+		_, err := dial.AddTLS(mock.NewConn(""), nil)
+		assert.DeepEqual(t, errNotSupportTLS, err)
+		_, err = dial.DialConnection("tcp", "localhost:10102", time.Microsecond, &tls.Config{})
+		assert.DeepEqual(t, errNotSupportTLS, err)
+		_, err = dial.DialTimeout("tcp", "localhost:10102", time.Microsecond, &tls.Config{})
+		assert.DeepEqual(t, errNotSupportTLS, err)
+	})
+}
diff --git a/pkg/network/netpoll/transport_test.go b/pkg/network/netpoll/transport_test.go
new file mode 100644
index 0000000..d8a0609
--- /dev/null
+++ b/pkg/network/netpoll/transport_test.go
@@ -0,0 +1,100 @@
+// Copyright 2023 CloudWeGo Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+//go:build !windows
+// +build !windows
+
+package netpoll
+
+import (
+	"context"
+	"net"
+	"sync/atomic"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/cloudwego/hertz/pkg/common/config"
+	"github.com/cloudwego/hertz/pkg/common/test/assert"
+	"github.com/cloudwego/hertz/pkg/network"
+	"golang.org/x/sys/unix"
+)
+
+func TestTransport(t *testing.T) {
+	const nw = "tcp"
+	const addr = "localhost:10103"
+	t.Run("TestDefault", func(t *testing.T) {
+		var onConnFlag, onAcceptFlag, onDataFlag int32
+		transporter := NewTransporter(&config.Options{
+			Addr:    addr,
+			Network: nw,
+			OnConnect: func(ctx context.Context, conn network.Conn) context.Context {
+				atomic.StoreInt32(&onConnFlag, 1)
+				return ctx
+			},
+			WriteTimeout: time.Second,
+			OnAccept: func(conn net.Conn) context.Context {
+				atomic.StoreInt32(&onAcceptFlag, 1)
+				return context.Background()
+			},
+		})
+		go transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
+			atomic.StoreInt32(&onDataFlag, 1)
+			return nil
+		})
+		defer transporter.Close()
+		time.Sleep(100 * time.Millisecond)
+
+		dial := NewDialer()
+		conn, err := dial.DialConnection(nw, addr, time.Second, nil)
+		assert.Nil(t, err)
+		_, err = conn.Write([]byte("123"))
+		assert.Nil(t, err)
+		time.Sleep(100 * time.Millisecond)
+
+		assert.Assert(t, atomic.LoadInt32(&onConnFlag) == 1)
+		assert.Assert(t, atomic.LoadInt32(&onAcceptFlag) == 1)
+		assert.Assert(t, atomic.LoadInt32(&onDataFlag) == 1)
+	})
+
+	t.Run("TestListenConfig", func(t *testing.T) {
+		listenCfg := &net.ListenConfig{Control: func(network, address string, c syscall.RawConn) error {
+			return c.Control(func(fd uintptr) {
+				syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEADDR, 1)
+				syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1)
+			})
+		}}
+		transporter := NewTransporter(&config.Options{
+			Addr:         addr,
+			Network:      nw,
+			ListenConfig: listenCfg,
+		})
+		go transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
+			return nil
+		})
+		defer transporter.Close()
+	})
+
+	t.Run("TestExceptionCase", func(t *testing.T) {
+		assert.Panic(t, func() { // listen err
+			transporter := NewTransporter(&config.Options{
+				Network: "unknow",
+			})
+			transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
+				return nil
+			})
+		})
+	})
+}
diff --git a/pkg/network/standard/dial_test.go b/pkg/network/standard/dial_test.go
new file mode 100644
index 0000000..7d9f3a0
--- /dev/null
+++ b/pkg/network/standard/dial_test.go
@@ -0,0 +1,174 @@
+/*
+ * Copyright 2023 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package standard
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"math/big"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/cloudwego/hertz/pkg/common/config"
+	"github.com/cloudwego/hertz/pkg/common/test/assert"
+)
+
+func TestDial(t *testing.T) {
+	const nw = "tcp"
+	const addr = "localhost:10104"
+	transporter := NewTransporter(&config.Options{
+		Addr:    addr,
+		Network: nw,
+	})
+
+	go transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
+		return nil
+	})
+	defer transporter.Close()
+	time.Sleep(time.Millisecond * 100)
+
+	dial := NewDialer()
+	_, err := dial.DialConnection(nw, addr, time.Second, nil)
+	assert.Nil(t, err)
+
+	nConn, err := dial.DialTimeout(nw, addr, time.Second, nil)
+	assert.Nil(t, err)
+	defer nConn.Close()
+}
+
+func TestDialTLS(t *testing.T) {
+	const nw = "tcp"
+	const addr = "localhost:10105"
+	data := []byte("abcdefg")
+	listened := make(chan struct{})
+	go func() {
+		mockTLSServe(nw, addr, func(conn net.Conn) {
+			defer conn.Close()
+			_, err := conn.Write(data)
+			assert.Nil(t, err)
+		}, listened)
+	}()
+
+	select {
+	case <-listened:
+	case <-time.After(time.Second * 5):
+		t.Fatalf("timeout")
+	}
+
+	dial := NewDialer()
+	_, err := dial.DialConnection(nw, addr, time.Second, &tls.Config{
+		InsecureSkipVerify: true,
+	})
+	assert.Nil(t, err)
+
+	conn, err := dial.DialConnection(nw, addr, time.Second, nil)
+	assert.Nil(t, err)
+	nConn, err := dial.AddTLS(conn, &tls.Config{
+		InsecureSkipVerify: true,
+	})
+	assert.Nil(t, err)
+
+	buf := make([]byte, len(data))
+	_, err = nConn.Read(buf)
+	assert.Nil(t, err)
+	assert.DeepEqual(t, string(data), string(buf))
+}
+
+func mockTLSServe(nw, addr string, handle func(conn net.Conn), listened chan struct{}) (err error) {
+	certData, keyData, err := generateTestCertificate("")
+	if err != nil {
+		return
+	}
+
+	cert, err := tls.X509KeyPair(certData, keyData)
+	if err != nil {
+		return
+	}
+
+	tlsConfig := &tls.Config{
+		Certificates: []tls.Certificate{cert},
+	}
+	ln, err := tls.Listen(nw, addr, tlsConfig)
+	if err != nil {
+		return
+	}
+	defer ln.Close()
+
+	listened <- struct{}{}
+	for {
+		conn, err := ln.Accept()
+		if err != nil {
+			continue
+		}
+		go handle(conn)
+	}
+}
+
+// generateTestCertificate generates a test certificate and private key based on the given host.
+func generateTestCertificate(host string) ([]byte, []byte, error) {
+	priv, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
+	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	cert := &x509.Certificate{
+		SerialNumber: serialNumber,
+		Subject: pkix.Name{
+			Organization: []string{"fasthttp test"},
+		},
+		NotBefore:             time.Now(),
+		NotAfter:              time.Now().Add(365 * 24 * time.Hour),
+		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
+		SignatureAlgorithm:    x509.SHA256WithRSA,
+		DNSNames:              []string{host},
+		BasicConstraintsValid: true,
+		IsCA:                  true,
+	}
+
+	certBytes, err := x509.CreateCertificate(
+		rand.Reader, cert, cert, &priv.PublicKey, priv,
+	)
+
+	p := pem.EncodeToMemory(
+		&pem.Block{
+			Type:  "PRIVATE KEY",
+			Bytes: x509.MarshalPKCS1PrivateKey(priv),
+		},
+	)
+
+	b := pem.EncodeToMemory(
+		&pem.Block{
+			Type:  "CERTIFICATE",
+			Bytes: certBytes,
+		},
+	)
+
+	return b, p, err
+}