<commit_before>// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "gbdt.h"
#include "vm_forest.h"
#include <vespa/eval/eval/basic_nodes.h>
#include <vespa/eval/eval/call_nodes.h>
#include <vespa/eval/eval/operator_nodes.h>

namespace vespalib {
namespace eval {
namespace gbdt {

namespace {

//-----------------------------------------------------------------------------

constexpr uint32_t LEAF     = 0;
constexpr uint32_t LESS     = 1;
constexpr uint32_t IN       = 2;
constexpr uint32_t INVERTED = 3;

// layout:
//
// <feature+types>: [feature ref|my type|left child type|right child type]
// bits:                      20       4               4                4
//
// LEAF:    [const]
// bits:         32
//
// LESS:    [<feature+types>][const][skip]
// bits                    32     32    32
//
// IN:      [<feature+types>][skip|set size](set size)X[const]
// bits                    32    24        8                64

// Note: We need to use double for set membership checks (IN) due to
// string hashing.

const double *as_double_ptr(const uint32_t *pos) {
    return reinterpret_cast<const double*>(pos);
}

const float *as_float_ptr(const uint32_t *pos) {
    return reinterpret_cast<const float*>(pos);
}

bool find_in(double value, const double *set, const double *end) {
    for (; set < end; ++set) {
        if (value == *set) {
            return true;
        }
    }
    return false;
}

double less_only_find_leaf(const double *input, const uint32_t *pos, uint32_t node_type) {
    for (;;) {
        if (input[pos[0] >> 12] < *as_float_ptr(pos + 1)) {
            node_type = (pos[0] & 0xf0) >> 4;
            pos += 3;
        } else {
            node_type = (pos[0] & 0xf);
            pos += 3 + pos[2];
        }
        if (node_type == LEAF) {
            return *as_float_ptr(pos);
        }
    }
}

double general_find_leaf(const double *input, const uint32_t *pos, uint32_t node_type) {
    for (;;) {
        if (node_type == LESS) {
            if (input[pos[0] >> 12] < *as_float_ptr(pos + 1)) {
                node_type = (pos[0] & 0xf0) >> 4;
                pos += 3;
            } else {
                node_type = (pos[0] & 0xf);
                pos += 3 + pos[2];
            }
            if (node_type == LEAF) {
                return *as_float_ptr(pos);
            }
        } else if (node_type == IN) {
            if (find_in(input[pos[0] >> 12], as_double_ptr(pos + 2),
                        as_double_ptr(pos + 2 + (2 * (pos[1] & 0xff)))))
            {
                node_type = (pos[0] & 0xf0) >> 4;
                pos += 2 + (2 * (pos[1] & 0xff));
            } else {
                node_type = (pos[0] & 0xf);
                pos += (2 + (2 * (pos[1] & 0xff))) + (pos[1] >> 8);
            }
            if (node_type == LEAF) {
                return *as_float_ptr(pos);
            }
        } else {
            if (input[pos[0] >> 12] >= *as_float_ptr(pos + 1)) {
                node_type = (pos[0] & 0xf);
                pos += 3 + pos[2];
            } else {
                node_type = (pos[0] & 0xf0) >> 4;
                pos += 3;
            }
            if (node_type == LEAF) {
                return *as_float_ptr(pos);
            }
        }
    }
}

//-----------------------------------------------------------------------------

void encode_large_const(double value, std::vector<uint32_t> &model_out) {
    union {
        double   d[1];
        uint32_t i[2];
    } buf;
    assert(sizeof(buf) == sizeof(double));
    buf.d[0] = value;
    model_out.push_back(buf.i[0]);
    model_out.push_back(buf.i[1]);
}

void encode_const(float value, std::vector<uint32_t> &model_out) {
    union {
        float    f[1];
        uint32_t i[1];
    } buf;
    assert(sizeof(buf) == sizeof(float));
    buf.f[0] = value;
    model_out.push_back(buf.i[0]);
}

uint32_t encode_node(const nodes::Node &node_in, std::vector<uint32_t> &model_out);

void encode_less(const nodes::Less &less,
                 const nodes::Node &left_child, const nodes::Node &right_child,
                 std::vector<uint32_t> &model_out)
{
    size_t meta_idx = model_out.size();
    auto symbol = nodes::as<nodes::Symbol>(less.lhs());
    assert(symbol);
    model_out.push_back(uint32_t(symbol->id()) << 12);
    assert(less.rhs().is_const());
    encode_const(less.rhs().get_const_value(), model_out);
    size_t skip_idx = model_out.size();
    model_out.push_back(0); // left child size placeholder
    uint32_t left_type = encode_node(left_child, model_out);
    model_out[skip_idx] = (model_out.size() - (skip_idx + 1));
    uint32_t right_type = encode_node(right_child, model_out);
    model_out[meta_idx] |= ((LESS << 8) | (left_type << 4) | right_type);
}

void encode_in(const nodes::In &in,
               const nodes::Node &left_child, const nodes::Node &right_child,
               std::vector<uint32_t> &model_out)
{
    size_t meta_idx = model_out.size();
    auto symbol = nodes::as<nodes::Symbol>(in.child());
    assert(symbol);
    model_out.push_back(uint32_t(symbol->id()) << 12);
    size_t set_size_idx = model_out.size();
    model_out.push_back(in.num_entries());
    for (size_t i = 0; i < in.num_entries(); ++i) {
        encode_large_const(in.get_entry(i).get_const_value(), model_out);
    }
    size_t left_idx = model_out.size();
    uint32_t left_type = encode_node(left_child, model_out);
    model_out[set_size_idx] |= (model_out.size() - left_idx) << 8;
    uint32_t right_type = encode_node(right_child, model_out);
    model_out[meta_idx] |= ((IN << 8) | (left_type << 4) | right_type);
}

void encode_inverted(const nodes::Not &inverted,
                     const nodes::Node &left_child, const nodes::Node &right_child,
                     std::vector<uint32_t> &model_out)
{
    size_t meta_idx = model_out.size();
    auto ge = nodes::as<nodes::GreaterEqual>(inverted.child());
    assert(ge);
    auto symbol = nodes::as<nodes::Symbol>(ge->lhs());
    assert(symbol);
    model_out.push_back(uint32_t(symbol->id()) << 12);
    assert(ge->rhs().is_const());
    encode_const(ge->rhs().get_const_value(), model_out);
    size_t skip_idx = model_out.size();
    model_out.push_back(0); // left child size placeholder
    uint32_t left_type = encode_node(left_child, model_out);
    model_out[skip_idx] = (model_out.size() - (skip_idx + 1));
    uint32_t right_type = encode_node(right_child, model_out);
    model_out[meta_idx] |= ((INVERTED << 8) | (left_type << 4) | right_type);
}

uint32_t encode_node(const nodes::Node &node_in, std::vector<uint32_t> &model_out) {
    auto if_node = nodes::as<nodes::If>(node_in);
    if (if_node) {
        auto less = nodes::as<nodes::Less>(if_node->cond());
        auto in = nodes::as<nodes::In>(if_node->cond());
        auto inverted = nodes::as<nodes::Not>(if_node->cond());
        if (less) {
            encode_less(*less, if_node->true_expr(), if_node->false_expr(), model_out);
            return LESS;
        } else if (in) {
            encode_in(*in, if_node->true_expr(), if_node->false_expr(), model_out);
            return IN;
        } else {
            assert(inverted);
            encode_inverted(*inverted, if_node->true_expr(), if_node->false_expr(), model_out);
            return INVERTED;
        }
    } else {
        assert(node_in.is_const());
        encode_const(node_in.get_const_value(), model_out);
        return LEAF;
    }
}

void encode_tree(const nodes::Node &root_in, std::vector<uint32_t> &model_out) {
    size_t size_idx = model_out.size();
    model_out.push_back(0); // tree size placeholder
    encode_node(root_in, model_out);
    model_out[size_idx] = (model_out.size() - (size_idx + 1));
}

//-----------------------------------------------------------------------------

Optimize::Result optimize(const std::vector<const nodes::Node *> &trees,
                          Forest::eval_function eval)
{
    std::vector<uint32_t> model;
    for (const nodes::Node *tree: trees) {
        encode_tree(*tree, model);
    }
    return Optimize::Result(Forest::UP(new VMForest(std::move(model))), eval);
}

//-----------------------------------------------------------------------------

} // namespace vespalib::eval::gbdt::<unnamed>

//-----------------------------------------------------------------------------

Optimize::Result
VMForest::less_only_optimize(const ForestStats &stats,
                             const std::vector<const nodes::Node *> &trees)
{
    if ((stats.total_in_checks > 0) || (stats.total_inverted_checks > 0)) {
        return Optimize::Result();
    }
    return optimize(trees, less_only_eval);
}

double
VMForest::less_only_eval(const Forest *forest, const double *input)
{
    const VMForest &self = *((const VMForest *)forest);
    const uint32_t *pos = &self._model[0];
    const uint32_t *end = pos + self._model.size();
    double sum = 0.0;
    while (pos < end) {
        uint32_t tree_size = *pos++;
        sum += less_only_find_leaf(input, pos, (*pos & 0xf00) >> 8);
        pos += tree_size;
    }
    return sum;
}

Optimize::Result
VMForest::general_optimize(const ForestStats &stats,
                           const std::vector<const nodes::Node *> &trees)
{
    if (stats.max_set_size > 255) {
        return Optimize::Result();
    }
    return optimize(trees, general_eval);
}

double
VMForest::general_eval(const Forest *forest, const double *input)
{
    const VMForest &self = *((const VMForest *)forest);
    const uint32_t *pos = &self._model[0];
    const uint32_t *end = pos + self._model.size();
    double sum = 0.0;
    while (pos < end) {
        uint32_t tree_size = *pos++;
        sum += general_find_leaf(input, pos, (*pos & 0xf00) >> 8);
        pos += tree_size;
    }
    return sum;
}

Optimize::Chain VMForest::optimize_chain({less_only_optimize, general_optimize});

//-----------------------------------------------------------------------------

} // namespace vespalib::eval::gbdt
} // namespace vespalib::eval
} // namespace vespalib
<commit_msg>avoid using union for type conversion<commit_after>// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "gbdt.h"
#include "vm_forest.h"
#include <vespa/eval/eval/basic_nodes.h>
#include <vespa/eval/eval/call_nodes.h>
#include <vespa/eval/eval/operator_nodes.h>

namespace vespalib {
namespace eval {
namespace gbdt {

namespace {

//-----------------------------------------------------------------------------

constexpr uint32_t LEAF     = 0;
constexpr uint32_t LESS     = 1;
constexpr uint32_t IN       = 2;
constexpr uint32_t INVERTED = 3;

// layout:
//
// <feature+types>: [feature ref|my type|left child type|right child type]
// bits:                      20       4               4                4
//
// LEAF:    [const]
// bits:         32
//
// LESS:    [<feature+types>][const][skip]
// bits                    32     32    32
//
// IN:      [<feature+types>][skip|set size](set size)X[const]
// bits                    32    24        8                64

// Note: We need to use double for set membership checks (IN) due to
// string hashing.

const double *as_double_ptr(const uint32_t *pos) {
    return reinterpret_cast<const double*>(pos);
}

const float *as_float_ptr(const uint32_t *pos) {
    return reinterpret_cast<const float*>(pos);
}

bool find_in(double value, const double *set, const double *end) {
    for (; set < end; ++set) {
        if (value == *set) {
            return true;
        }
    }
    return false;
}

double less_only_find_leaf(const double *input, const uint32_t *pos, uint32_t node_type) {
    for (;;) {
        if (input[pos[0] >> 12] < *as_float_ptr(pos + 1)) {
            node_type = (pos[0] & 0xf0) >> 4;
            pos += 3;
        } else {
            node_type = (pos[0] & 0xf);
            pos += 3 + pos[2];
        }
        if (node_type == LEAF) {
            return *as_float_ptr(pos);
        }
    }
}

double general_find_leaf(const double *input, const uint32_t *pos, uint32_t node_type) {
    for (;;) {
        if (node_type == LESS) {
            if (input[pos[0] >> 12] < *as_float_ptr(pos + 1)) {
                node_type = (pos[0] & 0xf0) >> 4;
                pos += 3;
            } else {
                node_type = (pos[0] & 0xf);
                pos += 3 + pos[2];
            }
            if (node_type == LEAF) {
                return *as_float_ptr(pos);
            }
        } else if (node_type == IN) {
            if (find_in(input[pos[0] >> 12], as_double_ptr(pos + 2),
                        as_double_ptr(pos + 2 + (2 * (pos[1] & 0xff)))))
            {
                node_type = (pos[0] & 0xf0) >> 4;
                pos += 2 + (2 * (pos[1] & 0xff));
            } else {
                node_type = (pos[0] & 0xf);
                pos += (2 + (2 * (pos[1] & 0xff))) + (pos[1] >> 8);
            }
            if (node_type == LEAF) {
                return *as_float_ptr(pos);
            }
        } else {
            if (input[pos[0] >> 12] >= *as_float_ptr(pos + 1)) {
                node_type = (pos[0] & 0xf);
                pos += 3 + pos[2];
            } else {
                node_type = (pos[0] & 0xf0) >> 4;
                pos += 3;
            }
            if (node_type == LEAF) {
                return *as_float_ptr(pos);
            }
        }
    }
}

//-----------------------------------------------------------------------------

void encode_large_const(double value, std::vector<uint32_t> &model_out) {
    uint32_t buf[2];
    static_assert(sizeof(buf) == sizeof(value));
    memcpy(buf, &value, sizeof(value));
    model_out.push_back(buf[0]);
    model_out.push_back(buf[1]);
}

void encode_const(float value, std::vector<uint32_t> &model_out) {
    uint32_t buf;
    static_assert(sizeof(buf) == sizeof(value));
    memcpy(&buf, &value, sizeof(value));
    model_out.push_back(buf);
}

uint32_t encode_node(const nodes::Node &node_in, std::vector<uint32_t> &model_out);

void encode_less(const nodes::Less &less,
                 const nodes::Node &left_child, const nodes::Node &right_child,
                 std::vector<uint32_t> &model_out)
{
    size_t meta_idx = model_out.size();
    auto symbol = nodes::as<nodes::Symbol>(less.lhs());
    assert(symbol);
    model_out.push_back(uint32_t(symbol->id()) << 12);
    assert(less.rhs().is_const());
    encode_const(less.rhs().get_const_value(), model_out);
    size_t skip_idx = model_out.size();
    model_out.push_back(0); // left child size placeholder
    uint32_t left_type = encode_node(left_child, model_out);
    model_out[skip_idx] = (model_out.size() - (skip_idx + 1));
    uint32_t right_type = encode_node(right_child, model_out);
    model_out[meta_idx] |= ((LESS << 8) | (left_type << 4) | right_type);
}

void encode_in(const nodes::In &in,
               const nodes::Node &left_child, const nodes::Node &right_child,
               std::vector<uint32_t> &model_out)
{
    size_t meta_idx = model_out.size();
    auto symbol = nodes::as<nodes::Symbol>(in.child());
    assert(symbol);
    model_out.push_back(uint32_t(symbol->id()) << 12);
    size_t set_size_idx = model_out.size();
    model_out.push_back(in.num_entries());
    for (size_t i = 0; i < in.num_entries(); ++i) {
        encode_large_const(in.get_entry(i).get_const_value(), model_out);
    }
    size_t left_idx = model_out.size();
    uint32_t left_type = encode_node(left_child, model_out);
    model_out[set_size_idx] |= (model_out.size() - left_idx) << 8;
    uint32_t right_type = encode_node(right_child, model_out);
    model_out[meta_idx] |= ((IN << 8) | (left_type << 4) | right_type);
}

void encode_inverted(const nodes::Not &inverted,
                     const nodes::Node &left_child, const nodes::Node &right_child,
                     std::vector<uint32_t> &model_out)
{
    size_t meta_idx = model_out.size();
    auto ge = nodes::as<nodes::GreaterEqual>(inverted.child());
    assert(ge);
    auto symbol = nodes::as<nodes::Symbol>(ge->lhs());
    assert(symbol);
    model_out.push_back(uint32_t(symbol->id()) << 12);
    assert(ge->rhs().is_const());
    encode_const(ge->rhs().get_const_value(), model_out);
    size_t skip_idx = model_out.size();
    model_out.push_back(0); // left child size placeholder
    uint32_t left_type = encode_node(left_child, model_out);
    model_out[skip_idx] = (model_out.size() - (skip_idx + 1));
    uint32_t right_type = encode_node(right_child, model_out);
    model_out[meta_idx] |= ((INVERTED << 8) | (left_type << 4) | right_type);
}

uint32_t encode_node(const nodes::Node &node_in, std::vector<uint32_t> &model_out) {
    auto if_node = nodes::as<nodes::If>(node_in);
    if (if_node) {
        auto less = nodes::as<nodes::Less>(if_node->cond());
        auto in = nodes::as<nodes::In>(if_node->cond());
        auto inverted = nodes::as<nodes::Not>(if_node->cond());
        if (less) {
            encode_less(*less, if_node->true_expr(), if_node->false_expr(), model_out);
            return LESS;
        } else if (in) {
            encode_in(*in, if_node->true_expr(), if_node->false_expr(), model_out);
            return IN;
        } else {
            assert(inverted);
            encode_inverted(*inverted, if_node->true_expr(), if_node->false_expr(), model_out);
            return INVERTED;
        }
    } else {
        assert(node_in.is_const());
        encode_const(node_in.get_const_value(), model_out);
        return LEAF;
    }
}

void encode_tree(const nodes::Node &root_in, std::vector<uint32_t> &model_out) {
    size_t size_idx = model_out.size();
    model_out.push_back(0); // tree size placeholder
    encode_node(root_in, model_out);
    model_out[size_idx] = (model_out.size() - (size_idx + 1));
}

//-----------------------------------------------------------------------------

Optimize::Result optimize(const std::vector<const nodes::Node *> &trees,
                          Forest::eval_function eval)
{
    std::vector<uint32_t> model;
    for (const nodes::Node *tree: trees) {
        encode_tree(*tree, model);
    }
    return Optimize::Result(Forest::UP(new VMForest(std::move(model))), eval);
}

//-----------------------------------------------------------------------------

} // namespace vespalib::eval::gbdt::<unnamed>

//-----------------------------------------------------------------------------

Optimize::Result
VMForest::less_only_optimize(const ForestStats &stats,
                             const std::vector<const nodes::Node *> &trees)
{
    if ((stats.total_in_checks > 0) || (stats.total_inverted_checks > 0)) {
        return Optimize::Result();
    }
    return optimize(trees, less_only_eval);
}

double
VMForest::less_only_eval(const Forest *forest, const double *input)
{
    const VMForest &self = *((const VMForest *)forest);
    const uint32_t *pos = &self._model[0];
    const uint32_t *end = pos + self._model.size();
    double sum = 0.0;
    while (pos < end) {
        uint32_t tree_size = *pos++;
        sum += less_only_find_leaf(input, pos, (*pos & 0xf00) >> 8);
        pos += tree_size;
    }
    return sum;
}

Optimize::Result
VMForest::general_optimize(const ForestStats &stats,
                           const std::vector<const nodes::Node *> &trees)
{
    if (stats.max_set_size > 255) {
        return Optimize::Result();
    }
    return optimize(trees, general_eval);
}

double
VMForest::general_eval(const Forest *forest, const double *input)
{
    const VMForest &self = *((const VMForest *)forest);
    const uint32_t *pos = &self._model[0];
    const uint32_t *end = pos + self._model.size();
    double sum = 0.0;
    while (pos < end) {
        uint32_t tree_size = *pos++;
        sum += general_find_leaf(input, pos, (*pos & 0xf00) >> 8);
        pos += tree_size;
    }
    return sum;
}

Optimize::Chain VMForest::optimize_chain({less_only_optimize, general_optimize});

//-----------------------------------------------------------------------------

} // namespace vespalib::eval::gbdt
} // namespace vespalib::eval
} // namespace vespalib
<|endoftext|>


<commit_before>/*
 * (C) Copyright 2015 ETH Zurich Systems Group (http://www.systems.ethz.ch/) and others.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Contributors:
 *     Markus Pilman <mpilman@inf.ethz.ch>
 *     Simon Loesing <sloesing@inf.ethz.ch>
 *     Thomas Etter <etterth@gmail.com>
 *     Kevin Bocksrocker <kevin.bocksrocker@gmail.com>
 *     Lucas Braun <braunl@inf.ethz.ch>
 */
/*
 * This class serves as a simple wrapper for RPC calls in Boost Asio. You can
 * define different classes of commands by calling GEN_COMMANDS(class-name, command-list)
 * once for each class.
 *
 */
#pragma once
#include <tuple>
#include <cstdint>
#include <type_traits>

#include <boost/system/error_code.hpp>
#include <boost/asio.hpp>
#include <boost/preprocessor.hpp>

#include <crossbow/Serializer.hpp>
#include <crossbow/string.hpp>

#define GEN_CASE(r, data, elem)\
    case data::elem:\
        static_cast<Derived*>(this)->template execute<data::elem>();\
        break;

#define SWITCH_CASE(Name, cmd, tuple) switch(cmd) {\
    BOOST_PP_SEQ_FOR_EACH(GEN_CASE, Name, BOOST_PP_TUPLE_TO_SEQ(tuple))\
}

#define GEN_COMMANDS(Name, tuple) enum class Name {\
    BOOST_PP_TUPLE_ELEM(0, tuple) = 1, \
    BOOST_PP_TUPLE_ENUM(BOOST_PP_TUPLE_POP_FRONT(tuple)) \
};\
template<class Derived>\
struct Name ## _Switch {\
    using type = Name;\
    void execute_impl(Name cmd) {\
        SWITCH_CASE(Name, cmd, tuple)\
    }\
}

#define SERVER_TYPE_S(Name, Signature) crossbow::protocol::Server<Name ## _Switch, Name, Signature, Impl>
#define SERVER_TYPE(Name, Impl) crossbow::protocol::Server<Name ## _Switch, Name, Signature, Impl>

namespace crossbow {
namespace protocol {

namespace impl {

template<class... Args>
struct ArgSerializer;

template<class Head, class... Tail>
struct ArgSerializer<Head, Tail...> {
    ArgSerializer<Tail...> rest;

    template<class C>
    void exec(C& c, const Head& head, const Tail&... tail) const {
        c & head;
        rest.exec(c, tail...);
    }
};

template<>
struct ArgSerializer<> {
    template<class C>
    void exec(C&) const {}
};

} // namespace impl

template<class... Args>
struct argsType;
template<class A, class B, class... Tail>
struct argsType<A, B, Tail...> {
    using type = std::tuple<A, B, Tail...>;
};
template<class Arg>
struct argsType<Arg> {
    using type = Arg;
};
template<>
struct argsType<> {
    using type = void;
};

template<class Command, template <Command> class Signature>
class Client {
    boost::asio::ip::tcp::socket& mSocket;
    size_t mCurrSize = 1024;
    std::unique_ptr<uint8_t[]> mCurrentRequest;
public:
    Client(boost::asio::ip::tcp::socket& socket)
        : mSocket(socket), mCurrentRequest(new uint8_t[mCurrSize])
    {
    }

    template<class Callback, class Result>
    typename std::enable_if<std::is_void<Result>::value, void>::type
    readResponse(const Callback& callback, size_t bytes_read = 0) {
        assert(bytes_read <= 1);
        if (bytes_read) {
            boost::system::error_code noError;
            callback(noError);
        }
        mSocket.async_read_some(boost::asio::buffer(mCurrentRequest.get(), mCurrSize),
                [this, callback, bytes_read](const boost::system::error_code& ec, size_t br){
                    if (ec) {
                        error<Result>(ec, callback);
                        return;
                    }
                    readResponse<Callback, Result>(callback, bytes_read + br);
                });
    }

    template<class Callback, class Result>
    typename std::enable_if<!std::is_void<Result>::value, void>::type
    readResponse(const Callback& callback, size_t bytes_read = 0) {
        auto respSize = *reinterpret_cast<size_t*>(mCurrentRequest.get());
        if (bytes_read >= 8 && respSize == bytes_read) {
            // response read
            Result res;
            boost::system::error_code noError;
            crossbow::deserializer ser(mCurrentRequest.get() + sizeof(size_t));
            ser & res;
            callback(noError, res);
            return;
        } else if (bytes_read >= 8 && respSize > mCurrSize) {
            std::unique_ptr<uint8_t[]> newBuf(new uint8_t[respSize]);
            memcpy(newBuf.get(), mCurrentRequest.get(), mCurrSize);
            mCurrentRequest.swap(newBuf);
        }
        mSocket.async_read_some(boost::asio::buffer(mCurrentRequest.get() + bytes_read, mCurrSize - bytes_read),
                [this, callback, bytes_read](const boost::system::error_code& ec, size_t br){
                    if (ec) {
                        Result res;
                        callback(ec, res);
                        return;
                    }
                    readResponse<Callback, Result>(callback, bytes_read + br);
                });
    }

    template<class Res, class Callback>
    typename std::enable_if<std::is_void<Res>::value, void>::type
    error(const boost::system::error_code& ec, const Callback& callback) {
        callback(ec);
    }

    template<class Res, class Callback>
    typename std::enable_if<!std::is_void<Res>::value, void>::type
    error(const boost::system::error_code& ec, const Callback& callback) {
        Res res;
        callback(ec, res);
    }

    template<Command C, class Callback, class... Args>
    void execute(const Callback& callback, const Args&... args) {
        static_assert(
                (std::is_void<typename Signature<C>::arguments>::value && std::is_void<argsType<Args...>>::value) ||
                std::is_same<typename Signature<C>::arguments, typename argsType<Args...>::type>::value,
                "Wrong function arguments");
        using ResType = typename Signature<C>::result;
        crossbow::sizer sizer;
        sizer & sizer.size;
        sizer & C;
        impl::ArgSerializer<Args...> argSerializer;
        argSerializer.exec(sizer, args...);
        if (mCurrSize < sizer.size) {
            mCurrentRequest.reset(new uint8_t[sizer.size]);
            mCurrSize = sizer.size;
        }
        crossbow::serializer ser(mCurrentRequest.get());
        ser & sizer.size;
        ser & C;
        argSerializer.exec(ser, args...);
        ser.buffer.release();
        boost::asio::async_write(mSocket, boost::asio::buffer(mCurrentRequest.get(), sizer.size),
                    [this, callback](const boost::system::error_code& ec, size_t){
                        if (ec) {
                            error<ResType>(ec, callback);
                            return;
                        }
                        readResponse<Callback, ResType>(callback);
                    });
    }
};

template<template <typename> class Cmd_Switch,
         class Command,
         template <Command> class Signature,
         class Implementation>
class Server : public Cmd_Switch<Server<Cmd_Switch, Command, Signature, Implementation>> {
    friend struct Cmd_Switch<Server<Cmd_Switch, Command, Signature, Implementation>>;
    Implementation& mImpl;
    boost::asio::ip::tcp::socket& mSocket;
    size_t mBufSize = 1024;
    std::unique_ptr<uint8_t[]> mBuffer;
    using error_code = boost::system::error_code;
    bool doQuit = false;
public:
    Server(Implementation& impl, boost::asio::ip::tcp::socket& socket)
        : mImpl(impl)
        , mSocket(socket)
        , mBuffer(new uint8_t[mBufSize])
    {}
    void run() {
        read();
    }
    void quit() {
        doQuit = true;
    }
private:
    template<Command C, class Callback>
    typename std::enable_if<std::is_void<typename Signature<C>::arguments>::value, void>::type
    execute(Callback callback) {
        mImpl.template execute<C>(callback);
    }

    template<Command C, class Callback>
    typename std::enable_if<!std::is_void<typename Signature<C>::arguments>::value, void>::type
    execute(Callback callback) {
        using Args = typename Signature<C>::arguments;
        Args args;
        crossbow::deserializer des(mBuffer.get() + sizeof(size_t) + sizeof(Command));
        des & args;
        mImpl.template execute<C>(args, callback);
    }

    template<Command C>
    typename std::enable_if<std::is_void<typename Signature<C>::result>::value, void>::type execute() {
        execute<C>([this]() {
            // send the result back
            mBuffer[0] = 1;
            boost::asio::async_write(mSocket,
                    boost::asio::buffer(mBuffer.get(), 1),
                    [this](const error_code& ec, size_t bytes_written) {
                        if (ec) {
                            std::cerr << ec.message() << std::endl;
                            return;
                        }
                        read(0);
                    }
            );
        });
    }

    template<Command C>
    typename std::enable_if<!std::is_void<typename Signature<C>::result>::value, void>::type execute() {
        using Res = typename Signature<C>::result;
        execute<C>([this](const Res& result) {
            // Serialize result
            crossbow::sizer sizer;
            sizer & sizer.size;
            sizer & result;
            if (mBufSize < sizer.size) {
                mBuffer.reset(new uint8_t[sizer.size]);
            }
            crossbow::serializer ser(mBuffer.get());
            ser & sizer.size;
            ser & result;
            ser.buffer.release();
            // send the result back
            boost::asio::async_write(mSocket,
                    boost::asio::buffer(mBuffer.get(), sizer.size),
                    [this](const error_code& ec, size_t bytes_written) {
                        if (ec) {
                            std::cerr << ec.message() << std::endl;
                            mSocket.close();
                            mImpl.close();
                            return;
                        }
                        read(0);
                    }
            );
        });
    }
    void read(size_t bytes_read = 0) {
        if (bytes_read == 0 && doQuit) {
            mSocket.get_io_service().stop();
        }
        size_t reqSize = 0;
        if (bytes_read != 0) {
            reqSize = *reinterpret_cast<size_t*>(mBuffer.get());
        }
        if (bytes_read != 0 && reqSize == bytes_read) {
            // done reading
            auto cmd = *reinterpret_cast<Command*>(mBuffer.get() + sizeof(size_t));
            this->execute_impl(cmd);
            return;
        } else if (bytes_read >= 8 && reqSize > mBufSize) {
            std::unique_ptr<uint8_t[]> newBuf(new uint8_t[reqSize]);
            memcpy(newBuf.get(), mBuffer.get(), mBufSize);
            mBuffer.swap(newBuf);
            mBufSize = reqSize;
        }
        mSocket.async_read_some(boost::asio::buffer(mBuffer.get() + bytes_read, mBufSize - bytes_read),
                [this, bytes_read](const error_code& ec, size_t br){
                    if (ec) {
                        std::cerr << ec.message() << std::endl;
                        mSocket.close();
                        mImpl.close();
                        return;
                    }
                    read(bytes_read + br);
                });
    }
};

} // namespace protocol
} // namespace crossbow

<commit_msg>added some comments to Protocol<commit_after>/*
 * (C) Copyright 2015 ETH Zurich Systems Group (http://www.systems.ethz.ch/) and others.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Contributors:
 *     Markus Pilman <mpilman@inf.ethz.ch>
 *     Simon Loesing <sloesing@inf.ethz.ch>
 *     Thomas Etter <etterth@gmail.com>
 *     Kevin Bocksrocker <kevin.bocksrocker@gmail.com>
 *     Lucas Braun <braunl@inf.ethz.ch>
 */
/*
 * This class serves as a simple wrapper for RPC calls in Boost Asio. You can
 * define different classes of commands by calling GEN_COMMANDS(class-name, command-list)
 * once for each class. After that you can define the signatures for each command
 * defining an input and an output type. Input and Output types can either be primitive
 * or nested types derived from structs or vectors using crossbow::serializable as a
 * serialization method. This serialization method then creates message buffers in the
 * form of:
 * |8 bytes: total buffer size|4 bytes: command-id (>= 1)|command args ...|
 *
 */
#pragma once
#include <tuple>
#include <cstdint>
#include <type_traits>

#include <boost/system/error_code.hpp>
#include <boost/asio.hpp>
#include <boost/preprocessor.hpp>

#include <crossbow/Serializer.hpp>
#include <crossbow/string.hpp>

#define GEN_CASE(r, data, elem)\
    case data::elem:\
        static_cast<Derived*>(this)->template execute<data::elem>();\
        break;

#define SWITCH_CASE(Name, cmd, tuple) switch(cmd) {\
    BOOST_PP_SEQ_FOR_EACH(GEN_CASE, Name, BOOST_PP_TUPLE_TO_SEQ(tuple))\
}

#define GEN_COMMANDS(Name, tuple) enum class Name {\
    BOOST_PP_TUPLE_ELEM(0, tuple) = 1, \
    BOOST_PP_TUPLE_ENUM(BOOST_PP_TUPLE_POP_FRONT(tuple)) \
};\
template<class Derived>\
struct Name ## _Switch {\
    using type = Name;\
    void execute_impl(Name cmd) {\
        SWITCH_CASE(Name, cmd, tuple)\
    }\
}

#define SERVER_TYPE_S(Name, Signature) crossbow::protocol::Server<Name ## _Switch, Name, Signature, Impl>
#define SERVER_TYPE(Name, Impl) crossbow::protocol::Server<Name ## _Switch, Name, Signature, Impl>

namespace crossbow {
namespace protocol {

namespace impl {

template<class... Args>
struct ArgSerializer;

template<class Head, class... Tail>
struct ArgSerializer<Head, Tail...> {
    ArgSerializer<Tail...> rest;

    template<class C>
    void exec(C& c, const Head& head, const Tail&... tail) const {
        c & head;
        rest.exec(c, tail...);
    }
};

template<>
struct ArgSerializer<> {
    template<class C>
    void exec(C&) const {}
};

} // namespace impl

template<class... Args>
struct argsType;
template<class A, class B, class... Tail>
struct argsType<A, B, Tail...> {
    using type = std::tuple<A, B, Tail...>;
};
template<class Arg>
struct argsType<Arg> {
    using type = Arg;
};
template<>
struct argsType<> {
    using type = void;
};

template<class Command, template <Command> class Signature>
class Client {
    boost::asio::ip::tcp::socket& mSocket;
    size_t mCurrSize = 1024;
    std::unique_ptr<uint8_t[]> mCurrentRequest;
public:
    Client(boost::asio::ip::tcp::socket& socket)
        : mSocket(socket), mCurrentRequest(new uint8_t[mCurrSize])
    {
    }

    template<class Callback, class Result>
    typename std::enable_if<std::is_void<Result>::value, void>::type
    readResponse(const Callback& callback, size_t bytes_read = 0) {
        assert(bytes_read <= 1);
        if (bytes_read) {
            boost::system::error_code noError;
            callback(noError);
        }
        mSocket.async_read_some(boost::asio::buffer(mCurrentRequest.get(), mCurrSize),
                [this, callback, bytes_read](const boost::system::error_code& ec, size_t br){
                    if (ec) {
                        error<Result>(ec, callback);
                        return;
                    }
                    readResponse<Callback, Result>(callback, bytes_read + br);
                });
    }

    template<class Callback, class Result>
    typename std::enable_if<!std::is_void<Result>::value, void>::type
    readResponse(const Callback& callback, size_t bytes_read = 0) {
        auto respSize = *reinterpret_cast<size_t*>(mCurrentRequest.get());
        if (bytes_read >= 8 && respSize == bytes_read) {
            // response read
            Result res;
            boost::system::error_code noError;
            crossbow::deserializer ser(mCurrentRequest.get() + sizeof(size_t));
            ser & res;
            callback(noError, res);
            return;
        } else if (bytes_read >= 8 && respSize > mCurrSize) {
            std::unique_ptr<uint8_t[]> newBuf(new uint8_t[respSize]);
            memcpy(newBuf.get(), mCurrentRequest.get(), mCurrSize);
            mCurrentRequest.swap(newBuf);
        }
        mSocket.async_read_some(boost::asio::buffer(mCurrentRequest.get() + bytes_read, mCurrSize - bytes_read),
                [this, callback, bytes_read](const boost::system::error_code& ec, size_t br){
                    if (ec) {
                        Result res;
                        callback(ec, res);
                        return;
                    }
                    readResponse<Callback, Result>(callback, bytes_read + br);
                });
    }

    template<class Res, class Callback>
    typename std::enable_if<std::is_void<Res>::value, void>::type
    error(const boost::system::error_code& ec, const Callback& callback) {
        callback(ec);
    }

    template<class Res, class Callback>
    typename std::enable_if<!std::is_void<Res>::value, void>::type
    error(const boost::system::error_code& ec, const Callback& callback) {
        Res res;
        callback(ec, res);
    }

    template<Command C, class Callback, class... Args>
    void execute(const Callback& callback, const Args&... args) {
        static_assert(
                (std::is_void<typename Signature<C>::arguments>::value && std::is_void<argsType<Args...>>::value) ||
                std::is_same<typename Signature<C>::arguments, typename argsType<Args...>::type>::value,
                "Wrong function arguments");
        using ResType = typename Signature<C>::result;
        crossbow::sizer sizer;
        sizer & sizer.size;
        sizer & C;
        impl::ArgSerializer<Args...> argSerializer;
        argSerializer.exec(sizer, args...);
        if (mCurrSize < sizer.size) {
            mCurrentRequest.reset(new uint8_t[sizer.size]);
            mCurrSize = sizer.size;
        }
        crossbow::serializer ser(mCurrentRequest.get());
        ser & sizer.size;
        ser & C;
        argSerializer.exec(ser, args...);
        ser.buffer.release();
        boost::asio::async_write(mSocket, boost::asio::buffer(mCurrentRequest.get(), sizer.size),
                    [this, callback](const boost::system::error_code& ec, size_t){
                        if (ec) {
                            error<ResType>(ec, callback);
                            return;
                        }
                        readResponse<Callback, ResType>(callback);
                    });
    }
};

template<template <typename> class Cmd_Switch,
         class Command,
         template <Command> class Signature,
         class Implementation>
class Server : public Cmd_Switch<Server<Cmd_Switch, Command, Signature, Implementation>> {
    friend struct Cmd_Switch<Server<Cmd_Switch, Command, Signature, Implementation>>;
    Implementation& mImpl;
    boost::asio::ip::tcp::socket& mSocket;
    size_t mBufSize = 1024;
    std::unique_ptr<uint8_t[]> mBuffer;
    using error_code = boost::system::error_code;
    bool doQuit = false;
public:
    Server(Implementation& impl, boost::asio::ip::tcp::socket& socket)
        : mImpl(impl)
        , mSocket(socket)
        , mBuffer(new uint8_t[mBufSize])
    {}
    void run() {
        read();
    }
    void quit() {
        doQuit = true;
    }
private:
    template<Command C, class Callback>
    typename std::enable_if<std::is_void<typename Signature<C>::arguments>::value, void>::type
    execute(Callback callback) {
        mImpl.template execute<C>(callback);
    }

    template<Command C, class Callback>
    typename std::enable_if<!std::is_void<typename Signature<C>::arguments>::value, void>::type
    execute(Callback callback) {
        using Args = typename Signature<C>::arguments;
        Args args;
        crossbow::deserializer des(mBuffer.get() + sizeof(size_t) + sizeof(Command));
        des & args;
        mImpl.template execute<C>(args, callback);
    }

    template<Command C>
    typename std::enable_if<std::is_void<typename Signature<C>::result>::value, void>::type execute() {
        execute<C>([this]() {
            // send the result back
            mBuffer[0] = 1;
            boost::asio::async_write(mSocket,
                    boost::asio::buffer(mBuffer.get(), 1),
                    [this](const error_code& ec, size_t bytes_written) {
                        if (ec) {
                            std::cerr << ec.message() << std::endl;
                            return;
                        }
                        read(0);
                    }
            );
        });
    }

    template<Command C>
    typename std::enable_if<!std::is_void<typename Signature<C>::result>::value, void>::type execute() {
        using Res = typename Signature<C>::result;
        execute<C>([this](const Res& result) {
            // Serialize result
            crossbow::sizer sizer;
            sizer & sizer.size;
            sizer & result;
            if (mBufSize < sizer.size) {
                mBuffer.reset(new uint8_t[sizer.size]);
            }
            crossbow::serializer ser(mBuffer.get());
            ser & sizer.size;
            ser & result;
            ser.buffer.release();
            // send the result back
            boost::asio::async_write(mSocket,
                    boost::asio::buffer(mBuffer.get(), sizer.size),
                    [this](const error_code& ec, size_t bytes_written) {
                        if (ec) {
                            std::cerr << ec.message() << std::endl;
                            mSocket.close();
                            mImpl.close();
                            return;
                        }
                        read(0);
                    }
            );
        });
    }
    void read(size_t bytes_read = 0) {
        if (bytes_read == 0 && doQuit) {
            mSocket.get_io_service().stop();
        }
        size_t reqSize = 0;
        if (bytes_read != 0) {
            reqSize = *reinterpret_cast<size_t*>(mBuffer.get());
        }
        if (bytes_read != 0 && reqSize == bytes_read) {
            // done reading
            auto cmd = *reinterpret_cast<Command*>(mBuffer.get() + sizeof(size_t));
            this->execute_impl(cmd);
            return;
        } else if (bytes_read >= 8 && reqSize > mBufSize) {
            std::unique_ptr<uint8_t[]> newBuf(new uint8_t[reqSize]);
            memcpy(newBuf.get(), mBuffer.get(), mBufSize);
            mBuffer.swap(newBuf);
            mBufSize = reqSize;
        }
        mSocket.async_read_some(boost::asio::buffer(mBuffer.get() + bytes_read, mBufSize - bytes_read),
                [this, bytes_read](const error_code& ec, size_t br){
                    if (ec) {
                        std::cerr << ec.message() << std::endl;
                        mSocket.close();
                        mImpl.close();
                        return;
                    }
                    read(bytes_read + br);
                });
    }
};

} // namespace protocol
} // namespace crossbow

<|endoftext|>


<commit_before>// Copyright (c) 2006-2009 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.

// This test is POSIX only.

#include <unistd.h>
#include <fcntl.h>

#include "base/basictypes.h"
#include "base/eintr_wrapper.h"
#include "chrome/common/file_descriptor_set_posix.h"
#include "testing/gtest/include/gtest/gtest.h"

// The FileDescriptorSet will try and close some of the descriptor numbers
// which we given it. This is the base descriptor value. It's great enough such
// that no real descriptor will accidently be closed.
static const int kFDBase = 50000;

TEST(FileDescriptorSet, BasicAdd) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_EQ(set->size(), 0u);
  ASSERT_TRUE(set->empty());
  ASSERT_TRUE(set->Add(kFDBase));
  ASSERT_EQ(set->size(), 1u);
  ASSERT_TRUE(!set->empty());

  // Empties the set and stops a warning about deleting a set with unconsumed
  // descriptors
  set->CommitAll();
}

TEST(FileDescriptorSet, BasicAddAndClose) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_EQ(set->size(), 0u);
  ASSERT_TRUE(set->empty());
  ASSERT_TRUE(set->AddAndAutoClose(kFDBase));
  ASSERT_EQ(set->size(), 1u);
  ASSERT_TRUE(!set->empty());

  set->CommitAll();
}

TEST(FileDescriptorSet, MaxSize) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  for (unsigned i = 0;
       i < FileDescriptorSet::MAX_DESCRIPTORS_PER_MESSAGE; ++i) {
    ASSERT_TRUE(set->Add(kFDBase + 1 + i));
  }

  ASSERT_TRUE(!set->Add(kFDBase));

  set->CommitAll();
}

TEST(FileDescriptorSet, SetDescriptors) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_TRUE(set->empty());
  set->SetDescriptors(NULL, 0);
  ASSERT_TRUE(set->empty());

  static const int fds[] = {kFDBase};
  set->SetDescriptors(fds, 1);
  ASSERT_TRUE(!set->empty());
  ASSERT_EQ(set->size(), 1u);

  set->CommitAll();
}

TEST(FileDescriptorSet, GetDescriptors) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  set->GetDescriptors(NULL);
  ASSERT_TRUE(set->Add(kFDBase));

  int fds[1];
  fds[0] = 0;
  set->GetDescriptors(fds);
  ASSERT_EQ(fds[0], kFDBase);
  set->CommitAll();
  ASSERT_TRUE(set->empty());
}

TEST(FileDescriptorSet, WalkInOrder) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_TRUE(set->Add(kFDBase));
  ASSERT_TRUE(set->Add(kFDBase + 1));
  ASSERT_TRUE(set->Add(kFDBase + 2));

  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(1), kFDBase + 1);
  ASSERT_EQ(set->GetDescriptorAt(2), kFDBase + 2);

  set->CommitAll();
}

TEST(FileDescriptorSet, WalkWrongOrder) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_TRUE(set->Add(kFDBase));
  ASSERT_TRUE(set->Add(kFDBase + 1));
  ASSERT_TRUE(set->Add(kFDBase + 2));

  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(2), -1);

  set->CommitAll();
}

TEST(FileDescriptorSet, WalkCycle) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_TRUE(set->Add(kFDBase));
  ASSERT_TRUE(set->Add(kFDBase + 1));
  ASSERT_TRUE(set->Add(kFDBase + 2));

  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(1), kFDBase + 1);
  ASSERT_EQ(set->GetDescriptorAt(2), kFDBase + 2);
  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(1), kFDBase + 1);
  ASSERT_EQ(set->GetDescriptorAt(2), kFDBase + 2);
  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(1), kFDBase + 1);
  ASSERT_EQ(set->GetDescriptorAt(2), kFDBase + 2);

  set->CommitAll();
}

TEST(FileDescriptorSet, DontClose) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  const int fd = open("/dev/null", O_RDONLY);
  ASSERT_TRUE(set->Add(fd));
  set->CommitAll();

  const int duped = dup(fd);
  ASSERT_GE(duped, 0);
  HANDLE_EINTR(close(duped));
  HANDLE_EINTR(close(fd));
}

TEST(FileDescriptorSet, DoClose) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  const int fd = open("/dev/null", O_RDONLY);
  ASSERT_TRUE(set->AddAndAutoClose(fd));
  set->CommitAll();

  const int duped = dup(fd);
  ASSERT_EQ(duped, -1);
  HANDLE_EINTR(close(fd));
}
<commit_msg>Fix valigrind complaint about invalid file descriptor on close. Fixes FileDescriptorSet.BasicAddAndClose and FileDescriptorSet.SetDescriptors.<commit_after>// Copyright (c) 2006-2009 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.

// This test is POSIX only.

#include <unistd.h>
#include <fcntl.h>

#include "base/basictypes.h"
#include "base/eintr_wrapper.h"
#include "chrome/common/file_descriptor_set_posix.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace {

// Get a safe file descriptor for test purposes.
int GetSafeFd() {
  return open("/dev/null", O_RDONLY);
}

// Returns true if fd was already closed.  Closes fd if not closed.
bool VerifyClosed(int fd) {
  const int duped = dup(fd);
  if (duped != -1) {
    HANDLE_EINTR(close(duped));
    HANDLE_EINTR(close(fd));
    return false;
  }
  return true;
}

// The FileDescriptorSet will try and close some of the descriptor numbers
// which we given it. This is the base descriptor value. It's great enough such
// that no real descriptor will accidently be closed.
static const int kFDBase = 50000;

TEST(FileDescriptorSet, BasicAdd) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_EQ(set->size(), 0u);
  ASSERT_TRUE(set->empty());
  ASSERT_TRUE(set->Add(kFDBase));
  ASSERT_EQ(set->size(), 1u);
  ASSERT_TRUE(!set->empty());

  // Empties the set and stops a warning about deleting a set with unconsumed
  // descriptors
  set->CommitAll();
}

TEST(FileDescriptorSet, BasicAddAndClose) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_EQ(set->size(), 0u);
  ASSERT_TRUE(set->empty());
  const int fd = GetSafeFd();
  ASSERT_TRUE(set->AddAndAutoClose(fd));
  ASSERT_EQ(set->size(), 1u);
  ASSERT_TRUE(!set->empty());

  set->CommitAll();

  ASSERT_TRUE(VerifyClosed(fd));
}
TEST(FileDescriptorSet, MaxSize) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  for (unsigned i = 0;
       i < FileDescriptorSet::MAX_DESCRIPTORS_PER_MESSAGE; ++i) {
    ASSERT_TRUE(set->Add(kFDBase + 1 + i));
  }

  ASSERT_TRUE(!set->Add(kFDBase));

  set->CommitAll();
}

TEST(FileDescriptorSet, SetDescriptors) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_TRUE(set->empty());
  set->SetDescriptors(NULL, 0);
  ASSERT_TRUE(set->empty());

  const int fd = GetSafeFd();
  static const int fds[] = {fd};
  set->SetDescriptors(fds, 1);
  ASSERT_TRUE(!set->empty());
  ASSERT_EQ(set->size(), 1u);

  set->CommitAll();

  ASSERT_TRUE(VerifyClosed(fd));
}

TEST(FileDescriptorSet, GetDescriptors) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  set->GetDescriptors(NULL);
  ASSERT_TRUE(set->Add(kFDBase));

  int fds[1];
  fds[0] = 0;
  set->GetDescriptors(fds);
  ASSERT_EQ(fds[0], kFDBase);
  set->CommitAll();
  ASSERT_TRUE(set->empty());
}

TEST(FileDescriptorSet, WalkInOrder) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_TRUE(set->Add(kFDBase));
  ASSERT_TRUE(set->Add(kFDBase + 1));
  ASSERT_TRUE(set->Add(kFDBase + 2));

  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(1), kFDBase + 1);
  ASSERT_EQ(set->GetDescriptorAt(2), kFDBase + 2);

  set->CommitAll();
}

TEST(FileDescriptorSet, WalkWrongOrder) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_TRUE(set->Add(kFDBase));
  ASSERT_TRUE(set->Add(kFDBase + 1));
  ASSERT_TRUE(set->Add(kFDBase + 2));

  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(2), -1);

  set->CommitAll();
}

TEST(FileDescriptorSet, WalkCycle) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  ASSERT_TRUE(set->Add(kFDBase));
  ASSERT_TRUE(set->Add(kFDBase + 1));
  ASSERT_TRUE(set->Add(kFDBase + 2));

  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(1), kFDBase + 1);
  ASSERT_EQ(set->GetDescriptorAt(2), kFDBase + 2);
  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(1), kFDBase + 1);
  ASSERT_EQ(set->GetDescriptorAt(2), kFDBase + 2);
  ASSERT_EQ(set->GetDescriptorAt(0), kFDBase);
  ASSERT_EQ(set->GetDescriptorAt(1), kFDBase + 1);
  ASSERT_EQ(set->GetDescriptorAt(2), kFDBase + 2);

  set->CommitAll();
}

TEST(FileDescriptorSet, DontClose) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  const int fd = GetSafeFd();
  ASSERT_TRUE(set->Add(fd));
  set->CommitAll();

  ASSERT_FALSE(VerifyClosed(fd));
}

TEST(FileDescriptorSet, DoClose) {
  scoped_refptr<FileDescriptorSet> set = new FileDescriptorSet;

  const int fd = GetSafeFd();
  ASSERT_TRUE(set->AddAndAutoClose(fd));
  set->CommitAll();

  ASSERT_TRUE(VerifyClosed(fd));
}

}  // namespace
<|endoftext|>


<commit_before>#include <iostream>
#include <ncurses.h>
#include <string>
#include "maze.h"

using namespace std;

int main()
{
    const int MAX_SIZE = 71;
    bool maze[MAX_SIZE * MAX_SIZE];
    int nrows = 19;
    int ncols = 31;
    int start[2] = {1, 1};
    int finish[2] = {1, 1};
    char c;

    initscr();

    while (true) {
        // generate a new maze
        backtracking_maze_gen(maze, MAX_SIZE, nrows, ncols);
        gen_entrances_opposites(start, finish, nrows, ncols);
        maze_print(maze, MAX_SIZE, nrows, ncols, start, finish);

        // parse user input
        c = getch();
        if (c == 'q'){
            break;
        }

        // increase the size of the maze (to a limit)
        if (ncols < MAX_SIZE) {
            ncols += 2;
        }
        if (nrows < (MAX_SIZE / 2)) {
            nrows += 2;
        }
    }
    endwin();

    return 0;
}
<commit_msg>unused import<commit_after>#include <iostream>
#include <ncurses.h>
#include "maze.h"

using namespace std;

int main()
{
    const int MAX_SIZE = 71;
    bool maze[MAX_SIZE * MAX_SIZE];
    int nrows = 19;
    int ncols = 31;
    int start[2] = {1, 1};
    int finish[2] = {1, 1};
    char c;

    initscr();

    while (true) {
        // generate a new maze
        backtracking_maze_gen(maze, MAX_SIZE, nrows, ncols);
        gen_entrances_opposites(start, finish, nrows, ncols);
        maze_print(maze, MAX_SIZE, nrows, ncols, start, finish);

        // parse user input
        c = getch();
        if (c == 'q'){
            break;
        }

        // increase the size of the maze (to a limit)
        if (ncols < MAX_SIZE) {
            ncols += 2;
        }
        if (nrows < (MAX_SIZE / 2)) {
            nrows += 2;
        }
    }
    endwin();

    return 0;
}
<|endoftext|>


<commit_before>// Copyright (c) 2008, Google Inc.
// All rights reserved.
//
// 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.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// 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 COPYRIGHT
// OWNER 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 "client/windows/crash_generation/minidump_generator.h"
#include <cassert>
#include "client/windows/common/auto_critical_section.h"
#include "common/windows/guid_string.h"

using std::wstring;

namespace google_breakpad {

MinidumpGenerator::MinidumpGenerator(const wstring& dump_path)
    : dbghelp_module_(NULL),
      rpcrt4_module_(NULL),
      dump_path_(dump_path),
      write_dump_(NULL),
      create_uuid_(NULL) {
  InitializeCriticalSection(&module_load_sync_);
  InitializeCriticalSection(&get_proc_address_sync_);
}

MinidumpGenerator::~MinidumpGenerator() {
  if (dbghelp_module_) {
    FreeLibrary(dbghelp_module_);
  }

  if (rpcrt4_module_) {
    FreeLibrary(rpcrt4_module_);
  }

  DeleteCriticalSection(&get_proc_address_sync_);
  DeleteCriticalSection(&module_load_sync_);
}

bool MinidumpGenerator::WriteMinidump(HANDLE process_handle,
                                      DWORD process_id,
                                      DWORD thread_id,
                                      DWORD requesting_thread_id,
                                      EXCEPTION_POINTERS* exception_pointers,
                                      MDRawAssertionInfo* assert_info,
                                      MINIDUMP_TYPE dump_type,
                                      bool is_client_pointers,
                                      wstring* dump_path) {
  // Just call the full WriteMinidump with NULL as the full_dump_path.
  return this->WriteMinidump(process_handle, process_id, thread_id,
                             requesting_thread_id, exception_pointers,
                             assert_info, dump_type, is_client_pointers,
                             dump_path, NULL);
}

bool MinidumpGenerator::WriteMinidump(HANDLE process_handle,
                                      DWORD process_id,
                                      DWORD thread_id,
                                      DWORD requesting_thread_id,
                                      EXCEPTION_POINTERS* exception_pointers,
                                      MDRawAssertionInfo* assert_info,
                                      MINIDUMP_TYPE dump_type,
                                      bool is_client_pointers,
                                      wstring* dump_path,
                                      wstring* full_dump_path) {
  MiniDumpWriteDumpType write_dump = GetWriteDump();
  if (!write_dump) {
    return false;
  }

  wstring dump_file_path;
  if (!GenerateDumpFilePath(&dump_file_path)) {
    return false;
  }

  // If the client requests a full memory dump, we will write a normal mini
  // dump and a full memory dump. Both dump files use the same uuid as file
  // name prefix.
  bool full_memory_dump = (dump_type & MiniDumpWithFullMemory) != 0;
  wstring full_dump_file_path;
  if (full_memory_dump) {
    full_dump_file_path.assign(dump_file_path);
    full_dump_file_path.resize(full_dump_file_path.size() - 4);  // strip .dmp
    full_dump_file_path.append(TEXT("-full.dmp"));
  }

  HANDLE dump_file = CreateFile(dump_file_path.c_str(),
                                GENERIC_WRITE,
                                0,
                                NULL,
                                CREATE_NEW,
                                FILE_ATTRIBUTE_NORMAL,
                                NULL);

  if (dump_file == INVALID_HANDLE_VALUE) {
    return false;
  }

  HANDLE full_dump_file = INVALID_HANDLE_VALUE;
  if (full_memory_dump) {
    full_dump_file = CreateFile(full_dump_file_path.c_str(),
                                GENERIC_WRITE,
                                0,
                                NULL,
                                CREATE_NEW,
                                FILE_ATTRIBUTE_NORMAL,
                                NULL);

    if (full_dump_file == INVALID_HANDLE_VALUE) {
      CloseHandle(dump_file);
      return false;
    }
  }

  MINIDUMP_EXCEPTION_INFORMATION* dump_exception_pointers = NULL;
  MINIDUMP_EXCEPTION_INFORMATION dump_exception_info;

  // Setup the exception information object only if it's a dump
  // due to an exception.
  if (exception_pointers) {
    dump_exception_pointers = &dump_exception_info;
    dump_exception_info.ThreadId = thread_id;
    dump_exception_info.ExceptionPointers = exception_pointers;
    dump_exception_info.ClientPointers = is_client_pointers;
  }

  // Add an MDRawBreakpadInfo stream to the minidump, to provide additional
  // information about the exception handler to the Breakpad processor.
  // The information will help the processor determine which threads are
  // relevant. The Breakpad processor does not require this information but
  // can function better with Breakpad-generated dumps when it is present.
  // The native debugger is not harmed by the presence of this information.
  MDRawBreakpadInfo breakpad_info = {0};
  if (!is_client_pointers) {
    // Set the dump thread id and requesting thread id only in case of
    // in-process dump generation.
    breakpad_info.validity = MD_BREAKPAD_INFO_VALID_DUMP_THREAD_ID |
                             MD_BREAKPAD_INFO_VALID_REQUESTING_THREAD_ID;
    breakpad_info.dump_thread_id = thread_id;
    breakpad_info.requesting_thread_id = requesting_thread_id;
  }

  // Leave room in user_stream_array for a possible assertion info stream.
  MINIDUMP_USER_STREAM user_stream_array[2];
  user_stream_array[0].Type = MD_BREAKPAD_INFO_STREAM;
  user_stream_array[0].BufferSize = sizeof(breakpad_info);
  user_stream_array[0].Buffer = &breakpad_info;

  MINIDUMP_USER_STREAM_INFORMATION user_streams;
  user_streams.UserStreamCount = 1;
  user_streams.UserStreamArray = user_stream_array;

  MDRawAssertionInfo* actual_assert_info = assert_info;
  MDRawAssertionInfo client_assert_info = {0};

  if (assert_info) {
    // If the assertion info object lives in the client process,
    // read the memory of the client process.
    if (is_client_pointers) {
      SIZE_T bytes_read = 0;
      if (!ReadProcessMemory(process_handle,
                             assert_info,
                             &client_assert_info,
                             sizeof(client_assert_info),
                             &bytes_read)) {
        CloseHandle(dump_file);
        if (full_dump_file != INVALID_HANDLE_VALUE)
          CloseHandle(full_dump_file);
        return false;
      }

      if (bytes_read != sizeof(client_assert_info)) {
        CloseHandle(dump_file);
        if (full_dump_file != INVALID_HANDLE_VALUE)
          CloseHandle(full_dump_file);
        return false;
      }

      actual_assert_info  = &client_assert_info;
    }

    user_stream_array[1].Type = MD_ASSERTION_INFO_STREAM;
    user_stream_array[1].BufferSize = sizeof(MDRawAssertionInfo);
    user_stream_array[1].Buffer = actual_assert_info;
    ++user_streams.UserStreamCount;
  }

  bool result_minidump = write_dump(
      process_handle,
      process_id,
      dump_file,
      static_cast<MINIDUMP_TYPE>((dump_type & (~MiniDumpWithFullMemory))
                                  | MiniDumpNormal),
      exception_pointers ? &dump_exception_info : NULL,
      &user_streams,
      NULL) != FALSE;

  bool result_full_memory = true;
  if (full_memory_dump) {
    result_full_memory = write_dump(
        process_handle,
        process_id,
        full_dump_file,
        static_cast<MINIDUMP_TYPE>(dump_type & (~MiniDumpNormal)),
        exception_pointers ? &dump_exception_info : NULL,
        &user_streams,
        NULL) != FALSE;
  }

  bool result = result_minidump && result_full_memory;

  CloseHandle(dump_file);
  if (full_dump_file != INVALID_HANDLE_VALUE)
    CloseHandle(full_dump_file);

  // Store the path of the dump file in the out parameter if dump generation
  // succeeded.
  if (result && dump_path) {
    *dump_path = dump_file_path;
  }
  if (result && full_memory_dump && full_dump_path) {
    *full_dump_path = full_dump_file_path;
  }

  return result;
}

HMODULE MinidumpGenerator::GetDbghelpModule() {
  AutoCriticalSection lock(&module_load_sync_);
  if (!dbghelp_module_) {
    dbghelp_module_ = LoadLibrary(TEXT("dbghelp.dll"));
  }

  return dbghelp_module_;
}

MinidumpGenerator::MiniDumpWriteDumpType MinidumpGenerator::GetWriteDump() {
  AutoCriticalSection lock(&get_proc_address_sync_);
  if (!write_dump_) {
    HMODULE module = GetDbghelpModule();
    if (module) {
      FARPROC proc = GetProcAddress(module, "MiniDumpWriteDump");
      write_dump_ = reinterpret_cast<MiniDumpWriteDumpType>(proc);
    }
  }

  return write_dump_;
}

HMODULE MinidumpGenerator::GetRpcrt4Module() {
  AutoCriticalSection lock(&module_load_sync_);
  if (!rpcrt4_module_) {
    rpcrt4_module_ = LoadLibrary(TEXT("rpcrt4.dll"));
  }

  return rpcrt4_module_;
}

MinidumpGenerator::UuidCreateType MinidumpGenerator::GetCreateUuid() {
  AutoCriticalSection lock(&module_load_sync_);
  if (!create_uuid_) {
    HMODULE module = GetRpcrt4Module();
    if (module) {
      FARPROC proc = GetProcAddress(module, "UuidCreate");
      create_uuid_ = reinterpret_cast<UuidCreateType>(proc);
    }
  }

  return create_uuid_;
}

bool MinidumpGenerator::GenerateDumpFilePath(wstring* file_path) {
  UUID id = {0};

  UuidCreateType create_uuid = GetCreateUuid();
  if (!create_uuid) {
    return false;
  }

  create_uuid(&id);
  wstring id_str = GUIDString::GUIDToWString(&id);

  *file_path = dump_path_ + TEXT("\\") + id_str + TEXT(".dmp");
  return true;
}

}  // namespace google_breakpad
<commit_msg>Allow the minidump writer to collect handle data so that resulting dump contains information about opened handles (!handle) and handle operations trace (!htrace).<commit_after>// Copyright (c) 2008, Google Inc.
// All rights reserved.
//
// 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.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// 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 COPYRIGHT
// OWNER 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 "client/windows/crash_generation/minidump_generator.h"

#include <assert.h>
#include <avrfsdk.h>

#include <algorithm>
#include <list>
#include <vector>

#include "client/windows/common/auto_critical_section.h"
#include "common/windows/guid_string.h"

using std::wstring;

// Disable C4996: 'std::copy': Function call with parameters that may be unsafe.
#pragma warning( disable : 4996 )

namespace {

// A helper class used to collect handle operations data. Unlike
// |MiniDumpWithHandleData| it records the operations for a single handle value
// only, making it possible to include this information to a minidump.
class HandleTraceData {
 public:
  HandleTraceData();
  ~HandleTraceData();

  // Collects the handle operations data and formats a user stream to be added
  // to the minidump.
  bool CollectHandleData(HANDLE process_handle,
                         EXCEPTION_POINTERS* exception_pointers);

  // Fills the user dump entry with a pointer to the collected handle operations
  // data. Returns |true| if the entry was initialized successfully, or |false|
  // if no trace data is available.
  bool GetUserStream(MINIDUMP_USER_STREAM* user_stream);

 private:
  // Reads the exception code from the client process's address space.
  // This routine assumes that the client process's pointer width matches ours.
  static bool ReadExceptionCode(HANDLE process_handle,
                                EXCEPTION_POINTERS* exception_pointers,
                                DWORD* exception_code);

  // Stores handle operations retrieved by VerifierEnumerateResource().
  static ULONG CALLBACK RecordHandleOperations(void* resource_description,
                                               void* enumeration_context,
                                               ULONG* enumeration_level);

  // Function pointer type for VerifierEnumerateResource, which is looked up
  // dynamically.
  typedef BOOL (WINAPI* VerifierEnumerateResourceType)(
      HANDLE Process,
      ULONG Flags,
      ULONG ResourceType,
      AVRF_RESOURCE_ENUMERATE_CALLBACK ResourceCallback,
      PVOID EnumerationContext);

  // Handle to dynamically loaded verifier.dll.
  HMODULE verifier_module_;

  // Pointer to the VerifierEnumerateResource function.
  VerifierEnumerateResourceType enumerate_resource_;

  // Handle value to look for.
  ULONG64 handle_;

  // List of handle operations for |handle_|.
  std::list<AVRF_HANDLE_OPERATION> operations_;

  // Minidump stream data.
  std::vector<char> stream_;
};

HandleTraceData::HandleTraceData()
    : verifier_module_(NULL),
      enumerate_resource_(NULL),
      handle_(NULL) {
}

HandleTraceData::~HandleTraceData() {
  if (verifier_module_) {
    FreeLibrary(verifier_module_);
  }
}

bool HandleTraceData::CollectHandleData(
    HANDLE process_handle,
    EXCEPTION_POINTERS* exception_pointers) {
  DWORD exception_code;
  if (!ReadExceptionCode(process_handle, exception_pointers, &exception_code)) {
    return false;
  }

  // Verify whether the execption is STATUS_INVALID_HANDLE. Do not record any
  // handle information if it is a different exception to keep the minidump
  // small.
  if (exception_code != STATUS_INVALID_HANDLE) {
    return true;
  }

  // Load verifier!VerifierEnumerateResource() dynamically.
  verifier_module_ = LoadLibrary(TEXT("verifier.dll"));
  if (!verifier_module_) {
    return false;
  }

  enumerate_resource_ = reinterpret_cast<VerifierEnumerateResourceType>(
      GetProcAddress(verifier_module_, "VerifierEnumerateResource"));
  if (!enumerate_resource_) {
    return false;
  }

  // STATUS_INVALID_HANDLE does not provide the offending handle value in
  // the exception parameters so we have to guess. At the moment we scan
  // the handle operations trace looking for the last invalid handle operation
  // and record only the operations for that handle value.
  if (enumerate_resource_(process_handle,
                          0,
                          AvrfResourceHandleTrace,
                          &RecordHandleOperations,
                          this) != ERROR_SUCCESS) {
    return false;
  }

  // Now that |handle_| is initialized, purge all irrelevant operations.
  std::list<AVRF_HANDLE_OPERATION>::iterator i = operations_.begin();
  std::list<AVRF_HANDLE_OPERATION>::iterator i_end = operations_.end();
  while (i != i_end) {
    if (i->Handle == handle_) {
      ++i;
    } else {
      i = operations_.erase(i);
    }
  }

  // Convert the list of recorded operations to a minidump stream.
  stream_.resize(sizeof(MINIDUMP_HANDLE_OPERATION_LIST) +
      sizeof(AVRF_HANDLE_OPERATION) * operations_.size());

  MINIDUMP_HANDLE_OPERATION_LIST* stream_data =
      reinterpret_cast<MINIDUMP_HANDLE_OPERATION_LIST*>(
          &stream_.front());
  stream_data->SizeOfHeader = sizeof(MINIDUMP_HANDLE_OPERATION_LIST);
  stream_data->SizeOfEntry = sizeof(AVRF_HANDLE_OPERATION);
  stream_data->NumberOfEntries = static_cast<ULONG32>(operations_.size());
  stream_data->Reserved = 0;
  std::copy(operations_.begin(),
            operations_.end(),
            reinterpret_cast<AVRF_HANDLE_OPERATION*>(stream_data + 1));

  return true;
}

bool HandleTraceData::GetUserStream(MINIDUMP_USER_STREAM* user_stream) {
  if (stream_.empty()) {
    return false;
  } else {
    user_stream->Type = HandleOperationListStream;
    user_stream->BufferSize = static_cast<ULONG>(stream_.size());
    user_stream->Buffer = &stream_.front();
    return true;
  }
}

bool HandleTraceData::ReadExceptionCode(
    HANDLE process_handle,
    EXCEPTION_POINTERS* exception_pointers,
    DWORD* exception_code) {
  EXCEPTION_POINTERS pointers;
  if (!ReadProcessMemory(process_handle,
                         exception_pointers,
                         &pointers,
                         sizeof(pointers),
                         NULL)) {
    return false;
  }

  if (!ReadProcessMemory(process_handle,
                         pointers.ExceptionRecord,
                         exception_code,
                         sizeof(*exception_code),
                         NULL)) {
    return false;
  }

  return true;
}

ULONG CALLBACK HandleTraceData::RecordHandleOperations(
    void* resource_description,
    void* enumeration_context,
    ULONG* enumeration_level) {
  AVRF_HANDLE_OPERATION* description =
      reinterpret_cast<AVRF_HANDLE_OPERATION*>(resource_description);
  HandleTraceData* self =
      reinterpret_cast<HandleTraceData*>(enumeration_context);

  // Remember the last invalid handle operation.
  if (description->OperationType == OperationDbBADREF) {
    self->handle_ = description->Handle;
  }

  // Record all handle operations.
  self->operations_.push_back(*description);

  *enumeration_level = HeapEnumerationEverything;
  return ERROR_SUCCESS;
}

}  // namespace

namespace google_breakpad {

MinidumpGenerator::MinidumpGenerator(const wstring& dump_path)
    : dbghelp_module_(NULL),
      rpcrt4_module_(NULL),
      dump_path_(dump_path),
      write_dump_(NULL),
      create_uuid_(NULL) {
  InitializeCriticalSection(&module_load_sync_);
  InitializeCriticalSection(&get_proc_address_sync_);
}

MinidumpGenerator::~MinidumpGenerator() {
  if (dbghelp_module_) {
    FreeLibrary(dbghelp_module_);
  }

  if (rpcrt4_module_) {
    FreeLibrary(rpcrt4_module_);
  }

  DeleteCriticalSection(&get_proc_address_sync_);
  DeleteCriticalSection(&module_load_sync_);
}

bool MinidumpGenerator::WriteMinidump(HANDLE process_handle,
                                      DWORD process_id,
                                      DWORD thread_id,
                                      DWORD requesting_thread_id,
                                      EXCEPTION_POINTERS* exception_pointers,
                                      MDRawAssertionInfo* assert_info,
                                      MINIDUMP_TYPE dump_type,
                                      bool is_client_pointers,
                                      wstring* dump_path) {
  // Just call the full WriteMinidump with NULL as the full_dump_path.
  return this->WriteMinidump(process_handle, process_id, thread_id,
                             requesting_thread_id, exception_pointers,
                             assert_info, dump_type, is_client_pointers,
                             dump_path, NULL);
}

bool MinidumpGenerator::WriteMinidump(HANDLE process_handle,
                                      DWORD process_id,
                                      DWORD thread_id,
                                      DWORD requesting_thread_id,
                                      EXCEPTION_POINTERS* exception_pointers,
                                      MDRawAssertionInfo* assert_info,
                                      MINIDUMP_TYPE dump_type,
                                      bool is_client_pointers,
                                      wstring* dump_path,
                                      wstring* full_dump_path) {
  MiniDumpWriteDumpType write_dump = GetWriteDump();
  if (!write_dump) {
    return false;
  }

  wstring dump_file_path;
  if (!GenerateDumpFilePath(&dump_file_path)) {
    return false;
  }

  // If the client requests a full memory dump, we will write a normal mini
  // dump and a full memory dump. Both dump files use the same uuid as file
  // name prefix.
  bool full_memory_dump = (dump_type & MiniDumpWithFullMemory) != 0;
  wstring full_dump_file_path;
  if (full_memory_dump) {
    full_dump_file_path.assign(dump_file_path);
    full_dump_file_path.resize(full_dump_file_path.size() - 4);  // strip .dmp
    full_dump_file_path.append(TEXT("-full.dmp"));
  }

  HANDLE dump_file = CreateFile(dump_file_path.c_str(),
                                GENERIC_WRITE,
                                0,
                                NULL,
                                CREATE_NEW,
                                FILE_ATTRIBUTE_NORMAL,
                                NULL);

  if (dump_file == INVALID_HANDLE_VALUE) {
    return false;
  }

  HANDLE full_dump_file = INVALID_HANDLE_VALUE;
  if (full_memory_dump) {
    full_dump_file = CreateFile(full_dump_file_path.c_str(),
                                GENERIC_WRITE,
                                0,
                                NULL,
                                CREATE_NEW,
                                FILE_ATTRIBUTE_NORMAL,
                                NULL);

    if (full_dump_file == INVALID_HANDLE_VALUE) {
      CloseHandle(dump_file);
      return false;
    }
  }

  MINIDUMP_EXCEPTION_INFORMATION* dump_exception_pointers = NULL;
  MINIDUMP_EXCEPTION_INFORMATION dump_exception_info;

  // Setup the exception information object only if it's a dump
  // due to an exception.
  if (exception_pointers) {
    dump_exception_pointers = &dump_exception_info;
    dump_exception_info.ThreadId = thread_id;
    dump_exception_info.ExceptionPointers = exception_pointers;
    dump_exception_info.ClientPointers = is_client_pointers;
  }

  // Add an MDRawBreakpadInfo stream to the minidump, to provide additional
  // information about the exception handler to the Breakpad processor.
  // The information will help the processor determine which threads are
  // relevant. The Breakpad processor does not require this information but
  // can function better with Breakpad-generated dumps when it is present.
  // The native debugger is not harmed by the presence of this information.
  MDRawBreakpadInfo breakpad_info = {0};
  if (!is_client_pointers) {
    // Set the dump thread id and requesting thread id only in case of
    // in-process dump generation.
    breakpad_info.validity = MD_BREAKPAD_INFO_VALID_DUMP_THREAD_ID |
                             MD_BREAKPAD_INFO_VALID_REQUESTING_THREAD_ID;
    breakpad_info.dump_thread_id = thread_id;
    breakpad_info.requesting_thread_id = requesting_thread_id;
  }

  // Leave room in user_stream_array for possible assertion info and handle
  // operations streams.
  MINIDUMP_USER_STREAM user_stream_array[3];
  user_stream_array[0].Type = MD_BREAKPAD_INFO_STREAM;
  user_stream_array[0].BufferSize = sizeof(breakpad_info);
  user_stream_array[0].Buffer = &breakpad_info;

  MINIDUMP_USER_STREAM_INFORMATION user_streams;
  user_streams.UserStreamCount = 1;
  user_streams.UserStreamArray = user_stream_array;

  MDRawAssertionInfo* actual_assert_info = assert_info;
  MDRawAssertionInfo client_assert_info = {0};

  if (assert_info) {
    // If the assertion info object lives in the client process,
    // read the memory of the client process.
    if (is_client_pointers) {
      SIZE_T bytes_read = 0;
      if (!ReadProcessMemory(process_handle,
                             assert_info,
                             &client_assert_info,
                             sizeof(client_assert_info),
                             &bytes_read)) {
        CloseHandle(dump_file);
        if (full_dump_file != INVALID_HANDLE_VALUE)
          CloseHandle(full_dump_file);
        return false;
      }

      if (bytes_read != sizeof(client_assert_info)) {
        CloseHandle(dump_file);
        if (full_dump_file != INVALID_HANDLE_VALUE)
          CloseHandle(full_dump_file);
        return false;
      }

      actual_assert_info  = &client_assert_info;
    }

    user_stream_array[1].Type = MD_ASSERTION_INFO_STREAM;
    user_stream_array[1].BufferSize = sizeof(MDRawAssertionInfo);
    user_stream_array[1].Buffer = actual_assert_info;
    ++user_streams.UserStreamCount;
  }

  // If the process is terminated by STATUS_INVALID_HANDLE exception store
  // the trace of operatios for the offending handle value. Do nothing special
  // if the client already requested the handle trace to be stored in the dump.
  HandleTraceData handle_trace_data;
  if (exception_pointers && (dump_type & MiniDumpWithHandleData) == 0) {
    if (!handle_trace_data.CollectHandleData(process_handle,
                                             exception_pointers)) {
      CloseHandle(dump_file);
      if (full_dump_file != INVALID_HANDLE_VALUE)
        CloseHandle(full_dump_file);
      return false;
    }
  }

  bool result_full_memory = true;
  if (full_memory_dump) {
    result_full_memory = write_dump(
        process_handle,
        process_id,
        full_dump_file,
        static_cast<MINIDUMP_TYPE>((dump_type & (~MiniDumpNormal))
                                    | MiniDumpWithHandleData),
        exception_pointers ? &dump_exception_info : NULL,
        &user_streams,
        NULL) != FALSE;
  }

  // Add handle operations trace stream to the minidump if it was collected.
  if (handle_trace_data.GetUserStream(
          &user_stream_array[user_streams.UserStreamCount])) {
    ++user_streams.UserStreamCount;
  }

  bool result_minidump = write_dump(
      process_handle,
      process_id,
      dump_file,
      static_cast<MINIDUMP_TYPE>((dump_type & (~MiniDumpWithFullMemory))
                                  | MiniDumpNormal),
      exception_pointers ? &dump_exception_info : NULL,
      &user_streams,
      NULL) != FALSE;

  bool result = result_minidump && result_full_memory;

  CloseHandle(dump_file);
  if (full_dump_file != INVALID_HANDLE_VALUE)
    CloseHandle(full_dump_file);

  // Store the path of the dump file in the out parameter if dump generation
  // succeeded.
  if (result && dump_path) {
    *dump_path = dump_file_path;
  }
  if (result && full_memory_dump && full_dump_path) {
    *full_dump_path = full_dump_file_path;
  }

  return result;
}

HMODULE MinidumpGenerator::GetDbghelpModule() {
  AutoCriticalSection lock(&module_load_sync_);
  if (!dbghelp_module_) {
    dbghelp_module_ = LoadLibrary(TEXT("dbghelp.dll"));
  }

  return dbghelp_module_;
}

MinidumpGenerator::MiniDumpWriteDumpType MinidumpGenerator::GetWriteDump() {
  AutoCriticalSection lock(&get_proc_address_sync_);
  if (!write_dump_) {
    HMODULE module = GetDbghelpModule();
    if (module) {
      FARPROC proc = GetProcAddress(module, "MiniDumpWriteDump");
      write_dump_ = reinterpret_cast<MiniDumpWriteDumpType>(proc);
    }
  }

  return write_dump_;
}

HMODULE MinidumpGenerator::GetRpcrt4Module() {
  AutoCriticalSection lock(&module_load_sync_);
  if (!rpcrt4_module_) {
    rpcrt4_module_ = LoadLibrary(TEXT("rpcrt4.dll"));
  }

  return rpcrt4_module_;
}

MinidumpGenerator::UuidCreateType MinidumpGenerator::GetCreateUuid() {
  AutoCriticalSection lock(&module_load_sync_);
  if (!create_uuid_) {
    HMODULE module = GetRpcrt4Module();
    if (module) {
      FARPROC proc = GetProcAddress(module, "UuidCreate");
      create_uuid_ = reinterpret_cast<UuidCreateType>(proc);
    }
  }

  return create_uuid_;
}

bool MinidumpGenerator::GenerateDumpFilePath(wstring* file_path) {
  UUID id = {0};

  UuidCreateType create_uuid = GetCreateUuid();
  if (!create_uuid) {
    return false;
  }

  create_uuid(&id);
  wstring id_str = GUIDString::GUIDToWString(&id);

  *file_path = dump_path_ + TEXT("\\") + id_str + TEXT(".dmp");
  return true;
}

}  // namespace google_breakpad
<|endoftext|>