transfer_test.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459
  1. // Copyright (C) 2019-2023 Nicola Murino
  2. //
  3. // This program is free software: you can redistribute it and/or modify
  4. // it under the terms of the GNU Affero General Public License as published
  5. // by the Free Software Foundation, version 3.
  6. //
  7. // This program is distributed in the hope that it will be useful,
  8. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  9. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  10. // GNU Affero General Public License for more details.
  11. //
  12. // You should have received a copy of the GNU Affero General Public License
  13. // along with this program. If not, see <https://www.gnu.org/licenses/>.
  14. package common
  15. import (
  16. "errors"
  17. "os"
  18. "path/filepath"
  19. "testing"
  20. "time"
  21. "github.com/pkg/sftp"
  22. "github.com/sftpgo/sdk"
  23. "github.com/stretchr/testify/assert"
  24. "github.com/stretchr/testify/require"
  25. "github.com/drakkan/sftpgo/v2/internal/dataprovider"
  26. "github.com/drakkan/sftpgo/v2/internal/kms"
  27. "github.com/drakkan/sftpgo/v2/internal/vfs"
  28. )
  29. func TestTransferUpdateQuota(t *testing.T) {
  30. conn := NewBaseConnection("", ProtocolSFTP, "", "", dataprovider.User{})
  31. transfer := BaseTransfer{
  32. Connection: conn,
  33. transferType: TransferUpload,
  34. Fs: vfs.NewOsFs("", os.TempDir(), "", nil),
  35. }
  36. transfer.BytesReceived.Store(123)
  37. errFake := errors.New("fake error")
  38. transfer.TransferError(errFake)
  39. err := transfer.Close()
  40. if assert.Error(t, err) {
  41. assert.EqualError(t, err, errFake.Error())
  42. }
  43. mappedPath := filepath.Join(os.TempDir(), "vdir")
  44. vdirPath := "/vdir"
  45. conn.User.VirtualFolders = append(conn.User.VirtualFolders, vfs.VirtualFolder{
  46. BaseVirtualFolder: vfs.BaseVirtualFolder{
  47. MappedPath: mappedPath,
  48. },
  49. VirtualPath: vdirPath,
  50. QuotaFiles: -1,
  51. QuotaSize: -1,
  52. })
  53. transfer.ErrTransfer = nil
  54. transfer.BytesReceived.Store(1)
  55. transfer.requestPath = "/vdir/file"
  56. assert.True(t, transfer.updateQuota(1, 0))
  57. err = transfer.Close()
  58. assert.NoError(t, err)
  59. transfer.ErrTransfer = errFake
  60. transfer.Fs = newMockOsFs(true, "", "", "S3Fs fake", nil)
  61. assert.False(t, transfer.updateQuota(1, 0))
  62. }
  63. func TestTransferThrottling(t *testing.T) {
  64. u := dataprovider.User{
  65. BaseUser: sdk.BaseUser{
  66. Username: "test",
  67. UploadBandwidth: 50,
  68. DownloadBandwidth: 40,
  69. },
  70. }
  71. fs := vfs.NewOsFs("", os.TempDir(), "", nil)
  72. testFileSize := int64(131072)
  73. wantedUploadElapsed := 1000 * (testFileSize / 1024) / u.UploadBandwidth
  74. wantedDownloadElapsed := 1000 * (testFileSize / 1024) / u.DownloadBandwidth
  75. // some tolerance
  76. wantedUploadElapsed -= wantedDownloadElapsed / 10
  77. wantedDownloadElapsed -= wantedDownloadElapsed / 10
  78. conn := NewBaseConnection("id", ProtocolSCP, "", "", u)
  79. transfer := NewBaseTransfer(nil, conn, nil, "", "", "", TransferUpload, 0, 0, 0, 0, true, fs, dataprovider.TransferQuota{})
  80. transfer.BytesReceived.Store(testFileSize)
  81. transfer.Connection.UpdateLastActivity()
  82. startTime := transfer.Connection.GetLastActivity()
  83. transfer.HandleThrottle()
  84. elapsed := time.Since(startTime).Nanoseconds() / 1000000
  85. assert.GreaterOrEqual(t, elapsed, wantedUploadElapsed, "upload bandwidth throttling not respected")
  86. err := transfer.Close()
  87. assert.NoError(t, err)
  88. transfer = NewBaseTransfer(nil, conn, nil, "", "", "", TransferDownload, 0, 0, 0, 0, true, fs, dataprovider.TransferQuota{})
  89. transfer.BytesSent.Store(testFileSize)
  90. transfer.Connection.UpdateLastActivity()
  91. startTime = transfer.Connection.GetLastActivity()
  92. transfer.HandleThrottle()
  93. elapsed = time.Since(startTime).Nanoseconds() / 1000000
  94. assert.GreaterOrEqual(t, elapsed, wantedDownloadElapsed, "download bandwidth throttling not respected")
  95. err = transfer.Close()
  96. assert.NoError(t, err)
  97. }
  98. func TestRealPath(t *testing.T) {
  99. testFile := filepath.Join(os.TempDir(), "afile.txt")
  100. fs := vfs.NewOsFs("123", os.TempDir(), "", nil)
  101. u := dataprovider.User{
  102. BaseUser: sdk.BaseUser{
  103. Username: "user",
  104. HomeDir: os.TempDir(),
  105. },
  106. }
  107. u.Permissions = make(map[string][]string)
  108. u.Permissions["/"] = []string{dataprovider.PermAny}
  109. file, err := os.Create(testFile)
  110. require.NoError(t, err)
  111. conn := NewBaseConnection(fs.ConnectionID(), ProtocolSFTP, "", "", u)
  112. transfer := NewBaseTransfer(file, conn, nil, testFile, testFile, "/transfer_test_file",
  113. TransferUpload, 0, 0, 0, 0, true, fs, dataprovider.TransferQuota{})
  114. rPath := transfer.GetRealFsPath(testFile)
  115. assert.Equal(t, testFile, rPath)
  116. rPath = conn.getRealFsPath(testFile)
  117. assert.Equal(t, testFile, rPath)
  118. err = transfer.Close()
  119. assert.NoError(t, err)
  120. err = file.Close()
  121. assert.NoError(t, err)
  122. transfer.File = nil
  123. rPath = transfer.GetRealFsPath(testFile)
  124. assert.Equal(t, testFile, rPath)
  125. rPath = transfer.GetRealFsPath("")
  126. assert.Empty(t, rPath)
  127. err = os.Remove(testFile)
  128. assert.NoError(t, err)
  129. assert.Len(t, conn.GetTransfers(), 0)
  130. }
  131. func TestTruncate(t *testing.T) {
  132. testFile := filepath.Join(os.TempDir(), "transfer_test_file")
  133. fs := vfs.NewOsFs("123", os.TempDir(), "", nil)
  134. u := dataprovider.User{
  135. BaseUser: sdk.BaseUser{
  136. Username: "user",
  137. HomeDir: os.TempDir(),
  138. },
  139. }
  140. u.Permissions = make(map[string][]string)
  141. u.Permissions["/"] = []string{dataprovider.PermAny}
  142. file, err := os.Create(testFile)
  143. if !assert.NoError(t, err) {
  144. assert.FailNow(t, "unable to open test file")
  145. }
  146. _, err = file.Write([]byte("hello"))
  147. assert.NoError(t, err)
  148. conn := NewBaseConnection(fs.ConnectionID(), ProtocolSFTP, "", "", u)
  149. transfer := NewBaseTransfer(file, conn, nil, testFile, testFile, "/transfer_test_file", TransferUpload, 0, 5,
  150. 100, 0, false, fs, dataprovider.TransferQuota{})
  151. err = conn.SetStat("/transfer_test_file", &StatAttributes{
  152. Size: 2,
  153. Flags: StatAttrSize,
  154. })
  155. assert.NoError(t, err)
  156. assert.Equal(t, int64(103), transfer.MaxWriteSize)
  157. err = transfer.Close()
  158. assert.NoError(t, err)
  159. err = file.Close()
  160. assert.NoError(t, err)
  161. fi, err := os.Stat(testFile)
  162. if assert.NoError(t, err) {
  163. assert.Equal(t, int64(2), fi.Size())
  164. }
  165. transfer = NewBaseTransfer(file, conn, nil, testFile, testFile, "/transfer_test_file", TransferUpload, 0, 0,
  166. 100, 0, true, fs, dataprovider.TransferQuota{})
  167. // file.Stat will fail on a closed file
  168. err = conn.SetStat("/transfer_test_file", &StatAttributes{
  169. Size: 2,
  170. Flags: StatAttrSize,
  171. })
  172. assert.Error(t, err)
  173. err = transfer.Close()
  174. assert.NoError(t, err)
  175. transfer = NewBaseTransfer(nil, conn, nil, testFile, testFile, "", TransferUpload, 0, 0, 0, 0, true,
  176. fs, dataprovider.TransferQuota{})
  177. _, err = transfer.Truncate("mismatch", 0)
  178. assert.EqualError(t, err, errTransferMismatch.Error())
  179. _, err = transfer.Truncate(testFile, 0)
  180. assert.NoError(t, err)
  181. _, err = transfer.Truncate(testFile, 1)
  182. assert.EqualError(t, err, vfs.ErrVfsUnsupported.Error())
  183. err = transfer.Close()
  184. assert.NoError(t, err)
  185. err = os.Remove(testFile)
  186. assert.NoError(t, err)
  187. assert.Len(t, conn.GetTransfers(), 0)
  188. }
  189. func TestTransferErrors(t *testing.T) {
  190. isCancelled := false
  191. cancelFn := func() {
  192. isCancelled = true
  193. }
  194. testFile := filepath.Join(os.TempDir(), "transfer_test_file")
  195. fs := vfs.NewOsFs("id", os.TempDir(), "", nil)
  196. u := dataprovider.User{
  197. BaseUser: sdk.BaseUser{
  198. Username: "test",
  199. HomeDir: os.TempDir(),
  200. },
  201. }
  202. err := os.WriteFile(testFile, []byte("test data"), os.ModePerm)
  203. assert.NoError(t, err)
  204. file, err := os.Open(testFile)
  205. if !assert.NoError(t, err) {
  206. assert.FailNow(t, "unable to open test file")
  207. }
  208. conn := NewBaseConnection("id", ProtocolSFTP, "", "", u)
  209. transfer := NewBaseTransfer(file, conn, nil, testFile, testFile, "/transfer_test_file", TransferUpload,
  210. 0, 0, 0, 0, true, fs, dataprovider.TransferQuota{})
  211. err = transfer.ConvertError(os.ErrNotExist)
  212. assert.ErrorIs(t, err, sftp.ErrSSHFxNoSuchFile)
  213. err = transfer.ConvertError(os.ErrPermission)
  214. assert.ErrorIs(t, err, sftp.ErrSSHFxPermissionDenied)
  215. assert.Nil(t, transfer.cancelFn)
  216. assert.Equal(t, testFile, transfer.GetFsPath())
  217. transfer.SetCancelFn(cancelFn)
  218. errFake := errors.New("err fake")
  219. transfer.BytesReceived.Store(9)
  220. transfer.TransferError(ErrQuotaExceeded)
  221. assert.True(t, isCancelled)
  222. transfer.TransferError(errFake)
  223. assert.Error(t, transfer.ErrTransfer, ErrQuotaExceeded.Error())
  224. // the file is closed from the embedding struct before to call close
  225. err = file.Close()
  226. assert.NoError(t, err)
  227. err = transfer.Close()
  228. if assert.Error(t, err) {
  229. assert.Error(t, err, ErrQuotaExceeded.Error())
  230. }
  231. assert.NoFileExists(t, testFile)
  232. err = os.WriteFile(testFile, []byte("test data"), os.ModePerm)
  233. assert.NoError(t, err)
  234. file, err = os.Open(testFile)
  235. if !assert.NoError(t, err) {
  236. assert.FailNow(t, "unable to open test file")
  237. }
  238. fsPath := filepath.Join(os.TempDir(), "test_file")
  239. transfer = NewBaseTransfer(file, conn, nil, fsPath, file.Name(), "/test_file", TransferUpload, 0, 0, 0, 0, true,
  240. fs, dataprovider.TransferQuota{})
  241. transfer.BytesReceived.Store(9)
  242. transfer.TransferError(errFake)
  243. assert.Error(t, transfer.ErrTransfer, errFake.Error())
  244. // the file is closed from the embedding struct before to call close
  245. err = file.Close()
  246. assert.NoError(t, err)
  247. err = transfer.Close()
  248. if assert.Error(t, err) {
  249. assert.Error(t, err, errFake.Error())
  250. }
  251. assert.NoFileExists(t, testFile)
  252. err = os.WriteFile(testFile, []byte("test data"), os.ModePerm)
  253. assert.NoError(t, err)
  254. file, err = os.Open(testFile)
  255. if !assert.NoError(t, err) {
  256. assert.FailNow(t, "unable to open test file")
  257. }
  258. transfer = NewBaseTransfer(file, conn, nil, fsPath, file.Name(), "/test_file", TransferUpload, 0, 0, 0, 0, true,
  259. fs, dataprovider.TransferQuota{})
  260. transfer.BytesReceived.Store(9)
  261. // the file is closed from the embedding struct before to call close
  262. err = file.Close()
  263. assert.NoError(t, err)
  264. err = transfer.Close()
  265. assert.NoError(t, err)
  266. assert.NoFileExists(t, testFile)
  267. assert.FileExists(t, fsPath)
  268. err = os.Remove(fsPath)
  269. assert.NoError(t, err)
  270. assert.Len(t, conn.GetTransfers(), 0)
  271. }
  272. func TestRemovePartialCryptoFile(t *testing.T) {
  273. testFile := filepath.Join(os.TempDir(), "transfer_test_file")
  274. fs, err := vfs.NewCryptFs("id", os.TempDir(), "", vfs.CryptFsConfig{Passphrase: kms.NewPlainSecret("secret")})
  275. require.NoError(t, err)
  276. u := dataprovider.User{
  277. BaseUser: sdk.BaseUser{
  278. Username: "test",
  279. HomeDir: os.TempDir(),
  280. },
  281. }
  282. conn := NewBaseConnection(fs.ConnectionID(), ProtocolSFTP, "", "", u)
  283. transfer := NewBaseTransfer(nil, conn, nil, testFile, testFile, "/transfer_test_file", TransferUpload,
  284. 0, 0, 0, 0, true, fs, dataprovider.TransferQuota{})
  285. transfer.ErrTransfer = errors.New("test error")
  286. _, _, err = transfer.getUploadFileSize()
  287. assert.Error(t, err)
  288. err = os.WriteFile(testFile, []byte("test data"), os.ModePerm)
  289. assert.NoError(t, err)
  290. size, deletedFiles, err := transfer.getUploadFileSize()
  291. assert.NoError(t, err)
  292. assert.Equal(t, int64(0), size)
  293. assert.Equal(t, 1, deletedFiles)
  294. assert.NoFileExists(t, testFile)
  295. }
  296. func TestFTPMode(t *testing.T) {
  297. conn := NewBaseConnection("", ProtocolFTP, "", "", dataprovider.User{})
  298. transfer := BaseTransfer{
  299. Connection: conn,
  300. transferType: TransferUpload,
  301. Fs: vfs.NewOsFs("", os.TempDir(), "", nil),
  302. }
  303. transfer.BytesReceived.Store(123)
  304. assert.Empty(t, transfer.ftpMode)
  305. transfer.SetFtpMode("active")
  306. assert.Equal(t, "active", transfer.ftpMode)
  307. }
  308. func TestTransferQuota(t *testing.T) {
  309. user := dataprovider.User{
  310. BaseUser: sdk.BaseUser{
  311. TotalDataTransfer: 3,
  312. UploadDataTransfer: 2,
  313. DownloadDataTransfer: 1,
  314. },
  315. }
  316. ul, dl, total := user.GetDataTransferLimits()
  317. assert.Equal(t, int64(2*1048576), ul)
  318. assert.Equal(t, int64(1*1048576), dl)
  319. assert.Equal(t, int64(3*1048576), total)
  320. user.TotalDataTransfer = -1
  321. user.UploadDataTransfer = -1
  322. user.DownloadDataTransfer = -1
  323. ul, dl, total = user.GetDataTransferLimits()
  324. assert.Equal(t, int64(0), ul)
  325. assert.Equal(t, int64(0), dl)
  326. assert.Equal(t, int64(0), total)
  327. transferQuota := dataprovider.TransferQuota{}
  328. assert.True(t, transferQuota.HasDownloadSpace())
  329. assert.True(t, transferQuota.HasUploadSpace())
  330. transferQuota.TotalSize = -1
  331. transferQuota.ULSize = -1
  332. transferQuota.DLSize = -1
  333. assert.True(t, transferQuota.HasDownloadSpace())
  334. assert.True(t, transferQuota.HasUploadSpace())
  335. transferQuota.TotalSize = 100
  336. transferQuota.AllowedTotalSize = 10
  337. assert.True(t, transferQuota.HasDownloadSpace())
  338. assert.True(t, transferQuota.HasUploadSpace())
  339. transferQuota.AllowedTotalSize = 0
  340. assert.False(t, transferQuota.HasDownloadSpace())
  341. assert.False(t, transferQuota.HasUploadSpace())
  342. transferQuota.TotalSize = 0
  343. transferQuota.DLSize = 100
  344. transferQuota.ULSize = 50
  345. transferQuota.AllowedTotalSize = 0
  346. assert.False(t, transferQuota.HasDownloadSpace())
  347. assert.False(t, transferQuota.HasUploadSpace())
  348. transferQuota.AllowedDLSize = 1
  349. transferQuota.AllowedULSize = 1
  350. assert.True(t, transferQuota.HasDownloadSpace())
  351. assert.True(t, transferQuota.HasUploadSpace())
  352. transferQuota.AllowedDLSize = -10
  353. transferQuota.AllowedULSize = -1
  354. assert.False(t, transferQuota.HasDownloadSpace())
  355. assert.False(t, transferQuota.HasUploadSpace())
  356. conn := NewBaseConnection("", ProtocolSFTP, "", "", user)
  357. transfer := NewBaseTransfer(nil, conn, nil, "file.txt", "file.txt", "/transfer_test_file", TransferUpload,
  358. 0, 0, 0, 0, true, vfs.NewOsFs("", os.TempDir(), "", nil), dataprovider.TransferQuota{})
  359. err := transfer.CheckRead()
  360. assert.NoError(t, err)
  361. err = transfer.CheckWrite()
  362. assert.NoError(t, err)
  363. transfer.transferQuota = dataprovider.TransferQuota{
  364. AllowedTotalSize: 10,
  365. }
  366. transfer.BytesReceived.Store(5)
  367. transfer.BytesSent.Store(4)
  368. err = transfer.CheckRead()
  369. assert.NoError(t, err)
  370. err = transfer.CheckWrite()
  371. assert.NoError(t, err)
  372. transfer.BytesSent.Store(6)
  373. err = transfer.CheckRead()
  374. if assert.Error(t, err) {
  375. assert.Contains(t, err.Error(), ErrReadQuotaExceeded.Error())
  376. }
  377. err = transfer.CheckWrite()
  378. assert.True(t, conn.IsQuotaExceededError(err))
  379. transferQuota = dataprovider.TransferQuota{
  380. AllowedTotalSize: 0,
  381. AllowedULSize: 10,
  382. AllowedDLSize: 5,
  383. }
  384. transfer.transferQuota = transferQuota
  385. assert.Equal(t, transferQuota, transfer.GetTransferQuota())
  386. err = transfer.CheckRead()
  387. if assert.Error(t, err) {
  388. assert.Contains(t, err.Error(), ErrReadQuotaExceeded.Error())
  389. }
  390. err = transfer.CheckWrite()
  391. assert.NoError(t, err)
  392. transfer.BytesReceived.Store(11)
  393. err = transfer.CheckRead()
  394. if assert.Error(t, err) {
  395. assert.Contains(t, err.Error(), ErrReadQuotaExceeded.Error())
  396. }
  397. err = transfer.CheckWrite()
  398. assert.True(t, conn.IsQuotaExceededError(err))
  399. }
  400. func TestUploadOutsideHomeRenameError(t *testing.T) {
  401. oldTempPath := Config.TempPath
  402. conn := NewBaseConnection("", ProtocolSFTP, "", "", dataprovider.User{})
  403. transfer := BaseTransfer{
  404. Connection: conn,
  405. transferType: TransferUpload,
  406. Fs: vfs.NewOsFs("", filepath.Join(os.TempDir(), "home"), "", nil),
  407. }
  408. transfer.BytesReceived.Store(123)
  409. fileName := filepath.Join(os.TempDir(), "_temp")
  410. err := os.WriteFile(fileName, []byte(`data`), 0644)
  411. assert.NoError(t, err)
  412. transfer.effectiveFsPath = fileName
  413. res := transfer.checkUploadOutsideHomeDir(os.ErrPermission)
  414. assert.Equal(t, 0, res)
  415. Config.TempPath = filepath.Clean(os.TempDir())
  416. res = transfer.checkUploadOutsideHomeDir(nil)
  417. assert.Equal(t, 0, res)
  418. assert.Greater(t, transfer.BytesReceived.Load(), int64(0))
  419. res = transfer.checkUploadOutsideHomeDir(os.ErrPermission)
  420. assert.Equal(t, 1, res)
  421. assert.Equal(t, int64(0), transfer.BytesReceived.Load())
  422. assert.NoFileExists(t, fileName)
  423. Config.TempPath = oldTempPath
  424. }