My own impl of secret-sharing (Shamir scheme)

Shamir.hs

{- Example: How to use------------------------------------------------------------

$ stack ghci
-- Import or load the below module: Shamir
> :l Shamir

-- Prepare any string "secret" less than 32-byte
> secret = "stop COVID-19"

-- Prepare parameter (n, k)
-- where n := "total number of token"
--       k := "minimum number of tokens required for decoding"
> (n, k) = (5, 3)

-- Encode a secret given
> share <- encode secret n k

-- Print the share generated
-- these shares go to friends in trust
> mapM_ (putStrLn . show) share

-- Recover secret from shares
> secret == decode share

-- Successfuly decoded only if the number of token is greater than or equal to k
> decode $ take 1 share    -- fail
> decode $ take 2 share    -- fail
> decode $ take 3 share    -- ok
> decode $ take 4 share    -- ok

-- With k >= 3, enough regardless configuration of shares
> decode $ drop 1 . take 4 $ share    -- ok
> decode $ last share : head share : share !! 3 : []    -- ok

-} -------------------------------------------------------------------------------

module Shamir
  ( encode
  , decode
  , evalLagrangePoly
  )
where

import           Data.Word                      ( Word8 )
import           Data.Bits
import           System.Random
import           Control.Monad                  ( replicateM )
import qualified Data.ByteString               as BS
import qualified Data.ByteString.Char8         as BC
import qualified Data.ByteString.Base16        as Hex

type Point = (Integer, Integer)

-- | Mersenne prime, where p = 127
-- fp = 2 ^ 127 - 1 :: Integer

-- | SECP256K1P = 2^256 - 2^32 - 977
fp =
  0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f :: Integer

-- | Modular power with base, exponent, and modulus
modPow :: Integer -> Integer -> Integer -> Integer
modPow b e m | m <= 0    = error "Non-positive modulo"
             | e < 0     = error "Negative exponent"
             | e == 0    = 1
             | otherwise = t * modPow ((b * b) `mod` m) (shiftR e 1) m `mod` m
  where t = if testBit e 0 then b `mod` m else 1

-- | Perform a `div` b in Finite Field, where a, b in Fp
(//) :: Integer -> Integer -> Integer
num // dem = (num * modPow dem (fp - 2) fp) `mod` fp

-- | Get a list of clamped-random-number within [lo, hi]
rollDice :: Random a => (a, a) -> Int -> IO [a]
rollDice (lo, hi) n = replicateM n (randomRIO (lo, hi))

-- | Evaluate a Lagrange interpolation polynomial
evalLagrangePoly :: [Point] -> Integer -> Integer
evalLagrangePoly xys xi = sum (zipWith (*) ys $ basis <$> xs) `mod` fp where
  xs = fst <$> xys
  ys = snd <$> xys
  basis xj = product (frac xj <$> xs) `mod` fp
  frac xj xm = if xj == xm then 1 else (xi - xm) // (xj - xm)

-- | Generate Shares ([Point]) from a secret given (String)
encode :: String -> Int -> Int -> IO [Point]
encode secret n k
  | i > fp - 1 = ioError $ userError "Secret is too long"
  | otherwise = do
    cs <- (++) [i] <$> rollDice (1, fp - 1) (k - 1)
    let point x = (,) x $ foldr f 0 cs where f a b = (a + x * b) `mod` fp
    return [ point $ fromIntegral i | i <- [1 .. n] ]
  where i = integerFromBytes . bytesFromString $ secret

-- | Recover the secret from the Shares
decode :: [Point] -> String
decode []    = undefined
decode share = stringFromBytes . bytesFromInteger $ evalLagrangePoly share 0

-- | Converters among String, ByteString, and Integer
bytesFromString :: String -> BC.ByteString
bytesFromString = BC.pack

stringFromBytes :: BC.ByteString -> String
stringFromBytes = BC.unpack

u8sFromBytes :: BC.ByteString -> [Word8]
u8sFromBytes = BS.unpack

integerFromBytes :: BS.ByteString -> Integer
integerFromBytes = BS.foldl' f 0 where f a b = shiftL a 8 .|. fromIntegral b

bytesFromInteger :: (Integral a, Bits a) => a -> BC.ByteString
bytesFromInteger x = BS.pack $ fromIntegral <$> u8s where
  u8s = u8 <$> reverse ((8 *) <$> [0 .. n])
  u8 s = shiftR x s `mod` 0x100
  n = floor . logBase 0x100 . fromIntegral $ x