From 9128454750109a86d9168df150bafcd6ca5ef0fe Mon Sep 17 00:00:00 2001
From: Simon Tatham <anakin@pobox.com>
Date: Fri, 11 Jan 2019 06:32:00 +0000
Subject: [PATCH] Localise AESContext into sshaes.c.

All access to AES throughout the code is now done via the ssh2_cipher
vtable interface. All code that previously made direct calls to the
underlying functions (for encrypting and decrypting private key files)
now does it by instantiating an ssh2_cipher.

This removes constraints on the AES module's internal structure, and
allows me to reorganise it as much as I like.
---
 import.c | 45 +++++++++++++++++++--------------------------
 ssh.h    | 11 -----------
 sshaes.c | 25 +++++++++++++++----------
 3 files changed, 34 insertions(+), 47 deletions(-)

diff --git a/import.c b/import.c
index 99baf928..22ec9e81 100644
--- a/import.c
+++ b/import.c
@@ -547,13 +547,11 @@ static ssh2_userkey *openssh_pem_read(
             des3_decrypt_pubkey_ossh(keybuf, key->iv,
                                      key->keyblob->u, key->keyblob->len);
         else {
-            AESContext *ctx;
-            assert(key->encryption == OP_E_AES);
-            ctx = aes_make_context();
-            aes128_key(ctx, keybuf);
-            aes_iv(ctx, key->iv);
-            aes_ssh2_decrypt_blk(ctx, key->keyblob->u, key->keyblob->len);
-            aes_free_context(ctx);
+            ssh2_cipher *cipher = ssh2_cipher_new(&ssh_aes128);
+            ssh2_cipher_setkey(cipher, keybuf);
+            ssh2_cipher_setiv(cipher, key->iv);
+            ssh2_cipher_decrypt(cipher, key->keyblob->u, key->keyblob->len);
+            ssh2_cipher_free(cipher);
         }
 
         smemclr(&md5c, sizeof(md5c));
@@ -1390,20 +1388,16 @@ static ssh2_userkey *openssh_new_read(
                 goto error;
             }
             {
-                void *ctx = aes_make_context();
-                aes256_key(ctx, keybuf);
-                aes_iv(ctx, keybuf + 32);
+                ssh2_cipher *cipher = ssh2_cipher_new(
+                    key->cipher == ON_E_AES256CBC ?
+                    &ssh_aes256 : &ssh_aes256_ctr);
+                ssh2_cipher_setkey(cipher, keybuf);
+                ssh2_cipher_setiv(cipher, keybuf + 32);
                 /* Decrypt the private section in place, casting away
                  * the const from key->private being a ptrlen */
-                if (key->cipher == ON_E_AES256CBC) {
-                    aes_ssh2_decrypt_blk(ctx, (char *)key->private.ptr,
-                                         key->private.len);
-                }
-                else {
-                    aes_ssh2_sdctr(ctx, (char *)key->private.ptr,
-                                   key->private.len);
-                }
-                aes_free_context(ctx);
+                ssh2_cipher_decrypt(cipher, (char *)key->private.ptr,
+                                    key->private.len);
+                ssh2_cipher_free(cipher);
             }
             break;
           default:
@@ -1594,18 +1588,17 @@ static bool openssh_new_write(
              * material: 32 bytes AES key + 16 bytes iv.
              */
             unsigned char keybuf[48];
-            void *ctx;
+            ssh2_cipher *cipher;
 
             openssh_bcrypt(passphrase,
                            bcrypt_salt, sizeof(bcrypt_salt), bcrypt_rounds,
                            keybuf, sizeof(keybuf));
 
-            ctx = aes_make_context();
-            aes256_key(ctx, keybuf);
-            aes_iv(ctx, keybuf + 32);
-            aes_ssh2_sdctr(ctx, cpblob->u,
-                           cpblob->len);
-            aes_free_context(ctx);
+            cipher = ssh2_cipher_new(&ssh_aes256_ctr);
+            ssh2_cipher_setkey(cipher, keybuf);
+            ssh2_cipher_setiv(cipher, keybuf + 32);
+            ssh2_cipher_encrypt(cipher, cpblob->u, cpblob->len);
+            ssh2_cipher_free(cipher);
 
             smemclr(keybuf, sizeof(keybuf));
         }
diff --git a/ssh.h b/ssh.h
index 2fdd5978..8d63af13 100644
--- a/ssh.h
+++ b/ssh.h
@@ -903,17 +903,6 @@ extern const ssh2_macalg ssh_hmac_sha256;
 extern const ssh2_macalg ssh2_poly1305;
 extern const ssh_compression_alg ssh_zlib;
 
-typedef struct AESContext AESContext;
-AESContext *aes_make_context(void);
-void aes_free_context(AESContext *ctx);
-void aes128_key(AESContext *ctx, const void *key);
-void aes192_key(AESContext *ctx, const void *key);
-void aes256_key(AESContext *ctx, const void *key);
-void aes_iv(AESContext *ctx, const void *iv);
-void aes_ssh2_encrypt_blk(AESContext *ctx, void *blk, int len);
-void aes_ssh2_decrypt_blk(AESContext *ctx, void *blk, int len);
-void aes_ssh2_sdctr(AESContext *ctx, void *blk, int len);
-
 /*
  * PuTTY version number formatted as an SSH version string. 
  */
diff --git a/sshaes.c b/sshaes.c
index 4e2ada56..df466f28 100644
--- a/sshaes.c
+++ b/sshaes.c
@@ -48,6 +48,7 @@
 #    define INLINE
 #endif
 
+typedef struct AESContext AESContext;
 struct AESContext {
     uint32_t keysched_buf[(MAX_NR + 1) * NB + 3];
     uint32_t invkeysched_buf[(MAX_NR + 1) * NB + 3];
@@ -1024,20 +1025,24 @@ void aes_ssh2_sdctr(AESContext *ctx, void *blk, int len)
 
 void aes256_encrypt_pubkey(const void *key, void *blk, int len)
 {
-    AESContext ctx;
-    aes_setup(&ctx, key, 32);
-    memset(ctx.iv, 0, sizeof(ctx.iv));
-    aes_encrypt_cbc(blk, len, &ctx);
-    smemclr(&ctx, sizeof(ctx));
+    char iv[16];
+    memset(iv, 0, 16);
+    ssh2_cipher *cipher = ssh2_cipher_new(&ssh_aes256);
+    ssh2_cipher_setkey(cipher, key);
+    ssh2_cipher_setiv(cipher, iv);
+    ssh2_cipher_encrypt(cipher, blk, len);
+    ssh2_cipher_free(cipher);
 }
 
 void aes256_decrypt_pubkey(const void *key, void *blk, int len)
 {
-    AESContext ctx;
-    aes_setup(&ctx, key, 32);
-    memset(ctx.iv, 0, sizeof(ctx.iv));
-    aes_decrypt_cbc(blk, len, &ctx);
-    smemclr(&ctx, sizeof(ctx));
+    char iv[16];
+    memset(iv, 0, 16);
+    ssh2_cipher *cipher = ssh2_cipher_new(&ssh_aes256);
+    ssh2_cipher_setkey(cipher, key);
+    ssh2_cipher_setiv(cipher, iv);
+    ssh2_cipher_decrypt(cipher, blk, len);
+    ssh2_cipher_free(cipher);
 }
 
 struct aes_ssh2_ctx {