From 1a57942a0bc9f4d2ed23140d5f3d66d6a8ac939f Mon Sep 17 00:00:00 2001 From: Rahul Rudragoudar Date: Thu, 13 Jun 2019 00:48:29 +0530 Subject: [PATCH 1/3] User validation for every request Implemented a separate class for db connection, for the ease of accessing queries Signed-off-by: Rahul Rudragoudar --- src/main/scala/lc/Main.scala | 62 ++++++++++++++++++++++++++++-------- 1 file changed, 49 insertions(+), 13 deletions(-) diff --git a/src/main/scala/lc/Main.scala b/src/main/scala/lc/Main.scala index 18e4753..4b8a236 100644 --- a/src/main/scala/lc/Main.scala +++ b/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 json = parse(body) - val param = json.extract[Parameters] - val id = captcha.getChallenge(param) + val accessToken = if(req.getHeaders().get("access-token") != null){ + req.getHeaders().get("access-token").toInt + } else 0 + 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 From a63cf3976d90a0e508318bba068cb73f092d94cb Mon Sep 17 00:00:00 2001 From: Rahul Rudragoudar Date: Mon, 17 Jun 2019 00:20:18 +0530 Subject: [PATCH 2/3] Add rate limiter Add user validation for every request Signed-off-by: Rahul Rudragoudar --- src/main/scala/lc/Main.scala | 48 ++++++++++++++++++++++++++++-------- 1 file changed, 38 insertions(+), 10 deletions(-) diff --git a/src/main/scala/lc/Main.scala b/src/main/scala/lc/Main.scala index 4b8a236..e53f46b 100644 --- a/src/main/scala/lc/Main.scala +++ b/src/main/scala/lc/Main.scala @@ -175,20 +175,48 @@ case class Secret(token: Int) class RateLimiter extends DBConn { val stmt = getConn() - val userActive = collection.mutable.Map[Int, Int]() + val userLastActive = collection.mutable.Map[Int, Long]() + val userAllowance = collection.mutable.Map[Int, Double]() + val rate = 2.0 + val per = 45.0 + val allowance = rate 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 + val allow = if(userLastActive.contains(user)){ true } else { - false + 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 } - validated + allow } + + def checkLimit(user: Int): Boolean = { + synchronized { + val current = System.currentTimeMillis() + val time_passed = (current - userLastActive(user)) / 1000000000 + 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 + } + } + } class Server(port: Int){ @@ -203,13 +231,13 @@ class Server(port: Int){ val accessToken = if(req.getHeaders().get("access-token") != null){ req.getHeaders().get("access-token").toInt } else 0 - val id = if(true == rateLimiter.validateUser(accessToken)){ + val id = if(true == rateLimiter.validateUser(accessToken) && true == rateLimiter.checkLimit(accessToken)){ val body = req.getJson() val json = parse(body) val param = json.extract[Parameters] captcha.getChallenge(param) } else { - "Not a valid user! Please register." + "Not a valid user or rate limit reached!" } resp.getHeaders().add("Content-Type","application/json") resp.send(200, write(id)) From d65d05003835cf301db45a251da365f6f7d06941 Mon Sep 17 00:00:00 2001 From: Rahul Rudragoudar Date: Thu, 20 Jun 2019 11:45:29 +0530 Subject: [PATCH 3/3] Minor fixes Signed-off-by: Rahul Rudragoudar --- src/main/scala/lc/Main.scala | 33 +++++++++++++++++---------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/src/main/scala/lc/Main.scala b/src/main/scala/lc/Main.scala index e53f46b..e37c75e 100644 --- a/src/main/scala/lc/Main.scala +++ b/src/main/scala/lc/Main.scala @@ -182,22 +182,24 @@ class RateLimiter extends DBConn { val allowance = rate 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 + synchronized { + val allow = if(userLastActive.contains(user)){ true } else { - false + 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 } - validated + allow } - allow } def checkLimit(user: Int): Boolean = { @@ -228,10 +230,9 @@ class Server(port: Int){ implicit val formats = DefaultFormats host.addContext("/v1/captcha",(req, resp) => { - val accessToken = if(req.getHeaders().get("access-token") != null){ - req.getHeaders().get("access-token").toInt - } else 0 - val id = if(true == rateLimiter.validateUser(accessToken) && true == rateLimiter.checkLimit(accessToken)){ + val accessToken = Option(req.getHeaders().get("access-token")).map(_.toInt) + val access = accessToken.map(t => rateLimiter.validateUser(t) && rateLimiter.checkLimit(t)).getOrElse(false) + val id = if(access){ val body = req.getJson() val json = parse(body) val param = json.extract[Parameters]