4 files changed, 6 insertions, 16 deletions
diff --git a/lib/plum/rack/server.rb b/lib/plum/rack/server.rb
index 5e34cd3..37ae166 100644
--- a/lib/plum/rack/server.rb
+++ b/lib/plum/rack/server.rb
@@ -52,6 +52,7 @@ module Plum
 
             con = Session.new(app: @app,
                               plum: plum,
+                              sock: sock,
                               logger: @logger,
                               server_push: @config[:server_push],
                               remote_addr: sock.peeraddr.last)
diff --git a/lib/plum/rack/session.rb b/lib/plum/rack/session.rb
index 8266b96..d0aedc5 100644
--- a/lib/plum/rack/session.rb
+++ b/lib/plum/rack/session.rb
@@ -8,9 +8,10 @@ module Plum
     class Session
       attr_reader :app, :plum
 
-      def initialize(app:, plum:, logger:, server_push: true, remote_addr: "127.0.0.1")
+      def initialize(app:, plum:, sock:, logger:, server_push: true, remote_addr: "127.0.0.1")
         @app = app
         @plum = plum
+        @sock = sock
         @logger = logger
         @server_push = server_push
         @remote_addr = remote_addr
@@ -24,7 +25,9 @@ module Plum
 
       def run
         begin
-          @plum.run
+          while !@sock.closed? && !@sock.eof?
+            @plum << @sock.readpartial(1024)
+          end
         rescue Errno::EPIPE, Errno::ECONNRESET => e
         rescue StandardError => e
           @logger.error("#{e.class}: #{e.message}\n#{e.backtrace.map { |b| "\t#{b}" }.join("\n")}")
diff --git a/lib/plum/server/http_connection.rb b/lib/plum/server/http_connection.rb
index f864691..0e9462c 100644
--- a/lib/plum/server/http_connection.rb
+++ b/lib/plum/server/http_connection.rb
@@ -14,13 +14,6 @@ module Plum
       super(@sock.method(:write), local_settings)
     end
 
-    # Starts communication with the peer. It blocks until the io is closed, or reaches EOF.
-    def run
-      while !@sock.closed? && !@sock.eof?
-        self << @sock.readpartial(1024)
-      end
-    end
-
     # Closes the socket.
     def close
       super
diff --git a/lib/plum/server/https_connection.rb b/lib/plum/server/https_connection.rb
index 1903bc1..09e360f 100644
--- a/lib/plum/server/https_connection.rb
+++ b/lib/plum/server/https_connection.rb
@@ -6,10 +6,7 @@ module Plum
     def initialize(sock, local_settings = {})
       @sock = sock
       super(@sock.method(:write), local_settings)
-    end
 
-    # Starts communication with the peer. It blocks until the io is closed, or reaches EOF.
-    def run
       if @sock.respond_to?(:cipher) # OpenSSL::SSL::SSLSocket-like
         if CIPHER_BLACKLIST.include?(@sock.cipher.first) # [cipher-suite, ssl-version, keylen, alglen]
           on(:negotiated) {
@@ -17,10 +14,6 @@ module Plum
           }
         end
       end
-
-      while !@sock.closed? && !@sock.eof?
-        self << @sock.readpartial(1024)
-      end
     end
 
     # Closes the socket.