DanaDanger-equity 0.4 → 0.5

data/bin/equity

@@ -1,6 +1,6 @@
 #!/usr/bin/env ruby
 #
-#  equity 0.4 - queueing software load balancer with transmission inspection
+#  equity 0.5 - queueing software load balancer with transmission inspection
 #               and simple HTTP header rewriting
 #
 #  Usage: equity [-dp] [-h "Header: Value"] <listen-port> [<node-address>:]<node-port> ...
@@ -25,7 +25,7 @@ end
 
 # Prints a debug message if debugging is enabled.
 def debug(socket, message)
-  return unless $debugging  
+  return unless $debugging
   if socket
     if node = Equity::Node.with_socket(socket)
       client_id = client_id(node)
@@ -124,7 +124,7 @@ options.each do |option, argument|
 end
 
 if ARGV.length < 2
-  STDERR.puts 'equity 0.4'
+  STDERR.puts 'equity 0.5'
   STDERR.puts 'Usage: equity [<options>] <listen-port> [<node-address>:]<node-port> ...'
   STDERR.print "\n"
   STDERR.puts 'Node addresses default to localhost.'
@@ -185,7 +185,7 @@ trap 'SIGINT', Proc.new {
 }
 
 # Start up control server.
-$control_server = Equity::Manager.new($nodes, $listen_port)
+$control_server = Equity::Manager.new($nodes, $listen_port, $debugging)
 
 # Start up server.
 $client_queue = []
@@ -221,8 +221,8 @@ while true
           mate.write(data)
         rescue Errno::EPIPE, Errno::ECONNRESET
           # Connection closed. Disconnect this node.
-          node.disconnect
           debug(socket, 'disconnected')
+          node.disconnect
         end
       end
     end

data/bin/equityctl

@@ -15,9 +15,20 @@
 #
 
 require 'equity/controller'
+require 'getoptlong'
 
-# Process arguments.
-if ARGV.length != 2
+# Process options.
+options = GetoptLong.new(
+  ["--genkey", GetoptLong::NO_ARGUMENT]
+)
+options.each do |option, argument|
+  case option
+  when "--genkey"
+    $genkey = true
+  end
+end
+
+if ($genkey && !ARGV.empty?) || (!$genkey && ARGV.length != 2)
   STDERR.puts 'Usage: equityctl [<host>:]<port> <command>'
   STDERR.print "\n"
   STDERR.puts 'The host defaults to localhost.'
@@ -25,9 +36,38 @@ if ARGV.length != 2
   STDERR.puts 'Command    Description'
   STDERR.puts 'status     Displays each node\'s address, port, connection status,'
   STDERR.puts '           connection counter, and failure counter.'
+  STDERR.puts 'shutdown   Close connections immediately and shut down.'
+  STDERR.print "\n\n"
+  STDERR.puts 'Usage: equityctl --genkey'
+  STDERR.print "\n"
+  STDERR.puts 'With the --genkey option, equityctl generates a public/private key pair for'
+  STDERR.puts 'authenticating itself to equity.'
   exit(64) # EX_USAGE
 end
 
+# If --genkey was specified, generate a key pair and exit.
+if $genkey
+  puts "Generating key..."
+  begin
+    key = Equity::Controller::Key.generate
+  rescue Exception => e
+    STDERR.puts e.message
+    exit(1)
+  end
+  puts "Done! For reference, your key is stored here:"
+  key.paths.each {|path| puts "  #{path}"}
+  # Tell the user where they need to put their public key in order for it to be
+  # useful.
+  puts "\nTo be authorized to use equityctl with your own equity processes, your public"
+  puts "key must be appended to:"
+  puts "  " + Equity::Controller::Key::PERSONAL_AUTHORIZED_KEYS_PATH  
+  puts "\nTo be authorized to use equityctl with other people's equity processes, your"
+  puts "public key must be appended to:"
+  puts "  " + Equity::Controller::Key::SYSTEM_AUTHORIZED_KEYS_PATH
+  exit
+end
+
+# Extract address, port, and command from arguments.
 address, port = ARGV.shift.split(':', 2)
 if port.nil?
   port = address
@@ -47,6 +87,8 @@ begin
     nodes.each do |node|
       puts "#{node} #{node.connected? ? 'C' : '-'} #{node.counter} #{node.failure_counter}"
     end
+  when 'shutdown'
+    controller.shutdown || raise(NoResponseError)
   else
     STDERR.puts "Unrecognized command: #{command}"
     STDERR.puts 'Run equityctl with no arguments for a list of commands.'

data/lib/equity/controller.rb

@@ -2,13 +2,15 @@ require 'socket'
 require 'dl/import'
 require 'equity/manager'
 require 'equity/node'
+require 'equity/controller/key'
 
 module Equity
   class Controller
-    def initialize(host, port)
+    def initialize(host, port, key = Key.new)
       @host = host.dup
       @port = port.to_i
       @socket = UDPSocket.new
+      @key = key
     end
     
     def node_status
@@ -28,6 +30,11 @@ module Equity
       nodes
     end
     
+    def shutdown
+      reply = send(Manager::CMD_SHUTDOWN)
+      !!reply
+    end
+    
     module Alarm
       extend DL::Importable
       if RUBY_PLATFORM =~ /darwin/
@@ -42,6 +49,7 @@ module Equity
   private
     
     def send(message)
+      message = @key.sign(message)
       @socket.send(message, 0, @host, @port)
       begin
         trap('ALRM') { raise Errno::EINTR }

data/lib/equity/controller/key.rb

@@ -0,0 +1,177 @@
+require 'openssl'
+require 'fileutils'
+require 'base64'
+
+module Equity
+  class Controller
+    class Key
+      
+      # Length, in bits, of generated keys.
+      LENGTH = 2048
+      
+      # Default file system path where the user's key is stored. This path is
+      # the private half of the key. The public key can be found by appending
+      # ".pub" to this path.
+      DEFAULT_IDENTITY_PATH = File.join(ENV['HOME'], ".equity", "id_dsa")
+      
+      # Array of paths that contain the lists of public keys that identify
+      # people who are authorized to control equity.
+      PERSONAL_AUTHORIZED_KEYS_PATH = File.join(ENV['HOME'], ".equity", "authorized_keys")
+      SYSTEM_AUTHORIZED_KEYS_PATH = "/etc/equity/authorized_keys"
+      
+      # Byte sequence that separates the data and the signature in signed data.
+      SIGNATURE_SEPARATOR = 0.chr
+      
+      # Generates a new DSA key pair and stores it in a pair of files at the
+      # default identity path. Returns the generated key.
+      def self.generate
+        # Generate a DSA key pair.
+        dsa = OpenSSL::PKey::DSA.generate(LENGTH)
+        # Write the private key to disk.
+        FileUtils.mkdir_p(File.dirname(DEFAULT_IDENTITY_PATH))
+        File.open(DEFAULT_IDENTITY_PATH, "w") do |io|
+          io.chmod(0600)
+          io.write(dsa.to_pem)
+        end
+        # Write the public key to disk.
+        File.open(DEFAULT_IDENTITY_PATH + ".pub", "w") do |io|
+          io.write(dsa.public_key.to_pem)
+        end
+        # Return a Key object with the newly generated key.
+        new
+      end
+      
+      # Loads a key from disk or a string. If the loaded key includes the
+      # private half of the key and the key has been loaded from disk, the
+      # file's permissions are verified to ensure that only the owner can
+      # access it. Raises a SecurityError if the file's permissions are unsafe.
+      def initialize(path_or_pem = DEFAULT_IDENTITY_PATH)
+        # Try to interpret path_or_pem as a PEM string first.
+        begin
+          @dsa = OpenSSL::PKey::DSA.new(path_or_pem)
+        rescue
+          # That didn't work. It must be a path.
+          @path = path_or_pem
+          @dsa = OpenSSL::PKey::DSA.new(IO.read(@path))
+          # If this is a private key, make sure nobody else can read it.
+          if @dsa.private?
+            mode = File.stat(@path).mode
+            unless mode & 077 == 0
+              raise SecurityError, "unsafe permissions on #{@path} - must not allow any access by group or world"
+            end
+          end
+        end
+      end
+      
+      # Returns an array of paths to all files associated with this key.
+      def paths
+        return [] if @path.nil?
+        if @dsa.private?
+          [@path, @path + ".pub"]
+        else
+          [@path + ".pub"]
+        end
+      end
+      
+      # Appends a cryptographic signature to +data+.
+      def sign(data)
+        sig = @dsa.sign(OpenSSL::Digest::DSS1.new, data)
+        a1sig = OpenSSL::ASN1.decode( sig )
+
+        sig_r = a1sig.value[0].value.to_s(2)
+        sig_s = a1sig.value[1].value.to_s(2)
+
+        if sig_r.length > 20 || sig_s.length > 20
+          raise OpenSSL::PKey::DSAError, "bad sig size"
+        end
+
+        sig_r = "\0" * ( 20 - sig_r.length ) + sig_r if sig_r.length < 20
+        sig_s = "\0" * ( 20 - sig_s.length ) + sig_s if sig_s.length < 20
+
+        return [
+          data,
+          Base64.encode64(sig_r + sig_s),
+          @dsa.public_key.to_pem
+        ].join(SIGNATURE_SEPARATOR)
+      end
+      
+      # Verifies the cryptographic signature appended to +data_with_sig+. If
+      # the signature is legitimate, the data is returned. Otherwise this
+      # method returns +false+.
+      def self.verify(data_with_sig)
+        data, sig, public_pem = data_with_sig.split(/#{SIGNATURE_SEPARATOR}/)
+        sig = Base64.decode64(sig)
+        dsa = OpenSSL::PKey::DSA.new(public_pem)
+        sig_r = sig[0,20].unpack("H*")[0].to_i(16)
+        sig_s = sig[20,20].unpack("H*")[0].to_i(16)
+        a1sig = OpenSSL::ASN1::Sequence([
+           OpenSSL::ASN1::Integer(sig_r),
+           OpenSSL::ASN1::Integer(sig_s)
+        ])
+        if dsa.verify(OpenSSL::Digest::DSS1.new, a1sig.to_der, data)
+          data
+        else
+          false
+        end
+      end
+      
+      # Returns true if this key is authorized to control equity.
+      def authorized?
+        self.class.authorized_keys.include?(self)
+      end
+      
+      # Returns an array of keys which are authorized to control equity. Raises
+      # a SecurityError if either of the authorized key files are writable by
+      # group or world, or if a private key is found in the files.
+      def self.authorized_keys
+        # Return cached keys if they've already been loaded.
+        return @authorized_keys if @authorized_keys
+        # Read the keys and verify file permissions.
+        pem = ""
+        if File.exist?(PERSONAL_AUTHORIZED_KEYS_PATH)
+          mode = File.stat(PERSONAL_AUTHORIZED_KEYS_PATH).mode
+          unless mode & 022 == 0
+            raise SecurityError, "unsafe permissions on #{PERSONAL_AUTHORIZED_KEYS_PATH} - must not be writable by group or world"
+          end
+          pem = IO.read(PERSONAL_AUTHORIZED_KEYS_PATH) + "\n"
+        end
+        if File.exist?(SYSTEM_AUTHORIZED_KEYS_PATH)
+          mode = File.stat(SYSTEM_AUTHORIZED_KEYS_PATH).mode
+          unless mode & 022 == 0
+            raise SecurityError, "unsafe permissions on #{SYSTEM_AUTHORIZED_KEYS_PATH} - must not be writable by group or world"
+          end
+          pem += IO.read(SYSTEM_AUTHORIZED_KEYS_PATH)
+        end
+        # Instantiate the keys and verify that they're all public.
+        pems = pem.split(/-----END DSA PUBLIC KEY-----/)
+        pems.pop
+        pems.collect! do |pem|
+          key = new(pem + "-----END DSA PUBLIC KEY-----")
+          if key.private?
+            raise SecurityError, "private key found in authorized keys - only public keys should be authorized"
+          end
+          key
+        end
+        # Cache the loaded keys.
+        @authorized_keys = pems
+      end
+      
+      # Returns true if this key includes the private half of the key.
+      def private?
+        @dsa.private?
+      end
+      
+      # Two keys are equal if they have the same public key.
+      def ==(other)
+        @dsa.public_key.to_s == other.instance_eval {@dsa.public_key.to_s}
+      end
+      
+      # Returns the public half of the key that signed +data_with_sig+.
+      def self.signature_key(data_with_sig)
+        data, sig, public_pem = data_with_sig.split(/#{SIGNATURE_SEPARATOR}/)
+        new(public_pem)
+      end
+      
+    end
+  end
+end

data/lib/equity/manager.rb

@@ -1,24 +1,51 @@
+require "equity/controller/key"
+
 module Equity
   class Manager
     
-    CMD_NODE_STATUS = "\000"
+    CMD_NODE_STATUS = "S"
+    CMD_SHUTDOWN = "K"
     
-    def initialize(nodes, port)
+    CMD_NAMES = {
+      CMD_NODE_STATUS => "status",
+      CMD_SHUTDOWN    => "shutdown"
+    }
+    
+    def initialize(nodes, port, debugging = false)
       @nodes = nodes
       @socket = UDPSocket.new
       @socket.bind(nil, port)
+      @debugging = debugging
     end
     
     def process_commands
+      # Read commands from the network.
       begin
-        data, address = @socket.recvfrom_nonblock(1024)
+        data, address = @socket.recvfrom_nonblock(4096)
       rescue Errno::EAGAIN
       end
       return unless data
-      
+      command = data.split(/#{Controller::Key::SIGNATURE_SEPARATOR}/).first
+      debug("received manager command: #{command_name(command)}", address)
+      # Verify the signature and ensure that the key used to sign the command
+      # is authorized.
+      unless verified_data = Controller::Key.verify(data)
+        debug("signature verification failed - ignoring command", address)
+        return
+      end
+      unless Controller::Key.signature_key(data).authorized?
+        debug("command signed by unauthorized key - ignoring command", address)
+        return
+      end
+      data = verified_data
+      # Process the command.
+      debug("signature verified - processing command", address)
       case data
       when CMD_NODE_STATUS
         reply = node_status_message
+      when CMD_SHUTDOWN
+        UDPSocket.new.send(generic_ok_message, 0, address[2], address[1])
+        exit
       else
         return
       end
@@ -41,5 +68,20 @@ module Equity
       message
     end
     
+    def generic_ok_message
+      "OK"
+    end
+    
+    def debug(message, sockaddr)
+      return unless @debugging
+      nameinfo = Socket.getnameinfo(sockaddr,
+        Socket::NI_NUMERICHOST | Socket::NI_NUMERICSERV)
+      puts "[#{nameinfo[0]}:#{nameinfo[1]}] #{message}"
+    end
+    
+    def command_name(command)
+      CMD_NAMES[command] || "unknown"
+    end
+    
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification 
 name: DanaDanger-equity
 version: !ruby/object:Gem::Version 
-  version: "0.4"
+  version: "0.5"
 platform: ruby
 authors: 
 - DanaDanger
@@ -26,6 +26,7 @@ files:
 - bin/equity
 - bin/equityctl
 - lib/equity/controller.rb
+- lib/equity/controller/key.rb
 - lib/equity/manager.rb
 - lib/equity/node.rb
 - lib/equity/socket_pairing.rb