105 lines
3.9 KiB
C++
105 lines
3.9 KiB
C++
#pragma once
|
|
|
|
#include <Sched/Scheduler.hpp>
|
|
#include <Ipc/Ipc.hpp>
|
|
|
|
namespace montauk::abi {
|
|
|
|
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::HandleSnapshot& snapshot) {
|
|
if (proc == nullptr || handle < 0) return false;
|
|
|
|
int slot = Ipc::SlotForPid(proc->pid);
|
|
if (slot < 0) return false;
|
|
|
|
if (!snapshot.Capture(slot, handle)) return false;
|
|
return snapshot.type == expectedType && snapshot.object != nullptr;
|
|
}
|
|
|
|
static Ipc::Stream* GetRedirOutStream(Sched::Process* proc, Ipc::HandleSnapshot& snapshot) {
|
|
proc = GetRedirTarget(proc);
|
|
if (proc == nullptr) return nullptr;
|
|
|
|
if (!ResolveProcessHandle(proc, proc->ioOutHandle, Ipc::HandleType::Stream, snapshot)) return nullptr;
|
|
return (Ipc::Stream*)snapshot.object;
|
|
}
|
|
|
|
static Ipc::Stream* GetRedirInStream(Sched::Process* proc, Ipc::HandleSnapshot& snapshot) {
|
|
proc = GetRedirTarget(proc);
|
|
if (proc == nullptr) return nullptr;
|
|
|
|
if (!ResolveProcessHandle(proc, proc->ioInHandle, Ipc::HandleType::Stream, snapshot)) return nullptr;
|
|
return (Ipc::Stream*)snapshot.object;
|
|
}
|
|
|
|
static Ipc::Mailbox* GetRedirKeyMailbox(Sched::Process* proc, Ipc::HandleSnapshot& snapshot) {
|
|
proc = GetRedirTarget(proc);
|
|
if (proc == nullptr) return nullptr;
|
|
|
|
if (!ResolveProcessHandle(proc, proc->ioKeyHandle, Ipc::HandleType::Mailbox, snapshot)) return nullptr;
|
|
return (Ipc::Mailbox*)snapshot.object;
|
|
}
|
|
|
|
static int DuplicateHandleBetweenSlots(int srcSlot, int handle, int dstSlot) {
|
|
if (srcSlot < 0 || dstSlot < 0 || handle < 0) return -1;
|
|
|
|
Ipc::HandleSnapshot snapshot;
|
|
if (!snapshot.Capture(srcSlot, handle)) return -1;
|
|
if ((snapshot.rights & Ipc::RightDup) == 0) return -1;
|
|
return Ipc::InstallHandleForSlot(dstSlot, snapshot.object, snapshot.type, snapshot.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) {
|
|
uint64_t observedWake = Sched::ObserveObjectWake(stream);
|
|
int written = Ipc::StreamWrite(stream, data + total, len - total, false);
|
|
if (written < 0) {
|
|
return (total > 0) ? total : -1;
|
|
}
|
|
if (written == 0) {
|
|
Sched::BlockOnObjectSince(stream, 0, observedWake);
|
|
continue;
|
|
}
|
|
total += written;
|
|
}
|
|
return total;
|
|
}
|
|
}
|