first commit

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
+		}
+	}
+}