Files
overlord/Overlord-Client/cmd/agent/capture/amfbridge/bridge_windows.cc
T

219 lines
10 KiB
C++

//go:build windows && cgo
#include "bridge_windows.h"
#include <windows.h>
#include <algorithm>
#include <chrono>
#include <cstring>
#include <string>
#include <thread>
#include <vector>
#include "../../../../third_party/amf/include/core/Factory.h"
#include "../../../../third_party/amf/include/core/Buffer.h"
#include "../../../../third_party/amf/include/components/VideoConverter.h"
#include "../../../../third_party/amf/include/components/VideoEncoderVCE.h"
namespace {
struct AMFEncoder {
HMODULE runtime = nullptr;
amf::AMFFactory* factory = nullptr;
amf::AMFContext* context = nullptr;
amf::AMFComponent* converter = nullptr;
amf::AMFComponent* encoder = nullptr;
std::vector<uint8_t> pending;
};
void set_error(char* dst, int capacity, const std::string& value) {
if (!dst || capacity <= 0) return;
const size_t count = std::min(value.size(), static_cast<size_t>(capacity - 1));
std::memcpy(dst, value.data(), count);
dst[count] = 0;
}
std::string amf_error(const char* stage, AMF_RESULT result) {
return std::string(stage) + " failed (AMF_RESULT=" + std::to_string(static_cast<int>(result)) + ")";
}
void destroy(AMFEncoder* state) {
if (!state) return;
if (state->encoder) {
state->encoder->Terminate();
state->encoder->Release();
}
if (state->converter) {
state->converter->Terminate();
state->converter->Release();
}
if (state->context) {
state->context->Terminate();
state->context->Release();
}
if (state->runtime) FreeLibrary(state->runtime);
delete state;
}
HMODULE load_runtime() {
return LoadLibraryW(sizeof(void*) == 8 ? L"amfrt64.dll" : L"amfrt32.dll");
}
AMF_RESULT query_with_wait(amf::AMFComponent* component, amf::AMFData** data, int timeout_ms) {
const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms);
AMF_RESULT result = AMF_REPEAT;
do {
result = component->QueryOutput(data);
if (result != AMF_REPEAT && result != AMF_NEED_MORE_INPUT) return result;
if (*data) return AMF_OK;
std::this_thread::sleep_for(std::chrono::milliseconds(1));
} while (std::chrono::steady_clock::now() < deadline);
return result;
}
AMF_RESULT collect_encoder_output(AMFEncoder* state, int timeout_ms) {
amf::AMFData* data = nullptr;
AMF_RESULT result = query_with_wait(state->encoder, &data, timeout_ms);
if (!data) {
return (result == AMF_REPEAT || result == AMF_NEED_MORE_INPUT) ? AMF_OK : result;
}
amf::AMFBufferPtr buffer(data);
data->Release();
if (!buffer) return AMF_INVALID_DATA_TYPE;
const auto* bytes = static_cast<const uint8_t*>(buffer->GetNative());
const size_t size = static_cast<size_t>(buffer->GetSize());
if (bytes && size) state->pending.insert(state->pending.end(), bytes, bytes + size);
return AMF_OK;
}
} // namespace
extern "C" int overlord_amf_probe(char* error_text, int error_capacity) {
HMODULE runtime = load_runtime();
if (!runtime) {
set_error(error_text, error_capacity, "AMF runtime DLL is not installed");
return 0;
}
auto init = reinterpret_cast<AMFInit_Fn>(GetProcAddress(runtime, AMF_INIT_FUNCTION_NAME));
amf::AMFFactory* factory = nullptr;
AMF_RESULT result = init ? init(AMF_FULL_VERSION, &factory) : AMF_FAIL;
FreeLibrary(runtime);
if (result != AMF_OK || !factory) {
set_error(error_text, error_capacity, amf_error("AMFInit", result));
return 0;
}
return 1;
}
extern "C" int overlord_amf_create(void* d3d11_device, int input_width, int input_height,
int encode_width, int encode_height, int fps,
uint32_t dxgi_format, int bitrate,
overlord_amf_encoder* output,
char* error_text, int error_capacity) {
if (!output || !d3d11_device) {
set_error(error_text, error_capacity, "invalid AMF create arguments");
return 0;
}
*output = nullptr;
auto* state = new AMFEncoder();
state->runtime = load_runtime();
if (!state->runtime) {
set_error(error_text, error_capacity, "AMF runtime DLL is not installed");
destroy(state);
return 0;
}
auto init = reinterpret_cast<AMFInit_Fn>(GetProcAddress(state->runtime, AMF_INIT_FUNCTION_NAME));
AMF_RESULT result = init ? init(AMF_FULL_VERSION, &state->factory) : AMF_FAIL;
if (result != AMF_OK || !state->factory) goto fail_init;
result = state->factory->CreateContext(&state->context);
if (result != AMF_OK || !state->context) goto fail_context;
result = state->context->InitDX11(d3d11_device, amf::AMF_DX11_0);
if (result != AMF_OK) goto fail_dx11;
result = state->factory->CreateComponent(state->context, AMFVideoConverter, &state->converter);
if (result != AMF_OK || !state->converter) goto fail_converter;
state->converter->SetProperty(AMF_VIDEO_CONVERTER_OUTPUT_FORMAT, static_cast<amf_int64>(amf::AMF_SURFACE_NV12));
state->converter->SetProperty(AMF_VIDEO_CONVERTER_MEMORY_TYPE, static_cast<amf_int64>(amf::AMF_MEMORY_DX11));
state->converter->SetProperty(AMF_VIDEO_CONVERTER_OUTPUT_SIZE, AMFConstructSize(encode_width, encode_height));
{
const amf::AMF_SURFACE_FORMAT input_format = dxgi_format == 28 ? amf::AMF_SURFACE_RGBA : amf::AMF_SURFACE_BGRA;
result = state->converter->Init(input_format, input_width, input_height);
}
if (result != AMF_OK) goto fail_converter_init;
result = state->factory->CreateComponent(state->context, AMFVideoEncoderVCE_AVC, &state->encoder);
if (result != AMF_OK || !state->encoder) goto fail_encoder;
state->encoder->SetProperty(AMF_VIDEO_ENCODER_USAGE, static_cast<amf_int64>(AMF_VIDEO_ENCODER_USAGE_ULTRA_LOW_LATENCY));
state->encoder->SetProperty(AMF_VIDEO_ENCODER_PROFILE, static_cast<amf_int64>(AMF_VIDEO_ENCODER_PROFILE_MAIN));
state->encoder->SetProperty(AMF_VIDEO_ENCODER_FRAMESIZE, AMFConstructSize(encode_width, encode_height));
state->encoder->SetProperty(AMF_VIDEO_ENCODER_FRAMERATE, AMFConstructRate(fps, 1));
state->encoder->SetProperty(AMF_VIDEO_ENCODER_TARGET_BITRATE, static_cast<amf_int64>(bitrate));
state->encoder->SetProperty(AMF_VIDEO_ENCODER_B_PIC_PATTERN, static_cast<amf_int64>(0));
state->encoder->SetProperty(AMF_VIDEO_ENCODER_LOWLATENCY_MODE, true);
state->encoder->SetProperty(AMF_VIDEO_ENCODER_HEADER_INSERTION_SPACING, static_cast<amf_int64>(60));
result = state->encoder->Init(amf::AMF_SURFACE_NV12, encode_width, encode_height);
if (result != AMF_OK) goto fail_encoder_init;
*output = state;
return 1;
fail_encoder_init: set_error(error_text, error_capacity, amf_error("H.264 encoder Init", result)); destroy(state); return 0;
fail_encoder: set_error(error_text, error_capacity, amf_error("CreateComponent encoder", result)); destroy(state); return 0;
fail_converter_init: set_error(error_text, error_capacity, amf_error("video converter Init", result)); destroy(state); return 0;
fail_converter: set_error(error_text, error_capacity, amf_error("CreateComponent converter", result)); destroy(state); return 0;
fail_dx11: set_error(error_text, error_capacity, amf_error("InitDX11", result)); destroy(state); return 0;
fail_context: set_error(error_text, error_capacity, amf_error("CreateContext", result)); destroy(state); return 0;
fail_init: set_error(error_text, error_capacity, amf_error("AMFInit", result)); destroy(state); return 0;
}
extern "C" int overlord_amf_encode(overlord_amf_encoder opaque, void* d3d11_texture,
int force_idr, uint8_t* output, int output_capacity,
int* output_size, char* error_text, int error_capacity) {
auto* state = static_cast<AMFEncoder*>(opaque);
if (!state || !d3d11_texture || !output_size) return 0;
*output_size = 0;
if (!state->pending.empty()) {
*output_size = static_cast<int>(state->pending.size());
if (output_capacity < *output_size || !output) return 2;
std::memcpy(output, state->pending.data(), state->pending.size());
state->pending.clear();
return 1;
}
amf::AMFSurface* source = nullptr;
AMF_RESULT result = state->context->CreateSurfaceFromDX11Native(d3d11_texture, &source, nullptr);
if (result != AMF_OK || !source) goto fail_source;
result = state->converter->SubmitInput(source);
source->Release();
if (result != AMF_OK) goto fail_converter_submit;
{
amf::AMFData* converted_data = nullptr;
result = query_with_wait(state->converter, &converted_data, 100);
if (result != AMF_OK || !converted_data) goto fail_converter_output;
amf::AMFSurfacePtr converted(converted_data);
converted_data->Release();
if (!converted) { result = AMF_INVALID_DATA_TYPE; goto fail_converter_output; }
if (force_idr) {
converted->SetProperty(AMF_VIDEO_ENCODER_FORCE_PICTURE_TYPE, static_cast<amf_int64>(AMF_VIDEO_ENCODER_PICTURE_TYPE_IDR));
converted->SetProperty(AMF_VIDEO_ENCODER_INSERT_SPS, true);
converted->SetProperty(AMF_VIDEO_ENCODER_INSERT_PPS, true);
}
result = state->encoder->SubmitInput(converted);
if (result != AMF_OK) goto fail_encoder_submit;
}
result = collect_encoder_output(state, 20);
if (result != AMF_OK) goto fail_encoder_output;
if (!state->pending.empty()) {
*output_size = static_cast<int>(state->pending.size());
if (output_capacity < *output_size || !output) return 2;
std::memcpy(output, state->pending.data(), state->pending.size());
state->pending.clear();
}
return 1;
fail_encoder_output: set_error(error_text, error_capacity, amf_error("encoder QueryOutput", result)); return 0;
fail_encoder_submit: set_error(error_text, error_capacity, amf_error("encoder SubmitInput", result)); return 0;
fail_converter_output: set_error(error_text, error_capacity, amf_error("converter QueryOutput", result)); return 0;
fail_converter_submit: set_error(error_text, error_capacity, amf_error("converter SubmitInput", result)); return 0;
fail_source: set_error(error_text, error_capacity, amf_error("CreateSurfaceFromDX11Native", result)); return 0;
}
extern "C" void overlord_amf_destroy(overlord_amf_encoder encoder) {
destroy(static_cast<AMFEncoder*>(encoder));
}