[test] WebPush gem-like integration test for EC + AES GCM

kares · kares · commit ff658dbba522 · 2016-06-07T11:08:58.000+02:00
diff --git a/src/test/ruby/ec/ece.rb b/src/test/ruby/ec/ece.rb
@@ -0,0 +1,136 @@
+class ECE
+
+  KEY_LENGTH=16
+  TAG_LENGTH=16
+  NONCE_LENGTH=12
+  SHA256_LENGTH=32
+
+  def self.hmac_hash(key, input)
+    digest = OpenSSL::Digest.new('sha256')
+    OpenSSL::HMAC.digest(digest, key, input)
+  end
+
+  def self.hkdf_extract(salt, ikm) #ikm stays for input keying material
+    hmac_hash(salt,ikm)
+  end
+
+  def self.get_info(type, client_public, server_public)
+    cl_len_no = [client_public.size].pack('n')
+    sv_len_no = [server_public.size].pack('n')
+    "Content-Encoding: #{type}\x00P-256\x00#{cl_len_no}#{client_public}#{sv_len_no}#{server_public}"
+  end
+
+  def self.extract_key(params)
+    raise "Salt must be 16-bytes long" unless params[:salt].length==16
+
+    input_key = params[:key]
+    auth = false
+    if params.has_key?(:auth) # Encrypted Content Encoding, March 11 2016, http://httpwg.org/http-extensions/draft-ietf-httpbis-encryption-encoding.html
+      auth = true
+      input = HKDF.new(input_key, {salt: params[:auth] , algorithm: 'sha256', info: "Content-Encoding: auth\x00"})
+      input_key = input.next_bytes(SHA256_LENGTH)
+      secret =  HKDF.new(input_key, {salt: params[:salt], algorithm: 'sha256', info: get_info("aesgcm", params[:user_public_key], params[:server_public_key])})
+      nonce =  HKDF.new(input_key, salt: params[:salt], algorithm: 'sha256', info: get_info("nonce", params[:user_public_key], params[:server_public_key]))
+    else
+      secret =  HKDF.new(input_key, {salt: params[:salt], algorithm: 'sha256', info: "Content-Encoding: aesgcm128"})
+      nonce =  HKDF.new(input_key, salt: params[:salt], algorithm: 'sha256', info: "Content-Encoding: nonce")
+    end
+
+    {key: secret.next_bytes(KEY_LENGTH), nonce: nonce.next_bytes(NONCE_LENGTH), auth: auth}
+  end
+
+  def self.generate_nonce(nonce, counter)
+    raise "Nonce must be #{NONCE_LENGTH} bytes long." unless nonce.length == NONCE_LENGTH
+    output = nonce.dup
+    integer = nonce[-6..-1].unpack('B*')[0].to_i(2) #taking last 6 bytes, treating as integer
+    x = ((integer ^ counter) & 0xffffff) + ((((integer / 0x1000000) ^ (counter / 0x1000000)) & 0xffffff) * 0x1000000)
+    bytestring = x.to_s(16).length < 12 ? "0"*(12-x.to_s(16).length)+x.to_s(16) : x.to_s(16) #it's for correct handling of cases when generated integer is less than 6 bytes
+    output[-6..-1] = [bytestring].pack('H*')                                                 #without it packing would produce less than 6 bytes
+    output                                                                                   #I didn't find pack directive for such usage, so there is a such solution
+  end
+
+  def self.encrypt(data, params)
+    key = extract_key(params)
+    rs = params[:rs] ? params [:rs] : 4096
+    padsize = params[:padsize] ? params [:padsize] : 0
+    raise "The rs parameter must be greater than 1." if rs <= 1
+    rs -=1 #this ensures encrypted data cannot be truncated
+    result = ""
+    pad_bytes = 1
+    if params[:auth] # old spec used 1 byte for padding, latest one always uses 2 bytes
+      pad_bytes = 2
+    end
+
+    counter = 0
+    (0..data.length).step(rs-pad_bytes+1) do |i|
+      block = encrypt_record(key, counter, data[i..i+rs-pad_bytes], padsize)
+      result += block
+      counter +=1
+    end
+    result
+  end
+
+  def self.decrypt(data, params)
+    key = extract_key(params)
+    rs = params[:rs] ? params [:rs] : 4096
+    raise "The rs parameter must be greater than 1." if rs <= 1
+    rs += TAG_LENGTH
+    raise "Message is truncated" if data.length % rs == 0
+    result = ""
+    counter = 0
+    (0..data.length).step(rs) do |i|
+      block = decrypt_record(key, counter, data[i..i+rs-1])
+      result += block
+      counter +=1
+    end
+    result
+  end
+
+  def self.decrypt_record(params, counter, buffer, pad=0)
+    gcm = OpenSSL::Cipher.new('aes-128-gcm')
+    gcm.decrypt
+    gcm.key = params[:key]
+    gcm.iv = generate_nonce(params[:nonce], counter)
+    pad_bytes = 1
+    if params[:auth] # old spec used 1 byte for padding, latest one always uses 2 bytes
+      pad_bytes = 2
+    end
+    raise "Block is too small" if buffer.length <= TAG_LENGTH+pad_bytes
+    gcm.auth_tag = buffer[-TAG_LENGTH..-1]
+    decrypted = gcm.update(buffer[0..-TAG_LENGTH-1]) + gcm.final
+
+    if params[:auth]
+      padding_length = decrypted[0..1].unpack("n")[0]
+      raise "Padding is too big" if padding_length+2 > decrypted.length
+      padding = decrypted[2..padding_length]
+      raise "Wrong padding"  unless padding = "\x00"*padding_length
+      return decrypted[2+padding_length..-1]
+    else
+      padding_length = decrypted[0].unpack("C")[0]
+      raise "Padding is too big" if padding_length+1 > decrypted.length
+      padding = decrypted[1..padding_length]
+      raise "Wrong padding"  unless padding = "\x00"*padding_length
+      return decrypted[1..-1]
+    end
+  end
+
+  def self.encrypt_record(params, counter, buffer, pad=0)
+    gcm = OpenSSL::Cipher.new('aes-128-gcm')
+    gcm.encrypt
+    gcm.key = params[:key]
+    gcm.iv = generate_nonce(params[:nonce], counter)
+    gcm.auth_data = ""
+    padding = ""
+    if params[:auth]
+      padding = [pad].pack('n') + "\x00"*pad # 2 bytes, big endian, then n zero bytes of padding
+      buf = padding+buffer
+      record = gcm.update(buf)
+    else
+      record = gcm.update("\x00"+buffer) # 1 padding byte, not fully implemented
+    end
+    enc = record + gcm.final + gcm.auth_tag
+    enc
+  end
+
+
+end
diff --git a/src/test/ruby/ec/hkdf.rb b/src/test/ruby/ec/hkdf.rb
@@ -0,0 +1,74 @@
+require 'stringio'
+
+class HKDF
+  DefaultAlgorithm = 'SHA256'
+  DefaultReadSize = 512 * 1024
+
+  def initialize(source, options = {})
+    source = StringIO.new(source) if source.is_a?(String)
+
+    algorithm = options.fetch(:algorithm, DefaultAlgorithm)
+    @digest = OpenSSL::Digest.new(algorithm)
+    @info = options.fetch(:info, '')
+
+    salt = options[:salt]
+    salt = 0.chr * @digest.digest_length if salt.nil? or salt.empty?
+    read_size = options.fetch(:read_size, DefaultReadSize)
+
+    @prk = _generate_prk(salt, source, read_size)
+    @position = 0
+    @blocks = []
+    @blocks << ''
+  end
+
+  def algorithm
+    @digest.name
+  end
+
+  def max_length
+    @max_length ||= @digest.digest_length * 255
+  end
+
+  def seek(position)
+    raise RangeError.new("cannot seek past #{max_length}") if position > max_length
+
+    @position = position
+  end
+
+  def rewind
+    seek(0)
+  end
+
+  def next_bytes(length)
+    new_position = length + @position
+    raise RangeError.new("requested #{length} bytes, only #{max_length} available") if new_position > max_length
+
+    _generate_blocks(new_position)
+
+    start = @position
+    @position = new_position
+
+    @blocks.join('').slice(start, length)
+  end
+
+  def next_hex_bytes(length)
+    next_bytes(length).unpack('H*').first
+  end
+
+  def _generate_prk(salt, source, read_size)
+    hmac = OpenSSL::HMAC.new(salt, @digest)
+    while block = source.read(read_size)
+      hmac.update(block)
+    end
+    hmac.digest
+  end
+
+  def _generate_blocks(length)
+    start = @blocks.size
+    block_count = (length.to_f / @digest.digest_length).ceil
+    start.upto(block_count) do |n|
+      @blocks << OpenSSL::HMAC.digest(@digest, @prk, @blocks[n - 1] + @info + n.chr)
+    end
+  end
+end
+
diff --git a/src/test/ruby/ec/test_ec.rb b/src/test/ruby/ec/test_ec.rb
@@ -81,6 +81,116 @@ def test_encrypt
     assert_equal expected, server.dh_compute_key(client_public_key)
   end
 
+  def test_encrypt_integration # inspired by WebPush
+    require File.expand_path('ece.rb', File.dirname(__FILE__)) unless defined? ECE
+    require File.expand_path('hkdf.rb', File.dirname(__FILE__)) unless defined? HKDF
+
+    p256dh = Base64.urlsafe_encode64 generate_ecdh_key
+    auth = Base64.urlsafe_encode64 Random.new.bytes(16)
+
+    payload = Encryption.encrypt("Hello World", p256dh, auth)
+
+    encrypted = payload.fetch(:ciphertext)
+
+    decrypted_data = ECE.decrypt(encrypted,
+      :key => payload.fetch(:shared_secret),
+      :salt => payload.fetch(:salt),
+      :server_public_key => payload.fetch(:server_public_key_bn),
+      :user_public_key => Base64.urlsafe_decode64(p256dh),
+      :auth => Base64.urlsafe_decode64(auth))
+
+    assert_equal "Hello World", decrypted_data
+  end
+
+  def generate_ecdh_key(group = 'prime256v1')
+    curve = OpenSSL::PKey::EC.new(group)
+    curve.generate_key
+    str = curve.public_key.to_bn.to_s(2)
+    puts "curve.public_key.to_bn.to_s(2): #{str.inspect}" if $VERBOSE
+    str
+  end
+  private :generate_ecdh_key
+
+  module Encryption # EC + (symmetric) AES GCM AAED encryption
+    extend self
+
+    def encrypt(message, p256dh, auth)
+
+      group_name = "prime256v1"
+      salt = Random.new.bytes(16)
+
+      server = OpenSSL::PKey::EC.new(group_name)
+      server.generate_key
+      server_public_key_bn = server.public_key.to_bn
+
+      group = OpenSSL::PKey::EC::Group.new(group_name)
+      client_public_key_bn = OpenSSL::BN.new(Base64.urlsafe_decode64(p256dh), 2)
+
+      #puts client_public_key_bn.to_s if $VERBOSE
+
+      client_public_key = OpenSSL::PKey::EC::Point.new(group, client_public_key_bn)
+
+      shared_secret = server.dh_compute_key(client_public_key)
+
+      client_auth_token = Base64.urlsafe_decode64(auth)
+
+      prk = HKDF.new(shared_secret, :salt => client_auth_token, :algorithm => 'SHA256', :info => "Content-Encoding: auth\0").next_bytes(32)
+
+      context = create_context(client_public_key_bn, server_public_key_bn)
+
+      content_encryption_key_info = create_info('aesgcm', context)
+      content_encryption_key = HKDF.new(prk, :salt => salt, :info => content_encryption_key_info).next_bytes(16)
+
+      nonce_info = create_info('nonce', context)
+      nonce = HKDF.new(prk, :salt => salt, :info => nonce_info).next_bytes(12)
+
+      ciphertext = encrypt_payload(message, content_encryption_key, nonce)
+
+      {
+        :ciphertext => ciphertext, :salt => salt, :shared_secret => shared_secret,
+        :server_public_key_bn => convert16bit(server_public_key_bn)
+      }
+    end
+
+    private
+
+    def create_context(client_public_key, server_public_key)
+      c = convert16bit(client_public_key)
+      s = convert16bit(server_public_key)
+      context = "\0"
+      context += [c.bytesize].pack("n*")
+      context += c
+      context += [s.bytesize].pack("n*")
+      context += s
+      context
+    end
+
+    def encrypt_payload(plaintext, content_encryption_key, nonce)
+      cipher = OpenSSL::Cipher.new('aes-128-gcm')
+      cipher.encrypt
+      cipher.key = content_encryption_key
+      cipher.iv = nonce
+      padding = cipher.update("\0\0")
+      text = cipher.update(plaintext)
+
+      e_text = padding + text + cipher.final
+      e_tag = cipher.auth_tag
+
+      e_text + e_tag
+    end
+
+    def create_info(type, context)
+      info = "Content-Encoding: "
+      info += type; info += "\0"; info += "P-256"; info += context
+      info
+    end
+
+    def convert16bit(key)
+      [key.to_s(16)].pack("H*")
+    end
+
+  end
+
   def setup
     super
     self.class.disable_security_restrictions!
