{- Copyright 2016 Joey Hess <id@joeyh.name>
 -
 - Licensed under the GNU AGPL version 3 or higher.
 -}

{-# LANGUAGE DeriveGeneric, OverloadedStrings #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module HTTP.ProofOfWork where

import Types
import Types.Cost
import ExpensiveHash
import Tunables
import ByteStrings
import GHC.Generics
import qualified Data.Text as T
import qualified Data.ByteString as B
import Data.Text.Encoding (encodeUtf8)
import Raaz.Core.Encode
import qualified Raaz
import Data.BloomFilter.Hash
import Control.Monad
import Control.DeepSeq
import Data.Word
import Data.Bits
import Data.Monoid
import Prelude

-- | A value that the client has to do some work to calculate.
data ProofOfWork = ProofOfWork B.ByteString RequestID
	deriving (Show, Generic)

instance NFData ProofOfWork

data ProofOfWorkRequirement = ProofOfWorkRequirement
	{ leadingZeros :: Int
	, addedArgon2Iterations :: Word32
	, requestID :: RequestID
	}
	deriving (Generic, Show)

-- | A request ID has two parts, a RandomSalt and a HMAC.
-- The server can verify if a request ID is one it generated.
data RequestID = RequestID
	{ randomSalt :: RandomSalt
	, requestHMAC :: T.Text
	}
	deriving (Generic, Show, Eq)

instance NFData RequestID

instance Hashable RequestID where
	hashIO32 = hashIO32 . hashRequestID
	hashIO64 = hashIO64 . hashRequestID

hashRequestID :: RequestID -> B.ByteString
hashRequestID rid = encodeUtf8 (fromRandomSalt (randomSalt rid)) 
	<> ":" <> encodeUtf8 (requestHMAC rid)

-- | Using Text and not ByteString so that ProofOfWorkRequirement can have a
-- JSON instance.
newtype RandomSalt = RandomSalt { fromRandomSalt :: T.Text }
	deriving (Generic, Show, Eq)

instance NFData RandomSalt

-- | Servers should never demand a proof of work that takes longer than
-- this to generate. Note that if a server changes its mind and doubles
-- the proof of work, a client counts that cumulatively. So, a server
-- should avoid any single proof of work requirement taking more than half
-- this long.
maxProofOfWork :: Seconds
maxProofOfWork = Seconds (16*60)

-- | How long it will take to generate a proof of work meeting the
-- requirement, maximum.
--
-- Of course, a client can get lucky and find a value that works
-- on the very first try. On average, the client will need to work for half
-- as long as the returned number of Seconds.
generationTime :: ProofOfWorkRequirement -> Seconds
generationTime req = 
	let UseArgon2 (CPUCost (Seconds s) _) _ = proofOfWorkHashTunable (addedArgon2Iterations req)
	in Seconds ((2^(leadingZeros req)) * s)

mkProofOfWorkRequirement :: Seconds -> Maybe (RequestID -> ProofOfWorkRequirement)
mkProofOfWorkRequirement (Seconds n)
	| lz < 1 || n <= 1 = Nothing
	| otherwise = Just $ ProofOfWorkRequirement lz its
  where
	lz = floor (logBase 2 (fromRational (max 1 (n / s))) :: Double)
	UseArgon2 (CPUCost (Seconds s) _) _ = proofOfWorkHashTunable its
	its = 0

newtype RequestIDSecret = RequestIDSecret (Raaz.Key (Raaz.HMAC Raaz.SHA256))

newRequestIDSecret :: IO RequestIDSecret
newRequestIDSecret = do
	prg <- Raaz.newPRG () :: IO Raaz.SystemPRG
	RequestIDSecret <$> Raaz.random prg

mkRequestID :: RequestIDSecret -> IO RequestID
mkRequestID secret = mkRequeestID' secret <$> mkRandomSalt

mkRequeestID' :: RequestIDSecret -> RandomSalt -> RequestID
mkRequeestID' (RequestIDSecret key) salt =
	let hmac = Raaz.hmacSha256 key (encodeUtf8 $ fromRandomSalt salt)
	in RequestID salt (T.pack (showBase16 hmac))

validRequestID :: RequestIDSecret -> RequestID -> Bool
validRequestID secret rid =
	let rid' = mkRequeestID' secret (randomSalt rid)
	in requestHMAC rid == requestHMAC rid'

mkRandomSalt :: IO RandomSalt
mkRandomSalt = do
	prg <- Raaz.newPRG () :: IO Raaz.SystemPRG
	rs <- replicateM 16 (Raaz.random prg :: IO Word8)
	return $ RandomSalt $ T.pack $ concatMap show rs

class POWIdent p where
	getPOWIdent :: p -> B.ByteString

instance POWIdent StorableObjectIdent where
	getPOWIdent (StorableObjectIdent i) = i

data NoPOWIdent = NoPOWIdent

instance POWIdent NoPOWIdent where
	getPOWIdent NoPOWIdent = B.empty

instance POWIdent Int where
	getPOWIdent = encodeUtf8 . T.pack . show

-- Note that this does not check validRequestID.
isValidProofOfWork :: POWIdent p => ProofOfWork -> ProofOfWorkRequirement -> p -> Bool
isValidProofOfWork (ProofOfWork pow rid) req p = samerequestids && enoughzeros
  where
	samerequestids = rid == requestID req
	enoughzeros = all (== False) (take (leadingZeros req) (setBits b))
	tunable = proofOfWorkHashTunable (addedArgon2Iterations req)
	salt = Salt $ POWSalt $
		encodeUtf8 (fromRandomSalt (randomSalt (requestID req))) <> pow
	ExpensiveHash _ hash = expensiveHash tunable salt (getPOWIdent p)
	-- Since expensiveHash generates an ascii encoded hash that
	-- includes the parameters, take the sha256 of it to get the
	-- bytestring that is what's checked for the neccesary number
	-- of leading 0 bits.
	b = Raaz.toByteString $ Raaz.sha256 $ encodeUtf8 hash

setBits :: B.ByteString -> [Bool]
setBits = concatMap go . B.unpack
  where
	go byte = map (uncurry testBit) (zip (repeat byte) [0..7])

newtype POWSalt = POWSalt B.ByteString

instance Encodable POWSalt where
	toByteString (POWSalt n) = n
	fromByteString = Just . POWSalt

genProofOfWork :: POWIdent p => ProofOfWorkRequirement -> p -> ProofOfWork
genProofOfWork req p = go allByteStrings
  where
	go [] = error "failed to generate Proof Of Work. This should be impossible!"
	go (b:bs)
		| isValidProofOfWork candidate req p = candidate
		| otherwise = go bs
	  where
		candidate = ProofOfWork b (requestID req)