<commit_before>#include "src/cryptography/Crypto.hpp"
#include "src/cryptography/CryptoState.hpp"
#include "src/cryptography/BotanCrypto.hpp"
#include "src/cryptography/constants.h"
#include "src/utility/pimpl_impl.h"
#include <QFileInfo>
#include <QDir>
#include <QStringRef>
#include <QStringBuilder>
#include <string>
#include <memory>

#if defined(_MSC_VER)
#include "src/utility/make_unique.h"
#endif

QString outputFilePath(const QString& inputFilePath,
                       const QString& inputFileName,
                       const QString& outputPath)
{
  auto path = QString{inputFilePath % QStringLiteral(".") % kExtension};

  if (!outputPath.isEmpty())
  {
    path = QDir::cleanPath(outputPath) % QDir::separator() % inputFileName %
           QStringLiteral(".") % kExtension;
  }

  return path;
}

class Crypto::CryptoPrivate {
 public:
  /*!
   * \brief CryptoPrivate Constructs the Crypto private implementation
   */
  CryptoPrivate();

  /*!
   * \brief errorMessage Returns error message
   * \return String containing error message
   */
  QString errorMessage(const int index) const;

  std::unique_ptr<CryptoState> state;
  std::unique_ptr<BotanCrypto> botanCrypto;

  // The list of status messages that can be displayed to the user
  const QStringList messages;
};

Crypto::Crypto(QObject* parent)
  : QObject{parent}
{
  // Subscribe to provider's signals
  connect(m->botanCrypto.get(), &BotanCrypto::progress,
          this, &Crypto::progress);
  connect(m->botanCrypto.get(), &BotanCrypto::statusMessage,
          this, &Crypto::statusMessage);
  connect(m->botanCrypto.get(), &BotanCrypto::errorMessage,
          this, &Crypto::errorMessage);
}

Crypto::~Crypto() {}

void Crypto::encrypt(const QString& passphrase,
                     const QStringList& inputFilePaths,
                     const QString& outputPath,
                     const QString& cipher,
                     const std::size_t& inputKeySize,
                     const QString& modeOfOperation,
                     const bool compress,
                     const bool container)
{
  m->state->busy(true);
  emit busyStatus(m->state->isBusy());

  auto keySize = std::size_t{128};

  if (inputKeySize > 0)
  {
    keySize = inputKeySize;
  }

  auto algorithm = QStringLiteral("AES-128/GCM");

  if (QStringLiteral("AES") == cipher)
  {
    algorithm = cipher % QStringLiteral("-") % QString::number(keySize) %
                QStringLiteral("/") % modeOfOperation;
  }
  else
  {
    algorithm = cipher % QStringLiteral("/") % modeOfOperation;
  }

  auto outputFilePaths = QStringList{};

  for (const auto& inputFilePath : inputFilePaths)
  {
    const QFileInfo inputFileInfo{inputFilePath};
    const auto inFileName = inputFileInfo.fileName();
    const auto inFilePath = inputFileInfo.absoluteFilePath();
    const auto outFilePath = outputFilePath(inFilePath,inFileName,
                                            outputPath);

    outputFilePaths << outFilePath;

    try
    {
      m->botanCrypto->encryptFile(passphrase, inFilePath, outFilePath,
                                  algorithm, keySize, compress);

      if (m->state->isAborted() || m->state->isStopped(inFilePath))
      {
        emit errorMessage(m->messages[2], inFilePath);

        if (m->state->isAborted())
        {
          m->state->abort(false);
          break;
        }
      }
    }
    catch (const Botan::Stream_IO_Error& e)
    {
      emit errorMessage(m->messages[7], inFilePath);
    }
    catch (const std::exception& e)
    {
      const auto error = QString{e.what()};
      emit errorMessage(QStringLiteral("Error: ") % error, inFilePath);
    }
  } // End file loop

  m->state->reset();

  m->state->busy(false);
  emit busyStatus(m->state->isBusy());
}

void Crypto::decrypt(const QString& passphrase,
                     const QStringList& inputFilePaths,
                     const QString& outputPath)
{
  m->state->busy(true);
  emit busyStatus(m->state->isBusy());

  for (const auto& inputFilePath : inputFilePaths)
  {
    const QFileInfo inputFileInfo{inputFilePath};
    const auto inFileName = inputFileInfo.fileName();
    const auto inFilePath = inputFileInfo.absoluteFilePath();
    const auto outFilePath = outputFilePath(inFilePath,inFileName, outputPath);

    try
    {
      m->botanCrypto->decryptFile(passphrase, inFilePath, outFilePath);

      if (m->state->isAborted())
      { // Reset abort flag
        m->state->abort(false);
        emit errorMessage(m->messages[3], inFilePath);
        break;
      }

      if (m->state->isStopped(inFilePath))
      {
        emit errorMessage(m->messages[3], inFilePath);
      }
    }
    catch (const Botan::Decoding_Error& e)
    {
      emit errorMessage(m->messages[5], inFilePath);
    }
    catch (const Botan::Integrity_Failure& e)
    {
      emit errorMessage(m->messages[5], inFilePath);
    }
    catch (const Botan::Invalid_Argument& e)
    {
      emit errorMessage(m->messages[6], inFilePath);
    }
    catch (const std::invalid_argument& e)
    {
      emit errorMessage(m->messages[6], inFilePath);
    }
    catch (const std::runtime_error& e)
    {
      emit errorMessage(m->messages[4], inFilePath);
    }
    catch (const std::exception& e)
    {
      const auto error = QString{e.what()};

      if (QStringLiteral("zlib inflate error -3") == error)
      {
        emit errorMessage(m->messages[5], inFilePath);
      }
      else
      {
        emit errorMessage(QStringLiteral("Error: ") % error, inFilePath);
      }
    }
  }

  m->state->reset();

  m->state->busy(false);
  emit busyStatus(m->state->isBusy());
}

void Crypto::abort()
{
  if (m->state->isBusy())
  {
    m->state->abort(true);
  }
}

void Crypto::pause(bool pause)
{
  m->state->pause(pause);
}

void Crypto::stop(const QString& filePath)
{
  if (m->state->isBusy())
  {
    m->state->stop(filePath, true);
  }
}

Crypto::CryptoPrivate::CryptoPrivate()
  : state{std::make_unique<CryptoState>()},
    botanCrypto{std::make_unique<BotanCrypto>(state.get())}
{}
<commit_msg>Use lambdas to increase const correctness<commit_after>#include "src/cryptography/Crypto.hpp"
#include "src/cryptography/CryptoState.hpp"
#include "src/cryptography/BotanCrypto.hpp"
#include "src/cryptography/constants.h"
#include "src/utility/pimpl_impl.h"
#include <QFileInfo>
#include <QDir>
#include <QStringRef>
#include <QStringBuilder>
#include <string>
#include <memory>

#if defined(_MSC_VER)
#include "src/utility/make_unique.h"
#endif

QString outputFilePath(const QString& inputFilePath,
                       const QString& inputFileName,
                       const QString& outputPath)
{
  auto path = QString{inputFilePath % QStringLiteral(".") % kExtension};

  if (!outputPath.isEmpty())
  {
    path = QDir::cleanPath(outputPath) % QDir::separator() % inputFileName %
           QStringLiteral(".") % kExtension;
  }

  return path;
}

class Crypto::CryptoPrivate {
 public:
  /*!
   * \brief CryptoPrivate Constructs the Crypto private implementation
   */
  CryptoPrivate();

  /*!
   * \brief errorMessage Returns error message
   * \return String containing error message
   */
  QString errorMessage(const int index) const;

  std::unique_ptr<CryptoState> state;
  std::unique_ptr<BotanCrypto> botanCrypto;

  // The list of status messages that can be displayed to the user
  const QStringList messages;
};

Crypto::Crypto(QObject* parent)
  : QObject{parent}
{
  // Subscribe to provider's signals
  connect(m->botanCrypto.get(), &BotanCrypto::progress,
          this, &Crypto::progress);
  connect(m->botanCrypto.get(), &BotanCrypto::statusMessage,
          this, &Crypto::statusMessage);
  connect(m->botanCrypto.get(), &BotanCrypto::errorMessage,
          this, &Crypto::errorMessage);
}

Crypto::~Crypto() {}

void Crypto::encrypt(const QString& passphrase,
                     const QStringList& inputFilePaths,
                     const QString& outputPath,
                     const QString& cipher,
                     const std::size_t& inputKeySize,
                     const QString& modeOfOperation,
                     const bool compress,
                     const bool container)
{
  m->state->busy(true);
  emit busyStatus(m->state->isBusy());

  const auto keySize = [&inputKeySize] {
    auto size = std::size_t{128};

    if (inputKeySize > 0) {
      size = inputKeySize;
    }

    return size;
  }();

  const auto algorithm = [&cipher, &keySize, &modeOfOperation] {
    auto algo = QStringLiteral("AES-128/GCM");

    if (QStringLiteral("AES") == cipher) {
      algo = cipher % QStringLiteral("-") % QString::number(keySize) %
             QStringLiteral("/") % modeOfOperation;
    }
    else {
      algo = cipher % QStringLiteral("/") % modeOfOperation;
    }

    return algo;
  }();

  auto outputFilePaths = QStringList{};

  for (const auto& inputFilePath : inputFilePaths)
  {
    const QFileInfo inputFileInfo{inputFilePath};
    const auto inFileName = inputFileInfo.fileName();
    const auto inFilePath = inputFileInfo.absoluteFilePath();
    const auto outFilePath = outputFilePath(inFilePath,inFileName,
                                            outputPath);

    outputFilePaths << outFilePath;

    try
    {
      m->botanCrypto->encryptFile(passphrase, inFilePath, outFilePath,
                                  algorithm, keySize, compress);

      if (m->state->isAborted() || m->state->isStopped(inFilePath))
      {
        emit errorMessage(m->messages[2], inFilePath);

        if (m->state->isAborted())
        {
          m->state->abort(false);
          break;
        }
      }
    }
    catch (const Botan::Stream_IO_Error& e)
    {
      emit errorMessage(m->messages[7], inFilePath);
    }
    catch (const std::exception& e)
    {
      const auto error = QString{e.what()};
      emit errorMessage(QStringLiteral("Error: ") % error, inFilePath);
    }
  } // End file loop

  m->state->reset();

  m->state->busy(false);
  emit busyStatus(m->state->isBusy());
}

void Crypto::decrypt(const QString& passphrase,
                     const QStringList& inputFilePaths,
                     const QString& outputPath)
{
  m->state->busy(true);
  emit busyStatus(m->state->isBusy());

  for (const auto& inputFilePath : inputFilePaths)
  {
    const QFileInfo inputFileInfo{inputFilePath};
    const auto inFileName = inputFileInfo.fileName();
    const auto inFilePath = inputFileInfo.absoluteFilePath();
    const auto outFilePath = outputFilePath(inFilePath,inFileName, outputPath);

    try
    {
      m->botanCrypto->decryptFile(passphrase, inFilePath, outFilePath);

      if (m->state->isAborted())
      { // Reset abort flag
        m->state->abort(false);
        emit errorMessage(m->messages[3], inFilePath);
        break;
      }

      if (m->state->isStopped(inFilePath))
      {
        emit errorMessage(m->messages[3], inFilePath);
      }
    }
    catch (const Botan::Decoding_Error& e)
    {
      emit errorMessage(m->messages[5], inFilePath);
    }
    catch (const Botan::Integrity_Failure& e)
    {
      emit errorMessage(m->messages[5], inFilePath);
    }
    catch (const Botan::Invalid_Argument& e)
    {
      emit errorMessage(m->messages[6], inFilePath);
    }
    catch (const std::invalid_argument& e)
    {
      emit errorMessage(m->messages[6], inFilePath);
    }
    catch (const std::runtime_error& e)
    {
      emit errorMessage(m->messages[4], inFilePath);
    }
    catch (const std::exception& e)
    {
      const auto error = QString{e.what()};

      if (QStringLiteral("zlib inflate error -3") == error)
      {
        emit errorMessage(m->messages[5], inFilePath);
      }
      else
      {
        emit errorMessage(QStringLiteral("Error: ") % error, inFilePath);
      }
    }
  }

  m->state->reset();

  m->state->busy(false);
  emit busyStatus(m->state->isBusy());
}

void Crypto::abort()
{
  if (m->state->isBusy())
  {
    m->state->abort(true);
  }
}

void Crypto::pause(bool pause)
{
  m->state->pause(pause);
}

void Crypto::stop(const QString& filePath)
{
  if (m->state->isBusy())
  {
    m->state->stop(filePath, true);
  }
}

Crypto::CryptoPrivate::CryptoPrivate()
  : state{std::make_unique<CryptoState>()},
    botanCrypto{std::make_unique<BotanCrypto>(state.get())}
{}
<|endoftext|>


<commit_before>#include "cvm.h"

#include "defaults_shared.h"

DabClass &DabVM::define_builtin_class(const std::string &name, dab_class_t class_index,
                                      dab_class_t superclass_index)
{
    DabClass klass;
    klass.index            = class_index;
    klass.name             = name;
    klass.builtin          = true;
    klass.superclass_index = superclass_index;
    classes[class_index]   = klass;
    return classes[class_index];
}

void DabVM::predefine_default_classes()
{
    fprintf(stderr, "vm: predefine default classes\n");

    define_builtin_class("Object", CLASS_OBJECT);
    define_builtin_class("String", CLASS_STRING);
    define_builtin_class("Fixnum", CLASS_FIXNUM);
    define_builtin_class("Boolean", CLASS_BOOLEAN);
    define_builtin_class("NilClass", CLASS_NILCLASS);
    define_builtin_class("Array", CLASS_ARRAY);
    define_builtin_class("Uint8", CLASS_UINT8, CLASS_FIXNUM);
    define_builtin_class("Uint16", CLASS_UINT16, CLASS_FIXNUM);
    define_builtin_class("Uint32", CLASS_UINT32, CLASS_FIXNUM);
    define_builtin_class("Uint64", CLASS_UINT64, CLASS_FIXNUM);
    define_builtin_class("Int8", CLASS_INT8, CLASS_FIXNUM);
    define_builtin_class("Int16", CLASS_INT16, CLASS_FIXNUM);
    define_builtin_class("Int32", CLASS_INT32, CLASS_FIXNUM);
    define_builtin_class("Int64", CLASS_INT64, CLASS_FIXNUM);
    define_builtin_class("Method", CLASS_METHOD);
    define_builtin_class("IntPtr", CLASS_INTPTR);
    define_builtin_class("ByteBuffer", CLASS_BYTEBUFFER);
    define_builtin_class("LiteralString", CLASS_LITERALSTRING, CLASS_STRING);
    define_builtin_class("DynamicString", CLASS_DYNAMICSTRING, CLASS_STRING);
}

void DabVM::define_default_classes()
{
    fprintf(stderr, "vm: define default classes\n");

    auto &object_class = get_class(CLASS_OBJECT);
    object_class.add_static_reg_function("new", [this](DabValue self, std::vector<DabValue>) {
        auto arg = self;
        assert(arg.data.type == TYPE_CLASS);
        auto instance = arg.create_instance();
        cinstcall(instance, "__construct");
        return instance;
    });
    object_class.add_reg_function(
        "class", [](DabValue self, std::vector<DabValue>) { return self.get_class(); });
    object_class.add_reg_function("to_s", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return self.print_value();
    });
    object_class.add_reg_function("__construct", [](DabValue, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return nullptr;
    });
    object_class.add_reg_function("__destruct", [](DabValue, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return nullptr;
    });
    object_class.add_reg_function("is", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = args[0];
        return arg0.is_a(arg1.get_class());
    });
    object_class.add_static_reg_function("to_s", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return self.class_name();
    });
    object_class.add_static_reg_function("==", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = args[0];
        assert(arg0.data.type == TYPE_CLASS);
        assert(arg1.data.type == TYPE_CLASS);
        return DabValue(arg0.data.fixnum == arg1.data.fixnum);
    });

    auto &string_class = get_class(CLASS_STRING);
    string_class.add_reg_function("upcase", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        auto s = self.string();
        std::transform(s.begin(), s.end(), s.begin(), ::toupper);
        return s;
    });
    string_class.add_reg_function("length", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        auto     s = self.string();
        DabValue ret(CLASS_INT32, (int)s.size());
        return ret;
    });
    string_class.add_reg_function("[]", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg = $VM->cast(args[0], CLASS_FIXNUM);
        auto s   = self.string();
        auto n   = arg.data.fixnum;
        return DabValue(s.substr(n, 1));
    });
    string_class.add_reg_function("+", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = $VM->cast(args[0], CLASS_STRING);
        auto s0   = arg0.string();
        auto s1   = arg1.string();
        return DabValue(s0 + s1);
    });
    string_class.add_static_reg_function("new", [](DabValue, std::vector<DabValue> args) {
        auto argc = args.size();
        assert(argc <= 2);
        std::string s;
        if (argc == 2)
        {
            auto        arg1   = $VM->cast(args[0], CLASS_BYTEBUFFER);
            auto        arg2   = $VM->cast(args[1], CLASS_FIXNUM);
            const auto &buffer = arg1.bytebuffer();
            const auto  length = arg2.data.fixnum;
            if (length)
            {
                s = std::string((const char *)&buffer[0], length);
            }
        }
        else if (argc == 1)
        {
            s = args[0].string();
        }
        return DabValue::allocate_dynstr(s.c_str());
    });
    string_class.add_reg_function("to_s", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        if (!$VM->options.autorelease)
        {
            self.retain();
        }
        return self;
    });
    DAB_MEMBER_EQUALS_OPERATORS(string_class, CLASS_STRING, .string());
    DAB_MEMBER_COMPARE_OPERATORS(string_class, CLASS_STRING, .string());

    auto &fixnum_class = get_class(CLASS_FIXNUM);
    fixnum_class.add_static_reg_function("new", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() <= 1);
        auto     klass = self;
        DabValue ret_value;
        ret_value.data.type   = TYPE_FIXNUM;
        ret_value.data.fixnum = 0;
        if (args.size() == 1)
        {
            auto arg = args[0];
            assert(arg.data.type == TYPE_FIXNUM);
            ret_value.data.fixnum = arg.data.fixnum;
        }
        return ret_value;
    });
    DAB_MEMBER_NUMERIC_OPERATORS(fixnum_class, CLASS_FIXNUM, uint64_t, .data.fixnum);

    define_default_classes_int();

    auto &intptr_class = get_class(CLASS_INTPTR);
    intptr_class.add_reg_function("fetch_int32", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        auto     ptr   = self.data.intptr;
        auto     iptr  = (int32_t *)ptr;
        auto     value = *iptr;
        DabValue ret(CLASS_INT32, value);
        return ret;
    });

    auto &boolean_class = get_class(CLASS_BOOLEAN);
    DAB_MEMBER_EQUALS_OPERATORS(boolean_class, CLASS_BOOLEAN, .data.boolean);

    auto &array_class = get_class(CLASS_ARRAY);
    array_class.add_reg_function("count", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        assert(self.data.type == TYPE_ARRAY);
        return (uint64_t)self.array().size();
    });
    array_class.add_reg_function("shift", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        assert(self.data.type == TYPE_ARRAY);
        auto &a = self.array();
        if (a.size() == 0)
        {
            return DabValue(nullptr);
        }
        else
        {
            auto ret = a[0];
            a.erase(a.begin());
            return ret;
        }
    });
    array_class.add_reg_function("insert", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg = args[0];
        assert(self.data.type == TYPE_ARRAY);
        auto &a = self.array();
        a.push_back(arg);
        return nullptr;
    });
    array_class.add_reg_function("[]", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = args[0];
        assert(arg0.data.type == TYPE_ARRAY);
        assert(arg1.data.type == TYPE_FIXNUM);
        auto &a = arg0.array();
        auto  n = arg1.data.fixnum;
        if (n < 0)
            n = a.size() + n;
        if (n < 0 || n >= (int64_t)a.size())
            return DabValue(nullptr);
        else
            return a[n];
    });
    array_class.add_reg_function("[]=", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 2);
        auto arg0 = self;
        auto arg1 = args[0];
        auto arg2 = args[1];
        assert(arg0.data.type == TYPE_ARRAY);
        assert(arg1.data.type == TYPE_FIXNUM);
        auto &a = arg0.array();
        auto  n = arg1.data.fixnum;
        if (n < 0)
            n = a.size() + n;
        if (n < 0 || n >= (int64_t)a.size())
        {
            fprintf(stderr, "vm: index outside of array bounds.\n");
            exit(1);
        }
        else
        {
            a[n] = arg2;
        }
        return nullptr;
    });
    array_class.add_reg_function("join", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        std::string ret;
        auto &      a = self.array();
        for (size_t i = 0; i < a.size(); i++)
        {
            if (i)
                ret += ", ";
            auto val = cinstcall(a[i], "to_s");
            ret += val.string();
            if (!$VM->options.autorelease)
            {
                val.release();
            }
        }
        return ret;
    });
    array_class.add_reg_function("to_s", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        auto inner = cinstcall(self, "join");
        auto ret   = std::string("[" + inner.string() + "]");
        if (!$VM->options.autorelease)
        {
            inner.release();
        }
        return ret;
    });
    array_class.add_reg_function("+", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg = $VM->cast(args[0], CLASS_ARRAY);
        return $VM->merge_arrays(self, arg);
    });

    auto &method_class = get_class(CLASS_METHOD);
    method_class.add_reg_function("to_s", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return std::string("@method(" + self.string() + ")");
    });

    auto &bytebuffer_class = get_class(CLASS_BYTEBUFFER);
    bytebuffer_class.add_static_reg_function("new", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto klass = self;
        assert(klass.data.type == TYPE_CLASS);

        auto instance = klass.create_instance();

        auto _size = $VM->cast(args[0], CLASS_FIXNUM);
        auto size  = _size.data.fixnum;
        instance.bytebuffer().resize(size);
        return instance;
    });
    bytebuffer_class.add_reg_function("[]", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = args[0];
        assert(arg0.data.type == TYPE_BYTEBUFFER);
        assert(arg1.data.type == TYPE_FIXNUM);
        auto &a = arg0.bytebuffer();
        auto  n = arg1.data.fixnum;
        if (n < 0)
            n = a.size() + n;
        if (n < 0 || n >= (int64_t)a.size())
            return DabValue(nullptr);
        else
            return DabValue(CLASS_UINT8, a[n]);
    });
    bytebuffer_class.add_reg_function("[]=", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 2);
        auto arg0 = self;
        auto arg1 = args[0];
        auto arg2 = args[1];
        assert(arg0.data.type == TYPE_BYTEBUFFER);
        assert(arg1.data.type == TYPE_FIXNUM);
        auto &a = arg0.bytebuffer();
        auto  n = arg1.data.fixnum;
        if (n < 0)
            n = a.size() + n;
        if (n >= 0 && n < (int64_t)a.size())
        {
            a[n] = $VM->cast(arg2, CLASS_UINT8).data.num_uint8;
        }
        return nullptr;
    });
    bytebuffer_class.add_reg_function("to_s", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        std::string ret;
        auto &      a = self.bytebuffer();
        for (size_t i = 0; i < a.size(); i++)
        {
            if (i)
                ret += ", ";
            char string[32];
            sprintf(string, "%d", (int)a[i]);
            ret += string;
        }
        return "[" + ret + "]";
    });
}
<commit_msg>vm: Array.[]= throws instead of stderr/exit<commit_after>#include "cvm.h"

#include "defaults_shared.h"

DabClass &DabVM::define_builtin_class(const std::string &name, dab_class_t class_index,
                                      dab_class_t superclass_index)
{
    DabClass klass;
    klass.index            = class_index;
    klass.name             = name;
    klass.builtin          = true;
    klass.superclass_index = superclass_index;
    classes[class_index]   = klass;
    return classes[class_index];
}

void DabVM::predefine_default_classes()
{
    fprintf(stderr, "vm: predefine default classes\n");

    define_builtin_class("Object", CLASS_OBJECT);
    define_builtin_class("String", CLASS_STRING);
    define_builtin_class("Fixnum", CLASS_FIXNUM);
    define_builtin_class("Boolean", CLASS_BOOLEAN);
    define_builtin_class("NilClass", CLASS_NILCLASS);
    define_builtin_class("Array", CLASS_ARRAY);
    define_builtin_class("Uint8", CLASS_UINT8, CLASS_FIXNUM);
    define_builtin_class("Uint16", CLASS_UINT16, CLASS_FIXNUM);
    define_builtin_class("Uint32", CLASS_UINT32, CLASS_FIXNUM);
    define_builtin_class("Uint64", CLASS_UINT64, CLASS_FIXNUM);
    define_builtin_class("Int8", CLASS_INT8, CLASS_FIXNUM);
    define_builtin_class("Int16", CLASS_INT16, CLASS_FIXNUM);
    define_builtin_class("Int32", CLASS_INT32, CLASS_FIXNUM);
    define_builtin_class("Int64", CLASS_INT64, CLASS_FIXNUM);
    define_builtin_class("Method", CLASS_METHOD);
    define_builtin_class("IntPtr", CLASS_INTPTR);
    define_builtin_class("ByteBuffer", CLASS_BYTEBUFFER);
    define_builtin_class("LiteralString", CLASS_LITERALSTRING, CLASS_STRING);
    define_builtin_class("DynamicString", CLASS_DYNAMICSTRING, CLASS_STRING);
}

void DabVM::define_default_classes()
{
    fprintf(stderr, "vm: define default classes\n");

    auto &object_class = get_class(CLASS_OBJECT);
    object_class.add_static_reg_function("new", [this](DabValue self, std::vector<DabValue>) {
        auto arg = self;
        assert(arg.data.type == TYPE_CLASS);
        auto instance = arg.create_instance();
        cinstcall(instance, "__construct");
        return instance;
    });
    object_class.add_reg_function(
        "class", [](DabValue self, std::vector<DabValue>) { return self.get_class(); });
    object_class.add_reg_function("to_s", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return self.print_value();
    });
    object_class.add_reg_function("__construct", [](DabValue, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return nullptr;
    });
    object_class.add_reg_function("__destruct", [](DabValue, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return nullptr;
    });
    object_class.add_reg_function("is", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = args[0];
        return arg0.is_a(arg1.get_class());
    });
    object_class.add_static_reg_function("to_s", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return self.class_name();
    });
    object_class.add_static_reg_function("==", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = args[0];
        assert(arg0.data.type == TYPE_CLASS);
        assert(arg1.data.type == TYPE_CLASS);
        return DabValue(arg0.data.fixnum == arg1.data.fixnum);
    });

    auto &string_class = get_class(CLASS_STRING);
    string_class.add_reg_function("upcase", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        auto s = self.string();
        std::transform(s.begin(), s.end(), s.begin(), ::toupper);
        return s;
    });
    string_class.add_reg_function("length", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        auto     s = self.string();
        DabValue ret(CLASS_INT32, (int)s.size());
        return ret;
    });
    string_class.add_reg_function("[]", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg = $VM->cast(args[0], CLASS_FIXNUM);
        auto s   = self.string();
        auto n   = arg.data.fixnum;
        return DabValue(s.substr(n, 1));
    });
    string_class.add_reg_function("+", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = $VM->cast(args[0], CLASS_STRING);
        auto s0   = arg0.string();
        auto s1   = arg1.string();
        return DabValue(s0 + s1);
    });
    string_class.add_static_reg_function("new", [](DabValue, std::vector<DabValue> args) {
        auto argc = args.size();
        assert(argc <= 2);
        std::string s;
        if (argc == 2)
        {
            auto        arg1   = $VM->cast(args[0], CLASS_BYTEBUFFER);
            auto        arg2   = $VM->cast(args[1], CLASS_FIXNUM);
            const auto &buffer = arg1.bytebuffer();
            const auto  length = arg2.data.fixnum;
            if (length)
            {
                s = std::string((const char *)&buffer[0], length);
            }
        }
        else if (argc == 1)
        {
            s = args[0].string();
        }
        return DabValue::allocate_dynstr(s.c_str());
    });
    string_class.add_reg_function("to_s", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        if (!$VM->options.autorelease)
        {
            self.retain();
        }
        return self;
    });
    DAB_MEMBER_EQUALS_OPERATORS(string_class, CLASS_STRING, .string());
    DAB_MEMBER_COMPARE_OPERATORS(string_class, CLASS_STRING, .string());

    auto &fixnum_class = get_class(CLASS_FIXNUM);
    fixnum_class.add_static_reg_function("new", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() <= 1);
        auto     klass = self;
        DabValue ret_value;
        ret_value.data.type   = TYPE_FIXNUM;
        ret_value.data.fixnum = 0;
        if (args.size() == 1)
        {
            auto arg = args[0];
            assert(arg.data.type == TYPE_FIXNUM);
            ret_value.data.fixnum = arg.data.fixnum;
        }
        return ret_value;
    });
    DAB_MEMBER_NUMERIC_OPERATORS(fixnum_class, CLASS_FIXNUM, uint64_t, .data.fixnum);

    define_default_classes_int();

    auto &intptr_class = get_class(CLASS_INTPTR);
    intptr_class.add_reg_function("fetch_int32", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        auto     ptr   = self.data.intptr;
        auto     iptr  = (int32_t *)ptr;
        auto     value = *iptr;
        DabValue ret(CLASS_INT32, value);
        return ret;
    });

    auto &boolean_class = get_class(CLASS_BOOLEAN);
    DAB_MEMBER_EQUALS_OPERATORS(boolean_class, CLASS_BOOLEAN, .data.boolean);

    auto &array_class = get_class(CLASS_ARRAY);
    array_class.add_reg_function("count", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        assert(self.data.type == TYPE_ARRAY);
        return (uint64_t)self.array().size();
    });
    array_class.add_reg_function("shift", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        assert(self.data.type == TYPE_ARRAY);
        auto &a = self.array();
        if (a.size() == 0)
        {
            return DabValue(nullptr);
        }
        else
        {
            auto ret = a[0];
            a.erase(a.begin());
            return ret;
        }
    });
    array_class.add_reg_function("insert", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg = args[0];
        assert(self.data.type == TYPE_ARRAY);
        auto &a = self.array();
        a.push_back(arg);
        return nullptr;
    });
    array_class.add_reg_function("[]", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = args[0];
        assert(arg0.data.type == TYPE_ARRAY);
        assert(arg1.data.type == TYPE_FIXNUM);
        auto &a = arg0.array();
        auto  n = arg1.data.fixnum;
        if (n < 0)
            n = a.size() + n;
        if (n < 0 || n >= (int64_t)a.size())
            return DabValue(nullptr);
        else
            return a[n];
    });
    array_class.add_reg_function("[]=", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 2);
        auto arg0 = self;
        auto arg1 = args[0];
        auto arg2 = args[1];
        assert(arg0.data.type == TYPE_ARRAY);
        assert(arg1.data.type == TYPE_FIXNUM);
        auto &a = arg0.array();
        auto  n = arg1.data.fixnum;
        if (n < 0)
            n = a.size() + n;
        if (n < 0 || n >= (int64_t)a.size())
        {
            throw DabRuntimeError("vm: index outside of array bounds.");
        }
        else
        {
            a[n] = arg2;
        }
        return nullptr;
    });
    array_class.add_reg_function("join", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        std::string ret;
        auto &      a = self.array();
        for (size_t i = 0; i < a.size(); i++)
        {
            if (i)
                ret += ", ";
            auto val = cinstcall(a[i], "to_s");
            ret += val.string();
            if (!$VM->options.autorelease)
            {
                val.release();
            }
        }
        return ret;
    });
    array_class.add_reg_function("to_s", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        auto inner = cinstcall(self, "join");
        auto ret   = std::string("[" + inner.string() + "]");
        if (!$VM->options.autorelease)
        {
            inner.release();
        }
        return ret;
    });
    array_class.add_reg_function("+", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg = $VM->cast(args[0], CLASS_ARRAY);
        return $VM->merge_arrays(self, arg);
    });

    auto &method_class = get_class(CLASS_METHOD);
    method_class.add_reg_function("to_s", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        return std::string("@method(" + self.string() + ")");
    });

    auto &bytebuffer_class = get_class(CLASS_BYTEBUFFER);
    bytebuffer_class.add_static_reg_function("new", [](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto klass = self;
        assert(klass.data.type == TYPE_CLASS);

        auto instance = klass.create_instance();

        auto _size = $VM->cast(args[0], CLASS_FIXNUM);
        auto size  = _size.data.fixnum;
        instance.bytebuffer().resize(size);
        return instance;
    });
    bytebuffer_class.add_reg_function("[]", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 1);
        auto arg0 = self;
        auto arg1 = args[0];
        assert(arg0.data.type == TYPE_BYTEBUFFER);
        assert(arg1.data.type == TYPE_FIXNUM);
        auto &a = arg0.bytebuffer();
        auto  n = arg1.data.fixnum;
        if (n < 0)
            n = a.size() + n;
        if (n < 0 || n >= (int64_t)a.size())
            return DabValue(nullptr);
        else
            return DabValue(CLASS_UINT8, a[n]);
    });
    bytebuffer_class.add_reg_function("[]=", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 2);
        auto arg0 = self;
        auto arg1 = args[0];
        auto arg2 = args[1];
        assert(arg0.data.type == TYPE_BYTEBUFFER);
        assert(arg1.data.type == TYPE_FIXNUM);
        auto &a = arg0.bytebuffer();
        auto  n = arg1.data.fixnum;
        if (n < 0)
            n = a.size() + n;
        if (n >= 0 && n < (int64_t)a.size())
        {
            a[n] = $VM->cast(arg2, CLASS_UINT8).data.num_uint8;
        }
        return nullptr;
    });
    bytebuffer_class.add_reg_function("to_s", [this](DabValue self, std::vector<DabValue> args) {
        assert(args.size() == 0);
        std::string ret;
        auto &      a = self.bytebuffer();
        for (size_t i = 0; i < a.size(); i++)
        {
            if (i)
                ret += ", ";
            char string[32];
            sprintf(string, "%d", (int)a[i]);
            ret += string;
        }
        return "[" + ret + "]";
    });
}
<|endoftext|>


<commit_before>// Copyright (c) 2012 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "printing/backend/print_backend.h"

#include <objidl.h>
#include <winspool.h>

#include "base/memory/scoped_ptr.h"
#include "base/strings/string_piece.h"
#include "base/strings/utf_string_conversions.h"
#include "base/win/scoped_bstr.h"
#include "base/win/scoped_comptr.h"
#include "base/win/scoped_hglobal.h"
#include "printing/backend/print_backend_consts.h"
#include "printing/backend/printing_info_win.h"
#include "printing/backend/win_helper.h"

namespace printing {

namespace {

HRESULT StreamOnHGlobalToString(IStream* stream, std::string* out) {
  DCHECK(stream);
  DCHECK(out);
  HGLOBAL hdata = NULL;
  HRESULT hr = GetHGlobalFromStream(stream, &hdata);
  if (SUCCEEDED(hr)) {
    DCHECK(hdata);
    base::win::ScopedHGlobal<char> locked_data(hdata);
    out->assign(locked_data.release(), locked_data.Size());
  }
  return hr;
}

template <class T>
void GetDeviceCapabilityArray(const wchar_t* printer,
                              const wchar_t* port,
                              WORD id,
                              std::vector<T>* result) {
  int count = DeviceCapabilities(printer, port, id, NULL, NULL);
  if (count <= 0)
    return;
  result->resize(count);
  CHECK_EQ(count,
           DeviceCapabilities(printer, port, id,
                              reinterpret_cast<LPTSTR>(result->data()), NULL));
}

void LoadPaper(const wchar_t* printer,
               const wchar_t* port,
               const DEVMODE* devmode,
               PrinterSemanticCapsAndDefaults* caps) {
  static const size_t kToUm = 100;  // Windows uses 0.1mm.
  static const size_t kMaxPaperName = 64;

  struct PaperName {
    wchar_t chars[kMaxPaperName];
  };

  DCHECK_EQ(sizeof(PaperName), sizeof(wchar_t) * kMaxPaperName);

  // Paper
  std::vector<PaperName> names;
  GetDeviceCapabilityArray(printer, port, DC_PAPERNAMES, &names);

  std::vector<POINT> sizes;
  GetDeviceCapabilityArray(printer, port, DC_PAPERSIZE, &sizes);

  std::vector<WORD> ids;
  GetDeviceCapabilityArray(printer, port, DC_PAPERS, &ids);

  DCHECK_EQ(ids.size(), sizes.size());
  DCHECK_EQ(names.size(), sizes.size());

  if (ids.size() != sizes.size())
    ids.clear();
  if (names.size() != sizes.size())
    names.clear();

  for (size_t i = 0; i < sizes.size(); ++i) {
    PrinterSemanticCapsAndDefaults::Paper paper;
    paper.size_um.SetSize(sizes[i].x * kToUm, sizes[i].y * kToUm);
    if (!names.empty()) {
      const wchar_t* name_start = names[i].chars;
      base::string16 tmp_name(name_start, kMaxPaperName);
      // Trim trailing zeros.
      tmp_name = tmp_name.c_str();
      paper.name = base::WideToUTF8(tmp_name);
    }
    caps->papers.push_back(paper);
  }

  if (devmode) {
    short default_id = 0;
    gfx::Size default_size;

    if (devmode->dmFields & DM_PAPERSIZE)
      default_id = devmode->dmPaperSize;
    if (devmode->dmFields & DM_PAPERWIDTH)
      default_size.set_width(devmode->dmPaperWidth * kToUm);
    if (devmode->dmFields & DM_PAPERLENGTH)
      default_size.set_height(devmode->dmPaperLength * kToUm);

    if (default_size.IsEmpty()) {
      for (size_t i = 0; i < ids.size(); ++i) {
        if (ids[i] == default_id) {
          PrinterSemanticCapsAndDefaults::Paper paper;
          paper.size_um.SetSize(sizes[i].x * kToUm, sizes[i].y * kToUm);
          if (!names.empty()) {
            const wchar_t* name_start = names[i].chars;
            base::string16 tmp_name(name_start, kMaxPaperName);
            // Trim trailing zeros.
            tmp_name = tmp_name.c_str();
            paper.name = base::WideToUTF8(tmp_name);
          }
          caps->default_paper = paper;
          break;
        }
      }
    } else {
      caps->default_paper.size_um = default_size;
    }
  }
}

void LoadDpi(const wchar_t* printer,
             const wchar_t* port,
             const DEVMODE* devmode,
             PrinterSemanticCapsAndDefaults* caps) {
  std::vector<POINT> dpis;
  GetDeviceCapabilityArray(printer, port, DC_ENUMRESOLUTIONS, &dpis);

  for (size_t i = 0; i < dpis.size() ; ++i)
    caps->dpis.push_back(gfx::Size(dpis[i].x, dpis[i].y));

  if (devmode) {
    if ((devmode->dmFields & DM_PRINTQUALITY) && devmode->dmPrintQuality > 0) {
      caps->default_dpi.SetSize(devmode->dmPrintQuality,
                                devmode->dmPrintQuality);
      if (devmode->dmFields & DM_YRESOLUTION) {
        caps->default_dpi.set_height(devmode->dmYResolution);
      }
    }
  }
}

}  // namespace

class PrintBackendWin : public PrintBackend {
 public:
  PrintBackendWin() {}

  // PrintBackend implementation.
  virtual bool EnumeratePrinters(PrinterList* printer_list) OVERRIDE;
  virtual std::string GetDefaultPrinterName() OVERRIDE;
  virtual bool GetPrinterSemanticCapsAndDefaults(
      const std::string& printer_name,
      PrinterSemanticCapsAndDefaults* printer_info) OVERRIDE;
  virtual bool GetPrinterCapsAndDefaults(
      const std::string& printer_name,
      PrinterCapsAndDefaults* printer_info) OVERRIDE;
  virtual std::string GetPrinterDriverInfo(
      const std::string& printer_name) OVERRIDE;
  virtual bool IsValidPrinter(const std::string& printer_name) OVERRIDE;

 protected:
  virtual ~PrintBackendWin() {}
};

bool PrintBackendWin::EnumeratePrinters(PrinterList* printer_list) {
  DCHECK(printer_list);
  DWORD bytes_needed = 0;
  DWORD count_returned = 0;
  const DWORD kLevel = 4;
  BOOL ret = EnumPrinters(PRINTER_ENUM_LOCAL|PRINTER_ENUM_CONNECTIONS, NULL,
                          kLevel, NULL, 0, &bytes_needed, &count_returned);
  if (!bytes_needed)
    return false;
  scoped_ptr<BYTE[]> printer_info_buffer(new BYTE[bytes_needed]);
  ret = EnumPrinters(PRINTER_ENUM_LOCAL|PRINTER_ENUM_CONNECTIONS, NULL, kLevel,
                     printer_info_buffer.get(), bytes_needed, &bytes_needed,
                     &count_returned);
  DCHECK(ret);
  if (!ret)
    return false;

  std::string default_printer = GetDefaultPrinterName();
  PRINTER_INFO_4* printer_info =
      reinterpret_cast<PRINTER_INFO_4*>(printer_info_buffer.get());
  for (DWORD index = 0; index < count_returned; index++) {
    ScopedPrinterHandle printer;
    PrinterBasicInfo info;
    if (printer.OpenPrinter(printer_info[index].pPrinterName) &&
        InitBasicPrinterInfo(printer, &info)) {
      info.is_default = (info.printer_name == default_printer);
      printer_list->push_back(info);
    }
  }
  return true;
}

std::string PrintBackendWin::GetDefaultPrinterName() {
  DWORD size = MAX_PATH;
  TCHAR default_printer_name[MAX_PATH];
  if (!::GetDefaultPrinter(default_printer_name, &size))
    return std::string();
  return base::WideToUTF8(default_printer_name);
}

bool PrintBackendWin::GetPrinterSemanticCapsAndDefaults(
    const std::string& printer_name,
    PrinterSemanticCapsAndDefaults* printer_info) {
  ScopedPrinterHandle printer_handle;
  if (!printer_handle.OpenPrinter(base::UTF8ToWide(printer_name).c_str())) {
    LOG(WARNING) << "Failed to open printer, error = " << GetLastError();
    return false;
  }

  PrinterInfo5 info_5;
  if (!info_5.Init(printer_handle))
    return false;
  const wchar_t* name = info_5.get()->pPrinterName;
  const wchar_t* port = info_5.get()->pPortName;
  DCHECK_EQ(name, base::UTF8ToUTF16(printer_name));

  PrinterSemanticCapsAndDefaults caps;
  UserDefaultDevMode user_settings;
  if (user_settings.Init(printer_handle)) {
    if (user_settings.get()->dmFields & DM_COLOR)
      caps.color_default = (user_settings.get()->dmColor == DMCOLOR_COLOR);

    if (user_settings.get()->dmFields & DM_DUPLEX) {
      switch (user_settings.get()->dmDuplex) {
      case DMDUP_SIMPLEX:
        caps.duplex_default = SIMPLEX;
        break;
      case DMDUP_VERTICAL:
        caps.duplex_default = LONG_EDGE;
        break;
      case DMDUP_HORIZONTAL:
        caps.duplex_default = SHORT_EDGE;
        break;
      default:
        NOTREACHED();
      }
    }

    if (user_settings.get()->dmFields & DM_COLLATE)
      caps.collate_default = (user_settings.get()->dmCollate == DMCOLLATE_TRUE);
  } else {
    LOG(WARNING) << "Fallback to color/simplex mode.";
    caps.color_default = caps.color_changeable;
    caps.duplex_default = SIMPLEX;
  }

  // Get printer capabilities. For more info see here:
  // http://msdn.microsoft.com/en-us/library/windows/desktop/dd183552(v=vs.85).aspx
  caps.color_changeable =
      (DeviceCapabilities(name, port, DC_COLORDEVICE, NULL, NULL) == 1);

  caps.duplex_capable =
      (DeviceCapabilities(name, port, DC_DUPLEX, NULL, NULL) == 1);

  caps.collate_capable =
      (DeviceCapabilities(name, port, DC_COLLATE, NULL, NULL) == 1);

  caps.copies_capable =
      (DeviceCapabilities(name, port, DC_COPIES, NULL, NULL) > 1);

  LoadPaper(name, port, user_settings.get(), &caps);
  LoadDpi(name, port, user_settings.get(), &caps);

  *printer_info = caps;
  return true;
}

bool PrintBackendWin::GetPrinterCapsAndDefaults(
    const std::string& printer_name,
    PrinterCapsAndDefaults* printer_info) {
  ScopedXPSInitializer xps_initializer;
  if (!xps_initializer.initialized()) {
    // TODO(sanjeevr): Handle legacy proxy case (with no prntvpt.dll)
    return false;
  }
  if (!IsValidPrinter(printer_name)) {
    return false;
  }
  DCHECK(printer_info);
  HPTPROVIDER provider = NULL;
  std::wstring printer_name_wide = base::UTF8ToWide(printer_name);
  HRESULT hr = XPSModule::OpenProvider(printer_name_wide, 1, &provider);
  if (provider) {
    base::win::ScopedComPtr<IStream> print_capabilities_stream;
    hr = CreateStreamOnHGlobal(NULL, TRUE,
                               print_capabilities_stream.Receive());
    DCHECK(SUCCEEDED(hr));
    if (print_capabilities_stream) {
      base::win::ScopedBstr error;
      hr = XPSModule::GetPrintCapabilities(provider,
                                           NULL,
                                           print_capabilities_stream,
                                           error.Receive());
      DCHECK(SUCCEEDED(hr));
      if (FAILED(hr)) {
        return false;
      }
      hr = StreamOnHGlobalToString(print_capabilities_stream.get(),
                                   &printer_info->printer_capabilities);
      DCHECK(SUCCEEDED(hr));
      printer_info->caps_mime_type = "text/xml";
    }
    ScopedPrinterHandle printer_handle;
    if (printer_handle.OpenPrinter(printer_name_wide.c_str())) {
      scoped_ptr<DEVMODE[]> devmode_out(CreateDevMode(printer_handle, NULL));
      if (!devmode_out)
        return false;
      base::win::ScopedComPtr<IStream> printer_defaults_stream;
      hr = CreateStreamOnHGlobal(NULL, TRUE,
                                 printer_defaults_stream.Receive());
      DCHECK(SUCCEEDED(hr));
      if (printer_defaults_stream) {
        DWORD dm_size = devmode_out.get()->dmSize +
                        devmode_out.get()->dmDriverExtra;
        hr = XPSModule::ConvertDevModeToPrintTicket(provider, dm_size,
            devmode_out.get(), kPTJobScope, printer_defaults_stream);
        DCHECK(SUCCEEDED(hr));
        if (SUCCEEDED(hr)) {
          hr = StreamOnHGlobalToString(printer_defaults_stream.get(),
                                       &printer_info->printer_defaults);
          DCHECK(SUCCEEDED(hr));
          printer_info->defaults_mime_type = "text/xml";
        }
      }
    }
    XPSModule::CloseProvider(provider);
  }
  return true;
}

// Gets the information about driver for a specific printer.
std::string PrintBackendWin::GetPrinterDriverInfo(
    const std::string& printer_name) {
  ScopedPrinterHandle printer;
  if (!printer.OpenPrinter(base::UTF8ToWide(printer_name).c_str())) {
    return std::string();
  }
  return GetDriverInfo(printer);
}

bool PrintBackendWin::IsValidPrinter(const std::string& printer_name) {
  ScopedPrinterHandle printer_handle;
  return printer_handle.OpenPrinter(base::UTF8ToWide(printer_name).c_str());
}

scoped_refptr<PrintBackend> PrintBackend::CreateInstance(
    const base::DictionaryValue* print_backend_settings) {
  return new PrintBackendWin;
}

}  // namespace printing
<commit_msg>Less strict CHECK for DeviceCapabilities results. We should crash only if returned result is bigger that buffer we provided.<commit_after>// Copyright (c) 2012 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "printing/backend/print_backend.h"

#include <objidl.h>
#include <winspool.h>

#include "base/memory/scoped_ptr.h"
#include "base/numerics/safe_conversions.h"
#include "base/strings/string_piece.h"
#include "base/strings/utf_string_conversions.h"
#include "base/win/scoped_bstr.h"
#include "base/win/scoped_comptr.h"
#include "base/win/scoped_hglobal.h"
#include "printing/backend/print_backend_consts.h"
#include "printing/backend/printing_info_win.h"
#include "printing/backend/win_helper.h"

namespace printing {

namespace {

HRESULT StreamOnHGlobalToString(IStream* stream, std::string* out) {
  DCHECK(stream);
  DCHECK(out);
  HGLOBAL hdata = NULL;
  HRESULT hr = GetHGlobalFromStream(stream, &hdata);
  if (SUCCEEDED(hr)) {
    DCHECK(hdata);
    base::win::ScopedHGlobal<char> locked_data(hdata);
    out->assign(locked_data.release(), locked_data.Size());
  }
  return hr;
}

template <class T>
void GetDeviceCapabilityArray(const wchar_t* printer,
                              const wchar_t* port,
                              WORD id,
                              std::vector<T>* result) {
  int count = DeviceCapabilities(printer, port, id, NULL, NULL);
  if (count <= 0)
    return;
  std::vector<T> tmp;
  tmp.resize(count * 2);
  count = DeviceCapabilities(printer, port, id,
                             reinterpret_cast<LPTSTR>(tmp.data()), NULL);
  if (count <= 0)
    return;
  CHECK_LE(count, base::checked_cast<int>(tmp.size()));
  tmp.resize(count);
  result->swap(tmp);
}

void LoadPaper(const wchar_t* printer,
               const wchar_t* port,
               const DEVMODE* devmode,
               PrinterSemanticCapsAndDefaults* caps) {
  static const size_t kToUm = 100;  // Windows uses 0.1mm.
  static const size_t kMaxPaperName = 64;

  struct PaperName {
    wchar_t chars[kMaxPaperName];
  };

  DCHECK_EQ(sizeof(PaperName), sizeof(wchar_t) * kMaxPaperName);

  // Paper
  std::vector<PaperName> names;
  GetDeviceCapabilityArray(printer, port, DC_PAPERNAMES, &names);

  std::vector<POINT> sizes;
  GetDeviceCapabilityArray(printer, port, DC_PAPERSIZE, &sizes);

  std::vector<WORD> ids;
  GetDeviceCapabilityArray(printer, port, DC_PAPERS, &ids);

  DCHECK_EQ(ids.size(), sizes.size());
  DCHECK_EQ(names.size(), sizes.size());

  if (ids.size() != sizes.size())
    ids.clear();
  if (names.size() != sizes.size())
    names.clear();

  for (size_t i = 0; i < sizes.size(); ++i) {
    PrinterSemanticCapsAndDefaults::Paper paper;
    paper.size_um.SetSize(sizes[i].x * kToUm, sizes[i].y * kToUm);
    if (!names.empty()) {
      const wchar_t* name_start = names[i].chars;
      base::string16 tmp_name(name_start, kMaxPaperName);
      // Trim trailing zeros.
      tmp_name = tmp_name.c_str();
      paper.name = base::WideToUTF8(tmp_name);
    }
    caps->papers.push_back(paper);
  }

  if (devmode) {
    short default_id = 0;
    gfx::Size default_size;

    if (devmode->dmFields & DM_PAPERSIZE)
      default_id = devmode->dmPaperSize;
    if (devmode->dmFields & DM_PAPERWIDTH)
      default_size.set_width(devmode->dmPaperWidth * kToUm);
    if (devmode->dmFields & DM_PAPERLENGTH)
      default_size.set_height(devmode->dmPaperLength * kToUm);

    if (default_size.IsEmpty()) {
      for (size_t i = 0; i < ids.size(); ++i) {
        if (ids[i] == default_id) {
          PrinterSemanticCapsAndDefaults::Paper paper;
          paper.size_um.SetSize(sizes[i].x * kToUm, sizes[i].y * kToUm);
          if (!names.empty()) {
            const wchar_t* name_start = names[i].chars;
            base::string16 tmp_name(name_start, kMaxPaperName);
            // Trim trailing zeros.
            tmp_name = tmp_name.c_str();
            paper.name = base::WideToUTF8(tmp_name);
          }
          caps->default_paper = paper;
          break;
        }
      }
    } else {
      caps->default_paper.size_um = default_size;
    }
  }
}

void LoadDpi(const wchar_t* printer,
             const wchar_t* port,
             const DEVMODE* devmode,
             PrinterSemanticCapsAndDefaults* caps) {
  std::vector<POINT> dpis;
  GetDeviceCapabilityArray(printer, port, DC_ENUMRESOLUTIONS, &dpis);

  for (size_t i = 0; i < dpis.size() ; ++i)
    caps->dpis.push_back(gfx::Size(dpis[i].x, dpis[i].y));

  if (devmode) {
    if ((devmode->dmFields & DM_PRINTQUALITY) && devmode->dmPrintQuality > 0) {
      caps->default_dpi.SetSize(devmode->dmPrintQuality,
                                devmode->dmPrintQuality);
      if (devmode->dmFields & DM_YRESOLUTION) {
        caps->default_dpi.set_height(devmode->dmYResolution);
      }
    }
  }
}

}  // namespace

class PrintBackendWin : public PrintBackend {
 public:
  PrintBackendWin() {}

  // PrintBackend implementation.
  virtual bool EnumeratePrinters(PrinterList* printer_list) OVERRIDE;
  virtual std::string GetDefaultPrinterName() OVERRIDE;
  virtual bool GetPrinterSemanticCapsAndDefaults(
      const std::string& printer_name,
      PrinterSemanticCapsAndDefaults* printer_info) OVERRIDE;
  virtual bool GetPrinterCapsAndDefaults(
      const std::string& printer_name,
      PrinterCapsAndDefaults* printer_info) OVERRIDE;
  virtual std::string GetPrinterDriverInfo(
      const std::string& printer_name) OVERRIDE;
  virtual bool IsValidPrinter(const std::string& printer_name) OVERRIDE;

 protected:
  virtual ~PrintBackendWin() {}
};

bool PrintBackendWin::EnumeratePrinters(PrinterList* printer_list) {
  DCHECK(printer_list);
  DWORD bytes_needed = 0;
  DWORD count_returned = 0;
  const DWORD kLevel = 4;
  BOOL ret = EnumPrinters(PRINTER_ENUM_LOCAL|PRINTER_ENUM_CONNECTIONS, NULL,
                          kLevel, NULL, 0, &bytes_needed, &count_returned);
  if (!bytes_needed)
    return false;
  scoped_ptr<BYTE[]> printer_info_buffer(new BYTE[bytes_needed]);
  ret = EnumPrinters(PRINTER_ENUM_LOCAL|PRINTER_ENUM_CONNECTIONS, NULL, kLevel,
                     printer_info_buffer.get(), bytes_needed, &bytes_needed,
                     &count_returned);
  DCHECK(ret);
  if (!ret)
    return false;

  std::string default_printer = GetDefaultPrinterName();
  PRINTER_INFO_4* printer_info =
      reinterpret_cast<PRINTER_INFO_4*>(printer_info_buffer.get());
  for (DWORD index = 0; index < count_returned; index++) {
    ScopedPrinterHandle printer;
    PrinterBasicInfo info;
    if (printer.OpenPrinter(printer_info[index].pPrinterName) &&
        InitBasicPrinterInfo(printer, &info)) {
      info.is_default = (info.printer_name == default_printer);
      printer_list->push_back(info);
    }
  }
  return true;
}

std::string PrintBackendWin::GetDefaultPrinterName() {
  DWORD size = MAX_PATH;
  TCHAR default_printer_name[MAX_PATH];
  if (!::GetDefaultPrinter(default_printer_name, &size))
    return std::string();
  return base::WideToUTF8(default_printer_name);
}

bool PrintBackendWin::GetPrinterSemanticCapsAndDefaults(
    const std::string& printer_name,
    PrinterSemanticCapsAndDefaults* printer_info) {
  ScopedPrinterHandle printer_handle;
  if (!printer_handle.OpenPrinter(base::UTF8ToWide(printer_name).c_str())) {
    LOG(WARNING) << "Failed to open printer, error = " << GetLastError();
    return false;
  }

  PrinterInfo5 info_5;
  if (!info_5.Init(printer_handle))
    return false;
  const wchar_t* name = info_5.get()->pPrinterName;
  const wchar_t* port = info_5.get()->pPortName;
  DCHECK_EQ(name, base::UTF8ToUTF16(printer_name));

  PrinterSemanticCapsAndDefaults caps;
  UserDefaultDevMode user_settings;
  if (user_settings.Init(printer_handle)) {
    if (user_settings.get()->dmFields & DM_COLOR)
      caps.color_default = (user_settings.get()->dmColor == DMCOLOR_COLOR);

    if (user_settings.get()->dmFields & DM_DUPLEX) {
      switch (user_settings.get()->dmDuplex) {
      case DMDUP_SIMPLEX:
        caps.duplex_default = SIMPLEX;
        break;
      case DMDUP_VERTICAL:
        caps.duplex_default = LONG_EDGE;
        break;
      case DMDUP_HORIZONTAL:
        caps.duplex_default = SHORT_EDGE;
        break;
      default:
        NOTREACHED();
      }
    }

    if (user_settings.get()->dmFields & DM_COLLATE)
      caps.collate_default = (user_settings.get()->dmCollate == DMCOLLATE_TRUE);
  } else {
    LOG(WARNING) << "Fallback to color/simplex mode.";
    caps.color_default = caps.color_changeable;
    caps.duplex_default = SIMPLEX;
  }

  // Get printer capabilities. For more info see here:
  // http://msdn.microsoft.com/en-us/library/windows/desktop/dd183552(v=vs.85).aspx
  caps.color_changeable =
      (DeviceCapabilities(name, port, DC_COLORDEVICE, NULL, NULL) == 1);

  caps.duplex_capable =
      (DeviceCapabilities(name, port, DC_DUPLEX, NULL, NULL) == 1);

  caps.collate_capable =
      (DeviceCapabilities(name, port, DC_COLLATE, NULL, NULL) == 1);

  caps.copies_capable =
      (DeviceCapabilities(name, port, DC_COPIES, NULL, NULL) > 1);

  LoadPaper(name, port, user_settings.get(), &caps);
  LoadDpi(name, port, user_settings.get(), &caps);

  *printer_info = caps;
  return true;
}

bool PrintBackendWin::GetPrinterCapsAndDefaults(
    const std::string& printer_name,
    PrinterCapsAndDefaults* printer_info) {
  ScopedXPSInitializer xps_initializer;
  if (!xps_initializer.initialized()) {
    // TODO(sanjeevr): Handle legacy proxy case (with no prntvpt.dll)
    return false;
  }
  if (!IsValidPrinter(printer_name)) {
    return false;
  }
  DCHECK(printer_info);
  HPTPROVIDER provider = NULL;
  std::wstring printer_name_wide = base::UTF8ToWide(printer_name);
  HRESULT hr = XPSModule::OpenProvider(printer_name_wide, 1, &provider);
  if (provider) {
    base::win::ScopedComPtr<IStream> print_capabilities_stream;
    hr = CreateStreamOnHGlobal(NULL, TRUE,
                               print_capabilities_stream.Receive());
    DCHECK(SUCCEEDED(hr));
    if (print_capabilities_stream) {
      base::win::ScopedBstr error;
      hr = XPSModule::GetPrintCapabilities(provider,
                                           NULL,
                                           print_capabilities_stream,
                                           error.Receive());
      DCHECK(SUCCEEDED(hr));
      if (FAILED(hr)) {
        return false;
      }
      hr = StreamOnHGlobalToString(print_capabilities_stream.get(),
                                   &printer_info->printer_capabilities);
      DCHECK(SUCCEEDED(hr));
      printer_info->caps_mime_type = "text/xml";
    }
    ScopedPrinterHandle printer_handle;
    if (printer_handle.OpenPrinter(printer_name_wide.c_str())) {
      scoped_ptr<DEVMODE[]> devmode_out(CreateDevMode(printer_handle, NULL));
      if (!devmode_out)
        return false;
      base::win::ScopedComPtr<IStream> printer_defaults_stream;
      hr = CreateStreamOnHGlobal(NULL, TRUE,
                                 printer_defaults_stream.Receive());
      DCHECK(SUCCEEDED(hr));
      if (printer_defaults_stream) {
        DWORD dm_size = devmode_out.get()->dmSize +
                        devmode_out.get()->dmDriverExtra;
        hr = XPSModule::ConvertDevModeToPrintTicket(provider, dm_size,
            devmode_out.get(), kPTJobScope, printer_defaults_stream);
        DCHECK(SUCCEEDED(hr));
        if (SUCCEEDED(hr)) {
          hr = StreamOnHGlobalToString(printer_defaults_stream.get(),
                                       &printer_info->printer_defaults);
          DCHECK(SUCCEEDED(hr));
          printer_info->defaults_mime_type = "text/xml";
        }
      }
    }
    XPSModule::CloseProvider(provider);
  }
  return true;
}

// Gets the information about driver for a specific printer.
std::string PrintBackendWin::GetPrinterDriverInfo(
    const std::string& printer_name) {
  ScopedPrinterHandle printer;
  if (!printer.OpenPrinter(base::UTF8ToWide(printer_name).c_str())) {
    return std::string();
  }
  return GetDriverInfo(printer);
}

bool PrintBackendWin::IsValidPrinter(const std::string& printer_name) {
  ScopedPrinterHandle printer_handle;
  return printer_handle.OpenPrinter(base::UTF8ToWide(printer_name).c_str());
}

scoped_refptr<PrintBackend> PrintBackend::CreateInstance(
    const base::DictionaryValue* print_backend_settings) {
  return new PrintBackendWin;
}

}  // namespace printing
<|endoftext|>


<commit_before>/*
 * Copyright 2007-2020 CM4all GmbH
 * All rights reserved.
 *
 * author: Max Kellermann <mk@cm4all.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * - Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *
 * - Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the
 * distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
 * FOUNDATION OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "IstreamSpawn.hxx"
#include "spawn/Interface.hxx"
#include "spawn/Prepared.hxx"
#include "spawn/ExitListener.hxx"
#include "istream/UnusedPtr.hxx"
#include "istream/istream.hxx"
#include "istream/Pointer.hxx"
#include "istream/Handler.hxx"
#include "io/Splice.hxx"
#include "io/SpliceSupport.hxx"
#include "io/Buffered.hxx"
#include "io/UniqueFileDescriptor.hxx"
#include "io/Logger.hxx"
#include "net/SocketDescriptor.hxx"
#include "event/SocketEvent.hxx"
#include "pool/pool.hxx"
#include "fb_pool.hxx"
#include "SliceFifoBuffer.hxx"
#include "system/Error.hxx"
#include "util/Cast.hxx"

#ifdef __linux
#include <fcntl.h>
#endif

#include <assert.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <sched.h>
#include <signal.h>
#include <limits.h>

struct SpawnIstream final : Istream, IstreamHandler, ExitListener {
	const LLogger logger;
	SpawnService &spawn_service;

	UniqueFileDescriptor output_fd;
	SocketEvent output_event;

	SliceFifoBuffer buffer;

	IstreamPointer input;
	UniqueFileDescriptor input_fd;
	SocketEvent input_event;

	int pid;

	bool direct = false;

	SpawnIstream(SpawnService &_spawn_service, EventLoop &event_loop,
		     struct pool &p,
		     UnusedIstreamPtr _input, UniqueFileDescriptor _input_fd,
		     UniqueFileDescriptor _output_fd,
		     pid_t _pid) noexcept;

	bool CheckDirect() const noexcept {
		return direct;
	}

	void Cancel() noexcept;

	void FreeBuffer() noexcept {
		buffer.FreeIfDefined();
	}

	/**
	 * Send data from the buffer.  Invokes the "eof" callback when the
	 * buffer becomes empty and the pipe has been closed already.
	 *
	 * @return true if the caller shall read more data from the pipe
	 */
	bool SendFromBuffer() noexcept;

	void ReadFromOutput() noexcept;

	void InputEventCallback(unsigned) noexcept {
		input_event.Cancel(); // TODO: take advantage of EV_PERSIST

		input.Read();
	}

	void OutputEventCallback(unsigned) noexcept {
		output_event.Cancel(); // TODO: take advantage of EV_PERSIST

		ReadFromOutput();
	}

	/* virtual methods from class Istream */

	void _SetDirect(FdTypeMask mask) noexcept override {
		direct = (mask & ISTREAM_TO_PIPE) != 0;
	}

	void _Read() noexcept override;
	// TODO: implement int AsFd() override;
	void _Close() noexcept override;

	/* virtual methods from class IstreamHandler */
	size_t OnData(const void *data, size_t length) noexcept override;
	ssize_t OnDirect(FdType type, int fd, size_t max_length) noexcept override;
	void OnEof() noexcept override;
	void OnError(std::exception_ptr ep) noexcept override;

	/* virtual methods from class ExitListener */
	void OnChildProcessExit(int status) noexcept override;
};

void
SpawnIstream::Cancel() noexcept
{
	assert(output_fd.IsDefined());

	if (input.IsDefined()) {
		assert(input_fd.IsDefined());

		input_event.Cancel();
		input_fd.Close();
		input.ClearAndClose();
	}

	output_event.Cancel();

	output_fd.Close();

	if (pid >= 0) {
		spawn_service.KillChildProcess(pid);
		pid = -1;
	}
}

inline bool
SpawnIstream::SendFromBuffer() noexcept
{
	assert(buffer.IsDefined());

	if (Istream::SendFromBuffer(buffer) == 0)
		return false;

	if (!output_fd.IsDefined()) {
		if (buffer.empty()) {
			FreeBuffer();
			DestroyEof();
		}

		return false;
	}

	buffer.FreeIfEmpty();

	return true;
}

/*
 * input handler
 *
 */

size_t
SpawnIstream::OnData(const void *data, size_t length) noexcept
{
	assert(input_fd.IsDefined());

	ssize_t nbytes = input_fd.Write(data, length);
	if (nbytes > 0)
		input_event.ScheduleWrite();
	else if (nbytes < 0) {
		if (errno == EAGAIN) {
			input_event.ScheduleWrite();
			return 0;
		}

		logger(1, "write() to subprocess failed: ", strerror(errno));
		input_event.Cancel();
		input_fd.Close();
		input.ClearAndClose();
		return 0;
	}

	return (size_t)nbytes;
}

ssize_t
SpawnIstream::OnDirect(gcc_unused FdType type, int fd, size_t max_length) noexcept
{
	assert(input_fd.IsDefined());

	ssize_t nbytes = SpliceToPipe(fd, input_fd.Get(), max_length);
	if (nbytes > 0)
		input_event.ScheduleWrite();
	else if (nbytes < 0) {
		if (errno == EAGAIN) {
			if (!input_fd.IsReadyForWriting()) {
				input_event.ScheduleWrite();
				return ISTREAM_RESULT_BLOCKING;
			}

			/* try again, just in case connection->fd has become ready
			   between the first splice() call and
			   fd_ready_for_writing() */
			nbytes = SpliceToPipe(fd, input_fd.Get(), max_length);
		}
	}

	return nbytes;
}

inline void
SpawnIstream::OnEof() noexcept
{
	assert(input.IsDefined());
	assert(input_fd.IsDefined());

	input_event.Cancel();
	input_fd.Close();

	input.Clear();
}

void
SpawnIstream::OnError(std::exception_ptr ep) noexcept
{
	assert(input.IsDefined());
	assert(input_fd.IsDefined());

	FreeBuffer();

	input_event.Cancel();
	input_fd.Close();
	input.Clear();

	Cancel();
	DestroyError(ep);
}

/*
 * event for fork.output_fd
 */

void
SpawnIstream::ReadFromOutput() noexcept
{
	assert(output_fd.IsDefined());

	if (!CheckDirect()) {
		buffer.AllocateIfNull(fb_pool_get());

		ssize_t nbytes = read_to_buffer(output_fd.Get(), buffer, INT_MAX);
		if (nbytes == -2) {
			/* XXX should not happen */
		} else if (nbytes > 0) {
			if (Istream::SendFromBuffer(buffer) > 0) {
				buffer.FreeIfEmpty();
				output_event.ScheduleRead();
			}
		} else if (nbytes == 0) {
			Cancel();

			if (buffer.empty()) {
				FreeBuffer();
				DestroyEof();
			}
		} else if (errno == EAGAIN) {
			buffer.FreeIfEmpty();
			output_event.ScheduleRead();

			if (input.IsDefined())
				/* the CGI may be waiting for more data from stdin */
				input.Read();
		} else {
			auto error = MakeErrno("failed to read from sub process");
			FreeBuffer();
			Cancel();
			DestroyError(std::make_exception_ptr(error));
		}
	} else {
		if (Istream::ConsumeFromBuffer(buffer) > 0)
			/* there's data left in the buffer, which must be consumed
			   before we can switch to "direct" transfer */
			return;

		buffer.FreeIfDefined();

		/* at this point, the handler might have changed inside
		   Istream::ConsumeFromBuffer(), and the new handler might not
		   support "direct" transfer - check again */
		if (!CheckDirect()) {
			output_event.ScheduleRead();
			return;
		}

		ssize_t nbytes = InvokeDirect(FdType::FD_PIPE, output_fd.Get(),
					      INT_MAX);
		if (nbytes == ISTREAM_RESULT_BLOCKING ||
		    nbytes == ISTREAM_RESULT_CLOSED) {
			/* -2 means the callback wasn't able to consume any data right
			   now */
		} else if (nbytes > 0) {
			output_event.ScheduleRead();
		} else if (nbytes == ISTREAM_RESULT_EOF) {
			FreeBuffer();
			Cancel();
			DestroyEof();
		} else if (errno == EAGAIN) {
			output_event.ScheduleRead();

			if (input.IsDefined())
				/* the CGI may be waiting for more data from stdin */
				input.Read();
		} else {
			auto error = MakeErrno("failed to read from sub process");
			FreeBuffer();
			Cancel();
			DestroyError(std::make_exception_ptr(error));
		}
	}
}

/*
 * istream implementation
 *
 */

void
SpawnIstream::_Read() noexcept
{
	if (buffer.empty() || SendFromBuffer())
		ReadFromOutput();
}

void
SpawnIstream::_Close() noexcept
{
	FreeBuffer();

	if (output_fd.IsDefined())
		Cancel();

	Destroy();
}

/*
 * child callback
 *
 */

void
SpawnIstream::OnChildProcessExit(gcc_unused int status) noexcept
{
	assert(pid >= 0);

	pid = -1;
}

/*
 * constructor
 *
 */

inline
SpawnIstream::SpawnIstream(SpawnService &_spawn_service, EventLoop &event_loop,
			   struct pool &p,
			   UnusedIstreamPtr _input,
			   UniqueFileDescriptor _input_fd,
			   UniqueFileDescriptor _output_fd,
			   pid_t _pid) noexcept
	:Istream(p),
	 logger("spawn"),
	 spawn_service(_spawn_service),
	 output_fd(std::move(_output_fd)),
	 output_event(event_loop, BIND_THIS_METHOD(OutputEventCallback),
		      SocketDescriptor::FromFileDescriptor(output_fd)),
	 input(std::move(_input), *this),
	 input_fd(std::move(_input_fd)),
	 input_event(event_loop, BIND_THIS_METHOD(InputEventCallback),
		     SocketDescriptor::FromFileDescriptor(input_fd)),
	 pid(_pid)
{
	if (input.IsDefined()) {
		input.SetDirect(ISTREAM_TO_PIPE);
		input_event.ScheduleWrite();
	}

	spawn_service.SetExitListener(pid, this);
}

UnusedIstreamPtr
SpawnChildProcess(EventLoop &event_loop, struct pool *pool, const char *name,
		  UnusedIstreamPtr input,
		  PreparedChildProcess &&prepared,
		  SpawnService &spawn_service)
{
	if (input) {
		int fd = input.AsFd();
		if (fd >= 0)
			prepared.SetStdin(fd);
	}

	UniqueFileDescriptor stdin_pipe;
	if (input) {
		UniqueFileDescriptor stdin_r;
		if (!UniqueFileDescriptor::CreatePipe(stdin_r, stdin_pipe))
			throw MakeErrno("pipe() failed");

		prepared.SetStdin(std::move(stdin_r));

		stdin_pipe.SetNonBlocking();
	}

	UniqueFileDescriptor stdout_pipe, stdout_w;
	if (!UniqueFileDescriptor::CreatePipe(stdout_pipe, stdout_w))
		throw MakeErrno("pipe() failed");

	prepared.SetStdout(std::move(stdout_w));

	stdout_pipe.SetNonBlocking();

	const int pid = spawn_service.SpawnChildProcess(name, std::move(prepared),
							SocketDescriptor::Undefined(),
							nullptr);
	auto f = NewFromPool<SpawnIstream>(*pool, spawn_service, event_loop,
					   *pool,
					   std::move(input), std::move(stdin_pipe),
					   std::move(stdout_pipe),
					   pid);

	/* XXX CLOEXEC */

	return UnusedIstreamPtr(f);
}
<commit_msg>spawn/istream: derive from class IstreamSink<commit_after>/*
 * Copyright 2007-2020 CM4all GmbH
 * All rights reserved.
 *
 * author: Max Kellermann <mk@cm4all.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * - Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *
 * - Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the
 * distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
 * FOUNDATION OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "IstreamSpawn.hxx"
#include "spawn/Interface.hxx"
#include "spawn/Prepared.hxx"
#include "spawn/ExitListener.hxx"
#include "istream/UnusedPtr.hxx"
#include "istream/istream.hxx"
#include "istream/Sink.hxx"
#include "io/Splice.hxx"
#include "io/SpliceSupport.hxx"
#include "io/Buffered.hxx"
#include "io/UniqueFileDescriptor.hxx"
#include "io/Logger.hxx"
#include "net/SocketDescriptor.hxx"
#include "event/SocketEvent.hxx"
#include "pool/pool.hxx"
#include "fb_pool.hxx"
#include "SliceFifoBuffer.hxx"
#include "system/Error.hxx"
#include "util/Cast.hxx"

#ifdef __linux
#include <fcntl.h>
#endif

#include <assert.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <sched.h>
#include <signal.h>
#include <limits.h>

struct SpawnIstream final : Istream, IstreamSink, ExitListener {
	const LLogger logger;
	SpawnService &spawn_service;

	UniqueFileDescriptor output_fd;
	SocketEvent output_event;

	SliceFifoBuffer buffer;

	UniqueFileDescriptor input_fd;
	SocketEvent input_event;

	int pid;

	bool direct = false;

	SpawnIstream(SpawnService &_spawn_service, EventLoop &event_loop,
		     struct pool &p,
		     UnusedIstreamPtr _input, UniqueFileDescriptor _input_fd,
		     UniqueFileDescriptor _output_fd,
		     pid_t _pid) noexcept;

	bool CheckDirect() const noexcept {
		return direct;
	}

	void Cancel() noexcept;

	void FreeBuffer() noexcept {
		buffer.FreeIfDefined();
	}

	/**
	 * Send data from the buffer.  Invokes the "eof" callback when the
	 * buffer becomes empty and the pipe has been closed already.
	 *
	 * @return true if the caller shall read more data from the pipe
	 */
	bool SendFromBuffer() noexcept;

	void ReadFromOutput() noexcept;

	void InputEventCallback(unsigned) noexcept {
		input_event.Cancel(); // TODO: take advantage of EV_PERSIST

		input.Read();
	}

	void OutputEventCallback(unsigned) noexcept {
		output_event.Cancel(); // TODO: take advantage of EV_PERSIST

		ReadFromOutput();
	}

	/* virtual methods from class Istream */

	void _SetDirect(FdTypeMask mask) noexcept override {
		direct = (mask & ISTREAM_TO_PIPE) != 0;
	}

	void _Read() noexcept override;
	// TODO: implement int AsFd() override;
	void _Close() noexcept override;

	/* virtual methods from class IstreamHandler */
	size_t OnData(const void *data, size_t length) noexcept override;
	ssize_t OnDirect(FdType type, int fd, size_t max_length) noexcept override;
	void OnEof() noexcept override;
	void OnError(std::exception_ptr ep) noexcept override;

	/* virtual methods from class ExitListener */
	void OnChildProcessExit(int status) noexcept override;
};

void
SpawnIstream::Cancel() noexcept
{
	assert(output_fd.IsDefined());

	if (HasInput()) {
		assert(input_fd.IsDefined());

		input_event.Cancel();
		input_fd.Close();
		ClearAndCloseInput();
	}

	output_event.Cancel();

	output_fd.Close();

	if (pid >= 0) {
		spawn_service.KillChildProcess(pid);
		pid = -1;
	}
}

inline bool
SpawnIstream::SendFromBuffer() noexcept
{
	assert(buffer.IsDefined());

	if (Istream::SendFromBuffer(buffer) == 0)
		return false;

	if (!output_fd.IsDefined()) {
		if (buffer.empty()) {
			FreeBuffer();
			DestroyEof();
		}

		return false;
	}

	buffer.FreeIfEmpty();

	return true;
}

/*
 * input handler
 *
 */

size_t
SpawnIstream::OnData(const void *data, size_t length) noexcept
{
	assert(input_fd.IsDefined());

	ssize_t nbytes = input_fd.Write(data, length);
	if (nbytes > 0)
		input_event.ScheduleWrite();
	else if (nbytes < 0) {
		if (errno == EAGAIN) {
			input_event.ScheduleWrite();
			return 0;
		}

		logger(1, "write() to subprocess failed: ", strerror(errno));
		input_event.Cancel();
		input_fd.Close();
		ClearAndCloseInput();
		return 0;
	}

	return (size_t)nbytes;
}

ssize_t
SpawnIstream::OnDirect(gcc_unused FdType type, int fd, size_t max_length) noexcept
{
	assert(input_fd.IsDefined());

	ssize_t nbytes = SpliceToPipe(fd, input_fd.Get(), max_length);
	if (nbytes > 0)
		input_event.ScheduleWrite();
	else if (nbytes < 0) {
		if (errno == EAGAIN) {
			if (!input_fd.IsReadyForWriting()) {
				input_event.ScheduleWrite();
				return ISTREAM_RESULT_BLOCKING;
			}

			/* try again, just in case connection->fd has become ready
			   between the first splice() call and
			   fd_ready_for_writing() */
			nbytes = SpliceToPipe(fd, input_fd.Get(), max_length);
		}
	}

	return nbytes;
}

inline void
SpawnIstream::OnEof() noexcept
{
	assert(HasInput());
	assert(input_fd.IsDefined());

	input_event.Cancel();
	input_fd.Close();

	ClearInput();
}

void
SpawnIstream::OnError(std::exception_ptr ep) noexcept
{
	assert(HasInput());
	assert(input_fd.IsDefined());

	FreeBuffer();

	input_event.Cancel();
	input_fd.Close();
	ClearInput();

	Cancel();
	DestroyError(ep);
}

/*
 * event for fork.output_fd
 */

void
SpawnIstream::ReadFromOutput() noexcept
{
	assert(output_fd.IsDefined());

	if (!CheckDirect()) {
		buffer.AllocateIfNull(fb_pool_get());

		ssize_t nbytes = read_to_buffer(output_fd.Get(), buffer, INT_MAX);
		if (nbytes == -2) {
			/* XXX should not happen */
		} else if (nbytes > 0) {
			if (Istream::SendFromBuffer(buffer) > 0) {
				buffer.FreeIfEmpty();
				output_event.ScheduleRead();
			}
		} else if (nbytes == 0) {
			Cancel();

			if (buffer.empty()) {
				FreeBuffer();
				DestroyEof();
			}
		} else if (errno == EAGAIN) {
			buffer.FreeIfEmpty();
			output_event.ScheduleRead();

			if (HasInput())
				/* the CGI may be waiting for more data from stdin */
				input.Read();
		} else {
			auto error = MakeErrno("failed to read from sub process");
			FreeBuffer();
			Cancel();
			DestroyError(std::make_exception_ptr(error));
		}
	} else {
		if (Istream::ConsumeFromBuffer(buffer) > 0)
			/* there's data left in the buffer, which must be consumed
			   before we can switch to "direct" transfer */
			return;

		buffer.FreeIfDefined();

		/* at this point, the handler might have changed inside
		   Istream::ConsumeFromBuffer(), and the new handler might not
		   support "direct" transfer - check again */
		if (!CheckDirect()) {
			output_event.ScheduleRead();
			return;
		}

		ssize_t nbytes = InvokeDirect(FdType::FD_PIPE, output_fd.Get(),
					      INT_MAX);
		if (nbytes == ISTREAM_RESULT_BLOCKING ||
		    nbytes == ISTREAM_RESULT_CLOSED) {
			/* -2 means the callback wasn't able to consume any data right
			   now */
		} else if (nbytes > 0) {
			output_event.ScheduleRead();
		} else if (nbytes == ISTREAM_RESULT_EOF) {
			FreeBuffer();
			Cancel();
			DestroyEof();
		} else if (errno == EAGAIN) {
			output_event.ScheduleRead();

			if (HasInput())
				/* the CGI may be waiting for more data from stdin */
				input.Read();
		} else {
			auto error = MakeErrno("failed to read from sub process");
			FreeBuffer();
			Cancel();
			DestroyError(std::make_exception_ptr(error));
		}
	}
}

/*
 * istream implementation
 *
 */

void
SpawnIstream::_Read() noexcept
{
	if (buffer.empty() || SendFromBuffer())
		ReadFromOutput();
}

void
SpawnIstream::_Close() noexcept
{
	FreeBuffer();

	if (output_fd.IsDefined())
		Cancel();

	Destroy();
}

/*
 * child callback
 *
 */

void
SpawnIstream::OnChildProcessExit(gcc_unused int status) noexcept
{
	assert(pid >= 0);

	pid = -1;
}

/*
 * constructor
 *
 */

inline
SpawnIstream::SpawnIstream(SpawnService &_spawn_service, EventLoop &event_loop,
			   struct pool &p,
			   UnusedIstreamPtr _input,
			   UniqueFileDescriptor _input_fd,
			   UniqueFileDescriptor _output_fd,
			   pid_t _pid) noexcept
	:Istream(p),
	 IstreamSink(std::move(_input)),
	 logger("spawn"),
	 spawn_service(_spawn_service),
	 output_fd(std::move(_output_fd)),
	 output_event(event_loop, BIND_THIS_METHOD(OutputEventCallback),
		      SocketDescriptor::FromFileDescriptor(output_fd)),
	 input_fd(std::move(_input_fd)),
	 input_event(event_loop, BIND_THIS_METHOD(InputEventCallback),
		     SocketDescriptor::FromFileDescriptor(input_fd)),
	 pid(_pid)
{
	if (HasInput()) {
		input.SetDirect(ISTREAM_TO_PIPE);
		input_event.ScheduleWrite();
	}

	spawn_service.SetExitListener(pid, this);
}

UnusedIstreamPtr
SpawnChildProcess(EventLoop &event_loop, struct pool *pool, const char *name,
		  UnusedIstreamPtr input,
		  PreparedChildProcess &&prepared,
		  SpawnService &spawn_service)
{
	if (input) {
		int fd = input.AsFd();
		if (fd >= 0)
			prepared.SetStdin(fd);
	}

	UniqueFileDescriptor stdin_pipe;
	if (input) {
		UniqueFileDescriptor stdin_r;
		if (!UniqueFileDescriptor::CreatePipe(stdin_r, stdin_pipe))
			throw MakeErrno("pipe() failed");

		prepared.SetStdin(std::move(stdin_r));

		stdin_pipe.SetNonBlocking();
	}

	UniqueFileDescriptor stdout_pipe, stdout_w;
	if (!UniqueFileDescriptor::CreatePipe(stdout_pipe, stdout_w))
		throw MakeErrno("pipe() failed");

	prepared.SetStdout(std::move(stdout_w));

	stdout_pipe.SetNonBlocking();

	const int pid = spawn_service.SpawnChildProcess(name, std::move(prepared),
							SocketDescriptor::Undefined(),
							nullptr);
	auto f = NewFromPool<SpawnIstream>(*pool, spawn_service, event_loop,
					   *pool,
					   std::move(input), std::move(stdin_pipe),
					   std::move(stdout_pipe),
					   pid);

	/* XXX CLOEXEC */

	return UnusedIstreamPtr(f);
}
<|endoftext|>