Add auth token support to the control server

lib/puma/cli.rb

@@ -142,6 +142,11 @@ module Puma
           end
         end
 
+        o.on "--control-token TOKEN",
+             "The token to use as authentication for the control server" do |arg|
+          @options[:control_auth_token] = arg
+        end
+
         o.on '-t', '--threads INT', "min:max threads to use (default 0:16)" do |arg|
           min, max = arg.split(":")
           if max
@@ -213,10 +218,7 @@ module Puma
         @options[:rackup] = @argv.shift || DefaultRackup
       end
 
-      if @options[:control_url] == "auto"
-        path = @temp_status_path = Configuration.temp_path
-        @options[:control_url] = "unix://#{path}"
-      end
+      @temp_status_path = @options[:control_path_temp]
     end
 
     # Parse the options, load the rackup, start the server and wait
@@ -273,6 +275,11 @@ module Puma
         uri = URI.parse str
 
         app = Puma::App::Status.new server, self
+
+        if token = @options[:control_auth_token]
+          app.auth_token = token unless token.empty? or token == :none
+        end
+
         status = Puma::Server.new app, @events
         status.min_threads = 0
         status.max_threads = 1

lib/puma/config.rb

@@ -25,6 +25,44 @@ module Puma
       if @options[:binds].empty?
         @options[:binds] << "tcp://#{DefaultTCPHost}:#{DefaultTCPPort}"
       end
+
+      if @options[:control_url] == "auto"
+        path = Configuration.temp_path
+        @options[:control_url] = "unix://#{path}"
+        @options[:control_url_temp] = path
+      end
+
+      unless @options[:control_auth_token]
+        setup_random_token
+      end
+    end
+
+    def setup_random_token
+      begin
+        require 'openssl'
+      rescue LoadError
+      end
+
+      count = 16
+
+      bytes = nil
+
+      if defined? OpenSSL::Random
+        bytes = OpenSSL::Random.random_bytes(count)
+      elsif File.exists?("/dev/urandom")
+        File.open("/dev/urandom") do |f|
+          bytes = f.read(count)
+        end
+      end
+
+      if bytes
+        token = ""
+        bytes.each_byte { |b| token << b.to_s(16) }
+      else
+        token = (0..count).to_a.map { rand(255).to_s(16) }.join
+      end
+
+      @options[:control_auth_token] = token
     end
 
     def self.temp_path
@@ -48,8 +86,18 @@ module Puma
     # Start the Puma control rack app on +url+. This app can be communicated
     # with to control the main server.
     #
-    def activate_control_app(url="auto")
+    def activate_control_app(url="auto", opts=nil)
       @options[:control_url] = url
+
+      if opts
+        if tok = opts[:auth_token]
+          @options[:control_auth_token] = tok
+        end
+
+        if opts[:no_token]
+          @options[:control_auth_token] = :none
+        end
+      end
     end
 
     # Bind the server to +url+. tcp:// and unix:// are the only accepted

lib/puma/control_cli.rb

@@ -59,16 +59,34 @@ module Puma
       end
     end
 
+    def request(sock, url)
+      token = @config.options[:control_auth_token]
+      if token
+        url = "#{url}?token=#{token}"
+      end
+
+      sock << "GET #{url} HTTP/1.0\r\n\r\n"
+
+      rep = sock.read.split("\r\n")
+
+      m = %r!HTTP/1.\d (\d+)!.match(rep.first)
+      if m[1] == "403"
+        raise "Unauthorized access to server (wrong auth token)"
+      elsif m[1] != "200"
+        raise "Bad response code from server: #{m[1]}"
+      end
+
+      return rep.last
+    end
+
     def command_pid
       puts "#{@state['pid']}"
     end
 
     def command_stop
       sock = connect
-      sock << "GET /stop HTTP/1.0\r\n\r\n"
-      rep = sock.read
+      body = request sock, "/stop"
 
-      body = rep.split("\r\n").last
       if body != '{ "status": "ok" }'
         raise "Invalid response: '#{body}'"
       else
@@ -78,10 +96,8 @@ module Puma
 
     def command_halt
       sock = connect
-      s << "GET /halt HTTP/1.0\r\n\r\n"
-      rep = s.read
+      body = request sock, "/halt"
 
-      body = rep.split("\r\n").last
       if body != '{ "status": "ok" }'
         raise "Invalid response: '#{body}'"
       else
@@ -91,10 +107,8 @@ module Puma
 
     def command_restart
       sock = connect
-      sock << "GET /restart HTTP/1.0\r\n\r\n"
-      rep = sock.read
+      body = request sock, "/restart"
 
-      body = rep.split("\r\n").last
       if body != '{ "status": "ok" }'
         raise "Invalid response: '#{body}'"
       else
@@ -104,10 +118,7 @@ module Puma
 
     def command_stats
       sock = connect
-      sock << "GET /stats HTTP/1.0\r\n\r\n"
-      rep = sock.read
-
-      body = rep.split("\r\n").last
+      body = request sock, "/stats"
 
       puts body
     end

test/test_cli.rb

@@ -34,7 +34,10 @@ class TestCLI < Test::Unit::TestCase
     sin = StringIO.new
     sout = StringIO.new
 
-    cli = Puma::CLI.new ["-b", "unix://#{@tmp_path2}", "--control", url, "test/lobster.ru"], sin, sout
+    cli = Puma::CLI.new ["-b", "unix://#{@tmp_path2}",
+                         "--control", url,
+                         "--control-token", "",
+                         "test/lobster.ru"], sin, sout
     cli.parse_options
 
     t = Thread.new { cli.run }
@@ -57,7 +60,10 @@ class TestCLI < Test::Unit::TestCase
     sin = StringIO.new
     sout = StringIO.new
 
-    cli = Puma::CLI.new ["-b", "unix://#{@tmp_path2}", "--control", url, "test/lobster.ru"], sin, sout
+    cli = Puma::CLI.new ["-b", "unix://#{@tmp_path2}",
+                         "--control", url,
+                         "--control-token", "",
+                         "test/lobster.ru"], sin, sout
     cli.parse_options
 
     t = Thread.new { cli.run }