Lorenzob's picture
Upload folder using huggingface_hub
19605ab verified
raw
history blame
11.9 kB
// Copyright (c) 2011-2012 Ryan Prichard
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to
// deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
#include <string.h>
#include <algorithm>
#include "EventLoop.h"
#include "NamedPipe.h"
#include "../shared/DebugClient.h"
#include "../shared/StringUtil.h"
#include "../shared/WindowsSecurity.h"
#include "../shared/WinptyAssert.h"
// Returns true if anything happens (data received, data sent, pipe error).
bool NamedPipe::serviceIo(std::vector<HANDLE> *waitHandles)
{
bool justConnected = false;
const auto kError = ServiceResult::Error;
const auto kProgress = ServiceResult::Progress;
const auto kNoProgress = ServiceResult::NoProgress;
if (m_handle == NULL) {
return false;
}
if (m_connectEvent.get() != nullptr) {
// We're still connecting this server pipe. Check whether the pipe is
// now connected. If it isn't, add the pipe to the list of handles to
// wait on.
DWORD actual = 0;
BOOL success =
GetOverlappedResult(m_handle, &m_connectOver, &actual, FALSE);
if (!success && GetLastError() == ERROR_PIPE_CONNECTED) {
// I'm not sure this can happen, but it's easy to handle if it
// does.
success = TRUE;
}
if (!success) {
ASSERT(GetLastError() == ERROR_IO_INCOMPLETE &&
"Pended ConnectNamedPipe call failed");
waitHandles->push_back(m_connectEvent.get());
} else {
TRACE("Server pipe [%s] connected",
utf8FromWide(m_name).c_str());
m_connectEvent.dispose();
startPipeWorkers();
justConnected = true;
}
}
const auto readProgress = m_inputWorker ? m_inputWorker->service() : kNoProgress;
const auto writeProgress = m_outputWorker ? m_outputWorker->service() : kNoProgress;
if (readProgress == kError || writeProgress == kError) {
closePipe();
return true;
}
if (m_inputWorker && m_inputWorker->getWaitEvent() != nullptr) {
waitHandles->push_back(m_inputWorker->getWaitEvent());
}
if (m_outputWorker && m_outputWorker->getWaitEvent() != nullptr) {
waitHandles->push_back(m_outputWorker->getWaitEvent());
}
return justConnected
|| readProgress == kProgress
|| writeProgress == kProgress;
}
// manual reset, initially unset
static OwnedHandle createEvent() {
HANDLE ret = CreateEventW(nullptr, TRUE, FALSE, nullptr);
ASSERT(ret != nullptr && "CreateEventW failed");
return OwnedHandle(ret);
}
NamedPipe::IoWorker::IoWorker(NamedPipe &namedPipe) :
m_namedPipe(namedPipe),
m_event(createEvent())
{
}
NamedPipe::ServiceResult NamedPipe::IoWorker::service()
{
ServiceResult progress = ServiceResult::NoProgress;
if (m_pending) {
DWORD actual = 0;
BOOL ret = GetOverlappedResult(m_namedPipe.m_handle, &m_over, &actual, FALSE);
if (!ret) {
if (GetLastError() == ERROR_IO_INCOMPLETE) {
// There is a pending I/O.
return progress;
} else {
// Pipe error.
return ServiceResult::Error;
}
}
ResetEvent(m_event.get());
m_pending = false;
completeIo(actual);
m_currentIoSize = 0;
progress = ServiceResult::Progress;
}
DWORD nextSize = 0;
bool isRead = false;
while (shouldIssueIo(&nextSize, &isRead)) {
m_currentIoSize = nextSize;
DWORD actual = 0;
memset(&m_over, 0, sizeof(m_over));
m_over.hEvent = m_event.get();
BOOL ret = isRead
? ReadFile(m_namedPipe.m_handle, m_buffer, nextSize, &actual, &m_over)
: WriteFile(m_namedPipe.m_handle, m_buffer, nextSize, &actual, &m_over);
if (!ret) {
if (GetLastError() == ERROR_IO_PENDING) {
// There is a pending I/O.
m_pending = true;
return progress;
} else {
// Pipe error.
return ServiceResult::Error;
}
}
ResetEvent(m_event.get());
completeIo(actual);
m_currentIoSize = 0;
progress = ServiceResult::Progress;
}
return progress;
}
// This function is called after CancelIo has returned. We need to block until
// the I/O operations have completed, which should happen very quickly.
// https://blogs.msdn.microsoft.com/oldnewthing/20110202-00/?p=11613
void NamedPipe::IoWorker::waitForCanceledIo()
{
if (m_pending) {
DWORD actual = 0;
GetOverlappedResult(m_namedPipe.m_handle, &m_over, &actual, TRUE);
m_pending = false;
}
}
HANDLE NamedPipe::IoWorker::getWaitEvent()
{
return m_pending ? m_event.get() : NULL;
}
void NamedPipe::InputWorker::completeIo(DWORD size)
{
m_namedPipe.m_inQueue.append(m_buffer, size);
}
bool NamedPipe::InputWorker::shouldIssueIo(DWORD *size, bool *isRead)
{
*isRead = true;
ASSERT(!m_namedPipe.isConnecting());
if (m_namedPipe.isClosed()) {
return false;
} else if (m_namedPipe.m_inQueue.size() < m_namedPipe.readBufferSize()) {
*size = kIoSize;
return true;
} else {
return false;
}
}
void NamedPipe::OutputWorker::completeIo(DWORD size)
{
ASSERT(size == m_currentIoSize);
}
bool NamedPipe::OutputWorker::shouldIssueIo(DWORD *size, bool *isRead)
{
*isRead = false;
if (!m_namedPipe.m_outQueue.empty()) {
auto &out = m_namedPipe.m_outQueue;
const DWORD writeSize = std::min<size_t>(out.size(), kIoSize);
std::copy(&out[0], &out[writeSize], m_buffer);
out.erase(0, writeSize);
*size = writeSize;
return true;
} else {
return false;
}
}
DWORD NamedPipe::OutputWorker::getPendingIoSize()
{
return m_pending ? m_currentIoSize : 0;
}
void NamedPipe::openServerPipe(LPCWSTR pipeName, OpenMode::t openMode,
int outBufferSize, int inBufferSize) {
ASSERT(isClosed());
ASSERT((openMode & OpenMode::Duplex) != 0);
const DWORD winOpenMode =
((openMode & OpenMode::Reading) ? PIPE_ACCESS_INBOUND : 0)
| ((openMode & OpenMode::Writing) ? PIPE_ACCESS_OUTBOUND : 0)
| FILE_FLAG_FIRST_PIPE_INSTANCE
| FILE_FLAG_OVERLAPPED;
const auto sd = createPipeSecurityDescriptorOwnerFullControl();
ASSERT(sd && "error creating data pipe SECURITY_DESCRIPTOR");
SECURITY_ATTRIBUTES sa = {};
sa.nLength = sizeof(sa);
sa.lpSecurityDescriptor = sd.get();
HANDLE handle = CreateNamedPipeW(
pipeName,
/*dwOpenMode=*/winOpenMode,
/*dwPipeMode=*/rejectRemoteClientsPipeFlag(),
/*nMaxInstances=*/1,
/*nOutBufferSize=*/outBufferSize,
/*nInBufferSize=*/inBufferSize,
/*nDefaultTimeOut=*/30000,
&sa);
TRACE("opened server pipe [%s], handle == %p",
utf8FromWide(pipeName).c_str(), handle);
ASSERT(handle != INVALID_HANDLE_VALUE && "Could not open server pipe");
m_name = pipeName;
m_handle = handle;
m_openMode = openMode;
// Start an asynchronous connection attempt.
m_connectEvent = createEvent();
memset(&m_connectOver, 0, sizeof(m_connectOver));
m_connectOver.hEvent = m_connectEvent.get();
BOOL success = ConnectNamedPipe(m_handle, &m_connectOver);
const auto err = GetLastError();
if (!success && err == ERROR_PIPE_CONNECTED) {
success = TRUE;
}
if (success) {
TRACE("Server pipe [%s] connected", utf8FromWide(pipeName).c_str());
m_connectEvent.dispose();
startPipeWorkers();
} else if (err != ERROR_IO_PENDING) {
ASSERT(false && "ConnectNamedPipe call failed");
}
}
void NamedPipe::connectToServer(LPCWSTR pipeName, OpenMode::t openMode)
{
ASSERT(isClosed());
ASSERT((openMode & OpenMode::Duplex) != 0);
HANDLE handle = CreateFileW(
pipeName,
GENERIC_READ | GENERIC_WRITE,
0,
NULL,
OPEN_EXISTING,
SECURITY_SQOS_PRESENT | SECURITY_IDENTIFICATION | FILE_FLAG_OVERLAPPED,
NULL);
TRACE("connected to [%s], handle == %p",
utf8FromWide(pipeName).c_str(), handle);
ASSERT(handle != INVALID_HANDLE_VALUE && "Could not connect to pipe");
m_name = pipeName;
m_handle = handle;
m_openMode = openMode;
startPipeWorkers();
}
void NamedPipe::startPipeWorkers()
{
if (m_openMode & OpenMode::Reading) {
m_inputWorker.reset(new InputWorker(*this));
}
if (m_openMode & OpenMode::Writing) {
m_outputWorker.reset(new OutputWorker(*this));
}
}
size_t NamedPipe::bytesToSend()
{
ASSERT(m_openMode & OpenMode::Writing);
auto ret = m_outQueue.size();
if (m_outputWorker != NULL) {
ret += m_outputWorker->getPendingIoSize();
}
return ret;
}
void NamedPipe::write(const void *data, size_t size)
{
ASSERT(m_openMode & OpenMode::Writing);
m_outQueue.append(reinterpret_cast<const char*>(data), size);
}
void NamedPipe::write(const char *text)
{
write(text, strlen(text));
}
size_t NamedPipe::readBufferSize()
{
ASSERT(m_openMode & OpenMode::Reading);
return m_readBufferSize;
}
void NamedPipe::setReadBufferSize(size_t size)
{
ASSERT(m_openMode & OpenMode::Reading);
m_readBufferSize = size;
}
size_t NamedPipe::bytesAvailable()
{
ASSERT(m_openMode & OpenMode::Reading);
return m_inQueue.size();
}
size_t NamedPipe::peek(void *data, size_t size)
{
ASSERT(m_openMode & OpenMode::Reading);
const auto out = reinterpret_cast<char*>(data);
const size_t ret = std::min(size, m_inQueue.size());
std::copy(&m_inQueue[0], &m_inQueue[ret], out);
return ret;
}
size_t NamedPipe::read(void *data, size_t size)
{
size_t ret = peek(data, size);
m_inQueue.erase(0, ret);
return ret;
}
std::string NamedPipe::readToString(size_t size)
{
ASSERT(m_openMode & OpenMode::Reading);
size_t retSize = std::min(size, m_inQueue.size());
std::string ret = m_inQueue.substr(0, retSize);
m_inQueue.erase(0, retSize);
return ret;
}
std::string NamedPipe::readAllToString()
{
ASSERT(m_openMode & OpenMode::Reading);
std::string ret = m_inQueue;
m_inQueue.clear();
return ret;
}
void NamedPipe::closePipe()
{
if (m_handle == NULL) {
return;
}
CancelIo(m_handle);
if (m_connectEvent.get() != nullptr) {
DWORD actual = 0;
GetOverlappedResult(m_handle, &m_connectOver, &actual, TRUE);
m_connectEvent.dispose();
}
if (m_inputWorker) {
m_inputWorker->waitForCanceledIo();
m_inputWorker.reset();
}
if (m_outputWorker) {
m_outputWorker->waitForCanceledIo();
m_outputWorker.reset();
}
CloseHandle(m_handle);
m_handle = NULL;
}