Merge branch 'ajk-30516-change-session-serialization' into 'master'

Use Gitlab::Json to serialize sessions

See merge request gitlab-org/gitlab!72701

app/models/active_session.rb

@@ -4,7 +4,7 @@
 #
 # The raw session information is stored by the Rails session store
 # (config/initializers/session_store.rb). These entries are accessible by the
-# rack_key_name class method and consistute the base of the session data
+# rack_key_name class method and constitute the base of the session data
 # entries. All other entries in the session store can be traced back to these
 # entries.
 #
@@ -26,10 +26,19 @@ class ActiveSession
   SESSION_BATCH_SIZE = 200
   ALLOWED_NUMBER_OF_ACTIVE_SESSIONS = 100
 
-  attr_accessor :created_at, :updated_at,
-    :ip_address, :browser, :os,
-    :device_name, :device_type,
-    :is_impersonated, :session_id, :session_private_id
+  attr_accessor :ip_address, :browser, :os,
+                :device_name, :device_type,
+                :is_impersonated, :session_id, :session_private_id
+
+  attr_reader :created_at, :updated_at
+
+  def created_at=(time)
+    @created_at = time.is_a?(String) ? Time.zone.parse(time) : time
+  end
+
+  def updated_at=(time)
+    @updated_at = time.is_a?(String) ? Time.zone.parse(time) : time
+  end
 
   def current?(rack_session)
     return false if session_private_id.nil? || rack_session.id.nil?
@@ -39,6 +48,19 @@ class ActiveSession
     session_private_id == rack_session.id.private_id
   end
 
+  def eql?(other)
+    other.is_a?(self.class) && id == other.id
+  end
+  alias_method :==, :eql?
+
+  def id
+    session_private_id.presence || session_id
+  end
+
+  def ids
+    [session_private_id, session_id].compact
+  end
+
   def human_device_type
     device_type&.titleize
   end
@@ -48,6 +70,7 @@ class ActiveSession
       session_private_id = request.session.id.private_id
       client = DeviceDetector.new(request.user_agent)
       timestamp = Time.current
+      expiry = Settings.gitlab['session_expire_delay'] * 60
 
       active_user_session = new(
         ip_address: request.remote_ip,
@@ -64,7 +87,14 @@ class ActiveSession
       redis.pipelined do
         redis.setex(
           key_name(user.id, session_private_id),
-          Settings.gitlab['session_expire_delay'] * 60,
+          expiry,
+          active_user_session.dump
+        )
+
+        # Deprecated legacy format - temporary to support mixed deployments
+        redis.setex(
+          key_name_v1(user.id, session_private_id),
+          expiry,
           Marshal.dump(active_user_session)
         )
 
@@ -92,7 +122,10 @@ class ActiveSession
   end
 
   def self.destroy_sessions(redis, user, session_ids)
+    return if session_ids.empty?
+
     key_names = session_ids.map { |session_id| key_name(user.id, session_id) }
+    key_names += session_ids.map { |session_id| key_name_v1(user.id, session_id) }
 
     redis.srem(lookup_key_name(user.id), session_ids)
 
@@ -115,20 +148,25 @@ class ActiveSession
     sessions.reject! { |session| session.current?(current_rack_session) } if current_rack_session
 
     redis_store_class.with do |redis|
-      session_ids = (sessions.map(&:session_id) | sessions.map(&:session_private_id)).compact
+      session_ids = sessions.flat_map(&:ids)
       destroy_sessions(redis, user, session_ids) if session_ids.any?
     end
   end
 
-  def self.not_impersonated(user)
+  private_class_method def self.not_impersonated(user)
     list(user).reject(&:is_impersonated)
   end
 
-  def self.rack_key_name(session_id)
+  private_class_method def self.rack_key_name(session_id)
     "#{Gitlab::Redis::Sessions::SESSION_NAMESPACE}:#{session_id}"
   end
 
   def self.key_name(user_id, session_id = '*')
+    "#{Gitlab::Redis::Sessions::USER_SESSIONS_NAMESPACE}::v2:#{user_id}:#{session_id}"
+  end
+
+  # Deprecated
+  def self.key_name_v1(user_id, session_id = '*')
     "#{Gitlab::Redis::Sessions::USER_SESSIONS_NAMESPACE}:#{user_id}:#{session_id}"
   end
 
@@ -170,73 +208,102 @@ class ActiveSession
     end
   end
 
-  # Deserializes a session Hash object from Redis.
-  #
+  def dump
+    "v2:#{Gitlab::Json.dump(self)}"
+  end
+
+  # Private:
+
   # raw_session - Raw bytes from Redis
   #
-  # Returns an ActiveSession object
-  def self.load_raw_session(raw_session)
-    # rubocop:disable Security/MarshalLoad
-    # Explanation of why this Marshal.load call is OK:
-    # https://gitlab.com/gitlab-com/gl-security/appsec/appsec-reviews/-/issues/124#note_744576714
-    Marshal.load(raw_session)
-    # rubocop:enable Security/MarshalLoad
+  # Returns an instance of this class
+  private_class_method def self.load_raw_session(raw_session)
+    return unless raw_session
+
+    if raw_session.start_with?('v2:')
+      session_data = Gitlab::Json.parse(raw_session[3..]).symbolize_keys
+      new(**session_data)
+    else
+      # Deprecated legacy format. To be removed in 15.0
+      # See: https://gitlab.com/gitlab-org/gitlab/-/issues/30516
+      # Explanation of why this Marshal.load call is OK:
+      # https://gitlab.com/gitlab-com/gl-security/appsec/appsec-reviews/-/issues/124#note_744576714
+      # rubocop:disable Security/MarshalLoad
+      Marshal.load(raw_session)
+      # rubocop:enable Security/MarshalLoad
+    end
   end
 
-  def self.rack_session_keys(rack_session_ids)
-    rack_session_ids.map { |session_id| rack_key_name(session_id)}
+  private_class_method def self.rack_session_keys(rack_session_ids)
+    rack_session_ids.map { |session_id| rack_key_name(session_id) }
   end
 
-  def self.raw_active_session_entries(redis, session_ids, user_id)
-    return [] if session_ids.empty?
+  private_class_method def self.raw_active_session_entries(redis, session_ids, user_id)
+    return {} if session_ids.empty?
+
+    found = Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+      entry_keys = session_ids.map { |session_id| key_name(user_id, session_id) }
+      session_ids.zip(redis.mget(entry_keys)).to_h
+    end
 
-    entry_keys = session_ids.map { |session_id| key_name(user_id, session_id) }
+    found.compact!
+    missing = session_ids - found.keys
+    return found if missing.empty?
 
-    Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
-      redis.mget(entry_keys)
+    fallbacks = Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands do
+      entry_keys = missing.map { |session_id| key_name_v1(user_id, session_id) }
+      missing.zip(redis.mget(entry_keys)).to_h
     end
+
+    fallbacks.merge(found.compact)
   end
 
-  def self.active_session_entries(session_ids, user_id, redis)
+  private_class_method def self.active_session_entries(session_ids, user_id, redis)
     return [] if session_ids.empty?
 
-    entry_keys = raw_active_session_entries(redis, session_ids, user_id)
-
-    entry_keys.compact.map do |raw_session|
-      load_raw_session(raw_session)
-    end
+    raw_active_session_entries(redis, session_ids, user_id)
+      .values
+      .compact
+      .map { load_raw_session(_1) }
   end
 
-  def self.clean_up_old_sessions(redis, user)
+  private_class_method def self.clean_up_old_sessions(redis, user)
     session_ids = session_ids_for_user(user.id)
 
     return if session_ids.count <= ALLOWED_NUMBER_OF_ACTIVE_SESSIONS
 
-    # remove sessions if there are more than ALLOWED_NUMBER_OF_ACTIVE_SESSIONS.
     sessions = active_session_entries(session_ids, user.id, redis)
-    sessions.sort_by! {|session| session.updated_at }.reverse!
-    destroyable_sessions = sessions.drop(ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
-    destroyable_session_ids = destroyable_sessions.flat_map { |session| [session.session_id, session.session_private_id] }.compact
-    destroy_sessions(redis, user, destroyable_session_ids) if destroyable_session_ids.any?
+    sessions.sort_by!(&:updated_at).reverse!
+
+    # remove sessions if there are more than ALLOWED_NUMBER_OF_ACTIVE_SESSIONS.
+    destroyable_session_ids = sessions
+      .drop(ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
+      .flat_map(&:ids)
+
+    destroy_sessions(redis, user, destroyable_session_ids)
   end
 
   # Cleans up the lookup set by removing any session IDs that are no longer present.
   #
   # Returns an array of marshalled ActiveModel objects that are still active.
-  def self.cleaned_up_lookup_entries(redis, user)
+  # Records removed keys in the optional `removed` argument array.
+  def self.cleaned_up_lookup_entries(redis, user, removed = [])
+    lookup_key = lookup_key_name(user.id)
     session_ids = session_ids_for_user(user.id)
-    entries = raw_active_session_entries(redis, session_ids, user.id)
+    session_ids_and_entries = raw_active_session_entries(redis, session_ids, user.id)
 
     # remove expired keys.
     # only the single key entries are automatically expired by redis, the
     # lookup entries in the set need to be removed manually.
-    session_ids_and_entries = session_ids.zip(entries)
     redis.pipelined do |pipeline|
-      session_ids_and_entries.reject { |_session_id, entry| entry }.each do |session_id, _entry|
-        pipeline.srem(lookup_key_name(user.id), session_id)
+      session_ids_and_entries.each do |session_id, entry|
+        next if entry
+
+        pipeline.srem(lookup_key, session_id)
+        removed << session_id
       end
     end
 
-    entries.compact
+    session_ids_and_entries.values.compact
   end
 end

lib/tasks/gitlab/cleanup.rake

@@ -121,27 +121,16 @@ namespace :gitlab do
               last_save_check = Time.now
             end
 
+            user = Struct.new(:id)
+
             keys.each do |key|
               user_id = key.split(':').last
 
-              lookup_key_count = redis.scard(key)
-
-              session_ids = ActiveSession.session_ids_for_user(user_id)
-              entries = ActiveSession.raw_active_session_entries(redis, session_ids, user_id)
-              session_ids_and_entries = session_ids.zip(entries)
-
-              inactive_session_ids = session_ids_and_entries.map do |session_id, session|
-                session_id if session.nil?
-              end.compact
-
-              redis.pipelined do |conn|
-                inactive_session_ids.each do |session_id|
-                  conn.srem(key, session_id)
-                end
-              end
+              removed = []
+              active = ActiveSession.cleaned_up_lookup_entries(redis, user.new(user_id), removed)
 
-              if inactive_session_ids
-                puts "deleted #{inactive_session_ids.count} out of #{lookup_key_count} lookup keys for User ##{user_id}"
+              if removed.any?
+                puts "deleted #{removed.count} out of #{active.count + removed.count} lookup keys for User ##{user_id}"
               end
             end
           end while cursor.to_i != 0

spec/models/active_session_spec.rb

@@ -3,6 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
+  let(:lookup_key) { described_class.lookup_key_name(user.id) }
   let(:user) do
     create(:user).tap do |user|
       user.current_sign_in_at = Time.current
@@ -43,46 +44,82 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
   end
 
   describe '.list' do
+    def make_session(id)
+      described_class.new(session_id: id)
+    end
+
     it 'returns all sessions by user' do
       Gitlab::Redis::Sessions.with do |redis|
-        redis.set("session:user:gitlab:#{user.id}:6919a6f1bb119dd7396fadc38fd18d0d", Marshal.dump({ session_id: 'a' }))
-        redis.set("session:user:gitlab:#{user.id}:59822c7d9fcdfa03725eff41782ad97d", Marshal.dump({ session_id: 'b' }))
-        redis.set("session:user:gitlab:9999:5c8611e4f9c69645ad1a1492f4131358", '')
+        # Some deprecated sessions
+        redis.set(described_class.key_name_v1(user.id, "6919a6f1bb119dd7396fadc38fd18d0d"), Marshal.dump(make_session('a')))
+        redis.set(described_class.key_name_v1(user.id, "59822c7d9fcdfa03725eff41782ad97d"), Marshal.dump(make_session('b')))
+        # Some new sessions
+        redis.set(described_class.key_name(user.id, 'some-unique-id-x'), make_session('c').dump)
+        redis.set(described_class.key_name(user.id, 'some-unique-id-y'), make_session('d').dump)
+        # Some red herrings
+        redis.set(described_class.key_name(9999, "5c8611e4f9c69645ad1a1492f4131358"), 'irrelevant')
+        redis.set(described_class.key_name_v1(9999, "5c8611e4f9c69645ad1a1492f4131358"), 'irrelevant')
 
         redis.sadd(
-          "session:lookup:user:gitlab:#{user.id}",
+          lookup_key,
           %w[
             6919a6f1bb119dd7396fadc38fd18d0d
             59822c7d9fcdfa03725eff41782ad97d
+            some-unique-id-x
+            some-unique-id-y
           ]
         )
       end
 
-      expect(ActiveSession.list(user)).to match_array [{ session_id: 'a' }, { session_id: 'b' }]
+      expect(described_class.list(user)).to contain_exactly(
+        have_attributes(session_id: 'a'),
+        have_attributes(session_id: 'b'),
+        have_attributes(session_id: 'c'),
+        have_attributes(session_id: 'd')
+      )
     end
 
-    it 'does not return obsolete entries and cleans them up' do
-      Gitlab::Redis::Sessions.with do |redis|
-        redis.set("session:user:gitlab:#{user.id}:6919a6f1bb119dd7396fadc38fd18d0d", Marshal.dump({ session_id: 'a' }))
+    shared_examples 'ignoring obsolete entries' do
+      let(:session_id) { '6919a6f1bb119dd7396fadc38fd18d0d' }
+      let(:session) { described_class.new(session_id: 'a') }
 
-        redis.sadd(
-          "session:lookup:user:gitlab:#{user.id}",
-          %w[
-            6919a6f1bb119dd7396fadc38fd18d0d
-            59822c7d9fcdfa03725eff41782ad97d
-          ]
-        )
-      end
+      it 'does not return obsolete entries and cleans them up' do
+        Gitlab::Redis::Sessions.with do |redis|
+          redis.set(session_key, serialized_session)
+
+          redis.sadd(
+            lookup_key,
+            [
+              session_id,
+              '59822c7d9fcdfa03725eff41782ad97d'
+            ]
+          )
+        end
 
-      expect(ActiveSession.list(user)).to eq [{ session_id: 'a' }]
+        expect(ActiveSession.list(user)).to contain_exactly(session)
 
-      Gitlab::Redis::Sessions.with do |redis|
-        expect(redis.sscan_each("session:lookup:user:gitlab:#{user.id}").to_a).to eq ['6919a6f1bb119dd7396fadc38fd18d0d']
+        Gitlab::Redis::Sessions.with do |redis|
+          expect(redis.sscan_each(lookup_key)).to contain_exactly session_id
+        end
       end
     end
 
-    it 'returns an empty array if the use does not have any active session' do
-      expect(ActiveSession.list(user)).to eq []
+    context 'when the current session is in the old format' do
+      let(:session_key) { described_class.key_name_v1(user.id, session_id) }
+      let(:serialized_session) { Marshal.dump(session) }
+
+      it_behaves_like 'ignoring obsolete entries'
+    end
+
+    context 'when the current session is in the new format' do
+      let(:session_key) { described_class.key_name(user.id, session_id) }
+      let(:serialized_session) { session.dump }
+
+      it_behaves_like 'ignoring obsolete entries'
+    end
+
+    it 'returns an empty array if the user does not have any active session' do
+      expect(ActiveSession.list(user)).to be_empty
     end
   end
 
@@ -107,13 +144,11 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
 
   describe '.session_ids_for_user' do
     it 'uses the user lookup table to return session ids' do
-      session_ids = ['59822c7d9fcdfa03725eff41782ad97d']
-
       Gitlab::Redis::Sessions.with do |redis|
-        redis.sadd("session:lookup:user:gitlab:#{user.id}", session_ids)
+        redis.sadd(lookup_key, %w[a b c])
       end
 
-      expect(ActiveSession.session_ids_for_user(user.id).map(&:to_s)).to eq(session_ids)
+      expect(ActiveSession.session_ids_for_user(user.id).map(&:to_s)).to match_array(%w[a b c])
     end
   end
 
@@ -151,49 +186,67 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
     it 'sets a new redis entry for the user session and a lookup entry' do
       ActiveSession.set(user, request)
 
+      session_id = "2::418729c72310bbf349a032f0bb6e3fce9f5a69df8f000d8ae0ac5d159d8f21ae"
+
       Gitlab::Redis::Sessions.with do |redis|
         expect(redis.scan_each.to_a).to include(
-          "session:user:gitlab:#{user.id}:2::418729c72310bbf349a032f0bb6e3fce9f5a69df8f000d8ae0ac5d159d8f21ae",
-          "session:lookup:user:gitlab:#{user.id}"
+          described_class.key_name(user.id, session_id),    # current session
+          described_class.key_name_v1(user.id, session_id), # support for mixed deployment
+          lookup_key
         )
       end
     end
 
     it 'adds timestamps and information from the request' do
-      Timecop.freeze(Time.zone.parse('2018-03-12 09:06')) do
-        ActiveSession.set(user, request)
+      time = Time.zone.parse('2018-03-12 09:06')
 
-        session = ActiveSession.list(user)
+      travel_to(time) do
+        described_class.set(user, request)
 
-        expect(session.count).to eq 1
-        expect(session.first).to have_attributes(
+        sessions = described_class.list(user)
+
+        expect(sessions).to contain_exactly have_attributes(
           ip_address: '127.0.0.1',
           browser: 'Mobile Safari',
           os: 'iOS',
           device_name: 'iPhone 6',
           device_type: 'smartphone',
-          created_at: Time.zone.parse('2018-03-12 09:06'),
-          updated_at: Time.zone.parse('2018-03-12 09:06')
+          created_at: eq(time),
+          updated_at: eq(time)
         )
       end
     end
 
+    it 'is possible to log in only using the old session key' do
+      session_id = "2::418729c72310bbf349a032f0bb6e3fce9f5a69df8f000d8ae0ac5d159d8f21ae"
+      ActiveSession.set(user, request)
+
+      Gitlab::Redis::SharedState.with do |redis|
+        redis.del(described_class.key_name(user.id, session_id))
+      end
+
+      sessions = ActiveSession.list(user)
+
+      expect(sessions).to be_present
+    end
+
     it 'keeps the created_at from the login on consecutive requests' do
-      now = Time.zone.parse('2018-03-12 09:06')
+      created_at = Time.zone.parse('2018-03-12 09:06')
+      updated_at = created_at + 1.minute
 
-      Timecop.freeze(now) do
+      travel_to(created_at) do
         ActiveSession.set(user, request)
+      end
 
-        Timecop.freeze(now + 1.minute) do
-          ActiveSession.set(user, request)
+      travel_to(updated_at) do
+        ActiveSession.set(user, request)
 
-          session = ActiveSession.list(user)
+        session = ActiveSession.list(user)
 
-          expect(session.first).to have_attributes(
-            created_at: Time.zone.parse('2018-03-12 09:06'),
-            updated_at: Time.zone.parse('2018-03-12 09:07')
-          )
-        end
+        expect(session.first).to have_attributes(
+          created_at: eq(created_at),
+          updated_at: eq(updated_at)
+        )
       end
     end
   end
@@ -206,10 +259,8 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
           # Emulate redis-rack: https://github.com/redis-store/redis-rack/blob/c75f7f1a6016ee224e2615017fbfee964f23a837/lib/rack/session/redis.rb#L88
           redis.set("session:gitlab:#{rack_session.private_id}", '')
 
-          redis.set(described_class.key_name(user.id, active_session_lookup_key),
-                    Marshal.dump(active_session))
-          redis.sadd(described_class.lookup_key_name(user.id),
-                     active_session_lookup_key)
+          redis.set(session_key, serialized_session)
+          redis.sadd(lookup_key, active_session_lookup_key)
         end
       end
 
@@ -225,7 +276,7 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
         subject
 
         Gitlab::Redis::Sessions.with do |redis|
-          expect(redis.scan_each(match: "session:lookup:user:gitlab:#{user.id}").to_a).to be_empty
+          expect(redis.scan_each(match: lookup_key).to_a).to be_empty
         end
       end
 
@@ -253,7 +304,19 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
         let(:active_session) { ActiveSession.new(session_private_id: rack_session.private_id) }
         let(:active_session_lookup_key) { rack_session.private_id }
 
-        include_examples 'removes all session data'
+        context 'when using old session key serialization' do
+          let(:session_key) { described_class.key_name_v1(user.id, active_session_lookup_key) }
+          let(:serialized_session) { Marshal.dump(active_session) }
+
+          include_examples 'removes all session data'
+        end
+
+        context 'when using new session key serialization' do
+          let(:session_key) { described_class.key_name(user.id, active_session_lookup_key) }
+          let(:serialized_session) { active_session.dump }
+
+          include_examples 'removes all session data'
+        end
       end
     end
   end
@@ -265,7 +328,7 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
       ActiveSession.destroy_all_but_current(user, nil)
     end
 
-    context 'with user sessions' do
+    shared_examples 'with user sessions' do
       let(:current_session_id) { '6919a6f1bb119dd7396fadc38fd18d0d' }
 
       before do
@@ -274,10 +337,8 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
           [current_session_id, '59822c7d9fcdfa03725eff41782ad97d'].each do |session_public_id|
             session_private_id = Rack::Session::SessionId.new(session_public_id).private_id
             active_session = ActiveSession.new(session_private_id: session_private_id)
-            redis.set(described_class.key_name(user.id, session_private_id),
-                      Marshal.dump(active_session))
-            redis.sadd(described_class.lookup_key_name(user.id),
-                       session_private_id)
+            redis.set(key_name(user.id, session_private_id), dump_session(active_session))
+            redis.sadd(lookup_key, session_private_id)
           end
 
           # setup for unrelated user
@@ -285,10 +346,8 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
           session_private_id = Rack::Session::SessionId.new('5c8611e4f9c69645ad1a1492f4131358').private_id
           active_session = ActiveSession.new(session_private_id: session_private_id)
 
-          redis.set(described_class.key_name(unrelated_user_id, session_private_id),
-                    Marshal.dump(active_session))
-          redis.sadd(described_class.lookup_key_name(unrelated_user_id),
-                     session_private_id)
+          redis.set(key_name(unrelated_user_id, session_private_id), dump_session(active_session))
+          redis.sadd(described_class.lookup_key_name(unrelated_user_id), session_private_id)
         end
       end
 
@@ -304,18 +363,16 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
         ActiveSession.destroy_all_but_current(user, request.session)
 
         Gitlab::Redis::Sessions.with do |redis|
-          expect(
-            redis.smembers(described_class.lookup_key_name(user.id))
-          ).to eq([session_private_id])
+          expect(redis.smembers(lookup_key)).to contain_exactly session_private_id
         end
       end
 
       it 'does not remove impersonated sessions' do
         impersonated_session_id = '6919a6f1bb119dd7396fadc38fd18eee'
         Gitlab::Redis::Sessions.with do |redis|
-          redis.set(described_class.key_name(user.id, impersonated_session_id),
-            Marshal.dump(ActiveSession.new(session_id: Rack::Session::SessionId.new(impersonated_session_id), is_impersonated: true)))
-          redis.sadd(described_class.lookup_key_name(user.id), impersonated_session_id)
+          redis.set(key_name(user.id, impersonated_session_id),
+            dump_session(ActiveSession.new(session_id: Rack::Session::SessionId.new(impersonated_session_id), is_impersonated: true)))
+          redis.sadd(lookup_key, impersonated_session_id)
         end
 
         expect { ActiveSession.destroy_all_but_current(user, request.session) }.to change { ActiveSession.session_ids_for_user(user.id).size }.from(3).to(2)
@@ -323,155 +380,289 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_sessions do
         expect(ActiveSession.session_ids_for_user(9999).size).to eq(1)
       end
     end
-  end
 
-  describe '.cleanup' do
-    before do
-      stub_const("ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS", 5)
-    end
+    context 'with legacy sessions' do
+      def key_name(user_id, id)
+        described_class.key_name_v1(user_id, id)
+      end
 
-    it 'removes obsolete lookup entries' do
-      Gitlab::Redis::Sessions.with do |redis|
-        redis.set("session:user:gitlab:#{user.id}:6919a6f1bb119dd7396fadc38fd18d0d", '')
-        redis.sadd("session:lookup:user:gitlab:#{user.id}", '6919a6f1bb119dd7396fadc38fd18d0d')
-        redis.sadd("session:lookup:user:gitlab:#{user.id}", '59822c7d9fcdfa03725eff41782ad97d')
+      def dump_session(session)
+        Marshal.dump(session)
       end
 
-      ActiveSession.cleanup(user)
+      it_behaves_like 'with user sessions'
+    end
 
-      Gitlab::Redis::Sessions.with do |redis|
-        expect(redis.smembers("session:lookup:user:gitlab:#{user.id}")).to eq ['6919a6f1bb119dd7396fadc38fd18d0d']
+    context 'with new sessions' do
+      def key_name(user_id, id)
+        described_class.key_name(user_id, id)
+      end
+
+      def dump_session(session)
+        session.dump
       end
+
+      it_behaves_like 'with user sessions'
     end
+  end
 
-    it 'does not bail if there are no lookup entries' do
-      ActiveSession.cleanup(user)
+  describe '.cleanup' do
+    before do
+      stub_const("ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS", 5)
     end
 
-    context 'cleaning up old sessions' do
-      let(:max_number_of_sessions_plus_one) { ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS + 1 }
-      let(:max_number_of_sessions_plus_two) { ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS + 2 }
+    shared_examples 'cleaning up' do
+      context 'when removing obsolete sessions' do
+        let(:current_session_id) { '6919a6f1bb119dd7396fadc38fd18d0d' }
 
-      before do
-        Gitlab::Redis::Sessions.with do |redis|
-          (1..max_number_of_sessions_plus_two).each do |number|
-            redis.set(
-              "session:user:gitlab:#{user.id}:#{number}",
-              Marshal.dump(ActiveSession.new(session_id: number.to_s, updated_at: number.days.ago))
-            )
-            redis.sadd(
-              "session:lookup:user:gitlab:#{user.id}",
-              "#{number}"
-            )
+        it 'removes obsolete lookup entries' do
+          Gitlab::Redis::Sessions.with do |redis|
+            redis.set(session_key, '')
+            redis.sadd(lookup_key, current_session_id)
+            redis.sadd(lookup_key, '59822c7d9fcdfa03725eff41782ad97d')
+          end
+
+          ActiveSession.cleanup(user)
+
+          Gitlab::Redis::Sessions.with do |redis|
+            expect(redis.smembers(lookup_key)).to contain_exactly current_session_id
           end
         end
       end
 
-      it 'removes obsolete active sessions entries' do
+      it 'does not bail if there are no lookup entries' do
         ActiveSession.cleanup(user)
+      end
 
-        Gitlab::Redis::Sessions.with do |redis|
-          sessions = redis.scan_each(match: "session:user:gitlab:#{user.id}:*").to_a
+      context 'cleaning up old sessions' do
+        let(:max_number_of_sessions_plus_one) { ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS + 1 }
+        let(:max_number_of_sessions_plus_two) { ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS + 2 }
 
-          expect(sessions.count).to eq(ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
-          expect(sessions).not_to include("session:user:gitlab:#{user.id}:#{max_number_of_sessions_plus_one}", "session:user:gitlab:#{user.id}:#{max_number_of_sessions_plus_two}")
+        before do
+          Gitlab::Redis::Sessions.with do |redis|
+            max_number_of_sessions_plus_two.times do |number|
+              redis.set(
+                key_name(user.id, number),
+                dump_session(ActiveSession.new(session_id: number.to_s, updated_at: number.days.ago))
+              )
+              redis.sadd(lookup_key, number.to_s)
+            end
+          end
         end
-      end
 
-      it 'removes obsolete lookup entries' do
-        ActiveSession.cleanup(user)
+        it 'removes obsolete active sessions entries' do
+          ActiveSession.cleanup(user)
 
-        Gitlab::Redis::Sessions.with do |redis|
-          lookup_entries = redis.smembers("session:lookup:user:gitlab:#{user.id}")
+          Gitlab::Redis::Sessions.with do |redis|
+            sessions = described_class.list(user)
 
-          expect(lookup_entries.count).to eq(ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
-          expect(lookup_entries).not_to include(max_number_of_sessions_plus_one.to_s, max_number_of_sessions_plus_two.to_s)
+            expect(sessions.count).to eq(ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
+            expect(sessions).not_to include(
+              have_attributes(session_id: max_number_of_sessions_plus_one),
+              have_attributes(session_id: max_number_of_sessions_plus_two)
+            )
+          end
         end
-      end
 
-      it 'removes obsolete lookup entries even without active session' do
-        Gitlab::Redis::Sessions.with do |redis|
-          redis.sadd(
-            "session:lookup:user:gitlab:#{user.id}",
-            "#{max_number_of_sessions_plus_two + 1}"
-          )
+        it 'removes obsolete lookup entries' do
+          ActiveSession.cleanup(user)
+
+          Gitlab::Redis::Sessions.with do |redis|
+            lookup_entries = redis.smembers(lookup_key)
+
+            expect(lookup_entries.count).to eq(ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
+            expect(lookup_entries).not_to include(max_number_of_sessions_plus_one.to_s, max_number_of_sessions_plus_two.to_s)
+          end
         end
 
-        ActiveSession.cleanup(user)
+        it 'removes obsolete lookup entries even without active session' do
+          Gitlab::Redis::Sessions.with do |redis|
+            redis.sadd(lookup_key, "#{max_number_of_sessions_plus_two + 1}")
+          end
 
-        Gitlab::Redis::Sessions.with do |redis|
-          lookup_entries = redis.smembers("session:lookup:user:gitlab:#{user.id}")
+          ActiveSession.cleanup(user)
 
-          expect(lookup_entries.count).to eq(ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
-          expect(lookup_entries).not_to include(
-            max_number_of_sessions_plus_one.to_s,
-            max_number_of_sessions_plus_two.to_s,
-            (max_number_of_sessions_plus_two + 1).to_s
-          )
+          Gitlab::Redis::Sessions.with do |redis|
+            lookup_entries = redis.smembers(lookup_key)
+
+            expect(lookup_entries.count).to eq(ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
+            expect(lookup_entries).not_to include(
+              max_number_of_sessions_plus_one.to_s,
+              max_number_of_sessions_plus_two.to_s,
+              (max_number_of_sessions_plus_two + 1).to_s
+            )
+          end
         end
-      end
 
-      context 'when the number of active sessions is lower than the limit' do
-        before do
-          Gitlab::Redis::Sessions.with do |redis|
-            ((max_number_of_sessions_plus_two - 4)..max_number_of_sessions_plus_two).each do |number|
-              redis.del("session:user:gitlab:#{user.id}:#{number}")
+        context 'when the number of active sessions is lower than the limit' do
+          before do
+            Gitlab::Redis::Sessions.with do |redis|
+              ((max_number_of_sessions_plus_two - 4)..max_number_of_sessions_plus_two).each do |number|
+                redis.del(key_name(user.id, number))
+              end
             end
           end
-        end
 
-        it 'does not remove active session entries, but removes lookup entries' do
-          lookup_entries_before_cleanup = Gitlab::Redis::Sessions.with do |redis|
-            redis.smembers("session:lookup:user:gitlab:#{user.id}")
+          it 'does not remove active session entries, but removes lookup entries' do
+            lookup_entries_before_cleanup = Gitlab::Redis::Sessions.with do |redis|
+              redis.smembers(lookup_key)
+            end
+
+            sessions_before_cleanup = described_class.list(user)
+
+            described_class.cleanup(user)
+
+            Gitlab::Redis::Sessions.with do |redis|
+              lookup_entries = redis.smembers(lookup_key)
+              sessions = described_class.list(user)
+
+              expect(sessions.count).to eq(sessions_before_cleanup.count)
+              expect(lookup_entries.count).to be < lookup_entries_before_cleanup.count
+            end
           end
+        end
+      end
+
+      context 'cleaning up old sessions stored by Rack::Session::SessionId#private_id' do
+        let(:max_number_of_sessions_plus_one) { ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS + 1 }
+        let(:max_number_of_sessions_plus_two) { ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS + 2 }
 
-          sessions_before_cleanup = Gitlab::Redis::Sessions.with do |redis|
-            redis.scan_each(match: "session:user:gitlab:#{user.id}:*").to_a
+        before do
+          Gitlab::Redis::Sessions.with do |redis|
+            (1..max_number_of_sessions_plus_two).each do |number|
+              redis.set(
+                key_name(user.id, number),
+                dump_session(ActiveSession.new(session_private_id: number.to_s, updated_at: number.days.ago))
+              )
+              redis.sadd(lookup_key, number.to_s)
+            end
           end
+        end
 
-          ActiveSession.cleanup(user)
+        it 'removes obsolete active sessions entries' do
+          described_class.cleanup(user)
 
           Gitlab::Redis::Sessions.with do |redis|
-            lookup_entries = redis.smembers("session:lookup:user:gitlab:#{user.id}")
-            sessions = redis.scan_each(match: "session:user:gitlab:#{user.id}:*").to_a
-            expect(sessions.count).to eq(sessions_before_cleanup.count)
-            expect(lookup_entries.count).to be < lookup_entries_before_cleanup.count
+            sessions = described_class.list(user)
+
+            expect(sessions.count).to eq(described_class::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
+            expect(sessions).not_to include(
+              key_name(user.id, max_number_of_sessions_plus_one),
+              key_name(user.id, max_number_of_sessions_plus_two)
+            )
           end
         end
       end
     end
 
-    context 'cleaning up old sessions stored by Rack::Session::SessionId#private_id' do
-      let(:max_number_of_sessions_plus_one) { ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS + 1 }
-      let(:max_number_of_sessions_plus_two) { ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS + 2 }
+    context 'with legacy sessions' do
+      let(:session_key) { described_class.key_name_v1(user.id, current_session_id) }
+
+      def key_name(user_id, session_id)
+        described_class.key_name_v1(user_id, session_id)
+      end
+
+      def dump_session(session)
+        Marshal.dump(session)
+      end
+
+      it_behaves_like 'cleaning up'
+    end
+
+    context 'with new sessions' do
+      let(:session_key) { described_class.key_name(user.id, current_session_id) }
+
+      def key_name(user_id, session_id)
+        described_class.key_name(user_id, session_id)
+      end
+
+      def dump_session(session)
+        session.dump
+      end
+
+      it_behaves_like 'cleaning up'
+    end
+  end
+
+  describe '.cleaned_up_lookup_entries' do
+    before do
+      stub_const("ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS", 5)
+    end
+
+    shared_examples 'cleaning up lookup entries' do
+      let(:current_session_id) { '6919a6f1bb119dd7396fadc38fd18d0d' }
+      let(:active_count) { 3 }
 
       before do
-        Gitlab::Redis::Sessions.with do |redis|
-          (1..max_number_of_sessions_plus_two).each do |number|
+        Gitlab::Redis::SharedState.with do |redis|
+          active_count.times do |number|
             redis.set(
-              "session:user:gitlab:#{user.id}:#{number}",
-              Marshal.dump(ActiveSession.new(session_private_id: number.to_s, updated_at: number.days.ago))
-            )
-            redis.sadd(
-              "session:lookup:user:gitlab:#{user.id}",
-              "#{number}"
+              key_name(user.id, number),
+              dump_session(ActiveSession.new(session_id: number.to_s, updated_at: number.days.ago))
             )
+
+            redis.sadd(lookup_key, number.to_s)
           end
+
+          redis.sadd(lookup_key, (active_count + 1).to_s)
+          redis.sadd(lookup_key, (active_count + 2).to_s)
         end
       end
 
-      it 'removes obsolete active sessions entries' do
-        ActiveSession.cleanup(user)
+      it 'removes obsolete lookup entries' do
+        active = Gitlab::Redis::SharedState.with do |redis|
+          ActiveSession.cleaned_up_lookup_entries(redis, user)
+        end
 
-        Gitlab::Redis::Sessions.with do |redis|
-          sessions = redis.scan_each(match: "session:user:gitlab:#{user.id}:*").to_a
+        expect(active.count).to eq(active_count)
+
+        Gitlab::Redis::SharedState.with do |redis|
+          lookup_entries = redis.smembers(lookup_key)
 
-          expect(sessions.count).to eq(ActiveSession::ALLOWED_NUMBER_OF_ACTIVE_SESSIONS)
-          expect(sessions).not_to(
-            include("session:user:gitlab:#{user.id}:#{max_number_of_sessions_plus_one}",
-                    "session:user:gitlab:#{user.id}:#{max_number_of_sessions_plus_two}"))
+          expect(lookup_entries.count).to eq(active_count)
+          expect(lookup_entries).not_to include(
+            (active_count + 1).to_s,
+            (active_count + 2).to_s
+          )
         end
       end
+
+      it 'reports the removed entries' do
+        removed = []
+        Gitlab::Redis::SharedState.with do |redis|
+          ActiveSession.cleaned_up_lookup_entries(redis, user, removed)
+        end
+
+        expect(removed.count).to eq(2)
+      end
+    end
+
+    context 'with legacy sessions' do
+      let(:session_key) { described_class.key_name_v1(user.id, current_session_id) }
+
+      def key_name(user_id, session_id)
+        described_class.key_name_v1(user_id, session_id)
+      end
+
+      def dump_session(session)
+        Marshal.dump(session)
+      end
+
+      it_behaves_like 'cleaning up lookup entries'
+    end
+
+    context 'with new sessions' do
+      let(:session_key) { described_class.key_name(user.id, current_session_id) }
+
+      def key_name(user_id, session_id)
+        described_class.key_name(user_id, session_id)
+      end
+
+      def dump_session(session)
+        session.dump
+      end
+
+      it_behaves_like 'cleaning up lookup entries'
     end
   end
 end

spec/tasks/gitlab/cleanup_rake_spec.rb

@@ -174,9 +174,9 @@ RSpec.describe 'gitlab:cleanup rake tasks', :silence_stdout do
 
       before do
         Gitlab::Redis::Sessions.with do |redis|
-          redis.set("session:user:gitlab:#{user.id}:#{existing_session_id}",
-                    Marshal.dump(true))
-          redis.sadd("session:lookup:user:gitlab:#{user.id}", (1..10).to_a)
+          redis.set(ActiveSession.key_name(user.id, existing_session_id),
+                    ActiveSession.new(session_id: 'x').dump)
+          redis.sadd(ActiveSession.lookup_key_name(user.id), (1..10).to_a)
         end
       end
 
@@ -186,10 +186,10 @@ RSpec.describe 'gitlab:cleanup rake tasks', :silence_stdout do
 
       it 'removes expired active session lookup keys' do
         Gitlab::Redis::Sessions.with do |redis|
-          lookup_key = "session:lookup:user:gitlab:#{user.id}"
+          lookup_key = ActiveSession.lookup_key_name(user.id)
+
           expect { subject }.to change { redis.scard(lookup_key) }.from(10).to(1)
-          expect(redis.smembers("session:lookup:user:gitlab:#{user.id}")).to(
-            eql([existing_session_id]))
+          expect(redis.smembers(lookup_key)).to contain_exactly existing_session_id
         end
       end
     end