lib/rack/websocket/handler/base.rb in websocket-rack-noodles-0.4.1 vs lib/rack/websocket/handler/base.rb in websocket-rack-noodles-0.4.2

@@ -4,36 +4,46 @@
       class Base
                 
         autoload :Connection, "#{ROOT_PATH}/websocket/handler/base/connection"
 
         def on_open
-          set_env_instance_variable
+          set_parent_instance_variables
           @parent.on_open(@env)
         end # Fired when a client is connected.
 
         def on_message(msg)
-          set_env_instance_variable
+          set_parent_instance_variables
           @parent.on_message(@env, msg)
         end # Fired when a message from a client is received.
 
         def on_close
-          set_env_instance_variable
+          set_parent_instance_variables
           @parent.on_close(@env)
         end # Fired when a client is disconnected.
 
         def on_error(error)
-          set_env_instance_variable
+          set_parent_instance_variables
           @parent.on_error(@env, error)
         end # Fired when error occurs.
 
         # Set application as parent and forward options
         def initialize(parent, options = {})
           @parent = parent
           @options = options[:backend] || {}
         end
 
+        def set_parent_instance_variables
+          set_env_instance_variable
+          set_session_instance_variable
+        end
+
         def set_env_instance_variable
           @parent.instance_variable_set("@env", @env)
+        end
+
+        def set_session_instance_variable
+          session = Noodles.use_memached_as_session_storage ? Noodles::MemcachedSession.new(@env) : @env['rack.session']
+          @parent.instance_variable_set("@session", session)
         end
 
         # Implemented in subclass
         def call(env)
           raise 'Not implemented'
\ No newline at end of file