Refine the codec function

Signed-off-by: JmPotato <[email protected]>

client/http/codec.go

@@ -15,6 +15,8 @@
 package http
 
 import (
+	"encoding/hex"
+
 	"github.com/pingcap/errors"
 )
 
@@ -64,11 +66,11 @@ func encodeBytes(data []byte) []byte {
 	return result
 }
 
-func decodeBytes(b []byte) ([]byte, []byte, error) {
+func decodeBytes(b []byte) ([]byte, error) {
 	buf := make([]byte, 0, len(b))
 	for {
 		if len(b) < encGroupSize+1 {
-			return nil, nil, errors.New("insufficient bytes to decode value")
+			return nil, errors.New("insufficient bytes to decode value")
 		}
 
 		groupBytes := b[:encGroupSize+1]
@@ -78,7 +80,7 @@ func decodeBytes(b []byte) ([]byte, []byte, error) {
 
 		padCount := encMarker - marker
 		if padCount > encGroupSize {
-			return nil, nil, errors.Errorf("invalid marker byte, group bytes %q", groupBytes)
+			return nil, errors.Errorf("invalid marker byte, group bytes %q", groupBytes)
 		}
 
 		realGroupSize := encGroupSize - padCount
@@ -89,11 +91,31 @@ func decodeBytes(b []byte) ([]byte, []byte, error) {
 			// Check validity of padding bytes.
 			for _, v := range group[realGroupSize:] {
 				if v != encPad {
-					return nil, nil, errors.Errorf("invalid padding byte, group bytes %q", groupBytes)
+					return nil, errors.Errorf("invalid padding byte, group bytes %q", groupBytes)
 				}
 			}
 			break
 		}
 	}
-	return b, buf, nil
+	return buf, nil
+}
+
+// keyToKeyHexStr converts a raw key to a hex string after encoding.
+func rawKeyToKeyHexStr(key []byte) string {
+	if len(key) == 0 {
+		return ""
+	}
+	return hex.EncodeToString(encodeBytes(key))
+}
+
+// keyHexStrToRawKey converts a hex string to a raw key after decoding.
+func keyHexStrToRawKey(hexKey string) ([]byte, error) {
+	if len(hexKey) == 0 {
+		return make([]byte, 0), nil
+	}
+	key, err := hex.DecodeString(hexKey)
+	if err != nil {
+		return nil, err
+	}
+	return decodeBytes(key)
 }

client/http/codec_test.go

@@ -39,7 +39,7 @@ func TestBytesCodec(t *testing.T) {
 		b := encodeBytes(input.enc)
 		require.Equal(t, input.dec, b)
 
-		_, d, err := decodeBytes(b)
+		d, err := decodeBytes(b)
 		require.NoError(t, err)
 		require.Equal(t, input.enc, d)
 	}
@@ -58,7 +58,7 @@ func TestBytesCodec(t *testing.T) {
 	}
 
 	for _, input := range errInputs {
-		_, _, err := decodeBytes(input)
+		_, err := decodeBytes(input)
 		require.Error(t, err)
 	}
 }

client/http/types.go

@@ -304,37 +304,79 @@ var (
 	_ json.Unmarshaler = (*Rule)(nil)
 )
 
+// This is a helper struct used to customizing the JSON marshal/unmarshal methods of `Rule`.
+type rule struct {
+	GroupID          string            `json:"group_id"`
+	ID               string            `json:"id"`
+	Index            int               `json:"index,omitempty"`
+	Override         bool              `json:"override,omitempty"`
+	StartKeyHex      string            `json:"start_key"`
+	EndKeyHex        string            `json:"end_key"`
+	Role             PeerRoleType      `json:"role"`
+	IsWitness        bool              `json:"is_witness"`
+	Count            int               `json:"count"`
+	LabelConstraints []LabelConstraint `json:"label_constraints,omitempty"`
+	LocationLabels   []string          `json:"location_labels,omitempty"`
+	IsolationLevel   string            `json:"isolation_level,omitempty"`
+}
+
 // MarshalJSON implements `json.Marshaler` interface to make sure we could set the correct start/end key.
 func (r *Rule) MarshalJSON() ([]byte, error) {
-	r.StartKeyHex = hex.EncodeToString(encodeBytes(r.StartKey))
-	r.EndKeyHex = hex.EncodeToString(encodeBytes(r.EndKey))
-	return json.Marshal(r)
+	tempRule := &rule{
+		GroupID:          r.GroupID,
+		ID:               r.ID,
+		Index:            r.Index,
+		Override:         r.Override,
+		StartKeyHex:      r.StartKeyHex,
+		EndKeyHex:        r.EndKeyHex,
+		Role:             r.Role,
+		IsWitness:        r.IsWitness,
+		Count:            r.Count,
+		LabelConstraints: r.LabelConstraints,
+		LocationLabels:   r.LocationLabels,
+		IsolationLevel:   r.IsolationLevel,
+	}
+	// Converts the start/end key to hex format if the corresponding hex field is empty.
+	if len(r.StartKey) > 0 && len(r.StartKeyHex) == 0 {
+		tempRule.StartKeyHex = rawKeyToKeyHexStr(r.StartKey)
+	}
+	if len(r.EndKey) > 0 && len(r.EndKeyHex) == 0 {
+		tempRule.EndKeyHex = rawKeyToKeyHexStr(r.EndKey)
+	}
+	return json.Marshal(tempRule)
 }
 
 // UnmarshalJSON implements `json.Unmarshaler` interface to make sure we could get the correct start/end key.
 func (r *Rule) UnmarshalJSON(bytes []byte) error {
-	if err := json.Unmarshal(bytes, r); err != nil {
-		return err
-	}
-
-	startKey, err := hex.DecodeString(r.StartKeyHex)
+	var tempRule rule
+	err := json.Unmarshal(bytes, &tempRule)
 	if err != nil {
 		return err
 	}
-
-	endKey, err := hex.DecodeString(r.EndKeyHex)
+	newRule := Rule{
+		GroupID:          tempRule.GroupID,
+		ID:               tempRule.ID,
+		Index:            tempRule.Index,
+		Override:         tempRule.Override,
+		StartKeyHex:      tempRule.StartKeyHex,
+		EndKeyHex:        tempRule.EndKeyHex,
+		Role:             tempRule.Role,
+		IsWitness:        tempRule.IsWitness,
+		Count:            tempRule.Count,
+		LabelConstraints: tempRule.LabelConstraints,
+		LocationLabels:   tempRule.LocationLabels,
+		IsolationLevel:   tempRule.IsolationLevel,
+	}
+	newRule.StartKey, err = keyHexStrToRawKey(newRule.StartKeyHex)
 	if err != nil {
 		return err
 	}
-
-	_, r.StartKey, err = decodeBytes(startKey)
+	newRule.EndKey, err = keyHexStrToRawKey(newRule.EndKeyHex)
 	if err != nil {
 		return err
 	}
-
-	_, r.EndKey, err = decodeBytes(endKey)
-
-	return err
+	*r = newRule
+	return nil
 }
 
 // RuleOpType indicates the operation type
@@ -365,37 +407,87 @@ var (
 	_ json.Unmarshaler = (*RuleOp)(nil)
 )
 
+// This is a helper struct used to customizing the JSON marshal/unmarshal methods of `RuleOp`.
+type ruleOp struct {
+	GroupID          string            `json:"group_id"`
+	ID               string            `json:"id"`
+	Index            int               `json:"index,omitempty"`
+	Override         bool              `json:"override,omitempty"`
+	StartKeyHex      string            `json:"start_key"`
+	EndKeyHex        string            `json:"end_key"`
+	Role             PeerRoleType      `json:"role"`
+	IsWitness        bool              `json:"is_witness"`
+	Count            int               `json:"count"`
+	LabelConstraints []LabelConstraint `json:"label_constraints,omitempty"`
+	LocationLabels   []string          `json:"location_labels,omitempty"`
+	IsolationLevel   string            `json:"isolation_level,omitempty"`
+	Action           RuleOpType        `json:"action"`
+	DeleteByIDPrefix bool              `json:"delete_by_id_prefix"`
+}
+
 // MarshalJSON implements `json.Marshaler` interface to make sure we could set the correct start/end key.
 func (r *RuleOp) MarshalJSON() ([]byte, error) {
-	r.StartKeyHex = hex.EncodeToString(encodeBytes(r.StartKey))
-	r.EndKeyHex = hex.EncodeToString(encodeBytes(r.EndKey))
-	return json.Marshal(r)
+	tempRuleOp := &ruleOp{
+		GroupID:          r.GroupID,
+		ID:               r.ID,
+		Index:            r.Index,
+		Override:         r.Override,
+		StartKeyHex:      r.StartKeyHex,
+		EndKeyHex:        r.EndKeyHex,
+		Role:             r.Role,
+		IsWitness:        r.IsWitness,
+		Count:            r.Count,
+		LabelConstraints: r.LabelConstraints,
+		LocationLabels:   r.LocationLabels,
+		IsolationLevel:   r.IsolationLevel,
+		Action:           r.Action,
+		DeleteByIDPrefix: r.DeleteByIDPrefix,
+	}
+	// Converts the start/end key to hex format if the corresponding hex field is empty.
+	if len(r.StartKey) > 0 && len(r.StartKeyHex) == 0 {
+		tempRuleOp.StartKeyHex = rawKeyToKeyHexStr(r.StartKey)
+	}
+	if len(r.EndKey) > 0 && len(r.EndKeyHex) == 0 {
+		tempRuleOp.EndKeyHex = rawKeyToKeyHexStr(r.EndKey)
+	}
+	return json.Marshal(tempRuleOp)
 }
 
 // UnmarshalJSON implements `json.Unmarshaler` interface to make sure we could get the correct start/end key.
 func (r *RuleOp) UnmarshalJSON(bytes []byte) error {
-	if err := json.Unmarshal(bytes, r); err != nil {
-		return err
-	}
-
-	startKey, err := hex.DecodeString(r.StartKeyHex)
+	var tempRuleOp ruleOp
+	err := json.Unmarshal(bytes, &tempRuleOp)
 	if err != nil {
 		return err
 	}
-
-	endKey, err := hex.DecodeString(r.EndKeyHex)
+	newRuleOp := RuleOp{
+		Rule: &Rule{
+			GroupID:          tempRuleOp.GroupID,
+			ID:               tempRuleOp.ID,
+			Index:            tempRuleOp.Index,
+			Override:         tempRuleOp.Override,
+			StartKeyHex:      tempRuleOp.StartKeyHex,
+			EndKeyHex:        tempRuleOp.EndKeyHex,
+			Role:             tempRuleOp.Role,
+			IsWitness:        tempRuleOp.IsWitness,
+			Count:            tempRuleOp.Count,
+			LabelConstraints: tempRuleOp.LabelConstraints,
+			LocationLabels:   tempRuleOp.LocationLabels,
+			IsolationLevel:   tempRuleOp.IsolationLevel,
+		},
+		Action:           tempRuleOp.Action,
+		DeleteByIDPrefix: tempRuleOp.DeleteByIDPrefix,
+	}
+	newRuleOp.StartKey, err = keyHexStrToRawKey(newRuleOp.StartKeyHex)
 	if err != nil {
 		return err
 	}
-
-	_, r.StartKey, err = decodeBytes(startKey)
+	newRuleOp.EndKey, err = keyHexStrToRawKey(newRuleOp.EndKeyHex)
 	if err != nil {
 		return err
 	}
-
-	_, r.EndKey, err = decodeBytes(endKey)
-
-	return err
+	*r = newRuleOp
+	return nil
 }
 
 // RuleGroup defines properties of a rule group.

client/http/types_test.go

@@ -15,6 +15,7 @@
 package http
 
 import (
+	"encoding/json"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -47,3 +48,65 @@ func TestMergeRegionsInfo(t *testing.T) {
 	re.Equal(2, len(regionsInfo.Regions))
 	re.Equal(append(regionsInfo1.Regions, regionsInfo2.Regions...), regionsInfo.Regions)
 }
+
+func TestRuleStartEndKey(t *testing.T) {
+	re := require.New(t)
+	// Empty start/end key and key hex.
+	ruleToMarshal := &Rule{}
+	rule := mustMarshalAndUnmarshal(re, ruleToMarshal)
+	re.Equal("", rule.StartKeyHex)
+	re.Equal("", rule.EndKeyHex)
+	re.Equal([]byte(""), rule.StartKey)
+	re.Equal([]byte(""), rule.EndKey)
+	// Empty start/end key and non-empty key hex.
+	ruleToMarshal = &Rule{
+		StartKeyHex: rawKeyToKeyHexStr([]byte("a")),
+		EndKeyHex:   rawKeyToKeyHexStr([]byte("b")),
+	}
+	rule = mustMarshalAndUnmarshal(re, ruleToMarshal)
+	re.Equal([]byte("a"), rule.StartKey)
+	re.Equal([]byte("b"), rule.EndKey)
+	re.Equal(ruleToMarshal.StartKeyHex, rule.StartKeyHex)
+	re.Equal(ruleToMarshal.EndKeyHex, rule.EndKeyHex)
+	// Non-empty start/end key and empty key hex.
+	ruleToMarshal = &Rule{
+		StartKey: []byte("a"),
+		EndKey:   []byte("b"),
+	}
+	rule = mustMarshalAndUnmarshal(re, ruleToMarshal)
+	re.Equal(ruleToMarshal.StartKey, rule.StartKey)
+	re.Equal(ruleToMarshal.EndKey, rule.EndKey)
+	re.Equal(rawKeyToKeyHexStr(ruleToMarshal.StartKey), rule.StartKeyHex)
+	re.Equal(rawKeyToKeyHexStr(ruleToMarshal.EndKey), rule.EndKeyHex)
+	// Non-empty start/end key and non-empty key hex.
+	ruleToMarshal = &Rule{
+		StartKey:    []byte("a"),
+		EndKey:      []byte("b"),
+		StartKeyHex: rawKeyToKeyHexStr([]byte("c")),
+		EndKeyHex:   rawKeyToKeyHexStr([]byte("d")),
+	}
+	rule = mustMarshalAndUnmarshal(re, ruleToMarshal)
+	re.Equal([]byte("c"), rule.StartKey)
+	re.Equal([]byte("d"), rule.EndKey)
+	re.Equal(ruleToMarshal.StartKeyHex, rule.StartKeyHex)
+	re.Equal(ruleToMarshal.EndKeyHex, rule.EndKeyHex)
+	// Half of each pair of keys is empty.
+	ruleToMarshal = &Rule{
+		StartKey:  []byte("a"),
+		EndKeyHex: rawKeyToKeyHexStr([]byte("d")),
+	}
+	rule = mustMarshalAndUnmarshal(re, ruleToMarshal)
+	re.Equal(ruleToMarshal.StartKey, rule.StartKey)
+	re.Equal([]byte("d"), rule.EndKey)
+	re.Equal(rawKeyToKeyHexStr(ruleToMarshal.StartKey), rule.StartKeyHex)
+	re.Equal(ruleToMarshal.EndKeyHex, rule.EndKeyHex)
+}
+
+func mustMarshalAndUnmarshal(re *require.Assertions, rule *Rule) *Rule {
+	ruleJSON, err := json.Marshal(rule)
+	re.NoError(err)
+	var newRule *Rule
+	err = json.Unmarshal(ruleJSON, &newRule)
+	re.NoError(err)
+	return newRule
+}