From 09a397ec34b0a86e74bc8435dacb87dce081cef4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E5=85=B5=E5=85=B5?= <liubingbing@qimai.com>
Date: Thu, 27 Aug 2020 13:31:04 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=BB=BB=E5=8A=A1=E8=BF=87?=
 =?UTF-8?q?=E5=A4=9A=E5=AF=BC=E5=85=A5=E5=A4=B1=E8=B4=A5=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 bin/csctl/cmd/restore.go | 33 +++++++++++++++++++++++----------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/bin/csctl/cmd/restore.go b/bin/csctl/cmd/restore.go
index 3218994..570e80a 100644
--- a/bin/csctl/cmd/restore.go
+++ b/bin/csctl/cmd/restore.go
@@ -71,10 +71,23 @@ var (
 	valBuf    = make([]byte, 1024)
 )
 
+func fixRead(r io.Reader, p []byte) (int, error)  {
+	valLen, readLen := len(p), 0
+	for readLen != valLen {
+		n, err := r.Read(p[readLen:])
+		readLen += n
+		if err != nil {
+			return readLen, err
+		}
+	}
+	return readLen, nil
+}
+
+
 func restoreKvs(r io.Reader, keyPrefix string, storeChan chan *kv, wg *sync.WaitGroup) error {
 	for {
 		// read length of key
-		n, err := r.Read(keyLenBuf)
+		n, err := fixRead(r, keyLenBuf)
 		if err == io.EOF && n != 0 {
 			return fmt.Errorf("unexcepted data, the file may broken")
 		} else if err == io.EOF && n == 0 {
@@ -82,28 +95,28 @@ func restoreKvs(r io.Reader, keyPrefix string, storeChan chan *kv, wg *sync.Wait
 		} else if err != nil {
 			return err
 		}
-		keylen := binary.LittleEndian.Uint16(keyLenBuf)
+		keyLen := binary.LittleEndian.Uint16(keyLenBuf)
 
 		// read key
-		if n, err = r.Read(keyBuf[:keylen]); err != nil {
+		if n, err = fixRead(r, keyBuf[:keyLen]); err != nil {
 			return err
 		}
-		key := keyBuf[:keylen]
+		key := keyBuf[:keyLen]
 
 		// read length of value
-		if n, err = r.Read(valLenBuf); err != nil {
+		if n, err = fixRead(r, valLenBuf); err != nil {
 			return err
 		}
-		vallen := binary.LittleEndian.Uint32(valLenBuf)
+		valLen := binary.LittleEndian.Uint32(valLenBuf)
 
 		// read value
-		if len(valBuf) < int(vallen) {
-			valBuf = make([]byte, vallen*2)
+		if len(valBuf) < int(valLen) {
+			valBuf = make([]byte, valLen*2)
 		}
-		if n, err = r.Read(valBuf[:vallen]); err != nil && err != io.EOF {
+		if n, err = fixRead(r, valBuf[:valLen]); err != nil && err != io.EOF {
 			return err
 		}
-		value := valBuf[:vallen]
+		value := valBuf[:valLen]
 
 		wg.Add(1)
 		storeChan <- &kv{