Improve redirect uri state and fix all remaining tests

2 files changed, 21 insertions, 12 deletions

diff --git a/app/controllers/google_api/authorizations_controller.rb b/app/controllers/google_api/authorizations_controller.rb
index 709d1d34796..5551057ff55 100644
--- a/app/controllers/google_api/authorizations_controller.rb
+++ b/app/controllers/google_api/authorizations_controller.rb
@@ -9,16 +9,21 @@ module GoogleApi
       session[GoogleApi::CloudPlatform::Client.session_key_for_expires_at] =
         expires_at.to_s
 
-      key, _ = GoogleApi::CloudPlatform::Client
-        .session_key_for_second_redirect_uri(secure: params[:state])
+      state_redirect_uri = redirect_uri_from_session_key(params[:state])
 
-      second_redirect_uri = session[key]
-
-      if second_redirect_uri.present?
-        redirect_to second_redirect_uri
+      if state_redirect_uri
+        redirect_to state_redirect_uri
       else
         redirect_to root_path
       end
     end
+
+    private
+
+    def redirect_uri_from_session_key(state)
+      key = GoogleApi::CloudPlatform::Client
+        .session_key_for_redirect_uri(params[:state])
+      session[key] if key
+    end
   end
 end
diff --git a/app/controllers/projects/clusters_controller.rb b/app/controllers/projects/clusters_controller.rb
index 2f7364f4abf..03019b0becc 100644
--- a/app/controllers/projects/clusters_controller.rb
+++ b/app/controllers/projects/clusters_controller.rb
@@ -16,13 +16,11 @@ class Projects::ClustersController < Projects::ApplicationController
 
   def login
     begin
-      GoogleApi::CloudPlatform::Client.session_key_for_second_redirect_uri.tap do |key, secure|
-        session[key] = namespace_project_clusters_url.to_s
+      state = generate_session_key_redirect(namespace_project_clusters_url.to_s)
 
-        @authorize_url = GoogleApi::CloudPlatform::Client.new(
-          nil, callback_google_api_auth_url,
-          state: secure).authorize_url
-      end
+      @authorize_url = GoogleApi::CloudPlatform::Client.new(
+        nil, callback_google_api_auth_url,
+        state: state).authorize_url
     rescue GoogleApi::Auth::ConfigMissingError
       # no-op
     end
@@ -122,6 +120,12 @@ class Projects::ClustersController < Projects::ApplicationController
       session[GoogleApi::CloudPlatform::Client.session_key_for_expires_at]
   end
 
+  def generate_session_key_redirect(uri)
+    GoogleApi::CloudPlatform::Client.new_session_key_for_redirect_uri do |key|
+      session[key] = uri
+    end
+  end
+
   def authorize_update_cluster!
     access_denied! unless can?(current_user, :update_cluster, cluster)
   end