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 static com.android.cobalt.crypto.PublicKeys.ANALYZER_CONTEXT_INFO_BYTES;
20 import static com.android.cobalt.crypto.PublicKeys.ANALYZER_KEY_DEV;
21 import static com.android.cobalt.crypto.PublicKeys.ANALYZER_KEY_INDEX_DEV;
22 import static com.android.cobalt.crypto.PublicKeys.ANALYZER_KEY_INDEX_PROD;
23 import static com.android.cobalt.crypto.PublicKeys.ANALYZER_KEY_PROD;
24 import static com.android.cobalt.crypto.PublicKeys.SHUFFLER_CONTEXT_INFO_BYTES;
25 import static com.android.cobalt.crypto.PublicKeys.SHUFFLER_KEY_DEV;
26 import static com.android.cobalt.crypto.PublicKeys.SHUFFLER_KEY_INDEX_DEV;
27 import static com.android.cobalt.crypto.PublicKeys.SHUFFLER_KEY_INDEX_PROD;
28 import static com.android.cobalt.crypto.PublicKeys.SHUFFLER_KEY_PROD;
29 import static com.android.cobalt.crypto.PublicKeys.X25519_PUBLIC_VALUE_LEN;
30 
31 import androidx.annotation.NonNull;
32 
33 import com.android.cobalt.CobaltPipelineType;
34 import com.android.internal.annotations.VisibleForTesting;
35 
36 import com.google.cobalt.EncryptedMessage;
37 import com.google.cobalt.Envelope;
38 import com.google.cobalt.Observation;
39 import com.google.cobalt.ObservationToEncrypt;
40 import com.google.protobuf.ByteString;
41 import com.google.protobuf.MessageLite;
42 
43 import java.util.Objects;
44 import java.util.Optional;
45 
46 /** Handler for encryption of {@link Envelope} and {@link Observation} via {@link HpkeEncrypt}. */
47 public final class HpkeEncrypter implements Encrypter {
48     private final HpkeEncrypt mEncrypter;
49 
50     @VisibleForTesting final int mShufflerKeyIndex;
51     @VisibleForTesting final int mAnalyzerKeyIndex;
52 
53     private final byte[] mShufflerKey;
54     private final byte[] mAnalyzerKey;
55 
56     /** Creates a HpkeEncrypter compatible with the specified Cobalt environment */
createForEnvironment( @onNull HpkeEncrypt encrypter, @NonNull CobaltPipelineType type)57     public static HpkeEncrypter createForEnvironment(
58             @NonNull HpkeEncrypt encrypter, @NonNull CobaltPipelineType type) {
59         Objects.requireNonNull(type);
60 
61         switch (type) {
62             case PROD:
63                 return new HpkeEncrypter(
64                         encrypter,
65                         SHUFFLER_KEY_PROD,
66                         SHUFFLER_KEY_INDEX_PROD,
67                         ANALYZER_KEY_PROD,
68                         ANALYZER_KEY_INDEX_PROD);
69             case DEV:
70                 return new HpkeEncrypter(
71                         encrypter,
72                         SHUFFLER_KEY_DEV,
73                         SHUFFLER_KEY_INDEX_DEV,
74                         ANALYZER_KEY_DEV,
75                         ANALYZER_KEY_INDEX_DEV);
76         }
77 
78         throw new IllegalArgumentException("Unknown Cobalt environment");
79     }
80 
HpkeEncrypter( @onNull HpkeEncrypt encrypter, @NonNull byte[] shufflerKey, int shufflerKeyIndex, @NonNull byte[] analyzerKey, int analyzerKeyIndex)81     HpkeEncrypter(
82             @NonNull HpkeEncrypt encrypter,
83             @NonNull byte[] shufflerKey,
84             int shufflerKeyIndex,
85             @NonNull byte[] analyzerKey,
86             int analyzerKeyIndex) {
87         this.mEncrypter = Objects.requireNonNull(encrypter);
88         this.mShufflerKey = Objects.requireNonNull(shufflerKey);
89         this.mShufflerKeyIndex = shufflerKeyIndex;
90         this.mAnalyzerKey = Objects.requireNonNull(analyzerKey);
91         this.mAnalyzerKeyIndex = analyzerKeyIndex;
92     }
93 
94     /**
95      * Encrypts the provided {@link Envelope} with the key for the shuffler and wraps it into an
96      * {@link EncryptedMessage}.
97      *
98      * @return {@link EncryptedMessage} wrapped in an Optional if the {@link Envelope} is
99      *     successfully encrypted. Optional will be empty if the {@link Envelope} is empty
100      * @throws EncryptionFailedException if encryption fails
101      */
102     @Override
encryptEnvelope(@onNull Envelope envelope)103     public Optional<EncryptedMessage> encryptEnvelope(@NonNull Envelope envelope)
104             throws EncryptionFailedException {
105         Objects.requireNonNull(envelope);
106 
107         return encrypt(
108                 envelope,
109                 mShufflerKey,
110                 mShufflerKeyIndex,
111                 SHUFFLER_CONTEXT_INFO_BYTES,
112                 ByteString.EMPTY);
113     }
114 
115     /**
116      * Extract and encrypts {@link Observation} from the provided {@link ObservationToEncrypt} with
117      * the key for the analyzer and wraps it into an {@link EncryptedMessage}.
118      *
119      * @return {@link EncryptedMessage} wrapped in an Optional if the {@link Observation} is
120      *     successfully encrypted. Optional will be empty if the {@link Observation} is empty
121      * @throws EncryptionFailedException if encryption fails
122      */
123     @Override
encryptObservation( @onNull ObservationToEncrypt observationToEncrypt)124     public Optional<EncryptedMessage> encryptObservation(
125             @NonNull ObservationToEncrypt observationToEncrypt) throws EncryptionFailedException {
126         Objects.requireNonNull(observationToEncrypt);
127 
128         return encrypt(
129                 observationToEncrypt.getObservation(),
130                 mAnalyzerKey,
131                 mAnalyzerKeyIndex,
132                 ANALYZER_CONTEXT_INFO_BYTES,
133                 observationToEncrypt.getContributionId());
134     }
135 
136     /**
137      * Encrypt the given message and wraps it into an {@link EncryptedMessage.Builder}
138      *
139      * @param publicKey used by the encryption algorithm, must satisfies the encryption scheme
140      *     required key length
141      * @param contextInfoBytes used by the encryption algorithm, intended to provide additional data
142      *     keeping the message integrity. Cannot be empty
143      * @param contributionId passed by the Message to encrypt, used to set contributionId in {@link
144      *     EncryptedMessage}. This field should only be set when encrypting an {@link Observation}
145      *     that should be counted towards the shuffler threshold. All other Messages should pass a
146      *     ByteString.EMPTY
147      * @return {@link EncryptedMessage} wrapped in an Optional if the {@link MessageLite} is
148      *     successfully encrypted. Optional will be empty if the {@link MessageLite} is empty
149      * @throws EncryptionFailedException if encryption fails
150      */
encrypt( MessageLite message, byte[] publicKey, int keyIndex, byte[] contextInfoBytes, ByteString contributionId)151     private Optional<EncryptedMessage> encrypt(
152             MessageLite message,
153             byte[] publicKey,
154             int keyIndex,
155             byte[] contextInfoBytes,
156             ByteString contributionId)
157             throws EncryptionFailedException {
158         // Assert the public key length matches the X25519 public key requirement, and
159         // contextInfoBytes.
160         if (publicKey.length != X25519_PUBLIC_VALUE_LEN || contextInfoBytes.length == 0) {
161             throw new AssertionError(
162                     String.format(
163                             "Invalid HPKE parameters. Expected public key length of %d, got %d. "
164                                     + "Expected non-zero context info length, got %d",
165                             X25519_PUBLIC_VALUE_LEN, publicKey.length, contextInfoBytes.length));
166         }
167 
168         byte[] plainText = message.toByteArray();
169         if (plainText.length == 0) {
170             return Optional.empty();
171         }
172 
173         byte[] encryptedMessageBytes =
174                 mEncrypter.encrypt(publicKey, message.toByteArray(), contextInfoBytes);
175         if (encryptedMessageBytes.length == 0) {
176             throw new EncryptionFailedException("Message couldn't be encrypted.");
177         }
178 
179         return Optional.of(
180                 EncryptedMessage.newBuilder()
181                         .setCiphertext(ByteString.copyFrom(encryptedMessageBytes))
182                         .setKeyIndex(keyIndex)
183                         .setContributionId(contributionId)
184                         .build());
185     }
186 }
187