1package pgs
  2
  3import (
  4	"context"
  5	"fmt"
  6	"io"
  7	"log/slog"
  8	"os"
  9	"path/filepath"
 10	"strings"
 11	"sync"
 12	"time"
 13
 14	"github.com/picosh/pico/pkg/db"
 15	"github.com/picosh/pico/pkg/pssh"
 16	sendutils "github.com/picosh/pico/pkg/send/utils"
 17	"github.com/picosh/pico/pkg/shared"
 18	"github.com/picosh/pico/pkg/storage"
 19	ignore "github.com/sabhiram/go-gitignore"
 20)
 21
 22type ctxUploadStateKey struct{}
 23
 24// uploadState is what a session's uploads share. It is set on the session
 25// once because every SetValue nests the session context one level deeper,
 26// and SFTP can run several writes at a time.
 27type uploadState struct {
 28	mu          sync.Mutex
 29	bucket      storage.Bucket
 30	featureFlag *db.FeatureFlag
 31	storageSize int64
 32	project     *db.Project
 33	denylist    *ignore.GitIgnore
 34}
 35
 36func getUploadState(s *pssh.SSHServerConnSession) (*uploadState, error) {
 37	state, ok := s.Context().Value(ctxUploadStateKey{}).(*uploadState)
 38	if !ok {
 39		return nil, fmt.Errorf("upload state not set on `ssh.Context()` for connection")
 40	}
 41	return state, nil
 42}
 43
 44func (u *uploadState) getStorageSize() int64 {
 45	u.mu.Lock()
 46	defer u.mu.Unlock()
 47	return u.storageSize
 48}
 49
 50// addStorageSize adds delta bytes to the bucket's size and returns the new
 51// size.
 52func (u *uploadState) addStorageSize(delta int64) int64 {
 53	u.mu.Lock()
 54	defer u.mu.Unlock()
 55	u.storageSize = max(u.storageSize+delta, 0)
 56	return u.storageSize
 57}
 58
 59type FileData struct {
 60	*sendutils.FileEntry
 61	User     *db.User
 62	Bucket   storage.Bucket
 63	Project  *db.Project
 64	DenyList *ignore.GitIgnore
 65}
 66
 67type UploadAssetHandler struct {
 68	Cfg                *Config
 69	CacheClearingQueue chan string
 70}
 71
 72func NewUploadAssetHandler(cfg *Config, ch chan string, ctx context.Context) *UploadAssetHandler {
 73	go runCacheQueue(cfg, ctx)
 74	return &UploadAssetHandler{
 75		Cfg:                cfg,
 76		CacheClearingQueue: ch,
 77	}
 78}
 79
 80func (h *UploadAssetHandler) GetLogger(s *pssh.SSHServerConnSession) *slog.Logger {
 81	return pssh.GetLogger(s)
 82}
 83
 84func (h *UploadAssetHandler) Read(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) (os.FileInfo, sendutils.ReadAndReaderAtCloser, error) {
 85	logger := pssh.GetLogger(s)
 86	user := pssh.GetUser(s)
 87
 88	if user == nil {
 89		err := fmt.Errorf("could not get user from ctx")
 90		logger.Error("error getting user from ctx", "err", err)
 91		return nil, nil, err
 92	}
 93
 94	fileInfo := &sendutils.VirtualFile{
 95		FName:    filepath.Base(entry.Filepath),
 96		FIsDir:   false,
 97		FSize:    entry.Size,
 98		FModTime: time.Unix(entry.Mtime, 0),
 99	}
100
101	bucket, err := h.Cfg.Storage.GetBucket(shared.GetAssetBucketName(user.ID))
102	if err != nil {
103		return nil, nil, err
104	}
105
106	fname := shared.GetAssetFileName(entry)
107	contents, info, err := h.Cfg.Storage.GetObject(bucket, fname)
108	if err != nil {
109		return nil, nil, err
110	}
111
112	fileInfo.FSize = info.Size
113	fileInfo.FModTime = info.LastModified
114
115	return fileInfo, contents, nil
116}
117
118func (h *UploadAssetHandler) List(s *pssh.SSHServerConnSession, fpath string, isDir bool, recursive bool) ([]os.FileInfo, error) {
119	var fileList []os.FileInfo
120
121	logger := pssh.GetLogger(s)
122	user := pssh.GetUser(s)
123
124	if user == nil {
125		err := fmt.Errorf("could not get user from ctx")
126		logger.Error("error getting user from ctx", "err", err)
127		return fileList, err
128	}
129
130	cleanFilename := fpath
131
132	bucketName := shared.GetAssetBucketName(user.ID)
133	bucket, err := h.Cfg.Storage.GetBucket(bucketName)
134	if err != nil {
135		return fileList, err
136	}
137
138	if cleanFilename == "" || cleanFilename == "." {
139		name := cleanFilename
140		if name == "" {
141			name = "/"
142		}
143
144		info := &sendutils.VirtualFile{
145			FName:  name,
146			FIsDir: true,
147		}
148
149		fileList = append(fileList, info)
150	} else {
151		if cleanFilename != "/" && isDir {
152			cleanFilename += "/"
153		}
154
155		foundList, err := h.Cfg.Storage.ListObjects(bucket, cleanFilename, recursive)
156		if err != nil {
157			return fileList, err
158		}
159
160		fileList = append(fileList, foundList...)
161	}
162
163	return fileList, nil
164}
165
166// IsImplicitDir reports top-level project names. A project holds its files
167// whether or not it exists yet.
168func (h *UploadAssetHandler) IsImplicitDir(_ *pssh.SSHServerConnSession, fpath string) bool {
169	return shared.IsValidSubdomain(strings.Trim(fpath, "/"))
170}
171
172func (h *UploadAssetHandler) Validate(s *pssh.SSHServerConnSession) error {
173	logger := pssh.GetLogger(s)
174	user := pssh.GetUser(s)
175
176	if user == nil {
177		err := fmt.Errorf("could not get user from ctx")
178		logger.Error("error getting user from ctx", "err", err)
179		return err
180	}
181
182	ff, err := findFeatureFlag(h.Cfg.DB, h.Cfg, user.ID)
183	if err != nil {
184		return err
185	}
186
187	assetBucket := shared.GetAssetBucketName(user.ID)
188	bucket, err := h.Cfg.Storage.UpsertBucket(assetBucket)
189	if err != nil {
190		return err
191	}
192
193	totalStorageSize, err := h.Cfg.Storage.GetBucketQuota(bucket)
194	if err != nil {
195		return err
196	}
197	s.SetValue(ctxUploadStateKey{}, &uploadState{
198		bucket:      bucket,
199		featureFlag: ff,
200		storageSize: int64(totalStorageSize),
201	})
202
203	logger.Info(
204		"bucket size",
205		"user", user.Name,
206		"bytes", totalStorageSize,
207	)
208
209	logger.Info(
210		"attempting to upload files",
211		"user", user.Name,
212		"txtPrefix", h.Cfg.TxtPrefix,
213	)
214
215	return nil
216}
217
218// sessionProject returns the project being uploaded to and its compiled
219// _pgs_ignore, finding or creating the project when it changes. A session
220// can upload to many projects because SFTP connections are kept alive.
221func (h *UploadAssetHandler) sessionProject(state *uploadState, user *db.User, projectName string, logger *slog.Logger) (*db.Project, *ignore.GitIgnore, error) {
222	state.mu.Lock()
223	defer state.mu.Unlock()
224	if state.project != nil && state.project.Name == projectName {
225		return state.project, state.denylist, nil
226	}
227
228	project, err := h.Cfg.DB.UpsertProject(user.ID, projectName, projectName)
229	if err != nil {
230		logger.Error("upsert project", "err", err.Error())
231		return nil, nil, err
232	}
233
234	dlist, err := h.findDenylist(state.bucket, project, logger)
235	if err != nil {
236		logger.Info("failed to get denylist, setting default (.*)", "err", err.Error())
237		dlist = ".*"
238	}
239
240	state.project = project
241	state.denylist = ignore.CompileIgnoreLines(strings.Split(dlist, "\n")...)
242	return state.project, state.denylist, nil
243}
244
245func (h *UploadAssetHandler) findDenylist(bucket storage.Bucket, project *db.Project, logger *slog.Logger) (string, error) {
246	fp, _, err := h.Cfg.Storage.GetObject(bucket, filepath.Join(project.ProjectDir, "_pgs_ignore"))
247	if err != nil {
248		return "", fmt.Errorf("_pgs_ignore not found")
249	}
250	defer func() {
251		_ = fp.Close()
252	}()
253
254	buf := new(strings.Builder)
255	_, err = io.Copy(buf, fp)
256	if err != nil {
257		logger.Error("io copy", "err", err.Error())
258		return "", err
259	}
260
261	str := buf.String()
262	return str, nil
263}
264
265func mtimeToTime(entry *sendutils.FileEntry) time.Time {
266	var mtime time.Time
267	if entry.Mtime > 0 {
268		return time.Unix(entry.Mtime, 0)
269	}
270	return mtime
271}
272
273func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) (string, error) {
274	return h.write(s, entry, false)
275}
276
277// ValidateWrite runs the checks Write makes before reading the file.
278func (h *UploadAssetHandler) ValidateWrite(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) error {
279	if entry.Mode.IsDir() {
280		return nil
281	}
282	_, err := h.write(s, entry, true)
283	return err
284}
285
286func (h *UploadAssetHandler) write(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry, validateOnly bool) (string, error) {
287	logger := pssh.GetLogger(s)
288	user := pssh.GetUser(s)
289
290	if user == nil {
291		err := fmt.Errorf("could not get user from ctx")
292		logger.Error("error getting user from ctx", "err", err)
293		return "", err
294	}
295
296	if entry.Mode.IsDir() && strings.Count(entry.Filepath, "/") == 1 {
297		entry.Filepath = strings.TrimPrefix(entry.Filepath, "/")
298	}
299
300	logger = logger.With(
301		"file", entry.Filepath,
302		"size", entry.Size,
303	)
304
305	state, err := getUploadState(s)
306	if err != nil {
307		logger.Error("could not find upload state in ctx", "err", err.Error())
308		return "", err
309	}
310	bucket := state.bucket
311
312	projectName := shared.GetProjectName(entry)
313	logger = logger.With("project", projectName)
314	if projectName == "" || projectName == "/" || projectName == "." {
315		return "", fmt.Errorf("ERROR: invalid project name, you must copy files to a non-root folder (e.g. pgs.sh:/project-name)")
316	}
317
318	project, denylist, err := h.sessionProject(state, user, projectName, logger)
319	if err != nil {
320		return "", err
321	}
322
323	if project.Blocked != "" {
324		msg := "project has been blocked and cannot upload files: %s"
325		return "", fmt.Errorf(msg, project.Blocked)
326	}
327
328	if entry.Mode.IsDir() {
329		return "", h.Cfg.Storage.PutDir(bucket, shared.GetAssetFileName(entry))
330	}
331
332	// calculate the filsize difference between the same file already
333	// stored and the updated file being uploaded
334	assetFilename := shared.GetAssetFileName(entry)
335	var curFileSize int64
336	if info, err := h.Cfg.Storage.StatObject(bucket, assetFilename); err == nil {
337		curFileSize = info.Size
338	}
339
340	data := &FileData{
341		FileEntry: entry,
342		User:      user,
343		Bucket:    bucket,
344		DenyList:  denylist,
345		Project:   project,
346	}
347
348	valid, err := h.validateAsset(data)
349	if !valid {
350		return "", err
351	}
352
353	featureFlag := state.featureFlag
354	// SFTP does not report file size so the more performant way to
355	//   check filesize constraints is to try and upload the file to s3
356	//	 with a specialized reader that raises an error if the filesize limit
357	//	 has been reached
358	storageMax := featureFlag.Data.StorageMax
359	fileMax := featureFlag.Data.FileMax
360	curStorageSize := state.getStorageSize()
361	remaining := int64(storageMax) - curStorageSize
362	sizeRemaining := min(remaining+curFileSize, fileMax)
363	if sizeRemaining <= 0 {
364		return "", fmt.Errorf("storage quota reached")
365	}
366	if validateOnly {
367		return "", nil
368	}
369	logger = logger.With(
370		"storageMax", storageMax,
371		"currentStorageMax", curStorageSize,
372		"fileMax", fileMax,
373		"sizeRemaining", sizeRemaining,
374	)
375
376	specialFileMax := featureFlag.Data.SpecialFileMax
377	if isSpecialFile(entry.Filepath) {
378		sizeRemaining = min(sizeRemaining, specialFileMax)
379	}
380
381	fsize, err := h.writeAsset(
382		s,
383		shared.NewMaxBytesReader(data.Reader, int64(sizeRemaining)),
384		data,
385	)
386	if err != nil {
387		logger.Error("could not write asset", "err", err.Error())
388		cerr := fmt.Errorf(
389			"%s: storage size %.2fmb, storage max %.2fmb, file max %.2fmb, special file max %.4fmb",
390			err,
391			shared.BytesToMB(int(curStorageSize)),
392			shared.BytesToMB(int(storageMax)),
393			shared.BytesToMB(int(fileMax)),
394			shared.BytesToMB(int(specialFileMax)),
395		)
396		return "", cerr
397	}
398
399	nextStorageSize := state.addStorageSize(fsize - curFileSize)
400
401	url := h.Cfg.AssetURL(
402		user.Name,
403		projectName,
404		strings.Replace(data.Filepath, "/"+projectName+"/", "", 1),
405	)
406
407	str := url + shared.StorageUsage(int(nextStorageSize), int(featureFlag.Data.StorageMax))
408
409	surrogate := getSurrogateKey(user.Name, projectName)
410	h.Cfg.CacheClearingQueue <- surrogate
411
412	return str, err
413}
414
415func isSpecialFile(entry string) bool {
416	fname := filepath.Base(entry)
417	return fname == "_headers" || fname == "_redirects" || fname == "_pgs_ignore"
418}
419
420func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) error {
421	logger := pssh.GetLogger(s)
422	user := pssh.GetUser(s)
423
424	if user == nil {
425		err := fmt.Errorf("could not get user from ctx")
426		logger.Error("error getting user from ctx", "err", err)
427		return err
428	}
429
430	if entry.Mode.IsDir() && strings.Count(entry.Filepath, "/") == 1 {
431		entry.Filepath = strings.TrimPrefix(entry.Filepath, "/")
432	}
433
434	assetFilepath := shared.GetAssetFileName(entry)
435
436	logger = logger.With(
437		"file", assetFilepath,
438	)
439
440	state, err := getUploadState(s)
441	if err != nil {
442		logger.Error("could not find upload state in ctx", "err", err.Error())
443		return err
444	}
445	bucket := state.bucket
446
447	projectName := shared.GetProjectName(entry)
448	logger = logger.With("project", projectName)
449
450	logger.Info("deleting file")
451
452	info, statErr := h.Cfg.Storage.StatObject(bucket, assetFilepath)
453	err = h.Cfg.Storage.DeleteObject(bucket, assetFilepath)
454	if err == nil && statErr == nil {
455		state.addStorageSize(-info.Size)
456	}
457
458	surrogate := getSurrogateKey(user.Name, projectName)
459	h.Cfg.CacheClearingQueue <- surrogate
460
461	if err != nil {
462		return err
463	}
464
465	return err
466}
467
468func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) {
469	fname := filepath.Base(data.Filepath)
470	projectName := shared.GetProjectName(data.FileEntry)
471
472	// special files we use for custom routing
473	if isSpecialFile(fname) {
474		return true, nil
475	}
476
477	fpath := strings.Replace(data.Filepath, "/"+projectName, "", 1)
478	if data.DenyList.MatchesPath(fpath) {
479		err := fmt.Errorf(
480			"%w: (%s), https://pico.sh/pgs#-pgs-ignore",
481			sendutils.ErrFileIgnored,
482			data.Filepath,
483		)
484		return false, err
485	}
486
487	return true, nil
488}
489
490func (h *UploadAssetHandler) writeAsset(s *pssh.SSHServerConnSession, reader io.Reader, data *FileData) (int64, error) {
491	assetFilepath := shared.GetAssetFileName(data.FileEntry)
492
493	logger := h.GetLogger(s)
494	logger.Info(
495		"uploading file to bucket",
496		"bucket", data.Bucket.Name,
497		"filename", assetFilepath,
498	)
499
500	info := &storage.ObjectInfo{
501		LastModified: mtimeToTime(data.FileEntry),
502	}
503	_, fsize, err := h.Cfg.Storage.PutObject(
504		data.Bucket,
505		assetFilepath,
506		reader,
507		info,
508	)
509	return fsize, err
510}
511
512// runCacheQueue processes requests to purge the cache for a single site.
513// One message arrives per file that is written/deleted during uploads.
514// Repeated messages for the same site are grouped so that we only flush once
515// per site per 5 seconds.
516func runCacheQueue(cfg *Config, ctx context.Context) {
517	var pendingFlushes sync.Map
518	tick := time.NewTicker(5 * time.Second)
519	defer tick.Stop()
520	for {
521		select {
522		case <-ctx.Done():
523			return
524		case host := <-cfg.CacheClearingQueue:
525			pendingFlushes.Store(host, host)
526		case <-tick.C:
527			go func() {
528				pendingFlushes.Range(func(key, value any) bool {
529					pendingFlushes.Delete(key)
530					err := purgeCache(cfg, cfg.Pubsub, key.(string))
531					if err != nil {
532						cfg.Logger.Error("failed to clear cache", "err", err.Error())
533					}
534					return true
535				})
536			}()
537		}
538	}
539}