[wip] feat (csas-migration): Check & backfill secret in trans AKV

src/Temporary/CallbackRetryPolicy.php

@@ -0,0 +1,33 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Keboola\ObjectEncryptor\Temporary;
+
+use Closure;
+use Retry\Policy\SimpleRetryPolicy;
+use Retry\RetryContextInterface;
+
+class CallbackRetryPolicy extends SimpleRetryPolicy
+{
+    private Closure $shouldRetryCallback;
+
+    public function __construct(
+        callable $shouldRetryCallback,
+        int $maxAttempts = 3,
+    ) {
+        parent::__construct($maxAttempts);
+        $this->shouldRetryCallback = $shouldRetryCallback(...);
+    }
+
+    public function canRetry(RetryContextInterface $context): bool
+    {
+        $e = $context->getLastException();
+
+        if (($this->shouldRetryCallback)($e, $context) !== true) {
+            return false;
+        }
+
+        return parent::canRetry($context);
+    }
+}

src/Wrapper/GenericAKVWrapper.php

@@ -8,17 +8,20 @@
 use Defuse\Crypto\Key;
 use Keboola\AzureKeyVaultClient\Authentication\AuthenticatorFactory;
 use Keboola\AzureKeyVaultClient\Client;
+use Keboola\AzureKeyVaultClient\Exception\ClientException;
 use Keboola\AzureKeyVaultClient\GuzzleClientFactory;
 use Keboola\AzureKeyVaultClient\Requests\SecretAttributes;
 use Keboola\AzureKeyVaultClient\Requests\SetSecretRequest;
 use Keboola\AzureKeyVaultClient\Responses\SecretBundle;
 use Keboola\ObjectEncryptor\EncryptorOptions;
 use Keboola\ObjectEncryptor\Exception\ApplicationException;
 use Keboola\ObjectEncryptor\Exception\UserException;
+use Keboola\ObjectEncryptor\Temporary\CallbackRetryPolicy;
 use Keboola\ObjectEncryptor\Temporary\TransClient;
 use Keboola\ObjectEncryptor\Temporary\TransClientNotAvailableException;
 use Psr\Log\NullLogger;
 use Retry\BackOff\ExponentialBackOffPolicy;
+use Retry\Policy\RetryPolicyInterface;
 use Retry\Policy\SimpleRetryPolicy;
 use Retry\RetryProxy;
 use Throwable;
@@ -81,9 +84,14 @@ public function getTransClient(): ?TransClient
         return $this->transClient ?: null;
     }
 
-    private function getRetryProxy(): RetryProxy
+    private static function getTransStackId(): ?string
     {
-        $retryPolicy = new SimpleRetryPolicy(3);
+        return (string) getenv('TRANS_ENCRYPTOR_STACK_ID') ?: null;
+    }
+
+    private function getRetryProxy(?RetryPolicyInterface $retryPolicy = null): RetryProxy
+    {
+        $retryPolicy ??= new SimpleRetryPolicy(3);
         $backOffPolicy = new ExponentialBackOffPolicy(1000);
         return new RetryProxy($retryPolicy, $backOffPolicy);
     }
@@ -171,8 +179,36 @@ public function decrypt(string $encryptedData): string
         ) {
             throw new UserException('Deciphering failed.');
         }
+
+        $metadata = $this->metadata;
+        $doBackfill = false;
+
+        // try retrieve secret from trans AKV
+        if ($this->getTransClient() !== null) {
+            $retryPolicy = new CallbackRetryPolicy(fn($e) =>
+                // do not retry if trans AKV response is 404
+                !$e instanceof ClientException || $e->getCode() !== 404);
+            try {
+                $decryptedContext = $this->getRetryProxy($retryPolicy)->call(function () use ($encrypted) {
+                    return $this->getTransClient()
+                        ?->getSecret($encrypted[self::SECRET_NAME])
+                        ->getValue();
+                });
+                if ($decryptedContext !== null && isset($this->metadata['stackId']) && self::getTransStackId()) {
+                    $metadata['stackId'] = self::getTransStackId();
+                }
+            } catch (ClientException $e) {
+                if ($e->getCode() === 404) {
+                    $doBackfill = true;
+                }
+            } catch (Throwable) {
+                // intentionally suppress all errors to prevent decrypt() from failing
+            }
+        }
+
         try {
-            $decryptedContext = $this->getRetryProxy()->call(function () use ($encrypted) {
+            // retrieve only if not found at trans AKV
+            $decryptedContext ??= $this->getRetryProxy()->call(function () use ($encrypted) {
                 return $this->getClient()
                     ->getSecret($encrypted[self::SECRET_NAME])
                     ->getValue();
@@ -188,12 +224,30 @@ public function decrypt(string $encryptedData): string
         } catch (Throwable $e) {
             throw new ApplicationException('Deciphering failed.', $e->getCode(), $e);
         }
-        $this->verifyMetadata($decryptedContext[self::METADATA_INDEX], $this->metadata);
+        $this->verifyMetadata($decryptedContext[self::METADATA_INDEX], $metadata);
         try {
             $key = Key::loadFromAsciiSafeString($decryptedContext[self::KEY_INDEX]);
             return Crypto::decrypt($encrypted[self::PAYLOAD_INDEX], $key, true);
         } catch (Throwable $e) {
+            $doBackfill = false;
             throw new ApplicationException('Deciphering failed.', $e->getCode(), $e);
+        } finally {
+            if ($doBackfill) {
+                if (isset($this->metadata['stackId']) && self::getTransStackId()) {
+                    $decryptedContext[self::METADATA_INDEX]['stackId'] = self::getTransStackId();
+                }
+                try {
+                    $this->getRetryProxy()->call(function () use ($encrypted, $decryptedContext) {
+                        $context = $this->encode($decryptedContext);
+                        $this->getTransClient()?->setSecret(
+                            new SetSecretRequest($context, new SecretAttributes()),
+                            $encrypted[self::SECRET_NAME],
+                        );
+                    });
+                } catch (Throwable) {
+                    // intentionally suppress all errors to prevent decrypt() from failing
+                }
+            }
         }
     }