Skip to content

Commit 345491d

Browse files
committed
Refine the codec function
Signed-off-by: JmPotato <[email protected]>
1 parent eeaf477 commit 345491d

File tree

5 files changed

+242
-47
lines changed

5 files changed

+242
-47
lines changed

Diff for: client/http/codec.go

+27-5
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
package http
1616

1717
import (
18+
"encoding/hex"
19+
1820
"github.com/pingcap/errors"
1921
)
2022

@@ -64,11 +66,11 @@ func encodeBytes(data []byte) []byte {
6466
return result
6567
}
6668

67-
func decodeBytes(b []byte) ([]byte, []byte, error) {
69+
func decodeBytes(b []byte) ([]byte, error) {
6870
buf := make([]byte, 0, len(b))
6971
for {
7072
if len(b) < encGroupSize+1 {
71-
return nil, nil, errors.New("insufficient bytes to decode value")
73+
return nil, errors.New("insufficient bytes to decode value")
7274
}
7375

7476
groupBytes := b[:encGroupSize+1]
@@ -78,7 +80,7 @@ func decodeBytes(b []byte) ([]byte, []byte, error) {
7880

7981
padCount := encMarker - marker
8082
if padCount > encGroupSize {
81-
return nil, nil, errors.Errorf("invalid marker byte, group bytes %q", groupBytes)
83+
return nil, errors.Errorf("invalid marker byte, group bytes %q", groupBytes)
8284
}
8385

8486
realGroupSize := encGroupSize - padCount
@@ -89,11 +91,31 @@ func decodeBytes(b []byte) ([]byte, []byte, error) {
8991
// Check validity of padding bytes.
9092
for _, v := range group[realGroupSize:] {
9193
if v != encPad {
92-
return nil, nil, errors.Errorf("invalid padding byte, group bytes %q", groupBytes)
94+
return nil, errors.Errorf("invalid padding byte, group bytes %q", groupBytes)
9395
}
9496
}
9597
break
9698
}
9799
}
98-
return b, buf, nil
100+
return buf, nil
101+
}
102+
103+
// keyToKeyHexStr converts a raw key to a hex string after encoding.
104+
func rawKeyToKeyHexStr(key []byte) string {
105+
if len(key) == 0 {
106+
return ""
107+
}
108+
return hex.EncodeToString(encodeBytes(key))
109+
}
110+
111+
// keyHexStrToRawKey converts a hex string to a raw key after decoding.
112+
func keyHexStrToRawKey(hexKey string) ([]byte, error) {
113+
if len(hexKey) == 0 {
114+
return make([]byte, 0), nil
115+
}
116+
key, err := hex.DecodeString(hexKey)
117+
if err != nil {
118+
return nil, err
119+
}
120+
return decodeBytes(key)
99121
}

Diff for: client/http/codec_test.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ func TestBytesCodec(t *testing.T) {
3939
b := encodeBytes(input.enc)
4040
require.Equal(t, input.dec, b)
4141

42-
_, d, err := decodeBytes(b)
42+
d, err := decodeBytes(b)
4343
require.NoError(t, err)
4444
require.Equal(t, input.enc, d)
4545
}
@@ -58,7 +58,7 @@ func TestBytesCodec(t *testing.T) {
5858
}
5959

6060
for _, input := range errInputs {
61-
_, _, err := decodeBytes(input)
61+
_, err := decodeBytes(input)
6262
require.Error(t, err)
6363
}
6464
}

Diff for: client/http/types.go

+124-32
Original file line numberDiff line numberDiff line change
@@ -304,37 +304,79 @@ var (
304304
_ json.Unmarshaler = (*Rule)(nil)
305305
)
306306

307+
// This is a helper struct used to customizing the JSON marshal/unmarshal methods of `Rule`.
308+
type rule struct {
309+
GroupID string `json:"group_id"`
310+
ID string `json:"id"`
311+
Index int `json:"index,omitempty"`
312+
Override bool `json:"override,omitempty"`
313+
StartKeyHex string `json:"start_key"`
314+
EndKeyHex string `json:"end_key"`
315+
Role PeerRoleType `json:"role"`
316+
IsWitness bool `json:"is_witness"`
317+
Count int `json:"count"`
318+
LabelConstraints []LabelConstraint `json:"label_constraints,omitempty"`
319+
LocationLabels []string `json:"location_labels,omitempty"`
320+
IsolationLevel string `json:"isolation_level,omitempty"`
321+
}
322+
307323
// MarshalJSON implements `json.Marshaler` interface to make sure we could set the correct start/end key.
308324
func (r *Rule) MarshalJSON() ([]byte, error) {
309-
r.StartKeyHex = hex.EncodeToString(encodeBytes(r.StartKey))
310-
r.EndKeyHex = hex.EncodeToString(encodeBytes(r.EndKey))
311-
return json.Marshal(r)
325+
tempRule := &rule{
326+
GroupID: r.GroupID,
327+
ID: r.ID,
328+
Index: r.Index,
329+
Override: r.Override,
330+
StartKeyHex: r.StartKeyHex,
331+
EndKeyHex: r.EndKeyHex,
332+
Role: r.Role,
333+
IsWitness: r.IsWitness,
334+
Count: r.Count,
335+
LabelConstraints: r.LabelConstraints,
336+
LocationLabels: r.LocationLabels,
337+
IsolationLevel: r.IsolationLevel,
338+
}
339+
// Converts the start/end key to hex format if the corresponding hex field is empty.
340+
if len(r.StartKey) > 0 && len(r.StartKeyHex) == 0 {
341+
tempRule.StartKeyHex = rawKeyToKeyHexStr(r.StartKey)
342+
}
343+
if len(r.EndKey) > 0 && len(r.EndKeyHex) == 0 {
344+
tempRule.EndKeyHex = rawKeyToKeyHexStr(r.EndKey)
345+
}
346+
return json.Marshal(tempRule)
312347
}
313348

314349
// UnmarshalJSON implements `json.Unmarshaler` interface to make sure we could get the correct start/end key.
315350
func (r *Rule) UnmarshalJSON(bytes []byte) error {
316-
if err := json.Unmarshal(bytes, r); err != nil {
317-
return err
318-
}
319-
320-
startKey, err := hex.DecodeString(r.StartKeyHex)
351+
var tempRule rule
352+
err := json.Unmarshal(bytes, &tempRule)
321353
if err != nil {
322354
return err
323355
}
324-
325-
endKey, err := hex.DecodeString(r.EndKeyHex)
356+
newRule := Rule{
357+
GroupID: tempRule.GroupID,
358+
ID: tempRule.ID,
359+
Index: tempRule.Index,
360+
Override: tempRule.Override,
361+
StartKeyHex: tempRule.StartKeyHex,
362+
EndKeyHex: tempRule.EndKeyHex,
363+
Role: tempRule.Role,
364+
IsWitness: tempRule.IsWitness,
365+
Count: tempRule.Count,
366+
LabelConstraints: tempRule.LabelConstraints,
367+
LocationLabels: tempRule.LocationLabels,
368+
IsolationLevel: tempRule.IsolationLevel,
369+
}
370+
newRule.StartKey, err = keyHexStrToRawKey(newRule.StartKeyHex)
326371
if err != nil {
327372
return err
328373
}
329-
330-
_, r.StartKey, err = decodeBytes(startKey)
374+
newRule.EndKey, err = keyHexStrToRawKey(newRule.EndKeyHex)
331375
if err != nil {
332376
return err
333377
}
334-
335-
_, r.EndKey, err = decodeBytes(endKey)
336-
337-
return err
378+
*r = newRule
379+
return nil
338380
}
339381

340382
// RuleOpType indicates the operation type
@@ -365,37 +407,87 @@ var (
365407
_ json.Unmarshaler = (*RuleOp)(nil)
366408
)
367409

410+
// This is a helper struct used to customizing the JSON marshal/unmarshal methods of `RuleOp`.
411+
type ruleOp struct {
412+
GroupID string `json:"group_id"`
413+
ID string `json:"id"`
414+
Index int `json:"index,omitempty"`
415+
Override bool `json:"override,omitempty"`
416+
StartKeyHex string `json:"start_key"`
417+
EndKeyHex string `json:"end_key"`
418+
Role PeerRoleType `json:"role"`
419+
IsWitness bool `json:"is_witness"`
420+
Count int `json:"count"`
421+
LabelConstraints []LabelConstraint `json:"label_constraints,omitempty"`
422+
LocationLabels []string `json:"location_labels,omitempty"`
423+
IsolationLevel string `json:"isolation_level,omitempty"`
424+
Action RuleOpType `json:"action"`
425+
DeleteByIDPrefix bool `json:"delete_by_id_prefix"`
426+
}
427+
368428
// MarshalJSON implements `json.Marshaler` interface to make sure we could set the correct start/end key.
369429
func (r *RuleOp) MarshalJSON() ([]byte, error) {
370-
r.StartKeyHex = hex.EncodeToString(encodeBytes(r.StartKey))
371-
r.EndKeyHex = hex.EncodeToString(encodeBytes(r.EndKey))
372-
return json.Marshal(r)
430+
tempRuleOp := &ruleOp{
431+
GroupID: r.GroupID,
432+
ID: r.ID,
433+
Index: r.Index,
434+
Override: r.Override,
435+
StartKeyHex: r.StartKeyHex,
436+
EndKeyHex: r.EndKeyHex,
437+
Role: r.Role,
438+
IsWitness: r.IsWitness,
439+
Count: r.Count,
440+
LabelConstraints: r.LabelConstraints,
441+
LocationLabels: r.LocationLabels,
442+
IsolationLevel: r.IsolationLevel,
443+
Action: r.Action,
444+
DeleteByIDPrefix: r.DeleteByIDPrefix,
445+
}
446+
// Converts the start/end key to hex format if the corresponding hex field is empty.
447+
if len(r.StartKey) > 0 && len(r.StartKeyHex) == 0 {
448+
tempRuleOp.StartKeyHex = rawKeyToKeyHexStr(r.StartKey)
449+
}
450+
if len(r.EndKey) > 0 && len(r.EndKeyHex) == 0 {
451+
tempRuleOp.EndKeyHex = rawKeyToKeyHexStr(r.EndKey)
452+
}
453+
return json.Marshal(tempRuleOp)
373454
}
374455

375456
// UnmarshalJSON implements `json.Unmarshaler` interface to make sure we could get the correct start/end key.
376457
func (r *RuleOp) UnmarshalJSON(bytes []byte) error {
377-
if err := json.Unmarshal(bytes, r); err != nil {
378-
return err
379-
}
380-
381-
startKey, err := hex.DecodeString(r.StartKeyHex)
458+
var tempRuleOp ruleOp
459+
err := json.Unmarshal(bytes, &tempRuleOp)
382460
if err != nil {
383461
return err
384462
}
385-
386-
endKey, err := hex.DecodeString(r.EndKeyHex)
463+
newRuleOp := RuleOp{
464+
Rule: &Rule{
465+
GroupID: tempRuleOp.GroupID,
466+
ID: tempRuleOp.ID,
467+
Index: tempRuleOp.Index,
468+
Override: tempRuleOp.Override,
469+
StartKeyHex: tempRuleOp.StartKeyHex,
470+
EndKeyHex: tempRuleOp.EndKeyHex,
471+
Role: tempRuleOp.Role,
472+
IsWitness: tempRuleOp.IsWitness,
473+
Count: tempRuleOp.Count,
474+
LabelConstraints: tempRuleOp.LabelConstraints,
475+
LocationLabels: tempRuleOp.LocationLabels,
476+
IsolationLevel: tempRuleOp.IsolationLevel,
477+
},
478+
Action: tempRuleOp.Action,
479+
DeleteByIDPrefix: tempRuleOp.DeleteByIDPrefix,
480+
}
481+
newRuleOp.StartKey, err = keyHexStrToRawKey(newRuleOp.StartKeyHex)
387482
if err != nil {
388483
return err
389484
}
390-
391-
_, r.StartKey, err = decodeBytes(startKey)
485+
newRuleOp.EndKey, err = keyHexStrToRawKey(newRuleOp.EndKeyHex)
392486
if err != nil {
393487
return err
394488
}
395-
396-
_, r.EndKey, err = decodeBytes(endKey)
397-
398-
return err
489+
*r = newRuleOp
490+
return nil
399491
}
400492

401493
// RuleGroup defines properties of a rule group.

Diff for: client/http/types_test.go

+63
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
package http
1616

1717
import (
18+
"encoding/json"
1819
"testing"
1920

2021
"github.com/stretchr/testify/require"
@@ -47,3 +48,65 @@ func TestMergeRegionsInfo(t *testing.T) {
4748
re.Equal(2, len(regionsInfo.Regions))
4849
re.Equal(append(regionsInfo1.Regions, regionsInfo2.Regions...), regionsInfo.Regions)
4950
}
51+
52+
func TestRuleStartEndKey(t *testing.T) {
53+
re := require.New(t)
54+
// Empty start/end key and key hex.
55+
ruleToMarshal := &Rule{}
56+
rule := mustMarshalAndUnmarshal(re, ruleToMarshal)
57+
re.Equal("", rule.StartKeyHex)
58+
re.Equal("", rule.EndKeyHex)
59+
re.Equal([]byte(""), rule.StartKey)
60+
re.Equal([]byte(""), rule.EndKey)
61+
// Empty start/end key and non-empty key hex.
62+
ruleToMarshal = &Rule{
63+
StartKeyHex: rawKeyToKeyHexStr([]byte("a")),
64+
EndKeyHex: rawKeyToKeyHexStr([]byte("b")),
65+
}
66+
rule = mustMarshalAndUnmarshal(re, ruleToMarshal)
67+
re.Equal([]byte("a"), rule.StartKey)
68+
re.Equal([]byte("b"), rule.EndKey)
69+
re.Equal(ruleToMarshal.StartKeyHex, rule.StartKeyHex)
70+
re.Equal(ruleToMarshal.EndKeyHex, rule.EndKeyHex)
71+
// Non-empty start/end key and empty key hex.
72+
ruleToMarshal = &Rule{
73+
StartKey: []byte("a"),
74+
EndKey: []byte("b"),
75+
}
76+
rule = mustMarshalAndUnmarshal(re, ruleToMarshal)
77+
re.Equal(ruleToMarshal.StartKey, rule.StartKey)
78+
re.Equal(ruleToMarshal.EndKey, rule.EndKey)
79+
re.Equal(rawKeyToKeyHexStr(ruleToMarshal.StartKey), rule.StartKeyHex)
80+
re.Equal(rawKeyToKeyHexStr(ruleToMarshal.EndKey), rule.EndKeyHex)
81+
// Non-empty start/end key and non-empty key hex.
82+
ruleToMarshal = &Rule{
83+
StartKey: []byte("a"),
84+
EndKey: []byte("b"),
85+
StartKeyHex: rawKeyToKeyHexStr([]byte("c")),
86+
EndKeyHex: rawKeyToKeyHexStr([]byte("d")),
87+
}
88+
rule = mustMarshalAndUnmarshal(re, ruleToMarshal)
89+
re.Equal([]byte("c"), rule.StartKey)
90+
re.Equal([]byte("d"), rule.EndKey)
91+
re.Equal(ruleToMarshal.StartKeyHex, rule.StartKeyHex)
92+
re.Equal(ruleToMarshal.EndKeyHex, rule.EndKeyHex)
93+
// Half of each pair of keys is empty.
94+
ruleToMarshal = &Rule{
95+
StartKey: []byte("a"),
96+
EndKeyHex: rawKeyToKeyHexStr([]byte("d")),
97+
}
98+
rule = mustMarshalAndUnmarshal(re, ruleToMarshal)
99+
re.Equal(ruleToMarshal.StartKey, rule.StartKey)
100+
re.Equal([]byte("d"), rule.EndKey)
101+
re.Equal(rawKeyToKeyHexStr(ruleToMarshal.StartKey), rule.StartKeyHex)
102+
re.Equal(ruleToMarshal.EndKeyHex, rule.EndKeyHex)
103+
}
104+
105+
func mustMarshalAndUnmarshal(re *require.Assertions, rule *Rule) *Rule {
106+
ruleJSON, err := json.Marshal(rule)
107+
re.NoError(err)
108+
var newRule *Rule
109+
err = json.Unmarshal(ruleJSON, &newRule)
110+
re.NoError(err)
111+
return newRule
112+
}

0 commit comments

Comments
 (0)