5 files changed, 95 insertions, 120 deletions

diff --git a/src/ext/Util/ca/precomp.h b/src/ext/Util/ca/precomp.h
index efde32a6..9456c691 100644
--- a/src/ext/Util/ca/precomp.h
+++ b/src/ext/Util/ca/precomp.h
@@ -15,7 +15,7 @@
 
 #include <msxml2.h>
 #include <Iads.h>
-#include <activeds.h> 
+#include <activeds.h>
 #include <lm.h>        // NetApi32.lib
 #include <Ntsecapi.h>
 #include <Dsgetdc.h>
@@ -50,5 +50,6 @@
 #include "scauser.h"
 #include "scasmb.h"
 #include "scasmbexec.h"
+#include "utilca.h"
 
 #include "..\..\caDecor.h"
diff --git a/src/ext/Util/ca/scaexec.cpp b/src/ext/Util/ca/scaexec.cpp
index 5119bc11..8579b8bb 100644
--- a/src/ext/Util/ca/scaexec.cpp
+++ b/src/ext/Util/ca/scaexec.cpp
@@ -613,8 +613,7 @@ static HRESULT RemoveUserInternal(
     LPWSTR pwz = NULL;
     LPWSTR pwzGroup = NULL;
     LPWSTR pwzGroupDomain = NULL;
-    LPCWSTR wz = NULL;
-    PDOMAIN_CONTROLLER_INFOW pDomainControllerInfo = NULL;
+    LPWSTR pwzDomainName = NULL;
 
     //
     // Remove the logon as service privilege.
@@ -644,30 +643,10 @@ static HRESULT RemoveUserInternal(
     //
     if (!(SCAU_DONT_CREATE_USER & iAttributes))
     {
-        if (wzDomain && *wzDomain)
-        {
-            er = ::DsGetDcNameW(NULL, (LPCWSTR)wzDomain, NULL, NULL, NULL, &pDomainControllerInfo);
-            if (RPC_S_SERVER_UNAVAILABLE == er)
-            {
-                // MSDN says, if we get the above error code, try again with the "DS_FORCE_REDISCOVERY" flag
-                er = ::DsGetDcNameW(NULL, (LPCWSTR)wzDomain, NULL, NULL, DS_FORCE_REDISCOVERY, &pDomainControllerInfo);
-            }
-            if (ERROR_SUCCESS == er)
-            {
-                if (2 <= wcslen(pDomainControllerInfo->DomainControllerName))
-                {
-                    wz = pDomainControllerInfo->DomainControllerName + 2; // Add 2 so that we don't get the \\ prefix.
-                                                                          // Pass the entire string if it is too short
-                                                                          // to have a \\ prefix.
-                }
-            }
-            else
-            {
-                wz = wzDomain;
-            }
-        }
+        hr = GetDomainFromServerName(&pwzDomainName, wzDomain, 0);
+        ExitOnFailure(hr, "Failed to get domain to remove user from server name: %ls", wzDomain);
 
-        er = ::NetUserDel(wz, wzName);
+        er = ::NetUserDel(pwzDomainName, wzName);
         if (NERR_UserNotFound == er)
         {
             er = NERR_Success;
@@ -707,52 +686,13 @@ static HRESULT RemoveUserInternal(
     }
 
 LExit:
-    if (pDomainControllerInfo)
-    {
-        ::NetApiBufferFree(static_cast<LPVOID>(pDomainControllerInfo));
-    }
+    ReleaseStr(pwzDomainName);
+    ReleaseStr(pwzGroupDomain);
+    ReleaseStr(pwzGroup);
 
     return hr;
 }
 
-static void GetServerName(LPWSTR pwzDomain, LPWSTR* ppwzServerName)
-{
-    DWORD er = ERROR_SUCCESS;
-    PDOMAIN_CONTROLLER_INFOW pDomainControllerInfo = NULL;
-
-    if (pwzDomain && *pwzDomain)
-    {
-        er = ::DsGetDcNameW(NULL, (LPCWSTR)pwzDomain, NULL, NULL, NULL, &pDomainControllerInfo);
-        if (RPC_S_SERVER_UNAVAILABLE == er)
-        {
-            // MSDN says, if we get the above error code, try again with the "DS_FORCE_REDISCOVERY" flag
-            er = ::DsGetDcNameW(NULL, (LPCWSTR)pwzDomain, NULL, NULL, DS_FORCE_REDISCOVERY, &pDomainControllerInfo);
-        }
-
-        if (ERROR_SUCCESS == er && pDomainControllerInfo->DomainControllerName)
-        {
-            // Skip the \\ prefix if present.
-            if ('\\' == *pDomainControllerInfo->DomainControllerName && '\\' == *pDomainControllerInfo->DomainControllerName + 1)
-            {
-                *ppwzServerName = pDomainControllerInfo->DomainControllerName + 2;
-            }
-            else
-            {
-                *ppwzServerName = pDomainControllerInfo->DomainControllerName;
-            }
-        }
-        else
-        {
-            *ppwzServerName = pwzDomain;
-        }
-    }
-
-    if (pDomainControllerInfo)
-    {
-        ::NetApiBufferFree((LPVOID)pDomainControllerInfo);
-    }
-}
-
 /********************************************************************
  CreateUser - CUSTOM ACTION ENTRY POINT for creating users
 
@@ -776,6 +716,7 @@ extern "C" UINT __stdcall CreateUser(
     LPWSTR pwzPassword = NULL;
     LPWSTR pwzGroup = NULL;
     LPWSTR pwzGroupDomain = NULL;
+    LPWSTR pwzDomainName = NULL;
     int iAttributes = 0;
     BOOL fInitializedCom = FALSE;
 
@@ -786,7 +727,6 @@ extern "C" UINT __stdcall CreateUser(
     USER_INFO_1 userInfo1;
     USER_INFO_1* pUserInfo1 = NULL;
     DWORD dw;
-    LPWSTR pwzServerName = NULL;
 
     hr = WcaInitialize(hInstall, "CreateUser");
     ExitOnFailure(hr, "failed to initialize");
@@ -845,9 +785,10 @@ extern "C" UINT __stdcall CreateUser(
         //
         // Create the User
         //
-        GetServerName(pwzDomain, &pwzServerName);
+        hr = GetDomainFromServerName(&pwzDomainName, pwzDomain, 0);
+        ExitOnFailure(hr, "Failed to get domain from server name: %ls", pwzDomain);
 
-        er = ::NetUserAdd(pwzServerName, 1, reinterpret_cast<LPBYTE>(pUserInfo1), &dw);
+        er = ::NetUserAdd(pwzDomainName, 1, reinterpret_cast<LPBYTE>(pUserInfo1), &dw);
         if (NERR_UserExists == er)
         {
             if (SCAU_FAIL_IF_EXISTS & iAttributes)
@@ -862,7 +803,7 @@ extern "C" UINT __stdcall CreateUser(
             if (SCAU_UPDATE_IF_EXISTS & iAttributes)
             {
                 pUserInfo1 = NULL;
-                er = ::NetUserGetInfo(pwzServerName, pwzName, 1, reinterpret_cast<LPBYTE*>(&pUserInfo1));
+                er = ::NetUserGetInfo(pwzDomainName, pwzName, 1, reinterpret_cast<LPBYTE*>(&pUserInfo1));
                 if (ERROR_SUCCESS == er)
                 {
                     // There is no rollback scheduled if the key is empty.
@@ -922,28 +863,28 @@ extern "C" UINT __stdcall CreateUser(
 
                 if (ERROR_SUCCESS == er)
                 {
-                    hr = SetUserPassword(pwzServerName, pwzName, pwzPassword);
+                    hr = SetUserPassword(pwzDomainName, pwzName, pwzPassword);
                     if (FAILED(hr))
                     {
-                        WcaLogError(hr, "failed to set user password for user %ls\\%ls, continuing anyway.", pwzServerName, pwzName);
+                        WcaLogError(hr, "failed to set user password for user %ls\\%ls, continuing anyway.", pwzDomainName, pwzName);
                         hr = S_OK;
                     }
 
                     if (SCAU_REMOVE_COMMENT & iAttributes)
                     {
-                        hr = SetUserComment(pwzServerName, pwzName, L"");
+                        hr = SetUserComment(pwzDomainName, pwzName, L"");
                         if (FAILED(hr))
                         {
-                            WcaLogError(hr, "failed to clear user comment for user %ls\\%ls, continuing anyway.", pwzServerName, pwzName);
+                            WcaLogError(hr, "failed to clear user comment for user %ls\\%ls, continuing anyway.", pwzDomainName, pwzName);
                             hr = S_OK;
                         }
                     }
                     else if (pwzComment && *pwzComment)
                     {
-                        hr = SetUserComment(pwzServerName, pwzName, pwzComment);
+                        hr = SetUserComment(pwzDomainName, pwzName, pwzComment);
                         if (FAILED(hr))
                         {
-                            WcaLogError(hr, "failed to set user comment to %ls for user %ls\\%ls, continuing anyway.", pwzComment, pwzServerName, pwzName);
+                            WcaLogError(hr, "failed to set user comment to %ls for user %ls\\%ls, continuing anyway.", pwzComment, pwzDomainName, pwzName);
                             hr = S_OK;
                         }
                     }
@@ -952,10 +893,10 @@ extern "C" UINT __stdcall CreateUser(
 
                     ApplyAttributes(iAttributes, &flags);
 
-                    hr = SetUserFlags(pwzServerName, pwzName, flags);
+                    hr = SetUserFlags(pwzDomainName, pwzName, flags);
                     if (FAILED(hr))
                     {
-                        WcaLogError(hr, "failed to set user flags for user %ls\\%ls, continuing anyway.", pwzServerName, pwzName);
+                        WcaLogError(hr, "failed to set user flags for user %ls\\%ls, continuing anyway.", pwzDomainName, pwzName);
                         hr = S_OK;
                     }
                 }
@@ -1018,6 +959,7 @@ LExit:
     ReleaseStr(pwzPassword);
     ReleaseStr(pwzGroup);
     ReleaseStr(pwzGroupDomain);
+    ReleaseStr(pwzDomainName)
 
     if (fInitializedCom)
     {
diff --git a/src/ext/Util/ca/scauser.cpp b/src/ext/Util/ca/scauser.cpp
index b643a842..79da155f 100644
--- a/src/ext/Util/ca/scauser.cpp
+++ b/src/ext/Util/ca/scauser.cpp
@@ -487,7 +487,7 @@ HRESULT ScaUserExecute(
 {
     HRESULT hr = S_OK;
     DWORD er = 0;
-    PDOMAIN_CONTROLLER_INFOW pDomainControllerInfo = NULL;
+    LPWSTR pwzDomainName = NULL;
 
     LPWSTR pwzBaseScriptKey = NULL;
     DWORD cScriptKey = 0;
@@ -518,36 +518,11 @@ HRESULT ScaUserExecute(
         ExitOnFailure(hr, "Failed to add user comment to custom action data: %ls", psu->wzComment);
 
         // Check to see if the user already exists since we have to be very careful when adding
-        // and removing users.  Note: MSDN says that it is safe to call these APIs from any
-        // user, so we should be safe calling it during immediate mode.
-        er = ::NetApiBufferAllocate(sizeof(USER_INFO_0), reinterpret_cast<LPVOID*>(&pUserInfo));
-        hr = HRESULT_FROM_WIN32(er);
-        ExitOnFailure(hr, "Failed to allocate memory to check existence of user: %ls", psu->wzName);
-
-        LPCWSTR wzDomain = psu->wzDomain;
-        if (wzDomain && *wzDomain)
-        {
-            er = ::DsGetDcNameW(NULL, wzDomain, NULL, NULL, NULL, &pDomainControllerInfo);
-            if (RPC_S_SERVER_UNAVAILABLE == er)
-            {
-                // MSDN says, if we get the above error code, try again with the "DS_FORCE_REDISCOVERY" flag
-                er = ::DsGetDcNameW(NULL, wzDomain, NULL, NULL, DS_FORCE_REDISCOVERY, &pDomainControllerInfo);
-            }
-            if (ERROR_SUCCESS == er && pDomainControllerInfo->DomainControllerName)
-            {
-                // If the  \\ prefix on the queried domain was present, skip it.
-                if ('\\' == *pDomainControllerInfo->DomainControllerName && '\\' == *pDomainControllerInfo->DomainControllerName + 1)
-                {
-                    wzDomain = pDomainControllerInfo->DomainControllerName + 2;
-                }
-                else
-                {
-                    wzDomain = pDomainControllerInfo->DomainControllerName;
-                }
-            }
-        }
+        // and removing users.
+        hr = GetDomainFromServerName(&pwzDomainName, psu->wzDomain, 0);
+        ExitOnFailure(hr, "Failed to get domain from server name: %ls", psu->wzDomain);
 
-        er = ::NetUserGetInfo(wzDomain, psu->wzName, 0, reinterpret_cast<LPBYTE*>(pUserInfo));
+        er = ::NetUserGetInfo(pwzDomainName, psu->wzName, 0, reinterpret_cast<LPBYTE*>(&pUserInfo));
         if (NERR_Success == er)
         {
             ueUserExists = USER_EXISTS_YES;
@@ -560,7 +535,7 @@ HRESULT ScaUserExecute(
         {
             ueUserExists = USER_EXISTS_INDETERMINATE;
             hr = HRESULT_FROM_WIN32(er);
-            WcaLog(LOGMSG_VERBOSE, "Failed to check existence of domain: %ls, user: %ls (error code 0x%x) - continuing", wzDomain, psu->wzName, hr);
+            WcaLog(LOGMSG_VERBOSE, "Failed to check existence of domain: %ls, user: %ls (error code 0x%x) - continuing", pwzDomainName, psu->wzName, hr);
             hr = S_OK;
             er = ERROR_SUCCESS;
         }
@@ -685,11 +660,6 @@ HRESULT ScaUserExecute(
             ::NetApiBufferFree(static_cast<LPVOID>(pUserInfo));
             pUserInfo = NULL;
         }
-        if (pDomainControllerInfo)
-        {
-            ::NetApiBufferFree(static_cast<LPVOID>(pDomainControllerInfo));
-            pDomainControllerInfo = NULL;
-        }
     }
 
 LExit:
@@ -697,14 +667,12 @@ LExit:
     ReleaseStr(pwzScriptKey);
     ReleaseStr(pwzActionData);
     ReleaseStr(pwzRollbackData);
+    ReleaseStr(pwzDomainName);
+
     if (pUserInfo)
     {
         ::NetApiBufferFree(static_cast<LPVOID>(pUserInfo));
     }
-    if (pDomainControllerInfo)
-    {
-        ::NetApiBufferFree(static_cast<LPVOID>(pDomainControllerInfo));
-    }
 
     return hr;
 }
diff --git a/src/ext/Util/ca/utilca.cpp b/src/ext/Util/ca/utilca.cpp
index 37664a1c..d41f00c2 100644
--- a/src/ext/Util/ca/utilca.cpp
+++ b/src/ext/Util/ca/utilca.cpp
@@ -1,3 +1,59 @@
 // Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
 
 #include "precomp.h"
+
+HRESULT GetDomainFromServerName(
+    __deref_out_z LPWSTR* ppwzDomainName,
+    __in_z LPCWSTR wzServerName,
+    __in DWORD dwFlags
+    )
+{
+    HRESULT hr = S_OK;
+    DWORD er = ERROR_SUCCESS;
+    PDOMAIN_CONTROLLER_INFOW pDomainControllerInfo = NULL;
+    LPCWSTR wz = wzServerName ? wzServerName : L""; // initialize the domain to the provided server name (or empty string).
+
+    // If the server name was not empty, try to get the domain name out of it.
+    if (*wz)
+    {
+        er = ::DsGetDcNameW(NULL, wz, NULL, NULL, dwFlags, &pDomainControllerInfo);
+        if (RPC_S_SERVER_UNAVAILABLE == er)
+        {
+            // MSDN says, if we get the above error code, try again with the "DS_FORCE_REDISCOVERY" flag.
+            er = ::DsGetDcNameW(NULL, wz, NULL, NULL, dwFlags | DS_FORCE_REDISCOVERY, &pDomainControllerInfo);
+        }
+        ExitOnWin32Error(er, hr, "Could not get domain name from server name: %ls", wz);
+
+        if (pDomainControllerInfo->DomainControllerName)
+        {
+            // Skip the \\ prefix if present.
+            if ('\\' == *pDomainControllerInfo->DomainControllerName && '\\' == *(pDomainControllerInfo->DomainControllerName + 1))
+            {
+                wz = pDomainControllerInfo->DomainControllerName + 2;
+            }
+            else
+            {
+                wz = pDomainControllerInfo->DomainControllerName;
+            }
+        }
+    }
+
+LExit:
+    // Note: we overwrite the error code here as failure to contact domain controller above is not a fatal error.
+    if (wz && *wz)
+    {
+        hr = StrAllocString(ppwzDomainName, wz, 0);
+    }
+    else // return NULL the server name ended up empty.
+    {
+        ReleaseNullStr(*ppwzDomainName);
+        hr = S_OK;
+    }
+
+    if (pDomainControllerInfo)
+    {
+        ::NetApiBufferFree((LPVOID)pDomainControllerInfo);
+    }
+
+    return hr;
+}
diff --git a/src/ext/Util/ca/utilca.h b/src/ext/Util/ca/utilca.h
new file mode 100644
index 00000000..97c54561
--- /dev/null
+++ b/src/ext/Util/ca/utilca.h
@@ -0,0 +1,8 @@
+#pragma once
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+HRESULT GetDomainFromServerName(
+    __deref_out_z LPWSTR* ppwzDomainName,
+    __in_z LPCWSTR wzServerName,
+    __in DWORD dwFlags
+    );