From ea99c1d1ca40d64adb835fed4ff4165c81101ee0 Mon Sep 17 00:00:00 2001
From: Flamefire <Flamefire@users.noreply.github.com>
Date: Fri, 13 Dec 2019 20:30:55 +0100
Subject: [PATCH] Avoid memory leak in fix_args/fix_env

Throwing bad_alloc from vector::resize could leak memory
---
 include/boost/nowide/args.hpp | 39 ++++++++++++++++++++++++++---------
 1 file changed, 29 insertions(+), 10 deletions(-)

diff --git a/include/boost/nowide/args.hpp b/include/boost/nowide/args.hpp
index 95bfd3d..feb436d 100644
--- a/include/boost/nowide/args.hpp
+++ b/include/boost/nowide/args.hpp
@@ -86,26 +86,46 @@ namespace nowide {
                 *old_env_ptr_ = old_env_;
         }
     private:
+        class wargv_ptr
+        {
+            wchar_t **p;
+            int argc;
+            wargv_ptr(const wargv_ptr &); // Non-copyable
+        public:
+            explicit wargv_ptr(const wchar_t *cmd_line) : p(CommandLineToArgvW(GetCommandLineW(), &argc)) {}
+            ~wargv_ptr() { if(p) LocalFree(p); }
+            int size() const { return argc; }
+            operator bool() const { return p != NULL; }
+            const wchar_t* operator[](size_t i) const { return p[i]; }
+        };
+        class wenv_ptr
+        {
+            wchar_t *p;
+            wenv_ptr(const wenv_ptr &); // Non-copyable
+        public:
+            wenv_ptr() : p(GetEnvironmentStringsW()) {}
+            ~wenv_ptr() { if(p) FreeEnvironmentStringsW(p); }
+            operator const wchar_t*() const { return p; }
+        };
+
         void fix_args(int &argc,char **&argv)
         {
-            int wargc;
-            wchar_t **wargv = CommandLineToArgvW(GetCommandLineW(),&wargc);
+            const wargv_ptr wargv(GetCommandLineW());
             if(!wargv)
                 throw std::runtime_error("Could not get command line!");
-            args_.resize(wargc+1,0);
-            arg_values_.resize(wargc);
-            for(int i=0;i<wargc;i++) 
+            args_.resize(wargv.size()+1,0);
+            arg_values_.resize(wargv.size());
+            for(int i=0;i<wargv.size();i++) 
                 args_[i] = arg_values_[i].convert(wargv[i]);
-            argc = wargc;
+            argc = wargv.size();
             argv = &args_[0];
-            LocalFree(wargv);
         }
         void fix_env(char **&env)
         {
-            wchar_t *wstrings = GetEnvironmentStringsW();
+            const wenv_ptr wstrings;
             if(!wstrings)
                 throw std::runtime_error("Could not get environment strings!");
-            wchar_t *wstrings_end = 0;
+            const wchar_t *wstrings_end = 0;
             int count = 0;
             for(wstrings_end = wstrings;*wstrings_end;wstrings_end+=wcslen(wstrings_end)+1)
                     count++;
@@ -119,7 +139,6 @@ namespace nowide {
                 p+=strlen(p)+1;
             }
             env = &envp_[0];
-            FreeEnvironmentStringsW(wstrings);
         }
 
         std::vector<char *> args_;