From 47c26669441603306ae223d8bc5376b4c214eb53 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Sat, 11 Feb 2023 21:29:10 +0800 Subject: [PATCH 1/2] check if key is not API V1 first, then decode the key Signed-off-by: iosmanthus --- tablecodec/BUILD.bazel | 1 + tablecodec/tablecodec.go | 13 +++++++------ tablecodec/tablecodec_test.go | 11 +++++++++++ 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/tablecodec/BUILD.bazel b/tablecodec/BUILD.bazel index db89b945b183d..2cf712bd02d34 100644 --- a/tablecodec/BUILD.bazel +++ b/tablecodec/BUILD.bazel @@ -49,6 +49,7 @@ go_test( "//util/rowcodec", "@com_github_pingcap_failpoint//:failpoint", "@com_github_stretchr_testify//require", + "@com_github_tikv_client_go_v2//tikv", "@org_uber_go_goleak//:goleak", ], ) diff --git a/tablecodec/tablecodec.go b/tablecodec/tablecodec.go index 394277ba2768c..f41b9ed2270e0 100644 --- a/tablecodec/tablecodec.go +++ b/tablecodec/tablecodec.go @@ -278,13 +278,14 @@ func DecodeKeyHead(key kv.Key) (tableID int64, indexID int64, isRecordKey bool, // DecodeTableID decodes the table ID of the key, if the key is not table key, returns 0. func DecodeTableID(key kv.Key) int64 { - // If the key is in API V2, then ignore the prefix - _, k, err := tikv.DecodeKey(key, kvrpcpb.APIVersion_V2) - if err == nil { - key = k - } if !key.HasPrefix(tablePrefix) { - return 0 + // If the key is in API V2, then ignore the prefix + _, k, err := tikv.DecodeKey(key, kvrpcpb.APIVersion_V2) + if err != nil { + terror.Log(errors.Trace(err)) + return 0 + } + key = k } key = key[len(tablePrefix):] _, tableID, err := codec.DecodeInt(key) diff --git a/tablecodec/tablecodec_test.go b/tablecodec/tablecodec_test.go index adc4ccc78c13b..1010669801049 100644 --- a/tablecodec/tablecodec_test.go +++ b/tablecodec/tablecodec_test.go @@ -30,6 +30,7 @@ import ( "github.com/pingcap/tidb/util/collate" "github.com/pingcap/tidb/util/rowcodec" "github.com/stretchr/testify/require" + "github.com/tikv/client-go/v2/tikv" ) // TestTableCodec tests some functions in package tablecodec @@ -698,3 +699,13 @@ func TestTempIndexValueCodec(t *testing.T) { require.Equal(t, result[1].Handle.IntValue(), int64(100)) require.Equal(t, result[2].Handle.IntValue(), int64(101)) } + +func TestV2TableCodec(t *testing.T) { + const tableID int64 = 31415926 + key := EncodeTablePrefix(tableID) + c, err := tikv.NewCodecV2(tikv.ModeTxn, 271828) + require.NoError(t, err) + key = c.EncodeKey(key) + tbid := DecodeTableID(key) + require.Equal(t, tableID, tbid) +} From db12bdd3cc98732b1edd21248c51c1654c9db2b1 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Sat, 11 Feb 2023 21:52:08 +0800 Subject: [PATCH 2/2] recheck table prefix Signed-off-by: iosmanthus --- tablecodec/tablecodec.go | 3 +++ tablecodec/tablecodec_test.go | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/tablecodec/tablecodec.go b/tablecodec/tablecodec.go index f41b9ed2270e0..e5adfa6f91b32 100644 --- a/tablecodec/tablecodec.go +++ b/tablecodec/tablecodec.go @@ -286,6 +286,9 @@ func DecodeTableID(key kv.Key) int64 { return 0 } key = k + if !key.HasPrefix(tablePrefix) { + return 0 + } } key = key[len(tablePrefix):] _, tableID, err := codec.DecodeInt(key) diff --git a/tablecodec/tablecodec_test.go b/tablecodec/tablecodec_test.go index 1010669801049..45020e034938c 100644 --- a/tablecodec/tablecodec_test.go +++ b/tablecodec/tablecodec_test.go @@ -708,4 +708,12 @@ func TestV2TableCodec(t *testing.T) { key = c.EncodeKey(key) tbid := DecodeTableID(key) require.Equal(t, tableID, tbid) + + key = []byte("x001HelloWorld") + tbid = DecodeTableID(key) + require.Equal(t, int64(0), tbid) + + key = []byte("x001x001t123") + tbid = DecodeTableID(key) + require.Equal(t, int64(0), tbid) }