CString: Introduce more creation options

Signed-off-by: Dmitry Fleytman <[email protected]>
xiayibing · Mar 10, 2015 · a1affb2 · a1affb2
1 parent 189d479
commit a1affb2
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 15 deletions.
diff --git a/UsbDk/UsbDkUtil.cpp b/UsbDk/UsbDkUtil.cpp
@@ -72,31 +72,25 @@ NTSTATUS CString::Create(NTSTRSAFE_PCWSTR String)
         return status;
     }
 
-    m_String.MaximumLength = StringHolder.MaximumLength;
+    return Create(&StringHolder);
+}
+
+NTSTATUS CString::Create(PCUNICODE_STRING String)
+{
+    m_String.MaximumLength = String->MaximumLength;
     m_String.Buffer = static_cast<PWCH>(ExAllocatePoolWithTag(NonPagedPool,
-        StringHolder.MaximumLength,
+        String->MaximumLength,
         'SUHR'));
 
     if (m_String.Buffer == nullptr)
     {
         return STATUS_INSUFFICIENT_RESOURCES;
     }
 
-    RtlCopyUnicodeString(&m_String, &StringHolder);
+    RtlCopyUnicodeString(&m_String, String);
     return STATUS_SUCCESS;
 }
 
-NTSTATUS CString::Create(NTSTRSAFE_PCWSTR Prefix, ULONG Postfix)
-{
-    auto status = Create(Prefix);
-    if (!NT_SUCCESS(status))
-    {
-        return status;
-    }
-
-    return Append(Postfix);
-}
-
 NTSTATUS CString::Append(PCUNICODE_STRING String)
 {
     auto status = Resize(m_String.Length + String->Length);

diff --git a/UsbDk/UsbDkUtil.h b/UsbDk/UsbDkUtil.h
@@ -433,7 +433,20 @@ class CString : public CStringBase
 {
 public:
     NTSTATUS Create(NTSTRSAFE_PCWSTR String);
-    NTSTATUS Create(NTSTRSAFE_PCWSTR Prefix, ULONG Postfix);
+    NTSTATUS Create(PCUNICODE_STRING String);
+
+    template <typename TPrefix, typename TPostfix>
+    NTSTATUS Create(const TPrefix &Prefix, const TPostfix &Postfix)
+    {
+        auto status = Create(Prefix);
+        if (!NT_SUCCESS(status))
+        {
+            return status;
+        }
+
+        return Append(Postfix);
+    }
+
     NTSTATUS Append(PCUNICODE_STRING String);
     NTSTATUS Append(ULONG Num, ULONG Base = 10);
     void Destroy();