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.

124 lines
2.1 KiB

  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. "sync"
  12. )
  13. // result is a message or error payload
  14. type result struct {
  15. msg string
  16. err error
  17. }
  18. // hashr exists so that we can make a thing that can return valid hash.Hash
  19. // interfaces.
  20. type hashr func() hash.Hash
  21. // hsh figures out which hash algo to use, and distributes the work of hashing
  22. func hsh(files []string) chan result {
  23. var h hashr
  24. switch *algo {
  25. case "sha1", "1":
  26. h = sha1.New
  27. case "sha256", "256":
  28. h = sha256.New
  29. case "sha512", "512":
  30. h = sha512.New
  31. case "md5":
  32. h = md5.New
  33. default:
  34. r := make(chan result)
  35. go func() {
  36. r <- result{err: fmt.Errorf("unsupported algorithm: %v", *algo)}
  37. }()
  38. return r
  39. }
  40. if len(files) == 0 {
  41. r := make(chan result)
  42. go func() {
  43. hsh := h()
  44. _, err := io.Copy(hsh, os.Stdin)
  45. if err != nil {
  46. fmt.Fprintf(os.Stderr, "%v\n", err)
  47. os.Exit(1)
  48. }
  49. r <- result{msg: fmt.Sprintf("%x -", hsh.Sum(nil))}
  50. close(r)
  51. }()
  52. return r
  53. }
  54. jobs := make(chan checksum)
  55. go func() {
  56. for _, name := range files {
  57. jobs <- checksum{filename: name}
  58. }
  59. close(jobs)
  60. }()
  61. res := []<-chan result{}
  62. for w := 0; w < *ngo; w++ {
  63. res = append(res, compute(h, jobs))
  64. }
  65. return rmerge(res)
  66. }
  67. // compute is the checksumming workhorse
  68. func compute(h hashr, jobs chan checksum) chan result {
  69. hsh := h()
  70. r := make(chan result)
  71. go func() {
  72. for job := range jobs {
  73. f, err := os.Open(job.filename)
  74. if err != nil {
  75. r <- result{err: err}
  76. continue
  77. }
  78. hsh.Reset()
  79. _, err = io.Copy(hsh, f)
  80. f.Close()
  81. if err != nil {
  82. r <- result{err: err}
  83. continue
  84. }
  85. r <- result{msg: fmt.Sprintf("%x %s", hsh.Sum(nil), job.filename)}
  86. }
  87. close(r)
  88. }()
  89. return r
  90. }
  91. // rmerge implements fan-in
  92. func rmerge(cs []<-chan result) chan result {
  93. out := make(chan result)
  94. var wg sync.WaitGroup
  95. output := func(c <-chan result) {
  96. for n := range c {
  97. out <- n
  98. }
  99. wg.Done()
  100. }
  101. wg.Add(len(cs))
  102. for _, c := range cs {
  103. go output(c)
  104. }
  105. go func() {
  106. wg.Wait()
  107. close(out)
  108. }()
  109. return out
  110. }