aes_cryptor.h

// Copyright 2016 Google LLC. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd
 
#ifndef PACKAGER_MEDIA_BASE_AES_CRYPTOR_H_
#define PACKAGER_MEDIA_BASE_AES_CRYPTOR_H_
 
#include <memory>
#include <string>
#include <vector>
 
#include <mbedtls/cipher.h>
 
#include <packager/macros/classes.h>
#include <packager/media/base/fourccs.h>
 
namespace shaka {
namespace media {
 
// AES cryptor interface. Inherited by various AES encryptor and decryptor
// implementations.
class AesCryptor {
 public:
  enum ConstantIvFlag {
    kUseConstantIv,
    kDontUseConstantIv,
  };
 
  explicit AesCryptor(ConstantIvFlag constant_iv_flag);
  virtual ~AesCryptor();
 
  virtual bool InitializeWithIv(const std::vector<uint8_t>& key,
                                const std::vector<uint8_t>& iv) = 0;
 
  virtual size_t RequiredOutputSize(size_t plaintext_size) {
    return plaintext_size;
  }
 
  bool Crypt(const std::vector<uint8_t>& text,
             std::vector<uint8_t>* crypt_text);
  bool Crypt(const std::string& text, std::string* crypt_text);
  bool Crypt(const uint8_t* text, size_t text_size, uint8_t* crypt_text) {
    size_t crypt_text_size = text_size;
    return Crypt(text, text_size, crypt_text, &crypt_text_size);
  }
  bool Crypt(const uint8_t* text,
             size_t text_size,
             uint8_t* crypt_text,
             size_t* crypt_text_size) {
    if (constant_iv_flag_ == kUseConstantIv)
      SetIvInternal();
    else
      num_crypt_bytes_ += text_size;
    return CryptInternal(text, text_size, crypt_text, crypt_text_size);
  }
 
  bool SetIv(const std::vector<uint8_t>& iv);
 
  void UpdateIv();
 
  const std::vector<uint8_t>& iv() const { return iv_; }
 
  bool use_constant_iv() const { return constant_iv_flag_ == kUseConstantIv; }
 
  static bool GenerateRandomIv(FourCC protection_scheme,
                               std::vector<uint8_t>* iv);
 
 protected:
  enum CipherMode {
    kCtrMode,
    kCbcMode,
  };
 
  // mbedTLS cipher context.
  mbedtls_cipher_context_t cipher_ctx_;
 
  bool SetupCipher(size_t key_size, CipherMode mode);
 
 private:
  // Internal implementation of crypt function.
  // |text| points to the input text.
  // |text_size| is the size of input text.
  // |crypt_text| points to the output encrypted or decrypted text, depends on
  // whether it is an encryption or decryption. |text| and |crypt_text| can
  // point to the same address for in place encryption/decryption.
  // |crypt_text_size| contains the size of |crypt_text| and it will be updated
  // to contain the actual encrypted/decrypted size for |crypt_text| on success.
  // Return false if the input |crypt_text_size| is not large enough to hold the
  // output |crypt_text| or if there is any error in encryption/decryption.
  virtual bool CryptInternal(const uint8_t* text,
                             size_t text_size,
                             uint8_t* crypt_text,
                             size_t* crypt_text_size) = 0;
 
  // Internal implementation of SetIv, which setup internal iv.
  virtual void SetIvInternal() = 0;
 
  // |size| specifies the input text size.
  // Return the number of padding bytes needed.
  // Note: No paddings should be needed except for pkcs5-cbc encryptor.
  virtual size_t NumPaddingBytes(size_t size) const;
 
  // Indicates whether a constant iv is used. Internal iv will be reset to
  // |iv_| before calling Crypt if that is the case.
  const ConstantIvFlag constant_iv_flag_;
  // Initialization vector from by SetIv or InitializeWithIv, with size 8 or 16
  // bytes.
  std::vector<uint8_t> iv_;
  // Tracks number of crypt bytes. It is used to calculate how many blocks
  // should iv advance in UpdateIv(). It will be reset to 0 after iv is updated.
  size_t num_crypt_bytes_;
 
  DISALLOW_COPY_AND_ASSIGN(AesCryptor);
};
 
}  // namespace media
}  // namespace shaka
 
#endif  // PACKAGER_MEDIA_BASE_AES_CRYPTOR_H_