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