Fix referers

6 files changed, 47 insertions, 18 deletions

diff --git a/src/invidious.cr b/src/invidious.cr
index 4b3c473e..1f33c2db 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -131,6 +131,19 @@ before_all do |env|
       end
     end
   end
+
+  current_page = env.request.path
+  if env.request.query
+    query = HTTP::Params.parse(env.request.query.not_nil!)
+
+    if query["referer"]?
+      query["referer"] = get_referer(env, "/")
+    end
+
+    current_page += "?#{query}"
+  end
+
+  env.set "current_page", URI.escape(current_page)
 end
 
 get "/" do |env|
@@ -411,8 +424,7 @@ end
 
 # See https://github.com/rg3/youtube-dl/blob/master/youtube_dl/extractor/youtube.py#L79
 post "/login" do |env|
-  referer = env.params.query["referer"]?
-  referer ||= get_referer(env, "/feed/subscriptions")
+  referer = get_referer(env, "/feed/subscriptions")
 
   email = env.params.body["email"]?
   password = env.params.body["password"]?
@@ -506,7 +518,7 @@ post "/login" do |env|
           end
 
           if !tfa_code
-            next env.redirect "/login?tfa=true&type=google"
+            next env.redirect "/login?tfa=true&type=google&referer=#{URI.escape(referer)}"
           end
 
           tl = challenge_results[1][2]
@@ -677,7 +689,7 @@ get "/signout" do |env|
   end
 
   env.request.cookies.add_response_headers(env.response.headers)
-  env.redirect referer
+  env.redirect URI.unescape(referer)
 end
 
 get "/preferences" do |env|
diff --git a/src/invidious/helpers/utils.cr b/src/invidious/helpers/utils.cr
index 9ff411cd..4690a405 100644
--- a/src/invidious/helpers/utils.cr
+++ b/src/invidious/helpers/utils.cr
@@ -150,10 +150,27 @@ def make_host_url(ssl, host)
 end
 
 def get_referer(env, fallback = "/")
-  referer = env.request.headers["referer"]?
+  referer = env.params.query["referer"]?
+  referer ||= env.request.headers["referer"]?
   referer ||= fallback
 
-  referer = URI.parse(referer).full_path
+  referer = URI.parse(referer)
+
+  # "Unroll" nested referers
+  loop do
+    if referer.query
+      params = HTTP::Params.parse(referer.query.not_nil!)
+      if params["referer"]?
+        referer = URI.parse(URI.unescape(params["referer"]))
+      else
+        break
+      end
+    else
+      break
+    end
+  end
+
+  referer = referer.full_path
 
   if referer == env.request.path
     referer = fallback
diff --git a/src/invidious/views/channel.ecr b/src/invidious/views/channel.ecr
index ab8c0b2e..3e2df0ba 100644
--- a/src/invidious/views/channel.ecr
+++ b/src/invidious/views/channel.ecr
@@ -16,16 +16,16 @@
 <p class="h-box">
 <% if user %>
     <% if subscriptions.includes? ucid %>
-    <a href="/subscription_ajax?action_remove_subscriptions=1&c=<%= ucid %>">
+    <a href="/subscription_ajax?action_remove_subscriptions=1&c=<%= ucid %>&referer=<%= env.get("current_page") %>">
         <b>Unsubscribe from <%= author %></b>
     </a>
     <% else %>
-    <a href="/subscription_ajax?action_create_subscription_to_channel=1&c=<%= ucid %>">
+    <a href="/subscription_ajax?action_create_subscription_to_channel=1&c=<%= ucid %>&referer=<%= env.get("current_page") %>">
         <b>Subscribe to <%= author %></b>
     </a>
     <% end %>
 <% else %>
-    <a href="/login">
+    <a href="/login?referer=<%= env.get("current_page") %>">
         <b>Login to subscribe to <%= author %></b>
     </a>
 <% end %>
diff --git a/src/invidious/views/login.ecr b/src/invidious/views/login.ecr
index 3f19ba53..dc88379f 100644
--- a/src/invidious/views/login.ecr
+++ b/src/invidious/views/login.ecr
@@ -16,7 +16,7 @@
             </div>
             <hr>
             <% if account_type == "invidious" %>
-            <form class="pure-form pure-form-stacked" action="/login?referer=<%= referer %>&type=invidious" method="post">
+            <form class="pure-form pure-form-stacked" action="/login?referer=<%= URI.escape(referer) %>&type=invidious" method="post">
                 <fieldset>
                     <label for="email">User ID:</label>
                     <input required class="pure-input-1" name="email" type="text" placeholder="User ID">
@@ -34,7 +34,7 @@
                 </fieldset>
             </form>
             <% elsif account_type == "google" %>
-            <form class="pure-form pure-form-stacked" action="/login?referer=<%= referer %>" method="post">
+            <form class="pure-form pure-form-stacked" action="/login?referer=<%= URI.escape(referer) %>" method="post">
                 <fieldset>
                     <label for="email">Email:</label>
                     <input required class="pure-input-1" name="email" type="email" placeholder="Email">
diff --git a/src/invidious/views/template.ecr b/src/invidious/views/template.ecr
index 2f565541..871a5f78 100644
--- a/src/invidious/views/template.ecr
+++ b/src/invidious/views/template.ecr
@@ -34,7 +34,7 @@
         <div class="pure-u-1 pure-u-md-8-24 user-field">
         <% if env.get? "user" %>
         <div class="pure-u-1-4">
-          <a href="/toggle_theme" class="pure-menu-heading">
+          <a href="/toggle_theme?referer=<%= env.get("current_page") %>" class="pure-menu-heading">
             <% preferences = env.get("user").as(User).preferences %>
             <% if preferences.dark_mode %>
             <i class="icon ion-ios-sunny"></i>
@@ -54,15 +54,15 @@
           </a>
         </div>
         <div class="pure-u-1-4">
-          <a href="/preferences" class="pure-menu-heading">
+          <a href="/preferences?referer=<%= env.get("current_page") %>" class="pure-menu-heading">
             <i class="icon ion-ios-cog"></i>
           </a>
         </div>
         <div class="pure-u-1-4">
-          <a href="/signout" class="pure-menu-heading">Sign out</a>
+          <a href="/signout?referer=<%= env.get("current_page") %>" class="pure-menu-heading">Sign out</a>
         </div>
         <% else %>
-          <a href="/login" class="pure-menu-heading">Login</a>
+          <a href="/login?referer=<%= env.get("current_page") %>" class="pure-menu-heading">Login</a>
         <% end %>
         </div>
       </div>
diff --git a/src/invidious/views/watch.ecr b/src/invidious/views/watch.ecr
index c09da851..ca5e9372 100644
--- a/src/invidious/views/watch.ecr
+++ b/src/invidious/views/watch.ecr
@@ -232,20 +232,20 @@ get_youtube_comments();
         <% if user %>
             <% if subscriptions.includes? video.ucid %>
             <p>
-                <a href="/subscription_ajax?action_remove_subscriptions=1&c=<%= video.ucid %>">
+                <a href="/subscription_ajax?action_remove_subscriptions=1&c=<%= video.ucid %>&referer=<%= env.get("current_page") %>">
                     <b>Unsubscribe from <%= video.author %></b>
                 </a>
             </p>
             <% else %>
             <p>
-                <a href="/subscription_ajax?action_create_subscription_to_channel=1&c=<%= video.ucid %>">
+                <a href="/subscription_ajax?action_create_subscription_to_channel=1&c=<%= video.ucid %>&referer=<%= env.get("current_page") %>">
                     <b>Subscribe to <%= video.author %></b>
                 </a>
             </p>
             <% end %>
         <% else %>
             <p>
-                <a href="/login">
+                <a href="/login?referer=<%= env.get("current_page") %>">
                     <b>Login to subscribe to <%= video.author %></b>
                 </a>
             </p>