rmm2/agent/ws/tunnel.go

package ws

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"
)

// TunnelSession repraesentiert eine einzelne TCP-Verbindung zu einem Ziel
type TunnelSession struct {
	ID         string
	TunnelID   string
	TargetHost string
	TargetPort int
	Conn       net.Conn
	Active     bool
}

// TunnelManager verwaltet alle Sessions auf Agent-Seite
type TunnelManager struct {
	client   *Client
	sessions map[string]*TunnelSession
	mutex    sync.RWMutex
}

func NewTunnelManager(client *Client) *TunnelManager {
	return &TunnelManager{
		client:   client,
		sessions: make(map[string]*TunnelSession),
	}
}

// ConnectSession oeffnet eine neue TCP-Verbindung fuer eine Session
func (tm *TunnelManager) ConnectSession(sessionID, tunnelID, targetHost string, targetPort int) error {
	target := fmt.Sprintf("%s:%d", targetHost, targetPort)
	conn, err := net.DialTimeout("tcp", target, 10*time.Second)
	if err != nil {
		return fmt.Errorf("Verbindung zu %s fehlgeschlagen: %w", target, err)
	}

	session := &TunnelSession{
		ID:         sessionID,
		TunnelID:   tunnelID,
		TargetHost: targetHost,
		TargetPort: targetPort,
		Conn:       conn,
		Active:     true,
	}

	tm.mutex.Lock()
	tm.sessions[sessionID] = session
	tm.mutex.Unlock()

	// Ziel -> Backend Forwarding starten
	go tm.forwardFromTarget(session)

	log.Printf("Session %s: Verbindung zu %s hergestellt", sessionID, target)
	return nil
}

// DisconnectSession schliesst eine Session
func (tm *TunnelManager) DisconnectSession(sessionID string) {
	tm.mutex.Lock()
	session, exists := tm.sessions[sessionID]
	if !exists {
		tm.mutex.Unlock()
		return
	}

	session.Active = false
	delete(tm.sessions, sessionID)
	tm.mutex.Unlock()

	if session.Conn != nil {
		session.Conn.Close()
	}

	log.Printf("Session %s: geschlossen", sessionID)
}

// SendDataToSession schreibt Daten in die Ziel-Verbindung einer Session
func (tm *TunnelManager) SendDataToSession(sessionID string, data []byte) error {
	tm.mutex.RLock()
	session, exists := tm.sessions[sessionID]
	tm.mutex.RUnlock()

	if !exists || !session.Active {
		return fmt.Errorf("Session %s nicht aktiv", sessionID)
	}

	_, err := session.Conn.Write(data)
	if err != nil {
		tm.DisconnectSession(sessionID)
		return fmt.Errorf("Session %s: Write fehlgeschlagen: %w", sessionID, err)
	}

	return nil
}

// forwardFromTarget liest Daten vom Ziel und schickt sie als Binary-Message ans Backend
func (tm *TunnelManager) forwardFromTarget(session *TunnelSession) {
	defer func() {
		if session.Active {
			tm.DisconnectSession(session.ID)
		}
		log.Printf("Session %s: Forwarding beendet", session.ID)
	}()

	buffer := make([]byte, 32768)

	for session.Active {
		session.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))

		n, err := session.Conn.Read(buffer)
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			if err == io.EOF {
				log.Printf("Session %s: Ziel-Verbindung geschlossen", session.ID)
			} else if session.Active {
				log.Printf("Session %s: Read-Fehler: %v", session.ID, err)
			}
			break
		}

		if n > 0 {
			if err := tm.client.messageHandler.SendSessionData(session.ID, buffer[:n]); err != nil {
				log.Printf("Session %s: WebSocket senden fehlgeschlagen: %v", session.ID, err)
				break
			}
		}
	}
}

// ProcessBinaryMessage verarbeitet eingehende Binary-Messages vom Backend (Proxy -> Ziel)
func (tm *TunnelManager) ProcessBinaryMessage(data []byte) error {
	sessionID, payload, err := ParseBinaryMessage(data)
	if err != nil {
		return fmt.Errorf("Binary-Message parsen: %w", err)
	}

	return tm.SendDataToSession(sessionID, payload)
}

// GetActiveSessions gibt die Anzahl aktiver Sessions zurueck
func (tm *TunnelManager) GetActiveSessions() int {
	tm.mutex.RLock()
	defer tm.mutex.RUnlock()
	return len(tm.sessions)
}

func (tm *TunnelManager) generateID() string {
	bytes := make([]byte, 8)
	rand.Read(bytes)
	return hex.EncodeToString(bytes)
}

// ParseBinaryMessage parst das Binary-Format: [id_len:1][id][data]
func ParseBinaryMessage(data []byte) (id string, payload []byte, err error) {
	if len(data) < 1 {
		return "", nil, fmt.Errorf("Daten zu kurz")
	}

	idLen := int(data[0])
	if len(data) < 1+idLen {
		return "", nil, fmt.Errorf("ID zu kurz")
	}

	id = string(data[1 : 1+idLen])
	payload = data[1+idLen:]
	return id, payload, nil
}