Implement CSRF support in the UI (#12354)

Co-authored-by: Nicholas Brown <znicholasbrown@gmail.com>
Co-authored-by: Craig Harshbarger <pleek91@gmail.com>

src/prefect/server/api/middleware.py

@@ -67,6 +67,7 @@ async def dispatch(
                     return JSONResponse(
                         {"detail": "Invalid CSRF token or client identifier."},
                         status_code=status.HTTP_403_FORBIDDEN,
+                        headers={"Access-Control-Allow-Origin": "*"},
                     )
 
         return await call_next(request)

ui/src/maps/csrfToken.ts

@@ -0,0 +1,12 @@
+import { CsrfToken } from "@/models/CsrfToken"
+import { CsrfTokenResponse } from "@/types/csrfTokenResponse"
+import { MapFunction } from '@/services/mapper'
+
+export const mapCsrfTokenResponseToCsrfToken: MapFunction<CsrfTokenResponse, CsrfToken> = function(source) {
+    return {
+        token: source.token,
+        expiration: this.map('string', source.expiration, 'Date'),
+        issued: new Date(),
+    }
+}
+

ui/src/maps/index.ts

@@ -1,9 +1,11 @@
 import { maps as designMaps } from '@prefecthq/prefect-ui-library'
 import { mapFlagResponseToFeatureFlag } from '@/maps/featureFlag'
 import { mapSettingsResponseToSettings } from '@/maps/uiSettings'
+import { mapCsrfTokenResponseToCsrfToken } from '@/maps/csrfToken'
 
 export const maps = {
   ...designMaps,
   FlagResponse: { FeatureFlag: mapFlagResponseToFeatureFlag },
   SettingsResponse: { Settings: mapSettingsResponseToSettings },
-}
+  CsrfTokenResponse: { CsrfToken: mapCsrfTokenResponseToCsrfToken },
+}

ui/src/models/CsrfToken.ts

@@ -0,0 +1,5 @@
+export type CsrfToken = {
+    token: string
+    expiration: Date
+    issued: Date
+}

ui/src/services/csrfTokenApi.ts

@@ -0,0 +1,134 @@
+import { AxiosError, AxiosInstance, InternalAxiosRequestConfig, isAxiosError } from 'axios'
+import { randomId } from '@prefecthq/prefect-design'
+import { Api, AxiosInstanceSetupHook, PrefectConfig } from '@prefecthq/prefect-ui-library'
+import { CreateActions } from '@prefecthq/vue-compositions'
+import { CsrfToken } from '@/models/CsrfToken'
+import { CsrfTokenResponse } from '@/types/csrfTokenResponse'
+import { mapper } from '@/services/mapper'
+
+export class CsrfTokenApi extends Api {
+    public csrfToken?: CsrfToken
+    public clientId: string = randomId()
+    public csrfSupportEnabled = true
+    private refreshTimeout: ReturnType<typeof setTimeout> | null = null
+    private ongoingRefresh: Promise<void> | null = null
+
+    public constructor(apiConfig: PrefectConfig, instanceSetupHook: AxiosInstanceSetupHook | null = null) {
+        super(apiConfig, instanceSetupHook)
+        this.startBackgroundTokenRefresh()
+    }
+
+    public async addCsrfHeaders(config: InternalAxiosRequestConfig) {
+        if (this.csrfSupportEnabled) {
+            const csrfToken = await this.getCsrfToken()
+            config.headers = config.headers || {}
+            config.headers['Prefect-Csrf-Token'] = csrfToken.token
+            config.headers['Prefect-Csrf-Client'] = this.clientId
+        }
+    }
+
+    private async getCsrfToken(): Promise<CsrfToken> {
+        if (this.shouldRefreshToken()) {
+            await this.refreshCsrfToken()
+        }
+
+        if (!this.csrfToken) {
+            throw new Error('CSRF token not available')
+        }
+
+        return this.csrfToken
+    }
+
+    private async refreshCsrfToken(force: boolean = false) {
+        if (!force && !this.shouldRefreshToken()) {
+            return this.ongoingRefresh ?? Promise.resolve()
+        }
+
+        if (this.ongoingRefresh) {
+            return this.ongoingRefresh
+        }
+
+        const refresh = async () => {
+            try {
+
+                const response = await this.get<CsrfTokenResponse>(`/csrf-token?client=${this.clientId}`)
+                this.csrfToken = mapper.map('CsrfTokenResponse', response.data, 'CsrfToken')
+                this.ongoingRefresh = null
+            } catch (error) {
+                this.ongoingRefresh = null
+                if (isAxiosError(error)) {
+                    if (this.isUnconfiguredServer(error)) {
+                        this.disableCsrfSupport()
+                    } else {
+                        console.error('Failed to refresh CSRF token:', error)
+                        throw new Error('Failed to refresh CSRF token')
+                    }
+                }
+            }
+        }
+
+        this.ongoingRefresh = refresh()
+        return this.ongoingRefresh
+    }
+
+    private shouldRefreshToken(): boolean {
+        if (!this.csrfSupportEnabled) {
+            return false
+        }
+
+        if (!this.csrfToken) {
+            return true;
+        }
+
+        if (!this.csrfToken.token || !this.csrfToken.expiration) {
+            return true
+        }
+
+        return new Date() > this.csrfToken.expiration
+    }
+
+    private isUnconfiguredServer(error: AxiosError<any, any>): boolean {
+        return error.response?.status === 422 && error.response?.data.detail.includes('CSRF protection is disabled.')
+    }
+
+    private disableCsrfSupport() {
+        this.csrfSupportEnabled = false
+    }
+
+    private startBackgroundTokenRefresh() {
+        const calculateTimeoutDuration = () => {
+            if (this.csrfToken) {
+                const now = new Date()
+                const expiration = new Date(this.csrfToken.expiration)
+                const issuedAt = this.csrfToken.issued
+                const lifetime = expiration.getTime() - issuedAt.getTime()
+                const refreshThreshold = issuedAt.getTime() + lifetime * 0.75
+                const durationUntilRefresh = refreshThreshold - now.getTime()
+
+                return durationUntilRefresh
+            }
+            return 0 // If we don't have token data cause an immediate refresh
+        }
+
+        const refreshTask = async () => {
+            if (this.csrfSupportEnabled) {
+                await this.refreshCsrfToken(true)
+                this.refreshTimeout = setTimeout(refreshTask, calculateTimeoutDuration())
+            }
+        }
+
+        this.refreshTimeout = setTimeout(refreshTask, calculateTimeoutDuration())
+    }
+}
+
+export function setupCsrfInterceptor(csrfTokenApi: CreateActions<CsrfTokenApi>, axiosInstance: AxiosInstance) {
+    axiosInstance.interceptors.request.use(async (config: InternalAxiosRequestConfig): Promise<InternalAxiosRequestConfig> => {
+        const method = config.method?.toLowerCase()
+
+        if (method && ['post', 'patch', 'put', 'delete'].includes(method)) {
+            await csrfTokenApi.addCsrfHeaders(config)
+        }
+
+        return config
+    })
+}

ui/src/types/csrfTokenResponse.ts

@@ -0,0 +1,5 @@
+export type CsrfTokenResponse = {
+    token: string
+    expiration: Date
+}
+

ui/src/utilities/api.ts

@@ -2,14 +2,24 @@ import { createApi, PrefectConfig } from '@prefecthq/prefect-ui-library'
 import { createActions } from '@prefecthq/vue-compositions'
 import { InjectionKey } from 'vue'
 import { AdminApi } from '@/services/adminApi'
+import { CsrfTokenApi, setupCsrfInterceptor } from '@/services/csrfTokenApi'
+import { AxiosInstance } from 'axios'
+
+
 
 // eslint-disable-next-line @typescript-eslint/explicit-function-return-type
 export function createPrefectApi(config: PrefectConfig) {
-  const workspaceApi = createApi(config)
+  const csrfTokenApi = createActions(new CsrfTokenApi(config))
+
+  function axiosInstanceSetupHook(axiosInstance: AxiosInstance) {
+    setupCsrfInterceptor(csrfTokenApi, axiosInstance)
+  };
 
+  const workspaceApi = createApi(config, axiosInstanceSetupHook)
   return {
     ...workspaceApi,
-    admin: createActions(new AdminApi(config)),
+    csrf: csrfTokenApi,
+    admin: createActions(new AdminApi(config, axiosInstanceSetupHook)),
   }
 }