Merge pull request #3 from Muimi-Chat/staging

Use new key management service for encryptions

.env.example

@@ -14,10 +14,6 @@ OPENAI_ORGANIZATION_ID=org-id
 # python -c "import secrets; print('Pepper:', secrets.token_hex(16))"
 HASHING_PEPPER=asdasdasdasdasdsadas
 
-# Used for encryption
-# python -c 'import secrets; print("AES-128 Key:", secrets.token_bytes(16).hex())'
-ENCRYPTION_KEY=asdasdasdasdasdasdasd
-
 # Default starting token amount by users.
 DEFAULT_STARTING_TOKEN=25000
 
@@ -45,4 +41,9 @@ QDRANT_API_KEY=asdasdasd
 # For communication with User service; Should be same as whats set in user service's
 USER_SERVICE_API_TOKEN=abcdefg
 USER_SERVICE_API_DOMAIN=user-api:8000
-USER_SERVICE_SSL_ENABLED=false
+USER_SERVICE_SSL_ENABLED=false
+
+# API Auth Key from Cappu Crypt
+CAPPU_CRYPT_API_KEY=changeme
+CAPPU_CRYPT_HOST=cappu-api
+CAPPU_CRYPT_PORT=8080

src/api/consumers/chatMessageConsumer.ts

@@ -26,6 +26,9 @@ import insertReturningMessage from '../repositories/insertReturningMessage';
 import selectAccountByUUID from '../repositories/selectAccountByUUID';
 import { db } from 'src/db';
 import updateAccountTokenByUUID from '../repositories/updateAccountTokenByUUID';
+import requestEncrypt from '../services/crypt/requestEncrypt';
+import { REDIS_CONNECTION_STRING } from 'src/configs/redisConnectionString';
+import { createClient } from 'redis';
 
 function getSenderUserUUID(data: any): Promise<string | null> {
     const sessionToken: string = data.session_token
@@ -37,7 +40,17 @@ function getSenderUserUUID(data: any): Promise<string | null> {
 
 // TODO: Support custom system messages...
 export default async function chatMessageConsumer(socketClient: WebSocket, content: RawData) {
+    const redisClient = createClient({
+        url: REDIS_CONNECTION_STRING
+    })
+    .on('error', async err => {
+        console.error('Connecting Redis Client Error', err);
+        await insertLog(`Error connecting to Redis :: ${err}`);
+    })
+
     try {
+        await redisClient.connect();
+
         const data = JSON.parse(content.toString())
 
         const userUUID = await getSenderUserUUID(data)
@@ -46,6 +59,7 @@ export default async function chatMessageConsumer(socketClient: WebSocket, conte
             socketClient.send(JSON.stringify({ status: "BAD_SESSION", message: "Bad session token, user should login again!" }));
             return
         }
+
         const matchingAccounts = await selectAccountByUUID(userUUID)
         if (matchingAccounts.length <= 0) {
             const warningMessage = `Client tried to login as ${data.username} (${userUUID}). But failed to match UUID in database.`
@@ -80,16 +94,28 @@ export default async function chatMessageConsumer(socketClient: WebSocket, conte
         }
 
         const conversationID: number | null = data.conversation_id
-        const attachments: MessageAttachment[] = data.attachments
-
-        // TODO: Handle attachment size/input validation.
+
+        // TODO: Handle attachments 
 
         const openAiClient = CreateOpenAiClient()
         if (openAiClient === null) {
             console.error("An incoming message on websocket consumer was ignored due to unset OpenAI keys!")
             return
         }
 
+        // validate from cache if user is currently sending message
+        const userSendingMessage = await redisClient.get(`${userUUID}_is_messaging`)
+        if (userSendingMessage != null) {
+            await insertLog(`User ${data.username} (${userUUID}) tried to send message but is already sending one.`, "WARNING")
+            console.warn(`User (${data.username}) tried to send message but is already sending one.`)
+            socketClient.send(JSON.stringify({ status: "ALREADY_SENDING", message: "User is already sending a message!" }));
+            return
+        }
+        await redisClient.set(`${userUUID}_is_messaging`, "true", { 
+            EX: 20,
+            NX: true
+        })
+
         // Create conversation if new conversation,
         // otherwise, try to fetch from database.
         let conversationObject;
@@ -128,7 +154,7 @@ export default async function chatMessageConsumer(socketClient: WebSocket, conte
         // Form a message chain, to let the AI know histories + current message.
         // TODO: Handle attachments, images, document (RAG)
 
-        const messageChain = await createMessageHistoryChain(conversationObject.id)
+        const messageChain = await createMessageHistoryChain(userUUID, conversationObject.id)
         messageChain.push({
             role: "user",
             content: message
@@ -170,8 +196,16 @@ export default async function chatMessageConsumer(socketClient: WebSocket, conte
             if (!userAccountModel.freeTokenUsage) {
                 await updateAccountTokenByUUID(userUUID, userAccountModel.token - promptTokenCost - completionTokenCost)
             }
-            userMessageObject = (await insertReturningMessage(conversationObject.id, message, promptTokenCost, "USER"))[0]
-            botMessageObject = (await insertReturningMessage(conversationObject.id, botReply, completionTokenCost, "BOT", chatModel))[0]
+
+            const encryptPromises = [
+                requestEncrypt(userUUID, message, userUUID),
+                requestEncrypt(userUUID, botReply, userUUID)
+            ];
+
+            const [encryptedMessage, encryptedBotReply] = await Promise.all(encryptPromises);
+
+            userMessageObject = (await insertReturningMessage(conversationObject.id, encryptedMessage, promptTokenCost, "USER"))[0]
+            botMessageObject = (await insertReturningMessage(conversationObject.id, encryptedBotReply, promptTokenCost, "BOT", chatModel))[0]
         })
 
         socketClient.send(JSON.stringify({ 
@@ -187,9 +221,13 @@ export default async function chatMessageConsumer(socketClient: WebSocket, conte
                 token_cost: completionTokenCost
             }
         }));
+
+        await redisClient.del(`${userUUID}_is_messaging`)
     } catch (err) {
         const logMessage = `Failed to handle chat message consumer :: ${err}`
         console.error(logMessage)
         await insertLog(logMessage,  "CRITICAL")
+    } finally {
+        await redisClient.disconnect()
     }
 }

src/api/controllers/getConversationMessages.ts

@@ -5,6 +5,7 @@ import fetchUserInformation from "../services/fetchUserInformation";
 import { REDIS_CONNECTION_STRING } from "src/configs/redisConnectionString";
 import { selectMessagesByConversationID } from "../repositories/selectMessagesByConversationID";
 import selectConversationByID from "../repositories/selectConversationByID";
+import requestDecrypt from "../services/crypt/requestDecrypt";
 
 export default async function getConversationMessages(req: Request, res: Response) {
   const client = createClient({
@@ -86,18 +87,19 @@ export default async function getConversationMessages(req: Request, res: Respons
       const targetConversation = matchingConversations[0]
       const messages = await selectMessagesByConversationID(
         targetConversation.id,
-        100
+        300
       );
 
-      const result = messages.map((message) => ({
+      const resultPromises = messages.reverse().map(async (message) => ({
         id: message.id,
-        content: message.content,
+        content: await requestDecrypt(userUUID, message.encryptedContent, userUUID),
         author: message.sender,
         token_cost: message.tokenCost,
         created_at: message.creationDate,
         bot_model: message.botModel,
       }));
 
+      const result = await Promise.all(resultPromises);
       client.quit();
 
       return res.status(200).json({

src/api/helpers/createMessageHistoryChain.ts

@@ -1,25 +1,28 @@
 import { selectMessagesByConversationID } from "../repositories/selectMessagesByConversationID";
+import requestDecrypt from "../services/crypt/requestDecrypt";
 
-export default async function createMessageHistoryChain(conversationID: number): Promise<{
+/**
+ * Create a message history chain from the given conversation ID.
+ * (Each of the output message in the chain is decrypted)
+ * 
+ */
+export default async function createDecryptedMessageHistoryChain(userUUID: string, conversationID: number): Promise<{
     role: "assistant" | "user" | "system",
     content: string
 }[]> {
     const messages = await selectMessagesByConversationID(conversationID)
 
-    const messageChain = []
-    for (let i = messages.length - 1; i >= 0; i--) {
-        const message = messages[i];
-
-        let author = message.sender!.toLowerCase()
-        if (author == "bot") {
-            author = "assistant"
+    const decryptionPromises = messages.map(async (message) => {
+        let author = message.sender!.toLowerCase();
+        if (author === "bot") {
+            author = "assistant";
         }
 
-        messageChain.push({
+        return {
             role: author as ("assistant" | "user" | "system"),
-            content: message.content
-        })
-    }
+            content: await requestDecrypt(userUUID, message.encryptedContent, userUUID)
+        };
+    });
 
-    return messageChain
+    return Promise.all(decryptionPromises);
 }

src/api/repositories/insertReturningMessage.ts

@@ -10,7 +10,7 @@ export default async function insertReturningMessage(
 ) {
     return await db.insert(message).values({
         conversationID: conversationID,
-        content: content,
+        encryptedContent: content,
         tokenCost: tokenCost,
         sender: sender,
         botModel: botModel

src/api/repositories/selectMessagesByConversationID.ts

@@ -17,7 +17,7 @@ export async function selectMessagesByConversationID(conversationId: number, lim
             .select()
             .from(message)
             .where(eq(message.conversationID, conversationId))
-            .orderBy(desc(message.creationDate))
+            .orderBy(desc(message.creationDate), desc(message.id))
             .limit(limit)
             .offset(offset)
     } catch (error) {

src/api/services/crypt/requestDecrypt.ts

@@ -0,0 +1,33 @@
+import insertLog from "src/api/repositories/insertLog";
+import { CAPPU_AUTH_KEY, CAPPU_BASE_HTTP } from "src/configs/cryptorConfig";
+
+export default async function requestDecrypt(id: string, content: string, metadata: string = "") {
+    try {
+        const formData = new FormData();
+        formData.append('content', content);
+        formData.append('id', id);
+        formData.append('metadata', metadata);
+
+        const response = await fetch(`${CAPPU_BASE_HTTP}/crypt/decrypt`, {
+            method: 'POST',
+            headers: {
+                'Authorization': CAPPU_AUTH_KEY
+            },
+            body: formData
+        })
+
+        const data = await response.json() as any
+        if (!response.ok) {
+            const errorMessage = `Non Success during decryption of content on id (${id}) :: ${data}`
+            console.error(errorMessage)
+            await insertLog(errorMessage, "ERROR")
+            return
+        }
+
+        return data.decryptedContent;
+    } catch (error) {
+        const errorMessage = `Error decrypting content on id (${id}) :: ${error}`
+        console.error(errorMessage)
+        await insertLog(errorMessage, "CRITICAL")
+    }
+}

src/api/services/crypt/requestEncrypt.ts

@@ -0,0 +1,33 @@
+import insertLog from "src/api/repositories/insertLog";
+import { CAPPU_AUTH_KEY, CAPPU_BASE_HTTP } from "src/configs/cryptorConfig";
+
+export default async function requestEncrypt(id: string, content: string, metadata: string = "") {
+    try {
+        const formData = new FormData();
+        formData.append('content', content);
+        formData.append('id', id);
+        formData.append('metadata', metadata);
+
+        const response = await fetch(`${CAPPU_BASE_HTTP}/crypt/encrypt`, {
+            method: 'POST',
+            headers: {
+                'Authorization': CAPPU_AUTH_KEY
+            },
+            body: formData
+        })
+
+        const data = await response.json() as any
+        if (!response.ok) {
+            const errorMessage = `Non Success during encryption of content on id (${id}) :: ${data}`
+            console.error(errorMessage)
+            await insertLog(errorMessage, "ERROR")
+            return
+        }
+
+        return data.encryptedContent;
+    } catch (error) {
+        const errorMessage = `Error encrypting content on id (${id}) :: ${error}`
+        console.error(errorMessage)
+        await insertLog(errorMessage, "CRITICAL")
+    }
+}

src/configs/chatInputLimitation.ts

@@ -1,5 +1,5 @@
 export const CHAT_INPUT_LIMITATION = {
-    MAX_MESSAGE_LENGTH: 2500,
+    MAX_MESSAGE_LENGTH: 2000,
     MAX_MESSAGE_ATTACHMENTS_COUNT: 5,
     MAX_ATTACHMENT_SIZE_MB: 15
 }

src/configs/cryptorConfig.ts

@@ -0,0 +1,7 @@
+const authApiKey = process.env.CAPPU_CRYPT_API_KEY || "CRYPT API KEY NOT SET"
+
+
+const cryptHost = process.env.CAPPU_CRYPT_HOST || "CRYPT HOST NOT SET"
+const crypyPort = process.env.CAPPU_CRYPT_PORT || "CRYPT PORT NOT SET"
+export const CAPPU_AUTH_KEY = authApiKey
+export const CAPPU_BASE_HTTP = `http://${cryptHost}:${crypyPort}`

src/schema.ts

@@ -39,7 +39,7 @@ export const messageAuthorEnum = pgEnum('message_author_enum',
 export const message = pgTable('message', {
 	id: serial('id').primaryKey(),
 	conversationID: serial('conversation_id').references(() => conversation.id, {onDelete: 'restrict'}).notNull(),
-	content: text('content').notNull(),
+	encryptedContent: text('encrypted_content').notNull(),
 	tokenCost: integer('token_cost').notNull(),
 	sender: messageAuthorEnum('sender').notNull(),
 	botModel: text('bot_model'),