rdp-proxy/agents/rap-node-agent/internal/mesh/fabric_quic_transport.go

package mesh

import (
	"context"
	"crypto/tls"
	"fmt"
	"sync"
	"time"

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

const fabricQUICNextProto = "rap-fabric-data-session-v1"

type QUICFabricTransport struct {
	Config *quic.Config
}

type quicFabricSession struct {
	conn       *quic.Conn
	stream     *quic.Stream
	inbound    chan fabricproto.Frame
	errors     chan error
	done       chan struct{}
	closeOnce  sync.Once
	writeMu    sync.Mutex
	maxPayload int
	timeout    time.Duration
}

func NewQUICFabricTransport(config *quic.Config) *QUICFabricTransport {
	return &QUICFabricTransport{Config: config}
}

func (t *QUICFabricTransport) Connect(ctx context.Context, target FabricTransportTarget) (FabricTransportSession, error) {
	if target.Endpoint == "" {
		return nil, fmt.Errorf("quic fabric endpoint is required")
	}
	tlsConfig := target.TLSConfig
	if tlsConfig == nil {
		tlsConfig = &tls.Config{NextProtos: []string{fabricQUICNextProto}}
	} else {
		tlsConfig = tlsConfig.Clone()
		if len(tlsConfig.NextProtos) == 0 {
			tlsConfig.NextProtos = []string{fabricQUICNextProto}
		}
	}
	conn, err := quic.DialAddr(ctx, target.Endpoint, tlsConfig, t.Config)
	if err != nil {
		return nil, err
	}
	stream, err := conn.OpenStreamSync(ctx)
	if err != nil {
		_ = conn.CloseWithError(1, "open stream failed")
		return nil, err
	}
	maxPayload := target.MaxPayload
	if maxPayload <= 0 {
		maxPayload = fabricproto.DefaultMaxPayload
	}
	inboundBuffer := target.InboundBuffer
	if inboundBuffer <= 0 {
		inboundBuffer = 64
	}
	errorBuffer := target.ErrorBuffer
	if errorBuffer <= 0 {
		errorBuffer = 8
	}
	session := &quicFabricSession{
		conn:       conn,
		stream:     stream,
		inbound:    make(chan fabricproto.Frame, inboundBuffer),
		errors:     make(chan error, errorBuffer),
		done:       make(chan struct{}),
		maxPayload: maxPayload,
		timeout:    target.Timeout,
	}
	go session.readLoop(ctx)
	return session, nil
}

func (t *QUICFabricTransport) Close() error {
	return nil
}

func (s *quicFabricSession) Send(ctx context.Context, frame fabricproto.Frame) error {
	if s == nil || s.stream == nil {
		return fmt.Errorf("quic fabric session is closed")
	}
	select {
	case <-s.done:
		return fmt.Errorf("quic fabric session is closed")
	default:
	}
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	s.applyWriteDeadline(ctx)
	return fabricproto.WriteFrame(s.stream, frame)
}

func (s *quicFabricSession) Frames() <-chan fabricproto.Frame {
	if s == nil {
		return nil
	}
	return s.inbound
}

func (s *quicFabricSession) Errors() <-chan error {
	if s == nil {
		return nil
	}
	return s.errors
}

func (s *quicFabricSession) Close() error {
	if s == nil {
		return nil
	}
	var err error
	s.closeOnce.Do(func() {
		close(s.done)
		if s.stream != nil {
			err = s.stream.Close()
		}
		if s.conn != nil {
			_ = s.conn.CloseWithError(0, "closed")
		}
	})
	return err
}

func (s *quicFabricSession) Closed() bool {
	if s == nil {
		return true
	}
	select {
	case <-s.done:
		return true
	default:
		return false
	}
}

func (s *quicFabricSession) readLoop(ctx context.Context) {
	defer s.Close()
	for {
		s.applyReadDeadline(ctx)
		frame, err := fabricproto.ReadFrame(s.stream, s.maxPayload)
		if err != nil {
			s.reportError(err)
			return
		}
		select {
		case <-ctx.Done():
			s.reportError(ctx.Err())
			return
		case <-s.done:
			return
		case s.inbound <- frame:
		}
	}
}

func (s *quicFabricSession) reportError(err error) {
	if err == nil {
		return
	}
	select {
	case s.errors <- err:
	default:
	}
}

func (s *quicFabricSession) applyReadDeadline(ctx context.Context) {
	if deadline, ok := ctx.Deadline(); ok {
		_ = s.stream.SetReadDeadline(deadline)
	} else if s.timeout > 0 {
		_ = s.stream.SetReadDeadline(time.Now().Add(s.timeout))
	}
}

func (s *quicFabricSession) applyWriteDeadline(ctx context.Context) {
	if deadline, ok := ctx.Deadline(); ok {
		_ = s.stream.SetWriteDeadline(deadline)
	} else if s.timeout > 0 {
		_ = s.stream.SetWriteDeadline(time.Now().Add(s.timeout))
	}
}