From 4b97a95470df03eca392c6e675324a86e29574c1 Mon Sep 17 00:00:00 2001
From: Akash Hiremath <akashh246@gmail.com>
Date: Thu, 27 Oct 2022 20:47:55 +0530
Subject: [PATCH] Fix segfault when database connections timeout (#218)

---
 .github/workflows/ci.yml               |  2 +-
 c_src/sqlite3_nif.c                    | 25 ++++++++++++++++
 test/exqlite/timeout_segfault_test.exs | 41 ++++++++++++++++++++++++++
 test/test_helper.exs                   |  2 +-
 4 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 test/exqlite/timeout_segfault_test.exs

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index e9cfe32c..5a5701e8 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -49,7 +49,7 @@ jobs:
       matrix:
         os: ["ubuntu-20.04", "windows-2019"]
         elixir: ["1.14", "1.13", "1.12", "1.11"]
-        otp: ["25", 24", "23"]
+        otp: ["25", "24", "23"]
         exclude:
           - otp: "25"
             os: "windows-2019"
diff --git a/c_src/sqlite3_nif.c b/c_src/sqlite3_nif.c
index eb9f1f80..174c4f79 100644
--- a/c_src/sqlite3_nif.c
+++ b/c_src/sqlite3_nif.c
@@ -22,6 +22,7 @@ static sqlite3_mem_methods default_alloc_methods = {0};
 typedef struct connection
 {
     sqlite3* db;
+    ErlNifMutex* mutex;
 } connection_t;
 
 typedef struct statement
@@ -219,6 +220,12 @@ exqlite_open(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
         return make_error_tuple(env, "database_open_failed");
     }
 
+    conn->mutex = enif_mutex_create("exqlite:connection");
+    if (conn->mutex == NULL) {
+        enif_release_resource(conn);
+        return make_error_tuple(env, "failed_to_create_mutex");
+    }
+
     sqlite3_busy_timeout(conn->db, 2000);
 
     result = enif_make_resource(env, conn);
@@ -256,6 +263,11 @@ exqlite_close(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
         }
     }
 
+    // close connection in critical section to avoid race-condition
+    // cases. Cases such as query timeout and connection pooling
+    // attempting to close the connection
+    enif_mutex_lock(conn->mutex);
+
     // note: _v2 may not fully close the connection, hence why we check if
     // any transaction is open above, to make sure other connections aren't blocked.
     // v1 is guaranteed to close or error, but will return error if any
@@ -263,10 +275,12 @@ exqlite_close(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     // to later run to clean those up
     rc = sqlite3_close_v2(conn->db);
     if (rc != SQLITE_OK) {
+        enif_mutex_unlock(conn->mutex);
         return make_sqlite3_error_tuple(env, rc, conn->db);
     }
 
     conn->db = NULL;
+    enif_mutex_unlock(conn->mutex);
 
     return make_atom(env, "ok");
 }
@@ -357,11 +371,21 @@ exqlite_prepare(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     if (!statement) {
         return make_error_tuple(env, "out_of_memory");
     }
+    statement->statement = NULL;
 
     enif_keep_resource(conn);
     statement->conn = conn;
 
+    // ensure connection is not getting closed by parallel thread
+    enif_mutex_lock(conn->mutex);
+    if (conn->db == NULL) {
+        enif_mutex_unlock(conn->mutex);
+        enif_release_resource(statement);
+        return make_error_tuple(env, "connection closed");
+    }
     rc = sqlite3_prepare_v3(conn->db, (char*)bin.data, bin.size, 0, &statement->statement, NULL);
+    enif_mutex_unlock(conn->mutex);
+
     if (rc != SQLITE_OK) {
         enif_release_resource(statement);
         return make_sqlite3_error_tuple(env, rc, conn->db);
@@ -857,6 +881,7 @@ connection_type_destructor(ErlNifEnv* env, void* arg)
     if (conn->db) {
         sqlite3_close_v2(conn->db);
         conn->db = NULL;
+        enif_mutex_destroy(conn->mutex);
     }
 }
 
diff --git a/test/exqlite/timeout_segfault_test.exs b/test/exqlite/timeout_segfault_test.exs
new file mode 100644
index 00000000..35776d50
--- /dev/null
+++ b/test/exqlite/timeout_segfault_test.exs
@@ -0,0 +1,41 @@
+defmodule Exqlite.TimeoutSegfaultTest do
+  use ExUnit.Case
+
+  @moduletag :slow_test
+
+  setup do
+    {:ok, path} = Temp.path()
+    on_exit(fn -> File.rm(path) end)
+
+    %{path: path}
+  end
+
+  test "segfault", %{path: path} do
+    {:ok, conn} =
+      DBConnection.start_link(Exqlite.Connection,
+        busy_timeout: 50_000,
+        pool_size: 50,
+        timeout: 1,
+        database: path,
+        journal_mode: :wal
+      )
+
+    query = %Exqlite.Query{statement: "create table foo(id integer, val integer)"}
+    {:ok, _, _} = DBConnection.execute(conn, query, [])
+
+    values = for i <- 1..1000, do: "(#{i}, #{i})"
+    statement = "insert into foo(id, val) values #{Enum.join(values, ",")}"
+    insert_query = %Exqlite.Query{statement: statement}
+
+    1..5000
+    |> Task.async_stream(fn _ ->
+      try do
+        DBConnection.execute(conn, insert_query, [], timeout: 1)
+      catch
+        kind, reason ->
+          IO.puts("Error: #{inspect(kind)} reason: #{inspect(reason)}")
+      end
+    end)
+    |> Stream.run()
+  end
+end
diff --git a/test/test_helper.exs b/test/test_helper.exs
index 6a0af57d..6682d142 100644
--- a/test/test_helper.exs
+++ b/test/test_helper.exs
@@ -1 +1 @@
-ExUnit.start(capture_log: true)
+ExUnit.start(capture_log: true, timeout: 120_000, exclude: [:slow_test])