From 8ebd8aacc2343fae0374d07d66aea16c05abc5e2 Mon Sep 17 00:00:00 2001
From: Jouni Malinen <j@w1.fi>
Date: Fri, 7 Jan 2022 13:47:16 +0200
Subject: [PATCH] SAE: Move sqrt() implementation into a helper function

This implementation within SSWU can be helpful for other users of the
dragonfly design, so move it into a shared helper function.

Signed-off-by: Jouni Malinen <j@w1.fi>
---
 src/common/dragonfly.c | 34 ++++++++++++++++++++++++++++++++++
 src/common/dragonfly.h |  2 ++
 src/common/sae.c       | 14 ++------------
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/src/common/dragonfly.c b/src/common/dragonfly.c
index 547be66f1..1e8427166 100644
--- a/src/common/dragonfly.c
+++ b/src/common/dragonfly.c
@@ -213,3 +213,37 @@ int dragonfly_generate_scalar(const struct crypto_bignum *order,
 		   "dragonfly: Unable to get randomness for own scalar");
 	return -1;
 }
+
+
+/* res = sqrt(val) */
+int dragonfly_sqrt(struct crypto_ec *ec, const struct crypto_bignum *val,
+		   struct crypto_bignum *res)
+{
+	const struct crypto_bignum *prime;
+	struct crypto_bignum *tmp, *one;
+	int ret = 0;
+	u8 prime_bin[DRAGONFLY_MAX_ECC_PRIME_LEN];
+	size_t prime_len;
+
+	/* For prime p such that p = 3 mod 4, sqrt(w) = w^((p+1)/4) mod p */
+
+	prime = crypto_ec_get_prime(ec);
+	prime_len = crypto_ec_prime_len(ec);
+	tmp = crypto_bignum_init();
+	one = crypto_bignum_init_uint(1);
+
+	if (crypto_bignum_to_bin(prime, prime_bin, sizeof(prime_bin),
+				 prime_len) < 0 ||
+	    (prime_bin[prime_len - 1] & 0x03) != 3 ||
+	    !tmp || !one ||
+	    /* tmp = (p+1)/4 */
+	    crypto_bignum_add(prime, one, tmp) < 0 ||
+	    crypto_bignum_rshift(tmp, 2, tmp) < 0 ||
+	    /* res = sqrt(val) */
+	    crypto_bignum_exptmod(val, tmp, prime, res) < 0)
+		ret = -1;
+
+	crypto_bignum_deinit(tmp, 0);
+	crypto_bignum_deinit(one, 0);
+	return ret;
+}
diff --git a/src/common/dragonfly.h b/src/common/dragonfly.h
index ec3dd593e..84d67f575 100644
--- a/src/common/dragonfly.h
+++ b/src/common/dragonfly.h
@@ -27,5 +27,7 @@ int dragonfly_generate_scalar(const struct crypto_bignum *order,
 			      struct crypto_bignum *_rand,
 			      struct crypto_bignum *_mask,
 			      struct crypto_bignum *scalar);
+int dragonfly_sqrt(struct crypto_ec *ec, const struct crypto_bignum *val,
+		   struct crypto_bignum *res);
 
 #endif /* DRAGONFLY_H */
diff --git a/src/common/sae.c b/src/common/sae.c
index b768c22fa..65e99d176 100644
--- a/src/common/sae.c
+++ b/src/common/sae.c
@@ -747,19 +747,9 @@ static struct crypto_ec_point * sswu(struct crypto_ec *ec, int group,
 	const_time_select_bin(is_qr, bin1, bin2, prime_len, x_y);
 	wpa_hexdump_key(MSG_DEBUG, "SSWU: x = CSEL(l, x1, x2)", x_y, prime_len);
 
-	/* y = sqrt(v)
-	 * For prime p such that p = 3 mod 4 --> v^((p+1)/4) */
-	if (crypto_bignum_to_bin(prime, bin1, sizeof(bin1), prime_len) < 0)
-		goto fail;
-	if ((bin1[prime_len - 1] & 0x03) != 3) {
-		wpa_printf(MSG_DEBUG, "SSWU: prime does not have p = 3 mod 4");
-		goto fail;
-	}
+	/* y = sqrt(v) */
 	y = crypto_bignum_init();
-	if (!y ||
-	    crypto_bignum_add(prime, one, t1) < 0 ||
-	    crypto_bignum_rshift(t1, 2, t1) < 0 ||
-	    crypto_bignum_exptmod(v, t1, prime, y) < 0)
+	if (!y || dragonfly_sqrt(ec, v, y) < 0)
 		goto fail;
 	debug_print_bignum("SSWU: y = sqrt(v)", y, prime_len);