门禁拆分与提速

chusql-server/src/ChuSQL/Server/Privileges.hs

@@ -7,6 +7,7 @@ module ChuSQL.Server.Privileges (
     PrivilegeError (..),
     Privileges,
     RoleView (..),
+    affectedAccounts,
     authorize,
     authorizeTables,
     filterTables,
@@ -111,6 +112,33 @@ runPrivilegeCommand service principal database command
         RevokePrivilegesCommand _ object _ -> T.strip object /= "*"
         _ -> False
 
+-- | 一条命令会影响的账号：角色连它的成员一起展开，账号名原样留下
+affectedAccounts :: Privileges -> PrivilegeCommand -> IO [Text]
+affectedAccounts service command = do
+    known <- roleNames service
+    case known of
+        Left _ -> pure (direct command)
+        Right roles -> walk roles [] (direct command)
+  where
+    -- | 命令直接点到的名字
+    direct cmd = case cmd of
+        CreateRoleCommand _ -> []
+        DropRoleCommand role -> [role]
+        GrantPrivilegesCommand _ _ role _ -> [role]
+        RevokePrivilegesCommand _ _ role -> [role]
+        GrantRoleCommand _ members -> members
+        RevokeRoleCommand _ members -> members
+    -- | 是角色就接着往下走它的成员，走过的名字不再走第二遍
+    walk roles seen [] = pure (reverse seen)
+    walk roles seen (name : rest)
+        | name `elem` seen = walk roles seen rest
+        | normalizeRole name `elem` roles = do
+            members <- membersOfRole service name
+            case members of
+                Left _ -> walk roles (name : seen) rest
+                Right found -> walk roles (name : seen) (found ++ rest)
+        | otherwise = walk roles (name : seen) rest
+
 -- | 语句鉴权：管理员全通，普通身份按需查
 authorize :: Privileges -> Principal -> Text -> Statement -> IO (Either PrivilegeError ())
 authorize service principal database stmt

chusql-server/src/ChuSQL/Server/Session.hs

@@ -6,6 +6,7 @@ module ChuSQL.Server.Session
     , SessionError (..)
     , engineErrorCode
     , newSession
+    , newSessionWith
     , sessionUser
     , sessionIsAdmin
     , sessionDatabase
@@ -36,7 +37,7 @@ import ChuSQL.Interface.Auth (defaultSessionPolicy, newSessionStore)
 import ChuSQL.Server.Accounts
 import ChuSQL.Server.Backend (Backend (..), StatementResult (..), columnsFromStatement)
 import ChuSQL.Server.Policy (configurePasswordPolicy)
-import ChuSQL.Server.Privileges (PrivilegeCommand, PrivilegeError (..), Privileges, RoleView, authorize, filterTables, listRoleViews, newPrivileges, privilegeCommand, runPrivilegeCommand)
+import ChuSQL.Server.Privileges (PrivilegeCommand, PrivilegeError (..), Privileges, RoleView, affectedAccounts, authorize, filterTables, listRoleViews, newPrivileges, privilegeCommand, runPrivilegeCommand)
 import Control.Exception (IOException, try)
 import Data.Bifunctor (first)
 import Data.Char (isAlpha, isAlphaNum)
@@ -67,6 +68,7 @@ data Session = Session
     , ssPrincipal :: IORef (Maybe Principal)
     , ssSettingsFile :: FilePath
     , ssTransaction :: IORef (Maybe Transaction)
+    , ssChanged :: [Text] -> IO ()
     }
 
 -- | 一条会话里的显式事务：开事务时的快照、事务内改到的状态与保存点栈
@@ -78,7 +80,11 @@ data Transaction = Transaction
 
 -- | 开一个会话：账号与角色服务、配置策略都和 Web 端一致
 newSession :: Backend -> Text -> FilePath -> IO Session
-newSession backend name settingsFile = do
+newSession backend name settingsFile = newSessionWith backend name settingsFile (const (pure ()))
+
+-- | 开会话并接上账号变更通知回调
+newSessionWith :: Backend -> Text -> FilePath -> ([Text] -> IO ()) -> IO Session
+newSessionWith backend name settingsFile changed = do
     sessions <- newSessionStore getCurrentTime defaultSessionPolicy
     accountsService <- newAccounts backend sessions name
     privileges <- newPrivileges backend
@@ -87,7 +93,7 @@ newSession backend name settingsFile = do
     current <- newIORef ""
     principal <- newIORef Nothing
     transaction <- newIORef Nothing
-    pure (Session backend accountsService privileges current principal settingsFile transaction)
+    pure (Session backend accountsService privileges current principal settingsFile transaction changed)
 
 -- | 当前生效的口令策略（只有管理员能看）
 policyOf :: Session -> IO (Either Text PasswordPolicy)
@@ -529,7 +535,7 @@ currentBackend session = do
     name <- readIORef (ssCurrent session)
     pure (beWithDatabase (ssBackend session) (T.unpack name))
 
--- | 账号语句走账号服务，普通账号一律拒绝
+-- | 账号语句走账号服务，普通账号一律拒绝；改完通知受影响账号
 runAccount :: Session -> AccountCommand -> IO (Either SessionError QueryResult)
 runAccount session command = do
     who <- readIORef (ssPrincipal session)
@@ -537,9 +543,20 @@ runAccount session command = do
         Nothing -> pure (Left (SessionError "unauthorized" "sign in first"))
         Just principal -> do
             result <- runAccountCommand (ssAccounts session) principal command
-            pure (either (Left . accountErrorOf) (const (Right (emptyResult Nothing))) result)
+            case result of
+                Left err -> pure (Left (accountErrorOf err))
+                Right () -> do
+                    notifyChanged session (accountCommandUsers command)
+                    pure (Right (emptyResult Nothing))
 
--- | 角色语句走权限服务，普通账号拒绝
+-- | 一条账号命令影响的账号
+accountCommandUsers :: AccountCommand -> [Text]
+accountCommandUsers command = case command of
+    CreateAccount _ _ -> []
+    ResetAccountPassword name _ -> [name]
+    DropAccount name -> [name]
+
+-- | 角色语句走权限服务，普通账号拒绝；改完通知受影响账号
 runPrivilege :: Session -> PrivilegeCommand -> IO (Either SessionError QueryResult)
 runPrivilege session command = do
     who <- readIORef (ssPrincipal session)
@@ -547,8 +564,18 @@ runPrivilege session command = do
         Nothing -> pure (Left (SessionError "unauthorized" "sign in first"))
         Just principal -> do
             database <- readIORef (ssCurrent session)
+            -- 受影响账号要趁命令落地之前算：删角色之后成员就查不到了
+            affected <- affectedAccounts (ssPrivileges session) command
             result <- runPrivilegeCommand (ssPrivileges session) principal database command
-            pure (either (Left . privilegeErrorOf) (const (Right (emptyResult Nothing))) result)
+            case result of
+                Left err -> pure (Left (privilegeErrorOf err))
+                Right () -> do
+                    notifyChanged session affected
+                    pure (Right (emptyResult Nothing))
+
+-- | 通知变更：把受影响的账号交给会话的变更钩子
+notifyChanged :: Session -> [Text] -> IO ()
+notifyChanged session names = ssChanged session (filter (not . T.null) names)
 
 -- | 账号错误转成会话错误
 accountErrorOf :: AccountError -> SessionError

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

@@ -5,6 +5,8 @@ module ChuSQL.Server.TCP (
     defaultServerConfig,
     loadServerConfigAt,
     ServerEnv (..),
+    LiveConnections,
+    newLiveConnections,
     newServerEnv,
     ServerHandle (..),
     startServer,
@@ -17,13 +19,15 @@ import ChuSQL.Interface.Protocol (ClientRequest (..), ServerResponse (..), decod
 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 ChuSQL.Server.Session (QueryResult (..), Session, SessionError (..), accounts, authenticateSessionCoded, catalog, databases, newSessionWith, 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 Control.Exception (IOException, finally, try)
+import Control.Monad (filterM, void)
 import qualified Data.ByteString.Char8 as BSC
 import qualified Data.ByteString.Lazy.Char8 as BLC
+import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef)
+import qualified Data.Map.Strict as Map
 import Data.Maybe (fromMaybe)
 import Data.Text (Text)
 import qualified Data.Text as T
@@ -45,6 +49,8 @@ import Network.Socket (
     getSocketName,
     listen,
     setSocketOption,
+    shutdown,
+    ShutdownCmd (ShutdownBoth),
     socket,
     socketToHandle,
   )
@@ -53,6 +59,41 @@ import System.IO.Error (isEOFError)
 
 -- TCP 服务器：监听端口，一连接一会话，按一行一 JSON 收发协议消息。
 
+-- 在线连接与服务器环境
+
+-- | 在线连接登记：账号名加一个「关掉这条连接」的动作
+data LiveConnections = LiveConnections
+    { lcNext :: IORef Int
+    , lcLive :: MVar (Map.Map Int (IORef (Maybe Text), IO ()))
+    }
+
+-- | 建一张空的在线连接登记表
+newLiveConnections :: IO LiveConnections
+newLiveConnections = LiveConnections <$> newIORef 0 <*> newMVar Map.empty
+
+-- | 登记一条连接，返回注销动作
+registerConnection :: LiveConnections -> IORef (Maybe Text) -> IO () -> IO (IO ())
+registerConnection live who kick = do
+    n <- atomicModifyIORef' (lcNext live) (\i -> (i + 1, i))
+    modifyMVar_ (lcLive live) (pure . Map.insert n (who, kick))
+    pure (modifyMVar_ (lcLive live) (pure . Map.delete n))
+
+-- | 断开这些账号的在线连接，返回断开的条数
+dropConnectionsOf :: LiveConnections -> [Text] -> IO Int
+dropConnectionsOf live names = do
+    let wanted = map T.toLower (filter (not . T.null) names)
+    entries <- readMVar (lcLive live)
+    hits <- filterM (isOf wanted) (Map.elems entries)
+    mapM_ (tryKick . snd) hits
+    pure (length hits)
+  where
+    -- | 这条连接登录的账号在名单里吗
+    isOf wanted (who, _) = do
+        current <- readIORef who
+        pure (maybe False (\name -> T.toLower name `elem` wanted) current)
+    -- | 踢人失败不当错误：对端可能已经自己走了
+    tryKick kick = void (try kick :: IO (Either IOException ()))
+
 -- 服务器
 
 data ServerEnv = ServerEnv
@@ -61,12 +102,14 @@ data ServerEnv = ServerEnv
     , srvSettingsFile :: FilePath
     , srvConfig :: ServerConfig
     , srvLimiter :: RateLimiter
+    , srvLive :: LiveConnections
     }
 
 -- | 组装服务器环境：每个连接再各开一个会话
 newServerEnv :: Backend -> Text -> FilePath -> ServerConfig -> RateLimiter -> IO ServerEnv
-newServerEnv backend rootName settingsFile config limiter =
-    pure (ServerEnv backend rootName settingsFile config limiter)
+newServerEnv backend rootName settingsFile config limiter = do
+    live <- newLiveConnections
+    pure (ServerEnv backend rootName settingsFile config limiter live)
 
 data ServerHandle = ServerHandle
     { shPort :: Int
@@ -148,31 +191,48 @@ reportConnection env conn = do
         Left err -> hPutStrLn stderr ("connection failed: " ++ show err)
         Right () -> pure ()
 
+-- | 一条连接上跨请求的状态：登录账号、已被断开与注销动作
+data Connection = Connection
+    { cnWho :: IORef (Maybe Text)
+    , cnDropped :: IORef Bool
+    }
+
 -- | 一个连接一个会话：一行一条 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 ())
+    who <- newIORef Nothing
+    dropped <- newIORef False
+    let kick = do
+            writeIORef dropped True
+            void (try (shutdown conn ShutdownBoth) :: IO (Either IOException ()))
+    unregister <- registerConnection (srvLive env) who kick
+    let connection = Connection who dropped
+    session <- newSessionWith (srvBackend env) (srvRootName env) (srvSettingsFile env) (\names -> void (dropConnectionsOf (srvLive env) names))
+    outcome <- try (loop h session connection `finally` unregister) :: 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
+    loop h session connection = 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
+            then send h (RespError "too_large" "request line too large") >> loop h session connection
             else case decodeRequest cleaned of
-                Left message -> send h (RespError "bad_request" message) >> loop h session
+                Left message -> send h (RespError "bad_request" message) >> loop h session connection
                 Right request -> do
-                    (response, keepGoing) <- dispatch env session request
-                    send h response
-                    if keepGoing then loop h session else pure ()
+                    expired <- readIORef (cnDropped connection)
+                    if expired
+                        then send h (RespError "unauthorized" "session was dropped by an account change")
+                        else do
+                            (response, keepGoing) <- dispatch env session connection request
+                            send h response
+                            if keepGoing then loop h session connection else pure ()
     -- | 发一条响应并冲缓冲
     send h response = do
         BLC.hPutStr h (encodeResponse response)
@@ -180,13 +240,16 @@ serveConnection env conn = do
         hFlush h
 
 -- | 一条请求一条响应，quit 后停连接
-dispatch :: ServerEnv -> Session -> ClientRequest -> IO (ServerResponse, Bool)
-dispatch env session request = case request of
+dispatch :: ServerEnv -> Session -> Connection -> ClientRequest -> IO (ServerResponse, Bool)
+dispatch env session connection 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
+        case response of
+            RespLogin who _ -> writeIORef (cnWho connection) (Just who)
+            _ -> pure ()
         pure (response, True)
     ReqQuery sql -> do
         who <- sessionUser session

chusql-server/test/Spec.hs

@@ -23,6 +23,7 @@ import ChuSQL.Server.Privileges (
     PrivilegeCommand (..),
     PrivilegeError (..),
     Privileges,
+    affectedAccounts,
     authorize,
     listRoleViews,
     newPrivileges,
@@ -480,6 +481,67 @@ tcpServerSpec = do
                         blocked <- tcpQuery other "select * from users"
                         asText (at "code" blocked) `shouldBe` "unauthorized"
 
+    describe "TCP server: account changes drop live sessions" $ do
+        it "disconnects a member whose role was revoked" $
+            withTcpServer defaultServerConfig $ \_ handle ->
+                withConnection handle $ \admin -> do
+                    _ <- tcpLogin admin "admin" "s3cret"
+                    _ <- tcpQuery admin "create user bob identified by 'bob-password-1'"
+                    _ <- tcpQuery admin "create role reader"
+                    _ <- tcpQuery admin "use test"
+                    _ <- tcpQuery admin "grant select on users to reader"
+                    _ <- tcpQuery admin "grant reader to bob"
+                    withConnection handle $ \bob -> do
+                        _ <- tcpLogin bob "bob" "bob-password-1"
+                        _ <- tcpQuery bob "use test"
+                        allowed <- tcpQuery bob "select name from users"
+                        asText (at "status" allowed) `shouldBe` "result"
+                        _ <- tcpQuery admin "revoke reader from bob"
+                        tcpClosed bob
+        it "disconnects every member of a role that was dropped" $
+            withTcpServer defaultServerConfig $ \_ handle ->
+                withConnection handle $ \admin -> do
+                    _ <- tcpLogin admin "admin" "s3cret"
+                    _ <- tcpQuery admin "create user bob identified by 'bob-password-1'"
+                    _ <- tcpQuery admin "create role reader"
+                    _ <- tcpQuery admin "grant reader to bob"
+                    withConnection handle $ \bob -> do
+                        _ <- tcpLogin bob "bob" "bob-password-1"
+                        _ <- tcpQuery admin "drop role reader"
+                        tcpClosed bob
+        it "disconnects an account whose password was reset" $
+            withTcpServer defaultServerConfig $ \_ handle ->
+                withConnection handle $ \admin -> do
+                    _ <- tcpLogin admin "admin" "s3cret"
+                    _ <- tcpQuery admin "create user bob identified by 'bob-password-1'"
+                    withConnection handle $ \bob -> do
+                        _ <- tcpLogin bob "bob" "bob-password-1"
+                        _ <- tcpQuery admin "alter user bob identified by 'bob-password-2'"
+                        tcpClosed bob
+                    withConnection handle $ \again -> do
+                        relogin <- tcpLogin again "bob" "bob-password-2"
+                        asText (at "status" relogin) `shouldBe` "ok"
+        it "disconnects an account that was dropped" $
+            withTcpServer defaultServerConfig $ \_ handle ->
+                withConnection handle $ \admin -> do
+                    _ <- tcpLogin admin "admin" "s3cret"
+                    _ <- tcpQuery admin "create user bob identified by 'bob-password-1'"
+                    withConnection handle $ \bob -> do
+                        _ <- tcpLogin bob "bob" "bob-password-1"
+                        _ <- tcpQuery admin "drop user bob"
+                        tcpClosed bob
+        it "leaves the connections of untouched accounts alone" $
+            withTcpServer defaultServerConfig $ \_ handle ->
+                withConnection handle $ \admin -> do
+                    _ <- tcpLogin admin "admin" "s3cret"
+                    _ <- tcpQuery admin "create user bob identified by 'bob-password-1'"
+                    withConnection handle $ \bob -> do
+                        _ <- tcpLogin bob "bob" "bob-password-1"
+                        _ <- tcpQuery admin "create user carol identified by 'carol-password-1'"
+                        _ <- tcpQuery admin "alter user carol identified by 'carol-password-2'"
+                        still <- tcpQuery bob "use test"
+                        asText (at "status" still) `shouldBe` "result"
+
 -- 权限服务夹具：内存后端上的系统库视角
 
 -- | 一个普通账号
@@ -695,6 +757,17 @@ privilegeSpec = describe "server privileges" $ do
         authorizeSql service (Ordinary testAccount) "SELECT name FROM users WHERE id > ANY (SELECT id FROM orders)"
             `shouldReturn` Left (PrivilegeError "forbidden" "permission denied: SELECT ON orders")
 
+    it "lists the accounts a role change touches, members included" $ withPrivileges $ \service root -> do
+        _ <- runPrivilegeCommand service root testDatabase (CreateRoleCommand "reader")
+        _ <- runPrivilegeCommand service root testDatabase (CreateRoleCommand "inner")
+        _ <- runPrivilegeCommand service root testDatabase (GrantRoleCommand "inner" ["reader"])
+        _ <- runPrivilegeCommand service root testDatabase (GrantRoleCommand "reader" ["alice"])
+        directly <- affectedAccounts service (RevokeRoleCommand "inner" ["bob"])
+        directly `shouldMatchList` ["inner", "bob", "reader", "alice"]
+        dropped <- affectedAccounts service (DropRoleCommand "inner")
+        dropped `shouldMatchList` ["inner", "reader", "alice"]
+        affectedAccounts service (CreateRoleCommand "other") `shouldReturn` []
+
     it "keeps the role tables on the real storage" $ withIpcStorage $ do
         base <- ipcBackend
         service <- newPrivileges base