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 PTY sessions.
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()

	idx := 0
	if sessionNum == "2" {
		idx = 1
	}

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

	session := s.console.getSession(idx)

	// Wait for PTY session to be available (it gets set when a command is executed)
	session.mu.Lock()
	ptySess := session.ptySess
	session.mu.Unlock()

	if ptySess == nil {
		log.Printf("WS[%s]: no PTY yet, polling...", sessionNum)
		conn.WriteMessage(websocket.TextMessage, []byte("\r\n\033[33m[Čekam pokretanje sesije...]\033[0m\r\n"))

		// Poll for session to start (up to 30s)
		ticker := time.NewTicker(300 * time.Millisecond)
		defer ticker.Stop()
		timeout := time.After(30 * time.Second)
		for {
			select {
			case <-ticker.C:
				session.mu.Lock()
				ptySess = session.ptySess
				session.mu.Unlock()
				if ptySess != nil {
					log.Printf("WS[%s]: PTY found after polling", sessionNum)
					goto connected
				}
			case <-timeout:
				log.Printf("WS[%s]: timeout waiting for PTY", sessionNum)
				conn.WriteMessage(websocket.TextMessage, []byte("\r\n\033[31m[Timeout — sesija nije pokrenuta]\033[0m\r\n"))
				return
			case <-c.Request.Context().Done():
				log.Printf("WS[%s]: client disconnected while polling", sessionNum)
				return
			}
		}
	}

connected:
	log.Printf("WS[%s]: subscribing to PTY", sessionNum)
	subID := fmt.Sprintf("ws-%d", time.Now().UnixNano())
	outputCh := ptySess.Subscribe(subID)
	defer ptySess.Unsubscribe(subID)

	// Send buffered output for reconnect
	buffered := ptySess.GetBuffer()
	log.Printf("WS[%s]: sending buffer (%d bytes)", sessionNum, len(buffered))
	if len(buffered) > 0 {
		conn.WriteMessage(websocket.BinaryMessage, buffered)
	}

	// Check if already done
	select {
	case <-ptySess.Done():
		log.Printf("WS[%s]: PTY already done, sending buffer only", sessionNum)
		conn.WriteMessage(websocket.TextMessage, []byte("\r\n\033[33m[Sesija završena]\033[0m\r\n"))
		return
	default:
	}

	// 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:
			}
		}
	}()

	// Watch for process exit
	go func() {
		<-ptySess.Done()
		log.Printf("WS[%s]: PTY process exited", sessionNum)
		select {
		case writeCh <- []byte("\r\n\033[33m[Sesija završena]\033[0m\r\n"):
		default:
		}
		time.Sleep(500 * time.Millisecond)
		conn.WriteMessage(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "done"))
	}()

	// 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]: PTY write error: %v", sessionNum, err)
			break
		}
	}

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