Rate limit apply for gateway connections

Author: shanakama

platform-api/src/config/config.go

@@ -66,9 +66,10 @@ type JWT struct {
 
 // WebSocket holds WebSocket-specific configuration
 type WebSocket struct {
-	MaxConnections    int `envconfig:"WS_MAX_CONNECTIONS" default:"1000"`
-	ConnectionTimeout int `envconfig:"WS_CONNECTION_TIMEOUT" default:"30"` // seconds
-	RateLimitPerMin   int `envconfig:"WS_RATE_LIMIT_PER_MINUTE" default:"10"`
+	MaxConnections       int `envconfig:"WS_MAX_CONNECTIONS" default:"1000"`
+	ConnectionTimeout    int `envconfig:"WS_CONNECTION_TIMEOUT" default:"30"` // seconds
+	RateLimitPerMin      int `envconfig:"WS_RATE_LIMIT_PER_MINUTE" default:"10"`
+	MaxConnectionsPerOrg int `envconfig:"WS_MAX_CONNECTIONS_PER_ORG" default:"3"`
 }
 
 // Database holds database-specific configuration

platform-api/src/internal/handler/websocket.go

@@ -105,16 +105,34 @@ func (h *WebSocketHandler) Connect(c *gin.Context) {
 	transport := ws.NewWebSocketTransport(conn)
 
 	// Register connection with manager
-	connection, err := h.manager.Register(gateway.ID, transport, apiKey)
+	connection, err := h.manager.Register(gateway.ID, transport, apiKey, gateway.OrganizationID)
 	if err != nil {
-		log.Printf("[ERROR] Connection registration failed: gatewayID=%s error=%v", gateway.ID, err)
-		// Send error message before closing
-		errorMsg := map[string]string{
-			"type":    "error",
-			"message": err.Error(),
-		}
-		if jsonErr, _ := json.Marshal(errorMsg); jsonErr != nil {
-			conn.WriteMessage(websocket.TextMessage, jsonErr)
+		log.Printf("[ERROR] Connection registration failed: gatewayID=%s orgID=%s error=%v",
+			gateway.ID, gateway.OrganizationID, err)
+
+		// Check if this is an org connection limit error
+		if orgLimitErr, ok := err.(*ws.OrgConnectionLimitError); ok {
+			errorMsg := map[string]interface{}{
+				"type":         "error",
+				"code":         "ORG_CONNECTION_LIMIT_EXCEEDED",
+				"message":      "Organization connection limit reached",
+				"currentCount": orgLimitErr.CurrentCount,
+				"maxAllowed":   orgLimitErr.MaxAllowed,
+			}
+			if jsonErr, _ := json.Marshal(errorMsg); jsonErr != nil {
+				conn.WriteMessage(websocket.TextMessage, jsonErr)
+			}
+			log.Printf("[WARN] Organization connection limit exceeded: orgID=%s count=%d max=%d",
+				orgLimitErr.OrganizationID, orgLimitErr.CurrentCount, orgLimitErr.MaxAllowed)
+		} else {
+			// Generic error
+			errorMsg := map[string]string{
+				"type":    "error",
+				"message": err.Error(),
+			}
+			if jsonErr, _ := json.Marshal(errorMsg); jsonErr != nil {
+				conn.WriteMessage(websocket.TextMessage, jsonErr)
+			}
 		}
 		conn.Close()
 		return

platform-api/src/internal/repository/organization.go

@@ -46,6 +46,7 @@ func (r *OrganizationRepo) CreateOrganization(org *model.Organization) error {
 		VALUES (?, ?, ?, ?, ?, ?)
 	`
 	_, err := r.db.Exec(r.db.Rebind(query), org.ID, org.Handle, org.Name, org.Region, org.CreatedAt, org.UpdatedAt)
+
 	return err
 }
 
@@ -118,6 +119,7 @@ func (r *OrganizationRepo) UpdateOrganization(org *model.Organization) error {
 		WHERE uuid = ?
 	`
 	_, err := r.db.Exec(r.db.Rebind(query), org.Handle, org.Name, org.Region, org.UpdatedAt, org.ID)
+
 	return err
 }
 

platform-api/src/internal/server/server.go

@@ -79,9 +79,10 @@ func StartPlatformAPIServer(cfg *config.Server) (*Server, error) {
 
 	// Initialize WebSocket manager first (needed for GatewayEventsService)
 	wsConfig := websocket.ManagerConfig{
-		MaxConnections:    cfg.WebSocket.MaxConnections,
-		HeartbeatInterval: 20 * time.Second,
-		HeartbeatTimeout:  time.Duration(cfg.WebSocket.ConnectionTimeout) * time.Second,
+		MaxConnections:       cfg.WebSocket.MaxConnections,
+		HeartbeatInterval:    20 * time.Second,
+		HeartbeatTimeout:     time.Duration(cfg.WebSocket.ConnectionTimeout) * time.Second,
+		MaxConnectionsPerOrg: cfg.WebSocket.MaxConnectionsPerOrg,
 	}
 	wsManager := websocket.NewManager(wsConfig)
 
@@ -145,8 +146,9 @@ func StartPlatformAPIServer(cfg *config.Server) (*Server, error) {
 	gitHandler.RegisterRoutes(router)
 	deploymentHandler.RegisterRoutes(router)
 
-	log.Printf("[INFO] WebSocket manager initialized: maxConnections=%d heartbeatTimeout=%ds rateLimitPerMin=%d",
-		cfg.WebSocket.MaxConnections, cfg.WebSocket.ConnectionTimeout, cfg.WebSocket.RateLimitPerMin)
+	log.Printf("[INFO] WebSocket manager initialized: maxConnections=%d heartbeatTimeout=%ds rateLimitPerMin=%d maxConnectionsPerOrg=%d",
+		cfg.WebSocket.MaxConnections, cfg.WebSocket.ConnectionTimeout, cfg.WebSocket.RateLimitPerMin,
+		cfg.WebSocket.MaxConnectionsPerOrg)
 
 	return &Server{
 		router:      router,

platform-api/src/internal/websocket/connection.go

@@ -36,6 +36,10 @@ type Connection struct {
 	// Used to distinguish between multiple connections from the same gateway (clustering).
 	ConnectionID string
 
+	// OrganizationID identifies the organization that owns this gateway connection.
+	// Used for per-organization connection limit tracking and cleanup.
+	OrganizationID string
+
 	// ConnectedAt records when the connection was established
 	ConnectedAt time.Time
 
@@ -68,19 +72,21 @@ type Connection struct {
 //   - connectionID: Unique identifier for this connection instance
 //   - transport: Transport implementation (e.g., WebSocketTransport)
 //   - authToken: API key used for authentication
+//   - orgID: UUID of the organization that owns the gateway
 //
 // Returns a fully initialized Connection ready for message delivery.
-func NewConnection(gatewayID, connectionID string, transport Transport, authToken string) *Connection {
+func NewConnection(gatewayID, connectionID string, transport Transport, authToken string, orgID string) *Connection {
 	now := time.Now()
 	return &Connection{
-		GatewayID:     gatewayID,
-		ConnectionID:  connectionID,
-		ConnectedAt:   now,
-		LastHeartbeat: now,
-		Transport:     transport,
-		AuthToken:     authToken,
-		DeliveryStats: &DeliveryStats{},
-		closed:        false,
+		GatewayID:      gatewayID,
+		ConnectionID:   connectionID,
+		OrganizationID: orgID,
+		ConnectedAt:    now,
+		LastHeartbeat:  now,
+		Transport:      transport,
+		AuthToken:      authToken,
+		DeliveryStats:  &DeliveryStats{},
+		closed:         false,
 	}
 }
 

platform-api/src/internal/websocket/errors.go

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2026, WSO2 LLC. (http://www.wso2.org) All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package websocket
+
+import "fmt"
+
+// OrgConnectionLimitError is returned when an organization has reached its connection limit
+type OrgConnectionLimitError struct {
+	OrganizationID string
+	CurrentCount   int
+	MaxAllowed     int
+}
+
+func (e *OrgConnectionLimitError) Error() string {
+	return fmt.Sprintf("organization %s has reached maximum connection limit: %d/%d",
+		e.OrganizationID, e.CurrentCount, e.MaxAllowed)
+}
+
+// IsOrgConnectionLimitError checks if an error is an OrgConnectionLimitError
+func IsOrgConnectionLimitError(err error) bool {
+	_, ok := err.(*OrgConnectionLimitError)
+	return ok
+}

platform-api/src/internal/websocket/manager.go

@@ -53,6 +53,9 @@ type Manager struct {
 	// heartbeatTimeout specifies when to consider a connection dead (default 30s)
 	heartbeatTimeout time.Duration
 
+	// orgLimiter enforces per-organization connection limits
+	orgLimiter *OrgConnectionLimiter
+
 	// shutdownCtx is used to signal graceful shutdown to all connection goroutines
 	shutdownCtx context.Context
 	shutdownFn  context.CancelFunc
@@ -63,17 +66,19 @@ type Manager struct {
 
 // ManagerConfig contains configuration parameters for the connection manager
 type ManagerConfig struct {
-	MaxConnections    int           // Maximum concurrent connections (default 1000)
-	HeartbeatInterval time.Duration // Ping interval (default 20s)
-	HeartbeatTimeout  time.Duration // Pong timeout (default 30s)
+	MaxConnections       int           // Maximum concurrent connections (default 1000)
+	HeartbeatInterval    time.Duration // Ping interval (default 20s)
+	HeartbeatTimeout     time.Duration // Pong timeout (default 30s)
+	MaxConnectionsPerOrg int           // Maximum connections per organization (default 3)
 }
 
 // DefaultManagerConfig returns sensible default configuration values
 func DefaultManagerConfig() ManagerConfig {
 	return ManagerConfig{
-		MaxConnections:    1000,
-		HeartbeatInterval: 20 * time.Second,
-		HeartbeatTimeout:  30 * time.Second,
+		MaxConnections:       1000,
+		HeartbeatInterval:    20 * time.Second,
+		HeartbeatTimeout:     30 * time.Second,
+		MaxConnectionsPerOrg: 3,
 	}
 }
 
@@ -86,6 +91,7 @@ func NewManager(config ManagerConfig) *Manager {
 		maxConnections:    config.MaxConnections,
 		heartbeatInterval: config.HeartbeatInterval,
 		heartbeatTimeout:  config.HeartbeatTimeout,
+		orgLimiter:        NewOrgConnectionLimiter(config.MaxConnectionsPerOrg),
 		shutdownCtx:       ctx,
 		shutdownFn:        cancel,
 	}
@@ -98,25 +104,37 @@ func NewManager(config ManagerConfig) *Manager {
 //   - gatewayID: UUID of the authenticated gateway
 //   - transport: Transport implementation for message delivery
 //   - authToken: API key used for authentication
+//   - orgID: UUID of the organization that owns the gateway
 //
 // Returns the Connection instance and any error encountered.
 //
 // Design decision: Support multiple connections per gateway ID by storing
 // connections in a slice. This enables gateway clustering where multiple
 // instances share the same gateway identity.
-func (m *Manager) Register(gatewayID string, transport Transport, authToken string) (*Connection, error) {
-	// Check connection limit
+func (m *Manager) Register(gatewayID string, transport Transport, authToken string,
+	orgID string) (*Connection, error) {
+
+	// Create connection ID early so we can use it for org limiter
+	connectionID := uuid.New().String()
+
+	// Check per-org limit first
+	if err := m.orgLimiter.AddConnection(orgID, connectionID); err != nil {
+		return nil, err
+	}
+
+	// Check global connection limit
 	m.mu.Lock()
 	if m.connectionCount >= m.maxConnections {
 		m.mu.Unlock()
+		// Rollback org limiter
+		m.orgLimiter.RemoveConnection(orgID, connectionID)
 		return nil, fmt.Errorf("maximum connection limit reached (%d)", m.maxConnections)
 	}
 	m.connectionCount++
 	m.mu.Unlock()
 
-	// Create connection with unique connection ID
-	connectionID := uuid.New().String()
-	conn := NewConnection(gatewayID, connectionID, transport, authToken)
+	// Create connection
+	conn := NewConnection(gatewayID, connectionID, transport, authToken, orgID)
 
 	// Add connection to registry
 	connsInterface, _ := m.connections.LoadOrStore(gatewayID, []*Connection{})
@@ -128,8 +146,8 @@ func (m *Manager) Register(gatewayID string, transport Transport, authToken stri
 	m.wg.Add(1)
 	go m.monitorHeartbeat(conn)
 
-	log.Printf("[INFO] Gateway connected: gatewayID=%s connectionID=%s totalConnections=%d",
-		gatewayID, connectionID, m.GetConnectionCount())
+	log.Printf("[INFO] Gateway connected: gatewayID=%s connectionID=%s orgID=%s totalConnections=%d orgConnections=%d",
+		gatewayID, connectionID, orgID, m.GetConnectionCount(), m.orgLimiter.GetOrgConnectionCount(orgID))
 
 	return conn, nil
 }
@@ -176,13 +194,18 @@ func (m *Manager) Unregister(gatewayID, connectionID string) {
 			gatewayID, connectionID, err)
 	}
 
+	// Remove from org limiter
+	if removed.OrganizationID != "" {
+		m.orgLimiter.RemoveConnection(removed.OrganizationID, connectionID)
+	}
+
 	// Decrement connection count
 	m.mu.Lock()
 	m.connectionCount--
 	m.mu.Unlock()
 
-	log.Printf("[INFO] Gateway disconnected: gatewayID=%s connectionID=%s totalConnections=%d",
-		gatewayID, connectionID, m.GetConnectionCount())
+	log.Printf("[INFO] Gateway disconnected: gatewayID=%s connectionID=%s orgID=%s totalConnections=%d",
+		gatewayID, connectionID, removed.OrganizationID, m.GetConnectionCount())
 }
 
 // GetConnections retrieves all connections for a specific gateway ID.
@@ -301,3 +324,13 @@ func (m *Manager) Shutdown() {
 
 	log.Println("[INFO] WebSocket manager shutdown complete")
 }
+
+// GetOrgConnectionStats returns connection statistics for a specific organization
+func (m *Manager) GetOrgConnectionStats(orgID string) OrgConnectionStats {
+	return m.orgLimiter.GetOrgStats(orgID)
+}
+
+// GetAllOrgConnectionStats returns connection counts for all organizations
+func (m *Manager) GetAllOrgConnectionStats() map[string]int {
+	return m.orgLimiter.GetAllOrgConnectionCounts()
+}