82 lines
2.3 KiB
Go
82 lines
2.3 KiB
Go
package artifact
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"mime"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
)
|
|
|
|
// TaskStyleReference is the style-only reference used as Image 2 for RIFT
|
|
// generation. Custom style assets live below <ARTIFACT_DIR>/_styles; an empty
|
|
// task reference resolves to the bundled default image.
|
|
type TaskStyleReference struct {
|
|
Bytes []byte `json:"-"`
|
|
Name string `json:"name"`
|
|
ContentType string `json:"content_type"`
|
|
Custom bool `json:"custom"`
|
|
SHA256 string `json:"sha256"`
|
|
}
|
|
|
|
func styleContentType(name string, b []byte) string {
|
|
ct := http.DetectContentType(b)
|
|
if ct == "application/octet-stream" {
|
|
if ext := filepath.Ext(name); ext != "" {
|
|
if v := mime.TypeByExtension(ext); v != "" {
|
|
ct = v
|
|
}
|
|
}
|
|
}
|
|
return strings.Split(ct, ";")[0]
|
|
}
|
|
|
|
func (w *Worker) loadStyleReference(ref string) (TaskStyleReference, error) {
|
|
ref = strings.TrimSpace(ref)
|
|
if ref == "" {
|
|
b := DefaultStyleReferenceJPEG()
|
|
h := sha256.Sum256(b)
|
|
return TaskStyleReference{
|
|
Bytes: b,
|
|
Name: "default_style_reference.jpg",
|
|
ContentType: "image/jpeg",
|
|
Custom: false,
|
|
SHA256: hex.EncodeToString(h[:]),
|
|
}, nil
|
|
}
|
|
if filepath.Base(ref) != ref || strings.ContainsAny(ref, `/\\`) {
|
|
return TaskStyleReference{}, fmt.Errorf("invalid task style reference %q", ref)
|
|
}
|
|
p := filepath.Join(w.dir, "_styles", ref)
|
|
b, err := os.ReadFile(p)
|
|
if err != nil {
|
|
return TaskStyleReference{}, fmt.Errorf("read task style reference %q: %w", ref, err)
|
|
}
|
|
if len(b) < 128 {
|
|
return TaskStyleReference{}, fmt.Errorf("task style reference %q is empty or invalid", ref)
|
|
}
|
|
h := sha256.Sum256(b)
|
|
return TaskStyleReference{
|
|
Bytes: b,
|
|
Name: ref,
|
|
ContentType: styleContentType(ref, b),
|
|
Custom: true,
|
|
SHA256: hex.EncodeToString(h[:]),
|
|
}, nil
|
|
}
|
|
|
|
// TaskStyleReference reads the configured style for a task. It is used by the
|
|
// UI preview endpoints and intentionally returns the bundled default when the
|
|
// task has no custom style assigned.
|
|
func (w *Worker) TaskStyleReference(ctx context.Context, taskID string) (TaskStyleReference, error) {
|
|
var ref string
|
|
if err := w.db.QueryRowContext(ctx, `SELECT nft_style_reference FROM tasks WHERE id=?`, taskID).Scan(&ref); err != nil {
|
|
return TaskStyleReference{}, err
|
|
}
|
|
return w.loadStyleReference(ref)
|
|
}
|