fix: put as task from web

AlistGo · Jul 1, 2022 · 4340a48 · 4340a48
1 parent 4d0ae6b
commit 4340a48
Show file tree

Hide file tree

Showing 8 changed files with 45 additions and 7 deletions.
diff --git a/internal/aria2/monitor.go b/internal/aria2/monitor.go
@@ -147,7 +147,7 @@ func (m *Monitor) Complete() error {
 				if err != nil {
 					return errors.Wrapf(err, "failed to open file %s", file.Path)
 				}
-				stream := model.FileStream{
+				stream := &model.FileStream{
 					Obj: model.Object{
 						Name:     path.Base(file.Path),
 						Size:     size,

diff --git a/internal/fs/put.go b/internal/fs/put.go
@@ -7,6 +7,7 @@ import (
 	"github.com/alist-org/alist/v3/internal/model"
 	"github.com/alist-org/alist/v3/internal/operations"
 	"github.com/alist-org/alist/v3/pkg/task"
+	"github.com/alist-org/alist/v3/pkg/utils"
 	"github.com/pkg/errors"
 	"sync/atomic"
 )
@@ -24,6 +25,13 @@ func putAsTask(dstDirPath string, file model.FileStreamer) error {
 	if err != nil {
 		return errors.WithMessage(err, "failed get account")
 	}
+	if file.NeedStore() {
+		tempFile, err := utils.CreateTempFile(file)
+		if err != nil {
+			return errors.Wrapf(err, "failed to create temp file")
+		}
+		file.SetReadCloser(tempFile)
+	}
 	UploadTaskManager.Submit(task.WithCancelCtx(&task.Task[uint64]{
 		Name: fmt.Sprintf("upload %s to [%s](%s)", file.GetName(), account.GetAccount().VirtualPath, dstDirActualPath),
 		Func: func(task *task.Task[uint64]) error {

diff --git a/internal/fs/util.go b/internal/fs/util.go
@@ -65,7 +65,7 @@ func getFileStreamFromLink(file model.Obj, link *model.Link) (model.FileStreamer
 	if mimetype == "" {
 		mimetype = "application/octet-stream"
 	}
-	stream := model.FileStream{
+	stream := &model.FileStream{
 		Obj:        file,
 		ReadCloser: rc,
 		Mimetype:   mimetype,

diff --git a/internal/model/obj.go b/internal/model/obj.go
@@ -19,6 +19,9 @@ type FileStreamer interface {
 	io.ReadCloser
 	Obj
 	GetMimetype() string
+	SetReadCloser(io.ReadCloser)
+	NeedStore() bool
+	GetReadCloser() io.ReadCloser
 }
 
 type URL interface {

diff --git a/internal/model/stream.go b/internal/model/stream.go
@@ -7,9 +7,22 @@ import (
 type FileStream struct {
 	Obj
 	io.ReadCloser
-	Mimetype string
+	Mimetype     string
+	WebPutAsTask bool
 }
 
 func (f FileStream) GetMimetype() string {
 	return f.Mimetype
 }
+
+func (f FileStream) NeedStore() bool {
+	return f.WebPutAsTask
+}
+
+func (f *FileStream) GetReadCloser() io.ReadCloser {
+	return f.ReadCloser
+}
+
+func (f *FileStream) SetReadCloser(rc io.ReadCloser) {
+	f.ReadCloser = rc
+}
diff --git a/internal/operations/fs.go b/internal/operations/fs.go
@@ -5,6 +5,7 @@ import (
 	"github.com/alist-org/alist/v3/internal/conf"
 	"github.com/alist-org/alist/v3/internal/errs"
 	log "github.com/sirupsen/logrus"
+	"os"
 	stdpath "path"
 	"strings"
 	"time"
@@ -222,6 +223,14 @@ func Remove(ctx context.Context, account driver.Driver, path string) error {
 }
 
 func Put(ctx context.Context, account driver.Driver, dstDirPath string, file model.FileStreamer, up driver.UpdateProgress) error {
+	defer func() {
+		if f, ok := file.GetReadCloser().(*os.File); ok {
+			err := os.RemoveAll(f.Name())
+			if err != nil {
+				log.Errorf("failed to remove file [%s]", f.Name())
+			}
+		}
+	}()
 	defer func() {
 		if err := file.Close(); err != nil {
 			log.Errorf("failed to close file streamer, %v", err)
@@ -241,6 +250,7 @@ func Put(ctx context.Context, account driver.Driver, dstDirPath string, file mod
 		up = func(p int) {}
 	}
 	err = account.Put(ctx, parentDir, file, up)
+	log.Debugf("put file [%s] done", file.GetName())
 	if err == nil {
 		// clear cache
 		key := stdpath.Join(account.GetAccount().VirtualPath, dstDirPath)

diff --git a/server/controllers/fsmanage.go b/server/controllers/fsmanage.go
@@ -9,6 +9,7 @@ import (
 	"github.com/alist-org/alist/v3/internal/sign"
 	"github.com/alist-org/alist/v3/server/common"
 	"github.com/gin-gonic/gin"
+	log "github.com/sirupsen/logrus"
 	stdpath "path"
 	"strconv"
 	"time"
@@ -205,14 +206,17 @@ func FsPut(c *gin.Context) {
 		common.ErrorResp(c, err, 400)
 		return
 	}
-	if err := fs.PutAsTask(dir, model.FileStream{
+	log.Debugf("c.Request.Close, %v", c.Request.Close)
+	c.Request.Close = false
+	if err := fs.PutAsTask(dir, &model.FileStream{
 		Obj: model.Object{
 			Name:     name,
 			Size:     size,
 			Modified: time.Now(),
 		},
-		ReadCloser: c.Request.Body,
-		Mimetype:   c.GetHeader("Content-Type"),
+		ReadCloser:   c.Request.Body,
+		Mimetype:     c.GetHeader("Content-Type"),
+		WebPutAsTask: true,
 	}); err != nil {
 		common.ErrorResp(c, err, 500)
 		return

diff --git a/server/webdav/webdav.go b/server/webdav/webdav.go
@@ -291,7 +291,7 @@ func (h *Handler) handlePut(w http.ResponseWriter, r *http.Request) (status int,
 		Size:     r.ContentLength,
 		Modified: time.Now(),
 	}
-	stream := model.FileStream{
+	stream := &model.FileStream{
 		Obj:        obj,
 		ReadCloser: r.Body,
 		Mimetype:   r.Header.Get("Content-Type"),