feat: Add rate limiting to routes

ReVanced · Jun 6, 2024 · 4a3a872 · 4a3a872
1 parent 19a7085
commit 4a3a872
Show file tree

Hide file tree

Showing 7 changed files with 104 additions and 71 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -62,9 +62,9 @@ dependencies {
     implementation(libs.ktor.server.auth.jwt)
     implementation(libs.ktor.server.cors)
     implementation(libs.ktor.server.caching.headers)
+    implementation(libs.ktor.server.rate.limit)
     implementation(libs.ktor.server.host.common)
     implementation(libs.ktor.server.jetty)
-    implementation(libs.ktor.server.conditional.headers)
     implementation(libs.ktor.serialization.kotlinx.json)
     implementation(libs.koin.ktor)
     implementation(libs.h2)

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -20,13 +20,13 @@ ktor-client-okhttp = { module = "io.ktor:ktor-client-okhttp" }
 ktor-client-resources = { module = "io.ktor:ktor-client-resources" }
 ktor-client-auth = { module = "io.ktor:ktor-client-auth" }
 ktor-client-content-negotiation = { module = "io.ktor:ktor-client-content-negotiation" }
-ktor-server-conditional-headers = { module = "io.ktor:ktor-server-conditional-headers" }
 ktor-server-core = { module = "io.ktor:ktor-server-core" }
 ktor-server-content-negotiation = { module = "io.ktor:ktor-server-content-negotiation" }
 ktor-server-auth = { module = "io.ktor:ktor-server-auth" }
 ktor-server-auth-jwt = { module = "io.ktor:ktor-server-auth-jwt" }
 ktor-server-cors = { module = "io.ktor:ktor-server-cors" }
 ktor-server-caching-headers = { module = "io.ktor:ktor-server-caching-headers" }
+ktor-server-rate-limit = { module = "io.ktor:ktor-server-rate-limit" }
 ktor-server-host-common = { module = "io.ktor:ktor-server-host-common" }
 ktor-server-jetty = { module = "io.ktor:ktor-server-jetty" }
 ktor-serialization-kotlinx-json = { module = "io.ktor:ktor-serialization-kotlinx-json" }

diff --git a/src/main/kotlin/app/revanced/api/configuration/HTTP.kt b/src/main/kotlin/app/revanced/api/configuration/HTTP.kt
@@ -3,15 +3,15 @@ package app.revanced.api.configuration
 import io.ktor.http.*
 import io.ktor.http.content.*
 import io.ktor.server.application.*
+import io.ktor.server.plugins.*
 import io.ktor.server.plugins.cachingheaders.*
-import io.ktor.server.plugins.conditionalheaders.*
 import io.ktor.server.plugins.cors.routing.*
+import io.ktor.server.plugins.ratelimit.*
 import kotlin.time.Duration.Companion.minutes
 
 fun Application.configureHTTP(
     allowedHost: String,
 ) {
-    install(ConditionalHeaders)
     install(CORS) {
         allowMethod(HttpMethod.Options)
         allowMethod(HttpMethod.Put)
@@ -23,4 +23,14 @@ fun Application.configureHTTP(
     install(CachingHeaders) {
         options { _, _ -> CachingOptions(CacheControl.MaxAge(maxAgeSeconds = 5.minutes.inWholeSeconds.toInt())) }
     }
+    install(RateLimit) {
+        register(RateLimitName("weak")) {
+            rateLimiter(limit = 30, refillPeriod = 2.minutes)
+            requestKey { it.request.origin.remoteAddress }
+        }
+        register(RateLimitName("strong")) {
+            rateLimiter(limit = 5, refillPeriod = 1.minutes)
+            requestKey { it.request.origin.remoteHost }
+        }
+    }
 }
diff --git a/src/main/kotlin/app/revanced/api/configuration/routing/routes/Announcements.kt b/src/main/kotlin/app/revanced/api/configuration/routing/routes/Announcements.kt
@@ -6,6 +6,7 @@ import app.revanced.api.configuration.schema.APIAnnouncementArchivedAt
 import app.revanced.api.configuration.services.AnnouncementService
 import io.ktor.server.application.*
 import io.ktor.server.auth.*
+import io.ktor.server.plugins.ratelimit.*
 import io.ktor.server.request.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
@@ -15,69 +16,78 @@ import org.koin.ktor.ext.get as koinGet
 internal fun Route.announcementsRoute() = route("announcements") {
     val announcementService = koinGet<AnnouncementService>()
 
-    route("{channel}/latest") {
-        get("id") {
-            val channel: String by call.parameters
+    rateLimit(RateLimitName("weak")) {
+        route("{channel}/latest") {
+            get("id") {
+                val channel: String by call.parameters
 
-            call.respondOrNotFound(announcementService.latestId(channel))
-        }
+                call.respondOrNotFound(announcementService.latestId(channel))
+            }
 
-        get {
-            val channel: String by call.parameters
+            get {
+                val channel: String by call.parameters
 
-            call.respondOrNotFound(announcementService.latest(channel))
+                call.respondOrNotFound(announcementService.latest(channel))
+            }
         }
     }
 
-    get("{channel}") {
-        val channel: String by call.parameters
+    rateLimit(RateLimitName("strong")) {
+        get("{channel}") {
+            val channel: String by call.parameters
 
-        call.respond(announcementService.all(channel))
+            call.respond(announcementService.all(channel))
+        }
     }
-
-    route("latest") {
-        get("id") {
-            call.respondOrNotFound(announcementService.latestId())
+    rateLimit(RateLimitName("strong")) {
+        route("latest") {
+            get("id") {
+                call.respondOrNotFound(announcementService.latestId())
+            }
+
+            get {
+                call.respondOrNotFound(announcementService.latest())
+            }
         }
+    }
 
+    rateLimit(RateLimitName("strong")) {
         get {
-            call.respondOrNotFound(announcementService.latest())
+            call.respond(announcementService.all())
         }
     }
 
-    get {
-        call.respond(announcementService.all())
-    }
+    rateLimit(RateLimitName("strong")) {
+        authenticate("jwt") {
+            post {
+                announcementService.new(call.receive<APIAnnouncement>())
+            }
 
-    authenticate("jwt") {
-        post {
-            announcementService.new(call.receive<APIAnnouncement>())
-        }
+            post("{id}/archive") {
+                val id: Int by call.parameters
+                val archivedAt = call.receiveNullable<APIAnnouncementArchivedAt>()?.archivedAt
 
-        post("{id}/archive") {
-            val id: Int by call.parameters
-            val archivedAt = call.receiveNullable<APIAnnouncementArchivedAt>()?.archivedAt
+                announcementService.archive(id, archivedAt)
+            }
 
-            announcementService.archive(id, archivedAt)
-        }
+            post("{id}/unarchive") {
+                val id: Int by call.parameters
 
-        post("{id}/unarchive") {
-            val id: Int by call.parameters
+                announcementService.unarchive(id)
+            }
 
-            announcementService.unarchive(id)
-        }
-
-        patch("{id}") {
-            val id: Int by call.parameters
-            val announcement = call.receive<APIAnnouncement>()
+            patch("{id}") {
+                val id: Int by call.parameters
+                val announcement = call.receive<APIAnnouncement>()
 
-            announcementService.update(id, announcement)
-        }
+                announcementService.update(id, announcement)
+            }
 
-        delete("{id}") {
-            val id: Int by call.parameters
+            delete("{id}") {
+                val id: Int by call.parameters
 
-            announcementService.delete(id)
+                announcementService.delete(id)
+            }
         }
     }
 }
diff --git a/src/main/kotlin/app/revanced/api/configuration/routing/routes/ApiRoute.kt b/src/main/kotlin/app/revanced/api/configuration/routing/routes/ApiRoute.kt
@@ -7,6 +7,7 @@ import io.ktor.http.*
 import io.ktor.server.application.*
 import io.ktor.server.auth.*
 import io.ktor.server.http.content.*
+import io.ktor.server.plugins.ratelimit.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
 import org.koin.ktor.ext.get
@@ -15,12 +16,20 @@ internal fun Route.rootRoute() {
     val apiService = get<ApiService>()
     val authService = get<AuthService>()
 
-    get("contributors") {
-        call.respond(apiService.contributors())
-    }
+    rateLimit(RateLimitName("strong")) {
+        authenticate("basic") {
+            get("token") {
+                call.respond(authService.newToken())
+            }
+        }
 
-    get("team") {
-        call.respond(apiService.team())
+        get("contributors") {
+            call.respond(apiService.contributors())
+        }
+
+        get("team") {
+            call.respond(apiService.team())
+        }
     }
 
     route("ping") {
@@ -29,18 +38,14 @@ internal fun Route.rootRoute() {
         }
     }
 
-    get("backend/rate_limit") {
-        call.respondOrNotFound(apiService.rateLimit())
-    }
-
-    authenticate("basic") {
-        get("token") {
-            call.respond(authService.newToken())
+    rateLimit(RateLimitName("weak")) {
+        get("backend/rate_limit") {
+            call.respondOrNotFound(apiService.rateLimit())
         }
-    }
 
-    staticResources("/", "/app/revanced/api/static") {
-        contentType { ContentType.Application.Json }
-        extensions("json")
+        staticResources("/", "/app/revanced/api/static") {
+            contentType { ContentType.Application.Json }
+            extensions("json")
+        }
     }
 }
diff --git a/src/main/kotlin/app/revanced/api/configuration/routing/routes/OldApi.kt b/src/main/kotlin/app/revanced/api/configuration/routing/routes/OldApi.kt
@@ -2,15 +2,18 @@ package app.revanced.api.configuration.routing.routes
 
 import app.revanced.api.configuration.services.OldApiService
 import io.ktor.server.application.*
+import io.ktor.server.plugins.ratelimit.*
 import io.ktor.server.routing.*
 import org.koin.ktor.ext.get
 
 internal fun Route.oldApiRoute() {
     val oldApiService = get<OldApiService>()
 
-    route(Regex("/(v2|tools|contributors).*")) {
-        handle {
-            oldApiService.proxy(call)
+    rateLimit(RateLimitName("weak")) {
+        route(Regex("/(v2|tools|contributors).*")) {
+            handle {
+                oldApiService.proxy(call)
+            }
         }
     }
 }
diff --git a/src/main/kotlin/app/revanced/api/configuration/routing/routes/PatchesRoute.kt b/src/main/kotlin/app/revanced/api/configuration/routing/routes/PatchesRoute.kt
@@ -3,6 +3,7 @@ package app.revanced.api.configuration.routing.routes
 import app.revanced.api.configuration.services.PatchesService
 import io.ktor.http.*
 import io.ktor.server.application.*
+import io.ktor.server.plugins.ratelimit.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
 import org.koin.ktor.ext.get as koinGet
@@ -11,16 +12,20 @@ internal fun Route.patchesRoute() = route("patches") {
     val patchesService = koinGet<PatchesService>()
 
     route("latest") {
-        get {
-            call.respond(patchesService.latestRelease())
-        }
+        rateLimit(RateLimitName("weak")) {
+            get {
+                call.respond(patchesService.latestRelease())
+            }
 
-        get("version") {
-            call.respond(patchesService.latestVersion())
+            get("version") {
+                call.respond(patchesService.latestVersion())
+            }
         }
 
-        get("list") {
-            call.respondBytes(ContentType.Application.Json) { patchesService.list() }
+        rateLimit(RateLimitName("strong")) {
+            get("list") {
+                call.respondBytes(ContentType.Application.Json) { patchesService.list() }
+            }
         }
     }
 }