修复客户端断联未释放端口问题

server/fwd/dispatcher/dispatch.go

@@ -9,6 +9,7 @@ import (
 	"proxy-server/server/fwd/http"
 	"proxy-server/server/fwd/socks"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/pkg/errors"
@@ -54,36 +55,43 @@ func (s *Server) Run() error {
 	m.SetReadTimeout(5 * time.Second)
 	defer m.Close()
 
-	go func() {
-		<-s.ctx.Done()
-		close(s.Conn)
-		m.Close()
-	}()
-
 	socksLs := m.Match(cmux.PrefixMatcher(string([]byte{0x05})))
-	defer utils.Close(socksLs)
 	go func() {
 		err = s.acceptSocks(socksLs)
 		if err != nil {
-			slog.Error("dispatcher socks accept error", "err", err)
+			if strings.Contains(err.Error(), "mux: server closed") {
+				return
+			}
+			slog.Warn("dispatcher socks accept error", "err", err)
 		}
 	}()
 
 	httpLs := m.Match(cmux.HTTP1Fast("PATCH"))
-	defer utils.Close(httpLs)
 	go func() {
 		err = s.acceptHttp(httpLs)
 		if err != nil {
-			slog.Error("dispatcher http accept error", "err", err)
+			if strings.Contains(err.Error(), "mux: server closed") {
+				return
+			}
+			slog.Warn("dispatcher http accept error", "err", err)
 		}
 	}()
 
-	err = m.Serve()
-	if err != nil {
-		return errors.Wrap(err, "dispatcher serve error")
-	}
+	errCh := make(chan error)
+	go func() {
+		err = m.Serve()
+		if err != nil {
+			err = errors.Wrap(err, "dispatcher serve error")
+		}
+		errCh <- err
+	}()
 
-	return nil
+	select {
+	case <-s.ctx.Done():
+		return nil
+	case err := <-errCh:
+		return err
+	}
 }
 
 func (s *Server) acceptHttp(ls net.Listener) error {

server/fwd/fwd.go

@@ -12,6 +12,7 @@ import (
 	"proxy-server/server/fwd/dispatcher"
 	"proxy-server/server/pkg/env"
 	"strconv"
+	"strings"
 	"sync"
 
 	"github.com/pkg/errors"
@@ -167,7 +168,6 @@ func (s *Service) startCtrlTun() error {
 
 func (s *Service) processCtrlConn(conn net.Conn) error {
 	slog.Debug("客户端连入", "addr", conn.RemoteAddr().String())
-
 	reader := bufio.NewReader(conn)
 
 	// 获取转发端口
@@ -177,20 +177,58 @@ func (s *Service) processCtrlConn(conn net.Conn) error {
 	}
 	port := binary.BigEndian.Uint16(portBuf)
 
-	// 开放转发端口
+	// 启动转发服务
+	proxy, err := dispatcher.New(port)
+	if err != nil {
+		return errors.Wrap(err, "创建 socks 转发服务失败")
+	}
+	defer proxy.Close()
+
 	s.fwdLesWg.Add(1)
 	go func() {
 		defer s.fwdLesWg.Done()
-		err := s.startFwdTun(port)
+		err := proxy.Run()
 		if err != nil {
-			slog.Error("代理服务启动失败", "err", err)
+			slog.Error("代理服务运行失败", "err", err)
 			return
 		}
 	}()
 
-	// 记录控制连接
-	s.ctrlConnMap.Store(port, conn)
-	return nil
+	// 监听客户端连接
+	errCh := make(chan error)
+	defer close(errCh)
+	go func() {
+		_, err := reader.ReadByte()
+		errCh <- err
+	}()
+
+	// 处理连接
+	for {
+		select {
+		case <-s.ctx.Done():
+			return nil
+		case err := <-errCh:
+			switch {
+			case strings.Contains(err.Error(), "An existing connection was forcibly closed by the remote host."):
+				slog.Debug("客户端主动断开连接")
+				return nil
+			case err == nil:
+				return errors.New("客户端握手失败")
+			default:
+				return errors.Wrap(err, "客户端意外断开连接")
+			}
+		case user := <-proxy.Conn:
+			s.userConnWg.Add(1)
+			go func() {
+				defer s.userConnWg.Done()
+				err := s.processUserConn(user, conn)
+				if err != nil {
+					slog.Error("处理用户连接失败", "err", err)
+					utils.Close(user)
+				}
+			}()
+		}
+	}
 }
 
 func (s *Service) startDataTun() error {
@@ -312,71 +350,22 @@ func (s *Service) processDataConn(client net.Conn) error {
 	return nil
 }
 
-func (s *Service) startFwdTun(port uint16) error {
-	slog.Debug("监听转发通道", "port", port)
-
-	proxy, err := dispatcher.New(port)
-	if err != nil {
-		return errors.Wrap(err, "创建 socks 转发服务失败")
-	}
-	defer proxy.Close()
-
-	errCh := make(chan error)
-	defer close(errCh)
-	go func() {
-		err := proxy.Run()
-		errCh <- err
-	}()
-
-	for {
-		select {
-		case <-s.ctx.Done():
-			return nil
-		case err := <-errCh:
-			if err != nil {
-				return errors.Wrap(err, "转发服务发生错误")
-			}
-		case conn := <-proxy.Conn:
-			s.userConnWg.Add(1)
-			go func() {
-				defer s.userConnWg.Done()
-				err := s.processUserConn(conn, port)
-				if err != nil {
-					slog.Error("处理用户连接失败", "err", err)
-				}
-			}()
-		}
-	}
-}
-
-func (s *Service) processUserConn(conn *core.Conn, port uint16) error {
-
-	// 记录用户连接
-	s.userConnMap.Store(conn.Tag, conn)
-
-	// 通知客户端建立数据通道
-	ctrlConnAny, ok := s.ctrlConnMap.Load(port)
-	if !ok {
-		return errors.New("查找控制连接失败")
-	}
-	ctrlConn := ctrlConnAny.(net.Conn)
+func (s *Service) processUserConn(user *core.Conn, ctrl net.Conn) error {
 
 	// 发送 tag
-	select {
-	case <-s.ctx.Done():
-		return nil
-	default:
-		tag := conn.Tag
-		tagLen := len(tag)
-		tagBuf := make([]byte, 1+tagLen)
-		tagBuf[0] = byte(tagLen)
-		copy(tagBuf[1:], tag)
-		_, err := ctrlConn.Write(tagBuf)
-		if err != nil {
-			return errors.Wrap(err, "向控制通道发送 tag 失败")
-		}
+	tag := user.Tag
+	tagLen := len(tag)
+	tagBuf := make([]byte, 1+tagLen)
+	tagBuf[0] = byte(tagLen)
+	copy(tagBuf[1:], tag)
+
+	_, err := ctrl.Write(tagBuf)
+	if err != nil {
+		return errors.Wrap(err, "向控制通道发送 tag 失败")
 	}
 
-	return nil
+	// 记录用户连接
+	s.userConnMap.Store(user.Tag, user)
 
+	return nil
 }