lpra/internal/server/forward.go

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": over the agent's own link to the panel - its WebSocket, or the
//     HTTP long polls while there is none (works anywhere nginx reaches,
//     highest latency);
//   - "tcp":  direct TCP multiplexed over the tunnel port (low latency);
//   - "udp":  direct UDP over the tunnel port.
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 direct 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 the tunnel port unreachable).
		conn.Close()
		return
	}
	at.openStream(rule.ClientPort, conn)
}

// bridgeHTTP wires a connection through a hub session: the agent carries it
// over its WebSocket, or by long polls while it has none.
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
			}
		}
	}()

	// write passes one message on; false: the stream is over.
	write := func(out proto.OutMessage) bool {
		if out.Exit {
			return false
		}
		if out.Data != "" {
			b, err := base64.StdEncoding.DecodeString(out.Data)
			if err != nil {
				return false
			}
			if _, err := conn.Write(b); err != nil {
				return false
			}
		}
		return true
	}
	for {
		select {
		case out, ok := <-sess.Out():
			if !ok || !write(out) {
				return
			}
		case <-sess.Done():
			// The agent's end of stream closes the session at once, while
			// what it sent before may still be queued (a download read
			// faster than the browser takes it): that goes out first.
			for {
				select {
				case out := <-sess.Out():
					if !write(out) {
						return
					}
				default:
					return
				}
			}
		}
	}
}