implemented dynamic rate limiting

7 files changed, 117 insertions, 46 deletions

diff --git a/CmdLine.hs b/CmdLine.hs
index 11dd34e..67af2da 100644
--- a/CmdLine.hs
+++ b/CmdLine.hs
@@ -33,7 +33,7 @@ data Mode = Backup | Restore | UploadQueued | Server | Benchmark | Test
 data ServerConfig = ServerConfig
 	{ serverPort :: Port
 	, serverAddress :: String
-	, monthsToFillHalfDisk :: Int
+	, monthsToFillHalfDisk :: Integer
 	}
 
 parse :: Parser CmdLine
diff --git a/HTTP/RateLimit.hs b/HTTP/RateLimit.hs
index da22b92..d9ec752 100644
--- a/HTTP/RateLimit.hs
+++ b/HTTP/RateLimit.hs
@@ -9,15 +9,22 @@ import Types.Cost
 import HTTP
 import HTTP.ProofOfWork
 import Tunables
+import CmdLine (ServerConfig(..))
+import Types.Storage
+import Storage.Local
 import Servant
+import Control.Concurrent
 import Control.Concurrent.STM
 import Control.Concurrent.TokenBucket
+import Control.Concurrent.Thread.Delay
 import qualified Data.BloomFilter.Mutable as BloomFilter
 import qualified Data.BloomFilter.Hash as BloomFilter
 import Data.BloomFilter.Easy (suggestSizing)
+import Control.Monad
 import Control.Monad.ST
 import Control.Exception.Lifted (bracket)
 import System.IO
+import System.DiskSpace
 import Data.Maybe
 import Data.Word
 import Control.Monad.IO.Class
@@ -38,6 +45,7 @@ data RateLimiter = RateLimiter
 	, numRandomSalts :: TMVar Int
 	, randomSaltGenerationLimiter :: TokenBucket
 	, blockedRequestQueue :: TMVar [()]
+	, requestCounter :: TMVar Integer
 	}
 
 type BloomFilter = TMVar (BloomFilter.MBloom RealWorld RandomSalt)
@@ -61,17 +69,21 @@ minFillInterval = 2 * 60 * 1000000 -- 1 token every other minute
 burstSize :: Word64
 burstSize = 4 -- 256 kb immediate storage
 
-newRateLimiter :: IO RateLimiter
-newRateLimiter = RateLimiter 
-	<$> (newTMVarIO =<< mkbuckets (sdiv maxProofOfWork 2) [])
-	<*> newTMVarIO []
-	<*> mkBloomFilter
-	<*> mkBloomFilter
-	<*> mkBloomFilter
-	<*> mkBloomFilter
-	<*> newTMVarIO 0
-	<*> newTokenBucket
-	<*> newTMVarIO []
+newRateLimiter :: ServerConfig -> Maybe LocalStorageDirectory -> IO RateLimiter
+newRateLimiter cfg storedir = do
+	rl <- RateLimiter 
+		<$> (newTMVarIO =<< mkbuckets (sdiv maxProofOfWork 2) [])
+		<*> newTMVarIO []
+		<*> mkBloomFilter
+		<*> mkBloomFilter
+		<*> mkBloomFilter
+		<*> mkBloomFilter
+		<*> newTMVarIO 0
+		<*> newTokenBucket
+		<*> newTMVarIO []
+		<*> newTMVarIO 0
+	_ <- forkIO (adjusterThread cfg storedir rl)
+	return rl
   where
 	-- The last bucket takes half of maxProofOfWork to access, and
 	-- each earlier bucket quarters that time, down to the first bucket,
@@ -136,7 +148,7 @@ rateLimit ratelimiter mpow p a = do
 		allowed <- liftIO $ tokenBucketTryAlloc (tokenBucket b) 
 			burstSize (fillInterval b) 1
 		if allowed
-			then Result <$> a
+			then allowRequest ratelimiter a
 			else go bs
 
 checkValidSalt :: RateLimiter -> Maybe ProofOfWork -> IO Bool
@@ -255,7 +267,7 @@ allBucketsEmpty ratelimiter a = bracket (liftIO addq) (liftIO . removeq) go
 	go True = do
 		ok <- liftIO waitlast
 		if ok
-			then Result <$> a
+			then allowRequest ratelimiter a
 			else giveup
 	
 	giveup = do
@@ -265,13 +277,14 @@ allBucketsEmpty ratelimiter a = bracket (liftIO addq) (liftIO . removeq) go
 -- | How much data could be stored, in bytes per second, assuming all
 -- buckets in the rate limiter being constantly drained by requests,
 -- and all requests store objects.
-maximumStorageRate :: RateLimiter -> IO Int
+maximumStorageRate :: RateLimiter -> IO Integer
 maximumStorageRate ratelimiter = do
 	bs <- getBuckets ratelimiter
 	return $ sum $ map calc bs
   where
 	storesize = maximum knownObjectSizes
-	calc b = (storesize * 1000000) `div` fromIntegral (fillInterval b)
+	calc b = fromIntegral $ 
+		(storesize * 1000000) `div` fromIntegral (fillInterval b)
 
 describeRateLimiter :: RateLimiter -> IO String
 describeRateLimiter ratelimiter = do
@@ -280,13 +293,14 @@ describeRateLimiter ratelimiter = do
 	return $ concat
 		[ "rate limiter buckets: " ++ show bs
 		, " ; maximum allowed storage rate: "
-		, showrate (storerate * 60 * 60 * 24 * 31) ++ "/month"
+		, showBytes (storerate * 60 * 60 * 24 * 31) ++ "/month"
 		]
-  where
-	showrate n
-		| n < 1024*1024 = show (n `div` 1024) ++ " KiB"
-		| n < 1024*1024*1024 = show (n `div` (1024 * 1024)) ++ " MiB"
-		| otherwise = show (n `div` (1024 * 1024 * 1024)) ++ " GiB"
+
+showBytes :: Integer -> String
+showBytes n
+	| n <= 1024*1024 = show (n `div` 1024) ++ " KiB"
+	| n <= 1024*1024*1024 = show (n `div` (1024 * 1024)) ++ " MiB"
+	| otherwise = show (n `div` (1024 * 1024 * 1024)) ++ " GiB"
 
 instance Show Bucket where
 	show b = show (fillInterval b `div` (60 * 1000000)) ++ " Second/Request"
@@ -297,11 +311,13 @@ increaseDifficulty ratelimiter = do
 	bs <- getBuckets ratelimiter
 	case bs of
 		[] -> unable
-		(b:[]) -> do
-			-- Make the remaining bucket take longer to fill.
-			let b' = b { fillInterval = fillInterval b * 2 }
-			putBuckets ratelimiter [b']
-			done
+		(b:[])
+			| fillInterval b < maxBound `div` 2 -> do
+				-- Make the remaining bucket take longer to fill.
+				let b' = b { fillInterval = fillInterval b * 2 }
+				putBuckets ratelimiter [b']
+				done
+			| otherwise -> unable
 		(b:rest) -> do
 			-- Remove less expensive to access buckets,
 			-- so that clients have to do some work.
@@ -313,7 +329,7 @@ increaseDifficulty ratelimiter = do
 			putBuckets ratelimiter rest
 			done
   where
-	unable = putStrLn "unable to increase difficulty; out of buckets"
+	unable = putStrLn "Unable to increase difficulty any further!"
 	done = do
 		desc <- describeRateLimiter ratelimiter
 		putStrLn $ "increased difficulty -- " ++ desc
@@ -348,3 +364,49 @@ reduceDifficulty ratelimiter = do
 	done = do
 		desc <- describeRateLimiter ratelimiter
 		putStrLn $ "reduced difficulty -- " ++ desc
+
+allowRequest :: RateLimiter -> Handler a -> Handler (POWGuarded a)
+allowRequest ratelimiter a = do
+	liftIO $ addRequest ratelimiter 1
+	Result <$> a
+
+addRequest :: RateLimiter -> Integer -> IO ()
+addRequest ratelimiter n = liftIO $ atomically $ do
+	v <- takeTMVar c
+	putTMVar c (v + n)
+  where
+	c = requestCounter ratelimiter
+
+-- Thread that wakes up periodically and checks the request rate
+-- against the available disk space. If the disk is filling too quickly,
+-- the difficulty is increased.
+adjusterThread :: ServerConfig -> Maybe LocalStorageDirectory -> RateLimiter -> IO ()
+adjusterThread cfg storedir ratelimiter = forever $ do
+	delay (1000000 * intervalsecs)
+	checkRequestRate cfg storedir ratelimiter intervalsecs
+  where
+	intervalsecs = 60*60
+
+checkRequestRate :: ServerConfig -> Maybe LocalStorageDirectory -> RateLimiter -> Integer -> IO ()
+checkRequestRate cfg storedir ratelimiter intervalsecs = do
+	let storesize = maximum knownObjectSizes
+	n <- liftIO $ atomically $ swapTMVar (requestCounter ratelimiter) 0
+	let maxstoredinterval = n * fromIntegral storesize
+	let maxstoredthismonth = maxstoredinterval * (intervalsecs `div` (60*60)) * 24 * 31
+	freespace <- diskFree <$> localDiskUsage storedir
+	let target = monthsToFillHalfDisk cfg
+	let estimate = if maxstoredthismonth <= 0 
+		then 10000
+		else freespace `div` maxstoredthismonth `div` 2
+	putStrLn $ unlines
+		[ "rate limit check"
+		, "  free disk space:" ++ showBytes freespace
+		, "  number of requests since last check: " ++ show n
+		, "  estimated max incoming data in the next month: " ++ showBytes maxstoredthismonth
+		, "  estimate min " ++ show estimate ++ " months to fill half of disk"
+		]
+	if estimate > target * 2
+		then reduceDifficulty ratelimiter
+		else if estimate < target
+			then increaseDifficulty ratelimiter
+			else return ()
diff --git a/HTTP/Server.hs b/HTTP/Server.hs
index 65d3d32..aab3dab 100644
--- a/HTTP/Server.hs
+++ b/HTTP/Server.hs
@@ -13,6 +13,7 @@ import HTTP.RateLimit
 import Types
 import Types.Storage
 import Tunables
+import CmdLine (ServerConfig(..))
 import Storage.Local
 import Serialization ()
 import Servant
@@ -20,6 +21,7 @@ import Network.Wai
 import Network.Wai.Handler.Warp
 import Control.Monad.IO.Class
 import Control.Concurrent
+import Control.Concurrent.Thread.Delay
 import Control.Concurrent.STM
 import Data.String
 import qualified Data.ByteString as B
@@ -30,20 +32,20 @@ data ServerState = ServerState
 	, rateLimiter :: RateLimiter
 	}
 
-newServerState :: Maybe LocalStorageDirectory -> IO ServerState
-newServerState d = ServerState
+newServerState :: Maybe LocalStorageDirectory -> ServerConfig -> IO ServerState
+newServerState d cfg = ServerState
 	<$> newEmptyTMVarIO
 	<*> pure d
-	<*> newRateLimiter
+	<*> newRateLimiter cfg d
 
-runServer :: Maybe LocalStorageDirectory -> String -> Port -> IO ()
-runServer d bindaddress port = do
-	st <- newServerState d
+runServer :: Maybe LocalStorageDirectory -> ServerConfig -> IO ()
+runServer d cfg = do
+	st <- newServerState d cfg
 	_ <- forkIO $ obscurerThread st
 	runSettings settings (app st)
   where
-	settings = setHost host $ setPort port $ defaultSettings 
-	host = fromString bindaddress
+	settings = setHost host $ setPort (serverPort cfg) $ defaultSettings 
+	host = fromString (serverAddress cfg)
 
 serverStorage :: ServerState -> Storage
 serverStorage st = localStorage (storageDir $ storageDirectory st) "server"
@@ -100,7 +102,7 @@ obscurerThread :: ServerState -> IO ()
 obscurerThread st = do
 	_ <- obscureShares (serverStorage st)
 	putStrLn "obscured shares"
-	threadDelay (1000000*60*30)
+	delay (1000000*60*30)
 	_ <- atomically $ takeTMVar (obscurerRequest st)
 	obscurerThread st
 
diff --git a/Storage/Local.hs b/Storage/Local.hs
index b9f0f3e..71d5aa7 100644
--- a/Storage/Local.hs
+++ b/Storage/Local.hs
@@ -3,7 +3,13 @@
  - Licensed under the GNU AGPL version 3 or higher.
  -}
 
-module Storage.Local (localStorage, storageDir, testStorageDir, uploadQueue) where
+module Storage.Local
+	( localStorage
+	, storageDir
+	, testStorageDir
+	, uploadQueue
+	, localDiskUsage
+	) where
 
 import Types
 import Types.Storage
@@ -22,6 +28,7 @@ import Raaz.Core.Encode
 import Control.DeepSeq
 import Control.Exception
 import Control.Monad
+import System.DiskSpace
 
 type GetShareDir = Section -> IO FilePath
 
@@ -127,6 +134,11 @@ storageDir (Just (LocalStorageDirectory d)) (Section section) =
 testStorageDir :: FilePath -> GetShareDir
 testStorageDir tmpdir = storageDir (Just (LocalStorageDirectory tmpdir))
 
+localDiskUsage :: Maybe LocalStorageDirectory -> IO DiskUsage
+localDiskUsage lsd = do
+	dir <- storageDir lsd (Section ".")
+	getDiskUsage dir
+
 -- | The takeFileName ensures that, if the StorableObjectIdent somehow
 -- contains a path (eg starts with "../" or "/"), it is not allowed
 -- to point off outside the shareDir.
diff --git a/TODO b/TODO
index 8bfb3f0..0bb8c65 100644
--- a/TODO
+++ b/TODO
@@ -1,7 +1,6 @@
 Soon:
 
-* test client/server Proof Of Work
-* finish --months-to-fill-half-disk
+* test rate limiting
 * Add some random padding to http requests and responses, to make it
   harder for traffic analysis to tell that it's keysafe traffic.
 * Implement the different categories of servers in the server list.
@@ -14,10 +13,6 @@ Soon:
 
 Later:
 
-* server: Keep track of number of objects stored per hour, and project
-  forward to see how long it would take to fill the disk at the current
-  rate. If less than some number of years, need to adjust the rate limiting
-  knobs to require more work be done to store objects.
 * improve restore progress bar points (update after every hash try)
 * If we retrieved enough shares successfully, but decrypt failed, must
   be a wrong password, so prompt for re-entry and retry with those shares.
diff --git a/keysafe.cabal b/keysafe.cabal
index a1e6fe2..b249b37 100644
--- a/keysafe.cabal
+++ b/keysafe.cabal
@@ -64,6 +64,7 @@ Executable keysafe
     , bloomfilter == 2.0.*
     , disk-free-space == 0.1.*
     , lifted-base == 0.2.*
+    , unbounded-delays == 0.1.*
     -- Temporarily inlined due to https://github.com/ocharles/argon2/issues/3
     -- argon2 == 1.1.*
   Extra-Libraries: argon2
diff --git a/keysafe.hs b/keysafe.hs
index 98e1ecb..e2b112c 100644
--- a/keysafe.hs
+++ b/keysafe.hs
@@ -67,8 +67,7 @@ dispatch cmdline ui storagelocations tunables possibletunables = do
 	go (CmdLine.Server) _ =
 		runServer
 			(CmdLine.localstoragedirectory cmdline)
-			(CmdLine.serverAddress $ CmdLine.serverConfig cmdline)
-			(CmdLine.serverPort $ CmdLine.serverConfig cmdline)
+			(CmdLine.serverConfig cmdline)
 	go CmdLine.Benchmark _ =
 		benchmarkTunables tunables
 	go CmdLine.Test _ =