Synchronize ByteString and String modules

System/Posix/Env/ByteString.hsc

@@ -1,5 +1,5 @@
 {-# LANGUAGE CApiFFI #-}
-{-# LANGUAGE Safe #-}
+{-# LANGUAGE Trustworthy #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -21,23 +21,28 @@ module System.Posix.Env.ByteString (
         , getEnvDefault
         , getEnvironmentPrim
         , getEnvironment
+        , setEnvironment
         , putEnv
         , setEnv
-       , unsetEnv
+        , unsetEnv
+        , clearEnv
 
        -- * Program arguments
        , getArgs
 ) where
 
 #include "HsUnix.h"
 
+import Control.Monad
 import Foreign
 import Foreign.C
 import Data.Maybe       ( fromMaybe )
 
+import System.Posix.Env ( clearEnv )
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Char8 as BC
 import Data.ByteString (ByteString)
+import Data.ByteString.Internal (ByteString (PS), memcpy)
 
 -- |'getEnv' looks up a variable in the environment.
 
@@ -96,6 +101,18 @@ getEnvironment = do
       | BC.head y == '=' = (x,B.tail y)
       | otherwise       = error $ "getEnvironment: insane variable " ++ BC.unpack x
 
+-- |'setEnvironment' resets the entire environment to the given list of
+-- @(key,value)@ pairs.
+--
+-- @since 2.8.0.0
+setEnvironment ::
+  [(ByteString,ByteString)] {- ^ @[(key,value)]@ -} ->
+  IO ()
+setEnvironment env = do
+  clearEnv
+  forM_ env $ \(key,value) ->
+    setEnv key value True {-overwrite-}
+
 -- |The 'unsetEnv' function deletes all instances of the variable name
 -- from the environment.
 
@@ -116,15 +133,25 @@ foreign import capi unsafe "HsUnix.h unsetenv"
    c_unsetenv :: CString -> IO ()
 # endif
 #else
-unsetEnv name = putEnv (name ++ "=")
+unsetEnv name = putEnv (BC.snoc name '=')
 #endif
 
 -- |'putEnv' function takes an argument of the form @name=value@
 -- and is equivalent to @setEnv(key,value,True{-overwrite-})@.
 
 putEnv :: ByteString {- ^ "key=value" -} -> IO ()
-putEnv keyvalue = B.useAsCString keyvalue $ \s ->
-  throwErrnoIfMinus1_ "putenv" (c_putenv s)
+putEnv (PS fp o l) = withForeignPtr fp $ \p -> do
+  -- https://pubs.opengroup.org/onlinepubs/009696899/functions/putenv.html
+  --
+  -- "the string pointed to by string shall become part of the environment,
+  -- so altering the string shall change the environment. The space used by
+  -- string is no longer used once a new string which defines name is passed to putenv()."
+  --
+  -- hence we must not free the buffer
+  buf <- mallocBytes (l+1)
+  memcpy buf (p `plusPtr` o) l
+  pokeByteOff buf l (0::Word8)
+  throwErrnoIfMinus1_ "putenv" (c_putenv (castPtr buf))
 
 foreign import ccall unsafe "putenv"
    c_putenv :: CString -> IO CInt

tests/PutEnv001.hs

@@ -0,0 +1,32 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# OPTIONS_GHC -O0 -Wno-name-shadowing #-}
+
+module Main (main) where
+
+import Data.String ( fromString )
+import System.Mem
+import System.Posix.Env.ByteString
+import Test.Tasty
+import Test.Tasty.HUnit
+
+-- test regression of incorrect 'free': https://github.com/haskell/unix/issues/68#issue-170072591
+main :: IO ()
+main = do
+  putEnv "foo=bar"
+  defaultMain $ testGroup "All" [ test ]
+
+test :: TestTree
+test = testCase "putEnv" $ do
+  performMinorGC
+  env <- System.Posix.Env.ByteString.getEnv (fromString "foo")
+  performMinorGC
+  print env
+  env <- System.Posix.Env.ByteString.getEnv (fromString "foo")
+  performMinorGC
+  print env
+  env <- System.Posix.Env.ByteString.getEnv (fromString "foo")
+  performMinorGC
+  print env
+  env <- System.Posix.Env.ByteString.getEnv (fromString "foo")
+  print env
+  env @?= Just (fromString "bar")

unix.cabal

@@ -242,3 +242,11 @@ test-suite Terminal
     default-language: Haskell2010
     build-depends: base, unix, tasty-hunit
     ghc-options: -Wall
+
+test-suite PutEnv001
+    hs-source-dirs: tests
+    main-is: PutEnv001.hs
+    type: exitcode-stdio-1.0
+    default-language: Haskell2010
+    build-depends: base, unix, tasty, tasty-hunit
+    ghc-options: -Wall -with-rtsopts=-V0 -O0