OpenSSL non-blocking read/write example with select

ssl_nb_read_write.cpp

// Reformatted the code available at https://github.com/Andersbakken/openssl-examples/blob/master/read_write.c
// The code is hopefully easier to understand now

#include <fcntl.h>
#include <openssl/ssl.h>
#include <unistd.h>

void err_exit(const char *s)
{
    /* TODO: Implement an exit function */
    (void) s;
}

void bio_err_exit(const char *s)
{
    /* TODO: Implement an exit function by printing the error strings to BIO */
    (void) s;
}

void set_non_blocking(int sock)
{
    /*First we make the socket non-blocking*/
    int mode = fcntl(sock, F_GETFL, 0);
    if (mode < 0)
        err_exit("Couldn't get socket flags");
    mode |= O_NDELAY;
    if (fcntl(sock, F_SETFL, mode) < 0)
        err_exit("Couldn't make socket non-blocking");
}

const size_t BUF_SIZE = 1024;

void read_write(SSL *ssl)
{

    bool read_blocked_on_write = false;
    bool write_blocked_on_read = false;

    // Boolean set if we initiate a SSL shutdown and then shutdown
    // did not complete in one-shot. If set to true, then we're
    // waiting for the peer to acknowledge and send alert_notify
    // of its own.
    bool shutdown_wait = false;

    // Buffer to hold data to be sent to SSL server
    char c2s[BUF_SIZE];
    size_t c2s_len = 0, c2s_offset = 0;

    int sock = SSL_get_fd(ssl);
    set_non_blocking(sock);

    while (true)
    {
        fd_set read_fds, write_fds;
        FD_ZERO(&read_fds);
        FD_ZERO(&write_fds);

        FD_SET(sock, &read_fds);

        /* If we're waiting for a read on the socket don't try to write to the server */
        if (!write_blocked_on_read)
        {
            /* If we have data in the write queue don't try to read from standard input */
            if (c2s_len || read_blocked_on_write)
            {
                FD_SET(sock, &write_fds);
            }
            else
            {
                FD_SET(fileno(stdin), &read_fds);
            }
        }

        int ret_val = select(sock + 1, &read_fds, &write_fds, nullptr, nullptr);
        if (ret_val < 0)
            err_exit("select system call failed");
        if (ret_val == 0)
            continue;

        /* Check for input on the console*/
        if (FD_ISSET(fileno(stdin), &read_fds))
        {
            int ret_code = read(fileno(stdin), c2s, BUF_SIZE);
            if (ret_code < 0)
            {
                err_exit("Read from standard input failed\n");
            }
            else if (ret_code == 0)
            {
                shutdown_wait = true;
                ret_code = SSL_shutdown(ssl);
                if (ret_code < 0)
                {
                    // Failed to send shutdown on the wire
                    bio_err_exit("SSL shutdown failed");
                }
                else if (ret_code == 1)
                {
                    // Shutdown successful in one-shot
                    goto end;
                }
                else
                {
                    // Wait for an acknowledgement from the peer
                    shutdown_wait = true;
                }
            }
            c2s_len = ret_code;
            c2s_offset = 0;
        }

        /* Now check if there's data to read */
        if ((!write_blocked_on_read && FD_ISSET(sock, &read_fds)) || (read_blocked_on_write && FD_ISSET(sock, &write_fds)))
        {
            bool read_blocked;
            do
            {
                read_blocked = read_blocked_on_write = false;

                char s2c[BUF_SIZE];
                int ret_code = SSL_read(ssl, s2c, BUF_SIZE);
                switch (SSL_get_error(ssl, ret_code))
                {
                    case SSL_ERROR_NONE:
                        /* Note: this call could block, which blocks the
                           entire application. It's arguable this is the
                           right behavior since this is essentially a terminal
                           client. However, in some other applications you
                           would have to prevent this condition */
                        fwrite(s2c, 1, ret_code, stdout);
                        break;
                    case SSL_ERROR_ZERO_RETURN:
                        /* End of data */
                        if (!shutdown_wait)
                        {
                            SSL_shutdown(ssl);
                        }
                        goto end;
                    case SSL_ERROR_WANT_READ:
                        read_blocked = true;
                        break;

                        /* We get a WANT_WRITE if we're
                           trying to re-handshake and we block on
                           a write during that re-handshake.
                           We need to wait on the socket to be
                           writeable but re-initiate the read
                           when it is */
                    case SSL_ERROR_WANT_WRITE:
                        read_blocked_on_write = true;
                        break;
                    default:
                        bio_err_exit("SSL read problem");
                }

                /* We need a check for read_blocked here because
                   SSL_pending() doesn't work properly during the
                   handshake. This check prevents a busy-wait
                   loop around SSL_read() */
            } while (SSL_pending(ssl) && !read_blocked);
        }

        /* If the socket is writeable... */
        if ((c2s_len && FD_ISSET(sock, &write_fds)) || (write_blocked_on_read && FD_ISSET(sock, &read_fds)))
        {
            write_blocked_on_read = false;

            /* Try to write */
            int ret_code = SSL_write(ssl, c2s + c2s_offset, c2s_len);
            switch (SSL_get_error(ssl, ret_code))
            {
                /* We wrote something, always the full bytes for now */
                case SSL_ERROR_NONE:
                    c2s_len -= ret_code;
                    c2s_offset += ret_code;
                    break;

                    /* We would have blocked */
                case SSL_ERROR_WANT_WRITE:
                    break;

                    /* We get a WANT_READ if we're
                       trying to re-handshake and we block on
                       write during the current connection.

                       We need to wait on the socket to be readable
                       but re-initiate our write when it is */
                case SSL_ERROR_WANT_READ:
                    write_blocked_on_read = true;
                    break;

                    /* Some other error */
                default:
                    bio_err_exit("SSL write problem");
            }
        }
    }

end:
    SSL_free(ssl);
    close(sock);
}