Add comments to Python ECDSA implementation

1 files changed, 56 insertions, 16 deletions

diff --git a/test/functional/test_framework/key.py b/test/functional/test_framework/key.py
index 996de5f94f..912c0ca978 100644
--- a/test/functional/test_framework/key.py
+++ b/test/functional/test_framework/key.py
@@ -1,18 +1,17 @@
 # Copyright (c) 2019 Pieter Wuille
-
+# Distributed under the MIT software license, see the accompanying
+# file COPYING or http://www.opensource.org/licenses/mit-license.php.
 """Test-only secp256k1 elliptic curve implementation
 
 WARNING: This code is slow, uses bad randomness, does not properly protect
 keys, and is trivially vulnerable to side channel attacks. Do not use for
-anything but tests.
-"""
-
+anything but tests."""
 import random
 
 def modinv(a, n):
     """Compute the modular inverse of a modulo n
 
-    See https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm#Modular_integers
+    See https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm#Modular_integers.
     """
     t1, t2 = 0, 1
     r1, r2 = n, a
@@ -30,8 +29,9 @@ def jacobi_symbol(n, k):
     """Compute the Jacobi symbol of n modulo k
 
     See http://en.wikipedia.org/wiki/Jacobi_symbol
-    """
-    assert k > 0 and k & 1
+
+    For our application k is always prime, so this is the same as the Legendre symbol."""
+    assert k > 0 and k & 1, "jacobi symbol is only defined for positive odd k"
     n %= k
     t = 0
     while n != 0:
@@ -47,11 +47,18 @@ def jacobi_symbol(n, k):
     return 0
 
 def modsqrt(a, p):
-    """Compute the square root of a modulo p
+    """Compute the square root of a modulo p when p % 4 = 3.
 
-    For p = 3 mod 4, if a square root exists, it is equal to a**((p+1)/4) mod p.
+    The Tonelli-Shanks algorithm can be used. See https://en.wikipedia.org/wiki/Tonelli-Shanks_algorithm
+
+    Limiting this function to only work for p % 4 = 3 means we don't need to
+    iterate through the loop. The highest n such that p - 1 = 2^n Q with Q odd
+    is n = 1. Therefore Q = (p-1)/2 and sqrt = a^((Q+1)/2) = a^((p+1)/4)
+
+    secp256k1's is defined over field of size 2**256 - 2**32 - 977, which is 3 mod 4.
     """
-    assert(p % 4 == 3) # Only p = 3 mod 4 is implemented
+    if p % 4 != 3:
+        raise NotImplementedError("modsqrt only implemented for p % 4 = 3")
     sqrt = pow(a, (p + 1)//4, p)
     if pow(sqrt, 2, p) == a % p:
         return sqrt
@@ -65,7 +72,9 @@ class EllipticCurve:
         self.b = b % p
 
     def affine(self, p1):
-        """Convert a Jacobian point tuple p1 to affine form, or None if at infinity."""
+        """Convert a Jacobian point tuple p1 to affine form, or None if at infinity.
+
+        An affine point is represented as the Jacobian (x, y, 1)"""
         x1, y1, z1 = p1
         if z1 == 0:
             return None
@@ -101,7 +110,9 @@ class EllipticCurve:
         return (x, y, 1)
 
     def double(self, p1):
-        """Double a Jacobian tuple p1"""
+        """Double a Jacobian tuple p1
+
+        See https://en.wikibooks.org/wiki/Cryptography/Prime_Curve/Jacobian_Coordinates - Point Doubling"""
         x1, y1, z1 = p1
         if z1 == 0:
             return (0, 1, 0)
@@ -119,10 +130,13 @@ class EllipticCurve:
         return (x2, y2, z2)
 
     def add_mixed(self, p1, p2):
-        """Add a Jacobian tuple p1 and an affine tuple p2"""
+        """Add a Jacobian tuple p1 and an affine tuple p2
+
+        See https://en.wikibooks.org/wiki/Cryptography/Prime_Curve/Jacobian_Coordinates - Point Addition (with affine point)"""
         x1, y1, z1 = p1
         x2, y2, z2 = p2
         assert(z2 == 1)
+        # Adding to the point at infinity is a no-op
         if z1 == 0:
             return p2
         z1_2 = (z1**2) % self.p
@@ -131,7 +145,9 @@ class EllipticCurve:
         s2 = (y2 * z1_3) % self.p
         if x1 == u2:
             if (y1 != s2):
+                # p1 and p2 are inverses. Return the point at infinity.
                 return (0, 1, 0)
+            # p1 == p2. The formulas below fail when the two points are equal.
             return self.double(p1)
         h = u2 - x1
         r = s2 - y1
@@ -144,13 +160,17 @@ class EllipticCurve:
         return (x3, y3, z3)
 
     def add(self, p1, p2):
-        """Add two Jacobian tuples p1 and p2"""
+        """Add two Jacobian tuples p1 and p2
+
+        See https://en.wikibooks.org/wiki/Cryptography/Prime_Curve/Jacobian_Coordinates - Point Addition"""
         x1, y1, z1 = p1
         x2, y2, z2 = p2
+        # Adding the point at infinity is a no-op
         if z1 == 0:
             return p2
         if z2 == 0:
             return p1
+        # Adding an Affine to a Jacobian is more efficient since we save field multiplications and squarings when z = 1
         if z1 == 1:
             return self.add_mixed(p2, p1)
         if z2 == 1:
@@ -165,7 +185,9 @@ class EllipticCurve:
         s2 = (y2 * z1_3) % self.p
         if u1 == u2:
             if (s1 != s2):
+                # p1 and p2 are inverses. Return the point at infinity.
                 return (0, 1, 0)
+            # p1 == p2. The formulas below fail when the two points are equal.
             return self.double(p1)
         h = u2 - u1
         r = s2 - s1
@@ -214,6 +236,8 @@ class ECPubKey():
             x = int.from_bytes(data[1:33], 'big')
             if SECP256K1.is_x_coord(x):
                 p = SECP256K1.lift_x(x)
+                # if the oddness of the y co-ord isn't correct, find the other
+                # valid y
                 if (p[1] & 1) != (data[0] & 1):
                     p = SECP256K1.negate(p)
                 self.p = p
@@ -243,8 +267,14 @@ class ECPubKey():
             return bytes([0x04]) + p[0].to_bytes(32, 'big') + p[1].to_bytes(32, 'big')
 
     def verify_ecdsa(self, sig, msg, low_s=True):
-        """Verify a strictly DER-encoded ECDSA signature against this pubkey."""
+        """Verify a strictly DER-encoded ECDSA signature against this pubkey.
+
+        See https://en.wikipedia.org/wiki/Elliptic_Curve_Digital_Signature_Algorithm for the
+        ECDSA verifier algorithm"""
         assert(self.valid)
+
+        # Extract r and s from the DER formatted signature. Return false for
+        # any DER encoding errors.
         if (sig[1] + 2 != len(sig)):
             return False
         if (len(sig) < 4):
@@ -275,11 +305,15 @@ class ECPubKey():
         if (slen > 1 and (sig[6+rlen] == 0) and not (sig[7+rlen] & 0x80)):
             return False
         s = int.from_bytes(sig[6+rlen:6+rlen+slen], 'big')
+
+        # Verify that r and s are within the group order
         if r < 1 or s < 1 or r >= SECP256K1_ORDER or s >= SECP256K1_ORDER:
             return False
         if low_s and s >= SECP256K1_ORDER_HALF:
             return False
         z = int.from_bytes(msg, 'big')
+
+        # Run verifier algorithm on r, s
         w = modinv(s, SECP256K1_ORDER)
         u1 = z*w % SECP256K1_ORDER
         u2 = r*w % SECP256K1_ORDER
@@ -331,7 +365,10 @@ class ECKey():
         return ret
 
     def sign_ecdsa(self, msg, low_s=True):
-        """Construct a DER-encoded ECDSA signature with this key."""
+        """Construct a DER-encoded ECDSA signature with this key.
+
+        See https://en.wikipedia.org/wiki/Elliptic_Curve_Digital_Signature_Algorithm for the
+        ECDSA signer algorithm."""
         assert(self.valid)
         z = int.from_bytes(msg, 'big')
         # Note: no RFC6979, but a simple random nonce (some tests rely on distinct transactions for the same operation)
@@ -341,6 +378,9 @@ class ECKey():
         s = (modinv(k, SECP256K1_ORDER) * (z + self.secret * r)) % SECP256K1_ORDER
         if low_s and s > SECP256K1_ORDER_HALF:
             s = SECP256K1_ORDER - s
+        # Represent in DER format. The byte representations of r and s have
+        # length rounded up (255 bits becomes 32 bytes and 256 bits becomes 33
+        # bytes).
         rb = r.to_bytes((r.bit_length() + 8) // 8, 'big')
         sb = s.to_bytes((s.bit_length() + 8) // 8, 'big')
         return b'\x30' + bytes([4 + len(rb) + len(sb), 2, len(rb)]) + rb + bytes([2, len(sb)]) + sb