Refactor HMAC implementations

Similarly to the cryptohash implementations, this refactors the existing
HMAC code into a single set of APIs that can be plugged with any crypto
libraries PostgreSQL is built with (only OpenSSL currently).  If there
is no such libraries, a fallback implementation is available.  Those new
APIs are designed similarly to the existing cryptohash layer, so there
is no real new design here, with the same logic around buffer bound
checks and memory handling.

HMAC has a dependency on cryptohashes, so all the cryptohash types
supported by cryptohash{_openssl}.c can be used with HMAC.  This
refactoring is an advantage mainly for SCRAM, that included its own
implementation of HMAC with SHA256 without relying on the existing
crypto libraries even if PostgreSQL was built with their support.

This code has been tested on Windows and Linux, with and without
OpenSSL, across all the versions supported on HEAD from 1.1.1 down to
1.0.1.  I have also checked that the implementations are working fine
using some sample results, a custom extension of my own, and doing
cross-checks across different major versions with SCRAM with the client
and the backend.

Author: Michael Paquier
Reviewed-by: Bruce Momjian
Discussion: https://postgr.es/m/X9m0nkEJEzIPXjeZ@paquier.xyz

4 files changed, 555 insertions, 126 deletions

diff --git a/src/common/Makefile b/src/common/Makefile
index 5422579a6a2..38a85993370 100644
--- a/src/common/Makefile
+++ b/src/common/Makefile
@@ -83,10 +83,12 @@ OBJS_COMMON = \
 ifeq ($(with_ssl),openssl)
 OBJS_COMMON += \
 	protocol_openssl.o \
-	cryptohash_openssl.o
+	cryptohash_openssl.o \
+	hmac_openssl.o
 else
 OBJS_COMMON += \
 	cryptohash.o \
+	hmac.o \
 	md5.o \
 	sha1.o \
 	sha2.o
diff --git a/src/common/hmac.c b/src/common/hmac.c
new file mode 100644
index 00000000000..f1b8555143a
--- /dev/null
+++ b/src/common/hmac.c
@@ -0,0 +1,263 @@
+/*-------------------------------------------------------------------------
+ *
+ * hmac.c
+ *	  Implements Keyed-Hashing for Message Authentication (HMAC)
+ *
+ * Fallback implementation of HMAC, as specified in RFC 2104.
+ *
+ * Portions Copyright (c) 1996-2020, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  src/common/hmac.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#ifndef FRONTEND
+#include "postgres.h"
+#else
+#include "postgres_fe.h"
+#endif
+
+#include "common/cryptohash.h"
+#include "common/hmac.h"
+#include "common/md5.h"
+#include "common/sha1.h"
+#include "common/sha2.h"
+
+/*
+ * In backend, use palloc/pfree to ease the error handling.  In frontend,
+ * use malloc to be able to return a failure status back to the caller.
+ */
+#ifndef FRONTEND
+#define ALLOC(size) palloc(size)
+#define FREE(ptr) pfree(ptr)
+#else
+#define ALLOC(size) malloc(size)
+#define FREE(ptr) free(ptr)
+#endif
+
+/*
+ * Internal structure for pg_hmac_ctx->data with this implementation.
+ */
+struct pg_hmac_ctx
+{
+	pg_cryptohash_ctx *hash;
+	pg_cryptohash_type type;
+	int			block_size;
+	int			digest_size;
+
+	/*
+	 * Use the largest block size among supported options.  This wastes some
+	 * memory but simplifies the allocation logic.
+	 */
+	uint8		k_ipad[PG_SHA512_BLOCK_LENGTH];
+	uint8		k_opad[PG_SHA512_BLOCK_LENGTH];
+};
+
+#define HMAC_IPAD 0x36
+#define HMAC_OPAD 0x5C
+
+/*
+ * pg_hmac_create
+ *
+ * Allocate a hash context.  Returns NULL on failure for an OOM.  The
+ * backend issues an error, without returning.
+ */
+pg_hmac_ctx *
+pg_hmac_create(pg_cryptohash_type type)
+{
+	pg_hmac_ctx *ctx;
+
+	ctx = ALLOC(sizeof(pg_hmac_ctx));
+	if (ctx == NULL)
+		return NULL;
+	memset(ctx, 0, sizeof(pg_hmac_ctx));
+	ctx->type = type;
+
+	/*
+	 * Initialize the context data.  This requires to know the digest and
+	 * block lengths, that depend on the type of hash used.
+	 */
+	switch (type)
+	{
+		case PG_MD5:
+			ctx->digest_size = MD5_DIGEST_LENGTH;
+			ctx->block_size = MD5_BLOCK_SIZE;
+			break;
+		case PG_SHA1:
+			ctx->digest_size = SHA1_DIGEST_LENGTH;
+			ctx->block_size = SHA1_BLOCK_SIZE;
+			break;
+		case PG_SHA224:
+			ctx->digest_size = PG_SHA224_DIGEST_LENGTH;
+			ctx->block_size = PG_SHA224_BLOCK_LENGTH;
+			break;
+		case PG_SHA256:
+			ctx->digest_size = PG_SHA256_DIGEST_LENGTH;
+			ctx->block_size = PG_SHA256_BLOCK_LENGTH;
+			break;
+		case PG_SHA384:
+			ctx->digest_size = PG_SHA384_DIGEST_LENGTH;
+			ctx->block_size = PG_SHA384_BLOCK_LENGTH;
+			break;
+		case PG_SHA512:
+			ctx->digest_size = PG_SHA512_DIGEST_LENGTH;
+			ctx->block_size = PG_SHA512_BLOCK_LENGTH;
+			break;
+	}
+
+	ctx->hash = pg_cryptohash_create(type);
+	if (ctx->hash == NULL)
+	{
+		explicit_bzero(ctx, sizeof(pg_hmac_ctx));
+		FREE(ctx);
+		return NULL;
+	}
+
+	return ctx;
+}
+
+/*
+ * pg_hmac_init
+ *
+ * Initialize a HMAC context.  Returns 0 on success, -1 on failure.
+ */
+int
+pg_hmac_init(pg_hmac_ctx *ctx, const uint8 *key, size_t len)
+{
+	int			i;
+	int			digest_size;
+	int			block_size;
+	uint8	   *shrinkbuf = NULL;
+
+	if (ctx == NULL)
+		return -1;
+
+	digest_size = ctx->digest_size;
+	block_size = ctx->block_size;
+
+	memset(ctx->k_opad, HMAC_OPAD, ctx->block_size);
+	memset(ctx->k_ipad, HMAC_IPAD, ctx->block_size);
+
+	/*
+	 * If the key is longer than the block size, pass it through the hash once
+	 * to shrink it down.
+	 */
+	if (len > block_size)
+	{
+		pg_cryptohash_ctx *hash_ctx;
+
+		/* temporary buffer for one-time shrink */
+		shrinkbuf = ALLOC(digest_size);
+		if (shrinkbuf == NULL)
+			return -1;
+		memset(shrinkbuf, 0, digest_size);
+
+		hash_ctx = pg_cryptohash_create(ctx->type);
+		if (hash_ctx == NULL)
+		{
+			FREE(shrinkbuf);
+			return -1;
+		}
+
+		if (pg_cryptohash_init(hash_ctx) < 0 ||
+			pg_cryptohash_update(hash_ctx, key, len) < 0 ||
+			pg_cryptohash_final(hash_ctx, shrinkbuf, digest_size) < 0)
+		{
+			pg_cryptohash_free(hash_ctx);
+			FREE(shrinkbuf);
+			return -1;
+		}
+
+		key = shrinkbuf;
+		len = digest_size;
+		pg_cryptohash_free(hash_ctx);
+	}
+
+	for (i = 0; i < len; i++)
+	{
+		ctx->k_ipad[i] ^= key[i];
+		ctx->k_opad[i] ^= key[i];
+	}
+
+	/* tmp = H(K XOR ipad, text) */
+	if (pg_cryptohash_init(ctx->hash) < 0 ||
+		pg_cryptohash_update(ctx->hash, ctx->k_ipad, ctx->block_size) < 0)
+	{
+		if (shrinkbuf)
+			FREE(shrinkbuf);
+		return -1;
+	}
+
+	if (shrinkbuf)
+		FREE(shrinkbuf);
+	return 0;
+}
+
+/*
+ * pg_hmac_update
+ *
+ * Update a HMAC context.  Returns 0 on success, -1 on failure.
+ */
+int
+pg_hmac_update(pg_hmac_ctx *ctx, const uint8 *data, size_t len)
+{
+	if (ctx == NULL)
+		return -1;
+
+	if (pg_cryptohash_update(ctx->hash, data, len) < 0)
+		return -1;
+
+	return 0;
+}
+
+/*
+ * pg_hmac_final
+ *
+ * Finalize a HMAC context.  Returns 0 on success, -1 on failure.
+ */
+int
+pg_hmac_final(pg_hmac_ctx *ctx, uint8 *dest, size_t len)
+{
+	uint8	   *h;
+
+	if (ctx == NULL)
+		return -1;
+
+	h = ALLOC(ctx->digest_size);
+	if (h == NULL)
+		return -1;
+	memset(h, 0, ctx->digest_size);
+
+	if (pg_cryptohash_final(ctx->hash, h, ctx->digest_size) < 0)
+		return -1;
+
+	/* H(K XOR opad, tmp) */
+	if (pg_cryptohash_init(ctx->hash) < 0 ||
+		pg_cryptohash_update(ctx->hash, ctx->k_opad, ctx->block_size) < 0 ||
+		pg_cryptohash_update(ctx->hash, h, ctx->digest_size) < 0 ||
+		pg_cryptohash_final(ctx->hash, dest, len) < 0)
+	{
+		return -1;
+	}
+
+	return 0;
+}
+
+/*
+ * pg_hmac_free
+ *
+ * Free a HMAC context.
+ */
+void
+pg_hmac_free(pg_hmac_ctx *ctx)
+{
+	if (ctx == NULL)
+		return;
+
+	pg_cryptohash_free(ctx->hash);
+	explicit_bzero(ctx, sizeof(pg_hmac_ctx));
+	FREE(ctx);
+}
diff --git a/src/common/hmac_openssl.c b/src/common/hmac_openssl.c
new file mode 100644
index 00000000000..b5e3065d1a9
--- /dev/null
+++ b/src/common/hmac_openssl.c
@@ -0,0 +1,256 @@
+/*-------------------------------------------------------------------------
+ *
+ * hmac_openssl.c
+ *	  Implementation of HMAC with OpenSSL.
+ *
+ * This should only be used if code is compiled with OpenSSL support.
+ *
+ * Portions Copyright (c) 1996-2020, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  src/common/hmac_openssl.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#ifndef FRONTEND
+#include "postgres.h"
+#else
+#include "postgres_fe.h"
+#endif
+
+#include <openssl/hmac.h>
+
+#include "common/hmac.h"
+#include "common/md5.h"
+#include "common/sha1.h"
+#include "common/sha2.h"
+#ifndef FRONTEND
+#include "utils/memutils.h"
+#include "utils/resowner.h"
+#include "utils/resowner_private.h"
+#endif
+
+/*
+ * In backend, use an allocation in TopMemoryContext to count for resowner
+ * cleanup handling if necesary.  For versions of OpenSSL where HMAC_CTX is
+ * known, just use palloc().  In frontend, use malloc to be able to return
+ * a failure status back to the caller.
+ */
+#ifndef FRONTEND
+#ifdef HAVE_HMAC_CTX_NEW
+#define ALLOC(size) MemoryContextAlloc(TopMemoryContext, size)
+#else
+#define ALLOC(size) palloc(size)
+#endif
+#define FREE(ptr) pfree(ptr)
+#else							/* FRONTEND */
+#define ALLOC(size) malloc(size)
+#define FREE(ptr) free(ptr)
+#endif							/* FRONTEND */
+
+/*
+ * Internal structure for pg_hmac_ctx->data with this implementation.
+ */
+struct pg_hmac_ctx
+{
+	HMAC_CTX   *hmacctx;
+	pg_cryptohash_type type;
+
+#ifndef FRONTEND
+	ResourceOwner resowner;
+#endif
+};
+
+/*
+ * pg_hmac_create
+ *
+ * Allocate a hash context.  Returns NULL on failure for an OOM.  The
+ * backend issues an error, without returning.
+ */
+pg_hmac_ctx *
+pg_hmac_create(pg_cryptohash_type type)
+{
+	pg_hmac_ctx *ctx;
+
+	ctx = ALLOC(sizeof(pg_hmac_ctx));
+	if (ctx == NULL)
+		return NULL;
+	memset(ctx, 0, sizeof(pg_hmac_ctx));
+
+	ctx->type = type;
+
+	/*
+	 * Initialization takes care of assigning the correct type for OpenSSL.
+	 */
+#ifdef HAVE_HMAC_CTX_NEW
+#ifndef FRONTEND
+	ResourceOwnerEnlargeHMAC(CurrentResourceOwner);
+#endif
+	ctx->hmacctx = HMAC_CTX_new();
+#else
+	ctx->hmacctx = ALLOC(sizeof(HMAC_CTX));
+#endif
+
+	if (ctx->hmacctx == NULL)
+	{
+		explicit_bzero(ctx, sizeof(pg_hmac_ctx));
+		FREE(ctx);
+#ifndef FRONTEND
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory")));
+#endif
+		return NULL;
+	}
+
+#ifdef HAVE_HMAC_CTX_NEW
+#ifndef FRONTEND
+	ctx->resowner = CurrentResourceOwner;
+	ResourceOwnerRememberHMAC(CurrentResourceOwner, PointerGetDatum(ctx));
+#endif
+#else
+	memset(ctx->hmacctx, 0, sizeof(HMAC_CTX));
+#endif							/* HAVE_HMAC_CTX_NEW */
+
+	return ctx;
+}
+
+/*
+ * pg_hmac_init
+ *
+ * Initialize a HMAC context.  Returns 0 on success, -1 on failure.
+ */
+int
+pg_hmac_init(pg_hmac_ctx *ctx, const uint8 *key, size_t len)
+{
+	int			status = 0;
+
+	if (ctx == NULL)
+		return -1;
+
+	switch (ctx->type)
+	{
+		case PG_MD5:
+			status = HMAC_Init_ex(ctx->hmacctx, key, len, EVP_md5(), NULL);
+			break;
+		case PG_SHA1:
+			status = HMAC_Init_ex(ctx->hmacctx, key, len, EVP_sha1(), NULL);
+			break;
+		case PG_SHA224:
+			status = HMAC_Init_ex(ctx->hmacctx, key, len, EVP_sha224(), NULL);
+			break;
+		case PG_SHA256:
+			status = HMAC_Init_ex(ctx->hmacctx, key, len, EVP_sha256(), NULL);
+			break;
+		case PG_SHA384:
+			status = HMAC_Init_ex(ctx->hmacctx, key, len, EVP_sha384(), NULL);
+			break;
+		case PG_SHA512:
+			status = HMAC_Init_ex(ctx->hmacctx, key, len, EVP_sha512(), NULL);
+			break;
+	}
+
+	/* OpenSSL internals return 1 on success, 0 on failure */
+	if (status <= 0)
+		return -1;
+
+	return 0;
+}
+
+/*
+ * pg_hmac_update
+ *
+ * Update a HMAC context.  Returns 0 on success, -1 on failure.
+ */
+int
+pg_hmac_update(pg_hmac_ctx *ctx, const uint8 *data, size_t len)
+{
+	int			status = 0;
+
+	if (ctx == NULL)
+		return -1;
+
+	status = HMAC_Update(ctx->hmacctx, data, len);
+
+	/* OpenSSL internals return 1 on success, 0 on failure */
+	if (status <= 0)
+		return -1;
+	return 0;
+}
+
+/*
+ * pg_hmac_final
+ *
+ * Finalize a HMAC context.  Returns 0 on success, -1 on failure.
+ */
+int
+pg_hmac_final(pg_hmac_ctx *ctx, uint8 *dest, size_t len)
+{
+	int			status = 0;
+	uint32		outlen;
+
+	if (ctx == NULL)
+		return -1;
+
+	switch (ctx->type)
+	{
+		case PG_MD5:
+			if (len < MD5_DIGEST_LENGTH)
+				return -1;
+			break;
+		case PG_SHA1:
+			if (len < SHA1_DIGEST_LENGTH)
+				return -1;
+			break;
+		case PG_SHA224:
+			if (len < PG_SHA224_DIGEST_LENGTH)
+				return -1;
+			break;
+		case PG_SHA256:
+			if (len < PG_SHA256_DIGEST_LENGTH)
+				return -1;
+			break;
+		case PG_SHA384:
+			if (len < PG_SHA384_DIGEST_LENGTH)
+				return -1;
+			break;
+		case PG_SHA512:
+			if (len < PG_SHA512_DIGEST_LENGTH)
+				return -1;
+			break;
+	}
+
+	status = HMAC_Final(ctx->hmacctx, dest, &outlen);
+
+	/* OpenSSL internals return 1 on success, 0 on failure */
+	if (status <= 0)
+		return -1;
+	return 0;
+}
+
+/*
+ * pg_hmac_free
+ *
+ * Free a HMAC context.
+ */
+void
+pg_hmac_free(pg_hmac_ctx *ctx)
+{
+	if (ctx == NULL)
+		return;
+
+#ifdef HAVE_HMAC_CTX_FREE
+	HMAC_CTX_free(ctx->hmacctx);
+#ifndef FRONTEND
+	ResourceOwnerForgetHMAC(ctx->resowner, PointerGetDatum(ctx));
+#endif
+#else
+	explicit_bzero(ctx->hmacctx, sizeof(HMAC_CTX));
+	FREE(ctx->hmacctx);
+#endif
+
+	explicit_bzero(ctx, sizeof(pg_hmac_ctx));
+	FREE(ctx);
+}
diff --git a/src/common/scram-common.c b/src/common/scram-common.c
index 0b9557376e9..69a96f65f65 100644
--- a/src/common/scram-common.c
+++ b/src/common/scram-common.c
@@ -20,118 +20,10 @@
 #endif
 
 #include "common/base64.h"
+#include "common/hmac.h"
 #include "common/scram-common.h"
 #include "port/pg_bswap.h"
 
-#define HMAC_IPAD 0x36
-#define HMAC_OPAD 0x5C
-
-/*
- * Calculate HMAC per RFC2104.
- *
- * The hash function used is SHA-256.  Returns 0 on success, -1 on failure.
- */
-int
-scram_HMAC_init(scram_HMAC_ctx *ctx, const uint8 *key, int keylen)
-{
-	uint8		k_ipad[SHA256_HMAC_B];
-	int			i;
-	uint8		keybuf[SCRAM_KEY_LEN];
-
-	/*
-	 * If the key is longer than the block size (64 bytes for SHA-256), pass
-	 * it through SHA-256 once to shrink it down.
-	 */
-	if (keylen > SHA256_HMAC_B)
-	{
-		pg_cryptohash_ctx *sha256_ctx;
-
-		sha256_ctx = pg_cryptohash_create(PG_SHA256);
-		if (sha256_ctx == NULL)
-			return -1;
-		if (pg_cryptohash_init(sha256_ctx) < 0 ||
-			pg_cryptohash_update(sha256_ctx, key, keylen) < 0 ||
-			pg_cryptohash_final(sha256_ctx, keybuf, sizeof(keybuf)) < 0)
-		{
-			pg_cryptohash_free(sha256_ctx);
-			return -1;
-		}
-		key = keybuf;
-		keylen = SCRAM_KEY_LEN;
-		pg_cryptohash_free(sha256_ctx);
-	}
-
-	memset(k_ipad, HMAC_IPAD, SHA256_HMAC_B);
-	memset(ctx->k_opad, HMAC_OPAD, SHA256_HMAC_B);
-
-	for (i = 0; i < keylen; i++)
-	{
-		k_ipad[i] ^= key[i];
-		ctx->k_opad[i] ^= key[i];
-	}
-
-	ctx->sha256ctx = pg_cryptohash_create(PG_SHA256);
-	if (ctx->sha256ctx == NULL)
-		return -1;
-
-	/* tmp = H(K XOR ipad, text) */
-	if (pg_cryptohash_init(ctx->sha256ctx) < 0 ||
-		pg_cryptohash_update(ctx->sha256ctx, k_ipad, SHA256_HMAC_B) < 0)
-	{
-		pg_cryptohash_free(ctx->sha256ctx);
-		return -1;
-	}
-
-	return 0;
-}
-
-/*
- * Update HMAC calculation
- * The hash function used is SHA-256.  Returns 0 on success, -1 on failure.
- */
-int
-scram_HMAC_update(scram_HMAC_ctx *ctx, const char *str, int slen)
-{
-	Assert(ctx->sha256ctx != NULL);
-	if (pg_cryptohash_update(ctx->sha256ctx, (const uint8 *) str, slen) < 0)
-	{
-		pg_cryptohash_free(ctx->sha256ctx);
-		return -1;
-	}
-	return 0;
-}
-
-/*
- * Finalize HMAC calculation.
- * The hash function used is SHA-256.  Returns 0 on success, -1 on failure.
- */
-int
-scram_HMAC_final(uint8 *result, scram_HMAC_ctx *ctx)
-{
-	uint8		h[SCRAM_KEY_LEN];
-
-	Assert(ctx->sha256ctx != NULL);
-
-	if (pg_cryptohash_final(ctx->sha256ctx, h, sizeof(h)) < 0)
-	{
-		pg_cryptohash_free(ctx->sha256ctx);
-		return -1;
-	}
-
-	/* H(K XOR opad, tmp) */
-	if (pg_cryptohash_init(ctx->sha256ctx) < 0 ||
-		pg_cryptohash_update(ctx->sha256ctx, ctx->k_opad, SHA256_HMAC_B) < 0 ||
-		pg_cryptohash_update(ctx->sha256ctx, h, SCRAM_KEY_LEN) < 0 ||
-		pg_cryptohash_final(ctx->sha256ctx, result, SCRAM_KEY_LEN) < 0)
-	{
-		pg_cryptohash_free(ctx->sha256ctx);
-		return -1;
-	}
-
-	pg_cryptohash_free(ctx->sha256ctx);
-	return 0;
-}
-
 /*
  * Calculate SaltedPassword.
  *
@@ -149,7 +41,10 @@ scram_SaltedPassword(const char *password,
 				j;
 	uint8		Ui[SCRAM_KEY_LEN];
 	uint8		Ui_prev[SCRAM_KEY_LEN];
-	scram_HMAC_ctx hmac_ctx;
+	pg_hmac_ctx *hmac_ctx = pg_hmac_create(PG_SHA256);
+
+	if (hmac_ctx == NULL)
+		return -1;
 
 	/*
 	 * Iterate hash calculation of HMAC entry using given salt.  This is
@@ -158,11 +53,12 @@ scram_SaltedPassword(const char *password,
 	 */
 
 	/* First iteration */
-	if (scram_HMAC_init(&hmac_ctx, (uint8 *) password, password_len) < 0 ||
-		scram_HMAC_update(&hmac_ctx, salt, saltlen) < 0 ||
-		scram_HMAC_update(&hmac_ctx, (char *) &one, sizeof(uint32)) < 0 ||
-		scram_HMAC_final(Ui_prev, &hmac_ctx) < 0)
+	if (pg_hmac_init(hmac_ctx, (uint8 *) password, password_len) < 0 ||
+		pg_hmac_update(hmac_ctx, (uint8 *) salt, saltlen) < 0 ||
+		pg_hmac_update(hmac_ctx, (uint8 *) &one, sizeof(uint32)) < 0 ||
+		pg_hmac_final(hmac_ctx, Ui_prev, sizeof(Ui_prev)) < 0)
 	{
+		pg_hmac_free(hmac_ctx);
 		return -1;
 	}
 
@@ -171,10 +67,11 @@ scram_SaltedPassword(const char *password,
 	/* Subsequent iterations */
 	for (i = 2; i <= iterations; i++)
 	{
-		if (scram_HMAC_init(&hmac_ctx, (uint8 *) password, password_len) < 0 ||
-			scram_HMAC_update(&hmac_ctx, (const char *) Ui_prev, SCRAM_KEY_LEN) < 0 ||
-			scram_HMAC_final(Ui, &hmac_ctx) < 0)
+		if (pg_hmac_init(hmac_ctx, (uint8 *) password, password_len) < 0 ||
+			pg_hmac_update(hmac_ctx, (uint8 *) Ui_prev, SCRAM_KEY_LEN) < 0 ||
+			pg_hmac_final(hmac_ctx, Ui, sizeof(Ui)) < 0)
 		{
+			pg_hmac_free(hmac_ctx);
 			return -1;
 		}
 
@@ -183,6 +80,7 @@ scram_SaltedPassword(const char *password,
 		memcpy(Ui_prev, Ui, SCRAM_KEY_LEN);
 	}
 
+	pg_hmac_free(hmac_ctx);
 	return 0;
 }
 
@@ -218,15 +116,20 @@ scram_H(const uint8 *input, int len, uint8 *result)
 int
 scram_ClientKey(const uint8 *salted_password, uint8 *result)
 {
-	scram_HMAC_ctx ctx;
+	pg_hmac_ctx *ctx = pg_hmac_create(PG_SHA256);
+
+	if (ctx == NULL)
+		return -1;
 
-	if (scram_HMAC_init(&ctx, salted_password, SCRAM_KEY_LEN) < 0 ||
-		scram_HMAC_update(&ctx, "Client Key", strlen("Client Key")) < 0 ||
-		scram_HMAC_final(result, &ctx) < 0)
+	if (pg_hmac_init(ctx, salted_password, SCRAM_KEY_LEN) < 0 ||
+		pg_hmac_update(ctx, (uint8 *) "Client Key", strlen("Client Key")) < 0 ||
+		pg_hmac_final(ctx, result, SCRAM_KEY_LEN) < 0)
 	{
+		pg_hmac_free(ctx);
 		return -1;
 	}
 
+	pg_hmac_free(ctx);
 	return 0;
 }
 
@@ -236,15 +139,20 @@ scram_ClientKey(const uint8 *salted_password, uint8 *result)
 int
 scram_ServerKey(const uint8 *salted_password, uint8 *result)
 {
-	scram_HMAC_ctx ctx;
+	pg_hmac_ctx *ctx = pg_hmac_create(PG_SHA256);
+
+	if (ctx == NULL)
+		return -1;
 
-	if (scram_HMAC_init(&ctx, salted_password, SCRAM_KEY_LEN) < 0 ||
-		scram_HMAC_update(&ctx, "Server Key", strlen("Server Key")) < 0 ||
-		scram_HMAC_final(result, &ctx) < 0)
+	if (pg_hmac_init(ctx, salted_password, SCRAM_KEY_LEN) < 0 ||
+		pg_hmac_update(ctx, (uint8 *) "Server Key", strlen("Server Key")) < 0 ||
+		pg_hmac_final(ctx, result, SCRAM_KEY_LEN) < 0)
 	{
+		pg_hmac_free(ctx);
 		return -1;
 	}
 
+	pg_hmac_free(ctx);
 	return 0;
 }