agent: LARGE cleanup

interfaces are like a good thing.
This commit is contained in:
Lily Tsuru 2024-11-23 21:29:59 -05:00
parent 4a37a97fee
commit c1bb08c31f
11 changed files with 330 additions and 224 deletions

View file

@ -8,10 +8,8 @@
#include <string>
#include <vector>
#include "capture.hpp"
#include "Utils.hpp"
#include "NvFBCLibrary.h"
#include <NvFBC/nvFBCToSys.h>
// clang-format on
#pragma pack(push, 1)
@ -74,7 +72,6 @@ class cStreamClient {
sockaddr_in clientSvc {};
clientSvc.sin_family = AF_INET;
inet_pton(AF_INET, address, &clientSvc.sin_addr.s_addr);
// clientSvc.sin_addr.s_addr = inet_addr(address);
clientSvc.sin_port = htons(port);
if(connect(tcpSocket, (SOCKADDR*)&clientSvc, sizeof(clientSvc)) == SOCKET_ERROR) {
@ -97,18 +94,19 @@ class cStreamClient {
send(tcpSocket, (const char*)&resize, sizeof(resize), 0);
}
void SendData(const unique_buffer<u32>& data, u32 width, u32 height, const std::vector<tileRect>& tiles) {
void SendData(const u32* pData, u32 width, u32 height, const std::vector<tileRect>& tiles) {
tMessageHeader header;
header.type = MessageType::Data;
// header.datalen = data.get_size() * sizeof(UINT32);
header.datalen = sizeof(tDataMessage);
auto* pData = data.data();
// send tile header
bool sendFull = false;
tDataMessage dm;
if(tiles.empty()) {
// Send the full screen as a "tile"
dm.tileCount = 1;
sendFull = true;
} else {
@ -125,17 +123,17 @@ class cStreamClient {
for(auto& tile : tiles) {
tTile tile_wire { tile.x, tile.y, tile.width, tile.height };
std::vector<u32> data;
data.resize(tile.width * tile.height * 4);
std::vector<u32> 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(&data[y * tile.width], pTileLineStart, tile.width * sizeof(u32));
memcpy(&tileData[y * tile.width], pTileLineStart, tile.width * sizeof(u32));
}
// send header
send(tcpSocket, (const char*)&tile_wire, (int)sizeof(tile_wire), 0);
send(tcpSocket, (const char*)&data[0], (int)data.size(), 0);
send(tcpSocket, (const char*)&tileData[0], tileData.size() * 4, 0);
// send data now
/*for (auto y = tile.y; y < tile.y + tile.height; ++y) {
@ -145,65 +143,14 @@ class cStreamClient {
}
} else {
tTile tDummyTile { 0, 0, width, height };
send(tcpSocket, (const char*)&tDummyTile, sizeof(tDummyTile), 0);
send(tcpSocket, (const char*)&data.data()[0], (i32)(data.get_size() * 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);
}
};
// FIXME: make this a class, I've clearly got the ability to pull it into one
// stuff
NvFBCLibrary* nvfbcLibrary = nullptr;
NvFBCToSys* nvfbcToSys = nullptr;
// filled in by NVFBC
unsigned char* frameBuffer = nullptr;
unsigned char* diffMap = nullptr;
bool nvfbcSetup() {
//! Setup the frame grab
NVFBC_TOSYS_SETUP_PARAMS fbcSysSetupParams = { 0 };
fbcSysSetupParams.dwVersion = NVFBC_TOSYS_SETUP_PARAMS_VER;
fbcSysSetupParams.eMode = NVFBC_TOSYS_ARGB;
fbcSysSetupParams.bWithHWCursor = true;
fbcSysSetupParams.ppBuffer = (void**)&frameBuffer;
// enable a 32x32 difference map
fbcSysSetupParams.bDiffMap = TRUE;
fbcSysSetupParams.ppDiffMap = (void**)&diffMap;
fbcSysSetupParams.eDiffMapBlockSize = NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32;
auto status = nvfbcToSys->NvFBCToSysSetUp(&fbcSysSetupParams);
return status == NVFBC_SUCCESS;
}
void nvfbcDestroy() {
if(nvfbcToSys) {
nvfbcToSys->NvFBCToSysRelease();
diffMap = nullptr;
frameBuffer = nullptr;
}
}
bool nvfbcCreate() {
if(nvfbcToSys) {
nvfbcDestroy();
}
DWORD maxDisplayWidth = -1, maxDisplayHeight = -1;
//! Create an instance of NvFBCToSys
nvfbcToSys = (NvFBCToSys*)nvfbcLibrary->create(NVFBC_TO_SYS, &maxDisplayWidth, &maxDisplayHeight);
if(!nvfbcToSys) {
return false;
}
return true;
}
int main(int argc, char** argv) {
WSADATA data;
@ -211,94 +158,54 @@ int main(int argc, char** argv) {
return 1;
}
cStreamClient* client = new cStreamClient();
cStreamClient client;
if(!client->Connect("192.168.1.149", 9438)) {
if(!client.Connect("192.168.1.149", 9438)) {
printf("conn failed\n");
return 1;
}
nvfbcLibrary = new NvFBCLibrary();
//! Load NvFBC
if(!nvfbcLibrary->load()) {
fprintf(stderr, "Unable to load the NvFBC library\n");
return -1;
// Create a capture interface
auto capture = hazelnut::CreateFramebufferCapture(hazelnut::GuessBestCaptureInterface());
if (!capture) {
printf("Failed to create a capture interface\n");
}
if(!nvfbcCreate()) {
fprintf(stderr, "Unable to create an instance of NvFBC\n");
return -1;
}
printf("Successfully created a framebuffer capture interface\n");
u32 width = 0;
u32 height = 0;
bool firstFrame = true;
NVFBCRESULT status = NVFBC_SUCCESS;
BOOL bRecoveryDone = FALSE;
NvFBCFrameGrabInfo grabInfo;
unique_buffer<u32> buffer;
if(nvfbcSetup()) {
// Sleep so that ToSysSetUp forces a framebuffer update
if(capture->Initialize()) {
// Sleep
Sleep(100);
NVFBC_TOSYS_GRAB_FRAME_PARAMS fbcSysGrabParams = { 0 };
bool firstFrame = true;
std::vector<tileRect> tiles {};
// set up grab parameters
fbcSysGrabParams.dwVersion = NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER;
fbcSysGrabParams.dwFlags = NVFBC_TOSYS_WAIT_WITH_TIMEOUT;
fbcSysGrabParams.dwTargetWidth = 0;
fbcSysGrabParams.dwTargetHeight = 0;
fbcSysGrabParams.dwStartX = 0;
fbcSysGrabParams.dwStartY = 0;
fbcSysGrabParams.eGMode = NVFBC_TOSYS_SOURCEMODE_FULL;
fbcSysGrabParams.pNvFBCFrameGrabInfo = &grabInfo;
fbcSysGrabParams.dwWaitTime = 16;
hazelnut::FramebufferInformation framebuffer{};
hazelnut::DiffInformation diff{};
while(true) {
// Grab the frame
status = nvfbcToSys->NvFBCToSysGrabFrame(&fbcSysGrabParams);
if(status == NVFBC_SUCCESS) {
bRecoveryDone = FALSE;
// handle resizing the buffer
if(width != grabInfo.dwWidth || height != grabInfo.dwHeight) {
width = grabInfo.dwWidth;
height = grabInfo.dwHeight;
buffer.resize(grabInfo.dwWidth * grabInfo.dwHeight);
auto result = capture->CaptureFrame();
if (result == hazelnut::DisplayCaptureResult::Ok || result == hazelnut::DisplayCaptureResult::OkButResized) {
// Check for resize.
if (result == hazelnut::DisplayCaptureResult::OkButResized) {
framebuffer = capture->GetFramebufferInformation();
diff = capture->GetDiffInformation();
firstFrame = true;
client->SendResize(tResizeMessage { width, height });
}
// splat the data into an unpadded buffer
// REMOVE THIS
for(u32 y = 0; y < grabInfo.dwHeight; ++y) {
memcpy(&buffer.data()[y * grabInfo.dwWidth], &frameBuffer[(y * grabInfo.dwBufferWidth) * 4], grabInfo.dwWidth * 4);
client.SendResize({ framebuffer.width, framebuffer.height });
}
tiles.clear();
// diffmap
if(firstFrame == false) {
u32 dwDiffMapWidth = (u32)ceil((f32)width / 32);
u32 dwDiffMapHeight = (u32)ceil((f32)height / 32);
for(u32 y = 0; y < dwDiffMapHeight; ++y) {
for(u32 x = 0; x < dwDiffMapWidth; ++x) {
auto& bl = diffMap[y * dwDiffMapWidth + x];
for(u32 y = 0; y < diff.diffMapHeight; ++y) {
for(u32 x = 0; x < diff.diffmapWidth; ++x) {
auto& bl = diff.pDiffMap[y * diff.diffmapWidth + x];
if(bl != 0) {
tiles.push_back(tileRect {
x * (width / dwDiffMapWidth), // x
y * (height / dwDiffMapHeight), // y
width / dwDiffMapWidth, // width
height / dwDiffMapHeight // height
x * (framebuffer.width / diff.diffmapWidth), // x
y * (framebuffer.height / diff.diffMapHeight), // y
framebuffer.width / diff.diffmapWidth, // width
framebuffer.height / diff.diffMapHeight // height
});
}
}
@ -306,49 +213,16 @@ int main(int argc, char** argv) {
}
// send that to the server
client->SendData(buffer, width, height, tiles);
client.SendData(framebuffer.pFramebuffer, framebuffer.width, framebuffer.height, tiles);
if(firstFrame)
firstFrame = false;
} else {
if(bRecoveryDone == TRUE) {
fprintf(stderr, "Unable to recover from NvFBC Frame grab failure.\n");
break;
}
if(status == NVFBC_ERROR_DYNAMIC_DISABLE) {
fprintf(stderr, "NvFBC disabled. Quitting\n");
break;
}
// Try to recover the session
if(status == NVFBC_ERROR_INVALIDATED_SESSION) {
fprintf(stderr, "Session Invalidated. Attempting recovery\n");
if(!nvfbcCreate()) {
fprintf(stderr, "Unable to re-create NvFBC\n");
break;
}
if(nvfbcSetup()) {
bRecoveryDone = TRUE;
} else {
fprintf(stderr, "Unable to recover from NvFBC Frame grab failure.\n");
break;
}
}
printf("Failed to capture\n");
break;
}
}
}
if(status != NVFBC_SUCCESS) {
fprintf(stderr, "Unable to setup frame grab.\n");
}
nvfbcDestroy();
delete client;
delete nvfbcLibrary;
return 0;
}

View file

@ -146,19 +146,14 @@
<Text Include="ReadMe.txt" />
</ItemGroup>
<ItemGroup>
<ClInclude Include="capture.hpp" />
<ClInclude Include="NvFBCLibrary.h" />
<ClInclude Include="stdafx.h" />
<ClInclude Include="targetver.h" />
<ClInclude Include="Utils.hpp" />
</ItemGroup>
<ItemGroup>
<ClCompile Include="capture.cpp" />
<ClCompile Include="capture_nvfbc.cpp" />
<ClCompile Include="CollabVMFbcAgent.cpp" />
<ClCompile Include="stdafx.cpp">
<PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">Create</PrecompiledHeader>
<PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">Create</PrecompiledHeader>
<PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">Create</PrecompiledHeader>
<PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Release|x64'">Create</PrecompiledHeader>
</ClCompile>
</ItemGroup>
<Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
<ImportGroup Label="ExtensionTargets">

View file

@ -5,37 +5,29 @@
<UniqueIdentifier>{4FC737F1-C7A5-4376-A066-2A32D752A2FF}</UniqueIdentifier>
<Extensions>cpp;c;cc;cxx;def;odl;idl;hpj;bat;asm;asmx</Extensions>
</Filter>
<Filter Include="Header Files">
<UniqueIdentifier>{93995380-89BD-4b04-88EB-625FBE52EBFB}</UniqueIdentifier>
<Extensions>h;hh;hpp;hxx;hm;inl;inc;xsd</Extensions>
</Filter>
<Filter Include="Resource Files">
<UniqueIdentifier>{67DA6AB6-F800-4c08-8B7A-83BB121AAD01}</UniqueIdentifier>
<Extensions>rc;ico;cur;bmp;dlg;rc2;rct;bin;rgs;gif;jpg;jpeg;jpe;resx;tiff;tif;png;wav;mfcribbon-ms</Extensions>
</Filter>
</ItemGroup>
<ItemGroup>
<Text Include="ReadMe.txt" />
</ItemGroup>
<ItemGroup>
<ClInclude Include="stdafx.h">
<Filter>Header Files</Filter>
</ClInclude>
<ClInclude Include="targetver.h">
<Filter>Header Files</Filter>
</ClInclude>
<ClInclude Include="NvFBCLibrary.h">
<Filter>Header Files</Filter>
<ClInclude Include="capture.hpp">
<Filter>Source Files</Filter>
</ClInclude>
<ClInclude Include="Utils.hpp">
<Filter>Header Files</Filter>
<Filter>Source Files</Filter>
</ClInclude>
<ClInclude Include="NvFBCLibrary.h">
<Filter>Source Files</Filter>
</ClInclude>
</ItemGroup>
<ItemGroup>
<ClCompile Include="stdafx.cpp">
<ClCompile Include="CollabVMFbcAgent.cpp">
<Filter>Source Files</Filter>
</ClCompile>
<ClCompile Include="CollabVMFbcAgent.cpp">
<ClCompile Include="capture.cpp">
<Filter>Source Files</Filter>
</ClCompile>
<ClCompile Include="capture_nvfbc.cpp">
<Filter>Source Files</Filter>
</ClCompile>
</ItemGroup>

View file

@ -11,7 +11,7 @@ using u32 = std::uint32_t;
using i32 = std::int32_t;
using u64 = std::uint64_t;
using i64 = std::int64_t;
using usize = std::size_t;
using usize = size_t;
using f32 = float;
using f64 = double;

View file

@ -0,0 +1,30 @@
#include "capture.hpp"
namespace hazelnut {
// capture_nvfbc.cpp
IFramebufferCapture* CreateFramebufferCapture_NVFBC();
DisplayCaptureInterface GuessBestCaptureInterface() {
// The only one we support.
return DisplayCaptureInterface::NVFBC;
}
std::unique_ptr<IFramebufferCapture> CreateFramebufferCapture(DisplayCaptureInterface type) {
IFramebufferCapture* pCapture = nullptr;
switch(type) {
case DisplayCaptureInterface::NVFBC:
pCapture = CreateFramebufferCapture_NVFBC();
break;
default: return nullptr;
}
// Initalize capture.
if (pCapture->Initialize())
return nullptr;
return std::unique_ptr<IFramebufferCapture>(pCapture);
}
} // namespace hazelnut

View file

@ -0,0 +1,61 @@
#include "Utils.hpp"
#include <memory>
namespace hazelnut {
enum class DisplayCaptureResult {
Ok,
OkButResized, // OK, but grab the buffer information again. It's different due to a resize
Fail, // epic fail
};
// Framebuffer information
struct FramebufferInformation {
u32* pFramebuffer;
u32 width;
u32 height;
};
// Difference tile information
struct DiffInformation {
u8* pDiffMap;
u32 diffmapWidth;
u32 diffMapHeight;
};
/// Interface used for framebuffer capture independence.
class IFramebufferCapture {
public:
virtual ~IFramebufferCapture() = default;
virtual bool Initialize() = 0;
/// Performs the capture.
virtual DisplayCaptureResult CaptureFrame() = 0;
/// Get framebuffer information.
virtual FramebufferInformation GetFramebufferInformation() = 0;
/// Gets the tile difference block information.
virtual DiffInformation GetDiffInformation() = 0;
};
enum class DisplayCaptureInterface : u32 {
Invalid,
NVFBC, // NVFBC capture
// FIXME: DXGI support
};
/// Returns a guess for the best capture interface.
/// On a NVIDIA GPU on Windows > 8.1, we prefer NVFBC.
/// Otherwise, we should prefer DXGI.
///
/// If no API works Invalid is returned.
DisplayCaptureInterface GuessBestCaptureInterface();
/// Creates a instance of the framebuffer capture interface for a particular interface.
/// Returns a null pointer on failure to create.
std::unique_ptr<IFramebufferCapture> CreateFramebufferCapture(DisplayCaptureInterface type);
} // namespace hazelnut

View file

@ -0,0 +1,182 @@
#include "capture.hpp"
// clang-format off
#include "Utils.hpp"
#include "NvFBCLibrary.h"
#include <NvFBC/nvFBCToSys.h>
// clang-format on
namespace hazelnut {
class FramebufferCapture_NVFBC final : public IFramebufferCapture {
NVFBCRESULT nvfbcStatus;
NvFBCLibrary nvfbcLib;
NVFBC_TOSYS_GRAB_FRAME_PARAMS fbcSysGrabParams;
NvFBCFrameGrabInfo grabInfo;
NvFBCToSys* nvfbc;
u32 width;
u32 height;
u32* pRawFramebuffer;
unique_buffer<u32> convertedFramebuffer;
u8* pDiffMap;
u32 diffmapWidth;
u32 diffmapHeight;
public:
virtual ~FramebufferCapture_NVFBC() { Shutdown(); }
void Shutdown() {
NvfbcDestroyInstance();
nvfbcLib.close();
}
bool Initialize() override {
if (!nvfbcLib.load()) {
return false;
}
if (!NvfbcInitSession())
return false;
return true;
}
void NvfbcDestroyInstance() {
if (nvfbc) {
nvfbc->NvFBCToSysRelease();
nvfbc = nullptr;
// reset state while we're here I guess
pRawFramebuffer = nullptr;
pDiffMap = nullptr;
width = 0;
height = 0;
diffmapWidth = 0;
diffmapHeight = 0;
}
}
bool NvfbcCreateInstance() {
// Destroy an existing NVFBC session to avoid memory leak
if (nvfbc) {
NvfbcDestroyInstance();
}
DWORD maxDisplayWidth = -1;
DWORD maxDisplayHeight = -1;
nvfbc = (NvFBCToSys*)nvfbcLib.create(NVFBC_TO_SYS, &maxDisplayWidth, &maxDisplayHeight);
if (!nvfbc) {
return false;
}
return true;
}
bool NvfbcInitSession() {
if (!NvfbcCreateInstance())
return false;
// set up a session
NVFBC_TOSYS_SETUP_PARAMS fbcSysSetupParams = { 0 };
fbcSysSetupParams.dwVersion = NVFBC_TOSYS_SETUP_PARAMS_VER;
fbcSysSetupParams.eMode = NVFBC_TOSYS_ARGB;
fbcSysSetupParams.bWithHWCursor = true;
fbcSysSetupParams.ppBuffer = (void**)&pRawFramebuffer;
// enable a 32x32 difference map
fbcSysSetupParams.bDiffMap = TRUE;
fbcSysSetupParams.ppDiffMap = (void**)&pDiffMap;
fbcSysSetupParams.eDiffMapBlockSize = NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32;
auto status = nvfbc->NvFBCToSysSetUp(&fbcSysSetupParams);
if(status != NVFBC_SUCCESS)
return false;
memset(&fbcSysGrabParams, 0, sizeof(NVFBC_TOSYS_GRAB_FRAME_PARAMS));
// set up grab params
fbcSysGrabParams.dwVersion = NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER;
fbcSysGrabParams.dwFlags = NVFBC_TOSYS_WAIT_WITH_TIMEOUT;
fbcSysGrabParams.dwTargetWidth = 0;
fbcSysGrabParams.dwTargetHeight = 0;
fbcSysGrabParams.dwStartX = 0;
fbcSysGrabParams.dwStartY = 0;
fbcSysGrabParams.eGMode = NVFBC_TOSYS_SOURCEMODE_FULL;
fbcSysGrabParams.pNvFBCFrameGrabInfo = &grabInfo;
fbcSysGrabParams.dwWaitTime = 16;
return true;
}
DisplayCaptureResult CaptureFrame() override {
auto nvStatus = nvfbc->NvFBCToSysGrabFrame(&this->fbcSysGrabParams);
switch (nvStatus) {
case NVFBC_SUCCESS:
break;
// Need to recreate the session. If it fails then we fail too.
case NVFBC_ERROR_INVALIDATED_SESSION: {
if (!NvfbcInitSession())
return DisplayCaptureResult::Fail;
// Recurse. This looks naughty, but will allow us to directly retry
// the capture. (Plus if this causes issues whatever has happened is probably beyond saving)
return CaptureFrame();
} break;
default:
return DisplayCaptureResult::Fail;
}
if (width != grabInfo.dwWidth || height != grabInfo.dwHeight) {
width = grabInfo.dwWidth;
height = grabInfo.dwHeight;
diffmapWidth = (u32)ceil((f32)width / 32);
diffmapHeight = (u32)ceil((f32)height / 32);
convertedFramebuffer.resize(width * height);
return DisplayCaptureResult::OkButResized;
}
// Splat to the converted framebuffer. It doesn't have padding on it.
auto* pBufferData = convertedFramebuffer.data();
for (u32 y = 0; y < grabInfo.dwHeight; ++y) {
memcpy(&pBufferData[y * width], &pRawFramebuffer[(y * grabInfo.dwBufferWidth)], grabInfo.dwWidth * 4);
}
return DisplayCaptureResult::Ok;
}
FramebufferInformation GetFramebufferInformation() override {
return FramebufferInformation{
convertedFramebuffer.data(),
width,
height
};
}
DiffInformation GetDiffInformation() override {
return DiffInformation{
pDiffMap,
diffmapWidth,
diffmapHeight
};
}
};
IFramebufferCapture* CreateFramebufferCapture_NVFBC() {
return new FramebufferCapture_NVFBC();
}
} // namespace hazelnut

View file

@ -1,8 +0,0 @@
// stdafx.cpp : source file that includes just the standard includes
// CollabVMFbcAgent.pch will be the pre-compiled header
// stdafx.obj will contain the pre-compiled type information
#include "stdafx.h"
// TODO: reference any additional headers you need in STDAFX.H
// and not in this file

View file

@ -1,15 +0,0 @@
// stdafx.h : include file for standard system include files,
// or project specific include files that are used frequently, but
// are changed infrequently
//
#pragma once
#include "targetver.h"
#include <stdio.h>
#include <tchar.h>
// TODO: reference additional headers your program requires here

View file

@ -1,8 +0,0 @@
#pragma once
// Including SDKDDKVer.h defines the highest available Windows platform.
// If you wish to build your application for a previous Windows platform, include WinSDKVer.h and
// set the _WIN32_WINNT macro to the platform you wish to support before including SDKDDKVer.h.
#include <SDKDDKVer.h>

View file

@ -37,7 +37,7 @@ fn read_message(stream: &mut TcpStream, argb_buffer: &mut Vec<u32>, width: u32)
MESSAGETYPE_DATA => {
let tile_count = stream.read_u32::<LittleEndian>().expect("fuck");
println!("{tile_count} tiles");
//println!("{tile_count} tiles");
for i in 0..tile_count {
// tile rect
@ -80,6 +80,9 @@ fn main() {
let listener = TcpListener::bind("192.168.1.149:9438").expect("fuck");
let (mut socket, client_addr) = listener.accept().expect("FUCK!");
// disable nagles garbage bullshit
socket.set_nodelay(true).expect("fuck tcp");
let mut window = Window::new("FbcServer", 320, 200, WindowOptions::default())
.expect("you banned forever: rules do not");