/*
 * Copyright (c) 2013-2014 Galois, Inc.
 * Distributed under the terms of the BSD3 license (see LICENSE file)
 */

// see http://cr.yp.to/snuffle/spec.pdf

module Salsa20 where

quarterround : [4][32] -> [4][32]
quarterround [y0, y1, y2, y3] = [z0, z1, z2, z3]
  where
    z1 = y1 ^ ((y0 + y3) <<< 0x7)
    z2 = y2 ^ ((z1 + y0) <<< 0x9)
    z3 = y3 ^ ((z2 + z1) <<< 0xd)
    z0 = y0 ^ ((z3 + z2) <<< 0x12)

property quarterround_passes_tests =
  (quarterround [0x00000000, 0x00000000, 0x00000000, 0x00000000] == [0x00000000, 0x00000000, 0x00000000, 0x00000000]) &&
  (quarterround [0x00000001, 0x00000000, 0x00000000, 0x00000000] == [0x08008145, 0x00000080, 0x00010200, 0x20500000]) &&
  (quarterround [0x00000000, 0x00000001, 0x00000000, 0x00000000] == [0x88000100, 0x00000001, 0x00000200, 0x00402000]) &&
  (quarterround [0x00000000, 0x00000000, 0x00000001, 0x00000000] == [0x80040000, 0x00000000, 0x00000001, 0x00002000]) &&
  (quarterround [0x00000000, 0x00000000, 0x00000000, 0x00000001] == [0x00048044, 0x00000080, 0x00010000, 0x20100001]) &&
  (quarterround [0xe7e8c006, 0xc4f9417d, 0x6479b4b2, 0x68c67137] == [0xe876d72b, 0x9361dfd5, 0xf1460244, 0x948541a3]) &&
  (quarterround [0xd3917c5b, 0x55f1c407, 0x52a58a7a, 0x8f887a3b] == [0x3e2f308c, 0xd90a8f36, 0x6ab2a923, 0x2883524c])

// rowround : [16][32] -> [16][32]
rowround [y0, y1, y2, y3, y4, y5, y6, y7, y8, y9, y10, y11, y12, y13, y14, y15] =
    [z0, z1, z2, z3, z4, z5, z6, z7, z8, z9, z10, z11, z12, z13, z14, z15]
  where
    [ z0,  z1,  z2,  z3] = quarterround [ y0,  y1,  y2,  y3]
    [ z5,  z6,  z7,  z4] = quarterround [ y5,  y6,  y7,  y4]
    [z10, z11,  z8,  z9] = quarterround [y10, y11,  y8,  y9]
    [z15, z12, z13, z14] = quarterround [y15, y12, y13, y14]

property rowround_passes_tests =
  (rowround [0x00000001, 0x00000000, 0x00000000, 0x00000000,
             0x00000001, 0x00000000, 0x00000000, 0x00000000,
             0x00000001, 0x00000000, 0x00000000, 0x00000000,
             0x00000001, 0x00000000, 0x00000000, 0x00000000] ==
            [0x08008145, 0x00000080, 0x00010200, 0x20500000,
             0x20100001, 0x00048044, 0x00000080, 0x00010000,
             0x00000001, 0x00002000, 0x80040000, 0x00000000,
             0x00000001, 0x00000200, 0x00402000, 0x88000100]) &&
  (rowround [0x08521bd6, 0x1fe88837, 0xbb2aa576, 0x3aa26365,
             0xc54c6a5b, 0x2fc74c2f, 0x6dd39cc3, 0xda0a64f6,
             0x90a2f23d, 0x067f95a6, 0x06b35f61, 0x41e4732e,
             0xe859c100, 0xea4d84b7, 0x0f619bff, 0xbc6e965a] ==
            [0xa890d39d, 0x65d71596, 0xe9487daa, 0xc8ca6a86,
             0x949d2192, 0x764b7754, 0xe408d9b9, 0x7a41b4d1,
             0x3402e183, 0x3c3af432, 0x50669f96, 0xd89ef0a8,
             0x0040ede5, 0xb545fbce, 0xd257ed4f, 0x1818882d])

/* commented out for now -- not called, and type checker not yet happy
rowround_opt : [16][32] -> [16][32]
rowround_opt ys = join [ (quarterround (yi>>>i))<<<i | yi <- split ys | i <- [0 .. 0x3] ]

rowround_opt_is_rowround ys = rowround ys == rowround_opt ys
*/

columnround : [16][32] -> [16][32]
columnround [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15] =
    [y0, y1, y2, y3, y4, y5, y6, y7, y8, y9, y10, y11, y12, y13, y14, y15]
  where
    [ y0,  y4,  y8, y12] = quarterround [ x0,  x4,  x8, x12]
    [ y5,  y9, y13,  y1] = quarterround [ x5,  x9, x13,  x1]
    [y10, y14,  y2,  y6] = quarterround [x10, x14,  x2,  x6]
    [y15,  y3,  y7, y11] = quarterround [x15,  x3,  x7, x11]

property columnround_passes_tests =
  (columnround [0x00000001, 0x00000000, 0x00000000, 0x00000000,
                0x00000001, 0x00000000, 0x00000000, 0x00000000,
                0x00000001, 0x00000000, 0x00000000, 0x00000000,
                0x00000001, 0x00000000, 0x00000000, 0x00000000] ==
               [0x10090288, 0x00000000, 0x00000000, 0x00000000,
                0x00000101, 0x00000000, 0x00000000, 0x00000000,
                0x00020401, 0x00000000, 0x00000000, 0x00000000,
                0x40a04001, 0x00000000, 0x00000000, 0x00000000]) &&
  (columnround [0x08521bd6, 0x1fe88837, 0xbb2aa576, 0x3aa26365,
                0xc54c6a5b, 0x2fc74c2f, 0x6dd39cc3, 0xda0a64f6,
                0x90a2f23d, 0x067f95a6, 0x06b35f61, 0x41e4732e,
                0xe859c100, 0xea4d84b7, 0x0f619bff, 0xbc6e965a] ==
               [0x8c9d190a, 0xce8e4c90, 0x1ef8e9d3, 0x1326a71a,
                0x90a20123, 0xead3c4f3, 0x63a091a0, 0xf0708d69,
                0x789b010c, 0xd195a681, 0xeb7d5504, 0xa774135c,
                0x481c2027, 0x53a8e4b5, 0x4c1f89c5, 0x3f78c9c8])

/* commented out for now -- not called, and type checker not yet happy
columnround_opt : [16][32] -> [16][32]
columnround_opt xs = join (transpose [ (quarterround (xi<<<i))>>>i | xi <- transpose(split xs) | i <- [0 .. 3] ])

columnround_opt_is_columnround xs = columnround xs == columnround_opt xs
*/

property columnround_is_transpose_of_rowround ys =
  rowround ys == join(transpose(split`{4}(columnround(xs))))
  where xs = join(transpose(split`{4}ys))

doubleround : [16][32] -> [16][32]
doubleround(xs) = rowround(columnround(xs))

property doubleround_passes_tests =
  (doubleround [0x00000001, 0x00000000, 0x00000000, 0x00000000,
                0x00000000, 0x00000000, 0x00000000, 0x00000000,
                0x00000000, 0x00000000, 0x00000000, 0x00000000,
                0x00000000, 0x00000000, 0x00000000, 0x00000000] ==
               [0x8186a22d, 0x0040a284, 0x82479210, 0x06929051,
                0x08000090, 0x02402200, 0x00004000, 0x00800000,
                0x00010200, 0x20400000, 0x08008104, 0x00000000,
                0x20500000, 0xa0000040, 0x0008180a, 0x612a8020]) &&
  (doubleround [0xde501066, 0x6f9eb8f7, 0xe4fbbd9b, 0x454e3f57,
                0xb75540d3, 0x43e93a4c, 0x3a6f2aa0, 0x726d6b36,
                0x9243f484, 0x9145d1e8, 0x4fa9d247, 0xdc8dee11,
                0x054bf545, 0x254dd653, 0xd9421b6d, 0x67b276c1] ==
               [0xccaaf672, 0x23d960f7, 0x9153e63a, 0xcd9a60d0,
                0x50440492, 0xf07cad19, 0xae344aa0, 0xdf4cfdfc,
                0xca531c29, 0x8e7943db, 0xac1680cd, 0xd503ca00,
                0xa74b2ad6, 0xbc331c5c, 0x1dda24c7, 0xee928277])

littleendian : [4][8] -> [32]
littleendian b = join(reverse b)

property littleendian_passes_tests =
  (littleendian [  0,   0,   0,   0] == 0x00000000) &&
  (littleendian [ 86,  75,  30,   9] == 0x091e4b56) &&
  (littleendian [255, 255, 255, 250] == 0xfaffffff)

littleendian_inverse : [32] -> [4][8]
littleendian_inverse b = reverse(split b)

littleendian_is_invertable b = littleendian_inverse(littleendian(b)) == b

Salsa20 : [64][8] -> [64][8]
Salsa20 xs = join(ar)
  where
    ar = [ littleendian_inverse words | words <- xw + (zs@0xa) ]
    xw = [ littleendian xi    | xi <- split xs ]
    zs = [xw] # [ doubleround(zi) | zi <- zs ]

property Salsa20_passes_tests =
  (Salsa20 [  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,   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,   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,   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,   0,   0,   0,   0]) &&
  (Salsa20 [211, 159,  13, 115,  76,  55,  82, 183,   3, 117, 222,  37, 191, 187, 234, 136,
             49, 237, 179,  48,   1, 106, 178, 219, 175, 199, 166,  48,  86,  16, 179, 207,
             31, 240,  32,  63,  15,  83,  93, 161, 116, 147,  48, 113, 238,  55, 204,  36,
             79, 201, 235,  79,   3,  81, 156,  47, 203,  26, 244, 243,  88, 118, 104,  54] ==
           [109,  42, 178, 168, 156, 240, 248, 238, 168, 196, 190, 203,  26, 110, 170, 154,
             29,  29, 150,  26, 150,  30, 235, 249, 190, 163, 251,  48,  69, 144,  51,  57,
            118,  40, 152, 157, 180,  57,  27,  94, 107,  42, 236,  35,  27, 111, 114, 114,
            219, 236, 232, 135, 111, 155, 110,  18,  24, 232,  95, 158, 179,  19,  48, 202]) &&
  (Salsa20 [ 88, 118, 104,  54,  79, 201, 235,  79,   3,  81, 156,  47, 203,  26, 244, 243,
            191, 187, 234, 136, 211, 159,  13, 115,  76,  55,  82, 183,   3, 117, 222,  37,
             86,  16, 179, 207,  49, 237, 179,  48,   1, 106, 178, 219, 175, 199, 166,  48,
            238,  55, 204,  36,  31, 240,  32,  63,  15,  83,  93, 161, 116, 147,  48, 113] ==
           [179,  19,  48, 202, 219, 236, 232, 135, 111, 155, 110,  18,  24, 232,  95, 158,
             26, 110, 170, 154, 109,  42, 178, 168, 156, 240, 248, 238, 168, 196, 190, 203,
             69, 144,  51,  57,  29,  29, 150,  26, 150,  30, 235, 249, 190, 163, 251,  48,
             27, 111, 114, 114, 118,  40, 152, 157, 180,  57,  27,  94, 107,  42, 236,  35])

property Salsa20_has_no_collisions x1 x2 =
  if(x1 != x2) then doubleround(x1) != doubleround(x2) else True

// Salsa 20 supports two key sizes, [16][8] and [32][8]
Salsa20_expansion : {a} (a >= 1, 2 >= a) => ([16*a][8], [16][8]) -> [64][8]
Salsa20_expansion(k, n) = z
  where
    [s0, s1, s2, s3] = split("expand 32-byte k") : [4][4][8]
    [t0, t1, t2, t3] = split("expand 16-byte k") : [4][4][8]
    x = if(`a == 2) then s0 # k0 # s1 # n # s2 # k1 # s3
                    else t0 # k0 # t1 # n # t2 # k0 # t3
    z = Salsa20(x)
    [k0, k1] = (split(k#zero)):[2][16][8]

Salsa20_encrypt : {a, b} (a >= 1, 2 >= a, fin b) => ([16*a][8], [8][8], [b][8]) -> [b][8]
Salsa20_encrypt(k, v, m) = c
  where
    salsa = take`{b} (join [ Salsa20_expansion(k, v#(split i)) | i <- [0, 1 ... ] ])
    c = m ^ salsa