ncbi_pipe.cpp

ncbi源码

/*
 * ===========================================================================
 * PRODUCTION $Log: ncbi_pipe.cpp,v $
 * PRODUCTION Revision 1000.4  2004/06/01 18:45:13  gouriano
 * PRODUCTION PRODUCTION: UPGRADED [GCC34_MSVC7] Dev-tree R1.37
 * PRODUCTION
 * ===========================================================================
 */

/*  $Id: ncbi_pipe.cpp,v 1000.4 2004/06/01 18:45:13 gouriano Exp $
 * ===========================================================================
 *
 *                            PUBLIC DOMAIN NOTICE
 *               National Center for Biotechnology Information
 *
 *  This software/database is a "United States Government Work" under the
 *  terms of the United States Copyright Act.  It was written as part of
 *  the author's official duties as a United States Government employee and
 *  thus cannot be copyrighted.  This software/database is freely available
 *  to the public for use. The National Library of Medicine and the U.S.
 *  Government have not placed any restriction on its use or reproduction.
 *
 *  Although all reasonable efforts have been taken to ensure the accuracy
 *  and reliability of the software and data, the NLM and the U.S.
 *  Government do not and cannot warrant the performance or results that
 *  may be obtained by using this software or data. The NLM and the U.S.
 *  Government disclaim all warranties, express or implied, including
 *  warranties of performance, merchantability or fitness for any particular
 *  purpose.
 *
 *  Please cite the author in any work or product based on this material.
 *
 * ===========================================================================
 *
 * Author:  Anton Lavrentiev, Mike DiCuccio, Vladimir Ivanov
 *
 *
 */

#include <ncbi_pch.hpp>
#include <connect/ncbi_pipe.hpp>
#include <corelib/ncbi_system.hpp>
#include <memory>
#include <stdio.h>

#ifdef NCBI_OS_MSWIN
#  include <windows.h>
#elif defined NCBI_OS_UNIX
#  include <unistd.h>
#  include <errno.h>
#  include <sys/time.h>
#  include <sys/types.h>
#  include <sys/wait.h>
#  include <signal.h>
#  include <fcntl.h>
#  ifdef NCBI_COMPILER_MW_MSL
#    include <ncbi_mslextras.h>
#  endif
#else
#  error "Class CPipe is supported only on Windows and Unix"
#endif


BEGIN_NCBI_SCOPE


// Sleep time for timeouts
const unsigned int kSleepTime = 100;


//////////////////////////////////////////////////////////////////////////////
//
// Auxiliary functions
//

#define IS_SET(flags, mask) (((flags) & (mask)) == (mask))


static STimeout* s_SetTimeout(const STimeout* from, STimeout* to)
{
    if ( !from ) {
        return const_cast<STimeout*>(kInfiniteTimeout);
    }
    to->sec  = from->usec / 1000000 + from->sec;
    to->usec = from->usec % 1000000;
    return to;
}

static string s_FormatErrorMessage(const string& where, const string& what)
{
    return "[CPipe::" + where + "]  " + what + ".";
}


//////////////////////////////////////////////////////////////////////////////
//
// Class CNamedPipeHandle handles forwarded requests from CPipe.
// This class is reimplemented in a platform-specific fashion where needed.
//

#if defined(NCBI_OS_MSWIN)

//////////////////////////////////////////////////////////////////////////////
//
// CPipeHandle -- MS Windows version
//

class CPipeHandle
{
public:
    CPipeHandle();
    ~CPipeHandle();
    EIO_Status Open(const string& cmd, const vector<string>& args,
                    CPipe::TCreateFlags create_flags);
    EIO_Status Close(int* exitcode, const STimeout* timeout);
    EIO_Status CloseHandle (CPipe::EChildIOHandle handle);
    EIO_Status Read(void* buf, size_t count, size_t* n_read,
                    const CPipe::EChildIOHandle from_handle,
                    const STimeout* timeout);
    EIO_Status Write(const void* buf, size_t count, size_t* written,
                     const STimeout* timeout);
    TProcessHandle GetProcessHandle(void) const { return m_ProcHandle; };

private:
    // Clear object state.
    void x_Clear(void);
    // Get child's I/O handle.
    HANDLE x_GetHandle(CPipe::EChildIOHandle from_handle) const;
    // Convert STimeout value to number of milliseconds.
    long   x_TimeoutToMSec(const STimeout* timeout) const;
    // Trigger blocking mode on specified I/O handle.
    bool   x_SetNonBlockingMode(HANDLE fd, bool nonblock = true) const;

private:
    // I/O handles for child process.
    HANDLE  m_ChildStdIn;
    HANDLE  m_ChildStdOut;
    HANDLE  m_ChildStdErr;

    // Child process descriptor.
    HANDLE  m_ProcHandle;
    // Child process pid.
    TPid    m_Pid;

    // Pipe flags
    CPipe::TCreateFlags m_Flags;
};


CPipeHandle::CPipeHandle()
    : m_ProcHandle(INVALID_HANDLE_VALUE),
      m_ChildStdIn(INVALID_HANDLE_VALUE),
      m_ChildStdOut(INVALID_HANDLE_VALUE),
      m_ChildStdErr(INVALID_HANDLE_VALUE),
      m_Flags(0)
{
    return;
}


CPipeHandle::~CPipeHandle()
{
    static const STimeout kZeroTimeout = {0, 0};
    Close(0, &kZeroTimeout);
    x_Clear();
}


EIO_Status CPipeHandle::Open(const string&         cmd,
                             const vector<string>& args,
                             CPipe::TCreateFlags   create_flags)
{
    x_Clear();

    bool need_restore_handles = false; 
    m_Flags = create_flags;

    HANDLE stdin_handle       = INVALID_HANDLE_VALUE;
    HANDLE stdout_handle      = INVALID_HANDLE_VALUE;
    HANDLE stderr_handle      = INVALID_HANDLE_VALUE;
    HANDLE child_stdin_read   = INVALID_HANDLE_VALUE;
    HANDLE child_stdin_write  = INVALID_HANDLE_VALUE;
    HANDLE child_stdout_read  = INVALID_HANDLE_VALUE;
    HANDLE child_stdout_write = INVALID_HANDLE_VALUE;
    HANDLE child_stderr_read  = INVALID_HANDLE_VALUE;
    HANDLE child_stderr_write = INVALID_HANDLE_VALUE;

    try {
        if (m_ProcHandle != INVALID_HANDLE_VALUE) {
            throw string("Pipe is already open");
        }

        // Save current I/O handles
        stdin_handle  = GetStdHandle(STD_INPUT_HANDLE);
        stdout_handle = GetStdHandle(STD_OUTPUT_HANDLE);
        stderr_handle = GetStdHandle(STD_ERROR_HANDLE);
        
        // Flush std.output buffers before remap
        FlushFileBuffers(stdout_handle);
        FlushFileBuffers(stderr_handle);

        // Set base security attributes
        SECURITY_ATTRIBUTES attr;
        attr.nLength = sizeof(attr);
        attr.bInheritHandle = TRUE;
        attr.lpSecurityDescriptor = NULL;

        HANDLE current_process = GetCurrentProcess();
        need_restore_handles = true; 

        // Create pipe for child's stdin
        assert(CPipe::fStdIn_Close);
        if ( !IS_SET(create_flags, CPipe::fStdIn_Close) ) {
            if ( !CreatePipe(&child_stdin_read,
                             &child_stdin_write, &attr, 0) ) {
                throw string("CreatePipe() failed");
            }
            if ( !SetStdHandle(STD_INPUT_HANDLE, child_stdin_read) ) {
                throw string("Failed to remap stdin for child process");
            }
            // Duplicate the handle
            if ( !DuplicateHandle(current_process, child_stdin_write,
                                  current_process, &m_ChildStdIn,
                                  0, FALSE, DUPLICATE_SAME_ACCESS) ) {
                throw string("DuplicateHandle() failed on "
                             "child's stdin handle");
            }
            ::CloseHandle(child_stdin_write);
            x_SetNonBlockingMode(m_ChildStdIn);
        }

        // Create pipe for child's stdout
        assert(CPipe::fStdOut_Close);
        if ( !IS_SET(create_flags, CPipe::fStdOut_Close) ) {
            if ( !CreatePipe(&child_stdout_read,
                             &child_stdout_write, &attr, 0)) {
                throw string("CreatePipe() failed");
            }
            if ( !SetStdHandle(STD_OUTPUT_HANDLE, child_stdout_write) ) {
                throw string("Failed to remap stdout for child process");
            }
            // Duplicate the handle
            if ( !DuplicateHandle(current_process, child_stdout_read,
                                  current_process, &m_ChildStdOut,
                                  0, FALSE, DUPLICATE_SAME_ACCESS) ) {
                throw string("DuplicateHandle() failed on "
                             "child's stdout handle");
            }
            ::CloseHandle(child_stdout_read);
            x_SetNonBlockingMode(m_ChildStdOut);
        }

        // Create pipe for child's stderr
        assert(CPipe::fStdErr_Open);
        if ( IS_SET(create_flags, CPipe::fStdErr_Open) ) {
            if ( !CreatePipe(&child_stderr_read,
                             &child_stderr_write, &attr, 0)) {
                throw string("CreatePipe() failed");
            }
            if ( !SetStdHandle(STD_ERROR_HANDLE, child_stderr_write) ) {
                throw string("Failed to remap stderr for child process");
            }
            // Duplicate the handle
            if ( !DuplicateHandle(current_process, child_stderr_read,
                                  current_process, &m_ChildStdErr,
                                  0, FALSE, DUPLICATE_SAME_ACCESS) ) {
                throw string("DuplicateHandle() failed on "
                             "child's stderr handle");
            }
            ::CloseHandle(child_stderr_read);
            x_SetNonBlockingMode(m_ChildStdErr);
        }

        // Prepare command line to run
        string cmd_line(cmd);
        ITERATE (vector<string>, iter, args) {
            if ( !cmd_line.empty() ) {
                cmd_line += " ";
            }
            cmd_line += *iter;
        }

        // Create child process
        PROCESS_INFORMATION pinfo;
        STARTUPINFO sinfo;
        ZeroMemory(&pinfo, sizeof(PROCESS_INFORMATION));
        ZeroMemory(&sinfo, sizeof(STARTUPINFO));
        sinfo.cb = sizeof(sinfo);

        if ( !CreateProcess(NULL,
                            const_cast<char*> (cmd_line.c_str()),
                            NULL, NULL, TRUE, 0,
                            NULL, NULL, &sinfo, &pinfo) ) {
            throw "CreateProcess() for \"" + cmd_line + "\" failed";
        }
        ::CloseHandle(pinfo.hThread);
        m_ProcHandle = pinfo.hProcess;
        m_Pid        = pinfo.dwProcessId; 

        assert(m_ProcHandle != INVALID_HANDLE_VALUE);

        // Restore remapped handles back to their original states
        if ( !SetStdHandle(STD_INPUT_HANDLE,  stdin_handle) ) {
            throw string("Failed to remap stdin for parent process");
        }
        if ( !SetStdHandle(STD_OUTPUT_HANDLE, stdout_handle) ) {
            throw string("Failed to remap stdout for parent process");
        }
        if ( !SetStdHandle(STD_ERROR_HANDLE,  stderr_handle) ) {
            throw string("Failed to remap stderr for parent process");
        }
        // Close unused pipe handles
        ::CloseHandle(child_stdin_read);
        ::CloseHandle(child_stdout_write);
        ::CloseHandle(child_stderr_write);

        return eIO_Success;
     }
    catch (string& what) {
        // Restore all standard handles on error
        if ( need_restore_handles ) {
            SetStdHandle(STD_INPUT_HANDLE,  stdin_handle);
            SetStdHandle(STD_OUTPUT_HANDLE, stdout_handle);
            SetStdHandle(STD_ERROR_HANDLE,  stderr_handle);

            ::CloseHandle(child_stdin_read);
            ::CloseHandle(child_stdin_write);
            ::CloseHandle(child_stdout_read);
            ::CloseHandle(child_stdout_write);
            ::CloseHandle(child_stderr_read);
            ::CloseHandle(child_stderr_write);
        }
        const STimeout kZeroZimeout = {0,0};
        Close(0, &kZeroZimeout);
        ERR_POST(s_FormatErrorMessage("Open", what));
        return eIO_Unknown;
    }
}


EIO_Status CPipeHandle::Close(int* exitcode, const STimeout* timeout)
{
    if (m_ProcHandle == INVALID_HANDLE_VALUE) {
        return eIO_Closed;
    }

    DWORD      x_exitcode = -1;
    EIO_Status status     = eIO_Unknown;

    // Get exit code of child process
    if ( GetExitCodeProcess(m_ProcHandle, &x_exitcode) ) {
        if (x_exitcode == STILL_ACTIVE) {
            // Wait for the child process to exit
            DWORD ws = WaitForSingleObject(m_ProcHandle,
                                           x_TimeoutToMSec(timeout));
            if (ws == WAIT_OBJECT_0) {
                // Get exit code of child process over again
                if ( GetExitCodeProcess(m_ProcHandle, &x_exitcode) ) {
                    status = (x_exitcode == STILL_ACTIVE) ? 
                        eIO_Timeout : eIO_Success;
                }
            } else if (ws == WAIT_TIMEOUT) {
                status = eIO_Timeout;
            }
        } else {
            status = eIO_Success;
        }
    }

    // Is the process running?
    if (status == eIO_Timeout) {
        x_exitcode = -1;
        assert(CPipe::fKillOnClose);
        if ( IS_SET(m_Flags, CPipe::fKillOnClose) ) {
            status = CProcess(m_Pid, CProcess::ePid).Kill() ?
                              eIO_Success : eIO_Unknown;
        }
        // Active by default.
        assert(!CPipe::fCloseOnClose);
        if ( !IS_SET(m_Flags, CPipe::fCloseOnClose) ) {
            status = eIO_Success;
        }
        // fKeepOnClose -- nothing to do.
        assert(CPipe::fKeepOnClose);
    }

    // Is the process still running? Nothing to do.
    if (status != eIO_Timeout) {
        x_Clear();
    }
    if ( exitcode ) {
        *exitcode = (int) x_exitcode;
    }
    return status;
}


EIO_Status CPipeHandle::CloseHandle(CPipe::EChildIOHandle handle)
{
    switch (handle) {
    case CPipe::eStdIn:
        if (m_ChildStdIn == INVALID_HANDLE_VALUE) {
            return eIO_Closed;
        }
        ::CloseHandle(m_ChildStdIn);
        m_ChildStdIn = INVALID_HANDLE_VALUE;
        break;
    case CPipe::eStdOut:
        if (m_ChildStdOut == INVALID_HANDLE_VALUE) {
            return eIO_Closed;
        }
        ::CloseHandle(m_ChildStdOut);
        m_ChildStdOut = INVALID_HANDLE_VALUE;
        break;
    case CPipe::eStdErr:
        if (m_ChildStdErr == INVALID_HANDLE_VALUE) {
            return eIO_Closed;
        }
        ::CloseHandle(m_ChildStdErr);
        m_ChildStdErr = INVALID_HANDLE_VALUE;
        break;
    default:
        return eIO_InvalidArg;
    }
    return eIO_Success;
}


EIO_Status CPipeHandle::Read(void* buf, size_t count, size_t* read, 
                             const CPipe::EChildIOHandle from_handle,
                             const STimeout* timeout)
{
    EIO_Status status = eIO_Unknown;

    try {
        if (m_ProcHandle == INVALID_HANDLE_VALUE) {
            status = eIO_Closed;
            throw string("Pipe is closed");
        }
        HANDLE fd = x_GetHandle(from_handle);
        if (fd == INVALID_HANDLE_VALUE) {
            throw string("Pipe I/O handle is closed");
        }
        if ( !count ) {
            return eIO_Success;
        }

        DWORD x_timeout   = x_TimeoutToMSec(timeout);
        DWORD bytes_avail = 0;
        DWORD bytes_read  = 0;

        // Wait for data from the pipe with timeout.
        // NOTE:  The function WaitForSingleObject() does not work with pipes.
        do {
            if ( !PeekNamedPipe(fd, NULL, 0, NULL, &bytes_avail, NULL) ) {
                // Has peer closed the connection?
                if (GetLastError() == ERROR_BROKEN_PIPE) {
                    return eIO_Closed;
                }
                throw string("PeekNamedPipe() failed");
            }
            if ( bytes_avail ) {
                break;
            }
            DWORD x_sleep = kSleepTime;
            if (x_timeout != INFINITE) {
                if (x_timeout < kSleepTime) {
                    x_sleep = x_timeout;
                }
                x_timeout -= x_sleep;
            }
            SleepMilliSec(x_sleep);
        } while (x_timeout == INFINITE  ||  x_timeout);

        // Data is available to read or read request has timed out
        if ( !bytes_avail ) {
            return eIO_Timeout;
        }
        // We must read only "count" bytes of data regardless of
        // the amount available to read
        if (bytes_avail > count) {
            bytes_avail = count;
        }
        if ( !ReadFile(fd, buf, count, &bytes_avail, NULL) ) {
            throw string("Failed to read data from pipe");
        }
        if ( read ) {
            *read = bytes_avail;
        }
        status = eIO_Success;
    }
    catch (string& what) {
        ERR_POST(s_FormatErrorMessage("Read", what));
    }
    return status;