Repository: randombit/botan
Pull request #3258 (2023-02-07 22:10:23)

Change Credentials_Manager::private_key_for to return a shared_ptr
The previous design introduced serious ownership management problems because there was no way for the Credentials_Manager to know when (or even if) the TLS layer had dropped the pointer, so it had to be kept alive forever.


Patch:
diff --git a/doc/api_ref/credentials_manager.rst b/doc/api_ref/credentials_manager.rst
index 5f0db2efe..be1fbf6d7 100644
--- a/doc/api_ref/credentials_manager.rst
+++ b/doc/api_ref/credentials_manager.rst
@@ -81,13 +81,13 @@ stores credentials. The main user is the :doc:`tls` implementation.
          the connection with an alert. This can be done by throwing an exception
          from the implementation of this function.
 
-   .. cpp:function:: Private_Key* private_key_for(const X509_Certificate& cert, \
+   .. cpp:function:: std::shared_ptr<Private_Key> private_key_for(const X509_Certificate& cert, \
                                                   const std::string& type, \
                                                   const std::string& context)
 
-      Return the private key for this certificate. The *cert* will be
-      the leaf cert of a chain returned previously by ``cert_chain``
-      or ``cert_chain_single_type``.
+      Return a shared pointer to the private key for this certificate. The
+      *cert* will be the leaf cert of a chain returned previously by
+      ``cert_chain`` or ``cert_chain_single_type``.
 
 In versions before 1.11.34, there was an additional function on `Credentials_Manager`
 
diff --git a/doc/migration_guide.rst b/doc/migration_guide.rst
index b835c7aea..68dcf5e11 100644
--- a/doc/migration_guide.rst
+++ b/doc/migration_guide.rst
@@ -182,3 +182,12 @@ DLIES Constructors
 
 Previously the constructors to the DLIES classes took raw pointers,
 and retained ownership of them. They now consume std::unique_ptrs
+
+Credentials_Manager::private_key_for
+-------------------------------------
+
+Previously this function returned a raw pointer, which the Credentials_Manager
+implementation had to keep alive "forever", since there was no way for it to
+know when or if the TLS layer had completed using the returned key.
+
+Now this function returns std::shared_ptr<Private_Key>
diff --git a/src/bogo_shim/bogo_shim.cpp b/src/bogo_shim/bogo_shim.cpp
index c126a4a88..77146d666 100644
--- a/src/bogo_shim/bogo_shim.cpp
+++ b/src/bogo_shim/bogo_shim.cpp
@@ -1239,7 +1239,7 @@ class Shim_Credentials final : public Botan::Credentials_Manager
          if(m_args.option_used("key-file") && m_args.option_used("cert-file"))
             {
             Botan::DataSource_Stream key_stream(m_args.get_string_opt("key-file"));
-            m_key = Botan::PKCS8::load_key(key_stream);
+            m_key.reset(Botan::PKCS8::load_key(key_stream).release());
 
             Botan::DataSource_Stream cert_stream(m_args.get_string_opt("cert-file"));
 
@@ -1308,19 +1308,20 @@ class Shim_Credentials final : public Botan::Credentials_Manager
          return {};
          }
 
-      Botan::Private_Key* private_key_for(const Botan::X509_Certificate& /*cert*/,
-                                          const std::string& /*type*/,
-                                          const std::string& /*context*/) override
+      std::shared_ptr<Botan::Private_Key>
+      private_key_for(const Botan::X509_Certificate& /*cert*/,
+                      const std::string& /*type*/,
+                      const std::string& /*context*/) override
          {
          // assumes cert == m_cert
-         return m_key.get();
+         return m_key;
          }
 
    private:
       const Shim_Arguments& m_args;
       Botan::SymmetricKey m_psk;
       std::string m_psk_identity;
-      std::unique_ptr<Botan::Private_Key> m_key;
+      std::shared_ptr<Botan::Private_Key> m_key;
       std::vector<Botan::X509_Certificate> m_cert_chain;
    };
 
diff --git a/src/cli/tls_helpers.h b/src/cli/tls_helpers.h
index f29ed1ae2..55cbbee08 100644
--- a/src/cli/tls_helpers.h
+++ b/src/cli/tls_helpers.h
@@ -156,15 +156,16 @@ class Basic_Credentials_Manager : public Botan::Credentials_Manager
          return {};
          }
 
-      Botan::Private_Key* private_key_for(const Botan::X509_Certificate& cert,
-                                          const std::string& /*type*/,
-                                          const std::string& /*context*/) override
+      std::shared_ptr<Botan::Private_Key>
+      private_key_for(const Botan::X509_Certificate& cert,
+                      const std::string& /*type*/,
+                      const std::string& /*context*/) override
          {
          for(auto const& i : m_creds)
             {
             if(cert == i.certs[0])
                {
-               return i.key.get();
+               return i.key;
                }
             }
 
diff --git a/src/examples/tls_client.cpp b/src/examples/tls_client.cpp
index 82e0534d8..5c0d878fb 100644
--- a/src/examples/tls_client.cpp
+++ b/src/examples/tls_client.cpp
@@ -62,8 +62,9 @@ public:
     return {};
   }
 
-  Botan::Private_Key *private_key_for(const Botan::X509_Certificate &cert, const std::string &type,
-                                      const std::string &context) override {
+  std::shared_ptr<Botan::Private_Key>
+  private_key_for(const Botan::X509_Certificate &cert, const std::string &type,
+                  const std::string &context) override {
     // when returning a chain in cert_chain(), return the private key
     // associated with the leaf certificate here
     return nullptr;
diff --git a/src/examples/tls_custom_curves_client.cpp b/src/examples/tls_custom_curves_client.cpp
index 0e3bd4eaa..13ec8f671 100644
--- a/src/examples/tls_custom_curves_client.cpp
+++ b/src/examples/tls_custom_curves_client.cpp
@@ -75,8 +75,9 @@ public:
     return {};
   }
 
-  Botan::Private_Key *private_key_for(const Botan::X509_Certificate &cert, const std::string &type,
-                                      const std::string &context) override {
+  std::shared_ptr<Botan::Private_Key>
+  private_key_for(const Botan::X509_Certificate &cert, const std::string &type,
+                  const std::string &context) override {
     // when returning a chain in cert_chain(), return the private key
     // associated with the leaf certificate here
     return nullptr;
diff --git a/src/examples/tls_custom_curves_server.cpp b/src/examples/tls_custom_curves_server.cpp
index e233d9d63..fa9cfde83 100644
--- a/src/examples/tls_custom_curves_server.cpp
+++ b/src/examples/tls_custom_curves_server.cpp
@@ -67,7 +67,7 @@ class Server_Credentials : public Botan::Credentials_Manager {
 public:
   Server_Credentials() {
     Botan::DataSource_Stream in("botan.randombit.net.key");
-    m_key = Botan::PKCS8::load_key(in);
+    m_key.reset(Botan::PKCS8::load_key(in).release());
   }
 
   std::vector<Botan::Certificate_Store *>
@@ -87,15 +87,16 @@ public:
     return {Botan::X509_Certificate("botan.randombit.net.crt")};
   }
 
-  Botan::Private_Key *private_key_for(const Botan::X509_Certificate &cert, const std::string &type,
-                                      const std::string &context) override {
+  std::shared_ptr<Botan::Private_Key>
+  private_key_for(const Botan::X509_Certificate &cert, const std::string &type,
+                  const std::string &context) override {
     // return the private key associated with the leaf certificate,
     // in this case the one associated with "botan.randombit.net.crt"
-    return m_key.get();
+    return m_key;
   }
 
 private:
-  std::unique_ptr<Botan::Private_Key> m_key;
+  std::shared_ptr<Botan::Private_Key> m_key;
 };
 
 class Server_Policy : public Botan::TLS::Strict_Policy {
diff --git a/src/examples/tls_proxy.cpp b/src/examples/tls_proxy.cpp
index 62666ab96..5046ac6a3 100644
--- a/src/examples/tls_proxy.cpp
+++ b/src/examples/tls_proxy.cpp
@@ -52,7 +52,7 @@ class Server_Credentials : public Botan::Credentials_Manager {
 public:
   Server_Credentials() {
     Botan::DataSource_Stream in("botan.randombit.net.key");
-    m_key = Botan::PKCS8::load_key(in);
+    m_key.reset(Botan::PKCS8::load_key(in).release());
   }
 
   std::vector<Botan::Certificate_Store *>
@@ -72,15 +72,16 @@ public:
     return {Botan::X509_Certificate("botan.randombit.net.crt")};
   }
 
-  Botan::Private_Key *private_key_for(const Botan::X509_Certificate &cert, const std::string &type,
-                                      const std::string &context) override {
+  std::shared_ptr<Botan::Private_Key>
+  private_key_for(const Botan::X509_Certificate &cert, const std::string &type,
+                  const std::string &context) override {
     // return the private key associated with the leaf certificate,
     // in this case the one associated with "botan.randombit.net.crt"
-    return m_key.get();
+    return m_key;
   }
 
 private:
-  std::unique_ptr<Botan::Private_Key> m_key;
+  std::shared_ptr<Botan::Private_Key> m_key;
 };
 
 int main() {
diff --git a/src/fuzzer/tls_server.cpp b/src/fuzzer/tls_server.cpp
index b9e0cee74..0fad047d0 100644
--- a/src/fuzzer/tls_server.cpp
+++ b/src/fuzzer/tls_server.cpp
@@ -7,6 +7,7 @@
 #include "fuzzers.h"
 #include <botan/tls_server.h>
 #include <botan/data_src.h>
+#include <botan/pkcs8.h>
 
 const char* fixed_rsa_key =
    "-----BEGIN PRIVATE KEY-----\n"
@@ -66,10 +67,10 @@ class Fuzzer_TLS_Server_Creds : public Botan::Credentials_Manager
       Fuzzer_TLS_Server_Creds()
          {
          Botan::DataSource_Memory cert_in(fixed_rsa_cert);
-         Botan::DataSource_Memory key_in(fixed_rsa_key);
-
          m_rsa_cert.reset(new Botan::X509_Certificate(cert_in));
-         //m_rsa_key.reset(Botan::PKCS8::load_key(key_in, fuzzer_rng());
+
+         Botan::DataSource_Memory key_in(fixed_rsa_key);
+         m_rsa_key.reset(Botan::PKCS8::load_key(key_in).release());
          }
 
       std::vector<Botan::X509_Certificate> cert_chain(
@@ -92,11 +93,14 @@ class Fuzzer_TLS_Server_Creds : public Botan::Credentials_Manager
          return v;
          }
 
-      Botan::Private_Key* private_key_for(const Botan::X509_Certificate& /*cert*/,
-                                          const std::string& /*type*/,
-                                          const std::string& /*context*/) override
+      std::shared_ptr<Botan::Private_Key>
+      private_key_for(const Botan::X509_Certificate& /*cert*/,
+                      const std::string& type,
+                      const std::string& /*context*/) override
          {
-         return m_rsa_key.get();
+         if(type == "RSA")
+            return m_rsa_key;
+         return nullptr;
          }
 
       std::string psk_identity_hint(const std::string&, const std::string&) override { return "psk_hint"; }
@@ -107,7 +111,7 @@ class Fuzzer_TLS_Server_Creds : public Botan::Credentials_Manager
          }
    private:
       std::unique_ptr<Botan::X509_Certificate> m_rsa_cert;
-      std::unique_ptr<Botan::Private_Key> m_rsa_key;
+      std::shared_ptr<Botan::Private_Key> m_rsa_key;
    };
 
 class Fuzzer_TLS_Policy : public Botan::TLS::Policy
diff --git a/src/lib/tls/credentials_manager.cpp b/src/lib/tls/credentials_manager.cpp
index cadd237db..fe3186c82 100644
--- a/src/lib/tls/credentials_manager.cpp
+++ b/src/lib/tls/credentials_manager.cpp
@@ -58,11 +58,12 @@ std::vector<X509_Certificate> Credentials_Manager::cert_chain_single_type(
    return find_cert_chain({cert_key_type}, cert_signature_schemes, std::vector<X509_DN>(), type, context);
    }
 
-Private_Key* Credentials_Manager::private_key_for(const X509_Certificate& /*unused*/,
-                                                  const std::string& /*unused*/,
-                                                  const std::string& /*unused*/)
+std::shared_ptr<Private_Key>
+Credentials_Manager::private_key_for(const X509_Certificate& /*unused*/,
+                                     const std::string& /*unused*/,
+                                     const std::string& /*unused*/)
    {
-   return nullptr;
+   return std::shared_ptr<Private_Key>();
    }
 
 std::vector<Certificate_Store*>
diff --git a/src/lib/tls/credentials_manager.h b/src/lib/tls/credentials_manager.h
index f03a67a57..a916b10b8 100644
--- a/src/lib/tls/credentials_manager.h
+++ b/src/lib/tls/credentials_manager.h
@@ -127,12 +127,13 @@ class BOTAN_PUBLIC_API(2,0) Credentials_Manager
       /**
       * @return private key associated with this certificate if we should
       *         use it with this context. cert was returned by cert_chain
-      * @note this object should retain ownership of the returned key;
-      *       it should not be deleted by the caller.
+      * This function should either return null or throw an exception if
+      * the key is unavailable.
       */
-      virtual Private_Key* private_key_for(const X509_Certificate& cert,
-                                           const std::string& type,
-                                           const std::string& context);
+      virtual std::shared_ptr<Private_Key>
+         private_key_for(const X509_Certificate& cert,
+                         const std::string& type,
+                         const std::string& context);
 
       /**
       * @param type specifies the type of operation occurring
diff --git a/src/lib/tls/msg_cert_verify.cpp b/src/lib/tls/msg_cert_verify.cpp
index 7559f4ffb..908ceadb7 100644
--- a/src/lib/tls/msg_cert_verify.cpp
+++ b/src/lib/tls/msg_cert_verify.cpp
@@ -164,10 +164,11 @@ Certificate_Verify_13::Certificate_Verify_13(
    {
    BOTAN_ASSERT_NOMSG(!certificate_msg.empty());
 
-   const auto* private_key = creds_mgr.private_key_for(
-                                       certificate_msg.leaf(),
-                                       m_side == Connection_Side::Client ? "tls-client" : "tls-server",
-                                       hostname);
+   const auto private_key = creds_mgr.private_key_for(
+      certificate_msg.leaf(),
+      m_side == Connection_Side::Client ? "tls-client" : "tls-server",
+      hostname);
+
    if(!private_key)
       { throw TLS_Exception(Alert::InternalError, "Application did not provide a private key for its certificate"); }
 
diff --git a/src/lib/tls/tls12/tls_client_impl_12.cpp b/src/lib/tls/tls12/tls_client_impl_12.cpp
index ee219ec0d..25750158f 100644
--- a/src/lib/tls/tls12/tls_client_impl_12.cpp
+++ b/src/lib/tls/tls12/tls_client_impl_12.cpp
@@ -699,17 +699,20 @@ void Client_Impl_12::process_handshake_msg(const Handshake_State* active_state,
       if(state.received_handshake_msg(Handshake_Type::CertificateRequest) &&
          !state.client_certs()->empty())
          {
-         Private_Key* private_key =
+         auto private_key =
             m_creds.private_key_for(state.client_certs()->cert_chain()[0],
                                     "tls-client",
                                     m_info.hostname());
 
+         if(!private_key)
+            throw TLS_Exception(Alert::InternalError, "Failed to get private key for signing");
+
          state.client_verify(
             new Certificate_Verify_12(state.handshake_io(),
                                       state,
                                       policy(),
                                       rng(),
-                                      private_key)
+                                      private_key.get())
             );
          }
 
diff --git a/src/lib/tls/tls12/tls_server_impl_12.cpp b/src/lib/tls/tls12/tls_server_impl_12.cpp
index 051a856ac..7bf1bb8b0 100644
--- a/src/lib/tls/tls12/tls_server_impl_12.cpp
+++ b/src/lib/tls/tls12/tls_server_impl_12.cpp
@@ -24,8 +24,8 @@ class Server_Handshake_State final : public Handshake_State
       Server_Handshake_State(std::unique_ptr<Handshake_IO> io, Callbacks& cb)
          : Handshake_State(std::move(io), cb) {}
 
-      Private_Key* server_rsa_kex_key() { return m_server_rsa_kex_key; }
-      void set_server_rsa_kex_key(Private_Key* key)
+      Private_Key* server_rsa_kex_key() { return m_server_rsa_kex_key.get(); }
+      void set_server_rsa_kex_key(std::shared_ptr<Private_Key> key)
          { m_server_rsa_kex_key = key; }
 
       bool allow_session_resumption() const
@@ -44,8 +44,8 @@ class Server_Handshake_State final : public Handshake_State
       bool is_a_resumption() const { return m_is_a_resumption; }
 
    private:
-      // Used by the server only, in case of RSA key exchange. Not owned
-      Private_Key* m_server_rsa_kex_key = nullptr;
+      // Used by the server only, in case of RSA key exchange.
+      std::shared_ptr<Private_Key> m_server_rsa_kex_key;
 
       /*
       * Used by the server to know if resumption should be allowed on
@@ -864,7 +864,7 @@ void Server_Impl_12::session_create(Server_Handshake_State& pending_state,
 
    const Ciphersuite& pending_suite = pending_state.ciphersuite();
 
-   Private_Key* private_key = nullptr;
+   std::shared_ptr<Private_Key> private_key;
 
    if(pending_suite.signature_used() || pending_suite.kex_method() == Kex_Algo::STATIC_RSA)
       {
@@ -911,7 +911,7 @@ void Server_Impl_12::session_create(Server_Handshake_State& pending_state,
       {
       pending_state.server_kex(new Server_Key_Exchange(pending_state.handshake_io(),
                                                        pending_state, policy(),
-                                                       m_creds, rng(), private_key));
+                                                       m_creds, rng(), private_key.get()));
       }
 
    auto trusted_CAs = m_creds.trusted_certificate_authorities("tls-server", sni_hostname);


Test patch:
diff --git a/src/tests/test_tls_rfc8448.cpp b/src/tests/test_tls_rfc8448.cpp
index 5bdd0f674..c34bc40aa 100644
--- a/src/tests/test_tls_rfc8448.cpp
+++ b/src/tests/test_tls_rfc8448.cpp
@@ -63,23 +63,6 @@ void add_entropy(Botan_Tests::Fixed_Output_RNG& rng, const std::vector<uint8_t>&
    rng.add_entropy(bin.data(), bin.size());
    }
 
-Botan::Private_Key* server_private_key()
-   {
-   static Botan::DataSource_Memory in(Test::read_data_file("tls_13_rfc8448/server_key.pem"));
-   static auto key = Botan::PKCS8::load_key(in);
-   return key.get();
-   }
-
-Botan::Private_Key* bogus_alternative_server_private_key()
-   {
-   // RFC 8448 does not actually provide that key. Hence we generate one on the
-   // fly as a stand-in. Instead of actually using it, the signatures generated
-   // by this private key must be hard-coded in `Callbacks::sign_message()`; see
-   // `MockSignature_Fn` for more details.
-   static auto key = Botan::ECDSA_PrivateKey(Test::rng(), Botan::EC_Group("secp256r1"));
-   return &key;
-   }
-
 Botan::X509_Certificate server_certificate()
    {
    // self-signed certificate with an RSA1024 public key valid until:
@@ -99,16 +82,6 @@ Botan::X509_Certificate alternative_server_certificate()
    return Botan::X509_Certificate(in);
    }
 
-Botan::Private_Key* client_private_key()
-   {
-   // RFC 8448 does not actually provide that key. Hence we generate one on the
-   // fly as a stand-in. Instead of actually using it, the signatures generated
-   // by this private key must be hard-coded in `Callbacks::sign_message()`; see
-   // `MockSignature_Fn` for more details.
-   static auto key = create_private_key("RSA", Test::rng(), "1024");
-   return key.get();
-   }
-
 Botan::X509_Certificate client_certificate()
    {
    // self-signed certificate with an RSA1024 public key valid until:
@@ -456,7 +429,20 @@ class Test_Credentials : public Botan::Credentials_Manager
    {
    public:
       explicit Test_Credentials(bool use_alternative_server_certificate)
-         : m_alternative_server_certificate(use_alternative_server_certificate) {}
+         : m_alternative_server_certificate(use_alternative_server_certificate)
+         {
+         Botan::DataSource_Memory in(Test::read_data_file("tls_13_rfc8448/server_key.pem"));
+         m_server_private_key.reset(Botan::PKCS8::load_key(in).release());
+
+         // RFC 8448 does not actually provide these keys. Hence we generate one on the
+         // fly as a stand-in. Instead of actually using it, the signatures generated
+         // by this private key must be hard-coded in `Callbacks::sign_message()`; see
+         // `MockSignature_Fn` for more details.
+         m_bogus_alternative_server_private_key.reset(
+            create_private_key("ECDSA", Test::rng(), "secp256r1").release());
+
+         m_client_private_key.reset(create_private_key("RSA", Test::rng(), "1024").release());
+         }
 
       std::vector<Botan::X509_Certificate> cert_chain(
          const std::vector<std::string>& cert_key_types,
@@ -475,21 +461,27 @@ class Test_Credentials : public Botan::Credentials_Manager
             };
          }
 
-      Botan::Private_Key* private_key_for(const Botan::X509_Certificate& cert,
-                                          const std::string& type,
-                                          const std::string& context) override
+      std::shared_ptr<Botan::Private_Key>
+      private_key_for(const Botan::X509_Certificate& cert,
+                      const std::string& type,
+                      const std::string& context) override
          {
          BOTAN_UNUSED(cert, context);
 
-         return (type == "tls-client")
-            ? client_private_key()
-            : ((m_alternative_server_certificate)
-               ? bogus_alternative_server_private_key()
-               : server_private_key());
+         if(type == "tls-client")
+            return m_client_private_key;
+
+         if(m_alternative_server_certificate)
+            return m_bogus_alternative_server_private_key;
+
+         return m_server_private_key;
          }
 
    private:
       bool m_alternative_server_certificate;
+      std::shared_ptr<Private_Key> m_client_private_key;
+      std::shared_ptr<Private_Key> m_bogus_alternative_server_private_key;
+      std::shared_ptr<Private_Key> m_server_private_key;
    };
 
 class RFC8448_Text_Policy : public Botan::TLS::Text_Policy
@@ -1322,7 +1314,12 @@ class Test_TLS_RFC8448_Client : public Test_TLS_RFC8448
             Botan_Tests::CHECK("Client Hello", [&](Test::Result& result)
                {
 
-               ctx = std::make_unique<Client_Context>(std::move(rng), read_tls_policy("rfc8448_1rtt"), vars.get_req_u64("CurrentTimestamp"), add_extensions_and_sort, std::nullopt, make_mock_signatures(vars));
+               ctx = std::make_unique<Client_Context>(std::move(rng),
+                                                      read_tls_policy("rfc8448_1rtt"),
+                                                      vars.get_req_u64("CurrentTimestamp"),
+                                                      add_extensions_and_sort,
+                                                      std::nullopt,
+                                                      make_mock_signatures(vars));
 
                ctx->check_callback_invocations(result, "initial callbacks", {
                   "tls_emit_data",
@@ -1588,12 +1585,14 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Verify generated Hello Retry Request message", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                result.test_eq("Server's Hello Retry Request record", ctx->pull_send_buffer(), vars.get_req_bin("Record_HelloRetryRequest"));
                result.confirm("TLS handshake not yet finished", !ctx->server.is_active());
                }),
 
             Botan_Tests::CHECK("Receive updated Client Hello message", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.received_data(vars.get_req_bin("Record_ClientHello_2"));
 
                ctx->check_callback_invocations(result, "updated client hello received", {
@@ -1615,6 +1614,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Verify generated messages in server's second flight", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                const auto& msgs = ctx->observed_handshake_messages();
 
                result.test_eq("Server Hello", msgs.at("server_hello")[0], strip_message_header(vars.get_opt_bin("Message_ServerHello")));
@@ -1630,6 +1630,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Receive Client Finished", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.received_data(vars.get_req_bin("Record_ClientFinished"));
 
                ctx->check_callback_invocations(result, "client finished received", {
@@ -1642,6 +1643,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Receive Client close_notify", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.received_data(vars.get_req_bin("Record_Client_CloseNotify"));
 
                ctx->check_callback_invocations(result, "client finished received", {
@@ -1655,6 +1657,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Expect Server close_notify", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.close();
 
                result.confirm("connection is now inactive", !ctx->server.is_active());
@@ -1703,6 +1706,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Verify server's generated handshake messages", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                const auto& msgs = ctx->observed_handshake_messages();
 
                result.test_eq("Server Hello", msgs.at("server_hello")[0], strip_message_header(vars.get_opt_bin("Message_ServerHello")));
@@ -1721,6 +1725,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Receive Client's second flight", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                // This encrypted message contains the following messages:
                // * client's Certificate message
                // * client's Certificate_Verify message
@@ -1745,6 +1750,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Receive Client close_notify", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.received_data(vars.get_req_bin("Record_Client_CloseNotify"));
 
                ctx->check_callback_invocations(result, "client finished received", {
@@ -1758,6 +1764,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Expect Server close_notify", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.close();
 
                result.confirm("connection is now inactive", !ctx->server.is_active());
@@ -1805,6 +1812,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Verify server's generated handshake messages", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                const auto& msgs = ctx->observed_handshake_messages();
 
                result.test_eq("Server Hello", msgs.at("server_hello")[0], strip_message_header(vars.get_opt_bin("Message_ServerHello")));
@@ -1822,6 +1830,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Receive Client Finished", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.received_data(vars.get_req_bin("Record_ClientFinished"));
 
                ctx->check_callback_invocations(result, "client finished received", {
@@ -1834,6 +1843,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Receive Client close_notify", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.received_data(vars.get_req_bin("Record_Client_CloseNotify"));
 
                ctx->check_callback_invocations(result, "client finished received", {
@@ -1847,6 +1857,7 @@ class Test_TLS_RFC8448_Server : public Test_TLS_RFC8448
 
             Botan_Tests::CHECK("Expect Server close_notify", [&](Test::Result& result)
                {
+               BOTAN_ASSERT_NONNULL(ctx);
                ctx->server.close();
 
                result.confirm("connection is now inactive", !ctx->server.is_active());
diff --git a/src/tests/unit_tls.cpp b/src/tests/unit_tls.cpp
index cb7ab2f38..875d895eb 100644
--- a/src/tests/unit_tls.cpp
+++ b/src/tests/unit_tls.cpp
@@ -125,17 +125,18 @@ class Credentials_Manager_Test final : public Botan::Credentials_Manager
          return chain;
          }
 
-      Botan::Private_Key* private_key_for(const Botan::X509_Certificate& crt,
-                                          const std::string& /*type*/,
-                                          const std::string& /*context*/) override
+      std::shared_ptr<Botan::Private_Key>
+      private_key_for(const Botan::X509_Certificate& crt,
+                      const std::string& /*type*/,
+                      const std::string& /*context*/) override
          {
          if(crt == m_rsa_cert)
             {
-            return m_rsa_key.get();
+            return m_rsa_key;
             }
          if(crt == m_ecdsa_cert)
             {
-            return m_ecdsa_key.get();
+            return m_ecdsa_key;
             }
          return nullptr;
          }
@@ -171,10 +172,10 @@ class Credentials_Manager_Test final : public Botan::Credentials_Manager
 
    private:
       Botan::X509_Certificate m_rsa_cert, m_rsa_ca;
-      std::unique_ptr<Botan::Private_Key> m_rsa_key;
+      std::shared_ptr<Botan::Private_Key> m_rsa_key;
 
       Botan::X509_Certificate m_ecdsa_cert, m_ecdsa_ca;
-      std::unique_ptr<Botan::Private_Key> m_ecdsa_key;
+      std::shared_ptr<Botan::Private_Key> m_ecdsa_key;
 
       std::vector<std::unique_ptr<Botan::Certificate_Store>> m_stores;
       bool m_provides_client_certs;