fix(download): fixed encryption header issues

code/go/0chain.net/blobbercore/handler/chunk_encoder.go

@@ -0,0 +1,88 @@
+package handler
+
+import (
+	"bytes"
+	"errors"
+	"strings"
+
+	zencryption "github.com/0chain/gosdk/zboxcore/encryption"
+)
+
+// ChunkEncoder encode/decode chunk data
+type ChunkEncoder interface {
+	// Encode encode chunk data if it is necessary
+	Encode(chunkSize int, data []byte) ([]byte, error)
+}
+
+// RawChunkEncoder raw chunk data
+type RawChunkEncoder struct {
+}
+
+// Encode read chunk data
+func (r *RawChunkEncoder) Encode(chunkSize int, data []byte) ([]byte, error) {
+	return data, nil
+}
+
+// PREChunkEncoder encode and decode chunk data with PREEncryptionScheme if file is shared with auth token
+type PREChunkEncoder struct {
+	EncryptedKey              string
+	ReEncryptionKey           string
+	ClientEncryptionPublicKey string
+}
+
+// Encode encode chunk data with PREEncryptionScheme for subscriber to download it
+func (r *PREChunkEncoder) Encode(chunkSize int, data []byte) ([]byte, error) {
+	encscheme := zencryption.NewEncryptionScheme()
+	if _, err := encscheme.Initialize(""); err != nil {
+		return nil, err
+	}
+
+	if err := encscheme.InitForDecryption("filetype:audio", r.EncryptedKey); err != nil {
+		return nil, err
+	}
+
+	totalSize := len(data)
+	result := []byte{}
+
+	for i := 0; i < totalSize; i += chunkSize {
+		encMsg := new(zencryption.EncryptedMessage)
+
+		nextIndex := i + chunkSize
+		var chunkData []byte
+		if nextIndex > totalSize {
+			chunkData = make([]byte, totalSize-i)
+			nextIndex = totalSize
+		} else {
+			chunkData = make([]byte, chunkSize)
+		}
+
+		copy(chunkData, data[i:nextIndex])
+
+		encMsg.EncryptedData = chunkData[EncryptionHeaderSize:]
+
+		headerBytes := chunkData[:EncryptionHeaderSize]
+		headerBytes = bytes.Trim(headerBytes, "\x00")
+		headerString := string(headerBytes)
+
+		headerChecksums := strings.Split(headerString, ",")
+		if len(headerChecksums) != 2 {
+			return nil, errors.New("Block has invalid encryption header")
+		}
+
+		encMsg.MessageChecksum, encMsg.OverallChecksum = headerChecksums[0], headerChecksums[1]
+		encMsg.EncryptedKey = encscheme.GetEncryptedKey()
+
+		reEncMsg, err := encscheme.ReEncrypt(encMsg, r.ReEncryptionKey, r.ClientEncryptionPublicKey)
+		if err != nil {
+			return nil, err
+		}
+
+		encData, err := reEncMsg.Marshal()
+		if err != nil {
+			return nil, err
+		}
+		// 256 bytes to save ReEncryption header instead of 2048 EncryptionHeader
+		result = append(result, encData...)
+	}
+	return result, nil
+}

code/go/0chain.net/blobbercore/handler/handler_test.go

@@ -1876,7 +1876,7 @@ func TestHandlers_Requiring_Signature(t *testing.T) {
 					t.Fatal(err)
 				}
 
-				header := make([]byte, HeaderChecksumSize)
+				header := make([]byte, EncryptionHeaderSize)
 				copy(header, encMsg.MessageChecksum+","+encMsg.OverallChecksum)
 				data := append(header, encMsg.EncryptedData...)
 				fmt.Println("Encrypted data: ", string(data))
@@ -2028,7 +2028,7 @@ func TestHandlers_Requiring_Signature(t *testing.T) {
 					t.Fatal(err)
 				}
 
-				header := make([]byte, HeaderChecksumSize)
+				header := make([]byte, EncryptionHeaderSize)
 				copy(header, encMsg.MessageChecksum+","+encMsg.OverallChecksum)
 				data := append(header, encMsg.EncryptedData...)
 				setMockFileBlock(data)
@@ -2183,7 +2183,7 @@ func TestHandlers_Requiring_Signature(t *testing.T) {
 					t.Fatal(err)
 				}
 
-				header := make([]byte, HeaderChecksumSize)
+				header := make([]byte, EncryptionHeaderSize)
 				copy(header, encMsg.MessageChecksum+","+encMsg.OverallChecksum)
 				data := append(header, encMsg.EncryptedData...)
 				setMockFileBlock(data)
@@ -2338,7 +2338,7 @@ func TestHandlers_Requiring_Signature(t *testing.T) {
 					t.Fatal(err)
 				}
 
-				header := make([]byte, HeaderChecksumSize)
+				header := make([]byte, EncryptionHeaderSize)
 				copy(header, encMsg.MessageChecksum+","+encMsg.OverallChecksum)
 				data := append(header, encMsg.EncryptedData...)
 				setMockFileBlock(data)

code/go/0chain.net/blobbercore/handler/object_operation_handler.go

@@ -1,17 +1,13 @@
 package handler
 
 import (
-	"bytes"
 	"context"
 	"encoding/hex"
 	"encoding/json"
 	"errors"
-	"math"
-	"strings"
 
 	"github.com/0chain/blobber/code/go/0chain.net/blobbercore/blobberhttp"
 	"github.com/0chain/blobber/code/go/0chain.net/blobbercore/stats"
-	zencryption "github.com/0chain/gosdk/zboxcore/encryption"
 
 	"net/http"
 	"path/filepath"
@@ -38,10 +34,10 @@ import (
 )
 
 const (
-	// EncryptionOverHead takes blockSize increment when data is incremented.
-	// messageCheckSum(128) + overallChecksum(128) + ","(1) + data-size-increment(16)
-	EncryptionOverHead = 273
-	HeaderChecksumSize = 2048 // Change it to 257 after header size is fixed in chunked upload as well
+	// EncryptionHeaderSize encryption header size in chunk
+	EncryptionHeaderSize = 2 * 1024
+	// ReEncryptionHeaderSize re-encryption header size in chunk
+	ReEncryptionHeaderSize = 256
 )
 
 func readPreRedeem(ctx context.Context, alloc *allocation.Allocation, numBlocks, pendNumBlocks int64, payerID string) (err error) {
@@ -371,53 +367,25 @@ func (fsh *StorageHandler) DownloadFile(ctx context.Context, r *http.Request) (r
 		return nil, common.NewErrorf("download_file", "couldn't save latest read marker")
 	}
 
-	if fileref.EncryptedKey != "" && authToken != nil {
-		encscheme := zencryption.NewEncryptionScheme()
-		if _, err := encscheme.Initialize(""); err != nil {
-			return nil, err
+	var chunkEncoder ChunkEncoder
+	if len(fileref.EncryptedKey) > 0 && authToken != nil {
+		chunkEncoder = &PREChunkEncoder{
+			EncryptedKey:              fileref.EncryptedKey,
+			ReEncryptionKey:           shareInfo.ReEncryptionKey,
+			ClientEncryptionPublicKey: shareInfo.ClientEncryptionPublicKey,
 		}
+	} else {
+		chunkEncoder = &RawChunkEncoder{}
+	}
 
-		if err := encscheme.InitForDecryption("filetype:audio", fileref.EncryptedKey); err != nil {
-			return nil, err
-		}
+	chunkData, err := chunkEncoder.Encode(int(fileref.ChunkSize), respData)
 
-		totalSize := len(respData)
-		result := []byte{}
-		for i := 0; i < totalSize; i += int(fileref.ChunkSize) {
-			encMsg := new(zencryption.EncryptedMessage)
-			chunkData := respData[i:int64(math.Min(float64(i+int(fileref.ChunkSize)), float64(totalSize)))]
-
-			encMsg.EncryptedData = chunkData[HeaderChecksumSize:]
-
-			headerBytes := chunkData[:HeaderChecksumSize]
-			headerBytes = bytes.Trim(headerBytes, "\x00")
-			headerString := string(headerBytes)
-
-			headerChecksums := strings.Split(headerString, ",")
-			if len(headerChecksums) != 2 {
-				Logger.Error("Block has invalid header", zap.String("request Url", r.URL.String()))
-				return nil, errors.New("Block has invalid header for request " + r.URL.String())
-			}
-
-			encMsg.MessageChecksum, encMsg.OverallChecksum = headerChecksums[0], headerChecksums[1]
-			encMsg.EncryptedKey = encscheme.GetEncryptedKey()
-
-			reEncMsg, err := encscheme.ReEncrypt(encMsg, shareInfo.ReEncryptionKey, shareInfo.ClientEncryptionPublicKey)
-			if err != nil {
-				return nil, err
-			}
-
-			encData, err := reEncMsg.Marshal()
-			if err != nil {
-				return nil, err
-			}
-			result = append(result, encData...)
-		}
-		respData = result
+	if err != nil {
+		return nil, err
 	}
 
 	stats.FileBlockDownloaded(ctx, fileref.ID)
-	return respData, nil
+	return chunkData, nil
 }
 
 func (fsh *StorageHandler) CommitWrite(ctx context.Context, r *http.Request) (*blobberhttp.CommitResult, error) {