Support for splitting sync storage items into chunks, to allow synchronization of big policies across devices.

1 files changed, 132 insertions, 31 deletions

diff --git a/src/common/Storage.js b/src/common/Storage.js
index 0e8181e..c534bfb 100644
--- a/src/common/Storage.js
+++ b/src/common/Storage.js
@@ -1,64 +1,165 @@
-var Storage = {
+"use strict";
+var Storage = (() => {
 
-  async safeOp(op, type, keys) {
+  let chunksKey = k => `${k}/CHUNKS`;
+
+  async function safeOp(op, type, keys) {
     let sync = type === "sync";
-    if (sync && op === "get") {
-      let localFallback = await this.localFallback();
-      if (localFallback.size) {
-        for (let k of Array.isArray(keys) ? keys : [keys]) {
-          if (localFallback.has(k)) {
-            type = "local";
-            break;
+
+    try {
+      if (sync) {
+        let remove = op === "remove";
+        if (remove || op === "get") {
+          keys = [].concat(keys); // don't touch the passed argument
+          let mergeResults = {};
+          let localFallback = await getLocalFallback();
+          if (localFallback.size) {
+            let localKeys = keys.filter(k => localFallback.has(k));
+            if (localKeys.length) {
+              if (remove) {
+                await browser.storage.local.remove(localKeys);
+                for (let k of localKeys) {
+                  localFallback.delete(k);
+                }
+                await setLocalFallback(localFallback);
+              } else {
+                mergeResults = await browser.storage.local.get(localKeys);
+              }
+              keys = keys.filter(k => !localFallback.has(k));
+            }
+          }
+
+          if (keys.length) { // we may not have non-fallback keys anymore
+            let chunkCounts = Object.entries(await browser.storage.sync.get(
+                keys.map(chunksKey)))
+                  .map(([k, count]) => [k.split("/")[0], count]);
+            if (chunkCounts.length) {
+              let chunkedKeys = [];
+              for (let [k, count] of chunkCounts) {
+                // prepare to fetch all the chunks at once
+                while (count-- > 0) chunkedKeys.push(`${k}/${count}`);
+              }
+              if (remove) {
+                let doomedKeys = keys
+                  .concat(chunkCounts.map(([k, count]) => chunksKey(k)))
+                  .concat(chunkedKeys);
+                return await browser.storage.sync.remove(doomedKeys);
+              } else {
+                let chunks = await browser.storage.sync.get(chunkedKeys);
+                for (let [k, count] of chunkCounts) {
+                  let orderedChunks = [];
+                  for (let j = 0; j < count; j++) {
+                    orderedChunks.push(chunks[`${k}/${j}`]);
+                  }
+                  let whole = orderedChunks.join('');
+                  try {
+                    mergeResults[k] = JSON.parse(whole);
+                    keys.splice(keys.indexOf(k), 1); // remove from "main" keys
+                  } catch (e) {
+                    error(e, "Could not parse chunked storage key %s (%s).", k, whole);
+                  }
+                }
+              }
+            }
+          }
+          return keys.length ?
+            Object.assign(mergeResults, await browser.storage.sync[op](keys))
+            : mergeResults;
+        } else if (op === "set") {
+          keys = Object.assign({}, keys); // don't touch the passed argument
+          const MAX_ITEM_SIZE = 4096;
+          // Firefox Sync's max object BYTEs size is 16384, Chrome's 8192.
+          // Rather than mesuring actual bytes, we play it safe by halving then
+          // lowest to cope with escapes / multibyte characters.
+          for (let k of Object.keys(keys)) {
+            let s = JSON.stringify(keys[k]);
+            if (s.length > MAX_ITEM_SIZE) {
+              let count = Math.ceil(s.length / MAX_ITEM_SIZE);
+              let chunksCountKey = chunksKey(k);
+              let oldCount = await browser.storage.sync.get(chunksCountKey);
+              let chunks = {
+                [chunksCountKey]: count
+              };
+              for(let j = 0, o = 0; j < count; ++j, o += MAX_ITEM_SIZE) {
+                chunks[`${k}/${j}`] = s.substr(o, MAX_ITEM_SIZE);
+              }
+              await browser.storage.sync.set(chunks);
+              keys[k] = "[CHUNKED]";
+              if (oldCount-- > count) {
+                let oldChunks = [];
+                do {
+                  oldChunks.push(`${k}${oldCount}`);
+                } while(oldCount-- > count);
+                await browser.storage.sync.remove(oldChunks);
+              }
+            }
           }
         }
       }
-    }
-    try {
+
       let ret = await browser.storage[type][op](keys);
       if (sync && op === "set") {
-        let localFallback = await this.localFallback();
+        let localFallback = await getLocalFallback();
         let size = localFallback.size;
         if (size > 0) {
           for (let k of Object.keys(keys)) {
             localFallback.delete(k);
           }
-          if (size > localFallback.size) this.localFallback(localFallback);
+          if (size > localFallback.size) {
+            await setLocalFallback(localFallback);
+          }
         }
       }
       return ret;
     } catch (e) {
+      error(e, "%s.%s(%o)", type, op, keys);
       if (sync) {
         debug("Sync disabled? Falling back to local storage (%s %o)", op, keys);
-        let localFallback = await this.localFallback();
+        let localFallback = await getLocalFallback();
         let failedKeys = Array.isArray(keys) ? keys
           : typeof keys === "string" ? [keys] : Object.keys(keys);
         for (let k of failedKeys) {
           localFallback.add(k);
         }
-        await this.localFallback(localFallback);
+        await setLocalFallback(localFallback);
       } else {
-        error(e);
         throw e;
       }
     }
 
     return await browser.storage.local[op](keys);
-  },
+  }
 
-  async get(type, keys) {
-    return await this.safeOp("get", type, keys);
-  },
+  const LFK_NAME = "__fallbackKeys";
+  async function setLocalFallback(keys) {
+    return await browser.storage.local.set({[LFK_NAME]: [...keys]});
+  }
+  async function getLocalFallback() {
+    let keys = (await browser.storage.local.get(LFK_NAME))[LFK_NAME];
+    return new Set(Array.isArray(keys) ? keys : []);
+  }
 
-  async set(type, keys) {
-    return await this.safeOp("set", type, keys);
-  },
+  return {
+    async get(type, keys) {
+      return await safeOp("get", type, keys);
+    },
 
-  async localFallback(keys) {
-    let name = "__fallbackKeys";
-    if (keys) {
-      return await browser.storage.local.set({[name]: [...keys]});
+    async set(type, keys) {
+      return await safeOp("set", type, keys);
+    },
+
+    async remove(type, keys) {
+      return await safeOp("remove", type, keys);
+    },
+
+    async hasLocalFallback(key) {
+      return (await getLocalFallback()).has(key);
+    },
+
+    async isChunked(key) {
+      let ccKey = chunksKey(key);
+      let data = await browser.storage.sync.get([key, ccKey]);
+      return data[key] === "[CHUNKED]" && parseInt(data[ccKey]);
     }
-    let fallbackKeys = (await browser.storage.local.get(name))[name];
-    return new Set(Array.isArray(fallbackKeys) ? fallbackKeys : []);
-  }
-}
+  };
+})()