;;; Curve25519: the elliptic curve y^2 = x^3 + 486662 x^2 + x over the
;;; prime field Z/(2^255 - 19)Z
;;;
;;; X25519: the Diffie-Hellman function based on x-coordinate scalar
;;; multiplication in Curve25519.
;;;
;;; DO NOT USE THIS CODE IN PRACTICE!  It is for experiments and for
;;; validating the results of other implementations.  It does not
;;; defend against side channel attacks, and I am not totally sure it
;;; is correct.
;;;
;;; Public domain, or if you can't handle that then cc0
;;; <https://creativecommons.org/publicdomain/zero/1.0/>.

(define (modexp a e n)
  (if (zero? e)
      1
      (let loop ((r 1) (b (modulo a n)) (f e))
        ;; r b^f = a^e (mod n)
        (let ((r (if (even? f) r (modulo (* r b) n)))
              (f (quotient f 2)))
          (if (zero? f)
              r
              (loop r (modulo (* b b) n) f))))))

(define p25519 (- (expt 2 255) 19))

(define (gf25519:zero? a) (zero? a))
(define (gf25519:one? a) (= a 1))

(define (gf25519:+ a b) (modulo (+ a b) p25519))
(define (gf25519:- a b) (modulo (- a b) p25519))
(define (gf25519:* a b) (modulo (* a b) p25519))
(define (gf25519:/ a b) (gf25519:* a (gf25519:recip b)))

(define (gf25519:negate a) (gf25519:- 0 a))
(define (gf25519:recip a) (gf25519:expt a (- p25519 2)))

(define (gf25519:double a) (gf25519:+ a a))

(define (gf25519:expt a n) (modexp a n p25519))
(define (gf25519:square a) (gf25519:* a a))

(define (gf25519:nonsquare? a)
  (= (gf25519:negate 1) (gf25519:expt a (/ (- p25519 1) 2))))
(assert (not (gf25519:nonsquare? 9)))

(assert (= 5 (modulo p25519 8)))
(assert (= 0 (modulo (+ p25519 3) 8)))
(define gf25519:i (gf25519:expt 2 (/ (- p25519 1) 4)))
(define (gf25519:sqrt a)
  (assert (not (gf25519:nonsquare? a)))
  (let ((r (gf25519:expt a (/ (+ p25519 3) 8))))
    (if (= (gf25519:square r) (gf25519:negate a))
        (gf25519:* gf25519:i r)
        r)))
(assert
 (= (gf25519:sqrt
     (gf25519:+ (gf25519:expt 9 3)
                (gf25519:+ (gf25519:* 486662 (gf25519:square 9))
                           9)))
    #x20ae19a1b8a086b4e01edd2c7748d14c923d4d7e6d7c61b229e9c5a27eced3d9))

(define (make-point x y) (cons x y))
(define (point-x p) (car p))
(define (point-y p) (cdr p))
(define infinity 'INFINITY)
(define (infinity? p) (eq? p infinity))

(define (curve25519? p)
  (or (infinity? p)
      (gf25519:zero?
       (gf25519:-
        (gf25519:square (point-y p))
        (gf25519:+
         (gf25519:expt (point-x p) 3)
         (gf25519:+
          (gf25519:* 486662 (gf25519:square (point-x p)))
          (point-x p)))))))

(define (curve25519:+ p q)
  (assert (curve25519? p))
  (assert (curve25519? q))
  (cond ((and (infinity? p) (infinity? q)) infinity)
        ((infinity? p) q)
        ((infinity? q) p)
        (else
         (let ((px (point-x p))
               (py (point-y p))
               (qx (point-x q))
               (qy (point-y q)))
           (assert (= px (modulo px p25519)))
           (assert (= qx (modulo qx p25519)))
           (assert (= py (modulo py p25519)))
           (assert (= qy (modulo qy p25519)))
           (cond ((= py (gf25519:negate qy)) infinity)
                 ((and (= px qx) (= py qy))
                  (let ((x px) (y py))
                    (let* ((l
                            (gf25519:/
                             (gf25519:+
                              (gf25519:* 3 (gf25519:square x))
                              (gf25519:+
                               (gf25519:double (gf25519:* 486662 x))
                               1))
                             (gf25519:double y)))
                           (x*
                            (gf25519:- (gf25519:- (gf25519:square l) 486662)
                                       (gf25519:double x))))
                      (make-point
                       x*
                       (gf25519:- (gf25519:* l (gf25519:- x x*)) y)))))
                 ((= px qx)
                  (assert (not (= py (gf25519:negate qy))))
                  (error "Not on curve25519:" p q))
                 (else
                  (let* ((l (gf25519:/ (gf25519:- qy py)
                                       (gf25519:- qx px)))
                         (x*
                          (gf25519:-
                           (gf25519:-
                            (gf25519:- (gf25519:square l) 486662)
                            px)
                           qx)))
                    (make-point x*
                                (gf25519:- (gf25519:* l (gf25519:- px x*))
                                           py)))))))))

;;; Naive scalar multiplication using the horrible addition formula
;;; with add-and-double steps.
(define (curve25519:* n p)
  (let loop ((r infinity) (q p) (m n))
    ;; r + m q = n p
    (let ((r (if (even? m) r (curve25519:+ r q)))
          (m (quotient m 2)))
      (if (zero? m)
          r
          (loop r (curve25519:+ q q) m)))))

;;; Montgomery ladder as written in the Explicit Formulas Database
;;; <http://hyperelliptic.org/EFD/g1p/auto-montgom-xz.html> and
;;; <https://tools.ietf.org/html/draft-irtf-cfrg-curves-01>.
(define (curve25519:ladder-efd n x)
  (define (cswap b x y)
    (let ((d (gf25519:* b (gf25519:- x y))))
      (values (gf25519:- x d)
              (gf25519:+ y d))))
  (let ((a24 (/ (- 486662 2) 4)))
    (let loop ((t 254) (x1 x) (x2 1) (z2 0) (x3 x) (z3 1))
      (let ((bit (extract-bit t n)))
        (receive (x2 x3) (cswap bit x2 x3)
          (receive (z2 z3) (cswap bit z2 z3)
            (let* ((a (gf25519:+ x2 z2))
                   (aa (gf25519:square a))
                   (b (gf25519:- x2 z2))
                   (bb (gf25519:square b))
                   (e (gf25519:- aa bb))
                   (c (gf25519:+ x3 z3))
                   (d (gf25519:- x3 z3))
                   (da (gf25519:* d a))
                   (cb (gf25519:* c b))
                   (x3 (gf25519:square (gf25519:+ da cb)))
                   (z3 (gf25519:* x1 (gf25519:square (gf25519:- da cb))))
                   (x2 (gf25519:* aa bb))
                   (z2 (gf25519:* e (gf25519:+ aa (gf25519:* a24 e)))))
              (receive (x2 x3) (cswap bit x2 x3)
                (receive (z2 z3) (cswap bit z2 z3)
                  (if (zero? t)
                      (gf25519:* x2 (gf25519:expt z2 (- p25519 2)))
                      (loop (- t 1) x1 x2 z2 x3 z3)))))))))))

;;; Montgomery ladder as written in NaCl's crypto_scalarmult_curve25519
;;; reference implementation.
(define (curve25519:ladder-nacl n x)
  (define (cswap b x y)
    (let ((d (gf25519:* b (gf25519:- x y))))
      (values (gf25519:- x d)
              (gf25519:+ y d))))
  (let ((a24 (/ (- 486662 2) 4)))
    (let loop ((t 254) (xm1 x) (zm1 1) (xm 1) (zm 0))
      (let ((bit (extract-bit t n)))
        (receive (xmb xm1b) (cswap bit xm xm1)
          (receive (zmb zm1b) (cswap bit zm zm1)
            (let* ((a0x (gf25519:+ xmb zmb))
                   (a0z (gf25519:- xmb zmb))
                   (a1x (gf25519:+ xm1b zm1b))
                   (a1z (gf25519:- xm1b zm1b))
                   (b0x (gf25519:square a0x))
                   (b0z (gf25519:square a0z))
                   (b1x (gf25519:* a1x a0z))
                   (b1z (gf25519:* a1z a0x))
                   (c1x (gf25519:+ b1x b1z))
                   (c1z (gf25519:- b1x b1z))
                   (rx (gf25519:square c1z))
                   (sx (gf25519:- b0x b0z))
                   (tx (gf25519:* a24 sx))
                   (ux (gf25519:+ tx b0x))
                   (xnb (gf25519:* b0x b0z))
                   (znb (gf25519:* sx ux))
                   (xn1b (gf25519:square c1x))
                   (zn1b (gf25519:* rx x)))
              (receive (xm xm1) (cswap bit xnb xn1b)
                (receive (zm zm1) (cswap bit znb zn1b)
                  (if (zero? t)
                      (gf25519:* xm (gf25519:recip zm))
                      (loop (- t 1) xm1 zm1 xm zm)))))))))))

(define (mont25519:double x z)
  (let ((a24 (/ (- 486662 2) 4)))
    (let ((x+z^2 (gf25519:square (gf25519:+ x z)))
          (x-z^2 (gf25519:square (gf25519:- x z))))
      (values
       (gf25519:* x+z^2 x-z^2)
       (gf25519:*
        (gf25519:- x+z^2 x-z^2)
        (gf25519:+ x+z^2
                   (gf25519:* a24
                              (gf25519:- x+z^2 x-z^2))))))))

(define (mont25519:diff+ qx qz px pz q-px q-pz)
  (let ((qx-qz (gf25519:- qx qz))
        (qx+qz (gf25519:+ qx qz))
        (px+pz (gf25519:+ px pz))
        (px-pz (gf25519:- px pz)))
    (let ((qx-qz*px+pz (gf25519:* qx-qz px+pz))
          (qx+qz*px-pz (gf25519:* qx+qz px-pz)))
      (values
       (gf25519:*
        q-pz
        (gf25519:square (gf25519:+ qx-qz*px+pz qx+qz*px-pz)))
       (gf25519:*
        q-px
        (gf25519:square (gf25519:- qx-qz*px+pz qx+qz*px-pz)))))))

(define (curve25519-decompress x sign)
  (make-point x
              (gf25519:* sign
                         (gf25519:sqrt
                          (gf25519:+
                           (gf25519:expt x 3)
                           (gf25519:+
                            (gf25519:* 486662 (gf25519:square x))
                            x))))))

(assert
 (equal?
  (curve25519-decompress 9 1)
  (make-point
   9
   #x20ae19a1b8a086b4e01edd2c7748d14c923d4d7e6d7c61b229e9c5a27eced3d9)))

(define (le256-decode v)
  (assert (= 32 (vector-length v)))
  (let loop ((i 0) (x 0))
    (if (< i 32)
        (loop (+ i 1)
              (+ x (* (vector-ref v i) (expt 2 (* 8 i)))))
        x)))

(define (le256-encode x)
  (let ((v (make-vector 32)))
    (let loop ((i 0) (x x))
      (if (< i 32)
          (begin
            (vector-set! v i (modulo x 256))
            (loop (+ i 1) (quotient x 256)))))
    v))

(define (x25519-clampc v)
  (assert (= 32 (vector-length v)))
  (define (vector-modify! v i f)
    (vector-set! v i (f (vector-ref v i))))
  (let ((v (let ((v0 (make-vector 32)))
             (do ((i 0 (+ i 1))) ((>= i 32))
               (vector-set! v0 i (vector-ref v i)))
             v0)))
    (vector-modify! v 0 (lambda (a0) (- a0 (modulo a0 8))))
    (vector-modify! v 31 (lambda (a31) (+ 64 (modulo a31 64))))
    v))

(define x25519-base
  (vector 9 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0))

(define (x25519 n p)
  (le256-encode
   (point-x
    (curve25519:* (le256-decode (x25519-clampc n))
                  (curve25519-decompress (le256-decode p) 1)))))

(define (x25519-nacl n p)
  (le256-encode
   (curve25519:ladder-nacl (le256-decode (x25519-clampc n))
                           (le256-decode p))))

(define (x25519-efd n p)
  (le256-encode
   (curve25519:ladder-efd (le256-decode (x25519-clampc n))
                          (le256-decode p))))

(let ((secret
       '#(#x77 #x07 #x6d #x0a #x73 #x18 #xa5 #x7d
          #x3c #x16 #xc1 #x72 #x51 #xb2 #x66 #x45
          #xdf #x4c #x2f #x87 #xeb #xc0 #x99 #x2a
          #xb1 #x77 #xfb #xa5 #x1d #xb9 #x2c #x2a))
      (public
       '#(#x85 #x20 #xf0 #x09 #x89 #x30 #xa7 #x54
          #x74 #x8b #x7d #xdc #xb4 #x3e #xf7 #x5a
          #x0d #xbf #x3a #x0d #x26 #x38 #x1a #xf4
          #xeb #xa4 #xa9 #x8e #xaa #x9b #x4e #x6a)))
  (assert (equal? public (x25519 secret x25519-base)))
  (assert (equal? public (x25519-nacl secret x25519-base)))
  (assert (equal? public (x25519-efd secret x25519-base))))