优化 socks 解析流程

template/service/service.go

@@ -0,0 +1,135 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"io"
+	"log"
+	"net"
+	"strconv"
+	"sync"
+	"time"
+)
+
+type Config struct {
+	Host      string
+	Port      uint16
+	CloseWait time.Duration
+}
+
+type Server struct {
+	config *Config
+	ctx    context.Context
+	cancel context.CancelFunc
+	wg     sync.WaitGroup
+}
+
+func New(conf *Config) (*Server, error) {
+
+	if conf.Host == "" {
+		conf.Host = "localhost"
+	}
+	if conf.Port == 0 {
+		return nil, errors.New("port is required")
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	return &Server{
+		conf,
+		ctx,
+		cancel,
+		sync.WaitGroup{},
+	}, nil
+}
+
+func (s *Server) Run() error {
+
+	// start listen
+	addr := net.JoinHostPort(s.config.Host, strconv.Itoa(int(s.config.Port)))
+	ls, err := net.Listen("tcp", addr)
+	if err != nil {
+		return errors.New("failed to listen")
+	}
+	defer closeRes(ls)
+
+	// wait accept
+	connCh := make(chan net.Conn)
+	defer close(connCh)
+	go func() {
+		for {
+			conn, err := ls.Accept()
+			if err != nil {
+				if !errors.Is(err, net.ErrClosed) {
+					log.Println("accept failed", err)
+				}
+				// retry on temporary error
+				var ne net.Error
+				if errors.As(err, &ne) && ne.Temporary() {
+					continue
+				}
+				return
+			}
+			select {
+			case <-s.ctx.Done():
+				closeRes(conn)
+				return
+			case connCh <- conn:
+			}
+		}
+	}()
+
+	// handle accept
+	func() {
+		for {
+			select {
+			case <-s.ctx.Done():
+				return
+			case conn, ok := <-connCh:
+				if !ok {
+					return
+				}
+				s.wg.Add(1)
+				go func() {
+					defer s.wg.Done()
+					_ = s.handle(conn)
+				}()
+			}
+		}
+	}()
+
+	// close
+	timeout, cancel := context.WithTimeout(context.Background(), s.config.CloseWait)
+	defer cancel()
+
+	waitCh := make(chan struct{})
+	defer close(waitCh)
+	go func() {
+		s.wg.Wait()
+		select {
+		case <-timeout.Done(): // waitCh may be closed
+		case waitCh <- struct{}{}:
+		}
+	}()
+
+	err = nil
+	select {
+	case <-timeout.Done():
+		err = errors.New("close timeout")
+	case <-waitCh:
+	}
+
+	return err
+}
+
+func (s *Server) Close() {
+	s.cancel()
+}
+
+func (s *Server) handle(conn net.Conn) error {
+	defer closeRes(conn)
+	return nil
+}
+
+func closeRes[T io.Closer](res T) {
+	_ = res.Close()
+}