feat: refactor state, store it securely (#48)

Author: nirinchev

package-lock.json

@@ -59,6 +59,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.credentials.auth.status) {
         case "not_auth":
             return false;
         case "requested":
             try {
-                if (!state.auth.code) {
+                if (!state.credentials.auth.code) {
                     return false;
                 }
-                await apiClient.retrieveToken(state.auth.code.device_code);
-                return !!state.auth.token;
+                await apiClient.retrieveToken(state.credentials.auth.code.device_code);
+                return !!state.credentials.auth.token;
             } catch {
                 return false;
             }
         case "issued":
-            if (!state.auth.token) {
+            if (!state.credentials.auth.token) {
                 return false;
             }
             return await apiClient.validateToken();

src/common/atlas/auth.ts

@@ -1,6 +1,6 @@
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { ApiClient } from "./common/atlas/apiClient.js";
-import { State, saveState, loadState } from "./state.js";
+import defaultState, { State } from "./state.js";
 import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import { registerAtlasTools } from "./tools/atlas/tools.js";
 import { registerMongoDBTools } from "./tools/mongodb/index.js";
@@ -9,26 +9,27 @@ import logger, { initializeLogger } from "./logger.js";
 import { mongoLogId } from "mongodb-log-writer";
 
 export class Server {
-    state: State | undefined = undefined;
+    state: State = defaultState;
     apiClient: ApiClient | undefined = undefined;
     initialized: boolean = false;
 
     private async init() {
         if (this.initialized) {
             return;
         }
-        this.state = await loadState();
+
+        await this.state.loadCredentials();
 
         this.apiClient = new ApiClient({
-            token: this.state?.auth.token,
+            token: this.state.credentials.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";
-                await saveState(this.state);
+                this.state.credentials.auth.code = undefined;
+                this.state.credentials.auth.token = token;
+                this.state.credentials.auth.status = "issued";
+                await this.state.persistCredentials();
             },
         });
 
@@ -43,8 +44,8 @@ export class Server {
 
         server.server.registerCapabilities({ logging: {} });
 
-        registerAtlasTools(server, this.state!, this.apiClient!);
-        registerMongoDBTools(server, this.state!);
+        registerAtlasTools(server, this.state, this.apiClient!);
+        registerMongoDBTools(server, this.state);
 
         return server;
     }

src/server.ts

@@ -1,8 +1,10 @@
-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";
 
-export interface State {
+interface Credentials {
     auth: {
         status: "not_auth" | "requested" | "issued";
         code?: OauthDeviceCode;
@@ -11,23 +13,33 @@ export interface State {
     connectionString?: string;
 }
 
-export async function saveState(state: State): Promise<void> {
-    await fs.writeFile(config.stateFile, JSON.stringify(state), { encoding: "utf-8" });
-}
+export class State {
+    private entry = new AsyncEntry("mongodb-mcp", "credentials");
+    credentials: Credentials = {
+        auth: {
+            status: "not_auth",
+        },
+    };
+    serviceProvider?: NodeDriverServiceProvider;
 
-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",
-                },
-            };
-        }
+    public async persistCredentials(): Promise<void> {
+        await this.entry.setPassword(JSON.stringify(this.credentials));
+    }
 
-        throw err;
+    public async loadCredentials(): Promise<boolean> {
+        try {
+            const data = await this.entry.getPassword();
+            if (data) {
+                this.credentials = JSON.parse(data);
+            }
+
+            return true;
+        } catch (err: unknown) {
+            logger.error(mongoLogId(1_000_007), "state", `Failed to load state: ${err}`);
+            return false;
+        }
     }
 }
+
+const defaultState = new State();
+export default defaultState;