4 files changed, 15 insertions, 8 deletions
diff --git a/gui/src/renderer/components/AdvancedSettings.tsx b/gui/src/renderer/components/AdvancedSettings.tsx
index 3f3440242d..42f48b1b87 100644
--- a/gui/src/renderer/components/AdvancedSettings.tsx
+++ b/gui/src/renderer/components/AdvancedSettings.tsx
@@ -25,6 +25,7 @@ import {
 } from './NavigationBar';
 import { ISelectorItem } from './cell/Selector';
 import SettingsHeader, { HeaderTitle } from './SettingsHeader';
+import Switch from './Switch';
 
 type OptionalTunnelProtocol = TunnelProtocol | undefined;
 
@@ -51,6 +52,8 @@ export default class AdvancedSettings extends React.Component<IProps, IState> {
     showConfirmBlockWhenDisconnectedAlert: false,
   };
 
+  private blockWhenDisconnectedRef = React.createRef<Switch>();
+
   public render() {
     const hasWireguardKey = this.props.wireguardKeyState.type === 'key-set';
 
@@ -118,6 +121,7 @@ export default class AdvancedSettings extends React.Component<IProps, IState> {
                     </AriaLabel>
                     <AriaInput>
                       <Cell.Switch
+                        ref={this.blockWhenDisconnectedRef}
                         isOn={this.props.blockWhenDisconnected}
                         onChange={this.setBlockWhenDisconnected}
                       />
@@ -257,7 +261,6 @@ export default class AdvancedSettings extends React.Component<IProps, IState> {
 
   private setBlockWhenDisconnected = (newValue: boolean) => {
     if (newValue) {
-      this.props.setBlockWhenDisconnected(true);
       this.setState({ showConfirmBlockWhenDisconnectedAlert: true });
     } else {
       this.props.setBlockWhenDisconnected(false);
@@ -265,13 +268,13 @@ export default class AdvancedSettings extends React.Component<IProps, IState> {
   };
 
   private hideConfirmBlockWhenDisconnectedAlert = () => {
-    this.props.setBlockWhenDisconnected(false);
     this.setState({ showConfirmBlockWhenDisconnectedAlert: false });
+    this.blockWhenDisconnectedRef.current?.setOn(this.props.blockWhenDisconnected);
   };
 
   private confirmEnableBlockWhenDisconnected = () => {
-    this.props.setBlockWhenDisconnected(true);
     this.setState({ showConfirmBlockWhenDisconnectedAlert: false });
+    this.props.setBlockWhenDisconnected(true);
   };
 
   private onSelectTunnelProtocol = (protocol?: TunnelProtocol) => {
diff --git a/gui/src/renderer/components/CustomDnsSettings.tsx b/gui/src/renderer/components/CustomDnsSettings.tsx
index c0d1a66f5c..d70b9d0eba 100644
--- a/gui/src/renderer/components/CustomDnsSettings.tsx
+++ b/gui/src/renderer/components/CustomDnsSettings.tsx
@@ -194,7 +194,7 @@ export default function CustomDnsSettings() {
           </AriaLabel>
           <AriaInput>
             <Cell.Switch
-              ref={switchRef}
+              innerRef={switchRef}
               isOn={dns.state === 'custom' || inputVisible}
               onChange={setCustomDnsEnabled}
             />
diff --git a/gui/src/renderer/components/Switch.tsx b/gui/src/renderer/components/Switch.tsx
index e3b722c721..c6464510f1 100644
--- a/gui/src/renderer/components/Switch.tsx
+++ b/gui/src/renderer/components/Switch.tsx
@@ -11,7 +11,7 @@ interface IProps {
   onChange?: (isOn: boolean) => void;
   className?: string;
   disabled?: boolean;
-  forwardedRef?: React.Ref<HTMLDivElement>;
+  innerRef?: React.Ref<HTMLDivElement>;
 }
 
 interface IState {
@@ -103,9 +103,13 @@ export default class Switch extends React.PureComponent<IProps, IState> {
     );
   }
 
+  public setOn(isOn: boolean) {
+    this.setState({ isOn });
+  }
+
   private refCallback = (element: HTMLDivElement | null) => {
     assignToRef(element, this.containerRef);
-    assignToRef(element, this.props.forwardedRef);
+    assignToRef(element, this.props.innerRef);
   };
 
   private onTransitionEnd = (event: React.TransitionEvent) => {
diff --git a/gui/src/renderer/components/cell/Input.tsx b/gui/src/renderer/components/cell/Input.tsx
index 1dd240b3ce..c57d552866 100644
--- a/gui/src/renderer/components/cell/Input.tsx
+++ b/gui/src/renderer/components/cell/Input.tsx
@@ -8,10 +8,10 @@ import ImageView from '../ImageView';
 
 export const Switch = React.forwardRef(function SwitchT(
   props: StandaloneSwitch['props'],
-  ref: React.Ref<HTMLDivElement>,
+  ref: React.Ref<StandaloneSwitch>,
 ) {
   const disabled = useContext(CellDisabledContext);
-  return <StandaloneSwitch forwardedRef={ref} disabled={disabled} {...props} />;
+  return <StandaloneSwitch ref={ref} disabled={disabled} {...props} />;
 });
 
 export const InputFrame = styled.div({