diff --git a/README.md b/README.md index 3a1441f..48260ae 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,15 @@ $ ./mok --keyspace-id 255 --table-id 43 --row-value 81934 built key: 780000FF74800000FF000000002B5F7280FF0000000001400E00FE ``` +Build a RawKV key for a given raw key under given keyspace +``` +$ ./mok --keyspace-id 0 --key-mode rawkv --raw-key test +built key: 7200000074657374FF0000000000000000F7 + +$ ./mok --keyspace-id 1 --key-mode rawkv --raw-key '' +built key: 7200000100000000FB +``` + ## TODO - [x] build keys diff --git a/main.go b/main.go index e51aaae..c08dcf6 100644 --- a/main.go +++ b/main.go @@ -14,23 +14,35 @@ import ( var keyFormat = flag.String("format", "proto", "output format (go/hex/base64/proto)") var NullSpaceID = int64(0xffffffff) +var keyMode = flag.String("key-mode", "txnkv", "key mode (txnkv/rawkv)") var keyspaceID = flag.Int64("keyspace-id", NullSpaceID, "keyspace ID") var tableID = flag.Int64("table-id", 0, "table ID") var indexID = flag.Int64("index-id", 0, "index ID") var rowValue = flag.String("row-value", "", "row value") var indexValue = flag.String("index-value", "", "index value") +var rawKey = flag.String("raw-key", "", "raw key (rawkv only)") +var rawKeyFormat = flag.String("raw-key-format", "str", "input format (str/hex, rawkv only)") -func getKeyPrefix(keyspaceID int64) ([]byte, error) { +func getKeyPrefix(keyModeStr string, keyspaceID int64) (*KeyMode, []byte, error) { if keyspaceID == NullSpaceID { - return []byte{'t'}, nil + return nil, []byte{'t'}, nil } if keyspaceID > 0xffffff { - return nil, fmt.Errorf("invalid keyspace value: %d", keyspaceID) + return nil, nil, fmt.Errorf("invalid keyspace value: %d", keyspaceID) } + + keyMode := FromStringToKeyMode(keyModeStr) + if keyMode == nil { + return nil, nil, fmt.Errorf("invalid key mode: %s", keyModeStr) + } + var prefix [4]byte binary.BigEndian.PutUint32(prefix[:], uint32(keyspaceID)) - prefix[0] = 'x' - return append(prefix[:], 't'), nil + prefix[0] = byte(*keyMode) + if *keyMode == KeyModeRaw { + return keyMode, prefix[:], nil + } + return keyMode, append(prefix[:], 't'), nil } func main() { @@ -40,42 +52,51 @@ func main() { n := N("key", []byte(flag.Arg(0))) n.Expand().Print() } else if flag.NArg() == 0 { // Build a key with given flags. - key, err := getKeyPrefix(*keyspaceID) + keyMode, key, err := getKeyPrefix(*keyMode, *keyspaceID) if err != nil { fmt.Println(err.Error()) os.Exit(1) } - key = codec.EncodeInt(key, *tableID) - if *tableID == 0 { - fmt.Println("table ID shouldn't be 0") - os.Exit(1) - } - - if *indexID == 0 { - if *rowValue != "" { - key = append(key, []byte("_r")...) - rowValueInt, err := strconv.ParseInt(*rowValue, 10, 64) - if err != nil { - fmt.Printf("invalid row value: %s\n", *rowValue) - os.Exit(1) - } - key = codec.EncodeInt(key, rowValueInt) + if keyMode != nil && *keyMode == KeyModeRaw { + key, err = buildRawKVKey(key, *rawKey, *rawKeyFormat) + if err != nil { + fmt.Println(err.Error()) + os.Exit(1) } } else { - key = append(key, []byte("_i")...) - key = codec.EncodeInt(key, *indexID) - if *indexValue != "" { - indexValueInt, err := strconv.ParseInt(*indexValue, 10, 64) - if err != nil { - fmt.Printf("invalid index value: %s\n", *indexValue) - os.Exit(1) + key = codec.EncodeInt(key, *tableID) + if *tableID == 0 { + fmt.Println("table ID shouldn't be 0") + os.Exit(1) + } + + if *indexID == 0 { + if *rowValue != "" { + key = append(key, []byte("_r")...) + rowValueInt, err := strconv.ParseInt(*rowValue, 10, 64) + if err != nil { + fmt.Printf("invalid row value: %s\n", *rowValue) + os.Exit(1) + } + key = codec.EncodeInt(key, rowValueInt) + } + } else { + key = append(key, []byte("_i")...) + key = codec.EncodeInt(key, *indexID) + if *indexValue != "" { + indexValueInt, err := strconv.ParseInt(*indexValue, 10, 64) + if err != nil { + fmt.Printf("invalid index value: %s\n", *indexValue) + os.Exit(1) + } key = codec.EncodeInt(key, indexValueInt) } } + + key = codec.EncodeBytes([]byte{}, key) } - key = codec.EncodeBytes([]byte{}, key) fmt.Printf("built key: %s\n", strings.ToUpper(hex.EncodeToString(key))) } else { fmt.Println("usage:\nmok {flags} [key]") @@ -83,3 +104,12 @@ func main() { os.Exit(1) } } + +func buildRawKVKey(key []byte, rawKey string, format string) ([]byte, error) { + parsedRawKey, err := ParseRawKey(rawKey, format) + if err != nil { + return nil, err + } + key = append(key, parsedRawKey...) + return codec.EncodeBytes([]byte{}, key), nil +} diff --git a/main_test.go b/main_test.go new file mode 100644 index 0000000..dc01002 --- /dev/null +++ b/main_test.go @@ -0,0 +1,136 @@ +package main + +import ( + "encoding/hex" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBuildRawKVKey(t *testing.T) { + assert := require.New(t) + + keyMode, prefix, err := getKeyPrefix("rawkv", 0) + assert.NoError(err) + assert.Equal(KeyModeRaw, *keyMode) + + tests := []struct { + rawKey string + format string + expected string + err bool + }{ + { + rawKey: "", + format: "str", + expected: "7200000000000000fb", + err: false, + }, + + { + rawKey: "test", + format: "str", + expected: "7200000074657374ff0000000000000000f7", + err: false, + }, + { + rawKey: "testtest", + format: "str", + expected: "7200000074657374ff7465737400000000fb", + err: false, + }, + { + rawKey: "74657374", + format: "hex", + expected: "7200000074657374ff0000000000000000f7", + err: false, + }, + { + rawKey: "test", + format: "invalid", + expected: "", + err: true, + }, + } + + for _, tt := range tests { + actualKey, err := buildRawKVKey(prefix, tt.rawKey, tt.format) + + if tt.err && err == nil { + t.Errorf("expected error, but got nil") + } + + if !tt.err && err != nil { + t.Errorf("expected nil error, but got %v", err) + } + + if hex.EncodeToString(actualKey) != tt.expected { + t.Errorf("expected key %s, but got %s", tt.expected, hex.EncodeToString(actualKey)) + } + } +} + +func TestParseRawKVKey(t *testing.T) { + // ❯ ./mok 7200000074657374ff0000000000000000f7 + // "7200000074657374ff0000000000000000f7" + // └─## decode hex key + // └─"r\000\000\000test\377\000\000\000\000\000\000\000\000\367" + // ├─## decode mvcc key + // │ └─"r\000\000\000test" + // │ └─## decode keyspace + // │ ├─key mode: rawkv + // │ ├─keyspace: 0 + // │ └─"test" + // └─## decode keyspace + // ├─key mode: rawkv + // ├─keyspace: 0 + // └─"test\377\000\000\000\000\000\000\000\000\367" + assert := require.New(t) + + tests := []struct { + encoded string + expected string + }{ + { + encoded: "7200000000000000fb", + expected: "", + }, + { + encoded: "7200000074657374ff0000000000000000f7", + expected: "test", + }, + { + encoded: "7200000074657374ff7465737400000000fb", + expected: "testtest", + }, + } + + for _, tt := range tests { + n := N("key", []byte(tt.encoded)).Expand() + assert.Len(n.variants, 1) // decode hex key + assert.Equal("decode hex key", n.variants[0].method) + assert.Len(n.variants[0].children, 1) + + n = n.variants[0].children[0] + assert.Len(n.variants, 2) // with and without mvcc + + variant := n.variants[0] + assert.Equal("decode mvcc key", variant.method) + assert.Len(variant.children, 1) + + n = variant.children[0] + assert.Len(n.variants, 1) + variant = n.variants[0] + assert.Equal("decode keyspace", variant.method) + assert.Len(variant.children, 3) + + assert.Equal("key_mode", variant.children[0].typ) + assert.Equal("r", string(variant.children[0].val)) + + assert.Equal("keyspace_id", variant.children[1].typ) + assert.Equal([]byte{0, 0, 0}, variant.children[1].val) + + assert.Equal("raw_key", variant.children[2].typ) + assert.Equal(tt.expected, string(variant.children[2].val)) + } +} diff --git a/mok.go b/mok.go index 673ca7d..982797d 100644 --- a/mok.go +++ b/mok.go @@ -27,7 +27,7 @@ func N(t string, v []byte) *Node { func (n *Node) String() string { switch n.typ { - case "key", "index_values": + case "key", "raw_key", "index_values": switch *keyFormat { case "hex": return `"` + strings.ToUpper(hex.EncodeToString(n.val)) + `"` @@ -38,6 +38,8 @@ func (n *Node) String() string { default: return fmt.Sprintf("%q", n.val) } + case "key_mode": + return fmt.Sprintf("key mode: %s", KeyMode(n.val[0])) case "keyspace_id": tmp := []byte{'\x00'} t := append(tmp, n.val...) diff --git a/rules.go b/rules.go index a16eb13..af24cd4 100644 --- a/rules.go +++ b/rules.go @@ -81,10 +81,14 @@ func DecodeRocksDBKey(n *Node) *Variant { } func DecodeKeyspace(n *Node) *Variant { - if n.typ == "key" && n.val[0] == 'x' && len(n.val) >= 4 { + if n.typ == "key" && IsValidKeyMode(n.val[0]) && len(n.val) >= 4 { + keyType := "key" + if IsRawKeyMode(n.val[0]) { + keyType = "raw_key" + } return &Variant{ method: "decode keyspace", - children: []*Node{N("keyspace_id", n.val[1:4]), N("key", n.val[4:])}, + children: []*Node{N("key_mode", n.val[0:1]), N("keyspace_id", n.val[1:4]), N(keyType, n.val[4:])}, } } return nil diff --git a/util.go b/util.go index fafdf28..1402371 100644 --- a/util.go +++ b/util.go @@ -2,6 +2,7 @@ package main import ( "bytes" + "encoding/hex" "fmt" "io" "strings" @@ -76,3 +77,52 @@ func GetTimeFromTS(ts uint64) time.Time { ms := int64(ts >> 18) return time.Unix(ms/1e3, (ms%1e3)*1e6) } + +type KeyMode byte + +const ( + KeyModeTxn KeyMode = 'x' + KeyModeRaw KeyMode = 'r' +) + +func IsValidKeyMode(b byte) bool { + return b == byte(KeyModeTxn) || b == byte(KeyModeRaw) +} + +func IsRawKeyMode(b byte) bool { + return b == byte(KeyModeRaw) +} + +func (k KeyMode) String() string { + switch k { + case KeyModeTxn: + return "txnkv" + case KeyModeRaw: + return "rawkv" + default: + return "other" + } +} + +func FromStringToKeyMode(s string) *KeyMode { + var keyMode KeyMode + switch s { + case "txnkv": + keyMode = KeyModeTxn + case "rawkv": + keyMode = KeyModeRaw + default: + } + return &keyMode +} + +func ParseRawKey(s string, format string) ([]byte, error) { + switch format { + case "hex": + return hex.DecodeString(s) + case "str": // for `s` with all characters printable. + return []byte(s), nil + default: + return nil, fmt.Errorf("invalid raw key format: %s", format) + } +}