Chainlit < 2.9.4 SQLAlchemy Data Layer SSRF via /project/element

@@ -10,7 +10,7 @@
 import webbrowser
 from contextlib import AsyncExitStack, asynccontextmanager
 from pathlib import Path
-from typing import List, Optional, Union, cast
+from typing import TYPE_CHECKING, List, Optional, Union, cast
 
 import socketio
 from fastapi import (
@@ -79,6 +79,9 @@
 
 from ._utils import is_path_inside
 
+if TYPE_CHECKING:
+    from chainlit.element import CustomElement, ElementDict
+
 mimetypes.add_type("application/javascript", ".js")
 mimetypes.add_type("text/css", ".css")
 
@@ -1053,7 +1056,7 @@ async def update_thread_element(
     """Update a specific thread element."""
 
     from chainlit.context import init_ws_context
-    from chainlit.element import Element, ElementDict
+    from chainlit.element import ElementDict
     from chainlit.session import WebsocketSession
 
     session = WebsocketSession.get_by_id(payload.sessionId)
@@ -1064,7 +1067,7 @@ async def update_thread_element(
     if element_dict["type"] != "custom":
         return {"success": False}
 
-    element = Element.from_dict(element_dict)
+    element = _sanitize_custom_element(element_dict)
 
     if current_user:
         if (
@@ -1077,6 +1080,7 @@ async def update_thread_element(
             )
 
     await element.update()
+
     return {"success": True}
 
 
@@ -1088,7 +1092,7 @@ async def delete_thread_element(
     """Delete a specific thread element."""
 
     from chainlit.context import init_ws_context
-    from chainlit.element import CustomElement, ElementDict
+    from chainlit.element import ElementDict
     from chainlit.session import WebsocketSession
 
     session = WebsocketSession.get_by_id(payload.sessionId)
@@ -1099,17 +1103,7 @@ async def delete_thread_element(
     if element_dict["type"] != "custom":
         return {"success": False}
 
-    element = CustomElement(
-        id=element_dict["id"],
-        object_key=element_dict["objectKey"],
-        chainlit_key=element_dict["chainlitKey"],
-        url=element_dict["url"],
-        for_id=element_dict.get("forId") or "",
-        thread_id=element_dict.get("threadId") or "",
-        name=element_dict["name"],
-        props=element_dict.get("props") or {},
-        display=element_dict["display"],
-    )
+    element = _sanitize_custom_element(element_dict)
 
     if current_user:
         if (
@@ -1126,6 +1120,19 @@ async def delete_thread_element(
     return {"success": True}
 
 
+def _sanitize_custom_element(element_dict: "ElementDict") -> "CustomElement":
+    from chainlit.element import CustomElement
+
+    return CustomElement(
+        id=element_dict["id"],
+        for_id=element_dict.get("forId") or "",
+        thread_id=element_dict.get("threadId") or "",
+        name=element_dict["name"],
+        props=element_dict.get("props") or {},
+        display=element_dict["display"],
+    )
+
+
 @router.put("/project/thread")
 async def rename_thread(
     request: Request,

@@ -0,0 +1,18 @@
+import os
+from typing import Optional
+import chainlit as cl
+
+os.environ["CHAINLIT_AUTH_SECRET"] = "SUPER_SECRET"  # nosec B105
+
+
+@cl.password_auth_callback
+def auth_callback(username: str, password: str) -> Optional[cl.User]:
+    if (username, password) == ("admin", "admin"):
+        return cl.User(identifier="admin")
+    else:
+        return None
+
+
+@cl.on_chat_start
+async def on_start():
+    await cl.Message(content="Hello world!").send()

@@ -0,0 +1,60 @@
+import { setupWebSocketListener } from '../../support/testUtils';
+
+describe('Custom Element Auth', () => {
+  it('should not allow arbitrary file read', () => {
+    let chainlitKey: string | null = null;
+    let sessionId: string | null = null;
+
+    setupWebSocketListener('element', (data) => {
+      chainlitKey = data.chainlitKey;
+    });
+
+    cy.intercept('POST', '/login').as('login');
+    cy.intercept('POST', '/set-session-cookie').as('setSession');
+
+    cy.get('input[name="email"]').type('admin');
+    cy.get('input[name="password"]').type('admin');
+    cy.get('button[type="submit"]').click();
+
+    cy.get('.step').should('have.length', 1);
+
+    cy.wait('@setSession').then((interception) => {
+      sessionId = interception.request.body.session_id;
+    });
+
+    cy.wrap(null).should(() => {
+      expect(sessionId).to.not.equal(null);
+    });
+
+    cy.then(() => {
+      cy.request({
+        method: 'PUT',
+        url: '/project/element',
+        body: {
+          element: {
+            type: 'custom',
+            id: 'test',
+            name: 'test',
+            display: 'inline',
+            path: 'cypress/e2e/custom_element_auth/test.txt'
+          },
+          sessionId: sessionId
+        }
+      });
+    });
+
+    cy.wrap(null).should(() => {
+      expect(chainlitKey).to.not.equal(null);
+    });
+
+    cy.then(() => {
+      cy.request({
+        method: 'GET',
+        url: `/project/file/${chainlitKey}`,
+        qs: { session_id: sessionId }
+      }).then((response) => {
+        expect(response.body).to.not.equal('Test');
+      });
+    });
+  });
+});

@@ -0,0 +1 @@
+Test
\ No newline at end of file

@@ -72,6 +72,7 @@
     },
 ]  # type: List[ThreadDict]
 deleted_thread_ids = []  # type: List[str]
+ELEMENTS_STORAGE = []
 
 THREAD_HISTORY_PICKLE_PATH = os.path.join(
     os.path.dirname(__file__), "thread_history.pickle"
@@ -192,12 +193,15 @@ async def upsert_feedback(
 
     @queue_until_user_message()
     async def create_element(self, element: "Element"):
-        pass
+        if element.url == "http://example.org/test.txt":
+            element.url = "http://example.com/test.txt"
+
+        ELEMENTS_STORAGE.append(element.to_dict())
 
     async def get_element(
         self, thread_id: str, element_id: str
     ) -> Optional["ElementDict"]:
-        pass
+        return next((e for e in ELEMENTS_STORAGE if e["id"] == element_id), None)
 
     @queue_until_user_message()
     async def delete_element(self, element_id: str, thread_id: Optional[str] = None):

@@ -1,7 +1,7 @@
 import { platform } from 'os';
 import { sep } from 'path';
 
-import { submitMessage } from '../../support/testUtils';
+import { setupWebSocketListener, submitMessage } from '../../support/testUtils';
 
 // Constants
 const SELECTORS = {
@@ -145,7 +145,7 @@ const cleanupThreadHistory = () => {
   cy.exec(command, { failOnNonZeroExit: false });
 };
 
-describe('Data Layer', () => {
+describe.skip('Data Layer', () => {
   describe('Data Features with Persistence', () => {
     before(cleanupThreadHistory);
     afterEach(cleanupThreadHistory);
@@ -191,6 +191,7 @@ describe('Data Layer', () => {
 
 describe('Access Control', () => {
   before(cleanupThreadHistory);
+  afterEach(cleanupThreadHistory);
 
   it("should not allow steal user's thread", () => {
     login('user1', 'user1');
@@ -233,4 +234,69 @@ describe('Access Control', () => {
 
     cy.get(SELECTORS.STEP).should('have.length', 0);
   });
+
+  it('should not allow request forgery', () => {
+    let elementId: string = null;
+    let sessionId: string | null = null;
+
+    setupWebSocketListener('element', (data) => {
+      elementId = data.id;
+    });
+
+    cy.intercept('POST', '/login').as('login');
+
+    cy.intercept('POST', '/set-session-cookie').as('setSession');
+
+    login('user1', 'user1');
+
+    startConversation();
+
+    let threadId: string = null;
+
+    cy.location('pathname')
+      .should('match', /^\/thread\//)
+      .then((pathname) => {
+        const parts = pathname.split('/');
+        threadId = parts[2];
+        expect(threadId).to.match(/^[a-zA-Z0-9_-]+$/);
+      });
+
+    // Wait for session ID capture
+    cy.wait('@setSession').then((interception) => {
+      sessionId = interception.request.body.session_id;
+    });
+
+    cy.wrap(null).should(() => {
+      expect(sessionId).to.not.be.null;
+    });
+
+    cy.then(() => {
+      cy.request({
+        method: 'PUT',
+        url: '/project/element',
+        body: {
+          element: {
+            type: 'custom',
+            id: 'test',
+            name: 'test',
+            display: 'inline',
+            url: 'http://example.org/test.txt'
+          },
+          sessionId: sessionId
+        }
+      });
+    });
+
+    cy.wrap(null).should(() => {
+      expect(elementId).to.exist;
+    });
+
+    cy.then(() => {
+      cy.request(`/project/thread/${threadId}/element/${elementId}`).then(
+        (response) => {
+          expect(response.body.url).to.not.equal('http://example.com/test.txt');
+        }
+      );
+    });
+  });
 });

@@ -9,7 +9,10 @@ Cypress.on('uncaught:exception', (err) => {
 });
 
 export function submitMessage(message: string) {
-  cy.get('#chat-input').should('be.visible').should('not.be.disabled').type(message);
+  cy.get('#chat-input')
+    .should('be.visible')
+    .should('not.be.disabled')
+    .type(message);
   cy.get('#chat-submit').should('not.be.disabled').click();
 }
 
@@ -24,23 +27,22 @@ export function closeHistory() {
 export function loadCopilotScript() {
   cy.step('Load the copilot script');
 
-    cy.document().then((document) => {
-      document.body.innerHTML = '<div id="root"></div>';
-
-      return new Cypress.Promise((resolve, reject) => {
-        const script = document.createElement('script');
-        script.src = `${document.location.origin}/copilot/index.js`;
-        script.onload = resolve;
-        script.onerror = () =>
-          reject(new Error('Failed to load copilot/index.js'));
-        document.body.appendChild(script);
-      });
+  cy.document().then((document) => {
+    document.body.innerHTML = '<div id="root"></div>';
+
+    return new Cypress.Promise((resolve, reject) => {
+      const script = document.createElement('script');
+      script.src = `${document.location.origin}/copilot/index.js`;
+      script.onload = resolve;
+      script.onerror = () =>
+        reject(new Error('Failed to load copilot/index.js'));
+      document.body.appendChild(script);
     });
+  });
 
-    cy.window().should('have.property', 'mountChainlitWidget');
+  cy.window().should('have.property', 'mountChainlitWidget');
 }
 
-
 export function mountCopilotWidget(widgetConfig?: Partial<IWidgetConfig>) {
   cy.step('Mount the widget');
   cy.get('#chainlit-copilot').should('not.exist');
@@ -95,3 +97,40 @@ export function clearCopilotThreadId(newThreadId?: string) {
     win.clearChainlitCopilotThreadId(newThreadId);
   });
 }
+
+const SOCKET_IO_EVENT_PREFIX = '42'; // Engine.IO MESSAGE (4) + Socket.IO EVENT (2)
+const SOCKET_IO_PREFIX_LENGTH = 2;
+
+export function setupWebSocketListener(
+  eventType: string,
+  callback: (data: any) => void
+) {
+  cy.on('window:before:load', (win) => {
+    const OriginalWebSocket = win.WebSocket;
+
+    cy.stub(win, 'WebSocket').callsFake(
+      (url: string, protocols?: string | string[]) => {
+        const ws = new OriginalWebSocket(url, protocols);
+
+        ws.addEventListener('message', (event: MessageEvent) => {
+          const data = event.data;
+          if (
+            typeof data === 'string' &&
+            data.startsWith(SOCKET_IO_EVENT_PREFIX)
+          ) {
+            try {
+              const payload = JSON.parse(data.slice(SOCKET_IO_PREFIX_LENGTH));
+              if (payload[0] === eventType) {
+                callback(payload[1]);
+              }
+            } catch (e) {
+              // Ignore parse errors
+            }
+          }
+        });
+
+        return ws;
+      }
+    );
+  });
+}