src/secure_layer.cpp - master - isachkov/crypto_server

src/secure_layer.cpp
214 строк5 KB
Ivan Sachkov
CryptoServer
15 май 2024, 20:26
15 май 2024, 20:269de0f1f
#include <crypto_chat/secure_layer.hpp>

#include <iostream>

SecureLayer::SecureLayer(SSLInstancePtr ssl) : ssl_{std::move(ssl)} {
}

bool SecureLayer::check_connection_state() {
  // In case we have ongoing handshakes, send some exchange messages
  if (SSL_is_init_finished(ssl_.get())) {
    return true;
  }

  perform_ssl_handshake();

  // Still not finished! Getting out with it.
  if (!SSL_is_init_finished(ssl_.get())) {
    return false;
  }

  return true;
}

int SecureLayer::receive_decoded(Buffer& buffer) {
  constexpr auto ssl_buffer_size = 4096;
  static std::array<std::uint8_t, ssl_buffer_size> ssl_buffer;

  // Read as much as SSL allows
  auto ssl_decrypt_size = 0;
  do {
    ssl_decrypt_size = SSL_read(ssl_.get(), ssl_buffer.data(), ssl_buffer.size());
    if (ssl_decrypt_size) {
      // Efficiency, yeah, you know already
      std::copy(ssl_buffer.begin(), ssl_buffer.begin() + ssl_decrypt_size, std::back_inserter(buffer));
    }
  } while (ssl_decrypt_size > 0);

  return ssl_decrypt_size;
}

int SecureLayer::receive_encoded(Buffer& buffer) {
  constexpr auto ssl_buffer_size = 4096;
  static std::array<std::uint8_t, ssl_buffer_size> ssl_buffer;

  auto encrypted_size = 0;
  do {
    encrypted_size = BIO_read(ssl_to_socket_, ssl_buffer.data(), ssl_buffer.size());
    if (encrypted_size > 0) {
      std::copy(ssl_buffer.begin(), ssl_buffer.begin() + encrypted_size, std::back_inserter(buffer));
    } else if (!BIO_should_retry(ssl_to_socket_)) {
      return -1;
    }
  } while (encrypted_size > 0);

  return encrypted_size;
}

void SecureLayer::receive(Buffer buffer) {
  if (!on_decrypted_) {
    std::cerr << "OnDecrypted is not set for the secure layer\n";
    return;
  }

  Buffer output;
  std::size_t pos = 0;
  while (pos != buffer.size()) {
    // Feed data into SSL object
    const auto processed = BIO_write(socket_to_ssl_, buffer.data() + pos, buffer.size() - pos);

    if (processed <= 0) {
      std::cerr << "Secure layer could not process input sequence!\n";
      return;
    }
    pos += processed;

    const auto is_stable = check_connection_state();
    if (!is_stable) {
      // Wait until ssl connection/state machine is stabilized
      return;
    }

    const auto decoded = receive_decoded(output);

    // If renegotiation is required
    if (is_ssl_io_requested(decoded)) {
      perform_ssl_exchange();
    }

    if (is_ssl_failure(decoded)) {
      std::cerr << "SSL failure, aborting\n";
      return;
    }
  }

  if (output.empty()) {
    return;
  }

  on_decrypted_(std::move(output));
}

void SecureLayer::send(Buffer buffer) {
  if (!on_encrypted_) {
    std::cerr << "OnEncrypted is not set for secure layer\n";
    return;
  }

  Buffer output;
  std::size_t pos = 0;
  while(pos != buffer.size()) {
    auto ssl_write_size = SSL_write(ssl_.get(), buffer.data() + pos, buffer.size() - pos);

    if(ssl_write_size > 0) {
      pos += ssl_write_size;
    }

    const auto encoded = receive_encoded(output);

    if(is_ssl_failure(ssl_write_size)) {
      std::cerr << "SSL failure: " << SSL_get_error(ssl_.get(), ssl_write_size) << "\n";
      return;
    }
  }

  on_encrypted_(std::move(output));
}

void SecureLayer::on_encrypted(OnEncrypted handler) {
  on_encrypted_ = std::move(handler);
}

void SecureLayer::on_decrypted(OnDecrypted handler) {
  on_decrypted_ = std::move(handler);
}

void SecureLayer::perform_ssl_handshake() {
  auto handshake = SSL_do_handshake(ssl_.get());

  if (!is_ssl_io_requested(handshake)) {
    return;
  }

  perform_ssl_exchange();
}

void SecureLayer::perform_ssl_exchange() {
  constexpr auto ssl_buffer_size = 4096;
  static std::array<std::uint8_t, ssl_buffer_size> ssl_buffer;

  Buffer response_data;

  // 1. Read response data from the SSL's bio
  auto ssl_read_size = 0;
  do {
    ssl_read_size = BIO_read(ssl_to_socket_, ssl_buffer.data(), ssl_buffer.size());
    if (ssl_read_size > 0) {
      std::copy(ssl_buffer.begin(), ssl_buffer.begin() + ssl_read_size, std::back_inserter(response_data));
    } else if (!BIO_should_retry(ssl_to_socket_)) {
      std::cerr << "No data has been received in SSL response and should_retry failed\n";
      return;
    }
  } while (ssl_read_size > 0);

  if(response_data.empty()) {
    return;
  }

  // 2. Send response to the peer
  if (on_encrypted_) {
    on_encrypted_(std::move(response_data));
  } else {
    std::cerr << "OnEncrypted is not set!\n";
  }
}

bool SecureLayer::init(bool is_server) {
  // Would be freed with it's associated SSL instance
  ssl_to_socket_ = BIO_new(BIO_s_mem());
  socket_to_ssl_ = BIO_new(BIO_s_mem());
  
  if(is_server) {
    SSL_set_accept_state(ssl_.get());
  } else {
    SSL_set_connect_state(ssl_.get());
  }

  SSL_set_bio(ssl_.get(), socket_to_ssl_, ssl_to_socket_);

  if (!is_server) {
    // Initiate it immideately, don't wait for first send
    perform_ssl_handshake();
  }

  return true;
}

bool SecureLayer::is_ssl_io_requested(int code) {
  auto status = SSL_get_error(ssl_.get(), code);
  // Unfortunatelly OpenSSL is very vague regards it's state and sometimes 
  // it requests for "READ" when it actually wants to send a token
  return (status == SSL_ERROR_WANT_WRITE) || (status == SSL_ERROR_WANT_READ);
}

bool SecureLayer::is_ssl_failure(int code) {
  auto status = SSL_get_error(ssl_.get(), code);
  switch (status) {
    case SSL_ERROR_NONE:
    case SSL_ERROR_WANT_WRITE:
    case SSL_ERROR_WANT_READ:
      return false;
  }

  return true;
}