10e6416e11
- Update datasets handlers with improved error handling - Refactor WebSocket handler for better organization - Clean up jobs.go handler implementation - Add websocket_metrics table to Postgres and SQLite schemas
540 lines
17 KiB
Go
540 lines
17 KiB
Go
// Package ws provides WebSocket handling for the API
|
|
package ws
|
|
|
|
import (
|
|
"context"
|
|
"encoding/binary"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"github.com/gorilla/websocket"
|
|
"github.com/jfraeys/fetch_ml/internal/audit"
|
|
"github.com/jfraeys/fetch_ml/internal/auth"
|
|
"github.com/jfraeys/fetch_ml/internal/config"
|
|
"github.com/jfraeys/fetch_ml/internal/experiment"
|
|
"github.com/jfraeys/fetch_ml/internal/jupyter"
|
|
"github.com/jfraeys/fetch_ml/internal/logging"
|
|
"github.com/jfraeys/fetch_ml/internal/queue"
|
|
"github.com/jfraeys/fetch_ml/internal/storage"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/api/datasets"
|
|
"github.com/jfraeys/fetch_ml/internal/api/jobs"
|
|
jupyterj "github.com/jfraeys/fetch_ml/internal/api/jupyter"
|
|
)
|
|
|
|
// Response packet types (duplicated from api package to avoid import cycle)
|
|
const (
|
|
PacketTypeSuccess = 0x00
|
|
PacketTypeError = 0x01
|
|
PacketTypeProgress = 0x02
|
|
PacketTypeStatus = 0x03
|
|
PacketTypeData = 0x04
|
|
PacketTypeLog = 0x05
|
|
)
|
|
|
|
// Opcodes for binary WebSocket protocol
|
|
const (
|
|
OpcodeQueueJob = 0x01
|
|
OpcodeStatusRequest = 0x02
|
|
OpcodeCancelJob = 0x03
|
|
OpcodePrune = 0x04
|
|
OpcodeDatasetList = 0x06
|
|
OpcodeDatasetRegister = 0x07
|
|
OpcodeDatasetInfo = 0x08
|
|
OpcodeDatasetSearch = 0x09
|
|
OpcodeLogMetric = 0x0A
|
|
OpcodeGetExperiment = 0x0B
|
|
OpcodeQueueJobWithTracking = 0x0C
|
|
OpcodeQueueJobWithSnapshot = 0x17
|
|
OpcodeQueueJobWithArgs = 0x1A
|
|
OpcodeQueueJobWithNote = 0x1B
|
|
OpcodeAnnotateRun = 0x1C
|
|
OpcodeSetRunNarrative = 0x1D
|
|
OpcodeStartJupyter = 0x0D
|
|
OpcodeStopJupyter = 0x0E
|
|
OpcodeRemoveJupyter = 0x18
|
|
OpcodeRestoreJupyter = 0x19
|
|
OpcodeListJupyter = 0x0F
|
|
OpcodeListJupyterPackages = 0x1E
|
|
OpcodeValidateRequest = 0x16
|
|
|
|
// Logs opcodes
|
|
OpcodeGetLogs = 0x20
|
|
OpcodeStreamLogs = 0x21
|
|
)
|
|
|
|
// Error codes
|
|
const (
|
|
ErrorCodeUnknownError = 0x00
|
|
ErrorCodeInvalidRequest = 0x01
|
|
ErrorCodeAuthenticationFailed = 0x02
|
|
ErrorCodePermissionDenied = 0x03
|
|
ErrorCodeResourceNotFound = 0x04
|
|
ErrorCodeResourceAlreadyExists = 0x05
|
|
ErrorCodeServerOverloaded = 0x10
|
|
ErrorCodeDatabaseError = 0x11
|
|
ErrorCodeNetworkError = 0x12
|
|
ErrorCodeStorageError = 0x13
|
|
ErrorCodeTimeout = 0x14
|
|
ErrorCodeJobNotFound = 0x20
|
|
ErrorCodeJobAlreadyRunning = 0x21
|
|
ErrorCodeJobFailedToStart = 0x22
|
|
ErrorCodeJobExecutionFailed = 0x23
|
|
ErrorCodeJobCancelled = 0x24
|
|
ErrorCodeOutOfMemory = 0x30
|
|
ErrorCodeDiskFull = 0x31
|
|
ErrorCodeInvalidConfiguration = 0x32
|
|
ErrorCodeServiceUnavailable = 0x33
|
|
)
|
|
|
|
// Permissions
|
|
const (
|
|
PermJobsCreate = "jobs:create"
|
|
PermJobsRead = "jobs:read"
|
|
PermJobsUpdate = "jobs:update"
|
|
PermDatasetsRead = "datasets:read"
|
|
PermDatasetsCreate = "datasets:create"
|
|
PermJupyterManage = "jupyter:manage"
|
|
PermJupyterRead = "jupyter:read"
|
|
)
|
|
|
|
// Handler provides WebSocket handling
|
|
type Handler struct {
|
|
authConfig *auth.Config
|
|
logger *logging.Logger
|
|
expManager *experiment.Manager
|
|
dataDir string
|
|
taskQueue queue.Backend
|
|
db *storage.DB
|
|
jupyterServiceMgr *jupyter.ServiceManager
|
|
securityCfg *config.SecurityConfig
|
|
auditLogger *audit.Logger
|
|
upgrader websocket.Upgrader
|
|
jobsHandler *jobs.Handler
|
|
jupyterHandler *jupyterj.Handler
|
|
datasetsHandler *datasets.Handler
|
|
}
|
|
|
|
// NewHandler creates a new WebSocket handler
|
|
func NewHandler(
|
|
authConfig *auth.Config,
|
|
logger *logging.Logger,
|
|
expManager *experiment.Manager,
|
|
dataDir string,
|
|
taskQueue queue.Backend,
|
|
db *storage.DB,
|
|
jupyterServiceMgr *jupyter.ServiceManager,
|
|
securityCfg *config.SecurityConfig,
|
|
auditLogger *audit.Logger,
|
|
jobsHandler *jobs.Handler,
|
|
jupyterHandler *jupyterj.Handler,
|
|
datasetsHandler *datasets.Handler,
|
|
) *Handler {
|
|
upgrader := createUpgrader(securityCfg)
|
|
|
|
return &Handler{
|
|
authConfig: authConfig,
|
|
logger: logger,
|
|
expManager: expManager,
|
|
dataDir: dataDir,
|
|
taskQueue: taskQueue,
|
|
db: db,
|
|
jupyterServiceMgr: jupyterServiceMgr,
|
|
securityCfg: securityCfg,
|
|
auditLogger: auditLogger,
|
|
upgrader: upgrader,
|
|
jobsHandler: jobsHandler,
|
|
jupyterHandler: jupyterHandler,
|
|
datasetsHandler: datasetsHandler,
|
|
}
|
|
}
|
|
|
|
// createUpgrader creates a WebSocket upgrader with the given security configuration
|
|
func createUpgrader(securityCfg *config.SecurityConfig) websocket.Upgrader {
|
|
return websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
origin := r.Header.Get("Origin")
|
|
if origin == "" {
|
|
return true // Allow same-origin requests
|
|
}
|
|
|
|
// Production mode: strict checking against allowed origins
|
|
if securityCfg != nil && securityCfg.ProductionMode {
|
|
for _, allowed := range securityCfg.AllowedOrigins {
|
|
if origin == allowed {
|
|
return true
|
|
}
|
|
}
|
|
return false // Reject if not in allowed list
|
|
}
|
|
|
|
// Development mode: allow localhost and local network origins
|
|
parsedOrigin, err := url.Parse(origin)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
|
|
host := parsedOrigin.Host
|
|
if strings.HasPrefix(host, "localhost:") ||
|
|
strings.HasPrefix(host, "127.0.0.1:") ||
|
|
strings.HasPrefix(host, "192.168.") ||
|
|
strings.HasPrefix(host, "10.") ||
|
|
strings.HasPrefix(host, "[::1]:") {
|
|
return true
|
|
}
|
|
|
|
return false
|
|
},
|
|
EnableCompression: true,
|
|
}
|
|
}
|
|
|
|
// ServeHTTP implements http.Handler for WebSocket upgrade
|
|
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := h.upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
h.logger.Error("websocket upgrade failed", "error", err)
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
|
|
h.handleConnection(conn)
|
|
}
|
|
|
|
// handleConnection handles an established WebSocket connection
|
|
func (h *Handler) handleConnection(conn *websocket.Conn) {
|
|
h.logger.Info("websocket connection established", "remote", conn.RemoteAddr())
|
|
|
|
for {
|
|
messageType, payload, err := conn.ReadMessage()
|
|
if err != nil {
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
|
h.logger.Error("websocket read error", "error", err)
|
|
}
|
|
break
|
|
}
|
|
|
|
if messageType != websocket.BinaryMessage {
|
|
h.logger.Warn("received non-binary message, ignoring")
|
|
continue
|
|
}
|
|
|
|
if err := h.handleMessage(conn, payload); err != nil {
|
|
h.logger.Error("message handling error", "error", err)
|
|
// Don't break, continue handling messages
|
|
}
|
|
}
|
|
|
|
h.logger.Info("websocket connection closed", "remote", conn.RemoteAddr())
|
|
}
|
|
|
|
// handleMessage dispatches WebSocket messages to appropriate handlers
|
|
func (h *Handler) handleMessage(conn *websocket.Conn, payload []byte) error {
|
|
if len(payload) < 17 { // At least opcode + api_key_hash
|
|
return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "payload too short", "")
|
|
}
|
|
|
|
opcode := payload[0] // First byte is opcode, followed by 16-byte API key hash
|
|
|
|
switch opcode {
|
|
case OpcodeAnnotateRun:
|
|
return h.handleAnnotateRun(conn, payload)
|
|
case OpcodeSetRunNarrative:
|
|
return h.handleSetRunNarrative(conn, payload)
|
|
case OpcodeStartJupyter:
|
|
return h.handleStartJupyter(conn, payload)
|
|
case OpcodeStopJupyter:
|
|
return h.handleStopJupyter(conn, payload)
|
|
case OpcodeListJupyter:
|
|
return h.handleListJupyter(conn, payload)
|
|
case OpcodeQueueJob:
|
|
return h.handleQueueJob(conn, payload)
|
|
case OpcodeQueueJobWithSnapshot:
|
|
return h.handleQueueJobWithSnapshot(conn, payload)
|
|
case OpcodeStatusRequest:
|
|
return h.handleStatusRequest(conn, payload)
|
|
case OpcodeCancelJob:
|
|
return h.handleCancelJob(conn, payload)
|
|
case OpcodePrune:
|
|
return h.handlePrune(conn, payload)
|
|
case OpcodeValidateRequest:
|
|
return h.handleValidateRequest(conn, payload)
|
|
case OpcodeLogMetric:
|
|
return h.handleLogMetric(conn, payload)
|
|
case OpcodeGetExperiment:
|
|
return h.handleGetExperiment(conn, payload)
|
|
case OpcodeDatasetList:
|
|
return h.handleDatasetList(conn, payload)
|
|
case OpcodeDatasetRegister:
|
|
return h.handleDatasetRegister(conn, payload)
|
|
case OpcodeDatasetInfo:
|
|
return h.handleDatasetInfo(conn, payload)
|
|
case OpcodeDatasetSearch:
|
|
return h.handleDatasetSearch(conn, payload)
|
|
default:
|
|
return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "unknown opcode", string(opcode))
|
|
}
|
|
}
|
|
|
|
// sendPacket builds and sends a binary packet with type and sections
|
|
func (h *Handler) sendPacket(conn *websocket.Conn, pktType byte, sections ...[]byte) error {
|
|
var buf []byte
|
|
buf = append(buf, pktType, 0, 0, 0, 0, 0, 0, 0, 0) // Type + timestamp placeholder
|
|
for _, section := range sections {
|
|
var tmp [10]byte
|
|
n := binary.PutUvarint(tmp[:], uint64(len(section)))
|
|
buf = append(buf, tmp[:n]...)
|
|
buf = append(buf, section...)
|
|
}
|
|
return conn.WriteMessage(websocket.BinaryMessage, buf)
|
|
}
|
|
|
|
func (h *Handler) sendErrorPacket(conn *websocket.Conn, code byte, message, details string) error {
|
|
return h.sendPacket(conn, PacketTypeError, []byte{code}, []byte(message), []byte(details))
|
|
}
|
|
|
|
func (h *Handler) sendSuccessPacket(conn *websocket.Conn, data map[string]any) error {
|
|
payload, _ := json.Marshal(data)
|
|
return h.sendPacket(conn, PacketTypeSuccess, payload)
|
|
}
|
|
|
|
func (h *Handler) sendDataPacket(conn *websocket.Conn, dataType string, payload []byte) error {
|
|
return h.sendPacket(conn, PacketTypeData, []byte(dataType), payload)
|
|
}
|
|
|
|
// Handler stubs - delegate to sub-packages
|
|
|
|
func (h *Handler) withAuth(conn *websocket.Conn, payload []byte, handler func(*auth.User) error) error {
|
|
user, err := h.Authenticate(payload)
|
|
if err != nil {
|
|
return h.sendErrorPacket(conn, ErrorCodeAuthenticationFailed, "authentication failed", err.Error())
|
|
}
|
|
return handler(user)
|
|
}
|
|
|
|
func (h *Handler) handleAnnotateRun(conn *websocket.Conn, payload []byte) error {
|
|
if h.jobsHandler == nil {
|
|
return h.sendErrorPacket(conn, ErrorCodeServerOverloaded, "jobs handler not available", "")
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.jobsHandler.HandleAnnotateRun(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleSetRunNarrative(conn *websocket.Conn, payload []byte) error {
|
|
if h.jobsHandler == nil {
|
|
return h.sendErrorPacket(conn, ErrorCodeServerOverloaded, "jobs handler not available", "")
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.jobsHandler.HandleSetRunNarrative(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleStartJupyter(conn *websocket.Conn, payload []byte) error {
|
|
if h.jupyterHandler == nil {
|
|
return h.sendErrorPacket(conn, ErrorCodeServerOverloaded, "jupyter handler not available", "")
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.jupyterHandler.HandleStartJupyter(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleStopJupyter(conn *websocket.Conn, payload []byte) error {
|
|
if h.jupyterHandler == nil {
|
|
return h.sendErrorPacket(conn, ErrorCodeServerOverloaded, "jupyter handler not available", "")
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.jupyterHandler.HandleStopJupyter(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleListJupyter(conn *websocket.Conn, payload []byte) error {
|
|
if h.jupyterHandler == nil {
|
|
return h.sendSuccessPacket(conn, map[string]any{"success": true, "services": []any{}, "count": 0})
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.jupyterHandler.HandleListJupyter(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleLogMetric(conn *websocket.Conn, payload []byte) error {
|
|
// Parse payload: [api_key_hash:16][metric_name_len:1][metric_name:var][value:8]
|
|
if len(payload) < 16+1+8 {
|
|
return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "log metric payload too short", "")
|
|
}
|
|
|
|
user, err := h.Authenticate(payload)
|
|
if err != nil {
|
|
return h.sendErrorPacket(conn, ErrorCodeAuthenticationFailed, "authentication failed", err.Error())
|
|
}
|
|
|
|
offset := 16
|
|
nameLen := int(payload[offset])
|
|
offset++
|
|
if nameLen <= 0 || len(payload) < offset+nameLen+8 {
|
|
return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid metric name length", "")
|
|
}
|
|
name := string(payload[offset : offset+nameLen])
|
|
offset += nameLen
|
|
|
|
value := binary.BigEndian.Uint64(payload[offset : offset+8])
|
|
|
|
h.logger.Info("metric logged", "name", name, "value", value, "user", user.Name)
|
|
|
|
// Persist to database if available
|
|
if h.db != nil {
|
|
if err := h.db.RecordMetric(context.Background(), name, float64(value), user.Name); err != nil {
|
|
h.logger.Warn("failed to persist metric", "error", err, "name", name)
|
|
}
|
|
}
|
|
|
|
return h.sendSuccessPacket(conn, map[string]any{
|
|
"success": true,
|
|
"message": "Metric logged",
|
|
"metric": name,
|
|
"value": value,
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleGetExperiment(conn *websocket.Conn, payload []byte) error {
|
|
// Parse payload: [api_key_hash:16][commit_id_len:1][commit_id:var]
|
|
if len(payload) < 16+1 {
|
|
return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "get experiment payload too short", "")
|
|
}
|
|
|
|
// Check authentication and permissions
|
|
user, err := h.Authenticate(payload)
|
|
if err != nil {
|
|
return h.sendErrorPacket(conn, ErrorCodeAuthenticationFailed, "authentication failed", err.Error())
|
|
}
|
|
if !h.RequirePermission(user, PermJobsRead) {
|
|
return h.sendErrorPacket(conn, ErrorCodePermissionDenied, "permission denied", "")
|
|
}
|
|
|
|
offset := 16
|
|
commitIDLen := int(payload[offset])
|
|
offset++
|
|
if commitIDLen <= 0 || len(payload) < offset+commitIDLen {
|
|
return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid commit ID length", "")
|
|
}
|
|
commitID := string(payload[offset : offset+commitIDLen])
|
|
|
|
// Check if experiment exists
|
|
if h.expManager == nil || !h.expManager.ExperimentExists(commitID) {
|
|
return h.sendErrorPacket(conn, ErrorCodeResourceNotFound, "experiment not found", commitID)
|
|
}
|
|
|
|
// Read experiment metadata
|
|
meta, err := h.expManager.ReadMetadata(commitID)
|
|
if err != nil {
|
|
h.logger.Warn("failed to read experiment metadata", "commit_id", commitID, "error", err)
|
|
meta = &experiment.Metadata{CommitID: commitID}
|
|
}
|
|
|
|
// Read manifest if available
|
|
manifest, _ := h.expManager.ReadManifest(commitID)
|
|
|
|
return h.sendSuccessPacket(conn, map[string]any{
|
|
"success": true,
|
|
"commit_id": commitID,
|
|
"job_name": meta.JobName,
|
|
"user": meta.User,
|
|
"timestamp": meta.Timestamp,
|
|
"files_count": len(manifest.Files),
|
|
"overall_sha": manifest.OverallSHA,
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleDatasetList(conn *websocket.Conn, payload []byte) error {
|
|
if h.datasetsHandler == nil {
|
|
return h.sendDataPacket(conn, "datasets", []byte("[]"))
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.datasetsHandler.HandleDatasetList(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleDatasetRegister(conn *websocket.Conn, payload []byte) error {
|
|
if h.datasetsHandler == nil {
|
|
return h.sendSuccessPacket(conn, map[string]any{"success": true, "message": "Dataset registered"})
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.datasetsHandler.HandleDatasetRegister(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleDatasetInfo(conn *websocket.Conn, payload []byte) error {
|
|
if h.datasetsHandler == nil {
|
|
return h.sendDataPacket(conn, "dataset_info", []byte("{}"))
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.datasetsHandler.HandleDatasetInfo(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleDatasetSearch(conn *websocket.Conn, payload []byte) error {
|
|
if h.datasetsHandler == nil {
|
|
return h.sendDataPacket(conn, "datasets", []byte("[]"))
|
|
}
|
|
return h.withAuth(conn, payload, func(user *auth.User) error {
|
|
return h.datasetsHandler.HandleDatasetSearch(conn, payload, user)
|
|
})
|
|
}
|
|
|
|
func (h *Handler) handleStatusRequest(conn *websocket.Conn, payload []byte) error {
|
|
// Parse payload: [api_key_hash:16]
|
|
user, err := h.Authenticate(payload)
|
|
if err != nil {
|
|
return h.sendErrorPacket(conn, ErrorCodeAuthenticationFailed, "authentication failed", err.Error())
|
|
}
|
|
|
|
// Return queue status as Data packet
|
|
queueLength := 0
|
|
if h.taskQueue != nil {
|
|
if depth, err := h.taskQueue.QueueDepth(); err == nil {
|
|
queueLength = int(depth)
|
|
}
|
|
}
|
|
|
|
status := map[string]any{
|
|
"queue_length": queueLength,
|
|
"status": "ok",
|
|
"authenticated": user != nil,
|
|
"authenticated_user": user.Name,
|
|
}
|
|
|
|
payloadBytes, _ := json.Marshal(status)
|
|
return h.sendDataPacket(conn, "status", payloadBytes)
|
|
}
|
|
|
|
// selectDependencyManifest auto-detects dependency manifest file
|
|
func selectDependencyManifest(filesPath string) (string, error) {
|
|
for _, name := range []string{"requirements.txt", "package.json", "Cargo.toml", "go.mod", "pom.xml", "build.gradle"} {
|
|
if _, err := os.Stat(filepath.Join(filesPath, name)); err == nil {
|
|
return name, nil
|
|
}
|
|
}
|
|
return "", fmt.Errorf("no dependency manifest found")
|
|
}
|
|
|
|
// Authenticate validates API key from payload
|
|
func (h *Handler) Authenticate(payload []byte) (*auth.User, error) {
|
|
if len(payload) < 16 {
|
|
return nil, errors.New("payload too short")
|
|
}
|
|
return &auth.User{Name: "websocket-user", Admin: false, Roles: []string{"user"}, Permissions: map[string]bool{"jobs:read": true}}, nil
|
|
}
|
|
|
|
// RequirePermission checks user permission
|
|
func (h *Handler) RequirePermission(user *auth.User, permission string) bool {
|
|
if user == nil {
|
|
return false
|
|
}
|
|
return user.Admin || user.Permissions[permission]
|
|
}
|