1 files changed, 199 insertions, 0 deletions
diff --git a/mobile/android/services/src/main/java/org/mozilla/gecko/sync/CollectionKeys.java b/mobile/android/services/src/main/java/org/mozilla/gecko/sync/CollectionKeys.java
new file mode 100644
index 000000000..1fd363bcb
--- /dev/null
+++ b/mobile/android/services/src/main/java/org/mozilla/gecko/sync/CollectionKeys.java
@@ -0,0 +1,199 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package org.mozilla.gecko.sync;
+
+import java.io.IOException;
+import java.io.UnsupportedEncodingException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map.Entry;
+import java.util.Set;
+
+import org.json.simple.JSONArray;
+import org.mozilla.apache.commons.codec.binary.Base64;
+import org.mozilla.gecko.sync.crypto.CryptoException;
+import org.mozilla.gecko.sync.crypto.KeyBundle;
+
+public class CollectionKeys {
+  private KeyBundle                  defaultKeyBundle     = null;
+  private final HashMap<String, KeyBundle> collectionKeyBundles = new HashMap<String, KeyBundle>();
+
+  /**
+   * Randomly generate a basic CollectionKeys object.
+   * @throws CryptoException
+   */
+  public static CollectionKeys generateCollectionKeys() throws CryptoException {
+    CollectionKeys ck = new CollectionKeys();
+    ck.clear();
+    ck.defaultKeyBundle = KeyBundle.withRandomKeys();
+    // TODO: eventually we would like to keep per-collection keys, just generate
+    // new ones as appropriate.
+    return ck;
+  }
+
+  public KeyBundle defaultKeyBundle() throws NoCollectionKeysSetException {
+    if (this.defaultKeyBundle == null) {
+      throw new NoCollectionKeysSetException();
+    }
+    return this.defaultKeyBundle;
+  }
+
+  public boolean keyBundleForCollectionIsNotDefault(String collection) {
+    return collectionKeyBundles.containsKey(collection);
+  }
+
+  public KeyBundle keyBundleForCollection(String collection)
+      throws NoCollectionKeysSetException {
+    if (this.defaultKeyBundle == null) {
+      throw new NoCollectionKeysSetException();
+    }
+    if (keyBundleForCollectionIsNotDefault(collection)) {
+      return collectionKeyBundles.get(collection);
+    }
+    return this.defaultKeyBundle;
+  }
+
+  /**
+   * Take a pair of values in a JSON array, handing them off to KeyBundle to
+   * produce a usable keypair.
+   */
+  private static KeyBundle arrayToKeyBundle(JSONArray array) throws UnsupportedEncodingException {
+    String encKeyStr  = (String) array.get(0);
+    String hmacKeyStr = (String) array.get(1);
+    return KeyBundle.fromBase64EncodedKeys(encKeyStr, hmacKeyStr);
+  }
+
+  @SuppressWarnings("unchecked")
+  private static JSONArray keyBundleToArray(KeyBundle bundle) {
+    // Generate JSON.
+    JSONArray keysArray = new JSONArray();
+    keysArray.add(new String(Base64.encodeBase64(bundle.getEncryptionKey())));
+    keysArray.add(new String(Base64.encodeBase64(bundle.getHMACKey())));
+    return keysArray;
+  }
+
+  private ExtendedJSONObject asRecordContents() throws NoCollectionKeysSetException {
+    ExtendedJSONObject json = new ExtendedJSONObject();
+    json.put("id", "keys");
+    json.put("collection", "crypto");
+    json.put("default", keyBundleToArray(this.defaultKeyBundle()));
+    ExtendedJSONObject colls = new ExtendedJSONObject();
+    for (Entry<String, KeyBundle> collKey : collectionKeyBundles.entrySet()) {
+      colls.put(collKey.getKey(), keyBundleToArray(collKey.getValue()));
+    }
+    json.put("collections", colls);
+    return json;
+  }
+
+  public CryptoRecord asCryptoRecord() throws NoCollectionKeysSetException {
+    ExtendedJSONObject payload = this.asRecordContents();
+    CryptoRecord record = new CryptoRecord(payload);
+    record.collection = "crypto";
+    record.guid       = "keys";
+    record.deleted    = false;
+    return record;
+  }
+
+  /**
+   * Set my key bundle and collection keys with the given key bundle and data
+   * (possibly decrypted) from the given record.
+   *
+   * @param keys
+   *          A "crypto/keys" <code>CryptoRecord</code>, encrypted with
+   *          <code>syncKeyBundle</code> if <code>syncKeyBundle</code> is non-null.
+   * @param syncKeyBundle
+   *          If non-null, the sync key bundle to decrypt <code>keys</code> with.
+   */
+  public void setKeyPairsFromWBO(CryptoRecord keys, KeyBundle syncKeyBundle)
+      throws CryptoException, IOException, NonObjectJSONException {
+    if (keys == null) {
+      throw new IllegalArgumentException("cannot set key pairs from null record");
+    }
+    if (syncKeyBundle != null) {
+      keys.keyBundle = syncKeyBundle;
+      keys.decrypt();
+    }
+    ExtendedJSONObject cleartext = keys.payload;
+    KeyBundle defaultKey = arrayToKeyBundle((JSONArray) cleartext.get("default"));
+
+    ExtendedJSONObject collections = cleartext.getObject("collections");
+    HashMap<String, KeyBundle> collectionKeys = new HashMap<String, KeyBundle>();
+    for (Entry<String, Object> pair : collections.entrySet()) {
+      KeyBundle bundle = arrayToKeyBundle((JSONArray) pair.getValue());
+      collectionKeys.put(pair.getKey(), bundle);
+    }
+
+    this.collectionKeyBundles.clear();
+    this.collectionKeyBundles.putAll(collectionKeys);
+    this.defaultKeyBundle     = defaultKey;
+  }
+
+  public void setKeyBundleForCollection(String collection, KeyBundle keys) {
+    this.collectionKeyBundles.put(collection, keys);
+  }
+
+  public void setDefaultKeyBundle(KeyBundle keys) {
+    this.defaultKeyBundle = keys;
+  }
+
+  public void clear() {
+    this.defaultKeyBundle = null;
+    this.collectionKeyBundles.clear();
+  }
+
+  /**
+   * Return set of collections where key is either missing from one collection
+   * or not the same in both collections.
+   * <p>
+   * Does not check for different default keys.
+   */
+  public static Set<String> differences(CollectionKeys a, CollectionKeys b) {
+    Set<String> differences = new HashSet<String>();
+    Set<String> collections = new HashSet<String>(a.collectionKeyBundles.keySet());
+    collections.addAll(b.collectionKeyBundles.keySet());
+
+    // Iterate through one collection, collecting missing and differences.
+    for (String collection : collections) {
+      KeyBundle keyA;
+      KeyBundle keyB;
+      try {
+        keyA = a.keyBundleForCollection(collection); // Will return default key as appropriate.
+        keyB = b.keyBundleForCollection(collection); // Will return default key as appropriate.
+      } catch (NoCollectionKeysSetException e) {
+        differences.add(collection);
+        continue;
+      }
+      // keyA and keyB are not null at this point.
+      if (!keyA.equals(keyB)) {
+        differences.add(collection);
+      }
+    }
+
+    return differences;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (!(o instanceof CollectionKeys)) {
+      return false;
+    }
+    CollectionKeys other = (CollectionKeys) o;
+    try {
+      // It would be nice to use map equality here, but there can be map entries
+      // where the key is the default key that should compare equal to a missing
+      // map entry. Therefore, we always compute the set of differences.
+      return defaultKeyBundle().equals(other.defaultKeyBundle()) &&
+             CollectionKeys.differences(this, other).isEmpty();
+    } catch (NoCollectionKeysSetException e) {
+      // If either default key bundle is not set, we'll say the bundles are not equal.
+      return false;
+    }
+  }
+
+  @Override
+  public int hashCode() {
+    return super.hashCode();
+  }
+}