github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/ccl/storageccl/engineccl/ctr_stream.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Licensed as a CockroachDB Enterprise file under the Cockroach Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
     8  
     9  package engineccl
    10  
    11  import (
    12  	"context"
    13  	"crypto/aes"
    14  	"crypto/cipher"
    15  	"crypto/rand"
    16  	"encoding/binary"
    17  	"fmt"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/ccl/storageccl/engineccl/enginepbccl"
    20  	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
    21  )
    22  
    23  // FileCipherStreamCreator wraps the KeyManager interface and provides functions to create a
    24  // FileStream for either a new file (using the active key provided by the KeyManager) or an
    25  // existing file (by looking up the key in the KeyManager).
    26  type FileCipherStreamCreator struct {
    27  	envType    enginepb.EnvType
    28  	keyManager PebbleKeyManager
    29  }
    30  
    31  const (
    32  	// The difference is 4 bytes, which are supplied by the counter.
    33  	ctrBlockSize = 16
    34  	ctrNonceSize = 12
    35  )
    36  
    37  // CreateNew creates a FileStream for a new file using the currently active key. It returns the
    38  // settings used, so that the caller can record these in a file registry.
    39  func (c *FileCipherStreamCreator) CreateNew(
    40  	ctx context.Context,
    41  ) (*enginepbccl.EncryptionSettings, FileStream, error) {
    42  	key, err := c.keyManager.ActiveKey(ctx)
    43  	if err != nil {
    44  		return nil, nil, err
    45  	}
    46  	settings := &enginepbccl.EncryptionSettings{}
    47  	if key == nil || key.Info.EncryptionType == enginepbccl.EncryptionType_Plaintext {
    48  		settings.EncryptionType = enginepbccl.EncryptionType_Plaintext
    49  		stream := &filePlainStream{}
    50  		return settings, stream, nil
    51  	}
    52  	settings.EncryptionType = key.Info.EncryptionType
    53  	settings.KeyId = key.Info.KeyId
    54  	settings.Nonce = make([]byte, ctrNonceSize)
    55  	_, err = rand.Read(settings.Nonce)
    56  	if err != nil {
    57  		return nil, nil, err
    58  	}
    59  	counterBytes := make([]byte, 4)
    60  	if _, err = rand.Read(counterBytes); err != nil {
    61  		return nil, nil, err
    62  	}
    63  	// Does not matter how we convert 4 random bytes into uint32
    64  	settings.Counter = binary.LittleEndian.Uint32(counterBytes)
    65  	ctrCS, err := newCTRBlockCipherStream(key, settings.Nonce, settings.Counter)
    66  	if err != nil {
    67  		return nil, nil, err
    68  	}
    69  	return settings, &fileCipherStream{bcs: ctrCS}, nil
    70  }
    71  
    72  // CreateExisting creates a FileStream for an existing file by looking up the key described by
    73  // settings in the key manager.
    74  func (c *FileCipherStreamCreator) CreateExisting(
    75  	settings *enginepbccl.EncryptionSettings,
    76  ) (FileStream, error) {
    77  	if settings == nil || settings.EncryptionType == enginepbccl.EncryptionType_Plaintext {
    78  		return &filePlainStream{}, nil
    79  	}
    80  	key, err := c.keyManager.GetKey(settings.KeyId)
    81  	if err != nil {
    82  		return nil, err
    83  	}
    84  	ctrCS, err := newCTRBlockCipherStream(key, settings.Nonce, settings.Counter)
    85  	if err != nil {
    86  		return nil, err
    87  	}
    88  	return &fileCipherStream{bcs: ctrCS}, nil
    89  }
    90  
    91  // FileStream encrypts/decrypts byte slices at arbitrary file offsets.
    92  //
    93  // There are two implementations: a noop filePlainStream and a fileCipherStream that wraps
    94  // a ctrBlockCipherStream. The ctrBlockCipherStream does AES in counter mode (CTR). CTR
    95  // allows us to encrypt/decrypt at arbitrary byte offsets in a file (including partial
    96  // blocks) without caring about what preceded the bytes.
    97  type FileStream interface {
    98  	// Encrypt encrypts the data to be written at fileOffset.
    99  	Encrypt(fileOffset int64, data []byte)
   100  	// Decrypt decrypts the data that has been read from fileOffset.
   101  	Decrypt(fileOffset int64, data []byte)
   102  }
   103  
   104  // Implements a noop FileStream.
   105  type filePlainStream struct{}
   106  
   107  func (s *filePlainStream) Encrypt(fileOffset int64, data []byte) {}
   108  func (s *filePlainStream) Decrypt(fileOffset int64, data []byte) {}
   109  
   110  // Implements a FileStream with AES-CTR.
   111  type fileCipherStream struct {
   112  	bcs *cTRBlockCipherStream
   113  }
   114  
   115  func (s *fileCipherStream) Encrypt(fileOffset int64, data []byte) {
   116  	if len(data) == 0 {
   117  		return
   118  	}
   119  	blockIndex := uint64(fileOffset / int64(ctrBlockSize))
   120  	blockOffset := int(fileOffset % int64(ctrBlockSize))
   121  	// TODO(sbhola): Use sync.Pool for these temporary buffers.
   122  	var buf struct {
   123  		dataScratch [ctrBlockSize]byte
   124  		ivScratch   [ctrBlockSize]byte
   125  	}
   126  	for len(data) > 0 {
   127  		// The num bytes that must be encrypted in this block.
   128  		byteCount := ctrBlockSize - blockOffset
   129  		if byteCount > len(data) {
   130  			// The data ends before the end of this block.
   131  			byteCount = len(data)
   132  		}
   133  		if byteCount < int(ctrBlockSize) {
   134  			// Need to copy into dataScratch, starting at blockOffset. NB: in CTR mode it does
   135  			// not matter what is contained in the other bytes in the block (the ones we are not
   136  			// initializing using this copy()).
   137  			copy(buf.dataScratch[blockOffset:blockOffset+byteCount], data[:byteCount])
   138  			s.bcs.transform(blockIndex, buf.dataScratch[:], buf.ivScratch[:])
   139  			// Copy the transformed data back into data.
   140  			copy(data[:byteCount], buf.dataScratch[blockOffset:blockOffset+byteCount])
   141  			blockOffset = 0
   142  		} else {
   143  			s.bcs.transform(blockIndex, data[:byteCount], buf.ivScratch[:])
   144  		}
   145  		blockIndex++
   146  		data = data[byteCount:]
   147  	}
   148  }
   149  
   150  // For CTR, decryption and encryption are the same
   151  func (s *fileCipherStream) Decrypt(fileOffset int64, data []byte) {
   152  	s.Encrypt(fileOffset, data)
   153  }
   154  
   155  // AES in CTR mode.
   156  type cTRBlockCipherStream struct {
   157  	key     *enginepbccl.SecretKey
   158  	nonce   [ctrNonceSize]byte
   159  	counter uint32
   160  
   161  	cBlock cipher.Block
   162  }
   163  
   164  func newCTRBlockCipherStream(
   165  	key *enginepbccl.SecretKey, nonce []byte, counter uint32,
   166  ) (*cTRBlockCipherStream, error) {
   167  	switch key.Info.EncryptionType {
   168  	case enginepbccl.EncryptionType_AES128_CTR:
   169  	case enginepbccl.EncryptionType_AES192_CTR:
   170  	case enginepbccl.EncryptionType_AES256_CTR:
   171  	default:
   172  		return nil, fmt.Errorf("unknown EncryptionType: %d", key.Info.EncryptionType)
   173  	}
   174  	stream := &cTRBlockCipherStream{key: key, counter: counter}
   175  	// Copy the nonce since the caller may overwrite it in the future.
   176  	copy(stream.nonce[:], nonce)
   177  	var err error
   178  	if stream.cBlock, err = aes.NewCipher(key.Key); err != nil {
   179  		return nil, err
   180  	}
   181  	if stream.cBlock.BlockSize() != ctrBlockSize {
   182  		return nil, fmt.Errorf("unexpected block size: %d", stream.cBlock.BlockSize())
   183  	}
   184  	return stream, nil
   185  }
   186  
   187  // For CTR, decryption and encryption are the same. data must have length equal to
   188  // the block size, and scratch must have length >= block size.
   189  func (s *cTRBlockCipherStream) transform(blockIndex uint64, data []byte, scratch []byte) {
   190  	iv := append(scratch[:0], s.nonce[:]...)
   191  	var blockCounter = uint32(uint64(s.counter) + blockIndex)
   192  	binary.BigEndian.PutUint32(iv[len(iv):len(iv)+4], blockCounter)
   193  	iv = iv[0 : len(iv)+4]
   194  	s.cBlock.Encrypt(iv, iv)
   195  	for i := 0; i < ctrBlockSize; i++ {
   196  		data[i] = data[i] ^ iv[i]
   197  	}
   198  }