Fix session recovery after deletion

CHANGELOG.md

@@ -4,6 +4,7 @@ Dalli Changelog
 Unreleased
 ==========
 
+- Fix session recovery after deletion (stengineering0)
 - Fix cannot read response data included terminator `\r\n` when use meta protocol (matsubara0507)
 - Support SERVER_ERROR response from Memcached as per the [memcached spec](https://github.com/memcached/memcached/blob/e43364402195c8e822bb8f88755a60ab8bbed62a/doc/protocol.txt#L172) (grcooper)
 - Update Socket timeout handling to use Socket#timeout= when available (nickamorim)

lib/rack/session/dalli.rb

@@ -9,6 +9,10 @@ module Rack
   module Session
     # Rack::Session::Dalli provides memcached based session management.
     class Dalli < Abstract::PersistedSecure
+      class MissingSessionError < StandardError; end
+
+      RACK_SESSION_PERSISTED = 'rack.session.persisted'
+
       attr_reader :data
 
       # Don't freeze this until we fix the specs/implementation
@@ -70,23 +74,37 @@ def initialize(app, options = {})
         @data = build_data_source(options)
       end
 
-      def find_session(_req, sid)
+      def call(*_args)
+        super
+      rescue MissingSessionError
+        [401, {}, ['Wrong session ID']]
+      end
+
+      def find_session(req, sid)
         with_dalli_client([nil, {}]) do |dc|
           existing_session = existing_session_for_sid(dc, sid)
-          return [sid, existing_session] unless existing_session.nil?
+          if existing_session.nil?
+            sid = create_sid_with_empty_session(dc)
+            existing_session = {}
+          end
 
-          [create_sid_with_empty_session(dc), {}]
+          update_session_persisted_data(req, { id: sid })
+          return [sid, existing_session]
         end
       end
 
-      def write_session(_req, sid, session, options)
+      def write_session(req, sid, session, options)
         return false unless sid
 
         key = memcached_key_from_sid(sid)
         return false unless key
 
         with_dalli_client(false) do |dc|
-          dc.set(memcached_key_from_sid(sid), session, ttl(options[:expire_after]))
+          write_session_safely!(
+            dc, sid, session_persisted_data(req),
+            write_args: [memcached_key_from_sid(sid), session, ttl(options[:expire_after])]
+          )
+
           sid
         end
       end
@@ -145,6 +163,15 @@ def build_data_source(options)
         end
       end
 
+      def write_session_safely!(dalli_client, sid, persisted_data, write_args:)
+        if persisted_data && persisted_data[:id] == sid # That means that we update the existing session
+          # Override the session only if it still exists in the store!
+          raise MissingSessionError unless dalli_client.replace(*write_args)
+        else
+          dalli_client.set(*write_args)
+        end
+      end
+
       def extract_dalli_options(options)
         raise 'Rack::Session::Dalli no longer supports the :cache option.' if options[:cache]
 
@@ -190,6 +217,14 @@ def with_dalli_client(result_on_error = nil, &block)
       def ttl(expire_after)
         expire_after.nil? ? 0 : expire_after + 1
       end
+
+      def session_persisted_data(req)
+        req.get_header RACK_SESSION_PERSISTED
+      end
+
+      def update_session_persisted_data(req, data)
+        req.set_header RACK_SESSION_PERSISTED, data
+      end
     end
   end
 end

test/test_rack_session.rb

@@ -51,6 +51,25 @@
                      incrementor_proc.call(env)
                    end)
   end
+  let(:user_id_session) do
+    Rack::Lint.new(proc do |env|
+                     session = env['rack.session']
+
+                     case env['PATH_INFO']
+                     when '/login'
+                       session[:user_id] = 1
+                     when '/logout'
+                       raise 'User not logged in' if session[:user_id].nil?
+
+                       session.delete(:user_id)
+                       session.options[:renew] = true
+                     when '/slow'
+                       Fiber.yield
+                     end
+
+                     Rack::Response.new(session.inspect).to_a
+                   end)
+  end
   let(:incrementor) { Rack::Lint.new(incrementor_proc) }
 
   it 'faults on no connection' do
@@ -346,4 +365,37 @@
 
     refute_equal ses0, ses1
   end
+
+  it "doesn't allow session id to be reused" do
+    rsd = Rack::Session::Dalli.new(user_id_session)
+
+    login_response = Rack::MockRequest.new(rsd).get('/login')
+    login_cookie = login_response['Set-Cookie']
+
+    slow_request = Fiber.new do
+      Rack::MockRequest.new(rsd).get('/slow', 'HTTP_COOKIE' => login_cookie)
+    end
+    slow_request.resume
+
+    # Check that the session is valid:
+    response = Rack::MockRequest.new(rsd).get('/', 'HTTP_COOKIE' => login_cookie)
+
+    assert_equal response.body, { 'user_id' => 1 }.to_s
+
+    logout_response = Rack::MockRequest.new(rsd).get('/logout', 'HTTP_COOKIE' => login_cookie)
+    logout_cookie = logout_response['Set-Cookie']
+
+    # Check that the session id is different after logout:
+    refute_equal login_cookie[session_match], logout_cookie[session_match]
+
+    slow_response = slow_request.resume
+
+    assert_equal 401, slow_response.status
+    assert_equal 'Wrong session ID', slow_response.body
+
+    # Check that the cookie can't be reused:
+    response = Rack::MockRequest.new(rsd).get('/', 'HTTP_COOKIE' => login_cookie)
+
+    assert_equal '{}', response.body
+  end
 end