m/rdp-proxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor RDP proxy handling and update related tests

Browse Source
This commit is contained in:
Mikhail
2026-05-17 20:38:35 +03:00
parent 8e9402580f
commit d551e57fd5
172 changed files with 22117 additions and 2509 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agents/rap-node-agent/internal/mesh/fabric_quic_transport.go
+473 -12
View File
@@ -6,7 +6,9 @@ import (
"crypto/tls"
"crypto/x509"
"encoding/hex"
"encoding/json"
"fmt"
"net"
"sort"
"strings"
"sync"
@@ -17,6 +19,7 @@ import (
)
const fabricQUICNextProto = "rap-fabric-data-session-v1"
const fabricQUICReverseHelloPrefix = "rap-fabric-reverse-hello-v1:"
const defaultQUICFabricConnIdleTTL = 5 * time.Minute
const defaultQUICFabricMaxStreamsPerConn = 64
const ErrQUICFabricStreamLimitReached = quicFabricError("quic fabric stream limit reached")
@@ -28,17 +31,29 @@ func (e quicFabricError) Error() string {
}
type QUICFabricTransport struct {
Config *quic.Config
IdleTTL time.Duration
MaxStreamsPerConn int
mu sync.Mutex
conns map[string]*quicFabricConnEntry
stats QUICFabricTransportStats
Config *quic.Config
LocalPeerID string
IdleTTL time.Duration
MaxStreamsPerConn int
DialAddr func(context.Context, string, *tls.Config, *quic.Config) (*quic.Conn, error)
mu sync.Mutex
conns map[string]*quicFabricConnEntry
reverseConns map[string]*quicFabricConnEntry
inboundProductionHandler func(context.Context, ProductionEnvelope) (ProductionForwardResult, error)
inboundWebIngressHandler func(context.Context, []byte) ([]byte, error)
inboundFabricControlHandler func(context.Context, []byte) ([]byte, error)
inboundSyntheticHandler func(context.Context, SyntheticEnvelope) (SyntheticEnvelope, error)
logger FabricSessionEventLogger
stats QUICFabricTransportStats
}
type QUICFabricTransportStats struct {
Opens uint64 `json:"opens"`
Reuses uint64 `json:"reuses"`
ReverseHelloSent uint64 `json:"reverse_hello_sent"`
ReverseHelloFailed uint64 `json:"reverse_hello_failed"`
ReverseRegisters uint64 `json:"reverse_registers"`
ReverseReuses uint64 `json:"reverse_reuses"`
OpenFailures uint64 `json:"open_failures"`
ClosedEvicted uint64 `json:"closed_evicted"`
CloseAllCalls uint64 `json:"close_all_calls"`
@@ -50,6 +65,7 @@ type QUICFabricTransportStats struct {
type QUICFabricTransportSnapshot struct {
SchemaVersion string `json:"schema_version"`
LocalPeerID string `json:"local_peer_id,omitempty"`
ActiveCount int `json:"active_count"`
ActiveStreams int `json:"active_streams"`
MaxStreamsPerConn int `json:"max_streams_per_conn"`
@@ -63,6 +79,7 @@ type QUICFabricConnSnapshot struct {
PeerID string `json:"peer_id,omitempty"`
Endpoint string `json:"endpoint,omitempty"`
CertSHA256 string `json:"cert_sha256,omitempty"`
Direction string `json:"direction,omitempty"`
ActiveStreams int `json:"active_streams"`
MaxStreams int `json:"max_streams"`
CapacityPressurePercent int `json:"capacity_pressure_percent"`
@@ -92,7 +109,41 @@ type quicFabricConnEntry struct {
}
func NewQUICFabricTransport(config *quic.Config) *QUICFabricTransport {
return &QUICFabricTransport{Config: config, IdleTTL: defaultQUICFabricConnIdleTTL, MaxStreamsPerConn: defaultQUICFabricMaxStreamsPerConn, conns: map[string]*quicFabricConnEntry{}}
return &QUICFabricTransport{Config: config, IdleTTL: defaultQUICFabricConnIdleTTL, MaxStreamsPerConn: defaultQUICFabricMaxStreamsPerConn, conns: map[string]*quicFabricConnEntry{}, reverseConns: map[string]*quicFabricConnEntry{}}
}
func (t *QUICFabricTransport) SetInboundHandlers(production func(context.Context, ProductionEnvelope) (ProductionForwardResult, error), synthetic func(context.Context, SyntheticEnvelope) (SyntheticEnvelope, error), logger FabricSessionEventLogger) {
t.SetInboundHandlersWithWebIngress(production, nil, synthetic, logger)
}
func (t *QUICFabricTransport) SetInboundHandlersWithWebIngress(production func(context.Context, ProductionEnvelope) (ProductionForwardResult, error), webIngress func(context.Context, []byte) ([]byte, error), synthetic func(context.Context, SyntheticEnvelope) (SyntheticEnvelope, error), logger FabricSessionEventLogger) {
if t == nil {
return
}
t.mu.Lock()
t.inboundProductionHandler = production
t.inboundWebIngressHandler = webIngress
t.inboundSyntheticHandler = synthetic
t.logger = logger
t.mu.Unlock()
}
func (t *QUICFabricTransport) SetInboundFabricControlHandler(handler func(context.Context, []byte) ([]byte, error)) {
if t == nil {
return
}
t.mu.Lock()
t.inboundFabricControlHandler = handler
t.mu.Unlock()
}
func (t *QUICFabricTransport) SetLocalPeerID(peerID string) {
if t == nil {
return
}
t.mu.Lock()
t.LocalPeerID = strings.TrimSpace(peerID)
t.mu.Unlock()
}
func quicTLSConfigForTarget(target FabricTransportTarget) *tls.Config {
@@ -186,9 +237,12 @@ func (t *QUICFabricTransport) connectConn(ctx context.Context, target FabricTran
conn, err := quic.DialAddr(ctx, target.Endpoint, tlsConfig, nil)
return conn, "", true, err
}
if conn, key, ok := t.reverseConnForTarget(target); ok {
return conn, key, false, nil
}
key := quicFabricConnKey(target)
if key == "" {
conn, err := quic.DialAddr(ctx, target.Endpoint, tlsConfig, t.Config)
conn, err := t.dialAddr(ctx, target.Endpoint, tlsConfig)
return conn, "", true, err
}
t.mu.Lock()
@@ -207,7 +261,7 @@ func (t *QUICFabricTransport) connectConn(ctx context.Context, target FabricTran
}
t.mu.Unlock()
conn, err := quic.DialAddr(ctx, target.Endpoint, tlsConfig, t.Config)
conn, err := t.dialAddr(ctx, target.Endpoint, tlsConfig)
if err != nil {
t.mu.Lock()
t.stats.OpenFailures++
@@ -235,16 +289,339 @@ func (t *QUICFabricTransport) connectConn(ctx context.Context, target FabricTran
t.conns[key] = &quicFabricConnEntry{conn: conn, lastUsed: time.Now()}
t.stats.Opens++
t.mu.Unlock()
go t.acceptInboundStreams(context.Background(), conn)
go t.sendReverseHello(context.Background(), conn)
return conn, key, false, nil
}
func (t *QUICFabricTransport) dialAddr(ctx context.Context, endpoint string, tlsConfig *tls.Config) (*quic.Conn, error) {
if t != nil && t.DialAddr != nil {
return t.DialAddr(ctx, endpoint, tlsConfig, t.Config)
}
return quic.DialAddr(ctx, endpoint, tlsConfig, t.Config)
}
func DialQUICAddrWithPacketConn(ctx context.Context, endpoint string, packetConn net.PacketConn, tlsConfig *tls.Config, config *quic.Config) (*quic.Conn, error) {
if packetConn == nil {
return nil, fmt.Errorf("quic packet connection is required")
}
addr, err := net.ResolveUDPAddr("udp", strings.TrimPrefix(strings.TrimSpace(endpoint), "quic://"))
if err != nil {
_ = packetConn.Close()
return nil, err
}
transport := &quic.Transport{Conn: packetConn}
conn, err := transport.Dial(ctx, addr, tlsConfig, config)
if err != nil {
_ = transport.Close()
return nil, err
}
go func() {
<-conn.Context().Done()
_ = transport.Close()
}()
return conn, nil
}
func (t *QUICFabricTransport) sendReverseHello(ctx context.Context, conn *quic.Conn) {
if t == nil || conn == nil {
return
}
localPeerID := t.localPeerID()
if localPeerID == "" {
t.mu.Lock()
t.stats.ReverseHelloFailed++
t.mu.Unlock()
return
}
helloCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
defer cancel()
stream, err := conn.OpenStreamSync(helloCtx)
if err != nil {
t.mu.Lock()
t.stats.ReverseHelloFailed++
t.mu.Unlock()
return
}
defer func() { _ = stream.Close() }()
if err := fabricproto.WriteFrame(stream, fabricproto.Frame{
Type: fabricproto.FramePing,
Sequence: 1,
Payload: []byte(fabricQUICReverseHelloPrefix + localPeerID),
}); err != nil {
t.mu.Lock()
t.stats.ReverseHelloFailed++
t.mu.Unlock()
return
}
t.mu.Lock()
t.stats.ReverseHelloSent++
t.mu.Unlock()
_, _ = fabricproto.ReadFrame(stream, fabricproto.DefaultMaxPayload)
}
func (t *QUICFabricTransport) acceptInboundStreams(ctx context.Context, conn *quic.Conn) {
if t == nil || conn == nil {
return
}
for {
stream, err := conn.AcceptStream(ctx)
if err != nil {
return
}
go t.handleInboundStream(ctx, conn, stream)
}
}
func (t *QUICFabricTransport) handleInboundStream(ctx context.Context, conn *quic.Conn, stream *quic.Stream) {
session := fabricproto.NewSession(fabricproto.SessionConfig{})
defer func() { _ = stream.Close() }()
t.logFabricSession(FabricSessionEventLogEntry{
Event: "fabric_session_quic_reverse_stream_opened",
AcceptedBy: "quic_reverse",
RemoteAddr: conn.RemoteAddr().String(),
})
defer t.logFabricSession(FabricSessionEventLogEntry{
Event: "fabric_session_quic_reverse_stream_closed",
AcceptedBy: "quic_reverse",
RemoteAddr: conn.RemoteAddr().String(),
})
for {
select {
case <-ctx.Done():
_ = stream.Close()
return
default:
}
frame, err := fabricproto.ReadFrame(stream, fabricproto.DefaultMaxPayload)
if err != nil {
return
}
t.registerReverseHelloFrame(conn, frame)
if t.handleInboundProductionForwardFrame(ctx, stream, frame) {
continue
}
if t.handleInboundWebIngressForwardFrame(ctx, stream, frame) {
continue
}
if t.handleInboundFabricControlForwardFrame(ctx, stream, frame) {
continue
}
if t.handleInboundSyntheticForwardFrame(ctx, stream, frame) {
continue
}
event, responses, err := session.HandleFrame(frame)
if err != nil {
_ = stream.Close()
return
}
if event.Type != fabricproto.SessionEventNone {
t.logFabricSession(FabricSessionEventLogEntry{
Event: "fabric_session_reverse_event",
SessionEvent: event.Type,
StreamID: event.StreamID,
Sequence: event.Sequence,
TrafficClass: event.TrafficClass,
AcceptedBy: "quic_reverse",
RemoteAddr: conn.RemoteAddr().String(),
})
}
for _, response := range responses {
if err := fabricproto.WriteFrame(stream, response); err != nil {
return
}
}
}
}
func (t *QUICFabricTransport) registerReverseHelloFrame(conn *quic.Conn, frame fabricproto.Frame) {
if t == nil || conn == nil || frame.Type != fabricproto.FramePing {
return
}
payload := string(frame.Payload)
if !strings.HasPrefix(payload, fabricQUICReverseHelloPrefix) {
return
}
peerID := strings.TrimPrefix(payload, fabricQUICReverseHelloPrefix)
t.RegisterReverseConn(peerID, conn)
t.logFabricSession(FabricSessionEventLogEntry{
Event: "fabric_session_quic_reverse_registered",
AcceptedBy: "quic_reverse_hello",
RemoteAddr: conn.RemoteAddr().String(),
PeerID: peerID,
})
}
func (t *QUICFabricTransport) handleInboundProductionForwardFrame(ctx context.Context, stream *quic.Stream, frame fabricproto.Frame) bool {
if frame.Type != fabricproto.FrameData || frame.StreamID != ProductionForwardQUICStreamID {
return false
}
response := quicProductionForwardResponse{}
productionHandler, _, _, _, _ := t.inboundHandlers()
if productionHandler == nil {
response.Error = ErrForwardRuntimeUnavailable.Error()
} else {
var envelope ProductionEnvelope
if err := json.Unmarshal(frame.Payload, &envelope); err != nil {
response.Error = "invalid production mesh envelope"
} else if result, err := productionHandler(ctx, envelope); err != nil {
response.Error = err.Error()
} else {
response.Result = result
}
}
payload, err := json.Marshal(response)
if err == nil {
_ = fabricproto.WriteFrame(stream, fabricproto.Frame{Type: fabricproto.FrameData, TrafficClass: fabricproto.TrafficClassReliable, StreamID: ProductionForwardQUICStreamID, Sequence: frame.Sequence, Payload: payload})
}
return true
}
func (t *QUICFabricTransport) handleInboundWebIngressForwardFrame(ctx context.Context, stream *quic.Stream, frame fabricproto.Frame) bool {
if frame.Type != fabricproto.FrameData || frame.StreamID != WebIngressForwardQUICStreamID {
return false
}
response := quicWebIngressForwardResponse{}
_, webIngressHandler, _, _, _ := t.inboundHandlers()
if webIngressHandler == nil {
response.Error = ErrForwardRuntimeUnavailable.Error()
} else if payload, err := webIngressHandler(ctx, append([]byte(nil), frame.Payload...)); err != nil {
response.Error = err.Error()
} else {
response.Payload = append(json.RawMessage(nil), payload...)
}
payload, err := json.Marshal(response)
if err == nil {
_ = fabricproto.WriteFrame(stream, fabricproto.Frame{Type: fabricproto.FrameData, TrafficClass: fabricproto.TrafficClassReliable, StreamID: WebIngressForwardQUICStreamID, Sequence: frame.Sequence, Payload: payload})
}
return true
}
func (t *QUICFabricTransport) handleInboundFabricControlForwardFrame(ctx context.Context, stream *quic.Stream, frame fabricproto.Frame) bool {
if frame.Type != fabricproto.FrameData || frame.StreamID != FabricControlForwardQUICStreamID {
return false
}
response := quicFabricControlForwardResponse{}
_, _, fabricControlHandler, _, _ := t.inboundHandlers()
if fabricControlHandler == nil {
response.Error = ErrForwardRuntimeUnavailable.Error()
} else if payload, err := fabricControlHandler(ctx, append([]byte(nil), frame.Payload...)); err != nil {
response.Error = err.Error()
} else {
response.Payload = append(json.RawMessage(nil), payload...)
}
payload, err := json.Marshal(response)
if err == nil {
_ = fabricproto.WriteFrame(stream, fabricproto.Frame{Type: fabricproto.FrameData, TrafficClass: fabricproto.TrafficClassReliable, StreamID: FabricControlForwardQUICStreamID, Sequence: frame.Sequence, Payload: payload})
}
return true
}
func (t *QUICFabricTransport) handleInboundSyntheticForwardFrame(ctx context.Context, stream *quic.Stream, frame fabricproto.Frame) bool {
if frame.Type != fabricproto.FrameData || frame.StreamID != SyntheticForwardQUICStreamID {
return false
}
response := quicSyntheticForwardResponse{}
_, _, _, syntheticHandler, _ := t.inboundHandlers()
if syntheticHandler == nil {
response.Error = ErrMeshRuntimeDisabled.Error()
} else {
var envelope SyntheticEnvelope
if err := json.Unmarshal(frame.Payload, &envelope); err != nil {
response.Error = "invalid synthetic mesh envelope"
} else if ack, err := syntheticHandler(ctx, envelope); err != nil {
response.Error = err.Error()
} else {
response.Envelope = ack
}
}
payload, err := json.Marshal(response)
if err == nil {
_ = fabricproto.WriteFrame(stream, fabricproto.Frame{Type: fabricproto.FrameData, TrafficClass: fabricproto.TrafficClassReliable, StreamID: SyntheticForwardQUICStreamID, Sequence: frame.Sequence, Payload: payload})
}
return true
}
func (t *QUICFabricTransport) inboundHandlers() (func(context.Context, ProductionEnvelope) (ProductionForwardResult, error), func(context.Context, []byte) ([]byte, error), func(context.Context, []byte) ([]byte, error), func(context.Context, SyntheticEnvelope) (SyntheticEnvelope, error), FabricSessionEventLogger) {
if t == nil {
return nil, nil, nil, nil, nil
}
t.mu.Lock()
defer t.mu.Unlock()
return t.inboundProductionHandler, t.inboundWebIngressHandler, t.inboundFabricControlHandler, t.inboundSyntheticHandler, t.logger
}
func (t *QUICFabricTransport) localPeerID() string {
if t == nil {
return ""
}
t.mu.Lock()
defer t.mu.Unlock()
return strings.TrimSpace(t.LocalPeerID)
}
func (t *QUICFabricTransport) logFabricSession(entry FabricSessionEventLogEntry) {
_, _, _, _, logger := t.inboundHandlers()
if logger != nil {
logger(entry)
}
}
func (t *QUICFabricTransport) RegisterReverseConn(peerID string, conn *quic.Conn) {
if t == nil || conn == nil {
return
}
peerID = strings.TrimSpace(peerID)
if peerID == "" {
return
}
t.mu.Lock()
defer t.mu.Unlock()
if t.reverseConns == nil {
t.reverseConns = map[string]*quicFabricConnEntry{}
}
if existing := t.reverseConns[peerID]; existing != nil && existing.conn != nil && existing.conn != conn {
select {
case <-existing.conn.Context().Done():
default:
_ = existing.conn.CloseWithError(0, "reverse connection replaced")
}
}
t.reverseConns[peerID] = &quicFabricConnEntry{conn: conn, lastUsed: time.Now()}
t.stats.ReverseRegisters++
}
func (t *QUICFabricTransport) reverseConnForTarget(target FabricTransportTarget) (*quic.Conn, string, bool) {
peerID := strings.TrimSpace(target.PeerID)
if t == nil || peerID == "" || !fabricTransportPrefersReverseConn(target.Transport) {
return nil, "", false
}
t.mu.Lock()
defer t.mu.Unlock()
t.pruneIdleLocked(time.Now())
entry := t.reverseConns[peerID]
if entry == nil || entry.conn == nil {
return nil, "", false
}
select {
case <-entry.conn.Context().Done():
delete(t.reverseConns, peerID)
t.stats.ClosedEvicted++
return nil, "", false
default:
entry.lastUsed = time.Now()
t.stats.ReverseReuses++
return entry.conn, quicFabricReverseConnKey(peerID), true
}
}
func (t *QUICFabricTransport) reserveStream(key string, conn *quic.Conn) error {
if t == nil || key == "" {
return nil
}
t.mu.Lock()
defer t.mu.Unlock()
entry := t.conns[key]
entry := t.connEntryLocked(key)
if entry == nil || entry.conn != conn {
return fmt.Errorf("quic fabric connection is not cached")
}
@@ -267,16 +644,26 @@ func (t *QUICFabricTransport) releaseStream(key string) {
return
}
t.mu.Lock()
if entry := t.conns[key]; entry != nil {
if entry := t.connEntryLocked(key); entry != nil {
if entry.activeStreams > 0 {
entry.activeStreams--
}
entry.lastUsed = time.Now()
t.stats.StreamCloses++
}
t.stats.StreamCloses++
t.mu.Unlock()
}
func (t *QUICFabricTransport) connEntryLocked(key string) *quicFabricConnEntry {
if t == nil || key == "" {
return nil
}
if strings.HasPrefix(key, "reverse\x00") {
return t.reverseConns[strings.TrimPrefix(key, "reverse\x00")]
}
return t.conns[key]
}
func (t *QUICFabricTransport) evictConn(target FabricTransportTarget, conn *quic.Conn) {
if t == nil || conn == nil {
return
@@ -315,6 +702,20 @@ func (t *QUICFabricTransport) pruneIdleLocked(now time.Time) {
t.stats.IdleEvicted++
}
}
for peerID, entry := range t.reverseConns {
if entry == nil || entry.conn == nil {
delete(t.reverseConns, peerID)
continue
}
if !entry.lastUsed.IsZero() && now.Sub(entry.lastUsed) > ttl {
if entry.activeStreams > 0 {
continue
}
_ = entry.conn.CloseWithError(0, "idle reverse")
delete(t.reverseConns, peerID)
t.stats.IdleEvicted++
}
}
}
func quicFabricConnKey(target FabricTransportTarget) string {
@@ -340,6 +741,23 @@ func parseQUICFabricConnKey(key string) (peerID string, endpoint string, certSHA
return peerID, endpoint, certSHA256
}
func quicFabricReverseConnKey(peerID string) string {
peerID = strings.TrimSpace(peerID)
if peerID == "" {
return ""
}
return "reverse\x00" + peerID
}
func fabricTransportPrefersReverseConn(transport string) bool {
switch strings.ToLower(strings.TrimSpace(transport)) {
case "reverse_quic", "relay_quic":
return true
default:
return false
}
}
func (t *QUICFabricTransport) Close() error {
if t == nil {
return nil
@@ -348,12 +766,19 @@ func (t *QUICFabricTransport) Close() error {
t.stats.CloseAllCalls++
conns := t.conns
t.conns = map[string]*quicFabricConnEntry{}
reverseConns := t.reverseConns
t.reverseConns = map[string]*quicFabricConnEntry{}
t.mu.Unlock()
for _, entry := range conns {
if entry != nil && entry.conn != nil {
_ = entry.conn.CloseWithError(0, "closed")
}
}
for _, entry := range reverseConns {
if entry != nil && entry.conn != nil {
_ = entry.conn.CloseWithError(0, "closed")
}
}
return nil
}
@@ -370,6 +795,7 @@ func (t *QUICFabricTransport) Snapshot() QUICFabricTransportSnapshot {
}
snapshot := QUICFabricTransportSnapshot{
SchemaVersion: "rap.quic_fabric_transport.v1",
LocalPeerID: strings.TrimSpace(t.LocalPeerID),
MaxStreamsPerConn: limit,
Stats: t.stats,
}
@@ -391,6 +817,40 @@ func (t *QUICFabricTransport) Snapshot() QUICFabricTransportSnapshot {
PeerID: peerID,
Endpoint: endpoint,
CertSHA256: certSHA256,
Direction: "outbound",
ActiveStreams: entry.activeStreams,
MaxStreams: limit,
Saturated: entry.activeStreams >= limit,
}
if !entry.lastUsed.IsZero() {
connSnapshot.LastUsedUnixSec = entry.lastUsed.UTC().Unix()
}
if limit > 0 {
connSnapshot.CapacityPressurePercent = (entry.activeStreams * 100) / limit
}
snapshot.Connections = append(snapshot.Connections, connSnapshot)
if entry.activeStreams >= limit {
snapshot.SaturatedConnections++
}
}
}
for peerID, entry := range t.reverseConns {
if entry == nil || entry.conn == nil {
delete(t.reverseConns, peerID)
continue
}
select {
case <-entry.conn.Context().Done():
delete(t.reverseConns, peerID)
t.stats.ClosedEvicted++
snapshot.Stats.ClosedEvicted++
default:
snapshot.ActiveCount++
snapshot.ActiveStreams += entry.activeStreams
connSnapshot := QUICFabricConnSnapshot{
PeerID: peerID,
Endpoint: entry.conn.RemoteAddr().String(),
Direction: "reverse",
ActiveStreams: entry.activeStreams,
MaxStreams: limit,
Saturated: entry.activeStreams >= limit,
@@ -462,6 +922,7 @@ func (s *quicFabricSession) Close() error {
s.closeOnce.Do(func() {
close(s.done)
if s.stream != nil {
s.stream.CancelRead(0)
err = s.stream.Close()
}
if s.transport != nil {