[Feature] Add aes-rng PRF to libottery

7 files changed, 208 insertions, 8 deletions

diff --git a/contrib/libottery/CMakeLists.txt b/contrib/libottery/CMakeLists.txt
index 7e504605b..1e9bca143 100644
--- a/contrib/libottery/CMakeLists.txt
+++ b/contrib/libottery/CMakeLists.txt
@@ -2,7 +2,9 @@ SET(OTTERYSRC            chacha_merged.c
                          ottery.c
                          ottery_cpuinfo.c
                          ottery_entropy.c
-                         ottery_global.c chacha_cryptobox.c)
+                         ottery_global.c
+                         chacha_cryptobox.c
+                         aes_cryptobox.c)
 ADD_LIBRARY(ottery STATIC ${OTTERYSRC})
 
 SET(OTTERY_CFLAGS "-DBUILD_RSPAMD")
diff --git a/contrib/libottery/aes_cryptobox.c b/contrib/libottery/aes_cryptobox.c
new file mode 100644
index 000000000..0e5434442
--- /dev/null
+++ b/contrib/libottery/aes_cryptobox.c
@@ -0,0 +1,179 @@
+/*
+ * Copyright (c) 2017, Vsevolod Stakhov
+ * Copyright (c) 2017, Frank Denis
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *	 * Redistributions of source code must retain the above copyright
+ *	   notice, this list of conditions and the following disclaimer.
+ *	 * Redistributions in binary form must reproduce the above copyright
+ *	   notice, this list of conditions and the following disclaimer in the
+ *	   documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY AUTHOR ''AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL AUTHOR BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#include "config.h"
+#include "ottery-internal.h"
+#include "cryptobox.h"
+
+#if defined(__x86_64__) && defined(RSPAMD_HAS_TARGET_ATTR)
+#pragma GCC push_options
+#pragma GCC target("aes")
+#ifndef __SSE2__
+#define __SSE2__
+#endif
+#ifndef __SSE__
+#define __SSE__
+#endif
+#ifndef __AES__
+#define __AES__
+#endif
+#include <immintrin.h>
+#define ROUNDS 10
+
+typedef struct RSPAMD_ALIGNED(16) aes_rng_state {
+	__m128i round_keys[ROUNDS + 1];
+	__m128i counter;
+} aes_stream_state;
+
+
+#define STATE_LEN   sizeof(aes_stream_state)
+#define STATE_BYTES 16
+
+#define OUTPUT_LEN  1024
+
+static void
+aes_key_expand (__m128i round_keys[ROUNDS + 1], __m128i t) __attribute__((target("aes")));
+
+static void
+aes_key_expand (__m128i round_keys[ROUNDS + 1], __m128i t)
+{
+	__m128i t1;
+
+#define DO_ROUND_KEY(ROUND, RC)                           \
+    do {                                                   \
+        t1 = _mm_aeskeygenassist_si128(t, (RC));           \
+        round_keys[ROUND] = t;                             \
+        t = _mm_xor_si128(t, _mm_slli_si128(t, 4));        \
+        t = _mm_xor_si128(t, _mm_slli_si128(t, 8));        \
+        t = _mm_xor_si128(t, _mm_shuffle_epi32(t1, 0xff)); \
+    } while (0)
+
+	DO_ROUND_KEY(0, 1);
+	DO_ROUND_KEY(1, 2);
+	DO_ROUND_KEY(2, 4);
+	DO_ROUND_KEY(3, 8);
+	DO_ROUND_KEY(4, 16);
+	DO_ROUND_KEY(5, 32);
+	DO_ROUND_KEY(6, 64);
+	DO_ROUND_KEY(7, 128);
+	DO_ROUND_KEY(8, 27);
+	DO_ROUND_KEY(9, 54);
+	round_keys[10] = t;
+}
+
+/*
+ * Computes one 128 bytes block and refresh keys
+ */
+static void
+aes_round(unsigned char *buf, struct aes_rng_state *st) __attribute__((target("aes")));
+static void
+aes_round(unsigned char *buf, struct aes_rng_state *st)
+{
+	const __m128i  one = _mm_set_epi64x(0, 1);
+	__m128i *round_keys = st->round_keys;
+	__m128i c0, c1, c2, c3, c4, c5, c6, c7;
+	__m128i r0, r1, r2, r3, r4, r5, r6, r7;
+	__m128i s0, s1, s2, s3, s4, s5, s6, s7;
+	size_t i;
+
+#define COMPUTE_ROUNDS(N)                                                              \
+    do {                                                                               \
+        r##N = _mm_aesenc_si128(   _mm_xor_si128(c##N, round_keys[0]), round_keys[1]); \
+        r##N = _mm_aesenc_si128(_mm_aesenc_si128(r##N, round_keys[2]), round_keys[3]); \
+        r##N = _mm_aesenc_si128(_mm_aesenc_si128(r##N, round_keys[4]), round_keys[5]); \
+        s##N = r##N;                                                                   \
+        r##N = _mm_aesenc_si128(_mm_aesenc_si128(r##N, round_keys[6]), round_keys[7]); \
+        r##N = _mm_aesenc_si128(_mm_aesenc_si128(r##N, round_keys[8]), round_keys[9]); \
+        r##N = _mm_xor_si128(s##N, _mm_aesenclast_si128(r##N, round_keys[10]));        \
+    } while (0)
+
+	c0 = st->counter;
+
+	for (i = 0; i < OUTPUT_LEN / 128; i ++) {
+		c1 = _mm_add_epi64 (c0, one);
+		c2 = _mm_add_epi64 (c1, one);
+		c3 = _mm_add_epi64 (c2, one);
+		c4 = _mm_add_epi64 (c3, one);
+		c5 = _mm_add_epi64 (c4, one);
+		c6 = _mm_add_epi64 (c5, one);
+		c7 = _mm_add_epi64 (c6, one);
+		COMPUTE_ROUNDS(0);
+		COMPUTE_ROUNDS(1);
+		COMPUTE_ROUNDS(2);
+		COMPUTE_ROUNDS(3);
+		COMPUTE_ROUNDS(4);
+		COMPUTE_ROUNDS(5);
+		COMPUTE_ROUNDS(6);
+		COMPUTE_ROUNDS(7);
+		c0 = _mm_add_epi64 (c7, one);
+		_mm_storeu_si128 ((__m128i *) (void *) (buf + 0), r0);
+		_mm_storeu_si128 ((__m128i *) (void *) (buf + 16), r1);
+		_mm_storeu_si128 ((__m128i *) (void *) (buf + 32), r2);
+		_mm_storeu_si128 ((__m128i *) (void *) (buf + 48), r3);
+		_mm_storeu_si128 ((__m128i *) (void *) (buf + 64), r4);
+		_mm_storeu_si128 ((__m128i *) (void *) (buf + 80), r5);
+		_mm_storeu_si128 ((__m128i *) (void *) (buf + 96), r6);
+		_mm_storeu_si128 ((__m128i *) (void *) (buf + 112), r7);
+		buf += 128;
+	}
+
+	st->counter = c0;
+	c0 = _mm_setzero_si128();
+	COMPUTE_ROUNDS(0);
+	aes_key_expand(round_keys, r0);
+}
+
+
+static void
+aes_cryptobox_state_setup (void *state_, const uint8_t *bytes)
+{
+	struct aes_rng_state *x = state_;
+
+	aes_key_expand (x->round_keys,
+			_mm_loadu_si128((const __m128i *) (const void *)bytes));
+}
+
+static void
+aes_cryptobox_generate (void *state_, uint8_t *output, uint32_t idx)
+{
+	struct aes_rng_state *x = state_;
+
+	aes_round(output, x);
+}
+
+#define PRF_AES(r) {                         \
+  "AES-" #r,                                  \
+  "AES-" #r "-NOSIMD",                        \
+  "AES-" #r "-NOSIMD-DEFAULT",                \
+  STATE_LEN,                                    \
+  STATE_BYTES,                                  \
+  OUTPUT_LEN,                                   \
+  OTTERY_CPUCAP_AES,                                            \
+  aes_cryptobox_state_setup,                    \
+  aes_cryptobox_generate               \
+}
+
+const struct ottery_prf ottery_prf_aes_cryptobox_ = PRF_AES(128);
+#endif /* x86_64 */
diff --git a/contrib/libottery/chacha_cryptobox.c b/contrib/libottery/chacha_cryptobox.c
index ecbcccaf3..4e9cdaef2 100644
--- a/contrib/libottery/chacha_cryptobox.c
+++ b/contrib/libottery/chacha_cryptobox.c
@@ -33,7 +33,7 @@
 #define OUTPUT_LEN  (IDX_STEP * 64)
 
 static void
-chacha_cryptobox_state_setup (void *state_, const uint8_t *bytes)
+chacha20_cryptobox_state_setup (void *state_, const uint8_t *bytes)
 {
 	chacha_state *x = state_;
 	chacha_init (x, (chacha_key *)bytes, (chacha_iv *)(bytes + 32), 20);
@@ -50,15 +50,15 @@ chacha20_cryptobox_generate (void *state_, uint8_t *output, uint32_t idx)
 }
 
 #define PRF_CHACHA(r) {                         \
-  "CHACHA" #r,                                  \
-  "CHACHA" #r "-NOSIMD",                        \
-  "CHACHA" #r "-NOSIMD-DEFAULT",                \
+  "CHACHA" #r "-CRYPTOBOX",                    \
+  "CHACHA" #r "-CRYPTOBOX",                    \
+  "CHACHA" #r "-CRYPTOBOX",                    \
   STATE_LEN,                                    \
   STATE_BYTES,                                  \
   OUTPUT_LEN,                                   \
-  0,                                            \
-  chacha_cryptobox_state_setup,                    \
-  chacha ## r ## _cryptobox_generate               \
+  0,                                             \
+  chacha ## r ## _cryptobox_state_setup,         \
+  chacha ## r ## _cryptobox_generate             \
 }
 
 const struct ottery_prf ottery_prf_chacha20_cryptobox_ = PRF_CHACHA(20);
diff --git a/contrib/libottery/ottery-internal.h b/contrib/libottery/ottery-internal.h
index 2c38a52b0..cc047f810 100644
--- a/contrib/libottery/ottery-internal.h
+++ b/contrib/libottery/ottery-internal.h
@@ -306,6 +306,9 @@ extern const struct ottery_prf ottery_prf_chacha12_merged_;
 extern const struct ottery_prf ottery_prf_chacha20_merged_;
 
 #ifdef BUILD_RSPAMD
+#ifdef __x86_64__
+extern const struct ottery_prf ottery_prf_aes_cryptobox_;
+#endif
 extern const struct ottery_prf ottery_prf_chacha20_cryptobox_;
 #endif
 /**@}*/
diff --git a/contrib/libottery/ottery.c b/contrib/libottery/ottery.c
index 0b466c641..c58a90173 100644
--- a/contrib/libottery/ottery.c
+++ b/contrib/libottery/ottery.c
@@ -33,6 +33,10 @@
 #define OTTERY_NO_PID_CHECK
 #endif
 
+#ifdef BUILD_RSPAMD
+#include "cryptobox.h"
+#endif
+
 /** Magic number for deciding whether an ottery_state is initialized. */
 #define MAGIC_BASIS 0x11b07734
 
@@ -183,7 +187,11 @@ ottery_get_impl(const char *impl)
     &ottery_prf_chacha12_krovetz_1_,
     &ottery_prf_chacha8_krovetz_1_,
 #endif
+
 #ifdef BUILD_RSPAMD
+#if defined(__x86_64__) && defined(RSPAMD_HAS_TARGET_ATTR)
+    &ottery_prf_aes_cryptobox_,
+#endif
     &ottery_prf_chacha20_cryptobox_,
 #endif
 	&ottery_prf_chacha20_merged_,
diff --git a/contrib/libottery/ottery_common.h b/contrib/libottery/ottery_common.h
index 5b3359c97..bac6f04ca 100644
--- a/contrib/libottery/ottery_common.h
+++ b/contrib/libottery/ottery_common.h
@@ -346,4 +346,6 @@ uint32_t ottery_get_version(void);
  */
 const char *ottery_get_version_string(void);
 
+const char *ottery_get_impl_name(void);
+
 #endif
diff --git a/contrib/libottery/ottery_global.c b/contrib/libottery/ottery_global.c
index 788b5076c..dd1efc55a 100644
--- a/contrib/libottery/ottery_global.c
+++ b/contrib/libottery/ottery_global.c
@@ -108,3 +108,9 @@ ottery_rand_range64(uint64_t top)
   CHECK_INIT(0);
   return ottery_st_rand_range64(&ottery_global_state_, top);
 }
+
+const char *ottery_get_impl_name(void)
+{
+	CHECK_INIT(0);
+	return ottery_global_state_.prf.name;
+}
\ No newline at end of file