1 /* 2 * Copyright (C) 2023 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package com.android.cobalt.crypto; 18 19 import com.android.cobalt.CobaltPipelineType; 20 import com.android.internal.annotations.VisibleForTesting; 21 22 import com.google.cobalt.EncryptedMessage; 23 import com.google.cobalt.Envelope; 24 import com.google.cobalt.Observation; 25 import com.google.cobalt.ObservationToEncrypt; 26 import com.google.protobuf.ByteString; 27 import com.google.protobuf.MessageLite; 28 29 import java.util.Objects; 30 import java.util.Optional; 31 32 /** Handler for encryption of {@link Envelope} and {@link Observation} via {@link HpkeEncrypt}. */ 33 public final class HpkeEncrypter implements Encrypter { 34 private final HpkeEncrypt mEncrypter; 35 private final PublicEncryptionKeys mPublicEncryptionKeys; 36 37 @VisibleForTesting final int mShufflerKeyIndex; 38 @VisibleForTesting final int mAnalyzerKeyIndex; 39 40 private final byte[] mShufflerKey; 41 private final byte[] mAnalyzerKey; 42 43 /** Creates a HpkeEncrypter compatible with the specified Cobalt environment */ createForEnvironment( HpkeEncrypt encrypter, CobaltPipelineType type, PublicEncryptionKeys publicEncryptionKeys)44 public static HpkeEncrypter createForEnvironment( 45 HpkeEncrypt encrypter, 46 CobaltPipelineType type, 47 PublicEncryptionKeys publicEncryptionKeys) { 48 Objects.requireNonNull(encrypter, "HpkeEncrypt cannot be null"); 49 Objects.requireNonNull(type, "CobaltPipelineType cannot be null"); 50 Objects.requireNonNull(publicEncryptionKeys, "PublicEncryptionKeys cannot be null"); 51 52 switch (type) { 53 case PROD: 54 return new HpkeEncrypter( 55 encrypter, 56 publicEncryptionKeys, 57 publicEncryptionKeys.getShufflerKeyProd(), 58 publicEncryptionKeys.getShufflerKeyIndexProd(), 59 publicEncryptionKeys.getAnalyzerKeyProd(), 60 publicEncryptionKeys.getAnalyzerKeyIndexProd()); 61 case DEV: 62 return new HpkeEncrypter( 63 encrypter, 64 publicEncryptionKeys, 65 publicEncryptionKeys.getShufflerKeyDev(), 66 publicEncryptionKeys.getShufflerKeyIndexDev(), 67 publicEncryptionKeys.getAnalyzerKeyDev(), 68 publicEncryptionKeys.getAnalyzerKeyIndexDev()); 69 } 70 71 throw new IllegalArgumentException("Unknown Cobalt environment" + type); 72 } 73 74 @VisibleForTesting HpkeEncrypter( HpkeEncrypt encrypter, PublicEncryptionKeys publicEncryptionKeys, byte[] shufflerKey, int shufflerKeyIndex, byte[] analyzerKey, int analyzerKeyIndex)75 HpkeEncrypter( 76 HpkeEncrypt encrypter, 77 PublicEncryptionKeys publicEncryptionKeys, 78 byte[] shufflerKey, 79 int shufflerKeyIndex, 80 byte[] analyzerKey, 81 int analyzerKeyIndex) { 82 this.mEncrypter = Objects.requireNonNull(encrypter, "HpkeEncrypt cannot be null"); 83 this.mPublicEncryptionKeys = 84 Objects.requireNonNull(publicEncryptionKeys, "PublicEncryptionKeys cannot be null"); 85 this.mShufflerKey = Objects.requireNonNull(shufflerKey, "Shuffler key cannot be null"); 86 this.mShufflerKeyIndex = shufflerKeyIndex; 87 this.mAnalyzerKey = Objects.requireNonNull(analyzerKey, "Analyzer key cannot be null"); 88 this.mAnalyzerKeyIndex = analyzerKeyIndex; 89 } 90 91 /** 92 * Encrypts the provided {@link Envelope} with the key for the shuffler and wraps it into an 93 * {@link EncryptedMessage}. 94 * 95 * @return {@link EncryptedMessage} wrapped in an Optional if the {@link Envelope} is 96 * successfully encrypted. Optional will be empty if the {@link Envelope} is empty 97 * @throws EncryptionFailedException if encryption fails 98 */ 99 @Override encryptEnvelope(Envelope envelope)100 public Optional<EncryptedMessage> encryptEnvelope(Envelope envelope) 101 throws EncryptionFailedException { 102 Objects.requireNonNull(envelope, "Envelope cannot be null"); 103 104 return encrypt( 105 envelope, 106 mShufflerKey, 107 mShufflerKeyIndex, 108 mPublicEncryptionKeys.getShufflerContextInfoBytes(), 109 ByteString.EMPTY); 110 } 111 112 /** 113 * Extracts and encrypts {@link Observation} from the provided {@link ObservationToEncrypt} with 114 * the key for the analyzer and wraps it into an {@link EncryptedMessage}. 115 * 116 * @return {@link EncryptedMessage} wrapped in an Optional if the {@link Observation} is 117 * successfully encrypted. Optional will be empty if the {@link Observation} is empty 118 * @throws EncryptionFailedException if encryption fails 119 */ 120 @Override encryptObservation(ObservationToEncrypt observationToEncrypt)121 public Optional<EncryptedMessage> encryptObservation(ObservationToEncrypt observationToEncrypt) 122 throws EncryptionFailedException { 123 Objects.requireNonNull(observationToEncrypt, "ObservationToEncrypt cannot be null"); 124 125 return encrypt( 126 observationToEncrypt.getObservation(), 127 mAnalyzerKey, 128 mAnalyzerKeyIndex, 129 mPublicEncryptionKeys.getAnalyzerContextInfoBytes(), 130 observationToEncrypt.getContributionId()); 131 } 132 133 /** 134 * Encrypts the given message and wraps it into an {@link EncryptedMessage.Builder} 135 * 136 * @param publicKey used by the encryption algorithm, must satisfies the encryption scheme 137 * required key length 138 * @param contextInfoBytes used by the encryption algorithm, intended to provide additional data 139 * keeping the message integrity. Cannot be empty 140 * @param contributionId passed by the Message to encrypt, used to set contributionId in {@link 141 * EncryptedMessage}. This field should only be set when encrypting an {@link Observation} 142 * that should be counted towards the shuffler threshold. All other Messages should pass a 143 * ByteString.EMPTY 144 * @return {@link EncryptedMessage} wrapped in an Optional if the {@link MessageLite} is 145 * successfully encrypted. Optional will be empty if the {@link MessageLite} is empty 146 * @throws EncryptionFailedException if encryption fails 147 */ encrypt( MessageLite message, byte[] publicKey, int keyIndex, byte[] contextInfoBytes, ByteString contributionId)148 private Optional<EncryptedMessage> encrypt( 149 MessageLite message, 150 byte[] publicKey, 151 int keyIndex, 152 byte[] contextInfoBytes, 153 ByteString contributionId) 154 throws EncryptionFailedException { 155 // Assert the public key length matches the X25519 public key requirement, and 156 // contextInfoBytes. 157 int x25519PublicValueLen = mPublicEncryptionKeys.getX25519PublicValueLen(); 158 if (publicKey.length != x25519PublicValueLen || contextInfoBytes.length == 0) { 159 throw new AssertionError( 160 String.format( 161 "Invalid HPKE parameters. Expected public key length of %d, got %d. " 162 + "Expected non-zero context info length, got %d", 163 x25519PublicValueLen, publicKey.length, contextInfoBytes.length)); 164 } 165 166 byte[] plainText = message.toByteArray(); 167 if (plainText.length == 0) { 168 return Optional.empty(); 169 } 170 171 byte[] encryptedMessageBytes = 172 mEncrypter.encrypt(publicKey, message.toByteArray(), contextInfoBytes); 173 if (encryptedMessageBytes.length == 0) { 174 throw new EncryptionFailedException("Message couldn't be encrypted."); 175 } 176 177 return Optional.of( 178 EncryptedMessage.newBuilder() 179 .setCiphertext(ByteString.copyFrom(encryptedMessageBytes)) 180 .setKeyIndex(keyIndex) 181 .setContributionId(contributionId) 182 .build()); 183 } 184 } 185