Commit 62179329 by ale

fix ReadRecord to avoid short reads

1 parent 562f0457
Showing with 11 additions and 4 deletions
......@@ -4,8 +4,8 @@ import (
"bytes"
"encoding/binary"
"encoding/gob"
"errors"
"flag"
"fmt"
"io"
"log"
"os"
......@@ -19,6 +19,7 @@ var (
configFile = flag.String("config", "/etc/djrandom/server.conf", "Config file location")
dump = flag.Bool("dump", false, "Dump the database")
restore = flag.Bool("restore", false, "Restore a database dump")
verbose = flag.Bool("verbose", false, "Print statistics on the dump")
)
// List of database buckets that we should dump. Other data can be
......@@ -47,12 +48,12 @@ func ReadRecord(in io.Reader) (*Record, error) {
}
data := make([]byte, sz)
n, err := in.Read(data)
n, err := io.ReadFull(in, data)
if err != nil {
return nil, err
}
if int64(n) != sz {
return nil, errors.New("short read")
return nil, fmt.Errorf("short read (got %d, want %d)", n, sz)
}
var record Record
......@@ -68,8 +69,8 @@ func WriteRecord(out io.Writer, record Record) error {
if err := gob.NewEncoder(&buf).Encode(&record); err != nil {
return err
}
sz := int64(buf.Len())
enc := buf.Bytes()
sz := int64(len(enc))
binary.Write(out, binary.LittleEndian, sz)
_, err := out.Write(enc)
......@@ -85,6 +86,7 @@ func scanner(db *db_client.DatabaseRpcClient, bucketCh chan string, outCh chan R
for bucket := range bucketCh {
// Call Scan() in batches, to avoid overloading the server.
total := 0
batchSize := 1000
startKey, endKey := "", "\xff"
for startKey < endKey {
......@@ -109,6 +111,7 @@ func scanner(db *db_client.DatabaseRpcClient, bucketCh chan string, outCh chan R
Data: data,
}
n++
total++
startKey = key
}
// An empty result means we've reached the end
......@@ -119,6 +122,9 @@ func scanner(db *db_client.DatabaseRpcClient, bucketCh chan string, outCh chan R
// Start the next scan from the next key.
startKey = db_client.IncrementKey(startKey)
}
if *verbose {
log.Printf("bucket %s: %d entries", bucket, total)
}
}
}
......@@ -198,6 +204,7 @@ func doRestore(db *db_client.DatabaseRpcClient) {
}
func main() {
log.SetFlags(0)
flag.Parse()
config.Parse(*configFile)
......
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!