大改，整理项目结构，完成rust ffi调用，写好tcp接口
This commit is contained in:
LIm1tZz7 committed
1 parent
df6d348995
commit
9758439b82
118 files changed
+8103
-6321
No files matched your search
@@ -0,0 +1,288 @@
|
||||
{-# LANGUAGE OverloadedStrings #-}
|
||||
|
||||
module ChuSQL.Server.TCP (
|
||||
ServerConfig (..),
|
||||
defaultServerConfig,
|
||||
loadServerConfigAt,
|
||||
ServerEnv (..),
|
||||
newServerEnv,
|
||||
ServerHandle (..),
|
||||
startServer,
|
||||
runServer,
|
||||
) where
|
||||
|
||||
import ChuSQL.Interface.Auth (Credential)
|
||||
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 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
|
||||
, srvCredential :: Credential
|
||||
, srvSettingsFile :: FilePath
|
||||
, srvConfig :: ServerConfig
|
||||
, srvLimiter :: RateLimiter
|
||||
}
|
||||
|
||||
-- | 组装服务器环境：每个连接再各开一个会话
|
||||
newServerEnv :: Backend -> Credential -> FilePath -> ServerConfig -> RateLimiter -> IO ServerEnv
|
||||
newServerEnv backend credential settingsFile config limiter =
|
||||
pure (ServerEnv backend credential 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) (srvCredential 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 message
|
||||
| message == "no database selected" = "no_database"
|
||||
| message == "administrator required" = "forbidden"
|
||||
| otherwise = "query_error"
|
||||
|
||||
-- | 登录：按账号名限流，成功清账
|
||||
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
|
||||
Reference in new issue
Block a user