Wire.hs

{-# LANGUAGE OverloadedStrings #-}

module YMonad.Protocol.Wayland.Wire (
  encodeWireMessage,
  encodeWireBatch,
  decodeRawFrames,
  decodeMessageArgs,
) where

import Data.Binary.Get
import Data.Binary.Put
import Data.Bits ((.&.))
import Data.ByteString qualified as BS

import YMonad.Protocol.Wayland.Types

encodeWireMessage :: WireMessage -> BS.ByteString
encodeWireMessage msg = toStrict $ runPut $ do
  let payload = toStrict (runPut (mapM_ putArgValue (wireArgs msg)))
      size = fromIntegral (8 + BS.length payload) :: Word16
  putWord32le (wireObjectId msg)
  putWord16le (wireOpcode msg)
  putWord16le size
  putByteString payload

encodeWireBatch :: [WireMessage] -> BS.ByteString
encodeWireBatch = BS.concat . map encodeWireMessage

decodeRawFrames :: BS.ByteString -> Either String ([RawFrame], BS.ByteString)
decodeRawFrames input = go input id
  where
    go bytes acc
      | BS.length bytes < 8 = Right (acc [], bytes)
      | otherwise =
          let header = runGet getHeader (fromStrict (BS.take 8 bytes))
              payloadLen = fromIntegral (frameSize header) - 8
           in if frameSize header < 8
                then Left "Wayland frame shorter than header"
                else
                  if BS.length bytes < fromIntegral (frameSize header)
                    then Right (acc [], bytes)
                    else
                      let payload = BS.take payloadLen (BS.drop 8 bytes)
                          rest = BS.drop (fromIntegral (frameSize header)) bytes
                          frame = RawFrame (frameObjectId header) (frameOpcode header) payload
                       in go rest (acc . (frame :))

    getHeader = do
      objectId <- getWord32le
      opcode <- getWord16le
      Header objectId opcode <$> getWord16le

decodeMessageArgs :: MessageSpec -> BS.ByteString -> Either String [ArgValue]
decodeMessageArgs spec payload =
  case runGetOrFail (mapM getArgValue (messageSpecArgs spec)) (fromStrict payload) of
    Left (_, _, err) -> Left err
    Right (_, _, args) -> Right args

putArgValue :: ArgValue -> Put
putArgValue arg = case arg of
  ArgIntValue n -> putInt32le n
  ArgUintValue n -> putWord32le n
  ArgObjectValue Nothing -> putWord32le 0
  ArgObjectValue (Just n) -> putWord32le n
  ArgNewIdValue n -> putWord32le n
  ArgStringValue Nothing -> putWord32le 0
  ArgStringValue (Just txt) -> do
    let bytes = encodeUtf8 txt
        total = BS.length bytes + 1
    putWord32le (fromIntegral total)
    putByteString bytes
    putWord8 0
    putPadding total

getArgValue :: ArgSpec -> Get ArgValue
getArgValue spec = case argSpecType spec of
  ArgInt -> ArgIntValue <$> getInt32le
  ArgUint -> ArgUintValue <$> getWord32le
  ArgObject _ -> do
    oid <- getWord32le
    pure $ ArgObjectValue (if oid == 0 then Nothing else Just oid)
  ArgNewId _ -> ArgNewIdValue <$> getWord32le
  ArgString -> do
    len <- getWord32le
    if len == 0
      then pure (ArgStringValue Nothing)
      else do
        let total = fromIntegral len
            contentLen = max 0 (total - 1)
        bytes <- getByteString contentLen
        _nul <- getWord8
        skipPadding total
        case decodeUtf8' bytes of
          Left err -> fail (show err)
          Right txt -> pure (ArgStringValue (Just txt))

putPadding :: Int -> Put
putPadding payloadLen = replicateM_ (padLength payloadLen) (putWord8 0)

skipPadding :: Int -> Get ()
skipPadding payloadLen = do
  _ <- getByteString (padLength payloadLen)
  pass

padLength :: Int -> Int
padLength n = (4 - (n .&. 3)) .&. 3

data Header = Header !Word32 !Word16 !Word16

frameObjectId :: Header -> Word32
frameObjectId (Header objectId _ _) = objectId

frameOpcode :: Header -> Word16
frameOpcode (Header _ opcode _) = opcode

frameSize :: Header -> Word16
frameSize (Header _ _ size) = size