edubox/agent/websocket.go

package main

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type WSMessage struct {
	Action           string `json:"action"`
	NodeID           string `json:"nodeId,omitempty"`
	Code             string `json:"code,omitempty"`
	InstanceID       string `json:"instanceId,omitempty"`
	Type             string `json:"type,omitempty"`
	Port             int    `json:"port,omitempty"`
	ComposeConfig    string `json:"composeConfig,omitempty"`
	StudentId        string `json:"studentId,omitempty"`
	StudentName      string `json:"studentName,omitempty"`
	Error            string `json:"error,omitempty"`
	TailscaleIP      string `json:"tailscaleIp,omitempty"`
	HeadscaleURL     string `json:"headscaleUrl,omitempty"`
	HeadscaleAuthKey string `json:"headscaleAuthKey,omitempty"`
}

var (
	mainConn   *websocket.Conn
	mainConnMu sync.Mutex
)

// headscale config received from the server during activation.
// These are mutable because activation may happen after the agent starts.
var (
	currentHeadscaleURL     string
	currentHeadscaleAuthKey string
	headscaleConfigMu       sync.Mutex
)

func setHeadscaleConfig(url, authKey string) {
	headscaleConfigMu.Lock()
	currentHeadscaleURL = url
	currentHeadscaleAuthKey = authKey
	headscaleConfigMu.Unlock()
}

func getHeadscaleConfig() (string, string) {
	headscaleConfigMu.Lock()
	defer headscaleConfigMu.Unlock()
	return currentHeadscaleURL, currentHeadscaleAuthKey
}

func sendMessage(msg WSMessage) error {
	mainConnMu.Lock()
	defer mainConnMu.Unlock()
	if mainConn == nil {
		return fmt.Errorf("not connected to server")
	}
	if msg.Action != "heartbeat" {
		log.Printf("sendMessage: sending %+v", msg)
	}
	return mainConn.WriteJSON(msg)
}

// UI notifier system: broadcast activation results to all connected UI clients
type uiNotifier func(msg map[string]interface{})

var (
	uiNotifiers   = make(map[int]uiNotifier)
	uiNotifiersMu sync.Mutex
	uiNotifierID  int
)

func registerUINotifier(fn uiNotifier) int {
	uiNotifiersMu.Lock()
	defer uiNotifiersMu.Unlock()
	id := uiNotifierID
	uiNotifierID++
	uiNotifiers[id] = fn
	log.Printf("registerUINotifier: registered ID %d (total: %d)", id, len(uiNotifiers))
	return id
}

func unregisterUINotifier(id int) {
	uiNotifiersMu.Lock()
	defer uiNotifiersMu.Unlock()
	delete(uiNotifiers, id)
	log.Printf("unregisterUINotifier: removed ID %d (total: %d)", id, len(uiNotifiers))
}

func notifyUI(msg map[string]interface{}) {
	uiNotifiersMu.Lock()
	notifiers := make([]uiNotifier, 0, len(uiNotifiers))
	for _, fn := range uiNotifiers {
		notifiers = append(notifiers, fn)
	}
	uiNotifiersMu.Unlock()

	log.Printf("notifyUI: broadcasting to %d UI clients", len(notifiers))
	for _, fn := range notifiers {
		go fn(msg)
	}
}

func startWebSocket(serverAddr, nodeID, dataDir, headscaleURL, headscaleAuthKey string) {
	setHeadscaleConfig(headscaleURL, headscaleAuthKey)

	for {
		conn, _, err := websocket.DefaultDialer.Dial(serverAddr, nil)
		if err != nil {
			log.Printf("WS connect error: %v, retrying in 5s...", err)
			time.Sleep(5 * time.Second)
			continue
		}

		log.Printf("WS connected to %s", serverAddr)

		mainConnMu.Lock()
		mainConn = conn
		mainConnMu.Unlock()

		// Register
		if err := conn.WriteJSON(WSMessage{Action: "register", NodeID: nodeID}); err != nil {
			log.Printf("WS register error: %v", err)
			conn.Close()
			mainConnMu.Lock()
			mainConn = nil
			mainConnMu.Unlock()
			continue
		}

		// Activation flow
		act, err := loadActivation(dataDir)
		if err != nil || !act.Activated {
			log.Println("Waiting for activation...")
		} else {
			log.Printf("Already activated as %s", act.StudentName)
			// If already activated and we have credentials, ensure VPN is up.
			hsURL, hsKey := getHeadscaleConfig()
			if hsURL != "" && hsKey != "" {
				go startTailscaleAndReport(dataDir, nodeID, hsURL, hsKey)
			}
		}

		// Heartbeat goroutine
		done := make(chan struct{})
		go func() {
			ticker := time.NewTicker(10 * time.Second)
			defer ticker.Stop()
			for {
				select {
				case <-ticker.C:
					if err := sendMessage(WSMessage{Action: "heartbeat", NodeID: nodeID}); err != nil {
						return
					}
				case <-done:
					return
				}
			}
		}()

		// Read loop
		for {
			var msg WSMessage
			if err := conn.ReadJSON(&msg); err != nil {
				log.Printf("WS read error: %v", err)
				break
			}
			log.Printf("WS received from server: action=%s", msg.Action)
			handleMessage(conn, msg, dataDir, nodeID)
		}

		close(done)
		conn.Close()
		mainConnMu.Lock()
		mainConn = nil
		mainConnMu.Unlock()
		log.Println("WS disconnected, reconnecting in 5s...")
		time.Sleep(5 * time.Second)
	}
}

func handleMessage(conn *websocket.Conn, msg WSMessage, dataDir, nodeID string) {
	switch msg.Action {
	case "activated":
		log.Printf("handleMessage: activated received, student=%s", msg.StudentName)
		if msg.StudentName != "" {
			act := &Activation{Activated: true, StudentId: msg.StudentId, StudentName: msg.StudentName, Code: msg.Code}
			if err := saveActivation(dataDir, act); err != nil {
				log.Printf("saveActivation error: %v", err)
			} else {
				log.Printf("Activated as %s", act.StudentName)
			}
		}

		// The server also sends Headscale credentials on activation.
		if msg.HeadscaleURL != "" && msg.HeadscaleAuthKey != "" {
			setHeadscaleConfig(msg.HeadscaleURL, msg.HeadscaleAuthKey)
			cfg, _, err := loadOrCreateConfig(dataDir)
			if err != nil {
				log.Printf("loadOrCreateConfig error: %v", err)
			} else {
				cfg.HeadscaleURL = msg.HeadscaleURL
				cfg.HeadscaleAuthKey = msg.HeadscaleAuthKey
				if err := saveConfig(dataDir, cfg); err != nil {
					log.Printf("saveConfig error: %v", err)
				} else {
					log.Printf("Saved Headscale config received from server")
				}
			}
			go startTailscaleAndReport(dataDir, nodeID, msg.HeadscaleURL, msg.HeadscaleAuthKey)
		}

		notifyUI(map[string]interface{}{
			"action":      "activated",
			"studentName": msg.StudentName,
		})
	case "registered":
		// Server acknowledged our register message; nothing to do.
		return
	case "start_vpn":
		log.Printf("Server requested VPN start")
		hsURL, hsKey := getHeadscaleConfig()
		if hsURL == "" || hsKey == "" {
			log.Printf("Cannot start VPN: headscale config missing")
			sendMessage(WSMessage{Action: "vpn_error", NodeID: nodeID, Error: "headscale config missing"})
			return
		}
		go func() {
			ip, err := startTailscale(dataDir, nodeID, hsURL, hsKey)
			if err != nil {
				log.Printf("start_vpn error: %v", err)
				sendMessage(WSMessage{Action: "vpn_error", NodeID: nodeID, Error: err.Error()})
				return
			}
			for {
				if err := sendMessage(WSMessage{Action: "tailscale_ip", NodeID: nodeID, TailscaleIP: ip}); err != nil {
					log.Printf("Waiting for WebSocket to send tailscale_ip...")
					time.Sleep(1 * time.Second)
					continue
				}
				log.Printf("Sent tailscale_ip to server: %s", ip)
				break
			}
		}()
	case "stop_vpn":
		log.Printf("Server requested VPN stop")
		stopTailscale()
		sendMessage(WSMessage{Action: "vpn_stopped", NodeID: nodeID})
	case "activation_failed":
		log.Printf("handleMessage: activation_failed received, error=%s", msg.Error)
		notifyUI(map[string]interface{}{
			"action": "activation_failed",
			"error":  msg.Error,
		})
	case "start":
		log.Printf("Start instance %s on port %d", msg.InstanceID, msg.Port)
		if err := upsertInstance(dataDir, &InstanceInfo{
			ID:           msg.InstanceID,
			TemplateName: msg.Type,
			Port:         msg.Port,
			Status:       "starting",
		}); err != nil {
			log.Printf("upsertInstance error: %v", err)
		}
		if err := writeCompose(dataDir, msg.InstanceID, msg.ComposeConfig, msg.Port); err != nil {
			log.Printf("writeCompose error: %v", err)
			_ = upsertInstance(dataDir, &InstanceInfo{ID: msg.InstanceID, TemplateName: msg.Type, Port: msg.Port, Status: "error"})
			sendMessage(WSMessage{Action: "instance_error", InstanceID: msg.InstanceID, Error: err.Error()})
			return
		}
		if err := dockerComposeUp(dataDir, msg.InstanceID); err != nil {
			log.Printf("dockerComposeUp error: %v", err)
			_ = upsertInstance(dataDir, &InstanceInfo{ID: msg.InstanceID, TemplateName: msg.Type, Port: msg.Port, Status: "error"})
			sendMessage(WSMessage{Action: "instance_error", InstanceID: msg.InstanceID, Error: err.Error()})
			return
		}
		// Repair older WordPress instances: remove hardcoded WP_HOME/WP_SITEURL
		// so the studioE5 mu-plugin can compute the public URL from the Host header.
		go func() {
			// Give the container a moment to be ready before touching wp-config.php
			time.Sleep(2 * time.Second)
			if err := stripWordPressHardcodedURLs(dataDir, msg.InstanceID); err != nil {
				log.Printf("stripWordPressHardcodedURLs error: %v", err)
			}
		}()
		// Ensure Tailscale is running so the server can reach the node
		go ensureTailscale(dataDir, nodeID, msg.Port)

		status := getInstanceStatus(dataDir, msg.InstanceID)
		_ = upsertInstance(dataDir, &InstanceInfo{ID: msg.InstanceID, TemplateName: msg.Type, Port: msg.Port, Status: status})
		sendMessage(WSMessage{Action: "instance_started", InstanceID: msg.InstanceID, Port: msg.Port})
		notifyUI(map[string]interface{}{"action": "instances_updated"})
	case "stop":
		log.Printf("Stop instance %s", msg.InstanceID)
		if err := dockerComposeDown(dataDir, msg.InstanceID); err != nil {
			log.Printf("dockerComposeDown error: %v", err)
		}
		if inst, _ := loadInstances(dataDir); inst[msg.InstanceID] != nil {
			inst[msg.InstanceID].Status = "stopped"
			_ = saveInstances(dataDir, inst)
		}
		notifyUI(map[string]interface{}{"action": "instances_updated"})
	case "delete":
		log.Printf("Delete instance %s", msg.InstanceID)
		dockerComposeRm(dataDir, msg.InstanceID)
		removeInstance(dataDir, msg.InstanceID)
		notifyUI(map[string]interface{}{"action": "instances_updated"})
	case "reset":
		log.Printf("Reset instance %s", msg.InstanceID)
		dockerComposeRm(dataDir, msg.InstanceID)
		if err := writeCompose(dataDir, msg.InstanceID, msg.ComposeConfig, msg.Port); err != nil {
			log.Printf("writeCompose error: %v", err)
			_ = upsertInstance(dataDir, &InstanceInfo{ID: msg.InstanceID, TemplateName: msg.Type, Port: msg.Port, Status: "error"})
			sendMessage(WSMessage{Action: "instance_error", InstanceID: msg.InstanceID, Error: err.Error()})
			return
		}
		if err := dockerComposeUp(dataDir, msg.InstanceID); err != nil {
			log.Printf("dockerComposeUp error: %v", err)
			_ = upsertInstance(dataDir, &InstanceInfo{ID: msg.InstanceID, TemplateName: msg.Type, Port: msg.Port, Status: "error"})
			sendMessage(WSMessage{Action: "instance_error", InstanceID: msg.InstanceID, Error: err.Error()})
			return
		}
		// Repair older WordPress instances: remove hardcoded WP_HOME/WP_SITEURL
		// so the studioE5 mu-plugin can compute the public URL from the Host header.
		go func() {
			// Give the container a moment to be ready before touching wp-config.php
			time.Sleep(2 * time.Second)
			if err := stripWordPressHardcodedURLs(dataDir, msg.InstanceID); err != nil {
				log.Printf("stripWordPressHardcodedURLs error: %v", err)
			}
		}()
		// Ensure Tailscale is running so the server can reach the node
		go ensureTailscale(dataDir, nodeID, msg.Port)

		status := getInstanceStatus(dataDir, msg.InstanceID)
		_ = upsertInstance(dataDir, &InstanceInfo{ID: msg.InstanceID, TemplateName: msg.Type, Port: msg.Port, Status: status})
		sendMessage(WSMessage{Action: "instance_started", InstanceID: msg.InstanceID, Port: msg.Port})
		notifyUI(map[string]interface{}{"action": "instances_updated"})
	default:
		log.Printf("Unknown action: %s", msg.Action)
	}
}

func ensureTailscale(dataDir, nodeID string, port int) {
	hsURL, hsKey := getHeadscaleConfig()
	if hsURL == "" || hsKey == "" {
		log.Printf("Cannot ensure Tailscale: headscale config missing")
		return
	}
	if isTailscaleRunning() {
		return
	}
	log.Printf("Tailscale not running, starting it for instance port %d", port)
	ip, err := startTailscale(dataDir, nodeID, hsURL, hsKey)
	if err != nil {
		log.Printf("ensureTailscale start error: %v", err)
		return
	}
	for {
		if err := sendMessage(WSMessage{Action: "tailscale_ip", NodeID: nodeID, TailscaleIP: ip}); err != nil {
			log.Printf("Waiting for WebSocket to send tailscale_ip...")
			time.Sleep(1 * time.Second)
			continue
		}
		log.Printf("Sent tailscale_ip to server: %s", ip)
		break
	}
}