{-# LANGUAGE CPP #-}

-- | 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 qualified Data.Sequence as Seq
import qualified Data.Set as Set
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
import System.Environment (lookupEnv)
#if defined(mingw32_HOST_OS)
import System.Win32.Info.Computer (getUserName)
#else
import System.Posix.User (getEffectiveUserName)
#endif

-- | Parsed connection parameters (the @key=value@ subset we support).
data ConnInfo = ConnInfo
  { ConnInfo -> ByteString
host :: ByteString,
    ConnInfo -> Int
port :: Int,
    ConnInfo -> ByteString
user :: ByteString,
    ConnInfo -> ByteString
database :: ByteString,
    ConnInfo -> ByteString
password :: ByteString,
    -- | Every other recognized @key=value@ pair (e.g. @application_name@,
    -- @options@), forwarded verbatim in the startup message so the server
    -- sees them, the way libpq does.
    ConnInfo -> Map ByteString ByteString
extraParams :: Map.Map ByteString ByteString
  }
  deriving stock (ConnInfo -> ConnInfo -> Bool
(ConnInfo -> ConnInfo -> Bool)
-> (ConnInfo -> ConnInfo -> Bool) -> Eq ConnInfo
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ConnInfo -> ConnInfo -> Bool
== :: ConnInfo -> ConnInfo -> Bool
$c/= :: ConnInfo -> ConnInfo -> Bool
/= :: ConnInfo -> ConnInfo -> Bool
Eq, Int -> ConnInfo -> ShowS
[ConnInfo] -> ShowS
ConnInfo -> String
(Int -> ConnInfo -> ShowS)
-> (ConnInfo -> String) -> ([ConnInfo] -> ShowS) -> Show ConnInfo
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ConnInfo -> ShowS
showsPrec :: Int -> ConnInfo -> ShowS
$cshow :: ConnInfo -> String
show :: ConnInfo -> String
$cshowList :: [ConnInfo] -> ShowS
showList :: [ConnInfo] -> ShowS
Show)

-- | Parse a conninfo string in either @key=value@ or @postgresql:\/\/@ URI
-- format. Unquoted key=value values only; URI values are percent-decoded.
-- @.pgpass@ is not supported.
--
-- The @dfltUser@ argument is the already-resolved default user (see
-- 'resolveDefaultUser'), used whenever the conninfo omits @user@. It is passed
-- in explicitly because resolving it is the one IO effect this otherwise-pure
-- parser needs.
parseConnInfo :: ByteString -> ByteString -> ConnInfo
parseConnInfo :: ByteString -> ByteString -> ConnInfo
parseConnInfo ByteString
raw ByteString
dfltUser =
  if
    | ByteString
"postgresql://" ByteString -> ByteString -> Bool
`ByteString.isPrefixOf` ByteString
raw -> ByteString -> ByteString -> ConnInfo
parseUri ByteString
dfltUser (Int -> ByteString -> ByteString
ByteString.drop Int
13 ByteString
raw)
    | ByteString
"postgres://" ByteString -> ByteString -> Bool
`ByteString.isPrefixOf` ByteString
raw -> ByteString -> ByteString -> ConnInfo
parseUri ByteString
dfltUser (Int -> ByteString -> ByteString
ByteString.drop Int
11 ByteString
raw)
    | Bool
otherwise -> ByteString -> ByteString -> ConnInfo
parseKeyValue ByteString
dfltUser ByteString
raw

-- | Resolve the default @user@ the way libpq does (@conninfo_add_defaults@ /
-- @pg_fe_getauthname@ in @fe-connect.c@): the @PGUSER@ environment variable if
-- set and non-empty, otherwise the operating-system login name
-- (@getpwuid(geteuid())->pw_name@ on Unix, @GetUserName@ on Windows).
--
-- Returns @Left msg@ if neither is available. Mirroring libpq, a lookup failure
-- must be surfaced by the caller as a 'ConnectionBad' connection (see
-- 'establish') rather than attempting to connect.
resolveDefaultUser :: IO (Either String ByteString)
resolveDefaultUser :: IO (Either String ByteString)
resolveDefaultUser = do
  Maybe String
pguser <- String -> IO (Maybe String)
lookupEnv String
"PGUSER"
  case Maybe String
pguser of
    Just String
u | Bool -> Bool
not (String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
u) -> Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either String ByteString
forall a b. b -> Either a b
Right (String -> ByteString
ByteString.Char8.pack String
u))
    Maybe String
_ -> do
      Either SomeException ByteString
result <- forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException (String -> ByteString
ByteString.Char8.pack (String -> ByteString) -> IO String -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO String
platformUserName)
      Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure case Either SomeException ByteString
result of
        Right ByteString
name -> ByteString -> Either String ByteString
forall a b. b -> Either a b
Right ByteString
name
        Left SomeException
err -> String -> Either String ByteString
forall a b. a -> Either a b
Left (String
platformUserNameLookupFailureMessage String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
": " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> SomeException -> String
forall a. Show a => a -> String
show SomeException
err)

-- | The operating-system login name, via the same call libpq uses:
-- 'getEffectiveUserName' (@getpwuid(geteuid())@) on Unix, 'getUserName'
-- (@GetUserName@) on Windows.
platformUserName :: IO String
#if defined(mingw32_HOST_OS)
platformUserName = getUserName
#else
platformUserName :: IO String
platformUserName = IO String
getEffectiveUserName
#endif

#if defined(mingw32_HOST_OS)
platformUserNameLookupFailureMessage :: String
platformUserNameLookupFailureMessage = "user name lookup failure"
#else
platformUserNameLookupFailureMessage :: String
platformUserNameLookupFailureMessage :: String
platformUserNameLookupFailureMessage = String
"could not look up local user name"
#endif

parseKeyValue :: ByteString -> ByteString -> ConnInfo
parseKeyValue :: ByteString -> ByteString -> ConnInfo
parseKeyValue ByteString
dfltUser ByteString
raw =
  ConnInfo
    { host :: ByteString
host = ByteString -> ByteString -> ByteString
get ByteString
"host" ByteString
"localhost",
      port :: Int
port = Int -> ((Int, ByteString) -> Int) -> Maybe (Int, ByteString) -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
5432 (Int, ByteString) -> Int
forall a b. (a, b) -> a
fst (ByteString -> Maybe (Int, ByteString)
ByteString.Char8.readInt (ByteString -> ByteString -> ByteString
get ByteString
"port" ByteString
"5432")),
      user :: ByteString
user = ByteString
theUser,
      database :: ByteString
database = ByteString -> ByteString -> ByteString
get ByteString
"dbname" ByteString
theUser,
      password :: ByteString
password = ByteString -> ByteString -> ByteString
get ByteString
"password" ByteString
"",
      extraParams :: Map ByteString ByteString
extraParams = Map ByteString ByteString
-> Set ByteString -> Map ByteString ByteString
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys Map ByteString ByteString
settings Set ByteString
reservedKeys
    }
  where
    pairs :: [(ByteString, ByteString)]
pairs = (ByteString -> Maybe (ByteString, ByteString))
-> [ByteString] -> [(ByteString, ByteString)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ByteString -> Maybe (ByteString, ByteString)
toPair (ByteString -> [ByteString]
ByteString.Char8.words ByteString
raw)
    settings :: Map ByteString ByteString
settings = [(ByteString, ByteString)] -> Map ByteString ByteString
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(ByteString, ByteString)]
pairs
    get :: ByteString -> ByteString -> ByteString
get ByteString
key ByteString
def = ByteString -> ByteString -> Map ByteString ByteString -> ByteString
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault ByteString
def ByteString
key Map ByteString ByteString
settings
    theUser :: ByteString
theUser = ByteString -> ByteString -> ByteString
get ByteString
"user" ByteString
dfltUser
    toPair :: ByteString -> Maybe (ByteString, ByteString)
toPair ByteString
token = case (Char -> Bool) -> ByteString -> (ByteString, ByteString)
ByteString.Char8.break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'=') ByteString
token of
      (ByteString
key, ByteString
value)
        | Bool -> Bool
not (ByteString -> Bool
ByteString.null ByteString
value) -> (ByteString, ByteString) -> Maybe (ByteString, ByteString)
forall a. a -> Maybe a
Just (ByteString
key, Int -> ByteString -> ByteString
ByteString.drop Int
1 ByteString
value)
      (ByteString, ByteString)
_ -> Maybe (ByteString, ByteString)
forall a. Maybe a
Nothing

-- | Conninfo keys already surfaced via their own 'ConnInfo' fields, so they're
-- excluded from 'extraParams' rather than duplicated there.
reservedKeys :: Set.Set ByteString
reservedKeys :: Set ByteString
reservedKeys = [ByteString] -> Set ByteString
forall a. Ord a => [a] -> Set a
Set.fromList [ByteString
"host", ByteString
"port", ByteString
"user", ByteString
"dbname", ByteString
"password"]

-- | 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 -> ByteString -> ConnInfo
parseUri :: ByteString -> ByteString -> ConnInfo
parseUri ByteString
dfltUser ByteString
withoutScheme =
  ConnInfo {ByteString
host :: ByteString
host :: ByteString
host, Int
port :: Int
port :: Int
port, ByteString
user :: ByteString
user :: ByteString
user, ByteString
database :: ByteString
database :: ByteString
database, ByteString
password :: ByteString
password :: ByteString
password, Map ByteString ByteString
extraParams :: Map ByteString ByteString
extraParams :: Map ByteString ByteString
extraParams}
  where
    -- Split off optional "userinfo@" prefix. The '@' is unambiguous in this
    -- position: hosts do not contain '@' in practice.
    (Maybe ByteString
userinfoMay, ByteString
afterAt) = case Word8 -> ByteString -> Maybe Int
ByteString.elemIndex Word8
0x40 ByteString
withoutScheme of
      Just Int
i -> (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (Int -> ByteString -> ByteString
ByteString.take Int
i ByteString
withoutScheme), Int -> ByteString -> ByteString
ByteString.drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ByteString
withoutScheme)
      Maybe Int
Nothing -> (Maybe ByteString
forall a. Maybe a
Nothing, ByteString
withoutScheme)

    (ByteString
user, ByteString
password) = case Maybe ByteString
userinfoMay of
      Maybe ByteString
Nothing -> (ByteString
dfltUser, ByteString
"")
      Just ByteString
ui -> case Word8 -> ByteString -> Maybe Int
ByteString.elemIndex Word8
0x3a ByteString
ui of
        Just Int
c -> (ByteString -> ByteString
pctDecode (Int -> ByteString -> ByteString
ByteString.take Int
c ByteString
ui), ByteString -> ByteString
pctDecode (Int -> ByteString -> ByteString
ByteString.drop (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ByteString
ui))
        Maybe Int
Nothing -> (ByteString -> ByteString
pctDecode ByteString
ui, ByteString
"")

    -- Split host[:port] from /dbname?params.
    (ByteString
hostport, ByteString
pathAndQuery) = (Word8 -> Bool) -> ByteString -> (ByteString, ByteString)
ByteString.break (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x2f) ByteString
afterAt

    -- Strip optional leading '/' to get just the dbname (ignoring ?query).
    rawDbname :: ByteString
rawDbname =
      let p :: ByteString
p = Int -> ByteString -> ByteString
ByteString.drop Int
1 ByteString
pathAndQuery
       in (Word8 -> Bool) -> ByteString -> ByteString
ByteString.takeWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0x3f) ByteString
p

    database :: ByteString
database =
      if ByteString -> Bool
ByteString.null ByteString
rawDbname
        then ByteString
user
        else ByteString -> ByteString
pctDecode ByteString
rawDbname

    -- Everything after the first '?', parsed as '&'-separated key=value pairs
    -- (e.g. ?application_name=foo&sslmode=disable).
    rawQuery :: ByteString
rawQuery = Int -> ByteString -> ByteString
ByteString.drop Int
1 ((Word8 -> Bool) -> ByteString -> ByteString
ByteString.dropWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0x3f) ByteString
pathAndQuery)

    extraParams :: Map ByteString ByteString
extraParams =
      Map ByteString ByteString
-> Set ByteString -> Map ByteString ByteString
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys
        ([(ByteString, ByteString)] -> Map ByteString ByteString
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((ByteString -> Maybe (ByteString, ByteString))
-> [ByteString] -> [(ByteString, ByteString)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ByteString -> Maybe (ByteString, ByteString)
toQueryPair (Word8 -> ByteString -> [ByteString]
ByteString.split Word8
0x26 ByteString
rawQuery)))
        Set ByteString
reservedKeys

    toQueryPair :: ByteString -> Maybe (ByteString, ByteString)
toQueryPair ByteString
token = case Word8 -> ByteString -> Maybe Int
ByteString.elemIndex Word8
0x3d ByteString
token of
      Just Int
i ->
        let (ByteString
k, ByteString
v) = Int -> ByteString -> (ByteString, ByteString)
ByteString.splitAt Int
i ByteString
token
         in if ByteString -> Bool
ByteString.null ByteString
k then Maybe (ByteString, ByteString)
forall a. Maybe a
Nothing else (ByteString, ByteString) -> Maybe (ByteString, ByteString)
forall a. a -> Maybe a
Just (ByteString -> ByteString
pctDecode ByteString
k, ByteString -> ByteString
pctDecode (Int -> ByteString -> ByteString
ByteString.drop Int
1 ByteString
v))
      Maybe Int
Nothing -> Maybe (ByteString, ByteString)
forall a. Maybe a
Nothing

    -- Parse host and port from "host:port", handling IPv6 "[::1]:port".
    (ByteString
host, Int
port)
      | Bool -> Bool
not (ByteString -> Bool
ByteString.null ByteString
hostport) Bool -> Bool -> Bool
&& HasCallStack => ByteString -> Word8
ByteString -> Word8
ByteString.head ByteString
hostport Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x5b =
          case Word8 -> ByteString -> Maybe Int
ByteString.elemIndex Word8
0x5d ByteString
hostport of
            Just Int
close ->
              let h :: ByteString
h = Int -> ByteString -> ByteString
ByteString.take Int
close (Int -> ByteString -> ByteString
ByteString.drop Int
1 ByteString
hostport)
                  portStr :: ByteString
portStr = Int -> ByteString -> ByteString
ByteString.drop Int
2 (Int -> ByteString -> ByteString
ByteString.drop Int
close ByteString
hostport)
               in (ByteString
h, ByteString -> Int
readPort ByteString
portStr)
            Maybe Int
Nothing -> (ByteString
hostport, Int
5432)
      | Bool
otherwise = case Word8 -> ByteString -> Maybe Int
ByteString.elemIndexEnd Word8
0x3a ByteString
hostport of
          Just Int
c ->
            let h :: ByteString
h = Int -> ByteString -> ByteString
ByteString.take Int
c ByteString
hostport
                p :: ByteString
p = Int -> ByteString -> ByteString
ByteString.drop (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ByteString
hostport
             in if ByteString -> Bool
ByteString.null ByteString
h then (ByteString
"localhost", ByteString -> Int
readPort ByteString
p) else (ByteString
h, ByteString -> Int
readPort ByteString
p)
          Maybe Int
Nothing ->
            (if ByteString -> Bool
ByteString.null ByteString
hostport then ByteString
"localhost" else ByteString
hostport, Int
5432)

    readPort :: ByteString -> Int
readPort ByteString
bs = Int -> ((Int, ByteString) -> Int) -> Maybe (Int, ByteString) -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
5432 (Int, ByteString) -> Int
forall a b. (a, b) -> a
fst (ByteString -> Maybe (Int, ByteString)
ByteString.Char8.readInt ByteString
bs)

-- | Decode @%XX@ percent-encoding in a URI component.
pctDecode :: ByteString -> ByteString
pctDecode :: ByteString -> ByteString
pctDecode ByteString
bs
  | ByteString -> Bool
ByteString.null ByteString
bs = ByteString
ByteString.empty
  | HasCallStack => ByteString -> Word8
ByteString -> Word8
ByteString.head ByteString
bs Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x25,
    ByteString -> Int
ByteString.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
3,
    Just Int
hi <- Word8 -> Maybe Int
forall {a}. Integral a => a -> Maybe Int
hexVal (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
ByteString.index ByteString
bs Int
1),
    Just Int
lo <- Word8 -> Maybe Int
forall {a}. Integral a => a -> Maybe Int
hexVal (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
ByteString.index ByteString
bs Int
2) =
      Word8 -> ByteString
ByteString.singleton (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
16 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lo)) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString
pctDecode (Int -> ByteString -> ByteString
ByteString.drop Int
3 ByteString
bs)
  | Bool
otherwise = Word8 -> ByteString
ByteString.singleton (HasCallStack => ByteString -> Word8
ByteString -> Word8
ByteString.head ByteString
bs) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString
pctDecode (HasCallStack => ByteString -> ByteString
ByteString -> ByteString
ByteString.tail ByteString
bs)
  where
    hexVal :: a -> Maybe Int
hexVal a
w
      | a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0x30 Bool -> Bool -> Bool
&& a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0x39 = Int -> Maybe Int
forall a. a -> Maybe a
Just (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
0x30 :: Int)
      | a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0x41 Bool -> Bool -> Bool
&& a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0x46 = Int -> Maybe Int
forall a. a -> Maybe a
Just (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
0x41 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
10 :: Int)
      | a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0x61 Bool -> Bool -> Bool
&& a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0x66 = Int -> Maybe Int
forall a. a -> Maybe a
Just (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
0x61 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
10 :: Int)
      | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing

-- | A native connection and its mutable state.
data Connection = Connection
  { Connection -> IORef Transport
transport :: IORef Transport,
    Connection -> ConnInfo
info :: ConnInfo,
    Connection -> Bool
isNull :: Bool,
    Connection -> IORef (Map ByteString ByteString)
parameters :: IORef (Map.Map ByteString ByteString),
    Connection -> IORef (Maybe (Int32, Int32))
backendKey :: IORef (Maybe (Int32, Int32)),
    Connection -> IORef Word8
txStatus :: IORef Word8,
    Connection -> IORef ConnStatus
connStatus :: IORef ConnStatus,
    Connection -> IORef (Maybe ByteString)
lastError :: IORef (Maybe ByteString),
    Connection -> IORef [ByteString]
notices :: IORef [ByteString],
    Connection -> IORef [Notify]
pendingNotifications :: IORef [Notify],
    Connection -> IORef Bool
noticeReporting :: IORef Bool,
    Connection -> IORef Bool
asyncPending :: IORef Bool,
    Connection -> IORef Bool
nonblocking :: IORef Bool,
    Connection -> IORef PipelineStatus
pipelineStatus :: IORef PipelineStatus,
    Connection -> IORef Bool
singleRowMode :: IORef Bool,
    Connection -> IORef [FieldDescription]
singleRowFields :: IORef [FieldDescription],
    Connection -> IORef Bool
pipelineSeparatorPending :: IORef Bool,
    Connection -> IORef Int
pendingSyncs :: IORef Int,
    Connection -> IORef Int
pendingCommands :: IORef Int,
    -- | FIFO, one entry per in-flight pipeline command that sends a @Parse@
    -- ('sendQueryParams' or 'sendPrepare'), recording whether that command's
    -- @ParseComplete@ is itself the terminal result. 'sendPrepare' pushes
    -- @True@ (its @ParseComplete@ ends the command as 'CommandOk');
    -- 'sendQueryParams' pushes @False@ (its @ParseComplete@ must fold into the
    -- command's accumulating result, like every other extended-protocol
    -- message, and never terminate it early). Popping the head at the right
    -- @ParseComplete@ keeps the origins in order even when several commands are
    -- pipelined together - a plain counter cannot, which is what let a
    -- pipelined 'sendPrepare' steal the 'ParseComplete' of an earlier
    -- 'sendQueryParams' and shift every later result by one.
    Connection -> IORef (Seq Bool)
pendingParseOrigins :: IORef (Seq.Seq Bool),
    Connection -> IORef Verbosity
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.
    Connection -> IORef ByteString
currentQuery :: IORef ByteString
  }

-- | Send a serialized frontend message.
sendMessage :: Connection -> Poker.Write -> IO ()
sendMessage :: Connection -> Write -> IO ()
sendMessage Connection
connection Write
write = do
  Transport
transport <- IORef Transport -> IO Transport
forall a. IORef a -> IO a
readIORef (Connection -> IORef Transport
transport Connection
connection)
  Transport -> Write -> IO ()
Transport.send Transport
transport Write
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 -> IO BackendMessage
nextMessage Connection
connection = do
  Transport
transport <- IORef Transport -> IO Transport
forall a. IORef a -> IO a
readIORef (Connection -> IORef Transport
transport Connection
connection)
  (Word8
typeByte, ByteString
body) <- Transport -> IO (Word8, ByteString)
Transport.receiveFrame Transport
transport
  case Word8 -> ByteString -> Either DecodingError BackendMessage
decodeBackendMessage Word8
typeByte ByteString
body of
    Left DecodingError
err -> IOException -> IO BackendMessage
forall a. IOException -> IO a
ioError (String -> IOException
userError (String
"pqi-native: protocol decode error: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> DecodingError -> String
forall a. Show a => a -> String
show DecodingError
err))
    Right BackendMessage
message -> case BackendMessage
message of
      ParameterStatus ByteString
key ByteString
value -> do
        IORef (Map ByteString ByteString)
-> (Map ByteString ByteString -> Map ByteString ByteString)
-> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef (Map ByteString ByteString)
parameters Connection
connection) (ByteString
-> ByteString
-> Map ByteString ByteString
-> Map ByteString ByteString
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ByteString
key ByteString
value)
        Connection -> IO BackendMessage
nextMessage Connection
connection
      NoticeResponse [(Word8, ByteString)]
fields -> do
        Bool
reporting <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (Connection -> IORef Bool
noticeReporting Connection
connection)
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
reporting (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
          let noticeText :: ByteString
noticeText = Map Word8 ByteString -> ByteString
formatErrorFields ([(Word8, ByteString)] -> Map Word8 ByteString
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Word8, ByteString)]
fields)
          Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
ByteString.null ByteString
noticeText)
            (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef [ByteString] -> ([ByteString] -> [ByteString]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef [ByteString]
notices Connection
connection) (ByteString
noticeText ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:)
        Connection -> IO BackendMessage
nextMessage Connection
connection
      NotificationResponse Int32
pid ByteString
channel ByteString
payload -> do
        IORef [Notify] -> ([Notify] -> [Notify]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef [Notify]
pendingNotifications Connection
connection) (ByteString -> Int32 -> ByteString -> Notify
Notify ByteString
channel Int32
pid ByteString
payload Notify -> [Notify] -> [Notify]
forall a. a -> [a] -> [a]
:)
        Connection -> IO BackendMessage
nextMessage Connection
connection
      BackendMessage
other -> BackendMessage -> IO BackendMessage
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BackendMessage
other

-- | Look up an error\/notice field by its single-byte code.
fieldValue :: Word8 -> [(Word8, ByteString)] -> Maybe ByteString
fieldValue :: Word8 -> [(Word8, ByteString)] -> Maybe ByteString
fieldValue Word8
code = Word8 -> [(Word8, ByteString)] -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Word8
code

-- | Record a flat error message and mark the connection bad.
setError :: Connection -> ByteString -> IO ()
setError :: Connection -> ByteString -> IO ()
setError Connection
connection ByteString
message = do
  IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
message)
  IORef ConnStatus -> ConnStatus -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef ConnStatus
connStatus Connection
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.
-- | Open a connection: resolve the default user, resolve and connect the
-- socket, send the startup message, and run the authentication\/startup
-- handshake. Like libpq, a failed connection - whether due to a user-name
-- lookup failure, a network error or a rejected handshake - yields a
-- 'ConnectionBad' connection rather than throwing.
establish :: ByteString -> IO Connection
establish :: ByteString -> IO Connection
establish ByteString
conninfo = do
  Either String ByteString
userResult <- IO (Either String ByteString)
resolveDefaultUser
  case Either String ByteString
userResult of
    Left String
message -> do
      Transport
transport <- IO Transport
Transport.unconnected
      Connection
connection <- Bool -> Transport -> ConnInfo -> IO Connection
newConnection Bool
False Transport
transport (ByteString -> ByteString -> ConnInfo
parseConnInfo ByteString
conninfo ByteString
"")
      Connection -> ByteString -> IO ()
setError Connection
connection (String -> ByteString
ByteString.Char8.pack String
message)
      Connection -> IO Connection
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Connection
connection
    Right ByteString
dfltUser -> do
      let info :: ConnInfo
info = ByteString -> ByteString -> ConnInfo
parseConnInfo ByteString
conninfo ByteString
dfltUser
      Either IOException Transport
transportResult <- forall e a. Exception e => IO a -> IO (Either e a)
try @IOException (ByteString -> Int -> IO Transport
Transport.connect (ConnInfo -> ByteString
host ConnInfo
info) (ConnInfo -> Int
port ConnInfo
info))
      case Either IOException Transport
transportResult of
        Left IOException
err -> do
          Transport
transport <- IO Transport
Transport.unconnected
          Connection
connection <- Bool -> Transport -> ConnInfo -> IO Connection
newConnection Bool
False Transport
transport ConnInfo
info
          Connection -> ByteString -> IO ()
setError Connection
connection (ByteString
"could not connect to server: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
ByteString.Char8.pack (IOException -> String
forall a. Show a => a -> String
show IOException
err))
          Connection -> IO Connection
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Connection
connection
        Right Transport
transport -> do
          Connection
connection <- Bool -> Transport -> ConnInfo -> IO Connection
newConnection Bool
False Transport
transport ConnInfo
info
          Connection -> Write -> IO ()
sendMessage Connection
connection ([(ByteString, ByteString)] -> Write
startupMessage (ConnInfo -> [(ByteString, ByteString)]
startupParams ConnInfo
info))
          Connection -> IO ()
handshake Connection
connection
          Connection -> IO Connection
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Connection
connection

-- | A \"null\" sentinel connection (the analogue of @PQnewNullConnection@): no
-- live socket, permanently in the 'ConnectionBad' state.
nullConnection :: IO Connection
nullConnection :: IO Connection
nullConnection = do
  Transport
transport <- IO Transport
Transport.unconnected
  let info :: ConnInfo
info = ByteString -> ByteString -> ConnInfo
parseConnInfo ByteString
"" ByteString
""
  Connection
conn <- Bool -> Transport -> ConnInfo -> IO Connection
newConnection Bool
True Transport
transport ConnInfo
info
  IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
conn) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"connection pointer is NULL\n")
  Connection -> IO Connection
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Connection
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 -> IO ()
reconnect Connection
connection = do
  Transport
oldTransport <- IORef Transport -> IO Transport
forall a. IORef a -> IO a
readIORef (Connection -> IORef Transport
transport Connection
connection)
  Transport -> IO ()
Transport.close Transport
oldTransport
  Transport
newTransport <- ByteString -> Int -> IO Transport
Transport.connect (ConnInfo -> ByteString
host (Connection -> ConnInfo
info Connection
connection)) (ConnInfo -> Int
port (Connection -> ConnInfo
info Connection
connection))
  IORef Transport -> Transport -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Transport
transport Connection
connection) Transport
newTransport
  IORef (Map ByteString ByteString)
-> Map ByteString ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Map ByteString ByteString)
parameters Connection
connection) Map ByteString ByteString
forall k a. Map k a
Map.empty
  IORef (Maybe (Int32, Int32)) -> Maybe (Int32, Int32) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe (Int32, Int32))
backendKey Connection
connection) Maybe (Int32, Int32)
forall a. Maybe a
Nothing
  IORef Word8 -> Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Word8
txStatus Connection
connection) Word8
0x49
  IORef ConnStatus -> ConnStatus -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef ConnStatus
connStatus Connection
connection) ConnStatus
ConnectionBad
  IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
  Connection -> Write -> IO ()
sendMessage Connection
connection ([(ByteString, ByteString)] -> Write
startupMessage (ConnInfo -> [(ByteString, ByteString)]
startupParams (Connection -> ConnInfo
info Connection
connection)))
  Connection -> IO ()
handshake Connection
connection

-- | The startup message parameter list: @user@ and @database@, plus any
-- extra conninfo params (e.g. @application_name@) forwarded verbatim.
startupParams :: ConnInfo -> [(ByteString, ByteString)]
startupParams :: ConnInfo -> [(ByteString, ByteString)]
startupParams ConnInfo
info =
  (ByteString
"user", ConnInfo -> ByteString
user ConnInfo
info) (ByteString, ByteString)
-> [(ByteString, ByteString)] -> [(ByteString, ByteString)]
forall a. a -> [a] -> [a]
: (ByteString
"database", ConnInfo -> ByteString
database ConnInfo
info) (ByteString, ByteString)
-> [(ByteString, ByteString)] -> [(ByteString, ByteString)]
forall a. a -> [a] -> [a]
: Map ByteString ByteString -> [(ByteString, ByteString)]
forall k a. Map k a -> [(k, a)]
Map.toList (ConnInfo -> Map ByteString ByteString
extraParams ConnInfo
info)

newConnection :: Bool -> Transport -> ConnInfo -> IO Connection
newConnection :: Bool -> Transport -> ConnInfo -> IO Connection
newConnection Bool
isNull Transport
transport ConnInfo
info = do
  IORef Transport
transportRef <- Transport -> IO (IORef Transport)
forall a. a -> IO (IORef a)
newIORef Transport
transport
  IORef Transport
-> ConnInfo
-> Bool
-> IORef (Map ByteString ByteString)
-> IORef (Maybe (Int32, Int32))
-> IORef Word8
-> IORef ConnStatus
-> IORef (Maybe ByteString)
-> IORef [ByteString]
-> IORef [Notify]
-> IORef Bool
-> IORef Bool
-> IORef Bool
-> IORef PipelineStatus
-> IORef Bool
-> IORef [FieldDescription]
-> IORef Bool
-> IORef Int
-> IORef Int
-> IORef (Seq Bool)
-> IORef Verbosity
-> IORef ByteString
-> Connection
Connection IORef Transport
transportRef ConnInfo
info Bool
isNull
    (IORef (Map ByteString ByteString)
 -> IORef (Maybe (Int32, Int32))
 -> IORef Word8
 -> IORef ConnStatus
 -> IORef (Maybe ByteString)
 -> IORef [ByteString]
 -> IORef [Notify]
 -> IORef Bool
 -> IORef Bool
 -> IORef Bool
 -> IORef PipelineStatus
 -> IORef Bool
 -> IORef [FieldDescription]
 -> IORef Bool
 -> IORef Int
 -> IORef Int
 -> IORef (Seq Bool)
 -> IORef Verbosity
 -> IORef ByteString
 -> Connection)
-> IO (IORef (Map ByteString ByteString))
-> IO
     (IORef (Maybe (Int32, Int32))
      -> IORef Word8
      -> IORef ConnStatus
      -> IORef (Maybe ByteString)
      -> IORef [ByteString]
      -> IORef [Notify]
      -> IORef Bool
      -> IORef Bool
      -> IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map ByteString ByteString -> IO (IORef (Map ByteString ByteString))
forall a. a -> IO (IORef a)
newIORef Map ByteString ByteString
forall k a. Map k a
Map.empty
    IO
  (IORef (Maybe (Int32, Int32))
   -> IORef Word8
   -> IORef ConnStatus
   -> IORef (Maybe ByteString)
   -> IORef [ByteString]
   -> IORef [Notify]
   -> IORef Bool
   -> IORef Bool
   -> IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef (Maybe (Int32, Int32)))
-> IO
     (IORef Word8
      -> IORef ConnStatus
      -> IORef (Maybe ByteString)
      -> IORef [ByteString]
      -> IORef [Notify]
      -> IORef Bool
      -> IORef Bool
      -> IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (Int32, Int32) -> IO (IORef (Maybe (Int32, Int32)))
forall a. a -> IO (IORef a)
newIORef Maybe (Int32, Int32)
forall a. Maybe a
Nothing
    IO
  (IORef Word8
   -> IORef ConnStatus
   -> IORef (Maybe ByteString)
   -> IORef [ByteString]
   -> IORef [Notify]
   -> IORef Bool
   -> IORef Bool
   -> IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef Word8)
-> IO
     (IORef ConnStatus
      -> IORef (Maybe ByteString)
      -> IORef [ByteString]
      -> IORef [Notify]
      -> IORef Bool
      -> IORef Bool
      -> IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Word8 -> IO (IORef Word8)
forall a. a -> IO (IORef a)
newIORef Word8
0x49 -- 'I'
    IO
  (IORef ConnStatus
   -> IORef (Maybe ByteString)
   -> IORef [ByteString]
   -> IORef [Notify]
   -> IORef Bool
   -> IORef Bool
   -> IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef ConnStatus)
-> IO
     (IORef (Maybe ByteString)
      -> IORef [ByteString]
      -> IORef [Notify]
      -> IORef Bool
      -> IORef Bool
      -> IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ConnStatus -> IO (IORef ConnStatus)
forall a. a -> IO (IORef a)
newIORef ConnStatus
ConnectionBad
    IO
  (IORef (Maybe ByteString)
   -> IORef [ByteString]
   -> IORef [Notify]
   -> IORef Bool
   -> IORef Bool
   -> IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef (Maybe ByteString))
-> IO
     (IORef [ByteString]
      -> IORef [Notify]
      -> IORef Bool
      -> IORef Bool
      -> IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe ByteString -> IO (IORef (Maybe ByteString))
forall a. a -> IO (IORef a)
newIORef (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
    IO
  (IORef [ByteString]
   -> IORef [Notify]
   -> IORef Bool
   -> IORef Bool
   -> IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef [ByteString])
-> IO
     (IORef [Notify]
      -> IORef Bool
      -> IORef Bool
      -> IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef []
    IO
  (IORef [Notify]
   -> IORef Bool
   -> IORef Bool
   -> IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef [Notify])
-> IO
     (IORef Bool
      -> IORef Bool
      -> IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Notify] -> IO (IORef [Notify])
forall a. a -> IO (IORef a)
newIORef []
    IO
  (IORef Bool
   -> IORef Bool
   -> IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef Bool)
-> IO
     (IORef Bool
      -> IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
    IO
  (IORef Bool
   -> IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef Bool)
-> IO
     (IORef Bool
      -> IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
    IO
  (IORef Bool
   -> IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef Bool)
-> IO
     (IORef PipelineStatus
      -> IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
    IO
  (IORef PipelineStatus
   -> IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef PipelineStatus)
-> IO
     (IORef Bool
      -> IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PipelineStatus -> IO (IORef PipelineStatus)
forall a. a -> IO (IORef a)
newIORef PipelineStatus
PipelineOff
    IO
  (IORef Bool
   -> IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef Bool)
-> IO
     (IORef [FieldDescription]
      -> IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
    IO
  (IORef [FieldDescription]
   -> IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef [FieldDescription])
-> IO
     (IORef Bool
      -> IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [FieldDescription] -> IO (IORef [FieldDescription])
forall a. a -> IO (IORef a)
newIORef []
    IO
  (IORef Bool
   -> IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef Bool)
-> IO
     (IORef Int
      -> IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
    IO
  (IORef Int
   -> IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef Int)
-> IO
     (IORef Int
      -> IORef (Seq Bool)
      -> IORef Verbosity
      -> IORef ByteString
      -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
    IO
  (IORef Int
   -> IORef (Seq Bool)
   -> IORef Verbosity
   -> IORef ByteString
   -> Connection)
-> IO (IORef Int)
-> IO
     (IORef (Seq Bool)
      -> IORef Verbosity -> IORef ByteString -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
    IO
  (IORef (Seq Bool)
   -> IORef Verbosity -> IORef ByteString -> Connection)
-> IO (IORef (Seq Bool))
-> IO (IORef Verbosity -> IORef ByteString -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Seq Bool -> IO (IORef (Seq Bool))
forall a. a -> IO (IORef a)
newIORef Seq Bool
forall a. Seq a
Seq.empty
    IO (IORef Verbosity -> IORef ByteString -> Connection)
-> IO (IORef Verbosity) -> IO (IORef ByteString -> Connection)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Verbosity -> IO (IORef Verbosity)
forall a. a -> IO (IORef a)
newIORef Verbosity
ErrorsDefault
    IO (IORef ByteString -> Connection)
-> IO (IORef ByteString) -> IO Connection
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ByteString -> IO (IORef ByteString)
forall a. a -> IO (IORef a)
newIORef ByteString
""

-- | The startup\/authentication state machine, ending at the first
-- @ReadyForQuery@ (success) or @ErrorResponse@ (failure).
handshake :: Connection -> IO ()
handshake :: Connection -> IO ()
handshake Connection
connection = IO ()
authenticating
  where
    authenticating :: IO ()
authenticating = do
      BackendMessage
message <- Connection -> IO BackendMessage
nextMessage Connection
connection
      case BackendMessage
message of
        BackendMessage
AuthenticationOk -> IO ()
startingUp
        BackendMessage
AuthenticationCleartextPassword -> do
          Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> Write
passwordMessage (ConnInfo -> ByteString
password (Connection -> ConnInfo
info Connection
connection)))
          IO ()
authenticating
        AuthenticationMD5Password ByteString
salt -> do
          let response :: ByteString
response = ByteString -> ByteString -> ByteString -> ByteString
Auth.md5Password (ConnInfo -> ByteString
user (Connection -> ConnInfo
info Connection
connection)) (ConnInfo -> ByteString
password (Connection -> ConnInfo
info Connection
connection)) ByteString
salt
          Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> Write
passwordMessage ByteString
response)
          IO ()
authenticating
        AuthenticationSASL [ByteString]
mechanisms ->
          ByteString
-> ByteString
-> [ByteString]
-> SaslStep
-> IO (Either ByteString ())
Auth.scram (ConnInfo -> ByteString
user (Connection -> ConnInfo
info Connection
connection)) (ConnInfo -> ByteString
password (Connection -> ConnInfo
info Connection
connection)) [ByteString]
mechanisms (Connection -> SaslStep
saslExchange Connection
connection) IO (Either ByteString ())
-> (Either ByteString () -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Left ByteString
problem -> Connection -> ByteString -> IO ()
setError Connection
connection ByteString
problem
            Right () -> IO ()
startingUp
        ErrorResponse [(Word8, ByteString)]
fields -> [(Word8, ByteString)] -> IO ()
failWith [(Word8, ByteString)]
fields
        BackendMessage
other -> [(Word8, ByteString)] -> IO ()
failWith [(Word8
0x4d, ByteString
"unexpected authentication message: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
ByteString.Char8.pack (BackendMessage -> String
forall a. Show a => a -> String
show BackendMessage
other))]
    startingUp :: IO ()
startingUp = do
      BackendMessage
message <- Connection -> IO BackendMessage
nextMessage Connection
connection
      case BackendMessage
message of
        BackendKeyData Int32
pid Int32
secret -> do
          IORef (Maybe (Int32, Int32)) -> Maybe (Int32, Int32) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe (Int32, Int32))
backendKey Connection
connection) ((Int32, Int32) -> Maybe (Int32, Int32)
forall a. a -> Maybe a
Just (Int32
pid, Int32
secret))
          IO ()
startingUp
        ReadyForQuery Word8
txState -> do
          IORef Word8 -> Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Word8
txStatus Connection
connection) Word8
txState
          IORef ConnStatus -> ConnStatus -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef ConnStatus
connStatus Connection
connection) ConnStatus
ConnectionOk
        ErrorResponse [(Word8, ByteString)]
fields -> [(Word8, ByteString)] -> IO ()
failWith [(Word8, ByteString)]
fields
        BackendMessage
_ -> IO ()
startingUp
    failWith :: [(Word8, ByteString)] -> IO ()
failWith [(Word8, ByteString)]
fields = do
      let fmtFields :: ByteString
fmtFields = Map Word8 ByteString -> ByteString
formatErrorFields ([(Word8, ByteString)] -> Map Word8 ByteString
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Word8, ByteString)]
fields)
      Transport
transport <- IORef Transport -> IO Transport
forall a. IORef a -> IO a
readIORef (Connection -> IORef Transport
transport Connection
connection)
      Maybe ByteString
mIp <- IO (Maybe ByteString)
-> (SomeException -> IO (Maybe ByteString))
-> IO (Maybe ByteString)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString)
-> IO ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Transport -> IO ByteString
Transport.peerIp Transport
transport) (\(SomeException
_ :: SomeException) -> Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ByteString
forall a. Maybe a
Nothing)
      Connection -> ByteString -> IO ()
setError Connection
connection (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ case Maybe ByteString
mIp of
        Maybe ByteString
Nothing -> ByteString
fmtFields
        Just ByteString
ip ->
          ByteString
"connection to server at \""
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ConnInfo -> ByteString
host (Connection -> ConnInfo
info Connection
connection)
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"\" ("
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
ip
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"), port "
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
ByteString.Char8.pack (Int -> String
forall a. Show a => a -> String
show (ConnInfo -> Int
port (Connection -> ConnInfo
info Connection
connection)))
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
" failed: "
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
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 -> SaslStep
saslExchange Connection
connection =
  Auth.SaslStep
    { sendInitial :: ByteString -> ByteString -> IO ()
Auth.sendInitial = \ByteString
mechanism ByteString
initial ->
        Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> ByteString -> Write
saslInitialResponse ByteString
mechanism ByteString
initial),
      sendResponse :: ByteString -> IO ()
Auth.sendResponse = \ByteString
payload ->
        Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> Write
saslResponse ByteString
payload),
      receive :: IO SaslMessage
Auth.receive = do
        BackendMessage
message <- Connection -> IO BackendMessage
nextMessage Connection
connection
        SaslMessage -> IO SaslMessage
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SaslMessage -> IO SaslMessage) -> SaslMessage -> IO SaslMessage
forall a b. (a -> b) -> a -> b
$ case BackendMessage
message of
          AuthenticationSASLContinue ByteString
payload -> ByteString -> SaslMessage
Auth.SaslContinue ByteString
payload
          AuthenticationSASLFinal ByteString
payload -> ByteString -> SaslMessage
Auth.SaslFinal ByteString
payload
          BackendMessage
AuthenticationOk -> SaslMessage
Auth.SaslOk
          ErrorResponse [(Word8, ByteString)]
fields -> ByteString -> SaslMessage
Auth.SaslError (ByteString -> Maybe ByteString -> ByteString
forall a. a -> Maybe a -> a
fromMaybe ByteString
"SASL error" (Word8 -> [(Word8, ByteString)] -> Maybe ByteString
fieldValue Word8
0x4d [(Word8, ByteString)]
fields))
          BackendMessage
other -> ByteString -> SaslMessage
Auth.SaslError (ByteString
"unexpected SASL message: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
ByteString.Char8.pack (BackendMessage -> String
forall a. Show a => a -> String
show BackendMessage
other))
    }