Subject: Optimize scram_SaltedPassword performance

From: Zixuan Fu <[email-masked]>
Date: Mon, 3 Feb 2025 15:11:48 +0800

Hi Hackers,  

While profiling a program with `perf`, I noticed that `scram_SaltedPassword`
consumed more CPU time than expected. After some investigation, I found
that the function performs many HMAC iterations (4096 rounds for
SCRAM-SHA-256), and each iteration reinitializes the HMAC context, causing
excessive overhead.

OpenSSL has an optimization for this case: when the key remains the
same, the HMAC context can be reused with a lightweight state reset by
passing NULL as the key. To take advantage of this, I introduced
`pg_hmac_reuse()`, which replaces the key with NULL when OpenSSL is used.

With this change, the performance improves by approximately **4x** (reducing
execution time from 4ms to 1ms). The built-in PostgreSQL HMAC implementation
does not support context reuse, and modifying it would require substantial
changes. Therefore, `pg_hmac_reuse()` simply calls `pg_hmac_init()` in that
case, maintaining the existing logic.

Regards,

Zixuan
>From e13d9e9c674d1cf4f671ca81b950ae010fd958a4 Mon Sep 17 00:00:00 2001
From: Zi-Xuan Fu <[email-masked]>
Date: Mon, 3 Feb 2025 13:33:43 +0800
Subject: [PATCH v1] Optimize "scram_SaltedPassword" with OpenSSL HMAC context reuse  

`scram_SaltedPassword()` requires a large number of HMAC iterations, with 
SCRAM-SHA-256 defaulting to 4096 rounds. Previously, each iteration 
reinitialized the HMAC context, incurring significant overhead. However, 
OpenSSL optimizes this case by allowing context reuse when the key remains 
unchanged—requiring only a lightweight state reset by passing NULL as the key.  

To leverage this, I introduced `pg_hmac_reuse()`, which sets the key to 
NULL when OpenSSL is used. This optimization improves performance by 
approximately 4x (reducing execution time from 4ms to 1ms).  

The native PostgreSQL HMAC implementation does not support context reuse, and 
modifying it would require substantial changes. Therefore, `pg_hmac_reuse()` 
simply calls `pg_hmac_init()` in that case, maintaining the existing logic.  

---
 src/common/hmac.c         | 11 +++++++++++
 src/common/hmac_openssl.c | 12 ++++++++++++
 src/common/scram-common.c |  2 +-
 src/include/common/hmac.h |  1 +
 4 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/common/hmac.c b/src/common/hmac.c
index 9b85910672..6de95c44ee 100644
--- a/src/common/hmac.c
+++ b/src/common/hmac.c
@@ -214,6 +214,17 @@ pg_hmac_init(pg_hmac_ctx *ctx, const uint8 *key, size_t len)
 	return 0;
 }
 
+/*
+ * pg_hmac_reuse
+ *
+ * Reuse a HMAC context with the same key.  Returns 0 on success, -1 on failure.
+ */
+int
+pg_hmac_reuse(pg_hmac_ctx *ctx, const uint8 *key, size_t len)
+{
+	return pg_hmac_init(ctx, key, len);
+}
+
 /*
  * pg_hmac_update
  *
diff --git a/src/common/hmac_openssl.c b/src/common/hmac_openssl.c
index 31cd188904..9e7a92884d 100644
--- a/src/common/hmac_openssl.c
+++ b/src/common/hmac_openssl.c
@@ -208,6 +208,18 @@ pg_hmac_init(pg_hmac_ctx *ctx, const uint8 *key, size_t len)
 	return 0;
 }
 
+/*
+ * pg_hmac_reuse
+ *
+ * Reuse a HMAC context with the same key.  Returns 0 on success, -1 on failure.
+ */
+int
+pg_hmac_reuse(pg_hmac_ctx *ctx, const uint8 *key, size_t len)
+{
+	/* OpenSSL skips unnecessary reinitialization if the key is NULL */
+	return pg_hmac_init(ctx, NULL, len);
+}
+
 /*
  * pg_hmac_update
  *
diff --git a/src/common/scram-common.c b/src/common/scram-common.c
index 400900c51c..5717daff28 100644
--- a/src/common/scram-common.c
+++ b/src/common/scram-common.c
@@ -84,7 +84,7 @@ scram_SaltedPassword(const char *password,
 		CHECK_FOR_INTERRUPTS();
 #endif
 
-		if (pg_hmac_init(hmac_ctx, (uint8 *) password, password_len) < 0 ||
+		if (pg_hmac_reuse(hmac_ctx, (uint8 *) password, password_len) < 0 ||
 			pg_hmac_update(hmac_ctx, (uint8 *) Ui_prev, key_length) < 0 ||
 			pg_hmac_final(hmac_ctx, Ui, key_length) < 0)
 		{
diff --git a/src/include/common/hmac.h b/src/include/common/hmac.h
index c4d069e49a..e9a0366d6e 100644
--- a/src/include/common/hmac.h
+++ b/src/include/common/hmac.h
@@ -22,6 +22,7 @@ typedef struct pg_hmac_ctx pg_hmac_ctx;
 
 extern pg_hmac_ctx *pg_hmac_create(pg_cryptohash_type type);
 extern int	pg_hmac_init(pg_hmac_ctx *ctx, const uint8 *key, size_t len);
+extern int	pg_hmac_reuse(pg_hmac_ctx *ctx, const uint8 *key, size_t len);
 extern int	pg_hmac_update(pg_hmac_ctx *ctx, const uint8 *data, size_t len);
 extern int	pg_hmac_final(pg_hmac_ctx *ctx, uint8 *dest, size_t len);
 extern void pg_hmac_free(pg_hmac_ctx *ctx);
-- 
2.34.1