create-upload-payload-for-load-testing

#!/usr/bin/env ruby
# frozen_string_literal: true

$LOAD_PATH.unshift(File.expand_path('../test/lib/protocol', __dir__))

require('bundler/setup')

require('covidshield_pb')
require('open3')
require('faraday')
require('rbnacl')
require('mysql2')

module CreateUploadPayloadForLoadTesting
  extend(self)

  def call(url, token)
    @conn = Faraday.new(url: url)
    @token = token
    sleep(0.1)

    req = encrypted_request(dummy_payload, new_valid_keyset)
    puts 'local body = "'+req.to_proto.each_byte.to_a.map { |f| format("\\%03d", f) }.join+'"'
  end

  def new_valid_one_time_code
    resp = @conn.post do |req|
      req.url('/new-key-claim')
      req.headers['Authorization'] = "Bearer #{@token}"
    end
    assert_response(resp, 200, 'text/plain; charset=utf-8')
    resp.body.chomp
  end

  def tek(key_data: '1' * 16, transmission_risk_level: 3, rolling_period: 144, rolling_start_interval_number: Time.now.to_i / 86400)
    Covidshield::TemporaryExposureKey.new(
      keyData: key_data,
      transmissionRiskLevel: transmission_risk_level,
      rollingPeriod: rolling_period,
      rollingStartIntervalNumber: rolling_start_interval_number
    )
  end

  def new_valid_keyset
    otc = new_valid_one_time_code

    app_private_key = RbNaCl::PrivateKey.generate
    app_public_key  = app_private_key.public_key

    kcq = Covidshield::KeyClaimRequest.new(
      oneTimeCode: otc,
      appPublicKey: app_public_key.to_s,
    )
    resp = @conn.post('/claim-key', kcq.to_proto)
    assert_response(resp, 200, 'application/x-protobuf')
    kcr = Covidshield::KeyClaimResponse.decode(resp.body)
    assert_equal(:NONE, kcr.error)
    assert_equal(32, kcr.serverPublicKey.each_byte.size)

    {
      app_public: app_public_key,
      app_private: app_private_key,
      server_public: kcr.serverPublicKey
    }
  end

  def assert_equal(exp, act)
    raise("assertion failed (wanted:#{exp.inspect}; got:#{act.inspect})") unless exp == act
  end

  def assert_response(resp, status, content_type, body: nil)
    assert_equal(status, resp.status)
    assert_equal(resp.headers['Content-Type'], content_type)
    case body
    when String
      assert_equal(body, resp.body)
    when Regexp
      assert_match(/\A[0-9]{8}\n\z/m, resp.body)
    end
  end

  def dummy_payload(nkeys=1, timestamp: Time.now)
    Covidshield::Upload.new(timestamp: timestamp, keys: [tek]*nkeys).to_proto
  end

  def encrypted_request(
    payload, keyset, server_public: keyset[:server_public], app_private: keyset[:app_private],
    app_public: keyset[:app_public], app_public_to_send: app_public,
    server_public_to_send: server_public,
    box: RbNaCl::Box.new(server_public, app_private),
    nonce: RbNaCl::Random.random_bytes(box.nonce_bytes),
    nonce_to_send: nonce,
    encrypted_payload: box.encrypt(nonce, payload)
  )
    Covidshield::EncryptedUploadRequest.new(
      serverPublicKey: server_public_to_send.to_s,
      appPublicKey: app_public_to_send.to_s,
      nonce: nonce_to_send,
      payload: encrypted_payload,
    )
  end
end

if $PROGRAM_NAME == __FILE__
  CreateUploadPayloadForLoadTesting.call("https://submission.covidshield.app", 'test')
end