boost.asio + sendfile(2)

asio_sendfile.cpp

// how to use boost.asio + sendfile(2)
// Michał Papierski <[email protected]>
#include <iostream>
#include <string>
#include <boost/asio.hpp>
#include <boost/make_shared.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/ref.hpp>
#include <boost/bind.hpp>
#include <boost/enable_shared_from_this.hpp>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/sendfile.h>
#include <fcntl.h>

// file to send
static
std::string g_filename;

struct client
	: boost::enable_shared_from_this<client>
{
	client(boost::asio::io_service & io_svc)
		: io_svc_(io_svc)
		, socket_(io_svc_)
		, fd_(open(g_filename.c_str(), O_RDONLY))
		, start_(0)
	{
		// insert better error checking here
		assert(fd_ != 0);
		struct stat statbuf;
		int result = stat(g_filename.c_str(), &statbuf);
		assert(result != -1);
		count_ = statbuf.st_size;
	}
	void start()
	{
		ssize_t result = sendfile(socket_.native(), fd_, &start_, count_ - start_);
		if (result == -1)
		{
			// something bad happens
			std::cerr << "error: " << errno << std::endl;
		}
		else if (result == 0)
		{
			// socket will disconnect now
			std::cout << "done" << std::endl;
		}
		else
		{
			// this will do nothing besides checking if socket
			// is writable again.
			socket_.async_write_some(boost::asio::null_buffers(),
				boost::bind(&client::handle_write, shared_from_this(), _1, result));
		}
	}
	void handle_write(const boost::system::error_code & error, std::size_t bytes_transferred)
	{
		if (error)
		{
			std::cerr << "handle_write: " << error.message() << std::endl;
			return;
		}
		start();
	}
	boost::asio::io_service & io_svc_;
	boost::asio::ip::tcp::socket & socket() { return socket_; }
	boost::asio::ip::tcp::socket socket_;
	int fd_;
	off_t start_;
	size_t count_;
};

struct server
{
	server(boost::asio::io_service & io_svc)
		: io_svc_(io_svc)
		, acceptor_(io_svc_, boost::asio::ip::tcp::endpoint(boost::asio::ip::tcp::v4(), 9999))
	{
		start_accept();
	}
	void start_accept()
	{
		boost::shared_ptr<client> new_conn = boost::make_shared<client>(boost::ref(io_svc_));
		acceptor_.async_accept(new_conn->socket(),
			boost::bind(&server::handle_accept, this, new_conn, boost::asio::placeholders::error()));
	}
	void handle_accept(boost::shared_ptr<client> new_conn, const boost::system::error_code & error)
	{
		if (error)
		{
			std::cerr << "accept error: " << error.message() << std::endl;
			return;
		}
		// send O_NONBLOCK here. sendfile(2) will block otherwise.
		boost::asio::socket_base::non_blocking_io option(true);
		new_conn->socket().io_control(option);
		// start writing
		new_conn->start();
		start_accept();
	}
	boost::asio::io_service & io_svc_;
	boost::asio::ip::tcp::acceptor acceptor_;
};

int
main(int argc, char * argv[])
{
	if (argc < 2)
	{
		std::cerr << "usage: " << argv[0] << " [filename]" << std::endl;
		return 1; 
	}
	g_filename = argv[1];
	boost::asio::io_service io_svc;
	server srv(io_svc);
	io_svc.run();
}