summaryrefslogtreecommitdiff
path: root/internal
diff options
context:
space:
mode:
Diffstat (limited to 'internal')
-rw-r--r--internal/image/download.go38
-rw-r--r--internal/image/download_test.go77
-rw-r--r--internal/image/gemini.go40
-rw-r--r--internal/image/gemini_test.go6
-rw-r--r--internal/image/registry.go41
-rw-r--r--internal/image/types_test.go21
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 }
generated by cgit v1.2.3 (git 2.25.1) at 2026-08-07 04:00:50 +0000