diff --git a/ebpf/connections/factory.go b/ebpf/connections/factory.go index 19718f68..99719f2b 100644 --- a/ebpf/connections/factory.go +++ b/ebpf/connections/factory.go @@ -40,7 +40,7 @@ func NewFactory() *Factory { } } -func convertToSingleByteArr(bufMap map[int][]byte) []byte { +func convertToSingleByteArr(bufMap map[int][]byte, isWebsocket bool) []byte { if len(bufMap) == 0 { return make([]byte, 0) @@ -59,17 +59,17 @@ func convertToSingleByteArr(bufMap map[int][]byte) []byte { for _, k := range keys { if kPrev == -1 { // C sets read, write event count=0 only on new connection open - // For requests arriving after a time gap on the same underlying connection the + // For requests arriving after a time gap on the same underlying connection the // read,write count will not be 1, they will simply continue from the last request // This can only be replicated when there is a time gap/inactivityThreshold between requests // on the same underlying connection - if !sequenceCheckSkip && k != 1 { + if !isWebsocket && !sequenceCheckSkip && k != 1 { utils.LogProcessing("Bad start sequence", "key", k, "value", string(bufMap[k])) break } kPrev = k } else { - if kPrev+1 != k { + if !isWebsocket && kPrev+1 != k { utils.LogProcessing("Missing sequence", "prev", kPrev, "current", k, "value", string(bufMap[k]), "prevValue", string(bufMap[kPrev])) break } @@ -105,6 +105,39 @@ func init() { utils.InitVar("SOCKET_DATA_EVENT_BYTES_THRESHOLD", &socketDataEventBytesThreshold) } +func isKnownWebSocketConnection(connID structs.ConnID) bool { + if kafkaUtil.WSConnectionManager == nil { + return false + } + return kafkaUtil.WSConnectionManager.IsRegisteredByConnID(connID.Id, connID.Fd) +} + +// hasCloseFrame checks if any of the frames contain a WebSocket close frame (opcode 8) +func hasCloseFrame(frames []kafkaUtil.WebSocketMessage) bool { + for _, f := range frames { + if f.Opcode == 0x8 { + return true + } + } + return false +} + +// extractHTTPHeadersOnly returns the buffer up to and including \r\n\r\n (HTTP headers only) +// This strips any WebSocket frames that may follow the headers +func extractHTTPHeadersOnly(buffer []byte) []byte { + if len(buffer) < 4 { + return buffer + } + + headerEnd := bytes.Index(buffer, []byte("\r\n\r\n")) + if headerEnd == -1 { + + return buffer + } + + return buffer[:headerEnd+4] +} + func ProcessTrackerData(connID structs.ConnID, tracker *Tracker, isComplete bool) { tracker.mutex.Lock() defer tracker.mutex.Unlock() @@ -112,8 +145,12 @@ func ProcessTrackerData(connID structs.ConnID, tracker *Tracker, isComplete bool if len(tracker.sentBuf) == 0 || len(tracker.recvBuf) == 0 { return } - receiveBuffer := convertToSingleByteArr(tracker.recvBuf) - sentBuffer := convertToSingleByteArr(tracker.sentBuf) + + // Check if this is already a known WebSocket connection - if so, skip sequence checks + // since binary WebSocket frames won't have sequential packet numbering + isKnownWebSocket := isKnownWebSocketConnection(connID) + receiveBuffer := convertToSingleByteArr(tracker.recvBuf, isKnownWebSocket) + sentBuffer := convertToSingleByteArr(tracker.sentBuf, isKnownWebSocket) originalInt := uint32(connID.Ip) // Convert integer to little-endian byte slice @@ -134,10 +171,33 @@ func ProcessTrackerData(connID structs.ConnID, tracker *Tracker, isComplete bool hostName = kafkaUtil.PodInformerInstance.GetPodNameByProcessId(int32(connID.Id >> 32)) } - if len(sentBuffer) >= len(httpBytes) && (bytes.Equal(sentBuffer[:len(httpBytes)], httpBytes)) { + if isKnownWebSocket { + processWebSocketConnection(connID, receiveBuffer, sentBuffer, isComplete, hostName) + return + } + + isWebSocket := isWebSocketUpgradeData(sentBuffer, receiveBuffer) + + if isWebSocket { + slog.Info("Detected WebSocket upgrade, registering connection", + "fd", connID.Fd, + "id", connID.Id, + "ssl", tracker.ssl, + ) + + headers := extractHeadersFromHTTPBuffer(receiveBuffer) + + kafkaUtil.WSConnectionManager.RegisterConnectionNumeric(connID.Id, connID.Fd, connID.Ip, connID.Port, tracker.srcIp, tracker.srcPort, headers) + + // Send the initial handshake as a normal HTTP request/response pair + // Strip WebSocket frames from buffers first - only pass HTTP headers to tryReadFromBD + httpReceiveBuffer := extractHTTPHeadersOnly(receiveBuffer) + httpSentBuffer := extractHTTPHeadersOnly(sentBuffer) + tryReadFromBD(destIpStr, srcIpStr, httpReceiveBuffer, httpSentBuffer, isComplete, 1, connID.Id, connID.Fd, uniqueDaemonsetId, hostName) + } else if len(sentBuffer) >= len(httpBytes) && (bytes.Equal(sentBuffer[:len(httpBytes)], httpBytes)) { tryReadFromBD(destIpStr, srcIpStr, receiveBuffer, sentBuffer, isComplete, 1, connID.Id, connID.Fd, uniqueDaemonsetId, hostName) } - if !disableEgress { + if !disableEgress && !isWebSocket { // attempt to parse the egress as well by switching the recv and sent buffers. if len(receiveBuffer) >= len(httpBytes) && (bytes.Equal(receiveBuffer[:len(httpBytes)], httpBytes)) { tryReadFromBD(srcIpStr, destIpStr, sentBuffer, receiveBuffer, isComplete, 2, connID.Id, connID.Fd, uniqueDaemonsetId, hostName) @@ -145,6 +205,93 @@ func ProcessTrackerData(connID structs.ConnID, tracker *Tracker, isComplete bool } } +func isWebSocketUpgradeData(sentBuffer, receiveBuffer []byte) bool { + if len(sentBuffer) == 0 || len(receiveBuffer) == 0 { + return false + } + + if bytes.Contains(sentBuffer, []byte(":9092")) || bytes.Contains(receiveBuffer, []byte(":9092")) { + return false + } + + // Check for 101 Switching Protocols in sent buffer (response) + has101 := bytes.Contains(sentBuffer, []byte("101 Switching Protocols")) + + hasUpgradeResponseHeader := bytes.Contains(sentBuffer, []byte("Upgrade: websocket")) || + bytes.Contains(sentBuffer, []byte("upgrade: websocket")) + + hasUpgradeRequestHeader := bytes.Contains(receiveBuffer, []byte("Upgrade: websocket")) || + bytes.Contains(receiveBuffer, []byte("upgrade: websocket")) + + hasConnectionResponseHeader := bytes.Contains(sentBuffer, []byte("Connection: Upgrade")) || + bytes.Contains(sentBuffer, []byte("Connection: upgrade")) || + bytes.Contains(sentBuffer, []byte("connection: Upgrade")) || + bytes.Contains(sentBuffer, []byte("connection: upgrade")) + + hasConnectionRequestHeader := bytes.Contains(receiveBuffer, []byte("Connection: Upgrade")) || + bytes.Contains(receiveBuffer, []byte("Connection: upgrade")) || + bytes.Contains(receiveBuffer, []byte("connection: Upgrade")) || + bytes.Contains(receiveBuffer, []byte("connection: upgrade")) + + return has101 && hasUpgradeResponseHeader && hasUpgradeRequestHeader && + hasConnectionResponseHeader && hasConnectionRequestHeader +} + +func extractHeadersFromHTTPBuffer(buffer []byte) map[string]string { + headers := make(map[string]string) + if len(buffer) == 0 { + return headers + } + + lines := bytes.Split(buffer, []byte("\r\n")) + + for i := 1; i < len(lines); i++ { + line := lines[i] + + if len(line) == 0 { + break + } + + parts := bytes.SplitN(line, []byte(":"), 2) + if len(parts) == 2 { + key := string(bytes.TrimSpace(parts[0])) + value := string(bytes.TrimSpace(parts[1])) + headers[key] = value + } + } + + return headers +} + +func processWebSocketConnection(connID structs.ConnID, receiveBuffer, sentBuffer []byte, isComplete bool, hostName string) { + connectionClosed := false + + if len(sentBuffer) > 0 { + msgs, err := kafkaUtil.WSConnectionManager.AccumulatePayloadNumeric(connID.Id, connID.Fd, sentBuffer, "outgoing") + if err != nil { + slog.Debug("Failed to accumulate WebSocket payload from sent buffer", "error", err) + } else if hasCloseFrame(msgs) { + connectionClosed = true + } + } + + if len(receiveBuffer) > 0 { + msgs, err := kafkaUtil.WSConnectionManager.AccumulatePayloadNumeric(connID.Id, connID.Fd, receiveBuffer, "incoming") + if err != nil { + slog.Debug("Failed to accumulate WebSocket payload from receive buffer", "error", err) + } else if hasCloseFrame(msgs) { + connectionClosed = true + } + } + + // If a close frame was seen, unmark so the fd can be reused cleanly by a new connection + if connectionClosed { + kafkaUtil.WSConnectionManager.RemoveConnectionByConnID(connID.Id, connID.Fd) + } + + slog.Debug("WebSocket connection processed", "connID", connID, "isComplete", isComplete) +} + func (factory *Factory) CanBeFilled() bool { factory.mutex.RLock() defer factory.mutex.RUnlock() @@ -310,6 +457,11 @@ func (factory *Factory) DeleteWorker(connectionID structs.ConnID) { } if _, exists := factory.connections[connectionID]; exists { + + if _, ok := factory.connections[connectionID]; ok { + // Don't unmark WebSocket connections - they maintain their type across the connection lifetime + // even if factory.connections is cleaned up during inactivity + } delete(factory.connections, connectionID) utils.LogProcessing("Deleted connection", "fd", connectionID.Fd, "id", connectionID.Id, "timestamp", connectionID.Conn_start_ns, "ip", connectionID.Ip, "port", connectionID.Port) requestProcessCount++ @@ -332,7 +484,13 @@ func (factory *Factory) DeleteWorker(connectionID structs.ConnID) { close(ch) delete(factory.processor, key) } + if _, ok := factory.connections[key]; ok { + // Don't unmark WebSocket connections - they maintain their type across the connection lifetime + } delete(factory.connections, key) + if kafkaUtil.WSConnectionManager != nil { + kafkaUtil.WSConnectionManager.RemoveConnectionByConnID(key.Id, key.Fd) + } } } } diff --git a/trafficUtil/kafkaUtil/kafka.go b/trafficUtil/kafkaUtil/kafka.go index 653d2837..5b33ef58 100644 --- a/trafficUtil/kafkaUtil/kafka.go +++ b/trafficUtil/kafkaUtil/kafka.go @@ -143,11 +143,105 @@ func InitKafka() { // Start heartbeat routine go sendKafkaHeartbeat() slog.Info("Started Kafka heartbeat routine", "interval_seconds", heartbeatIntervalSeconds) + + go sendWebSocketBatches() + slog.Info("Started WebSocket batch sender routine", "interval_minutes", 1) break } } } +func sendWebSocketBatches() { + ticker := time.NewTicker(1 * time.Minute) + defer ticker.Stop() + + for range ticker.C { + if WSConnectionManager == nil { + continue + } + + connections := WSConnectionManager.GetAllConnections() + slog.Debug("WebSocket batch sender tick", "active_connections", len(connections)) + + for _, key := range connections { + + parts := strings.Split(key, ":") + if len(parts) != 2 { + slog.Warn("Invalid WebSocket connection key format", "key", key) + continue + } + + connIDId, err := strconv.ParseUint(parts[0], 10, 64) + if err != nil { + slog.Warn("Failed to parse ConnID Id", "key", key, "error", err) + continue + } + connIDFd, err := strconv.ParseUint(parts[1], 10, 32) + if err != nil { + slog.Warn("Failed to parse ConnID Fd", "key", key, "error", err) + continue + } + + batch, err := WSConnectionManager.GetAndClearBatchByConnID(connIDId, uint32(connIDFd)) + if err != nil { + slog.Debug("Failed to get WebSocket batch", "key", key, "error", err) + continue + } + + if len(batch.Messages) == 0 { + continue + } + + eventsJSON, _ := json.Marshal(convertMessagesToEventList(batch.Messages)) + headersJSON, _ := json.Marshal(batch.Connection.Headers) + + payload := map[string]string{ + "ip": batch.Connection.SourceIP, + "destIp": batch.Connection.DestIP, + "time": fmt.Sprint(batch.BatchTime.Unix()), + "akto_account_id": fmt.Sprint(1000000), + "akto_vxlan_id": fmt.Sprint(0), + "source": "MIRRORING", + "connection_type": "WEBSOCKET", + "headers": string(headersJSON), + "events": string(eventsJSON), + } + + out, err := json.Marshal(payload) + if err != nil { + slog.Error("Failed to marshal WebSocket batch payload", "error", err) + continue + } + + ctx := context.Background() + sourceIP := batch.Connection.SourceIP + err = ProduceStr(ctx, string(out), key, sourceIP, "WEBSOCKET") + if err != nil { + slog.Error("Failed to write WebSocket batch to Kafka", "error", err) + } else { + slog.Debug("WebSocket batch sent to Kafka", "messages", len(batch.Messages), "source", sourceIP) + } + } + } +} + +// convertMessagesToEventList converts WebSocketMessage slice to a list of event maps +func convertMessagesToEventList(messages []WebSocketMessage) []map[string]interface{} { + eventList := make([]map[string]interface{}, 0) + for _, msg := range messages { + eventList = append(eventList, map[string]interface{}{ + "event_type": msg.EventType, + "payload": msg.Payload, + "direction": msg.Direction, + "timestamp": msg.Timestamp.Unix(), + "fin": msg.FIN, + "opcode": msg.Opcode, + "masked": msg.Masked, + }) + } + return eventList +} + func kafkaCompletion() func(messages []kafka.Message, err error) { return func(messages []kafka.Message, err error) { if err != nil { diff --git a/trafficUtil/kafkaUtil/parser.go b/trafficUtil/kafkaUtil/parser.go index dd9e2adc..76c3c5f2 100644 --- a/trafficUtil/kafkaUtil/parser.go +++ b/trafficUtil/kafkaUtil/parser.go @@ -608,6 +608,220 @@ func parseHTTPTraffic(reqBuffer, respBuffer []byte, shouldPrint bool) *ParsedTra } } +// WSFragmentAssembler reassembles fragmented WebSocket messages per RFC 6455 ยง5.4. +type WSFragmentAssembler struct { + active bool + opcode uint8 + buf bytes.Buffer +} + +func ParseWebSocketFrames(payload []byte, direction string, asm *WSFragmentAssembler) []WebSocketMessage { + messages := []WebSocketMessage{} + + if len(payload) == 0 { + return messages + } + + emitMessage := func(opcode uint8, fin bool, masked bool, framePayload []byte) { + + messages = append(messages, WebSocketMessage{ + Payload: string(framePayload), + Direction: direction, + Timestamp: time.Now(), + FIN: fin, + Opcode: opcode, + Masked: masked, + EventType: "", + }) + } + + resetAssembler := func() { + if asm != nil { + asm.active = false + asm.opcode = 0 + asm.buf.Reset() + } + } + + i := 0 + for i < len(payload) { + if i+2 > len(payload) { + break + } + + byte0 := payload[i] + fin := (byte0 & 0x80) != 0 + opcode := byte0 & 0x0F + i++ + + byte1 := payload[i] + masked := (byte1 & 0x80) != 0 + payloadLen := int(byte1 & 0x7F) + i++ + + if payloadLen == 126 { + if i+2 > len(payload) { + break + } + payloadLen = int(payload[i])<<8 | int(payload[i+1]) + i += 2 + } else if payloadLen == 127 { + if i+8 > len(payload) { + break + } + payloadLen = int(payload[i])<<56 | int(payload[i+1])<<48 | + int(payload[i+2])<<40 | int(payload[i+3])<<32 | + int(payload[i+4])<<24 | int(payload[i+5])<<16 | + int(payload[i+6])<<8 | int(payload[i+7]) + i += 8 + } + + var maskingKey [4]byte + if masked { + if i+4 > len(payload) { + break + } + copy(maskingKey[:], payload[i:i+4]) + i += 4 + } + + if i+payloadLen > len(payload) { + break + } + framePayload := payload[i : i+payloadLen] + i += payloadLen + + if masked { + for j := 0; j < len(framePayload); j++ { + framePayload[j] ^= maskingKey[j%4] + } + } + + if asm == nil { + msg := WebSocketMessage{ + Payload: string(framePayload), + Direction: direction, + Timestamp: time.Now(), + FIN: fin, + Opcode: opcode, + Masked: masked, + } + messages = append(messages, msg) + continue + } + + slog.Debug("WebSocket frame decoded", + "direction", direction, + "opcode", opcode, + "fin", fin, + "masked", masked, + "payloadLen", len(framePayload), + "asmActive", asm.active, + "asmOpcode", asm.opcode, + "asmBufLen", asm.buf.Len()) + + switch { + case opcode >= 0x8: + if !fin { + resetAssembler() + continue + } + + emitMessage(opcode, fin, masked, framePayload) + + case opcode == 0x1 || opcode == 0x2: + if asm.active { + + resetAssembler() + } + if fin { + emitMessage(opcode, true, masked, framePayload) + } else { + + asm.active = true + asm.opcode = opcode + asm.buf.Reset() + asm.buf.Write(framePayload) + } + + case opcode == 0x0: + if !asm.active { + + continue + } + asm.buf.Write(framePayload) + + if fin { + emitMessage(asm.opcode, true, masked, asm.buf.Bytes()) + resetAssembler() + } + + default: + slog.Debug("WebSocket unsupported opcode", + "direction", direction, "opcode", opcode, "fin", fin) + resetAssembler() + } + } + + if i < len(payload) { + slog.Debug("WebSocket parse stopped with trailing bytes", + "direction", direction, "consumed", i, "total", len(payload), "remaining", len(payload)-i) + } + + slog.Debug("WebSocket parse complete", + "direction", direction, "inputLen", len(payload), "messagesEmitted", len(messages)) + + return messages +} + +func detectWebSocketUpgrade(resp *http.Response) bool { + if resp.StatusCode != http.StatusSwitchingProtocols { + return false + } + + upgrade := strings.ToLower(resp.Header.Get("Upgrade")) + connection := strings.ToLower(resp.Header.Get("Connection")) + + isUpgrade := upgrade == "websocket" + isConnection := strings.Contains(connection, "upgrade") + + return isUpgrade && isConnection +} + +func ExtractIP(ipPort string) string { + if ipPort == "" { + return "" + } + + if strings.HasPrefix(ipPort, "[") { + if idx := strings.LastIndex(ipPort, "]"); idx > 0 { + return ipPort[1:idx] + } + } + + if idx := strings.LastIndex(ipPort, ":"); idx > 0 { + return ipPort[:idx] + } + return ipPort +} + +func ExtractPort(ipPort string) string { + if ipPort == "" { + return "0" + } + + if strings.HasPrefix(ipPort, "[") { + if idx := strings.LastIndex(ipPort, "]"); idx > 0 && idx+1 < len(ipPort) && ipPort[idx+1] == ':' { + return ipPort[idx+2:] + } + } + + if idx := strings.LastIndex(ipPort, ":"); idx > 0 { + return ipPort[idx+1:] + } + return "0" +} + func ParseAndProduce(receiveBuffer []byte, sentBuffer []byte, ctx TrafficContext) { if KafkaDisabled { return @@ -667,6 +881,61 @@ func ParseAndProduce(receiveBuffer []byte, sentBuffer []byte, ctx TrafficContext continue } + // Check for WebSocket upgrade + if detectWebSocketUpgrade(resp) { + // Extract ports from context IPs (format: "IP:Port") + // In eBPF flow: ctx.SourceIP is dest IP in "IP:Port" format, ctx.DestIP is source IP + sourcePort := ExtractPort(ctx.SourceIP) + destPort := ExtractPort(ctx.DestIP) + + // Register the WebSocket connection + err := WSConnectionManager.RegisterConnection( + ExtractIP(ctx.SourceIP), + sourcePort, + ExtractIP(ctx.DestIP), + destPort, + req, + ) + if err != nil { + slog.Warn("Failed to register WebSocket connection", "error", err) + } else { + slog.Info("WebSocket connection upgraded", "sourceIP", ctx.SourceIP, "destIP", ctx.DestIP) + + // Send initial WebSocket handshake message to Kafka + sourceIP := ExtractIP(ctx.SourceIP) + sourcePort := ExtractPort(ctx.SourceIP) + destIP := ExtractIP(ctx.DestIP) + + // Convert headers to JSON + reqHeadersJSON, _ := json.Marshal(headers.Request.StringMap) + respHeadersJSON, _ := json.Marshal(headers.Response.StringMap) + + handshakePayload := map[string]string{ + "ip": sourceIP, + "destIp": destIP, + "time": fmt.Sprint(time.Now().Unix()), + "akto_account_id": fmt.Sprint(1000000), + "akto_vxlan_id": fmt.Sprint(ctx.VxlanID), + "source": ctx.TrafficSource, + "connection_type": "WEBSOCKET", + "path": req.URL.String(), + "method": req.Method, + "statusCode": fmt.Sprint(resp.StatusCode), + "requestHeaders": string(reqHeadersJSON), + "responseHeaders": string(respHeadersJSON), + } + + out, err := json.Marshal(handshakePayload) + if err == nil { + bgCtx := context.Background() + connectionID := sourceIP + ":" + sourcePort + ProduceStr(bgCtx, string(out), connectionID, sourceIP, "WEBSOCKET_UPGRADE") + } + + continue + } + } + // Get source IP from headers ip := GetSourceIp(headers.Request.Protobuf, ctx.SourceIP) diff --git a/trafficUtil/kafkaUtil/websocket_manager.go b/trafficUtil/kafkaUtil/websocket_manager.go new file mode 100644 index 00000000..cd5f84a0 --- /dev/null +++ b/trafficUtil/kafkaUtil/websocket_manager.go @@ -0,0 +1,321 @@ +package kafkaUtil + +import ( + "fmt" + "log/slog" + "net/http" + "sync" + "time" +) + +// WebSocketConnection stores metadata about an active WebSocket connection. +type WebSocketConnection struct { + SourceIP string + SourcePort string + DestIP string + DestPort string + Headers map[string]string + EstablishedAt time.Time + LastMessageAt time.Time + Messages []WebSocketMessage + incomingAsm WSFragmentAssembler + outgoingAsm WSFragmentAssembler +} + +// WebSocketMessage represents a single WebSocket frame or message. +type WebSocketMessage struct { + EventType string // Custom event type (populated during post-processing, not from eBPF) + Payload string // Message payload + Direction string // "incoming" or "outgoing" + Timestamp time.Time + FIN bool // Final frame flag + Opcode uint8 // WebSocket opcode (1=text, 2=binary, etc.) + Masked bool // Whether payload was masked +} + +// WebSocketBatch represents accumulated messages ready to send to Kafka. +type WebSocketBatch struct { + ConnectionKey string + Connection WebSocketConnection + Messages []WebSocketMessage + BatchTime time.Time +} + +// WebSocketConnectionManager manages active WebSocket connections and message accumulation. +type WebSocketConnectionManager struct { + connections map[string]*WebSocketConnection + mutex sync.RWMutex + cleanupTTL time.Duration +} + +// NewWebSocketConnectionManager creates a new WebSocket connection manager. +func NewWebSocketConnectionManager() *WebSocketConnectionManager { + return &WebSocketConnectionManager{ + connections: make(map[string]*WebSocketConnection), + cleanupTTL: 5 * time.Minute, + } +} + +// buildConnectionKey creates a unique key for a connection. +func buildConnectionKey(sourceIP, sourcePort, destIP, destPort string) string { + return fmt.Sprintf("%s:%s:%s:%s", sourceIP, sourcePort, destIP, destPort) +} + +func buildConnIDKey(connIDId uint64, connIDFd uint32) string { + return fmt.Sprintf("%d:%d", connIDId, connIDFd) +} + +// IsRegisteredByConnID reports whether an eBPF ConnID is tracked as WebSocket. +func (wcm *WebSocketConnectionManager) IsRegisteredByConnID(connIDId uint64, connIDFd uint32) bool { + wcm.mutex.RLock() + defer wcm.mutex.RUnlock() + _, exists := wcm.connections[buildConnIDKey(connIDId, connIDFd)] + return exists +} + +// RegisterConnection registers a new WebSocket connection with initial HTTP handshake headers. +func (wcm *WebSocketConnectionManager) RegisterConnection(sourceIP, sourcePort, destIP, destPort string, req *http.Request) error { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnectionKey(sourceIP, sourcePort, destIP, destPort) + + if _, exists := wcm.connections[key]; exists { + slog.Warn("WebSocket connection already registered", "key", key) + return nil + } + + headers := make(map[string]string) + for name, values := range req.Header { + if len(values) > 0 { + headers[name] = values[0] + } + } + + wcm.connections[key] = &WebSocketConnection{ + SourceIP: sourceIP, + SourcePort: sourcePort, + DestIP: destIP, + DestPort: destPort, + Headers: headers, + EstablishedAt: time.Now(), + LastMessageAt: time.Now(), + Messages: []WebSocketMessage{}, + } + + slog.Debug("WebSocket connection registered", "key", key) + return nil +} + +func (wcm *WebSocketConnectionManager) RegisterConnectionNumeric(connIDId uint64, connIDFd uint32, destIPNumeric uint32, destPort uint16, srcIPNumeric uint32, srcPort uint16, headers map[string]string) error { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnIDKey(connIDId, connIDFd) + + if _, exists := wcm.connections[key]; exists { + slog.Debug("WebSocket connection already registered", "key", key) + return nil + } + + if headers == nil { + headers = make(map[string]string) + } + + wcm.connections[key] = &WebSocketConnection{ + SourceIP: fmt.Sprintf("%d", srcIPNumeric), + SourcePort: fmt.Sprintf("%d", srcPort), + DestIP: fmt.Sprintf("%d", destIPNumeric), + DestPort: fmt.Sprintf("%d", destPort), + Headers: headers, + EstablishedAt: time.Now(), + LastMessageAt: time.Now(), + Messages: []WebSocketMessage{}, + } + + slog.Info("WebSocket connection registered (numeric)", "key", key) + return nil +} + +// AccumulateMessages adds WebSocket messages to the connection's batch. +func (wcm *WebSocketConnectionManager) AccumulateMessages(sourceIP, sourcePort, destIP, destPort string, messages []WebSocketMessage) error { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnectionKey(sourceIP, sourcePort, destIP, destPort) + + conn, exists := wcm.connections[key] + if !exists { + slog.Warn("WebSocket connection not found for accumulation", "key", key) + return fmt.Errorf("connection not found: %s", key) + } + + conn.Messages = append(conn.Messages, messages...) + conn.LastMessageAt = time.Now() + + return nil +} + +func (wcm *WebSocketConnectionManager) AccumulateMessagesNumeric(connIDId uint64, connIDFd uint32, messages []WebSocketMessage) error { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnIDKey(connIDId, connIDFd) + + conn, exists := wcm.connections[key] + if !exists { + slog.Warn("WebSocket connection not found for accumulation", "key", key) + return fmt.Errorf("connection not found: %s", key) + } + + conn.Messages = append(conn.Messages, messages...) + conn.LastMessageAt = time.Now() + + return nil +} + +func (wcm *WebSocketConnectionManager) AccumulatePayloadNumeric(connIDId uint64, connIDFd uint32, payload []byte, direction string) ([]WebSocketMessage, error) { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnIDKey(connIDId, connIDFd) + conn, exists := wcm.connections[key] + if !exists { + slog.Warn("WebSocket connection not found for payload accumulation", "key", key) + return nil, fmt.Errorf("connection not found: %s", key) + } + + var asm *WSFragmentAssembler + if direction == "outgoing" { + asm = &conn.outgoingAsm + } else { + asm = &conn.incomingAsm + } + + msgs := ParseWebSocketFrames(payload, direction, asm) + if len(msgs) > 0 { + conn.Messages = append(conn.Messages, msgs...) + conn.LastMessageAt = time.Now() + slog.Info("WebSocket messages accumulated", + "key", key, "direction", direction, + "newMsgs", len(msgs), "totalPending", len(conn.Messages), + "asmActive", asm.active, "asmBufLen", asm.buf.Len()) + } else { + slog.Info("WebSocket parse yielded no messages", + "key", key, "direction", direction, "payloadLen", len(payload), + "asmActive", asm.active, "asmBufLen", asm.buf.Len()) + } + + return msgs, nil +} + +// GetAndClearBatch retrieves accumulated messages and resets the batch. +func (wcm *WebSocketConnectionManager) GetAndClearBatch(sourceIP, sourcePort, destIP, destPort string) (*WebSocketBatch, error) { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnectionKey(sourceIP, sourcePort, destIP, destPort) + + conn, exists := wcm.connections[key] + if !exists { + return nil, fmt.Errorf("connection not found: %s", key) + } + + batch := &WebSocketBatch{ + ConnectionKey: key, + Connection: *conn, + Messages: conn.Messages, + BatchTime: time.Now(), + } + + conn.Messages = []WebSocketMessage{} + + return batch, nil +} + +func (wcm *WebSocketConnectionManager) GetAndClearBatchByConnID(connIDId uint64, connIDFd uint32) (*WebSocketBatch, error) { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnIDKey(connIDId, connIDFd) + + conn, exists := wcm.connections[key] + if !exists { + return nil, fmt.Errorf("connection not found: %s", key) + } + + batch := &WebSocketBatch{ + ConnectionKey: key, + Connection: *conn, + Messages: conn.Messages, + BatchTime: time.Now(), + } + + conn.Messages = []WebSocketMessage{} + + return batch, nil +} + +// GetAllConnections returns all active WebSocket connections. +func (wcm *WebSocketConnectionManager) GetAllConnections() []string { + wcm.mutex.RLock() + defer wcm.mutex.RUnlock() + + keys := make([]string, 0, len(wcm.connections)) + for k := range wcm.connections { + keys = append(keys, k) + } + return keys +} + +// RemoveConnection removes a connection from tracking. +func (wcm *WebSocketConnectionManager) RemoveConnection(sourceIP, sourcePort, destIP, destPort string) { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnectionKey(sourceIP, sourcePort, destIP, destPort) + delete(wcm.connections, key) + slog.Debug("WebSocket connection removed", "key", key) +} + +func (wcm *WebSocketConnectionManager) RemoveConnectionByConnID(connIDId uint64, connIDFd uint32) { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + key := buildConnIDKey(connIDId, connIDFd) + if _, exists := wcm.connections[key]; !exists { + return + } + delete(wcm.connections, key) + slog.Info("WebSocket connection removed (closed)", "key", key) +} + +// CleanupStaleConnections removes connections that haven't been active within the TTL. +func (wcm *WebSocketConnectionManager) CleanupStaleConnections() { + wcm.mutex.Lock() + defer wcm.mutex.Unlock() + + cutoff := time.Now().Add(-wcm.cleanupTTL) + for key, conn := range wcm.connections { + if conn.LastMessageAt.Before(cutoff) { + delete(wcm.connections, key) + slog.Debug("WebSocket connection cleaned up (stale)", "key", key) + } + } +} + +// Global WebSocket connection manager instance +var WSConnectionManager *WebSocketConnectionManager + +func init() { + WSConnectionManager = NewWebSocketConnectionManager() + + // Start cleanup routine - runs every 5 minutes + go func() { + ticker := time.NewTicker(5 * time.Minute) + defer ticker.Stop() + for range ticker.C { + WSConnectionManager.CleanupStaleConnections() + } + }() +}