网络编程：Socket¶

import Network.Socket.ByteString (recv, sendAll)

sendAll :: Socket -> ByteString -> IO ()
recv    :: Socket -> Int -> IO ByteString

import Data.Text.Encoding (decodeUtf8, encodeUtf8)

sendAll sock (encodeUtf8 "你好\n")
bytes <- recv sock 4096
let text = decodeUtf8 bytes

import System.IO

h <- socketToHandle sock ReadWriteMode
hSetBuffering h LineBuffering   -- 每写完一行就发送，不用手动 hFlush
line <- hGetLine h              -- 读到换行符为止
hPutStrLn h "OK"
hClose h                        -- 关闭 Handle 时 Socket 也随之关闭

{-# LANGUAGE OverloadedStrings #-}

import Control.Exception (bracket)
import qualified Data.Text as T
import Data.Text.Encoding (decodeUtf8, encodeUtf8)
import qualified Data.Text.IO as TIO
import Network.Socket
import Network.Socket.ByteString (recv, sendAll)
import System.Environment (getArgs)
import System.Exit (exitFailure)
import System.IO (hPutStrLn, stderr)

main :: IO ()
main = withSocketsDo $ do
  args <- getArgs
  case args of
    [host, port, message] -> runClient host port message
    _ -> do
      hPutStrLn stderr "用法: client HOST PORT MESSAGE"
      exitFailure

runClient :: String -> String -> String -> IO ()
runClient host port message = do
  -- 1. 解析地址，取第一个结果
  let hints = defaultHints { addrSocketType = Stream }
  addr <- firstAddr <$> getAddrInfo (Just hints) (Just host) (Just port)
  -- 2. 打开 Socket 并连接，bracket 保证连接一定会被关闭
  bracket (openSocket addr) close $ \sock -> do
    connect sock (addrAddress addr)
    -- 3. 发送一行文本，然后等待回复
    sendAll sock (encodeUtf8 (T.pack message <> "\n"))
    reply <- recv sock 4096
    TIO.putStr ("服务器回复: " <> decodeUtf8 reply)

-- getAddrInfo 解析失败时会直接抛出异常，返回的列表不会为空
firstAddr :: [AddrInfo] -> AddrInfo
firstAddr (addr : _) = addr
firstAddr [] = error "getAddrInfo 没有返回任何地址"

{-# LANGUAGE OverloadedStrings #-}

import Control.Concurrent (forkFinally)
import Control.Exception (bracket)
import Control.Monad (forever, unless)
import qualified Data.ByteString as BS
import Network.Socket
import Network.Socket.ByteString (recv, sendAll)
import System.IO (BufferMode (..), hSetBuffering, stdout)

main :: IO ()
main = withSocketsDo $ do
  -- 日志输出到管道或文件时也按行刷新
  hSetBuffering stdout LineBuffering
  let hints = defaultHints { addrSocketType = Stream, addrFlags = [AI_PASSIVE] }
  addr <- firstAddr <$> getAddrInfo (Just hints) Nothing (Just "4000")
  bracket (openSocket addr) close $ \sock -> do
    -- 允许程序重启后立即复用端口
    setSocketOption sock ReuseAddr 1
    bind sock (addrAddress addr)
    listen sock 1024
    putStrLn "回显服务器监听 4000 端口"
    forever $ do
      (conn, peer) <- accept sock
      putStrLn $ "新连接: " ++ show peer
      -- 每个连接一个绿色线程，结束时（无论正常还是异常）关闭连接
      _ <- forkFinally (handleConn conn) (const $ gracefulClose conn 5000)
      return ()

handleConn :: Socket -> IO ()
handleConn conn = do
  chunk <- recv conn 4096
  -- recv 返回空串表示对方已经关闭连接
  unless (BS.null chunk) $ do
    sendAll conn chunk
    handleConn conn

firstAddr :: [AddrInfo] -> AddrInfo
firstAddr (addr : _) = addr
firstAddr [] = error "getAddrInfo 没有返回任何地址"

{-# LANGUAGE OverloadedStrings #-}

import Control.Concurrent (forkFinally)
import Control.Concurrent.STM
import Control.Exception (bracket)
import Control.Monad (forever)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import Network.Socket
import System.IO

type Store = TVar (Map.Map T.Text T.Text)

main :: IO ()
main = withSocketsDo $ do
  -- 日志输出到管道或文件时也按行刷新
  hSetBuffering stdout LineBuffering
  store <- newTVarIO Map.empty
  let hints = defaultHints { addrSocketType = Stream, addrFlags = [AI_PASSIVE] }
  addr <- firstAddr <$> getAddrInfo (Just hints) Nothing (Just "4001")
  bracket (openSocket addr) close $ \sock -> do
    setSocketOption sock ReuseAddr 1
    bind sock (addrAddress addr)
    listen sock 1024
    putStrLn "KV 服务器监听 4001 端口"
    forever $ do
      (conn, _) <- accept sock
      -- 把 Socket 转成 Handle，之后按行读写
      h <- socketToHandle conn ReadWriteMode
      hSetBuffering h LineBuffering
      _ <- forkFinally (serve store h) (const $ hClose h)
      return ()

serve :: Store -> Handle -> IO ()
serve store h = do
  eof <- hIsEOF h
  if eof
    then return ()
    else do
      line <- TIO.hGetLine h
      reply <- handleCommand store (T.words line)
      TIO.hPutStrLn h reply
      serve store h

handleCommand :: Store -> [T.Text] -> IO T.Text
handleCommand store cmd = case cmd of
  ["SET", key, value] -> do
    atomically $ modifyTVar' store (Map.insert key value)
    return "OK"
  ["GET", key] -> do
    m <- readTVarIO store
    return $ fromMaybe "(nil)" (Map.lookup key m)
  ["DEL", key] -> do
    atomically $ modifyTVar' store (Map.delete key)
    return "OK"
  _ -> return "ERR unknown command"

firstAddr :: [AddrInfo] -> AddrInfo
firstAddr (addr : _) = addr
firstAddr [] = error "getAddrInfo 没有返回任何地址"