KAOS/code/internal/server/ws.go

package server

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

var wsUpgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

// wsResizeMsg is sent from the browser when the terminal size changes.
type wsResizeMsg struct {
	Type string `json:"type"`
	Cols uint16 `json:"cols"`
	Rows uint16 `json:"rows"`
}

// handleConsoleWS handles WebSocket connections for console terminals.
// Each connection gets its own independent shell PTY session.
func (s *Server) handleConsoleWS(c *gin.Context) {
	sessionNum := c.Param("session")
	if sessionNum != "1" && sessionNum != "2" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "sesija mora biti 1 ili 2"})
		return
	}

	conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		return
	}
	defer conn.Close()

	log.Printf("WS[%s]: connected, spawning shell", sessionNum)

	// Spawn a fresh interactive shell for this connection
	ptySess, err := spawnShellPTY(s.projectRoot())
	if err != nil {
		log.Printf("WS[%s]: shell spawn error: %v", sessionNum, err)
		conn.WriteMessage(websocket.TextMessage,
			[]byte(fmt.Sprintf("\r\n\033[31m[Greška: %v]\033[0m\r\n", err)))
		return
	}
	defer ptySess.Close()

	log.Printf("WS[%s]: shell started (PID %d)", sessionNum, ptySess.Cmd.Process.Pid)

	subID := fmt.Sprintf("ws-%d", time.Now().UnixNano())
	outputCh := ptySess.Subscribe(subID)
	defer ptySess.Unsubscribe(subID)

	// Serialized write channel
	writeCh := make(chan []byte, 256)
	writeDone := make(chan struct{})
	go func() {
		defer close(writeDone)
		for data := range writeCh {
			if err := conn.WriteMessage(websocket.BinaryMessage, data); err != nil {
				return
			}
		}
	}()

	// PTY output → WebSocket
	go func() {
		for data := range outputCh {
			select {
			case writeCh <- data:
			default:
			}
		}
	}()

	// Signal to stop goroutines when read pump exits
	stopCh := make(chan struct{})

	// Watch for process exit
	go func() {
		select {
		case <-ptySess.Done():
			log.Printf("WS[%s]: shell exited", sessionNum)
			conn.WriteMessage(websocket.CloseMessage,
				websocket.FormatCloseMessage(websocket.CloseNormalClosure, "done"))
		case <-stopCh:
		}
	}()

	// WebSocket → PTY (read pump)
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				log.Printf("WS[%s]: read error: %v", sessionNum, err)
			}
			break
		}

		// Check for resize message
		var resize wsResizeMsg
		if json.Unmarshal(msg, &resize) == nil && resize.Type == "resize" && resize.Cols > 0 && resize.Rows > 0 {
			if err := ptySess.Resize(resize.Rows, resize.Cols); err != nil {
				log.Printf("WS[%s]: resize error: %v", sessionNum, err)
			}
			continue
		}

		// Regular keyboard input → PTY
		if _, err := ptySess.WriteInput(msg); err != nil {
			log.Printf("WS[%s]: write error: %v", sessionNum, err)
			break
		}
	}

	close(stopCh)
	close(writeCh)
	<-writeDone
	log.Printf("WS[%s]: disconnected", sessionNum)
}