From e389037711959401f08287a192d1a259c21d0465 Mon Sep 17 00:00:00 2001
From: Ehco1996 <zh19960202@gmail.com>
Date: Mon, 18 Dec 2023 11:29:47 +0800
Subject: [PATCH] *: add  DialTimeOut to ws Dialer #248

---
 internal/config/config.go   |  4 ++--
 internal/relay/relay.go     |  4 +++-
 internal/transporter/raw.go |  3 +--
 internal/transporter/ws.go  | 13 +++++++++++--
 4 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/internal/config/config.go b/internal/config/config.go
index f51d3881c..15d42d302 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -105,7 +105,7 @@ func (c *Config) readFromFile() error {
 	if err != nil {
 		return err
 	}
-	c.L.Infof("Load Config From File: %s", c.PATH)
+	c.L.Debugf("Load Config From File: %s", c.PATH)
 	if err != nil {
 		return err
 	}
@@ -119,7 +119,7 @@ func (c *Config) readFromHttp() error {
 		return err
 	}
 	defer r.Body.Close()
-	c.L.Infof("Load Config From HTTP: %s", c.PATH)
+	c.L.Debugf("Load Config From HTTP: %s", c.PATH)
 	return json.NewDecoder(r.Body).Decode(&c)
 }
 
diff --git a/internal/relay/relay.go b/internal/relay/relay.go
index ae1ea9a7d..f332cc23d 100644
--- a/internal/relay/relay.go
+++ b/internal/relay/relay.go
@@ -153,7 +153,9 @@ func (r *Relay) RunLocalTCPServer() error {
 			defer web.CurConnectionCount.WithLabelValues(remote.Label, web.METRIC_CONN_TYPE_TCP).Dec()
 			defer c.Close()
 			if err := r.TP.HandleTCPConn(c, remote); err != nil {
-				r.L.Errorf("HandleTCPConn meet error from:%s to:%s err:%s", c.RemoteAddr(), remote.Address, err)
+				r.L.Errorf("HandleTCPConn meet error tp:%s from:%s to:%s err:%s",
+					r.TransportType,
+					c.RemoteAddr(), remote.Address, err)
 			}
 		}(c)
 	}
diff --git a/internal/transporter/raw.go b/internal/transporter/raw.go
index 202de2e16..06694b43e 100644
--- a/internal/transporter/raw.go
+++ b/internal/transporter/raw.go
@@ -120,7 +120,6 @@ func (raw *Raw) dialRemote(remote *lb.Node) (net.Conn, error) {
 	d := net.Dialer{Timeout: constant.DialTimeOut}
 	rc, err := d.Dial("tcp", remote.Address)
 	if err != nil {
-		raw.L.Errorf("dial error: %s", err)
 		return nil, err
 	}
 	return rc, nil
@@ -133,10 +132,10 @@ func (raw *Raw) HandleTCPConn(c net.Conn, remote *lb.Node) error {
 	defer c.Close()
 	t1 := time.Now()
 	rc, err := raw.dialRemote(remote)
-	web.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(time.Since(t1).Milliseconds()))
 	if err != nil {
 		return err
 	}
+	web.HandShakeDuration.WithLabelValues(remote.Label).Observe(float64(time.Since(t1).Milliseconds()))
 	raw.L.Infof("HandleTCPConn from %s to %s", c.RemoteAddr(), remote.Address)
 	defer rc.Close()
 	return transport(rc, c, remote.Label)
diff --git a/internal/transporter/ws.go b/internal/transporter/ws.go
index a9e0743e4..39677ac05 100644
--- a/internal/transporter/ws.go
+++ b/internal/transporter/ws.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/Ehco1996/ehco/internal/constant"
 	"github.com/Ehco1996/ehco/internal/lb"
 	"github.com/Ehco1996/ehco/internal/web"
 	"github.com/gobwas/ws"
@@ -17,10 +18,18 @@ type Ws struct {
 	*Raw
 }
 
+func (s *Ws) dialRemote(remote *lb.Node) (net.Conn, error) {
+	d := ws.Dialer{Timeout: constant.DialTimeOut}
+	wsc, _, _, err := d.Dial(context.TODO(), remote.Address+"/ws/")
+	if err != nil {
+		return nil, err
+	}
+	return wsc, nil
+}
+
 func (s *Ws) HandleTCPConn(c net.Conn, remote *lb.Node) error {
 	defer c.Close()
-
-	wsc, _, _, err := ws.Dial(context.TODO(), remote.Address+"/ws/")
+	wsc, err := s.dialRemote(remote)
 	if err != nil {
 		return err
 	}