viper-app/pkg/md5/md5.go

package md5

import (
	"crypto/md5"
	"fmt"
	"io"
)

func Sum(r io.Reader, size, blockSize, chunkSize int64) (string, error) {
	if size > chunkSize {
		// 分段计算
		return partitionSum(r, size, blockSize)
	}

	data, err := io.ReadAll(r)
	if err != nil {
		return "", err
	}

	fileMD5 := fmt.Sprintf("%x", md5.Sum(data))
	return fileMD5, nil
}

func partitionSum(r io.Reader, size, blockSize int64) (string, error) {
	buf := make([]byte, blockSize)
	hash := md5.New()
	var err error
	for size > blockSize {
		size -= blockSize
		if _, err = r.Read(buf); err != nil {
			return "", err
		}
		hash.Write(buf)
	}

	buf = buf[:size]
	if _, err = r.Read(buf); err != nil {
		return "", err
	}
	hash.Write(buf)
	fileMD5 := fmt.Sprintf("%x", hash.Sum(nil))
	return fileMD5, nil
}

func SumWithHandler(r io.Reader, size, blockSize, chunkSize int64, h func(buf []byte) error) (string, error) {
	if size > chunkSize {
		// 分段计算
		return partitionSumWithHandler(r, size, blockSize, h)
	}

	data, err := io.ReadAll(r)
	if err != nil {
		return "", err
	}
	h(data)
	fileMD5 := fmt.Sprintf("%x", md5.Sum(data))
	return fileMD5, nil
}

func partitionSumWithHandler(r io.Reader, size, blockSize int64, h func(buf []byte) error) (string, error) {
	buf := make([]byte, blockSize)
	hash := md5.New()
	var err error
	for size > blockSize {
		size -= blockSize
		if _, err = r.Read(buf); err != nil {
			return "", err
		}
		hash.Write(buf)
		if err = h(buf); err != nil {
			return "", err
		}
	}

	buf = buf[:size]
	if _, err = r.Read(buf); err != nil {
		return "", err
	}
	hash.Write(buf)
	h(buf)
	fileMD5 := fmt.Sprintf("%x", hash.Sum(nil))
	return fileMD5, nil
}