@@ -3,6 +3,7 @@ package ftp
3
3
import (
4
4
"bytes"
5
5
"context"
6
+ "fmt"
6
7
"io"
7
8
"net/http"
8
9
"os"
@@ -62,10 +63,10 @@ func (f *FileUploadProxy) Read(p []byte) (n int, err error) {
62
63
func (f * FileUploadProxy ) Write (p []byte ) (n int , err error ) {
63
64
n , err = f .buffer .Write (p )
64
65
if err != nil {
65
- return
66
+ return n , err
66
67
}
67
68
err = stream .ClientUploadLimit .WaitN (f .ctx , n )
68
- return
69
+ return n , err
69
70
}
70
71
71
72
func (f * FileUploadProxy ) Seek (offset int64 , whence int ) (int64 , error ) {
@@ -89,6 +90,25 @@ func (f *FileUploadProxy) Close() error {
89
90
if _ , err := f .buffer .Seek (0 , io .SeekStart ); err != nil {
90
91
return err
91
92
}
93
+ user := f .ctx .Value (conf .UserKey ).(* model.User )
94
+ sf , borrow , err := MakeStage (f .ctx , f .buffer , size , f .path , func (target string ) {
95
+ ctx := context .WithValue (context .Background (), conf .UserKey , user )
96
+ dstDir , dstBase := stdpath .Split (target )
97
+ if dir == dstDir {
98
+ _ = fs .Rename (ctx , f .path , dstBase )
99
+ } else {
100
+ if name != dstBase {
101
+ e := fs .Rename (ctx , f .path , dstBase , true )
102
+ if e != nil {
103
+ return
104
+ }
105
+ }
106
+ _ , _ = fs .Move (ctx , stdpath .Join (dir , dstBase ), dstDir )
107
+ }
108
+ })
109
+ if err != nil {
110
+ return fmt .Errorf ("failed make stage for [%s]: %+v" , f .path , err )
111
+ }
92
112
if f .trunc {
93
113
_ = fs .Remove (f .ctx , f .path )
94
114
}
@@ -100,10 +120,18 @@ func (f *FileUploadProxy) Close() error {
100
120
},
101
121
Mimetype : contentType ,
102
122
WebPutAsTask : true ,
123
+ Reader : f .buffer ,
103
124
}
104
- s .SetTmpFile (f .buffer )
105
- _ , err = fs .PutAsTask (f .ctx , dir , s )
106
- return err
125
+ s .Add (borrow )
126
+ task , err := fs .PutAsTask (f .ctx , dir , s )
127
+ if err != nil {
128
+ _ = s .Close ()
129
+ return err
130
+ }
131
+ sf .SetRemoveCallback (func () {
132
+ fs .UploadTaskManager .Cancel (task .GetID ())
133
+ })
134
+ return nil
107
135
}
108
136
109
137
type FileUploadWithLengthProxy struct {
@@ -182,10 +210,10 @@ func (f *FileUploadWithLengthProxy) write(p []byte) (n int, err error) {
182
210
func (f * FileUploadWithLengthProxy ) Write (p []byte ) (n int , err error ) {
183
211
n , err = f .write (p )
184
212
if err != nil {
185
- return
213
+ return n , err
186
214
}
187
215
err = stream .ClientUploadLimit .WaitN (f .ctx , n )
188
- return
216
+ return n , err
189
217
}
190
218
191
219
func (f * FileUploadWithLengthProxy ) Seek (offset int64 , whence int ) (int64 , error ) {
@@ -214,6 +242,6 @@ func (f *FileUploadWithLengthProxy) Close() error {
214
242
WebPutAsTask : false ,
215
243
Reader : bytes .NewReader (data ),
216
244
}
217
- return fs .PutDirectly (f .ctx , dir , s , true )
245
+ return fs .PutDirectly (f .ctx , dir , s )
218
246
}
219
247
}
0 commit comments