Bug 1185544 - Add data delivery to the WebSocket backend. r=dragana,nsm
authorKit Cambridge <kcambridge@mozilla.com>
Thu, 17 Sep 2015 05:08:50 -0700
changeset 294421 d0a7044bb280b26897746b69829ea230e2e5ffec
parent 294420 0e0f3104478fcb2ed29be54f3e5ee113245378b8
child 294422 a39aa7c8e1ee5ed2d5e4c7de4a50f5d1da96ba85
push id5595
push usersteffen.wilberg@web.de
push dateSun, 20 Sep 2015 16:40:48 +0000
reviewersdragana, nsm
bugs1185544
milestone43.0a1
Bug 1185544 - Add data delivery to the WebSocket backend. r=dragana,nsm
dom/push/PushCrypto.jsm
dom/push/PushRecord.jsm
dom/push/PushService.jsm
dom/push/PushServiceHttp2.jsm
dom/push/PushServiceHttp2Crypto.jsm
dom/push/PushServiceWebSocket.jsm
dom/push/moz.build
dom/push/test/xpcshell/test_updateRecordNoEncryptionKeys.js
dom/push/test/xpcshell/test_updateRecordNoEncryptionKeys_http2.js
dom/push/test/xpcshell/test_updateRecordNoEncryptionKeys_ws.js
dom/push/test/xpcshell/xpcshell.ini
new file mode 100644
--- /dev/null
+++ b/dom/push/PushCrypto.jsm
@@ -0,0 +1,221 @@
+/* jshint moz: true, esnext: true */
+/* 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/. */
+
+'use strict';
+
+const Cu = Components.utils;
+
+Cu.importGlobalProperties(['crypto']);
+
+this.EXPORTED_SYMBOLS = ['PushCrypto', 'concatArray',
+                         'getEncryptionKeyParams', 'getEncryptionParams',
+                         'base64UrlDecode'];
+
+var ENCRYPT_INFO = new TextEncoder('utf-8').encode('Content-Encoding: aesgcm128');
+var NONCE_INFO = new TextEncoder('utf-8').encode('Content-Encoding: nonce');
+
+this.getEncryptionKeyParams = function(encryptKeyField) {
+  var params = encryptKeyField.split(',');
+  return params.reduce((m, p) => {
+    var pmap = p.split(';').reduce(parseHeaderFieldParams, {});
+    if (pmap.keyid && pmap.dh) {
+      m[pmap.keyid] = pmap.dh;
+    }
+    return m;
+  }, {});
+};
+
+this.getEncryptionParams = function(encryptField) {
+  var p = encryptField.split(',', 1)[0];
+  if (!p) {
+    return null;
+  }
+  return p.split(';').reduce(parseHeaderFieldParams, {});
+};
+
+var parseHeaderFieldParams = (m, v) => {
+  var i = v.indexOf('=');
+  if (i >= 0) {
+    // A quoted string with internal quotes is invalid for all the possible
+    // values of this header field.
+    m[v.substring(0, i).trim()] = v.substring(i + 1).trim()
+                                    .replace(/^"(.*)"$/, '$1');
+  }
+  return m;
+};
+
+function chunkArray(array, size) {
+  var start = array.byteOffset || 0;
+  array = array.buffer || array;
+  var index = 0;
+  var result = [];
+  while(index + size <= array.byteLength) {
+    result.push(new Uint8Array(array, start + index, size));
+    index += size;
+  }
+  if (index < array.byteLength) {
+    result.push(new Uint8Array(array, start + index));
+  }
+  return result;
+}
+
+this.base64UrlDecode = function(s) {
+  s = s.replace(/-/g, '+').replace(/_/g, '/');
+
+  // Replace padding if it was stripped by the sender.
+  // See http://tools.ietf.org/html/rfc4648#section-4
+  switch (s.length % 4) {
+    case 0:
+      break; // No pad chars in this case
+    case 2:
+      s += '==';
+      break; // Two pad chars
+    case 3:
+      s += '=';
+      break; // One pad char
+    default:
+      throw new Error('Illegal base64url string!');
+  }
+
+  // With correct padding restored, apply the standard base64 decoder
+  var decoded = atob(s);
+
+  var array = new Uint8Array(new ArrayBuffer(decoded.length));
+  for (var i = 0; i < decoded.length; i++) {
+    array[i] = decoded.charCodeAt(i);
+  }
+  return array;
+};
+
+this.concatArray = function(arrays) {
+  var size = arrays.reduce((total, a) => total + a.byteLength, 0);
+  var index = 0;
+  return arrays.reduce((result, a) => {
+    result.set(new Uint8Array(a), index);
+    index += a.byteLength;
+    return result;
+  }, new Uint8Array(size));
+};
+
+var HMAC_SHA256 = { name: 'HMAC', hash: 'SHA-256' };
+
+function hmac(key) {
+  this.keyPromise = crypto.subtle.importKey('raw', key, HMAC_SHA256,
+                                            false, ['sign']);
+}
+
+hmac.prototype.hash = function(input) {
+  return this.keyPromise.then(k => crypto.subtle.sign('HMAC', k, input));
+};
+
+function hkdf(salt, ikm) {
+  this.prkhPromise = new hmac(salt).hash(ikm)
+    .then(prk => new hmac(prk));
+}
+
+hkdf.prototype.generate = function(info, len) {
+  var input = concatArray([info, new Uint8Array([1])]);
+  return this.prkhPromise
+    .then(prkh => prkh.hash(input))
+    .then(h => {
+      if (h.byteLength < len) {
+        throw new Error('Length is too long');
+      }
+      return h.slice(0, len);
+    });
+};
+
+/* generate a 96-bit IV for use in GCM, 48-bits of which are populated */
+function generateNonce(base, index) {
+  if (index >= Math.pow(2, 48)) {
+    throw new Error('Error generating IV - index is too large.');
+  }
+  var nonce = base.slice(0, 12);
+  nonce = new Uint8Array(nonce);
+  for (var i = 0; i < 6; ++i) {
+    nonce[nonce.byteLength - 1 - i] ^= (index / Math.pow(256, i)) & 0xff;
+  }
+  return nonce;
+}
+
+this.PushCrypto = {
+
+  generateKeys: function() {
+    return crypto.subtle.generateKey({ name: 'ECDH', namedCurve: 'P-256'},
+                                     true,
+                                     ['deriveBits'])
+      .then(cryptoKey =>
+         Promise.all([
+           crypto.subtle.exportKey('raw', cryptoKey.publicKey),
+           // TODO: change this when bug 1048931 lands.
+           crypto.subtle.exportKey('jwk', cryptoKey.privateKey)
+         ]));
+  },
+
+  decodeMsg: function(aData, aPrivateKey, aRemotePublicKey, aSalt, aRs) {
+
+    if (aData.byteLength === 0) {
+      // Zero length messages will be passed as null.
+      return Promise.resolve(null);
+    }
+
+    // The last chunk of data must be less than aRs, if it is not return an
+    // error.
+    if (aData.byteLength % (aRs + 16) === 0) {
+      return Promise.reject(new Error('Data truncated'));
+    }
+
+    return Promise.all([
+      crypto.subtle.importKey('raw', base64UrlDecode(aRemotePublicKey),
+                              { name: 'ECDH', namedCurve: 'P-256' },
+                              false,
+                              ['deriveBits']),
+      crypto.subtle.importKey('jwk', aPrivateKey,
+                              { name: 'ECDH', namedCurve: 'P-256' },
+                              false,
+                              ['deriveBits'])
+    ])
+    .then(keys =>
+      crypto.subtle.deriveBits({ name: 'ECDH', public: keys[0] }, keys[1], 256))
+    .then(rawKey => {
+      var kdf = new hkdf(base64UrlDecode(aSalt), new Uint8Array(rawKey));
+      return Promise.all([
+        kdf.generate(ENCRYPT_INFO, 16)
+          .then(gcmBits =>
+                crypto.subtle.importKey('raw', gcmBits, 'AES-GCM', false,
+                                        ['decrypt'])),
+        kdf.generate(NONCE_INFO, 12)
+      ])
+    })
+    .then(r =>
+      // AEAD_AES_128_GCM expands ciphertext to be 16 octets longer.
+      Promise.all(chunkArray(aData, aRs + 16).map((slice, index) =>
+        this._decodeChunk(slice, index, r[1], r[0]))))
+    .then(r => concatArray(r));
+  },
+
+  _decodeChunk: function(aSlice, aIndex, aNonce, aKey) {
+    return crypto.subtle.decrypt({name: 'AES-GCM',
+                                  iv: generateNonce(aNonce, aIndex)
+                                 },
+                                 aKey, aSlice)
+      .then(decoded => {
+        decoded = new Uint8Array(decoded);
+        if (decoded.length == 0) {
+          return Promise.reject(new Error('Decoded array is too short!'));
+        } else if (decoded[0] > decoded.length) {
+          return Promise.reject(new Error ('Padding is wrong!'));
+        } else {
+          // All padded bytes must be zero except the first one.
+          for (var i = 1; i <= decoded[0]; i++) {
+            if (decoded[i] != 0) {
+              return Promise.reject(new Error('Padding is wrong!'));
+            }
+          }
+          return decoded.slice(decoded[0] + 1);
+        }
+      });
+  }
+};
--- a/dom/push/PushRecord.jsm
+++ b/dom/push/PushRecord.jsm
@@ -39,16 +39,18 @@ const QUOTA_REFRESH_TRANSITIONS_SQL = [
 ].join(",");
 
 function PushRecord(props) {
   this.pushEndpoint = props.pushEndpoint;
   this.scope = props.scope;
   this.originAttributes = props.originAttributes;
   this.pushCount = props.pushCount || 0;
   this.lastPush = props.lastPush || 0;
+  this.p256dhPublicKey = props.p256dhPublicKey;
+  this.p256dhPrivateKey = props.p256dhPrivateKey;
   this.setQuota(props.quota);
   this.ctime = (typeof props.ctime === "number") ? props.ctime : 0;
 }
 
 PushRecord.prototype = {
   setQuota(suggestedQuota) {
     this.quota = (!isNaN(suggestedQuota) && suggestedQuota >= 0) ?
                  suggestedQuota : prefs.get("maxQuotaPerSubscription");
@@ -186,22 +188,24 @@ PushRecord.prototype = {
     return this.quota === 0;
   },
 
   toRegistration() {
     return {
       pushEndpoint: this.pushEndpoint,
       lastPush: this.lastPush,
       pushCount: this.pushCount,
+      p256dhKey: this.p256dhPublicKey,
     };
   },
 
   toRegister() {
     return {
       pushEndpoint: this.pushEndpoint,
+      p256dhKey: this.p256dhPublicKey,
     };
   },
 };
 
 // Define lazy getters for the principal and scope URI. IndexedDB can't store
 // `nsIPrincipal` objects, so we keep them in a private weak map.
 var principals = new WeakMap();
 Object.defineProperties(PushRecord.prototype, {
--- a/dom/push/PushService.jsm
+++ b/dom/push/PushService.jsm
@@ -23,16 +23,17 @@ const {PushDB} = Cu.import("resource://g
 Cu.import("resource://gre/modules/XPCOMUtils.jsm");
 Cu.import("resource://gre/modules/Services.jsm");
 Cu.import("resource://gre/modules/Timer.jsm");
 Cu.import("resource://gre/modules/Preferences.jsm");
 Cu.import("resource://gre/modules/Promise.jsm");
 
 const {PushServiceWebSocket} = Cu.import("resource://gre/modules/PushServiceWebSocket.jsm");
 const {PushServiceHttp2} = Cu.import("resource://gre/modules/PushServiceHttp2.jsm");
+const {PushCrypto} = Cu.import("resource://gre/modules/PushCrypto.jsm");
 
 // Currently supported protocols: WebSocket.
 const CONNECTION_PROTOCOLS = [PushServiceWebSocket, PushServiceHttp2];
 
 XPCOMUtils.defineLazyModuleGetter(this, "AlarmService",
                                   "resource://gre/modules/AlarmService.jsm");
 
 this.EXPORTED_SYMBOLS = ["PushService"];
@@ -705,62 +706,85 @@ this.PushService = {
     }
   },
 
   // Fires a push-register system message to all applications that have
   // registration.
   _notifyAllAppsRegister: function() {
     debug("notifyAllAppsRegister()");
     // records are objects describing the registration as stored in IndexedDB.
-    return this._db.getAllUnexpired().then(records =>
-      records.forEach(record =>
-        this._notifySubscriptionChangeObservers(record)
-      )
-    );
+    return this._db.getAllUnexpired().then(records => {
+      records.forEach(record => {
+        this._notifySubscriptionChangeObservers(record);
+      });
+    });
   },
 
   dropRegistrationAndNotifyApp: function(aKeyId) {
     return this._db.getByKeyID(aKeyId).then(record => {
       this._notifySubscriptionChangeObservers(record);
       return this._db.delete(aKeyId);
     });
   },
 
   updateRegistrationAndNotifyApp: function(aOldKey, aRecord) {
     return this._db.delete(aOldKey)
       .then(_ => this._db.put(aRecord))
       .then(record => this._notifySubscriptionChangeObservers(record));
   },
 
+  ensureP256dhKey: function(record) {
+    if (record.p256dhPublicKey && record.p256dhPrivateKey) {
+      return Promise.resolve(record);
+    }
+    // We do not have a encryption key. so we need to generate it. This
+    // is only going to happen on db upgrade from version 4 to higher.
+    return PushCrypto.generateKeys()
+      .then(exportedKeys => {
+        return this.updateRecordAndNotifyApp(record.keyID, record => {
+          record.p256dhPublicKey = exportedKeys[0];
+          record.p256dhPrivateKey = exportedKeys[1];
+          return record;
+        });
+      }, error => {
+        return this.dropRegistrationAndNotifyApp(record.keyID).then(
+          () => Promise.reject(error));
+      });
+  },
+
   updateRecordAndNotifyApp: function(aKeyID, aUpdateFunc) {
     return this._db.update(aKeyID, aUpdateFunc)
-      .then(record => this._notifySubscriptionChangeObservers(record));
+      .then(record => {
+        this._notifySubscriptionChangeObservers(record);
+        return record;
+      });
   },
 
   _recordDidNotNotify: function(reason) {
     Services.telemetry.
       getHistogramById("PUSH_API_NOTIFICATION_RECEIVED_BUT_DID_NOT_NOTIFY").
       add(reason);
   },
 
   /**
    * Dispatches an incoming message to a service worker, recalculating the
    * quota for the associated push registration. If the quota is exceeded,
    * the registration and message will be dropped, and the worker will not
    * be notified.
    *
    * @param {String} keyID The push registration ID.
    * @param {String} message The message contents.
+   * @param {Object} cryptoParams The message encryption settings.
    * @param {Function} updateFunc A function that receives the existing
    *  registration record as its argument, and returns a new record. If the
    *  function returns `null` or `undefined`, the record will not be updated.
    *  `PushServiceWebSocket` uses this to drop incoming updates with older
    *  versions.
    */
-  receivedPushMessage: function(keyID, message, updateFunc) {
+  receivedPushMessage: function(keyID, message, cryptoParams, updateFunc) {
     debug("receivedPushMessage()");
     Services.telemetry.getHistogramById("PUSH_API_NOTIFICATION_RECEIVED").add();
 
     let shouldNotify = false;
     return this.getByKeyID(keyID).then(record => {
       if (!record) {
         this._recordDidNotNotify(kDROP_NOTIFICATION_REASON_KEY_NOT_FOUND);
         throw new Error("No record for key ID " + keyID);
@@ -774,45 +798,59 @@ this.PushService = {
           this._recordDidNotNotify(kDROP_NOTIFICATION_REASON_NO_HISTORY);
       }
       return this._db.update(keyID, record => {
         let newRecord = updateFunc(record);
         if (!newRecord) {
           this._recordDidNotNotify(kDROP_NOTIFICATION_REASON_NO_VERSION_INCREMENT);
           return null;
         }
-        // FIXME(nsm): WHY IS expired checked here but then also checked in the next case?
+        // Because `unregister` is advisory only, we can still receive messages
+        // for stale Simple Push registrations from the server. To work around
+        // this, we check if the record has expired before *and* after updating
+        // the quota.
         if (newRecord.isExpired()) {
-          // Because `unregister` is advisory only, we can still receive messages
-          // for stale registrations from the server.
           debug("receivedPushMessage: Ignoring update for expired key ID " + keyID);
           return null;
         }
         newRecord.receivedPush(lastVisit);
         return newRecord;
       });
     }).then(record => {
       var notified = false;
       if (!record) {
         return notified;
       }
-
-      if (shouldNotify) {
-        notified = this._notifyApp(record, message);
+      let decodedPromise;
+      if (cryptoParams) {
+        decodedPromise = PushCrypto.decodeMsg(
+          message,
+          record.p256dhPrivateKey,
+          cryptoParams.dh,
+          cryptoParams.salt,
+          cryptoParams.rs
+        ).then(bytes => new TextDecoder("utf-8").decode(bytes));
+      } else {
+        decodedPromise = Promise.resolve("");
       }
-      if (record.isExpired()) {
-        this._recordDidNotNotify(kDROP_NOTIFICATION_REASON_EXPIRED);
-        // Drop the registration in the background. If the user returns to the
-        // site, the service worker will be notified on the next `idle-daily`
-        // event.
-        this._sendUnregister(record).catch(error => {
-          debug("receivedPushMessage: Unregister error: " + error);
-        });
-      }
-      return notified;
+      return decodedPromise.then(message => {
+        if (shouldNotify) {
+          notified = this._notifyApp(record, message);
+        }
+        if (record.isExpired()) {
+          this._recordDidNotNotify(kDROP_NOTIFICATION_REASON_EXPIRED);
+          // Drop the registration in the background. If the user returns to the
+          // site, the service worker will be notified on the next `idle-daily`
+          // event.
+          this._sendUnregister(record).catch(error => {
+            debug("receivedPushMessage: Unregister error: " + error);
+          });
+        }
+        return notified;
+      });
     }).catch(error => {
       debug("receivedPushMessage: Error notifying app: " + error);
     });
   },
 
   _notifyApp: function(aPushRecord, message) {
     if (!aPushRecord || !aPushRecord.scope ||
         aPushRecord.originAttributes === undefined) {
--- a/dom/push/PushServiceHttp2.jsm
+++ b/dom/push/PushServiceHttp2.jsm
@@ -14,18 +14,22 @@ const {PushDB} = Cu.import("resource://g
 const {PushRecord} = Cu.import("resource://gre/modules/PushRecord.jsm");
 Cu.import("resource://gre/modules/XPCOMUtils.jsm");
 Cu.import("resource://gre/modules/Services.jsm");
 Cu.import("resource://gre/modules/IndexedDBHelper.jsm");
 Cu.import("resource://gre/modules/Timer.jsm");
 Cu.import("resource://gre/modules/Preferences.jsm");
 Cu.import("resource://gre/modules/Promise.jsm");
 
-const {PushServiceHttp2Crypto, concatArray} =
-  Cu.import("resource://gre/modules/PushServiceHttp2Crypto.jsm");
+const {
+  PushCrypto,
+  concatArray,
+  getEncryptionKeyParams,
+  getEncryptionParams,
+} = Cu.import("resource://gre/modules/PushCrypto.jsm");
 
 this.EXPORTED_SYMBOLS = ["PushServiceHttp2"];
 
 const prefs = new Preferences("dom.push.");
 
 // Don't modify this, instead set dom.push.debug.
 // Set debug first so that all debugging actually works.
 var gDebuggingEnabled = prefs.get("debug");
@@ -171,52 +175,30 @@ PushChannelListener.prototype = {
                                                          msg,
                                                          dh,
                                                          salt,
                                                          rs);
     }
   }
 };
 
-var parseHeaderFieldParams = (m, v) => {
-  var i = v.indexOf('=');
-  if (i >= 0) {
-    // A quoted string with internal quotes is invalid for all the possible
-    // values of this header field.
-    m[v.substring(0, i).trim()] = v.substring(i + 1).trim()
-                                    .replace(/^"(.*)"$/, '$1');
-  }
-  return m;
-};
-
 function encryptKeyFieldParser(aRequest) {
   try {
     var encryptKeyField = aRequest.getRequestHeader("Encryption-Key");
-
-    var params = encryptKeyField.split(',');
-    return params.reduce((m, p) => {
-      var pmap = p.split(';').reduce(parseHeaderFieldParams, {});
-      if (pmap.keyid && pmap.dh) {
-        m[pmap.keyid] = pmap.dh;
-      }
-      return m;
-    }, {});
-
+    return getEncryptionKeyParams(encryptKeyField);
   } catch(e) {
     // getRequestHeader can throw.
     return null;
   }
 }
 
 function encryptFieldParser(aRequest) {
   try {
-    return aRequest.getRequestHeader("Encryption")
-             .split(',', 1)[0]
-             .split(';')
-             .reduce(parseHeaderFieldParams, {});
+    var encryptField = aRequest.getRequestHeader("Encryption");
+    return getEncryptionParams(encryptField);
   } catch(e) {
     // getRequestHeader can throw.
     return null;
   }
 }
 
 var PushServiceDelete = function(resolve, reject) {
   this._resolve = resolve;
@@ -528,17 +510,17 @@ this.PushServiceHttp2 = {
   _subscribeResource: function(aRecord) {
     debug("subscribeResource()");
 
     return this._subscribeResourceInternal({
       record: aRecord,
       retries: 0
     })
     .then(result =>
-      PushServiceHttp2Crypto.generateKeys()
+      PushCrypto.generateKeys()
       .then(exportedKeys => {
         result.p256dhPublicKey = exportedKeys[0];
         result.p256dhPrivateKey = exportedKeys[1];
         this._conns[result.subscriptionUri] = {
           channel: null,
           listener: null,
           countUnableToConnect: 0,
           lastStartListening: 0,
@@ -719,44 +701,21 @@ this.PushServiceHttp2 = {
   },
 
   // Start listening if subscriptions present.
   startConnections: function(aSubscriptions) {
     debug("startConnections() " + aSubscriptions.length);
 
     for (let i = 0; i < aSubscriptions.length; i++) {
       let record = aSubscriptions[i];
-      if (record.p256dhPublicKey && record.p256dhPrivateKey) {
+      this._mainPushService.ensureP256dhKey(record).then(record => {
         this._startSingleConnection(record);
-      } else {
-        // We do not have a encryption key. so we need to generate it. This
-        // is only going to happen on db upgrade from version 4 to higher.
-        PushServiceHttp2Crypto.generateKeys()
-          .then(exportedKeys => {
-            if (this._mainPushService) {
-              return this._mainPushService
-                .updateRecordAndNotifyApp(record.subscriptionUri, record => {
-                  record.p256dhPublicKey = exportedKeys[0];
-                  record.p256dhPrivateKey = exportedKeys[1];
-                  return record;
-                });
-            }
-          }, error => {
-            record = null;
-            if (this._mainPushService) {
-              this._mainPushService
-                .dropRegistrationAndNotifyApp(record.subscriptionUri);
-            }
-          })
-          .then(_ => {
-            if (record) {
-              this._startSingleConnection(record);
-            }
-          });
-      }
+      }, error => {
+        debug("startConnections: Error updating record " + record.keyID);
+      });
     }
   },
 
   _startSingleConnection: function(record) {
     debug("_startSingleConnection()");
     if (typeof this._conns[record.subscriptionUri] != "object") {
       this._conns[record.subscriptionUri] = {channel: null,
                                              listener: null,
@@ -870,65 +829,55 @@ this.PushServiceHttp2 = {
     } else {
       this._retryAfterBackoff(aSubscriptionUri, -1);
     }
   },
 
   _pushChannelOnStop: function(aUri, aAckUri, aMessage, dh, salt, rs) {
     debug("pushChannelOnStop() ");
 
-    this._mainPushService.getByKeyID(aUri)
-    .then(aPushRecord =>
-      PushServiceHttp2Crypto.decodeMsg(aMessage, aPushRecord.p256dhPrivateKey,
-                                       dh, salt, rs)
-      .then(msg => {
-        var msgString = '';
-        for (var i=0; i<msg.length; i++) {
-          msgString += String.fromCharCode(msg[i]);
-        }
-        return this._mainPushService.receivedPushMessage(aUri,
-                                                         msgString,
-                                                         record => {
-          // Always update the stored record.
-          return record;
-        });
-      })
+    let cryptoParams = {
+      dh: dh,
+      salt: salt,
+      rs: rs,
+    };
+    this._mainPushService.receivedPushMessage(
+      aUri, aMessage, cryptoParams, record => {
+        // Always update the stored record.
+        return record;
+      }
     )
     .then(_ => this._ackMsgRecv(aAckUri))
     .catch(err => {
       debug("Error receiving message: " + err);
     });
   },
 
   onAlarmFired: function() {
     this._startConnectionsWaitingForAlarm();
   },
 };
 
 function PushRecordHttp2(record) {
   PushRecord.call(this, record);
   this.subscriptionUri = record.subscriptionUri;
   this.pushReceiptEndpoint = record.pushReceiptEndpoint;
-  this.p256dhPublicKey = record.p256dhPublicKey;
-  this.p256dhPrivateKey = record.p256dhPrivateKey;
 }
 
 PushRecordHttp2.prototype = Object.create(PushRecord.prototype, {
   keyID: {
     get() {
       return this.subscriptionUri;
     },
   },
 });
 
 PushRecordHttp2.prototype.toRegistration = function() {
   let registration = PushRecord.prototype.toRegistration.call(this);
   registration.pushReceiptEndpoint = this.pushReceiptEndpoint;
-  registration.p256dhKey = this.p256dhPublicKey;
   return registration;
 };
 
 PushRecordHttp2.prototype.toRegister = function() {
   let register = PushRecord.prototype.toRegister.call(this);
   register.pushReceiptEndpoint = this.pushReceiptEndpoint;
-  register.p256dhKey = this.p256dhPublicKey;
   return register;
 };
deleted file mode 100644
--- a/dom/push/PushServiceHttp2Crypto.jsm
+++ /dev/null
@@ -1,189 +0,0 @@
-/* jshint moz: true, esnext: true */
-/* 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/. */
-
-'use strict';
-
-const Cu = Components.utils;
-
-Cu.importGlobalProperties(['crypto']);
-
-this.EXPORTED_SYMBOLS = ['PushServiceHttp2Crypto', 'concatArray'];
-
-var ENCRYPT_INFO = new TextEncoder('utf-8').encode('Content-Encoding: aesgcm128');
-var NONCE_INFO = new TextEncoder('utf-8').encode('Content-Encoding: nonce');
-
-function chunkArray(array, size) {
-  var start = array.byteOffset || 0;
-  array = array.buffer || array;
-  var index = 0;
-  var result = [];
-  while(index + size <= array.byteLength) {
-    result.push(new Uint8Array(array, start + index, size));
-    index += size;
-  }
-  if (index < array.byteLength) {
-    result.push(new Uint8Array(array, start + index));
-  }
-  return result;
-}
-
-function base64UrlDecode(s) {
-  s = s.replace(/-/g, '+').replace(/_/g, '/');
-
-  // Replace padding if it was stripped by the sender.
-  // See http://tools.ietf.org/html/rfc4648#section-4
-  switch (s.length % 4) {
-    case 0:
-      break; // No pad chars in this case
-    case 2:
-      s += '==';
-      break; // Two pad chars
-    case 3:
-      s += '=';
-      break; // One pad char
-    default:
-      throw new Error('Illegal base64url string!');
-  }
-
-  // With correct padding restored, apply the standard base64 decoder
-  var decoded = atob(s);
-
-  var array = new Uint8Array(new ArrayBuffer(decoded.length));
-  for (var i = 0; i < decoded.length; i++) {
-    array[i] = decoded.charCodeAt(i);
-  }
-  return array;
-}
-
-this.concatArray = function(arrays) {
-  var size = arrays.reduce((total, a) => total + a.byteLength, 0);
-  var index = 0;
-  return arrays.reduce((result, a) => {
-    result.set(new Uint8Array(a), index);
-    index += a.byteLength;
-    return result;
-  }, new Uint8Array(size));
-};
-
-var HMAC_SHA256 = { name: 'HMAC', hash: 'SHA-256' };
-
-function hmac(key) {
-  this.keyPromise = crypto.subtle.importKey('raw', key, HMAC_SHA256,
-                                            false, ['sign']);
-}
-
-hmac.prototype.hash = function(input) {
-  return this.keyPromise.then(k => crypto.subtle.sign('HMAC', k, input));
-};
-
-function hkdf(salt, ikm) {
-  this.prkhPromise = new hmac(salt).hash(ikm)
-    .then(prk => new hmac(prk));
-}
-
-hkdf.prototype.generate = function(info, len) {
-  var input = concatArray([info, new Uint8Array([1])]);
-  return this.prkhPromise
-    .then(prkh => prkh.hash(input))
-    .then(h => {
-      if (h.byteLength < len) {
-        throw new Error('Length is too long');
-      }
-      return h.slice(0, len);
-    });
-};
-
-/* generate a 96-bit IV for use in GCM, 48-bits of which are populated */
-function generateNonce(base, index) {
-  if (index >= Math.pow(2, 48)) {
-    throw new Error('Error generating IV - index is too large.');
-  }
-  var nonce = base.slice(0, 12);
-  nonce = new Uint8Array(nonce);
-  for (var i = 0; i < 6; ++i) {
-    nonce[nonce.byteLength - 1 - i] ^= (index / Math.pow(256, i)) & 0xff;
-  }
-  return nonce;
-}
-
-this.PushServiceHttp2Crypto = {
-
-  generateKeys: function() {
-    return crypto.subtle.generateKey({ name: 'ECDH', namedCurve: 'P-256'},
-                                     true,
-                                     ['deriveBits'])
-      .then(cryptoKey =>
-         Promise.all([
-           crypto.subtle.exportKey('raw', cryptoKey.publicKey),
-           // TODO: change this when bug 1048931 lands.
-           crypto.subtle.exportKey('jwk', cryptoKey.privateKey)
-         ]));
-  },
-
-  decodeMsg: function(aData, aPrivateKey, aRemotePublicKey, aSalt, aRs) {
-
-    if (aData.byteLength === 0) {
-      // Zero length messages will be passed as null.
-      return Promise.resolve(null);
-    }
-
-    // The last chunk of data must be less than aRs, if it is not return an
-    // error.
-    if (aData.byteLength % (aRs + 16) === 0) {
-      return Promise.reject(new Error('Data truncated'));
-    }
-
-    return Promise.all([
-      crypto.subtle.importKey('raw', base64UrlDecode(aRemotePublicKey),
-                              { name: 'ECDH', namedCurve: 'P-256' },
-                              false,
-                              ['deriveBits']),
-      crypto.subtle.importKey('jwk', aPrivateKey,
-                              { name: 'ECDH', namedCurve: 'P-256' },
-                              false,
-                              ['deriveBits'])
-    ])
-    .then(keys =>
-      crypto.subtle.deriveBits({ name: 'ECDH', public: keys[0] }, keys[1], 256))
-    .then(rawKey => {
-      var kdf = new hkdf(base64UrlDecode(aSalt), new Uint8Array(rawKey));
-      return Promise.all([
-        kdf.generate(ENCRYPT_INFO, 16)
-          .then(gcmBits =>
-                crypto.subtle.importKey('raw', gcmBits, 'AES-GCM', false,
-                                        ['decrypt'])),
-        kdf.generate(NONCE_INFO, 12)
-      ])
-    })
-    .then(r =>
-      // AEAD_AES_128_GCM expands ciphertext to be 16 octets longer.
-      Promise.all(chunkArray(aData, aRs + 16).map((slice, index) =>
-        this._decodeChunk(slice, index, r[1], r[0]))))
-    .then(r => concatArray(r));
-  },
-
-  _decodeChunk: function(aSlice, aIndex, aNonce, aKey) {
-    return crypto.subtle.decrypt({name: 'AES-GCM',
-                                  iv: generateNonce(aNonce, aIndex)
-                                 },
-                                 aKey, aSlice)
-      .then(decoded => {
-        decoded = new Uint8Array(decoded);
-        if (decoded.length == 0) {
-          return Promise.reject(new Error('Decoded array is too short!'));
-        } else if (decoded[0] > decoded.length) {
-          return Promise.reject(new Error ('Padding is wrong!'));
-        } else {
-          // All padded bytes must be zero except the first one.
-          for (var i = 1; i <= decoded[0]; i++) {
-            if (decoded[i] != 0) {
-              return Promise.reject(new Error('Padding is wrong!'));
-            }
-          }
-          return decoded.slice(decoded[0] + 1);
-        }
-      });
-  }
-};
--- a/dom/push/PushServiceWebSocket.jsm
+++ b/dom/push/PushServiceWebSocket.jsm
@@ -7,16 +7,22 @@
 
 const Cc = Components.classes;
 const Ci = Components.interfaces;
 const Cu = Components.utils;
 const Cr = Components.results;
 
 const {PushDB} = Cu.import("resource://gre/modules/PushDB.jsm");
 const {PushRecord} = Cu.import("resource://gre/modules/PushRecord.jsm");
+const {
+  PushCrypto,
+  base64UrlDecode,
+  getEncryptionKeyParams,
+  getEncryptionParams,
+} = Cu.import("resource://gre/modules/PushCrypto.jsm");
 Cu.import("resource://gre/modules/Preferences.jsm");
 Cu.import("resource://gre/modules/Timer.jsm");
 Cu.import("resource://gre/modules/Promise.jsm");
 Cu.import("resource://gre/modules/XPCOMUtils.jsm");
 Cu.import("resource://gre/modules/Services.jsm");
 
 
 XPCOMUtils.defineLazyServiceGetter(this, "gDNSService",
@@ -28,30 +34,51 @@ XPCOMUtils.defineLazyServiceGetter(this,
                                    "@mozilla.org/power/powermanagerservice;1",
                                    "nsIPowerManagerService");
 #endif
 
 var threadManager = Cc["@mozilla.org/thread-manager;1"]
                       .getService(Ci.nsIThreadManager);
 
 const kPUSHWSDB_DB_NAME = "pushapi";
-const kPUSHWSDB_DB_VERSION = 4; // Change this if the IndexedDB format changes
+const kPUSHWSDB_DB_VERSION = 5; // Change this if the IndexedDB format changes
 const kPUSHWSDB_STORE_NAME = "pushapi";
 
 const kUDP_WAKEUP_WS_STATUS_CODE = 4774;  // WebSocket Close status code sent
                                           // by server to signal that it can
                                           // wake client up using UDP.
 
 const prefs = new Preferences("dom.push.");
 
 this.EXPORTED_SYMBOLS = ["PushServiceWebSocket"];
 
 // Don't modify this, instead set dom.push.debug.
 var gDebuggingEnabled = true;
 
+function getCryptoParams(headers) {
+  if (!headers) {
+    return null;
+  }
+  var keymap = getEncryptionKeyParams(headers.encryption_key);
+  if (!keymap) {
+    return null;
+  }
+  var enc = getEncryptionParams(headers.encryption);
+  if (!enc || !enc.keyid) {
+    return null;
+  }
+  var dh = keymap[enc.keyid];
+  var salt = enc.salt;
+  var rs = (enc.rs)? parseInt(enc.rs, 10) : 4096;
+  if (!dh || !salt || isNaN(rs) || (rs <= 1)) {
+    return null;
+  }
+  return {dh, salt, rs};
+}
+
 function debug(s) {
   if (gDebuggingEnabled) {
     dump("-*- PushServiceWebSocket.jsm: " + s + "\n");
   }
 }
 
 // Set debug first so that all debugging actually works.
 gDebuggingEnabled = prefs.get("debug");
@@ -131,21 +158,21 @@ this.PushServiceWebSocket = {
 
   disconnect: function() {
     this._shutdownWS();
   },
 
   observe: function(aSubject, aTopic, aData) {
 
     switch (aTopic) {
-      case "nsPref:changed":
-        if (aData == "dom.push.debug") {
-          gDebuggingEnabled = prefs.get("debug");
-        }
-        break;
+    case "nsPref:changed":
+      if (aData == "dom.push.debug") {
+        gDebuggingEnabled = prefs.get("debug");
+      }
+      break;
     case "timer-callback":
       if (aSubject == this._requestTimeoutTimer) {
         if (Object.keys(this._pendingRequests).length === 0) {
           this._requestTimeoutTimer.cancel();
         }
 
         // Set to true if at least one request timed out.
         let requestTimedOut = false;
@@ -260,16 +287,19 @@ this.PushServiceWebSocket = {
    */
   _lastGoodPingInterval: 0,
 
   /**
    * Maximum ping interval that we can reach.
    */
   _upperLimit: 0,
 
+  /** Indicates whether the server supports Web Push-style message delivery. */
+  _dataEnabled: false,
+
   /**
    * Sends a message to the Push Server through an open websocket.
    * typeof(msg) shall be an object
    */
   _wsSendMessage: function(msg) {
     if (!this._ws) {
       debug("No WebSocket initialized. Cannot send a message.");
       return;
@@ -351,16 +381,18 @@ this.PushServiceWebSocket = {
     // or receiving notifications.
     this._shutdownWS();
 
     if (this._requestTimeoutTimer) {
       this._requestTimeoutTimer.cancel();
     }
 
     this._mainPushService = null;
+
+    this._dataEnabled = false;
   },
 
   /**
    * How retries work:  The goal is to ensure websocket is always up on
    * networks not supporting UDP. So the websocket should only be shutdown if
    * onServerClose indicates UDP wakeup.  If WS is closed due to socket error,
    * _reconnectAfterBackoff() is called.  The retry alarm is started and when
    * it times out, beginWSSetup() is called again.
@@ -744,22 +776,37 @@ this.PushServiceWebSocket = {
     // To avoid sticking extra large values sent by an evil server into prefs.
     if (reply.uaid.length > 128) {
       debug("UAID received from server was too long: " +
             reply.uaid);
       this._shutdownWS();
       return;
     }
 
+    let notifyRequestQueue = () => {
+      if (this._notifyRequestQueue) {
+        this._notifyRequestQueue();
+        this._notifyRequestQueue = null;
+      }
+    };
+
     function finishHandshake() {
       this._UAID = reply.uaid;
       this._currentState = STATE_READY;
-      if (this._notifyRequestQueue) {
-        this._notifyRequestQueue();
-        this._notifyRequestQueue = null;
+      this._dataEnabled = !!reply.use_webpush;
+      if (this._dataEnabled) {
+        this._mainPushService.getAllUnexpired().then(records =>
+          Promise.all(records.map(record =>
+            this._mainPushService.ensureP256dhKey(record).catch(error => {
+              debug("finishHandshake: Error updating record " + record.keyID);
+            })
+          ))
+        ).then(notifyRequestQueue);
+      } else {
+        notifyRequestQueue();
       }
     }
 
     // By this point we've got a UAID from the server that we are ready to
     // accept.
     //
     // If we already had a valid UAID before, we have to ask apps to
     // re-register.
@@ -816,21 +863,58 @@ this.PushServiceWebSocket = {
       dump("PushWebSocket " +  JSON.stringify(record));
       Services.telemetry.getHistogramById("PUSH_API_SUBSCRIBE_WS_TIME").add(Date.now() - tmp.ctime);
       tmp.resolve(record);
     } else {
       tmp.reject(reply);
     }
   },
 
+  _handleDataUpdate: function(update) {
+    let promise;
+    if (typeof update.channelID != "string") {
+      debug("handleDataUpdate: Discarding message without channel ID");
+      return;
+    }
+    if (typeof update.data != "string") {
+      promise = this._mainPushService.receivedPushMessage(
+        update.channelID,
+        null,
+        null,
+        record => record
+      );
+    } else {
+      let params = getCryptoParams(update.headers);
+      if (!params) {
+        debug("handleDataUpdate: Discarding invalid encrypted message");
+        return;
+      }
+      let message = base64UrlDecode(update.data);
+      promise = this._mainPushService.receivedPushMessage(
+        update.channelID,
+        message,
+        params,
+        record => record
+      );
+    }
+    promise.then(() => this._sendAck(update.channelID)).catch(err => {
+      debug("handleDataUpdate: Error delivering message: " + err);
+    });
+  },
+
   /**
    * Protocol handler invoked by server message.
    */
   _handleNotificationReply: function(reply) {
     debug("handleNotificationReply()");
+    if (this._dataEnabled) {
+      this._handleDataUpdate(reply);
+      return;
+    }
+
     if (typeof reply.updates !== 'object') {
       debug("No 'updates' field in response. Type = " + typeof reply.updates);
       return;
     }
 
     debug("Reply updates: " + reply.updates.length);
     for (let i = 0; i < reply.updates.length; i++) {
       let update = reply.updates[i];
@@ -897,16 +981,26 @@ this.PushServiceWebSocket = {
 
       return new Promise((resolve, reject) => {
         this._pendingRequests[data.channelID] = {record: record,
                                                  resolve: resolve,
                                                  reject: reject,
                                                  ctime: Date.now()
                                                 };
         this._queueRequest(data);
+      }).then(record => {
+        if (!this._dataEnabled) {
+          return record;
+        }
+        return PushCrypto.generateKeys()
+          .then(([publicKey, privateKey]) => {
+            record.p256dhPublicKey = publicKey;
+            record.p256dhPrivateKey = privateKey;
+            return record;
+          });
       });
     }
 
     this._queueRequest({channelID: record.channelID,
                         messageType: action});
     return Promise.resolve();
   },
 
@@ -956,17 +1050,17 @@ this.PushServiceWebSocket = {
         this._notifyRequestQueue = null;
       }
     }
   },
 
   _receivedUpdate: function(aChannelID, aLatestVersion) {
     debug("Updating: " + aChannelID + " -> " + aLatestVersion);
 
-    this._mainPushService.receivedPushMessage(aChannelID, "", record => {
+    this._mainPushService.receivedPushMessage(aChannelID, null, null, record => {
       if (record.version === null ||
           record.version < aLatestVersion) {
         debug("Version changed for " + aChannelID + ": " + aLatestVersion);
         record.version = aLatestVersion;
         return record;
       }
       debug("No significant version change for " + aChannelID + ": " +
             aLatestVersion);
@@ -985,16 +1079,17 @@ this.PushServiceWebSocket = {
       return;
     }
 
     // Since we've had a successful connection reset the retry fail count.
     this._retryFailCount = 0;
 
     let data = {
       messageType: "hello",
+      use_webpush: true,
     };
 
     if (this._UAID) {
       data.uaid = this._UAID;
     }
 
     function sendHelloMessage(ids) {
       // On success, ids is an array, on error its not.
--- a/dom/push/moz.build
+++ b/dom/push/moz.build
@@ -10,22 +10,22 @@ EXTRA_COMPONENTS += [
     'PushNotificationService.js',
 ]
 
 EXTRA_PP_JS_MODULES += [
     'PushServiceWebSocket.jsm',
 ]
 
 EXTRA_JS_MODULES += [
+    'PushCrypto.jsm',
     'PushDB.jsm',
     'PushRecord.jsm',
     'PushService.jsm',
     'PushServiceChildPreload.jsm',
     'PushServiceHttp2.jsm',
-    'PushServiceHttp2Crypto.jsm',
 ]
 
 MOCHITEST_MANIFESTS += [
     'test/mochitest.ini',
 ]
 
 XPCSHELL_TESTS_MANIFESTS += [
     'test/xpcshell/xpcshell.ini',
deleted file mode 100644
--- a/dom/push/test/xpcshell/test_updateRecordNoEncryptionKeys.js
+++ /dev/null
@@ -1,80 +0,0 @@
-/* Any copyright is dedicated to the Public Domain.
-   http://creativecommons.org/publicdomain/zero/1.0/ */
-
-'use strict';
-
-Cu.import("resource://gre/modules/Services.jsm");
-Cu.import("resource://testing-common/httpd.js");
-
-const {PushDB, PushService, PushServiceHttp2} = serviceExports;
-
-var httpServer = null;
-
-XPCOMUtils.defineLazyGetter(this, "serverPort", function() {
-  return httpServer.identity.primaryPort;
-});
-
-function listenHandler(metadata, response) {
-  do_check_true(true, "Start listening");
-  httpServer.stop(do_test_finished);
-  response.setHeader("Retry-After", "10");
-  response.setStatusLine(metadata.httpVersion, 500, "Retry");
-}
-
-httpServer = new HttpServer();
-httpServer.registerPathHandler("/subscriptionNoKey", listenHandler);
-httpServer.start(-1);
-
-function run_test() {
-
-  do_get_profile();
-  setPrefs({
-    'http2.retryInterval': 1000,
-    'http2.maxRetries': 2
-  });
-  disableServiceWorkerEvents(
-    'https://example.com/page'
-  );
-
-  run_next_test();
-}
-
-add_task(function* test1() {
-
-  let db = PushServiceHttp2.newPushDB();
-  do_register_cleanup(_ => {
-    return db.drop().then(_ => db.close());
-  });
-
-  do_test_pending();
-
-  var serverURL = "http://localhost:" + httpServer.identity.primaryPort;
-
-  let record = {
-    subscriptionUri: serverURL + '/subscriptionNoKey',
-    pushEndpoint: serverURL + '/pushEndpoint',
-    pushReceiptEndpoint: serverURL + '/pushReceiptEndpoint',
-    scope: 'https://example.com/page',
-    originAttributes: '',
-    quota: Infinity,
-  };
-
-  yield db.put(record);
-
-  let notifyPromise = promiseObserverNotification('push-subscription-change',
-                                                  _ => true);
-
-  PushService.init({
-    serverURI: serverURL + "/subscribe",
-    service: PushServiceHttp2,
-    db
-  });
-
-  yield waitForPromise(notifyPromise, DEFAULT_TIMEOUT,
-    'Timed out waiting for notifications');
-
-  let aRecord = yield db.getByKeyID(serverURL + '/subscriptionNoKey');
-  ok(aRecord, 'The record should still be there');
-  ok(aRecord.p256dhPublicKey, 'There should be a public key');
-  ok(aRecord.p256dhPrivateKey, 'There should be a private key');
-});
new file mode 100644
--- /dev/null
+++ b/dom/push/test/xpcshell/test_updateRecordNoEncryptionKeys_http2.js
@@ -0,0 +1,80 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+'use strict';
+
+Cu.import("resource://gre/modules/Services.jsm");
+Cu.import("resource://testing-common/httpd.js");
+
+const {PushDB, PushService, PushServiceHttp2} = serviceExports;
+
+var httpServer = null;
+
+XPCOMUtils.defineLazyGetter(this, "serverPort", function() {
+  return httpServer.identity.primaryPort;
+});
+
+function listenHandler(metadata, response) {
+  do_check_true(true, "Start listening");
+  httpServer.stop(do_test_finished);
+  response.setHeader("Retry-After", "10");
+  response.setStatusLine(metadata.httpVersion, 500, "Retry");
+}
+
+httpServer = new HttpServer();
+httpServer.registerPathHandler("/subscriptionNoKey", listenHandler);
+httpServer.start(-1);
+
+function run_test() {
+
+  do_get_profile();
+  setPrefs({
+    'http2.retryInterval': 1000,
+    'http2.maxRetries': 2
+  });
+  disableServiceWorkerEvents(
+    'https://example.com/page'
+  );
+
+  run_next_test();
+}
+
+add_task(function* test1() {
+
+  let db = PushServiceHttp2.newPushDB();
+  do_register_cleanup(_ => {
+    return db.drop().then(_ => db.close());
+  });
+
+  do_test_pending();
+
+  var serverURL = "http://localhost:" + httpServer.identity.primaryPort;
+
+  let record = {
+    subscriptionUri: serverURL + '/subscriptionNoKey',
+    pushEndpoint: serverURL + '/pushEndpoint',
+    pushReceiptEndpoint: serverURL + '/pushReceiptEndpoint',
+    scope: 'https://example.com/page',
+    originAttributes: '',
+    quota: Infinity,
+  };
+
+  yield db.put(record);
+
+  let notifyPromise = promiseObserverNotification('push-subscription-change',
+                                                  _ => true);
+
+  PushService.init({
+    serverURI: serverURL + "/subscribe",
+    service: PushServiceHttp2,
+    db
+  });
+
+  yield waitForPromise(notifyPromise, DEFAULT_TIMEOUT,
+    'Timed out waiting for notifications');
+
+  let aRecord = yield db.getByKeyID(serverURL + '/subscriptionNoKey');
+  ok(aRecord, 'The record should still be there');
+  ok(aRecord.p256dhPublicKey, 'There should be a public key');
+  ok(aRecord.p256dhPrivateKey, 'There should be a private key');
+});
new file mode 100644
--- /dev/null
+++ b/dom/push/test/xpcshell/test_updateRecordNoEncryptionKeys_ws.js
@@ -0,0 +1,92 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+'use strict';
+
+const {PushDB, PushService, PushServiceWebSocket, PushCrypto} = serviceExports;
+
+const userAgentID = '4dffd396-6582-471d-8c0c-84f394e9f7db';
+
+function run_test() {
+  do_get_profile();
+  setPrefs({
+    userAgentID,
+  });
+  disableServiceWorkerEvents(
+    'https://example.com/page/1',
+    'https://example.com/page/2',
+    'https://example.com/page/3'
+  );
+  run_next_test();
+}
+
+add_task(function* test_with_data_enabled() {
+  let db = PushServiceWebSocket.newPushDB();
+  do_register_cleanup(() => {return db.drop().then(_ => db.close());});
+
+  let [publicKey, privateKey] = yield PushCrypto.generateKeys();
+  let records = [{
+    channelID: 'eb18f12a-cc42-4f14-accb-3bfc1227f1aa',
+    pushEndpoint: 'https://example.org/push/no-key/1',
+    scope: 'https://example.com/page/1',
+    originAttributes: '',
+    quota: Infinity,
+  }, {
+    channelID: '0d8886b9-8da1-4778-8f5d-1cf93a877ed6',
+    pushEndpoint: 'https://example.org/push/key',
+    scope: 'https://example.com/page/2',
+    originAttributes: '',
+    p256dhPublicKey: publicKey,
+    p256dhPrivateKey: privateKey,
+    quota: Infinity,
+  }];
+  for (let record of records) {
+    yield db.put(record);
+  }
+
+  PushService.init({
+    serverURI: "wss://push.example.org/",
+    networkInfo: new MockDesktopNetworkInfo(),
+    db,
+    makeWebSocket(uri) {
+      return new MockWebSocket(uri, {
+        onHello(request) {
+          ok(request.use_webpush,
+            'Should use Web Push if data delivery is enabled');
+          this.serverSendMsg(JSON.stringify({
+            messageType: 'hello',
+            status: 200,
+            uaid: request.uaid,
+            use_webpush: true,
+          }));
+        },
+        onRegister(request) {
+          this.serverSendMsg(JSON.stringify({
+            messageType: 'register',
+            status: 200,
+            uaid: userAgentID,
+            channelID: request.channelID,
+            pushEndpoint: 'https://example.org/push/new',
+          }));
+        }
+      });
+    },
+  });
+
+  let newRecord = yield PushNotificationService.register(
+    'https://example.com/page/3',
+    ChromeUtils.originAttributesToSuffix({ appId: Ci.nsIScriptSecurityManager.NO_APP_ID, inBrowser: false })
+  );
+  ok(newRecord.p256dhPublicKey, 'Should generate public keys for new records');
+  ok(newRecord.p256dhPrivateKey, 'Should generate private keys for new records');
+
+  let record = yield db.getByKeyID('eb18f12a-cc42-4f14-accb-3bfc1227f1aa');
+  ok(record.p256dhPublicKey, 'Should add public key to partial record');
+  ok(record.p256dhPrivateKey, 'Should add private key to partial record');
+
+  record = yield db.getByKeyID('0d8886b9-8da1-4778-8f5d-1cf93a877ed6');
+  deepEqual(record.p256dhPublicKey, publicKey,
+    'Should leave existing public key');
+  deepEqual(record.p256dhPrivateKey, privateKey,
+    'Should leave existing private key');
+});
--- a/dom/push/test/xpcshell/xpcshell.ini
+++ b/dom/push/test/xpcshell/xpcshell.ini
@@ -28,22 +28,23 @@ skip-if = toolkit == 'android'
 [test_registration_none.js]
 [test_registration_success.js]
 [test_unregister_empty_scope.js]
 [test_unregister_error.js]
 [test_unregister_invalid_json.js]
 [test_unregister_not_found.js]
 [test_unregister_success.js]
 [test_webapps_cleardata.js]
+[test_updateRecordNoEncryptionKeys_ws.js]
 #http2 test
 [test_resubscribe_4xxCode_http2.js]
 [test_resubscribe_5xxCode_http2.js]
 [test_resubscribe_listening_for_msg_error_http2.js]
 [test_register_5xxCode_http2.js]
-[test_updateRecordNoEncryptionKeys.js]
+[test_updateRecordNoEncryptionKeys_http2.js]
 [test_register_success_http2.js]
 skip-if = !hasNode
 run-sequentially = node server exceptions dont replay well
 [test_register_error_http2.js]
 skip-if = !hasNode
 run-sequentially = node server exceptions dont replay well
 [test_unregister_success_http2.js]
 skip-if = !hasNode