zogo/spmux/spmux.go

// Package spmux 单端口多路复用：同一端口上按报文特征把流量分发给 HTTP / HTTPS / TCP 三个 Listener。
//
// 迁移自 go-hua/spmux：原依赖的 go-hua/log 已弃用，日志改为两方法接口 Logger
// （*logrus.Logger 天然满足），传 nil 时静默；协议识别改用 httpprobe 包。
package spmux

import (
	"fmt"
	"net"
	"sync"
	"time"

	"git.zeroonesoft.cn/golib/zogo/httpprobe"
)

// Logger spmux 内部日志接口：仅需 Debugf / Errorf 两个方法。
// *logrus.Logger（zogo/logger 初始化的实例）天然满足该接口。
type Logger interface {
	Debugf(format string, a ...any)
	Errorf(format string, a ...any)
}

// AcceptTimeout 分发连接时等待消费方接收的超时（秒）
const AcceptTimeout = 100

// ChannelBufferSize 连接分发通道缓冲大小
const ChannelBufferSize = 1

// FirstPeekSize 协议识别时窥探的首包字节数
const FirstPeekSize = 1

// PortMux 单端口多路复用器，按协议特征将连接分发到 HTTP/HTTPS/TCP 三个通道
type PortMux struct {
	listener  net.Listener
	logger    Logger
	once      sync.Once
	port      int
	tcpConn   chan *PortConn
	httpConn  chan *PortConn
	httpsConn chan *PortConn
}

// NewPortMux 创建端口复用器，logger 用于内部错误输出，可为 nil（静默）
func NewPortMux(port int, logger Logger) *PortMux {
	p := &PortMux{
		port:      port,
		logger:    logger,
		tcpConn:   make(chan *PortConn),
		httpConn:  make(chan *PortConn),
		httpsConn: make(chan *PortConn),
	}
	return p
}

// Start 开始监听指定端口并后台分发连接
func (p *PortMux) Start() error {
	addr := fmt.Sprintf("0.0.0.0:%d", p.port)
	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return err
	}
	p.listener, err = net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return err
	}
	go p.listen()
	return nil
}

// listen 接受连接循环
func (p *PortMux) listen() {
	for {
		conn, err := p.listener.Accept()
		if err != nil {
			p.logErrorf("接受连接失败 端口:%d 错误:%s", p.port, err.Error())
			p.Close()
			break
		}
		p.logDebugf("新连接 地址:%s", conn.RemoteAddr().String())
		go p.process(conn)
	}
}

// process 识别单个连接的协议并分发到对应通道
func (p *PortMux) process(conn net.Conn) {
	newConn := newPortConn(conn)
	buf, err := newConn.Peek(FirstPeekSize)
	if err != nil {
		conn.Close()
		return
	}

	n := newConn.Buffered()
	if n >= 8 {
		buf, err = newConn.Peek(n)
		if err != nil {
			conn.Close()
			return
		}
	}

	if n >= 8 && httpprobe.IsHttp(buf) {
		timer := time.NewTimer(AcceptTimeout)
		select {
		case <-timer.C:
			conn.Close()
			p.logErrorf("接受超时 端口:%d 协议:%s", p.port, "http")
		case p.httpConn <- newConn:
		}
	} else if n >= 8 && httpprobe.IsHttps(buf) {
		timer := time.NewTimer(AcceptTimeout)
		select {
		case <-timer.C:
			conn.Close()
			p.logErrorf("接受超时 端口:%d 协议:%s", p.port, "https")
		case p.httpsConn <- newConn:
		}
	} else {
		timer := time.NewTimer(AcceptTimeout)
		select {
		case <-timer.C:
			conn.Close()
			p.logErrorf("接受超时 端口:%d 协议:%s", p.port, "tcp")
		case p.tcpConn <- newConn:
		}
	}
}

// Close 关闭监听与所有分发通道（幂等）
func (p *PortMux) Close() (err error) {
	p.once.Do(func() {
		close(p.tcpConn)
		close(p.httpsConn)
		close(p.httpConn)
		err = p.listener.Close()
	})
	return
}

// TcpListener TCP监听
// Deprecated: 一些网络通讯(RFB协议)是由服务器先发送数据的，会导致无法识别成TCP
func (p *PortMux) TcpListener() net.Listener {
	return NewPortListener(p.tcpConn, p.listener.Addr())
}

// HttpListener 返回 HTTP 流量的监听器
func (p *PortMux) HttpListener() net.Listener {
	return NewPortListener(p.httpConn, p.listener.Addr())
}

// HttpsListener 返回 HTTPS 流量的监听器
func (p *PortMux) HttpsListener() net.Listener {
	return NewPortListener(p.httpsConn, p.listener.Addr())
}

// logErrorf nil 安全的错误日志
func (p *PortMux) logErrorf(format string, a ...any) {
	if p.logger != nil {
		p.logger.Errorf(format, a...)
	}
}

// logDebugf nil 安全的调试日志
func (p *PortMux) logDebugf(format string, a ...any) {
	if p.logger != nil {
		p.logger.Debugf(format, a...)
	}
}