fix encrypting multipart upload issue

Signed-off-by: Dweb Fan <[email protected]>

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)
@@ -90,6 +92,9 @@ func mkISO(ctx *cli.Context) error {
 			len(files)-len(leftFiles), datasize.ByteSize(size).HR(), filename,
 			len(leftFiles), datasize.ByteSize(currentSizeNotInISO-size).HR())
 
+		if len(leftFiles) == 0 {
+			return nil
+		}
 		if len(ctx.Args()) > 0 {
 			fmt.Println("Please supply another filename")
 			return nil

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/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
 	}

cmd/lomob/upload-iso.go

@@ -341,6 +341,11 @@ func uploadEncryptParts(cli *clients.AWSClient, region, bucket, storageClass, is
 	}
 
 	salt := decoded[:crypto.SaltLen()]
+
+	// Derive key from passphrase using Argon2
+	// TODO: Using IV as salt for simplicity, change to different salt?
+	encryptKey := crypto.DeriveKeyFromMasterKey([]byte(masterKey), salt)
+
 	// iso size need add salt block size so as to compare with remote size
 	isoInfo.Size += crypto.SaltLen()
 	isoInfo.HashRemote = ""
@@ -356,8 +361,12 @@ func uploadEncryptParts(cli *clients.AWSClient, region, bucket, storageClass, is
 
 	partsHash := [][]byte{}
 
-	var start, end int64
-	var failParts []int
+	var (
+		start, end int64
+		failParts  []int
+		encryptor  *crypto.Encryptor
+		prs        *lomoio.FilePartReadSeeker
+	)
 	for i, p := range parts {
 		// add salt len for the last part
 		if i == len(parts)-1 {
@@ -392,18 +401,47 @@ func uploadEncryptParts(cli *clients.AWSClient, region, bucket, storageClass, is
 		defer os.Remove(tmpFilename)
 		defer tmpFile.Close()
 
-		prs := lomoio.NewFilePartReadSeeker(isoFile, start, end)
-		hl, hr, err := encryptLocalFile(prs, tmpFile, []byte(masterKey), salt, i == 0)
+		if prs == nil {
+			prs = lomoio.NewFilePartReadSeeker(isoFile, start, end)
+		} else {
+			prs.SetStartEnd(start, end)
+		}
+
+		hr := sha256.New()
+		mw := io.MultiWriter(hr, tmpFile)
+
+		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
 		}
-		p.SetHashLocal(hl)
-		p.SetHashRemote(hr)
 
+		if n != end-start {
+			return fmt.Errorf("write %d byte salt while expecting %d btw [%d, %d]", n, end-start, start, end)
+		}
+
+		hrData := hr.Sum(nil)
+		p.SetHashRemote(hrData)
+
+		// seek to beginning for upload
 		_, err = tmpFile.Seek(0, io.SeekStart)
 		if err != nil {
 			return err
 		}
+
 		p.Etag, err = cli.Upload(int64(p.PartNo), int64(p.Size), request, tmpFile, p.HashRemote)
 		if err != nil {
 			failParts = append(failParts, p.PartNo)
@@ -415,7 +453,7 @@ func uploadEncryptParts(cli *clients.AWSClient, region, bucket, storageClass, is
 			}
 			continue
 		}
-		partsHash = append(partsHash, hr)
+		partsHash = append(partsHash, hrData)
 		err = db.UpdatePartEtagAndStatusHash(p.IsoID, p.PartNo, p.Etag, p.HashLocal, p.HashRemote, types.PartUploaded)
 		if err != nil {
 			logrus.Infof("Update %s's part number %d status %s:%s", isoFilename, p.PartNo,
@@ -479,13 +517,17 @@ func uploadISO(accessKeyID, accessKey, region, bucket, storageClass, isoFilename
 }
 
 func uploadISOs(ctx *cli.Context) error {
-	partSize, err := datasize.ParseString(ctx.String("part-size"))
+	ps, err := datasize.ParseString(ctx.String("part-size"))
 	if err != nil {
 		return err
 	}
+	partSize := int(ps)
 	if partSize < 5*1024*1024 {
 		return errors.New("part size must be larger than 5*1024*1024=5242880")
 	}
+	if partSize%crypto.SaltLen() != 0 || (partSize-crypto.SaltLen())%crypto.SaltLen() != 0 {
+		return errors.Errorf("part size must be able to divided by salt length '%d'", crypto.SaltLen())
+	}
 
 	err = initDB(ctx.GlobalString("db"))
 	if err != nil {
@@ -525,7 +567,7 @@ func uploadISOs(ctx *cli.Context) error {
 
 	for _, isoFilename := range ctx.Args() {
 		err = uploadISO(accessKeyID, secretAccessKey, region, bucket, storageClass,
-			isoFilename, masterKey, int(partSize), saveParts, force)
+			isoFilename, masterKey, partSize, saveParts, force)
 		if err != nil {
 			return err
 		}

common/io/crypto.go

@@ -96,7 +96,7 @@ func (r *CryptoStreamReader) Read(p []byte) (n int, err error) {
 
 	r.offset += n
 
-	r.stream.XORKeyStream(p, buf)
+	r.stream.XORKeyStream(p[:n], buf[:n])
 
 	_, err = r.hashEncrypt.Write(p[:n])
 	if err != nil {

common/io/crypto_test.go

@@ -468,6 +468,67 @@ func TestCryptoStreamReaderSeekReadEncrypt(t *testing.T) {
 	verifyCryptoReadSeek(t, expectFile, r, 101, -1, -1, io.SeekCurrent, expectStream)
 }
 
+func TestCryptoStreamReaderEncryptLargeBuffer(t *testing.T) {
+	// use large buffer to read multiple times, and value should be same
+	nl := 16
+	nonce := make([]byte, nl)
+	for i := 0; i < nl; i++ {
+		nonce[i] = byte(i)
+	}
+
+	f, err := os.Open(testFilename)
+	require.Nil(t, err)
+	defer f.Close()
+
+	key, _ := hex.DecodeString("6368616e676520746869732070617373")
+
+	stream := getCryptoStream(t, key, nonce)
+
+	prs := NewFilePartReadSeeker(f, 0, 100)
+	r, err := NewCryptoStreamReader(prs, nonce, stream)
+	require.Nil(t, err)
+
+	expectFile, err := os.Open(testFilename)
+	require.Nil(t, err)
+	defer expectFile.Close()
+
+	// initial read will return nonce
+	buf := make([]byte, 200)
+	n, err := r.Read(buf)
+	require.Nil(t, err)
+	require.EqualValues(t, len(nonce), n)
+	require.EqualValues(t, nonce, buf[:n])
+
+	expectStream := getCryptoStream(t, key, nonce)
+
+	verifyCryptoLargeBuffer(t, expectFile, expectStream, 100, r)
+
+	prs.SetStartEnd(100, 200)
+	verifyCryptoLargeBuffer(t, expectFile, expectStream, 100, r)
+
+	prs.SetStartEnd(200, 201)
+	verifyCryptoLargeBuffer(t, expectFile, expectStream, 1, r)
+
+	prs.SetStartEnd(201, 300)
+	verifyCryptoLargeBuffer(t, expectFile, expectStream, 99, r)
+}
+
+func verifyCryptoLargeBuffer(t *testing.T, expectReader io.Reader, expectStream cipher.Stream,
+	expectLen int, stream *CryptoStreamReader) {
+	expectReadBuffer := make([]byte, expectLen)
+	expectBuffer := make([]byte, expectLen)
+	expectSize, err := expectReader.Read(expectReadBuffer)
+	require.Nil(t, err)
+	expectStream.XORKeyStream(expectBuffer, expectReadBuffer)
+
+	buffer := make([]byte, expectLen+100)
+	size, err := stream.Read(buffer)
+	require.Nil(t, err, "read lengh: %d", size)
+
+	require.Equal(t, expectSize, size)
+	require.Equal(t, expectBuffer, buffer[:size], "expect len: %d", expectLen)
+}
+
 func TestCryptoStreamReaderSeekReadEncryptNoNonce(t *testing.T) {
 	nl := 16
 	nonce := make([]byte, nl)