hpgsql 0.1.0.0 → 0.2.0.0

raw patch · 8 files changed

+347/−58 lines, 8 filesdep +cryptondep +memorydep −cryptohash-md5PVP ok

version bump matches the API change (PVP)

Dependencies added: crypton, memory

Dependencies removed: cryptohash-md5

API changes (from Hackage documentation)

Files

• hpgsql.cabal +4/−2
• src/Hpgsql/Connection.hs +2/−1
• src/Hpgsql/Internal.hs +70/−15
• src/Hpgsql/InternalTypes.hs +13/−2
• src/Hpgsql/Msgs.hs +44/−19
• src/Hpgsql/Pipeline.hs +3/−1
• src/Hpgsql/ScramSHA256.hs +192/−0
• src/Hpgsql/Transaction.hs +19/−18

@@ -1,7 +1,7 @@ cabal-version: 2.0  name:           hpgsql-version:        0.1.0.0+version:        0.2.0.0 synopsis:       Pure Haskell PostgreSQL driver (no libpq) description:    hpgsql is a pure Haskell implementation of a PostgreSQL driver category:       Database@@ -47,6 +47,7 @@       Hpgsql.Msgs       Hpgsql.Networking       Hpgsql.QueryInternal+      Hpgsql.ScramSHA256       Hpgsql.SimpleParser       Hpgsql.TransactionStatusInternal       Paths_hpgsql@@ -99,7 +100,8 @@     case-insensitive >= 1.2 && < 1.3,     cereal >= 0.5 && < 0.6,     containers >= 0.6 && < 0.8,-    cryptohash-md5 >= 0.11.7.1 && < 0.12,+    crypton >= 1.0.0 && < 1.1,+    memory >= 0.18.0 && < 0.19,     hashable >= 1.5 && < 1.6,     haskell-src-meta >= 0.8 && < 0.9,     network >= 3.2 && < 3.3,

@@ -6,6 +6,7 @@     withConnectionOpts,     closeGracefully,     closeForcefully,+    connectionIsClosed,     ConnectionString (..),     parseLibpqConnectionString,     ResetConnectionOpts (..),@@ -48,7 +49,7 @@ import Data.Text (Text) import qualified Data.Text as Text import Data.Text.Encoding (encodeUtf8)-import Hpgsql.Internal (closeForcefully, closeGracefully, connect, connectOpts, defaultConnectOpts, getBackendPid, getParameterStatus, refreshTypeInfoCache, resetConnectionState, resetTypeInfoCache, withConnection, withConnectionOpts)+import Hpgsql.Internal (closeForcefully, closeGracefully, connect, connectOpts, connectionIsClosed, defaultConnectOpts, getBackendPid, getParameterStatus, refreshTypeInfoCache, resetConnectionState, resetTypeInfoCache, withConnection, withConnectionOpts) import Hpgsql.InternalTypes (ConnectionString (..), ResetConnectionOpts (..)) import Network.URI   ( URI (..),

@@ -7,6 +7,7 @@     withConnectionOpts,     closeGracefully,     closeForcefully,+    connectionIsClosed,      -- * Query     query,@@ -97,7 +98,6 @@     debugPrint,     _globalDebugLock,     timeDebugNonBlockingOperation,-    whenNotClosed,     chunkBuildersBySize,     pipelineExecInternal,     pipelineM,@@ -112,7 +112,7 @@ import Control.Concurrent.MVar (MVar, newMVar) import Control.Concurrent.STM (STM, TVar) import qualified Control.Concurrent.STM as STM-import Control.Exception.Safe (MonadThrow, SomeException, bracket, bracketOnError, finally, handle, mask, mask_, onException, throw, toException, tryJust)+import Control.Exception.Safe (Exception (..), MonadThrow, SomeException, bracket, bracketOnError, finally, handleJust, mask, mask_, onException, throw, toException, tryJust) import Control.Monad (forM, forM_, join, unless, void, when) import Data.ByteString (ByteString) import Data.ByteString.Internal (w2c)@@ -139,10 +139,11 @@ import Hpgsql.Encoding.RowDecoderMonadic (ConversionState (..), RowDecoderMonadic (..)) import Hpgsql.InternalTypes (BindComplete (..), CommandComplete (..), ConnectOpts (..), ConnectionString (..), CopyInResponse (..), CopyQueryState (..), DataRow (..), Either3 (..), EncodingContext (..), ErrorDetail (..), ErrorResponse (..), HPgConnection (..), InternalConnectionState (..), IrrecoverableHpgsqlError (..), NoData (..), NotificationResponse (..), ParseComplete (..), Pipeline (..), PostgresError (..), Query (..), QueryId (..), QueryProtocol (..), QueryState (..), ReadyForQuery (..), ResetConnectionOpts (..), ResponseMsg (..), ResponseMsgsReceived (..), RowDescription (..), SingleQuery (..), TransactionStatus (..), WeakThreadId (..), mkMutex, queryToByteString, throwIrrecoverableError) import Hpgsql.Locking (getMyWeakThreadId, withMutex)-import Hpgsql.Msgs (AuthenticationMethod (..), AuthenticationResponse (..), BackendKeyData (..), Bind (..), CancelRequest (..), CopyData (..), CopyDone (..), Describe (..), Execute (..), FromPgMessage (..), NoticeResponse (..), ParameterStatus (..), Parse (..), PasswordMessage (..), PgMsgParser (..), StartupMessage (..), Sync (..), Terminate (..), ToPgMessage (..), parsePgMessage)+import Hpgsql.Msgs (AuthenticationMethod (..), AuthenticationResponse (..), BackendKeyData (..), Bind (..), CancelRequest (..), CopyData (..), CopyDone (..), Describe (..), Execute (..), FromPgMessage (..), NoticeResponse (..), ParameterStatus (..), Parse (..), PasswordMessage (..), PgMsgParser (..), SASLInitialResponse (..), SASLResponse (..), StartupMessage (..), Sync (..), Terminate (..), ToPgMessage (..), parsePgMessage) import qualified Hpgsql.Msgs as Msgs import Hpgsql.Networking (recvNonBlocking, sendNonBlocking, socketWaitRead, socketWaitWrite) import Hpgsql.Query (breakQueryIntoStatements)+import qualified Hpgsql.ScramSHA256 as ScramSHA256 import qualified Hpgsql.SimpleParser as Parser import Hpgsql.TypeInfo (ArrayTypeDetails (..), TypeDetails (..), TypeInfo (..), buildTypeInfoCache, builtinPgTypesMap) import Network.Socket (AddrInfo (..))@@ -284,6 +285,43 @@             AuthMD5Password _ -> do               nonAtomicSendMsg hpgConnPartialDoNotReturn $ PasswordMessage authMethod (Text.unpack user) (Text.unpack password)               receiveAuthOkOrThrow hpgConnPartialDoNotReturn+            AuthSASL mechanisms+              | "SCRAM-SHA-256" `elem` mechanisms -> do+                  -- 1. Generate and send client-first-message+                  clientFirstMsg <- ScramSHA256.generateClientFirstMessage+                  nonAtomicSendMsg hpgConnPartialDoNotReturn $ SASLInitialResponse "SCRAM-SHA-256" clientFirstMsg.fullMessage++                  -- 2. Receive server-first-message+                  saslContinueMsg <-+                    receiveNextMsgUnsafe+                      hpgConnPartialDoNotReturn+                      (Right <$> msgParser @AuthenticationResponse <|> Left <$> msgParser @ErrorResponse)+                  serverFirst <- case saslContinueMsg of+                    Right (AuthenticationResponse (AuthSASLContinue d)) -> pure d+                    Left errResp -> throw $ IrrecoverableHpgsqlError {hpgsqlDetails = "SCRAM-SHA-256 authentication failed", innerException = Just $ toException $ mkPostgresError "" errResp, relatedStatement = Nothing}+                    _ -> throwIrrecoverableError "Expected AuthSASLContinue during SCRAM-SHA-256 authentication"++                  -- 3. Generate and send client-final-message+                  clientFinalMsg <- case ScramSHA256.handleServerFirstMsg password clientFirstMsg serverFirst of+                    Left err -> throwIrrecoverableError $ "SCRAM-SHA-256 error: " <> Text.pack err+                    Right r -> pure r+                  nonAtomicSendMsg hpgConnPartialDoNotReturn $ SASLResponse clientFinalMsg++                  -- 4. Receive server-final-message+                  saslFinalMsg <-+                    receiveNextMsgUnsafe+                      hpgConnPartialDoNotReturn+                      (Right <$> msgParser @AuthenticationResponse <|> Left <$> msgParser @ErrorResponse)+                  case saslFinalMsg of+                    Right (AuthenticationResponse (AuthSASLFinal serverFinalData)) ->+                      case ScramSHA256.verifyServerFinal clientFinalMsg serverFinalData of+                        Left err -> throwIrrecoverableError $ "SCRAM-SHA-256 server verification failed: " <> Text.pack err+                        Right () -> pure ()+                    Left errResp -> throw $ IrrecoverableHpgsqlError {hpgsqlDetails = "SCRAM-SHA-256 authentication failed", innerException = Just $ toException $ mkPostgresError "" errResp, relatedStatement = Nothing}+                    _ -> throwIrrecoverableError "Expected AuthSASLFinal during SCRAM-SHA-256 authentication"++                  -- 5. Receive AuthOk+                  receiveAuthOkOrThrow hpgConnPartialDoNotReturn             _ -> throwIrrecoverableError $ "Hpgsql does not yet support authenticating with method " <> Text.pack (show authMethod)           errorOrBackendKeyData <- receiveNextMsgUnsafe hpgConnPartialDoNotReturn $ Right <$> msgParser @BackendKeyData <|> Left <$> msgParser @ErrorResponse           case errorOrBackendKeyData of@@ -312,7 +350,7 @@             AuthOk -> pure ()             _ -> throwIrrecoverableError "Failed to authenticate user." -    getConnectedSocket = do+    getConnectedSocket = rethrowAsIrrecoverable $ do       addrInfo <- case connectOrCancel of         CancelNotConnect _ addrInfo -> pure addrInfo         Connect ->@@ -398,11 +436,19 @@   -- running "UNLISTEN *" those would be received, so this might be fine as long as we   -- clear the internal queue _after_ "UNLISTEN *".   do-    let qs = (if unlistenAll cleanOpts then ["UNLISTEN *"] else []) ++ (if resetAll cleanOpts then ["RESET ALL", "RESET ROLE"] else [])+    let qs =+          (if cleanOpts.unlistenAll then ["UNLISTEN *"] else [])+            ++ (if cleanOpts.closeAllCursors then ["CLOSE ALL"] else [])+            ++ (if cleanOpts.resetAll then ["RESET ALL"] else [])+            ++ (if cleanOpts.resetRole then ["SET SESSION AUTHORIZATION DEFAULT"] else [])+            ++ (if cleanOpts.deallocatePreparedStmts then ["DEALLOCATE ALL; DISCARD PLANS;"] else [])+            ++ (if cleanOpts.releaseAdvisoryLocks then ["SELECT pg_advisory_unlock_all()"] else [])+            ++ (if cleanOpts.discardTempTables then ["DISCARD TEMP"] else [])+            ++ (if cleanOpts.discardSequences then ["DISCARD SEQUENCES"] else [])     runPipeline conn (traverse pipelineExec_ qs) >>= sequence_   when (unlistenAll cleanOpts) clearInternalNotificationQueue   where-    cleanOpts = fromMaybe ResetConnectionOpts {resetAll = True, unlistenAll = True, checkTransactionState = True} mCleanOpts+    cleanOpts = fromMaybe ResetConnectionOpts {closeAllCursors = True, resetAll = True, resetRole = True, deallocatePreparedStmts = False, unlistenAll = True, releaseAdvisoryLocks = True, discardTempTables = True, discardSequences = True, checkTransactionState = True} mCleanOpts     clearInternalNotificationQueue = STM.atomically $ do       st <- STM.readTVar internalConnectionState       emptyQueue <- STM.newTQueue@@ -411,7 +457,7 @@ -- | Closes the connection with postgres. Do not use this in exception handlers; use `closeForcefully` -- instead. closeGracefully :: HPgConnection -> IO ()-closeGracefully conn@(HPgConnection {socket}) = whenNotClosed conn $ flip finally (Socket.close socket >> modifyMVar_ conn.socketClosed (const $ pure True)) $ do+closeGracefully conn = closeInternalFinally conn $ do   withControlMsgsLock     conn     (const $ pure ())@@ -419,21 +465,23 @@     $ \() -> do       nonAtomicSendMsg conn Terminate --- | Closes the connection with postgres as quickly as possible without+-- | Closes the connection with postgres as quickly as possible, without -- the proper postgres protocol handshake procedures. This is equivalent to -- closing the connection's socket in the kernel without making postgres -- aware of it. -- Use this if you need to close the connection in exception handlers or -- if you received an irrecoverable Hpgsql exception. closeForcefully :: HPgConnection -> IO ()-closeForcefully conn@(HPgConnection {socket}) = whenNotClosed conn $ do-  Socket.close socket-  modifyMVar_ conn.socketClosed $ const $ pure True+closeForcefully conn = closeInternalFinally conn $ pure () -whenNotClosed :: HPgConnection -> IO () -> IO ()-whenNotClosed conn f = do+-- | Returns `True` if the connection has already been closed.+connectionIsClosed :: HPgConnection -> IO Bool+connectionIsClosed conn = readMVar conn.socketClosed++closeInternalFinally :: HPgConnection -> IO () -> IO ()+closeInternalFinally conn f = do   isClosed <- readMVar conn.socketClosed-  unless isClosed f+  unless isClosed $ finally f (Socket.close conn.socket >> modifyMVar_ conn.socketClosed (const $ pure True))  withConnection :: ConnectionString -> DiffTime -> (HPgConnection -> IO a) -> IO a withConnection connstr conntimeout f = bracketOnError (connect connstr conntimeout) closeForcefully $ \conn -> do@@ -597,6 +645,10 @@   IO c withControlMsgsLock conn@HPgConnection {socket, socketMutex} acqStm relStm f = do   withMutex socketMutex $ do+    -- Check the connection status here as this function is in almost+    -- every code path+    connIsClosed <- connectionIsClosed conn+    when connIsClosed $ throwIrrecoverableError "Attempting to use a database connection that has been closed"     -- The lock is acquired, so now we run flushSendBuffer     -- so the caller will only see internal state after previous     -- messages were sent@@ -1644,8 +1696,11 @@  -- | Wraps an IO action to rethrow any exception as a IrrecoverableHpgsqlError. rethrowAsIrrecoverable :: IO a -> IO a-rethrowAsIrrecoverable = handle (throw . asIrrec)+rethrowAsIrrecoverable = handleJust isNotIrrec (throw . asIrrec)   where+    isNotIrrec (ex :: SomeException) = case fromException ex of+      Just (_ :: IrrecoverableHpgsqlError) -> Nothing -- Let irrecoverable errors blow up+      Nothing -> Just ex -- Catch others to rethrow wrapped     asIrrec (ex :: SomeException) = IrrecoverableHpgsqlError {hpgsqlDetails = "An inner exception was thrown", innerException = Just ex, relatedStatement = Nothing}  {-# NOINLINE _globalDebugLock #-}

@@ -301,13 +301,24 @@   deriving stock (Eq, Show)  data ResetConnectionOpts = ResetConnectionOpts-  { -- | Runs `RESET ALL` and `RESET ROLE` on the connection. Defaults to True.+  { -- | Runs `CLOSE ALL` to close all cursors. Defaults to True.+    closeAllCursors :: Bool,+    -- | Runs `RESET ALL` on the connection. Defaults to True.     resetAll :: Bool,+    -- | Runs `SET SESSION AUTHORIZATION DEFAULT` on the connection. Defaults to True.+    resetRole :: Bool,+    -- | Runs `DEALLOCATE ALL; DISCARD PLANS;` to deallocate all prepared statements and cached query plans. Defaults to False to reap more benefits from prepared statements.+    deallocatePreparedStmts :: Bool,     -- | Runs `UNLISTEN *` on the connection and clears the internal queue of notifications. Defaults to True.     unlistenAll :: Bool,+    -- | Runs `SELECT pg_advisory_unlock_all()` to release all session-level advisory locks. Defaults to True.+    releaseAdvisoryLocks :: Bool,+    -- | Runs `DISCARD TEMP` to drop temporary tables created in this session. Defaults to True.+    discardTempTables :: Bool,+    -- | Runs `DISCARD SEQUENCES` to discard all cached sequence-related state. Defaults to True.+    discardSequences :: Bool,     -- | Throws an exception if there is an open transaction or if there's a transaction in error state. Defaults to True.     checkTransactionState :: Bool-    -- TODO: Check for any temporary tables and throw?   }  -- | An error coming from PostgreSQL. You can safely handle this and continue using the connection.

@@ -1,12 +1,13 @@-module Hpgsql.Msgs (AuthenticationResponse (..), AuthenticationMethod (..), BackendKeyData (..), Bind (..), BindComplete (..), CancelRequest (..), CommandComplete (..), CopyData (..), CopyDone (..), CopyFail (..), CopyInResponse (..), DataRow (..), Describe (..), ErrorDetail (..), ErrorResponse (..), Execute (..), Flush (..), NoData (..), ParameterStatus (..), Query (..), ReadyForQuery (..), RowDescription (..), StartupMessage (..), ToPgMessage (..), FromPgMessage (..), PgMsgParser (..), Terminate (..), TransactionStatus (..), NoticeResponse (..), NotificationResponse (..), Parse (..), ParseComplete (..), PasswordMessage (..), Sync (..), parsePgMessage, nulTermCString) where+module Hpgsql.Msgs (AuthenticationResponse (..), AuthenticationMethod (..), BackendKeyData (..), Bind (..), BindComplete (..), CancelRequest (..), CommandComplete (..), CopyData (..), CopyDone (..), CopyFail (..), CopyInResponse (..), DataRow (..), Describe (..), ErrorDetail (..), ErrorResponse (..), Execute (..), Flush (..), NoData (..), ParameterStatus (..), Query (..), ReadyForQuery (..), RowDescription (..), SASLInitialResponse (..), SASLResponse (..), StartupMessage (..), ToPgMessage (..), FromPgMessage (..), PgMsgParser (..), Terminate (..), TransactionStatus (..), NoticeResponse (..), NotificationResponse (..), Parse (..), ParseComplete (..), PasswordMessage (..), Sync (..), parsePgMessage, nulTermCString) where  import Control.Applicative (Alternative (..)) import Control.Monad (replicateM)-import qualified Crypto.Hash.MD5 as MD5+import qualified Crypto.Hash as Crypto import qualified Data.Attoparsec.ByteString as Parsec import qualified Data.Attoparsec.ByteString.Lazy as LazyParsec import qualified Data.Attoparsec.Text as TextParsec import Data.Bifunctor (first)+import Data.ByteArray (convert) import Data.ByteString (ByteString) import qualified Data.ByteString as BS import Data.ByteString.Internal (w2c)@@ -18,11 +19,12 @@ import Data.Maybe (fromMaybe, mapMaybe) import qualified Data.Serialize as Cereal import Data.Text (Text)-import Data.Text.Encoding (decodeASCII, decodeUtf8)+import Data.Text.Encoding (decodeASCII, decodeUtf8, encodeUtf8) import Data.Word (Word8) import Hpgsql.Builder (BinaryField, Builder, builderLength) import qualified Hpgsql.Builder as Builder import Hpgsql.InternalTypes (BindComplete (..), CommandComplete (..), CopyInResponse (..), DataRow (..), ErrorDetail (..), ErrorResponse (..), NoData (..), NotificationResponse (..), ParseComplete (..), ReadyForQuery (..), RowDescription (..), TransactionStatus (..))+import Hpgsql.ScramSHA256 (ScramClientFinalMessage (..), ScramServerFirstMessage (..)) import Hpgsql.TypeInfo (Oid (..))  class ToPgMessage a where@@ -71,9 +73,15 @@ newtype AuthenticationResponse = AuthenticationResponse AuthenticationMethod   deriving stock (Show) -data AuthenticationMethod = AuthOk | AuthKerberosV5 | AuthCleartextPassword | AuthMD5Password LBS.ByteString | AuthGSS | AuthGSSContinue | AuthSSPI | AuthSASL | AuthSASLContinue | AuthSASLFinal+data AuthenticationMethod = AuthOk | AuthKerberosV5 | AuthCleartextPassword | AuthMD5Password LBS.ByteString | AuthGSS | AuthGSSContinue | AuthSSPI | AuthSASL [Text] | AuthSASLContinue !ScramServerFirstMessage | AuthSASLFinal ByteString   deriving stock (Show) +data SASLInitialResponse = SASLInitialResponse Text ByteString+  deriving stock (Show)++newtype SASLResponse = SASLResponse ScramClientFinalMessage+  deriving stock (Show)+ data PasswordMessage   = -- | Two last are username and password     PasswordMessage AuthenticationMethod String String@@ -136,19 +144,19 @@  instance FromPgMessage AuthenticationResponse where   msgParser = PgMsgParser $ \c restOfMsg -> case c of-    'R' -> case LBS.length restOfMsg of-      4 ->-        case Cereal.decodeLazy @Int32 restOfMsg of-          Right 0 -> Just $ AuthenticationResponse AuthOk-          Right 2 -> Just $ AuthenticationResponse AuthKerberosV5-          Right 3 -> Just $ AuthenticationResponse AuthCleartextPassword-          Right 7 -> Just $ AuthenticationResponse AuthGSS-          Right 9 -> Just $ AuthenticationResponse AuthSSPI-          _ -> Nothing-      8 -> case first (Cereal.decodeLazy @Int32) $ LBS.splitAt 4 restOfMsg of-        (Right 5, salt) -> Just $ AuthenticationResponse (AuthMD5Password salt)-        _ -> Nothing -- We don't care too much about parsing every possible response yet-      _ -> Nothing -- We don't care too much about parsing every possible response yet+    'R' -> case first (Cereal.decodeLazy @Int32) $ LBS.splitAt 4 restOfMsg of+      (Right 0, _) -> Just $ AuthenticationResponse AuthOk+      (Right 2, _) -> Just $ AuthenticationResponse AuthKerberosV5+      (Right 3, _) -> Just $ AuthenticationResponse AuthCleartextPassword+      (Right 5, salt) -> Just $ AuthenticationResponse (AuthMD5Password salt)+      (Right 7, _) -> Just $ AuthenticationResponse AuthGSS+      (Right 9, _) -> Just $ AuthenticationResponse AuthSSPI+      (Right 10, saslMechanismName) -> case Parsec.parseOnly (Parsec.many1 nulTerminatedCStringParser) (LBS.toStrict saslMechanismName) of+        Right sm -> Just $ AuthenticationResponse (AuthSASL sm)+        Left _ -> Nothing+      (Right 11, saslData) -> Just $ AuthenticationResponse (AuthSASLContinue (ScramServerFirstMessage $ LBS.toStrict saslData))+      (Right 12, saslData) -> Just $ AuthenticationResponse (AuthSASLFinal (LBS.toStrict saslData))+      _ -> Nothing     _ -> Nothing  instance FromPgMessage BackendKeyData where@@ -250,12 +258,29 @@           AuthMD5Password (LBS.toStrict -> salt) ->             let passwordBytes = Builder.toStrictByteString (Builder.string7 password)                 usernameBytes = Builder.toStrictByteString (Builder.string7 username)-                innerHex = bytestringToHex (MD5.hash (passwordBytes <> usernameBytes))-                outerHex = bytestringToHex (MD5.hash (innerHex <> salt))+                innerHex = bytestringToHex (md5Hash (passwordBytes <> usernameBytes))+                outerHex = bytestringToHex (md5Hash (innerHex <> salt))              in Builder.byteString "md5" <> Builder.byteString outerHex <> Builder.word8 0           _ -> error "PasswordMessage method not implemented"         msgLen = builderLength passwordBs + 4      in Builder.char7 'p' <> Builder.int32BE msgLen <> passwordBs+    where+      md5Hash :: ByteString -> ByteString+      md5Hash = convert @(Crypto.Digest Crypto.MD5) . Crypto.hash++instance ToPgMessage SASLInitialResponse where+  toPgMessage (SASLInitialResponse mechanism clientFirstMsg) =+    let mechanismBs = Builder.byteString (encodeUtf8 mechanism) <> Builder.word8 0+        dataLen = fromIntegral (BS.length clientFirstMsg) :: Int32+        contents = mechanismBs <> Builder.int32BE dataLen <> Builder.byteString clientFirstMsg+        msgLen = builderLength contents + 4+     in Builder.char7 'p' <> Builder.int32BE msgLen <> contents++instance ToPgMessage SASLResponse where+  toPgMessage (SASLResponse responseData) =+    let contents = Builder.byteString responseData.clientFinalMessage+        msgLen = builderLength contents + 4+     in Builder.char7 'p' <> Builder.int32BE msgLen <> contents  instance ToPgMessage Query where   toPgMessage (Query bs) =

@@ -39,7 +39,9 @@     pipelineExec_,     pipeline1,     pipeline1With,+    pipelineMay,+    pipelineMayWith,   ) where -import Hpgsql.Internal (Pipeline, pipeline, pipeline1, pipeline1With, pipelineExec, pipelineExec_, pipelineS, pipelineSMWith, pipelineSWith, pipelineWith, runPipeline)+import Hpgsql.Internal (Pipeline, pipeline, pipeline1, pipeline1With, pipelineExec, pipelineExec_, pipelineMay, pipelineMayWith, pipelineS, pipelineSMWith, pipelineSWith, pipelineWith, runPipeline)

@@ -0,0 +1,192 @@+module Hpgsql.ScramSHA256+  ( ScramClientFirstMessage (..),+    ScramServerFirstMessage (..),+    ScramClientFinalMessage (..),+    generateClientFirstMessage,+    handleServerFirstMsg,+    verifyServerFinal,+  )+where++-- https://datatracker.ietf.org/doc/html/rfc5802#section-3+-- The RFC for SCRAM-SHA-1 (which SCRAM-SHA-256 follows except it uses SHA256)+-- says this:+--+-- "Informative Note: Implementors are encouraged to create test cases+--    that use both usernames and passwords with non-ASCII codepoints.  In+--    particular, it's useful to test codepoints whose "Unicode+--    Normalization Form C" and "Unicode Normalization Form KC" are+--    different.  Some examples of such codepoints include Vulgar Fraction+--    One Half (U+00BD) and Acute Accent (U+00B4)."+--+-- TODO: We should test these scenarios, but we don't yet.++import Control.Applicative ((<|>))+import Control.Monad (void)+import Crypto.Hash (Digest, SHA256, hash)+import Crypto.KDF.PBKDF2 (Parameters (..), fastPBKDF2_SHA256)+import Crypto.MAC.HMAC (HMAC, hmac, hmacGetDigest)+import Crypto.Random (getRandomBytes)+import Data.Attoparsec.ByteString (parseOnly, string)+import qualified Data.Attoparsec.ByteString.Char8 as Parsec+import Data.ByteArray (convert, xor)+import Data.ByteArray.Encoding (Base (Base64), convertFromBase, convertToBase)+import Data.ByteString (ByteString)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as BS8+import Data.Text (Text)+import Data.Text.Encoding (encodeUtf8)++-- From the RFC:+-- SCRAM is a SASL mechanism whose client response and server challenge+--    messages are text-based messages containing one or more attribute-+--    value pairs separated by commas.  Each attribute has a one-letter+--    name.  The messages and their attributes are described in+--    Section 5.1, and defined in Section 7.++--    SCRAM is a client-first SASL mechanism (see [RFC4422], Section 5,+--    item 2a), and returns additional data together with a server's+--    indication of a successful outcome.++--    This is a simple example of a SCRAM-SHA-1 authentication exchange+--    when the client doesn't support channel bindings (username 'user' and+--    password 'pencil' are used):++--    C: n,,n=user,r=fyko+d2lbbFgONRv9qkxdawL+--    S: r=fyko+d2lbbFgONRv9qkxdawL3rfcNHYJY1ZVvWVs7j,s=QSXCR+Q6sek8bf92,+--       i=4096+--    C: c=biws,r=fyko+d2lbbFgONRv9qkxdawL3rfcNHYJY1ZVvWVs7j,+--       p=v0X8v3Bz2T0CJGbJQyF0X+HI4Ts=+--    S: v=rmF9pqV8S7suAoZWja4dJRkFsKQ=+--+--+-- And from Postgres (https://www.postgresql.org/docs/current/sasl-authentication.html):+-- 1. The server sends an AuthenticationSASL message. It includes a list of SASL authentication mechanisms that the server can accept. This will be SCRAM-SHA-256-PLUS and SCRAM-SHA-256 if the server is built with SSL support, or else just the latter.+-- 2. The client responds by sending a SASLInitialResponse message, which indicates the chosen mechanism, SCRAM-SHA-256 or SCRAM-SHA-256-PLUS. (A client is free to choose either mechanism, but for better security it should choose the channel-binding variant if it can support it.) In the Initial Client response field, the message contains the SCRAM client-first-message. The client-first-message also contains the channel binding type chosen by the client.+--      TODO: hpgsql right now only supports SCRAM-SHA-256, not the -PLUS variant.+-- 3. Server sends an AuthenticationSASLContinue message, with a SCRAM server-first-message as the content.+-- 4. Client sends a SASLResponse message, with SCRAM client-final-message as the content.+-- 5. Server sends an AuthenticationSASLFinal message, with the SCRAM server-final-message, followed immediately by an AuthenticationOk message.++-- | The "client-first-message", something like `n,,n=user,r=nonce`, where+-- the nonce is something like `fyko+d2lbbFgONRv9qkxdawL`.+data ScramClientFirstMessage = ScramClientFirstMessage+  { -- | This is the string "n,,"+    gs2Header :: !ByteString,+    -- | This is the string "n=user,r=nonce", except that "user" is the+    -- empty string because postgres ignores it anyway.+    clientFirstBare :: !ByteString,+    -- | This is just the concatenation of the two strings above.+    fullMessage :: !ByteString+  }++-- | Generate the client-first-message for SCRAM-SHA-256.+generateClientFirstMessage :: IO ScramClientFirstMessage+generateClientFirstMessage = do+  -- We don't send a username in this message even if it's necessary according+  -- to the RFC, because postgres says+  -- "When SCRAM-SHA-256 is used in PostgreSQL, the server will ignore the user name that the client sends in the client-first-message. The user name that was already sent in the startup message is used instead."+  nonce <- generateNonce+  let gs2Header = "n,,"+      clientFirstBare = "n=,r=" <> nonce+      fullMessage = gs2Header <> clientFirstBare+  pure ScramClientFirstMessage {..}+  where+    generateNonce :: IO ByteString+    generateNonce = convertToBase Base64 <$> (getRandomBytes 24 :: IO ByteString)++newtype ScramServerFirstMessage = ScramServerFirstMessage+  { -- | As per the RFC, this looks something like (notice the concatenation+    -- of client and server nonces, and some salt):+    --    S: r=ClientNonceServerNonce,s=QSXCR+Q6sek8bf92,+    --       i=4096+    fullMessage :: ByteString+  }+  deriving stock (Show)++data ScramClientFinalMessage = ScramClientFinalMessage+  { -- | This looks something like:+    -- c=biws,r=fyko+d2lbbFgONRv9qkxdawL3rfcNHYJY1ZVvWVs7j,+    --       p=v0X8v3Bz2T0CJGbJQyF0X+HI4Ts=+    clientFinalMessage :: !ByteString,+    serverSignature :: !ByteString+  }+  deriving stock (Show)++-- | From the RFC:+-- In response, the server sends a "server-first-message" containing the+-- user's iteration count i and the user's salt, and appends its own+-- nonce to the client-specified one.++-- The client then responds by sending a "client-final-message" with the+-- same nonce and a ClientProof computed using the selected hash+-- function as explained earlier.+handleServerFirstMsg ::+  Text ->+  ScramClientFirstMessage ->+  ScramServerFirstMessage ->+  Either String ScramClientFinalMessage+handleServerFirstMsg (encodeUtf8 -> password) ScramClientFirstMessage {gs2Header, clientFirstBare} serverFirst = do+  (serverNonce, salt, iterations) <- parseServerFirst serverFirst.fullMessage+  let saltedPassword = hi salt iterations -- "hi" takes the password implicitly+      clientKey = hmacSHA256 saltedPassword "Client Key"+      storedKey = hashSHA256 clientKey+      channelBinding = "c=" <> convertToBase Base64 gs2Header+      clientFinalWithoutProof = channelBinding <> ",r=" <> serverNonce+      authMessage = clientFirstBare <> "," <> serverFirst.fullMessage <> "," <> clientFinalWithoutProof+      clientSignature = hmacSHA256 storedKey authMessage+      clientProof = Data.ByteArray.xor clientKey clientSignature :: ByteString+      clientFinalMessage = clientFinalWithoutProof <> ",p=" <> convertToBase Base64 clientProof+      serverKey = hmacSHA256 saltedPassword "Server Key"+      serverSignature = convertToBase Base64 (hmacSHA256 serverKey authMessage) :: ByteString+  pure ScramClientFinalMessage {clientFinalMessage, serverSignature}+  where+    parseServerFirst :: ByteString -> Either String (ByteString, ByteString, Int)+    parseServerFirst = parseOnly $ do+      -- As per the RFC, this looks something like (notice the concatenation+      -- of client and server nonces, and some salt):+      --    S: r=ClientNonceServerNonce,s=QSXCR+Q6sek8bf92,+      --       i=4096+      void $ string "r=" <|> fail "Missing nonce in server-first-message (1)"+      nonce <- Parsec.takeWhile1 (/= ',') <|> fail "Missing nonce in server-first-message (2)"+      void $ string ","+      void $ string "s=" <|> fail "Missing salt in server-first-message (1)"+      saltB64 <- Parsec.takeWhile1 (/= ',') <|> fail "Missing salt in server-first-message (2)"+      void $ string ","+      salt <- case convertFromBase Base64 saltB64 of+        Left e -> fail $ "Error converting salt from base64: " ++ e+        Right s -> pure s+      void $ string "i=" <|> fail "Missing iteration count in server-first-message"+      iters <- Parsec.decimal <|> fail "Iteration count not a number in server-first-message"+      Parsec.endOfInput+      pure (nonce, salt, iters)+    hi :: ByteString -> Int -> ByteString+    hi salt iterations =+      fastPBKDF2_SHA256+        Parameters {iterCounts = iterations, outputLength = 32}+        password+        salt++-- | Verify the server-final-message matches the expected server signature.+verifyServerFinal :: ScramClientFinalMessage -> ByteString -> Either String ()+verifyServerFinal finalMsg serverFinal = do+  actualSig <- parseServerFinal serverFinal+  if actualSig == finalMsg.serverSignature+    then Right ()+    else Left "Server signature mismatch"+  where+    parseServerFinal :: ByteString -> Either String ByteString+    parseServerFinal msg+      -- This looks like:+      --   v=rmF9pqV8S7suAoZWja4dJRkFsKQ=+      -- And v is a base64-encoded ServerSignature, or it looks like+      --   e=some-error+      | BS.isPrefixOf "v=" msg = Right (BS.drop 2 msg)+      | BS.isPrefixOf "e=" msg = Left $ "Server error: " <> BS8.unpack (BS.drop 2 msg)+      | otherwise = Left $ "Invalid server-final-message: " <> BS8.unpack msg++hmacSHA256 :: ByteString -> ByteString -> ByteString+hmacSHA256 key msg = convert $ hmacGetDigest (hmac key msg :: HMAC SHA256)++hashSHA256 :: ByteString -> ByteString+hashSHA256 bs = convert (hash bs :: Digest SHA256)

@@ -1,7 +1,7 @@ module Hpgsql.Transaction (withTransaction, withTransactionMode, begin, beginMode, commit, rollback, transactionStatus, IsolationLevel (..), ReadWriteMode (..), TransactionStatus (..)) where  import qualified Control.Concurrent.STM as STM-import Control.Exception.Safe (Exception (..), bracketWithError, throw, tryAny)+import Control.Exception.Safe (Exception (..), bracketWithError, throw, tryJust) import Control.Monad (unless) import Hpgsql.Internal (execute_, fullTransactionStatus, transactionStatus) import Hpgsql.InternalTypes (HPgConnection (..), InternalConnectionState (..), IrrecoverableHpgsqlError)@@ -73,24 +73,22 @@ -- will be issued before your next query on this connection. withTransactionMode :: HPgConnection -> IsolationLevel -> ReadWriteMode -> IO a -> IO a withTransactionMode conn il rw f = bracketWithError (beginMode conn il rw) cleanup $ \() -> do-  res <- tryAny f+  res <- tryJust notIrrecoverableError f   case res of-    Left ex -> case fromException ex of-      Just (_ :: IrrecoverableHpgsqlError) -> throw ex -- Rethrow internal errors-      Nothing -> do-        -- In case of a synchronous exception, we rollback synchronously.-        -- If this is interrupted:-        -- - Before ROLLBACK is sent, `cleanup` will enqueue a "ROLLBACK".-        -- - After ROLLBACK is sent but before it finishes, a ROLLBACK will still be enqueued-        --   (check `cleanup`), which means:-        --      - If ROLLBACK has already completed by the time the new ROLLBACK is meant to be sent,-        --        the new ROLLBACK will produce a "WARNING: there is no transaction in progress". This-        --        isn't great, but is mostly harmless.-        --      - If ROLLBACK is cancelled by the future ROLLBACK, all is good as well.-        -- - After ROLLBACK is sent and completed, `cleanup` won't enqueue a new ROLLBACK, and all is well.-        ---        -- We rollback here, not in `cleanup`, because that runs with async exceptions masked-        rollback conn >> throw ex -- Reminder that this can be interrupted before/after "rollback" is sent+    Left ex -> do+      -- In case of a synchronous exception, we rollback synchronously.+      -- If this is interrupted:+      -- - Before ROLLBACK is sent, `cleanup` will enqueue a "ROLLBACK".+      -- - After ROLLBACK is sent but before it finishes, a ROLLBACK will still be enqueued+      --   (check `cleanup`), which means:+      --      - If ROLLBACK has already completed by the time the new ROLLBACK is meant to be sent,+      --        the new ROLLBACK will produce a "WARNING: there is no transaction in progress". This+      --        isn't great, but is mostly harmless.+      --      - If ROLLBACK is cancelled by the future ROLLBACK, all is good as well.+      -- - After ROLLBACK is sent and completed, `cleanup` won't enqueue a new ROLLBACK, and all is well.+      --+      -- We rollback here, not in `cleanup`, because that runs with async exceptions masked+      rollback conn >> throw ex -- Reminder that this can be interrupted before/after "rollback" is sent     Right v -> do       -- If this is interrupted:       -- - Before COMMIT is sent, it's equivalent to failing in the middle of the user-supplied action@@ -105,6 +103,9 @@       commit conn       pure v   where+    notIrrecoverableError e = case fromException e of+      Just (_ :: IrrecoverableHpgsqlError) -> Nothing+      _ -> Just e     cleanup mEx () = case mEx of       Nothing -> pure ()       Just (fromException -> (Just (_ :: IrrecoverableHpgsqlError))) -> pure () -- Do nothing if an internal error was thrown, just let it propagate