aria2/src/CookieStorage.cc

/* <!-- copyright */
/*
 * aria2 - The high speed download utility
 *
 * Copyright (C) 2013 Tatsuhiro Tsujikawa
 *
 * 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 *
 * In addition, as a special exception, the copyright holders give
 * permission to link the code of portions of this program with the
 * OpenSSL library under certain conditions as described in each
 * individual source file, and distribute linked combinations
 * including the two.
 * You must obey the GNU General Public License in all respects
 * for all of the code used other than OpenSSL.  If you modify
 * file(s) with this exception, you may extend this exception to your
 * version of the file(s), but you are not obligated to do so.  If you
 * do not wish to do so, delete this exception statement from your
 * version.  If you delete this exception statement from all source
 * files in the program, then also delete it here.
 */
/* copyright --> */
#include "CookieStorage.h"

#include <cstring>
#include <cstdio>
#include <algorithm>

#include "util.h"
#include "LogFactory.h"
#include "Logger.h"
#include "DlAbortEx.h"
#include "fmt.h"
#include "NsCookieParser.h"
#include "File.h"
#include "a2functional.h"
#include "A2STR.h"
#include "message.h"
#include "cookie_helper.h"
#include "BufferedFile.h"
#ifdef HAVE_SQLITE3
# include "Sqlite3CookieParserImpl.h"
#endif // HAVE_SQLITE3

namespace aria2 {

DomainNode::DomainNode(std::string label, DomainNode* parent)
  : label_{std::move(label)}, parent_{parent}, lastAccessTime_{0},
    lruAccessTime_{0}, inLru_{false}
{}

void DomainNode::findCookie
(std::vector<const Cookie*>& out,
 const std::string& requestHost,
 const std::string& requestPath,
 time_t now, bool secure)
{
  if(cookies_) {
    for(auto& c : *cookies_) {
      if(c->match(requestHost, requestPath, now, secure)) {
        c->setLastAccessTime(now);
        out.push_back(c.get());
      }
    }
  }
}

bool DomainNode::addCookie(std::unique_ptr<Cookie> cookie, time_t now)
{
  using namespace std::placeholders;
  setLastAccessTime(now);
  if(!cookies_) {
    if(cookie->isExpired(now)) {
      return false;
    } else {
      cookies_ = make_unique<std::deque<std::unique_ptr<Cookie>>>();
      cookies_->push_back(std::move(cookie));
      return true;
    }
  }

  auto i = std::find_if(std::begin(*cookies_), std::end(*cookies_),
                        [&](const std::unique_ptr<Cookie>& c)
                        {
                          return *c == *cookie;
                        });
  if(i == std::end(*cookies_)) {
    if(cookie->isExpired(now)) {
      return false;
    } else {
      if(cookies_->size() >= CookieStorage::MAX_COOKIE_PER_DOMAIN) {
        cookies_->erase
          (std::remove_if(std::begin(*cookies_), std::end(*cookies_),
                          std::bind(&Cookie::isExpired, _1, now)),
           std::end(*cookies_));
        if(cookies_->size() >= CookieStorage::MAX_COOKIE_PER_DOMAIN) {
          auto m = std::min_element(std::begin(*cookies_), std::end(*cookies_),
                                    [](const std::unique_ptr<Cookie>& lhs,
                                       const std::unique_ptr<Cookie>& rhs)
                                    {
                                      return lhs->getLastAccessTime() <
                                      rhs->getLastAccessTime();
                                    });
          *m = std::move(cookie);
        } else {
          cookies_->push_back(std::move(cookie));
        }
      } else {
        cookies_->push_back(std::move(cookie));
      }
      return true;
    }
  } else if(cookie->isExpired(now)) {
    cookies_->erase(i);
    return false;
  } else {
    *i = std::move(cookie);
    return true;
  }
}

bool DomainNode::contains(const Cookie& cookie) const
{
  if(cookies_) {
    for(auto& i : *cookies_) {
      if(*i == cookie) {
        return true;
      }
    }
  }
  return false;
}

bool DomainNode::writeCookie(BufferedFile& fp) const
{
  if(cookies_) {
    for(const auto& c : *cookies_) {
      std::string data = c->toNsCookieFormat();
      data += "\n";
      if(fp.write(data.data(), data.size()) != data.size()) {
        return false;
      }
    }
  }
  return true;
}

size_t DomainNode::countCookie() const
{
  if(cookies_) {
    return cookies_->size();
  } else {
    return 0;
  }
}

void DomainNode::clearCookie()
{
  cookies_->clear();
}

void DomainNode::setLastAccessTime(time_t lastAccessTime)
{
  lastAccessTime_ = lastAccessTime;
}

time_t DomainNode::getLastAccessTime() const
{
  return lastAccessTime_;
}

void DomainNode::setLruAccessTime(time_t t)
{
  lruAccessTime_ = t;
}

time_t DomainNode::getLruAccessTime() const
{
  return lruAccessTime_;
}

bool DomainNode::empty() const
{
  return !cookies_ || cookies_->empty();
}

bool DomainNode::hasNext() const
{
  return !next_.empty();
}

DomainNode* DomainNode::getParent() const
{
  return parent_;
}

void DomainNode::removeNode(DomainNode* node)
{
  next_.erase(node->getLabel());
}

DomainNode* DomainNode::findNext(const std::string& label) const
{
  auto i = next_.find(label);
  if(i == std::end(next_)) {
    return nullptr;
  } else {
    return (*i).second.get();
  }
}

DomainNode* DomainNode::addNext(std::string label,
                                std::unique_ptr<DomainNode> node)
{
  auto& res = next_[std::move(label)] = std::move(node);
  return res.get();
}

const std::string& DomainNode::getLabel() const
{
  return label_;
}

bool DomainNode::getInLru() const
{
  return inLru_;
}

void DomainNode::setInLru(bool f)
{
  inLru_ = f;
}

CookieStorage::CookieStorage()
  : rootNode_{make_unique<DomainNode>("", nullptr)}
{}

namespace {
// See CookieStorageTest::testDomainIsFull() in CookieStorageTest.cc
const size_t DOMAIN_EVICTION_TRIGGER = 2000;

const double DOMAIN_EVICTION_RATE = 0.1;
} // namespace

namespace {
std::vector<std::string> splitDomainLabel(const std::string& domain)
{
  auto labels = std::vector<std::string>{};
  if(util::isNumericHost(domain)) {
    labels.push_back(domain);
  } else {
    util::split(std::begin(domain), std::end(domain),
                std::back_inserter(labels), '.');
  }
  return labels;
}
} // namespace


size_t CookieStorage::getLruTrackerSize() const
{
  return lruTracker_.size();
}

void CookieStorage::evictNode(size_t delnum)
{
  for(; delnum > 0 && !lruTracker_.empty(); --delnum) {
    auto node = (*lruTracker_.begin()).second;
    lruTracker_.erase(lruTracker_.begin());
    node->setInLru(false);
    node->clearCookie();
    while(node->empty() && !node->hasNext()) {
      auto parent = node->getParent();
      parent->removeNode(node);
      if(!parent->empty() || parent->hasNext() ||
         parent == rootNode_.get()) {
        break;
      }
      node = parent;
      if(node->getInLru()) {
        lruTracker_.erase({node->getLruAccessTime(), node});
        node->setInLru(false);
      }
    }
  }
}

const DomainNode* CookieStorage::getRootNode() const
{
  return rootNode_.get();
}

bool CookieStorage::store(std::unique_ptr<Cookie> cookie, time_t now)
{
  if(lruTracker_.size() >= DOMAIN_EVICTION_TRIGGER) {
    auto delnum = size_t(lruTracker_.size()*DOMAIN_EVICTION_RATE);
    evictNode(delnum);
  }
  auto labels = splitDomainLabel(cookie->getDomain());
  auto node = rootNode_.get();
  for(auto i = labels.rbegin(), eoi = labels.rend(); i != eoi; ++i) {
    auto nextNode = node->findNext(*i);
    if(nextNode) {
      node = nextNode;
    } else {
      node = node->addNext(*i, make_unique<DomainNode>(*i, node));
    }
  }
  bool ok = node->addCookie(std::move(cookie), now);
  if(ok) {
    updateLru(node, now);
  }
  return ok;
}

void CookieStorage::updateLru(DomainNode* node, time_t now)
{
  if(node->getInLru()) {
    lruTracker_.erase({node->getLruAccessTime(), node});
  } else {
    node->setInLru(true);
  }
  node->setLruAccessTime(now);
  lruTracker_.insert({node->getLruAccessTime(), node});
}

bool CookieStorage::parseAndStore
(const std::string& setCookieString,
 const std::string& requestHost,
 const std::string& defaultPath,
 time_t now)
{
  auto cookie = cookie::parse(setCookieString, requestHost, defaultPath, now);
  return cookie && store(std::move(cookie), now);
}

namespace {
struct CookiePathDivider {
  const Cookie* cookie_;
  int pathDepth_;
  CookiePathDivider(const Cookie* cookie):cookie_(cookie), pathDepth_(0)
  {
    const std::string& path = cookie_->getPath();
    if(!path.empty()) {
      for(size_t i = 1, len = path.size(); i < len; ++i) {
        if(path[i] == '/' && path[i-1] != '/') {
          ++pathDepth_;
        }
      }
      if(path[path.size()-1] != '/') {
        ++pathDepth_;
      }
    }
  }
};
} // namespace

namespace {
class CookiePathDividerConverter {
public:
  CookiePathDivider operator()(const Cookie* cookie) const
  {
    return CookiePathDivider(cookie);
  }

  const Cookie* operator()(const CookiePathDivider& cookiePathDivider) const
  {
    return cookiePathDivider.cookie_;
  }
};
} // namespace

namespace {
class OrderByPathDepthDesc:public std::binary_function<Cookie, Cookie, bool> {
public:
  bool operator()
  (const CookiePathDivider& lhs, const CookiePathDivider& rhs) const
  {
    // From http://tools.ietf.org/html/rfc6265#section-5.4:
    // 2.  The user agent SHOULD sort the cookie-list in the following
    //    order:
    //
    //    *  Cookies with longer paths are listed before cookies with
    //       shorter paths.
    //
    //    *  Among cookies that have equal-length path fields, cookies with
    //       earlier creation-times are listed before cookies with later
    //       creation-times.
    return lhs.pathDepth_ > rhs.pathDepth_ ||
      (!(rhs.pathDepth_ > lhs.pathDepth_) &&
       lhs.cookie_->getCreationTime() < rhs.cookie_->getCreationTime());
  }
};
} // namespace

namespace {
DomainNode* findNode(const std::string& domain, DomainNode* node)
{
  auto labels = splitDomainLabel(domain);
  for(auto i = labels.rbegin(), eoi = labels.rend(); i != eoi && node; ++i) {
    node = node->findNext(*i);
  }
  return node;
}
} // namespace

bool CookieStorage::contains(const Cookie& cookie) const
{
  auto node = findNode(cookie.getDomain(), rootNode_.get());
  return node && node->contains(cookie);
}

std::vector<const Cookie*> CookieStorage::criteriaFind
(const std::string& requestHost,
 const std::string& requestPath,
 time_t now,
 bool secure)
{
  auto res = std::vector<const Cookie*>{};
  if(requestPath.empty()) {
    return res;
  }
  auto labels = splitDomainLabel(requestHost);
  auto node = rootNode_.get();
  for (auto i = labels.rbegin(), eoi = labels.rend(); i != eoi; ++i) {
    auto nextNode = node->findNext(*i);
    if(!nextNode) {
      break;
    }
    nextNode->setLastAccessTime(now);
    if(nextNode->getInLru()) {
      updateLru(nextNode, now);
    }
    nextNode->findCookie(res, requestHost, requestPath, now, secure);
    node = nextNode;
  }
  auto divs = std::vector<CookiePathDivider>{};
  std::transform(std::begin(res), std::end(res), std::back_inserter(divs),
                 CookiePathDividerConverter{});
  std::sort(std::begin(divs), std::end(divs), OrderByPathDepthDesc{});
  std::transform(std::begin(divs), std::end(divs), std::begin(res),
                 CookiePathDividerConverter{});
  return res;
}

size_t CookieStorage::size() const
{
  size_t n = 0;
  for(auto& p : lruTracker_) {
    n += p.second->countCookie();
  }
  return n;
}

bool CookieStorage::load(const std::string& filename, time_t now)
{
  char header[16]; // "SQLite format 3" plus \0
  size_t headlen;
  {
    BufferedFile fp{filename.c_str(), BufferedFile::READ};
    if(!fp) {
      A2_LOG_ERROR(fmt("Failed to open cookie file %s", filename.c_str()));
      return false;
    }
    headlen = fp.read(header, sizeof(header));
  }
  try {
    if(headlen == 16 && memcmp(header, "SQLite format 3\0", 16) == 0) {
#ifdef HAVE_SQLITE3
      try {
        auto cookies = Sqlite3MozCookieParser(filename).parse();
        storeCookies(std::make_move_iterator(std::begin(cookies)),
                     std::make_move_iterator(std::end(cookies)), now);
      } catch(RecoverableException& e) {
        A2_LOG_INFO_EX(EX_EXCEPTION_CAUGHT, e);
        A2_LOG_INFO("This does not look like Firefox3 cookie file."
                    " Retrying, assuming it is Chromium cookie file.");
        // Try chrome cookie format
        auto cookies = Sqlite3ChromiumCookieParser(filename).parse();
        storeCookies(std::make_move_iterator(std::begin(cookies)),
                     std::make_move_iterator(std::end(cookies)), now);
      }
#else // !HAVE_SQLITE3
      throw DL_ABORT_EX
        ("Cannot read SQLite3 database because SQLite3 support is disabled by"
         " configuration.");
#endif // !HAVE_SQLITE3
    } else {
      auto cookies = NsCookieParser().parse(filename, now);
      storeCookies(std::make_move_iterator(std::begin(cookies)),
                   std::make_move_iterator(std::end(cookies)), now);
    }
    return true;
  } catch(RecoverableException& e) {
    A2_LOG_ERROR(fmt("Failed to load cookies from %s", filename.c_str()));
    return false;
  }
}

bool CookieStorage::saveNsFormat(const std::string& filename)
{
  auto tempfilename = filename;
  tempfilename += "__temp";
  {
    BufferedFile fp{tempfilename.c_str(), BufferedFile::WRITE};
    if(!fp) {
      A2_LOG_ERROR(fmt("Cannot create cookie file %s", filename.c_str()));
      return false;
    }
    for(auto& p : lruTracker_) {
      if(!p.second->writeCookie(fp)) {
        A2_LOG_ERROR(fmt("Failed to save cookies to %s", filename.c_str()));
        return false;
      }
    }
    if(fp.close() == EOF) {
      A2_LOG_ERROR(fmt("Failed to save cookies to %s", filename.c_str()));
      return false;
    }
  }
  if(File(tempfilename).renameTo(filename)) {
    return true;
  } else {
    A2_LOG_ERROR(fmt("Could not rename file %s as %s",
                     tempfilename.c_str(),
                     filename.c_str()));
    return false;
  }
}

} // namespace aria2