1 // Copyright 2020 Google LLC 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 // 15 //////////////////////////////////////////////////////////////////////////////// 16 package com.google.crypto.tink.prf; 17 18 import com.google.crypto.tink.internal.LegacyProtoKey; 19 import com.google.crypto.tink.internal.MonitoringClient; 20 import com.google.crypto.tink.internal.MonitoringKeysetInfo; 21 import com.google.crypto.tink.internal.MonitoringUtil; 22 import com.google.crypto.tink.internal.MutableMonitoringRegistry; 23 import com.google.crypto.tink.internal.MutablePrimitiveRegistry; 24 import com.google.crypto.tink.internal.PrimitiveConstructor; 25 import com.google.crypto.tink.internal.PrimitiveRegistry; 26 import com.google.crypto.tink.internal.PrimitiveSet; 27 import com.google.crypto.tink.internal.PrimitiveWrapper; 28 import com.google.crypto.tink.prf.internal.LegacyFullPrf; 29 import com.google.errorprone.annotations.Immutable; 30 import java.security.GeneralSecurityException; 31 import java.util.Collections; 32 import java.util.HashMap; 33 import java.util.List; 34 import java.util.Map; 35 36 /** 37 * PrfSetWrapper is the implementation of PrimitiveWrapper for the PrfSet primitive. 38 * 39 * <p>The returned primitive has instances of {@code Prf} for each key in the KeySet. The individual 40 * Prf instances can then be used to compute psuedo-random sequences from the underlying key. 41 */ 42 @Immutable 43 public class PrfSetWrapper implements PrimitiveWrapper<Prf, PrfSet> { 44 45 private static final PrfSetWrapper WRAPPER = new PrfSetWrapper(); 46 private static final PrimitiveConstructor<LegacyProtoKey, Prf> 47 LEGACY_FULL_PRF_PRIMITIVE_CONSTRUCTOR = 48 PrimitiveConstructor.create(LegacyFullPrf::create, LegacyProtoKey.class, Prf.class); 49 50 private static class WrappedPrfSet extends PrfSet { 51 // This map is constructed using Collections.unmodifiableMap 52 @SuppressWarnings("Immutable") 53 private final Map<Integer, Prf> keyIdToPrfMap; 54 55 private final int primaryKeyId; 56 57 @Immutable 58 private static class PrfWithMonitoring implements Prf { 59 private final Prf prf; 60 private final int keyId; 61 62 @SuppressWarnings("Immutable") 63 private final MonitoringClient.Logger logger; 64 65 @Override compute(byte[] input, int outputLength)66 public byte[] compute(byte[] input, int outputLength) throws GeneralSecurityException { 67 try { 68 byte[] output = prf.compute(input, outputLength); 69 logger.log(keyId, input.length); 70 return output; 71 } catch (GeneralSecurityException e) { 72 logger.logFailure(); 73 throw e; 74 } 75 } 76 PrfWithMonitoring(Prf prf, int keyId, MonitoringClient.Logger logger)77 public PrfWithMonitoring(Prf prf, int keyId, MonitoringClient.Logger logger) { 78 this.prf = prf; 79 this.keyId = keyId; 80 this.logger = logger; 81 } 82 } 83 WrappedPrfSet(PrimitiveSet<Prf> primitives)84 private WrappedPrfSet(PrimitiveSet<Prf> primitives) throws GeneralSecurityException { 85 if (primitives.getRawPrimitives().isEmpty()) { 86 throw new GeneralSecurityException("No primitives provided."); 87 } 88 if (primitives.getPrimary() == null) { 89 throw new GeneralSecurityException("Primary key not set."); 90 } 91 MonitoringClient.Logger logger; 92 if (primitives.hasAnnotations()) { 93 MonitoringClient client = MutableMonitoringRegistry.globalInstance().getMonitoringClient(); 94 MonitoringKeysetInfo keysetInfo = MonitoringUtil.getMonitoringKeysetInfo(primitives); 95 logger = client.createLogger(keysetInfo, "prf", "compute"); 96 } else { 97 logger = MonitoringUtil.DO_NOTHING_LOGGER; 98 } 99 100 primaryKeyId = primitives.getPrimary().getKeyId(); 101 List<PrimitiveSet.Entry<Prf>> entries = primitives.getRawPrimitives(); 102 Map<Integer, Prf> mutablePrfMap = new HashMap<>(); 103 for (PrimitiveSet.Entry<Prf> entry : entries) { 104 // Likewise, the key IDs of the PrfSet passed 105 mutablePrfMap.put( 106 entry.getKeyId(), 107 new PrfWithMonitoring(entry.getFullPrimitive(), entry.getKeyId(), logger)); 108 } 109 keyIdToPrfMap = Collections.unmodifiableMap(mutablePrfMap); 110 } 111 112 @Override getPrimaryId()113 public int getPrimaryId() { 114 return primaryKeyId; 115 } 116 117 @Override getPrfs()118 public Map<Integer, Prf> getPrfs() throws GeneralSecurityException { 119 return keyIdToPrfMap; 120 } 121 } 122 123 @Override wrap(PrimitiveSet<Prf> set)124 public PrfSet wrap(PrimitiveSet<Prf> set) throws GeneralSecurityException { 125 return new WrappedPrfSet(set); 126 } 127 128 @Override getPrimitiveClass()129 public Class<PrfSet> getPrimitiveClass() { 130 return PrfSet.class; 131 } 132 133 @Override getInputPrimitiveClass()134 public Class<Prf> getInputPrimitiveClass() { 135 return Prf.class; 136 } 137 register()138 public static void register() throws GeneralSecurityException { 139 MutablePrimitiveRegistry.globalInstance().registerPrimitiveWrapper(WRAPPER); 140 MutablePrimitiveRegistry.globalInstance() 141 .registerPrimitiveConstructor(LEGACY_FULL_PRF_PRIMITIVE_CONSTRUCTOR); 142 } 143 144 /** 145 * registerToInternalPrimitiveRegistry is a non-public method (it takes an argument of an 146 * internal-only type) registering an instance of {@code PrfSetWrapper} to the provided {@code 147 * PrimitiveRegistry.Builder}. 148 */ registerToInternalPrimitiveRegistry( PrimitiveRegistry.Builder primitiveRegistryBuilder)149 public static void registerToInternalPrimitiveRegistry( 150 PrimitiveRegistry.Builder primitiveRegistryBuilder) throws GeneralSecurityException { 151 primitiveRegistryBuilder.registerPrimitiveWrapper(WRAPPER); 152 } 153 } 154