statsd: Abstract the Secure Client

Vicent Marti · Vicent Marti · commit 80ac4de0deee · 2016-01-12T14:42:26.000+01:00
All the secure client implementation has been moved into the
SecureUDPClient class; the `StatsD` core no longer has any
secure-related features. Hence, the secure client must be initialized
directly with the shared key, or with `StatsD#add_shard`, which has been
updated to transparently forward its arguments to the client
constructor.

The hashing/HMAC is now performed right before writing to the socket,
which allows for greater composability.
diff --git a/lib/statsd.rb b/lib/statsd.rb
@@ -14,14 +14,52 @@
 #   statsd = Statsd.new('localhost').tap{|sd| sd.namespace = 'account'}
 #   statsd.increment 'activate'
 class Statsd
-  class RubyUdpClient
-    attr_reader :key, :sock
+  class SecureUDPClient < UDPClient
+    def initialize(address, port, key)
+      super(address, port)
+      @key = key
+    end
+
+    def send(msg)
+      super(signed_payload(msg))
+    end
+
+    private
+    # defer loading openssl and securerandom unless needed. this shaves ~10ms off
+    # of baseline require load time for environments that don't require message signing.
+    def self.setup_openssl
+      @sha256 ||= begin
+        require 'securerandom'
+        require 'openssl'
+        OpenSSL::Digest::SHA256.new
+      end
+    end
+
+    def signed_payload(message)
+      sha256 = SecureUDPClient.setup_openssl
+      payload = timestamp + nonce + message
+      signature = OpenSSL::HMAC.digest(sha256, @key, payload)
+      signature + payload
+    end
+
+    def timestamp
+      [Time.now.to_i].pack("Q<")
+    end
+
+    def nonce
+      SecureRandom.random_bytes(4)
+    end
+  end
+
+  class UDPClient
+    attr_reader :sock
 
-    def initialize(address, port, key = nil)
+    def initialize(address, port = nil)
+      address, port = address.split(':') if address.include?(':')
       addrinfo = Addrinfo.ip(address)
+
       @sock = UDPSocket.new(addrinfo.pfamily)
       @sock.connect(addrinfo.ip_address, port)
-      @key = key
     end
 
     def send(msg)
@@ -52,17 +90,16 @@ def namespace=(namespace)
 
   def initialize(client_class = nil)
     @shards = []
-    @client_class = client_class || RubyUdpClient
+    @client_class = client_class || UDPClient
     self.namespace = nil
   end
 
   def self.simple(addr, port = nil)
     self.new.add_shard(addr, port)
   end
 
-  def add_shard(addr, port = nil, key = nil)
-    addr, port = addr.split(':') if addr.include?(':')
-    @shards << @client_class.new(addr, port.to_i, key)
+  def add_shard(*args)
+    @shards << @client_class.new(*args)
     self
   end
 
@@ -129,7 +166,6 @@ def time(stat, sample_rate=1)
   def histogram(stat, value, sample_rate=1); send stat, value, HISTOGRAM_TYPE, sample_rate end
 
   private
-
   def sampled(sample_rate)
     yield unless sample_rate < 1 and rand > sample_rate
   end
@@ -140,7 +176,7 @@ def send(stat, delta, type, sample_rate=1)
       stat.gsub!(/::/, ".".freeze)
       stat.gsub!(RESERVED_CHARS_REGEX, "_".freeze)
 
-      msg = ""
+      msg = String.new
       msg << @prefix
       msg << stat
       msg << ":".freeze
@@ -153,7 +189,7 @@ def send(stat, delta, type, sample_rate=1)
       end
 
       shard = select_shard(stat)
-      shard.send(shard.key ? signed_payload(shard.key, msg) : msg)
+      shard.send(msg)
     end
   end
 
@@ -164,29 +200,4 @@ def select_shard(stat)
       @shards[Zlib.crc32(stat) % @shards.size]
     end
   end
-
-  def signed_payload(key, message)
-    sha256 = Statsd.setup_openssl
-    payload = timestamp + nonce + message
-    signature = OpenSSL::HMAC.digest(sha256, key, payload)
-    signature + payload
-  end
-
-  # defer loading openssl and securerandom unless needed. this shaves ~10ms off
-  # of baseline require load time for environments that don't require message signing.
-  def self.setup_openssl
-    @sha256 ||= begin
-      require 'securerandom'
-      require 'openssl'
-      OpenSSL::Digest::SHA256.new
-    end
-  end
-
-  def timestamp
-    [Time.now.to_i].pack("Q<")
-  end
-
-  def nonce
-    SecureRandom.random_bytes(4)
-  end
 end
