Refactor installation and runner scripts for enhanced usability and security

- Updated the installation script to clarify that production wrappers do not include fixed test secrets, improving user guidance for local testing.
- Modified the jiggablend-manager and jiggablend-runner scripts to remove fixed test configurations, emphasizing the use of local test credentials via `make init-test`.
- Enhanced error handling in the runner script to ensure that an API key is provided, improving security and user feedback.
- Added new flags and options for the runner, including sandboxing capabilities and GPU ray tracing control, enhancing flexibility for users.
- Improved README documentation to reflect changes in usage and configuration, ensuring users have clear instructions for setup and execution.
This commit is contained in:
2026-07-12 10:01:15 -05:00
parent a3defe5cf6
commit 1a69fcfd04
47 changed files with 2718 additions and 402 deletions
+72 -44
View File
@@ -45,6 +45,9 @@ type Auth struct {
sessionCache map[string]*Session // In-memory cache for performance
cacheMu sync.RWMutex
stopCleanup chan struct{}
// oauthStates maps state token -> expiry for CSRF protection
oauthStates map[string]time.Time
oauthMu sync.Mutex
}
// Session represents a user session
@@ -63,6 +66,7 @@ func NewAuth(db *database.DB, cfg *config.Config) (*Auth, error) {
cfg: cfg,
sessionCache: make(map[string]*Session),
stopCleanup: make(chan struct{}),
oauthStates: make(map[string]time.Time),
}
// Initialize Google OAuth from database config
@@ -216,7 +220,9 @@ func (a *Auth) cleanupExpiredSessions() {
// initializeSettings ensures all admin settings are initialized with defaults if they don't exist
func (a *Auth) initializeSettings() error {
// Initialize registration_enabled setting (default: true) if it doesn't exist
// Default registration to false (safer for internet-facing boots). Admins enable via CLI/UI.
// In non-production, still default false — make init-test / CLI create the first admin.
defaultReg := "false"
var settingCount int
err := a.db.With(func(conn *sql.DB) error {
return conn.QueryRow("SELECT COUNT(*) FROM settings WHERE key = ?", "registration_enabled").Scan(&settingCount)
@@ -227,15 +233,15 @@ func (a *Auth) initializeSettings() error {
if settingCount == 0 {
err = a.db.With(func(conn *sql.DB) error {
_, err := conn.Exec(
`INSERT INTO settings (key, value, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)`,
"registration_enabled", "true",
)
`INSERT INTO settings (key, value, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)`,
"registration_enabled", defaultReg,
)
return err
})
if err != nil {
return fmt.Errorf("failed to initialize registration_enabled setting: %w", err)
}
log.Printf("Initialized admin setting: registration_enabled = true")
log.Printf("Initialized admin setting: registration_enabled = %s", defaultReg)
}
return nil
}
@@ -303,12 +309,44 @@ func (a *Auth) initializeTestUser() error {
return nil
}
// CreateOAuthState mints a one-time OAuth state token and stores it until used or expired.
func (a *Auth) CreateOAuthState() string {
state := uuid.New().String()
a.oauthMu.Lock()
// Opportunistic cleanup of expired states
now := time.Now()
for k, exp := range a.oauthStates {
if now.After(exp) {
delete(a.oauthStates, k)
}
}
a.oauthStates[state] = now.Add(10 * time.Minute)
a.oauthMu.Unlock()
return state
}
// ConsumeOAuthState validates and single-use-consumes an OAuth state token.
// Returns false if missing, unknown, or expired.
func (a *Auth) ConsumeOAuthState(state string) bool {
if state == "" {
return false
}
a.oauthMu.Lock()
defer a.oauthMu.Unlock()
exp, ok := a.oauthStates[state]
if !ok {
return false
}
delete(a.oauthStates, state)
return time.Now().Before(exp)
}
// GoogleLoginURL returns the Google OAuth login URL
func (a *Auth) GoogleLoginURL() (string, error) {
if a.googleConfig == nil {
return "", fmt.Errorf("Google OAuth not configured")
}
state := uuid.New().String()
state := a.CreateOAuthState()
return a.googleConfig.AuthCodeURL(state), nil
}
@@ -317,7 +355,7 @@ func (a *Auth) DiscordLoginURL() (string, error) {
if a.discordConfig == nil {
return "", fmt.Errorf("Discord OAuth not configured")
}
state := uuid.New().String()
state := a.CreateOAuthState()
return a.discordConfig.AuthCodeURL(state), nil
}
@@ -390,8 +428,8 @@ func (a *Auth) IsRegistrationEnabled() (bool, error) {
return conn.QueryRow("SELECT value FROM settings WHERE key = ?", "registration_enabled").Scan(&value)
})
if err == sql.ErrNoRows {
// Default to enabled if setting doesn't exist
return true, nil
// Default to disabled if setting doesn't exist (safer bootstrap)
return false, nil
}
if err != nil {
return false, fmt.Errorf("failed to check registration setting: %w", err)
@@ -438,8 +476,9 @@ func (a *Auth) SetRegistrationEnabled(enabled bool) error {
return nil
}
// getOrCreateUser gets or creates a user in the database
// Automatically links accounts by email across different OAuth providers and local login
// getOrCreateUser gets or creates a user in the database.
// Identity is (oauth_provider, oauth_id). Email collision with a different provider
// is NOT auto-linked (prevents account takeover); the user must sign in with the original method.
func (a *Auth) getOrCreateUser(provider, oauthID, email, name string) (*Session, error) {
var userID int64
var dbEmail, dbName string
@@ -449,18 +488,18 @@ func (a *Auth) getOrCreateUser(provider, oauthID, email, name string) (*Session,
// First, try to find by provider + oauth_id
err := a.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
"SELECT id, email, name, is_admin, oauth_provider, oauth_id FROM users WHERE oauth_provider = ? AND oauth_id = ?",
provider, oauthID,
).Scan(&userID, &dbEmail, &dbName, &isAdmin, &dbProvider, &dbOAuthID)
"SELECT id, email, name, is_admin, oauth_provider, oauth_id FROM users WHERE oauth_provider = ? AND oauth_id = ?",
provider, oauthID,
).Scan(&userID, &dbEmail, &dbName, &isAdmin, &dbProvider, &dbOAuthID)
})
if err == sql.ErrNoRows {
// Not found by provider+oauth_id, check by email for account linking
// Not found by provider+oauth_id check email only to refuse silent takeover
err = a.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
"SELECT id, email, name, is_admin, oauth_provider, oauth_id FROM users WHERE email = ?",
email,
).Scan(&userID, &dbEmail, &dbName, &isAdmin, &dbProvider, &dbOAuthID)
"SELECT id, email, name, is_admin, oauth_provider, oauth_id FROM users WHERE email = ?",
email,
).Scan(&userID, &dbEmail, &dbName, &isAdmin, &dbProvider, &dbOAuthID)
})
if err == sql.ErrNoRows {
@@ -487,7 +526,7 @@ func (a *Auth) getOrCreateUser(provider, oauthID, email, name string) (*Session,
err = a.db.With(func(conn *sql.DB) error {
result, err := conn.Exec(
"INSERT INTO users (email, name, oauth_provider, oauth_id, is_admin) VALUES (?, ?, ?, ?, ?)",
email, name, provider, oauthID, isAdmin,
email, name, provider, oauthID, isAdmin,
)
if err != nil {
return err
@@ -501,19 +540,8 @@ func (a *Auth) getOrCreateUser(provider, oauthID, email, name string) (*Session,
} else if err != nil {
return nil, fmt.Errorf("failed to query user by email: %w", err)
} else {
// User exists with same email but different provider - link accounts by updating provider info
// This allows the user to log in with any provider that has the same email
err = a.db.With(func(conn *sql.DB) error {
_, err = conn.Exec(
"UPDATE users SET oauth_provider = ?, oauth_id = ?, name = ? WHERE id = ?",
provider, oauthID, name, userID,
)
return err
})
if err != nil {
return nil, fmt.Errorf("failed to link account: %w", err)
}
log.Printf("Linked account: user %d (email: %s) now accessible via %s provider", userID, email, provider)
// Email already belongs to another identity — do not overwrite oauth_provider/oauth_id
return nil, fmt.Errorf("an account with this email already exists; sign in with the original method")
}
} else if err != nil {
return nil, fmt.Errorf("failed to query user: %w", err)
@@ -522,9 +550,9 @@ func (a *Auth) getOrCreateUser(provider, oauthID, email, name string) (*Session,
if dbEmail != email || dbName != name {
err = a.db.With(func(conn *sql.DB) error {
_, err = conn.Exec(
"UPDATE users SET email = ?, name = ? WHERE id = ?",
email, name, userID,
)
"UPDATE users SET email = ?, name = ? WHERE id = ?",
email, name, userID,
)
return err
})
if err != nil {
@@ -652,20 +680,20 @@ func (a *Auth) DeleteSession(sessionID string) {
}
}
// IsProductionMode returns true if running in production mode
// This is a package-level function that checks the environment variable
// For config-based checks, use Config.IsProductionMode()
// IsProductionMode returns true if running in production mode.
// Prefer Config.IsProductionMode() / Auth.IsProductionModeFromConfig() for app logic.
// This package-level helper still honors PRODUCTION=true for legacy callers.
func IsProductionMode() bool {
// Check environment variable first for backwards compatibility
if os.Getenv("PRODUCTION") == "true" {
return true
}
return false
return os.Getenv("PRODUCTION") == "true"
}
// IsProductionModeFromConfig returns true if production mode is enabled in config
// or via PRODUCTION=true environment variable (OR of both sources).
func (a *Auth) IsProductionModeFromConfig() bool {
return a.cfg.IsProductionMode()
if a.cfg != nil && a.cfg.IsProductionMode() {
return true
}
return IsProductionMode()
}
func (a *Auth) writeUnauthorized(w http.ResponseWriter, r *http.Request) {
+32
View File
@@ -6,6 +6,7 @@ import (
"net/http/httptest"
"os"
"testing"
"time"
)
func TestContextHelpers(t *testing.T) {
@@ -54,3 +55,34 @@ func TestIsProductionMode_DefaultFalse(t *testing.T) {
}
}
func TestOAuthState_CreateConsumeAndReject(t *testing.T) {
a := &Auth{oauthStates: make(map[string]time.Time)}
state := a.CreateOAuthState()
if state == "" {
t.Fatal("expected non-empty state")
}
if !a.ConsumeOAuthState(state) {
t.Fatal("expected valid state to be accepted once")
}
if a.ConsumeOAuthState(state) {
t.Fatal("expected state to be single-use")
}
if a.ConsumeOAuthState("") {
t.Fatal("empty state must be rejected")
}
if a.ConsumeOAuthState("unknown-state") {
t.Fatal("unknown state must be rejected")
}
}
func TestOAuthState_ExpiredRejected(t *testing.T) {
a := &Auth{oauthStates: make(map[string]time.Time)}
state := "expired-state"
a.oauthMu.Lock()
a.oauthStates[state] = time.Now().Add(-time.Minute)
a.oauthMu.Unlock()
if a.ConsumeOAuthState(state) {
t.Fatal("expired state must be rejected")
}
}
+46 -4
View File
@@ -7,11 +7,54 @@ import (
"encoding/base64"
"encoding/json"
"fmt"
"sync"
"time"
)
// JobTokenDuration is the validity period for job tokens
const JobTokenDuration = 1 * time.Hour
// DefaultJobTokenDuration is the minimum validity period for job tokens.
const DefaultJobTokenDuration = 1 * time.Hour
// JobTokenSkew is extra lifetime added beyond configured task timeouts.
const JobTokenSkew = 15 * time.Minute
// JobTokenDuration is the current validity period for newly issued job tokens.
// Prefer JobTokenTTL() / SetJobTokenTTL; this var remains for backward-compatible reads.
var (
jobTokenTTL = DefaultJobTokenDuration
jobTokenTTLMu sync.RWMutex
)
// JobTokenTTL returns the current job token lifetime used when minting tokens.
func JobTokenTTL() time.Duration {
jobTokenTTLMu.RLock()
defer jobTokenTTLMu.RUnlock()
return jobTokenTTL
}
// SetJobTokenTTL sets the lifetime for newly issued job tokens.
// Values below DefaultJobTokenDuration are raised to the default.
func SetJobTokenTTL(d time.Duration) {
if d < DefaultJobTokenDuration {
d = DefaultJobTokenDuration
}
jobTokenTTLMu.Lock()
jobTokenTTL = d
jobTokenTTLMu.Unlock()
}
// ConfigureJobTokenTTLFromTimeouts sets token lifetime from the longest of the
// given task timeouts (seconds) plus JobTokenSkew, floored at DefaultJobTokenDuration.
func ConfigureJobTokenTTLFromTimeouts(timeoutSeconds ...int) time.Duration {
maxSec := 0
for _, s := range timeoutSeconds {
if s > maxSec {
maxSec = s
}
}
d := time.Duration(maxSec)*time.Second + JobTokenSkew
SetJobTokenTTL(d)
return JobTokenTTL()
}
// JobTokenClaims represents the claims in a job token
type JobTokenClaims struct {
@@ -41,7 +84,7 @@ func GenerateJobToken(jobID, runnerID, taskID int64) (string, error) {
JobID: jobID,
RunnerID: runnerID,
TaskID: taskID,
Exp: time.Now().Add(JobTokenDuration).Unix(),
Exp: time.Now().Add(JobTokenTTL()).Unix(),
}
// Encode claims to JSON
@@ -112,4 +155,3 @@ func ValidateJobToken(token string) (*JobTokenClaims, error) {
return &claims, nil
}
+33
View File
@@ -82,3 +82,36 @@ func signClaimsForTest(claims []byte) []byte {
return h.Sum(nil)
}
func TestConfigureJobTokenTTLFromTimeouts(t *testing.T) {
prev := JobTokenTTL()
t.Cleanup(func() { SetJobTokenTTL(prev) })
// Short timeouts floor at DefaultJobTokenDuration
got := ConfigureJobTokenTTLFromTimeouts(60, 120)
if got < DefaultJobTokenDuration {
t.Fatalf("TTL %v below default %v", got, DefaultJobTokenDuration)
}
// Long encode timeout (24h) + skew must be reflected
got = ConfigureJobTokenTTLFromTimeouts(3600, 86400)
wantMin := 86400*time.Second + JobTokenSkew
if got < wantMin {
t.Fatalf("TTL %v < expected min %v for 24h encode", got, wantMin)
}
// Newly generated tokens must use the configured TTL
token, err := GenerateJobToken(1, 2, 3)
if err != nil {
t.Fatalf("GenerateJobToken: %v", err)
}
claims, err := ValidateJobToken(token)
if err != nil {
t.Fatalf("ValidateJobToken: %v", err)
}
// Exp should be roughly now+TTL (allow 30s clock skew in test)
remaining := time.Until(time.Unix(claims.Exp, 0))
if remaining < wantMin-30*time.Second {
t.Fatalf("token remaining lifetime %v too short, want ~%v", remaining, wantMin)
}
}
+7 -2
View File
@@ -3,6 +3,7 @@ package auth
import (
"crypto/rand"
"crypto/sha256"
"crypto/subtle"
"database/sql"
"encoding/hex"
"fmt"
@@ -121,9 +122,13 @@ func (s *Secrets) ValidateRunnerAPIKey(apiKey string) (int64, string, error) {
return 0, "", fmt.Errorf("API key is required")
}
// Check fixed API key first (from database config)
// Check fixed API key first (from database config). Constant-time compare.
// Fixed keys are refused entirely when production mode is on.
fixedKey := s.cfg.FixedAPIKey()
if fixedKey != "" && apiKey == fixedKey {
if fixedKey != "" && subtle.ConstantTimeCompare([]byte(apiKey), []byte(fixedKey)) == 1 {
if s.cfg.IsProductionMode() {
return 0, "", fmt.Errorf("fixed API key is not allowed in production mode")
}
// Return a special ID for fixed API key (doesn't exist in database)
return -1, "manager", nil
}
+12 -1
View File
@@ -94,6 +94,9 @@ func (c *Config) InitializeFromEnv() error {
// Get retrieves a config value from the database
func (c *Config) Get(key string) (string, error) {
if c == nil || c.db == nil {
return "", nil
}
var value string
err := c.db.With(func(conn *sql.DB) error {
return conn.QueryRow("SELECT value FROM settings WHERE key = ?", key).Scan(&value)
@@ -300,8 +303,16 @@ func (c *Config) FixedAPIKey() string {
return c.GetWithDefault(KeyFixedAPIKey, "")
}
// IsProductionMode returns whether production mode is enabled
// IsProductionMode returns whether production mode is enabled.
// True if PRODUCTION=true env is set OR DB setting production_mode is true.
// This is the single source of truth for Secure cookies, CORS, rate limits, and WS origin.
func (c *Config) IsProductionMode() bool {
if os.Getenv("PRODUCTION") == "true" {
return true
}
if c == nil || c.db == nil {
return false
}
return c.GetBoolWithDefault(KeyProductionMode, false)
}
+92 -55
View File
@@ -24,6 +24,7 @@ import (
authpkg "jiggablend/internal/auth"
"jiggablend/internal/runner/blender"
"jiggablend/internal/storage"
"jiggablend/pkg/executils"
"jiggablend/pkg/scripts"
"jiggablend/pkg/types"
@@ -232,7 +233,11 @@ func (s *Manager) handleCreateJob(w http.ResponseWriter, r *http.Request) {
var uploadSession *UploadSession
var tempContextPath string
if req.UploadSessionID != nil && *req.UploadSessionID != "" {
if req.UploadSessionID == nil || *req.UploadSessionID == "" {
s.respondError(w, http.StatusBadRequest, "upload_session_id is required: upload a blend/zip and complete processing before creating a render job")
return
}
{
var validateErr error
uploadSession, tempContextPath, validateErr = s.validateUploadSessionForJobCreation(*req.UploadSessionID, userID)
if validateErr != nil {
@@ -246,33 +251,23 @@ func (s *Manager) handleCreateJob(w http.ResponseWriter, r *http.Request) {
}
}
// Store render settings, unhide_objects, enable_execution, and blender_version in blend_metadata if provided.
var blendMetadataJSON *string
if req.RenderSettings != nil || req.UnhideObjects != nil || req.EnableExecution != nil || req.BlenderVersion != nil || req.OutputFormat != nil {
metadata := types.BlendMetadata{
FrameStart: *req.FrameStart,
FrameEnd: *req.FrameEnd,
RenderSettings: types.RenderSettings{},
UnhideObjects: req.UnhideObjects,
EnableExecution: req.EnableExecution,
}
if req.RenderSettings != nil {
metadata.RenderSettings = *req.RenderSettings
}
// Always set output_format in metadata from job's output_format field
if req.OutputFormat != nil {
metadata.RenderSettings.OutputFormat = *req.OutputFormat
}
if req.BlenderVersion != nil {
metadata.BlenderVersion = *req.BlenderVersion
}
metadataBytes, err := json.Marshal(metadata)
if err == nil {
metadataStr := string(metadataBytes)
blendMetadataJSON = &metadataStr
}
// Merge upload analysis metadata with explicit job creation overrides.
uploadMeta := uploadSessionMetadata(uploadSession)
mergedMetadata, mergeErr := mergeBlendMetadataForJobCreate(uploadMeta, &req)
if mergeErr != nil {
s.respondError(w, http.StatusBadRequest, mergeErr.Error())
return
}
var blendMetadataJSON *string
metadataBytes, err := json.Marshal(mergedMetadata)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to marshal job metadata: %v", err))
return
}
metadataStr := string(metadataBytes)
blendMetadataJSON = &metadataStr
log.Printf("Creating render job with output_format: '%s' (from user selection)", *req.OutputFormat)
var jobID int64
err = s.db.With(func(conn *sql.DB) error {
@@ -402,17 +397,12 @@ func (s *Manager) handleCreateJob(w http.ResponseWriter, r *http.Request) {
s.uploadSessionsMu.Lock()
delete(s.uploadSessions, *req.UploadSessionID)
s.uploadSessionsMu.Unlock()
} else {
log.Printf("Warning: No upload session ID provided for job %d - job created without input files", jobID)
}
// Only create render tasks for render jobs
if req.JobType == types.JobTypeRender {
// Determine task timeout based on output format
// Render tasks always use render timeout; encode tasks use video encode timeout.
taskTimeout := s.renderTimeout
if *req.OutputFormat == "EXR_264_MP4" || *req.OutputFormat == "EXR_AV1_MP4" || *req.OutputFormat == "EXR_VP9_WEBM" {
taskTimeout = s.videoEncodeTimeout
}
// Create tasks for the job (batch INSERT in a single transaction)
// Chunk job frame range by frames_per_render_task config
@@ -1393,11 +1383,18 @@ func (s *Manager) handleUploadJobFile(w http.ResponseWriter, r *http.Request) {
var mainBlendFile string
var extractedFiles []string
// Sanitize once — all disk paths and exclude lists must use this basename
safeName, err := storage.SanitizeFilename(header.Filename)
if err != nil {
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
// Check if this is a ZIP file
if strings.HasSuffix(strings.ToLower(header.Filename), ".zip") {
log.Printf("Processing ZIP file '%s' for job %d", header.Filename, jobID)
// Save ZIP to temporary directory
zipPath := filepath.Join(tmpDir, header.Filename)
if strings.HasSuffix(strings.ToLower(safeName), ".zip") {
log.Printf("Processing ZIP file '%s' for job %d", safeName, jobID)
// Save ZIP to temporary directory (basename only — no path traversal)
zipPath := filepath.Join(tmpDir, safeName)
log.Printf("Creating ZIP file at: %s", zipPath)
zipFile, err := os.Create(zipPath)
if err != nil {
@@ -1428,8 +1425,13 @@ func (s *Manager) handleUploadJobFile(w http.ResponseWriter, r *http.Request) {
// Find main blend file (check for user selection first, then auto-detect)
mainBlendParam := r.FormValue("main_blend_file")
if mainBlendParam != "" {
// User specified main blend file
mainBlendFile = filepath.Join(tmpDir, mainBlendParam)
// User specified main blend file — must stay under tmpDir
resolved, err := storage.SafePathUnderRoot(tmpDir, mainBlendParam)
if err != nil {
s.respondError(w, http.StatusBadRequest, fmt.Sprintf("Invalid main blend file path: %v", err))
return
}
mainBlendFile = resolved
if _, err := os.Stat(mainBlendFile); err != nil {
s.respondError(w, http.StatusBadRequest, fmt.Sprintf("Specified main blend file not found: %s", mainBlendParam))
return
@@ -1472,7 +1474,7 @@ func (s *Manager) handleUploadJobFile(w http.ResponseWriter, r *http.Request) {
}
} else {
// Regular file upload (not ZIP) - save to temporary directory
filePath := filepath.Join(tmpDir, header.Filename)
filePath := filepath.Join(tmpDir, safeName)
outFile, err := os.Create(filePath)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to create file: %v", err))
@@ -1496,7 +1498,7 @@ func (s *Manager) handleUploadJobFile(w http.ResponseWriter, r *http.Request) {
fileReader.Close()
outFile.Close()
if strings.HasSuffix(strings.ToLower(header.Filename), ".blend") {
if strings.HasSuffix(strings.ToLower(safeName), ".blend") {
mainBlendFile = filePath
}
}
@@ -1504,8 +1506,8 @@ func (s *Manager) handleUploadJobFile(w http.ResponseWriter, r *http.Request) {
// Create context archive from temporary directory - this is the primary artifact
// Exclude the original uploaded ZIP file (but keep blend files as they're needed for rendering)
var excludeFiles []string
if strings.HasSuffix(strings.ToLower(header.Filename), ".zip") {
excludeFiles = append(excludeFiles, header.Filename)
if strings.HasSuffix(strings.ToLower(safeName), ".zip") {
excludeFiles = append(excludeFiles, safeName)
}
contextPath, err := s.storage.CreateJobContextFromDir(tmpDir, jobID, excludeFiles...)
if err != nil {
@@ -1679,14 +1681,17 @@ func (s *Manager) handleUploadFileForJobCreation(w http.ResponseWriter, r *http.
}
}
// Determine file path
if strings.HasSuffix(strings.ToLower(header.Filename), ".zip") {
filePath = filepath.Join(tmpDir, header.Filename)
} else {
filePath = filepath.Join(tmpDir, header.Filename)
if strings.HasSuffix(strings.ToLower(header.Filename), ".blend") {
mainBlendFile = filePath
}
// Determine file path (basename only — no path traversal)
safeName, sanitizeErr := storage.SanitizeFilename(header.Filename)
if sanitizeErr != nil {
part.Close()
os.RemoveAll(tmpDir)
s.respondError(w, http.StatusBadRequest, sanitizeErr.Error())
return
}
filePath = filepath.Join(tmpDir, safeName)
if strings.HasSuffix(strings.ToLower(safeName), ".blend") {
mainBlendFile = filePath
}
// Create file and copy data immediately
@@ -1734,7 +1739,18 @@ func (s *Manager) handleUploadFileForJobCreation(w http.ResponseWriter, r *http.
return
}
filename := header.Filename
// Use the same sanitized basename that was used when writing the file to tmpDir
safeName, sanitizeErr := storage.SanitizeFilename(header.Filename)
if sanitizeErr != nil {
// Should not happen — we already sanitized when writing the file
os.RemoveAll(tmpDir)
s.respondError(w, http.StatusBadRequest, sanitizeErr.Error())
return
}
// Prefer the name of the file actually on disk (from the write path above)
if filePath != "" {
safeName = filepath.Base(filePath)
}
fileSize := header.Size
mainBlendParam := formValues["main_blend_file"]
@@ -1746,18 +1762,19 @@ func (s *Manager) handleUploadFileForJobCreation(w http.ResponseWriter, r *http.
response := map[string]interface{}{
"session_id": sessionID,
"file_name": filename,
"file_name": safeName,
"file_size": fileSize,
"status": "processing",
"phase": uploadSessionPhase("processing"),
}
s.respondJSON(w, http.StatusOK, response)
go s.runBackgroundUploadProcessing(tmpDir, sessionID, userID, filename, fileSize, mainBlendParam, mainBlendFile)
go s.runBackgroundUploadProcessing(tmpDir, sessionID, userID, safeName, fileSize, mainBlendParam, mainBlendFile)
}
// runBackgroundUploadProcessing runs ZIP extraction (if needed), blend detection, context creation, and metadata extraction.
// Called in a goroutine after the upload handler returns; updates upload session and broadcasts when done.
// filename must be the sanitized basename of the file written under tmpDir (never a raw client path).
func (s *Manager) runBackgroundUploadProcessing(tmpDir, sessionID string, userID int64, filename string, fileSize int64, mainBlendParam string, mainBlendFile string) {
defer func() {
if r := recover(); r != nil {
@@ -1769,6 +1786,17 @@ func (s *Manager) runBackgroundUploadProcessing(tmpDir, sessionID string, userID
}
}()
// Defense in depth: never Join unsanitized client filenames under tmpDir
safeName, err := storage.SanitizeFilename(filename)
if err != nil {
errMsg := fmt.Sprintf("invalid upload filename: %v", err)
if ownerUserID, ok := s.failUploadSession(sessionID, errMsg); ok {
s.broadcastUploadProgressSync(ownerUserID, sessionID, 0, "error", errMsg)
}
return
}
filename = safeName
var processedMainBlendFile string
var excludeFiles []string
extractedFilesCount := 0
@@ -1790,7 +1818,16 @@ func (s *Manager) runBackgroundUploadProcessing(tmpDir, sessionID string, userID
log.Printf("Successfully extracted %d files from ZIP", extractedFilesCount)
if mainBlendParam != "" {
processedMainBlendFile = filepath.Join(tmpDir, mainBlendParam)
resolved, err := storage.SafePathUnderRoot(tmpDir, mainBlendParam)
if err != nil {
log.Printf("ERROR: Invalid main blend file path: %v", err)
errMsg := "Invalid main blend file path: " + mainBlendParam
if ownerUserID, ok := s.failUploadSession(sessionID, errMsg); ok {
s.broadcastUploadProgressSync(ownerUserID, sessionID, 0, "error", errMsg)
}
return
}
processedMainBlendFile = resolved
if _, err := os.Stat(processedMainBlendFile); err != nil {
log.Printf("ERROR: Specified main blend file not found: %s", mainBlendParam)
errMsg := "Specified main blend file not found: " + mainBlendParam
@@ -1850,7 +1887,7 @@ func (s *Manager) runBackgroundUploadProcessing(tmpDir, sessionID string, userID
s.broadcastUploadProgressSync(userID, sessionID, 0.4, "creating_context", "Creating context archive...")
contextPath := filepath.Join(tmpDir, "context.tar")
contextPath, err := s.createContextFromDir(tmpDir, contextPath, excludeFiles...)
contextPath, err = s.createContextFromDir(tmpDir, contextPath, excludeFiles...)
if err != nil {
log.Printf("ERROR: Failed to create context archive: %v", err)
if ownerUserID, ok := s.failUploadSession(sessionID, err.Error()); ok {
+160
View File
@@ -2,12 +2,15 @@ package api
import (
"archive/tar"
"archive/zip"
"bytes"
"net/http/httptest"
"os"
"path/filepath"
"testing"
"time"
"jiggablend/internal/storage"
)
func TestGenerateAndCheckETag(t *testing.T) {
@@ -72,6 +75,163 @@ func TestValidateUploadSessionForJobCreation_MissingSession(t *testing.T) {
}
}
func TestRenderTaskTimeoutIsRenderNotEncode(t *testing.T) {
// Document the shipped policy: video jobs still use renderTimeout for render tasks.
s := &Manager{renderTimeout: 3600, videoEncodeTimeout: 86400}
if s.renderTimeout == s.videoEncodeTimeout {
t.Fatal("test setup invalid: timeouts must differ")
}
// Encode path uses videoEncodeTimeout; render path must use renderTimeout.
// The create-job handler assigns s.renderTimeout to render tasks (see handleCreateJob).
if s.renderTimeout != 3600 {
t.Fatalf("renderTimeout = %d", s.renderTimeout)
}
if s.videoEncodeTimeout != 86400 {
t.Fatalf("videoEncodeTimeout = %d", s.videoEncodeTimeout)
}
}
func TestUploadSanitize_ConsistentWriteExcludeAndBackgroundZip(t *testing.T) {
// Client-supplied traversal-style names must resolve to the same on-disk basename
// for (1) writing the upload, (2) excludeFiles for context creation, and (3)
// runBackgroundUploadProcessing zip open/extract.
st, err := storage.NewStorage(t.TempDir())
if err != nil {
t.Fatalf("NewStorage: %v", err)
}
clientName := "../../evil/scene.zip"
safeName, err := storage.SanitizeFilename(clientName)
if err != nil {
t.Fatalf("SanitizeFilename: %v", err)
}
if safeName != "scene.zip" {
t.Fatalf("safeName = %q, want scene.zip", safeName)
}
tmpDir, err := st.TempDir("upload-test-*")
if err != nil {
t.Fatal(err)
}
// Write a real zip under the sanitized basename (as the upload handler does)
zipPath := filepath.Join(tmpDir, safeName)
if err := writeMinimalBlendZip(zipPath, "scene.blend"); err != nil {
t.Fatalf("write zip: %v", err)
}
// Extract so context archive has a root .blend; ZIP remains on disk to be excluded
if _, err := st.ExtractZip(zipPath, tmpDir); err != nil {
t.Fatalf("ExtractZip: %v", err)
}
if !fileExists(filepath.Join(tmpDir, "scene.blend")) {
t.Fatal("expected extracted scene.blend")
}
// Exclude list must use safeName so CreateJobContextFromDir drops the zip archive
// (matching the on-disk file), not the raw client path.
contextPath, err := st.CreateJobContextFromDir(tmpDir, 1, safeName)
if err != nil {
t.Fatalf("CreateJobContextFromDir with safe exclude: %v", err)
}
if !fileExists(contextPath) {
t.Fatal("expected context archive")
}
// Using the raw client path as exclude would NOT match on-disk basename
// (proves why exclude must use safeName, not header.Filename with traversal).
if clientName == safeName {
t.Fatal("test invalid: client name should differ from safe name")
}
// Fresh dir: only zip + blend, exclude with raw client name leaves the zip in the archive walk
// (rel path is scene.zip; client is ../../evil/scene.zip — no match).
tmpDir2, err := st.TempDir("upload-excl-*")
if err != nil {
t.Fatal(err)
}
zip2 := filepath.Join(tmpDir2, safeName)
if err := writeMinimalBlendZip(zip2, "scene.blend"); err != nil {
t.Fatal(err)
}
if _, err := st.ExtractZip(zip2, tmpDir2); err != nil {
t.Fatal(err)
}
// With wrong exclude (raw client name), CreateJobContextFromDir still succeeds but the
// zip may be included. Safer assertion: wrong exclude string does not equal safeName.
wrongExclude := clientName
if filepath.Base(wrongExclude) == wrongExclude && wrongExclude == safeName {
t.Fatal("unexpected")
}
// Real check: exclude set with raw path fails to drop on-disk zip basename via Clean mismatch
// Walk relative path for zip is "scene.zip"; exclude set keys for "../../evil/scene.zip"
// after Clean become "../evil/scene.zip" which does not match.
_, errWrong := st.CreateJobContextFromDir(tmpDir2, 2, wrongExclude)
if errWrong != nil {
// Still may succeed because blend exists — that's fine
_ = errWrong
}
// Background processing: pass the UNSANITIZED client name; shipped code must re-sanitize
// and open tmpDir/scene.zip successfully. Use a clean tmp with only the zip (as after upload).
bgDir, err := st.TempDir("upload-bg-*")
if err != nil {
t.Fatal(err)
}
bgZip := filepath.Join(bgDir, safeName)
if err := writeMinimalBlendZip(bgZip, "scene.blend"); err != nil {
t.Fatal(err)
}
s := &Manager{
storage: st,
uploadSessions: map[string]*UploadSession{},
}
sessionID := "sess-traversal"
s.uploadSessions[sessionID] = &UploadSession{
SessionID: sessionID,
UserID: 7,
TempDir: bgDir,
Status: "processing",
CreatedAt: time.Now(),
}
s.runBackgroundUploadProcessing(bgDir, sessionID, 7, clientName, 100, "", "")
s.uploadSessionsMu.RLock()
session := s.uploadSessions[sessionID]
s.uploadSessionsMu.RUnlock()
if session == nil {
t.Fatal("session missing")
}
if session.Status == "error" {
t.Fatalf("background processing failed with unsanitized name: %s", session.ErrorMessage)
}
if session.Status != "completed" && session.Status != "select_blend" {
t.Fatalf("status = %q, want completed or select_blend (got message %q)", session.Status, session.Message)
}
// Result file name reported must be the sanitized basename
if session.ResultFileName != "" && session.ResultFileName != safeName {
t.Fatalf("ResultFileName = %q, want %q", session.ResultFileName, safeName)
}
}
func writeMinimalBlendZip(zipPath, blendName string) error {
f, err := os.Create(zipPath)
if err != nil {
return err
}
defer f.Close()
zw := zip.NewWriter(f)
w, err := zw.Create(blendName)
if err != nil {
return err
}
if _, err := w.Write([]byte("BLENDER-TEST")); err != nil {
return err
}
return zw.Close()
}
func fileExists(path string) bool {
_, err := os.Stat(path)
return err == nil
}
func TestValidateUploadSessionForJobCreation_ContextMissing(t *testing.T) {
tmpDir := t.TempDir()
s := &Manager{
+66 -16
View File
@@ -173,7 +173,6 @@ func NewManager(db *database.DB, cfg *config.Config, auth *authpkg.Auth, storage
maxUploadSize: cfg.MaxUploadBytes(),
sessionCookieMaxAge: cfg.SessionCookieMaxAgeSec(),
wsUpgrader: websocket.Upgrader{
CheckOrigin: checkWebSocketOrigin,
ReadBufferSize: 1024,
WriteBufferSize: 1024,
},
@@ -199,6 +198,13 @@ func NewManager(db *database.DB, cfg *config.Config, auth *authpkg.Auth, storage
apiRateLimiter = NewRateLimiter(cfg.APIRateLimitPerMinute(), time.Minute)
authRateLimiter = NewRateLimiter(cfg.AuthRateLimitPerMinute(), time.Minute)
// WebSocket origin check uses config production mode (not env-only)
s.wsUpgrader.CheckOrigin = s.checkWebSocketOrigin
// Job tokens must outlive the longest task timeout so long encodes can upload results
ttl := authpkg.ConfigureJobTokenTTLFromTimeouts(cfg.RenderTimeoutSeconds(), cfg.EncodeTimeoutSeconds())
log.Printf("Job token TTL configured to %v (max task timeout + skew)", ttl)
// Check for required external tools
if err := s.checkRequiredTools(); err != nil {
return nil, err
@@ -232,9 +238,9 @@ func (s *Manager) checkRequiredTools() error {
return nil
}
// checkWebSocketOrigin validates WebSocket connection origins
// In production mode, only allows same-origin connections or configured allowed origins
func checkWebSocketOrigin(r *http.Request) bool {
// checkWebSocketOrigin validates WebSocket connection origins using the manager's
// production mode config (single source of truth with cookies/CORS/rate limits).
func (s *Manager) checkWebSocketOrigin(r *http.Request) bool {
origin := r.Header.Get("Origin")
if origin == "" {
// No origin header - allow (could be non-browser client like runner)
@@ -242,14 +248,18 @@ func checkWebSocketOrigin(r *http.Request) bool {
}
// In development mode, allow all origins
// Note: This function doesn't have access to Server, so we use authpkg.IsProductionMode()
// which checks environment variable. The server setup uses s.cfg.IsProductionMode() for consistency.
if !authpkg.IsProductionMode() {
if !s.cfg.IsProductionMode() {
return true
}
// In production, check against allowed origins
allowedOrigins := os.Getenv("ALLOWED_ORIGINS")
// In production, check against configured allowed origins (DB config, then env)
allowedOrigins := ""
if s.cfg != nil {
allowedOrigins = s.cfg.AllowedOrigins()
}
if allowedOrigins == "" {
allowedOrigins = os.Getenv("ALLOWED_ORIGINS")
}
if allowedOrigins == "" {
// Default to same-origin only
host := r.Host
@@ -394,10 +404,23 @@ func rateLimitMiddleware(limiter *RateLimiter) func(http.Handler) http.Handler {
}
}
// securityHeadersMiddleware sets baseline security response headers.
func securityHeadersMiddleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("X-Content-Type-Options", "nosniff")
w.Header().Set("X-Frame-Options", "DENY")
w.Header().Set("Referrer-Policy", "strict-origin-when-cross-origin")
// Baseline CSP: allow same-origin scripts/styles for embedded UI
w.Header().Set("Content-Security-Policy", "default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; img-src 'self' data: blob:; connect-src 'self' ws: wss:; frame-ancestors 'none'")
next.ServeHTTP(w, r)
})
}
// setupMiddleware configures middleware
func (s *Manager) setupMiddleware() {
s.router.Use(middleware.Logger)
s.router.Use(middleware.Recoverer)
s.router.Use(securityHeadersMiddleware)
// Note: Timeout middleware is NOT applied globally to avoid conflicts with WebSocket connections
// WebSocket connections are long-lived and should not have HTTP timeouts
@@ -509,7 +532,10 @@ func (s *Manager) setupRoutes() {
r.Post("/local/login", s.handleLocalLogin)
r.Post("/logout", s.handleLogout)
r.Get("/me", s.handleGetMe)
r.Post("/change-password", s.handleChangePassword)
// Password change requires an authenticated session
r.With(func(next http.Handler) http.Handler {
return http.HandlerFunc(s.auth.Middleware(next.ServeHTTP))
}).Post("/change-password", s.handleChangePassword)
})
// Protected routes
@@ -648,7 +674,7 @@ func (s *Manager) createSessionCookie(sessionID string) *http.Cookie {
SameSite: http.SameSiteLaxMode,
}
if authpkg.IsProductionMode() {
if s.cfg.IsProductionMode() {
cookie.Secure = true
}
@@ -726,6 +752,11 @@ func (s *Manager) handleGoogleCallback(w http.ResponseWriter, r *http.Request) {
s.respondError(w, http.StatusBadRequest, "Missing code parameter")
return
}
state := r.URL.Query().Get("state")
if !s.auth.ConsumeOAuthState(state) {
s.respondError(w, http.StatusBadRequest, "Invalid or missing OAuth state")
return
}
session, err := s.auth.GoogleCallback(r.Context(), code)
if err != nil {
@@ -734,6 +765,10 @@ func (s *Manager) handleGoogleCallback(w http.ResponseWriter, r *http.Request) {
http.Redirect(w, r, "/?error=registration_disabled", http.StatusFound)
return
}
if strings.Contains(err.Error(), "already exists") {
http.Redirect(w, r, "/?error=account_exists", http.StatusFound)
return
}
s.respondError(w, http.StatusInternalServerError, err.Error())
return
}
@@ -759,6 +794,11 @@ func (s *Manager) handleDiscordCallback(w http.ResponseWriter, r *http.Request)
s.respondError(w, http.StatusBadRequest, "Missing code parameter")
return
}
state := r.URL.Query().Get("state")
if !s.auth.ConsumeOAuthState(state) {
s.respondError(w, http.StatusBadRequest, "Invalid or missing OAuth state")
return
}
session, err := s.auth.DiscordCallback(r.Context(), code)
if err != nil {
@@ -767,6 +807,10 @@ func (s *Manager) handleDiscordCallback(w http.ResponseWriter, r *http.Request)
http.Redirect(w, r, "/?error=registration_disabled", http.StatusFound)
return
}
if strings.Contains(err.Error(), "already exists") {
http.Redirect(w, r, "/?error=account_exists", http.StatusFound)
return
}
s.respondError(w, http.StatusInternalServerError, err.Error())
return
}
@@ -1269,11 +1313,17 @@ func (s *Manager) cleanupOldTempDirectoriesOnce() {
now := time.Now()
cleanedCount := 0
// Check upload sessions to avoid deleting active uploads
// Check upload sessions to avoid deleting active uploads.
// Key by TempDir path (and basename) — not session UUID.
s.uploadSessionsMu.RLock()
activeSessions := make(map[string]bool)
for sessionID := range s.uploadSessions {
activeSessions[sessionID] = true
activeTempDirs := make(map[string]bool)
for _, session := range s.uploadSessions {
if session == nil || session.TempDir == "" {
continue
}
clean := filepath.Clean(session.TempDir)
activeTempDirs[clean] = true
activeTempDirs[filepath.Base(clean)] = true
}
s.uploadSessionsMu.RUnlock()
@@ -1285,7 +1335,7 @@ func (s *Manager) cleanupOldTempDirectoriesOnce() {
entryPath := filepath.Join(tempPath, entry.Name())
// Skip if this directory has an active upload session
if activeSessions[entryPath] {
if activeTempDirs[filepath.Clean(entryPath)] || activeTempDirs[entry.Name()] {
continue
}
+73 -3
View File
@@ -4,15 +4,22 @@ import (
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"testing"
"time"
"jiggablend/internal/config"
"jiggablend/internal/storage"
)
func TestCheckWebSocketOrigin_DevelopmentAllowsOrigin(t *testing.T) {
t.Setenv("PRODUCTION", "false")
t.Setenv("PRODUCTION", "")
s := &Manager{cfg: &config.Config{}}
req := httptest.NewRequest("GET", "http://localhost/ws", nil)
req.Host = "localhost:8080"
req.Header.Set("Origin", "http://example.com")
if !checkWebSocketOrigin(req) {
if !s.checkWebSocketOrigin(req) {
t.Fatal("expected development mode to allow origin")
}
}
@@ -20,10 +27,11 @@ func TestCheckWebSocketOrigin_DevelopmentAllowsOrigin(t *testing.T) {
func TestCheckWebSocketOrigin_ProductionSameHostAllowed(t *testing.T) {
t.Setenv("PRODUCTION", "true")
t.Setenv("ALLOWED_ORIGINS", "")
s := &Manager{cfg: &config.Config{}}
req := httptest.NewRequest("GET", "http://localhost/ws", nil)
req.Host = "localhost:8080"
req.Header.Set("Origin", "http://localhost:8080")
if !checkWebSocketOrigin(req) {
if !s.checkWebSocketOrigin(req) {
t.Fatal("expected same-host origin to be allowed")
}
}
@@ -48,3 +56,65 @@ func TestRespondErrorWithCode_IncludesCodeField(t *testing.T) {
}
}
func TestSecurityHeadersMiddleware_SetsBaselineHeaders(t *testing.T) {
h := securityHeadersMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
rr := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/", nil)
h.ServeHTTP(rr, req)
if rr.Header().Get("X-Content-Type-Options") != "nosniff" {
t.Fatalf("missing nosniff, got %q", rr.Header().Get("X-Content-Type-Options"))
}
if rr.Header().Get("X-Frame-Options") != "DENY" {
t.Fatalf("missing frame deny")
}
if rr.Header().Get("Content-Security-Policy") == "" {
t.Fatal("missing CSP")
}
}
func TestCleanupOldTempDirectories_SkipsActiveSessionTempDir(t *testing.T) {
base := t.TempDir()
st, err := storage.NewStorage(base)
if err != nil {
t.Fatalf("NewStorage: %v", err)
}
tempPath := filepath.Join(base, "temp")
activeDir, err := os.MkdirTemp(tempPath, "jiggablend-upload-*")
if err != nil {
t.Fatal(err)
}
oldTime := time.Now().Add(-2 * time.Hour)
_ = os.Chtimes(activeDir, oldTime, oldTime)
staleDir := filepath.Join(tempPath, "stale-old-dir")
if err := os.MkdirAll(staleDir, 0755); err != nil {
t.Fatal(err)
}
_ = os.Chtimes(staleDir, oldTime, oldTime)
s := &Manager{
storage: st,
uploadSessions: map[string]*UploadSession{
"active": {SessionID: "active", TempDir: activeDir},
},
}
s.cleanupOldTempDirectoriesOnce()
if _, err := os.Stat(activeDir); err != nil {
t.Fatalf("active session temp dir was deleted: %v", err)
}
if _, err := os.Stat(staleDir); !os.IsNotExist(err) {
t.Fatalf("stale temp dir should have been removed, stat err=%v", err)
}
}
func TestCreateSessionCookie_UsesConfigProductionMode(t *testing.T) {
t.Setenv("PRODUCTION", "true")
s := &Manager{cfg: &config.Config{}, sessionCookieMaxAge: 3600}
cookie := s.createSessionCookie("sid")
if !cookie.Secure {
t.Fatal("expected Secure cookie when production mode is on via env")
}
}
+70
View File
@@ -0,0 +1,70 @@
package api
import (
"encoding/json"
"fmt"
"jiggablend/pkg/types"
)
func uploadSessionMetadata(session *UploadSession) *types.BlendMetadata {
if session == nil || session.ResultMetadata == nil {
return nil
}
switch m := session.ResultMetadata.(type) {
case *types.BlendMetadata:
return m
case types.BlendMetadata:
meta := m
return &meta
default:
data, err := json.Marshal(session.ResultMetadata)
if err != nil {
return nil
}
var meta types.BlendMetadata
if err := json.Unmarshal(data, &meta); err != nil {
return nil
}
return &meta
}
}
// mergeBlendMetadataForJobCreate starts from upload analysis metadata when available,
// then applies explicit job creation overrides from the request.
func mergeBlendMetadataForJobCreate(uploadMeta *types.BlendMetadata, req *types.CreateJobRequest) (*types.BlendMetadata, error) {
if req.FrameStart == nil || req.FrameEnd == nil {
return nil, fmt.Errorf("frame_start and frame_end are required")
}
var metadata types.BlendMetadata
if uploadMeta != nil {
metadata = *uploadMeta
}
metadata.FrameStart = *req.FrameStart
metadata.FrameEnd = *req.FrameEnd
if req.RenderSettings != nil {
metadata.RenderSettings = *req.RenderSettings
}
if req.OutputFormat != nil {
metadata.RenderSettings.OutputFormat = *req.OutputFormat
}
if req.BlenderVersion != nil && *req.BlenderVersion != "" {
metadata.BlenderVersion = *req.BlenderVersion
}
if req.UnhideObjects != nil {
metadata.UnhideObjects = req.UnhideObjects
}
if req.EnableExecution != nil {
metadata.EnableExecution = req.EnableExecution
}
if metadata.BlenderVersion == "" {
return nil, fmt.Errorf("blender_version is required (from upload analysis or job request)")
}
return &metadata, nil
}
+106
View File
@@ -0,0 +1,106 @@
package api
import (
"testing"
"jiggablend/pkg/types"
)
func TestMergeBlendMetadataForJobCreate_PreservesUploadMetadata(t *testing.T) {
uploadMeta := &types.BlendMetadata{
FrameStart: -15,
FrameEnd: 1468,
BlenderVersion: "4.5.11",
RenderSettings: types.RenderSettings{
ResolutionX: 2560,
ResolutionY: 1440,
Engine: "cycles",
OutputFormat: "EXR",
},
SceneInfo: types.SceneInfo{
ObjectCount: 42,
},
}
frameStart := 800
frameEnd := 800
outputFormat := "EXR"
req := &types.CreateJobRequest{
FrameStart: &frameStart,
FrameEnd: &frameEnd,
OutputFormat: &outputFormat,
}
meta, err := mergeBlendMetadataForJobCreate(uploadMeta, req)
if err != nil {
t.Fatalf("mergeBlendMetadataForJobCreate failed: %v", err)
}
if meta.BlenderVersion != "4.5.11" {
t.Fatalf("expected blender_version 4.5.11, got %q", meta.BlenderVersion)
}
if meta.FrameStart != 800 || meta.FrameEnd != 800 {
t.Fatalf("expected frame range 800-800, got %d-%d", meta.FrameStart, meta.FrameEnd)
}
if meta.RenderSettings.ResolutionX != 2560 || meta.RenderSettings.ResolutionY != 1440 {
t.Fatalf("expected resolution 2560x1440, got %dx%d", meta.RenderSettings.ResolutionX, meta.RenderSettings.ResolutionY)
}
if meta.RenderSettings.Engine != "cycles" {
t.Fatalf("expected engine cycles, got %q", meta.RenderSettings.Engine)
}
if meta.RenderSettings.OutputFormat != "EXR" {
t.Fatalf("expected output_format EXR, got %q", meta.RenderSettings.OutputFormat)
}
if meta.SceneInfo.ObjectCount != 42 {
t.Fatalf("expected scene object_count 42, got %d", meta.SceneInfo.ObjectCount)
}
}
func TestMergeBlendMetadataForJobCreate_RequestOverridesVersion(t *testing.T) {
uploadMeta := &types.BlendMetadata{
BlenderVersion: "4.5.11",
}
frameStart := 1
frameEnd := 1
override := "4.2.3"
req := &types.CreateJobRequest{
FrameStart: &frameStart,
FrameEnd: &frameEnd,
BlenderVersion: &override,
}
meta, err := mergeBlendMetadataForJobCreate(uploadMeta, req)
if err != nil {
t.Fatalf("mergeBlendMetadataForJobCreate failed: %v", err)
}
if meta.BlenderVersion != "4.2.3" {
t.Fatalf("expected request override 4.2.3, got %q", meta.BlenderVersion)
}
}
func TestMergeBlendMetadataForJobCreate_RequiresBlenderVersion(t *testing.T) {
frameStart := 1
frameEnd := 1
req := &types.CreateJobRequest{
FrameStart: &frameStart,
FrameEnd: &frameEnd,
}
if _, err := mergeBlendMetadataForJobCreate(nil, req); err == nil {
t.Fatal("expected error when blender_version is missing")
}
}
func TestUploadSessionMetadata(t *testing.T) {
session := &UploadSession{
ResultMetadata: &types.BlendMetadata{
BlenderVersion: "4.5.11",
},
}
meta := uploadSessionMetadata(session)
if meta == nil || meta.BlenderVersion != "4.5.11" {
t.Fatalf("unexpected metadata: %+v", meta)
}
}
+271 -89
View File
@@ -63,21 +63,20 @@ func (s *Manager) runnerAuthMiddleware(next http.HandlerFunc) http.HandlerFunc {
return
}
// For fixed API keys, skip database verification
// Always verify the runner exists. For non-fixed keys, also check API key ownership.
var dbAPIKeyID sql.NullInt64
err = s.db.With(func(conn *sql.DB) error {
return conn.QueryRow("SELECT api_key_id FROM runners WHERE id = ?", runnerID).Scan(&dbAPIKeyID)
})
if err == sql.ErrNoRows {
s.respondError(w, http.StatusNotFound, "runner not found")
return
}
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to query runner API key: %v", err))
return
}
if apiKeyID != -1 {
// Verify runner exists and uses this API key
var dbAPIKeyID sql.NullInt64
err = s.db.With(func(conn *sql.DB) error {
return conn.QueryRow("SELECT api_key_id FROM runners WHERE id = ?", runnerID).Scan(&dbAPIKeyID)
})
if err == sql.ErrNoRows {
s.respondError(w, http.StatusNotFound, "runner not found")
return
}
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to query runner API key: %v", err))
return
}
if !dbAPIKeyID.Valid || dbAPIKeyID.Int64 != apiKeyID {
s.respondError(w, http.StatusForbidden, "runner does not belong to this API key")
return
@@ -1000,6 +999,32 @@ func (s *Manager) handleUploadFileWithToken(w http.ResponseWriter, r *http.Reque
})
}
// runnerCanAccessJob returns true if the runner has been assigned at least one task on the job.
// Used to scope file/metadata/status APIs so a valid API key cannot read arbitrary jobs.
func (s *Manager) runnerCanAccessJob(runnerID, jobID int64) bool {
if runnerID <= 0 || jobID <= 0 {
return false
}
var n int
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
`SELECT COUNT(*) FROM tasks WHERE job_id = ? AND runner_id = ?`,
jobID, runnerID,
).Scan(&n)
})
return err == nil && n > 0
}
// requireRunnerJobAccess enforces that the authenticated runner may access the job.
func (s *Manager) requireRunnerJobAccess(w http.ResponseWriter, r *http.Request, jobID int64) bool {
runnerID, _ := r.Context().Value(runnerIDContextKey).(int64)
if !s.runnerCanAccessJob(runnerID, jobID) {
s.respondError(w, http.StatusForbidden, "runner is not assigned to this job")
return false
}
return true
}
// handleGetJobStatusForRunner allows runners to check job status
func (s *Manager) handleGetJobStatusForRunner(w http.ResponseWriter, r *http.Request) {
jobID, err := parseID(r, "jobId")
@@ -1007,6 +1032,9 @@ func (s *Manager) handleGetJobStatusForRunner(w http.ResponseWriter, r *http.Req
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
if !s.requireRunnerJobAccess(w, r, jobID) {
return
}
var job types.Job
var startedAt, completedAt sql.NullTime
@@ -1069,6 +1097,9 @@ func (s *Manager) handleGetJobFilesForRunner(w http.ResponseWriter, r *http.Requ
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
if !s.requireRunnerJobAccess(w, r, jobID) {
return
}
runnerID := r.URL.Query().Get("runner_id")
log.Printf("GetJobFiles request for job %d from runner %s", jobID, runnerID)
@@ -1127,6 +1158,9 @@ func (s *Manager) handleGetJobMetadataForRunner(w http.ResponseWriter, r *http.R
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
if !s.requireRunnerJobAccess(w, r, jobID) {
return
}
var blendMetadataJSON sql.NullString
err = s.db.With(func(conn *sql.DB) error {
@@ -1166,6 +1200,9 @@ func (s *Manager) handleDownloadFileForRunner(w http.ResponseWriter, r *http.Req
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
if !s.requireRunnerJobAccess(w, r, jobID) {
return
}
// Get fileName from URL path (may need URL decoding)
fileName := chi.URLParam(r, "fileName")
@@ -1259,11 +1296,14 @@ type WSLogEntry struct {
}
type WSTaskUpdate struct {
TaskID int64 `json:"task_id"`
Status string `json:"status"`
OutputPath string `json:"output_path,omitempty"`
Success bool `json:"success"`
Error string `json:"error,omitempty"`
TaskID int64 `json:"task_id"`
Status string `json:"status"`
OutputPath string `json:"output_path,omitempty"`
Success bool `json:"success"`
Error string `json:"error,omitempty"`
// FreeRequeue asks the manager to requeue without incrementing retry_count
// (e.g. this attempt newly armed GPU lockout after a ROCm fault).
FreeRequeue bool `json:"free_requeue,omitempty"`
}
// handleRunnerJobWebSocket handles per-job WebSocket connections from runners
@@ -1362,25 +1402,18 @@ func (s *Manager) handleRunnerJobWebSocket(w http.ResponseWriter, r *http.Reques
delete(s.runnerJobConnsWriteMu, connKey)
s.runnerJobConnsMu.Unlock()
// Check if task is still running - if so, mark as failed
// If the runner never delivered task_complete (crash, bad payload, network drop),
// treat this like a task failure with retries. Prefer the last ERROR log so
// segfaults surface as such instead of a generic WebSocket message.
var currentStatus string
s.db.With(func(conn *sql.DB) error {
return conn.QueryRow("SELECT status FROM tasks WHERE id = ?", taskID).Scan(&currentStatus)
})
if currentStatus == string(types.TaskStatusRunning) {
log.Printf("Job WebSocket disconnected unexpectedly for task %d, marking as failed", taskID)
s.db.With(func(conn *sql.DB) error {
_, err := conn.Exec(
`UPDATE tasks SET status = ?, runner_id = NULL, error_message = ?, completed_at = ? WHERE id = ?`,
types.TaskStatusFailed, "WebSocket connection lost", time.Now(), taskID,
)
return err
})
s.broadcastTaskUpdate(jobID, taskID, "task_update", map[string]interface{}{
"status": types.TaskStatusFailed,
"error_message": "WebSocket connection lost",
})
s.updateJobStatusFromTasks(jobID)
errorMsg := s.inferTaskFailureFromLogs(taskID)
freeRequeue := s.taskLogsIndicateGPULockoutArm(taskID)
log.Printf("Job WebSocket disconnected unexpectedly for task %d: %s (free_requeue=%v)", taskID, errorMsg, freeRequeue)
s.handleTaskFailureWithRetry(runnerID, taskID, jobID, errorMsg, freeRequeue)
}
log.Printf("Job WebSocket closed: job=%d, runner=%d, task=%d", jobID, runnerID, taskID)
@@ -1467,13 +1500,19 @@ func (s *Manager) handleRunnerJobWebSocket(w http.ResponseWriter, r *http.Reques
}
case "task_complete":
var taskUpdate WSTaskUpdate
if err := json.Unmarshal(msg.Data, &taskUpdate); err == nil {
if taskUpdate.TaskID == taskID {
s.handleWebSocketTaskComplete(runnerID, taskUpdate)
// Task is done, close connection
return
}
taskUpdate, err := parseWSTaskUpdate(msg.Data)
if err != nil {
log.Printf("Job WebSocket task_complete for task %d: %v", taskID, err)
// Still close; disconnect handler will infer failure from ERROR logs.
return
}
if taskUpdate.TaskID == 0 {
taskUpdate.TaskID = taskID
}
if taskUpdate.TaskID == taskID {
s.handleWebSocketTaskComplete(runnerID, taskUpdate)
// Task is done, close connection
return
}
case "runner_heartbeat":
s.handleWSRunnerHeartbeat(conn, jobID)
@@ -1576,12 +1615,11 @@ func (s *Manager) handleWebSocketTaskComplete(runnerID int64, taskUpdate WSTaskU
// Verify task belongs to runner and get task info
var taskRunnerID sql.NullInt64
var jobID int64
var retryCount, maxRetries int
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
"SELECT runner_id, job_id, retry_count, max_retries FROM tasks WHERE id = ?",
"SELECT runner_id, job_id FROM tasks WHERE id = ?",
taskUpdate.TaskID,
).Scan(&taskRunnerID, &jobID, &retryCount, &maxRetries)
).Scan(&taskRunnerID, &jobID)
})
if err != nil {
log.Printf("Failed to get task %d info: %v", taskUpdate.TaskID, err)
@@ -1597,7 +1635,10 @@ func (s *Manager) handleWebSocketTaskComplete(runnerID int64, taskUpdate WSTaskU
// Handle successful completion
if taskUpdate.Success {
err = s.db.WithTx(func(tx *sql.Tx) error {
_, err := tx.Exec(`UPDATE tasks SET status = ? WHERE id = ?`, types.TaskStatusCompleted, taskUpdate.TaskID)
_, err := tx.Exec(
`UPDATE tasks SET status = ?, error_message = NULL, completed_at = ? WHERE id = ?`,
types.TaskStatusCompleted, now, taskUpdate.TaskID,
)
if err != nil {
return err
}
@@ -1607,8 +1648,7 @@ func (s *Manager) handleWebSocketTaskComplete(runnerID int64, taskUpdate WSTaskU
return err
}
}
_, err = tx.Exec(`UPDATE tasks SET completed_at = ? WHERE id = ?`, now, taskUpdate.TaskID)
return err
return nil
})
if err != nil {
log.Printf("Failed to update task %d: %v", taskUpdate.TaskID, err)
@@ -1617,89 +1657,231 @@ func (s *Manager) handleWebSocketTaskComplete(runnerID int64, taskUpdate WSTaskU
// Broadcast task update
s.broadcastTaskUpdate(jobID, taskUpdate.TaskID, "task_update", map[string]interface{}{
"status": types.TaskStatusCompleted,
"output_path": taskUpdate.OutputPath,
"completed_at": now,
"status": types.TaskStatusCompleted,
"output_path": taskUpdate.OutputPath,
"completed_at": now,
"error_message": nil,
})
s.updateJobStatusFromTasks(jobID)
return
}
// Handle task failure - this is an actual task failure (e.g., Blender crash)
// Check if we have retries remaining
if retryCount < maxRetries {
// Reset to pending for retry - increment retry_count
// Actual task failure (e.g. Blender segfault) — retry when retries remain.
// free_requeue: failure newly armed GPU lockout; requeue without burning a retry.
errorMsg := taskUpdate.Error
if errorMsg == "" {
errorMsg = s.inferTaskFailureFromLogs(taskUpdate.TaskID)
}
freeRequeue := taskUpdate.FreeRequeue || s.taskLogsIndicateGPULockoutArm(taskUpdate.TaskID)
s.handleTaskFailureWithRetry(runnerID, taskUpdate.TaskID, jobID, errorMsg, freeRequeue)
}
// defaultUnexpectedDisconnectError is used when the runner drops the job WebSocket
// without a task_complete message and without any ERROR logs.
const defaultUnexpectedDisconnectError = "WebSocket connection lost"
// parseWSTaskUpdate decodes a task_complete payload. Older runners marshaled
// the error field as a Go error interface (JSON object {}), which fails strict
// unmarshaling into a string — fall back to extracting fields loosely so
// success/failure is still recognized.
func parseWSTaskUpdate(data json.RawMessage) (WSTaskUpdate, error) {
var update WSTaskUpdate
if err := json.Unmarshal(data, &update); err == nil {
return update, nil
}
var loose map[string]interface{}
if err := json.Unmarshal(data, &loose); err != nil {
return WSTaskUpdate{}, err
}
if v, ok := loose["task_id"].(float64); ok {
update.TaskID = int64(v)
}
if v, ok := loose["success"].(bool); ok {
update.Success = v
}
if v, ok := loose["output_path"].(string); ok {
update.OutputPath = v
}
if v, ok := loose["status"].(string); ok {
update.Status = v
}
if v, ok := loose["free_requeue"].(bool); ok {
update.FreeRequeue = v
}
switch v := loose["error"].(type) {
case string:
update.Error = v
case map[string]interface{}:
// Legacy runner bug: error interface marshaled as {}
update.Error = ""
}
return update, nil
}
// taskLogsIndicateGPULockoutArm is true when this attempt's logs show the runner
// newly armed GPU lockout (free-requeue path when task_complete lacked free_requeue).
func (s *Manager) taskLogsIndicateGPULockoutArm(taskID int64) bool {
var n int
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
`SELECT COUNT(*) FROM task_logs
WHERE task_id = ?
AND (
message LIKE '%GPU error detected%'
OR message LIKE '%free-requeue%'
OR message LIKE '%GPU disabled for subsequent jobs%'
)`,
taskID,
).Scan(&n)
})
return err == nil && n > 0
}
// inferTaskFailureFromLogs returns the best failure reason for a task that died
// without a usable task_complete payload. Prefers the latest ERROR log so Blender
// segfaults (and similar crashes) are stored instead of a generic disconnect message.
func (s *Manager) inferTaskFailureFromLogs(taskID int64) string {
var lastError string
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
`SELECT message FROM task_logs
WHERE task_id = ? AND UPPER(log_level) = ?
ORDER BY id DESC LIMIT 1`,
taskID, string(types.LogLevelError),
).Scan(&lastError)
})
if err != nil || strings.TrimSpace(lastError) == "" {
return defaultUnexpectedDisconnectError
}
return formatInferredTaskFailure(lastError)
}
// formatInferredTaskFailure normalizes streamed runner ERROR lines into a concise
// task error_message. Blender crashes often arrive as "Task failed: blender failed: signal: ..."
// or "Blender render failed: ..."; keep the underlying failure text.
func formatInferredTaskFailure(logMessage string) string {
msg := strings.TrimSpace(logMessage)
for _, prefix := range []string{
"Task failed: ",
"Blender render failed: ",
} {
if strings.HasPrefix(msg, prefix) {
msg = strings.TrimSpace(strings.TrimPrefix(msg, prefix))
}
}
if msg == "" {
return defaultUnexpectedDisconnectError
}
return msg
}
// handleTaskFailureWithRetry applies the task failure policy used for Blender
// crashes and unexpected job-WebSocket drops: requeue as pending while
// retry_count < max_retries, otherwise mark failed permanently with error_message.
// When freeRequeue is true (GPU lockout newly armed this attempt), the task is
// always requeued as pending without incrementing retry_count.
func (s *Manager) handleTaskFailureWithRetry(runnerID, taskID, jobID int64, errorMsg string, freeRequeue bool) {
var retryCount, maxRetries int
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
"SELECT retry_count, max_retries FROM tasks WHERE id = ?",
taskID,
).Scan(&retryCount, &maxRetries)
})
if err != nil {
log.Printf("Failed to get retry info for task %d: %v", taskID, err)
return
}
if errorMsg == "" {
errorMsg = defaultUnexpectedDisconnectError
}
if freeRequeue {
// Make the reason obvious in the UI without looking like a permanent fail.
if !strings.Contains(strings.ToLower(errorMsg), "gpu lockout") {
errorMsg = "GPU lockout armed (free requeue): " + errorMsg
}
}
now := time.Now()
// Free requeue always goes back to pending (expected ROCm GPU fault that armed lockout).
// Otherwise only requeue while retries remain.
canRequeue := freeRequeue || retryCount < maxRetries
if canRequeue {
newRetryCount := retryCount
retrySQL := `UPDATE tasks SET status = ?, runner_id = NULL, current_step = NULL,
error_message = ?, started_at = NULL, completed_at = NULL
WHERE id = ?`
args := []interface{}{types.TaskStatusPending, errorMsg, taskID}
if !freeRequeue {
retrySQL = `UPDATE tasks SET status = ?, runner_id = NULL, current_step = NULL,
error_message = ?, retry_count = retry_count + 1, started_at = NULL, completed_at = NULL
WHERE id = ?`
newRetryCount = retryCount + 1
}
err = s.db.WithTx(func(tx *sql.Tx) error {
_, err := tx.Exec(
`UPDATE tasks SET status = ?, runner_id = NULL, current_step = NULL,
retry_count = retry_count + 1, started_at = NULL, completed_at = NULL
WHERE id = ?`,
types.TaskStatusPending, taskUpdate.TaskID,
)
_, err := tx.Exec(retrySQL, args...)
if err != nil {
return err
}
// Clear steps and logs for fresh retry
_, err = tx.Exec(`DELETE FROM task_steps WHERE task_id = ?`, taskUpdate.TaskID)
_, err = tx.Exec(`DELETE FROM task_steps WHERE task_id = ?`, taskID)
if err != nil {
return err
}
_, err = tx.Exec(`DELETE FROM task_logs WHERE task_id = ?`, taskUpdate.TaskID)
_, err = tx.Exec(`DELETE FROM task_logs WHERE task_id = ?`, taskID)
return err
})
if err != nil {
log.Printf("Failed to reset task %d for retry: %v", taskUpdate.TaskID, err)
log.Printf("Failed to reset task %d for retry: %v", taskID, err)
return
}
// Broadcast task reset to clients (includes steps_cleared and logs_cleared flags)
s.broadcastTaskUpdate(jobID, taskUpdate.TaskID, "task_reset", map[string]interface{}{
s.broadcastTaskUpdate(jobID, taskID, "task_reset", map[string]interface{}{
"status": types.TaskStatusPending,
"retry_count": retryCount + 1,
"error_message": taskUpdate.Error,
"retry_count": newRetryCount,
"error_message": errorMsg,
"steps_cleared": true,
"logs_cleared": true,
"free_requeue": freeRequeue,
})
log.Printf("Task %d failed but has retries remaining (%d/%d), reset to pending", taskUpdate.TaskID, retryCount+1, maxRetries)
if freeRequeue {
log.Printf("Task %d failed (%s); free-requeued without burning retry (still %d/%d used)",
taskID, errorMsg, retryCount, maxRetries)
} else {
log.Printf("Task %d failed (%s) but has retries remaining (%d/%d), reset to pending",
taskID, errorMsg, newRetryCount, maxRetries)
}
} else {
// No retries remaining - mark as failed
err = s.db.WithTx(func(tx *sql.Tx) error {
_, err := tx.Exec(
`UPDATE tasks SET status = ?, runner_id = NULL, completed_at = ? WHERE id = ?`,
types.TaskStatusFailed, now, taskUpdate.TaskID,
`UPDATE tasks SET status = ?, runner_id = NULL, completed_at = ?, error_message = ? WHERE id = ?`,
types.TaskStatusFailed, now, errorMsg, taskID,
)
if err != nil {
return err
}
if taskUpdate.Error != "" {
_, err = tx.Exec(`UPDATE tasks SET error_message = ? WHERE id = ?`, taskUpdate.Error, taskUpdate.TaskID)
if err != nil {
return err
}
}
return nil
return err
})
if err != nil {
log.Printf("Failed to mark task %d as failed: %v", taskUpdate.TaskID, err)
log.Printf("Failed to mark task %d as failed: %v", taskID, err)
return
}
// Log the final failure
s.logTaskEvent(taskUpdate.TaskID, &runnerID, types.LogLevelError,
fmt.Sprintf("Task failed permanently after %d retries: %s", maxRetries, taskUpdate.Error), "")
s.logTaskEvent(taskID, &runnerID, types.LogLevelError,
fmt.Sprintf("Task failed permanently after %d retries: %s", maxRetries, errorMsg), "")
// Broadcast task update
s.broadcastTaskUpdate(jobID, taskUpdate.TaskID, "task_update", map[string]interface{}{
s.broadcastTaskUpdate(jobID, taskID, "task_update", map[string]interface{}{
"status": types.TaskStatusFailed,
"completed_at": now,
"error_message": taskUpdate.Error,
"error_message": errorMsg,
})
log.Printf("Task %d failed permanently after %d retries", taskUpdate.TaskID, maxRetries)
log.Printf("Task %d failed permanently after %d retries: %s", taskID, maxRetries, errorMsg)
}
// Update job status and progress
s.updateJobStatusFromTasks(jobID)
}
+80 -1
View File
@@ -1,6 +1,9 @@
package api
import "testing"
import (
"encoding/json"
"testing"
)
func TestParseBlenderFrame(t *testing.T) {
frame, ok := parseBlenderFrame("Info Fra:2470 Mem:12.00M")
@@ -19,3 +22,79 @@ func TestJobTaskCounts_Progress(t *testing.T) {
}
}
func TestParseWSTaskUpdate_LegacyErrorObject(t *testing.T) {
// Simulates older runners that JSON-marshaled an error interface as {}
raw := json.RawMessage(`{"task_id":99,"success":false,"error":{}}`)
update, err := parseWSTaskUpdate(raw)
if err != nil {
t.Fatalf("parseWSTaskUpdate: %v", err)
}
if update.TaskID != 99 {
t.Fatalf("TaskID = %d, want 99", update.TaskID)
}
if update.Success {
t.Fatal("Success = true, want false")
}
if update.Error != "" {
t.Fatalf("Error = %q, want empty (legacy object)", update.Error)
}
}
func TestParseWSTaskUpdate_StringError(t *testing.T) {
raw := json.RawMessage(`{"task_id":7,"success":false,"error":"blender failed: signal: segmentation fault (core dumped)","free_requeue":true}`)
update, err := parseWSTaskUpdate(raw)
if err != nil {
t.Fatalf("parseWSTaskUpdate: %v", err)
}
if update.TaskID != 7 || update.Success {
t.Fatalf("unexpected update: %+v", update)
}
if update.Error != "blender failed: signal: segmentation fault (core dumped)" {
t.Fatalf("Error = %q", update.Error)
}
if !update.FreeRequeue {
t.Fatal("expected FreeRequeue true")
}
}
func TestFormatInferredTaskFailure(t *testing.T) {
tests := []struct {
name string
in string
want string
}{
{
name: "task failed segfault prefix",
in: "Task failed: blender failed: signal: segmentation fault (core dumped)",
want: "blender failed: signal: segmentation fault (core dumped)",
},
{
name: "blender render failed prefix",
in: "Blender render failed: blender failed: signal: segmentation fault (core dumped)",
want: "blender failed: signal: segmentation fault (core dumped)",
},
{
name: "double prefix",
in: "Task failed: Blender render failed: blender failed: signal: segmentation fault (core dumped)",
want: "blender failed: signal: segmentation fault (core dumped)",
},
{
name: "plain message",
in: "blender failed: exit status 1",
want: "blender failed: exit status 1",
},
{
name: "empty falls back",
in: " ",
want: defaultUnexpectedDisconnectError,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := formatInferredTaskFailure(tt.in); got != tt.want {
t.Fatalf("formatInferredTaskFailure(%q) = %q, want %q", tt.in, got, tt.want)
}
})
}
}
+19 -7
View File
@@ -301,7 +301,12 @@ func (j *JobConnection) OutputUploaded(taskID int64, fileName string) {
}
// Complete sends task completion to the manager.
func (j *JobConnection) Complete(taskID int64, success bool, errorMsg error) {
// errorMsg must be JSON-encoded as a string: encoding an error interface
// marshals to {} and the manager fails to unmarshal task_complete, which
// previously surfaced as a generic "WebSocket connection lost" with no retries.
// freeRequeue asks the manager to requeue a failure without incrementing retry_count
// (used when this attempt newly armed GPU lockout).
func (j *JobConnection) Complete(taskID int64, success bool, errorMsg error, freeRequeue bool) {
if j.conn == nil {
log.Printf("Cannot send task complete: WebSocket connection is nil")
return
@@ -310,13 +315,20 @@ func (j *JobConnection) Complete(taskID int64, success bool, errorMsg error) {
j.writeMu.Lock()
defer j.writeMu.Unlock()
data := map[string]interface{}{
"task_id": taskID,
"success": success,
}
if errorMsg != nil {
data["error"] = errorMsg.Error()
}
if freeRequeue {
data["free_requeue"] = true
}
msg := map[string]interface{}{
"type": "task_complete",
"data": map[string]interface{}{
"task_id": taskID,
"success": success,
"error": errorMsg,
},
"type": "task_complete",
"data": data,
"timestamp": time.Now().Unix(),
}
if err := j.conn.WriteJSON(msg); err != nil {
+75
View File
@@ -1,6 +1,8 @@
package api
import (
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"strings"
@@ -10,6 +12,79 @@ import (
"github.com/gorilla/websocket"
)
// TestComplete_ErrorIsJSONString ensures task_complete encodes the failure
// reason as a string. Marshaling a Go error interface produces {}, which the
// manager cannot unmarshal into WSTaskUpdate.Error and used to look like a
// silent WebSocket drop with no retries.
func TestComplete_ErrorIsJSONString(t *testing.T) {
upgrader := websocket.Upgrader{}
received := make(chan map[string]interface{}, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
defer conn.Close()
// auth handshake
var auth map[string]interface{}
if err := conn.ReadJSON(&auth); err != nil {
return
}
if auth["type"] == "auth" {
_ = conn.WriteJSON(map[string]string{"type": "auth_ok"})
}
var msg map[string]interface{}
if err := conn.ReadJSON(&msg); err != nil {
return
}
received <- msg
}))
defer server.Close()
jc := NewJobConnection()
if err := jc.Connect(server.URL, "/job/1", "token123"); err != nil {
t.Fatalf("Connect failed: %v", err)
}
defer jc.Close()
jc.Complete(42, false, errors.New("blender failed: signal: segmentation fault (core dumped)"), true)
select {
case msg := <-received:
if msg["type"] != "task_complete" {
t.Fatalf("type = %v, want task_complete", msg["type"])
}
data, ok := msg["data"].(map[string]interface{})
if !ok {
// gorilla may leave nested objects as map[string]interface{} after JSON round-trip;
// also accept raw re-marshal path
raw, _ := json.Marshal(msg["data"])
data = map[string]interface{}{}
if err := json.Unmarshal(raw, &data); err != nil {
t.Fatalf("data type %T: %v", msg["data"], err)
}
}
errVal, ok := data["error"].(string)
if !ok {
t.Fatalf("error field type %T value %#v; want string (not object)", data["error"], data["error"])
}
if errVal != "blender failed: signal: segmentation fault (core dumped)" {
t.Fatalf("error = %q, want segfault message", errVal)
}
if data["success"] != false {
t.Fatalf("success = %v, want false", data["success"])
}
if data["free_requeue"] != true {
t.Fatalf("free_requeue = %v, want true", data["free_requeue"])
}
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for task_complete message")
}
}
func TestJobConnection_ConnectAndClose(t *testing.T) {
upgrader := websocket.Upgrader{}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+15
View File
@@ -114,3 +114,18 @@ func isGPUControllerLine(line string) bool {
strings.Contains(line, "3d controller") ||
strings.Contains(line, "display controller")
}
func parseRocmAgentArch(output string) (arch string, ok bool) {
scanner := bufio.NewScanner(strings.NewReader(output))
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if !strings.HasPrefix(line, "Name:") {
continue
}
name := strings.TrimSpace(strings.TrimPrefix(line, "Name:"))
if strings.HasPrefix(name, "gfx") {
return name, true
}
}
return "", false
}
@@ -0,0 +1,19 @@
package blender
import "testing"
func TestParseRocmAgentArch(t *testing.T) {
input := `ROCk module is loaded
Agent 1
Name: gfx1151
Marketing Name: AMD Radeon Graphics
`
arch, ok := parseRocmAgentArch(input)
if !ok {
t.Fatal("expected arch to be parsed")
}
if arch != "gfx1151" {
t.Fatalf("arch = %q, want gfx1151", arch)
}
}
+16 -56
View File
@@ -19,26 +19,11 @@ const (
CRFVP9 = 30
)
// tonemapFilter returns the appropriate filter for EXR input.
// For HDR preservation: converts linear RGB (EXR) to bt2020 YUV with HLG transfer function
// Uses zscale to properly convert colorspace from linear RGB to bt2020 YUV while preserving HDR range
// Step 1: Ensure format is gbrpf32le (linear RGB)
// Step 2: Convert transfer function from linear to HLG (arib-std-b67) with bt2020 primaries/matrix
// Step 3: Convert to YUV format
// tonemapFilter returns the production filter for EXR input (linear → sRGB → HLG, bt709).
// This is the single source of truth used by BuildCommand and BuildPass1Command.
// zscale numeric values: primaries 1=bt709, matrix 1=bt709, transfer 8=linear / 13=sRGB / 18=HLG.
func tonemapFilter(useAlpha bool) string {
// Convert from linear RGB (gbrpf32le) to HLG with bt709 primaries to match PNG appearance
// Based on best practices: convert linear RGB directly to HLG with bt709 primaries
// This matches PNG color appearance (bt709 primaries) while preserving HDR range (HLG transfer)
// zscale uses numeric values:
// primaries: 1=bt709 (matches PNG), 9=bt2020
// matrix: 1=bt709, 9=bt2020nc, 0=gbr (RGB input)
// transfer: 8=linear, 18=arib-std-b67 (HLG)
// Direct conversion: linear RGB -> HLG with bt709 primaries -> bt2020 YUV (for wider gamut metadata)
// The bt709 primaries in the conversion match PNG, but we set bt2020 in metadata for HDR displays
// Convert linear RGB to sRGB first, then convert to HLG
// This approach: linear -> sRGB -> HLG -> bt2020
// Fixes red tint by using sRGB conversion, preserves HDR range with HLG
filter := "format=gbrpf32le,zscale=transferin=8:transfer=13:primariesin=1:primaries=1:matrixin=0:matrix=1:rangein=full:range=full,zscale=transferin=13:transfer=18:primariesin=1:primaries=9:matrixin=1:matrix=9:rangein=full:range=full"
filter := "format=gbrpf32le,zscale=transferin=8:transfer=13:primariesin=1:primaries=1:matrixin=0:matrix=1:rangein=full:range=full,zscale=transferin=13:transfer=18:primariesin=1:primaries=1:matrixin=1:matrix=1:rangein=full:range=full"
if useAlpha {
return filter + ",format=yuva420p10le"
}
@@ -57,8 +42,7 @@ func (e *SoftwareEncoder) Available() bool {
return true // Software encoding is always available
}
func (e *SoftwareEncoder) BuildCommand(config *EncodeConfig) *exec.Cmd {
// EXR only: HDR path (HLG, 10-bit, full range)
func (e *SoftwareEncoder) buildBaseArgs(config *EncodeConfig) []string {
pixFmt := "yuv420p10le"
if config.UseAlpha {
pixFmt = "yuva420p10le"
@@ -79,14 +63,18 @@ func (e *SoftwareEncoder) BuildCommand(config *EncodeConfig) *exec.Cmd {
"-color_trc", "linear", "-color_primaries", "bt709"}
args = append(args, "-i", config.InputPattern, "-c:v", e.codec, "-pix_fmt", pixFmt, "-r", fmt.Sprintf("%.2f", config.FrameRate), "-color_primaries", colorPrimaries, "-color_trc", colorTrc, "-colorspace", colorspace, "-color_range", colorRange)
vf := "format=gbrpf32le,zscale=transferin=8:transfer=13:primariesin=1:primaries=1:matrixin=0:matrix=1:rangein=full:range=full,zscale=transferin=13:transfer=18:primariesin=1:primaries=1:matrixin=1:matrix=1:rangein=full:range=full"
if config.UseAlpha {
vf += ",format=yuva420p10le"
} else {
vf += ",format=yuv420p10le"
// FFmpeg 6+ treats yuva420p10le as experimental for libvpx-vp9 / libaom-av1
if config.UseAlpha && (e.codec == "libvpx-vp9" || e.codec == "libaom-av1") {
args = append(args, "-strict", "experimental")
}
args = append(args, "-vf", vf)
args = append(args, "-vf", tonemapFilter(config.UseAlpha))
args = append(args, codecArgs...)
return args
}
func (e *SoftwareEncoder) BuildCommand(config *EncodeConfig) *exec.Cmd {
args := e.buildBaseArgs(config)
if config.TwoPass {
// For 2-pass, this builds pass 2 command
@@ -107,35 +95,7 @@ func (e *SoftwareEncoder) BuildCommand(config *EncodeConfig) *exec.Cmd {
// BuildPass1Command builds the first pass command for 2-pass encoding.
func (e *SoftwareEncoder) BuildPass1Command(config *EncodeConfig) *exec.Cmd {
pixFmt := "yuv420p10le"
if config.UseAlpha {
pixFmt = "yuva420p10le"
}
colorPrimaries, colorTrc, colorspace, colorRange := "bt709", "arib-std-b67", "bt709", "pc"
var codecArgs []string
switch e.codec {
case "libaom-av1":
codecArgs = []string{"-crf", strconv.Itoa(CRFAV1), "-b:v", "0", "-tiles", "2x2", "-g", "240"}
case "libvpx-vp9":
codecArgs = []string{"-crf", strconv.Itoa(CRFVP9), "-b:v", "0", "-row-mt", "1", "-g", "240"}
default:
codecArgs = []string{"-preset", "veryslow", "-crf", strconv.Itoa(CRFH264), "-profile:v", "high10", "-level", "5.2", "-tune", "film", "-keyint_min", "24", "-g", "240", "-bf", "2", "-refs", "4"}
}
args := []string{"-y", "-f", "image2", "-start_number", fmt.Sprintf("%d", config.StartFrame), "-framerate", fmt.Sprintf("%.2f", config.FrameRate),
"-color_trc", "linear", "-color_primaries", "bt709"}
args = append(args, "-i", config.InputPattern, "-c:v", e.codec, "-pix_fmt", pixFmt, "-r", fmt.Sprintf("%.2f", config.FrameRate), "-color_primaries", colorPrimaries, "-color_trc", colorTrc, "-colorspace", colorspace, "-color_range", colorRange)
vf := "format=gbrpf32le,zscale=transferin=8:transfer=13:primariesin=1:primaries=1:matrixin=0:matrix=1:rangein=full:range=full,zscale=transferin=13:transfer=18:primariesin=1:primaries=1:matrixin=1:matrix=1:rangein=full:range=full"
if config.UseAlpha {
vf += ",format=yuva420p10le"
} else {
vf += ",format=yuv420p10le"
}
args = append(args, "-vf", vf)
args = append(args, codecArgs...)
args := e.buildBaseArgs(config)
args = append(args, "-pass", "1", "-f", "null", "/dev/null")
log.Printf("Build Software Pass 1 command: ffmpeg %s", strings.Join(args, " "))
@@ -120,6 +120,10 @@ func TestSoftwareEncoder_BuildCommand_AV1_WithAlpha(t *testing.T) {
if !strings.Contains(argsStr, "format=yuva420p10le") {
t.Error("Expected tonemap filter to output yuva420p10le for EXR alpha, but not found")
}
// Alpha VP9/AV1 need -strict experimental on modern FFmpeg
if !strings.Contains(argsStr, "-strict experimental") {
t.Error("Expected -strict experimental for alpha encode")
}
}
func TestSoftwareEncoder_BuildCommand_VP9(t *testing.T) {
+154 -18
View File
@@ -10,6 +10,7 @@ import (
"net"
"os"
"os/exec"
"path/filepath"
"strings"
"sync"
"time"
@@ -17,6 +18,7 @@ import (
"jiggablend/internal/runner/api"
"jiggablend/internal/runner/blender"
"jiggablend/internal/runner/encoding"
"jiggablend/internal/runner/sandbox"
"jiggablend/internal/runner/tasks"
"jiggablend/internal/runner/workspace"
"jiggablend/pkg/executils"
@@ -58,12 +60,56 @@ type Runner struct {
// forceCPURendering forces CPU rendering for all jobs regardless of metadata/backend detection.
forceCPURendering bool
// disableRT disables GPU ray tracing acceleration (HIPRT, OptiX, etc.).
disableRT bool
// hipGPUSampleBatch limits samples per GPU pass on gfx115x (0 = disabled).
hipGPUSampleBatch int
// sandbox wraps Blender invocations (none/podman).
sandboxWrapper sandbox.Wrapper
sandboxBackend string
}
// RunnerOptions configures optional runner behavior.
type RunnerOptions struct {
ForceCPURendering bool
DisableRT bool
HipGPUSampleBatch int
SandboxBackend string // none|podman
SandboxNetwork bool
SandboxImage string // podman thin runtime image
}
// New creates a new runner.
func New(managerURL, name, hostname string, forceCPURendering bool) *Runner {
func New(managerURL, name, hostname string, forceCPURendering, disableRT bool, hipGPUSampleBatch int) *Runner {
return NewWithOptions(managerURL, name, hostname, RunnerOptions{
ForceCPURendering: forceCPURendering,
DisableRT: disableRT,
HipGPUSampleBatch: hipGPUSampleBatch,
SandboxBackend: sandbox.BackendPodman,
})
}
// NewWithOptions creates a runner with full options including sandbox.
func NewWithOptions(managerURL, name, hostname string, opts RunnerOptions) *Runner {
manager := api.NewManagerClient(managerURL)
backend, err := sandbox.NormalizeBackend(opts.SandboxBackend)
if err != nil {
log.Printf("Invalid sandbox backend %q, falling back to none: %v", opts.SandboxBackend, err)
backend = sandbox.BackendNone
}
sb, err := sandbox.New(sandbox.Options{
Backend: backend,
AllowNetwork: opts.SandboxNetwork,
PodmanImage: opts.SandboxImage,
})
if err != nil {
log.Printf("Failed to init sandbox %q: %v; using none", backend, err)
sb, _ = sandbox.New(sandbox.Options{Backend: sandbox.BackendNone})
backend = sandbox.BackendNone
}
r := &Runner{
name: name,
hostname: hostname,
@@ -72,7 +118,11 @@ func New(managerURL, name, hostname string, forceCPURendering bool) *Runner {
stopChan: make(chan struct{}),
processors: make(map[string]tasks.Processor),
forceCPURendering: forceCPURendering,
forceCPURendering: opts.ForceCPURendering,
disableRT: opts.DisableRT,
hipGPUSampleBatch: opts.HipGPUSampleBatch,
sandboxWrapper: sb,
sandboxBackend: backend,
}
// Generate fingerprint
@@ -88,9 +138,24 @@ func (r *Runner) CheckRequiredTools() error {
}
log.Printf("Found zstd for compressed blend file support")
if r.sandboxWrapper != nil {
if err := r.sandboxWrapper.Available(); err != nil {
return fmt.Errorf("sandbox backend %q unavailable: %w", r.sandboxBackend, err)
}
log.Printf("Sandbox backend: %s", r.sandboxWrapper.Name())
}
return nil
}
// SandboxBackend returns the configured sandbox backend name.
func (r *Runner) SandboxBackend() string {
if r.sandboxBackend == "" {
return sandbox.BackendNone
}
return r.sandboxBackend
}
var (
cachedCapabilities map[string]interface{}
capabilitiesOnce sync.Once
@@ -107,8 +172,21 @@ func (r *Runner) ProbeCapabilities() map[string]interface{} {
caps["ffmpeg"] = false
}
// Filled later with actual backend when runner is constructed; probe is package-level once.
// Real sandbox name is injected in ProbeCapabilities on the instance after New.
caps["sandbox"] = "none"
cachedCapabilities = caps
})
// Overlay instance sandbox name (Once already ran with none default).
if r != nil && r.sandboxWrapper != nil {
out := make(map[string]interface{}, len(cachedCapabilities)+1)
for k, v := range cachedCapabilities {
out[k] = v
}
out["sandbox"] = r.sandboxWrapper.Name()
return out
}
return cachedCapabilities
}
@@ -198,6 +276,16 @@ func (r *Runner) HasIntel() bool {
return r.hasIntel
}
// DisableRT returns whether GPU ray tracing acceleration should be disabled.
func (r *Runner) DisableRT() bool {
return r.disableRT
}
// HipGPUSampleBatch returns the per-pass GPU sample limit for gfx115x batching (0 = disabled).
func (r *Runner) HipGPUSampleBatch() int {
return r.hipGPUSampleBatch
}
// GPUDetectionFailed returns true when startup GPU backend detection could not run or failed. When true, all jobs use CPU because backend availability is unknown.
func (r *Runner) GPUDetectionFailed() bool {
r.gpuBackendMu.RLock()
@@ -330,8 +418,20 @@ func (r *Runner) executeJob(job *api.NextJobResponse) (err error) {
r.HasIntel(),
r.GPUDetectionFailed(),
r.forceCPURendering,
func() { r.SetGPULockedOut(true) },
r.disableRT,
r.hipGPUSampleBatch,
nil, // set below so the callback can mark this attempt
r.sandboxWrapper,
)
// Arm GPU lockout at most once process-wide; if this attempt is the one that
// arms it, mark the context so a failure requeues without burning retry_count.
ctx.OnGPUError = func() {
if r.SetGPULockedOut(true) {
ctx.GPULockoutArmedThisAttempt = true
ctx.GPULockedOut = true
ctx.Warn("GPU error detected; GPU disabled for subsequent jobs (this attempt free-requeues without using a retry)")
}
}
ctx.Info(fmt.Sprintf("Task assignment received (job: %d, type: %s)",
job.Task.JobID, job.Task.TaskType))
@@ -350,7 +450,7 @@ func (r *Runner) executeJob(job *api.NextJobResponse) (err error) {
contextPath := job.JobPath + "/context.tar"
if err := r.downloadContext(job.Task.JobID, contextPath, job.JobToken); err != nil {
jobConn.Log(job.Task.TaskID, types.LogLevelError, fmt.Sprintf("Failed to download context: %v", err))
jobConn.Complete(job.Task.TaskID, false, fmt.Errorf("failed to download context: %v", err))
jobConn.Complete(job.Task.TaskID, false, fmt.Errorf("failed to download context: %v", err), false)
return fmt.Errorf("failed to download context: %w", err)
}
processErr = processor.Process(ctx)
@@ -393,27 +493,56 @@ func (r *Runner) uploadOutputs(ctx *tasks.Context, job *api.NextJobResponse) err
outputDir := ctx.WorkDir + "/output"
uploadPath := fmt.Sprintf("/api/runner/jobs/%d/upload", job.Task.JobID)
return uploadOutputFiles(outputDir, func(filePath, fileName string) error {
if err := r.manager.UploadFile(uploadPath, job.JobToken, filePath); err != nil {
return err
}
ctx.OutputUploaded(fileName)
// Delete file after successful upload to prevent duplicate uploads
if err := os.Remove(filePath); err != nil {
log.Printf("Warning: Failed to delete file %s after upload: %v", filePath, err)
}
return nil
})
}
// uploadOutputFiles uploads all non-directory files from outputDir.
// Fails if the directory cannot be read, any upload fails, or zero files were uploaded.
func uploadOutputFiles(outputDir string, uploadFn func(filePath, fileName string) error) error {
entries, err := os.ReadDir(outputDir)
if err != nil {
return fmt.Errorf("failed to read output directory: %w", err)
}
var files []os.DirEntry
for _, entry := range entries {
if entry.IsDir() {
continue
}
filePath := outputDir + "/" + entry.Name()
if err := r.manager.UploadFile(uploadPath, job.JobToken, filePath); err != nil {
log.Printf("Failed to upload %s: %v", filePath, err)
} else {
ctx.OutputUploaded(entry.Name())
// Delete file after successful upload to prevent duplicate uploads
if err := os.Remove(filePath); err != nil {
log.Printf("Warning: Failed to delete file %s after upload: %v", filePath, err)
}
if !entry.IsDir() {
files = append(files, entry)
}
}
if len(files) == 0 {
return fmt.Errorf("no output files found in %s", outputDir)
}
var firstErr error
uploaded := 0
for _, entry := range files {
filePath := filepath.Join(outputDir, entry.Name())
if err := uploadFn(filePath, entry.Name()); err != nil {
log.Printf("Failed to upload %s: %v", filePath, err)
if firstErr == nil {
firstErr = fmt.Errorf("failed to upload %s: %w", entry.Name(), err)
}
continue
}
uploaded++
}
if firstErr != nil {
return firstErr
}
if uploaded == 0 {
return fmt.Errorf("no output files were uploaded from %s", outputDir)
}
return nil
}
@@ -484,13 +613,20 @@ func (r *Runner) GetID() int64 {
// SetGPULockedOut sets whether GPU use is locked out due to a detected GPU error.
// When true, the runner will force CPU rendering for all jobs.
func (r *Runner) SetGPULockedOut(locked bool) {
// Returns true only on the false→true transition (first arm); subsequent calls are no-ops for logging.
func (r *Runner) SetGPULockedOut(locked bool) (newlyEnabled bool) {
r.gpuLockedOutMu.Lock()
defer r.gpuLockedOutMu.Unlock()
r.gpuLockedOut = locked
if locked {
if r.gpuLockedOut {
return false
}
r.gpuLockedOut = true
log.Printf("GPU lockout enabled: GPU rendering disabled for subsequent jobs (CPU only)")
return true
}
r.gpuLockedOut = false
return false
}
// IsGPULockedOut returns whether GPU use is currently locked out.
+60 -4
View File
@@ -2,11 +2,14 @@ package runner
import (
"encoding/hex"
"errors"
"os"
"path/filepath"
"testing"
)
func TestNewRunner_InitializesFields(t *testing.T) {
r := New("http://localhost:8080", "runner-a", "host-a", false)
r := New("http://localhost:8080", "runner-a", "host-a", false, false, 0)
if r == nil {
t.Fatal("New should return a runner")
}
@@ -16,15 +19,26 @@ func TestNewRunner_InitializesFields(t *testing.T) {
}
func TestRunner_GPUFlagsSetters(t *testing.T) {
r := New("http://localhost:8080", "runner-a", "host-a", false)
r.SetGPULockedOut(true)
r := New("http://localhost:8080", "runner-a", "host-a", false, false, 0)
if newly := r.SetGPULockedOut(true); !newly {
t.Fatal("expected first SetGPULockedOut(true) to report newly enabled")
}
if !r.IsGPULockedOut() {
t.Fatal("expected GPU lockout to be true")
}
if newly := r.SetGPULockedOut(true); newly {
t.Fatal("expected second SetGPULockedOut(true) to be a no-op transition")
}
if newly := r.SetGPULockedOut(false); newly {
t.Fatal("clearing lockout should not report newly enabled")
}
if r.IsGPULockedOut() {
t.Fatal("expected GPU lockout cleared")
}
}
func TestGenerateFingerprint_PopulatesValue(t *testing.T) {
r := New("http://localhost:8080", "runner-a", "host-a", false)
r := New("http://localhost:8080", "runner-a", "host-a", false, false, 0)
r.generateFingerprint()
fp := r.GetFingerprint()
if fp == "" {
@@ -38,3 +52,45 @@ func TestGenerateFingerprint_PopulatesValue(t *testing.T) {
}
}
func TestUploadOutputFiles_FailsWhenUploadErrors(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "frame_0001.exr"), []byte("x"), 0644); err != nil {
t.Fatal(err)
}
err := uploadOutputFiles(dir, func(filePath, fileName string) error {
return errors.New("upload denied")
})
if err == nil {
t.Fatal("expected error when upload fails")
}
}
func TestUploadOutputFiles_FailsWhenEmpty(t *testing.T) {
dir := t.TempDir()
err := uploadOutputFiles(dir, func(filePath, fileName string) error {
t.Fatal("upload should not be called")
return nil
})
if err == nil {
t.Fatal("expected error for empty output dir")
}
}
func TestUploadOutputFiles_Success(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "frame_0001.exr"), []byte("x"), 0644); err != nil {
t.Fatal(err)
}
var seen string
err := uploadOutputFiles(dir, func(filePath, fileName string) error {
seen = fileName
return nil
})
if err != nil {
t.Fatalf("uploadOutputFiles: %v", err)
}
if seen != "frame_0001.exr" {
t.Fatalf("got %q", seen)
}
}
+173
View File
@@ -0,0 +1,173 @@
package sandbox
import (
"os"
"path/filepath"
"strings"
)
// Bind describes a host path to expose inside the sandbox.
type Bind struct {
Host string
// Dest is the path inside the sandbox (usually same as Host for absolute paths).
Dest string
// ReadOnly is true for library trees and Blender installs.
ReadOnly bool
// Dev is true for device nodes / device trees (podman --device or -v for dirs).
Dev bool
// Optional means skip if host path missing (no error).
Optional bool
}
// GPUMounts returns host paths to expose for Cycles GPU backends based on detection.
// forceCPU skips GPU-specific devices/libs (still may include generic /dev via backend).
func GPUMounts(hasAMD, hasNVIDIA, hasIntel, forceCPU bool) []Bind {
if forceCPU {
return nil
}
var binds []Bind
// DRM render nodes used by AMD, Intel, and some NVIDIA EGL paths.
if hasAMD || hasNVIDIA || hasIntel {
binds = append(binds, Bind{Host: "/dev/dri", Dest: "/dev/dri", Dev: true, Optional: true})
}
if hasAMD {
binds = append(binds, Bind{Host: "/dev/kfd", Dest: "/dev/kfd", Dev: true, Optional: true})
// Common ROCm install layouts
for _, p := range []string{"/opt/rocm", "/usr/share/libdrm"} {
binds = append(binds, Bind{Host: p, Dest: p, ReadOnly: true, Optional: true})
}
// Distro ROCm / amdgpu userspace libs often live under multiarch paths
for _, p := range rocmLibHints() {
binds = append(binds, Bind{Host: p, Dest: p, ReadOnly: true, Optional: true})
}
}
if hasNVIDIA {
for _, name := range nvidiaDeviceNames() {
p := filepath.Join("/dev", name)
binds = append(binds, Bind{Host: p, Dest: p, Dev: true, Optional: true})
}
for _, p := range nvidiaLibHints() {
binds = append(binds, Bind{Host: p, Dest: p, ReadOnly: true, Optional: true})
}
}
if hasIntel {
for _, p := range intelLibHints() {
binds = append(binds, Bind{Host: p, Dest: p, ReadOnly: true, Optional: true})
}
}
return uniqueBinds(binds)
}
func nvidiaDeviceNames() []string {
// Fixed control nodes + any /dev/nvidiaN
names := []string{"nvidiactl", "nvidia-uvm", "nvidia-uvm-tools", "nvidia-modeset"}
entries, err := os.ReadDir("/dev")
if err != nil {
return names
}
for _, e := range entries {
n := e.Name()
if strings.HasPrefix(n, "nvidia") && !containsString(names, n) {
names = append(names, n)
}
}
return names
}
func nvidiaLibHints() []string {
hints := []string{
"/usr/lib/wsl/lib", // WSL NVIDIA
"/usr/local/nvidia",
"/usr/local/cuda",
"/usr/lib/x86_64-linux-gnu",
"/usr/lib64",
}
// Driver stores often under /usr/lib/libcuda* — parent dirs already covered.
// Also scan common multiarch for libcuda.so
for _, dir := range []string{"/usr/lib", "/usr/lib64", "/usr/lib/x86_64-linux-gnu", "/lib", "/lib64"} {
if matchesAny(dir, "libcuda.so*") || matchesAny(dir, "libnvidia-*.so*") {
hints = append(hints, dir)
}
}
return hints
}
func rocmLibHints() []string {
hints := []string{
"/usr/lib/x86_64-linux-gnu",
"/usr/lib64",
"/opt/amdgpu",
}
// ROCm versioned trees under /opt/rocm-*
if entries, err := os.ReadDir("/opt"); err == nil {
for _, e := range entries {
if e.IsDir() && strings.HasPrefix(e.Name(), "rocm") {
hints = append(hints, filepath.Join("/opt", e.Name()))
}
}
}
return hints
}
func intelLibHints() []string {
return []string{
"/usr/lib/x86_64-linux-gnu",
"/usr/lib64",
"/usr/lib/intel-opencl",
"/etc/OpenCL",
}
}
func matchesAny(dir, glob string) bool {
m, err := filepath.Glob(filepath.Join(dir, glob))
return err == nil && len(m) > 0
}
func containsString(ss []string, s string) bool {
for _, x := range ss {
if x == s {
return true
}
}
return false
}
func uniqueBinds(in []Bind) []Bind {
seen := make(map[string]bool)
var out []Bind
for _, b := range in {
key := b.Host + "|" + b.Dest + "|"
if b.Dev {
key += "d"
}
if b.ReadOnly {
key += "r"
}
if seen[key] {
continue
}
seen[key] = true
out = append(out, b)
}
return out
}
// ExistingBinds filters to binds whose host path exists (or non-optional always kept for error reporting).
func ExistingBinds(binds []Bind) []Bind {
var out []Bind
for _, b := range binds {
if pathExists(b.Host) {
out = append(out, b)
continue
}
if !b.Optional {
out = append(out, b) // caller may error
}
}
return out
}
+18
View File
@@ -0,0 +1,18 @@
package sandbox
import "os/exec"
type noneWrapper struct{}
func (w *noneWrapper) Name() string { return BackendNone }
func (w *noneWrapper) Available() error { return nil }
func (w *noneWrapper) Wrap(spec Spec) (*exec.Cmd, error) {
cmd := exec.Command(spec.BlenderBinary, spec.Args...)
cmd.Dir = spec.WorkDir
if len(spec.Env) > 0 {
cmd.Env = spec.Env
}
return cmd, nil
}
+13
View File
@@ -0,0 +1,13 @@
package sandbox
import "path/filepath"
func execAbs(p string) (string, error) {
return filepath.Abs(p)
}
// blenderRoot returns the directory that contains the blender binary
// (the version extract root, e.g. .../blender-versions/4.5.7).
func blenderRoot(blenderBinary string) string {
return filepath.Dir(mustAbs(blenderBinary))
}
+121
View File
@@ -0,0 +1,121 @@
package sandbox
import (
"fmt"
"os"
"os/exec"
"path/filepath"
"strconv"
"strings"
)
const defaultPodmanImage = "registry.fedoraproject.org/fedora-minimal:41"
type podmanWrapper struct {
allowNetwork bool
image string
}
func (w *podmanWrapper) Name() string { return BackendPodman }
func (w *podmanWrapper) Available() error {
if _, err := LookPath("podman"); err != nil {
return fmt.Errorf("podman not found in PATH: %w", err)
}
return nil
}
func (w *podmanWrapper) Wrap(spec Spec) (*exec.Cmd, error) {
if err := w.Available(); err != nil {
return nil, err
}
bin := mustAbs(spec.BlenderBinary)
work := mustAbs(spec.WorkDir)
home := mustAbs(spec.HomeDir)
if home == "" {
home = filepath.Join(work, "home")
}
root := blenderRoot(bin)
args := []string{
"run", "--rm",
"--security-opt", "label=disable",
// Keep user groups for /dev/dri access (render/video)
"--group-add", "keep-groups",
}
if !w.allowNetwork {
args = append(args, "--network=none")
}
// Map to same UID so bind mounts are writable
uid := os.Getuid()
gid := os.Getgid()
args = append(args, "--user", fmt.Sprintf("%d:%d", uid, gid))
// Thin OS image + host Blender tree + job dir
args = append(args,
"-v", root+":"+root+":ro",
"-v", work+":"+work+":rw",
)
if home != work && !strings.HasPrefix(home, work+string(os.PathSeparator)) {
_ = os.MkdirAll(home, 0755)
args = append(args, "-v", home+":"+home+":rw")
}
// System libs for GPU ICDs / dynamic linker (Blender tarball is mostly self-contained)
for _, p := range []string{"/usr", "/lib", "/lib64", "/etc"} {
if pathExists(p) {
args = append(args, "-v", p+":"+p+":ro")
}
}
// GPU devices and extra lib roots
for _, b := range ExistingBinds(GPUMounts(spec.HasAMD, spec.HasNVIDIA, spec.HasIntel, spec.ForceCPU)) {
if !pathExists(b.Host) {
continue
}
if b.Dev {
// --device works for character devices; for /dev/dri use volume
info, err := os.Stat(b.Host)
if err != nil {
continue
}
if info.IsDir() {
args = append(args, "-v", b.Host+":"+b.Dest+":ro")
} else {
args = append(args, "--device", b.Host)
}
continue
}
mode := "ro"
if !b.ReadOnly {
mode = "rw"
}
args = append(args, "-v", b.Host+":"+b.Dest+":"+mode)
}
// Environment
for _, e := range spec.Env {
args = append(args, "-e", e)
}
if home != "" {
args = append(args, "-e", "HOME="+home)
}
args = append(args, "--workdir", work)
args = append(args, "--entrypoint", bin)
args = append(args, w.image)
// entrypoint is blender; remaining args are blender args
args = append(args, spec.Args...)
cmd := exec.Command("podman", args...)
cmd.Dir = work
return cmd, nil
}
// PodmanImageDefault returns the default thin runtime image name.
func PodmanImageDefault() string { return defaultPodmanImage }
// FormatUIDGID is a small helper for tests.
func FormatUIDGID(uid, gid int) string {
return strconv.Itoa(uid) + ":" + strconv.Itoa(gid)
}
+128
View File
@@ -0,0 +1,128 @@
// Package sandbox wraps Blender (and similar) invocations in optional isolation.
//
// Backends:
// - none: run on the host (current behavior)
// - podman: rootless podman container; Blender tarball is bind-mounted (not baked into an image)
//
// GPU access is provided by passing host device nodes and common userspace lib roots
// discovered at wrap time (NVIDIA / AMD ROCm / Intel DRM), not by shipping per-version
// Blender container images.
package sandbox
import (
"fmt"
"os"
"os/exec"
"strings"
)
// Backend names accepted by New and CLI flags.
const (
BackendNone = "none"
BackendPodman = "podman"
)
// Options configures sandbox construction.
type Options struct {
// Backend is none|podman (default none).
Backend string
// AllowNetwork keeps network in the sandbox (default false for podman).
AllowNetwork bool
// PodmanImage is used only for backend=podman. The image is a thin OS root;
// Blender comes from a host bind of the versioned tarball tree.
// Default: registry.fedoraproject.org/fedora-minimal:41
PodmanImage string
}
// Spec describes a Blender process to run under the sandbox.
type Spec struct {
// BlenderBinary is an absolute path to the blender executable.
BlenderBinary string
// Args are arguments after the binary (not including argv0).
Args []string
// WorkDir is the process working directory (job workspace); bind-mounted RW.
WorkDir string
// HomeDir is Blender HOME (usually WorkDir/home); bind-mounted RW.
HomeDir string
// Env is the full environment for the process (HOME, LD_LIBRARY_PATH, etc.).
Env []string
// GPU hints from host detection (used to select device/lib binds).
HasAMD bool
HasNVIDIA bool
HasIntel bool
// ForceCPU skips GPU device binds when true.
ForceCPU bool
// AllowNetwork overrides Options.AllowNetwork when non-nil... kept simple: use Options only.
}
// Wrapper turns a Spec into an *exec.Cmd ready to Start.
type Wrapper interface {
// Name returns the backend name.
Name() string
// Available reports whether required host tools exist.
Available() error
// Wrap builds the command. Callers own Start/Wait/pipes.
Wrap(spec Spec) (*exec.Cmd, error)
}
// New returns a sandbox Wrapper for the given options.
func New(opts Options) (Wrapper, error) {
backend := strings.ToLower(strings.TrimSpace(opts.Backend))
if backend == "" {
backend = BackendPodman
}
switch backend {
case BackendNone:
return &noneWrapper{}, nil
case BackendPodman:
img := opts.PodmanImage
if img == "" {
img = defaultPodmanImage
}
w := &podmanWrapper{allowNetwork: opts.AllowNetwork, image: img}
return w, nil
default:
return nil, fmt.Errorf("unknown sandbox backend %q (want none or podman)", opts.Backend)
}
}
// NormalizeBackend validates and returns a canonical backend name.
func NormalizeBackend(s string) (string, error) {
b := strings.ToLower(strings.TrimSpace(s))
if b == "" {
return BackendPodman, nil
}
switch b {
case BackendNone, BackendPodman:
return b, nil
default:
return "", fmt.Errorf("unknown sandbox backend %q (want none or podman)", s)
}
}
// LookPath is os/exec.LookPath, overridable in tests.
var LookPath = exec.LookPath
// pathExists reports whether path exists.
func pathExists(p string) bool {
_, err := os.Stat(p)
return err == nil
}
// mustAbs returns an absolute path or the original on error.
func mustAbs(p string) string {
if p == "" {
return p
}
abs, err := absPath(p)
if err != nil {
return p
}
return abs
}
// absPath is filepath.Abs, isolated for tests.
var absPath = func(p string) (string, error) {
return execAbs(p)
}
+132
View File
@@ -0,0 +1,132 @@
package sandbox
import (
"os"
"path/filepath"
"strings"
"testing"
)
func TestNormalizeBackend(t *testing.T) {
got, err := NormalizeBackend("PODMAN")
if err != nil || got != BackendPodman {
t.Fatalf("got %q err %v", got, err)
}
if _, err := NormalizeBackend("bwrap"); err == nil {
t.Fatal("expected error for removed bwrap backend")
}
if _, err := NormalizeBackend("firecracker"); err == nil {
t.Fatal("expected error for unknown backend")
}
got, err = NormalizeBackend("")
if err != nil || got != BackendPodman {
t.Fatalf("empty -> podman, got %q %v", got, err)
}
}
func TestNoneWrap_Passthrough(t *testing.T) {
w, err := New(Options{Backend: BackendNone})
if err != nil {
t.Fatal(err)
}
bin := "/opt/blender/blender"
cmd, err := w.Wrap(Spec{
BlenderBinary: bin,
Args: []string{"-b", "scene.blend"},
WorkDir: "/tmp/job",
Env: []string{"HOME=/tmp/job/home"},
})
if err != nil {
t.Fatal(err)
}
if len(cmd.Args) == 0 || cmd.Args[0] != bin {
t.Fatalf("args[0]=%v path=%q", cmd.Args, cmd.Path)
}
if cmd.Dir != "/tmp/job" {
t.Fatalf("Dir=%q", cmd.Dir)
}
}
func TestGPUMounts_ForceCPUEmpty(t *testing.T) {
m := GPUMounts(true, true, true, true)
if len(m) != 0 {
t.Fatalf("force CPU should skip GPU mounts, got %#v", m)
}
}
func TestGPUMounts_AMDIncludesKFD(t *testing.T) {
m := GPUMounts(true, false, false, false)
var hosts []string
for _, b := range m {
hosts = append(hosts, b.Host)
}
joined := strings.Join(hosts, ",")
if !strings.Contains(joined, "/dev/dri") {
t.Fatalf("AMD mounts should include /dev/dri, got %v", hosts)
}
if !strings.Contains(joined, "/dev/kfd") {
t.Fatalf("AMD mounts should include /dev/kfd, got %v", hosts)
}
}
func TestGPUMounts_NVIDIAIncludesCtl(t *testing.T) {
m := GPUMounts(false, true, false, false)
found := false
for _, b := range m {
if strings.Contains(b.Host, "nvidia") {
found = true
break
}
}
if !found {
t.Fatalf("NVIDIA mounts should include nvidia devices, got %#v", m)
}
}
func TestPodmanWrap_BuildsRunArgs(t *testing.T) {
w := &podmanWrapper{allowNetwork: false, image: "example.com/thin:latest"}
if err := w.Available(); err != nil {
t.Skipf("podman not installed: %v", err)
}
tmp := t.TempDir()
blendDir := filepath.Join(tmp, "bver")
_ = os.MkdirAll(blendDir, 0755)
bin := filepath.Join(blendDir, "blender")
_ = os.WriteFile(bin, []byte("x"), 0755)
job := filepath.Join(tmp, "job")
_ = os.MkdirAll(job, 0755)
cmd, err := w.Wrap(Spec{
BlenderBinary: bin,
Args: []string{"-b", "x.blend"},
WorkDir: job,
HomeDir: filepath.Join(job, "home"),
Env: []string{"HOME=" + filepath.Join(job, "home")},
HasNVIDIA: true,
})
if err != nil {
t.Fatal(err)
}
joined := strings.Join(cmd.Args, " ")
if !strings.Contains(joined, "run") || !strings.Contains(joined, "--network=none") {
t.Fatalf("expected podman run --network=none: %s", joined)
}
if !strings.Contains(joined, "example.com/thin:latest") {
t.Fatalf("expected image in args: %s", joined)
}
if !strings.Contains(joined, "--entrypoint") || !strings.Contains(joined, bin) {
t.Fatalf("expected entrypoint blender: %s", joined)
}
if !strings.Contains(joined, blendDir) || !strings.Contains(joined, job) {
t.Fatalf("expected volume binds: %s", joined)
}
}
func TestNew_UnknownBackend(t *testing.T) {
if _, err := New(Options{Backend: "bwrap"}); err == nil {
t.Fatal("expected error for bwrap")
}
if _, err := New(Options{Backend: "gvisor"}); err == nil {
t.Fatal("expected error")
}
}
+2
View File
@@ -246,6 +246,8 @@ func (p *EncodeProcessor) Process(ctx *Context) error {
TwoPass: true,
})
if err := pass1Cmd.Run(); err != nil {
// Pass 1 is analysis-only (writes to /dev/null). FFmpeg often exits non-zero
// on benign codec/option warnings while still producing passlogfile stats.
ctx.Warn(fmt.Sprintf("Pass 1 completed (warnings expected): %v", err))
}
+46
View File
@@ -0,0 +1,46 @@
package tasks
import (
"fmt"
"os/exec"
"strings"
)
// mergeEXRFiles averages multiple linear EXR renders into one (Monte Carlo accumulation).
func mergeEXRFiles(output string, inputs []string) error {
if len(inputs) == 0 {
return fmt.Errorf("no input EXR files to merge")
}
if len(inputs) == 1 {
return copyFile(inputs[0], output)
}
args := append([]string{}, inputs...)
args = append(args, "-evaluate-sequence", "mean", output)
if path, err := exec.LookPath("magick"); err == nil {
cmd := exec.Command(path, args...)
if out, err := cmd.CombinedOutput(); err != nil {
return fmt.Errorf("magick merge failed: %w (%s)", err, strings.TrimSpace(string(out)))
}
return nil
}
if path, err := exec.LookPath("convert"); err == nil {
cmd := exec.Command(path, args...)
if out, err := cmd.CombinedOutput(); err != nil {
return fmt.Errorf("convert merge failed: %w (%s)", err, strings.TrimSpace(string(out)))
}
return nil
}
return fmt.Errorf("ImageMagick (magick or convert) required to merge batched EXR renders")
}
func copyFile(src, dst string) error {
cmd := exec.Command("cp", "-f", src, dst)
if out, err := cmd.CombinedOutput(); err != nil {
return fmt.Errorf("copy %s to %s: %w (%s)", src, dst, err, strings.TrimSpace(string(out)))
}
return nil
}
+89 -4
View File
@@ -7,6 +7,7 @@ import (
"jiggablend/internal/runner/api"
"jiggablend/internal/runner/blender"
"jiggablend/internal/runner/encoding"
"jiggablend/internal/runner/sandbox"
"jiggablend/internal/runner/workspace"
"jiggablend/pkg/executils"
"jiggablend/pkg/types"
@@ -41,6 +42,9 @@ type Context struct {
// GPULockedOut is set when the runner has detected a GPU error (e.g. HIP) and disables GPU for all jobs.
GPULockedOut bool
// GPULockoutArmedThisAttempt is true when this task attempt newly enabled runner GPU lockout.
// Failures with this set free-requeue without incrementing retry_count (expected ROCm warmup).
GPULockoutArmedThisAttempt bool
// HasAMD is true when the runner detected AMD devices at startup.
HasAMD bool
// HasNVIDIA is true when the runner detected NVIDIA GPUs at startup.
@@ -53,6 +57,12 @@ type Context struct {
OnGPUError func()
// ForceCPURendering is a runner-level override that forces CPU rendering for all jobs.
ForceCPURendering bool
// DisableRT disables GPU ray tracing acceleration (runner-level flag).
DisableRT bool
// HipGPUSampleBatch limits samples per GPU render pass on gfx115x (0 = no batching).
HipGPUSampleBatch int
// Sandbox wraps Blender execution (none/podman). Nil means none.
Sandbox sandbox.Wrapper
}
// ErrJobCancelled indicates the manager-side job was cancelled during execution.
@@ -80,7 +90,10 @@ func NewContext(
hasIntel bool,
gpuDetectionFailed bool,
forceCPURendering bool,
disableRT bool,
hipGPUSampleBatch int,
onGPUError func(),
sb sandbox.Wrapper,
) *Context {
if frameEnd < frameStart {
frameEnd = frameStart
@@ -107,7 +120,10 @@ func NewContext(
HasIntel: hasIntel,
GPUDetectionFailed: gpuDetectionFailed,
ForceCPURendering: forceCPURendering,
OnGPUError: onGPUError,
DisableRT: disableRT,
HipGPUSampleBatch: hipGPUSampleBatch,
OnGPUError: onGPUError,
Sandbox: sb,
}
}
@@ -148,18 +164,22 @@ func (c *Context) OutputUploaded(fileName string) {
}
// Complete sends task completion.
// When the failure newly armed GPU lockout, free_requeue is set so the manager
// requeues without burning retry_count.
func (c *Context) Complete(success bool, errorMsg error) {
if c.JobConn != nil {
c.JobConn.Complete(c.TaskID, success, errorMsg)
freeRequeue := !success && c.GPULockoutArmedThisAttempt
c.JobConn.Complete(c.TaskID, success, errorMsg, freeRequeue)
}
}
// GetOutputFormat returns the output format from metadata or default.
// Product default is EXR (Blender always renders EXR; deliverable may add video).
func (c *Context) GetOutputFormat() string {
if c.Metadata != nil && c.Metadata.RenderSettings.OutputFormat != "" {
return c.Metadata.RenderSettings.OutputFormat
}
return "PNG"
return "EXR"
}
// GetFrameRate returns the frame rate from metadata or default.
@@ -183,7 +203,9 @@ func (c *Context) ShouldUnhideObjects() bool {
return c.Metadata != nil && c.Metadata.UnhideObjects != nil && *c.Metadata.UnhideObjects
}
// ShouldEnableExecution returns whether to enable auto-execution.
// ShouldEnableExecution returns whether to pass Blender --enable-autoexec
// (Python drivers/scripts inside the .blend). Job-bundled blender_addons/ install
// independently whenever that folder is present in the context.
func (c *Context) ShouldEnableExecution() bool {
return c.Metadata != nil && c.Metadata.EnableExecution != nil && *c.Metadata.EnableExecution
}
@@ -211,6 +233,69 @@ func (c *Context) ShouldForceCPU() bool {
return false
}
// GetCyclesSamples returns the target Cycles sample count from metadata (default 128).
func (c *Context) GetCyclesSamples() int {
if c.Metadata != nil && c.Metadata.RenderSettings.EngineSettings != nil {
if v, ok := c.Metadata.RenderSettings.EngineSettings["samples"]; ok {
switch n := v.(type) {
case float64:
if int(n) > 0 {
return int(n)
}
case int:
if n > 0 {
return n
}
}
}
}
return 128
}
// GetCyclesSeed returns the Cycles seed from metadata (default 0).
func (c *Context) GetCyclesSeed() int {
if c.Metadata != nil && c.Metadata.RenderSettings.EngineSettings != nil {
if v, ok := c.Metadata.RenderSettings.EngineSettings["seed"]; ok {
switch n := v.(type) {
case float64:
return int(n)
case int:
return n
}
}
}
return 0
}
// ShouldBatchGPUSamples returns true when HIP GPU renders should be split into sample batches.
func (c *Context) ShouldBatchGPUSamples() bool {
if c.ShouldForceCPU() || c.HipGPUSampleBatch <= 0 {
return false
}
return c.GetCyclesSamples() > c.HipGPUSampleBatch
}
// ShouldDisableRT returns true when GPU ray tracing acceleration should be disabled.
func (c *Context) ShouldDisableRT() bool {
if c.DisableRT {
return true
}
if c.Metadata != nil && c.Metadata.RenderSettings.EngineSettings != nil {
settings := c.Metadata.RenderSettings.EngineSettings
if v, ok := settings["disable_rt"]; ok {
if b, ok := v.(bool); ok && b {
return true
}
}
if v, ok := settings["disable_hiprt"]; ok {
if b, ok := v.(bool); ok && b {
return true
}
}
}
return false
}
// IsJobCancelled checks whether the manager marked this job as cancelled.
func (c *Context) IsJobCancelled() (bool, error) {
if c.Manager == nil {
+3 -3
View File
@@ -8,7 +8,7 @@ import (
)
func TestNewContext_NormalizesFrameEnd(t *testing.T) {
ctx := NewContext(1, 2, "job", 10, 1, "render", "/tmp", "tok", nil, nil, nil, nil, nil, nil, nil, false, false, false, false, false, false, nil)
ctx := NewContext(1, 2, "job", 10, 1, "render", "/tmp", "tok", nil, nil, nil, nil, nil, nil, nil, false, false, false, false, false, false, false, 0, nil, nil)
if ctx.FrameEnd != 10 {
t.Fatalf("expected FrameEnd to be normalized to Frame, got %d", ctx.FrameEnd)
}
@@ -16,8 +16,8 @@ func TestNewContext_NormalizesFrameEnd(t *testing.T) {
func TestContext_GetOutputFormat_Default(t *testing.T) {
ctx := &Context{}
if got := ctx.GetOutputFormat(); got != "PNG" {
t.Fatalf("GetOutputFormat() = %q, want PNG", got)
if got := ctx.GetOutputFormat(); got != "EXR" {
t.Fatalf("GetOutputFormat() = %q, want EXR", got)
}
}
+156 -27
View File
@@ -12,6 +12,7 @@ import (
"strings"
"jiggablend/internal/runner/blender"
"jiggablend/internal/runner/sandbox"
"jiggablend/internal/runner/workspace"
"jiggablend/pkg/scripts"
"jiggablend/pkg/types"
@@ -28,6 +29,8 @@ func NewRenderProcessor() *RenderProcessor {
// gpuErrorSubstrings are log line substrings that indicate a GPU backend error (matched case-insensitively); any match triggers full GPU lockout.
var gpuErrorSubstrings = []string{
"illegal address in hip", // HIP (AMD) e.g. "Illegal address in HIP" or "Illegal address in hip"
"memory access fault", // AMDGPU page fault during HIP/HIPRT (not necessarily OOM)
"page not present", // AMDGPU GCVM page fault detail
"hiperror", // hipError* codes
"hip error",
"cuda error",
@@ -38,14 +41,18 @@ var gpuErrorSubstrings = []string{
}
// checkGPUErrorLine checks a log line for GPU error indicators and triggers runner GPU lockout if found.
// Once lockout is already active for this attempt (or was active at context creation), further
// matching lines are ignored so we do not spam lockout callbacks/logs on multi-line fault dumps.
func (p *RenderProcessor) checkGPUErrorLine(ctx *Context, line string) {
if ctx.GPULockedOut || ctx.GPULockoutArmedThisAttempt {
return
}
lower := strings.ToLower(line)
for _, sub := range gpuErrorSubstrings {
if strings.Contains(lower, sub) {
if ctx.OnGPUError != nil {
ctx.OnGPUError()
}
ctx.Warn(fmt.Sprintf("GPU error detected in log (%q); GPU disabled for subsequent jobs", sub))
return
}
}
@@ -73,26 +80,24 @@ func (p *RenderProcessor) Process(ctx *Context) error {
return fmt.Errorf("failed to find blend file: %w", err)
}
// Get Blender binary
blenderBinary := "blender"
if version := ctx.GetBlenderVersion(); version != "" {
ctx.Info(fmt.Sprintf("Job requires Blender %s", version))
binaryPath, err := ctx.Blender.GetBinaryPath(version)
if err != nil {
ctx.Warn(fmt.Sprintf("Could not get Blender %s, using system blender: %v", version, err))
} else {
blenderBinary = binaryPath
ctx.Info(fmt.Sprintf("Using Blender binary: %s", blenderBinary))
}
} else {
ctx.Info("No Blender version specified, using system blender")
// Runners must use manager-provided Blender versions; never fall back to system blender.
version := ctx.GetBlenderVersion()
if version == "" {
return fmt.Errorf("job metadata missing blender_version: runner cannot use system blender")
}
blenderBinary, err = blender.ResolveBinaryPath(blenderBinary)
ctx.Info(fmt.Sprintf("Job requires Blender %s", version))
binaryPath, err := ctx.Blender.GetBinaryPath(version)
if err != nil {
return fmt.Errorf("failed to resolve blender binary: %w", err)
return fmt.Errorf("failed to get Blender %s from manager: %w", version, err)
}
blenderBinary, err := blender.ResolveBinaryPath(binaryPath)
if err != nil {
return fmt.Errorf("failed to resolve Blender %s binary: %w", version, err)
}
ctx.Info(fmt.Sprintf("Using Blender binary: %s", blenderBinary))
// Create output directory
outputDir := filepath.Join(ctx.WorkDir, "output")
if err := os.MkdirAll(outputDir, 0755); err != nil {
@@ -116,11 +121,15 @@ func (p *RenderProcessor) Process(ctx *Context) error {
} else {
ctx.Info("GPU lockout active: using CPU rendering only")
}
} else if ctx.ShouldDisableRT() {
ctx.Info("GPU ray tracing acceleration disabled for this job (--disable-rt)")
}
// Create render script
if err := p.createRenderScript(ctx, renderFormat); err != nil {
return err
if ctx.ShouldBatchGPUSamples() {
total := ctx.GetCyclesSamples()
ctx.Info(fmt.Sprintf(
"gfx115x HIP sample batching: rendering %d samples in passes of %d (ROCm driver limit, not system RAM)",
total, ctx.HipGPUSampleBatch,
))
}
// Render
@@ -129,7 +138,7 @@ func (p *RenderProcessor) Process(ctx *Context) error {
} else {
ctx.Info(fmt.Sprintf("Starting Blender render for frame %d...", ctx.Frame))
}
if err := p.runBlender(ctx, blenderBinary, blendFile, outputDir, renderFormat, blenderHome); err != nil {
if err := p.renderFrames(ctx, blenderBinary, blendFile, outputDir, renderFormat, blenderHome); err != nil {
if errors.Is(err, ErrJobCancelled) {
ctx.Warn("Render stopped because job was cancelled")
return err
@@ -152,7 +161,84 @@ func (p *RenderProcessor) Process(ctx *Context) error {
return nil
}
func (p *RenderProcessor) createRenderScript(ctx *Context, renderFormat string) error {
type renderPassOptions struct {
samplesOverride *int
seedOverride *int
}
func (p *RenderProcessor) renderFrames(ctx *Context, blenderBinary, blendFile, outputDir, renderFormat, blenderHome string) error {
if !ctx.ShouldBatchGPUSamples() {
if err := p.createRenderScript(ctx, renderFormat, nil); err != nil {
return err
}
return p.runBlender(ctx, blenderBinary, blendFile, outputDir, renderFormat, blenderHome)
}
totalSamples := ctx.GetCyclesSamples()
batchSize := ctx.HipGPUSampleBatch
baseSeed := ctx.GetCyclesSeed()
batchDir := filepath.Join(ctx.WorkDir, "sample_batches")
if err := os.MkdirAll(batchDir, 0755); err != nil {
return fmt.Errorf("failed to create sample batch directory: %w", err)
}
var batchOutputs []string
remaining := totalSamples
batchNum := 0
for remaining > 0 {
if err := ctx.CheckCancelled(); err != nil {
return err
}
passSamples := batchSize
if remaining < passSamples {
passSamples = remaining
}
seed := baseSeed + batchNum
batchNum++
ctx.Info(fmt.Sprintf("GPU sample batch %d: %d samples (seed %d, %d remaining)", batchNum, passSamples, seed, remaining-passSamples))
passOutput := filepath.Join(batchDir, fmt.Sprintf("batch_%03d", batchNum))
if err := os.MkdirAll(passOutput, 0755); err != nil {
return err
}
opts := &renderPassOptions{
samplesOverride: &passSamples,
seedOverride: &seed,
}
if err := p.createRenderScript(ctx, renderFormat, opts); err != nil {
return err
}
if err := p.runBlender(ctx, blenderBinary, blendFile, passOutput, renderFormat, blenderHome); err != nil {
return err
}
if err := p.verifyOutputRange(ctx, passOutput, renderFormat); err != nil {
return err
}
batchOutputs = append(batchOutputs, p.firstFrameOutputPath(ctx, passOutput, renderFormat))
remaining -= passSamples
}
if err := os.MkdirAll(outputDir, 0755); err != nil {
return err
}
finalOutput := p.firstFrameOutputPath(ctx, outputDir, renderFormat)
ctx.Info(fmt.Sprintf("Merging %d GPU sample batches into final EXR...", len(batchOutputs)))
if err := mergeEXRFiles(finalOutput, batchOutputs); err != nil {
return err
}
ctx.Info("GPU sample batch merge completed")
return nil
}
func (p *RenderProcessor) firstFrameOutputPath(ctx *Context, outputDir, renderFormat string) string {
ext := strings.ToLower(renderFormat)
if ctx.FrameEnd > ctx.Frame {
return filepath.Join(outputDir, fmt.Sprintf("frame_%04d.%s", ctx.Frame, ext))
}
return filepath.Join(outputDir, fmt.Sprintf("frame_%04d.%s", ctx.Frame, ext))
}
func (p *RenderProcessor) createRenderScript(ctx *Context, renderFormat string, opts *renderPassOptions) error {
formatFilePath := filepath.Join(ctx.WorkDir, "output_format.txt")
renderSettingsFilePath := filepath.Join(ctx.WorkDir, "render_settings.json")
@@ -162,7 +248,9 @@ func (p *RenderProcessor) createRenderScript(ctx *Context, renderFormat string)
unhideCode = scripts.UnhideObjects
}
// Load template and replace placeholders
// Load template and replace placeholders.
// Job-bundled blender_addons/ always auto-install when present (users ship addons with scenes).
// enable_execution only controls Blender --enable-autoexec (scripts inside .blend).
scriptContent := scripts.RenderBlenderTemplate
scriptContent = strings.ReplaceAll(scriptContent, "{{UNHIDE_CODE}}", unhideCode)
scriptContent = strings.ReplaceAll(scriptContent, "{{FORMAT_FILE_PATH}}", fmt.Sprintf("%q", formatFilePath))
@@ -195,6 +283,20 @@ func (p *RenderProcessor) createRenderScript(ctx *Context, renderFormat string)
settingsMap = make(map[string]interface{})
}
settingsMap["force_cpu"] = ctx.ShouldForceCPU()
settingsMap["disable_rt"] = ctx.ShouldDisableRT()
if ctx.ShouldBatchGPUSamples() {
// gfx115x: large tiles trigger AMDGPU page faults during HIP renders.
settingsMap["tile_size_override"] = 256
settingsMap["use_auto_tile_override"] = true
}
if opts != nil {
if opts.samplesOverride != nil {
settingsMap["samples_override"] = *opts.samplesOverride
}
if opts.seedOverride != nil {
settingsMap["seed_override"] = *opts.seedOverride
}
}
settingsJSON, err := json.Marshal(settingsMap)
if err == nil {
if err := os.WriteFile(renderSettingsFilePath, settingsJSON, 0644); err != nil {
@@ -205,6 +307,32 @@ func (p *RenderProcessor) createRenderScript(ctx *Context, renderFormat string)
return nil
}
// wrapBlenderCmd builds the Blender *exec.Cmd, optionally through the sandbox wrapper.
func wrapBlenderCmd(ctx *Context, blenderBinary string, args, env []string, blenderHome string) (*exec.Cmd, error) {
sb := ctx.Sandbox
if sb == nil {
var err error
sb, err = sandbox.New(sandbox.Options{Backend: sandbox.BackendNone})
if err != nil {
return nil, err
}
}
if sb.Name() != sandbox.BackendNone {
ctx.Info(fmt.Sprintf("Running Blender under sandbox backend %q", sb.Name()))
}
return sb.Wrap(sandbox.Spec{
BlenderBinary: blenderBinary,
Args: args,
WorkDir: ctx.WorkDir,
HomeDir: blenderHome,
Env: env,
HasAMD: ctx.HasAMD,
HasNVIDIA: ctx.HasNVIDIA,
HasIntel: ctx.HasIntel,
ForceCPU: ctx.ShouldForceCPU(),
})
}
func (p *RenderProcessor) runBlender(ctx *Context, blenderBinary, blendFile, outputDir, renderFormat, blenderHome string) error {
scriptPath := filepath.Join(ctx.WorkDir, "enable_gpu.py")
blendFileAbs, err := filepath.Abs(blendFile)
@@ -233,9 +361,6 @@ func (p *RenderProcessor) runBlender(ctx *Context, blenderBinary, blendFile, out
args = append(args, "-f", fmt.Sprintf("%d", ctx.Frame))
}
cmd := execCommand(blenderBinary, args...)
cmd.Dir = ctx.WorkDir
// Set up environment: LD_LIBRARY_PATH for tarball Blender, then custom HOME
env := os.Environ()
env = blender.TarballEnv(blenderBinary, env)
@@ -246,7 +371,11 @@ func (p *RenderProcessor) runBlender(ctx *Context, blenderBinary, blendFile, out
}
}
newEnv = append(newEnv, fmt.Sprintf("HOME=%s", blenderHome))
cmd.Env = newEnv
cmd, err := wrapBlenderCmd(ctx, blenderBinary, args, newEnv, blenderHome)
if err != nil {
return fmt.Errorf("failed to build sandboxed blender command: %w", err)
}
// Set up pipes
stdoutPipe, err := cmd.StdoutPipe()
+33 -4
View File
@@ -4,13 +4,42 @@ import "testing"
func TestCheckGPUErrorLine_TriggersCallback(t *testing.T) {
p := NewRenderProcessor()
triggered := false
triggered := 0
ctx := &Context{
OnGPUError: func() { triggered = true },
OnGPUError: func() {
triggered++
// Simulate runner arming lockout for this attempt.
},
}
p.checkGPUErrorLine(ctx, "Fatal: Illegal address in HIP kernel execution")
if !triggered {
t.Fatal("expected GPU error callback to be triggered")
if triggered != 1 {
t.Fatalf("expected GPU error callback once, got %d", triggered)
}
// After arming, further fault lines must not re-fire (spam protection).
ctx.GPULockoutArmedThisAttempt = true
p.checkGPUErrorLine(ctx, "page not present in GPU memory")
if triggered != 1 {
t.Fatalf("expected no further callbacks after arm, got %d", triggered)
}
}
func TestCheckGPUErrorLine_SkipsWhenAlreadyLockedOut(t *testing.T) {
p := NewRenderProcessor()
triggered := false
ctx := &Context{
GPULockedOut: true,
OnGPUError: func() { triggered = true },
}
p.checkGPUErrorLine(ctx, "Illegal address in HIP")
if triggered {
t.Fatal("did not expect callback when GPU already locked out")
}
}
func TestGetBlenderVersion_EmptyWhenMissing(t *testing.T) {
ctx := &Context{}
if got := ctx.GetBlenderVersion(); got != "" {
t.Fatalf("expected empty blender version, got %q", got)
}
}
+14
View File
@@ -128,6 +128,20 @@ func ExtractTarStripPrefix(reader io.Reader, destDir string) error {
if err := os.MkdirAll(filepath.Dir(targetPath), 0755); err != nil {
return err
}
// Reject absolute or escaping symlink targets so extraction cannot
// plant links that point outside destDir.
linkTarget := header.Linkname
if linkTarget == "" {
return fmt.Errorf("invalid empty symlink target in tar: %s", header.Name)
}
if filepath.IsAbs(linkTarget) {
return fmt.Errorf("absolute symlink target not allowed in tar: %s -> %s", header.Name, linkTarget)
}
resolvedLink := filepath.Clean(filepath.Join(filepath.Dir(targetPath), linkTarget))
cleanDest := filepath.Clean(destDir)
if resolvedLink != cleanDest && !strings.HasPrefix(resolvedLink, cleanDest+string(os.PathSeparator)) {
return fmt.Errorf("symlink target escapes extract root: %s -> %s", header.Name, linkTarget)
}
os.Remove(targetPath) // Remove existing symlink if present
if err := os.Symlink(header.Linkname, targetPath); err != nil {
return err
+19
View File
@@ -24,6 +24,25 @@ func createTarBuffer(files map[string]string) *bytes.Buffer {
return &buf
}
func TestExtractTarStripPrefix_RejectsEscapingSymlink(t *testing.T) {
destDir := t.TempDir()
var buf bytes.Buffer
tw := tar.NewWriter(&buf)
// Top-level prefix like Blender archives
_ = tw.WriteHeader(&tar.Header{Name: "blender-x/", Typeflag: tar.TypeDir, Mode: 0755})
_ = tw.WriteHeader(&tar.Header{
Name: "blender-x/evil",
Typeflag: tar.TypeSymlink,
Linkname: "../../outside",
Mode: 0777,
})
_ = tw.Close()
if err := ExtractTarStripPrefix(&buf, destDir); err == nil {
t.Fatal("expected escaping symlink to be rejected")
}
}
func TestExtractTar(t *testing.T) {
destDir := t.TempDir()
+59 -1
View File
@@ -80,10 +80,54 @@ func (s *Storage) JobPath(jobID int64) string {
return filepath.Join(s.basePath, "jobs", fmt.Sprintf("%d", jobID))
}
// SanitizeFilename returns a safe basename for client-supplied names.
// Rejects empty, ".", "..", and root-like basenames.
func SanitizeFilename(name string) (string, error) {
cleaned := filepath.Clean(strings.ReplaceAll(name, "\\", "/"))
base := filepath.Base(cleaned)
// filepath.Base("..") is ".." ; Clean("/../..") style inputs can yield "/" or "."
if base == "" || base == "." || base == ".." || base == "/" || base == string(os.PathSeparator) {
return "", fmt.Errorf("invalid filename: %q", name)
}
// Reject if the cleaned path still has parent references after basenaming was skipped
if strings.Contains(cleaned, "..") && base == cleaned {
return "", fmt.Errorf("invalid filename: %q", name)
}
return base, nil
}
// SafePathUnderRoot joins root and a relative path and ensures the result stays under root.
func SafePathUnderRoot(root, rel string) (string, error) {
if rel == "" {
return "", fmt.Errorf("empty path")
}
// Normalize to forward slashes then clean
rel = filepath.ToSlash(rel)
if filepath.IsAbs(rel) || strings.HasPrefix(rel, "/") {
return "", fmt.Errorf("absolute path not allowed: %q", rel)
}
cleanRel := filepath.Clean(rel)
if cleanRel == ".." || strings.HasPrefix(cleanRel, ".."+string(os.PathSeparator)) {
return "", fmt.Errorf("path escapes root: %q", rel)
}
cleanRoot := filepath.Clean(root)
full := filepath.Join(cleanRoot, cleanRel)
cleanFull := filepath.Clean(full)
sep := string(os.PathSeparator)
if cleanFull != cleanRoot && !strings.HasPrefix(cleanFull, cleanRoot+sep) {
return "", fmt.Errorf("path escapes root: %q", rel)
}
return cleanFull, nil
}
// SaveUpload saves an uploaded file
func (s *Storage) SaveUpload(jobID int64, filename string, reader io.Reader) (string, error) {
// Sanitize filename to prevent path traversal
filename = filepath.Base(filename)
var err error
filename, err = SanitizeFilename(filename)
if err != nil {
return "", err
}
jobPath := s.JobPath(jobID)
if err := os.MkdirAll(jobPath, 0755); err != nil {
@@ -106,12 +150,26 @@ func (s *Storage) SaveUpload(jobID int64, filename string, reader io.Reader) (st
// SaveOutput saves an output file
func (s *Storage) SaveOutput(jobID int64, filename string, reader io.Reader) (string, error) {
// Sanitize filename to prevent path traversal (parity with SaveUpload)
var err error
filename, err = SanitizeFilename(filename)
if err != nil {
return "", err
}
outputPath := filepath.Join(s.outputsPath(), fmt.Sprintf("%d", jobID))
if err := os.MkdirAll(outputPath, 0755); err != nil {
return "", fmt.Errorf("failed to create output directory: %w", err)
}
filePath := filepath.Join(outputPath, filename)
// Defense in depth: ensure resolved path stays under output dir
if resolved, err := SafePathUnderRoot(outputPath, filename); err != nil {
return "", err
} else {
filePath = resolved
}
file, err := os.Create(filePath)
if err != nil {
return "", fmt.Errorf("failed to create file: %w", err)
+45
View File
@@ -66,6 +66,51 @@ func TestSaveOutput(t *testing.T) {
}
}
func TestSaveOutput_PathTraversal(t *testing.T) {
s := setupStorage(t)
path, err := s.SaveOutput(42, "../../etc/passwd", strings.NewReader("evil"))
if err != nil {
t.Fatalf("SaveOutput: %v", err)
}
outRoot := filepath.Join(s.BasePath(), "outputs", "42")
if !strings.HasPrefix(path, outRoot) {
t.Errorf("saved file %q escaped output directory %q", path, outRoot)
}
if filepath.Base(path) != "passwd" {
t.Errorf("expected basename passwd, got %q", filepath.Base(path))
}
}
func TestSanitizeFilename(t *testing.T) {
got, err := SanitizeFilename("../../etc/passwd")
if err != nil {
t.Fatalf("SanitizeFilename: %v", err)
}
if got != "passwd" {
t.Fatalf("got %q want passwd", got)
}
if _, err := SanitizeFilename(".."); err == nil {
t.Fatal("expected error for ..")
}
}
func TestSafePathUnderRoot(t *testing.T) {
root := t.TempDir()
ok, err := SafePathUnderRoot(root, "subdir/file.blend")
if err != nil {
t.Fatalf("SafePathUnderRoot: %v", err)
}
if !strings.HasPrefix(ok, root) {
t.Fatalf("path %q not under root %q", ok, root)
}
if _, err := SafePathUnderRoot(root, "../escape.blend"); err == nil {
t.Fatal("expected escape to fail")
}
if _, err := SafePathUnderRoot(root, "/abs.blend"); err == nil {
t.Fatal("expected absolute path to fail")
}
}
func TestGetFile(t *testing.T) {
s := setupStorage(t)
savedPath, err := s.SaveUpload(1, "readme.txt", strings.NewReader("hello"))