socketconnection.h

// Copyright (c) 2018 IIS (The Internet Foundation in Sweden)
// Written by Göran Andersson <initgoran@gmail.com>
 
#pragma once
 
#include <algorithm>
#include "socket.h"
 
#ifdef USE_GNUTLS
#include <gnutls/gnutls.h>
#endif
 
class SocketConnection : public Socket {
    friend class Engine;
public:
    SocketConnection(const std::string &label, Task *owner,
                     const std::string &hostname, uint16_t port,
                     uint16_t iptype = 0, struct addrinfo *local_addr=nullptr);
 
#ifdef USE_GNUTLS
    virtual ~SocketConnection() override {
        if (tlsInitialized()) {
            gnutls_deinit(session);
        }
    }
 
    void enableTLS() {
        use_tls = true;
    }
 
    bool is_tls() const {
        return use_tls;
    }
 
    gnutls_session_t cache_session() {
        session_initialized = false;
        return session;
    }
 
    void insert_cached_session(gnutls_session_t &old_session) {
        session_initialized = true;
        use_tls = true;
        session = old_session;
        gnutls_transport_set_ptr(session,
                                 static_cast<gnutls_transport_ptr_t>(this));
    }
#endif
 
    virtual PollState connected() {
        return PollState::CLOSE;
    }
 
    virtual void connectionFailed(const std::string &err_msg) {
        log() << "connection failed: " << err_msg;
    }
 
    virtual void closedByPeer();
 
    virtual PollState readData(char *, size_t ) {
        return PollState::CLOSE;
    }
 
    virtual PollState unexpectedData(char *buf, size_t len);
 
    virtual PollState writeData() {
        return PollState::CLOSE;
    }
 
    size_t sendData(const char *buf, size_t len);
 
#ifdef USE_WEBROOT
    size_t sendFileData(int fd, size_t len);
#endif
 
    void asyncSendData(const char *buf, size_t len);
 
    void asyncSendData(const std::string data) {
        asyncSendData(data.c_str(), data.size());
    }
 
    size_t asyncBufferSize() const {
        return to_send.size();
    }
 
    static uint64_t totBytesSent() {
        return tot_bytes_sent;
    }
 
    static uint64_t totBytesReceived() {
        return tot_bytes_received;
    }
 
    static void resetByteCounter() {
        tot_bytes_sent = 0;
        tot_bytes_received = 0;
    }
 
    const std::string &peerIp() const {
        return peer_ip;
    }
 
    uint16_t peerPort() const {
        return peer_port;
    }
 
    void dbgOn(bool b = true) {
        debugging = b;
    }
 
    bool dbgIsOn() {
        return debugging;
    }
protected:
 
    SocketConnection(const std::string &label, Task *owner, int fd,
                     const char *ip, uint16_t port);
 
    PollState tellOwner(const std::string &msg);
 
private:
    // Create a connection to the given host. Will be executed asynchronously,
    // then one of the callbacks connected / connectionFailed will be called.
    SocketConnection(Task *owner, const std::string &hostname,
                     unsigned int port);
 
    SocketConnection(const SocketConnection &);
 
    // Will return false if connection fails immediately, e.g. if DNS
    // lookup fails. Otherwise callback "connected" or "connectionFailed"
    // will be called - perhaps even before this call returns:
    bool asyncConnect();
 
    PollState doRead(int fd);
 
    // In states READ and READ_BLOCKED, this will be called to check if
    // we the connection should be checked for writability too.
    // If you override this, make sure to return true if asyncBufferSize() > 0.
    bool wantToSend() override {
        return !to_send.empty();
    }
 
    PollState doWrite() {
        if (to_send.empty())
            return this->writeData();
        if (size_t written = sendData(to_send.c_str(), to_send.size()))
            to_send.erase(0, written);
        return state();
    }
 
    std::string to_send;
    char socket_buffer[100000];
 
    std::string peer_ip;
    struct addrinfo *local_ip;
    uint16_t peer_port;
    uint16_t prefer_ip_type;
 
#ifdef USE_THREADS
    thread_local
#endif
    // Per thread byte counters.
    static uint64_t tot_bytes_sent, tot_bytes_received;
    bool debugging = false;
#ifdef USE_GNUTLS
    bool use_tls = false;
    bool session_initialized = false;
    bool tls_send_pending = false;
    void setSessionInitialized() {
        session_initialized = true;
    }
    bool tlsInitialized() {
        return session_initialized;
    }
    bool init_tls_server(gnutls_certificate_credentials_t &x509_cred,
                         gnutls_priority_t &priority_cache);
    bool init_tls_client(gnutls_certificate_credentials_t &x509_cred,
                         bool verify_cert);
    int try_tls_handshake() {
        dbg_log() << "TLS handshake socket " << socket();
        return gnutls_handshake(session);
    }
    static ssize_t tls_pull_static(gnutls_transport_ptr_t self,
                                   void *buf, size_t len) {
        return static_cast<SocketConnection *>(self)->tls_pull(buf, len);
    }
    static ssize_t tls_push_static(gnutls_transport_ptr_t self,
                                   const void *buf, size_t len) {
        return static_cast<SocketConnection *>(self)->tls_push(buf, len);
    }
    ssize_t tls_pull(void *buf, size_t len);
    ssize_t tls_push(const void *buf, size_t len);
    gnutls_session_t session;
#endif
};