From 7d274257bc740f9bd750757066796c61114b872d Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@gmail.com>
Date: Sat, 15 Feb 2020 20:55:19 +0100
Subject: [PATCH] Change storing of the data to use multiple chunks

Data needs to be stored in multiple chunks with a maximum size of 1MB in order for image viewers to still handle the files. Especially libpng had problems with too large chunks sizes.
---
 cryptpng.go | 20 +++++++++++++++-----
 pngUtils.go | 11 +++++++++++
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/cryptpng.go b/cryptpng.go
index c03c8ef..e8cfbe1 100644
--- a/cryptpng.go
+++ b/cryptpng.go
@@ -14,6 +14,7 @@ import (
 	"log"
 	"os"
 	"syscall"
+	"math"
 
 	"golang.org/x/crypto/ssh/terminal"
 )
@@ -25,6 +26,7 @@ func check(err error) {
 }
 
 const chunkName = "crPt"
+const chunkSize = 0x100000
 
 var inputFile string
 var outputFile string
@@ -69,8 +71,13 @@ func EncryptDataPng(f *os.File, fin *os.File, fout *os.File) {
 	check(err)
 	inputData, err = encryptData(inputData)
 	check(err)
-	cryptChunk := CreateChunk(inputData, chunkName)
-	png.AddMetaChunk(cryptChunk)
+	chunkCount := int(math.Ceil(float64(len(inputData)) / chunkSize))
+	for i := 0; i < chunkCount; i++ {
+		dataStart := i * chunkSize
+		dataEnd := dataStart + int(math.Min(chunkSize, float64(len(inputData[dataStart:]))))
+		cryptChunk := CreateChunk(inputData[dataStart:dataEnd], chunkName)
+		png.AddMetaChunk(cryptChunk)
+	}
 	err = png.Write(fout)
 	check(err)
 }
@@ -80,9 +87,12 @@ func DecryptDataPng(f *os.File, fout *os.File) {
 	png := PngData{}
 	err := png.Read(f)
 	check(err)
-	cryptChunk := png.GetChunk(chunkName)
-	if cryptChunk != nil {
-		data, err := decryptData(cryptChunk.data)
+	var data []byte
+	for _, cryptChunk := range png.GetChunksByName(chunkName) {
+		data = append(data, cryptChunk.data...)
+	}
+	if len(data) > 0 {
+		data, err = decryptData(data)
 		if err != nil {
 			log.Println("\nThe provided password is probably incorrect.")
 		}
diff --git a/pngUtils.go b/pngUtils.go
index d9e27a1..51844f8 100644
--- a/pngUtils.go
+++ b/pngUtils.go
@@ -110,6 +110,17 @@ func (p *PngData) GetChunk(name string) *ChunkData {
 	return nil
 }
 
+// returns all chunks with a given name
+func (p *PngData) GetChunksByName(name string) []ChunkData {
+	var chunks []ChunkData
+	for _, chunk := range p.chunks {
+		if chunk.name == name {
+			chunks = append(chunks, chunk)
+		}
+	}
+	return chunks
+}
+
 // validates the png by reading the header of the file
 func ValidatePng(f *os.File) (bool, []byte) {
 	headerBytes := make([]byte, 8)