// Package jupyter provides WebSocket handlers for Jupyter-related operations package jupyter import ( "encoding/binary" "net/http" "time" "github.com/gorilla/websocket" "github.com/jfraeys/fetch_ml/internal/auth" "github.com/jfraeys/fetch_ml/internal/container" "github.com/jfraeys/fetch_ml/internal/jupyter" "github.com/jfraeys/fetch_ml/internal/logging" ) // Handler provides Jupyter-related WebSocket handlers type Handler struct { logger *logging.Logger jupyterMgr *jupyter.ServiceManager authConfig *auth.Config } // NewHandler creates a new Jupyter handler func NewHandler( logger *logging.Logger, jupyterMgr *jupyter.ServiceManager, authConfig *auth.Config, ) *Handler { return &Handler{ logger: logger, jupyterMgr: jupyterMgr, authConfig: authConfig, } } // Error codes const ( ErrorCodeInvalidRequest = 0x01 ErrorCodeAuthenticationFailed = 0x02 ErrorCodePermissionDenied = 0x03 ErrorCodeResourceNotFound = 0x04 ErrorCodeServiceUnavailable = 0x33 ) // Permissions const ( PermJupyterManage = "jupyter:manage" PermJupyterRead = "jupyter:read" ) // sendErrorPacket sends an error response packet to the client func (h *Handler) sendErrorPacket(conn *websocket.Conn, code byte, message, details string) error { err := map[string]interface{}{ "error": true, "code": code, "message": message, "details": details, } return conn.WriteJSON(err) } // sendSuccessPacket sends a success response packet func (h *Handler) sendSuccessPacket(conn *websocket.Conn, data map[string]interface{}) error { return conn.WriteJSON(data) } // HandleStartJupyter handles starting a Jupyter service // Protocol: [api_key_hash:16][workspace_len:1][workspace:var][config_len:2][config:var] func (h *Handler) HandleStartJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1+2 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "start jupyter payload too short", "") } offset := 16 workspaceLen := int(payload[offset]) offset += 1 if workspaceLen <= 0 || len(payload) < offset+workspaceLen+2 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace length", "") } workspace := string(payload[offset : offset+workspaceLen]) offset += workspaceLen configLen := int(binary.BigEndian.Uint16(payload[offset : offset+2])) offset += 2 if configLen < 0 || len(payload) < offset+configLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid config length", "") } if err := container.ValidateJobName(workspace); err != nil { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace name", err.Error()) } h.logger.Info("starting jupyter service", "workspace", workspace, "user", user.Name) // Start Jupyter service if h.jupyterMgr == nil { return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "") } return h.sendSuccessPacket(conn, map[string]interface{}{ "success": true, "workspace": workspace, "timestamp": time.Now().UTC(), }) } // HandleStopJupyter handles stopping a Jupyter service // Protocol: [api_key_hash:16][service_id_len:1][service_id:var] func (h *Handler) HandleStopJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "stop jupyter payload too short", "") } offset := 16 serviceIDLen := int(payload[offset]) offset += 1 if serviceIDLen <= 0 || len(payload) < offset+serviceIDLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service ID length", "") } serviceID := string(payload[offset : offset+serviceIDLen]) h.logger.Info("stopping jupyter service", "service_id", serviceID, "user", user.Name) if h.jupyterMgr == nil { return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "") } return h.sendSuccessPacket(conn, map[string]interface{}{ "success": true, "service_id": serviceID, "timestamp": time.Now().UTC(), }) } // HandleListJupyter handles listing Jupyter services // Protocol: [api_key_hash:16] func (h *Handler) HandleListJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { h.logger.Info("listing jupyter services", "user", user.Name) if h.jupyterMgr == nil { return h.sendSuccessPacket(conn, map[string]interface{}{ "success": true, "services": []interface{}{}, "count": 0, }) } services := h.jupyterMgr.ListServices() return h.sendSuccessPacket(conn, map[string]interface{}{ "success": true, "services": services, "count": len(services), }) } // HandleListJupyterPackages handles listing packages in a Jupyter service // Protocol: [api_key_hash:16][service_name_len:1][service_name:var] func (h *Handler) HandleListJupyterPackages(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "list packages payload too short", "") } offset := 16 serviceNameLen := int(payload[offset]) offset += 1 if serviceNameLen <= 0 || len(payload) < offset+serviceNameLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service name length", "") } serviceName := string(payload[offset : offset+serviceNameLen]) h.logger.Info("listing jupyter packages", "service", serviceName, "user", user.Name) return h.sendSuccessPacket(conn, map[string]interface{}{ "success": true, "service_name": serviceName, "packages": []interface{}{}, "count": 0, }) } // HandleRemoveJupyter handles removing a Jupyter service // Protocol: [api_key_hash:16][service_id_len:1][service_id:var][purge:1] func (h *Handler) HandleRemoveJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "remove jupyter payload too short", "") } offset := 16 serviceIDLen := int(payload[offset]) offset += 1 if serviceIDLen <= 0 || len(payload) < offset+serviceIDLen+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service ID length", "") } serviceID := string(payload[offset : offset+serviceIDLen]) offset += serviceIDLen purge := payload[offset] != 0 h.logger.Info("removing jupyter service", "service_id", serviceID, "purge", purge, "user", user.Name) if h.jupyterMgr == nil { return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "") } return h.sendSuccessPacket(conn, map[string]interface{}{ "success": true, "service_id": serviceID, "purged": purge, }) } // HandleRestoreJupyter handles restoring a Jupyter workspace // Protocol: [api_key_hash:16][workspace_len:1][workspace:var] func (h *Handler) HandleRestoreJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "restore jupyter payload too short", "") } offset := 16 workspaceLen := int(payload[offset]) offset += 1 if workspaceLen <= 0 || len(payload) < offset+workspaceLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace length", "") } workspace := string(payload[offset : offset+workspaceLen]) h.logger.Info("restoring jupyter workspace", "workspace", workspace, "user", user.Name) if h.jupyterMgr == nil { return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "") } return h.sendSuccessPacket(conn, map[string]interface{}{ "success": true, "workspace": workspace, "restored": true, }) } // HTTP Handlers for REST API // ListServicesHTTP handles HTTP requests for listing Jupyter services func (h *Handler) ListServicesHTTP(w http.ResponseWriter, r *http.Request) { http.Error(w, "Not implemented", http.StatusNotImplemented) } // StartServiceHTTP handles HTTP requests for starting Jupyter service func (h *Handler) StartServiceHTTP(w http.ResponseWriter, r *http.Request) { http.Error(w, "Not implemented", http.StatusNotImplemented) }