SynchRoGazer/SynchRoGazer.go

package main

import (
	"bufio"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/minio/md5-simd"
	"github.com/minio/sha256-simd"
	"hash"
	"io"
	"os"
	"runtime"
	"sync/atomic"
)

type HashFileResult struct {
	Error  error
	Path   string
	SHA256 string
	MD5    string
}

func HashFile(results chan<- HashFileResult, md5hasher *md5simd.Hasher, sha256hasher *hash.Hash, path string) {
	fh, err := os.Open(path)
	if err != nil {
		results <- HashFileResult{
			Error: err,
			Path:  path,
		}
		return
	}
	defer fh.Close()

	io.Copy(io.MultiWriter(*sha256hasher, *md5hasher), fh)

	results <- HashFileResult{
		Error:  nil,
		Path:   path,
		SHA256: hex.EncodeToString((*sha256hasher).Sum(nil)),
		MD5:    hex.EncodeToString((*md5hasher).Sum(nil)),
	}
}

func PrintHashFileResult(result *HashFileResult, format string) {
	switch {
	case format == "json":
		var jsonData []byte
		jsonData, err := json.Marshal(*result)
		if err != nil {
			fmt.Fprintln(os.Stderr, result.Path, "Error: ", result.Error)
		} else {
			if result.Error != nil {
				fmt.Fprintln(os.Stderr, string(jsonData))
			} else {
				fmt.Println(string(jsonData))
			}

		}
	case format == "text":
		if result.Error != nil {
			fmt.Fprintln(os.Stderr, result.Path, "Error: ", result.Error)
		} else {
			fmt.Println(result.SHA256, result.MD5, result.Path)
		}

	}
}

func main() {
	taskLimit := flag.Int("tasklimit", func() int {
		result := runtime.NumCPU() * 16
		if result > 128 {
			return 128
		}
		return result
	}(), "Maximum number of concurrent hashing tasks. Change to avoid fdlimit issues. Defaults to number of min(128, CPU cores * 16)")

	outputFormat := flag.String("format", "text", "Output formats. Allowed: text, json")

	flag.Parse()

	var taskCount int64

	scanner := bufio.NewScanner(os.Stdin)
	defer os.Stdin.Close()

	var md5servers []md5simd.Server
	md5hashers := make(chan md5simd.Hasher, *taskLimit)
	sha256hashers := make(chan hash.Hash, *taskLimit)

	for j := 0; j < *taskLimit; j++ {
		serverIndex := j / 16

		if (serverIndex + 1) > len(md5servers) {
			md5servers = append(md5servers, md5simd.NewServer())
		}

		hasher := md5servers[serverIndex].NewHash()
		md5hashers <- hasher
		sha256hashers <- sha256.New()
	}

	resultChannel := make(chan HashFileResult)

	atomic.AddInt64(&taskCount, 1)
	go func() {
		defer atomic.AddInt64(&taskCount, -1)
		for scanner.Scan() {
			path := scanner.Text()

			atomic.AddInt64(&taskCount, 1)
			go func() {
				defer atomic.AddInt64(&taskCount, -1)
				md5hasher := <-md5hashers
				sha256hasher := <-sha256hashers
				md5hasher.Reset()
				sha256hasher.Reset()
				HashFile(resultChannel, &md5hasher, &sha256hasher, path)
				md5hashers <- md5hasher
				sha256hashers <- sha256hasher
			}()
		}

	}()

	//Already print before finishing, use atomic ints instead of a WaitGroup
	for atomic.LoadInt64(&taskCount) > 0 {
		result := <-resultChannel
		PrintHashFileResult(&result, *outputFormat)
	}

	close(resultChannel)

	for result := range resultChannel {
		PrintHashFileResult(&result, *outputFormat)
	}

	close(md5hashers)
	close(sha256hashers)

	for md5hasher := range md5hashers {
		md5hasher.Close()
	}

	for _, md5server := range md5servers {
		md5server.Close()
	}
}