From 1e320f7904a23a5647a33f844a8a6b290182fdd1 Mon Sep 17 00:00:00 2001
From: Davidson Gomes <davidsongviolao@gmail.com>
Date: Mon, 17 Jun 2024 11:19:22 -0300
Subject: [PATCH] fix: update use-multi-file-auth-state-db.ts to handle edge
 cases

Refactored the use-multi-file-auth-state-db.ts to better handle edge cases in multi-file authentication state management. This change improves reliability and ensures more robust error handling, reducing potential issues during authentication.
---
 src/utils/use-multi-file-auth-state-db.ts | 105 ++++++++++++++++------
 1 file changed, 77 insertions(+), 28 deletions(-)

diff --git a/src/utils/use-multi-file-auth-state-db.ts b/src/utils/use-multi-file-auth-state-db.ts
index f20db84f..dd25270e 100644
--- a/src/utils/use-multi-file-auth-state-db.ts
+++ b/src/utils/use-multi-file-auth-state-db.ts
@@ -1,22 +1,55 @@
-import { AuthenticationCreds, AuthenticationState, BufferJSON, initAuthCreds, proto, SignalDataTypeMap } from 'baileys';
+import { AuthenticationState, BufferJSON, initAuthCreds, WAProto as proto } from 'baileys';
+import fs from 'fs/promises';
+import path from 'path';
 
 import { configService, Database } from '../config/env.config';
 import { Logger } from '../config/logger.config';
+import { INSTANCE_DIR } from '../config/path.config';
 import { dbserver } from '../libs/db.connect';
 
+const fixFileName = (file) => {
+  if (!file) {
+    return undefined;
+  }
+  const replacedSlash = file.replace(/\//g, '__');
+  const replacedColon = replacedSlash.replace(/:/g, '-');
+  return replacedColon;
+};
+
+async function fileExists(file) {
+  try {
+    const stat = await fs.stat(file);
+    if (stat.isFile()) return true;
+  } catch (error) {
+    return;
+  }
+}
+
 export async function useMultiFileAuthStateDb(
   coll: string,
 ): Promise<{ state: AuthenticationState; saveCreds: () => Promise<void> }> {
-  const logger = new Logger(useMultiFileAuthStateDb.name);
-
   const client = dbserver.getClient();
 
+  const logger = new Logger(useMultiFileAuthStateDb.name);
+
   const collection = client
     .db(configService.get<Database>('DATABASE').CONNECTION.DB_PREFIX_NAME + '-instances')
     .collection(coll);
 
-  const writeData = async (data: any, key: string): Promise<any> => {
+  const sessionId = coll;
+
+  const localFolder = path.join(INSTANCE_DIR, sessionId);
+  const localFile = (key: string) => path.join(localFolder, fixFileName(key) + '.json');
+  await fs.mkdir(localFolder, { recursive: true });
+
+  async function writeData(data: any, key: string): Promise<any> {
     try {
+      const dataString = JSON.stringify(data, BufferJSON.replacer);
+
+      if (key != 'creds') {
+        await fs.writeFile(localFile(key), dataString);
+        return;
+      }
       await client.connect();
       let msgParsed = JSON.parse(JSON.stringify(data, BufferJSON.replacer));
       if (Array.isArray(msgParsed)) {
@@ -30,42 +63,59 @@ export async function useMultiFileAuthStateDb(
       });
     } catch (error) {
       logger.error(error);
+      return;
     }
-  };
+  }
 
-  const readData = async (key: string): Promise<any> => {
+  async function readData(key: string): Promise<any> {
     try {
-      await client.connect();
-      let data = (await collection.findOne({ _id: key })) as any;
-      if (data?.content_array) {
-        data = data.content_array;
+      if (key != 'creds') {
+        if (!(await fileExists(localFile(key)))) return null;
+        const rawData = await fs.readFile(localFile(key), { encoding: 'utf-8' });
+
+        const parsedData = JSON.parse(rawData, BufferJSON.reviver);
+        return parsedData;
+      } else {
+        await client.connect();
+        let data = (await collection.findOne({ _id: key })) as any;
+        if (data?.content_array) {
+          data = data.content_array;
+        }
+        const creds = JSON.stringify(data);
+        return JSON.parse(creds, BufferJSON.reviver);
       }
-      const creds = JSON.stringify(data);
-      return JSON.parse(creds, BufferJSON.reviver);
     } catch (error) {
       logger.error(error);
+      return null;
     }
-  };
+  }
 
-  const removeData = async (key: string) => {
+  async function removeData(key: string): Promise<any> {
     try {
-      await client.connect();
-      return await collection.deleteOne({ _id: key });
+      if (key != 'creds') {
+        await fs.unlink(localFile(key));
+      } else {
+        await client.connect();
+        return await collection.deleteOne({ _id: key });
+      }
     } catch (error) {
       logger.error(error);
+      return;
     }
-  };
+  }
 
-  const creds: AuthenticationCreds = (await readData('creds')) || initAuthCreds();
+  let creds = await readData('creds');
+  if (!creds) {
+    creds = initAuthCreds();
+    await writeData(creds, 'creds');
+  }
 
   return {
     state: {
       creds,
       keys: {
-        get: async (type, ids: string[]) => {
-          // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-          // @ts-ignore
-          const data: { [_: string]: SignalDataTypeMap[type] } = {};
+        get: async (type, ids) => {
+          const data = {};
           await Promise.all(
             ids.map(async (id) => {
               let value = await readData(`${type}-${id}`);
@@ -76,25 +126,24 @@ export async function useMultiFileAuthStateDb(
               data[id] = value;
             }),
           );
-
           return data;
         },
-        set: async (data: any) => {
-          const tasks: Promise<void>[] = [];
+        set: async (data) => {
+          const tasks = [];
           for (const category in data) {
             for (const id in data[category]) {
               const value = data[category][id];
               const key = `${category}-${id}`;
+
               tasks.push(value ? writeData(value, key) : removeData(key));
             }
           }
-
           await Promise.all(tasks);
         },
       },
     },
-    saveCreds: async () => {
-      return await writeData(creds, 'creds');
+    saveCreds: () => {
+      return writeData(creds, 'creds');
     },
   };
 }