116 lines
4.1 KiB
C++
116 lines
4.1 KiB
C++
#pragma once
|
|
|
|
#include <Sched/Scheduler.hpp>
|
|
#include <Ipc/Ipc.hpp>
|
|
|
|
namespace Montauk {
|
|
|
|
static Sched::Process* GetRedirTarget(Sched::Process* proc) {
|
|
if (!proc || !proc->redirected) return nullptr;
|
|
if (proc->ioOutHandle >= 0 || proc->ioInHandle >= 0 || proc->ioKeyHandle >= 0) {
|
|
return proc;
|
|
}
|
|
return (proc->parentPid >= 0) ? Sched::GetProcessByPid(proc->parentPid) : nullptr;
|
|
}
|
|
|
|
static bool ResolveProcessHandle(Sched::Process* proc, int handle, Ipc::HandleType expectedType,
|
|
Ipc::Object*& outObject, uint32_t* outRights = nullptr) {
|
|
if (proc == nullptr || handle < 0) return false;
|
|
|
|
int slot = Ipc::SlotForPid(proc->pid);
|
|
if (slot < 0) return false;
|
|
|
|
Ipc::HandleType type = Ipc::HandleType::None;
|
|
Ipc::Object* object = nullptr;
|
|
uint32_t rights = 0;
|
|
if (!Ipc::SnapshotHandleForSlot(slot, handle, type, object, rights)) return false;
|
|
if (type != expectedType || object == nullptr) return false;
|
|
|
|
outObject = object;
|
|
if (outRights != nullptr) *outRights = rights;
|
|
return true;
|
|
}
|
|
|
|
static Ipc::Stream* GetRedirOutStream(Sched::Process* proc) {
|
|
proc = GetRedirTarget(proc);
|
|
if (proc == nullptr) return nullptr;
|
|
|
|
Ipc::Object* object = nullptr;
|
|
if (!ResolveProcessHandle(proc, proc->ioOutHandle, Ipc::HandleType::Stream, object)) return nullptr;
|
|
return (Ipc::Stream*)object;
|
|
}
|
|
|
|
static Ipc::Stream* GetRedirInStream(Sched::Process* proc) {
|
|
proc = GetRedirTarget(proc);
|
|
if (proc == nullptr) return nullptr;
|
|
|
|
Ipc::Object* object = nullptr;
|
|
if (!ResolveProcessHandle(proc, proc->ioInHandle, Ipc::HandleType::Stream, object)) return nullptr;
|
|
return (Ipc::Stream*)object;
|
|
}
|
|
|
|
static Ipc::Mailbox* GetRedirKeyMailbox(Sched::Process* proc) {
|
|
proc = GetRedirTarget(proc);
|
|
if (proc == nullptr) return nullptr;
|
|
|
|
Ipc::Object* object = nullptr;
|
|
if (!ResolveProcessHandle(proc, proc->ioKeyHandle, Ipc::HandleType::Mailbox, object)) return nullptr;
|
|
return (Ipc::Mailbox*)object;
|
|
}
|
|
|
|
static int DuplicateHandleBetweenSlots(int srcSlot, int handle, int dstSlot) {
|
|
if (srcSlot < 0 || dstSlot < 0 || handle < 0) return -1;
|
|
|
|
Ipc::HandleType type = Ipc::HandleType::None;
|
|
Ipc::Object* object = nullptr;
|
|
uint32_t rights = 0;
|
|
if (!Ipc::SnapshotHandleForSlot(srcSlot, handle, type, object, rights)) return -1;
|
|
if ((rights & Ipc::RightDup) == 0) return -1;
|
|
return Ipc::InstallHandleForSlot(dstSlot, object, type, rights);
|
|
}
|
|
|
|
static bool ConfigureRedirWaitsetForSlot(int slot, Sched::Process* proc) {
|
|
if (slot < 0 || proc == nullptr) return false;
|
|
|
|
int waitset = Ipc::CreateWaitsetHandleForSlot(slot);
|
|
if (waitset < 0) return false;
|
|
|
|
if (proc->ioInHandle >= 0) {
|
|
if (Ipc::WaitsetAddHandleForSlot(slot, waitset, proc->ioInHandle,
|
|
Ipc::SignalReadable | Ipc::SignalPeerClosed) < 0) {
|
|
Ipc::CloseHandleForSlot(slot, waitset);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
if (proc->ioKeyHandle >= 0) {
|
|
if (Ipc::WaitsetAddHandleForSlot(slot, waitset, proc->ioKeyHandle,
|
|
Ipc::SignalReadable | Ipc::SignalPeerClosed) < 0) {
|
|
Ipc::CloseHandleForSlot(slot, waitset);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
proc->ioWaitsetHandle = waitset;
|
|
return true;
|
|
}
|
|
|
|
static int WriteAllToStream(Ipc::Stream* stream, const uint8_t* data, int len) {
|
|
if (stream == nullptr || data == nullptr || len <= 0) return -1;
|
|
|
|
int total = 0;
|
|
while (total < len) {
|
|
int written = Ipc::StreamWrite(stream, data + total, len - total, false);
|
|
if (written < 0) {
|
|
return (total > 0) ? total : -1;
|
|
}
|
|
if (written == 0) {
|
|
Sched::BlockOnObject(stream, 0);
|
|
continue;
|
|
}
|
|
total += written;
|
|
}
|
|
return total;
|
|
}
|
|
}
|