add timeout for simple dialer

Signed-off-by: Asutorufa <16442314+Asutorufa@users.noreply.github.com>
yuhaiin · Jun 13, 2023 · 440a60b · 440a60b
1 parent d7f0f3d
commit 440a60b
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 62 deletions.
diff --git a/pkg/net/latency/tcp.go b/pkg/net/latency/tcp.go
@@ -24,7 +24,7 @@ func HTTP(p proxy.Proxy, target string) (time.Duration, error) {
 	defer tr.CloseIdleConnections()
 
 	start := time.Now()
-	resp, err := (&http.Client{Transport: tr, Timeout: 4 * time.Second}).Get(target)
+	resp, err := http.Get(target)
 	if err != nil {
 		return 0, err
 	}

diff --git a/pkg/net/latency/udp.go b/pkg/net/latency/udp.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/Asutorufa/yuhaiin/pkg/components/inbound"
 	"github.com/Asutorufa/yuhaiin/pkg/net/dns"
 	proxy "github.com/Asutorufa/yuhaiin/pkg/net/interfaces"
 	pdns "github.com/Asutorufa/yuhaiin/pkg/protos/config/dns"
@@ -22,7 +23,7 @@ func DNS(p proxy.Proxy, host, target string) (time.Duration, error) {
 
 	start := time.Now()
 
-	ctx, cancel := context.WithTimeout(context.TODO(), time.Second*5)
+	ctx, cancel := context.WithTimeout(context.TODO(), inbound.Timeout)
 	defer cancel()
 
 	_, err = d.LookupIP(ctx, target)
@@ -48,7 +49,7 @@ func DNSOverQuic(p proxy.Proxy, host, target string) (time.Duration, error) {
 
 	start := time.Now()
 
-	ctx, cancel := context.WithTimeout(context.TODO(), time.Second*5)
+	ctx, cancel := context.WithTimeout(context.TODO(), inbound.Timeout)
 	defer cancel()
 
 	_, err = d.LookupIP(ctx, target)

diff --git a/pkg/net/proxy/http/client_test.go b/pkg/net/proxy/http/client_test.go
diff --git a/pkg/net/proxy/simple/simple.go b/pkg/net/proxy/simple/simple.go
@@ -25,6 +25,8 @@ type Simple struct {
 
 	index      int
 	updateTime time.Time
+
+	timeout time.Duration
 }
 
 func New(c *protocol.Protocol_Simple) protocol.WrapProxy {
@@ -44,16 +46,31 @@ func New(c *protocol.Protocol_Simple) protocol.WrapProxy {
 			addrs = append(addrs, proxy.ParseAddressPort(0, v.GetHost(), proxy.ParsePort(v.GetPort())))
 		}
 
+		timeout := time.Duration(0)
+
+		if c.Simple.Timeout > 0 {
+			timeout = time.Millisecond * time.Duration(c.Simple.Timeout)
+		}
+
 		return &Simple{
 			addrs:        addrs,
 			packetDirect: c.Simple.PacketConnDirect,
 			tlsConfig:    tls,
 			serverNames:  servernames,
+			timeout:      timeout,
 		}, nil
 	}
 }
 
 func (c *Simple) dial(ctx context.Context, addr proxy.Address) (net.Conn, error) {
+	var cancel context.CancelFunc
+	if c.timeout > 0 {
+		ctx, cancel = context.WithTimeout(context.TODO(), c.timeout)
+	} else {
+		ctx, cancel = context.WithTimeout(ctx, time.Second*3)
+	}
+	defer cancel()
+
 	ip, err := addr.IP(ctx)
 	if err != nil {
 		return nil, err
@@ -81,16 +98,12 @@ func (c *Simple) Conn(ctx context.Context, d proxy.Address) (net.Conn, error) {
 
 	if conn == nil {
 		for i, addr := range c.addrs {
-			ctx, cancel := context.WithTimeout(ctx, time.Second*4)
 			con, er := c.dial(ctx, addr)
 			if er != nil {
 				err = errors.Join(err, er)
-				cancel()
 				continue
 			}
 
-			cancel()
-
 			conn = con
 			c.index = i
 

diff --git a/pkg/protos/node/protocol/protocol.pb.go b/pkg/protos/node/protocol/protocol.pb.go
diff --git a/pkg/protos/node/protocol/protocol.proto b/pkg/protos/node/protocol/protocol.proto
@@ -94,6 +94,7 @@ message none {}
 message simple {
   string host = 1 [ json_name = "host" ];
   int32 port = 2 [ json_name = "port" ];
+  uint64 timeout = 6 [ json_name = "timeout" ];
   repeated host alternate_host = 5 [ json_name = "alternate_host" ];
   // udp will write to every packet target instead of only write to host:port
   bool packet_conn_direct = 3 [ json_name = "packet_conn_direct" ];