<commit_before>/*
** Copyright 2013-2014 Centreon
**
** Licensed under the Apache License, Version 2.0 (the "License");
** you may not use this file except in compliance with the License.
** You may obtain a copy of the License at
**
**     http://www.apache.org/licenses/LICENSE-2.0
**
** Unless required by applicable law or agreed to in writing, software
** distributed under the License is distributed on an "AS IS" BASIS,
** WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
** See the License for the specific language governing permissions and
** limitations under the License.
**
** For more information : contact@centreon.com
*/

#include <cerrno>
#include <cstdio>
#include <cstring>
#include <fcntl.h>
#include <rrd.h>
#include <sstream>
#ifdef __linux__
#  include <sys/sendfile.h>
#endif // Linux
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include "com/centreon/broker/exceptions/msg.hh"
#include "com/centreon/broker/logging/logging.hh"
#include "com/centreon/broker/rrd/creator.hh"
#include "com/centreon/broker/rrd/exceptions/open.hh"
#include "com/centreon/broker/storage/perfdata.hh"

using namespace com::centreon::broker;
using namespace com::centreon::broker::rrd;

/**
 *  Constructor.
 *
 *  @param[in] tmpl_path  The template path.
 *  @param[in] cache_size The maximum number of cache element.
 */
creator::creator(
           std::string const& tmpl_path,
           unsigned int cache_size)
  : _cache_size(cache_size), _tmpl_path(tmpl_path) {
  logging::debug(logging::medium)
    << "RRD: file creator will maintain at most " << _cache_size
    << " templates in '" << _tmpl_path << "'";
}

/**
 *  Copy constructor.
 *
 *  @param[in] right The object to copy.
 */
creator::creator(creator const& right) {
  operator=(right);
}

/**
 *  Destructor.
 */
creator::~creator() {
  clear();
}

/**
 *  Copy operator.
 *
 *  @param[in] right The object to copy.
 *
 *  @return This object.
 */
creator& creator::operator=(creator const& right) {
  if (this != &right) {
    _cache_size = right._cache_size;
    _fds = right._fds;
    _tmpl_path = right._tmpl_path;
  }
  return (*this);
}

/**
 *  Clear cache and remove template file.
 */
void creator::clear() {
  for (std::map<tmpl_info, fd_info>::const_iterator
         it(_fds.begin()), end(_fds.end());
       it != end;
       ++it) {
    tmpl_info info(it->first);
    ::close(it->second.fd);
    std::ostringstream oss;
    oss << _tmpl_path << "/tmpl_" << info.length << "_" << info.step
        << "_" << info.value_type << ".rrd";
    ::remove(oss.str().c_str());
  }
  _fds.clear();
  return ;
}

/**
 *  Create a RRD file if it does not exists.
 *
 *  @param[in] filename   Path to the RRD file.
 *  @param[in] length     Number of recording in the RRD file.
 *  @param[in] from       Timestamp of the first record.
 *  @param[in] step       Specifies the base interval in seconds with
 *                        which data will be fed into the RRD.
 *  @param[in] value_type Type of the metric.
 */
void creator::create(
                std::string const& filename,
                unsigned int length,
                time_t from,
                unsigned int step,
                short value_type) {
  // Fill template informations.
  if (!step)
    step = 5 * 60 * 60; // Default to every 5 minutes.
  if (!length)
    length = 31 * 24 * 60 * 60 / step; // Default to one month long.
  tmpl_info info;
  info.length = length;
  info.step = step;
  info.value_type = value_type;

  // Find fd informations.
  std::map<tmpl_info, fd_info>::const_iterator it(_fds.find(info));
  // Is in the cache, just duplicate file.
  if (it != _fds.end())
    _duplicate(filename, it->second);
  // Not is the cache, but we have enough space in the cache.
  // Create new entry.
  else if (_fds.size() < _cache_size) {
    std::ostringstream oss;
    oss << _tmpl_path << "/tmpl_" << length << "_" << step
        << "_" << value_type << ".rrd";
    std::string tmpl_filename(oss.str());

    // Create new template.
    _open(tmpl_filename, length, from, step, value_type);

    // Get template file size.
    struct stat s;
    if (stat(tmpl_filename.c_str(), &s) < 0) {
      char const* msg(strerror(errno));
      throw (exceptions::open() << "RRD: could not create template file '"
             << tmpl_filename << "': " << msg);
    }

    // Get template file fd.
    int in_fd(open(tmpl_filename.c_str(), O_RDONLY));
    if (in_fd < 0) {
      char const* msg(strerror(errno));
      throw (exceptions::open() << "RRD: could not open template file '"
             << tmpl_filename << "': " << msg);
    }

    // Store fd informations into the cache.
    fd_info fdinfo;
    fdinfo.fd = in_fd;
    fdinfo.size = s.st_size;
    _fds[info] = fdinfo;

    _duplicate(filename, fdinfo);
  }
  // No more space in the cache, juste create rrd file.
  else
    _open(filename, length, from, step, value_type);
}

/**
 *  Duplicate file.
 *
 *  @param[in] filename The file name to save data.
 *  @param[in] in_fd    The fd informations to duplicate file.
 */
void creator::_duplicate(std::string const& filename, fd_info const& in_fd) {
  // Remove previous file.
  remove(filename.c_str());

  int out_fd(open(
               filename.c_str(),
               O_CREAT | O_TRUNC | O_WRONLY,
               S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH));
  if (out_fd < 0) {
    char const* msg(strerror(errno));
    throw (exceptions::open() << "RRD: could not create file '"
           << filename << "': " << msg);
  }

#ifdef __linux__
  // First call(s) to sendfile detect if the kernel support the syscall
  // on any FD (Linux < 2.6.33 only supports writing to socket FD).
  off_t offset(0);
  ssize_t ret;
  do {
    ret = ::sendfile(out_fd, in_fd.fd, &offset, in_fd.size);
  } while ((ret < 0) && (errno == EAGAIN));
  bool fallback((ret < 0) && ((errno == EINVAL) || (errno == ENOSYS)));
  if (!fallback) {
    if (ret < 0) {
      char const* msg(strerror(errno));
      throw (exceptions::open() << "RRD: could not create file '"
             << filename << "': " << msg);
    }
    // Good to go with the sendfile syscall.
    _sendfile(out_fd, in_fd.fd, ret, in_fd.size, filename);
  }
  else
#endif // Linux
    // We must fallback to the read/write combo.
    _read_write(out_fd, in_fd.fd, in_fd.size, filename);

  ::close(out_fd);
}

/**
 *  Open a RRD file and create it if it does not exists.
 *
 *  @param[in] filename   Path to the RRD file.
 *  @param[in] length     Number of recording in the RRD file.
 *  @param[in] from       Timestamp of the first record.
 *  @param[in] step       Time interval between each record.
 *  @param[in] value_type Type of the metric.
 */
void creator::_open(
                std::string const& filename,
                unsigned int length,
                time_t from,
                unsigned int step,
                short value_type) {
  /* Find step of RRD file if already existing. */
  /* XXX : why is it here ?
  rrd_info_t* rrdinfo(rrd_info_r(_filename));
  time_t interval_offset(0);
  for (rrd_info_t* tmp = rrdinfo; tmp; tmp = tmp->next)
    if (!strcmp(rrdinfo->key, "step"))
      if (interval < static_cast<time_t>(rrdinfo->value.u_cnt))
        interval_offset = rrdinfo->value.u_cnt / interval - 1;
  rrd_info_free(rrdinfo);
  */

  // Set parameters.
  std::ostringstream ds_oss;
  std::ostringstream rra1_oss;
  std::ostringstream rra2_oss;
  ds_oss << "DS:value:";
  switch (value_type) {
  case storage::perfdata::absolute:
    ds_oss << "ABSOLUTE";
    break ;
  case storage::perfdata::counter:
    ds_oss << "COUNTER";
    break ;
  case storage::perfdata::derive:
    ds_oss << "DERIVE";
    break ;
  default:
    ds_oss << "GAUGE";
  };
  ds_oss << ":"<< step * 10 << ":U:U";
  rra1_oss << "RRA:AVERAGE:0.5:1:" << length + 1;
  rra2_oss << "RRA:AVERAGE:0.5:12:" << length / 12 + 1;
  std::string ds(ds_oss.str());
  std::string rra1(rra1_oss.str());
  std::string rra2(rra2_oss.str());
  char const* argv[5];
  argv[0] = ds.c_str();
  argv[1] = rra1.c_str();
  argv[2] = rra2.c_str();
  argv[3] = NULL;

  // Debug message.
  logging::debug(logging::high) << "RRD: opening file '" << filename
    << "' (" << argv[0] << ", " << argv[1] << ", " << argv[2]
    << ", step " << step << ", from " << from << ")";

  // Create RRD file.
  rrd_clear_error();
  if (rrd_create_r(
        filename.c_str(),
        step,
        from,
        3,
        argv))
    throw (exceptions::open() << "RRD: could not create file '"
             << filename << "': " << rrd_get_error());

  return ;
}

/**
 *  Transfer file between two FDs using read and write.
 *
 *  @param[in] out_fd   Output FD.
 *  @param[in] in_fd    Input FD.
 *  @param[in] size     Size to transfer.
 *  @param[in] filename Path to the file being created.
 */
void creator::_read_write(
                int out_fd,
                int in_fd,
                ssize_t size,
                std::string const& filename) {
  // Reset position of in_fd.
  if (lseek(in_fd, 0, SEEK_SET) == (off_t)-1) {
    char const* msg(strerror(errno));
    throw (exceptions::open() << "RRD: could not create file '"
           << filename << "': " << msg);
  }

  char buffer[4096];
  ssize_t transfered(0);
  while (transfered < size) {
    // Read from in_fd.
    ssize_t rb(::read(in_fd, buffer, sizeof(buffer)));
    if (rb <= 0) {
      if (errno != EAGAIN) {
        char const* msg(strerror(errno));
        throw (exceptions::open() << "RRD: could not create file '"
               << filename << "': " << msg);
      }
      continue ;
    }

    // Write to out_fd.
    ssize_t wb(0);
    while (wb < rb) {
      ssize_t ret(::write(out_fd, buffer + wb, rb - wb));
      if (ret <= 0) {
        if (errno != EAGAIN) {
          char const* msg(strerror(errno));
          throw (exceptions::open() << "RRD: could not create file '"
                 << filename << "': " << msg);
        }
      }
      else
        wb += ret;
    }

    // Update total transfered bytes.
    transfered += wb;
  }
  return ;
}

#ifdef __linux__
/**
 *  Transfer file between two FDs using sendfile.
 *
 *  @param[in] out_fd             Output FD.
 *  @param[in] in_fd              Input FD.
 *  @param[in] already_transfered Number of bytes already transfered.
 *  @param[in] size               Total size to transfer.
 *  @param[in] filename           Path to the file being created.
 */
void creator::_sendfile(
                 int out_fd,
                 int in_fd,
                 off_t already_transfered,
                 ssize_t size,
                 std::string const& filename) {
  ssize_t total(already_transfered);
  while (total < size) {
    already_transfered = total;
    ssize_t ret = ::sendfile(
                      out_fd,
                      in_fd,
                      &already_transfered,
                      size - already_transfered);
    if ((ret <= 0) && (errno != EAGAIN)) {
      char const* msg(strerror(errno));
      throw (exceptions::open() << "RRD: could not create file '"
             << filename << "': " << msg);
    }
    else if (ret > 0)
      total += ret;
  }
  return ;
}
#endif // Linux
<commit_msg>RRD: set default step to 5 minutes.<commit_after>/*
** Copyright 2013-2016 Centreon
**
** Licensed under the Apache License, Version 2.0 (the "License");
** you may not use this file except in compliance with the License.
** You may obtain a copy of the License at
**
**     http://www.apache.org/licenses/LICENSE-2.0
**
** Unless required by applicable law or agreed to in writing, software
** distributed under the License is distributed on an "AS IS" BASIS,
** WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
** See the License for the specific language governing permissions and
** limitations under the License.
**
** For more information : contact@centreon.com
*/

#include <cerrno>
#include <cstdio>
#include <cstring>
#include <fcntl.h>
#include <rrd.h>
#include <sstream>
#ifdef __linux__
#  include <sys/sendfile.h>
#endif // Linux
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include "com/centreon/broker/exceptions/msg.hh"
#include "com/centreon/broker/logging/logging.hh"
#include "com/centreon/broker/rrd/creator.hh"
#include "com/centreon/broker/rrd/exceptions/open.hh"
#include "com/centreon/broker/storage/perfdata.hh"

using namespace com::centreon::broker;
using namespace com::centreon::broker::rrd;

/**
 *  Constructor.
 *
 *  @param[in] tmpl_path  The template path.
 *  @param[in] cache_size The maximum number of cache element.
 */
creator::creator(
           std::string const& tmpl_path,
           unsigned int cache_size)
  : _cache_size(cache_size), _tmpl_path(tmpl_path) {
  logging::debug(logging::medium)
    << "RRD: file creator will maintain at most " << _cache_size
    << " templates in '" << _tmpl_path << "'";
}

/**
 *  Copy constructor.
 *
 *  @param[in] right The object to copy.
 */
creator::creator(creator const& right) {
  operator=(right);
}

/**
 *  Destructor.
 */
creator::~creator() {
  clear();
}

/**
 *  Copy operator.
 *
 *  @param[in] right The object to copy.
 *
 *  @return This object.
 */
creator& creator::operator=(creator const& right) {
  if (this != &right) {
    _cache_size = right._cache_size;
    _fds = right._fds;
    _tmpl_path = right._tmpl_path;
  }
  return (*this);
}

/**
 *  Clear cache and remove template file.
 */
void creator::clear() {
  for (std::map<tmpl_info, fd_info>::const_iterator
         it(_fds.begin()), end(_fds.end());
       it != end;
       ++it) {
    tmpl_info info(it->first);
    ::close(it->second.fd);
    std::ostringstream oss;
    oss << _tmpl_path << "/tmpl_" << info.length << "_" << info.step
        << "_" << info.value_type << ".rrd";
    ::remove(oss.str().c_str());
  }
  _fds.clear();
  return ;
}

/**
 *  Create a RRD file if it does not exists.
 *
 *  @param[in] filename   Path to the RRD file.
 *  @param[in] length     Number of recording in the RRD file.
 *  @param[in] from       Timestamp of the first record.
 *  @param[in] step       Specifies the base interval in seconds with
 *                        which data will be fed into the RRD.
 *  @param[in] value_type Type of the metric.
 */
void creator::create(
                std::string const& filename,
                unsigned int length,
                time_t from,
                unsigned int step,
                short value_type) {
  // Fill template informations.
  if (!step)
    step = 5 * 60; // Default to every 5 minutes.
  if (!length)
    length = 31 * 24 * 60 * 60 / step; // Default to one month long.
  tmpl_info info;
  info.length = length;
  info.step = step;
  info.value_type = value_type;

  // Find fd informations.
  std::map<tmpl_info, fd_info>::const_iterator it(_fds.find(info));
  // Is in the cache, just duplicate file.
  if (it != _fds.end())
    _duplicate(filename, it->second);
  // Not is the cache, but we have enough space in the cache.
  // Create new entry.
  else if (_fds.size() < _cache_size) {
    std::ostringstream oss;
    oss << _tmpl_path << "/tmpl_" << length << "_" << step
        << "_" << value_type << ".rrd";
    std::string tmpl_filename(oss.str());

    // Create new template.
    _open(tmpl_filename, length, from, step, value_type);

    // Get template file size.
    struct stat s;
    if (stat(tmpl_filename.c_str(), &s) < 0) {
      char const* msg(strerror(errno));
      throw (exceptions::open() << "RRD: could not create template file '"
             << tmpl_filename << "': " << msg);
    }

    // Get template file fd.
    int in_fd(open(tmpl_filename.c_str(), O_RDONLY));
    if (in_fd < 0) {
      char const* msg(strerror(errno));
      throw (exceptions::open() << "RRD: could not open template file '"
             << tmpl_filename << "': " << msg);
    }

    // Store fd informations into the cache.
    fd_info fdinfo;
    fdinfo.fd = in_fd;
    fdinfo.size = s.st_size;
    _fds[info] = fdinfo;

    _duplicate(filename, fdinfo);
  }
  // No more space in the cache, juste create rrd file.
  else
    _open(filename, length, from, step, value_type);
}

/**
 *  Duplicate file.
 *
 *  @param[in] filename The file name to save data.
 *  @param[in] in_fd    The fd informations to duplicate file.
 */
void creator::_duplicate(std::string const& filename, fd_info const& in_fd) {
  // Remove previous file.
  remove(filename.c_str());

  int out_fd(open(
               filename.c_str(),
               O_CREAT | O_TRUNC | O_WRONLY,
               S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH));
  if (out_fd < 0) {
    char const* msg(strerror(errno));
    throw (exceptions::open() << "RRD: could not create file '"
           << filename << "': " << msg);
  }

#ifdef __linux__
  // First call(s) to sendfile detect if the kernel support the syscall
  // on any FD (Linux < 2.6.33 only supports writing to socket FD).
  off_t offset(0);
  ssize_t ret;
  do {
    ret = ::sendfile(out_fd, in_fd.fd, &offset, in_fd.size);
  } while ((ret < 0) && (errno == EAGAIN));
  bool fallback((ret < 0) && ((errno == EINVAL) || (errno == ENOSYS)));
  if (!fallback) {
    if (ret < 0) {
      char const* msg(strerror(errno));
      throw (exceptions::open() << "RRD: could not create file '"
             << filename << "': " << msg);
    }
    // Good to go with the sendfile syscall.
    _sendfile(out_fd, in_fd.fd, ret, in_fd.size, filename);
  }
  else
#endif // Linux
    // We must fallback to the read/write combo.
    _read_write(out_fd, in_fd.fd, in_fd.size, filename);

  ::close(out_fd);
}

/**
 *  Open a RRD file and create it if it does not exists.
 *
 *  @param[in] filename   Path to the RRD file.
 *  @param[in] length     Number of recording in the RRD file.
 *  @param[in] from       Timestamp of the first record.
 *  @param[in] step       Time interval between each record.
 *  @param[in] value_type Type of the metric.
 */
void creator::_open(
                std::string const& filename,
                unsigned int length,
                time_t from,
                unsigned int step,
                short value_type) {
  /* Find step of RRD file if already existing. */
  /* XXX : why is it here ?
  rrd_info_t* rrdinfo(rrd_info_r(_filename));
  time_t interval_offset(0);
  for (rrd_info_t* tmp = rrdinfo; tmp; tmp = tmp->next)
    if (!strcmp(rrdinfo->key, "step"))
      if (interval < static_cast<time_t>(rrdinfo->value.u_cnt))
        interval_offset = rrdinfo->value.u_cnt / interval - 1;
  rrd_info_free(rrdinfo);
  */

  // Set parameters.
  std::ostringstream ds_oss;
  std::ostringstream rra1_oss;
  std::ostringstream rra2_oss;
  ds_oss << "DS:value:";
  switch (value_type) {
  case storage::perfdata::absolute:
    ds_oss << "ABSOLUTE";
    break ;
  case storage::perfdata::counter:
    ds_oss << "COUNTER";
    break ;
  case storage::perfdata::derive:
    ds_oss << "DERIVE";
    break ;
  default:
    ds_oss << "GAUGE";
  };
  ds_oss << ":"<< step * 10 << ":U:U";
  rra1_oss << "RRA:AVERAGE:0.5:1:" << length + 1;
  rra2_oss << "RRA:AVERAGE:0.5:12:" << length / 12 + 1;
  std::string ds(ds_oss.str());
  std::string rra1(rra1_oss.str());
  std::string rra2(rra2_oss.str());
  char const* argv[5];
  argv[0] = ds.c_str();
  argv[1] = rra1.c_str();
  argv[2] = rra2.c_str();
  argv[3] = NULL;

  // Debug message.
  logging::debug(logging::high) << "RRD: opening file '" << filename
    << "' (" << argv[0] << ", " << argv[1] << ", " << argv[2]
    << ", step " << step << ", from " << from << ")";

  // Create RRD file.
  rrd_clear_error();
  if (rrd_create_r(
        filename.c_str(),
        step,
        from,
        3,
        argv))
    throw (exceptions::open() << "RRD: could not create file '"
             << filename << "': " << rrd_get_error());

  return ;
}

/**
 *  Transfer file between two FDs using read and write.
 *
 *  @param[in] out_fd   Output FD.
 *  @param[in] in_fd    Input FD.
 *  @param[in] size     Size to transfer.
 *  @param[in] filename Path to the file being created.
 */
void creator::_read_write(
                int out_fd,
                int in_fd,
                ssize_t size,
                std::string const& filename) {
  // Reset position of in_fd.
  if (lseek(in_fd, 0, SEEK_SET) == (off_t)-1) {
    char const* msg(strerror(errno));
    throw (exceptions::open() << "RRD: could not create file '"
           << filename << "': " << msg);
  }

  char buffer[4096];
  ssize_t transfered(0);
  while (transfered < size) {
    // Read from in_fd.
    ssize_t rb(::read(in_fd, buffer, sizeof(buffer)));
    if (rb <= 0) {
      if (errno != EAGAIN) {
        char const* msg(strerror(errno));
        throw (exceptions::open() << "RRD: could not create file '"
               << filename << "': " << msg);
      }
      continue ;
    }

    // Write to out_fd.
    ssize_t wb(0);
    while (wb < rb) {
      ssize_t ret(::write(out_fd, buffer + wb, rb - wb));
      if (ret <= 0) {
        if (errno != EAGAIN) {
          char const* msg(strerror(errno));
          throw (exceptions::open() << "RRD: could not create file '"
                 << filename << "': " << msg);
        }
      }
      else
        wb += ret;
    }

    // Update total transfered bytes.
    transfered += wb;
  }
  return ;
}

#ifdef __linux__
/**
 *  Transfer file between two FDs using sendfile.
 *
 *  @param[in] out_fd             Output FD.
 *  @param[in] in_fd              Input FD.
 *  @param[in] already_transfered Number of bytes already transfered.
 *  @param[in] size               Total size to transfer.
 *  @param[in] filename           Path to the file being created.
 */
void creator::_sendfile(
                 int out_fd,
                 int in_fd,
                 off_t already_transfered,
                 ssize_t size,
                 std::string const& filename) {
  ssize_t total(already_transfered);
  while (total < size) {
    already_transfered = total;
    ssize_t ret = ::sendfile(
                      out_fd,
                      in_fd,
                      &already_transfered,
                      size - already_transfered);
    if ((ret <= 0) && (errno != EAGAIN)) {
      char const* msg(strerror(errno));
      throw (exceptions::open() << "RRD: could not create file '"
             << filename << "': " << msg);
    }
    else if (ret > 0)
      total += ret;
  }
  return ;
}
#endif // Linux
<|endoftext|>


<commit_before>/*-------------------------------------------------------------------
 *
 * Copyright (C) 2012 The PECOS Development Team
 *
 * Please see http://pecos.ices.utexas.edu for more information.
 *
 * This file is part of the QUESO Library (Quantification of Uncertainty
 * for Estimation, Simulation and Optimization).
 *
 * QUESO 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.
 *
 * QUESO 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 QUESO. If not, see <http://www.gnu.org/licenses/>.
 *
 *-------------------------------------------------------------------
 *
 */
 /*------------------------------------------------------------------
 * This file is divided in two parts:
 * - the first one handles the statistical inverse problem (SIP) for estimating the magnitude
 * 'g' of gravity acceleration; and
 * - the second part handles the statistical forward problem (SFP) for predicting the maximum
 * distance traveled by a projectile.
 *
 * The SIP definition requires a user defined likelihood function; refer to files
 * 'gravity_likelihood.h' and 'gravity_likelihood.C'. The SFP definition requires a user
 * defined qoi function; refer to files 'gravity_qoi.h' and 'gravity_qoi.C'.
 *-----------------------------------------------------------------*/

#include <gravity_compute.h>
#include <gravity_likelihood.h>
#include <gravity_qoi.h>
#include <queso/GslMatrix.h>
#include <queso/GenericScalarFunction.h>
#include <queso/GenericVectorFunction.h>
#include <queso/GaussianVectorRV.h>
#include <queso/UniformVectorRV.h>
#include <queso/GenericVectorRV.h>
#include <queso/StatisticalInverseProblem.h>
#include <queso/StatisticalForwardProblem.h>
#include <sys/time.h>
#include <cmath>

//================================================================
// If PRIOR_IS_GAUSSIAN is defined, then:
// --> Gaussian prior for gravity.
// Otherwise:
// --> Uniform prior for gravity.
//================================================================

//#define PRIOR_IS_GAUSSIAN

void computeGravityAndTraveledDistance(const QUESO::FullEnvironment& env) {
  struct timeval timevalNow;

  gettimeofday(&timevalNow, NULL);
  if (env.fullRank() == 0) {
    std::cout << "\nBeginning run of 'Gravity + Projectile motion' example at "
              << ctime(&timevalNow.tv_sec)
              << "\n my fullRank = "         << env.fullRank()
              << "\n my subEnvironmentId = " << env.subId()
              << "\n my subRank = "          << env.subRank()
              << "\n my interRank = "        << env.inter0Rank()
               << std::endl << std::endl;
  }

  // Just examples of possible calls
  if ((env.subDisplayFile()       ) &&
      (env.displayVerbosity() >= 2)) {
    *env.subDisplayFile() << "Beginning run of 'Gravity + Projectile motion' example at "
                          << ctime(&timevalNow.tv_sec)
                          << std::endl;
  }
  env.fullComm().Barrier();
  env.subComm().Barrier();  // Just an example of a possible call

  //================================================================
  // Statistical inverse problem (SIP): find posterior PDF for 'g'
  //================================================================
  gettimeofday(&timevalNow, NULL);
  if (env.fullRank() == 0) {
    std::cout << "Beginning 'SIP -> Gravity estimation' at "
              << ctime(&timevalNow.tv_sec)
              << std::endl;
  }

  //------------------------------------------------------
  // SIP Step 1 of 6: Instantiate the parameter space
  //------------------------------------------------------
  QUESO::VectorSpace<QUESO::GslVector,QUESO::GslMatrix> paramSpace(env, "param_", 1, NULL);

  //------------------------------------------------------
  // SIP Step 2 of 6: Instantiate the parameter domain
  //------------------------------------------------------
  QUESO::GslVector paramMinValues(paramSpace.zeroVector());
  QUESO::GslVector paramMaxValues(paramSpace.zeroVector());

  paramMinValues[0] = 8.;
  paramMaxValues[0] = 11.;

  QUESO::BoxSubset<QUESO::GslVector,QUESO::GslMatrix>
    paramDomain("param_",
      paramSpace,
      paramMinValues,
      paramMaxValues);

  //------------------------------------------------------
  // SIP Step 3 of 6: Instantiate the likelihood function
  // object to be used by QUESO.
  //------------------------------------------------------
  likelihoodRoutine_Data likelihoodRoutine_Data(env);
  QUESO::GenericScalarFunction<QUESO::GslVector,QUESO::GslMatrix>
    likelihoodFunctionObj("like_",
			  paramDomain,
			  likelihoodRoutine,
			  (void *) &likelihoodRoutine_Data,
			  true); // the routine computes [ln(function)]

  //------------------------------------------------------
  // SIP Step 4 of 6: Define the prior RV
  //------------------------------------------------------

#ifdef PRIOR_IS_GAUSSIAN
  QUESO::GslVector meanVector( paramSpace.zeroVector() );
  meanVector[0] = 9;

  QUESO::GslMatrix covMatrix = QUESO::GslMatrix(paramSpace.zeroVector());
  covMatrix(0,0) = 1.;

  // Create a Gaussian prior RV
  QUESO::GaussianVectorRV<QUESO::GslVector,QUESO::GslMatrix> priorRv("prior_",paramDomain,meanVector,covMatrix);

#else
  // Create an uniform prior RV
  QUESO::UniformVectorRV<QUESO::GslVector,QUESO::GslMatrix> priorRv("prior_",paramDomain);

#endif

  //------------------------------------------------------
  // SIP Step 5 of 6: Instantiate the inverse problem
  //------------------------------------------------------
  QUESO::GenericVectorRV<QUESO::GslVector,QUESO::GslMatrix>
    postRv("post_",  // Extra prefix before the default "rv_" prefix
           paramSpace);

  QUESO::StatisticalInverseProblem<QUESO::GslVector,QUESO::GslMatrix>
    ip("",          // No extra prefix before the default "ip_" prefix
       NULL,
       priorRv,
       likelihoodFunctionObj,
       postRv);

  //------------------------------------------------------
  // SIP Step 6 of 6: Solve the inverse problem, that is,
  // set the 'pdf' and the 'realizer' of the posterior RV
  //------------------------------------------------------
  std::cout << "Solving the SIP with Metropolis Hastings"
	    << std::endl << std::endl;

  QUESO::GslVector paramInitials(paramSpace.zeroVector());
  priorRv.realizer().realization(paramInitials);

  QUESO::GslMatrix proposalCovMatrix(paramSpace.zeroVector());
  proposalCovMatrix(0,0) = std::pow( fabs(paramInitials[0])/20. , 2. );

  ip.solveWithBayesMetropolisHastings(NULL, paramInitials, &proposalCovMatrix);

  //================================================================
  // Statistical forward problem (SFP): find the max distance
  // traveled by an object in projectile motion; input pdf for 'g'
  // is the solution of the SIP above.
  //================================================================
  gettimeofday(&timevalNow, NULL);
  std::cout << "Beginning 'SFP -> Projectile motion' at "
            << ctime(&timevalNow.tv_sec)
            << std::endl;

  //------------------------------------------------------
  // SFP Step 1 of 6: Instantiate the parameter *and* qoi spaces.
  // SFP input RV = FIP posterior RV, so SFP parameter space
  // has been already defined.
  //------------------------------------------------------
  QUESO::VectorSpace<QUESO::GslVector,QUESO::GslMatrix> qoiSpace(env, "qoi_", 1, NULL);

  //------------------------------------------------------
  // SFP Step 2 of 6: Instantiate the parameter domain
  //------------------------------------------------------

  // Not necessary because input RV of the SFP = output RV of SIP.
  // Thus, the parameter domain has been already defined.

  //------------------------------------------------------
  // SFP Step 3 of 6: Instantiate the qoi function object
  // to be used by QUESO.
  //------------------------------------------------------
  qoiRoutine_Data qoiRoutine_Data;
  qoiRoutine_Data.m_angle          = M_PI/4.0; //45 degrees (radians)
  qoiRoutine_Data.m_initialVelocity= 5.;      //initial speed (m/s)
  qoiRoutine_Data.m_initialHeight  = 0.;       //initial height (m)

  QUESO::GenericVectorFunction<QUESO::GslVector,QUESO::GslMatrix,QUESO::GslVector,QUESO::GslMatrix>
    qoiFunctionObj("qoi_",
                   paramDomain,
                   qoiSpace,
                   qoiRoutine,
                   (void *) &qoiRoutine_Data);

  //------------------------------------------------------
  // SFP Step 4 of 6: Define the input RV
  //------------------------------------------------------

  // Not necessary because input RV of SFP = output RV of SIP
  // (postRv).

  //------------------------------------------------------
  // SFP Step 5 of 6: Instantiate the forward problem
  //------------------------------------------------------
  QUESO::GenericVectorRV<QUESO::GslVector,QUESO::GslMatrix> qoiRv("qoi_", qoiSpace);

  QUESO::StatisticalForwardProblem<QUESO::GslVector,QUESO::GslMatrix,QUESO::GslVector,QUESO::GslMatrix>
    fp("",
       NULL,
       postRv,
       qoiFunctionObj,
       qoiRv);

  //------------------------------------------------------
  // SFP Step 6 of 6: Solve the forward problem
  //------------------------------------------------------
  std::cout << "Solving the SFP with Monte Carlo"
            << std::endl << std::endl;
  fp.solveWithMonteCarlo(NULL);

  //------------------------------------------------------
  gettimeofday(&timevalNow, NULL);
  if ((env.subDisplayFile()       ) &&
      (env.displayVerbosity() >= 2)) {
    *env.subDisplayFile() << "Ending run of 'Gravity + Projectile motion' example at "
                          << ctime(&timevalNow.tv_sec)
                          << std::endl;
  }
  if (env.fullRank() == 0) {
    std::cout << "Ending run of 'Gravity + Projectile motion' example at "
              << ctime(&timevalNow.tv_sec)
              << std::endl;
  }

  return;
}
<commit_msg>Remove trailing whitespace and update licence<commit_after>//-----------------------------------------------------------------------bl-
//--------------------------------------------------------------------------
//
// QUESO - a library to support the Quantification of Uncertainty
// for Estimation, Simulation and Optimization
//
// Copyright (C) 2008,2009,2010,2011,2012,2013 The PECOS Development Team
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the Version 2.1 GNU Lesser General
// Public License as published by the Free Software Foundation.
//
// This library 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
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc. 51 Franklin Street, Fifth Floor,
// Boston, MA  02110-1301  USA
//
//-----------------------------------------------------------------------el-

/*
 * This file is divided in two parts:
 * - the first one handles the statistical inverse problem (SIP) for estimating
 *   the magnitude 'g' of gravity acceleration; and
 * - the second part handles the statistical forward problem (SFP) for
 *   predicting the maximum distance traveled by a projectile.
 *
 * The SIP definition requires a user defined likelihood function; refer to
 * files 'gravity_likelihood.h' and 'gravity_likelihood.C'. The SFP definition
 * requires a user defined qoi function; refer to files 'gravity_qoi.h' and
 * 'gravity_qoi.C'.
 */

#include <cmath>
#include <sys/time.h>

#include <queso/GslMatrix.h>
#include <queso/GenericScalarFunction.h>
#include <queso/GenericVectorFunction.h>
#include <queso/GaussianVectorRV.h>
#include <queso/UniformVectorRV.h>
#include <queso/GenericVectorRV.h>
#include <queso/StatisticalInverseProblem.h>
#include <queso/StatisticalForwardProblem.h>

#include <gravity_compute.h>
#include <gravity_likelihood.h>
#include <gravity_qoi.h>

void computeGravityAndTraveledDistance(const QUESO::FullEnvironment& env) {
  struct timeval timevalNow;

  gettimeofday(&timevalNow, NULL);
  if (env.fullRank() == 0) {
    std::cout << "\nBeginning run of 'Gravity + Projectile motion' example at "
              << ctime(&timevalNow.tv_sec)
              << "\n my fullRank = "         << env.fullRank()
              << "\n my subEnvironmentId = " << env.subId()
              << "\n my subRank = "          << env.subRank()
              << "\n my interRank = "        << env.inter0Rank()
               << std::endl << std::endl;
  }

  // Just examples of possible calls
  if ((env.subDisplayFile()       ) &&
      (env.displayVerbosity() >= 2)) {
    *env.subDisplayFile() << "Beginning run of 'Gravity + Projectile motion' example at "
                          << ctime(&timevalNow.tv_sec)
                          << std::endl;
  }
  env.fullComm().Barrier();
  env.subComm().Barrier();  // Just an example of a possible call

  //================================================================
  // Statistical inverse problem (SIP): find posterior PDF for 'g'
  //================================================================
  gettimeofday(&timevalNow, NULL);
  if (env.fullRank() == 0) {
    std::cout << "Beginning 'SIP -> Gravity estimation' at "
              << ctime(&timevalNow.tv_sec)
              << std::endl;
  }

  //------------------------------------------------------
  // SIP Step 1 of 6: Instantiate the parameter space
  //------------------------------------------------------
  QUESO::VectorSpace<QUESO::GslVector,QUESO::GslMatrix> paramSpace(env,
      "param_", 1, NULL);

  //------------------------------------------------------
  // SIP Step 2 of 6: Instantiate the parameter domain
  //------------------------------------------------------
  QUESO::GslVector paramMinValues(paramSpace.zeroVector());
  QUESO::GslVector paramMaxValues(paramSpace.zeroVector());

  paramMinValues[0] = 8.;
  paramMaxValues[0] = 11.;

  QUESO::BoxSubset<QUESO::GslVector,QUESO::GslMatrix> paramDomain("param_",
      paramSpace, paramMinValues, paramMaxValues);

  //------------------------------------------------------
  // SIP Step 3 of 6: Instantiate the likelihood function
  // object to be used by QUESO.
  //------------------------------------------------------
  Likelihood<QUESO::GslVector, QUESO::GslMatrix> lhood("like_", paramDomain);

  //------------------------------------------------------
  // SIP Step 4 of 6: Define the prior RV
  //------------------------------------------------------
  QUESO::UniformVectorRV<QUESO::GslVector,QUESO::GslMatrix> priorRv("prior_",
      paramDomain);

  //------------------------------------------------------
  // SIP Step 5 of 6: Instantiate the inverse problem
  //------------------------------------------------------
  QUESO::GenericVectorRV<QUESO::GslVector,QUESO::GslMatrix>
    postRv("post_",  // Extra prefix before the default "rv_" prefix
           paramSpace);

  QUESO::StatisticalInverseProblem<QUESO::GslVector,QUESO::GslMatrix>
    ip("",          // No extra prefix before the default "ip_" prefix
       NULL,
       priorRv,
       lhood,
       postRv);

  //------------------------------------------------------
  // SIP Step 6 of 6: Solve the inverse problem, that is,
  // set the 'pdf' and the 'realizer' of the posterior RV
  //------------------------------------------------------

  QUESO::GslVector paramInitials(paramSpace.zeroVector());
  priorRv.realizer().realization(paramInitials);

  QUESO::GslMatrix proposalCovMatrix(paramSpace.zeroVector());
  proposalCovMatrix(0,0) = std::pow(std::abs(paramInitials[0]) / 20.0, 2.0);

  ip.solveWithBayesMetropolisHastings(NULL, paramInitials, &proposalCovMatrix);

  //================================================================
  // Statistical forward problem (SFP): find the max distance
  // traveled by an object in projectile motion; input pdf for 'g'
  // is the solution of the SIP above.
  //================================================================
  gettimeofday(&timevalNow, NULL);
  std::cout << "Beginning 'SFP -> Projectile motion' at "
            << ctime(&timevalNow.tv_sec)
            << std::endl;

  //------------------------------------------------------
  // SFP Step 1 of 6: Instantiate the parameter *and* qoi spaces.
  // SFP input RV = FIP posterior RV, so SFP parameter space
  // has been already defined.
  //------------------------------------------------------
  QUESO::VectorSpace<QUESO::GslVector,QUESO::GslMatrix> qoiSpace(env, "qoi_",
      1, NULL);

  //------------------------------------------------------
  // SFP Step 2 of 6: Instantiate the parameter domain
  //------------------------------------------------------

  // Not necessary because input RV of the SFP = output RV of SIP.
  // Thus, the parameter domain has been already defined.

  //------------------------------------------------------
  // SFP Step 3 of 6: Instantiate the qoi object
  // to be used by QUESO.
  //------------------------------------------------------
  Qoi<QUESO::GslVector, QUESO::GslMatrix, QUESO::GslVector, QUESO::GslMatrix>
    qoi("qoi_", paramDomain, qoiSpace);

  //------------------------------------------------------
  // SFP Step 4 of 6: Define the input RV
  //------------------------------------------------------

  // Not necessary because input RV of SFP = output RV of SIP
  // (postRv).

  //------------------------------------------------------
  // SFP Step 5 of 6: Instantiate the forward problem
  //------------------------------------------------------
  QUESO::GenericVectorRV<QUESO::GslVector, QUESO::GslMatrix> qoiRv("qoi_",
      qoiSpace);

  QUESO::StatisticalForwardProblem<QUESO::GslVector, QUESO::GslMatrix,
    QUESO::GslVector, QUESO::GslMatrix> fp("", NULL, postRv, qoi, qoiRv);

  //------------------------------------------------------
  // SFP Step 6 of 6: Solve the forward problem
  //------------------------------------------------------
  std::cout << "Solving the SFP with Monte Carlo"
            << std::endl << std::endl;
  fp.solveWithMonteCarlo(NULL);

  gettimeofday(&timevalNow, NULL);
  if ((env.subDisplayFile()       ) &&
      (env.displayVerbosity() >= 2)) {
    *env.subDisplayFile() << "Ending run of 'Gravity + Projectile motion' example at "
                          << ctime(&timevalNow.tv_sec)
                          << std::endl;
  }
  if (env.fullRank() == 0) {
    std::cout << "Ending run of 'Gravity + Projectile motion' example at "
              << ctime(&timevalNow.tv_sec)
              << std::endl;
  }
}
<|endoftext|>


<commit_before>/*
Request handlers for WSGI apps and static files

Copyright (c) 2016 Roman Miroshnychenko <romanvm@yandex.ua>
License: MIT, see License.txt
*/

#include "request_handlers.h"

#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/copy.hpp>

#include <fstream>
#include <iostream>
#include <sstream>
#include <utility>

using namespace std;
namespace py = boost::python;
namespace fs = boost::filesystem;
namespace sys = boost::system;
namespace alg = boost::algorithm;

namespace wsgi_boost
{
#pragma region StaticRequestHandler

	void StaticRequestHandler::handle()
	{
		const auto content_dir_path = fs::path{ m_request.content_dir };
		if (!fs::exists(content_dir_path))
		{
			m_response.send_html("500 Internal Server Error",
				"Error 500",
				"Internal Server Error",
				"Invalid static content directory is configured.");
			return;
		}
		if (m_request.method != "GET" && m_request.method != "HEAD")
		{
			m_response.send_mesage("405 Method Not Allowed", "");
			return;
		}
		open_file(content_dir_path);
	}

	void StaticRequestHandler::open_file(const fs::path& content_dir_path)
	{
		fs::path path = content_dir_path;
		path /= boost::regex_replace(m_request.path, m_request.path_regex, "");
		if (fs::exists(path))
		{
			path = fs::canonical(path);
			// Checking if path is inside content_dir
			if (distance(content_dir_path.begin(), content_dir_path.end()) <= distance(path.begin(), path.end()) &&
				equal(content_dir_path.begin(), content_dir_path.end(), path.begin()))
			{
				if (fs::is_directory(path))
				{
					path /= "index.html";
				}
				if (fs::exists(path) && fs::is_regular_file(path))
				{
					ifstream ifs;
					ifs.open(path.string(), ifstream::in | ios::binary);
					if (ifs)
					{
						headers_type out_headers;
						out_headers.emplace_back("Cache-Control", m_cache_control);
						time_t last_modified = fs::last_write_time(path);
						out_headers.emplace_back("Last-Modified", time_to_header(last_modified));
						// Use hex representation of the last modified POSIX timestamp as ETag
						string etag = "\"" + hex(static_cast<size_t>(last_modified)) + "\"";
						out_headers.emplace_back("ETag", etag);
						string ims = m_request.get_header("If-Modified-Since");
						if (m_request.get_header("If-None-Match") == etag || (ims != "" && header_to_time(ims) >= last_modified))
						{
							out_headers.emplace_back("Content-Length", "0");
							m_response.send_header("304 Not Modified", out_headers);
							ifs.close();
							return;
						}
						string ext = path.extension().string();
						out_headers.emplace_back("Content-Type", get_mime(ext));
						if (m_use_gzip && m_request.check_header("Accept-Encoding", "gzip") && is_compressable(ext))
						{
							boost::iostreams::filtering_istream gzstream;
							gzstream.push(boost::iostreams::gzip_compressor());
							gzstream.push(ifs);
							stringstream compressed;
							boost::iostreams::copy(gzstream, compressed);
							out_headers.emplace_back("Content-Encoding", "gzip");
							send_file(compressed, out_headers);
						}
						else
						{
							out_headers.emplace_back("Accept-Ranges", "bytes");
							send_file(ifs, out_headers);
						}
						ifs.close();
						return;
					}
				}
			}
		}
		m_response.send_html("404 Not Found",
			"Error 404", "Not Found",
			"The requested path " + m_request.path + " was not found on this server.");
	}

	void StaticRequestHandler::send_file(istream& content_stream, headers_type& headers)
	{
		content_stream.seekg(0, ios::end);
		size_t length = content_stream.tellg();
		size_t start_pos = 0;
		size_t end_pos = length - 1;
		string requested_range = m_request.get_header("Range");
		pair<string, string> range;
		if (requested_range != "" && ((range = parse_range(requested_range)) != pair<string, string>()))
		{
			if (range.first != string())
				start_pos = stoull(range.first);
			else
				range.first = to_string(start_pos);
			if (range.second != string())
				end_pos = stoull(range.second);
			else
				range.second = to_string(end_pos);
			if (start_pos > end_pos || start_pos >= length || end_pos >= length)
			{
				m_response.send_mesage("416 Range Not Satisfiable", "");
				return;
			}
			else
			{
				headers.emplace_back("Content-Length", to_string(end_pos - start_pos));
				headers.emplace_back("Content-Range", "bytes " + range.first + "-" + range.second + "/" + to_string(length));
				m_response.send_header("206 Partial Content", headers);
			}
		}
		else
		{
			headers.emplace_back("Content-Length", to_string(length));
			m_response.send_header("200 OK", headers);
		}
		if (m_request.method == "GET")
		{
			if (start_pos > 0)
				content_stream.seekg(start_pos);
			else
				content_stream.seekg(0, ios::beg);
			const size_t buffer_size = 131072;
			vector<char> buffer(buffer_size);
			size_t read_length;
			size_t bytes_left = end_pos - start_pos + 1;
			while (bytes_left > 0 &&
				((read_length = content_stream.read(&buffer[0], min(bytes_left, buffer_size)).gcount()) > 0))
			{
				sys::error_code ec = m_response.send_data(string{ &buffer[0], read_length });
				if (ec)
					return;
				bytes_left -= read_length;
			}
		}
	}

#pragma endregion

#pragma region WsgiRequestHandler

	WsgiRequestHandler::WsgiRequestHandler(Request& request, Response& response, py::object& app,
		string& scheme, string& host, unsigned short local_port, bool multithread) :
		BaseRequestHandler(request, response), m_app{ app },
		m_url_scheme{ scheme }, m_host_name{ host },
		m_local_port{ local_port }, m_multithread{ multithread }
	{
		m_write = create_write();
		m_start_response = create_start_response();
	}

	py::object WsgiRequestHandler::create_write()
	{
		function<void(py::object)> wr{
			[this](py::object data)
			{
				string cpp_data = py::extract<string>(data);
				GilRelease release_gil;
				// Read/write operations executed from inside Python must be syncronous!
				sys::error_code ec = this->m_response.send_header(this->m_status, this->m_out_headers, false);
				if (ec)
					return;
				size_t data_len = cpp_data.length();
				if (data_len > 0)
				{
					if (this->m_content_length == -1)
						cpp_data = hex(data_len) + "\r\n" + cpp_data + "\r\n";
					// Read/write operations executed from inside Python must be syncronous!
					this->m_response.send_data(cpp_data, false);
				}
			}
		};
		return py::make_function(wr, py::default_call_policies(),
			py::args("data"), boost::mpl::vector<void, py::object>());
	}

	py::object WsgiRequestHandler::create_start_response()
	{
		function<py::object(py::str, py::list, py::object)> sr{
			[this](py::str status, py::list headers, py::object exc_info = py::object())
			{
				if (!exc_info.is_none() && this->m_response.header_sent())
				{
					py::object t = exc_info[0];
					py::object v = exc_info[1];
					py::object tb = exc_info[2];
					cerr << "The WSGI app passed exc_info after sending headers to the client!\n";
					py::import("traceback").attr("print_exception")(t, v, tb);
					PyErr_Restore(t.ptr(), v.ptr(), tb.ptr());
					py::throw_error_already_set();
				}
				this->m_status = py::extract<string>(status);
				this->m_out_headers.clear();
				for (py::ssize_t i = 0; i < py::len(headers); ++i)
				{
					string header_name = py::extract<string>(headers[i][0]);
					string header_value = py::extract<string>(headers[i][1]);
					if (alg::iequals(header_name, "Content-Length"))
					{
						try
						{
							this->m_content_length = stoll(header_value);
						}
						catch (const logic_error&) {}
					}
					this->m_out_headers.emplace_back(header_name, header_value);
				}
				if (this->m_content_length == -1)
				{
					// If a WSGI app does not provide Content-Length header (e.g. Django)
					// we use Transfer-Encoding: chunked
					this->m_out_headers.emplace_back("Transfer-Encoding", "chunked");
				}
				return this->m_write;
			}
		};
		return py::make_function(sr, py::default_call_policies(),
			(py::arg("status"), py::arg("headers"), py::arg("exc_info") = py::object()),
			boost::mpl::vector<py::object, py::str, py::list, py::object>());
	}

	void WsgiRequestHandler::prepare_environ()
	{
		m_environ["REQUEST_METHOD"] = m_request.method;
		m_environ["SCRIPT_NAME"] = "";
		pair<string, string> path_and_query = split_path(m_request.path);
		m_environ["PATH_INFO"] = path_and_query.first;
		m_environ["QUERY_STRING"] = path_and_query.second;
		m_environ["CONTENT_TYPE"] = m_request.get_header("Content-Type");
		m_environ["CONTENT_LENGTH"] = m_request.get_header("Content-Length");
		m_environ["SERVER_NAME"] = m_host_name;
		m_environ["SERVER_PORT"] = to_string(m_local_port);
		m_environ["SERVER_PROTOCOL"] = m_request.http_version;
		m_environ["REMOTE_ADDR"] = m_environ["REMOTE_HOST"] = m_request.remote_address();
		m_environ["REMOTE_PORT"] = to_string(m_request.remote_port());
		for (const auto& header : m_request.headers)
		{
			string env_header = header.first;
			transform_header(env_header);
			if (env_header == "HTTP_CONTENT_TYPE" || env_header == "HTTP_CONTENT_LENGTH")
				continue;
			// Headers are already checked for duplicates during parsing
			m_environ[env_header] = header.second;
		}
		m_environ["wsgi.version"] = py::make_tuple<int, int>(1, 0);
		m_environ["wsgi.url_scheme"] = m_url_scheme;
		m_environ["wsgi.input"] = InputStream{ m_request.connection() };
		m_environ["wsgi.errors"] = ErrorStream{};
		m_environ["wsgi.file_wrapper"] = FileWrapper{};
		m_environ["wsgi.multithread"] = m_multithread;
		m_environ["wsgi.multiprocess"] = false;
		m_environ["wsgi.run_once"] = false;
	}

	void WsgiRequestHandler::handle()
	{
		if (m_app.is_none())
		{
			m_response.send_html("500 Internal Server Error",
				"Error 500",  "Internal Server Error",
				"A WSGI application is not set.");
			return;
		}
		prepare_environ();
		Iterable iterable{ move(m_app(m_environ, m_start_response)) };
		send_iterable(iterable);
	}

	void WsgiRequestHandler::send_iterable(Iterable& iterable)
	{
		py::object iterator = iterable.attr("__iter__")();
		while (true)
		{
			try
			{
#if PY_MAJOR_VERSION < 3
				std::string chunk = py::extract<string>(iterator.attr("next")());
#else
				std::string chunk = py::extract<string>(iterator.attr("__next__")());
#endif
				// Releasing GIL around async operations not only allows other Python threads to run
				// but also allows io_service to safely transfer control to another coroutine
				// that may acquire GIL again.
				// I found this scheme by accident and if we do not release GIL at this point
				// Python will crash!
				GilRelease release_gil;
				sys::error_code ec;
				if (!m_response.header_sent())
				{
					ec = m_response.send_header(m_status, m_out_headers);
					if (ec)
						break;
				}
				if (m_content_length == -1)
				{
					size_t length = chunk.length();
					// Skip 0-length chunks, if any
					if (length == 0)
						continue;
					chunk = hex(length) + "\r\n" + chunk + "\r\n";
				}
				ec = m_response.send_data(chunk);
				if (ec)
					break;
			}
			catch (const py::error_already_set&)
			{
				if (PyErr_ExceptionMatches(PyExc_StopIteration))
				{
					PyErr_Clear();
					if (m_content_length == -1)
						m_response.send_data("0\r\n\r\n");
					break;
				}
				throw;
			}
		}
	}
#pragma endregion
}
<commit_msg>Check http method first when serving static files<commit_after>/*
Request handlers for WSGI apps and static files

Copyright (c) 2016 Roman Miroshnychenko <romanvm@yandex.ua>
License: MIT, see License.txt
*/

#include "request_handlers.h"

#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/copy.hpp>

#include <fstream>
#include <iostream>
#include <sstream>
#include <utility>

using namespace std;
namespace py = boost::python;
namespace fs = boost::filesystem;
namespace sys = boost::system;
namespace alg = boost::algorithm;

namespace wsgi_boost
{
#pragma region StaticRequestHandler

	void StaticRequestHandler::handle()
	{
		if (m_request.method != "GET" && m_request.method != "HEAD")
		{
			m_response.send_mesage("405 Method Not Allowed", "Invalid HTTP method! Only GET and HEAD are allowed.");
			return;
		}
		const auto content_dir_path = fs::path{ m_request.content_dir };
		if (!fs::exists(content_dir_path))
		{
			m_response.send_html("500 Internal Server Error",
				"Error 500",
				"Internal Server Error",
				"Invalid static content directory is configured.");
			return;
		}
		open_file(content_dir_path);
	}

	void StaticRequestHandler::open_file(const fs::path& content_dir_path)
	{
		fs::path path = content_dir_path;
		path /= boost::regex_replace(m_request.path, m_request.path_regex, "");
		if (fs::exists(path))
		{
			path = fs::canonical(path);
			// Checking if path is inside content_dir
			if (distance(content_dir_path.begin(), content_dir_path.end()) <= distance(path.begin(), path.end()) &&
				equal(content_dir_path.begin(), content_dir_path.end(), path.begin()))
			{
				if (fs::is_directory(path))
				{
					path /= "index.html";
				}
				if (fs::exists(path) && fs::is_regular_file(path))
				{
					ifstream ifs;
					ifs.open(path.string(), ifstream::in | ios::binary);
					if (ifs)
					{
						headers_type out_headers;
						out_headers.emplace_back("Cache-Control", m_cache_control);
						time_t last_modified = fs::last_write_time(path);
						out_headers.emplace_back("Last-Modified", time_to_header(last_modified));
						// Use hex representation of the last modified POSIX timestamp as ETag
						string etag = "\"" + hex(static_cast<size_t>(last_modified)) + "\"";
						out_headers.emplace_back("ETag", etag);
						string ims = m_request.get_header("If-Modified-Since");
						if (m_request.get_header("If-None-Match") == etag || (ims != "" && header_to_time(ims) >= last_modified))
						{
							out_headers.emplace_back("Content-Length", "0");
							m_response.send_header("304 Not Modified", out_headers);
							ifs.close();
							return;
						}
						string ext = path.extension().string();
						out_headers.emplace_back("Content-Type", get_mime(ext));
						if (m_use_gzip && m_request.check_header("Accept-Encoding", "gzip") && is_compressable(ext))
						{
							boost::iostreams::filtering_istream gzstream;
							gzstream.push(boost::iostreams::gzip_compressor());
							gzstream.push(ifs);
							stringstream compressed;
							boost::iostreams::copy(gzstream, compressed);
							out_headers.emplace_back("Content-Encoding", "gzip");
							send_file(compressed, out_headers);
						}
						else
						{
							out_headers.emplace_back("Accept-Ranges", "bytes");
							send_file(ifs, out_headers);
						}
						ifs.close();
						return;
					}
				}
			}
		}
		m_response.send_html("404 Not Found",
			"Error 404", "Not Found",
			"The requested path " + m_request.path + " was not found on this server.");
	}

	void StaticRequestHandler::send_file(istream& content_stream, headers_type& headers)
	{
		content_stream.seekg(0, ios::end);
		size_t length = content_stream.tellg();
		size_t start_pos = 0;
		size_t end_pos = length - 1;
		string requested_range = m_request.get_header("Range");
		pair<string, string> range;
		if (requested_range != "" && ((range = parse_range(requested_range)) != pair<string, string>()))
		{
			if (range.first != string())
				start_pos = stoull(range.first);
			else
				range.first = to_string(start_pos);
			if (range.second != string())
				end_pos = stoull(range.second);
			else
				range.second = to_string(end_pos);
			if (start_pos > end_pos || start_pos >= length || end_pos >= length)
			{
				m_response.send_mesage("416 Range Not Satisfiable", "");
				return;
			}
			else
			{
				headers.emplace_back("Content-Length", to_string(end_pos - start_pos));
				headers.emplace_back("Content-Range", "bytes " + range.first + "-" + range.second + "/" + to_string(length));
				m_response.send_header("206 Partial Content", headers);
			}
		}
		else
		{
			headers.emplace_back("Content-Length", to_string(length));
			m_response.send_header("200 OK", headers);
		}
		if (m_request.method == "GET")
		{
			if (start_pos > 0)
				content_stream.seekg(start_pos);
			else
				content_stream.seekg(0, ios::beg);
			const size_t buffer_size = 131072;
			vector<char> buffer(buffer_size);
			size_t read_length;
			size_t bytes_left = end_pos - start_pos + 1;
			while (bytes_left > 0 &&
				((read_length = content_stream.read(&buffer[0], min(bytes_left, buffer_size)).gcount()) > 0))
			{
				sys::error_code ec = m_response.send_data(string{ &buffer[0], read_length });
				if (ec)
					return;
				bytes_left -= read_length;
			}
		}
	}

#pragma endregion

#pragma region WsgiRequestHandler

	WsgiRequestHandler::WsgiRequestHandler(Request& request, Response& response, py::object& app,
		string& scheme, string& host, unsigned short local_port, bool multithread) :
		BaseRequestHandler(request, response), m_app{ app },
		m_url_scheme{ scheme }, m_host_name{ host },
		m_local_port{ local_port }, m_multithread{ multithread }
	{
		m_write = create_write();
		m_start_response = create_start_response();
	}

	py::object WsgiRequestHandler::create_write()
	{
		function<void(py::object)> wr{
			[this](py::object data)
			{
				string cpp_data = py::extract<string>(data);
				GilRelease release_gil;
				// Read/write operations executed from inside Python must be syncronous!
				sys::error_code ec = this->m_response.send_header(this->m_status, this->m_out_headers, false);
				if (ec)
					return;
				size_t data_len = cpp_data.length();
				if (data_len > 0)
				{
					if (this->m_content_length == -1)
						cpp_data = hex(data_len) + "\r\n" + cpp_data + "\r\n";
					// Read/write operations executed from inside Python must be syncronous!
					this->m_response.send_data(cpp_data, false);
				}
			}
		};
		return py::make_function(wr, py::default_call_policies(),
			py::args("data"), boost::mpl::vector<void, py::object>());
	}

	py::object WsgiRequestHandler::create_start_response()
	{
		function<py::object(py::str, py::list, py::object)> sr{
			[this](py::str status, py::list headers, py::object exc_info = py::object())
			{
				if (!exc_info.is_none() && this->m_response.header_sent())
				{
					py::object t = exc_info[0];
					py::object v = exc_info[1];
					py::object tb = exc_info[2];
					cerr << "The WSGI app passed exc_info after sending headers to the client!\n";
					py::import("traceback").attr("print_exception")(t, v, tb);
					PyErr_Restore(t.ptr(), v.ptr(), tb.ptr());
					py::throw_error_already_set();
				}
				this->m_status = py::extract<string>(status);
				this->m_out_headers.clear();
				for (py::ssize_t i = 0; i < py::len(headers); ++i)
				{
					string header_name = py::extract<string>(headers[i][0]);
					string header_value = py::extract<string>(headers[i][1]);
					if (alg::iequals(header_name, "Content-Length"))
					{
						try
						{
							this->m_content_length = stoll(header_value);
						}
						catch (const logic_error&) {}
					}
					this->m_out_headers.emplace_back(header_name, header_value);
				}
				if (this->m_content_length == -1)
				{
					// If a WSGI app does not provide Content-Length header (e.g. Django)
					// we use Transfer-Encoding: chunked
					this->m_out_headers.emplace_back("Transfer-Encoding", "chunked");
				}
				return this->m_write;
			}
		};
		return py::make_function(sr, py::default_call_policies(),
			(py::arg("status"), py::arg("headers"), py::arg("exc_info") = py::object()),
			boost::mpl::vector<py::object, py::str, py::list, py::object>());
	}

	void WsgiRequestHandler::prepare_environ()
	{
		m_environ["REQUEST_METHOD"] = m_request.method;
		m_environ["SCRIPT_NAME"] = "";
		pair<string, string> path_and_query = split_path(m_request.path);
		m_environ["PATH_INFO"] = path_and_query.first;
		m_environ["QUERY_STRING"] = path_and_query.second;
		m_environ["CONTENT_TYPE"] = m_request.get_header("Content-Type");
		m_environ["CONTENT_LENGTH"] = m_request.get_header("Content-Length");
		m_environ["SERVER_NAME"] = m_host_name;
		m_environ["SERVER_PORT"] = to_string(m_local_port);
		m_environ["SERVER_PROTOCOL"] = m_request.http_version;
		m_environ["REMOTE_ADDR"] = m_environ["REMOTE_HOST"] = m_request.remote_address();
		m_environ["REMOTE_PORT"] = to_string(m_request.remote_port());
		for (const auto& header : m_request.headers)
		{
			string env_header = header.first;
			transform_header(env_header);
			if (env_header == "HTTP_CONTENT_TYPE" || env_header == "HTTP_CONTENT_LENGTH")
				continue;
			// Headers are already checked for duplicates during parsing
			m_environ[env_header] = header.second;
		}
		m_environ["wsgi.version"] = py::make_tuple<int, int>(1, 0);
		m_environ["wsgi.url_scheme"] = m_url_scheme;
		m_environ["wsgi.input"] = InputStream{ m_request.connection() };
		m_environ["wsgi.errors"] = ErrorStream{};
		m_environ["wsgi.file_wrapper"] = FileWrapper{};
		m_environ["wsgi.multithread"] = m_multithread;
		m_environ["wsgi.multiprocess"] = false;
		m_environ["wsgi.run_once"] = false;
	}

	void WsgiRequestHandler::handle()
	{
		if (m_app.is_none())
		{
			m_response.send_html("500 Internal Server Error",
				"Error 500",  "Internal Server Error",
				"A WSGI application is not set.");
			return;
		}
		prepare_environ();
		Iterable iterable{ move(m_app(m_environ, m_start_response)) };
		send_iterable(iterable);
	}

	void WsgiRequestHandler::send_iterable(Iterable& iterable)
	{
		py::object iterator = iterable.attr("__iter__")();
		while (true)
		{
			try
			{
#if PY_MAJOR_VERSION < 3
				std::string chunk = py::extract<string>(iterator.attr("next")());
#else
				std::string chunk = py::extract<string>(iterator.attr("__next__")());
#endif
				// Releasing GIL around async operations not only allows other Python threads to run
				// but also allows io_service to safely transfer control to another coroutine
				// that may acquire GIL again.
				// I found this scheme by accident and if we do not release GIL at this point
				// Python will crash!
				GilRelease release_gil;
				sys::error_code ec;
				if (!m_response.header_sent())
				{
					ec = m_response.send_header(m_status, m_out_headers);
					if (ec)
						break;
				}
				if (m_content_length == -1)
				{
					size_t length = chunk.length();
					// Skip 0-length chunks, if any
					if (length == 0)
						continue;
					chunk = hex(length) + "\r\n" + chunk + "\r\n";
				}
				ec = m_response.send_data(chunk);
				if (ec)
					break;
			}
			catch (const py::error_already_set&)
			{
				if (PyErr_ExceptionMatches(PyExc_StopIteration))
				{
					PyErr_Clear();
					if (m_content_length == -1)
						m_response.send_data("0\r\n\r\n");
					break;
				}
				throw;
			}
		}
	}
#pragma endregion
}
<|endoftext|>


<commit_before>/*
* Copyright (C) 2008-2011 J-P Nurmi jpnurmi@gmail.com
*
* 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.
*/

#include "messagehandler.h"
#include <qabstractsocket.h>
#include <ircsession.h>
#include <qvariant.h>
#include <qdebug.h>
#include <irc.h>

MessageHandler::MessageHandler(QObject* parent) : QObject(parent)
{
    d.session = 0;
    d.defaultReceiver = 0;
    d.currentReceiver = 0;
    setSession(qobject_cast<IrcSession*>(parent));
}

MessageHandler::~MessageHandler()
{
    d.defaultReceiver = 0;
    d.currentReceiver = 0;
    d.receivers.clear();
    d.channelUsers.clear();
    if (d.session && d.session->isActive())
        d.session->socket()->waitForDisconnected(500);
}

IrcSession* MessageHandler::session() const
{
    return d.session;
}

void MessageHandler::setSession(IrcSession* session)
{
    if (d.session != session)
    {
        if (d.session)
            disconnect(d.session, SIGNAL(messageReceived(IrcMessage*)), this, SLOT(handleMessage(IrcMessage*)));

        if (session)
            connect(session, SIGNAL(messageReceived(IrcMessage*)), this, SLOT(handleMessage(IrcMessage*)));

        d.session = session;
    }
}

QObject* MessageHandler::defaultReceiver() const
{
    return d.defaultReceiver;
}

void MessageHandler::setDefaultReceiver(QObject* receiver)
{
    d.defaultReceiver = receiver;
}

QObject* MessageHandler::currentReceiver() const
{
    return d.currentReceiver;
}

void MessageHandler::setCurrentReceiver(QObject* receiver)
{
    d.currentReceiver = receiver;
}

void MessageHandler::addReceiver(const QString& name, QObject* receiver)
{
    d.receivers.insert(name.toLower(), receiver);
}

QObject* MessageHandler::getReceiver(const QString& name) const
{
    return d.receivers.value(name.toLower());
}

void MessageHandler::removeReceiver(const QString& name)
{
    const QString lower = name.toLower();
    if (d.receivers.contains(lower))
    {
        d.receivers.remove(name.toLower());
        emit receiverToBeRemoved(name);
    }
}

void MessageHandler::handleMessage(IrcMessage* message)
{
    switch (message->type())
    {
    case IrcMessage::Invite:
        handleInviteMessage(static_cast<IrcInviteMessage*>(message));
        break;
    case IrcMessage::Join:
        handleJoinMessage(static_cast<IrcJoinMessage*>(message));
        break;
    case IrcMessage::Kick:
        handleKickMessage(static_cast<IrcKickMessage*>(message));
        break;
    case IrcMessage::Mode:
        handleModeMessage(static_cast<IrcModeMessage*>(message));
        break;
    case IrcMessage::Nick:
        handleNickMessage(static_cast<IrcNickMessage*>(message));
        break;
    case IrcMessage::Notice:
        handleNoticeMessage(static_cast<IrcNoticeMessage*>(message));
        break;
    case IrcMessage::Numeric:
        handleNumericMessage(static_cast<IrcNumericMessage*>(message));
        break;
    case IrcMessage::Part:
        handlePartMessage(static_cast<IrcPartMessage*>(message));
        break;
    case IrcMessage::Pong:
        handlePongMessage(static_cast<IrcPongMessage*>(message));
        break;
    case IrcMessage::Private:
        handlePrivateMessage(static_cast<IrcPrivateMessage*>(message));
        break;
    case IrcMessage::Quit:
        handleQuitMessage(static_cast<IrcQuitMessage*>(message));
        break;
    case IrcMessage::Topic:
        handleTopicMessage(static_cast<IrcTopicMessage*>(message));
        break;
    case IrcMessage::Unknown:
        handleUnknownMessage(static_cast<IrcMessage*>(message));
        break;
    default:
        break;
    }
}

void MessageHandler::handleInviteMessage(IrcInviteMessage* message)
{
    sendMessage(message, d.currentReceiver);
}

void MessageHandler::handleJoinMessage(IrcJoinMessage* message)
{
    sendMessage(message, message->channel());
    d.addChannelUser(message->channel(), message->sender().name());
}

void MessageHandler::handleKickMessage(IrcKickMessage* message)
{
    sendMessage(message, message->channel());
    d.removeChannelUser(message->channel(), message->user());
}

void MessageHandler::handleModeMessage(IrcModeMessage* message)
{
    if (message->sender().name() == message->target())
        sendMessage(message, d.defaultReceiver);
    else
        sendMessage(message, message->target());
}

void MessageHandler::handleNickMessage(IrcNickMessage* message)
{
    QString nick = message->sender().name().toLower();
    foreach (const QString& channel, d.userChannels(nick))
    {
        sendMessage(message, channel);
        d.removeChannelUser(channel, nick);
        d.addChannelUser(channel, message->nick());
    }
    foreach (const QString& receiver, d.receivers.keys())
    {
        if (!nick.compare(receiver, Qt::CaseInsensitive))
        {
            emit receiverToBeRenamed(receiver, message->nick());
            QObject* object = d.receivers.take(nick);
            d.receivers.insert(nick, object);
            sendMessage(message, object);
        }
    }
}

void MessageHandler::handleNoticeMessage(IrcNoticeMessage* message)
{
    sendMessage(message, d.currentReceiver);
}

void MessageHandler::handleNumericMessage(IrcNumericMessage* message)
{
    switch (message->code())
    {
    case Irc::RPL_ENDOFWHO:
    case Irc::RPL_WHOREPLY:
    case Irc::RPL_UNAWAY:
    case Irc::RPL_NOWAWAY:
    case Irc::RPL_AWAY:
    case Irc::RPL_WHOISOPERATOR:
    case Irc::RPL_WHOISHELPOP: // "is available for help"
    case Irc::RPL_WHOISSPECIAL: // "is identified to services"
    case Irc::RPL_WHOISHOST: // nick is connecting from <...>
    case Irc::RPL_WHOISSECURE: // nick is using a secure connection
    case Irc::RPL_WHOISUSER:
    case Irc::RPL_WHOISSERVER:
    case Irc::RPL_WHOISACCOUNT: // nick user is logged in as
    case Irc::RPL_WHOWASUSER:
    case Irc::RPL_WHOISIDLE:
    case Irc::RPL_WHOISCHANNELS:
    case Irc::RPL_ENDOFWHOIS:
    case Irc::RPL_INVITING:
    case Irc::RPL_VERSION:
    case Irc::RPL_TIME:
        sendMessage(message, d.currentReceiver);
        break;

    case Irc::RPL_ENDOFBANLIST:
    case Irc::RPL_ENDOFEXCEPTLIST:
    case Irc::RPL_ENDOFINFO:
    case Irc::RPL_ENDOFINVITELIST:
    case Irc::RPL_ENDOFLINKS:
    case Irc::RPL_ENDOFMOTD:
    case Irc::RPL_ENDOFSTATS:
    case Irc::RPL_ENDOFUSERS:
    case Irc::RPL_ENDOFWHOWAS:
        break; // ignore

    case Irc::RPL_ENDOFNAMES:
    case Irc::RPL_CHANNELMODEIS:
    case Irc::RPL_CHANNEL_URL:
    case Irc::RPL_CREATIONTIME:
    case Irc::RPL_NOTOPIC:
    case Irc::RPL_TOPIC:
    case Irc::RPL_TOPICWHOTIME:
        sendMessage(message, message->parameters().value(1));
        break;

    case Irc::RPL_NAMREPLY: {
        const int count = message->parameters().count();
        const QString channel = message->parameters().value(count - 2);
        const QStringList names = message->parameters().value(count - 1).split(" ", QString::SkipEmptyParts);
        foreach (QString name, names)
        {
            if (name.startsWith("@") || name.startsWith("+"))
                name.remove(0, 1);
            d.addChannelUser(channel, name);
        }
        sendMessage(message, channel);
        break;
        }

    default:
        sendMessage(message, d.defaultReceiver);
        break;
    }
}

void MessageHandler::handlePartMessage(IrcPartMessage* message)
{
    d.removeChannelUser(message->channel(), message->sender().name());
    if (message->sender().name() == d.session->nickName())
        removeReceiver(message->channel());
    else
        sendMessage(message, message->channel());
}

void MessageHandler::handlePongMessage(IrcPongMessage* message)
{
    sendMessage(message, d.currentReceiver);
}

void MessageHandler::handlePrivateMessage(IrcPrivateMessage* message)
{
    if (message->isRequest())
        sendMessage(message, d.currentReceiver);
    else if (message->target() == d.session->nickName())
        sendMessage(message, message->sender().name());
    else
        sendMessage(message, message->target());
}

void MessageHandler::handleQuitMessage(IrcQuitMessage* message)
{
    QString nick = message->sender().name();
    foreach (const QString& channel, d.userChannels(nick))
    {
        sendMessage(message, channel);
        d.removeChannelUser(channel, nick);
    }

    if (d.receivers.contains(nick.toLower()))
        sendMessage(message, nick);

    if (nick == d.session->nickName())
        removeReceiver(d.session->host());
}

void MessageHandler::handleTopicMessage(IrcTopicMessage* message)
{
    sendMessage(message, message->channel());
}

void MessageHandler::handleUnknownMessage(IrcMessage* message)
{
    sendMessage(message, d.defaultReceiver);
}

void MessageHandler::sendMessage(IrcMessage* message, QObject* receiver)
{
    QMetaObject::invokeMethod(receiver, "receiveMessage", Q_ARG(IrcMessage*, message));

    // if the receiver was a QML element, argument type would be QVariant(QObject*):
    // QMetaObject::invokeMethod(receiver, "receiveMessage", Q_ARG(QVariant, QVariant::fromValue((QObject*) message)));
}

void MessageHandler::sendMessage(IrcMessage* message, const QString& receiver)
{
    if (!d.receivers.contains(receiver.toLower()))
        emit receiverToBeAdded(receiver);
    QObject* object = getReceiver(receiver);
    Q_ASSERT(object);
    sendMessage(message, object);
}

void MessageHandler::onSessionDestroyed()
{
    setSession(0);
}

QStringList MessageHandler::Private::userChannels(const QString& user) const
{
    QStringList channels;
    QHash<QString, QSet<QString> >::const_iterator it;
    for (it = channelUsers.constBegin(); it != channelUsers.constEnd(); ++it)
        if (it.value().contains(user.toLower()))
            channels += it.key();
    return channels;
}

void MessageHandler::Private::addChannelUser(QString channel, const QString& user)
{
    channel = channel.toLower();
    if (!channelUsers.value(channel).contains(user.toLower()))
    {
        channelUsers[channel].insert(user.toLower());
        QObject* receiver = receivers.value(channel);
        if (receiver)
            QMetaObject::invokeMethod(receiver, "addUser", Q_ARG(QString, user));
    }
}

void MessageHandler::Private::removeChannelUser(QString channel, const QString& user)
{
    channel = channel.toLower();
    if (channelUsers.value(channel).contains(user.toLower()))
    {
        channelUsers[channel].remove(user.toLower());
        QObject* receiver = receivers.value(channel);
        if (receiver)
            QMetaObject::invokeMethod(receiver, "removeUser", Q_ARG(QString, user));
    }
}
<commit_msg>MessageHandler: cleanup channelUsers when parting a channel<commit_after>/*
* Copyright (C) 2008-2011 J-P Nurmi jpnurmi@gmail.com
*
* 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.
*/

#include "messagehandler.h"
#include <qabstractsocket.h>
#include <ircsession.h>
#include <qvariant.h>
#include <qdebug.h>
#include <irc.h>

MessageHandler::MessageHandler(QObject* parent) : QObject(parent)
{
    d.session = 0;
    d.defaultReceiver = 0;
    d.currentReceiver = 0;
    setSession(qobject_cast<IrcSession*>(parent));
}

MessageHandler::~MessageHandler()
{
    d.defaultReceiver = 0;
    d.currentReceiver = 0;
    d.receivers.clear();
    d.channelUsers.clear();
    if (d.session && d.session->isActive())
        d.session->socket()->waitForDisconnected(500);
}

IrcSession* MessageHandler::session() const
{
    return d.session;
}

void MessageHandler::setSession(IrcSession* session)
{
    if (d.session != session)
    {
        if (d.session)
            disconnect(d.session, SIGNAL(messageReceived(IrcMessage*)), this, SLOT(handleMessage(IrcMessage*)));

        if (session)
            connect(session, SIGNAL(messageReceived(IrcMessage*)), this, SLOT(handleMessage(IrcMessage*)));

        d.session = session;
    }
}

QObject* MessageHandler::defaultReceiver() const
{
    return d.defaultReceiver;
}

void MessageHandler::setDefaultReceiver(QObject* receiver)
{
    d.defaultReceiver = receiver;
}

QObject* MessageHandler::currentReceiver() const
{
    return d.currentReceiver;
}

void MessageHandler::setCurrentReceiver(QObject* receiver)
{
    d.currentReceiver = receiver;
}

void MessageHandler::addReceiver(const QString& name, QObject* receiver)
{
    d.receivers.insert(name.toLower(), receiver);
}

QObject* MessageHandler::getReceiver(const QString& name) const
{
    return d.receivers.value(name.toLower());
}

void MessageHandler::removeReceiver(const QString& name)
{
    const QString lower = name.toLower();
    if (d.receivers.contains(lower))
    {
        d.receivers.remove(name.toLower());
        emit receiverToBeRemoved(name);
    }
}

void MessageHandler::handleMessage(IrcMessage* message)
{
    switch (message->type())
    {
    case IrcMessage::Invite:
        handleInviteMessage(static_cast<IrcInviteMessage*>(message));
        break;
    case IrcMessage::Join:
        handleJoinMessage(static_cast<IrcJoinMessage*>(message));
        break;
    case IrcMessage::Kick:
        handleKickMessage(static_cast<IrcKickMessage*>(message));
        break;
    case IrcMessage::Mode:
        handleModeMessage(static_cast<IrcModeMessage*>(message));
        break;
    case IrcMessage::Nick:
        handleNickMessage(static_cast<IrcNickMessage*>(message));
        break;
    case IrcMessage::Notice:
        handleNoticeMessage(static_cast<IrcNoticeMessage*>(message));
        break;
    case IrcMessage::Numeric:
        handleNumericMessage(static_cast<IrcNumericMessage*>(message));
        break;
    case IrcMessage::Part:
        handlePartMessage(static_cast<IrcPartMessage*>(message));
        break;
    case IrcMessage::Pong:
        handlePongMessage(static_cast<IrcPongMessage*>(message));
        break;
    case IrcMessage::Private:
        handlePrivateMessage(static_cast<IrcPrivateMessage*>(message));
        break;
    case IrcMessage::Quit:
        handleQuitMessage(static_cast<IrcQuitMessage*>(message));
        break;
    case IrcMessage::Topic:
        handleTopicMessage(static_cast<IrcTopicMessage*>(message));
        break;
    case IrcMessage::Unknown:
        handleUnknownMessage(static_cast<IrcMessage*>(message));
        break;
    default:
        break;
    }
}

void MessageHandler::handleInviteMessage(IrcInviteMessage* message)
{
    sendMessage(message, d.currentReceiver);
}

void MessageHandler::handleJoinMessage(IrcJoinMessage* message)
{
    sendMessage(message, message->channel());
    d.addChannelUser(message->channel(), message->sender().name());
}

void MessageHandler::handleKickMessage(IrcKickMessage* message)
{
    sendMessage(message, message->channel());
    d.removeChannelUser(message->channel(), message->user());
}

void MessageHandler::handleModeMessage(IrcModeMessage* message)
{
    if (message->sender().name() == message->target())
        sendMessage(message, d.defaultReceiver);
    else
        sendMessage(message, message->target());
}

void MessageHandler::handleNickMessage(IrcNickMessage* message)
{
    QString nick = message->sender().name().toLower();
    foreach (const QString& channel, d.userChannels(nick))
    {
        sendMessage(message, channel);
        d.removeChannelUser(channel, nick);
        d.addChannelUser(channel, message->nick());
    }
    foreach (const QString& receiver, d.receivers.keys())
    {
        if (!nick.compare(receiver, Qt::CaseInsensitive))
        {
            emit receiverToBeRenamed(receiver, message->nick());
            QObject* object = d.receivers.take(nick);
            d.receivers.insert(nick, object);
            sendMessage(message, object);
        }
    }
}

void MessageHandler::handleNoticeMessage(IrcNoticeMessage* message)
{
    sendMessage(message, d.currentReceiver);
}

void MessageHandler::handleNumericMessage(IrcNumericMessage* message)
{
    switch (message->code())
    {
    case Irc::RPL_ENDOFWHO:
    case Irc::RPL_WHOREPLY:
    case Irc::RPL_UNAWAY:
    case Irc::RPL_NOWAWAY:
    case Irc::RPL_AWAY:
    case Irc::RPL_WHOISOPERATOR:
    case Irc::RPL_WHOISHELPOP: // "is available for help"
    case Irc::RPL_WHOISSPECIAL: // "is identified to services"
    case Irc::RPL_WHOISHOST: // nick is connecting from <...>
    case Irc::RPL_WHOISSECURE: // nick is using a secure connection
    case Irc::RPL_WHOISUSER:
    case Irc::RPL_WHOISSERVER:
    case Irc::RPL_WHOISACCOUNT: // nick user is logged in as
    case Irc::RPL_WHOWASUSER:
    case Irc::RPL_WHOISIDLE:
    case Irc::RPL_WHOISCHANNELS:
    case Irc::RPL_ENDOFWHOIS:
    case Irc::RPL_INVITING:
    case Irc::RPL_VERSION:
    case Irc::RPL_TIME:
        sendMessage(message, d.currentReceiver);
        break;

    case Irc::RPL_ENDOFBANLIST:
    case Irc::RPL_ENDOFEXCEPTLIST:
    case Irc::RPL_ENDOFINFO:
    case Irc::RPL_ENDOFINVITELIST:
    case Irc::RPL_ENDOFLINKS:
    case Irc::RPL_ENDOFMOTD:
    case Irc::RPL_ENDOFSTATS:
    case Irc::RPL_ENDOFUSERS:
    case Irc::RPL_ENDOFWHOWAS:
        break; // ignore

    case Irc::RPL_ENDOFNAMES:
    case Irc::RPL_CHANNELMODEIS:
    case Irc::RPL_CHANNEL_URL:
    case Irc::RPL_CREATIONTIME:
    case Irc::RPL_NOTOPIC:
    case Irc::RPL_TOPIC:
    case Irc::RPL_TOPICWHOTIME:
        sendMessage(message, message->parameters().value(1));
        break;

    case Irc::RPL_NAMREPLY: {
        const int count = message->parameters().count();
        const QString channel = message->parameters().value(count - 2);
        const QStringList names = message->parameters().value(count - 1).split(" ", QString::SkipEmptyParts);
        foreach (QString name, names)
        {
            if (name.startsWith("@") || name.startsWith("+"))
                name.remove(0, 1);
            d.addChannelUser(channel, name);
        }
        sendMessage(message, channel);
        break;
        }

    default:
        sendMessage(message, d.defaultReceiver);
        break;
    }
}

void MessageHandler::handlePartMessage(IrcPartMessage* message)
{
    d.removeChannelUser(message->channel(), message->sender().name());
    if (message->sender().name() == d.session->nickName())
    {
        removeReceiver(message->channel());
        d.channelUsers.remove(message->channel().toLower());
    }
    else
    {
        sendMessage(message, message->channel());
    }
}

void MessageHandler::handlePongMessage(IrcPongMessage* message)
{
    sendMessage(message, d.currentReceiver);
}

void MessageHandler::handlePrivateMessage(IrcPrivateMessage* message)
{
    if (message->isRequest())
        sendMessage(message, d.currentReceiver);
    else if (message->target() == d.session->nickName())
        sendMessage(message, message->sender().name());
    else
        sendMessage(message, message->target());
}

void MessageHandler::handleQuitMessage(IrcQuitMessage* message)
{
    QString nick = message->sender().name();
    foreach (const QString& channel, d.userChannels(nick))
    {
        sendMessage(message, channel);
        d.removeChannelUser(channel, nick);
    }

    if (d.receivers.contains(nick.toLower()))
        sendMessage(message, nick);

    if (nick == d.session->nickName())
        removeReceiver(d.session->host());
}

void MessageHandler::handleTopicMessage(IrcTopicMessage* message)
{
    sendMessage(message, message->channel());
}

void MessageHandler::handleUnknownMessage(IrcMessage* message)
{
    sendMessage(message, d.defaultReceiver);
}

void MessageHandler::sendMessage(IrcMessage* message, QObject* receiver)
{
    QMetaObject::invokeMethod(receiver, "receiveMessage", Q_ARG(IrcMessage*, message));

    // if the receiver was a QML element, argument type would be QVariant(QObject*):
    // QMetaObject::invokeMethod(receiver, "receiveMessage", Q_ARG(QVariant, QVariant::fromValue((QObject*) message)));
}

void MessageHandler::sendMessage(IrcMessage* message, const QString& receiver)
{
    if (!d.receivers.contains(receiver.toLower()))
        emit receiverToBeAdded(receiver);
    QObject* object = getReceiver(receiver);
    Q_ASSERT(object);
    sendMessage(message, object);
}

void MessageHandler::onSessionDestroyed()
{
    setSession(0);
}

QStringList MessageHandler::Private::userChannels(const QString& user) const
{
    QStringList channels;
    QHash<QString, QSet<QString> >::const_iterator it;
    for (it = channelUsers.constBegin(); it != channelUsers.constEnd(); ++it)
        if (it.value().contains(user.toLower()))
            channels += it.key();
    return channels;
}

void MessageHandler::Private::addChannelUser(QString channel, const QString& user)
{
    channel = channel.toLower();
    if (!channelUsers.value(channel).contains(user.toLower()))
    {
        channelUsers[channel].insert(user.toLower());
        QObject* receiver = receivers.value(channel);
        if (receiver)
            QMetaObject::invokeMethod(receiver, "addUser", Q_ARG(QString, user));
    }
}

void MessageHandler::Private::removeChannelUser(QString channel, const QString& user)
{
    channel = channel.toLower();
    if (channelUsers.value(channel).contains(user.toLower()))
    {
        channelUsers[channel].remove(user.toLower());
        QObject* receiver = receivers.value(channel);
        if (receiver)
            QMetaObject::invokeMethod(receiver, "removeUser", Q_ARG(QString, user));
    }
}
<|endoftext|>