From 3369fa75a36a09756177a386f1d2336edc9e9b50 Mon Sep 17 00:00:00 2001
From: Sergey Matveev <stargrave@stargrave.org>
Date: Sun, 4 Jun 2023 12:58:29 +0300
Subject: [PATCH] Faster and parallalizable integrity checking

---
 integrity.go | 151 +++++++++++++++++++++++++++++++--------------------
 main.go      |   2 +-
 2 files changed, 94 insertions(+), 59 deletions(-)

diff --git a/integrity.go b/integrity.go
index 26fe6a4..3f9455b 100644
--- a/integrity.go
+++ b/integrity.go
@@ -21,88 +21,123 @@ import (
 	"bufio"
 	"bytes"
 	"crypto/sha256"
+	"errors"
 	"fmt"
-	"hash"
 	"io"
+	"io/fs"
 	"log"
 	"os"
 	"path/filepath"
+	"runtime"
 	"strings"
-
-	"golang.org/x/crypto/blake2b"
+	"sync"
 )
 
-func checkFile(
-	pkgName, fn, fnHash, hasherName string,
-	hasher hash.Hash, digest []byte,
-) bool {
-	expected, err := os.ReadFile(fnHash)
-	if err != nil {
-		log.Fatal(err)
-	}
-	fd, err := os.Open(fn)
-	if err != nil {
-		if os.IsNotExist(err) {
-			return true
+func checker(jobs chan string, isGood *bool, workers *sync.WaitGroup) {
+	hasherBLAKE2b256 := blake2b256New()
+	hasherSHA256 := sha256.New()
+	digestBuf := make([]byte, 32)
+	defer workers.Done()
+	var fd *os.File
+	var br bufio.Reader
+	var err error
+	var digest []byte
+	for fn := range jobs {
+		fd, err = os.Open(fn)
+		if err != nil {
+			fmt.Println("ERR", fn, err)
+			*isGood = false
+			continue
+		}
+
+		digest, err = os.ReadFile(fn + "." + HashAlgoBLAKE2b256)
+		if err != nil {
+			if !errors.Is(err, fs.ErrNotExist) {
+				fmt.Println("ERR", fn+"."+HashAlgoBLAKE2b256, err)
+				*isGood = false
+			}
+			goto NoBLAKE2b256
+		}
+		hasherBLAKE2b256.Reset()
+		br.Reset(fd)
+		_, err = io.Copy(hasherBLAKE2b256, &br)
+		fd.Seek(0, io.SeekStart)
+		if err != nil {
+			fmt.Println("ERR", fn+"."+HashAlgoBLAKE2b256, err)
+			*isGood = false
+			goto NoBLAKE2b256
+		}
+		if bytes.Equal(hasherBLAKE2b256.Sum(digestBuf[:0]), digest) {
+			fmt.Println("GOOD", fn+"."+HashAlgoBLAKE2b256)
+		} else {
+			fmt.Println("BAD", fn+"."+HashAlgoBLAKE2b256)
+			*isGood = false
+		}
+
+	NoBLAKE2b256:
+		digest, err = os.ReadFile(fn + "." + HashAlgoSHA256)
+		if err != nil {
+			if !errors.Is(err, fs.ErrNotExist) {
+				fmt.Println("ERR", fn+"."+HashAlgoSHA256, err)
+				*isGood = false
+			}
+			fd.Close()
+			continue
+		}
+		hasherSHA256.Reset()
+		br.Reset(fd)
+		_, err = io.Copy(hasherSHA256, &br)
+		fd.Close()
+		if err != nil {
+			fmt.Println("ERR", fn+"."+HashAlgoBLAKE2b256, err)
+			*isGood = false
+			continue
+		}
+		if bytes.Equal(hasherSHA256.Sum(digestBuf[:0]), digest) {
+			fmt.Println("GOOD", fn+"."+HashAlgoSHA256)
+		} else {
+			fmt.Println("BAD", fn+"."+HashAlgoSHA256)
+			*isGood = false
 		}
-		log.Fatal(err)
-	}
-	_, err = io.Copy(hasher, bufio.NewReader(fd))
-	fd.Close()
-	if err != nil {
-		log.Fatal(err)
-	}
-	isEqual := bytes.Equal(hasher.Sum(digest[:0]), expected)
-	hasher.Reset()
-	if isEqual {
-		fmt.Println("GOOD", hasherName, pkgName)
-		return true
 	}
-	fmt.Println("BAD", hasherName, pkgName)
-	return false
 }
 
 func goodIntegrity() bool {
-	dirs, err := os.ReadDir(Root)
+	isGood := true
+	jobs := make(chan string, 1<<10)
+	var workers sync.WaitGroup
+	workers.Add(runtime.NumCPU())
+	for i := 0; i < runtime.NumCPU(); i++ {
+		go checker(jobs, &isGood, &workers)
+	}
+	fd, err := os.Open(Root)
+	if err != nil {
+		log.Fatal(err)
+	}
+	dirs, err := fd.ReadDir(0)
+	fd.Close()
 	if err != nil {
 		log.Fatal(err)
 	}
-	hasherSHA256 := sha256.New()
-	hasherBLAKE2b256 := blake2b256New()
-	digestSHA256 := make([]byte, sha256.Size)
-	digestBLAKE2b256 := make([]byte, blake2b.Size256)
-	isGood := true
-	var pkgName string
 	for _, dir := range dirs {
-		files, err := os.ReadDir(filepath.Join(Root, dir.Name()))
+		fd, err = os.Open(filepath.Join(Root, dir.Name()))
+		if err != nil {
+			log.Fatal(err)
+		}
+		files, err := fd.ReadDir(0)
+		fd.Close()
 		if err != nil {
 			log.Fatal(err)
 		}
 		for _, file := range files {
-			if strings.HasSuffix(file.Name(), "."+HashAlgoSHA256) {
-				pkgName = strings.TrimSuffix(file.Name(), "."+HashAlgoSHA256)
-				if !checkFile(
-					pkgName,
-					filepath.Join(Root, dir.Name(), pkgName),
-					filepath.Join(Root, dir.Name(), file.Name()),
-					"SHA256", hasherSHA256, digestSHA256,
-				) {
-					isGood = false
-				}
-				continue
-			}
-			if strings.HasSuffix(file.Name(), "."+HashAlgoBLAKE2b256) {
-				pkgName = strings.TrimSuffix(file.Name(), "."+HashAlgoBLAKE2b256)
-				if !checkFile(
-					pkgName,
-					filepath.Join(Root, dir.Name(), pkgName),
-					filepath.Join(Root, dir.Name(), file.Name()),
-					"BLAKE2b-256", hasherBLAKE2b256, digestBLAKE2b256,
-				) {
-					isGood = false
+			for _, ext := range KnownExts {
+				if strings.HasSuffix(file.Name(), ext) {
+					jobs <- filepath.Join(Root, dir.Name(), file.Name())
 				}
 			}
 		}
 	}
+	close(jobs)
+	workers.Wait()
 	return isGood
 }
diff --git a/main.go b/main.go
index c3342c7..f0ef969 100644
--- a/main.go
+++ b/main.go
@@ -44,7 +44,7 @@ import (
 )
 
 const (
-	Version   = "4.0.0"
+	Version   = "4.1.0"
 	UserAgent = "GoCheese/" + Version
 )
 
-- 
2.44.0