Implement streaming aes+hmac operations to try and optimize mem

Somewhat slower?

page_swap.c

@@ -172,53 +172,53 @@ pswap_establish_hmac_key(void) {
 }
 #endif  // USE_PAGE_CRYPTO
 
-static void
-pswap_encrypt(const void* addr, void* dst, uint64_t pageout_ctr) {
-  size_t len = RISCV_PAGE_SIZE;
-
-#ifdef USE_PAGE_CRYPTO
-  pswap_establish_boot_key();
-  uint8_t iv[32] = {0};
+typedef struct {
   WORD key_sched[80];
-  aes_key_setup(pswap_boot_key, key_sched, 256);
-
-  memcpy(iv + 8, &pageout_ctr, 8);
+  uint8_t ctr[32];
+} aes_ctx_t;
 
-  aes_encrypt_ctr((uint8_t*)addr, len, (uint8_t*)dst, key_sched, 256, iv);
-#else
-  memcpy(dst, addr, len);
-#endif
-}
-
-static void
-pswap_decrypt(const void* addr, void* dst, uint64_t pageout_ctr) {
-  size_t len = RISCV_PAGE_SIZE;
-
-#ifdef USE_PAGE_CRYPTO
+static aes_ctx_t
+aes_ctx_init(uint64_t pageout_ctr) {
   pswap_establish_boot_key();
-  uint8_t iv[32] = {0};
-  WORD key_sched[80];
-  aes_key_setup(pswap_boot_key, key_sched, 256);
-
-  memcpy(iv + 8, &pageout_ctr, 8);
-
-  aes_decrypt_ctr((uint8_t*)addr, len, (uint8_t*)dst, key_sched, 256, iv);
-#else
-  memcpy(dst, addr, len);
-#endif
+  aes_ctx_t ctx = {};
+  aes_key_setup(pswap_boot_key, ctx.key_sched, 256);
+  memcpy(ctx.ctr + 8, &pageout_ctr, 8);
+  return ctx;
 }
 
-static void
-pswap_hmac(uint8_t* hash, void* page_addr, uint64_t pageout_ctr) {
-#ifdef USE_PAGE_HASH
+static HMACContext
+hmac_ctx_init(uint64_t pageout_ctr) {
   pswap_establish_hmac_key();
   HMACContext ctx;
 
   hmacReset(&ctx, pswap_hmac_key, sizeof pswap_hmac_key);
-  hmacInput(&ctx, page_addr, RISCV_PAGE_SIZE);
-  hmacInput(&ctx, (uint8_t*)&pageout_ctr, sizeof(pageout_ctr));
-  hmacResult(&ctx, hash);
-#endif
+  hmacInput(&ctx, (uint8_t*)&pageout_ctr, sizeof pageout_ctr);
+  return ctx;
+}
+
+#define PSWAP_CHUNK_SIZE 32
+void
+pswap_process_chunk(
+    void* back, void* front, bool swap, aes_ctx_t* out_aes,
+    HMACContext* out_hmac, aes_ctx_t* in_aes, HMACContext* in_hmac) {
+  uint8_t tmp[PSWAP_CHUNK_SIZE];
+
+  if (swap) {
+    memcpy(tmp, back, PSWAP_CHUNK_SIZE);
+  }
+  hmacInput(out_hmac, front, PSWAP_CHUNK_SIZE);
+  aes_encrypt_ctr(
+      front, PSWAP_CHUNK_SIZE, back, out_aes->key_sched, 256, out_aes->ctr);
+  for (int i = 0; i < PSWAP_CHUNK_SIZE / AES_BLOCK_SIZE; i++)
+    increment_iv(out_aes->ctr, 8);
+
+  if (swap) {
+    aes_decrypt_ctr(
+        tmp, PSWAP_CHUNK_SIZE, front, in_aes->key_sched, 256, in_aes->ctr);
+    for (int i = 0; i < PSWAP_CHUNK_SIZE / AES_BLOCK_SIZE; i++)
+      increment_iv(in_aes->ctr, 8);
+    hmacInput(in_hmac, front, PSWAP_CHUNK_SIZE);
+  }
 }
 
 /* evict a page from EPM and store it to the backing storage
@@ -230,36 +230,30 @@ page_swap_epm(uintptr_t back_page, uintptr_t epm_page, uintptr_t swap_page) {
   assert(paging_epm_inbounds(epm_page));
   assert(paging_backpage_inbounds(back_page));
 
-  char buffer[RISCV_PAGE_SIZE] = {};
-  if (swap_page) {
-    assert(swap_page == back_page);
-    memcpy(buffer, (void*)swap_page, RISCV_PAGE_SIZE);
-  }
+  assert(!swap_page || swap_page == back_page);
 
   uint64_t* pageout_ctr    = pswap_pageout_ctr(back_page);
   hash_t* hmac             = pswap_pageout_hmac(back_page);
   uint64_t old_pageout_ctr = *pageout_ctr;
   uint64_t new_pageout_ctr = old_pageout_ctr + 1;
 
-  uint8_t new_hmac[32];
-  pswap_hmac(new_hmac, (void*)epm_page, new_pageout_ctr);
-  pswap_encrypt((void*)epm_page, (void*)back_page, new_pageout_ctr);
-
-  if (swap_page) {
-    uint8_t old_hmac[32];
-    pswap_decrypt((void*)buffer, (void*)epm_page, old_pageout_ctr);
-    pswap_hmac(old_hmac, (void*)epm_page, old_pageout_ctr);
+  aes_ctx_t out_aes    = aes_ctx_init(new_pageout_ctr);
+  HMACContext out_hmac = hmac_ctx_init(new_pageout_ctr);
+  aes_ctx_t in_aes     = aes_ctx_init(old_pageout_ctr);
+  HMACContext in_hmac  = hmac_ctx_init(old_pageout_ctr);
 
-#ifdef USE_PAGE_HASH
-    if (memcmp(hmac, old_hmac, sizeof(hash_t))) {
-      return -1;
-    }
-#endif
+  for (size_t cpos = 0; cpos < RISCV_PAGE_SIZE; cpos += PSWAP_CHUNK_SIZE) {
+    pswap_process_chunk(
+        (uint8_t*)back_page + cpos, (uint8_t*)epm_page + cpos, !!swap_page,
+        &out_aes, &out_hmac, &in_aes, &in_hmac);
   }
 
-#ifdef USE_PAGE_HASH
-  *hmac = *(hash_t*)new_hmac;
-#endif
+  uint8_t old_hmac[32];
+  hmacResult(&in_hmac, old_hmac);
+  if (swap_page && memcmp(hmac, old_hmac, sizeof(hash_t))) {
+    return -1;
+  }
+  hmacResult(&out_hmac, hmac->val);
 
   *pageout_ctr = new_pageout_ctr;