Refactor RDP proxy handling and update related tests
This commit is contained in:
@@ -2,42 +2,369 @@ package mesh
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/example/remote-access-platform/agents/rap-node-agent/internal/fabricproto"
|
||||
)
|
||||
|
||||
type ProductionForwardTransport interface {
|
||||
SendProduction(ctx context.Context, nextNodeID string, envelope ProductionEnvelope) (ProductionForwardResult, error)
|
||||
}
|
||||
|
||||
type HTTPProductionForwardTransport struct {
|
||||
PeerURLs map[string]string
|
||||
HTTPClient *http.Client
|
||||
type QUICProductionForwardTransport struct {
|
||||
Targets map[string]FabricTransportTarget
|
||||
RouteSets map[string]FabricRouteSet
|
||||
Transport FabricTransport
|
||||
Router FabricChannelRouter
|
||||
Timeout time.Duration
|
||||
Pressure *FabricRoutePressureTracker
|
||||
Health *FabricRouteHealthTracker
|
||||
sequence atomic.Uint64
|
||||
}
|
||||
|
||||
func NewHTTPProductionForwardTransport(peerURLs map[string]string) *HTTPProductionForwardTransport {
|
||||
normalized := make(map[string]string, len(peerURLs))
|
||||
for nodeID, baseURL := range peerURLs {
|
||||
type QUICProductionForwardTransportSnapshot struct {
|
||||
RoutePressure FabricRoutePressureSnapshot `json:"route_pressure"`
|
||||
RouteHealth FabricRouteHealthSnapshot `json:"route_health,omitempty"`
|
||||
}
|
||||
|
||||
func NewQUICProductionForwardTransport(targets map[string]FabricTransportTarget, transport *QUICFabricTransport) *QUICProductionForwardTransport {
|
||||
routeSets := make(map[string]FabricRouteSet, len(targets))
|
||||
for nodeID, target := range targets {
|
||||
nodeID = strings.TrimSpace(nodeID)
|
||||
baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
||||
if nodeID != "" && baseURL != "" {
|
||||
normalized[nodeID] = baseURL
|
||||
target.Endpoint = strings.TrimRight(strings.TrimSpace(target.Endpoint), "/")
|
||||
target.Transport = strings.TrimSpace(target.Transport)
|
||||
if nodeID != "" && target.Endpoint != "" {
|
||||
target.PeerID = firstNonEmpty(strings.TrimSpace(target.PeerID), nodeID)
|
||||
routeSets[nodeID] = FabricRouteSetForTransportTargets("", "", nodeID, []FabricTransportTarget{target})
|
||||
}
|
||||
}
|
||||
return &HTTPProductionForwardTransport{PeerURLs: normalized}
|
||||
if transport == nil {
|
||||
transport = NewQUICFabricTransport(nil)
|
||||
}
|
||||
return NewQUICProductionForwardTransportFromRouteSets(routeSets, transport)
|
||||
}
|
||||
|
||||
func (t *HTTPProductionForwardTransport) SendProduction(ctx context.Context, nextNodeID string, envelope ProductionEnvelope) (ProductionForwardResult, error) {
|
||||
if t == nil {
|
||||
return ProductionForwardResult{}, ErrForwardPeerUnavailable
|
||||
func NewQUICProductionForwardTransportFromRouteSets(routeSets map[string]FabricRouteSet, transport FabricTransport) *QUICProductionForwardTransport {
|
||||
normalizedRouteSets := make(map[string]FabricRouteSet, len(routeSets))
|
||||
targets := make(map[string]FabricTransportTarget, len(routeSets))
|
||||
for nodeID, routeSet := range routeSets {
|
||||
nodeID = strings.TrimSpace(nodeID)
|
||||
if nodeID == "" {
|
||||
continue
|
||||
}
|
||||
normalizedRouteSets[nodeID] = routeSet
|
||||
if target, err := FabricTransportTargetForRoute(routeSet.Primary); err == nil {
|
||||
targets[nodeID] = target
|
||||
}
|
||||
}
|
||||
baseURL := strings.TrimRight(strings.TrimSpace(t.PeerURLs[nextNodeID]), "/")
|
||||
if baseURL == "" {
|
||||
return ProductionForwardResult{}, ErrForwardPeerUnavailable
|
||||
if transport == nil {
|
||||
transport = NewQUICFabricTransport(nil)
|
||||
}
|
||||
client := NewClient(baseURL)
|
||||
if t.HTTPClient != nil {
|
||||
client.HTTPClient = t.HTTPClient
|
||||
return &QUICProductionForwardTransport{
|
||||
Targets: targets,
|
||||
RouteSets: normalizedRouteSets,
|
||||
Transport: transport,
|
||||
Router: NewFabricChannelRouter(FabricChannelRouterConfig{
|
||||
MaxAckLatencyMs: 2000,
|
||||
MinRerouteInterval: 50 * time.Millisecond,
|
||||
}),
|
||||
Timeout: 30 * time.Second,
|
||||
Pressure: NewFabricRoutePressureTracker(),
|
||||
Health: NewFabricRouteHealthTracker(30 * time.Second),
|
||||
}
|
||||
return client.SendProduction(ctx, envelope)
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) SendProduction(ctx context.Context, nextNodeID string, envelope ProductionEnvelope) (ProductionForwardResult, error) {
|
||||
if t == nil || t.Transport == nil {
|
||||
return ProductionForwardResult{}, ErrForwardPeerUnavailable
|
||||
}
|
||||
nextNodeID = strings.TrimSpace(nextNodeID)
|
||||
routeSet, ok := t.RouteSets[nextNodeID]
|
||||
if !ok {
|
||||
target, targetOK := t.Targets[nextNodeID]
|
||||
if !targetOK || strings.TrimSpace(target.Endpoint) == "" {
|
||||
return ProductionForwardResult{}, ErrForwardPeerUnavailable
|
||||
}
|
||||
routeSet = FabricRouteSetForTransportTargets(envelope.ClusterID, envelope.CurrentHopNodeID, nextNodeID, []FabricTransportTarget{target})
|
||||
}
|
||||
spec := FabricChannelSpec{
|
||||
ChannelID: firstNonEmpty(strings.TrimSpace(envelope.MessageID), fmt.Sprintf("production-%d", t.sequence.Add(1))),
|
||||
ClusterID: envelope.ClusterID,
|
||||
SourceNodeID: firstNonEmpty(productionRouteSetSourceNodeID(routeSet), envelope.CurrentHopNodeID),
|
||||
TargetKind: FabricChannelTargetNode,
|
||||
TargetID: nextNodeID,
|
||||
TrafficClass: FabricServiceChannelReliable,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
payload, err := json.Marshal(envelope)
|
||||
if err != nil {
|
||||
return ProductionForwardResult{}, err
|
||||
}
|
||||
result, err := t.sendProductionWithRouteSet(ctx, spec, routeSet, payload)
|
||||
if err != nil {
|
||||
return ProductionForwardResult{}, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func productionRouteSetSourceNodeID(routeSet FabricRouteSet) string {
|
||||
for _, route := range flattenFabricRouteSet(routeSet) {
|
||||
if sourceNodeID := strings.TrimSpace(route.SourceNodeID); sourceNodeID != "" {
|
||||
return sourceNodeID
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) sendProductionWithRouteSet(ctx context.Context, spec FabricChannelSpec, routeSet FabricRouteSet, payload []byte) (ProductionForwardResult, error) {
|
||||
router := t.Router
|
||||
if router.Config.MaxRoutePressure == 0 {
|
||||
router = NewFabricChannelRouter(FabricChannelRouterConfig{MaxAckLatencyMs: 2000, MinRerouteInterval: 50 * time.Millisecond})
|
||||
}
|
||||
routeSet = t.routeSetForScheduling(routeSet)
|
||||
channel, _, err := router.OpenChannel(spec, routeSet, time.Now().UTC())
|
||||
if err != nil {
|
||||
return ProductionForwardResult{}, err
|
||||
}
|
||||
timeout := t.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = 30 * time.Second
|
||||
}
|
||||
for {
|
||||
routeSet = t.routeSetForScheduling(routeSet)
|
||||
route, ok := findFabricRoute(routeSet, channel.RouteID)
|
||||
if !ok {
|
||||
return ProductionForwardResult{}, ErrFabricRouteNotFound
|
||||
}
|
||||
target, err := FabricTransportTargetForRoute(route)
|
||||
if err != nil {
|
||||
return ProductionForwardResult{}, err
|
||||
}
|
||||
target.PeerID = firstNonEmpty(strings.TrimSpace(target.PeerID), spec.TargetID)
|
||||
target.MaxPayload = fabricproto.DefaultMaxPayload
|
||||
releaseRoute := t.acquireProductionRoute(route.RouteID)
|
||||
session, err := t.Transport.Connect(ctx, target)
|
||||
if err != nil {
|
||||
releaseRoute()
|
||||
t.markProductionRouteFailure(route.RouteID, err)
|
||||
updated, event, rerouteErr := router.ObserveChannel(channel, routeSet, FabricChannelObservation{
|
||||
ChannelID: spec.ChannelID,
|
||||
RouteID: route.RouteID,
|
||||
Failed: true,
|
||||
Reason: "connect_failed",
|
||||
ObservedAt: time.Now().UTC(),
|
||||
}, time.Now().UTC())
|
||||
channel = updated
|
||||
if event.Type == FabricChannelRouteEventReroute {
|
||||
continue
|
||||
}
|
||||
if rerouteErr != nil {
|
||||
return ProductionForwardResult{}, rerouteErr
|
||||
}
|
||||
return ProductionForwardResult{}, err
|
||||
}
|
||||
response, ackMs, err := t.sendProductionOnSession(ctx, session, payload, timeout)
|
||||
_ = session.Close()
|
||||
releaseRoute()
|
||||
if err == nil {
|
||||
t.markProductionRouteSuccess(route.RouteID)
|
||||
_, _, _ = router.ObserveChannel(channel, routeSet, FabricChannelObservation{
|
||||
ChannelID: spec.ChannelID,
|
||||
RouteID: route.RouteID,
|
||||
AckLatencyMs: ackMs,
|
||||
BytesSent: uint64(len(payload)),
|
||||
FramesSent: 1,
|
||||
BytesRecv: uint64(len(response.Payload)),
|
||||
FramesRecv: 1,
|
||||
ObservedAt: time.Now().UTC(),
|
||||
}, time.Now().UTC())
|
||||
return decodeQUICProductionForwardResponse(response.Payload)
|
||||
}
|
||||
t.markProductionRouteFailure(route.RouteID, err)
|
||||
updated, event, rerouteErr := router.ObserveChannel(channel, routeSet, FabricChannelObservation{
|
||||
ChannelID: spec.ChannelID,
|
||||
RouteID: route.RouteID,
|
||||
Failed: true,
|
||||
Reason: "response_failed",
|
||||
ObservedAt: time.Now().UTC(),
|
||||
}, time.Now().UTC())
|
||||
channel = updated
|
||||
if event.Type == FabricChannelRouteEventReroute {
|
||||
continue
|
||||
}
|
||||
if rerouteErr != nil {
|
||||
return ProductionForwardResult{}, rerouteErr
|
||||
}
|
||||
return ProductionForwardResult{}, err
|
||||
}
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) routeSetWithActiveChannels(routeSet FabricRouteSet) FabricRouteSet {
|
||||
if t == nil || t.Pressure == nil {
|
||||
return routeSet
|
||||
}
|
||||
return t.Pressure.Apply(routeSet)
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) routeSetForScheduling(routeSet FabricRouteSet) FabricRouteSet {
|
||||
if t != nil && t.Health != nil {
|
||||
routeSet = t.Health.Apply(routeSet, time.Now().UTC())
|
||||
}
|
||||
return t.routeSetWithActiveChannels(routeSet)
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) acquireProductionRoute(routeID string) func() {
|
||||
if t == nil || t.Pressure == nil {
|
||||
return func() {}
|
||||
}
|
||||
return t.Pressure.Acquire(routeID)
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) markProductionRouteFailure(routeID string, err error) {
|
||||
if t == nil || t.Health == nil || err == nil {
|
||||
return
|
||||
}
|
||||
t.Health.MarkFailure(routeID, err.Error(), time.Now().UTC())
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) markProductionRouteSuccess(routeID string) {
|
||||
if t == nil || t.Health == nil {
|
||||
return
|
||||
}
|
||||
t.Health.MarkSuccess(routeID)
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) Snapshot() QUICProductionForwardTransportSnapshot {
|
||||
if t == nil {
|
||||
return QUICProductionForwardTransportSnapshot{}
|
||||
}
|
||||
var pressure FabricRoutePressureSnapshot
|
||||
if t.Pressure != nil {
|
||||
pressure = t.Pressure.SnapshotPressure()
|
||||
}
|
||||
var health FabricRouteHealthSnapshot
|
||||
if t.Health != nil {
|
||||
health = t.Health.Snapshot(time.Now().UTC())
|
||||
}
|
||||
return QUICProductionForwardTransportSnapshot{RoutePressure: pressure, RouteHealth: health}
|
||||
}
|
||||
|
||||
func (t *QUICProductionForwardTransport) sendProductionOnSession(ctx context.Context, session FabricTransportSession, payload []byte, timeout time.Duration) (fabricproto.Frame, int64, error) {
|
||||
sequence := t.sequence.Add(1)
|
||||
if err := session.Send(ctx, fabricproto.Frame{
|
||||
Type: fabricproto.FrameData,
|
||||
TrafficClass: fabricproto.TrafficClassReliable,
|
||||
StreamID: ProductionForwardQUICStreamID,
|
||||
Sequence: sequence,
|
||||
Payload: payload,
|
||||
}); err != nil {
|
||||
return fabricproto.Frame{}, 0, err
|
||||
}
|
||||
waitCtx := ctx
|
||||
if timeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
waitCtx, cancel = context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
}
|
||||
started := time.Now()
|
||||
for {
|
||||
select {
|
||||
case <-waitCtx.Done():
|
||||
return fabricproto.Frame{}, 0, waitCtx.Err()
|
||||
case err, ok := <-session.Errors():
|
||||
if !ok {
|
||||
return fabricproto.Frame{}, 0, ErrForwardPeerUnavailable
|
||||
}
|
||||
if err != nil {
|
||||
return fabricproto.Frame{}, 0, err
|
||||
}
|
||||
case frame, ok := <-session.Frames():
|
||||
if !ok {
|
||||
return fabricproto.Frame{}, 0, ErrForwardPeerUnavailable
|
||||
}
|
||||
if frame.Type != fabricproto.FrameData || frame.StreamID != ProductionForwardQUICStreamID || frame.Sequence != sequence {
|
||||
continue
|
||||
}
|
||||
return frame, time.Since(started).Milliseconds(), nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func decodeQUICProductionForwardResponse(payload []byte) (ProductionForwardResult, error) {
|
||||
var response quicProductionForwardResponse
|
||||
if err := json.Unmarshal(payload, &response); err != nil {
|
||||
return ProductionForwardResult{}, err
|
||||
}
|
||||
if strings.TrimSpace(response.Error) != "" {
|
||||
return ProductionForwardResult{}, fmt.Errorf("%w: %s", ErrForwardPeerUnavailable, response.Error)
|
||||
}
|
||||
return response.Result, nil
|
||||
}
|
||||
|
||||
func FabricRouteSetForTransportTargets(clusterID string, sourceNodeID string, targetNodeID string, targets []FabricTransportTarget) FabricRouteSet {
|
||||
routeSet := FabricRouteSet{TargetKind: FabricChannelTargetNode, TargetID: strings.TrimSpace(targetNodeID)}
|
||||
routes := make([]FabricRoute, 0, len(targets))
|
||||
for index, target := range targets {
|
||||
target.Endpoint = strings.TrimRight(strings.TrimSpace(target.Endpoint), "/")
|
||||
if strings.TrimSpace(target.Endpoint) == "" {
|
||||
continue
|
||||
}
|
||||
peerID := firstNonEmpty(strings.TrimSpace(target.PeerID), strings.TrimSpace(targetNodeID))
|
||||
routeID := strings.TrimSpace(target.EndpointID)
|
||||
if routeID == "" {
|
||||
routeID = fmt.Sprintf("%s-quic-%d", peerID, index)
|
||||
}
|
||||
routes = append(routes, FabricRoute{
|
||||
RouteID: routeID,
|
||||
ClusterID: strings.TrimSpace(clusterID),
|
||||
SourceNodeID: strings.TrimSpace(sourceNodeID),
|
||||
DestinationNodeID: peerID,
|
||||
Hops: []FabricRouteHop{{
|
||||
NodeID: peerID,
|
||||
Mode: fabricRouteModeForTransportTarget(target),
|
||||
EndpointID: strings.TrimSpace(target.EndpointID),
|
||||
Address: target.Endpoint,
|
||||
PeerCertSHA256: strings.TrimSpace(target.PeerCertSHA256),
|
||||
}},
|
||||
BaseLatencyMs: routeLatencyForIndex(index),
|
||||
Capacity: 100,
|
||||
ActiveChannels: 0,
|
||||
Healthy: true,
|
||||
LastUpdatedAt: time.Now().UTC(),
|
||||
})
|
||||
}
|
||||
if len(routes) == 0 {
|
||||
return routeSet
|
||||
}
|
||||
routeSet.Primary = routes[0]
|
||||
if len(routes) > 1 {
|
||||
routeSet.WarmStandby = append(routeSet.WarmStandby, routes[1:]...)
|
||||
}
|
||||
return routeSet
|
||||
}
|
||||
|
||||
func fabricRouteModeForTransportTarget(target FabricTransportTarget) FabricRouteMode {
|
||||
switch strings.ToLower(strings.TrimSpace(target.Transport)) {
|
||||
case string(FabricRouteLAN):
|
||||
return FabricRouteLAN
|
||||
case string(FabricRouteReverse):
|
||||
return FabricRouteReverse
|
||||
case string(FabricRouteRelay):
|
||||
return FabricRouteRelay
|
||||
case string(FabricRouteICE):
|
||||
return FabricRouteICE
|
||||
default:
|
||||
return FabricRouteDirect
|
||||
}
|
||||
}
|
||||
|
||||
func routeLatencyForIndex(index int) int {
|
||||
if index <= 0 {
|
||||
return 10
|
||||
}
|
||||
return 10 + index
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user