react-mention-input/src/useProtectedImage.ts

A React component for input with @mention functionality.

import { useState, useEffect, useRef } from 'react';

interface UseProtectedImageOptions {
  url: string | null;
  isProtected?: boolean | ((url: string) => boolean);
  getAuthHeaders?: () => Record<string, string> | Promise<Record<string, string>>;
}

// Module-level cache to persist blob URLs across component re-renders
const blobUrlCache = new Map<string, string>();
const fetchingUrls = new Set<string>();

/**
 * Custom hook to handle protected image URLs that require authentication tokens in headers.
 * For protected URLs, it fetches the image with auth headers and converts it to a blob URL.
 * For non-protected URLs, it returns the URL as-is.
 * Uses a module-level cache to prevent blinking on re-renders.
 */
export const useProtectedImage = ({
  url,
  isProtected,
  getAuthHeaders,
}: UseProtectedImageOptions): string | null => {
  const [blobUrl, setBlobUrl] = useState<string | null>(() => {
    // Initialize from cache if available
    return url ? blobUrlCache.get(url) || null : null;
  });
  const previousUrlRef = useRef<string | null>(null);
  const isProtectedRef = useRef(isProtected);
  const getAuthHeadersRef = useRef(getAuthHeaders);
  const mountedRef = useRef(true);

  // Update refs when props change (but don't trigger re-fetch)
  useEffect(() => {
    isProtectedRef.current = isProtected;
    getAuthHeadersRef.current = getAuthHeaders;
  }, [isProtected, getAuthHeaders]);

  useEffect(() => {
    mountedRef.current = true;
    return () => {
      mountedRef.current = false;
    };
  }, []);

  useEffect(() => {
    // Always check cache first and restore if needed (synchronous)
    if (url && blobUrlCache.has(url) && !blobUrl) {
      const cachedBlobUrl = blobUrlCache.get(url)!;
      setBlobUrl(cachedBlobUrl);
    }

    // If URL hasn't changed, keep using cached blob URL
    if (url === previousUrlRef.current) {
      return;
    }

    const oldUrl = previousUrlRef.current;
    previousUrlRef.current = url;

    if (!url) {
      setBlobUrl(null);
      return;
    }

    // Check if URL is protected
    const shouldUseAuth = typeof isProtectedRef.current === 'boolean' 
      ? isProtectedRef.current 
      : isProtectedRef.current ? isProtectedRef.current(url) : false;

    // If not protected, use original URL
    if (!shouldUseAuth || !getAuthHeadersRef.current) {
      // Clear any cached blob URL for this URL
      if (blobUrlCache.has(url)) {
        const cachedBlobUrl = blobUrlCache.get(url)!;
        URL.revokeObjectURL(cachedBlobUrl);
        blobUrlCache.delete(url);
      }
      setBlobUrl(null);
      return;
    }

    // Check cache first - if we have a cached blob URL, use it immediately
    if (blobUrlCache.has(url)) {
      const cachedBlobUrl = blobUrlCache.get(url)!;
      setBlobUrl(cachedBlobUrl);
      return;
    }

    // If already fetching this URL, don't fetch again
    if (fetchingUrls.has(url)) {
      return;
    }

    // Fetch protected image with auth headers
    fetchingUrls.add(url);
    const fetchProtectedImage = async () => {
      try {
        const headers = await Promise.resolve(getAuthHeadersRef.current!());
        const response = await fetch(url, {
          method: 'GET',
          headers: {
            ...headers,
          },
        });

        if (!response.ok) {
          throw new Error(`Failed to fetch protected image: ${response.statusText}`);
        }

        const blob = await response.blob();
        const newBlobUrl = URL.createObjectURL(blob);
        
        // Only update if URL hasn't changed during fetch and component is still mounted
        if (previousUrlRef.current === url && mountedRef.current) {
          // Cache the blob URL
          blobUrlCache.set(url, newBlobUrl);
          setBlobUrl(newBlobUrl);
        } else {
          // URL changed during fetch or component unmounted, revoke this blob URL
          URL.revokeObjectURL(newBlobUrl);
        }
      } catch (error) {
        console.error('Error fetching protected image:', error);
        // On error, only clear if URL hasn't changed and component is mounted
        if (previousUrlRef.current === url && mountedRef.current) {
          setBlobUrl(null);
        }
      } finally {
        fetchingUrls.delete(url);
      }
    };

    fetchProtectedImage();
  }, [url]); // Only depend on url

  // Cleanup: Don't revoke blob URLs on unmount - keep them in cache for reuse
  // They will be cleaned up when the URL changes or when the cache is cleared

  // Return blob URL if available, otherwise return original URL (or null)
  if (!url) {
    return null;
  }
  
  const isUrlProtected = typeof isProtectedRef.current === 'boolean' 
    ? isProtectedRef.current 
    : isProtectedRef.current ? isProtectedRef.current(url) : false;
  
  // For protected URLs, return cached blob URL or current blobUrl state
  // For non-protected URLs, return the original URL
  if (isUrlProtected) {
    return blobUrl || blobUrlCache.get(url) || null;
  }
  
  return url;
};