123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167 |
- package main
- import (
- "flag"
- "fmt"
- "io"
- "os"
- "path/filepath"
- "github.com/klauspost/reedsolomon"
- )
- var dataShards = flag.Int("data", 4, "Number of shards to split the data into")
- var parShards = flag.Int("par", 2, "Number of parity shards")
- var outFile = flag.String("out", "", "Alternative output path/file")
- func init() {
- flag.Usage = func() {
- fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
- fmt.Fprintf(os.Stderr, " %s [-flags] basefile.ext\nDo not add the number to the filename.\n", os.Args[0])
- fmt.Fprintf(os.Stderr, "Valid flags:\n")
- flag.PrintDefaults()
- }
- }
- func main() {
-
- flag.Parse()
- args := flag.Args()
- if len(args) != 1 {
- fmt.Fprintf(os.Stderr, "Error: No filenames given\n")
- flag.Usage()
- os.Exit(1)
- }
- fname := args[0]
-
- enc, err := reedsolomon.NewStream(*dataShards, *parShards)
- checkErr(err)
-
- shards, size, err := openInput(*dataShards, *parShards, fname)
- checkErr(err)
-
- ok, err := enc.Verify(shards)
- if ok {
- fmt.Println("No reconstruction needed")
- } else {
- fmt.Println("Verification failed. Reconstructing data")
- shards, size, err = openInput(*dataShards, *parShards, fname)
- checkErr(err)
-
- out := make([]io.Writer, len(shards))
- for i := range out {
- if shards[i] == nil {
- dir, _ := filepath.Split(fname)
- outfn := fmt.Sprintf("%s.%d", fname, i)
- fmt.Println("Creating", outfn)
- out[i], err = os.Create(filepath.Join(dir, outfn))
- checkErr(err)
- }
- }
- err = enc.Reconstruct(shards, out)
- if err != nil {
- fmt.Println("Reconstruct failed -", err)
- os.Exit(1)
- }
-
- for i := range out {
- if out[i] != nil {
- err := out[i].(*os.File).Close()
- checkErr(err)
- }
- }
- shards, size, err = openInput(*dataShards, *parShards, fname)
- ok, err = enc.Verify(shards)
- if !ok {
- fmt.Println("Verification failed after reconstruction, data likely corrupted:", err)
- os.Exit(1)
- }
- checkErr(err)
- }
-
- outfn := *outFile
- if outfn == "" {
- outfn = fname
- }
- fmt.Println("Writing data to", outfn)
- f, err := os.Create(outfn)
- checkErr(err)
- shards, size, err = openInput(*dataShards, *parShards, fname)
- checkErr(err)
-
- err = enc.Join(f, shards, int64(*dataShards)*size)
- checkErr(err)
- }
- func openInput(dataShards, parShards int, fname string) (r []io.Reader, size int64, err error) {
-
- shards := make([]io.Reader, dataShards+parShards)
- for i := range shards {
- infn := fmt.Sprintf("%s.%d", fname, i)
- fmt.Println("Opening", infn)
- f, err := os.Open(infn)
- if err != nil {
- fmt.Println("Error reading file", err)
- shards[i] = nil
- continue
- } else {
- shards[i] = f
- }
- stat, err := f.Stat()
- checkErr(err)
- if stat.Size() > 0 {
- size = stat.Size()
- } else {
- shards[i] = nil
- }
- }
- return shards, size, nil
- }
- func checkErr(err error) {
- if err != nil {
- fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
- os.Exit(2)
- }
- }
|