{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE DataKinds #-}
{-|
Module      : Netw.Ancillary
Description : Ancillary data
Portabilty  : Linux

This module contains ancillary data types. Ancillary data can be send and received via the sendmsg and recvmsg functions.
-}
module Netw.Ancillary where
import Netw.Internal.Ancillary
import qualified Netw.Internal.Cmsg as I

import Control.Monad.Primitive
import Control.Monad
import Control.Exception

import Data.Dynamic
import Data.Foldable
import Data.Function
import Data.Primitive hiding (indexByteArray, writeByteArray, readByteArray)
import Data.Primitive.ByteArray.Unaligned

import System.Posix.Types

import GHC.TypeLits

class (Typeable a, KnownNatPair (Anci a)) => Ancillary a where
  -- | This type family enforces an unique mapping between cmsg_level and cmsg_type pair and the ancillary data type
  type Anci a = (r :: (Nat, Nat)) | r -> a
  -- ^ the order is (cmsg_level, cmsg_type)

  -- | Get the size of the cmsg payload (just cmsg_data[])
  cmsgDataSize :: a -> Int

  -- | Write the payload onto the data section. Take the offset of the data section.
  cmsgPokeData :: PrimMonad m => a -> MutableByteArray (PrimState m) -> Int -> m ()
  -- ^ This function assumes enough space has been allocated

  -- | Read the payload from the data section. Take the offset and size of the payload section.
  cmsgPeekData :: ByteArray -> Int -> Int -> a

cmsgLevel :: forall a. Ancillary a => I.CmsghdrLevel
cmsgLevel :: forall a. Ancillary a => CmsghdrLevel
cmsgLevel = Integer -> CmsghdrLevel
forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (p :: (Nat, Nat)). KnownNatPair p => Integer
fstNatVal @(Anci a))

cmsgType :: forall a. Ancillary a => I.CmsghdrType
cmsgType :: forall a. Ancillary a => CmsghdrLevel
cmsgType = Integer -> CmsghdrLevel
forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (p :: (Nat, Nat)). KnownNatPair p => Integer
sndNatVal @(Anci a))

-- | Find the space a control message element occupies in cmsg buffer given the size of its payload section.
cmsgSpace :: Ancillary a => a -> Int
cmsgSpace :: forall a. Ancillary a => a -> Int
cmsgSpace = Int -> Int
I.cmsgSpace (Int -> Int) -> (a -> Int) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int
forall a. Ancillary a => a -> Int
cmsgDataSize

-- | Read a control message element, assuming type a. Return Nothing if a is not the correct type.
peekAncillary :: forall a. Ancillary a => ByteArray -> Int -> Maybe a
peekAncillary :: forall a. Ancillary a => ByteArray -> Int -> Maybe a
peekAncillary ByteArray
cmsgs Int
hdroffs = ByteArray
-> Int
-> (CmsghdrLevel
    -> CmsghdrLevel -> ByteArray -> Int -> Int -> Maybe a)
-> Maybe a
forall a.
ByteArray
-> Int
-> (CmsghdrLevel -> CmsghdrLevel -> ByteArray -> Int -> Int -> a)
-> a
I.peekCmsg ByteArray
cmsgs Int
hdroffs CmsghdrLevel -> CmsghdrLevel -> ByteArray -> Int -> Int -> Maybe a
reader
  where reader :: CmsghdrLevel -> CmsghdrLevel -> ByteArray -> Int -> Int -> Maybe a
reader CmsghdrLevel
l CmsghdrLevel
t ByteArray
b Int
o Int
s
          | CmsghdrLevel
l CmsghdrLevel -> CmsghdrLevel -> Bool
forall a. Eq a => a -> a -> Bool
== forall a. Ancillary a => CmsghdrLevel
cmsgLevel @a Bool -> Bool -> Bool
&& CmsghdrLevel
t CmsghdrLevel -> CmsghdrLevel -> Bool
forall a. Eq a => a -> a -> Bool
== forall a. Ancillary a => CmsghdrLevel
cmsgType @a =
            a -> Maybe a
forall a. a -> Maybe a
Just (ByteArray -> Int -> Int -> a
forall a. Ancillary a => ByteArray -> Int -> Int -> a
cmsgPeekData ByteArray
b Int
o Int
s)
          | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing
-- | Store a control message element at offset in buffer
pokeAncillary :: forall a m. (Ancillary a, PrimMonad m) => MutableByteArray (PrimState m) -> Int -> a -> m ()
pokeAncillary :: forall a (m :: * -> *).
(Ancillary a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
pokeAncillary MutableByteArray (PrimState m)
cmsgs Int
hdroffs a
andat = MutableByteArray (PrimState m)
-> Int
-> CmsghdrLevel
-> CmsghdrLevel
-> Int
-> (MutableByteArray (PrimState m) -> Int -> m ())
-> m ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int
-> CmsghdrLevel
-> CmsghdrLevel
-> Int
-> (MutableByteArray (PrimState m) -> Int -> m ())
-> m ()
I.pokeCmsg MutableByteArray (PrimState m)
cmsgs Int
hdroffs CmsghdrLevel
l CmsghdrLevel
t (a -> Int
forall a. Ancillary a => a -> Int
cmsgDataSize a
andat) (a -> MutableByteArray (PrimState m) -> Int -> m ()
forall a (m :: * -> *).
(Ancillary a, PrimMonad m) =>
a -> MutableByteArray (PrimState m) -> Int -> m ()
forall (m :: * -> *).
PrimMonad m =>
a -> MutableByteArray (PrimState m) -> Int -> m ()
cmsgPokeData a
andat)
  where l :: CmsghdrLevel
l = forall a. Ancillary a => CmsghdrLevel
cmsgLevel @a
        t :: CmsghdrLevel
t = forall a. Ancillary a => CmsghdrLevel
cmsgType  @a
-- ^ Warning: The buffer must be big enough to store the control message element

-- | This data type holds an abitrary ancillary data type.
data AncillaryData = AncillaryData
  { AncillaryData -> Int
ancillarySpace :: !Int
  , AncillaryData -> forall a. Ancillary a => Maybe a
ancillaryData  :: !(forall a. Ancillary a => Maybe a)
  , AncillaryData
-> forall (m :: * -> *).
   PrimMonad m =>
   MutableByteArray (PrimState m) -> Int -> m ()
ancillaryWrite :: !(forall m. PrimMonad m => MutableByteArray (PrimState m) -> Int -> m ())
  }

mkAncillaryData :: Ancillary a => a -> AncillaryData
mkAncillaryData :: forall a. Ancillary a => a -> AncillaryData
mkAncillaryData a
a = AncillaryData { ancillarySpace :: Int
ancillarySpace = a -> Int
forall a. Ancillary a => a -> Int
cmsgSpace a
a
                                  , ancillaryData :: forall a. Ancillary a => Maybe a
ancillaryData  = Dynamic -> Maybe a
forall a. Typeable a => Dynamic -> Maybe a
fromDynamic (a -> Dynamic
forall a. Typeable a => a -> Dynamic
toDyn a
a)
                                  , ancillaryWrite :: forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> m ()
ancillaryWrite = \ MutableByteArray (PrimState m)
buffer Int
offset -> MutableByteArray (PrimState m) -> Int -> a -> m ()
forall a (m :: * -> *).
(Ancillary a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
pokeAncillary MutableByteArray (PrimState m)
buffer Int
offset a
a
                                  }
-- | ancillaryWrite but without the impredicative type
writeAncillary :: PrimMonad m => AncillaryData -> MutableByteArray (PrimState m) -> Int -> m ()
writeAncillary :: forall (m :: * -> *).
PrimMonad m =>
AncillaryData -> MutableByteArray (PrimState m) -> Int -> m ()
writeAncillary AncillaryData
a = AncillaryData
-> forall (m :: * -> *).
   PrimMonad m =>
   MutableByteArray (PrimState m) -> Int -> m ()
ancillaryWrite AncillaryData
a

-- | ancillaryData but without the impredicative type
recoverAncillary :: Ancillary a => AncillaryData -> Maybe a
recoverAncillary :: forall a. Ancillary a => AncillaryData -> Maybe a
recoverAncillary AncillaryData
a = AncillaryData -> forall a. Ancillary a => Maybe a
ancillaryData AncillaryData
a

{-# INLINE encodeAncillaryData #-}
-- | Encode a list of control message elements
encodeAncillaryData :: [AncillaryData] -> ByteArray
encodeAncillaryData :: [AncillaryData] -> ByteArray
encodeAncillaryData [AncillaryData]
cmsgs = (forall s. ST s (MutableByteArray s)) -> ByteArray
runByteArray ((forall s. ST s (MutableByteArray s)) -> ByteArray)
-> (forall s. ST s (MutableByteArray s)) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
buffer <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
size
  (AncillaryData -> Int -> ST s ())
-> [AncillaryData] -> [Int] -> ST s ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (AncillaryData
-> MutableByteArray (PrimState (ST s)) -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
AncillaryData -> MutableByteArray (PrimState m) -> Int -> m ()
`writeAncillary` MutableByteArray s
MutableByteArray (PrimState (ST s))
buffer) [AncillaryData]
cmsgs [Int]
offsets
  MutableByteArray s -> ST s (MutableByteArray s)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MutableByteArray s
buffer
  where ([Int]
offsets, Int
size) = (\ [Int]
a -> ([Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
init [Int]
a, [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
a)) ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ((AncillaryData -> Int) -> [AncillaryData] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map AncillaryData -> Int
ancillarySpace [AncillaryData]
cmsgs))
-- ^ Used internally

{-# INLINE decodeAncillaryData #-}
-- | Decode a buffer of control message elements. NOTE: The bytearray must be shrunk to `msg_controllen` bytes before being passed to this function.
decodeAncillaryData :: ByteArray -> [AncillaryData]
decodeAncillaryData :: ByteArray -> [AncillaryData]
decodeAncillaryData ByteArray
cmsgs = ((Int -> [AncillaryData]) -> Int -> [AncillaryData])
-> Int -> [AncillaryData]
forall a. (a -> a) -> a
fix (\ Int -> [AncillaryData]
as Int
offs ->
  if Int
offs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
size then
    let space :: Int
space = ByteArray
-> Int
-> (CmsghdrLevel -> CmsghdrLevel -> ByteArray -> Int -> Int -> Int)
-> Int
forall a.
ByteArray
-> Int
-> (CmsghdrLevel -> CmsghdrLevel -> ByteArray -> Int -> Int -> a)
-> a
I.peekCmsg ByteArray
cmsgs Int
offs (\ CmsghdrLevel
_ CmsghdrLevel
_ ByteArray
_ Int
_ -> Int -> Int
I.cmsgSpace)
        andat :: AncillaryData
andat = AncillaryData { ancillarySpace :: Int
ancillarySpace = Int
space 
                              , ancillaryData :: forall a. Ancillary a => Maybe a
ancillaryData  = ByteArray -> Int -> Maybe a
forall a. Ancillary a => ByteArray -> Int -> Maybe a
peekAncillary ByteArray
cmsgs Int
offs
                              , ancillaryWrite :: forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> m ()
ancillaryWrite = \ MutableByteArray (PrimState m)
dest Int
destOffs -> MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray (PrimState m)
dest Int
destOffs ByteArray
cmsgs Int
offs Int
space
                              }
    in  AncillaryData
andatAncillaryData -> [AncillaryData] -> [AncillaryData]
forall a. a -> [a] -> [a]
:Int -> [AncillaryData]
as (Int
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
space)
  else []) Int
0
  where size :: Int
size = ByteArray -> Int
sizeofByteArray ByteArray
cmsgs
-- ^ Used internally. Every AncillaryData element holds a reference to the bytearray and keeps it alive.

-- | Transfer file descriptors between sockets
newtype ScmRights = ScmRights [Fd]
instance Ancillary ScmRights where
  type Anci ScmRights = UNIX_CONTROL_MESSAGE_RIGHTS
  cmsgDataSize :: ScmRights -> Int
cmsgDataSize (ScmRights [Fd]
fds) = forall a. Prim a => Int
sizeOfType @Fd Int -> Int -> Int
forall a. Num a => a -> a -> a
* [Fd] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Fd]
fds
  cmsgPokeData :: forall (m :: * -> *).
PrimMonad m =>
ScmRights -> MutableByteArray (PrimState m) -> Int -> m ()
cmsgPokeData (ScmRights [Fd]
fds) MutableByteArray (PrimState m)
buffer Int
initOffs = m Int -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void ((Int -> Fd -> m Int) -> Int -> [Fd] -> m Int
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM (\ Int
offs Fd
fd -> MutableByteArray (PrimState m) -> Int -> Fd -> m ()
forall (m :: * -> *) a.
(PrimMonad m, PrimUnaligned a) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeUnalignedByteArray MutableByteArray (PrimState m)
buffer Int
offs Fd
fd m () -> m Int -> m Int
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ forall a. Prim a => Int
sizeOfType @Fd)) Int
initOffs [Fd]
fds)
  cmsgPeekData :: ByteArray -> Int -> Int -> ScmRights
cmsgPeekData ByteArray
buffer Int
o Int
s = Bool -> ScmRights -> ScmRights
forall a. HasCallStack => Bool -> a -> a
assert (Int
s Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` forall a. Prim a => Int
sizeOfType @Fd Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (ScmRights -> ScmRights) -> ScmRights -> ScmRights
forall a b. (a -> b) -> a -> b
$ [Fd] -> ScmRights
ScmRights ([Fd] -> ScmRights) -> [Fd] -> ScmRights
forall a b. (a -> b) -> a -> b
$ ((Int -> Int -> [Fd]) -> Int -> Int -> [Fd]) -> Int -> Int -> [Fd]
forall a. (a -> a) -> a
fix (\ Int -> Int -> [Fd]
as Int
offs Int
size ->
    if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then
      ByteArray -> Int -> Fd
forall a. PrimUnaligned a => ByteArray -> Int -> a
indexUnalignedByteArray ByteArray
buffer Int
offsFd -> [Fd] -> [Fd]
forall a. a -> [a] -> [a]
:Int -> Int -> [Fd]
as (Int
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ forall a. Prim a => Int
sizeOfType @Fd) (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- forall a. Prim a => Int
sizeOfType @Fd)
    else
      []) Int
o Int
s