util/mem/shared_memory_client.hh - public/gem5 - Git at Google

 /*
  * Copyright 2022 Google, Inc.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
  * met: redistributions of source code must retain the above copyright
  * notice, this list of conditions and the following disclaimer;
  * redistributions in binary form must reproduce the above copyright
  * notice, this list of conditions and the following disclaimer in the
  * documentation and/or other materials provided with the distribution;
  * neither the name of the copyright holders nor the names of its
  * contributors may be used to endorse or promote products derived from
  * this software without specific prior written permission.
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */

 #ifndef __UTIL_MEM_SHARED_MEMORY_CLIENT_HH__
 #define __UTIL_MEM_SHARED_MEMORY_CLIENT_HH__

 #include <cinttypes>
 #include <cstdint>
 #include <cstdlib>
 #include <string>
 #include <unordered_map>
 #include <utility>

 #include <err.h>
 #include <errno.h>
 #include <sys/mman.h>
 #include <sys/socket.h>
 #include <sys/types.h>
 #include <sys/un.h>
 #include <unistd.h>

 namespace gem5
 {
 namespace util
 {
 namespace memory
 {

 class SharedMemoryClient
 {
   public:
     enum RequestType : int
     {
         kGetPhysRange = 0
     };

     explicit SharedMemoryClient(const std::string& server_path);

     // Request to access the range [start, end] of physical memory from the
     // viewpoint of the gem5 system providing the shared memory service.
     // There is no guarantee that the physical memory range is not used by
     // others. It is the user's responsibility to make sure not to break other
     // services or IPs accessing the same range. For example, you might want to
     // configure the kernel running in gem5 simulator to reserve such range.
     void* MapMemory(uint64_t start, uint64_t end);

     // Unmap previous mapped region, no client is needed here.
     static bool UnmapMemory(void* mem);

   private:
     using AllocRecordStorage = std::unordered_map<void*, size_t>;

     int GetConnection();
     bool SendGetPhysRangeRequest(int sock_fd, uint64_t start, uint64_t end);
     bool RecvGetPhysRangeResponse(int sock_fd, int* ptr_fd, off_t* ptr_offset);
     void* DoMap(int shm_fd, off_t shm_offset, size_t size);

     bool SendAll(int sock_fd, const void* buffer, size_t size);

     static AllocRecordStorage& GetAllocRecordStorage();

     std::string server_path_;
 };

 inline SharedMemoryClient::SharedMemoryClient(const std::string& server_path)
     : server_path_(server_path)
 {
 }


 inline void*
 SharedMemoryClient::MapMemory(uint64_t start, uint64_t end)
 {
     void* mem = nullptr;
     int sock_fd = -1;
     int shm_fd = -1;
     off_t shm_offset;

     do {
         if (start > end) {
             warnx("invalid range %" PRIu64 "-%" PRIu64, start, end);
             break;
         }
         sock_fd = GetConnection();
         if (sock_fd < 0) {
             warnx("cannot connect to shared memory server");
             break;
         }
         if (!SendGetPhysRangeRequest(sock_fd, start, end)) {
             warnx("cannot send request to shared memory server");
             break;
         }
         if (!RecvGetPhysRangeResponse(sock_fd, &shm_fd, &shm_offset)) {
             warnx("failed to read shared memory server response");
             break;
         }
         mem = DoMap(shm_fd, shm_offset, end - start + 1);
         if (mem == nullptr) {
             warnx("failed to create memory mapping");
             break;
         }
     } while (false);

     if (sock_fd >= 0) {
         close(sock_fd);
     }
     if (shm_fd >= 0) {
         close(shm_fd);
     }

     return mem;
 }

 inline bool
 SharedMemoryClient::UnmapMemory(void* mem)
 {
     auto& storage = GetAllocRecordStorage();
     auto it = storage.find(mem);
     if (it == storage.end()) {
         return false;
     }
     if (munmap(mem, it->second) < 0) {
         warn("munmap failed");
         return false;
     }
     storage.erase(it);
     return true;
 }

 inline int
 SharedMemoryClient::GetConnection()
 {
     int sock_fd = socket(AF_UNIX, SOCK_STREAM, 0);
     if (sock_fd < 0) {
         warn("create unix socket failed");
         return -1;
     }

     sockaddr_un serv_addr;
     memset(&serv_addr, 0, sizeof(serv_addr));
     serv_addr.sun_family = AF_UNIX;
     strncpy(serv_addr.sun_path, server_path_.c_str(),
             sizeof(serv_addr.sun_path) - 1);
     if (strlen(serv_addr.sun_path) != server_path_.size()) {
         warnx("server address truncated");
         close(sock_fd);
         return -1;
     }
     if (connect(sock_fd, reinterpret_cast<sockaddr*>(&serv_addr),
                 sizeof(serv_addr)) < 0) {
         warn("connect failed");
         close(sock_fd);
         return -1;
     }
     return sock_fd;
 }

 inline bool
 SharedMemoryClient::SendGetPhysRangeRequest(int sock_fd, uint64_t start,
                                             uint64_t end)
 {
     int req_type = RequestType::kGetPhysRange;
     struct
     {
         uint64_t start;
         uint64_t end;
     } request = {start, end};
     return SendAll(sock_fd, &req_type, sizeof(req_type)) &&
            SendAll(sock_fd, &request, sizeof(request));
 }

 inline bool
 SharedMemoryClient::RecvGetPhysRangeResponse(int sock_fd, int* ptr_fd,
                                              off_t* ptr_offset)
 {
     if (!ptr_fd || !ptr_offset) {
         return false;
     }

     msghdr msg = {};
     // Setup ptr_offset as buffer.
     iovec io = {.iov_base = ptr_offset, .iov_len = sizeof(*ptr_offset)};
     msg.msg_iov = &io;
     msg.msg_iovlen = 1;
     // Setup buffer for fd.
     union
     {
         char buffer[CMSG_SPACE(sizeof(*ptr_fd))];
         struct cmsghdr align;
     } cmsgs;
     msg.msg_control = cmsgs.buffer;
     msg.msg_controllen = sizeof(cmsgs.buffer);
     cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
     cmsg->cmsg_level = SOL_SOCKET;
     cmsg->cmsg_type = SCM_RIGHTS;
     cmsg->cmsg_len = CMSG_LEN(sizeof(*ptr_fd));
     // Try receive the message.
     ssize_t retv = recvmsg(sock_fd, &msg, 0);
     if (retv < 0) {
         warn("recvmsg failed");
         return false;
     }
     if (retv != sizeof(*ptr_offset)) {
         warnx("cannot receive all response");
         return false;
     }
     memcpy(ptr_fd, CMSG_DATA(cmsg), sizeof(*ptr_fd));
     return true;
 }

 inline void*
 SharedMemoryClient::DoMap(int shm_fd, off_t shm_offset, size_t size)
 {
     void* mem = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd,
                      shm_offset);
     if (mem == MAP_FAILED) {
         warn("mmap failed");
         return nullptr;
     }
     // If we cannot record a new mapping, our mapping are probably corrupted.
     if (!GetAllocRecordStorage().emplace(mem, size).second) {
         errx(EXIT_FAILURE, "cannot register memory mapping!");
     }
     return mem;
 }

 inline bool
 SharedMemoryClient::SendAll(int sock_fd, const void* buffer, size_t size)
 {
     const char* char_buffer = reinterpret_cast<const char*>(buffer);
     for (size_t offset = 0; offset < size;) {
         ssize_t retv = send(sock_fd, char_buffer + offset, size - offset, 0);
         if (retv >= 0) {
             offset += retv;
         } else if (errno != EINTR) {
             warn("send failed");
             return false;
         }
     }
     return true;
 }

 inline SharedMemoryClient::AllocRecordStorage&
 SharedMemoryClient::GetAllocRecordStorage()
 {
     static auto storage = new SharedMemoryClient::AllocRecordStorage();
     return *storage;
 }

 } // namespace memory
 } // namespace util
 } // namespace gem5

 #endif  // __UTIL_MEM_SHARED_MEMORY_CLIENT_HH__
	/*
	* Copyright 2022 Google, Inc.
	*
	* Redistribution and use in source and binary forms, with or without
	* modification, are permitted provided that the following conditions are
	* met: redistributions of source code must retain the above copyright
	* notice, this list of conditions and the following disclaimer;
	* redistributions in binary form must reproduce the above copyright
	* notice, this list of conditions and the following disclaimer in the
	* documentation and/or other materials provided with the distribution;
	* neither the name of the copyright holders nor the names of its
	* contributors may be used to endorse or promote products derived from
	* this software without specific prior written permission.
	*
	* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
	* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
	* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
	* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
	* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
	* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
	* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
	* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
	* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
	* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
	* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
	*/

	#ifndef __UTIL_MEM_SHARED_MEMORY_CLIENT_HH__
	#define __UTIL_MEM_SHARED_MEMORY_CLIENT_HH__

	#include <cinttypes>
	#include <cstdint>
	#include <cstdlib>
	#include <string>
	#include <unordered_map>
	#include <utility>

	#include <err.h>
	#include <errno.h>
	#include <sys/mman.h>
	#include <sys/socket.h>
	#include <sys/types.h>
	#include <sys/un.h>
	#include <unistd.h>

	namespace gem5
	{
	namespace util
	{
	namespace memory
	{

	class SharedMemoryClient
	{
	public:
	enum RequestType : int
	{
	kGetPhysRange = 0
	};

	explicit SharedMemoryClient(const std::string& server_path);

	// Request to access the range [start, end] of physical memory from the
	// viewpoint of the gem5 system providing the shared memory service.
	// There is no guarantee that the physical memory range is not used by
	// others. It is the user's responsibility to make sure not to break other
	// services or IPs accessing the same range. For example, you might want to
	// configure the kernel running in gem5 simulator to reserve such range.
	void* MapMemory(uint64_t start, uint64_t end);

	// Unmap previous mapped region, no client is needed here.
	static bool UnmapMemory(void* mem);

	private:
	using AllocRecordStorage = std::unordered_map<void*, size_t>;

	int GetConnection();
	bool SendGetPhysRangeRequest(int sock_fd, uint64_t start, uint64_t end);
	bool RecvGetPhysRangeResponse(int sock_fd, int* ptr_fd, off_t* ptr_offset);
	void* DoMap(int shm_fd, off_t shm_offset, size_t size);

	bool SendAll(int sock_fd, const void* buffer, size_t size);

	static AllocRecordStorage& GetAllocRecordStorage();

	std::string server_path_;
	};

	inline SharedMemoryClient::SharedMemoryClient(const std::string& server_path)
	: server_path_(server_path)
	{
	}


	inline void*
	SharedMemoryClient::MapMemory(uint64_t start, uint64_t end)
	{
	void* mem = nullptr;
	int sock_fd = -1;
	int shm_fd = -1;
	off_t shm_offset;

	do {
	if (start > end) {
	warnx("invalid range %" PRIu64 "-%" PRIu64, start, end);
	break;
	}
	sock_fd = GetConnection();
	if (sock_fd < 0) {
	warnx("cannot connect to shared memory server");
	break;
	}
	if (!SendGetPhysRangeRequest(sock_fd, start, end)) {
	warnx("cannot send request to shared memory server");
	break;
	}
	if (!RecvGetPhysRangeResponse(sock_fd, &shm_fd, &shm_offset)) {
	warnx("failed to read shared memory server response");
	break;
	}
	mem = DoMap(shm_fd, shm_offset, end - start + 1);
	if (mem == nullptr) {
	warnx("failed to create memory mapping");
	break;
	}
	} while (false);

	if (sock_fd >= 0) {
	close(sock_fd);
	}
	if (shm_fd >= 0) {
	close(shm_fd);
	}

	return mem;
	}

	inline bool
	SharedMemoryClient::UnmapMemory(void* mem)
	{
	auto& storage = GetAllocRecordStorage();
	auto it = storage.find(mem);
	if (it == storage.end()) {
	return false;
	}
	if (munmap(mem, it->second) < 0) {
	warn("munmap failed");
	return false;
	}
	storage.erase(it);
	return true;
	}

	inline int
	SharedMemoryClient::GetConnection()
	{
	int sock_fd = socket(AF_UNIX, SOCK_STREAM, 0);
	if (sock_fd < 0) {
	warn("create unix socket failed");
	return -1;
	}

	sockaddr_un serv_addr;
	memset(&serv_addr, 0, sizeof(serv_addr));
	serv_addr.sun_family = AF_UNIX;
	strncpy(serv_addr.sun_path, server_path_.c_str(),
	sizeof(serv_addr.sun_path) - 1);
	if (strlen(serv_addr.sun_path) != server_path_.size()) {
	warnx("server address truncated");
	close(sock_fd);
	return -1;
	}
	if (connect(sock_fd, reinterpret_cast<sockaddr*>(&serv_addr),
	sizeof(serv_addr)) < 0) {
	warn("connect failed");
	close(sock_fd);
	return -1;
	}
	return sock_fd;
	}

	inline bool
	SharedMemoryClient::SendGetPhysRangeRequest(int sock_fd, uint64_t start,
	uint64_t end)
	{
	int req_type = RequestType::kGetPhysRange;
	struct
	{
	uint64_t start;
	uint64_t end;
	} request = {start, end};
	return SendAll(sock_fd, &req_type, sizeof(req_type)) &&
	SendAll(sock_fd, &request, sizeof(request));
	}

	inline bool
	SharedMemoryClient::RecvGetPhysRangeResponse(int sock_fd, int* ptr_fd,
	off_t* ptr_offset)
	{
	if (!ptr_fd \|\| !ptr_offset) {
	return false;
	}

	msghdr msg = {};
	// Setup ptr_offset as buffer.
	iovec io = {.iov_base = ptr_offset, .iov_len = sizeof(*ptr_offset)};
	msg.msg_iov = &io;
	msg.msg_iovlen = 1;
	// Setup buffer for fd.
	union
	{
	char buffer[CMSG_SPACE(sizeof(*ptr_fd))];
	struct cmsghdr align;
	} cmsgs;
	msg.msg_control = cmsgs.buffer;
	msg.msg_controllen = sizeof(cmsgs.buffer);
	cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	cmsg->cmsg_len = CMSG_LEN(sizeof(*ptr_fd));
	// Try receive the message.
	ssize_t retv = recvmsg(sock_fd, &msg, 0);
	if (retv < 0) {
	warn("recvmsg failed");
	return false;
	}
	if (retv != sizeof(*ptr_offset)) {
	warnx("cannot receive all response");
	return false;
	}
	memcpy(ptr_fd, CMSG_DATA(cmsg), sizeof(*ptr_fd));
	return true;
	}

	inline void*
	SharedMemoryClient::DoMap(int shm_fd, off_t shm_offset, size_t size)
	{
	void* mem = mmap(nullptr, size, PROT_READ \| PROT_WRITE, MAP_SHARED, shm_fd,
	shm_offset);
	if (mem == MAP_FAILED) {
	warn("mmap failed");
	return nullptr;
	}
	// If we cannot record a new mapping, our mapping are probably corrupted.
	if (!GetAllocRecordStorage().emplace(mem, size).second) {
	errx(EXIT_FAILURE, "cannot register memory mapping!");
	}
	return mem;
	}

	inline bool
	SharedMemoryClient::SendAll(int sock_fd, const void* buffer, size_t size)
	{
	const char* char_buffer = reinterpret_cast<const char*>(buffer);
	for (size_t offset = 0; offset < size;) {
	ssize_t retv = send(sock_fd, char_buffer + offset, size - offset, 0);
	if (retv >= 0) {
	offset += retv;
	} else if (errno != EINTR) {
	warn("send failed");
	return false;
	}
	}
	return true;
	}

	inline SharedMemoryClient::AllocRecordStorage&
	SharedMemoryClient::GetAllocRecordStorage()
	{
	static auto storage = new SharedMemoryClient::AllocRecordStorage();
	return *storage;
	}

	} // namespace memory
	} // namespace util
	} // namespace gem5

	#endif // __UTIL_MEM_SHARED_MEMORY_CLIENT_HH__