YugabyteDB (2.13.1.0-b60, 21121d69985fbf76aa6958d8f04a9bfa936293b5)

// Copyright (c) YugaByte, Inc.

//

// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except

// in compliance with the License.  You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software distributed under the License

// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express

// or implied.  See the License for the specific language governing permissions and limitations

// under the License.

//

#include "yb/rocksutil/rocksdb_encrypted_file_factory.h"

#include "yb/encryption/cipher_stream.h"

#include "yb/encryption/encrypted_file.h"

#include "yb/encryption/encryption_util.h"

#include "yb/encryption/header_manager.h"

namespace yb {

// An encrypted file implementation for sequential reads.

class EncryptedSequentialFile : public SequentialFileWrapper {

 public:

  static Status Create(std::unique_ptr<rocksdb::SequentialFile>* result,

                       encryption::HeaderManager* header_manager,

                       std::unique_ptr<rocksdb::SequentialFile> underlying_seq,

                       std::unique_ptr<rocksdb::RandomAccessFile> underlying_ra) {

    result->reset();

    std::unique_ptr<encryption::BlockAccessCipherStream> stream;

    uint32_t header_size;

    const auto file_encrypted = VERIFY_RESULT(encryption::GetEncryptionInfoFromFile<uint8_t>(

        header_manager, underlying_ra.get(), &stream, &header_size));

    if (!file_encrypted) {

      *result = std::move(underlying_seq);

      return Status::OK();

    }

    RETURN_NOT_OK(underlying_seq->Skip(header_size));

    *result = std::make_unique<EncryptedSequentialFile>(

        std::move(underlying_seq), std::move(stream));

    return Status::OK();

  }

  // Default constructor.

  EncryptedSequentialFile(std::unique_ptr<rocksdb::SequentialFile> file,

                          std::unique_ptr<encryption::BlockAccessCipherStream> stream)

      : SequentialFileWrapper(std::move(file)), stream_(std::move(stream)) {}

  ~EncryptedSequentialFile() {}

  Status Read(size_t n, Slice* result, uint8_t* scratch) override {

    if (!scratch) {

      return STATUS(InvalidArgument, "scratch argument is null.");

    }

    uint8_t* buf = static_cast<uint8_t*>(encryption::EncryptionBuffer::Get()->GetBuffer(n));

    RETURN_NOT_OK(SequentialFileWrapper::Read(n, result, buf));

    RETURN_NOT_OK(stream_->Decrypt(offset_, *result, scratch));

    *result = Slice(scratch, result->size());

    offset_ += result->size();

    return Status::OK();

  }

  Status Skip(uint64_t n) override {

    RETURN_NOT_OK(SequentialFileWrapper::Skip(n));

    offset_ += n;

    return Status::OK();

  }

 private:

  std::unique_ptr<encryption::BlockAccessCipherStream> stream_;

  uint64_t offset_ = 0;

};

// An encrypted file implementation for a writable file.

class RocksDBEncryptedWritableFile : public rocksdb::WritableFileWrapper {

 public:

  static Status Create(std::unique_ptr<rocksdb::WritableFile>* result,

                       encryption::HeaderManager* header_manager,

                       std::unique_ptr<rocksdb::WritableFile> underlying) {

    return encryption::CreateWritableFile<RocksDBEncryptedWritableFile>(

        result, header_manager, std::move(underlying));

  }

  // Default constructor.

  RocksDBEncryptedWritableFile(std::unique_ptr<rocksdb::WritableFile> file,

                               std::unique_ptr<encryption::BlockAccessCipherStream> stream,

                               uint32_t header_size)

      : WritableFileWrapper(std::move(file)), stream_(std::move(stream)),

        header_size_(header_size) {}

  ~RocksDBEncryptedWritableFile() {}

  Status Append(const Slice& data) override {

    if (data.size() > 0) {

      char* buf = static_cast<char*>(encryption::EncryptionBuffer::Get()->GetBuffer(data.size()));

      RETURN_NOT_OK(stream_->Encrypt(GetFileSize() - header_size_, data, buf));

      RETURN_NOT_OK(WritableFileWrapper::Append(Slice(buf, data.size())));

    }

    return Status::OK();

  }

 private:

  std::unique_ptr<encryption::BlockAccessCipherStream> stream_;

  uint32_t header_size_;

};

// EncryptedEnv implements an Env wrapper that adds encryption to files stored on disk.

class RocksDBEncryptedFileFactory : public rocksdb::RocksDBFileFactoryWrapper {

 public:

  explicit RocksDBEncryptedFileFactory(

      rocksdb::RocksDBFileFactory* factory,

      std::unique_ptr<encryption::HeaderManager> header_manager) :

      RocksDBFileFactoryWrapper(factory),

      header_manager_(std::move(header_manager)) {

    VLOG
(1) << "Created RocksDB encrypted env"0
;

  }

  ~RocksDBEncryptedFileFactory() {}

  // NewRandomAccessFile opens a file for random read access.

  Status NewRandomAccessFile(const std::string& fname,

                             std::unique_ptr<rocksdb::RandomAccessFile>* result,

                             const rocksdb::EnvOptions& options) override {

    std::unique_ptr<rocksdb::RandomAccessFile> underlying;

    RETURN_NOT_OK(RocksDBFileFactoryWrapper::NewRandomAccessFile(fname, &underlying, options));

    return encryption::EncryptedRandomAccessFile::Create(

        result, header_manager_.get(), std::move(underlying));

  }

  Status NewWritableFile(const std::string& fname, std::unique_ptr<rocksdb::WritableFile>* result,

                         const rocksdb::EnvOptions& options) override {

    std::unique_ptr<rocksdb::WritableFile> underlying;

    RETURN_NOT_OK(RocksDBFileFactoryWrapper::NewWritableFile(fname, &underlying, options));

    return RocksDBEncryptedWritableFile::Create(

        result, header_manager_.get(), std::move(underlying));

  }

  Status NewSequentialFile(const std::string& fname,

                           std::unique_ptr<rocksdb::SequentialFile>* result,

                           const rocksdb::EnvOptions& options) override {

    // Open file using underlying Env implementation

    std::unique_ptr<rocksdb::RandomAccessFile> underlying_ra;

    RETURN_NOT_OK(RocksDBFileFactoryWrapper::NewRandomAccessFile(fname, &underlying_ra, options));

    std::unique_ptr<rocksdb::SequentialFile> underlying_seq;

    RETURN_NOT_OK(RocksDBFileFactoryWrapper::NewSequentialFile(fname, &underlying_seq, options));

    return EncryptedSequentialFile::Create(

        result, header_manager_.get(), std::move(underlying_seq), std::move(underlying_ra));

  }

  Status GetFileSize(const std::string& fname, uint64_t* size) override {

    RETURN_NOT_OK(RocksDBFileFactoryWrapper::GetFileSize(fname, size));

    // Determine if file is encrypted.

    std::unique_ptr<rocksdb::SequentialFile> underlying;

    RETURN_NOT_OK(RocksDBFileFactoryWrapper::NewSequentialFile(

        fname, &underlying, rocksdb::EnvOptions()));

    *size -= VERIFY_RESULT(encryption::GetHeaderSize(underlying.get(), header_manager_.get()));

    return Status::OK();

  }

  bool IsPlainText() const override {

    return false;

  }

 private:

  std::unique_ptr<encryption::HeaderManager> header_manager_;

};

std::unique_ptr<rocksdb::Env> NewRocksDBEncryptedEnv(

    std::unique_ptr<encryption::HeaderManager> header_manager) {

  auto file_factory = rocksdb::Env::DefaultFileFactory();

  auto encrypted_file_factory = std::make_unique<RocksDBEncryptedFileFactory>(

      file_factory, std::move(header_manager));

  auto encrypted_env = rocksdb::Env::NewRocksDBDefaultEnv(std::move(encrypted_file_factory));

  return encrypted_env;

}

} // namespace yb