From 7dc6fa1ac982e1b4960c54d012005c7ebcf124a5 Mon Sep 17 00:00:00 2001
From: alivecow <mail@ta-lotz.de>
Date: Fri, 15 Nov 2024 20:14:54 +0100
Subject: [PATCH] fix: Fix handling of special cases in powmod

---
 src/utils/field.rs | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/src/utils/field.rs b/src/utils/field.rs
index 859978b..37502f6 100644
--- a/src/utils/field.rs
+++ b/src/utils/field.rs
@@ -75,6 +75,18 @@ impl Polynomial {
             polynomial_2_block(vec![0], "gcm").unwrap(),
         )]);
 
+        if exponent == 1 {
+            return self;
+        }
+
+        if exponent == 0 {
+            Polynomial::new(vec![FieldElement::new(
+                polynomial_2_block(vec![1], "gcm").unwrap(),
+            )])
+            .div(&modulus)
+            .1;
+        }
+
         //eprintln!("Initial result: {:?}", result);
         while exponent > 0 {
             //eprintln!("Current exponent: {:02X}", exponent);
@@ -1008,4 +1020,30 @@ mod tests {
         eprintln!("Result is: {:02X?}", result);
         assert_eq!(result.to_c_array(), vec!["oNXl5P8xq2WpUTP92u25zg=="]);
     }
+
+    #[test]
+    fn test_field_poly_powmod_k1() {
+        let json1 = json!(["JAAAAAAAAAAAAAAAAAAAAA==",]);
+        let json2 = json!(["KryptoanalyseAAAAAAAAA=="]);
+        let element1: Polynomial = Polynomial::from_c_array(&json1);
+        let modulus: Polynomial = Polynomial::from_c_array(&json2);
+
+        let result = element1.pow_mod(1, modulus);
+
+        eprintln!("Result is: {:02X?}", result);
+        assert_eq!(result.to_c_array(), vec!["JAAAAAAAAAAAAAAAAAAAAA=="]);
+    }
+
+    #[test]
+    fn test_field_poly_powmod_k0() {
+        let json1 = json!(["JAAAAAAAAAAAAAAAAAAAAA==",]);
+        let json2 = json!(["KryptoanalyseAAAAAAAAA=="]);
+        let element1: Polynomial = Polynomial::from_c_array(&json1);
+        let modulus: Polynomial = Polynomial::from_c_array(&json2);
+
+        let result = element1.pow_mod(0, modulus);
+
+        eprintln!("Result is: {:02X?}", result);
+        assert_eq!(result.to_c_array(), vec!["gAAAAAAAAAAAAAAAAAAAAA=="]);
+    }
 }