codespace/internal/api/shell_handler.go

package api

import (
	"fmt"
	"net/http"
	"sync"
	"time"

	"codespace/internal/model"
	"codespace/internal/service"
	"codespace/internal/shell"
	"codespace/internal/util"

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

type shellHandler struct {
	svc *service.ShellService
}

func shellExitBanner(exit shell.ExitInfo) string {
	if exit.Signal != "" {
		return fmt.Sprintf("\r\n[process exited: signal %s]\r\n", exit.Signal)
	}
	return fmt.Sprintf("\r\n[process exited with code %d]\r\n", exit.Code)
}

func (h *shellHandler) start(c *gin.Context) {
	id := c.Param("id")
	if err := h.svc.Start(id); err != nil {
		writeError(c, err)
		return
	}
	c.Status(http.StatusNoContent)
}

func (h *shellHandler) stop(c *gin.Context) {
	id := c.Param("id")
	if err := h.svc.Stop(id); err != nil {
		writeError(c, err)
		return
	}
	c.Status(http.StatusNoContent)
}

func (h *shellHandler) restart(c *gin.Context) {
	id := c.Param("id")
	if err := h.svc.Restart(id); err != nil {
		writeError(c, err)
		return
	}
	c.Status(http.StatusNoContent)
}

func (h *shellHandler) resize(c *gin.Context) {
	id := c.Param("id")
	var req model.ShellResizeRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		writeBadRequest(c, "invalid json")
		return
	}
	if req.Cols <= 0 || req.Rows <= 0 {
		writeBadRequest(c, "cols and rows must be positive")
		return
	}
	if err := h.svc.Resize(id, req.Cols, req.Rows); err != nil {
		writeError(c, err)
		return
	}
	c.Status(http.StatusNoContent)
}

func (h *shellHandler) status(c *gin.Context) {
	id := c.Param("id")
	status, err := h.svc.Status(id)
	if err != nil {
		writeError(c, err)
		return
	}
	c.JSON(http.StatusOK, model.ShellStatusResponse{
		WorkspaceID: status.WorkspaceID,
		Running:     status.Running,
		PID:         status.PID,
	})
}

func (h *shellHandler) ws(c *gin.Context) {
	id := c.Param("id")

	stdin, err := h.svc.Input(id)
	if err != nil {
		if util.CodeOf(err) == util.CodeNotFound {
			c.JSON(409, gin.H{"error": "shell not running"})
			return
		}
		writeError(c, err)
		return
	}

	sub, err := h.svc.Subscribe(id)
	if err != nil {
		writeError(c, err)
		return
	}

	exit, _ := h.svc.ExitStatus(id)

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

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}

	done := make(chan struct{})
	var closeOnce sync.Once
	closeAll := func() {
		closeOnce.Do(func() {
			sub.Close()
			conn.Close()
			stdin.Close()
			close(done)
		})
	}

	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-done:
				return
			case <-ticker.C:
				if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(5*time.Second)); err != nil {
					closeAll()
					return
				}
			case chunk, ok := <-sub.Output():
				if !ok {
					banner := shellExitBanner(exit)
					conn.WriteMessage(websocket.TextMessage, []byte(banner))
					conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
					closeAll()
					return
				}
				if err := conn.WriteMessage(websocket.TextMessage, chunk); err != nil {
					closeAll()
					return
				}
			}
		}
	}()

	conn.SetReadLimit(1 << 20)
	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})
	for {
		mt, data, err := conn.ReadMessage()
		if err != nil {
			break
		}
		if mt == websocket.TextMessage {
			stdin.Write(data)
		}
	}
	closeAll()
}