content: {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE FlexibleContexts #-}
  
  {-| Implementation of the Ganeti Unix Domain Socket JSON server interface.
  
  -}
  
  {-
  
  Copyright (C) 2013 Google Inc.
  All rights reserved.
  
  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are
  met:
  
  1. Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
  
  2. Redistributions in binary form must reproduce the above copyright
  notice, this list of conditions and the following disclaimer in the
  documentation and/or other materials provided with the distribution.
  
  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
  TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  -}
  
  module Ganeti.UDSServer
    ( ConnectConfig(..)
    , ServerConfig(..)
    , Client
    , Server
    , RecvResult(..)
    , MsgKeys(..)
    , strOfKey
    -- * Unix sockets
    , openClientSocket
    , closeClientSocket
    , openServerSocket
    , closeServerSocket
    , acceptSocket
    -- * Client and server
    , connectClient
    , connectServer
    , pipeClient
    , acceptClient
    , closeClient
    , clientToFd
    , clientToHandle
    , closeServer
    , buildResponse
    , parseResponse
    , buildCall
    , parseCall
    , recvMsg
    , recvMsgExt
    , sendMsg
    -- * Client handler
    , Handler(..)
    , HandlerResult
    , listener
    ) where
  
  import Control.Concurrent.Lifted (fork, yield)
  import Control.Monad.Base
  import Control.Monad.Trans.Control
  import Control.Monad
  import qualified Data.ByteString as B
  import qualified Data.ByteString.UTF8 as UTF8
  import Data.IORef
  import Data.List
  import Data.Word (Word8)
  import qualified Network.Socket as S
  import System.Directory (removeFile)
  import System.IO ( hClose, hFlush, hPutStr, hWaitForInput, Handle, IOMode(..)
                   , hSetBuffering, BufferMode(..))
  import System.IO.Error (catchIOError, isEOFError)
  import System.Posix.Types (Fd)
  import System.Posix.IO (createPipe, fdToHandle, handleToFd)
  import System.Timeout
  import Text.JSON (encodeStrict, decodeStrict)
  import qualified Text.JSON as J
  import Text.JSON.Types
  
  import Ganeti.BasicTypes
  import Ganeti.Errors (GanetiException(..), ErrorResult)
  import Ganeti.JSON (fromJResult, fromJVal, fromJResultE, fromObj)
  import Ganeti.Logging
  import Ganeti.THH
  import Ganeti.Utils
  import Ganeti.Utils.Time (getCurrentTimeUSec)
  import Ganeti.Constants (privateParametersBlacklist)
  
  -- * Utility functions
  
  -- | Wrapper over System.Timeout.timeout that fails in the IO monad.
  withTimeout :: Int -> String -> IO a -> IO a
  withTimeout secs descr action = do
    result <- timeout (secs * 1000000) action
    case result of
      Nothing -> fail $ "Timeout in " ++ descr
      Just v -> return v
  
  
  -- * Generic protocol functionality
  
  -- | Result of receiving a message from the socket.
  data RecvResult = RecvConnClosed    -- ^ Connection closed
                  | RecvError String  -- ^ Any other error
                  | RecvOk String     -- ^ Successfull receive
                    deriving (Show, Eq)
  
  
  -- | The end-of-message separator.
  eOM :: Word8
  eOM = 3
  
  -- | The end-of-message encoded as a ByteString.
  bEOM :: B.ByteString
  bEOM = B.singleton eOM
  
  -- | Valid keys in the requests and responses.
  data MsgKeys = Method
               | Args
               | Success
               | Result
  
  -- | The serialisation of MsgKeys into strings in messages.
  $(genStrOfKey ''MsgKeys "strOfKey")
  
  
  -- Information required for creating a server connection.
  data ServerConfig = ServerConfig
                      { connPermissions :: FilePermissions
                      , connConfig :: ConnectConfig
                      }
  
  -- Information required for creating a client or server connection.
  data ConnectConfig = ConnectConfig
                       { recvTmo :: Int
                       , sendTmo :: Int
                       }
  
  -- | A client encapsulation. Note that it has separate read and write handle.
  -- For sockets it is the same handle. It is required for bi-directional
  -- inter-process pipes though.
  data Client = Client { rsocket :: Handle          -- ^ The read part of
                                                    -- the client socket
                       , wsocket :: Handle          -- ^ The write part of
                                                    -- the client socket
                       , rbuf :: IORef B.ByteString -- ^ Already received buffer
                       , clientConfig :: ConnectConfig
                       }
  
  -- | A server encapsulation.
  data Server = Server { sSocket :: S.Socket        -- ^ The bound server socket
                       , sPath :: FilePath          -- ^ The scoket's path
                       , serverConfig :: ConnectConfig
                       }
  
  -- * Unix sockets
  
  -- | Creates a Unix socket and connects it to the specified @path@,
  -- where @timeout@ specifies the connection timeout.
  openClientSocket
    :: Int              -- ^ connection timeout
    -> FilePath         -- ^ socket path
    -> IO Handle
  openClientSocket tmo path = do
    sock <- S.socket S.AF_UNIX S.Stream S.defaultProtocol
    withTimeout tmo "creating a connection" $
                S.connect sock (S.SockAddrUnix path)
    S.socketToHandle sock ReadWriteMode
  
  -- | Closes the handle.
  -- Performing the operation on a handle that has already been closed has no
  -- effect; doing so is not an error.
  -- All other operations on a closed handle will fail.
  closeClientSocket :: Handle -> IO ()
  closeClientSocket = hClose
  
  -- | Creates a Unix socket and binds it to the specified @path@.
  openServerSocket :: FilePath -> IO S.Socket
  openServerSocket path = do
    sock <- S.socket S.AF_UNIX S.Stream S.defaultProtocol
    S.bind sock (S.SockAddrUnix path)
    return sock
  
  closeServerSocket :: S.Socket -> FilePath -> IO ()
  closeServerSocket sock path = do
    S.close sock
    removeFile path
  
  acceptSocket :: S.Socket -> IO Handle
  acceptSocket sock = do
    -- ignore client socket address
    (clientSock, _) <- S.accept sock
    S.socketToHandle clientSock ReadWriteMode
  
  -- * Client and server
  
  -- | Connects to the master daemon and returns a Client.
  connectClient
    :: ConnectConfig    -- ^ configuration for the client
    -> Int              -- ^ connection timeout
    -> FilePath         -- ^ socket path
    -> IO Client
  connectClient conf tmo path = do
    h <- openClientSocket tmo path
    rf <- newIORef B.empty
    return Client { rsocket=h, wsocket=h, rbuf=rf, clientConfig=conf }
  
  -- | Creates and returns a server endpoint.
  connectServer :: ServerConfig -> Bool -> FilePath -> IO Server
  connectServer sconf setOwner path = do
    s <- openServerSocket path
    when setOwner $ do
      res <- ensurePermissions path (connPermissions sconf)
      exitIfBad "Error - could not set socket properties" res
  
    S.listen s 5 -- 5 is the max backlog
    return Server { sSocket = s, sPath = path, serverConfig = connConfig sconf }
  
  -- | Creates a new bi-directional client pipe. The two returned clients
  -- talk to each other through the pipe.
  pipeClient :: ConnectConfig -> IO (Client, Client)
  pipeClient conf =
    let newClient r w = do
          rf <- newIORef B.empty
          rh <- fdToHandle r
          wh <- fdToHandle w
          return Client { rsocket = rh, wsocket = wh
                        , rbuf = rf, clientConfig = conf }
    in do
      (r1, w1) <- createPipe
      (r2, w2) <- createPipe
      (,) <$> newClient r1 w2 <*> newClient r2 w1
  
  -- | Closes a server endpoint.
  closeServer :: (MonadBase IO m) => Server -> m ()
  closeServer server =
    liftBase $ closeServerSocket (sSocket server) (sPath server)
  
  -- | Accepts a client
  acceptClient :: Server -> IO Client
  acceptClient s = do
    handle <- acceptSocket (sSocket s)
    new_buffer <- newIORef B.empty
    return Client { rsocket=handle
                  , wsocket=handle
                  , rbuf=new_buffer
                  , clientConfig=serverConfig s
                  }
  
  -- | Closes the client socket.
  -- Performing the operation on a client that has already been closed has no
  -- effect; doing so is not an error.
  -- All other operations on a closed client will fail with an exception.
  closeClient :: Client -> IO ()
  closeClient client = do
    closeClientSocket . wsocket $ client
    closeClientSocket . rsocket $ client
  
  -- | Extracts the read (the first) and the write (the second) file descriptor
  -- of a client. This closes the underlying 'Handle's, therefore the original
  -- client is closed and unusable after the call.
  --
  -- The purpose of this function is to keep the communication channel open,
  -- while replacing a 'Client' with some other means.
  clientToFd :: Client -> IO (Fd, Fd)
  clientToFd client | rh == wh  = join (,) <$> handleToFd rh
                    | otherwise = (,) <$> handleToFd rh <*> handleToFd wh
    where
      rh = rsocket client
      wh = wsocket client
  
  -- | Extracts the read (first) and the write (second) handles of a client.
  -- The purpose of this function is to allow using a client's handles as
  -- input/output streams elsewhere.
  clientToHandle :: Client -> (Handle, Handle)
  clientToHandle client = (rsocket client, wsocket client)
  
  -- | Sends a message over a transport.
  sendMsg :: Client -> String -> IO ()
  sendMsg s buf = withTimeout (sendTmo $ clientConfig s) "sending a message" $ do
    t1 <- getCurrentTimeUSec
    let handle = wsocket s
    -- Allow buffering (up to 1MiB) when writing to the socket. Note that
    -- otherwise we get the default of sending each byte in a separate
    -- system call, resulting in very poor performance.
    hSetBuffering handle (BlockBuffering . Just $ 1024 * 1024)
    hPutStr handle buf
    B.hPut handle bEOM
    hFlush handle
    t2 <- getCurrentTimeUSec
    logDebug $ "sendMsg: " ++ show ((t2 - t1) `div` 1000) ++ "ms"
  
  -- | Given a current buffer and the handle, it will read from the
  -- network until we get a full message, and it will return that
  -- message and the leftover buffer contents.
  recvUpdate :: ConnectConfig -> Handle -> B.ByteString
             -> IO (B.ByteString, B.ByteString)
  recvUpdate conf handle obuf = do
    nbuf <- withTimeout (recvTmo conf) "reading a response" $ do
              _ <- hWaitForInput handle (-1)
              B.hGetNonBlocking handle 4096
    let (msg, remaining) = B.break (eOM ==) nbuf
        newbuf = B.append obuf msg
    if B.null remaining
      then recvUpdate conf handle newbuf
      else return (newbuf, B.copy (B.tail remaining))
  
  -- | Waits for a message over a transport.
  recvMsg :: Client -> IO String
  recvMsg s = do
    cbuf <- readIORef $ rbuf s
    let (imsg, ibuf) = B.break (eOM ==) cbuf
    (msg, nbuf) <-
      if B.null ibuf      -- if old buffer didn't contain a full message
                          -- then we read from network:
        then recvUpdate (clientConfig s) (rsocket s) cbuf
        -- else we return data from our buffer, copying it so that the whole
        -- message isn't retained and can be garbage collected
        else return (imsg, B.copy (B.tail ibuf))
    writeIORef (rbuf s) nbuf
    return $ UTF8.toString msg
  
  -- | Extended wrapper over recvMsg.
  recvMsgExt :: Client -> IO RecvResult
  recvMsgExt s =
    catchIOError (liftM RecvOk (recvMsg s)) $ \e ->
      return $ if isEOFError e
                 then RecvConnClosed
                 else RecvError (show e)
  
  
  -- | Serialize a request to String.
  buildCall :: (J.JSON mth, J.JSON args)
            => mth    -- ^ The method
            -> args   -- ^ The arguments
            -> String -- ^ The serialized form
  buildCall mth args =
    let keyToObj :: (J.JSON a) => MsgKeys -> a -> (String, J.JSValue)
        keyToObj k v = (strOfKey k, J.showJSON v)
    in encodeStrict $ toJSObject [ keyToObj Method mth, keyToObj Args args ]
  
  -- | Parse the required keys out of a call.
  parseCall :: (J.JSON mth, J.JSON args) => String -> Result (mth, args)
  parseCall s = do
    arr <- fromJResult "parsing top-level JSON message" $
             decodeStrict s :: Result (JSObject JSValue)
    let keyFromObj :: (J.JSON a) => MsgKeys -> Result a
        keyFromObj = fromObj (fromJSObject arr) . strOfKey
    (,) <$> keyFromObj Method <*> keyFromObj Args
  
  
  -- | Serialize the response to String.
  buildResponse :: Bool    -- ^ Success
                -> JSValue -- ^ The arguments
                -> String  -- ^ The serialized form
  buildResponse success args =
    let ja = [ (strOfKey Success, JSBool success)
             , (strOfKey Result, args)]
        jo = toJSObject ja
    in encodeStrict jo
  
  -- | Try to decode an error from the server response. This function
  -- will always fail, since it's called only on the error path (when
  -- status is False).
  decodeError :: JSValue -> ErrorResult JSValue
  decodeError val =
    case fromJVal val of
      Ok e -> Bad e
      Bad msg -> Bad $ GenericError msg
  
  -- | Check that luxi responses contain the required keys and that the
  -- call was successful.
  parseResponse :: String -> ErrorResult JSValue
  parseResponse s = do
    when (UTF8.replacement_char `elem` s) $
        failError "Failed to decode UTF-8,\
                  \ detected replacement char after decoding"
    oarr <- fromJResultE "Parsing LUXI response" (decodeStrict s)
    let arr = J.fromJSObject oarr
    status <- fromObj arr (strOfKey Success)
    result <- fromObj arr (strOfKey Result)
    if status
      then return result
      else decodeError result
  
  -- | Logs an outgoing message.
  logMsg
      :: (Show e, J.JSON e, MonadLog m)
      => Handler i m o
      -> i                          -- ^ the received request (used for logging)
      -> GenericResult e J.JSValue  -- ^ A message to be sent
      -> m ()
  logMsg handler req (Bad err) =
    logWarning $ "Failed to execute request " ++ hInputLogLong handler req ++ ": "
                 ++ show err
  logMsg handler req (Ok result) = do
    -- only log the first 2,000 chars of the result
    logDebug $ "Result (truncated): " ++ take 2000 (J.encode result)
    logDebug $ "Successfully handled " ++ hInputLogShort handler req
  
  -- | Prepares an outgoing message.
  prepareMsg
      :: (J.JSON e)
      => GenericResult e J.JSValue  -- ^ A message to be sent
      -> (Bool, J.JSValue)
  prepareMsg (Bad err)   = (False, J.showJSON err)
  prepareMsg (Ok result) = (True, result)
  
  
  -- * Processing client requests
  
  type HandlerResult m o = m (Bool, GenericResult GanetiException o)
  
  data Handler i m o = Handler
    { hParse         :: J.JSValue -> J.JSValue -> Result i
      -- ^ parses method and its arguments into the input type
    , hInputLogShort :: i -> String
      -- ^ short description of an input, for the INFO logging level
    , hInputLogLong  :: i -> String
      -- ^ long description of an input, for the DEBUG logging level
    , hExec          :: i -> HandlerResult m o
      -- ^ executes the handler on an input
    }
  
  
  handleJsonMessage
      :: (J.JSON o, Monad m)
      => Handler i m o            -- ^ handler
      -> i                        -- ^ parsed input
      -> HandlerResult m J.JSValue
  handleJsonMessage handler req = do
    (close, call_result) <- hExec handler req
    return (close, fmap J.showJSON call_result)
  
  -- | Takes a request as a 'String', parses it, passes it to a handler and
  -- formats its response.
  handleRawMessage
      :: (J.JSON o, MonadLog m)
      => Handler i m o            -- ^ handler
      -> String                   -- ^ raw unparsed input
      -> m (Bool, String)
  handleRawMessage handler payload =
    case parseCall payload >>= uncurry (hParse handler) of
      Bad err -> do
           let errmsg = "Failed to parse request: " ++ err
           logWarning errmsg
           return (False, buildResponse False (J.showJSON errmsg))
      Ok req -> do
          logDebug $ "Request: " ++ hInputLogLong handler req
          (close, call_result_json) <- handleJsonMessage handler req
          logMsg handler req call_result_json
          let (status, response) = prepareMsg call_result_json
          return (close, buildResponse status response)
  
  isRisky :: RecvResult -> Bool
  isRisky msg = case msg of
    RecvOk payload -> any (`isInfixOf` payload) privateParametersBlacklist
    _ -> False
  
  -- | Reads a request, passes it to a handler and sends a response back to the
  -- client.
  handleClient
      :: (J.JSON o, MonadBase IO m, MonadLog m)
      => Handler i m o
      -> Client
      -> m Bool
  handleClient handler client = do
    msg <- liftBase $ recvMsgExt client
  
    debugMode <- liftBase isDebugMode
    when (debugMode && isRisky msg) $
      logAlert "POSSIBLE LEAKING OF CONFIDENTIAL PARAMETERS. \
               \Daemon is running in debug mode. \
               \The text of the request has been logged."
    logDebug $ "Received message (truncated): " ++ take 500 (show msg)
  
    case msg of
      RecvConnClosed -> logDebug "Connection closed" >>
                        return False
      RecvError err -> logWarning ("Error during message receiving: " ++ err) >>
                       return False
      RecvOk payload -> do
        (close, outMsg) <- handleRawMessage handler payload
        liftBase $ sendMsg client outMsg
        return close
  
  
  -- | Main client loop: runs one loop of 'handleClient', and if that
  -- doesn't report a finished (closed) connection, restarts itself.
  clientLoop
      :: (J.JSON o, MonadBase IO m, MonadLog m)
      => Handler i m o
      -> Client
      -> m ()
  clientLoop handler client = do
    result <- handleClient handler client
    {- It's been observed sometimes that reading immediately after sending
       a response leads to worse performance, as there is nothing to read and
       the system calls are just wasted. Thus yielding before reading gives
       other threads a chance to proceed and provides a natural pause, leading
       to a bit more efficient communication.
    -}
    if result
      then yield >> clientLoop handler client
      else liftBase $ closeClient client
  
  -- | Main listener loop: accepts clients, forks an I/O thread to handle
  -- that client.
  listener
      :: (J.JSON o, MonadBaseControl IO m, MonadLog m)
      => Handler i m o
      -> Server
      -> m ()
  listener handler server = do
    client <- liftBase $ acceptClient server
    _ <- fork $ clientLoop handler client
    return ()
  