3 files changed, 47 insertions, 21 deletions
diff --git a/src/pages/background/background.js b/src/pages/background/background.js
index 0d558b01..1fcba190 100644
--- a/src/pages/background/background.js
+++ b/src/pages/background/background.js
@@ -67,16 +67,22 @@ browser.webRequest.onBeforeRequest.addListener(
       tabIdRedirects[details.tabId]
     )
 
-    if (newUrl && newUrl.startsWith("https://no-instance.libredirect.invalid")) {
-      const url = new URL(newUrl)
-      const frontend = url.searchParams.get("frontend")
-      const oldUrl = new URL(url.searchParams.get("url"))
+    if (
+      (newUrl && newUrl.startsWith("https://no-instance.libredirect.invalid")) ||
+      (!newUrl && url.href.startsWith("https://no-instance.libredirect.invalid"))
+    ) {
+      newUrl = newUrl ? new URL(newUrl) : url
+      const frontend = newUrl.searchParams.get("frontend")
+      const oldUrl = new URL(newUrl.searchParams.get("url"))
       const params = new URLSearchParams({
         message: "no_instance",
         url: oldUrl,
         frontend: frontend,
       })
-      newUrl = browser.runtime.getURL(`/pages/messages/index.html?${params.toString()}`)
+      browser.tabs.update({
+        url: browser.runtime.getURL(`/pages/messages/index.html?${params.toString()}`),
+      })
+      return { cancel: true }
     }
 
     if (!newUrl) {
@@ -85,7 +91,10 @@ browser.webRequest.onBeforeRequest.addListener(
           message: "disabled",
           url: url.href,
         })
-        newUrl = browser.runtime.getURL(`/pages/messages/index.html?${params.toString()}`)
+        browser.tabs.update({
+          url: browser.runtime.getURL(`/pages/messages/index.html?${params.toString()}`),
+        })
+        return { cancel: true }
       }
     }
 
@@ -121,9 +130,11 @@ browser.webRequest.onHeadersReceived.addListener(
         frontend: frontend,
         service: service,
       })
-      browser.tabs.update({
-        url: browser.runtime.getURL(`/pages/messages/index.html?${params.toString()}`),
-      })
+      setTimeout(() => {
+        browser.tabs.update({
+          url: browser.runtime.getURL(`/pages/messages/index.html?${params.toString()}`),
+        })
+      }, 2000)
     }
   },
   { urls: ["<all_urls>"] }
@@ -365,3 +376,8 @@ browser.runtime.getPlatformInfo(r => {
     })
   }
 })
+
+browser.runtime.onMessage.addListener(r => {
+  if (r.message == "reverse") tabIdRedirects[r.tabId] = false
+  else if (r.message == "redirect") tabIdRedirects[r.tabId] = true
+})
diff --git a/src/pages/messages_src/App.svelte b/src/pages/messages_src/App.svelte
index 6d68ff5f..1c5170dd 100644
--- a/src/pages/messages_src/App.svelte
+++ b/src/pages/messages_src/App.svelte
@@ -33,7 +33,6 @@
       await servicesHelper.initDefaults()
       opts = await utils.getOptions()
     }
-    redirects = await utils.getList(opts)
     options.set(opts)
     config.set(await utils.getConfig())
   })
@@ -45,7 +44,6 @@
   $: if (_options) style = utils.style(_options, window)
 
   let autoPicking = false
-  let redirects
 
   const params = new URLSearchParams(window.location.search)
   const oldUrl = new URL(params.get("url"))
@@ -53,10 +51,14 @@
   async function autoPick() {
     const frontend = params.get("frontend")
     autoPicking = true
+    const redirects = await utils.getList(_options)
     const instances = utils.randomInstances(redirects[frontend]["clearnet"], 5)
-    const pings = await Promise.all([...instances.map(async instance => [instance, await utils.ping(instance)])])
+    const pings = await Promise.all([
+      ...instances.map(async instance => {
+        return [instance, await utils.ping(instance)]
+      }),
+    ])
     pings.sort((a, b) => a[1] - b[1])
-    console.log(pings)
     _options[frontend].push(pings[0][0])
     options.set(_options)
     autoPicking = false
@@ -85,8 +87,9 @@
   }
 
   async function removeInstance() {
+    const service = await servicesHelper.computeService(oldUrl)
     const frontend = params.get("frontend")
-    const i = _options[frontend].indexOf(utils.protocolHost(oldUrl))
+    const i = _options[frontend].findIndex(instance => oldUrl.href.startsWith(instance))
     _options[frontend].splice(i, 1)
     options.set(_options)
     const newUrl = await servicesHelper.switchInstance(oldUrl, service)
@@ -94,8 +97,10 @@
   }
 
   async function removeAndAutoPickInstance() {
+    const service = await servicesHelper.computeService(oldUrl)
+
     const frontend = params.get("frontend")
-    const i = _options[frontend].indexOf(utils.protocolHost(oldUrl))
+    const i = _options[frontend].findIndex(instance => oldUrl.href.startsWith(instance))
     _options[frontend].splice(i, 1)
     options.set(_options)
     await autoPick()
diff --git a/src/pages/popup_src/Buttons.svelte b/src/pages/popup_src/Buttons.svelte
index c0644ff7..ab5682dc 100644
--- a/src/pages/popup_src/Buttons.svelte
+++ b/src/pages/popup_src/Buttons.svelte
@@ -44,8 +44,10 @@
     <Row
       class="interactive"
       on:click={() => {
-        browser.tabs.update({ url: redirect }, () => {
-          window.close()
+        browser.tabs.query({ active: true, currentWindow: true }, tabs => {
+          browser.runtime.sendMessage({ message: "redirect", tabId: tabs[0].id }, () => {
+            browser.tabs.update({ url: redirect })
+          })
         })
       }}
     >
@@ -74,10 +76,13 @@
     </Row>
     <Row
       class="interactive"
-      on:click={() =>
-        browser.tabs.update({ url: redirectToOriginal }, () => {
-          window.close()
-        })}
+      on:click={() => {
+        browser.tabs.query({ active: true, currentWindow: true }, tabs => {
+          browser.runtime.sendMessage({ message: "reverse", tabId: tabs[0].id }, () => {
+            browser.tabs.update({ url: redirectToOriginal })
+          })
+        })
+      }}
     >
       <Label>{browser.i18n.getMessage("redirectToOriginal" || "Redirect to Original")}</Label>
       <RedirectToOriginalIcon />