konnosDashboard
role-impersonation-state.tsx162 lines · main
1import { ident, literal, safeSql } from '@supabase/pg-meta/src/pg-format'
2import { useConstant } from 'common'
3import { createContext, PropsWithChildren, useCallback, useContext, useEffect } from 'react'
4import { proxy, snapshot, subscribe, useSnapshot } from 'valtio'
5
6import { CustomAccessTokenHookDetails } from '../hooks/misc/useCustomAccessTokenHookDetails'
7import { executeSql } from '@/data/sql/execute-sql-query'
8import useLatest from '@/hooks/misc/useLatest'
9import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
10import { getPostgrestClaims, ImpersonationRole } from '@/lib/role-impersonation'
11
12export function createRoleImpersonationState(
13 projectRef: string,
14 customizeAccessTokenRef: {
15 current: (args: {
16 schema: string
17 functionName: string
18 claims: ReturnType<typeof getPostgrestClaims>
19 }) => Promise<any>
20 }
21) {
22 const roleImpersonationState = proxy({
23 projectRef,
24 role: undefined as ImpersonationRole | undefined,
25 claims: undefined as ReturnType<typeof getPostgrestClaims> | undefined,
26
27 setRole: async (
28 role: ImpersonationRole | undefined,
29 customAccessTokenHookDetails?: CustomAccessTokenHookDetails
30 ) => {
31 let claims = role?.type === 'postgrest' ? getPostgrestClaims(projectRef, role) : undefined
32
33 if (customAccessTokenHookDetails?.type === 'postgres' && claims !== undefined) {
34 const { schema, functionName } = customAccessTokenHookDetails
35 const updatedClaims = await customizeAccessTokenRef.current({
36 schema,
37 functionName,
38 claims,
39 })
40 if (updatedClaims) {
41 claims = updatedClaims
42 }
43 }
44
45 roleImpersonationState.role = role
46 if (claims) {
47 roleImpersonationState.claims = claims
48 }
49 },
50 })
51
52 return roleImpersonationState
53}
54
55export type RoleImpersonationState = ReturnType<typeof createRoleImpersonationState>
56
57export const RoleImpersonationStateContext = createContext<RoleImpersonationState>(
58 createRoleImpersonationState('', { current: async () => {} })
59)
60
61export const RoleImpersonationStateContextProvider = ({ children }: PropsWithChildren) => {
62 const { data: project } = useSelectedProjectQuery()
63 async function customizeAccessToken({
64 schema,
65 functionName,
66 claims,
67 }: {
68 schema: string
69 functionName: string
70 claims: ReturnType<typeof getPostgrestClaims>
71 }) {
72 const event = { user_id: claims.sub, claims, authentication_method: 'password' }
73
74 const result = await executeSql({
75 projectRef: project?.ref,
76 connectionString: project?.connectionString,
77 sql: safeSql`select ${ident(schema)}.${ident(functionName)}(${literal(JSON.stringify(event))}::jsonb) as event;`,
78 queryKey: ['customize-access-token', project?.ref],
79 })
80
81 return result?.result?.[0]?.event?.claims
82 }
83
84 const customizeAccessTokenRef = useLatest(customizeAccessToken)
85
86 const state = useConstant(() =>
87 createRoleImpersonationState(project?.ref ?? '', customizeAccessTokenRef)
88 )
89
90 return (
91 <RoleImpersonationStateContext.Provider value={state}>
92 {children}
93 </RoleImpersonationStateContext.Provider>
94 )
95}
96
97export function useRoleImpersonationStateSnapshot(options?: Parameters<typeof useSnapshot>[1]) {
98 const roleImpersonationState = useContext(RoleImpersonationStateContext)
99
100 return useSnapshot(roleImpersonationState, options)
101}
102
103export function useGetImpersonatedRoleState() {
104 const roleImpersonationState = useContext(RoleImpersonationStateContext)
105
106 return useCallback(
107 // [Alaister]: typeof roleImpersonationState is needed to avoid readonly type errors everywhere
108 () => snapshot(roleImpersonationState) as typeof roleImpersonationState,
109 [roleImpersonationState]
110 )
111}
112
113export function useSubscribeToImpersonatedRole(
114 onChange: (role: ImpersonationRole | undefined) => void
115) {
116 const roleImpersonationState = useContext(RoleImpersonationStateContext)
117 const onChangeRef = useLatest(onChange)
118
119 useEffect(() => {
120 return subscribe(roleImpersonationState, () => {
121 onChangeRef.current(snapshot(roleImpersonationState).role)
122 })
123 }, [roleImpersonationState])
124}
125
126export function isRoleImpersonationEnabled(impersonationRole?: ImpersonationRole) {
127 return impersonationRole?.type === 'postgrest'
128}
129
130export const useIsImpersonatingAnon = () => {
131 const state = useRoleImpersonationStateSnapshot()
132 return state.role?.type === 'postgrest' && state.role.role === 'anon'
133}
134
135export const useImpersonatedUser = () => {
136 const state = useRoleImpersonationStateSnapshot()
137 return state.role?.type === 'postgrest' &&
138 state.role.role === 'authenticated' &&
139 state.role.userType === 'native'
140 ? state.role.user
141 : undefined
142}
143
144export const useImpersonatedExternalAuth = () => {
145 const state = useRoleImpersonationStateSnapshot()
146 return state.role?.type === 'postgrest' &&
147 state.role.role === 'authenticated' &&
148 state.role.userType === 'external' &&
149 state.role.externalAuth
150 ? state.role.externalAuth.sub
151 : undefined
152}
153
154export const useImpersonatedAAL = () => {
155 const state = useRoleImpersonationStateSnapshot()
156 return (
157 state.role?.type === 'postgrest' &&
158 state.role.role === 'authenticated' &&
159 state.role.userType === 'external' &&
160 state.role.aal
161 )
162}