feat(gateway): repair native Apollo provider path

This commit is contained in:
sechmachine
2026-07-29 17:24:34 +07:00
parent 67eeb9e898
commit 236c7c2973
23 changed files with 5088 additions and 345 deletions
+393 -67
View File
@@ -19,9 +19,10 @@ import (
)
const (
defaultHelloLimit = 16 * 1024
defaultControlLimit = 128 * 1024
applicationError = quic.ApplicationErrorCode(0x100)
defaultHelloLimit = 16 * 1024
defaultControlLimit = 128 * 1024
clientControlBacklog = 64
applicationError = quic.ApplicationErrorCode(0x100)
)
var (
@@ -53,25 +54,31 @@ type ProviderStateReporter interface {
ReportProviderState(context.Context, protocol.ProviderState) error
}
type ClipboardAuditReporter interface {
ReportClipboardAudit(context.Context, protocol.GatewayClipboardAudit) error
}
type ServerConfig struct {
ListenAddress string
TLSConfig *tls.Config
QUICConfig *quic.Config
GatewayID string
Capabilities protocol.CapabilityProfile
ProviderCapabilities protocol.CapabilityProfile
Admission Admission
ProviderStateReporter ProviderStateReporter
Provider Provider
ProviderProfile string
ProviderIdentity string
PacerKbps int64
ListenAddress string
TLSConfig *tls.Config
QUICConfig *quic.Config
GatewayID string
Capabilities protocol.CapabilityProfile
ProviderCapabilities protocol.CapabilityProfile
Admission Admission
ProviderStateReporter ProviderStateReporter
ClipboardAuditReporter ClipboardAuditReporter
Provider Provider
ProviderProfile string
ProviderIdentity string
PacerKbps int64
}
type Server struct {
listener *quic.Listener
config ServerConfig
metrics *Metrics
pacer *fairPacer
mu sync.Mutex
sessions map[*gatewaySession]struct{}
draining atomic.Bool
@@ -115,7 +122,7 @@ func NewServer(config ServerConfig) (*Server, error) {
if err != nil {
return nil, err
}
return &Server{listener: listener, config: config, metrics: &Metrics{}, sessions: make(map[*gatewaySession]struct{})}, nil
return &Server{listener: listener, config: config, metrics: &Metrics{}, pacer: newFairPacer(config.PacerKbps), sessions: make(map[*gatewaySession]struct{})}, nil
}
func validateServerTLS(config *tls.Config) error {
@@ -130,7 +137,9 @@ func (s *Server) Metrics() MetricsSnapshot { return s.metrics.Snapshot() }
func (s *Server) Draining() bool { return s.draining.Load() }
func (s *Server) BeginDrain() {
s.draining.Store(true)
if s.draining.CompareAndSwap(false, true) {
s.metrics.DrainTransitions.Add(1)
}
}
func (s *Server) Serve(ctx context.Context) error {
@@ -228,6 +237,17 @@ func (s *Server) handleConnection(parent context.Context, connection *quic.Conn)
_ = writeStableError(stream, "no_capability_overlap", err, false)
return
}
clipboard, err := newClipboardGate(work.ClipboardPolicy, time.Now)
if err != nil {
_ = s.config.Admission.Release(context.Background(), authority)
_ = writeStableError(stream, "provider_work_unavailable", ErrAdmissionRejected, false)
return
}
if (work.ClipboardPolicy.ClientToProviderEnabled || work.ClipboardPolicy.ProviderToClientEnabled) && s.config.ClipboardAuditReporter == nil {
_ = s.config.Admission.Release(context.Background(), authority)
_ = writeStableError(stream, "clipboard_audit_unavailable", ErrAdmissionRejected, true)
return
}
if err := s.reportProviderState(ctx, protocol.ProviderState{Version: "1", SessionID: request.SessionID, State: ProviderStateStarting, CleanupPending: false, Channels: []string{"video", "audio", "input", "feedback"}}); err != nil {
_ = s.config.Admission.Release(context.Background(), authority)
_ = writeStableError(stream, "provider_state_unavailable", err, true)
@@ -256,9 +276,13 @@ func (s *Server) handleConnection(parent context.Context, connection *quic.Conn)
_ = s.config.Admission.Release(context.Background(), authority)
return
}
session := newGatewaySession(s, connection, stream, request, authority, providerSession)
session := newGatewaySession(s, connection, stream, request, authority, providerSession, clipboard)
s.addSession(session)
s.metrics.ActiveSessions.Add(1)
s.metrics.AdmittedSessions.Add(1)
if authority.ReconnectSequence > 0 {
s.metrics.Reconnects.Add(1)
}
defer func() {
s.removeSession(session)
s.metrics.ActiveSessions.Add(-1)
@@ -267,12 +291,13 @@ func (s *Server) handleConnection(parent context.Context, connection *quic.Conn)
}
func (s *Server) reportProviderState(ctx context.Context, state protocol.ProviderState) error {
if s.config.ProviderStateReporter == nil {
return nil
}
if err := state.Validate(); err != nil {
return err
}
s.metrics.observeProviderState(state.State)
if s.config.ProviderStateReporter == nil {
return nil
}
return s.config.ProviderStateReporter.ReportProviderState(ctx, state)
}
@@ -315,25 +340,27 @@ func (s *Server) removeSession(session *gatewaySession) {
}
type gatewaySession struct {
server *Server
connection *quic.Conn
control *quic.Stream
request protocol.TunnelAdmissionRequest
authority protocol.SessionAuthority
provider ProviderSession
pacer *Pacer
ctx context.Context
cancel context.CancelFunc
cleanupOnce sync.Once
inputMu sync.Mutex
pressed map[string]struct{}
sequence atomic.Uint32
result chan error
server *Server
connection *quic.Conn
control *quic.Stream
request protocol.TunnelAdmissionRequest
authority protocol.SessionAuthority
provider ProviderSession
clipboard *clipboardGate
ctx context.Context
cancel context.CancelFunc
cleanupOnce sync.Once
inputMu sync.Mutex
controlWriteMu sync.Mutex
pressed map[string]struct{}
sequence atomic.Uint32
mediaDrops uint64
result chan error
}
func newGatewaySession(server *Server, connection *quic.Conn, control *quic.Stream, request protocol.TunnelAdmissionRequest, authority protocol.SessionAuthority, provider ProviderSession) *gatewaySession {
func newGatewaySession(server *Server, connection *quic.Conn, control *quic.Stream, request protocol.TunnelAdmissionRequest, authority protocol.SessionAuthority, provider ProviderSession, clipboard *clipboardGate) *gatewaySession {
ctx, cancel := context.WithCancel(context.Background())
return &gatewaySession{server: server, connection: connection, control: control, request: request, authority: authority, provider: provider, pacer: NewPacer(server.config.PacerKbps), ctx: ctx, cancel: cancel, pressed: make(map[string]struct{}), result: make(chan error, 3)}
return &gatewaySession{server: server, connection: connection, control: control, request: request, authority: authority, provider: provider, clipboard: clipboard, ctx: ctx, cancel: cancel, pressed: make(map[string]struct{}), result: make(chan error, 3)}
}
func (s *gatewaySession) run() {
@@ -350,6 +377,11 @@ func (s *gatewaySession) run() {
go s.controlLoop()
go s.datagramLoop()
go s.mediaLoop()
go s.providerEventLoop()
go s.providerTelemetryLoop()
if s.clipboard != nil && s.clipboard.policy.ProviderToClientEnabled {
go s.clipboardLoop()
}
select {
case <-timer.C:
s.server.metrics.InputRejected.Add(1)
@@ -359,6 +391,96 @@ func (s *gatewaySession) run() {
s.cancel()
}
func (s *gatewaySession) providerEventLoop() {
events := s.provider.Events()
for events != nil {
select {
case <-s.ctx.Done():
return
case event, ok := <-events:
if !ok {
return
}
payload, err := EncodeProviderEvent(event)
if err == nil {
err = s.sendControl(s.sequence.Add(1), payload)
}
if err != nil {
s.result <- err
return
}
}
}
}
func (s *gatewaySession) clipboardLoop() {
ticker := time.NewTicker(500 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-s.ctx.Done():
return
case <-ticker.C:
text, err := s.provider.ReadClipboard(s.ctx)
if err != nil {
if auditErr := s.reportClipboardAudit("provider_to_client", "rejected", clipboardAuditTextBytes(text), clipboardAuditReason(err)); auditErr != nil {
s.result <- auditErr
return
}
s.result <- err
return
}
value, suppress, err := s.clipboard.fromProvider(text)
if err != nil {
if auditErr := s.reportClipboardAudit("provider_to_client", "rejected", clipboardAuditTextBytes(text), clipboardAuditReason(err)); auditErr != nil {
s.result <- auditErr
return
}
s.result <- err
return
}
if suppress {
if err := s.reportClipboardAudit("provider_to_client", "suppressed", clipboardAuditTextBytes(text), "loop"); err != nil {
s.result <- err
return
}
continue
}
if err := s.reportClipboardAudit(value.Direction, "forwarded", clipboardAuditTextBytes(value.Text), "forwarded"); err != nil {
s.result <- err
return
}
if err := s.sendClipboard(value); err != nil {
s.result <- err
return
}
}
}
}
func (s *gatewaySession) providerTelemetryLoop() {
s.observeProviderTelemetry()
ticker := time.NewTicker(500 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-s.ctx.Done():
return
case <-ticker.C:
s.observeProviderTelemetry()
}
}
}
func (s *gatewaySession) observeProviderTelemetry() {
telemetry := s.provider.Telemetry()
if telemetry.MediaDrops >= s.mediaDrops {
s.server.metrics.MediaDrops.Add(telemetry.MediaDrops - s.mediaDrops)
s.mediaDrops = telemetry.MediaDrops
}
s.server.metrics.observeProviderTelemetry(telemetry)
}
func (s *gatewaySession) controlLoop() {
for {
data, err := readWire(s.control, defaultControlLimit)
@@ -378,12 +500,32 @@ func (s *gatewaySession) controlLoop() {
}
switch frame.FlowID {
case "control":
if err := s.handleControl(payload); err != nil {
sequence, sequenceErr := channelSequence(frame.Sequence)
if sequenceErr != nil {
s.result <- sequenceErr
return
}
if err := s.handleControl(payload, sequence); err != nil {
s.result <- err
return
}
case "input":
if err := s.handleInput(payload); err != nil {
sequence, sequenceErr := channelSequence(frame.Sequence)
if sequenceErr != nil {
s.result <- sequenceErr
return
}
if err := s.handleInput(payload, sequence); err != nil {
s.result <- err
return
}
case "clipboard":
value, decodeErr := protocol.DecodeGatewayClipboardText(payload)
if decodeErr != nil {
s.result <- ErrProviderMalformed
return
}
if err := s.handleClipboard(value); err != nil {
s.result <- err
return
}
@@ -408,19 +550,13 @@ func (s *gatewaySession) datagramLoop() {
}
switch frame.Channel {
case ChannelInput:
if err := s.handleInput(frame.Payload); err != nil {
if err := s.handleInput(frame.Payload, frame.Sequence); err != nil {
s.result <- err
return
}
case ChannelText:
if len(frame.Payload) > 4096 {
s.result <- ErrFramePayloadLimit
return
}
if err := s.provider.Feedback(s.ctx, Feedback{Sequence: frame.Sequence, Payload: append([]byte(nil), frame.Payload...)}); err != nil {
s.result <- err
return
}
s.result <- ErrFrameChannel
return
default:
s.result <- ErrFrameChannel
return
@@ -458,6 +594,7 @@ func (s *gatewaySession) mediaLoop() {
}
func (s *gatewaySession) sendMedia(channel byte, payload []byte) error {
processingStarted := time.Now()
frames, err := FragmentPayload(channel, s.sequence.Add(1), uint64(time.Now().UnixMilli()), payload)
if err != nil {
return err
@@ -467,37 +604,122 @@ func (s *gatewaySession) sendMedia(channel byte, payload []byte) error {
if err != nil {
return err
}
if err := s.pacer.Wait(s.ctx, len(encoded)); err != nil {
pacingStarted := time.Now()
if err := s.server.pacer.wait(s.ctx, s.authority.SessionID, len(encoded)); err != nil {
return err
}
s.server.metrics.PacingDelayNanos.Add(uint64(time.Since(pacingStarted)))
s.server.metrics.QueueDelayNanos.Add(uint64(time.Since(pacingStarted)))
if err := s.connection.SendDatagram(encoded); err != nil {
return err
}
s.server.metrics.MediaPackets.Add(1)
s.server.metrics.MediaBytes.Add(uint64(len(encoded)))
s.server.metrics.ProcessingDelayNanos.Add(uint64(time.Since(processingStarted)))
s.server.metrics.ProcessingSamples.Add(1)
}
return nil
}
func (s *gatewaySession) handleControl(payload []byte) error {
if len(payload) < 4 {
return ErrProviderMalformed
func (s *gatewaySession) sendControl(sequence uint32, payload []byte) error {
if len(payload) > 1024 {
return ErrFramePayloadLimit
}
switch string(payload[:4]) {
case "TERM":
return errors.New("client requested termination")
case "RECN":
return s.provider.Reconnect(s.ctx)
case "FBRK":
return s.provider.Feedback(s.ctx, Feedback{Payload: append([]byte(nil), payload[4:]...)})
frame := protocol.ChannelFrame{Version: "1", FlowID: "control", Sequence: int64(sequence), Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
encoded, err := protocol.EncodeChannelFrame(frame)
if err != nil {
return err
}
s.controlWriteMu.Lock()
defer s.controlWriteMu.Unlock()
return writeWire(s.control, encoded, defaultControlLimit)
}
func (s *gatewaySession) sendClipboard(value protocol.GatewayClipboardText) error {
payload, err := protocol.EncodeGatewayClipboardText(value)
if err != nil {
return err
}
frame := protocol.ChannelFrame{Version: "1", FlowID: "clipboard", Sequence: int64(s.sequence.Add(1)), Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
encoded, err := protocol.EncodeChannelFrame(frame)
if err != nil {
return err
}
s.controlWriteMu.Lock()
defer s.controlWriteMu.Unlock()
return writeWire(s.control, encoded, defaultControlLimit)
}
func (s *gatewaySession) handleControl(payload []byte, sequence uint32) error {
feedback, err := DecodeClientFeedback(payload)
if err != nil {
return err
}
feedback.Sequence = sequence
return s.provider.Feedback(s.ctx, feedback)
}
func (s *gatewaySession) handleClipboard(value protocol.GatewayClipboardText) error {
if s.clipboard == nil {
return ErrClipboardDenied
}
suppress, err := s.clipboard.fromClient(value)
if err != nil {
if auditErr := s.reportClipboardAudit(value.Direction, "rejected", clipboardAuditTextBytes(value.Text), clipboardAuditReason(err)); auditErr != nil {
return auditErr
}
return err
}
if suppress {
return s.reportClipboardAudit(value.Direction, "suppressed", clipboardAuditTextBytes(value.Text), "loop")
}
if err := s.provider.WriteClipboard(s.ctx, value.Text); err != nil {
s.clipboard.retractClient(value)
if auditErr := s.reportClipboardAudit(value.Direction, "rejected", clipboardAuditTextBytes(value.Text), "provider"); auditErr != nil {
return auditErr
}
return err
}
return s.reportClipboardAudit(value.Direction, "forwarded", clipboardAuditTextBytes(value.Text), "forwarded")
}
func (s *gatewaySession) reportClipboardAudit(direction, outcome string, textBytes int, reason string) error {
if s.server.config.ClipboardAuditReporter == nil {
return ErrClipboardDenied
}
ctx, cancel := context.WithTimeout(s.ctx, 5*time.Second)
defer cancel()
return s.server.config.ClipboardAuditReporter.ReportClipboardAudit(ctx, protocol.GatewayClipboardAudit{
Version: "1", SessionID: s.authority.SessionID, Direction: direction, Outcome: outcome, TextBytes: int64(textBytes), Reason: reason,
})
}
func clipboardAuditTextBytes(text string) int {
if len(text) > 65536 {
return 65536
}
return len(text)
}
func clipboardAuditReason(err error) string {
switch {
case errors.Is(err, ErrClipboardRate):
return "rate"
case errors.Is(err, ErrProviderMalformed):
return "malformed"
case errors.Is(err, ErrClipboardDenied):
return "policy"
default:
return ErrProviderMalformed
return "provider"
}
}
func (s *gatewaySession) handleInput(payload []byte) error {
func (s *gatewaySession) handleInput(payload []byte, sequence uint32) error {
event, err := DecodeInputEvent(payload)
if err != nil {
return err
}
event.Sequence = sequence
if err := s.provider.Input(s.ctx, event); err != nil {
s.server.metrics.InputRejected.Add(1)
return err
@@ -513,9 +735,17 @@ func (s *gatewaySession) handleInput(payload []byte) error {
return nil
}
func channelSequence(sequence int64) (uint32, error) {
if sequence < 0 || sequence > int64(^uint32(0)) {
return 0, ErrProviderMalformed
}
return uint32(sequence), nil
}
func (s *gatewaySession) cleanup() {
s.cleanupOnce.Do(func() {
s.cancel()
s.server.pacer.remove(s.authority.SessionID)
cleanupCtx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
releaseInputsErr := s.provider.ReleaseAll(cleanupCtx)
@@ -613,9 +843,12 @@ func readWire(reader io.Reader, max int) ([]byte, error) {
}
type Client struct {
connection *quic.Conn
control *quic.Stream
Authority protocol.SessionAuthority
connection *quic.Conn
control *quic.Stream
controlReadMu sync.Mutex
controlWriteMu sync.Mutex
pendingControl map[string][][]byte
Authority protocol.SessionAuthority
}
func Dial(ctx context.Context, address string, tlsConfig *tls.Config, request protocol.TunnelAdmissionRequest) (*Client, error) {
@@ -658,7 +891,7 @@ func Dial(ctx context.Context, address string, tlsConfig *tls.Config, request pr
}
return nil, authorityErr
}
return &Client{connection: connection, control: stream, Authority: authority}, nil
return &Client{connection: connection, control: stream, pendingControl: make(map[string][][]byte), Authority: authority}, nil
}
func (c *Client) SendInput(event InputEvent) error {
@@ -671,15 +904,46 @@ func (c *Client) SendInput(event InputEvent) error {
if err != nil {
return err
}
return writeWire(c.control, encoded, defaultControlLimit)
return c.writeControl(encoded)
}
func (c *Client) SendControl(payload []byte) error {
frame := protocol.ChannelFrame{Version: "1", FlowID: "control", Sequence: 0, Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
return c.sendControl(0, payload)
}
func (c *Client) SendFeedback(feedback Feedback) error {
payload, err := EncodeClientFeedback(feedback)
if err != nil {
return err
}
return c.sendControl(feedback.Sequence, payload)
}
func (c *Client) SendClipboard(value protocol.GatewayClipboardText) error {
payload, err := protocol.EncodeGatewayClipboardText(value)
if err != nil {
return err
}
frame := protocol.ChannelFrame{Version: "1", FlowID: "clipboard", Sequence: 0, Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
encoded, err := protocol.EncodeChannelFrame(frame)
if err != nil {
return err
}
return c.writeControl(encoded)
}
func (c *Client) sendControl(sequence uint32, payload []byte) error {
frame := protocol.ChannelFrame{Version: "1", FlowID: "control", Sequence: int64(sequence), Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
encoded, err := protocol.EncodeChannelFrame(frame)
if err != nil {
return err
}
return c.writeControl(encoded)
}
func (c *Client) writeControl(encoded []byte) error {
c.controlWriteMu.Lock()
defer c.controlWriteMu.Unlock()
return writeWire(c.control, encoded, defaultControlLimit)
}
@@ -691,6 +955,68 @@ func (c *Client) ReceiveFrame(ctx context.Context) (Frame, error) {
return DecodeFrame(data)
}
func (c *Client) ReceiveProviderEvent(ctx context.Context) (ProviderEvent, error) {
payload, err := c.receiveControlPayload(ctx, "control")
if err != nil {
return ProviderEvent{}, err
}
if len(payload) > 1024 {
return ProviderEvent{}, ErrProviderMalformed
}
return DecodeProviderEvent(payload)
}
func (c *Client) ReceiveClipboard(ctx context.Context) (protocol.GatewayClipboardText, error) {
payload, err := c.receiveControlPayload(ctx, "clipboard")
if err != nil {
return protocol.GatewayClipboardText{}, err
}
return protocol.DecodeGatewayClipboardText(payload)
}
func (c *Client) receiveControlPayload(ctx context.Context, flowID string) ([]byte, error) {
if c == nil || c.control == nil || (flowID != "control" && flowID != "clipboard") {
return nil, ErrProviderMalformed
}
c.controlReadMu.Lock()
defer c.controlReadMu.Unlock()
if err := ctx.Err(); err != nil {
return nil, err
}
if queued := c.pendingControl[flowID]; len(queued) > 0 {
payload := queued[0]
c.pendingControl[flowID] = queued[1:]
return payload, nil
}
if deadline, ok := ctx.Deadline(); ok {
if err := c.control.SetReadDeadline(deadline); err != nil {
return nil, err
}
defer c.control.SetReadDeadline(time.Time{})
}
for {
data, err := readWire(c.control, defaultControlLimit)
if err != nil {
return nil, err
}
frame, err := protocol.DecodeChannelFrame(data)
if err != nil || (frame.FlowID != "control" && frame.FlowID != "clipboard") {
return nil, ErrProviderMalformed
}
payload, err := base64.StdEncoding.DecodeString(frame.Payload)
if err != nil || len(payload) > maxFrameSize {
return nil, ErrProviderMalformed
}
if frame.FlowID == flowID {
return payload, nil
}
if len(c.pendingControl[frame.FlowID]) >= clientControlBacklog {
return nil, ErrFramePayloadLimit
}
c.pendingControl[frame.FlowID] = append(c.pendingControl[frame.FlowID], payload)
}
}
func (c *Client) Close() error {
return c.connection.CloseWithError(applicationError, "client closed")
}