pay attention to '\0's at the end

1 files changed, 85 insertions, 16 deletions

diff --git a/src/Windows/Registry.hs b/src/Windows/Registry.hs
index e6b37de..c584b90 100644
--- a/src/Windows/Registry.hs
+++ b/src/Windows/Registry.hs
@@ -33,6 +33,8 @@ module Windows.Registry
 import           Data.Bits             ((.|.))
 import qualified Data.ByteString       as B
 import           Data.List             (intercalate)
+import           Data.Maybe            (fromJust)
+import           Data.Tuple            (swap)
 import qualified Data.Text             as T
 import           Data.Text.Encoding    (decodeUtf16LE, encodeUtf16LE)
 import           Control.Exception     (bracket)
@@ -82,14 +84,52 @@ instance Show KeyPath where
     show (KeyPath root path) = intercalate pathSep $ show root : path
 
 type ValueName = String
-type ValueType = WinAPI.DWORD
+
+data ValueType = TypeNone
+               | TypeBinary
+               | TypeDWord
+               | TypeDWordBE
+               | TypeQWord
+               | TypeString
+               | TypeMultiString
+               | TypeExpandableString
+               | TypeLink
+               deriving (Eq, Show)
+
+instance Enum ValueType where
+    fromEnum = fromJust . flip lookup valueTypeTable
+    toEnum = fromJust . flip lookup (map swap valueTypeTable)
+
+valueTypeTable :: [(ValueType, Int)]
+valueTypeTable =
+    [ (TypeNone,     0)
+    , (TypeBinary,   3)
+    , (TypeDWord,    4)
+    , (TypeDWordBE,  5)
+    , (TypeQWord,   11)
+    , (TypeString,   1)
+    , (TypeMultiString,      7)
+    , (TypeExpandableString, 2)
+    , (TypeLink,     6)
+    ]
+
 type ValueData = (ValueType, B.ByteString)
 
 encodeString :: String -> B.ByteString
-encodeString = encodeUtf16LE . T.pack
+encodeString str = encodeUtf16LE addLastZero
+  where
+    text = T.pack str
+    addLastZero | T.null text = text
+                | T.last text == '\0' = text
+                | otherwise = T.snoc text '\0'
 
 decodeString :: ValueData -> String
-decodeString (_, valueData) = T.unpack . decodeUtf16LE $ valueData
+decodeString (_, bytes) = T.unpack dropLastZero
+  where
+    text = decodeUtf16LE bytes
+    dropLastZero | T.null text = text
+                 | T.last text == '\0' = T.init text
+                 | otherwise = text
 
 openCloseCatch :: IsKeyPath a => a -> (Handle -> IO b) -> IO (Either IOError b)
 openCloseCatch keyPath f = catchIOError (fmap Right openClose) $ return . Left
@@ -119,28 +159,57 @@ queryValue keyPath valueName =
                 WinAPI.failUnlessSuccess "RegQueryValueExW" $ c_RegQueryValueEx keyHandlePtr valueNamePtr WinAPI.nullPtr dataTypePtr bufferPtr dataSizePtr
                 buffer <- peekArray dataSize bufferPtr
                 dataType <- peek dataTypePtr
-                return (dataType, B.pack buffer)
-
-getValue :: IsKeyPath a => a -> ValueName -> [ValueType] -> IO (Either IOError ValueData)
-getValue keyPath valueName allowedTypes =
+                return (toEnum $ fromIntegral dataType, B.pack buffer)
+
+data GetValueFlag = RestrictAny
+                  | RestrictNone
+                  | RestrictBinary
+                  | RestrictDWord
+                  | RestrictQWord
+                  | RestrictString
+                  | RestrictMultiString
+                  | RestrictExpandableString
+                  | DoNotExpand
+                  deriving (Eq, Show)
+
+instance Enum GetValueFlag where
+    fromEnum = fromJust . flip lookup getValueFlagsTable
+    toEnum = fromJust . flip lookup (map swap getValueFlagsTable)
+
+getValueFlagsTable :: [(GetValueFlag, Int)]
+getValueFlagsTable =
+    [ (RestrictAny,    0x0000ffff)
+    , (RestrictNone,   0x00000001)
+    , (RestrictBinary, 0x00000008)
+    , (RestrictDWord,  0x00000010)
+    , (RestrictQWord,  0x00000040)
+    , (RestrictString, 0x00000002)
+    , (RestrictMultiString,      0x00000020)
+    , (RestrictExpandableString, 0x00000004)
+    , (DoNotExpand,    0x10000000)
+    ]
+
+getValue :: IsKeyPath a => a -> ValueName -> [GetValueFlag] -> IO (Either IOError ValueData)
+getValue keyPath valueName flagList =
     openCloseCatch keyPath $ \keyHandle ->
     withForeignPtr keyHandle $ \keyHandlePtr ->
     WinAPI.withTString valueName $ \valueNamePtr ->
     alloca $ \dataSizePtr -> do
         poke dataSizePtr 0
-        let flags = foldr (.|.) 0 allowedTypes
+        let flags = fromIntegral $ foldr (.|.) 0 $ map fromEnum flagList
         WinAPI.failUnlessSuccess "RegGetValueW" $ c_RegGetValue keyHandlePtr WinAPI.nullPtr valueNamePtr flags WinAPI.nullPtr WinAPI.nullPtr dataSizePtr
-        dataSize <- fromIntegral <$> peek dataSizePtr
+        bufferCapacity <- fromIntegral <$> peek dataSizePtr
         alloca $ \dataTypePtr ->
-            allocaBytes dataSize $ \bufferPtr -> do
+            allocaBytes bufferCapacity $ \bufferPtr -> do
                 WinAPI.failUnlessSuccess "RegGetValueW" $ c_RegGetValue keyHandlePtr WinAPI.nullPtr valueNamePtr flags dataTypePtr bufferPtr dataSizePtr
-                buffer <- peekArray dataSize bufferPtr
+                bufferSize <- fromIntegral <$> peek dataSizePtr
+                buffer <- peekArray bufferSize bufferPtr
                 dataType <- peek dataTypePtr
-                return (dataType, B.pack buffer)
+                return (toEnum $ fromIntegral dataType, B.pack buffer)
 
 getExpandedString :: IsKeyPath a => a -> ValueName -> IO (Either IOError String)
 getExpandedString keyPath valueName = do
-    valueData <- getValue keyPath valueName [WinAPI.rEG_SZ, WinAPI.rEG_EXPAND_SZ]
+    valueData <- getValue keyPath valueName [RestrictString, RestrictExpandableString]
     return $ fmap decodeString valueData
 
 setValue :: IsKeyPath a => a -> ValueName -> ValueData -> IO (Either IOError ())
@@ -152,15 +221,15 @@ setValue keyPath valueName (valueType, valueData) =
         let dataSize = B.length valueData
         allocaBytes dataSize $ \bufferPtr -> do
             pokeArray bufferPtr buffer
-            WinAPI.failUnlessSuccess "RegSetValueExW" $ c_RegSetValueEx keyHandlePtr valueNamePtr 0 valueType bufferPtr (fromIntegral dataSize)
+            WinAPI.failUnlessSuccess "RegSetValueExW" $ c_RegSetValueEx keyHandlePtr valueNamePtr 0 (fromIntegral $ fromEnum valueType) bufferPtr (fromIntegral dataSize)
 
 setString :: IsKeyPath a => a -> ValueName -> String -> IO (Either IOError ())
 setString keyPath valueName valueData =
-    setValue keyPath valueName (WinAPI.rEG_SZ, encodeString valueData)
+    setValue keyPath valueName (TypeString, encodeString valueData)
 
 setExpandableString :: IsKeyPath a => a -> ValueName -> String -> IO (Either IOError ())
 setExpandableString keyPath valueName valueData =
-    setValue keyPath valueName (WinAPI.rEG_EXPAND_SZ, encodeString valueData)
+    setValue keyPath valueName (TypeExpandableString, encodeString valueData)
 
 deleteValue :: IsKeyPath a => a -> ValueName -> IO (Either IOError ())
 deleteValue keyPath valueName =