@wener/console/src/foundation/auth/AuthStore.ts

Base console UI toolkit

import { use, useEffect, useRef, type FC } from 'react';
import { createReactContext } from '@wener/reaction';
import { useNetworkStatus } from '@wener/reaction/store';
import { createBoundedUseStore } from '@wener/reaction/zustand';
import { getGlobalStates } from '@wener/utils';
import { createStore } from 'zustand';
import { mutative } from 'zustand-mutative';

export const AuthStatus = {
	Init: 'Init',
	Authenticated: 'Authenticated',
	Unauthenticated: 'Unauthenticated',
	Expired: 'Expired',
	Loading: 'Loading',
	Error: 'Error',
	Locked: 'Locked',
} as const;
type AuthStatusCode = (typeof AuthStatus)[keyof typeof AuthStatus];

interface SetAuthOptions {
	accessToken: string;
	refreshToken?: string;
	expiresIn?: number;
	expiresAt?: Date | string;
}

interface AuthStoreState {
	status: AuthStatusCode;
	accessToken?: string;
	refreshToken?: string;
	expiresIn?: number;
	expiresAt?: Date;
	error?: any;

	setAuth(o: SetAuthOptions): void;

	reset(): void;
}

export type AuthStore = ReturnType<typeof createAuthStore>;

export function createAuthStore(init: Partial<AuthStoreState> = {}) {
	return createStore(
		mutative<AuthStoreState>((setState, getState, store) => {
			return {
				...init,
				status: AuthStatus.Init,
				setAuth(o: SetAuthOptions) {
					setState((s) => {
						Object.assign(s, {
							status: AuthStatus.Authenticated,
							...o,
							expiresAt: o.expiresAt ? new Date(o.expiresAt) : undefined,
						});
					});
				},
				reset() {
					setState((s) => {
						Object.assign(s, {
							status: AuthStatus.Unauthenticated,
							accessToken: undefined,
							refreshToken: undefined,
							expiresIn: undefined,
							expiresAt: undefined,
							error: undefined,
						});
					});
				},
			};
		}),
	);
}

type AuthSidecarProps = {
	store: AuthStore;
	actions: {
		refresh: (o: { accessToken: string; refreshToken?: string }) => Promise<{
			accessToken: string;
			refreshToken?: string;
			expiresAt: Date | string;
		}>;
	};
	storage?: Storage;
};

function useAuthSidecar({ store, actions: { refresh }, storage = localStorage }: AuthSidecarProps) {
	const { online } = useNetworkStatus();
	// todo watch storage ?
	const checkAuth = async () => {
		const accessToken = storage.getItem('accessToken');
		const refreshToken = storage.getItem('refreshToken') ?? undefined;
		if (accessToken) {
			try {
				const out = await refresh({ accessToken, refreshToken });
				store.getState().setAuth(out);
				return true;
			} catch (e) {
				console.error('Failed to refresh token', e);
			}
		}
		// not authenticated or server error
		store.getState().reset();
		return false;
	};

	const authRef = useRef<Promise<boolean> | undefined>(undefined);

	const doAuthCheck = () => {
		const state = store.getState();
		// only check for init and authenticated
		switch (state.status) {
			case AuthStatus.Init:
			case AuthStatus.Authenticated:
				break;
			default:
				return;
		}
		// avoid race
		let current = authRef.current;
		if (current) {
			return current.then((v) => {
				// authed, skip for now, will check for next
				if (v) {
					authRef.current = undefined;
				} else {
					// check again
					return (current = checkAuth());
				}
				return v;
			});
		} else {
			return (authRef.current = checkAuth().then((v) => {
				// done
				authRef.current = undefined;
				return v;
			}));
		}
	};

	// auth check
	useEffect(() => {
		if (!online) return;
		doAuthCheck();
		const timer = setInterval(doAuthCheck, 5 * 60 * 1000);
		return () => {
			clearInterval(timer);
		};
	}, [store, online]);

	//
	useAuthTokenPersist(store, storage);
}

const AuthStoreStateKey = 'AuthStore';

export const AuthSidecar: FC<Omit<AuthSidecarProps, 'store'>> = (props) => {
	const store = useAuthStoreContext();
	useAuthSidecar({
		store,
		...props,
	});
	return null;
};

function useAuthTokenPersist(store: AuthStore, storage: Storage) {
	const ref = useRef(storage);
	ref.current = storage;
	// persist token
	useEffect(() => {
		return store.subscribe((s) => {
			if (s.status !== AuthStatus.Authenticated) {
				return;
			}
			const { accessToken = '', refreshToken } = s;

			const storage = ref.current;
			if ((storage.getItem('accessToken') ?? '') === accessToken) {
				return;
			}

			accessToken ? storage.setItem('accessToken', accessToken) : deleteItem(storage, 'accessToken');
			refreshToken ? storage.setItem('refreshToken', refreshToken) : deleteItem(storage, 'refreshToken');
		});
	}, [store]);
}

export function getAuthStore(): AuthStore {
	return getGlobalStates(AuthStoreStateKey, () => {
		return createAuthStore();
	});
}

export function getAuthState() {
	return getAuthStore().getState();
}

export function getAccessToken() {
	return getAuthState().accessToken;
}

export const AuthStoreContext = createReactContext<undefined | AuthStore>('AuthStore', undefined);

export function useAuthStoreContext() {
	return use(AuthStoreContext) || getAuthStore();
}

function deleteItem(s: any, key: string) {
	if ('removeItem' in s) {
		s.removeItem(key);
	} else {
		delete s[key];
	}
}

export const useAuthStore = createBoundedUseStore(useAuthStoreContext);