Просмотр исходного кода

fix: Max of 16 clients in an event stream

jamesread 3 недель назад
Родитель
Сommit
7c4eefa378
2 измененных файлов с 64 добавлено и 4 удалено
  1. 24 4
      service/internal/api/api.go
  2. 40 0
      service/internal/api/api_test.go

+ 24 - 4
service/internal/api/api.go

@@ -43,6 +43,11 @@ type oliveTinAPI struct {
 	streamingClientsMutex sync.RWMutex
 }
 
+// Caps concurrent EventStream connections to limit memory/FD/goroutine exhaustion.
+const maxEventStreamClients = 16
+
+var errEventStreamClientLimit = errors.New("too many concurrent event stream clients")
+
 // This is used to avoid race conditions when iterating over the connectedClients map.
 // and holds the lock for as minimal time as possible to avoid blocking the API for too long.
 func (api *oliveTinAPI) copyOfStreamingClients() []*streamingClient {
@@ -1028,14 +1033,14 @@ func (api *oliveTinAPI) EventStream(ctx ctx.Context, req *connect.Request[apiv1.
 		heartbeatDone:     make(chan struct{}),
 	}
 
+	if err := api.registerStreamingClient(client); err != nil {
+		return connect.NewError(connect.CodeResourceExhausted, err)
+	}
+
 	log.WithFields(log.Fields{
 		"authenticatedUser": user.Username,
 	}).Debugf("EventStream: client connected")
 
-	api.streamingClientsMutex.Lock()
-	api.streamingClients[client] = struct{}{}
-	api.streamingClientsMutex.Unlock()
-
 	go api.sendEventStreamHeartbeats(client)
 
 	// loop over client channel and send events to connectedClient
@@ -1054,6 +1059,21 @@ func (api *oliveTinAPI) EventStream(ctx ctx.Context, req *connect.Request[apiv1.
 	return nil
 }
 
+func (api *oliveTinAPI) registerStreamingClient(client *streamingClient) error {
+	api.streamingClientsMutex.Lock()
+	defer api.streamingClientsMutex.Unlock()
+
+	if len(api.streamingClients) >= maxEventStreamClients {
+		log.WithFields(log.Fields{
+			"limit": maxEventStreamClients,
+		}).Warn("EventStream: rejecting client; concurrent client limit reached")
+		return errEventStreamClientLimit
+	}
+
+	api.streamingClients[client] = struct{}{}
+	return nil
+}
+
 func (api *oliveTinAPI) sendEventStreamHeartbeats(client *streamingClient) {
 	defer close(client.heartbeatDone)
 

+ 40 - 0
service/internal/api/api_test.go

@@ -782,6 +782,46 @@ func TestEventStreamACLNoLeakToUnauthorizedUser(t *testing.T) {
 	assertEventStreamAdminReceivesSecretActionEvents(t, adminEvents)
 }
 
+func TestRegisterStreamingClientEnforcesLimit(t *testing.T) {
+	cfg := config.DefaultConfig()
+	ex := executor.DefaultExecutor(cfg)
+	api := newServer(ex)
+	user := &authpublic.AuthenticatedUser{Username: "limit-test"}
+
+	clients := make([]*streamingClient, 0, maxEventStreamClients)
+	for i := 0; i < maxEventStreamClients; i++ {
+		client := &streamingClient{
+			channel:           make(chan *apiv1.EventStreamResponse, 1),
+			AuthenticatedUser: user,
+			heartbeatStop:     make(chan struct{}),
+			heartbeatDone:     make(chan struct{}),
+		}
+		close(client.heartbeatDone)
+		require.NoError(t, api.registerStreamingClient(client))
+		clients = append(clients, client)
+	}
+
+	overflow := &streamingClient{
+		channel:           make(chan *apiv1.EventStreamResponse, 1),
+		AuthenticatedUser: user,
+		heartbeatStop:     make(chan struct{}),
+		heartbeatDone:     make(chan struct{}),
+	}
+	close(overflow.heartbeatDone)
+	err := api.registerStreamingClient(overflow)
+	assert.ErrorIs(t, err, errEventStreamClientLimit)
+	assert.Equal(t, maxEventStreamClients, len(api.streamingClients))
+
+	api.removeClient(clients[0])
+	require.NoError(t, api.registerStreamingClient(overflow))
+	assert.Equal(t, maxEventStreamClients, len(api.streamingClients))
+
+	for _, client := range clients[1:] {
+		api.removeClient(client)
+	}
+	api.removeClient(overflow)
+}
+
 func addEventStreamTestClients(t *testing.T, api *oliveTinAPI, lowUser, adminUser *authpublic.AuthenticatedUser) (*streamingClient, *streamingClient) {
 	t.Helper()
 	clientLow := &streamingClient{