stafprocess.cpp

Software Testing Automation Framework (STAF)的开发代码

                // Create an UNICODE environment block (contains environment 
                // variables for the specified user in a form that can be
                // passed directly to CreateProcessAsUser
                rc = sCreateEnvironmentBlockFunc(
                     reinterpret_cast<void**>(&environment), hUsrToken, FALSE);
            }

            BOOL unicodeEnvBlock = TRUE;

            if (rc != TRUE)
            {
                if (!errorFound)
                {
                    errorFound = TRUE;
                    systemErrorCode = GetLastError();
                    failedFunction = "CreateEnvironmentBlock";
                }
            }
            else
            {
                char *newEnvBuf = 0;
                
                if (startData->environment != 0)
                {
                    // Merge environment variables (specified on the PROCESS
                    // START request via the ENV parameters) that are in UTF-8
                    // map with the UNICODE environment block for the user.
                    // Resulting environment block is the current code page.
                    
                    unicodeEnvBlock = FALSE;
                    
                    // Find the end of the user's UNICODE environment block

                    int numWideChars = 0;

                    for (int i = 0; !numWideChars; i += 2)
                    {
                        if (environment[i]   == 0 && environment[i+1] == 0 && 
                            environment[i+2] == 0 && environment[i+3] == 0)
                        {
                            numWideChars = (i + 4) / 2;  // Get # of wide chars
                        }
                    }
                    
                    // Determine the # of bytes required for MultiByte block

                    int numBytes = WideCharToMultiByte(CP_UTF8, 0,
                                   reinterpret_cast<wchar_t *>(environment),
                                   numWideChars, 0, 0, NULL, NULL);

                    // Allocate new multi-byte environment block

                    char *newEnvBlock = new char[numBytes];

                    // Convert the UNICODE env block to MultiByte (UTF-8)

                    WideCharToMultiByte(CP_UTF8, 0,
                        reinterpret_cast<wchar_t *>(environment),
                        numWideChars, newEnvBlock, numBytes, NULL, NULL);
                    
                    // Break up the env block string (in UTF-8) into a map

                    STAFString envBlock(newEnvBlock, numBytes, 
                                        STAFString::kUTF8);

                    delete [] newEnvBlock;  // Delete since already copied

                    // Create a STAFProcessEnvMap from the env block string so
                    // that the key (the environment variable name) can be
                    // case-insensitive if designated by the operating system.

                    STAFProcessEnvMap envMap;
                    BOOL moreEnvVars = TRUE;

                    while (moreEnvVars)
                    {
                        unsigned int endPos = envBlock.find(kUTF8_NULL);
                        if (endPos == 0)
                        {
                            moreEnvVars = FALSE;
                        }
                        else
                        {
                            STAFString envStr(envBlock.subString(0, endPos));
                            unsigned int equalPos = envStr.find(kUTF8_EQUAL);
                            if (equalPos == 0)
                            {
                                moreEnvVars = FALSE;
                            }
                            else
                            {
                                STAFString envName(
                                    envStr.subString(0, equalPos));
                                STAFString envNameKey(envName.toUpperCase());
                                STAFString envValue;
                                envValue = envStr.subString(equalPos + 1);

                                envMap[envNameKey] = STAFProcessEnvData(
                                    envName, envValue);

                                if (endPos + 1 > numBytes)
                                    moreEnvVars = FALSE;
                                else
                                    envBlock = STAFString(
                                        envBlock.subString(endPos+1));
                            }
                        }
                    }

                    // Merge the env vars specified on the process start
                    // request (stored in startData->userEnvMap) into envMap

                    for (unsigned int userEnvIndex = 0;
                         userEnvIndex < startData->userEnvCount;
                         ++userEnvIndex)
                    {
                        STAFString userEnv(startData->userEnvList[userEnvIndex]);
                        unsigned int equalPos = userEnv.find(kUTF8_EQUAL);

                        if (equalPos != STAFString::kNPos)
                        {
                            STAFString aEnvName(userEnv.subString(0, equalPos));
                            STAFString aEnvNameKey(aEnvName.toUpperCase());
                            STAFString aEnvValue(
                                userEnv.subString(equalPos + 1));

                            // Retain original env variable name if exists
                            if (envMap.find(aEnvNameKey) != envMap.end())
                                envMap[aEnvNameKey].envValue = aEnvValue;
                            else
                                envMap[aEnvNameKey] = STAFProcessEnvData(
                                    aEnvName, aEnvValue);
                        }
                    }

                    // Iterate thru the merged map (envMap): combine entries
                    // back into Name=Value form, get the current code page
                    // representation and figure out how big a buffer we need.
                    
                    int size = 0;
                    std::deque<STAFStringBufferPtr> envList;

                    for (STAFProcessEnvMap::iterator iter2 = envMap.begin();
                        iter2 != envMap.end(); ++iter2)
                    {
                        STAFProcessEnvData envData = iter2->second;

                        STAFString envCombo = envData.envName + kUTF8_EQUAL +
                            envData.envValue;

                        STAFStringBufferPtr envComboPtr = 
                            envCombo.toCurrentCodePage();

                        // Add 1 for the null byte
                        size += envComboPtr->length() + 1; 

                        envList.push_back(envComboPtr);
                    }

                    // Allocate the buffer

                    size += 1;                // Add one for the trailing null
                    newEnvBuf = new char[size];

                    // Walk the list and add the entries to the buffer

                    size = 0;

                    for (std::deque<STAFStringBufferPtr>::iterator iter = 
                         envList.begin(); iter != envList.end(); ++iter)
                    {
                        memcpy(newEnvBuf + size, (*iter)->buffer(), 
                               (*iter)->length());
                        newEnvBuf[size + (*iter)->length()] = 0;
                        size += (*iter)->length() + 1;
                    }

                    // Add the trailing null

                    newEnvBuf[size] = 0;

                }  // End merge of environment block

                ranCreateProcessAsUser = TRUE;

                // Call CreateProcessAsUser function to create a new process
                // which executes a specified executable file.  The new process
                // runs in the security context of the user represented by the
                // hToken parameter.  The process inherits the environment
                // associated with the specified user and any additional
                // environment variables specified.  Run process in the user's
                // interactive desktop (via startInfo).

                rc = CreateProcessAsUser(hUsrToken, 0, command, 0, 0, TRUE,
                     ((startData->consoleMode == kSTAFProcessNewConsole) ?
                      CREATE_NEW_CONSOLE : 0) | CREATE_NEW_PROCESS_GROUP |
                      (unicodeEnvBlock ? CREATE_UNICODE_ENVIRONMENT : 0),
                      (unicodeEnvBlock ? environment : newEnvBuf),
                      workdir, &startInfo, &processInfo);

                if (rc != TRUE)
                {
                    systemErrorCode = GetLastError();
                    failedFunction = "Error starting the process. "
                        "CreateProcessAsUser";
                }
                
                if (unicodeEnvBlock)
                    delete [] newEnvBuf;
                
                // Free the environment block (it's already been copied)
                if (!sDestroyEnvironmentBlockFunc(environment))
                    WriteTraceError("DestroyEnvironmentBlock()");
            }
        }

        if (rc != TRUE)
        {
            STAFString_t systemErrorMsg = 0;

            STAFUtilWin32LookupSystemErrorMessage(
                systemErrorCode, &systemErrorMsg);

            returnErrorMsg = failedFunction + " failed with OS RC " +
                STAFString(systemErrorCode);
             
            STAFString systemErrorString = STAFString(
                systemErrorMsg, STAFString::kShallow);

            if (systemErrorString.length() != 0)
                returnErrorMsg += ": " + STAFString(systemErrorString);

            if ((systemErrorCode == ERROR_FILE_NOT_FOUND) ||
                (systemErrorCode == ERROR_PATH_NOT_FOUND))
            {
                returnErrorMsg += "\nInvalid command: " +
                        STAFString(command);
            }
            else if (ERROR_DIRECTORY && startData->workdir != 0)
            {
                returnErrorMsg += "\nInvalid working directory: " +
                    STAFString(workdir);
            }
            else if (ranCreateProcessAsUser &&
                (systemErrorCode == ERROR_PRIVILEGE_NOT_HELD))
            {
                // Add additional information to the error message

                HANDLE processToken = 0;

                if (OpenProcessToken(GetCurrentProcess(),
                                     TOKEN_ALL_ACCESS, &processToken))
                {
                    STAFString privilegeErrorMsg = "";
                    GetPrivilegeError(processToken, privilegeErrorMsg);
                    returnErrorMsg += "\n" + privilegeErrorMsg;
                }
            }

            ret  = kSTAFBaseOSError;
            if (osRC) *osRC = systemErrorCode;
            if (errorBuffer) *errorBuffer = returnErrorMsg.adoptImpl();
        }
    }

    CloseHandle(processInfo.hThread);

    CloseHandle(startInfo.hStdInput);
    CloseHandle(startInfo.hStdOutput);
    CloseHandle(startInfo.hStdError);

    if (startData->stdinMode != kSTAFProcessIONoRedirect)
        CloseHandle(newInpHandle);
    if (startData->stdoutMode != kSTAFProcessIONoRedirect)
        CloseHandle(newOutHandle);      
    if (startData->stderrMode != kSTAFProcessIONoRedirect)
        CloseHandle(newErrHandle);

    if (ret != kSTAFOk) return ret;

    if (processID)     *processID     = processInfo.dwProcessId;
    if (processHandle) *processHandle = processInfo.hProcess;

    STAFProcessEndCallbackLevel1 dummyCallback = { 0 };
    
    AddProcessMonitor(ProcessMonitorInfo(processInfo.hProcess, 
                                         processInfo.dwProcessId,
                                         (startData->callback != 0) ?
                                             *startData->callback :
                                             dummyCallback,
                                         hUsrToken, profinfo.hProfile));

    return kSTAFOk;
}


BOOL CALLBACK sendWM_CLOSE(HWND winHandle, LPARAM pid)
{
    DWORD winPID = 0;
    DWORD winTID = GetWindowThreadProcessId(winHandle, &winPID);

    if (winPID == pid)
    {
        BOOL rc = PostMessage(winHandle, WM_CLOSE, 0, 0);

        if (rc == FALSE)
        {
            // XXX: What to do?  Maybe STAFTrace?
        }
    }

    return TRUE;
}


STAFRC_t STAFProcessStop(STAFProcessID_t pid, STAFProcessStopMethod_t stopMethod,
                         unsigned int *osRC)
{
    if ((stopMethod != kSTAFProcessStopWithSigKill) &&
        (stopMethod != kSTAFProcessStopWithSigTerm) &&
        (stopMethod != kSTAFProcessStopWithSigInt) &&
        (stopMethod != kSTAFProcessStopWithWM_CLOSE))
    {
        return kSTAFInvalidValue;
    }

    STAFMutexSemLock monitorListLock(sMonitorDataSem);

    ProcessMonitorInfo info;
    bool found = false;

    for (ProcessMonitorThreadList::iterator threadIter = sThreadList.begin();
         threadIter != sThreadList.end(); ++threadIter)
    {
        for (ProcessMonitorMap::iterator mapIter =
             (*threadIter)->monitorMap.begin();
             mapIter != (*threadIter)->monitorMap.end(); ++mapIter)
        {
            for (ProcessMonitorList::iterator listIter = 
                 mapIter->second.begin();
                 listIter != mapIter->second.end(); ++ listIter)
            {
                if (listIter->pid == pid)
                {
                    info = *listIter;
                    found = true;
                    break;
                }
            }
        }
    }

    if (found)
    {
        if (stopMethod == kSTAFProcessStopWithSigKill)
        {
            BOOL rc = TerminateProcess(info.handle, 0xFFFFFFFF);
            if (rc == FALSE)
            {
                if (osRC) *osRC = GetLastError();
                return kSTAFBaseOSError;
            }
        }
        else if (stopMethod == kSTAFProcessStopWithWM_CLOSE)
        {
            BOOL rc = EnumWindows(sendWM_CLOSE, info.pid);

            if (rc == FALSE)
            {
                if (osRC) *osRC = GetLastError();
                return kSTAFBaseOSError;
            }
        }
        else if ((stopMethod == kSTAFProcessStopWithSigInt) ||
                 (stopMethod == kSTAFProcessStopWithSigTerm))
        {
            BOOL rc = GenerateConsoleCtrlEvent(
                      (stopMethod == kSTAFProcessStopWithSigInt) ?