services/api/app/middlewares/rack_socket.rb

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: AGPL-3.0
   4
   5 require 'rack'
   6 require 'faye/websocket'
   7 require 'eventmachine'
   8
   9 # A Rack middleware to handle inbound websocket connection requests and hand
  10 # them over to the faye websocket library.
  11 class RackSocket
  12
  13   DEFAULT_ENDPOINT  = '/websocket'
  14
  15   # Stop EventMachine on signal, this should give it a chance to to unwind any
  16   # open connections.
  17   def die_gracefully_on_signal
  18     Signal.trap("INT") { EM.stop }
  19     Signal.trap("TERM") { EM.stop }
  20   end
  21
  22   # Create a new RackSocket handler
  23   # +app+  The next layer of the Rack stack.
  24   #
  25   # Accepts options:
  26   # +:handler+ (Required) A class to handle new connections.  #initialize will
  27   # call handler.new to create the actual handler instance object.  When a new
  28   # websocket connection is established, #on_connect on the handler instance
  29   # object will be called with the new connection.
  30   #
  31   # +:mount+ The HTTP request path that will be recognized for websocket
  32   # connect requests, defaults to '/websocket'.
  33   #
  34   # +:websocket_only+  If true, the server will only handle websocket requests,
  35   # and all other requests will result in an error.  If false, unhandled
  36   # non-websocket requests will be passed along on to 'app' in the usual Rack
  37   # way.
  38   def initialize(app = nil, options = nil)
  39     @app = app if app.respond_to?(:call)
  40     @options = [app, options].grep(Hash).first || {}
  41     @endpoint = @options[:mount] || DEFAULT_ENDPOINT
  42     @websocket_only = @options[:websocket_only] || false
  43
  44     # from https://gist.github.com/eatenbyagrue/1338545#file-eventmachine-rb
  45     if defined?(PhusionPassenger)
  46       PhusionPassenger.on_event(:starting_worker_process) do |forked|
  47         # for passenger, we need to avoid orphaned threads
  48         if forked && EM.reactor_running?
  49           EM.stop
  50         end
  51         Thread.new do
  52           begin
  53             EM.run
  54           ensure
  55             ActiveRecord::Base.connection.close
  56           end
  57         end
  58         die_gracefully_on_signal
  59       end
  60     else
  61       # faciliates debugging
  62       Thread.abort_on_exception = true
  63       # just spawn a thread and start it up
  64       Thread.new do
  65         begin
  66           EM.run
  67         ensure
  68           ActiveRecord::Base.connection.close
  69         end
  70       end
  71     end
  72
  73     # Create actual handler instance object from handler class.
  74     @handler = @options[:handler].new
  75   end
  76
  77   # Handle websocket connection request, or pass on to the next middleware
  78   # supplied in +app+ initialize (unless +:websocket_only+ option is true, in
  79   # which case return an error response.)
  80   # +env+ the Rack environment with information about the request.
  81   def call env
  82     request = Rack::Request.new(env)
  83     if request.path_info == @endpoint and Faye::WebSocket.websocket?(env)
  84       if @handler.overloaded?
  85         return [503, {"Content-Type" => "text/plain"}, ["Too many connections, try again later."]]
  86       end
  87
  88       ws = Faye::WebSocket.new(env, nil, :ping => 30)
  89
  90       # Notify handler about new connection
  91       @handler.on_connect ws
  92
  93       # Return async Rack response
  94       ws.rack_response
  95     elsif not @websocket_only
  96       @app.call env
  97     else
  98       [406, {"Content-Type" => "text/plain"}, ["Only websocket connections are permitted on this port."]]
  99     end
 100   end
 101
 102 end