implement client-server Proof Of Work

Mashed up a argon2-based PoW with token buckets and bloom filters.
This is intended to prevent a few abuses including:

* Using a keysafe server for general file storage, by storing a whole
  lot of chunks.
* An attacker guessing names that people will use, and uploading junk
  to keysafe servers under those names, to make it harder for others to use
  keysafe later.
* An attacker trying to guess the names used for objects on keysafe
  servers in order to download them and start password cracking.
  (As a second level of defense, since the name generation hash
  is expensive already.)

Completely untested, but it builds!

This commit was sponsored by Andreas on Patreon.

4 files changed, 325 insertions, 21 deletions

diff --git a/HTTP/Client.hs b/HTTP/Client.hs
index 0c28005..8f81db4 100644
--- a/HTTP/Client.hs
+++ b/HTTP/Client.hs
@@ -18,7 +18,7 @@ httpAPI :: Proxy HttpAPI
 httpAPI = Proxy
 
 motd :: Manager -> BaseUrl -> ClientM Motd
-getObject :: StorableObjectIdent -> Maybe ProofOfWork -> Manager -> BaseUrl -> ClientM (ProofOfWorkRequirement StorableObject)
-putObject :: StorableObjectIdent -> Maybe ProofOfWork -> StorableObject -> Manager -> BaseUrl -> ClientM (ProofOfWorkRequirement StoreResult)
-countObjects :: Maybe ProofOfWork -> Manager -> BaseUrl -> ClientM (ProofOfWorkRequirement CountResult)
+getObject :: StorableObjectIdent -> Maybe ProofOfWork -> Manager -> BaseUrl -> ClientM (POWGuarded StorableObject)
+putObject :: StorableObjectIdent -> Maybe ProofOfWork -> StorableObject -> Manager -> BaseUrl -> ClientM (POWGuarded StoreResult)
+countObjects :: Manager -> BaseUrl -> ClientM CountResult
 motd :<|> getObject :<|> putObject :<|> countObjects = client httpAPI
diff --git a/HTTP/ProofOfWork.hs b/HTTP/ProofOfWork.hs
index 45cc96d..04aec57 100644
--- a/HTTP/ProofOfWork.hs
+++ b/HTTP/ProofOfWork.hs
@@ -4,19 +4,117 @@
  -}
 
 {-# LANGUAGE DeriveGeneric #-}
+{-# OPTIONS_GHC -fno-warn-orphans #-}
 
 module HTTP.ProofOfWork where
 
 import Types
-import qualified Data.Text as T
+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 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 RandomSalt
+	deriving (Show)
+
+data ProofOfWorkRequirement = ProofOfWorkRequirement
+	{ leadingZeros :: Int
+	, addedArgon2Iterations :: Word32
+	, randomSalt :: RandomSalt
+	}
+	deriving (Generic, Show)
+
+-- | Using Text and not ByteString so that ProofOfWorkRequirement can have a
+-- JSON instance.
+newtype RandomSalt = RandomSalt { fromRandomSalt :: T.Text }
+	deriving (Generic, Show, Eq)
+
+instance Hashable RandomSalt where
+	hashIO32 = hashIO32 . encodeUtf8 . fromRandomSalt
+	hashIO64 = hashIO64 . encodeUtf8 . fromRandomSalt
+
+-- | 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 (RandomSalt -> ProofOfWorkRequirement)
+mkProofOfWorkRequirement (Seconds n)
+	| lz < 1 = Nothing
+	| otherwise = Just $ ProofOfWorkRequirement lz its
+  where
+	lz = floor (logBase 2 (max 1 (fromIntegral n / fromIntegral s)) :: Double)
+	UseArgon2 (CPUCost (Seconds s) _) _ = proofOfWorkHashTunable its
+	its = 0
+
+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
+
+instance Raaz.Random Word8
+
+isValidProofOfWork :: ProofOfWork -> ProofOfWorkRequirement -> StorableObjectIdent -> Bool
+isValidProofOfWork (ProofOfWork pow rsalt) req (StorableObjectIdent n) =
+	samesalts && enoughzeros
+  where
+	samesalts = rsalt == randomSalt req
+	enoughzeros = all (== False) (take (leadingZeros req) (setBits b))
+	tunable = proofOfWorkHashTunable (addedArgon2Iterations req)
+	salt = Salt $ POWSalt $
+		encodeUtf8 (fromRandomSalt (randomSalt req)) <> pow
+	ExpensiveHash _ hash = expensiveHash tunable salt n
+	-- 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
 
-data ProofOfWorkRequirement t
-	= Result t
-	| ProofOfWorkRequirement
-		{ leadingZeros :: Int
-		, argon2Iterations :: Int
-		}
-	deriving (Generic)
+instance Encodable POWSalt where
+	toByteString (POWSalt n) = n
+	fromByteString = Just . POWSalt
 
-newtype ProofOfWork = ProofOfWork T.Text
+genProofOfWork :: ProofOfWorkRequirement -> StorableObjectIdent -> ProofOfWork
+genProofOfWork req i = go allByteStrings
+  where
+	go [] = error "failed to generate Proof Of Work. This should be impossible!"
+	go (b:bs)
+		| isValidProofOfWork candidate req i = candidate
+		| otherwise = go bs
+	  where
+		candidate = ProofOfWork b (randomSalt req)
diff --git a/HTTP/RateLimit.hs b/HTTP/RateLimit.hs
new file mode 100644
index 0000000..9153664
--- /dev/null
+++ b/HTTP/RateLimit.hs
@@ -0,0 +1,203 @@
+{- Copyright 2016 Joey Hess <id@joeyh.name>
+ -
+ - Licensed under the GNU AGPL version 3 or higher.
+ -}
+
+module HTTP.RateLimit where
+
+import Types
+import Types.Cost
+import HTTP
+import HTTP.ProofOfWork
+import Servant
+import Control.Concurrent
+import Control.Concurrent.STM
+import Control.Concurrent.TokenBucket
+import qualified Data.BloomFilter.Mutable as BloomFilter
+import qualified Data.BloomFilter.Hash as BloomFilter
+import Data.BloomFilter.Easy (suggestSizing)
+import Control.Monad.ST
+import System.IO
+import Data.Maybe
+import Data.Word
+import Control.Monad.IO.Class
+
+-- | A rate limiter is a series of buckets. Each bucket has a
+-- successively more difficult proof of work access requirement.
+-- 
+-- To guard against DOS attacks that reuse the same proof of work,
+-- RandomSalt values are used, and bloom filters keep track of
+-- the ones that have been assigned and used.
+data RateLimiter = RateLimiter
+	{ buckets :: TMVar [Bucket]
+	, assignedRandomSalts :: BloomFilter
+	, assignedRandomSaltsOld :: BloomFilter
+	, usedRandomSalts :: BloomFilter
+	, usedRandomSaltsOld :: BloomFilter
+	, numRandomSalts :: TMVar Int
+	, accessLock :: TMVar ()
+	}
+
+type BloomFilter = TMVar (BloomFilter.MBloom RealWorld RandomSalt)
+
+-- | Buckets fill up at a fixed rate (which can be tuned by the server
+-- operator), and accessing a bucket removes one unit from it.
+data Bucket = Bucket
+	{ tokenBucket :: TokenBucket
+	, mkProofReq :: Maybe (RandomSalt -> ProofOfWorkRequirement)
+	}
+
+newRateLimiter :: IO RateLimiter
+newRateLimiter = RateLimiter 
+	<$> (newTMVarIO =<< mkbuckets maxProofOfWork [])
+	<*> mkBloomFilter
+	<*> mkBloomFilter
+	<*> mkBloomFilter
+	<*> mkBloomFilter
+	<*> newTMVarIO 0
+	<*> newTMVarIO ()
+  where
+	-- The last bucket takes half of maxProofOfWork to access, and
+	-- each earlier bucket halves that time, down to the first bucket,
+	-- which needs no proof of work. This ensures that in the edge case
+	-- where a client keeps getting bumped up to more and more expensive
+	-- buckets, it doesn't need to do more than maxProofOfWork total work.
+	mkbuckets (Seconds n) bs
+		| n <= 0 = return bs
+		| otherwise = do
+			let s = Seconds (n `div` 2)
+			let mkreq = mkProofOfWorkRequirement s
+			b <- Bucket
+				<$> newTokenBucket
+				<*> pure mkreq
+			case mkreq of
+				Nothing -> return (b:bs)
+				Just _ -> mkbuckets s (b:bs)
+
+mkBloomFilter :: IO BloomFilter
+mkBloomFilter = do
+	b <- stToIO $ BloomFilter.new (BloomFilter.cheapHashes bloomhashes) bloomsize
+	newTMVarIO b
+  where
+	-- Size the bloom filter to hold 1 million items, with a false
+	-- positive rate of 1 in 100 thousand. This will use around 32 mb
+	-- of memory.
+	(bloomhashes, bloomsize) = suggestSizing bloomMaxSize (1/100000)
+
+-- | Maximum number of RandomSalts that can be stored in a bloom filter
+-- without the false positive rate getting bad.
+bloomMaxSize :: Int
+bloomMaxSize = 1000000
+
+-- | Size of the bucket. This allows a burst of accesses after an idle
+-- period, which is especially useful when retrieving keys that were
+-- split into multiple chunks. However, setting this too high lets clients
+-- cheaply store lots of data, so keep the objectSize in mind.
+burstSize :: Word64
+burstSize = 4 -- allow 128 kb of data to be stored/retrieved w/o proof of work
+
+-- | Rate that the bucket is filled.
+fillRate :: Word64
+fillRate = 60000000 -- 1 token per minute
+
+-- A request is tried in each bucket in turn which its proof of work allows
+-- access to. If all accessible token buckets are empty, generate a
+-- new ProofOfWorkRequirement for the client.
+--
+-- If all buckets are tried and are empty, we must be very overloaded. 
+-- In this case, the request is still processed, since the client has done
+-- quite a lot of work.
+rateLimit :: RateLimiter -> Maybe ProofOfWork -> StorableObjectIdent -> Handler a -> Handler (POWGuarded a)
+rateLimit ratelimiter mpow i a = do
+	validsalt <- liftIO $ checkValidSalt ratelimiter mpow
+	bs <- liftIO $ atomically $ readTMVar (buckets ratelimiter)
+	if validsalt
+		then go bs
+		else assignWork ratelimiter bs
+  where
+	go [] = do
+		liftIO $ hPutStrLn stderr "** warning: all token buckets are empty; possible DOS attack?"
+		Result <$> a
+	go (b:bs) = case mkProofReq b of
+		Nothing -> checkbucket b bs
+		Just mkreq -> case mpow of
+			Nothing -> assignWork ratelimiter (b:bs)
+			Just pow@(ProofOfWork _ salt) -> 
+				if isValidProofOfWork pow (mkreq salt) i
+					then checkbucket b bs
+					else assignWork ratelimiter (b:bs)
+	checkbucket b bs = do
+		allowed <- liftIO $ tokenBucketTryAlloc (tokenBucket b) 
+			burstSize fillRate 1
+		if allowed
+			then Result <$> a
+			else go bs
+
+checkValidSalt :: RateLimiter -> Maybe ProofOfWork -> IO Bool
+checkValidSalt _ Nothing = return True
+checkValidSalt rl (Just (ProofOfWork _ salt)) = do
+	assigned <- iselem assignedRandomSalts
+	oldassigned <- iselem assignedRandomSaltsOld
+	used <- iselem usedRandomSalts
+	oldused <- iselem usedRandomSaltsOld
+	if assigned && not oldassigned && not used && not oldused
+		then do
+			withBloomFilter rl usedRandomSalts
+				(`BloomFilter.insert` salt)
+			return True
+		else return False
+  where
+	iselem f = withBloomFilter rl f (BloomFilter.elem salt)
+
+assignWork :: RateLimiter -> [Bucket] -> Handler (POWGuarded a)
+assignWork ratelimiter bs = case mapMaybe mkProofReq bs of
+	[] -> throwError err404
+	(mkreq:_) -> liftIO $ do
+		-- 10 second pause here, with a lock held so that only
+		-- this thread can run. This prevents an attacker
+		-- flooding requests that cause new random salts to be
+		-- assigned, in order to fill up the bloom table and cause
+		-- salts assigned to other clients to be rejected.
+		-- Since the bloom filters hold 1 million salts,
+		-- the attacker would need to send requests for over 10
+		-- hours to force a bloom filter rotation, so would not
+		-- impact many users.
+		atomically $ takeTMVar (accessLock ratelimiter)
+		threadDelay 10000000
+		atomically $ putTMVar (accessLock ratelimiter) ()
+
+		salt <- liftIO mkRandomSalt
+		withBloomFilter ratelimiter assignedRandomSalts
+			(`BloomFilter.insert` salt)
+		needrot <- atomically $ do
+			n <- takeTMVar (numRandomSalts ratelimiter)
+			if n > bloomMaxSize `div` 2
+				then return Nothing
+				else do
+					putTMVar (numRandomSalts ratelimiter) (n+1)
+					return (Just n)
+		handlerotation needrot
+		return $ NeedProofOfWork $ mkreq salt
+  where
+	handlerotation Nothing = return ()
+	handlerotation (Just n) = do
+		hPutStrLn stderr $ "rotating bloom filters after processing " ++ show n ++ " requests"
+		newassigned <- mkBloomFilter
+		newused <- mkBloomFilter
+		atomically $ do
+			oldassigned <- takeTMVar (assignedRandomSalts ratelimiter)
+			oldused <- takeTMVar (usedRandomSalts ratelimiter)
+			putTMVar (assignedRandomSaltsOld ratelimiter) oldassigned
+			putTMVar (usedRandomSaltsOld ratelimiter) oldused
+			putTMVar (assignedRandomSalts ratelimiter) =<< takeTMVar newassigned
+			putTMVar (usedRandomSalts ratelimiter) =<< takeTMVar newused
+			putTMVar (numRandomSalts ratelimiter) 0
+
+withBloomFilter
+	:: RateLimiter
+	-> (RateLimiter -> BloomFilter)
+	-> (BloomFilter.MBloom RealWorld RandomSalt -> ST RealWorld a)
+	-> IO a
+withBloomFilter rl field a = do
+	b <- atomically $ readTMVar (field rl)
+	stToIO (a b)
diff --git a/HTTP/Server.hs b/HTTP/Server.hs
index ab27aaa..dd35d1c 100644
--- a/HTTP/Server.hs
+++ b/HTTP/Server.hs
@@ -9,6 +9,7 @@ module HTTP.Server (runServer) where
 
 import HTTP
 import HTTP.ProofOfWork
+import HTTP.RateLimit
 import Types
 import Types.Storage
 import Tunables
@@ -26,12 +27,14 @@ import qualified Data.ByteString as B
 data ServerState = ServerState
 	{ obscurerRequest :: TMVar ()
 	, storageDirectory :: Maybe LocalStorageDirectory
+	, rateLimiter :: RateLimiter
 	}
 
 newServerState :: Maybe LocalStorageDirectory -> IO ServerState
 newServerState d = ServerState
 	<$> newEmptyTMVarIO
 	<*> pure d
+	<*> newRateLimiter
 
 runServer :: Maybe LocalStorageDirectory -> String -> Port -> IO ()
 runServer d bindaddress port = do
@@ -60,30 +63,30 @@ server st = motd
 motd :: Handler Motd
 motd = return $ Motd "Hello World!"
 
-getObject :: ServerState -> StorableObjectIdent -> Maybe ProofOfWork -> Handler (ProofOfWorkRequirement StorableObject)
-getObject st i _pow = do
+getObject :: ServerState -> StorableObjectIdent -> Maybe ProofOfWork -> Handler (POWGuarded StorableObject)
+getObject st i pow = rateLimit (rateLimiter st) pow i $ do
 	r <- liftIO $ retrieveShare (serverStorage st) dummyShareNum i
 	liftIO $ requestObscure st
 	case r of
-		RetrieveSuccess (Share _n o) -> return $ Result o
+		RetrieveSuccess (Share _n o) -> return o
 		RetrieveFailure _ -> throwError err404
 
-putObject :: ServerState -> StorableObjectIdent -> Maybe ProofOfWork -> StorableObject -> Handler (ProofOfWorkRequirement StoreResult)
-putObject st i _pow o = do
+putObject :: ServerState -> StorableObjectIdent -> Maybe ProofOfWork -> StorableObject -> Handler (POWGuarded StoreResult)
+putObject st i pow o = rateLimit (rateLimiter st) pow i $ do
 	if validObjectsize o
 		then do
 			r <- liftIO $ storeShare (serverStorage st) i (Share dummyShareNum o)
 			liftIO $ requestObscure st
-			return $ Result r
-		else return $ Result $ StoreFailure "invalid object size"
+			return r
+		else return $ StoreFailure "invalid object size"
 
 validObjectsize :: StorableObject -> Bool
 validObjectsize o = any (sz ==) knownObjectSizes
   where
 	sz = B.length (fromStorableObject o)
 
-countObjects :: ServerState -> Maybe ProofOfWork -> Handler (ProofOfWorkRequirement CountResult)
-countObjects st _pow = liftIO $ Result <$> countShares (serverStorage st)
+countObjects :: ServerState -> Handler CountResult
+countObjects = liftIO . countShares . serverStorage
 
 -- | 1 is a dummy value; the server does not know the actual share numbers.
 dummyShareNum :: ShareNum