Value decoding/parsing: reject currencies/tokens longer than 32 bytes (#7371)

Author: zliu41

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

@@ -20,7 +20,6 @@ import PlutusCore.Core
 import PlutusCore.Data (Data)
 import PlutusCore.DeBruijn
 import PlutusCore.Name.Unique
-import PlutusCore.Value (Value)
 
 import Data.Proxy
 import PlutusCore.Flat
@@ -122,8 +121,6 @@ decodeConstant = dBEBits8 constantWidth
 
 deriving via FlatViaSerialise Data instance Flat Data
 
-deriving via FlatViaSerialise Value instance Flat Value
-
 decodeKindedUniFlat :: Closed uni => Get (SomeTypeIn (Kinded uni))
 decodeKindedUniFlat =
     go . decodeKindedUni . map (fromIntegral :: Word8 -> Int)

plutus-core/plutus-core/src/PlutusCore/Parser/Builtin.hs

@@ -18,7 +18,7 @@ import PlutusCore.Value qualified as PLC (Value)
 import PlutusCore.Value qualified as Value
 
 import Control.Monad.Combinators
-import Data.ByteString (ByteString, pack)
+import Data.ByteString (ByteString, pack, unpack)
 import Data.Map.Strict qualified as Map
 import Data.Text qualified as T
 import Data.Text.Internal.Read (hexDigitToInt)
@@ -90,7 +90,16 @@ conArray uniA = Vector.fromList <$> conList uniA
 
 -- | Parser for values.
 conValue :: Parser PLC.Value
-conValue = Value.fromList <$> conList knownUni
+conValue = do
+  Value.fromList <$> (traverse validateKeys =<< conList knownUni)
+ where
+  validateToken (token, amt) = do
+    tk <- maybe (fail $ "Invalid token: " <> show (unpack token)) pure (Value.k token)
+    pure (tk, amt)
+  validateKeys (currency, tokens) = do
+    ck <- maybe (fail $ "Invalid currency: " <> show (unpack currency)) pure (Value.k currency)
+    tks <- traverse validateToken tokens
+    pure (ck, tks)
 
 -- | Parser for pairs.
 conPair :: DefaultUni (Esc a) -> DefaultUni (Esc b) -> Parser (a, b)

plutus-core/plutus-core/src/PlutusCore/Pretty/PrettyConst.hs

@@ -159,6 +159,9 @@ instance PrettyBy ConstConfig Data where
         I i         ->  ("I" <+> prettyArg i) :| []
         B b         ->  ("B" <+> prettyArg b) :| []
 
+instance PrettyBy ConstConfig Value.K where
+    prettyBy config = prettyBy config . Value.unK
+
 instance PrettyBy ConstConfig Value where
     prettyBy config = prettyBy config . Value.toList
 

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

@@ -5,6 +5,10 @@
 
 module PlutusCore.Value (
   Value, -- Do not expose data constructor
+  K, -- Do not expose data constructor
+  k,
+  unK,
+  maxKeyLen,
   NestedMap,
   unpack,
   pack,
@@ -23,11 +27,12 @@ module PlutusCore.Value (
   unValueData,
 ) where
 
-import Codec.Serialise (Serialise)
+import Codec.Serialise qualified as CBOR
 import Control.DeepSeq (NFData)
 import Data.Bifunctor
 import Data.Bitraversable
 import Data.ByteString (ByteString)
+import Data.ByteString qualified as B
 import Data.ByteString.Base64 qualified as Base64
 import Data.Functor
 import Data.Hashable (Hashable)
@@ -39,11 +44,45 @@ import Data.Map.Strict qualified as Map
 import Data.Maybe
 import Data.Text.Encoding qualified as Text
 import GHC.Generics
+
 import PlutusCore.Builtin.Result
 import PlutusCore.Data (Data (..))
+import PlutusCore.Flat qualified as Flat
 import PlutusPrelude (Pretty (..))
 
-type NestedMap = Map ByteString (Map ByteString Integer)
+-- Max length (in bytes) for currency symbols and token names in `Value`,
+-- both of which cannot exceed 32 bytes. Currency symbols are in fact either
+-- empty or 28 bytes, but for simplicity we allow anything between 0 and 32 bytes.
+maxKeyLen :: Int
+maxKeyLen = 32
+{-# INLINE maxKeyLen #-}
+
+-- | A `ByteString` with maximum length of `maxKeyLen` bytes.
+newtype K = UnsafeK {unK :: ByteString}
+  deriving newtype (Eq, Ord, Show, Hashable, NFData)
+  deriving stock (Generic)
+
+k :: ByteString -> Maybe K
+k b = if B.length b <= maxKeyLen then Just (UnsafeK b) else Nothing
+{-# INLINEABLE k #-}
+
+instance Flat.Flat K where
+  encode (UnsafeK b) = Flat.encode b
+  {-# INLINE encode #-}
+  decode = do
+    b <- Flat.decode
+    maybe (fail $ "Invalid Value key: " <> show (B.unpack b)) pure (k b)
+  {-# INLINEABLE decode #-}
+
+instance CBOR.Serialise K where
+  encode (UnsafeK b) = CBOR.encode b
+  {-# INLINE encode #-}
+  decode = do
+    b <- CBOR.decode
+    maybe (fail $ "Invalid Value key: " <> show (B.unpack b)) pure (k b)
+  {-# INLINEABLE decode #-}
+
+type NestedMap = Map K (Map K Integer)
 
 -- | The underlying type of the UPLC built-in type @Value@.
 data Value
@@ -61,11 +100,23 @@ data Value
       Invariant: all values are positive.
       -}
       {-# UNPACK #-} !Int
-      {-^ Total size, i.e., sum total of inner map sizes. This avoids recomputing
+      {- ^ Total size, i.e., sum total of inner map sizes. This avoids recomputing
       the total size during the costing of operations like `unionValue`.
       -}
   deriving stock (Eq, Show, Generic)
-  deriving anyclass (Hashable, Serialise, NFData)
+  deriving anyclass (Hashable, NFData)
+
+instance CBOR.Serialise Value where
+  encode (Value v _ _) = CBOR.encode v
+  {-# INLINE encode #-}
+  decode = pack <$> CBOR.decode
+  {-# INLINE decode #-}
+
+instance Flat.Flat Value where
+  encode (Value v _ _) = Flat.encode v
+  {-# INLINE encode #-}
+  decode = pack <$> Flat.decode
+  {-# INLINE decode #-}
 
 {-| Unpack a `Value` into a map from (currency symbol, token name) to amount.
 
@@ -110,19 +161,20 @@ empty :: Value
 empty = Value mempty mempty 0
 {-# INLINE empty #-}
 
-toList :: Value -> [(ByteString, [(ByteString, Integer)])]
+toList :: Value -> [(K, [(K, Integer)])]
 toList = Map.toList . Map.map Map.toList . unpack
 {-# INLINEABLE toList #-}
 
-toFlatList :: Value -> [(ByteString, ByteString, Integer)]
+toFlatList :: Value -> [(K, K, Integer)]
 toFlatList (toList -> xs) = [(c, t, a) | (c, ys) <- xs, (t, a) <- ys]
 {-# INLINEABLE toFlatList #-}
 
-fromList :: [(ByteString, [(ByteString, Integer)])] -> Value
+fromList :: [(K, [(K, Integer)])] -> Value
 fromList =
   pack
     . Map.fromListWith (Map.unionWith (+))
     . fmap (second (Map.fromListWith (+)))
+{-# INLINEABLE fromList #-}
 
 normalize :: NestedMap -> NestedMap
 normalize = Map.filter (not . Map.null) . Map.map (Map.filter (/= 0))
@@ -131,7 +183,7 @@ normalize = Map.filter (not . Map.null) . Map.map (Map.filter (/= 0))
 instance Pretty Value where
   pretty = pretty . fmap (bimap toText (fmap (first toText))) . toList
    where
-    toText = Text.decodeLatin1 . Base64.encode
+    toText = Text.decodeLatin1 . Base64.encode . unK
 
 {-| \(O(\log \max(m, k))\), where \(m\) is the size of the outer map, and \(k\) is
 the size of the largest inner map.
@@ -140,38 +192,40 @@ insertCoin :: ByteString -> ByteString -> Integer -> Value -> Value
 insertCoin currency token amt v@(Value outer sizes size)
   | amt == 0 = deleteCoin currency token v
   | otherwise =
-      let (mold, outer') = Map.alterF f currency outer
+      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 ByteString Integer)
+    :: Maybe (Map K Integer)
     -> ( -- Just (old size of inner map) if the total size grows by 1, otherwise Nothing
          Maybe Int
-       , Maybe (Map ByteString Integer)
+       , Maybe (Map K Integer)
        )
   f = \case
-    Nothing -> (Just 0, Just (Map.singleton token amt))
+    Nothing -> (Just 0, Just (Map.singleton (UnsafeK token) amt))
     Just inner ->
-      let (isJust -> exists, inner') = Map.insertLookupWithKey (\_ _ _ -> amt) token amt inner
+      let (isJust -> exists, inner') =
+            Map.insertLookupWithKey (\_ _ _ -> amt) (UnsafeK token) amt inner
        in (if exists then Nothing else Just (Map.size inner), Just inner')
 {-# INLINEABLE insertCoin #-}
 
 -- | \(O(\log \max(m, k))\)
 deleteCoin :: ByteString -> ByteString -> Value -> Value
-deleteCoin currency token (Value outer sizes size) = Value outer' sizes' size'
+deleteCoin (UnsafeK -> currency) (UnsafeK -> token) (Value outer sizes size) =
+  Value outer' sizes' size'
  where
   (mold, outer') = Map.alterF f currency outer
   (sizes', size') = case mold of
     Just old -> (updateSizes old (old - 1) sizes, size - 1)
     Nothing  -> (sizes, size)
   f
-    :: Maybe (Map ByteString Integer)
+    :: Maybe (Map K Integer)
     -> ( -- Just (old size of inner map) if the total size shrinks by 1, otherwise Nothing
          Maybe Int
-       , Maybe (Map ByteString Integer)
+       , Maybe (Map K Integer)
        )
   f = \case
     Nothing -> (Nothing, Nothing)
@@ -181,9 +235,10 @@ deleteCoin currency token (Value outer sizes size) = Value outer' sizes' size'
 
 -- | \(O(\log \max(m, k))\)
 lookupCoin :: ByteString -> ByteString -> Value -> Integer
-lookupCoin currency token (unpack -> outer) = case Map.lookup currency outer of
-  Nothing    -> 0
-  Just inner -> Map.findWithDefault 0 token inner
+lookupCoin (UnsafeK -> currency) (UnsafeK -> token) (unpack -> outer) =
+  case Map.lookup currency outer of
+    Nothing    -> 0
+    Just inner -> Map.findWithDefault 0 token inner
 
 {-| \(O(n_{2}\log \max(m_{1}, k_{1}))\), where \(n_{2}\) is the total size of the second
 `Value`, \(m_{1}\) is the size of the outer map in the first `Value` and \(k_{1}\) is
@@ -200,7 +255,7 @@ valueContains v = Map.foldrWithKey' go True . unpack
    where
     goInner t a2 =
       (&&)
-        ( let a1 = lookupCoin c t v
+        ( let a1 = lookupCoin (unK c) (unK t) v
            in if a2 > 0
                 then a1 >= a2
                 else a1 == a2
@@ -240,10 +295,10 @@ unionValue (unpack -> vA) (unpack -> vB) =
 This is the denotation of @ValueData@ in Plutus V1, V2 and V3.
 -}
 valueData :: Value -> Data
-valueData = Map . fmap (bimap B tokensData) . Map.toList . unpack
+valueData = Map . fmap (bimap (B . unK) tokensData) . Map.toList . unpack
  where
-  tokensData :: Map ByteString Integer -> Data
-  tokensData = Map . fmap (bimap B I) . Map.toList
+  tokensData :: Map K Integer -> Data
+  tokensData = Map . fmap (bimap (B . unK) I) . Map.toList
 {-# INLINEABLE valueData #-}
 
 {-| \(O(n \log n)\). Decodes `Data` into `Value`, in the same way as non-builtin @Value@.
@@ -255,17 +310,17 @@ unValueData =
     Map cs -> fmap (Map.fromListWith (Map.unionWith (+))) (traverse (bitraverse unB unTokens) cs)
     _ -> fail "unValueData: non-Map constructor"
  where
-  unB :: Data -> BuiltinResult ByteString
+  unB :: Data -> BuiltinResult K
   unB = \case
-    B b -> pure b
+    B b -> pure (UnsafeK b)
     _ -> fail "unValueData: non-B constructor"
 
   unI :: Data -> BuiltinResult Integer
   unI = \case
     I i -> pure i
     _ -> fail "unValueData: non-I constructor"
 
-  unTokens :: Data -> BuiltinResult (Map ByteString Integer)
+  unTokens :: Data -> BuiltinResult (Map K Integer)
   unTokens = \case
     Map ts -> fmap (Map.fromListWith (+)) (traverse (bitraverse unB unI) ts)
     _ -> fail "unValueData: non-Map constructor"

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

@@ -15,6 +15,7 @@ import Hedgehog.Gen qualified as Gen
 import Hedgehog.Range qualified as Range
 import Test.Tasty
 import Test.Tasty.Hedgehog
+import Test.Tasty.HUnit
 
 -- | The `SrcSpan` of a parsed `Term` should not including trailing whitespaces.
 propTermSrcSpan :: Property
@@ -30,6 +31,46 @@ propTermSrcSpan = property $ do
              in (srcSpanELine sp, srcSpanECol sp) === (endingLine, endingCol + 1)
         Left err -> annotate (display err) >> failure
 
+expectParserSuccess :: T.Text -> Assertion
+expectParserSuccess code = case runQuoteT (parseTerm code) of
+  Right _ -> pure ()
+  Left _  -> assertFailure $ "Unexpected failure when parsing term: " <> T.unpack code
+
+expectParserFailure :: T.Text -> Assertion
+expectParserFailure code = case runQuoteT (parseTerm code) of
+  Right _ -> assertFailure $ "Unexpected success when parsing term: " <> T.unpack code
+  Left _  -> pure ()
+
+parseValueInvalidCurrency :: Assertion
+parseValueInvalidCurrency = do
+  expectParserFailure code
+  where
+    -- Currency is 33 bytes
+    code = "(con value \
+       \[ ( #616161616161616161616161616161616161616161616161616161616161616161\
+       \, [ ( #6161616161616161616161616161616161616161616161616161616161616161\
+       \    , -100 ) ] ) ])"
+
+parseValueInvalidToken :: Assertion
+parseValueInvalidToken = do
+  expectParserFailure code
+  where
+    -- Token is 33 bytes
+    code = "(con value \
+       \[ ( #6161616161616161616161616161616161616161616161616161616161616161\
+       \, [ ( #616161616161616161616161616161616161616161616161616161616161616161\
+       \    , -100 ) ] ) ])"
+
+parseValueValid :: Assertion
+parseValueValid = do
+  expectParserSuccess code
+  where
+    -- Both currency and token are 32 bytes
+    code = "(con value \
+       \[ ( #6161616161616161616161616161616161616161616161616161616161616161\
+       \, [ ( #6161616161616161616161616161616161616161616161616161616161616161\
+       \    , -100 ) ] ) ])"
+
 tests :: TestTree
 tests =
     testGroup
@@ -38,4 +79,13 @@ tests =
             "parser captures ending positions correctly"
             "propTermSrcSpan"
             propTermSrcSpan
+        , testCase
+            "parser of Value should fail upon invalid currency"
+            parseValueInvalidCurrency
+        , testCase
+            "parser of Value should fail upon invalid token"
+            parseValueInvalidToken
+        , testCase
+            "parser of Value should succeed"
+            parseValueValid
         ]