diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,22 @@ +Copyright (c) 2026, Nikita Volkov + +Permission is hereby granted, free of charge, to any person +obtaining a copy of this software and associated documentation +files (the "Software"), to deal in the Software without +restriction, including without limitation the rights to use, +copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES +OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT +HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, +WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR +OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,29 @@ +# pqi-native + +[![Hackage](https://img.shields.io/hackage/v/pqi-native.svg)](https://hackage.haskell.org/package/pqi-native) +[![Continuous Haddock](https://img.shields.io/badge/haddock-master-blue)](https://nikita-volkov.github.io/pqi-native/) + +A pure-Haskell [`pqi`](https://github.com/nikita-volkov/pqi) adapter +that speaks the PostgreSQL frontend/backend wire protocol directly — no +dependency on the C `libpq` library. + +`pqi-native` is an LLM-generated port of the PostgreSQL C client library, [`libpq`](https://www.postgresql.org/docs/current/libpq.html). The upstream [`libpq` source](https://github.com/postgres/postgres/tree/master/src/interfaces/libpq) is the direct reference for the implementation. + +## Fidelity goal + +The goal is **byte-identical output to `libpq`** for every protocol-derived +value — error message strings, notice text, result status, field metadata, +cell data, and all structured error fields. Fidelity is continuously enforced +by [`pqi-conformance`](https://github.com/nikita-volkov/pqi-conformance), which +runs every operation on both this adapter and a direct +[`postgresql-libpq`](https://hackage.haskell.org/package/postgresql-libpq) +reference connection against the same database and asserts exact equality. + +## Status + +All classes are implemented. Verified against the `postgresql-libpq` reference +via the conformance differential suite. + +Authentication: **trust**, **MD5**, and **SCRAM-SHA-256** are implemented. SCRAM +is verified against a password-auth PostgreSQL 17 container (which defaults to +`scram-sha-256`). diff --git a/pqi-native.cabal b/pqi-native.cabal new file mode 100644 --- /dev/null +++ b/pqi-native.cabal @@ -0,0 +1,154 @@ +cabal-version: 3.0 +name: pqi-native +version: 0.0.1.0 +category: Database, PostgreSQL +synopsis: Native (pure-Haskell) adapter for pqi +description: + A @pqi@ adapter that speaks the PostgreSQL frontend\/backend wire + protocol directly, with no dependency on the C @libpq@ library. The transport + layer (framing, message serialization\/deserialization, socket I\/O) is + isolated in an internal @transport@ sub-library built on @ptr-poker@ and + @ptr-peeker@; the public library is an abstraction over it that implements the + @pqi@ capability classes. + +homepage: https://github.com/nikita-volkov/pqi-native +bug-reports: https://github.com/nikita-volkov/pqi-native/issues +author: Nikita Volkov +maintainer: Nikita Volkov +copyright: (c) 2026, Nikita Volkov +license: MIT +license-file: LICENSE +extra-doc-files: + CHANGELOG.md + LICENSE + README.md + +source-repository head + type: git + location: https://github.com/nikita-volkov/pqi-native + +common base + default-language: Haskell2010 + default-extensions: + ApplicativeDo + BangPatterns + BinaryLiterals + BlockArguments + ConstraintKinds + DataKinds + DefaultSignatures + DeriveDataTypeable + DeriveFoldable + DeriveFunctor + DeriveTraversable + DerivingStrategies + DerivingVia + DuplicateRecordFields + EmptyDataDecls + ExistentialQuantification + FlexibleContexts + FlexibleInstances + FunctionalDependencies + GADTs + GeneralizedNewtypeDeriving + LambdaCase + LiberalTypeSynonyms + MagicHash + MultiParamTypeClasses + MultiWayIf + NamedFieldPuns + NoFieldSelectors + NoImplicitPrelude + NoMonomorphismRestriction + NumericUnderscores + OverloadedRecordDot + OverloadedStrings + ParallelListComp + PatternGuards + QuasiQuotes + RankNTypes + RecordWildCards + ScopedTypeVariables + StandaloneDeriving + StrictData + TemplateHaskell + TupleSections + TypeApplications + TypeFamilies + TypeOperators + UnboxedTuples + ViewPatterns + +common test + import: base + ghc-options: + -threaded + -with-rtsopts=-N + +library comms + import: base + visibility: public + hs-source-dirs: src/comms + exposed-modules: + Pqi.Native.Comms + + build-depends: + base >=4.11 && <5, + bytestring >=0.10 && <0.13, + ptr-peeker ^>=0.2, + ptr-poker ^>=0.1.3, + transformers >=0.5 && <0.7, + +library transport + import: base + visibility: public + hs-source-dirs: src/transport + exposed-modules: + Pqi.Native.Transport + Pqi.Native.Transport.Message + + other-modules: + Pqi.Native.Transport.Prelude + + build-depends: + base >=4.11 && <5, + bytestring >=0.10 && <0.13, + network >=3.1 && <3.3, + pqi-native:comms, + ptr-poker ^>=0.1.3, + transformers >=0.5 && <0.7, + +library + import: base + hs-source-dirs: src/library + exposed-modules: + Pqi.Native + + other-modules: + Pqi.Native.Auth + Pqi.Native.Connection + Pqi.Native.LargeObject + Pqi.Native.Prelude + Pqi.Native.Query + Pqi.Native.Types + + build-depends: + base >=4.11 && <5, + base64-bytestring >=1.2 && <1.3, + bytestring >=0.10 && <0.13, + containers >=0.6 && <0.9, + crypton >=0.34 && <2, + pqi ^>=0, + pqi-native:transport, + ptr-poker ^>=0.1.3, + +test-suite native-test + import: test + type: exitcode-stdio-1.0 + hs-source-dirs: src/test + main-is: Spec.hs + build-depends: + base >=4.11 && <5, + hspec >=2.11 && <2.12, + pqi-conformance, + pqi-native, diff --git a/src/comms/Pqi/Native/Comms.hs b/src/comms/Pqi/Native/Comms.hs new file mode 100644 --- /dev/null +++ b/src/comms/Pqi/Native/Comms.hs @@ -0,0 +1,121 @@ +-- | The serialization abstraction shared by the wire-protocol messages: a type +-- that can be both written (via @ptr-poker@) and decoded (via @ptr-peeker@, +-- with semantic failures carried in an 'ExceptT' layer, since peekers +-- themselves have no failure mechanism). +module Pqi.Native.Comms + ( Comms (..), + Decoder, + DecodingError (..), + runDecoder, + + -- * Decoder primitives + liftVariable, + liftFixed, + int16, + int32, + word8, + word32, + cstring, + bytes, + remaining, + + -- * Serializable primitives + CString (..), + ) +where + +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.Except (ExceptT (..), runExceptT) +import Data.ByteString (ByteString) +import Data.Int (Int16, Int32) +import Data.Word (Word32, Word8) +import qualified PtrPeeker as Peeker +import qualified PtrPoker.Write as Poker +import Prelude + +-- | A semantic decoding failure (as opposed to a \"need more bytes\" framing +-- shortfall, which the transport handles separately). +data DecodingError + = -- | A backend message had an unrecognized type byte. + UnexpectedMessageType Word8 + | -- | An @Authentication@ message had an unrecognized sub-type. + UnexpectedAuthType Int32 + | -- | A field could not be interpreted (context, offending tag). + BadField ByteString + | -- | The message body was shorter than expected. + TruncatedInput + | -- | A general protocol violation. + ProtocolViolation ByteString + deriving stock (Eq, Show) + +-- | A decoder of a message body: a @ptr-peeker@ 'Peeker.Variable' parser with +-- semantic failures layered on top. +type Decoder = ExceptT DecodingError Peeker.Variable + +-- | A type that can be serialized to the wire and decoded back from it. +class Comms a where + toWrite :: a -> Poker.Write + decoderOf :: Decoder a + +-- | Run a decoder against a complete message body. A framing shortfall (the +-- body being shorter than the decoder consumes) is reported as 'TruncatedInput'. +runDecoder :: Decoder a -> ByteString -> Either DecodingError a +runDecoder decoder body = + case Peeker.runVariableOnByteString (runExceptT decoder) body of + Left _needed -> Left TruncatedInput + Right (Left err) -> Left err + Right (Right value) -> Right value + +liftVariable :: Peeker.Variable a -> Decoder a +liftVariable = lift + +liftFixed :: Peeker.Fixed a -> Decoder a +liftFixed = lift . Peeker.fixed + +int16 :: Decoder Int16 +int16 = liftFixed Peeker.beSignedInt2 + +int32 :: Decoder Int32 +int32 = liftFixed Peeker.beSignedInt4 + +word8 :: Decoder Word8 +word8 = liftFixed Peeker.unsignedInt1 + +word32 :: Decoder Word32 +word32 = liftFixed Peeker.beUnsignedInt4 + +-- | A null-terminated string. +cstring :: Decoder ByteString +cstring = liftVariable Peeker.nullTerminatedStringAsByteString + +-- | Exactly @n@ bytes. +bytes :: Int -> Decoder ByteString +bytes n = liftFixed (Peeker.byteArrayAsByteString n) + +-- | All remaining bytes of the body. +remaining :: Decoder ByteString +remaining = liftVariable Peeker.remainderAsByteString + +-- | A null-terminated string, as a 'Comms' building block. +newtype CString = CString ByteString + deriving stock (Eq, Show) + +instance Comms CString where + toWrite (CString value) = Poker.byteString value <> Poker.word8 0 + decoderOf = CString <$> cstring + +instance Comms Int16 where + toWrite = Poker.bInt16 + decoderOf = int16 + +instance Comms Int32 where + toWrite = Poker.bInt32 + decoderOf = int32 + +instance Comms Word32 where + toWrite = Poker.bWord32 + decoderOf = word32 + +instance Comms Word8 where + toWrite = Poker.word8 + decoderOf = word8 diff --git a/src/library/Pqi/Native.hs b/src/library/Pqi/Native.hs new file mode 100644 --- /dev/null +++ b/src/library/Pqi/Native.hs @@ -0,0 +1,251 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} + +-- | The native (pure-Haskell) @pqi@ adapter. +-- +-- 'Connection' speaks the PostgreSQL wire protocol directly. +-- Provides the 'IsConnection' instance. +module Pqi.Native + ( Connection, + ) +where + +import qualified Data.ByteString as ByteString +import qualified Data.ByteString.Char8 as ByteString.Char8 +import qualified Data.Map.Strict as Map +import Pqi +import Pqi.Native.Connection (Connection (..)) +import qualified Pqi.Native.Connection as Connection +import qualified Pqi.Native.LargeObject as LargeObject +import Pqi.Native.Prelude +import qualified Pqi.Native.Query as Query +import qualified Pqi.Native.Transport as Transport +import Pqi.Native.Transport.Message + ( BackendMessage (..), + copyDataMessage, + copyDoneMessage, + copyFailMessage, + flushMessage, + syncMessage, + ) +import Pqi.Native.Types (NativeCancel (..), NativeResult (..)) +import System.Posix.Types (Fd) + +instance IsConnection Connection where + type ResultOf Connection = NativeResult + type CancelOf Connection = NativeCancel + + connectdb = Connection.establish + connectStart = Connection.establish + connectPoll _ = pure PollingOk + newNullConnection = Connection.nullConnection + isNullConnection connection = connection.isNull + finish connection = readIORef connection.transport >>= Transport.close + reset = Connection.reconnect + resetStart connection = Connection.reconnect connection $> True + resetPoll _ = pure PollingOk + db connection = pure (Just connection.info.database) + user connection = pure (Just connection.info.user) + pass connection = pure (Just connection.info.password) + host connection = pure (Just connection.info.host) + port connection = pure (Just (ByteString.Char8.pack (show connection.info.port))) + options _ = pure (Just "") + status connection = readIORef connection.connStatus + transactionStatus connection = transactionStatusOf <$> readIORef connection.txStatus + parameterStatus connection name = Map.lookup name <$> readIORef connection.parameters + protocolVersion _ = pure 3 + serverVersion connection = + maybe 0 parseServerVersion . Map.lookup "server_version" <$> readIORef connection.parameters + errorMessage connection = readIORef connection.lastError + socket connection = do + transport <- readIORef connection.transport + fd <- Transport.socketFd transport + pure (Just (fromIntegral fd :: Fd)) + backendPID connection = maybe 0 fst <$> readIORef connection.backendKey + connectionNeedsPassword _ = pure False + connectionUsedPassword connection = pure (not (ByteString.null connection.info.password)) + + exec connection sql = Query.exec connection sql + execParams connection sql params resultFormat = + Query.execParams connection sql params resultFormat + prepare connection name sql parameterTypes = + Query.prepare connection name sql parameterTypes + execPrepared connection name params resultFormat = + Query.execPrepared connection name params resultFormat + describePrepared connection name = Query.describePrepared connection name + describePortal connection name = Query.describePortal connection name + + escapeStringConn _ value + | isValidUtf8 value = + pure (Just (ByteString.intercalate "''" (ByteString.split 0x27 value))) + | otherwise = pure Nothing + escapeByteaConn _ value = + pure (Just ("\\x" <> hexEncode value)) + escapeIdentifier _ value + | isValidUtf8 value = + pure (Just ("\"" <> ByteString.intercalate "\"\"" (ByteString.split 0x22 value) <> "\"")) + | otherwise = pure Nothing + + sendQuery = Query.sendQuery + sendQueryParams = Query.sendQueryParams + sendPrepare = Query.sendPrepare + sendQueryPrepared = Query.sendQueryPrepared + sendDescribePrepared = Query.sendDescribePrepared + sendDescribePortal = Query.sendDescribePortal + getResult connection = Query.getNextResult connection + consumeInput _ = pure True + isBusy _ = pure False + setnonblocking connection flag = writeIORef connection.nonblocking flag $> True + isnonblocking connection = readIORef connection.nonblocking + setSingleRowMode connection = do + pending <- readIORef connection.asyncPending + if pending + then writeIORef connection.singleRowMode True $> True + else pure False + flush _ = pure FlushOk + + pipelineStatus connection = readIORef connection.pipelineStatus + enterPipelineMode connection = writeIORef connection.pipelineStatus PipelineOn $> True + exitPipelineMode connection = do + pending <- readIORef connection.asyncPending + if pending + then pure False + else writeIORef connection.pipelineStatus PipelineOff $> True + pipelineSync connection = do + Connection.sendMessage connection syncMessage + modifyIORef' connection.pendingSyncs (+ 1) + writeIORef connection.asyncPending True + pure True + sendFlushRequest connection = Connection.sendMessage connection flushMessage $> True + + getCancel connection = do + key <- readIORef connection.backendKey + pure + $ fmap + ( \(pid, secret) -> + NativeCancel + { host = connection.info.host, + port = connection.info.port, + pid, + secret, + asyncPendingRef = connection.asyncPending, + pipelineStatusRef = connection.pipelineStatus, + pendingCommandsRef = connection.pendingCommands + } + ) + key + + notifies connection = popFirst connection.pendingNotifications + disableNoticeReporting connection = writeIORef connection.noticeReporting False + enableNoticeReporting connection = writeIORef connection.noticeReporting True + getNotice connection = popFirst connection.notices + + putCopyData connection payload = Connection.sendMessage connection (copyDataMessage payload) $> CopyInOk + putCopyEnd connection reason = do + Connection.sendMessage connection (maybe copyDoneMessage copyFailMessage reason) + writeIORef connection.asyncPending True + pure CopyInOk + getCopyData connection _nonBlocking = do + message <- Connection.nextMessage connection + case message of + CopyData payload -> pure (CopyOutRow payload) + CopyDone -> do + writeIORef connection.asyncPending True + pure CopyOutDone + CommandComplete _ -> drainToReady connection $> CopyOutDone + ErrorResponse _ -> drainToReady connection $> CopyOutError + ReadyForQuery txState -> writeIORef connection.txStatus txState $> CopyOutDone + _ -> getCopyData connection _nonBlocking + + loCreat = LargeObject.loCreat + loCreate = LargeObject.loCreate + loImport = LargeObject.loImport + loImportWithOid = LargeObject.loImportWithOid + loExport = LargeObject.loExport + loOpen = LargeObject.loOpen + loWrite = LargeObject.loWrite + loRead = LargeObject.loRead + loSeek = LargeObject.loSeek + loTell = LargeObject.loTell + loTruncate = LargeObject.loTruncate + loClose = LargeObject.loClose + loUnlink = LargeObject.loUnlink + + clientEncoding connection = + fromMaybe "SQL_ASCII" . Map.lookup "client_encoding" <$> readIORef connection.parameters + setClientEncoding connection encoding = do + result <- Query.exec connection ("SET client_encoding TO '" <> encoding <> "'") + pure (maybe False (\value -> value.status /= FatalError) result) + setErrorVerbosity connection verbosity = do + previous <- readIORef connection.errorVerbosity + writeIORef connection.errorVerbosity verbosity + pure previous + +-- | Read messages until @ReadyForQuery@, recording the transaction status. +drainToReady :: Connection -> IO () +drainToReady connection = do + message <- Connection.nextMessage connection + case message of + ReadyForQuery txState -> writeIORef connection.txStatus txState + _ -> drainToReady connection + +-- | Pop the oldest element of a list stored newest-first. +popFirst :: IORef [a] -> IO (Maybe a) +popFirst ref = + atomicModifyIORef' ref \xs -> case reverse xs of + [] -> ([], Nothing) + oldest : rest -> (reverse rest, Just oldest) + +transactionStatusOf :: Word8 -> TransactionStatus +transactionStatusOf = \case + 0x49 -> TransIdle -- 'I' + 0x54 -> TransInTrans -- 'T' + 0x45 -> TransInError -- 'E' + _ -> TransUnknown + +-- | Parse the @server_version@ parameter into libpq's @MMmmpp@ integer form +-- (e.g. @\"17.2\"@ -> @170002@, @\"9.6.3\"@ -> @90603@). +parseServerVersion :: ByteString -> Int +parseServerVersion raw = + case ByteString.Char8.readInt raw of + Nothing -> 0 + Just (major, rest) + | major >= 10 -> major * 10000 + nextInt rest + | otherwise -> + let minor = nextInt rest + patch = nextInt (dropInt rest) + in major * 10000 + minor * 100 + patch + where + nextInt bs = case ByteString.Char8.uncons bs of + Just ('.', remainder) -> maybe 0 fst (ByteString.Char8.readInt remainder) + _ -> 0 + dropInt bs = case ByteString.Char8.uncons bs of + Just ('.', remainder) -> case ByteString.Char8.readInt remainder of + Just (_, leftover) -> leftover + Nothing -> remainder + _ -> bs + +isValidUtf8 :: ByteString -> Bool +isValidUtf8 = go . ByteString.unpack + where + go [] = True + go (b : bs) + | b < 0x80 = go bs + | b < 0xc2 = False + | b < 0xe0 = cont bs 1 + | b < 0xf0 = cont bs 2 + | b < 0xf5 = cont bs 3 + | otherwise = False + cont bs (0 :: Int) = go bs + cont [] _ = False + cont (b : bs) n + | b .&. 0xc0 == 0x80 = cont bs (n - 1) + | otherwise = False + +hexEncode :: ByteString -> ByteString +hexEncode = ByteString.Char8.pack . concatMap toHex . ByteString.unpack + where + toHex byte = [digit (byte `div` 16), digit (byte `mod` 16)] + digit n + | n < 10 = toEnum (fromIntegral n + fromEnum '0') + | otherwise = toEnum (fromIntegral n - 10 + fromEnum 'a') diff --git a/src/library/Pqi/Native/Auth.hs b/src/library/Pqi/Native/Auth.hs new file mode 100644 --- /dev/null +++ b/src/library/Pqi/Native/Auth.hs @@ -0,0 +1,125 @@ +-- | Authentication helpers: MD5 password hashing and the SASL\/SCRAM-SHA-256 +-- exchange. +-- +-- The pure crypto lives here; the message round-trip is abstracted as 'SaslStep' +-- so the connection module owns the actual socket I\/O. +module Pqi.Native.Auth + ( md5Password, + SaslStep (..), + SaslMessage (..), + scram, + ) +where + +import Crypto.Hash (Digest, MD5 (..), SHA256 (..), hashWith) +import Crypto.KDF.PBKDF2 (Parameters (..), fastPBKDF2_SHA256) +import Crypto.MAC.HMAC (HMAC, hmac, hmacGetDigest) +import Crypto.Random (getRandomBytes) +import qualified Data.ByteString as ByteString +import qualified Data.ByteString.Base64 as Base64 +import qualified Data.ByteString.Char8 as ByteString.Char8 +import qualified Data.List as List +import Pqi.Native.Prelude + +-- | Compute the response to an @AuthenticationMD5Password@ challenge: +-- @"md5" <> md5hex (md5hex (password <> user) <> salt)@. +md5Password :: ByteString -> ByteString -> ByteString -> ByteString +md5Password user password salt = + "md5" <> md5Hex (md5Hex (password <> user) <> salt) + +md5Hex :: ByteString -> ByteString +md5Hex = ByteString.Char8.pack . show . hashWith MD5 + +-- | One server SASL\/authentication message, as the SCRAM logic sees it. +data SaslMessage + = SaslContinue ByteString + | SaslFinal ByteString + | SaslOk + | SaslError ByteString + +-- | The message round-trip the SCRAM exchange drives, supplied by the +-- connection module. +data SaslStep = SaslStep + { sendInitial :: ByteString -> ByteString -> IO (), + sendResponse :: ByteString -> IO (), + receive :: IO SaslMessage + } + +mechanismName :: ByteString +mechanismName = "SCRAM-SHA-256" + +-- | Run the SCRAM-SHA-256 exchange (without channel binding). Returns @Right ()@ +-- once the server accepts the client proof, or @Left@ with a problem +-- description. +scram :: ByteString -> ByteString -> [ByteString] -> SaslStep -> IO (Either ByteString ()) +scram _user password mechanisms step + | mechanismName `notElem` mechanisms = + pure (Left "server did not offer SCRAM-SHA-256") + | otherwise = do + clientNonce <- Base64.encode <$> getRandomBytes 18 + let clientFirstBare = "n=,r=" <> clientNonce + step.sendInitial mechanismName ("n,," <> clientFirstBare) + step.receive >>= \case + SaslError problem -> pure (Left problem) + SaslContinue serverFirst -> + case parseServerFirst serverFirst of + Nothing -> pure (Left "malformed SCRAM server-first message") + Just (serverNonce, salt, iterations) -> do + let saltedPassword = fastPBKDF2_SHA256 (Parameters iterations 32) password salt :: ByteString + clientKey = hmacSha256 saltedPassword "Client Key" + storedKey = sha256 clientKey + clientFinalWithoutProof = "c=biws,r=" <> serverNonce + authMessage = + ByteString.intercalate "," [clientFirstBare, serverFirst, clientFinalWithoutProof] + clientSignature = hmacSha256 storedKey authMessage + clientProof = xorBytes clientKey clientSignature + clientFinal = clientFinalWithoutProof <> ",p=" <> Base64.encode clientProof + step.sendResponse clientFinal + step.receive >>= \case + SaslFinal _ -> pure (Right ()) + SaslOk -> pure (Right ()) + SaslError problem -> pure (Left problem) + _ -> pure (Left "unexpected SCRAM server-final message") + _ -> pure (Left "unexpected SCRAM message") + +-- | Parse @r=,s=,i=@ into the server nonce, the +-- decoded salt, and the iteration count. +parseServerFirst :: ByteString -> Maybe (ByteString, ByteString, Int) +parseServerFirst message = do + let attributes = ByteString.Char8.split ',' message + nonce <- attributeValue "r=" attributes + saltEncoded <- attributeValue "s=" attributes + salt <- either (const Nothing) Just (Base64.decode saltEncoded) + iterationsText <- attributeValue "i=" attributes + (iterations, _) <- ByteString.Char8.readInt iterationsText + pure (nonce, salt, iterations) + where + attributeValue prefix = + fmap (ByteString.drop (ByteString.length prefix)) + . List.find (ByteString.isPrefixOf prefix) + +hmacSha256 :: ByteString -> ByteString -> ByteString +hmacSha256 key message = digestBytes (hmacGetDigest (hmac key message :: HMAC SHA256)) + +sha256 :: ByteString -> ByteString +sha256 = digestBytes . hashWith SHA256 + +-- | Extract the raw bytes of a digest via its hexadecimal 'Show' instance, +-- avoiding a direct @memory@ dependency (whose @ByteArrayAccess@ instance for +-- crypton's @Digest@ conflicts across package versions). +digestBytes :: Digest a -> ByteString +digestBytes = hexToBytes . ByteString.Char8.pack . show + +hexToBytes :: ByteString -> ByteString +hexToBytes = ByteString.pack . pairs . ByteString.unpack + where + pairs (hi : lo : rest) = (hexValue hi * 16 + hexValue lo) : pairs rest + pairs _ = [] + hexValue w + | w >= 0x30 && w <= 0x39 = w - 0x30 + | w >= 0x61 && w <= 0x66 = w - 0x57 + | w >= 0x41 && w <= 0x46 = w - 0x37 + | otherwise = 0 + +xorBytes :: ByteString -> ByteString -> ByteString +xorBytes a b = ByteString.pack (ByteString.zipWith xor a b) diff --git a/src/library/Pqi/Native/Connection.hs b/src/library/Pqi/Native/Connection.hs new file mode 100644 --- /dev/null +++ b/src/library/Pqi/Native/Connection.hs @@ -0,0 +1,347 @@ +-- | The native connection: its mutable state, conninfo parsing, the +-- startup\/authentication handshake, and the interleave-aware receive loop that +-- the higher-level query code is built on. +module Pqi.Native.Connection + ( Connection (..), + ConnInfo (..), + parseConnInfo, + establish, + nullConnection, + reconnect, + nextMessage, + sendMessage, + fieldValue, + setError, + ) +where + +import Control.Exception (IOException, SomeException, catch, try) +import qualified Data.ByteString as ByteString +import qualified Data.ByteString.Char8 as ByteString.Char8 +import qualified Data.Map.Strict as Map +import Pqi (ConnStatus (..), Notify (..), PipelineStatus (..), Verbosity (..)) +import qualified Pqi.Native.Auth as Auth +import Pqi.Native.Prelude +import Pqi.Native.Transport (Transport) +import qualified Pqi.Native.Transport as Transport +import Pqi.Native.Transport.Message +import Pqi.Native.Types (formatErrorFields) +import qualified PtrPoker.Write as Poker + +-- | Parsed connection parameters (the @key=value@ subset we support). +data ConnInfo = ConnInfo + { host :: ByteString, + port :: Int, + user :: ByteString, + database :: ByteString, + password :: ByteString + } + deriving stock (Eq, Show) + +-- | Parse a conninfo string in either @key=value@ or @postgresql:\/\/@ URI +-- format. Unquoted key=value values only; URI values are percent-decoded. +-- Environment variables and @.pgpass@ are not supported. +parseConnInfo :: ByteString -> ConnInfo +parseConnInfo raw + | "postgresql://" `ByteString.isPrefixOf` raw = parseUri (ByteString.drop 13 raw) + | "postgres://" `ByteString.isPrefixOf` raw = parseUri (ByteString.drop 11 raw) + | otherwise = parseKeyValue raw + +parseKeyValue :: ByteString -> ConnInfo +parseKeyValue raw = + ConnInfo + { host = get "host" "localhost", + port = maybe 5432 fst (ByteString.Char8.readInt (get "port" "5432")), + user = theUser, + database = get "dbname" theUser, + password = get "password" "" + } + where + pairs = mapMaybe toPair (ByteString.Char8.words raw) + settings = Map.fromList pairs + get key def = Map.findWithDefault def key settings + theUser = get "user" "postgres" + toPair token = case ByteString.Char8.break (== '=') token of + (key, value) + | not (ByteString.null value) -> Just (key, ByteString.drop 1 value) + _ -> Nothing + +-- | Parse the authority+path portion of a @postgresql://@ URI (scheme already +-- stripped). Handles @[user[:password]@][host[:port]][/dbname]@; ignores query +-- parameters other than what appears in those components. +parseUri :: ByteString -> ConnInfo +parseUri withoutScheme = + ConnInfo {host, port, user, database, password} + where + -- Split off optional "userinfo@" prefix. The '@' is unambiguous in this + -- position: hosts do not contain '@' in practice. + (userinfoMay, afterAt) = case ByteString.elemIndex 0x40 withoutScheme of + Just i -> (Just (ByteString.take i withoutScheme), ByteString.drop (i + 1) withoutScheme) + Nothing -> (Nothing, withoutScheme) + + (user, password) = case userinfoMay of + Nothing -> ("postgres", "") + Just ui -> case ByteString.elemIndex 0x3a ui of + Just c -> (pctDecode (ByteString.take c ui), pctDecode (ByteString.drop (c + 1) ui)) + Nothing -> (pctDecode ui, "") + + -- Split host[:port] from /dbname?params. + (hostport, pathAndQuery) = ByteString.break (== 0x2f) afterAt + + -- Strip optional leading '/' to get just the dbname (ignoring ?query). + rawDbname = + let p = ByteString.drop 1 pathAndQuery + in ByteString.takeWhile (/= 0x3f) p + + database = + if ByteString.null rawDbname + then user + else pctDecode rawDbname + + -- Parse host and port from "host:port", handling IPv6 "[::1]:port". + (host, port) + | not (ByteString.null hostport) && ByteString.head hostport == 0x5b = + case ByteString.elemIndex 0x5d hostport of + Just close -> + let h = ByteString.take close (ByteString.drop 1 hostport) + portStr = ByteString.drop 2 (ByteString.drop close hostport) + in (h, readPort portStr) + Nothing -> (hostport, 5432) + | otherwise = case ByteString.elemIndexEnd 0x3a hostport of + Just c -> + let h = ByteString.take c hostport + p = ByteString.drop (c + 1) hostport + in if ByteString.null h then ("localhost", readPort p) else (h, readPort p) + Nothing -> + (if ByteString.null hostport then "localhost" else hostport, 5432) + + readPort bs = maybe 5432 fst (ByteString.Char8.readInt bs) + +-- | Decode @%XX@ percent-encoding in a URI component. +pctDecode :: ByteString -> ByteString +pctDecode bs + | ByteString.null bs = ByteString.empty + | ByteString.head bs == 0x25, + ByteString.length bs >= 3, + Just hi <- hexVal (ByteString.index bs 1), + Just lo <- hexVal (ByteString.index bs 2) = + ByteString.singleton (fromIntegral (hi * 16 + lo)) <> pctDecode (ByteString.drop 3 bs) + | otherwise = ByteString.singleton (ByteString.head bs) <> pctDecode (ByteString.tail bs) + where + hexVal w + | w >= 0x30 && w <= 0x39 = Just (fromIntegral w - 0x30 :: Int) + | w >= 0x41 && w <= 0x46 = Just (fromIntegral w - 0x41 + 10 :: Int) + | w >= 0x61 && w <= 0x66 = Just (fromIntegral w - 0x61 + 10 :: Int) + | otherwise = Nothing + +-- | A native connection and its mutable state. +data Connection = Connection + { transport :: IORef Transport, + info :: ConnInfo, + isNull :: Bool, + parameters :: IORef (Map.Map ByteString ByteString), + backendKey :: IORef (Maybe (Int32, Int32)), + txStatus :: IORef Word8, + connStatus :: IORef ConnStatus, + lastError :: IORef (Maybe ByteString), + notices :: IORef [ByteString], + pendingNotifications :: IORef [Notify], + noticeReporting :: IORef Bool, + asyncPending :: IORef Bool, + nonblocking :: IORef Bool, + pipelineStatus :: IORef PipelineStatus, + singleRowMode :: IORef Bool, + singleRowFields :: IORef [FieldDescription], + pipelineSeparatorPending :: IORef Bool, + pendingSyncs :: IORef Int, + pendingCommands :: IORef Int, + -- | Count of standalone @Parse@ messages sent in pipeline mode (each from + -- 'sendPrepare'), for which only @ParseComplete@ (no @CommandComplete@) is + -- expected as a terminal response. + pendingParses :: IORef Int, + errorVerbosity :: IORef Verbosity, + -- | The SQL text of the most recently sent query (set by sendQuery / + -- sendQueryParams). Used when formatting error messages for async results + -- so that @LINE N:@ position context can be reproduced. + currentQuery :: IORef ByteString + } + +-- | Send a serialized frontend message. +sendMessage :: Connection -> Poker.Write -> IO () +sendMessage connection write = do + transport <- readIORef connection.transport + Transport.send transport write + +-- | Receive the next /protocol-relevant/ backend message, transparently +-- consuming and recording the asynchronous messages the backend may interleave +-- at any time: @ParameterStatus@ (updates the parameter map), @NoticeResponse@ +-- (collected when notice reporting is on), and @NotificationResponse@ (queued). +nextMessage :: Connection -> IO BackendMessage +nextMessage connection = do + transport <- readIORef connection.transport + (typeByte, body) <- Transport.receiveFrame transport + case decodeBackendMessage typeByte body of + Left err -> ioError (userError ("pqi-native: protocol decode error: " <> show err)) + Right message -> case message of + ParameterStatus key value -> do + modifyIORef' connection.parameters (Map.insert key value) + nextMessage connection + NoticeResponse fields -> do + reporting <- readIORef connection.noticeReporting + when reporting $ do + let noticeText = formatErrorFields (Map.fromList fields) + unless (ByteString.null noticeText) + $ modifyIORef' connection.notices (noticeText :) + nextMessage connection + NotificationResponse pid channel payload -> do + modifyIORef' connection.pendingNotifications (Notify channel pid payload :) + nextMessage connection + other -> pure other + +-- | Look up an error\/notice field by its single-byte code. +fieldValue :: Word8 -> [(Word8, ByteString)] -> Maybe ByteString +fieldValue code = lookup code + +-- | Record a flat error message and mark the connection bad. +setError :: Connection -> ByteString -> IO () +setError connection message = do + writeIORef connection.lastError (Just message) + writeIORef connection.connStatus ConnectionBad + +-- | Open a connection: resolve and connect the socket, send the startup +-- message, and run the authentication\/startup handshake. Like libpq, a failed +-- connection (whether due to a network error or a rejected handshake) yields a +-- 'ConnectionBad' connection rather than throwing. +establish :: ByteString -> IO Connection +establish conninfo = do + let info = parseConnInfo conninfo + transportResult <- try @IOException (Transport.connect info.host info.port) + case transportResult of + Left err -> do + transport <- Transport.unconnected + connection <- newConnection False transport info + setError connection ("could not connect to server: " <> ByteString.Char8.pack (show err)) + pure connection + Right transport -> do + connection <- newConnection False transport info + sendMessage connection (startupMessage [("user", info.user), ("database", info.database)]) + handshake connection + pure connection + +-- | A \"null\" sentinel connection (the analogue of @PQnewNullConnection@): no +-- live socket, permanently in the 'ConnectionBad' state. +nullConnection :: IO Connection +nullConnection = do + transport <- Transport.unconnected + conn <- newConnection True transport (parseConnInfo "") + writeIORef conn.lastError (Just "connection pointer is NULL\n") + pure conn + +-- | Close the current socket and run the startup handshake again on a fresh +-- one, reusing the stored conninfo (the analogue of @PQreset@). +reconnect :: Connection -> IO () +reconnect connection = do + oldTransport <- readIORef connection.transport + Transport.close oldTransport + newTransport <- Transport.connect connection.info.host connection.info.port + writeIORef connection.transport newTransport + writeIORef connection.parameters Map.empty + writeIORef connection.backendKey Nothing + writeIORef connection.txStatus 0x49 + writeIORef connection.connStatus ConnectionBad + writeIORef connection.lastError (Just "") + sendMessage connection (startupMessage [("user", connection.info.user), ("database", connection.info.database)]) + handshake connection + +newConnection :: Bool -> Transport -> ConnInfo -> IO Connection +newConnection isNull transport info = do + transportRef <- newIORef transport + Connection transportRef info isNull + <$> newIORef Map.empty + <*> newIORef Nothing + <*> newIORef 0x49 -- 'I' + <*> newIORef ConnectionBad + <*> newIORef (Just "") + <*> newIORef [] + <*> newIORef [] + <*> newIORef False + <*> newIORef False + <*> newIORef False + <*> newIORef PipelineOff + <*> newIORef False + <*> newIORef [] + <*> newIORef False + <*> newIORef 0 + <*> newIORef 0 + <*> newIORef 0 + <*> newIORef ErrorsDefault + <*> newIORef "" + +-- | The startup\/authentication state machine, ending at the first +-- @ReadyForQuery@ (success) or @ErrorResponse@ (failure). +handshake :: Connection -> IO () +handshake connection = authenticating + where + authenticating = do + message <- nextMessage connection + case message of + AuthenticationOk -> startingUp + AuthenticationCleartextPassword -> do + sendMessage connection (passwordMessage connection.info.password) + authenticating + AuthenticationMD5Password salt -> do + let response = Auth.md5Password connection.info.user connection.info.password salt + sendMessage connection (passwordMessage response) + authenticating + AuthenticationSASL mechanisms -> + Auth.scram connection.info.user connection.info.password mechanisms (saslExchange connection) >>= \case + Left problem -> setError connection problem + Right () -> startingUp + ErrorResponse fields -> failWith fields + other -> failWith [(0x4d, "unexpected authentication message: " <> ByteString.Char8.pack (show other))] + startingUp = do + message <- nextMessage connection + case message of + BackendKeyData pid secret -> do + writeIORef connection.backendKey (Just (pid, secret)) + startingUp + ReadyForQuery txState -> do + writeIORef connection.txStatus txState + writeIORef connection.connStatus ConnectionOk + ErrorResponse fields -> failWith fields + _ -> startingUp + failWith fields = do + let fmtFields = formatErrorFields (Map.fromList fields) + transport <- readIORef connection.transport + mIp <- catch (Just <$> Transport.peerIp transport) (\(_ :: SomeException) -> pure Nothing) + setError connection $ case mIp of + Nothing -> fmtFields + Just ip -> + "connection to server at \"" + <> connection.info.host + <> "\" (" + <> ip + <> "), port " + <> ByteString.Char8.pack (show connection.info.port) + <> " failed: " + <> fmtFields + +-- | The SASL message round-trip used by 'Auth.scram': send a client message and +-- receive the next server SASL\/auth message, projected to the bytes the SCRAM +-- logic needs. +saslExchange :: Connection -> Auth.SaslStep +saslExchange connection = + Auth.SaslStep + { Auth.sendInitial = \mechanism initial -> + sendMessage connection (saslInitialResponse mechanism initial), + Auth.sendResponse = \payload -> + sendMessage connection (saslResponse payload), + Auth.receive = do + message <- nextMessage connection + pure $ case message of + AuthenticationSASLContinue payload -> Auth.SaslContinue payload + AuthenticationSASLFinal payload -> Auth.SaslFinal payload + AuthenticationOk -> Auth.SaslOk + ErrorResponse fields -> Auth.SaslError (fromMaybe "SASL error" (fieldValue 0x4d fields)) + other -> Auth.SaslError ("unexpected SASL message: " <> ByteString.Char8.pack (show other)) + } diff --git a/src/library/Pqi/Native/LargeObject.hs b/src/library/Pqi/Native/LargeObject.hs new file mode 100644 --- /dev/null +++ b/src/library/Pqi/Native/LargeObject.hs @@ -0,0 +1,131 @@ +-- | The large-object interface, implemented over the server's @lo_*@ SQL +-- functions (rather than libpq's fast-path protocol) — identical results, far +-- less machinery. As with libpq, the open\/read\/write\/close operations must be +-- run inside a transaction managed by the caller. +module Pqi.Native.LargeObject + ( loCreat, + loCreate, + loImport, + loImportWithOid, + loExport, + loOpen, + loWrite, + loRead, + loSeek, + loTell, + loTruncate, + loClose, + loUnlink, + ) +where + +import qualified Data.ByteString.Char8 as ByteString.Char8 +import Pqi (Format (..)) +import Pqi.Native.Connection (Connection) +import Pqi.Native.Prelude +import qualified Pqi.Native.Query as Query +import Pqi.Native.Types (NativeResult (..)) +import System.IO (IOMode (..), SeekMode (..)) + +loCreat :: Connection -> IO (Maybe Word32) +loCreat connection = + (>>= parseOid) <$> callText connection "select lo_creat($1 :: integer)" [intParam (-1)] + +loCreate :: Connection -> Word32 -> IO (Maybe Word32) +loCreate connection oid = + (>>= parseOid) <$> callText connection "select lo_create($1 :: oid)" [oidParam oid] + +loImport :: Connection -> FilePath -> IO (Maybe Word32) +loImport connection path = + (>>= parseOid) <$> callText connection "select lo_import($1)" [textParam (ByteString.Char8.pack path)] + +loImportWithOid :: Connection -> FilePath -> Word32 -> IO (Maybe Word32) +loImportWithOid connection path oid = + (>>= parseOid) + <$> callText connection "select lo_import($1, $2 :: oid)" [textParam (ByteString.Char8.pack path), oidParam oid] + +loExport :: Connection -> Word32 -> FilePath -> IO (Maybe ()) +loExport connection oid path = + succeeded <$> callText connection "select lo_export($1 :: oid, $2)" [oidParam oid, textParam (ByteString.Char8.pack path)] + +loOpen :: Connection -> Word32 -> IOMode -> IO (Maybe Int32) +loOpen connection oid mode = + fmap fromIntegral + . (>>= parseInt) + <$> callText connection "select lo_open($1 :: oid, $2 :: integer)" [oidParam oid, intParam (ioModeFlag mode)] + +loWrite :: Connection -> Int32 -> ByteString -> IO (Maybe Int) +loWrite connection fd payload = + (>>= parseInt) <$> callText connection "select lowrite($1 :: integer, $2 :: bytea)" [intParam (fromIntegral fd), byteaParam payload] + +loRead :: Connection -> Int32 -> Int -> IO (Maybe ByteString) +loRead connection fd len = + callBinary connection "select loread($1 :: integer, $2 :: integer)" [intParam (fromIntegral fd), intParam len] + +loSeek :: Connection -> Int32 -> SeekMode -> Int -> IO (Maybe Int) +loSeek connection fd whence offset = + (>>= parseInt) + <$> callText connection "select lo_lseek($1 :: integer, $2 :: integer, $3 :: integer)" [intParam (fromIntegral fd), intParam offset, intParam (seekFlag whence)] + +loTell :: Connection -> Int32 -> IO (Maybe Int) +loTell connection fd = + (>>= parseInt) <$> callText connection "select lo_tell($1 :: integer)" [intParam (fromIntegral fd)] + +loTruncate :: Connection -> Int32 -> Int -> IO (Maybe ()) +loTruncate connection fd len = + succeeded <$> callText connection "select lo_truncate($1 :: integer, $2 :: integer)" [intParam (fromIntegral fd), intParam len] + +loClose :: Connection -> Int32 -> IO (Maybe ()) +loClose connection fd = + succeeded <$> callText connection "select lo_close($1 :: integer)" [intParam (fromIntegral fd)] + +loUnlink :: Connection -> Word32 -> IO (Maybe ()) +loUnlink connection oid = + succeeded <$> callText connection "select lo_unlink($1 :: oid)" [oidParam oid] + +-- * Helpers + +callText :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> IO (Maybe ByteString) +callText connection sql params = (>>= firstValue) <$> Query.execParams connection sql params Text + +callBinary :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> IO (Maybe ByteString) +callBinary connection sql params = (>>= firstValue) <$> Query.execParams connection sql params Binary + +firstValue :: NativeResult -> Maybe ByteString +firstValue result = case result.rows of + (cell : _) : _ -> cell + _ -> Nothing + +succeeded :: Maybe ByteString -> Maybe () +succeeded = fmap (const ()) + +intParam :: Int -> Maybe (Word32, ByteString, Format) +intParam value = Just (0, ByteString.Char8.pack (show value), Text) + +oidParam :: Word32 -> Maybe (Word32, ByteString, Format) +oidParam value = Just (0, ByteString.Char8.pack (show value), Text) + +textParam :: ByteString -> Maybe (Word32, ByteString, Format) +textParam value = Just (0, value, Text) + +byteaParam :: ByteString -> Maybe (Word32, ByteString, Format) +byteaParam value = Just (17, value, Binary) + +parseInt :: ByteString -> Maybe Int +parseInt = fmap fst . ByteString.Char8.readInt + +parseOid :: ByteString -> Maybe Word32 +parseOid = fmap (fromIntegral . fst) . ByteString.Char8.readInt + +ioModeFlag :: IOMode -> Int +ioModeFlag = \case + ReadMode -> 0x40000 -- INV_READ + WriteMode -> 0x20000 -- INV_WRITE + AppendMode -> 0x20000 + ReadWriteMode -> 0x60000 -- INV_READ | INV_WRITE + +seekFlag :: SeekMode -> Int +seekFlag = \case + AbsoluteSeek -> 0 + RelativeSeek -> 1 + SeekFromEnd -> 2 diff --git a/src/library/Pqi/Native/Prelude.hs b/src/library/Pqi/Native/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/library/Pqi/Native/Prelude.hs @@ -0,0 +1,15 @@ +module Pqi.Native.Prelude + ( module Exports, + ) +where + +import Control.Applicative as Exports +import Control.Monad as Exports +import Data.Bits as Exports +import Data.ByteString as Exports (ByteString) +import Data.Functor as Exports (($>), (<&>)) +import Data.IORef as Exports +import Data.Int as Exports +import Data.Maybe as Exports +import Data.Word as Exports +import Prelude as Exports diff --git a/src/library/Pqi/Native/Query.hs b/src/library/Pqi/Native/Query.hs new file mode 100644 --- /dev/null +++ b/src/library/Pqi/Native/Query.hs @@ -0,0 +1,446 @@ +-- | Command execution: the simple- and extended-query flows, and the +-- materialization of the backend message stream into a 'NativeResult'. +module Pqi.Native.Query + ( exec, + execParams, + prepare, + execPrepared, + describePrepared, + describePortal, + sendQuery, + sendQueryParams, + sendPrepare, + sendQueryPrepared, + sendDescribePrepared, + sendDescribePortal, + getNextResult, + ) +where + +import qualified Data.Map.Strict as Map +import Pqi (ConnStatus (..), ExecStatus (..), Format (..), PipelineStatus (..)) +import Pqi.Native.Connection +import Pqi.Native.Prelude +import Pqi.Native.Transport.Message +import Pqi.Native.Types (NativeResult (..), formatResultError) +import qualified PtrPoker.Write as Poker + +-- * Message construction + +-- Sync-inclusive variants used by the synchronous exec* functions. + +paramsWrite :: ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> Poker.Write +paramsWrite sql params resultFormat = + asyncParamsWrite sql params resultFormat <> syncMessage + +preparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Poker.Write +preparedWrite name params resultFormat = + asyncPreparedWrite name params resultFormat <> syncMessage + +prepareWrite :: ByteString -> ByteString -> Maybe [Word32] -> Poker.Write +prepareWrite name sql parameterTypes = + parseMessage name sql (fromMaybe [] parameterTypes) <> syncMessage + +-- Sync-free variants used by the async send* functions. +-- In non-pipeline mode sendAsync appends syncMessage; in pipeline mode it does not. + +asyncParamsWrite :: ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> Poker.Write +asyncParamsWrite sql params resultFormat = + parseMessage "" sql (fmap paramOid params) + <> bindMessage "" "" (fmap paramFormat params) (fmap paramValue params) [formatCode resultFormat] + <> describePortalMessage "" + <> executeMessage "" 0 + +asyncPreparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Poker.Write +asyncPreparedWrite name params resultFormat = + bindMessage "" name (fmap boundFormat params) (fmap boundValue params) [formatCode resultFormat] + <> describePortalMessage "" + <> executeMessage "" 0 + +-- * Synchronous flows + +-- | Simple query. Returns the last result, mirroring @PQexec@. +exec :: Connection -> ByteString -> IO (Maybe NativeResult) +exec connection sql = withReady connection do + sendMessage connection (queryMessage sql) + lastMaybe <$> collectSimple connection sql + +-- | Parameterized query via the extended protocol. +execParams :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> IO (Maybe NativeResult) +execParams connection sql params resultFormat = withReady connection do + sendMessage connection (paramsWrite sql params resultFormat) + Just <$> collectExtended connection sql + +-- | Prepare a named statement. +prepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO (Maybe NativeResult) +prepare connection name sql parameterTypes = withReady connection do + sendMessage connection (prepareWrite name sql parameterTypes) + Just <$> collectExtended connection sql + +-- | Execute a previously prepared statement. +execPrepared :: Connection -> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO (Maybe NativeResult) +execPrepared connection name params resultFormat = withReady connection do + sendMessage connection (preparedWrite name params resultFormat) + Just <$> collectExtended connection "" + +-- * Asynchronous flows + +-- | Send a write in async mode, tracking pending commands for pipeline abort. +sendAsync :: Connection -> ByteString -> Poker.Write -> IO Bool +sendAsync connection sql write = do + status <- readIORef connection.connStatus + case status of + ConnectionOk -> do + sendMessage connection write + writeIORef connection.currentQuery sql + writeIORef connection.asyncPending True + pipeStatus <- readIORef connection.pipelineStatus + when (pipeStatus /= PipelineOff) $ modifyIORef' connection.pendingCommands (+ 1) + pure True + _ -> pure False + +-- | Whether the connection is in pipeline mode. +inPipeline :: Connection -> IO Bool +inPipeline connection = (/= PipelineOff) <$> readIORef connection.pipelineStatus + +-- Simple query protocol: no Sync needed (server sends ReadyForQuery on its own). +sendQuery :: Connection -> ByteString -> IO Bool +sendQuery connection sql = sendAsync connection sql (queryMessage sql) + +-- Extended query: include Sync when not in pipeline mode; omit Sync in +-- pipeline mode (the caller drives sync boundaries via 'pipelineSync'). +sendQueryParams :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> IO Bool +sendQueryParams connection sql params resultFormat = do + pipeline <- inPipeline connection + sendAsync connection sql + $ if pipeline + then asyncParamsWrite sql params resultFormat + else paramsWrite sql params resultFormat + +sendPrepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO Bool +sendPrepare connection name sql parameterTypes = do + pipeline <- inPipeline connection + ok <- + sendAsync connection sql + $ if pipeline + then parseMessage name sql (fromMaybe [] parameterTypes) + else prepareWrite name sql parameterTypes + when (ok && pipeline) $ modifyIORef' connection.pendingParses (+ 1) + pure ok + +sendQueryPrepared :: Connection -> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO Bool +sendQueryPrepared connection name params resultFormat = do + pipeline <- inPipeline connection + sendAsync connection "" + $ if pipeline + then asyncPreparedWrite name params resultFormat + else preparedWrite name params resultFormat + +sendDescribePrepared :: Connection -> ByteString -> IO Bool +sendDescribePrepared connection name = do + pipeline <- inPipeline connection + sendAsync connection "" + $ if pipeline + then describeStatementMessage name + else describeStatementMessage name <> syncMessage + +sendDescribePortal :: Connection -> ByteString -> IO Bool +sendDescribePortal connection name = do + pipeline <- inPipeline connection + sendAsync connection "" + $ if pipeline + then describePortalMessage name + else describePortalMessage name <> syncMessage + +-- | Read the next result of an in-flight asynchronous command, or 'Nothing' +-- once @ReadyForQuery@ is reached (clearing the pending flag), mirroring +-- @PQgetResult@. +-- +-- In pipeline mode a separator 'Nothing' is returned between each command's +-- result set, and a 'PipelineSync' result is returned for each @Sync@ +-- boundary. In single-row mode each data row is delivered as a separate +-- 'SingleTuple' result followed by a final 'TuplesOk' with no rows. +getNextResult :: Connection -> IO (Maybe NativeResult) +getNextResult connection = do + pending <- readIORef connection.asyncPending + if not pending + then pure Nothing + else do + sepPending <- readIORef connection.pipelineSeparatorPending + if sepPending + then do + writeIORef connection.pipelineSeparatorPending False + pure Nothing + else do + singleRow <- readIORef connection.singleRowMode + cachedFields <- readIORef connection.singleRowFields + let initBuilder = + if singleRow && not (null cachedFields) + then emptyBuilder {accFields = cachedFields, accSawRowDescription = True} + else emptyBuilder + go singleRow initBuilder + where + -- Decrement the pending-command counter and set the separator flag when in + -- pipeline mode. Called when a "terminal" result is about to be returned. + finishCommand pipeStatus = do + when (pipeStatus /= PipelineOff) $ do + modifyIORef' connection.pendingCommands (subtract 1) + writeIORef connection.pipelineSeparatorPending True + + go singleRow builder = do + pipeStatus <- readIORef connection.pipelineStatus + -- In aborted pipeline mode, if the server has already sent nothing for + -- the remaining commands (it discards them after the first error), we + -- generate synthetic PipelineAbort results for each outstanding command + -- rather than blocking on a wire read that will never come. + pending <- readIORef connection.pendingCommands + if pipeStatus == PipelineAborted && pending > 0 + then do + modifyIORef' connection.pendingCommands (subtract 1) + writeIORef connection.pipelineSeparatorPending True + pure (Just (NativeResult PipelineAbort [] [] Nothing Map.empty [] "")) + else readAndProcess singleRow builder pipeStatus + + readAndProcess singleRow builder pipeStatus = do + message <- nextMessage connection + case message of + RowDescription fs -> + go singleRow builder {accFields = fs, accSawRowDescription = True, accHadResponse = True} + ParameterDescription oids -> + go singleRow builder {accParamOids = oids, accHadResponse = True} + NoData -> + go singleRow builder {accHadResponse = True} + DataRow values -> + if singleRow + then do + writeIORef connection.singleRowFields builder.accFields + pure (Just (NativeResult SingleTuple builder.accFields [values] Nothing Map.empty [] "")) + else go singleRow builder {accRevRows = values : builder.accRevRows} + ParseComplete -> do + parses <- readIORef connection.pendingParses + if parses > 0 && pipeStatus /= PipelineOff + then do + modifyIORef' connection.pendingParses (subtract 1) + finishCommand pipeStatus + pure (Just (NativeResult CommandOk [] [] Nothing Map.empty [] "")) + else go singleRow builder {accHadResponse = True} + BindComplete -> + go singleRow builder {accHadResponse = True} + CloseComplete -> + go singleRow builder {accHadResponse = True} + CommandComplete tag -> do + if singleRow + then do + writeIORef connection.singleRowMode False + writeIORef connection.singleRowFields [] + writeIORef connection.lastError (Just "") + pure (Just (NativeResult TuplesOk builder.accFields [] (Just tag) Map.empty [] "")) + else do + writeIORef connection.lastError (Just "") + finishCommand pipeStatus + pure (Just (commandResult builder (Just tag))) + EmptyQueryResponse -> do + writeIORef connection.lastError (Just "") + finishCommand pipeStatus + pure (Just (NativeResult EmptyQuery [] [] Nothing Map.empty [] "")) + ErrorResponse fs -> do + let errMap = Map.fromList fs + case pipeStatus of + PipelineAborted -> do + -- Should not normally happen (server discards commands in abort + -- mode) but handle defensively. + finishCommand pipeStatus + pure (Just (NativeResult PipelineAbort [] [] Nothing Map.empty [] "")) + PipelineOn -> do + writeIORef connection.pipelineStatus PipelineAborted + finishCommand PipelineOn + pure (Just (NativeResult FatalError [] [] Nothing errMap [] "")) + PipelineOff -> do + sql <- readIORef connection.currentQuery + writeIORef connection.lastError (Just (formatResultError sql errMap)) + pure (Just (NativeResult FatalError [] [] Nothing errMap [] sql)) + PortalSuspended -> + pure (Just (commandResult builder Nothing)) + ReadyForQuery txState -> do + writeIORef connection.txStatus txState + case pipeStatus of + PipelineOff -> do + writeIORef connection.asyncPending False + if builder.accHadResponse + then pure (Just (describeResult builder)) + else pure Nothing + _ -> do + writeIORef connection.pipelineStatus PipelineOn + -- A PipelineSync result is its own command boundary: unlike a + -- normal command result, libpq does not emit a separating NULL + -- after it, so consecutive syncs are reported back-to-back. We + -- therefore never set 'pipelineSeparatorPending' here. Only the + -- final sync clears 'asyncPending'; an earlier one leaves it set + -- so the next 'getNextResult' reads straight on to the next sync. + remaining <- atomicModifyIORef' connection.pendingSyncs (\n -> (n - 1, n - 1)) + when (remaining == 0) $ writeIORef connection.asyncPending False + pure (Just (NativeResult PipelineSync [] [] Nothing Map.empty [] "")) + _ -> go singleRow builder + +-- | Describe a prepared statement. +describePrepared :: Connection -> ByteString -> IO (Maybe NativeResult) +describePrepared connection name = withReady connection do + sendMessage connection (describeStatementMessage name <> syncMessage) + Just <$> collectExtended connection "" + +-- | Describe a portal. +describePortal :: Connection -> ByteString -> IO (Maybe NativeResult) +describePortal connection name = withReady connection do + sendMessage connection (describePortalMessage name <> syncMessage) + Just <$> collectExtended connection "" + +-- * Parameter projections + +paramOid :: Maybe (Word32, ByteString, Format) -> Word32 +paramOid = maybe 0 (\(oid, _, _) -> oid) + +paramFormat :: Maybe (Word32, ByteString, Format) -> Int16 +paramFormat = maybe 0 (\(_, _, format) -> formatCode format) + +paramValue :: Maybe (Word32, ByteString, Format) -> Maybe ByteString +paramValue = fmap (\(_, value, _) -> value) + +boundFormat :: Maybe (ByteString, Format) -> Int16 +boundFormat = maybe 0 (formatCode . snd) + +boundValue :: Maybe (ByteString, Format) -> Maybe ByteString +boundValue = fmap fst + +formatCode :: Format -> Int16 +formatCode = \case + Text -> 0 + Binary -> 1 + +-- * Result collection + +-- | Only run a flow on a ready connection; mirror libpq returning no result +-- when the connection is not usable. +withReady :: Connection -> IO (Maybe a) -> IO (Maybe a) +withReady connection action = do + status <- readIORef connection.connStatus + case status of + ConnectionOk -> action + _ -> pure Nothing + +-- accumulator for a result under construction +data Builder = Builder + { accFields :: [FieldDescription], + accRevRows :: [[Maybe ByteString]], + accParamOids :: [Word32], + accSawRowDescription :: Bool, + accHadResponse :: Bool + } + +emptyBuilder :: Builder +emptyBuilder = Builder [] [] [] False False + +-- | Collect the (possibly several) results of a simple query, up to +-- @ReadyForQuery@. The last is what @PQexec@ returns. +-- @CopyInResponse@ and @CopyOutResponse@ terminate the loop immediately, +-- returning a synthetic result so the caller can enter the copy sub-protocol. +collectSimple :: Connection -> ByteString -> IO [NativeResult] +collectSimple connection sql = go emptyBuilder [] + where + go builder acc = do + message <- nextMessage connection + case message of + RowDescription fs -> go builder {accFields = fs, accSawRowDescription = True} acc + DataRow values -> go builder {accRevRows = values : builder.accRevRows} acc + CommandComplete tag -> do + writeIORef connection.lastError (Just "") + go emptyBuilder (commandResult builder (Just tag) : acc) + EmptyQueryResponse -> do + writeIORef connection.lastError (Just "") + go emptyBuilder (NativeResult EmptyQuery [] [] Nothing Map.empty [] "" : acc) + ErrorResponse fs -> do + let errMap = Map.fromList fs + writeIORef connection.lastError (Just (formatResultError sql errMap)) + go emptyBuilder (NativeResult FatalError [] [] Nothing errMap [] sql : acc) + CopyInResponse _ formats -> + let fields = map copyField formats + in pure (reverse (NativeResult CopyIn fields [] Nothing Map.empty [] "" : acc)) + CopyOutResponse _ formats -> + let fields = map copyField formats + in pure (reverse (NativeResult CopyOut fields [] Nothing Map.empty [] "" : acc)) + ReadyForQuery txState -> do + writeIORef connection.txStatus txState + pure (reverse acc) + _ -> go builder acc + +-- | Collect the single result of an extended-protocol command. +collectExtended :: Connection -> ByteString -> IO NativeResult +collectExtended connection sql = go emptyBuilder Nothing + where + go builder finished = do + message <- nextMessage connection + case message of + RowDescription fs -> go builder {accFields = fs, accSawRowDescription = True} finished + ParameterDescription oids -> go builder {accParamOids = oids} finished + NoData -> go builder finished + DataRow values -> go builder {accRevRows = values : builder.accRevRows} finished + ParseComplete -> go builder finished + BindComplete -> go builder finished + CloseComplete -> go builder finished + PortalSuspended -> go emptyBuilder (finished <|> Just (commandResult builder Nothing)) + CommandComplete tag -> do + writeIORef connection.lastError (Just "") + go emptyBuilder (Just (commandResult builder (Just tag))) + EmptyQueryResponse -> do + writeIORef connection.lastError (Just "") + go emptyBuilder (Just (NativeResult EmptyQuery [] [] Nothing Map.empty [] "")) + ErrorResponse fs -> do + let errMap = Map.fromList fs + writeIORef connection.lastError (Just (formatResultError sql errMap)) + go emptyBuilder (Just (NativeResult FatalError [] [] Nothing errMap [] sql)) + ReadyForQuery txState -> do + writeIORef connection.txStatus txState + pure (fromMaybe (describeResult builder) finished) + _ -> go builder finished + +-- | A result terminated by @CommandComplete@\/@PortalSuspended@: 'TuplesOk' if a +-- row description was seen, else 'CommandOk'. +commandResult :: Builder -> Maybe ByteString -> NativeResult +commandResult builder tag = + NativeResult + (if builder.accSawRowDescription then TuplesOk else CommandOk) + builder.accFields + (reverse builder.accRevRows) + tag + Map.empty + builder.accParamOids + "" + +-- | A result with no command completion (a @Describe@\/@Parse@-only flow): +-- 'CommandOk', carrying any column descriptions and parameter OIDs. +describeResult :: Builder -> NativeResult +describeResult builder = + NativeResult + CommandOk + builder.accFields + (reverse builder.accRevRows) + Nothing + Map.empty + builder.accParamOids + "" + +lastMaybe :: [a] -> Maybe a +lastMaybe = foldl (\_ x -> Just x) Nothing + +-- | Build a synthetic 'FieldDescription' from a COPY format code (0=text, +-- 1=binary). COPY results have no column names, table OID, type OID, etc. +copyField :: Int16 -> FieldDescription +copyField fmt = + FieldDescription + { name = "", + tableOid = 0, + columnAttributeNumber = 0, + typeOid = 0, + typeSize = 0, + typeModifier = 0, + formatCode = fmt + } diff --git a/src/library/Pqi/Native/Types.hs b/src/library/Pqi/Native/Types.hs new file mode 100644 --- /dev/null +++ b/src/library/Pqi/Native/Types.hs @@ -0,0 +1,252 @@ +-- | Internal types for the native adapter, separated from 'Connection.hs' to +-- avoid orphan-instance warnings for 'IsResult' and 'IsCancel'. +module Pqi.Native.Types + ( NativeResult (..), + NativeCancel (..), + formatErrorFields, + formatResultError, + ) +where + +import Control.Exception (IOException, try) +import qualified Data.ByteString as ByteString +import qualified Data.ByteString.Char8 as ByteString.Char8 +import Data.Char (isDigit) +import Data.List (findIndex) +import qualified Data.Map.Strict as Map +import Pqi + ( ExecStatus (..), + FieldCode (..), + Format (..), + IsCancel (..), + IsResult (..), + PipelineStatus (..), + ) +import Pqi.Native.Prelude +import qualified Pqi.Native.Transport as Transport +import Pqi.Native.Transport.Message (FieldDescription (..), cancelRequest) + +-- | A fully materialized result. The native adapter buffers the entire result +-- in memory, so the accessors are pure lookups and 'unsafeFreeResult' is a +-- no-op. +data NativeResult = NativeResult + { status :: ExecStatus, + fields :: [FieldDescription], + rows :: [[Maybe ByteString]], + commandTag :: Maybe ByteString, + errorFields :: Map.Map Word8 ByteString, + paramOids :: [Word32], + -- | The query text that produced this result (used to format @LINE N:@ + -- position context in 'resultErrorMessage', matching libpq's behaviour). + -- Empty for non-error results and for async results where the query is + -- unavailable. + queryText :: ByteString + } + deriving stock (Eq, Show) + +-- | A standalone cancellation handle for the native adapter. +-- +-- Carries references to connection state so the cancel implementation can +-- decide whether a network round-trip is necessary: +-- +-- * @asyncPendingRef@: False when nothing is in flight at all — skip the +-- round-trip entirely. +-- * @pipelineStatusRef@ + @pendingCommandsRef@: in pipeline mode, +-- @asyncPending@ stays True even after all @CommandComplete@ messages +-- have arrived (while @ReadyForQuery@ is still unread). Sending a cancel +-- in that window produces a stale signal that can interrupt the very next +-- command (e.g. the @ABORT@ issued during clean-up). Checking +-- @pendingCommands > 0@ instead avoids the stale cancel: once all command +-- completions are in, @pendingCommands@ is 0 and the server is idle. +data NativeCancel = NativeCancel + { host :: ByteString, + port :: Int, + pid :: Int32, + secret :: Int32, + asyncPendingRef :: IORef Bool, + pipelineStatusRef :: IORef PipelineStatus, + pendingCommandsRef :: IORef Int + } + +-- | Like 'formatResultError' but without a client query text (for +-- connection-level errors, which never carry a statement-position field). +formatErrorFields :: Map.Map Word8 ByteString -> ByteString +formatErrorFields = formatResultError "" + +-- | Format error\/notice fields into a message string matching libpq's +-- @PQresultErrorMessage@ at DEFAULT verbosity with @SHOW_CONTEXT_ERRORS@ +-- visibility. +-- +-- Field order (matching libpq): +-- 1. @\: \\\n@ +-- 2. @LINE N: \\\n \^\\n@ (statement position via @queryText@ + @\'P\'@ field) +-- 3. @LINE N: \\\n \^\\nQUERY: \\\n@ (@\'q\'@+@\'p\'@ fields) +-- 4. @DETAIL: \\\n@ +-- 5. @HINT: \\\n@ +-- 6. @CONTEXT: \\\n@ (ERROR\/FATAL\/PANIC only) +formatResultError :: ByteString -> Map.Map Word8 ByteString -> ByteString +formatResultError queryText fields = + case Map.lookup 0x4d fields of + Nothing -> "" + Just msg -> + let sev = Map.findWithDefault "" 0x53 fields + isError = sev `elem` ["ERROR", "FATAL", "PANIC"] + line1 = (if ByteString.null sev then msg else sev <> ": " <> msg) <> "\n" + -- Statement position ('P', 0x50): needs client query text + stmtCtx = case Map.lookup 0x50 fields of + Just posStr + | not (ByteString.null queryText) -> + maybe "" (positionContext queryText) (readPositiveInt posStr) + _ -> "" + -- Internal query+position ('q'=0x71, 'p'=0x70): both in wire fields + intCtx = case (Map.lookup 0x71 fields, Map.lookup 0x70 fields) of + (Just intQuery, Just posStr) -> + maybe + "" + (\p -> positionContext intQuery p <> "QUERY: " <> intQuery <> "\n") + (readPositiveInt posStr) + _ -> "" + detLine = case Map.lookup 0x44 fields of + Just det -> "DETAIL: " <> det <> "\n" + _ -> "" + hntLine = case Map.lookup 0x48 fields of + Just hnt -> "HINT: " <> hnt <> "\n" + _ -> "" + ctxLine = case Map.lookup 0x57 fields of + Just ctx | isError -> "CONTEXT: " <> ctx <> "\n" + _ -> "" + in line1 <> stmtCtx <> intCtx <> detLine <> hntLine <> ctxLine + +-- | Build the @LINE N: \\\n \^\\n@ block for a 1-indexed position +-- within a query string, matching libpq's formatting exactly. +positionContext :: ByteString -> Int -> ByteString +positionContext query pos = + let pos0 = max 0 (pos - 1) + before = ByteString.take pos0 query + lineNum = 1 + ByteString.length (ByteString.filter (== 0x0a) before) + lineStart = maybe 0 (+ 1) (ByteString.elemIndexEnd 0x0a before) + col = pos0 - lineStart + rest = ByteString.drop lineStart query + lineText = ByteString.takeWhile (/= 0x0a) rest + prefix = ByteString.Char8.pack ("LINE " <> show lineNum <> ": ") + caret = ByteString.replicate (ByteString.length prefix + col) 0x20 <> "^\n" + in prefix <> lineText <> "\n" <> caret + +readPositiveInt :: ByteString -> Maybe Int +readPositiveInt bs = case ByteString.Char8.readInt bs of + Just (n, _) | n > 0 -> Just n + _ -> Nothing + +instance IsResult NativeResult where + resultStatus result = pure result.status + resultErrorMessage result = pure (Just (formatResultError result.queryText result.errorFields)) + resultErrorField result field = pure (Map.lookup (fieldCodeByte field) result.errorFields) + unsafeFreeResult _ = pure () + ntuples result = pure (fromIntegral (length result.rows)) + nfields result = pure (fromIntegral (length result.fields)) + fname result column = pure $ do + fd <- atMay result.fields column + if ByteString.null fd.name then Nothing else Just fd.name + fnumber result name = + pure (fromIntegral <$> findIndex (\field -> field.name == folded) result.fields) + where + folded = foldIdentifier name + ftable result column = pure (maybe 0 (.tableOid) (atMay result.fields column)) + ftablecol result column = + pure (maybe 0 (\field -> fromIntegral (field.columnAttributeNumber :: Int16)) (atMay result.fields column)) + fformat result column = + pure (maybe Text (\field -> formatOf field.formatCode) (atMay result.fields column)) + ftype result column = pure (maybe 0 (.typeOid) (atMay result.fields column)) + fmod result column = pure (maybe 0 (\field -> fromIntegral (field.typeModifier :: Int32)) (atMay result.fields column)) + fsize result column = pure (maybe 0 (\field -> fromIntegral (field.typeSize :: Int16)) (atMay result.fields column)) + getvalue result row column = pure (join (cellAt result row column)) + getvalue' result row column = pure (join (cellAt result row column)) + getisnull result row column = pure (maybe True isNothing (cellAt result row column)) + getlength result row column = + pure (maybe 0 (maybe 0 ByteString.length) (cellAt result row column)) + nparams result = pure (fromIntegral (length result.paramOids)) + paramtype result index = pure (fromMaybe 0 (atMay result.paramOids index)) + cmdStatus result = pure (Just (fromMaybe "" result.commandTag)) + cmdTuples result = pure (Just (maybe "" affectedRows result.commandTag)) + +instance IsCancel NativeCancel where + cancel nc = do + pending <- readIORef nc.asyncPendingRef + if not pending + then pure (Right ()) + else do + transport <- Transport.connect nc.host nc.port + Transport.send transport (cancelRequest nc.pid nc.secret) + -- Read until EOF to ensure the server has processed the cancel request + -- before we close the connection. This matches libpq's PQcancel behavior + -- and prevents the cancel signal from racing with the next query. + _ <- try @IOException (Transport.readUntilClosed transport) + Transport.close transport + pure (Right ()) + +-- * Helpers for the 'IsResult' instance + +atMay :: [a] -> Int32 -> Maybe a +atMay xs i + | i < 0 = Nothing + | otherwise = case drop (fromIntegral i) xs of + x : _ -> Just x + [] -> Nothing + +cellAt :: NativeResult -> Int32 -> Int32 -> Maybe (Maybe ByteString) +cellAt result row column = do + rowValues <- atMay result.rows row + atMay rowValues column + +formatOf :: Int16 -> Format +formatOf = \case + 1 -> Binary + _ -> Text + +foldIdentifier :: ByteString -> ByteString +foldIdentifier = ByteString.pack . outside . ByteString.unpack + where + quote = 0x22 + outside = \case + [] -> [] + c : rest + | c == quote -> inside rest + | otherwise -> asciiToLower c : outside rest + inside = \case + [] -> [] + c : rest + | c == quote -> case rest of + c' : rest' | c' == quote -> quote : inside rest' + _ -> outside rest + | otherwise -> c : inside rest + asciiToLower c + | c >= 0x41 && c <= 0x5a = c + 0x20 + | otherwise = c + +fieldCodeByte :: FieldCode -> Word8 +fieldCodeByte = \case + DiagSeverity -> 0x53 -- 'S' + DiagSqlstate -> 0x43 -- 'C' + DiagMessagePrimary -> 0x4d -- 'M' + DiagMessageDetail -> 0x44 -- 'D' + DiagMessageHint -> 0x48 -- 'H' + DiagStatementPosition -> 0x50 -- 'P' + DiagInternalPosition -> 0x70 -- 'p' + DiagInternalQuery -> 0x71 -- 'q' + DiagContext -> 0x57 -- 'W' + DiagSourceFile -> 0x46 -- 'F' + DiagSourceLine -> 0x4c -- 'L' + DiagSourceFunction -> 0x52 -- 'R' + +-- | The affected-row count from a @CommandComplete@ tag, matching +-- @PQcmdTuples@: the last whitespace-delimited token if it is all digits (which +-- covers @INSERT oid rows@, @UPDATE n@, @SELECT n@, …), else empty. +affectedRows :: ByteString -> ByteString +affectedRows tag = + case ByteString.Char8.words tag of + [] -> "" + tokens -> + let final = last tokens + in if not (ByteString.null final) && ByteString.Char8.all isDigit final + then final + else "" diff --git a/src/test/Spec.hs b/src/test/Spec.hs new file mode 100644 --- /dev/null +++ b/src/test/Spec.hs @@ -0,0 +1,14 @@ +-- | The native adapter's conformance test suite: a single delegate to +-- 'Pqi.Conformance.specs', which brings up a throwaway PostgreSQL +-- container, runs the full differential battery (core, capabilities, and +-- SCRAM) against the FFI reference, and tears the container down again. +module Main (main) where + +import Data.Proxy (Proxy (..)) +import Pqi.Conformance (specs) +import Pqi.Native (Connection) +import Test.Hspec +import Prelude + +main :: IO () +main = hspec (specs (Proxy @Connection)) diff --git a/src/transport/Pqi/Native/Transport.hs b/src/transport/Pqi/Native/Transport.hs new file mode 100644 --- /dev/null +++ b/src/transport/Pqi/Native/Transport.hs @@ -0,0 +1,119 @@ +-- | The byte-level transport: a TCP socket with a read buffer, plus the framing +-- that turns the stream into discrete @[type byte][Int32 length][body]@ +-- messages. +module Pqi.Native.Transport + ( Transport, + connect, + unconnected, + close, + send, + receiveFrame, + socketFd, + peerIp, + readUntilClosed, + ) +where + +import qualified Data.ByteString as ByteString +import qualified Data.ByteString.Char8 as ByteString.Char8 +import Data.IORef +import qualified Network.Socket as Socket +import qualified Network.Socket.ByteString as Socket.ByteString +import Pqi.Native.Transport.Prelude +import qualified PtrPoker.Write as Poker +import System.IO.Error (eofErrorType, mkIOError) + +-- | An open connection's byte transport. +data Transport = Transport + { socket :: Socket.Socket, + readBuffer :: IORef ByteString + } + +-- | Open a TCP connection to the given host and port. +connect :: ByteString -> Int -> IO Transport +connect host port = do + let hints = Socket.defaultHints {Socket.addrSocketType = Socket.Stream} + addresses <- + Socket.getAddrInfo (Just hints) (Just (ByteString.Char8.unpack host)) (Just (show port)) + case addresses of + [] -> ioError (userError ("Could not resolve host: " <> ByteString.Char8.unpack host)) + address : _ -> do + sock <- Socket.socket (Socket.addrFamily address) (Socket.addrSocketType address) (Socket.addrProtocol address) + Socket.connect sock (Socket.addrAddress address) + buffer <- newIORef ByteString.empty + pure Transport {socket = sock, readBuffer = buffer} + +-- | An unconnected transport, for representing a \"null\" connection. Its +-- socket is allocated but never connected; it must not be used for I\/O. +unconnected :: IO Transport +unconnected = do + sock <- Socket.socket Socket.AF_INET Socket.Stream Socket.defaultProtocol + buffer <- newIORef ByteString.empty + pure Transport {socket = sock, readBuffer = buffer} + +-- | Close the connection. +close :: Transport -> IO () +close transport = Socket.close transport.socket + +-- | The underlying socket file descriptor. +socketFd :: Transport -> IO Int32 +socketFd transport = fromIntegral <$> Socket.unsafeFdSocket transport.socket + +-- | Send a serialized message. +send :: Transport -> Poker.Write -> IO () +send transport write = Socket.ByteString.sendAll transport.socket (Poker.toByteString write) + +-- | Read exactly @n@ bytes, looping over @recv@ (which yields up to @n@) and +-- buffering any overshoot. Throws on EOF before @n@ bytes arrive. +receiveExactly :: Transport -> Int -> IO ByteString +receiveExactly transport n = do + buffered <- readIORef transport.readBuffer + go buffered + where + go accumulated + | ByteString.length accumulated >= n = do + let (result, rest) = ByteString.splitAt n accumulated + writeIORef transport.readBuffer rest + pure result + | otherwise = do + chunk <- Socket.ByteString.recv transport.socket (max 4096 (n - ByteString.length accumulated)) + if ByteString.null chunk + then ioError (mkIOError eofErrorType "pqi-native: connection closed by server" Nothing Nothing) + else go (accumulated <> chunk) + +-- | Receive one framed message: its type byte and its body (the length prefix, +-- which counts itself, is consumed). +receiveFrame :: Transport -> IO (Word8, ByteString) +receiveFrame transport = do + header <- receiveExactly transport 5 + let typeByte = ByteString.head header + frameLength = decodeInt32BE (ByteString.drop 1 header) + bodyLength = frameLength - 4 + body <- + if bodyLength > 0 + then receiveExactly transport bodyLength + else pure ByteString.empty + pure (typeByte, body) + +-- | Read bytes until the peer closes the connection (EOF), discarding them. +-- Used by the cancel path to mirror libpq's behaviour: keep the socket open +-- until the server has read the cancel request and closed its end. +readUntilClosed :: Transport -> IO () +readUntilClosed transport = go + where + go = do + chunk <- Socket.ByteString.recv transport.socket 4096 + if ByteString.null chunk + then pure () + else go + +-- | The numeric IP address of the connected peer (e.g. @\"::1\"@ or +-- @\"127.0.0.1\"@). Throws if the socket has no peer (unconnected). +peerIp :: Transport -> IO ByteString +peerIp transport = do + addr <- Socket.getPeerName transport.socket + (Just ip, _) <- Socket.getNameInfo [Socket.NI_NUMERICHOST] True False addr + pure (ByteString.Char8.pack ip) + +decodeInt32BE :: ByteString -> Int +decodeInt32BE = ByteString.foldl' (\acc w -> acc * 256 + fromIntegral w) 0 . ByteString.take 4 diff --git a/src/transport/Pqi/Native/Transport/Message.hs b/src/transport/Pqi/Native/Transport/Message.hs new file mode 100644 --- /dev/null +++ b/src/transport/Pqi/Native/Transport/Message.hs @@ -0,0 +1,302 @@ +-- | The PostgreSQL frontend\/backend wire messages: builders ('Poker.Write') for +-- the frontend messages we send, and a decoder for the backend messages we +-- receive. Message framing (@[type byte][Int32 length][body]@, with the +-- startup message lacking a type byte) lives here too. +module Pqi.Native.Transport.Message + ( -- * Frontend messages + startupMessage, + sslRequest, + cancelRequest, + passwordMessage, + saslInitialResponse, + saslResponse, + queryMessage, + parseMessage, + bindMessage, + describeStatementMessage, + describePortalMessage, + executeMessage, + closeStatementMessage, + closePortalMessage, + syncMessage, + flushMessage, + terminateMessage, + copyDataMessage, + copyDoneMessage, + copyFailMessage, + + -- * Backend messages + BackendMessage (..), + FieldDescription (..), + decodeBackendMessage, + ) +where + +import qualified Data.ByteString as ByteString +import Pqi.Native.Comms +import Pqi.Native.Transport.Prelude +import qualified PtrPoker.Write as Poker + +-- * Framing helpers + +-- | Frame a typed frontend message: type byte, then a length covering the +-- length field and the body, then the body. +framed :: Word8 -> Poker.Write -> Poker.Write +framed typeByte body = + Poker.word8 typeByte + <> Poker.bInt32 (fromIntegral (4 + Poker.writeSize body)) + <> body + +-- | A null-terminated string. +cstr :: ByteString -> Poker.Write +cstr value = Poker.byteString value <> Poker.word8 0 + +formatCodes :: [Int16] -> Poker.Write +formatCodes codes = + Poker.bInt16 (fromIntegral (length codes)) <> foldMap Poker.bInt16 codes + +-- * Frontend messages + +-- | The startup message (no type byte), requesting protocol 3.0 with the given +-- parameters (e.g. @user@, @database@). +startupMessage :: [(ByteString, ByteString)] -> Poker.Write +startupMessage parameters = + let body = + Poker.bInt32 196608 -- protocol version 3.0 == (3 << 16) + <> foldMap (\(k, v) -> cstr k <> cstr v) parameters + <> Poker.word8 0 + in Poker.bInt32 (fromIntegral (4 + Poker.writeSize body)) <> body + +-- | The SSL negotiation request (no type byte). +sslRequest :: Poker.Write +sslRequest = Poker.bInt32 8 <> Poker.bInt32 80877103 + +-- | A cancel request (no type byte), sent on a fresh connection: the cancel +-- request code followed by the target backend's process ID and secret key. +cancelRequest :: Int32 -> Int32 -> Poker.Write +cancelRequest processId secretKey = + Poker.bInt32 16 <> Poker.bInt32 80877102 <> Poker.bInt32 processId <> Poker.bInt32 secretKey + +-- | A cleartext or MD5 password response (message type @p@). +passwordMessage :: ByteString -> Poker.Write +passwordMessage password = framed 0x70 (cstr password) + +-- | A SASL initial response (message type @p@): the chosen mechanism and the +-- client-first message. +saslInitialResponse :: ByteString -> ByteString -> Poker.Write +saslInitialResponse mechanism initialData = + framed 0x70 + $ cstr mechanism + <> Poker.bInt32 (fromIntegral (ByteString.length initialData)) + <> Poker.byteString initialData + +-- | A subsequent SASL response (message type @p@). +saslResponse :: ByteString -> Poker.Write +saslResponse responseData = framed 0x70 (Poker.byteString responseData) + +-- | A simple @Query@ message. +queryMessage :: ByteString -> Poker.Write +queryMessage sql = framed 0x51 (cstr sql) + +-- | A @Parse@ message: statement name, SQL, and parameter type OIDs. +parseMessage :: ByteString -> ByteString -> [Word32] -> Poker.Write +parseMessage name sql oids = + framed 0x50 + $ cstr name + <> cstr sql + <> Poker.bInt16 (fromIntegral (length oids)) + <> foldMap Poker.bWord32 oids + +-- | A @Bind@ message: portal, statement, per-parameter format codes, parameter +-- values ('Nothing' for SQL @NULL@), and result format codes. +bindMessage :: ByteString -> ByteString -> [Int16] -> [Maybe ByteString] -> [Int16] -> Poker.Write +bindMessage portal statement parameterFormats values resultFormats = + framed 0x42 + $ cstr portal + <> cstr statement + <> formatCodes parameterFormats + <> Poker.bInt16 (fromIntegral (length values)) + <> foldMap bindValue values + <> formatCodes resultFormats + where + bindValue = \case + Nothing -> Poker.bInt32 (-1) + Just value -> Poker.bInt32 (fromIntegral (ByteString.length value)) <> Poker.byteString value + +-- | A @Describe@ message targeting a prepared statement. +describeStatementMessage :: ByteString -> Poker.Write +describeStatementMessage name = framed 0x44 (Poker.word8 0x53 <> cstr name) + +-- | A @Describe@ message targeting a portal. +describePortalMessage :: ByteString -> Poker.Write +describePortalMessage name = framed 0x44 (Poker.word8 0x50 <> cstr name) + +-- | An @Execute@ message: portal name and a maximum row count (0 = unlimited). +executeMessage :: ByteString -> Int32 -> Poker.Write +executeMessage portal maxRows = framed 0x45 (cstr portal <> Poker.bInt32 maxRows) + +-- | A @Close@ message targeting a prepared statement. +closeStatementMessage :: ByteString -> Poker.Write +closeStatementMessage name = framed 0x43 (Poker.word8 0x53 <> cstr name) + +-- | A @Close@ message targeting a portal. +closePortalMessage :: ByteString -> Poker.Write +closePortalMessage name = framed 0x43 (Poker.word8 0x50 <> cstr name) + +-- | A @Sync@ message. +syncMessage :: Poker.Write +syncMessage = framed 0x53 mempty + +-- | A @Flush@ message. +flushMessage :: Poker.Write +flushMessage = framed 0x48 mempty + +-- | A @Terminate@ message. +terminateMessage :: Poker.Write +terminateMessage = framed 0x58 mempty + +-- | A @CopyData@ message. +copyDataMessage :: ByteString -> Poker.Write +copyDataMessage payload = framed 0x64 (Poker.byteString payload) + +-- | A @CopyDone@ message. +copyDoneMessage :: Poker.Write +copyDoneMessage = framed 0x63 mempty + +-- | A @CopyFail@ message with an error string. +copyFailMessage :: ByteString -> Poker.Write +copyFailMessage reason = framed 0x66 (cstr reason) + +-- * Backend messages + +-- | The description of one result column, from a @RowDescription@ message. +data FieldDescription = FieldDescription + { name :: ByteString, + tableOid :: Word32, + columnAttributeNumber :: Int16, + typeOid :: Word32, + typeSize :: Int16, + typeModifier :: Int32, + formatCode :: Int16 + } + deriving stock (Eq, Show) + +-- | A decoded backend message. Only the variants the adapter acts on are +-- distinguished; anything else becomes 'UnknownMessage'. +data BackendMessage + = AuthenticationOk + | AuthenticationCleartextPassword + | AuthenticationMD5Password ByteString + | AuthenticationSASL [ByteString] + | AuthenticationSASLContinue ByteString + | AuthenticationSASLFinal ByteString + | ParameterStatus ByteString ByteString + | BackendKeyData Int32 Int32 + | ReadyForQuery Word8 + | RowDescription [FieldDescription] + | DataRow [Maybe ByteString] + | CommandComplete ByteString + | EmptyQueryResponse + | ErrorResponse [(Word8, ByteString)] + | NoticeResponse [(Word8, ByteString)] + | ParseComplete + | BindComplete + | CloseComplete + | ParameterDescription [Word32] + | NoData + | PortalSuspended + | NotificationResponse Int32 ByteString ByteString + | CopyInResponse Int16 [Int16] + | CopyOutResponse Int16 [Int16] + | CopyData ByteString + | CopyDone + | FunctionCallResponse (Maybe ByteString) + | UnknownMessage Word8 ByteString + deriving stock (Eq, Show) + +-- | Decode a backend message from its type byte and body. +decodeBackendMessage :: Word8 -> ByteString -> Either DecodingError BackendMessage +decodeBackendMessage typeByte body = + case typeByte of + 0x52 -> runDecoder authentication body + 0x53 -> runDecoder (ParameterStatus <$> cstring <*> cstring) body + 0x4b -> runDecoder (BackendKeyData <$> int32 <*> int32) body + 0x5a -> runDecoder (ReadyForQuery <$> word8) body + 0x54 -> runDecoder (RowDescription <$> repeatedInt16 fieldDescription) body + 0x44 -> runDecoder (DataRow <$> repeatedInt16 columnValue) body + 0x43 -> runDecoder (CommandComplete <$> cstring) body + 0x49 -> Right EmptyQueryResponse + 0x45 -> runDecoder (ErrorResponse <$> noticeFields) body + 0x4e -> runDecoder (NoticeResponse <$> noticeFields) body + 0x31 -> Right ParseComplete + 0x32 -> Right BindComplete + 0x33 -> Right CloseComplete + 0x74 -> runDecoder (ParameterDescription <$> repeatedInt16 word32) body + 0x6e -> Right NoData + 0x73 -> Right PortalSuspended + 0x41 -> runDecoder (NotificationResponse <$> int32 <*> cstring <*> cstring) body + 0x47 -> runDecoder (uncurry CopyInResponse <$> copyResponse) body + 0x48 -> runDecoder (uncurry CopyOutResponse <$> copyResponse) body + 0x64 -> runDecoder (CopyData <$> remaining) body + 0x63 -> Right CopyDone + 0x56 -> runDecoder (FunctionCallResponse <$> columnValue) body + other -> Right (UnknownMessage other body) + +-- | Repeat a decoder an @Int16@-prefixed number of times. +repeatedInt16 :: Decoder a -> Decoder [a] +repeatedInt16 element = do + count <- int16 + replicateM (fromIntegral count) element + +fieldDescription :: Decoder FieldDescription +fieldDescription = + FieldDescription + <$> cstring + <*> word32 + <*> int16 + <*> word32 + <*> int16 + <*> int32 + <*> int16 + +columnValue :: Decoder (Maybe ByteString) +columnValue = do + len <- int32 + if len < 0 + then pure Nothing + else Just <$> bytes (fromIntegral len) + +-- | The fields of an @ErrorResponse@\/@NoticeResponse@: @(code, value)@ pairs, +-- terminated by a zero code byte. +noticeFields :: Decoder [(Word8, ByteString)] +noticeFields = do + code <- word8 + if code == 0 + then pure [] + else do + value <- cstring + ((code, value) :) <$> noticeFields + +copyResponse :: Decoder (Int16, [Int16]) +copyResponse = do + overall <- fromIntegral <$> word8 + columns <- repeatedInt16 int16 + pure (overall, columns) + +authentication :: Decoder BackendMessage +authentication = do + subtype <- int32 + case subtype of + 0 -> pure AuthenticationOk + 3 -> pure AuthenticationCleartextPassword + 5 -> AuthenticationMD5Password <$> bytes 4 + 10 -> AuthenticationSASL <$> saslMechanisms + 11 -> AuthenticationSASLContinue <$> remaining + 12 -> AuthenticationSASLFinal <$> remaining + other -> throwE (UnexpectedAuthType other) + where + saslMechanisms = do + mechanism <- cstring + if ByteString.null mechanism + then pure [] + else (mechanism :) <$> saslMechanisms diff --git a/src/transport/Pqi/Native/Transport/Prelude.hs b/src/transport/Pqi/Native/Transport/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/transport/Pqi/Native/Transport/Prelude.hs @@ -0,0 +1,19 @@ +module Pqi.Native.Transport.Prelude + ( module Exports, + ) +where + +import Control.Applicative as Exports +import Control.Monad as Exports +import Control.Monad.Trans.Class as Exports (lift) +import Control.Monad.Trans.Except as Exports + ( ExceptT (..), + except, + runExceptT, + throwE, + withExceptT, + ) +import Data.ByteString as Exports (ByteString) +import Data.Int as Exports +import Data.Word as Exports +import Prelude as Exports