cr50: optimize AES/GCM/App cipher implementation

1. Introduced AES register structure and replaced register access with
   accesses to fields. In many cases it reduce code size and number of
   instructions.
2. Deduplication between AES implementation and App Cipher which use
   AES engine with key coming from key ladder. Added internal function
   dcrypto_aes_process() which applies current AES configuration to
   aligned data in highly optimized manner, same as previous outer_loop
   and inner_loop() functions.

Overall it saves 322 bytes with gcc 8.3

BUG=none
TEST=make BOARD=cr50 CRYPTO_TEST=1; test/tpm_test/tpm_test.py
In ccd: cipher [to test app_cipher];
TCG Tests.

Signed-off-by: Vadim Sukhomlinov <sukhomlinov@google.com>
Change-Id: I6551e21e5e8798aa4691cb6ba476d565778cea3d
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/3213610
Reviewed-by: Vadim Sukhomlinov <sukhomlinov@chromium.org>
Reviewed-by: Andrey Pronin <apronin@chromium.org>
Tested-by: Vadim Sukhomlinov <sukhomlinov@chromium.org>
Auto-Submit: Vadim Sukhomlinov <sukhomlinov@chromium.org>
Commit-Queue: Vadim Sukhomlinov <sukhomlinov@chromium.org>

5 files changed, 189 insertions, 153 deletions

diff --git a/board/cr50/dcrypto/aes.c b/board/cr50/dcrypto/aes.c
index a4895b7a6f..e2d80a0795 100644
--- a/board/cr50/dcrypto/aes.c
+++ b/board/cr50/dcrypto/aes.c
@@ -3,26 +3,33 @@
  * found in the LICENSE file.
  */
 
-#include "dcrypto.h"
 #include "internal.h"
-#include "registers.h"
+#include "dcrypto_regs.h"
+/**
+ * Define KEYMGR AES access structure.
+ */
+static volatile struct keymgr_aes *reg_aes = (void *)(GC_KEYMGR_BASE_ADDR);
 
-static void set_control_register(
-	unsigned mode, unsigned key_size, unsigned encrypt)
+static void set_control_register(enum cipher_mode mode, uint32_t key_size,
+				 enum encrypt_mode encrypt)
 {
-	GWRITE_FIELD(KEYMGR, AES_CTRL, RESET, CTRL_NO_SOFT_RESET);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, KEYSIZE, key_size);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, CIPHER_MODE, mode);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, ENC_MODE, encrypt);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, CTR_ENDIAN, CTRL_CTR_BIG_ENDIAN);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, ENABLE, CTRL_ENABLE);
+	/* Make sure we don't use key coming from keyladder. */
+	reg_aes->use_hidden_key = GC_KEYMGR_AES_USE_HIDDEN_KEY_ENABLE_DEFAULT;
+	/* Bring AES engine FIFO into known state. */
+	reg_aes->ctrl = GC_KEYMGR_AES_CTRL_RESET_MASK;
+	reg_aes->ctrl =
+		(CTRL_NO_SOFT_RESET << GC_KEYMGR_AES_CTRL_RESET_LSB) |
+		(key_size << GC_KEYMGR_AES_CTRL_KEYSIZE_LSB) |
+		(mode << GC_KEYMGR_AES_CTRL_CIPHER_MODE_LSB) |
+		(encrypt << GC_KEYMGR_AES_CTRL_ENC_MODE_LSB) |
+		(CTRL_CTR_BIG_ENDIAN << GC_KEYMGR_AES_CTRL_CTR_ENDIAN_LSB) |
+		(CTRL_ENABLE << GC_KEYMGR_AES_CTRL_ENABLE_LSB);
 
 	/* Turn off random nops (which are enabled by default). */
-	GWRITE_FIELD(KEYMGR, AES_RAND_STALL_CTL, STALL_EN, 0);
-	/* Configure random nop percentage at 25%. */
-	GWRITE_FIELD(KEYMGR, AES_RAND_STALL_CTL, FREQ, 1);
-	/* Now turn on random nops. */
-	GWRITE_FIELD(KEYMGR, AES_RAND_STALL_CTL, STALL_EN, 1);
+	reg_aes->rand_stall = 0;
+	/* Configure random nop percentage at 25%, turn on random nops.  */
+	reg_aes->rand_stall = (1 << GC_KEYMGR_AES_RAND_STALL_CTL_FREQ_LSB) |
+			      GC_KEYMGR_AES_RAND_STALL_CTL_STALL_EN_MASK;
 }
 
 static int wait_read_data(volatile uint32_t *addr)
@@ -75,12 +82,12 @@ enum dcrypto_result dcrypto_aes_init(const uint8_t *key, size_t key_len,
 	/* Initialize hardware with AES key */
 	p = (struct access_helper *) key;
 	for (i = 0; i < (key_len >> 5); i++)
-		GR_KEYMGR_AES_KEY(i) = p[i].udata;
+		reg_aes->key[i] = p[i].udata;
 	/* Trigger key expansion. */
-	GREG32(KEYMGR, AES_KEY_START) = 1;
+	reg_aes->key_start = 1;
 
 	/* Wait for key expansion. */
-	if (!wait_read_data(GREG32_ADDR(KEYMGR, AES_KEY_START))) {
+	if (!wait_read_data(&reg_aes->key_start)) {
 		/* Should not happen. */
 		return DCRYPTO_FAIL;
 	}
@@ -103,10 +110,9 @@ enum dcrypto_result DCRYPTO_aes_init(const uint8_t *key, size_t key_len,
 
 enum dcrypto_result DCRYPTO_aes_block(const uint8_t *in, uint8_t *out)
 {
-	int i;
 	uint32_t buf[4];
 	const uint32_t *inw;
-	uint32_t *outw;
+	uint32_t *outw, *outw2;
 
 	if (is_not_aligned(in)) {
 		memcpy(buf, in, sizeof(buf));
@@ -114,27 +120,16 @@ enum dcrypto_result DCRYPTO_aes_block(const uint8_t *in, uint8_t *out)
 	} else
 		inw = (const uint32_t *)in;
 
-	/* Write plaintext. */
-	for (i = 0; i < 4; i++)
-		GREG32(KEYMGR, AES_WFIFO_DATA) = inw[i];
-
-	/* Wait for the result. */
-	if (!wait_read_data(GREG32_ADDR(KEYMGR, AES_RFIFO_EMPTY))) {
-		/* Should not happen, ciphertext not ready. */
-		return DCRYPTO_FAIL;
-	}
-
-	/* Read ciphertext. */
 	if (is_not_aligned(out))
 		outw = buf;
 	else
 		outw = (uint32_t *)out;
 
-	for (i = 0; i < 4; i++)
-		outw[i] = GREG32(KEYMGR, AES_RFIFO_DATA);
+	outw2 = outw;
+	dcrypto_aes_process(&outw, &inw, 16);
 
-	if (out != (uint8_t *)outw)
-		memcpy(out, outw, sizeof(buf));
+	if (out != (uint8_t *)outw2)
+		memcpy(out, outw2, sizeof(buf));
 
 	return DCRYPTO_OK;
 }
@@ -152,7 +147,7 @@ void DCRYPTO_aes_write_iv(const uint8_t *iv)
 		ivw = (uint32_t *)iv;
 
 	for (i = 0; i < 4; i++)
-		GR_KEYMGR_AES_CTR(i) = ivw[i];
+		reg_aes->counter[i] = ivw[i];
 }
 
 void DCRYPTO_aes_read_iv(uint8_t *iv)
@@ -167,7 +162,7 @@ void DCRYPTO_aes_read_iv(uint8_t *iv)
 		ivw = (uint32_t *)iv;
 
 	for (i = 0; i < 4; i++)
-		ivw[i] = GR_KEYMGR_AES_CTR(i);
+		ivw[i] = reg_aes->counter[i];
 
 	if (iv != (uint8_t *)ivw)
 		memcpy(iv, ivw, sizeof(buf));
diff --git a/board/cr50/dcrypto/app_cipher.c b/board/cr50/dcrypto/app_cipher.c
index 004db6bd6c..cf862d246e 100644
--- a/board/cr50/dcrypto/app_cipher.c
+++ b/board/cr50/dcrypto/app_cipher.c
@@ -6,123 +6,121 @@
 #include "crypto_api.h"
 #include "internal.h"
 #include "registers.h"
-
-/* The default build options compile for size (-Os); instruct the
- * compiler to optimize for speed here.  Incidentally -O produces
- * faster code than -O2!
+#include "dcrypto_regs.h"
+/**
+ * Define KEYMGR AES access structure.
  */
-static int __optimize("O") inner_loop(uint32_t **out, const uint32_t **in,
-				      size_t len)
-{
-	uint32_t *outw = *out;
-	const uint32_t *inw = *in;
-
-	while (len >= 16) {
-		uint32_t w0, w1, w2, w3;
-
-		w0 = inw[0];
-		w1 = inw[1];
-		w2 = inw[2];
-		w3 = inw[3];
-		GREG32(KEYMGR, AES_WFIFO_DATA) = w0;
-		GREG32(KEYMGR, AES_WFIFO_DATA) = w1;
-		GREG32(KEYMGR, AES_WFIFO_DATA) = w2;
-		GREG32(KEYMGR, AES_WFIFO_DATA) = w3;
+static volatile struct keymgr_aes *reg_aes = (void *)(GC_KEYMGR_BASE_ADDR);
 
-		while (GREG32(KEYMGR, AES_RFIFO_EMPTY))
-			;
-
-		w0 = GREG32(KEYMGR, AES_RFIFO_DATA);
-		w1 = GREG32(KEYMGR, AES_RFIFO_DATA);
-		w2 = GREG32(KEYMGR, AES_RFIFO_DATA);
-		w3 = GREG32(KEYMGR, AES_RFIFO_DATA);
-		outw[0] = w0;
-		outw[1] = w1;
-		outw[2] = w2;
-		outw[3] = w3;
-
-		inw += 4;
-		outw += 4;
-		len -= 16;
-	}
+static inline void read_aes_fifo(uint32_t *outw)
+{
+	uint32_t w0, w1, w2, w3;
+
+	w0 = reg_aes->rfifo_data;
+	w1 = reg_aes->rfifo_data;
+	w2 = reg_aes->rfifo_data;
+	w3 = reg_aes->rfifo_data;
+	outw[0] = w0;
+	outw[1] = w1;
+	outw[2] = w2;
+	outw[3] = w3;
+}
 
-	*in = inw;
-	*out = outw;
-	return len;
+static inline void write_aes_fifo(const uint32_t *inw)
+{
+	uint32_t w0, w1, w2, w3;
+
+	w0 = inw[0];
+	w1 = inw[1];
+	w2 = inw[2];
+	w3 = inw[3];
+	reg_aes->wfifo_data = w0;
+	reg_aes->wfifo_data = w1;
+	reg_aes->wfifo_data = w2;
+	reg_aes->wfifo_data = w3;
 }
 
-static int outer_loop(uint32_t **out, const uint32_t **in, size_t len)
+/* The default build options compile for size (-Os); instruct the
+ * compiler to optimize for speed here.  Incidentally -O produces
+ * faster code than -O2!
+ *
+ * AES FIFOs are 32-bits wide by 16 entries deep (64-Bytes), so to
+ * get maximum throughput try to load it at least at 50% and that's why we don't
+ * read after the first 4 words are loaded, unless they are the only 4 words
+ */
+size_t __optimize("O")
+	dcrypto_aes_process(uint32_t **out, const uint32_t **in, size_t len)
 {
 	uint32_t *outw = *out;
 	const uint32_t *inw = *in;
 
 	if (len >= 16) {
-		GREG32(KEYMGR, AES_WFIFO_DATA) = inw[0];
-		GREG32(KEYMGR, AES_WFIFO_DATA) = inw[1];
-		GREG32(KEYMGR, AES_WFIFO_DATA) = inw[2];
-		GREG32(KEYMGR, AES_WFIFO_DATA) = inw[3];
+		write_aes_fifo(inw);
 		inw += 4;
 		len -= 16;
 
-		len = inner_loop(&outw, &inw, len);
+		while (len >= 16) {
+			write_aes_fifo(inw);
+			inw += 4;
+			len -= 16;
+			while (reg_aes->rfifo_empty)
+				;
 
-		while (GREG32(KEYMGR, AES_RFIFO_EMPTY))
+			read_aes_fifo(outw);
+			outw += 4;
+		}
+		while (reg_aes->rfifo_empty)
 			;
 
-		outw[0] = GREG32(KEYMGR, AES_RFIFO_DATA);
-		outw[1] = GREG32(KEYMGR, AES_RFIFO_DATA);
-		outw[2] = GREG32(KEYMGR, AES_RFIFO_DATA);
-		outw[3] = GREG32(KEYMGR, AES_RFIFO_DATA);
+		read_aes_fifo(outw);
 		outw += 4;
 	}
-
-	*in =  inw;
+	*in = inw;
 	*out = outw;
 	return len;
 }
 
-static int aes_init(enum dcrypto_appid appid, const uint32_t iv[4])
+static int aes_init(enum dcrypto_appid appid)
 {
+	uint32_t aes_config;
+
 	/* Setup USR-based application key. */
 	if (!DCRYPTO_appkey_init(appid))
 		return 0;
 
 	/* Configure AES engine. */
-	GWRITE_FIELD(KEYMGR, AES_CTRL, RESET, CTRL_NO_SOFT_RESET);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, KEYSIZE, 2 /* AES-256 */);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, CIPHER_MODE, CIPHER_MODE_CTR);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, ENC_MODE, ENCRYPT_MODE);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, CTR_ENDIAN, CTRL_CTR_BIG_ENDIAN);
+	aes_config = (CTRL_NO_SOFT_RESET << GC_KEYMGR_AES_CTRL_RESET_LSB) |
+		     (2 << GC_KEYMGR_AES_CTRL_KEYSIZE_LSB) |
+		     (CIPHER_MODE_CTR << GC_KEYMGR_AES_CTRL_CIPHER_MODE_LSB) |
+		     (ENCRYPT_MODE << GC_KEYMGR_AES_CTRL_ENC_MODE_LSB) |
+		     (CTRL_CTR_BIG_ENDIAN << GC_KEYMGR_AES_CTRL_CTR_ENDIAN_LSB);
+	reg_aes->ctrl = aes_config;
 
 	/*
 	 * For fixed-key, bulk ciphering, turn off random nops (which
 	 * are enabled by default).
 	 */
-	GWRITE_FIELD(KEYMGR, AES_RAND_STALL_CTL, STALL_EN, 0);
+	reg_aes->rand_stall = 0;
 
 	/* Enable hidden key usage, each appid gets its own
 	 * USR, with USR0 starting at 0x2a0.
 	 */
-	GWRITE_FIELD(KEYMGR, AES_USE_HIDDEN_KEY, INDEX,
-		0x2a0 + (appid * 2));
-	GWRITE_FIELD(KEYMGR, AES_USE_HIDDEN_KEY, ENABLE, 1);
-	GWRITE_FIELD(KEYMGR, AES_CTRL, ENABLE, CTRL_ENABLE);
+	reg_aes->use_hidden_key = GC_KEYMGR_AES_USE_HIDDEN_KEY_ENABLE_MASK |
+				  ((0x2a0 + (appid * 2))
+				   << GC_KEYMGR_AES_USE_HIDDEN_KEY_INDEX_LSB);
+
+	reg_aes->ctrl = aes_config |
+			(CTRL_ENABLE << GC_KEYMGR_AES_CTRL_ENABLE_LSB);
 
 	/* Wait for key-expansion. */
-	GREG32(KEYMGR, AES_KEY_START) = 1;
-	while (GREG32(KEYMGR, AES_KEY_START))
+	reg_aes->key_start = 1;
+	while (reg_aes->key_start)
 		;
 
 	/* Check for errors (e.g. USR not correctly setup. */
 	if (GREG32(KEYMGR, HKEY_ERR_FLAGS))
 		return 0;
 
-	/* Set IV. */
-	GR_KEYMGR_AES_CTR(0) = iv[0];
-	GR_KEYMGR_AES_CTR(1) = iv[1];
-	GR_KEYMGR_AES_CTR(2) = iv[2];
-	GR_KEYMGR_AES_CTR(3) = iv[3];
-
 	return 1;
 }
 
@@ -138,15 +136,14 @@ int DCRYPTO_app_cipher(enum dcrypto_appid appid, const void *salt,
 
 	{
 		/* Initialize key, and AES engine. */
-		uint32_t iv[4];
-
-		BUILD_ASSERT(sizeof(iv) == CIPHER_SALT_SIZE);
-		memcpy(iv, salt, sizeof(iv));
-		if (!aes_init(appid, iv))
+		BUILD_ASSERT(CIPHER_SALT_SIZE == 16);
+		if (!aes_init(appid))
 			return 0;
+		/* Set IV. */
+		DCRYPTO_aes_write_iv(salt);
 	}
 
-	len = outer_loop(&outw, &inw, len);
+	len = dcrypto_aes_process(&outw, &inw, len);
 
 	if (len) {
 		/* Cipher the final partial block */
@@ -159,7 +156,7 @@ int DCRYPTO_app_cipher(enum dcrypto_appid appid, const void *salt,
 		tmpoutw = tmpout;
 
 		memcpy(tmpin, inw, len);
-		outer_loop(&tmpoutw, &tmpinw, 16);
+		dcrypto_aes_process(&tmpoutw, &tmpinw, 16);
 		memcpy(outw, tmpout, len);
 	}
 
@@ -283,7 +280,7 @@ static int prepare_running(struct test_info *pinfo)
 static int basic_check(struct test_info *pinfo)
 {
 	size_t half;
-	int i;
+	size_t i;
 	uint32_t *p;
 
 	ccprintf("original data  %ph\n", HEX_BUF(pinfo->p, 16));
@@ -291,7 +288,7 @@ static int basic_check(struct test_info *pinfo)
 	half = (pinfo->test_blob_size/2) & ~3;
 	if (!DCRYPTO_app_cipher(NVMEM, pinfo->p, pinfo->p,
 				pinfo->p + half, half)) {
-		ccprintf("first ecnryption run failed\n");
+		ccprintf("first encryption run failed\n");
 		return EC_ERROR_UNKNOWN;
 	}
 
@@ -341,7 +338,7 @@ static int command_loop(struct test_info *pinfo)
 
 		*p_last_byte = last_byte;
 
-		if (!(iteration % 500))
+		if (!(iteration % 512))
 			watchdog_reload();
 
 		tstamp = get_time().val;
diff --git a/board/cr50/dcrypto/dcrypto_regs.h b/board/cr50/dcrypto/dcrypto_regs.h
index 7a524220fe..6d118f750a 100644
--- a/board/cr50/dcrypto/dcrypto_regs.h
+++ b/board/cr50/dcrypto/dcrypto_regs.h
@@ -29,7 +29,42 @@ extern "C" {
 #endif
 #include "registers.h"
 
-
+/**
+ * AES/GCM part of KEYMGR starting offset 0x000
+ */
+struct keymgr_aes {
+	uint32_t ctrl; /* KEYMGR_AES_CTRL */
+	const uint32_t _pad0;
+	uint32_t wfifo_data; /*KEYMGR_AES_WFIFO_DATA */
+	const uint32_t rfifo_data; /* KEYMGR_AES_RFIFO_DATA */
+	const uint32_t _pad1[7];
+	uint32_t key[8]; /* KEYMGR_AES_KEY0 .. 7 */
+	uint32_t key_start; /* KEYMGR_AES_KEY_START */
+	uint32_t counter[4]; /* KEYMGR_AES_CTR0 .. 3 */
+	uint32_t rand_stall; /* KEYMGR_AES_RAND_STALL_CTL */
+	uint32_t wfifo_level; /* KEYMGR_AES_WFIFO_LEVEL */
+	uint32_t wfifo_full; /* KEYMGR_AES_WFIFO_FULL */
+	uint32_t rfifo_level; /* KEYMGR_AES_RFIFO_LEVEL */
+	uint32_t rfifo_empty; /* KEYMGR_AES_RFIFO_EMPTY */
+	const uint32_t execute_count_state; /* KEYMGR_AES_EXECUTE_COUNT_STATE */
+	uint32_t execute_count_max; /* KEYMGR_AES_EXECUTE_COUNT_MAX */
+	uint32_t gcm_do_acc; /* KEYMGR_GCM_DO_ACC */
+	uint32_t gcm_h[4]; /* KEYMGR_GCM_H0 .. 3 */
+	uint32_t gcm_mac[4]; /* KEYMGR_GCM_MAC0 .. 3 */
+	uint32_t gcm_hash_in[4]; /* KEYMGR_GCM_HASH_IN0 .. 3 */
+	uint32_t wipe_secrets; /* KEYMGR_AES_WIPE_SECRETS */
+	uint32_t int_enable; /* KEYMGR_AES_INT_ENABLE */
+	uint32_t int_state; /* KEYMGR_AES_INT_STATE */
+	uint32_t int_test; /* KEYMGR_AES_INT_TEST */
+	uint32_t use_hidden_key; /* KEYMGR_AES_USE_HIDDEN_KEY */
+};
+BUILD_ASSERT(offsetof(struct keymgr_aes, wfifo_data) ==
+	     GC_KEYMGR_AES_WFIFO_DATA_OFFSET);
+BUILD_ASSERT(offsetof(struct keymgr_aes, key) == GC_KEYMGR_AES_KEY0_OFFSET);
+BUILD_ASSERT(offsetof(struct keymgr_aes, counter) == GC_KEYMGR_AES_CTR0_OFFSET);
+BUILD_ASSERT(offsetof(struct keymgr_aes, gcm_h) == GC_KEYMGR_GCM_H0_OFFSET);
+BUILD_ASSERT(offsetof(struct keymgr_aes, use_hidden_key) ==
+	     GC_KEYMGR_AES_USE_HIDDEN_KEY_OFFSET);
 
 /**
  * SHA/HMAC part of KEYMGR starting offset 0x400
diff --git a/board/cr50/dcrypto/gcm.c b/board/cr50/dcrypto/gcm.c
index ff57a1e635..96e8b1dbae 100644
--- a/board/cr50/dcrypto/gcm.c
+++ b/board/cr50/dcrypto/gcm.c
@@ -3,39 +3,38 @@
  * found in the LICENSE file.
  */
 
-#include "dcrypto.h"
 #include "internal.h"
 #include "registers.h"
-
+#include "dcrypto_regs.h"
 #include "endian.h"
 
+/**
+ * Define KEYMGR AES access structure.
+ */
+static volatile struct keymgr_aes *reg_aes = (void *)(GC_KEYMGR_BASE_ADDR);
+
 static void gcm_mul(uint32_t *counter)
 {
-	int i;
-	volatile uint32_t *p;
+	size_t i;
 
 	/* Set HASH to zero. */
-	p = GREG32_ADDR(KEYMGR, GCM_HASH_IN0);
 	for (i = 0; i < 4; i++)
-		*p++ = 0;
+		reg_aes->gcm_hash_in[i] = 0;
 
 	/* Initialize GMAC. */
-	p = GREG32_ADDR(KEYMGR, GCM_MAC0);
 	for (i = 0; i < 4; i++)
-		*p++ = counter[i];
+		reg_aes->gcm_mac[i] = counter[i];
 
 	/* Crank GMAC. */
-	GREG32(KEYMGR, GCM_DO_ACC) = 1;
+	reg_aes->gcm_do_acc = 1;
 
 	/* Read GMAC. */
-	p = GREG32_ADDR(KEYMGR, GCM_MAC0);
 	for (i = 0; i < 4; i++)
-		counter[i] = *p++;
+		counter[i] = reg_aes->gcm_mac[i];
 
 	/* Reset GMAC. */
-	p = GREG32_ADDR(KEYMGR, GCM_MAC0);
 	for (i = 0; i < 4; ++i)
-		*p++ = 0;
+		reg_aes->gcm_mac[i] = 0;
 }
 
 static void gcm_init_iv(
@@ -83,8 +82,8 @@ static void gcm_init_iv(
 void DCRYPTO_gcm_init(struct GCM_CTX *ctx, uint32_t key_bits,
 		const uint8_t *key, const uint8_t *iv, size_t iv_len)
 {
-	int i;
-	const uint32_t zero[4] = {0, 0, 0, 0};
+	size_t i;
+	static const uint32_t zero[4] = {0, 0, 0, 0};
 	uint32_t H[4];
 	uint32_t counter[4];
 
@@ -98,38 +97,45 @@ void DCRYPTO_gcm_init(struct GCM_CTX *ctx, uint32_t key_bits,
 
 	/* Initialize the GMAC accumulator to ZERO. */
 	for (i = 0; i < 4; i++)
-		GR_KEYMGR_GCM_MAC(i) = zero[i];
+		reg_aes->gcm_mac[i] = 0;
 
 	/* Initialize H. */
 	for (i = 0; i < 4; i++)
-		GR_KEYMGR_GCM_H(i) = H[i];
+		reg_aes->gcm_h[i] = H[i];
 
 	/* Map the IV to a 128-bit counter. */
 	gcm_init_iv(iv, iv_len, counter);
 
 	/* Re-initialize the IV counter. */
 	for (i = 0; i < 4; i++)
-		GR_KEYMGR_AES_CTR(i) = counter[i];
+		reg_aes->counter[i] = counter[i];
 
 	/* Calculate Ej0: encrypt IV counter XOR ZERO. */
 	DCRYPTO_aes_block((const uint8_t *) zero, ctx->Ej0.c);
 }
 
-static void gcm_aad_block(const struct GCM_CTX *ctx, const uint32_t *block)
+static void gcm_aad_block(const struct GCM_CTX *ctx, const void *block)
 {
-	int i;
-	const struct access_helper *p = (struct access_helper *) block;
+	size_t i;
+	uint32_t buf[4];
+	const uint32_t *p;
+
+	if (is_not_aligned(block)) {
+		memcpy(buf, block, 16);
+		p = buf;
+	} else
+		p = block;
 
 	if (ctx->aad_len == 0 && ctx->count <= 16) {
 		/* Update GMAC. */
 		for (i = 0; i < 4; i++)
-			GR_KEYMGR_GCM_MAC(i) = p[i].udata;
+			reg_aes->gcm_mac[i] = p[i];
 	} else {
 		for (i = 0; i < 4; i++)
-			GR_KEYMGR_GCM_HASH_IN(i) = p[i].udata;
+			reg_aes->gcm_hash_in[i] = p[i];
 
 		/* Crank GMAC. */
-		GREG32(KEYMGR, GCM_DO_ACC) = 1;
+		reg_aes->gcm_do_acc = 1;
 	}
 }
 
@@ -173,7 +179,7 @@ int DCRYPTO_gcm_encrypt(struct GCM_CTX *ctx, uint8_t *out, size_t out_len,
 
 		DCRYPTO_aes_block(ctx->block.c, outp);
 		ctx->count += 16;
-		gcm_aad_block(ctx, (uint32_t *) outp);
+		gcm_aad_block(ctx, outp);
 		ctx->remainder = 0;
 		in += count;
 		in_len -= count;
@@ -184,7 +190,7 @@ int DCRYPTO_gcm_encrypt(struct GCM_CTX *ctx, uint8_t *out, size_t out_len,
 		DCRYPTO_aes_block(in, outp);
 		ctx->count += 16;
 
-		gcm_aad_block(ctx, (uint32_t *) outp);
+		gcm_aad_block(ctx, outp);
 
 		in_len -= 16;
 		in += 16;
@@ -304,23 +310,23 @@ static void dcrypto_gcm_len_vector(
 
 static void dcrypto_gcm_tag(const struct GCM_CTX *ctx,
 			const uint32_t *len_vector, uint32_t *tag) {
-	int i;
+	size_t i;
 
 	for (i = 0; i < 4; i++)
-		GR_KEYMGR_GCM_HASH_IN(i) = len_vector[i];
+		reg_aes->gcm_hash_in[i] = len_vector[i];
 
 	/* Crank GMAC. */
-	GREG32(KEYMGR, GCM_DO_ACC) = 1;
+	reg_aes->gcm_do_acc = 1;
 
 	for (i = 0; i < 4; i++)
-		GR_KEYMGR_GCM_HASH_IN(i) = ctx->Ej0.d[i];
+		reg_aes->gcm_hash_in[i] = ctx->Ej0.d[i];
 
 	/* Crank GMAC. */
-	GREG32(KEYMGR, GCM_DO_ACC) = 1;
+	reg_aes->gcm_do_acc = 1;
 
 	/* Read tag. */
 	for (i = 0; i < 4; i++)
-		tag[i] = GR_KEYMGR_GCM_MAC(i);
+		tag[i] = reg_aes->gcm_mac[i];
 }
 
 int DCRYPTO_gcm_tag(struct GCM_CTX *ctx, uint8_t *tag, size_t tag_len)
@@ -339,5 +345,5 @@ int DCRYPTO_gcm_tag(struct GCM_CTX *ctx, uint8_t *tag, size_t tag_len)
 void DCRYPTO_gcm_finish(struct GCM_CTX *ctx)
 {
 	always_memset(ctx, 0, sizeof(struct GCM_CTX));
-	GREG32(KEYMGR, AES_WIPE_SECRETS) = 1;
+	reg_aes->wipe_secrets = 1;
 }
diff --git a/board/cr50/dcrypto/internal.h b/board/cr50/dcrypto/internal.h
index 706baeaa17..f69ce4d059 100644
--- a/board/cr50/dcrypto/internal.h
+++ b/board/cr50/dcrypto/internal.h
@@ -778,6 +778,9 @@ __always_inline enum dcrypto_result dcrypto_ok_if_zero(uintptr_t val)
 							    DCRYPTO_OK);
 }
 
+/* Process blocks using configured AES engine. */
+size_t dcrypto_aes_process(uint32_t **out, const uint32_t **in, size_t len);
+
 /*
  * Key ladder.
  */