DanaDanger-equity 0.3 → 0.4

data/bin/equity

@@ -1,41 +1,145 @@
 #!/usr/bin/env ruby
 #
-#  equity 0.3 - simple queueing software load balancer
+#  equity 0.4 - queueing software load balancer with transmission inspection
+#               and simple HTTP header rewriting
 #
-#  Usage: equity [-d] <listen-port> [<node-address>:]<node-port> ...
-#
-#  The -d option turns on debug mode. Equity will remain in the foreground and
-#  print status messages.
-#
-#  Node addresses default to localhost.
+#  Usage: equity [-dp] [-h "Header: Value"] <listen-port> [<node-address>:]<node-port> ...
+#  Run equity with no arguments for detailed usage information.
 #
 
 require 'socket'
 require 'equity/node'
 require 'equity/manager'
+require 'getoptlong'
+
+# Returns a string that can be used to identify a particular node in debugging
+# messages. At the moment this returns the address and port of the client being
+# served by the node.
+def client_id(node)
+  if node.connected?
+    "#{node.client_address}:#{node.client_port}"
+  else
+    "not connected"
+  end
+end
 
 # Prints a debug message if debugging is enabled.
 def debug(socket, message)
-  return unless $debugging
-  print "[#{socket.object_id.to_s(16)}] " if socket
+  return unless $debugging  
+  if socket
+    if node = Equity::Node.with_socket(socket)
+      client_id = client_id(node)
+    else
+      client_id = nil
+      begin
+        nameinfo = Socket.getnameinfo(socket.getpeername,
+          Socket::NI_NUMERICHOST | Socket::NI_NUMERICSERV)
+        client_id = "#{nameinfo[0]}:#{nameinfo[1]}"
+      rescue
+        client_id = "???:???"
+      end
+    end
+    print "[#{client_id}] "
+  end
   puts message
 end
 
+# Prints data in hex and ASCII if packet dumping is enabled.
+BYTES_PER_PACKET_LINE = 16
+SAFE_PACKET_RX = (32..126).to_a.collect! {|n| n.chr}.join("").sub!(/[\[\]]/, '\\\1')
+def dump_packet(data, src, dst)
+  return unless $dump_packets
+  
+  node = Equity::Node.with_socket(src) || return
+  to_server = (src == node.sockets.first)
+  arrow = (to_server ? "--->" : "<---")
+  print "[#{client_id(node)}] #{arrow} #{node.address}:#{node.port}\n\n"
+  
+  data = data.dup
+  until data.empty?
+    bytes = data.slice!(0, BYTES_PER_PACKET_LINE)
+    hex = bytes.unpack("H*").first + ("  " * (BYTES_PER_PACKET_LINE - bytes.length))
+    hex1 = hex.slice!(0, BYTES_PER_PACKET_LINE)
+    hex2 = hex
+    hex1.gsub!(/(....)/, '\1 ')
+    hex2.gsub!(/(....)/, '\1 ')
+    bytes.gsub!(/[^#{SAFE_PACKET_RX}]/, ".")
+    puts "  #{hex1}  #{hex2}  #{bytes}"
+  end
+  print "\n"
+end
+
+# Handles rewriting of HTTP request headers. This method isn't very smart about
+# how it detects HTTP requests, but should work for most situations.
+def rewrite_headers!(data, src)
+  return if $header_rewrites.empty?
+  return unless http_slash_idx = data.index(" HTTP/")
+  node = Equity::Node.with_socket(src) || return
+  return unless (src == node.sockets.first)
+  
+  debug(src, "rewriting HTTP request headers")
+  
+  before_headers = data.slice!(0, data.index("\r\n", http_slash_idx))
+  headers = data.slice!(0, data.index("\r\n\r\n"))
+  after_headers = data
+  
+  $header_rewrites.each do |header, value|
+    if headers.sub!(/\r\n#{header}: (.+?)\r\n/m, "\r\n#{header}: #{value}\r\n")
+      debug(src, "  #{header}: #{$1} => #{value}")
+    end
+  end
+  data.replace(before_headers + headers + after_headers)
+rescue
+  # If an exception occurs, ignore it. Rewriting the headers is far less
+  # important than keeping the program alive.
+end
+
 
-# Process arguments and instantiate nodes.
-$debugging = !!ARGV.delete('-d')
+# Process arguments.
+$debugging       = !!ARGV.delete('-d')
+$dump_packets    = !!ARGV.delete('-p')
+$header_rewrites = {}
+
+options = GetoptLong.new(
+  ["-d", "--debug",   GetoptLong::NO_ARGUMENT],
+  ["-p", "--packets", GetoptLong::NO_ARGUMENT],
+  ["-h", "--header",  GetoptLong::REQUIRED_ARGUMENT]
+)
+options.each do |option, argument|
+  case option
+  when "-d"
+    $debugging = true
+  when "-p"
+    $dump_packets = true
+  when "-h"
+    header = argument.split(/: +/)
+    unless header.length == 2
+      STDERR.puts "I don't understand this header rewrite:"
+      STDERR.puts "  #{argument}"
+      STDERR.puts "Run equity with no arguments for help."
+      exit(64) # EX_USAGE
+    end
+    $header_rewrites[header.first] = header.last
+  end
+end
 
 if ARGV.length < 2
-  STDERR.puts 'equity 0.3'
-  STDERR.puts 'Usage: equity [-d] <listen-port> [<node-address>:]<node-port> ...'
-  STDERR.print "\n"
-  STDERR.puts 'The -d option turns on debug mode. Equity will remain in the foreground and'
-  STDERR.puts 'print status messages.'
+  STDERR.puts 'equity 0.4'
+  STDERR.puts 'Usage: equity [<options>] <listen-port> [<node-address>:]<node-port> ...'
   STDERR.print "\n"
   STDERR.puts 'Node addresses default to localhost.'
+  STDERR.print "\n"
+  STDERR.puts 'The following options can be used:'
+  STDERR.puts '-d Debug mode. Stays in the foreground and prints status messages.'
+  STDERR.puts '-p Packet dumper. Stays in the foreground and displays the raw data being'
+  STDERR.puts '   transferred.'
+  STDERR.puts '-h "<header>: <value>" Rewrite an HTTP request header before retransmitting.'
+  STDERR.puts '   Can be used more than once to rewrite several headers.'
+  STDERR.print "\n"
   exit(64) # EX_USAGE
 end
 
+# Instantiate nodes.
 $listen_port = ARGV.shift.to_i
 
 $nodes = []
@@ -49,7 +153,7 @@ ARGV.each do |node_spec|
 end
 
 # Daemonize.
-unless $debugging
+unless $debugging || $dump_packets
   fork && exit
   Process.setsid
   trap 'SIGHUP', 'IGNORE'
@@ -66,6 +170,11 @@ unless $debugging
   STDERR.reopen STDOUT
 end
 
+# Print debugging info for header rewrites.
+$header_rewrites.each do |header, value|
+  debug(nil, "Will rewrite value of HTTP request header `#{header}' to `#{value}'")
+end
+
 # Set up SIGINT handler to print node counters.
 trap 'SIGINT', Proc.new {
   debug(nil, "\nNode Counters")
@@ -106,7 +215,10 @@ while true
         begin
           data = socket.recvfrom(65536)[0]
           raise Errno::EPIPE if data.empty?
-          socket.mate.write(data)
+          mate = socket.mate
+          rewrite_headers!(data, socket)
+          dump_packet(data, socket, mate)
+          mate.write(data)
         rescue Errno::EPIPE, Errno::ECONNRESET
           # Connection closed. Disconnect this node.
           node.disconnect
@@ -127,8 +239,8 @@ while true
       next if node.connected?
       begin
         node.connect(client)
-      rescue
-        debug(client, "connection to #{node} failed")
+      rescue Exception => e
+        debug(client, "connection to #{node} failed: #{e.message}")
         failure_count += 1
         next
       end

data/bin/equityctl

@@ -23,9 +23,9 @@ if ARGV.length != 2
   STDERR.puts 'The host defaults to localhost.'
   STDERR.print "\n"
   STDERR.puts 'Command    Description'
-  STDERR.puts 'status     Displays each node\'s address, port, connection status, and'
-  STDERR.puts '           connection counter.'
-  exit(64)
+  STDERR.puts 'status     Displays each node\'s address, port, connection status,'
+  STDERR.puts '           connection counter, and failure counter.'
+  exit(64) # EX_USAGE
 end
 
 address, port = ARGV.shift.split(':', 2)
@@ -45,7 +45,7 @@ begin
   when 'status'
     nodes = controller.node_status || raise(NoResponseError)
     nodes.each do |node|
-      puts "#{node} #{node.connected? ? 'C' : '-'} #{node.counter}"
+      puts "#{node} #{node.connected? ? 'C' : '-'} #{node.counter} #{node.failure_counter}"
     end
   else
     STDERR.puts "Unrecognized command: #{command}"

data/lib/equity/node.rb

@@ -4,14 +4,18 @@ module Equity
   # Represents a cluster node.
   class Node
     attr_reader :counter
+    attr_reader :failure_counter
     attr_reader :address
     attr_reader :port
+    attr_reader :client_address
+    attr_reader :client_port
   
     def initialize(address, port)
-      @client  = nil
-      @counter = 0
-      @address = address.dup
-      @port    = port.to_i
+      @client          = nil
+      @counter         = 0
+      @failure_counter = 0
+      @address         = address.dup
+      @port            = port.to_i
     end
   
     def connect(client)
@@ -22,10 +26,20 @@ module Equity
         @server.extend SocketPairing
         @server.pair_with(client)
         @counter += 1
+        
+        begin
+          nameinfo = Socket.getnameinfo(@client.getpeername,
+            Socket::NI_NUMERICHOST | Socket::NI_NUMERICSERV)
+          @client_address, @client_port = nameinfo
+        rescue
+          @client_address, @client_port = "???", "???"
+        end
+        
         @server
       rescue Exception => e
         @client = nil
         @server = nil
+        @failure_counter += 1
         raise e
       end
     end
@@ -35,6 +49,8 @@ module Equity
       @server.close
       @client = nil
       @server = nil
+      @client_address = nil
+      @client_port = nil
     end
   
     def connected?
@@ -53,6 +69,13 @@ module Equity
       "#{@address}:#{@port}"
     end
     
+    def self.with_socket(socket)
+      ObjectSpace.each_object(self) do |node|
+        return node if node.owns_socket?(socket)
+      end
+      nil
+    end
+    
     class Static < Node
       attr_writer :counter
       

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification 
 name: DanaDanger-equity
 version: !ruby/object:Gem::Version 
-  version: "0.3"
+  version: "0.4"
 platform: ruby
 authors: 
 - DanaDanger
@@ -9,7 +9,7 @@ autorequire:
 bindir: bin
 cert_chain: []
 
-date: 2008-08-13 00:00:00 -07:00
+date: 2009-04-08 00:00:00 -07:00
 default_executable: 
 dependencies: []