calculate checksums
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

149 lines
2.7 KiB

3 years ago
  1. package main
  2. import (
  3. "crypto/md5"
  4. "crypto/sha1"
  5. "crypto/sha256"
  6. "crypto/sha512"
  7. "fmt"
  8. "hash"
  9. "io"
  10. "os"
  11. "sort"
  12. "sync"
  13. )
  14. // result is a message or error payload
  15. type result struct {
  16. f string
  17. cs string
  18. err error
  19. }
  20. // results exists to sort a slice of result
  21. type results []result
  22. func (r results) Len() int { return len(r) }
  23. func (r results) Swap(i, j int) { r[i], r[j] = r[j], r[i] }
  24. func (r results) Less(i, j int) bool { return r[i].f < r[j].f }
  25. // hashr exists so that we can make a thing that can return valid hash.Hash
  26. // interfaces.
  27. type hashr func() hash.Hash
  28. // hsh figures out which hash algo to use, and distributes the work of hashing
  29. func hsh(files []string, verbose bool) chan result {
  30. var h hashr
  31. switch *algo {
  32. case "sha1", "1":
  33. h = sha1.New
  34. case "sha256", "256":
  35. h = sha256.New
  36. case "sha512", "512":
  37. h = sha512.New
  38. case "md5":
  39. h = md5.New
  40. default:
  41. r := make(chan result)
  42. go func() {
  43. r <- result{err: fmt.Errorf("unsupported algorithm: %v (supported: md5, sha1, sha256, sha512)", *algo)}
  44. close(r)
  45. }()
  46. return r
  47. }
  48. if len(files) == 0 {
  49. r := make(chan result)
  50. go func() {
  51. hsh := h()
  52. _, err := io.Copy(hsh, os.Stdin)
  53. if err != nil {
  54. fmt.Fprintf(os.Stderr, "%v\n", err)
  55. os.Exit(1)
  56. }
  57. r <- result{cs: fmt.Sprintf("%x", hsh.Sum(nil)), f: "-"}
  58. close(r)
  59. }()
  60. return r
  61. }
  62. jobs := make(chan checksum)
  63. go func() {
  64. for _, name := range files {
  65. jobs <- checksum{filename: name}
  66. }
  67. close(jobs)
  68. }()
  69. res := []<-chan result{}
  70. for w := 0; w < *ngo; w++ {
  71. res = append(res, compute(h, jobs, verbose))
  72. }
  73. o := make(chan result)
  74. go func() {
  75. rs := results{}
  76. for r := range rmerge(res) {
  77. rs = append(rs, r)
  78. }
  79. sort.Sort(rs)
  80. for _, r := range rs {
  81. o <- r
  82. }
  83. close(o)
  84. }()
  85. return o
  86. }
  87. // compute is the checksumming workhorse
  88. func compute(h hashr, jobs chan checksum, verbose bool) chan result {
  89. hsh := h()
  90. r := make(chan result)
  91. go func() {
  92. for job := range jobs {
  93. f, err := os.Open(job.filename)
  94. if err != nil {
  95. r <- result{err: err}
  96. continue
  97. }
  98. hsh.Reset()
  99. _, err = io.Copy(hsh, f)
  100. f.Close()
  101. if err != nil {
  102. r <- result{err: err}
  103. continue
  104. }
  105. if verbose {
  106. fmt.Fprintf(os.Stderr, "%v\n", job.filename)
  107. }
  108. r <- result{f: job.filename, cs: fmt.Sprintf("%x", hsh.Sum(nil))}
  109. }
  110. close(r)
  111. }()
  112. return r
  113. }
  114. // rmerge implements fan-in
  115. func rmerge(cs []<-chan result) chan result {
  116. out := make(chan result)
  117. var wg sync.WaitGroup
  118. output := func(c <-chan result) {
  119. for n := range c {
  120. out <- n
  121. }
  122. wg.Done()
  123. }
  124. wg.Add(len(cs))
  125. for _, c := range cs {
  126. go output(c)
  127. }
  128. go func() {
  129. wg.Wait()
  130. close(out)
  131. }()
  132. return out
  133. }