Diffie-Hellman key exchange using Botan

dh_botan.cpp

/**
 * Diffie-Hellman key exchange using botan
 *
 * This code will fork two processes: one being the server (A) and other being the client (B)
 *
 * They will only share the exponent, any other value is independent and must be passed through
 * the socket. The goal is for both to get the same key
 *
 * Compiled using
 * $ g++ dh_botan.cpp --std=c++20 -Wall -Wextra -I/usr/include/botan-2 -lbotan-2 -o dh_botan
 */

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <memory>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

#include <botan/auto_rng.h>
#include <botan/dh.h>
#include <botan/pubkey.h>

#include <botan/numthry.h> // for random_prime()

#include <queue>
#include <vector>

const Botan::BigInt g("65537");

/**
 * Diffie-Hellman key creation server test using Botan
 */

int create_socket_server() {
  auto listfd = socket(AF_UNIX, SOCK_STREAM, 0);
  if (listfd == -1)
    exit(-1);

  struct sockaddr_un saddr {
    AF_UNIX, "sserver"
  };
  auto saddrlen = sizeof(struct sockaddr_un);
  auto r = bind(listfd, (struct sockaddr *)&saddr, saddrlen);
  if (r < 0) {
    printf("bind() error: %s\n", strerror(errno));
    exit(1);
  }

  listen(listfd, 10);

  fflush(stdout);
  printf("Running...\n");

  return listfd;
}

int create_socket_client() {
  int sock;
  int conn;

  struct sockaddr_un saddr {
    AF_UNIX, "sserver"
  };
  auto saddrlen = sizeof(struct sockaddr_un);

  sock = socket(AF_UNIX, SOCK_STREAM, 0);

  conn = connect(sock, (struct sockaddr *)&saddr, saddrlen);

  return sock;
}

void send_number(int socket, Botan::BigInt &num) {
  uint8_t numdata[num.bytes()];
  uint32_t size = num.bytes();
  num.binary_encode(numdata, size);
  write(socket, &size, sizeof(size));
  write(socket, numdata, size);
}

Botan::BigInt receive_number(int socket) {
  uint32_t size = 0;
  read(socket, &size, sizeof(size));
  if (size > 512) {
    printf("size too big %d\n", size);
    close(socket);
    exit(1);
  }

  uint8_t numdata[size];
  auto readsize = read(socket, numdata, size);
  if (readsize < 0) {
    printf("server error: %s\n", strerror(errno));
    close(socket);
    exit(1);
  }
  if (readsize != size) {
    printf("sizes differ (%zu != %d)\n", readsize, size);
    close(socket);
    exit(1);
  }
  Botan::BigInt num(numdata, size);
  return num;
}

int run_server() {
  printf("Running server\n");

  int sock = create_socket_server();
  int cli = accept(sock, nullptr, nullptr);

  auto rng = Botan::AutoSeeded_RNG();
  Botan::BigInt p = Botan::random_prime(rng, 256);

  send_number(cli, p);
  
  const Botan::DL_Group grp(p, g);
  printf("\tg=%s, p=%s\n", g.to_dec_string().c_str(),
         p.to_dec_string().c_str());

  const Botan::BigInt a{rng.random_vec(128/8)};

  printf("\ta=%s\n", a.to_dec_string().c_str());

  auto privkey = std::make_unique<Botan::DH_PrivateKey>(rng, grp, a);

  auto A = privkey->get_y();
  printf("\tpublic part: (g^a mod p)=%s\n", A.to_dec_string().c_str());

  send_number(cli, A);
  auto B = receive_number(cli);
  
  printf("\tvalue from client part: (B = g^b mod p)=%s\n",
         B.to_dec_string().c_str());

  Botan::DH_PublicKey key(grp, B);
  auto public_value = key.public_value();

  auto kas = std::make_unique<Botan::PK_Key_Agreement>(*privkey, rng, "Raw");
  auto derived_key = kas->derive_key(16, public_value);

  printf("\t final key at the server: %s\n", derived_key.to_string().c_str());

  close(cli);
  close(sock);

  return 0;
}

int run_client() {
  printf("Running client\n");

  int cli = create_socket_client();

  Botan::BigInt p = receive_number(cli);
  const Botan::DL_Group grp(p, g);
  printf("\tg=%s, p=%s\n", g.to_dec_string().c_str(),
         p.to_dec_string().c_str());

  auto rng = Botan::AutoSeeded_RNG();
  const Botan::BigInt b{rng.random_vec(128/8)};

  printf("\tb=%s\n", b.to_dec_string().c_str());

  auto privkey = std::make_unique<Botan::DH_PrivateKey>(rng, grp, b);

  auto B = privkey->get_y();
  printf("\tpublic part: (g^b mod p)=%s\n", B.to_dec_string().c_str());

  send_number(cli, B);
  auto A = receive_number(cli);

  printf("\tvalue from server part: (A = g^a mod p)=%s\n",
         A.to_dec_string().c_str());

  Botan::DH_PublicKey key(grp, A);
  auto public_value = key.public_value();

  auto kas = std::make_unique<Botan::PK_Key_Agreement>(*privkey, rng, "Raw");
  auto derived_key = kas->derive_key(16, public_value);

  printf("\t final key at the client: %s\n", derived_key.to_string().c_str());

  close(cli);

  return 0;
}

int main(int argc, char *argv[]) {
  auto pid = fork();
  if (pid == 0)
    return run_client();
  else if (pid > 0)
    return run_server();

  printf("fork() failed");
  return 1;
}