Eric Bower
·
20 Sep 24
handler.go
1package uploadassets
2
3import (
4 "bytes"
5 "fmt"
6 "io"
7 "io/fs"
8 "log/slog"
9 "os"
10 "path"
11 "path/filepath"
12 "slices"
13 "strings"
14 "time"
15
16 "github.com/charmbracelet/ssh"
17 "github.com/charmbracelet/wish"
18 "github.com/picosh/pico/db"
19 "github.com/picosh/pico/filehandlers/util"
20 futil "github.com/picosh/pico/filehandlers/util"
21 "github.com/picosh/pico/shared"
22 "github.com/picosh/pico/shared/storage"
23 "github.com/picosh/pobj"
24 sst "github.com/picosh/pobj/storage"
25 "github.com/picosh/send/send/utils"
26 ignore "github.com/sabhiram/go-gitignore"
27)
28
29type ctxBucketKey struct{}
30type ctxStorageSizeKey struct{}
31type ctxProjectKey struct{}
32type ctxDenylistKey struct{}
33
34type DenyList struct {
35 Denylist string
36}
37
38func getDenylist(s ssh.Session) *DenyList {
39 v := s.Context().Value(ctxDenylistKey{})
40 if v == nil {
41 return nil
42 }
43 denylist := s.Context().Value(ctxDenylistKey{}).(*DenyList)
44 return denylist
45}
46
47func setDenylist(s ssh.Session, denylist string) {
48 s.Context().SetValue(ctxDenylistKey{}, &DenyList{Denylist: denylist})
49}
50
51func getProject(s ssh.Session) *db.Project {
52 v := s.Context().Value(ctxProjectKey{})
53 if v == nil {
54 return nil
55 }
56 project := s.Context().Value(ctxProjectKey{}).(*db.Project)
57 return project
58}
59
60func setProject(s ssh.Session, project *db.Project) {
61 s.Context().SetValue(ctxProjectKey{}, project)
62}
63
64func getBucket(s ssh.Session) (sst.Bucket, error) {
65 bucket := s.Context().Value(ctxBucketKey{}).(sst.Bucket)
66 if bucket.Name == "" {
67 return bucket, fmt.Errorf("bucket not set on `ssh.Context()` for connection")
68 }
69 return bucket, nil
70}
71
72func getStorageSize(s ssh.Session) uint64 {
73 return s.Context().Value(ctxStorageSizeKey{}).(uint64)
74}
75
76func incrementStorageSize(s ssh.Session, fileSize int64) uint64 {
77 curSize := getStorageSize(s)
78 var nextStorageSize uint64
79 if fileSize < 0 {
80 nextStorageSize = curSize - uint64(fileSize)
81 } else {
82 nextStorageSize = curSize + uint64(fileSize)
83 }
84 s.Context().SetValue(ctxStorageSizeKey{}, nextStorageSize)
85 return nextStorageSize
86}
87
88func shouldIgnoreFile(fp, ignoreStr string) bool {
89 object := ignore.CompileIgnoreLines(strings.Split(ignoreStr, "\n")...)
90 return object.MatchesPath(fp)
91}
92
93type FileData struct {
94 *utils.FileEntry
95 User *db.User
96 Bucket sst.Bucket
97 Project *db.Project
98 DenyList string
99}
100
101type UploadAssetHandler struct {
102 DBPool db.DB
103 Cfg *shared.ConfigSite
104 Storage storage.StorageServe
105}
106
107func NewUploadAssetHandler(dbpool db.DB, cfg *shared.ConfigSite, storage storage.StorageServe) *UploadAssetHandler {
108 return &UploadAssetHandler{
109 DBPool: dbpool,
110 Cfg: cfg,
111 Storage: storage,
112 }
113}
114
115func (h *UploadAssetHandler) GetLogger() *slog.Logger {
116 return h.Cfg.Logger
117}
118
119func (h *UploadAssetHandler) Read(s ssh.Session, entry *utils.FileEntry) (os.FileInfo, utils.ReaderAtCloser, error) {
120 user, err := futil.GetUser(s.Context())
121 if err != nil {
122 return nil, nil, err
123 }
124
125 fileInfo := &utils.VirtualFile{
126 FName: filepath.Base(entry.Filepath),
127 FIsDir: false,
128 FSize: entry.Size,
129 FModTime: time.Unix(entry.Mtime, 0),
130 }
131
132 bucket, err := h.Storage.GetBucket(shared.GetAssetBucketName(user.ID))
133 if err != nil {
134 return nil, nil, err
135 }
136
137 fname := shared.GetAssetFileName(entry)
138 contents, size, modTime, err := h.Storage.GetObject(bucket, fname)
139 if err != nil {
140 return nil, nil, err
141 }
142
143 fileInfo.FSize = size
144 fileInfo.FModTime = modTime
145
146 reader := pobj.NewAllReaderAt(contents)
147
148 return fileInfo, reader, nil
149}
150
151func (h *UploadAssetHandler) List(s ssh.Session, fpath string, isDir bool, recursive bool) ([]os.FileInfo, error) {
152 var fileList []os.FileInfo
153
154 user, err := futil.GetUser(s.Context())
155 if err != nil {
156 return fileList, err
157 }
158
159 cleanFilename := fpath
160
161 bucketName := shared.GetAssetBucketName(user.ID)
162 bucket, err := h.Storage.GetBucket(bucketName)
163 if err != nil {
164 return fileList, err
165 }
166
167 if cleanFilename == "" || cleanFilename == "." {
168 name := cleanFilename
169 if name == "" {
170 name = "/"
171 }
172
173 info := &utils.VirtualFile{
174 FName: name,
175 FIsDir: true,
176 }
177
178 fileList = append(fileList, info)
179 } else {
180 if cleanFilename != "/" && isDir {
181 cleanFilename += "/"
182 }
183
184 foundList, err := h.Storage.ListObjects(bucket, cleanFilename, recursive)
185 if err != nil {
186 return fileList, err
187 }
188
189 fileList = append(fileList, foundList...)
190 }
191
192 return fileList, nil
193}
194
195func (h *UploadAssetHandler) Validate(s ssh.Session) error {
196 user, err := util.GetUser(s.Context())
197 if err != nil {
198 return err
199 }
200
201 assetBucket := shared.GetAssetBucketName(user.ID)
202 bucket, err := h.Storage.UpsertBucket(assetBucket)
203 if err != nil {
204 return err
205 }
206 s.Context().SetValue(ctxBucketKey{}, bucket)
207
208 totalStorageSize, err := h.Storage.GetBucketQuota(bucket)
209 if err != nil {
210 return err
211 }
212 s.Context().SetValue(ctxStorageSizeKey{}, totalStorageSize)
213 h.Cfg.Logger.Info(
214 "bucket size",
215 "user", user.Name,
216 "bytes", totalStorageSize,
217 )
218
219 h.Cfg.Logger.Info(
220 "attempting to upload files",
221 "user", user.Name,
222 "space", h.Cfg.Space,
223 )
224
225 return nil
226}
227
228func (h *UploadAssetHandler) findDenylist(bucket sst.Bucket, project *db.Project, logger *slog.Logger) (string, error) {
229 fp, _, _, err := h.Storage.GetObject(bucket, filepath.Join(project.ProjectDir, "_pgs_ignore"))
230 if err != nil {
231 return "", fmt.Errorf("_pgs_ignore not found")
232 }
233
234 defer fp.Close()
235 buf := new(strings.Builder)
236 _, err = io.Copy(buf, fp)
237 if err != nil {
238 logger.Error("io copy", "err", err.Error())
239 return "", err
240 }
241
242 str := buf.String()
243 return str, nil
244}
245
246func (h *UploadAssetHandler) Write(s ssh.Session, entry *utils.FileEntry) (string, error) {
247 user, err := futil.GetUser(s.Context())
248 if err != nil {
249 h.Cfg.Logger.Error("user not found in ctx", "err", err.Error())
250 return "", err
251 }
252
253 if entry.Mode.IsDir() && strings.Count(entry.Filepath, "/") == 1 {
254 entry.Filepath = strings.TrimPrefix(entry.Filepath, "/")
255 }
256
257 logger := h.GetLogger()
258 logger = shared.LoggerWithUser(logger, user)
259 logger = logger.With(
260 "file", entry.Filepath,
261 "size", entry.Size,
262 )
263
264 bucket, err := getBucket(s)
265 if err != nil {
266 logger.Error("could not find bucket in ctx", "err", err.Error())
267 return "", err
268 }
269
270 project := getProject(s)
271 projectName := shared.GetProjectName(entry)
272 logger = logger.With("project", projectName)
273
274 // find, create, or update project if we haven't already done it
275 if project == nil {
276 project, err = h.DBPool.FindProjectByName(user.ID, projectName)
277 if err == nil {
278 err = h.DBPool.UpdateProject(user.ID, projectName)
279 if err != nil {
280 logger.Error("could not update project", "err", err.Error())
281 return "", err
282 }
283 } else {
284 _, err = h.DBPool.InsertProject(user.ID, projectName, projectName)
285 if err != nil {
286 logger.Error("could not create project", "err", err.Error())
287 return "", err
288 }
289 project, err = h.DBPool.FindProjectByName(user.ID, projectName)
290 if err != nil {
291 logger.Error("could not find project", "err", err.Error())
292 return "", err
293 }
294 }
295 setProject(s, project)
296 }
297
298 if project.Blocked != "" {
299 msg := "project has been blocked and cannot upload files: %s"
300 return "", fmt.Errorf(msg, project.Blocked)
301 }
302
303 if entry.Mode.IsDir() {
304 _, _, err := h.Storage.PutObject(
305 bucket,
306 path.Join(shared.GetAssetFileName(entry), "._pico_keep_dir"),
307 bytes.NewReader([]byte{}),
308 entry,
309 )
310 return "", err
311 }
312
313 featureFlag, err := futil.GetFeatureFlag(s.Context())
314 if err != nil {
315 return "", err
316 }
317
318 // calculate the filsize difference between the same file already
319 // stored and the updated file being uploaded
320 assetFilename := shared.GetAssetFileName(entry)
321 curFileSize, _ := h.Storage.GetObjectSize(bucket, assetFilename)
322
323 denylist := getDenylist(s)
324 if denylist == nil {
325 dlist, err := h.findDenylist(bucket, project, logger)
326 if err != nil {
327 logger.Info("failed to get denylist, setting default (.*)", "err", err.Error())
328 dlist = ".*"
329 }
330 setDenylist(s, dlist)
331 denylist = &DenyList{Denylist: dlist}
332 }
333
334 data := &FileData{
335 FileEntry: entry,
336 User: user,
337 Bucket: bucket,
338 DenyList: denylist.Denylist,
339 Project: project,
340 }
341
342 valid, err := h.validateAsset(data)
343 if !valid {
344 return "", err
345 }
346
347 // SFTP does not report file size so the more performant way to
348 // check filesize constraints is to try and upload the file to s3
349 // with a specialized reader that raises an error if the filesize limit
350 // has been reached
351 storageMax := featureFlag.Data.StorageMax
352 fileMax := featureFlag.Data.FileMax
353 curStorageSize := getStorageSize(s)
354 remaining := int64(storageMax) - int64(curStorageSize)
355 sizeRemaining := min(remaining+curFileSize, fileMax)
356 if sizeRemaining <= 0 {
357 wish.Fatalln(s, "storage quota reached")
358 return "", fmt.Errorf("storage quota reached")
359 }
360 logger = logger.With(
361 "storage max", storageMax,
362 "current storage max", curStorageSize,
363 "file max", fileMax,
364 "sizeRemaining", sizeRemaining,
365 )
366
367 fsize, err := h.writeAsset(
368 shared.NewMaxBytesReader(data.Reader, int64(sizeRemaining)),
369 data,
370 )
371 if err != nil {
372 logger.Error("could not write asset", "err", err.Error())
373 cerr := fmt.Errorf(
374 "%s: storage size %.2fmb, storage max %.2fmb, file max %.2fmb",
375 err,
376 shared.BytesToMB(int(curStorageSize)),
377 shared.BytesToMB(int(storageMax)),
378 shared.BytesToMB(int(fileMax)),
379 )
380 return "", cerr
381 }
382
383 deltaFileSize := curFileSize - fsize
384 nextStorageSize := incrementStorageSize(s, deltaFileSize)
385
386 url := h.Cfg.AssetURL(
387 user.Name,
388 projectName,
389 strings.Replace(data.Filepath, "/"+projectName+"/", "", 1),
390 )
391
392 maxSize := int(featureFlag.Data.StorageMax)
393 str := fmt.Sprintf(
394 "%s (space: %.2f/%.2fGB, %.2f%%)",
395 url,
396 shared.BytesToGB(int(nextStorageSize)),
397 shared.BytesToGB(maxSize),
398 (float32(nextStorageSize)/float32(maxSize))*100,
399 )
400
401 return str, nil
402}
403
404func (h *UploadAssetHandler) Delete(s ssh.Session, entry *utils.FileEntry) error {
405 user, err := futil.GetUser(s.Context())
406 if err != nil {
407 h.Cfg.Logger.Error("user not found in ctx", "err", err.Error())
408 return err
409 }
410
411 if entry.Mode.IsDir() && strings.Count(entry.Filepath, "/") == 1 {
412 entry.Filepath = strings.TrimPrefix(entry.Filepath, "/")
413 }
414
415 assetFilepath := shared.GetAssetFileName(entry)
416
417 logger := h.GetLogger()
418 logger = shared.LoggerWithUser(logger, user)
419 logger = logger.With(
420 "file", assetFilepath,
421 )
422
423 bucket, err := getBucket(s)
424 if err != nil {
425 logger.Error("could not find bucket in ctx", "err", err.Error())
426 return err
427 }
428
429 projectName := shared.GetProjectName(entry)
430 logger = logger.With("project", projectName)
431
432 if assetFilepath == filepath.Join("/", projectName, "._pico_keep_dir") {
433 return os.ErrPermission
434 }
435
436 logger.Info("deleting file")
437
438 pathDir := filepath.Dir(assetFilepath)
439 fileName := filepath.Base(assetFilepath)
440
441 sibs, err := h.Storage.ListObjects(bucket, pathDir+"/", false)
442 if err != nil {
443 return err
444 }
445
446 sibs = slices.DeleteFunc(sibs, func(sib fs.FileInfo) bool {
447 return sib.Name() == fileName
448 })
449
450 if len(sibs) == 0 {
451 _, _, err := h.Storage.PutObject(
452 bucket,
453 filepath.Join(pathDir, "._pico_keep_dir"),
454 bytes.NewReader([]byte{}),
455 entry,
456 )
457 if err != nil {
458 return err
459 }
460 }
461
462 return h.Storage.DeleteObject(bucket, assetFilepath)
463}
464
465func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) {
466 fname := filepath.Base(data.Filepath)
467
468 projectName := shared.GetProjectName(data.FileEntry)
469 if projectName == "" || projectName == "/" || projectName == "." {
470 return false, fmt.Errorf("ERROR: invalid project name, you must copy files to a non-root folder (e.g. pgs.sh:/project-name)")
471 }
472
473 // special files we use for custom routing
474 if fname == "_pgs_ignore" || fname == "_redirects" || fname == "_headers" {
475 return true, nil
476 }
477
478 fpath := strings.Replace(data.Filepath, "/"+projectName, "", 1)
479 if shouldIgnoreFile(fpath, data.DenyList) {
480 err := fmt.Errorf(
481 "ERROR: (%s) file rejected, https://pico.sh/pgs#file-denylist",
482 data.Filepath,
483 )
484 return false, err
485 }
486
487 return true, nil
488}
489
490func (h *UploadAssetHandler) writeAsset(reader io.Reader, data *FileData) (int64, error) {
491 assetFilepath := shared.GetAssetFileName(data.FileEntry)
492
493 logger := shared.LoggerWithUser(h.Cfg.Logger, data.User)
494 logger.Info(
495 "uploading file to bucket",
496 "bucket", data.Bucket.Name,
497 "filename", assetFilepath,
498 )
499
500 _, fsize, err := h.Storage.PutObject(
501 data.Bucket,
502 assetFilepath,
503 reader,
504 data.FileEntry,
505 )
506 return fsize, err
507}