From e325de9908edf51a3a8f3ad6079998ac8bcae2cc Mon Sep 17 00:00:00 2001 From: Eduardo Miranda Date: Mon, 4 May 2026 18:02:15 +0100 Subject: [PATCH] Add password-protected shutdown/restart --- .../providers/global-system-state/index.tsx | 31 +++- .../providers/global-system-state/restart.tsx | 19 ++- .../global-system-state/shutdown.tsx | 19 ++- packages/ui/src/routes/login.tsx | 132 +++++++++++++++++- .../umbreld/source/modules/system/routes.ts | 44 ++++++ 5 files changed, 235 insertions(+), 10 deletions(-) diff --git a/packages/ui/src/providers/global-system-state/index.tsx b/packages/ui/src/providers/global-system-state/index.tsx index a7b8ec511..22b27c96f 100644 --- a/packages/ui/src/providers/global-system-state/index.tsx +++ b/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 restart: () => void + restartWithPassword: (input: RouterInput['system']['restartWithPassword']) => Promise 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 ( {children} {debugInfo} diff --git a/packages/ui/src/providers/global-system-state/restart.tsx b/packages/ui/src/providers/global-system-state/restart.tsx index e95e5c805..4d73b8ed2 100644 --- a/packages/ui/src/providers/global-system-state/restart.tsx +++ b/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 ( diff --git a/packages/ui/src/providers/global-system-state/shutdown.tsx b/packages/ui/src/providers/global-system-state/shutdown.tsx index 801a3470b..c6fd779bb 100644 --- a/packages/ui/src/providers/global-system-state/shutdown.tsx +++ b/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 ( diff --git a/packages/ui/src/routes/login.tsx b/packages/ui/src/routes/login.tsx index 7cbca4cc2..081713750 100644 --- a/packages/ui/src/routes/login.tsx +++ b/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() { ) + const powerFooter = ( + <> + + + + ) + switch (step) { case 'password': { return ( <> - +
- + @@ -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('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) => { + 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 ( + + + + + + {step === 'password' ? ( +
+ + {t(titleKey)} + +
+ { + setPasswordError('') + setPassword(value) + }} + error={passwordError} + /> +
+ + + {t(submitKey)} + + {t('cancel')} + +
+ ) : ( +
+ + {t(titleKey)} + {t('login-2fa.subtitle')} + +
+ +
+ + {t('cancel')} + +
+ )} + + + ) +} diff --git a/packages/umbreld/source/modules/system/routes.ts b/packages/umbreld/source/modules/system/routes.ts index ec7e9613e..f5024de43 100644 --- a/packages/umbreld/source/modules/system/routes.ts +++ b/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) { + 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'