Fix OCM websocket proxy lifecycle and headers

service/ocm/service.go

@@ -139,7 +139,9 @@ type Service struct {
 	userManager    *UserManager
 	accessMutex    sync.RWMutex
 	usageTracker   *AggregatedUsage
-	trackingGroup  sync.WaitGroup
+	webSocketMutex sync.Mutex
+	webSocketGroup sync.WaitGroup
+	webSocketConns map[*webSocketSession]struct{}
 	shuttingDown   bool
 }
 
@@ -197,8 +199,9 @@ func NewService(ctx context.Context, logger log.ContextLogger, tag string, optio
 			Network: []string{N.NetworkTCP},
 			Listen:  options.ListenOptions,
 		}),
-		userManager:  userManager,
-		usageTracker: usageTracker,
+		userManager:    userManager,
+		usageTracker:   usageTracker,
+		webSocketConns: make(map[*webSocketSession]struct{}),
 	}
 
 	if options.TLS != nil {
@@ -631,11 +634,17 @@ func (s *Service) handleResponseWithTracking(writer http.ResponseWriter, respons
 }
 
 func (s *Service) Close() error {
+	webSocketSessions := s.startWebSocketShutdown()
+
 	err := common.Close(
 		common.PtrOrNil(s.httpServer),
 		common.PtrOrNil(s.listener),
 		s.tlsConfig,
 	)
+	for _, session := range webSocketSessions {
+		session.Close()
+	}
+	s.webSocketGroup.Wait()
 
 	if s.usageTracker != nil {
 		s.usageTracker.cancelPendingSave()
@@ -647,3 +656,48 @@ func (s *Service) Close() error {
 
 	return err
 }
+
+func (s *Service) registerWebSocketSession(session *webSocketSession) bool {
+	s.webSocketMutex.Lock()
+	defer s.webSocketMutex.Unlock()
+
+	if s.shuttingDown {
+		return false
+	}
+
+	s.webSocketConns[session] = struct{}{}
+	s.webSocketGroup.Add(1)
+	return true
+}
+
+func (s *Service) unregisterWebSocketSession(session *webSocketSession) {
+	s.webSocketMutex.Lock()
+	_, loaded := s.webSocketConns[session]
+	if loaded {
+		delete(s.webSocketConns, session)
+	}
+	s.webSocketMutex.Unlock()
+
+	if loaded {
+		s.webSocketGroup.Done()
+	}
+}
+
+func (s *Service) isShuttingDown() bool {
+	s.webSocketMutex.Lock()
+	defer s.webSocketMutex.Unlock()
+	return s.shuttingDown
+}
+
+func (s *Service) startWebSocketShutdown() []*webSocketSession {
+	s.webSocketMutex.Lock()
+	defer s.webSocketMutex.Unlock()
+
+	s.shuttingDown = true
+
+	webSocketSessions := make([]*webSocketSession, 0, len(s.webSocketConns))
+	for session := range s.webSocketConns {
+		webSocketSessions = append(webSocketSessions, session)
+	}
+	return webSocketSessions
+}