aes-neon.c 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360
  1. /* ----------------------------------------------------------------------
  2. * Hardware-accelerated implementation of AES using Arm NEON.
  3. */
  4. #include "ssh.h"
  5. #include "aes.h"
  6. #if USE_ARM64_NEON_H
  7. #include <arm64_neon.h>
  8. #else
  9. #include <arm_neon.h>
  10. #endif
  11. static bool aes_neon_available(void)
  12. {
  13. /*
  14. * For Arm, we delegate to a per-platform AES detection function,
  15. * because it has to be implemented by asking the operating system
  16. * rather than directly querying the CPU.
  17. *
  18. * That's because Arm systems commonly have multiple cores that
  19. * are not all alike, so any method of querying whether NEON
  20. * crypto instructions work on the _current_ CPU - even one as
  21. * crude as just trying one and catching the SIGILL - wouldn't
  22. * give an answer that you could still rely on the first time the
  23. * OS migrated your process to another CPU.
  24. */
  25. return platform_aes_neon_available();
  26. }
  27. /*
  28. * Core NEON encrypt/decrypt functions, one per length and direction.
  29. */
  30. #define NEON_CIPHER(len, repmacro) \
  31. static inline uint8x16_t aes_neon_##len##_e( \
  32. uint8x16_t v, const uint8x16_t *keysched) \
  33. { \
  34. repmacro(v = vaesmcq_u8(vaeseq_u8(v, *keysched++));); \
  35. v = vaeseq_u8(v, *keysched++); \
  36. return veorq_u8(v, *keysched); \
  37. } \
  38. static inline uint8x16_t aes_neon_##len##_d( \
  39. uint8x16_t v, const uint8x16_t *keysched) \
  40. { \
  41. repmacro(v = vaesimcq_u8(vaesdq_u8(v, *keysched++));); \
  42. v = vaesdq_u8(v, *keysched++); \
  43. return veorq_u8(v, *keysched); \
  44. }
  45. NEON_CIPHER(128, REP9)
  46. NEON_CIPHER(192, REP11)
  47. NEON_CIPHER(256, REP13)
  48. /*
  49. * The main key expansion.
  50. */
  51. static void aes_neon_key_expand(
  52. const unsigned char *key, size_t key_words,
  53. uint8x16_t *keysched_e, uint8x16_t *keysched_d)
  54. {
  55. size_t rounds = key_words + 6;
  56. size_t sched_words = (rounds + 1) * 4;
  57. /*
  58. * Store the key schedule as 32-bit integers during expansion, so
  59. * that it's easy to refer back to individual previous words. We
  60. * collect them into the final uint8x16_t form at the end.
  61. */
  62. uint32_t sched[MAXROUNDKEYS * 4];
  63. unsigned rconpos = 0;
  64. for (size_t i = 0; i < sched_words; i++) {
  65. if (i < key_words) {
  66. sched[i] = GET_32BIT_LSB_FIRST(key + 4 * i);
  67. } else {
  68. uint32_t temp = sched[i - 1];
  69. bool rotate_and_round_constant = (i % key_words == 0);
  70. bool sub = rotate_and_round_constant ||
  71. (key_words == 8 && i % 8 == 4);
  72. if (rotate_and_round_constant)
  73. temp = (temp << 24) | (temp >> 8);
  74. if (sub) {
  75. uint32x4_t v32 = vdupq_n_u32(temp);
  76. uint8x16_t v8 = vreinterpretq_u8_u32(v32);
  77. v8 = vaeseq_u8(v8, vdupq_n_u8(0));
  78. v32 = vreinterpretq_u32_u8(v8);
  79. temp = vget_lane_u32(vget_low_u32(v32), 0);
  80. }
  81. if (rotate_and_round_constant) {
  82. assert(rconpos < lenof(aes_key_setup_round_constants));
  83. temp ^= aes_key_setup_round_constants[rconpos++];
  84. }
  85. sched[i] = sched[i - key_words] ^ temp;
  86. }
  87. }
  88. /*
  89. * Combine the key schedule words into uint8x16_t vectors and
  90. * store them in the output context.
  91. */
  92. for (size_t round = 0; round <= rounds; round++)
  93. keysched_e[round] = vreinterpretq_u8_u32(vld1q_u32(sched + 4*round));
  94. smemclr(sched, sizeof(sched));
  95. /*
  96. * Now prepare the modified keys for the inverse cipher.
  97. */
  98. for (size_t eround = 0; eround <= rounds; eround++) {
  99. size_t dround = rounds - eround;
  100. uint8x16_t rkey = keysched_e[eround];
  101. if (eround && dround) /* neither first nor last */
  102. rkey = vaesimcq_u8(rkey);
  103. keysched_d[dround] = rkey;
  104. }
  105. }
  106. /*
  107. * Auxiliary routine to reverse the byte order of a vector, so that
  108. * the SDCTR IV can be made big-endian for feeding to the cipher.
  109. *
  110. * In fact we don't need to reverse the vector _all_ the way; we leave
  111. * the two lanes in MSW,LSW order, because that makes no difference to
  112. * the efficiency of the increment. That way we only have to reverse
  113. * bytes within each lane in this function.
  114. */
  115. static inline uint8x16_t aes_neon_sdctr_reverse(uint8x16_t v)
  116. {
  117. return vrev64q_u8(v);
  118. }
  119. /*
  120. * Auxiliary routine to increment the 128-bit counter used in SDCTR
  121. * mode. There's no instruction to treat a 128-bit vector as a single
  122. * long integer, so instead we have to increment the bottom half
  123. * unconditionally, and the top half if the bottom half started off as
  124. * all 1s (in which case there was about to be a carry).
  125. */
  126. static inline uint8x16_t aes_neon_sdctr_increment(uint8x16_t in)
  127. {
  128. #ifdef __aarch64__
  129. /* There will be a carry if the low 64 bits are all 1s. */
  130. uint64x1_t all1 = vcreate_u64(0xFFFFFFFFFFFFFFFF);
  131. uint64x1_t carry = vceq_u64(vget_high_u64(vreinterpretq_u64_u8(in)), all1);
  132. /* Make a word whose bottom half is unconditionally all 1s, and
  133. * the top half is 'carry', i.e. all 0s most of the time but all
  134. * 1s if we need to increment the top half. Then that word is what
  135. * we need to _subtract_ from the input counter. */
  136. uint64x2_t subtrahend = vcombine_u64(carry, all1);
  137. #else
  138. /* AArch32 doesn't have comparisons that operate on a 64-bit lane,
  139. * so we start by comparing each 32-bit half of the low 64 bits
  140. * _separately_ to all-1s. */
  141. uint32x2_t all1 = vdup_n_u32(0xFFFFFFFF);
  142. uint32x2_t carry = vceq_u32(
  143. vget_high_u32(vreinterpretq_u32_u8(in)), all1);
  144. /* Swap the 32-bit words of the compare output, and AND with the
  145. * unswapped version. Now carry is all 1s iff the bottom half of
  146. * the input counter was all 1s, and all 0s otherwise. */
  147. carry = vand_u32(carry, vrev64_u32(carry));
  148. /* Now make the vector to subtract in the same way as above. */
  149. uint64x2_t subtrahend = vreinterpretq_u64_u32(vcombine_u32(carry, all1));
  150. #endif
  151. return vreinterpretq_u8_u64(
  152. vsubq_u64(vreinterpretq_u64_u8(in), subtrahend));
  153. }
  154. /*
  155. * Much simpler auxiliary routine to increment the counter for GCM
  156. * mode. This only has to increment the low word.
  157. */
  158. static inline uint8x16_t aes_neon_gcm_increment(uint8x16_t in)
  159. {
  160. uint32x4_t inw = vreinterpretq_u32_u8(in);
  161. uint32x4_t ONE = vcombine_u32(vcreate_u32(0), vcreate_u32(1));
  162. inw = vaddq_u32(inw, ONE);
  163. return vreinterpretq_u8_u32(inw);
  164. }
  165. /*
  166. * The SSH interface and the cipher modes.
  167. */
  168. typedef struct aes_neon_context aes_neon_context;
  169. struct aes_neon_context {
  170. uint8x16_t keysched_e[MAXROUNDKEYS], keysched_d[MAXROUNDKEYS], iv;
  171. ssh_cipher ciph;
  172. };
  173. static ssh_cipher *aes_neon_new(const ssh_cipheralg *alg)
  174. {
  175. const struct aes_extra *extra = (const struct aes_extra *)alg->extra;
  176. if (!check_availability(extra))
  177. return NULL;
  178. aes_neon_context *ctx = snew(aes_neon_context);
  179. ctx->ciph.vt = alg;
  180. return &ctx->ciph;
  181. }
  182. static void aes_neon_free(ssh_cipher *ciph)
  183. {
  184. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  185. smemclr(ctx, sizeof(*ctx));
  186. sfree(ctx);
  187. }
  188. static void aes_neon_setkey(ssh_cipher *ciph, const void *vkey)
  189. {
  190. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  191. const unsigned char *key = (const unsigned char *)vkey;
  192. aes_neon_key_expand(key, ctx->ciph.vt->real_keybits / 32,
  193. ctx->keysched_e, ctx->keysched_d);
  194. }
  195. static void aes_neon_setiv_cbc(ssh_cipher *ciph, const void *iv)
  196. {
  197. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  198. ctx->iv = vld1q_u8(iv);
  199. }
  200. static void aes_neon_setiv_sdctr(ssh_cipher *ciph, const void *iv)
  201. {
  202. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  203. uint8x16_t counter = vld1q_u8(iv);
  204. ctx->iv = aes_neon_sdctr_reverse(counter);
  205. }
  206. static void aes_neon_setiv_gcm(ssh_cipher *ciph, const void *iv)
  207. {
  208. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  209. uint8x16_t counter = vld1q_u8(iv);
  210. ctx->iv = aes_neon_sdctr_reverse(counter);
  211. ctx->iv = vreinterpretq_u8_u32(vsetq_lane_u32(
  212. 1, vreinterpretq_u32_u8(ctx->iv), 2));
  213. }
  214. static void aes_neon_next_message_gcm(ssh_cipher *ciph)
  215. {
  216. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  217. uint32x4_t iv = vreinterpretq_u32_u8(ctx->iv);
  218. uint64_t msg_counter = vgetq_lane_u32(iv, 0);
  219. msg_counter = (msg_counter << 32) | vgetq_lane_u32(iv, 3);
  220. msg_counter++;
  221. iv = vsetq_lane_u32(msg_counter >> 32, iv, 0);
  222. iv = vsetq_lane_u32(msg_counter, iv, 3);
  223. iv = vsetq_lane_u32(1, iv, 2);
  224. ctx->iv = vreinterpretq_u8_u32(iv);
  225. }
  226. typedef uint8x16_t (*aes_neon_fn)(uint8x16_t v, const uint8x16_t *keysched);
  227. static inline void aes_cbc_neon_encrypt(
  228. ssh_cipher *ciph, void *vblk, int blklen, aes_neon_fn encrypt)
  229. {
  230. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  231. for (uint8_t *blk = (uint8_t *)vblk, *finish = blk + blklen;
  232. blk < finish; blk += 16) {
  233. uint8x16_t plaintext = vld1q_u8(blk);
  234. uint8x16_t cipher_input = veorq_u8(plaintext, ctx->iv);
  235. uint8x16_t ciphertext = encrypt(cipher_input, ctx->keysched_e);
  236. vst1q_u8(blk, ciphertext);
  237. ctx->iv = ciphertext;
  238. }
  239. }
  240. static inline void aes_cbc_neon_decrypt(
  241. ssh_cipher *ciph, void *vblk, int blklen, aes_neon_fn decrypt)
  242. {
  243. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  244. for (uint8_t *blk = (uint8_t *)vblk, *finish = blk + blklen;
  245. blk < finish; blk += 16) {
  246. uint8x16_t ciphertext = vld1q_u8(blk);
  247. uint8x16_t decrypted = decrypt(ciphertext, ctx->keysched_d);
  248. uint8x16_t plaintext = veorq_u8(decrypted, ctx->iv);
  249. vst1q_u8(blk, plaintext);
  250. ctx->iv = ciphertext;
  251. }
  252. }
  253. static inline void aes_sdctr_neon(
  254. ssh_cipher *ciph, void *vblk, int blklen, aes_neon_fn encrypt)
  255. {
  256. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  257. for (uint8_t *blk = (uint8_t *)vblk, *finish = blk + blklen;
  258. blk < finish; blk += 16) {
  259. uint8x16_t counter = aes_neon_sdctr_reverse(ctx->iv);
  260. uint8x16_t keystream = encrypt(counter, ctx->keysched_e);
  261. uint8x16_t input = vld1q_u8(blk);
  262. uint8x16_t output = veorq_u8(input, keystream);
  263. vst1q_u8(blk, output);
  264. ctx->iv = aes_neon_sdctr_increment(ctx->iv);
  265. }
  266. }
  267. static inline void aes_encrypt_ecb_block_neon(
  268. ssh_cipher *ciph, void *blk, aes_neon_fn encrypt)
  269. {
  270. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  271. uint8x16_t plaintext = vld1q_u8(blk);
  272. uint8x16_t ciphertext = encrypt(plaintext, ctx->keysched_e);
  273. vst1q_u8(blk, ciphertext);
  274. }
  275. static inline void aes_gcm_neon(
  276. ssh_cipher *ciph, void *vblk, int blklen, aes_neon_fn encrypt)
  277. {
  278. aes_neon_context *ctx = container_of(ciph, aes_neon_context, ciph);
  279. for (uint8_t *blk = (uint8_t *)vblk, *finish = blk + blklen;
  280. blk < finish; blk += 16) {
  281. uint8x16_t counter = aes_neon_sdctr_reverse(ctx->iv);
  282. uint8x16_t keystream = encrypt(counter, ctx->keysched_e);
  283. uint8x16_t input = vld1q_u8(blk);
  284. uint8x16_t output = veorq_u8(input, keystream);
  285. vst1q_u8(blk, output);
  286. ctx->iv = aes_neon_gcm_increment(ctx->iv);
  287. }
  288. }
  289. #define NEON_ENC_DEC(len) \
  290. static void aes##len##_neon_cbc_encrypt( \
  291. ssh_cipher *ciph, void *vblk, int blklen) \
  292. { aes_cbc_neon_encrypt(ciph, vblk, blklen, aes_neon_##len##_e); } \
  293. static void aes##len##_neon_cbc_decrypt( \
  294. ssh_cipher *ciph, void *vblk, int blklen) \
  295. { aes_cbc_neon_decrypt(ciph, vblk, blklen, aes_neon_##len##_d); } \
  296. static void aes##len##_neon_sdctr( \
  297. ssh_cipher *ciph, void *vblk, int blklen) \
  298. { aes_sdctr_neon(ciph, vblk, blklen, aes_neon_##len##_e); } \
  299. static void aes##len##_neon_gcm( \
  300. ssh_cipher *ciph, void *vblk, int blklen) \
  301. { aes_gcm_neon(ciph, vblk, blklen, aes_neon_##len##_e); } \
  302. static void aes##len##_neon_encrypt_ecb_block( \
  303. ssh_cipher *ciph, void *vblk) \
  304. { aes_encrypt_ecb_block_neon(ciph, vblk, aes_neon_##len##_e); }
  305. NEON_ENC_DEC(128)
  306. NEON_ENC_DEC(192)
  307. NEON_ENC_DEC(256)
  308. AES_EXTRA(_neon);
  309. AES_ALL_VTABLES(_neon, "NEON accelerated");