From d93f4494d0a4e289dd04e436ba3f37ebd7049e43 Mon Sep 17 00:00:00 2001
From: Eileen Li <eileen@dragonscale.ai>
Date: Wed, 4 Jun 2025 16:38:52 -0700
Subject: [PATCH] fix: support authentication in image

---
 src/components/image/image.stories.tsx |  6 ++++
 src/components/image/image.tsx         | 48 ++++++++++++++++++++------
 src/components/types.ts                |  2 ++
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/src/components/image/image.stories.tsx b/src/components/image/image.stories.tsx
index 9ee27689..61f7bcf0 100644
--- a/src/components/image/image.stories.tsx
+++ b/src/components/image/image.stories.tsx
@@ -17,6 +17,12 @@ export const Default = {
   args: {
     alt: 'A curved facade covered in white latticework',
     src: 'images/image-component-example.png',
+    getAuthHeaders: () =>
+      Promise.resolve({
+        headers: {
+          Authorization: 'Bearer example-token',
+        },
+      }),
   },
 }
 
diff --git a/src/components/image/image.tsx b/src/components/image/image.tsx
index f22be5c2..9249ecd2 100644
--- a/src/components/image/image.tsx
+++ b/src/components/image/image.tsx
@@ -3,7 +3,7 @@ import '../../index.css'
 
 import CircularProgress from '@mui/material/CircularProgress'
 import Typography from '@mui/material/Typography'
-import { useState } from 'react'
+import { useEffect, useState } from 'react'
 import React from 'react'
 
 import { getSizeStyles } from '../helper'
@@ -17,6 +17,32 @@ export default function Image({
 }: ImageFormat) {
   const [isLoading, setIsLoading] = useState(true)
   const [errorMessage, setErrorMessage] = useState<string>('')
+  const [objectUrl, setObjectUrl] = useState<string>('')
+
+  useEffect(() => {
+    if (props.getAuthHeaders) {
+      setIsLoading(true)
+
+      props
+        .getAuthHeaders()
+        .then((authResult) => {
+          return fetch(props.src, {
+            headers: authResult.headers,
+          })
+        })
+        .then((response) => {
+          return response.blob()
+        })
+        .then((blob) => {
+          const newObjectUrl = URL.createObjectURL(blob)
+          setObjectUrl(newObjectUrl)
+        })
+        .catch(() => {
+          setErrorMessage('Image failed to load')
+          setIsLoading(false)
+        })
+    }
+  }, [props.src, props.getAuthHeaders])
 
   function handleImageError(): void {
     setErrorMessage('Image failed to load')
@@ -31,15 +57,17 @@ export default function Image({
     <figure>
       {isLoading && <CircularProgress data-cy="spinner" />}
       {props.title && <Typography variant="h6">{props.title}</Typography>}
-      <img
-        {...getSizeStyles(props.width, props.height)}
-        src={props.src}
-        alt={alt}
-        onError={handleImageError}
-        onLoad={() => {
-          setIsLoading(false)
-        }}
-      />
+      {(!props.getAuthHeaders || objectUrl) && (
+        <img
+          {...getSizeStyles(props.width, props.height)}
+          src={props.getAuthHeaders ? objectUrl : props.src}
+          alt={alt}
+          onError={handleImageError}
+          onLoad={() => {
+            setIsLoading(false)
+          }}
+        />
+      )}
       {props.description && (
         <figcaption>
           <MarkedMarkdown text={props.description} />
diff --git a/src/components/types.ts b/src/components/types.ts
index d48411f6..78e4243c 100644
--- a/src/components/types.ts
+++ b/src/components/types.ts
@@ -123,6 +123,8 @@ export interface ImageFormat extends VisualizationFormat {
   width?: number
   /** Height rendered in pixels. */
   height?: number
+  /** A function that can be used to get the headers for the url requests. */
+  getAuthHeaders?: GetAuthHeaders
 }
 
 export interface TableHeader {