From e9dfca12190dbf72e3fd1c81f01fdeb593635e1c Mon Sep 17 00:00:00 2001
From: Fred Gylys-Colwell <fredgc@google.com>
Date: Tue, 19 May 2020 10:58:05 -0700
Subject: [PATCH] The Entitlement key map was not being cleaned up.

Merge from Widevine repo of http://go/wvgerrit/100385

The map now contains unique_ptr instead of raw pointers
to ensure the memory is released.

Bug: 156780432 OEMCrypto Fuzzing: Fix OEMCrypto Memory Leak.
Test: oemcrypto reference code only
Change-Id: I78054f9207399f052d6e4bfdfa96824f6e050bac
---
 .../ref/src/oemcrypto_session_key_table.cpp   | 19 +++++--------------
 .../ref/src/oemcrypto_session_key_table.h     | 11 ++++++-----
 2 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/libwvdrmengine/oemcrypto/ref/src/oemcrypto_session_key_table.cpp b/libwvdrmengine/oemcrypto/ref/src/oemcrypto_session_key_table.cpp
index e7131975..00c16ed8 100644
--- a/libwvdrmengine/oemcrypto/ref/src/oemcrypto_session_key_table.cpp
+++ b/libwvdrmengine/oemcrypto/ref/src/oemcrypto_session_key_table.cpp
@@ -11,17 +11,9 @@
 
 namespace wvoec_ref {
 
-SessionKeyTable::~SessionKeyTable() {
-  for (KeyMap::iterator i = keys_.begin(); i != keys_.end(); ++i) {
-    if (nullptr != i->second) {
-      delete i->second;
-    }
-  }
-}
-
 bool SessionKeyTable::Insert(const KeyId key_id, const Key& key_data) {
   if (keys_.find(key_id) != keys_.end()) return false;
-  keys_[key_id] = new Key(key_data);
+  keys_[key_id] = std::unique_ptr<Key>(new Key(key_data));
   return true;
 }
 
@@ -29,12 +21,11 @@ Key* SessionKeyTable::Find(const KeyId key_id) {
   if (keys_.find(key_id) == keys_.end()) {
     return nullptr;
   }
-  return keys_[key_id];
+  return keys_[key_id].get();
 }
 
 void SessionKeyTable::Remove(const KeyId key_id) {
   if (keys_.find(key_id) != keys_.end()) {
-    delete keys_[key_id];
     keys_.erase(key_id);
   }
 }
@@ -49,7 +40,7 @@ bool EntitlementKeyTable::Insert(const KeyId key_id, const Key& key_data) {
   // |key_id| and |key_data| are for an entitlement key. Insert a new
   // entitlement key entry.
   if (keys_.find(key_id) != keys_.end()) return false;
-  keys_[key_id] = new EntitlementKey(key_data);
+  keys_[key_id] = std::unique_ptr<EntitlementKey>(new EntitlementKey(key_data));
   // If this is a new insertion, we don't have a content key assigned yet.
   return true;
 }
@@ -65,7 +56,7 @@ Key* EntitlementKeyTable::Find(const KeyId key_id) {
   if (keys_.find(it->second) == keys_.end()) {
     return nullptr;
   }
-  return keys_[it->second];
+  return keys_[it->second].get();
 }
 
 void EntitlementKeyTable::Remove(const KeyId key_id) {
@@ -108,7 +99,7 @@ EntitlementKey* EntitlementKeyTable::GetEntitlementKey(
   if (it == keys_.end()) {
     return nullptr;
   }
-  return it->second;
+  return it->second.get();
 }
 
 }  // namespace wvoec_ref
diff --git a/libwvdrmengine/oemcrypto/ref/src/oemcrypto_session_key_table.h b/libwvdrmengine/oemcrypto/ref/src/oemcrypto_session_key_table.h
index c84096be..ad74b873 100644
--- a/libwvdrmengine/oemcrypto/ref/src/oemcrypto_session_key_table.h
+++ b/libwvdrmengine/oemcrypto/ref/src/oemcrypto_session_key_table.h
@@ -9,6 +9,7 @@
 
 #include <stdint.h>
 #include <map>
+#include <memory>
 #include <vector>
 
 #include "disallow_copy_and_assign.h"
@@ -23,18 +24,18 @@ class UsageTable;
 class UsageTableEntry;
 
 typedef std::vector<uint8_t> KeyId;
-typedef std::map<KeyId, Key*> KeyMap;
-typedef std::map<KeyId, EntitlementKey*> EntitlementKeyMap;
+typedef std::map<KeyId, std::unique_ptr<Key>> KeyMap;
+typedef std::map<KeyId, std::unique_ptr<EntitlementKey>> EntitlementKeyMap;
 
 // SessionKeyTable holds the keys for the current session
 class SessionKeyTable {
  public:
   SessionKeyTable() {}
-  ~SessionKeyTable();
+  ~SessionKeyTable() {}
 
   bool Insert(const KeyId key_id, const Key& key_data);
   Key* Find(const KeyId key_id);
-  Key* FirstKey() { return keys_.begin()->second; }
+  Key* FirstKey() { return keys_.begin()->second.get(); }
   void Remove(const KeyId key_id);
   void UpdateDuration(const KeyControlBlock& control);
   size_t size() const { return keys_.size(); }
@@ -53,7 +54,7 @@ class EntitlementKeyTable {
   ~EntitlementKeyTable() {}
   bool Insert(const KeyId key_id, const Key& key_data);
   Key* Find(const KeyId key_id);
-  Key* FirstKey() { return keys_.begin()->second; }
+  Key* FirstKey() { return keys_.begin()->second.get(); }
   void Remove(const KeyId key_id);
   void UpdateDuration(const KeyControlBlock& control);
   size_t size() const { return contentid_to_entitlementid_.size(); }