Pty.hs

{-# LANGUAGE CApiFFI #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE Trustworthy #-}
-------------------------------------------------------------------------------
-- |
-- Module      :  System.Posix.Pty
-- Copyright   :  (C) 2013 Merijn Verstraaten
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Merijn Verstraaten <merijn@inconsistent.nl>
-- Stability   :  experimental
-- Portability :  haha
--
-- A module for interacting with subprocesses through a pseudo terminal (pty).
-- Provides functions for reading from, writing to and resizing pseudo
-- terminals. Re-exports most of "System.Posix.Terminal", providing wrappers
-- that work with the 'Pty' type where necessary.
-------------------------------------------------------------------------------
module System.Posix.Pty (
    -- * Subprocess Creation
      spawnWithPty
    -- * Data Structures
    , Pty
    , PtyControlCode (..)
    -- * Pty Interaction Functions
    , createPty
    , tryReadPty
    , readPty
    , writePty
    , resizePty
    , ptyDimensions
    -- * Re-exports of "System.Posix.Terminal"
    -- $posix-reexport
    , getTerminalAttributes
    , setTerminalAttributes
    , sendBreak
    , drainOutput
    , discardData
    , controlFlow
    , getTerminalProcessGroupID
    , getTerminalName
    , getSlaveTerminalName
    , module System.Posix.Terminal
    ) where

import Control.Applicative
import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS

import Foreign
import Foreign.C.String (CString, newCString, peekCString)
import Foreign.C.Types

import System.IO (Handle)
import System.IO.Error (mkIOError, eofErrorType)
import System.Posix.IO.ByteString (fdToHandle)
import System.Posix.Types

import qualified System.Posix.Terminal as T
import System.Posix.Terminal hiding
        ( getTerminalAttributes
        , setTerminalAttributes
        , sendBreak
        , drainOutput
        , discardData
        , controlFlow
        , getTerminalProcessGroupID
        , setTerminalProcessGroupID
        , queryTerminal
        , getTerminalName
        , openPseudoTerminal
        , getSlaveTerminalName)

-- | Abstract pseudo terminal type.
data Pty = Pty !Fd !Handle

-- | Pseudo terminal control information.
--
-- [Terminal read queue] The terminal read queue contains the data that was
-- written from the master terminal to the slave terminal, which was not read
-- from the slave yet.
--
-- [Terminal write queue] The terminal write queue contains the data that was
-- written from the slave terminal, which was not sent to the master yet.
data PtyControlCode = FlushRead     -- ^ Terminal read queue was flushed.
                    | FlushWrite    -- ^ Terminal write queue was flushed.
                    | OutputStopped -- ^ Terminal output was stopped.
                    | OutputStarted -- ^ Terminal output was restarted.
                    | DoStop        -- ^ Terminal stop and start characters are
                                    --   @^S@ and @^Q@ respectively.
                    | NoStop        -- ^ Terminal stop and start characters are
                                    --   NOT @^S@ and @^Q@.
                    deriving (Eq, Read, Show)

-- | Produces a 'Pty' if the file descriptor is associated with a terminal and
-- Nothing if not.
createPty :: Fd -> IO (Maybe Pty)
createPty fd = do
    isTerm <- T.queryTerminal fd
    if isTerm
       then Just . Pty fd <$> fdToHandle fd
       else return Nothing

-- | Attempt to read data from a pseudo terminal. Produces either the data read
-- or a list of 'PtyControlCode'@s@ indicating which control status events that
-- have happened on the slave terminal.
--
-- Throws an 'IOError' of type 'eofErrorType' when the terminal has been
-- closed, for example when the subprocess has terminated.
tryReadPty :: Pty -> IO (Either [PtyControlCode] ByteString)
tryReadPty (Pty _ hnd) = do
    result <- BS.hGetSome hnd 1024
    case BS.uncons result of
         Nothing -> ioError ptyClosed
         Just (byte, rest)
            | byte == 0    -> return (Right rest)
            | BS.null rest -> return $ Left (byteToControlCode byte)
            | otherwise    -> ioError can'tHappen
  where
    ptyClosed = mkIOError eofErrorType "pty terminated" Nothing Nothing
    can'tHappen = userError "Uh-oh! Something different went horribly wrong!"

-- | The same as 'tryReadPty', but discards any control status events.
readPty :: Pty -> IO ByteString
readPty pty = tryReadPty pty >>= \case
                   Left _ -> readPty pty
                   Right bs -> return bs

-- | Write a 'ByteString' to the pseudo terminal, throws an 'IOError' when the
-- terminal has been closed, for example when the subprocess has terminated.
writePty :: Pty -> ByteString -> IO ()
writePty (Pty _ hnd) = BS.hPut hnd

-- | Set the pseudo terminal's dimensions to the specified width and height.
resizePty :: Pty -> (Int, Int) -> IO ()
resizePty (Pty fd _) (x, y) =
    set_pty_size fd x y >>= throwCErrorOnMinus1 "unable to set pty dimensions"

-- | Produces the pseudo terminal's current dimensions.
ptyDimensions :: Pty -> IO (Int, Int)
ptyDimensions (Pty fd _) = alloca $ \x -> alloca $ \y -> do
    get_pty_size fd x y >>= throwCErrorOnMinus1 "unable to get pty size"
    (,) <$> peek x <*> peek y

-- | Create a new process that is connected to the current process through a
-- pseudo terminal. If an environment is specified, then only the specified
-- environment variables will be set. If no environment is specified the
-- process will inherit its environment from the current process. Example:
--
-- > pty <- spawnWithPty (Just [("SHELL", "tcsh")]) True "ls" ["-l"] (20, 10)
--
-- This searches the user's PATH for a binary called @ls@, then runs this
-- binary with the commandline argument @-l@ in a terminal that is 20
-- characters wide and 10 characters high. The environment of @ls@ will
-- contains one variable, SHELL, which will be set to the value \"tcsh\".
spawnWithPty :: Maybe [(String, String)]    -- ^ Optional environment for the
                                            --   new process.
             -> Bool                        -- ^ Search for the executable in
                                            --   PATH?
             -> FilePath                    -- ^ Program's name.
             -> [String]                    -- ^ Command line arguments for the
                                            --   program.
             -> (Int, Int)                  -- ^ Initial dimensions for the
                                            --   pseudo terminal.
             -> IO Pty
spawnWithPty env' search path' argv' (x, y) = do
    path <- newCString path'
    argv <- mapM newCString argv'
    env <- maybe (return []) (mapM fuse) env'

    result <- forkExecWithPty x y path (fromBool search) argv env

    mapM_ free (env ++ argv)
    free path

    throwCErrorOnMinus1 "unable to fork or open new pty" result

    hnd <- fdToHandle result
    return (Pty result hnd)
  where
    fuse (key, val) = newCString (key ++ "=" ++ val)

-- Module internal functions

getFd :: Pty -> Fd
getFd (Pty fd _) = fd

throwCErrorOnMinus1 :: (Eq a, Num a) => String -> a -> IO ()
throwCErrorOnMinus1 s i = when (i == -1) $ do
    errnoMsg <- errno >>= peekCString . strerror
    ioError . userError $ s ++ ": " ++ errnoMsg

forkExecWithPty :: Int
                -> Int
                -> CString
                -> CInt
                -> [CString]
                -> [CString]
                -> IO Fd
forkExecWithPty x y path search argv' env' = do
    argv <- newArray0 nullPtr (path:argv')
    env <- case env' of
                [] -> return nullPtr
                _  -> newArray0 nullPtr env'

    result <- fork_exec_with_pty x y search path argv env
    free argv >> free env
    return result

byteToControlCode :: Word8 -> [PtyControlCode]
byteToControlCode i = map snd $ filter ((/=0) . (.&.i) . fst) codeMapping
    where codeMapping :: [(Word8, PtyControlCode)]
          codeMapping =
            [ (tiocPktFlushRead,  FlushRead)
            , (tiocPktFlushWrite, FlushWrite)
            , (tiocPktStop,       OutputStopped)
            , (tiocPktStart,      OutputStarted)
            , (tiocPktDoStop,     DoStop)
            , (tiocPktNoStop,     NoStop)
            ]

-- Foreign imports

foreign import capi unsafe "sys/termios.h value TIOCPKT_FLUSHREAD"
    tiocPktFlushRead :: Word8
foreign import capi unsafe "sys/termios.h value TIOCPKT_FLUSHWRITE"
    tiocPktFlushWrite :: Word8
foreign import capi unsafe "sys/termios.h value TIOCPKT_STOP"
    tiocPktStop :: Word8
foreign import capi unsafe "sys/termios.h value TIOCPKT_START"
    tiocPktStart :: Word8
foreign import capi unsafe "sys/termios.h value TIOCPKT_DOSTOP"
    tiocPktDoStop :: Word8
foreign import capi unsafe "sys/termios.h value TIOCPKT_NOSTOP"
    tiocPktNoStop :: Word8

foreign import ccall unsafe "errno.h"
    errno :: IO CInt

foreign import ccall unsafe "string.h"
    strerror :: CInt -> CString

foreign import ccall "pty_size.h"
    set_pty_size :: Fd -> Int -> Int -> IO CInt

foreign import ccall "pty_size.h"
    get_pty_size :: Fd -> Ptr Int -> Ptr Int -> IO CInt

foreign import ccall "fork_exec_with_pty.h"
    fork_exec_with_pty :: Int
                       -> Int
                       -> CInt
                       -> CString
                       -> Ptr CString
                       -> Ptr CString
                       -> IO Fd

-- Pty specialised re-exports of System.Posix.Terminal

{- $posix-reexport
This module re-exports the entirety of "System.Posix.Terminal", with the
exception of the following functions:

[setTerminalProcessGroupID] This function can't be used after a process using
the slave terminal has been created, rendering it mostly useless for working
with 'Pty'@s@ created by this module.

[queryTerminal] Useless, 'Pty' is always a terminal.

[openPseudoTerminal] Only useful for the kind of tasks this module is supposed
abstract away.

In addition, some functions from "System.Posix.Terminal" work directly with
'Fd'@s@, these have been hidden and instead the following replacements working
on 'Pty'@s@ are exported.
-}

-- | See 'System.Posix.Terminal.getTerminalAttributes'.
getTerminalAttributes :: Pty -> IO TerminalAttributes
getTerminalAttributes = T.getTerminalAttributes . getFd

-- | See 'System.Posix.Terminal.setTerminalAttributes'.
setTerminalAttributes :: Pty -> TerminalAttributes -> TerminalState -> IO ()
setTerminalAttributes = T.setTerminalAttributes . getFd

-- | See 'System.Posix.Terminal.sendBreak'.
sendBreak :: Pty -> Int -> IO ()
sendBreak = T.sendBreak . getFd

-- | See 'System.Posix.Terminal.drainOutput'.
drainOutput :: Pty -> IO ()
drainOutput = T.drainOutput . getFd

-- | See 'System.Posix.Terminal.discardData'.
discardData :: Pty -> QueueSelector -> IO ()
discardData = T.discardData . getFd

-- | See 'System.Posix.Terminal.controlFlow'.
controlFlow :: Pty -> FlowAction -> IO ()
controlFlow = T.controlFlow . getFd

-- | See 'System.Posix.Terminal.getTerminalProcessGroupID'.
getTerminalProcessGroupID :: Pty -> IO ProcessGroupID
getTerminalProcessGroupID = T.getTerminalProcessGroupID . getFd

-- | See 'System.Posix.Terminal.getTerminalName'.
getTerminalName :: Pty -> IO FilePath
getTerminalName = T.getTerminalName . getFd

-- | See 'System.Posix.Terminal.getSlaveTerminalName'.
getSlaveTerminalName :: Pty -> IO FilePath
getSlaveTerminalName = T.getSlaveTerminalName . getFd