Reactor/lib/Socket.hpp

#pragma once

#include "InetAddress.hpp"
#include "Utilities.hpp"
#include <sys/socket.h>
#include <netinet/tcp.h>
#include <unistd.h>
#include <fcntl.h>
#include <system_error>
#include <optional>

namespace reactor
{

class Socket : public NonCopyable
{
private:
	int fd_;

public:
	/*
	 * Constructs a Socket by taking ownership of a file descriptor.
	 */
	explicit Socket(int fd)
		: fd_(fd)
	{
		LOG_TRACE << "Socket created with fd=" << fd_;
	}

	/*
	 * Destructor, closes the socket file descriptor.
	 */
	~Socket()
	{
		if (fd_ >= 0)
		{
			close(fd_);
			LOG_TRACE << "Socket fd=" << fd_ << " closed";
		}
	}

	/*
	 * Move constructor.
	 */
	Socket(Socket &&other) noexcept
		: fd_(other.fd_)
	{
		other.fd_ = -1;
		LOG_TRACE << "Socket moved fd=" << fd_;
	}

	/*
	 * Move assignment operator.
	 */
	Socket &operator=(Socket &&other) noexcept
	{
		if (this != &other)
		{
			if (fd_ >= 0)
			{
				close(fd_);
				LOG_TRACE << "Socket fd=" << fd_ << " closed in move assignment";
			}
			fd_ = other.fd_;
			other.fd_ = -1;
		}
		return *this;
	}

	/*
	 * Creates a non-blocking TCP socket.
	 * Throws std::runtime_error on failure.
	 */
	static Socket createTcp(bool ipv6 = false)
	{
		int fd = socket(ipv6 ? AF_INET6 : AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
		if (fd < 0)
		{
			throw std::runtime_error("Failed to create TCP socket: " + std::string(strerror(errno)));
		}
		LOG_DEBUG << "Created TCP socket fd=" << fd << " ipv6=" << ipv6;
		return Socket(fd);
	}

	/*
	 * Creates a non-blocking UDP socket.
	 * Throws std::runtime_error on failure.
	 */
	static Socket createUdp(bool ipv6 = false)
	{
		int fd = socket(ipv6 ? AF_INET6 : AF_INET, SOCK_DGRAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
		if (fd < 0)
		{
			throw std::runtime_error("Failed to create UDP socket: " + std::string(strerror(errno)));
		}
		LOG_DEBUG << "Created UDP socket fd=" << fd << " ipv6=" << ipv6;
		return Socket(fd);
	}

	/*
	 * Binds the socket to a specific address.
	 * Throws std::runtime_error on failure.
	 */
	void bind(const InetAddress &addr)
	{
		if (::bind(fd_, addr.getSockAddr(), addr.getSockLen()) < 0)
		{
			throw std::runtime_error("Socket bind to " + addr.toIpPort() + " failed: " + std::string(strerror(errno)));
		}
		LOG_INFO << "Socket fd=" << fd_ << " bound to " << addr.toIpPort();
	}

	/*
	 * Puts the socket in listening mode for incoming connections.
	 * Throws std::runtime_error on failure.
	 */
	void listen(int backlog = SOMAXCONN)
	{
		if (::listen(fd_, backlog) < 0)
		{
			throw std::runtime_error("Socket listen failed: " + std::string(strerror(errno)));
		}
		LOG_INFO << "Socket fd=" << fd_ << " listening with backlog=" << backlog;
	}

	/*
	 * Accepts a new connection.
	 * Returns an optional Socket. Returns nullopt if no pending connection.
	 */
	std::optional<Socket> accept(InetAddress &peerAddr)
	{
		sockaddr_in6 addr;
		socklen_t len = sizeof(addr);
		int connfd = accept4(fd_, reinterpret_cast<sockaddr *>(&addr), &len, SOCK_NONBLOCK | SOCK_CLOEXEC);

		if (connfd >= 0)
		{
			if (addr.sin6_family == AF_INET)
			{
				peerAddr = InetAddress(*reinterpret_cast<sockaddr_in *>(&addr));
			}
			else
			{
				peerAddr = InetAddress(addr);
			}
			LOG_DEBUG << "Socket fd=" << fd_ << " accepted connection fd=" << connfd
					  << " from " << peerAddr.toIpPort();
			return std::optional<Socket>(Socket(connfd));
		}
		else if (errno != EAGAIN && errno != EWOULDBLOCK)
		{
			LOG_ERROR << "Socket accept failed: " << strerror(errno);
		}

		return std::nullopt;
	}

	/*
	 * Establishes a connection to a specified address.
	 */
	int connect(const InetAddress &addr)
	{
		int ret = ::connect(fd_, addr.getSockAddr(), addr.getSockLen());
		if (ret < 0 && errno != EINPROGRESS)
		{
			LOG_ERROR << "Socket connect to " << addr.toIpPort() << " failed: " << strerror(errno);
		}
		else
		{
			LOG_DEBUG << "Socket fd=" << fd_ << " connecting to " << addr.toIpPort();
		}
		return ret;
	}

	ssize_t read(void *buf, size_t len)
	{
		ssize_t n = ::read(fd_, buf, len);
		if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR)
		{
			LOG_ERROR << "Socket read failed: " << strerror(errno);
		}
		return n;
	}

	ssize_t write(const void *buf, size_t len)
	{
		ssize_t n = ::write(fd_, buf, len);
		if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR)
		{
			LOG_ERROR << "Socket write failed: " << strerror(errno);
		}
		return n;
	}

	ssize_t sendTo(const void *buf, size_t len, const InetAddress &addr)
	{
		ssize_t n = ::sendto(fd_, buf, len, 0, addr.getSockAddr(), addr.getSockLen());
		if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK)
		{
			LOG_ERROR << "Socket sendto failed: " << strerror(errno);
		}
		return n;
	}

	ssize_t recvFrom(void *buf, size_t len, InetAddress &addr)
	{
		sockaddr_in6 sockaddr;
		socklen_t addrlen = sizeof(sockaddr);
		ssize_t n = ::recvfrom(fd_, buf, len, 0, reinterpret_cast<struct sockaddr *>(&sockaddr), &addrlen);

		if (n >= 0)
		{
			if (sockaddr.sin6_family == AF_INET)
			{
				addr = InetAddress(*reinterpret_cast<sockaddr_in *>(&sockaddr));
			}
			else
			{
				addr = InetAddress(sockaddr);
			}
		}
		else if (errno != EAGAIN && errno != EWOULDBLOCK)
		{
			LOG_ERROR << "Socket recvfrom failed: " << strerror(errno);
		}

		return n;
	}

	void setReuseAddr(bool on = true)
	{
		int optval = on ? 1 : 0;
		if (setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval)) < 0)
		{
			LOG_ERROR << "setsockopt SO_REUSEADDR failed: " << strerror(errno);
		}
		else
		{
			LOG_TRACE << "Socket fd=" << fd_ << " SO_REUSEADDR=" << on;
		}
	}

	void setReusePort(bool on = true)
	{
		int optval = on ? 1 : 0;
		if (setsockopt(fd_, SOL_SOCKET, SO_REUSEPORT, &optval, sizeof(optval)) < 0)
		{
			LOG_ERROR << "setsockopt SO_REUSEPORT failed: " << strerror(errno);
		}
		else
		{
			LOG_TRACE << "Socket fd=" << fd_ << " SO_REUSEPORT=" << on;
		}
	}

	void setTcpNoDelay(bool on = true)
	{
		int optval = on ? 1 : 0;
		if (setsockopt(fd_, IPPROTO_TCP, TCP_NODELAY, &optval, sizeof(optval)) < 0)
		{
			LOG_ERROR << "setsockopt TCP_NODELAY failed: " << strerror(errno);
		}
		else
		{
			LOG_TRACE << "Socket fd=" << fd_ << " TCP_NODELAY=" << on;
		}
	}

	void setKeepAlive(bool on = true)
	{
		int optval = on ? 1 : 0;
		if (setsockopt(fd_, SOL_SOCKET, SO_KEEPALIVE, &optval, sizeof(optval)) < 0)
		{
			LOG_ERROR << "setsockopt SO_KEEPALIVE failed: " << strerror(errno);
		}
		else
		{
			LOG_TRACE << "Socket fd=" << fd_ << " SO_KEEPALIVE=" << on;
		}
	}

	void setTcpKeepAlive(int idle, int interval, int count)
	{
		if (setsockopt(fd_, IPPROTO_TCP, TCP_KEEPIDLE, &idle, sizeof(idle)) < 0 ||
			setsockopt(fd_, IPPROTO_TCP, TCP_KEEPINTVL, &interval, sizeof(interval)) < 0 ||
			setsockopt(fd_, IPPROTO_TCP, TCP_KEEPCNT, &count, sizeof(count)) < 0)
		{
			LOG_ERROR << "setsockopt TCP_KEEP* failed: " << strerror(errno);
		}
		else
		{
			LOG_TRACE << "Socket fd=" << fd_ << " TCP keepalive: idle=" << idle
					  << " interval=" << interval << " count=" << count;
		}
	}

	void setRecvBuffer(int size)
	{
		if (setsockopt(fd_, SOL_SOCKET, SO_RCVBUF, &size, sizeof(size)) < 0)
		{
			LOG_ERROR << "setsockopt SO_RCVBUF failed: " << strerror(errno);
		}
		else
		{
			LOG_TRACE << "Socket fd=" << fd_ << " SO_RCVBUF=" << size;
		}
	}

	void setSendBuffer(int size)
	{
		if (setsockopt(fd_, SOL_SOCKET, SO_SNDBUF, &size, sizeof(size)) < 0)
		{
			LOG_ERROR << "setsockopt SO_SNDBUF failed: " << strerror(errno);
		}
		else
		{
			LOG_TRACE << "Socket fd=" << fd_ << " SO_SNDBUF=" << size;
		}
	}

	void shutdownWrite()
	{
		if (shutdown(fd_, SHUT_WR) < 0)
		{
			LOG_ERROR << "Socket shutdown write failed: " << strerror(errno);
		}
		else
		{
			LOG_DEBUG << "Socket fd=" << fd_ << " shutdown write";
		}
	}

	void shutdownRead()
	{
		if (shutdown(fd_, SHUT_RD) < 0)
		{
			LOG_ERROR << "Socket shutdown read failed: " << strerror(errno);
		}
		else
		{
			LOG_DEBUG << "Socket fd=" << fd_ << " shutdown read";
		}
	}

	int getSocketError() const
	{
		int optval;
		socklen_t optlen = sizeof(optval);
		if (getsockopt(fd_, SOL_SOCKET, SO_ERROR, &optval, &optlen) < 0)
		{
			return errno;
		}
		return optval;
	}

	int fd() const { return fd_; }

	InetAddress getLocalAddr() const
	{
		sockaddr_in6 addr;
		socklen_t addrlen = sizeof(addr);
		if (getsockname(fd_, reinterpret_cast<sockaddr *>(&addr), &addrlen) < 0)
		{
			LOG_ERROR << "getsockname failed: " << strerror(errno);
			// Return a default-constructed address on failure
			return InetAddress();
		}

		if (addr.sin6_family == AF_INET)
		{
			return InetAddress(*reinterpret_cast<sockaddr_in *>(&addr));
		}
		return InetAddress(addr);
	}

	InetAddress getPeerAddr() const
	{
		sockaddr_in6 addr;
		socklen_t addrlen = sizeof(addr);
		if (getpeername(fd_, reinterpret_cast<sockaddr *>(&addr), &addrlen) < 0)
		{
			LOG_ERROR << "getpeername failed: " << strerror(errno);
			// Return a default-constructed address on failure
			return InetAddress();
		}

		if (addr.sin6_family == AF_INET)
		{
			return InetAddress(*reinterpret_cast<sockaddr_in *>(&addr));
		}
		return InetAddress(addr);
	}

	bool isSelfConnected() const
	{
		// This is only meaningful for a connected TCP socket
		return getLocalAddr() == getPeerAddr();
	}
};

} // namespace reactor