Fix GH issue #104 guest-side wrfsbase 0

Author: fwsGonzo

lib/tinykvm/machine.hpp

@@ -437,6 +437,11 @@ struct Machine
 	mutable std::unique_ptr<FileDescriptors> m_fds = nullptr;
 
 	Machine* m_remote = nullptr;
+	/* True only while a remote connection is live (between activate and
+	   disconnect). Must not be derived from guest-visible state: the guest
+	   can write FSBASE (CR4.FSGSBASE is enabled), so using its TLS base as
+	   the connection token lets a guest skip the disconnect cleanup. */
+	bool     m_remote_active = false;
 	uint32_t m_remote_connections = 0;
 
 	std::unique_ptr<MachineProfiling> m_profiling = nullptr;

lib/tinykvm/remote.cpp

@@ -178,16 +178,26 @@ void Machine::ipre_remote_resume_now(bool save_all, std::function<void(Machine&)
 	remote_vm.cpu().get_special_registers().fs.base =
 		this->get_special_registers().fs.base;
 	// After disconnect, access is no longer serialized (don't touch remote anymore)
-	const auto our_fsbase = this->remote_disconnect();
-	if (our_fsbase == 0)
+	// NOTE: The restored FSBASE is guest-writable and may be zero, so the
+	// connection state itself decides whether the disconnect succeeded.
+	if (!this->is_remote_connected())
 		throw std::runtime_error("ipre_resume_storage: Remote disconnect failed");
+	const auto our_fsbase = this->remote_disconnect();
 
 	// 6. When returning, restore original register state
 	copy_callee_saved_registers(save_all, this->registers(), saved_gprs);
 	this->registers().rax = this->registers().rsi; // RAX is return value
 	this->registers().rip += 2; // Skip over OUT instruction
 	if (save_all)
 		this->set_fpu_registers(saved_fprs);
+	if (our_fsbase == 0) {
+		// The resume stub skips wrfsbase when the pushed value is zero, so
+		// restore a zeroed guest FSBASE here instead of leaving the remote
+		// VM's FSBASE loaded after the remote memory is unmapped again.
+		auto& local_sprs = vcpu.get_special_registers();
+		local_sprs.fs.base = 0;
+		this->set_special_registers(local_sprs);
+	}
 	this->prepare_vmresume(our_fsbase, true);
 	vcpu.stopped = false;
 }
@@ -312,8 +322,10 @@ Machine::address_t Machine::remote_activate_now()
 
 	this->remote_connect(*this->m_remote, true);
 	this->m_remote_connections++;
+	this->m_remote_active = true;
 
-	// Set current FSBASE to remote original FSBASE
+	// Remember our own FSBASE, so that it can be restored on disconnect.
+	// NOTE: This is guest-writable and may legitimately be zero.
 	vcpu.remote_original_tls_base = get_fsgs().first;
 
 	auto& remote = *this->m_remote;
@@ -356,6 +368,7 @@ Machine::address_t Machine::remote_disconnect()
 {
 	if (!this->is_remote_connected())
 		return 0;
+	this->m_remote_active = false;
 
 	auto& remote = *this->m_remote;
 	remote.m_remote = nullptr; // Clear halfway state
@@ -397,7 +410,7 @@ Machine::address_t Machine::remote_disconnect()
 }
 bool Machine::is_remote_connected() const noexcept
 {
-	return this->m_remote != nullptr && this->vcpu.remote_original_tls_base != 0;
+	return this->m_remote != nullptr && this->m_remote_active;
 }
 
 bool Machine::is_foreign_address(address_t addr) const noexcept

tests/unit/remote.cpp

@@ -76,6 +76,112 @@ int main() {
 	REQUIRE(machine.remote_connection_count() == 1);
 }
 
+TEST_CASE("Remote call with guest-zeroed FSBASE must not wedge fork", "[Remote]")
+{
+	const auto storage_binary = build_and_load(R"M(
+extern long write(int, const void*, unsigned long);
+int main() {
+	return 1234;
+}
+extern void remote_hello_world() {
+	write(1, "Hello Remote World!", 19);
+}
+)M", "-Wl,-Ttext-segment=0x40400000");
+
+	// Extract storage remote symbols
+	const std::string command = "objcopy -w --extract-symbol --strip-symbol=!remote* --strip-symbol=* " + storage_binary.first + " storage.syms";
+	FILE* f = popen(command.c_str(), "r");
+	if (f == nullptr) {
+		throw std::runtime_error("Unable to extract remote symbols");
+	}
+	pclose(f);
+
+	const auto main_binary = build_and_load(R"M(
+extern void remote_hello_world();
+int main() {
+	return 2345;
+}
+extern long test_remote_call() {
+	remote_hello_world();
+	return 1;
+}
+extern long test_zero_fsbase() {
+	unsigned long fsbase;
+	__asm__ volatile("rdfsbase %0" : "=r"(fsbase));
+	/* CR4.FSGSBASE is enabled for guests, so a guest can zero its own
+	   FSBASE right before a remote call. That must not affect the
+	   connection state kept by the host. */
+	__asm__ volatile("wrfsbase %0" :: "r"(0UL));
+	remote_hello_world();
+	__asm__ volatile("wrfsbase %0" :: "r"(fsbase));
+	return 1;
+}
+)M", "-Wl,--just-symbols=storage.syms");
+
+	tinykvm::Machine storage { storage_binary.second, {
+		.max_mem = 16ULL << 20, // MB
+		.vmem_base_address = 1ULL << 30, // 1GB
+	} };
+	storage.setup_linux({"storage"}, env);
+	storage.run(4.0f);
+	REQUIRE(storage.return_value() == 1234);
+
+	unsigned remote_writes = 0;
+	storage.set_printer([&] (const char* data, size_t size) {
+		if (std::string_view{data, size} == "Hello Remote World!")
+			remote_writes ++;
+	});
+
+	tinykvm::Machine machine { main_binary.second, {
+		.max_mem = MAX_MEMORY
+	} };
+	machine.setup_linux({"main"}, env);
+	machine.remote_connect(storage);
+	machine.set_remote_allow_page_faults(true);
+	machine.run(4.0f);
+	REQUIRE(machine.return_value() == 2345);
+	REQUIRE(!machine.is_remote_connected());
+
+	machine.prepare_copy_on_write(1UL << 20);
+	const tinykvm::MachineOptions fork_options {
+		.max_mem = MAX_MEMORY,
+		.max_cow_mem = MAX_COWMEM,
+		.split_hugepages = true
+	};
+
+	// A guest that zeroes its FSBASE still performs a real remote call, and
+	// the disconnect afterwards must still happen.
+	tinykvm::Machine fork(machine, fork_options);
+	fork.set_remote_allow_page_faults(true);
+
+	fork.vmcall("test_zero_fsbase");
+	REQUIRE(fork.return_value() == 1);
+	REQUIRE(remote_writes == 1);
+	REQUIRE(!fork.is_remote_connected());
+	REQUIRE(fork.remote_connection_count() == 1);
+
+	// The fork must be recyclable: an ordinary remote call after reset works.
+	fork.reset_to(machine, fork_options);
+	fork.vmcall("test_remote_call");
+	REQUIRE(fork.return_value() == 1);
+	REQUIRE(remote_writes == 2);
+	REQUIRE(!fork.is_remote_connected());
+	REQUIRE(fork.remote_connection_count() == 2);
+
+	// Control: an identical sequence without wrfsbase.
+	tinykvm::Machine fork2(machine, fork_options);
+	fork2.set_remote_allow_page_faults(true);
+
+	fork2.vmcall("test_remote_call");
+	REQUIRE(fork2.return_value() == 1);
+	fork2.reset_to(machine, fork_options);
+	fork2.vmcall("test_remote_call");
+	REQUIRE(fork2.return_value() == 1);
+	REQUIRE(remote_writes == 4);
+	REQUIRE(!fork2.is_remote_connected());
+	REQUIRE(fork2.remote_connection_count() == 2);
+}
+
 TEST_CASE("Fail accessing remote VM directly", "[Remote]")
 {
 	const auto storage_binary = build_and_load(R"M(