Add password-protected shutdown/restart

packages/ui/src/providers/global-system-state/index.tsx

@@ -10,9 +10,9 @@ import {toast} from '@/components/ui/toast'
 import {usePrefixedLocalStorage} from '@/hooks/use-prefixed-local-storage'
 import {useJwt} from '@/modules/auth/use-auth'
 import {MigratingCover, useMigrate} from '@/providers/global-system-state/migrate'
-import {RestartingCover, useRestart} from '@/providers/global-system-state/restart'
-import {ShuttingDownCover, useShutdown} from '@/providers/global-system-state/shutdown'
-import {RouterError, RouterOutput, trpcReact} from '@/trpc/trpc'
+import {RestartingCover, useRestart, useRestartWithPassword} from '@/providers/global-system-state/restart'
+import {ShuttingDownCover, useShutdown, useShutdownWithPassword} from '@/providers/global-system-state/shutdown'
+import {RouterError, RouterInput, RouterOutput, trpcReact} from '@/trpc/trpc'
 import {MS_PER_SECOND} from '@/utils/date-time'
 import {assertUnreachable, IS_DEV} from '@/utils/misc'
 
@@ -24,7 +24,9 @@ type SystemStatus = RouterOutput['system']['status']
 
 const GlobalSystemStateContext = createContext<{
 	shutdown: () => void
+	shutdownWithPassword: (input: RouterInput['system']['shutdownWithPassword']) => Promise<boolean>
 	restart: () => void
+	restartWithPassword: (input: RouterInput['system']['restartWithPassword']) => Promise<boolean>
 	update: () => void
 	migrate: () => void
 	reset: (password: string) => void
@@ -73,6 +75,14 @@ export function GlobalSystemStateProvider({children}: {children: ReactNode}) {
 		// Prevent logout/redirect when error occurs
 		setShouldLogoutOnRunning(false)
 	}
+
+	const onPowerError = async () => {
+		setTriggered(false)
+		setShouldLogoutOnRunning(false)
+		setStartShutdownTimer(false)
+		setShutdownComplete(false)
+		setRouterError(null)
+	}
 	const getError = () => routerError
 	const clearError = () => setRouterError(null)
 	// Allow external code to suppress errors (e.g., RAID setup doing its own restart flow)
@@ -97,7 +107,9 @@ export function GlobalSystemStateProvider({children}: {children: ReactNode}) {
 
 	// TODO: handle `onError` for other actions than reset?
 	const restart = useRestart({onMutate, onSuccess})
+	const restartWithPassword = useRestartWithPassword({onMutate, onSuccess, onError: onPowerError})
 	const shutdown = useShutdown({onMutate, onSuccess})
+	const shutdownWithPassword = useShutdownWithPassword({onMutate, onSuccess, onError: onPowerError})
 	const update = useUpdate({onMutate, onSuccess})
 	const migrate = useMigrate({onMutate, onSuccess})
 	const reset = useReset({onMutate, onError})
@@ -248,7 +260,18 @@ export function GlobalSystemStateProvider({children}: {children: ReactNode}) {
 		case 'running': {
 			return (
 				<GlobalSystemStateContext
-					value={{shutdown, restart, update, migrate, reset, getError, clearError, suppressErrors}}
+					value={{
+						shutdown,
+						shutdownWithPassword,
+						restart,
+						restartWithPassword,
+						update,
+						migrate,
+						reset,
+						getError,
+						clearError,
+						suppressErrors,
+					}}
 				>
 					{children}
 					{debugInfo}

packages/ui/src/providers/global-system-state/restart.tsx

@@ -2,7 +2,7 @@ import {useTranslation} from 'react-i18next'
 
 import {CoverMessage, CoverMessageParagraph} from '@/components/ui/cover-message'
 import {Loading} from '@/components/ui/loading'
-import {trpcReact} from '@/trpc/trpc'
+import {type RouterError, trpcReact} from '@/trpc/trpc'
 
 export function useRestart({onMutate, onSuccess}: {onMutate?: () => void; onSuccess?: (didWork: boolean) => void}) {
 	const restartMut = trpcReact.system.restart.useMutation({
@@ -14,6 +14,23 @@ export function useRestart({onMutate, onSuccess}: {onMutate?: () => void; onSucc
 	return restart
 }
 
+export function useRestartWithPassword({
+	onMutate,
+	onSuccess,
+	onError,
+}: {
+	onMutate?: () => void
+	onSuccess?: (didWork: boolean) => void
+	onError?: (error: RouterError) => void
+}) {
+	const restartMut = trpcReact.system.restartWithPassword.useMutation({
+		onMutate,
+		onSuccess,
+		onError,
+	})
+	return restartMut.mutateAsync
+}
+
 export function RestartingCover() {
 	const {t} = useTranslation()
 	return (

packages/ui/src/providers/global-system-state/shutdown.tsx

@@ -2,7 +2,7 @@ import {useTranslation} from 'react-i18next'
 
 import {CoverMessage, CoverMessageParagraph} from '@/components/ui/cover-message'
 import {Loading} from '@/components/ui/loading'
-import {trpcReact} from '@/trpc/trpc'
+import {type RouterError, trpcReact} from '@/trpc/trpc'
 
 export function useShutdown({onMutate, onSuccess}: {onMutate?: () => void; onSuccess?: (didWork: boolean) => void}) {
 	const shutdownMut = trpcReact.system.shutdown.useMutation({
@@ -14,6 +14,23 @@ export function useShutdown({onMutate, onSuccess}: {onMutate?: () => void; onSuc
 	return shutdown
 }
 
+export function useShutdownWithPassword({
+	onMutate,
+	onSuccess,
+	onError,
+}: {
+	onMutate?: () => void
+	onSuccess?: (didWork: boolean) => void
+	onError?: (error: RouterError) => void
+}) {
+	const shutdownMut = trpcReact.system.shutdownWithPassword.useMutation({
+		onMutate,
+		onSuccess,
+		onError,
+	})
+	return shutdownMut.mutateAsync
+}
+
 export function ShuttingDownCover() {
 	const {t} = useTranslation()
 	return (

packages/ui/src/routes/login.tsx

@@ -1,25 +1,31 @@
-import {useState} from 'react'
+import {useEffect, useState} from 'react'
 import {useTranslation} from 'react-i18next'
 import {TbCircleCheckFilled} from 'react-icons/tb'
+import {RiRestartLine, RiShutDownLine} from 'react-icons/ri'
 import {useLocation} from 'react-router-dom'
 
 import {
 	AlertDialog,
 	AlertDialogAction,
+	AlertDialogCancel,
 	AlertDialogContent,
 	AlertDialogDescription,
 	AlertDialogFooter,
 	AlertDialogHeader,
 	AlertDialogTitle,
+	AlertDialogTrigger,
 } from '@/components/ui/alert-dialog'
 import {PasswordInput} from '@/components/ui/input'
 import {PinInput} from '@/components/ui/pin-input'
-import {formGroupClass, Layout, primaryButtonProps} from '@/layouts/bare/shared'
+import {formGroupClass, Layout, primaryButtonProps, secondaryButtonClasss} from '@/layouts/bare/shared'
 import {cn} from '@/lib/utils'
 import {useAuth} from '@/modules/auth/use-auth'
+import {useGlobalSystemState} from '@/providers/global-system-state/index'
 import {trpcReact} from '@/trpc/trpc'
 
 type Step = 'password' | '2fa'
+type PowerStep = 'password' | '2fa'
+type PowerAction = 'shutdown' | 'restart'
 
 export default function Login() {
 	const {t} = useTranslation()
@@ -68,11 +74,18 @@ export default function Login() {
 		</AlertDialog>
 	)
 
+	const powerFooter = (
+		<>
+			<PowerActionDialog action='restart' />
+			<PowerActionDialog action='shutdown' />
+		</>
+	)
+
 	switch (step) {
 		case 'password': {
 			return (
 				<>
-					<Layout title={t('login.title')} subTitle={t('login.subtitle')}>
+					<Layout title={t('login.title')} subTitle={t('login.subtitle')} footer={powerFooter}>
 						<form className='flex w-full flex-col items-center gap-5 px-4 md:px-0' onSubmit={handleSubmitPassword}>
 							<div className={cn(formGroupClass, 'max-w-[280px]')}>
 								<PasswordInput
@@ -95,7 +108,7 @@ export default function Login() {
 		case '2fa': {
 			return (
 				<>
-					<Layout title={t('login-2fa.title')} subTitle={t('login-2fa.subtitle')}>
+					<Layout title={t('login-2fa.title')} subTitle={t('login-2fa.subtitle')} footer={powerFooter}>
 						<form className='flex w-full flex-col items-center gap-5 px-4 md:px-0' onSubmit={handleSubmitPassword}>
 							<PinInput autoFocus length={6} onCodeCheck={handleSubmit2fa} />
 						</form>
@@ -106,3 +119,114 @@ export default function Login() {
 		}
 	}
 }
+
+function PowerActionDialog({action}: {action: PowerAction}) {
+	const {t} = useTranslation()
+	const {shutdownWithPassword, restartWithPassword} = useGlobalSystemState()
+	const powerAction = action === 'shutdown' ? shutdownWithPassword : restartWithPassword
+	const [open, setOpen] = useState(false)
+	const [step, setStep] = useState<PowerStep>('password')
+	const [password, setPassword] = useState('')
+	const [passwordError, setPasswordError] = useState('')
+	const [isPending, setIsPending] = useState(false)
+
+	useEffect(() => {
+		if (!open) {
+			setStep('password')
+			setPassword('')
+			setPasswordError('')
+			setIsPending(false)
+		}
+	}, [open])
+
+	const titleKey = action === 'shutdown' ? 'shut-down.confirm.title' : 'restart.confirm.title'
+	const submitKey = action === 'shutdown' ? 'shut-down.confirm.submit' : 'restart.confirm.submit'
+	const triggerKey = action === 'shutdown' ? 'shut-down' : 'restart'
+	const ActionIcon = action === 'shutdown' ? RiShutDownLine : RiRestartLine
+
+	const handlePasswordSubmit = async (event: React.FormEvent<HTMLFormElement>) => {
+		event.preventDefault()
+		if (!password) return
+		setPasswordError('')
+		setIsPending(true)
+		try {
+			await powerAction({password})
+			setOpen(false)
+		} catch (error) {
+			const message = (error as {message?: string})?.message ?? ''
+			if (message === 'Missing 2FA code') {
+				setPasswordError('')
+				setStep('2fa')
+				return
+			}
+			setPasswordError(message || t('something-went-wrong'))
+		} finally {
+			setIsPending(false)
+		}
+	}
+
+	const handleSubmit2fa = async (totpToken: string) => {
+		try {
+			await powerAction({password, totpToken})
+			setOpen(false)
+			return true
+		} catch (error) {
+			const message = (error as {message?: string})?.message ?? ''
+			if (message === 'Incorrect password') {
+				setPasswordError(message)
+				setStep('password')
+			}
+			return false
+		}
+	}
+
+	return (
+		<AlertDialog open={open} onOpenChange={setOpen}>
+			<AlertDialogTrigger asChild>
+				<button className={secondaryButtonClasss} type='button'>
+					{t(triggerKey)}
+				</button>
+			</AlertDialogTrigger>
+			<AlertDialogContent>
+				{step === 'password' ? (
+					<form className='flex flex-col gap-5' onSubmit={handlePasswordSubmit}>
+						<AlertDialogHeader icon={ActionIcon}>
+							<AlertDialogTitle>{t(titleKey)}</AlertDialogTitle>
+						</AlertDialogHeader>
+						<div className={cn(formGroupClass, 'mx-auto w-full max-w-[280px]')}>
+							<PasswordInput
+								autoFocus
+								label={t('login.password-label')}
+								value={password}
+								onValueChange={(value) => {
+									setPasswordError('')
+									setPassword(value)
+								}}
+								error={passwordError}
+							/>
+						</div>
+						<AlertDialogFooter>
+							<AlertDialogAction variant='destructive' type='submit' disabled={!password || isPending}>
+								{t(submitKey)}
+							</AlertDialogAction>
+							<AlertDialogCancel>{t('cancel')}</AlertDialogCancel>
+						</AlertDialogFooter>
+					</form>
+				) : (
+					<div className='flex flex-col gap-5'>
+						<AlertDialogHeader icon={ActionIcon}>
+							<AlertDialogTitle>{t(titleKey)}</AlertDialogTitle>
+							<AlertDialogDescription>{t('login-2fa.subtitle')}</AlertDialogDescription>
+						</AlertDialogHeader>
+						<div className='mx-auto'>
+							<PinInput autoFocus length={6} onCodeCheck={handleSubmit2fa} />
+						</div>
+						<AlertDialogFooter>
+							<AlertDialogCancel>{t('cancel')}</AlertDialogCancel>
+						</AlertDialogFooter>
+					</div>
+				)}
+			</AlertDialogContent>
+		</AlertDialog>
+	)
+}

packages/umbreld/source/modules/system/routes.ts

@@ -30,10 +30,32 @@ import {
 } from './system.js'
 
 import {privateProcedure, publicProcedure, publicProcedureWhenNoUserExists, router} from '../server/trpc/trpc.js'
+import type {Context} from '../server/trpc/context.js'
 
 type SystemStatus = 'running' | 'updating' | 'shutting-down' | 'restarting' | 'migrating' | 'resetting' | 'restoring'
 let systemStatus: SystemStatus = 'running'
 
+const powerActionInput = z.object({
+	password: z.string(),
+	totpToken: z.string().optional(),
+})
+
+async function validatePowerActionCredentials(ctx: Context, input: z.infer<typeof powerActionInput>) {
+	const userExists = await ctx.user.exists()
+	if (!userExists) return
+	if (!(await ctx.user.validatePassword(input.password))) {
+		throw new TRPCError({code: 'UNAUTHORIZED', message: 'Incorrect password'})
+	}
+	if (await ctx.user.is2faEnabled()) {
+		if (!input.totpToken) {
+			throw new TRPCError({code: 'UNAUTHORIZED', message: 'Missing 2FA code'})
+		}
+		if (!(await ctx.user.validate2faToken(input.totpToken))) {
+			throw new TRPCError({code: 'UNAUTHORIZED', message: 'Incorrect 2FA code'})
+		}
+	}
+}
+
 // Quick hack so we can set system status from migration module until we refactor this
 export function setSystemStatus(status: SystemStatus) {
 	systemStatus = status
@@ -155,6 +177,28 @@ export default router({
 			}),
 		)
 		.mutation(async ({ctx, input}) => clearStaticIp(ctx.umbreld, input)),
+	// Public on login screen, but requires password (and 2FA when enabled)
+	shutdownWithPassword: publicProcedure
+		.input(powerActionInput)
+		.mutation(async ({ctx, input}) => {
+			await validatePowerActionCredentials(ctx, input)
+			systemStatus = 'shutting-down'
+			await ctx.umbreld.stop()
+			await shutdown()
+
+			return true
+		}),
+	// Public on login screen, but requires password (and 2FA when enabled)
+	restartWithPassword: publicProcedure
+		.input(powerActionInput)
+		.mutation(async ({ctx, input}) => {
+			await validatePowerActionCredentials(ctx, input)
+			systemStatus = 'restarting'
+			await ctx.umbreld.stop()
+			await reboot()
+
+			return true
+		}),
 	// Public during onboarding and recovery mode so users can shut down during RAID setup or mount failure
 	shutdown: publicProcedureWhenNoUserExists.mutation(async ({ctx}) => {
 		systemStatus = 'shutting-down'