packages feed

morley-client-0.2.1: test/Test/Util.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Module with various helpers that are used in morley-client fake tests.
module Test.Util
  ( chainOperationHandlers
  , dumbContractState
  , dumbImplicitContractState
  , dumbManagerKey
  , mapToContractStateBigMap
  , handleGetBigMapValue

    -- * Internals
  , handleRunOperationInternal
  , assertHeadBlockId
  ) where

import Control.Exception.Safe (throwString)
import Control.Lens (at, (?~))
import Data.Aeson (encode)
import Data.ByteArray (ScrubbedBytes)
import Data.ByteString.Lazy qualified as LBS (toStrict)
import Data.Map as Map (elems, insert, lookup, toList)
import Data.Singletons (demote)
import Fmt ((+|), (|+))
import Network.HTTP.Types.Status (status404)
import Network.HTTP.Types.Version (http20)
import Servant.Client.Core
  (BaseUrl(..), ClientError(..), RequestF(..), ResponseF(..), Scheme(..), defaultRequest)
import Text.Hex (encodeHex)

import Lorentz as L (compileLorentz, drop)
import Lorentz.Constraints
import Lorentz.Pack
import Morley.Client.RPC.Types
import Morley.Client.Types
import Morley.Micheline
import Morley.Michelson.Typed
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Address.Kinds
import Morley.Tezos.Crypto
import Morley.Tezos.Crypto.Ed25519 qualified as Ed25519
import Morley.Util.ByteString
import TestM

-- | Function to convert given map to big map representation
-- used in fake state.
mapToContractStateBigMap
  :: forall k v. (NicePackedValue k, NicePackedValue v)
  => BigMapId k v -> Map k v -> ContractStateBigMap
mapToContractStateBigMap (BigMapId bigMapId) map' = ContractStateBigMap
  { csbmKeyType = toExpression $ demote @(ToT k)
  , csbmValueType = toExpression $ demote @(ToT v)
  , csbmMap = fromList $
    map (bimap (encodeBase58Check . valueToScriptExpr) lEncodeValue) $
    Map.toList map'
  , csbmId = bigMapId
  }

-- | Initial simple contract fake state.
dumbContractState :: ContractState 'AddressKindContract
dumbContractState = ContractState
  { csCounter = 100500
  , csAlias = "genesis2"
  , csContractData = ContractData
      OriginationScript
        { osCode    = toExpression $ compileLorentz L.drop
        , osStorage = toExpression $ toVal ()
        }
      Nothing
  }

dumbImplicitContractState :: ContractState 'AddressKindImplicit
dumbImplicitContractState = ContractState
  { csCounter = 100500
  , csAlias = "genesis1"
  , csContractData = ImplicitContractData Nothing
  }

-- | Fake handlers used for transaction sending and contract origination.
chainOperationHandlers :: Monad m => Handlers (TestT m)
chainOperationHandlers = defaultHandlers
  { hGetBlockHash = handleGetBlockHash
  , hGetCounter = handleGetCounter
  , hGetBlockConstants = handleGetBlockConstants
  , hGetProtocolParameters = handleGetProtocolParameters
  , hRunOperation = handleRunOperation
  , hPreApplyOperations = mapM . handlePreApplyOperation
  , hForgeOperation = handleForgeOperation
  , hInjectOperation = pure . OperationHash . (<> "_injected") . encodeHex . unHexJSONByteString
  , hGetContractScript = handleGetContractScript
  , hSignBytes =
    \_ _ -> pure . SignatureEd25519 . Ed25519.sign testSecretKey
  , hWaitForOperation = id
  , hGetAlias = handleGetAlias
  , hResolveAddressMaybe = handleResolveAddressMaybe
  , hRememberContract = handleRememberContract
  , hGetKeyPassword = \_ -> pure Nothing
  , hGenKey = handleGenKey
  , hGetManagerKey = handleGetManagerKey
  , hRevealKey = handleRevealKey
  }
  where
    testSecretKey :: Ed25519.SecretKey
    testSecretKey = Ed25519.detSecretKey "\001\002\003\004"

mkRunOperationResult :: [ContractAddress] -> RunOperationResult
mkRunOperationResult originatedContracts = RunOperationResult
  { rrOperationContents =
    one $ OperationContent $ RunMetadata
      { rmOperationResult = OperationApplied $
        AppliedResult 100500 100500 100500 originatedContracts 0
      , rmInternalOperationResults = []
      }
  }

handleGetBlockHash :: Monad m => BlockId -> TestT m BlockHash
handleGetBlockHash blkId = do
  unless (blkId == FinalHeadId) do
    throwString "Expected `getBlockHash` to be called with `head~2`."

  FakeState{..} <- get
  pure fsFinalHeadBlock

handleGetCounter
  :: ( MonadState FakeState m
     , MonadThrow m
     )
  => BlockId -> ImplicitAddress -> m TezosInt64
handleGetCounter blk addr = do
  assertHeadBlockId blk
  FakeState{..} <- get
  case lookup addr fsImplicits of
    Nothing -> throwM $ UnknownContract $ SomeAddressOrAlias $ AddressResolved addr
    Just ContractState{..} -> pure $ csCounter

handleGetBlockConstants
  :: MonadState FakeState m
  => BlockId -> m BlockConstants
handleGetBlockConstants blkId = do
  FakeState{..} <- get
  pure $ fsBlockConstants blkId

handleGetProtocolParameters
  :: (MonadState FakeState m, MonadThrow m)
  => BlockId -> m ProtocolParameters
handleGetProtocolParameters blk = do
  assertHeadBlockId blk
  FakeState{..} <- get
  pure $ fsProtocolParameters

handleRunOperation :: Monad m => BlockId -> RunOperation -> TestT m RunOperationResult
handleRunOperation blk RunOperation{..} = do
  assertHeadBlockId blk
  FakeState{..} <- get
  -- Ensure that passed chain id matches with one that fake state has
  unless (roChainId == bcChainId (fsBlockConstants blk)) (throwM $ InvalidChainId)
  -- As of release of the ithaca protocol, the "branch" field should be "head~2".
  -- https://web.archive.org/web/20220305165609/https://tezos.gitlab.io/protocols/tenderbake.html
  unless (roiBranch roOperation == fsFinalHeadBlock) do
    throwM $ InvalidBranch $ roiBranch roOperation
  originatedContracts <- handleRunOperationInternal roOperation
  pure $ mkRunOperationResult originatedContracts

handlePreApplyOperation :: Monad m => BlockId -> PreApplyOperation -> TestT m RunOperationResult
handlePreApplyOperation blk PreApplyOperation{..} = do
  assertHeadBlockId blk
  FakeState{..} <- get
  -- Ensure that passed protocol matches with one that mock state has
  unless (paoProtocol == bcProtocol (fsBlockConstants blk)) $
    throwM InvalidProtocol
  -- As of release of the ithaca protocol, the "branch" field should be "head~2".
  -- https://web.archive.org/web/20220305165609/https://tezos.gitlab.io/protocols/tenderbake.html
  unless (paoBranch == fsFinalHeadBlock) do
    throwM $ InvalidBranch paoBranch
  originatedContracts <- concatMapM handleTransactionOrOrigination paoContents
  pure $ mkRunOperationResult originatedContracts

handleForgeOperation :: Monad m => BlockId -> ForgeOperation -> TestT m HexJSONByteString
handleForgeOperation blkId op = do
  assertHeadBlockId blkId
  ms <- get
  -- As of release of the ithaca protocol, the "branch" field should be "head~2".
  -- https://web.archive.org/web/20220305165609/https://tezos.gitlab.io/protocols/tenderbake.html
  unless (foBranch op == fsFinalHeadBlock ms) do
    throwM $ InvalidBranch $ foBranch op
  pure . HexJSONByteString . LBS.toStrict . encode $ op

handleRunOperationInternal :: Monad m => RunOperationInternal -> TestT m [ContractAddress]
handleRunOperationInternal RunOperationInternal{..} = do
  concatMapM handleTransactionOrOrigination roiContents

handleTransactionOrOrigination
  :: (Monad m, HasCallStack) => OperationInput -> TestT m [ContractAddress]
handleTransactionOrOrigination op = do
  FakeState{..} <- get
  case wcoCustom op of
    -- Ensure that transaction sender exists
    OpTransfer TransactionOperation{..} -> case lookup codSource fsImplicits of
      Nothing -> throwM $ UnknownContract $ SomeAddressOrAlias $ AddressResolved codSource
      Just ContractState{..} -> do
        -- Ensure that sender counter matches
        unless (csCounter + 1 == codCounter) (throwM CounterMismatch)
        case toDestination of
          MkAddress dest@ContractAddress{} ->
            case lookup dest fsContracts of
              Nothing -> throwM $ UnknownContract $ SomeAddressOrAlias $ AddressResolved dest
              Just _ -> pure []
          MkAddress dest@ImplicitAddress{} ->
            case lookup dest fsImplicits of
              Nothing -> throwM $ UnknownContract $ SomeAddressOrAlias $ AddressResolved dest
              Just _ -> pure []
          MkAddress TxRollupAddress{} -> error "tx rollup unsupported"
    -- Ensure that originator exists
    OpOriginate _ -> case lookup codSource fsImplicits of
      Nothing -> throwM $ UnknownContract $ SomeAddressOrAlias $ AddressResolved codSource
      Just ContractState{..} -> do
        -- Ensure that originator counter matches
        unless (csCounter + 1 == codCounter) (throwM CounterMismatch)
        pure [dummyContractAddr]
      where
        dummyContractAddr = [ta|KT1LZwEZqbqtLYhwzaidBp6So9LgYDpkpEv7|]
    OpReveal _ ->
      -- We do not care about reveals at the moment
      return []
  where
    CommonOperationData{..} = wcoCommon op

-- | In most places, @morley-client@ executes operations against the @head@ block.
assertHeadBlockId :: (HasCallStack, MonadThrow m) => BlockId -> m ()
assertHeadBlockId blockId = unless (blockId == HeadId) $
  throwString "Accessing non-head block is not supported in tests"

handleGetContractScript
  :: ( MonadState FakeState m
     , MonadThrow m
     )
  => BlockId
  -> ContractAddress
  -> m OriginationScript
handleGetContractScript blockId addr = do
  assertHeadBlockId blockId
  FakeState{..} <- get
  case lookup addr fsContracts of
    Nothing -> throwM $ err404 path
    Just ContractState{..} -> case csContractData of
      ContractData script _ -> pure script
  where
    path = "/chains/main/blocks/head/context/contracts/" <> formatAddress addr <> "/script"

handleGetBigMapValue :: Monad m => BlockId -> Natural -> Text -> TestT m Expression
handleGetBigMapValue blockId bigMapId scriptExpr = do
  assertHeadBlockId blockId
  st <- get

  let allBigMaps :: [ContractStateBigMap] =
        catMaybes $
          Map.elems (fsContracts st) <&> \cs -> case (csContractData cs) of
            ContractData _ bigMapMaybe -> bigMapMaybe

  -- Check if a big_map with the given ID exists and, if so, check
  -- whether the giv en key exists.
  case find (\bigMap -> csbmId bigMap == bigMapId) allBigMaps of
    Nothing -> throwM $ err404 path
    Just bigMap ->
      case lookup scriptExpr (csbmMap bigMap ) of
        Nothing -> throwM $ err404 path
        Just serializedValue -> pure $ decodeExpression serializedValue
  where
    path = "/chains/main/blocks/head/context/big_maps/" <> show bigMapId <> "/" <> scriptExpr

handleRememberContract :: Monad m => Bool -> ContractAddress -> ContractAlias -> TestT m ()
handleRememberContract replaceExisting addr alias = do
  let
    cs = dumbContractState { csAlias = alias }
    remember addr' cs' FakeState{..} =
      modify $ \s -> s { fsContracts = insert addr' cs' fsContracts }

  st@FakeState{..} <- get
  case lookup addr fsContracts of
    Nothing -> remember addr cs st
    _       -> bool pass (remember addr cs st) replaceExisting

handleGenKey :: Monad m => ImplicitAlias -> TestT m ImplicitAddress
handleGenKey alias = do
  let
    addr = detGenKeyAddress (encodeUtf8 $ unAlias alias)
    newContractState = dumbImplicitContractState { csAlias =  alias }
  modify $ \s ->
    s & fsImplicitsL . at addr ?~ newContractState
  pure addr

handleGetAlias
  :: forall m kind. (Monad m, HasCallStack)
  => AddressOrAlias kind -> TestT m (Alias kind)
handleGetAlias = \case
  AddressAlias alias -> pure alias
  AddressResolved addr -> do
    FakeState{..} <- get
    case addr of
      ContractAddress{} -> go addr fsContracts
      ImplicitAddress{} -> go addr fsImplicits
      TxRollupAddress{} -> error "tx rollup unsupported"
  where
    go :: KindedAddress kind -> Map (KindedAddress kind) (ContractState kind) -> TestT m (Alias kind)
    go addr m = case lookup addr m of
      Nothing -> throwM $ UnknownContract $ SomeAddressOrAlias $ AddressResolved addr
      Just ContractState{..} -> pure $ csAlias

handleGetManagerKey :: (Monad m) => BlockId -> ImplicitAddress -> TestT m (Maybe PublicKey)
handleGetManagerKey blk addr = do
  assertHeadBlockId blk
  s <- get
  let mbCs = s ^. fsImplicitsL . at addr
  case mbCs of
    Just ContractState{..} -> case csContractData of
      ImplicitContractData mbManagerKey -> pure mbManagerKey
    Nothing -> throwM $ UnknownContract $ SomeAddressOrAlias $ AddressResolved addr

-- In scenarios where the system under test checks for 404 errors, we
-- use this function to fake and simulate those errors.
err404 :: Text -> ClientError
err404 path = FailureResponse
  (defaultRequest { requestBody = Nothing, requestPath = (baseUrl , "") })
  response
  where
    baseUrl = BaseUrl
      { baseUrlScheme = Http
      , baseUrlHost = "localhost"
      , baseUrlPort = 8732
      , baseUrlPath = toString path
      }
    response = Response
      { responseStatusCode = status404
      , responseHeaders = mempty
      , responseHttpVersion = http20
      , responseBody = "Contract with given address not found"
      }

handleResolveAddressMaybe
  :: forall m kind. Monad m
  => AddressOrAlias kind
  -> TestT m (Maybe (KindedAddress kind))
handleResolveAddressMaybe = \case
  AddressResolved addr -> pure (Just addr)
  AddressAlias alias -> do
    FakeState{..} <- get
    case alias of
      ImplicitAlias{} -> go fsImplicits
      ContractAlias{} -> go fsContracts
    where
      checkAlias (_, ContractState { csAlias = alias' }) = alias == alias'
      go :: Map (KindedAddress kind) (ContractState kind) -> TestT m (Maybe (KindedAddress kind))
      go m = case find checkAlias $ Map.toList m of
        Just (addr, _) -> pure (Just addr)
        Nothing -> pure Nothing

handleRevealKey :: Monad m => ImplicitAlias -> Maybe ScrubbedBytes -> TestT m ()
handleRevealKey alias _ = do
  contracts <- gets (Map.toList . fsImplicits)
  let contracts' = filter (\(_, ContractState{..}) -> csAlias == alias) contracts
  case contracts' of
    []  -> throwM $ UnknownContract $ SomeAddressOrAlias $ AddressAlias alias
    [(addr, cs@ContractState{..})] ->
      case csContractData of
        ImplicitContractData (Just _) ->
          throwM $ AlreadyRevealed $ MkAddress addr
        ImplicitContractData Nothing ->
            -- We don't care about the public key itself, but only its presence.
            let newContractState = cs { csContractData = ImplicitContractData $ Just dumbManagerKey }
            in modify $ \s -> s & fsImplicitsL . at addr ?~ newContractState
    _   ->
      error $ "Multiple contracts have alias '" +| alias |+
        "'. This is most likely a bug in tests."

-- | Dummy public key used in fake tests.
dumbManagerKey :: PublicKey
dumbManagerKey = fromRight (error "impossible") $ parsePublicKey
  "edpkuwTWKgQNnhR5v17H2DYHbfcxYepARyrPGbf1tbMoGQAj8Ljr3V"