Address prime_field_inv(a: int, n: int) a == n case

py_ecc/utils.py

@@ -22,6 +22,12 @@ def prime_field_inv(a: int, n: int) -> int:
     """
     Extended euclidean algorithm to find modular inverses for integers
     """
+    # To address a == n edge case.
+    # https://tools.ietf.org/html/draft-irtf-cfrg-hash-to-curve-09#section-4
+    # inv0(x): This function returns the multiplicative inverse of x in
+    # F, extended to all of F by fixing inv0(0) == 0.
+    a %= n
+
     if a == 0:
         return 0
     lm, hm = 1, 0

tests/test_utils.py

@@ -0,0 +1,16 @@
+import pytest
+
+from py_ecc.utils import prime_field_inv
+
+
+@pytest.mark.parametrize(
+    'a,n,result',
+    [
+        (0, 7, 0),
+        (7, 7, 0),
+        (2, 7, 4),
+        (10, 7, 5),
+    ]
+)
+def test_prime_field_inv(a, n, result):
+    assert prime_field_inv(a, n) % n == result