fix(gateway): secure control and terminal ownership

gateway/transport.go

@@ -23,6 +23,7 @@ const (
 	defaultHelloLimit    = 16 * 1024
 	defaultControlLimit  = 128 * 1024
 	clientControlBacklog = 64
+	terminalAckTimeout   = 2 * time.Second
 	applicationError     = quic.ApplicationErrorCode(0x100)
 	controlFlowID        = "control.ack.v1"
 	inputFlowID          = "input.sequenced.v1"
@@ -394,18 +395,21 @@ type gatewaySession struct {
 	inputMu        sync.Mutex
 	controlWriteMu sync.Mutex
 	outputMu       sync.Mutex
+	terminalMu     sync.Mutex
 	pressed        map[string]struct{}
 	sequence       atomic.Uint32
 	mediaDrops     uint64
 	mediaQuiesced  bool
 	terminalSent   atomic.Bool
+	terminalAwait  bool
+	terminalAck    chan struct{}
 	endReason      error
 	result         chan error
 }
 
 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, clipboard: clipboard, 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{}), terminalAck: make(chan struct{}, 1), result: make(chan error, 3)}
 }
 
 func (s *gatewaySession) run() {
@@ -436,10 +440,6 @@ func (s *gatewaySession) run() {
 	s.cancel()
 	if s.terminalSent.Load() {
 		s.cleanup()
-		select {
-		case <-s.connection.Context().Done():
-		case <-timer.C:
-		}
 	}
 }
 
@@ -457,12 +457,15 @@ func (s *gatewaySession) providerEventLoop() {
 			if terminal {
 				s.outputMu.Lock()
 				s.mediaQuiesced = true
+				s.terminalMu.Lock()
 			}
 			payload, err := EncodeProviderEvent(event)
 			if err == nil {
 				err = s.sendControl(s.sequence.Add(1), payload)
 			}
 			if terminal {
+				s.terminalAwait = err == nil
+				s.terminalMu.Unlock()
 				s.outputMu.Unlock()
 			}
 			if err != nil {
@@ -471,6 +474,23 @@ func (s *gatewaySession) providerEventLoop() {
 			}
 			if terminal {
 				s.terminalSent.Store(true)
+				timer := time.NewTimer(terminalAckTimeout)
+				select {
+				case <-s.terminalAck:
+				case <-timer.C:
+					s.terminalMu.Lock()
+					s.terminalAwait = false
+					s.terminalMu.Unlock()
+					s.result <- context.DeadlineExceeded
+					return
+				case <-s.ctx.Done():
+					timer.Stop()
+					s.terminalMu.Lock()
+					s.terminalAwait = false
+					s.terminalMu.Unlock()
+					return
+				}
+				timer.Stop()
 			}
 			switch event.Kind {
 			case ProviderEventTerminated:
@@ -747,6 +767,20 @@ func (s *gatewaySession) handleControl(payload []byte, sequence uint32) error {
 	if err != nil {
 		return err
 	}
+	if feedback.Kind == FeedbackTerminalReceipt {
+		s.terminalMu.Lock()
+		defer s.terminalMu.Unlock()
+		if !s.terminalAwait {
+			return ErrProviderMalformed
+		}
+		s.terminalAwait = false
+		select {
+		case s.terminalAck <- struct{}{}:
+			return nil
+		default:
+			return ErrProviderMalformed
+		}
+	}
 	feedback.Sequence = sequence
 	return s.provider.Feedback(s.ctx, feedback)
 }
@@ -857,9 +891,7 @@ func (s *gatewaySession) cleanup() {
 			} else if err := s.server.config.Admission.Release(cleanupCtx, s.authority); err != nil {
 				s.server.metrics.ProviderErrors.Add(1)
 			}
-			if !s.terminalSent.Load() {
-				_ = s.connection.CloseWithError(applicationError, "session closed")
-			}
+			_ = s.connection.CloseWithError(applicationError, "session closed")
 			return
 		}
 		if errors.Is(s.endReason, ErrProviderDisconnected) {
@@ -872,9 +904,7 @@ func (s *gatewaySession) cleanup() {
 		if err := s.server.reportProviderState(cleanupCtx, state); err != nil {
 			s.server.metrics.ProviderErrors.Add(1)
 		}
-		if !s.terminalSent.Load() {
-			_ = s.connection.CloseWithError(applicationError, "session closed")
-		}
+		_ = s.connection.CloseWithError(applicationError, "session closed")
 	})
 }
 
@@ -1065,7 +1095,7 @@ func (c *Client) ReceiveProviderEvent(ctx context.Context) (ProviderEvent, error
 	}
 	event, err := DecodeProviderEvent(payload)
 	if err == nil && (event.Kind == ProviderEventTerminated || event.Kind == ProviderEventDisconnected) {
-		_ = c.Close()
+		err = c.SendFeedback(Feedback{Kind: FeedbackTerminalReceipt})
 	}
 	return event, err
 }