|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, software |
| 13 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | + * See the License for the specific language governing permissions and |
| 16 | + * limitations under the License. |
| 17 | + */ |
| 18 | +package org.apache.beam.sdk.transforms; |
| 19 | + |
| 20 | +import java.util.Arrays; |
| 21 | +import javax.crypto.Cipher; |
| 22 | +import javax.crypto.Mac; |
| 23 | +import javax.crypto.spec.GCMParameterSpec; |
| 24 | +import javax.crypto.spec.SecretKeySpec; |
| 25 | +import org.apache.beam.sdk.coders.Coder; |
| 26 | +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; |
| 27 | +import org.apache.beam.sdk.coders.IterableCoder; |
| 28 | +import org.apache.beam.sdk.coders.KvCoder; |
| 29 | +import org.apache.beam.sdk.util.Secret; |
| 30 | +import org.apache.beam.sdk.values.KV; |
| 31 | +import org.apache.beam.sdk.values.PCollection; |
| 32 | + |
| 33 | +/** |
| 34 | + * A {@link PTransform} that provides a secure alternative to {@link |
| 35 | + * org.apache.beam.sdk.transforms.GroupByKey}. |
| 36 | + * |
| 37 | + * <p>This transform encrypts the keys of the input {@link PCollection}, performs a {@link |
| 38 | + * org.apache.beam.sdk.transforms.GroupByKey} on the encrypted keys, and then decrypts the keys in |
| 39 | + * the output. This is useful when the keys contain sensitive data that should not be stored at rest |
| 40 | + * by the runner. |
| 41 | + * |
| 42 | + * <p>The transform requires a {@link Secret} which returns a 32 byte secret which can be used to |
| 43 | + * generate a {@link SecretKeySpec} object using the HmacSHA256 algorithm. |
| 44 | + * |
| 45 | + * <p>Note the following caveats: 1) Runners can implement arbitrary materialization steps, so this |
| 46 | + * does not guarantee that the whole pipeline will not have unencrypted data at rest by itself. 2) |
| 47 | + * If using this transform in streaming mode, this transform may not properly handle update |
| 48 | + * compatibility checks around coders. This means that an improper update could lead to invalid |
| 49 | + * coders, causing pipeline failure or data corruption. If you need to update, make sure that the |
| 50 | + * input type passed into this transform does not change. |
| 51 | + */ |
| 52 | +public class GroupByEncryptedKey<K, V> |
| 53 | + extends PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>> { |
| 54 | + |
| 55 | + private final Secret hmacKey; |
| 56 | + |
| 57 | + private GroupByEncryptedKey(Secret hmacKey) { |
| 58 | + this.hmacKey = hmacKey; |
| 59 | + } |
| 60 | + |
| 61 | + /** |
| 62 | + * Creates a {@link GroupByEncryptedKey} transform. |
| 63 | + * |
| 64 | + * @param hmacKey The {@link Secret} key to use for encryption. |
| 65 | + * @param <K> The type of the keys in the input PCollection. |
| 66 | + * @param <V> The type of the values in the input PCollection. |
| 67 | + * @return A {@link GroupByEncryptedKey} transform. |
| 68 | + */ |
| 69 | + public static <K, V> GroupByEncryptedKey<K, V> create(Secret hmacKey) { |
| 70 | + return new GroupByEncryptedKey<>(hmacKey); |
| 71 | + } |
| 72 | + |
| 73 | + @Override |
| 74 | + public PCollection<KV<K, Iterable<V>>> expand(PCollection<KV<K, V>> input) { |
| 75 | + Coder<KV<K, V>> inputCoder = input.getCoder(); |
| 76 | + if (!(inputCoder instanceof KvCoder)) { |
| 77 | + throw new IllegalStateException("GroupByEncryptedKey requires its input to use KvCoder"); |
| 78 | + } |
| 79 | + KvCoder<K, V> inputKvCoder = (KvCoder<K, V>) inputCoder; |
| 80 | + Coder<K> keyCoder = inputKvCoder.getKeyCoder(); |
| 81 | + |
| 82 | + try { |
| 83 | + keyCoder.verifyDeterministic(); |
| 84 | + } catch (NonDeterministicException e) { |
| 85 | + throw new IllegalStateException( |
| 86 | + "the keyCoder of a GroupByEncryptedKey must be deterministic", e); |
| 87 | + } |
| 88 | + |
| 89 | + Coder<V> valueCoder = inputKvCoder.getValueCoder(); |
| 90 | + |
| 91 | + PCollection<KV<byte[], Iterable<KV<byte[], byte[]>>>> grouped = |
| 92 | + input |
| 93 | + .apply( |
| 94 | + "EncryptMessage", |
| 95 | + ParDo.of(new EncryptMessage<>(this.hmacKey, keyCoder, valueCoder))) |
| 96 | + .apply(GroupByKey.create()); |
| 97 | + |
| 98 | + return grouped |
| 99 | + .apply("DecryptMessage", ParDo.of(new DecryptMessage<>(this.hmacKey, keyCoder, valueCoder))) |
| 100 | + .setCoder(KvCoder.of(keyCoder, IterableCoder.of(valueCoder))); |
| 101 | + } |
| 102 | + |
| 103 | + /** |
| 104 | + * A {@link PTransform} that encrypts the key and value of an element. |
| 105 | + * |
| 106 | + * <p>The resulting PCollection will be a KV pair with the key being the HMAC of the encoded key, |
| 107 | + * and the value being a KV pair of the encrypted key and value. |
| 108 | + */ |
| 109 | + @SuppressWarnings("initialization.fields.uninitialized") |
| 110 | + private static class EncryptMessage<K, V> extends DoFn<KV<K, V>, KV<byte[], KV<byte[], byte[]>>> { |
| 111 | + private final Secret hmacKey; |
| 112 | + private final Coder<K> keyCoder; |
| 113 | + private final Coder<V> valueCoder; |
| 114 | + private transient Mac mac; |
| 115 | + private transient Cipher cipher; |
| 116 | + private transient SecretKeySpec secretKeySpec; |
| 117 | + private transient java.security.SecureRandom generator; |
| 118 | + |
| 119 | + EncryptMessage(Secret hmacKey, Coder<K> keyCoder, Coder<V> valueCoder) { |
| 120 | + this.hmacKey = hmacKey; |
| 121 | + this.keyCoder = keyCoder; |
| 122 | + this.valueCoder = valueCoder; |
| 123 | + } |
| 124 | + |
| 125 | + @Setup |
| 126 | + public void setup() { |
| 127 | + try { |
| 128 | + byte[] secretBytes = this.hmacKey.getSecretBytes(); |
| 129 | + this.mac = Mac.getInstance("HmacSHA256"); |
| 130 | + this.mac.init(new SecretKeySpec(secretBytes, "HmacSHA256")); |
| 131 | + this.cipher = Cipher.getInstance("AES/GCM/NoPadding"); |
| 132 | + this.secretKeySpec = new SecretKeySpec(secretBytes, "AES"); |
| 133 | + } catch (Exception ex) { |
| 134 | + throw new RuntimeException( |
| 135 | + "Failed to initialize cryptography libraries needed for GroupByEncryptedKey", ex); |
| 136 | + } |
| 137 | + this.generator = new java.security.SecureRandom(); |
| 138 | + } |
| 139 | + |
| 140 | + @ProcessElement |
| 141 | + public void processElement(ProcessContext c) throws Exception { |
| 142 | + byte[] encodedKey = encode(this.keyCoder, c.element().getKey()); |
| 143 | + byte[] encodedValue = encode(this.valueCoder, c.element().getValue()); |
| 144 | + |
| 145 | + byte[] hmac = this.mac.doFinal(encodedKey); |
| 146 | + |
| 147 | + byte[] keyIv = new byte[12]; |
| 148 | + byte[] valueIv = new byte[12]; |
| 149 | + this.generator.nextBytes(keyIv); |
| 150 | + this.generator.nextBytes(valueIv); |
| 151 | + GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(128, keyIv); |
| 152 | + this.cipher.init(Cipher.ENCRYPT_MODE, this.secretKeySpec, gcmParameterSpec); |
| 153 | + byte[] encryptedKey = this.cipher.doFinal(encodedKey); |
| 154 | + gcmParameterSpec = new GCMParameterSpec(128, valueIv); |
| 155 | + this.cipher.init(Cipher.ENCRYPT_MODE, this.secretKeySpec, gcmParameterSpec); |
| 156 | + byte[] encryptedValue = this.cipher.doFinal(encodedValue); |
| 157 | + |
| 158 | + c.output( |
| 159 | + KV.of( |
| 160 | + hmac, |
| 161 | + KV.of( |
| 162 | + com.google.common.primitives.Bytes.concat(keyIv, encryptedKey), |
| 163 | + com.google.common.primitives.Bytes.concat(valueIv, encryptedValue)))); |
| 164 | + } |
| 165 | + |
| 166 | + private <T> byte[] encode(Coder<T> coder, T value) throws Exception { |
| 167 | + java.io.ByteArrayOutputStream os = new java.io.ByteArrayOutputStream(); |
| 168 | + coder.encode(value, os); |
| 169 | + return os.toByteArray(); |
| 170 | + } |
| 171 | + } |
| 172 | + |
| 173 | + /** |
| 174 | + * A {@link PTransform} that decrypts the key and values of an element. |
| 175 | + * |
| 176 | + * <p>The input PCollection will be a KV pair with the key being the HMAC of the encoded key, and |
| 177 | + * the value being a list of KV pairs of the encrypted key and value. |
| 178 | + * |
| 179 | + * <p>This will return a tuple containing the decrypted key and a list of decrypted values. |
| 180 | + * |
| 181 | + * <p>Since there is some loss of precision in the HMAC encoding of the key (but not the key |
| 182 | + * encryption), there is some extra work done here to ensure that all key/value pairs are mapped |
| 183 | + * out appropriately. |
| 184 | + */ |
| 185 | + @SuppressWarnings("initialization.fields.uninitialized") |
| 186 | + private static class DecryptMessage<K, V> |
| 187 | + extends DoFn<KV<byte[], Iterable<KV<byte[], byte[]>>>, KV<K, Iterable<V>>> { |
| 188 | + private final Secret hmacKey; |
| 189 | + private final Coder<K> keyCoder; |
| 190 | + private final Coder<V> valueCoder; |
| 191 | + private transient Cipher cipher; |
| 192 | + private transient SecretKeySpec secretKeySpec; |
| 193 | + |
| 194 | + DecryptMessage(Secret hmacKey, Coder<K> keyCoder, Coder<V> valueCoder) { |
| 195 | + this.hmacKey = hmacKey; |
| 196 | + this.keyCoder = keyCoder; |
| 197 | + this.valueCoder = valueCoder; |
| 198 | + } |
| 199 | + |
| 200 | + @Setup |
| 201 | + public void setup() { |
| 202 | + try { |
| 203 | + this.cipher = Cipher.getInstance("AES/GCM/NoPadding"); |
| 204 | + this.secretKeySpec = new SecretKeySpec(this.hmacKey.getSecretBytes(), "AES"); |
| 205 | + } catch (Exception ex) { |
| 206 | + throw new RuntimeException( |
| 207 | + "Failed to initialize cryptography libraries needed for GroupByEncryptedKey", ex); |
| 208 | + } |
| 209 | + } |
| 210 | + |
| 211 | + @ProcessElement |
| 212 | + public void processElement(ProcessContext c) throws Exception { |
| 213 | + java.util.Map<K, java.util.List<V>> decryptedKvs = new java.util.HashMap<>(); |
| 214 | + for (KV<byte[], byte[]> encryptedKv : c.element().getValue()) { |
| 215 | + byte[] iv = Arrays.copyOfRange(encryptedKv.getKey(), 0, 12); |
| 216 | + GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(128, iv); |
| 217 | + this.cipher.init(Cipher.DECRYPT_MODE, this.secretKeySpec, gcmParameterSpec); |
| 218 | + |
| 219 | + byte[] encryptedKey = |
| 220 | + Arrays.copyOfRange(encryptedKv.getKey(), 12, encryptedKv.getKey().length); |
| 221 | + byte[] decryptedKeyBytes = this.cipher.doFinal(encryptedKey); |
| 222 | + K key = decode(this.keyCoder, decryptedKeyBytes); |
| 223 | + |
| 224 | + if (key != null) { |
| 225 | + if (!decryptedKvs.containsKey(key)) { |
| 226 | + decryptedKvs.put(key, new java.util.ArrayList<>()); |
| 227 | + } |
| 228 | + |
| 229 | + iv = Arrays.copyOfRange(encryptedKv.getValue(), 0, 12); |
| 230 | + gcmParameterSpec = new GCMParameterSpec(128, iv); |
| 231 | + this.cipher.init(Cipher.DECRYPT_MODE, this.secretKeySpec, gcmParameterSpec); |
| 232 | + |
| 233 | + byte[] encryptedValue = |
| 234 | + Arrays.copyOfRange(encryptedKv.getValue(), 12, encryptedKv.getValue().length); |
| 235 | + byte[] decryptedValueBytes = this.cipher.doFinal(encryptedValue); |
| 236 | + V value = decode(this.valueCoder, decryptedValueBytes); |
| 237 | + decryptedKvs.get(key).add(value); |
| 238 | + } else { |
| 239 | + throw new RuntimeException( |
| 240 | + "Found null key when decoding " + Arrays.toString(decryptedKeyBytes)); |
| 241 | + } |
| 242 | + } |
| 243 | + |
| 244 | + for (java.util.Map.Entry<K, java.util.List<V>> entry : decryptedKvs.entrySet()) { |
| 245 | + c.output(KV.of(entry.getKey(), entry.getValue())); |
| 246 | + } |
| 247 | + } |
| 248 | + |
| 249 | + private <T> T decode(Coder<T> coder, byte[] bytes) throws Exception { |
| 250 | + java.io.ByteArrayInputStream is = new java.io.ByteArrayInputStream(bytes); |
| 251 | + return coder.decode(is); |
| 252 | + } |
| 253 | + } |
| 254 | +} |
0 commit comments