imagetransfer.cpp

/*******************************************************************************
 * Copyright (c) 2018 Nerian Vision GmbH
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *******************************************************************************/

#include <cstdio>
#include <iostream>
#include <cstring>
#include <memory>
#include <fcntl.h>
#include <string>
#include <vector>
#include <mutex>
#include "visiontransfer/imagetransfer.h"
#include "visiontransfer/exceptions.h"
#include "visiontransfer/datablockprotocol.h"
#include "visiontransfer/networking.h"

using namespace std;

/*************** Pimpl class containing all private members ***********/

class ImageTransfer::Pimpl {
public:
    Pimpl(const char* address, const char* service, ImageProtocol::ProtocolType protType,
        bool server, int bufferSize, int maxUdpPacketSize);
    ~Pimpl();

    // Redeclaration of public members
    void setRawTransferData(const ImagePair& metaData, unsigned char* rawData,
        int secondTileWidth = 0, int validBytes = 0x7FFFFFFF);
    void setRawValidBytes(int validBytes);
    void setTransferImagePair(const ImagePair& imagePair);
    TransferStatus transferData();
    bool receiveImagePair(ImagePair& imagePair);
    bool receivePartialImagePair(ImagePair& imagePair, int& validRows, bool& complete);
    int getNumDroppedFrames() const;
    bool isConnected() const;
    void disconnect();
    std::string getRemoteAddress() const;
    bool tryAccept();

private:
    // Configuration parameters
    ImageProtocol::ProtocolType protType;
    bool isServer;
    int bufferSize;
    int maxUdpPacketSize;

    // Thread synchronization
    std::recursive_mutex receiveMutex;
    std::recursive_mutex sendMutex;

    // Transfer related members
    SOCKET clientSocket;
    SOCKET tcpServerSocket;
    sockaddr_in remoteAddress;

    // Object for encoding and decoding the network protocol
    std::unique_ptr<ImageProtocol> protocol;

    // Outstanding network message that still has to be transferred
    int currentMsgLen;
    int currentMsgOffset;
    const unsigned char* currentMsg;

    // Socket configuration
    void setSocketOptions();
    void setSocketBlocking(bool blocking, SOCKET socket);

    // Network socket initialization
    void initTcpServer(const addrinfo* addressInfo);
    void initTcpClient(const addrinfo* addressInfo);
    void initUdp(const addrinfo* addressInfo);

    // Data reception
    bool receiveNetworkData(bool block);

    // Data transmission
    bool sendNetworkMessage(const unsigned char* msg, int length);
    void sendPendingControlMessages();

    bool selectSocket(bool read, bool wait);
    void closeSocket(SOCKET& socket);
};

/******************** Stubs for all public members ********************/

ImageTransfer::ImageTransfer(const char* address, const char* service,
        ImageProtocol::ProtocolType protType, bool server, int bufferSize, int maxUdpPacketSize):
        pimpl(new Pimpl(address, service, protType, server, bufferSize, maxUdpPacketSize)) {
    // All initialization in the pimpl class
}

ImageTransfer::ImageTransfer(const DeviceInfo& device, int bufferSize, int maxUdpPacketSize):
        pimpl(new Pimpl(device.getIpAddress().c_str(), "7681", static_cast<ImageProtocol::ProtocolType>(device.getNetworkProtocol()),
        false, bufferSize, maxUdpPacketSize)) {
    // All initialization in the pimpl class
}

ImageTransfer::~ImageTransfer() {
    delete pimpl;
}

void ImageTransfer::setRawTransferData(const ImagePair& metaData, unsigned char* rawData,
        int secondTileWidth, int validBytes) {
    pimpl->setRawTransferData(metaData, rawData, secondTileWidth, validBytes);
}

void ImageTransfer::setRawValidBytes(int validBytes) {
    pimpl->setRawValidBytes(validBytes);
}

void ImageTransfer::setTransferImagePair(const ImagePair& imagePair) {
    pimpl->setTransferImagePair(imagePair);
}

ImageTransfer::TransferStatus ImageTransfer::transferData() {
    return pimpl->transferData();
}

bool ImageTransfer::receiveImagePair(ImagePair& imagePair) {
    return pimpl->receiveImagePair(imagePair);
}

bool ImageTransfer::receivePartialImagePair(ImagePair& imagePair, int& validRows, bool& complete) {
    return pimpl->receivePartialImagePair(imagePair, validRows, complete);
}

int ImageTransfer::getNumDroppedFrames() const {
    return pimpl->getNumDroppedFrames();
}

bool ImageTransfer::isConnected() const {
    return pimpl->isConnected();
}

void ImageTransfer::disconnect() {
    pimpl->disconnect();
}

std::string ImageTransfer::getRemoteAddress() const {
    return pimpl->getRemoteAddress();
}

bool ImageTransfer::tryAccept() {
    return pimpl->tryAccept();
}

/******************** Implementation in pimpl class *******************/
ImageTransfer::Pimpl::Pimpl(const char* address, const char* service,
        ImageProtocol::ProtocolType protType, bool server, int
        bufferSize, int maxUdpPacketSize)
        : protType(protType), isServer(server), bufferSize(bufferSize),
        maxUdpPacketSize(maxUdpPacketSize),
        clientSocket(INVALID_SOCKET), tcpServerSocket(INVALID_SOCKET),
        currentMsgLen(0), currentMsgOffset(0), currentMsg(nullptr) {

#ifdef _WIN32
    // In windows, we first have to initialize winsock
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        throw TransferException("WSAStartup failed!");
    }
#else
    // We don't want to be interrupted by the pipe signal
    signal(SIGPIPE, SIG_IGN);
#endif

    memset(&remoteAddress, 0, sizeof(remoteAddress));

    // If address is null we use the any address
    if(address == nullptr || string(address) == "") {
        address = "0.0.0.0";
    }

    // Resolve address
    addrinfo hints;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_INET; // Use IPv4
    hints.ai_socktype = (protType == ImageProtocol::PROTOCOL_TCP ? SOCK_STREAM : SOCK_DGRAM);
    hints.ai_flags = 0;
    hints.ai_protocol = 0;

    addrinfo* addressInfo = nullptr;
    if(service != nullptr && string(service) != "") {
        if(getaddrinfo(address, service, &hints, &addressInfo) != 0 || addressInfo == nullptr) {
            TransferException ex("Error resolving address: " + string(strerror(errno)));
            throw ex;
        }
    }

    if(addressInfo->ai_addrlen != sizeof(remoteAddress)) {
        throw TransferException("Illegal address length");
    }

    try {
        if(protType == ImageProtocol::PROTOCOL_UDP) {
            initUdp(addressInfo);
        } else if(protType == ImageProtocol::PROTOCOL_TCP && isServer) {
            initTcpServer(addressInfo);
        } else {
            initTcpClient(addressInfo);
        }
    } catch(...) {
        freeaddrinfo(addressInfo);
        throw;
    }

    if(addressInfo != nullptr) {
        freeaddrinfo(addressInfo);
    }
}

ImageTransfer::Pimpl::~Pimpl() {
    if(clientSocket != INVALID_SOCKET) {
        closeSocket(clientSocket);
    }
    if(tcpServerSocket != INVALID_SOCKET) {
        closeSocket(tcpServerSocket);
    }
}

void ImageTransfer::Pimpl::initTcpClient(const addrinfo* addressInfo) {
    protocol.reset(new ImageProtocol(isServer, ImageProtocol::PROTOCOL_TCP));

    // Connect
    clientSocket = ::socket(addressInfo->ai_family, addressInfo->ai_socktype,
        addressInfo->ai_protocol);
    if(clientSocket == INVALID_SOCKET) {
        TransferException ex("Error creating socket: " + string(strerror(errno)));
        throw ex;
    }

    // Perform connection
    if(connect(clientSocket, addressInfo->ai_addr, static_cast<int>(addressInfo->ai_addrlen)) < 0) {
        TransferException ex("Error connection to destination address: " + string(strerror(errno)));
        throw ex;
    }

    memcpy(&remoteAddress, addressInfo->ai_addr, sizeof(remoteAddress));

    // Set special socket options
    setSocketOptions();
}

void ImageTransfer::Pimpl::initTcpServer(const addrinfo* addressInfo) {
    protocol.reset(new ImageProtocol(isServer, ImageProtocol::PROTOCOL_TCP));

    // Create socket
    tcpServerSocket = ::socket(addressInfo->ai_family, addressInfo->ai_socktype,
        addressInfo->ai_protocol);
    if (tcpServerSocket == INVALID_SOCKET)  {
        TransferException ex("Error opening socket: " + string(strerror(errno)));
        throw ex;
    }

    // Enable reuse address
    int enable = 1;
    setsockopt(tcpServerSocket, SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<char*>(&enable), sizeof(int));

    // Open a server port
    if (::bind(tcpServerSocket, addressInfo->ai_addr, static_cast<int>(addressInfo->ai_addrlen)) < 0)  {
        TransferException ex("Error binding socket: " + string(strerror(errno)));
        throw ex;
    }

    clientSocket = INVALID_SOCKET;

    // Make the server socket non-blocking
    setSocketBlocking(false, tcpServerSocket);

    // Listen on port
    listen(tcpServerSocket, 1);
}

void ImageTransfer::Pimpl::initUdp(const addrinfo* addressInfo) {
    protocol.reset(new ImageProtocol(isServer, ImageProtocol::PROTOCOL_UDP, maxUdpPacketSize));
    // Create sockets
    clientSocket = socket(AF_INET, SOCK_DGRAM, 0);
    if(clientSocket == INVALID_SOCKET) {
        TransferException ex("Error creating receive socket: " + string(strerror(errno)));
        throw ex;
    }

    // Enable reuse address
    int enable = 1;
    setsockopt(clientSocket, SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<char*>(&enable), sizeof(int));

    // Bind socket to port
    if(isServer && addressInfo != nullptr) {
        if (::bind(clientSocket, addressInfo->ai_addr, static_cast<int>(addressInfo->ai_addrlen)) < 0)  {
            TransferException ex("Error binding socket: " + string(strerror(errno)));
            throw ex;
        }
    }

    if(!isServer) {
        memcpy(&remoteAddress, addressInfo->ai_addr, sizeof(remoteAddress));
    }

    // Set special socket options
    setSocketOptions();
}

bool ImageTransfer::Pimpl::tryAccept() {
    if(protType != ImageProtocol::PROTOCOL_TCP || ! isServer) {
        throw TransferException("Connections can only be accepted in tcp server mode");
    }

    unique_lock<recursive_mutex> recvLock(receiveMutex);
    unique_lock<recursive_mutex> sendLock(sendMutex);

    // Accept one connection
    socklen_t clientAddressLength = sizeof(remoteAddress);

    SOCKET newSocket = accept(tcpServerSocket,
        reinterpret_cast<sockaddr *>(&remoteAddress),
        &clientAddressLength);

    if(clientAddressLength != sizeof(remoteAddress)) {
        throw TransferException("Received network address with invalid length");
    }

    if(newSocket == INVALID_SOCKET) {
        if(errno == EWOULDBLOCK || errno == ETIMEDOUT) {
            // No connection
            return false;
        } else {
            TransferException ex("Error accepting connection: " + string(strerror(errno)));
            throw ex;
        }
    }

    // Close old socket and set new socket
    if(clientSocket != INVALID_SOCKET) {
        closeSocket(clientSocket);
    }
    clientSocket = newSocket;

    // Set special socket options
    setSocketOptions();

    // Reset connection data
    protocol->resetTransfer();
    protocol->resetReception();
    currentMsg = nullptr;

    return true;
}

std::string ImageTransfer::Pimpl::getRemoteAddress() const {
    unique_lock<recursive_mutex> lock(const_cast<recursive_mutex&>(sendMutex)); // either mutex will work

    if(remoteAddress.sin_family != AF_INET) {
        return "";
    }

    char strPort[11];
    snprintf(strPort, sizeof(strPort), ":%d", remoteAddress.sin_port);

    return string(inet_ntoa(remoteAddress.sin_addr)) + strPort;
}

void ImageTransfer::Pimpl::setSocketOptions() {
    // Set the socket buffer sizes
    if(bufferSize > 0) {
        setsockopt(clientSocket, SOL_SOCKET, SO_RCVBUF, reinterpret_cast<char*>(&bufferSize), sizeof(bufferSize));
        setsockopt(clientSocket, SOL_SOCKET, SO_SNDBUF, reinterpret_cast<char*>(&bufferSize), sizeof(bufferSize));
    }

    // Set sending and receive timeouts
#ifdef _WIN32
    unsigned int timeout = 500;
#else
    struct timeval timeout;
    timeout.tv_sec = 0;
    timeout.tv_usec = 500000;
#endif

    setsockopt(clientSocket, SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<char*>(&timeout), sizeof(timeout));
    setsockopt(clientSocket, SOL_SOCKET, SO_SNDTIMEO, reinterpret_cast<char*>(&timeout), sizeof(timeout));
    setSocketBlocking(true, clientSocket);
}

void ImageTransfer::Pimpl::setSocketBlocking(bool blocking, SOCKET socket) {
#ifdef _WIN32
    unsigned long on = (blocking ? 0 : 1);
    ioctlsocket(socket, FIONBIO, &on);
#else
    int flags = fcntl(socket, F_GETFL, 0);
    if(flags != -1) {
        if(blocking) {
            flags &= ~O_NONBLOCK;
        } else {
            flags |= O_NONBLOCK;
        }
        fcntl(socket, F_SETFL, flags);
    }
#endif
}

void ImageTransfer::Pimpl::setRawTransferData(const ImagePair& metaData,
        unsigned char* rawData, int secondTileWidth, int validBytes) {
    unique_lock<recursive_mutex> sendLock(sendMutex);
    protocol->setRawTransferData(metaData, rawData, secondTileWidth, validBytes);
    currentMsg = nullptr;
}

void ImageTransfer::Pimpl::setRawValidBytes(int validBytes) {
    unique_lock<recursive_mutex> sendLock(sendMutex);
    protocol->setRawValidBytes(validBytes);
}

void ImageTransfer::Pimpl::setTransferImagePair(const ImagePair& imagePair) {
    unique_lock<recursive_mutex> sendLock(sendMutex);
    protocol->setTransferImagePair(imagePair);
    currentMsg = nullptr;
}

ImageTransfer::TransferStatus ImageTransfer::Pimpl::transferData() {
    unique_lock<recursive_mutex> lock(sendMutex);

    // First receive data in case a control message arrives
    if(protType == ImageProtocol::PROTOCOL_UDP) {
        receiveNetworkData(false);
    }

    if(remoteAddress.sin_family != AF_INET || !protocol->isConnected()) {
        return NOT_CONNECTED;
    }

    // Get first message to transfer
    if(currentMsg == nullptr) {
        currentMsgOffset = 0;
        currentMsg = protocol->getTransferMessage(currentMsgLen);

        if(currentMsg == nullptr) {
            if(protocol->transferComplete()) {
                return ALL_TRANSFERRED;
            } else {
                return NO_VALID_DATA;
            }
        }
    }

    // Try transferring messages
    bool dataTransferred = (currentMsg != nullptr);
    while(currentMsg != nullptr) {
        int writing = (int)(currentMsgLen - currentMsgOffset);

        if(sendNetworkMessage(&currentMsg[currentMsgOffset], writing)) {
            // Get next message
            currentMsgOffset = 0;
            currentMsg = protocol->getTransferMessage(currentMsgLen);
        } else {
            return WOULD_BLOCK;
        }
    }

    if(dataTransferred && protType == ImageProtocol::PROTOCOL_TCP && protocol->transferComplete()) {
        // Force a flush for TCP by turning the nagle algorithm off and on
        int flag = 1;
        setsockopt(clientSocket, IPPROTO_TCP, TCP_NODELAY, (char *) &flag, sizeof(int));
        flag = 0;
        setsockopt(clientSocket, IPPROTO_TCP, TCP_NODELAY, (char *) &flag, sizeof(int));
    }

    // Also check for control messages at the end
    if(protType == ImageProtocol::PROTOCOL_UDP) {
        receiveNetworkData(false);
    }

    if(protocol->transferComplete()) {
        return ALL_TRANSFERRED;
    } else {
        return PARTIAL_TRANSFER;
    }
}

bool ImageTransfer::Pimpl::receiveImagePair(ImagePair& imagePair) {
    int validRows = 0;
    bool complete = false;

    std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
    while(!complete) {
        if(!receivePartialImagePair(imagePair, validRows, complete)) {
            return false;
        }

        unsigned int time = static_cast<unsigned int>(std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now() - startTime).count());
        if(time > 1000) {
            return false;
        }
    }

    return true;
}

bool ImageTransfer::Pimpl::receivePartialImagePair(ImagePair& imagePair,
        int& validRows, bool& complete) {
    unique_lock<recursive_mutex> lock(receiveMutex);

    // Try to receive further image data if needed
    bool block = true;
    while(!protocol->imagesReceived() && receiveNetworkData(block)) {
        block = false;
    }

    // Get received image
    return protocol->getPartiallyReceivedImagePair(imagePair, validRows, complete);
}

bool ImageTransfer::Pimpl::receiveNetworkData(bool block) {
    unique_lock<recursive_mutex> lock = block ?
        unique_lock<recursive_mutex>(receiveMutex) : unique_lock<recursive_mutex>(receiveMutex, std::try_to_lock);

    if(clientSocket == INVALID_SOCKET) {
        return false; // Not connected
    }

    // First send control messages if necessary
    sendPendingControlMessages();

    if(!lock.owns_lock()) {
        // Waiting for the lock would block this call
        return false;
    }

    // Test if the socket has data available
    if(!block && !selectSocket(true, false)) {
        return 0;
    }

    int maxLength = 0;
    char* buffer = reinterpret_cast<char*>(protocol->getNextReceiveBuffer(maxLength));

    // Receive data
    sockaddr_in fromAddress;
    socklen_t fromSize = sizeof(fromAddress);

    int bytesReceived = recvfrom(clientSocket, buffer, maxLength,
        0, reinterpret_cast<sockaddr*>(&fromAddress), &fromSize);

    if(bytesReceived == 0 || (protType == ImageProtocol::PROTOCOL_TCP && bytesReceived < 0 && errno == WSAECONNRESET)) {
        // Connection closed
        disconnect();
    } else if(bytesReceived < 0 && errno != EWOULDBLOCK && errno != EINTR &&
            errno != ETIMEDOUT && errno != WSA_IO_PENDING && errno != WSAECONNRESET) {
        TransferException ex("Error reading from socket: " + string(strerror(errno)));
        throw ex;
    } else if(bytesReceived > 0) {
            protocol->processReceivedMessage(bytesReceived);
        if(protocol->newClientConnected()) {
            // We have just established a new connection
            memcpy(&remoteAddress, &fromAddress, sizeof(remoteAddress));
        }
    }

    return bytesReceived > 0;
}

void ImageTransfer::Pimpl::disconnect() {
    // We just need to forget the remote address in order to
    // disconnect
    unique_lock<recursive_mutex> recvLock(receiveMutex);
    unique_lock<recursive_mutex> sendLock(sendMutex);

    if(clientSocket != INVALID_SOCKET && protType == ImageProtocol::PROTOCOL_TCP) {
        closeSocket(clientSocket);
    }
    memset(&remoteAddress, 0, sizeof(remoteAddress));
}

bool ImageTransfer::Pimpl::isConnected() const {
    unique_lock<recursive_mutex> lock(const_cast<recursive_mutex&>(sendMutex)); //either mutex will work

    return remoteAddress.sin_family == AF_INET && protocol->isConnected();
}

bool ImageTransfer::Pimpl::sendNetworkMessage(const unsigned char* msg, int length) {
    int written = 0;
    if(protType == ImageProtocol::PROTOCOL_UDP) {
        sockaddr_in destAddr;
        SOCKET destSocket;
        {
            unique_lock<recursive_mutex> lock(sendMutex);
            destAddr = remoteAddress;
            destSocket = clientSocket;
        }

        if(destAddr.sin_family != AF_INET) {
            return false; // Not connected
        }

        written = sendto(destSocket, reinterpret_cast<const char*>(msg), length, 0,
            reinterpret_cast<sockaddr*>(&destAddr), sizeof(destAddr));
    } else {
        SOCKET destSocket;
        {
            unique_lock<recursive_mutex> lock(sendMutex);
            destSocket = clientSocket;
        }
        written = send(destSocket, reinterpret_cast<const char*>(msg), length, 0);
    }

    unsigned long sendError = errno;

    if(written < 0) {
        if(sendError == EAGAIN || sendError == EWOULDBLOCK || sendError == ETIMEDOUT) {
            // The socket is not yet ready for a new transfer
            return false;
        } else if(sendError == EPIPE) {
            // The connection has been closed
            disconnect();
            return false;
        } else {
            TransferException ex("Error sending network packet: " + string(strerror(errno)));
            throw ex;
        }
    } else if(written != length) {
        if(protType == ImageProtocol::PROTOCOL_UDP) {
            // The message has been transmitted partially
            throw TransferException("Unable to transmit complete UDP message");
        } else {
            // For TCP we can transmit the remaining data later
            currentMsgOffset += written;
            return false;
        }
    } else {
        return true;
    }
}

void ImageTransfer::Pimpl::sendPendingControlMessages() {
    const unsigned char* controlMsgData = nullptr;
    int controlMsgLen = 0;

    while(true) {
        unique_lock<recursive_mutex> lock(sendMutex);
        if(remoteAddress.sin_family != AF_INET) {
            return;
        }

        controlMsgData = protocol->getNextControlMessage(controlMsgLen);

        if(controlMsgData != nullptr) {
            sendNetworkMessage(controlMsgData, controlMsgLen);
        } else {
            break;
        }
    }
}

int ImageTransfer::Pimpl::getNumDroppedFrames() const {
    return protocol->getNumDroppedFrames();
}

bool ImageTransfer::Pimpl::selectSocket(bool read, bool wait) {
    SOCKET sock;
    {
        unique_lock<recursive_mutex> lock(sendMutex); // Either mutex will do
        sock = clientSocket;
    }

    fd_set fds;
    struct timeval tv;
    FD_ZERO(&fds);
    FD_SET(sock, &fds);
    tv.tv_sec = 0;
    if(wait) {
        tv.tv_usec = 100000;
    } else {
        tv.tv_usec = 0;
    }

    if(select(sock+1, (read ? &fds : nullptr), (!read ? &fds : nullptr), nullptr, &tv) <= 0) {
        // The socket is currently not ready
        return false;
    } else {
        return true;
    }
}

void ImageTransfer::Pimpl::closeSocket(SOCKET& socket) {
    shutdown(socket, SHUT_WR);

    // Receive remaining data
    char buffer[1024];
    for(int i=0; i<3; i++) {
        int received = recv(clientSocket, buffer, sizeof(buffer), 0);
        if(received <= 0) {
            break;
        }
    }

    close(socket);
    socket = INVALID_SOCKET;
}