diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,25 @@ +# Changelog + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [0.1.0.0] - 2026-04-25 + +Initial release. Pure-Haskell implementation of the PostgreSQL v3 wire +protocol. + +- Connection setup over TCP and Unix sockets, TLS 1.2/1.3, + cleartext / MD5 / SCRAM-SHA-256 authentication, multi-host failover + (`target_session_attrs`, `load_balance_hosts`), startup parameters. +- Connection pool with idle reaping, health checks, jittered max-life, + queue mode (LIFO/FIFO), recycling modes (always/on-error), lifecycle + hooks, and `nothunks`-checked invariants. +- Wire frames: extended query (parse / bind / describe / execute), + simple query, COPY in/out, large objects, LISTEN/NOTIFY, query + cancellation. +- Binary format type classes (`PgEncode`, `PgDecode`) and OID metadata. + Concrete codecs for PostgreSQL types live in the + [`valiant`](https://hackage.haskell.org/package/valiant) runtime. +- Nine structured error types with column-by-column diagnostics. + +[0.1.0.0]: https://github.com/joshburgess/valiant/releases/tag/v0.1.0.0 diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,27 @@ +Copyright (c) 2026 Josh Burgess + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from this + software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,51 @@ +# pg-wire + +Pure-Haskell driver for the PostgreSQL v3 wire protocol. No `libpq`, no +C dependencies. + +## What you get + +- Connection setup: TCP and Unix sockets, TLS 1.2/1.3, MD5 and + SCRAM-SHA-256 auth, multi-host failover (`target_session_attrs`, + `load_balance_hosts`), startup parameters. +- Connection pool: idle reaping, health checks, jittered max-life, + queue mode (LIFO/FIFO), recycling modes (always/on-error), lifecycle + hooks, `nothunks` invariants. +- Wire frames: extended query (parse / bind / describe / execute), + simple query, COPY in/out, large objects, LISTEN/NOTIFY, query + cancellation. +- Binary format type classes: `PgEncode`, `PgDecode`, OID metadata. + Codecs for individual types live in the `valiant` runtime library. +- Errors: 8 structured error constructors (`ConnectionError`, + `AuthError`, `ProtocolError`, `QueryError`, `DecodeError`, + `PoolTimeout`, `PoolClosed`, `ConnectionDead`). + +## When to use this directly + +`pg-wire` is the transport layer. Most users want the higher-level +[`valiant`](https://hackage.haskell.org/package/valiant) runtime, which +re-exports everything from `pg-wire` and adds compile-time SQL safety +via the `valiant-plugin` GHC source plugin. + +Use `pg-wire` directly only when you want raw wire-protocol access +without the compile-time validation layer (e.g. for building a +different ORM or for testing the protocol itself). + +## Installation + +``` +cabal install pg-wire +``` + +## Documentation + +- [Top-level README](https://github.com/joshburgess/valiant#readme): + the project overview and how the four packages fit together. +- [`docs/`](https://github.com/joshburgess/valiant/tree/main/docs): + tutorial, performance notes, async architecture, gap analysis. +- [`docs/benchmark-results/`](https://github.com/joshburgess/valiant/tree/main/docs/benchmark-results): + archived comparative numbers. + +## License + +BSD-3-Clause. See `LICENSE`. diff --git a/mockserver/PgWire/MockServer.hs b/mockserver/PgWire/MockServer.hs new file mode 100644 --- /dev/null +++ b/mockserver/PgWire/MockServer.hs @@ -0,0 +1,320 @@ +-- | A minimal mock PostgreSQL server for testing. +-- +-- Speaks enough of the PG v3 wire protocol to test client behavior +-- without a real database: authentication, simple queries, extended +-- query protocol, error responses, and connection lifecycle. +-- +-- @ +-- withMockServer defaultMockConfig $ \\port -> do +-- conn <- connectString (\"postgres:\/\/user:pass\@localhost:\" <> BS8.pack (show port) <> \"\/testdb\") +-- (rows, _) <- simpleQuery conn \"SELECT 1\" +-- close conn +-- @ +module PgWire.MockServer + ( -- * Server lifecycle + withMockServer + , MockConfig (..) + , defaultMockConfig + -- * Query handlers + , QueryHandler + , simpleHandler + , errorHandler + , failNTimes + -- * Response builders + , sendBackendMsg + , buildBackendMsg + , buildErrorResponse + , buildCommandComplete + , buildRowDescription + , buildDataRow + , buildReadyForQuery + ) where + +import Control.Concurrent (forkIO, killThread) +import Control.Exception (SomeException, catch, finally) +import Data.IORef +import Data.Bits (shiftL, (.|.)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Builder (Builder) +import Data.ByteString.Builder qualified as B +import Data.ByteString.Char8 qualified as BS8 +import Data.ByteString.Lazy qualified as LBS +import Data.Int (Int32) +import Data.Word (Word8) +import Network.Socket qualified as NS +import Network.Socket.ByteString qualified as NSB + +-- | Configuration for the mock server. +data MockConfig = MockConfig + { mockUser :: ByteString + -- ^ Expected username. Default: @\"testuser\"@. + , mockPassword :: ByteString + -- ^ Expected password. Default: @\"testpass\"@. + , mockDatabase :: ByteString + -- ^ Expected database name. Default: @\"testdb\"@. + , mockQueryHandler :: QueryHandler + -- ^ Handler for simple queries. Default: returns empty result. + } + +-- | Handler called for each simple query. Receives the SQL text and +-- a function to send backend messages. +type QueryHandler = ByteString -> (Builder -> IO ()) -> IO () + +-- | Default config: accepts auth, returns empty results for all queries. +defaultMockConfig :: MockConfig +defaultMockConfig = MockConfig + { mockUser = "testuser" + , mockPassword = "testpass" + , mockDatabase = "testdb" + , mockQueryHandler = \_ send -> do + send $ buildCommandComplete "SELECT 0" + } + +-- | A simple query handler that returns fixed rows for a specific query. +simpleHandler :: [(ByteString, [[(ByteString, ByteString)]])] -> QueryHandler +simpleHandler table sql send = case lookup sql table of + Nothing -> send $ buildCommandComplete "SELECT 0" + Just rows -> do + case rows of + [] -> send $ buildCommandComplete "SELECT 0" + (first : _) -> do + -- Send RowDescription + let cols = map fst first + send $ buildRowDescription cols + -- Send DataRows + mapM_ (\row -> send $ buildDataRow (map snd row)) rows + send $ buildCommandComplete ("SELECT " <> BS8.pack (show (length rows))) + +-- | A handler that returns a PG error with the given SQLSTATE and message. +-- +-- @ +-- errorHandler \"23505\" \"duplicate key value violates unique constraint \\\"users_email_key\\\"\" +-- @ +errorHandler :: ByteString -> ByteString -> QueryHandler +errorHandler sqlstate msg _ send = + send $ buildErrorResponse sqlstate msg + +-- | A handler backed by an IORef counter. Returns an error for the first +-- @n@ calls, then delegates to the fallback handler. Useful for testing +-- retry logic. +-- +-- @ +-- ref <- newIORef 0 +-- let handler = failNTimes ref 2 \"40001\" \"serialization failure\" defaultHandler +-- @ +failNTimes :: IORef Int -> Int -> ByteString -> ByteString -> QueryHandler -> QueryHandler +failNTimes ref n sqlstate msg fallback sql send = do + count <- atomicModifyIORef' ref (\c -> (c + 1, c)) + if count < n + then send $ buildErrorResponse sqlstate msg + else fallback sql send + +-- | Start a mock server on a random port, run an action with the port, +-- then shut down. +withMockServer :: MockConfig -> (NS.PortNumber -> IO a) -> IO a +withMockServer cfg action = do + -- Create a listening socket + sock <- NS.socket NS.AF_INET NS.Stream NS.defaultProtocol + NS.setSocketOption sock NS.ReuseAddr 1 + NS.bind sock (NS.SockAddrInet 0 (NS.tupleToHostAddress (127, 0, 0, 1))) + NS.listen sock 5 + port <- NS.socketPort sock + + -- Start acceptor thread + tid <- forkIO $ acceptLoop sock cfg + + -- Run the action, then clean up + action port `finally` do + killThread tid + NS.close sock + +-- | Accept loop: handle one connection at a time. +acceptLoop :: NS.Socket -> MockConfig -> IO () +acceptLoop listenSock cfg = go + where + go = do + (clientSock, _) <- NS.accept listenSock + _ <- forkIO (handleClient clientSock cfg `catch` \(_ :: SomeException) -> pure ()) + go + +-- | Handle a single client connection. +handleClient :: NS.Socket -> MockConfig -> IO () +handleClient sock cfg = do + -- Read startup message (no tag byte) + startupBytes <- recvN sock 4 + let startupLen = decodeInt32BE startupBytes 0 + _payload <- recvN sock (fromIntegral startupLen - 4) + + -- Send AuthOk (cleartext auth skipped for simplicity) + sendBuilder sock $ buildAuth 0 -- AuthOk + + -- Send server parameters + sendBuilder sock $ buildParameterStatus "server_version" "16.0" + sendBuilder sock $ buildParameterStatus "server_encoding" "UTF8" + sendBuilder sock $ buildParameterStatus "client_encoding" "UTF8" + sendBuilder sock $ buildParameterStatus "is_superuser" "off" + + -- Send BackendKeyData + sendBuilder sock $ buildBackendKeyData 1234 5678 + + -- Send ReadyForQuery (Idle) + sendBuilder sock $ buildReadyForQuery 'I' + + -- Message loop + let loop = do + mTag <- recvMaybe sock 1 + case mTag of + Nothing -> pure () -- client disconnected + Just tagBs -> do + let tag = BS.index tagBs 0 + lenBs <- recvN sock 4 + let len = decodeInt32BE lenBs 0 + payload <- if len > 4 then recvN sock (fromIntegral len - 4) else pure BS.empty + handleMessage sock cfg tag payload + loop + + loop `catch` \(_ :: SomeException) -> pure () + NS.close sock + +handleMessage :: NS.Socket -> MockConfig -> Word8 -> ByteString -> IO () +handleMessage sock cfg tag payload = case tag of + -- 'Q' = Simple Query + 0x51 -> do + let sql = BS.takeWhile (/= 0) payload -- strip NUL terminator + mockQueryHandler cfg sql (sendBuilder sock) + sendBuilder sock $ buildReadyForQuery 'I' + + -- 'P' = Parse + 0x50 -> do + sendBuilder sock $ buildSimpleTag '1' -- ParseComplete + + -- 'B' = Bind + 0x42 -> do + sendBuilder sock $ buildSimpleTag '2' -- BindComplete + + -- 'D' = Describe + 0x44 -> do + sendBuilder sock $ buildNoData + + -- 'E' = Execute + 0x45 -> do + sendBuilder sock $ buildCommandComplete "SELECT 0" + + -- 'S' = Sync + 0x53 -> do + sendBuilder sock $ buildReadyForQuery 'I' + + -- 'H' = Flush + 0x48 -> pure () + + -- 'C' = Close + 0x43 -> do + sendBuilder sock $ buildSimpleTag '3' -- CloseComplete + + -- 'X' = Terminate + 0x58 -> pure () -- connection will close + + -- Unknown + _ -> do + sendBuilder sock $ buildErrorResponse "08P01" ("Unknown message tag: " <> BS8.pack (show tag)) + sendBuilder sock $ buildReadyForQuery 'E' + +------------------------------------------------------------------------ +-- Backend message builders +------------------------------------------------------------------------ + +-- | Build a complete backend message as a Builder. +buildBackendMsg :: Word8 -> Builder -> Builder +buildBackendMsg tag payload = + let payloadBs = LBS.toStrict (B.toLazyByteString payload) + len = fromIntegral (BS.length payloadBs + 4) :: Int32 + in B.word8 tag <> B.int32BE len <> B.byteString payloadBs + +-- | Send a backend message builder to a socket. +sendBackendMsg :: NS.Socket -> Word8 -> Builder -> IO () +sendBackendMsg sock tag payload = sendBuilder sock (buildBackendMsg tag payload) + +buildAuth :: Int32 -> Builder +buildAuth authType = buildBackendMsg 0x52 (B.int32BE authType) + +buildParameterStatus :: ByteString -> ByteString -> Builder +buildParameterStatus name value = + buildBackendMsg 0x53 (B.byteString name <> B.word8 0 <> B.byteString value <> B.word8 0) + +buildBackendKeyData :: Int32 -> Int32 -> Builder +buildBackendKeyData pid key = + buildBackendMsg 0x4B (B.int32BE pid <> B.int32BE key) + +buildReadyForQuery :: Char -> Builder +buildReadyForQuery status = + buildBackendMsg 0x5A (B.word8 (fromIntegral (fromEnum status))) + +buildSimpleTag :: Char -> Builder +buildSimpleTag c = buildBackendMsg (fromIntegral (fromEnum c)) mempty + +buildNoData :: Builder +buildNoData = buildSimpleTag 'n' + +buildCommandComplete :: ByteString -> Builder +buildCommandComplete tag = + buildBackendMsg 0x43 (B.byteString tag <> B.word8 0) + +buildRowDescription :: [ByteString] -> Builder +buildRowDescription cols = + buildBackendMsg 0x54 $ do + B.int16BE (fromIntegral (length cols)) + <> foldMap buildFieldInfo cols + where + buildFieldInfo name = + B.byteString name <> B.word8 0 -- name + NUL + <> B.word32BE 0 -- table OID + <> B.int16BE 0 -- column number + <> B.word32BE 25 -- type OID (text) + <> B.int16BE (-1) -- type size + <> B.int32BE (-1) -- type modifier + <> B.int16BE 0 -- format (text) + +buildDataRow :: [ByteString] -> Builder +buildDataRow cols = + buildBackendMsg 0x44 $ + B.int16BE (fromIntegral (length cols)) + <> foldMap (\v -> B.int32BE (fromIntegral (BS.length v)) <> B.byteString v) cols + +buildErrorResponse :: ByteString -> ByteString -> Builder +buildErrorResponse sqlState msg = + buildBackendMsg 0x45 $ + B.word8 (fromIntegral (fromEnum 'S')) <> B.byteString "ERROR" <> B.word8 0 + <> B.word8 (fromIntegral (fromEnum 'C')) <> B.byteString sqlState <> B.word8 0 + <> B.word8 (fromIntegral (fromEnum 'M')) <> B.byteString msg <> B.word8 0 + <> B.word8 0 -- terminator + +------------------------------------------------------------------------ +-- Socket helpers +------------------------------------------------------------------------ + +sendBuilder :: NS.Socket -> Builder -> IO () +sendBuilder sock b = NSB.sendAll sock (LBS.toStrict (B.toLazyByteString b)) + +recvN :: NS.Socket -> Int -> IO ByteString +recvN sock n = go n [] + where + go 0 acc = pure (BS.concat (reverse acc)) + go remaining acc = do + chunk <- NSB.recv sock (min 65536 remaining) + if BS.null chunk + then pure (BS.concat (reverse acc)) + else go (remaining - BS.length chunk) (chunk : acc) + +recvMaybe :: NS.Socket -> Int -> IO (Maybe ByteString) +recvMaybe sock n = do + bs <- NSB.recv sock n + if BS.null bs then pure Nothing else pure (Just bs) + +decodeInt32BE :: ByteString -> Int -> Int32 +decodeInt32BE bs off = + let b0 = fromIntegral (BS.index bs off) :: Int32 + b1 = fromIntegral (BS.index bs (off + 1)) :: Int32 + b2 = fromIntegral (BS.index bs (off + 2)) :: Int32 + b3 = fromIntegral (BS.index bs (off + 3)) :: Int32 + in b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3 diff --git a/pg-wire.cabal b/pg-wire.cabal new file mode 100644 --- /dev/null +++ b/pg-wire.cabal @@ -0,0 +1,168 @@ +cabal-version: 3.0 +name: pg-wire +version: 0.1.0.0 +synopsis: Pure Haskell PostgreSQL wire protocol driver +description: + A complete implementation of the PostgreSQL v3 wire protocol in + pure Haskell. Includes connection management, authentication + (SCRAM-SHA-256, MD5, cleartext), TLS, connection pooling, and + binary format type classes. No dependency on libpq or any C library. + +license: BSD-3-Clause +license-file: LICENSE +author: Josh Burgess +maintainer: joshburgess.webdev@gmail.com +category: Database +homepage: https://github.com/joshburgess/valiant +bug-reports: https://github.com/joshburgess/valiant/issues +build-type: Simple +extra-doc-files: + README.md + CHANGELOG.md +tested-with: GHC ==9.10.3 + +source-repository head + type: git + location: https://github.com/joshburgess/valiant + subdir: wire + +flag werror + description: Enable -Werror for development builds. + default: False + manual: True + +common warnings + ghc-options: -Wall -Wcompat -Wno-unticked-promoted-constructors -funbox-strict-fields -fspecialise-aggressively + if flag(werror) + ghc-options: -Werror + +library + import: warnings + hs-source-dirs: src + default-language: GHC2021 + default-extensions: + DerivingStrategies + LambdaCase + OverloadedStrings + RecordWildCards + StrictData + + exposed-modules: + PgWire.Async + PgWire.Auth.Cleartext + PgWire.Auth.MD5 + PgWire.Auth.ScramSHA256 + PgWire.Binary.Types + PgWire.Cancel + PgWire.Connection + PgWire.Connection.Config + PgWire.Error + PgWire.Pool + PgWire.Pool.Config + PgWire.Pool.Observation + PgWire.Protocol.Backend + PgWire.Protocol.Builders + PgWire.Protocol.Frontend + PgWire.Protocol.Oid + PgWire.Protocol.Parsers + PgWire.TypeCache + PgWire.Wire + + build-depends: + -- Core + , base >=4.17 && <5 + , bytestring >=0.11 && <0.13 + , text >=2.0 && <2.2 + , containers >=0.6 && <0.8 + , vector >=0.13 && <0.14 + , time >=1.12 && <1.15 + , deepseq >=1.4 && <1.6 + + -- Networking + , network >=3.1 && <3.3 + + -- TLS + , tls >=1.7 && <2.2 + , crypton-x509-store >=1.6 && <1.7 + , crypton-x509-system >=1.6 && <1.7 + + -- Crypto (auth) + , crypton >=0.34 && <1.1 + , cryptohash-md5 >=0.11 && <0.12 + , memory >=0.18 && <0.19 + , base64-bytestring >=1.2 && <1.3 + + -- Misc + , directory >=1.3 && <1.4 + , hashable >=1.4 && <1.6 + , nothunks >=0.1 && <0.3 + , psqueues >=0.2.8 && <0.3 + , random >=1.2 && <1.4 + + -- Concurrency / pool + , stm >=2.5 && <2.6 + , async >=2.2 && <2.3 + +-- Mock PostgreSQL server for testing client behaviour without a real +-- database. Sub-library so this isn't part of the public pg-wire API. +library mockserver + import: warnings + visibility: public + hs-source-dirs: mockserver + default-language: GHC2021 + default-extensions: + DerivingStrategies + LambdaCase + OverloadedStrings + RecordWildCards + StrictData + exposed-modules: + PgWire.MockServer + build-depends: + , base >=4.17 && <5 + , bytestring >=0.11 && <0.13 + , network >=3.1 && <3.3 + , pg-wire + +test-suite pg-wire-test + import: warnings + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + default-language: GHC2021 + default-extensions: + OverloadedStrings + + other-modules: + PgWire.Auth.MD5Spec + PgWire.Auth.ScramFieldsSpec + PgWire.CancelSpec + PgWire.Connection.ConfigSpec + PgWire.Connection.EscapingSpec + PgWire.Connection.FeaturesSpec + PgWire.ErrorSpec + PgWire.MockServerSpec + PgWire.Pool.ConfigSpec + PgWire.Pool.NoThunksSpec + PgWire.Pool.PropertySpec + PgWire.Pool.StateMachineSpec + PgWire.Protocol.BuildersSpec + PgWire.Protocol.FuzzSpec + PgWire.Protocol.OidSpec + PgWire.Protocol.ParsersSpec + + ghc-options: -rtsopts "-with-rtsopts=-K8K" -finfo-table-map -fdistinct-constructor-tables + build-depends: + , async >=2.2 && <2.3 + , base >=4.17 && <5 + , bytestring + , crypton + , deepseq + , hedgehog >=1.2 && <1.6 + , hspec >=2.11 && <2.13 + , hspec-hedgehog >=0.1 && <0.3 + , memory + , nothunks + , pg-wire + , pg-wire:mockserver + , vector diff --git a/src/PgWire/Async.hs b/src/PgWire/Async.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Async.hs @@ -0,0 +1,675 @@ +{-# OPTIONS_GHC -fno-full-laziness #-} + +-- | Asynchronous sender/receiver split for PostgreSQL connections. +-- +-- __Audience:__ this module is exposed for downstream library authors +-- (the @valiant@ runtime, custom adapters). Most application code +-- should reach for "PgWire.Connection" instead. The types and +-- functions here are public and stable, but they sit below the +-- request/response abstraction and require knowledge of the v3 wire +-- protocol to use safely. +-- +-- Each connection spawns two green threads after startup: +-- +-- * __Writer thread__ — drains the send queue, batches messages via +-- vectored I/O ('sendMany'), and enqueues response MVars. +-- * __Reader thread__ — parses backend messages from the socket and +-- fills MVars in FIFO order. +-- +-- Application threads submit 'Request's and block on the returned 'MVar'. +-- PostgreSQL processes messages in strict FIFO order, so the i-th response +-- always corresponds to the i-th pending request — no correlation IDs needed. +-- +-- This architecture enables automatic pipelining: when N threads query +-- concurrently, their messages are coalesced into a single 'sendMany' +-- syscall by the writer, and responses are demultiplexed by the reader. +module PgWire.Async + ( -- * Types + AsyncWireConn (..) + , Request (..) + , Response (..) + , ResponseCollector (..) + -- * Lifecycle + , spawnAsyncWireConn + , shutdownAsyncWireConn + -- * Submitting requests + , submitRequest + , submitExclusive + ) where + +import Control.Concurrent (forkIO) +import Control.Concurrent.Async (Async, async, cancel, link2) +import Control.Concurrent.MVar +import Control.Concurrent.STM +import Control.Exception (SomeException, catch, mask, onException, throwIO, try) +import Control.Monad (void) +import Data.ByteString (ByteString) +import Data.ByteString.Char8 qualified as BS8 +import Data.IORef +import Data.Int (Int32, Int64) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Vector (Vector) +import Data.Vector qualified as V +import Data.Vector.Mutable qualified as VM +import PgWire.Error (PgWireError (..), throwPgWire) +import PgWire.Protocol.Backend +import PgWire.Protocol.Frontend (FrontendMsg (..)) +import PgWire.Wire (WireConn, recvBackendMsg, sendFrontendMsgs) + +------------------------------------------------------------------------ +-- Types +------------------------------------------------------------------------ + +-- | What application threads submit to the async connection. +data Request + = -- | Extended query: pre-built messages (Bind+Execute+Sync) with a collector. + ReqExtendedQuery ![FrontendMsg] !ResponseCollector + | -- | Simple text query. + ReqSimpleQuery !ByteString + | -- | Parse a prepared statement (Parse+Sync). + ReqPrepare !FrontendMsg + | -- | Close a prepared statement (Close+Sync). + ReqClose !FrontendMsg + +-- | What comes back to the caller via MVar. +data Response + = -- | Rows collected from DataRow messages. + RespRows ![Vector (Maybe ByteString)] + | -- | Command tag from CommandComplete. + RespCommand !CommandTag + | -- | Simple query result: rows + optional tag. + RespSimple ![[Maybe ByteString]] !(Maybe CommandTag) + | -- | ParseComplete acknowledged. + RespParsed + | -- | CloseComplete acknowledged. + RespClosed + | -- | Multiple sub-results (one per Bind+Execute in a batch). + RespBatchRows ![[Vector (Maybe ByteString)]] + | -- | Batch command result: total rows affected. + RespBatchCommand !Int64 + | -- | Rows collected from DataRow messages AND the command tag. + RespRowsAndCommand ![Vector (Maybe ByteString)] !CommandTag + | -- | Rows collected as a Vector (not a list). Used by fetchAllVec. + RespRowsVec !(Vector (Vector (Maybe ByteString))) + | -- | At most one row. Used by fetchOne/fetchScalar/fetchExists. + RespFirstRow !(Maybe (Vector (Maybe ByteString))) + +-- | Tells the reader thread how to interpret backend messages for a request. +data ResponseCollector + = -- | Collect DataRow until CommandComplete, then ReadyForQuery. + CollectRows + | -- | Expect CommandComplete + ReadyForQuery, extract row count. + CollectCommand + | -- | Expect ParseComplete + ReadyForQuery. + CollectParseComplete + | -- | Expect CloseComplete + ReadyForQuery. + CollectCloseComplete + | -- | Collect N sub-results (Bind+Execute repeated N times), one ReadyForQuery at end. + CollectBatch !Int + | -- | Collect N command sub-results, one ReadyForQuery at end. + CollectBatchCommand !Int + | -- | Collect DataRow until CommandComplete, returning both rows and the command tag. + CollectRowsAndCommand + | -- | Like CollectRows but collects into a Vector instead of a list. + CollectRowsVec + | -- | Collect at most the first row. Discards remaining DataRows. + CollectFirstRow + | -- | Simple query protocol: text rows + optional tag + ReadyForQuery. + CollectSimple + +-- | Enqueued by the writer, dequeued by the reader. +data PendingResponse = PendingResponse + { prCollector :: !ResponseCollector + , prMVar :: !(MVar (Either PgWireError Response)) + } + +-- | The async connection state. Owns the writer and reader threads. +data AsyncWireConn = AsyncWireConn + { awcWire :: !WireConn + , awcSendQueue :: !(TBQueue (Request, MVar (Either PgWireError Response))) + , awcPending :: !(TQueue PendingResponse) + , awcExclusive :: !(TMVar (MVar (Either PgWireError Response))) + -- ^ When set, the writer pauses the pipeline for exclusive access. + , awcAlive :: !(TVar Bool) + , awcTxStatus :: !(IORef TxStatus) + , awcNotifyHandler :: !(IORef (Int32 -> ByteString -> ByteString -> IO ())) + , awcNoticeHandler :: !(IORef (PgNotice -> IO ())) + , awcParamStatus :: !(IORef (Map ByteString ByteString)) + , awcWriterThread :: !(Async ()) + , awcReaderThread :: !(Async ()) + , awcSendLock :: !(TMVar ()) + -- ^ Fast-path send lock. When the writer is idle, a caller can take + -- this lock and send directly on the wire, bypassing the TBQueue. + -- The reader still collects the response normally via the pending queue. + } + +-- | Shared state between writer/reader threads. Doesn't include the +-- thread handles themselves (avoids circular dependency with StrictData). +data AsyncCore = AsyncCore + { acWire :: !WireConn + , acSendQueue :: !(TBQueue (Request, MVar (Either PgWireError Response))) + , acPending :: !(TQueue PendingResponse) + , acExclusive :: !(TMVar (MVar (Either PgWireError Response))) + , acAlive :: !(TVar Bool) + , acTxStatus :: !(IORef TxStatus) + , acNotifyHandler :: !(IORef (Int32 -> ByteString -> ByteString -> IO ())) + , acNoticeHandler :: !(IORef (PgNotice -> IO ())) + , acParamStatus :: !(IORef (Map ByteString ByteString)) + , acSendLock :: !(TMVar ()) + } + +------------------------------------------------------------------------ +-- Lifecycle +------------------------------------------------------------------------ + +-- | Spawn writer and reader threads for an established connection. +-- Call this after the serial startup handshake completes. +spawnAsyncWireConn + :: WireConn + -> IORef TxStatus + -> IORef (Map ByteString ByteString) + -> IO AsyncWireConn +spawnAsyncWireConn wc txRef paramRef = do + sendQueue <- newTBQueueIO 64 + pending <- newTQueueIO + exclusive <- newEmptyTMVarIO + alive <- newTVarIO True + sendLock <- newTMVarIO () -- initially available + notifyHandler <- newIORef (\_ _ _ -> pure ()) + noticeHandler <- newIORef (\_ -> pure ()) + + -- The writer and reader threads need access to the shared STM/IORef state + -- but NOT to each other's Async handles. We build a "core" with just the + -- shared state, spawn threads from that, then assemble the final record. + let core = AsyncCore + { acWire = wc + , acSendQueue = sendQueue + , acPending = pending + , acExclusive = exclusive + , acAlive = alive + , acTxStatus = txRef + , acNotifyHandler = notifyHandler + , acNoticeHandler = noticeHandler + , acParamStatus = paramRef + , acSendLock = sendLock + } + + writer <- async (writerThread core) + reader <- async (readerThread core) + + -- Link threads: if either dies, both die. + link2 writer reader + + pure AsyncWireConn + { awcWire = wc + , awcSendQueue = sendQueue + , awcPending = pending + , awcExclusive = exclusive + , awcAlive = alive + , awcTxStatus = txRef + , awcNotifyHandler = notifyHandler + , awcNoticeHandler = noticeHandler + , awcParamStatus = paramRef + , awcWriterThread = writer + , awcReaderThread = reader + , awcSendLock = sendLock + } + +-- | Shut down the async connection. Signals threads to stop and cancels them. +shutdownAsyncWireConn :: AsyncWireConn -> IO () +shutdownAsyncWireConn awc = do + atomically $ writeTVar (awcAlive awc) False + cancel (awcWriterThread awc) `catch` \(_ :: SomeException) -> pure () + cancel (awcReaderThread awc) `catch` \(_ :: SomeException) -> pure () + drainPendingWithError awc ConnectionDead + +------------------------------------------------------------------------ +-- Submitting requests +------------------------------------------------------------------------ + +-- | Submit a request and block until the response arrives. +-- Throws 'ConnectionDead' if the async threads have died. +-- | Submit a request and block until the response arrives. +-- Fast path: if the send lock is available (no contention), send directly +-- on the wire without going through the writer thread's TBQueue. The reader +-- still collects the response normally. This eliminates ~2-4μs of thread +-- coordination overhead for single-threaded workloads. +-- Slow path: enqueue to the writer thread (enables automatic pipelining +-- under concurrency). +submitRequest :: AsyncWireConn -> Request -> IO Response +submitRequest awc req = do + respVar <- newEmptyMVar + -- Try fast path: take the send lock and send directly. + -- If the writer is busy (lock unavailable), fall back to the queue. + path <- atomically $ do + a <- readTVar (awcAlive awc) + if not a + then pure Nothing -- dead + else do + mLock <- tryTakeTMVar (awcSendLock awc) + case mLock of + Just () -> pure (Just True) -- fast path: we have the lock + Nothing -> do + writeTBQueue (awcSendQueue awc) (req, respVar) + pure (Just False) -- slow path: queued + case path of + Nothing -> throwPgWire ConnectionDead + Just False -> do + -- Slow path: wait for response from reader via writer + result <- takeMVar respVar + case result of + Left err -> throwIO err + Right resp -> pure resp + Just True -> do + -- Fast path: send directly, enqueue pending for reader + let (msgs, collector) = requestToMsgs req + releaseLock = atomically $ putTMVar (awcSendLock awc) () + sendFrontendMsgs (awcWire awc) msgs + `onException` releaseLock + -- Enqueue pending and release lock atomically so another fast-path + -- caller cannot send and enqueue before our pending is visible to + -- the reader — preserving the FIFO invariant. + atomically $ do + writeTQueue (awcPending awc) (PendingResponse collector respVar) + putTMVar (awcSendLock awc) () + -- Wait for response from reader + result <- takeMVar respVar + case result of + Left err -> throwIO err + Right resp -> pure resp + +-- | Submit an exclusive request. The writer thread will: +-- 1. Stop accepting new requests +-- 2. Wait for all pending responses to be collected by the reader +-- 3. Hand control to the caller +-- 4. Resume normal operation when the caller finishes +-- +-- Used for COPY, cursors, folds — operations that need direct socket access. +submitExclusive :: AsyncWireConn -> (WireConn -> IORef TxStatus -> IO a) -> IO a +submitExclusive awc action = mask $ \restore -> do + respVar <- newEmptyMVar + -- Combine alive check with exclusive signal in one STM transaction. + signaled <- atomically $ do + a <- readTVar (awcAlive awc) + if a + then putTMVar (awcExclusive awc) respVar >> pure True + else pure False + if not signaled + then throwPgWire ConnectionDead + else do + + -- Wait for writer to drain pending and hand us control + result <- takeMVar respVar + case result of + Left err -> throwIO err + Right _ -> do + -- We now have exclusive access to the WireConn. + -- Reader is blocked on empty awcPending. + -- Writer is blocked waiting for us to signal done. + let signalDone = do + doneVar <- newEmptyMVar + atomically $ putTMVar (awcExclusive awc) doneVar + putMVar doneVar (Right RespParsed) + a <- restore (action (awcWire awc) (awcTxStatus awc)) + `onException` signalDone + signalDone + pure a + +-- | Extract messages and collector from a request (used by fast path). +requestToMsgs :: Request -> ([FrontendMsg], ResponseCollector) +requestToMsgs (ReqExtendedQuery msgs collector) = (msgs, collector) +requestToMsgs (ReqSimpleQuery sql) = ([Query sql], CollectSimple) +requestToMsgs (ReqPrepare parseMsg) = ([parseMsg, Flush], CollectParseComplete) +requestToMsgs (ReqClose closeMsg) = ([closeMsg, Sync], CollectCloseComplete) + +------------------------------------------------------------------------ +-- Writer thread +------------------------------------------------------------------------ + +data WriterAction + = WriterExclusive !(MVar (Either PgWireError Response)) + | WriterBatch ![(Request, MVar (Either PgWireError Response))] + +writerThread :: AsyncCore -> IO () +writerThread ac = go `catch` onDeath + where + go :: IO () + go = do + action <- atomically $ tryExclusive `orElse` drainQueue + case action of + WriterExclusive respVar -> handleExclusive respVar + WriterBatch items -> handleBatch items + go + + tryExclusive :: STM WriterAction + tryExclusive = do + respVar <- takeTMVar (acExclusive ac) + pure (WriterExclusive respVar) + + drainQueue :: STM WriterAction + drainQueue = do + first <- readTBQueue (acSendQueue ac) + -- Take the send lock so fast-path callers can't send while we're sending. + takeTMVar (acSendLock ac) + rest <- drainTBQueue (acSendQueue ac) + pure (WriterBatch (first : rest)) + + handleBatch :: [(Request, MVar (Either PgWireError Response))] -> IO () + handleBatch items = do + let (allMsgs, pendings) = unzip (map buildItem items) + atomically $ mapM_ (writeTQueue (acPending ac)) pendings + sendFrontendMsgs (acWire ac) (concat allMsgs) + -- Release the send lock so fast-path callers can send again. + atomically $ putTMVar (acSendLock ac) () + + buildItem :: (Request, MVar (Either PgWireError Response)) -> ([FrontendMsg], PendingResponse) + buildItem (req, respVar) = case req of + ReqExtendedQuery msgs collector -> + (msgs, PendingResponse collector respVar) + ReqSimpleQuery sql -> + ([Query sql], PendingResponse CollectSimple respVar) + ReqPrepare parseMsg -> + ([parseMsg, Flush], PendingResponse CollectParseComplete respVar) + ReqClose closeMsg -> + ([closeMsg, Sync], PendingResponse CollectCloseComplete respVar) + + handleExclusive :: MVar (Either PgWireError Response) -> IO () + handleExclusive respVar = do + atomically $ do + empty <- isEmptyTQueue (acPending ac) + check empty + putMVar respVar (Right RespParsed) + doneVar <- atomically $ takeTMVar (acExclusive ac) + _ <- takeMVar doneVar + pure () + + onDeath :: SomeException -> IO () + onDeath _ = do + atomically $ writeTVar (acAlive ac) False + drainPendingWithErrorCore ac ConnectionDead + +------------------------------------------------------------------------ +-- Reader thread +------------------------------------------------------------------------ + +readerThread :: AsyncCore -> IO () +readerThread ac = go `catch` onDeath + where + go :: IO () + go = do + pr <- atomically $ readTQueue (acPending ac) + -- Catch query errors so they don't kill the reader thread. + -- A QueryError is per-request, not connection-fatal. + result <- try (collectResponse (prCollector pr)) + case result of + Right resp -> putMVar (prMVar pr) (Right resp) + Left (err :: PgWireError) -> do + -- After ErrorResponse with Sync, Postgres sends ReadyForQuery. + -- After ErrorResponse with Flush (preparation), it does NOT. + case prCollector pr of + CollectParseComplete -> pure () + _ -> drainUntilReady + putMVar (prMVar pr) (Left err) + go + + -- Drain messages until ReadyForQuery (error recovery). + drainUntilReady :: IO () + drainUntilReady = do + msg <- recvAndDispatch + case msg of + ReadyForQuery status -> writeIORef (acTxStatus ac) status + _ -> drainUntilReady + + collectResponse :: ResponseCollector -> IO Response + collectResponse CollectRows = do + rows <- collectRowsLoop + waitReadyForQuery + pure (RespRows rows) + collectResponse CollectCommand = do + tag <- collectCommandLoop + waitReadyForQuery + pure (RespCommand tag) + collectResponse CollectParseComplete = do + waitParseCompleteLoop + -- No waitReadyForQuery: Flush doesn't trigger ReadyForQuery. + pure RespParsed + collectResponse CollectCloseComplete = do + waitCloseCompleteLoop + waitReadyForQuery + pure RespClosed + collectResponse CollectRowsAndCommand = do + (rows, tag) <- collectRowsAndCommandLoop + waitReadyForQuery + pure (RespRowsAndCommand rows tag) + collectResponse CollectRowsVec = do + vec <- collectRowsVecLoop + waitReadyForQuery + pure (RespRowsVec vec) + collectResponse CollectFirstRow = do + mRow <- collectFirstRowLoop + waitReadyForQuery + pure (RespFirstRow mRow) + collectResponse (CollectBatch n) = do + results <- collectBatchLoop n + waitReadyForQuery + pure (RespBatchRows results) + collectResponse (CollectBatchCommand n) = do + total <- collectBatchCommandLoop n 0 + waitReadyForQuery + pure (RespBatchCommand total) + collectResponse CollectSimple = do + (rows, tag) <- collectSimpleLoop [] Nothing + pure (RespSimple rows tag) + + -- Collectors skip both ParseComplete and BindComplete so that + -- Parse+Bind+Execute can be coalesced into a single request. + collectRowsLoop :: IO [Vector (Maybe ByteString)] + collectRowsLoop = loop id + where + loop !acc = do + msg <- recvAndDispatch + case msg of + ParseComplete -> loop acc + BindComplete -> loop acc + DataRow vals -> loop (acc . (vals :)) + CommandComplete _ -> pure (acc []) + EmptyQueryResponse -> pure (acc []) + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Unexpected in rows: " <> BS8.pack (show other))) + + -- | Collect rows into a growable mutable vector. Starts at capacity 64, + -- doubles when full. Final freeze+slice produces an exact-size immutable Vector. + collectRowsVecLoop :: IO (Vector (Vector (Maybe ByteString))) + collectRowsVecLoop = do + mv <- VM.new 64 + (finalMv, !n) <- loop mv 0 + V.unsafeFreeze (VM.slice 0 n finalMv) + where + loop !mv !i = do + msg <- recvAndDispatch + case msg of + ParseComplete -> loop mv i + BindComplete -> loop mv i + DataRow vals -> do + mv' <- if i >= VM.length mv + then VM.grow mv (VM.length mv) -- double capacity + else pure mv + VM.write mv' i vals + loop mv' (i + 1) + CommandComplete _ -> pure (mv, i) + EmptyQueryResponse -> pure (mv, i) + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Unexpected in rows vec: " <> BS8.pack (show other))) + + -- | Collect at most the first row, discard the rest. + collectFirstRowLoop :: IO (Maybe (Vector (Maybe ByteString))) + collectFirstRowLoop = loop + where + loop = do + msg <- recvAndDispatch + case msg of + ParseComplete -> loop + BindComplete -> loop + DataRow vals -> do + -- Got first row, now drain remaining rows + drainRows + pure (Just vals) + CommandComplete _ -> pure Nothing + EmptyQueryResponse -> pure Nothing + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Unexpected in first row: " <> BS8.pack (show other))) + drainRows = do + msg <- recvAndDispatch + case msg of + DataRow _ -> drainRows -- discard + CommandComplete _ -> pure () + EmptyQueryResponse -> pure () + NoticeResponse _ -> drainRows + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Unexpected draining rows: " <> BS8.pack (show other))) + + collectRowsAndCommandLoop :: IO ([Vector (Maybe ByteString)], CommandTag) + collectRowsAndCommandLoop = loop id + where + loop !acc = do + msg <- recvAndDispatch + case msg of + ParseComplete -> loop acc + BindComplete -> loop acc + DataRow vals -> loop (acc . (vals :)) + CommandComplete tag -> pure (acc [], tag) + EmptyQueryResponse -> pure (acc [], OtherTag "EMPTY") + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Unexpected in rows+command: " <> BS8.pack (show other))) + + collectCommandLoop :: IO CommandTag + collectCommandLoop = do + msg <- recvAndDispatch + case msg of + ParseComplete -> collectCommandLoop + BindComplete -> collectCommandLoop + CommandComplete tag -> pure tag + EmptyQueryResponse -> pure (OtherTag "EMPTY") + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Unexpected in command: " <> BS8.pack (show other))) + + waitParseCompleteLoop :: IO () + waitParseCompleteLoop = do + msg <- recvAndDispatch + case msg of + ParseComplete -> pure () + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Expected ParseComplete, got: " <> BS8.pack (show other))) + + waitCloseCompleteLoop :: IO () + waitCloseCompleteLoop = do + msg <- recvAndDispatch + case msg of + CloseComplete -> pure () + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Expected CloseComplete, got: " <> BS8.pack (show other))) + + collectBatchLoop :: Int -> IO [[Vector (Maybe ByteString)]] + collectBatchLoop 0 = pure [] + collectBatchLoop n = do + rows <- collectRowsLoop + rest <- collectBatchLoop (n - 1) + pure (rows : rest) + + collectBatchCommandLoop :: Int -> Int64 -> IO Int64 + collectBatchCommandLoop 0 !total = pure total + collectBatchCommandLoop n !total = do + msg <- recvAndDispatch + case msg of + ParseComplete -> collectBatchCommandLoop n total + BindComplete -> collectBatchCommandLoop n total + CommandComplete tag -> collectBatchCommandLoop (n - 1) (total + tagRows tag) + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Unexpected in batch cmd: " <> BS8.pack (show other))) + + collectSimpleLoop :: [[Maybe ByteString]] -> Maybe CommandTag -> IO ([[Maybe ByteString]], Maybe CommandTag) + collectSimpleLoop !rows !tag = do + msg <- recvAndDispatch + case msg of + RowDescription _ -> collectSimpleLoop rows tag + DataRow vals -> collectSimpleLoop (V.toList vals : rows) tag + CommandComplete ct -> collectSimpleLoop rows (Just ct) + EmptyQueryResponse -> collectSimpleLoop rows tag + ReadyForQuery status -> do + writeIORef (acTxStatus ac) status + pure (reverse rows, tag) + ErrorResponse err -> throwIO (QueryError err) + other -> throwIO (ProtocolError ("Unexpected in simple query: " <> BS8.pack (show other))) + + waitReadyForQuery :: IO () + waitReadyForQuery = do + msg <- recvAndDispatch + case msg of + ReadyForQuery status -> writeIORef (acTxStatus ac) status + _ -> waitReadyForQuery + + recvAndDispatch :: IO BackendMsg + recvAndDispatch = do + msg <- recvBackendMsg (acWire ac) + case msg of + NotificationResponse pid channel payload -> do + handler <- readIORef (acNotifyHandler ac) + _ <- forkIO (handler pid channel payload `catch` \(_ :: SomeException) -> pure ()) + recvAndDispatch + NoticeResponse notice -> do + handler <- readIORef (acNoticeHandler ac) + _ <- forkIO (handler notice `catch` \(_ :: SomeException) -> pure ()) + recvAndDispatch + ParameterStatus key value -> do + modifyIORef' (acParamStatus ac) (Map.insert key value) + recvAndDispatch + _ -> pure msg + + onDeath :: SomeException -> IO () + onDeath _ = do + atomically $ writeTVar (acAlive ac) False + drainPendingWithErrorCore ac ConnectionDead + +------------------------------------------------------------------------ +-- Helpers +------------------------------------------------------------------------ + +drainTBQueue :: TBQueue a -> STM [a] +drainTBQueue q = loop id + where + loop acc = do + mItem <- tryReadTBQueue q + case mItem of + Just item -> loop (acc . (item :)) + Nothing -> pure (acc []) + +drainPendingWithError :: AsyncWireConn -> PgWireError -> IO () +drainPendingWithError awc err = do + pendings <- atomically $ flushTQueue (awcPending awc) + mapM_ (\pr -> void (tryPutMVar (prMVar pr) (Left err))) pendings + queued <- atomically $ drainTBQueue (awcSendQueue awc) + mapM_ (\(_, mv) -> void (tryPutMVar mv (Left err))) queued + mExcl <- atomically $ tryTakeTMVar (awcExclusive awc) + case mExcl of + Just mv -> void (tryPutMVar mv (Left err)) + Nothing -> pure () + +-- | Like 'drainPendingWithError' but operates on 'AsyncCore' (used by threads). +drainPendingWithErrorCore :: AsyncCore -> PgWireError -> IO () +drainPendingWithErrorCore ac err = do + pendings <- atomically $ flushTQueue (acPending ac) + mapM_ (\pr -> void (tryPutMVar (prMVar pr) (Left err))) pendings + queued <- atomically $ drainTBQueue (acSendQueue ac) + mapM_ (\(_, mv) -> void (tryPutMVar mv (Left err))) queued + mExcl <- atomically $ tryTakeTMVar (acExclusive ac) + case mExcl of + Just mv -> void (tryPutMVar mv (Left err)) + Nothing -> pure () + +tagRows :: CommandTag -> Int64 +tagRows (InsertTag n) = n +tagRows (UpdateTag n) = n +tagRows (DeleteTag n) = n +tagRows (SelectTag n) = n +tagRows (OtherTag _) = 0 diff --git a/src/PgWire/Auth/Cleartext.hs b/src/PgWire/Auth/Cleartext.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Auth/Cleartext.hs @@ -0,0 +1,17 @@ +-- | Cleartext password authentication. +-- +-- __Audience:__ exposed for downstream library authors. The standard +-- authentication flow is driven from "PgWire.Connection"; this module +-- is only useful if you are reimplementing handshake. +module PgWire.Auth.Cleartext + ( cleartextAuth + ) where + +import Data.ByteString (ByteString) +import PgWire.Protocol.Frontend (FrontendMsg (..)) +import PgWire.Wire (WireConn, sendFrontendMsg) + +-- | Handle cleartext password authentication. +cleartextAuth :: WireConn -> ByteString -> IO () +cleartextAuth wc password = + sendFrontendMsg wc (PasswordMessage password) diff --git a/src/PgWire/Auth/MD5.hs b/src/PgWire/Auth/MD5.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Auth/MD5.hs @@ -0,0 +1,39 @@ +-- | MD5 password authentication (legacy PostgreSQL auth method). +-- +-- __Audience:__ exposed for downstream library authors. The standard +-- authentication flow is driven from "PgWire.Connection"; this module +-- is only useful if you are reimplementing handshake. +module PgWire.Auth.MD5 + ( md5Auth + , md5Password + ) where + +import Crypto.Hash.MD5 qualified as MD5 +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.Word (Word8) +import PgWire.Protocol.Frontend (FrontendMsg (..)) +import PgWire.Wire (WireConn, sendFrontendMsg) +import Numeric (showHex) + +-- | Handle MD5 password authentication. +-- PG expects: @"md5" ++ md5(md5(password ++ user) ++ salt)@ +md5Auth :: WireConn -> ByteString -> ByteString -> ByteString -> IO () +md5Auth wc user password salt = + sendFrontendMsg wc (PasswordMessage (md5Password user password salt)) + +-- | Compute the MD5 password hash. +md5Password :: ByteString -> ByteString -> ByteString -> ByteString +md5Password user password salt = + let inner = hexMd5 (password <> user) + outer = hexMd5 (inner <> salt) + in "md5" <> outer + +hexMd5 :: ByteString -> ByteString +hexMd5 = BS.pack . concatMap toHexBytes . BS.unpack . MD5.hash + +toHexBytes :: Word8 -> [Word8] +toHexBytes w = + let s = showHex w "" + padded = if length s == 1 then '0' : s else s + in map (fromIntegral . fromEnum) padded diff --git a/src/PgWire/Auth/ScramSHA256.hs b/src/PgWire/Auth/ScramSHA256.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Auth/ScramSHA256.hs @@ -0,0 +1,163 @@ +-- | SCRAM-SHA-256 authentication (RFC 5802 / RFC 7677), with optional +-- @tls-server-end-point@ channel binding. +-- +-- __Audience:__ exposed for downstream library authors. The standard +-- authentication flow is driven from "PgWire.Connection"; this module +-- is only useful if you are reimplementing handshake. +module PgWire.Auth.ScramSHA256 + ( scramAuth + , scramAuthWithChannelBinding + ) where + +import Control.Monad (unless) +import Text.Read (readMaybe) +import Crypto.Hash (SHA256 (..), hashWith) +import Crypto.KDF.PBKDF2 qualified as PBKDF2 +import Crypto.MAC.HMAC (HMAC, hmac) +import Crypto.Random (getRandomBytes) +import Data.ByteArray qualified as BA +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Base64 qualified as B64 +import Data.ByteString.Char8 qualified as BS8 +import Data.Bits (xor) +import PgWire.Error (PgWireError (..), throwPgWire) +import PgWire.Protocol.Backend (AuthType (..), BackendMsg (..), PgError (..)) +import PgWire.Protocol.Frontend (FrontendMsg (..)) +import PgWire.Wire (WireConn, recvBackendMsg, sendFrontendMsg) + +-- | Perform SCRAM-SHA-256 authentication without channel binding. +scramAuth :: WireConn -> ByteString -> ByteString -> IO () +scramAuth wc user password = scramAuthInternal wc user password Nothing + +-- | Perform SCRAM-SHA-256-PLUS authentication with tls-server-end-point +-- channel binding. The ByteString argument is the SHA-256 hash of the +-- server's TLS certificate (DER encoded). +scramAuthWithChannelBinding :: WireConn -> ByteString -> ByteString -> ByteString -> IO () +scramAuthWithChannelBinding wc user password certHash = + scramAuthInternal wc user password (Just certHash) + +scramAuthInternal :: WireConn -> ByteString -> ByteString -> Maybe ByteString -> IO () +scramAuthInternal wc user password mCertHash = do + nonceBytes <- getRandomBytes 18 :: IO ByteString + let clientNonce = B64.encode nonceBytes + + -- Channel binding header: + -- "n,," for no channel binding + -- "p=tls-server-end-point,," for tls-server-end-point + (gs2Header, cbData, mechName) = case mCertHash of + Nothing -> + ("n,,", "n,,", "SCRAM-SHA-256") + Just certHash -> + let hdr = "p=tls-server-end-point,," + in (hdr, hdr <> certHash, "SCRAM-SHA-256-PLUS") + + clientFirstBare = "n=" <> user <> ",r=" <> clientNonce + clientFirstMsg = gs2Header <> clientFirstBare + + -- Send SASL initial response + sendFrontendMsg wc (SASLInitialResponse mechName clientFirstMsg) + + -- Receive server first message + msg1 <- recvBackendMsg wc + serverFirstMsg <- case msg1 of + Authentication (AuthSASLContinue serverData) -> pure serverData + Authentication AuthOk -> pure "" + ErrorResponse err -> throwPgWire (AuthError (pgMessage err)) + other -> throwPgWire (AuthError ("Unexpected message during SCRAM: " <> BS8.pack (show other))) + + -- Parse server first message: r=,s=,i= + let serverFields = parseScramFields serverFirstMsg + serverNonce <- lookupField "r" serverFields + saltB64 <- lookupField "s" serverFields + iterStr <- lookupField "i" serverFields + + iterations <- case readMaybe (BS8.unpack iterStr) :: Maybe Int of + Just n | n > 0 -> pure n + _ -> throwPgWire (AuthError ("Bad SCRAM iteration count: " <> iterStr)) + salt <- case B64.decode saltB64 of + Left err -> throwPgWire (AuthError ("Bad salt base64: " <> BS8.pack err)) + Right s -> pure s + + -- Verify server nonce starts with our client nonce + unless (clientNonce `BS.isPrefixOf` serverNonce) $ + throwPgWire (AuthError "Server nonce doesn't start with client nonce") + + -- Compute proofs + let saltedPassword = hi password salt iterations + clientKey = hmacSHA256 saltedPassword "Client Key" + storedKey = hashSHA256 clientKey + serverKey = hmacSHA256 saltedPassword "Server Key" + + channelBinding = B64.encode cbData + clientFinalWithoutProof = "c=" <> channelBinding <> ",r=" <> serverNonce + authMessage = clientFirstBare <> "," <> serverFirstMsg <> "," <> clientFinalWithoutProof + + clientSignature = hmacSHA256 storedKey authMessage + clientProof = BS.pack (BS.zipWith xor clientKey clientSignature) + serverSignature = hmacSHA256 serverKey authMessage + + clientFinalMsg = clientFinalWithoutProof <> ",p=" <> B64.encode clientProof + + -- Send client final message + sendFrontendMsg wc (SASLResponse clientFinalMsg) + + -- Receive server final message + msg2 <- recvBackendMsg wc + case msg2 of + Authentication (AuthSASLFinal serverData) -> do + let sFields = parseScramFields serverData + case lookup "v" sFields of + Nothing -> throwPgWire (AuthError "Missing server signature in SASL final") + Just sigB64 -> case B64.decode sigB64 of + Left err -> throwPgWire (AuthError ("Bad server sig base64: " <> BS8.pack err)) + Right sig + | not (BA.constEq sig serverSignature) -> + throwPgWire (AuthError "Server signature mismatch") + | otherwise -> pure () + Authentication AuthOk -> pure () + ErrorResponse err -> throwPgWire (AuthError (pgMessage err)) + other -> throwPgWire (AuthError ("Unexpected message during SCRAM final: " <> BS8.pack (show other))) + + -- Wait for AuthOk + msg3 <- recvBackendMsg wc + case msg3 of + Authentication AuthOk -> pure () + ErrorResponse err -> throwPgWire (AuthError (pgMessage err)) + _ -> throwPgWire (AuthError "Expected AuthOk after SCRAM") + +-- Crypto helpers ---------------------------------------------------------- + +hmacSHA256 :: ByteString -> ByteString -> ByteString +hmacSHA256 key msg = BA.convert (hmac key msg :: HMAC SHA256) +{-# INLINE hmacSHA256 #-} + +hashSHA256 :: ByteString -> ByteString +hashSHA256 bs = BA.convert (hashWith SHA256 bs) +{-# INLINE hashSHA256 #-} + +-- | PBKDF2-SHA256 (called "Hi" in SCRAM spec). +hi :: ByteString -> ByteString -> Int -> ByteString +hi password salt iterations = + PBKDF2.generate + (PBKDF2.prfHMAC SHA256) + (PBKDF2.Parameters iterations 32) + password + salt + +-- SCRAM field parsing ----------------------------------------------------- + +parseScramFields :: ByteString -> [(ByteString, ByteString)] +parseScramFields bs = + [ (key, val) + | field <- BS8.split ',' bs + , let (key, rest) = BS.break (== 61) field -- '=' + , not (BS.null rest) + , let val = BS.drop 1 rest + ] + +lookupField :: ByteString -> [(ByteString, ByteString)] -> IO ByteString +lookupField key fields = + case lookup key fields of + Nothing -> throwPgWire (AuthError ("Missing SCRAM field: " <> key)) + Just v -> pure v diff --git a/src/PgWire/Binary/Types.hs b/src/PgWire/Binary/Types.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Binary/Types.hs @@ -0,0 +1,25 @@ +-- | Type classes for encoding and decoding Haskell values to and from +-- PostgreSQL binary format. +-- +-- Instances are provided for common types ('Int32', 'Text', 'Bool', +-- 'UTCTime', etc.) in "Valiant.Binary.Encode" and "Valiant.Binary.Decode". +-- Custom types can provide their own instances. +module PgWire.Binary.Types + ( PgEncode (..) + , PgDecode (..) + ) where + +import Data.ByteString (ByteString) +import Data.Proxy (Proxy) +import PgWire.Protocol.Oid (Oid) + +-- | Encode a Haskell value to PostgreSQL binary format. +class PgEncode a where + pgEncode :: a -> ByteString + pgOid :: Proxy a -> Oid + {-# MINIMAL pgEncode, pgOid #-} + +-- | Decode a Haskell value from PostgreSQL binary format. +class PgDecode a where + pgDecode :: ByteString -> Either String a + {-# MINIMAL pgDecode #-} diff --git a/src/PgWire/Cancel.hs b/src/PgWire/Cancel.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Cancel.hs @@ -0,0 +1,88 @@ +-- | Query cancellation via the CancelRequest protocol. +-- +-- PostgreSQL's cancel mechanism works by opening a separate TCP connection +-- and sending a CancelRequest message containing the backend PID and +-- secret key from the original connection's BackendKeyData. The server +-- then signals the backend process to cancel its current query. +-- +-- @ +-- -- Cancel a long-running query from another thread: +-- 'cancelQuery' conn +-- +-- -- Or use a timeout wrapper: +-- 'withQueryTimeout' conn 5.0 $ do +-- fetchAll conn expensiveQuery () +-- @ +module PgWire.Cancel + ( cancelQuery + , withQueryTimeout + ) where + +import Control.Concurrent.Async (race) +import Control.Concurrent (threadDelay) +import Control.Exception (SomeException, try) +import Data.ByteString.Builder qualified as B +import Data.ByteString.Lazy qualified as LBS +import Data.Int (Int32) +import Data.Time (NominalDiffTime) +import Network.Socket qualified as NS +import Network.Socket.ByteString qualified as NSB +import PgWire.Connection (Connection (..)) +import PgWire.Connection.Config (ConnConfig (..)) +import PgWire.Error (PgWireError (..), throwPgWire) + +-- | Cancel the currently executing query on a connection. +-- +-- Opens a separate TCP connection to the same Postgres server and sends +-- a CancelRequest message. The original query will receive an ErrorResponse +-- with SQLSTATE 57014 (query_canceled). +-- +-- This is safe to call from any thread. +cancelQuery :: Connection -> IO () +cancelQuery conn = do + let cfg = connConfig conn + pid = connBackendPid conn + key = connBackendKey conn + result <- try @SomeException $ sendCancelRequest (ccHost cfg) (ccPort cfg) pid key + case result of + Right () -> pure () + Left _ -> pure () -- Best-effort: cancel may fail if server is unreachable + +-- | Run an IO action with a timeout. If the timeout expires, the query +-- is cancelled and a 'PoolTimeout' error is thrown. +-- +-- > withQueryTimeout conn 5.0 $ do +-- > fetchAll conn someExpensiveQuery () +withQueryTimeout :: Connection -> NominalDiffTime -> IO a -> IO a +withQueryTimeout conn seconds action = do + let micros = round (seconds * 1000000) :: Int + result <- race (threadDelay micros) action + case result of + Left () -> do + cancelQuery conn + throwPgWire (ConnectionError "Query timed out") + Right a -> pure a + +-- | Send a CancelRequest to the given host/port. +-- +-- CancelRequest format (not a normal frontend message): +-- [length :: Int32 = 16] +-- [cancel code :: Int32 = 80877102] +-- [pid :: Int32] +-- [key :: Int32] +sendCancelRequest :: NS.HostName -> NS.PortNumber -> Int32 -> Int32 -> IO () +sendCancelRequest host port pid key = do + let hints = NS.defaultHints {NS.addrSocketType = NS.Stream} + addrs <- NS.getAddrInfo (Just hints) (Just host) (Just (show port)) + case addrs of + [] -> throwPgWire (ConnectionError "Cancel: no address found") + (addr : _) -> do + sock <- NS.socket (NS.addrFamily addr) NS.Stream NS.defaultProtocol + NS.connect sock (NS.addrAddress addr) + let msg = LBS.toStrict . B.toLazyByteString $ + B.int32BE 16 + <> B.int32BE 80877102 + <> B.int32BE pid + <> B.int32BE key + NSB.sendAll sock msg + NS.close sock diff --git a/src/PgWire/Connection.hs b/src/PgWire/Connection.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Connection.hs @@ -0,0 +1,598 @@ +-- | PostgreSQL connection management. +-- +-- A 'Connection' represents an active TCP (or TLS) session with a +-- PostgreSQL server. It handles authentication, parameter negotiation, +-- and prepared statement caching. +-- +-- After startup, each connection spawns dedicated writer and reader +-- threads (see 'PgWire.Async') that enable automatic pipelining when +-- multiple green threads share a connection. +-- +-- For production use, prefer 'PgWire.Pool' over direct connections. +-- +-- @ +-- conn <- 'connectString' \"postgres:\/\/user:pass\@localhost:5432\/mydb\" +-- (rows, _) <- 'simpleQuery' conn \"SELECT 1\" +-- 'close' conn +-- @ +module PgWire.Connection + ( -- * Connection type + Connection (..) + , connWire + -- * Connecting + , connect + , connectString + , reset + -- * Closing + , close + , withConnection + -- * Simple queries + , simpleQuery + -- * Connection status + , connectionStatus + , transactionStatus + , parameterStatus + , serverVersion + , backendPid + , isSslInUse + -- * Notifications (non-blocking) + , checkNotification + -- * Notice handling + , setNoticeHandler + -- * SQL escaping + , escapeLiteral + , escapeIdentifier + -- * Statement introspection + , describePrepared + , ParamDescription (..) + , ColumnDescription (..) + -- * Server ping + , ping + -- * Password utilities + , lookupPgpass + , encryptPassword + -- * Protocol tracing + , setTraceHandler + ) where + +import Control.Exception (SomeException, bracket, catch, try) +import Crypto.Hash (MD5, hash, Digest) +import Data.ByteArray qualified as BA +import Data.ByteString qualified as BS +import System.Directory (doesFileExist) +import System.Environment (lookupEnv) +import System.Random (randomRIO) +import Data.ByteString (ByteString) +import Data.ByteString.Char8 qualified as BS8 +import Data.IORef +import Data.HashPSQ (HashPSQ) +import Data.HashPSQ qualified as PSQ +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Int (Int32) +import Data.Vector qualified as V +import Data.Vector (Vector) +import Data.Word (Word32, Word64) +import PgWire.Async (AsyncWireConn (..), Request (..), Response (..), spawnAsyncWireConn, shutdownAsyncWireConn, submitRequest, submitExclusive) +import PgWire.Auth.Cleartext (cleartextAuth) +import PgWire.Auth.MD5 (md5Auth) +import PgWire.Auth.ScramSHA256 (scramAuth) +import PgWire.Connection.Config (ConnConfig (..), TargetSessionAttrs (..), TlsMode (..), parseConnString, parseHosts) +import PgWire.Error (PgWireError (..), throwPgWire) +import PgWire.Protocol.Backend +import PgWire.Protocol.Builders (buildStartup) +import PgWire.Protocol.Frontend (DescribeTarget (..), FrontendMsg (..), StartupParams (..)) +import PgWire.Wire (TlsConfig (..), TraceDirection (..), WireConn (..), connectTcpTimeout, recvBackendMsg, sendFrontendMsg, sendRawBytes, upgradeTls) + +-- | A connection to a PostgreSQL database. +-- +-- After the initial handshake, dedicated writer and reader threads handle +-- all wire I/O. Use 'submitRequest' or the higher-level functions in +-- "Valiant.Execute" to interact with the connection. +data Connection = Connection + { connAsync :: !AsyncWireConn + -- ^ The async wire connection (writer + reader threads). + , connConfig :: !ConnConfig + , connBackendPid :: {-# UNPACK #-} !Int32 + , connBackendKey :: {-# UNPACK #-} !Int32 + , connStmtCache :: !(IORef (HashPSQ ByteString Word64 ByteString)) + -- ^ LRU prepared statement cache. Keys are SQL text, priorities are + -- monotonic access ticks (lower = least recently used), values are + -- server-side statement names. + , connStmtCounter :: !(IORef Word64) + , connStmtTick :: !(IORef Word64) + -- ^ Monotonic tick counter for LRU cache priorities. + , connSslActive :: !Bool + , connPreparedStatements :: !Bool + -- ^ Whether to use named prepared statements (default: True). + -- Set to False for compatibility with PgBouncer in transaction mode. + -- When False, all queries use the unnamed statement (parsed + discarded + -- each time, never cached). + , connMaxPreparedStatements :: !Int + -- ^ Maximum number of prepared statements cached on this connection. + } + +-- | Access the underlying 'WireConn' for direct wire operations. +-- Only use this inside exclusive mode ('submitExclusive'). +connWire :: Connection -> WireConn +connWire = awcWire . connAsync + +-- | Access the transaction status IORef. +connTxStatus :: Connection -> IORef TxStatus +connTxStatus = awcTxStatus . connAsync + +-- | Access the server parameters IORef. +connParams :: Connection -> IORef (Map ByteString ByteString) +connParams = awcParamStatus . connAsync + +-- | Access the notice handler IORef. +connNoticeHandler :: Connection -> IORef (PgNotice -> IO ()) +connNoticeHandler = awcNoticeHandler . connAsync + +-- | Parameter type information returned by 'describePrepared'. +data ParamDescription = ParamDescription + { pdOids :: Vector Word32 + -- ^ OIDs of each parameter's type (one per @$N@ placeholder). + } + deriving stock (Show, Eq) + +-- | Column metadata returned by 'describePrepared'. +data ColumnDescription = ColumnDescription + { cdFields :: Vector FieldInfo + -- ^ Name, type OID, table OID, column number, and format for each result column. + } + deriving stock (Show, Eq) + +-- | Connect using a 'ConnConfig'. Supports multi-host failover: +-- if @ccHost@ contains comma-separated hosts (@\"host1,host2\"@), +-- tries each in order until one succeeds. If @ccTargetSessionAttrs@ +-- is set, verifies the server matches (e.g., primary vs standby). +connect :: ConnConfig -> IO Connection +connect cfg0 = do + -- Auto-lookup password from ~/.pgpass if not provided + cfg <- if BS8.null (ccPassword cfg0) + then do + mPw <- lookupPgpass cfg0 + pure $ case mPw of + Just pw -> cfg0 { ccPassword = pw } + Nothing -> cfg0 + else pure cfg0 + let hosts = parseHosts (ccHost cfg) + orderedHosts <- if ccLoadBalanceHosts cfg + then shuffleHosts hosts + else pure hosts + tryHosts cfg orderedHosts + +tryHosts :: ConnConfig -> [String] -> IO Connection +tryHosts _ [] = throwPgWire (ConnectionError "All hosts failed") +tryHosts cfg [h] = connectSingleHost cfg { ccHost = h } +tryHosts cfg (h : hs) = do + result <- try @SomeException (connectSingleHost cfg { ccHost = h }) + case result of + Right conn -> do + ok <- checkSessionAttrs conn (ccTargetSessionAttrs cfg) + if ok + then pure conn + else do + close conn + tryHosts cfg hs + Left _ -> tryHosts cfg hs + +checkSessionAttrs :: Connection -> TargetSessionAttrs -> IO Bool +checkSessionAttrs _ SessionAny = pure True +checkSessionAttrs conn attr = do + (rows, _) <- simpleQuery conn "SHOW transaction_read_only" + case rows of + [[Just val]] -> + let readOnly = val == "on" + in pure $ case attr of + SessionReadWrite -> not readOnly + SessionReadOnly -> readOnly + SessionPrimary -> not readOnly + SessionStandby -> readOnly + SessionPreferStandby -> True + _ -> pure True + +-- | Fisher-Yates shuffle for load balancing +shuffleHosts :: [a] -> IO [a] +shuffleHosts [] = pure [] +shuffleHosts [x] = pure [x] +shuffleHosts xs = do + let arr = zip [0 :: Int ..] xs + n = length xs + go (reverse arr) (n - 1) + where + go [] _ = pure [] + go ((_, x) : _rest) 0 = pure [x] + go items i = do + j <- randomRIO (0, i) + let (before, rest) = splitAt j items + case rest of + [] -> pure (map snd before) -- unreachable: j <= length items - 1 + (_, picked) : after -> (picked :) <$> go (before ++ after) (i - 1) + +connectSingleHost :: ConnConfig -> IO Connection +connectSingleHost cfg = do + wc0 <- connectTcpTimeout (ccConnectTimeout cfg) (ccHost cfg) (ccPort cfg) + + let tlsCfg = TlsConfig + { tlsHostname = ccHost cfg + , tlsVerify = ccTls cfg `elem` [TlsVerifyCa, TlsVerifyFull] + , tlsVerifyHostname = ccTls cfg == TlsVerifyFull + , tlsClientCert = ccSslCert cfg + , tlsClientKey = ccSslKey cfg + , tlsCaCert = ccSslRootCert cfg + } + + -- Optionally upgrade to TLS + wc <- case ccTls cfg of + TlsDisable -> pure wc0 + TlsRequire -> upgradeTls wc0 tlsCfg + TlsVerifyCa -> upgradeTls wc0 tlsCfg + TlsVerifyFull -> upgradeTls wc0 tlsCfg + TlsPrefer -> do + result <- try @SomeException (upgradeTls wc0 tlsCfg) + case result of + Right tlsWc -> pure tlsWc + Left _ -> do + wcClose wc0 + connectTcpTimeout (ccConnectTimeout cfg) (ccHost cfg) (ccPort cfg) + + -- Send startup message + let extraParams = + [ ("client_encoding", ccClientEncoding cfg) | ccClientEncoding cfg /= "UTF8" ] + ++ [ ("options", ccOptions cfg) | not (BS8.null (ccOptions cfg)) ] + startup = + StartupParams + { spUser = ccUser cfg + , spDatabase = ccDatabase cfg + , spAppName = ccAppName cfg + , spExtraParams = extraParams + } + sendRawBytes wc (buildStartup startup) + + -- Handle authentication + handleAuth wc cfg + + -- Collect server parameters until ReadyForQuery + paramsRef <- newIORef Map.empty + pidRef <- newIORef 0 + keyRef <- newIORef 0 + txRef <- newIORef TxIdle + + let loop = do + msg <- recvBackendMsg wc + case msg of + ParameterStatus name value -> do + modifyIORef' paramsRef (Map.insert name value) + loop + BackendKeyData pid key -> do + writeIORef pidRef pid + writeIORef keyRef key + loop + ReadyForQuery status -> do + writeIORef txRef status + pure () + NoticeResponse _ -> loop + other -> + throwPgWire (ProtocolError ("Unexpected message during startup: " <> BS8.pack (show other))) + + loop + + pid <- readIORef pidRef + key <- readIORef keyRef + stmtCache <- newIORef PSQ.empty + stmtCounter <- newIORef 0 + stmtTick <- newIORef 0 + let sslActive = case ccTls cfg of + TlsDisable -> False + _ -> True + + -- Spawn async writer + reader threads + asyncConn <- spawnAsyncWireConn wc txRef paramsRef + + pure + Connection + { connAsync = asyncConn + , connConfig = cfg + , connBackendPid = pid + , connBackendKey = key + , connStmtCache = stmtCache + , connStmtCounter = stmtCounter + , connStmtTick = stmtTick + , connSslActive = sslActive + , connPreparedStatements = ccPreparedStatements cfg + , connMaxPreparedStatements = max 1 (ccMaxPreparedStatements cfg) + } + +-- | Connect using a connection string. +connectString :: ByteString -> IO Connection +connectString bs = case parseConnString bs of + Left err -> throwPgWire (ConnectionError (BS8.pack err)) + Right cfg -> connect cfg + +-- | Close a connection. +close :: Connection -> IO () +close conn = do + shutdownAsyncWireConn (connAsync conn) + -- Send Terminate and close the socket (best effort, threads are already dead) + sendFrontendMsg (connWire conn) Terminate `catch_` pure () + wcClose (connWire conn) `catch_` pure () + where + catch_ :: IO a -> IO a -> IO a + catch_ action fallback = action `catch` \(_ :: SomeException) -> fallback + +-- | Bracket-style connection management. +withConnection :: ConnConfig -> (Connection -> IO a) -> IO a +withConnection cfg = bracket (connect cfg) close + +-- | Execute a simple text query (no parameters). Returns all result rows +-- as lists of nullable bytestrings, plus the command tag. +simpleQuery :: Connection -> ByteString -> IO ([[Maybe ByteString]], Maybe CommandTag) +simpleQuery conn sql = do + resp <- submitRequest (connAsync conn) (ReqSimpleQuery sql) + case resp of + RespSimple rows tag -> pure (rows, tag) + _ -> throwPgWire (ProtocolError "simpleQuery: unexpected response type") + +-- | Reset the connection: close and reconnect using the same config. +reset :: Connection -> IO Connection +reset conn = do + close conn + connect (connConfig conn) + +-- Connection status ------------------------------------------------------- + +-- | Check if the connection is alive by sending an empty query. +connectionStatus :: Connection -> IO Bool +connectionStatus conn = do + result <- try @SomeException (simpleQuery conn "") + pure $ case result of + Right _ -> True + Left _ -> False + +-- | Get the current transaction status. +transactionStatus :: Connection -> IO TxStatus +transactionStatus conn = readIORef (connTxStatus conn) + +-- | Look up a server parameter (e.g., @\"server_version\"@, @\"server_encoding\"@). +parameterStatus :: Connection -> ByteString -> IO (Maybe ByteString) +parameterStatus conn key = do + params <- readIORef (connParams conn) + pure (Map.lookup key params) + +-- | Get the server version as an integer (e.g., 160004 for 16.4). +-- Parses from the @server_version@ parameter. +serverVersion :: Connection -> IO (Maybe Int) +serverVersion conn = do + mVersion <- parameterStatus conn "server_version" + pure (mVersion >>= parseServerVersion) + where + parseServerVersion bs = + -- Strip distribution suffix: "16.6 (Ubuntu 16.6-0ubuntu0.24.04.1)" → "16.6" + let versionOnly = BS8.takeWhile (\c -> c == '.' || (c >= '0' && c <= '9')) bs + parts = BS8.split '.' versionOnly + in case parts of + [major, minor] -> do + maj <- readInt major + mn <- readInt minor + pure (maj * 10000 + mn) + [major, minor, patch] -> do + maj <- readInt major + mn <- readInt minor + p <- readInt patch + pure (maj * 10000 + mn * 100 + p) + _ -> Nothing + readInt s = case BS8.readInt s of + Just (n, _) -> Just n + Nothing -> Nothing + +-- | Get the backend process ID. +backendPid :: Connection -> Int32 +backendPid = connBackendPid + +-- | Check if SSL/TLS is in use on this connection. +isSslInUse :: Connection -> Bool +isSslInUse = connSslActive + +-- Notifications (non-blocking) -------------------------------------------- + +-- | Check for a pending notification without blocking. +-- Returns 'Nothing' if no notification is available. +checkNotification :: Connection -> IO (Maybe (Int32, ByteString, ByteString)) +checkNotification conn = do + -- Send empty query to flush pending notifications from the server + _ <- simpleQuery conn "" + -- Notifications are dispatched by the reader thread to the handler. + -- For non-blocking check, we'd need a different mechanism. + -- For now, return Nothing — notifications arrive via the handler. + pure Nothing + +-- Notice handling --------------------------------------------------------- + +-- | Set a callback for server notice messages (warnings, info, etc.). +-- The default handler discards all notices. +setNoticeHandler :: Connection -> (PgNotice -> IO ()) -> IO () +setNoticeHandler conn handler = writeIORef (connNoticeHandler conn) handler + +-- SQL escaping ------------------------------------------------------------ + +-- | Escape a string for use as a SQL literal. Returns a properly quoted +-- and escaped string including the surrounding single quotes. +-- +-- @ +-- escapeLiteral conn "O'Brien" == "'O''Brien'" +-- @ +escapeLiteral :: Connection -> ByteString -> ByteString +escapeLiteral _conn bs = + "'" <> BS8.concatMap escapeChar bs <> "'" + where + escapeChar '\'' = "''" + escapeChar '\\' = "\\\\" + escapeChar c = BS8.singleton c + +-- | Escape a string for use as a SQL identifier (table, column, function name). +-- Returns a properly quoted identifier with surrounding double quotes. +-- +-- @ +-- escapeIdentifier conn "user table" == "\"user table\"" +-- @ +escapeIdentifier :: Connection -> ByteString -> ByteString +escapeIdentifier _conn bs = + "\"" <> BS8.concatMap escapeChar bs <> "\"" + where + escapeChar '"' = "\"\"" + escapeChar c = BS8.singleton c + +-- Statement introspection ------------------------------------------------- + +-- | Describe a prepared statement. Returns parameter OIDs and column metadata. +-- The statement must already be prepared on this connection. +describePrepared :: Connection -> ByteString -> IO (ParamDescription, ColumnDescription) +describePrepared conn stmtName = do + -- describePrepared needs direct wire access (Describe + Sync protocol) + submitExclusive (connAsync conn) $ \wc txRef -> do + sendFrontendMsg wc (Describe DescribeStatement stmtName) + sendFrontendMsg wc Sync + (params, cols) <- collectDescribe wc txRef + pure (ParamDescription params, ColumnDescription cols) + where + collectDescribe wc txRef = do + pds <- collectParams wc + cds <- collectColumns wc + waitDescribeReady wc txRef + pure (pds, cds) + + collectParams wc = do + msg <- recvBackendMsg wc + case msg of + ParameterDescription oids -> pure oids + ErrorResponse err -> throwPgWire (QueryError err) + _ -> collectParams wc + + collectColumns wc = do + msg <- recvBackendMsg wc + case msg of + RowDescription fields -> pure fields + NoData -> pure V.empty + ErrorResponse err -> throwPgWire (QueryError err) + _ -> collectColumns wc + + waitDescribeReady wc txRef = do + msg <- recvBackendMsg wc + case msg of + ReadyForQuery status -> writeIORef txRef status + _ -> waitDescribeReady wc txRef + +-- Server ping ------------------------------------------------------------- + +-- | Check if a PostgreSQL server is accepting connections, without +-- fully authenticating. Attempts a TCP connection and checks if the +-- server responds to the startup sequence. +ping :: ConnConfig -> IO Bool +ping cfg = do + result <- try @SomeException (connect cfg >>= close) + pure $ case result of + Right () -> True + Left _ -> False + +-- Password file ----------------------------------------------------------- + +-- | Look up a password from @~/.pgpass@ file. +-- Format: @hostname:port:database:username:password@ (one per line). +-- @*@ matches any value in a field. +lookupPgpass :: ConnConfig -> IO (Maybe ByteString) +lookupPgpass cfg = do + home <- lookupEnv "HOME" + case home of + Nothing -> pure Nothing + Just h -> do + let path = h <> "/.pgpass" + exists <- doesFileExist path + if not exists + then pure Nothing + else do + contents <- BS8.readFile path + let host = BS8.pack (ccHost cfg) + port = BS8.pack (show (ccPort cfg)) + db = ccDatabase cfg + user = ccUser cfg + pure (findMatch host port db user (BS8.lines contents)) + where + findMatch _ _ _ _ [] = Nothing + findMatch host port db user (line : rest) + | BS8.null line || BS8.isPrefixOf "#" line = findMatch host port db user rest + | otherwise = case BS8.split ':' line of + [h, p, d, u, pw] + | matches h host && matches p port && matches d db && matches u user -> + Just pw + _ -> findMatch host port db user rest + matches "*" _ = True + matches pattern value = pattern == value + +-- Password encryption ----------------------------------------------------- + +-- | Encrypt a password for storage, using the same algorithm as +-- @PQencryptPasswordConn@. Supports MD5 format. +encryptPassword :: ByteString -> ByteString -> ByteString +encryptPassword user password = + -- MD5 format: "md5" + hex(md5(password + user)) + let digest = BA.convert (hash (password <> user) :: Digest MD5) :: ByteString + in "md5" <> toHex digest + +toHex :: ByteString -> ByteString +toHex bs = BS8.pack (concatMap byteToHex (BS.unpack bs)) + where + byteToHex w = + let (hi, lo) = w `divMod` 16 + in [hexDigit hi, hexDigit lo] + hexDigit n + | n < 10 = toEnum (fromEnum '0' + fromIntegral n) + | otherwise = toEnum (fromEnum 'a' + fromIntegral n - 10) + +-- Protocol tracing -------------------------------------------------------- + +-- | Enable protocol tracing on this connection. +-- The callback receives 'True' for send, 'False' for recv, plus the raw bytes. +-- Set to @Nothing@ to disable tracing. +-- +-- @ +-- setTraceHandler conn $ \\isSend bytes -> +-- BS8.putStrLn $ (if isSend then \">>> \" else \"<<< \") <> BS8.take 40 bytes +-- @ +setTraceHandler :: Connection -> (Bool -> ByteString -> IO ()) -> IO () +setTraceHandler conn handler = + writeIORef (wcTrace (connWire conn)) (Just wrapper) + where + wrapper TraceSend bs = handler True bs + wrapper TraceRecv bs = handler False bs + +-- Authentication ---------------------------------------------------------- + +handleAuth :: WireConn -> ConnConfig -> IO () +handleAuth wc cfg = do + msg <- recvBackendMsg wc + case msg of + Authentication AuthOk -> pure () + Authentication AuthCleartextPassword -> do + cleartextAuth wc (ccPassword cfg) + waitForAuthOk wc + Authentication (AuthMD5Password salt) -> do + md5Auth wc (ccUser cfg) (ccPassword cfg) salt + waitForAuthOk wc + Authentication (AuthSASL mechs) + | "SCRAM-SHA-256" `elem` mechs -> do + scramAuth wc (ccUser cfg) (ccPassword cfg) + -- scramAuth handles waiting for AuthOk internally + | otherwise -> + throwPgWire (AuthError ("Unsupported SASL mechanisms: " <> BS8.pack (show mechs))) + ErrorResponse err -> throwPgWire (AuthError (pgMessage err)) + other -> throwPgWire (AuthError ("Unexpected auth message: " <> BS8.pack (show other))) + +waitForAuthOk :: WireConn -> IO () +waitForAuthOk wc = do + msg <- recvBackendMsg wc + case msg of + Authentication AuthOk -> pure () + ErrorResponse err -> throwPgWire (AuthError (pgMessage err)) + other -> throwPgWire (AuthError ("Expected AuthOk, got: " <> BS8.pack (show other))) diff --git a/src/PgWire/Connection/Config.hs b/src/PgWire/Connection/Config.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Connection/Config.hs @@ -0,0 +1,298 @@ +-- | Configuration types for PostgreSQL connections. +-- +-- Supports both URI format (@postgres:\/\/user:pass\@host:port\/db@) and +-- key=value format (@host=localhost port=5432 dbname=mydb@). Multi-host +-- failover, TLS modes, and session attribute targeting are configurable. +module PgWire.Connection.Config + ( ConnConfig (..) + , TlsMode (..) + , TargetSessionAttrs (..) + , defaultConnConfig + , parseConnString + , parseHosts + ) where + +import Data.ByteString (ByteString) +import Data.ByteString.Char8 qualified as BS8 +import Data.Char (isDigit) +import Data.Maybe (fromMaybe) +import Data.Time (NominalDiffTime) +import Data.Word (Word16) +import Network.Socket (HostName, PortNumber) + +-- | TLS connection mode, matching PostgreSQL's @sslmode@ parameter. +data TlsMode + = TlsDisable -- ^ No SSL + | TlsPrefer -- ^ Try SSL, fall back to plain + | TlsRequire -- ^ Require SSL, don't verify server cert + | TlsVerifyCa -- ^ Require SSL, verify server cert against CA + | TlsVerifyFull -- ^ Require SSL, verify cert + hostname match + deriving stock (Show, Eq) + +-- | Session attribute requirements for multi-host failover. +data TargetSessionAttrs + = SessionAny -- ^ Connect to any server + | SessionReadWrite -- ^ Must be a read-write server (primary) + | SessionReadOnly -- ^ Must be a read-only server (standby) + | SessionPrimary -- ^ Must be the primary + | SessionStandby -- ^ Must be a standby + | SessionPreferStandby -- ^ Prefer standby, fall back to primary + deriving stock (Show, Eq) + +-- | Connection configuration for a single PostgreSQL server. +data ConnConfig = ConnConfig + { ccHost :: HostName + -- ^ Primary host. For multi-host, use comma-separated: @\"host1,host2\"@. + , ccPort :: PortNumber + , ccDatabase :: ByteString + , ccUser :: ByteString + , ccPassword :: ByteString + , ccTls :: TlsMode + , ccAppName :: ByteString + , ccConnectTimeout :: NominalDiffTime + -- ^ Timeout for establishing the TCP connection (seconds). 0 = no timeout. + , ccQueryTimeout :: NominalDiffTime + -- ^ Default timeout for query execution (seconds). 0 = no timeout. + -- Can be overridden per-query with 'withQueryTimeout'. + , ccSslCert :: Maybe FilePath + -- ^ Path to client SSL certificate file (@sslcert@). + , ccSslKey :: Maybe FilePath + -- ^ Path to client SSL private key file (@sslkey@). + , ccSslRootCert :: Maybe FilePath + -- ^ Path to SSL certificate authority (CA) file (@sslrootcert@). + -- If @\"system\"@, uses the system certificate store. + , ccKeepalives :: Bool + -- ^ Enable TCP keepalives (default: True). + , ccKeepalivesIdle :: Int + -- ^ Seconds before first keepalive probe (default: 0 = system default). + , ccKeepalivesInterval :: Int + -- ^ Seconds between keepalive probes (default: 0 = system default). + , ccKeepalivesCount :: Int + -- ^ Number of failed probes before disconnect (default: 0 = system default). + , ccTargetSessionAttrs :: TargetSessionAttrs + -- ^ Required session attributes for multi-host failover (default: 'SessionAny'). + , ccClientEncoding :: ByteString + -- ^ Client encoding (default: @\"UTF8\"@). Set to @\"auto\"@ for auto-detection. + , ccOptions :: ByteString + -- ^ Extra command-line options to send to the server at startup. + , ccLoadBalanceHosts :: Bool + -- ^ Randomize host order for multi-host connections (default: False). + , ccPreparedStatements :: Bool + -- ^ Use named prepared statements (default: True). Set to False for + -- compatibility with PgBouncer in transaction mode or pgpool. + -- When False, all queries use the unnamed statement which is + -- parsed and discarded each time. + , ccMaxPreparedStatements :: !Int + -- ^ Maximum number of prepared statements cached per connection + -- (default: 256). When the cache is full, the least recently used + -- statement is evicted. Higher values reduce re-parses for workloads + -- with many distinct queries, at the cost of server-side memory + -- (each cached statement holds a parse tree in PostgreSQL). + -- Ignored when 'ccPreparedStatements' is False. + } + deriving stock (Show) + +-- | Default connection configuration: @localhost:5432@, no TLS, 10s connect timeout. +defaultConnConfig :: ConnConfig +defaultConnConfig = + ConnConfig + { ccHost = "localhost" + , ccPort = 5432 + , ccDatabase = "" + , ccUser = "" + , ccPassword = "" + , ccTls = TlsDisable + , ccAppName = "pg-wire" + , ccConnectTimeout = 10 + , ccQueryTimeout = 0 + , ccSslCert = Nothing + , ccSslKey = Nothing + , ccSslRootCert = Nothing + , ccKeepalives = True + , ccKeepalivesIdle = 0 + , ccKeepalivesInterval = 0 + , ccKeepalivesCount = 0 + , ccTargetSessionAttrs = SessionAny + , ccClientEncoding = "UTF8" + , ccOptions = "" + , ccLoadBalanceHosts = False + , ccPreparedStatements = True + , ccMaxPreparedStatements = 256 + } + +-- | Parse a PostgreSQL connection string. +-- Supports both URI format (@postgres://user:pass\@host:port/db@) +-- and key=value format (@host=localhost port=5432 dbname=mydb@). +parseConnString :: ByteString -> Either String ConnConfig +parseConnString bs + | "postgres://" `BS8.isPrefixOf` bs || "postgresql://" `BS8.isPrefixOf` bs = + parseUri bs + | otherwise = + parseKeyValue bs + +parseUri :: ByteString -> Either String ConnConfig +parseUri bs = do + -- Strip scheme: "postgres://..." -> drop up to first '/', then drop "//" + let afterScheme = BS8.drop 2 (BS8.dropWhile (/= '/') bs) -- drop "scheme://" + afterSlash = afterScheme + + -- Split user:pass@host:port/db?params + let (authHost, pathQuery) = case BS8.break (== '/') afterSlash of + (ah, pq) -> (ah, BS8.drop 1 pq) + dbAndParams = pathQuery + (db, queryStr) = BS8.break (== '?') dbAndParams + + -- Split auth@host + let (auth, hostPort) = case BS8.breakEnd (== '@') authHost of + (a, _) | BS8.null a -> ("", authHost) + (a, hp) -> (BS8.init a, hp) -- drop trailing '@' + + -- Split user:pass + let (user, pass) = case BS8.break (== ':') auth of + (u, p) + | BS8.null p -> (u, "") + | otherwise -> (u, BS8.drop 1 p) + + -- Split host:port + let (host, portStr) = case BS8.break (== ':') hostPort of + (h, p) + | BS8.null p -> (h, "5432") + | otherwise -> (h, BS8.drop 1 p) + port = maybe 5432 fromIntegral (readPort portStr) + + -- Parse query params for sslmode + let params = parseQueryParams (BS8.drop 1 queryStr) -- drop '?' + tlsMode = case lookup "sslmode" params of + Just "require" -> TlsRequire + Just "prefer" -> TlsPrefer + Just "disable" -> TlsDisable + Just "verify-ca" -> TlsVerifyCa + Just "verify-full" -> TlsVerifyFull + _ -> TlsDisable + sslCert = BS8.unpack <$> lookup "sslcert" params + sslKey = BS8.unpack <$> lookup "sslkey" params + sslRootCert = BS8.unpack <$> lookup "sslrootcert" params + + Right + ConnConfig + { ccHost = BS8.unpack host + , ccPort = port + , ccDatabase = db + , ccUser = user + , ccPassword = pass + , ccTls = tlsMode + , ccAppName = "pg-wire" + , ccConnectTimeout = readTimeout (lookup "connect_timeout" params) + , ccQueryTimeout = 0 + , ccSslCert = sslCert + , ccSslKey = sslKey + , ccSslRootCert = sslRootCert + , ccKeepalives = True + , ccKeepalivesIdle = 0 + , ccKeepalivesInterval = 0 + , ccKeepalivesCount = 0 + , ccTargetSessionAttrs = maybe SessionAny parseSessionAttrs + (lookup "target_session_attrs" params) + , ccClientEncoding = fromMaybe "UTF8" (lookup "client_encoding" params) + , ccOptions = fromMaybe "" (lookup "options" params) + , ccLoadBalanceHosts = lookup "load_balance_hosts" params == Just "random" + , ccPreparedStatements = lookup "prepared_statements" params /= Just "false" + , ccMaxPreparedStatements = maybe 256 (\v -> readIntDef 256 v) (lookup "statement_cache_size" params) + } + +readTimeout :: Maybe ByteString -> NominalDiffTime +readTimeout Nothing = 10 +readTimeout (Just bs) = case BS8.readInt bs of + Just (n, _) | n > 0 -> fromIntegral n + _ -> 10 + +parseKeyValue :: ByteString -> Either String ConnConfig +parseKeyValue bs = + let pairs = map parsePair (BS8.words bs) + get key def = fromMaybe def (lookup key pairs) + portStr = get "port" "5432" + port = maybe 5432 fromIntegral (readPort portStr) + tlsMode = case get "sslmode" "disable" of + "require" -> TlsRequire + "prefer" -> TlsPrefer + "verify-ca" -> TlsVerifyCa + "verify-full" -> TlsVerifyFull + _ -> TlsDisable + getMaybe key = let v = get key "" in if BS8.null v then Nothing else Just (BS8.unpack v) + in Right + ConnConfig + { ccHost = BS8.unpack (get "host" "localhost") + , ccPort = port + , ccDatabase = get "dbname" (get "database" "") + , ccUser = get "user" "" + , ccPassword = get "password" "" + , ccTls = tlsMode + , ccAppName = get "application_name" "pg-wire" + , ccConnectTimeout = case BS8.readInt (get "connect_timeout" "10") of + Just (n, _) | n > 0 -> fromIntegral n + _ -> 10 + , ccQueryTimeout = case BS8.readInt (get "query_timeout" "0") of + Just (n, _) | n > 0 -> fromIntegral n + _ -> 0 + , ccSslCert = getMaybe "sslcert" + , ccSslKey = getMaybe "sslkey" + , ccSslRootCert = getMaybe "sslrootcert" + , ccKeepalives = get "keepalives" "1" /= "0" + , ccKeepalivesIdle = readIntDef 0 (get "keepalives_idle" "0") + , ccKeepalivesInterval = readIntDef 0 (get "keepalives_interval" "0") + , ccKeepalivesCount = readIntDef 0 (get "keepalives_count" "0") + , ccTargetSessionAttrs = parseSessionAttrs (get "target_session_attrs" "any") + , ccClientEncoding = get "client_encoding" "UTF8" + , ccOptions = get "options" "" + , ccLoadBalanceHosts = get "load_balance_hosts" "disable" == "random" + , ccPreparedStatements = get "prepared_statements" "true" /= "false" + , ccMaxPreparedStatements = readIntDef 256 (get "statement_cache_size" "256") + } + where + parsePair p = + let (k, v) = BS8.break (== '=') p + in (k, BS8.drop 1 v) + +readIntDef :: Int -> ByteString -> Int +readIntDef d bs = case BS8.readInt bs of + Just (n, _) -> n + Nothing -> d + +readPort :: ByteString -> Maybe Word16 +readPort bs + | BS8.null bs = Nothing + | BS8.all (\c -> isDigit c) bs = + let n = BS8.foldl' (\acc c -> acc * 10 + fromIntegral (fromEnum c - 48)) 0 bs :: Int + in if n > 0 && n <= 65535 then Just (fromIntegral n) else Nothing + | otherwise = Nothing + +parseQueryParams :: ByteString -> [(ByteString, ByteString)] +parseQueryParams bs + | BS8.null bs = [] + | otherwise = + [ (k, BS8.drop 1 v) + | param <- BS8.split '&' bs + , let (k, v) = BS8.break (== '=') param + , not (BS8.null k) + ] + +parseSessionAttrs :: ByteString -> TargetSessionAttrs +parseSessionAttrs "any" = SessionAny +parseSessionAttrs "read-write" = SessionReadWrite +parseSessionAttrs "read-only" = SessionReadOnly +parseSessionAttrs "primary" = SessionPrimary +parseSessionAttrs "standby" = SessionStandby +parseSessionAttrs "prefer-standby" = SessionPreferStandby +parseSessionAttrs _ = SessionAny + +-- | Parse a comma-separated host string into individual hosts. +-- @\"host1,host2,host3\"@ → @[\"host1\", \"host2\", \"host3\"]@ +parseHosts :: String -> [String] +parseHosts [] = ["localhost"] +parseHosts s = go s + where + go [] = [] + go str = let (h, rest) = break (== ',') str + in h : case rest of + [] -> [] + (_ : rs) -> go rs diff --git a/src/PgWire/Error.hs b/src/PgWire/Error.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Error.hs @@ -0,0 +1,44 @@ +-- | Error types for pg-wire runtime exceptions. +-- +-- All errors are thrown as 'PgWireError' via 'throwIO' and can be caught +-- with the standard @Control.Exception@ machinery. +module PgWire.Error + ( PgWireError (..) + , throwPgWire + ) where + +import Control.Exception (Exception, throwIO) +import Data.ByteString (ByteString) +import GHC.Generics (Generic) +import NoThunks.Class (NoThunks) +import PgWire.Protocol.Backend (PgError) + +-- | All runtime errors thrown by pg-wire. +data PgWireError + = ConnectionError ByteString + -- ^ Failed to establish a TCP\/TLS connection to the server. + | AuthError ByteString + -- ^ Authentication failed (wrong password, unsupported mechanism, etc.). + | ProtocolError ByteString + -- ^ Unexpected message in the wire protocol (indicates a driver bug or + -- incompatible server). + | QueryError PgError + -- ^ The server returned an error in response to a query. The 'PgError' + -- contains SQLSTATE, message, detail, hint, and position fields. + | DecodeError ByteString + -- ^ Failed to decode a result row (type mismatch, unexpected NULL, or + -- no rows returned to 'fetchOneOrThrow'). + | PoolTimeout + -- ^ Timed out waiting to acquire a connection from the pool. + | PoolClosed + -- ^ Attempted to use a pool after 'PgWire.Pool.closePool' was called. + | ConnectionDead + -- ^ The connection's reader thread has died (network failure, server crash). + deriving stock (Show, Eq, Generic) + +instance NoThunks PgWireError +instance Exception PgWireError + +-- | Throw a 'PgWireError' as an exception. +throwPgWire :: PgWireError -> IO a +throwPgWire = throwIO diff --git a/src/PgWire/Pool.hs b/src/PgWire/Pool.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Pool.hs @@ -0,0 +1,796 @@ +{-# LANGUAGE DataKinds #-} + +-- | Thread-safe connection pool for PostgreSQL. +-- +-- Manages a set of reusable 'Connection's with configurable pool size, +-- idle timeout, max lifetime, health checking, and lifecycle hooks. +-- Inspired by deadpool-postgres (Rust). +-- +-- @ +-- pool <- 'newPool' 'PgWire.Pool.Config.defaultPoolConfig' +-- { poolConnString = \"postgres:\/\/...\" +-- , poolSize = 10 +-- } +-- 'withResource' pool $ \\conn -> ... +-- 'closePool' pool +-- @ +module PgWire.Pool + ( Pool (..) + , PoolStats (..) + , newPool + , closePool + , drainPool + , withResource + , poolStats + , poolIsAlive + , resize + , retain + , setPostCreateHook + , setOnAcquireHook + , setPreReleaseHook + , withResourceTimeout + ) where + +import GHC.Generics (Generic) +import NoThunks.Class (NoThunks (..), OnlyCheckWhnfNamed (..), allNoThunks) +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async (Async, async, cancel, race) +import Control.Concurrent.STM +import Control.Exception (AsyncException, SomeException, catch, mask, onException, throwIO, try) +import Data.ByteString.Char8 qualified as BS8 +import Data.IORef +import Data.Int (Int32) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Sequence (Seq (..)) +import Data.Sequence qualified as Seq +import Data.Time (NominalDiffTime, UTCTime, addUTCTime, diffUTCTime, getCurrentTime) +import Data.HashPSQ qualified as PSQ +import PgWire.Async (AsyncWireConn (..)) +import PgWire.Connection (Connection (..), close, connectString, simpleQuery) +import PgWire.Error (PgWireError (..), throwPgWire) +import PgWire.Pool.Config (PoolConfig (..), QueueMode (..), RecyclingMethod (..)) +import PgWire.Pool.Observation (PoolEvent (ConnectionAcquired, ConnectionCreated, ConnectionDestroyed, ConnectionRecycled, ConnectionReleased, HealthCheckFailed, AcquireTimeout, ReaperSwept, WarmerCreated, PoolResized, PoolShutdown)) +import PgWire.TypeCache (TypeCache, newTypeCache) +import System.Random (randomRIO) + +-- | A thread-safe connection pool for PostgreSQL. +-- +-- Manages a bounded set of reusable 'Connection's with idle timeout, max lifetime +-- (with jitter to avoid thundering-herd reconnects), health checking, and lifecycle hooks. +-- The pool is safe to share across threads. Use 'newPool' to create and 'closePool' to shut down. +data Pool = Pool + { pConfig :: PoolConfig + , pIdle :: TVar (Seq PoolEntry) + , pActive :: TVar Int + , pWaiters :: TVar (Seq (TMVar (Either PgWireError Connection))) + , pClosed :: TVar Bool + , pReaper :: IORef (Async ()) + , pWarmer :: IORef (Maybe (Async ())) + , pEffectiveSize :: TVar Int + -- ^ Mutable pool size for runtime 'resize'. + , pTotalCreated :: TVar Int + , pTotalDestroyed :: TVar Int + , pTotalTimeouts :: TVar Int + , pConnMeta :: TVar (Map Int32 ConnMeta) + -- ^ Per-connection metadata keyed by backend PID. + -- Tracks creation time and jittered max-life for each connection. + , pOnCreate :: IORef (Connection -> IO ()) + , pOnAcquire :: IORef (Connection -> IO ()) + , pOnRelease :: IORef (Connection -> IO ()) + , pTypeCache :: TypeCache + -- ^ Pool-level cache for resolved PG type metadata (OID → TypeInfo). + -- Avoids redundant @pg_type@ round-trips across connections. + } + +-- | Per-connection metadata, stored by backend PID. +data ConnMeta = ConnMeta + { cmCreatedAt :: !UTCTime + , cmDeadline :: !UTCTime + -- ^ Creation time + jittered max-life. Connection is expired after this. + } + deriving stock (Generic) + +-- | 'UTCTime' is treated as opaque (its internal 'Day'/'DiffTime' are +-- lazy in the @time@ package). Bounded size, not a leak target. +instance NoThunks ConnMeta where + showTypeOf _ = "ConnMeta" + wNoThunks ctx (ConnMeta createdAt deadline) = allNoThunks + [ noThunks ctx (OnlyCheckWhnfNamed @"UTCTime" createdAt) + , noThunks ctx (OnlyCheckWhnfNamed @"UTCTime" deadline) + ] + +data PoolEntry = PoolEntry + { peConn :: Connection + , peCreatedAt :: UTCTime + , peLastUsed :: IORef UTCTime + } + +-- | Treats 'Connection', the @IORef UTCTime@, and 'UTCTime' itself as +-- opaque references (their internal state is not pool-bookkeeping and +-- 'UTCTime' has lazy fields in the @time@ package). +instance NoThunks PoolEntry where + showTypeOf _ = "PoolEntry" + wNoThunks ctx (PoolEntry conn createdAt lastUsed) = allNoThunks + [ noThunks ctx (OnlyCheckWhnfNamed @"Connection" conn) + , noThunks ctx (OnlyCheckWhnfNamed @"UTCTime" createdAt) + , noThunks ctx (OnlyCheckWhnfNamed @"IORef UTCTime" lastUsed) + ] + +-- | Snapshot of pool statistics at a point in time, obtained via 'poolStats'. +data PoolStats = PoolStats + { psIdle :: !Int + -- ^ Number of connections currently idle in the pool. + , psInUse :: !Int + -- ^ Number of connections currently checked out by callers. + , psWaiters :: !Int + -- ^ Number of threads blocked waiting to acquire a connection. + , psMaxSize :: !Int + -- ^ Current maximum pool size (may differ from initial config after 'resize'). + , psTotalCreated :: !Int + -- ^ Cumulative count of connections created since pool creation. + , psTotalDestroyed :: !Int + -- ^ Cumulative count of connections destroyed since pool creation. + , psTotalTimeouts :: !Int + -- ^ Cumulative count of acquire attempts that timed out. + } + deriving stock (Show, Eq, Generic) + +instance NoThunks PoolStats + +-- | Create a new connection pool from the given configuration. +-- +-- Starts background threads for reaping idle\/expired connections and (if +-- @poolMinIdle > 0@) warming the pool to maintain a minimum number of idle +-- connections. No connections are created eagerly; the first 'withResource' +-- call triggers the initial connection. +-- +-- Call 'closePool' when the pool is no longer needed to release all resources. +newPool :: PoolConfig -> IO Pool +newPool cfg = do + idle <- newTVarIO Seq.empty + active <- newTVarIO 0 + waiters <- newTVarIO Seq.empty + closed <- newTVarIO False + reaperRef <- newIORef (error "reaper not started") + warmerRef <- newIORef Nothing + effectiveSize <- newTVarIO (poolSize cfg) + totalCreated <- newTVarIO 0 + totalDestroyed <- newTVarIO 0 + totalTimeouts <- newTVarIO 0 + connMeta <- newTVarIO Map.empty + onCreate <- newIORef (\_ -> pure ()) + onAcquire <- newIORef (\_ -> pure ()) + onRelease <- newIORef (\_ -> pure ()) + typeCache <- newTypeCache + let pool = + Pool + { pConfig = cfg + , pIdle = idle + , pActive = active + , pWaiters = waiters + , pClosed = closed + , pReaper = reaperRef + , pWarmer = warmerRef + , pEffectiveSize = effectiveSize + , pTotalCreated = totalCreated + , pTotalDestroyed = totalDestroyed + , pTotalTimeouts = totalTimeouts + , pConnMeta = connMeta + , pOnCreate = onCreate + , pOnAcquire = onAcquire + , pOnRelease = onRelease + , pTypeCache = typeCache + } + reaper <- async (reaperThread pool) + writeIORef reaperRef reaper + -- Spawn warmer if poolMinIdle > 0 + when (poolMinIdle cfg > 0) $ do + w <- async (warmerThread pool) + writeIORef warmerRef (Just w) + pure pool + +-- | Close the pool, releasing all resources. +-- +-- Closes every idle connection, cancels background threads (reaper and warmer), +-- and wakes all blocked waiters with a 'PoolClosed' error. After this call, +-- any subsequent 'withResource' will throw 'PoolClosed'. Idempotent: calling +-- 'closePool' on an already-closed pool is a no-op for the connection teardown. +closePool :: Pool -> IO () +closePool pool = do + -- Atomically: mark closed, drain idle, drain waiters + (entries, blockedWaiters) <- atomically $ do + writeTVar (pClosed pool) True + idle <- readTVar (pIdle pool) + writeTVar (pIdle pool) Seq.empty + ws <- readTVar (pWaiters pool) + writeTVar (pWaiters pool) Seq.empty + pure (idle, ws) + -- Cancel background threads + reaper <- readIORef (pReaper pool) + cancel reaper + mWarmer <- readIORef (pWarmer pool) + mapM_ cancel mWarmer + -- Wake all blocked waiters with PoolClosed + mapM_ (\w -> atomically $ tryPutTMVar w (Left PoolClosed)) blockedWaiters + -- Close all idle connections (properly tracked) + mapM_ (destroyEntry pool "pool closed") entries + observe pool PoolShutdown + +-- | Gracefully drain the pool for rolling deployments. +-- +-- Marks the pool as closed (no new acquires), then polls until all +-- in-use connections are returned, then closes everything. Blocks +-- until the drain is complete or the timeout expires. +-- +-- @ +-- -- During rolling deployment: +-- drainPool pool 30 -- wait up to 30 seconds for in-flight queries +-- @ +drainPool :: Pool -> NominalDiffTime -> IO () +drainPool pool timeout = do + -- Mark as closed — new withResource calls will get PoolClosed + atomically $ writeTVar (pClosed pool) True + -- Wake blocked waiters immediately + ws <- atomically $ do + waiters <- readTVar (pWaiters pool) + writeTVar (pWaiters pool) Seq.empty + pure waiters + mapM_ (\w -> atomically $ tryPutTMVar w (Left PoolClosed)) ws + -- Poll until all active connections are returned or timeout + let timeoutMicros = round (timeout * 1000000) :: Int + pollInterval = 100000 -- 100ms + pollLoop !remaining + | remaining <= 0 = pure () -- timeout, force close + | otherwise = do + active <- readTVarIO (pActive pool) + idle <- Seq.length <$> readTVarIO (pIdle pool) + if active <= idle + then pure () -- all connections returned to idle + else do + threadDelay pollInterval + pollLoop (remaining - pollInterval) + pollLoop timeoutMicros + -- Now close everything (reuse closePool's cleanup logic) + entries <- atomically $ do + idle <- readTVar (pIdle pool) + writeTVar (pIdle pool) Seq.empty + pure idle + reaper <- readIORef (pReaper pool) + cancel reaper + mWarmer <- readIORef (pWarmer pool) + mapM_ cancel mWarmer + mapM_ (destroyEntry pool "drain") entries + observe pool PoolShutdown + +-- | Acquire a connection from the pool, run an action, and return the connection. +-- +-- The connection is guaranteed to be returned to the pool (or destroyed) even +-- if the action throws an exception. Uses 'mask' internally so the acquire and +-- release cannot be interrupted by async exceptions. +-- +-- @ +-- withResource pool $ \\conn -> +-- simpleQuery conn \"SELECT 1\" +-- @ +withResource :: Pool -> (Connection -> IO a) -> IO a +withResource pool action = mask $ \restore -> do + t0 <- getCurrentTime + conn <- acquire pool + t1 <- getCurrentTime + observe pool (ConnectionAcquired (diffUTCTime t1 t0)) + result <- restore (action conn) `onException` destroyConn pool conn "exception" + release pool conn + observe pool ConnectionReleased + pure result + +-- | Like 'withResource' but with a custom acquire timeout. +-- +-- Overrides 'poolAcquireTimeout' for this single acquisition. Useful when +-- certain operations can tolerate longer (or shorter) waits than the pool +-- default. +-- +-- @ +-- -- Wait up to 30 seconds for this particular query +-- withResourceTimeout pool 30 $ \\conn -> +-- simpleQuery conn \"SELECT expensive_function()\" +-- @ +withResourceTimeout :: Pool -> NominalDiffTime -> (Connection -> IO a) -> IO a +withResourceTimeout pool timeout action = mask $ \restore -> do + t0 <- getCurrentTime + conn <- acquireWithTimeout pool timeout + t1 <- getCurrentTime + observe pool (ConnectionAcquired (diffUTCTime t1 t0)) + result <- restore (action conn) `onException` destroyConn pool conn "exception" + release pool conn + observe pool ConnectionReleased + pure result + +-- | Get an atomic snapshot of the pool's current statistics. +-- +-- The returned 'PoolStats' is consistent (all fields read in a single STM +-- transaction) but represents a point-in-time view that may be stale by +-- the time you inspect it. +poolStats :: Pool -> IO PoolStats +poolStats pool = atomically $ do + idle <- Seq.length <$> readTVar (pIdle pool) + active <- readTVar (pActive pool) + waiters <- Seq.length <$> readTVar (pWaiters pool) + maxSz <- readTVar (pEffectiveSize pool) + created <- readTVar (pTotalCreated pool) + destroyed <- readTVar (pTotalDestroyed pool) + timeouts <- readTVar (pTotalTimeouts pool) + pure PoolStats + { psIdle = idle + , psInUse = active - idle + , psWaiters = waiters + , psMaxSize = maxSz + , psTotalCreated = created + , psTotalDestroyed = destroyed + , psTotalTimeouts = timeouts + } + +-- | Check if the pool is still open (not closed). +-- +-- Returns 'True' if the pool is alive and accepting connections, +-- 'False' if 'closePool' has been called. +poolIsAlive :: Pool -> IO Bool +poolIsAlive pool = not <$> readTVarIO (pClosed pool) + +-- | Resize the pool at runtime to a new maximum connection count. +-- +-- If shrinking, excess idle connections are closed immediately. In-use +-- connections are not interrupted and will finish naturally; they are simply +-- not returned to the pool once the count exceeds the new limit. If growing, +-- new connections are created on demand by subsequent 'withResource' calls. +resize :: Pool -> Int -> IO () +resize pool newSize = do + (oldSize, excess) <- atomically $ do + old <- readTVar (pEffectiveSize pool) + writeTVar (pEffectiveSize pool) newSize + idle <- readTVar (pIdle pool) + let idleCount = Seq.length idle + excessCount = max 0 (idleCount - newSize) + if excessCount > 0 + then do + let (keep, toDrop) = Seq.splitAt (idleCount - excessCount) idle + writeTVar (pIdle pool) keep + pure (old, seqToList toDrop) + else pure (old, []) + logPool pool "info" ("resized from " <> BS8.pack (show oldSize) <> " to " <> BS8.pack (show newSize)) + observe pool (PoolResized oldSize newSize) + mapM_ (destroyEntry pool "resize") excess + +-- | Filter idle connections, keeping only those that satisfy the predicate. +-- +-- Atomically drains the idle queue, applies the predicate to each connection +-- in IO, returns matching ones to the pool, and destroys the rest. Useful for +-- invalidating connections after a credential rotation or schema migration. +-- +-- @ +-- -- Drop all connections that were created before the migration +-- retain pool (\\conn -> isPostMigration conn) +-- @ +retain :: Pool -> (Connection -> IO Bool) -> IO () +retain pool predicate = do + entries <- atomically $ do + idle <- readTVar (pIdle pool) + writeTVar (pIdle pool) Seq.empty + pure (seqToList idle) + (kept, dropped) <- partitionM (\e -> predicate (peConn e)) entries + atomically $ modifyTVar' (pIdle pool) (Seq.fromList kept Seq.><) + mapM_ (destroyEntry pool "retain") dropped + +-- | Set a hook called immediately after a new connection is created. +-- +-- Useful for running session-level setup such as @SET search_path@ or +-- @SET statement_timeout@. If the hook throws, the exception is silently +-- caught and the connection is still placed in the pool. +setPostCreateHook :: Pool -> (Connection -> IO ()) -> IO () +setPostCreateHook pool hook = writeIORef (pOnCreate pool) $! hook + +-- | Set a hook called each time a connection is handed out to a caller via 'withResource'. +-- +-- Runs after any recycling\/health checks have passed. If the hook throws, the +-- exception is silently caught and the connection is still provided. +setOnAcquireHook :: Pool -> (Connection -> IO ()) -> IO () +setOnAcquireHook pool hook = writeIORef (pOnAcquire pool) $! hook + +-- | Set a hook called before a connection is returned to the idle pool. +-- +-- Runs after the caller's action completes and before the connection is placed +-- back in the idle queue or delivered to a waiting thread. Useful for cleanup +-- such as @RESET ALL@ or @DISCARD TEMP@. If the hook throws, the exception is +-- silently caught. +setPreReleaseHook :: Pool -> (Connection -> IO ()) -> IO () +setPreReleaseHook pool hook = writeIORef (pOnRelease pool) $! hook + +-- Internal ---------------------------------------------------------------- + +-- | What the single STM transaction decided we should do. +data AcquireAction + = GotIdle !PoolEntry + | CreateNew + | MustWait !(TMVar (Either PgWireError Connection)) + | PoolIsClosed + +-- | Atomically decide what to do: take idle, create new, or enqueue waiter. +-- This is the core fix for the acquire race condition — a single STM +-- transaction replaces the old two-step check. +acquireAction :: Pool -> STM AcquireAction +acquireAction pool = do + closed <- readTVar (pClosed pool) + if closed + then pure PoolIsClosed + else do + idle <- readTVar (pIdle pool) + case takeIdle (poolQueueMode (pConfig pool)) idle of + Just (entry, rest) -> do + writeTVar (pIdle pool) rest + pure (GotIdle entry) + Nothing -> do + active <- readTVar (pActive pool) + maxSz <- readTVar (pEffectiveSize pool) + if active < maxSz + then do + writeTVar (pActive pool) (active + 1) + pure CreateNew + else do + waiter <- newEmptyTMVar + modifyTVar' (pWaiters pool) (Seq.|> waiter) + pure (MustWait waiter) + +-- | Take from the appropriate end based on QueueMode. +takeIdle :: QueueMode -> Seq PoolEntry -> Maybe (PoolEntry, Seq PoolEntry) +takeIdle QueueLIFO s = case Seq.viewr s of + Seq.EmptyR -> Nothing + rest Seq.:> e -> Just (e, rest) +takeIdle QueueFIFO s = case Seq.viewl s of + Seq.EmptyL -> Nothing + e Seq.:< rest -> Just (e, rest) + +acquire :: Pool -> IO Connection +acquire pool = acquireWithTimeout pool (poolAcquireTimeout (pConfig pool)) + +acquireWithTimeout :: Pool -> NominalDiffTime -> IO Connection +acquireWithTimeout pool timeout = do + act <- atomically (acquireAction pool) + case act of + PoolIsClosed -> throwPgWire PoolClosed + GotIdle entry -> tryRecycle pool entry + CreateNew -> createConnection pool + MustWait waiter -> waitForConnectionWith pool waiter timeout + +-- | Try to recycle an idle connection. If it's expired or unhealthy, +-- destroy it and retry. +tryRecycle :: Pool -> PoolEntry -> IO Connection +tryRecycle pool entry = do + now <- getCurrentTime + expired <- isExpiredIO pool entry now + if expired + then do + destroyEntry pool "expired" entry + acquire pool + else do + ok <- recycleCheck pool entry now + if ok + then do + logPool pool "debug" "recycled connection" + observe pool ConnectionRecycled + writeIORef (peLastUsed entry) now + runHook (pOnAcquire pool) (peConn entry) + pure (peConn entry) + else do + observe pool HealthCheckFailed + destroyEntry pool "unhealthy" entry + acquire pool + +-- | Recycling check based on the configured method. +recycleCheck :: Pool -> PoolEntry -> UTCTime -> IO Bool +recycleCheck pool entry now = case poolRecyclingMethod (pConfig pool) of + RecycleFast -> do + readTVarIO (awcAlive (connAsync (peConn entry))) + RecycleVerified -> do + alive <- readTVarIO (awcAlive (connAsync (peConn entry))) + if not alive + then pure False + else do + lastUsed <- readIORef (peLastUsed entry) + if diffUTCTime now lastUsed <= poolHealthCheckAge (pConfig pool) + then pure True + else checkHealthWithTimeout (poolAcquireTimeout (pConfig pool)) (peConn entry) + RecycleClean -> do + alive <- readTVarIO (awcAlive (connAsync (peConn entry))) + if not alive + then pure False + else do + let timeoutSecs = poolAcquireTimeout (pConfig pool) + timeoutMicros = round (timeoutSecs * 1000000) :: Int + result <- race (threadDelay timeoutMicros) + (try @SomeException (simpleQuery (peConn entry) "DISCARD ALL")) + case result of + Right (Right _) -> do + -- DISCARD ALL deallocates all prepared statements on the server, + -- so the client-side cache must be cleared to stay in sync. + writeIORef (connStmtCache (peConn entry)) PSQ.empty + pure True + _ -> pure False + +createConnection :: Pool -> IO Connection +createConnection pool = do + conn <- connectWithRetry (poolConnectionRetries (pConfig pool)) 100000 + `onException` atomically (modifyTVar' (pActive pool) (subtract 1)) + now <- getCurrentTime + deadline <- jitteredDeadline (pConfig pool) now + atomically $ do + modifyTVar' (pTotalCreated pool) (+ 1) + modifyTVar' (pConnMeta pool) (Map.insert (connBackendPid conn) (ConnMeta now deadline)) + logPool pool "debug" "created connection" + observe pool ConnectionCreated + runHook (pOnCreate pool) conn + runHook (pOnAcquire pool) conn + pure conn + where + -- Exponential backoff: 100ms, 200ms, 400ms, ... + connectWithRetry :: Int -> Int -> IO Connection + connectWithRetry 0 _ = connectString (poolConnString (pConfig pool)) + connectWithRetry !retriesLeft !delayMicros = do + result <- try @SomeException (connectString (poolConnString (pConfig pool))) + case result of + Right c -> pure c + Left _ -> do + logPool pool "warn" ("connection failed, retrying in " <> BS8.pack (show (delayMicros `div` 1000)) <> "ms") + threadDelay delayMicros + connectWithRetry (retriesLeft - 1) (delayMicros * 2) + +-- | Compute a jittered deadline for a connection. +-- deadline = now + maxLife + uniform(-jitter, +jitter) +jitteredDeadline :: PoolConfig -> UTCTime -> IO UTCTime +jitteredDeadline cfg now = do + let jitter = poolMaxLifeJitter cfg + offset <- if jitter > 0 + then do + let jitterMicros = round (jitter * 1000000) :: Int + r <- randomRIO (negate jitterMicros, jitterMicros) + pure (fromIntegral r / 1000000 :: NominalDiffTime) + else pure 0 + pure (addUTCTime (poolMaxLife cfg + offset) now) + +waitForConnectionWith :: Pool -> TMVar (Either PgWireError Connection) -> NominalDiffTime -> IO Connection +waitForConnectionWith pool waiter timeout = do + let timeoutMicros = round (timeout * 1000000) :: Int + result <- race + (threadDelay timeoutMicros) + (atomically $ takeTMVar waiter) + case result of + Left () -> do + atomically $ modifyTVar' (pTotalTimeouts pool) (+ 1) + logPool pool "warn" "acquire timeout" + observe pool AcquireTimeout + throwPgWire PoolTimeout + Right (Left err) -> throwPgWire err + Right (Right conn) -> do + runHook (pOnAcquire pool) conn + pure conn + +release :: Pool -> Connection -> IO () +release pool conn = do + runHook (pOnRelease pool) conn + now <- getCurrentTime + -- Try to deliver to a waiter first. Use tryPutTMVar and loop + -- to handle dead waiters (waiter died between enqueue and delivery). + delivered <- deliverToWaiter pool conn + if delivered + then pure () + else do + -- Look up creation time from the tracking map (fixes maxLife bug) + meta <- Map.lookup (connBackendPid conn) <$> readTVarIO (pConnMeta pool) + let createdAt = maybe now cmCreatedAt meta + lastUsed <- newIORef now + let !entry = PoolEntry conn createdAt lastUsed + atomically $ modifyTVar' (pIdle pool) (Seq.|> entry) + +-- | Try to deliver a connection to a waiting thread. Loops through waiters, +-- skipping any that have died (timed out). Returns True if delivered. +deliverToWaiter :: Pool -> Connection -> IO Bool +deliverToWaiter pool conn = atomically $ do + ws <- readTVar (pWaiters pool) + case Seq.viewl ws of + Seq.EmptyL -> pure False + w Seq.:< rest -> do + delivered <- tryPutTMVar w (Right conn) + if delivered + then do + writeTVar (pWaiters pool) rest + pure True + else do + -- Waiter died (timeout). Skip it and try next. + writeTVar (pWaiters pool) rest + -- Retry with remaining waiters (re-enter STM) + deliverLoop pool rest conn + +-- | STM loop to deliver to next available waiter. +deliverLoop :: Pool -> Seq (TMVar (Either PgWireError Connection)) -> Connection -> STM Bool +deliverLoop _ Empty _ = pure False +deliverLoop pool (w :<| rest) conn = do + delivered <- tryPutTMVar w (Right conn) + if delivered + then do + writeTVar (pWaiters pool) rest + pure True + else do + writeTVar (pWaiters pool) rest + deliverLoop pool rest conn + +destroyConn :: Pool -> Connection -> BS8.ByteString -> IO () +destroyConn pool conn reason = do + safeClose conn + atomically $ do + modifyTVar' (pActive pool) (subtract 1) + modifyTVar' (pTotalDestroyed pool) (+ 1) + modifyTVar' (pConnMeta pool) (Map.delete (connBackendPid conn)) + logPool pool "debug" ("destroyed connection: " <> reason) + observe pool (ConnectionDestroyed reason) + +destroyEntry :: Pool -> BS8.ByteString -> PoolEntry -> IO () +destroyEntry pool reason entry = destroyConn pool (peConn entry) reason + +-- | Check if a connection has exceeded its (jittered) max-life. +isExpiredIO :: Pool -> PoolEntry -> UTCTime -> IO Bool +isExpiredIO pool entry now = do + meta <- Map.lookup (connBackendPid (peConn entry)) <$> readTVarIO (pConnMeta pool) + pure $ case meta of + Just cm -> now >= cmDeadline cm + Nothing -> diffUTCTime now (peCreatedAt entry) > poolMaxLife (pConfig pool) + +safeClose :: Connection -> IO () +safeClose conn = close conn `catch` \(_ :: SomeException) -> pure () + +-- | Lightweight health check: send an empty query and see if we get a response. +-- Times out after the given duration to prevent a hung backend from blocking acquire. +checkHealthWithTimeout :: NominalDiffTime -> Connection -> IO Bool +checkHealthWithTimeout timeout conn = do + let timeoutMicros = round (timeout * 1000000) :: Int + result <- race (threadDelay timeoutMicros) (try @SomeException (simpleQuery conn "")) + pure $ case result of + Left () -> False -- timed out + Right (Right _) -> True + Right (Left _) -> False + +-- | Run a lifecycle hook, catching and ignoring exceptions so a hook +-- failure doesn't break the pool. +runHook :: IORef (Connection -> IO ()) -> Connection -> IO () +runHook ref conn = do + hook <- readIORef ref + hook conn `catch` \(_ :: SomeException) -> pure () + +-- | Log a pool event using the configured logger. +logPool :: Pool -> BS8.ByteString -> BS8.ByteString -> IO () +logPool pool = poolLogger (pConfig pool) + +-- | Emit a structured observation event, catching and ignoring exceptions. +observe :: Pool -> PoolEvent -> IO () +observe pool event = + poolObserver (pConfig pool) event `catch` \(_ :: SomeException) -> pure () + +-- | Background thread that periodically reaps idle/expired connections. +-- +-- Atomically drains the idle queue, filters in IO (reading IORefs), +-- then atomically merges kept entries back. This avoids the race where +-- acquire/release modify the idle queue between read and write. +reaperThread :: Pool -> IO () +reaperThread pool = go + where + intervalMicros = round (poolReaperInterval (pConfig pool) * 1000000) :: Int + + go = do + threadDelay intervalMicros + reap + `catch` \(e :: AsyncException) -> throwIO e + `catch` \(_ :: SomeException) -> + logPool pool "warn" "reaper: exception in sweep cycle (continuing)" + go + + reap = do + now <- getCurrentTime + -- Atomically drain the idle queue + entries <- atomically $ do + idle <- readTVar (pIdle pool) + writeTVar (pIdle pool) Seq.empty + pure (seqToList idle) + (keep, toReap) <- partitionM (shouldKeep now) entries + -- Merge kept entries back (prepend; entries added by release + -- during filtering are already at the tail) + atomically $ modifyTVar' (pIdle pool) (Seq.fromList keep Seq.><) + let reapCount = length toReap + mapM_ (destroyEntry pool "reaped") toReap + when (reapCount > 0) $ do + logPool pool "info" ("reaper swept " <> BS8.pack (show reapCount) <> " connections") + observe pool (ReaperSwept reapCount) + + shouldKeep now entry = do + expired <- isExpiredIO pool entry now + if expired + then pure False + else do + lastUsed <- readIORef (peLastUsed entry) + pure (diffUTCTime now lastUsed <= poolIdleTime (pConfig pool)) + +-- | Background thread that maintains 'poolMinIdle' connections. +warmerThread :: Pool -> IO () +warmerThread pool = go + where + intervalMicros = round (poolReaperInterval (pConfig pool) * 1000000) :: Int + + go = do + threadDelay intervalMicros + warm + `catch` \(e :: AsyncException) -> throwIO e + `catch` \(_ :: SomeException) -> + logPool pool "warn" "warmer: exception in warm cycle (continuing)" + go + + warm = do + deficit <- atomically $ do + closed <- readTVar (pClosed pool) + if closed + then pure 0 + else do + idle <- Seq.length <$> readTVar (pIdle pool) + active <- readTVar (pActive pool) + maxSz <- readTVar (pEffectiveSize pool) + let minIdle = poolMinIdle (pConfig pool) + canCreate = maxSz - active + need = max 0 (minIdle - idle) + pure (min need canCreate) + created <- warmN 0 deficit + when (created > 0) $ do + logPool pool "info" ("warmer created " <> BS8.pack (show created) <> " connections") + observe pool (WarmerCreated created) + + warmN :: Int -> Int -> IO Int + warmN !acc 0 = pure acc + warmN !acc n = do + result <- try @SomeException $ do + -- Reserve a slot + reserved <- atomically $ do + active <- readTVar (pActive pool) + maxSz <- readTVar (pEffectiveSize pool) + if active < maxSz + then do + writeTVar (pActive pool) (active + 1) + pure True + else pure False + if reserved + then do + conn <- connectString (poolConnString (pConfig pool)) + `onException` atomically (modifyTVar' (pActive pool) (subtract 1)) + now <- getCurrentTime + deadline <- jitteredDeadline (pConfig pool) now + atomically $ do + modifyTVar' (pTotalCreated pool) (+ 1) + modifyTVar' (pConnMeta pool) (Map.insert (connBackendPid conn) (ConnMeta now deadline)) + runHook (pOnCreate pool) conn + lastUsed <- newIORef now + let !entry = PoolEntry conn now lastUsed + atomically $ modifyTVar' (pIdle pool) (Seq.|> entry) + pure True + else pure False + case result of + Left _ -> pure acc -- Connection failed, stop warming this round + Right True -> warmN (acc + 1) (n - 1) + Right False -> pure acc + +when :: Bool -> IO () -> IO () +when True f = f +when False _ = pure () + +seqToList :: Seq a -> [a] +seqToList = foldr (:) [] + +partitionM :: (a -> IO Bool) -> [a] -> IO ([a], [a]) +partitionM _ [] = pure ([], []) +partitionM p (x : xs) = do + b <- p x + (!ys, !zs) <- partitionM p xs + pure (if b then (x : ys, zs) else (ys, x : zs)) diff --git a/src/PgWire/Pool/Config.hs b/src/PgWire/Pool/Config.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Pool/Config.hs @@ -0,0 +1,141 @@ +-- | Configuration types for the connection pool. +-- +-- Use 'defaultPoolConfig' as a starting point and override fields: +-- +-- @ +-- cfg = 'defaultPoolConfig' +-- { poolConnString = \"postgres:\/\/user:pass\@localhost\/mydb\" +-- , poolSize = 20 +-- , poolRecyclingMethod = 'RecycleVerified' +-- } +-- @ +module PgWire.Pool.Config + ( PoolConfig (..) + , defaultPoolConfig + , RecyclingMethod (..) + , QueueMode (..) + , PoolLogger + , nullPoolLogger + ) where + +import Data.ByteString (ByteString) +import Data.Time (NominalDiffTime) +import PgWire.Pool.Observation (PoolObserver, nullObserver) + +-- | How to validate a connection before handing it to the caller. +data RecyclingMethod + = RecycleFast + -- ^ Check @awcAlive@ TVar only — free, no I/O. + | RecycleVerified + -- ^ Send an empty query if the connection has been idle longer than + -- 'poolHealthCheckAge'. Falls back to 'RecycleFast' for recently used + -- connections. + | RecycleClean + -- ^ Send @DISCARD ALL@ before every reuse (resets all session state). + deriving stock (Show, Eq) + +-- | Which end of the idle queue to draw from. +data QueueMode + = QueueFIFO + -- ^ Oldest idle connection first — spreads load across backends. + | QueueLIFO + -- ^ Newest idle connection first — keeps fewer total connections alive. + deriving stock (Show, Eq) + +-- | A pool logging callback. The pool calls this with a severity string +-- (@\"debug\"@, @\"info\"@, @\"warn\"@) and a message. +type PoolLogger = ByteString -> ByteString -> IO () + +-- | Logger that discards all events. +nullPoolLogger :: PoolLogger +nullPoolLogger _ _ = pure () + +-- | Configuration for a connection pool. Use 'defaultPoolConfig' and +-- override the fields you need: +-- +-- @ +-- pool <- 'PgWire.Pool.newPool' defaultPoolConfig +-- { poolConnString = \"postgres:\/\/user:pass\@localhost:5432\/mydb\" +-- , poolSize = 20 +-- } +-- @ +data PoolConfig = PoolConfig + { poolConnString :: ByteString + -- ^ PostgreSQL connection string. Required. + , poolSize :: !Int + -- ^ Maximum number of connections. Default: 10. + , poolIdleTime :: !NominalDiffTime + -- ^ Close idle connections after this duration. Default: 600s (10 min). + , poolMaxLife :: !NominalDiffTime + -- ^ Close connections after this total lifetime. Default: 3600s (1 hour). + , poolAcquireTimeout :: !NominalDiffTime + -- ^ Maximum time to wait for a connection. Throws 'PgWire.Error.PoolTimeout' + -- if exceeded. Default: 10s. + , poolMinIdle :: !Int + -- ^ Minimum idle connections to maintain. A background warmer thread + -- creates connections to fill the deficit. Default: 0 (no warming). + , poolReaperInterval :: !NominalDiffTime + -- ^ How often the reaper sweeps idle/expired connections. Default: 30s. + , poolRecyclingMethod :: !RecyclingMethod + -- ^ How to validate connections before reuse. Default: 'RecycleFast'. + , poolQueueMode :: !QueueMode + -- ^ Which end of the idle queue to draw from. Default: 'QueueLIFO'. + , poolHealthCheckAge :: !NominalDiffTime + -- ^ Only verify (empty query) if a connection has been idle longer than + -- this. Only relevant when 'poolRecyclingMethod' is 'RecycleVerified'. + -- Default: 5s. + , poolMaxLifeJitter :: !NominalDiffTime + -- ^ Randomize 'poolMaxLife' ± this amount to avoid thundering herd + -- expiration. Default: 60s. + , poolConnectionRetries :: !Int + -- ^ Number of retries when creating a connection fails. Retries use + -- exponential backoff (100ms, 200ms, 400ms, ...). Default: 3. + , poolLogger :: PoolLogger + -- ^ Logging callback. Default: 'nullPoolLogger' (discards all events). + , poolObserver :: PoolObserver + -- ^ Structured event callback for metrics and monitoring. + -- Called for every significant pool state transition. + -- Default: 'nullObserver' (discards all events). + -- See "PgWire.Pool.Observation" for the event types. + } + +-- Manual Show instance that omits the function field. +instance Show PoolConfig where + show cfg = unwords + [ "PoolConfig" + , "{poolConnString=" <> show (poolConnString cfg) + , ",poolSize=" <> show (poolSize cfg) + , ",poolIdleTime=" <> show (poolIdleTime cfg) + , ",poolMaxLife=" <> show (poolMaxLife cfg) + , ",poolAcquireTimeout=" <> show (poolAcquireTimeout cfg) + , ",poolMinIdle=" <> show (poolMinIdle cfg) + , ",poolReaperInterval=" <> show (poolReaperInterval cfg) + , ",poolRecyclingMethod=" <> show (poolRecyclingMethod cfg) + , ",poolQueueMode=" <> show (poolQueueMode cfg) + , ",poolHealthCheckAge=" <> show (poolHealthCheckAge cfg) + , ",poolMaxLifeJitter=" <> show (poolMaxLifeJitter cfg) + , ",poolConnectionRetries=" <> show (poolConnectionRetries cfg) + , ",poolLogger=" + , ",poolObserver=}" + ] + +-- | Sensible defaults: 10 connections, 10-minute idle timeout, 1-hour max life, +-- 10-second acquire timeout, 'RecycleFast', 'QueueLIFO', no warming, no logging. +defaultPoolConfig :: PoolConfig +defaultPoolConfig = + PoolConfig + { poolConnString = "" + , poolSize = 10 + , poolIdleTime = 600 + , poolMaxLife = 3600 + , poolAcquireTimeout = 10 + , poolMinIdle = 0 + , poolReaperInterval = 30 + , poolRecyclingMethod = RecycleFast + , poolQueueMode = QueueLIFO + , poolHealthCheckAge = 5 + , poolMaxLifeJitter = 60 + , poolConnectionRetries = 3 + , poolLogger = nullPoolLogger + , poolObserver = nullObserver + } diff --git a/src/PgWire/Pool/Observation.hs b/src/PgWire/Pool/Observation.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Pool/Observation.hs @@ -0,0 +1,60 @@ +-- | Structured pool lifecycle events for observability. +-- +-- Set an observation handler via 'PgWire.Pool.Config.poolObserver' to +-- receive events for every significant pool state transition. Use these +-- for metrics, structured logging, alerting, or debugging. +-- +-- @ +-- pool <- newPool defaultPoolConfig +-- { poolObserver = \\event -> do +-- case event of +-- ConnectionCreated -> incrementCounter \"pool.created\" +-- ConnectionDestroyed reason -> logInfo $ \"destroyed: \" <> reason +-- AcquireTimeout -> incrementCounter \"pool.timeouts\" +-- _ -> pure () +-- } +-- @ +module PgWire.Pool.Observation + ( PoolEvent (..) + , PoolObserver + , nullObserver + ) where + +import Data.ByteString (ByteString) +import Data.Time (NominalDiffTime) + +-- | A structured event emitted by the connection pool. +data PoolEvent + = ConnectionCreated + -- ^ A new connection was established. + | ConnectionDestroyed !ByteString + -- ^ A connection was closed. The 'ByteString' is the reason + -- (e.g., @\"expired\"@, @\"unhealthy\"@, @\"reaped\"@, @\"pool closed\"@, + -- @\"resize\"@, @\"exception\"@, @\"retain\"@). + | ConnectionAcquired !NominalDiffTime + -- ^ A connection was handed to a caller. The 'NominalDiffTime' is + -- how long the caller waited (0 if a connection was immediately available). + | ConnectionReleased + -- ^ A connection was returned to the pool. + | ConnectionRecycled + -- ^ An idle connection passed its recycling check and was reused. + | HealthCheckFailed + -- ^ A connection failed its health check and was destroyed. + | AcquireTimeout + -- ^ A caller timed out waiting for a connection. + | ReaperSwept !Int + -- ^ The background reaper closed N idle/expired connections. + | WarmerCreated !Int + -- ^ The background warmer created N connections to fill min-idle. + | PoolResized !Int !Int + -- ^ Pool was resized from old size to new size. + | PoolShutdown + -- ^ 'closePool' was called. + deriving stock (Show, Eq) + +-- | Callback for pool events. Set in 'PgWire.Pool.Config.PoolConfig'. +type PoolObserver = PoolEvent -> IO () + +-- | Observer that discards all events. +nullObserver :: PoolObserver +nullObserver _ = pure () diff --git a/src/PgWire/Protocol/Backend.hs b/src/PgWire/Protocol/Backend.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Protocol/Backend.hs @@ -0,0 +1,121 @@ +-- | Decoded backend (server-to-client) message types from the +-- PostgreSQL v3 wire protocol. +-- +-- __Audience:__ exposed for downstream library authors that want to +-- decode raw protocol responses or build alternative pipeline +-- abstractions. Most application code only sees these types +-- indirectly via "PgWire.Connection" results and "PgWire.Error". +module PgWire.Protocol.Backend + ( BackendMsg (..) + , FieldInfo (..) + , PgError (..) + , PgNotice (..) + , TxStatus (..) + , CommandTag (..) + , AuthType (..) + ) where + +import Data.ByteString (ByteString) +import Data.Int (Int16, Int32, Int64) +import Data.Vector (Vector) +import Data.Word (Word32) +import GHC.Generics (Generic) +import NoThunks.Class (NoThunks) + +-- | Transaction status indicator. +data TxStatus + = TxIdle + | TxInTransaction + | TxFailed + deriving stock (Show, Eq) + +-- | Parsed command completion tag. +data CommandTag + = SelectTag Int64 + | InsertTag Int64 + | UpdateTag Int64 + | DeleteTag Int64 + | OtherTag ByteString + deriving stock (Show, Eq) + +-- | Authentication sub-types. +data AuthType + = AuthOk + | AuthCleartextPassword + | AuthMD5Password ByteString -- 4-byte salt + | AuthSASL [ByteString] -- mechanism names + | AuthSASLContinue ByteString -- server data + | AuthSASLFinal ByteString -- server signature + deriving stock (Show, Eq) + +-- | Column metadata from RowDescription. +data FieldInfo = FieldInfo + { fiName :: ByteString + , fiTableOid :: {-# UNPACK #-} !Word32 + , fiColumnNum :: {-# UNPACK #-} !Int16 + , fiTypeOid :: {-# UNPACK #-} !Word32 + , fiTypeSize :: {-# UNPACK #-} !Int16 + , fiTypeMod :: {-# UNPACK #-} !Int32 + , fiFormatCode :: {-# UNPACK #-} !Int16 + } + deriving stock (Show, Eq) + +-- | Structured error from ErrorResponse. +-- See https://www.postgresql.org/docs/current/protocol-error-fields.html +data PgError = PgError + { pgSeverity :: ByteString + , pgCode :: ByteString + , pgMessage :: ByteString + , pgDetail :: Maybe ByteString + , pgHint :: Maybe ByteString + , pgPosition :: !(Maybe Int32) + , pgInternalPosition :: !(Maybe Int32) + , pgInternalQuery :: Maybe ByteString + , pgWhere :: Maybe ByteString + , pgSchema :: Maybe ByteString + , pgTable :: Maybe ByteString + , pgColumn :: Maybe ByteString + , pgDataType :: Maybe ByteString + , pgConstraint :: Maybe ByteString + , pgFile :: Maybe ByteString + , pgLine :: !(Maybe Int32) + , pgRoutine :: Maybe ByteString + } + deriving stock (Show, Eq, Generic) + +instance NoThunks PgError + +-- | Structured notice from NoticeResponse. +newtype PgNotice = PgNotice {unPgNotice :: PgError} + deriving stock (Show, Eq) + +-- | Server-to-client (backend) wire protocol messages. +data BackendMsg + = Authentication AuthType + | ParameterStatus ByteString ByteString + | BackendKeyData Int32 Int32 + | ReadyForQuery TxStatus + | ParseComplete + | BindComplete + | CloseComplete + | NoData + | EmptyQueryResponse + | ParameterDescription (Vector Word32) + | RowDescription (Vector FieldInfo) + | DataRow (Vector (Maybe ByteString)) + | CommandComplete CommandTag + | ErrorResponse PgError + | NoticeResponse PgNotice + | NotificationResponse Int32 ByteString ByteString + | -- | CopyInResponse: server is ready to receive COPY data + CopyInResponse Int8Format (Vector Int16) + | -- | CopyOutResponse: server will send COPY data + CopyOutResponse Int8Format (Vector Int16) + | -- | CopyData: a chunk of COPY data from the server + CopyDataMsg ByteString + | -- | CopyDone: end of COPY data from server + CopyDoneMsg + deriving stock (Show, Eq) + +-- | COPY format code (0 = text, 1 = binary) +type Int8Format = Int16 diff --git a/src/PgWire/Protocol/Builders.hs b/src/PgWire/Protocol/Builders.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Protocol/Builders.hs @@ -0,0 +1,177 @@ +-- | Encoders for frontend (client-to-server) messages. +-- +-- __Audience:__ exposed for downstream library authors that need to +-- emit raw protocol bytes (e.g. custom pipelining or proxies). Most +-- application code drives these encoders indirectly through +-- "PgWire.Connection". +module PgWire.Protocol.Builders + ( buildFrontendMsg + , buildFrontendMsgsConcat + , buildStartup + ) where + +import Data.Bits (shiftL) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Builder (Builder) +import Data.ByteString.Builder qualified as B +import Data.ByteString.Lazy qualified as LBS +import Data.Int (Int16, Int32) +import Data.Vector (Vector) +import Data.Vector qualified as V +import PgWire.Protocol.Frontend + +-- | Build a complete frontend message as a strict 'ByteString'. +buildFrontendMsg :: FrontendMsg -> ByteString +buildFrontendMsg = LBS.toStrict . B.toLazyByteString . encodeFrontendMsg +{-# INLINE buildFrontendMsg #-} + +-- | Fuse multiple frontend messages into a single strict 'ByteString'. +-- Uses one 'Builder' pass and one 'toStrict' call, eliminating N-1 +-- intermediate allocations compared to mapping 'buildFrontendMsg'. +buildFrontendMsgsConcat :: [FrontendMsg] -> ByteString +buildFrontendMsgsConcat = LBS.toStrict . B.toLazyByteString . foldMap encodeFrontendMsg +{-# INLINE buildFrontendMsgsConcat #-} + +-- | Build the startup message (special: no tag byte). +buildStartup :: StartupParams -> ByteString +buildStartup = LBS.toStrict . B.toLazyByteString . encodeStartup + +-- Encoding ---------------------------------------------------------------- +-- +-- Each message is [tag :: Word8] [length :: Int32] [payload]. +-- The length field includes itself (4 bytes) but not the tag byte. +-- +-- We pre-compute the payload size from the message fields, then write +-- tag + length + payload in a single Builder pass. This avoids the +-- previous approach of materializing the payload to a ByteString just +-- to compute its length (which caused a double-copy). + +encodeFrontendMsg :: FrontendMsg -> Builder +encodeFrontendMsg = \case + Startup params -> encodeStartup params + Parse name sql oids -> + let !sz = cstringSize name + cstringSize sql + 2 + V.length oids * 4 + in tag 'P' sz <> cstring name <> cstring sql + <> B.int16BE (fromIntegral (V.length oids) :: Int16) + <> V.foldl' (\b oid -> b <> B.word32BE oid) mempty oids + Bind portal stmt pfmts vals rfmts -> + let !sz = cstringSize portal + cstringSize stmt + + formatCodesSize pfmts + + 2 + paramValuesSize vals + + formatCodesSize rfmts + in tag 'B' sz <> cstring portal <> cstring stmt + <> encodeFormatCodes pfmts + <> B.int16BE (fromIntegral (V.length vals) :: Int16) + <> V.foldl' (\b mv -> b <> encodeParamValue mv) mempty vals + <> encodeFormatCodes rfmts + Describe target name -> + let !sz = 1 + cstringSize name + in tag 'D' sz <> B.char8 (targetChar target) <> cstring name + Execute portal maxRows -> + let !sz = cstringSize portal + 4 + in tag 'E' sz <> cstring portal <> B.int32BE maxRows + Close target name -> + let !sz = 1 + cstringSize name + in tag 'C' sz <> B.char8 (targetChar target) <> cstring name + Sync -> tag 'S' 0 + Flush -> tag 'H' 0 + Terminate -> tag 'X' 0 + Query sql -> + let !sz = cstringSize sql + in tag 'Q' sz <> cstring sql + PasswordMessage pw -> + let !sz = cstringSize pw + in tag 'p' sz <> cstring pw + SASLInitialResponse mech clientFirst -> + let !sz = cstringSize mech + 4 + BS.length clientFirst + in tag 'p' sz <> cstring mech + <> B.int32BE (fromIntegral (BS.length clientFirst)) + <> B.byteString clientFirst + SASLResponse msg -> + let !sz = BS.length msg + in tag 'p' sz <> B.byteString msg + CopyData dat -> + let !sz = BS.length dat + in tag 'd' sz <> B.byteString dat + CopyDone -> tag 'c' 0 + CopyFail msg -> + let !sz = cstringSize msg + in tag 'f' sz <> cstring msg +{-# INLINE encodeFrontendMsg #-} + +-- | Write tag byte + length (payload size + 4 for the length field itself). +tag :: Char -> Int -> Builder +tag c payloadSize = + B.char8 c <> B.int32BE (fromIntegral (payloadSize + 4) :: Int32) +{-# INLINE tag #-} + +-- Size computation helpers ------------------------------------------------ + +-- | Size of a NUL-terminated C string: data bytes + 1 NUL byte. +cstringSize :: ByteString -> Int +cstringSize bs = BS.length bs + 1 +{-# INLINE cstringSize #-} + +-- | Size of format codes: 2 (count) + 2 per code. +formatCodesSize :: Vector FormatCode -> Int +formatCodesSize fmts = 2 + V.length fmts * 2 +{-# INLINE formatCodesSize #-} + +-- | Size of parameter values: sum of (4 + data) per non-NULL, 4 per NULL. +paramValuesSize :: Vector (Maybe ByteString) -> Int +paramValuesSize = V.foldl' (\acc mv -> acc + paramValueSize mv) 0 +{-# INLINE paramValuesSize #-} + +paramValueSize :: Maybe ByteString -> Int +paramValueSize Nothing = 4 +paramValueSize (Just bs) = 4 + BS.length bs +{-# INLINE paramValueSize #-} + +-- Payload encoding helpers ------------------------------------------------ + +encodeFormatCodes :: Vector FormatCode -> Builder +encodeFormatCodes fmts = + B.int16BE (fromIntegral (V.length fmts) :: Int16) + <> V.foldl' (\b fc -> b <> B.int16BE (formatCodeToInt16 fc)) mempty fmts +{-# INLINE encodeFormatCodes #-} + +formatCodeToInt16 :: FormatCode -> Int16 +formatCodeToInt16 TextFormat = 0 +formatCodeToInt16 BinaryFormat = 1 +{-# INLINE formatCodeToInt16 #-} + +encodeParamValue :: Maybe ByteString -> Builder +encodeParamValue Nothing = B.int32BE (-1) +encodeParamValue (Just bs) = + B.int32BE (fromIntegral (BS.length bs)) + <> B.byteString bs +{-# INLINE encodeParamValue #-} + +targetChar :: DescribeTarget -> Char +targetChar DescribeStatement = 'S' +targetChar DescribePortal = 'P' +{-# INLINE targetChar #-} + +-- | NUL-terminated C string. +cstring :: ByteString -> Builder +cstring bs = B.byteString bs <> B.word8 0 +{-# INLINE cstring #-} + +-- Startup message --------------------------------------------------------- +-- Startup has a different format: [length :: Int32] [protocol] [params] [NUL] +-- No tag byte. We still need to materialize to compute the length here, +-- but startup is called once per connection, not on the hot path. + +encodeStartup :: StartupParams -> Builder +encodeStartup StartupParams {..} = + let params = + cstring "user" <> cstring spUser + <> cstring "database" <> cstring spDatabase + <> (if BS.null spAppName then mempty else cstring "application_name" <> cstring spAppName) + <> mconcat [cstring k <> cstring v | (k, v) <- spExtraParams] + <> B.word8 0 -- terminator + paramsBs = LBS.toStrict (B.toLazyByteString params) + len = fromIntegral (4 + 4 + BS.length paramsBs) :: Int32 + protocolVersion = (3 :: Int32) `shiftL` 16 -- 3.0 + in B.int32BE len <> B.int32BE protocolVersion <> B.byteString paramsBs diff --git a/src/PgWire/Protocol/Frontend.hs b/src/PgWire/Protocol/Frontend.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Protocol/Frontend.hs @@ -0,0 +1,94 @@ +-- | Frontend (client-to-server) message types from the PostgreSQL v3 +-- wire protocol. +-- +-- __Audience:__ exposed for downstream library authors that need to +-- assemble protocol messages directly. Most application code uses +-- "PgWire.Connection" instead. +module PgWire.Protocol.Frontend + ( FrontendMsg (..) + , StartupParams (..) + , DescribeTarget (..) + , FormatCode (..) + ) where + +import Data.ByteString (ByteString) +import Data.Int (Int32) +import Data.Vector (Vector) +import Data.Word (Word32) + +-- | Format codes for parameter/result encoding. +data FormatCode = TextFormat | BinaryFormat + deriving stock (Show, Eq) + +-- | Target for a Describe message. +data DescribeTarget = DescribeStatement | DescribePortal + deriving stock (Show, Eq) + +-- | Parameters for the startup message. +data StartupParams = StartupParams + { spUser :: ByteString + , spDatabase :: ByteString + , spAppName :: ByteString + , spExtraParams :: [(ByteString, ByteString)] + } + deriving stock (Show) + +-- | Client-to-server (frontend) wire protocol messages. +data FrontendMsg + = -- | Initial startup (no tag byte) + Startup StartupParams + | -- | Parse a prepared statement + Parse + ByteString + -- ^ Statement name (empty = unnamed) + ByteString + -- ^ SQL query + (Vector Word32) + -- ^ Parameter type OIDs (0 = let server infer) + | -- | Bind parameters to a prepared statement + Bind + ByteString + -- ^ Portal name (empty = unnamed) + ByteString + -- ^ Statement name + (Vector FormatCode) + -- ^ Parameter format codes + (Vector (Maybe ByteString)) + -- ^ Parameter values (Nothing = NULL) + (Vector FormatCode) + -- ^ Result format codes + | -- | Describe a prepared statement or portal + Describe DescribeTarget ByteString + | -- | Execute a portal + Execute + ByteString + -- ^ Portal name + Int32 + -- ^ Max rows (0 = unlimited) + | -- | Close a prepared statement or portal + Close DescribeTarget ByteString + | -- | Sync — end of an extended query sequence + Sync + | -- | Flush — request server to send pending output + Flush + | -- | Terminate the connection + Terminate + | -- | Simple query protocol + Query ByteString + | -- | Password message (cleartext or MD5) + PasswordMessage ByteString + | -- | SASL initial response + SASLInitialResponse + ByteString + -- ^ Mechanism name + ByteString + -- ^ Client first message + | -- | SASL response (subsequent messages) + SASLResponse ByteString + | -- | CopyData — a chunk of COPY data + CopyData ByteString + | -- | CopyDone — end of COPY data from client + CopyDone + | -- | CopyFail — client aborts COPY + CopyFail ByteString + deriving stock (Show) diff --git a/src/PgWire/Protocol/Oid.hs b/src/PgWire/Protocol/Oid.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Protocol/Oid.hs @@ -0,0 +1,143 @@ +-- | PostgreSQL type OIDs and convenience constants. +-- +-- Used when writing custom 'PgWire.Binary.Types.PgEncode' / +-- 'PgWire.Binary.Types.PgDecode' instances and when inspecting +-- raw 'PgWire.Protocol.Backend.FieldInfo' descriptions. This is +-- part of the public extensibility surface, not an internal module. +module PgWire.Protocol.Oid + ( Oid (..) + , oidBool + , oidBytea + , oidInt8 + , oidInt2 + , oidInt4 + , oidText + , oidJson + , oidFloat4 + , oidFloat8 + , oidVarchar + , oidDate + , oidTime + , oidTimetz + , oidTimestamp + , oidTimestamptz + , oidNumeric + , oidUuid + , oidJsonb + , oidInterval + , oidInet + , oidCidr + , oidMacAddr + , oidMacAddr8 + -- * Array OIDs + , oidBoolArray + , oidByteaArray + , oidInt2Array + , oidInt4Array + , oidInt8Array + , oidTextArray + , oidVarcharArray + , oidFloat4Array + , oidFloat8Array + , oidTimestampArray + , oidTimestamptzArray + , oidDateArray + , oidTimeArray + , oidUuidArray + , oidJsonArray + , oidJsonbArray + , isArrayOid + , arrayElementOid + ) where + +import Data.Word (Word32) + +-- | A PostgreSQL type OID. +newtype Oid = Oid {unOid :: Word32} + deriving stock (Show, Eq, Ord) + +oidBool, oidBytea, oidInt8, oidInt2, oidInt4 :: Oid +oidBool = Oid 16 +oidBytea = Oid 17 +oidInt8 = Oid 20 +oidInt2 = Oid 21 +oidInt4 = Oid 23 + +oidText, oidJson, oidFloat4, oidFloat8 :: Oid +oidText = Oid 25 +oidJson = Oid 114 +oidFloat4 = Oid 700 +oidFloat8 = Oid 701 + +oidVarchar, oidDate, oidTime, oidTimetz, oidTimestamp, oidTimestamptz :: Oid +oidVarchar = Oid 1043 +oidDate = Oid 1082 +oidTime = Oid 1083 +oidTimetz = Oid 1266 +oidTimestamp = Oid 1114 +oidTimestamptz = Oid 1184 + +oidNumeric, oidUuid, oidJsonb, oidInterval, oidInet, oidCidr, oidMacAddr, oidMacAddr8 :: Oid +oidNumeric = Oid 1700 +oidUuid = Oid 2950 +oidJsonb = Oid 3802 +oidInterval = Oid 1186 +oidInet = Oid 869 +oidCidr = Oid 650 +oidMacAddr = Oid 829 +oidMacAddr8 = Oid 774 + +-- Array OIDs ---------------------------------------------------------------- + +oidBoolArray, oidByteaArray, oidInt2Array, oidInt4Array, oidInt8Array :: Oid +oidBoolArray = Oid 1000 +oidByteaArray = Oid 1001 +oidInt2Array = Oid 1005 +oidInt4Array = Oid 1007 +oidInt8Array = Oid 1016 + +oidTextArray, oidVarcharArray :: Oid +oidTextArray = Oid 1009 +oidVarcharArray = Oid 1015 + +oidFloat4Array, oidFloat8Array :: Oid +oidFloat4Array = Oid 1021 +oidFloat8Array = Oid 1022 + +oidTimestampArray, oidTimestamptzArray, oidDateArray, oidTimeArray :: Oid +oidTimestampArray = Oid 1115 +oidTimestamptzArray = Oid 1185 +oidDateArray = Oid 1182 +oidTimeArray = Oid 1183 + +oidUuidArray, oidJsonArray, oidJsonbArray :: Oid +oidUuidArray = Oid 2951 +oidJsonArray = Oid 199 +oidJsonbArray = Oid 3807 + +-- | Check whether an OID is a known array type. +isArrayOid :: Oid -> Bool +isArrayOid oid = case arrayElementOid oid of + Just _ -> True + Nothing -> False + +-- | For an array OID, return the element OID. +arrayElementOid :: Oid -> Maybe Oid +arrayElementOid (Oid o) = case o of + 1000 -> Just oidBool + 1001 -> Just oidBytea + 1005 -> Just oidInt2 + 1007 -> Just oidInt4 + 1009 -> Just oidText + 1015 -> Just oidVarchar + 1016 -> Just oidInt8 + 1021 -> Just oidFloat4 + 1022 -> Just oidFloat8 + 1115 -> Just oidTimestamp + 1182 -> Just oidDate + 1183 -> Just oidTime + 1185 -> Just oidTimestamptz + 2951 -> Just oidUuid + 199 -> Just oidJson + 3807 -> Just oidJsonb + _ -> Nothing diff --git a/src/PgWire/Protocol/Parsers.hs b/src/PgWire/Protocol/Parsers.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Protocol/Parsers.hs @@ -0,0 +1,355 @@ +-- | Parsers for backend (server-to-client) protocol messages. +-- +-- __Audience:__ exposed for downstream library authors that decode +-- raw protocol bytes themselves. Most application code consumes +-- decoded results via "PgWire.Connection". +module PgWire.Protocol.Parsers + ( parseBackendMsg + , parseCommandTag + ) where + +import Data.Bits (shiftL, (.|.)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Unsafe qualified as BU +import Data.Int (Int16, Int32, Int64) +import Data.Vector (Vector) +import Data.Vector qualified as V +import Data.Vector.Mutable qualified as MV +import Data.Word (Word8, Word32) +import PgWire.Protocol.Backend + +-- | Parse a complete backend message from tag + payload bytes. +-- The caller must have already read the tag byte and the 4-byte length, +-- and pass the tag and the remaining payload (length - 4 bytes). +parseBackendMsg :: Word8 -> ByteString -> Either String BackendMsg +parseBackendMsg tag payload = case tag of + 82 {- R -} -> parseAuth payload + 75 {- K -} -> parseBackendKeyData payload + 83 {- S -} -> parseParameterStatus payload + 90 {- Z -} -> parseReadyForQuery payload + 49 {- 1 -} -> Right ParseComplete + 50 {- 2 -} -> Right BindComplete + 51 {- 3 -} -> Right CloseComplete + 67 {- C -} -> parseCommandComplete payload + 68 {- D -} -> parseDataRow payload + 69 {- E -} -> ErrorResponse <$> parseErrorFields payload + 71 {- G -} -> parseCopyInResponse payload + 72 {- H -} -> parseCopyOutResponse payload + 73 {- I -} -> Right EmptyQueryResponse + 99 {- c -} -> Right CopyDoneMsg + 100 {- d -} -> Right (CopyDataMsg payload) + 78 {- N -} -> NoticeResponse . PgNotice <$> parseErrorFields payload + 84 {- T -} -> parseRowDescription payload + 110 {- n -} -> Right NoData + 116 {- t -} -> parseParameterDescription payload + 65 {- A -} -> parseNotification payload + _ -> Left $ "Unknown backend message tag: " <> show tag + +-- Authentication ---------------------------------------------------------- + +parseAuth :: ByteString -> Either String BackendMsg +parseAuth bs = do + authType <- getInt32 bs 0 + case authType of + 0 -> Right (Authentication AuthOk) + 3 -> Right (Authentication AuthCleartextPassword) + 5 -> do + salt <- getBytes bs 4 4 + Right (Authentication (AuthMD5Password salt)) + 10 -> do + -- SASL: list of mechanism names, each NUL-terminated, final NUL + let mechData = BS.drop 4 bs + mechs = parseCStringList mechData + Right (Authentication (AuthSASL mechs)) + 11 -> Right (Authentication (AuthSASLContinue (BS.drop 4 bs))) + 12 -> Right (Authentication (AuthSASLFinal (BS.drop 4 bs))) + _ -> Left $ "Unknown auth type: " <> show authType + +-- ParameterStatus --------------------------------------------------------- + +parseParameterStatus :: ByteString -> Either String BackendMsg +parseParameterStatus bs = + let (name, rest) = BS.break (== 0) bs + value = BS.takeWhile (/= 0) (BS.drop 1 rest) + in Right (ParameterStatus name value) + +-- BackendKeyData ---------------------------------------------------------- + +parseBackendKeyData :: ByteString -> Either String BackendMsg +parseBackendKeyData bs = do + pid <- getInt32 bs 0 + key <- getInt32 bs 4 + Right (BackendKeyData pid key) + +-- ReadyForQuery ----------------------------------------------------------- + +parseReadyForQuery :: ByteString -> Either String BackendMsg +parseReadyForQuery bs + | BS.length bs < 1 = Left "ReadyForQuery: empty payload" + | otherwise = case BS.index bs 0 of + 73 {- I -} -> Right (ReadyForQuery TxIdle) + 84 {- T -} -> Right (ReadyForQuery TxInTransaction) + 69 {- E -} -> Right (ReadyForQuery TxFailed) + c -> Left $ "ReadyForQuery: unknown status " <> show c + +-- CommandComplete --------------------------------------------------------- + +parseCommandComplete :: ByteString -> Either String BackendMsg +parseCommandComplete bs = + let tag = BS.takeWhile (/= 0) bs + in Right (CommandComplete (parseCommandTag tag)) + +parseCommandTag :: ByteString -> CommandTag +parseCommandTag bs + | "SELECT " `BS.isPrefixOf` bs = SelectTag (readInt64 (BS.drop 7 bs)) + | "INSERT " `BS.isPrefixOf` bs = + -- INSERT oid count — we want the count (after the second space) + let afterInsert = BS.drop 7 bs + afterOid = BS.drop 1 (BS.dropWhile (/= 32) afterInsert) + in InsertTag (readInt64 afterOid) + | "UPDATE " `BS.isPrefixOf` bs = UpdateTag (readInt64 (BS.drop 7 bs)) + | "DELETE " `BS.isPrefixOf` bs = DeleteTag (readInt64 (BS.drop 7 bs)) + | otherwise = OtherTag bs + +readInt64 :: ByteString -> Int64 +readInt64 bs = + BS.foldl' (\acc w -> acc * 10 + fromIntegral (w - 48)) 0 (BS.takeWhile (\w -> w >= 48 && w <= 57) bs) + +-- RowDescription ---------------------------------------------------------- + +parseRowDescription :: ByteString -> Either String BackendMsg +parseRowDescription bs = do + nFields <- getInt16 bs 0 + (fields, _) <- parseFields (fromIntegral nFields) (BS.drop 2 bs) + Right (RowDescription (V.fromList fields)) + +parseFields :: Int -> ByteString -> Either String ([FieldInfo], ByteString) +parseFields 0 rest = Right ([], rest) +parseFields n bs = do + let (name, afterName) = BS.break (== 0) bs + rest = BS.drop 1 afterName -- skip NUL + tableOid <- getWord32 rest 0 + colNum <- getInt16 rest 4 + typeOid <- getWord32 rest 6 + typeSize <- getInt16 rest 10 + typeMod <- getInt32 rest 12 + fmtCode <- getInt16 rest 16 + let field = + FieldInfo + { fiName = name + , fiTableOid = tableOid + , fiColumnNum = colNum + , fiTypeOid = typeOid + , fiTypeSize = typeSize + , fiTypeMod = typeMod + , fiFormatCode = fmtCode + } + (fields, rest') <- parseFields (n - 1) (BS.drop 18 rest) + Right (field : fields, rest') + +-- DataRow ----------------------------------------------------------------- + +-- | Parse a DataRow message directly into a 'Vector', skipping the +-- intermediate list. Column value slices share the underlying buffer +-- via 'BS.take'/'BS.drop' (zero-copy). +parseDataRow :: ByteString -> Either String BackendMsg +parseDataRow bs = do + nCols <- getInt16 bs 0 + let !n = fromIntegral nCols + vals <- parseColsDirect n bs 2 + Right (DataRow vals) + +parseColsDirect :: Int -> ByteString -> Int -> Either String (Vector (Maybe ByteString)) +parseColsDirect n bs startOff + | BS.length bs < startOff = Left "DataRow: buffer too short for column data" + | otherwise = + let !bsLen = BS.length bs + -- First pass: validate all lengths fit within the buffer. + validate !i !off + | i >= n = Right off + | off + 4 > bsLen = Left "DataRow: truncated column length" + | otherwise = + let !len = getInt32Unsafe bs off + in if len == -1 + then validate (i + 1) (off + 4) + else + let !end = off + 4 + fromIntegral len + in if end > bsLen + then Left "DataRow: truncated column data" + else validate (i + 1) end + in case validate 0 startOff of + Left err -> Left err + Right _ -> Right $ V.create $ do + mv <- MV.new n + let go !i !off + | i >= n = pure () + | otherwise = + let !len = getInt32Unsafe bs off + in if len == -1 + then do + MV.write mv i Nothing + go (i + 1) (off + 4) + else do + let !dataLen = fromIntegral len + !val = BU.unsafeTake dataLen (BU.unsafeDrop (off + 4) bs) + MV.write mv i (Just val) + go (i + 1) (off + 4 + dataLen) + go 0 startOff + pure mv + where + -- Unchecked Int32 decode — bounds validated by the validation pass. + getInt32Unsafe :: ByteString -> Int -> Int32 + getInt32Unsafe b off = + let !b0 = fromIntegral (BU.unsafeIndex b off) :: Int32 + !b1 = fromIntegral (BU.unsafeIndex b (off + 1)) :: Int32 + !b2 = fromIntegral (BU.unsafeIndex b (off + 2)) :: Int32 + !b3 = fromIntegral (BU.unsafeIndex b (off + 3)) :: Int32 + in (b0 `shiftL` 24) .|. (b1 `shiftL` 16) .|. (b2 `shiftL` 8) .|. b3 + {-# INLINE getInt32Unsafe #-} + +-- ParameterDescription ---------------------------------------------------- + +parseParameterDescription :: ByteString -> Either String BackendMsg +parseParameterDescription bs = do + nParams <- getInt16 bs 0 + oids <- parseOids (fromIntegral nParams) (BS.drop 2 bs) + Right (ParameterDescription (V.fromList oids)) + +parseOids :: Int -> ByteString -> Either String [Word32] +parseOids 0 _ = Right [] +parseOids n bs = do + oid <- getWord32 bs 0 + oids <- parseOids (n - 1) (BS.drop 4 bs) + Right (oid : oids) + +-- ErrorResponse / NoticeResponse ------------------------------------------ + +parseErrorFields :: ByteString -> Either String PgError +parseErrorFields = go emptyPgError + where + go !err bs + | BS.null bs = Right err + | BS.index bs 0 == 0 = Right err + | otherwise = + let fieldType = BS.index bs 0 + (value, rest) = BS.break (== 0) (BS.drop 1 bs) + err' = case fieldType of + 83 {- S -} -> err {pgSeverity = value} + 67 {- C -} -> err {pgCode = value} + 77 {- M -} -> err {pgMessage = value} + 68 {- D -} -> err {pgDetail = Just value} + 72 {- H -} -> err {pgHint = Just value} + 80 {- P -} -> err {pgPosition = Just (fromIntegral (readInt64 value))} + 112 {- p -} -> err {pgInternalPosition = Just (fromIntegral (readInt64 value))} + 113 {- q -} -> err {pgInternalQuery = Just value} + 87 {- W -} -> err {pgWhere = Just value} + 115 {- s -} -> err {pgSchema = Just value} + 116 {- t -} -> err {pgTable = Just value} + 99 {- c -} -> err {pgColumn = Just value} + 100 {- d -} -> err {pgDataType = Just value} + 110 {- n -} -> err {pgConstraint = Just value} + 70 {- F -} -> err {pgFile = Just value} + 76 {- L -} -> err {pgLine = Just (fromIntegral (readInt64 value))} + 82 {- R -} -> err {pgRoutine = Just value} + _ -> err + in go err' (BS.drop 1 rest) -- skip NUL after value + + emptyPgError = + PgError + { pgSeverity = "" + , pgCode = "" + , pgMessage = "" + , pgDetail = Nothing + , pgHint = Nothing + , pgPosition = Nothing + , pgInternalPosition = Nothing + , pgInternalQuery = Nothing + , pgWhere = Nothing + , pgSchema = Nothing + , pgTable = Nothing + , pgColumn = Nothing + , pgDataType = Nothing + , pgConstraint = Nothing + , pgFile = Nothing + , pgLine = Nothing + , pgRoutine = Nothing + } + +-- CopyInResponse / CopyOutResponse ---------------------------------------- + +parseCopyInResponse :: ByteString -> Either String BackendMsg +parseCopyInResponse bs = do + fmt <- getInt8 bs 0 + nCols <- getInt16 bs 1 + fmts <- mapM (\i -> getInt16 bs (3 + i * 2)) [0 .. fromIntegral nCols - 1] + Right (CopyInResponse (fromIntegral fmt) (V.fromList fmts)) + +parseCopyOutResponse :: ByteString -> Either String BackendMsg +parseCopyOutResponse bs = do + fmt <- getInt8 bs 0 + nCols <- getInt16 bs 1 + fmts <- mapM (\i -> getInt16 bs (3 + i * 2)) [0 .. fromIntegral nCols - 1] + Right (CopyOutResponse (fromIntegral fmt) (V.fromList fmts)) + +getInt8 :: ByteString -> Int -> Either String Word8 +getInt8 bs off + | BS.length bs < off + 1 = Left "getInt8: buffer too short" + | otherwise = Right (BS.index bs off) + +-- Notification ------------------------------------------------------------ + +parseNotification :: ByteString -> Either String BackendMsg +parseNotification bs = do + pid <- getInt32 bs 0 + let (channel, rest1) = BS.break (== 0) (BS.drop 4 bs) + payload = BS.takeWhile (/= 0) (BS.drop 1 rest1) + Right (NotificationResponse pid channel payload) + +-- Helpers ----------------------------------------------------------------- + +parseCStringList :: ByteString -> [ByteString] +parseCStringList bs + | BS.null bs = [] + | BS.index bs 0 == 0 = [] + | otherwise = + let (s, rest) = BS.break (== 0) bs + in s : parseCStringList (BS.drop 1 rest) + +getInt16 :: ByteString -> Int -> Either String Int16 +getInt16 bs off + | BS.length bs < off + 2 = Left "getInt16: buffer too short" + | otherwise = + -- After bounds check, use unsafeIndex for the hot path + let hi = fromIntegral (BU.unsafeIndex bs off) :: Int16 + lo = fromIntegral (BU.unsafeIndex bs (off + 1)) :: Int16 + in Right (hi `shiftL` 8 .|. lo) +{-# INLINE getInt16 #-} + +getInt32 :: ByteString -> Int -> Either String Int32 +getInt32 bs off + | BS.length bs < off + 4 = Left "getInt32: buffer too short" + | otherwise = + let b0 = fromIntegral (BU.unsafeIndex bs off) :: Int32 + b1 = fromIntegral (BU.unsafeIndex bs (off + 1)) :: Int32 + b2 = fromIntegral (BU.unsafeIndex bs (off + 2)) :: Int32 + b3 = fromIntegral (BU.unsafeIndex bs (off + 3)) :: Int32 + in Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3) +{-# INLINE getInt32 #-} + +getWord32 :: ByteString -> Int -> Either String Word32 +getWord32 bs off + | BS.length bs < off + 4 = Left "getWord32: buffer too short" + | otherwise = + let b0 = fromIntegral (BU.unsafeIndex bs off) :: Word32 + b1 = fromIntegral (BU.unsafeIndex bs (off + 1)) :: Word32 + b2 = fromIntegral (BU.unsafeIndex bs (off + 2)) :: Word32 + b3 = fromIntegral (BU.unsafeIndex bs (off + 3)) :: Word32 + in Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3) +{-# INLINE getWord32 #-} + +getBytes :: ByteString -> Int -> Int -> Either String ByteString +getBytes bs off len + | BS.length bs < off + len = Left "getBytes: buffer too short" + | otherwise = Right (BS.take len (BS.drop off bs)) +{-# INLINE getBytes #-} diff --git a/src/PgWire/TypeCache.hs b/src/PgWire/TypeCache.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/TypeCache.hs @@ -0,0 +1,60 @@ +-- | Pool-level cache for resolved PostgreSQL type metadata. +-- +-- When multiple connections encounter the same custom OID (e.g., a PG +-- enum), this cache avoids redundant @pg_type@ round-trips. The cache +-- is shared across all connections in a pool via a 'TVar'. +-- +-- The cache grows without bound but is bounded in practice by the number +-- of distinct custom types in the database (typically tens, not thousands). +-- Applications that dynamically create and drop types at high volume may +-- want to monitor 'cachedTypeCount' for unexpected growth. +module PgWire.TypeCache + ( TypeCache + , TypeInfo (..) + , newTypeCache + , lookupTypeInfo + , cacheTypeInfo + , cachedTypeCount + ) where + +import Control.Concurrent.STM (TVar, atomically, modifyTVar', newTVarIO, readTVarIO) +import Data.ByteString (ByteString) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Word (Word32) + +-- | Metadata about a PG type, cached at the pool level. +data TypeInfo = TypeInfo + { tiName :: !ByteString + -- ^ Type name from @pg_type.typname@. + , tiCategory :: !Char + -- ^ Type category: @\'E\'@ enum, @\'C\'@ composite, @\'R\'@ range, + -- @\'D\'@ domain, @\'B\'@ base, @\'P\'@ pseudo, etc. + , tiArrayOid :: !Word32 + -- ^ OID of the array form of this type (0 if none). + , tiBaseOid :: !Word32 + -- ^ For domains: the underlying type OID. 0 otherwise. + } + deriving stock (Show, Eq) + +-- | A thread-safe cache mapping OIDs to their resolved type info. +newtype TypeCache = TypeCache (TVar (Map Word32 TypeInfo)) + +-- | Create an empty type cache. +newTypeCache :: IO TypeCache +newTypeCache = TypeCache <$> newTVarIO Map.empty + +-- | Look up cached type info for an OID. +lookupTypeInfo :: TypeCache -> Word32 -> IO (Maybe TypeInfo) +lookupTypeInfo (TypeCache tv) oid = do + m <- readTVarIO tv + pure (Map.lookup oid m) + +-- | Cache type info for an OID. +cacheTypeInfo :: TypeCache -> Word32 -> TypeInfo -> IO () +cacheTypeInfo (TypeCache tv) oid info = + atomically $ modifyTVar' tv (Map.insert oid info) + +-- | Return the number of cached entries (for diagnostics). +cachedTypeCount :: TypeCache -> IO Int +cachedTypeCount (TypeCache tv) = Map.size <$> readTVarIO tv diff --git a/src/PgWire/Wire.hs b/src/PgWire/Wire.hs new file mode 100644 --- /dev/null +++ b/src/PgWire/Wire.hs @@ -0,0 +1,390 @@ +-- | Low-level PostgreSQL wire protocol I\/O. +-- +-- 'WireConn' abstracts over a TCP or TLS connection, providing +-- send\/recv operations for protocol messages. This module handles +-- message framing (tag + length + payload), buffering, and +-- TLS upgrades via the SSLRequest subprotocol. +-- +-- Most users should use 'PgWire.Connection' instead of this module. +module PgWire.Wire + ( -- * Wire connection + WireConn (..) + -- * Connecting + , connectTcp + , connectTcpTimeout + -- * TLS + , TlsConfig (..) + , upgradeTls + -- * Tracing + , TraceDirection (..) + -- * Sending messages + , sendFrontendMsg + , sendFrontendMsgs + , sendRawBytes + -- * Receiving messages + , recvBackendMsg + ) where + +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Char8 qualified as BS8 +import Data.ByteString.Unsafe qualified as BU +import Data.ByteString.Builder qualified as B +import Data.ByteString.Lazy qualified as LBS +import Data.IORef +import Data.Maybe (fromMaybe) +import Data.Time (NominalDiffTime) +import PgWire.Error (PgWireError (..), throwPgWire) +import System.Timeout (timeout) +import PgWire.Protocol.Backend (BackendMsg) +import PgWire.Protocol.Builders (buildFrontendMsg, buildFrontendMsgsConcat) +import PgWire.Protocol.Frontend (FrontendMsg) +import PgWire.Protocol.Parsers (parseBackendMsg) +import Network.Socket (Socket) +import Network.Socket qualified as NS +import Network.Socket.ByteString qualified as NSB +import Network.Socket (setSocketOption, SocketOption(..)) +import Network.TLS qualified as TLS +import Network.TLS (ClientParams (..), ClientHooks (..), Supported (..), Shared (..), Credentials (..)) +import Data.X509.CertificateStore (readCertificateStore) +import System.X509 (getSystemCertificateStore) + +-- | Abstraction over a Postgres wire connection (plain TCP or TLS). +-- | Direction for protocol trace callback. +data TraceDirection = TraceSend | TraceRecv + deriving stock (Show, Eq) + +data WireConn = WireConn + { wcSend :: ByteString -> IO () + , wcSendMany :: [ByteString] -> IO () + -- ^ Vectored I/O: send multiple chunks in one syscall (writev). + , wcRecv :: Int -> IO ByteString + , wcClose :: IO () + , wcBuffer :: IORef ByteString + , wcTrace :: IORef (Maybe (TraceDirection -> ByteString -> IO ())) + } + +-- | Connect via TCP to the given host and port (no timeout). +connectTcp :: NS.HostName -> NS.PortNumber -> IO WireConn +connectTcp = connectTcpTimeout 0 + +-- | Connect via TCP with a timeout in seconds (0 = no timeout). +connectTcpTimeout :: NominalDiffTime -> NS.HostName -> NS.PortNumber -> IO WireConn +connectTcpTimeout timeoutSecs host port = do + let hints = NS.defaultHints {NS.addrSocketType = NS.Stream} + addrs <- NS.getAddrInfo (Just hints) (Just host) (Just (show port)) + case addrs of + [] -> throwPgWire (ConnectionError "No address found") + (addr : _) -> do + sock <- NS.socket (NS.addrFamily addr) NS.Stream NS.defaultProtocol + let doConnect = NS.connect sock (NS.addrAddress addr) + if timeoutSecs > 0 + then do + let micros = round (timeoutSecs * 1000000) :: Int + result <- timeout micros doConnect + case result of + Nothing -> do + NS.close sock + throwPgWire (ConnectionError "Connect timed out") + Just () -> pure () + else doConnect + -- Disable Nagle's algorithm for lower latency on small messages + setSocketOption sock NoDelay 1 + mkWireConn sock + +mkWireConn :: Socket -> IO WireConn +mkWireConn sock = do + buf <- newIORef BS.empty + traceRef <- newIORef Nothing + pure + WireConn + { wcSend = sendAll sock + , wcSendMany = NSB.sendMany sock + , wcRecv = recvExact sock buf + , wcClose = NS.close sock + , wcBuffer = buf + , wcTrace = traceRef + } + +-- | TLS configuration for upgradeTls. +data TlsConfig = TlsConfig + { tlsHostname :: NS.HostName + , tlsVerify :: Bool + -- ^ Whether to verify the server certificate (verify-ca / verify-full) + , tlsVerifyHostname :: Bool + -- ^ Whether to verify hostname matches cert (verify-full only) + , tlsClientCert :: Maybe FilePath + -- ^ Path to client certificate file + , tlsClientKey :: Maybe FilePath + -- ^ Path to client private key file + , tlsCaCert :: Maybe FilePath + -- ^ Path to CA certificate file (Nothing = use system store) + } + +-- | Upgrade a TCP WireConn to TLS. Sends the SSLRequest message, +-- checks the server response, and performs the TLS handshake. +-- Returns a new WireConn that sends/receives over TLS. +upgradeTls :: WireConn -> TlsConfig -> IO WireConn +upgradeTls wc tlsCfg = do + let sslRequest = LBS.toStrict . B.toLazyByteString $ + B.int32BE 8 <> B.int32BE 80877103 + wcSend wc sslRequest + + resp <- wcRecv wc 1 + case BS.index resp 0 of + 83 {- S -} -> do + -- Load CA certificate store + systemStore <- getSystemCertificateStore + caStore <- case tlsCaCert tlsCfg of + Nothing -> pure systemStore + Just "system" -> pure systemStore + Just path -> do + mStore <- readCertificateStore path + pure (fromMaybe systemStore mStore) + + -- Load client certificate + key if provided + clientCred <- case (tlsClientCert tlsCfg, tlsClientKey tlsCfg) of + (Just certPath, Just keyPath) -> do + result <- TLS.credentialLoadX509 certPath keyPath + case result of + Right cred -> pure (Credentials [cred]) + Left _err -> do + -- Non-fatal: warn but continue without client cert + pure (Credentials []) + _ -> pure (Credentials []) + + let hostname = tlsHostname tlsCfg + baseParams = TLS.defaultParamsClient hostname "" + -- Configure certificate validation based on mode + hooks = (clientHooks baseParams) + { onServerCertificate = + if tlsVerify tlsCfg + then onServerCertificate (clientHooks baseParams) -- default validation + else \_ _ _ _ -> pure [] -- skip validation (sslmode=require) + } + clientParams = baseParams + { clientSupported = (clientSupported baseParams) + { supportedVersions = [TLS.TLS13, TLS.TLS12] + } + , clientShared = (clientShared baseParams) + { sharedCAStore = caStore + , sharedCredentials = clientCred + } + , clientHooks = hooks + } + ctx <- TLS.contextNew (wireBackend wc) clientParams + TLS.handshake ctx + mkTlsWireConn ctx (wcBuffer wc) + + 78 {- N -} -> + throwPgWire (ConnectionError "Server does not support SSL") + + other -> + throwPgWire (ConnectionError ("Unexpected SSL response: " <> BS8.pack (show other))) + +-- | Wrap a WireConn as a TLS backend (for the handshake, which needs +-- raw socket I/O). +wireBackend :: WireConn -> TLS.Backend +wireBackend wc = TLS.Backend + { TLS.backendFlush = pure () + , TLS.backendClose = wcClose wc + , TLS.backendSend = wcSend wc + , TLS.backendRecv = \n -> do + -- TLS library expects exactly n bytes + wcRecv wc n + } + +mkTlsWireConn :: TLS.Context -> IORef ByteString -> IO WireConn +mkTlsWireConn ctx bufRef = do + writeIORef bufRef BS.empty + tlsBuf <- newIORef BS.empty + traceRef <- newIORef Nothing + let tlsSend = TLS.sendData ctx . LBS.fromStrict + pure + WireConn + { wcSend = tlsSend + , wcSendMany = \chunks -> tlsSend (BS.concat chunks) + , wcRecv = tlsRecvExact ctx tlsBuf + , wcClose = TLS.bye ctx >> TLS.contextClose ctx + , wcBuffer = tlsBuf + , wcTrace = traceRef + } + +-- | Receive exactly n bytes from a TLS context, buffering leftovers. +-- Same optimization as 'recvExact': fast path for buffered data, Builder +-- accumulation for multi-chunk reads. +tlsRecvExact :: TLS.Context -> IORef ByteString -> Int -> IO ByteString +tlsRecvExact ctx bufRef n = do + buf <- readIORef bufRef + if BS.length buf >= n + then do + let (!taken, !rest) = BS.splitAt n buf + writeIORef bufRef rest + pure taken + else if BS.null buf + then do + chunk <- TLS.recvData ctx + if BS.null chunk + then throwPgWire (ConnectionError "TLS connection closed") + else if BS.length chunk >= n + then do + let (!taken, !rest) = BS.splitAt n chunk + writeIORef bufRef rest + pure taken + else tlsAccumulate ctx bufRef chunk (n - BS.length chunk) (B.byteString chunk) + else do + let !remaining = n - BS.length buf + !builder0 = B.byteString buf + chunk <- TLS.recvData ctx + if BS.null chunk + then throwPgWire (ConnectionError "TLS connection closed") + else if BS.length chunk >= remaining + then do + let (!taken, !rest) = BS.splitAt remaining chunk + writeIORef bufRef rest + pure (LBS.toStrict (B.toLazyByteString (builder0 <> B.byteString taken))) + else tlsAccumulate ctx bufRef chunk (remaining - BS.length chunk) (builder0 <> B.byteString chunk) + +tlsAccumulate :: TLS.Context -> IORef ByteString -> ByteString -> Int -> B.Builder -> IO ByteString +tlsAccumulate ctx bufRef chunk remaining !builder + | BS.length chunk >= remaining = do + let (!taken, !rest) = BS.splitAt remaining chunk + writeIORef bufRef rest + pure (LBS.toStrict (B.toLazyByteString (builder <> B.byteString taken))) + | otherwise = do + let !remaining' = remaining - BS.length chunk + !builder' = builder <> B.byteString chunk + next <- TLS.recvData ctx + if BS.null next + then throwPgWire (ConnectionError "TLS connection closed") + else tlsAccumulate ctx bufRef next remaining' builder' + +-- | Send a frontend message over the wire. +sendFrontendMsg :: WireConn -> FrontendMsg -> IO () +sendFrontendMsg wc msg = do + let bytes = buildFrontendMsg msg + traceIfEnabled wc TraceSend bytes + wcSend wc bytes +{-# INLINE sendFrontendMsg #-} + +-- | Send multiple frontend messages in a single allocation and syscall. +-- All messages are fused into one 'Builder', materialized once with +-- 'buildFrontendMsgsConcat', and sent with a single @send()@. +sendFrontendMsgs :: WireConn -> [FrontendMsg] -> IO () +sendFrontendMsgs wc msgs = do + let !bytes = buildFrontendMsgsConcat msgs + traceIfEnabled wc TraceSend bytes + wcSend wc bytes +{-# INLINE sendFrontendMsgs #-} + +-- | Send raw bytes (for startup message which has a different format). +sendRawBytes :: WireConn -> ByteString -> IO () +sendRawBytes wc bs = do + traceIfEnabled wc TraceSend bs + wcSend wc bs + +-- | Call the trace handler if one is set. +traceIfEnabled :: WireConn -> TraceDirection -> ByteString -> IO () +traceIfEnabled wc dir bytes = do + mHandler <- readIORef (wcTrace wc) + case mHandler of + Nothing -> pure () + Just handler -> handler dir bytes +{-# INLINE traceIfEnabled #-} + +-- | Receive and parse a single backend message. +recvBackendMsg :: WireConn -> IO BackendMsg +recvBackendMsg wc = do + -- Read tag (1 byte) + length (4 bytes) together in a single recv + header <- wcRecv wc 5 + let !tag = BU.unsafeIndex header 0 + !len = decodeInt32At header 1 + !payloadLen = len - 4 + -- Read payload + payload <- + if payloadLen > 0 + then wcRecv wc (fromIntegral payloadLen) + else if payloadLen < 0 + then throwPgWire (ProtocolError ("Invalid message length: " <> BS8.pack (show len))) + else pure BS.empty + -- Trace the raw bytes (header + payload) + traceIfEnabled wc TraceRecv (header <> payload) + case parseBackendMsg tag payload of + Left err -> throwPgWire (ProtocolError (BS8.pack err)) + Right msg -> pure msg + where + decodeInt32At :: ByteString -> Int -> Int + decodeInt32At bs off = + let !b0 = fromIntegral (BU.unsafeIndex bs off) + !b1 = fromIntegral (BU.unsafeIndex bs (off + 1)) + !b2 = fromIntegral (BU.unsafeIndex bs (off + 2)) + !b3 = fromIntegral (BU.unsafeIndex bs (off + 3)) + in b0 * 16777216 + b1 * 65536 + b2 * 256 + b3 + +-- Socket helpers ---------------------------------------------------------- + +sendAll :: Socket -> ByteString -> IO () +sendAll sock bs + | BS.null bs = pure () + | otherwise = do + sent <- NSB.send sock bs + sendAll sock (BS.drop sent bs) + +-- | Socket recv buffer size. Larger values reduce syscalls for big result sets. +recvChunkSize :: Int +recvChunkSize = 32768 +{-# INLINE recvChunkSize #-} + +-- | Receive exactly @n@ bytes, using the buffer for leftovers. +-- Fast path: if the buffer already has enough data, just splitAt (zero allocation +-- beyond the two slices). Slow path: accumulate chunks via Builder to avoid +-- the O(n) reverse + concat of a list accumulator. +recvExact :: Socket -> IORef ByteString -> Int -> IO ByteString +recvExact sock bufRef n = do + buf <- readIORef bufRef + if BS.length buf >= n + then do + -- Fast path: buffer has enough data, no syscall needed. + let (!taken, !rest) = BS.splitAt n buf + writeIORef bufRef rest + pure taken + else if BS.null buf + then do + -- Buffer empty: read fresh chunk from socket. + chunk <- NSB.recv sock (max recvChunkSize n) + if BS.null chunk + then throwPgWire (ConnectionError "Connection closed by server") + else if BS.length chunk >= n + then do + let (!taken, !rest) = BS.splitAt n chunk + writeIORef bufRef rest + pure taken + else accumulate chunk (n - BS.length chunk) (B.byteString chunk) + else do + -- Buffer has partial data: use it and read more. + let !remaining = n - BS.length buf + !builder0 = B.byteString buf + chunk <- NSB.recv sock (max recvChunkSize remaining) + if BS.null chunk + then throwPgWire (ConnectionError "Connection closed by server") + else if BS.length chunk >= remaining + then do + let (!taken, !rest) = BS.splitAt remaining chunk + writeIORef bufRef rest + pure (LBS.toStrict (B.toLazyByteString (builder0 <> B.byteString taken))) + else accumulate chunk (remaining - BS.length chunk) (builder0 <> B.byteString chunk) + where + -- Slow path: need multiple recv calls. Uses Builder for O(1) append. + accumulate :: ByteString -> Int -> B.Builder -> IO ByteString + accumulate chunk remaining !builder + | BS.length chunk >= remaining = do + let (!taken, !rest) = BS.splitAt remaining chunk + writeIORef bufRef rest + pure (LBS.toStrict (B.toLazyByteString (builder <> B.byteString taken))) + | otherwise = do + let !remaining' = remaining - BS.length chunk + !builder' = builder <> B.byteString chunk + next <- NSB.recv sock (max recvChunkSize remaining') + if BS.null next + then throwPgWire (ConnectionError "Connection closed by server") + else accumulate next remaining' builder' diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,38 @@ +module Main where + +import PgWire.Auth.MD5Spec qualified as MD5Spec +import PgWire.Auth.ScramFieldsSpec qualified as ScramFieldsSpec +import PgWire.CancelSpec qualified as CancelSpec +import PgWire.Connection.ConfigSpec qualified as ConfigSpec +import PgWire.ErrorSpec qualified as ErrorSpec +import PgWire.MockServerSpec qualified as MockServerSpec +import PgWire.Pool.ConfigSpec qualified as PoolConfigSpec +import PgWire.Pool.NoThunksSpec qualified as NoThunksSpec +import PgWire.Pool.PropertySpec qualified as PoolPropertySpec +import PgWire.Pool.StateMachineSpec qualified as PoolStateMachineSpec +import PgWire.Protocol.BuildersSpec qualified as BuildersSpec +import PgWire.Protocol.OidSpec qualified as OidSpec +import PgWire.Connection.EscapingSpec qualified as EscapingSpec +import PgWire.Connection.FeaturesSpec qualified as FeaturesSpec +import PgWire.Protocol.FuzzSpec qualified as FuzzSpec +import PgWire.Protocol.ParsersSpec qualified as ParsersSpec +import Test.Hspec + +main :: IO () +main = hspec $ do + describe "PgWire.Protocol.Builders" BuildersSpec.spec + describe "PgWire.Protocol.Parsers" ParsersSpec.spec + describe "PgWire.Protocol.Fuzz" FuzzSpec.spec + describe "PgWire.Protocol.Oid" OidSpec.spec + describe "PgWire.Connection.Config" ConfigSpec.spec + describe "PgWire.Pool.Config" PoolConfigSpec.spec + describe "PgWire.Error" ErrorSpec.spec + describe "PgWire.Auth.MD5" MD5Spec.spec + describe "PgWire.Auth.ScramFields" ScramFieldsSpec.spec + describe "PgWire.Cancel" CancelSpec.spec + describe "PgWire.Connection.Features" FeaturesSpec.spec + describe "PgWire.Connection.Escaping" EscapingSpec.spec + describe "PgWire.MockServer" MockServerSpec.spec + describe "PgWire.Pool.NoThunks" NoThunksSpec.spec + describe "PgWire.Pool.Property" PoolPropertySpec.spec + describe "PgWire.Pool.StateMachine" PoolStateMachineSpec.spec diff --git a/test/PgWire/Auth/MD5Spec.hs b/test/PgWire/Auth/MD5Spec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Auth/MD5Spec.hs @@ -0,0 +1,44 @@ +module PgWire.Auth.MD5Spec (spec) where + +import Crypto.Hash (MD5 (..), hashWith) +import Data.ByteArray qualified as BA +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Char8 qualified as BS8 +import Test.Hspec + +spec :: Spec +spec = do + describe "MD5 password hash" $ do + it "produces correct md5(md5(password+user)+salt) format" $ do + -- PG MD5 auth: md5(md5(password + username) + salt) + -- where md5() returns lowercase hex string + let user = "testuser" :: ByteString + password = "testpass" :: ByteString + salt = BS.pack [0x01, 0x02, 0x03, 0x04] + inner = md5Hex (password <> user) + outer = "md5" <> md5Hex (inner <> salt) + -- The hash should be deterministic + BS.length outer `shouldBe` 35 -- "md5" + 32 hex chars + BS8.isPrefixOf "md5" outer `shouldBe` True + + it "different passwords produce different hashes" $ do + let user = "user" :: ByteString + salt = BS.pack [0, 0, 0, 0] + h1 = md5Hex (md5Hex ("pass1" <> user) <> salt) + h2 = md5Hex (md5Hex ("pass2" <> user) <> salt) + h1 `shouldNotBe` h2 + +md5Hex :: ByteString -> ByteString +md5Hex bs = BS8.pack $ concatMap (pad2 . showHex') (BS.unpack (BA.convert (hashWith MD5 bs))) + where + showHex' n = showHexByte n "" + pad2 s = replicate (2 - length s) '0' ++ s + showHexByte n = showHex n + showHex 0 = ('0' :) + showHex n + | n < 16 = (hexDigit (fromIntegral n) :) + | otherwise = showHex (n `div` 16) . (hexDigit (fromIntegral (n `mod` 16)) :) + hexDigit i + | i < 10 = toEnum (fromEnum '0' + i) + | otherwise = toEnum (fromEnum 'a' + i - 10) diff --git a/test/PgWire/Auth/ScramFieldsSpec.hs b/test/PgWire/Auth/ScramFieldsSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Auth/ScramFieldsSpec.hs @@ -0,0 +1,44 @@ +module PgWire.Auth.ScramFieldsSpec (spec) where + +import Data.ByteString qualified as BS +import Test.Hspec + +-- Test the SCRAM field parser (internal, but we replicate the logic here) +spec :: Spec +spec = do + describe "SCRAM field parsing" $ do + it "parses r=nonce,s=salt,i=4096" $ do + let input = "r=clientnonce+servernonce,s=c2FsdA==,i=4096" + fields = parseScramFields input + lookup "r" fields `shouldBe` Just "clientnonce+servernonce" + lookup "s" fields `shouldBe` Just "c2FsdA==" + lookup "i" fields `shouldBe` Just "4096" + + it "parses v=signature" $ do + let input = "v=dGVzdHNpZw==" + fields = parseScramFields input + lookup "v" fields `shouldBe` Just "dGVzdHNpZw==" + + it "handles empty input" $ do + let fields = parseScramFields "" + fields `shouldBe` [] + + it "handles single field" $ do + let fields = parseScramFields "r=nonce" + lookup "r" fields `shouldBe` Just "nonce" + + it "handles values containing equals signs" $ do + -- base64 can contain '=' padding + let fields = parseScramFields "s=abc=,i=100" + -- The key is "s", value starts after first '=' + lookup "s" fields `shouldBe` Just "abc=" + +-- Replicated from ScramSHA256 module for testing +parseScramFields :: BS.ByteString -> [(BS.ByteString, BS.ByteString)] +parseScramFields bs = + [ (key, val) + | field <- BS.split 44 bs -- ',' + , let (key, rest) = BS.break (== 61) field -- '=' + , not (BS.null rest) + , let val = BS.drop 1 rest + ] diff --git a/test/PgWire/CancelSpec.hs b/test/PgWire/CancelSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/CancelSpec.hs @@ -0,0 +1,52 @@ +module PgWire.CancelSpec (spec) where + +import Data.ByteString.Builder qualified as B +import Data.ByteString.Lazy qualified as LBS +import Data.ByteString qualified as BS +import Data.Int (Int32) +import Test.Hspec + +spec :: Spec +spec = do + describe "CancelRequest message format" $ do + it "is exactly 16 bytes" $ do + let msg = buildCancelRequest 12345 67890 + BS.length msg `shouldBe` 16 + + it "starts with length 16" $ do + let msg = buildCancelRequest 1 2 + -- First 4 bytes: length = 16 = 0x00000010 + BS.take 4 msg `shouldBe` BS.pack [0, 0, 0, 16] + + it "contains cancel code 80877102" $ do + let msg = buildCancelRequest 1 2 + -- Bytes 4-7: cancel code = 80877102 = 0x04D2162E + let code = fromIntegral (BS.index msg 4) * 256 * 256 * 256 + + fromIntegral (BS.index msg 5) * 256 * 256 + + fromIntegral (BS.index msg 6) * 256 + + fromIntegral (BS.index msg 7) :: Int + code `shouldBe` 80877102 + + it "contains PID and key" $ do + let msg = buildCancelRequest 42 99 + -- Bytes 8-11: PID + let pid = fromIntegral (BS.index msg 8) * 256 * 256 * 256 + + fromIntegral (BS.index msg 9) * 256 * 256 + + fromIntegral (BS.index msg 10) * 256 + + fromIntegral (BS.index msg 11) :: Int + pid `shouldBe` 42 + -- Bytes 12-15: key + let key = fromIntegral (BS.index msg 12) * 256 * 256 * 256 + + fromIntegral (BS.index msg 13) * 256 * 256 + + fromIntegral (BS.index msg 14) * 256 + + fromIntegral (BS.index msg 15) :: Int + key `shouldBe` 99 + +-- Replicate the cancel message builder for testing +buildCancelRequest :: Int32 -> Int32 -> BS.ByteString +buildCancelRequest pid key = + LBS.toStrict . B.toLazyByteString $ + B.int32BE 16 + <> B.int32BE 80877102 + <> B.int32BE pid + <> B.int32BE key diff --git a/test/PgWire/Connection/ConfigSpec.hs b/test/PgWire/Connection/ConfigSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Connection/ConfigSpec.hs @@ -0,0 +1,71 @@ +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} +module PgWire.Connection.ConfigSpec (spec) where + +import PgWire.Connection.Config +import Test.Hspec + +spec :: Spec +spec = do + describe "parseConnString (URI format)" $ do + it "parses basic postgres:// URL" $ do + let Right cfg = parseConnString "postgres://user:pass@localhost:5432/mydb" + ccHost cfg `shouldBe` "localhost" + ccPort cfg `shouldBe` 5432 + ccUser cfg `shouldBe` "user" + ccPassword cfg `shouldBe` "pass" + ccDatabase cfg `shouldBe` "mydb" + + it "parses postgresql:// scheme" $ do + let Right cfg = parseConnString "postgresql://u:p@host:1234/db" + ccHost cfg `shouldBe` "host" + ccPort cfg `shouldBe` 1234 + + it "defaults port to 5432" $ do + let Right cfg = parseConnString "postgres://u:p@myhost/db" + ccPort cfg `shouldBe` 5432 + + it "parses sslmode=require" $ do + let Right cfg = parseConnString "postgres://u:p@h:5432/db?sslmode=require" + ccTls cfg `shouldBe` TlsRequire + + it "parses connect_timeout from query params" $ do + let Right cfg = parseConnString "postgres://u:p@h:5432/db?connect_timeout=30" + ccConnectTimeout cfg `shouldBe` 30 + + it "defaults connect timeout to 10" $ do + let Right cfg = parseConnString "postgres://u:p@h:5432/db" + ccConnectTimeout cfg `shouldBe` 10 + + it "defaults query timeout to 0 (no timeout)" $ do + let Right cfg = parseConnString "postgres://u:p@h:5432/db" + ccQueryTimeout cfg `shouldBe` 0 + + describe "parseConnString (key=value format)" $ do + it "parses key=value pairs" $ do + let Right cfg = parseConnString "host=myhost port=9999 user=alice password=secret dbname=testdb" + ccHost cfg `shouldBe` "myhost" + ccPort cfg `shouldBe` 9999 + ccUser cfg `shouldBe` "alice" + ccPassword cfg `shouldBe` "secret" + ccDatabase cfg `shouldBe` "testdb" + + it "defaults missing fields" $ do + let Right cfg = parseConnString "host=myhost dbname=db" + ccPort cfg `shouldBe` 5432 + ccUser cfg `shouldBe` "" + + it "parses connect_timeout" $ do + let Right cfg = parseConnString "host=h connect_timeout=5" + ccConnectTimeout cfg `shouldBe` 5 + + it "parses sslmode" $ do + let Right cfg = parseConnString "host=h sslmode=prefer" + ccTls cfg `shouldBe` TlsPrefer + + describe "defaultConnConfig" $ do + it "has sensible defaults" $ do + ccHost defaultConnConfig `shouldBe` "localhost" + ccPort defaultConnConfig `shouldBe` 5432 + ccTls defaultConnConfig `shouldBe` TlsDisable + ccConnectTimeout defaultConnConfig `shouldBe` 10 + ccQueryTimeout defaultConnConfig `shouldBe` 0 diff --git a/test/PgWire/Connection/EscapingSpec.hs b/test/PgWire/Connection/EscapingSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Connection/EscapingSpec.hs @@ -0,0 +1,38 @@ +module PgWire.Connection.EscapingSpec (spec) where + +import PgWire.Connection (escapeLiteral, escapeIdentifier) +import Test.Hspec + +spec :: Spec +spec = do + describe "escapeLiteral" $ do + it "wraps in single quotes" $ + escapeLiteral undefined "hello" `shouldBe` "'hello'" + + it "escapes single quotes" $ + escapeLiteral undefined "O'Brien" `shouldBe` "'O''Brien'" + + it "escapes backslashes" $ + escapeLiteral undefined "back\\slash" `shouldBe` "'back\\\\slash'" + + it "handles empty string" $ + escapeLiteral undefined "" `shouldBe` "''" + + it "handles multiple quotes" $ + escapeLiteral undefined "it's a 'test'" `shouldBe` "'it''s a ''test'''" + + describe "escapeIdentifier" $ do + it "wraps in double quotes" $ + escapeIdentifier undefined "column" `shouldBe` "\"column\"" + + it "escapes double quotes" $ + escapeIdentifier undefined "my\"col" `shouldBe` "\"my\"\"col\"" + + it "handles spaces" $ + escapeIdentifier undefined "user table" `shouldBe` "\"user table\"" + + it "handles empty string" $ + escapeIdentifier undefined "" `shouldBe` "\"\"" + + it "handles reserved words" $ + escapeIdentifier undefined "select" `shouldBe` "\"select\"" diff --git a/test/PgWire/Connection/FeaturesSpec.hs b/test/PgWire/Connection/FeaturesSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Connection/FeaturesSpec.hs @@ -0,0 +1,102 @@ +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} +module PgWire.Connection.FeaturesSpec (spec) where + +import PgWire.Connection.Config +import Test.Hspec + +spec :: Spec +spec = do + describe "parseHosts" $ do + it "parses single host" $ + parseHosts "localhost" `shouldBe` ["localhost"] + + it "parses multiple hosts" $ + parseHosts "host1,host2,host3" `shouldBe` ["host1", "host2", "host3"] + + it "returns localhost for empty string" $ + parseHosts "" `shouldBe` ["localhost"] + + it "handles two hosts" $ + parseHosts "primary,replica" `shouldBe` ["primary", "replica"] + + describe "TlsMode parsing" $ do + it "parses verify-ca from URI" $ do + let Right cfg = parseConnString "postgres://u:p@h:5432/db?sslmode=verify-ca" + ccTls cfg `shouldBe` TlsVerifyCa + + it "parses verify-full from URI" $ do + let Right cfg = parseConnString "postgres://u:p@h:5432/db?sslmode=verify-full" + ccTls cfg `shouldBe` TlsVerifyFull + + it "parses verify-ca from key-value" $ do + let Right cfg = parseConnString "host=h sslmode=verify-ca" + ccTls cfg `shouldBe` TlsVerifyCa + + it "parses verify-full from key-value" $ do + let Right cfg = parseConnString "host=h sslmode=verify-full" + ccTls cfg `shouldBe` TlsVerifyFull + + describe "SSL cert params" $ do + it "parses sslcert from URI" $ do + let Right cfg = parseConnString "postgres://u:p@h:5432/db?sslcert=/path/cert.pem" + ccSslCert cfg `shouldBe` Just "/path/cert.pem" + + it "parses sslkey from key-value" $ do + let Right cfg = parseConnString "host=h sslkey=/path/key.pem" + ccSslKey cfg `shouldBe` Just "/path/key.pem" + + it "parses sslrootcert" $ do + let Right cfg = parseConnString "host=h sslrootcert=/path/ca.pem" + ccSslRootCert cfg `shouldBe` Just "/path/ca.pem" + + it "defaults to Nothing" $ do + let Right cfg = parseConnString "host=h" + ccSslCert cfg `shouldBe` Nothing + ccSslKey cfg `shouldBe` Nothing + ccSslRootCert cfg `shouldBe` Nothing + + describe "keepalive config" $ do + it "parses keepalives=0 as False" $ do + let Right cfg = parseConnString "host=h keepalives=0" + ccKeepalives cfg `shouldBe` False + + it "defaults keepalives to True" $ do + let Right cfg = parseConnString "host=h" + ccKeepalives cfg `shouldBe` True + + it "parses keepalives_idle" $ do + let Right cfg = parseConnString "host=h keepalives_idle=60" + ccKeepalivesIdle cfg `shouldBe` 60 + + it "parses keepalives_interval" $ do + let Right cfg = parseConnString "host=h keepalives_interval=10" + ccKeepalivesInterval cfg `shouldBe` 10 + + it "parses keepalives_count" $ do + let Right cfg = parseConnString "host=h keepalives_count=5" + ccKeepalivesCount cfg `shouldBe` 5 + + describe "target_session_attrs" $ do + it "parses read-write from URI" $ do + let Right cfg = parseConnString "postgres://u:p@h:5432/db?target_session_attrs=read-write" + ccTargetSessionAttrs cfg `shouldBe` SessionReadWrite + + it "parses primary from key-value" $ do + let Right cfg = parseConnString "host=h target_session_attrs=primary" + ccTargetSessionAttrs cfg `shouldBe` SessionPrimary + + it "parses standby" $ do + let Right cfg = parseConnString "host=h target_session_attrs=standby" + ccTargetSessionAttrs cfg `shouldBe` SessionStandby + + it "parses prefer-standby" $ do + let Right cfg = parseConnString "host=h target_session_attrs=prefer-standby" + ccTargetSessionAttrs cfg `shouldBe` SessionPreferStandby + + it "parses read-only" $ do + let Right cfg = parseConnString "host=h target_session_attrs=read-only" + ccTargetSessionAttrs cfg `shouldBe` SessionReadOnly + + it "defaults to any" $ do + let Right cfg = parseConnString "host=h" + ccTargetSessionAttrs cfg `shouldBe` SessionAny diff --git a/test/PgWire/ErrorSpec.hs b/test/PgWire/ErrorSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/ErrorSpec.hs @@ -0,0 +1,62 @@ +module PgWire.ErrorSpec (spec) where + +import Control.Exception (catch) +import PgWire.Error +import PgWire.Protocol.Backend (PgError (..)) +import Test.Hspec + +spec :: Spec +spec = do + describe "PgWireError constructors" $ do + it "ConnectionError carries message" $ do + let err = ConnectionError "host unreachable" + show err `shouldContain` "host unreachable" + + it "AuthError carries message" $ do + let err = AuthError "bad password" + show err `shouldContain` "bad password" + + it "ProtocolError carries message" $ do + let err = ProtocolError "unexpected tag" + show err `shouldContain` "unexpected tag" + + it "DecodeError carries message" $ do + let err = DecodeError "int32: wrong size" + show err `shouldContain` "int32: wrong size" + + it "PoolTimeout has no payload" $ + show PoolTimeout `shouldContain` "PoolTimeout" + + it "PoolClosed has no payload" $ + show PoolClosed `shouldContain` "PoolClosed" + + it "QueryError wraps PgError" $ do + let pgErr = emptyPgError {pgMessage = "table not found", pgCode = "42P01"} + err = QueryError pgErr + show err `shouldContain` "table not found" + show err `shouldContain` "42P01" + + describe "throwPgWire" $ do + it "throws as an Exception" $ do + let action = throwPgWire (ConnectionError "test") :: IO () + action `shouldThrow` (== ConnectionError "test") + + it "can be caught with specific pattern" $ do + result <- (throwPgWire PoolTimeout >> pure "no") `catch` \(e :: PgWireError) -> + case e of + PoolTimeout -> pure "caught" + _ -> pure "wrong" + result `shouldBe` ("caught" :: String) + + describe "Eq instance" $ do + it "equal errors are equal" $ + ConnectionError "x" `shouldBe` ConnectionError "x" + + it "different errors are not equal" $ + ConnectionError "x" `shouldNotBe` ConnectionError "y" + + it "different constructors are not equal" $ + ConnectionError "x" `shouldNotBe` AuthError "x" + +emptyPgError :: PgError +emptyPgError = PgError "" "" "" Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing diff --git a/test/PgWire/MockServerSpec.hs b/test/PgWire/MockServerSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/MockServerSpec.hs @@ -0,0 +1,113 @@ +module PgWire.MockServerSpec (spec) where + +import Control.Exception (try) +import Data.ByteString.Char8 qualified as BS8 +import PgWire.Connection (Connection, close, connectString, simpleQuery) +import PgWire.Error (PgWireError (..)) +import PgWire.MockServer +import Test.Hspec + +connectToMock :: MockConfig -> (Connection -> IO a) -> IO a +connectToMock cfg action = + withMockServer cfg $ \port -> do + let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb" + conn <- connectString connStr + result <- action conn + close conn + pure result + +spec :: Spec +spec = do + describe "MockServer" $ do + describe "basic protocol" $ do + it "accepts a connection and responds to simple query" $ do + connectToMock defaultMockConfig $ \conn -> do + (rows, _) <- simpleQuery conn "SELECT 1" + rows `shouldBe` [] + + it "returns rows from a custom handler" $ do + let cfg = defaultMockConfig + { mockQueryHandler = simpleHandler + [ ("SELECT name FROM test", + [ [("name", "Alice")] + , [("name", "Bob")] + ]) + ] + } + connectToMock cfg $ \conn -> do + (rows, _) <- simpleQuery conn "SELECT name FROM test" + length rows `shouldBe` 2 + + it "handles unknown queries gracefully" $ do + let cfg = defaultMockConfig + { mockQueryHandler = simpleHandler [] + } + connectToMock cfg $ \conn -> do + (rows, _) <- simpleQuery conn "SELECT nothing" + rows `shouldBe` [] + + it "supports multiple sequential queries" $ do + connectToMock defaultMockConfig $ \conn -> do + _ <- simpleQuery conn "SELECT 1" + _ <- simpleQuery conn "SELECT 2" + (rows, _) <- simpleQuery conn "SELECT 3" + rows `shouldBe` [] + + describe "error responses" $ do + it "propagates server errors as QueryError" $ do + let cfg = defaultMockConfig + { mockQueryHandler = errorHandler "42P01" "relation \"foo\" does not exist" + } + connectToMock cfg $ \conn -> do + result <- try @PgWireError (simpleQuery conn "SELECT * FROM foo") + case result of + Left (QueryError _) -> pure () + Left other -> expectationFailure $ "Expected QueryError, got: " <> show other + Right _ -> expectationFailure "Expected an error" + + it "recovers after an error and accepts subsequent queries" $ do + let cfg = defaultMockConfig + { mockQueryHandler = \sql send -> + if sql == "BAD" + then send $ buildErrorResponse "42601" "syntax error" + else send $ buildCommandComplete "SELECT 0" + } + connectToMock cfg $ \conn -> do + -- First query errors + result1 <- try @PgWireError (simpleQuery conn "BAD") + case result1 of + Left (QueryError _) -> pure () + _ -> expectationFailure "Expected QueryError for BAD query" + -- Second query succeeds on the same connection + (rows, _) <- simpleQuery conn "GOOD" + rows `shouldBe` [] + + describe "multi-row results" $ do + it "returns correct number of rows for large result sets" $ do + let manyRows = [ [("id", BS8.pack (show i)), ("val", "x")] + | i <- [1 :: Int .. 100] + ] + cfg = defaultMockConfig + { mockQueryHandler = simpleHandler + [("SELECT id, val FROM big", manyRows)] + } + connectToMock cfg $ \conn -> do + (rows, _) <- simpleQuery conn "SELECT id, val FROM big" + length rows `shouldBe` 100 + + it "returns multiple columns correctly" $ do + let cfg = defaultMockConfig + { mockQueryHandler = simpleHandler + [("SELECT a, b, c FROM t", + [ [("a", "1"), ("b", "hello"), ("c", "true")] + , [("a", "2"), ("b", "world"), ("c", "false")] + ])] + } + connectToMock cfg $ \conn -> do + (rows, _) <- simpleQuery conn "SELECT a, b, c FROM t" + length rows `shouldBe` 2 + case rows of + [row1, row2] -> do + length row1 `shouldBe` 3 + length row2 `shouldBe` 3 + _ -> expectationFailure "Expected exactly 2 rows" diff --git a/test/PgWire/Pool/ConfigSpec.hs b/test/PgWire/Pool/ConfigSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Pool/ConfigSpec.hs @@ -0,0 +1,22 @@ +module PgWire.Pool.ConfigSpec (spec) where + +import PgWire.Pool.Config +import Test.Hspec + +spec :: Spec +spec = do + describe "defaultPoolConfig" $ do + it "has empty connection string" $ + poolConnString defaultPoolConfig `shouldBe` "" + + it "has pool size 10" $ + poolSize defaultPoolConfig `shouldBe` 10 + + it "has idle time 600 seconds" $ + poolIdleTime defaultPoolConfig `shouldBe` 600 + + it "has max life 3600 seconds" $ + poolMaxLife defaultPoolConfig `shouldBe` 3600 + + it "has acquire timeout 10 seconds" $ + poolAcquireTimeout defaultPoolConfig `shouldBe` 10 diff --git a/test/PgWire/Pool/NoThunksSpec.hs b/test/PgWire/Pool/NoThunksSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Pool/NoThunksSpec.hs @@ -0,0 +1,23 @@ +module PgWire.Pool.NoThunksSpec (spec) where + +import NoThunks.Class (noThunks) +import PgWire.Pool (PoolStats (..)) +import Test.Hspec + +spec :: Spec +spec = do + describe "NoThunks" $ do + it "PoolStats has no thunks when fully evaluated" $ do + let stats = PoolStats + { psIdle = 3 + , psInUse = 5 + , psWaiters = 0 + , psMaxSize = 10 + , psTotalCreated = 20 + , psTotalDestroyed = 12 + , psTotalTimeouts = 1 + } + result <- noThunks [] stats + case result of + Nothing -> pure () + Just thunkInfo -> expectationFailure $ "Found thunk: " <> show thunkInfo diff --git a/test/PgWire/Pool/PropertySpec.hs b/test/PgWire/Pool/PropertySpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Pool/PropertySpec.hs @@ -0,0 +1,128 @@ +module PgWire.Pool.PropertySpec (spec) where + +import Control.Concurrent.Async (mapConcurrently) +import Control.Exception (try) +import Data.ByteString.Char8 qualified as BS8 +import Hedgehog +import Hedgehog.Gen qualified as Gen +import Hedgehog.Range qualified as Range +import PgWire.Connection (simpleQuery) +import PgWire.Error (PgWireError (..)) +import PgWire.MockServer (defaultMockConfig, withMockServer) +import PgWire.Pool (Pool, PoolStats (..), closePool, newPool, poolIsAlive, poolStats, resize, withResource) +import PgWire.Pool.Config (PoolConfig (..), defaultPoolConfig) +import Test.Hspec +import Test.Hspec.Hedgehog (hedgehog) + +-- | Create a pool backed by the mock server. +withMockPool :: Int -> (Pool -> IO a) -> IO a +withMockPool size action = + withMockServer defaultMockConfig $ \port -> do + let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb" + pool <- newPool defaultPoolConfig + { poolConnString = connStr + , poolSize = size + , poolAcquireTimeout = 5 + } + result <- action pool + closePool pool + pure result + +spec :: Spec +spec = do + describe "Pool invariants" $ do + it "idle + inUse never exceeds maxSize" $ hedgehog $ do + size <- forAll (Gen.int (Range.linear 1 10)) + nOps <- forAll (Gen.int (Range.linear 1 20)) + result <- evalIO $ withMockPool size $ \pool -> do + -- Do several acquire/release cycles + mapM_ (\_ -> withResource pool $ \conn -> + simpleQuery conn "SELECT 1") [1 .. nOps] + stats <- poolStats pool + pure (psIdle stats + psInUse stats <= psMaxSize stats) + assert result + + it "all connections returned after withResource" $ hedgehog $ do + size <- forAll (Gen.int (Range.linear 1 8)) + nOps <- forAll (Gen.int (Range.linear 1 15)) + stats <- evalIO $ withMockPool size $ \pool -> do + mapM_ (\_ -> withResource pool $ \conn -> + simpleQuery conn "SELECT 1") [1 .. nOps] + poolStats pool + psInUse stats === 0 + + it "concurrent acquire/release maintains invariant" $ hedgehog $ do + size <- forAll (Gen.int (Range.linear 2 8)) + nThreads <- forAll (Gen.int (Range.linear 2 10)) + stats <- evalIO $ withMockPool size $ \pool -> do + _ <- mapConcurrently (\_ -> + withResource pool $ \conn -> + simpleQuery conn "SELECT 1") [1 .. nThreads] + poolStats pool + -- After all threads finish, no connections should be in use + psInUse stats === 0 + -- idle + inUse should not exceed max + assert (psIdle stats <= size) + + it "totalCreated >= totalDestroyed" $ hedgehog $ do + size <- forAll (Gen.int (Range.linear 1 5)) + nOps <- forAll (Gen.int (Range.linear 1 10)) + stats <- evalIO $ withMockPool size $ \pool -> do + mapM_ (\_ -> withResource pool $ \conn -> + simpleQuery conn "SELECT 1") [1 .. nOps] + poolStats pool + assert (psTotalCreated stats >= psTotalDestroyed stats) + + it "closePool makes pool not alive" $ hedgehog $ do + size <- forAll (Gen.int (Range.linear 1 5)) + alive <- evalIO $ withMockServer defaultMockConfig $ \port -> do + let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb" + pool <- newPool defaultPoolConfig + { poolConnString = connStr + , poolSize = size + , poolAcquireTimeout = 5 + } + closePool pool + poolIsAlive pool + assert (not alive) + + it "closePool rejects new acquires" $ hedgehog $ do + size <- forAll (Gen.int (Range.linear 1 5)) + gotError <- evalIO $ withMockServer defaultMockConfig $ \port -> do + let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb" + pool <- newPool defaultPoolConfig + { poolConnString = connStr + , poolSize = size + , poolAcquireTimeout = 5 + } + closePool pool + result <- try (withResource pool $ \_ -> pure ()) + pure $ case result of + Left PoolClosed -> True + Left _ -> False + Right _ -> False + assert gotError + + it "resize changes maxSize" $ hedgehog $ do + initial <- forAll (Gen.int (Range.linear 2 10)) + newSize <- forAll (Gen.int (Range.linear 1 20)) + stats <- evalIO $ withMockPool initial $ \pool -> do + resize pool newSize + poolStats pool + psMaxSize stats === newSize + + it "pool stats are consistent" $ hedgehog $ do + size <- forAll (Gen.int (Range.linear 1 8)) + nOps <- forAll (Gen.int (Range.linear 0 10)) + stats <- evalIO $ withMockPool size $ \pool -> do + mapM_ (\_ -> withResource pool $ \conn -> + simpleQuery conn "SELECT 1") [1 .. nOps] + poolStats pool + -- Basic consistency checks + assert (psIdle stats >= 0) + assert (psInUse stats >= 0) + assert (psWaiters stats >= 0) + assert (psMaxSize stats > 0) + assert (psTotalCreated stats >= 0) + assert (psTotalDestroyed stats >= 0) + assert (psTotalTimeouts stats >= 0) diff --git a/test/PgWire/Pool/StateMachineSpec.hs b/test/PgWire/Pool/StateMachineSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Pool/StateMachineSpec.hs @@ -0,0 +1,178 @@ +-- | Lightweight state-machine tests for the connection pool. +-- +-- Existing 'PgWire.Pool.PropertySpec' tests cover individual scenarios +-- (a burst of acquires, a single resize, a single close). This module +-- generates arbitrary sequences of Burst/Resize/Close and checks the +-- model-vs-real invariants after every single step. The model is +-- hand-rolled rather than using 'Hedgehog.Internal.State' so the test +-- drops straight into the existing hspec-hedgehog runner. +module PgWire.Pool.StateMachineSpec (spec) where + +import Control.Concurrent.Async (mapConcurrently_) +import Control.Exception (try) +import Data.ByteString.Char8 qualified as BS8 +import Hedgehog hiding (Command) +import Hedgehog.Gen qualified as Gen +import Hedgehog.Range qualified as Range +import PgWire.Connection (simpleQuery) +import PgWire.Error (PgWireError (..)) +import PgWire.MockServer (defaultMockConfig, withMockServer) +import PgWire.Pool + ( Pool + , PoolStats (..) + , closePool + , newPool + , poolIsAlive + , poolStats + , resize + , withResource + ) +import PgWire.Pool.Config (PoolConfig (..), defaultPoolConfig) +import Test.Hspec +import Test.Hspec.Hedgehog (hedgehog) + +data Command + = CmdBurst !Int + | CmdResize !Int + | CmdClose + deriving (Show) + +data Model = Model + { mMaxSize :: !Int + , mClosed :: !Bool + } + deriving (Show) + +initialModel :: Int -> Model +initialModel size = Model {mMaxSize = size, mClosed = False} + +stepModel :: Model -> Command -> Model +stepModel m cmd = case cmd of + CmdBurst _ -> m + CmdResize n -> m {mMaxSize = n} + CmdClose -> m {mClosed = True} + +-- | Burst issues n concurrent acquire/release cycles. Errors are +-- swallowed: if the pool is closed or the acquire times out, the +-- model already reflects that and the invariant pass after the +-- command settles is what matters. +runCommand :: Pool -> Command -> IO () +runCommand pool cmd = case cmd of + CmdBurst n -> + mapConcurrently_ + ( \_ -> do + _ <- + try @PgWireError $ + withResource pool $ \conn -> simpleQuery conn "SELECT 1" + pure () + ) + [1 .. n] + CmdResize n -> resize pool n + CmdClose -> closePool pool + +genCommand :: Gen Command +genCommand = + Gen.frequency + [ (6, CmdBurst <$> Gen.int (Range.linear 1 6)) + , (2, CmdResize <$> Gen.int (Range.linear 1 8)) + , (1, pure CmdClose) + ] + +-- | Generate a program with at most one Close so we don't burn cycles +-- on post-close no-ops. Close may appear at any position or not at all. +genCommands :: Gen [Command] +genCommands = do + len <- Gen.int (Range.linear 1 12) + cmds <- Gen.list (Range.singleton len) (Gen.filterT (not . isClose) genCommand) + includeClose <- Gen.bool + if includeClose + then do + pos <- Gen.int (Range.linear 0 (length cmds)) + let (hd, tl) = splitAt pos cmds + pure (hd <> [CmdClose] <> tl) + else pure cmds + where + isClose CmdClose = True + isClose _ = False + +-- | A per-step observation collected while running the program, then +-- re-checked inside the Hedgehog property. +data Step = Step + { stCmd :: !Command + , stModel :: !Model + , stStats :: !PoolStats + , stAlive :: !Bool + , stClosedProbe :: !(Maybe (Either PgWireError ())) + -- ^ If the model says closed, probe the real pool by trying to + -- acquire a connection. Expected: Left PoolClosed. + } + deriving (Show) + +runProgram :: Pool -> Int -> [Command] -> IO [Step] +runProgram pool initialSize = go (initialModel initialSize) [] + where + go _ acc [] = pure (reverse acc) + go m acc (c : cs) = do + runCommand pool c + let m' = stepModel m c + stats <- poolStats pool + alive <- poolIsAlive pool + probe <- + if mClosed m' + then do + r <- try @PgWireError (withResource pool $ \_ -> pure ()) + pure (Just r) + else pure Nothing + let step = Step c m' stats alive probe + go m' (step : acc) cs + +withMockPool :: Int -> (Pool -> IO a) -> IO a +withMockPool size action = + withMockServer defaultMockConfig $ \port -> do + let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb" + pool <- + newPool + defaultPoolConfig + { poolConnString = connStr + , poolSize = size + , poolAcquireTimeout = 2 + } + r <- action pool + closePool pool -- idempotent; safe even if the program closed it + pure r + +checkStep :: (MonadTest m) => Step -> m () +checkStep Step {stCmd, stModel, stStats, stAlive, stClosedProbe} = do + annotateShow stCmd + annotateShow stModel + annotateShow stStats + stAlive === not (mClosed stModel) + psMaxSize stStats === mMaxSize stModel + -- After a command has settled, nothing should remain in flight. + psInUse stStats === 0 + -- Idle never exceeds the current cap (resize-down drops excess idle + -- eagerly, so even a just-executed shrink should satisfy this). + assert (psIdle stStats <= mMaxSize stModel) + assert (psIdle stStats >= 0) + assert (psWaiters stStats >= 0) + assert (psTotalCreated stStats >= 0) + assert (psTotalDestroyed stStats >= 0) + assert (psTotalTimeouts stStats >= 0) + -- Destroy count can never exceed create count. + assert (psTotalCreated stStats >= psTotalDestroyed stStats) + case stClosedProbe of + Nothing -> pure () + Just (Left PoolClosed) -> pure () + Just other -> do + annotateShow other + failure + +spec :: Spec +spec = do + describe "Pool state machine" $ do + it "model-vs-pool invariants hold after every command" $ hedgehog $ do + initialSize <- forAll (Gen.int (Range.linear 1 6)) + program <- forAll genCommands + steps <- evalIO $ withMockPool initialSize $ \pool -> + runProgram pool initialSize program + mapM_ checkStep steps diff --git a/test/PgWire/Protocol/BuildersSpec.hs b/test/PgWire/Protocol/BuildersSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Protocol/BuildersSpec.hs @@ -0,0 +1,181 @@ +module PgWire.Protocol.BuildersSpec (spec) where + +import Data.ByteString qualified as BS +import Data.Int (Int32) +import Data.Vector qualified as V +import PgWire.Protocol.Builders (buildFrontendMsg, buildStartup) +import PgWire.Protocol.Frontend +import Test.Hspec + +spec :: Spec +spec = do + describe "zero-payload messages" $ do + it "Sync is tag 'S' + length 4" $ do + let bs = buildFrontendMsg Sync + BS.length bs `shouldBe` 5 + BS.index bs 0 `shouldBe` 83 + + it "Terminate is tag 'X' + length 4" $ do + let bs = buildFrontendMsg Terminate + BS.length bs `shouldBe` 5 + BS.index bs 0 `shouldBe` 88 + + it "Flush is tag 'H' + length 4" $ do + let bs = buildFrontendMsg Flush + BS.length bs `shouldBe` 5 + BS.index bs 0 `shouldBe` 72 + + it "CopyDone is tag 'c' + length 4" $ do + let bs = buildFrontendMsg CopyDone + BS.length bs `shouldBe` 5 + BS.index bs 0 `shouldBe` 99 + + describe "Query" $ do + it "has tag 'Q' and NUL-terminated SQL" $ do + let bs = buildFrontendMsg (Query "SELECT 1") + BS.index bs 0 `shouldBe` 81 + BS.last bs `shouldBe` 0 + "SELECT 1" `BS.isInfixOf` bs `shouldBe` True + + it "length includes itself + sql + NUL" $ do + let bs = buildFrontendMsg (Query "SEL") + -- tag(1) + len(4) + "SEL"(3) + NUL(1) = 9 + -- length field = 4 + 3 + 1 = 8 + BS.length bs `shouldBe` 9 + let lenBytes = BS.take 4 (BS.drop 1 bs) + decodeInt32 lenBytes `shouldBe` 8 + + it "handles empty query" $ do + let bs = buildFrontendMsg (Query "") + -- tag(1) + len(4) + NUL(1) = 6 + BS.length bs `shouldBe` 6 + + describe "Parse" $ do + it "has tag 'P'" $ do + let bs = buildFrontendMsg (Parse "s1" "SELECT $1" (V.singleton 23)) + BS.index bs 0 `shouldBe` 80 + + it "contains statement name and SQL" $ do + let bs = buildFrontendMsg (Parse "mystmt" "SELECT * FROM t" V.empty) + "mystmt" `BS.isInfixOf` bs `shouldBe` True + "SELECT * FROM t" `BS.isInfixOf` bs `shouldBe` True + + it "encodes param count and OIDs" $ do + let bs = buildFrontendMsg (Parse "" "SELECT $1, $2" (V.fromList [23, 25])) + -- Should contain the Int16 param count (2) and two Int32 OIDs + BS.length bs `shouldSatisfy` (> 15) + + describe "Bind" $ do + it "has tag 'B'" $ do + let bs = buildFrontendMsg (Bind "" "s1" V.empty V.empty V.empty) + BS.index bs 0 `shouldBe` 66 + + it "encodes portal and statement names" $ do + let bs = buildFrontendMsg (Bind "portal" "stmt" V.empty V.empty V.empty) + "portal" `BS.isInfixOf` bs `shouldBe` True + "stmt" `BS.isInfixOf` bs `shouldBe` True + + it "encodes parameter values including NULLs" $ do + let params = V.fromList [Just "hello", Nothing, Just "world"] + bs = buildFrontendMsg (Bind "" "" V.empty params V.empty) + "hello" `BS.isInfixOf` bs `shouldBe` True + "world" `BS.isInfixOf` bs `shouldBe` True + + describe "Execute" $ do + it "has tag 'E'" $ do + let bs = buildFrontendMsg (Execute "" 0) + BS.index bs 0 `shouldBe` 69 + + it "encodes max rows" $ do + let bs = buildFrontendMsg (Execute "" 100) + BS.length bs `shouldSatisfy` (> 5) + + describe "Describe" $ do + it "has tag 'D' for statement" $ do + let bs = buildFrontendMsg (Describe DescribeStatement "s1") + BS.index bs 0 `shouldBe` 68 + + it "has tag 'D' for portal" $ do + let bs = buildFrontendMsg (Describe DescribePortal "p1") + BS.index bs 0 `shouldBe` 68 + + describe "Close" $ do + it "has tag 'C'" $ do + let bs = buildFrontendMsg (Close DescribeStatement "s1") + BS.index bs 0 `shouldBe` 67 + + describe "PasswordMessage" $ do + it "has tag 'p'" $ do + let bs = buildFrontendMsg (PasswordMessage "secret") + BS.index bs 0 `shouldBe` 112 + "secret" `BS.isInfixOf` bs `shouldBe` True + + describe "CopyData" $ do + it "has tag 'd' and preserves payload" $ do + let bs = buildFrontendMsg (CopyData "row data") + BS.index bs 0 `shouldBe` 100 + "row data" `BS.isInfixOf` bs `shouldBe` True + + describe "CopyFail" $ do + it "has tag 'f' and includes error message" $ do + let bs = buildFrontendMsg (CopyFail "bad data") + BS.index bs 0 `shouldBe` 102 + "bad data" `BS.isInfixOf` bs `shouldBe` True + + describe "SASLInitialResponse" $ do + it "has tag 'p'" $ do + let bs = buildFrontendMsg (SASLInitialResponse "SCRAM-SHA-256" "client-first") + BS.index bs 0 `shouldBe` 112 + "SCRAM-SHA-256" `BS.isInfixOf` bs `shouldBe` True + + describe "SASLResponse" $ do + it "has tag 'p'" $ do + let bs = buildFrontendMsg (SASLResponse "client-final") + BS.index bs 0 `shouldBe` 112 + + describe "buildStartup" $ do + it "starts with length and protocol 3.0" $ do + let bs = buildStartup (StartupParams "u" "d" "" []) + BS.length bs `shouldSatisfy` (> 8) + let version = decodeInt32 (BS.take 4 (BS.drop 4 bs)) + version `shouldBe` 196608 + + it "includes user and database" $ do + let bs = buildStartup (StartupParams "alice" "mydb" "" []) + "user" `BS.isInfixOf` bs `shouldBe` True + "alice" `BS.isInfixOf` bs `shouldBe` True + "database" `BS.isInfixOf` bs `shouldBe` True + "mydb" `BS.isInfixOf` bs `shouldBe` True + + it "includes application_name when non-empty" $ do + let bs = buildStartup (StartupParams "u" "d" "myapp" []) + "application_name" `BS.isInfixOf` bs `shouldBe` True + "myapp" `BS.isInfixOf` bs `shouldBe` True + + it "omits application_name when empty" $ do + let bs = buildStartup (StartupParams "u" "d" "" []) + "application_name" `BS.isInfixOf` bs `shouldBe` False + + it "includes extra params" $ do + let bs = buildStartup (StartupParams "u" "d" "" [("extra_float_digits", "2")]) + "extra_float_digits" `BS.isInfixOf` bs `shouldBe` True + "2" `BS.isInfixOf` bs `shouldBe` True + + it "ends with NUL terminator" $ do + let bs = buildStartup (StartupParams "u" "d" "" []) + BS.last bs `shouldBe` 0 + + it "length field matches actual length" $ do + let bs = buildStartup (StartupParams "testuser" "testdb" "" []) + declaredLen = decodeInt32 (BS.take 4 bs) + declaredLen `shouldBe` fromIntegral (BS.length bs) + +-- Helpers ------------------------------------------------------------------- + +decodeInt32 :: BS.ByteString -> Int32 +decodeInt32 bs = + let b0 = fromIntegral (BS.index bs 0) :: Int32 + b1 = fromIntegral (BS.index bs 1) :: Int32 + b2 = fromIntegral (BS.index bs 2) :: Int32 + b3 = fromIntegral (BS.index bs 3) :: Int32 + in b0 * 256 * 256 * 256 + b1 * 256 * 256 + b2 * 256 + b3 diff --git a/test/PgWire/Protocol/FuzzSpec.hs b/test/PgWire/Protocol/FuzzSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Protocol/FuzzSpec.hs @@ -0,0 +1,360 @@ +{-# OPTIONS_GHC -Wno-orphans #-} + +module PgWire.Protocol.FuzzSpec (spec) where + +import Control.DeepSeq (NFData (..), deepseq) +import Control.Exception (evaluate, try, SomeException) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Builder qualified as B +import Data.ByteString.Lazy qualified as LBS +import Data.Int (Int16, Int32) +import Data.Word (Word8, Word32) +import Hedgehog +import Hedgehog.Gen qualified as Gen +import Hedgehog.Range qualified as Range +import PgWire.Protocol.Backend (CommandTag (..)) +import PgWire.Protocol.Parsers (parseBackendMsg, parseCommandTag) +import Test.Hspec +import Test.Hspec.Hedgehog (hedgehog) + +-- | NFData instance for CommandTag so we can deepseq it. +instance NFData CommandTag where + rnf (SelectTag n) = n `seq` () + rnf (InsertTag n) = n `seq` () + rnf (UpdateTag n) = n `seq` () + rnf (DeleteTag n) = n `seq` () + rnf (OtherTag b) = rnf b + +-- | Tags whose parsers are purely bounds-checked (no unbounded recursion). +-- These are safe to fuzz with completely random payloads even on a small +-- stack (the test suite runs with -K8K). +safeTags :: [Word8] +safeTags = + [ 82 -- R: Authentication + , 75 -- K: BackendKeyData + , 83 -- S: ParameterStatus + , 90 -- Z: ReadyForQuery + , 49 -- 1: ParseComplete + , 50 -- 2: BindComplete + , 51 -- 3: CloseComplete + , 73 -- I: EmptyQueryResponse + , 99 -- c: CopyDone + , 100 -- d: CopyData + , 110 -- n: NoData + , 65 -- A: NotificationResponse + ] + +-- | Tags whose parsers may recurse on their payload (ErrorResponse, +-- NoticeResponse) or iterate based on an encoded count (DataRow, +-- RowDescription, ParameterDescription, CopyIn/OutResponse, +-- CommandComplete). We fuzz these with payloads that contain NUL +-- terminators and bounded counts so they do not overflow the tiny +-- test-suite stack. +recursiveTags :: [Word8] +recursiveTags = + [ 67 -- C: CommandComplete + , 68 -- D: DataRow + , 69 -- E: ErrorResponse + , 71 -- G: CopyInResponse + , 72 -- H: CopyOutResponse + , 78 -- N: NoticeResponse + , 84 -- T: RowDescription + , 116 -- t: ParameterDescription + ] + +-- | Evaluate a parse result to WHNF (forcing the Either constructor). +-- The module uses StrictData and bang patterns internally, so forcing +-- the outer Either exercises the parser. Returns True when the parser +-- produces Left or Right without throwing. +doesNotCrash :: Word8 -> ByteString -> IO Bool +doesNotCrash tag payload = do + result <- try $ evaluate $ case parseBackendMsg tag payload of + Left _ -> () + Right _ -> () + case result of + Left (_ :: SomeException) -> pure False + Right () -> pure True + +spec :: Spec +spec = do + -- ── 1. No crashes on arbitrary input ────────────────────────────── + describe "No crashes on arbitrary input" $ do + it "safe tags with completely random payloads" $ hedgehog $ do + tag <- forAll (Gen.element safeTags) + payload <- forAll (Gen.bytes (Range.linear 0 128)) + ok <- evalIO (doesNotCrash tag payload) + assert ok + + it "unknown tags with random payloads" $ hedgehog $ do + -- Tags outside the known set always return Left immediately + tag <- forAll $ Gen.filter (\t -> t `notElem` safeTags && t `notElem` recursiveTags) + (Gen.word8 Range.linearBounded) + payload <- forAll (Gen.bytes (Range.linear 0 128)) + ok <- evalIO (doesNotCrash tag payload) + assert ok + + it "ErrorResponse with NUL-bearing payloads" $ hedgehog $ do + -- Generate payloads that contain NUL bytes so the field parser terminates + payload <- forAll genNulBearingPayload + ok <- evalIO (doesNotCrash 69 payload) + assert ok + + it "NoticeResponse with NUL-bearing payloads" $ hedgehog $ do + payload <- forAll genNulBearingPayload + ok <- evalIO (doesNotCrash 78 payload) + assert ok + + it "CommandComplete with arbitrary tag strings" $ hedgehog $ do + tagStr <- forAll (Gen.bytes (Range.linear 0 64)) + -- CommandComplete expects a NUL-terminated tag + let payload = tagStr <> BS.singleton 0 + ok <- evalIO (doesNotCrash 67 payload) + assert ok + + it "RowDescription with bounded field count" $ hedgehog $ do + payload <- forAll genRowDescPayload + ok <- evalIO (doesNotCrash 84 payload) + assert ok + + it "ParameterDescription with bounded param count" $ hedgehog $ do + nParams <- forAll (Gen.int16 (Range.linear 0 20)) + trailing <- forAll (Gen.bytes (Range.linear 0 128)) + let payload = int16 nParams <> trailing + ok <- evalIO (doesNotCrash 116 payload) + assert ok + + it "CopyInResponse with bounded column count" $ hedgehog $ do + payload <- forAll genCopyResponsePayload + ok <- evalIO (doesNotCrash 71 payload) + assert ok + + it "CopyOutResponse with bounded column count" $ hedgehog $ do + payload <- forAll genCopyResponsePayload + ok <- evalIO (doesNotCrash 72 payload) + assert ok + + it "medium payloads do not hang (safe tags)" $ hedgehog $ do + tag <- forAll (Gen.element safeTags) + payload <- forAll (Gen.bytes (Range.linear 64 256)) + ok <- evalIO (doesNotCrash tag payload) + assert ok + + -- ── 2. No crashes on truncated messages ─────────────────────────── + describe "No crashes on truncated messages" $ do + it "truncated AuthOk" $ hedgehog $ do + let fullPayload = int32 0 + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 82 (BS.take truncateAt fullPayload)) + assert ok + + it "truncated AuthMD5" $ hedgehog $ do + let fullPayload = int32 5 <> BS.pack [0xDE, 0xAD, 0xBE, 0xEF] + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 82 (BS.take truncateAt fullPayload)) + assert ok + + it "truncated BackendKeyData" $ hedgehog $ do + let fullPayload = int32 12345 <> int32 67890 + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 75 (BS.take truncateAt fullPayload)) + assert ok + + it "truncated ReadyForQuery" $ hedgehog $ do + ok <- evalIO (doesNotCrash 90 BS.empty) + assert ok + + it "truncated DataRow with columns" $ hedgehog $ do + nCols <- forAll (Gen.int (Range.linear 1 10)) + let colPayload = mconcat [int32 3 <> "abc" | _ <- [1 :: Int .. nCols]] + fullPayload = int16 (fromIntegral nCols) <> colPayload + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 68 (BS.take truncateAt fullPayload)) + assert ok + + it "truncated RowDescription" $ hedgehog $ do + let field = "col\0" <> word32 0 <> int16 1 <> word32 23 <> int16 4 <> int32 (-1) <> int16 0 + fullPayload = int16 1 <> field + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 84 (BS.take truncateAt fullPayload)) + assert ok + + it "truncated ParameterDescription" $ hedgehog $ do + nParams <- forAll (Gen.int (Range.linear 1 5)) + let fullPayload = int16 (fromIntegral nParams) + <> mconcat [word32 (fromIntegral i) | i <- [1 :: Int .. nParams]] + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 116 (BS.take truncateAt fullPayload)) + assert ok + + it "truncated CopyInResponse" $ hedgehog $ do + let fullPayload = BS.singleton 0 <> int16 2 <> int16 0 <> int16 1 + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 71 (BS.take truncateAt fullPayload)) + assert ok + + it "truncated ErrorResponse" $ hedgehog $ do + let fullPayload = BS.concat + [ BS.singleton 83, "ERROR\0" + , BS.singleton 67, "42P01\0" + , BS.singleton 77, "relation not found\0" + , BS.singleton 0 + ] + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 69 (BS.take truncateAt fullPayload)) + assert ok + + it "truncated NotificationResponse" $ hedgehog $ do + let fullPayload = int32 42 <> "my_channel\0" <> "hello\0" + truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1))) + ok <- evalIO (doesNotCrash 65 (BS.take truncateAt fullPayload)) + assert ok + + -- ── 3. CommandTag fuzz ──────────────────────────────────────────── + describe "CommandTag fuzz" $ do + it "parseCommandTag handles arbitrary bytestrings" $ hedgehog $ do + bs <- forAll (Gen.bytes (Range.linear 0 256)) + ok <- evalIO (commandTagDoesNotCrash bs) + assert ok + + it "parseCommandTag handles SELECT prefix with garbage" $ hedgehog $ do + suffix <- forAll (Gen.bytes (Range.linear 0 64)) + ok <- evalIO (commandTagDoesNotCrash ("SELECT " <> suffix)) + assert ok + + it "parseCommandTag handles INSERT prefix with garbage" $ hedgehog $ do + suffix <- forAll (Gen.bytes (Range.linear 0 64)) + ok <- evalIO (commandTagDoesNotCrash ("INSERT " <> suffix)) + assert ok + + it "parseCommandTag handles UPDATE prefix with garbage" $ hedgehog $ do + suffix <- forAll (Gen.bytes (Range.linear 0 64)) + ok <- evalIO (commandTagDoesNotCrash ("UPDATE " <> suffix)) + assert ok + + it "parseCommandTag handles DELETE prefix with garbage" $ hedgehog $ do + suffix <- forAll (Gen.bytes (Range.linear 0 64)) + ok <- evalIO (commandTagDoesNotCrash ("DELETE " <> suffix)) + assert ok + + -- ── 4. DataRow bounds safety ────────────────────────────────────── + describe "DataRow bounds safety" $ do + it "well-formed random columns" $ hedgehog $ do + nCols <- forAll (Gen.int (Range.linear 0 50)) + cols <- forAll (Gen.list (Range.singleton nCols) genColumnData) + let payload = int16 (fromIntegral nCols) <> mconcat cols + ok <- evalIO (doesNotCrash 68 payload) + assert ok + + it "declared column count exceeding actual data" $ hedgehog $ do + declaredCols <- forAll (Gen.int (Range.linear 1 100)) + actualCols <- forAll (Gen.int (Range.linear 0 (declaredCols - 1))) + cols <- forAll (Gen.list (Range.singleton actualCols) genColumnData) + let payload = int16 (fromIntegral declaredCols) <> mconcat cols + ok <- evalIO (doesNotCrash 68 payload) + assert ok + + it "zero-length columns" $ hedgehog $ do + nCols <- forAll (Gen.int (Range.linear 1 20)) + let payload = int16 (fromIntegral nCols) + <> mconcat (replicate nCols (int32 0)) + ok <- evalIO (doesNotCrash 68 payload) + assert ok + + it "all NULL columns" $ hedgehog $ do + nCols <- forAll (Gen.int (Range.linear 1 50)) + let payload = int16 (fromIntegral nCols) + <> mconcat (replicate nCols (int32 (-1))) + ok <- evalIO (doesNotCrash 68 payload) + assert ok + + it "very large declared column length" $ hedgehog $ do + bigLen <- forAll (Gen.int32 (Range.linear 1000 2147483647)) + let payload = int16 1 <> int32 bigLen <> "tiny" + ok <- evalIO (doesNotCrash 68 payload) + assert ok + + it "mixed valid and truncated columns" $ hedgehog $ do + nValid <- forAll (Gen.int (Range.linear 1 5)) + validCols <- forAll (Gen.list (Range.singleton nValid) genColumnData) + -- Append a partial column (length header but no data) + partialLen <- forAll (Gen.int32 (Range.linear 1 100)) + let payload = int16 (fromIntegral (nValid + 1)) + <> mconcat validCols + <> int32 partialLen + ok <- evalIO (doesNotCrash 68 payload) + assert ok + + it "DataRow with random trailing garbage" $ hedgehog $ do + nCols <- forAll (Gen.int (Range.linear 1 5)) + cols <- forAll (Gen.list (Range.singleton nCols) genColumnData) + garbage <- forAll (Gen.bytes (Range.linear 0 64)) + let payload = int16 (fromIntegral nCols) <> mconcat cols <> garbage + ok <- evalIO (doesNotCrash 68 payload) + assert ok + +-- | Evaluate parseCommandTag, fully forcing the result via deepseq. +commandTagDoesNotCrash :: ByteString -> IO Bool +commandTagDoesNotCrash bs = do + result <- try $ evaluate (parseCommandTag bs `deepseq` ()) + case result of + Left (_ :: SomeException) -> pure False + Right () -> pure True + +-- Generators ------------------------------------------------------------------ + +-- | Generate a payload with embedded NUL bytes so recursive parsers +-- (ErrorResponse / NoticeResponse) terminate promptly. +genNulBearingPayload :: Gen ByteString +genNulBearingPayload = do + -- Build a sequence of (fieldType, value) pairs terminated by NUL + nFields <- Gen.int (Range.linear 0 10) + fields <- Gen.list (Range.singleton nFields) $ do + fieldType <- Gen.word8 Range.linearBounded + value <- Gen.bytes (Range.linear 0 32) + -- Ensure value doesn't contain NUL (so NUL terminators are clear) + let safeValue = BS.map (\b -> if b == 0 then 1 else b) value + pure $ BS.singleton fieldType <> safeValue <> BS.singleton 0 + pure $ mconcat fields <> BS.singleton 0 + +-- | Generate a single DataRow column encoding: either NULL (-1) or +-- a length-prefixed bytestring. +genColumnData :: Gen ByteString +genColumnData = Gen.choice + [ pure (int32 (-1)) -- NULL + , do + colBytes <- Gen.bytes (Range.linear 0 128) + pure (int32 (fromIntegral (BS.length colBytes)) <> colBytes) + ] + +-- | Generate a RowDescription payload with a bounded field count. +genRowDescPayload :: Gen ByteString +genRowDescPayload = do + nFields <- Gen.int16 (Range.linear 0 5) + fields <- Gen.list (Range.singleton (fromIntegral nFields)) $ do + -- name (NUL-terminated) + 18 bytes of field metadata + name <- Gen.bytes (Range.linear 1 16) + let safeName = BS.map (\b -> if b == 0 then 65 else b) name + metadata <- Gen.bytes (Range.singleton 18) + pure $ safeName <> BS.singleton 0 <> metadata + pure $ int16 nFields <> mconcat fields + +-- | Generate a CopyInResponse / CopyOutResponse payload. +genCopyResponsePayload :: Gen ByteString +genCopyResponsePayload = do + fmt <- Gen.word8 (Range.linear 0 1) + nCols <- Gen.int16 (Range.linear 0 10) + colFmts <- Gen.list (Range.singleton (fromIntegral nCols)) $ + Gen.bytes (Range.singleton 2) + pure $ BS.singleton fmt <> int16 nCols <> mconcat colFmts + +-- Helpers --------------------------------------------------------------------- + +int16 :: Int16 -> ByteString +int16 n = LBS.toStrict . B.toLazyByteString $ B.int16BE n + +int32 :: Int32 -> ByteString +int32 n = LBS.toStrict . B.toLazyByteString $ B.int32BE n + +word32 :: Word32 -> ByteString +word32 n = LBS.toStrict . B.toLazyByteString $ B.word32BE n diff --git a/test/PgWire/Protocol/OidSpec.hs b/test/PgWire/Protocol/OidSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Protocol/OidSpec.hs @@ -0,0 +1,78 @@ +module PgWire.Protocol.OidSpec (spec) where + +import PgWire.Protocol.Oid +import Test.Hspec + +spec :: Spec +spec = do + describe "scalar OID constants" $ do + it "oidBool = 16" $ unOid oidBool `shouldBe` 16 + it "oidBytea = 17" $ unOid oidBytea `shouldBe` 17 + it "oidInt2 = 21" $ unOid oidInt2 `shouldBe` 21 + it "oidInt4 = 23" $ unOid oidInt4 `shouldBe` 23 + it "oidInt8 = 20" $ unOid oidInt8 `shouldBe` 20 + it "oidText = 25" $ unOid oidText `shouldBe` 25 + it "oidJson = 114" $ unOid oidJson `shouldBe` 114 + it "oidFloat4 = 700" $ unOid oidFloat4 `shouldBe` 700 + it "oidFloat8 = 701" $ unOid oidFloat8 `shouldBe` 701 + it "oidVarchar = 1043" $ unOid oidVarchar `shouldBe` 1043 + it "oidDate = 1082" $ unOid oidDate `shouldBe` 1082 + it "oidTime = 1083" $ unOid oidTime `shouldBe` 1083 + it "oidTimestamp = 1114" $ unOid oidTimestamp `shouldBe` 1114 + it "oidTimestamptz = 1184" $ unOid oidTimestamptz `shouldBe` 1184 + it "oidNumeric = 1700" $ unOid oidNumeric `shouldBe` 1700 + it "oidUuid = 2950" $ unOid oidUuid `shouldBe` 2950 + it "oidJsonb = 3802" $ unOid oidJsonb `shouldBe` 3802 + it "oidInterval = 1186" $ unOid oidInterval `shouldBe` 1186 + + describe "array OID constants" $ do + it "oidBoolArray = 1000" $ unOid oidBoolArray `shouldBe` 1000 + it "oidByteaArray = 1001" $ unOid oidByteaArray `shouldBe` 1001 + it "oidInt2Array = 1005" $ unOid oidInt2Array `shouldBe` 1005 + it "oidInt4Array = 1007" $ unOid oidInt4Array `shouldBe` 1007 + it "oidInt8Array = 1016" $ unOid oidInt8Array `shouldBe` 1016 + it "oidTextArray = 1009" $ unOid oidTextArray `shouldBe` 1009 + it "oidVarcharArray = 1015" $ unOid oidVarcharArray `shouldBe` 1015 + it "oidFloat4Array = 1021" $ unOid oidFloat4Array `shouldBe` 1021 + it "oidFloat8Array = 1022" $ unOid oidFloat8Array `shouldBe` 1022 + it "oidTimestampArray = 1115" $ unOid oidTimestampArray `shouldBe` 1115 + it "oidTimestamptzArray = 1185" $ unOid oidTimestamptzArray `shouldBe` 1185 + it "oidDateArray = 1182" $ unOid oidDateArray `shouldBe` 1182 + it "oidTimeArray = 1183" $ unOid oidTimeArray `shouldBe` 1183 + it "oidUuidArray = 2951" $ unOid oidUuidArray `shouldBe` 2951 + it "oidJsonArray = 199" $ unOid oidJsonArray `shouldBe` 199 + it "oidJsonbArray = 3807" $ unOid oidJsonbArray `shouldBe` 3807 + + describe "isArrayOid" $ do + it "True for all 16 array OIDs" $ do + let arrayOids = [1000, 1001, 1005, 1007, 1009, 1015, 1016, + 1021, 1022, 1115, 1182, 1183, 1185, 2951, 199, 3807] + mapM_ (\o -> isArrayOid (Oid o) `shouldBe` True) arrayOids + + it "False for all scalar OIDs" $ do + let scalarOids = [16, 17, 20, 21, 23, 25, 114, 700, 701, + 1043, 1082, 1083, 1114, 1184, 1186, 1700, 2950, 3802] + mapM_ (\o -> isArrayOid (Oid o) `shouldBe` False) scalarOids + + it "False for unknown OID" $ + isArrayOid (Oid 99999) `shouldBe` False + + describe "arrayElementOid" $ do + it "bool[] → bool" $ arrayElementOid oidBoolArray `shouldBe` Just oidBool + it "bytea[] → bytea" $ arrayElementOid oidByteaArray `shouldBe` Just oidBytea + it "int2[] → int2" $ arrayElementOid oidInt2Array `shouldBe` Just oidInt2 + it "int4[] → int4" $ arrayElementOid oidInt4Array `shouldBe` Just oidInt4 + it "int8[] → int8" $ arrayElementOid oidInt8Array `shouldBe` Just oidInt8 + it "text[] → text" $ arrayElementOid oidTextArray `shouldBe` Just oidText + it "varchar[] → varchar" $ arrayElementOid oidVarcharArray `shouldBe` Just oidVarchar + it "float4[] → float4" $ arrayElementOid oidFloat4Array `shouldBe` Just oidFloat4 + it "float8[] → float8" $ arrayElementOid oidFloat8Array `shouldBe` Just oidFloat8 + it "timestamp[] → timestamp" $ arrayElementOid oidTimestampArray `shouldBe` Just oidTimestamp + it "timestamptz[] → timestamptz" $ arrayElementOid oidTimestamptzArray `shouldBe` Just oidTimestamptz + it "date[] → date" $ arrayElementOid oidDateArray `shouldBe` Just oidDate + it "time[] → time" $ arrayElementOid oidTimeArray `shouldBe` Just oidTime + it "uuid[] → uuid" $ arrayElementOid oidUuidArray `shouldBe` Just oidUuid + it "json[] → json" $ arrayElementOid oidJsonArray `shouldBe` Just oidJson + it "jsonb[] → jsonb" $ arrayElementOid oidJsonbArray `shouldBe` Just oidJsonb + it "scalar → Nothing" $ arrayElementOid oidInt4 `shouldBe` Nothing + it "unknown → Nothing" $ arrayElementOid (Oid 99999) `shouldBe` Nothing diff --git a/test/PgWire/Protocol/ParsersSpec.hs b/test/PgWire/Protocol/ParsersSpec.hs new file mode 100644 --- /dev/null +++ b/test/PgWire/Protocol/ParsersSpec.hs @@ -0,0 +1,363 @@ +module PgWire.Protocol.ParsersSpec (spec) where + +import Data.ByteString qualified as BS +import Data.ByteString.Builder qualified as B +import Data.ByteString.Lazy qualified as LBS +import Data.Int (Int16, Int32) +import Data.Vector qualified as V +import Data.Word (Word32) +import PgWire.Protocol.Backend +import PgWire.Protocol.Parsers (parseBackendMsg, parseCommandTag) +import Test.Hspec + +spec :: Spec +spec = do + -- ── Authentication ────────────────────────────────────────────── + describe "Authentication" $ do + it "parses AuthOk" $ + parseBackendMsg 82 (int32 0) `shouldBe` Right (Authentication AuthOk) + + it "parses AuthCleartextPassword" $ + parseBackendMsg 82 (int32 3) `shouldBe` Right (Authentication AuthCleartextPassword) + + it "parses AuthMD5Password with 4-byte salt" $ do + let salt = BS.pack [0xDE, 0xAD, 0xBE, 0xEF] + parseBackendMsg 82 (int32 5 <> salt) + `shouldBe` Right (Authentication (AuthMD5Password salt)) + + it "parses AuthSASL with mechanism list" $ do + let payload = int32 10 <> "SCRAM-SHA-256\0\0" + case parseBackendMsg 82 payload of + Right (Authentication (AuthSASL mechs)) -> + mechs `shouldBe` ["SCRAM-SHA-256"] + other -> expectationFailure $ "Expected AuthSASL, got: " <> show other + + it "parses AuthSASLContinue" $ do + let serverData = "r=nonce,s=salt,i=4096" + parseBackendMsg 82 (int32 11 <> serverData) + `shouldBe` Right (Authentication (AuthSASLContinue serverData)) + + it "parses AuthSASLFinal" $ do + let serverData = "v=signature" + parseBackendMsg 82 (int32 12 <> serverData) + `shouldBe` Right (Authentication (AuthSASLFinal serverData)) + + it "rejects unknown auth type" $ + parseBackendMsg 82 (int32 999) `shouldSatisfy` isLeft + + -- ── ReadyForQuery ─────────────────────────────────────────────── + describe "ReadyForQuery" $ do + it "idle" $ parseBackendMsg 90 (BS.singleton 73) `shouldBe` Right (ReadyForQuery TxIdle) + it "in transaction" $ parseBackendMsg 90 (BS.singleton 84) `shouldBe` Right (ReadyForQuery TxInTransaction) + it "failed" $ parseBackendMsg 90 (BS.singleton 69) `shouldBe` Right (ReadyForQuery TxFailed) + it "rejects unknown status byte" $ parseBackendMsg 90 (BS.singleton 0) `shouldSatisfy` isLeft + it "rejects empty payload" $ parseBackendMsg 90 BS.empty `shouldSatisfy` isLeft + + -- ── Simple response messages ──────────────────────────────────── + describe "Simple responses" $ do + it "ParseComplete" $ parseBackendMsg 49 BS.empty `shouldBe` Right ParseComplete + it "BindComplete" $ parseBackendMsg 50 BS.empty `shouldBe` Right BindComplete + it "CloseComplete" $ parseBackendMsg 51 BS.empty `shouldBe` Right CloseComplete + it "NoData" $ parseBackendMsg 110 BS.empty `shouldBe` Right NoData + it "EmptyQueryResponse" $ parseBackendMsg 73 BS.empty `shouldBe` Right EmptyQueryResponse + it "CopyDoneMsg" $ parseBackendMsg 99 BS.empty `shouldBe` Right CopyDoneMsg + + -- ── CopyData ──────────────────────────────────────────────────── + describe "CopyDataMsg" $ do + it "preserves payload" $ + parseBackendMsg 100 "hello world" `shouldBe` Right (CopyDataMsg "hello world") + it "handles empty payload" $ + parseBackendMsg 100 BS.empty `shouldBe` Right (CopyDataMsg BS.empty) + + -- ── DataRow (hot path) ────────────────────────────────────────── + describe "DataRow" $ do + it "parses zero columns" $ do + let payload = int16 0 + parseBackendMsg 68 payload `shouldBe` Right (DataRow V.empty) + + it "parses single non-null column" $ do + let payload = int16 1 <> int32 5 <> "hello" + case parseBackendMsg 68 payload of + Right (DataRow v) -> do + V.length v `shouldBe` 1 + v V.! 0 `shouldBe` Just "hello" + other -> expectationFailure $ show other + + it "parses single NULL column" $ do + let payload = int16 1 <> int32 (-1) + case parseBackendMsg 68 payload of + Right (DataRow v) -> do + V.length v `shouldBe` 1 + v V.! 0 `shouldBe` Nothing + other -> expectationFailure $ show other + + it "parses multiple columns with mixed nulls" $ do + let payload = int16 3 + <> int32 2 <> "ab" -- col 0: "ab" + <> int32 (-1) -- col 1: NULL + <> int32 4 <> int32 42 -- col 2: 4 bytes (an int32) + case parseBackendMsg 68 payload of + Right (DataRow v) -> do + V.length v `shouldBe` 3 + v V.! 0 `shouldBe` Just "ab" + v V.! 1 `shouldBe` Nothing + v V.! 2 `shouldBe` Just (int32 42) + other -> expectationFailure $ show other + + it "parses empty string column (length 0)" $ do + let payload = int16 1 <> int32 0 + case parseBackendMsg 68 payload of + Right (DataRow v) -> do + V.length v `shouldBe` 1 + v V.! 0 `shouldBe` Just BS.empty + other -> expectationFailure $ show other + + it "parses many columns" $ do + let n = 20 + colPayload = int32 1 <> "x" + payload = int16 n <> BS.concat (replicate (fromIntegral n) colPayload) + case parseBackendMsg 68 payload of + Right (DataRow v) -> V.length v `shouldBe` fromIntegral n + other -> expectationFailure $ show other + + it "fails on truncated column data" $ + parseBackendMsg 68 (int16 1 <> int32 100 <> "short") `shouldSatisfy` isLeft + + it "fails on truncated header" $ + parseBackendMsg 68 (BS.singleton 0) `shouldSatisfy` isLeft + + -- ── RowDescription ────────────────────────────────────────────── + describe "RowDescription" $ do + it "parses zero fields" $ do + let payload = int16 0 + parseBackendMsg 84 payload `shouldBe` Right (RowDescription V.empty) + + it "parses single field" $ do + let payload = int16 1 + <> "id\0" -- name (NUL-terminated) + <> word32 16384 -- table OID + <> int16 1 -- column number + <> word32 23 -- type OID (int4) + <> int16 4 -- type size + <> int32 (-1) -- type modifier + <> int16 0 -- format code (text) + case parseBackendMsg 84 payload of + Right (RowDescription v) -> do + V.length v `shouldBe` 1 + let fi = v V.! 0 + fiName fi `shouldBe` "id" + fiTableOid fi `shouldBe` 16384 + fiColumnNum fi `shouldBe` 1 + fiTypeOid fi `shouldBe` 23 + fiTypeSize fi `shouldBe` 4 + fiFormatCode fi `shouldBe` 0 + other -> expectationFailure $ show other + + it "parses multiple fields" $ do + let mkField name oid = + name <> "\0" + <> word32 0 <> int16 0 <> word32 oid + <> int16 (-1) <> int32 (-1) <> int16 1 + payload = int16 3 + <> mkField "id" 23 + <> mkField "name" 25 + <> mkField "email" 25 + case parseBackendMsg 84 payload of + Right (RowDescription v) -> do + V.length v `shouldBe` 3 + fiName (v V.! 0) `shouldBe` "id" + fiName (v V.! 1) `shouldBe` "name" + fiName (v V.! 2) `shouldBe` "email" + fiTypeOid (v V.! 0) `shouldBe` 23 + fiTypeOid (v V.! 1) `shouldBe` 25 + other -> expectationFailure $ show other + + -- ── ParameterDescription ──────────────────────────────────────── + describe "ParameterDescription" $ do + it "parses zero params" $ + parseBackendMsg 116 (int16 0) `shouldBe` Right (ParameterDescription V.empty) + + it "parses multiple param OIDs" $ do + let payload = int16 2 <> word32 23 <> word32 25 + case parseBackendMsg 116 payload of + Right (ParameterDescription v) -> do + V.length v `shouldBe` 2 + v V.! 0 `shouldBe` 23 + v V.! 1 `shouldBe` 25 + other -> expectationFailure $ show other + + -- ── ParameterStatus ───────────────────────────────────────────── + describe "ParameterStatus" $ do + it "parses name=value pair" $ do + let payload = "server_version\0" <> "16.2\0" + parseBackendMsg 83 payload + `shouldBe` Right (ParameterStatus "server_version" "16.2") + + it "handles empty value" $ do + let payload = "param\0\0" + parseBackendMsg 83 payload + `shouldBe` Right (ParameterStatus "param" "") + + -- ── BackendKeyData ────────────────────────────────────────────── + describe "BackendKeyData" $ do + it "parses PID and key" $ do + let payload = int32 12345 <> int32 67890 + parseBackendMsg 75 payload `shouldBe` Right (BackendKeyData 12345 67890) + + it "handles zero values" $ + parseBackendMsg 75 (int32 0 <> int32 0) `shouldBe` Right (BackendKeyData 0 0) + + -- ── NotificationResponse ──────────────────────────────────────── + describe "NotificationResponse" $ do + it "parses PID, channel, payload" $ do + let payload = int32 42 <> "my_channel\0" <> "hello\0" + parseBackendMsg 65 payload + `shouldBe` Right (NotificationResponse 42 "my_channel" "hello") + + it "handles empty payload string" $ do + let payload = int32 1 <> "ch\0" <> "\0" + parseBackendMsg 65 payload + `shouldBe` Right (NotificationResponse 1 "ch" "") + + -- ── CopyInResponse / CopyOutResponse ──────────────────────────── + describe "CopyInResponse" $ do + it "parses format and column formats" $ do + let payload = BS.singleton 0 -- overall format: text + <> int16 2 -- 2 columns + <> int16 0 -- col 0: text + <> int16 1 -- col 1: binary + case parseBackendMsg 71 payload of + Right (CopyInResponse fmt cols) -> do + fmt `shouldBe` 0 + V.toList cols `shouldBe` [0, 1] + other -> expectationFailure $ show other + + describe "CopyOutResponse" $ do + it "parses format and column formats" $ do + let payload = BS.singleton 1 <> int16 1 <> int16 1 + case parseBackendMsg 72 payload of + Right (CopyOutResponse fmt cols) -> do + fmt `shouldBe` 1 + V.toList cols `shouldBe` [1] + other -> expectationFailure $ show other + + -- ── CommandComplete ───────────────────────────────────────────── + describe "CommandComplete" $ do + it "parses SELECT tag" $ do + let payload = "SELECT 42\0" + parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (SelectTag 42)) + + it "parses INSERT tag (with OID)" $ do + let payload = "INSERT 0 5\0" + parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (InsertTag 5)) + + it "parses UPDATE tag" $ do + let payload = "UPDATE 100\0" + parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (UpdateTag 100)) + + it "parses DELETE tag" $ do + let payload = "DELETE 0\0" + parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (DeleteTag 0)) + + it "parses CREATE TABLE as OtherTag" $ do + let payload = "CREATE TABLE\0" + parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (OtherTag "CREATE TABLE")) + + describe "parseCommandTag" $ do + it "SELECT 0" $ parseCommandTag "SELECT 0" `shouldBe` SelectTag 0 + it "SELECT 999999" $ parseCommandTag "SELECT 999999" `shouldBe` SelectTag 999999 + it "INSERT 0 1" $ parseCommandTag "INSERT 0 1" `shouldBe` InsertTag 1 + it "UPDATE 1" $ parseCommandTag "UPDATE 1" `shouldBe` UpdateTag 1 + it "DELETE 50" $ parseCommandTag "DELETE 50" `shouldBe` DeleteTag 50 + it "COPY 100" $ parseCommandTag "COPY 100" `shouldBe` OtherTag "COPY 100" + it "BEGIN" $ parseCommandTag "BEGIN" `shouldBe` OtherTag "BEGIN" + + -- ── ErrorResponse / NoticeResponse ────────────────────────────── + describe "ErrorResponse" $ do + it "parses all standard fields" $ do + let payload = BS.concat + [ BS.singleton 83, "ERROR\0" -- S: severity + , BS.singleton 67, "42P01\0" -- C: code + , BS.singleton 77, "relation not found\0" -- M: message + , BS.singleton 68, "some detail\0" -- D: detail + , BS.singleton 72, "try this\0" -- H: hint + , BS.singleton 80, "15\0" -- P: position + , BS.singleton 87, "at function foo\0" -- W: where + , BS.singleton 115, "public\0" -- s: schema + , BS.singleton 116, "users\0" -- t: table + , BS.singleton 99, "email\0" -- c: column + , BS.singleton 100, "text\0" -- d: data type + , BS.singleton 110, "users_pkey\0" -- n: constraint + , BS.singleton 70, "parse.c\0" -- F: file + , BS.singleton 76, "42\0" -- L: line + , BS.singleton 82, "exec_query\0" -- R: routine + , BS.singleton 0 -- terminator + ] + case parseBackendMsg 69 payload of + Right (ErrorResponse err) -> do + pgSeverity err `shouldBe` "ERROR" + pgCode err `shouldBe` "42P01" + pgMessage err `shouldBe` "relation not found" + pgDetail err `shouldBe` Just "some detail" + pgHint err `shouldBe` Just "try this" + pgPosition err `shouldBe` Just 15 + pgWhere err `shouldBe` Just "at function foo" + pgSchema err `shouldBe` Just "public" + pgTable err `shouldBe` Just "users" + pgColumn err `shouldBe` Just "email" + pgDataType err `shouldBe` Just "text" + pgConstraint err `shouldBe` Just "users_pkey" + pgFile err `shouldBe` Just "parse.c" + pgLine err `shouldBe` Just 42 + pgRoutine err `shouldBe` Just "exec_query" + other -> expectationFailure $ show other + + it "handles minimal error (severity + code + message only)" $ do + let payload = BS.concat + [ BS.singleton 83, "ERROR\0" + , BS.singleton 67, "00000\0" + , BS.singleton 77, "ok\0" + , BS.singleton 0 + ] + case parseBackendMsg 69 payload of + Right (ErrorResponse err) -> do + pgDetail err `shouldBe` Nothing + pgHint err `shouldBe` Nothing + pgPosition err `shouldBe` Nothing + pgSchema err `shouldBe` Nothing + other -> expectationFailure $ show other + + describe "NoticeResponse" $ do + it "parses like ErrorResponse wrapped in PgNotice" $ do + let payload = BS.concat + [ BS.singleton 83, "WARNING\0" + , BS.singleton 67, "01000\0" + , BS.singleton 77, "be careful\0" + , BS.singleton 0 + ] + case parseBackendMsg 78 payload of + Right (NoticeResponse (PgNotice err)) -> do + pgSeverity err `shouldBe` "WARNING" + pgMessage err `shouldBe` "be careful" + other -> expectationFailure $ show other + + -- ── Unknown tags ──────────────────────────────────────────────── + describe "Unknown tags" $ do + it "rejects tag 0" $ parseBackendMsg 0 BS.empty `shouldSatisfy` isLeft + it "rejects tag 255" $ parseBackendMsg 255 BS.empty `shouldSatisfy` isLeft + it "rejects tag 42" $ parseBackendMsg 42 BS.empty `shouldSatisfy` isLeft + +-- Helpers ------------------------------------------------------------------- + +int16 :: Int16 -> BS.ByteString +int16 n = LBS.toStrict . B.toLazyByteString $ B.int16BE n + +int32 :: Int32 -> BS.ByteString +int32 n = LBS.toStrict . B.toLazyByteString $ B.int32BE n + +word32 :: Word32 -> BS.ByteString +word32 n = LBS.toStrict . B.toLazyByteString $ B.word32BE n + +isLeft :: Either a b -> Bool +isLeft (Left _) = True +isLeft _ = False