diff --git a/src/Simplex/Messaging/Agent/Store/SQLite.hs b/src/Simplex/Messaging/Agent/Store/SQLite.hs index 45c1f26ad..d5b8f8290 100644 --- a/src/Simplex/Messaging/Agent/Store/SQLite.hs +++ b/src/Simplex/Messaging/Agent/Store/SQLite.hs @@ -67,8 +67,8 @@ import Simplex.Messaging.Agent.Store.Migrations (DBMigrate (..), sharedMigrateSc import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations import Simplex.Messaging.Agent.Store.SQLite.Common import qualified Simplex.Messaging.Agent.Store.SQLite.DB as DB -import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfig (..), MigrationError (..)) import Simplex.Messaging.Agent.Store.SQLite.Util (SQLiteFunc, createStaticFunction, mkSQLiteFunc) +import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfig (..), MigrationError (..)) import qualified Simplex.Messaging.Crypto as C import Simplex.Messaging.Util (ifM, safeDecodeUtf8) import System.Directory (copyFile, createDirectoryIfMissing, doesFileExist) @@ -77,10 +77,10 @@ import System.FilePath (takeDirectory, takeFileName, ()) -- * SQLite Store implementation createDBStore :: DBOpts -> [Migration] -> MigrationConfig -> IO (Either MigrationError DBStore) -createDBStore opts@DBOpts {dbFilePath, dbKey, keepKey, track} migrations migrationConfig = do +createDBStore opts@DBOpts {dbFilePath} migrations migrationConfig = do let dbDir = takeDirectory dbFilePath createDirectoryIfMissing True dbDir - st <- connectSQLiteStore dbFilePath dbKey keepKey track + st <- connectSQLiteStore opts r <- migrateDBSchema st opts Nothing migrations migrationConfig `onException` closeDBStore st case r of Right () -> pure $ Right st @@ -99,23 +99,24 @@ migrateDBSchema st DBOpts {dbFilePath, vacuum} migrationsTable migrations Migrat dbm = DBMigrate {initialize, getCurrent, run, backup} in sharedMigrateSchema dbm (dbNew st) migrations confirm -connectSQLiteStore :: FilePath -> ScrubbedBytes -> Bool -> DB.TrackQueries -> IO DBStore -connectSQLiteStore dbFilePath key keepKey track = do +connectSQLiteStore :: DBOpts -> IO DBStore +connectSQLiteStore DBOpts {dbFilePath, dbFunctions, dbKey = key, keepKey, track} = do dbNew <- not <$> doesFileExist dbFilePath - dbConn <- dbBusyLoop (connectDB dbFilePath key track) + dbConn <- dbBusyLoop $ connectDB dbFilePath dbFunctions key track dbConnection <- newMVar dbConn dbKey <- newTVarIO $! storeKey key keepKey dbClosed <- newTVarIO False dbSem <- newTVarIO 0 - pure DBStore {dbFilePath, dbKey, dbSem, dbConnection, dbNew, dbClosed} + pure DBStore {dbFilePath, dbFunctions, dbKey, dbSem, dbConnection, dbNew, dbClosed} -connectDB :: FilePath -> ScrubbedBytes -> DB.TrackQueries -> IO DB.Connection -connectDB path key track = do +connectDB :: FilePath -> [SQLiteFuncDef] -> ScrubbedBytes -> DB.TrackQueries -> IO DB.Connection +connectDB path functions key track = do db <- DB.open path track prepare db `onException` DB.close db -- _printPragmas db path pure db where + functions' = SQLiteFuncDef "simplex_xor_md5_combine" 2 True sqliteXorMd5CombinePtr : functions prepare db = do let db' = SQL.connectionHandle $ DB.conn db unless (BA.null key) . SQLite3.exec db' $ "PRAGMA key = " <> keyString key <> ";" @@ -127,8 +128,9 @@ connectDB path key track = do PRAGMA secure_delete = ON; PRAGMA auto_vacuum = FULL; |] - createStaticFunction db' "simplex_xor_md5_combine" 2 True sqliteXorMd5CombinePtr - >>= either (throwIO . userError . show) pure + forM_ functions' $ \SQLiteFuncDef {funcName, argCount, deterministic, funcPtr} -> + createStaticFunction db' funcName argCount deterministic funcPtr + >>= either (throwIO . userError . show) pure foreign export ccall "simplex_xor_md5_combine" sqliteXorMd5Combine :: SQLiteFunc @@ -155,12 +157,12 @@ openSQLiteStore st@DBStore {dbClosed} key keepKey = ifM (readTVarIO dbClosed) (openSQLiteStore_ st key keepKey) (putStrLn "openSQLiteStore: already opened") openSQLiteStore_ :: DBStore -> ScrubbedBytes -> Bool -> IO () -openSQLiteStore_ DBStore {dbConnection, dbFilePath, dbKey, dbClosed} key keepKey = +openSQLiteStore_ DBStore {dbConnection, dbFilePath, dbFunctions, dbKey, dbClosed} key keepKey = bracketOnError (takeMVar dbConnection) (tryPutMVar dbConnection) $ \DB.Connection {slow, track} -> do - DB.Connection {conn} <- connectDB dbFilePath key track + DB.Connection {conn} <- connectDB dbFilePath dbFunctions key track atomically $ do writeTVar dbClosed False writeTVar dbKey $! storeKey key keepKey diff --git a/src/Simplex/Messaging/Agent/Store/SQLite/Common.hs b/src/Simplex/Messaging/Agent/Store/SQLite/Common.hs index af70c41f5..0634360a2 100644 --- a/src/Simplex/Messaging/Agent/Store/SQLite/Common.hs +++ b/src/Simplex/Messaging/Agent/Store/SQLite/Common.hs @@ -6,6 +6,7 @@ module Simplex.Messaging.Agent.Store.SQLite.Common ( DBStore (..), DBOpts (..), + SQLiteFuncDef (..), withConnection, withConnection', withTransaction, @@ -20,9 +21,13 @@ import Control.Concurrent (threadDelay) import Control.Concurrent.STM (retry) import Data.ByteArray (ScrubbedBytes) import qualified Data.ByteArray as BA +import Data.ByteString (ByteString) import Database.SQLite.Simple (SQLError) import qualified Database.SQLite.Simple as SQL +import Database.SQLite3.Bindings +import Foreign.Ptr import qualified Simplex.Messaging.Agent.Store.SQLite.DB as DB +import Simplex.Messaging.Agent.Store.SQLite.Util import Simplex.Messaging.Util (ifM, unlessM) import qualified UnliftIO.Exception as E import UnliftIO.MVar @@ -33,6 +38,7 @@ storeKey key keepKey = if keepKey || BA.null key then Just key else Nothing data DBStore = DBStore { dbFilePath :: FilePath, + dbFunctions :: [SQLiteFuncDef], dbKey :: TVar (Maybe ScrubbedBytes), dbSem :: TVar Int, dbConnection :: MVar DB.Connection, @@ -42,12 +48,21 @@ data DBStore = DBStore data DBOpts = DBOpts { dbFilePath :: FilePath, + dbFunctions :: [SQLiteFuncDef], dbKey :: ScrubbedBytes, keepKey :: Bool, vacuum :: Bool, track :: DB.TrackQueries } +-- e.g. `SQLiteFuncDef "name" 2 True f` +data SQLiteFuncDef = SQLiteFuncDef + { funcName :: ByteString, + argCount :: CArgCount, + deterministic :: Bool, + funcPtr :: FunPtr SQLiteFunc + } + withConnectionPriority :: DBStore -> Bool -> (DB.Connection -> IO a) -> IO a withConnectionPriority DBStore {dbSem, dbConnection} priority action | priority = E.bracket_ signal release $ withMVar dbConnection action diff --git a/tests/AgentTests/FunctionalAPITests.hs b/tests/AgentTests/FunctionalAPITests.hs index cb74bc0b6..2a62deb45 100644 --- a/tests/AgentTests/FunctionalAPITests.hs +++ b/tests/AgentTests/FunctionalAPITests.hs @@ -3715,7 +3715,7 @@ insertUser :: DBStore -> IO () insertUser st = withTransaction st (`DB.execute_` "INSERT INTO users DEFAULT VALUES") #else createStore :: String -> IO (Either MigrationError DBStore) -createStore dbPath = createAgentStore (DBOpts dbPath "" False True DB.TQOff) (MigrationConfig MCError Nothing) +createStore dbPath = createAgentStore (DBOpts dbPath [] "" False True DB.TQOff) (MigrationConfig MCError Nothing) insertUser :: DBStore -> IO () insertUser st = withTransaction st (`DB.execute_` "INSERT INTO users (user_id) VALUES (1)") diff --git a/tests/AgentTests/MigrationTests.hs b/tests/AgentTests/MigrationTests.hs index 8245cfd51..321e112d7 100644 --- a/tests/AgentTests/MigrationTests.hs +++ b/tests/AgentTests/MigrationTests.hs @@ -230,6 +230,7 @@ createStore randSuffix migrations confirmMigrations = do let dbOpts = DBOpts { dbFilePath = testDB randSuffix, + dbFunctions = [], dbKey = "", keepKey = False, vacuum = True, diff --git a/tests/AgentTests/SQLiteTests.hs b/tests/AgentTests/SQLiteTests.hs index f66dfe5df..1cebc716b 100644 --- a/tests/AgentTests/SQLiteTests.hs +++ b/tests/AgentTests/SQLiteTests.hs @@ -72,7 +72,7 @@ withStore2 = before connect2 . after (removeStore . fst) connect2 :: IO (DBStore, DBStore) connect2 = do s1@DBStore {dbFilePath} <- createStore' - s2 <- connectSQLiteStore dbFilePath "" False DB.TQOff + s2 <- connectSQLiteStore $ DBOpts dbFilePath [] "" False False DB.TQOff pure (s1, s2) createStore' :: IO DBStore @@ -83,7 +83,7 @@ createEncryptedStore key keepKey = do -- Randomize DB file name to avoid SQLite IO errors supposedly caused by asynchronous -- IO operations on multiple similarly named files; error seems to be environment specific r <- randomIO :: IO Word32 - Right st <- createDBStore (DBOpts (testDB <> show r) key keepKey True DB.TQOff) appMigrations (MigrationConfig MCError Nothing) + Right st <- createDBStore (DBOpts (testDB <> show r) [] key keepKey True DB.TQOff) appMigrations (MigrationConfig MCError Nothing) withTransaction' st (`SQL.execute_` "INSERT INTO users (user_id) VALUES (1);") pure st diff --git a/tests/AgentTests/SchemaDump.hs b/tests/AgentTests/SchemaDump.hs index 1f83973e6..51aa59528 100644 --- a/tests/AgentTests/SchemaDump.hs +++ b/tests/AgentTests/SchemaDump.hs @@ -51,7 +51,7 @@ testVerifySchemaDump :: IO () testVerifySchemaDump = do savedSchema <- ifM (doesFileExist appSchema) (readFile appSchema) (pure "") savedSchema `deepseq` pure () - void $ createDBStore (DBOpts testDB "" False True TQOff) appMigrations (MigrationConfig MCConsole Nothing) + void $ createDBStore (DBOpts testDB [] "" False True TQOff) appMigrations (MigrationConfig MCConsole Nothing) getSchema testDB appSchema `shouldReturn` savedSchema removeFile testDB @@ -59,14 +59,14 @@ testVerifyLintFKeyIndexes :: IO () testVerifyLintFKeyIndexes = do savedLint <- ifM (doesFileExist appLint) (readFile appLint) (pure "") savedLint `deepseq` pure () - void $ createDBStore (DBOpts testDB "" False True TQOff) appMigrations (MigrationConfig MCConsole Nothing) + void $ createDBStore (DBOpts testDB [] "" False True TQOff) appMigrations (MigrationConfig MCConsole Nothing) getLintFKeyIndexes testDB "tests/tmp/agent_lint.sql" `shouldReturn` savedLint removeFile testDB testSchemaMigrations :: IO () testSchemaMigrations = do let noDownMigrations = dropWhileEnd (\Migration {down} -> isJust down) appMigrations - Right st <- createDBStore (DBOpts testDB "" False True TQOff) noDownMigrations (MigrationConfig MCError Nothing) + Right st <- createDBStore (DBOpts testDB [] "" False True TQOff) noDownMigrations (MigrationConfig MCError Nothing) mapM_ (testDownMigration st) $ drop (length noDownMigrations) appMigrations closeDBStore st removeFile testDB @@ -89,7 +89,7 @@ testSchemaMigrations = do testUsersMigrationNew :: IO () testUsersMigrationNew = do - Right st <- createDBStore (DBOpts testDB "" False True TQOff) appMigrations (MigrationConfig MCError Nothing) + Right st <- createDBStore (DBOpts testDB [] "" False True TQOff) appMigrations (MigrationConfig MCError Nothing) withTransaction' st (`SQL.query_` "SELECT user_id FROM users;") `shouldReturn` ([] :: [Only Int]) closeDBStore st @@ -97,11 +97,11 @@ testUsersMigrationNew = do testUsersMigrationOld :: IO () testUsersMigrationOld = do let beforeUsers = takeWhile (("m20230110_users" /=) . name) appMigrations - Right st <- createDBStore (DBOpts testDB "" False True TQOff) beforeUsers (MigrationConfig MCError Nothing) + Right st <- createDBStore (DBOpts testDB [] "" False True TQOff) beforeUsers (MigrationConfig MCError Nothing) withTransaction' st (`SQL.query_` "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'users';") `shouldReturn` ([] :: [Only String]) closeDBStore st - Right st' <- createDBStore (DBOpts testDB "" False True TQOff) appMigrations (MigrationConfig MCYesUp Nothing) + Right st' <- createDBStore (DBOpts testDB [] "" False True TQOff) appMigrations (MigrationConfig MCYesUp Nothing) withTransaction' st' (`SQL.query_` "SELECT user_id FROM users;") `shouldReturn` ([Only (1 :: Int)]) closeDBStore st'