From ac838dd5e14ee778202ba31ac5ba546992954d4e Mon Sep 17 00:00:00 2001
From: Chun-Hung Tseng <henrybear327@gmail.com>
Date: Wed, 28 Jun 2023 22:28:49 +0200
Subject: [PATCH] Add downloaded block hash verification

---
 README.md |  2 +-
 crypto.go | 14 +++++++++++++-
 error.go  |  3 ++-
 file.go   |  2 +-
 4 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index ce2ec38..6177010 100644
--- a/README.md
+++ b/README.md
@@ -57,8 +57,8 @@ Currently, the development are split into 2 versions. V1 supports the features [
         - [x] Download empty file
         - [x] Properly handle large files and empty files (check iOS codebase)
             - esp. large files, where buffering in-memory will screw up the runtime
+        - [x] Check signature and hash
         - [ ] Improve large file handling
-        - [ ] Check signature and hash
     - [x] Delete
     - [x] Upload
         - [x] Handle empty file        
diff --git a/crypto.go b/crypto.go
index 3bc1bba..b188be4 100644
--- a/crypto.go
+++ b/crypto.go
@@ -1,6 +1,7 @@
 package proton_api_bridge
 
 import (
+	"crypto/sha256"
 	"encoding/base64"
 	"io"
 
@@ -126,7 +127,7 @@ func getKeyRing(kr, addrKR *crypto.KeyRing, key, passphrase, passphraseSignature
 	return crypto.NewKeyRing(unlockedKey)
 }
 
-func decryptBlockIntoBuffer(sessionKey *crypto.SessionKey, addrKR, nodeKR *crypto.KeyRing, encSignature string, buffer io.ReaderFrom, block io.ReadCloser) error {
+func decryptBlockIntoBuffer(sessionKey *crypto.SessionKey, addrKR, nodeKR *crypto.KeyRing, originalHash, encSignature string, buffer io.ReaderFrom, block io.ReadCloser) error {
 	data, err := io.ReadAll(block)
 	if err != nil {
 		return err
@@ -152,5 +153,16 @@ func decryptBlockIntoBuffer(sessionKey *crypto.SessionKey, addrKR, nodeKR *crypt
 		return err
 	}
 
+	h := sha256.New()
+	h.Write(data)
+	hash := h.Sum(nil)
+	base64Hash := base64.StdEncoding.EncodeToString(hash)
+	if err != nil {
+		return err
+	}
+	if base64Hash != originalHash {
+		return ErrDownloadedBlockHashVerificationFailed
+	}
+
 	return nil
 }
diff --git a/error.go b/error.go
index 5c778f2..2a7f64e 100644
--- a/error.go
+++ b/error.go
@@ -10,6 +10,7 @@ var (
 	ErrFolderIsNotEmpty                      = errors.New("folder can't be deleted becuase it is not empty")
 	ErrInternalErrorOnFileUpload             = errors.New("either link or createFileResp must be not nil")
 	ErrMissingInputUploadAndCollectBlockData = errors.New("missing either session key or key ring")
-	ErrLinkMustNotBeNil                      = errors.New("Missing input proton link")
+	ErrLinkMustNotBeNil                      = errors.New("missing input proton link")
 	ErrLinkMustBeActive                      = errors.New("can not operate on link state other than active")
+	ErrDownloadedBlockHashVerificationFailed = errors.New("the hash of the downloaded block doesn't match the original hash")
 )
diff --git a/file.go b/file.go
index 97cd2ca..5745ebf 100644
--- a/file.go
+++ b/file.go
@@ -121,7 +121,7 @@ func (protonDrive *ProtonDrive) DownloadFile(ctx context.Context, link *proton.L
 		}
 		defer blockReader.Close()
 
-		err = decryptBlockIntoBuffer(sessionKey, protonDrive.AddrKR, nodeKR, revision.Blocks[i].EncSignature, buffer, blockReader)
+		err = decryptBlockIntoBuffer(sessionKey, protonDrive.AddrKR, nodeKR, revision.Blocks[i].Hash, revision.Blocks[i].EncSignature, buffer, blockReader)
 		if err != nil {
 			return nil, nil, err
 		}