pgs-filehandler.diff

· erock's pastes · raw

expires: 28 Aug, 2024

 1diff --git a/filehandlers/assets/handler.go b/filehandlers/assets/handler.go
 2index 39253a4..c9b9c8c 100644
 3--- a/filehandlers/assets/handler.go
 4+++ b/filehandlers/assets/handler.go
 5@@ -2,7 +2,6 @@ package uploadassets
 6 
 7 import (
 8 	"bytes"
 9-	"encoding/binary"
10 	"fmt"
11 	"io"
12 	"io/fs"
13@@ -91,7 +90,6 @@ func shouldIgnoreFile(fp, ignoreStr string) bool {
14 
15 type FileData struct {
16 	*utils.FileEntry
17-	Text          []byte
18 	User          *db.User
19 	Bucket        sst.Bucket
20 	StorageSize   uint64
21@@ -99,6 +97,7 @@ type FileData struct {
22 	DeltaFileSize int64
23 	Project       *db.Project
24 	DenyList      string
25+	Reader        utils.ReaderAtCloser
26 }
27 
28 type UploadAssetHandler struct {
29@@ -334,14 +333,14 @@ func (h *UploadAssetHandler) Write(s ssh.Session, entry *utils.FileEntry) (strin
30 		return "", err
31 	}
32 
33-	var origText []byte
34-	if b, err := io.ReadAll(entry.Reader); err == nil {
35-		origText = b
36+	buff := bytes.NewBuffer([]byte{})
37+	size, err := io.Copy(buff, entry.Reader)
38+	if err != nil {
39+		return "", err
40 	}
41-	fileSize := binary.Size(origText)
42-	// TODO: hack for now until I figure out how to get correct
43-	// filesize from sftp,scp,rsync
44-	entry.Size = int64(fileSize)
45+	entry.Size = size
46+
47+	reader := utils.NopReaderAtCloser(bytes.NewReader(buff.Bytes()))
48 
49 	storageSize := getStorageSize(s)
50 	featureFlag, err := futil.GetFeatureFlag(s)
51@@ -369,13 +368,13 @@ func (h *UploadAssetHandler) Write(s ssh.Session, entry *utils.FileEntry) (strin
52 	data := &FileData{
53 		FileEntry:     entry,
54 		User:          user,
55-		Text:          origText,
56 		Bucket:        bucket,
57 		StorageSize:   storageSize,
58 		FeatureFlag:   featureFlag,
59 		DeltaFileSize: deltaFileSize,
60 		DenyList:      denylist.Denylist,
61 		Project:       project,
62+		Reader:        reader,
63 	}
64 
65 	valid, err := h.validateAsset(data)
66@@ -499,11 +498,7 @@ func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) {
67 
68 	// ".well-known" is a special case
69 	if strings.Contains(data.Filepath, "/.well-known/") {
70-		if shared.IsTextFile(string(data.Text)) {
71-			return true, nil
72-		} else {
73-			return false, fmt.Errorf("(%s) not a utf-8 text file", data.Filepath)
74-		}
75+		return true, nil
76 	}
77 
78 	// special files we use for custom routing
79@@ -526,8 +521,6 @@ func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) {
80 func (h *UploadAssetHandler) writeAsset(data *FileData) error {
81 	assetFilepath := shared.GetAssetFileName(data.FileEntry)
82 
83-	reader := bytes.NewReader(data.Text)
84-
85 	h.Cfg.Logger.Info(
86 		"uploading file to bucket",
87 		"user", data.User.Name,
88@@ -538,7 +531,7 @@ func (h *UploadAssetHandler) writeAsset(data *FileData) error {
89 	_, err := h.Storage.PutObject(
90 		data.Bucket,
91 		assetFilepath,
92-		utils.NopReaderAtCloser(reader),
93+		data.Reader,
94 		data.FileEntry,
95 	)
96 	if err != nil {