cs/main.go

150 lines
2.3 KiB
Go
Raw Normal View History

2015-12-08 10:18:42 -08:00
package main
import (
2016-01-03 20:15:32 -08:00
"crypto/md5"
2015-12-08 10:18:42 -08:00
"crypto/sha1"
"crypto/sha256"
"crypto/sha512"
"flag"
"fmt"
2016-11-16 00:52:39 -08:00
"hash"
2015-12-08 10:18:42 -08:00
"io"
"os"
"runtime"
2016-11-16 00:52:39 -08:00
"sync"
2015-12-08 10:18:42 -08:00
)
2015-12-10 14:46:06 -08:00
var algo = flag.String("a", "sha1", "algorithm to use")
2016-11-15 20:24:15 -08:00
var mode = flag.Bool("c", false, "check")
var ngo = flag.Int("n", runtime.NumCPU(), "number of goroutines")
2015-12-08 10:18:42 -08:00
func main() {
flag.Parse()
files := flag.Args()
2016-11-15 20:24:15 -08:00
switch *mode {
case true:
2016-11-15 20:50:04 -08:00
c := 0
for err := range check(files) {
c++
2016-11-15 20:24:15 -08:00
fmt.Fprintf(os.Stderr, "%v\n", err)
2016-11-15 20:50:04 -08:00
}
if c > 0 {
2016-11-15 20:24:15 -08:00
os.Exit(1)
}
case false:
2016-11-16 00:52:39 -08:00
c := 0
for res := range hsh(files) {
if res.err != nil {
c++
fmt.Fprintf(os.Stderr, "%v\n", res.err)
} else {
fmt.Printf("%v\n", res.msg)
}
}
if c > 0 {
2016-11-15 20:24:15 -08:00
os.Exit(1)
}
}
}
2016-11-16 00:52:39 -08:00
type hashr func() hash.Hash
func hsh(files []string) chan result {
var h hashr
2015-12-08 10:18:42 -08:00
switch *algo {
case "sha1", "1":
2016-11-16 00:52:39 -08:00
h = sha1.New
2015-12-08 10:18:42 -08:00
case "sha256", "256":
2016-11-16 00:52:39 -08:00
h = sha256.New
2015-12-08 10:18:42 -08:00
case "sha512", "512":
2016-11-16 00:52:39 -08:00
h = sha512.New
2016-01-03 20:15:32 -08:00
case "md5":
2016-11-16 00:52:39 -08:00
h = md5.New
2015-12-08 10:18:42 -08:00
default:
2016-11-16 00:52:39 -08:00
r := make(chan result)
go func() {
r <- result{err: fmt.Errorf("unsupported algorithm: %v", *algo)}
}()
return r
2015-12-08 10:18:42 -08:00
}
if len(files) == 0 {
2016-11-16 00:52:39 -08:00
hsh := h()
_, err := io.Copy(hsh, os.Stdin)
2015-12-08 10:18:42 -08:00
if err != nil {
fmt.Fprintf(os.Stderr, "%v\n", err)
os.Exit(1)
}
2016-11-16 00:52:39 -08:00
fmt.Printf("%x -\n", hsh.Sum(nil))
return nil
}
jobs := make(chan work)
go func() {
2015-12-08 10:18:42 -08:00
for _, name := range files {
2016-11-16 00:52:39 -08:00
jobs <- work{cs: checksum{filename: name}}
}
close(jobs)
}()
res := []<-chan result{}
for w := 0; w < *ngo; w++ {
res = append(res, compute(h, jobs))
}
return rmerge(res)
}
type result struct {
msg string
err error
}
func compute(h hashr, jobs chan work) chan result {
hsh := h()
r := make(chan result)
go func() {
for job := range jobs {
f, err := os.Open(job.cs.filename)
2015-12-08 10:18:42 -08:00
if err != nil {
2016-11-16 00:52:39 -08:00
r <- result{err: err}
2015-12-08 10:18:42 -08:00
continue
}
2016-11-16 00:52:39 -08:00
hsh.Reset()
_, err = io.Copy(hsh, f)
2015-12-08 10:18:42 -08:00
f.Close()
if err != nil {
2016-11-16 00:52:39 -08:00
r <- result{err: err}
2015-12-08 10:18:42 -08:00
continue
}
2016-11-16 00:52:39 -08:00
r <- result{msg: fmt.Sprintf("%x %s", hsh.Sum(nil), job.cs.filename)}
}
close(r)
}()
return r
}
func rmerge(cs []<-chan result) chan result {
out := make(chan result)
var wg sync.WaitGroup
output := func(c <-chan result) {
for n := range c {
out <- n
2015-12-08 10:18:42 -08:00
}
2016-11-16 00:52:39 -08:00
wg.Done()
2015-12-08 10:18:42 -08:00
}
2016-11-16 00:52:39 -08:00
wg.Add(len(cs))
for _, c := range cs {
go output(c)
}
go func() {
wg.Wait()
close(out)
}()
return out
2015-12-08 10:18:42 -08:00
}