IO Multiplexing API with POLL() common Unix API

poll-server.cpp

// File: socket-echo-server.cpp 
// Desc: Simple socket server with a single thread. 

#include <iostream>
#include <string> 
#include <vector>
#include <cassert> 
#include <cstring>  // memcpy 

// ----- Unix/Linux headers -----// 
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>    
#include <sys/socket.h>
#include <netdb.h>
#include <limits.h>
#include <sys/ioctl.h>

#include <poll.h>

// -------------  MAIN() ------------------------------//

struct sresult{
    std::string msg;
    ssize_t     size;
};

sresult read_string(int fd, size_t buffer_size);
void    write_buffer(int fd, const void* buffer, size_t size);
void    write_buffer(int fd, const void* buffer, size_t size);
void    write_string(int fd, std::string const& text);

constexpr int UNIX_FAILURE = -1;
constexpr int BLOCK_INFINITE_TIMEOUT = -1;


int main(int argc, char** argv)
{
    std::fprintf(stderr, " [TRACE] Process Unique Identifier  PID = %d \n", ::getpid());

    // ======= Server configuration =========// 
    std::uint16_t port     = 9056;
     const char*  hostname = "0.0.0.0";

    // ====== Create Socket File Descriptor =============// 

     int sockfd = socket (  AF_INET       // domain:   IPv4
                          , SOCK_STREAM   // type:     TCP 
                          , 0             // protocol: (value 0) 
                         );

    std::fprintf(stderr, " [TRACE] sockfd = %d \n", sockfd);

    // Placeholder for future error handling. 
    // Returns (-1) on failure 
    assert( sockfd != - 1);

    // Note: t 'struct' is redundant. 
     struct hostent* h = gethostbyname(hostname);
     assert(h != nullptr);
     struct sockaddr_in sa; 
     memcpy(&sa.sin_addr.s_addr, h->h_addr, h->h_length);
     sa.sin_family = AF_INET;
     sa.sin_port   = htons(port); 

     // ----- Bind to Port and wait for client connections ---------//
     // Enables binding to the same address 
     int enable_reuseaddr = 1;
     if ( setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &enable_reuseaddr, sizeof(int)) < 0 )
     {
         fprintf(stderr, "Error: unable to set socket option. \n");
         return EXIT_FAILURE;
     }

     if( ::bind(sockfd, (sockaddr *) &sa, sizeof(sa)) == -1)
     {
         fprintf(stderr, "Error: unable to bind socket \n");
         return EXIT_FAILURE;
     }

    int backlog = 5;
    assert( listen(sockfd, backlog) != UNIX_FAILURE && " Non-recoverable failure. Check errno. " );

    // ================ Server Main Loop ===============//

    // This heap-allocated array contains all file 
    // descriptors that will be monitored. 
    auto pollv = std::vector<struct pollfd>{};    

    struct pollfd pfd;
    
    // Register STDIN file descriptor 
    pfd.fd      = STDIN_FILENO;
    pfd.events  = POLLIN;
    pfd.revents = 0;
    pollv.push_back(pfd);

    // Register socket server file descriptor 
    pfd.fd      = sockfd;
    pfd.events  = POLLIN;
    pfd.revents = 0;
    pollv.push_back(pfd);
    
    std::cout << " [TRACE] pollv.size() = " << pollv.size() << '\n';

    // ======= Server Event Loop / IO Multiplexing Loop =======//
    //
    for(;;)
    {

        std::fprintf(stderr, " [TRACE] Waiting IO events. \n");

        // poll() Blocks the current thread until any of the file descriptors become ready 
        // for IO (Input/Output) operations. On failure, the function returns -1 setting ERRNO.
        // On success, the function poll() returns the number of file descriptors (nfds)
        // ready to be written or read. 
        // 
        int nfds = ::poll( &pollv[0], pollv.size(), BLOCK_INFINITE_TIMEOUT);
        std::fprintf(stderr, " [TRACE] File descriptors ready. OK => nfds = %d \n", nfds);        
        
        if(nfds == UNIX_FAILURE)
        {
            std::fprintf(stderr, " [ERROR] errno = %d / %s \n", errno, strerror(errno) );
            exit(EXIT_FAILURE);
        }

        assert(nfds > 0 && "Nfds supposed to be greater than zero.");

        // Iterate over ready file descriptors 
        for(auto& pf : pollv)
        {   
            // Skip entries for which there is no IO event.          
            if(pf.revents == 0){ continue; }

            // Event that happens when user types something in the console.
            if(pf.fd == STDIN_FILENO)
            {
                std::string line;
                std::getline(std::cin, line);
                std::fprintf(stdout, " [EVENT] User entered: %s \n", line.c_str());

                // Send STDIN message to all connected socket clients.
                for(int i = 2; i < pollv.size(); i++){
                    int fd = pollv[i].fd;
                    if(fd == -1){ continue; }
                    write_string(fd, " [STDIN] " + line + "\n");
                }
                
                // Skip the parts below this statement. 
                continue;   
            }

            // Event happens when a socket client attempts to connect. 
            // // if(pf.fd == sockfd && ((pf.revents & POLLIN) == POLLIN)  )
            if( pf.fd == sockfd )
            {   
                sockaddr_in client_sa;     
                socklen_t addrlen = sizeof(sockaddr_in);
                int client_fd = accept(sockfd, nullptr, nullptr);

                //if(client_fd == UNIX_FAILURE){ break; }
                assert( client_fd != UNIX_FAILURE && " Non recoverable failure. \n" );

                std::fprintf(stderr, " [TRACE] Received client connection => client_fd = %d \n", client_fd);     

                // Register client socket file descriptor 
                struct pollfd pfd;
                pfd.fd      = client_fd;
                pfd.events  = POLLIN | POLLHUP;
                pfd.revents = 0;
                pollv.push_back(pfd);            

                // Greet client socket
                write_string(client_fd, " [SERVER] Hi, client socket. Welcome to the server. \n\n");

                continue;
            }

            // Event that happens when any client socket sends a message.
            if(pf.revents & POLLIN)
            {
                auto [msg, size] = read_string(pf.fd, 500);            
                assert( size >= 0 && "  Number of bytes returned supposed to be grater or equal to zero " );

                // Socket returns 0 bytes when disconnected 
                if(size == 0){
                    std::fprintf(stdout, " [EVENT] Client socket fd = %d closed connection \n", pf.fd);
                    // Disable file descriptor 
                    close(pf.fd);
                    pf.fd = -1;
                    continue;
                }
                std::fprintf(stdout, " [EVENT] fd = %d ; message = %s ", pf.fd, msg.c_str());
                
                write_string(pf.fd, "\n [ECHO] fd = " + std::to_string(pf.fd) + " / " + msg + "\n");
                continue;
            }

        } // ---- End of for() pf iterations --------//

    } // -------- End of server pf loop -----------------------//

    close(sockfd);
    return 0;
}

    // ------------------------------------------------------//
    //          I M P L E M E N T A T I O N S                //
    //-------------------------------------------------------//

sresult
read_string(int fd, size_t buffer_size)
{
    // Create a string initialized with null characters 
    std::string str(buffer_size, 0x00);
    ssize_t n  = read(fd, &str[0], buffer_size);
    if(n > 0){ str.resize(n); }
    return sresult{ str, n };
}


void write_buffer(int fd, const void* buffer, size_t size)
{
   // Pre-condition 
   assert( size <= SSIZE_MAX  );
   size_t  len  = size;
   ssize_t nr   = -1;
   auto pbuf = reinterpret_cast<const std::uint8_t*>(buffer); 

  while(  len != 0 && ( nr = write(fd, pbuf, len) ) != 0 ) 
  {
      if(nr == -1 && errno == EINTR){ continue; } 
      if(nr == -1) { 
          throw std::runtime_error("An error has happened, check ERRNO ");
          break;
      }
      len  = len - nr; 
      pbuf = pbuf + nr;
   }
}

void write_string(int fd, std::string const& text)
{
    write_buffer(fd, text.data(), text.length());
}