node.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net/url"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var queryScopes = []string{"ai_agents", "external_apis", "internal_users"}

var simulatedNodeNames = []string{
	"payments-db",
	"users-db",
	"analytics-db",
	"medical-records-db",
	"audit-log-db",
}

var queryPool = []string{
	"SELECT name, email FROM users WHERE id = 123",
	"SELECT ssn, name FROM users WHERE id = 456",
	"SELECT * FROM products WHERE category = 'electronics'",
	"SELECT credit_card, billing_address FROM payments WHERE order_id = 789",
	"SELECT title, description FROM articles WHERE published = true",
	"SELECT dob, medical_record FROM patients WHERE patient_id = 101",
	"INSERT INTO audit_log (user_id, action, timestamp) VALUES (42, 'login', NOW())",
	"INSERT INTO users (name, email, ssn) VALUES ('John', 'john@example.com', '123-45-6789')",
	"INSERT INTO payments (credit_card, amount) VALUES ('4111-1111-1111-1111', 99.99)",
	"UPDATE users SET password = 'newpass123' WHERE id = 42",
	"UPDATE patients SET medical_record = 'updated diagnosis' WHERE patient_id = 101",
	"UPDATE orders SET credit_card = '****' WHERE id = 1",
	"DELETE FROM sessions WHERE expired = true",
	"DELETE FROM users WHERE id = 456",
	"DELETE FROM audit_log WHERE user_id = 42",
}

func StartSimulatedNodes(count int, serverURL string) {
	for index := 1; index <= count; index++ {
		var nodeID string
		if index <= len(simulatedNodeNames) {
			nodeID = simulatedNodeNames[index-1]
		} else {
			nodeID = fmt.Sprintf("sidecar-%d", index)
		}
		go runSimulatedNode(nodeID, serverURL, index)
	}
}

func runSimulatedNode(nodeID string, serverURL string, index int) {
	parsedURL, err := url.Parse(serverURL)
	if err != nil {
		log.Printf("node %s: invalid server URL %s: %v", nodeID, serverURL, err)
		return
	}
	if parsedURL.Scheme == "https" {
		parsedURL.Scheme = "wss"
	} else {
		parsedURL.Scheme = "ws"
	}
	parsedURL.Path = "/ws/node"
	wsURL := parsedURL.String()

	for {
		conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
		if err != nil {
			log.Printf("node %s: disconnected, reconnecting in 2s...", nodeID)
			time.Sleep(2 * time.Second)
			continue
		}

		nodeNumber := fmt.Sprintf("node-%d", index)
		registerMsg := WSMessage{
			Type:    "register",
			Payload: mustMarshal(RegisterPayload{NodeID: nodeID, NodeNumber: nodeNumber}),
		}
		if err := conn.WriteJSON(registerMsg); err != nil {
			log.Printf("node %s: disconnected, reconnecting in 2s...", nodeID)
			conn.Close()
			time.Sleep(2 * time.Second)
			continue
		}

		var policiesMu sync.Mutex
		localPolicies := make([]Policy, 0)
		policyUpdating := false
		var connWriteMu sync.Mutex
		var queryPausedMu sync.Mutex
		queryPaused := false

		heartbeatDone := make(chan struct{})
		var disconnectOnce sync.Once
		disconnect := func() {
			disconnectOnce.Do(func() {
				close(heartbeatDone)
				conn.Close()
			})
		}

		go func() {
			ticker := time.NewTicker(5 * time.Second)
			defer ticker.Stop()
			for {
				select {
				case <-heartbeatDone:
					return
				case <-ticker.C:
					heartbeatMsg := WSMessage{
						Type:    "heartbeat",
						Payload: mustMarshal(HeartbeatPayload{NodeID: nodeID}),
					}
					connWriteMu.Lock()
					err := conn.WriteJSON(heartbeatMsg)
					connWriteMu.Unlock()
					if err != nil {
						log.Printf("node %s: heartbeat write failed: %v", nodeID, err)
						disconnect()
						return
					}
				}
			}
		}()

		go func() {
			ticker := time.NewTicker(time.Duration(2000+rand.Intn(1000)) * time.Millisecond)
			defer ticker.Stop()
			for {
				select {
				case <-heartbeatDone:
					return
				case <-ticker.C:
					queryPausedMu.Lock()
					paused := queryPaused
					queryPausedMu.Unlock()
					if paused {
						continue
					}
					policiesMu.Lock()
					if policyUpdating {
						policiesMu.Unlock()
						continue
					}
					query := queryPool[rand.Intn(len(queryPool))]
					policies := make([]Policy, len(localPolicies))
					copy(policies, localPolicies)
					policiesMu.Unlock()

					queryScope := queryScopes[rand.Intn(len(queryScopes))]
					action, matchedPolicy := EvaluateQuery(query, policies, queryScope)
					matchedPolicyName := ""
					if matchedPolicy != nil {
						matchedPolicyName = matchedPolicy.Name
					}

					payload := QueryResultPayload{
						NodeID:        nodeID,
						NodeNumber:    nodeNumber,
						Query:         query,
						Action:        action,
						MatchedPolicy: matchedPolicyName,
						Timestamp:     time.Now().Format(time.RFC3339),
						Scope:         queryScope,
					}
					msg := WSMessage{
						Type:    "query_result",
						Payload: mustMarshal(payload),
					}
					connWriteMu.Lock()
					err := conn.WriteJSON(msg)
					connWriteMu.Unlock()
					if err != nil {
						log.Printf("node %s: query_result write failed: %v", nodeID, err)
						disconnect()
						return
					}
				}
			}
		}()

		readLoop:
		for {
			var msg WSMessage
			if err := conn.ReadJSON(&msg); err != nil {
				log.Printf("node %s: disconnected, reconnecting in 2s...", nodeID)
				disconnect()
				break readLoop
			}
			if msg.Type == "policy_update" {
				var payload PolicyUpdatePayload
				if err := json.Unmarshal(msg.Payload, &payload); err == nil {
					log.Printf("node %s: received policy_update version %d with %d policies", nodeID, payload.Version, len(payload.Policies))
					policiesMu.Lock()
					localPolicies = make([]Policy, len(payload.Policies))
					copy(localPolicies, payload.Policies)
					policyUpdating = true
					policiesMu.Unlock()

					time.Sleep(time.Duration(200+rand.Intn(600)) * time.Millisecond)
					connWriteMu.Lock()
					ackMsg := WSMessage{
						Type:    "policy_ack",
						Payload: mustMarshal(PolicyAckPayload{NodeID: nodeID, Version: payload.Version}),
					}
					err := conn.WriteJSON(ackMsg)
					connWriteMu.Unlock()
					policiesMu.Lock()
					policyUpdating = false
					policiesMu.Unlock()
					if err != nil {
						log.Printf("node %s: FAILED to send policy_ack for version %d: %v", nodeID, payload.Version, err)
						log.Printf("node %s: disconnected, reconnecting in 2s...", nodeID)
						disconnect()
						break readLoop
					}
					log.Printf("node %s: sent policy_ack for version %d", nodeID, payload.Version)
				}
			}
			if msg.Type == "query_pause" {
				var payload QueryPausePayload
				if err := json.Unmarshal(msg.Payload, &payload); err == nil {
					queryPausedMu.Lock()
					queryPaused = payload.Paused
					queryPausedMu.Unlock()
					log.Printf("node %s: query_pause paused=%v", nodeID, payload.Paused)
				}
			}
		}

		time.Sleep(2 * time.Second)
	}
}

func mustMarshal(v interface{}) json.RawMessage {
	data, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	return data
}