User validation for every request

Implemented a separate class for db connection, for the ease of accessing queries

Signed-off-by: Rahul Rudragoudar <rr83019@gmail.com>

src/main/scala/lc/Main.scala

@@ -19,18 +19,32 @@ import java.util.concurrent._
 import java.util.UUID
 import scala.Array
 
-class Captcha(throttle: Int) {
+class DBConn(){
   val con: Connection = DriverManager.getConnection("jdbc:h2:./captcha", "sa", "")
-  val stmt: Statement = con.createStatement()
+
-  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 insertPstmt: PreparedStatement = con.prepareStatement("INSERT INTO challenge(token, id, secret, provider, contentType, image) VALUES (?, ?, ?, ?, ?, ?)")
   val mapPstmt: PreparedStatement = con.prepareStatement("INSERT INTO mapId(uuid, token) VALUES (?, ?)")
   val selectPstmt: PreparedStatement = con.prepareStatement("SELECT secret, provider FROM challenge WHERE token = ?")
   val imagePstmt: PreparedStatement = con.prepareStatement("SELECT image FROM challenge c, mapId m WHERE c.token=m.token AND m.uuid = ?")
   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 = ?)")
   val userPstmt: PreparedStatement = con.prepareStatement("INSERT INTO users(email, hash) VALUES (?,?)")
+  val validatePstmt: PreparedStatement = con.prepareStatement("SELECT hash FROM users WHERE hash = ? LIMIT 1")
+
+  def getConn(): Statement = {
+    con.createStatement()
+  }
+
+  def closeConnection(): Unit = {
+	  con.close()
+  } 
+}
+
+class Captcha(throttle: Int) extends DBConn {
+  
+  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,
@@ -151,10 +165,6 @@ class Captcha(throttle: Int) {
       println(s"${token}\t\t${id}\t\t${secret}\t\t${solved}")
     }
   }
-  
-  def closeConnection(): Unit = {
-	  con.close()
-  } 
 }
 
 case class Size(height: Int, width: Int)
@@ -163,18 +173,44 @@ case class Id(id: String)
 case class Answer(answer: String, id: String)
 case class Secret(token: Int)
 
+class RateLimiter extends DBConn {
+  val stmt = getConn()
+  val userActive = collection.mutable.Map[Int, Int]()
+
+  def validateUser(user: Int) : Boolean = {
+    validatePstmt.setInt(1, user)
+    val rs = validatePstmt.executeQuery()
+    val validated = if(rs.next()){
+      val hash = rs.getInt("hash")
+      userActive(hash) = 0
+      true
+    } else {
+      false
+    }
+    validated
+  }
+}
+
 class Server(port: Int){
 	val captcha = new Captcha(0)
+  val rateLimiter = new RateLimiter()
 	val server = new HTTPServer(port)
 	val host = server.getVirtualHost(null)
 
 	implicit val formats = DefaultFormats
 
 	host.addContext("/v1/captcha",(req, resp) => {
-    	val body = req.getJson()
+      val accessToken = if(req.getHeaders().get("access-token") != null){
-    	val json = parse(body)
+        req.getHeaders().get("access-token").toInt
-    	val param = json.extract[Parameters]
+      } else 0
-    	val id = captcha.getChallenge(param)
+      val id = if(true == rateLimiter.validateUser(accessToken)){
+        val body = req.getJson()
+      	val json = parse(body)
+      	val param = json.extract[Parameters]
+        captcha.getChallenge(param)
+      } else {
+        "Not a valid user! Please register."
+      }
     	resp.getHeaders().add("Content-Type","application/json")
     	resp.send(200, write(id))
     	0