feat: initial commit

internal/server/forward.go

@@ -0,0 +1,149 @@
+package server
+
+import (
+	"encoding/base64"
+	"fmt"
+	"log"
+	"net"
+	"sync"
+
+	"github.com/x-grk/longpolling-remote-access/internal/proto"
+)
+
+// forwardManager owns the host-side listeners for port-forward rules. It
+// supports three transports:
+//   - "http": tunneled over the HTTP long-poll channel (works anywhere nginx
+//     reaches, highest latency);
+//   - "tcp":  direct TCP multiplexed over the 8446 tunnel (low latency);
+//   - "udp":  direct UDP over the 8446 channel.
+type forwardManager struct {
+	hub *Hub
+	tun *tunnelHub
+	udp *udpRelay
+
+	mu  sync.Mutex
+	lst map[int64]net.Listener // TCP listeners for http/tcp rules
+}
+
+func newForwardManager(hub *Hub, tun *tunnelHub, udp *udpRelay) *forwardManager {
+	return &forwardManager{hub: hub, tun: tun, udp: udp, lst: map[int64]net.Listener{}}
+}
+
+// open starts the host-side listener(s) for a rule.
+func (m *forwardManager) open(rule ForwardRule) error {
+	m.close(rule.ID)
+	switch rule.Proto {
+	case "udp":
+		if err := m.udp.openRulePort(rule); err != nil {
+			return err
+		}
+		log.Printf("forward[udp]: %s:%d -> agent %s :%d", rule.ListenAddr, rule.HostPort, rule.AgentID, rule.ClientPort)
+		return nil
+	default: // "http" and "tcp" both listen TCP on the host port
+		addr := fmt.Sprintf("%s:%d", rule.ListenAddr, rule.HostPort)
+		ln, err := net.Listen("tcp", addr)
+		if err != nil {
+			return err
+		}
+		m.mu.Lock()
+		m.lst[rule.ID] = ln
+		m.mu.Unlock()
+		log.Printf("forward[%s]: %s -> agent %s :%d", rule.Proto, addr, rule.AgentID, rule.ClientPort)
+		go m.acceptTCP(rule, ln)
+		return nil
+	}
+}
+
+func (m *forwardManager) close(ruleID int64) {
+	m.mu.Lock()
+	ln := m.lst[ruleID]
+	delete(m.lst, ruleID)
+	m.mu.Unlock()
+	if ln != nil {
+		_ = ln.Close()
+	}
+	m.udp.closeRulePort(ruleID)
+}
+
+func (m *forwardManager) closeAll() {
+	m.mu.Lock()
+	lst := m.lst
+	m.lst = map[int64]net.Listener{}
+	m.mu.Unlock()
+	for _, ln := range lst {
+		_ = ln.Close()
+	}
+}
+
+func (m *forwardManager) acceptTCP(rule ForwardRule, ln net.Listener) {
+	for {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		if tc, ok := conn.(*net.TCPConn); ok {
+			_ = tc.SetNoDelay(true)
+		}
+		if rule.Proto == "tcp" {
+			go m.bridgeDirect(rule, conn)
+		} else {
+			go m.bridgeHTTP(rule, conn)
+		}
+	}
+}
+
+// bridgeDirect wires a connection to the agent's low-latency 8446 tunnel.
+func (m *forwardManager) bridgeDirect(rule ForwardRule, conn net.Conn) {
+	at := m.tun.get(rule.AgentID)
+	if at == nil {
+		// No direct tunnel (agent offline or 8446 unreachable).
+		conn.Close()
+		return
+	}
+	at.openStream(rule.ClientPort, conn)
+}
+
+// bridgeHTTP wires a connection through the HTTP long-poll hub session.
+func (m *forwardManager) bridgeHTTP(rule ForwardRule, conn net.Conn) {
+	sess := m.hub.CreateForward(rule.AgentID, rule.ClientPort)
+	defer m.hub.CloseSession(sess.ID)
+	defer conn.Close()
+
+	go func() {
+		defer m.hub.CloseSession(sess.ID)
+		buf := make([]byte, 32*1024)
+		for {
+			n, err := conn.Read(buf)
+			if n > 0 {
+				sess.PushIn(proto.InFrame{
+					Type: proto.FrameStdin,
+					Data: base64.StdEncoding.EncodeToString(buf[:n]),
+				})
+			}
+			if err != nil {
+				sess.PushIn(proto.InFrame{Type: proto.FrameClose})
+				return
+			}
+		}
+	}()
+
+	for {
+		select {
+		case <-sess.Done():
+			return
+		case out, ok := <-sess.Out():
+			if !ok || out.Exit {
+				return
+			}
+			if out.Data != "" {
+				b, err := base64.StdEncoding.DecodeString(out.Data)
+				if err != nil {
+					return
+				}
+				if _, err := conn.Write(b); err != nil {
+					return
+				}
+			}
+		}
+	}
+}