codespace/internal/api/shell_handler.go

package api

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

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

	"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 "\r\n[process exited: signal " + exit.Signal + "]\r\n"
	}
	return "\r\n[process exited with code " + strconv.Itoa(exit.Code) + "]\r\n"
}

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

func (h *shellHandler) stop(c *gin.Context) {
	id := c.Param("id")
	var req model.ShellStopRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		writeBadRequest(c, "invalid json")
		return
	}
	if req.ShellID == "" {
		writeBadRequest(c, "shellId is required")
		return
	}
	if err := h.svc.Stop(id, req.ShellID); err != nil {
		writeError(c, err)
		return
	}
	c.Status(http.StatusNoContent)
}

func (h *shellHandler) restart(c *gin.Context) {
	id := c.Param("id")
	var req model.ShellRestartRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		writeBadRequest(c, "invalid json")
		return
	}
	if req.ShellID == "" {
		writeBadRequest(c, "shellId is required")
		return
	}
	newShellID, err := h.svc.Restart(id, req.ShellID)
	if err != nil {
		writeError(c, err)
		return
	}
	status, err := h.svc.Status(id, newShellID)
	if err != nil {
		writeError(c, err)
		return
	}
	c.JSON(http.StatusCreated, model.ShellStartResponse{
		WorkspaceID: status.WorkspaceID,
		ShellID:     status.ShellID,
		Running:     status.Running,
		PID:         status.PID,
	})
}

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.ShellID == "" {
		writeBadRequest(c, "shellId is required")
		return
	}
	if req.Cols <= 0 || req.Rows <= 0 {
		writeBadRequest(c, "cols and rows must be positive")
		return
	}
	if err := h.svc.Resize(id, req.ShellID, 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")
	shellID := c.Query("shellId")
	if shellID == "" {
		writeBadRequest(c, "shellId query parameter is required")
		return
	}
	status, err := h.svc.Status(id, shellID)
	if err != nil {
		writeError(c, err)
		return
	}
	c.JSON(http.StatusOK, model.ShellStatusResponse{
		WorkspaceID: status.WorkspaceID,
		ShellID:     status.ShellID,
		Running:     status.Running,
		PID:         status.PID,
	})
}

func (h *shellHandler) list(c *gin.Context) {
	id := c.Param("id")
	shells, err := h.svc.List(id)
	if err != nil {
		writeError(c, err)
		return
	}

	infos := make([]model.ShellInfo, 0, len(shells))
	for _, sh := range shells {
		infos = append(infos, model.ShellInfo{
			WorkspaceID: sh.WorkspaceID,
			ShellID:     sh.ShellID,
			Running:     sh.Running,
			PID:         sh.PID,
			CreatedAt:   sh.CreatedAt.Format(time.RFC3339Nano),
		})
	}

	c.JSON(http.StatusOK, model.ShellListResponse{
		WorkspaceID: id,
		Shells:      infos,
	})
}

func (h *shellHandler) ws(c *gin.Context) {
	id := c.Param("id")
	shellID := c.Query("shellId")
	if shellID == "" {
		writeBadRequest(c, "shellId query parameter is required")
		return
	}

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

	stdin, err := h.svc.Input(id, shellID)
	if err != nil {
		writeError(c, err)
		return
	}

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

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