content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DuplicateRecordFields #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE NumericUnderscores #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  module Simplex.Messaging.Transport.Client
    ( runTransportClient,
      runTLSTransportClient,
      smpClientHandshake,
      defaultSMPPort,
      defaultTcpConnectTimeout,
      defaultTransportClientConfig,
      defaultSocksProxyWithAuth,
      defaultSocksProxy,
      defaultSocksHost,
      TransportClientConfig (..),
      SocksProxy (..),
      SocksProxyWithAuth (..),
      SocksAuth (..),
      TransportHost (..),
      TransportHosts (..),
      TransportHosts_ (..),
      validateCertificateChain,
    )
  where
  
  import Control.Applicative (optional, (<|>))
  import Control.Logger.Simple (logError)
  import Control.Monad
  import Data.Aeson (FromJSON (..), ToJSON (..))
  import qualified Data.Attoparsec.ByteString.Char8 as A
  import Data.ByteString.Char8 (ByteString)
  import qualified Data.ByteString.Char8 as B
  import Data.Char (isAsciiLower, isDigit, isHexDigit)
  import Data.Default (def)
  import Data.Functor (($>))
  import Data.IORef
  import Data.IP
  import Data.List.NonEmpty (NonEmpty (..))
  import qualified Data.List.NonEmpty as L
  import Data.Maybe (fromMaybe)
  import Data.String
  import Data.Word (Word32, Word8)
  import qualified Data.X509 as X
  import qualified Data.X509.CertificateStore as XS
  import Data.X509.Validation (Fingerprint (..))
  import qualified Data.X509.Validation as XV
  import GHC.IO.Exception (IOErrorType (..))
  import Network.Socket
  import Network.Socks5
  import qualified Network.TLS as T
  import qualified Simplex.Messaging.Crypto as C
  import Simplex.Messaging.Encoding
  import Simplex.Messaging.Encoding.String
  import Simplex.Messaging.Parsers (parseAll, parseString)
  import Simplex.Messaging.Transport
  import Simplex.Messaging.Transport.KeepAlive
  import Simplex.Messaging.Transport.Shared
  import Simplex.Messaging.Util (bshow, catchAll, catchAll_, tshow, (<$?>))
  import System.IO.Error
  import Text.Read (readMaybe)
  import UnliftIO.Exception (IOException)
  import qualified UnliftIO.Exception as E
  import UnliftIO.STM
  
  data TransportHost
    = THIPv4 (Word8, Word8, Word8, Word8)
    | THIPv6 (Word32, Word32, Word32, Word32)
    | THOnionHost ByteString
    | THDomainName HostName
    deriving (Eq, Ord, Show)
  
  instance Encoding TransportHost where
    smpEncode = smpEncode . strEncode
    smpP = parseAll strP <$?> smpP
  
  instance StrEncoding TransportHost where
    strEncode = \case
      THIPv4 (a1, a2, a3, a4) -> B.intercalate "." $ map bshow [a1, a2, a3, a4]
      THIPv6 addr -> bshow $ toIPv6w addr
      THOnionHost host -> host
      THDomainName host -> B.pack host
    strP =
      A.choice
        [ THIPv4 <$> ((,,,) <$> ipNum <*> ipNum <*> ipNum <*> A.decimal),
          maybe (Left "bad IPv6") (Right . THIPv6 . fromIPv6w) . readMaybe . B.unpack <$?> ipv6StrP,
          THOnionHost <$> ((<>) <$> A.takeWhile (\c -> isAsciiLower c || isDigit c) <*> A.string ".onion"),
          THDomainName . B.unpack <$> (notOnion <$?> A.takeWhile1 (A.notInClass ":#,;/ \n\r\t"))
        ]
      where
        ipNum = validIP <$?> (A.decimal <* A.char '.')
        validIP :: Int -> Either String Word8
        validIP n = if 0 <= n && n <= 255 then Right $ fromIntegral n else Left "invalid IP address"
        ipv6StrP =
          A.char '[' *> A.takeWhile1 (/= ']') <* A.char ']'
            <|> A.takeWhile1 (\c -> isHexDigit c || c == ':')
        notOnion s = if ".onion" `B.isSuffixOf` s then Left "invalid onion host" else Right s
  
  instance ToJSON TransportHost where
    toEncoding = strToJEncoding
    toJSON = strToJSON
  
  instance FromJSON TransportHost where
    parseJSON = strParseJSON "TransportHost"
  
  newtype TransportHosts = TransportHosts {thList :: NonEmpty TransportHost}
  
  instance StrEncoding TransportHosts where
    strEncode = strEncodeList . L.toList . thList
    strP = TransportHosts . L.fromList <$> strP `A.sepBy1'` A.char ','
  
  newtype TransportHosts_ = TransportHosts_ {thList_ :: [TransportHost]}
  
  instance StrEncoding TransportHosts_ where
    strEncode = strEncodeList . thList_
    strP = TransportHosts_ <$> strP `A.sepBy'` A.char ','
  
  instance IsString TransportHost where fromString = parseString strDecode
  
  instance IsString (NonEmpty TransportHost) where fromString = parseString strDecode
  
  data TransportClientConfig = TransportClientConfig
    { socksProxy :: Maybe SocksProxy,
      tcpConnectTimeout :: Int,
      tcpKeepAlive :: Maybe KeepAliveOpts,
      logTLSErrors :: Bool,
      clientCredentials :: Maybe T.Credential,
      clientALPN :: Maybe [ALPN],
      useSNI :: Bool
    }
    deriving (Eq, Show)
  
  -- time to resolve host, connect socket, set up TLS
  defaultTcpConnectTimeout :: Int
  defaultTcpConnectTimeout = 25_000_000
  
  defaultTransportClientConfig :: TransportClientConfig
  defaultTransportClientConfig =
    TransportClientConfig
      { socksProxy = Nothing,
        tcpConnectTimeout = defaultTcpConnectTimeout,
        tcpKeepAlive = Just defaultKeepAliveOpts,
        logTLSErrors = True,
        clientCredentials = Nothing,
        clientALPN = Nothing,
        useSNI = True
      }
  
  clientTransportConfig :: TransportClientConfig -> TransportConfig
  clientTransportConfig TransportClientConfig {logTLSErrors} =
    TransportConfig {logTLSErrors, transportTimeout = Nothing}
  
  -- | Connect to passed TCP host:port and pass handle to the client.
  runTransportClient :: Transport c => TransportClientConfig -> Maybe SocksCredentials -> TransportHost -> ServiceName -> Maybe C.KeyHash -> (c 'TClient -> IO a) -> IO a
  runTransportClient = runTLSTransportClient defaultSupportedParams Nothing
  
  data ConnectionHandle c
    = CHSocket Socket
    | CHContext T.Context
    | CHTransport (c 'TClient)
  
  runTLSTransportClient :: Transport c => T.Supported -> Maybe XS.CertificateStore -> TransportClientConfig -> Maybe SocksCredentials -> TransportHost -> ServiceName -> Maybe C.KeyHash -> (c 'TClient -> IO a) -> IO a
  runTLSTransportClient tlsParams caStore_ cfg@TransportClientConfig {socksProxy, tcpKeepAlive, clientCredentials, clientALPN, useSNI} socksCreds host port keyHash client = do
    serverCert <- newEmptyTMVarIO
    clientCredsSent <- newIORef False
    let hostName = B.unpack $ strEncode host
        clientParams = mkTLSClientParams tlsParams caStore_ hostName port keyHash clientCredentials clientCredsSent clientALPN useSNI serverCert
        connectTCP = case socksProxy of
          Just proxy -> connectSocksClient proxy socksCreds (hostAddr host)
          _ -> connectTCPClient hostName
    h <- newIORef Nothing
    let set hc = (>>= \c -> writeIORef h (Just $ hc c) $> c)
    E.bracket (set CHSocket $ connectTCP port) (\_ -> closeConn h) $ \sock -> do
      mapM_ (setSocketKeepAlive sock) tcpKeepAlive `catchAll` \e -> logError ("Error setting TCP keep-alive " <> tshow e)
      let tCfg = clientTransportConfig cfg
      -- No TLS timeout to avoid failing connections via SOCKS
      tls <- set CHContext $ connectTLS (Just hostName) tCfg clientParams sock
      chain <- takePeerCertChain serverCert
      sent <- readIORef clientCredsSent
      client =<< set CHTransport (getTransportConnection tCfg sent chain tls)
    where
      closeConn = readIORef >=> mapM_ (\c -> E.uninterruptibleMask_ $ closeConn_ c `catchAll_` pure ())
      closeConn_ = \case
        CHSocket sock -> close sock
        CHContext tls -> closeTLS tls
        CHTransport c -> closeConnection c
      hostAddr = \case
        THIPv4 addr -> SocksAddrIPV4 $ tupleToHostAddress addr
        THIPv6 addr -> SocksAddrIPV6 addr
        THOnionHost h -> SocksAddrDomainName h
        THDomainName h -> SocksAddrDomainName $ B.pack h
  
  connectTCPClient :: HostName -> ServiceName -> IO Socket
  connectTCPClient host port = withSocketsDo $ resolve >>= tryOpen err
    where
      err :: IOException
      err = mkIOError NoSuchThing "no address" Nothing Nothing
  
      resolve :: IO [AddrInfo]
      resolve =
        let hints = defaultHints {addrSocketType = Stream}
         in getAddrInfo (Just hints) (Just host) (Just port)
  
      tryOpen :: IOException -> [AddrInfo] -> IO Socket
      tryOpen e [] = E.throwIO e
      tryOpen _ (addr : as) =
        E.try (open addr) >>= either (`tryOpen` as) pure
  
      open :: AddrInfo -> IO Socket
      open addr =
        E.bracketOnError
          (socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr))
          close
          (\sock -> connect sock (addrAddress addr) $> sock)
  
  defaultSMPPort :: PortNumber
  defaultSMPPort = 5223
  
  connectSocksClient :: SocksProxy -> Maybe SocksCredentials -> SocksHostAddress -> ServiceName -> IO Socket
  connectSocksClient (SocksProxy addr) socksCreds hostAddr _port = do
    let port = if null _port then defaultSMPPort else fromMaybe defaultSMPPort $ readMaybe _port
    fst <$> case socksCreds of
      Just creds -> socksConnectAuth (defaultSocksConf addr) (SocksAddress hostAddr port) creds
      _ -> socksConnect (defaultSocksConf addr) (SocksAddress hostAddr port)
  
  defaultSocksHost :: (Word8, Word8, Word8, Word8)
  defaultSocksHost = (127, 0, 0, 1)
  
  defaultSocksProxyWithAuth :: SocksProxyWithAuth
  defaultSocksProxyWithAuth = SocksProxyWithAuth SocksIsolateByAuth defaultSocksProxy
  
  defaultSocksProxy :: SocksProxy
  defaultSocksProxy = SocksProxy $ SockAddrInet 9050 $ tupleToHostAddress defaultSocksHost
  
  newtype SocksProxy = SocksProxy SockAddr
    deriving (Eq)
  
  data SocksProxyWithAuth = SocksProxyWithAuth SocksAuth SocksProxy
    deriving (Eq, Show)
  
  data SocksAuth
    = SocksAuthUsername {username :: ByteString, password :: ByteString}
    | SocksAuthNull
    | SocksIsolateByAuth -- this is default
    deriving (Eq, Show)
  
  instance Show SocksProxy where show (SocksProxy addr) = show addr
  
  instance StrEncoding SocksProxy where
    strEncode = B.pack . show
    strP = do
      host <- fromMaybe (THIPv4 defaultSocksHost) <$> optional strP
      port <- fromMaybe 9050 <$> optional (A.char ':' *> (fromInteger <$> A.decimal))
      SocksProxy <$> socksAddr port host
      where
        socksAddr port = \case
          THIPv4 addr -> pure $ SockAddrInet port $ tupleToHostAddress addr
          THIPv6 addr -> pure $ SockAddrInet6 port 0 addr 0
          _ -> fail "SOCKS5 host should be IPv4 or IPv6 address"
  
  instance StrEncoding SocksProxyWithAuth where
    strEncode (SocksProxyWithAuth auth proxy) = strEncode auth <> strEncode proxy
    strP = SocksProxyWithAuth <$> strP <*> strP
  
  instance ToJSON SocksProxyWithAuth where
    toJSON = strToJSON
    toEncoding = strToJEncoding
  
  instance FromJSON SocksProxyWithAuth where
    parseJSON = strParseJSON "SocksProxyWithAuth"
  
  instance StrEncoding SocksAuth where
    strEncode = \case
      SocksAuthUsername {username, password} -> username <> ":" <> password <> "@"
      SocksAuthNull -> "@"
      SocksIsolateByAuth -> ""
    strP = usernameP <|> (SocksAuthNull <$ A.char '@') <|> pure SocksIsolateByAuth
      where
        usernameP = do
          username <- A.takeTill (== ':') <* A.char ':'
          password <- A.takeTill (== '@') <* A.char '@'
          pure SocksAuthUsername {username, password}
  
  mkTLSClientParams :: T.Supported -> Maybe XS.CertificateStore -> HostName -> ServiceName -> Maybe C.KeyHash -> Maybe T.Credential -> IORef Bool -> Maybe [ALPN] -> Bool -> TMVar (Maybe X.CertificateChain) -> T.ClientParams
  mkTLSClientParams supported caStore_ host port cafp_ clientCreds_ clientCredsSent alpn_ sni serverCerts =
    (T.defaultParamsClient host p)
      { T.clientUseServerNameIndication = sni,
        T.clientShared = def {T.sharedCAStore = fromMaybe (T.sharedCAStore def) caStore_},
        T.clientHooks =
          def
            { T.onServerCertificate = onServerCert,
              T.onCertificateRequest = onCertRequest,
              T.onSuggestALPN = pure alpn_
            },
        T.clientSupported = supported
      }
    where
      p = B.pack port
      onServerCert _ _ _ cc = do
        errs <- maybe def (\ca -> validateCertificateChain ca host p cc) cafp_
        atomically $ putTMVar serverCerts $ if null errs then Just cc else Nothing
        pure errs
      onCertRequest = case clientCreds_ of
        Just _ -> \_ -> clientCreds_ <$ writeIORef clientCredsSent True
        Nothing -> \_ -> pure Nothing
  
  validateCertificateChain :: C.KeyHash -> HostName -> ByteString -> X.CertificateChain -> IO [XV.FailedReason]
  validateCertificateChain (C.KeyHash kh) host port cc = case chainIdCaCerts cc of
    CCEmpty -> pure [XV.EmptyChain]
    CCSelf _ -> pure [XV.EmptyChain]
    CCValid {idCert, caCert} -> validate idCert caCert
    CCLong -> pure [XV.AuthorityTooDeep]
    where
      validate idCert caCert
        | Fingerprint kh == XV.getFingerprint idCert X.HashSHA256 = x509validate caCert (host, port) cc
        | otherwise = pure [XV.UnknownCA]
  