浏览代码

eventmanager: check disk quota before executing the compress action

Signed-off-by: Nicola Murino <nicola.murino@gmail.com>
Nicola Murino 2 年之前
父节点
当前提交
6cebc037a0
共有 5 个文件被更改,包括 259 次插入6 次删除
  1. 1 1
      go.mod
  2. 2 3
      go.sum
  3. 51 2
      internal/common/eventmanager.go
  4. 37 0
      internal/common/eventmanager_test.go
  5. 168 0
      internal/common/protocol_test.go

+ 1 - 1
go.mod

@@ -17,7 +17,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.29.6
 	github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.17.0
 	github.com/aws/aws-sdk-go-v2/service/sts v1.17.7
-	github.com/cockroachdb/cockroach-go/v2 v2.2.19
+	github.com/cockroachdb/cockroach-go/v2 v2.2.20
 	github.com/coreos/go-oidc/v3 v3.4.0
 	github.com/drakkan/webdav v0.0.0-20221101181759-17ed21f9337b
 	github.com/eikenb/pipeat v0.0.0-20210730190139-06b3e6902001

+ 2 - 3
go.sum

@@ -360,8 +360,8 @@ github.com/cncf/xds/go v0.0.0-20211001041855-01bcc9b48dfe/go.mod h1:eXthEFrGJvWH
 github.com/cncf/xds/go v0.0.0-20211011173535-cb28da3451f1/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
 github.com/cncf/xds/go v0.0.0-20220314180256-7f1daf1720fc/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
-github.com/cockroachdb/cockroach-go/v2 v2.2.19 h1:YIHyz17jZumBeXPuoZKq/0nrITsqDoDD8/KQt3/xiyc=
-github.com/cockroachdb/cockroach-go/v2 v2.2.19/go.mod h1:mzlIDDBALQfEjv/7DU12fb2AfQ/MUYTlychcMpWp9QI=
+github.com/cockroachdb/cockroach-go/v2 v2.2.20 h1:TLSzwdTdIwgsbdApHzaxunhSMrmbGf5YY6oxtaP2kvw=
+github.com/cockroachdb/cockroach-go/v2 v2.2.20/go.mod h1:73vQi5H/H7kE8SgOt+XA6729Tubvj5hxKIEgbQQhp4c=
 github.com/cockroachdb/datadriven v0.0.0-20190809214429-80d97fb3cbaa/go.mod h1:zn76sxSg3SzpJ0PPJaLDCu+Bu0Lg3sKTORVIj19EIF8=
 github.com/cockroachdb/datadriven v0.0.0-20200714090401-bf6692d28da5/go.mod h1:h6jFvWxBdQXxjopDMZyH2UVceIRfR84bdzbkoKrsWNo=
 github.com/cockroachdb/errors v1.2.4/go.mod h1:rQD95gz6FARkaKkQXUksEje/d9a6wBJoCr5oaCLELYA=
@@ -1337,7 +1337,6 @@ github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4/go.mod h1:N6UoU20jOqgg
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1-0.20171018195549-f15c970de5b7/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
-github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/profile v1.2.1/go.mod h1:hJw3o1OdXxsrSjjVksARp5W95eeEaEfptyVZyv6JUPA=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

+ 51 - 2
internal/common/eventmanager.go

@@ -1615,6 +1615,50 @@ func getArchiveBaseDir(paths []string) string {
 	return baseDir
 }
 
+func getSizeForPath(conn *BaseConnection, p string, info os.FileInfo) (int64, error) {
+	if info.IsDir() {
+		var dirSize int64
+		entries, err := conn.ListDir(p)
+		if err != nil {
+			return 0, err
+		}
+		for _, entry := range entries {
+			size, err := getSizeForPath(conn, path.Join(p, entry.Name()), entry)
+			if err != nil {
+				return 0, err
+			}
+			dirSize += size
+		}
+		return dirSize, nil
+	}
+	if info.Mode().IsRegular() {
+		return info.Size(), nil
+	}
+	return 0, nil
+}
+
+func estimateZipSize(conn *BaseConnection, zipPath string, paths []string) (int64, error) {
+	q, _ := conn.HasSpace(false, false, zipPath)
+	if q.HasSpace && q.GetRemainingSize() > 0 {
+		var size int64
+		for _, item := range paths {
+			info, err := conn.DoStat(item, 1, false)
+			if err != nil {
+				return size, err
+			}
+			itemSize, err := getSizeForPath(conn, item, info)
+			if err != nil {
+				return size, err
+			}
+			size += itemSize
+		}
+		eventManagerLog(logger.LevelDebug, "archive paths %v, archive name %q, size: %d", paths, zipPath, size)
+		// we assume the zip size will be half of the real size
+		return size / 2, nil
+	}
+	return -1, nil
+}
+
 func executeCompressFsActionForUser(c dataprovider.EventActionFsCompress, replacer *strings.Replacer,
 	user dataprovider.User,
 ) error {
@@ -1638,14 +1682,19 @@ func executeCompressFsActionForUser(c dataprovider.EventActionFsCompress, replac
 		}
 		paths = append(paths, p)
 	}
-	writer, numFiles, truncatedSize, cancelFn, err := getFileWriter(conn, name, -1)
+	paths = util.RemoveDuplicates(paths, false)
+	estimatedSize, err := estimateZipSize(conn, name, paths)
+	if err != nil {
+		eventManagerLog(logger.LevelError, "unable to estimate size for archive %q: %v", name, err)
+		return fmt.Errorf("unable to estimate archive size: %w", err)
+	}
+	writer, numFiles, truncatedSize, cancelFn, err := getFileWriter(conn, name, estimatedSize)
 	if err != nil {
 		eventManagerLog(logger.LevelError, "unable to create archive %q: %v", name, err)
 		return fmt.Errorf("unable to create archive: %w", err)
 	}
 	defer cancelFn()
 
-	paths = util.RemoveDuplicates(paths, false)
 	baseDir := getArchiveBaseDir(paths)
 	eventManagerLog(logger.LevelDebug, "creating archive %q for paths %+v", name, paths)
 

+ 37 - 0
internal/common/eventmanager_test.go

@@ -1720,6 +1720,43 @@ func TestReplacePathsPlaceholders(t *testing.T) {
 	assert.Equal(t, []string{"/path1", "/path2"}, paths)
 }
 
+func TestEstimateZipSizeErrors(t *testing.T) {
+	u := dataprovider.User{
+		BaseUser: sdk.BaseUser{
+			Username: "u",
+			HomeDir:  filepath.Join(os.TempDir(), "u"),
+			Status:   1,
+			Permissions: map[string][]string{
+				"/": {dataprovider.PermAny},
+			},
+			QuotaSize: 1000,
+		},
+	}
+	err := dataprovider.AddUser(&u, "", "", "")
+	assert.NoError(t, err)
+	err = os.MkdirAll(u.GetHomeDir(), os.ModePerm)
+	assert.NoError(t, err)
+	conn := NewBaseConnection("", ProtocolFTP, "", "", u)
+	_, err = getSizeForPath(conn, "/missing", vfs.NewFileInfo("missing", true, 0, time.Now(), false))
+	assert.True(t, conn.IsNotExistError(err))
+	if runtime.GOOS != osWindows {
+		err = os.MkdirAll(filepath.Join(u.HomeDir, "d1", "d2", "sub"), os.ModePerm)
+		assert.NoError(t, err)
+		err = os.WriteFile(filepath.Join(u.HomeDir, "d1", "d2", "sub", "file.txt"), []byte("data"), 0666)
+		assert.NoError(t, err)
+		err = os.Chmod(filepath.Join(u.HomeDir, "d1", "d2"), 0001)
+		assert.NoError(t, err)
+		size, err := estimateZipSize(conn, "/archive.zip", []string{"/d1"})
+		assert.Error(t, err, "size %d", size)
+		err = os.Chmod(filepath.Join(u.HomeDir, "d1", "d2"), os.ModePerm)
+		assert.NoError(t, err)
+	}
+	err = dataprovider.DeleteUser(u.Username, "", "", "")
+	assert.NoError(t, err)
+	err = os.RemoveAll(u.GetHomeDir())
+	assert.NoError(t, err)
+}
+
 func getErrorString(err error) string {
 	if err == nil {
 		return ""

+ 168 - 0
internal/common/protocol_test.go

@@ -4714,6 +4714,174 @@ func TestEventActionCompress(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+func TestEventActionCompressQuotaErrors(t *testing.T) {
+	smtpCfg := smtp.Config{
+		Host:          "127.0.0.1",
+		Port:          2525,
+		From:          "notify@example.com",
+		TemplatesPath: "templates",
+	}
+	err := smtpCfg.Initialize(configDir)
+	require.NoError(t, err)
+
+	testDir := "archiveDir"
+	zipPath := "/archive.zip"
+	a1 := dataprovider.BaseEventAction{
+		Name: "action1",
+		Type: dataprovider.ActionTypeFilesystem,
+		Options: dataprovider.BaseEventActionOptions{
+			FsConfig: dataprovider.EventActionFilesystemConfig{
+				Type: dataprovider.FilesystemActionCompress,
+				Compress: dataprovider.EventActionFsCompress{
+					Name:  zipPath,
+					Paths: []string{"/" + testDir},
+				},
+			},
+		},
+	}
+	action1, _, err := httpdtest.AddEventAction(a1, http.StatusCreated)
+	assert.NoError(t, err)
+	a2 := dataprovider.BaseEventAction{
+		Name: "action2",
+		Type: dataprovider.ActionTypeEmail,
+		Options: dataprovider.BaseEventActionOptions{
+			EmailConfig: dataprovider.EventActionEmailConfig{
+				Recipients: []string{"test@example.com"},
+				Subject:    `"Compress failed"`,
+				Body:       "Error: {{ErrorString}}",
+			},
+		},
+	}
+	action2, _, err := httpdtest.AddEventAction(a2, http.StatusCreated)
+	assert.NoError(t, err)
+	r1 := dataprovider.EventRule{
+		Name:    "test compress",
+		Trigger: dataprovider.EventTriggerFsEvent,
+		Conditions: dataprovider.EventConditions{
+			FsEvents: []string{"rename"},
+		},
+		Actions: []dataprovider.EventAction{
+			{
+				BaseEventAction: dataprovider.BaseEventAction{
+					Name: action1.Name,
+				},
+				Order: 1,
+			},
+			{
+				BaseEventAction: dataprovider.BaseEventAction{
+					Name: action2.Name,
+				},
+				Options: dataprovider.EventActionOptions{
+					IsFailureAction: true,
+				},
+				Order: 2,
+			},
+		},
+	}
+	rule1, _, err := httpdtest.AddEventRule(r1, http.StatusCreated)
+	assert.NoError(t, err)
+	fileSize := int64(100)
+	u := getTestUser()
+	u.QuotaSize = 10 * fileSize
+	user, _, err := httpdtest.AddUser(u, http.StatusCreated)
+	assert.NoError(t, err)
+	conn, client, err := getSftpClient(user)
+	if assert.NoError(t, err) {
+		defer conn.Close()
+		defer client.Close()
+
+		err = client.MkdirAll(path.Join(testDir, "1", "1"))
+		assert.NoError(t, err)
+		err = writeSFTPFile(path.Join(testDir, "1", testFileName), fileSize, client)
+		assert.NoError(t, err)
+		err = client.MkdirAll(path.Join(testDir, "2", "2"))
+		assert.NoError(t, err)
+		err = writeSFTPFile(path.Join(testDir, "2", testFileName), fileSize, client)
+		assert.NoError(t, err)
+		err = client.Symlink(path.Join(testDir, "2", testFileName), path.Join(testDir, "2", testFileName+"_link"))
+		assert.NoError(t, err)
+		// trigger the compress action
+		err = client.Mkdir("a")
+		assert.NoError(t, err)
+		err = client.Rename("a", "b")
+		assert.NoError(t, err)
+		assert.Eventually(t, func() bool {
+			_, err := client.Stat(zipPath)
+			return err == nil
+		}, 3*time.Second, 100*time.Millisecond)
+		err = client.Remove(zipPath)
+		assert.NoError(t, err)
+		// add other 6 file, the compress action should fail with a quota error
+		err = writeSFTPFile(path.Join(testDir, "1", "1", testFileName), fileSize, client)
+		assert.NoError(t, err)
+		err = writeSFTPFile(path.Join(testDir, "2", "2", testFileName), fileSize, client)
+		assert.NoError(t, err)
+		err = writeSFTPFile(path.Join(testDir, "1", "1", testFileName+"1"), fileSize, client)
+		assert.NoError(t, err)
+		err = writeSFTPFile(path.Join(testDir, "2", "2", testFileName+"2"), fileSize, client)
+		assert.NoError(t, err)
+		err = writeSFTPFile(path.Join(testDir, "1", testFileName+"1"), fileSize, client)
+		assert.NoError(t, err)
+		err = writeSFTPFile(path.Join(testDir, "2", testFileName+"2"), fileSize, client)
+		assert.NoError(t, err)
+		lastReceivedEmail.reset()
+		err = client.Rename("b", "a")
+		assert.NoError(t, err)
+		assert.Eventually(t, func() bool {
+			return lastReceivedEmail.get().From != ""
+		}, 3*time.Second, 100*time.Millisecond)
+		email := lastReceivedEmail.get()
+		assert.Len(t, email.To, 1)
+		assert.True(t, util.Contains(email.To, "test@example.com"))
+		assert.Contains(t, email.Data, `Subject: "Compress failed"`)
+		assert.Contains(t, email.Data, common.ErrQuotaExceeded.Error())
+		// update quota size so the user is already overquota
+		user.QuotaSize = 7 * fileSize
+		_, _, err = httpdtest.UpdateUser(user, http.StatusOK, "")
+		assert.NoError(t, err)
+		lastReceivedEmail.reset()
+		err = client.Rename("a", "b")
+		assert.NoError(t, err)
+		assert.Eventually(t, func() bool {
+			return lastReceivedEmail.get().From != ""
+		}, 3*time.Second, 100*time.Millisecond)
+		email = lastReceivedEmail.get()
+		assert.Len(t, email.To, 1)
+		assert.True(t, util.Contains(email.To, "test@example.com"))
+		assert.Contains(t, email.Data, `Subject: "Compress failed"`)
+		assert.Contains(t, email.Data, common.ErrQuotaExceeded.Error())
+		// remove the path to compress to trigger an error for size estimation
+		out, err := runSSHCommand(fmt.Sprintf("sftpgo-remove %s", testDir), user)
+		assert.NoError(t, err, string(out))
+		lastReceivedEmail.reset()
+		err = client.Rename("b", "a")
+		assert.NoError(t, err)
+		assert.Eventually(t, func() bool {
+			return lastReceivedEmail.get().From != ""
+		}, 3*time.Second, 100*time.Millisecond)
+		email = lastReceivedEmail.get()
+		assert.Len(t, email.To, 1)
+		assert.True(t, util.Contains(email.To, "test@example.com"))
+		assert.Contains(t, email.Data, `Subject: "Compress failed"`)
+		assert.Contains(t, email.Data, "unable to estimate archive size")
+	}
+
+	_, err = httpdtest.RemoveEventRule(rule1, http.StatusOK)
+	assert.NoError(t, err)
+	_, err = httpdtest.RemoveEventAction(action1, http.StatusOK)
+	assert.NoError(t, err)
+	_, err = httpdtest.RemoveEventAction(action2, http.StatusOK)
+	assert.NoError(t, err)
+	_, err = httpdtest.RemoveUser(user, http.StatusOK)
+	assert.NoError(t, err)
+	err = os.RemoveAll(user.GetHomeDir())
+	assert.NoError(t, err)
+
+	smtpCfg = smtp.Config{}
+	err = smtpCfg.Initialize(configDir)
+	require.NoError(t, err)
+}
+
 func TestEventActionCompressQuotaFolder(t *testing.T) {
 	testDir := "/folder"
 	a1 := dataprovider.BaseEventAction{