From 8f2fcbeebc61787975e9eebd996ca439cddbd233 Mon Sep 17 00:00:00 2001
From: Bruno Oliveira <bruno@esss.com.br>
Date: Fri, 18 Jan 2019 14:59:48 -0200
Subject: [PATCH] Minor improvements to C++ code

* Use a std::wstring directly instead of a temporary buffer
* Pass filename by reference in `load_impls_from_library`
---
 hookman/hookman_generator.py                | 18 ++++++++++--------
 tests/test_hookman_generator/HookCaller.hpp | 18 ++++++++++--------
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/hookman/hookman_generator.py b/hookman/hookman_generator.py
index b8c23ef..174de46 100644
--- a/hookman/hookman_generator.py
+++ b/hookman/hookman_generator.py
@@ -326,11 +326,9 @@ def _hook_caller_hpp_content(self) -> str:
             "#define _H_HOOKMAN_HOOK_CALLER",
             "",
             "#include <functional>",
-            "#include <memory>",
             "#include <stdexcept>",
             "#include <string>",
             "#include <vector>",
-            "#include <iostream>",
             "",
             "#ifdef _WIN32",
             f"    #include <windows.h>",
@@ -624,7 +622,7 @@ def _generate_windows_body(hooks):
 
     # generate load_impls_from_library()
     result += [
-        f"    void load_impls_from_library(const std::string utf8_filename) {{",
+        f"    void load_impls_from_library(const std::string& utf8_filename) {{",
         f'        std::wstring w_filename = utf8_to_wstring(utf8_filename);',
         f'        auto handle = LoadLibraryW(w_filename.c_str());',
         f'        if (handle == NULL) {{',
@@ -648,11 +646,15 @@ def _generate_windows_body(hooks):
         "",
         "",
         "private:",
-        f"    std::wstring utf8_to_wstring(const std::string &s) {{",
+        f"    std::wstring utf8_to_wstring(const std::string& s) {{",
         f"        int flags = 0;",
         f"        int required_size = MultiByteToWideChar(CP_UTF8, flags, s.c_str(), -1, nullptr, 0);",
-        f"        auto buffer = std::make_unique<WCHAR[]>(required_size);",
-        f"        int err = MultiByteToWideChar(CP_UTF8, flags, s.c_str(), -1, buffer.get(), required_size);",
+        f"        std::wstring result;",
+        f"        if (required_size == 0) {{",
+        f"            return result;",
+        f"        }}",
+        f"        result.resize(required_size);",
+        f"        int err = MultiByteToWideChar(CP_UTF8, flags, s.c_str(), -1, &result[0], required_size);",
         f"        if (err == 0) {{",
         f"            // error handling: https://docs.microsoft.com/en-us/windows/desktop/api/stringapiset/nf-stringapiset-multibytetowidechar#return-value",
         f"            switch (GetLastError()) {{",
@@ -663,7 +665,7 @@ def _generate_windows_body(hooks):
         f"                default: throw std::runtime_error(\"Undefined error: \" + std::to_string(GetLastError()));",
         f"            }}",
         f"        }}",
-        f"        return std::wstring(buffer.get(), required_size);",
+        f"        return result;",
         f"    }}",
         f"",
         f"",
@@ -695,7 +697,7 @@ def _generate_linux_body(hooks):
 
     # generate load_impls_from_library()
     result += [
-        f"    void load_impls_from_library(const std::string utf8_filename) {{",
+        f"    void load_impls_from_library(const std::string& utf8_filename) {{",
         f'        auto handle = dlopen(utf8_filename.c_str(), RTLD_LAZY);',
         f'        if (handle == nullptr) {{',
         f'            throw std::runtime_error("Error loading library " + utf8_filename + ": dlopen failed");',
diff --git a/tests/test_hookman_generator/HookCaller.hpp b/tests/test_hookman_generator/HookCaller.hpp
index aff2a94..9186ff4 100644
--- a/tests/test_hookman_generator/HookCaller.hpp
+++ b/tests/test_hookman_generator/HookCaller.hpp
@@ -2,11 +2,9 @@
 #define _H_HOOKMAN_HOOK_CALLER
 
 #include <functional>
-#include <memory>
 #include <stdexcept>
 #include <string>
 #include <vector>
-#include <iostream>
 
 #ifdef _WIN32
     #include <windows.h>
@@ -49,7 +47,7 @@ class HookCaller {
             FreeLibrary(handle);
         }
     }
-    void load_impls_from_library(const std::string utf8_filename) {
+    void load_impls_from_library(const std::string& utf8_filename) {
         std::wstring w_filename = utf8_to_wstring(utf8_filename);
         auto handle = LoadLibraryW(w_filename.c_str());
         if (handle == NULL) {
@@ -71,11 +69,15 @@ class HookCaller {
 
 
 private:
-    std::wstring utf8_to_wstring(const std::string &s) {
+    std::wstring utf8_to_wstring(const std::string& s) {
         int flags = 0;
         int required_size = MultiByteToWideChar(CP_UTF8, flags, s.c_str(), -1, nullptr, 0);
-        auto buffer = std::make_unique<WCHAR[]>(required_size);
-        int err = MultiByteToWideChar(CP_UTF8, flags, s.c_str(), -1, buffer.get(), required_size);
+        std::wstring result;
+        if (required_size == 0) {
+            return result;
+        }
+        result.resize(required_size);
+        int err = MultiByteToWideChar(CP_UTF8, flags, s.c_str(), -1, &result[0], required_size);
         if (err == 0) {
             // error handling: https://docs.microsoft.com/en-us/windows/desktop/api/stringapiset/nf-stringapiset-multibytetowidechar#return-value
             switch (GetLastError()) {
@@ -86,7 +88,7 @@ class HookCaller {
                 default: throw std::runtime_error("Undefined error: " + std::to_string(GetLastError()));
             }
         }
-        return std::wstring(buffer.get(), required_size);
+        return result;
     }
 
 
@@ -104,7 +106,7 @@ class HookCaller {
             dlclose(handle);
         }
     }
-    void load_impls_from_library(const std::string utf8_filename) {
+    void load_impls_from_library(const std::string& utf8_filename) {
         auto handle = dlopen(utf8_filename.c_str(), RTLD_LAZY);
         if (handle == nullptr) {
             throw std::runtime_error("Error loading library " + utf8_filename + ": dlopen failed");