From e845353bd22c1969b3da4d55286f17cc91d51c0c Mon Sep 17 00:00:00 2001
From: zoujiaqing <zoujiaqing@gmail.com>
Date: Wed, 22 May 2024 17:16:32 +0800
Subject: [PATCH] Add Client.BindLocalIP(string) for custom use local ip to
 connect.

---
 Client.go               | 67 ++++++++++++++++++++++++++++++++++++-----
 examples/client/main.go |  3 ++
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/Client.go b/Client.go
index 2c2a77b..7052e05 100644
--- a/Client.go
+++ b/Client.go
@@ -24,7 +24,8 @@ type Client struct {
 	pingInterval   time.Duration
 	pingTimer      *time.Timer
 	pingValue      int
-	timeout        uint // 超时时间（秒）
+	timeout        uint   // 超时时间（秒）
+	localIP        string // 本地网卡 IP
 }
 
 func NewClient(addr string, port uint) *Client {
@@ -54,25 +55,75 @@ func (c *Client) OnRequest(handle RequestMessageHandler) {
 	c.onMessageHanle = handle
 }
 
+func createCustomPacketConn(localIP string) (net.PacketConn, error) {
+	localAddr := &net.UDPAddr{
+		IP: net.ParseIP(localIP),
+	}
+	conn, err := net.ListenUDP("udp", localAddr)
+	if err != nil {
+		return nil, err
+	}
+	return conn, nil
+}
+
+func (c *Client) BindLocalIP(localIP string) {
+	c.localIP = localIP
+}
+
 func (c *Client) Connect() error {
 	tlsConf := &tls.Config{
 		InsecureSkipVerify: true,
 		NextProtos:         []string{"qrpc"},
 	}
+
 	quicConf := &quic.Config{
 		MaxIncomingStreams:    1e10, // bidirectional streams
 		MaxIncomingUniStreams: 1e10, // unidirectional streams
 	}
 
-	conn, err := quic.DialAddr(
-		context.Background(),
-		fmt.Sprintf("%s:%d", c.addr, c.port),
-		tlsConf,
-		quicConf)
+	serverAddr := fmt.Sprintf("%s:%d", c.addr, c.port)
 
-	if err != nil {
-		return err
+	// Create a context with timeout
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+
+	var conn quic.Connection
+	var err error
+
+	// Custom localIP
+	if c.localIP != "" {
+		log.Printf("Use bind localIP: %s", c.localIP)
+		// Use custom dialer to specify local address
+		packetConn, err := createCustomPacketConn(c.localIP)
+		if err != nil {
+			log.Printf("Failed to create packet connection: %v", err)
+			return err
+		}
+
+		// Resolve server address
+		serverUDPAddr, err := net.ResolveUDPAddr("udp", serverAddr)
+		if err != nil {
+			log.Printf("Failed to resolve server address: %v", err)
+			return err
+		}
+
+		// Dial QUIC session
+		conn, err = quic.Dial(ctx, packetConn, serverUDPAddr, tlsConf, quicConf)
+		if err != nil {
+			return err
+		}
+	} else {
+		conn, err = quic.DialAddr(
+			ctx,
+			serverAddr,
+			tlsConf,
+			quicConf)
+
+		if err != nil {
+			return err
+		}
 	}
+
 	c.conn = NewRpcConnection(1, conn.Context(), conn, c.timeout)
 	c.conn.OnClose(c.handleConnectionClosed)
 	c.conn.OnRequest(c.handleMessage)
diff --git a/examples/client/main.go b/examples/client/main.go
index 8fe9c4d..cdf57f8 100644
--- a/examples/client/main.go
+++ b/examples/client/main.go
@@ -26,6 +26,9 @@ func run() error {
 
 	client := qrpc.NewClient(*addr, *port)
 
+	// bind local IP
+	client.BindLocalIP("10.10.101.1")
+
 	if err := client.Connect(); err != nil {
 		return err
 	}