key_registry.go raw
1 /*
2 * SPDX-FileCopyrightText: © Hypermode Inc. <hello@hypermode.com>
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6 package badger
7
8 import (
9 "bytes"
10 "crypto/aes"
11 "crypto/rand"
12 "encoding/binary"
13 "hash/crc32"
14 "io"
15 "os"
16 "path/filepath"
17 "sync"
18 "time"
19
20 "github.com/dgraph-io/badger/v4/pb"
21 "github.com/dgraph-io/badger/v4/y"
22 "google.golang.org/protobuf/proto"
23 )
24
25 const (
26 // KeyRegistryFileName is the file name for the key registry file.
27 KeyRegistryFileName = "KEYREGISTRY"
28 // KeyRegistryRewriteFileName is the file name for the rewrite key registry file.
29 KeyRegistryRewriteFileName = "REWRITE-KEYREGISTRY"
30 )
31
32 // SanityText is used to check whether the given user provided storage key is valid or not
33 var sanityText = []byte("Hello Badger")
34
35 // KeyRegistry used to maintain all the data keys.
36 type KeyRegistry struct {
37 sync.RWMutex
38 dataKeys map[uint64]*pb.DataKey
39 lastCreated int64 //lastCreated is the timestamp(seconds) of the last data key generated.
40 nextKeyID uint64
41 fp *os.File
42 opt KeyRegistryOptions
43 }
44
45 type KeyRegistryOptions struct {
46 Dir string
47 ReadOnly bool
48 EncryptionKey []byte
49 EncryptionKeyRotationDuration time.Duration
50 InMemory bool
51 }
52
53 // newKeyRegistry returns KeyRegistry.
54 func newKeyRegistry(opt KeyRegistryOptions) *KeyRegistry {
55 return &KeyRegistry{
56 dataKeys: make(map[uint64]*pb.DataKey),
57 nextKeyID: 0,
58 opt: opt,
59 }
60 }
61
62 // OpenKeyRegistry opens key registry if it exists, otherwise it'll create key registry
63 // and returns key registry.
64 func OpenKeyRegistry(opt KeyRegistryOptions) (*KeyRegistry, error) {
65 // sanity check the encryption key length.
66 if len(opt.EncryptionKey) > 0 {
67 switch len(opt.EncryptionKey) {
68 default:
69 return nil, y.Wrapf(ErrInvalidEncryptionKey, "During OpenKeyRegistry")
70 case 16, 24, 32:
71 break
72 }
73 }
74 // If db is opened in InMemory mode, we don't need to write key registry to the disk.
75 if opt.InMemory {
76 return newKeyRegistry(opt), nil
77 }
78 path := filepath.Join(opt.Dir, KeyRegistryFileName)
79 var flags y.Flags
80 if opt.ReadOnly {
81 flags |= y.ReadOnly
82 } else {
83 flags |= y.Sync
84 }
85 fp, err := y.OpenExistingFile(path, flags)
86 // OpenExistingFile just open file.
87 // So checking whether the file exist or not. If not
88 // We'll create new keyregistry.
89 if os.IsNotExist(err) {
90 // Creating new registry file if not exist.
91 kr := newKeyRegistry(opt)
92 if opt.ReadOnly {
93 return kr, nil
94 }
95 // Writing the key registry to the file.
96 if err := WriteKeyRegistry(kr, opt); err != nil {
97 return nil, y.Wrapf(err, "Error while writing key registry.")
98 }
99 fp, err = y.OpenExistingFile(path, flags)
100 if err != nil {
101 return nil, y.Wrapf(err, "Error while opening newly created key registry.")
102 }
103 } else if err != nil {
104 return nil, y.Wrapf(err, "Error while opening key registry.")
105 }
106 kr, err := readKeyRegistry(fp, opt)
107 if err != nil {
108 // This case happens only if the file is opened properly and
109 // not able to read.
110 fp.Close()
111 return nil, err
112 }
113 if opt.ReadOnly {
114 // We'll close the file in readonly mode.
115 return kr, fp.Close()
116 }
117 kr.fp = fp
118 return kr, nil
119 }
120
121 // keyRegistryIterator reads all the datakey from the key registry
122 type keyRegistryIterator struct {
123 encryptionKey []byte
124 fp *os.File
125 // lenCrcBuf contains crc buf and data length to move forward.
126 lenCrcBuf [8]byte
127 }
128
129 // newKeyRegistryIterator returns iterator which will allow you to iterate
130 // over the data key of the key registry.
131 func newKeyRegistryIterator(fp *os.File, encryptionKey []byte) (*keyRegistryIterator, error) {
132 return &keyRegistryIterator{
133 encryptionKey: encryptionKey,
134 fp: fp,
135 lenCrcBuf: [8]byte{},
136 }, validRegistry(fp, encryptionKey)
137 }
138
139 // validRegistry checks that given encryption key is valid or not.
140 func validRegistry(fp *os.File, encryptionKey []byte) error {
141 iv := make([]byte, aes.BlockSize)
142 var err error
143 if _, err = fp.Read(iv); err != nil {
144 return y.Wrapf(err, "Error while reading IV for key registry.")
145 }
146 eSanityText := make([]byte, len(sanityText))
147 if _, err = fp.Read(eSanityText); err != nil {
148 return y.Wrapf(err, "Error while reading sanity text.")
149 }
150 if len(encryptionKey) > 0 {
151 // Decrypting sanity text.
152 if eSanityText, err = y.XORBlockAllocate(eSanityText, encryptionKey, iv); err != nil {
153 return y.Wrapf(err, "During validRegistry")
154 }
155 }
156 // Check the given key is valid or not.
157 if !bytes.Equal(eSanityText, sanityText) {
158 return ErrEncryptionKeyMismatch
159 }
160 return nil
161 }
162
163 func (kri *keyRegistryIterator) next() (*pb.DataKey, error) {
164 var err error
165 // Read crc buf and data length.
166 if _, err = kri.fp.Read(kri.lenCrcBuf[:]); err != nil {
167 // EOF means end of the iteration.
168 if err != io.EOF {
169 return nil, y.Wrapf(err, "While reading crc in keyRegistryIterator.next")
170 }
171 return nil, err
172 }
173 l := int64(binary.BigEndian.Uint32(kri.lenCrcBuf[0:4]))
174 // Read protobuf data.
175 data := make([]byte, l)
176 if _, err = kri.fp.Read(data); err != nil {
177 // EOF means end of the iteration.
178 if err != io.EOF {
179 return nil, y.Wrapf(err, "While reading protobuf in keyRegistryIterator.next")
180 }
181 return nil, err
182 }
183 // Check checksum.
184 if crc32.Checksum(data, y.CastagnoliCrcTable) != binary.BigEndian.Uint32(kri.lenCrcBuf[4:]) {
185 return nil, y.Wrapf(y.ErrChecksumMismatch, "Error while checking checksum for data key.")
186 }
187 dataKey := &pb.DataKey{}
188 if err = proto.Unmarshal(data, dataKey); err != nil {
189 return nil, y.Wrapf(err, "While unmarshal of datakey in keyRegistryIterator.next")
190 }
191 if len(kri.encryptionKey) > 0 {
192 // Decrypt the key if the storage key exists.
193 if dataKey.Data, err = y.XORBlockAllocate(dataKey.Data, kri.encryptionKey, dataKey.Iv); err != nil {
194 return nil, y.Wrapf(err, "While decrypting datakey in keyRegistryIterator.next")
195 }
196 }
197 return dataKey, nil
198 }
199
200 // readKeyRegistry will read the key registry file and build the key registry struct.
201 func readKeyRegistry(fp *os.File, opt KeyRegistryOptions) (*KeyRegistry, error) {
202 itr, err := newKeyRegistryIterator(fp, opt.EncryptionKey)
203 if err != nil {
204 return nil, err
205 }
206 kr := newKeyRegistry(opt)
207 var dk *pb.DataKey
208 dk, err = itr.next()
209 for err == nil && dk != nil {
210 if dk.KeyId > kr.nextKeyID {
211 // Set the maximum key ID for next key ID generation.
212 kr.nextKeyID = dk.KeyId
213 }
214 if dk.CreatedAt > kr.lastCreated {
215 // Set the last generated key timestamp.
216 kr.lastCreated = dk.CreatedAt
217 }
218 // No need to lock since we are building the initial state.
219 kr.dataKeys[dk.KeyId] = dk
220 // Forward the iterator.
221 dk, err = itr.next()
222 }
223 // We read all the key. So, Ignoring this error.
224 if err == io.EOF {
225 err = nil
226 }
227 return kr, err
228 }
229
230 /*
231 Structure of Key Registry.
232 +-------------------+---------------------+--------------------+--------------+------------------+
233 | IV | Sanity Text | DataKey1 | DataKey2 | ... |
234 +-------------------+---------------------+--------------------+--------------+------------------+
235 */
236
237 // WriteKeyRegistry will rewrite the existing key registry file with new one.
238 // It is okay to give closed key registry. Since, it's using only the datakey.
239 func WriteKeyRegistry(reg *KeyRegistry, opt KeyRegistryOptions) error {
240 buf := &bytes.Buffer{}
241 iv, err := y.GenerateIV()
242 y.Check(err)
243 // Encrypt sanity text if the encryption key is presents.
244 eSanity := sanityText
245 if len(opt.EncryptionKey) > 0 {
246 var err error
247 eSanity, err = y.XORBlockAllocate(eSanity, opt.EncryptionKey, iv)
248 if err != nil {
249 return y.Wrapf(err, "Error while encrpting sanity text in WriteKeyRegistry")
250 }
251 }
252 y.Check2(buf.Write(iv))
253 y.Check2(buf.Write(eSanity))
254 // Write all the datakeys to the buf.
255 for _, k := range reg.dataKeys {
256 // Writing the datakey to the given buffer.
257 if err := storeDataKey(buf, opt.EncryptionKey, k); err != nil {
258 return y.Wrapf(err, "Error while storing datakey in WriteKeyRegistry")
259 }
260 }
261 tmpPath := filepath.Join(opt.Dir, KeyRegistryRewriteFileName)
262 // Open temporary file to write the data and do atomic rename.
263 fp, err := y.OpenTruncFile(tmpPath, true)
264 if err != nil {
265 return y.Wrapf(err, "Error while opening tmp file in WriteKeyRegistry")
266 }
267 // Write buf to the disk.
268 if _, err = fp.Write(buf.Bytes()); err != nil {
269 // close the fd before returning error. We're not using defer
270 // because, for windows we need to close the fd explicitly before
271 // renaming.
272 fp.Close()
273 return y.Wrapf(err, "Error while writing buf in WriteKeyRegistry")
274 }
275 // In Windows the files should be closed before doing a Rename.
276 if err = fp.Close(); err != nil {
277 return y.Wrapf(err, "Error while closing tmp file in WriteKeyRegistry")
278 }
279 // Rename to the original file.
280 if err = os.Rename(tmpPath, filepath.Join(opt.Dir, KeyRegistryFileName)); err != nil {
281 return y.Wrapf(err, "Error while renaming file in WriteKeyRegistry")
282 }
283 // Sync Dir.
284 return syncDir(opt.Dir)
285 }
286
287 // DataKey returns datakey of the given key id.
288 func (kr *KeyRegistry) DataKey(id uint64) (*pb.DataKey, error) {
289 kr.RLock()
290 defer kr.RUnlock()
291 if id == 0 {
292 // nil represent plain text.
293 return nil, nil
294 }
295 dk, ok := kr.dataKeys[id]
296 if !ok {
297 return nil, y.Wrapf(ErrInvalidDataKeyID, "Error for the KEY ID %d", id)
298 }
299 return dk, nil
300 }
301
302 // LatestDataKey will give you the latest generated datakey based on the rotation
303 // period. If the last generated datakey lifetime exceeds the rotation period.
304 // It'll create new datakey.
305 func (kr *KeyRegistry) LatestDataKey() (*pb.DataKey, error) {
306 if len(kr.opt.EncryptionKey) == 0 {
307 // nil is for no encryption.
308 return nil, nil
309 }
310 // validKey return datakey if the last generated key duration less than
311 // rotation duration.
312 validKey := func() (*pb.DataKey, bool) {
313 // Time diffrence from the last generated time.
314 diff := time.Since(time.Unix(kr.lastCreated, 0))
315 if diff < kr.opt.EncryptionKeyRotationDuration {
316 return kr.dataKeys[kr.nextKeyID], true
317 }
318 return nil, false
319 }
320 kr.RLock()
321 key, valid := validKey()
322 kr.RUnlock()
323 if valid {
324 // If less than EncryptionKeyRotationDuration, returns the last generated key.
325 return key, nil
326 }
327 kr.Lock()
328 defer kr.Unlock()
329 // Key might have generated by another go routine. So,
330 // checking once again.
331 key, valid = validKey()
332 if valid {
333 return key, nil
334 }
335 k := make([]byte, len(kr.opt.EncryptionKey))
336 iv, err := y.GenerateIV()
337 if err != nil {
338 return nil, err
339 }
340 _, err = rand.Read(k)
341 if err != nil {
342 return nil, err
343 }
344 // Otherwise Increment the KeyID and generate new datakey.
345 kr.nextKeyID++
346 dk := &pb.DataKey{
347 KeyId: kr.nextKeyID,
348 Data: k,
349 CreatedAt: time.Now().Unix(),
350 Iv: iv,
351 }
352 // Don't store the datakey on file if badger is running in InMemory mode.
353 if !kr.opt.InMemory {
354 // Store the datekey.
355 buf := &bytes.Buffer{}
356 if err = storeDataKey(buf, kr.opt.EncryptionKey, dk); err != nil {
357 return nil, err
358 }
359 // Persist the datakey to the disk
360 if _, err = kr.fp.Write(buf.Bytes()); err != nil {
361 return nil, err
362 }
363 }
364 // storeDatakey encrypts the datakey So, placing un-encrypted key in the memory.
365 dk.Data = k
366 kr.lastCreated = dk.CreatedAt
367 kr.dataKeys[kr.nextKeyID] = dk
368 return dk, nil
369 }
370
371 // Close closes the key registry.
372 func (kr *KeyRegistry) Close() error {
373 if !(kr.opt.ReadOnly || kr.opt.InMemory) {
374 return kr.fp.Close()
375 }
376 return nil
377 }
378
379 // storeDataKey stores datakey in an encrypted format in the given buffer. If storage key preset.
380 func storeDataKey(buf *bytes.Buffer, storageKey []byte, k *pb.DataKey) error {
381 // xor will encrypt the IV and xor with the given data.
382 // It'll used for both encryption and decryption.
383 xor := func() error {
384 if len(storageKey) == 0 {
385 return nil
386 }
387 var err error
388 k.Data, err = y.XORBlockAllocate(k.Data, storageKey, k.Iv)
389 return err
390 }
391 // In memory datakey will be plain text so encrypting before storing to the disk.
392 var err error
393 if err = xor(); err != nil {
394 return y.Wrapf(err, "Error while encrypting datakey in storeDataKey")
395 }
396 var data []byte
397 if data, err = proto.Marshal(k); err != nil {
398 err = y.Wrapf(err, "Error while marshaling datakey in storeDataKey")
399 var err2 error
400 // decrypting the datakey back.
401 if err2 = xor(); err2 != nil {
402 return y.Wrapf(err,
403 y.Wrapf(err2, "Error while decrypting datakey in storeDataKey").Error())
404 }
405 return err
406 }
407 var lenCrcBuf [8]byte
408 binary.BigEndian.PutUint32(lenCrcBuf[0:4], uint32(len(data)))
409 binary.BigEndian.PutUint32(lenCrcBuf[4:8], crc32.Checksum(data, y.CastagnoliCrcTable))
410 y.Check2(buf.Write(lenCrcBuf[:]))
411 y.Check2(buf.Write(data))
412 // Decrypting the datakey back since we're using the pointer.
413 return xor()
414 }
415