From 7f8107c9a730a07927f7208d0ce41928863f06c0 Mon Sep 17 00:00:00 2001
From: videah <videah@selfish.systems>
Date: Tue, 16 Jan 2024 10:08:28 -0700
Subject: [PATCH] add rate limit response and reuse http client connections

---
 lib/auth.dart           | 13 +++++++++++--
 pubspec.lock            |  2 +-
 pubspec.yaml            |  1 +
 routes/_middleware.dart | 19 ++++++++++++++++++-
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/lib/auth.dart b/lib/auth.dart
index 1ee6f5d..f37b81f 100644
--- a/lib/auth.dart
+++ b/lib/auth.dart
@@ -3,12 +3,16 @@ import 'dart:convert';
 import 'package:bluesky/bluesky.dart' as bsky;
 import 'package:dart_frog/dart_frog.dart';
 import 'package:dart_jsonwebtoken/dart_jsonwebtoken.dart';
+import 'package:http/http.dart' as http;
 import 'package:sky_bridge/crypto.dart';
 import 'package:sky_bridge/database.dart';
 import 'package:sky_bridge/models/oauth/oauth_access_token.dart';
 import 'package:sky_bridge/models/preferences.dart';
 import 'package:sky_bridge/src/generated/prisma/prisma_client.dart';
 
+/// Global client
+final httpClient = http.Client();
+
 /// Takes an IP address and ensures that authentication attempts are not
 /// being made too frequently. Bluesky now has quite aggressive rate limiting
 /// for authentication attempts, so we have to do the same.
@@ -157,7 +161,7 @@ Future<bsky.Session?> sessionFromContext(RequestContext context) async {
 
         // Credentials are just straight up invalid. Bail.
         if (newSession == null) {
-          incrementFailedAuthAttempt(context);
+          await incrementFailedAuthAttempt(context);
           return null;
         }
 
@@ -170,6 +174,7 @@ Future<bsky.Session?> sessionFromContext(RequestContext context) async {
         // try to refresh the session.
         final refreshedSession = await bsky.refreshSession(
           refreshJwt: session.refreshJwt,
+          mockedPostClient: httpClient.post,
         );
 
         // Update the session in the database.
@@ -224,7 +229,10 @@ SkybridgePreferences preferencesFromContext(RequestContext context) {
 Future<bsky.Bluesky?> blueskyFromContext(RequestContext context) async {
   final session = await sessionFromContext(context);
   if (session == null) return null;
-  return bsky.Bluesky.fromSession(session);
+  return bsky.Bluesky.fromSession(session,
+    mockedGetClient: httpClient.get,
+    mockedPostClient: httpClient.post,
+  );
 }
 
 /// A helper function to return a 401 response for an invalid bearer token.
@@ -262,6 +270,7 @@ Future<bsky.Session?> createBlueskySession({
     final session = await bsky.createSession(
       identifier: identifier,
       password: appPassword,
+      mockedPostClient: httpClient.post,
     );
 
     // If we've gotten this far, the credentials are valid.
diff --git a/pubspec.lock b/pubspec.lock
index 1908eac..f8fa80e 100644
--- a/pubspec.lock
+++ b/pubspec.lock
@@ -458,7 +458,7 @@ packages:
     source: hosted
     version: "0.15.4"
   http:
-    dependency: "direct overridden"
+    dependency: "direct main"
     description:
       name: http
       sha256: "5895291c13fa8a3bd82e76d5627f69e0d85ca6a30dcac95c4ea19a5d555879c2"
diff --git a/pubspec.yaml b/pubspec.yaml
index 412c18c..1e701e3 100644
--- a/pubspec.yaml
+++ b/pubspec.yaml
@@ -18,6 +18,7 @@ dependencies:
   dcli_core: ^3.0.3
   dotenv: ^4.1.0
   html: ^0.15.3
+  http: ^0.13.6
   image_compression: ^1.0.4
   json_annotation: ^4.8.0
   orm: ^3.4.4
diff --git a/routes/_middleware.dart b/routes/_middleware.dart
index 30440d5..eeddca3 100644
--- a/routes/_middleware.dart
+++ b/routes/_middleware.dart
@@ -11,7 +11,6 @@ Handler middleware(Handler handler) {
   return (RequestContext context) async {
     try {
       final response = await handler(context);
-
       // Reset Cache-Control so that we don't cache responses.
       return response.copyWith(
         headers: Map.of(response.headers)
@@ -53,6 +52,24 @@ Handler middleware(Handler handler) {
             },
           ),
       );
+    } on bsky.XRPCException catch (e) {
+      if (e.response.status.code == HttpStatus.tooManyRequests) {
+        final response = Response.json(
+          statusCode: HttpStatus.tooManyRequests,
+          body: {'error': 'Rate limit exceeded'},
+        );
+
+        return response.copyWith(
+          headers: Map.of(response.headers)
+            ..addAll(
+              {
+                'Cache-Control': 'no-cache, no-store, must-revalidate',
+              },
+            ),
+        );
+      } else {
+        rethrow;
+      }
     }
   }.use(requestLogger()).use(
         // Temporary middleware that print requests.