redis-protocol-server/Redis/Rdb/BinaryReaderExtensions.cs at main · inner/redis-protocol-server · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
using Redis.Common;

namespace Redis.Rdb;

public static class BinaryReaderExtensions
{
    public static ulong ReadLength(this BinaryReader br)
    {
        var (len, _) = ReadLengthWithEncoding(br);
        return len;
    }

    public static byte[] ReadStr(this BinaryReader binaryReader)
    {
        // https://github.com/sripathikrishnan/redis-rdb-tools/wiki/Redis-RDB-Dump-File-Format#string-encoding
        var (length, isEncoded) = ReadLengthWithEncoding(binaryReader);

        if (!isEncoded)
        {
            return binaryReader.ReadBytes((int)length);
        }

        switch (length)
        {
            case RdbConstants.EncType.Int8:
            {
                var tmp = binaryReader.ReadBytes(RdbConstants.One);
                return ((sbyte)tmp[0]).ToString().AsBytes();
            }
            case RdbConstants.EncType.Int16:
            {
                var tmp = binaryReader.ReadBytes(RdbConstants.Two);
                return BitConverter.ToInt16(tmp).ToString().AsBytes();
            }
            case RdbConstants.EncType.Int32:
            {
                var tmp = binaryReader.ReadBytes(RdbConstants.Four);
                return BitConverter.ToInt32(tmp).ToString().AsBytes();
            }
            case RdbConstants.EncType.Lzf:
            {
                var compressedLen = ReadLength(binaryReader);
                var uncompressedLen = ReadLength(binaryReader);

                var compressed = binaryReader.ReadBytes((int)compressedLen);
                var decompressed = LzfDecompress(compressed, (int)uncompressedLen);

                if (decompressed.Length != (int)uncompressedLen)
                    throw new RdbReaderException(
                        $"decompressed string length {decompressed.Length} didn't match expected length {(int)uncompressedLen}");

                return decompressed;
            }
            default:
                throw new RdbReaderException($"Invalid string encoding {length}");
        }
    }

    private static (ulong Length, bool IsEncoded) ReadLengthWithEncoding(this BinaryReader binaryReader)
    {
        // https://github.com/sripathikrishnan/redis-rdb-tools/wiki/Redis-RDB-Dump-File-Format#length-encoding
        ulong len;
        var isEncoded = false;

        var @byte = binaryReader.ReadByte();

        // 8bit right shift 6bit, get the starting 2bit
        // 0xC0  11000000
        // 0x3F  00111111
        var encodingType = (@byte & 0xC0) >> 6;

        switch (encodingType)
        {
            case RdbConstants.LengthEncoding.Bit6:
                // starting bits are 00
                len = (ulong)(@byte & 0x3F);
                break;
            case RdbConstants.LengthEncoding.Bit14:
            {
                // starting bits are 01
                var b1 = binaryReader.ReadByte();
                len = (ulong)((@byte & 0x3F) << 8 | b1);
                break;
            }
            case RdbConstants.LengthEncoding.EncVal:
                // starting bits are 11
                len = (ulong)(@byte & 0x3F);
                isEncoded = true;
                break;
            default:
            {
                len = @byte switch
                {
                    RdbConstants.LengthEncoding.Bit32 or RdbConstants.LengthEncoding.Bit64 =>
                        // starting bits are 10
                        binaryReader.ReadUInt32BigEndian(),
                    _ => throw new RdbReaderException($"Invalid string encoding {encodingType} (encoding byte {@byte})")
                };

                break;
            }
        }

        return (len, isEncoded);
    }

    private static byte[] LzfDecompress(byte[] compressed, int uncompressedLength)
    {
        var outStream = new List<byte>(uncompressedLength);
        var outIndex = 0;

        var inLen = compressed.Length;
        var inIndex = 0;

        while (inIndex < inLen)
        {
            var ctrl = compressed[inIndex];

            inIndex += 1;

            if (ctrl < 32)
            {
                for (var i = 0; i < ctrl + 1; i++)
                {
                    outStream.Add(compressed[inIndex]);
                    inIndex += 1;
                    outIndex += 1;
                }
            }
            else
            {
                var length = ctrl >> 5;
                if (length == 7)
                {
                    length += compressed[inIndex];
                    inIndex += 1;
                }

                var @ref = outIndex - ((ctrl & 0x1f) << 8) - compressed[inIndex] - 1;
                inIndex += 1;

                for (var i = 0; i < length + 2; i++)
                {
                    outStream.Add(outStream[@ref]);
                    @ref += 1;
                    outIndex += 1;
                }
            }
        }

        return outStream.ToArray();
    }

    private static uint ReadUInt32BigEndian(this BinaryReader br)
    {
        var bytes = br.ReadBytes(RdbConstants.Four);
        Array.Reverse(bytes);
        return BitConverter.ToUInt32(bytes, 0);
    }
}