firecracker-orchestrator/orchestrator/serve.go

package orchestrator

import (
	"encoding/json"
	"fmt"
	_ "embed"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/gorilla/websocket"
	log "github.com/sirupsen/logrus"
)

//go:embed web/terminal.html
var terminalHTML []byte

var upgrader = websocket.Upgrader{
	// Allow all origins for local/dev use. Add an origin check for production.
	CheckOrigin: func(r *http.Request) bool { return true },
}

// Serve starts an HTTP server that exposes:
//
//	GET    /              — terminal UI (xterm.js); ?id=N selects a clone
//	GET    /clones        — JSON list of running clone IDs
//	POST   /clones        — spawn a new clone; returns {"id": N}
//	DELETE /clones/{id}   — destroy clone {id}
//	GET    /ws/{id}       — WebSocket console for clone {id}
//
// Binary WebSocket frames carry raw terminal bytes (stdin/stdout).
// Text WebSocket frames carry JSON resize commands: {"rows":N,"cols":M}.
func Serve(orch *Orchestrator, addr string) error {
	mux := http.NewServeMux()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Write(terminalHTML) //nolint:errcheck
	})

	// /clones — list (GET) or spawn (POST)
	mux.HandleFunc("/clones", func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet, "":
			ids := runningCloneIDs(orch.cfg)
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(ids) //nolint:errcheck
		case http.MethodPost:
			id, err := orch.SpawnSingle()
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusCreated)
			json.NewEncoder(w).Encode(map[string]int{"id": id}) //nolint:errcheck
		default:
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		}
	})

	// /clones/{id} — destroy (DELETE)
	mux.HandleFunc("/clones/", func(w http.ResponseWriter, r *http.Request) {
		idStr := strings.TrimPrefix(r.URL.Path, "/clones/")
		if idStr == "" {
			// redirect bare /clones/ to /clones
			http.Redirect(w, r, "/clones", http.StatusMovedPermanently)
			return
		}
		if r.Method != http.MethodDelete {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		id, err := strconv.Atoi(idStr)
		if err != nil {
			http.Error(w, "invalid clone id", http.StatusBadRequest)
			return
		}
		if err := orch.KillClone(id); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusNoContent)
	})

	mux.HandleFunc("/ws/", func(w http.ResponseWriter, r *http.Request) {
		idStr := strings.TrimPrefix(r.URL.Path, "/ws/")
		id, err := strconv.Atoi(idStr)
		if err != nil {
			http.Error(w, "invalid clone id", http.StatusBadRequest)
			return
		}
		handleConsoleWS(orch.cfg, id, w, r)
	})

	log.Infof("terminal UI: http://%s", addr)
	return http.ListenAndServe(addr, mux)
}

// handleConsoleWS upgrades the request to a WebSocket and bridges it to the
// clone's console.sock (I/O) and console-resize.sock (terminal resize).
func handleConsoleWS(cfg Config, id int, w http.ResponseWriter, r *http.Request) {
	cloneDir := filepath.Join(cfg.BaseDir, "clones", strconv.Itoa(id))
	consoleSock := filepath.Join(cloneDir, "console.sock")
	resizeSock := filepath.Join(cloneDir, "console-resize.sock")

	if _, err := os.Stat(consoleSock); err != nil {
		http.Error(w, fmt.Sprintf("clone %d is not running", id), http.StatusNotFound)
		return
	}

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	defer ws.Close()

	consoleConn, err := net.Dial("unix", consoleSock)
	if err != nil {
		writeWSError(ws, fmt.Sprintf("console unavailable: %v", err))
		return
	}
	defer consoleConn.Close()

	resizeConn, err := net.Dial("unix", resizeSock)
	if err != nil {
		log.Warnf("clone %d: could not connect to resize socket, resize disabled: %v", id, err)
		resizeConn = nil
	}
	if resizeConn != nil {
		defer resizeConn.Close()
	}

	log.Infof("ws: clone %d: client connected from %s", id, r.RemoteAddr)
	bridgeWS(ws, consoleConn, resizeConn)
	log.Infof("ws: clone %d: client disconnected", id)
}

// bridgeWS proxies between a WebSocket connection and a console Unix socket.
//
//   - Binary WS frames  → consoleConn (terminal input)
//   - Text WS frames    → resizeConn as a JSON line (resize command)
//   - consoleConn reads → binary WS frames (terminal output)
func bridgeWS(ws *websocket.Conn, consoleConn net.Conn, resizeConn net.Conn) {
	// console → WebSocket
	sockDone := make(chan struct{})
	go func() {
		defer close(sockDone)
		buf := make([]byte, 4096)
		for {
			n, err := consoleConn.Read(buf)
			if n > 0 {
				if werr := ws.WriteMessage(websocket.BinaryMessage, buf[:n]); werr != nil {
					return
				}
			}
			if err != nil {
				return
			}
		}
	}()

	// WebSocket → console / resize
	for {
		msgType, data, err := ws.ReadMessage()
		if err != nil {
			break
		}
		switch msgType {
		case websocket.BinaryMessage:
			consoleConn.Write(data) //nolint:errcheck
		case websocket.TextMessage:
			if resizeConn != nil {
				// Append newline so the scanner in handleResize sees a complete line.
				resizeConn.Write(append(data, '\n')) //nolint:errcheck
			}
		}
	}

	consoleConn.Close()
	<-sockDone
}

// runningCloneIDs returns clone IDs that have a live console socket.
func runningCloneIDs(cfg Config) []int {
	clonesDir := filepath.Join(cfg.BaseDir, "clones")
	entries, err := os.ReadDir(clonesDir)
	if err != nil {
		return nil
	}
	var ids []int
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		id, err := strconv.Atoi(e.Name())
		if err != nil {
			continue
		}
		sock := filepath.Join(clonesDir, e.Name(), "console.sock")
		if _, err := os.Stat(sock); err == nil {
			ids = append(ids, id)
		}
	}
	return ids
}

func writeWSError(ws *websocket.Conn, msg string) {
	ws.WriteMessage(websocket.TextMessage, //nolint:errcheck
		[]byte("\r\n\x1b[31m["+msg+"]\x1b[0m\r\n"))
}