Merge pull request #18 from lomorage/dev

fix encrypting multipart upload issue

cmd/lomob/crypt.go

@@ -8,10 +8,13 @@ import (
 	"io"
 	"os"
 	"reflect"
+	"strconv"
 	"syscall"
 
 	"github.com/lomorage/lomo-backup/common/crypto"
+	"github.com/lomorage/lomo-backup/common/datasize"
 	lomohash "github.com/lomorage/lomo-backup/common/hash"
+	lomoio "github.com/lomorage/lomo-backup/common/io"
 	"github.com/urfave/cli"
 	"golang.org/x/term"
 )
@@ -65,11 +68,7 @@ func encryptCmd(ctx *cli.Context) error {
 		return errors.New("usage: [input filename] [[output filename]]. If output filename is not given, it will be <intput filename>.enc")
 	}
 
-	salt, err := genSalt(ifilename)
-	if err != nil {
-		return err
-	}
-
+	var err error
 	masterKey := ctx.String("encrypt-key")
 	if masterKey == "" {
 		masterKey, err = getMasterKey()
@@ -78,6 +77,11 @@ func encryptCmd(ctx *cli.Context) error {
 		}
 	}
 
+	salt, err := genSalt(ifilename)
+	if err != nil {
+		return err
+	}
+
 	src, err := os.Open(ifilename)
 	if err != nil {
 		return err
@@ -91,30 +95,117 @@ func encryptCmd(ctx *cli.Context) error {
 	defer dst.Close()
 
 	fmt.Printf("Start encrypt '%s', and save output to '%s'\n", ifilename, ofilename)
-	_, _, err = encryptLocalFile(src, dst, []byte(masterKey), salt, true)
+
+	ps := ctx.String("part-size")
+	if ps == "" {
+		_, err = encryptLocalFile(src, dst, []byte(masterKey), salt, true)
+		if err != nil {
+			return err
+		}
+
+		fmt.Println("Finish encryption!")
+
+		return nil
+	}
+
+	// Derive key from passphrase using Argon2
+	// TODO: Using IV as salt for simplicity, change to different salt?
+	encryptKey := crypto.DeriveKeyFromMasterKey([]byte(masterKey), salt)
+
+	partSize, err := datasize.ParseString(ps)
+	if err != nil {
+		return err
+	}
+
+	stat, err := src.Stat()
 	if err != nil {
 		return err
 	}
 
+	index := 1
+	remaining := stat.Size()
+	var (
+		start, end, curr, partLength int64
+		encryptor                    *crypto.Encryptor
+		prs                          *lomoio.FilePartReadSeeker
+	)
+	for curr = 0; remaining != 0; curr += partLength {
+		if remaining < int64(partSize) {
+			partLength = remaining
+		} else {
+			partLength = int64(partSize)
+		}
+
+		if curr == 0 {
+			end = int64(int(partLength) - crypto.SaltLen())
+		} else {
+			start = end
+			end += partLength
+		}
+
+		// create a local tmpfile and save intermittent part
+		pf, err := os.Create(ofilename + ".part" + strconv.Itoa(index))
+		if err != nil {
+			return err
+		}
+		defer pf.Close()
+
+		mw := io.MultiWriter(dst, pf)
+
+		if prs == nil {
+			prs = lomoio.NewFilePartReadSeeker(src, start, end)
+		} else {
+			prs.SetStartEnd(start, end)
+		}
+
+		if encryptor == nil {
+			encryptor, err = crypto.NewEncryptor(prs, encryptKey, salt, false)
+			if err != nil {
+				return err
+			}
+			n, err := mw.Write(salt)
+			if err != nil {
+				return err
+			}
+			if n != len(salt) {
+				return fmt.Errorf("write %d byte salt while expecting %d", n, len(salt))
+			}
+		}
+
+		n, err := io.Copy(mw, encryptor)
+		if err != nil {
+			return err
+		}
+
+		if n != end-start {
+			return fmt.Errorf("write %d byte salt while expecting %d btw [%d, %d]", n, end-start, start, end)
+		}
+
+		fmt.Printf("Created '%s'\n", pf.Name())
+
+		index++
+		remaining -= end - start
+	}
+
 	fmt.Println("Finish encryption!")
 
 	return nil
 }
 
-func encryptLocalFile(src io.ReadSeeker, dst io.Writer, masterKey, iv []byte, hasHeader bool) ([]byte, []byte, error) {
+func encryptLocalFile(src io.ReadSeeker, dst io.Writer, masterKey, iv []byte, hasHeader bool) ([]byte, error) {
 	// Derive key from passphrase using Argon2
 	// TODO: Using IV as salt for simplicity, change to different salt?
 	encryptKey := crypto.DeriveKeyFromMasterKey(masterKey, iv)
 	encryptor, err := crypto.NewEncryptor(src, encryptKey, iv, hasHeader)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	_, err = io.Copy(dst, encryptor)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
-	return encryptor.GetHashOrig(), encryptor.GetHashEncrypt(), nil
+	return encryptor.GetHashEncrypt(), nil
 }
 
 func decryptLocalFile(ctx *cli.Context) error {

cmd/lomob/iso.go

@@ -63,6 +63,8 @@ func mkISO(ctx *cli.Context) error {
 		isoFilename = ctx.Args()[0]
 	}
 
+	logrus.Infof("Total %d files (%s)", len(files), datasize.ByteSize(currentSizeNotInISO).HR())
+
 	for {
 		if currentSizeNotInISO < isoSize.Bytes() {
 			currSize := datasize.ByteSize(currentSizeNotInISO)
@@ -82,14 +84,38 @@ func mkISO(ctx *cli.Context) error {
 				datasize.ByteSize(iso.Size).HR())
 		}
 
-		size, filename, leftFiles, err := createIso(isoSize.Bytes(), isoFilename, scanRootDirs, files)
+		size, filename, leftFiles, notExistFiles, err := createIso(isoSize.Bytes(), isoFilename, scanRootDirs, files)
 		if err != nil {
 			return err
 		}
-		logrus.Infof("%d files (%s) are added into %s, and %d files (%s) need to be added",
-			len(files)-len(leftFiles), datasize.ByteSize(size).HR(), filename,
-			len(leftFiles), datasize.ByteSize(currentSizeNotInISO-size).HR())
+		if len(notExistFiles) == 0 {
+			logrus.Infof("%d files (%s) are added into %s, and %d files (%s) need to be added",
+				len(files)-len(leftFiles), datasize.ByteSize(size).HR(), filename,
+				len(leftFiles), datasize.ByteSize(currentSizeNotInISO-size).HR())
+		} else {
+			fileIDs := bytes.Buffer{}
+			notExistSizes := 0
+			for _, f := range notExistFiles {
+				notExistSizes += f.Size
+				fileIDs.WriteString(strconv.Itoa(f.ID))
+				fileIDs.WriteString(",")
+			}
+			logrus.Infof("%d files (%s) are added into %s, %d files (%s) need to be added, %d files (%s) not exist",
+				len(files)-len(leftFiles), datasize.ByteSize(size).HR(), filename,
+				len(leftFiles), datasize.ByteSize(currentSizeNotInISO-size).HR(),
+				len(notExistFiles), datasize.ByteSize(notExistSizes).HR())
 
+			ids := fileIDs.String()
+			_, err = db.DeleteBatchFiles(strings.Trim(ids, ","))
+			if err != nil {
+				return err
+			}
+			size += uint64(notExistSizes)
+		}
+
+		if len(leftFiles) == 0 {
+			return nil
+		}
 		if len(ctx.Args()) > 0 {
 			fmt.Println("Please supply another filename")
 			return nil
@@ -125,18 +151,19 @@ func createFileInStaging(srcFile, dstFile string) error {
 }
 
 func createIso(maxSize uint64, isoFilename string, scanRootDirs map[int]string,
-	files []*types.FileInfo) (uint64, string, []*types.FileInfo, error) {
+	files []*types.FileInfo) (uint64, string, []*types.FileInfo, []*types.FileInfo, error) {
 	stagingDir, err := os.MkdirTemp("", "lomobackup-")
 	if err != nil {
-		return 0, "", nil, err
+		return 0, "", nil, nil, err
 	}
 	defer os.RemoveAll(stagingDir)
 
 	const seperater = ','
 	var (
-		fileCount int
-		filesSize uint64
-		end       time.Time
+		fileCount     int
+		filesSize     uint64
+		end           time.Time
+		notExistFiles []*types.FileInfo
 	)
 	start := futuretime
 	fileIDs := bytes.Buffer{}
@@ -164,6 +191,11 @@ func createIso(maxSize uint64, isoFilename string, scanRootDirs map[int]string,
 
 		err = createFileInStaging(srcFile, dstFile)
 		if err != nil {
+			if os.IsNotExist(err) {
+				notExistFiles = append(notExistFiles, f)
+				logrus.Warnf("'%s' not exist anymore", srcFile)
+				continue
+			}
 			logrus.Warnf("Add %s into %s:%s: %s", srcFile, isoFilename, dstFile, err)
 			continue
 		}
@@ -207,18 +239,18 @@ func createIso(maxSize uint64, isoFilename string, scanRootDirs map[int]string,
 			stagingDir).CombinedOutput()
 		if err != nil {
 			fmt.Println(string(out))
-			return 0, "", nil, err
+			return 0, "", nil, nil, err
 		}
 
 		fileInfo, err := os.Stat(isoFilename)
 		if err != nil {
-			return 0, "", nil, err
+			return 0, "", nil, nil, err
 		}
 		isoInfo := &types.ISOInfo{Name: isoFilename, Size: int(fileInfo.Size())}
 
 		hash, err := lomohash.CalculateHashFile(isoFilename)
 		if err != nil {
-			return 0, "", nil, err
+			return 0, "", nil, nil, err
 		}
 		isoInfo.SetHashLocal(hash)
 		// create db entry and update file info
@@ -230,10 +262,10 @@ func createIso(maxSize uint64, isoFilename string, scanRootDirs map[int]string,
 		}
 
 		logrus.Infof("Takes %s to update iso_id for %d files in DB", time.Since(start).Truncate(time.Second).String(), count)
-		return filesSize, isoFilename, files[idx+1:], err
+		return filesSize, isoFilename, files[idx+1:], notExistFiles, err
 	}
 
-	return filesSize, isoFilename, nil, nil
+	return filesSize, isoFilename, nil, nil, nil
 }
 
 func listISO(ctx *cli.Context) error {

cmd/lomob/main.go

@@ -390,6 +390,10 @@ func main() {
 							Usage:  "Master key to encrypt current upload file",
 							EnvVar: "LOMOB_MASTER_KEY",
 						},
+						cli.StringFlag{
+							Name:  "part-size,p",
+							Usage: "Size of each upload partition. KB=1000 Byte. 0 means no part. Mainly for local test purpose",
+						},
 					},
 				},
 				{

cmd/lomob/script/create_db.sh

@@ -1 +1,2 @@
+rm ../lomob.db
 sqlite3 ../lomob.db < ../../../common/dbx/schema/1.sql

cmd/lomob/upload-files.go

@@ -327,7 +327,7 @@ func uploadEncryptFileToS3(cli *clients.AWSClient, bucket, storageClass, filenam
 	tmpFileName := tmpFile.Name()
 	defer tmpFile.Close()
 
-	_, hash, err := encryptLocalFile(src, tmpFile, []byte(masterKey), salt, true)
+	hash, err := encryptLocalFile(src, tmpFile, []byte(masterKey), salt, true)
 	if err != nil {
 		return "", err
 	}