mysql-haskell-1.2.1: src/Database/MySQL/Connection.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE PackageImports #-}
{-|
Module : Database.MySQL.Connection
Description : Connection managment
Copyright : (c) Winterland, 2016
License : BSD
Maintainer : drkoster@qq.com
Stability : experimental
Portability : PORTABLE
This is an internal module, the 'MySQLConn' type should not directly acessed to user.
-}
module Database.MySQL.Connection
( module Database.MySQL.Connection
) where
import Control.Exception (Exception, bracketOnError,
throwIO, catch, SomeException)
import Control.Monad
import qualified Crypto.Hash as Crypto
import qualified Data.Binary as Binary
import qualified Data.Binary.Put as Binary
import Data.Bits
#if MIN_VERSION_crypton(1,1,0)
import qualified "ram" Data.ByteArray as BA
#else
import qualified "memory" Data.ByteArray as BA
#endif
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Unsafe as B
import Data.IORef (IORef, newIORef, readIORef,
writeIORef)
import Data.Typeable
import Data.Word
import Database.MySQL.Protocol.Auth
import Database.MySQL.Protocol.Command
import Database.MySQL.Protocol.Packet
import Network.Socket (HostName, PortNumber)
import qualified Network.Socket as N
import System.IO.Streams (InputStream)
import qualified System.IO.Streams as Stream
import qualified System.IO.Streams.TCP as TCP
import qualified Data.Connection as TCP
--------------------------------------------------------------------------------
-- | 'MySQLConn' wrap both 'InputStream' and 'OutputStream' for MySQL 'Packet'.
--
-- You shouldn't use one 'MySQLConn' in different thread, if you do that,
-- consider protecting it with a @MVar@.
--
data MySQLConn = MySQLConn {
mysqlRead :: {-# UNPACK #-} !(InputStream Packet)
, mysqlWrite :: (Packet -> IO ())
, mysqlCloseSocket :: IO ()
, isConsumed :: {-# UNPACK #-} !(IORef Bool)
}
-- | Everything you need to establish a MySQL connection.
--
-- To setup a TLS connection, use module "Database.MySQL.TLS" or "Database.MySQL.OpenSSL".
--
data ConnectInfo = ConnectInfo
{ ciHost :: HostName
, ciPort :: PortNumber
, ciDatabase :: ByteString
, ciUser :: ByteString
, ciPassword :: ByteString
, ciCharset :: Word8
} deriving Show
-- | A simple 'ConnectInfo' targeting localhost with @user=root@ and empty password.
--
-- Default charset is set to @utf8_general_ci@ to support older(< 5.5.3) MySQL versions,
-- but be aware this is a partial utf8 encoding, you may want to use 'defaultConnectInfoMB4'
-- instead to support full utf8 charset(emoji, etc.). You can query your server's support
-- with @SELECT id, collation_name FROM information_schema.collations ORDER BY id;@
--
defaultConnectInfo :: ConnectInfo
defaultConnectInfo = ConnectInfo "127.0.0.1" 3306 "" "root" "" utf8_general_ci
-- | 'defaultConnectInfo' with charset set to @utf8mb4_unicode_ci@
--
-- This is recommanded on any MySQL server version >= 5.5.3.
--
defaultConnectInfoMB4 :: ConnectInfo
defaultConnectInfoMB4 = ConnectInfo "127.0.0.1" 3306 "" "root" "" utf8mb4_unicode_ci
utf8_general_ci :: Word8
utf8_general_ci = 33
utf8mb4_unicode_ci :: Word8
utf8mb4_unicode_ci = 224
--------------------------------------------------------------------------------
-- | Socket buffer size.
--
-- maybe exposed to 'ConnectInfo' laster?
--
bUFSIZE :: Int
bUFSIZE = 16384
-- | Establish a MySQL connection.
--
connect :: ConnectInfo -> IO MySQLConn
connect = fmap snd . connectDetail
-- | Establish a MySQL connection with 'Greeting' back, so you can find server's version .etc.
--
connectDetail :: ConnectInfo -> IO (Greeting, MySQLConn)
connectDetail (ConnectInfo host port db user pass charset)
= bracketOnError open TCP.close go
where
open = connectWithBufferSize host port bUFSIZE
go c = do
let is = TCP.source c
is' <- decodeInputStream is
p <- readPacket is'
greet <- decodeFromPacket p
let auth = mkAuth db user pass charset greet
write c $ encodeToPacket 1 auth
q <- readPacket is'
completeAuth is' (write c) pass q plainFullAuth
consumed <- newIORef True
let waitNotMandatoryOK = catch
(void (waitCommandReply is')) -- server will either reply an OK packet
((\ _ -> return ()) :: SomeException -> IO ()) -- or directy close the connection
conn = MySQLConn is'
(write c)
(writeCommand COM_QUIT (write c) >> waitNotMandatoryOK >> TCP.close c)
consumed
return (greet, conn)
connectWithBufferSize h p bs = TCP.connectSocket h p >>= TCP.socketToConnection bs
write c a = TCP.send c $ Binary.runPut . Binary.put $ a
-- | Connect to MySQL via a Unix domain socket.
--
connectUnixSocket :: FilePath -> ConnectInfo -> IO MySQLConn
connectUnixSocket socketPath ci = fmap snd (connectUnixSocketDetail socketPath ci)
-- | Connect to MySQL via a Unix domain socket with 'Greeting' back,
-- so you can find server's version .etc.
--
connectUnixSocketDetail :: FilePath -> ConnectInfo -> IO (Greeting, MySQLConn)
connectUnixSocketDetail socketPath (ConnectInfo _host _port db user pass charset)
= bracketOnError open TCP.close go
where
open = bracketOnError
(N.socket N.AF_UNIX N.Stream 0)
N.close
(\sock -> do
N.connect sock (N.SockAddrUnix socketPath)
TCP.socketToConnection bUFSIZE (sock, N.SockAddrUnix socketPath)
)
go c = do
let is = TCP.source c
is' <- decodeInputStream is
p <- readPacket is'
greet <- decodeFromPacket p
let auth = mkAuth db user pass charset greet
write c $ encodeToPacket 1 auth
q <- readPacket is'
completeAuth is' (write c) pass q plainFullAuth
consumed <- newIORef True
let waitNotMandatoryOK = catch
(void (waitCommandReply is')) -- server will either reply an OK packet
((\_ -> return ()) :: SomeException -> IO ()) -- or directy close the connection
conn = MySQLConn is'
(write c)
(writeCommand COM_QUIT (write c) >> waitNotMandatoryOK >> TCP.close c)
consumed
return (greet, conn)
write c a = TCP.send c $ Binary.runPut . Binary.put $ a
mkAuth :: ByteString -> ByteString -> ByteString -> Word8 -> Greeting -> Auth
mkAuth db user pass charset greet =
let salt = greetingSalt1 greet `B.append` greetingSalt2 greet
plugin = greetingAuthPlugin greet
scambleBuf = scrambleForPlugin plugin salt pass
in Auth clientCap clientMaxPacketSize charset user scambleBuf db plugin
-- | Dispatch scramble based on the authentication plugin name.
scrambleForPlugin :: ByteString -> ByteString -> ByteString -> ByteString
scrambleForPlugin plugin salt pass
| plugin == "caching_sha2_password" = scrambleSHA256 salt pass
| otherwise = scrambleSHA1 salt pass
-- | SHA1-based scramble for @mysql_native_password@.
scrambleSHA1 :: ByteString -> ByteString -> ByteString
scrambleSHA1 salt pass
| B.null pass = B.empty
| otherwise = B.pack (B.zipWith xor sha1pass withSalt)
where sha1pass = sha1 pass
withSalt = sha1 (salt `B.append` sha1 sha1pass)
sha1 :: ByteString -> ByteString
sha1 = BA.convert . (Crypto.hash :: ByteString -> Crypto.Digest Crypto.SHA1)
-- | SHA256-based scramble for @caching_sha2_password@.
-- XOR(SHA256(password), SHA256(SHA256(SHA256(password)) + nonce))
scrambleSHA256 :: ByteString -> ByteString -> ByteString
scrambleSHA256 salt pass
| B.null pass = B.empty
| otherwise = B.pack (B.zipWith xor sha256pass withSalt)
where sha256pass = sha256 pass
withSalt = sha256 (sha256 sha256pass `B.append` salt)
sha256 :: ByteString -> ByteString
sha256 = BA.convert . (Crypto.hash :: ByteString -> Crypto.Digest Crypto.SHA256)
-- | Handle multi-step authentication after sending the initial auth response.
--
-- This handles OK, ERR, AuthMoreData (0x01), and AuthSwitchRequest (0xFE).
-- The @fullAuth@ callback is invoked when the server requests full authentication
-- (e.g., cleartext password over TLS).
completeAuth :: InputStream Packet -- ^ packet input stream
-> (Packet -> IO ()) -- ^ packet writer
-> ByteString -- ^ password
-> Packet -- ^ the first response packet from server
-> (Word8 -> ByteString -> (Packet -> IO ()) -> InputStream Packet -> IO ())
-- ^ full auth callback (seqN, password, writer, input)
-> IO ()
completeAuth is writePacket pass p fullAuth
| isOK p = return ()
| isERR p = decodeFromPacket p >>= throwIO . ERRException
| isAuthMoreData p = do
let body = L.toStrict (pBody p)
case B.index body 1 of
0x03 -> do -- fast auth success, read the final OK
ok <- readPacket is
if isOK ok
then return ()
else decodeFromPacket ok >>= throwIO . ERRException
0x04 -> do -- full auth required
fullAuth (pSeqN p + 1) pass writePacket is
_ -> throwIO (UnexpectedPacket p)
| isAuthSwitch p = do
-- Parse AuthSwitchRequest: 0xFE, plugin name (NUL), salt
let body = L.toStrict (pBody p)
rest = B.drop 1 body -- skip 0xFE
(newPlugin, rest') = B.break (== 0) rest
newSalt = B.drop 1 rest' -- skip NUL; trailing NUL may or may not be present
-- Remove trailing NUL from salt if present
newSalt' = if not (B.null newSalt) && B.last newSalt == 0
then B.init newSalt
else newSalt
scrambled = scrambleForPlugin newPlugin newSalt' pass
seqN = pSeqN p + 1
responseBody = L.fromStrict scrambled
responsePacket = Packet (fromIntegral (B.length scrambled)) seqN responseBody
writePacket responsePacket
q <- readPacket is
completeAuth is writePacket pass q fullAuth
| otherwise = throwIO (UnexpectedPacket p)
-- | Full auth handler for plain TCP connections: throws an error because
-- caching_sha2_password full authentication requires a secure connection.
plainFullAuth :: Word8 -> ByteString -> (Packet -> IO ()) -> InputStream Packet -> IO ()
plainFullAuth _ _ _ _ =
throwIO $ AuthException "caching_sha2_password full authentication requires a TLS connection. Use Database.MySQL.TLS to connect, or ensure the password verifier is cached (fast auth path)."
data AuthException = AuthException String deriving (Typeable, Show)
instance Exception AuthException
-- | A specialized 'decodeInputStream' here for speed
decodeInputStream :: InputStream ByteString -> IO (InputStream Packet)
decodeInputStream is = Stream.makeInputStream $ do
bs <- Stream.readExactly 4 is
let len = fromIntegral (bs `B.unsafeIndex` 0)
.|. fromIntegral (bs `B.unsafeIndex` 1) `shiftL` 8
.|. fromIntegral (bs `B.unsafeIndex` 2) `shiftL` 16
seqN = bs `B.unsafeIndex` 3
body <- loopRead [] len is
return . Just $ Packet len seqN body
where
loopRead acc 0 _ = return $! L.fromChunks (reverse acc)
loopRead acc k is' = do
bs <- Stream.read is'
case bs of Nothing -> throwIO NetworkException
Just bs' -> do let l = fromIntegral (B.length bs')
if l >= k
then do
let (a, rest) = B.splitAt (fromIntegral k) bs'
unless (B.null rest) (Stream.unRead rest is')
return $! L.fromChunks (reverse (a:acc))
else do
let k' = k - l
k' `seq` loopRead (bs':acc) k' is'
-- | Close a MySQL connection.
--
close :: MySQLConn -> IO ()
close (MySQLConn _ _ closeSocket _) = closeSocket
-- | Send a 'COM_PING'.
--
ping :: MySQLConn -> IO OK
ping = flip command COM_PING
--------------------------------------------------------------------------------
-- helpers
-- | Send a 'Command' which don't return a resultSet.
--
command :: MySQLConn -> Command -> IO OK
command conn@(MySQLConn is os _ _) cmd = do
guardUnconsumed conn
writeCommand cmd os
waitCommandReply is
{-# INLINE command #-}
waitCommandReply :: InputStream Packet -> IO OK
waitCommandReply is = do
p <- readPacket is
if | isERR p -> decodeFromPacket p >>= throwIO . ERRException
| isOK p -> decodeFromPacket p
| otherwise -> throwIO (UnexpectedPacket p)
{-# INLINE waitCommandReply #-}
waitCommandReplys :: InputStream Packet -> IO [OK]
waitCommandReplys is = do
p <- readPacket is
if | isERR p -> decodeFromPacket p >>= throwIO . ERRException
| isOK p -> do ok <- decodeFromPacket p
if isThereMore ok
then (ok :) <$> waitCommandReplys is
else return [ok]
| otherwise -> throwIO (UnexpectedPacket p)
{-# INLINE waitCommandReplys #-}
readPacket :: InputStream Packet -> IO Packet
readPacket is = Stream.read is >>= maybe
(throwIO NetworkException)
(\ p@(Packet len _ bs) -> if len < 16777215 then return p else go len [bs])
where
go len acc = Stream.read is >>= maybe
(throwIO NetworkException)
(\ (Packet len' seqN bs) -> do
let len'' = len + len'
acc' = bs:acc
if len' < 16777215
then return (Packet len'' seqN (L.concat . reverse $ acc'))
else len'' `seq` go len'' acc'
)
{-# INLINE readPacket #-}
writeCommand :: Command -> (Packet -> IO ()) -> IO ()
writeCommand a writePacket = let bs = Binary.runPut (putCommand a) in
go (L.length bs) 0 bs writePacket
where
go len seqN bs writePacket' = do
if len < 16777215
then writePacket (Packet len seqN bs)
else do
let (bs', rest) = L.splitAt 16777215 bs
seqN' = seqN + 1
len' = len - 16777215
writePacket (Packet 16777215 seqN bs')
seqN' `seq` len' `seq` go len' seqN' rest writePacket'
{-# INLINE writeCommand #-}
guardUnconsumed :: MySQLConn -> IO ()
guardUnconsumed (MySQLConn _ _ _ consumed) = do
c <- readIORef consumed
unless c (throwIO UnconsumedResultSet)
{-# INLINE guardUnconsumed #-}
writeIORef' :: IORef a -> a -> IO ()
writeIORef' ref x = x `seq` writeIORef ref x
{-# INLINE writeIORef' #-}
--------------------------------------------------------------------------------
-- Exceptions
data NetworkException = NetworkException deriving (Typeable, Show)
instance Exception NetworkException
data UnconsumedResultSet = UnconsumedResultSet deriving (Typeable, Show)
instance Exception UnconsumedResultSet
data ERRException = ERRException ERR deriving (Typeable, Show)
instance Exception ERRException
data UnexpectedPacket = UnexpectedPacket Packet deriving (Typeable, Show)
instance Exception UnexpectedPacket