Antonio Mika
·
2025-03-12
handler.go
1package uploadimgs
2
3import (
4 "bytes"
5 "encoding/binary"
6 "fmt"
7 "io"
8 "net/http"
9 "os"
10 "path/filepath"
11 "slices"
12 "strings"
13
14 exifremove "github.com/neurosnap/go-exif-remove"
15 "github.com/picosh/pico/pkg/db"
16 "github.com/picosh/pico/pkg/pobj"
17 sst "github.com/picosh/pico/pkg/pobj/storage"
18 "github.com/picosh/pico/pkg/pssh"
19 sendutils "github.com/picosh/pico/pkg/send/utils"
20 "github.com/picosh/pico/pkg/shared"
21 "github.com/picosh/pico/pkg/shared/storage"
22 "github.com/picosh/utils"
23)
24
25var Space = "imgs"
26
27type PostMetaData struct {
28 Text []byte
29 FileSize int
30 TotalFileSize int
31 Filename string
32 User *db.User
33 FeatureFlag *db.FeatureFlag
34 Bucket sst.Bucket
35}
36
37type UploadImgHandler struct {
38 DBPool db.DB
39 Cfg *shared.ConfigSite
40 Storage storage.StorageServe
41}
42
43func NewUploadImgHandler(dbpool db.DB, cfg *shared.ConfigSite, storage storage.StorageServe) *UploadImgHandler {
44 return &UploadImgHandler{
45 DBPool: dbpool,
46 Cfg: cfg,
47 Storage: storage,
48 }
49}
50
51func (h *UploadImgHandler) getObjectPath(fpath string) string {
52 return filepath.Join("prose", fpath)
53}
54
55func (h *UploadImgHandler) List(s *pssh.SSHServerConnSession, fpath string, isDir bool, recursive bool) ([]os.FileInfo, error) {
56 var fileList []os.FileInfo
57
58 logger := pssh.GetLogger(s)
59 user := pssh.GetUser(s)
60
61 if user == nil {
62 err := fmt.Errorf("could not get user from ctx")
63 logger.Error("error getting user from ctx", "err", err)
64 return fileList, err
65 }
66
67 cleanFilename := fpath
68
69 bucketName := shared.GetAssetBucketName(user.ID)
70 bucket, err := h.Storage.GetBucket(bucketName)
71 if err != nil {
72 return fileList, err
73 }
74
75 if cleanFilename == "" || cleanFilename == "." {
76 name := cleanFilename
77 if name == "" {
78 name = "/"
79 }
80
81 info := &sendutils.VirtualFile{
82 FName: name,
83 FIsDir: true,
84 }
85
86 fileList = append(fileList, info)
87 } else {
88 fp := h.getObjectPath(cleanFilename)
89 if fp != "/" && isDir {
90 fp += "/"
91 }
92
93 foundList, err := h.Storage.ListObjects(bucket, fp, recursive)
94 if err != nil {
95 return fileList, err
96 }
97
98 fileList = append(fileList, foundList...)
99 }
100
101 return fileList, nil
102}
103
104func (h *UploadImgHandler) Read(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) (os.FileInfo, sendutils.ReadAndReaderAtCloser, error) {
105 logger := pssh.GetLogger(s)
106 user := pssh.GetUser(s)
107
108 if user == nil {
109 err := fmt.Errorf("could not get user from ctx")
110 logger.Error("error getting user from ctx", "err", err)
111 return nil, nil, err
112 }
113
114 cleanFilename := filepath.Base(entry.Filepath)
115
116 if cleanFilename == "" || cleanFilename == "." {
117 return nil, nil, os.ErrNotExist
118 }
119
120 bucket, err := h.Storage.GetBucket(shared.GetAssetBucketName(user.ID))
121 if err != nil {
122 return nil, nil, err
123 }
124
125 contents, info, err := h.Storage.GetObject(bucket, h.getObjectPath(cleanFilename))
126 if err != nil {
127 return nil, nil, err
128 }
129 reader := pobj.NewAllReaderAt(contents)
130
131 fileInfo := &sendutils.VirtualFile{
132 FName: cleanFilename,
133 FIsDir: false,
134 FSize: info.Size,
135 FModTime: info.LastModified,
136 }
137
138 return fileInfo, reader, nil
139}
140
141func (h *UploadImgHandler) Write(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) (string, error) {
142 logger := pssh.GetLogger(s)
143 user := pssh.GetUser(s)
144
145 if user == nil {
146 err := fmt.Errorf("could not get user from ctx")
147 logger.Error("error getting user from ctx", "err", err)
148 return "", err
149 }
150
151 filename := filepath.Base(entry.Filepath)
152
153 var text []byte
154 if b, err := io.ReadAll(entry.Reader); err == nil {
155 text = b
156 }
157 mimeType := http.DetectContentType(text)
158 ext := filepath.Ext(filename)
159 if ext == ".svg" {
160 mimeType = "image/svg+xml"
161 }
162 // strip exif data
163 if slices.Contains([]string{"image/png", "image/jpg", "image/jpeg"}, mimeType) {
164 noExifBytes, err := exifremove.Remove(text)
165 if err == nil {
166 if len(noExifBytes) == 0 {
167 logger.Info("file silently failed to strip exif data", "filename", filename)
168 } else {
169 text = noExifBytes
170 logger.Info("stripped exif data", "filename", filename)
171 }
172 } else {
173 logger.Error("could not strip exif data", "err", err.Error())
174 }
175 }
176
177 fileSize := binary.Size(text)
178 featureFlag := shared.FindPlusFF(h.DBPool, h.Cfg, user.ID)
179
180 bucket, err := h.Storage.UpsertBucket(shared.GetAssetBucketName(user.ID))
181 if err != nil {
182 return "", err
183 }
184
185 totalFileSize, err := h.Storage.GetBucketQuota(bucket)
186 if err != nil {
187 logger.Error("bucket quota", "err", err)
188 return "", err
189 }
190
191 metadata := PostMetaData{
192 Filename: filename,
193 FileSize: fileSize,
194 Text: text,
195 User: user,
196 FeatureFlag: featureFlag,
197 Bucket: bucket,
198 TotalFileSize: int(totalFileSize),
199 }
200
201 err = h.writeImg(s, &metadata)
202 if err != nil {
203 logger.Error("could not write img", "err", err.Error())
204 return "", err
205 }
206
207 curl := shared.NewCreateURL(h.Cfg)
208 url := h.Cfg.FullPostURL(
209 curl,
210 user.Name,
211 metadata.Filename,
212 )
213 maxSize := int(featureFlag.Data.StorageMax)
214 str := fmt.Sprintf(
215 "%s (space: %.2f/%.2fGB, %.2f%%)",
216 url,
217 utils.BytesToGB(metadata.TotalFileSize+fileSize),
218 utils.BytesToGB(maxSize),
219 (float32(totalFileSize)/float32(maxSize))*100,
220 )
221 return str, nil
222}
223
224func (h *UploadImgHandler) Delete(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) error {
225 logger := pssh.GetLogger(s)
226 user := pssh.GetUser(s)
227
228 if user == nil {
229 err := fmt.Errorf("could not get user from ctx")
230 logger.Error("error getting user from ctx", "err", err)
231 return err
232 }
233
234 filename := filepath.Base(entry.Filepath)
235
236 logger = logger.With(
237 "filename", filename,
238 )
239
240 bucket, err := h.Storage.UpsertBucket(shared.GetAssetBucketName(user.ID))
241 if err != nil {
242 return err
243 }
244
245 logger.Info("deleting image")
246 err = h.Storage.DeleteObject(bucket, h.getObjectPath(filename))
247 if err != nil {
248 return err
249 }
250
251 return nil
252}
253
254func (h *UploadImgHandler) validateImg(data *PostMetaData) (bool, error) {
255 fileMax := data.FeatureFlag.Data.FileMax
256 if int64(data.FileSize) > fileMax {
257 return false, fmt.Errorf("ERROR: file (%s) has exceeded maximum file size (%d bytes)", data.Filename, fileMax)
258 }
259
260 storageMax := data.FeatureFlag.Data.StorageMax
261 if uint64(data.TotalFileSize+data.FileSize) > storageMax {
262 return false, fmt.Errorf("ERROR: user (%s) has exceeded (%d bytes) max (%d bytes)", data.User.Name, data.TotalFileSize, storageMax)
263 }
264
265 if !utils.IsExtAllowed(data.Filename, h.Cfg.AllowedExt) {
266 extStr := strings.Join(h.Cfg.AllowedExt, ",")
267 err := fmt.Errorf(
268 "ERROR: (%s) invalid file, format must be (%s), skipping",
269 data.Filename,
270 extStr,
271 )
272 return false, err
273 }
274
275 return true, nil
276}
277
278func (h *UploadImgHandler) metaImg(data *PostMetaData) error {
279 // if the file is empty that means we should delete it
280 // so we can skip all the meta info
281 if data.FileSize == 0 {
282 return nil
283 }
284
285 // make sure we have a bucket
286 bucket, err := h.Storage.UpsertBucket(shared.GetAssetBucketName(data.User.ID))
287 if err != nil {
288 return err
289 }
290
291 // make sure we have a prose project to upload to
292 _, err = h.DBPool.UpsertProject(data.User.ID, "prose", "prose")
293 if err != nil {
294 return err
295 }
296
297 reader := bytes.NewReader([]byte(data.Text))
298 _, _, err = h.Storage.PutObject(
299 bucket,
300 h.getObjectPath(data.Filename),
301 sendutils.NopReadAndReaderAtCloser(reader),
302 &sendutils.FileEntry{},
303 )
304 if err != nil {
305 return err
306 }
307
308 return nil
309}
310
311func (h *UploadImgHandler) writeImg(s *pssh.SSHServerConnSession, data *PostMetaData) error {
312 valid, err := h.validateImg(data)
313 if !valid {
314 return err
315 }
316
317 logger := pssh.GetLogger(s)
318 logger = logger.With(
319 "filename", data.Filename,
320 )
321
322 logger.Info("uploading image")
323 err = h.metaImg(data)
324 if err != nil {
325 logger.Error("meta img", "err", err)
326 return err
327 }
328
329 return nil
330}