feat: refactor state, store it securely

Author: nirinchev

package-lock.json

@@ -53,6 +53,7 @@
   "dependencies": {
     "@mongodb-js/devtools-connect": "^3.7.2",
     "@mongosh/service-provider-node-driver": "^3.6.0",
+    "@napi-rs/keyring": "^1.1.6",
     "@types/express": "^5.0.1",
     "bson": "^6.10.3",
     "mongodb": "^6.15.0",

package.json

@@ -8,21 +8,21 @@ export async function ensureAuthenticated(state: State, apiClient: ApiClient): P
 }
 
 export async function isAuthenticated(state: State, apiClient: ApiClient): Promise<boolean> {
-    switch (state.auth.status) {
+    switch (state.persistent.auth.status) {
         case "not_auth":
             return false;
         case "requested":
             try {
-                if (!state.auth.code) {
+                if (!state.persistent.auth.code) {
                     return false;
                 }
-                await apiClient.retrieveToken(state.auth.code.device_code);
-                return !!state.auth.token;
+                await apiClient.retrieveToken(state.persistent.auth.code.device_code);
+                return !!state.persistent.auth.token;
             } catch {
                 return false;
             }
         case "issued":
-            if (!state.auth.token) {
+            if (!state.persistent.auth.token) {
                 return false;
             }
             return await apiClient.validateToken();

src/common/atlas/auth.ts

@@ -20,14 +20,14 @@ export class Server {
         this.state = await loadState();
 
         this.apiClient = new ApiClient({
-            token: this.state?.auth.token,
+            token: this.state.persistent.auth.token,
             saveToken: async (token) => {
                 if (!this.state) {
                     throw new Error("State is not initialized");
                 }
-                this.state.auth.code = undefined;
-                this.state.auth.token = token;
-                this.state.auth.status = "issued";
+                this.state.persistent.auth.code = undefined;
+                this.state.persistent.auth.token = token;
+                this.state.persistent.auth.status = "issued";
                 await saveState(this.state);
             },
         });

src/server.ts

@@ -1,33 +1,51 @@
-import fs from "fs/promises";
-import config from "./config.js";
 import { OauthDeviceCode, OAuthToken } from "./common/atlas/apiClient.js";
+import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
+import { AsyncEntry } from "@napi-rs/keyring";
+import logger from "./logger.js";
+import { mongoLogId } from "mongodb-log-writer";
+
+const entry = new AsyncEntry("mongodb-mcp", "credentials");
 
 export interface State {
-    auth: {
-        status: "not_auth" | "requested" | "issued";
-        code?: OauthDeviceCode;
-        token?: OAuthToken;
+    persistent: {
+        auth: {
+            status: "not_auth" | "requested" | "issued";
+            code?: OauthDeviceCode;
+            token?: OAuthToken;
+        };
+        connectionString?: string;
+    };
+    session: {
+        serviceProvider?: NodeDriverServiceProvider;
     };
-    connectionString?: string;
 }
 
+const defaultState: State = {
+    persistent: {
+        auth: {
+            status: "not_auth",
+        },
+    },
+    session: {},
+};
+
 export async function saveState(state: State): Promise<void> {
-    await fs.writeFile(config.stateFile, JSON.stringify(state), { encoding: "utf-8" });
+    await entry.setPassword(JSON.stringify(state.persistent));
 }
 
 export async function loadState(): Promise<State> {
     try {
-        const data = await fs.readFile(config.stateFile, "utf-8");
-        return JSON.parse(data) as State;
-    } catch (err: unknown) {
-        if (err && typeof err === "object" && "code" in err && err.code === "ENOENT") {
-            return {
-                auth: {
-                    status: "not_auth",
-                },
-            };
+        const data = await entry.getPassword();
+        if (!data) {
+            return defaultState;
         }
 
-        throw err;
+        return {
+            persistent: JSON.parse(data),
+            session: {},
+        };
+    } catch (err: unknown) {
+        logger.error(mongoLogId(1_000_007), "state", `Failed to load state: ${err}`);
+        return defaultState;
     }
 }

src/state.ts

@@ -25,9 +25,9 @@ export class AuthTool extends AtlasToolBase {
         try {
             const code = await this.apiClient.authenticate();
 
-            this.state.auth.status = "requested";
-            this.state.auth.code = code;
-            this.state.auth.token = undefined;
+            this.state.persistent.auth.status = "requested";
+            this.state.persistent.auth.code = code;
+            this.state.persistent.auth.token = undefined;
 
             await saveState(this.state);
 

src/tools/atlas/auth.ts

@@ -21,7 +21,7 @@ export class ConnectTool extends MongoDBToolBase {
     protected async execute({
         connectionStringOrClusterName,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        connectionStringOrClusterName ??= this.state.connectionString;
+        connectionStringOrClusterName ??= this.state.persistent.connectionString;
         if (!connectionStringOrClusterName) {
             return {
                 content: [
@@ -71,8 +71,8 @@ export class ConnectTool extends MongoDBToolBase {
             productName: "MongoDB MCP",
         });
 
-        this.mongodbState.serviceProvider = provider;
-        this.state.connectionString = connectionString;
+        this.state.session.serviceProvider = provider;
+        this.state.persistent.connectionString = connectionString;
         await saveState(this.state);
     }
 }