fix ReadRecord to avoid short reads

services/database/backuptool/backuptool.go

@@ -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)