Initial commit

src/services/auth.test.ts

@@ -0,0 +1,282 @@
+import crypto from "crypto";
+import {
+  decryptString,
+  exchangeToken,
+  generateMcpTokens,
+  generatePKCEChallenge,
+  generateToken,
+  getClientRegistration,
+  readMcpInstallation,
+  readPendingAuthorization,
+  readRefreshToken,
+  revokeMcpInstallation,
+  saveClientRegistration,
+  saveMcpInstallation,
+  savePendingAuthorization,
+  saveRefreshToken,
+  saveTokenExchange,
+} from "./auth.js";
+import { MockRedisClient, setRedisClient } from "../redis.js";
+import { McpInstallation, PendingAuthorization, TokenExchange } from "../types.js";
+import { OAuthClientInformationFull } from "@modelcontextprotocol/sdk/shared/auth.js";
+
+
+describe("auth utils", () => {
+  let mockRedis: MockRedisClient;
+
+  beforeEach(() => {
+    mockRedis = new MockRedisClient();
+    setRedisClient(mockRedis);
+    
+    jest.resetAllMocks();
+  });
+
+  describe("generateToken", () => {
+    it("generates a 64-character hex string", () => {
+      const token = generateToken();
+      expect(token).toMatch(/^[0-9a-f]{64}$/);
+    });
+
+    it("generates unique tokens", () => {
+      const token1 = generateToken();
+      const token2 = generateToken();
+      expect(token1).not.toBe(token2);
+    });
+  });
+
+  describe("generateMcpTokens", () => {
+    it("generates valid token structure", () => {
+      const tokens = generateMcpTokens();
+      
+      expect(tokens.access_token).toBeDefined();
+      expect(tokens.access_token).toMatch(/^[0-9a-f]{64}$/);
+      
+      expect(tokens.refresh_token).toBeDefined();
+      expect(tokens.refresh_token).toMatch(/^[0-9a-f]{64}$/);
+      
+      expect(tokens.token_type).toBe("bearer");
+      expect(tokens.expires_in).toBeDefined();
+    });
+    
+    it("generates unique tokens", () => {
+      const tokens1 = generateMcpTokens();
+      const tokens2 = generateMcpTokens();
+      
+      expect(tokens1.access_token).not.toBe(tokens2.access_token);
+      expect(tokens1.refresh_token).not.toBe(tokens2.refresh_token);
+    });
+  });
+
+  describe("client registration", () => {
+    it("saves and retrieves client registration", async () => {
+      const clientId = "test-client-id";
+      const registration: OAuthClientInformationFull = {
+        client_id: clientId,
+        client_name: "Test Client",
+        client_uri: "https://example.com",
+        redirect_uris: ["https://example.com/callback"]
+      };
+      
+      await saveClientRegistration(clientId, registration);
+      const retrieved = await getClientRegistration(clientId);
+      
+      expect(retrieved).toEqual(registration);
+    });
+    
+    it("returns undefined for non-existent client", async () => {
+      const result = await getClientRegistration("non-existent");
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe("pending authorization", () => {
+    it("saves and retrieves pending authorization", async () => {
+      const authCode = generateToken();
+      const pendingAuth: PendingAuthorization = {
+        redirectUri: "https://example.com/callback",
+        codeChallenge: "test-challenge",
+        codeChallengeMethod: "S256",
+        clientId: "test-client-id"
+      };
+      
+      await savePendingAuthorization(authCode, pendingAuth);
+      const retrieved = await readPendingAuthorization(authCode);
+      
+      expect(retrieved).toEqual(pendingAuth);
+    });
+    
+    it("returns undefined for non-existent code", async () => {
+      const result = await readPendingAuthorization("non-existent");
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe("token exchange", () => {
+    it("saves and retrieves token exchange data", async () => {
+      const authCode = generateToken();
+      const tokenExchange: TokenExchange = {
+        mcpAccessToken: generateToken(),
+        alreadyUsed: false
+      };
+      
+      // For this test, we'll directly manipulate Redis to check the token
+      // instead of using exchangeToken which changes the value
+      await saveTokenExchange(authCode, tokenExchange);
+      
+      // Get the key used by saveTokenExchange
+      const key = "exch:" + crypto.createHash("sha256").update(authCode).digest("hex");
+      
+      // Get the encrypted data
+      const encryptedData = await mockRedis.get(key);
+      expect(encryptedData).not.toBeNull();
+      
+      // Decrypt it manually to verify
+      const decoded = JSON.parse(decryptString({
+        encryptedText: encryptedData!,
+        key: authCode
+      }));
+      
+      expect(decoded).toEqual(tokenExchange);
+    });
+    
+    it("prevents duplicate token exchange", async () => {
+      const authCode = generateToken();
+      const tokenExchange: TokenExchange = {
+        mcpAccessToken: generateToken(),
+        alreadyUsed: false
+      };
+      
+      await saveTokenExchange(authCode, tokenExchange);
+      
+      // First exchange succeeds
+      const first = await exchangeToken(authCode);
+      expect(first).toBeDefined();
+      
+      // Second exchange throws
+      await expect(exchangeToken(authCode)).rejects.toThrow(
+        "Duplicate use of authorization code detected"
+      );
+    });
+    
+    it("returns undefined for non-existent code", async () => {
+      const result = await exchangeToken("non-existent");
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe("MCP installation", () => {
+    it("returns undefined for missing tokens", async () => {
+      const accessToken = generateToken();
+      const result = await readMcpInstallation(accessToken);
+      expect(result).toBeUndefined();
+    });
+
+    it("saves and retrieves installation data", async () => {
+      const accessToken = generateToken();
+
+      const mcpInstallation: McpInstallation = {
+        fakeUpstreamInstallation: {
+          fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+          fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+        },
+        mcpTokens: {
+          access_token: accessToken,
+          token_type: "Bearer",
+          expires_in: 3600,
+        },
+        clientId: "client-id",
+        issuedAt: Date.now() / 1000,
+      }
+
+      await saveMcpInstallation(accessToken, mcpInstallation);
+
+      const result = await readMcpInstallation(accessToken);
+      expect(result).toEqual(mcpInstallation);
+    });
+  });
+
+  describe("refresh token", () => {
+    it("saves and retrieves refresh token mapping", async () => {
+      const refreshToken = generateToken();
+      const accessToken = generateToken();
+      
+      await saveRefreshToken(refreshToken, accessToken);
+      const retrieved = await readRefreshToken(refreshToken);
+      
+      expect(retrieved).toBe(accessToken);
+    });
+    
+    it("returns undefined for non-existent refresh token", async () => {
+      const result = await readRefreshToken("non-existent");
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe("revokeMcpInstallation", () => {
+    it("revokes token for valid installation", async () => {
+      // For this test, we'll simply test if the WebClient is called with our mocked token
+      
+      // Create mock installation
+      const accessToken = generateToken();
+      
+      // Save it to Redis with actual function
+      await saveMcpInstallation(accessToken, {
+        fakeUpstreamInstallation: {
+          fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+          fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+        },
+        mcpTokens: {
+          access_token: accessToken,
+          token_type: "Bearer",
+          expires_in: 3600,
+        },
+        clientId: "client-id",
+        issuedAt: Date.now() / 1000,
+      });
+      
+      const getDel = jest.spyOn(mockRedis, 'getDel').mockImplementationOnce(() => {
+        // Need to return encrypted data for successful decryption in the revoke function
+        // Create encrypted data using our access token
+        const mcpInstallation = {
+          mcpTokens: {
+            access_token: accessToken,
+            token_type: "Bearer",
+            expires_in: 3600,
+          },
+          clientId: "client-id",
+          issuedAt: Date.now() / 1000,
+        };
+        const value = JSON.stringify(mcpInstallation);
+        const iv = crypto.randomBytes(16);
+        const cipher = crypto.createCipheriv("aes-256-cbc", Buffer.from(accessToken, "hex"), iv);
+        let encrypted = cipher.update(value, "utf-8", "hex");
+        encrypted += cipher.final("hex");
+        return Promise.resolve(`${iv.toString("hex")}:${encrypted}`);
+      });
+
+      await revokeMcpInstallation(accessToken);
+      
+      // Should have called getDel with the correct key
+      expect(getDel).toHaveBeenCalledWith(expect.stringContaining("mcp:"));
+      
+    });
+    
+    it("handles non-existent installation without error", async () => {
+      await expect(revokeMcpInstallation("non-existent")).resolves.not.toThrow();
+    });
+  });
+
+  describe("generatePKCEChallenge", () => {
+    it("generates base64url-encoded SHA256 hash", () => {
+      const verifier = "test_verifier";
+      const challenge = generatePKCEChallenge(verifier);
+      expect(challenge).toBe("0Ku4rR8EgR1w3HyHLBCxVLtPsAAks5HOlpmTEt0XhVA");
+    });
+
+    it("generates different challenges for different verifiers", () => {
+      const challenge1 = generatePKCEChallenge("verifier1");
+      const challenge2 = generatePKCEChallenge("verifier2");
+      expect(challenge1).not.toBe(challenge2);
+    });
+  });
+});

src/services/auth.ts

@@ -0,0 +1,266 @@
+import { SetOptions } from "@redis/client";
+import crypto from "crypto";
+import { redisClient } from "../redis.js";
+import { McpInstallation, PendingAuthorization, TokenExchange } from "../types.js";
+import { OAuthClientInformationFull, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
+
+export function generatePKCEChallenge(verifier: string): string {
+  const buffer = Buffer.from(verifier);
+  const hash = crypto.createHash("sha256").update(buffer);
+  return hash.digest("base64url");
+}
+
+export function generateToken(): string {
+  return crypto.randomBytes(32).toString("hex");
+}
+
+function sha256(data: string): string {
+  return crypto.createHash("sha256").update(data).digest("hex");
+}
+
+function encryptString({ text, key }: { text: string; key: string }): string {
+  const iv = crypto.randomBytes(16);
+  const cipher = crypto.createCipheriv("aes-256-cbc", Buffer.from(key, "hex"), iv);
+  let encrypted = cipher.update(text, "utf-8", "hex");
+  encrypted += cipher.final("hex");
+  return `${iv.toString("hex")}:${encrypted}`;
+}
+
+export function decryptString({
+  encryptedText,
+  key,
+}: {
+  encryptedText: string;
+  key: string;
+}): string {
+  const [ivHex, encrypted] = encryptedText.split(":");
+  const iv = Buffer.from(ivHex, "hex");
+  const decipher = crypto.createDecipheriv("aes-256-cbc", Buffer.from(key, "hex"), iv);
+  let decrypted = decipher.update(encrypted, "hex", "utf-8");
+  decrypted += decipher.final("utf-8");
+  return decrypted;
+}
+
+const CLIENT_REGISTRATION_KEY_PREFIX = "client:";
+const PENDING_AUTHORIZATION_KEY_PREFIX = "pending:";
+const MCP_AUTHORIZATION_KEY_PREFIX = "mcp:";
+const TOKEN_EXCHANGE_KEY_PREFIX = "exch:";
+const REFRESH_TOKEN_KEY_PREFIX = "refresh:";
+
+// Timeouts of redis keys for different stages of the OAuth flow
+const REDIS_PENDING_AUTHORIZATION_EXPIRY_SEC = 10 * 60; // 10 minutes in seconds - authorization code -> PendingAuthorization
+const REDIS_TOKEN_EXCHANGE_EXPIRY_SEC = 10 * 60; // 10 minutes in seconds - authorization code -> MCP access token
+const REDIS_UPSTREAM_INSTALLATION_EXPIRY_SEC = 7 * 24 * 60 * 60; // 7 days in seconds - MCP access token -> UpstreamInstallation
+const REDIS_REFRESH_TOKEN_EXPIRY_SEC = 7 * 24 * 60 * 60; // 7 days in seconds - MCP refresh token -> access token
+
+// Access token expiry
+const ACCESS_TOKEN_EXPIRY_SEC = 60 * 60 // 1 hour in seconds
+
+async function saveEncrypted<T>({
+  prefix,
+  key,
+  data,
+  options,
+}: {
+  prefix: string;
+  key: string;
+  data: T;
+  options?: SetOptions;
+}) {
+  const value = encryptString({
+    text: JSON.stringify(data),
+    key: key,
+  });
+
+  return await redisClient.set(prefix + sha256(key), value, options);
+}
+
+async function readEncrypted<T>({
+  prefix,
+  key,
+  del = false,
+}: {
+  prefix: string;
+  key: string;
+  del?: boolean;
+}): Promise<T | undefined> {
+  const data = del
+    ? await redisClient.getDel(prefix + sha256(key))
+    : await redisClient.get(prefix + sha256(key));
+
+  if (!data) {
+    return undefined;
+  }
+
+  const decoded = decryptString({
+    encryptedText: data,
+    key: key,
+  });
+
+  return JSON.parse(decoded);
+}
+
+export function generateMcpTokens(): OAuthTokens {
+  // Generate MCP access token and store both tokens
+  const mcpAccessToken = generateToken();
+  const mcpRefreshToken = generateToken();
+  
+  return {
+    access_token: mcpAccessToken,
+    refresh_token: mcpRefreshToken,
+    expires_in: ACCESS_TOKEN_EXPIRY_SEC,
+    token_type: "bearer",
+  }
+}
+
+export async function saveClientRegistration(
+  clientId: string,
+  registration: OAuthClientInformationFull,
+) {
+  await redisClient.set(
+    CLIENT_REGISTRATION_KEY_PREFIX + clientId,
+    JSON.stringify(registration),
+  );
+}
+
+export async function getClientRegistration(
+  clientId: string,
+): Promise<OAuthClientInformationFull | undefined> {
+  const data = await redisClient.get(CLIENT_REGISTRATION_KEY_PREFIX + clientId);
+  if (!data) {
+    return undefined;
+  }
+  return JSON.parse(data);
+}
+
+export async function savePendingAuthorization(
+  authorizationCode: string,
+  pendingAuthorization: PendingAuthorization,
+) {
+  await saveEncrypted({
+    prefix: PENDING_AUTHORIZATION_KEY_PREFIX,
+    key: authorizationCode,
+    data: pendingAuthorization,
+    options: { EX: REDIS_PENDING_AUTHORIZATION_EXPIRY_SEC },
+  });
+}
+
+export async function readPendingAuthorization(
+  authorizationCode: string,
+): Promise<PendingAuthorization | undefined> {
+  return readEncrypted<PendingAuthorization>({
+    prefix: PENDING_AUTHORIZATION_KEY_PREFIX,
+    key: authorizationCode,
+  });
+}
+
+export async function saveMcpInstallation(
+  mcpAccessToken: string,
+  installation: McpInstallation,
+) {
+  await saveEncrypted({
+    prefix: MCP_AUTHORIZATION_KEY_PREFIX,
+    key: mcpAccessToken,
+    data: installation,
+    options: { EX: REDIS_UPSTREAM_INSTALLATION_EXPIRY_SEC },
+  });
+}
+
+export async function readMcpInstallation(
+  mcpAccessToken: string,
+): Promise<McpInstallation | undefined> {
+  return readEncrypted<McpInstallation>({
+    prefix: MCP_AUTHORIZATION_KEY_PREFIX,
+    key: mcpAccessToken,
+  });
+}
+
+// This just links the refresh token to the upstream installation + mcp access token
+export async function saveRefreshToken(
+  refreshToken: string,
+  mcpAccessToken: string,
+) {
+  saveEncrypted({
+    prefix: REFRESH_TOKEN_KEY_PREFIX,
+    key: refreshToken,
+    data: mcpAccessToken,
+    options: { EX: REDIS_REFRESH_TOKEN_EXPIRY_SEC },
+  })
+}
+
+export async function readRefreshToken(
+  refreshToken: string,
+): Promise<string | undefined> {
+  return readEncrypted<string>({
+    prefix: REFRESH_TOKEN_KEY_PREFIX,
+    key: refreshToken,
+  });
+}
+
+export async function revokeMcpInstallation(
+  mcpAccessToken: string,
+): Promise<void> {
+  const installation = await readEncrypted<McpInstallation>({
+    prefix: MCP_AUTHORIZATION_KEY_PREFIX,
+    key: mcpAccessToken,
+    del: true,
+  });
+
+  if (!installation) {
+    return;
+  }
+  // Revoke upstream tokens here
+}
+
+export async function saveTokenExchange(
+  authorizationCode: string,
+  tokenExchange: TokenExchange,
+) {
+  await saveEncrypted({
+    prefix: TOKEN_EXCHANGE_KEY_PREFIX,
+    key: authorizationCode,
+    data: tokenExchange,
+    options: { EX: REDIS_TOKEN_EXCHANGE_EXPIRY_SEC },
+  });
+}
+
+/**
+ * Exchanges a temporary authorization code for an MCP access token. Will only succeed the first time.
+ */
+export async function exchangeToken(
+  authorizationCode: string,
+): Promise<TokenExchange | undefined> {
+  const data = await redisClient.get(TOKEN_EXCHANGE_KEY_PREFIX + sha256(authorizationCode));
+
+  if (!data) {
+    return undefined;
+  }
+
+  const decoded = decryptString({
+    encryptedText: data,
+    key: authorizationCode,
+  });
+
+  const tokenExchange: TokenExchange = JSON.parse(decoded);
+  if (tokenExchange.alreadyUsed) {
+    console.error("Duplicate use of authorization code detected; revoking tokens");
+    await revokeMcpInstallation(tokenExchange.mcpAccessToken);
+    throw new Error("Duplicate use of authorization code detected; tokens revoked");
+  }
+
+  const rereadData = await saveEncrypted({
+    prefix: TOKEN_EXCHANGE_KEY_PREFIX,
+    key: authorizationCode,
+    data: { ...tokenExchange, alreadyUsed: true },
+    options: { KEEPTTL: true, GET: true },
+  });
+
+  if (rereadData !== data) {
+    // Data concurrently changed while we were updating it. This necessarily means a duplicate use.
+    console.error("Duplicate use of authorization code detected; revoking tokens");
+    await revokeMcpInstallation(tokenExchange.mcpAccessToken);
+    throw new Error("Duplicate use of authorization code detected; tokens revoked");
+  }
+
+  return tokenExchange;
+}