rdp-proxy/agents/rap-node-agent/internal/vpnruntime/fabric_session_registry.go

package vpnruntime

import (
	"context"
	"sync"
	"time"

	"github.com/example/remote-access-platform/agents/rap-node-agent/internal/fabricproto"
	"github.com/example/remote-access-platform/agents/rap-node-agent/internal/mesh"
)

type FabricSessionFrameWriter interface {
	SendFrame(context.Context, fabricproto.Frame) error
}

type FabricSessionPacketPeerRegistry struct {
	mu    sync.RWMutex
	peers map[string]FabricSessionPacketPeer
}

type FabricSessionPacketPeer struct {
	VPNConnectionID         string
	Sender                  FabricSessionFrameWriter
	StreamID                uint64
	StreamIDsByTrafficClass map[string][]uint64
	RegisteredAt            time.Time
	LastPacketAt            time.Time
}

type FabricSessionPacketPeerTransport struct {
	Registry        *FabricSessionPacketPeerRegistry
	Inbox           *FabricPacketInbox
	VPNConnectionID string
}

func NewFabricSessionPacketPeerRegistry() *FabricSessionPacketPeerRegistry {
	return &FabricSessionPacketPeerRegistry{peers: map[string]FabricSessionPacketPeer{}}
}

func (r *FabricSessionPacketPeerRegistry) RegisterFrame(ctx context.Context, sender FabricSessionFrameWriter, frame fabricproto.Frame) (bool, error) {
	if r == nil || sender == nil || frame.Type != fabricproto.FrameData || frame.StreamID == 0 {
		return false, nil
	}
	payload, err := DecodeFabricVPNPacketDataFrame(frame)
	if err != nil {
		return false, nil
	}
	if payload.VPNConnectionID == "" {
		return false, nil
	}
	now := time.Now().UTC()
	r.mu.Lock()
	if r.peers == nil {
		r.peers = map[string]FabricSessionPacketPeer{}
	}
	peer := r.peers[payload.VPNConnectionID]
	if peer.RegisteredAt.IsZero() {
		peer.RegisteredAt = now
	}
	peer.VPNConnectionID = payload.VPNConnectionID
	peer.Sender = sender
	peer.StreamID = frame.StreamID
	peer.LastPacketAt = now
	if peer.StreamIDsByTrafficClass == nil {
		peer.StreamIDsByTrafficClass = map[string][]uint64{}
	}
	trafficClass := fabricSessionTrafficClassName(frame.TrafficClass)
	if trafficClass != "" && !containsUint64(peer.StreamIDsByTrafficClass[trafficClass], frame.StreamID) {
		peer.StreamIDsByTrafficClass[trafficClass] = append(peer.StreamIDsByTrafficClass[trafficClass], frame.StreamID)
	}
	r.peers[payload.VPNConnectionID] = peer
	r.mu.Unlock()
	return true, nil
}

func (r *FabricSessionPacketPeerRegistry) TransportFor(vpnConnectionID string, inbox *FabricPacketInbox) PacketTransport {
	if r == nil || inbox == nil || vpnConnectionID == "" {
		return nil
	}
	r.mu.RLock()
	peer, ok := r.peers[vpnConnectionID]
	r.mu.RUnlock()
	if !ok || peer.Sender == nil || peer.StreamID == 0 {
		return nil
	}
	return &FabricSessionPacketTransport{
		Sender:                  fabricSessionFrameWriterAdapter{writer: peer.Sender},
		Inbox:                   inbox,
		StreamID:                peer.StreamID,
		StreamIDsByTrafficClass: copyStreamIDsByClass(peer.StreamIDsByTrafficClass),
		VPNConnectionID:         vpnConnectionID,
		SendDirection:           FabricDirectionGatewayToClient,
		ReceiveDirection:        FabricDirectionClientToGateway,
	}
}

func (t *FabricSessionPacketPeerTransport) SendGatewayPacketBatch(ctx context.Context, packets [][]byte) error {
	if t == nil || t.Registry == nil || t.Inbox == nil || t.VPNConnectionID == "" {
		return mesh.ErrForwardRuntimeUnavailable
	}
	transport := t.Registry.TransportFor(t.VPNConnectionID, t.Inbox)
	if transport == nil {
		return mesh.ErrForwardRuntimeUnavailable
	}
	return transport.SendGatewayPacketBatch(ctx, packets)
}

func (t *FabricSessionPacketPeerTransport) ReceiveGatewayPacketBatch(ctx context.Context, timeout time.Duration) ([][]byte, error) {
	if t == nil || t.Inbox == nil || t.VPNConnectionID == "" {
		return nil, mesh.ErrForwardRuntimeUnavailable
	}
	return t.Inbox.Receive(ctx, t.VPNConnectionID, FabricDirectionClientToGateway, timeout)
}

func (t *FabricSessionPacketPeerTransport) Snapshot() map[string]any {
	if t == nil {
		return map[string]any{
			"transport":  "fabric_session_peer_dynamic",
			"peer_ready": false,
		}
	}
	ready := 0
	if t.Registry != nil {
		if transport := t.Registry.TransportFor(t.VPNConnectionID, t.Inbox); transport != nil {
			ready = 1
		}
	}
	return map[string]any{
		"transport":         "fabric_session_peer_dynamic",
		"vpn_connection_id": t.VPNConnectionID,
		"peer_ready":        ready == 1,
	}
}

func (r *FabricSessionPacketPeerRegistry) Snapshot() map[string]any {
	if r == nil {
		return map[string]any{"ready": 0}
	}
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := map[string]any{"ready": len(r.peers)}
	items := make([]map[string]any, 0, len(r.peers))
	for _, peer := range r.peers {
		item := map[string]any{
			"vpn_connection_id": peer.VPNConnectionID,
			"stream_id":         peer.StreamID,
		}
		if !peer.RegisteredAt.IsZero() {
			item["registered_at"] = peer.RegisteredAt.Format(time.RFC3339Nano)
		}
		if !peer.LastPacketAt.IsZero() {
			item["last_packet_at"] = peer.LastPacketAt.Format(time.RFC3339Nano)
		}
		items = append(items, item)
	}
	out["peers"] = items
	return out
}

type fabricSessionFrameWriterAdapter struct {
	writer FabricSessionFrameWriter
}

func (a fabricSessionFrameWriterAdapter) Send(ctx context.Context, frame fabricproto.Frame) error {
	if a.writer == nil {
		return mesh.ErrForwardRuntimeUnavailable
	}
	return a.writer.SendFrame(ctx, frame)
}

func containsUint64(values []uint64, value uint64) bool {
	for _, item := range values {
		if item == value {
			return true
		}
	}
	return false
}

func copyStreamIDsByClass(values map[string][]uint64) map[string][]uint64 {
	if len(values) == 0 {
		return nil
	}
	out := make(map[string][]uint64, len(values))
	for key, ids := range values {
		out[key] = append([]uint64(nil), ids...)
	}
	return out
}