- commit
- 16e2800
- parent
- 5884532
- author
- Eric Bower
- date
- 2024-06-12 19:19:54 +0000 UTC
refactor(pgs): dont copy reader into memory
6 files changed,
+79,
-84
+41,
-76
1@@ -2,7 +2,6 @@ package uploadassets
2
3 import (
4 "bytes"
5- "encoding/binary"
6 "fmt"
7 "io"
8 "io/fs"
9@@ -15,6 +14,7 @@ import (
10 "time"
11
12 "github.com/charmbracelet/ssh"
13+ "github.com/charmbracelet/wish"
14 "github.com/picosh/pico/db"
15 futil "github.com/picosh/pico/filehandlers/util"
16 "github.com/picosh/pico/shared"
17@@ -91,14 +91,10 @@ func shouldIgnoreFile(fp, ignoreStr string) bool {
18
19 type FileData struct {
20 *utils.FileEntry
21- Text []byte
22- User *db.User
23- Bucket sst.Bucket
24- StorageSize uint64
25- FeatureFlag *db.FeatureFlag
26- DeltaFileSize int64
27- Project *db.Project
28- DenyList string
29+ User *db.User
30+ Bucket sst.Bucket
31+ Project *db.Project
32+ DenyList string
33 }
34
35 type UploadAssetHandler struct {
36@@ -325,25 +321,15 @@ func (h *UploadAssetHandler) Write(s ssh.Session, entry *utils.FileEntry) (strin
37 }
38
39 if entry.Mode.IsDir() {
40- _, err := h.Storage.PutObject(
41+ _, _, err := h.Storage.PutObject(
42 bucket,
43 path.Join(shared.GetAssetFileName(entry), "._pico_keep_dir"),
44- utils.NopReaderAtCloser(bytes.NewReader([]byte{})),
45+ bytes.NewReader([]byte{}),
46 entry,
47 )
48 return "", err
49 }
50
51- var origText []byte
52- if b, err := io.ReadAll(entry.Reader); err == nil {
53- origText = b
54- }
55- fileSize := binary.Size(origText)
56- // TODO: hack for now until I figure out how to get correct
57- // filesize from sftp,scp,rsync
58- entry.Size = int64(fileSize)
59-
60- storageSize := getStorageSize(s)
61 featureFlag, err := futil.GetFeatureFlag(s)
62 if err != nil {
63 return "", err
64@@ -353,7 +339,6 @@ func (h *UploadAssetHandler) Write(s ssh.Session, entry *utils.FileEntry) (strin
65 // stored and the updated file being uploaded
66 assetFilename := shared.GetAssetFileName(entry)
67 curFileSize, _ := h.Storage.GetObjectSize(bucket, assetFilename)
68- deltaFileSize := curFileSize - entry.Size
69
70 denylist := getDenylist(s)
71 if denylist == nil {
72@@ -367,15 +352,11 @@ func (h *UploadAssetHandler) Write(s ssh.Session, entry *utils.FileEntry) (strin
73 }
74
75 data := &FileData{
76- FileEntry: entry,
77- User: user,
78- Text: origText,
79- Bucket: bucket,
80- StorageSize: storageSize,
81- FeatureFlag: featureFlag,
82- DeltaFileSize: deltaFileSize,
83- DenyList: denylist.Denylist,
84- Project: project,
85+ FileEntry: entry,
86+ User: user,
87+ Bucket: bucket,
88+ DenyList: denylist.Denylist,
89+ Project: project,
90 }
91
92 valid, err := h.validateAsset(data)
93@@ -383,11 +364,31 @@ func (h *UploadAssetHandler) Write(s ssh.Session, entry *utils.FileEntry) (strin
94 return "", err
95 }
96
97- err = h.writeAsset(data)
98+ // SFTP does not report file size so the more performant way to
99+ // check filesize constraints is to try and upload the file to s3
100+ // with a specialized reader that raises an error if the filesize limit
101+ // has been reached
102+ storageMax := featureFlag.Data.StorageMax
103+ fileMax := featureFlag.Data.FileMax
104+ curStorageSize := getStorageSize(s)
105+ remaining := int64(storageMax) - int64(curStorageSize)
106+ sizeRemaining := min(remaining+curFileSize, fileMax)
107+ if sizeRemaining <= 0 {
108+ msg := "storage quota reached"
109+ wish.Fatalln(s, msg)
110+ return "", fmt.Errorf(msg)
111+ }
112+
113+ fsize, err := h.writeAsset(
114+ shared.NewMaxBytesReader(data.Reader, int64(sizeRemaining)),
115+ data,
116+ )
117 if err != nil {
118- logger.Error(err.Error())
119+ logger.Error("could not write asset", "err", err.Error())
120 return "", err
121 }
122+
123+ deltaFileSize := curFileSize - fsize
124 nextStorageSize := incrementStorageSize(s, deltaFileSize)
125
126 url := h.Cfg.AssetURL(
127@@ -454,10 +455,10 @@ func (h *UploadAssetHandler) Delete(s ssh.Session, entry *utils.FileEntry) error
128 })
129
130 if len(sibs) == 0 {
131- _, err := h.Storage.PutObject(
132+ _, _, err := h.Storage.PutObject(
133 bucket,
134 filepath.Join(pathDir, "._pico_keep_dir"),
135- utils.NopReaderAtCloser(bytes.NewReader([]byte{})),
136+ bytes.NewReader([]byte{}),
137 entry,
138 )
139 if err != nil {
140@@ -469,43 +470,13 @@ func (h *UploadAssetHandler) Delete(s ssh.Session, entry *utils.FileEntry) error
141 }
142
143 func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) {
144- storageMax := data.FeatureFlag.Data.StorageMax
145- var nextStorageSize uint64
146- if data.DeltaFileSize < 0 {
147- nextStorageSize = data.StorageSize - uint64(data.DeltaFileSize)
148- } else {
149- nextStorageSize = data.StorageSize + uint64(data.DeltaFileSize)
150- }
151- if nextStorageSize >= storageMax {
152- return false, fmt.Errorf(
153- "ERROR: user (%s) has exceeded (%d bytes) max (%d bytes)",
154- data.User.Name,
155- data.StorageSize,
156- storageMax,
157- )
158- }
159+ fname := filepath.Base(data.Filepath)
160
161 projectName := shared.GetProjectName(data.FileEntry)
162 if projectName == "" || projectName == "/" || projectName == "." {
163 return false, fmt.Errorf("ERROR: invalid project name, you must copy files to a non-root folder (e.g. pgs.sh:/project-name)")
164 }
165
166- fileSize := data.Size
167- fname := filepath.Base(data.Filepath)
168- fileMax := data.FeatureFlag.Data.FileMax
169- if fileSize > fileMax {
170- return false, fmt.Errorf("ERROR: file (%s) has exceeded maximum file size (%d bytes)", fname, fileMax)
171- }
172-
173- // ".well-known" is a special case
174- if strings.Contains(data.Filepath, "/.well-known/") {
175- if shared.IsTextFile(string(data.Text)) {
176- return true, nil
177- } else {
178- return false, fmt.Errorf("(%s) not a utf-8 text file", data.Filepath)
179- }
180- }
181-
182 // special files we use for custom routing
183 if fname == "_pgs_ignore" || fname == "_redirects" || fname == "_headers" {
184 return true, nil
185@@ -523,11 +494,9 @@ func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) {
186 return true, nil
187 }
188
189-func (h *UploadAssetHandler) writeAsset(data *FileData) error {
190+func (h *UploadAssetHandler) writeAsset(reader io.Reader, data *FileData) (int64, error) {
191 assetFilepath := shared.GetAssetFileName(data.FileEntry)
192
193- reader := bytes.NewReader(data.Text)
194-
195 h.Cfg.Logger.Info(
196 "uploading file to bucket",
197 "user", data.User.Name,
198@@ -535,15 +504,11 @@ func (h *UploadAssetHandler) writeAsset(data *FileData) error {
199 "filename", assetFilepath,
200 )
201
202- _, err := h.Storage.PutObject(
203+ _, fsize, err := h.Storage.PutObject(
204 data.Bucket,
205 assetFilepath,
206- utils.NopReaderAtCloser(reader),
207+ reader,
208 data.FileEntry,
209 )
210- if err != nil {
211- return err
212- }
213-
214- return nil
215+ return fsize, err
216 }
+1,
-1
1@@ -56,7 +56,7 @@ func (h *UploadImgHandler) metaImg(data *PostMetaData) error {
2
3 reader := bytes.NewReader([]byte(data.Text))
4
5- fname, err := h.Storage.PutObject(
6+ fname, _, err := h.Storage.PutObject(
7 bucket,
8 data.Filename,
9 utils.NopReaderAtCloser(reader),
M
go.mod
+2,
-2
1@@ -4,7 +4,7 @@ go 1.22
2
3 // replace github.com/picosh/ptun => ../ptun
4
5-// replace github.com/picosh/send => ../send
6+replace github.com/picosh/send => ../send
7
8 // replace github.com/picosh/pobj => ../pobj
9
10@@ -31,7 +31,7 @@ require (
11 github.com/mmcdole/gofeed v1.3.0
12 github.com/muesli/reflow v0.3.0
13 github.com/neurosnap/go-exif-remove v0.0.0-20221010134343-50d1e3c35577
14- github.com/picosh/pobj v0.0.0-20240529200402-7b5398cf8a9f
15+ github.com/picosh/pobj v0.0.0-20240613155002-f703e1356060
16 github.com/picosh/ptun v0.0.0-20240529133708-fcf1376b935e
17 github.com/picosh/send v0.0.0-20240529200640-3667d1ad154e
18 github.com/sabhiram/go-gitignore v0.0.0-20210923224102-525f6e181f06
M
go.sum
+2,
-4
1@@ -219,12 +219,10 @@ github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N
2 github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
3 github.com/philhofer/fwd v1.1.2 h1:bnDivRJ1EWPjUIRXV5KfORO897HTbpFAQddBdE8t7Gw=
4 github.com/philhofer/fwd v1.1.2/go.mod h1:qkPdfjR2SIEbspLqpe1tO4n5yICnr2DY7mqEx2tUTP0=
5-github.com/picosh/pobj v0.0.0-20240529200402-7b5398cf8a9f h1:9Y0xaTqq/7JiW7iUX4jSgJkN81X2gQucVqBPVXJDux0=
6-github.com/picosh/pobj v0.0.0-20240529200402-7b5398cf8a9f/go.mod h1:CViLWaCp2KP/zJdd7miNPkJb6i0v9HOgZ2wdbQuxCrQ=
7+github.com/picosh/pobj v0.0.0-20240613155002-f703e1356060 h1:lMVTxeTjp6TnuJOd7xvxTbOhk7rTrZhEWBkVwlkXpQQ=
8+github.com/picosh/pobj v0.0.0-20240613155002-f703e1356060/go.mod h1:CViLWaCp2KP/zJdd7miNPkJb6i0v9HOgZ2wdbQuxCrQ=
9 github.com/picosh/ptun v0.0.0-20240529133708-fcf1376b935e h1:Um9aCUg1ysiUaB0nh3400UHlFAnhd8BXBsawqePxxqQ=
10 github.com/picosh/ptun v0.0.0-20240529133708-fcf1376b935e/go.mod h1:WXCrhe0l9VL3ji0pdhvSJD6LLx99rJoAA/+PUQXf0Mo=
11-github.com/picosh/send v0.0.0-20240529200640-3667d1ad154e h1:2NMuieR/7GIjiGYNPQsh6KOJiz2WhzU5ispxQCXmOyU=
12-github.com/picosh/send v0.0.0-20240529200640-3667d1ad154e/go.mod h1:V418obz9YdzjS3/oFzyDFzmPDnLu1nvy3wkLaixiT84=
13 github.com/picosh/senpai v0.0.0-20240503200611-af89e73973b0 h1:pBRIbiCj7K6rGELijb//dYhyCo8A3fvxW5dijrJVtjs=
14 github.com/picosh/senpai v0.0.0-20240503200611-af89e73973b0/go.mod h1:QaBDtybFC5gz7EG/9c3bgzuyW7W5W2rYLFZxWNuWk3Q=
15 github.com/pkg/sftp v1.13.6 h1:JFZT4XbOU7l77xGSpOdW+pwIMqP044IyjXX6FGyEKFo=
+1,
-1
1@@ -5,7 +5,7 @@ import (
2 )
3
4 var maxSize = uint64(25 * shared.MB)
5-var maxAssetSize = int64(5 * shared.MB)
6+var maxAssetSize = int64(1 * shared.MB)
7
8 func NewConfigSite() *shared.ConfigSite {
9 debug := shared.GetEnv("PGS_DEBUG", "0")
1@@ -0,0 +1,32 @@
2+package shared
3+
4+import (
5+ "errors"
6+ "io"
7+)
8+
9+// Throws an error if the reader is bigger than limit.
10+var ErrSizeExceeded = errors.New("stream size exceeded")
11+
12+type MaxBytesReader struct {
13+ io.Reader // reader object
14+ N int64 // max bytes remaining.
15+}
16+
17+func NewMaxBytesReader(r io.Reader, limit int64) *MaxBytesReader {
18+ return &MaxBytesReader{r, limit}
19+}
20+
21+func (b *MaxBytesReader) Read(p []byte) (n int, err error) {
22+ if b.N <= 0 {
23+ return 0, ErrSizeExceeded
24+ }
25+
26+ if int64(len(p)) > b.N {
27+ p = p[0:b.N]
28+ }
29+
30+ n, err = b.Reader.Read(p)
31+ b.N -= int64(n)
32+ return
33+}