Track DomStorageMap usage and enforce a quota on byte size.

BUG=106763


Review URL: http://codereview.chromium.org/9361062

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@123046 0039d316-1c4b-4281-b951-d872f2087c98

5 files changed, 116 insertions, 20 deletions

diff --git a/webkit/dom_storage/dom_storage_area.cc b/webkit/dom_storage/dom_storage_area.cc
index c152f3f..e394f64 100644
--- a/webkit/dom_storage/dom_storage_area.cc
+++ b/webkit/dom_storage/dom_storage_area.cc
@@ -5,15 +5,18 @@
 #include "webkit/dom_storage/dom_storage_area.h"
 #include "webkit/dom_storage/dom_storage_map.h"
 #include "webkit/dom_storage/dom_storage_namespace.h"
+#include "webkit/dom_storage/dom_storage_types.h"
 
 namespace dom_storage {
 
 DomStorageArea::DomStorageArea(
     int64 namespace_id, const GURL& origin,
     const FilePath& directory, DomStorageTaskRunner* task_runner)
-    : namespace_id_(namespace_id), origin_(origin),
-      directory_(directory), task_runner_(task_runner),
-      map_(new DomStorageMap()) {
+    : namespace_id_(namespace_id),
+      origin_(origin),
+      directory_(directory),
+      task_runner_(task_runner),
+      map_(new DomStorageMap(kPerAreaQuota)) {
 }
 
 DomStorageArea::~DomStorageArea() {
@@ -50,7 +53,7 @@ bool DomStorageArea::RemoveItem(
 bool DomStorageArea::Clear() {
   if (map_->Length() == 0)
     return false;
-  map_ = new DomStorageMap();
+  map_ = new DomStorageMap(kPerAreaQuota);
   return true;
 }
 
diff --git a/webkit/dom_storage/dom_storage_map.cc b/webkit/dom_storage/dom_storage_map.cc
index 04ed753..aa34752 100644
--- a/webkit/dom_storage/dom_storage_map.cc
+++ b/webkit/dom_storage/dom_storage_map.cc
@@ -4,9 +4,32 @@
 
 #include "webkit/dom_storage/dom_storage_map.h"
 
+#include "base/logging.h"
+
+namespace {
+
+size_t size_of_item(const string16& key, const string16& value) {
+  return (key.length() + value.length()) * sizeof(char16);
+}
+
+size_t CountBytes(const dom_storage::ValuesMap& values) {
+  if (values.size() == 0)
+    return 0;
+
+  size_t count = 0;
+  dom_storage::ValuesMap::const_iterator it = values.begin();
+  for (; it != values.end(); ++it)
+    count += size_of_item(it->first, it->second.string());
+  return count;
+}
+
+}  // namespace
+
 namespace dom_storage {
 
-DomStorageMap::DomStorageMap() {
+DomStorageMap::DomStorageMap(size_t quota)
+    : bytes_used_(0),
+      quota_(quota) {
   ResetKeyIterator();
 }
 
@@ -41,14 +64,22 @@ NullableString16 DomStorageMap::GetItem(const string16& key) const {
 bool DomStorageMap::SetItem(
     const string16& key, const string16& value,
     NullableString16* old_value) {
-  // TODO(benm): check quota and return false if exceeded
   ValuesMap::const_iterator found = values_.find(key);
   if (found == values_.end())
     *old_value = NullableString16(true);
   else
     *old_value = found->second;
+
+  size_t old_item_size = old_value->is_null() ?
+      0 : size_of_item(key, old_value->string());
+  size_t new_size = bytes_used_ - old_item_size + size_of_item(key, value);
+  if (new_size > quota_)
+    return false;
+
   values_[key] = NullableString16(value, false);
   ResetKeyIterator();
+  bytes_used_ -= old_item_size;
+  bytes_used_ += size_of_item(key, value);
   return true;
 }
 
@@ -61,17 +92,25 @@ bool DomStorageMap::RemoveItem(
   *old_value = found->second.string();
   values_.erase(found);
   ResetKeyIterator();
+  bytes_used_ -= size_of_item(key, *old_value);
   return true;
 }
 
-void DomStorageMap::SwapValues(ValuesMap* values) {
+bool DomStorageMap::SwapValues(ValuesMap* values) {
+  size_t new_size = CountBytes(*values);
+  if (new_size > quota_)
+    return false;
   values_.swap(*values);
+  bytes_used_ = new_size;
   ResetKeyIterator();
+  return true;
 }
 
 DomStorageMap* DomStorageMap::DeepCopy() const {
-  DomStorageMap* copy = new DomStorageMap;
+  DCHECK(CountBytes(values_) <= quota_);
+  DomStorageMap* copy = new DomStorageMap(quota_);
   copy->values_ = values_;
+  copy->bytes_used_ = bytes_used_;
   copy->ResetKeyIterator();
   return copy;
 }
diff --git a/webkit/dom_storage/dom_storage_map.h b/webkit/dom_storage/dom_storage_map.h
index 5610824..d8151e6 100644
--- a/webkit/dom_storage/dom_storage_map.h
+++ b/webkit/dom_storage/dom_storage_map.h
@@ -13,18 +13,15 @@
 #include "base/string16.h"
 #include "webkit/dom_storage/dom_storage_types.h"
 
-class FilePath;
-class GURL;
-
 namespace dom_storage {
 
 // A wrapper around a std::map that adds refcounting and
-// imposes a limit on the total byte size of the keys/values.
+// tracks the size in bytes of the keys/values, enforcing a quota.
 // See class comments for DomStorageContext for a larger overview.
 class DomStorageMap
     : public base::RefCountedThreadSafe<DomStorageMap> {
  public:
-  DomStorageMap();
+  DomStorageMap(size_t quota);
 
   unsigned Length() const;
   NullableString16 Key(unsigned index);
@@ -33,9 +30,14 @@ class DomStorageMap
                NullableString16* old_value);
   bool RemoveItem(const string16& key, string16* old_value);
 
-  void SwapValues(ValuesMap* map);
+  // Replaces values_ with |map|. Returns true on success, false
+  // if the swap was prevented because |map| would exceed this
+  // DomStorageMap's quota_.
+  bool SwapValues(ValuesMap* map);
   DomStorageMap* DeepCopy() const;
 
+  size_t bytes_used() const { return bytes_used_; }
+
  private:
   friend class base::RefCountedThreadSafe<DomStorageMap>;
   ~DomStorageMap();
@@ -45,9 +47,10 @@ class DomStorageMap
   ValuesMap values_;
   ValuesMap::const_iterator key_iterator_;
   unsigned last_key_index_;
-  // TODO(benm): track usage and enforce a quota limit.
+  size_t bytes_used_;
+  size_t quota_;
 };
 
 }  // namespace dom_storage
 
-#endif  // WEBKIT_DOM_STORAGE_DOM_STORAGE_AREA_H_
+#endif  // WEBKIT_DOM_STORAGE_DOM_STORAGE_MAP_H_
diff --git a/webkit/dom_storage/dom_storage_map_unittest.cc b/webkit/dom_storage/dom_storage_map_unittest.cc
index 81043fd..67c5bb5 100644
--- a/webkit/dom_storage/dom_storage_map_unittest.cc
+++ b/webkit/dom_storage/dom_storage_map_unittest.cc
@@ -11,10 +11,17 @@ namespace dom_storage {
 TEST(DomStorageMapTest, DomStorageMapBasics) {
   const string16 kKey(ASCIIToUTF16("key"));
   const string16 kValue(ASCIIToUTF16("value"));
+  const size_t kValueBytes = kValue.size() * sizeof(char16);
+  const size_t kItemBytes =
+      (kKey.size() + kValue.size()) * sizeof(char16);
   const string16 kKey2(ASCIIToUTF16("key2"));
+  const size_t kKey2Bytes = kKey2.size() * sizeof(char16);
   const string16 kValue2(ASCIIToUTF16("value2"));
+  const size_t kItem2Bytes =
+      (kKey2.size() + kValue2.size()) * sizeof(char16);
+  const size_t kQuota = 1024;  // 1K quota for this test.
 
-  scoped_refptr<DomStorageMap> map(new DomStorageMap);
+  scoped_refptr<DomStorageMap> map(new DomStorageMap(kQuota));
   string16 old_value;
   NullableString16 old_nullable_value;
   ValuesMap swap;
@@ -26,9 +33,11 @@ TEST(DomStorageMapTest, DomStorageMapBasics) {
   EXPECT_TRUE(map->Key(100).is_null());
   EXPECT_TRUE(map->GetItem(kKey).is_null());
   EXPECT_FALSE(map->RemoveItem(kKey, &old_value));
+  EXPECT_EQ(0u, map->bytes_used());
   copy = map->DeepCopy();
   EXPECT_EQ(0u, copy->Length());
-  map->SwapValues(&swap);
+  EXPECT_EQ(0u, copy->bytes_used());
+  EXPECT_TRUE(map->SwapValues(&swap));
   EXPECT_TRUE(swap.empty());
 
   // Check the behavior of a map containing some values.
@@ -39,19 +48,24 @@ TEST(DomStorageMapTest, DomStorageMapBasics) {
   EXPECT_TRUE(map->Key(1).is_null());
   EXPECT_EQ(kValue, map->GetItem(kKey).string());
   EXPECT_TRUE(map->GetItem(kKey2).is_null());
+  EXPECT_EQ(kItemBytes, map->bytes_used());
   EXPECT_TRUE(map->RemoveItem(kKey, &old_value));
+  EXPECT_EQ(0u, map->bytes_used());
   EXPECT_EQ(kValue, old_value);
   old_value.clear();
 
   EXPECT_TRUE(map->SetItem(kKey, kValue, &old_nullable_value));
   EXPECT_TRUE(map->SetItem(kKey2, kValue, &old_nullable_value));
   EXPECT_TRUE(old_nullable_value.is_null());
+  EXPECT_EQ(kItemBytes + kKey2Bytes + kValueBytes, map->bytes_used());
   EXPECT_TRUE(map->SetItem(kKey2, kValue2, &old_nullable_value));
+  EXPECT_EQ(kItemBytes + kItem2Bytes, map->bytes_used());
   EXPECT_EQ(kValue, old_nullable_value.string());
   EXPECT_EQ(2u, map->Length());
   EXPECT_EQ(kKey, map->Key(0).string());
   EXPECT_EQ(kKey2, map->Key(1).string());
   EXPECT_EQ(kKey, map->Key(0).string());
+  EXPECT_EQ(kItemBytes + kItem2Bytes, map->bytes_used());
 
   copy = map->DeepCopy();
   EXPECT_EQ(2u, copy->Length());
@@ -60,10 +74,47 @@ TEST(DomStorageMapTest, DomStorageMapBasics) {
   EXPECT_EQ(kKey, copy->Key(0).string());
   EXPECT_EQ(kKey2, copy->Key(1).string());
   EXPECT_TRUE(copy->Key(2).is_null());
+  EXPECT_EQ(kItemBytes + kItem2Bytes, copy->bytes_used());
 
-  map->SwapValues(&swap);
+  EXPECT_TRUE(map->SwapValues(&swap));
   EXPECT_EQ(2ul, swap.size());
   EXPECT_EQ(0u, map->Length());
+  EXPECT_EQ(0u, map->bytes_used());
+}
+
+TEST(DomStorageMapTest, EnforcesQuota) {
+  const string16 kKey = ASCIIToUTF16("test_key");
+  const string16 kValue = ASCIIToUTF16("test_value");
+  const string16 kKey2 = ASCIIToUTF16("test_key_2");
+
+  // A 50 byte quota is too small to hold both keys, so we
+  // should see the DomStorageMap enforcing it.
+  const size_t kQuota = 50;
+
+  string16 old_value;
+  NullableString16 old_nullable_value;
+
+  scoped_refptr<DomStorageMap> map(new DomStorageMap(kQuota));
+  EXPECT_TRUE(map->SetItem(kKey, kValue, &old_nullable_value));
+  EXPECT_FALSE(map->SetItem(kKey2, kValue, &old_nullable_value));
+  EXPECT_EQ(1u, map->Length());
+
+  EXPECT_TRUE(map->RemoveItem(kKey, &old_value));
+  EXPECT_EQ(kValue, old_value);
+  EXPECT_EQ(0u, map->Length());
+  EXPECT_TRUE(map->SetItem(kKey2, kValue, &old_nullable_value));
+  EXPECT_EQ(1u, map->Length());
+
+  ValuesMap swap;
+  EXPECT_TRUE(map->SwapValues(&swap));
+  EXPECT_EQ(0u, map->Length());
+
+  swap[kKey] = NullableString16(kValue, false);
+  swap[kKey2] = NullableString16(kValue, false);
+
+  // swap is now too big to fit in the map, the swap should fail.
+  EXPECT_FALSE(map->SwapValues(&swap));
+  EXPECT_EQ(0u, map->Length());
 }
 
 }  // namespace dom_storage
diff --git a/webkit/dom_storage/dom_storage_types.h b/webkit/dom_storage/dom_storage_types.h
index 363cbe4..70ee69c 100644
--- a/webkit/dom_storage/dom_storage_types.h
+++ b/webkit/dom_storage/dom_storage_types.h
@@ -15,7 +15,7 @@
 namespace dom_storage {
 
 // The quota for each storage area. Suggested by the spec.
-const int kPerAreaQuota = 5 * 1024 * 1024;
+const size_t kPerAreaQuota = 5 * 1024 * 1024;
 
 // Value to indicate the localstorage namespace vs non-zero
 // values for sessionstorage namespaces.