/* doscan - Denial Of Service Capable Auditing of Networks       -*- C++ -*-
 * Copyright (C) 2003, 2005 Florian Weimer
 *
 * This program 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 2 of the License, or
 * (at your option) any later version.
 *
 * This program 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 this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

#include "config.h"
#include "opt.h"
#include "scan_tcp.h"
#include "results.h"

#include <cerrno>
#include <cstdio>
#include <fcntl.h>
#include <netinet/in.h>
#include <string>
#include <sys/fcntl.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>

#ifdef HAVE_SYS_FILIO_H
#include <sys/filio.h>
#endif

// tcp_handler

tcp_client_handler::tcp_client_handler(event_queue& q,
                                       ipv4_t host, unsigned short port)
  : fd_handler(q, make_connection(host, port), watch_write),
    the_host(host), the_port(port)
{
}

tcp_client_handler::~tcp_client_handler()
{
  if (fd() >= 0) {
    int d = fd();
    unwatch();
    close(d);
  }
}

int
tcp_client_handler::make_connection(ipv4_t host, unsigned short port)
{
  int sockfd = socket(AF_INET, SOCK_STREAM, 0);

  if (sockfd == -1) {
    int err = errno;
    ipv4_string_t a;

    // If we encounter an error at this point, it is not actually
    // network-related, so we bail out immediately.

    ipv4_address_to_string (host, a);
    fprintf (stderr, "%s: could not create socket for %s, error was: %s\n",
             opt_program, a, strerror (err));
    fprintf (stderr,
             "%s: try the '--connections' option with a smaller value\n",
             opt_program);
    exit (EXIT_FAILURE);
  }

  // Make socket non-blocking.

  int flags = fcntl (sockfd, F_GETFL, 0);
  if (fcntl (sockfd, F_SETFL, flags | O_NONBLOCK) == -1) {
    int err = errno;
    ipv4_string_t a;

    // Again, this error is not network-related.

    ipv4_address_to_string (host, a);
    fprintf (stderr,
             "%s: could not set non-blocking mode for %s, error was: %s\n",
             opt_program, a, strerror (err));
    exit (EXIT_FAILURE);
  }

  struct sockaddr_in sa;
  memset (&sa, 0, sizeof (sa));
  sa.sin_family = AF_INET;
  sa.sin_port = htons (port);
  sa.sin_addr.s_addr = htonl (host);

  if (connect (sockfd, (struct sockaddr *)&sa, sizeof (sa)) == -1) {
    int err = errno;

    if (err != EINPROGRESS) {
      close (sockfd);
      if (opt_net_errors) {
        results_add (ticks_get_cached (), host, err, 0, 0);
      }
      return -1;
    }
  }

  return sockfd;
}

int
tcp_client_handler::get_error(int fd)
{
    int error = 0;
    socklen_t len = sizeof(error);

    errno = 0;
    if (getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len) >= 0) {
      // Non-Solaris case: error is overwritten with the error code,
      // Nothing to do.
    } else {
      // Solaris case: errno has the error code.
      error = errno;
    }

    return error;
}

// tcp_half_duplex_handler

tcp_half_duplex_handler::tcp_half_duplex_handler(event_queue& queue,
                                                 ipv4_t host,
                                                 unsigned short port)
  : tcp_client_handler(queue, host, port),
    next_state(0), next_is_read(false), offset(0)
{
  if (fd() == -1) {
    // An error occured.  We want to collect the timeout immediately.
    set_immediate_timeout();
  }
}

bool
tcp_half_duplex_handler::on_timeout(ticks_t)
{
  //  There is no file descriptor, terminate immediately.

  if (fd() == -1) {
    return false;
  }

  if (next_is_read) {
    // Truncate string to the length which has actually been read.
    data.resize(offset);
  } else {
    data.clear();
  }

  ready(next_state, ETIMEDOUT, data);
  return false;
}

void
tcp_half_duplex_handler::request_data(int new_state, unsigned count)
{
  data.clear();
  offset = 0;
  request_more_data(new_state, count);
}

void
tcp_half_duplex_handler::request_more_data(int new_state, unsigned count)
{
  offset = data.size();
  data.resize(data.size() + count);
  next_is_read = true;
  stop = false;
  watch(watch_read);
}

void
tcp_half_duplex_handler::send(int new_state, const std::string& data_to_send)
{
  data = data_to_send;
  next_state = new_state;
  offset = 0;
  next_is_read = false;
  stop = false;
  watch(watch_write);
}

void
tcp_half_duplex_handler::reconnect(int new_state, bool first_read)
{
  next_state = new_state;
  data.clear();
  offset = 0;
  int new_fd = make_connection(host(), port());
  int old_fd = fd();
  unwatch();
  close(old_fd);
  if (new_fd != -1) {
    next_is_read = first_read;
    watch(new_fd, first_read ? watch_read : watch_write);
    stop = false;
  } else {
    set_immediate_timeout();
    stop = true;
  }
}

bool
tcp_half_duplex_handler::on_activity(activity act)
{
  // send(), reqest_data(), request_more_data() will set this variable
  // to true.
  stop = true;

  if (next_is_read) {
    return on_activity_read(act);
  } else {
    return on_activity_write(act);
  }
}

bool
tcp_half_duplex_handler::on_activity_read(activity act)
{
  switch (act) {
  case activity_read:
  case activity_read_write:
  case activity_error:
    break;
  case activity_write:
    abort();
  }

  int remaining = data.size() - offset;

  if (remaining <= 0) {
    // This means that the available data was requested.  First, we
    // have to discover the amount of data.

    int result = ioctl(fd(), FIONREAD, &remaining);
    if (result == -1) {
      int err = errno;
      // No resizing is necessary.
      ready(next_state, err, data);
      return false;
    } else {
      if (remaining == 0) {
        remaining = 1;
      }
      // Grow data as necessay.
      data.resize(data.size() + remaining);
    }
  }

  int result = ::read(fd(), &data[offset], remaining);

  if (result == -1) {
    int err = errno;
    data.resize(offset);
    ready(next_state, err, data);
    return false;

  } else if (result == 0) {
    // Empty reply, other side is closing.
    data.resize(offset);
    ready(next_state, 0, data);
    return false;

  } else {
    offset += result;
    if (offset == data.size()) {
      // Read operation has completed.

      std::string copy(data);
      next_is_read = false;
      offset = 0;
      ready(next_state, 0, copy);
      return !stop;
    } else {
      return true;
    }
  }
}

bool
tcp_half_duplex_handler::on_activity_write(activity act)
{
  switch (act) {
  case activity_write:
  case activity_read_write:
  case activity_error:
    break;
  case activity_read:
    abort();
  }

  if (data.size() == 0) {
    // There is no data, go and fetch some.
    std::string empty;
    ready(next_state, 0, empty);
    if (stop || data.size() == 0) {
      return false;
    }
  }

  unsigned remaining = data.size() - offset;
  int result = ::write(fd(), &data[offset], remaining);

  if (result == -1) {
    int err = errno;
    data.clear();
    offset = 0;
    ready(next_state, err, data);
    return false;
  }

  offset += remaining;
  if (offset == data.size()) {
    data.clear();
    next_is_read = false;
    offset = 0;
    std::string empty;
    ready(next_state, 0, empty);
    return !stop;
  } else {
    // Continue writing.

    return true;
  }
}

// arch-tag: 91207afc-6334-41e4-a6a0-3c2aac453597