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 295992 d0a7044bb280b26897746b69829ea230e2e5ffec
parent 295991 0e0f3104478fcb2ed29be54f3e5ee113245378b8
child 295993 a39aa7c8e1ee5ed2d5e4c7de4a50f5d1da96ba85
push id5245
push userraliiev@mozilla.com
push dateThu, 29 Oct 2015 11:30:51 +0000
treeherdermozilla-beta@dac831dc1bd0 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersdragana, nsm
bugs1185544
milestone43.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
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