From aec1f2b4d188c4e85c178c3c1dc97cbc7a149d81 Mon Sep 17 00:00:00 2001 From: walker1211 <13750528578@163.com> Date: Tue, 16 Jun 2026 09:02:52 +0800 Subject: [PATCH] =?UTF-8?q?fix(security):=20=E9=99=90=E5=88=B6=20submit=20?= =?UTF-8?q?=E5=9B=BE=E7=89=87=E8=B7=AF=E5=BE=84=E8=AE=BF=E9=97=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为修复 Code scanning 的路径注入告警，submit 会先将本地参考图暂存到受控目录，服务端只读取该目录内的相对路径。 --- cmd/imgen/main.go | 12 ++++- cmd/imgen/main_test.go | 24 ++++++++++ configs/config.example.yaml | 1 + internal/api/handlers.go | 2 +- internal/api/server_test.go | 2 +- internal/cli/client.go | 76 +++++++++++++++++++++++++++++++- internal/cli/client_test.go | 33 ++++++++++++++ internal/config/config.go | 5 ++- internal/config/load.go | 1 + internal/config/load_test.go | 23 ++++++++++ internal/service/service.go | 36 +++++++++++---- internal/service/service_test.go | 47 ++++++++++++++------ 12 files changed, 231 insertions(+), 31 deletions(-) diff --git a/cmd/imgen/main.go b/cmd/imgen/main.go index 64d2e26..13becc3 100644 --- a/cmd/imgen/main.go +++ b/cmd/imgen/main.go @@ -58,7 +58,7 @@ func main() { defer st.Close() hub := notify.NewWebSocketHub() serveGen := serveGenerator(generator, cfg) - svc := &service.Service{Store: st, Generator: serveGen, PromptPrefix: cfg.Backend.Prompt.Prefix, PromptPrelude: cfg.Backend.Prompt.Prelude, DefaultJobConcurrency: cfg.Scheduler.DefaultJobConcurrency, MaxJobConcurrency: cfg.Scheduler.MaxJobConcurrency, MaxCountPerJob: cfg.Scheduler.MaxCountPerJob, MaxAttempts: cfg.Scheduler.MaxAttempts, Publisher: hub} + svc := &service.Service{Store: st, Generator: serveGen, PromptPrefix: cfg.Backend.Prompt.Prefix, PromptPrelude: cfg.Backend.Prompt.Prelude, DefaultJobConcurrency: cfg.Scheduler.DefaultJobConcurrency, MaxJobConcurrency: cfg.Scheduler.MaxJobConcurrency, MaxCountPerJob: cfg.Scheduler.MaxCountPerJob, ImageInputDir: imageInputDir(dataDir, cfg), MaxAttempts: cfg.Scheduler.MaxAttempts, Publisher: hub} handler := api.NewServerWithOptions(svc, api.ServerOptions{ Notifier: hub, Realtime: realtimeOptions(serveGen, cfg), @@ -72,7 +72,8 @@ func main() { maintenance := scheduler.Maintenance{Store: st, Mailer: notify.NewMailer(cfg.Email), LeaseTimeout: cfg.Scheduler.TaskLeaseTimeout} app.ServerRunner = cli.HTTPServerRunner{Server: server, Maintenance: cli.MaintenanceAdapter{Maintenance: maintenance}, MaintenanceInterval: cfg.Scheduler.MaintenanceInterval} case "submit", "status", "get", "list", "cancel": - app.Client = &cli.Client{BaseURL: "http://" + cfg.Server.Listen} + dataDir, _ := storagePaths(home, cfg) + app.Client = &cli.Client{BaseURL: "http://" + cfg.Server.Listen, ImageInputDir: imageInputDir(dataDir, cfg)} case "doctor": app.OpenClawDoctor = doctor.NewOpenClawChecker(home) default: @@ -112,6 +113,13 @@ func storagePaths(home string, cfg config.Config) (string, string) { return dataDir, dbPath } +func imageInputDir(dataDir string, cfg config.Config) string { + if cfg.Storage.ImageInputDir != "" { + return cfg.Storage.ImageInputDir + } + return filepath.Join(dataDir, "input_images") +} + func mustOpenStore(dbPath string) *store.Store { _ = os.MkdirAll(filepath.Dir(dbPath), 0o755) st, err := store.Open(dbPath) diff --git a/cmd/imgen/main_test.go b/cmd/imgen/main_test.go index bbd391c..66cae86 100644 --- a/cmd/imgen/main_test.go +++ b/cmd/imgen/main_test.go @@ -2,6 +2,7 @@ package main import ( "context" + "path/filepath" "sync" "testing" "time" @@ -63,6 +64,29 @@ func TestRealtimeOptionsUsesConfig(t *testing.T) { } } +func TestImageInputDirDefaultsUnderDataDir(t *testing.T) { + cfg := config.Default() + dataDir := filepath.Join(t.TempDir(), "data") + + got := imageInputDir(dataDir, cfg) + + want := filepath.Join(dataDir, "input_images") + if got != want { + t.Fatalf("imageInputDir() = %q, want %q", got, want) + } +} + +func TestImageInputDirUsesStorageConfig(t *testing.T) { + cfg := config.Default() + cfg.Storage.ImageInputDir = filepath.Join(t.TempDir(), "custom-inputs") + + got := imageInputDir(filepath.Join(t.TempDir(), "data"), cfg) + + if got != cfg.Storage.ImageInputDir { + t.Fatalf("imageInputDir() = %q, want %q", got, cfg.Storage.ImageInputDir) + } +} + type blockingMainTestGenerator struct { entered chan backend.GenerateRequest release chan struct{} diff --git a/configs/config.example.yaml b/configs/config.example.yaml index 5f2ebe9..b348e73 100644 --- a/configs/config.example.yaml +++ b/configs/config.example.yaml @@ -6,6 +6,7 @@ server: storage: data_dir: "" # 服务数据目录；为空时使用系统用户数据目录 sqlite_path: "" # SQLite 数据库路径；为空时使用 data_dir/imgen.db + image_input_dir: "" # submit 图片暂存目录；为空时使用 data_dir/input_images scheduler: global_max_concurrency: 10 # serve 内底层生图队列总并发；submit async 和 WebSocket realtime 共用 diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 1a31dea..bdc5354 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -19,7 +19,7 @@ func handleCreateJob(w http.ResponseWriter, r *http.Request, service Service) { if err != nil { status := http.StatusInternalServerError code := "internal_error" - if err.Error() == "prompt is required" || strings.HasPrefix(err.Error(), "image path not found:") { + if err.Error() == "prompt is required" || strings.HasPrefix(err.Error(), "image path not found") || strings.HasPrefix(err.Error(), "image path must ") { status = http.StatusBadRequest code = "invalid_argument" } diff --git a/internal/api/server_test.go b/internal/api/server_test.go index 4d66538..321279c 100644 --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -107,7 +107,7 @@ func (invalidPromptService) CancelJob(jobID string) error { return n type invalidImageService struct{} func (invalidImageService) CreateJob(req CreateJobRequest) (CreateJobResult, error) { - return CreateJobResult{}, errors.New("image path not found: /tmp/missing.png") + return CreateJobResult{}, errors.New("image path not found") } func (invalidImageService) GetJob(jobID string) (JobStatus, error) { return JobStatus{}, nil } func (invalidImageService) ListJobs(limit int) ([]JobSummary, error) { return nil, nil } diff --git a/internal/cli/client.go b/internal/cli/client.go index 153735f..12b13ab 100644 --- a/internal/cli/client.go +++ b/internal/cli/client.go @@ -2,23 +2,95 @@ package cli import ( "context" + "crypto/sha256" + "encoding/hex" + "errors" + "os" + "path/filepath" + "strings" api "github.com/walker1211/codex-imgen/internal/api" ) type Client struct { - BaseURL string + BaseURL string + ImageInputDir string } func (c Client) CreateJob(ctx context.Context, prompt string, images []string, count int, concurrency int) (api.CreateJobResult, error) { + stagedImages, err := c.stageImages(images) + if err != nil { + return api.CreateJobResult{}, err + } return api.Client{BaseURL: c.BaseURL}.CreateJob(ctx, api.CreateJobRequest{ Prompt: prompt, - Images: images, + Images: stagedImages, Count: count, Concurrency: concurrency, }) } +func (c Client) stageImages(images []string) ([]string, error) { + if strings.TrimSpace(c.ImageInputDir) == "" { + return images, nil + } + staged := make([]string, 0, len(images)) + for _, image := range images { + trimmed := strings.TrimSpace(image) + if trimmed == "" { + continue + } + name, err := c.stageImage(trimmed) + if err != nil { + return nil, err + } + staged = append(staged, name) + } + return staged, nil +} + +func (c Client) stageImage(image string) (string, error) { + source, err := filepath.Abs(image) + if err != nil { + return "", err + } + info, err := os.Stat(source) + if err != nil { + if os.IsNotExist(err) { + return "", errors.New("image path not found") + } + return "", err + } + if !info.Mode().IsRegular() { + return "", errors.New("image path must be a regular file") + } + if err := os.MkdirAll(c.ImageInputDir, 0o700); err != nil { + return "", err + } + content, err := os.ReadFile(source) + if err != nil { + return "", err + } + name := stagedImageName(source) + if err := os.WriteFile(filepath.Join(c.ImageInputDir, name), content, 0o600); err != nil { + return "", err + } + return name, nil +} + +func stagedImageName(source string) string { + cleanPath := filepath.Clean(source) + base := filepath.Base(cleanPath) + ext := filepath.Ext(base) + stem := strings.TrimSuffix(base, ext) + sum := sha256.Sum256([]byte(cleanPath)) + hash := hex.EncodeToString(sum[:6]) + if stem == "" { + return hash + ext + } + return stem + "-" + hash + ext +} + func (c Client) GetJob(ctx context.Context, jobID string) (api.JobStatus, error) { return api.Client{BaseURL: c.BaseURL}.GetJob(ctx, jobID) } diff --git a/internal/cli/client_test.go b/internal/cli/client_test.go index c80ce2d..6a00895 100644 --- a/internal/cli/client_test.go +++ b/internal/cli/client_test.go @@ -6,6 +6,8 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "os" + "path/filepath" "testing" api "github.com/walker1211/codex-imgen/internal/api" @@ -204,3 +206,34 @@ func TestClientSubmitRequestIncludesImages(t *testing.T) { t.Fatalf("CreateJob returned error: %v", err) } } + +func TestClientSubmitStagesImagesIntoInputDir(t *testing.T) { + source := filepath.Join(t.TempDir(), "ref.png") + if err := os.WriteFile(source, []byte("png"), 0o644); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + inputDir := t.TempDir() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req api.CreateJobRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Fatalf("Decode returned error: %v", err) + } + if len(req.Images) != 1 { + t.Fatalf("images = %v", req.Images) + } + if !filepath.IsLocal(req.Images[0]) { + t.Fatalf("image path should be relative to input dir: %v", req.Images) + } + if _, err := os.Stat(filepath.Join(inputDir, req.Images[0])); err != nil { + t.Fatalf("staged image missing: %v", err) + } + _ = json.NewEncoder(w).Encode(map[string]any{"ok": true, "data": map[string]any{"job_id": "job_123", "status": "queued"}}) + })) + defer server.Close() + + _, err := (&Client{BaseURL: server.URL, ImageInputDir: inputDir}).CreateJob(context.Background(), "draw a dragon", []string{source}, 1, 1) + if err != nil { + t.Fatalf("CreateJob returned error: %v", err) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 9e09541..bf81dcf 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -18,8 +18,9 @@ type ServerConfig struct { } type StorageConfig struct { - DataDir string `yaml:"data_dir"` - SQLitePath string `yaml:"sqlite_path"` + DataDir string `yaml:"data_dir"` + SQLitePath string `yaml:"sqlite_path"` + ImageInputDir string `yaml:"image_input_dir"` } type SchedulerConfig struct { diff --git a/internal/config/load.go b/internal/config/load.go index 818ce73..2bd3710 100644 --- a/internal/config/load.go +++ b/internal/config/load.go @@ -33,6 +33,7 @@ func Load(path string) (Config, error) { } cfg.Storage.DataDir = expandHome(cfg.Storage.DataDir) cfg.Storage.SQLitePath = expandHome(cfg.Storage.SQLitePath) + cfg.Storage.ImageInputDir = expandHome(cfg.Storage.ImageInputDir) cfg.Backend.CWD = expandHome(cfg.Backend.CWD) cfg.Backend.DeliveryDir = expandHome(cfg.Backend.DeliveryDir) applyEnv(&cfg) diff --git a/internal/config/load_test.go b/internal/config/load_test.go index cbc528c..1d93e4b 100644 --- a/internal/config/load_test.go +++ b/internal/config/load_test.go @@ -120,6 +120,29 @@ func TestLoadParsesSchedulerJobConcurrencyFields(t *testing.T) { } } +func TestLoadParsesStorageImageInputDir(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.yaml") + content := []byte(`storage: + image_input_dir: ~/imgen-inputs +`) + if err := os.WriteFile(path, content, 0o644); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + cfg, err := Load(path) + if err != nil { + t.Fatalf("Load returned error: %v", err) + } + home, err := os.UserHomeDir() + if err != nil { + t.Fatalf("UserHomeDir returned error: %v", err) + } + want := filepath.Join(home, "imgen-inputs") + if cfg.Storage.ImageInputDir != want { + t.Fatalf("storage.image_input_dir = %q, want %q", cfg.Storage.ImageInputDir, want) + } +} + func TestLoadParsesBackendDeliveryDir(t *testing.T) { path := filepath.Join(t.TempDir(), "config.yaml") content := []byte(`backend: diff --git a/internal/service/service.go b/internal/service/service.go index a4cf1f6..a49b5ee 100644 --- a/internal/service/service.go +++ b/internal/service/service.go @@ -43,6 +43,7 @@ type Service struct { DefaultJobConcurrency int MaxJobConcurrency int MaxCountPerJob int + ImageInputDir string Now func() time.Time RetryDelays []time.Duration MaxAttempts int @@ -90,23 +91,40 @@ func (s *Service) publish(event notify.Event) { } } -func normalizeImagePaths(paths []string) ([]string, error) { +func normalizeImagePaths(paths []string, inputDir string) ([]string, error) { + root := strings.TrimSpace(inputDir) + if root == "" { + var err error + root, err = os.Getwd() + if err != nil { + return nil, err + } + } + root, err := filepath.Abs(root) + if err != nil { + return nil, err + } + var result []string - for _, path := range paths { - trimmed := strings.TrimSpace(path) + for _, imagePath := range paths { + trimmed := strings.TrimSpace(imagePath) if trimmed == "" { continue } - abs, err := filepath.Abs(trimmed) - if err != nil { - return nil, err + if !filepath.IsLocal(trimmed) { + return nil, errors.New("image path must stay within image input directory") } - if _, err := os.Stat(abs); err != nil { + abs := filepath.Join(root, filepath.Clean(trimmed)) + info, err := os.Stat(abs) + if err != nil { if os.IsNotExist(err) { - return nil, fmt.Errorf("image path not found: %s", abs) + return nil, errors.New("image path not found") } return nil, err } + if !info.Mode().IsRegular() { + return nil, errors.New("image path must be a regular file") + } result = append(result, abs) } return result, nil @@ -141,7 +159,7 @@ func (s *Service) CreateJob(req api.CreateJobRequest) (api.CreateJobResult, erro if strings.TrimSpace(req.Prompt) == "" { return api.CreateJobResult{}, errors.New("prompt is required") } - normalizedImages, err := normalizeImagePaths(req.Images) + normalizedImages, err := normalizeImagePaths(req.Images, s.ImageInputDir) if err != nil { return api.CreateJobResult{}, err } diff --git a/internal/service/service_test.go b/internal/service/service_test.go index c66e068..ca1ba44 100644 --- a/internal/service/service_test.go +++ b/internal/service/service_test.go @@ -848,16 +848,12 @@ func waitForServiceQueueTestEntry(t *testing.T, entered <-chan backend.GenerateR } } -func TestServiceCreateJobStoresAbsoluteImagePaths(t *testing.T) { - img := filepath.Join(t.TempDir(), "1.png") +func TestServiceCreateJobStoresInputDirImagePaths(t *testing.T) { + inputDir := t.TempDir() + img := filepath.Join(inputDir, "1.png") if err := os.WriteFile(img, []byte("png"), 0o644); err != nil { t.Fatalf("WriteFile failed: %v", err) } - cwd, _ := os.Getwd() - rel, err := filepath.Rel(cwd, img) - if err != nil { - t.Fatalf("Rel returned error: %v", err) - } dbPath := filepath.Join(t.TempDir(), "imgen.db") s, err := store.Open(dbPath) @@ -866,8 +862,8 @@ func TestServiceCreateJobStoresAbsoluteImagePaths(t *testing.T) { } defer s.Close() - svc := Service{Store: s} - created, err := svc.CreateJob(api.CreateJobRequest{Prompt: "draw a dragon", Images: []string{rel}}) + svc := Service{Store: s, ImageInputDir: inputDir} + created, err := svc.CreateJob(api.CreateJobRequest{Prompt: "draw a dragon", Images: []string{"1.png"}}) if err != nil { t.Fatalf("CreateJob returned error: %v", err) } @@ -880,6 +876,28 @@ func TestServiceCreateJobStoresAbsoluteImagePaths(t *testing.T) { } } +func TestServiceCreateJobRejectsImagePathOutsideInputDir(t *testing.T) { + outside := filepath.Join(t.TempDir(), "secret.png") + if err := os.WriteFile(outside, []byte("png"), 0o644); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + dbPath := filepath.Join(t.TempDir(), "imgen.db") + s, err := store.Open(dbPath) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + defer s.Close() + + svc := Service{Store: s, ImageInputDir: t.TempDir()} + _, err = svc.CreateJob(api.CreateJobRequest{Prompt: "draw a dragon", Images: []string{outside}}) + if err == nil || !strings.Contains(err.Error(), "image path must stay within image input directory") { + t.Fatalf("err = %v", err) + } + if strings.Contains(err.Error(), outside) { + t.Fatalf("error leaked rejected path: %v", err) + } +} + func TestServiceCreateJobRejectsMissingImagePath(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "imgen.db") s, err := store.Open(dbPath) @@ -888,15 +906,16 @@ func TestServiceCreateJobRejectsMissingImagePath(t *testing.T) { } defer s.Close() - svc := Service{Store: s} - _, err = svc.CreateJob(api.CreateJobRequest{Prompt: "draw a dragon", Images: []string{"./missing.png"}}) + svc := Service{Store: s, ImageInputDir: t.TempDir()} + _, err = svc.CreateJob(api.CreateJobRequest{Prompt: "draw a dragon", Images: []string{"missing.png"}}) if err == nil || !strings.Contains(err.Error(), "image path not found") { t.Fatalf("err = %v", err) } } func TestServiceRunJobPassesImagesToGenerator(t *testing.T) { - img := filepath.Join(t.TempDir(), "1.png") + inputDir := t.TempDir() + img := filepath.Join(inputDir, "1.png") if err := os.WriteFile(img, []byte("png"), 0o644); err != nil { t.Fatalf("WriteFile failed: %v", err) } @@ -907,8 +926,8 @@ func TestServiceRunJobPassesImagesToGenerator(t *testing.T) { t.Fatalf("Open returned error: %v", err) } defer s.Close() - svc := Service{Store: s, Generator: gen, PromptPrefix: "$imagegen", RetryDelays: []time.Duration{time.Millisecond}} - created, err := svc.CreateJob(api.CreateJobRequest{Prompt: "draw a dragon", Images: []string{img}}) + svc := Service{Store: s, Generator: gen, PromptPrefix: "$imagegen", RetryDelays: []time.Duration{time.Millisecond}, ImageInputDir: inputDir} + created, err := svc.CreateJob(api.CreateJobRequest{Prompt: "draw a dragon", Images: []string{"1.png"}}) if err != nil { t.Fatalf("CreateJob returned error: %v", err) }