git.giacomolli.uk - onechat/blob - image.go

   1 package main
   2
   3 import (
   4         "bytes"
   5         "crypto/sha256"
   6         "encoding/hex"
   7         "errors"
   8         "fmt"
   9         "image"
  10         _ "image/jpeg"
  11         _ "image/png"
  12         "io"
  13         "os"
  14         "path/filepath"
  15         "strings"
  16
  17         _ "golang.org/x/image/webp"
  18 )
  19
  20 const (
  21         imageDirName  = "images"
  22         imageMaxBytes = 20 << 20
  23 )
  24
  25 // ImageStore holds immutable images shared by the root, children, and
  26 // subagents. Names are the SHA-256 of the bytes with the decoded format as
  27 // extension. Publication is atomic and does not replace an existing image.
  28 type ImageStore struct {
  29         dir string
  30 }
  31
  32 // Save validates and stores data, and returns its filename. Only JPEG, PNG, and
  33 // WebP images up to 20 MB are accepted.
  34 func (s *ImageStore) Save(data []byte) (name string, err error) {
  35         format, err := imageFormat(data)
  36         if err != nil {
  37                 return "", fmt.Errorf("check image: %w", err)
  38         }
  39
  40         sum := sha256.Sum256(data)
  41         name = hex.EncodeToString(sum[:]) + "." + format
  42         err = os.MkdirAll(s.dir, 0o700)
  43         if err != nil {
  44                 return "", fmt.Errorf("create image directory: %w", err)
  45         }
  46
  47         file, err := os.CreateTemp(s.dir, ".image-*")
  48         if err != nil {
  49                 return "", fmt.Errorf("create image file: %w", err)
  50         }
  51
  52         defer func() {
  53                 err = errors.Join(err, os.Remove(file.Name()))
  54         }()
  55
  56         _, err = file.Write(data)
  57         if err == nil {
  58                 err = file.Sync()
  59         }
  60
  61         err = errors.Join(err, file.Close())
  62         if err != nil {
  63                 return "", fmt.Errorf("write image: %w", err)
  64         }
  65
  66         err = os.Link(file.Name(), filepath.Join(s.dir, name))
  67         if err != nil && !errors.Is(err, os.ErrExist) {
  68                 return "", fmt.Errorf("publish image: %w", err)
  69         }
  70
  71         return name, nil
  72 }
  73
  74 // Read returns the bytes of a stored image after checking its name and content.
  75 func (s *ImageStore) Read(name string) ([]byte, error) {
  76         if !imageName(name) {
  77                 return nil, errors.New("invalid image name")
  78         }
  79
  80         data, err := imageReadFile(filepath.Join(s.dir, name))
  81         if err != nil {
  82                 return nil, fmt.Errorf("read stored image: %w", err)
  83         }
  84
  85         format, err := imageFormat(data)
  86         if err != nil {
  87                 return nil, fmt.Errorf("check stored image: %w", err)
  88         }
  89
  90         sum := sha256.Sum256(data)
  91         if name != hex.EncodeToString(sum[:])+"."+format {
  92                 return nil, errors.New("image content does not match its name")
  93         }
  94
  95         return data, nil
  96 }
  97
  98 func imageFormat(data []byte) (string, error) {
  99         if len(data) > imageMaxBytes {
 100                 return "", errors.New("image exceeds 20 MB")
 101         }
 102
 103         config, format, err := image.DecodeConfig(bytes.NewReader(data))
 104         if err != nil || config.Width <= 0 || config.Height <= 0 {
 105                 return "", errors.New("use a JPEG, PNG, or WebP image")
 106         }
 107
 108         switch format {
 109         case "jpeg":
 110                 return "jpg", nil
 111         case "png", "webp":
 112                 return format, nil
 113         default:
 114                 return "", errors.New("use a JPEG, PNG, or WebP image")
 115         }
 116 }
 117
 118 func imageName(name string) bool {
 119         hash, ext, found := strings.Cut(name, ".")
 120         return found && len(hash) == 64 && strings.Trim(hash, "0123456789abcdef") == "" && (ext == "jpg" || ext == "png" || ext == "webp")
 121 }
 122
 123 func imageReadFile(path string) (data []byte, err error) {
 124         file, err := os.Open(path)
 125         if err != nil {
 126                 return nil, fmt.Errorf("open image: %w", err)
 127         }
 128
 129         defer func() {
 130                 err = errors.Join(err, file.Close())
 131         }()
 132
 133         data, err = io.ReadAll(io.LimitReader(file, imageMaxBytes+1))
 134         if err != nil {
 135                 return nil, fmt.Errorf("read image: %w", err)
 136         }
 137
 138         if len(data) > imageMaxBytes {
 139                 return nil, errors.New("image exceeds 20 MB")
 140         }
 141
 142         return data, nil
 143 }
 144
 145 // imageFiles reads only complete marker lines with valid store filenames.
 146 func imageFiles(text string) []string {
 147         var names []string
 148
 149         for line := range strings.SplitSeq(text, "\n") {
 150                 name, found := strings.CutPrefix(line, "[image ")
 151                 if !found || !strings.HasSuffix(name, "]") {
 152                         continue
 153                 }
 154
 155                 name = strings.TrimSuffix(name, "]")
 156                 if imageName(name) {
 157                         names = append(names, name)
 158                 }
 159         }
 160
 161         return names
 162 }
 163
 164 func imageText(names []string, caption string) string {
 165         var text strings.Builder
 166
 167         for _, name := range names {
 168                 text.WriteString("[image " + name + "]\n")
 169         }
 170
 171         text.WriteString(caption)
 172
 173         return strings.TrimSuffix(text.String(), "\n")
 174 }