diff --git a/justfile b/justfile index 9e53b3e33..daabed537 100644 --- a/justfile +++ b/justfile @@ -1,4 +1,4 @@ -set shell := ["powershell.exe", "-c"] +set windows-shell := ["powershell.exe", "-c"] dev: npx concurrently \ diff --git a/new-ui/src/pages/full/AddInstancePage/AddInstancePage.tsx b/new-ui/src/pages/full/AddInstancePage/AddInstancePage.tsx index 2ee2f08b4..3cc9c578a 100644 --- a/new-ui/src/pages/full/AddInstancePage/AddInstancePage.tsx +++ b/new-ui/src/pages/full/AddInstancePage/AddInstancePage.tsx @@ -1,6 +1,7 @@ import './style.scss'; import { useLoaderData, useNavigate, useSearch } from '@tanstack/react-router'; +import { error } from '@tauri-apps/plugin-log'; import { Fragment, useMemo } from 'react'; import z from 'zod'; import { Button } from '../../../shared/components/Button/Button'; @@ -68,16 +69,21 @@ export const AddInstancePage = () => { }); return; } + void error(`Failed to add instance: ${result.error ?? 'unknown error'}`); Snackbar.error('Communication error, contact administrator.'); return; } - if (result.cookie) { - await edgeApi.createDevice(value.url, result.cookie, value.name.trim()); + if (result.session_id) { + await edgeApi.createDevice(result.session_id, value.name.trim()); } - if (result.startResponse && !result.startResponse.user.enrolled && result.cookie) { + if ( + result.startResponse && + !result.startResponse.user.enrolled && + result.session_id + ) { useEnrollmentStore .getState() - .start(result.startResponse, value.url, result.cookie, undefined); + .start(result.startResponse, result.session_id, undefined); navigate({ to: '/full/enrollment', replace: true, diff --git a/new-ui/src/pages/full/EnrollmentPage/hooks/activateUser.ts b/new-ui/src/pages/full/EnrollmentPage/hooks/activateUser.ts index 78b6f06a7..df5412422 100644 --- a/new-ui/src/pages/full/EnrollmentPage/hooks/activateUser.ts +++ b/new-ui/src/pages/full/EnrollmentPage/hooks/activateUser.ts @@ -1,10 +1,10 @@ -import { edgeApi } from '../../../../shared/edge-api/api'; +import { api } from '../../../../shared/rust-api/api'; import { useEnrollmentStore } from './useEnrollmentStore'; /** - * Activate the enrolling user by calling the proxy activation API. + * Activate the enrolling user by calling the enrollment Tauri command. * - * Reads proxy URL, cookie, skipPassword, and userPassword from the store via + * Reads sessionId, skipPassword, and userPassword from the store via * {@linkcode useEnrollmentStore.getState} so every invocation uses the latest * values (no stale-closure risk after in-flight {@linkcode setState} updates). * @@ -12,9 +12,8 @@ import { useEnrollmentStore } from './useEnrollmentStore'; * body instead of sending `null` or an empty string. */ export const activateUser = async () => { - const { proxyUrl, sessionCookie, skipPassword, userPassword } = - useEnrollmentStore.getState(); - const body = skipPassword ? {} : { password: userPassword ?? '' }; - // biome-ignore lint/style/noNonNullAssertion: proxy and cookie are set in start() - return edgeApi.activateUser(proxyUrl!, sessionCookie!, body); + const { sessionId, skipPassword, userPassword } = useEnrollmentStore.getState(); + const password = skipPassword ? null : (userPassword ?? ''); + // biome-ignore lint/style/noNonNullAssertion: sessionId is set in start(), called before this + return api.enrollmentActivateUser(sessionId!, password, null); }; diff --git a/new-ui/src/pages/full/EnrollmentPage/hooks/useEnrollmentStore.tsx b/new-ui/src/pages/full/EnrollmentPage/hooks/useEnrollmentStore.tsx index 0f3cb0869..60b3bf8da 100644 --- a/new-ui/src/pages/full/EnrollmentPage/hooks/useEnrollmentStore.tsx +++ b/new-ui/src/pages/full/EnrollmentPage/hooks/useEnrollmentStore.tsx @@ -1,14 +1,16 @@ import dayjs from 'dayjs'; import { create } from 'zustand'; import { createJSONStorage, persist } from 'zustand/middleware'; -import type { EnrollmentStartResponse } from '../../../../shared/edge-api/types'; -import type { MfaMethodValue } from '../../../../shared/rust-api/types'; +import type { + EnrollmentStartResult, + MfaMethodValue, +} from '../../../../shared/rust-api/types'; import { EnrollmentStep, type EnrollmentStepValue } from '../types'; type StoreValues = { activeStep: EnrollmentStepValue; - startResponse: EnrollmentStartResponse | null; - proxyUrl: string | null; + startResponse: EnrollmentStartResult | null; + sessionId: string | null; skipMfa: boolean; skipPassword: boolean; skipMfaChoice: boolean; @@ -17,13 +19,11 @@ type StoreValues = { userTotpSecret: string | null; userRecoveryCodes: string[] | null; userMfaChoice: MfaMethodValue; - sessionCookie: string | null; }; const defaults: StoreValues = { activeStep: EnrollmentStep.Welcome, - sessionCookie: null, - proxyUrl: null, + sessionId: null, userRecoveryCodes: null, startResponse: null, skipMfa: false, @@ -37,9 +37,8 @@ const defaults: StoreValues = { interface Store extends StoreValues { start: ( - response: EnrollmentStartResponse, - url: string, - cookie: string, + response: EnrollmentStartResult, + sessionId: string, totpSecret?: string, ) => void; next: () => void; @@ -50,13 +49,12 @@ export const useEnrollmentStore = create()( persist( (set, get) => ({ ...defaults, - start: (response, url, cookie, secret) => { + start: (response, sessionId, secret) => { set({ ...defaults, - proxyUrl: url, + sessionId, activeStep: EnrollmentStep.Welcome, startResponse: response, - sessionCookie: cookie, // if smtp is not present then it's not possible to choose. skipMfaChoice: !response.settings.smtp_configured || !response.settings.mfa_required, diff --git a/new-ui/src/pages/full/EnrollmentPage/steps/FinishStep/FinishStep.tsx b/new-ui/src/pages/full/EnrollmentPage/steps/FinishStep/FinishStep.tsx index f1a453d14..7ca28f94f 100644 --- a/new-ui/src/pages/full/EnrollmentPage/steps/FinishStep/FinishStep.tsx +++ b/new-ui/src/pages/full/EnrollmentPage/steps/FinishStep/FinishStep.tsx @@ -4,6 +4,7 @@ import { Button } from '../../../../../shared/components/Button/Button'; import { ButtonVariant } from '../../../../../shared/components/Button/types'; import { Controls } from '../../../../../shared/components/Controls/Controls'; import { RenderMarkdown } from '../../../../../shared/components/RenderMarkdown/RenderMarkdown'; +import { api } from '../../../../../shared/rust-api/api'; import { isPresent } from '../../../../../shared/utils/isPresent'; import { useEnrollmentStore } from '../../hooks/useEnrollmentStore'; import bannerSrc from './assets/banner.png'; @@ -11,6 +12,7 @@ import bannerSrc from './assets/banner.png'; export const FinishStep = () => { const navigate = useNavigate(); const response = useEnrollmentStore((s) => s.startResponse); + const sessionId = useEnrollmentStore((s) => s.sessionId); const markdownContent = response?.final_page_content; const showMarkdown = isPresent(markdownContent) && markdownContent.length > 0; @@ -32,6 +34,12 @@ export const FinishStep = () => { variant={ButtonVariant.Primary} text="Got it" onClick={() => { + // Release the enrollment session now that the wizard is done. + // Best-effort: enrollment_finish errors if the session is already + // gone, so swallow failures and leave the wizard regardless. + if (isPresent(sessionId)) { + void api.enrollmentFinish(sessionId).catch(() => {}); + } navigate({ to: '/full/overview', replace: true }); }} /> diff --git a/new-ui/src/pages/full/EnrollmentPage/steps/MfaChoiceStep/MfaChoiceStep.tsx b/new-ui/src/pages/full/EnrollmentPage/steps/MfaChoiceStep/MfaChoiceStep.tsx index c4a2bb124..47952f297 100644 --- a/new-ui/src/pages/full/EnrollmentPage/steps/MfaChoiceStep/MfaChoiceStep.tsx +++ b/new-ui/src/pages/full/EnrollmentPage/steps/MfaChoiceStep/MfaChoiceStep.tsx @@ -1,8 +1,8 @@ -/** biome-ignore-all lint/style/noNonNullAssertion: this cannot be shown without those store values */ +/** biome-ignore-all lint/style/noNonNullAssertion: sessionId is set in start(), called before */ import { useMutation } from '@tanstack/react-query'; +import './style.scss'; import clsx from 'clsx'; import { useState } from 'react'; -import { useShallow } from 'zustand/shallow'; import { Icon, IconKind, @@ -10,28 +10,25 @@ import { } from '../../../../../shared/components/Icon'; import { RadioIndicator } from '../../../../../shared/components/RadioIndicator/RadioIndicator'; import { SizedBox } from '../../../../../shared/components/SizedBox/SizedBox'; -import { edgeApi } from '../../../../../shared/edge-api/api'; +import { api } from '../../../../../shared/rust-api/api'; import { MfaMethod, type MfaMethodValue } from '../../../../../shared/rust-api/types'; import { ThemeSpacing } from '../../../../../shared/types'; import { EnrollmentControls } from '../../components/EnrollmentControls/EnrollmentControls'; import { useEnrollmentStore } from '../../hooks/useEnrollmentStore'; export const MfaChoiceStep = () => { - const [proxyUrl, cookie] = useEnrollmentStore( - useShallow((s) => [s.proxyUrl!, s.sessionCookie!]), - ); + const sessionId = useEnrollmentStore((s) => s.sessionId); const [mfa, setMfa] = useState( useEnrollmentStore.getState().userMfaChoice, ); const { mutate, isPending } = useMutation({ - mutationFn: () => edgeApi.startMfaSetup(proxyUrl, cookie, mfa), + mutationFn: () => api.enrollmentRegisterMfaStart(sessionId!, mfa), onSuccess: (resp) => { if (mfa === MfaMethod.Totp) { - const secret = resp.result?.totp_secret; useEnrollmentStore.setState({ userMfaChoice: mfa, - userTotpSecret: secret, + userTotpSecret: resp.totp_secret, }); } else { useEnrollmentStore.setState({ userMfaChoice: mfa }); diff --git a/new-ui/src/pages/full/EnrollmentPage/steps/MfaConfigurationStep/MfaConfigurationStep.tsx b/new-ui/src/pages/full/EnrollmentPage/steps/MfaConfigurationStep/MfaConfigurationStep.tsx index e349378a0..d6061246f 100644 --- a/new-ui/src/pages/full/EnrollmentPage/steps/MfaConfigurationStep/MfaConfigurationStep.tsx +++ b/new-ui/src/pages/full/EnrollmentPage/steps/MfaConfigurationStep/MfaConfigurationStep.tsx @@ -1,13 +1,13 @@ import { Fragment, useCallback, useEffect, useMemo, useState } from 'react'; import './style.scss'; import { useMutation } from '@tanstack/react-query'; -import { useShallow } from 'zustand/shallow'; +import { error as logError } from '@tauri-apps/plugin-log'; import { CodeInput } from '../../../../../shared/components/CodeInput/CodeInput'; import { CopyField } from '../../../../../shared/components/CopyField/CopyField'; import { Divider } from '../../../../../shared/components/Divider/Divider'; import { QrCard } from '../../../../../shared/components/QrCard/QrCard'; import { SizedBox } from '../../../../../shared/components/SizedBox/SizedBox'; -import { edgeApi } from '../../../../../shared/edge-api/api'; +import { api } from '../../../../../shared/rust-api/api'; import { MfaMethod } from '../../../../../shared/rust-api/types'; import { ThemeSpacing } from '../../../../../shared/types'; import { isPresent } from '../../../../../shared/utils/isPresent'; @@ -15,43 +15,57 @@ import { EnrollmentControls } from '../../components/EnrollmentControls/Enrollme import { activateUser } from '../../hooks/activateUser'; import { useEnrollmentStore } from '../../hooks/useEnrollmentStore'; +/** + * Map a register-mfa-finish error to a user-facing message. The proxy rejects + * a wrong code with a `proxy_error` (HTTP 400) whose message contains "invalid" + * (core returns "Code invalid" / "Email code invalid"); anything else is + * unexpected and gets a generic message. + */ +const mfaFinishErrorMessage = (err: unknown): string => { + try { + const parsed = JSON.parse(String(err)) as { type?: string; message?: string }; + if ( + parsed.type === 'proxy_error' && + parsed.message?.toLowerCase().includes('invalid') + ) { + return 'Invalid code'; + } + } catch { + // Non-JSON error: fall through to the generic message. + } + return 'MFA setup failed. Please try again.'; +}; + export const MfaConfigurationStep = () => { + const sessionId = useEnrollmentStore((s) => s.sessionId); const method = useEnrollmentStore((s) => s.userMfaChoice); const [code, setCode] = useState(null); const [error, setError] = useState(null); - const [proxyUrl, cookie] = useEnrollmentStore( - // biome-ignore lint/style/noNonNullAssertion: safe - useShallow((s) => [s.proxyUrl!, s.sessionCookie!]), - ); const { mutate, isPending } = useMutation({ mutationFn: async () => { - const resp = await edgeApi.finishMfaSetup(proxyUrl, cookie, { - // biome-ignore lint/style/noNonNullAssertion: checked in handleSubmit - code: code!, - method, - }); + // biome-ignore lint/style/noNonNullAssertion: checked in handleSubmit + const resp = await api.enrollmentRegisterMfaFinish(sessionId!, code!, method); await activateUser(); return resp; }, - onError: () => {}, + onError: (err) => { + void logError(`MFA configuration failed: ${err}`); + setError(mfaFinishErrorMessage(err)); + }, onSuccess: (resp) => { - if (resp.result) { - useEnrollmentStore.setState({ - userRecoveryCodes: resp.result.recovery_codes, - deadline: null, - }); - useEnrollmentStore.getState().next(); - } - if (resp.error) { - setError('Enter a valid code'); - } + useEnrollmentStore.setState({ + userRecoveryCodes: resp.recovery_codes, + deadline: null, + }); + useEnrollmentStore.getState().next(); }, }); const handleSubmit = useCallback(() => { if (code?.trim().length !== 6) { - setError(''); + setError('Enter a valid code'); + return; } mutate(); }, [code, mutate]); diff --git a/new-ui/src/pages/full/EnrollmentPage/steps/PasswordStep/PasswordStep.tsx b/new-ui/src/pages/full/EnrollmentPage/steps/PasswordStep/PasswordStep.tsx index f1af01e37..a6f590b4f 100644 --- a/new-ui/src/pages/full/EnrollmentPage/steps/PasswordStep/PasswordStep.tsx +++ b/new-ui/src/pages/full/EnrollmentPage/steps/PasswordStep/PasswordStep.tsx @@ -1,16 +1,16 @@ /** biome-ignore-all lint/style/noNonNullAssertion: ensured by wizard page */ +import { error } from '@tauri-apps/plugin-log'; import './style.scss'; import { useStore } from '@tanstack/react-form'; import { useMutation } from '@tanstack/react-query'; import clsx from 'clsx'; import { useCallback } from 'react'; import z from 'zod'; -import { useShallow } from 'zustand/shallow'; import { Icon, type IconKindValue } from '../../../../../shared/components/Icon'; import { SizedBox } from '../../../../../shared/components/SizedBox/SizedBox'; -import { edgeApi } from '../../../../../shared/edge-api/api'; import { useAppForm, withForm } from '../../../../../shared/form'; import { formChangeLogic } from '../../../../../shared/formLogic'; +import { api } from '../../../../../shared/rust-api/api'; import { MfaMethod } from '../../../../../shared/rust-api/types'; import { ThemeSpacing } from '../../../../../shared/types'; import { EnrollmentControls } from '../../components/EnrollmentControls/EnrollmentControls'; @@ -59,11 +59,9 @@ const defaultValues: FormFields = { }; export const PasswordStep = () => { - const [proxyUrl, cookie] = useEnrollmentStore( - useShallow((s) => [s.proxyUrl!, s.sessionCookie!]), - ); + const sessionId = useEnrollmentStore((s) => s.sessionId); const { mutateAsync: startMfa } = useMutation({ - mutationFn: () => edgeApi.startMfaSetup(proxyUrl, cookie, MfaMethod.Totp), + mutationFn: () => api.enrollmentRegisterMfaStart(sessionId!, MfaMethod.Totp), }); const form = useAppForm({ @@ -74,25 +72,24 @@ export const PasswordStep = () => { onChange: formSchema, }, onSubmit: async ({ value }) => { - const { skipMfaChoice, skipMfa } = useEnrollmentStore.getState(); - if (skipMfaChoice) { - const mfaResponse = await startMfa(); - if (mfaResponse.result) { + try { + const { skipMfaChoice, skipMfa } = useEnrollmentStore.getState(); + if (skipMfaChoice) { + const mfaResponse = await startMfa(); useEnrollmentStore.setState({ - userTotpSecret: mfaResponse.result.totp_secret ?? null, + userTotpSecret: mfaResponse.totp_secret, }); - } else { - console.error(mfaResponse); - return; } + useEnrollmentStore.setState({ + userPassword: value.password.trim(), + }); + if (skipMfa) { + await activateUser(); + } + useEnrollmentStore.getState().next(); + } catch (err) { + void error(`Password step failed: ${err}`); } - useEnrollmentStore.setState({ - userPassword: value.password.trim(), - }); - if (skipMfa) { - await activateUser(); - } - useEnrollmentStore.getState().next(); }, }); diff --git a/new-ui/src/pages/full/OverviewPage/components/ConnectModal/hooks/useConnectModalMfaOidc.ts b/new-ui/src/pages/full/OverviewPage/components/ConnectModal/hooks/useConnectModalMfaOidc.ts index de844cd13..cb9addd9e 100644 --- a/new-ui/src/pages/full/OverviewPage/components/ConnectModal/hooks/useConnectModalMfaOidc.ts +++ b/new-ui/src/pages/full/OverviewPage/components/ConnectModal/hooks/useConnectModalMfaOidc.ts @@ -1,24 +1,22 @@ -import { useMutation, useQuery } from '@tanstack/react-query'; -import { fetch } from '@tauri-apps/plugin-http'; +import { useQuery } from '@tanstack/react-query'; +import type { UnlistenFn } from '@tauri-apps/api/event'; +import { listen } from '@tauri-apps/api/event'; import { error } from '@tauri-apps/plugin-log'; import { useCallback, useEffect, useRef, useState } from 'react'; import { useShallow } from 'zustand/shallow'; -import { - CLIENT_MFA_ENDPOINT, - MfaStartMethod, - shouldShowPostureError, - startClientMfaSession, -} from '../../../../../../shared/components/LocationCard/api/startClientMfaSession'; import { api } from '../../../../../../shared/rust-api/api'; +import { + isConnectFailure, + isMfaPostureError, + isSessionExpired, + isTimeout, + mfaErrorMessage, +} from '../../../../../../shared/rust-api/mfaError'; import { getInstancesQueryOptions } from '../../../../../../shared/rust-api/query'; +import type { MfaErrorPayload } from '../../../../../../shared/rust-api/types'; +import { MfaMethod, TauriEvent } from '../../../../../../shared/rust-api/types'; import { useConnectModal } from './useConnectModal'; -const POLL_INTERVAL_MS = 5_000; -const POLL_TIMEOUT_MS = 5 * 60 * 1_000; - -type MfaFinishResponse = { preshared_key: string }; -type MfaErrorResponse = { error: string }; - type Options = { onPostureError?: (msg: string) => void; onSessionExpired?: () => void; @@ -38,92 +36,25 @@ export const useConnectModalMfaOidc = ({ const { data: instances } = useQuery(getInstancesQueryOptions); const instance = instances?.find((i) => i.id === location?.instance_id); - const intervalRef = useRef | null>(null); - const timeoutRef = useRef | null>(null); - - const { mutate: connectMutate } = useMutation({ - mutationFn: api.connect, - onError: (err) => { - error(`Connect command failed after successful OIDC MFA\n${err}`); - setPollError('Failed to establish VPN connection'); - }, - }); - - const stopPolling = useCallback(() => { - if (intervalRef.current !== null) { - clearInterval(intervalRef.current); - intervalRef.current = null; - } - if (timeoutRef.current !== null) { - clearTimeout(timeoutRef.current); - timeoutRef.current = null; + const taskIdRef = useRef(null); + const unlistenRef = useRef(null); + + const cleanup = useCallback(() => { + if (unlistenRef.current !== null) { + unlistenRef.current(); + unlistenRef.current = null; } }, []); useEffect(() => { return () => { - stopPolling(); + cleanup(); + const taskId = taskIdRef.current; + if (taskId) { + void api.cancelMfa(taskId).catch(() => {}); + } }; - }, [stopPolling]); - - const startPolling = useCallback( - (token: string, proxyUrl: string, headers: Record) => { - if (!location) return; - - setIsPolling(true); - setPollError(null); - - const poll = async () => { - try { - const res = await fetch(`${proxyUrl}${CLIENT_MFA_ENDPOINT}/finish`, { - method: 'POST', - headers: { 'Content-Type': 'application/json', ...headers }, - body: JSON.stringify({ token }), - }); - - if (res.ok) { - stopPolling(); - setIsPolling(false); - const data = (await res.json()) as MfaFinishResponse; - connectMutate({ - locationId: location.id, - connectionType: location.connection_type, - presharedKey: data.preshared_key, - }); - return; - } - - if (res.status === 428) return; - - stopPolling(); - setIsPolling(false); - const data = (await res.json()) as MfaErrorResponse; - const msg = data.error; - if (msg === 'invalid token' || msg === 'login session not found') { - onSessionExpired?.(); - } else { - setPollError('Authentication failed. Please try again.'); - } - error(`OIDC MFA poll failed for location ${location.id}: ${msg}`); - } catch (e) { - stopPolling(); - setIsPolling(false); - setPollError('Failed to reach server. Please try again.'); - error(`OIDC MFA poll network error for location ${location.id}: ${e}`); - } - }; - - intervalRef.current = setInterval(poll, POLL_INTERVAL_MS); - - timeoutRef.current = setTimeout(() => { - stopPolling(); - setIsPolling(false); - setPollError('Authentication timed out. Please try again.'); - error(`OIDC MFA timed out for location ${location.id}`); - }, POLL_TIMEOUT_MS); - }, - [location, connectMutate, stopPolling, onSessionExpired], - ); + }, [cleanup]); const start = useCallback(async () => { if (!instance || !location) { @@ -134,29 +65,58 @@ export const useConnectModalMfaOidc = ({ setIsStarting(true); setStartError(null); setPollError(null); - stopPolling(); + cleanup(); try { - const { response, headers } = await startClientMfaSession({ - instance, - location, - method: MfaStartMethod.Oidc, + const info = await api.mfaStart(instance.id, location.id, MfaMethod.Oidc); + await api.openLink(`${instance.proxy_url}openid/mfa?token=${info.token}`); + + setIsStarting(false); + setIsPolling(true); + + const taskId = await api.mfaPollOpenId(instance.id, location.id, info.token); + taskIdRef.current = taskId; + + // The backend brings up the connection itself; completion means connected. + const completeUnlisten = await listen(TauriEvent.MfaOpenIdComplete, () => { + cleanup(); + setIsPolling(false); }); - await api.openLink(`${instance.proxy_url}openid/mfa?token=${response.token}`); - startPolling(response.token, instance.proxy_url, headers); + + const errorUnlisten = await listen( + TauriEvent.MfaOpenIdError, + (event) => { + cleanup(); + setIsPolling(false); + error(`OIDC MFA failed for location ${location.id}: ${event.payload.error}`); + const message = mfaErrorMessage(event.payload.error); + if (isTimeout(event.payload.error)) { + setPollError('Authentication timed out. Please try again.'); + } else if (isConnectFailure(message)) { + setPollError('Failed to establish VPN connection'); + } else if (isSessionExpired(message)) { + onSessionExpired?.(); + } else { + setPollError('Authentication failed. Please try again.'); + } + }, + ); + + unlistenRef.current = () => { + completeUnlisten(); + errorUnlisten(); + }; } catch (e) { - if (shouldShowPostureError(e, location)) { - onPostureError?.(e.message); + void error(`OIDC MFA start failed for location ${location.id}: ${e}`); + if (isMfaPostureError(e, location)) { + onPostureError?.(mfaErrorMessage(e)); return; } - setStartError( - e instanceof Error ? e.message : 'Failed to start OIDC authentication', - ); - error(`OIDC MFA start network error for location ${location.id}: ${e}`); + setStartError(mfaErrorMessage(e)); } finally { setIsStarting(false); } - }, [instance, location, startPolling, stopPolling, onPostureError]); + }, [instance, location, cleanup, onPostureError, onSessionExpired]); return { start, isStarting, startError, isPolling, pollError }; }; diff --git a/new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaEmail/ConnectModalMfaEmail.tsx b/new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaEmail/ConnectModalMfaEmail.tsx index f17cf6ef9..0f866ae49 100644 --- a/new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaEmail/ConnectModalMfaEmail.tsx +++ b/new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaEmail/ConnectModalMfaEmail.tsx @@ -4,8 +4,10 @@ import { Button } from '../../../../../../../shared/components/Button/Button'; import { ButtonVariant } from '../../../../../../../shared/components/Button/types'; import { CodeInput } from '../../../../../../../shared/components/CodeInput/CodeInput'; import { Controls } from '../../../../../../../shared/components/Controls/Controls'; -import { MfaStartMethod } from '../../../../../../../shared/components/LocationCard/api/startClientMfaSession'; -import { useMfaConnect } from '../../../../../../../shared/components/LocationCard/hooks/useMfaConnect'; +import { + MfaStartMethod, + useMfaConnect, +} from '../../../../../../../shared/components/LocationCard/hooks/useMfaConnect'; import type { LocationInfo } from '../../../../../../../shared/rust-api/types'; import { isPresent } from '../../../../../../../shared/utils/isPresent'; import { ConnectModalPostureCheckLoading } from '../../components/ConnectModalPostureCheckLoading/ConnectModalPostureCheckLoading'; diff --git a/new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaTotp/ConnectModalMfaTotp.tsx b/new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaTotp/ConnectModalMfaTotp.tsx index dba162116..d6beb7201 100644 --- a/new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaTotp/ConnectModalMfaTotp.tsx +++ b/new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaTotp/ConnectModalMfaTotp.tsx @@ -4,8 +4,10 @@ import { Button } from '../../../../../../../shared/components/Button/Button'; import { ButtonVariant } from '../../../../../../../shared/components/Button/types'; import { CodeInput } from '../../../../../../../shared/components/CodeInput/CodeInput'; import { Controls } from '../../../../../../../shared/components/Controls/Controls'; -import { MfaStartMethod } from '../../../../../../../shared/components/LocationCard/api/startClientMfaSession'; -import { useMfaConnect } from '../../../../../../../shared/components/LocationCard/hooks/useMfaConnect'; +import { + MfaStartMethod, + useMfaConnect, +} from '../../../../../../../shared/components/LocationCard/hooks/useMfaConnect'; import type { LocationInfo } from '../../../../../../../shared/rust-api/types'; import { isPresent } from '../../../../../../../shared/utils/isPresent'; import { ConnectModalPostureCheckLoading } from '../../components/ConnectModalPostureCheckLoading/ConnectModalPostureCheckLoading'; diff --git a/new-ui/src/pages/full/OverviewPage/components/UpdateTunnelModal/UpdateTunnelModal.tsx b/new-ui/src/pages/full/OverviewPage/components/UpdateTunnelModal/UpdateTunnelModal.tsx index 44b3ac98f..9aaa6950a 100644 --- a/new-ui/src/pages/full/OverviewPage/components/UpdateTunnelModal/UpdateTunnelModal.tsx +++ b/new-ui/src/pages/full/OverviewPage/components/UpdateTunnelModal/UpdateTunnelModal.tsx @@ -1,4 +1,5 @@ import { useMutation } from '@tanstack/react-query'; +import { error as logError } from '@tauri-apps/plugin-log'; import { useEffect, useMemo, useState } from 'react'; import z from 'zod'; import { Button } from '../../../../../shared/components/Button/Button'; @@ -133,7 +134,8 @@ const ModalContent = ({ data }: { data: OpenUpdateTunnelModalData }) => { Snackbar.default('Tunnel updated.'); closeModal(modalNameKey); } catch (e) { - Snackbar.error(String(e)); + void logError(`Tunnel update failed: ${e}`); + Snackbar.error('Failed to update tunnel.'); } }, }); diff --git a/new-ui/src/shared/components/LocationCard/api/startClientMfaSession.ts b/new-ui/src/shared/components/LocationCard/api/startClientMfaSession.ts deleted file mode 100644 index be527d949..000000000 --- a/new-ui/src/shared/components/LocationCard/api/startClientMfaSession.ts +++ /dev/null @@ -1,119 +0,0 @@ -import { fetch } from '@tauri-apps/plugin-http'; -import { edgeApi } from '../../../edge-api/api'; -import type { EdgeRequestHeaders } from '../../../edge-api/types'; -import { api } from '../../../rust-api/api'; -import type { InstanceInfo, LocationInfo } from '../../../rust-api/types'; - -export const CLIENT_MFA_ENDPOINT = 'api/v1/client-mfa'; - -/** Error raised when the MFA start request or its prerequisites fail. */ -export class MfaStartError extends Error { - public readonly status?: number; - - constructor(message: string, status?: number) { - super(message); - this.name = 'MfaStartError'; - this.status = status; - } -} - -/** MFA method identifiers expected by the desktop-client MFA API */ -export const MfaStartMethod = { - Totp: 0, - Email: 1, - Oidc: 2, - MobileApprove: 4, -} as const; - -export type MfaStartMethod = (typeof MfaStartMethod)[keyof typeof MfaStartMethod]; - -/** Successful MFA start response returned by the proxy. */ -export type MfaStartResponse = { - token: string; - challenge?: string; -}; - -/** Error response shape returned by the proxy for MFA start failures. */ -type MfaStartErrorResponse = { - error?: string; -}; - -/** Narrows MFA start errors that should open the posture failure view. */ -export const shouldShowPostureError = ( - err: unknown, - location: LocationInfo, -): err is MfaStartError => - err instanceof MfaStartError && err.status === 403 && location.posture_check_required; - -/** Input required to start a desktop-client MFA session. */ -type StartClientMfaSessionParams = { - instance: InstanceInfo; - location: LocationInfo; - method: MfaStartMethod; -}; - -/** MFA start response plus request headers required by later MFA calls. */ -type StartClientMfaSessionResult = { - response: MfaStartResponse; - headers: EdgeRequestHeaders; -}; - -/** Starts an MFA session, including posture data when the location requires it. */ -export const startClientMfaSession = async ({ - instance, - location, - method, -}: StartClientMfaSessionParams): Promise => { - let headers: EdgeRequestHeaders; - try { - headers = await edgeApi.getEdgeRequestHeaders(); - } catch { - throw new MfaStartError('Failed to load request headers'); - } - - let posture_data: unknown; - try { - posture_data = location.posture_check_required - ? await api.getPostureData() - : undefined; - } catch { - throw new MfaStartError('Failed to load posture data'); - } - - try { - const response = await fetch(`${instance.proxy_url}${CLIENT_MFA_ENDPOINT}/start`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - ...headers, - }, - body: JSON.stringify({ - method, - pubkey: instance.pubkey, - location_id: location.network_id, - posture_data, - }), - }); - - if (!response.ok) { - let message = 'Failed to start MFA'; - try { - const data = (await response.json()) as MfaStartErrorResponse; - message = data.error ?? message; - } catch { - // Keep the response status even if the proxy sends a malformed error body. - } - throw new MfaStartError(message, response.status); - } - - return { - response: (await response.json()) as MfaStartResponse, - headers, - }; - } catch (err) { - if (err instanceof MfaStartError) { - throw err; - } - throw new MfaStartError('Failed to reach server'); - } -}; diff --git a/new-ui/src/shared/components/LocationCard/hooks/handleMfaStartError.ts b/new-ui/src/shared/components/LocationCard/hooks/handleMfaStartError.ts deleted file mode 100644 index fbc929dbf..000000000 --- a/new-ui/src/shared/components/LocationCard/hooks/handleMfaStartError.ts +++ /dev/null @@ -1,26 +0,0 @@ -import type { LocationInfo } from '../../../rust-api/types'; -import { shouldShowPostureError } from '../api/startClientMfaSession'; -import { LocationCardViews, type LocationCardViewsValue } from '../context/types'; - -type HandleMfaStartErrorParams = { - err: unknown; - location: LocationInfo; - setPostureError: (error: string | null) => void; - setView: (view: LocationCardViewsValue) => void; -}; - -/** Handles MFA start posture failures and reports whether the error was consumed. */ -export const handleMfaStartError = ({ - err, - location, - setPostureError, - setView, -}: HandleMfaStartErrorParams): boolean => { - if (!shouldShowPostureError(err, location)) { - return false; - } - - setPostureError(err.message); - setView(LocationCardViews.PostureCheckFail); - return true; -}; diff --git a/new-ui/src/shared/components/LocationCard/hooks/useMfaConnect.ts b/new-ui/src/shared/components/LocationCard/hooks/useMfaConnect.ts index 97be47dbe..4ac1707d5 100644 --- a/new-ui/src/shared/components/LocationCard/hooks/useMfaConnect.ts +++ b/new-ui/src/shared/components/LocationCard/hooks/useMfaConnect.ts @@ -1,28 +1,36 @@ -import { useMutation, useQuery } from '@tanstack/react-query'; -import { fetch } from '@tauri-apps/plugin-http'; +import { useQuery } from '@tanstack/react-query'; import { error } from '@tauri-apps/plugin-log'; import { useCallback, useEffect, useRef, useState } from 'react'; -import type { EdgeRequestHeaders } from '../../../edge-api/types'; import { api } from '../../../rust-api/api'; +import { + isConnectFailure, + isInvalidCode, + isMfaPostureError, + isSessionExpired, + mfaErrorMessage, +} from '../../../rust-api/mfaError'; import { getInstancesQueryOptions } from '../../../rust-api/query'; import type { LocationInfo } from '../../../rust-api/types'; -import { - CLIENT_MFA_ENDPOINT, - type MfaStartMethod, - shouldShowPostureError, - startClientMfaSession, -} from '../api/startClientMfaSession'; - -type MfaFinishResponse = { - preshared_key: string; -}; +import { MfaMethod } from '../../../rust-api/types'; -type MfaErrorResponse = { - error: string; -}; +/** MFA method identifiers matching the proxy proto enum (numeric). + * Used by TOTP/email view callers that pass the method to useMfaConnect. */ +export const MfaStartMethod = { + Totp: 0, + Email: 1, + Oidc: 2, + MobileApprove: 4, +} as const; + +export type MfaStartMethod = (typeof MfaStartMethod)[keyof typeof MfaStartMethod]; type CodeMfaStartMethod = Extract; +const MFA_METHOD_MAP: Record = { + [MfaStartMethod.Totp]: MfaMethod.Totp, + [MfaStartMethod.Email]: MfaMethod.Email, +}; + type UseMfaConnectOptions = { debounceMs?: number; onConnected?: () => void; @@ -52,23 +60,11 @@ export const useMfaConnect = ( const [startError, setStartError] = useState(null); const [isVerifying, setIsVerifying] = useState(false); const [verifyError, setVerifyError] = useState(null); - const [requestHeaders, setRequestHeaders] = useState(null); const { data: instances } = useQuery(getInstancesQueryOptions); const instance = instances?.find((i) => i.id === location.instance_id); - const { mutate: connectMutate } = useMutation({ - mutationFn: api.connect, - meta: { invalidate: ['locations'] }, - onSuccess: () => { - onConnected?.(); - }, - onError: (err) => { - error(`Connect command failed after successful code verification\n${err}`); - }, - }); - // Fire the /start request exactly once when instance data is ready. const startCalled = useRef(false); @@ -81,23 +77,26 @@ export const useMfaConnect = ( setIsStarting(true); + const methodString = MFA_METHOD_MAP[method]; + if (!methodString) { + setStartError('Unsupported MFA method'); + setIsStarting(false); + return; + } + (async () => { try { - const { response, headers } = await startClientMfaSession({ - instance, - location, - method, - }); + const info = await api.mfaStart(instance.id, location.id, methodString); await waitForMinimumDuration(startedAt, debounceMs); - setRequestHeaders(headers); - setToken(response.token); + setToken(info.token); } catch (err) { + void error(`MFA start failed: ${err}`); await waitForMinimumDuration(startedAt, debounceMs); - if (shouldShowPostureError(err, location)) { - onPostureError?.(err.message); + if (isMfaPostureError(err, location)) { + onPostureError?.(mfaErrorMessage(err)); return; } - setStartError(err instanceof Error ? err.message : 'Failed to start MFA'); + setStartError(mfaErrorMessage(err)); } finally { setIsStarting(false); } @@ -106,51 +105,33 @@ export const useMfaConnect = ( const verifyCode = useCallback( async (code: string) => { - if (!token || !instance || !requestHeaders) return; + if (!token || !instance) return; setIsVerifying(true); setVerifyError(null); - const body = JSON.stringify({ token, code }); - try { - const res = await fetch(`${instance.proxy_url}${CLIENT_MFA_ENDPOINT}/finish`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - ...requestHeaders, - }, - body, - }); - - if (res.ok) { - const data = (await res.json()) as MfaFinishResponse; - connectMutate({ - locationId: location.id, - connectionType: location.connection_type, - presharedKey: data.preshared_key, - }); + // mfaFinishCode completes MFA and brings up the connection in the + // backend; the preshared key never reaches the frontend. + await api.mfaFinishCode(instance.id, location.id, token, code); + onConnected?.(); + } catch (err) { + void error(`MFA verification failed: ${err}`); + const message = mfaErrorMessage(err); + if (isConnectFailure(message)) { + setVerifyError('Failed to establish VPN connection'); + } else if (isInvalidCode(message)) { + setVerifyError('Invalid code'); + } else if (isSessionExpired(message)) { + onSessionExpired?.(); } else { - const data = (await res.json()) as MfaErrorResponse; - const { error: errorMessage } = data; - if (errorMessage === 'Unauthorized') { - setVerifyError('Invalid code'); - } else if ( - errorMessage === 'invalid token' || - errorMessage === 'login session not found' - ) { - onSessionExpired?.(); - } else { - setVerifyError('Verification failed'); - } + setVerifyError('Verification failed'); } - } catch { - setVerifyError('Failed to reach server'); } finally { setIsVerifying(false); } }, - [token, instance, requestHeaders, location, connectMutate, onSessionExpired], + [token, instance, location, onConnected, onSessionExpired], ); return { token, isStarting, startError, verifyCode, isVerifying, verifyError }; diff --git a/new-ui/src/shared/components/LocationCard/hooks/useMfaMobileConnect.ts b/new-ui/src/shared/components/LocationCard/hooks/useMfaMobileConnect.ts index 050a1a40d..696574f61 100644 --- a/new-ui/src/shared/components/LocationCard/hooks/useMfaMobileConnect.ts +++ b/new-ui/src/shared/components/LocationCard/hooks/useMfaMobileConnect.ts @@ -1,16 +1,18 @@ import { encode } from '@stablelib/base64'; -import { useMutation, useQuery } from '@tanstack/react-query'; +import { useQuery } from '@tanstack/react-query'; +import type { UnlistenFn } from '@tauri-apps/api/event'; +import { listen } from '@tauri-apps/api/event'; import { error } from '@tauri-apps/plugin-log'; import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { api } from '../../../rust-api/api'; -import { getInstancesQueryOptions } from '../../../rust-api/query'; -import type { LocationInfo } from '../../../rust-api/types'; import { - CLIENT_MFA_ENDPOINT, - MfaStartMethod, - shouldShowPostureError, - startClientMfaSession, -} from '../api/startClientMfaSession'; + isConnectFailure, + isMfaPostureError, + mfaErrorMessage, +} from '../../../rust-api/mfaError'; +import { getInstancesQueryOptions } from '../../../rust-api/query'; +import type { LocationInfo, MfaErrorPayload } from '../../../rust-api/types'; +import { MfaMethod, TauriEvent } from '../../../rust-api/types'; type TokenData = { token: string; @@ -34,100 +36,92 @@ export const useMfaMobileConnect = (location: LocationInfo, options?: Options) = const [isConnecting, setIsConnecting] = useState(false); const [connectionError, setConnectionError] = useState(null); - const wsRef = useRef(null); - const expectedCloseRef = useRef(false); - - const { mutate: connectMutate } = useMutation({ - mutationFn: api.connect, - onSuccess: () => { - onConnected?.(); - }, - onError: (err) => { - error(`Connect command failed after successful mobile MFA\n${err}`); - setConnectionError('Failed to establish VPN connection'); - }, - }); - - // Open WebSocket when tokenData is available - useEffect(() => { - if (!tokenData || !instance) return; - - const wsUrl = `${instance.proxy_url - .replace(/^http:/, 'ws:') - .replace( - /^https:/, - 'wss:', - )}${CLIENT_MFA_ENDPOINT}/remote?token=${encodeURIComponent(tokenData.token)}`; + const taskIdRef = useRef(null); + const unlistenRef = useRef(null); - expectedCloseRef.current = false; - const ws = new WebSocket(wsUrl); - wsRef.current = ws; + const cleanupListeners = useCallback(() => { + if (unlistenRef.current !== null) { + unlistenRef.current(); + unlistenRef.current = null; + } + }, []); - ws.onopen = () => { - setIsConnecting(true); - setConnectionError(null); + // Clean up on unmount + useEffect(() => { + return () => { + cleanupListeners(); + const taskId = taskIdRef.current; + if (taskId) { + void api.cancelMfa(taskId).catch(() => {}); + } }; + }, [cleanupListeners]); + + // Connect WebSocket via Rust when tokenData is available + useEffect(() => { + if (!tokenData || !instance) return; + + let cancelled = false; + cleanupListeners(); + setIsConnecting(true); + setConnectionError(null); - ws.onmessage = (event: MessageEvent) => { + (async () => { try { - const parsed = JSON.parse(event.data as string) as unknown; - if ( - typeof parsed === 'object' && - parsed !== null && - 'preshared_key' in parsed && - typeof (parsed as Record).preshared_key === 'string' - ) { - const presharedKey = (parsed as { preshared_key: string }).preshared_key; - expectedCloseRef.current = true; - connectMutate({ - locationId: location.id, - connectionType: location.connection_type, - presharedKey, - }); - } else { - error( - `Unexpected mobile MFA WS message for location ${location.id}: ${event.data}`, - ); + const taskId = await api.mfaConnectMobileApprove( + instance.id, + location.id, + tokenData.token, + ); + if (cancelled) { + void api.cancelMfa(taskId).catch(() => {}); + return; } + taskIdRef.current = taskId; + + // The backend brings up the connection itself; completion means connected. + const completeUnlisten = await listen(TauriEvent.MfaMobileComplete, () => { + cleanupListeners(); + setIsConnecting(false); + onConnected?.(); + }); + + const errorUnlisten = await listen( + TauriEvent.MfaMobileError, + (event) => { + cleanupListeners(); + setIsConnecting(false); + error( + `Mobile MFA failed for location ${location.id}: ${event.payload.error}`, + ); + const message = mfaErrorMessage(event.payload.error); + setConnectionError( + isConnectFailure(message) + ? 'Failed to establish VPN connection' + : 'Connection error. Please try again.', + ); + }, + ); + + unlistenRef.current = () => { + completeUnlisten(); + errorUnlisten(); + }; } catch (e) { - error(`Failed to parse mobile MFA WS message for location ${location.id}: ${e}`); - } - }; - - ws.onerror = () => { - if (!expectedCloseRef.current) { - setIsConnecting(false); - setConnectionError('Connection error. Please try again.'); - error(`Mobile MFA WebSocket error for location ${location.id}`); - } - }; - - ws.onclose = () => { - if (!expectedCloseRef.current) { - setIsConnecting(false); - setConnectionError('Connection closed unexpectedly. Please try again.'); - error(`Mobile MFA WebSocket closed unexpectedly for location ${location.id}`); + if (!cancelled) { + setIsConnecting(false); + setConnectionError('Failed to start mobile approval. Please try again.'); + error(`Mobile MFA connect failed for location ${location.id}: ${e}`); + } } - }; + })(); return () => { - expectedCloseRef.current = true; - ws.close(); - wsRef.current = null; + cancelled = true; + cleanupListeners(); setIsConnecting(false); }; - }, [tokenData, instance, connectMutate, location]); - - // Clean up WebSocket on unmount - useEffect(() => { - return () => { - if (wsRef.current) { - expectedCloseRef.current = true; - wsRef.current.close(); - wsRef.current = null; - } - }; - }, []); + }, [tokenData, instance, location, onConnected, cleanupListeners]); const qrValue = useMemo(() => { if (!tokenData || !instance) return null; @@ -148,47 +142,42 @@ export const useMfaMobileConnect = (location: LocationInfo, options?: Options) = setIsStarting(true); setStartError(null); setConnectionError(null); - // Clear previous token → triggers WS cleanup via effect + // Clear previous task via effect setTokenData(null); try { - const { response } = await startClientMfaSession({ - instance, - location, - method: MfaStartMethod.MobileApprove, - }); - if (!response.challenge) { + const info = await api.mfaStart(instance.id, location.id, MfaMethod.MobileApprove); + if (!info.challenge) { setStartError('Unsupported response from proxy'); return; } - setTokenData({ token: response.token, challenge: response.challenge }); + setTokenData({ token: info.token, challenge: info.challenge }); } catch (e) { - if (shouldShowPostureError(e, location)) { - onPostureError?.(e instanceof Error ? e.message : undefined); + void error(`Mobile MFA start failed for location ${location.id}: ${e}`); + if (isMfaPostureError(e, location)) { + onPostureError?.(mfaErrorMessage(e)); return; } - setStartError( - e instanceof Error ? e.message : 'Failed to start mobile authentication', - ); - error(`Mobile MFA start network error for location ${location.id}: ${e}`); + setStartError(mfaErrorMessage(e)); } finally { setIsStarting(false); } }, [instance, location, onPostureError]); const reset = useCallback(() => { - if (wsRef.current) { - expectedCloseRef.current = true; - wsRef.current.close(); - wsRef.current = null; + cleanupListeners(); + const taskId = taskIdRef.current; + if (taskId) { + void api.cancelMfa(taskId).catch(() => {}); + taskIdRef.current = null; } setTokenData(null); setIsStarting(false); setStartError(null); setIsConnecting(false); setConnectionError(null); - }, []); + }, [cleanupListeners]); return { start, isStarting, startError, qrValue, isConnecting, connectionError, reset }; }; diff --git a/new-ui/src/shared/components/LocationCard/hooks/useMfaOidcConnect.ts b/new-ui/src/shared/components/LocationCard/hooks/useMfaOidcConnect.ts index 938b6bb65..77248bea1 100644 --- a/new-ui/src/shared/components/LocationCard/hooks/useMfaOidcConnect.ts +++ b/new-ui/src/shared/components/LocationCard/hooks/useMfaOidcConnect.ts @@ -1,23 +1,21 @@ -import { useMutation, useQuery } from '@tanstack/react-query'; -import { fetch } from '@tauri-apps/plugin-http'; +import { useQuery } from '@tanstack/react-query'; +import type { UnlistenFn } from '@tauri-apps/api/event'; +import { listen } from '@tauri-apps/api/event'; import { error } from '@tauri-apps/plugin-log'; import { useCallback, useEffect, useRef, useState } from 'react'; import { api } from '../../../rust-api/api'; -import { getInstancesQueryOptions } from '../../../rust-api/query'; import { - CLIENT_MFA_ENDPOINT, - MfaStartMethod, - startClientMfaSession, -} from '../api/startClientMfaSession'; + isConnectFailure, + isMfaPostureError, + isSessionExpired, + isTimeout, + mfaErrorMessage, +} from '../../../rust-api/mfaError'; +import { getInstancesQueryOptions } from '../../../rust-api/query'; +import type { MfaErrorPayload } from '../../../rust-api/types'; +import { MfaMethod, TauriEvent } from '../../../rust-api/types'; import { useLocationCardContext } from '../context/context'; import { LocationCardViews } from '../context/types'; -import { handleMfaStartError } from './handleMfaStartError'; - -const POLL_INTERVAL_MS = 5_000; -const POLL_TIMEOUT_MS = 5 * 60 * 1_000; // 5 minutes - -type MfaFinishResponse = { preshared_key: string }; -type MfaErrorResponse = { error: string }; export const useMfaOidcConnect = () => { const { location, setPostureError, setView } = useLocationCardContext(); @@ -30,96 +28,26 @@ export const useMfaOidcConnect = () => { const { data: instances } = useQuery(getInstancesQueryOptions); const instance = instances?.find((i) => i.id === location.instance_id); - const intervalRef = useRef | null>(null); - const timeoutRef = useRef | null>(null); - - const { mutate: connectMutate } = useMutation({ - mutationFn: api.connect, - onSuccess: () => { - setView(LocationCardViews.Connected); - }, - onError: (err) => { - error(`Connect command failed after successful OIDC MFA\n${err}`); - setPollError('Failed to establish VPN connection'); - }, - }); + const taskIdRef = useRef(null); + const unlistenRef = useRef(null); - const stopPolling = useCallback(() => { - if (intervalRef.current !== null) { - clearInterval(intervalRef.current); - intervalRef.current = null; - } - if (timeoutRef.current !== null) { - clearTimeout(timeoutRef.current); - timeoutRef.current = null; + const cleanup = useCallback(() => { + if (unlistenRef.current !== null) { + unlistenRef.current(); + unlistenRef.current = null; } }, []); // Clean up on unmount useEffect(() => { return () => { - stopPolling(); + cleanup(); + const taskId = taskIdRef.current; + if (taskId) { + void api.cancelMfa(taskId).catch(() => {}); + } }; - }, [stopPolling]); - - const startPolling = useCallback( - (token: string, proxyUrl: string, headers: Record) => { - setIsPolling(true); - setPollError(null); - - const poll = async () => { - try { - const res = await fetch(`${proxyUrl}${CLIENT_MFA_ENDPOINT}/finish`, { - method: 'POST', - headers: { 'Content-Type': 'application/json', ...headers }, - body: JSON.stringify({ token }), - }); - - if (res.ok) { - stopPolling(); - setIsPolling(false); - const data = (await res.json()) as MfaFinishResponse; - connectMutate({ - locationId: location.id, - connectionType: location.connection_type, - presharedKey: data.preshared_key, - }); - return; - } - - // 428 Precondition Required: user hasn't completed browser auth yet, keep polling - if (res.status === 428) return; - - // Any other error: stop polling and surface the error - stopPolling(); - setIsPolling(false); - const data = (await res.json()) as MfaErrorResponse; - const msg = data.error; - if (msg === 'invalid token' || msg === 'login session not found') { - setPollError('Session expired. Please try again.'); - } else { - setPollError('Authentication failed. Please try again.'); - } - error(`OIDC MFA poll failed for location ${location.id}: ${msg}`); - } catch (e) { - stopPolling(); - setIsPolling(false); - setPollError('Failed to reach server. Please try again.'); - error(`OIDC MFA poll network error for location ${location.id}: ${e}`); - } - }; - - intervalRef.current = setInterval(poll, POLL_INTERVAL_MS); - - timeoutRef.current = setTimeout(() => { - stopPolling(); - setIsPolling(false); - setPollError('Authentication timed out. Please try again.'); - error(`OIDC MFA timed out for location ${location.id}`); - }, POLL_TIMEOUT_MS); - }, - [location, connectMutate, stopPolling], - ); + }, [cleanup]); const start = useCallback(async () => { if (!instance) { @@ -130,28 +58,60 @@ export const useMfaOidcConnect = () => { setIsStarting(true); setStartError(null); setPollError(null); - stopPolling(); + cleanup(); try { - const { response, headers } = await startClientMfaSession({ - instance, - location, - method: MfaStartMethod.Oidc, + const info = await api.mfaStart(instance.id, location.id, MfaMethod.Oidc); + await api.openLink(`${instance.proxy_url}openid/mfa?token=${info.token}`); + + setIsStarting(false); + setIsPolling(true); + + const taskId = await api.mfaPollOpenId(instance.id, location.id, info.token); + taskIdRef.current = taskId; + + // The backend brings up the connection itself; completion means connected. + const completeUnlisten = await listen(TauriEvent.MfaOpenIdComplete, () => { + cleanup(); + setIsPolling(false); + setView(LocationCardViews.Connected); }); - await api.openLink(`${instance.proxy_url}openid/mfa?token=${response.token}`); - startPolling(response.token, instance.proxy_url, headers); + + const errorUnlisten = await listen( + TauriEvent.MfaOpenIdError, + (event) => { + cleanup(); + setIsPolling(false); + error(`OIDC MFA failed for location ${location.id}: ${event.payload.error}`); + const message = mfaErrorMessage(event.payload.error); + if (isTimeout(event.payload.error)) { + setPollError('Authentication timed out. Please try again.'); + } else if (isConnectFailure(message)) { + setPollError('Failed to establish VPN connection'); + } else if (isSessionExpired(message)) { + setPollError('Session expired. Please try again.'); + } else { + setPollError('Authentication failed. Please try again.'); + } + }, + ); + + unlistenRef.current = () => { + completeUnlisten(); + errorUnlisten(); + }; } catch (e) { - if (handleMfaStartError({ err: e, location, setPostureError, setView })) { + void error(`OIDC MFA start failed for location ${location.id}: ${e}`); + if (isMfaPostureError(e, location)) { + setPostureError(mfaErrorMessage(e)); + setView(LocationCardViews.PostureCheckFail); return; } - setStartError( - e instanceof Error ? e.message : 'Failed to start OIDC authentication', - ); - error(`OIDC MFA start network error for location ${location.id}: ${e}`); + setStartError(mfaErrorMessage(e)); } finally { setIsStarting(false); } - }, [instance, location, setPostureError, setView, startPolling, stopPolling]); + }, [instance, location, setPostureError, setView, cleanup]); return { start, isStarting, startError, isPolling, pollError }; }; diff --git a/new-ui/src/shared/components/LocationCard/views/LocationCardMfaEmailView/LocationCardMfaEmailView.tsx b/new-ui/src/shared/components/LocationCard/views/LocationCardMfaEmailView/LocationCardMfaEmailView.tsx index e03ecb8cc..61c282f51 100644 --- a/new-ui/src/shared/components/LocationCard/views/LocationCardMfaEmailView/LocationCardMfaEmailView.tsx +++ b/new-ui/src/shared/components/LocationCard/views/LocationCardMfaEmailView/LocationCardMfaEmailView.tsx @@ -10,11 +10,10 @@ import { IconKind } from '../../../Icon'; import { IconButton } from '../../../IconButton/IconButton'; import { IconButtonVariant } from '../../../IconButton/types'; import { SizedBox } from '../../../SizedBox/SizedBox'; -import { MfaStartMethod } from '../../api/startClientMfaSession'; import { LocationViewHeader } from '../../components/LocationViewHeader/LocationViewHeader'; import { useLocationCardContext } from '../../context/context'; import { LocationCardViews } from '../../context/types'; -import { useMfaConnect } from '../../hooks/useMfaConnect'; +import { MfaStartMethod, useMfaConnect } from '../../hooks/useMfaConnect'; import { LocationCardMfaStartLoader } from '../LocationCardMfaStartLoader/LocationCardMfaStartLoader'; const MIN_POSTURE_LOADER_MS = 500; diff --git a/new-ui/src/shared/components/LocationCard/views/LocationCardMfaTotpView/LocationCardMfaTotpView.tsx b/new-ui/src/shared/components/LocationCard/views/LocationCardMfaTotpView/LocationCardMfaTotpView.tsx index 5a9e46f3a..083373b14 100644 --- a/new-ui/src/shared/components/LocationCard/views/LocationCardMfaTotpView/LocationCardMfaTotpView.tsx +++ b/new-ui/src/shared/components/LocationCard/views/LocationCardMfaTotpView/LocationCardMfaTotpView.tsx @@ -10,11 +10,10 @@ import { IconKind } from '../../../Icon'; import { IconButton } from '../../../IconButton/IconButton'; import { IconButtonVariant } from '../../../IconButton/types'; import { SizedBox } from '../../../SizedBox/SizedBox'; -import { MfaStartMethod } from '../../api/startClientMfaSession'; import { LocationViewHeader } from '../../components/LocationViewHeader/LocationViewHeader'; import { useLocationCardContext } from '../../context/context'; import { LocationCardViews } from '../../context/types'; -import { useMfaConnect } from '../../hooks/useMfaConnect'; +import { MfaStartMethod, useMfaConnect } from '../../hooks/useMfaConnect'; import { LocationCardMfaStartLoader } from '../LocationCardMfaStartLoader/LocationCardMfaStartLoader'; const MIN_POSTURE_LOADER_MS = 500; diff --git a/new-ui/src/shared/edge-api/api.ts b/new-ui/src/shared/edge-api/api.ts index b9d4769d9..24faa6350 100644 --- a/new-ui/src/shared/edge-api/api.ts +++ b/new-ui/src/shared/edge-api/api.ts @@ -1,34 +1,21 @@ -import { getVersion } from '@tauri-apps/api/app'; import { invoke } from '@tauri-apps/api/core'; -import { fetch } from '@tauri-apps/plugin-http'; +import { api } from '../rust-api/api'; import type { CreateDeviceResponse, InstanceInfo, - MfaMethodValue, SaveDeviceConfigResponse, } from '../rust-api/types'; import { TauriCommand } from '../rust-api/types'; import { generateWGKeys } from '../utils/generateWGKeys'; -import { mfaToApi } from '../utils/mfa'; import type { - ActivateUserRequest, - ActivateUserResponse, AddInstanceRequest, AddInstanceResult, - EdgeRequestHeaders, EnrollmentErrorKind, - EnrollmentStartResponse, - MfaSetupFinishRequest, - MfaSetupFinishResponse, - MfaSetupStartResponse, UpdateInstanceRequest, UpdateInstanceResult, } from './types'; -const getPlatformHeader = (): Promise => invoke(TauriCommand.GetPlatformHeader); const getInstances = (): Promise => invoke(TauriCommand.AllInstances); -const deleteInstance = (instanceId: number): Promise => - invoke(TauriCommand.DeleteInstance, { instanceId }); const updateInstanceRecord = (args: { instanceId: number; response: CreateDeviceResponse; @@ -38,271 +25,164 @@ const saveDeviceConfig = (args: { response: CreateDeviceResponse; }): Promise => invoke(TauriCommand.SaveDeviceConfig, args); -const buildProxyUrl = (url: string): string => { - const base = url.endsWith('/') ? url.slice(0, -1) : url; - return `${base}/api/v1`; +/// Extract the raw Rust error string from a Tauri command rejection. +/// Tauri 2 command errors may be plain objects (message on `message`) or the +/// returned `Err(String)` directly. +const rustErrorMessage = (err: unknown): string => + typeof err === 'object' && err !== null && 'message' in err + ? String((err as Record).message) + : String(err); + +/// Parse a Tauri command error that contains a serialized `EnrollmentError` +/// JSON string into an `EnrollmentErrorKind` and human-readable message. +const parseEnrollmentError = ( + err: unknown, +): { error?: string; errorKind: EnrollmentErrorKind } => { + const raw = rustErrorMessage(err); + try { + const parsed = JSON.parse(raw) as { type: string; message?: string; status?: number }; + switch (parsed.type) { + case 'token_expired': + return { errorKind: 'unauthorized' }; + case 'network_error': + return { errorKind: 'network' }; + case 'proxy_error': + return { error: parsed.message, errorKind: 'server' }; + default: + return { error: parsed.message ?? raw, errorKind: 'server' }; + } + } catch { + return { error: raw, errorKind: 'server' }; + } }; -const getEdgeRequestHeaders = async (): Promise => { - const platform = await getPlatformHeader(); - const version = await getVersion().catch(() => 'unknown'); - return { - 'defguard-client-platform': platform, - 'defguard-client-version': version, - }; +/// True when a command error is a proxy 404 - the device was deleted +/// server-side while a stale local record survived. +const isDeviceNotFound = (err: unknown): boolean => { + try { + const parsed = JSON.parse(rustErrorMessage(err)) as { + type?: string; + status?: number; + }; + return parsed.type === 'proxy_error' && parsed.status === 404; + } catch { + return false; + } }; const createDevice = async ( - proxyUrl: string, - cookie: string, + sessionId: string, name: string, ): Promise<{ error?: string }> => { - const edgeHeaders = await getEdgeRequestHeaders(); - const { publicKey, privateKey } = generateWGKeys(); - const url = buildProxyUrl(proxyUrl); - const deviceRes = await fetch(`${url}/enrollment/create_device`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - Cookie: cookie, - ...edgeHeaders, - }, - body: JSON.stringify({ name, pubkey: publicKey }), - }); - - if (!deviceRes.ok) { - const body = (await deviceRes.json()) as { error?: string }; - return { error: body.error ?? `create_device failed (${deviceRes.status})` }; - } - - await saveDeviceConfig({ privateKey, response: await deviceRes.json() }); - return {}; -}; - -type EnrollmentStartOutcome = - | { ok: true; response: Response } - | { ok: false; error?: string; errorKind: EnrollmentErrorKind }; - -const startEnrollment = async ( - proxyUrl: string, - token: string, - edgeHeaders: EdgeRequestHeaders, -): Promise => { - let res: Response; try { - res = await fetch(`${proxyUrl}/enrollment/start`, { - method: 'POST', - headers: { 'Content-Type': 'application/json', ...edgeHeaders }, - body: JSON.stringify({ token }), + const { publicKey, privateKey } = generateWGKeys(); + const deviceResponse = await api.enrollmentCreateDevice(sessionId, name, publicKey); + await saveDeviceConfig({ + privateKey, + response: deviceResponse as CreateDeviceResponse, }); - } catch { - return { ok: false, errorKind: 'network' }; - } - - if (!res.ok) { - if (res.status === 401) { - return { ok: false, errorKind: 'unauthorized' }; - } - const body = (await res.json().catch(() => ({}))) as { error?: string }; - return { - ok: false, - error: body.error ?? `Enrollment start failed (${res.status})`, - errorKind: 'server', - }; + return {}; + } catch (e) { + return { error: e instanceof Error ? e.message : String(e) }; } - - return { ok: true, response: res }; }; const addInstance = async (values: AddInstanceRequest): Promise => { + let sessionId: string | undefined; + // The successful new-enrollment path hands the session to the enrollment + // wizard, which owns its lifecycle from there. Every other exit (existing + // instance, name clash, error) must release the session it created. + let handOffSession = false; try { - const proxyUrl = buildProxyUrl(values.url); - - const edgeHeaders = await getEdgeRequestHeaders(); - - const startResult = await startEnrollment(proxyUrl, values.token, edgeHeaders); - if (!startResult.ok) { - return { error: startResult.error, errorKind: startResult.errorKind }; - } - const startRes = startResult.response; - - const cookie = startRes.headers - .getSetCookie() - .find((c) => c.startsWith('defguard_proxy=')); - if (!cookie) return { error: 'Auth cookie missing from enrollment response' }; - - const resp = (await startRes.json()) as EnrollmentStartResponse; + const startResult = await api.enrollmentStart(values.url.trim(), values.token.trim()); + sessionId = startResult.session_id; const instances = await getInstances(); - const existing = instances.find((i) => i.uuid === resp.instance.id); + const existing = instances.find((i) => i.uuid === startResult.instance.id); if (existing) { - const netRes = await fetch(`${proxyUrl}/enrollment/network_info`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - Cookie: cookie, - ...edgeHeaders, - }, - body: JSON.stringify({ pubkey: existing.pubkey }), - }); - // device no longer exists core side, clean it up - if (netRes.status === 404) { - await deleteInstance(existing.id); - } else { - if (!netRes.ok) return { error: `network_info failed (${netRes.status})` }; + try { + const netInfo = await api.enrollmentNetworkInfo( + startResult.session_id, + existing.pubkey, + ); await updateInstanceRecord({ instanceId: existing.id, - response: await netRes.json(), + response: netInfo as CreateDeviceResponse, }); return {}; + } catch (e) { + // Device was deleted server-side but the local record survived: drop + // the stale record and fall through to a fresh enrollment. + if (!isDeviceNotFound(e)) throw e; + await api.deleteInstance(existing.id); } } const normalizedName = values.name.trim().toLowerCase(); - if (resp.user.device_names.some((n) => n.trim().toLowerCase() === normalizedName)) { + if ( + startResult.user.device_names.some((n) => n.trim().toLowerCase() === normalizedName) + ) { return { error: `Device name '${values.name}' is already in use` }; } - return { startResponse: resp, proxyUrl, cookie }; + handOffSession = true; + return { startResponse: startResult, session_id: startResult.session_id }; } catch (e) { - return { error: e instanceof Error ? e.message : String(e) }; + const parsed = parseEnrollmentError(e); + return { error: parsed.error, errorKind: parsed.errorKind }; + } finally { + // Best-effort cleanup; enrollment_finish errors if the session is already + // gone, so swallow failures rather than mask the real result. + if (sessionId && !handOffSession) { + await api.enrollmentFinish(sessionId).catch(() => {}); + } } }; const updateExistingInstance = async ( values: UpdateInstanceRequest, ): Promise => { + // The update flow only needs the session for the network_info call; it never + // hands it to the wizard, so every exit after start() must release it. + let sessionId: string | undefined; try { - const proxyUrl = buildProxyUrl(values.url); - const edgeHeaders = await getEdgeRequestHeaders(); - const instances = await getInstances(); const existing = instances.find((i) => i.id === values.instanceId); if (!existing) return { error: 'Instance no longer exists.' }; - const startResult = await startEnrollment(proxyUrl, values.token, edgeHeaders); - if (!startResult.ok) { - return { error: startResult.error, errorKind: startResult.errorKind }; - } - const startRes = startResult.response; - - const cookie = startRes.headers - .getSetCookie() - .find((c) => c.startsWith('defguard_proxy=')); - if (!cookie) return { error: 'Auth cookie missing from enrollment response' }; + const startResult = await api.enrollmentStart(values.url, values.token); + sessionId = startResult.session_id; - const resp = (await startRes.json()) as EnrollmentStartResponse; - - if (resp.instance.id !== existing.uuid) { + if (startResult.instance.id !== existing.uuid) { return { error: 'Provided token belongs to a different instance.', errorKind: 'unauthorized', }; } - const netRes = await fetch(`${proxyUrl}/enrollment/network_info`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - Cookie: cookie, - ...edgeHeaders, - }, - body: JSON.stringify({ pubkey: existing.pubkey }), - }); - - if (!netRes.ok) { - const body = (await netRes.json()) as { error?: string }; - return { error: body.error ?? `network_info failed (${netRes.status})` }; - } - + const netInfo = await api.enrollmentNetworkInfo( + startResult.session_id, + existing.pubkey, + ); await updateInstanceRecord({ instanceId: existing.id, - response: await netRes.json(), + response: netInfo as CreateDeviceResponse, }); return {}; } catch (e) { - return { error: e instanceof Error ? e.message : String(e) }; - } -}; - -const startMfaSetup = async ( - proxyUrl: string, - cookie: string, - method: MfaMethodValue, -): Promise<{ result?: MfaSetupStartResponse; error?: string }> => { - try { - const base = buildProxyUrl(proxyUrl); - const edgeHeaders = await getEdgeRequestHeaders(); - const res = await fetch(`${base}/enrollment/register-mfa/code/start`, { - method: 'POST', - headers: { 'Content-Type': 'application/json', Cookie: cookie, ...edgeHeaders }, - body: JSON.stringify({ method: mfaToApi(method) }), - }); - if (!res.ok) { - const body = (await res.json()) as { error?: string }; - return { error: body.error ?? `MFA setup start failed (${res.status})` }; + const parsed = parseEnrollmentError(e); + return { error: parsed.error, errorKind: parsed.errorKind }; + } finally { + // Best-effort cleanup; guarded so we never finish a session that was never + // created (e.g. the instance-missing early return, or a start() failure). + if (sessionId) { + await api.enrollmentFinish(sessionId).catch(() => {}); } - return { result: (await res.json()) as MfaSetupStartResponse }; - } catch (e) { - return { error: e instanceof Error ? e.message : String(e) }; - } -}; - -const activateUser = async ( - proxyUrl: string, - cookie: string, - request: Omit, -): Promise<{ result?: ActivateUserResponse; error?: string }> => { - try { - const base = buildProxyUrl(proxyUrl); - const edgeHeaders = await getEdgeRequestHeaders(); - const res = await fetch(`${base}/enrollment/activate_user`, { - method: 'POST', - headers: { 'Content-Type': 'application/json', Cookie: cookie, ...edgeHeaders }, - body: JSON.stringify({ ...request }), - }); - if (!res.ok) { - const body = (await res.json()) as { error?: string }; - return { error: body.error ?? `activate_user failed (${res.status})` }; - } - return { result: (await res.json()) as ActivateUserResponse }; - } catch (e) { - return { error: e instanceof Error ? e.message : String(e) }; - } -}; - -const finishMfaSetup = async ( - proxyUrl: string, - cookie: string, - request: MfaSetupFinishRequest, -): Promise<{ result?: MfaSetupFinishResponse; error?: string }> => { - try { - const base = buildProxyUrl(proxyUrl); - const edgeHeaders = await getEdgeRequestHeaders(); - const res = await fetch(`${base}/enrollment/register-mfa/code/finish`, { - method: 'POST', - headers: { 'Content-Type': 'application/json', Cookie: cookie, ...edgeHeaders }, - body: JSON.stringify({ - code: request.code, - method: mfaToApi(request.method), - }), - }); - if (!res.ok) { - const body = (await res.json()) as { error?: string }; - return { error: body.error ?? `MFA setup finish failed (${res.status})` }; - } - return { result: (await res.json()) as MfaSetupFinishResponse }; - } catch (e) { - return { error: e instanceof Error ? e.message : String(e) }; } }; export const edgeApi = { - getEdgeRequestHeaders, createDevice, addInstance, updateInstance: updateExistingInstance, - startMfaSetup, - activateUser, - finishMfaSetup, }; diff --git a/new-ui/src/shared/edge-api/types.ts b/new-ui/src/shared/edge-api/types.ts index a4e050941..41ffde631 100644 --- a/new-ui/src/shared/edge-api/types.ts +++ b/new-ui/src/shared/edge-api/types.ts @@ -1,54 +1,4 @@ -import type { MfaMethodValue } from '../rust-api/types'; - -export type EnrollmentInstanceInfo = { - id: string; - name: string; - url: string; - proxy_url?: string; - username: string; - openid_display_name?: string; -}; - -export type AdminInfo = { - name: string; - email: string; - phone_number?: string; -}; - -export type UserInfo = { - first_name: string; - last_name: string; - login: string; - email: string; - is_active: boolean; - phone_number: string; - device_names: string[]; - enrolled: boolean; - password_management_disabled: boolean; -}; - -export type EnrollmentSettings = { - admin_device_management: boolean; - mfa_required: boolean; - only_client_activation: boolean; - smtp_configured: boolean; - vpn_setup_optional: boolean; -}; - -export type EnrollmentStartResponse = { - admin: AdminInfo; - user: UserInfo; - instance: EnrollmentInstanceInfo; - deadline_timestamp: number; - final_page_content: string; - vpn_setup_optional: boolean; - settings: EnrollmentSettings; -}; - -export type EdgeRequestHeaders = { - 'defguard-client-version': string; - 'defguard-client-platform': string; -}; +import type { EnrollmentStartResult } from '../rust-api/types'; /** `network`: the request could not be sent, most likely a bad URL. * `unauthorized`: the server responded 401, the token is invalid. @@ -57,9 +7,8 @@ export type EnrollmentErrorKind = 'network' | 'unauthorized' | 'server'; export type AddInstanceRequest = { url: string; token: string; name: string }; export type AddInstanceResult = { - startResponse?: EnrollmentStartResponse; - proxyUrl?: string; - cookie?: string; + startResponse?: EnrollmentStartResult; + session_id?: string; error?: string; errorKind?: EnrollmentErrorKind; }; @@ -69,18 +18,3 @@ export type UpdateInstanceResult = { error?: string; errorKind?: EnrollmentErrorKind; }; - -export type MfaSetupStartRequest = { method: MfaMethodValue }; -export type MfaSetupStartResponse = { totp_secret?: string }; - -export type MfaSetupFinishRequest = { code: string; method: MfaMethodValue }; -export type MfaSetupFinishResponse = { recovery_codes: string[] }; - -export type ActivateUserRequest = { - password?: string; - phone_number: string; -}; - -export type ActivateUserResponse = { - token: string; -}; diff --git a/new-ui/src/shared/rust-api/api.ts b/new-ui/src/shared/rust-api/api.ts index 9d2ca9373..88b1b8b94 100644 --- a/new-ui/src/shared/rust-api/api.ts +++ b/new-ui/src/shared/rust-api/api.ts @@ -1,15 +1,21 @@ import { invoke } from '@tauri-apps/api/core'; +import { mfaToApi } from '../utils/mfa'; import type { ActiveConnectionSummary, AppConfig, AppConfigPatch, Connection, ConnectionArgs, + EnrollmentMfaFinishResult, + EnrollmentMfaStartResult, + EnrollmentStartResult, InstanceInfo, LocationDetails, LocationDetailsArgs, LocationInfo, LocationStats, + MfaMethodValue, + MfaStartResult, NewAppVersionInfo, ProvisioningConfig, RoutingArgs, @@ -131,6 +137,90 @@ const getSessionState = (): Promise => invoke(TauriCommand.GetSess const patchSessionState = (patch: SessionStatePatch): Promise => invoke(TauriCommand.PatchSessionState, { patch }); +// Enrollment + +const enrollmentStart = ( + proxyUrl: string, + token: string, +): Promise => + invoke(TauriCommand.EnrollmentStart, { proxyUrl, token }); + +const enrollmentCreateDevice = ( + sessionId: string, + name: string, + pubkey: string, +): Promise => + invoke(TauriCommand.EnrollmentCreateDevice, { sessionId, name, pubkey }); + +const enrollmentActivateUser = ( + sessionId: string, + password?: string | null, + phoneNumber?: string | null, +): Promise => + invoke(TauriCommand.EnrollmentActivateUser, { sessionId, password, phoneNumber }); + +const enrollmentRegisterMfaStart = ( + sessionId: string, + method: MfaMethodValue, +): Promise => + invoke(TauriCommand.EnrollmentRegisterMfaStart, { + sessionId, + method: mfaToApi(method), + }); + +const enrollmentRegisterMfaFinish = ( + sessionId: string, + code: string, + method: MfaMethodValue, +): Promise => + invoke(TauriCommand.EnrollmentRegisterMfaFinish, { + sessionId, + code, + method: mfaToApi(method), + }); + +const enrollmentNetworkInfo = (sessionId: string, pubkey: string): Promise => + invoke(TauriCommand.EnrollmentNetworkInfo, { sessionId, pubkey }); + +const enrollmentFinish = (sessionId: string): Promise => + invoke(TauriCommand.EnrollmentFinish, { sessionId }); + +// MFA (connect-time) + +const mfaStart = ( + instanceId: number, + locationId: number, + method: string, +): Promise => + invoke(TauriCommand.MfaStart, { instanceId, locationId, method }); + +// Completes MFA and brings up the connection in the backend; the preshared key +// never crosses back to the frontend. +const mfaFinishCode = ( + instanceId: number, + locationId: number, + token: string, + code: string, +): Promise => + invoke(TauriCommand.MfaFinishCode, { instanceId, locationId, token, code }); + +const mfaPollOpenId = ( + instanceId: number, + locationId: number, + token: string, +): Promise => + invoke(TauriCommand.MfaPollOpenId, { instanceId, locationId, token }); + +const mfaConnectMobileApprove = ( + instanceId: number, + locationId: number, + token: string, +): Promise => + invoke(TauriCommand.MfaConnectMobileApprove, { instanceId, locationId, token }); + +const cancelMfa = (taskId: string): Promise => + invoke(TauriCommand.CancelMfa, { taskId }); + export const api = { // Instances getInstances, @@ -177,4 +267,18 @@ export const api = { // Session state getSessionState, patchSessionState, + // Enrollment + enrollmentStart, + enrollmentCreateDevice, + enrollmentActivateUser, + enrollmentRegisterMfaStart, + enrollmentRegisterMfaFinish, + enrollmentNetworkInfo, + enrollmentFinish, + // MFA + mfaStart, + mfaFinishCode, + mfaPollOpenId, + mfaConnectMobileApprove, + cancelMfa, }; diff --git a/new-ui/src/shared/rust-api/mfaError.ts b/new-ui/src/shared/rust-api/mfaError.ts new file mode 100644 index 000000000..67d79b148 --- /dev/null +++ b/new-ui/src/shared/rust-api/mfaError.ts @@ -0,0 +1,47 @@ +import type { LocationInfo } from './types'; + +/** Shape of the tagged `MfaError` the Rust backend serializes to JSON. */ +export type ParsedMfaError = { + type: string; + message?: string; + status?: number; +}; + +/** Parse a structured `MfaError` (JSON) thrown by a command or carried on an + * event payload. Returns null for plain-string errors. */ +export const parseMfaError = (err: unknown): ParsedMfaError | null => { + try { + const parsed = JSON.parse(String(err)) as ParsedMfaError; + return parsed && typeof parsed.type === 'string' ? parsed : null; + } catch { + return null; + } +}; + +/** Best-effort human-readable message: the structured `message` when present, + * otherwise the raw error string. */ +export const mfaErrorMessage = (err: unknown): string => + parseMfaError(err)?.message ?? String(err); + +/** True when the error is a posture rejection for a posture-gated location. + * The backend maps only HTTP 403 (a failed device posture check) to + * `posture_rejected`; ordinary MFA rejections stay `mfa_rejected`. */ +export const isMfaPostureError = (err: unknown, location: LocationInfo): boolean => + location.posture_check_required && parseMfaError(err)?.type === 'posture_rejected'; + +/** The proxy session/token is no longer valid. */ +export const isSessionExpired = (message: string): boolean => + message.includes('invalid token') || message.includes('login session not found'); + +/** The MFA operation timed out (the backend poll deadline was reached). */ +export const isTimeout = (err: unknown): boolean => + parseMfaError(err)?.type === 'timeout'; + +/** A submitted one-time code was rejected. */ +export const isInvalidCode = (message: string): boolean => + message.includes('Unauthorized'); + +/** MFA succeeded but bringing up the VPN connection afterwards failed + * (see `connect_after_mfa` in the Rust backend). */ +export const isConnectFailure = (message: string): boolean => + message.includes('VPN connection failed'); diff --git a/new-ui/src/shared/rust-api/types.ts b/new-ui/src/shared/rust-api/types.ts index 7ffb204b5..2feacf68b 100644 --- a/new-ui/src/shared/rust-api/types.ts +++ b/new-ui/src/shared/rust-api/types.ts @@ -81,6 +81,20 @@ export type ConnectionType = (typeof ConnectionType)[keyof typeof ConnectionType /** Typed enum for every Tauri command available on the backend. */ export const TauriCommand = { + // Enrollment + EnrollmentStart: 'enrollment_start', + EnrollmentCreateDevice: 'enrollment_create_device', + EnrollmentActivateUser: 'enrollment_activate_user', + EnrollmentRegisterMfaStart: 'enrollment_register_mfa_start', + EnrollmentRegisterMfaFinish: 'enrollment_register_mfa_finish', + EnrollmentNetworkInfo: 'enrollment_network_info', + EnrollmentFinish: 'enrollment_finish', + // MFA + MfaStart: 'mfa_start', + MfaFinishCode: 'mfa_finish_code', + MfaPollOpenId: 'mfa_poll_openid', + MfaConnectMobileApprove: 'mfa_connect_mobile_approve', + CancelMfa: 'cancel_mfa', // Instances AllInstances: 'all_instances', DeleteInstance: 'delete_instance', @@ -147,6 +161,10 @@ export const TauriEvent = { GlobalLogUpdate: 'log-update-global', WindowSwapped: 'window-swapped', SessionStateChanged: 'session-state-changed', + MfaOpenIdComplete: 'mfa-openid-complete', + MfaOpenIdError: 'mfa-openid-error', + MfaMobileComplete: 'mfa-mobile-complete', + MfaMobileError: 'mfa-mobile-error', } as const; export type TauriEventValue = (typeof TauriEvent)[keyof typeof TauriEvent]; @@ -319,7 +337,6 @@ export type SaveDeviceConfigResponse = { export type ConnectionArgs = { locationId: number; connectionType: ConnectionType; - presharedKey?: string; }; export type RoutingArgs = { @@ -389,3 +406,84 @@ export type SessionState = { }; export type SessionStatePatch = Partial; + +/** User information returned by enrollment_start. + * + * Mirrors `InitialUserInfo` from `client_types.proto`. */ +export type EnrollmentUserInfo = { + first_name: string; + last_name: string; + login: string; + email: string; + phone_number: string | null; + is_active: boolean; + device_names: string[]; + enrolled: boolean; + is_admin: boolean; + password_management_disabled: boolean; +}; + +/** Administrator information returned by enrollment_start. + * + * Mirrors `AdminInfo` from `client_types.proto`. */ +export type EnrollmentAdminInfo = { + name: string; + phone_number: string | null; + email: string; +}; + +/** Enrollment settings returned by enrollment_start. + * + * Mirrors `EnrollmentSettings` from `client_types.proto`. */ +export type EnrollmentSettings = { + vpn_setup_optional: boolean; + only_client_activation: boolean; + admin_device_management: boolean; + smtp_configured: boolean; + mfa_required: boolean; +}; + +/** Instance info returned by enrollment_start. + * + * Mirrors `InstanceInfo` from `client_types.proto`. */ +export type EnrollmentInstanceInfo = { + id: string; + name: string; + url: string; + proxy_url: string; + username: string; + enterprise_enabled: boolean; + openid_display_name: string | null; +}; + +/** Full result from the enrollment_start Tauri command. */ +export type EnrollmentStartResult = { + session_id: string; + user: EnrollmentUserInfo; + admin: EnrollmentAdminInfo; + settings: EnrollmentSettings; + instance: EnrollmentInstanceInfo; + deadline_timestamp: number; + final_page_content: string; +}; + +/** Result from enrollment_register_mfa_start. */ +export type EnrollmentMfaStartResult = { + totp_secret: string | null; +}; + +/** Result from enrollment_register_mfa_finish. */ +export type EnrollmentMfaFinishResult = { + recovery_codes: string[]; +}; + +/** Result from mfa_start Tauri command. */ +export type MfaStartResult = { + token: string; + challenge: string | null; +}; + +/** Payload for mfa-openid-error / mfa-mobile-error events. */ +export type MfaErrorPayload = { + error: string; +}; diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index c95350497..de82c8a42 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -272,6 +272,16 @@ dependencies = [ "winnow 0.7.15", ] +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-broadcast" version = "0.7.2" @@ -1582,6 +1592,24 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + [[package]] name = "defguard-cli" version = "2.1.0" @@ -1610,9 +1638,11 @@ dependencies = [ "tokio", "tokio-stream", "tokio-tungstenite", + "tokio-util", "tonic", "tracing", "tracing-subscriber", + "url", "webbrowser", ] @@ -1685,6 +1715,7 @@ dependencies = [ "tracing", "tracing-appender", "tracing-subscriber", + "uuid", "vergen-git2", "webbrowser", "windows 0.62.2", @@ -1732,6 +1763,7 @@ dependencies = [ "defguard-client-proto", "defguard_wireguard_rs", "dirs-next", + "futures-util", "hex", "hyper-util", "log", @@ -1754,10 +1786,13 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokio", + "tokio-tungstenite", + "tokio-util", "tonic", "tower", "tracing", "windows-sys 0.61.2", + "wiremock", "x25519-dalek", ] @@ -4411,6 +4446,16 @@ dependencies = [ "libm", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "num_enum" version = "0.7.6" @@ -9704,6 +9749,29 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "wiremock" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" +dependencies = [ + "assert-json-diff", + "base64 0.22.1", + "deadpool", + "futures", + "http", + "http-body-util", + "hyper", + "hyper-util", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", + "url", +] + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index beb94b9cc..49b25af55 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -38,7 +38,10 @@ http = "1" image = "0.25" qrcode = { version = "0.14", features = ["image"] } tokio-tungstenite = { version = "0.29", features = ["native-tls"] } +tokio-util = "0.7" +url = "2" vergen-git2 = { version = "9.1", features = ["build"] } +wiremock = "0.6" [workspace.package] authors = ["Defguard"] @@ -126,6 +129,7 @@ time = { version = "0.3", features = ["formatting", "macros"] } tokio.workspace = true tokio-util = "0.7" tonic.workspace = true +uuid = { version = "1", features = ["v4"] } tonic-prost.workspace = true tower = "0.5" tracing.workspace = true diff --git a/src-tauri/client-cli/Cargo.toml b/src-tauri/client-cli/Cargo.toml index 6b5a46cc6..660c38179 100644 --- a/src-tauri/client-cli/Cargo.toml +++ b/src-tauri/client-cli/Cargo.toml @@ -29,8 +29,10 @@ serde_json.workspace = true sqlx.workspace = true thiserror.workspace = true tokio.workspace = true +tokio-util.workspace = true tracing.workspace = true tracing-subscriber = { workspace = true, features = ["env-filter"] } +url.workspace = true webbrowser.workspace = true http.workspace = true qrcode.workspace = true diff --git a/src-tauri/client-cli/src/main.rs b/src-tauri/client-cli/src/main.rs index edf34dd80..5b4eab6c3 100644 --- a/src-tauri/client-cli/src/main.rs +++ b/src-tauri/client-cli/src/main.rs @@ -18,8 +18,6 @@ mod resolve; mod state; #[cfg(all(test, target_os = "linux"))] mod tests_daemon; -#[cfg(test)] -mod tests_proxy; use cli::{Cli, InstanceCommand, LocationCommand, TunnelCommand}; diff --git a/src-tauri/client-cli/src/mfa.rs b/src-tauri/client-cli/src/mfa.rs index dcc8d638f..66a3281b8 100644 --- a/src-tauri/client-cli/src/mfa.rs +++ b/src-tauri/client-cli/src/mfa.rs @@ -1,15 +1,13 @@ -//! Connect-time VPN MFA over `core::proxy` (HTTP). +//! Connect-time VPN MFA thin wrapper over `defguard_core::mfa`. //! //! Supports TOTP, email, OIDC, and mobile-approve methods. - -use std::time::Duration; +//! +//! CLI-specific code (method resolution from flags, browser-open, QR +//! rendering, TTY prompting) stays here; all HTTP, WebSocket, and poll +//! logic delegates to `defguard_core::mfa`. use defguard_client_proto::defguard::{ - client_types::{ - ClientMfaFinishRequest, ClientMfaFinishResponse, ClientMfaStartRequest, - ClientMfaStartResponse, MfaMethod, - }, - enterprise::posture::v2::DevicePostureData, + client_types::MfaMethod, enterprise::posture::v2::DevicePostureData, }; use defguard_core::{ database::{ @@ -21,21 +19,11 @@ use defguard_core::{ }, DbPool, }, - proxy::post_with_headers, + mfa, }; -use futures_util::StreamExt; -use reqwest::{StatusCode, Url}; use secrecy::{ExposeSecret, SecretString}; -use serde::Deserialize; -use serde_json::Value; -use tokio::{ - net::TcpStream, - select, - signal::ctrl_c, - time::{sleep, Instant}, -}; -use tokio_tungstenite::{connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream}; use tracing::{debug, info, warn}; +use url::Url; use crate::{ mfa_code::{obtain_code, CodeSource, MfaContext}, @@ -43,6 +31,20 @@ use crate::{ state::CliError, }; +/// Convert a `defguard_core::mfa::MfaError` into a [`CliError`]. +fn into_cli(err: mfa::MfaError) -> CliError { + let msg = err.to_string(); + match err { + mfa::MfaError::NetworkError { .. } => CliError::Other(msg), + mfa::MfaError::ProxyError { .. } => CliError::Other(msg), + mfa::MfaError::MfaRejected { .. } => CliError::MfaFailed(msg), + mfa::MfaError::PostureRejected { .. } => CliError::MfaFailed(msg), + mfa::MfaError::Timeout => CliError::MfaFailed(msg), + mfa::MfaError::Cancelled => CliError::Cancelled(msg), + mfa::MfaError::Other { .. } => CliError::Other(msg), + } +} + /// Resolve the effective MFA method for a location. /// /// When `method_override` is `Some`, parses it into [`MfaMethod`]; otherwise @@ -57,7 +59,7 @@ pub(crate) fn resolve_method( let method = if let Some(raw) = method_override { let method = parse_method(raw)?; // OIDC override on an Internal-mode location will be rejected by the - // server. Fail early to give the user a clear error before I/O. + // server. Fail early to give the user a clear error before I/O. if method == MfaMethod::Oidc && location.location_mfa_mode == LocationMfaMode::Internal { return Err(CliError::InvalidInput( "--mfa-method oidc is only valid for locations that use external (OIDC) MFA." @@ -101,19 +103,11 @@ pub(crate) fn validate_mfa_flags( Ok(()) } -/// Run the VPN MFA handshake for a location. +/// Run the VPN MFA handshake for a location (TOTP or email). /// -/// * `location` - the target location. -/// * `source` - how to obtain the code. -/// * `instance` - the instance this location belongs to (for proxy URL + pubkey). -/// * `method` - the resolved [`MfaMethod`] (use [`resolve_method`] to obtain it). -/// **Must not be [`MfaMethod::Oidc`]**; OIDC uses [`authorize_oidc`] instead. -/// * `posture_data` - device posture data; must be provided when the location also -/// requires posture checks. -/// -/// Returns the WireGuard preshared key (as a [`SecretString`]) that must -/// be passed to `bring_up`. -pub async fn authorize( +/// The HTTP calls are handled by `defguard_core::mfa`; this function +/// handles CLI-specific code sourcing (TTY / --code / --code-command). +pub(crate) async fn authorize( location: &Location, source: &CodeSource, instance: &Instance, @@ -122,10 +116,9 @@ pub async fn authorize( pool: &DbPool, ) -> Result { // Reject methods not yet supported by the CLI before doing any I/O. - // OIDC is not rejected as "not yet supported" because it IS supported - - // but it has its own dedicated code path (authorize_oidc). If callers - // accidentally invoke authorize() with OIDC, this catch-all is a - // defense-in-depth barrier that emits a clear error. + // OIDC/MobileApprove are not "unsupported" - they have dedicated code + // paths (authorize_oidc / authorize_mobile_approve). This catch-all is a + // defense-in-depth barrier that emits a clear error if they land here. match method { MfaMethod::Biometric => { return Err(CliError::MfaFailed(format!( @@ -156,194 +149,44 @@ pub async fn authorize( )) })?; - // Parse the proxy URL once and reuse it for both MFA requests. let proxy_url = Url::parse(&instance.proxy_url) .map_err(|e| CliError::Other(format!("Invalid proxy URL: {e}")))?; - - // The one-time MFA code and the returned preshared key are sensitive; warn (but do - // not block) if the proxy is not using HTTPS, since they would travel in cleartext. check_proxy_scheme(&proxy_url); - // Step 1: Start the MFA session. - let start_req = ClientMfaStartRequest { + debug!("Starting MFA session for location {}", location.name); + let request = defguard_client_proto::defguard::client_types::ClientMfaStartRequest { location_id: location.network_id, pubkey: wireguard_keys.pubkey, method: method as i32, posture_data, }; - - let start_url = proxy_url - .join("api/v1/client-mfa/start") - .map_err(|e| CliError::Other(format!("Failed to build MFA start URL: {e}")))?; - - debug!("Starting MFA session at {start_url}"); - let response = post_with_headers(start_url, &start_req) + let info = mfa::mfa_start(proxy_url.clone(), request) .await - .map_err(|e| CliError::Other(format!("Failed to reach proxy: {e}")))?; - - if !response.status().is_success() { - return Err(handle_mfa_error(response).await); - } + .map_err(into_cli)?; - let start_resp: ClientMfaStartResponse = response - .json() - .await - .map_err(|e| CliError::Other(format!("Invalid MFA start response: {e}")))?; - - let token = start_resp.token; - debug!("MFA session started, token obtained"); - - // Step 2: Obtain the code. let ctx = MfaContext { instance: instance.name.clone(), location: location.name.clone(), }; let code = obtain_code(source, &ctx)?; - // Step 3: Finish the MFA session. - let finish_req = ClientMfaFinishRequest { - token, + let finish_req = defguard_client_proto::defguard::client_types::ClientMfaFinishRequest { + token: info.token, code: Some(code.expose_secret().to_string()), auth_pub_key: None, }; - - let finish_url = proxy_url - .join("api/v1/client-mfa/finish") - .map_err(|e| CliError::Other(format!("Failed to build MFA finish URL: {e}")))?; - - debug!("Finishing MFA session at {finish_url}"); - let response = post_with_headers(finish_url, &finish_req) - .await - .map_err(|e| CliError::Other(format!("Failed to reach proxy: {e}")))?; - - if !response.status().is_success() { - return Err(handle_mfa_error(response).await); - } - - let finish_resp: ClientMfaFinishResponse = response - .json() + let psk = mfa::mfa_finish_code(proxy_url, finish_req) .await - .map_err(|e| CliError::Other(format!("Invalid MFA finish response: {e}")))?; + .map_err(into_cli)?; info!("MFA session completed, preshared key obtained"); - Ok(SecretString::from(finish_resp.preshared_key)) -} - -#[derive(Deserialize)] -struct ErrorBody { - error: Option, -} - -/// Map a non-2xx MFA response to a `CliError`. -async fn handle_mfa_error(response: reqwest::Response) -> CliError { - let status = response.status(); - - let message = response - .json::() - .await - .ok() - .and_then(|b| b.error) - .unwrap_or_else(|| format!("HTTP {status}")); - - match status { - StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED => { - CliError::MfaFailed(format!("MFA rejected: {message}")) - } - _ if status.is_client_error() => CliError::MfaFailed(format!("MFA error: {message}")), - _ => CliError::Other(format!("Proxy error (HTTP {status}): {message}")), - } -} - -/// Parse a `--mfa-method` flag string into the proto [`MfaMethod`] enum. -fn parse_method(raw: &str) -> Result { - match raw.to_lowercase().as_str() { - "totp" => Ok(MfaMethod::Totp), - "email" => Ok(MfaMethod::Email), - "oidc" => Ok(MfaMethod::Oidc), - "biometric" => Ok(MfaMethod::Biometric), - "mobile" | "mobile_approve" => Ok(MfaMethod::MobileApprove), - _ => Err(CliError::Usage(format!( - "Invalid --mfa-method '{raw}'. Valid: totp, email, oidc, biometric, mobile." - ))), - } -} - -/// Determine the MFA method to use for a location. -/// -/// Delegates to the core's [`infer_mfa_method`] so that -/// [`LocationMfaMode`] is respected - an External-mode location always -/// uses OIDC, while an Internal-mode location respects the stored -/// preference (defaulting to TOTP when unset). -fn infer_method(location: &Location) -> MfaMethod { - let method = infer_mfa_method(location.location_mfa_mode, location.mfa_method); - match method { - Some(LocationMfaMethod::Totp) => MfaMethod::Totp, - Some(LocationMfaMethod::Email) => MfaMethod::Email, - Some(LocationMfaMethod::Oidc) => MfaMethod::Oidc, - Some(LocationMfaMethod::Biometric) => MfaMethod::Biometric, - Some(LocationMfaMethod::MobileApprove) => MfaMethod::MobileApprove, - None => { - // `infer_mfa_method` only returns None for Disabled mode, but - // this function is only called when `mfa_enabled()` is true. - // Default to TOTP as a safe fallback. - MfaMethod::Totp - } - } + Ok(SecretString::from(psk.preshared_key)) } -/// Warn if the proxy is not using HTTPS. -/// -/// The one-time MFA code and the returned preshared key are sensitive and would -/// travel in cleartext over plain HTTP. -fn check_proxy_scheme(proxy_base: &Url) { - if proxy_base.scheme() != "https" { - warn!( - "Proxy URL '{}' is not HTTPS; secrets will be sent in cleartext.", - proxy_base.as_str() - ); - } -} - -/// Open a URL in the system browser. -/// -/// Production: calls [`webbrowser::open`]; prints a hint to stderr on failure. -/// When `json_mode` is true, the fallback message includes the URL itself since -/// it wasn't already printed above. -/// Tests: no-op (never spawn a browser). -#[cfg(not(test))] -fn open_url(url: &str, json_mode: bool) { - if webbrowser::open(url).is_err() { - if json_mode { - eprintln!("Could not open browser. Open this URL manually: {url}"); - } else { - eprintln!("Could not open browser. Open the URL above manually."); - } - } -} - -#[cfg(test)] -fn open_url(_url: &str, _json_mode: bool) { - // no-op: tests must not spawn a browser -} - -/// Fixed interval between OIDC MFA finish polls (shortened for tests). -#[cfg(not(test))] -const OIDC_POLL_INTERVAL: Duration = Duration::from_secs(5); -#[cfg(test)] -const OIDC_POLL_INTERVAL: Duration = Duration::from_millis(5); - -/// Total time the CLI will wait for the user to complete OIDC authentication -/// before giving up (shortened for tests). -#[cfg(not(test))] -const OIDC_POLL_TIMEOUT: Duration = Duration::from_mins(5); -#[cfg(test)] -const OIDC_POLL_TIMEOUT: Duration = Duration::from_millis(200); - /// Run the OIDC MFA flow for an external-IdP location. /// -/// Opens the system browser to `{proxy_url}openid/mfa?token=...` and polls -/// the proxy until the user completes authentication with the external -/// identity provider. Returns the WireGuard preshared key. +/// Opens the system browser and delegates the HTTP poll to +/// `defguard_core::mfa::poll_openid_mfa`. /// /// When `json_mode` is true, progress messages on stderr are suppressed so /// that `--json` output consumers only see the final result/error. @@ -364,44 +207,27 @@ pub(crate) async fn authorize_oidc( )) })?; - let proxy_base = Url::parse(&instance.proxy_url) + let proxy_url = Url::parse(&instance.proxy_url) .map_err(|e| CliError::Other(format!("Invalid proxy URL: {e}")))?; - check_proxy_scheme(&proxy_base); + check_proxy_scheme(&proxy_url); - // Step 1: Start the OIDC MFA session. - let start_req = ClientMfaStartRequest { + debug!("Starting OIDC MFA session for location {}", location.name); + let request = defguard_client_proto::defguard::client_types::ClientMfaStartRequest { location_id: location.network_id, pubkey: wireguard_keys.pubkey, method: MfaMethod::Oidc as i32, posture_data, }; - - let start_url = proxy_base - .join("api/v1/client-mfa/start") - .map_err(|e| CliError::Other(format!("Failed to build MFA start URL: {e}")))?; - - debug!("Starting OIDC MFA session at {start_url}"); - let response = post_with_headers(start_url, &start_req) + let info = mfa::mfa_start(proxy_url.clone(), request) .await - .map_err(|e| CliError::Other(format!("Failed to reach proxy: {e}")))?; + .map_err(into_cli)?; - if !response.status().is_success() { - return Err(handle_mfa_error(response).await); - } - - let start_resp: ClientMfaStartResponse = response - .json() - .await - .map_err(|e| CliError::Other(format!("Invalid MFA start response: {e}")))?; - - // Step 2: Open the browser for the user to authenticate. - // Never log the token-bearing URL via tracing. - let mut browser_url = proxy_base + let mut browser_url = proxy_url .join("openid/mfa") .map_err(|e| CliError::Other(format!("Failed to build OIDC MFA URL: {e}")))?; browser_url .query_pairs_mut() - .append_pair("token", &start_resp.token); + .append_pair("token", &info.token); if !json_mode { eprintln!("Open this URL to authenticate:"); @@ -410,99 +236,25 @@ pub(crate) async fn authorize_oidc( } open_url(browser_url.as_ref(), json_mode); - // Step 3: Poll until the user completes authentication or the deadline expires. - poll_finish(&proxy_base, &start_resp.token, json_mode).await -} + let cancel = tokio_util::sync::CancellationToken::new(); + let cancel_clone = cancel.clone(); + let ctrlc_handle = tokio::spawn(async move { + tokio::signal::ctrl_c().await.ok(); + cancel_clone.cancel(); + }); -/// Poll [`client-mfa/finish`] until the OIDC session is completed or the -/// deadline expires. -/// -/// Returns the preshared key on success. -async fn poll_finish( - proxy_base: &Url, - token: &str, - json_mode: bool, -) -> Result { - let finish_url = proxy_base - .join("api/v1/client-mfa/finish") - .map_err(|e| CliError::Other(format!("Failed to build MFA finish URL: {e}")))?; - - let finish_req = ClientMfaFinishRequest { - token: token.to_string(), - code: None, - auth_pub_key: None, - }; - - let deadline = Instant::now() + OIDC_POLL_TIMEOUT; - - loop { - // Check if we've exceeded the overall deadline. - let remaining = deadline - .checked_duration_since(Instant::now()) - .unwrap_or_default(); - if remaining.is_zero() { - if !json_mode { - eprintln!("MFA login timed out."); - } - return Err(CliError::MfaFailed("MFA login timed out".into())); - } + let result = mfa::poll_openid_mfa(proxy_url, info.token, cancel).await; + ctrlc_handle.abort(); - // Poll the proxy (first iteration is immediate, subsequent ones wait - // the poll interval at the end of each loop body). - let (status, body) = select! { - _ = ctrl_c() => { - if !json_mode { - eprintln!("MFA login cancelled."); - } - return Err(CliError::Cancelled("MFA login cancelled.".into())); - } - result = post_with_headers(finish_url.clone(), &finish_req) => { - let response = result - .map_err(|e| CliError::Other(format!("Failed to reach proxy: {e}")))?; - let status = response.status(); - if status == StatusCode::OK { - let finish_resp: ClientMfaFinishResponse = response - .json() - .await - .map_err(|e| CliError::Other(format!("Invalid MFA finish response: {e}")))?; - info!("OIDC MFA session completed, preshared key obtained"); - return Ok(SecretString::from(finish_resp.preshared_key)); - } - (status, response) - } - }; - - // Non-OK, non-428: report the error. - if status != StatusCode::PRECONDITION_REQUIRED { - return Err(handle_mfa_error(body).await); - } - - // 428: OIDC not complete yet. Wait the poll interval, yielding to - // Ctrl-C, then loop around for another check. - select! { - _ = ctrl_c() => { - if !json_mode { - eprintln!("MFA login cancelled."); - } - return Err(CliError::Cancelled("MFA login cancelled.".into())); - } - () = sleep(OIDC_POLL_INTERVAL) => {} - } - } + let psk = result.map_err(into_cli)?; + info!("OIDC MFA session completed, preshared key obtained"); + Ok(SecretString::from(psk.preshared_key)) } -/// How long the CLI waits for the user to approve MFA on their mobile device. -#[cfg(not(test))] -const MOBILE_APPROVE_TIMEOUT: Duration = Duration::from_mins(2); -#[cfg(test)] -const MOBILE_APPROVE_TIMEOUT: Duration = Duration::from_secs(5); - /// Run the mobile-approve MFA flow. /// -/// Displays a QR code (terminal and/or `--qr-file` PNG), opens a WebSocket -/// to the proxy, and waits for the mobile app to approve the authentication. -/// The CLI performs no cryptography; the mobile device signs the challenge -/// and the proxy pushes the resulting preshared key back over the WebSocket. +/// Displays a QR code (terminal and/or `--qr-file` PNG) and delegates the +/// WebSocket connection to `defguard_core::mfa::connect_mobile_approve`. /// /// When `json_mode` is true, progress messages on stderr are suppressed so /// that `--json` output consumers only see the final result/error. @@ -524,204 +276,121 @@ pub(crate) async fn authorize_mobile_approve( )) })?; - let proxy_base = Url::parse(&instance.proxy_url) + let proxy_url = Url::parse(&instance.proxy_url) .map_err(|e| CliError::Other(format!("Invalid proxy URL: {e}")))?; - check_proxy_scheme(&proxy_base); + check_proxy_scheme(&proxy_url); - // Step 1: Start the MFA session. - let start_req = ClientMfaStartRequest { + debug!( + "Starting mobile-approve MFA session for location {}", + location.name + ); + let request = defguard_client_proto::defguard::client_types::ClientMfaStartRequest { location_id: location.network_id, pubkey: wireguard_keys.pubkey, method: MfaMethod::MobileApprove as i32, posture_data, }; - - let start_url = proxy_base - .join("api/v1/client-mfa/start") - .map_err(|e| CliError::Other(format!("Failed to build MFA start URL: {e}")))?; - - debug!("Starting mobile-approve MFA session at {start_url}"); - let response = post_with_headers(start_url, &start_req) - .await - .map_err(|e| CliError::Other(format!("Failed to reach proxy: {e}")))?; - - if !response.status().is_success() { - return Err(handle_mobile_approve_start_error(response).await); - } - - let start_resp: ClientMfaStartResponse = response - .json() + let info = mfa::mfa_start(proxy_url.clone(), request) .await - .map_err(|e| CliError::Other(format!("Invalid MFA start response: {e}")))?; + .map_err(into_cli)?; - // The challenge is always present for MobileApprove (core always generates - // a BiometricChallenge for this method). - let challenge = start_resp.challenge.ok_or_else(|| { + let challenge = info.challenge.ok_or_else(|| { CliError::Other("Proxy did not return a challenge for mobile-approve MFA".into()) })?; - // Step 2: Build and render the QR code. - let payload = mfa_qr::build_qr_payload(&start_resp.token, &challenge, &instance.uuid); + let payload = mfa_qr::build_qr_payload(&info.token, &challenge, &instance.uuid); mfa_qr::render_qr(&payload, qr_file, json_mode)?; if !json_mode { eprintln!("Waiting for mobile approval... (Ctrl-C to cancel)"); } - // Step 3: Open a WebSocket and wait for the preshared key. - let ws_url = derive_ws_url(&proxy_base, &start_resp.token)?; - let (ws_stream, _response) = connect_async(&ws_url) - .await - .map_err(|e| CliError::Other(format!("Failed to connect to proxy: {e}")))?; + let ws_url = mfa::derive_ws_url(&proxy_url, &info.token).map_err(into_cli)?; + + let cancel = tokio_util::sync::CancellationToken::new(); + let cancel_clone = cancel.clone(); + let ctrlc_handle = tokio::spawn(async move { + tokio::signal::ctrl_c().await.ok(); + cancel_clone.cancel(); + }); - let psk = wait_for_mfa_success(ws_stream, MOBILE_APPROVE_TIMEOUT, json_mode).await?; + let result = mfa::connect_mobile_approve(&ws_url, cancel).await; + ctrlc_handle.abort(); + let psk = result.map_err(into_cli)?; info!("Mobile-approve MFA completed, preshared key obtained"); - Ok(SecretString::from(psk)) + Ok(SecretString::from(psk.preshared_key)) } -/// Handle a non-2xx response from /start during mobile-approve MFA. -/// -/// Rewraps the cryptic server error "selected MFA method is not available" -/// into actionable guidance telling the user to register a mobile authenticator. -async fn handle_mobile_approve_start_error(response: reqwest::Response) -> CliError { - let status = response.status(); - let error_body: Option = response.json().await.ok(); - let message = error_body - .and_then(|b| b.error) - .unwrap_or_else(|| format!("HTTP {status}")); - - if message.contains("selected MFA method is not available") { - return CliError::MfaFailed( - "No mobile authenticator is registered for your account. \ - Register one in the Defguard mobile app, then retry." - .into(), - ); - } - - if status.is_client_error() { - CliError::MfaFailed(format!("MFA error: {message}")) - } else { - CliError::Other(format!("Proxy error (HTTP {status}): {message}")) +/// Parse a `--mfa-method` flag string into the proto [`MfaMethod`] enum. +fn parse_method(raw: &str) -> Result { + match raw.to_lowercase().as_str() { + "totp" => Ok(MfaMethod::Totp), + "email" => Ok(MfaMethod::Email), + "oidc" => Ok(MfaMethod::Oidc), + "biometric" => Ok(MfaMethod::Biometric), + "mobile" | "mobile_approve" => Ok(MfaMethod::MobileApprove), + _ => Err(CliError::Usage(format!( + "Invalid --mfa-method '{raw}'. Valid: totp, email, oidc, biometric, mobile." + ))), } } -/// Derive the WebSocket URL from the proxy's base URL and the MFA token. +/// Determine the MFA method to use for a location. /// -/// Uses [`Url::join`] so that any path prefix is preserved. -fn derive_ws_url(proxy_base: &Url, token: &str) -> Result { - let mut ws_url = proxy_base - .join("api/v1/client-mfa/remote") - .map_err(|e| CliError::Other(format!("Failed to build WebSocket URL: {e}")))?; - - let ws_scheme = match proxy_base.scheme() { - "https" => "wss", - "http" => "ws", - other => { - return Err(CliError::Other(format!( - "Invalid proxy URL scheme '{other}'; expected http or https" - ))); +/// Delegates to the core's [`infer_mfa_method`] so that [`LocationMfaMode`] +/// is respected - an External-mode location always uses OIDC, while an +/// Internal-mode location respects the stored preference (defaulting to TOTP). +fn infer_method(location: &Location) -> MfaMethod { + let method = infer_mfa_method(location.location_mfa_mode, location.mfa_method); + match method { + Some(LocationMfaMethod::Totp) => MfaMethod::Totp, + Some(LocationMfaMethod::Email) => MfaMethod::Email, + Some(LocationMfaMethod::Oidc) => MfaMethod::Oidc, + Some(LocationMfaMethod::Biometric) => MfaMethod::Biometric, + Some(LocationMfaMethod::MobileApprove) => MfaMethod::MobileApprove, + None => { + // infer_mfa_method only returns None for Disabled mode, but this is + // only called when MFA is enabled. Default to TOTP as a safe fallback. + MfaMethod::Totp } - }; - - ws_url - .set_scheme(ws_scheme) - .map_err(|()| CliError::Other("Failed to set WebSocket URL scheme".into()))?; - ws_url.query_pairs_mut().append_pair("token", token); - - Ok(ws_url.to_string()) + } } -/// Wait on the WebSocket for a single `{"type":"mfa_success","preshared_key":"..."}` -/// text frame, or fail if the deadline expires or the user cancels. +/// Warn if the proxy is not using HTTPS. /// -/// Uses an absolute deadline (not a per-frame-gap timeout) so that stray -/// ping/pong traffic does not silently extend the wait. -async fn wait_for_mfa_success( - ws_stream: WebSocketStream>, - timeout: Duration, - json_mode: bool, -) -> Result { - let (_write, mut read) = ws_stream.split(); - - let deadline = Instant::now() + timeout; - - loop { - let remaining = deadline - .checked_duration_since(Instant::now()) - .unwrap_or_default(); - if remaining.is_zero() { - if !json_mode { - eprintln!("Mobile approval timed out."); - } - return Err(CliError::MfaFailed( - "mobile approval timed out; re-run to get a fresh QR".into(), - )); - } +/// The one-time MFA code and the returned preshared key are sensitive and +/// would travel in cleartext over plain HTTP. +fn check_proxy_scheme(proxy_base: &Url) { + if proxy_base.scheme() != "https" { + warn!( + "Proxy URL '{}' is not HTTPS; secrets will be sent in cleartext.", + proxy_base.as_str() + ); + } +} - let msg = select! { - () = sleep(remaining) => { - if !json_mode { - eprintln!("Mobile approval timed out."); - } - return Err(CliError::MfaFailed( - "mobile approval timed out; re-run to get a fresh QR".into(), - )); - } - _ = ctrl_c() => { - if !json_mode { - eprintln!("MFA cancelled."); - } - return Err(CliError::Cancelled("MFA cancelled.".into())); - } - msg = read.next() => { - match msg { - Some(Ok(msg)) => msg, - Some(Err(_)) => { - // Server closed or errored without sending mfa_success. - if !json_mode { - eprintln!("Mobile approval failed: connection closed by proxy."); - } - return Err(CliError::MfaFailed( - "mobile approval failed: connection closed by proxy".into(), - )); - } - None => { - // Server closed the connection without sending mfa_success. - if !json_mode { - eprintln!("Mobile approval failed: connection closed by proxy."); - } - return Err(CliError::MfaFailed( - "mobile approval failed: connection closed by proxy".into(), - )); - } - } - } - }; - - match msg { - Message::Text(text) => { - if let Ok(parsed) = serde_json::from_str::(&text) { - if parsed.get("type").and_then(|v| v.as_str()) == Some("mfa_success") { - if let Some(key) = parsed["preshared_key"].as_str() { - return Ok(key.to_string()); - } - } - } - // Ignore unrecognised text frames (non-JSON, wrong type, or missing preshared_key). - } - Message::Close(_) => { - if !json_mode { - eprintln!("Mobile approval failed: connection closed by proxy."); - } - return Err(CliError::MfaFailed( - "mobile approval failed: connection closed by proxy".into(), - )); - } - _ => {} // Ignore ping, pong, binary. +/// Open a URL in the system browser. +/// +/// Production: calls [`webbrowser::open`]; prints a hint to stderr on failure. +/// When `json_mode` is true, the fallback message includes the URL itself since +/// it wasn't already printed above. +/// Tests: no-op (never spawn a browser). +#[cfg(not(test))] +fn open_url(url: &str, json_mode: bool) { + if webbrowser::open(url).is_err() { + if json_mode { + eprintln!("Could not open browser. Open this URL manually: {url}"); + } else { + eprintln!("Could not open browser. Open the URL above manually."); } } } +#[cfg(test)] +fn open_url(_url: &str, _json_mode: bool) { + // no-op: tests must not spawn a browser +} + #[cfg(test)] mod tests { use defguard_core::database::models::location::ServiceLocationMode; @@ -808,7 +477,6 @@ mod tests { #[test] fn test_validate_flags_qr_file_only_for_mobile_approve() { - // qr-file on TOTP let err = validate_mfa_flags(MfaMethod::Totp, "office", None, None, Some("qr.png")).unwrap_err(); assert!(matches!(err, CliError::InvalidInput(_))); @@ -829,7 +497,6 @@ mod tests { #[test] fn test_validate_flags_pass_through_totp() { - // TOTP with --code should pass validation. validate_mfa_flags(MfaMethod::Totp, "office", Some("123456"), None, None).unwrap(); } diff --git a/src-tauri/client-cli/src/tests_proxy.rs b/src-tauri/client-cli/src/tests_proxy.rs deleted file mode 100644 index 071d5ada5..000000000 --- a/src-tauri/client-cli/src/tests_proxy.rs +++ /dev/null @@ -1,563 +0,0 @@ -//! Integration tests: mock MFA proxy (HTTP). -//! -//! Each test spawns a tiny tokio-based HTTP server on a random port so -//! `mfa::authorize` can make real HTTP calls against it. The database is -//! seeded via `#[sqlx::test]`. - -use std::{ - io::{ErrorKind, Read, Write}, - net::{SocketAddr, TcpListener, TcpStream}, - sync::{ - atomic::{AtomicBool, AtomicU32, Ordering}, - Arc, - }, - thread::{sleep, spawn, JoinHandle}, - time::Duration, -}; - -use base64::{prelude::BASE64_STANDARD, Engine as _}; -use defguard_client_proto::defguard::client_types::MfaMethod; -use defguard_core::database::{ - models::{ - instance::{ClientTrafficPolicy, Instance}, - location::{Location, LocationMfaMode, ServiceLocationMode}, - Id, NoId, - }, - DbPool, -}; -use secrecy::ExposeSecret; -use serde_json::json; -use sha1::{Digest, Sha1}; - -use crate::{mfa, mfa_code::CodeSource, state::CliError}; - -const READ_TIMEOUT: Duration = Duration::from_secs(5); -const CONNECT_TIMEOUT: Duration = Duration::from_millis(50); -const WAIT_TIMEOUT: Duration = Duration::from_millis(10); - -/// Response template for the mock proxy. -#[derive(Clone)] -struct MockResponse { - status: u16, - body: String, -} - -/// A tiny HTTP server that responds to MFA start/finish/remote requests. -struct MockProxy { - addr: SocketAddr, - shutdown: Arc, - handle: Option>, -} - -impl MockProxy { - /// Mobile-approve MFA: /start returns a token+challenge, /remote upgrades to - /// WebSocket and sends the given preshared key (if `Some`). When `None`, the - /// server closes the connection right after the handshake (used for timeout tests). - fn with_mobile_approve(start_response: MockResponse, preshared_key: Option<&str>) -> Self { - let psk = preshared_key.map(|s| s.to_string()); - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - listener.set_nonblocking(true).unwrap(); - let addr = listener.local_addr().unwrap(); - let shutdown = Arc::new(AtomicBool::new(false)); - let shutdown_thread = shutdown.clone(); - let handle = spawn(move || { - while !shutdown_thread.load(Ordering::Relaxed) { - let mut stream = match listener.accept() { - Ok((stream, _)) => stream, - Err(ref e) if e.kind() == ErrorKind::WouldBlock => { - sleep(WAIT_TIMEOUT); - continue; - } - Err(_) => break, - }; - stream.set_nonblocking(false).ok(); - stream.set_read_timeout(Some(READ_TIMEOUT)).ok(); - - // Read the full request head. - let mut data = Vec::new(); - let mut buf = [0u8; 4096]; - loop { - match stream.read(&mut buf) { - Ok(0) => break, - Ok(n) => { - data.extend_from_slice(&buf[..n]); - if data.windows(4).any(|w| w == b"\r\n\r\n") { - break; - } - } - Err(_) => break, - } - } - - let request = String::from_utf8_lossy(&data); - - if request.contains("/api/v1/client-mfa/remote") { - // WebSocket upgrade. - if let Some(key) = extract_header(&request, "Sec-WebSocket-Key") { - let accept = compute_ws_accept(&key); - let resp = format!( - "HTTP/1.1 101 Switching Protocols\r\n\ - Upgrade: websocket\r\n\ - Connection: Upgrade\r\n\ - Sec-WebSocket-Accept: {accept}\r\n\ - \r\n" - ); - let _ = stream.write_all(resp.as_bytes()); - - // Send the mfa_success text frame if we have a PSK. - if let Some(ref key) = psk { - let payload = json!({ - "type": "mfa_success", - "preshared_key": key, - }); - let payload_str = serde_json::to_string(&payload).unwrap(); - let frame = build_ws_text_frame(&payload_str); - let _ = stream.write_all(&frame); - } - // When psk is None, just close (no frame) to simulate timeout. - } - } else if request.contains("/api/v1/client-mfa/start") { - let body = format!( - "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\n\ - Content-Length: {}\r\nConnection: close\r\n\r\n{}", - start_response.status, - start_response.body.len(), - start_response.body, - ); - let _ = stream.write_all(body.as_bytes()); - } - // Unknown paths: no response, just close. - } - }); - MockProxy { - addr, - shutdown, - handle: Some(handle), - } - } - - /// Single-shot finish: each finish request gets the same response. - fn new(start_response: MockResponse, finish_response: MockResponse) -> Self { - Self::with_counter(start_response, finish_response, None) - } - - /// OIDC polling: the first `fail_count` calls to `finish` return HTTP 428; - /// subsequent calls return `finish_response` (HTTP 200). - fn with_poll( - start_response: MockResponse, - finish_response: MockResponse, - fail_count: u32, - ) -> Self { - Self::with_counter(start_response, finish_response, Some(fail_count)) - } - - fn with_counter( - start_response: MockResponse, - finish_response: MockResponse, - fail_count: Option, - ) -> Self { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - // Non-blocking accept so the loop can observe the shutdown flag and exit cleanly. - listener.set_nonblocking(true).unwrap(); - let addr = listener.local_addr().unwrap(); - let shutdown = Arc::new(AtomicBool::new(false)); - let shutdown_thread = shutdown.clone(); - let finish_hits = Arc::new(AtomicU32::new(0)); - let finish_hits_thread = finish_hits.clone(); - let handle = spawn(move || { - while !shutdown_thread.load(Ordering::Relaxed) { - let mut stream = match listener.accept() { - Ok((stream, _)) => stream, - Err(ref e) if e.kind() == ErrorKind::WouldBlock => { - sleep(WAIT_TIMEOUT); - continue; - } - Err(_) => break, - }; - stream.set_nonblocking(false).ok(); - stream.set_read_timeout(Some(READ_TIMEOUT)).ok(); - - // Read the full request head rather than assuming one read captures it. - let mut data = Vec::new(); - let mut buf = [0u8; 4096]; - loop { - match stream.read(&mut buf) { - Ok(0) => break, - Ok(n) => { - data.extend_from_slice(&buf[..n]); - if data.windows(4).any(|w| w == b"\r\n\r\n") { - break; - } - } - Err(_) => break, - } - } - - let request = String::from_utf8_lossy(&data); - let response = if request.contains("/api/v1/client-mfa/start") { - &start_response - } else if let Some(max) = fail_count { - let current = finish_hits_thread.fetch_add(1, Ordering::Relaxed); - if current < max { - // 428: not yet authenticated - &MockResponse { - status: 428, - body: r#"{"error":"OIDC authentication not completed yet"}"#.into(), - } - } else { - &finish_response - } - } else { - &finish_response - }; - let body = format!( - "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", - response.status, - response.body.len(), - response.body, - ); - let _ = stream.write_all(body.as_bytes()); - } - }); - MockProxy { - addr, - shutdown, - handle: Some(handle), - } - } - - fn url(&self) -> String { - format!("http://{}/", self.addr) - } - - /// Wait until the proxy is accepting connections. - fn wait_ready(&self) { - for _ in 0..50 { - if TcpStream::connect_timeout(&self.addr, CONNECT_TIMEOUT).is_ok() { - return; - } - sleep(WAIT_TIMEOUT); - } - panic!("MockProxy not ready after 500 ms"); - } -} - -impl Drop for MockProxy { - fn drop(&mut self) { - self.shutdown.store(true, Ordering::Relaxed); - if let Some(handle) = self.handle.take() { - let _ = handle.join(); - } - } -} - -/// Extract a header value from an HTTP request string. -fn extract_header(request: &str, name: &str) -> Option { - let prefix = format!("{name}: "); - request - .lines() - .find(|l| l.to_lowercase().starts_with(&prefix.to_lowercase())) - .and_then(|l| l.split_once(": ").map(|(_, v)| v.trim().to_string())) -} - -/// Compute the Sec-WebSocket-Accept value per RFC 6455 §4.2.2. -fn compute_ws_accept(key: &str) -> String { - let mut hasher = Sha1::new(); - hasher.update(key.as_bytes()); - hasher.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"); - BASE64_STANDARD.encode(hasher.finalize()) -} - -/// Build a WebSocket text frame (unmasked, for server→client). -fn build_ws_text_frame(payload: &str) -> Vec { - let len = payload.len(); - let mut frame = vec![0x81]; // FIN + text opcode - if len < 126 { - frame.push(len as u8); - } else { - // Our test payloads are small; 16-bit length is sufficient. - frame.push(126); - frame.extend_from_slice(&(len as u16).to_be_bytes()); - } - frame.extend_from_slice(payload.as_bytes()); - frame -} - -fn mfa_enabled_location(name: &str, instance_id: Id) -> Location { - Location { - id: NoId, - instance_id, - network_id: 1, - name: name.into(), - address: "10.0.0.2/24".into(), - pubkey: "pk-loc".into(), - endpoint: "1.2.3.4:51820".into(), - allowed_ips: "0.0.0.0/0".into(), - dns: None, - route_all_traffic: false, - keepalive_interval: 25, - location_mfa_mode: LocationMfaMode::Internal, - service_location_mode: ServiceLocationMode::Disabled, - mfa_method: None, - posture_check_required: false, - } -} - -/// Build a location that requires **external** (OIDC) MFA. -fn oidc_location(name: &str, instance_id: Id) -> Location { - Location { - location_mfa_mode: LocationMfaMode::External, - ..mfa_enabled_location(name, instance_id) - } -} - -async fn seed_db(pool: &DbPool) -> (Instance, Location) { - let inst = Instance { - id: NoId, - name: "acme".into(), - uuid: "uuid-1".into(), - url: "https://core.example".into(), - proxy_url: String::new(), // filled later - username: "alice".into(), - token: None, - client_traffic_policy: ClientTrafficPolicy::None, - enterprise_enabled: false, - openid_display_name: None, - } - .save(pool) - .await - .unwrap(); - - // Insert wireguard keys (required by mfa::authorize). - sqlx::query("INSERT INTO wireguard_keys (instance_id, pubkey, prvkey) VALUES ($1, $2, $3)") - .bind(inst.id) - .bind("wg-pubkey") - .bind("wg-prvkey") - .execute(pool) - .await - .unwrap(); - - let loc = mfa_enabled_location("office", inst.id) - .save(pool) - .await - .unwrap(); - - (inst, loc) -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_mfa_success_returns_psk(pool: DbPool) { - let (mut instance, location) = seed_db(&pool).await; - let mock = MockProxy::new( - MockResponse { - status: 200, - body: r#"{"token":"tok-123"}"#.into(), - }, - MockResponse { - status: 200, - body: r#"{"preshared_key":"secret-psk"}"#.into(), - }, - ); - mock.wait_ready(); - instance.proxy_url = mock.url(); - - let source = CodeSource::Literal("123456".into()); - let psk = mfa::authorize(&location, &source, &instance, MfaMethod::Totp, None, &pool) - .await - .unwrap(); - assert_eq!(psk.expose_secret(), "secret-psk"); -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_mfa_rejection_returns_mfa_failed(pool: DbPool) { - let (mut instance, location) = seed_db(&pool).await; - let mock = MockProxy::new( - MockResponse { - status: 200, - body: r#"{"token":"tok-456"}"#.into(), - }, - MockResponse { - status: 403, - body: r#"{"error":"invalid TOTP code"}"#.into(), - }, - ); - mock.wait_ready(); - instance.proxy_url = mock.url(); - - let source = CodeSource::Literal("000000".into()); - let err = mfa::authorize(&location, &source, &instance, MfaMethod::Totp, None, &pool) - .await - .unwrap_err(); - - assert!(matches!(err, CliError::MfaFailed(_))); - assert!(err.to_string().contains("invalid TOTP code")); -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_mfa_proxy_unreachable(pool: DbPool) { - let (mut instance, location) = seed_db(&pool).await; - // Point at a port where nothing is listening. - instance.proxy_url = "http://127.0.0.1:19999/".into(); - - let source = CodeSource::Literal("123456".into()); - let err = mfa::authorize(&location, &source, &instance, MfaMethod::Totp, None, &pool) - .await - .unwrap_err(); - - assert!( - matches!(err, CliError::Other(_)), - "expected CliError::Other, got {err:?}" - ); - assert!(err.to_string().contains("Failed to reach proxy")); -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_oidc_mfa_success_returns_psk(pool: DbPool) { - let (mut instance, _location) = seed_db(&pool).await; - let oidc_location = oidc_location("office-oidc", instance.id) - .save(&pool) - .await - .unwrap(); - - let mock = MockProxy::with_poll( - MockResponse { - status: 200, - body: r#"{"token":"tok-oidc"}"#.into(), - }, - MockResponse { - status: 200, - body: r#"{"preshared_key":"secret-oidc-psk"}"#.into(), - }, - 2, // first two finish calls return 428, third returns 200 - ); - mock.wait_ready(); - instance.proxy_url = mock.url(); - - let psk = mfa::authorize_oidc(&oidc_location, &instance, None, &pool, false) - .await - .unwrap(); - assert_eq!(psk.expose_secret(), "secret-oidc-psk"); -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_oidc_mfa_times_out_when_never_completed(pool: DbPool) { - let (mut instance, _location) = seed_db(&pool).await; - let oidc_location = oidc_location("office-oidc", instance.id) - .save(&pool) - .await - .unwrap(); - - // Proxy always returns 428 for finish: never authenticates. - let mock = MockProxy::with_poll( - MockResponse { - status: 200, - body: r#"{"token":"tok-timeout"}"#.into(), - }, - MockResponse { - status: 200, - body: r#"{"preshared_key":"unreachable"}"#.into(), - }, - u32::MAX, // never returns 200 - ); - mock.wait_ready(); - instance.proxy_url = mock.url(); - - let err = mfa::authorize_oidc(&oidc_location, &instance, None, &pool, false) - .await - .unwrap_err(); - - assert!(matches!(err, CliError::MfaFailed(_))); - assert!(err.to_string().contains("timed out")); -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_mobile_approve_success_returns_psk(pool: DbPool) { - let (mut instance, location) = seed_db(&pool).await; - - // /start must return a challenge — mobile-approve requires it. - let mock = MockProxy::with_mobile_approve( - MockResponse { - status: 200, - body: r#"{"token":"tok-mob","challenge":"chal-xyz"}"#.into(), - }, - Some("secret-mobile-psk"), - ); - mock.wait_ready(); - instance.proxy_url = mock.url(); - - let psk = mfa::authorize_mobile_approve( - &location, &instance, None, // no posture data - None, // no qr-file (test render_qr is a no-op) - &pool, false, // not json mode - ) - .await - .unwrap(); - assert_eq!(psk.expose_secret(), "secret-mobile-psk"); -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_mobile_approve_no_device_returns_guidance(pool: DbPool) { - let (mut instance, location) = seed_db(&pool).await; - - let mock = MockProxy::with_mobile_approve( - MockResponse { - status: 400, - body: r#"{"error":"selected MFA method is not available"}"#.into(), - }, - Some("unused"), - ); - mock.wait_ready(); - instance.proxy_url = mock.url(); - - let err = mfa::authorize_mobile_approve(&location, &instance, None, None, &pool, false) - .await - .unwrap_err(); - - assert!(matches!(err, CliError::MfaFailed(_))); - assert!(err.to_string().contains("No mobile authenticator")); -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_mobile_approve_unrelated_error_passes_through(pool: DbPool) { - let (mut instance, location) = seed_db(&pool).await; - - let mock = MockProxy::with_mobile_approve( - MockResponse { - status: 500, - body: r#"{"error":"internal server error"}"#.into(), - }, - Some("unused"), - ); - mock.wait_ready(); - instance.proxy_url = mock.url(); - - let err = mfa::authorize_mobile_approve(&location, &instance, None, None, &pool, false) - .await - .unwrap_err(); - - assert!(matches!(err, CliError::Other(_))); - assert!(err.to_string().contains("internal server error")); -} - -#[sqlx::test(migrations = "../migrations")] -async fn test_mobile_approve_server_close_without_frame_times_out(pool: DbPool) { - let (mut instance, location) = seed_db(&pool).await; - - // psk=None: the mock accepts the WS upgrade but closes without sending a frame. - let mock = MockProxy::with_mobile_approve( - MockResponse { - status: 200, - body: r#"{"token":"tok-close","challenge":"chal-xyz"}"#.into(), - }, - None, - ); - mock.wait_ready(); - instance.proxy_url = mock.url(); - - let err = mfa::authorize_mobile_approve(&location, &instance, None, None, &pool, false) - .await - .unwrap_err(); - - assert!(matches!(err, CliError::MfaFailed(_))); - assert!(err.to_string().contains("connection closed by proxy")); -} diff --git a/src-tauri/client-proto/build.rs b/src-tauri/client-proto/build.rs index 95f41e21f..8d982665f 100644 --- a/src-tauri/client-proto/build.rs +++ b/src-tauri/client-proto/build.rs @@ -12,6 +12,22 @@ fn main() -> Result<(), Box> { ) // Make all messages serde-serializable. .type_attribute(".", "#[derive(serde::Serialize,serde::Deserialize)]") + // Use proto defaults for missing fields in enrollment types that + // may differ across proxy versions. + .type_attribute(".defguard.client_types.AdminInfo", "#[serde(default)]") + .type_attribute( + ".defguard.client_types.InitialUserInfo", + "#[serde(default)]", + ) + .type_attribute( + ".defguard.client_types.EnrollmentSettings", + "#[serde(default)]", + ) + .type_attribute(".defguard.client_types.InstanceInfo", "#[serde(default)]") + .type_attribute( + ".defguard.client_types.EnrollmentStartResponse", + "#[serde(default)]", + ) .compile_protos( &[ "../proto/v1/client/client.proto", diff --git a/src-tauri/core/Cargo.toml b/src-tauri/core/Cargo.toml index 8b98cf965..6789a367e 100644 --- a/src-tauri/core/Cargo.toml +++ b/src-tauri/core/Cargo.toml @@ -16,6 +16,7 @@ defguard-client-common = { path = "../common" } defguard-client-proto = { path = "../client-proto" } defguard_wireguard_rs.workspace = true dirs-next.workspace = true +futures-util.workspace = true hyper-util = "0.1" log.workspace = true os_info = { version = "3.14", default-features = false } @@ -31,13 +32,18 @@ strum = { version = "0.28", features = ["derive"] } struct-patch = "0.12" thiserror.workspace = true tokio.workspace = true +tokio-tungstenite.workspace = true +tokio-util = "0.7" tonic.workspace = true tower = "0.5" tracing = { workspace = true } x25519-dalek = { version = "2", features = ["getrandom", "serde", "static_secrets"] } [dev-dependencies] +futures-util.workspace = true tempfile.workspace = true +tokio-tungstenite.workspace = true +wiremock.workspace = true [target.'cfg(unix)'.dependencies] nix = { version = "0.31", features = ["user", "fs"] } diff --git a/src-tauri/core/src/enrollment.rs b/src-tauri/core/src/enrollment.rs new file mode 100644 index 000000000..8c09dc4d4 --- /dev/null +++ b/src-tauri/core/src/enrollment.rs @@ -0,0 +1,615 @@ +//! Enrollment flow for adding a new Defguard instance. +//! +//! Handles the client-side enrollment protocol against the Edge proxy: +//! starting enrollment, creating a device, activating the user, registering +//! MFA, and finishing the enrollment. + +use defguard_client_proto::defguard::client_types::EnrollmentStartResponse; +use reqwest::{Client, Response, StatusCode, Url}; +use serde::{Deserialize, Serialize}; +use serde_json::Value as JsonValue; +use thiserror::Error; + +use crate::{ + proxy::construct_platform_header, + version::{CLIENT_PLATFORM_HEADER, CLIENT_VERSION_HEADER, PKG_VERSION}, +}; + +/// Error type returned by enrollment operations. +/// +/// Serialized as a tagged JSON union so the TypeScript frontend can +/// match on the `type` field to show context-specific messages. +#[derive(Debug, Error, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum EnrollmentError { + #[error("Enrollment token has expired or is invalid")] + TokenExpired, + + #[error("Session cookie not found in enrollment response")] + MissingCookie, + + #[error("{message}")] + NetworkError { message: String }, + + #[error("Proxy error (HTTP {status}): {message}")] + ProxyError { status: u16, message: String }, + + #[error("{message}")] + Other { message: String }, +} + +/// Holds the enrollment session state. +#[derive(Clone, Debug)] +pub struct EnrollmentSession { + pub cookie: String, + pub proxy_url: Url, + pub client: Client, +} + +/// Response from `POST /api/v1/enrollment/register-mfa/code/start`. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct MfaStartResponse { + pub totp_secret: Option, +} + +/// Response from `POST /api/v1/enrollment/register-mfa/code/finish`. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct MfaFinishResponse { + pub recovery_codes: Vec, +} + +/// Send a JSON POST request to an enrollment endpoint with the session +/// cookie and standard client headers. +async fn enrollment_post( + session: &EnrollmentSession, + path: &str, + body: serde_json::Value, +) -> Result { + let url = session + .proxy_url + .join(path) + .map_err(|e| EnrollmentError::Other { + message: format!("Failed to build URL '{path}': {e}"), + })?; + let response = session + .client + .post(url) + .json(&body) + .header("Cookie", &session.cookie) + .header(CLIENT_VERSION_HEADER, PKG_VERSION) + .header(CLIENT_PLATFORM_HEADER, construct_platform_header()) + .send() + .await + .map_err(|e| EnrollmentError::NetworkError { + message: format!("Failed to reach proxy: {e}"), + })?; + + check_enrollment_response(response).await +} + +/// Check an enrollment response status and map it to `EnrollmentError`. +async fn check_enrollment_response(response: Response) -> Result { + let status = response.status(); + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + return Err(EnrollmentError::TokenExpired); + } + if !status.is_success() { + let message = read_error_body(response).await; + return Err(EnrollmentError::ProxyError { + status: status.as_u16(), + message, + }); + } + Ok(response) +} + +/// Extract the `defguard_proxy` session cookie from a response's `Set-Cookie` +/// headers. +fn extract_defguard_cookie(response: &Response) -> Result { + for value in response.headers().get_all(reqwest::header::SET_COOKIE) { + let raw = value.to_str().unwrap_or_default(); + if raw.starts_with("defguard_proxy=") { + // Take everything up to the first `;`. + return Ok(raw.split(';').next().unwrap_or(raw).to_string()); + } + } + Err(EnrollmentError::MissingCookie) +} + +/// Read an error body from a non-2xx response. +async fn read_error_body(response: Response) -> String { + let status = response.status(); + response + .json::() + .await + .ok() + .and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from)) + .unwrap_or_else(|| format!("HTTP {status}")) +} + +/// Start the enrollment process. +/// +/// POSTs `{ "token": "" }` to `/api/v1/enrollment/start`, extracts +/// the `defguard_proxy` session cookie from the `Set-Cookie` response +/// header, and returns the session together with the full enrollment +/// start response (admin, user, settings, instance, deadline). +pub async fn enrollment_start( + proxy_url: Url, + token: String, +) -> Result<(EnrollmentSession, EnrollmentStartResponse), EnrollmentError> { + let client = Client::new(); + + let url = proxy_url + .join("api/v1/enrollment/start") + .map_err(|e| EnrollmentError::Other { + message: format!("Failed to build enrollment start URL: {e}"), + })?; + + let response = client + .post(url) + .json(&serde_json::json!({ "token": token })) + .header(CLIENT_VERSION_HEADER, PKG_VERSION) + .header(CLIENT_PLATFORM_HEADER, construct_platform_header()) + .send() + .await + .map_err(|e| EnrollmentError::NetworkError { + message: format!("Failed to reach proxy: {e}"), + })?; + + let status = response.status(); + + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + return Err(EnrollmentError::TokenExpired); + } + + if !status.is_success() { + let message = read_error_body(response).await; + return Err(EnrollmentError::ProxyError { + status: status.as_u16(), + message, + }); + } + + let cookie = extract_defguard_cookie(&response)?; + + let body: EnrollmentStartResponse = + response.json().await.map_err(|e| EnrollmentError::Other { + message: format!("Failed to parse enrollment response: {e}"), + })?; + + let session = EnrollmentSession { + cookie, + proxy_url, + client, + }; + + Ok((session, body)) +} + +/// Create a device during enrollment. +/// +/// POSTs `{ "name": "...", "pubkey": "..." }` to +/// `/api/v1/enrollment/create_device` and returns the full device +/// configuration response as a JSON value. +pub async fn enrollment_create_device( + session: EnrollmentSession, + name: String, + pubkey: String, +) -> Result { + let response = enrollment_post( + &session, + "api/v1/enrollment/create_device", + serde_json::json!({ "name": name, "pubkey": pubkey }), + ) + .await?; + + response.json().await.map_err(|e| EnrollmentError::Other { + message: format!("Failed to parse create_device response: {e}"), + }) +} + +/// Activate the user account during enrollment. +/// +/// POSTs `{ "password": "...", "phone_number": "..." }` to +/// `/api/v1/enrollment/activate_user`. Either field may be omitted when +/// the server does not require it (externally managed users skip password; +/// most users skip phone). +pub async fn enrollment_activate_user( + session: EnrollmentSession, + password: Option, + phone_number: Option, +) -> Result<(), EnrollmentError> { + enrollment_post( + &session, + "api/v1/enrollment/activate_user", + serde_json::json!({ + "password": password, + "phone_number": phone_number, + }), + ) + .await?; + + Ok(()) +} + +/// Start MFA registration during enrollment. +/// +/// POSTs `{ "method": "..." }` to +/// `/api/v1/enrollment/register-mfa/code/start`. Returns the TOTP secret +/// (for TOTP method) or an empty response (for email method). +pub async fn enrollment_register_mfa_start( + session: EnrollmentSession, + method: String, +) -> Result { + let response = enrollment_post( + &session, + "api/v1/enrollment/register-mfa/code/start", + serde_json::json!({ "method": method }), + ) + .await?; + + response.json().await.map_err(|e| EnrollmentError::Other { + message: format!("Failed to parse MFA start response: {e}"), + }) +} + +/// Finish MFA registration during enrollment. +/// +/// POSTs `{ "code": "...", "method": "..." }` to +/// `/api/v1/enrollment/register-mfa/code/finish`. Returns the recovery +/// codes that the user should save. +pub async fn enrollment_register_mfa_finish( + session: EnrollmentSession, + code: String, + method: String, +) -> Result { + let response = enrollment_post( + &session, + "api/v1/enrollment/register-mfa/code/finish", + serde_json::json!({ "code": code, "method": method }), + ) + .await?; + + response.json().await.map_err(|e| EnrollmentError::Other { + message: format!("Failed to parse MFA finish response: {e}"), + }) +} + +/// Fetch network configuration for an existing device during enrollment. +/// +/// POSTs `{ "pubkey": "..." }` to `/api/v1/enrollment/network_info`. +/// This is the fast path for re-enrolling a device whose WireGuard keys +/// already exist on the server. +pub async fn enrollment_network_info( + session: EnrollmentSession, + pubkey: String, +) -> Result { + let response = enrollment_post( + &session, + "api/v1/enrollment/network_info", + serde_json::json!({ "pubkey": pubkey }), + ) + .await?; + + response.json().await.map_err(|e| EnrollmentError::Other { + message: format!("Failed to parse network_info response: {e}"), + }) +} + +/// Mark the enrollment session as finished. +/// +/// There is no HTTP call for this step -- the session cookie is already +/// cleared by the `activate_user` call on the server side. This function +/// exists as an explicit anchor point for the Tauri command so the +/// implementation mirrors the UI flow. It simply drops the session. +pub fn enrollment_finish(_session: EnrollmentSession) { + // Session dropped; cookie is no longer needed. +} + +#[cfg(test)] +mod tests { + use reqwest::Url; + use serde_json::json; + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + + use super::*; + + fn mock_url(server: &MockServer) -> Url { + Url::parse(&server.uri()).expect("MockServer URI should be valid") + } + + fn user_json() -> serde_json::Value { + json!({ + "first_name": "John", + "last_name": "Doe", + "login": "jdoe", + "email": "john@example.com", + "phone_number": null, + "is_active": true, + "device_names": [], + "enrolled": false, + "is_admin": false, + "password_management_disabled": false, + }) + } + + fn full_start_json() -> serde_json::Value { + json!({ + "admin": { + "name": "Admin", + "phone_number": null, + "email": "admin@example.com", + }, + "user": user_json(), + "settings": { + "vpn_setup_optional": false, + "only_client_activation": false, + "admin_device_management": false, + "smtp_configured": true, + "mfa_required": true, + }, + "instance": { + "id": "inst-1", + "name": "Test Instance", + "url": "https://test.defguard.net", + "proxy_url": "https://proxy.defguard.net", + "username": "jdoe", + "enterprise_enabled": false, + "disable_all_traffic": false, + "openid_display_name": null, + }, + "deadline_timestamp": 1734567890, + "final_page_content": "Welcome!", + }) + } + + #[tokio::test] + async fn test_start_success() { + let server = MockServer::start().await; + let response_body = full_start_json(); + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/start")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(&response_body) + .insert_header( + "Set-Cookie", + "defguard_proxy=test-session-cookie; Path=/api/v1/enrollment", + ), + ) + .mount(&server) + .await; + + let url = mock_url(&server); + let (session, response) = enrollment_start(url, "valid-token".into()).await.unwrap(); + let user = response.user.expect("user must be present"); + let admin = response.admin.expect("admin must be present"); + let settings = response.settings.expect("settings must be present"); + let instance = response.instance.expect("instance must be present"); + + assert_eq!(session.cookie, "defguard_proxy=test-session-cookie"); + assert_eq!(user.first_name, "John"); + assert_eq!(user.login, "jdoe"); + assert_eq!(admin.name, "Admin"); + assert!(settings.mfa_required); + assert_eq!(instance.id, "inst-1"); + assert_eq!(response.deadline_timestamp, 1734567890); + assert_eq!(response.final_page_content, "Welcome!"); + } + + #[tokio::test] + async fn test_start_token_expired_401() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/start")) + .respond_with(ResponseTemplate::new(401)) + .mount(&server) + .await; + + let url = mock_url(&server); + let err = enrollment_start(url, "bad-token".into()).await.unwrap_err(); + + assert!(matches!(err, EnrollmentError::TokenExpired)); + } + + #[tokio::test] + async fn test_start_token_expired_403() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/start")) + .respond_with(ResponseTemplate::new(403)) + .mount(&server) + .await; + + let url = mock_url(&server); + let err = enrollment_start(url, "bad-token".into()).await.unwrap_err(); + + assert!(matches!(err, EnrollmentError::TokenExpired)); + } + + #[tokio::test] + async fn test_start_missing_cookie() { + let server = MockServer::start().await; + let response_body = json!({ "user": user_json() }); + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/start")) + .respond_with(ResponseTemplate::new(200).set_body_json(&response_body)) + .mount(&server) + .await; + + let url = mock_url(&server); + let err = enrollment_start(url, "token".into()).await.unwrap_err(); + + assert!(matches!(err, EnrollmentError::MissingCookie)); + } + + #[tokio::test] + async fn test_start_proxy_error() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/start")) + .respond_with( + ResponseTemplate::new(500).set_body_json(json!({ "error": "internal boom" })), + ) + .mount(&server) + .await; + + let url = mock_url(&server); + let err = enrollment_start(url, "token".into()).await.unwrap_err(); + + match err { + EnrollmentError::ProxyError { status, message } => { + assert_eq!(status, 500); + assert!(message.contains("internal boom")); + } + other => panic!("expected ProxyError, got {other:?}"), + } + } + + #[tokio::test] + async fn test_start_network_error() { + // Use an address where nothing is listening. + let url = "http://127.0.0.1:1".parse().unwrap(); + let err = enrollment_start(url, "token".into()).await.unwrap_err(); + + assert!(matches!(err, EnrollmentError::NetworkError { .. })); + } + + fn make_session(server: &MockServer) -> EnrollmentSession { + EnrollmentSession { + cookie: "defguard_proxy=test-session-cookie".into(), + proxy_url: mock_url(server), + client: Client::new(), + } + } + + #[tokio::test] + async fn test_create_device_success() { + let server = MockServer::start().await; + let response_body = json!({ "config": "wg0", "assignedIp": "10.0.0.1" }); + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/create_device")) + .respond_with(ResponseTemplate::new(200).set_body_json(&response_body)) + .mount(&server) + .await; + + let session = make_session(&server); + let result = enrollment_create_device(session, "my-device".into(), "pk".into()) + .await + .unwrap(); + + assert_eq!(result["config"], "wg0"); + } + + #[tokio::test] + async fn test_activate_user_success() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/activate_user")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + + let session = make_session(&server); + enrollment_activate_user(session, Some("p4ssw0rd".into()), None) + .await + .unwrap(); + } + + // enrollment_register_mfa_start + + #[tokio::test] + async fn test_register_mfa_start_success() { + let server = MockServer::start().await; + let response_body = json!({ "totp_secret": "SECRET123" }); + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/register-mfa/code/start")) + .respond_with(ResponseTemplate::new(200).set_body_json(&response_body)) + .mount(&server) + .await; + + let session = make_session(&server); + let result = enrollment_register_mfa_start(session, "totp".into()) + .await + .unwrap(); + + assert_eq!(result.totp_secret.as_deref(), Some("SECRET123")); + } + + // enrollment_register_mfa_finish + + #[tokio::test] + async fn test_register_mfa_finish_success() { + let server = MockServer::start().await; + let response_body = json!({ "recovery_codes": ["rc1", "rc2", "rc3"] }); + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/register-mfa/code/finish")) + .respond_with(ResponseTemplate::new(200).set_body_json(&response_body)) + .mount(&server) + .await; + + let session = make_session(&server); + let result = enrollment_register_mfa_finish(session, "123456".into(), "totp".into()) + .await + .unwrap(); + + assert_eq!(result.recovery_codes, vec!["rc1", "rc2", "rc3"]); + } + + // enrollment_network_info + + #[tokio::test] + async fn test_network_info_success() { + let server = MockServer::start().await; + let response_body = json!({ "config": "wg0", "assignedIp": "10.0.0.2" }); + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/network_info")) + .respond_with(ResponseTemplate::new(200).set_body_json(&response_body)) + .mount(&server) + .await; + + let session = make_session(&server); + let result = enrollment_network_info(session, "pk".into()).await.unwrap(); + + assert_eq!(result["assignedIp"], "10.0.0.2"); + } + + #[tokio::test] + async fn test_network_info_not_found() { + // A 404 means the device was deleted server-side. The status must be + // preserved as ProxyError { status: 404 } so the client can detect it + // and re-enroll (see the addInstance recovery path in the frontend). + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/enrollment/network_info")) + .respond_with( + ResponseTemplate::new(404).set_body_json(json!({ "error": "device not found" })), + ) + .mount(&server) + .await; + + let session = make_session(&server); + let err = enrollment_network_info(session, "pk".into()) + .await + .unwrap_err(); + + assert!(matches!( + err, + EnrollmentError::ProxyError { status: 404, .. } + )); + } +} diff --git a/src-tauri/core/src/events.rs b/src-tauri/core/src/events.rs index 71f859e31..2d3e1a69e 100644 --- a/src-tauri/core/src/events.rs +++ b/src-tauri/core/src/events.rs @@ -17,6 +17,10 @@ pub enum EventKey { WindowSwapped, SessionStateChanged, InstanceUpdated, + MfaOpenIdComplete, + MfaOpenIdError, + MfaMobileComplete, + MfaMobileError, } impl From for &'static str { @@ -37,6 +41,10 @@ impl From for &'static str { EventKey::WindowSwapped => "window-swapped", EventKey::SessionStateChanged => "session-state-changed", EventKey::InstanceUpdated => "instance-updated", + EventKey::MfaOpenIdComplete => "mfa-openid-complete", + EventKey::MfaOpenIdError => "mfa-openid-error", + EventKey::MfaMobileComplete => "mfa-mobile-complete", + EventKey::MfaMobileError => "mfa-mobile-error", } } } diff --git a/src-tauri/core/src/lib.rs b/src-tauri/core/src/lib.rs index e29ea56ad..c9a0ea2ce 100644 --- a/src-tauri/core/src/lib.rs +++ b/src-tauri/core/src/lib.rs @@ -16,9 +16,13 @@ use serde::{Deserialize, Serialize}; pub mod app_config; pub mod connection; pub mod database; +pub mod enrollment; pub mod error; pub mod events; +pub mod mfa; pub mod proxy; +#[cfg(test)] +mod test_helpers; pub mod version; pub mod wg_config; diff --git a/src-tauri/core/src/mfa.rs b/src-tauri/core/src/mfa.rs new file mode 100644 index 000000000..b72c3f36b --- /dev/null +++ b/src-tauri/core/src/mfa.rs @@ -0,0 +1,836 @@ +//! Connect-time VPN MFA over HTTP. +//! +//! Synchronous (request/response) MFA functions for TOTP and email methods, +//! plus long-running flows for OpenID (poll loop) and mobile approve (WebSocket). + +use std::time::Duration; + +use defguard_client_proto::defguard::client_types::{ + ClientMfaFinishRequest, ClientMfaFinishResponse, ClientMfaStartRequest, ClientMfaStartResponse, + MfaMethod, +}; +use futures_util::StreamExt; +use reqwest::{Client, Response, StatusCode, Url}; +use serde::Serialize; +use thiserror::Error; +use tokio::{ + net::TcpStream, + select, + time::{sleep, Instant}, +}; +use tokio_tungstenite::{ + connect_async, + tungstenite::{Error as WsError, Message}, + MaybeTlsStream, WebSocketStream, +}; +use tokio_util::sync::CancellationToken; + +use crate::{ + proxy::construct_platform_header, + version::{CLIENT_PLATFORM_HEADER, CLIENT_VERSION_HEADER, PKG_VERSION}, +}; + +/// Error type returned by MFA operations. +/// +/// Serialized as a tagged JSON union so the TypeScript frontend can +/// match on the `type` field to show context-specific messages. +#[derive(Debug, Error, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum MfaError { + #[error("{message}")] + NetworkError { message: String }, + + #[error("Proxy error (HTTP {status}): {message}")] + ProxyError { status: u16, message: String }, + + #[error("MFA rejected: {message}")] + MfaRejected { message: String }, + + #[error("Posture check failed: {message}")] + PostureRejected { message: String }, + + #[error("MFA operation timed out")] + Timeout, + + #[error("MFA operation cancelled")] + Cancelled, + + #[error("{message}")] + Other { message: String }, +} + +fn build_client() -> Client { + Client::new() +} + +fn standard_headers() -> Vec<(&'static str, String)> { + vec![ + (CLIENT_VERSION_HEADER, PKG_VERSION.to_string()), + (CLIENT_PLATFORM_HEADER, construct_platform_header()), + ] +} + +/// Check an MFA response status and map it to `MfaError`. +async fn check_mfa_response(response: Response) -> Result { + let status = response.status(); + if status.is_success() { + return Ok(response); + } + + let message = response + .json::() + .await + .ok() + .and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from)) + .unwrap_or_else(|| format!("HTTP {status}")); + + match status { + // The proxy returns 403 only for a failed device posture check + // (ApiError::PostureRejected); 401 and other 4xx are ordinary MFA + // rejections. Keeping them distinct lets the frontend route posture + // failures to the dedicated posture-check-failed view. + StatusCode::FORBIDDEN => Err(MfaError::PostureRejected { message }), + StatusCode::UNAUTHORIZED => Err(MfaError::MfaRejected { message }), + _ if status.is_client_error() => Err(MfaError::MfaRejected { message }), + _ => Err(MfaError::ProxyError { + status: status.as_u16(), + message, + }), + } +} + +/// Start an MFA handshake for a VPN location. +/// +/// POSTs a `ClientMfaStartRequest` (proto JSON) to +/// `/api/v1/client-mfa/start` and returns the session token (and +/// optionally the biometric challenge). +pub async fn mfa_start( + proxy_url: Url, + request: ClientMfaStartRequest, +) -> Result { + let client = build_client(); + + let url = proxy_url + .join("api/v1/client-mfa/start") + .map_err(|e| MfaError::Other { + message: format!("Failed to build MFA start URL: {e}"), + })?; + + let mut req = client.post(url).json(&request); + + for (k, v) in standard_headers() { + req = req.header(k, v); + } + + let response = req.send().await.map_err(|e| MfaError::NetworkError { + message: format!("Failed to reach proxy: {e}"), + })?; + + let response = match check_mfa_response(response).await { + Ok(response) => response, + Err(err) => return Err(rewrap_mobile_start_error(request.method, err)), + }; + response.json().await.map_err(|e| MfaError::Other { + message: format!("Invalid MFA start response: {e}"), + }) +} + +/// Turn the proxy's generic "selected MFA method is not available" rejection +/// into actionable guidance for mobile-approve MFA (the user has no registered +/// mobile authenticator). Restores the CLI behavior that was lost when this +/// logic moved into core; benefits the desktop client too. Non-mobile methods +/// keep the original message. +fn rewrap_mobile_start_error(method: i32, err: MfaError) -> MfaError { + if method == MfaMethod::MobileApprove as i32 { + if let MfaError::MfaRejected { message } = &err { + if message.contains("selected MFA method is not available") { + return MfaError::MfaRejected { + message: "No mobile authenticator is registered for your account. \ + Register one in the Defguard mobile app, then retry." + .into(), + }; + } + } + } + err +} + +/// Finish an MFA handshake using a one-time code (TOTP or email). +/// +/// POSTs a `ClientMfaFinishRequest` to `/api/v1/client-mfa/finish` +/// and returns the preshared key. +pub async fn mfa_finish_code( + proxy_url: Url, + request: ClientMfaFinishRequest, +) -> Result { + let client = build_client(); + + let url = proxy_url + .join("api/v1/client-mfa/finish") + .map_err(|e| MfaError::Other { + message: format!("Failed to build MFA finish URL: {e}"), + })?; + + let mut req = client.post(url).json(&request); + + for (k, v) in standard_headers() { + req = req.header(k, v); + } + + let response = req.send().await.map_err(|e| MfaError::NetworkError { + message: format!("Failed to reach proxy: {e}"), + })?; + + let response = check_mfa_response(response).await?; + response.json().await.map_err(|e| MfaError::Other { + message: format!("Invalid MFA finish response: {e}"), + }) +} + +#[cfg(not(test))] +const OIDC_POLL_INTERVAL: Duration = Duration::from_secs(5); +#[cfg(test)] +const OIDC_POLL_INTERVAL: Duration = Duration::from_millis(5); + +#[cfg(not(test))] +const OIDC_POLL_TIMEOUT: Duration = Duration::from_secs(300); +#[cfg(test)] +const OIDC_POLL_TIMEOUT: Duration = Duration::from_millis(200); + +#[cfg(not(test))] +const MOBILE_APPROVE_TIMEOUT: Duration = Duration::from_secs(120); +#[cfg(test)] +const MOBILE_APPROVE_TIMEOUT: Duration = Duration::from_secs(5); + +/// Poll the proxy for OpenID MFA completion. +/// +/// The caller must already have opened the browser to the OIDC provider +/// URL (the token from `mfa_start` encodes the redirect). This function +/// POSTs a `ClientMfaFinishRequest` to `/api/v1/client-mfa/finish` every +/// [`OIDC_POLL_INTERVAL`] until the server returns a 200 (success), +/// the deadline expires, or the [`CancellationToken`] is fired. +pub async fn poll_openid_mfa( + proxy_url: Url, + token: String, + cancel: CancellationToken, +) -> Result { + let client = build_client(); + let url = proxy_url + .join("api/v1/client-mfa/finish") + .map_err(|e| MfaError::Other { + message: format!("Failed to build MFA finish URL: {e}"), + })?; + + let deadline = Instant::now() + OIDC_POLL_TIMEOUT; + + let request = ClientMfaFinishRequest { + token, + code: None, + auth_pub_key: None, + }; + + loop { + let remaining = deadline + .checked_duration_since(Instant::now()) + .unwrap_or_default(); + if remaining.is_zero() { + return Err(MfaError::Timeout); + } + + let mut req = client.post(url.clone()).json(&request); + for (k, v) in standard_headers() { + req = req.header(k, v); + } + + select! { + () = cancel.cancelled() => { + return Err(MfaError::Cancelled); + } + result = req.send() => { + let response = result.map_err(|e| MfaError::NetworkError { + message: format!("Failed to reach proxy: {e}"), + })?; + + let status = response.status(); + if status == StatusCode::OK { + return response.json().await.map_err(|e| MfaError::Other { + message: format!("Invalid MFA finish response: {e}"), + }); + } + if status != StatusCode::PRECONDITION_REQUIRED { + return Err(check_mfa_response(response).await.err().unwrap_or( + MfaError::Other { + message: format!("Unexpected status: {status}"), + }, + )); + } + // 428: not complete yet — fall through to sleep. + } + } + + select! { + () = cancel.cancelled() => { + return Err(MfaError::Cancelled); + } + () = sleep(OIDC_POLL_INTERVAL) => {} + } + } +} + +/// Connect to a WebSocket endpoint and wait for mobile-approve MFA +/// completion. +/// +/// The caller must have already displayed the QR code to the user +/// (the token from `mfa_start` encodes the challenge). This function +/// opens a WebSocket to `ws_url` and waits for a +/// `{"type":"mfa_success","preshared_key":"..."}` text frame. +/// Returns [`MfaError::Cancelled`] if the token fires or +/// [`MfaError::Timeout`] if the deadline expires. +pub async fn connect_mobile_approve( + ws_url: &str, + cancel: CancellationToken, +) -> Result { + let (ws_stream, _response) = + connect_async(ws_url) + .await + .map_err(|e| MfaError::NetworkError { + // Never interpolate the raw error: `ws_url` carries the MFA + // token as a query parameter and can appear in the error's + // Display, which is surfaced to the frontend and logs. + message: match &e { + WsError::Io(io_err) => { + format!("Failed to connect to proxy ({})", io_err.kind()) + } + _ => "Failed to connect to proxy".to_string(), + }, + })?; + + wait_for_mfa_success(ws_stream, cancel).await +} + +/// Derive the WebSocket URL from the proxy's base URL and MFA token. +pub fn derive_ws_url(proxy_base: &Url, token: &str) -> Result { + let mut ws_url = proxy_base + .join("api/v1/client-mfa/remote") + .map_err(|e| MfaError::Other { + message: format!("Failed to build WebSocket URL: {e}"), + })?; + + let ws_scheme = match proxy_base.scheme() { + "https" => "wss", + "http" => "ws", + other => { + return Err(MfaError::Other { + message: format!("Invalid proxy URL scheme '{other}'; expected http or https"), + }); + } + }; + + ws_url.set_scheme(ws_scheme).map_err(|()| MfaError::Other { + message: "Failed to set WebSocket URL scheme".into(), + })?; + ws_url.query_pairs_mut().append_pair("token", token); + + Ok(ws_url.to_string()) +} + +/// Wait on the WebSocket for an `mfa_success` frame. +async fn wait_for_mfa_success( + ws_stream: WebSocketStream>, + cancel: CancellationToken, +) -> Result { + let (_write, mut read) = ws_stream.split(); + let deadline = Instant::now() + MOBILE_APPROVE_TIMEOUT; + + loop { + let remaining = deadline + .checked_duration_since(Instant::now()) + .unwrap_or_default(); + if remaining.is_zero() { + return Err(MfaError::Timeout); + } + + let msg = select! { + () = sleep(remaining) => { + return Err(MfaError::Timeout); + } + () = cancel.cancelled() => { + return Err(MfaError::Cancelled); + } + msg = read.next() => { + match msg { + Some(Ok(msg)) => msg, + Some(Err(_)) | None => { + return Err(MfaError::MfaRejected { + message: "mobile approval failed: connection closed by proxy" + .into(), + }); + } + } + } + }; + + if let Message::Text(text) = msg { + if let Ok(parsed) = serde_json::from_str::(&text) { + if parsed.get("type").and_then(|v| v.as_str()) == Some("mfa_success") { + if let Some(key) = parsed["preshared_key"].as_str() { + return Ok(ClientMfaFinishResponse { + preshared_key: key.to_string(), + token: None, + }); + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use reqwest::Url; + use serde_json::json; + use tokio_util::sync::CancellationToken; + use wiremock::{ + matchers::{body_partial_json, method, path}, + Mock, MockServer, ResponseTemplate, + }; + + use super::*; + use crate::test_helpers::{start_ws_stub, WsStubCommand}; + + fn mock_url(server: &MockServer) -> Url { + Url::parse(&server.uri()).expect("MockServer URI should be valid") + } + + fn start_request() -> ClientMfaStartRequest { + ClientMfaStartRequest { + location_id: 1, + pubkey: "pk".into(), + method: 0, // TOTP + posture_data: None, + } + } + + fn start_response_json(token: &str) -> serde_json::Value { + json!({ + "token": token, + "challenge": null, + }) + } + + fn finish_response_json(key: &str) -> serde_json::Value { + json!({ + "preshared_key": key, + }) + } + + #[tokio::test] + async fn test_mfa_start_success() { + let server = MockServer::start().await; + let body = start_response_json("mfa-token-1"); + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/start")) + .respond_with(ResponseTemplate::new(200).set_body_json(&body)) + .mount(&server) + .await; + + let url = mock_url(&server); + let info = mfa_start(url, start_request()).await.unwrap(); + assert_eq!(info.token, "mfa-token-1"); + assert!(info.challenge.is_none()); + } + + #[tokio::test] + async fn test_mfa_start_rejected() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/start")) + .respond_with( + ResponseTemplate::new(401).set_body_json(json!({ "error": "unauthorized" })), + ) + .mount(&server) + .await; + + let url = mock_url(&server); + let err = mfa_start(url, start_request()).await.unwrap_err(); + assert!(matches!(err, MfaError::MfaRejected { .. })); + } + + #[tokio::test] + async fn test_mfa_start_posture_rejected_on_403() { + // 403 is the proxy's posture-check-failure status; it must map to the + // dedicated PostureRejected variant, not the generic MfaRejected. + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/start")) + .respond_with( + ResponseTemplate::new(403).set_body_json(json!({ "error": "firewall enabled" })), + ) + .mount(&server) + .await; + + let url = mock_url(&server); + let err = mfa_start(url, start_request()).await.unwrap_err(); + assert!(matches!(err, MfaError::PostureRejected { .. })); + } + + #[tokio::test] + async fn test_mfa_start_sends_snake_case_numeric_body() { + // Guards the wire contract: the proxy expects snake_case fields and a + // *numeric* `method`. If serde ever serialized camelCase or a string + // enum, the body matcher fails, the mock returns nothing, and the call + // errors instead of silently sending a malformed request. + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/start")) + .and(body_partial_json( + json!({ "location_id": 1, "pubkey": "pk", "method": 0 }), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(start_response_json("t"))) + .mount(&server) + .await; + + let url = mock_url(&server); + mfa_start(url, start_request()) + .await + .expect("request body did not match the expected wire contract"); + } + + #[tokio::test] + async fn test_mfa_start_network_error() { + // Nothing listening on this port. + let url = "http://127.0.0.1:1".parse().unwrap(); + let err = mfa_start(url, start_request()).await.unwrap_err(); + assert!(matches!(err, MfaError::NetworkError { .. })); + } + + #[tokio::test] + async fn test_mfa_start_proxy_error_on_5xx() { + // 5xx is a server fault (ProxyError), distinct from a 4xx rejection. + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/start")) + .respond_with(ResponseTemplate::new(500).set_body_json(json!({ "error": "boom" }))) + .mount(&server) + .await; + + let url = mock_url(&server); + let err = mfa_start(url, start_request()).await.unwrap_err(); + assert!(matches!(err, MfaError::ProxyError { status: 500, .. })); + } + + #[tokio::test] + async fn test_mfa_start_mobile_no_authenticator_guidance() { + // Mobile-approve start rejected because no authenticator is registered: + // the generic proxy message becomes actionable guidance. + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/start")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(json!({ "error": "selected MFA method is not available" })), + ) + .mount(&server) + .await; + + let url = mock_url(&server); + let request = ClientMfaStartRequest { + location_id: 1, + pubkey: "pk".into(), + method: MfaMethod::MobileApprove as i32, + posture_data: None, + }; + match mfa_start(url, request).await.unwrap_err() { + MfaError::MfaRejected { message } => { + assert!( + message.contains("mobile app"), + "unexpected message: {message}" + ); + } + other => panic!("expected MfaRejected, got {other:?}"), + } + } + + #[tokio::test] + async fn test_mfa_start_non_mobile_not_rewrapped() { + // The mobile guidance must not leak into other methods. + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/start")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(json!({ "error": "selected MFA method is not available" })), + ) + .mount(&server) + .await; + + let url = mock_url(&server); + // start_request() uses method 0 (TOTP). + match mfa_start(url, start_request()).await.unwrap_err() { + MfaError::MfaRejected { message } => { + assert!( + !message.contains("mobile app"), + "TOTP got mobile guidance: {message}" + ); + } + other => panic!("expected MfaRejected, got {other:?}"), + } + } + + #[tokio::test] + async fn test_mfa_finish_code_success() { + let server = MockServer::start().await; + let body = finish_response_json("psk-123"); + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/finish")) + .respond_with(ResponseTemplate::new(200).set_body_json(&body)) + .mount(&server) + .await; + + let url = mock_url(&server); + let psk = mfa_finish_code( + url, + ClientMfaFinishRequest { + token: "token".into(), + code: Some("123456".into()), + auth_pub_key: None, + }, + ) + .await + .unwrap(); + assert_eq!(psk.preshared_key, "psk-123"); + } + + #[tokio::test] + async fn test_mfa_finish_code_rejected() { + // A wrong code is a 4xx rejection. + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/finish")) + .respond_with( + ResponseTemplate::new(401).set_body_json(json!({ "error": "Unauthorized" })), + ) + .mount(&server) + .await; + + let url = mock_url(&server); + let err = mfa_finish_code( + url, + ClientMfaFinishRequest { + token: "token".into(), + code: Some("000000".into()), + auth_pub_key: None, + }, + ) + .await + .unwrap_err(); + assert!(matches!(err, MfaError::MfaRejected { .. })); + } + + #[tokio::test] + async fn test_poll_openid_success() { + let server = MockServer::start().await; + let body = finish_response_json("oidc-psk"); + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/finish")) + .respond_with(ResponseTemplate::new(200).set_body_json(&body)) + .mount(&server) + .await; + + let url = mock_url(&server); + let cancel = CancellationToken::new(); + let psk = poll_openid_mfa(url, "token".into(), cancel).await.unwrap(); + assert_eq!(psk.preshared_key, "oidc-psk"); + } + + #[tokio::test] + async fn test_poll_openid_428_then_success() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/finish")) + .respond_with(ResponseTemplate::new(428)) + .up_to_n_times(2) + .mount(&server) + .await; + + let success_body = finish_response_json("oidc-psk"); + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/finish")) + .respond_with(ResponseTemplate::new(200).set_body_json(&success_body)) + .mount(&server) + .await; + + let url = mock_url(&server); + let cancel = CancellationToken::new(); + let psk = poll_openid_mfa(url, "token".into(), cancel).await.unwrap(); + assert_eq!(psk.preshared_key, "oidc-psk"); + } + + #[tokio::test] + async fn test_poll_openid_stops_on_error() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/finish")) + .respond_with(ResponseTemplate::new(500).set_body_json(json!({ "error": "boom" }))) + .mount(&server) + .await; + + let url = mock_url(&server); + let cancel = CancellationToken::new(); + let err = poll_openid_mfa(url, "token".into(), cancel) + .await + .unwrap_err(); + match err { + MfaError::ProxyError { status, message } => { + assert_eq!(status, 500); + assert!(message.contains("boom")); + } + other => panic!("expected ProxyError, got {other:?}"), + } + } + + #[tokio::test] + async fn test_poll_openid_timeout() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/finish")) + .respond_with(ResponseTemplate::new(428)) + .mount(&server) + .await; + + let url = mock_url(&server); + let cancel = CancellationToken::new(); + let err = poll_openid_mfa(url, "token".into(), cancel) + .await + .unwrap_err(); + assert!(matches!(err, MfaError::Timeout)); + } + + #[tokio::test] + async fn test_poll_openid_cancelled() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/api/v1/client-mfa/finish")) + .respond_with(ResponseTemplate::new(428)) + .mount(&server) + .await; + + let url = mock_url(&server); + let cancel = CancellationToken::new(); + cancel.cancel(); + let err = poll_openid_mfa(url, "token".into(), cancel) + .await + .unwrap_err(); + assert!(matches!(err, MfaError::Cancelled)); + } + + #[tokio::test] + async fn test_mobile_approve_success() { + let stub = start_ws_stub().await; + let addr = stub.addr; + let tx = stub.tx; + let ws_url = format!("ws://{addr}/test"); + + let cancel = CancellationToken::new(); + let handle = tokio::spawn(async move { connect_mobile_approve(&ws_url, cancel).await }); + + tx.send(WsStubCommand::SendMessage( + r#"{"type":"mfa_success","preshared_key":"mobile-psk"}"#.into(), + )) + .unwrap(); + tx.send(WsStubCommand::Close).unwrap(); + + let psk = handle.await.unwrap().unwrap(); + assert_eq!(psk.preshared_key, "mobile-psk"); + } + + #[tokio::test] + async fn test_mobile_approve_close_without_success() { + let stub = start_ws_stub().await; + let addr = stub.addr; + let tx = stub.tx; + let ws_url = format!("ws://{addr}/test"); + + let cancel = CancellationToken::new(); + let handle = tokio::spawn(async move { connect_mobile_approve(&ws_url, cancel).await }); + + tx.send(WsStubCommand::Close).unwrap(); + + let err = handle.await.unwrap().unwrap_err(); + assert!(matches!(err, MfaError::MfaRejected { .. })); + } + + #[tokio::test] + async fn test_mobile_approve_cancelled() { + let stub = start_ws_stub().await; + let addr = stub.addr; + let ws_url = format!("ws://{addr}/test"); + + let cancel = CancellationToken::new(); + cancel.cancel(); + let err = connect_mobile_approve(&ws_url, cancel).await.unwrap_err(); + assert!(matches!(err, MfaError::Cancelled)); + } + + #[tokio::test] + async fn test_mobile_approve_connect_error_does_not_leak_token() { + // Nothing is listening, so the WebSocket connect fails. The error must + // be a NetworkError whose message never contains the MFA token (the + // token rides in the ws_url query string). + let base: Url = "http://127.0.0.1:1".parse().unwrap(); + let token = "super-secret-mfa-token"; + let ws_url = derive_ws_url(&base, token).unwrap(); + + let cancel = CancellationToken::new(); + let err = connect_mobile_approve(&ws_url, cancel).await.unwrap_err(); + + assert!(matches!(err, MfaError::NetworkError { .. })); + assert!( + !err.to_string().contains(token), + "error leaked the MFA token: {err}" + ); + } + + #[test] + fn test_derive_ws_url_http_to_ws() { + let base = Url::parse("http://proxy.example.com/").unwrap(); + let ws = derive_ws_url(&base, "tok").unwrap(); + assert!(ws.starts_with("ws://proxy.example.com/api/v1/client-mfa/remote")); + assert!(ws.contains("token=tok")); + } + + #[test] + fn test_derive_ws_url_https_to_wss() { + let base = Url::parse("https://proxy.example.com/").unwrap(); + let ws = derive_ws_url(&base, "tok").unwrap(); + assert!(ws.starts_with("wss://proxy.example.com/api/v1/client-mfa/remote")); + } + + #[test] + fn test_derive_ws_url_preserves_path_prefix() { + let base = Url::parse("https://proxy.example.com/defguard/").unwrap(); + let ws = derive_ws_url(&base, "tok").unwrap(); + assert!(ws.starts_with("wss://proxy.example.com/defguard/api/v1/client-mfa/remote")); + } + + #[test] + fn test_derive_ws_url_rejects_non_http_scheme() { + let base = Url::parse("ftp://proxy.example.com/").unwrap(); + let err = derive_ws_url(&base, "tok").unwrap_err(); + assert!(matches!(err, MfaError::Other { .. })); + } +} diff --git a/src-tauri/core/src/test_helpers.rs b/src-tauri/core/src/test_helpers.rs new file mode 100644 index 000000000..bf7061dd8 --- /dev/null +++ b/src-tauri/core/src/test_helpers.rs @@ -0,0 +1,77 @@ +//! Shared test helpers for defguard-client-core tests. +//! +//! Provides a controllable WebSocket stub for MFA mobile-approve tests, +//! eliminating the need for Docker or external services. + +use std::net::SocketAddr; + +use futures_util::{SinkExt, StreamExt}; +use tokio::{ + net::TcpListener, + sync::mpsc::{unbounded_channel, UnboundedSender}, +}; +use tokio_tungstenite::{accept_async, tungstenite::Message}; + +/// Command to control the WebSocket stub's behavior after a client connects. +pub enum WsStubCommand { + /// Send a text frame to the connected client. + SendMessage(String), + /// Close the WebSocket connection gracefully. + Close, +} + +/// A controllable WebSocket stub for testing MFA mobile-approve flows. +/// +/// Binds to a random port on localhost. The test connects to [`WebSocketStub::addr`], +/// then sends [`WsStubCommand`] values through [`WebSocketStub::tx`] to control +/// what frames the stub emits. +pub struct WebSocketStub { + pub addr: SocketAddr, + pub tx: UnboundedSender, +} + +/// Start a controllable WebSocket stub on a random localhost port. +/// +/// The returned [`WebSocketStub`] spawns a Tokio task that accepts exactly one +/// TCP connection and upgrades it to a WebSocket. After the upgrade, the task +/// waits for commands on the returned `tx` sender. +/// +/// # Panics +/// +/// Panics if the Tokio runtime is not available. +pub async fn start_ws_stub() -> WebSocketStub { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("Failed to bind WebSocket stub"); + let addr = listener.local_addr().expect("Failed to get local address"); + let (tx, mut rx) = unbounded_channel::(); + + tokio::spawn(async move { + // Accept a single connection. + let (stream, _peer) = match listener.accept().await { + Ok(conn) => conn, + Err(_) => return, + }; + + let ws_stream = match accept_async(stream).await { + Ok(ws) => ws, + Err(_) => return, + }; + + let (mut write, mut _read) = ws_stream.split(); + + while let Some(cmd) = rx.recv().await { + match cmd { + WsStubCommand::SendMessage(text) => { + let _ = write.send(Message::Text(text.into())).await; + } + WsStubCommand::Close => { + let _ = write.close().await; + return; + } + } + } + }); + + WebSocketStub { addr, tx } +} diff --git a/src-tauri/permissions/default.toml b/src-tauri/permissions/default.toml index 18b028e8e..305b21e8a 100644 --- a/src-tauri/permissions/default.toml +++ b/src-tauri/permissions/default.toml @@ -2,6 +2,18 @@ identifier = "allow-app-commands" description = "Allow all application commands for both UI windows (old-ui and new-ui)." commands.allow = [ + "enrollment_start", + "enrollment_create_device", + "enrollment_activate_user", + "enrollment_register_mfa_start", + "enrollment_register_mfa_finish", + "enrollment_network_info", + "enrollment_finish", + "mfa_start", + "mfa_finish_code", + "mfa_poll_openid", + "mfa_connect_mobile_approve", + "cancel_mfa", "all_locations", "has_any_visible_locations", "save_device_config", diff --git a/src-tauri/src/appstate.rs b/src-tauri/src/appstate.rs index ec013e1cb..4e0416cc8 100644 --- a/src-tauri/src/appstate.rs +++ b/src-tauri/src/appstate.rs @@ -1,11 +1,15 @@ use std::{collections::HashMap, sync::Mutex}; -use defguard_client_core::connection::active_connections::ACTIVE_CONNECTIONS; +use defguard_client_core::{ + connection::active_connections::ACTIVE_CONNECTIONS, enrollment::EnrollmentSession, +}; +use defguard_client_provisioning::ProvisioningConfig; use tauri::{ async_runtime::{spawn, JoinHandle}, PhysicalPosition, }; use tokio_util::sync::CancellationToken; +use uuid::Uuid; use crate::{ app_config::AppConfig, @@ -14,10 +18,11 @@ use crate::{ utils::stats_handler, ConnectionType, }; -use defguard_client_provisioning::ProvisioningConfig; pub struct AppState { + pub enrollment_sessions: Mutex>, pub log_watchers: Mutex>, + pub mfa_tasks: Mutex>, pub app_config: Mutex, pub tray_click_position: Mutex>>, stat_threads: Mutex>>, // location ID is the key @@ -29,7 +34,9 @@ impl AppState { #[must_use] pub fn new(config: AppConfig, provisioning_config: Option) -> Self { Self { + enrollment_sessions: Mutex::new(HashMap::new()), log_watchers: Mutex::new(HashMap::new()), + mfa_tasks: Mutex::new(HashMap::new()), app_config: Mutex::new(config), tray_click_position: Mutex::new(None), stat_threads: Mutex::new(HashMap::new()), diff --git a/src-tauri/src/bin/defguard-client.rs b/src-tauri/src/bin/defguard-client.rs index 931d18e03..1157e3cd1 100644 --- a/src-tauri/src/bin/defguard-client.rs +++ b/src-tauri/src/bin/defguard-client.rs @@ -207,6 +207,18 @@ fn main() { close_tray_window, all_active_connections, disconnect_locations, + enrollment_start, + enrollment_create_device, + enrollment_activate_user, + enrollment_register_mfa_start, + enrollment_register_mfa_finish, + enrollment_network_info, + enrollment_finish, + mfa_start, + mfa_finish_code, + mfa_poll_openid, + mfa_connect_mobile_approve, + cancel_mfa, session_state::get_session_state, session_state::patch_session_state, ]) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 36847b72c..fd0d58816 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -4,9 +4,13 @@ use std::{collections::HashMap, env, str::FromStr}; use chrono::{DateTime, Duration, Utc}; #[cfg(not(target_os = "macos"))] use defguard_client_core::connection::daemon_client::DAEMON_CLIENT; -use defguard_client_core::connection::{ - active_connections::{find_connection, get_connection_id_by_type, ACTIVE_CONNECTIONS}, - disconnect_interface, +use defguard_client_core::{ + connection::{ + active_connections::{find_connection, get_connection_id_by_type, ACTIVE_CONNECTIONS}, + disconnect_interface, + }, + enrollment::{self, MfaFinishResponse, MfaStartResponse}, + mfa, }; use defguard_client_posture::authorize_posture_session; #[cfg(not(target_os = "macos"))] @@ -14,14 +18,22 @@ use defguard_client_proto::defguard::client::v1::{ DeleteServiceLocationsRequest, RemoveInterfaceRequest, SaveServiceLocationsRequest, }; use defguard_client_proto::defguard::{ - client_types::DeviceConfigResponse, enterprise::posture::v2::DevicePostureData, + client_types::{ + AdminInfo, ClientMfaFinishRequest, ClientMfaFinishResponse, ClientMfaStartRequest, + DeviceConfigResponse, EnrollmentSettings, InitialUserInfo, + InstanceInfo as ProtoInstanceInfo, MfaMethod, + }, + enterprise::posture::v2::DevicePostureData, }; use defguard_client_provisioning::ProvisioningConfig; #[cfg(not(target_os = "macos"))] use defguard_client_service_locations::to_service_location; +use reqwest::Url; use serde::{Deserialize, Serialize}; use struct_patch::Patch; use tauri::{AppHandle, Emitter, Manager, State}; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; const UPDATE_URL: &str = "https://pkgs.defguard.net/api/update/check"; @@ -85,12 +97,49 @@ impl From for ConnectError { } } +/// Serialize a structured error (e.g. `MfaError`, `EnrollmentError`) to JSON so +/// the frontend can match on its tagged `type`, falling back to the Display +/// string if serialization somehow fails. +fn err_to_json(e: E) -> String { + serde_json::to_string(&e).unwrap_or_else(|_| e.to_string()) +} + +/// Look up a cloned enrollment session by its opaque string id. Used by the +/// enrollment commands that need read access to the in-memory session. +fn get_enrollment_session( + state: &AppState, + session_id: &str, +) -> Result { + let uid = Uuid::parse_str(session_id).map_err(|e| format!("Invalid session ID: {e}"))?; + state + .enrollment_sessions + .lock() + .expect("enrollment_sessions mutex poisoned") + .get(&uid) + .cloned() + .ok_or_else(|| "Enrollment session not found".to_string()) +} + +/// Bring up a location connection with an already-obtained preshared key and +/// refresh the tray. Shared by `connect` and the MFA finish flows so the +/// preshared key never has to cross back into the frontend. +async fn connect_location_with_psk( + location: Location, + preshared_key: Option, + handle: &AppHandle, +) -> Result<(), Error> { + handle_connection_for_location(location.clone(), preshared_key, handle).await?; + reload_tray_menu(handle).await; + info!("Connected to location {location}"); + configure_tray_icon(handle).await?; + Ok(()) +} + /// Open new WireGuard connection. #[tauri::command(async)] pub async fn connect( location_id: Id, connection_type: ConnectionType, - preshared_key: Option, handle: AppHandle, ) -> Result<(), ConnectError> { debug!("Received a command to connect to a {connection_type} with ID {location_id}"); @@ -100,14 +149,15 @@ pub async fn connect( "Identified location with ID {location_id} as \"{}\", handling connection.", location.name ); - let preshared_key = if location.posture_check_required && preshared_key.is_none() { + // Connect-time MFA brings the tunnel up itself (keeping the preshared + // key backend-side), so the only preshared key resolved here is for + // posture-only locations. + let preshared_key = if location.posture_check_required { Some(authorize_posture_session(&location).await?) } else { - preshared_key + None }; - handle_connection_for_location(location.clone(), preshared_key, &handle).await?; - reload_tray_menu(&handle).await; - info!("Connected to location {location}"); + connect_location_with_psk(location, preshared_key, &handle).await?; } else { error!( "Location with ID {location_id} not found in the database, aborting connection \ @@ -122,14 +172,13 @@ pub async fn connect( ); handle_connection_for_tunnel(tunnel.clone(), &handle).await?; info!("Successfully connected to tunnel {tunnel}"); + // Update tray icon to reflect connection state. + configure_tray_icon(&handle).await?; } else { error!("Tunnel {location_id} not found"); return Err(Error::NotFound.into()); } - // Update tray icon to reflect connection state. - configure_tray_icon(&handle).await?; - Ok(()) } @@ -350,9 +399,9 @@ pub async fn save_device_config( debug!("Saving device configuration: {response:#?}."); let mut transaction = DB_POOL.begin().await?; - let instance_info = response - .instance - .expect("Missing instance info in device config response"); + let instance_info = response.instance.ok_or_else(|| { + Error::ResourceNotFound("instance info in device config response".to_string()) + })?; let mut instance = Instance::from(instance_info); if response.token.is_some() { debug!( @@ -372,9 +421,9 @@ pub async fn save_device_config( let instance = instance.save(&mut *transaction).await?; debug!("Saved instance {}", instance.name); - let device = response - .device - .expect("Missing device info in device config response"); + let device = response.device.ok_or_else(|| { + Error::ResourceNotFound("device info in device config response".to_string()) + })?; let keys = WireguardKeys::new(instance.id, device.pubkey, private_key); debug!( "Saving wireguard key {} for instance {}({})", @@ -1326,3 +1375,385 @@ pub async fn all_active_connections() -> Result, Er debug!("Returning {} active connections.", result.len()); Ok(result) } + +/// Returned by the `enrollment_start` Tauri command. +#[derive(Clone, Debug, Serialize)] +pub struct EnrollmentStartResult { + pub session_id: String, + pub user: InitialUserInfo, + pub admin: AdminInfo, + pub settings: EnrollmentSettings, + pub instance: ProtoInstanceInfo, + pub deadline_timestamp: i64, + pub final_page_content: String, +} + +#[tauri::command(async)] +pub async fn enrollment_start( + proxy_url: String, + token: String, + state: State<'_, AppState>, +) -> Result { + debug!("Starting enrollment at {proxy_url}"); + let url = Url::parse(&proxy_url).map_err(|e| format!("Invalid proxy URL: {e}"))?; + let (session, response) = enrollment::enrollment_start(url, token) + .await + .map_err(err_to_json)?; + let session_uuid = Uuid::new_v4(); + let session_id = session_uuid.to_string(); + state + .enrollment_sessions + .lock() + .expect("enrollment_sessions mutex poisoned") + .insert(session_uuid, session); + let login = response + .user + .as_ref() + .map(|u| u.login.as_str()) + .unwrap_or(""); + info!("Enrollment started for user {login}, session {session_id}"); + Ok(EnrollmentStartResult { + session_id, + user: response + .user + .ok_or_else(|| "Proxy did not return user info".to_string())?, + admin: response + .admin + .ok_or_else(|| "Proxy did not return admin info".to_string())?, + settings: response + .settings + .ok_or_else(|| "Proxy did not return enrollment settings".to_string())?, + instance: response + .instance + .ok_or_else(|| "Proxy did not return instance info".to_string())?, + deadline_timestamp: response.deadline_timestamp, + final_page_content: response.final_page_content, + }) +} + +#[tauri::command(async)] +pub async fn enrollment_create_device( + session_id: String, + name: String, + pubkey: String, + state: State<'_, AppState>, +) -> Result { + debug!("Creating device \"{name}\""); + let session = get_enrollment_session(&state, &session_id)?; + let result = enrollment::enrollment_create_device(session, name, pubkey) + .await + .map_err(err_to_json)?; + info!("Device created"); + Ok(result) +} + +#[tauri::command(async)] +pub async fn enrollment_activate_user( + session_id: String, + password: Option, + phone_number: Option, + state: State<'_, AppState>, +) -> Result<(), String> { + debug!("Activating user"); + let session = get_enrollment_session(&state, &session_id)?; + enrollment::enrollment_activate_user(session, password, phone_number) + .await + .map_err(err_to_json)?; + info!("User activated"); + Ok(()) +} + +#[tauri::command(async)] +pub async fn enrollment_register_mfa_start( + session_id: String, + method: String, + state: State<'_, AppState>, +) -> Result { + debug!("Starting MFA setup"); + let session = get_enrollment_session(&state, &session_id)?; + enrollment::enrollment_register_mfa_start(session, method) + .await + .map_err(err_to_json) +} + +#[tauri::command(async)] +pub async fn enrollment_register_mfa_finish( + session_id: String, + code: String, + method: String, + state: State<'_, AppState>, +) -> Result { + debug!("Finishing MFA setup"); + let session = get_enrollment_session(&state, &session_id)?; + enrollment::enrollment_register_mfa_finish(session, code, method) + .await + .map_err(err_to_json) +} + +#[tauri::command(async)] +pub async fn enrollment_network_info( + session_id: String, + pubkey: String, + state: State<'_, AppState>, +) -> Result { + debug!("Fetching network info"); + let session = get_enrollment_session(&state, &session_id)?; + enrollment::enrollment_network_info(session, pubkey) + .await + .map_err(err_to_json) +} + +#[tauri::command(async)] +pub async fn enrollment_finish( + session_id: String, + state: State<'_, AppState>, +) -> Result<(), String> { + debug!("Finishing enrollment"); + let session = { + let uid = Uuid::parse_str(&session_id).map_err(|e| format!("Invalid session ID: {e}"))?; + let mut sessions = state + .enrollment_sessions + .lock() + .expect("enrollment_sessions mutex poisoned"); + sessions + .remove(&uid) + .ok_or_else(|| "Enrollment session not found".to_string())? + }; + enrollment::enrollment_finish(session); + info!("Enrollment finished, session {session_id} removed"); + Ok(()) +} + +#[derive(Clone, Serialize)] +pub struct MfaErrorPayload { + pub error: String, +} + +/// Bring up a location connection with a preshared key obtained from a +/// completed MFA handshake. Keeps the preshared key inside the backend - it is +/// never returned to or emitted at the frontend. +async fn connect_after_mfa( + location_id: Id, + preshared_key: String, + handle: &AppHandle, +) -> Result<(), String> { + let location = Location::find_by_id(&*DB_POOL, location_id) + .await + .map_err(|e| e.to_string())? + .ok_or_else(|| "Location not found".to_string())?; + connect_location_with_psk(location, Some(preshared_key), handle) + .await + // Distinct prefix so the frontend can tell a post-MFA connection + // failure apart from an MFA/auth failure. + .map_err(|e| format!("VPN connection failed: {e}")) +} + +/// Map the frontend MFA method string to the proto `MfaMethod` enum the proxy +/// expects on the wire (a numeric enum, not a string). +fn parse_mfa_method(method: &str) -> Result { + match method { + "totp" => Ok(MfaMethod::Totp), + "email" => Ok(MfaMethod::Email), + "oidc" => Ok(MfaMethod::Oidc), + "biometric" => Ok(MfaMethod::Biometric), + "mobileapprove" => Ok(MfaMethod::MobileApprove), + other => Err(format!("Unsupported MFA method: {other}")), + } +} + +#[tauri::command(async)] +pub async fn mfa_start( + instance_id: Id, + location_id: Id, + method: String, +) -> Result { + debug!("Starting MFA session for location {location_id}"); + let method = parse_mfa_method(&method)?; + let instance = Instance::find_by_id(&*DB_POOL, instance_id) + .await + .map_err(|e| e.to_string())? + .ok_or_else(|| "Instance not found".to_string())?; + let keys = WireguardKeys::find_by_instance_id(&*DB_POOL, instance_id) + .await + .map_err(|e| e.to_string())? + .ok_or_else(|| "WireGuard keys not found".to_string())?; + let proxy_url = + Url::parse(&instance.proxy_url).map_err(|e| format!("Invalid proxy URL: {e}"))?; + let location = Location::find_by_id(&*DB_POOL, location_id) + .await + .map_err(|e| e.to_string())? + .ok_or_else(|| "Location not found".to_string())?; + let posture_data = if location.posture_check_required { + Some( + defguard_client_posture::get_posture_data() + .await + .map_err(|e| format!("Failed to collect posture data: {e}"))?, + ) + } else { + None + }; + let request = ClientMfaStartRequest { + location_id: location.network_id, + pubkey: keys.pubkey, + method: method as i32, + posture_data, + }; + mfa::mfa_start(proxy_url, request) + .await + .map_err(err_to_json) +} + +#[tauri::command(async)] +pub async fn mfa_finish_code( + instance_id: Id, + location_id: Id, + token: String, + code: String, + handle: AppHandle, +) -> Result<(), String> { + debug!("Finishing MFA with code for instance {instance_id}"); + let instance = Instance::find_by_id(&*DB_POOL, instance_id) + .await + .map_err(|e| e.to_string())? + .ok_or_else(|| "Instance not found".to_string())?; + let proxy_url = + Url::parse(&instance.proxy_url).map_err(|e| format!("Invalid proxy URL: {e}"))?; + let request = ClientMfaFinishRequest { + token, + code: Some(code), + auth_pub_key: None, + }; + let response = mfa::mfa_finish_code(proxy_url, request) + .await + .map_err(err_to_json)?; + connect_after_mfa(location_id, response.preshared_key, &handle).await +} + +/// Register a long-running MFA task, run its future in the background, and on +/// success bring up the connection Rust-side before emitting a payload-free +/// completion event (or an error event). Shared by the OpenID poll and mobile +/// approve flows so the preshared key never leaves the backend. Returns the +/// task id the frontend uses to cancel. +fn spawn_mfa_task( + handle: &AppHandle, + location_id: Id, + complete_event: EventKey, + error_event: EventKey, + run: R, +) -> String +where + R: FnOnce(CancellationToken) -> F + Send + 'static, + F: std::future::Future> + + Send + + 'static, +{ + let cancel = CancellationToken::new(); + let task_id = Uuid::new_v4().to_string(); + handle + .state::() + .mfa_tasks + .lock() + .expect("mfa_tasks mutex poisoned") + .insert(task_id.clone(), cancel.clone()); + + let task_id_for_task = task_id.clone(); + let listen_handle = handle.clone(); + tokio::spawn(async move { + let result = run(cancel).await; + listen_handle + .state::() + .mfa_tasks + .lock() + .expect("mfa_tasks mutex poisoned") + .remove(&task_id_for_task); + match result { + Ok(response) => { + info!("MFA completed for task {task_id_for_task}"); + match connect_after_mfa(location_id, response.preshared_key, &listen_handle).await { + Ok(()) => { + let _ = listen_handle.emit(complete_event.into(), ()); + } + Err(err) => { + warn!("Connect after MFA failed for task {task_id_for_task}: {err}"); + let _ = + listen_handle.emit(error_event.into(), MfaErrorPayload { error: err }); + } + } + } + Err(err) => { + warn!("MFA task {task_id_for_task} failed: {err}"); + // Emit the structured error as JSON so the frontend classifies + // it the same way as command errors. + let _ = listen_handle.emit( + error_event.into(), + MfaErrorPayload { + error: err_to_json(err), + }, + ); + } + } + }); + + task_id +} + +#[tauri::command(async)] +pub async fn mfa_poll_openid( + instance_id: Id, + location_id: Id, + token: String, + handle: AppHandle, +) -> Result { + debug!("Starting OpenID MFA poll for instance {instance_id}"); + let instance = Instance::find_by_id(&*DB_POOL, instance_id) + .await + .map_err(|e| e.to_string())? + .ok_or_else(|| "Instance not found".to_string())?; + let proxy_url = + Url::parse(&instance.proxy_url).map_err(|e| format!("Invalid proxy URL: {e}"))?; + Ok(spawn_mfa_task( + &handle, + location_id, + EventKey::MfaOpenIdComplete, + EventKey::MfaOpenIdError, + move |cancel| mfa::poll_openid_mfa(proxy_url, token, cancel), + )) +} + +#[tauri::command(async)] +pub async fn mfa_connect_mobile_approve( + instance_id: Id, + location_id: Id, + token: String, + handle: AppHandle, +) -> Result { + debug!("Starting mobile approve MFA for instance {instance_id}"); + let instance = Instance::find_by_id(&*DB_POOL, instance_id) + .await + .map_err(|e| e.to_string())? + .ok_or_else(|| "Instance not found".to_string())?; + let proxy_url = + Url::parse(&instance.proxy_url).map_err(|e| format!("Invalid proxy URL: {e}"))?; + let ws_url = mfa::derive_ws_url(&proxy_url, &token).map_err(|e| e.to_string())?; + Ok(spawn_mfa_task( + &handle, + location_id, + EventKey::MfaMobileComplete, + EventKey::MfaMobileError, + move |cancel| async move { mfa::connect_mobile_approve(&ws_url, cancel).await }, + )) +} + +#[tauri::command(async)] +pub async fn cancel_mfa(task_id: String, state: State<'_, AppState>) -> Result<(), String> { + debug!("Cancelling MFA task {task_id}"); + let cancel = { + let tasks = state.mfa_tasks.lock().expect("mfa_tasks mutex poisoned"); + tasks + .get(&task_id) + .cloned() + .ok_or_else(|| "MFA task not found".to_string())? + }; + cancel.cancel(); + Ok(()) +} diff --git a/src-tauri/src/periodic/config.rs b/src-tauri/src/periodic/config.rs index fd9134df1..30672e6c2 100644 --- a/src-tauri/src/periodic/config.rs +++ b/src-tauri/src/periodic/config.rs @@ -4,6 +4,9 @@ use std::{ time::Duration, }; +pub use defguard_client_config_sync::commands::{ + disable_enterprise_features, do_update_instance, locations_changed, +}; use defguard_client_config_sync::{ poll_instance, poll_instances, PollInstanceResult, VersionMismatchPayload, }; @@ -22,10 +25,6 @@ use sqlx::{Sqlite, Transaction}; use tauri::{AppHandle, Emitter}; use tokio::time::sleep; -pub use defguard_client_config_sync::commands::{ - disable_enterprise_features, do_update_instance, locations_changed, -}; - const INTERVAL_SECONDS: Duration = Duration::from_secs(30); /// Tracks instance IDs for which we already sent a version-mismatch notification, diff --git a/src-tauri/src/periodic/connection.rs b/src-tauri/src/periodic/connection.rs index a9f4c102d..96ebd4418 100644 --- a/src-tauri/src/periodic/connection.rs +++ b/src-tauri/src/periodic/connection.rs @@ -49,7 +49,7 @@ async fn reconnect( peer_alive_period: peer_alive_period.num_seconds(), }; payload.emit(app_handle); - match connect(con_id, con_type, None, app_handle.clone()).await { + match connect(con_id, con_type, app_handle.clone()).await { Ok(()) => { info!("Reconnect for {con_type} {con_interface_name} ({con_id}) succeeded."); } diff --git a/src-tauri/src/session_state.rs b/src-tauri/src/session_state.rs index ab9d2d6eb..09cc04385 100644 --- a/src-tauri/src/session_state.rs +++ b/src-tauri/src/session_state.rs @@ -1,11 +1,10 @@ use std::collections::HashMap; +use defguard_client_core::events::EventKey; use serde::{Deserialize, Serialize}; use struct_patch::Patch; use tauri::{AppHandle, Emitter, Manager, State}; -use defguard_client_core::events::EventKey; - use crate::appstate::AppState; #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] diff --git a/src-tauri/src/tray.rs b/src-tauri/src/tray.rs index f7d3f4591..3d468f75d 100644 --- a/src-tauri/src/tray.rs +++ b/src-tauri/src/tray.rs @@ -328,7 +328,7 @@ async fn handle_location_tray_menu(id: String, app: &AppHandle) { show_main_window(app); let _ = app.emit(EventKey::MfaTrigger.into(), &location); } else if let Err(err) = - connect(location_id, ConnectionType::Location, None, app.clone()).await + connect(location_id, ConnectionType::Location, app.clone()).await { info!("Unable to connect location with ID {id}, error: {err:?}"); } diff --git a/src-tauri/src/utils.rs b/src-tauri/src/utils.rs index 0d3a118f1..8056ddc08 100644 --- a/src-tauri/src/utils.rs +++ b/src-tauri/src/utils.rs @@ -8,6 +8,8 @@ use std::{path::Path, str::FromStr}; use defguard_client_common::{find_free_tcp_port, get_interface_name}; #[cfg(windows)] use defguard_client_core::connection::active_connections::find_connection; +#[cfg(target_os = "macos")] +use defguard_client_core::connection::apple::tunnel_stats; use defguard_client_core::connection::{bring_up, ConnectionTarget}; #[cfg(not(target_os = "macos"))] use defguard_client_core::{ @@ -46,8 +48,6 @@ use crate::{ log_watcher::service_log_watcher::spawn_log_watcher_task, ConnectionType, }; -#[cfg(target_os = "macos")] -use defguard_client_core::connection::apple::tunnel_stats; // Work-around MFA propagation delay. FIXME: remove once Core API is corrected. #[cfg(target_os = "macos")]