diff --git a/.gitignore b/.gitignore index def911c..b0f943f 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,6 @@ /agent/bin /agent/obj /agent/compile_commands.json + +# test stuff +/shared/src/test \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 82e536c..7dae6c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -47,6 +47,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "dlib" version = "0.5.2" @@ -83,7 +89,9 @@ name = "fbcserver" version = "0.1.0" dependencies = [ "byteorder", + "libc", "minifb", + "nix 0.29.0", ] [[package]] @@ -204,9 +212,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" -version = "0.2.164" +version = "0.2.167" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "433bfe06b8c75da9b2e3fbea6e5329ff87748f0b144ef75306e674c3f6f7c13f" +checksum = "09d6582e104315a817dff97f75133544b2e094ee22447d2acf4a74e189ba06fc" [[package]] name = "libloading" @@ -294,6 +302,18 @@ dependencies = [ "memoffset", ] +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags 2.6.0", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "once_cell" version = "1.20.2" @@ -558,7 +578,7 @@ dependencies = [ "bitflags 1.3.2", "downcast-rs", "libc", - "nix", + "nix 0.24.3", "scoped-tls", "wayland-commons", "wayland-scanner", @@ -571,7 +591,7 @@ version = "0.29.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8691f134d584a33a6606d9d717b95c4fa20065605f798a3f350d78dced02a902" dependencies = [ - "nix", + "nix 0.24.3", "once_cell", "smallvec", "wayland-sys", @@ -583,7 +603,7 @@ version = "0.29.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6865c6b66f13d6257bef1cd40cbfe8ef2f150fb8ebbdb1e8e873455931377661" dependencies = [ - "nix", + "nix 0.24.3", "wayland-client", "xcursor", ] diff --git a/Cargo.toml b/Cargo.toml index f5ec5fa..df4549c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,4 +5,6 @@ edition = "2021" [dependencies] byteorder = "1.5.0" +libc = "0.2.167" minifb = "0.27.0" +nix = { version = "0.29.0", features = [ "stat", "mman" ] } diff --git a/agent/Makefile b/agent/Makefile index b06d69d..d349aa1 100644 --- a/agent/Makefile +++ b/agent/Makefile @@ -9,6 +9,7 @@ OBJDIR = obj/$(ARCH)/$(CONFIG) VPATH = src/crt \ src/base \ src/ \ + ../shared/src # Objects for crt OBJS_STARTUP := $(OBJDIR)/crt0.o \ @@ -21,6 +22,7 @@ OBJS_BASELIB := OBJS := $(OBJDIR)/capture.o \ $(OBJDIR)/capture_nvfbc.o \ $(OBJDIR)/nvfbc_library.o \ + $(OBJDIR)/ivshmem.o \ $(OBJDIR)/main.o \ LINK_LIBS := $(VS2022_PATH)/ucrt/lib/$(ARCH)/libucrt$(D).lib \ @@ -30,6 +32,7 @@ LINK_LIBS := $(VS2022_PATH)/ucrt/lib/$(ARCH)/libucrt$(D).lib \ $(VS2022_PATH)/winsdk/lib/$(ARCH)/kernel32.lib \ $(VS2022_PATH)/winsdk/lib/$(ARCH)/user32.lib \ $(VS2022_PATH)/winsdk/lib/$(ARCH)/comctl32.lib \ + $(VS2022_PATH)/winsdk/lib/$(ARCH)/setupapi.lib \ $(VS2022_PATH)/winsdk/lib/$(ARCH)/ws2_32.lib .PHONY: all dumpinfo clean matrix @@ -40,6 +43,11 @@ publish: $(BINDIR)/$(NAME).exe cp -rv bin /data/sda/shit/Release/clangent chmod -Rvv 777 /data/sda/shit/Release/clangent + +publish2: $(BINDIR)/$(NAME).exe + cp -rv bin/ /data/sda/shit/Release/ivshtest/ + chmod -Rvv 777 /data/sda/shit/Release/ivshtest + dumpinfo: echo -e "\e[96mBuilding configuration $(CONFIG) for arch $(ARCH) \e[0m" diff --git a/agent/build/configs.mk b/agent/build/configs.mk index d36be51..cc4a3fd 100644 --- a/agent/build/configs.mk +++ b/agent/build/configs.mk @@ -8,7 +8,7 @@ SDK_INCLUDES := -isystem $(VS2022_PATH)/ucrt/include -isystem $(VS2022_PATH)/crt # Windows 6.1 BASE_CCFLAGS := $(SDK_INCLUDES) -D_WIN32_WINNT=0x0601 -Ires -std=c17 -BASE_CXXFLAGS := $(SDK_INCLUDES) -D_CRT_SECURE_NO_WARNINGS -fno-exceptions -fno-rtti -D_WIN32_WINNT=0x0601 -Isrc -Ires -std=c++23 +BASE_CXXFLAGS := $(SDK_INCLUDES) -D_CRT_SECURE_NO_WARNINGS -fno-exceptions -fno-rtti -D_WIN32_WINNT=0x0601 -I../shared/src -Isrc -Ires -std=c++23 Release_Valid = yes Release_CCFLAGS = -O3 -ffast-math -DNDEBUG diff --git a/agent/sdk/inc/virtio/ivshmem.h b/agent/sdk/inc/virtio/ivshmem.h new file mode 100644 index 0000000..3f98491 --- /dev/null +++ b/agent/sdk/inc/virtio/ivshmem.h @@ -0,0 +1,69 @@ +#include + +DEFINE_GUID (GUID_DEVINTERFACE_IVSHMEM, + 0xdf576976,0x569d,0x4672,0x95,0xa0,0xf5,0x7e,0x4e,0xa0,0xb2,0x10); +// {df576976-569d-4672-95a0-f57e4ea0b210} + +typedef UINT16 IVSHMEM_PEERID; +typedef UINT64 IVSHMEM_SIZE; + +#define IVSHMEM_CACHE_NONCACHED 0 +#define IVSHMEM_CACHE_CACHED 1 +#define IVSHMEM_CACHE_WRITECOMBINED 2 + +/* + This structure is for use with the IOCTL_IVSHMEM_REQUEST_MMAP IOCTL +*/ +typedef struct IVSHMEM_MMAP_CONFIG +{ + UINT8 cacheMode; // the caching mode of the mapping, see IVSHMEM_CACHE_* for options +} +IVSHMEM_MMAP_CONFIG, *PIVSHMEM_MMAP_CONFIG; + +/* + This structure is for use with the IOCTL_IVSHMEM_REQUEST_MMAP IOCTL +*/ +typedef struct IVSHMEM_MMAP +{ + IVSHMEM_PEERID peerID; // our peer id + IVSHMEM_SIZE size; // the size of the memory region + PVOID ptr; // pointer to the memory region + UINT16 vectors; // the number of vectors available +} +IVSHMEM_MMAP, *PIVSHMEM_MMAP; + +/* + This structure is for use with the IOCTL_IVSHMEM_RING_DOORBELL IOCTL +*/ +typedef struct IVSHMEM_RING +{ + IVSHMEM_PEERID peerID; // the id of the peer to ring + UINT16 vector; // the doorbell to ring +} +IVSHMEM_RING, *PIVSHMEM_RING; + +/* + This structure is for use with the IOCTL_IVSHMEM_REGISTER_EVENT IOCTL + + Please Note: + - The IVSHMEM driver has a hard limit of 32 events. + - Events that are singleShot are released after they have been set. + - At this time repeating events are only released when the driver device + handle is closed, closing the event handle doesn't release it from the + drivers list. While this won't cause a problem in the driver, it will + cause you to run out of event slots. + */ +typedef struct IVSHMEM_EVENT +{ + UINT16 vector; // the vector that triggers the event + HANDLE event; // the event to trigger + BOOLEAN singleShot; // set to TRUE if you want the driver to only trigger this event once +} +IVSHMEM_EVENT, *PIVSHMEM_EVENT; + +#define IOCTL_IVSHMEM_REQUEST_PEERID CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_BUFFERED, FILE_ANY_ACCESS) +#define IOCTL_IVSHMEM_REQUEST_SIZE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_BUFFERED, FILE_ANY_ACCESS) +#define IOCTL_IVSHMEM_REQUEST_MMAP CTL_CODE(FILE_DEVICE_UNKNOWN, 0x802, METHOD_BUFFERED, FILE_ANY_ACCESS) +#define IOCTL_IVSHMEM_RELEASE_MMAP CTL_CODE(FILE_DEVICE_UNKNOWN, 0x803, METHOD_BUFFERED, FILE_ANY_ACCESS) +#define IOCTL_IVSHMEM_RING_DOORBELL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x804, METHOD_BUFFERED, FILE_ANY_ACCESS) +#define IOCTL_IVSHMEM_REGISTER_EVENT CTL_CODE(FILE_DEVICE_UNKNOWN, 0x805, METHOD_BUFFERED, FILE_ANY_ACCESS) diff --git a/agent/src/main.cpp b/agent/src/main.cpp index 1e9033d..bdc6e1d 100755 --- a/agent/src/main.cpp +++ b/agent/src/main.cpp @@ -1,4 +1,5 @@ // clang-format off +#include "ivshmem.hpp" #pragma comment(lib, "ws2_32.lib") #include #include @@ -11,170 +12,69 @@ #include "Utils.hpp" // clang-format on -#pragma pack(push, 1) -enum class MessageType : u32 { - Resize, // tResizeMessage - Data, // tDataMessage -}; -struct tMessageHeader { - MessageType type; - u32 datalen; - // data -}; - -struct tResizeMessage { - u32 width; - u32 height; -}; - -struct tTile { - u32 tile_x; - u32 tile_y; - - u32 tile_width; - u32 tile_height; - - // u32 tile_rgba[tile_width*tile_height] -}; - -struct tDataMessage { - u32 tileCount; -}; - -#pragma pack(pop) +#include "atomic_spinlock.hpp" struct tileRect { u32 x, y, width, height; }; -// client for streamserver -class cStreamClient { - SOCKET tcpSocket; - - public: - cStreamClient() = default; - - ~cStreamClient() { - if(tcpSocket != -1) { - closesocket(tcpSocket); - tcpSocket = -1; - } - } - - bool Connect(const char* address, int port) { - tcpSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP); - if(tcpSocket == -1) { - return false; - } - - sockaddr_in clientSvc {}; - clientSvc.sin_family = AF_INET; - inet_pton(AF_INET, address, &clientSvc.sin_addr.s_addr); - clientSvc.sin_port = htons(port); - - if(connect(tcpSocket, (SOCKADDR*)&clientSvc, sizeof(clientSvc)) == SOCKET_ERROR) { - printf("No connection socket. Fuck you\n"); - return false; - } - - // disable Nagle's alrogithm - int nodelay = 1; - setsockopt(tcpSocket, IPPROTO_TCP, TCP_NODELAY, (const char*)&nodelay, sizeof(nodelay)); - return true; - } - - void SendResize(const tResizeMessage& resize) { - tMessageHeader header; - header.type = MessageType::Resize; - header.datalen = sizeof(tResizeMessage); - - send(tcpSocket, (const char*)&header, sizeof(header), 0); - send(tcpSocket, (const char*)&resize, sizeof(resize), 0); - } - - void SendData(const u32* pData, u32 width, u32 height, const std::vector& tiles, bool sendFull) { - tMessageHeader header; - header.type = MessageType::Data; - // header.datalen = data.get_size() * sizeof(UINT32); - header.datalen = sizeof(tDataMessage); - - // send tile header - // bool sendFull = false; - tDataMessage dm; - - // if(tiles.empty()) { - // sendFull = true; - // } - - if(sendFull) { - dm.tileCount = 1; - } else { - dm.tileCount = (u32)tiles.size(); - } - - /*if(tiles.empty()) { - // Send the full screen as a "tile" - dm.tileCount = 1; - sendFull = true; - } else { - - }*/ - - // we have writev() at home - // writev() at home: - send(tcpSocket, (const char*)&header, sizeof(header), 0); - send(tcpSocket, (const char*)&dm, sizeof(dm), 0); - - // send each tile - if(!sendFull) { - for(auto& tile : tiles) { - tTile tile_wire { tile.x, tile.y, tile.width, tile.height }; - -#if 0 - std::vector tileData; - tileData.resize(tile.width * tile.height); - - for(u32 y = 0; y < tile.height; ++y) { - auto* pTileLineStart = &pData[(tile.y + y) * width + tile.x]; - memcpy(&tileData[y * tile.width], pTileLineStart, tile.width * sizeof(u32)); - } -#endif - - // send header - send(tcpSocket, (const char*)&tile_wire, (i32)sizeof(tile_wire), 0); - // send(tcpSocket, (const char*)&tileData[0], (i32)tileData.size() * 4, 0); - - // send data now - /*for (auto y = tile.y; y < tile.y + tile.height; ++y) { - auto* pTileLineStart = &pData[y * width + tile.x]; - send(tcpSocket, (const char*)pTileLineStart, tile.width * sizeof(UINT32), 0); - }*/ - } - } else { - tTile tDummyTile { 0, 0, width, height }; - send(tcpSocket, (const char*)&tDummyTile, sizeof(tDummyTile), 0); - // send(tcpSocket, (const char*)&pData[0], (i32)((width * height) * sizeof(u32)), 0); - } - - send(tcpSocket, (const char*)&pData[0], (i32)((width * height) * sizeof(u32)), 0); - // send(tcpSocket, (const char*)&data.data()[0], data.get_size() * sizeof(UINT32), 0); - } +struct Test { + hazelnut::AtomicSpinlock lk{}; + std::atomic sessionId {}; + std::atomic pingPong{}; }; int main(int argc, char** argv) { - WSADATA data; - if(WSAStartup(MAKEWORD(2, 2), &data) != NO_ERROR) { + hazelnut::IvshmemDevice dev; + + if(!dev.Open()) { + printf("Failed to open ivshmem device\n"); return 1; } - cStreamClient client; + auto size = dev.GetSize(); + auto ptr = (u32*)dev.GetPointer(); - if(!client.Connect("192.168.1.149", 9438)) { - printf("conn failed\n"); - return 1; + printf("opened a %zu MB large ivshmem\n", (size / (1024 * 1024))); + + // wipe the first 1mb + memset(&ptr[0], 0, 1*(1024*1024)); + + printf("wiped memory\n"); + + // Setup a test struct at the start of ivshmem + auto* pTest = new(&ptr[0]) Test; + + // reset pingpong counter + pTest->pingPong.store(0); + + u32 tries = 0; + u32 curTries = 10; + + pTest->sessionId.store(rand()); + + while(true) { + // lock + { + auto guard = pTest->lk.lock(); + Sleep(5); + } + + + printf("pingpong %u\n", pTest->pingPong.load()); + + if(tries++ == curTries) { + tries = 0; + pTest->pingPong.fetch_add(1); + } } + + + return 0; + +#if 0 printf("Hazelnut agent\n"); // Create a capture interface @@ -216,8 +116,7 @@ int main(int argc, char** argv) { continue; } - // send that to the server - client.SendData(framebuffer.pFramebuffer, framebuffer.width, framebuffer.height, tiles, firstFrame == false); + if(firstFrame) firstFrame = false; @@ -227,16 +126,12 @@ int main(int argc, char** argv) { diff = capture->GetDiffInformation(); firstFrame = true; - client.SendResize({ framebuffer.width, framebuffer.height }); - - // send empty frame - tiles.clear(); - client.SendData(framebuffer.pFramebuffer, framebuffer.width, framebuffer.height, tiles, true); } else { printf("Failed to capture\n"); break; } } +#endif return 0; } diff --git a/shared/README.md b/shared/README.md new file mode 100644 index 0000000..b369858 --- /dev/null +++ b/shared/README.md @@ -0,0 +1,3 @@ +# shared code + +This is shared between the two ends. It is cross platform C++20. \ No newline at end of file diff --git a/agent/src/Utils.hpp b/shared/src/Utils.hpp similarity index 93% rename from agent/src/Utils.hpp rename to shared/src/Utils.hpp index 0cae605..80cede3 100755 --- a/agent/src/Utils.hpp +++ b/shared/src/Utils.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include // common types using u8 = std::uint8_t; diff --git a/shared/src/atomic_spinlock.hpp b/shared/src/atomic_spinlock.hpp new file mode 100644 index 0000000..28c753c --- /dev/null +++ b/shared/src/atomic_spinlock.hpp @@ -0,0 +1,49 @@ +#pragma once + +#include +#include //debugging + +#include "Utils.hpp" + +#pragma pack(push, 1) + +namespace hazelnut { + + /// A atomic spinlock. Can be held across threads, or virtual machines. + struct AtomicSpinlock { + struct LockGuard { + ~LockGuard() { lock.unlock(); } + + protected: + friend AtomicSpinlock; + LockGuard(AtomicSpinlock& lock) : lock(lock) {} + AtomicSpinlock& lock; + }; + + AtomicSpinlock() { __lock.store(0, std::memory_order::seq_cst); } + + LockGuard lock() { + lock_manually(); + return LockGuard(*this); + } + + void lock_manually() { + u32 expected = 0; + while(!__lock.compare_exchange_strong(expected, 1, std::memory_order::seq_cst)) { + //printf("LOCK CONTENDED\n"); + expected = 0; + } + } + + bool try_lock_manually() { + u32 expected = 0; + return __lock.compare_exchange_strong(expected, 1, std::memory_order::seq_cst); + } + + void unlock() { __lock.store(0, std::memory_order::seq_cst); } + + std::atomic_uint32_t __lock; + }; +} // namespace hazelnut + +#pragma pack(pop) \ No newline at end of file diff --git a/shared/src/build_test.sh b/shared/src/build_test.sh new file mode 100755 index 0000000..009a74e --- /dev/null +++ b/shared/src/build_test.sh @@ -0,0 +1,3 @@ +#!/bin/bash +set -xeuo pipefail +g++ -O3 -std=c++20 ivshmem.cpp test.cpp -o test diff --git a/shared/src/ivshmem.cpp b/shared/src/ivshmem.cpp new file mode 100644 index 0000000..a77b45c --- /dev/null +++ b/shared/src/ivshmem.cpp @@ -0,0 +1,40 @@ +#include "ivshmem.hpp" + +#include + +#ifdef _WIN32 +#include "ivshmem_windows.cpp" +#elif __linux__ +#include "ivshmem_linux.cpp" +#else +#error YOU fucked +#endif + +namespace hazelnut { + + IvshmemDevice::IvshmemDevice() { + pImpl = new Impl; + } + + IvshmemDevice::~IvshmemDevice() { + Close(); + delete pImpl; + } + + bool IvshmemDevice::Open(const char* devName) { + return pImpl->Open(devName); + } + + void IvshmemDevice::Close() { + pImpl->Close(); + } + + void* IvshmemDevice::GetPointer() { + return pImpl->GetPointer(); + } + + usize IvshmemDevice::GetSize() { + return pImpl->GetSize(); + } + +} // namespace hazelnut \ No newline at end of file diff --git a/shared/src/ivshmem.hpp b/shared/src/ivshmem.hpp new file mode 100644 index 0000000..197be0d --- /dev/null +++ b/shared/src/ivshmem.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include "Utils.hpp" + +namespace hazelnut { + + /// Cross-platform IVSHMEM. + struct IvshmemDevice { + IvshmemDevice(); + + // forbid copy/move + IvshmemDevice(const IvshmemDevice&) = delete; + IvshmemDevice(IvshmemDevice&&) = delete; + + ~IvshmemDevice(); + + /// Opens the device. On Linux [devName] MUST point to a SHMEM file. + bool Open(const char* devName = nullptr); + + void Close(); + + /// Gets pointer to IVSHMEM memory. + void* GetPointer(); + + usize GetSize(); + + private: + struct Impl; + Impl* pImpl {}; + }; +} // namespace hazelnut \ No newline at end of file diff --git a/shared/src/ivshmem_linux.cpp b/shared/src/ivshmem_linux.cpp new file mode 100644 index 0000000..96c02bb --- /dev/null +++ b/shared/src/ivshmem_linux.cpp @@ -0,0 +1,82 @@ + +#include +#include +#include +#include +#include +#include + +#include "Utils.hpp" + +namespace hazelnut { + + struct IvshmemDevice::Impl { + int fd; + void* pMap = nullptr; + usize size = 0; + + bool OpenDevice(const char* devName) { + // devName is not optional on Linux. + if(devName == nullptr) + return false; + + int devFd = open(devName, O_RDWR, static_cast(0600)); + + if(devFd == -1) + return false; + + fd = devFd; + return true; + } + + bool MapMemory() { + // Device was opened, let's map memory + + struct stat st {}; + if(fstat(fd, &st) != 0) { + return false; + } + + usize devSize = st.st_size; + + void* map = mmap(0, devSize, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + if(map == MAP_FAILED) { + printf("map failed\n"); + return false; + } + + pMap = map; + size = devSize; + return true; + } + + void UnmapMemory() { + // Unmap memory. + if(pMap) + munmap(pMap, size); + } + + bool Open(const char* pDevName) { + if(!OpenDevice(pDevName)) + return false; + + if(!MapMemory()) { + Close(); + return false; + } + + return true; // mapped + } + + void Close() { + UnmapMemory(); + close(fd); + fd = -1; + } + + void* GetPointer() { return pMap; } + + usize GetSize() { return size; } + }; + +} // namespace hazelnut \ No newline at end of file diff --git a/shared/src/ivshmem_protocol.hpp b/shared/src/ivshmem_protocol.hpp new file mode 100644 index 0000000..e5feb5a --- /dev/null +++ b/shared/src/ivshmem_protocol.hpp @@ -0,0 +1,40 @@ +#pragma once +#include +#include "atomic_spinlock.hpp" + +namespace hazelnut { + +#pragma pack(push, 4) + + + /// Header for Hazelnut ivshmem. At 0x0 in the ivshmem memory space. + struct alignas(4096) IvshHeader { + std::atomic magic; + + // When this lock is held by the host, the host can read the rest of memory freely. + // We wait for the host to release before updating memory ourselves. + AtomicSpinlock lock; + + // See the next page boundary for FrameHeader + }; + + /// Stored at a page boundary + struct alignas(4096) FrameHeader { + /// The serial of the frame. If this is unchanged between a + /// lock-grab-unlock cycle, the frame has not changed. + std::atomic serial; + + // size + std::atomic width; + std::atomic height; + + // obtain bits at the next page! + u32* bits() { + return reinterpret_cast(this + 1); + } + + }; + +#pragma pack(pop) + +} \ No newline at end of file diff --git a/shared/src/ivshmem_windows.cpp b/shared/src/ivshmem_windows.cpp new file mode 100644 index 0000000..e968bdf --- /dev/null +++ b/shared/src/ivshmem_windows.cpp @@ -0,0 +1,168 @@ + +// clang-format off +#include +#include + +// IVSHMEM ioctls +#include +// clang-format on + + +namespace hazelnut { + + template + struct free_deleter { + void operator()(T* ptr) { + if(ptr) + free(ptr); + } + }; + + // DOES NOT UNBLESS WITH ->~T() !!! + template + using unique_malloc_ptr = std::unique_ptr>; + + // does not bless with placement new + template + unique_malloc_ptr unique_malloc(usize len) { + return { reinterpret_cast(malloc(len)), free_deleter {} }; + } + + struct scoped_devinfo { + scoped_devinfo(HDEVINFO hDevInfo) : hDevInfo(hDevInfo) {} + + ~scoped_devinfo() { SetupDiDestroyDeviceInfoList(hDevInfo); } + + scoped_devinfo(const scoped_devinfo&) = delete; + scoped_devinfo(scoped_devinfo&&) = delete; + + operator HDEVINFO() { return hDevInfo; } + + private: + HDEVINFO hDevInfo; + }; + + struct IvshmemDevice::Impl { + HANDLE hIvshmem { INVALID_HANDLE_VALUE }; + IVSHMEM_MMAP ivshmemMmapStruct {}; + + bool OpenDevice() { + memset(&ivshmemMmapStruct, 0, sizeof(IVSHMEM_MMAP)); + + SP_DEVINFO_DATA devInfoData = { + .cbSize = sizeof(SP_DEVINFO_DATA), + }; + + SP_DEVICE_INTERFACE_DATA devInterfaceData { .cbSize = sizeof(SP_DEVICE_INTERFACE_DATA) }; + + auto hDevInfo = + scoped_devinfo { SetupDiGetClassDevsA(&GUID_DEVINTERFACE_IVSHMEM, nullptr, nullptr, DIGCF_PRESENT | DIGCF_DEVICEINTERFACE) }; + + // just do the first one for now + SetupDiEnumDeviceInfo(hDevInfo, 0, &devInfoData); + + if(auto le = GetLastError(); le != ERROR_SUCCESS && le != ERROR_NO_MORE_ITEMS) { + return false; + } + + if(SetupDiEnumDeviceInterfaces(hDevInfo, &devInfoData, &GUID_DEVINTERFACE_IVSHMEM, 0, &devInterfaceData) == FALSE) { + return false; + } + + DWORD dwReq = 0; + SetupDiGetDeviceInterfaceDetail(hDevInfo, &devInterfaceData, nullptr, 0, &dwReq, nullptr); + + if(!dwReq) { + return false; + } + + auto pData = unique_malloc(dwReq); + pData->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA); + + if(!SetupDiGetDeviceInterfaceDetail(hDevInfo, &devInterfaceData, pData.get(), dwReq, nullptr, nullptr)) { + return false; + } + + hIvshmem = CreateFile(pData->DevicePath, 0, 0, nullptr, OPEN_EXISTING, 0, nullptr); + + if(hIvshmem == INVALID_HANDLE_VALUE) { + return false; + } + + return true; + } + + template + bool Ioctl(u32 ctlCode, TInput& in, TOutput& out) { + // Needed because some versions of the IVSHMEM driver, notably the last one + // that works on Windows 7, assumes return length is not nullptr, and will crash + // the calling process if it is nullptr. Red Hat quality code everyone; it doesn't + // just envelope the Linux ecosystem - it even gets to take over the Windows one too! + DWORD dwDummyRetLength = 0; + + if(!DeviceIoControl(hIvshmem, ctlCode, reinterpret_cast(&in), sizeof(TInput), reinterpret_cast(&out), sizeof(TOutput), + &dwDummyRetLength, nullptr)) { + return false; + } + + return true; + } + + // Ioctl() overload for bare + bool Ioctl(u32 ctlCode) { + DWORD dwDummyRetLength = 0; + if(!DeviceIoControl(hIvshmem, ctlCode, nullptr, 0, nullptr, 0, &dwDummyRetLength, nullptr)) { + return false; + } + return true; + } + + bool MapMemory() { + // Device was opened, let's map memory + + DWORD dwDummyRetLength = 0; + + // Yay, one less syscall. + IVSHMEM_MMAP_CONFIG config = { .cacheMode = IVSHMEM_CACHE_WRITECOMBINED }; + return Ioctl(IOCTL_IVSHMEM_REQUEST_MMAP, config, this->ivshmemMmapStruct); + } + + void UnmapMemory() { + // Unmap memory. + if(ivshmemMmapStruct.ptr) { + // We do not care if this fails since the device will eventually be closed and this is done only + // on agent exit, so the map will be released by the kernel anyhow. + static_cast(Ioctl(IOCTL_IVSHMEM_RELEASE_MMAP)); + } + + memset(&ivshmemMmapStruct, 0, sizeof(IVSHMEM_MMAP)); + } + + bool Open(const char* pDevName) { + if(!OpenDevice()) + return false; + + if(!MapMemory()) { + Close(); + return false; + } + + return true; // mapped + } + + void Close() { + UnmapMemory(); + CloseHandle(hIvshmem); + hIvshmem = INVALID_HANDLE_VALUE; + } + + void* GetPointer() { + return ivshmemMmapStruct.ptr; + } + + usize GetSize() { + return ivshmemMmapStruct.size; + } + }; + +} // namespace hazelnut \ No newline at end of file diff --git a/shared/src/test.cpp b/shared/src/test.cpp new file mode 100644 index 0000000..8460c07 --- /dev/null +++ b/shared/src/test.cpp @@ -0,0 +1,78 @@ + +#include + +#include "atomic_spinlock.hpp" +#include "ivshmem.hpp" + +int main(int argc, char** argv) { + struct Test { + hazelnut::AtomicSpinlock lk {}; + std::atomic sessionId {}; + std::atomic pingPong {}; + }; + + hazelnut::IvshmemDevice dev; + + if(!dev.Open("/dev/shm/lg-win7")) { + printf("Failed to open ivshmem device\n"); + return 1; + } + + auto size = dev.GetSize(); + auto ptr = (u32*)dev.GetPointer(); + + printf("opened a %zu MB large ivshmem\n", (size / (1024 * 1024))); + + auto* pTest = (Test*)ptr; + u32 last = 0; + u32 nrSame = 0; + + u64 initalizedSessionId = pTest->sessionId.load(); + + while(true) { + //printf("%d\n", pTest->timestamp.load()); + + if(initalizedSessionId != pTest->sessionId.load()) { + printf("Agent restarted. Closing\n"); + break; + } + +// need monotonic way +#if 0 + if(auto curTimestamp = pTest->timestamp.load(); curTimestamp < lastTimestamp) { + //printf("Not connected or agent crashed %d, %d\n", pTest->timestamp.load() , time(nullptr)); + printf("Not connected or agent crashed. %zu\n", curTimestamp); + break; + } else { + lastTimestamp = curTimestamp; + } +#endif + + // lock for a bit + { + if(pTest->lk.try_lock_manually()) { + // failed to lock + continue; + } + + auto current = pTest->pingPong.load(); + + if(current == last) { + pTest->lk.unlock(); + continue; + } + + nrSame = 0; + + last = current; + printf("pingpong %u\n", current); + + pTest->lk.unlock(); + } + + // allow the vm some time to do whatever it is it wants to do + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + + return 0; +} \ No newline at end of file diff --git a/src/atomic_spinlock.rs b/src/atomic_spinlock.rs new file mode 100644 index 0000000..58ac3e0 --- /dev/null +++ b/src/atomic_spinlock.rs @@ -0,0 +1,37 @@ +use std::sync::atomic::AtomicU32; +use std::sync::atomic::Ordering; + +#[repr(C, packed(4))] +pub struct AtomicSpinlock { + lock: AtomicU32 +} + +pub struct AtomicSpinlockGuard<'a> { + lock: &'a mut AtomicSpinlock, +} + +impl Drop for AtomicSpinlockGuard { + fn drop(&mut self) { + self.lock.store(0, Ordering::SeqCst); + } +} + +impl AtomicSpinlock { + fn init(&mut self) { + self.lock.store(0, Ordering::SeqCst); + } + + + fn lock(&mut self) -> AtomicSpinlockGuard<'_> { + loop { + match self + .sync + .compare_exchange(cur, 1, Ordering::SeqCst, Ordering::SeqCst) + { + Ok(last) => return HzLockGuard { lock: self }, + Err(_) => {} + } + } + } +} + diff --git a/src/ivshmem.rs b/src/ivshmem.rs new file mode 100644 index 0000000..89a1874 --- /dev/null +++ b/src/ivshmem.rs @@ -0,0 +1,30 @@ +//! Utilities for ivshmem + +use std::sync::atomic::AtomicU32; +use std::sync::atomic::Ordering; + +use crate::atomic_spinlock::AtomicSpinlock; + +// at 0x0 in ivshmem map +#[repr(C, align(4096))] +pub struct HzHeader { + magic: AtomicU32, + sync: AtomicSpinlock, + + // frame serial. if this is unchanged from prev. load it's not the same + frame_serial: AtomicU32, +} + +impl HzHeader { + fn valid(&self) -> bool { + self.magic.load(std::sync::atomic::Ordering::Acquire) == 0xabcd1234 + } + +} + +pub struct HzFrameHeader { + width: AtomicU32, + height: AtomicU32, +} + +pub struct Ivshmem {} diff --git a/src/main.rs b/src/main.rs index d9639bd..afd1b10 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,9 @@ use std::{ use byteorder::{LittleEndian, ReadBytesExt}; use minifb::{Window, WindowOptions}; +mod atomic_spinlock; +mod ivshmem; + #[derive(Debug)] enum Message { Resize { width: u32, height: u32 },