Add windows specific drive sanitization while normalizing paths

composewell · Dec 6, 2024 · 94f985c · 94f985c
1 parent b81862d
commit 94f985c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 8 deletions.
diff --git a/core/src/Streamly/Internal/FileSystem/Path/Common.hs b/core/src/Streamly/Internal/FileSystem/Path/Common.hs
@@ -378,21 +378,21 @@ normalize os arr =
     if arrElemLen == 1
     then arr
     else Array.unsafeFreeze $ unsafePerformIO $ do
-        let workSliceMut = Array.unsafeThaw workSlice
-            workSliceStream = MutArray.read workSliceMut
+        let workSliceStream = MutArray.read workSliceMut
         (mid :: MutArray.MutArray a) <-
             Stream.indexOnSuffix (== sepElem) workSliceStream
                 & Stream.filter (not . shouldFilterOut)
-                & fmap (\(i, len) -> getSliceWithSepSuffix i len workSliceMut)
+                & Stream.mapM (\(i, len) -> getSliceWithSepSuffix i len)
                 & Stream.fold (Fold.foldlM' MutArray.unsafeSplice initBufferM)
         if startsWithDotSlash && MutArray.length mid == 0
-        then MutArray.fromListN 2 [fstElem, sndElem]
+        then MutArray.fromListN 2 [dotElem, sepElem]
         else pure mid
 
     where
 
-    sepElem = fromIntegral (ord (primarySeparator os))
-    dotElem = fromIntegral (ord '.')
+    (sepElem :: a) = fromIntegral (ord (primarySeparator os))
+    (dotElem :: a) = fromIntegral (ord '.')
+    (driveSepElem :: a) = fromIntegral (ord ':')
     arrElemLen = Array.length arr
 
     fstElem = Array.getIndexUnsafe 0 arr
@@ -405,15 +405,34 @@ normalize os arr =
         | startsWithSep = Array.getSliceUnsafe 1 (arrElemLen - 1) arr
         | startsWithDotSlash = Array.getSliceUnsafe 2 (arrElemLen - 2) arr
         | otherwise = arr
+    workSliceMut = Array.unsafeThaw workSlice
     workSliceElemLen = Array.length workSlice
 
     shouldFilterOut (off, len) =
         len == 0 ||
             (len == 1 && Array.getIndexUnsafe off workSlice == dotElem)
 
     getSliceWithSepSuffix i len
-        | i + len == workSliceElemLen = MutArray.unsafeGetSlice i len
-    getSliceWithSepSuffix i len = MutArray.unsafeGetSlice i (len + 1)
+        | i + len == workSliceElemLen =
+            sanitizeDrive os $ MutArray.unsafeGetSlice i len workSliceMut
+    getSliceWithSepSuffix i len =
+        sanitizeDrive os $ MutArray.unsafeGetSlice i (len + 1) workSliceMut
+
+    unsafeCapASCII x | x >= 97 = x - 32
+    unsafeCapASCII x = x
+
+    -- This is windows specific
+    sanitizeDrive Posix val = pure val
+    sanitizeDrive Windows val =
+        if MutArray.length val == 2
+        then do
+            b <- MutArray.unsafeGetIndex 1 val
+            if b == driveSepElem
+            then do
+                a <- MutArray.unsafeGetIndex 0 val
+                MutArray.fromListN 2 [unsafeCapASCII a, driveSepElem]
+            else pure val
+        else pure val
 
     initBufferM = do
         (newArr :: MutArray.MutArray a) <- MutArray.emptyOf arrElemLen

diff --git a/test/Streamly/Test/FileSystem/Path.hs b/test/Streamly/Test/FileSystem/Path.hs
@@ -32,6 +32,9 @@ main =
       describe moduleName $ do
         describe "normalize" $ do
             -- Primarily for Windows
+            testNormalize "C:\\"
+            testNormalize "C:"
+            testNormalize "\\\\?\\c:\\"
             testNormalize "c:\\file/bob\\"
             testNormalize "c:\\"
             testNormalize "c:\\\\\\\\"