config AgentConfig
}
+// Images returns the shared data-directory image store, also for a child.
+func (a *Agent) Images() *ImageStore {
+ root := a.dir
+ if a.config.Parent != "" {
+ root = filepath.Dir(filepath.Dir(root))
+ }
+
+ return &ImageStore{dir: filepath.Join(filepath.Dir(root), imageDirName)}
+}
+
// CacheKey returns the prompt cache key of every model call of the agent.
func (a *Agent) CacheKey() string {
return "onechat-" + a.config.ID
return true
}
- if filepath.Base(name) == "lock" {
+ if filepath.Base(name) == "lock" || (first == imageDirName && strings.HasPrefix(filepath.Base(name), ".image-")) {
return true
}
require (
github.com/stretchr/testify v1.12.0
go.mau.fi/whatsmeow v0.0.0-20261107160454-79dc2a8c33df
+ golang.org/x/image v0.47.0
golang.org/x/sys v0.49.0
golang.org/x/text v0.43.0
modernc.org/sqlite v1.62.0
golang.org/x/crypto v0.58.0/go.mod h1:Xh+kl5A+M0gMiWDytH2Hhr5Z7WbNCpv/eB/EFgfSk+w=
golang.org/x/exp v0.0.0-20261014143641-f0251eb4aa3e h1:HgI6W1OolKK6tcV74is+rl9SMT0njsHlp0TbQiGXI/A=
golang.org/x/exp v0.0.0-20261014143641-f0251eb4aa3e/go.mod h1:DBoxwWhjaW8DJJw/YeM4zpeOKMC1FMXSblsYy+1mFSw=
+golang.org/x/image v0.47.0 h1:za/8gVz++PPZqmZJBNSISICN5Rrj5zixP/8K++fgmqo=
+golang.org/x/image v0.47.0/go.mod h1:IkN5o4A7vq55vXN8FWjjnhg9Rs2GOtffqsSi9mWAKYE=
golang.org/x/mod v0.42.0 h1:ICeL3t2L30ljSPHHxDnR4fqVcBt/zR2zJAb9juavmQg=
golang.org/x/mod v0.42.0/go.mod h1:u5fzIrRInkL8MaO1cbRkUHT64v303YIT5KEqpNr7jRs=
golang.org/x/net v0.61.0 h1:RR+6j/BTBCrggOLugxa+66R8FvoN6MHlzDlZIbN9YFw=
--- /dev/null
+package main
+
+import (
+ "bytes"
+ "crypto/sha256"
+ "encoding/hex"
+ "errors"
+ "fmt"
+ "image"
+ _ "image/jpeg"
+ _ "image/png"
+ "io"
+ "os"
+ "path/filepath"
+ "strings"
+
+ _ "golang.org/x/image/webp"
+)
+
+const (
+ imageDirName = "images"
+ imageMaxBytes = 20 << 20
+)
+
+// ImageStore holds immutable images shared by the root, children, and
+// subagents. Names are the SHA-256 of the bytes with the decoded format as
+// extension. Publication is atomic and does not replace an existing image.
+type ImageStore struct {
+ dir string
+}
+
+// Save validates and stores data, and returns its filename. Only JPEG, PNG, and
+// WebP images up to 20 MB are accepted.
+func (s *ImageStore) Save(data []byte) (name string, err error) {
+ format, err := imageFormat(data)
+ if err != nil {
+ return "", fmt.Errorf("check image: %w", err)
+ }
+
+ sum := sha256.Sum256(data)
+ name = hex.EncodeToString(sum[:]) + "." + format
+ err = os.MkdirAll(s.dir, 0o700)
+ if err != nil {
+ return "", fmt.Errorf("create image directory: %w", err)
+ }
+
+ file, err := os.CreateTemp(s.dir, ".image-*")
+ if err != nil {
+ return "", fmt.Errorf("create image file: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, os.Remove(file.Name()))
+ }()
+
+ _, err = file.Write(data)
+ if err == nil {
+ err = file.Sync()
+ }
+
+ err = errors.Join(err, file.Close())
+ if err != nil {
+ return "", fmt.Errorf("write image: %w", err)
+ }
+
+ err = os.Link(file.Name(), filepath.Join(s.dir, name))
+ if err != nil && !errors.Is(err, os.ErrExist) {
+ return "", fmt.Errorf("publish image: %w", err)
+ }
+
+ return name, nil
+}
+
+// Read returns the bytes of a stored image after checking its name and content.
+func (s *ImageStore) Read(name string) ([]byte, error) {
+ if !imageName(name) {
+ return nil, errors.New("invalid image name")
+ }
+
+ data, err := imageReadFile(filepath.Join(s.dir, name))
+ if err != nil {
+ return nil, fmt.Errorf("read stored image: %w", err)
+ }
+
+ format, err := imageFormat(data)
+ if err != nil {
+ return nil, fmt.Errorf("check stored image: %w", err)
+ }
+
+ sum := sha256.Sum256(data)
+ if name != hex.EncodeToString(sum[:])+"."+format {
+ return nil, errors.New("image content does not match its name")
+ }
+
+ return data, nil
+}
+
+func imageFormat(data []byte) (string, error) {
+ if len(data) > imageMaxBytes {
+ return "", errors.New("image exceeds 20 MB")
+ }
+
+ config, format, err := image.DecodeConfig(bytes.NewReader(data))
+ if err != nil || config.Width <= 0 || config.Height <= 0 {
+ return "", errors.New("use a JPEG, PNG, or WebP image")
+ }
+
+ switch format {
+ case "jpeg":
+ return "jpg", nil
+ case "png", "webp":
+ return format, nil
+ default:
+ return "", errors.New("use a JPEG, PNG, or WebP image")
+ }
+}
+
+func imageName(name string) bool {
+ hash, ext, found := strings.Cut(name, ".")
+ return found && len(hash) == 64 && strings.Trim(hash, "0123456789abcdef") == "" && (ext == "jpg" || ext == "png" || ext == "webp")
+}
+
+func imageReadFile(path string) (data []byte, err error) {
+ file, err := os.Open(path)
+ if err != nil {
+ return nil, fmt.Errorf("open image: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, file.Close())
+ }()
+
+ data, err = io.ReadAll(io.LimitReader(file, imageMaxBytes+1))
+ if err != nil {
+ return nil, fmt.Errorf("read image: %w", err)
+ }
+
+ if len(data) > imageMaxBytes {
+ return nil, errors.New("image exceeds 20 MB")
+ }
+
+ return data, nil
+}
+
+// imageFiles reads only complete marker lines with valid store filenames.
+func imageFiles(text string) []string {
+ var names []string
+
+ for line := range strings.SplitSeq(text, "\n") {
+ name, found := strings.CutPrefix(line, "[image ")
+ if !found || !strings.HasSuffix(name, "]") {
+ continue
+ }
+
+ name = strings.TrimSuffix(name, "]")
+ if imageName(name) {
+ names = append(names, name)
+ }
+ }
+
+ return names
+}
+
+func imageText(names []string, caption string) string {
+ var text strings.Builder
+
+ for _, name := range names {
+ text.WriteString("[image " + name + "]\n")
+ }
+
+ text.WriteString(caption)
+
+ return strings.TrimSuffix(text.String(), "\n")
+}
--- /dev/null
+package main
+
+import (
+ "bytes"
+ "image"
+ "image/color"
+ "image/png"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func imageTestPNG(t *testing.T, width, height int) []byte {
+ t.Helper()
+
+ picture := image.NewRGBA(image.Rect(0, 0, width, height))
+ picture.Set(0, 0, color.RGBA{
+ R: 255,
+ A: 255,
+ })
+
+ var data bytes.Buffer
+
+ require.NoError(t, png.Encode(&data, picture))
+
+ return data.Bytes()
+}
--- /dev/null
+package main
+
+import (
+ "bytes"
+ "crypto/sha256"
+ "encoding/hex"
+ "os"
+ "path/filepath"
+ "sync"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestImagesArePrivateSharedAndWrittenOnce(t *testing.T) {
+ root, err := openAgent(filepath.Join(t.TempDir(), "agent"))
+ require.NoError(t, err)
+
+ child := &Agent{
+ dir: filepath.Join(root.dir, "children", "House"),
+ config: AgentConfig{Parent: root.config.Name},
+ }
+ require.Equal(t, root.Images().dir, child.Images().dir)
+
+ data := imageTestPNG(t, 3, 2)
+ sum := sha256.Sum256(data)
+ expected := hex.EncodeToString(sum[:]) + ".png"
+ store := root.Images()
+ name, err := store.Save(data)
+ require.NoError(t, err)
+ require.Equal(t, expected, name)
+
+ info, err := os.Stat(filepath.Join(store.dir, name))
+ require.NoError(t, err)
+ require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
+
+ var workers sync.WaitGroup
+
+ for range 8 {
+ workers.Go(func() {
+ saved, saveErr := child.Images().Save(data)
+ require.NoError(t, saveErr)
+ require.Equal(t, name, saved)
+ })
+ }
+
+ workers.Wait()
+
+ again, err := os.Stat(filepath.Join(store.dir, name))
+ require.NoError(t, err)
+ require.True(t, os.SameFile(info, again))
+
+ files, err := os.ReadDir(store.dir)
+ require.NoError(t, err)
+ require.Len(t, files, 1)
+
+ actual, err := store.Read(name)
+ require.NoError(t, err)
+ require.Equal(t, data, actual)
+}
+
+func TestImageStoreRejectsLargeUnsupportedAndUnsafeImages(t *testing.T) {
+ store := &ImageStore{dir: filepath.Join(t.TempDir(), imageDirName)}
+ _, err := store.Save(bytes.Repeat([]byte{0}, imageMaxBytes+1))
+ require.ErrorContains(t, err, "exceeds 20 MB")
+
+ _, err = store.Save([]byte("not an image"))
+ require.ErrorContains(t, err, "JPEG, PNG, or WebP")
+
+ _, err = store.Read("../secret.png")
+ require.ErrorContains(t, err, "invalid image name")
+
+ _, err = os.Stat(store.dir)
+ require.ErrorIs(t, err, os.ErrNotExist)
+}
+
+func TestBackupKeepsImagesAndTheirModes(t *testing.T) {
+ dir := t.TempDir()
+ store := &ImageStore{dir: filepath.Join(dir, imageDirName)}
+ data := imageTestPNG(t, 2, 2)
+ name, err := store.Save(data)
+ require.NoError(t, err)
+ require.NoError(t, os.WriteFile(filepath.Join(store.dir, ".image-incomplete"), []byte("partial"), 0o600))
+
+ stage := t.TempDir()
+ archive := filepath.Join(t.TempDir(), "images.tar.gz")
+ require.NoError(t, backupCopy(t.Context(), dir, stage, archive))
+
+ _, err = backupArchive(t.Context(), stage, archive)
+ require.NoError(t, err)
+
+ restored := t.TempDir()
+ require.NoError(t, backupRestore(t.Context(), archive, restored))
+
+ actual, err := (&ImageStore{dir: filepath.Join(restored, imageDirName)}).Read(name)
+ require.NoError(t, err)
+ require.Equal(t, data, actual)
+
+ info, err := os.Stat(filepath.Join(restored, imageDirName, name))
+ require.NoError(t, err)
+ require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
+
+ _, err = os.Stat(filepath.Join(restored, imageDirName, ".image-incomplete"))
+ require.ErrorIs(t, err, os.ErrNotExist)
+}
+
+func TestImageMarkersKeepCaptionAndOnlySafeNames(t *testing.T) {
+ name := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.png"
+ text := imageText([]string{name}, "Read the label.")
+ require.Equal(t, "[image "+name+"]\nRead the label.", text)
+ require.Equal(t, []string{name}, imageFiles(text+"\n[image ../secret.png]\ntext [image "+name+"]"))
+}
subPackages = [ "." ];
ldflags = [ "-X main.version=${version}" ];
- vendorHash = "sha256-HEYbx27w3s+x29/gDLUOCrkBJu9Ke1ue+Ms++LdHpyQ=";
+ vendorHash = "sha256-nRDUAGSXAK5/L2xyfDh3NbMOZmrGGEIm4vtgfWsxHYU=";
# The flake check runs the full replay.
checkFlags = [ "-short" ];
projects / onechat / commitdiff