add fixed private key for exchange (#1448)

world-federation-of-advertisers · Feb 5, 2024 · 77fd9b3 · 77fd9b3
1 parent 1a59063
commit 77fd9b3
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 1 deletion.
diff --git a/src/main/kotlin/org/wfanet/panelmatch/client/deploy/ExchangeWorkflowDaemonFromFlags.kt b/src/main/kotlin/org/wfanet/panelmatch/client/deploy/ExchangeWorkflowDaemonFromFlags.kt
@@ -80,6 +80,7 @@ abstract class ExchangeWorkflowDaemonFromFlags : ExchangeWorkflowDaemon() {
       algorithm = flags.certAlgorithm,
       certificateAuthority = certificateAuthority,
       localName = identity.toName(),
+      fallbackPrivateKeyBlobKey = flags.fallbackPrivateKeyBlobKey,
     )
   }
 

diff --git a/src/main/kotlin/org/wfanet/panelmatch/client/deploy/ExchangeWorkflowFlags.kt b/src/main/kotlin/org/wfanet/panelmatch/client/deploy/ExchangeWorkflowFlags.kt
@@ -124,4 +124,16 @@ class ExchangeWorkflowFlags {
   )
   var maxParallelClaimedExchangeSteps by Delegates.notNull<Int>()
     private set
+
+  @Option(
+    names = ["--fallback-private-key-blob-key"],
+    defaultValue = "",
+    description =
+      [
+        "Fallback blob key for a kms-encrypted private signing key when a workflow does not generate one."
+      ],
+    required = true,
+  )
+  lateinit var fallbackPrivateKeyBlobKey: String
+    private set
 }
diff --git a/src/main/kotlin/org/wfanet/panelmatch/client/launcher/ExchangeTaskExecutor.kt b/src/main/kotlin/org/wfanet/panelmatch/client/launcher/ExchangeTaskExecutor.kt
@@ -108,7 +108,9 @@ class ExchangeTaskExecutor(
     val privateStorageClient: StorageClient =
       privateStorageSelector.getStorageClient(exchangeDateKey)
     if (!isAlreadyComplete(step, privateStorageClient)) {
+      logger.fine { "Running Step" }
       runStep(privateStorageClient)
+      logger.fine { "Writing Done Blob" }
       writeDoneBlob(step, privateStorageClient)
     }
     // The Kingdom will be able to detect if it's handing out duplicate tasks because it will
@@ -119,9 +121,12 @@ class ExchangeTaskExecutor(
   private suspend fun ExchangeContext.runStep(privateStorage: StorageClient) {
     timeout.runWithTimeout {
       val exchangeTask: ExchangeTask = exchangeTaskMapper.getExchangeTaskForStep(this@runStep)
+      logger.fine { "Reading Inputs" }
       val taskInput: Map<String, Blob> =
         if (exchangeTask.skipReadInput()) emptyMap() else readInputs(step, privateStorage)
+      logger.fine { "Executing Exchange Task" }
       val taskOutput: Map<String, Flow<ByteString>> = exchangeTask.execute(taskInput)
+      logger.fine { "Writing Outputs" }
       writeOutputs(step, taskOutput, privateStorage)
     }
   }

diff --git a/src/main/kotlin/org/wfanet/panelmatch/common/certificates/V2AlphaCertificateManager.kt b/src/main/kotlin/org/wfanet/panelmatch/common/certificates/V2AlphaCertificateManager.kt
@@ -30,6 +30,7 @@ import org.wfanet.measurement.common.crypto.jceProvider
 import org.wfanet.measurement.common.crypto.readCertificate
 import org.wfanet.panelmatch.common.ExchangeDateKey
 import org.wfanet.panelmatch.common.certificates.CertificateManager.KeyPair
+import org.wfanet.panelmatch.common.loggerFor
 import org.wfanet.panelmatch.common.secrets.MutableSecretMap
 import org.wfanet.panelmatch.common.secrets.SecretMap
 
@@ -47,6 +48,7 @@ class V2AlphaCertificateManager(
   private val algorithm: String,
   private val certificateAuthority: CertificateAuthority,
   private val localName: String,
+  private val fallbackPrivateKeyBlobKey: String? = null,
 ) : CertificateManager {
 
   private val x509CertCache = ConcurrentHashMap<String, X509Certificate>()
@@ -76,7 +78,13 @@ class V2AlphaCertificateManager(
   }
 
   override suspend fun getExchangeKeyPair(exchange: ExchangeDateKey): KeyPair {
-    val signingKeys = requireNotNull(getSigningKeys(exchange.path)) { "Missing keys for $exchange" }
+    val keyFromPrimaryPath = getSigningKeys(exchange.path)
+    val signingKeys =
+      if (keyFromPrimaryPath == null) {
+        checkNotNull(getSigningKeys(fallbackPrivateKeyBlobKey!!))
+      } else {
+        keyFromPrimaryPath
+      }
     val x509Certificate = getCertificate(exchange, localName, signingKeys.certResourceName)
     val privateKey = parsePrivateKey(signingKeys.privateKey)
     return KeyPair(x509Certificate, privateKey, signingKeys.certResourceName)
@@ -106,9 +114,12 @@ class V2AlphaCertificateManager(
       this.privateKey = privateKey.encoded.toByteString()
     }
 
+    logger.fine { "Writing private key to SecretMap" }
     privateKeys.put(exchange.path, signingKeys.toByteString())
+    logger.fine { "Finish writing private key to SecretMap" }
     x509CertCache[certResourceName] = x509
     signingKeysCache[exchange.path] = Optional.of(signingKeys)
+    logger.fine { "Returning certResourceName: $certResourceName" }
 
     return certResourceName
   }
@@ -146,4 +157,8 @@ class V2AlphaCertificateManager(
       readCertificate(certBytes)
     }
   }
+
+  companion object {
+    private val logger by loggerFor()
+  }
 }