From f0927aac073bcbaafaae36c196101419bbc98805 Mon Sep 17 00:00:00 2001
From: Nikos Dragazis <nikolaos.dragazis@scylladb.com>
Date: Wed, 16 Apr 2025 19:28:29 +0300
Subject: [PATCH] encryption: azure: Delegate hosts to shard 0

As in the AWS and GCP hosts, make all Azure hosts delegate their traffic
to shard 0 to avoid creating too many data encryption keys and API
calls to Key Vault.

Signed-off-by: Nikos Dragazis <nikolaos.dragazis@scylladb.com>
---
 ent/encryption/azure_host.cc | 34 +++++++++++++++++++++++++++-------
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/ent/encryption/azure_host.cc b/ent/encryption/azure_host.cc
index daa2707c80..951f9d98c7 100644
--- a/ent/encryption/azure_host.cc
+++ b/ent/encryption/azure_host.cc
@@ -79,12 +79,13 @@ public:
     static inline constexpr std::chrono::milliseconds default_expiry = 600s;
     static inline constexpr std::chrono::milliseconds default_refresh = 1200s;
 
-    impl(const std::string& name, const host_options&);
+    impl(encryption_context*, const std::string& name, const host_options& options);
     future<> init();
     const host_options& options() const;
     future<key_and_id_type> get_or_create_key(const key_info&, const option_override* = nullptr);
     future<key_ptr> get_key_by_id(const id_type&, const key_info&);
 private:
+    encryption_context* _ctxt;
     const std::string _name;
     const std::string _log_prefix;
     const host_options _options;
@@ -147,8 +148,9 @@ private:
     future<bytes> find_key(const id_cache_key&);
 };
 
-azure_host::impl::impl(const std::string& name, const host_options& options)
-    : _name(name)
+azure_host::impl::impl(encryption_context* ctxt, const std::string& name, const azure_host::host_options& options)
+    : _ctxt(ctxt)
+    , _name(name)
     , _log_prefix(fmt::format("AzureVault:{}", name))
     , _options(options)
     , _credentials()
@@ -415,6 +417,17 @@ future<rjson::value> azure_host::impl::send_request(const sstring& host, unsigne
 future<azure_host::key_and_id_type> azure_host::impl::create_key(const attr_cache_key& k) {
     static const vault_log_filter filter{vault_log_filter::op_type::wrapkey};
     auto& info = k.info;
+    if (_ctxt && this_shard_id() != 0) {
+        auto [data, id] = co_await smp::submit_to(0, [this, k]() -> future<std::tuple<bytes, id_type>> {
+            auto host = _ctxt->get_azure_host(_name);
+            auto [key, id] = co_await host->_impl->_attr_cache.get(k);
+            co_return std::make_tuple(key != nullptr ? key->key() : bytes{}, id);
+        });
+        co_return key_and_id_type{
+            data.empty() ? nullptr : make_shared<symmetric_key>(info, data),
+            id
+        };
+    }
     azlog.debug("[{}] Creating new key: {}", _log_prefix, info);
     auto [vault, keyname] = parse_key(k.master_key);
     auto [scheme, host, port] = parse_vault(vault);
@@ -451,6 +464,13 @@ future<azure_host::key_and_id_type> azure_host::impl::create_key(const attr_cach
 
 future<bytes> azure_host::impl::find_key(const id_cache_key& k) {
     static const vault_log_filter filter{vault_log_filter::op_type::unwrapkey};
+    if (_ctxt && this_shard_id() != 0) {
+        co_return co_await smp::submit_to(0, [this, k]() -> future<bytes> {
+            auto host = _ctxt->get_azure_host(_name);
+            auto bytes = co_await host->_impl->_id_cache.get(k);
+            co_return bytes;
+        });
+    }
     const auto id = to_string_view(k.id);
     azlog.debug("[{}] Finding key: {}", _log_prefix, id);
 
@@ -492,14 +512,14 @@ future<bytes> azure_host::impl::find_key(const id_cache_key& k) {
 // ==================== azure_host class implementation ====================
 
 azure_host::azure_host(const std::string& name, const host_options& options)
-    : _impl(std::make_unique<impl>(name, options))
+    : _impl(std::make_unique<impl>(nullptr, name, options))
 {}
 
-azure_host::azure_host([[maybe_unused]] encryption_context& ctxt, const std::string& name, const host_options& options)
-    : _impl(std::make_unique<impl>(name, options))
+azure_host::azure_host(encryption_context& ctxt, const std::string& name, const host_options& options)
+    : _impl(std::make_unique<impl>(&ctxt, name, options))
 {}
 
-azure_host::azure_host([[maybe_unused]] encryption_context& ctxt, const std::string& name, const std::unordered_map<sstring, sstring>& map)
+azure_host::azure_host(encryption_context& ctxt, const std::string& name, const std::unordered_map<sstring, sstring>& map)
     : azure_host(ctxt, name, [&map] {
         host_options opts;
         map_wrapper<std::unordered_map<sstring, sstring>> m(map);