Fix memory leak from #12664

This fixes the leak with `setProgArgv`. The problem was
that `setProgArgv` would not free the objects pointed
to by `prog_argc`, `prog_argv` when the globals were
changed resulting in a leak.

The only strictly necessary change is in `rts/RtsFlags.c`, but
the code in `System.Environment` was a bit confusing and not
exception safe, so I refactored it.

Test Plan: ./validate

Reviewers: simonmar, ezyang, austin, hvr, bgamari, erikd

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D2576

GHC Trac Issues: #12664

5 files changed, 37 insertions, 18 deletions

diff --git a/libraries/base/GHC/Foreign.hs b/libraries/base/GHC/Foreign.hs
index e8553d8061..7d2f915920 100644
--- a/libraries/base/GHC/Foreign.hs
+++ b/libraries/base/GHC/Foreign.hs
@@ -32,6 +32,7 @@ module GHC.Foreign (
     --
     withCString,
     withCStringLen,
+    withCStringsLen,
 
     charIsRepresentable,
   ) where
@@ -134,6 +135,23 @@ withCString enc s act = withEncodedCString enc True s $ \(cp, _sz) -> act cp
 withCStringLen         :: TextEncoding -> String -> (CStringLen -> IO a) -> IO a
 withCStringLen enc = withEncodedCString enc False
 
+-- | Marshal a list of Haskell strings into an array of NUL terminated C strings
+-- using temporary storage.
+--
+-- * the Haskell strings may /not/ contain any NUL characters
+--
+-- * the memory is freed when the subcomputation terminates (either
+--   normally or via an exception), so the pointer to the temporary
+--   storage must /not/ be used after this.
+--
+withCStringsLen :: TextEncoding
+                -> [String]
+                -> (Int -> Ptr CString -> IO a)
+                -> IO a
+withCStringsLen enc strs f = go [] strs
+  where
+  go cs (s:ss) = withCString enc s $ \c -> go (c:cs) ss
+  go cs [] = withArrayLen (reverse cs) f
 
 -- | Determines whether a character can be accurately encoded in a 'CString'.
 --
diff --git a/libraries/base/System/Environment.hs b/libraries/base/System/Environment.hs
index 242845aae6..d8b3e03be5 100644
--- a/libraries/base/System/Environment.hs
+++ b/libraries/base/System/Environment.hs
@@ -32,12 +32,14 @@ module System.Environment
 import Foreign
 import Foreign.C
 import System.IO.Error (mkIOError)
-import Control.Exception.Base (bracket, throwIO)
+import Control.Exception.Base (bracket_, throwIO)
+#ifdef mingw32_HOST_OS
+import Control.Exception.Base (bracket)
+#endif
 -- import GHC.IO
 import GHC.IO.Exception
 import GHC.IO.Encoding (getFileSystemEncoding)
 import qualified GHC.Foreign as GHC
-import Data.List
 import Control.Monad
 #ifdef mingw32_HOST_OS
 import GHC.Environment
@@ -369,25 +371,17 @@ withProgArgv :: [String] -> IO a -> IO a
 withProgArgv new_args act = do
   pName <- System.Environment.getProgName
   existing_args <- System.Environment.getArgs
-  bracket (setProgArgv new_args)
-          (\argv -> do _ <- setProgArgv (pName:existing_args)
-                       freeProgArgv argv)
-          (const act)
-
-freeProgArgv :: Ptr CString -> IO ()
-freeProgArgv argv = do
-  size <- lengthArray0 nullPtr argv
-  sequence_ [ peek (argv `advancePtr` i) >>= free
-            | i <- [size - 1, size - 2 .. 0]]
-  free argv
-
-setProgArgv :: [String] -> IO (Ptr CString)
+  bracket_ (setProgArgv new_args)
+           (setProgArgv (pName:existing_args))
+           act
+
+setProgArgv :: [String] -> IO ()
 setProgArgv argv = do
   enc <- getFileSystemEncoding
-  vs <- mapM (GHC.newCString enc) argv >>= newArray0 nullPtr
-  c_setProgArgv (genericLength argv) vs
-  return vs
+  GHC.withCStringsLen enc argv $ \len css ->
+    c_setProgArgv (fromIntegral len) css
 
+-- setProgArgv copies the arguments
 foreign import ccall unsafe "setProgArgv"
   c_setProgArgv  :: CInt -> Ptr CString -> IO ()
 
diff --git a/libraries/base/tests/IO/environment001.hs b/libraries/base/tests/IO/environment001.hs
index 11d7912cdd..1d7a5c1c4a 100644
--- a/libraries/base/tests/IO/environment001.hs
+++ b/libraries/base/tests/IO/environment001.hs
@@ -14,3 +14,7 @@ main = do
     [arg1] <- withArgs ["你好!"] getArgs
     putStrLn arg1
     putStrLn ("Test 3: " ++ show (length arg1))
+
+    args2 <- withArgs ["a", "b"] getArgs
+    print args2
+    putStrLn ("Test 4: " ++ show (length args2))
diff --git a/libraries/base/tests/IO/environment001.stdout b/libraries/base/tests/IO/environment001.stdout
index 2434d0c14d..2d32a83370 100644
--- a/libraries/base/tests/IO/environment001.stdout
+++ b/libraries/base/tests/IO/environment001.stdout
@@ -4,3 +4,5 @@ Test 1: 3
 Test 2: 1
 你好!
 Test 3: 3
+["a","b"]
+Test 4: 2
diff --git a/rts/RtsFlags.c b/rts/RtsFlags.c
index c994a0ce88..4bd544ee29 100644
--- a/rts/RtsFlags.c
+++ b/rts/RtsFlags.c
@@ -1943,6 +1943,7 @@ getProgArgv(int *argc, char **argv[])
 void
 setProgArgv(int argc, char *argv[])
 {
+    freeArgv(prog_argc,prog_argv);
     prog_argc = argc;
     prog_argv = copyArgv(argc,argv);
     setProgName(prog_argv);