mysql-haskell-1.2.3: 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 Crypto.Hash.Algorithms (SHA1(..))
import qualified Crypto.Number.Serialize as Serialize
import qualified Crypto.PubKey.RSA as RSA
import qualified Crypto.PubKey.RSA.OAEP as OAEP
import qualified Data.ASN1.BinaryEncoding as ASN1B
import qualified Data.ASN1.BitArray as ASN1BA
import qualified Data.ASN1.Encoding as ASN1E
import qualified Data.ASN1.Types as ASN1
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 qualified Data.PEM as PEM
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 salt = greetingSalt1 greet `B.append` greetingSalt2 greet
auth = mkAuth db user pass charset greet
write c $ encodeToPacket 1 auth
q <- readPacket is'
completeAuth is' (write c) pass q (plainFullAuth salt)
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 salt = greetingSalt1 greet `B.append` greetingSalt2 greet
auth = mkAuth db user pass charset greet
write c $ encodeToPacket 1 auth
q <- readPacket is'
completeAuth is' (write c) pass q (plainFullAuth salt)
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 using RSA encryption.
--
-- The nonce is captured via partial application at the call site:
-- @completeAuth is writePacket pass q (plainFullAuth salt)@
plainFullAuth :: ByteString -> Word8 -> ByteString -> (Packet -> IO ()) -> InputStream Packet -> IO ()
plainFullAuth nonce seqN pass writePacket is = do
-- Request server's RSA public key
writePacket (putToPacket seqN (Binary.putWord8 0x02))
keyPkt <- readPacket is
-- Key packet body: 0x01 <PEM bytes>
let pemKey = L.toStrict (L.drop 1 (pBody keyPkt))
doRSAAuth writePacket (pSeqN keyPkt + 1) pass nonce pemKey
q <- readPacket is
if isOK q then pure ()
else decodeFromPacket q >>= throwIO . ERRException
-- | Parse a PEM-encoded SubjectPublicKeyInfo RSA public key (PKCS#8 / RFC 5480).
parseRSAPublicKey :: ByteString -> Either String RSA.PublicKey
parseRSAPublicKey pemBytes =
either (Left . show) Right (PEM.pemParseBS pemBytes)
>>= firstPEM
>>= decodeDER
>>= spkiBitString
>>= decodeDER
>>= rsaKey
where
firstPEM [] = Left "empty PEM"
firstPEM (x:_) = Right (PEM.pemContent x)
decodeDER bs = either (Left . show) Right (ASN1E.decodeASN1' ASN1B.DER bs)
spkiBitString (ASN1.Start ASN1.Sequence
: ASN1.Start ASN1.Sequence
: ASN1.OID _
: ASN1.Null
: ASN1.End ASN1.Sequence
: ASN1.BitString bs
: ASN1.End ASN1.Sequence
: [])
= Right (ASN1BA.bitArrayGetData bs)
spkiBitString xs = Left ("unexpected SPKI ASN1: " ++ show xs)
rsaKey (ASN1.Start ASN1.Sequence
: ASN1.IntVal n
: ASN1.IntVal e
: ASN1.End ASN1.Sequence
: [])
= Right RSA.PublicKey
{ RSA.public_size = B.length (Serialize.i2osp n :: ByteString)
, RSA.public_n = n
, RSA.public_e = e
}
rsaKey xs = Left ("unexpected RSA key ASN1: " ++ show xs)
-- | XOR @(password ++ NUL)@ with a cyclically-extended nonce.
xorPasswordNonce :: ByteString -> ByteString -> ByteString
xorPasswordNonce pass nonce =
let msg = pass `B.snoc` 0x00
nonceLen = B.length nonce
cycled = B.pack [nonce `B.index` (i `mod` nonceLen) | i <- [0 .. B.length msg - 1]]
in B.pack (B.zipWith xor msg cycled)
-- | Encrypt password with the server's RSA public key and send it.
doRSAAuth :: (Packet -> IO ()) -> Word8 -> ByteString -> ByteString -> ByteString -> IO ()
doRSAAuth writePacket seqN pass nonce pemKey = do
pubKey <- either (\e -> throwIO (userError ("RSA key parse: " ++ e))) pure
(parseRSAPublicKey pemKey)
result <- OAEP.encrypt (OAEP.defaultOAEPParams SHA1) pubKey
(xorPasswordNonce pass nonce)
ct <- either (\e -> throwIO (userError ("RSA encrypt: " ++ show e))) pure result
writePacket (putToPacket seqN (Binary.putByteString ct))
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