Merge pull request #426 from upstash/CTX7-272

src/index.ts

@@ -10,6 +10,7 @@ import { createServer } from "http";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
 import { Command } from "commander";
+import { IncomingMessage } from "http";
 
 const DEFAULT_MINIMUM_TOKENS = 10000;
 
@@ -46,8 +47,21 @@ const CLI_PORT = (() => {
 // Store SSE transports by session ID
 const sseTransports: Record<string, SSEServerTransport> = {};
 
+function getClientIp(req: IncomingMessage): string | undefined {
+  // Check for X-Forwarded-For header (set by AWS ELB and other load balancers)
+  const forwardedFor = req.headers["x-forwarded-for"];
+  if (forwardedFor) {
+    // X-Forwarded-For can contain multiple IPs, take the first one
+    const ips = Array.isArray(forwardedFor) ? forwardedFor[0] : forwardedFor;
+    return ips.split(",")[0].trim();
+  }
+
+  // Fall back to socket remote address
+  return req.socket?.remoteAddress || undefined;
+}
+
 // Function to create a new server instance with all tools registered
-function createServerInstance() {
+function createServerInstance(clientIp?: string) {
   const server = new McpServer(
     {
       name: "Context7",
@@ -87,7 +101,7 @@ For ambiguous queries, request clarification before proceeding with a best-guess
         .describe("Library name to search for and retrieve a Context7-compatible library ID."),
     },
     async ({ libraryName }) => {
-      const searchResponse: SearchResponse = await searchLibraries(libraryName);
+      const searchResponse: SearchResponse = await searchLibraries(libraryName, clientIp);
 
       if (!searchResponse.results || searchResponse.results.length === 0) {
         return {
@@ -151,10 +165,14 @@ ${resultsText}`,
         ),
     },
     async ({ context7CompatibleLibraryID, tokens = DEFAULT_MINIMUM_TOKENS, topic = "" }) => {
-      const fetchDocsResponse = await fetchLibraryDocumentation(context7CompatibleLibraryID, {
+      const fetchDocsResponse = await fetchLibraryDocumentation(
+        context7CompatibleLibraryID,
+        {
           tokens,
           topic,
-      });
+        },
+        clientIp
+      );
 
       if (!fetchDocsResponse) {
         return {
@@ -205,8 +223,11 @@ async function main() {
       }
 
       try {
+        // Extract client IP address using socket remote address (most reliable)
+        const clientIp = getClientIp(req);
+
         // Create new server instance for each request
-        const requestServer = createServerInstance();
+        const requestServer = createServerInstance(clientIp);
 
         if (url === "/mcp") {
           const transport = new StreamableHTTPServerTransport({

src/lib/api.ts

@@ -1,4 +1,5 @@
 import { SearchResponse } from "./types.js";
+import { generateHeaders } from "./encryption.js";
 
 const CONTEXT7_API_BASE_URL = "https://context7.com/api";
 const DEFAULT_TYPE = "txt";
@@ -6,13 +7,17 @@ const DEFAULT_TYPE = "txt";
 /**
  * Searches for libraries matching the given query
  * @param query The search query
+ * @param clientIp Optional client IP address to include in headers
  * @returns Search results or null if the request fails
  */
-export async function searchLibraries(query: string): Promise<SearchResponse> {
+export async function searchLibraries(query: string, clientIp?: string): Promise<SearchResponse> {
   try {
     const url = new URL(`${CONTEXT7_API_BASE_URL}/v1/search`);
     url.searchParams.set("query", query);
-    const response = await fetch(url);
+
+    const headers = generateHeaders(clientIp);
+
+    const response = await fetch(url, { headers });
     if (!response.ok) {
       const errorCode = response.status;
       if (errorCode === 429) {
@@ -39,6 +44,7 @@ export async function searchLibraries(query: string): Promise<SearchResponse> {
  * Fetches documentation context for a specific library
  * @param libraryId The library ID to fetch documentation for
  * @param options Options for the request
+ * @param clientIp Optional client IP address to include in headers
  * @returns The documentation text or null if the request fails
  */
 export async function fetchLibraryDocumentation(
@@ -46,7 +52,8 @@ export async function fetchLibraryDocumentation(
   options: {
     tokens?: number;
     topic?: string;
-  } = {}
+  } = {},
+  clientIp?: string
 ): Promise<string | null> {
   try {
     if (libraryId.startsWith("/")) {
@@ -56,11 +63,10 @@ export async function fetchLibraryDocumentation(
     if (options.tokens) url.searchParams.set("tokens", options.tokens.toString());
     if (options.topic) url.searchParams.set("topic", options.topic);
     url.searchParams.set("type", DEFAULT_TYPE);
-    const response = await fetch(url, {
-      headers: {
-        "X-Context7-Source": "mcp-server",
-      },
-    });
+
+    const headers = generateHeaders(clientIp, { "X-Context7-Source": "mcp-server" });
+
+    const response = await fetch(url, { headers });
     if (!response.ok) {
       const errorCode = response.status;
       if (errorCode === 429) {

src/lib/encryption.ts

@@ -0,0 +1,40 @@
+import { createCipheriv, randomBytes } from "crypto";
+
+export const ENCRYPTION_KEY =
+  process.env.CLIENT_IP_ENCRYPTION_KEY ||
+  "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f";
+export const ALGORITHM = "aes-256-cbc";
+
+export function validateEncryptionKey(key: string): boolean {
+  // Must be exactly 64 hex characters (32 bytes)
+  return /^[0-9a-fA-F]{64}$/.test(key);
+}
+
+export function encryptClientIp(clientIp: string): string {
+  if (!validateEncryptionKey(ENCRYPTION_KEY)) {
+    console.error("Invalid encryption key format. Must be 64 hex characters.");
+    return clientIp; // Fallback to unencrypted
+  }
+
+  try {
+    const iv = randomBytes(16);
+    const cipher = createCipheriv(ALGORITHM, Buffer.from(ENCRYPTION_KEY, "hex"), iv);
+    let encrypted = cipher.update(clientIp, "utf8", "hex");
+    encrypted += cipher.final("hex");
+    return iv.toString("hex") + ":" + encrypted;
+  } catch (error) {
+    console.error("Error encrypting client IP:", error);
+    return clientIp; // Fallback to unencrypted
+  }
+}
+
+export function generateHeaders(
+  clientIp?: string,
+  extraHeaders: Record<string, string> = {}
+): Record<string, string> {
+  const headers: Record<string, string> = { ...extraHeaders };
+  if (clientIp) {
+    headers["mcp-client-ip"] = encryptClientIp(clientIp);
+  }
+  return headers;
+}