Handshake: ECDH between Ruby on Rails and Web Crypto APIs

// Store our private key on IndexedDB
import {IndexedDBService} from './indexedDBHelpers';

// Setup constants for options
const EXTRACTABLE = true;
const CONNECTION_USAGE = ['encrypt'];
const KEY_USAGE = ['deriveKey', 'deriveBits'];
const ALGORITHM = {name: 'ECDH', namedCurve: 'P-256'};
const DERIVE_ALGORITHM = {name: 'AES-CTR', length: 256};
const ENCRYPT_ALGORITHM = {name: DERIVE_ALGORITHM.name, length: 128};

export const CriptographyService = {
  // Generate key pair
  generate: () => (
    window.crypto.subtle
      .generateKey(ALGORITHM, EXTRACTABLE, KEY_USAGE)
      .then(keys => (
        // Store private key and obtain keys
        IndexedDBService.set(1, keys.privateKey) && keys
      ))
  ),
  // Export key to base64
  export: (key) => (
    window.crypto.subtle
      .exportKey('spki', key)
      .then(buffer => (
        window.btoa(String.fromCharCode.apply(null, new Uint8Array(buffer)))
      ))
  ),
  // Import key from base64
  import: (key) => (
    window.crypto.subtle
      .importKey(
        'raw',
        Uint8Array.from(window.atob(key), (c) => c.charCodeAt(0)).buffer,
        ALGORITHM,
        EXTRACTABLE,
        KEY_USAGE,
      )
  ),
  // Encrypt data with key
  encrypt: (key, data) => {
    const counter = crypto.getRandomValues(new Uint8Array(16));
    return window.crypto.subtle
      .encrypt({...ENCRYPT_ALGORITHM, counter}, key, new TextEncoder().encode(data))
      .then(encryptedData => ({
        iv: window.btoa(String.fromCharCode.apply(null, counter)),
        encrypted: window.btoa(
          String.fromCharCode.apply(null, new Uint8Array(encryptedData))
        ),
      }));
  },
  // Derive key from server public key
  computeSecret: async (publicKey) => (
    window.crypto.subtle.deriveKey(
      {...ALGORITHM, public: publicKey},
      // Get private key
      await IndexedDBService.get(1),
      DERIVE_ALGORITHM,
      EXTRACTABLE,
      CONNECTION_USAGE,
    )
  ),
};

class HandshakeManagerService
  attr_accessor :initiated, :client_public_key

  EXPIRE_TIME = 1.minute
  CURVE_NAME = 'prime256v1'.freeze # P-256 on Web Crypto API
  DERIVE_NAME = 'aes-256-ctr'.freeze

  def initialize(params = {})
    @iv = params[:iv]
    @token = params[:token]
    @encrypted = params[:encrypted]
    @client_public_key = params[:public_key]
    @initiated = @client_public_key.present?
    @content = nil
    @shared_key = nil
  end

  def initiate
    # Import client public key
    client = OpenSSL::PKey::EC.new(Base64.strict_decode64(client_public_key))
    
    # Generate key pair
    server = OpenSSL::PKey::EC.new(CURVE_NAME)
    server.generate_key!

    # Convert generated key pair to base64
    server_public_key = Base64.strict_encode64(server.public_key.to_bn.to_s(2))

    # Derive shared key from client public key
    @shared_key = server.dh_compute_key(client.public_key)
    
    # Upon each request, unique to token shared key is stored
    Rails.cache.write(@token, Base64.strict_encode64(@shared_key), expires_in: EXPIRE_TIME)

    {token: @token, public_key: server_public_key}
  end

  def on_initialize(&block)
    return initiate if initiated

    if @token
      # Get private key particular to a given token
      @shared_key = Base64.strict_decode64(Rails.cache.read(@token))

      # Upon adding IV and Shared key, decipher encrypted data
      @iv = Base64.strict_decode64(@iv)
      decipher = OpenSSL::Cipher.new(DERIVE_NAME)
      decipher.decrypt
      decipher.key = @shared_key
      decipher.iv = @iv
      decrypted_content = decipher.update(Base64.strict_decode64(@encrypted)) + decipher.final
      
      # Parse decrypted json content
      @content = JSON.parse(decrypted_content, symbolize_names: true)
    end

    return block.call(@content) if @content

    block.call
  end
end

def handshake(&block)
  HandshakeManagerService.new(params).on_initialize do |parameters|
    block.call(parameters)
  end
end

// Function gets 2 flags: initiated, established and processes request
function* processHandshake(request, initiated, established) {
  if (initiated) {
    const {endpoint, method, actions} = request.payload[API_CALL];
    // Generate pair on initiated flag set
    const keyPair = yield call(() => CriptographyService.generate());
    // Export public key base64 
    const publicKey = yield call(() => CriptographyService.export(keyPair.publicKey));
    // Generate random token
    const token = window.crypto.randomUUID();
    // Add to initiated request and send to server
    const initiatedRequest = {
      actions,
      method,
      endpoint,
      action: request.payload,
      body: {token, public_key: publicKey},
    };
    yield fork(processRequest, {
      type: request.type,
      payload: {[API_CALL]: initiatedRequest},
    });
  }

  if (established) {
    const {endpoint, method, actions} = request.payload.action[API_CALL];
    const {payload: {response: {result: {token, public_key: publicKey}}, action: {[API_CALL]: {body}}}} = request;
    // Upon established flag set, import base64 public key
    const serverPublicKey = yield call(() => CriptographyService.import(publicKey));
    // Derive secret from server public key
    const sharedKey = yield call(() => CriptographyService.computeSecret(serverPublicKey));
    // Convert JSON to string and encrypt data
    const encrypted = yield call(() => CriptographyService.encrypt(sharedKey, JSON.stringify(body)));
    // Ending handshake on client side send encrypted data to server
    const establishedRequest = {
      ...request.payload.action,
      [API_CALL]: {
        actions,
        method,
        endpoint,
        body: {...encrypted, token},
      },
    };
    yield put(requestResponse(
      {type: actions.REQUEST, payload: establishedRequest})
    );
  }
}