;;;; Cryptographic hash functions in pure Scheme
;;; J. Welsh, November 2017 - May 2018
;;; Trimmed for gbw, March 2020

(lambda (fz-ops)
  (define word->ufz (fz-ops 'word->ufz))
  (define ufz+ (fz-ops 'ufz+))
  (define ufzshift (fz-ops 'ufzshift))
  (define fzrot (fz-ops 'fzrot))
  (define fznot (fz-ops 'fznot))
  (define fzior (fz-ops 'fzior))
  (define fzxor (fz-ops 'fzxor))
  (define fzif (fz-ops 'fzif))
  (define fzmaj (fz-ops 'fzmaj))
  (define bytes->fz (fz-ops 'bytes->fz))
  (define fz->bytes (fz-ops 'fz->bytes))
  (define hex->bytes (fz-ops 'hex->bytes))
  (define hex->fz (fz-ops 'hex->fz))

  (define (compose f g) (lambda (x) (f (g x))))

  (define null-byte (integer->char 0))
  (define initial-padding (string (integer->char 128)))

  ;; The Merkle-Damgard construction: turns a fixed size compression function into an arbitrary size hash function, by blockwise iteration with padding and length hardening. For simplicity, this implementation requires the full message to fit in a string.
  (define (merkle-damgard block-size length-bits big-endian compress offset iv)
    (define (finish hash tail packed-bit-length)
      (let ((length-size (string-length packed-bit-length))
            (space (- block-size (string-length tail))))
        (if (< space length-size)
            (finish (compress hash (string-append
                                     tail (make-string space null-byte)))
                    "" packed-bit-length)
            (compress hash (string-append
                             ;; assuming length goes at end of block
                             tail (make-string (- space length-size) null-byte)
                             packed-bit-length)))))
    (lambda (msg)
      (let ((len (string-length msg)))
        (do ((pos 0 next)
             (next block-size (+ next block-size))
             (hash iv (compress hash (substring msg pos next))))
          ((> next len)
           (finish hash
                   (string-append (substring msg pos len) initial-padding)
                   (fz->bytes (ufzshift (ufz+ (word->ufz length-bits len)
                                              (word->ufz length-bits offset))
                                        3)
                              big-endian)))))))

  (define (append-hex->bytes . args)
    (apply string-append (map hex->bytes args)))
  (define (hex->fz-vec . args)
    (list->vector (map hex->fz args)))

  ;; Load aligned word from byte string, little/big endian
  (define (slice s k len) (substring s k (+ k len)))
  (define (loadw32 s k) (bytes->fz (slice s (fxshift k 2) 4) #f))
  (define (loadbw32 s k) (bytes->fz (slice s (fxshift k 2) 4) #t))

  ;; Pack words as byte string, little/big endian
  (define (pack-words w)
    (apply string-append (map (lambda (w) (fz->bytes w #f)) w)))
  (define (packb-words w)
    (apply string-append (map (lambda (w) (fz->bytes w #t)) w)))

  ;; The various "sigma" functions from [SHS] for 256-bit hashes
  (define (ss0-32 x) (fzxor (fzrot x -2) (fzrot x -13) (fzrot x -22)))
  (define (ss1-32 x) (fzxor (fzrot x -6) (fzrot x -11) (fzrot x -25)))
  (define (s0-32 x) (fzxor (fzrot x -7) (fzrot x -18) (ufzshift x -3)))
  (define (s1-32 x) (fzxor (fzrot x -17) (fzrot x -19) (ufzshift x -10)))

  (let ((rmd160-iv (append-hex->bytes
                     "01234567" "89abcdef" "fedcba98" "76543210" "f0e1d2c3"))
        ;; First 32 fractional bits of square roots of first 8 primes
        (sha256-iv (append-hex->bytes
                     "6a09e667" "bb67ae85"
                     "3c6ef372" "a54ff53a"
                     "510e527f" "9b05688c"
                     "1f83d9ab" "5be0cd19")))

    (define rmd160-compress
      (letrec ((F fzxor)
               (G fzif)
               (H (lambda (x y z) (fzxor (fzior (fznot y) x) z)))
               (I (lambda (x y z) (G z x y)))
               (J (lambda (x y z) (H y z x)))
               (FF (lambda (a b c d e x s)
                     (ufz+ (fzrot (ufz+ a (F b c d) x) s) e)))
               (make-op (lambda (func const)
                          (let ((k (hex->fz const)))
                            (lambda (a b c d e x s)
                              (ufz+ (fzrot (ufz+ a (func b c d) x k) s) e))))))
        ;; Floor of 2^30 times square and cube roots of (2 3 5 7)
        (let ((ops1 (vector FF
                            (make-op G "5a827999")
                            (make-op H "6ed9eba1")
                            (make-op I "8f1bbcdc")
                            (make-op J "a953fd4e")))
              (ops2 (vector (make-op J "50a28be6")
                            (make-op I "5c4dd124")
                            (make-op H "6d703ef3")
                            (make-op G "7a6d76e9")
                            FF))
              (rho (lambda (i)
                     (vector-ref '#(7 4 13 1 10 6 15 3 12 0 9 5 2 14 11 8) i)))
              (pi (lambda (i) (fxand (+ (* 9 i) 5) 15)))
              (shifts '#(#(11 14 15 12 5 8 7 9 11 13 14 15 6 7 9 8)
                         #(12 13 11 15 6 9 9 7 12 15 11 13 7 8 7 7)
                         #(13 15 14 11 7 7 6 8 13 14 13 12 5 5 6 9)
                         #(14 11 12 14 8 6 5 5 15 12 15 14 9 9 8 6)
                         #(15 12 13 13 9 5 8 6 14 11 12 11 8 6 5 5)))
              (W1 (make-vector 80))
              (W2 (make-vector 80)))
          (do ((t 0 (+ t 1))) ((= t 16))
            (vector-set! W1 t t)
            (vector-set! W2 t (pi t)))
          (do ((t 16 (+ t 1))) ((= t 80))
            (vector-set! W1 t (rho (vector-ref W1 (- t 16))))
            (vector-set! W2 t (rho (vector-ref W2 (- t 16)))))

          (lambda (hash block)
            (let ((h0 (loadw32 hash 0))
                  (h1 (loadw32 hash 1))
                  (h2 (loadw32 hash 2))
                  (h3 (loadw32 hash 3))
                  (h4 (loadw32 hash 4))
                  (X (make-vector 16)))
              (do ((t 0 (+ t 1))) ((= t 16))
                (vector-set! X t (loadw32 block t)))
              (let loop ((t 0) (a1 h0) (b1 h1) (c1 h2) (d1 h3) (e1 h4)
                               (a2 h0) (b2 h1) (c2 h2) (d2 h3) (e2 h4))
                (if (= t 80)
                    (pack-words (map ufz+ (list h1 h2 h3 h4 h0)
                                     (list c1 d1 e1 a1 b1)
                                     (list d2 e2 a2 b2 c2)))
                    (let ((round (fxshift/unsigned t -4)))
                      (let ((shifts (vector-ref shifts round))
                            (op1 (vector-ref ops1 round))
                            (op2 (vector-ref ops2 round))
                            (i1 (vector-ref W1 t))
                            (i2 (vector-ref W2 t)))
                        (loop (+ t 1)
                              e1 (op1 a1 b1 c1 d1 e1 (vector-ref X i1)
                                      (vector-ref shifts i1))
                              b1 (fzrot c1 10) d1
                              e2 (op2 a2 b2 c2 d2 e2 (vector-ref X i2)
                                      (vector-ref shifts i2))
                              b2 (fzrot c2 10) d2))))))))))

    (define sha256-compress
      ;; First 32 fractional bits of cube roots of first 64 primes
      (let ((K (hex->fz-vec
                 "428a2f98" "71374491" "b5c0fbcf" "e9b5dba5"
                 "3956c25b" "59f111f1" "923f82a4" "ab1c5ed5"
                 "d807aa98" "12835b01" "243185be" "550c7dc3"
                 "72be5d74" "80deb1fe" "9bdc06a7" "c19bf174"
                 "e49b69c1" "efbe4786" "0fc19dc6" "240ca1cc"
                 "2de92c6f" "4a7484aa" "5cb0a9dc" "76f988da"
                 "983e5152" "a831c66d" "b00327c8" "bf597fc7"
                 "c6e00bf3" "d5a79147" "06ca6351" "14292967"
                 "27b70a85" "2e1b2138" "4d2c6dfc" "53380d13"
                 "650a7354" "766a0abb" "81c2c92e" "92722c85"
                 "a2bfe8a1" "a81a664b" "c24b8b70" "c76c51a3"
                 "d192e819" "d6990624" "f40e3585" "106aa070"
                 "19a4c116" "1e376c08" "2748774c" "34b0bcb5"
                 "391c0cb3" "4ed8aa4a" "5b9cca4f" "682e6ff3"
                 "748f82ee" "78a5636f" "84c87814" "8cc70208"
                 "90befffa" "a4506ceb" "bef9a3f7" "c67178f2")))
        (lambda (hash block)
          (let ((W (make-vector 64))
                (h0 (loadbw32 hash 0)) (h1 (loadbw32 hash 1))
                (h2 (loadbw32 hash 2)) (h3 (loadbw32 hash 3))
                (h4 (loadbw32 hash 4)) (h5 (loadbw32 hash 5))
                (h6 (loadbw32 hash 6)) (h7 (loadbw32 hash 7)))
            (do ((t 0 (+ t 1))) ((= t 16))
              (vector-set! W t (loadbw32 block t)))
            (do ((t 16 (+ t 1))) ((= t 64))
              (vector-set! W t (ufz+ (s1-32 (vector-ref W (- t 2)))
                                     (vector-ref W (- t 7))
                                     (s0-32 (vector-ref W (- t 15)))
                                     (vector-ref W (- t 16)))))
            (let loop ((t 0) (a h0) (b h1) (c h2) (d h3) (e h4) (f h5) (g h6)
                             (h h7))
              (if (= t 64)
                  (packb-words (map ufz+ (list a b c d e f g h)
                                    (list h0 h1 h2 h3 h4 h5 h6 h7)))
                  (let* ((temp1 (ufz+ (ss1-32 e)
                                      (fzif e f g)
                                      h
                                      (vector-ref K t)
                                      (vector-ref W t)))
                         (temp2 (ufz+ (ss0-32 a)
                                      (fzmaj a b c)
                                      temp1)))
                    (loop (+ t 1) temp2 a b c (ufz+ d temp1) e f g))))))))

    (let ((ripemd160 (merkle-damgard 64 64 #f rmd160-compress 0 rmd160-iv))
          (sha256 (merkle-damgard 64 64 #t sha256-compress 0 sha256-iv)))
      (export ripemd160 sha256))))

;;; References

;; H. Dobbertin, A. Bosselaers, B. Preneel 1996. "RIPEMD-160: A Strengthened Version of RIPEMD." http://homes.esat.kuleuven.be/bosselae/ripemd160/pdf/AB-9601/AB-9601.pdf
;; [SHS] Information Technology Laboratory 2015. "FIPS PUB 180-4: Secure Hash Standard." National Institute of Standards and Technology. http://dx.doi.org/10.6028/NIST.FIPS.180-4