From 53bdf096a45af2c98c72cf7fdedd35b8b219fec1 Mon Sep 17 00:00:00 2001
From: Ilya Rezvov <ilya.rezvov@revjet.com>
Date: Wed, 7 Apr 2021 22:38:34 -0700
Subject: [PATCH] fix nan with payload parsing and safely handle out of range
 type indecies

---
 src/Language/Wasm/Lexer.x  | 35 +++++++++++++++++++++++++++++------
 src/Language/Wasm/Parser.y |  8 +++++++-
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/Language/Wasm/Lexer.x b/src/Language/Wasm/Lexer.x
index 35c4299..8b3170c 100644
--- a/src/Language/Wasm/Lexer.x
+++ b/src/Language/Wasm/Lexer.x
@@ -14,9 +14,9 @@ import qualified Data.ByteString.Lazy as LBS
 import qualified Data.Char as Char
 import qualified Data.ByteString.Lazy.UTF8 as LBSUtf8
 import Control.Monad (when)
-import Numeric.IEEE (infinity, nan)
+import Numeric.IEEE (infinity, nan, nanWithPayload)
 import Language.Wasm.FloatUtils (makeNaN, doubleToFloat)
-import Data.Word (Word8)
+import Data.Word (Word8, Word64)
 import Data.List (isPrefixOf)
 import Text.Read (readEither)
 
@@ -62,8 +62,8 @@ tokens :-
 <0> "nan"                                 { constToken $ TFloatLit $ BinRep (abs nan) }
 <0> "+nan"                                { constToken $ TFloatLit $ BinRep (abs nan) }
 <0> "-nan"                                { constToken $ TFloatLit $ BinRep nan }
-<0> "nan:canonical"                       { constToken $ TFloatLit $ BinRep nan }
-<0> "nan:arithmetic"                      { constToken $ TFloatLit $ BinRep nan }
+<0> "nan:canonical"                       { constToken $ TFloatLit $ NanRep Canonical }
+<0> "nan:arithmetic"                      { constToken $ TFloatLit $ NanRep Arithmetic }
 <0> $sign? @nanhex                        { parseNanSigned }
 <0> "inf"                                 { constToken $ TFloatLit $ BinRep inf }
 <0> "+inf"                                { constToken $ TFloatLit $ BinRep inf }
@@ -131,9 +131,13 @@ parseHexalSignedInt = token $ \(pos, _, s, _) len ->
 
 parseNanSigned :: AlexAction Lexeme
 parseNanSigned = token $ \(pos, _, s, _) len -> 
-    let (sign, slen) = parseSign s in
+    let (sign, slen) = case LBSUtf8.decode s of
+            Just ('-', _) -> (False, 1)
+            Just ('+', _) -> (True, 1)
+            otherwise -> (True, 0)
+    in
     let num = readHexFromPrefix (len - 6 - slen) $ LBSUtf8.drop (6 + slen) s in
-    Lexeme (Just pos) $ TFloatLit $ BinRep $ sign $ makeNaN $ fromIntegral num
+    Lexeme (Just pos) $ TFloatLit $ NanRep $ NanHex sign $ fromIntegral num
 
 parseDecimalSignedInt :: AlexAction Lexeme
 parseDecimalSignedInt = token $ \(pos, _, s, _) len ->
@@ -218,11 +222,23 @@ asFloat :: FloatRep -> Either String Float
 asFloat (BinRep d) = Right $ doubleToFloat d
 asFloat (HexRep s) = doubleToFloat <$> readHexFloat 128 "ffffff" s
 asFloat (DecRep s) = readDecFloat s
+asFloat (NanRep Canonical) = Right nan
+asFloat (NanRep Arithmetic) = Right nan
+asFloat (NanRep (NanHex isPos payload)) =
+    if payload >= 1 && payload < 2 ^ 23
+    then return $ doubleToFloat $ (if isPos then id else negate) $ makeNaN payload
+    else Left "constant out of range"
 
 asDouble :: FloatRep -> Either String Double
 asDouble (BinRep d) = Right d
 asDouble (HexRep s) = readHexFloat 1024 "fffffffffffff8" s
 asDouble (DecRep s) = readDecDouble s
+asDouble (NanRep Canonical) = Right nan
+asDouble (NanRep Arithmetic) = Right nan
+asDouble (NanRep (NanHex isPos payload)) =
+    if payload >= 1 && payload < 2 ^ 52
+    then return $ (if isPos then id else negate) $ makeNaN payload
+    else Left "constant out of range"
 
 startBlockComment :: AlexAction Lexeme
 startBlockComment _inp _len = do
@@ -300,6 +316,13 @@ data FloatRep
     = BinRep Double
     | DecRep String
     | HexRep String
+    | NanRep NaN
+    deriving (Show, Eq)
+
+data NaN
+    = Canonical
+    | Arithmetic
+    | NanHex Bool Word64
     deriving (Show, Eq)
 
 data Token = TKeyword LBS.ByteString
diff --git a/src/Language/Wasm/Parser.y b/src/Language/Wasm/Parser.y
index 5a47b85..1769444 100644
--- a/src/Language/Wasm/Parser.y
+++ b/src/Language/Wasm/Parser.y
@@ -1461,6 +1461,11 @@ desugarize fields = do
             else (TypeDef Nothing funcType) : defs
         matchTypeUse defs _ = defs
 
+        nth :: Natural -> [a] -> Maybe a
+        nth 0 (x : xs) = Just x
+        nth n (_ : xs) = nth (n - 1) xs
+        nth _ _ = Nothing
+
         getTypeIndex :: [TypeDef] -> TypeUse -> Maybe Natural
         getTypeIndex defs (AnonimousTypeUse funcType) =
             fromIntegral <$> findIndex (matchTypeFunc funcType) defs
@@ -1471,7 +1476,8 @@ desugarize fields = do
         getTypeIndex defs (IndexedTypeUse (Named ident) Nothing) =
             fromIntegral <$> findIndex (\(TypeDef i _) -> i == Just ident) defs
         getTypeIndex defs (IndexedTypeUse (Index n) (Just funcType)) = do
-            guard $ matchTypeFunc funcType $ defs !! fromIntegral n
+            def <- nth n defs
+            guard $ matchTypeFunc funcType def
             return n
         getTypeIndex defs (IndexedTypeUse (Index n) Nothing) = return n