From efef7e13ae95a2f272e68fa038abd42dab98b470 Mon Sep 17 00:00:00 2001
From: vinllen <cvinllen@gmail.com>
Date: Wed, 16 Oct 2019 13:24:30 +0800
Subject: [PATCH] add rdb checksum check

---
 src/redis-shake/common/utils.go | 16 ++++++++++++++++
 src/redis-shake/main/main.go    | 22 ++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/src/redis-shake/common/utils.go b/src/redis-shake/common/utils.go
index 6fb9b35..edb21ae 100644
--- a/src/redis-shake/common/utils.go
+++ b/src/redis-shake/common/utils.go
@@ -923,6 +923,22 @@ func GetRedisVersion(target, authType, auth string, tlsEnable bool) (string, err
 	}
 }
 
+func GetRDBChecksum(target, authType, auth string, tlsEnable bool) (string, error) {
+	c := OpenRedisConn([]string{target}, authType, auth, false, tlsEnable)
+	defer c.Close()
+
+	content, err := c.Do("config", "get", "rdbchecksum")
+	if err != nil {
+		return "", err
+	}
+
+	conentList := content.([]interface{})
+	if len(conentList) != 2 {
+		return "", fmt.Errorf("return length != 2, return: %v", conentList)
+	}
+	return string(conentList[1].([]byte)), nil
+}
+
 func CheckHandleNetError(err error) bool {
 	if err == io.EOF {
 		return true
diff --git a/src/redis-shake/main/main.go b/src/redis-shake/main/main.go
index f615f50..115605a 100644
--- a/src/redis-shake/main/main.go
+++ b/src/redis-shake/main/main.go
@@ -434,6 +434,8 @@ func sanitizeOptions(tp string) error {
 			 * set 1 if target is target version can't be fetched just like twemproxy.
 			 */
 			conf.Options.BigKeyThreshold = 1
+			log.Warnf("target version[%v] given, set big_key_threshold = 1. see #173",
+				conf.Options.TargetVersion, conf.Options.SourceVersion)
 		}
 
 		if strings.HasPrefix(conf.Options.TargetVersion, "4.") ||
@@ -463,6 +465,8 @@ func sanitizeOptions(tp string) error {
 		// compare version. see github issue #173.
 		if ret := utils.CompareVersion(conf.Options.SourceVersion, conf.Options.TargetVersion, 2); ret != 0 && ret != 1 {
 			// target version is smaller than source version, or unknown
+			log.Warnf("target version[%v] < source version[%v], set big_key_threshold = 1. see #173",
+				conf.Options.TargetVersion, conf.Options.SourceVersion)
 			conf.Options.BigKeyThreshold = 1
 		}
 	}
@@ -492,6 +496,24 @@ func sanitizeOptions(tp string) error {
 		//}
 	}
 
+	// check rdbchecksum
+	if tp == conf.TypeDump || (tp == conf.TypeSync || tp == conf.TypeRump) && conf.Options.BigKeyThreshold > 1 {
+		for _, address := range conf.Options.SourceAddressList {
+			check, err := utils.GetRDBChecksum(address, conf.Options.SourceAuthType,
+				conf.Options.SourcePasswordRaw, conf.Options.SourceTLSEnable)
+			if err != nil {
+				// ignore
+				log.Warnf("fetch source rdb[%v] checksum failed[%v], ignore", address, err)
+				continue
+			}
+
+			log.Infof("source rdb[%v] checksum[%v]", address, check)
+			if check == "no" {
+				return fmt.Errorf("source rdb[%v] checksum should be open[config set rdbchecksum yes]", address)
+			}
+		}
+	}
+
 	return nil
 }