;;;; ecdsa.scm: Elliptic Curve Digital Signature Algorithm based on bn.scm
;;; J. Welsh, October 2017

(lambda (error bn)
  (define bytes->bn (bn 'bytes->bn))
  (define unshift (bn 'bn-unshift))
  (define nbits (bn 'bn-bits))
  (define /2 (bn 'bn/2))
  (define bn-1 (bn 'bn-1))
  (define bn0 (bn 'bn0))
  (define bn1 (bn 'bn1))
  (define bn->hex (bn 'bn->hex))
  (define hex->bn (bn 'hex->bn))
  (define rand-bn (bn 'rand-bn))
  (define fix- -)
  (define fix-zero? zero?)

  ;; Convert the leftmost max-bits bits of big-endian byte sequence b to bignum
  (define (left-bytes->bn b max-bits)
    ; Equivalent to [SEC1] section 4.1.3 step 5
    (let ((e (bytes->bn b))
          (bits (* 8 (vector-length b))))
      (if (<= bits max-bits) e
        (unshift e (- bits max-bits)))))

  ;; Operations on a Weierstrass elliptic curve, over the prime field of order
  ;; p, defined by:
  ;;
  ;;   y^2 == x^3 + ax + b mod p
  ;;
  ;; ("a == b mod m" means modular congruence: a-b is a multiple of m.)
  ;;
  ;; g -- base point (generator) (cons gx gy)
  ;; n -- order of g
  ;; h -- cofactor (nh = number of points on the curve)
  (define (curve p a b g n h)
    (let ((+ (bn 'bn+))
          (- (bn 'bn-))
          (* (bn 'bn*))
          (*fix (bn 'bn*fix))
          (*2 (bn 'bn*2))
          (^2 (bn 'bn^2))
          (zero? (bn 'bn-zero?))
          (even? (bn 'bn-even?))
          (= (bn 'bn=))
          (> (bn 'bn>))
          (< (bn 'bn<))
          (divrem-by-p (delay ((bn 'bn-divrem-by) p (+ 2 (* 2 (nbits p))))))
          (n-bits (nbits n))
          (n/2 (/2 n))
          (remainder (bn 'bn-remainder))
          (mod-inverse (bn 'bn-mod-inverse)))

      ;; Compute a mod p by multiplication with precomputed inverse,
      ;; assuming a < 4p^2
      (define (modp a)
        ((force divrem-by-p) a (lambda (q r) r)))
      ;; In case you don't trust the optimization (which doesn't seem to help
      ;; much anyway):
      ;; (remainder a p)

      ;; It's perhaps less obvious how many bits would suffice here
      (define (modn a)
        (remainder a n))

      ;; Compute a-b mod p, avoiding negatives; output reduced if inputs are
      (define (mod- a b)
        (if (< a b) (- (+ p a) b) (- a b)))

      ;; Compute 2a mod p; output is reduced if input is
      (define (mod*2 a)
        (let ((b (*2 a)))
          (if (< b p) b (- b p))))

      (define (mod^2 a) (modp (^2 a)))
      (define (mod* a b) (modp (* a b)))

      ;; Saves a reduction compared to (= (modp a) (modp b))
      (define (congruent? a b)
        (zero? (modp (if (< a b) (- b a) (- a b)))))

      (define (on-curve? point)
        (if (eq? point 'inf) #t
          (let ((x (car point)) (y (cdr point)))
            (and (< x p)
                 (< y p)
                 (congruent? (^2 y)
                             (+ (* (mod^2 x) x) (+ (* a x) b)))))))

      ;; EC group operation, per [SEC1] section 2.2.1
      (define (ec+ p1 p2)
        (if (eq? p1 'inf) p2 ; adding the identity
          (if (eq? p2 'inf) p1
            (let ((x1 (car p1))
                  (y1 (cdr p1))
                  (x2 (car p2))
                  (y2 (cdr p2)))
              (let ((x2-x1 (mod- x2 x1)))
                (if (and (zero? x2-x1)
                         (or (zero? y1) (not (= y1 y2))))
                  'inf ; adding the inverse
                  (let* ((slope (if (zero? x2-x1)
                                  ; same point (doubling)
                                  (mod* (modp (+ (*fix (^2 x1) 3) a))
                                        (mod-inverse (mod*2 y1) p))
                                  (mod* (mod- y2 y1)
                                        (mod-inverse x2-x1 p))))
                         (x3 (mod- (mod- (mod^2 slope) x1) x2)))
                    (cons x3 (mod- (mod* slope (mod- x1 x3)) y1)))))))))

      ;; Supposedly doubling can be faster than general addition but I'm not
      ;; seeing how...
      (define (ec*2 p1)
        (if (eq? p1 'inf) p1
          (let ((x (car p1))
                (y (cdr p1)))
            (if (zero? y) 'inf
              (let ((slope (mod* (modp (+ (*fix (^2 x) 3) a))
                                 (mod-inverse (mod*2 y) p))))
                (let ((x3 (mod- (mod^2 slope) (mod*2 x))))
                  (cons x3 (mod- (mod* slope (mod- x x3)) y))))))))

      ;; Scalar multiplication: computes the equivalent of k repeated additions
      ;; of point p, in O(log k) time.
      (define (scalar* k p)
        (do ((k k (/2 k))
             (acc 'inf (if (even? k) acc (ec+ acc p*2^bits)))
             (p*2^bits p (ec*2 p*2^bits)))
          ((zero? k) acc)))

      ;; Optimize products of g by precomputation
      (define (scalar*g k)
        (do ((k k (/2 k))
             (acc 'inf (if (even? k) acc (ec+ acc (car doublings))))
             (doublings (force doublings-of-g) (cdr doublings)))
          ((zero? k) acc)))

      (define doublings-of-g
        (delay
          (do ((k (nbits (bn-1 n)) (fix- k 1))
               (g*2^bits g (ec*2 g*2^bits))
               (doublings '() (cons g*2^bits doublings)))
            ((fix-zero? k) (reverse doublings)))))

      (define gen-priv-key
        (let ((get-rand-int (rand-bn n)))
          (lambda (rng-port)
            (let ((k (get-rand-int rng-port)))
              (if (zero? k)
                (error "generated all-zero key?!")
                k)))))

      (define (valid-pub-key? point)
        ;; Per [SEC1] section 3.2.2.1
        (and (not (eq? point 'inf))
             (on-curve? point)
             (or (= h bn1) (eq? (scalar* n point) 'inf))))

      ;; Return an ECDSA signature of a message hash via (cont r s). Always
      ;; succeeds, assuming valid inputs and properly functioning RNG.
      ;;
      ;; As (r, -s mod n) is also a valid signature, the result is
      ;; canonicalized to use the lesser of the two possible s-values.
      ;;
      ;; Note the sharp edge of the scheme: you do need a good RNG here, as
      ;; knowledge of the ephemeral key compromises the private key. Repeated
      ;; r-values make it trivial to compute, and other forms of
      ;; predictability may well do the same.
      (define (sign hash private-key rng-port cont)
        ; Per [SEC1] section 4.1.3
        (let ((e (left-bytes->bn hash n-bits)))
          (define (find-temp-key)
            (let* ((k (gen-priv-key rng-port))
                   (r (modn (car (scalar*g k)))))
              (if (zero? r) (find-temp-key)
                (let ((s (modn (* (mod-inverse k n)
                                  (+ e (modn (* r private-key)))))))
                  (if (zero? s) (find-temp-key)
                    (cont r (if (< n/2 s) (- n s) s)))))))
          (find-temp-key)))

      ;; Return whether (r, s) is a valid ECDSA signature of hash by pub-key.
      ;; WARNING: you might also have to ensure (valid-pub-key? pub-key)!
      (define (valid-sig? r s hash pub-key)
        ;; Per [SEC1] section 4.1.4
        (and
          (< bn0 r) (< r n) (< bn0 s) (< s n)
          (let ((e (left-bytes->bn hash n-bits))
                (s-inv (mod-inverse s n)))
            (let ((u1 (modn (* e s-inv)))
                  (u2 (modn (* r s-inv))))
              (let ((bigR (ec+ (scalar*g u1)
                               (scalar* u2 pub-key))))
                (and (not (eq? bigR 'inf))
                     (= (modn (car bigR)) r)))))))

      (define (point->hex p)
        (if (eq? p 'inf) p
          (cons (bn->hex (car p))
                (bn->hex (cdr p)))))

      (define (hex->point p)
        (if (eq? p 'inf) p
          (cons (hex->bn (car p))
                (hex->bn (cdr p)))))

      (define (read-cache file)
        ;; I'd rather this be automatic, but no standard way in R5RS to check
        ;; if a file exists or handle errors if not...
        (let* ((data (with-input-from-file file read))
               (doublings (map hex->point data)))
          (set! doublings-of-g (delay doublings))))

      (define (write-cache file)
        (with-output-to-file
          file (lambda () (write (map point->hex (force doublings-of-g))))))

      (lambda (message)
        (case message
          ((ec+) ec+)
          ((gen-priv-key) gen-priv-key)
          ((priv->pub) scalar*g)
          ((valid-pub-key?) valid-pub-key?)
          ((sign) sign)
          ((valid-sig?) valid-sig?)
          ((read-cache) read-cache)
          ((write-cache) write-cache)
          (else (error "bad message:" message))))))

  ;;; Well-known curve parameters

  (define (from-hex o)
    (cond ((string? o) (hex->bn o))
          ((pair? o) (cons (from-hex (car o)) (from-hex (cdr o))))
          (else o)))

  (define (curve-from-hex . args)
    (apply curve (from-hex args)))

  ;; Generalized Koblitz curve over a 256-bit prime field, per [SEC2]
  (let ((secp256k1
          (curve-from-hex
            "fffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f"
            "0"
            "7"
            '("79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798" . "483ada7726a3c4655da4fbfc0e1108a8fd17b448a68554199c47d08ffb10d4b8")
            "fffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141"
            "1")))

    (export curve secp256k1)))

;;; References
;;
;; [SEC1] Certicom Research 2009. "Standards for Efficient Cryptography: SEC 1:
;;      Elliptic Curve Cryptography". Certicom Corp. Version 2.0.
;;      http://www.secg.org/sec1-v2.pdf
;;
;; [SEC2] Certicom Research 2010. "Standards for Efficient Cryptography: SEC 2:
;;      Recommended Elliptic Curve Domain Parameters". Certicom Corp. Version
;;      2.0. http://www.secg.org/sec2-v2.pdf
;;
;; [HAC] Menezes, A., van Oorshot, P. and Vanstone, S. 1996. "Handbook of
;;      Applied Cryptography". CRC Press. http://www.cacr.math.uwaterloo.ca/hac