feat(agent/stream): remove STREAM_PORT and use tcp multiplexing on the same port

agent/cmd/main.go

@@ -1,15 +1,21 @@
 package main
 
 import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
 	"net"
+	"net/http"
 	"os"
+	"sync"
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"github.com/yusing/godoxy/agent/pkg/agent"
 	"github.com/yusing/godoxy/agent/pkg/agent/stream"
 	"github.com/yusing/godoxy/agent/pkg/env"
-	"github.com/yusing/godoxy/agent/pkg/server"
+	"github.com/yusing/godoxy/agent/pkg/handler"
 	"github.com/yusing/godoxy/internal/metrics/systeminfo"
 	socketproxy "github.com/yusing/godoxy/socketproxy/pkg"
 	gperr "github.com/yusing/goutils/errs"
@@ -19,6 +25,53 @@ import (
 	"github.com/yusing/goutils/version"
 )
 
+var errListenerClosed = errors.New("listener closed")
+
+type connQueueListener struct {
+	addr      net.Addr
+	conns     chan net.Conn
+	closed    chan struct{}
+	closeOnce sync.Once
+}
+
+func newConnQueueListener(addr net.Addr, buffer int) *connQueueListener {
+	return &connQueueListener{
+		addr:   addr,
+		conns:  make(chan net.Conn, buffer),
+		closed: make(chan struct{}),
+	}
+}
+
+func (l *connQueueListener) push(conn net.Conn) error {
+	select {
+	case <-l.closed:
+		_ = conn.Close()
+		return errListenerClosed
+	case l.conns <- conn:
+		return nil
+	}
+}
+
+func (l *connQueueListener) Accept() (net.Conn, error) {
+	conn, ok := <-l.conns
+	if !ok {
+		return nil, errListenerClosed
+	}
+	return conn, nil
+}
+
+func (l *connQueueListener) Close() error {
+	l.closeOnce.Do(func() {
+		close(l.closed)
+		close(l.conns)
+	})
+	return nil
+}
+
+func (l *connQueueListener) Addr() net.Addr {
+	return l.addr
+}
+
 func main() {
 	writer := zerolog.ConsoleWriter{
 		Out:        os.Stderr,
@@ -55,28 +108,129 @@ func main() {
 Tips:
 1. To change the agent name, you can set the AGENT_NAME environment variable.
 2. To change the agent port, you can set the AGENT_PORT environment variable.
-`)
+	`)
 
 	t := task.RootTask("agent", false)
-	opts := server.Options{
-		CACert:     caCert,
-		ServerCert: srvCert,
-		Port:       env.AgentPort,
-	}
 
-	server.StartAgentServer(t, opts)
-
-	tcpListener, err := net.ListenTCP("tcp", &net.TCPAddr{Port: env.AgentStreamPort})
+	// One TCP listener on AGENT_PORT, then multiplex by TLS ALPN:
+	// - Stream ALPN: route to TCP stream tunnel handler
+	// - Otherwise: route to HTTPS API handler
+	tcpListener, err := net.ListenTCP("tcp", &net.TCPAddr{Port: env.AgentPort})
 	if err != nil {
 		gperr.LogFatal("failed to listen on port", err)
 	}
-	tcpServer := stream.NewTCPServer(t.Context(), tcpListener, caCert.Leaf, srvCert)
-	go tcpServer.Start()
-	log.Info().Int("port", env.AgentStreamPort).Msg("TCP stream server started")
 
-	udpServer := stream.NewUDPServer(t.Context(), &net.UDPAddr{Port: env.AgentStreamPort}, caCert.Leaf, srvCert)
-	go udpServer.Start()
-	log.Info().Int("port", env.AgentStreamPort).Msg("UDP stream server started")
+	caCertPool := x509.NewCertPool()
+	caCertPool.AddCert(caCert.Leaf)
+
+	muxTLSConfig := &tls.Config{
+		Certificates: []tls.Certificate{*srvCert},
+		ClientCAs:    caCertPool,
+		ClientAuth:   tls.RequireAndVerifyClientCert,
+		MinVersion:   tls.VersionTLS12,
+		// Keep HTTP limited to HTTP/1.1 (matching current agent server behavior)
+		// and add the stream tunnel ALPN for multiplexing.
+		NextProtos: []string{"http/1.1", stream.StreamALPN},
+	}
+	if env.AgentSkipClientCertCheck {
+		muxTLSConfig.ClientAuth = tls.NoClientCert
+	}
+
+	httpLn := newConnQueueListener(tcpListener.Addr(), 128)
+	streamLn := newConnQueueListener(tcpListener.Addr(), 128)
+
+	httpSrv := &http.Server{
+		Handler: handler.NewAgentHandler(),
+		BaseContext: func(net.Listener) context.Context {
+			return t.Context()
+		},
+	}
+	{
+		subtask := t.Subtask("agent-http", true)
+		t.OnCancel("stop_http", func() {
+			_ = httpSrv.Shutdown(context.Background())
+			_ = httpLn.Close()
+		})
+		go func() {
+			err := httpSrv.Serve(httpLn)
+			if err != nil && !errors.Is(err, http.ErrServerClosed) {
+				log.Error().Err(err).Msg("agent HTTP server stopped with error")
+			}
+			subtask.Finish(err)
+		}()
+		log.Info().Int("port", env.AgentPort).Msg("HTTPS API server started")
+	}
+
+	{
+		tcpServer := stream.NewTCPServerFromListener(t.Context(), streamLn)
+		subtask := t.Subtask("agent-stream-tcp", true)
+		t.OnCancel("stop_stream_tcp", func() {
+			_ = tcpServer.Close()
+			_ = streamLn.Close()
+		})
+		go func() {
+			err := tcpServer.Start()
+			subtask.Finish(err)
+		}()
+		log.Info().Int("port", env.AgentPort).Msg("TCP stream server started")
+	}
+
+	{
+		udpServer := stream.NewUDPServer(t.Context(), &net.UDPAddr{Port: env.AgentPort}, caCert.Leaf, srvCert)
+		subtask := t.Subtask("agent-stream-udp", true)
+		t.OnCancel("stop_stream_udp", func() {
+			_ = udpServer.Close()
+		})
+		go func() {
+			err := udpServer.Start()
+			subtask.Finish(err)
+		}()
+		log.Info().Int("port", env.AgentPort).Msg("UDP stream server started")
+	}
+
+	// Accept raw TCP connections, terminate TLS once, and dispatch by ALPN.
+	{
+		subtask := t.Subtask("agent-tls-mux", true)
+		t.OnCancel("stop_mux", func() {
+			_ = tcpListener.Close()
+			_ = httpLn.Close()
+			_ = streamLn.Close()
+		})
+		go func() {
+			defer subtask.Finish(subtask.FinishCause())
+			for {
+				select {
+				case <-t.Context().Done():
+					return
+				default:
+				}
+
+				conn, err := tcpListener.Accept()
+				if err != nil {
+					if t.Context().Err() != nil {
+						return
+					}
+					log.Error().Err(err).Msg("failed to accept connection")
+					continue
+				}
+
+				tlsConn := tls.Server(conn, muxTLSConfig)
+				if err := tlsConn.HandshakeContext(t.Context()); err != nil {
+					_ = tlsConn.Close()
+					log.Debug().Err(err).Msg("TLS handshake failed")
+					continue
+				}
+
+				alpn := tlsConn.ConnectionState().NegotiatedProtocol
+				switch alpn {
+				case stream.StreamALPN:
+					_ = streamLn.push(tlsConn)
+				default:
+					_ = httpLn.push(tlsConn)
+				}
+			}
+		}()
+	}
 
 	if socketproxy.ListenAddr != "" {
 		runtime := strutils.Title(string(env.Runtime))