// Copyright 2020 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine Master
// License Agreement.

#include "oemcrypto_entitled_key_session.h"

namespace wvoec_ref {

const Key* EntitledKeySession::GetEntitlementKey(
    const KeyId& content_key_id) const {
  if (entitlement_key_map_.find(content_key_id) == entitlement_key_map_.end()) {
    return nullptr;
  }
  return entitlement_key_map_.at(content_key_id);
}

EntitledKey* EntitledKeySession::GetContentKey(
    const KeyId& content_key_id) const {
  if (content_key_map_.find(content_key_id) == content_key_map_.end()) {
    return nullptr;
  }
  return content_key_map_.at(content_key_id).get();
}

bool EntitledKeySession::AddOrUpdateContentKey(
    Key* entitlement_key, const KeyId& content_key_id,
    std::unique_ptr<EntitledKey> content_key) {
  if (entitlement_key == nullptr || content_key_id.empty() ||
      content_key == nullptr) {
    return false;
  }
  // Remove the entry if |entitlement_key| already exists. Each entitlement key
  // can only be referred by one content key within an entitled key session.
  for (auto const& content_id_entitlement : entitlement_key_map_) {
    if (content_id_entitlement.second == entitlement_key) {
      RemoveContentKey(content_id_entitlement.first);
      break;
    }
  }
  // |content_key_id| must be unique.
  if (content_key_map_.find(content_key_id) != content_key_map_.end()) {
    return false;
  }
  content_key_map_[content_key_id] = std::move(content_key);
  entitlement_key_map_[content_key_id] = entitlement_key;
  return true;
}

bool EntitledKeySession::RemoveContentKey(const KeyId& content_key_id) {
  if (content_key_map_.find(content_key_id) == content_key_map_.end()) {
    return false;
  }
  content_key_map_.erase(content_key_id);
  entitlement_key_map_.erase(content_key_id);
  return true;
}

/***************************************/

EntitledKeySession* EntitledKeySessionTable::CreateEntitledKeySession(
    SessionId oec_sid, SessionId key_sid) {
  std::unique_lock<std::mutex> lock(session_lock_);
  // The given |key_sid| already exists.
  if (entitled_key_sessions_.find(key_sid) != entitled_key_sessions_.end()) {
    return nullptr;
  }
  entitled_key_sessions_[key_sid] =
      std::unique_ptr<EntitledKeySession>(new EntitledKeySession(key_sid));

  oec_session_to_key_sessions_[oec_sid].insert(oec_sid);
  key_session_to_oec_session_[key_sid] = oec_sid;

  return entitled_key_sessions_[key_sid].get();
}

EntitledKeySession* EntitledKeySessionTable::FindEntitledKeySession(
    SessionId key_sid) {
  std::unique_lock<std::mutex> lock(session_lock_);
  if (entitled_key_sessions_.find(key_sid) == entitled_key_sessions_.end()) {
    return nullptr;
  }
  return entitled_key_sessions_.at(key_sid).get();
}

void EntitledKeySessionTable::RemoveEntitledKeySession(SessionId key_sid) {
  std::unique_lock<std::mutex> lock(session_lock_);
  if (entitled_key_sessions_.find(key_sid) == entitled_key_sessions_.end()) {
    return;
  }

  SessionId oec_sid = key_session_to_oec_session_[key_sid];
  key_session_to_oec_session_.erase(key_sid);
  oec_session_to_key_sessions_[oec_sid].erase(key_sid);

  entitled_key_sessions_.erase(key_sid);
}

SessionId EntitledKeySessionTable::GetOEMCryptoSessionId(SessionId key_sid) {
  std::unique_lock<std::mutex> lock(session_lock_);
  if (key_session_to_oec_session_.find(key_sid) ==
      key_session_to_oec_session_.end()) {
    return 0;
  }
  return key_session_to_oec_session_.at(key_sid);
}

void EntitledKeySessionTable::OnDestroyOEMCryptoSession(SessionId oec_sid) {
  std::unique_lock<std::mutex> lock(session_lock_);
  if (oec_session_to_key_sessions_.find(oec_sid) ==
      oec_session_to_key_sessions_.end()) {
    return;
  }
  for (auto entitled_key_sid : oec_session_to_key_sessions_[oec_sid]) {
    key_session_to_oec_session_.erase(entitled_key_sid);
    entitled_key_sessions_.erase(entitled_key_sid);
  }
  oec_session_to_key_sessions_.erase(oec_sid);
}

}  // namespace wvoec_ref