now with daemonize and additional options

other/websocket.rb

@@ -50,6 +50,10 @@ Sec-WebSocket-Accept: %s\r
 
 
   def initialize(opts)
+    @verbose = opts['verbose']
+    @opts = opts
+    @@client_id = 0  # Track client number total on class
+    
     vmsg "in WebSocketServer.initialize"
     port = opts['listen_port']
     host = opts['listen_host'] || GServer::DEFAULT_HOST
@@ -65,11 +69,6 @@ Sec-WebSocket-Accept: %s\r
       @sslContext.verify_mode = OpenSSL::SSL::VERIFY_NONE
       @sslContext.verify_depth = 0
     end
-
-    @@client_id = 0  # Track client number total on class
-
-    @verbose = opts['verbose']
-    @opts = opts
   end
   
   def serve(io)
@@ -119,7 +118,7 @@ Sec-WebSocket-Accept: %s\r
   end
 
   def msg(m)
-    printf("% 3d: %s\n", Thread.current[:my_client_id] || 0, m)
+    log sprintf("% 3d: %s\n", Thread.current[:my_client_id] || 0, m)
   end
 
   def vmsg(m)

other/websockify.rb

@@ -9,13 +9,15 @@ $: << "other"
 $: << "../other"
 require 'websocket'
 require 'optparse'
+require 'fileutils'
 
 # Proxy traffic to and from a WebSockets client to a normal TCP
 # socket server target. All traffic to/from the client is base64
 # encoded/decoded to allow binary data to be sent/received to/from
 # the target.
 class WebSocketProxy < WebSocketServer
-  attr_accessor :target_host, :target_port
+  attr_accessor 'target_host', 'target_port'
+  attr_reader :opts, :quit
 
   @@Traffic_legend = "
 Traffic Legend:
@@ -32,21 +34,140 @@ Traffic Legend:
 
   def initialize(opts)
     vmsg "in WebSocketProxy.initialize"
+    @opts = opts
     
     super(opts)
 
-    @target_host = opts["target_host"]
-    @target_port = opts["target_port"]
+    @target_host = opts['target_host']
+    @target_port = opts['target_port']
+    
+    #replace the logfile option with the expanded path
+    opts[:logfile] = File.expand_path(logfile) if logfile?
+    opts[:pidfile] = File.expand_path(pidfile) if pidfile?
+  end
+  
+  def daemonize?
+    opts[:daemonize]
+  end
+
+  def logfile
+    opts[:logfile]
+  end
+
+  def pidfile
+    opts[:pidfile]
+  end
+
+  def logfile?
+    !logfile.nil?
+  end
+
+  def pidfile?
+    !pidfile.nil?
+  end
+  
+  def write_pid
+    if pidfile?
+      begin
+        File.open(pidfile, ::File::CREAT | ::File::EXCL | ::File::WRONLY){|f| f.write("#{Process.pid}") }
+        at_exit { File.delete(pidfile) if File.exists?(pidfile) }
+      rescue Errno::EEXIST
+        check_pid
+        retry
+      end
+    end
+  end
+  
+  def check_pid
+    if pidfile?
+      case pid_status(pidfile)
+      when :running, :not_owned
+        log "A server is already running. Check #{pidfile}"
+        exit(1)
+      when :dead
+        File.delete(pidfile)
+      end
+    end
+  end
+
+  def pid_status(pidfile)
+    return :exited unless File.exists?(pidfile)
+    pid = ::File.read(pidfile).to_i
+    return :dead if pid == 0
+    Process.kill(0, pid)      # check process status
+    :running
+  rescue Errno::ESRCH
+    :dead
+  rescue Errno::EPERM
+    :not_owned
+  end
+  
+  def redirect_output
+    FileUtils.mkdir_p(File.dirname(logfile), :mode => 0755)
+    FileUtils.touch logfile
+    File.chmod(0644, logfile)
+    $stderr.reopen(logfile, 'a')
+    $stdout.reopen($stderr)
+    $stdout.sync = $stderr.sync = true
+  end
+  
+  def suppress_output
+    $stderr.reopen('/dev/null', 'a')
+    $stdout.reopen($stderr)
+  end
+  
+  def trap_signals
+    trap(:QUIT) do   # kill -9
+      log "Hard killing websockify server..."
+      stop
+      log "Bye!"
+    end
+    
+    trap(:TERM) do   # kill -15
+      log "Gracefully shutting down websockify server..."
+      Thread.new { shutdown; exit }
+      log "Bye!"
+    end
+  end
+  
+  def daemonize
+    exit if fork
+    Process.setsid
+    exit if fork
+    Dir.chdir "/"
+  end
+  
+  def run!(max_connections)
+    check_pid
+    daemonize if daemonize?
+    write_pid
+    trap_signals
+    
+    if logfile?
+      redirect_output
+    elsif daemonize?
+      suppress_output
+    end
+    
+    log "Starting Websockify server on #{opts['listen_host']}:#{opts['listen_port']} and proxying to #{target_host}:#{target_port}"
+    
+    #tell gserver to start
+    start(max_connections)
+    
+    #join the gserver thread to this one
+    join
   end
 
   # Echo back whatever is received
   def new_websocket_client(client)
     path = Thread.current[:path]
+    msg "path = #{path}"
     if path =~ /:/
-      msg "Path has a ip and port specified. Using those for target"
+      msg "path has a ip and port"
       possible_target_host = path[/(\d+\.\d+\.\d+\.\d+):(\d+)/,1]
       possible_target_port = path[/(\d+\.\d+\.\d+\.\d+):(\d+)/,2].to_i
     end
+    Thread.current[:path] = "/"
 
     if possible_target_host and possible_target_port
       @target_host = possible_target_host
@@ -56,7 +177,7 @@ Traffic Legend:
     msg "connecting to: %s:%s" % [@target_host, @target_port]
     tsock = TCPSocket.open(@target_host, @target_port)
 
-    if @verbose then puts @@Traffic_legend end
+    if @verbose then log @@Traffic_legend end
 
     begin
       do_proxy(client, tsock)
@@ -134,14 +255,48 @@ Traffic Legend:
 end
 
 # Parse parameters
+
 opts           = {}
-parser = OptionParser.new do |o|
-  o.on('--verbose', '-v') { |b| opts['verbose'] = b }
-  o.parse!
-end
+version        = "1.0.0"
+maxcon_help = "maximum number of connections allowed"
+daemonize_help = "run daemonized in the background (default: false)"
+pidfile_help   = "the pid filename"
+logfile_help   = "the log filename"
+include_help   = "an additional $LOAD_PATH"
+debug_help     = "set $DEBUG to true"
+warn_help      = "enable warnings"
+usage          = "Usage: server list_port target_host'target_port' [options]"
+
+op = OptionParser.new
+op.banner =  "Websockify Server #{version}"
+op.separator ""
+op.separator "#{usage}"
+op.separator ""
+
+op.separator "Process options:"
+op.on("-m", "--maxcon NUMBER", maxcon_help)    { |value| opts[:maxcon]    = value || 100 }
+op.on("-d", "--daemonize",     daemonize_help) {         opts[:daemonize] = true  }
+op.on("-p", "--pid PIDFILE",   pidfile_help)   { |value| opts[:pidfile]   = value }
+op.on("-l", "--log LOGFILE",   logfile_help)   { |value| opts[:logfile]   = value }
+op.separator ""
+
+op.separator "Ruby options:"
+op.on("-I", "--include PATH", include_help) { |value| $LOAD_PATH.unshift(*value.split(":").map{|v| File.expand_path(v)}) }
+op.on(      "--debug",        debug_help)   { $DEBUG = true }
+op.on(      "--warn",         warn_help)    { $-w = true    }
+op.separator ""
+
+op.separator "Common options:"
+op.on("-h", "--help")    { puts op.to_s; exit }
+op.on("-V", "--version") { puts version; exit }
+op.on("-v", "--verbose")   { |value| opts['verbose'] = true }
+op.separator ""
+
+op.parse!(ARGV)
 
 if ARGV.length < 2
-  puts "Too few arguments"
+  puts "Too few arguments."
+  puts op.to_s
   exit 2
 end
 
@@ -149,7 +304,7 @@ end
 if ARGV[0].count(":") > 0
   opts['listen_host'], _, opts['listen_port'] = ARGV[0].rpartition(':')
 else
-  opts['listen_host'], opts['listen_port'] = nil, ARGV[0]
+  opts['listen_host'], opts['listen_port'] = '0.0.0.0', ARGV[0]
 end
 
 begin
@@ -173,11 +328,7 @@ rescue
   exit 2
 end
 
-puts "Starting server on #{opts['listen_host']}:#{opts['listen_port']}"
 server = WebSocketProxy.new(opts)
-server.start(100)
-server.join
-
-puts "Server has been terminated"
+server.run!(opts[:maxcon].to_i)
 
 # vim: sw=2