-- | Authentication helpers: MD5 password hashing and the SASL\/SCRAM-SHA-256
-- exchange.
--
-- The pure crypto lives here; the message round-trip is abstracted as 'SaslStep'
-- so the connection module owns the actual socket I\/O.
module Pqi.Native.Auth
  ( md5Password,
    SaslStep (..),
    SaslMessage (..),
    scram,
  )
where

import Crypto.Hash (Digest, MD5 (..), SHA256 (..), hashWith)
import Crypto.KDF.PBKDF2 (Parameters (..), fastPBKDF2_SHA256)
import Crypto.MAC.HMAC (HMAC, hmac, hmacGetDigest)
import Crypto.Random (getRandomBytes)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Base64 as Base64
import qualified Data.ByteString.Char8 as ByteString.Char8
import qualified Data.List as List
import Pqi.Native.Prelude

-- | Compute the response to an @AuthenticationMD5Password@ challenge:
-- @"md5" <> md5hex (md5hex (password <> user) <> salt)@.
md5Password :: ByteString -> ByteString -> ByteString -> ByteString
md5Password :: ByteString -> ByteString -> ByteString -> ByteString
md5Password ByteString
user ByteString
password ByteString
salt =
  ByteString
"md5" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString
md5Hex (ByteString -> ByteString
md5Hex (ByteString
password ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
user) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
salt)

md5Hex :: ByteString -> ByteString
md5Hex :: ByteString -> ByteString
md5Hex = String -> ByteString
ByteString.Char8.pack (String -> ByteString)
-> (ByteString -> String) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Digest MD5 -> String
forall a. Show a => a -> String
show (Digest MD5 -> String)
-> (ByteString -> Digest MD5) -> ByteString -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MD5 -> ByteString -> Digest MD5
forall ba alg.
(ByteArrayAccess ba, HashAlgorithm alg) =>
alg -> ba -> Digest alg
hashWith MD5
MD5

-- | One server SASL\/authentication message, as the SCRAM logic sees it.
data SaslMessage
  = SaslContinue ByteString
  | SaslFinal ByteString
  | SaslOk
  | SaslError ByteString

-- | The message round-trip the SCRAM exchange drives, supplied by the
-- connection module.
data SaslStep = SaslStep
  { SaslStep -> ByteString -> ByteString -> IO ()
sendInitial :: ByteString -> ByteString -> IO (),
    SaslStep -> ByteString -> IO ()
sendResponse :: ByteString -> IO (),
    SaslStep -> IO SaslMessage
receive :: IO SaslMessage
  }

mechanismName :: ByteString
mechanismName :: ByteString
mechanismName = ByteString
"SCRAM-SHA-256"

-- | Run the SCRAM-SHA-256 exchange (without channel binding). Returns @Right ()@
-- once the server accepts the client proof, or @Left@ with a problem
-- description.
scram :: ByteString -> ByteString -> [ByteString] -> SaslStep -> IO (Either ByteString ())
scram :: ByteString
-> ByteString
-> [ByteString]
-> SaslStep
-> IO (Either ByteString ())
scram ByteString
_user ByteString
password [ByteString]
mechanisms SaslStep
step
  | ByteString
mechanismName ByteString -> [ByteString] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [ByteString]
mechanisms =
      Either ByteString () -> IO (Either ByteString ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either ByteString ()
forall a b. a -> Either a b
Left ByteString
"server did not offer SCRAM-SHA-256")
  | Bool
otherwise = do
      ByteString
clientNonce <- ByteString -> ByteString
Base64.encode (ByteString -> ByteString) -> IO ByteString -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO ByteString
forall byteArray. ByteArray byteArray => Int -> IO byteArray
forall (m :: * -> *) byteArray.
(MonadRandom m, ByteArray byteArray) =>
Int -> m byteArray
getRandomBytes Int
18
      let clientFirstBare :: ByteString
clientFirstBare = ByteString
"n=,r=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
clientNonce
      (SaslStep -> ByteString -> ByteString -> IO ()
sendInitial SaslStep
step) ByteString
mechanismName (ByteString
"n,," ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
clientFirstBare)
      (SaslStep -> IO SaslMessage
receive SaslStep
step) IO SaslMessage
-> (SaslMessage -> IO (Either ByteString ()))
-> IO (Either ByteString ())
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        SaslError ByteString
problem -> Either ByteString () -> IO (Either ByteString ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either ByteString ()
forall a b. a -> Either a b
Left ByteString
problem)
        SaslContinue ByteString
serverFirst ->
          case ByteString -> Maybe (ByteString, ByteString, Int)
parseServerFirst ByteString
serverFirst of
            Maybe (ByteString, ByteString, Int)
Nothing -> Either ByteString () -> IO (Either ByteString ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either ByteString ()
forall a b. a -> Either a b
Left ByteString
"malformed SCRAM server-first message")
            Just (ByteString
serverNonce, ByteString
salt, Int
iterations) -> do
              let saltedPassword :: ByteString
saltedPassword = Parameters -> ByteString -> ByteString -> ByteString
forall password salt out.
(ByteArrayAccess password, ByteArrayAccess salt, ByteArray out) =>
Parameters -> password -> salt -> out
fastPBKDF2_SHA256 (Int -> Int -> Parameters
Parameters Int
iterations Int
32) ByteString
password ByteString
salt :: ByteString
                  clientKey :: ByteString
clientKey = ByteString -> ByteString -> ByteString
hmacSha256 ByteString
saltedPassword ByteString
"Client Key"
                  storedKey :: ByteString
storedKey = ByteString -> ByteString
sha256 ByteString
clientKey
                  clientFinalWithoutProof :: ByteString
clientFinalWithoutProof = ByteString
"c=biws,r=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
serverNonce
                  authMessage :: ByteString
authMessage =
                    ByteString -> [ByteString] -> ByteString
ByteString.intercalate ByteString
"," [ByteString
clientFirstBare, ByteString
serverFirst, ByteString
clientFinalWithoutProof]
                  clientSignature :: ByteString
clientSignature = ByteString -> ByteString -> ByteString
hmacSha256 ByteString
storedKey ByteString
authMessage
                  clientProof :: ByteString
clientProof = ByteString -> ByteString -> ByteString
xorBytes ByteString
clientKey ByteString
clientSignature
                  clientFinal :: ByteString
clientFinal = ByteString
clientFinalWithoutProof ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
",p=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString
Base64.encode ByteString
clientProof
              (SaslStep -> ByteString -> IO ()
sendResponse SaslStep
step) ByteString
clientFinal
              (SaslStep -> IO SaslMessage
receive SaslStep
step) IO SaslMessage
-> (SaslMessage -> IO (Either ByteString ()))
-> IO (Either ByteString ())
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                SaslFinal ByteString
_ -> Either ByteString () -> IO (Either ByteString ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() -> Either ByteString ()
forall a b. b -> Either a b
Right ())
                SaslMessage
SaslOk -> Either ByteString () -> IO (Either ByteString ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() -> Either ByteString ()
forall a b. b -> Either a b
Right ())
                SaslError ByteString
problem -> Either ByteString () -> IO (Either ByteString ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either ByteString ()
forall a b. a -> Either a b
Left ByteString
problem)
                SaslMessage
_ -> Either ByteString () -> IO (Either ByteString ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either ByteString ()
forall a b. a -> Either a b
Left ByteString
"unexpected SCRAM server-final message")
        SaslMessage
_ -> Either ByteString () -> IO (Either ByteString ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either ByteString ()
forall a b. a -> Either a b
Left ByteString
"unexpected SCRAM message")

-- | Parse @r=<nonce>,s=<salt base64>,i=<iterations>@ into the server nonce, the
-- decoded salt, and the iteration count.
parseServerFirst :: ByteString -> Maybe (ByteString, ByteString, Int)
parseServerFirst :: ByteString -> Maybe (ByteString, ByteString, Int)
parseServerFirst ByteString
message = do
  let attributes :: [ByteString]
attributes = Char -> ByteString -> [ByteString]
ByteString.Char8.split Char
',' ByteString
message
  ByteString
nonce <- ByteString -> [ByteString] -> Maybe ByteString
forall {t :: * -> *}.
Foldable t =>
ByteString -> t ByteString -> Maybe ByteString
attributeValue ByteString
"r=" [ByteString]
attributes
  ByteString
saltEncoded <- ByteString -> [ByteString] -> Maybe ByteString
forall {t :: * -> *}.
Foldable t =>
ByteString -> t ByteString -> Maybe ByteString
attributeValue ByteString
"s=" [ByteString]
attributes
  ByteString
salt <- (String -> Maybe ByteString)
-> (ByteString -> Maybe ByteString)
-> Either String ByteString
-> Maybe ByteString
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe ByteString -> String -> Maybe ByteString
forall a b. a -> b -> a
const Maybe ByteString
forall a. Maybe a
Nothing) ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Either String ByteString
Base64.decode ByteString
saltEncoded)
  ByteString
iterationsText <- ByteString -> [ByteString] -> Maybe ByteString
forall {t :: * -> *}.
Foldable t =>
ByteString -> t ByteString -> Maybe ByteString
attributeValue ByteString
"i=" [ByteString]
attributes
  (Int
iterations, ByteString
_) <- ByteString -> Maybe (Int, ByteString)
ByteString.Char8.readInt ByteString
iterationsText
  (ByteString, ByteString, Int)
-> Maybe (ByteString, ByteString, Int)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
nonce, ByteString
salt, Int
iterations)
  where
    attributeValue :: ByteString -> t ByteString -> Maybe ByteString
attributeValue ByteString
prefix =
      (ByteString -> ByteString) -> Maybe ByteString -> Maybe ByteString
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> ByteString -> ByteString
ByteString.drop (ByteString -> Int
ByteString.length ByteString
prefix))
        (Maybe ByteString -> Maybe ByteString)
-> (t ByteString -> Maybe ByteString)
-> t ByteString
-> Maybe ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString -> Bool) -> t ByteString -> Maybe ByteString
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
List.find (ByteString -> ByteString -> Bool
ByteString.isPrefixOf ByteString
prefix)

hmacSha256 :: ByteString -> ByteString -> ByteString
hmacSha256 :: ByteString -> ByteString -> ByteString
hmacSha256 ByteString
key ByteString
message = Digest SHA256 -> ByteString
forall a. Digest a -> ByteString
digestBytes (HMAC SHA256 -> Digest SHA256
forall a. HMAC a -> Digest a
hmacGetDigest (ByteString -> ByteString -> HMAC SHA256
forall key message a.
(ByteArrayAccess key, ByteArrayAccess message, HashAlgorithm a) =>
key -> message -> HMAC a
hmac ByteString
key ByteString
message :: HMAC SHA256))

sha256 :: ByteString -> ByteString
sha256 :: ByteString -> ByteString
sha256 = Digest SHA256 -> ByteString
forall a. Digest a -> ByteString
digestBytes (Digest SHA256 -> ByteString)
-> (ByteString -> Digest SHA256) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SHA256 -> ByteString -> Digest SHA256
forall ba alg.
(ByteArrayAccess ba, HashAlgorithm alg) =>
alg -> ba -> Digest alg
hashWith SHA256
SHA256

-- | Extract the raw bytes of a digest via its hexadecimal 'Show' instance,
-- avoiding a direct @memory@ dependency (whose @ByteArrayAccess@ instance for
-- crypton's @Digest@ conflicts across package versions).
digestBytes :: Digest a -> ByteString
digestBytes :: forall a. Digest a -> ByteString
digestBytes = ByteString -> ByteString
hexToBytes (ByteString -> ByteString)
-> (Digest a -> ByteString) -> Digest a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ByteString
ByteString.Char8.pack (String -> ByteString)
-> (Digest a -> String) -> Digest a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Digest a -> String
forall a. Show a => a -> String
show

hexToBytes :: ByteString -> ByteString
hexToBytes :: ByteString -> ByteString
hexToBytes = [Word8] -> ByteString
ByteString.pack ([Word8] -> ByteString)
-> (ByteString -> [Word8]) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> [Word8]
forall {a}. (Ord a, Num a) => [a] -> [a]
pairs ([Word8] -> [Word8])
-> (ByteString -> [Word8]) -> ByteString -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [Word8]
ByteString.unpack
  where
    pairs :: [a] -> [a]
pairs (a
hi : a
lo : [a]
rest) = (a -> a
forall {a}. (Ord a, Num a) => a -> a
hexValue a
hi a -> a -> a
forall a. Num a => a -> a -> a
* a
16 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall {a}. (Ord a, Num a) => a -> a
hexValue a
lo) a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
pairs [a]
rest
    pairs [a]
_ = []
    hexValue :: a -> a
hexValue 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 = a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
0x30
      | 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 = a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
0x57
      | 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 = a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
0x37
      | Bool
otherwise = a
0

xorBytes :: ByteString -> ByteString -> ByteString
xorBytes :: ByteString -> ByteString -> ByteString
xorBytes ByteString
a ByteString
b = [Word8] -> ByteString
ByteString.pack ((Word8 -> Word8 -> Word8) -> ByteString -> ByteString -> [Word8]
forall a. (Word8 -> Word8 -> a) -> ByteString -> ByteString -> [a]
ByteString.zipWith Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
xor ByteString
a ByteString
b)