first commit

pkg/config/config.go

@@ -0,0 +1,62 @@
+package config
+
+import (
+	"os"
+
+	"github.com/gotunnel/pkg/protocol"
+	"gopkg.in/yaml.v3"
+)
+
+// ServerConfig 服务端配置
+type ServerConfig struct {
+	Server  ServerSettings `yaml:"server"`
+	Clients []ClientConfig `yaml:"clients"`
+}
+
+// ServerSettings 服务端设置
+type ServerSettings struct {
+	BindAddr      string `yaml:"bind_addr"`
+	BindPort      int    `yaml:"bind_port"`
+	Token         string `yaml:"token"`
+	HeartbeatSec  int    `yaml:"heartbeat_sec"`
+	HeartbeatTimeout int `yaml:"heartbeat_timeout"`
+}
+
+// ClientConfig 客户端配置（服务端维护）
+type ClientConfig struct {
+	ID    string              `yaml:"id"`
+	Rules []protocol.ProxyRule `yaml:"rules"`
+}
+
+// LoadServerConfig 加载服务端配置
+func LoadServerConfig(path string) (*ServerConfig, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+
+	var cfg ServerConfig
+	if err := yaml.Unmarshal(data, &cfg); err != nil {
+		return nil, err
+	}
+
+	// 设置默认值
+	if cfg.Server.HeartbeatSec == 0 {
+		cfg.Server.HeartbeatSec = 30
+	}
+	if cfg.Server.HeartbeatTimeout == 0 {
+		cfg.Server.HeartbeatTimeout = 90
+	}
+
+	return &cfg, nil
+}
+
+// GetClientRules 获取指定客户端的代理规则
+func (c *ServerConfig) GetClientRules(clientID string) []protocol.ProxyRule {
+	for _, client := range c.Clients {
+		if client.ID == clientID {
+			return client.Rules
+		}
+	}
+	return nil
+}

pkg/protocol/message.go

@@ -0,0 +1,118 @@
+package protocol
+
+import (
+	"encoding/binary"
+	"encoding/json"
+	"errors"
+	"io"
+)
+
+// 消息类型定义
+const (
+	MsgTypeAuth         uint8 = 1  // 认证请求
+	MsgTypeAuthResp     uint8 = 2  // 认证响应
+	MsgTypeProxyConfig  uint8 = 3  // 代理配置下发
+	MsgTypeHeartbeat    uint8 = 4  // 心跳
+	MsgTypeHeartbeatAck uint8 = 5  // 心跳响应
+	MsgTypeNewProxy     uint8 = 6  // 新建代理连接请求
+	MsgTypeProxyReady   uint8 = 7  // 代理就绪
+	MsgTypeError        uint8 = 8  // 错误消息
+)
+
+// Message 基础消息结构
+type Message struct {
+	Type    uint8  `json:"type"`
+	Payload []byte `json:"payload"`
+}
+
+// AuthRequest 认证请求
+type AuthRequest struct {
+	ClientID string `json:"client_id"`
+	Token    string `json:"token"`
+}
+
+// AuthResponse 认证响应
+type AuthResponse struct {
+	Success bool   `json:"success"`
+	Message string `json:"message"`
+}
+
+// ProxyRule 代理规则
+type ProxyRule struct {
+	Name       string `json:"name" yaml:"name"`
+	LocalIP    string `json:"local_ip" yaml:"local_ip"`
+	LocalPort  int    `json:"local_port" yaml:"local_port"`
+	RemotePort int    `json:"remote_port" yaml:"remote_port"`
+}
+
+// ProxyConfig 代理配置下发
+type ProxyConfig struct {
+	Rules []ProxyRule `json:"rules"`
+}
+
+// NewProxyRequest 新建代理连接请求
+type NewProxyRequest struct {
+	RemotePort int `json:"remote_port"`
+}
+
+// ErrorMessage 错误消息
+type ErrorMessage struct {
+	Code    int    `json:"code"`
+	Message string `json:"message"`
+}
+
+// WriteMessage 写入消息到 writer
+func WriteMessage(w io.Writer, msg *Message) error {
+	// 消息格式: [1字节类型][4字节长度][payload]
+	header := make([]byte, 5)
+	header[0] = msg.Type
+	binary.BigEndian.PutUint32(header[1:], uint32(len(msg.Payload)))
+
+	if _, err := w.Write(header); err != nil {
+		return err
+	}
+	if len(msg.Payload) > 0 {
+		if _, err := w.Write(msg.Payload); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// ReadMessage 从 reader 读取消息
+func ReadMessage(r io.Reader) (*Message, error) {
+	header := make([]byte, 5)
+	if _, err := io.ReadFull(r, header); err != nil {
+		return nil, err
+	}
+
+	msgType := header[0]
+	length := binary.BigEndian.Uint32(header[1:])
+
+	if length > 1024*1024 { // 最大 1MB
+		return nil, errors.New("message too large")
+	}
+
+	payload := make([]byte, length)
+	if length > 0 {
+		if _, err := io.ReadFull(r, payload); err != nil {
+			return nil, err
+		}
+	}
+
+	return &Message{Type: msgType, Payload: payload}, nil
+}
+
+// NewMessage 创建新消息
+func NewMessage(msgType uint8, data interface{}) (*Message, error) {
+	payload, err := json.Marshal(data)
+	if err != nil {
+		return nil, err
+	}
+	return &Message{Type: msgType, Payload: payload}, nil
+}
+
+// ParsePayload 解析消息载荷
+func (m *Message) ParsePayload(v interface{}) error {
+	return json.Unmarshal(m.Payload, v)
+}

pkg/tunnel/client.go

@@ -0,0 +1,182 @@
+package tunnel
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/gotunnel/pkg/protocol"
+	"github.com/google/uuid"
+	"github.com/hashicorp/yamux"
+)
+
+// Client 隧道客户端
+type Client struct {
+	ServerAddr string
+	Token      string
+	ID         string
+	session    *yamux.Session
+	rules      []protocol.ProxyRule
+	mu         sync.RWMutex
+}
+
+// NewClient 创建客户端
+func NewClient(serverAddr, token, id string) *Client {
+	if id == "" {
+		id = uuid.New().String()[:8]
+	}
+	return &Client{
+		ServerAddr: serverAddr,
+		Token:      token,
+		ID:         id,
+	}
+}
+
+// Run 启动客户端（带断线重连）
+func (c *Client) Run() error {
+	for {
+		if err := c.connect(); err != nil {
+			log.Printf("[Client] Connect error: %v", err)
+			log.Printf("[Client] Reconnecting in 5s...")
+			time.Sleep(5 * time.Second)
+			continue
+		}
+
+		c.handleSession()
+		log.Printf("[Client] Disconnected, reconnecting...")
+		time.Sleep(3 * time.Second)
+	}
+}
+
+// connect 连接到服务端并认证
+func (c *Client) connect() error {
+	conn, err := net.DialTimeout("tcp", c.ServerAddr, 10*time.Second)
+	if err != nil {
+		return err
+	}
+
+	// 发送认证
+	authReq := protocol.AuthRequest{ClientID: c.ID, Token: c.Token}
+	msg, _ := protocol.NewMessage(protocol.MsgTypeAuth, authReq)
+	if err := protocol.WriteMessage(conn, msg); err != nil {
+		conn.Close()
+		return err
+	}
+
+	// 读取响应
+	resp, err := protocol.ReadMessage(conn)
+	if err != nil {
+		conn.Close()
+		return err
+	}
+
+	var authResp protocol.AuthResponse
+	resp.ParsePayload(&authResp)
+	if !authResp.Success {
+		conn.Close()
+		return fmt.Errorf("auth failed: %s", authResp.Message)
+	}
+
+	log.Printf("[Client] Authenticated as %s", c.ID)
+
+	// 建立 Yamux 会话
+	session, err := yamux.Client(conn, nil)
+	if err != nil {
+		conn.Close()
+		return err
+	}
+
+	c.mu.Lock()
+	c.session = session
+	c.mu.Unlock()
+
+	return nil
+}
+
+// handleSession 处理会话
+func (c *Client) handleSession() {
+	defer c.session.Close()
+
+	for {
+		stream, err := c.session.Accept()
+		if err != nil {
+			return
+		}
+		go c.handleStream(stream)
+	}
+}
+
+// handleStream 处理流
+func (c *Client) handleStream(stream net.Conn) {
+	defer stream.Close()
+
+	msg, err := protocol.ReadMessage(stream)
+	if err != nil {
+		return
+	}
+
+	switch msg.Type {
+	case protocol.MsgTypeProxyConfig:
+		c.handleProxyConfig(msg)
+	case protocol.MsgTypeNewProxy:
+		c.handleNewProxy(stream, msg)
+	case protocol.MsgTypeHeartbeat:
+		c.handleHeartbeat(stream)
+	}
+}
+
+// handleProxyConfig 处理代理配置
+func (c *Client) handleProxyConfig(msg *protocol.Message) {
+	var cfg protocol.ProxyConfig
+	msg.ParsePayload(&cfg)
+
+	c.mu.Lock()
+	c.rules = cfg.Rules
+	c.mu.Unlock()
+
+	log.Printf("[Client] Received %d proxy rules", len(cfg.Rules))
+	for _, r := range cfg.Rules {
+		log.Printf("[Client]   %s: %s:%d", r.Name, r.LocalIP, r.LocalPort)
+	}
+}
+
+// handleNewProxy 处理新代理请求
+func (c *Client) handleNewProxy(stream net.Conn, msg *protocol.Message) {
+	var req protocol.NewProxyRequest
+	msg.ParsePayload(&req)
+
+	// 查找对应规则
+	var rule *protocol.ProxyRule
+	c.mu.RLock()
+	for _, r := range c.rules {
+		if r.RemotePort == req.RemotePort {
+			rule = &r
+			break
+		}
+	}
+	c.mu.RUnlock()
+
+	if rule == nil {
+		log.Printf("[Client] Unknown port %d", req.RemotePort)
+		return
+	}
+
+	// 连接本地服务
+	localAddr := fmt.Sprintf("%s:%d", rule.LocalIP, rule.LocalPort)
+	localConn, err := net.DialTimeout("tcp", localAddr, 5*time.Second)
+	if err != nil {
+		log.Printf("[Client] Connect %s error: %v", localAddr, err)
+		return
+	}
+
+	// 双向转发
+	relay(stream, localConn)
+}
+
+// handleHeartbeat 处理心跳
+func (c *Client) handleHeartbeat(stream net.Conn) {
+	msg := &protocol.Message{Type: protocol.MsgTypeHeartbeatAck}
+	protocol.WriteMessage(stream, msg)
+}

pkg/tunnel/server.go

@@ -0,0 +1,302 @@
+package tunnel
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/gotunnel/pkg/config"
+	"github.com/gotunnel/pkg/protocol"
+	"github.com/gotunnel/pkg/utils"
+	"github.com/hashicorp/yamux"
+)
+
+// Server 隧道服务端
+type Server struct {
+	config      *config.ServerConfig
+	portManager *utils.PortManager
+	clients     map[string]*ClientSession
+	mu          sync.RWMutex
+}
+
+// ClientSession 客户端会话
+type ClientSession struct {
+	ID        string
+	Session   *yamux.Session
+	Rules     []protocol.ProxyRule
+	Listeners map[int]net.Listener
+	LastPing  time.Time
+	mu        sync.Mutex
+}
+
+// NewServer 创建服务端
+func NewServer(cfg *config.ServerConfig) *Server {
+	return &Server{
+		config:      cfg,
+		portManager: utils.NewPortManager(),
+		clients:     make(map[string]*ClientSession),
+	}
+}
+
+// Run 启动服务端
+func (s *Server) Run() error {
+	addr := fmt.Sprintf("%s:%d", s.config.Server.BindAddr, s.config.Server.BindPort)
+	ln, err := net.Listen("tcp", addr)
+	if err != nil {
+		return fmt.Errorf("failed to listen on %s: %v", addr, err)
+	}
+	defer ln.Close()
+
+	log.Printf("[Server] Listening on %s", addr)
+
+	for {
+		conn, err := ln.Accept()
+		if err != nil {
+			log.Printf("[Server] Accept error: %v", err)
+			continue
+		}
+		go s.handleConnection(conn)
+	}
+}
+
+// handleConnection 处理客户端连接
+func (s *Server) handleConnection(conn net.Conn) {
+	defer conn.Close()
+
+	// 设置认证超时
+	conn.SetReadDeadline(time.Now().Add(10 * time.Second))
+
+	// 读取认证消息
+	msg, err := protocol.ReadMessage(conn)
+	if err != nil {
+		log.Printf("[Server] Read auth error: %v", err)
+		return
+	}
+
+	if msg.Type != protocol.MsgTypeAuth {
+		log.Printf("[Server] Expected auth, got %d", msg.Type)
+		return
+	}
+
+	var authReq protocol.AuthRequest
+	if err := msg.ParsePayload(&authReq); err != nil {
+		log.Printf("[Server] Parse auth error: %v", err)
+		return
+	}
+
+	// 验证 Token
+	if authReq.Token != s.config.Server.Token {
+		s.sendAuthResponse(conn, false, "invalid token")
+		return
+	}
+
+	// 获取客户端配置
+	rules := s.config.GetClientRules(authReq.ClientID)
+	if rules == nil {
+		s.sendAuthResponse(conn, false, "client not configured")
+		return
+	}
+
+	conn.SetReadDeadline(time.Time{})
+
+	if err := s.sendAuthResponse(conn, true, "ok"); err != nil {
+		return
+	}
+
+	log.Printf("[Server] Client %s authenticated", authReq.ClientID)
+	s.setupClientSession(conn, authReq.ClientID, rules)
+}
+
+// setupClientSession 建立客户端会话
+func (s *Server) setupClientSession(conn net.Conn, clientID string, rules []protocol.ProxyRule) {
+	session, err := yamux.Server(conn, nil)
+	if err != nil {
+		log.Printf("[Server] Yamux error: %v", err)
+		return
+	}
+
+	cs := &ClientSession{
+		ID:        clientID,
+		Session:   session,
+		Rules:     rules,
+		Listeners: make(map[int]net.Listener),
+		LastPing:  time.Now(),
+	}
+
+	s.registerClient(cs)
+	defer s.unregisterClient(cs)
+
+	if err := s.sendProxyConfig(session, rules); err != nil {
+		log.Printf("[Server] Send config error: %v", err)
+		return
+	}
+
+	s.startProxyListeners(cs)
+	go s.heartbeatLoop(cs)
+
+	<-session.CloseChan()
+	log.Printf("[Server] Client %s disconnected", clientID)
+}
+
+// sendAuthResponse 发送认证响应
+func (s *Server) sendAuthResponse(conn net.Conn, success bool, message string) error {
+	resp := protocol.AuthResponse{Success: success, Message: message}
+	msg, _ := protocol.NewMessage(protocol.MsgTypeAuthResp, resp)
+	return protocol.WriteMessage(conn, msg)
+}
+
+// sendProxyConfig 发送代理配置
+func (s *Server) sendProxyConfig(session *yamux.Session, rules []protocol.ProxyRule) error {
+	stream, err := session.Open()
+	if err != nil {
+		return err
+	}
+	defer stream.Close()
+
+	cfg := protocol.ProxyConfig{Rules: rules}
+	msg, _ := protocol.NewMessage(protocol.MsgTypeProxyConfig, cfg)
+	return protocol.WriteMessage(stream, msg)
+}
+
+// registerClient 注册客户端
+func (s *Server) registerClient(cs *ClientSession) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.clients[cs.ID] = cs
+}
+
+// unregisterClient 注销客户端
+func (s *Server) unregisterClient(cs *ClientSession) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// 关闭所有监听器
+	cs.mu.Lock()
+	for port, ln := range cs.Listeners {
+		ln.Close()
+		s.portManager.Release(port)
+	}
+	cs.mu.Unlock()
+
+	delete(s.clients, cs.ID)
+}
+
+// startProxyListeners 启动代理监听
+func (s *Server) startProxyListeners(cs *ClientSession) {
+	for _, rule := range cs.Rules {
+		if err := s.portManager.Reserve(rule.RemotePort, cs.ID); err != nil {
+			log.Printf("[Server] Port %d error: %v", rule.RemotePort, err)
+			continue
+		}
+
+		ln, err := net.Listen("tcp", fmt.Sprintf(":%d", rule.RemotePort))
+		if err != nil {
+			log.Printf("[Server] Listen %d error: %v", rule.RemotePort, err)
+			s.portManager.Release(rule.RemotePort)
+			continue
+		}
+
+		cs.mu.Lock()
+		cs.Listeners[rule.RemotePort] = ln
+		cs.mu.Unlock()
+
+		log.Printf("[Server] Proxy %s: :%d -> %s:%d",
+			rule.Name, rule.RemotePort, rule.LocalIP, rule.LocalPort)
+
+		go s.acceptProxyConns(cs, ln, rule)
+	}
+}
+
+// acceptProxyConns 接受代理连接
+func (s *Server) acceptProxyConns(cs *ClientSession, ln net.Listener, rule protocol.ProxyRule) {
+	for {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		go s.handleProxyConn(cs, conn, rule)
+	}
+}
+
+// handleProxyConn 处理代理连接
+func (s *Server) handleProxyConn(cs *ClientSession, conn net.Conn, rule protocol.ProxyRule) {
+	defer conn.Close()
+
+	// 打开到客户端的流
+	stream, err := cs.Session.Open()
+	if err != nil {
+		log.Printf("[Server] Open stream error: %v", err)
+		return
+	}
+	defer stream.Close()
+
+	// 发送新代理请求
+	req := protocol.NewProxyRequest{RemotePort: rule.RemotePort}
+	msg, _ := protocol.NewMessage(protocol.MsgTypeNewProxy, req)
+	if err := protocol.WriteMessage(stream, msg); err != nil {
+		return
+	}
+
+	// 双向转发
+	relay(conn, stream)
+}
+
+// relay 双向数据转发
+func relay(c1, c2 net.Conn) {
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	copy := func(dst, src net.Conn) {
+		defer wg.Done()
+		buf := make([]byte, 32*1024)
+		for {
+			n, err := src.Read(buf)
+			if n > 0 {
+				dst.Write(buf[:n])
+			}
+			if err != nil {
+				return
+			}
+		}
+	}
+
+	go copy(c1, c2)
+	go copy(c2, c1)
+	wg.Wait()
+}
+
+// heartbeatLoop 心跳检测循环
+func (s *Server) heartbeatLoop(cs *ClientSession) {
+	ticker := time.NewTicker(time.Duration(s.config.Server.HeartbeatSec) * time.Second)
+	defer ticker.Stop()
+
+	timeout := time.Duration(s.config.Server.HeartbeatTimeout) * time.Second
+
+	for {
+		select {
+		case <-ticker.C:
+			cs.mu.Lock()
+			if time.Since(cs.LastPing) > timeout {
+				cs.mu.Unlock()
+				log.Printf("[Server] Client %s heartbeat timeout", cs.ID)
+				cs.Session.Close()
+				return
+			}
+			cs.mu.Unlock()
+
+			// 发送心跳
+			stream, err := cs.Session.Open()
+			if err != nil {
+				return
+			}
+			msg := &protocol.Message{Type: protocol.MsgTypeHeartbeat}
+			protocol.WriteMessage(stream, msg)
+			stream.Close()
+
+		case <-cs.Session.CloseChan():
+			return
+		}
+	}
+}

pkg/utils/port.go

@@ -0,0 +1,78 @@
+package utils
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"time"
+)
+
+// PortManager 端口管理器
+type PortManager struct {
+	mu       sync.RWMutex
+	occupied map[int]string // port -> clientID
+}
+
+// NewPortManager 创建端口管理器
+func NewPortManager() *PortManager {
+	return &PortManager{
+		occupied: make(map[int]string),
+	}
+}
+
+// IsPortAvailable 检测端口是否可用（系统级）
+func IsPortAvailable(port int) bool {
+	addr := fmt.Sprintf(":%d", port)
+	ln, err := net.Listen("tcp", addr)
+	if err != nil {
+		return false
+	}
+	ln.Close()
+	return true
+}
+
+// Reserve 预留端口给指定客户端
+func (pm *PortManager) Reserve(port int, clientID string) error {
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+
+	if owner, exists := pm.occupied[port]; exists {
+		return fmt.Errorf("port %d already occupied by client %s", port, owner)
+	}
+
+	if !IsPortAvailable(port) {
+		return fmt.Errorf("port %d is occupied by system", port)
+	}
+
+	pm.occupied[port] = clientID
+	return nil
+}
+
+// Release 释放端口
+func (pm *PortManager) Release(port int) {
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+	delete(pm.occupied, port)
+}
+
+// ReleaseByClient 释放指定客户端的所有端口
+func (pm *PortManager) ReleaseByClient(clientID string) {
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+	for port, owner := range pm.occupied {
+		if owner == clientID {
+			delete(pm.occupied, port)
+		}
+	}
+}
+
+// CheckLocalService 检测本地服务是否可连接
+func CheckLocalService(ip string, port int, timeout time.Duration) error {
+	addr := fmt.Sprintf("%s:%d", ip, port)
+	conn, err := net.DialTimeout("tcp", addr, timeout)
+	if err != nil {
+		return fmt.Errorf("cannot connect to local service %s: %v", addr, err)
+	}
+	conn.Close()
+	return nil
+}