From edc0d094e2ff150d8636eb5b88881fd959b39c4c Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Sun, 30 Jun 2024 05:08:27 -0500
Subject: [PATCH] tries to disable quarantined aws keys

---
 .../middleware/response/error-generator.ts    | 38 +++++++++++++------
 src/proxy/middleware/response/index.ts        | 19 ++++++----
 2 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/src/proxy/middleware/response/error-generator.ts b/src/proxy/middleware/response/error-generator.ts
index c5c9735..8e96514 100644
--- a/src/proxy/middleware/response/error-generator.ts
+++ b/src/proxy/middleware/response/error-generator.ts
@@ -65,7 +65,7 @@ type ErrorGeneratorOptions = {
   format: APIFormat | "unknown";
   title: string;
   message: string;
-  obj?: object;
+  obj?: Record<string, any>;
   reqId: string | number | object;
   model?: string;
   statusCode?: number;
@@ -95,6 +95,23 @@ export function tryInferFormat(body: any): APIFormat | "unknown" {
   return "unknown";
 }
 
+// avoid leaking upstream hostname on dns resolution error
+function redactHostname(options: ErrorGeneratorOptions): ErrorGeneratorOptions {
+  if (!options.message.includes("getaddrinfo")) return options;
+
+  const redacted = { ...options };
+  redacted.message = "Could not resolve hostname";
+
+  if (typeof redacted.obj?.error === "object") {
+    redacted.obj = {
+      ...redacted.obj,
+      error: { message: "Could not resolve hostname" },
+    };
+  }
+
+  return redacted;
+}
+
 export function sendErrorToClient({
   options,
   req,
@@ -104,27 +121,26 @@ export function sendErrorToClient({
   req: express.Request;
   res: express.Response;
 }) {
-  const { format: inputFormat } = options;
+  const redactedOpts = redactHostname(options);
+  const { format: inputFormat } = redactedOpts;
 
-  // This is an error thrown before we know the format of the request, so we
-  // can't send a response in the format the client expects.
   const format =
     inputFormat === "unknown" ? tryInferFormat(req.body) : inputFormat;
   if (format === "unknown") {
-    return res.status(options.statusCode || 400).json({
-      error: options.message,
-      details: options.obj,
+    return res.status(redactedOpts.statusCode || 400).json({
+      error: redactedOpts.message,
+      details: redactedOpts.obj,
     });
   }
 
-  const completion = buildSpoofedCompletion({ ...options, format });
-  const event = buildSpoofedSSE({ ...options, format });
+  const completion = buildSpoofedCompletion({ ...redactedOpts, format });
+  const event = buildSpoofedSSE({ ...redactedOpts, format });
   const isStreaming =
     req.isStreaming || req.body.stream === true || req.body.stream === "true";
 
   if (!res.headersSent) {
-    res.setHeader("x-oai-proxy-error", options.title);
-    res.setHeader("x-oai-proxy-error-status", options.statusCode || 500);
+    res.setHeader("x-oai-proxy-error", redactedOpts.title);
+    res.setHeader("x-oai-proxy-error-status", redactedOpts.statusCode || 500);
   }
 
   if (isStreaming) {
diff --git a/src/proxy/middleware/response/index.ts b/src/proxy/middleware/response/index.ts
index 9662def..15aab49 100644
--- a/src/proxy/middleware/response/index.ts
+++ b/src/proxy/middleware/response/index.ts
@@ -1,4 +1,5 @@
 /* This file is fucking horrendous, sorry */
+// TODO: extract all per-service error response handling into its own modules
 import { Request, Response } from "express";
 import * as http from "http";
 import { config } from "../../../config";
@@ -194,7 +195,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
     { statusCode, type: errorType, errorPayload, key: req.key?.hash },
     `Received error response from upstream. (${proxyRes.statusMessage})`
   );
-  
+
   // TODO: split upstream error handling into separate modules for each service,
   // this is out of control.
 
@@ -225,7 +226,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
         break;
       case "anthropic":
       case "aws":
-        await handleAnthropicBadRequestError(req, errorPayload);
+        await handleAnthropicAwsBadRequestError(req, errorPayload);
         break;
       default:
         assertNever(service);
@@ -247,7 +248,9 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
           );
           keyPool.update(req.key!, { allowsMultimodality: false });
           await reenqueueRequest(req);
-          throw new RetryableError("Claude request re-enqueued because key does not support multimodality.");
+          throw new RetryableError(
+            "Claude request re-enqueued because key does not support multimodality."
+          );
         } else {
           keyPool.disable(req.key!, "revoked");
           errorPayload.proxy_note = `Assigned API key is invalid or revoked, please try again.`;
@@ -347,7 +350,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
   throw new HttpError(statusCode, errorPayload.error?.message);
 };
 
-async function handleAnthropicBadRequestError(
+async function handleAnthropicAwsBadRequestError(
   req: Request,
   errorPayload: ProxiedErrorPayload
 ) {
@@ -382,11 +385,13 @@ async function handleAnthropicBadRequestError(
     return;
   }
 
-  const isDisabled = error?.message?.match(/organization has been disabled/i);
+  const isDisabled =
+    error?.message?.match(/organization has been disabled/i) ||
+    error?.message?.match(/^operation not allowed/i);
   if (isDisabled) {
     req.log.warn(
       { key: req.key?.hash, message: error?.message },
-      "Anthropic key has been disabled."
+      "Anthropic/AWS key has been disabled."
     );
     keyPool.disable(req.key!, "revoked");
     errorPayload.proxy_note = `Assigned key has been disabled. (${error?.message})`;
@@ -512,7 +517,7 @@ async function handleOpenAIRateLimitError(
     //   keyPool.markRateLimited(req.key!);
     //   break;
     default:
-      errorPayload.proxy_note = `This is likely a temporary error with OpenAI. Try again in a few seconds.`;
+      errorPayload.proxy_note = `This is likely a temporary error with the API. Try again in a few seconds.`;
       break;
   }
   return errorPayload;