fix api v2 decode of region bucket's keys (#757)

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

internal/apicodec/codec.go

@@ -57,6 +57,8 @@ type Codec interface {
 	EncodeRegionKey(key []byte) []byte
 	// DecodeRegionKey decode region's key
 	DecodeRegionKey(encodedKey []byte) ([]byte, error)
+	// DecodeBucketKey decode region bucket's key
+	DecodeBucketKey(encodedKey []byte) ([]byte, error)
 	// EncodeRegionRange encode region's start and end.
 	EncodeRegionRange(start, end []byte) ([]byte, []byte)
 	// DecodeRegionRange decode region's start and end.

internal/apicodec/codec_v1.go

@@ -202,3 +202,7 @@ func (c *codecV1) DecodeRange(start, end []byte) ([]byte, []byte, error) {
 func (c *codecV1) DecodeKey(key []byte) ([]byte, error) {
 	return key, nil
 }
+
+func (c *codecV1) DecodeBucketKey(key []byte) ([]byte, error) {
+	return c.DecodeRegionKey(key)
+}

internal/apicodec/codec_v1_test.go

@@ -0,0 +1,23 @@
+package apicodec
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tikv/client-go/v2/util/codec"
+)
+
+func TestV1DecodeBucketKey(t *testing.T) {
+	c := NewCodecV1(ModeTxn)
+	raw := []byte("test")
+	encoded := codec.EncodeBytes(nil, raw)
+	key, err := c.DecodeBucketKey(encoded)
+	assert.Nil(t, err)
+	assert.Equal(t, raw, key)
+
+	raw = []byte{}
+	encoded = codec.EncodeBytes(nil, raw)
+	key, err = c.DecodeBucketKey(encoded)
+	assert.Nil(t, err)
+	assert.Equal(t, raw, key)
+}

internal/apicodec/codec_v2.go

@@ -953,3 +953,16 @@ func (c *codecV2) decodeLockInfos(locks []*kvrpcpb.LockInfo) ([]*kvrpcpb.LockInf
 	}
 	return locks, nil
 }
+
+func (c *codecV2) DecodeBucketKey(encodedKey []byte) ([]byte, error) {
+	key, err := c.memCodec.decodeKey(encodedKey)
+	if err != nil {
+		return nil, err
+	}
+
+	if bytes.Compare(key, c.endKey) >= 0 || bytes.Compare(key, c.prefix) < 0 {
+		return []byte{}, nil
+	}
+
+	return key[len(c.prefix):], nil
+}

internal/apicodec/codec_v2_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/mpp"
 	"github.com/stretchr/testify/suite"
 	"github.com/tikv/client-go/v2/tikvrpc"
+	"github.com/tikv/client-go/v2/util/codec"
 )
 
 var (
@@ -18,12 +19,13 @@ var (
 	// Keys below are ordered as following:
 	// beforePrefix, keyspacePrefix, insideLeft, insideRight, keyspaceEndKey, afterEndKey
 	// where valid keyspace range is [keyspacePrefix, keyspaceEndKey)
-	keyspacePrefix = []byte{'r', 0, 16, 146}
+	prevKeyspacePrefix = []byte{'r', 0, 16, 145}
-	keyspaceEndKey = []byte{'r', 0, 16, 147}
+	keyspacePrefix     = []byte{'r', 0, 16, 146}
-	beforePrefix   = []byte{'r', 0, 0, 1}
+	keyspaceEndKey     = []byte{'r', 0, 16, 147}
-	afterEndKey    = []byte{'r', 1, 0, 0}
+	beforePrefix       = []byte{'r', 0, 0, 1}
-	insideLeft     = []byte{'r', 0, 16, 146, 100}
+	afterEndKey        = []byte{'r', 1, 0, 0}
-	insideRight    = []byte{'r', 0, 16, 146, 200}
+	insideLeft         = []byte{'r', 0, 16, 146, 100}
+	insideRight        = []byte{'r', 0, 16, 146, 200}
 )
 
 type testCodecV2Suite struct {
@@ -293,3 +295,23 @@ func (suite *testCodecV2Suite) TestEncodeMPPRequest() {
 	suite.Equal(task.Regions[0].Ranges[0].Start, suite.codec.EncodeKey([]byte("a")))
 	suite.Equal(task.Regions[0].Ranges[0].End, suite.codec.EncodeKey([]byte("b")))
 }
+
+func (suite *testCodecV2Suite) TestDecodeBucketKey() {
+	raw := []byte("a")
+	key := suite.codec.EncodeRegionKey(raw)
+	bucketKey, err := suite.codec.DecodeBucketKey(key)
+	suite.Nil(err)
+	suite.Equal(raw, bucketKey)
+
+	raw = keyspaceEndKey
+	key = codec.EncodeBytes([]byte{}, raw)
+	bucketKey, err = suite.codec.DecodeBucketKey(key)
+	suite.Nil(err)
+	suite.Empty(bucketKey)
+
+	raw = append(prevKeyspacePrefix, []byte("a")...)
+	key = codec.EncodeBytes([]byte{}, raw)
+	bucketKey, err = suite.codec.DecodeBucketKey(key)
+	suite.Nil(err)
+	suite.Empty(bucketKey)
+}

internal/locate/pd_codec.go

@@ -167,7 +167,7 @@ func (c *CodecPDClient) decodeRegionKeyInPlace(r *pd.Region) error {
 			if len(k) == 0 {
 				continue
 			}
-			decoded, err := c.codec.DecodeRegionKey(k)
+			decoded, err := c.codec.DecodeBucketKey(k)
 			if err != nil {
 				return errors.WithStack(err)
 			}