Enforce currency and token lengths in insertCoin and unValueData (#7372)

Author: zliu41

plutus-core/plutus-core/src/PlutusCore/Default/Builtins.hs

@@ -2047,7 +2047,7 @@ instance uni ~ DefaultUni => ToBuiltinMeaning uni DefaultFun where
             (runCostingFunTwoArguments . paramBls12_381_G2_multiScalarMul)
 
     toBuiltinMeaning _semvar InsertCoin =
-      let insertCoinDenotation :: ByteString -> ByteString -> Integer -> Value -> Value
+      let insertCoinDenotation :: ByteString -> ByteString -> Integer -> Value -> BuiltinResult Value
           insertCoinDenotation = Value.insertCoin
           {-# INLINE insertCoinDenotation #-}
        in makeBuiltinMeaning

plutus-core/plutus-core/src/PlutusCore/Value.hs

@@ -188,28 +188,31 @@ instance Pretty Value where
 {-| \(O(\log \max(m, k))\), where \(m\) is the size of the outer map, and \(k\) is
 the size of the largest inner map.
 -}
-insertCoin :: ByteString -> ByteString -> Integer -> Value -> Value
+insertCoin :: ByteString -> ByteString -> Integer -> Value -> BuiltinResult Value
 insertCoin currency token amt v@(Value outer sizes size)
-  | amt == 0 = deleteCoin currency token v
-  | otherwise =
-      let (mold, outer') = Map.alterF f (UnsafeK currency) outer
-          (sizes', size') = case mold of
-            Just old -> (updateSizes old (old + 1) sizes, size + 1)
-            Nothing  -> (sizes, size)
-       in Value outer' sizes' size'
- where
-  f
-    :: Maybe (Map K Integer)
-    -> ( -- Just (old size of inner map) if the total size grows by 1, otherwise Nothing
-         Maybe Int
-       , Maybe (Map K Integer)
-       )
-  f = \case
-    Nothing -> (Just 0, Just (Map.singleton (UnsafeK token) amt))
-    Just inner ->
-      let (isJust -> exists, inner') =
-            Map.insertLookupWithKey (\_ _ _ -> amt) (UnsafeK token) amt inner
-       in (if exists then Nothing else Just (Map.size inner), Just inner')
+  | amt == 0 = pure $ deleteCoin currency token v
+  | otherwise = case (k currency, k token) of
+      (Nothing, _) -> fail $ "insertCoin: invalid currency: " <> show (B.unpack currency)
+      (_, Nothing) -> fail $ "insertCoin: invalid token: " <> show (B.unpack token)
+      (Just ck, Just tk) ->
+        let f
+              :: Maybe (Map K Integer)
+              -> ( -- Just (old size of inner map) if the total size grows by 1,
+                   -- otherwise Nothing
+                   Maybe Int
+                 , Maybe (Map K Integer)
+                 )
+            f = \case
+              Nothing -> (Just 0, Just (Map.singleton tk amt))
+              Just inner ->
+                let (isJust -> exists, inner') =
+                      Map.insertLookupWithKey (\_ _ _ -> amt) tk amt inner
+                 in (if exists then Nothing else Just (Map.size inner), Just inner')
+            (mold, outer') = Map.alterF f ck outer
+            (sizes', size') = case mold of
+              Just old -> (updateSizes old (old + 1) sizes, size + 1)
+              Nothing  -> (sizes, size)
+         in pure $ Value outer' sizes' size'
 {-# INLINEABLE insertCoin #-}
 
 -- | \(O(\log \max(m, k))\)
@@ -312,7 +315,7 @@ unValueData =
  where
   unB :: Data -> BuiltinResult K
   unB = \case
-    B b -> pure (UnsafeK b)
+    B b -> maybe (fail $ "unValueData: invalid key: " <> show (B.unpack b)) pure (k b)
     _ -> fail "unValueData: non-B constructor"
 
   unI :: Data -> BuiltinResult Integer

plutus-core/plutus-core/test/TypeSynthesis/Golden/Signatures/DefaultFun/InsertCoin.golden.sig

@@ -1 +1 @@
-ByteString -> ByteString -> Integer -> Value -> Value
+ByteString -> ByteString -> Integer -> Value -> BuiltinResult Value

plutus-core/plutus-core/test/Value/Spec.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TypeApplications  #-}
 {-# LANGUAGE ViewPatterns      #-}
+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 module Value.Spec (tests) where
@@ -11,15 +12,18 @@ import Data.Either
 import Data.Foldable qualified as F
 import Data.Map.Strict qualified as Map
 import Data.Maybe
-import PlutusCore.Flat qualified as Flat
-import PlutusCore.Generators.QuickCheck.Builtin (ValueAmount (..), genShortHex)
-import PlutusCore.Value (Value)
-import PlutusCore.Value qualified as V
 import Safe.Foldable (maximumMay)
 import Test.QuickCheck
 import Test.Tasty
 import Test.Tasty.QuickCheck
 
+import PlutusCore.Builtin (BuiltinResult (..))
+import PlutusCore.Data (Data (..))
+import PlutusCore.Flat qualified as Flat
+import PlutusCore.Generators.QuickCheck.Builtin (ValueAmount (..), genShortHex)
+import PlutusCore.Value (Value)
+import PlutusCore.Value qualified as V
+
 prop_packUnpackRoundtrip :: Value -> Property
 prop_packUnpackRoundtrip v = v === V.pack (V.unpack v)
 
@@ -38,15 +42,15 @@ prop_insertCoinBookkeeping :: Value -> ValueAmount -> Property
 prop_insertCoinBookkeeping v (ValueAmount amt) =
   forAll (genShortHex (V.totalSize v)) $ \currency ->
     forAll (genShortHex (V.totalSize v)) $ \token ->
-      let v' = V.insertCoin (V.unK currency) (V.unK token) amt v
+      let BuiltinSuccess v' = V.insertCoin (V.unK currency) (V.unK token) amt v
        in checkSizes v'
 
 -- | Verifies that @insertCoin@ preserves @Value@ invariants
 prop_insertCoinPreservesInvariants :: Value -> ValueAmount -> Property
 prop_insertCoinPreservesInvariants v (ValueAmount amt) =
   forAll (genShortHex (V.totalSize v)) $ \currency ->
     forAll (genShortHex (V.totalSize v)) $ \token ->
-      let v' = V.insertCoin (V.unK currency) (V.unK token) amt v
+      let BuiltinSuccess v' = V.insertCoin (V.unK currency) (V.unK token) amt v
        in checkInvariants v'
 
 prop_unionCommutative :: Value -> Value -> Property
@@ -60,15 +64,33 @@ prop_insertCoinIdempotent :: Value -> Property
 prop_insertCoinIdempotent v =
   v
     === F.foldl'
-      (\acc (c, t, a) -> let v' = V.insertCoin (V.unK c) (V.unK t) a acc in v')
+      (\acc (c, t, a) -> let BuiltinSuccess v' = V.insertCoin (V.unK c) (V.unK t) a acc in v')
       v
       (V.toFlatList v)
 
+prop_insertCoinValidatesCurrency :: Value -> Property
+prop_insertCoinValidatesCurrency v =
+  forAll gen33Bytes $ \c ->
+    forAll gen32BytesOrFewer $ \t ->
+      forAll (arbitrary `suchThat` (/= 0)) $ \amt ->
+        case V.insertCoin c t amt v of
+          BuiltinFailure{} -> property True
+          _                -> property False
+
+prop_insertCoinValidatesToken :: Value -> Property
+prop_insertCoinValidatesToken v =
+  forAll gen32BytesOrFewer $ \c ->
+    forAll gen33Bytes $ \t ->
+      forAll (arbitrary `suchThat` (/= 0)) $ \amt ->
+        case V.insertCoin c t amt v of
+          BuiltinFailure{} -> property True
+          _                -> property False
+
 prop_lookupAfterInsertion :: Value -> ValueAmount -> Property
 prop_lookupAfterInsertion v (ValueAmount amt) =
   forAll (genShortHex (V.totalSize v)) $ \currency ->
     forAll (genShortHex (V.totalSize v)) $ \token ->
-      let v' = V.insertCoin (V.unK currency) (V.unK token) amt v
+      let BuiltinSuccess v' = V.insertCoin (V.unK currency) (V.unK token) amt v
        in V.lookupCoin (V.unK currency) (V.unK token) v' === amt
 
 prop_lookupAfterDeletion :: Value -> Property
@@ -84,7 +106,7 @@ prop_deleteCoinIdempotent v0 =
     let v' = V.deleteCoin c t v
      in v' === V.deleteCoin c t v'
  where
-  v = if V.totalSize v0 > 0 then v0 else V.insertCoin "c" "t" 1 v0
+  BuiltinSuccess v = if V.totalSize v0 > 0 then pure v0 else V.insertCoin "c" "t" 1 v0
   fl = V.toFlatList v
 
 prop_containsReflexive :: Value -> Property
@@ -113,7 +135,7 @@ prop_flatDecodeSuccess = forAll (arbitrary `suchThat` (/= 0)) $ \amt ->
   forAll gen32BytesOrFewer $ \c ->
     forAll gen32BytesOrFewer $ \t ->
       let flat = Flat.flat $ Map.singleton c (Map.singleton t amt)
-          v = V.insertCoin c t amt V.empty
+          BuiltinSuccess v = V.insertCoin c t amt V.empty
        in Flat.unflat flat === Right v
 
 prop_flatDecodeInvalidCurrency :: Property
@@ -145,6 +167,24 @@ checkInvariants (V.unpack -> v) =
   property ((not . any Map.null) v)
     .&&. property ((not . any (elem 0)) v)
 
+prop_unValueDataValidatesCurrency :: ValueAmount -> Property
+prop_unValueDataValidatesCurrency (ValueAmount amt) =
+  forAll gen33Bytes $ \c ->
+    forAll gen32BytesOrFewer $ \t ->
+      let d = Map [(B c, Map [(B t, I amt)])]
+       in case V.unValueData d of
+            BuiltinFailure{} -> property True
+            _                -> property False
+
+prop_unValueDataValidatesToken :: ValueAmount -> Property
+prop_unValueDataValidatesToken (ValueAmount amt) =
+  forAll gen32BytesOrFewer $ \c ->
+    forAll gen33Bytes $ \t ->
+      let d = Map [(B c, Map [(B t, I amt)])]
+       in case V.unValueData d of
+            BuiltinFailure{} -> property True
+            _                -> property False
+
 tests :: TestTree
 tests =
   testGroup
@@ -173,6 +213,12 @@ tests =
     , testProperty
         "insertCoinIdempotent"
         prop_insertCoinIdempotent
+    , testProperty
+        "insertCoinValidatesCurrency"
+        prop_insertCoinValidatesCurrency
+    , testProperty
+        "insertCoinValidatesToken"
+        prop_insertCoinValidatesToken
     , testProperty
         "lookupAfterInsertion"
         prop_lookupAfterInsertion
@@ -188,6 +234,12 @@ tests =
     , testProperty
         "containsAfterDeletion"
         prop_containsAfterDeletion
+    , testProperty
+        "unValueDataValidatesCurrency"
+        prop_unValueDataValidatesCurrency
+    , testProperty
+        "unValueDataValidatesToken"
+        prop_unValueDataValidatesToken
     , testProperty
         "flatRoundtrip"
         prop_flatRoundtrip

plutus-tx/src/PlutusTx/Builtins/Internal.hs

@@ -1084,8 +1084,13 @@ insertCoin
   -> BuiltinInteger
   -> BuiltinValue
   -> BuiltinValue
-insertCoin (BuiltinByteString c) (BuiltinByteString t) amt (BuiltinValue v) =
-  BuiltinValue $ Value.insertCoin c t amt v
+insertCoin (BuiltinByteString c) (BuiltinByteString t) amt (BuiltinValue v0) =
+  case Value.insertCoin c t amt v0 of
+    BuiltinSuccess v -> BuiltinValue v
+    BuiltinSuccessWithLogs logs v -> traceAll logs (BuiltinValue v)
+    BuiltinFailure logs err ->
+      traceAll (logs <> pure (display err)) $
+        Haskell.error "insertCoin errored."
 {-# OPAQUE insertCoin #-}
 
 lookupCoin