Rb 5271 blacklisting and global ratelimiting

build.sbt

@@ -2,6 +2,14 @@ name := "ratelimiter2"
 
 organization := "org.tunnelbear"
 
-version := "0.1"
+version := "1.0"
 
 scalaVersion := "2.12.6"
+
+libraryDependencies += "org.specs2" %% "specs2-core" % "4.8.3" % "test"
+libraryDependencies += "org.specs2" %% "specs2-mock" % "4.8.3" % "test"
+
+// Cross-build to support Scala 2.11 projects (remembackend...), and Scala 2.12 projects (tbearDashboard2, polarbackend...)
+// See https://www.scala-sbt.org/1.x/docs/Cross-Build.html
+lazy val root = (project in file(".")).settings(
+  crossScalaVersions := List("2.12.6", "2.11.7"))

src/main/scala/RateLimiter/RateLimiterService.scala

@@ -1,32 +1,44 @@
 package RateLimiter
 
-import RateLimiter.RateLimiters.{AuthLimiter, IPLimiter, TagLimiter}
+import RateLimiter.RateLimiters._
 
 import scala.concurrent.ExecutionContext
 import scala.concurrent.duration.Duration
 
+// TODO: change blacklistOnBlock field to enableBlacklisting?
 trait RateLimiterService {
   implicit def storage: RateLimiterStorage
 
   def dictLimit: Long
   def dictExpiry: Duration
+  def dictBlacklist: Boolean
 
   def bruteLimit: Long
   def bruteExpiry: Duration
+  def bruteBlacklist: Boolean
 
   def ipLimit: Long
   def ipExpiry: Duration
+  def ipBlacklist: Boolean
 
+  def tagLimit(tag: String): Long
+  def tagExpiry(tag: String): Duration
+  def tagBlacklist(tag: String): Boolean
 
   def authLimiter(ip: String, userIdentifier: String)(implicit executionContext: ExecutionContext): AuthLimiter = {
-    AuthLimiter(ip, userIdentifier, dictLimit, dictExpiry.toMillis, bruteLimit, bruteExpiry.toMillis)
+    AuthLimiter(ip, userIdentifier, dictLimit, dictExpiry.toMillis, dictBlacklist, bruteLimit, bruteExpiry.toMillis, bruteBlacklist)
   }
 
-  def ipLimiter(ip: String)(implicit executionContext: ExecutionContext) : IPLimiter = {
-    IPLimiter(ip, ipLimit, ipExpiry.toMillis)
+  def ipLimiter(ip: String)(implicit executionContext: ExecutionContext): IPLimiter = {
+    IPLimiter(ip, ipLimit, ipExpiry.toMillis, ipBlacklist)
   }
 
-  def tagLimiter(tag: String, ip: String, limit: Long, expiry: Duration)(implicit executionContext: ExecutionContext): TagLimiter = {
-    TagLimiter(tag, ip, limit, expiry.toMillis)
+  def tagLimiter(tag: String, ip: String)(implicit executionContext: ExecutionContext): TagLimiter = {
+    TagLimiter(tag, ip, tagLimit(tag), tagExpiry(tag).toMillis, tagBlacklist(tag))
+  }
+
+  // TODO: define separate methods for this?
+  def globalTagLimiter(tag: String)(implicit executionContext: ExecutionContext): GlobalTagLimiter = {
+    GlobalTagLimiter(tag, tagLimit(tag), tagExpiry(tag).toMillis)
   }
 }

src/main/scala/RateLimiter/RateLimiterStatus.scala

@@ -0,0 +1,9 @@
+package RateLimiter
+
+object RateLimiterStatus extends Enumeration {
+  type RateLimiterStatus = Value
+
+  val Allow: RateLimiterStatus = Value("Allow")
+  val Block: RateLimiterStatus = Value("Block")
+  val Blacklist: RateLimiterStatus = Value("Blacklist")
+}

src/main/scala/RateLimiter/RateLimiters/AuthLimiter.scala

@@ -3,26 +3,24 @@ package RateLimiter.RateLimiters
 import RateLimiter.RateLimiterStorage
 import RateLimiter.Strategies.{BruteForceStrategy, DictionaryStrategy}
 
-import scala.concurrent.{ExecutionContext, Future}
-
-case class AuthLimiter(ip: String, userIdentifier: String, dictLimit: Long, dictExpiry: Long, bruteLimit: Long, bruteExpiry: Long)(implicit rateLimiterStorage: RateLimiterStorage, executionContext: ExecutionContext) extends BaseRateLimiter {
+import scala.concurrent.ExecutionContext
+
+case class AuthLimiter(
+  ip: String,
+  userIdentifier: String,
+  dictLimit: Long,
+  dictExpiry: Long,
+  dictBlacklist: Boolean,
+  bruteLimit: Long,
+  bruteExpiry: Long,
+  bruteBlacklist: Boolean
+)(implicit rateLimiterStorage: RateLimiterStorage, override val executionContext: ExecutionContext) extends StrategyRateLimiter {
 
   private final val DictIdentifier = "DictAuthLimiter"
   private final val BruteIdentifier = "BruteAuthLimiter"
 
-  private final val Strategies = List(
-    DictionaryStrategy(DictIdentifier, ip, userIdentifier, dictLimit, dictExpiry),
-    BruteForceStrategy(BruteIdentifier, ip, userIdentifier, bruteLimit, bruteExpiry)
+  protected final override def strategies = Seq(
+    DictionaryStrategy(DictIdentifier, ip, userIdentifier, dictLimit, dictExpiry, dictBlacklist),
+    BruteForceStrategy(BruteIdentifier, ip, userIdentifier, bruteLimit, bruteExpiry, bruteBlacklist)
   )
-
-  override def allow: Future[Boolean] = {
-    Future.traverse(Strategies)(strategy => strategy.allow)
-      .map(_.forall(identity))
-  }
-
-  override def increment: Future[Unit] = {
-    Future.traverse(Strategies)(strategy => strategy.increment())
-      .map(_.tail)
-  }
 }
-

src/main/scala/RateLimiter/RateLimiters/BaseRateLimiter.scala

@@ -1,8 +1,11 @@
 package RateLimiter.RateLimiters
 
+import RateLimiter.RateLimiterStatus.RateLimiterStatus
+
 import scala.concurrent.Future
 
 trait BaseRateLimiter {
-  def allow: Future[Boolean]
-  def increment: Future[Unit]
+  def status: Future[RateLimiterStatus]
+  def increment(): Future[Unit]
+  def statusWithIncrement(): Future[RateLimiterStatus]
 }

src/main/scala/RateLimiter/RateLimiters/GlobalTagLimiter.scala

@@ -0,0 +1,19 @@
+package RateLimiter.RateLimiters
+
+import RateLimiter.RateLimiterStorage
+import RateLimiter.Strategies.GlobalTagStrategy
+
+import scala.concurrent.ExecutionContext
+
+case class GlobalTagLimiter(
+  tag: String,
+  limit: Long,
+  expiry: Long
+)(implicit rateLimiterStorage: RateLimiterStorage, override val executionContext: ExecutionContext) extends StrategyRateLimiter {
+
+  private final val Identifier = "GlobalTagLimiter"
+
+  protected final override def strategies = Seq(
+    GlobalTagStrategy(Identifier, tag, limit, expiry)
+  )
+}

src/main/scala/RateLimiter/RateLimiters/IPLimiter.scala

@@ -3,18 +3,12 @@ package RateLimiter.RateLimiters
 import RateLimiter.RateLimiterStorage
 import RateLimiter.Strategies.IPStrategy
 
-import scala.concurrent.{ExecutionContext, Future}
+import scala.concurrent.ExecutionContext
 
-case class IPLimiter(ip: String, limit: Long, expiry: Long)(implicit rateLimiterStorage: RateLimiterStorage, executionContext: ExecutionContext) extends BaseRateLimiter {
-  private final val Identifier = "IPLimiter"
-
-  override def allow: Future[Boolean] = {
-    IPStrategy(Identifier, ip, limit, expiry).allow
-  }
-
-  override def increment: Future[Unit] = {
-    IPStrategy(Identifier, ip, limit, expiry).increment()
-  }
+case class IPLimiter(ip: String, limit: Long, expiry: Long, blacklistOnBlock: Boolean)(implicit rateLimiterStorage: RateLimiterStorage, override val executionContext: ExecutionContext) extends StrategyRateLimiter {
+  private final val Identifier = s"IPLimiter"
 
+  protected final override def strategies = Seq(
+    IPStrategy(Identifier, ip, limit, expiry, blacklistOnBlock)
+  )
 }
-

src/main/scala/RateLimiter/RateLimiters/StrategyRateLimiter.scala

@@ -0,0 +1,36 @@
+package RateLimiter.RateLimiters
+
+import RateLimiter.RateLimiterStatus._
+import RateLimiter.Strategies.BaseStrategy
+
+import scala.concurrent.{ExecutionContext, Future}
+
+trait StrategyRateLimiter extends BaseRateLimiter {
+  protected def strategies: Seq[BaseStrategy]
+  implicit val executionContext: ExecutionContext
+
+  override def status: Future[RateLimiterStatus] = {
+    Future
+      .traverse(strategies)(strategy => strategy.status)
+      .map(_.fold(Allow) {
+        case (Allow, status) => status
+        case (Block, status) => if (status != Allow) status else Block
+        case (Blacklist, _) => Blacklist
+      })
+  }
+
+  override def increment(): Future[Unit] = {
+    Future.traverse(strategies)(strategy => strategy.increment())
+      .map(_.tail)
+  }
+
+  // TODO: does this logic make sense, and is it intuitive? Should this logic live here?
+  override def statusWithIncrement(): Future[RateLimiterStatus] = {
+    status.map {
+      case Allow =>
+        increment()
+        Allow
+      case status => status
+    }
+  }
+}

src/main/scala/RateLimiter/RateLimiters/TagLimiter.scala

@@ -3,19 +3,18 @@ package RateLimiter.RateLimiters
 import RateLimiter.RateLimiterStorage
 import RateLimiter.Strategies.TagStrategy
 
-import scala.concurrent.{ExecutionContext, Future}
-
-case class TagLimiter(tag: String, ip: String, limit: Long, expiry: Long)(implicit rateLimiterStorage: RateLimiterStorage, executionContext: ExecutionContext) extends BaseRateLimiter {
-
+import scala.concurrent.ExecutionContext
+
+case class TagLimiter(
+  tag: String,
+  ip: String,
+  limit: Long,
+  expiry: Long,
+  blacklistOnBlock: Boolean
+)(implicit rateLimiterStorage: RateLimiterStorage, override val executionContext: ExecutionContext) extends StrategyRateLimiter {
   private final val Identifier = "TagLimiter"
 
-  override def allow: Future[Boolean] = {
-    TagStrategy(Identifier, tag, ip, limit, expiry).allow
-  }
-
-  override def increment: Future[Unit] = {
-    TagStrategy(Identifier, tag, ip, limit, expiry).increment()
-  }
-
+  protected final override def strategies = Seq(
+    TagStrategy(Identifier, tag, ip, limit, expiry, blacklistOnBlock)
+  )
 }
-

src/main/scala/RateLimiter/Strategies/BaseStrategy.scala

@@ -1,6 +1,7 @@
 package RateLimiter.Strategies
 
 import RateLimiter.RateLimiterStorage
+import RateLimiter.RateLimiterStatus._
 
 import scala.concurrent.{ExecutionContext, Future}
 
@@ -9,14 +10,18 @@ trait BaseStrategy {
   implicit def storage: RateLimiterStorage
 
   def identifier: String
-  def ip: String
   def limit: Long
   def expiry: Long
-
-  def key: String = s"$identifier:$ip"
-
-  def allow(implicit executionContext: ExecutionContext): Future[Boolean] = {
-    storage.getCount(key, expiry).map(_ < limit)
+  def key: String
+  def blacklistOnBlock: Boolean
+
+  def status(implicit executionContext: ExecutionContext): Future[RateLimiterStatus] = {
+    storage.getCount(key, expiry).map { count =>
+      println(s"CHECKING: $identifier, $count")
+      if (count < limit) Allow
+      else if (!blacklistOnBlock) Block
+      else Blacklist
+    }
   }
 
   def increment(): Future[Unit] = {

src/main/scala/RateLimiter/Strategies/BruteForceStrategy.scala

@@ -5,8 +5,8 @@ import RateLimiter.RateLimiterStorage
 /*
   Ratelimits based on number of attempts on a single user
  */
-case class BruteForceStrategy(identifier: String, ip: String, userIdentifier: String, limit: Long, expiry: Long)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
-  override def storage = rateLimiterStorage
+case class BruteForceStrategy(identifier: String, ip: String, userIdentifier: String, limit: Long, expiry: Long, blacklistOnBlock: Boolean)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
+  override implicit def storage: RateLimiterStorage = rateLimiterStorage
 
-  override def key: String = s"$identifier:$userIdentifier"
+  def key: String = s"$identifier:$userIdentifier"
 }

src/main/scala/RateLimiter/Strategies/DictionaryStrategy.scala

@@ -7,10 +7,12 @@ import scala.concurrent.Future
 /*
   Ratelimits based on a single IP attempting many different users
  */
-case class DictionaryStrategy(identifier: String, ip: String, userIdentifier: String, limit: Long, expiry: Long)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
+case class DictionaryStrategy(identifier: String, ip: String, userIdentifier: String, limit: Long, expiry: Long, blacklistOnBlock: Boolean)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
   override implicit def storage: RateLimiterStorage = rateLimiterStorage
 
-  override def increment: Future[Unit] = {
+  def key: String = s"$identifier:$ip"
+
+  override def increment(): Future[Unit] = {
     storage.incrementCount(key, userIdentifier, expiry)
   }
 

src/main/scala/RateLimiter/Strategies/GlobalTagStrategy.scala

@@ -0,0 +1,16 @@
+package RateLimiter.Strategies
+
+import RateLimiter.RateLimiterStorage
+
+/*
+  Ratelimits based on number of requests with this tag for all users
+  Can be used to ratelimit specific actions for example
+ */
+case class GlobalTagStrategy(identifier: String, tag: String, limit: Long, expiry: Long)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
+  override implicit def storage: RateLimiterStorage = rateLimiterStorage
+
+  def key: String = s"$identifier:$tag"
+
+  // Should never blacklist since that would effectively block all users
+  override def blacklistOnBlock = false
+}

src/main/scala/RateLimiter/Strategies/IPStrategy.scala

@@ -5,6 +5,8 @@ import RateLimiter.RateLimiterStorage
 /*
   Ratelimits based on number of requests for a single ip
  */
-case class IPStrategy(identifier: String, ip: String, limit: Long, expiry: Long)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
-  override def storage = rateLimiterStorage
+case class IPStrategy(identifier: String, ip: String, limit: Long, expiry: Long, blacklistOnBlock: Boolean)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
+  override implicit def storage: RateLimiterStorage = rateLimiterStorage
+
+  def key: String = s"$identifier:$ip"
 }

src/main/scala/RateLimiter/Strategies/TagStrategy.scala

@@ -6,8 +6,8 @@ import RateLimiter.RateLimiterStorage
   Ratelimits based on number of requests with this tag for a single ip.
   Can be used to ratelimit specific actions for example
  */
-case class TagStrategy(identifier: String, tag: String, ip: String, limit: Long, expiry: Long)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
+case class TagStrategy(identifier: String, tag: String, ip: String, limit: Long, expiry: Long, blacklistOnBlock: Boolean)(implicit rateLimiterStorage: RateLimiterStorage) extends BaseStrategy {
   override implicit def storage: RateLimiterStorage = rateLimiterStorage
 
-  override def key: String = s"$identifier:$tag:$ip"
+  def key: String = s"$identifier:$tag:$ip"
 }