Merge pull request #49 from librecaptcha/cleanup

Code Cleanup

build.sbt

@@ -2,16 +2,16 @@ lazy val root = (project in file(".")).
   settings(
     inThisBuild(List(
       organization := "com.example",
-      scalaVersion := "2.13.2",
+      scalaVersion := "2.13.3",
       version      := "0.1.0-SNAPSHOT")),
     name := "LibreCaptcha",
 
     libraryDependencies += "com.sksamuel.scrimage" % "scrimage-core" % "4.0.5",
 
     libraryDependencies += "com.sksamuel.scrimage" % "scrimage-filters" % "4.0.5",
-    
+
     libraryDependencies += "org.json4s" % "json4s-jackson_2.13" % "3.6.9"
-   
+
 )
 
 unmanagedResourceDirectories in Compile += {baseDirectory.value / "lib"}

project/build.properties

@@ -1 +1 @@
-sbt.version=1.3.12
+sbt.version=1.3.13

src/main/scala/lc/DB.scala

@@ -5,15 +5,7 @@ import java.sql._
 class DBConn(){
   val con: Connection = DriverManager.getConnection("jdbc:h2:./captcha", "sa", "")
 
-  lazy val insertPstmt: PreparedStatement = con.prepareStatement("INSERT INTO challenge(token, id, secret, provider, contentType, image) VALUES (?, ?, ?, ?, ?, ?)")
+  def getStatement(): Statement = {
-  lazy val mapPstmt: PreparedStatement = con.prepareStatement("INSERT INTO mapId(uuid, token) VALUES (?, ?)")
-  lazy val selectPstmt: PreparedStatement = con.prepareStatement("SELECT secret, provider FROM challenge WHERE token = (SELECT m.token FROM mapId m, challenge c WHERE m.token=c.token AND m.uuid = ?)")
-  lazy val imagePstmt: PreparedStatement = con.prepareStatement("SELECT image FROM challenge c, mapId m WHERE c.token=m.token AND m.uuid = ?")
-  lazy val updatePstmt: PreparedStatement = con.prepareStatement("UPDATE challenge SET solved = True WHERE token = (SELECT m.token FROM mapId m, challenge c WHERE m.token=c.token AND m.uuid = ?)")
-  lazy val userPstmt: PreparedStatement = con.prepareStatement("INSERT INTO users(email, hash) VALUES (?,?)")
-  lazy val validatePstmt: PreparedStatement = con.prepareStatement("SELECT hash FROM users WHERE hash = ? LIMIT 1")
-
-  def getConn(): Statement = {
     con.createStatement()
   }
 

src/main/scala/lc/Main.scala

@@ -14,13 +14,7 @@ case class Answer(answer: String, id: String)
 
 case class ProviderSecret(provider: String, secret: String)
 
-class Captcha(throttle: Int) extends DBConn {
+object CaptchaProviders {
-
-  val stmt = getConn()
-  stmt.execute("CREATE TABLE IF NOT EXISTS challenge(token varchar, id varchar, secret varchar, provider varchar, contentType varchar, image blob, solved boolean default False, PRIMARY KEY(token))")
-  stmt.execute("CREATE TABLE IF NOT EXISTS mapId(uuid varchar, token varchar, PRIMARY KEY(uuid), FOREIGN KEY(token) REFERENCES challenge(token))")
-  stmt.execute("CREATE TABLE IF NOT EXISTS users(email varchar, hash int)")
-
   val providers = Map(
     "FilterChallenge" -> new FilterChallenge,
     "FontFunCaptcha" -> new FontFunCaptcha,
@@ -30,6 +24,28 @@ class Captcha(throttle: Int) extends DBConn {
     "LabelCaptcha" -> new LabelCaptcha
     )
 
+  def generateChallengeSamples() = {
+    providers.map {case (key, provider) =>
+      (key, provider.returnChallenge())
+    }
+  }
+}
+
+class Captcha(throttle: Int, dbConn: DBConn) {
+  import CaptchaProviders._
+
+  private val stmt = dbConn.getStatement()
+  stmt.execute("CREATE TABLE IF NOT EXISTS challenge(token varchar, id varchar, secret varchar, provider varchar, contentType varchar, image blob, solved boolean default False, PRIMARY KEY(token))")
+  stmt.execute("CREATE TABLE IF NOT EXISTS mapId(uuid varchar, token varchar, PRIMARY KEY(uuid), FOREIGN KEY(token) REFERENCES challenge(token))")
+  stmt.execute("CREATE TABLE IF NOT EXISTS users(email varchar, hash int)")
+
+  private val insertPstmt = dbConn.con.prepareStatement("INSERT INTO challenge(token, id, secret, provider, contentType, image) VALUES (?, ?, ?, ?, ?, ?)")
+  private val mapPstmt = dbConn.con.prepareStatement("INSERT INTO mapId(uuid, token) VALUES (?, ?)")
+  private val selectPstmt = dbConn.con.prepareStatement("SELECT secret, provider FROM challenge WHERE token = (SELECT m.token FROM mapId m, challenge c WHERE m.token=c.token AND m.uuid = ?)")
+  private val imagePstmt = dbConn.con.prepareStatement("SELECT image FROM challenge c, mapId m WHERE c.token=m.token AND m.uuid = ?")
+  private val updatePstmt = dbConn.con.prepareStatement("UPDATE challenge SET solved = True WHERE token = (SELECT m.token FROM mapId m, challenge c WHERE m.token=c.token AND m.uuid = ?)")
+  private val userPstmt = dbConn.con.prepareStatement("INSERT INTO users(email, hash) VALUES (?,?)")
+
   def getProvider(): String = {
     val random = new scala.util.Random
     val keys = providers.keys
@@ -57,12 +73,6 @@ class Captcha(throttle: Int) extends DBConn {
     imageOpt
   }
 
-  def generateChallengeSamples() = {
-    providers.map {case (key, provider) =>
-      (key, provider.returnChallenge())
-    }
-  }
-
   private val uniqueIntCount = new AtomicInteger()
 
   def generateChallenge(param: Parameters): String = {
@@ -167,8 +177,9 @@ class Captcha(throttle: Int) extends DBConn {
 
 object LCFramework{
   def main(args: scala.Array[String]) {
-  	val captcha = new Captcha(2)
+    val dbConn = new DBConn()
-    val server = new Server(8888)
+  	val captcha = new Captcha(2, dbConn)
+    val server = new Server(8888, captcha, dbConn)
     captcha.beginThread(2)
     server.start()
   }
@@ -176,8 +187,7 @@ object LCFramework{
 
 object MakeSamples {
   def main(args: scala.Array[String]) {
-    val captcha = new Captcha(2)
+    val samples = CaptchaProviders.generateChallengeSamples()
-    val samples = captcha.generateChallengeSamples()
     samples.foreach {case (key, sample) =>
       val extensionMap = Map("image/png" -> "png", "image/gif" -> "gif")
       println(key + ": " + sample)

src/main/scala/lc/Server.scala

@@ -9,57 +9,62 @@ import lc.HTTPServer._
 
 case class Secret(token: Int)
 
-class RateLimiter extends DBConn {
+class RateLimiter(dbConn: DBConn) {
-  val stmt = getConn()
+  private val userLastActive = collection.mutable.Map[Int, Long]()
-  val userLastActive = collection.mutable.Map[Int, Long]()
+  private val userAllowance = collection.mutable.Map[Int, Double]()
-  val userAllowance = collection.mutable.Map[Int, Double]()
+  private val rate = 800000.0
-  val rate = 8.0
+  private val per = 45.0
-  val per = 45.0
+  private val allowance = rate
-  val allowance = rate
 
-  def validateUser(user: Int) : Boolean = {
+  private val validatePstmt = dbConn.con.prepareStatement("SELECT hash FROM users WHERE hash = ? LIMIT 1")
-    synchronized {
+
-      val allow = if(userLastActive.contains(user)){
+  private def validateUser(user: Int) : Boolean = {
+    val allow = if(userLastActive.contains(user)){
+      true
+    } else {
+      validatePstmt.setInt(1, user)
+      val rs = validatePstmt.executeQuery()
+      val validated = if(rs.next()){
+        val hash = rs.getInt("hash")
+        userLastActive(hash) = System.currentTimeMillis()
+        userAllowance(hash) = allowance
         true
       } else {
-        validatePstmt.setInt(1, user)
-        val rs = validatePstmt.executeQuery()
-        val validated = if(rs.next()){
-          val hash = rs.getInt("hash")
-          userLastActive(hash) = System.currentTimeMillis()
-          userAllowance(hash) = allowance
-          true
-        } else {
-          false
-        }
-        validated
-      }
-      allow
-    }
-  }
-
-  def checkLimit(user: Int): Boolean = {
-    synchronized {
-      val current = System.currentTimeMillis()
-      val time_passed = (current - userLastActive(user)) / 1000
-      userLastActive(user) = current
-      userAllowance(user) += time_passed * (rate/per)
-      if(userAllowance(user) > rate){ userAllowance(user) = rate }
-      val allow = if(userAllowance(user) < 1.0){
         false
-      } else {
-        userAllowance(user) -= 1.0
-        true
       }
-      allow
+      validated
     }
+    allow
   }
 
+  private def checkLimit(user: Int): Boolean = {
+    val current = System.currentTimeMillis()
+    val time_passed = (current - userLastActive(user)) / 1000
+    userLastActive(user) = current
+    userAllowance(user) += time_passed * (rate/per)
+    if(userAllowance(user) > rate){ userAllowance(user) = rate }
+    val allow = if(userAllowance(user) < 1.0){
+      false
+    } else {
+      userAllowance(user) -= 1.0
+      true
+    }
+    allow
+  }
+
+  def checkUserAccess(token: Int) : Boolean = {
+    synchronized {
+      if (validateUser(token)) {
+        return checkLimit(token)
+      } else {
+        return false
+      }
+    }
+  }
 }
 
-class Server(port: Int){
+class Server(port: Int, captcha: Captcha, dbConn: DBConn){
-	val captcha = new Captcha(0)
+  val rateLimiter = new RateLimiter(dbConn)
-  val rateLimiter = new RateLimiter()
 	val server = new HTTPServer(port)
 	val host = server.getVirtualHost(null)
 
@@ -67,7 +72,7 @@ class Server(port: Int){
 
 	host.addContext("/v1/captcha",(req, resp) => {
       val accessToken = Option(req.getHeaders().get("access-token")).map(_.toInt)
-      val access = accessToken.map(t => rateLimiter.validateUser(t) && rateLimiter.checkLimit(t)).getOrElse(false)
+      val access = accessToken.map(rateLimiter.checkUserAccess).getOrElse(false)
       if(access){
         val body = req.getJson()
       	val json = parse(body)