From 30b5fbcd8edb6632ccba2bbaf8c38da7f3f2cd45 Mon Sep 17 00:00:00 2001
From: "Brian R. Bondy" <netzen@gmail.com>
Date: Wed, 9 Apr 2025 14:01:47 -0400
Subject: [PATCH] Allow BYOM links and disallow HTTP links (#28560)

---
 .../components/assistant_response/index.tsx   | 10 ++++---
 .../components/conversation_entries/index.tsx |  1 +
 .../components/markdown_renderer/index.tsx    | 10 ++++---
 .../markdown_renderer/render_link.test.tsx    | 27 +++++++++++++++++++
 4 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/components/ai_chat/resources/untrusted_conversation_frame/components/assistant_response/index.tsx b/components/ai_chat/resources/untrusted_conversation_frame/components/assistant_response/index.tsx
index e08f68b56df..72e0c5bb422 100644
--- a/components/ai_chat/resources/untrusted_conversation_frame/components/assistant_response/index.tsx
+++ b/components/ai_chat/resources/untrusted_conversation_frame/components/assistant_response/index.tsx
@@ -52,9 +52,10 @@ function AssistantEvent(props: {
   event: Mojom.ConversationEntryEvent,
   hasCompletionStarted: boolean,
   isEntryInProgress: boolean,
-  allowedLinks: string[]
+  allowedLinks: string[],
+  isLeoModel: boolean
 }) {
-  const { allowedLinks, event, isEntryInProgress } = props;
+  const { allowedLinks, event, isEntryInProgress, isLeoModel } = props;
 
   if (event.completionEvent) {
     const numberedLinks =
@@ -75,6 +76,7 @@ function AssistantEvent(props: {
         shouldShowTextCursor={isEntryInProgress}
         text={fullText}
         allowedLinks={allowedLinks}
+        disableLinkRestrictions={!isLeoModel}
       />
     )
   }
@@ -103,7 +105,8 @@ function AssistantEvent(props: {
 export default function AssistantResponse(props: {
   entry: Mojom.ConversationTurn,
   isEntryInProgress: boolean,
-  allowedLinks: string[]
+  allowedLinks: string[],
+  isLeoModel: boolean
 }) {
   // Extract certain events which need to render at specific locations (e.g. end of the events)
   const searchQueriesEvent = props.entry.events?.find(event => event.searchQueriesEvent)?.searchQueriesEvent
@@ -121,6 +124,7 @@ export default function AssistantResponse(props: {
         hasCompletionStarted={hasCompletionStarted}
         isEntryInProgress={props.isEntryInProgress}
         allowedLinks={props.allowedLinks}
+        isLeoModel={props.isLeoModel}
       />
     )
   }
diff --git a/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/index.tsx b/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/index.tsx
index 4d7fd0e75e5..634084680fb 100644
--- a/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/index.tsx
+++ b/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/index.tsx
@@ -148,6 +148,7 @@ function ConversationEntries() {
                       entry={latestTurn}
                       isEntryInProgress={isEntryInProgress}
                       allowedLinks={allAllowedLinks}
+                      isLeoModel={conversationContext.isLeoModel}
                     />
                   )}
                   {isHuman && !turn.selectedText && !showEditInput && (
diff --git a/components/ai_chat/resources/untrusted_conversation_frame/components/markdown_renderer/index.tsx b/components/ai_chat/resources/untrusted_conversation_frame/components/markdown_renderer/index.tsx
index 2e231c3c8b0..81caad4771c 100644
--- a/components/ai_chat/resources/untrusted_conversation_frame/components/markdown_renderer/index.tsx
+++ b/components/ai_chat/resources/untrusted_conversation_frame/components/markdown_renderer/index.tsx
@@ -85,18 +85,20 @@ function CursorDecorator(props: CursorDecoratorProps) {
 interface RenderLinkProps {
   a: React.ComponentProps<'a'>
   allowedLinks?: string[]
+  disableLinkRestrictions?: boolean
 }
 
 export function RenderLink(props: RenderLinkProps) {
-  const { a, allowedLinks } = props
+  const { a, allowedLinks, disableLinkRestrictions} = props
   const { href, children } = a
 
   // Context
   const context = useUntrustedConversationContext()
 
   // Computed
-  const isLinkAllowed =
-    allowedLinks?.some((link) => href?.startsWith(link)) ?? false
+  const isHttps = href?.toLowerCase().startsWith('https://')
+  const isLinkAllowed = isHttps && (disableLinkRestrictions ||
+    (allowedLinks?.some((link) => href?.startsWith(link)) ?? false))
 
   const handleLinkClicked = React.useCallback(() => {
     if (href && isLinkAllowed) {
@@ -132,6 +134,7 @@ interface MarkdownRendererProps {
   text: string
   shouldShowTextCursor: boolean
   allowedLinks?: string[]
+  disableLinkRestrictions?: boolean
 }
 
 export default function MarkdownRenderer(mainProps: MarkdownRendererProps) {
@@ -203,6 +206,7 @@ export default function MarkdownRenderer(mainProps: MarkdownRendererProps) {
             <RenderLink
               a={props}
               allowedLinks={mainProps.allowedLinks}
+              disableLinkRestrictions={mainProps.disableLinkRestrictions}
             />
           )
         }}
diff --git a/components/ai_chat/resources/untrusted_conversation_frame/components/markdown_renderer/render_link.test.tsx b/components/ai_chat/resources/untrusted_conversation_frame/components/markdown_renderer/render_link.test.tsx
index e727183353b..0bdd2189e97 100644
--- a/components/ai_chat/resources/untrusted_conversation_frame/components/markdown_renderer/render_link.test.tsx
+++ b/components/ai_chat/resources/untrusted_conversation_frame/components/markdown_renderer/render_link.test.tsx
@@ -43,3 +43,30 @@ test('Test RenderLink component with citations.', async () => {
   expect(screen.getByText('1').tagName).toBe('A')
   expect(screen.getByText('1').className).toBe('conversationLink citation')
 })
+
+test('Test RenderLink component with disableLinkRestrictions.', async () => {
+  render(
+    <RenderLink
+      a={{ href: 'https://example.com', children: 'Test Link' }}
+      allowedLinks={[]}
+      disableLinkRestrictions={true}
+    />
+  )
+  expect(screen.getByText('Test Link')).toBeInTheDocument()
+  expect(screen.getByText('Test Link').tagName).toBe('A')
+  expect(screen.getByText('Test Link').className).toBe('conversationLink')
+})
+
+// HTTP links should never be allowed
+test('Test RenderLink component with http links.', async () => {
+  render(
+    <RenderLink
+      a={{ href: 'http://example.com', children: 'Test Link' }}
+      allowedLinks={['http://example.com']}
+      disableLinkRestrictions={true}
+    />
+  )
+  expect(screen.getByText('Test Link')).toBeInTheDocument()
+  expect(screen.getByText('Test Link').tagName).toBe('SPAN')
+  expect(screen.getByText('Test Link').className).toBe('')
+})