utils/checksum.go at main · distbuild/utils · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
package utils

import (
	"crypto/sha256"
	"encoding/hex"
	"io"
	"os"

	"github.com/pkg/errors"
)

const (
	size1K   = 1024
	size1M   = 1024 * 1024
	size20M  = 20 * 1024 * 1024
	size200M = 200 * 1024 * 1024

	seekSet = 0
)

func Checksum(name string) (string, error) {
	var size int64
	var sha string
	var err error

	size, err = getFileSize(name)
	if err != nil {
		return "", errors.Wrap(err, "failed to get file size\n")
	}

	if size < size20M {
		sha, err = checksumSmallFile(name)
	} else if size < size200M {
		sha, err = checksumMediumFile(name)
	} else {
		sha, err = checksumLargeFile(name)
	}

	if err != nil {
		return "", errors.Wrap(err, "failed to get file checksum\n")
	}

	return sha, nil
}

func getFileSize(name string) (int64, error) {
	info, err := os.Stat(name)
	if err != nil {
		return 0, errors.Wrap(err, "failed to stat file\n")
	}

	return info.Size(), nil
}

func checksumSmallFile(name string) (string, error) {
	file, err := os.Open(name)
	if err != nil {
		return "", errors.Wrap(err, "failed to open file\n")
	}

	defer func(file *os.File) {
		_ = file.Close()
	}(file)

	hash := sha256.New()
	if _, err := io.Copy(hash, file); err != nil {
		return "", errors.Wrap(err, "failed to hash file\n")
	}

	return hex.EncodeToString(hash.Sum(nil)), nil
}

func checksumMediumFile(name string) (string, error) {
	file, err := os.Open(name)
	if err != nil {
		return "", errors.Wrap(err, "failed to open file\n")
	}

	defer func(file *os.File) {
		_ = file.Close()
	}(file)

	chunkSize := int64(size1M)
	hash := sha256.New()

	for {
		buffer := make([]byte, chunkSize)
		n, err := file.Read(buffer)
		if err != nil {
			if err == io.EOF {
				_, _ = hash.Write(buffer[:n])
				break
			}
			return "", errors.Wrap(err, "failed to read file\n")
		}
		_, _ = hash.Write(buffer[:n])
	}

	return hex.EncodeToString(hash.Sum(nil)), nil
}

func checksumLargeFile(name string) (string, error) {
	file, err := os.Open(name)
	if err != nil {
		return "", errors.Wrap(err, "failed to open file\n")
	}

	defer func(file *os.File) {
		_ = file.Close()
	}(file)

	startBuffer := make([]byte, size1K)
	if _, err := io.ReadFull(file, startBuffer); err != nil {
		return "", errors.Wrap(err, "failed to read file\n")
	}

	hash := sha256.New()
	if _, err := hash.Write(startBuffer); err != nil {
		return "", errors.Wrap(err, "failed to hash file\n")
	}

	fileStat, err := file.Stat()
	if err != nil {
		return "", errors.Wrap(err, "failed to stat file\n")
	}

	if _, err := file.Seek(fileStat.Size()-int64(size1K), seekSet); err != nil {
		return "", errors.Wrap(err, "failed to seek file\n")
	}

	endBuffer := make([]byte, size1K)
	if _, err := io.ReadFull(file, endBuffer); err != nil {
		return "", errors.Wrap(err, "failed to read file\n")
	}

	if _, err := hash.Write(endBuffer); err != nil {
		return "", errors.Wrap(err, "failed to hash file\n")
	}

	return hex.EncodeToString(hash.Sum(nil)), nil
}