From: TheSaminator <TheSaminator@users.noreply.github.com>
Date: Mon, 21 Feb 2022 15:11:03 +0000 (-0500)
Subject: Obey Discord rate-limits
X-Git-Url: https://gitweb.starshipfights.net/?a=commitdiff_plain;h=178b0ed0f3a5a4da1ed6a0048c40bf708ff3308e;p=starship-fights

Obey Discord rate-limits
---

diff --git a/build.gradle.kts b/build.gradle.kts
index 5e3c20c..4db65c7 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -60,6 +60,10 @@ kotlin {
 				implementation("org.slf4j:slf4j-api:1.7.32")
 				implementation("ch.qos.logback:logback-classic:1.2.10")
 				
+				implementation("io.github.instantwebp2p:tweetnacl-java:1.1.2")
+				
+				implementation("com.aventrix.jnanoid:jnanoid:2.0.0")
+				
 				implementation("org.litote.kmongo:kmongo-coroutine-serialization:4.4.0") {
 					exclude("org.jetbrains.kotlin", "kotlin-reflect")
 					
@@ -68,9 +72,8 @@ kotlin {
 					exclude("org.jetbrains.kotlinx", "kotlinx-serialization-core-jvm")
 				}
 				
+				// development only
 				implementation("de.flapdoodle.embed:de.flapdoodle.embed.mongo:3.0.0")
-				
-				implementation("com.aventrix.jnanoid:jnanoid:2.0.0")
 			}
 		}
 		
diff --git a/src/jvmMain/kotlin/starshipfights/auth/providers.kt b/src/jvmMain/kotlin/starshipfights/auth/providers.kt
index 86c74f0..23a9714 100644
--- a/src/jvmMain/kotlin/starshipfights/auth/providers.kt
+++ b/src/jvmMain/kotlin/starshipfights/auth/providers.kt
@@ -4,6 +4,7 @@ import io.ktor.application.*
 import io.ktor.auth.*
 import io.ktor.client.*
 import io.ktor.client.engine.apache.*
+import io.ktor.client.features.*
 import io.ktor.client.request.*
 import io.ktor.features.*
 import io.ktor.html.*
@@ -43,9 +44,14 @@ interface AuthProvider {
 	fun installRouting(conf: Routing)
 	
 	companion object Installer {
-		private val currentProvider: AuthProvider
+		private val newCurrentProvider: AuthProvider
 			get() = CurrentConfiguration.discordClient?.let { ProductionAuthProvider(it) } ?: TestAuthProvider
 		
+		private var cachedCurrentProvider: AuthProvider? = null
+		
+		val currentProvider: AuthProvider
+			get() = cachedCurrentProvider ?: newCurrentProvider.also { cachedCurrentProvider = it }
+		
 		fun install(into: Application) {
 			currentProvider.installApplication(into)
 			
@@ -490,7 +496,13 @@ object TestAuthProvider : AuthProvider {
 }
 
 class ProductionAuthProvider(private val discordLogin: DiscordLogin) : AuthProvider {
-	private val httpClient = HttpClient(Apache)
+	val httpClient = HttpClient(Apache) {
+		install(UserAgent) {
+			agent = discordLogin.userAgent
+		}
+		
+		install(RateLimit)
+	}
 	
 	override fun installAuth(conf: Authentication.Configuration) {
 		conf.oauth("auth-oauth-discord") {
diff --git a/src/jvmMain/kotlin/starshipfights/auth/ratelimit.kt b/src/jvmMain/kotlin/starshipfights/auth/ratelimit.kt
new file mode 100644
index 0000000..987c3d5
--- /dev/null
+++ b/src/jvmMain/kotlin/starshipfights/auth/ratelimit.kt
@@ -0,0 +1,69 @@
+package starshipfights.auth
+
+import io.ktor.client.*
+import io.ktor.client.call.*
+import io.ktor.client.features.*
+import io.ktor.client.request.*
+import io.ktor.client.statement.*
+import io.ktor.http.*
+import io.ktor.util.*
+import kotlinx.coroutines.delay
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.json.Json
+import kotlin.math.roundToLong
+
+class RateLimit(
+	val jsonCodec: Json,
+	val remainingHeader: String,
+	val resetAfterHeader: String,
+) {
+	class Config {
+		var jsonCodec: Json = Json.Default
+		var remainingHeader: String = "X-RateLimit-Remaining"
+		var resetAfterHeader: String = "X-RateLimit-Reset-After"
+	}
+	
+	private var remainingRequests = -1
+	private var resetAfter = 0.0
+	
+	companion object Feature : HttpClientFeature<Config, RateLimit> {
+		override val key: AttributeKey<RateLimit> = AttributeKey("RateLimit")
+		override fun prepare(block: Config.() -> Unit): RateLimit = Config().apply(block).run {
+			RateLimit(jsonCodec, remainingHeader, resetAfterHeader)
+		}
+		
+		override fun install(feature: RateLimit, scope: HttpClient) {
+			scope.requestPipeline.intercept(HttpRequestPipeline.Before) {
+				feature.remainingRequests.takeIf { it >= 0 }?.let { remaining ->
+					delay((feature.resetAfter * 1000 / (remaining + 1)).roundToLong())
+				}
+			}
+			
+			scope.responsePipeline.intercept(HttpResponsePipeline.Receive) {
+				if (context.response.status == HttpStatusCode.TooManyRequests) {
+					feature.remainingRequests = 0
+					val jsonBody = context.response.receive<String>()
+					val rateLimitedResponse = feature.jsonCodec.decodeFromString(RateLimitedResponse.serializer(), jsonBody)
+					feature.resetAfter = rateLimitedResponse.retryAfter
+				} else {
+					context.response.headers[feature.remainingHeader]?.toIntOrNull()?.let {
+						feature.remainingRequests = it
+					}
+					
+					context.response.headers[feature.resetAfterHeader]?.toDoubleOrNull()?.let {
+						feature.resetAfter = it
+					}
+				}
+			}
+		}
+	}
+}
+
+@Serializable
+data class RateLimitedResponse(
+	val message: String,
+	@SerialName("retry_after")
+	val retryAfter: Double,
+	val global: Boolean
+)
diff --git a/src/jvmMain/kotlin/starshipfights/server_conf.kt b/src/jvmMain/kotlin/starshipfights/server_conf.kt
index c0a5883..fcc4097 100644
--- a/src/jvmMain/kotlin/starshipfights/server_conf.kt
+++ b/src/jvmMain/kotlin/starshipfights/server_conf.kt
@@ -21,8 +21,9 @@ data class Configuration(
 
 @Serializable
 data class DiscordLogin(
+	val userAgent: String,
+	
 	val clientId: String,
-	//val clientPubKey: String,
 	val clientSecret: String,
 	
 	val ownerId: String,