WebPush flow from start to finish (RFC 8291 & RFC 8188 (without the record chunking part over HTTP))

WebPush.java

package io.github.asvanberg.http.webpush;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.Mac;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.security.GeneralSecurityException;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.interfaces.ECPublicKey;
import java.security.spec.ECGenParameterSpec;
import java.util.*;

public class WebPush {

    private static final int AUTHENTICATION_TAG_LENGTH_IN_BITS = 16 * 8; // 128 bits

    public static void main(String[] args) throws GeneralSecurityException {
        KeyPairGenerator xdh = KeyPairGenerator.getInstance("EC");
        xdh.initialize(new ECGenParameterSpec("secp256r1"));
        
        // Step 0 - if your user agents push service requires VAPID authentication (RFC 8292)
        // Generate and store a ECDH keypair
        // Use the public key when subscribing (applicationServerKey)
        // Use the private key to sign your credentials (the JWT)
        // Send signed JWT and public key according to the RFC

        // Step 1 - Subscribe (https://developer.mozilla.org/en-US/docs/Web/API/PushManager/subscribe)
        // This key pair is generated by the user agent and the public key is exposed as keys.p256d
        // The user agent also generates an authenticationSecret (salt) exposed as keys.auth
        // Both values are encoded as URL-safe base64 without padding
        // Send both, plus the endpoint, to your application server and save them
        KeyPair userAgent = xdh.generateKeyPair();
        byte[] authenticationSecret = generateSalt();
        String endpoint = "absolute url is in the subscription along with the keys";
        
        // Step 2 - Generate a keypair (this is ephemeral, only for the one message, no need to save)
        final KeyPair applicationServer = xdh.generateKeyPair();

        // Step 3 - Send a push notification
        // The data here is the raw form of what would normally be encoded using
        // Encrypted Content-Encoding for HTTP (RFC 8188)
        String original = "Hello world!";
        ApplicationServerOutput data = applicationServerSendWebPush(userAgent.getPublic(), authenticationSecret, original, applicationServer);

        // Step 4 - This is done by the user agent
        // Missing is the step of decoding the data that would be in RFC 8188 format
        String roundTrippedPlaintext = userAgentReceiveWebPush(userAgent, authenticationSecret, data);

        // Step 5 - Display the user notification in some way
        System.out.println(roundTrippedPlaintext);
    }

    /**
     * The "output" of ECE encrypting {@code data}.
     * @param applicationServerPublic the "keyid" in RFC 8188
     * @param salt
     * @param data the combined data of all records
     */
    record ApplicationServerOutput(PublicKey applicationServerPublic, byte[] salt, byte[] data) {}

    private static ApplicationServerOutput applicationServerSendWebPush(PublicKey userAgentPublic, byte[] authenticationSecret, String data, KeyPair applicationServer) throws
            GeneralSecurityException
    {

        byte[] sharedSecret = sharedSecret(userAgentPublic, applicationServer.getPrivate());
        byte[] pseudoRandomKey = extract(sharedSecret, authenticationSecret);

        byte[] info = webPushInfo(userAgentPublic, applicationServer.getPublic());
        byte[] ikm = expand(pseudoRandomKey, info, 32);

        final ApplicationServerECE applicationServerECE = eceEncrypt(ikm, data);
        return new ApplicationServerOutput(applicationServer.getPublic(), applicationServerECE.salt(), applicationServerECE.data());
    }

    record ApplicationServerECE(byte[] salt, byte[] data) {}

    private static ApplicationServerECE eceEncrypt(byte[] ikm, String data) throws GeneralSecurityException {
        byte[] salt = generateSalt();
        byte[] pseudoRandomKey = extract(ikm, salt);
        KeyAndNonce keyAndNonce = deriveKeyAndNonce(pseudoRandomKey);

        final Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(keyAndNonce.contentEncryptionKey(), "AES"), new GCMParameterSpec(AUTHENTICATION_TAG_LENGTH_IN_BITS, keyAndNonce.nonce()));
        byte[] encrypted = cipher.doFinal(data.getBytes());
        return new ApplicationServerECE(salt, encrypted);
    }

    private static String userAgentReceiveWebPush(KeyPair userAgent, byte[] authenticationSecret, ApplicationServerOutput data) throws GeneralSecurityException {
        byte[] sharedSecret = sharedSecret(data.applicationServerPublic(), userAgent.getPrivate());
        byte[] pseudoRandomKey = extract(sharedSecret, authenticationSecret);
        byte[] ikm = expand(pseudoRandomKey, webPushInfo(userAgent.getPublic(), data.applicationServerPublic()), 32);

        return eceDecrypt(ikm, data.salt(), data.data());
    }

    private static String eceDecrypt(byte[] ikm, byte[] salt, byte[] data) throws GeneralSecurityException {
        byte[] pseudoRandomKey = extract(ikm, salt);
        KeyAndNonce keyAndNonce = deriveKeyAndNonce(pseudoRandomKey);

        final Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(keyAndNonce.contentEncryptionKey(), "AES"), new GCMParameterSpec(AUTHENTICATION_TAG_LENGTH_IN_BITS, keyAndNonce.nonce()));
        return new String(cipher.doFinal(data));
    }

    private static byte[] webPushInfo(PublicKey userAgentPublic, final PublicKey applicationServerPublic) {
        ByteBuffer buffer = ByteBuffer.allocate(15 + 65 + 65);
        buffer.put("WebPush: info\0".getBytes());
        encode(userAgentPublic, buffer);
        encode(applicationServerPublic, buffer);
        return buffer.array();
    }

    record KeyAndNonce(byte[] contentEncryptionKey, byte[] nonce) {}

    private static KeyAndNonce deriveKeyAndNonce(byte[] pseudoRandomKey) throws GeneralSecurityException {
        byte[] contentEncryptionKeyInfo = "Content-Encoding: aes128gcm\0".getBytes();
        byte[] contentEncryptionKey = expand(pseudoRandomKey, contentEncryptionKeyInfo, 16);
        byte[] nonceInfo = "Content-ENcoding: nonce\0".getBytes();
        byte[] nonce = expand(pseudoRandomKey, nonceInfo, 12);
        return new KeyAndNonce(contentEncryptionKey, nonce);
    }

    /**
     * Encodes the elliptic curve public key as uncompressed X9.62 into the specified buffer
     */
    private static void encode(PublicKey publicKey, ByteBuffer buffer) {
        if (publicKey instanceof ECPublicKey ecPublicKey) {
            buffer.put((byte) 4);
            byte[] x = ecPublicKey.getW().getAffineX().toByteArray();
            buffer.put(x, 1, x.length - 1);
            byte[] y = ecPublicKey.getW().getAffineY().toByteArray();
            buffer.put(y, 1, y.length - 1);
        }
        else {
            throw new IllegalArgumentException("Must be an elliptic curve key");
        }
    }

    /**
     * @return a shared secret computed using Elliptic-curve Diffie–Hellman
     */
    private static byte[] sharedSecret(PublicKey publicKey, PrivateKey privateKey) throws GeneralSecurityException {
        KeyAgreement ecdh = KeyAgreement.getInstance("ECDH");
        ecdh.init(privateKey);
        ecdh.doPhase(publicKey, true);
        return ecdh.generateSecret();
    }

    /**
     * @return a salt of length 16
     */
    private static byte[] generateSalt() {
        byte[] salt = new byte[16];
        SecureRandom secureRandom = new SecureRandom();
        secureRandom.nextBytes(salt);
        return salt;
    }

    /**
     * @return a pseudo random key
     */
    private static byte[] extract(byte[] ikm, byte[] salt) throws GeneralSecurityException {
        Mac hmacSHA256 = Mac.getInstance("HmacSHA256");
        hmacSHA256.init(new SecretKeySpec(salt, "HMacSHA256"));
        return hmacSHA256.doFinal(ikm);
    }

    /**
     * @return a cryptographically stronger key of the specified length
     */
    private static byte[] expand(byte[] prk, byte[] info, int length) throws GeneralSecurityException {
        Mac hmacSHA256 = Mac.getInstance("HmacSHA256");
        hmacSHA256.init(new SecretKeySpec(prk, "HMacSHA256"));
        hmacSHA256.update(info);
        hmacSHA256.update((byte) 1);
        return Arrays.copyOf(hmacSHA256.doFinal(), length);
    }
}