Initial commit

src/auth/provider.test.ts

@@ -0,0 +1,419 @@
+import { Response } from "express";
+import { AuthorizationParams } from "@modelcontextprotocol/sdk/server/auth/provider.js";
+import { OAuthClientInformationFull, OAuthTokenRevocationRequest, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
+import { InvalidTokenError } from "@modelcontextprotocol/sdk/server/auth/errors.js";
+import { MockRedisClient, setRedisClient } from "../redis.js";
+import { McpInstallation, PendingAuthorization, TokenExchange } from "../types.js";
+import { EverythingAuthProvider, EverythingOAuthClientsStore } from "./provider.js";
+import * as authService from "../services/auth.js";
+
+// Create mocks for the auth service functions
+jest.mock("../services/auth.js", () => {
+  // Use actual implementations for some functions
+  const originalModule = jest.requireActual("../services/auth.js");
+
+  return {
+    ...originalModule,
+    exchangeToken: jest.fn(),
+    generateToken: jest.fn(),
+    getClientRegistration: jest.fn(),
+    readMcpInstallation: jest.fn(),
+    readPendingAuthorization: jest.fn(),
+    readRefreshToken: jest.fn(),
+    revokeMcpInstallation: jest.fn(),
+    saveClientRegistration: jest.fn(),
+    savePendingAuthorization: jest.fn(),
+    saveMcpInstallation: jest.fn(),
+    generateMcpTokens: jest.fn(),
+    saveRefreshToken: jest.fn(),
+  };
+});
+
+// Helper function to create sample client
+function createTestClient(): OAuthClientInformationFull {
+  return {
+    client_id: "test-client-id",
+    client_name: "Test Client",
+    client_uri: "https://example.com",
+    redirect_uris: ["https://example.com/callback"]
+  };
+}
+
+// Helper function to create a mock Response object
+function createMockResponse() {
+  const res = {
+    redirect: jest.fn().mockReturnThis(),
+    status: jest.fn().mockReturnThis(),
+    json: jest.fn().mockReturnThis()
+  };
+  return res as unknown as jest.Mocked<Response>;
+}
+
+
+function getMockAuthValues() {
+  const client = createTestClient();
+  const accessToken = "test-access-token";
+  const newTokens: OAuthTokens = {
+    access_token: "new-access-token",
+    refresh_token: "new-refresh-token",
+    token_type: "bearer",
+    expires_in: 3600,
+  };
+  const mcpInstallation: McpInstallation = {
+    fakeUpstreamInstallation: {
+      fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+      fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+    },
+    mcpTokens: {
+      access_token: accessToken,
+      token_type: "Bearer",
+      refresh_token: "test-refresh-token",
+      expires_in: 3600,
+    },
+    clientId: client.client_id,
+    issuedAt: Date.now() / 1000,
+  };
+
+  return {
+    client,
+    accessToken,
+    newTokens,
+    mcpInstallation,
+  }
+}
+
+describe("EverythingOAuthClientsStore", () => {
+  let clientsStore: EverythingOAuthClientsStore;
+  
+  beforeEach(() => {
+    jest.resetAllMocks();
+    clientsStore = new EverythingOAuthClientsStore();
+  });
+  
+  describe("getClient", () => {
+    it("returns undefined for non-existent client", async () => {
+      (authService.getClientRegistration as jest.Mock).mockResolvedValueOnce(undefined);
+      
+      const result = await clientsStore.getClient("non-existent");
+      
+      expect(result).toBeUndefined();
+      expect(authService.getClientRegistration).toHaveBeenCalledWith("non-existent");
+    });
+    
+    it("returns client information for existing client", async () => {
+      const client = createTestClient();
+      (authService.getClientRegistration as jest.Mock).mockResolvedValueOnce(client);
+      
+      const result = await clientsStore.getClient(client.client_id);
+      
+      expect(result).toEqual(client);
+      expect(authService.getClientRegistration).toHaveBeenCalledWith(client.client_id);
+    });
+  });
+  
+  describe("registerClient", () => {
+    it("saves and returns client information", async () => {
+      const client = createTestClient();
+      (authService.saveClientRegistration as jest.Mock).mockResolvedValueOnce(undefined);
+      
+      const result = await clientsStore.registerClient(client);
+      
+      expect(result).toEqual(client);
+      expect(authService.saveClientRegistration).toHaveBeenCalledWith(client.client_id, client);
+    });
+  });
+});
+
+describe("EverythingAuthProvider", () => {
+  let provider: EverythingAuthProvider;
+  let mockRedis: MockRedisClient;
+  
+  beforeEach(() => {
+    jest.resetAllMocks();
+    
+    mockRedis = new MockRedisClient();
+    setRedisClient(mockRedis);
+    
+    // Set up token generator mock
+    (authService.generateToken as jest.Mock).mockImplementation(() => "mock-token");
+    
+    provider = new EverythingAuthProvider();
+  });
+  
+  describe("authorize", () => {
+    it("saves pending authorization and redirects to upstream install", async () => {
+      const client = createTestClient();
+      // Use a type assertion to make TypeScript ignore the mismatch
+      const params = {
+        redirectUri: "https://example.com/callback",
+        codeChallenge: "test-challenge",
+        codeChallengeMethod: "S256",
+      } as unknown as AuthorizationParams;
+      const res = createMockResponse();
+      
+      await provider.authorize(client, params, res);
+      
+      // Verify pending authorization is saved
+      expect(authService.savePendingAuthorization).toHaveBeenCalledWith("mock-token", {
+        redirectUri: params.redirectUri,
+        codeChallenge: params.codeChallenge,
+        codeChallengeMethod: "S256",
+        clientId: client.client_id,
+      });
+      
+      // Verify redirect to upstream installation
+      expect(res.redirect).toHaveBeenCalledWith(expect.stringContaining("/fakeupstreamauth/authorize?metadata=mock-token"));
+    });
+  });
+  
+  describe("challengeForAuthorizationCode", () => {
+    it("returns code challenge for valid authorization code", async () => {
+      const client = createTestClient();
+      const pendingAuth: PendingAuthorization = {
+        redirectUri: "https://example.com/callback",
+        codeChallenge: "test-challenge",
+        codeChallengeMethod: "S256",
+        clientId: client.client_id,
+      };
+      
+      (authService.readPendingAuthorization as jest.Mock).mockResolvedValueOnce(pendingAuth);
+      
+      const result = await provider.challengeForAuthorizationCode(client, "test-code");
+      
+      expect(result).toBe("test-challenge");
+      expect(authService.readPendingAuthorization).toHaveBeenCalledWith("test-code");
+    });
+    
+    it("throws error for non-existent authorization code", async () => {
+      const client = createTestClient();
+      
+      (authService.readPendingAuthorization as jest.Mock).mockResolvedValueOnce(undefined);
+      
+      await expect(provider.challengeForAuthorizationCode(client, "test-code"))
+        .rejects
+        .toThrow("Authorization code not found");
+    });
+    
+    it("throws error when client ID doesn't match", async () => {
+      const client = createTestClient();
+      const pendingAuth: PendingAuthorization = {
+        redirectUri: "https://example.com/callback",
+        codeChallenge: "test-challenge",
+        codeChallengeMethod: "S256",
+        clientId: "different-client-id",
+      };
+      
+      (authService.readPendingAuthorization as jest.Mock).mockResolvedValueOnce(pendingAuth);
+      
+      await expect(provider.challengeForAuthorizationCode(client, "test-code"))
+        .rejects
+        .toThrow("Authorization code does not match client");
+    });
+  });
+  
+  describe("exchangeAuthorizationCode", () => {
+    it("returns tokens when exchange is successful", async () => {
+      const client = createTestClient();
+      const tokenData: TokenExchange = {
+        mcpAccessToken: "test-access-token",
+        alreadyUsed: false,
+      };
+      const { mcpInstallation } = getMockAuthValues();
+      
+      (authService.exchangeToken as jest.Mock).mockResolvedValueOnce(tokenData);
+      (authService.readMcpInstallation as jest.Mock).mockResolvedValueOnce(mcpInstallation);
+      
+      const result = await provider.exchangeAuthorizationCode(client, "test-code");
+      
+      expect(result).toEqual({
+        access_token: tokenData.mcpAccessToken,
+        expires_in: mcpInstallation.mcpTokens.expires_in,
+        refresh_token: mcpInstallation.mcpTokens.refresh_token,
+        token_type: "Bearer",
+      });
+      expect(authService.exchangeToken).toHaveBeenCalledWith("test-code");
+    });
+    
+    it("throws error for invalid authorization code", async () => {
+      const client = createTestClient();
+      
+      (authService.exchangeToken as jest.Mock).mockResolvedValueOnce(undefined);
+      
+      await expect(provider.exchangeAuthorizationCode(client, "test-code"))
+        .rejects
+        .toThrow("Invalid authorization code");
+    });
+  });
+  
+  describe("exchangeRefreshToken", () => {
+    it("returns new tokens when refresh token is valid", async () => {
+      const {
+        client,
+        accessToken,
+        newTokens,
+        mcpInstallation,
+      } = getMockAuthValues();
+
+      
+      // Mock service calls
+      (authService.readRefreshToken as jest.Mock).mockResolvedValueOnce(accessToken);
+      (authService.readMcpInstallation as jest.Mock).mockResolvedValueOnce(mcpInstallation);
+      (authService.generateMcpTokens as jest.Mock).mockReturnValueOnce(newTokens);
+      
+      const result = await provider.exchangeRefreshToken(client, "test-refresh-token");
+      
+      expect(result).toEqual(newTokens);
+      expect(authService.readRefreshToken).toHaveBeenCalledWith("test-refresh-token");
+      expect(authService.readMcpInstallation).toHaveBeenCalledWith(accessToken);
+      expect(authService.saveRefreshToken).toHaveBeenCalledWith(newTokens.refresh_token, newTokens.access_token);
+      expect(authService.saveMcpInstallation).toHaveBeenCalledWith(
+        newTokens.access_token, 
+        expect.objectContaining({
+          ...mcpInstallation,
+          mcpTokens: newTokens,
+          issuedAt: expect.any(Number),
+        })
+      );
+    });
+    
+    it("throws error for invalid refresh token", async () => {
+      const client = createTestClient();
+      
+      (authService.readRefreshToken as jest.Mock).mockResolvedValueOnce(undefined);
+      
+      await expect(provider.exchangeRefreshToken(client, "test-refresh-token"))
+        .rejects
+        .toThrow("Invalid refresh token");
+    });
+    
+    it("throws error for refresh token with no installation", async () => {
+      const client = createTestClient();
+      
+      (authService.readRefreshToken as jest.Mock).mockResolvedValueOnce("test-access-token");
+      (authService.readMcpInstallation as jest.Mock).mockResolvedValueOnce(undefined);
+      
+      await expect(provider.exchangeRefreshToken(client, "test-refresh-token"))
+        .rejects
+        .toThrow("Invalid refresh token");
+    });
+    
+    it("throws error when client ID doesn't match", async () => {
+      const client = createTestClient();
+      const accessToken = "test-access-token";
+      
+      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: "different-client-id",
+        issuedAt: Date.now() / 1000,
+      };
+      
+      (authService.readRefreshToken as jest.Mock).mockResolvedValueOnce(accessToken);
+      (authService.readMcpInstallation as jest.Mock).mockResolvedValueOnce(mcpInstallation);
+      
+      await expect(provider.exchangeRefreshToken(client, "test-refresh-token"))
+        .rejects
+        .toThrow("Invalid client");
+    });
+
+    it("works correctly even when the refresh token has expired from Redis", async () => {
+      const client = createTestClient();
+      
+      // Simulate the refresh token not being found in Redis (expired)
+      (authService.readRefreshToken as jest.Mock).mockResolvedValueOnce(undefined);
+      
+      await expect(provider.exchangeRefreshToken(client, "expired-refresh-token"))
+        .rejects
+        .toThrow("Invalid refresh token");
+      
+      expect(authService.readRefreshToken).toHaveBeenCalledWith("expired-refresh-token");
+      // Should not proceed to check the installation since the refresh token lookup failed
+      expect(authService.readMcpInstallation).not.toHaveBeenCalled();
+    });
+  });
+  
+  describe("verifyAccessToken", () => {
+    it("returns auth info for valid token", async () => {
+      const accessToken = "test-access-token";
+      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,
+      };
+      
+      (authService.readMcpInstallation as jest.Mock).mockResolvedValueOnce(mcpInstallation);
+      
+      const result = await provider.verifyAccessToken(accessToken);
+      
+      expect(result).toEqual({
+        token: accessToken,
+        clientId: mcpInstallation.clientId,
+        scopes: ['mcp'],
+        expiresAt: mcpInstallation.mcpTokens.expires_in! + mcpInstallation.issuedAt,
+      });
+    });
+    
+    it("throws error for invalid token", async () => {
+      (authService.readMcpInstallation as jest.Mock).mockResolvedValueOnce(undefined);
+      
+      await expect(provider.verifyAccessToken("invalid-token"))
+        .rejects
+        .toThrow("Invalid access token");
+    });
+    
+    it("throws InvalidTokenError for expired token", async () => {
+      const accessToken = "test-access-token";
+      const oneDayInSeconds = 24 * 60 * 60;
+      const twoDaysAgoInSeconds = Math.floor(Date.now() / 1000) - (2 * oneDayInSeconds);
+      
+      const mcpInstallation: McpInstallation = {
+        fakeUpstreamInstallation: {
+          fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+          fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+        },
+        mcpTokens: {
+          access_token: accessToken,
+          token_type: "Bearer",
+          expires_in: oneDayInSeconds,
+        },
+        clientId: "client-id",
+        issuedAt: twoDaysAgoInSeconds, // 2 days ago, with 1-day expiry
+      };
+      
+      (authService.readMcpInstallation as jest.Mock).mockResolvedValueOnce(mcpInstallation);
+      
+      await expect(provider.verifyAccessToken(accessToken))
+        .rejects
+        .toThrow(InvalidTokenError);
+    });
+  });
+  
+  describe("revokeToken", () => {
+    it("calls revokeMcpInstallation with token", async () => {
+      const client = createTestClient();
+      const request: OAuthTokenRevocationRequest = {
+        token: "test-token",
+        token_type_hint: "access_token"
+      };
+      
+      await provider.revokeToken(client, request);
+      
+      expect(authService.revokeMcpInstallation).toHaveBeenCalledWith(request.token);
+    });
+  });
+});

src/auth/provider.ts

@@ -0,0 +1,196 @@
+import { Response } from 'express';
+import { OAuthServerProvider, AuthorizationParams } from '@modelcontextprotocol/sdk/server/auth/provider.js';
+import { OAuthRegisteredClientsStore } from '@modelcontextprotocol/sdk/server/auth/clients.js';
+import { OAuthClientInformationFull, OAuthTokenRevocationRequest, OAuthTokens } from '@modelcontextprotocol/sdk/shared/auth.js';
+import { AuthInfo } from '@modelcontextprotocol/sdk/server/auth/types.js';
+import { BASE_URI } from '../config.js';
+import {
+  exchangeToken,
+  generateToken,
+  getClientRegistration,
+  readPendingAuthorization,
+  readMcpInstallation,
+  revokeMcpInstallation,
+  saveClientRegistration,
+  savePendingAuthorization,
+  readRefreshToken,
+  generateMcpTokens,
+  saveMcpInstallation,
+  saveRefreshToken,
+} from '../services/auth.js';
+import { InvalidTokenError } from '@modelcontextprotocol/sdk/server/auth/errors.js';
+
+/**
+ * Implementation of the OAuthRegisteredClientsStore interface using the existing client registration system
+ */
+export class EverythingOAuthClientsStore implements OAuthRegisteredClientsStore {
+  async getClient(clientId: string): Promise<OAuthClientInformationFull | undefined> {
+    const registration = await getClientRegistration(clientId);
+    if (!registration) {
+      return undefined;
+    }
+    return registration;
+  }
+
+  async registerClient(client: OAuthClientInformationFull): Promise<OAuthClientInformationFull> {
+    await saveClientRegistration(client.client_id, client);
+    return client;
+  }
+}
+
+/**
+ * Implementation of the OAuthServerProvider interface for upstream authentication
+ */
+export class EverythingAuthProvider implements OAuthServerProvider {
+  private _clientsStore: EverythingOAuthClientsStore;
+
+  constructor() {
+    this._clientsStore = new EverythingOAuthClientsStore();
+  }
+
+  get clientsStore(): OAuthRegisteredClientsStore {
+    return this._clientsStore;
+  }
+
+  async authorize(client: OAuthClientInformationFull, params: AuthorizationParams, res: Response): Promise<void> {
+
+    // Client is validated by the MCP sdk.
+
+    // Generate authorization code
+    const authorizationCode = generateToken();
+
+    // Save the pending authorization with code challenge and state
+    await savePendingAuthorization(authorizationCode, {
+      redirectUri: params.redirectUri,
+      codeChallenge: params.codeChallenge,
+      codeChallengeMethod: 'S256', // Currently only support S256
+      clientId: client.client_id,
+      state: params.state,
+    });
+
+    // TODO: should we use a different key, other than the authorization code, to store the pending authorization?
+    
+    // You can redirect to another page, or you can send an html response directly
+    // res.redirect(new URL(`fakeupstreamauth/authorize?metadata=${authorizationCode}`, BASE_URI).href);
+
+    res.send(`
+      <html>
+        <head>
+          <title>MCP Auth Page</title>
+        </head>
+        <body>
+          <h1>MCP Server Auth Page</h1>
+          <p>
+            This page is the authorization page presented by the MCP server, routing the user upstream. This is only 
+            needed on 2025-03-26 Auth spec, where the MCP server acts as it's own authoriztion server. This page should
+            be present to avoid confused deputy attacks.
+          </p>
+          <p>
+            Click <a href="/fakeupstreamauth/authorize?redirect_uri=/fakeupstreamauth/callback&state=${authorizationCode}">here</a> 
+            to continue to the upstream auth
+          </p>
+        </body>
+      </html>
+    `);
+  }
+
+  async challengeForAuthorizationCode(client: OAuthClientInformationFull, authorizationCode: string): Promise<string> {
+    const pendingAuth = await readPendingAuthorization(authorizationCode);
+    if (!pendingAuth) {
+      throw new Error('Authorization code not found');
+    }
+
+    if (pendingAuth.clientId !== client.client_id) {
+      console.log('Pending auth clientId', pendingAuth.clientId);
+      console.log('ClientId', client.client_id);
+      throw new Error('Authorization code does not match client');
+    }
+
+    return pendingAuth.codeChallenge;
+  }
+
+  async exchangeAuthorizationCode(client: OAuthClientInformationFull, authorizationCode: string): Promise<OAuthTokens> {
+    const tokenData = await exchangeToken(authorizationCode);
+    if (!tokenData) {
+      throw new Error('Invalid authorization code');
+    }
+
+    // Get the MCP installation to retrieve the full token data including refresh token
+    const mcpInstallation = await readMcpInstallation(tokenData.mcpAccessToken);
+    if (!mcpInstallation) {
+      throw new Error('Failed to retrieve MCP installation');
+    }
+
+    // Return the full token data including refresh token
+    return {
+      access_token: mcpInstallation.mcpTokens.access_token,
+      refresh_token: mcpInstallation.mcpTokens.refresh_token,
+      expires_in: mcpInstallation.mcpTokens.expires_in,
+      token_type: 'Bearer',
+    };
+  }
+
+  async exchangeRefreshToken(client: OAuthClientInformationFull, refreshToken: string, _scopes?: string[]): Promise<OAuthTokens> {
+    const accessToken = await readRefreshToken(refreshToken);
+
+    if (!accessToken) {
+      throw new Error('Invalid refresh token');
+    }
+
+    const mcpInstallation = await readMcpInstallation(accessToken);
+
+    if (!mcpInstallation) {
+      throw new Error('Invalid refresh token');
+    }
+
+    // Check the client_id
+    if (mcpInstallation.clientId !== client.client_id) {
+      throw new Error('Invalid client');
+    }
+    
+    const newTokens = generateMcpTokens();
+
+    if (newTokens.refresh_token) {
+      await saveRefreshToken(newTokens.refresh_token, newTokens.access_token);
+    }
+
+    // Update the installation with the new tokens
+    await saveMcpInstallation(newTokens.access_token, {
+      ...mcpInstallation,
+      mcpTokens: newTokens,
+      issuedAt: Date.now() / 1000,
+    });
+
+    return newTokens;
+  }
+
+  async verifyAccessToken(token: string): Promise<AuthInfo> {
+    const installation = await readMcpInstallation(token);
+    if (!installation) {
+      throw new InvalidTokenError("Invalid access token");
+    }
+
+    const expiresAt = (
+      installation.mcpTokens.expires_in
+      ? installation.mcpTokens.expires_in + installation.issuedAt
+      : undefined
+    );
+
+    // This can be removed once in the SDK
+    // Check if the token is expired
+    if (!!expiresAt && expiresAt < Date.now() / 1000) {
+      throw new InvalidTokenError("Token has expired");
+    }
+    
+    return {
+      token,
+      clientId: installation.clientId,
+      scopes: ['mcp'],
+      expiresAt
+    };
+  }
+
+  async revokeToken(client: OAuthClientInformationFull, request: OAuthTokenRevocationRequest): Promise<void> {
+    await revokeMcpInstallation(request.token);
+  }
+}

src/config.ts

@@ -0,0 +1,6 @@
+import "dotenv/config";
+
+export const PORT = Number(process.env.PORT) || 3232;
+
+export const BASE_URI =
+  process.env.BASE_URI || "https://localhost:3232";

src/context.test.ts

@@ -0,0 +1,113 @@
+import { getFakeUpstreamInstallation, getMcpAccessToken, withContext } from "./context.js";
+import { generateToken } from "./services/auth.js";
+
+describe("context", () => {
+  const mockContext = {
+    mcpAccessToken: generateToken(),
+    fakeUpstreamInstallation: {
+      fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+      fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+    },
+  };
+
+  test("getFakeUpstreamInstallation throws when called outside context", () => {
+    expect(() => getFakeUpstreamInstallation()).toThrow(
+      "No request context found - are you calling this from within a request handler?"
+    );
+  });
+
+  test("getMcpAccessToken throws when called outside context", () => {
+    expect(() => getMcpAccessToken()).toThrow(
+      "No request context found - are you calling this from within a request handler?"
+    );
+  });
+
+  test("context functions return data within context", async () => {
+    await withContext(mockContext, () => {
+      const upstreamInstall = getFakeUpstreamInstallation();
+      expect(upstreamInstall).toEqual(mockContext.fakeUpstreamInstallation);
+
+      const mcpToken = getMcpAccessToken();
+      expect(mcpToken).toEqual(mockContext.mcpAccessToken);
+    });
+  });
+
+  test("nested contexts maintain isolation", async () => {
+    const context1 = {
+      mcpAccessToken: generateToken(),
+      fakeUpstreamInstallation: {
+        fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+        fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+      },
+    };
+    const context2 = {
+      mcpAccessToken: generateToken(),
+      fakeUpstreamInstallation: {
+        fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+        fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+      },
+    };
+
+    await withContext(context1, async () => {
+      expect(getFakeUpstreamInstallation()).toEqual(context1.fakeUpstreamInstallation);
+      expect(getMcpAccessToken()).toEqual(context1.mcpAccessToken);
+
+      await withContext(context2, () => {
+        expect(getFakeUpstreamInstallation()).toEqual(context2.fakeUpstreamInstallation);
+        expect(getMcpAccessToken()).toEqual(context2.mcpAccessToken);
+      });
+
+      // Outer context should be preserved
+      expect(getFakeUpstreamInstallation()).toEqual(context1.fakeUpstreamInstallation);
+      expect(getMcpAccessToken()).toEqual(context1.mcpAccessToken);
+    });
+  });
+
+  test("concurrent contexts maintain isolation", async () => {
+    const context1 = {
+      mcpAccessToken: generateToken(),
+      fakeUpstreamInstallation: {
+        fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+        fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+      },
+    };
+    const context2 = {
+      mcpAccessToken: generateToken(),
+      fakeUpstreamInstallation: {
+        fakeAccessTokenForDemonstration: "fake-upstream-access-token",
+        fakeRefreshTokenForDemonstration: "fake-upstream-refresh-token",
+      },
+    };
+
+    // Run two contexts concurrently
+    await Promise.all([
+      withContext(context1, async () => {
+        // Simulate async work
+        await new Promise(resolve => setTimeout(resolve, 10));
+        expect(getFakeUpstreamInstallation()).toEqual(context1.fakeUpstreamInstallation);
+        expect(getMcpAccessToken()).toEqual(context1.mcpAccessToken);
+      }),
+      withContext(context2, async () => {
+        // Simulate async work
+        await new Promise(resolve => setTimeout(resolve, 5));
+        expect(getFakeUpstreamInstallation()).toEqual(context2.fakeUpstreamInstallation);
+        expect(getMcpAccessToken()).toEqual(context2.mcpAccessToken);
+      }),
+    ]);
+  });
+
+  test("context is preserved across async operations", async () => {
+    await withContext(mockContext, async () => {
+      const initialUpstreamInstall = getFakeUpstreamInstallation();
+      const initialMcpToken = getMcpAccessToken();
+      
+      // Simulate async work
+      await new Promise(resolve => setTimeout(resolve, 10));
+      
+      const afterUpstreamInstall = getFakeUpstreamInstallation();
+      const afterMcpToken = getMcpAccessToken();
+      expect(afterUpstreamInstall).toEqual(initialUpstreamInstall);
+      expect(afterMcpToken).toEqual(initialMcpToken);
+    });
+  });
+});

src/context.ts

@@ -0,0 +1,36 @@
+import { AsyncLocalStorage } from "async_hooks";
+import { FakeUpstreamInstallation } from "./types.js";
+
+interface RequestContext {
+  mcpAccessToken: string;
+  fakeUpstreamInstallation: FakeUpstreamInstallation;
+}
+
+export const asyncLocalStorage = new AsyncLocalStorage<RequestContext>();
+
+export function getFakeUpstreamInstallation(): FakeUpstreamInstallation {
+  const context = asyncLocalStorage.getStore();
+  if (!context) {
+    throw new Error(
+      "No request context found - are you calling this from within a request handler?",
+    );
+  }
+  return context.fakeUpstreamInstallation;
+}
+
+export function getMcpAccessToken(): string {
+  const context = asyncLocalStorage.getStore();
+  if (!context) {
+    throw new Error(
+      "No request context found - are you calling this from within a request handler?",
+    );
+  }
+  return context.mcpAccessToken;
+}
+
+export function withContext<T>(
+  context: RequestContext,
+  fn: () => T | Promise<T>,
+): T | Promise<T> {
+  return asyncLocalStorage.run(context, fn);
+}

src/handlers/fakeauth.ts

@@ -0,0 +1,99 @@
+import { Request, Response } from "express";
+import { generateMcpTokens, readPendingAuthorization, saveMcpInstallation, saveRefreshToken, saveTokenExchange } from "../services/auth.js";
+import { McpInstallation } from "../types.js";
+
+// this module has a fake upstream auth server that returns a fake auth code, it also allows you to authorize or fail 
+// authorization, to test the different flows
+
+// TODO: make an implementation of this using the ProxyAuthProvider
+
+// This is mocking an upstream auth server. This wouldn't normally be in the same server as the MCP auth server
+export async function handleFakeAuthorize(req: Request, res: Response) {
+  // get the redirect_uri and state from the query params
+  const { redirect_uri, state } = req.query;
+
+  // TODO, mint actual codes?
+
+  res.send(`
+    <html>
+      <head>
+        <title>Fake Upstream Auth Provider!</title>
+      </head>
+      <body>
+        <h1>Fake Auth</h1>
+        <p>Fake auth page</p>
+        <p>Click <a href="${redirect_uri}?state=${state}&code=fakecode">here</a> to authorize</p>
+        <p>Click <a href="${redirect_uri}?state=${state}&code=fakecode">here</a> to fail authorization</p>
+      </body>
+    </html>
+  `);
+}
+
+
+// This is the callback URL that the upstream auth server will redirect to after authorization
+export async function handleFakeAuthorizeRedirect(req: Request, res: Response) {
+  const {
+    // The state returned from the upstream auth server is actually the authorization code
+    state: mcpAuthorizationCode,
+    code: upstreamAuthorizationCode,
+  } = req.query;
+
+  // This is where you'd exchange the upstreamAuthorizationCode for access/refresh tokens
+  // In this case, we're just going to fake it
+  const upstreamTokens = await fakeUpstreamTokenExchange(upstreamAuthorizationCode as string);
+
+  // Validate that it's a string
+  if (typeof mcpAuthorizationCode !== "string") {
+    throw new Error("Invalid authorization code");
+  }
+
+  const pendingAuth = await readPendingAuthorization(mcpAuthorizationCode);
+  if (!pendingAuth) {
+    throw new Error("No matching authorization found");
+  }
+
+
+  const mcpTokens = generateMcpTokens();
+
+  const mcpInstallation: McpInstallation = {
+    fakeUpstreamInstallation: {
+      fakeAccessTokenForDemonstration: upstreamTokens.access_token,
+      fakeRefreshTokenForDemonstration: upstreamTokens.refresh_token,
+    },
+    mcpTokens,
+    clientId: pendingAuth.clientId,
+    issuedAt: Date.now() / 1000,
+  }
+
+  // Store the upstream authorization data
+  await saveMcpInstallation(mcpTokens.access_token, mcpInstallation);
+
+  // Store the refresh token -> access token mapping
+  if (mcpTokens.refresh_token) {
+    await saveRefreshToken(mcpTokens.refresh_token, mcpTokens.access_token);
+  }
+
+  // Store the token exchange data
+  await saveTokenExchange(mcpAuthorizationCode, {
+    mcpAccessToken: mcpTokens.access_token,
+    alreadyUsed: false,
+  });
+
+  // Redirect back to the original application with the authorization code and state
+  const redirectUrl = pendingAuth.state ?
+    `${pendingAuth.redirectUri}?code=${mcpAuthorizationCode}&state=${pendingAuth.state}` :
+    `${pendingAuth.redirectUri}?code=${mcpAuthorizationCode}`;
+  res.redirect(redirectUrl);
+};
+
+function fakeUpstreamTokenExchange(
+  authorizationCode: string,
+): Promise<{ access_token: string; refresh_token: string }> {
+  // just return the authorization code with a suffix
+  return new Promise((resolve) => {
+    resolve({
+      access_token: `${authorizationCode}-exchanged-for-access-token`,
+      refresh_token: `${authorizationCode}-exchanged-for-refresh-token`,
+    });
+  });
+}

src/handlers/mcp.ts

@@ -0,0 +1,144 @@
+import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
+import contentType from "content-type";
+import { NextFunction, Request, Response } from "express";
+import getRawBody from "raw-body";
+import { readMcpInstallation } from "../services/auth.js";
+import { withContext } from "../context.js";
+import { createMcpServer } from "../services/mcp.js";
+import { redisClient } from "../redis.js";
+import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
+
+const MAXIMUM_MESSAGE_SIZE = "4mb";
+
+declare module "express-serve-static-core" {
+  interface Request {
+    /**
+     * Information about the validated access token, if the `requireBearerAuth` middleware was used.
+     */
+    auth?: AuthInfo;
+  }
+}
+
+export async function authContext(
+  req: Request,
+  res: Response,
+  next: NextFunction,
+) {
+  const authInfo = req.auth
+
+  if (!authInfo) {
+    res.set("WWW-Authenticate", 'Bearer error="invalid_token"');
+    res.status(401).json({ error: "Invalid access token" });
+    return;
+  }
+
+  const token = authInfo.token;
+
+  // Load UpstreamInstallation based on the access token
+  const mcpInstallation = await readMcpInstallation(token);
+  if (!mcpInstallation) {
+    res.set("WWW-Authenticate", 'Bearer error="invalid_token"');
+    res.status(401).json({ error: "Invalid access token" });
+    return;
+  }
+
+  // Wrap the rest of the request handling in the context
+  withContext({ mcpAccessToken: token, fakeUpstreamInstallation: mcpInstallation.fakeUpstreamInstallation }, () =>
+    next(),
+  );
+}
+
+function redisChannelForSession(sessionId: string): string {
+  return `mcp:${sessionId}`;
+}
+
+export async function handleSSEConnection(req: Request, res: Response) {
+  const { server: mcpServer, cleanup: mcpCleanup }  = createMcpServer();
+  const transport = new SSEServerTransport("/message", res);
+  console.info(`[session ${transport.sessionId}] Received MCP SSE connection`);
+
+  const redisCleanup = await redisClient.createSubscription(
+    redisChannelForSession(transport.sessionId),
+    (json) => {
+      const message = JSON.parse(json);
+
+      if (message.method) {
+        if (message.method === "tools/call") {
+          console.info(
+            `[session ${transport.sessionId}] Processing ${message.method}, for tool ${message.params?.name}`,
+          );
+        } else {
+          console.info(
+            `[session ${transport.sessionId}] Processing ${message.method} method`,
+          );
+      }
+      } else if (message.error) {
+        console.warn(
+          `[session ${transport.sessionId}] Received error message: ${message.error.message}, ${message.error.code}`,
+        )
+      }
+      transport.handleMessage(message).catch((error) => {
+        console.error(
+          `[session ${transport.sessionId}] Error handling message:`,
+          error,
+        );
+      });
+    },
+    (error) => {
+      console.error(
+        `[session ${transport.sessionId}] Disconnecting due to error in Redis subscriber:`,
+        error,
+      );
+      transport
+        .close()
+        .catch((error) =>
+          console.error(
+            `[session ${transport.sessionId}] Error closing transport:`,
+            error,
+          ),
+        );
+    },
+  );
+
+  const cleanup = () => {
+    void mcpCleanup();
+    redisCleanup().catch((error) =>
+      console.error(
+        `[session ${transport.sessionId}] Error disconnecting Redis subscriber:`,
+        error,
+      ),
+    );
+  }
+
+  // Clean up Redis subscription when the connection closes
+  mcpServer.onclose = cleanup
+
+  console.info(`[session ${transport.sessionId}] Listening on Redis channel`);
+  await mcpServer.connect(transport);
+}
+
+export async function handleMessage(req: Request, res: Response) {
+  const sessionId = req.query.sessionId;
+  let body: string;
+  try {
+    if (typeof sessionId !== "string") {
+      throw new Error("Only one sessionId allowed");
+    }
+
+    const ct = contentType.parse(req.headers["content-type"] ?? "");
+    if (ct.type !== "application/json") {
+      throw new Error(`Unsupported content-type: ${ct}`);
+    }
+
+    body = await getRawBody(req, {
+      limit: MAXIMUM_MESSAGE_SIZE,
+      encoding: ct.parameters.charset ?? "utf-8",
+    });
+  } catch (error) {
+    res.status(400).json(error);
+    console.error("Bad POST request:", error);
+    return;
+  }
+  await redisClient.publish(redisChannelForSession(sessionId), body);
+  res.status(202).end();
+}

src/index.ts

@@ -0,0 +1,107 @@
+import cors from "cors";
+import express from "express";
+import { BASE_URI, PORT } from "./config.js";
+import { AuthRouterOptions, mcpAuthRouter } from "@modelcontextprotocol/sdk/server/auth/router.js";
+import { EverythingAuthProvider } from "./auth/provider.js";
+import { handleMessage, handleSSEConnection, authContext } from "./handlers/mcp.js";
+import { handleFakeAuthorizeRedirect, handleFakeAuthorize } from "./handlers/fakeauth.js";
+import { redisClient } from "./redis.js";
+import { requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js";
+
+const app = express();
+
+// Base security middleware - applied to all routes
+const baseSecurityHeaders = (req: express.Request, res: express.Response, next: express.NextFunction) => {
+  // Basic security headers
+  res.setHeader('X-Content-Type-Options', 'nosniff');
+  res.setHeader('X-Frame-Options', 'SAMEORIGIN');
+  res.setHeader('Strict-Transport-Security', 'max-age=31536000; includeSubDomains');
+
+  // Content Security Policy
+  const csp = [
+    "default-src 'self'",
+    "object-src 'none'",      // Disable plugins
+    "frame-ancestors 'none'", // No embedding
+    "form-action 'self'",     // Only allow forms to submit to our domain
+    "base-uri 'self'",       // Restrict base tag
+    "upgrade-insecure-requests",
+    "block-all-mixed-content"
+  ].join('; ');
+
+  res.setHeader('Content-Security-Policy', csp);
+  next();
+};
+
+// simple logging middleware
+const logger = (req: express.Request, res: express.Response, next: express.NextFunction) => {
+  console.log(`${req.method} ${req.url}`);
+  next();
+  // Log the response status code
+  res.on('finish', () => {
+    console.log(`Response status code: ${res.statusCode}`);
+  });
+};
+
+
+// Sensitive data middleware - for routes with sensitive data
+const sensitiveDataHeaders = (req: express.Request, res: express.Response, next: express.NextFunction) => {
+  res.setHeader('Cache-Control', 'no-store, max-age=0');
+  next();
+};
+
+// SSE middleware - specific for SSE endpoint
+const sseHeaders = (req: express.Request, res: express.Response, next: express.NextFunction) => {
+  res.setHeader('Content-Type', 'text/event-stream');
+  res.setHeader('Cache-Control', 'no-store, max-age=0');
+  res.setHeader('Connection', 'keep-alive');
+  next();
+};
+
+// Configure CORS to allow any origin since this is a public API service
+const corsOptions = {
+  origin: true, // Allow any origin
+  methods: ['GET', 'POST'],
+  allowedHeaders: ['Content-Type', 'Authorization', "MCP-Protocol-Version"],
+  credentials: true
+};
+
+app.use(logger);
+
+// Apply base security headers to all routes
+app.use(baseSecurityHeaders);
+
+// Enable CORS pre-flight requests
+app.options('*', cors(corsOptions));
+
+// Auth configuration
+const options: AuthRouterOptions = {
+  provider: new EverythingAuthProvider(),
+  issuerUrl: new URL(BASE_URI),
+  tokenOptions: {
+    rateLimit: {
+      windowMs: 5 * 1000,
+      limit: 100,
+    }
+  }
+};
+app.use(mcpAuthRouter(options));
+const bearerAuth = requireBearerAuth(options);
+
+// MCP routes
+app.get("/sse", cors(corsOptions), bearerAuth, authContext, sseHeaders, handleSSEConnection);
+app.post("/message", cors(corsOptions), bearerAuth, authContext, sensitiveDataHeaders, handleMessage);
+
+// Upstream auth routes
+app.get("/fakeupstreamauth/authorize", cors(corsOptions), handleFakeAuthorize);
+app.get("/fakeupstreamauth/callback", cors(corsOptions), handleFakeAuthorizeRedirect);
+
+try {
+  await redisClient.connect();
+} catch (error) {
+  console.error("Could not connect to Redis:", error);
+  process.exit(1);
+}
+
+app.listen(PORT, () => {
+  console.log(`Server running on http://localhost:${PORT}`);
+});

src/redis.ts

@@ -0,0 +1,206 @@
+import { createClient, SetOptions } from "@redis/client";
+
+/**
+ * Describes the Redis primitives we use in this application, to be able to mock
+ * them in tests (so we don't need to actually hit Redis).
+ */
+export interface RedisClient {
+  get(key: string): Promise<string | null>;
+  set(key: string, value: string, options?: SetOptions): Promise<string | null>;
+  getDel(key: string): Promise<string | null>;
+  connect(): Promise<void>;
+  on(event: string, callback: (error: Error) => void): void;
+  options?: { url: string };
+
+  /**
+   * Creates a pub/sub subscription. Returns a cleanup function to unsubscribe.
+   * Handles Redis client duplication and error handling internally.
+   */
+  createSubscription(
+    channel: string,
+    onMessage: (message: string) => void,
+    onError: (error: Error) => void,
+  ): Promise<() => Promise<void>>;
+
+  /**
+   * Publishes a message to a channel.
+   */
+  publish(channel: string, message: string): Promise<void>;
+}
+
+export class RedisClientImpl implements RedisClient {
+  private redis = createClient({
+    url: process.env.REDIS_URL || "redis://localhost:6379",
+    password: process.env.REDIS_PASSWORD,
+    socket: {
+      tls: process.env.REDIS_TLS === "1",
+      ca: process.env.REDIS_TLS_CA,
+    }
+  });
+
+  constructor() {
+    this.redis.on("error", (error) =>
+      console.error("Redis client error:", error),
+    );
+  }
+
+  async get(key: string): Promise<string | null> {
+    return await this.redis.get(key);
+  }
+
+  async getDel(key: string): Promise<string | null> {
+    return await this.redis.getDel(key);
+  }
+
+  async set(key: string, value: string, options?: SetOptions): Promise<string | null> {
+    return await this.redis.set(
+      key,
+      value,
+      options,
+    );
+  }
+
+  async connect(): Promise<void> {
+    await this.redis.connect();
+  }
+
+  on(event: string, callback: (error: Error) => void): void {
+    this.redis.on(event, callback);
+  }
+
+  get options() {
+    return { url: process.env.REDIS_URL || "redis://localhost:6379" };
+  }
+
+  async createSubscription(
+    channel: string,
+    onMessage: (message: string) => void,
+    onError: (error: Error) => void,
+  ): Promise<() => Promise<void>> {
+    const subscriber = this.redis.duplicate();
+    subscriber.on("error", onError);
+    await subscriber.connect();
+    await subscriber.subscribe(channel, onMessage);
+
+    return async () => {
+      await subscriber.disconnect();
+    };
+  }
+
+  async publish(channel: string, message: string): Promise<void> {
+    await this.redis.publish(channel, message);
+  }
+}
+
+// Export a mutable reference that can be swapped in tests
+export let redisClient: RedisClient = new RedisClientImpl();
+
+// Function to replace the Redis client (used in tests)
+export function setRedisClient(client: RedisClient) {
+  redisClient = client;
+}
+
+export class MockRedisClient implements RedisClient {
+  options = { url: "redis://localhost:6379" };
+  private store = new Map<string, string>();
+  private subscribers = new Map<string, ((message: string) => void)[]>();
+  private errorCallbacks = new Map<string, ((error: Error) => void)[]>();
+
+  async get(key: string): Promise<string | null> {
+    return this.store.get(key) ?? null;
+  }
+
+  async getDel(key: string): Promise<string | null> {
+    const value = this.store.get(key) ?? null;
+    this.store.delete(key);
+    return value;
+  }
+
+  async set(key: string, value: string, options?: SetOptions): Promise<string | null> {
+    let oldValue: string | null = null;
+    if (options?.GET) {
+      oldValue = this.store.get(key) ?? null;
+    }
+    this.store.set(key, value);
+    return oldValue;
+  }
+
+  async connect(): Promise<void> {
+    // No-op in mock
+  }
+
+  on(event: string, callback: (error: Error) => void): void {
+    if (event === "error") {
+      const callbacks = this.errorCallbacks.get("global") ?? [];
+      callbacks.push(callback);
+      this.errorCallbacks.set("global", callbacks);
+    }
+  }
+
+  async createSubscription(
+    channel: string,
+    onMessage: (message: string) => void,
+    onError: (error: Error) => void,
+  ): Promise<() => Promise<void>> {
+    const callbacks = this.subscribers.get(channel) ?? [];
+    callbacks.push(onMessage);
+    this.subscribers.set(channel, callbacks);
+
+    const errorCallbacks = this.errorCallbacks.get(channel) ?? [];
+    errorCallbacks.push(onError);
+    this.errorCallbacks.set(channel, errorCallbacks);
+
+    return async () => {
+      const callbacks = this.subscribers.get(channel) ?? [];
+      const index = callbacks.indexOf(onMessage);
+      if (index !== -1) {
+        callbacks.splice(index, 1);
+      }
+      this.subscribers.set(channel, callbacks);
+
+      if (onError) {
+        const errorCallbacks = this.errorCallbacks.get(channel) ?? [];
+        const errorIndex = errorCallbacks.indexOf(onError);
+        if (errorIndex !== -1) {
+          errorCallbacks.splice(errorIndex, 1);
+        }
+        this.errorCallbacks.set(channel, errorCallbacks);
+      }
+    };
+  }
+
+  async publish(channel: string, message: string): Promise<void> {
+    const callbacks = this.subscribers.get(channel) ?? [];
+    for (const callback of callbacks) {
+      try {
+        callback(message);
+      } catch (error) {
+        const errorCallbacks = this.errorCallbacks.get(channel) ?? [];
+        for (const errorCallback of errorCallbacks) {
+          errorCallback(error as Error);
+        }
+      }
+    }
+  }
+
+  clear() {
+    this.store.clear();
+    this.subscribers.clear();
+    this.errorCallbacks.clear();
+  }
+
+  // Helper method for tests to simulate Redis errors
+  simulateError(error: Error, channel?: string): void {
+    if (channel) {
+      const callbacks = this.errorCallbacks.get(channel) ?? [];
+      for (const callback of callbacks) {
+        callback(error);
+      }
+    } else {
+      const callbacks = this.errorCallbacks.get("global") ?? [];
+      for (const callback of callbacks) {
+        callback(error);
+      }
+    }
+  }
+}

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;
+}

src/types.ts

@@ -0,0 +1,30 @@
+import { OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
+
+// authorization code -> PendingAuthorization
+export interface PendingAuthorization {
+  redirectUri: string;
+  codeChallenge: string;
+  codeChallengeMethod: string;
+  clientId: string;
+  state?: string;
+}
+
+// authorization code -> MCP access token (once authorized)
+export interface TokenExchange {
+  mcpAccessToken: string;
+  alreadyUsed: boolean;
+}
+
+export interface FakeUpstreamInstallation {
+  fakeAccessTokenForDemonstration: string;
+  fakeRefreshTokenForDemonstration: string;
+}
+
+// This is the object stored in Redis holding the upstream "Installation" + all the relevant MCP tokens
+// It is stored encrypted by the MCP access token
+export interface McpInstallation {
+  fakeUpstreamInstallation: FakeUpstreamInstallation;
+  mcpTokens: OAuthTokens;
+  clientId: string;
+  issuedAt: number;
+}