feat(server): Use mimalloc in SSL calls (#2710)

* feat(server): Use mimalloc in SSL calls

Until now, OpenSSL used `malloc()` directly. This PR overrides it to use
mimalloc.

Fixes #2709

* Add generate-tls-files.sh

src/facade/dragonfly_listener.cc

@@ -156,11 +156,38 @@ bool ConfigureKeepAlive(int fd) {
   return true;
 }
 
+thread_local size_t ssl_allocated_bytes = 0;
+
+void* OverriddenSSLMalloc(size_t size, const char* file, int line) {
+  void* res = mi_malloc(size);
+  ssl_allocated_bytes += mi_malloc_usable_size(res);
+  return res;
+}
+
+void* OverriddenSSLRealloc(void* addr, size_t size, const char* file, int line) {
+  size_t prev_size = mi_malloc_usable_size(addr);
+  void* res = mi_realloc(addr, size);
+  ssl_allocated_bytes += mi_malloc_usable_size(res);
+  ssl_allocated_bytes -= prev_size;
+  return res;
+}
+
+void OverriddenSSLFree(void* addr, const char* file, int line) {
+  ssl_allocated_bytes -= mi_malloc_usable_size(addr);
+  mi_free(addr);
+}
+
+std::once_flag g_set_mem_functions_flag;
+
 }  // namespace
 
 Listener::Listener(Protocol protocol, ServiceInterface* si, Role role)
     : service_(si), protocol_(protocol) {
 #ifdef DFLY_USE_SSL
+  call_once(g_set_mem_functions_flag, []() {
+    CRYPTO_set_mem_functions(&OverriddenSSLMalloc, &OverriddenSSLRealloc, &OverriddenSSLFree);
+  });
+
   // Always initialise OpenSSL so we can enable TLS at runtime.
   OPENSSL_init_ssl(OPENSSL_INIT_SSL_DEFAULT, nullptr);
   if (!ReconfigureTLS()) {
@@ -243,6 +270,10 @@ bool Listener::ReconfigureTLS() {
   return true;
 }
 
+size_t Listener::TLSUsedMemoryThreadLocal() {
+  return ssl_allocated_bytes;
+}
+
 void Listener::PreAcceptLoop(util::ProactorBase* pb) {
   per_thread_.resize(pool()->size());
 }

src/facade/dragonfly_listener.h

@@ -43,6 +43,9 @@ class Listener : public util::ListenerInterface {
   // ReconfigureTLS MUST be called from the same proactor as the listener.
   bool ReconfigureTLS();
 
+  // Returns thread-local dynamic memory usage by TLS.
+  static size_t TLSUsedMemoryThreadLocal();
+
   bool IsPrivilegedInterface() const;
   bool IsMainInterface() const;
 

src/server/memory_cmd.cc

@@ -13,6 +13,7 @@
 #include "core/allocation_tracker.h"
 #include "facade/cmd_arg_parser.h"
 #include "facade/dragonfly_connection.h"
+#include "facade/dragonfly_listener.h"
 #include "facade/error.h"
 #include "server/engine_shard_set.h"
 #include "server/main_service.h"
@@ -253,11 +254,15 @@ void MemoryCmd::Stats() {
                        &stats);
 
   atomic<size_t> serialization_memory = 0;
-  shard_set->pool()->AwaitFiberOnAll(
-      [&](auto*) { serialization_memory.fetch_add(SliceSnapshot::GetThreadLocalMemoryUsage()); });
+  atomic<size_t> tls_memory = 0;
+  shard_set->pool()->AwaitFiberOnAll([&](auto*) {
+    serialization_memory.fetch_add(SliceSnapshot::GetThreadLocalMemoryUsage());
+    tls_memory.fetch_add(Listener::TLSUsedMemoryThreadLocal());
+  });
 
   // Serialization stats, including both replication-related serialization and saving to RDB files.
   stats.push_back({"serialization", serialization_memory.load()});
+  stats.push_back({"tls", tls_memory.load()});
 
   auto* rb = static_cast<RedisReplyBuilder*>(cntx_->reply_builder());
   rb->StartCollection(stats.size(), RedisReplyBuilder::MAP);

tools/generate-tls-files.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+# This script generates locally-signed TLS files for development usage.
+# It's probably a good idea to run in an empty, temporary directory.
+#
+# Example usage:
+#
+# mkdir /tmp/dfly-tls
+# cd /tmp/dfly-tls
+# ~/dragonfly/tools/generate-tls-files.sh
+# ~/dragonfly/build-dbg/dragonfly \
+#      --dbfilename= \
+#      --logtostdout \
+#      --tls=true \
+#      --tls_key_file=/tmp/dfly-tls/df-key.pem \
+#      --tls_cert_file=/tmp/dfly-tls/df-cert.pem \
+#      --requirepass=XXX
+# redis-cli --tls --cacert /tmp/dfly-tls/ca-cert.pem -a XXX
+
+CA_KEY_PATH=ca-key.pem
+CA_CERTIFICATE_PATH=ca-cert.pem
+CERTIFICATE_REQUEST_PATH=df-req.pem
+PRIVATE_KEY_PATH=df-key.pem
+CERTIFICATE_PATH=df-cert.pem
+
+echo "Generating files in local directory (rm *.pem to cleanup)"
+
+openssl req -x509 -newkey rsa:4096 -days 1 -nodes \
+  -keyout ${CA_KEY_PATH} \
+  -out ${CA_CERTIFICATE_PATH} \
+  -subj "/C=GR/ST=SKG/L=Thessaloniki/O=KK/OU=AcmeStudios/CN=Gr/emailAddress=acme@gmail.com"
+
+openssl req -newkey rsa:4096 -nodes \
+  -keyout ${PRIVATE_KEY_PATH} \
+  -out ${CERTIFICATE_REQUEST_PATH} \
+  -subj "/C=GR/ST=SKG/L=Thessaloniki/O=KK/OU=Comp/CN=Gr/emailAddress=does_not_exist@gmail.com"
+
+openssl x509 -req \
+  -in ${CERTIFICATE_REQUEST_PATH} \
+  -days 1 \
+  -CA ${CA_CERTIFICATE_PATH} \
+  -CAkey ${CA_KEY_PATH} \
+  -CAcreateserial -out ${CERTIFICATE_PATH}
+
+echo "You can now run:"
+echo "dragonfly --tls=true --tls_key_file=${PRIVATE_KEY_PATH} --tls_cert_file=${CERTIFICATE_PATH} --requirepass=XXX"
+echo "redis-cli --tls --cacert ${CA_CERTIFICATE_PATH} -a XXX"