v4: fix image attribution and registry config

6 files changed, 166 insertions, 57 deletions

diff --git a/internal/image/download.go b/internal/image/download.go
index 7a196fa..d66b946 100644
--- a/internal/image/download.go
+++ b/internal/image/download.go
@@ -111,7 +111,7 @@ func (d *Downloader) DownloadImage(ctx context.Context, result *SearchResult, ou
 		return fmt.Errorf("sync output file %q: %w", outputPath, err)
 	}
 
-	if attribution := d.provider.GetAttribution(result); attribution != "" {
+	if attribution := strings.TrimSpace(result.Attribution); attribution != "" {
 		attrPath := strings.TrimSuffix(outputPath, filepath.Ext(outputPath)) + "_attribution.txt"
 		if err := os.WriteFile(attrPath, []byte(attribution), 0o644); err != nil {
 			fmt.Fprintf(os.Stderr, "Warning: failed to save attribution: %v\n", err)
@@ -154,7 +154,11 @@ func (d *Downloader) DownloadBestMatchWithOptions(ctx context.Context, opts *Sea
 		if d.options != nil && d.options.OutputDir != "" {
 			outputDir = d.options.OutputDir
 		}
-		outputPath := filepath.Join(outputDir, filename)
+		outputPath, err := joinWithinBaseDir(outputDir, filename)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Warning: refusing unsafe output path %q: %v\n", filename, err)
+			continue
+		}
 
 		if err := d.DownloadImage(ctx, &result, outputPath); err == nil {
 			return &result, outputPath, nil
@@ -177,8 +181,8 @@ func (d *Downloader) generateFileName(word string, result *SearchResult, index i
 
 	filename = strings.ReplaceAll(filename, "{word}", sanitizeFileName(word))
 	if result != nil {
-		filename = strings.ReplaceAll(filename, "{source}", result.Source)
-		filename = strings.ReplaceAll(filename, "{id}", result.ID)
+		filename = strings.ReplaceAll(filename, "{source}", sanitizeFileName(result.Source))
+		filename = strings.ReplaceAll(filename, "{id}", sanitizeFileName(result.ID))
 	}
 	filename = strings.ReplaceAll(filename, "{index}", fmt.Sprintf("%d", index))
 
@@ -199,6 +203,32 @@ func (d *Downloader) generateFileName(word string, result *SearchResult, index i
 	return filename
 }
 
+func joinWithinBaseDir(baseDir, name string) (string, error) {
+	if strings.TrimSpace(baseDir) == "" {
+		baseDir = "."
+	}
+
+	cleanBase, err := filepath.Abs(baseDir)
+	if err != nil {
+		return "", fmt.Errorf("resolve base dir: %w", err)
+	}
+
+	fullPath, err := filepath.Abs(filepath.Join(cleanBase, name))
+	if err != nil {
+		return "", fmt.Errorf("resolve output path: %w", err)
+	}
+
+	rel, err := filepath.Rel(cleanBase, fullPath)
+	if err != nil {
+		return "", fmt.Errorf("relativize output path: %w", err)
+	}
+	if rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
+		return "", fmt.Errorf("path escapes base dir")
+	}
+
+	return fullPath, nil
+}
+
 func sanitizeFileName(name string) string {
 	replacer := strings.NewReplacer(
 		"/", "_",
diff --git a/internal/image/download_test.go b/internal/image/download_test.go
index 588d875..b9a6e31 100644
--- a/internal/image/download_test.go
+++ b/internal/image/download_test.go
@@ -13,9 +13,9 @@ type mockDownloaderProvider struct {
 	results       []SearchResult
 	searchErr     error
 	payload       string
-	attribution   string
 	searchQueries []string
 	downloadURLs  []string
+	getAttrCalls  int
 }
 
 func (m *mockDownloaderProvider) Name() string { return "mock" }
@@ -36,7 +36,8 @@ func (m *mockDownloaderProvider) Download(_ context.Context, url string) (io.Rea
 }
 
 func (m *mockDownloaderProvider) GetAttribution(*SearchResult) string {
-	return m.attribution
+	m.getAttrCalls++
+	return "provider attribution"
 }
 
 func TestDownloaderGenerateFileName_DataURIUsesPNG(t *testing.T) {
@@ -55,12 +56,11 @@ func TestDownloaderGenerateFileName_DataURIUsesPNG(t *testing.T) {
 	}
 }
 
-func TestDownloadImageWritesAttribution(t *testing.T) {
+func TestDownloadImageWritesResultAttribution(t *testing.T) {
 	t.Parallel()
 
 	provider := &mockDownloaderProvider{
-		payload:     "image-bytes",
-		attribution: "attribution text",
+		payload: "image-bytes",
 	}
 	d := NewDownloader(provider, &DownloadOptions{
 		OutputDir:         t.TempDir(),
@@ -72,9 +72,10 @@ func TestDownloadImageWritesAttribution(t *testing.T) {
 
 	outputPath := filepath.Join(d.options.OutputDir, "ябълка_gemini.png")
 	if err := d.DownloadImage(context.Background(), &SearchResult{
-		URL:    "https://example.com/image.png",
-		Source: Gemini,
-		ID:     "1",
+		URL:         "https://example.com/image.png",
+		Source:      Gemini,
+		ID:          "1",
+		Attribution: "result attribution text",
 	}, outputPath); err != nil {
 		t.Fatalf("DownloadImage() error = %v", err)
 	}
@@ -86,14 +87,17 @@ func TestDownloadImageWritesAttribution(t *testing.T) {
 	if string(data) != "image-bytes" {
 		t.Fatalf("downloaded file = %q, want %q", string(data), "image-bytes")
 	}
+	if provider.getAttrCalls != 0 {
+		t.Fatalf("GetAttribution() calls = %d, want 0", provider.getAttrCalls)
+	}
 
 	attrPath := strings.TrimSuffix(outputPath, filepath.Ext(outputPath)) + "_attribution.txt"
 	attr, err := os.ReadFile(attrPath)
 	if err != nil {
 		t.Fatalf("ReadFile(attribution) error = %v", err)
 	}
-	if string(attr) != "attribution text" {
-		t.Fatalf("attribution = %q, want %q", string(attr), "attribution text")
+	if string(attr) != "result attribution text" {
+		t.Fatalf("attribution = %q, want %q", string(attr), "result attribution text")
 	}
 }
 
@@ -103,13 +107,13 @@ func TestDownloadBestMatchWithOptions(t *testing.T) {
 	provider := &mockDownloaderProvider{
 		results: []SearchResult{
 			{
-				ID:     "1",
-				URL:    "https://example.com/image1.jpg",
-				Source: Gemini,
+				ID:          "1",
+				URL:         "https://example.com/image1.jpg",
+				Source:      Gemini,
+				Attribution: "result attribution text",
 			},
 		},
-		payload:     "image-bytes",
-		attribution: "attribution text",
+		payload: "image-bytes",
 	}
 	d := NewDownloader(provider, &DownloadOptions{
 		OutputDir:         t.TempDir(),
@@ -134,6 +138,49 @@ func TestDownloadBestMatchWithOptions(t *testing.T) {
 	}
 }
 
+func TestDownloadBestMatchWithOptions_SanitizesUnsafeProviderFields(t *testing.T) {
+	t.Parallel()
+
+	outputDir := t.TempDir()
+	provider := &mockDownloaderProvider{
+		results: []SearchResult{
+			{
+				ID:          "../escape/..//id",
+				URL:         "https://example.com/image1.jpg",
+				Source:      "../../outside/path",
+				Attribution: "safe attribution",
+			},
+		},
+		payload: "image-bytes",
+	}
+	d := NewDownloader(provider, &DownloadOptions{
+		OutputDir:         outputDir,
+		CreateDir:         true,
+		OverwriteExisting: true,
+		FileNamePattern:   "{word}_{source}_{id}",
+		MaxSizeBytes:      10 * 1024 * 1024,
+	})
+
+	_, path, err := d.DownloadBestMatchWithOptions(context.Background(), &SearchOptions{Query: "ябълка"})
+	if err != nil {
+		t.Fatalf("DownloadBestMatchWithOptions() error = %v", err)
+	}
+
+	rel, err := filepath.Rel(outputDir, path)
+	if err != nil {
+		t.Fatalf("filepath.Rel() error = %v", err)
+	}
+	if strings.HasPrefix(rel, "..") {
+		t.Fatalf("path escaped output dir: %q (rel=%q)", path, rel)
+	}
+	if strings.Contains(path, "..") {
+		t.Fatalf("path contains path traversal elements: %q", path)
+	}
+	if _, err := os.Stat(path); err != nil {
+		t.Fatalf("downloaded file missing: %v", err)
+	}
+}
+
 func TestDownloadImageRejectsNilResult(t *testing.T) {
 	t.Parallel()
 
diff --git a/internal/image/gemini.go b/internal/image/gemini.go
index 2ccd6fa..57fa61a 100644
--- a/internal/image/gemini.go
+++ b/internal/image/gemini.go
@@ -158,9 +158,9 @@ func (c *GeminiProvider) Search(ctx context.Context, opts *SearchOptions) ([]Sea
 		Width:        width,
 		Height:       height,
 		Description:  description,
-		Attribution:  "Generated by Google Gemini Nano Banana",
 		Source:       geminiSource,
 	}
+	result.Attribution = c.buildAttribution(&result, prompt)
 
 	return []SearchResult{result}, nil
 }
@@ -193,24 +193,11 @@ func (c *GeminiProvider) Download(ctx context.Context, url string) (io.ReadClose
 
 // GetAttribution returns attribution text for the generated image.
 func (c *GeminiProvider) GetAttribution(result *SearchResult) string {
-	width, height := 0, 0
-	if result != nil {
-		width = result.Width
-		height = result.Height
+	if result == nil {
+		return ""
 	}
 
-	var attribution strings.Builder
-	attribution.WriteString("Image generated by Google Gemini Nano Banana\n\n")
-	fmt.Fprintf(&attribution, "Model: %s\n", c.modelName())
-	fmt.Fprintf(&attribution, "Text model: %s\n", c.textModelName())
-	fmt.Fprintf(&attribution, "Aspect ratio: %s\n", geminiAspectRatio)
-	fmt.Fprintf(&attribution, "Size: %dx%d\n", width, height)
-	if result != nil && result.Description != "" {
-		fmt.Fprintf(&attribution, "Result: %s\n", result.Description)
-	}
-	fmt.Fprintf(&attribution, "\nPrompt used:\n%s\n", c.lastPrompt)
-	fmt.Fprintf(&attribution, "\nGenerated at: %s\n", time.Now().Format("2006-01-02 15:04:05"))
-	return attribution.String()
+	return result.Attribution
 }
 
 // Name returns the provider name.
@@ -562,6 +549,25 @@ func (c *GeminiProvider) generateImageID(word string) string {
 	return hex.EncodeToString(hash[:])[:8]
 }
 
+func (c *GeminiProvider) buildAttribution(result *SearchResult, prompt string) string {
+	if result == nil {
+		return ""
+	}
+
+	var attribution strings.Builder
+	attribution.WriteString("Image generated by Google Gemini Nano Banana\n\n")
+	fmt.Fprintf(&attribution, "Model: %s\n", c.modelName())
+	fmt.Fprintf(&attribution, "Text model: %s\n", c.textModelName())
+	fmt.Fprintf(&attribution, "Aspect ratio: %s\n", geminiAspectRatio)
+	fmt.Fprintf(&attribution, "Size: %dx%d\n", result.Width, result.Height)
+	if result.Description != "" {
+		fmt.Fprintf(&attribution, "Result: %s\n", result.Description)
+	}
+	fmt.Fprintf(&attribution, "\nPrompt used:\n%s\n", prompt)
+	fmt.Fprintf(&attribution, "\nGenerated at: %s\n", time.Now().Format("2006-01-02 15:04:05"))
+	return attribution.String()
+}
+
 func (c *GeminiProvider) modelName() string {
 	if c == nil || c.config == nil || strings.TrimSpace(c.config.Model) == "" {
 		return DefaultGeminiImageModel
diff --git a/internal/image/gemini_test.go b/internal/image/gemini_test.go
index 8c34c6e..bb3affe 100644
--- a/internal/image/gemini_test.go
+++ b/internal/image/gemini_test.go
@@ -188,6 +188,12 @@ func TestGeminiProvider_Search_GeneratedPromptFlow(t *testing.T) {
 	if !strings.Contains(gotPrompt, "Scene: A bright apple sits centered on a wooden table.") {
 		t.Fatalf("Prompt = %q, want generated scene in prompt", gotPrompt)
 	}
+	if result.Attribution == "" {
+		t.Fatal("expected result attribution to be populated")
+	}
+	if got := client.GetAttribution(&result); got != result.Attribution {
+		t.Fatalf("GetAttribution() = %q, want result attribution %q", got, result.Attribution)
+	}
 
 	reader, err := client.Download(context.Background(), result.URL)
 	if err != nil {
diff --git a/internal/image/registry.go b/internal/image/registry.go
index 4ddb40b..0338115 100644
--- a/internal/image/registry.go
+++ b/internal/image/registry.go
@@ -2,11 +2,12 @@ package image
 
 import (
 	"fmt"
+	"reflect"
 	"sync"
 )
 
 // Factory builds a provider instance.
-type Factory func() (ImageProvider, error)
+type Factory[C any] func(C) (ImageProvider, error)
 
 // Config exposes the configured image provider name.
 type Config interface {
@@ -14,18 +15,18 @@ type Config interface {
 }
 
 // Registry resolves provider names to factories.
-type Registry struct {
+type Registry[C Config] struct {
 	mu        sync.RWMutex
-	factories map[string]Factory
+	factories map[string]Factory[C]
 }
 
 // NewRegistry creates an empty provider registry.
-func NewRegistry() *Registry {
-	return &Registry{factories: make(map[string]Factory)}
+func NewRegistry[C Config]() *Registry[C] {
+	return &Registry[C]{factories: make(map[string]Factory[C])}
 }
 
 // Register associates name with a factory. Later registrations replace earlier ones.
-func (r *Registry) Register(name string, factory Factory) {
+func (r *Registry[C]) Register(name string, factory Factory[C]) {
 	if r == nil || factory == nil {
 		return
 	}
@@ -38,13 +39,13 @@ func (r *Registry) Register(name string, factory Factory) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	if r.factories == nil {
-		r.factories = make(map[string]Factory)
+		r.factories = make(map[string]Factory[C])
 	}
 	r.factories[normalized] = factory
 }
 
 // Resolve returns the factory registered for name.
-func (r *Registry) Resolve(name string) (Factory, bool) {
+func (r *Registry[C]) Resolve(name string) (Factory[C], bool) {
 	if r == nil {
 		return nil, false
 	}
@@ -56,19 +57,33 @@ func (r *Registry) Resolve(name string) (Factory, bool) {
 }
 
 // New constructs a provider for name.
-func (r *Registry) New(name string) (ImageProvider, error) {
+func (r *Registry[C]) New(name string, cfg C) (ImageProvider, error) {
 	var zero ImageProvider
 	factory, ok := r.Resolve(name)
 	if !ok {
 		return zero, fmt.Errorf("%w: %s", ErrUnknownProvider, name)
 	}
-	return factory()
+	return factory(cfg)
 }
 
 // NewFromConfig resolves the provider name from cfg and constructs it.
-func (r *Registry) NewFromConfig(cfg Config) (ImageProvider, error) {
-	if cfg == nil {
+func (r *Registry[C]) NewFromConfig(cfg C) (ImageProvider, error) {
+	if isNilValue(cfg) {
 		return nil, fmt.Errorf("image config is required")
 	}
-	return r.New(cfg.ImageProviderName())
+	return r.New(cfg.ImageProviderName(), cfg)
+}
+
+func isNilValue[T any](value T) bool {
+	v := reflect.ValueOf(value)
+	if !v.IsValid() {
+		return true
+	}
+
+	switch v.Kind() {
+	case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice:
+		return v.IsNil()
+	default:
+		return false
+	}
 }
diff --git a/internal/image/types_test.go b/internal/image/types_test.go
index c4e04f9..3597ee3 100644
--- a/internal/image/types_test.go
+++ b/internal/image/types_test.go
@@ -51,25 +51,28 @@ func TestSearchError(t *testing.T) {
 func TestRegistryNewFromConfig(t *testing.T) {
 	t.Parallel()
 
-	registry := NewRegistry()
-	registry.Register(Gemini, func() (ImageProvider, error) {
-		return fakeProvider{name: Gemini}, nil
+	registry := NewRegistry[fakeConfig]()
+	registry.Register(Gemini, func(cfg fakeConfig) (ImageProvider, error) {
+		return fakeProvider{name: cfg.name, token: cfg.token}, nil
 	})
 
-	provider, err := registry.NewFromConfig(fakeConfig{name: Gemini})
+	provider, err := registry.NewFromConfig(fakeConfig{name: Gemini, token: "secret"})
 	if err != nil {
 		t.Fatalf("NewFromConfig() error = %v", err)
 	}
 	if got, want := provider.Name(), Gemini; got != want {
 		t.Fatalf("provider.Name() = %q, want %q", got, want)
 	}
+	if got, want := provider.(fakeProvider).token, "secret"; got != want {
+		t.Fatalf("provider token = %q, want %q", got, want)
+	}
 }
 
 func TestRegistryUnknownProvider(t *testing.T) {
 	t.Parallel()
 
-	registry := NewRegistry()
-	_, err := registry.New("missing")
+	registry := NewRegistry[fakeConfig]()
+	_, err := registry.New("missing", fakeConfig{})
 	if err == nil {
 		t.Fatal("expected unknown provider error")
 	}
@@ -79,13 +82,15 @@ func TestRegistryUnknownProvider(t *testing.T) {
 }
 
 type fakeConfig struct {
-	name string
+	name  string
+	token string
 }
 
 func (f fakeConfig) ImageProviderName() string { return f.name }
 
 type fakeProvider struct {
-	name string
+	name  string
+	token string
 }
 
 func (f fakeProvider) Name() string { return f.name }