From 0a202004c432b08a17ca5d0f43e087c3097a7f14 Mon Sep 17 00:00:00 2001
From: rui hildt <rui@mullvad.net>
Date: Thu, 1 Feb 2024 18:41:32 +0100
Subject: [PATCH] Fix selection of proxy server per host

---
 .../ConnectionDetails/ConnectionDetails.vue   |  6 +-
 src/components/Location.vue                   |  4 +-
 src/components/ProxyStatus/ProxyHost.vue      | 57 ++++----------
 src/composables/useSocksProxy.ts              | 75 +++++++++++--------
 src/helpers/socksProxy.ts                     | 23 +++---
 src/popup/App.vue                             | 16 +---
 src/popup/views/Home.vue                      | 20 ++++-
 7 files changed, 97 insertions(+), 104 deletions(-)

diff --git a/src/components/ConnectionDetails/ConnectionDetails.vue b/src/components/ConnectionDetails/ConnectionDetails.vue
index ba855f60..38ab9a11 100644
--- a/src/components/ConnectionDetails/ConnectionDetails.vue
+++ b/src/components/ConnectionDetails/ConnectionDetails.vue
@@ -14,7 +14,7 @@ import DnsLeakStatus from '@/components/ConnectionStatus/DnsLeakStatus.vue';
 import { ConnectionKey, defaultConnection } from '@/composables/useConnection';
 import useSocksProxy from '@/composables/useSocksProxy';
 
-const { globalProxyEnabled, activeTabProxyEnabled } = useSocksProxy();
+const { globalProxyEnabled } = useSocksProxy();
 
 const { isLoading, isError, connection } = inject(ConnectionKey, defaultConnection);
 const connected = computed(() => connection.value.isMullvad);
@@ -34,7 +34,7 @@ const connected = computed(() => connection.value.isMullvad);
         <IconLabel text="Checking connection" type="spinner" />
       </p>
 
-      <div v-if="globalProxyEnabled || activeTabProxyEnabled">
+      <div v-if="globalProxyEnabled">
         <p>The connection is set to use a proxy server.</p>
         <p>If you can't load any pages, try disabling the proxy to access the internet directly</p>
         <router-link to="/proxy">
@@ -45,7 +45,7 @@ const connected = computed(() => connection.value.isMullvad);
 
     <div v-else-if="isError">
       <IconLabel text="Couldn't get connection details" type="warning" />
-      <div v-if="globalProxyEnabled || activeTabProxyEnabled">
+      <div v-if="globalProxyEnabled">
         <p>The connection is set to use a proxy server.</p>
         <p>If you can't load any pages, try disabling the proxy to access the internet directly</p>
         <router-link to="/proxy">
diff --git a/src/components/Location.vue b/src/components/Location.vue
index 1ed8c38d..acbd90a1 100644
--- a/src/components/Location.vue
+++ b/src/components/Location.vue
@@ -15,7 +15,7 @@ import type { HistoricConnection } from '@/composables/useHistoricConnections/Hi
 const { activeTabHost } = useActiveTab();
 const { hostProxySelect, toggleLocations } = useLocations();
 const { proxiesList } = useListProxies();
-const { setHostProxy, setGlobalProxy } = useSocksProxy();
+const { setCurrentHostProxy, setGlobalProxy } = useSocksProxy();
 const { storeSocksProxyUsage } = useHistoricConnections();
 
 const currentOrAllWebsites = computed(() =>
@@ -33,7 +33,7 @@ const setProxy = (
   toggleLocations();
 
   if (hostProxySelect.value) {
-    setHostProxy({ country, city, hostname, ipv4_address, port }, activeTabHost.value);
+    setCurrentHostProxy({ country, city, hostname, ipv4_address, port }, activeTabHost.value);
   } else {
     setGlobalProxy({ country, city, hostname, ipv4_address, port });
   }
diff --git a/src/components/ProxyStatus/ProxyHost.vue b/src/components/ProxyStatus/ProxyHost.vue
index 7f3a5f2c..5cbf4d9c 100644
--- a/src/components/ProxyStatus/ProxyHost.vue
+++ b/src/components/ProxyStatus/ProxyHost.vue
@@ -1,10 +1,9 @@
 <script lang="ts" setup>
 import { computed } from 'vue';
-import { NButtonGroup, NCard, NTag } from 'naive-ui';
+import { NButtonGroup, NCard, NSwitch, NTooltip } from 'naive-ui';
 
 import Button from '@/components/Buttons/Button.vue';
 import CurrentProxyDetails from '@/components/CurrentProxyDetails.vue';
-import FeCheck from '@/components/Icons/FeCheck.vue';
 import TitleCategory from '@/components/TitleCategory.vue';
 
 import useActiveTab from '@/composables/useActiveTab';
@@ -13,26 +12,18 @@ import useLocations from '@/composables/useLocations';
 import useSocksProxy from '@/composables/useSocksProxy';
 import useStore from '@/composables/useStore';
 
+const { activeTabHost, activeTab } = useActiveTab();
 const { connection } = useConnection();
 const { hostProxySelect, toggleLocations } = useLocations();
-const { activeTabProxyEnabled } = useSocksProxy();
-const { hostProxies, hostProxiesDetails, excludedHosts } = useStore();
+const { currentHostProxyDetails, currentHostProxyEnabled, toggleCurrentHostProxy } =
+  useSocksProxy();
+const { excludedHosts } = useStore();
 
-const { activeTabHost, activeTab } = useActiveTab();
+const title = computed(() => `For this website (${activeTabHost.value})`);
 
-const customTitle = computed(() => `For this website (${activeTabHost.value})`);
-const currentProxyDetails = computed(() => {
-  return hostProxiesDetails.value[activeTabHost.value];
-});
 const isExcluded = computed(() => excludedHosts.value.includes(activeTabHost.value));
 const isWireGuard = computed(() => connection.value.protocol?.includes('WireGuard'));
 
-const removeHostProxy = () => {
-  const updatedHostProxies = { ...hostProxies.value };
-  delete updatedHostProxies[activeTabHost.value];
-  hostProxies.value = updatedHostProxies;
-};
-
 const neverProxyHost = () => {
   excludedHosts.value = [...excludedHosts.value, activeTabHost.value];
 };
@@ -51,7 +42,7 @@ const handleHostProxySelect = () => {
 // watchEffect(() => {
 //   let message = '';
 
-//   if (activeTabProxyEnabled.value) {
+//   if (currentHostProxyEnabled.value) {
 //     message = `Only this website is configured to use the following proxy.`;
 //   } else if (isExcluded.value) {
 //     message = 'The website is configured to never use a proxy.';
@@ -66,20 +57,16 @@ const handleHostProxySelect = () => {
 <template>
   <n-card v-if="activeTab" :bordered="false" class="mb-4">
     <div class="flex justify-between">
-      <TitleCategory :title="customTitle" />
-      <n-tag v-if="activeTabProxyEnabled" round :bordered="false" type="success">
-        active
-        <template #icon>
-          <FeCheck />
+      <TitleCategory :title="title" />
+      <n-tooltip>
+        <template #trigger>
+          <n-switch :value="currentHostProxyEnabled" @update-value="toggleCurrentHostProxy" />
         </template>
-      </n-tag>
-      <div v-else />
+        <span> {{ currentHostProxyEnabled ? 'enabled' : 'disabled' }}</span>
+      </n-tooltip>
     </div>
 
-    <CurrentProxyDetails
-      v-if="activeTabProxyEnabled && currentProxyDetails"
-      :proxy-details="currentProxyDetails"
-    />
+    <CurrentProxyDetails v-if="currentHostProxyDetails" :proxy-details="currentHostProxyDetails" />
 
     <!-- <IconLabel :text="instructions" type="info" class="mb-8 mt-6" /> -->
 
@@ -98,21 +85,7 @@ const handleHostProxySelect = () => {
           Switch location
         </Button>
 
-        <Button
-          v-if="activeTabProxyEnabled"
-          color="error"
-          class="flex items-center justify-center"
-          @click="removeHostProxy"
-        >
-          Disable proxy
-        </Button>
-
-        <Button
-          v-else
-          color="error"
-          class="flex items-center justify-center"
-          @click="neverProxyHost"
-        >
+        <Button color="error" class="flex items-center justify-center" @click="neverProxyHost">
           Never proxy
         </Button>
       </n-button-group>
diff --git a/src/composables/useSocksProxy.ts b/src/composables/useSocksProxy.ts
index 1183c3ae..acd2ff0f 100644
--- a/src/composables/useSocksProxy.ts
+++ b/src/composables/useSocksProxy.ts
@@ -16,10 +16,10 @@ const { activeTabHost } = useActiveTab();
 const { connection, updateConnection } = useConnection();
 const { globalProxy, globalProxyDetails, hostProxies, hostProxiesDetails } = useStore();
 
+const currentHostProxyDetails = computed(() => hostProxiesDetails.value[activeTabHost.value]);
+
 const globalProxyEnabled = computed(() => globalProxyDetails.value.socksEnabled);
-const activeTabProxyEnabled = computed(() =>
-  Object.keys(hostProxies.value).includes(activeTabHost.value),
-);
+const currentHostProxyEnabled = computed(() => currentHostProxyDetails.value.socksEnabled);
 
 const baseConfig: Partial<ProxyInfo> = {
   port: 1080,
@@ -37,33 +37,45 @@ const toggleGlobalProxy = () => {
   setProxyDetails(globalProxyDetails, !globalProxyDetails.value.socksEnabled);
 };
 
-const initGlobalProxy = async (
+const toggleCurrentHostProxy = () => {
+  hostProxiesDetails.value[activeTabHost.value] = {
+    ...hostProxiesDetails.value[activeTabHost.value],
+    socksEnabled: !currentHostProxyEnabled.value,
+  };
+};
+
+const createDefaultProxy = (protocol: string) => {
+  const socksIp = getSocksIpForProtocol(protocol);
+
+  return {
+    host: socksIp,
+    port: baseConfig.port!,
+    proxyDNS: baseConfig.proxyDNS,
+    type: ProxyInfoType.socks,
+  };
+};
+
+const initProxies = async (
   { hostname, country, city }: Partial<ProxyOperationArgs>,
   protocol: string,
 ) => {
-  // Create a default globalProxy if there's none yet
-  // We don't do it it in the background script, because we don't want to make a request on extension start
+  const newProxyDetails = {
+    city,
+    country,
+    protocol,
+    proxyDNS: baseConfig.proxyDNS,
+    server: hostname,
+    socksEnabled: false,
+  };
+
   if (Object.keys(globalProxy.value).length !== 4) {
-    const socksIp = getSocksIpForProtocol(protocol);
-
-    const defaultGlobalProxy: ProxyInfo = {
-      host: socksIp,
-      port: baseConfig.port!,
-      proxyDNS: baseConfig.proxyDNS,
-      type: ProxyInfoType.socks,
-    };
-
-    const newGlobalProxyDetails: ProxyDetails = {
-      socksEnabled: false,
-      protocol: connection.value.protocol,
-      server: hostname,
-      country: country,
-      city: city,
-      proxyDNS: baseConfig.proxyDNS,
-    };
-
-    globalProxy.value = defaultGlobalProxy;
-    globalProxyDetails.value = newGlobalProxyDetails;
+    globalProxy.value = createDefaultProxy(protocol);
+    globalProxyDetails.value = newProxyDetails;
+  }
+
+  if (!hostProxiesDetails.value[activeTabHost.value]) {
+    hostProxies.value[activeTabHost.value] = createDefaultProxy(protocol);
+    hostProxiesDetails.value[activeTabHost.value] = newProxyDetails;
   }
 };
 
@@ -96,7 +108,7 @@ const setGlobalProxy = ({
   updateConnection();
 };
 
-const setHostProxy = (
+const setCurrentHostProxy = (
   { country, city, hostname, ipv4_address, port = baseConfig.port }: Partial<ProxyOperationArgs>,
   host: string,
 ) => {
@@ -124,13 +136,14 @@ const setHostProxy = (
 
 const useSocksProxy = () => {
   return {
-    activeTabProxyEnabled,
-    hostProxiesDetails,
+    currentHostProxyDetails,
+    currentHostProxyEnabled,
     globalProxyDetails,
     globalProxyEnabled,
-    initGlobalProxy,
-    setHostProxy,
+    initProxies,
+    setCurrentHostProxy,
     setGlobalProxy,
+    toggleCurrentHostProxy,
     toggleGlobalProxy,
   };
 };
diff --git a/src/helpers/socksProxy.ts b/src/helpers/socksProxy.ts
index 144501eb..8ef690ba 100644
--- a/src/helpers/socksProxy.ts
+++ b/src/helpers/socksProxy.ts
@@ -21,6 +21,7 @@ const handleProxyRequest = async (details: browser.proxy._OnRequestDetails) => {
   const { globalProxyDetails } = await browser.storage.local.get('globalProxyDetails');
   const { excludedHosts } = await browser.storage.local.get('excludedHosts');
   const { hostProxies } = await browser.storage.local.get('hostProxies');
+  const { hostProxiesDetails } = await browser.storage.local.get('hostProxiesDetails');
 
   // TODO implement random proxy
   const { randomProxyActive } = await browser.storage.local.get('randomProxyActive');
@@ -29,6 +30,7 @@ const handleProxyRequest = async (details: browser.proxy._OnRequestDetails) => {
   const globalProxyDetailsParsed: ProxyDetails = JSON.parse(globalProxyDetails);
   const excludedHostsParsed: string[] = JSON.parse(excludedHosts);
   const hostProxiesParsed = JSON.parse(hostProxies);
+  const hostProxiesDetailsParsed = JSON.parse(hostProxiesDetails);
   const randomProxyActiveParsed = JSON.parse(randomProxyActive);
 
   const proxiedHosts = Object.keys(hostProxiesParsed);
@@ -42,24 +44,27 @@ const handleProxyRequest = async (details: browser.proxy._OnRequestDetails) => {
       console.log('_____________________________');
     }
     return { type: 'direct' };
-  } else if (proxiedHosts.includes(currentHost)) {
+  } else if (
+    proxiedHosts.includes(currentHost) &&
+    hostProxiesDetailsParsed[currentHost].socksEnabled
+  ) {
     // TODO Warn about ubO or other extension creation DNS leaks
     // Maybe check automatically?
 
-    console.log(`host proxy "${currentHost}" used for :`, details.url);
-    console.log('_____________________________');
+    // console.log(`host proxy "${currentHost}" used for :`, details.url);
+    // console.log('_____________________________');
     return hostProxiesParsed[currentHost];
   } else if (randomProxyActiveParsed) {
-    console.log('random proxy');
-    console.log('_____________________________');
+    // console.log('random proxy');
+    // console.log('_____________________________');
     return { type: 'direct' };
   } else if (globalProxyDetailsParsed.socksEnabled) {
-    console.log('global proxy', details.url);
-    console.log('_____________________________');
+    // console.log('global proxy', details.url);
+    // console.log('_____________________________');
     return globalConfigParsed;
   }
-  console.log('no proxy in use');
-  console.log('_____________________________');
+  // console.log('no proxy in use');
+  // console.log('_____________________________');
   return { type: 'direct' };
 };
 
diff --git a/src/popup/App.vue b/src/popup/App.vue
index 8d50bd22..e4a7626e 100644
--- a/src/popup/App.vue
+++ b/src/popup/App.vue
@@ -1,32 +1,18 @@
 <script lang="ts" setup>
-import { provide, watch, onMounted } from 'vue';
+import { provide, onMounted } from 'vue';
 import { useQueryProvider } from 'vue-query';
 import { NConfigProvider, GlobalThemeOverrides, darkTheme } from 'naive-ui';
 
 import Popup from '@/popup/Popup.vue';
 
 import useConnection, { ConnectionKey } from '@/composables/useConnection';
-import useSocksProxy from '@/composables/useSocksProxy';
 import useListProxies from '@/composables/useListProxies';
 
 const { isLoading, connection, isError } = useConnection();
-const { initGlobalProxy } = useSocksProxy();
 const { getSocksProxies } = useListProxies();
 
 provide(ConnectionKey, { connection, isLoading, isError });
 
-watch(connection, async () => {
-  // Make sure the proxy always has some value
-  await initGlobalProxy(
-    {
-      hostname: connection.value.server,
-      country: connection.value.country,
-      city: connection.value.city,
-    },
-    connection.value.protocol!,
-  );
-});
-
 const loadProxies = async () => {
   await getSocksProxies();
 };
diff --git a/src/popup/views/Home.vue b/src/popup/views/Home.vue
index 784953a5..81736c30 100644
--- a/src/popup/views/Home.vue
+++ b/src/popup/views/Home.vue
@@ -1,12 +1,28 @@
 <script lang="ts" setup>
-import { inject } from 'vue';
+import { inject, watch } from 'vue';
 
 import NotificationsCarousel from '@/components/NotificationsCarousel.vue';
 import ConnectionDetails from '@/components/ConnectionDetails/ConnectionDetails.vue';
 
-import { ConnectionKey, defaultConnection } from '@/composables/useConnection';
+import useConnection, { ConnectionKey, defaultConnection } from '@/composables/useConnection';
+import useSocksProxy from '@/composables/useSocksProxy';
+
+const { initProxies } = useSocksProxy();
+const { connection } = useConnection();
 
 const { isLoading } = inject(ConnectionKey, defaultConnection);
+
+watch(connection, async () => {
+  // Make sure the proxies always have some value on popup start
+  await initProxies(
+    {
+      hostname: connection.value.server,
+      country: connection.value.country,
+      city: connection.value.city,
+    },
+    connection.value.protocol!,
+  );
+});
 </script>
 
 <template>