udpraw.cpp

Go to the documentation of this file.

/* This file is part of Udpraw - A library for UDP transmission.
 *
 * Copyright (C) 2011 Trevor Dodds <@gmail.com trev.dodds>
 *
 * Udpraw is free software: you can redistribute it and/or modify it under the
 * terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 * 
 * Udpraw is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
 * details.
 * 
 * You should have received a copy of the GNU General Public License along with
 * Udpraw.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "udpraw.h"
#include <algorithm> // for max()
#include <ctime>
#include <iostream>
#include <string>
#ifndef _MSC_VER
  #include <sys/time.h>
  #include <stdio.h>
  #include <sys/errno.h>
  #include <arpa/inet.h>
  //#include <netdb.h> // temporary for testing socket properties
  //#include <sys/socket.h> // temporary for testing
#endif

#ifdef max
#undef max // get rid of annoying windows definitions that override std::max()
#endif
#ifdef min
#undef min
#endif

Udpraw::Udpraw()
{
  createdRecv = false;
  createdSend = false;
  FD_ZERO(&readSocks);
  FD_ZERO(&writeSocks);
  FD_ZERO(&exceptionSocks);
  recvSocket = 0; // for max() in doSelect()
  sendSocket = 0;
  //recvPort = 0;

  statSent = 0; // generate some stats
  statRecvd = 0;
  lastError = UDPRAW_NO_ERROR;
  errorPrepend = "Udpraw error: ";

  simulateLoss.send = 0.0;
  simulateLoss.recv = 0.0;
  srand(time(NULL));

#ifdef _MSC_VER
  initWSA = false;
#endif
}

Udpraw::~Udpraw()
{
  closeAndCleanup();
}

int Udpraw::getError()
  // Get and clear last error number.
{
  int temp = lastError;
  lastError = UDPRAW_NO_ERROR;
  return temp;
}

std::string Udpraw::getErrorStr()
  // Get last error message and clear last error.
{
  return getErrorStr(getError());
}

std::string Udpraw::getErrorStr(int error)
  // Get error message.
{
  std::string str;

  switch (error)
  {
    case UDPRAW_NO_ERROR:
      str = "No error.";
      break;
    case UDPRAW_BIND:
      str = "Error binding.";
      break;
    case UDPRAW_MSGSIZE:
      str = "Buffer not big enough for message, some data lost.";
      break;
    case UDPRAW_NOINIT:
      str = "Socket not initialised.";
      break;
    case UDPRAW_NOIP:
      str = "No ip address specified.";
      break;
    case UDPRAW_NOPORT:
      str = "No port specified.";
      break;
    case UDPRAW_RECV:
      str = "Error receiving (not message size).";
      break;
    case UDPRAW_SELECT:
      str = "Error with select.";
      break;
    case UDPRAW_SEND:
      str = "Error sending.";
      break;
    case UDPRAW_SOCKEXCEPTION:
      str = "Select detected a socket exception.";
      break;
    case UDPRAW_UNMATCHED:
      str = "> 1 IPs, > 1 ports, but the number of IPs does not match the number of ports.";
      break;
    default:
      str = "Error not found";
      break;
  }

  return str;
}

void Udpraw::writeError()
  // Write and clear last error.
{
  std::cerr << errorPrepend;
  writeError(getError());
}

void Udpraw::writeError(std::string prependMsg)
{
  if (prependMsg != "") std::cerr << prependMsg;
  writeError(getError());
}

void Udpraw::writeError(int error)
  // Write error message for number 'error'
{
  std::cerr << getErrorStr(error) << std::endl;
}

void Udpraw::init()
{
#ifdef _MSC_VER
  // Initialize Winsock
  WSAStartup(MAKEWORD(2,2), &wsaData); // version 2.2
  initWSA = true;
#endif
}

int Udpraw::initReceiver(int port)
// returns 0 success, -1 error binding
{
#ifdef _MSC_VER
  if (!initWSA) init();
#endif

  //recvPort = port;

  // Create a receiver socket to receive datagrams
  recvSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
  createdRecv = true;

  // Bind the socket to address (or INADDR_ANY) and the specified recvPort.
  recvAddr.sin_family = AF_INET;
  recvAddr.sin_port = htons(port);
  recvAddr.sin_addr.s_addr = htonl(INADDR_ANY);

#ifdef _MSC_VER
  if (bind(recvSocket, (SOCKADDR *) &recvAddr, sizeof(recvAddr)) != 0) {
#else
  if (bind(recvSocket, (struct sockaddr *) &recvAddr, sizeof(recvAddr)) != 0) {
#endif
    //std::cerr << "Error binding" << std::endl;
    lastError = UDPRAW_BIND;
    return -1;
  }

  return 0;

  //
  //
  // check socket properties (temporary for testing)
  //
  // send and recv buffers were 65535 at socket level (cygwin).

  //protoent *pent = getprotobyname("udp");
  //int protnum = pent->p_proto;
  //cout << "Protocol number: " << protnum << std::endl;
  //long int optval = 0;
  //int optvallen = sizeof(optval);
  //cout << "optvallen: " << optvallen << std::endl;
  //int getsockoptret = getsockopt(recvSocket, SOL_SOCKET, SO_SNDBUF, &optval, &optvallen);
  //if (getsockoptret < 0) {
  //  perror("perror getsockopt");
  //}
  //if (optvallen > 0) cout << "RCVBUF2: " << optval << std::endl;
  //cout << "getsockoptret: " << getsockoptret << std::endl;

}

void Udpraw::initSender(int port)
{
#ifdef _MSC_VER
  if (!initWSA) init();
#endif

  //sendPort = port;

  // Create a socket for sending data
  sendSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

  sendAddr.sin_family = AF_INET;
  //sendAddr.sin_port = htons(port);

  if (port > 0) setSendPort(port);

  createdSend = true;
}

void Udpraw::initSender(const char *ip, int port)
{
  setSendIp(ip);
  initSender(port);
}

void Udpraw::initSender(std::vector <std::string> ip_list, int port)
{
  setSendIp(ip_list);
  initSender(port);
}

void Udpraw::initSender(std::vector <std::string> ip_list, std::vector <int> port_list)
{
  setSendIp(ip_list);
  initSender(0);
  setSendPort(port_list);
}

void Udpraw::setSendIp(const char* ip)
{
  sendIp.clear();
  sendIp.push_back(ip);
}

void Udpraw::setSendIp(std::vector <std::string> ip_list)
{
  sendIp = ip_list;
}

void Udpraw::setSendPort(int port)
{
  sendPort.clear();
  sendPort.push_back(port);
}

void Udpraw::setSendPort(std::vector <int> port_list)
{
  sendPort = port_list;
}

int Udpraw::doSelect()
  // Select, check for errors
  // Returns 0 on success, -1 error, -2 if a socket exception was detected.
  // NOTE: doSelect() is called by the client, so it can be done once for send and recv
{
  struct timeval timeout;             // timeout for using select()
  timeout.tv_sec = 0; // don't halt
  timeout.tv_usec = 0;

  FD_ZERO(&readSocks);
  FD_ZERO(&writeSocks);
  FD_ZERO(&exceptionSocks);

  if (createdRecv) {
    FD_SET(recvSocket, &readSocks); // check for sockfd being readable
    FD_SET(recvSocket, &exceptionSocks);
  }
  if (createdSend) {
    FD_SET(sendSocket, &writeSocks); // check for sockfd being writable
    FD_SET(sendSocket, &exceptionSocks);
  }

  int sockfdMax = std::max( (int) recvSocket, (int) sendSocket ) + 1;

  // check if socket is readable/writable
  if ((numSocksReadable = select(sockfdMax, &readSocks, &writeSocks, &exceptionSocks, &timeout)) == -1) {
    lastError = UDPRAW_SELECT;
    return -1;
  }

  if (FD_ISSET(recvSocket, &exceptionSocks) || FD_ISSET(sendSocket, &exceptionSocks)) {
    lastError = UDPRAW_SOCKEXCEPTION;
    return -2;
  }

  return 0;
}

int Udpraw::recvRaw(char* buf, int len, bool select)
  // return number of bytes received and output pointer to data
  // 0 if socket will block
  // -1 if data didn't fit, -2 for any other error
{
  // error check
  if (!createdRecv) {
    lastError = UDPRAW_NOINIT;
    return -2;
  }

  if (!select || FD_ISSET(recvSocket, &readSocks)) {
#ifdef _MSC_VER
    int fromAddrSize = sizeof(fromAddr); // this should get filled in with the size of the fromAddr, but recvfrom seems to like it to already be defined
#else
    socklen_t fromAddrSize = sizeof(fromAddr);
#endif

    // here, data is extracted from the first enqueued message.
    // note if the datagram is larger than recvBufLen, the excess will be discarded!
    // see http://msdn.microsoft.com/en-us/library/ms740120%28VS.85%29.aspx
    //cout << "Receiving from port: " << recvAddr.sin_port << ", length: " << len << std::endl;

    // TODO may need to use MSG_WAITALL, see man recvfrom, to make sure
    // we get all the data we're looking for
    // TODO see also man 2 getsockopt, and find out the minimum send/recv
    // buffer sizes
#ifdef _MSC_VER
    int bytesRecvd = recvfrom(recvSocket, buf, len, 0, (SOCKADDR *)&fromAddr, &fromAddrSize);
#else
    int bytesRecvd = recvfrom(recvSocket, buf, len, 0, (struct sockaddr *)&fromAddr, &fromAddrSize);
#endif

    if (bytesRecvd > 0) statRecvd += bytesRecvd;

#ifdef _MSC_VER
    if (bytesRecvd == 0 || bytesRecvd == SOCKET_ERROR)
#else
    if (bytesRecvd < 1)
#endif
    {
      //cout << "Error occurred with recvfrom, possibly buffer to small, or socket closed" << std::endl;
#ifdef _MSC_VER
      int error = WSAGetLastError();
      if (error == WSAEMSGSIZE) {
        //cout << "Detected error: the message size was longer than the receive buffer." << std::endl;
        //cout << "Set last byte to null before attempting to output!" << std::endl;
        lastError = UDPRAW_MSGSIZE;
        return -1;
      }else{
        lastError = UDPRAW_RECV;
        return -2;
      }
#else
      // use cmsg:

      // TODO receive 'truncated' message here:
      /*struct msghdr msgh;
      struct cmsghdr *cmsg;
      int *ttlptr;
      int received_ttl;

      // Receive auxiliary data in msgh
      for (cmsg = CMSG_FIRSTHDR(&msgh); cmsg != NULL;
           cmsg = CMSG_NXTHDR(&msgh,cmsg)) {

        if (cmsg->cmsg_level == IPPROTO_IP && cmsg->cmsg_type == IP_TTL) {
                   ttlptr = (int *) CMSG_DATA(cmsg);
                   received_ttl = *ttlptr;
                   break;
        }
      }
      if (cmsg == NULL) {
               //
               // Error: IP_TTL not enabled or small buffer
               // or I/O error.
               //
      }*/

      //perror("perror recvfrom");
      //cout << "Error number: " << errno << std::endl;
      lastError = UDPRAW_RECV;
      return -2;
#endif
    }

    // simulate packet loss
    if (simulateLoss.recv > 0.0 && simulatePacketLoss(simulateLoss.recv)) return 0;

    return bytesRecvd;
  }else{
    return 0;
  }
}

int Udpraw::sendRaw(const char* buf, int len, bool select)
// returns the minimum bytes sent to each destination, 0 if dropped (simulated)
// or -1 on error or when socket would block
{
  int minbytes = -1, bytes = 0;

  std::vector<std::string>::iterator ip_iterator = sendIp.begin();
  std::vector<int>::iterator port_iterator = sendPort.begin();

  if (!createdSend) {
    lastError = UDPRAW_NOINIT;
    return -1;
  }

  // error check, either 1 ip and > 0 ports, 1 port and > 0 ips,
  // or > 1 ip and equal number of ports
  if (sendPort.size() == 0) {
    //std::cerr << "UDP sendRaw error: no port." << std::endl;
    lastError = UDPRAW_NOPORT;
    return -1;
  }
  if (sendIp.size() == 0) {
    //std::cerr << "UDP sendRaw error: no IP." << std::endl;
    lastError = UDPRAW_NOIP;
    return -1;
  }
  if (sendPort.size() > 1 && sendIp.size() > 1) {
    // in this case need a list of corresponding ports for each IP.
    if (sendPort.size() != sendIp.size()) {
      //std::cerr << "UDP sendRaw error: Number of send ports does not match number of IPs." << std::endl;
      lastError = UDPRAW_UNMATCHED;
      return -1;
    }
  }
  

  // simulate packet loss:
  if (simulateLoss.send > 0.0 && simulatePacketLoss(simulateLoss.send)) return 0;


  bool first = true;
  while (ip_iterator != sendIp.end() && port_iterator != sendPort.end())
  {
    if (first) first = false;
    else{
      if (select) doSelect(); // only repeat this if sending to more than one IP, otherwise let app do it
      // this is so they can do select once for send and recv
    }
    if (!select || FD_ISSET(sendSocket, &writeSocks)) {
      sendAddr.sin_port = htons(*port_iterator);
      sendAddr.sin_addr.s_addr = inet_addr(ip_iterator->c_str());
      //cout << "sending: " << buf << " (size: " << len << ") to ip: " << ip_iterator->c_str() << ", and port: " << *port_iterator << std::endl;
#ifdef _MSC_VER
      bytes = sendto(sendSocket, buf, len, 0, (SOCKADDR *) &sendAddr, sizeof(sendAddr));
      if (bytes == SOCKET_ERROR) bytes = -1; // compatible with below
#else
      bytes = sendto(sendSocket, buf, len, 0, (struct sockaddr *) &sendAddr, sizeof(sendAddr));
      // linux returns -1 on error
#endif
      if (bytes == -1) lastError = UDPRAW_SEND;
      if (bytes > 0) statSent += bytes;
      if (minbytes < 0 || bytes < minbytes) minbytes = bytes; // could be -1, for error
    }else break; // may mean packets don't go to all addresses, but we know that anyway since we're using UDP!

    // either increase one, or the other, or both, or break
    if (sendIp.size() > 1) ++ip_iterator;
    if (sendPort.size() > 1) ++port_iterator;
    if (sendIp.size() == 1 && sendPort.size() == 1) break;
  }

  return minbytes; // it's possible that sendto could send less than we asked it to
}

int Udpraw::getStatSent()
{
  return statSent;
}

int Udpraw::getStatRecvd()
{
  return statRecvd;
}

void Udpraw::setStatSent(int bytes)
{
  statSent = bytes;
}

void Udpraw::setStatRecvd(int bytes)
{
  statRecvd = bytes;
}

void Udpraw::setSimulateLossSend(float n)
{
  simulateLoss.send = n;
}

void Udpraw::setSimulateLossRecv(float n)
{
  simulateLoss.recv = n;
}

float Udpraw::getSimulateLossSend()
{
  return simulateLoss.send;
}

float Udpraw::getSimulateLossRecv()
{
  return simulateLoss.recv;
}

bool Udpraw::simulatePacketLoss(float prob)
// takes a probability from 0-1, 0 = never lose, 1 = always lose.
{
  if (rand() % 1000 < prob * 1000) return true; // loss
  return false; // no loss
}

int Udpraw::closeAndCleanup()
{
  int retval = 0;
#ifdef _MSC_VER
  if (createdRecv) { createdRecv = false; if (closesocket(recvSocket) != 0) retval = 1; }
  if (createdSend) { createdSend = false; if (closesocket(sendSocket) != 0) retval = 1; }
  if (initWSA) { initWSA = false; WSACleanup(); }
#else
  if (createdRecv) { createdRecv = false; if (close(recvSocket) != 0) retval = 1; }
  if (createdSend) { createdSend = false; if (close(sendSocket) != 0) retval = 1; }
#endif

  return retval;
}