Eric Bower
·
2026-08-17
1package pgs
2
3import (
4 "bytes"
5 "context"
6 "fmt"
7 "io"
8 "io/fs"
9 "log/slog"
10 "os"
11 "path"
12 "path/filepath"
13 "slices"
14 "strings"
15 "sync"
16 "time"
17
18 "github.com/picosh/pico/pkg/db"
19 "github.com/picosh/pico/pkg/pssh"
20 sendutils "github.com/picosh/pico/pkg/send/utils"
21 "github.com/picosh/pico/pkg/shared"
22 "github.com/picosh/pico/pkg/storage"
23 ignore "github.com/sabhiram/go-gitignore"
24)
25
26type ctxBucketKey struct{}
27type ctxStorageSizeKey struct{}
28type ctxProjectKey struct{}
29type ctxFeatureFlagKey struct{}
30type ctxDenylistKey struct{}
31
32type DenyList struct {
33 Denylist string
34}
35
36func getDenylist(s *pssh.SSHServerConnSession) *DenyList {
37 v := s.Context().Value(ctxDenylistKey{})
38 if v == nil {
39 return nil
40 }
41 denylist := s.Context().Value(ctxDenylistKey{}).(*DenyList)
42 return denylist
43}
44
45func setDenylist(s *pssh.SSHServerConnSession, denylist string) {
46 s.SetValue(ctxDenylistKey{}, &DenyList{Denylist: denylist})
47}
48
49func getProject(s *pssh.SSHServerConnSession) *db.Project {
50 v := s.Context().Value(ctxProjectKey{})
51 if v == nil {
52 return nil
53 }
54 project := s.Context().Value(ctxProjectKey{}).(*db.Project)
55 return project
56}
57
58func setProject(s *pssh.SSHServerConnSession, project *db.Project) {
59 s.SetValue(ctxProjectKey{}, project)
60}
61
62func getFeatureFlag(s *pssh.SSHServerConnSession) *db.FeatureFlag {
63 v := s.Context().Value(ctxFeatureFlagKey{})
64 if v == nil {
65 return nil
66 }
67 ff := s.Context().Value(ctxFeatureFlagKey{}).(*db.FeatureFlag)
68 return ff
69}
70
71func setFeatureFlag(s *pssh.SSHServerConnSession, ff *db.FeatureFlag) {
72 s.SetValue(ctxFeatureFlagKey{}, ff)
73}
74
75func getBucket(s *pssh.SSHServerConnSession) (storage.Bucket, error) {
76 bucket := s.Context().Value(ctxBucketKey{}).(storage.Bucket)
77 if bucket.Name == "" {
78 return bucket, fmt.Errorf("bucket not set on `ssh.Context()` for connection")
79 }
80 return bucket, nil
81}
82
83func getStorageSize(s *pssh.SSHServerConnSession) uint64 {
84 return s.Context().Value(ctxStorageSizeKey{}).(uint64)
85}
86
87func incrementStorageSize(s *pssh.SSHServerConnSession, fileSize int64) uint64 {
88 curSize := getStorageSize(s)
89 var nextStorageSize uint64
90 if fileSize < 0 {
91 nextStorageSize = curSize - uint64(fileSize)
92 } else {
93 nextStorageSize = curSize + uint64(fileSize)
94 }
95 s.SetValue(ctxStorageSizeKey{}, nextStorageSize)
96 return nextStorageSize
97}
98
99func shouldIgnoreFile(fp, ignoreStr string) bool {
100 object := ignore.CompileIgnoreLines(strings.Split(ignoreStr, "\n")...)
101 return object.MatchesPath(fp)
102}
103
104type FileData struct {
105 *sendutils.FileEntry
106 User *db.User
107 Bucket storage.Bucket
108 Project *db.Project
109 DenyList string
110}
111
112type UploadAssetHandler struct {
113 Cfg *PgsConfig
114 CacheClearingQueue chan string
115}
116
117func NewUploadAssetHandler(cfg *PgsConfig, ch chan string, ctx context.Context) *UploadAssetHandler {
118 go runCacheQueue(cfg, ctx)
119 return &UploadAssetHandler{
120 Cfg: cfg,
121 CacheClearingQueue: ch,
122 }
123}
124
125func (h *UploadAssetHandler) GetLogger(s *pssh.SSHServerConnSession) *slog.Logger {
126 return pssh.GetLogger(s)
127}
128
129func (h *UploadAssetHandler) Read(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) (os.FileInfo, sendutils.ReadAndReaderAtCloser, error) {
130 logger := pssh.GetLogger(s)
131 user := pssh.GetUser(s)
132
133 if user == nil {
134 err := fmt.Errorf("could not get user from ctx")
135 logger.Error("error getting user from ctx", "err", err)
136 return nil, nil, err
137 }
138
139 fileInfo := &sendutils.VirtualFile{
140 FName: filepath.Base(entry.Filepath),
141 FIsDir: false,
142 FSize: entry.Size,
143 FModTime: time.Unix(entry.Mtime, 0),
144 }
145
146 bucket, err := h.Cfg.Storage.GetBucket(shared.GetAssetBucketName(user.ID))
147 if err != nil {
148 return nil, nil, err
149 }
150
151 fname := shared.GetAssetFileName(entry)
152 contents, info, err := h.Cfg.Storage.GetObject(bucket, fname)
153 if err != nil {
154 return nil, nil, err
155 }
156
157 fileInfo.FSize = info.Size
158 fileInfo.FModTime = info.LastModified
159
160 return fileInfo, contents, nil
161}
162
163func (h *UploadAssetHandler) List(s *pssh.SSHServerConnSession, fpath string, isDir bool, recursive bool) ([]os.FileInfo, error) {
164 var fileList []os.FileInfo
165
166 logger := pssh.GetLogger(s)
167 user := pssh.GetUser(s)
168
169 if user == nil {
170 err := fmt.Errorf("could not get user from ctx")
171 logger.Error("error getting user from ctx", "err", err)
172 return fileList, err
173 }
174
175 cleanFilename := fpath
176
177 bucketName := shared.GetAssetBucketName(user.ID)
178 bucket, err := h.Cfg.Storage.GetBucket(bucketName)
179 if err != nil {
180 return fileList, err
181 }
182
183 if cleanFilename == "" || cleanFilename == "." {
184 name := cleanFilename
185 if name == "" {
186 name = "/"
187 }
188
189 info := &sendutils.VirtualFile{
190 FName: name,
191 FIsDir: true,
192 }
193
194 fileList = append(fileList, info)
195 } else {
196 if cleanFilename != "/" && isDir {
197 cleanFilename += "/"
198 }
199
200 foundList, err := h.Cfg.Storage.ListObjects(bucket, cleanFilename, recursive)
201 if err != nil {
202 return fileList, err
203 }
204
205 fileList = append(fileList, foundList...)
206 }
207
208 return fileList, nil
209}
210
211func (h *UploadAssetHandler) Validate(s *pssh.SSHServerConnSession) error {
212 logger := pssh.GetLogger(s)
213 user := pssh.GetUser(s)
214
215 if user == nil {
216 err := fmt.Errorf("could not get user from ctx")
217 logger.Error("error getting user from ctx", "err", err)
218 return err
219 }
220
221 ff, err := findFeatureFlag(h.Cfg.DB, h.Cfg, user.ID)
222 if err != nil {
223 return err
224 }
225 setFeatureFlag(s, ff)
226
227 assetBucket := shared.GetAssetBucketName(user.ID)
228 bucket, err := h.Cfg.Storage.UpsertBucket(assetBucket)
229 if err != nil {
230 return err
231 }
232 s.SetValue(ctxBucketKey{}, bucket)
233
234 totalStorageSize, err := h.Cfg.Storage.GetBucketQuota(bucket)
235 if err != nil {
236 return err
237 }
238 s.SetValue(ctxStorageSizeKey{}, totalStorageSize)
239
240 logger.Info(
241 "bucket size",
242 "user", user.Name,
243 "bytes", totalStorageSize,
244 )
245
246 logger.Info(
247 "attempting to upload files",
248 "user", user.Name,
249 "txtPrefix", h.Cfg.TxtPrefix,
250 )
251
252 return nil
253}
254
255func (h *UploadAssetHandler) findDenylist(bucket storage.Bucket, project *db.Project, logger *slog.Logger) (string, error) {
256 fp, _, err := h.Cfg.Storage.GetObject(bucket, filepath.Join(project.ProjectDir, "_pgs_ignore"))
257 if err != nil {
258 return "", fmt.Errorf("_pgs_ignore not found")
259 }
260 defer func() {
261 _ = fp.Close()
262 }()
263
264 buf := new(strings.Builder)
265 _, err = io.Copy(buf, fp)
266 if err != nil {
267 logger.Error("io copy", "err", err.Error())
268 return "", err
269 }
270
271 str := buf.String()
272 return str, nil
273}
274
275func mtimeToTime(entry *sendutils.FileEntry) time.Time {
276 var mtime time.Time
277 if entry.Mtime > 0 {
278 return time.Unix(entry.Mtime, 0)
279 }
280 return mtime
281}
282
283func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) (string, error) {
284 logger := pssh.GetLogger(s)
285 user := pssh.GetUser(s)
286
287 if user == nil {
288 err := fmt.Errorf("could not get user from ctx")
289 logger.Error("error getting user from ctx", "err", err)
290 return "", err
291 }
292
293 if entry.Mode.IsDir() && strings.Count(entry.Filepath, "/") == 1 {
294 entry.Filepath = strings.TrimPrefix(entry.Filepath, "/")
295 }
296
297 logger = logger.With(
298 "file", entry.Filepath,
299 "size", entry.Size,
300 )
301
302 bucket, err := getBucket(s)
303 if err != nil {
304 logger.Error("could not find bucket in ctx", "err", err.Error())
305 return "", err
306 }
307
308 project := getProject(s)
309 projectName := shared.GetProjectName(entry)
310 logger = logger.With("project", projectName)
311
312 // find, create, or update project if we haven't already done it
313 // we need to also check if the project stored in ctx is the same project
314 // being uploaded since users can keep an ssh connection alive via sftp
315 // and created many projects in a single session
316 if project == nil || project.Name != projectName {
317 project, err = h.Cfg.DB.UpsertProject(user.ID, projectName, projectName)
318 if err != nil {
319 logger.Error("upsert project", "err", err.Error())
320 return "", err
321 }
322 setProject(s, project)
323 }
324
325 if project.Blocked != "" {
326 msg := "project has been blocked and cannot upload files: %s"
327 return "", fmt.Errorf(msg, project.Blocked)
328 }
329
330 info := &storage.ObjectInfo{
331 LastModified: mtimeToTime(entry),
332 }
333 if entry.Mode.IsDir() {
334 _, _, err := h.Cfg.Storage.PutObject(
335 bucket,
336 path.Join(shared.GetAssetFileName(entry), "._pico_keep_dir"),
337 bytes.NewReader([]byte{}),
338 info,
339 )
340 return "", err
341 }
342
343 // calculate the filsize difference between the same file already
344 // stored and the updated file being uploaded
345 assetFilename := shared.GetAssetFileName(entry)
346 obj, info, _ := h.Cfg.Storage.GetObject(bucket, assetFilename)
347 var curFileSize int64
348 if info != nil {
349 curFileSize = info.Size
350 }
351 if obj != nil {
352 defer func() {
353 _ = obj.Close()
354 }()
355 }
356
357 denylist := getDenylist(s)
358 if denylist == nil {
359 dlist, err := h.findDenylist(bucket, project, logger)
360 if err != nil {
361 logger.Info("failed to get denylist, setting default (.*)", "err", err.Error())
362 dlist = ".*"
363 }
364 setDenylist(s, dlist)
365 denylist = &DenyList{Denylist: dlist}
366 }
367
368 data := &FileData{
369 FileEntry: entry,
370 User: user,
371 Bucket: bucket,
372 DenyList: denylist.Denylist,
373 Project: project,
374 }
375
376 valid, err := h.validateAsset(data)
377 if !valid {
378 return "", err
379 }
380
381 featureFlag := getFeatureFlag(s)
382 if featureFlag == nil {
383 return "", fmt.Errorf("pico+ feature flag ctx not set")
384 }
385 // SFTP does not report file size so the more performant way to
386 // check filesize constraints is to try and upload the file to s3
387 // with a specialized reader that raises an error if the filesize limit
388 // has been reached
389 storageMax := featureFlag.Data.StorageMax
390 fileMax := featureFlag.Data.FileMax
391 curStorageSize := getStorageSize(s)
392 remaining := int64(storageMax) - int64(curStorageSize)
393 sizeRemaining := min(remaining+curFileSize, fileMax)
394 if sizeRemaining <= 0 {
395 _, _ = fmt.Fprintln(s.Stderr(), "storage quota reached")
396 _, _ = fmt.Fprintf(s.Stderr(), "\r")
397 _ = s.Exit(1)
398 _ = s.Close()
399 return "", fmt.Errorf("storage quota reached")
400 }
401 logger = logger.With(
402 "storageMax", storageMax,
403 "currentStorageMax", curStorageSize,
404 "fileMax", fileMax,
405 "sizeRemaining", sizeRemaining,
406 )
407
408 specialFileMax := featureFlag.Data.SpecialFileMax
409 if isSpecialFile(entry.Filepath) {
410 sizeRemaining = min(sizeRemaining, specialFileMax)
411 }
412
413 fsize, err := h.writeAsset(
414 s,
415 shared.NewMaxBytesReader(data.Reader, int64(sizeRemaining)),
416 data,
417 )
418 if err != nil {
419 logger.Error("could not write asset", "err", err.Error())
420 cerr := fmt.Errorf(
421 "%s: storage size %.2fmb, storage max %.2fmb, file max %.2fmb, special file max %.4fmb",
422 err,
423 shared.BytesToMB(int(curStorageSize)),
424 shared.BytesToMB(int(storageMax)),
425 shared.BytesToMB(int(fileMax)),
426 shared.BytesToMB(int(specialFileMax)),
427 )
428 return "", cerr
429 }
430
431 deltaFileSize := curFileSize - fsize
432 nextStorageSize := incrementStorageSize(s, deltaFileSize)
433
434 url := h.Cfg.AssetURL(
435 user.Name,
436 projectName,
437 strings.Replace(data.Filepath, "/"+projectName+"/", "", 1),
438 )
439
440 maxSize := int(featureFlag.Data.StorageMax)
441 str := fmt.Sprintf(
442 "%s (space: %.2f/%.2fGB, %.2f%%)",
443 url,
444 shared.BytesToGB(int(nextStorageSize)),
445 shared.BytesToGB(maxSize),
446 (float32(nextStorageSize)/float32(maxSize))*100,
447 )
448
449 surrogate := getSurrogateKey(user.Name, projectName)
450 h.Cfg.CacheClearingQueue <- surrogate
451
452 return str, err
453}
454
455func isSpecialFile(entry string) bool {
456 fname := filepath.Base(entry)
457 return fname == "_headers" || fname == "_redirects" || fname == "_pgs_ignore"
458}
459
460func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) error {
461 logger := pssh.GetLogger(s)
462 user := pssh.GetUser(s)
463
464 if user == nil {
465 err := fmt.Errorf("could not get user from ctx")
466 logger.Error("error getting user from ctx", "err", err)
467 return err
468 }
469
470 if entry.Mode.IsDir() && strings.Count(entry.Filepath, "/") == 1 {
471 entry.Filepath = strings.TrimPrefix(entry.Filepath, "/")
472 }
473
474 assetFilepath := shared.GetAssetFileName(entry)
475
476 logger = logger.With(
477 "file", assetFilepath,
478 )
479
480 bucket, err := getBucket(s)
481 if err != nil {
482 logger.Error("could not find bucket in ctx", "err", err.Error())
483 return err
484 }
485
486 projectName := shared.GetProjectName(entry)
487 logger = logger.With("project", projectName)
488
489 logger.Info("deleting file")
490
491 // Check if this path represents a directory (has a . _pico_keep_dir marker)
492 keepDirPath := filepath.Join(assetFilepath, "._pico_keep_dir")
493 keepDirReader, _, keepDirErr := h.Cfg.Storage.GetObject(bucket, keepDirPath)
494 if keepDirReader != nil {
495 defer func() {
496 _ = keepDirReader.Close()
497 }()
498 }
499
500 if keepDirErr == nil {
501 // This is a directory being deleted. We must delete all nested
502 // . _pico_keep_dir files first, otherwise os.Remove() on the
503 // directory will fail with "directory not empty".
504 nested, err := h.Cfg.Storage.ListObjects(bucket, assetFilepath+"/", true)
505 if err != nil {
506 return err
507 }
508
509 baseDepth := strings.Count(assetFilepath, string(os.PathSeparator))
510 for _, nestedEntry := range nested {
511 if filepath.Base(nestedEntry.Name()) != "._pico_keep_dir" {
512 continue
513 }
514 // Only delete keep_dir files that are direct or nested children
515 nestedDepth := strings.Count(nestedEntry.Name(), string(os.PathSeparator))
516 if nestedDepth <= baseDepth {
517 continue
518 }
519 nestedKeepDirPath := filepath.Join(assetFilepath, nestedEntry.Name())
520 if delErr := h.Cfg.Storage.DeleteObject(bucket, nestedKeepDirPath); delErr != nil {
521 return delErr
522 }
523 }
524
525 // Delete this directory's own . _pico_keep_dir
526 err = h.Cfg.Storage.DeleteObject(bucket, keepDirPath)
527 if err != nil {
528 return err
529 }
530
531 // Delete the directory itself (no-op for S3-style storage, removes the dir for fs storage)
532 _ = h.Cfg.Storage.DeleteObject(bucket, assetFilepath)
533
534 surrogate := getSurrogateKey(user.Name, projectName)
535 h.Cfg.CacheClearingQueue <- surrogate
536 return nil
537 }
538
539 // Regular file deletion: create . _pico_keep_dir if the directory becomes empty
540 pathDir := filepath.Dir(assetFilepath)
541 fileName := filepath.Base(assetFilepath)
542
543 sibs, err := h.Cfg.Storage.ListObjects(bucket, pathDir+"/", false)
544 if err != nil {
545 return err
546 }
547
548 sibs = slices.DeleteFunc(sibs, func(sib fs.FileInfo) bool {
549 return sib.Name() == fileName
550 })
551
552 if len(sibs) == 0 {
553 info := &storage.ObjectInfo{
554 LastModified: mtimeToTime(entry),
555 }
556 _, _, err := h.Cfg.Storage.PutObject(
557 bucket,
558 filepath.Join(pathDir, "._pico_keep_dir"),
559 bytes.NewReader([]byte{}),
560 info,
561 )
562 if err != nil {
563 return err
564 }
565 }
566 err = h.Cfg.Storage.DeleteObject(bucket, assetFilepath)
567
568 surrogate := getSurrogateKey(user.Name, projectName)
569 h.Cfg.CacheClearingQueue <- surrogate
570
571 if err != nil {
572 return err
573 }
574
575 return err
576}
577
578func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) {
579 fname := filepath.Base(data.Filepath)
580
581 projectName := shared.GetProjectName(data.FileEntry)
582 if projectName == "" || projectName == "/" || projectName == "." {
583 return false, fmt.Errorf("ERROR: invalid project name, you must copy files to a non-root folder (e.g. pgs.sh:/project-name)")
584 }
585
586 // special files we use for custom routing
587 if isSpecialFile(fname) {
588 return true, nil
589 }
590
591 fpath := strings.Replace(data.Filepath, "/"+projectName, "", 1)
592 if shouldIgnoreFile(fpath, data.DenyList) {
593 err := fmt.Errorf(
594 "ERROR: (%s) file rejected, https://pico.sh/pgs#-pgs-ignore",
595 data.Filepath,
596 )
597 return false, err
598 }
599
600 return true, nil
601}
602
603func (h *UploadAssetHandler) writeAsset(s *pssh.SSHServerConnSession, reader io.Reader, data *FileData) (int64, error) {
604 assetFilepath := shared.GetAssetFileName(data.FileEntry)
605
606 logger := h.GetLogger(s)
607 logger.Info(
608 "uploading file to bucket",
609 "bucket", data.Bucket.Name,
610 "filename", assetFilepath,
611 )
612
613 info := &storage.ObjectInfo{
614 LastModified: mtimeToTime(data.FileEntry),
615 }
616 _, fsize, err := h.Cfg.Storage.PutObject(
617 data.Bucket,
618 assetFilepath,
619 reader,
620 info,
621 )
622 return fsize, err
623}
624
625// runCacheQueue processes requests to purge the cache for a single site.
626// One message arrives per file that is written/deleted during uploads.
627// Repeated messages for the same site are grouped so that we only flush once
628// per site per 5 seconds.
629func runCacheQueue(cfg *PgsConfig, ctx context.Context) {
630 var pendingFlushes sync.Map
631 tick := time.NewTicker(5 * time.Second)
632 defer tick.Stop()
633 for {
634 select {
635 case <-ctx.Done():
636 return
637 case host := <-cfg.CacheClearingQueue:
638 pendingFlushes.Store(host, host)
639 case <-tick.C:
640 go func() {
641 pendingFlushes.Range(func(key, value any) bool {
642 pendingFlushes.Delete(key)
643 err := purgeCache(cfg, cfg.Pubsub, key.(string))
644 if err != nil {
645 cfg.Logger.Error("failed to clear cache", "err", err.Error())
646 }
647 return true
648 })
649 }()
650 }
651 }
652}