From 5faa98f481e0f5e0ccb1758c34a104f523ab21d2 Mon Sep 17 00:00:00 2001
From: James Edwards-Jones <jedwardsjones@gitlab.com>
Date: Wed, 24 Apr 2019 17:50:00 +0700
Subject: [PATCH] Session stored globally per request

- This can be accessed with Session.current and is restored after.
- Data can be stored under a key with NamespacedSessionStore
---
 app/controllers/application_controller.rb     |  5 ++++
 lib/gitlab/namespaced_session_store.rb        | 22 +++++++++++++++
 lib/gitlab/session.rb                         | 27 +++++++++++++++++++
 .../gitlab/namespaced_session_store_spec.rb   | 22 +++++++++++++++
 spec/lib/gitlab/session_spec.rb               | 27 +++++++++++++++++++
 5 files changed, 103 insertions(+)
 create mode 100644 lib/gitlab/namespaced_session_store.rb
 create mode 100644 lib/gitlab/session.rb
 create mode 100644 spec/lib/gitlab/namespaced_session_store_spec.rb
 create mode 100644 spec/lib/gitlab/session_spec.rb

diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index ceaa84acaba..4cbab6811bc 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -27,6 +27,7 @@ class ApplicationController < ActionController::Base
   before_action :check_impersonation_availability
 
   around_action :set_locale
+  around_action :set_session_storage
 
   after_action :set_page_title_header, if: :json_request?
   after_action :limit_unauthenticated_session_times
@@ -434,6 +435,10 @@ class ApplicationController < ActionController::Base
     Gitlab::I18n.with_user_locale(current_user, &block)
   end
 
+  def set_session_storage(&block)
+    Gitlab::Session.with_session(session, &block)
+  end
+
   def set_page_title_header
     # Per https://tools.ietf.org/html/rfc5987, headers need to be ISO-8859-1, not UTF-8
     response.headers['Page-Title'] = URI.escape(page_title('GitLab'))
diff --git a/lib/gitlab/namespaced_session_store.rb b/lib/gitlab/namespaced_session_store.rb
new file mode 100644
index 00000000000..34520078bfb
--- /dev/null
+++ b/lib/gitlab/namespaced_session_store.rb
@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Gitlab
+  class NamespacedSessionStore
+    delegate :[], :[]=, to: :store
+
+    def initialize(key)
+      @key = key
+    end
+
+    def initiated?
+      !Session.current.nil?
+    end
+
+    def store
+      return unless Session.current
+
+      Session.current[@key] ||= {}
+      Session.current[@key]
+    end
+  end
+end
diff --git a/lib/gitlab/session.rb b/lib/gitlab/session.rb
new file mode 100644
index 00000000000..7487ba04a6d
--- /dev/null
+++ b/lib/gitlab/session.rb
@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Gitlab
+  class Session
+    STORE_KEY = :session_storage
+
+    class << self
+      def with_session(session)
+        old = self.current
+        self.current = session
+        yield
+      ensure
+        self.current = old
+      end
+
+      def current
+        Thread.current[STORE_KEY]
+      end
+
+      protected
+
+      def current=(value)
+        Thread.current[STORE_KEY] = value
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/namespaced_session_store_spec.rb b/spec/lib/gitlab/namespaced_session_store_spec.rb
new file mode 100644
index 00000000000..c0af2ede32a
--- /dev/null
+++ b/spec/lib/gitlab/namespaced_session_store_spec.rb
@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::NamespacedSessionStore do
+  let(:key) { :some_key }
+  subject { described_class.new(key) }
+
+  it 'stores data under the specified key' do
+    Gitlab::Session.with_session({}) do
+      subject[:new_data] = 123
+
+      expect(Thread.current[:session_storage][key]).to eq(new_data: 123)
+    end
+  end
+
+  it 'retrieves data from the given key' do
+    Thread.current[:session_storage] = { key => { existing_data: 123 } }
+
+    expect(subject[:existing_data]).to eq 123
+  end
+end
diff --git a/spec/lib/gitlab/session_spec.rb b/spec/lib/gitlab/session_spec.rb
new file mode 100644
index 00000000000..8db73f0ec7b
--- /dev/null
+++ b/spec/lib/gitlab/session_spec.rb
@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Session do
+  it 'uses the current thread as a data store' do
+    Thread.current[:session_storage] = { a: :b }
+
+    expect(described_class.current).to eq(a: :b)
+  ensure
+    Thread.current[:session_storage] = nil
+  end
+
+  describe '#with_session' do
+    it 'sets session hash' do
+      described_class.with_session(one: 1) do
+        expect(described_class.current).to eq(one: 1)
+      end
+    end
+
+    it 'restores current store after' do
+      described_class.with_session(two: 2) { }
+
+      expect(described_class.current).to eq nil
+    end
+  end
+end