blender-bridge / src / uds_server.h

#pragma once

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <queue>
#include <vector>

#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

// ── Message from a client ─────────────────────────────────────────────────────
struct UdsMessage {
	int client_id;
	uint32_t size;
	std::vector<uint8_t> data;
};

// ── Non-blocking UDS server ───────────────────────────────────────────────────
// Call poll() each frame from _process(). Length-prefixed messages (4-byte LE).
class UdsServer {
public:
	UdsServer() = default;

	~UdsServer() { stop(); }

	// Non-copyable
	UdsServer(const UdsServer &) = delete;
	UdsServer &operator=(const UdsServer &) = delete;

	bool start(const char *path) {
		// Remove stale socket
		unlink(path);

		_listen_fd = socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK, 0);
		if (_listen_fd < 0)
			return false;

		struct sockaddr_un addr = {};
		addr.sun_family = AF_UNIX;
		strncpy(addr.sun_path, path, sizeof(addr.sun_path) - 1);

		if (bind(_listen_fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
			close(_listen_fd);
			_listen_fd = -1;
			return false;
		}

		if (listen(_listen_fd, 4) < 0) {
			close(_listen_fd);
			_listen_fd = -1;
			return false;
		}

		strncpy(_path, path, sizeof(_path) - 1);
		return true;
	}

	void stop() {
		// Close all clients
		for (auto &c : _clients) {
			if (c.fd >= 0)
				close(c.fd);
		}
		_clients.clear();

		// Close listener
		if (_listen_fd >= 0) {
			close(_listen_fd);
			_listen_fd = -1;
		}

		// Remove socket file
		if (_path[0]) {
			unlink(_path);
			_path[0] = 0;
		}

		// Clear message queue
		while (!_messages.empty())
			_messages.pop();
	}

	bool listening() const { return _listen_fd >= 0; }

	int client_count() const { return (int)_clients.size(); }

	bool has_messages() const { return !_messages.empty(); }

	UdsMessage next_message() {
		UdsMessage msg = std::move(_messages.front());
		_messages.pop();
		return msg;
	}

	// Send length-prefixed message to all connected clients
	void broadcast(const void *data, uint32_t size) {
		for (auto &c : _clients) {
			if (c.fd >= 0) {
				_send_to(c.fd, data, size);
			}
		}
	}

	// Send length-prefixed message to a specific client
	void send_to(int client_id, const void *data, uint32_t size) {
		for (auto &c : _clients) {
			if (c.id == client_id && c.fd >= 0) {
				_send_to(c.fd, data, size);
				return;
			}
		}
	}

	// Poll for new connections and incoming messages. Call each frame.
	void poll() {
		if (_listen_fd < 0)
			return;

		_accept_clients();
		_read_clients();
		_prune_disconnected();
	}

private:
	struct ClientState {
		int fd = -1;
		int id = 0;
		// Receive buffer for partial messages
		std::vector<uint8_t> recv_buf;
	};

	int _listen_fd = -1;
	char _path[256] = {};
	int _next_client_id = 1;
	std::vector<ClientState> _clients;
	std::queue<UdsMessage> _messages;

	void _accept_clients() {
		// Accept all pending connections
		for (;;) {
			int fd = accept4(_listen_fd, nullptr, nullptr, SOCK_NONBLOCK);
			if (fd < 0)
				break;

			ClientState cs;
			cs.fd = fd;
			cs.id = _next_client_id++;
			_clients.push_back(std::move(cs));
		}
	}

	void _read_clients() {
		uint8_t tmp[4096];

		for (auto &c : _clients) {
			if (c.fd < 0)
				continue;

			// Read available data
			for (;;) {
				ssize_t n = read(c.fd, tmp, sizeof(tmp));
				if (n > 0) {
					c.recv_buf.insert(c.recv_buf.end(), tmp, tmp + n);
				} else if (n == 0) {
					// Client disconnected
					close(c.fd);
					c.fd = -1;
					break;
				} else {
					// EAGAIN/EWOULDBLOCK = no more data; EINTR = signal interrupted
					if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
						break;
					// Other error = disconnect
					close(c.fd);
					c.fd = -1;
					break;
				}
			}

			// Parse complete messages from buffer
			_parse_messages(c);
		}
	}

	void _parse_messages(ClientState &c) {
		while (c.recv_buf.size() >= 4) {
			uint32_t msg_len;
			memcpy(&msg_len, c.recv_buf.data(), 4);

			// Sanity: reject messages > 16 MB
			if (msg_len > 16 * 1024 * 1024) {
				close(c.fd);
				c.fd = -1;
				c.recv_buf.clear();
				return;
			}

			if (c.recv_buf.size() < 4 + msg_len)
				break; // incomplete

			UdsMessage msg;
			msg.client_id = c.id;
			msg.size = msg_len;
			msg.data.assign(
					c.recv_buf.begin() + 4,
					c.recv_buf.begin() + 4 + msg_len);

			_messages.push(std::move(msg));
			c.recv_buf.erase(c.recv_buf.begin(), c.recv_buf.begin() + 4 + msg_len);
		}
	}

	void _prune_disconnected() {
		_clients.erase(
				std::remove_if(_clients.begin(), _clients.end(),
						[](const ClientState &c) { return c.fd < 0; }),
				_clients.end());
	}

	static bool _write_all(int fd, const void *buf, size_t len) {
		const uint8_t *p = static_cast<const uint8_t *>(buf);
		size_t remaining = len;
		while (remaining > 0) {
			ssize_t n = write(fd, p, remaining);
			if (n > 0) {
				p += n;
				remaining -= (size_t)n;
			} else if (n < 0) {
				if (errno == EINTR)
					continue;
				if (errno == EAGAIN || errno == EWOULDBLOCK) {
					// Wait for socket to become writable (up to 2s)
					struct pollfd pfd = { fd, POLLOUT, 0 };
					int r = ::poll(&pfd, 1, 2000);
					if (r <= 0)
						return false; // Timeout or error
					continue;
				}
				return false; // Real error
			}
		}
		return true;
	}

	static void _send_to(int fd, const void *data, uint32_t size) {
		uint32_t le_size = size;
		if (_write_all(fd, &le_size, 4)) {
			_write_all(fd, data, size);
		}
	}
};