From 43e44cf0a97a15d1264f4b314aa6bf9235d36380 Mon Sep 17 00:00:00 2001
From: Ehco1996 <zh19960202@gmail.com>
Date: Sat, 24 Aug 2024 13:07:22 +0800
Subject: [PATCH] chore: opt logger

---
 internal/conn/relay_conn.go | 32 ++++++++++++++++++++++----------
 internal/relay/relay.go     |  7 ++-----
 internal/relay/server.go    | 17 +++++++----------
 internal/transporter/raw.go |  4 ++++
 test/relay_test.go          |  2 +-
 5 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/internal/conn/relay_conn.go b/internal/conn/relay_conn.go
index 363893337..6415ba5dc 100644
--- a/internal/conn/relay_conn.go
+++ b/internal/conn/relay_conn.go
@@ -17,6 +17,14 @@ import (
 	"go.uber.org/zap"
 )
 
+const (
+	shortHashLength = 7
+)
+
+var (
+	ErrIdleTimeout = errors.New("connection closed due to idle timeout")
+)
+
 // RelayConn is the interface that represents a relay connection.
 // it contains two connections: clientConn and remoteConn
 // clientConn is the connection from the client to the relay server
@@ -103,7 +111,12 @@ func WithRelayOptions(opts *conf.Options) RelayConnOption {
 }
 
 func (rc *relayConnImpl) Transport() error {
-	defer rc.Close() // nolint: errcheck
+	defer func() {
+		err := rc.Close()
+		if err != nil {
+			rc.l.Errorf("error closing relay connection: %s", err.Error())
+		}
+	}()
 	rc.l = rc.l.Named(shortHashSHA256(rc.GetFlow()))
 	rc.l.Debugf("transport start")
 	c1 := newInnerConn(rc.clientConn, rc)
@@ -124,7 +137,7 @@ func (rc *relayConnImpl) Close() error {
 	err1 := rc.clientConn.Close()
 	err2 := rc.remoteConn.Close()
 	rc.Closed = true
-	return combineErrorsAndMuteEOF(err1, err2)
+	return combineErrorsAndMuteIDLE(err1, err2)
 }
 
 // functions that for web ui
@@ -151,11 +164,11 @@ func (rc *relayConnImpl) GetConnType() string {
 	return rc.ConnType
 }
 
-func combineErrorsAndMuteEOF(err1, err2 error) error {
-	if err1 == io.EOF {
+func combineErrorsAndMuteIDLE(err1, err2 error) error {
+	if err1 == ErrIdleTimeout {
 		err1 = nil
 	}
-	if err2 == io.EOF {
+	if err2 == ErrIdleTimeout {
 		return nil
 	}
 	if err1 != nil && err2 != nil {
@@ -231,9 +244,8 @@ func (c *innerConn) Read(p []byte) (n int, err error) {
 			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
 				since := time.Since(c.lastActive)
 				if since > c.rc.Options.IdleTimeout {
-					c.l.Debugf("Read idle, close remote: %s, since: %s lastActive: %s",
-						c.rc.remote.Label, since, c.lastActive)
-					return 0, err
+					c.l.Debugf("Read idle, close remote: %s", c.rc.remote.Label)
+					return 0, ErrIdleTimeout
 				}
 				continue
 			}
@@ -274,7 +286,7 @@ func shortHashSHA256(input string) string {
 	hasher := sha256.New()
 	hasher.Write([]byte(input))
 	hash := hasher.Sum(nil)
-	return hex.EncodeToString(hash)[:7]
+	return hex.EncodeToString(hash)[:shortHashLength]
 }
 
 func copyConn(conn1, conn2 *innerConn) error {
@@ -298,5 +310,5 @@ func copyConn(conn1, conn2 *innerConn) error {
 	err2 := <-errCH
 	_ = conn1.CloseRead()
 	_ = conn2.CloseRead()
-	return combineErrorsAndMuteEOF(err, err2)
+	return combineErrorsAndMuteIDLE(err, err2)
 }
diff --git a/internal/relay/relay.go b/internal/relay/relay.go
index afccf0afe..c6e1a0fa3 100644
--- a/internal/relay/relay.go
+++ b/internal/relay/relay.go
@@ -44,9 +44,6 @@ func (r *Relay) ListenAndServe(ctx context.Context) error {
 	return <-errCh
 }
 
-func (r *Relay) Close() {
-	r.l.Infof("Close Relay Server:%s", r.cfg.DefaultLabel())
-	if err := r.relayServer.Close(); err != nil {
-		r.l.Errorf(err.Error())
-	}
+func (r *Relay) Stop() error {
+	return r.relayServer.Close()
 }
diff --git a/internal/relay/server.go b/internal/relay/server.go
index 022be14ca..799c2dab2 100644
--- a/internal/relay/server.go
+++ b/internal/relay/server.go
@@ -57,7 +57,7 @@ func (s *Server) startOneRelay(ctx context.Context, r *Relay) {
 }
 
 func (s *Server) stopOneRelay(r *Relay) {
-	r.Close()
+	_ = r.Stop()
 	s.relayM.Delete(r.UniqueID())
 }
 
@@ -85,23 +85,20 @@ func (s *Server) Start(ctx context.Context) error {
 		return err
 	case <-ctx.Done():
 		s.l.Info("ctx cancelled start to stop all relay servers")
-		s.relayM.Range(func(key, value interface{}) bool {
-			r := value.(*Relay)
-			r.Close()
-			return true
-		})
-		return nil
+		return s.Stop()
 	}
 }
 
 func (s *Server) Stop() error {
-	s.l.Info("relay server stop now")
+	var err error
 	s.relayM.Range(func(key, value interface{}) bool {
 		r := value.(*Relay)
-		r.Close()
+		if e := r.Stop(); e != nil {
+			err = errors.Join(err, e)
+		}
 		return true
 	})
-	return nil
+	return err
 }
 
 func (s *Server) TriggerReload() {
diff --git a/internal/transporter/raw.go b/internal/transporter/raw.go
index 0eddf2256..539a140d3 100644
--- a/internal/transporter/raw.go
+++ b/internal/transporter/raw.go
@@ -110,6 +110,10 @@ func (s *RawServer) listenUDP(ctx context.Context) error {
 	for {
 		c, err := s.udpLis.Accept()
 		if err != nil {
+			// Check if the error is due to context cancellation
+			if errors.Is(err, context.Canceled) {
+				return nil // Return without logging the error
+			}
 			s.l.Errorf("UDP accept error: %v", err)
 			return err
 		}
diff --git a/test/relay_test.go b/test/relay_test.go
index 228add58b..487ba36e4 100644
--- a/test/relay_test.go
+++ b/test/relay_test.go
@@ -56,7 +56,7 @@ func TestMain(m *testing.M) {
 	// Cleanup
 	echoServer.Stop()
 	for _, server := range relayServers {
-		server.Close()
+		server.Stop()
 	}
 
 	os.Exit(code)