ChuBlog/chusql-server/src/ChuSQL/Server/TCP.hs

{-# LANGUAGE OverloadedStrings #-}

module ChuSQL.Server.TCP (
    ServerConfig (..),
    defaultServerConfig,
    loadServerConfigAt,
    ServerEnv (..),
    newServerEnv,
    ServerHandle (..),
    startServer,
    runServer,
) where

import ChuSQL.Interface.Config (ServerConfig (..), defaultServerConfig, loadServerConfigAt)
import ChuSQL.Interface.Policy (PasswordPolicy (..))
import ChuSQL.Interface.Protocol (ClientRequest (..), ServerResponse (..), decodeRequest, encodeResponse, protocolVersion)
import ChuSQL.Interface.RateLimit (RateLimiter, rateLimitBlock, rateLimitClear, rateLimitRecord)
import qualified ChuSQL.Core.Engine.Error as E
import ChuSQL.Server.Backend (Backend (..))
import ChuSQL.Server.Session (QueryResult (..), Session, SessionError (..), accounts, authenticateSessionCoded, catalog, databases, newSession, policyOf, reloadPolicy, roleViews, runStatementCoded, sessionIsAdmin, sessionUser)
import Control.Concurrent (forkIO)
import Control.Concurrent.MVar (MVar, modifyMVar_, newMVar, readMVar)
import Control.Exception (IOException, try)
import Control.Monad (void)
import qualified Data.ByteString.Char8 as BSC
import qualified Data.ByteString.Lazy.Char8 as BLC
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import Network.Socket (
    AddrInfo (..),
    AddrInfoFlag (AI_PASSIVE),
    SockAddr (SockAddrInet, SockAddrInet6),
    Socket,
    SocketOption (ReuseAddr),
    SocketType (Stream),
    accept,
    addrAddress,
    addrFamily,
    bind,
    close,
    defaultHints,
    defaultProtocol,
    getAddrInfo,
    getSocketName,
    listen,
    setSocketOption,
    socket,
    socketToHandle,
  )
import System.IO (BufferMode (LineBuffering), IOMode (ReadWriteMode), hClose, hFlush, hPutStrLn, hSetBuffering, hSetNewlineMode, noNewlineTranslation, stderr)
import System.IO.Error (isEOFError)

-- TCP 服务器：监听端口，一连接一会话，按一行一 JSON 收发协议消息。

-- 服务器

data ServerEnv = ServerEnv
    { srvBackend :: Backend
    , srvRootName :: Text
    , srvSettingsFile :: FilePath
    , srvConfig :: ServerConfig
    , srvLimiter :: RateLimiter
    }

-- | 组装服务器环境：每个连接再各开一个会话
newServerEnv :: Backend -> Text -> FilePath -> ServerConfig -> RateLimiter -> IO ServerEnv
newServerEnv backend rootName settingsFile config limiter =
    pure (ServerEnv backend rootName settingsFile config limiter)

data ServerHandle = ServerHandle
    { shPort :: Int
    , shStop :: IO ()
    }

-- | 起监听并在线程里收连接（测试与嵌入用）
startServer :: ServerEnv -> IO ServerHandle
startServer env = do
    sock <- listenSocket (scHost (srvConfig env)) (scPort (srvConfig env))
    port <- actualPort sock
    stopped <- newMVar False
    _ <- forkIO (acceptLoop env sock stopped)
    pure (ServerHandle port (stopListening sock stopped))

-- | 在调用线程里收连接（可执行文件用）；绑好端口先回调一次
runServer :: ServerEnv -> (Int -> IO ()) -> IO ()
runServer env announce = do
    sock <- listenSocket (scHost (srvConfig env)) (scPort (srvConfig env))
    port <- actualPort sock
    announce port
    stopped <- newMVar False
    acceptLoop env sock stopped

-- | 置停止标记并关掉监听套接字
stopListening :: Socket -> MVar Bool -> IO ()
stopListening sock stopped = do
    modifyMVar_ stopped (const (pure True))
    void (try (close sock) :: IO (Either IOException ()))

-- | 建监听套接字；端口 0 表示由系统分配
listenSocket :: String -> Int -> IO Socket
listenSocket host port = do
    let hints = defaultHints{addrSocketType = Stream, addrFlags = [AI_PASSIVE]}
        node = if null host then Nothing else Just host
    infos <- getAddrInfo (Just hints) node (Just (show port))
    case infos of
        [] -> ioError (userError ("cannot resolve listen address: " ++ host))
        (info : _) -> do
            sock <- socket (addrFamily info) Stream defaultProtocol
            setSocketOption sock ReuseAddr 1
            bind sock (addrAddress info)
            listen sock 128
            pure sock

-- | 系统分配端口后读回真正的端口
actualPort :: Socket -> IO Int
actualPort sock = do
    address <- getSocketName sock
    pure $ case address of
        SockAddrInet port _ -> fromIntegral port
        SockAddrInet6 port _ _ _ -> fromIntegral port
        _ -> 0

-- | 循环 accept，把每个连接交给独立线程
acceptLoop :: ServerEnv -> Socket -> MVar Bool -> IO ()
acceptLoop env sock stopped = loop
  where
    -- | 等下一个连接
    loop = do
        stop <- readMVar stopped
        if stop
            then pure ()
            else do
                accepted <- try (accept sock) :: IO (Either IOException (Socket, SockAddr))
                case accepted of
                    Left err -> do
                        stopping <- readMVar stopped
                        if stopping then pure () else hPutStrLn stderr ("accept failed: " ++ show err)
                    Right (conn, _) -> do
                        void (forkIO (reportConnection env conn))
                        loop

-- | 连接线程：异常写到 stderr
reportConnection :: ServerEnv -> Socket -> IO ()
reportConnection env conn = do
    outcome <- try (serveConnection env conn) :: IO (Either IOException ())
    case outcome of
        Left err -> hPutStrLn stderr ("connection failed: " ++ show err)
        Right () -> pure ()

-- | 一个连接一个会话：一行一条 JSON，读到 EOF 就收工
serveConnection :: ServerEnv -> Socket -> IO ()
serveConnection env conn = do
    h <- socketToHandle conn ReadWriteMode
    hSetBuffering h LineBuffering
    hSetNewlineMode h noNewlineTranslation
    session <- newSession (srvBackend env) (srvRootName env) (srvSettingsFile env)
    outcome <- try (loop h session) :: IO (Either IOException ())
    case outcome of
        Left err | not (isEOFError err) -> hPutStrLn stderr ("connection error: " ++ show err)
        _ -> pure ()
    void (try (hClose h) :: IO (Either IOException ()))
  where
    -- | 逐行读请求并应答
    loop h session = do
        line <- BSC.hGetLine h
        let cleaned = BSC.dropWhileEnd (== '\r') line
        if BSC.length cleaned > scMaxMessage (srvConfig env)
            then send h (RespError "too_large" "request line too large") >> loop h session
            else case decodeRequest cleaned of
                Left message -> send h (RespError "bad_request" message) >> loop h session
                Right request -> do
                    (response, keepGoing) <- dispatch env session request
                    send h response
                    if keepGoing then loop h session else pure ()
    -- | 发一条响应并冲缓冲
    send h response = do
        BLC.hPutStr h (encodeResponse response)
        BLC.hPutStr h "\n"
        hFlush h

-- | 一条请求一条响应，quit 后停连接
dispatch :: ServerEnv -> Session -> ClientRequest -> IO (ServerResponse, Bool)
dispatch env session request = case request of
    ReqHello version
        | version == protocolVersion -> pure (RespHello protocolVersion, True)
        | otherwise -> pure (RespError "bad_request" ("unsupported protocol version: " <> tshow version), True)
    ReqLogin user password -> do
        response <- login env session user password
        pure (response, True)
    ReqQuery sql -> do
        who <- sessionUser session
        response <- case who of
            Nothing -> pure (RespError "unauthorized" "sign in first")
            Just _
                | T.null (T.strip sql) -> pure (RespError "bad_request" "sql must not be empty")
                | otherwise -> do
                    result <- runStatementCoded session sql
                    pure $ case result of
                        Left err -> RespError (sessCode err) (sessMessage err)
                        Right queryResult -> RespResult (trimResult (scMaxRows (srvConfig env)) queryResult)
        pure (response, True)
    ReqStorage payload -> do
        who <- sessionUser session
        response <- case who of
            Nothing -> pure (RespError "unauthorized" "sign in first")
            Just _ -> do
                admin <- sessionIsAdmin session
                if not admin
                    then pure (RespError "forbidden" "administrator required")
                    else do
                        result <- beStorage (srvBackend env) payload
                        pure $ case result of
                            Left err -> RespError "storage_error" (T.pack err)
                            Right value -> RespStorage value
        pure (response, True)
    ReqCatalog -> do
        response <- overSession session (catalog session) RespCatalog
        pure (response, True)
    ReqDatabases -> do
        response <- overSession session (databases session) RespDatabases
        pure (response, True)
    ReqRoles -> do
        response <- overSession session (roleViews session) RespRoles
        pure (response, True)
    ReqAccounts -> do
        response <- overSession session (accounts session) RespAccounts
        pure (response, True)
    ReqPolicy -> do
        response <- overSession session (policyOf session) policyResponse
        pure (response, True)
    ReqReloadPolicy -> do
        response <- overSession session (reloadPolicy session) policyResponse
        pure (response, True)
    ReqPing -> pure (RespPong, True)
    ReqQuit -> pure (RespBye, False)

-- | 口令策略在线上只有两个数字：最短长度与至少几类字符
policyResponse :: PasswordPolicy -> ServerResponse
policyResponse (PasswordPolicy minLength classes) = RespPolicy minLength classes

-- | 会话方法统一口径：没登录先挡，其余失败按码表落到响应
overSession :: Session -> IO (Either Text a) -> (a -> ServerResponse) -> IO ServerResponse
overSession session action wrap = do
    who <- sessionUser session
    case who of
        Nothing -> pure (RespError "unauthorized" "sign in first")
        Just _ -> do
            result <- action
            pure $ case result of
                Left message -> RespError (sessionCode message) message
                Right value -> wrap value

-- | 会话层的失败文案归一到协议错误码
sessionCode :: Text -> Text
sessionCode = T.pack . E.errorCode . T.unpack

-- | 登录：按账号名限流，成功清账
login :: ServerEnv -> Session -> Text -> Text -> IO ServerResponse
login env session user password = do
    let key = T.toLower (T.strip user)
    blocked <- rateLimitBlock (srvLimiter env) key
    case blocked of
        Just seconds -> pure (RespError "too_many_attempts" ("too many failed logins, retry in " <> tshow (ceiling seconds :: Integer) <> "s"))
        Nothing -> do
            result <- authenticateSessionCoded session key password
            case result of
                Left err -> do
                    rateLimitRecord (srvLimiter env) key
                    pure (RespError (sessCode err) (sessMessage err))
                Right () -> do
                    rateLimitClear (srvLimiter env) key
                    who <- sessionUser session
                    admin <- sessionIsAdmin session
                    pure (RespLogin (fromMaybe key who) admin)

-- | 按上限截行，总行数保持不变
trimResult :: Int -> QueryResult -> QueryResult
trimResult limit result
    | length (qrRows result) > limit = result{qrRows = take limit (qrRows result), qrTruncated = True}
    | otherwise = result

-- | show 成 Text
tshow :: Show a => a -> Text
tshow = T.pack . show