From 66860d3de8b173919226aed1d534dc752e3b5bc7 Mon Sep 17 00:00:00 2001
From: xjasonlyu <xjasonlyu@gmail.com>
Date: Mon, 27 Mar 2023 00:43:39 +0800
Subject: [PATCH] Improve: enhance error handling

---
 tunnel/tcp.go | 13 +++++++++----
 tunnel/udp.go | 14 ++++++++++----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/tunnel/tcp.go b/tunnel/tcp.go
index beb8167..f3c547e 100644
--- a/tunnel/tcp.go
+++ b/tunnel/tcp.go
@@ -47,18 +47,22 @@ func handleTCPConn(localConn adapter.TCPConn) {
 	defer targetConn.Close()
 
 	log.Infof("[TCP] %s <-> %s", metadata.SourceAddress(), metadata.DestinationAddress())
-	relay(localConn, targetConn) /* relay connections */
+	if err = relay(localConn, targetConn); err != nil {
+		log.Warnf("[TCP] %s <-> %s: %v", metadata.SourceAddress(), metadata.DestinationAddress(), err)
+	}
 }
 
 // relay copies between left and right bidirectionally.
-func relay(left, right net.Conn) {
+func relay(left, right net.Conn) error {
 	wg := sync.WaitGroup{}
 	wg.Add(2)
 
+	var leftErr, rightErr error
+
 	go func() {
 		defer wg.Done()
 		if err := copyBuffer(right, left); err != nil {
-			log.Warnf("[TCP] %v", err)
+			leftErr = errors.Join(leftErr, err)
 		}
 		right.SetReadDeadline(time.Now().Add(tcpWaitTimeout))
 	}()
@@ -66,12 +70,13 @@ func relay(left, right net.Conn) {
 	go func() {
 		defer wg.Done()
 		if err := copyBuffer(left, right); err != nil {
-			log.Warnf("[TCP] %v", err)
+			rightErr = errors.Join(rightErr, err)
 		}
 		left.SetReadDeadline(time.Now().Add(tcpWaitTimeout))
 	}()
 
 	wg.Wait()
+	return errors.Join(leftErr, rightErr)
 }
 
 func copyBuffer(dst io.Writer, src io.Reader) error {
diff --git a/tunnel/udp.go b/tunnel/udp.go
index 8b9c780..6aae098 100644
--- a/tunnel/udp.go
+++ b/tunnel/udp.go
@@ -1,6 +1,7 @@
 package tunnel
 
 import (
+	"errors"
 	"io"
 	"net"
 	"sync"
@@ -57,28 +58,33 @@ func handleUDPConn(uc adapter.UDPConn) {
 	pc = newSymmetricNATPacketConn(pc, metadata)
 
 	log.Infof("[UDP] %s <-> %s", metadata.SourceAddress(), metadata.DestinationAddress())
-	relayPacket(uc, pc, remote)
+	if err = relayPacket(uc, pc, remote); err != nil {
+		log.Warnf("[TCP] %s <-> %s: %v", metadata.SourceAddress(), metadata.DestinationAddress(), err)
+	}
 }
 
-func relayPacket(left net.PacketConn, right net.PacketConn, to net.Addr) {
+func relayPacket(left net.PacketConn, right net.PacketConn, to net.Addr) error {
 	wg := sync.WaitGroup{}
 	wg.Add(2)
 
+	var leftErr, rightErr error
+
 	go func() {
 		defer wg.Done()
 		if err := copyPacketBuffer(right, left, to, _udpSessionTimeout); err != nil {
-			log.Warnf("[UDP] %v", err)
+			leftErr = errors.Join(leftErr, err)
 		}
 	}()
 
 	go func() {
 		defer wg.Done()
 		if err := copyPacketBuffer(left, right, nil, _udpSessionTimeout); err != nil {
-			log.Warnf("[UDP] %v", err)
+			rightErr = errors.Join(rightErr, err)
 		}
 	}()
 
 	wg.Wait()
+	return errors.Join(leftErr, rightErr)
 }
 
 func copyPacketBuffer(dst net.PacketConn, src net.PacketConn, to net.Addr, timeout time.Duration) error {