<commit_before>/* <x0d/plugins/status.cpp>
 *
 * This file is part of the x0 web server project and is released under AGPL-3.
 * http://www.xzero.io/
 *
 * (c) 2009-2014 Christian Parpart <trapni@gmail.com>
 */

#include <x0d/XzeroPlugin.h>
#include <x0/http/HttpServer.h>
#include <x0/http/HttpRequest.h>
#include <x0/http/HttpConnection.h>
#include <x0/http/HttpHeader.h>
#include <x0/io/BufferSource.h>
#include <x0/TimeSpan.h>
#include <x0/strutils.h>
#include <x0/Types.h>

#include <cstring>
#include <cerrno>
#include <cstddef>

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <dirent.h>

#define TRACE(msg...) DEBUG("status: " msg)

struct Stats :
    public x0::CustomData
{
    unsigned long long connectionsAccepted;
    unsigned long long requestsAccepted;
    unsigned long long requestsProcessed;
    long long active;
    long long reading;
    long long writing;

    Stats() :
        connectionsAccepted(0),
        requestsAccepted(0),
        requestsProcessed(0),
        active(0),
        reading(0),
        writing(0)
    {
    }

    ~Stats()
    {
    }

    Stats& operator+=(const Stats& s)
    {
        connectionsAccepted += s.connectionsAccepted;
        requestsAccepted += s.requestsAccepted;
        requestsProcessed += s.requestsProcessed;
        active += s.active;
        reading += s.reading;
        writing += s.writing;

        return *this;
    }
};

/**
 * \ingroup plugins
 * \brief example content generator plugin
 */
class StatusPlugin :
    public x0d::XzeroPlugin
{
private:
    std::list<x0::HttpConnection*> connections_;

    Stats historical_;

    x0::HttpServer::WorkerHook::Connection onWorkerSpawn_;
    x0::HttpServer::WorkerHook::Connection onWorkerUnspawn_;
    x0::HttpServer::ConnectionHook::Connection onConnectionOpen_;
    x0::HttpServer::ConnectionStateHook::Connection onConnectionStateChanged_;
    x0::HttpServer::ConnectionHook::Connection onConnectionClose_;
    x0::HttpServer::RequestHook::Connection onPreProcess_;
    x0::HttpServer::RequestHook::Connection onPostProcess_;

public:
    StatusPlugin(x0d::XzeroDaemon* d, const std::string& name) :
        x0d::XzeroPlugin(d, name)
    {
        mainHandler("status", &StatusPlugin::handleRequest);
        mainHandler("status.nginx_compat", &StatusPlugin::nginx_compat);

        onWorkerSpawn_ = server().onWorkerSpawn.connect<StatusPlugin, &StatusPlugin::onWorkerSpawn>(this);
        onWorkerUnspawn_ = server().onWorkerUnspawn.connect<StatusPlugin, &StatusPlugin::onWorkerUnspawn>(this);

        onConnectionOpen_ = server().onConnectionOpen.connect<StatusPlugin, &StatusPlugin::onConnectionOpen>(this);
        onConnectionStateChanged_ = server().onConnectionStateChanged.connect<StatusPlugin, &StatusPlugin::onConnectionStateChanged>(this);
        onConnectionClose_ = server().onConnectionClose.connect<StatusPlugin, &StatusPlugin::onConnectionClose>(this);

        onPreProcess_ = server().onPreProcess.connect<StatusPlugin, &StatusPlugin::onPreProcess>(this);
        onPostProcess_ = server().onPostProcess.connect<StatusPlugin, &StatusPlugin::onPostProcess>(this);
    }

    ~StatusPlugin()
    {
        server().onWorkerSpawn.disconnect(onWorkerSpawn_);
        server().onWorkerUnspawn.disconnect(onWorkerUnspawn_);

        server().onConnectionOpen.disconnect(onConnectionOpen_);
        server().onConnectionStateChanged.disconnect(onConnectionStateChanged_);
        server().onConnectionClose.disconnect(onConnectionClose_);

        server().onPreProcess.disconnect(onPreProcess_);
        server().onPostProcess.disconnect(onPostProcess_);
    }

private:
    void onWorkerSpawn(x0::HttpWorker* worker) {
        worker->setCustomData<Stats>(this);
    }

    void onWorkerUnspawn(x0::HttpWorker* worker) {
        // XXX stats' active/reading/writing should be all zero at this point already.
        auto stats = worker->customData<Stats>(this);
        historical_ += *stats;
        worker->clearCustomData(this);
    }

    void onConnectionOpen(x0::HttpConnection* connection) {
        auto stats = connection->worker().customData<Stats>(this);
        ++stats->connectionsAccepted;
        ++stats->active;
    }

    void onConnectionStateChanged(x0::HttpConnection* connection, x0::HttpConnection::State lastState) {
        auto stats = connection->worker().customData<Stats>(this);

        switch (lastState) {
            case x0::HttpConnection::ReadingRequest:
                --stats->reading;
                break;
            case x0::HttpConnection::SendingReply:
                --stats->writing;
                break;
            default:
                break;
        }

        switch (connection->state()) {
            case x0::HttpConnection::ReadingRequest:
                ++stats->reading;
                break;
            case x0::HttpConnection::SendingReply:
                ++stats->writing;
                break;
            default:
                break;
        }
    }

    void onConnectionClose(x0::HttpConnection* connection) {
        auto stats = connection->worker().customData<Stats>(this);

        switch (connection->state()) {
            case x0::HttpConnection::ReadingRequest:
                --stats->reading;
                break;
            case x0::HttpConnection::SendingReply:
                --stats->writing;
                break;
            default:
                break;
        }

        --stats->active;
    }

    void onPreProcess(x0::HttpRequest* r) {
        auto stats = r->connection.worker().customData<Stats>(this);
        ++stats->requestsAccepted;
    }

    void onPostProcess(x0::HttpRequest* r) {
        auto stats = r->connection.worker().customData<Stats>(this);
        ++stats->requestsProcessed;
    }

    bool nginx_compat(x0::HttpRequest* r, x0::FlowVM::Params& args)
    {
        x0::Buffer nginxCompatStatus(1024);
        Stats sum;

        for (std::size_t i = 0, e = server().workers().size(); i != e; ++i) {
            const x0::HttpWorker *w = server().workers()[i];
            const auto stats = w->customData<Stats>(this);
            sum += *stats;
        }

        nginxCompatStatus
            << "Active connections: " << sum.active << "\n"
            << "server accepts handled requests\n"
            << sum.connectionsAccepted << ' ' << sum.requestsAccepted << ' ' << sum.requestsProcessed << "\n"
            << "Reading: " << sum.reading << " Writing: " << sum.writing << " Waiting: " << (sum.active - (sum.reading + sum.writing)) << "\n"
            ;

        char buf[80];
        snprintf(buf, sizeof(buf), "%zu", nginxCompatStatus.size());
        r->responseHeaders.push_back("Content-Length", (char *)buf);

        r->responseHeaders.push_back("Content-Type", "text/plain");

        r->write<x0::BufferSource>(nginxCompatStatus);

        r->finish();

        return true;
    }

    bool handleRequest(x0::HttpRequest *r, x0::FlowVM::Params& args)
    {
        // set response status code
        r->status = x0::HttpStatus::Ok;
        r->responseHeaders.push_back("Content-Type", "text/html; charset=utf-8");

        bool debug = true;

        r->write<x0::BufferSource>(createResponseBody(debug));

        r->finish();

        return true;
    }

    // TODO let this method return a movable instead of a full copy
    x0::Buffer createResponseBody(bool debug)
    {
        /*
         * process uptime
         * number of threads (total / active)
         * cpu load (process / per worker)
         * # requests
         * # connections
         * # requests since start
         */

        x0::TimeSpan uptime(server().uptime());
        std::size_t nconns = 0;
        unsigned long long numTotalRequests = 0;
        unsigned long long numTotalConns = 0;
        double p1 = 0, p5 = 0, p15 = 0;

        for (std::size_t i = 0, e = server().workers().size(); i != e; ++i) {
            const x0::HttpWorker *w = server().workers()[i];
            nconns += w->connectionLoad();
            numTotalRequests += w->requestCount();
            numTotalConns += w->connectionCount();
            w->fetchPerformanceCounts(&p1, &p5, &p15);
        }

        x0::Buffer buf;
        buf << "";
        buf << "<title>x0 status page</title>\n";
        buf <<
            "";
        buf << "";
        buf << "";
        buf << "
x0 status page
\n";
        buf << "
" << server().tag() << "
\n";
        buf << "
\n";
        buf << "process uptime: " << uptime << "\n";
        buf << "process generation: " << server().generation() << "\n";

        buf << "average requests per second: ";
        char tmp[80];
        snprintf(tmp, sizeof(tmp), "%.2f, %.2f, %.2f", p1, p5, p15);
        buf << ((char*)tmp) << "\n";

        buf << "# workers: " << server().workers().size() << "\n";
        buf << "# connections: " << nconns << "\n";
        buf << "# total requests: " << numTotalRequests << "\n";
        buf << "# total connections: " << numTotalConns << "\n";
        buf << "
\n";

        buf << "
\n";

        buf << "" << "cid" << "	";
        buf << "" << "wid" << "	";
        buf << "" << "r/n" << "	";
        buf << "" << "IP" << "	";
        buf << "" << "state" << "	";
        buf << "" << "age" << "	";
        buf << "" << "idle" << "	";
        buf << "" << "read" << "	";
        buf << "" << "written" << "	";
        buf << "" << "host" << "	";
        buf << "" << "method" << "	";
        buf << "" << "uri" << "	";
        buf << "" << "status" << "	";

        if (debug)
            buf << "" << "debug" << "	";

        // {{{ FIXME make this one being serialized better
        for (auto w: server().workers()) {
            w->eachConnection([&](x0::HttpConnection* c) -> bool {
                dump(buf, c, debug);
                return true;
            });
        }
        // FIXME end. }}}

        buf << "
\n";

        buf << "\n";

        return buf;
    }

    void dump(x0::Buffer& out, x0::HttpConnection* c, bool debug)
    {
        out << "
";

        out << "" << c->id() << "	";
        out << "" << c->worker().id() << "	";
        out << "" << c->requestCount() << "	";
        out << "" << c->remoteIP().str() << "	";

        out << "" << c->state_str();
        if (c->state() == x0::HttpConnection::ReadingRequest)
            out << " (" << c->state_str() << ")";
        out << "	";

        out << "" << (c->worker().now() - c->socket()->startedAt()).str() << "	";
        out << "" << (c->worker().now() - c->socket()->lastActivityAt()).str() << "	";
        out << "" << c->requestParserOffset() << "/" << c->requestBufferSize() << "	";

        const x0::HttpRequest* r = c->request();
        if (r && c->state() != x0::HttpConnection::KeepAliveRead) {
            out << "" << r->bytesTransmitted() << "	";
            out << "" << sanitize(r->hostname) << "	";
            out << "" << sanitize(r->method) << "	";
            out << "" << sanitize(r->unparsedUri) << "	";
            out << "" << r->statusStr(r->status) << "	";
        } else {
            out << "" << "	";
        }

        if (debug) {
            out << "";
            out << "refcount:" << c->refCount() << ", ";
            c->socket()->inspect(out);
            if (c->request()) {
                c->request()->inspect(out);
            }
            out << "	";
        }

        out << "
\n";
    }

    template<typename T>
    std::string sanitize(const T& value)
    {
        std::string out;
        char buf[32];
        for (auto i: value) {
            switch (i) {
                case '<':
                case '>':
                case '&':
                    snprintf(buf, sizeof(buf), "&#%d;", static_cast<unsigned>(i) & 0xFF);
                    out += buf;
                    break;
                default:
                    out += i;
                    break;
            }
        }
        return !out.empty() ? out : "(null)";
    }
};

X0_EXPORT_PLUGIN_CLASS(StatusPlugin)
<commit_msg>[plugins] status: minor improvements<commit_after>/* <x0d/plugins/status.cpp>
 *
 * This file is part of the x0 web server project and is released under AGPL-3.
 * http://www.xzero.io/
 *
 * (c) 2009-2014 Christian Parpart <trapni@gmail.com>
 */

#include <x0d/XzeroPlugin.h>
#include <x0/http/HttpServer.h>
#include <x0/http/HttpRequest.h>
#include <x0/http/HttpConnection.h>
#include <x0/http/HttpHeader.h>
#include <x0/io/BufferSource.h>
#include <x0/TimeSpan.h>
#include <x0/strutils.h>
#include <x0/Types.h>

#include <cstring>
#include <cerrno>
#include <cstddef>

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <dirent.h>

#define TRACE(msg...) DEBUG("status: " msg)

struct Stats :
    public x0::CustomData
{
    unsigned long long connectionsAccepted;
    unsigned long long requestsAccepted;
    unsigned long long requestsProcessed;
    long long active;
    long long reading;
    long long writing;

    Stats() :
        connectionsAccepted(0),
        requestsAccepted(0),
        requestsProcessed(0),
        active(0),
        reading(0),
        writing(0)
    {
    }

    ~Stats()
    {
    }

    Stats& operator+=(const Stats& s)
    {
        connectionsAccepted += s.connectionsAccepted;
        requestsAccepted += s.requestsAccepted;
        requestsProcessed += s.requestsProcessed;
        active += s.active;
        reading += s.reading;
        writing += s.writing;

        return *this;
    }
};

/**
 * \ingroup plugins
 * \brief example content generator plugin
 */
class StatusPlugin :
    public x0d::XzeroPlugin
{
private:
    std::list<x0::HttpConnection*> connections_;

    Stats historical_;

    x0::HttpServer::WorkerHook::Connection onWorkerSpawn_;
    x0::HttpServer::WorkerHook::Connection onWorkerUnspawn_;
    x0::HttpServer::ConnectionHook::Connection onConnectionOpen_;
    x0::HttpServer::ConnectionStateHook::Connection onConnectionStateChanged_;
    x0::HttpServer::ConnectionHook::Connection onConnectionClose_;
    x0::HttpServer::RequestHook::Connection onPreProcess_;
    x0::HttpServer::RequestHook::Connection onPostProcess_;

public:
    StatusPlugin(x0d::XzeroDaemon* d, const std::string& name) :
        x0d::XzeroPlugin(d, name)
    {
        mainHandler("status", &StatusPlugin::handleRequest);
        mainHandler("status.nginx_compat", &StatusPlugin::nginx_compat);

        onWorkerSpawn_ = server().onWorkerSpawn.connect<StatusPlugin, &StatusPlugin::onWorkerSpawn>(this);
        onWorkerUnspawn_ = server().onWorkerUnspawn.connect<StatusPlugin, &StatusPlugin::onWorkerUnspawn>(this);

        onConnectionOpen_ = server().onConnectionOpen.connect<StatusPlugin, &StatusPlugin::onConnectionOpen>(this);
        onConnectionStateChanged_ = server().onConnectionStateChanged.connect<StatusPlugin, &StatusPlugin::onConnectionStateChanged>(this);
        onConnectionClose_ = server().onConnectionClose.connect<StatusPlugin, &StatusPlugin::onConnectionClose>(this);

        onPreProcess_ = server().onPreProcess.connect<StatusPlugin, &StatusPlugin::onPreProcess>(this);
        onPostProcess_ = server().onPostProcess.connect<StatusPlugin, &StatusPlugin::onPostProcess>(this);
    }

    ~StatusPlugin()
    {
        server().onWorkerSpawn.disconnect(onWorkerSpawn_);
        server().onWorkerUnspawn.disconnect(onWorkerUnspawn_);

        server().onConnectionOpen.disconnect(onConnectionOpen_);
        server().onConnectionStateChanged.disconnect(onConnectionStateChanged_);
        server().onConnectionClose.disconnect(onConnectionClose_);

        server().onPreProcess.disconnect(onPreProcess_);
        server().onPostProcess.disconnect(onPostProcess_);
    }

private:
    void onWorkerSpawn(x0::HttpWorker* worker) {
        worker->setCustomData<Stats>(this);
    }

    void onWorkerUnspawn(x0::HttpWorker* worker) {
        // XXX stats' active/reading/writing should be all zero at this point already.
        auto stats = worker->customData<Stats>(this);
        historical_ += *stats;
        worker->clearCustomData(this);
    }

    void onConnectionOpen(x0::HttpConnection* connection) {
        auto stats = connection->worker().customData<Stats>(this);
        ++stats->connectionsAccepted;
        ++stats->active;
    }

    void onConnectionStateChanged(x0::HttpConnection* connection, x0::HttpConnection::State lastState) {
        auto stats = connection->worker().customData<Stats>(this);

        switch (lastState) {
            case x0::HttpConnection::ReadingRequest:
                --stats->reading;
                break;
            case x0::HttpConnection::SendingReply:
                --stats->writing;
                break;
            default:
                break;
        }

        switch (connection->state()) {
            case x0::HttpConnection::ReadingRequest:
                ++stats->reading;
                break;
            case x0::HttpConnection::SendingReply:
                ++stats->writing;
                break;
            default:
                break;
        }
    }

    void onConnectionClose(x0::HttpConnection* connection) {
        auto stats = connection->worker().customData<Stats>(this);

        switch (connection->state()) {
            case x0::HttpConnection::ReadingRequest:
                --stats->reading;
                break;
            case x0::HttpConnection::SendingReply:
                --stats->writing;
                break;
            default:
                break;
        }

        --stats->active;
    }

    void onPreProcess(x0::HttpRequest* r) {
        auto stats = r->connection.worker().customData<Stats>(this);
        ++stats->requestsAccepted;
    }

    void onPostProcess(x0::HttpRequest* r) {
        auto stats = r->connection.worker().customData<Stats>(this);
        ++stats->requestsProcessed;
    }

    bool nginx_compat(x0::HttpRequest* r, x0::FlowVM::Params& args)
    {
        x0::Buffer nginxCompatStatus(1024);
        Stats sum;

        for (std::size_t i = 0, e = server().workers().size(); i != e; ++i) {
            const x0::HttpWorker *w = server().workers()[i];
            const auto stats = w->customData<Stats>(this);
            sum += *stats;
        }

        nginxCompatStatus
            << "Active connections: " << sum.active << "\n"
            << "server accepts handled requests\n"
            << sum.connectionsAccepted << ' ' << sum.requestsAccepted << ' ' << sum.requestsProcessed << "\n"
            << "Reading: " << sum.reading << " Writing: " << sum.writing << " Waiting: " << (sum.active - (sum.reading + sum.writing)) << "\n"
            ;

        char buf[80];
        snprintf(buf, sizeof(buf), "%zu", nginxCompatStatus.size());
        r->responseHeaders.push_back("Content-Length", (char *)buf);

        r->responseHeaders.push_back("Content-Type", "text/plain");

        r->write<x0::BufferSource>(nginxCompatStatus);

        r->finish();

        return true;
    }

    bool handleRequest(x0::HttpRequest *r, x0::FlowVM::Params& args)
    {
        // set response status code
        r->status = x0::HttpStatus::Ok;
        r->responseHeaders.push_back("Content-Type", "text/html; charset=utf-8");

        bool debug = true;

        r->write<x0::BufferSource>(createResponseBody(debug));

        r->finish();

        return true;
    }

    // TODO let this method return a movable instead of a full copy
    x0::Buffer createResponseBody(bool debug)
    {
        /*
         * process uptime
         * number of threads (total / active)
         * cpu load (process / per worker)
         * # requests
         * # connections
         * # requests since start
         */

        x0::TimeSpan uptime(server().uptime());
        std::size_t nconns = 0;
        unsigned long long numTotalRequests = 0;
        unsigned long long numTotalConns = 0;
        double p1 = 0, p5 = 0, p15 = 0;

        for (std::size_t i = 0, e = server().workers().size(); i != e; ++i) {
            const x0::HttpWorker *w = server().workers()[i];
            nconns += w->connectionLoad();
            numTotalRequests += w->requestCount();
            numTotalConns += w->connectionCount();
            w->fetchPerformanceCounts(&p1, &p5, &p15);
        }

        x0::Buffer buf;
        buf << "";
        buf << "<title>x0 status page</title>\n";
        buf <<
            "";
        buf << "";
        buf << "";
        buf << "
x0 status page
\n";
        buf << "
" << server().tag() << "
\n";
        buf << "
\n";
        buf << "process uptime: " << uptime << "\n";
        buf << "process generation: " << server().generation() << "\n";

        buf << "average requests per second: ";
        char tmp[80];
        snprintf(tmp, sizeof(tmp), "%.2f, %.2f, %.2f", p1, p5, p15);
        buf << ((char*)tmp) << "\n";

        buf << "# workers: " << server().workers().size() << "\n";
        buf << "# connections: " << nconns << "\n";
        buf << "# total requests: " << numTotalRequests << "\n";
        buf << "# total connections: " << numTotalConns << "\n";
        buf << "
\n";

        buf << "
\n";

        buf << "" << "cid" << "	";
        buf << "" << "wid" << "	";
        buf << "" << "r/n" << "	";
        buf << "" << "IP" << "	";
        buf << "" << "state" << "	";
        buf << "" << "age" << "	";
        buf << "" << "idle" << "	";
        buf << "" << "read" << "	";
        buf << "" << "written" << "	";
        buf << "" << "host" << "	";
        buf << "" << "method" << "	";
        buf << "" << "uri" << "	";
        buf << "" << "status" << "	";

        if (debug)
            buf << "" << "debug" << "	";

        // {{{ FIXME make this one being serialized better
        for (auto w: server().workers()) {
            w->eachConnection([&](x0::HttpConnection* c) -> bool {
                dump(buf, c, debug);
                return true;
            });
        }
        // FIXME end. }}}

        buf << "
\n";

        buf << "\n";

        return buf;
    }

    void dump(x0::Buffer& out, x0::HttpConnection* c, bool debug)
    {
        out << "
";

        out << "" << c->id() << "	";
        out << "" << c->worker().id() << "	";
        out << "" << c->requestCount() << "	";
        out << "" << c->remoteIP().str() << "	";

        out << "" << c->state_str();
        if (c->state() == x0::HttpConnection::ReadingRequest)
            out << " (" << c->state_str() << ")";
        out << "	";

        out << "" << (c->worker().now() - c->socket()->startedAt()).str() << "	";
        out << "" << (c->worker().now() - c->socket()->lastActivityAt()).str() << "	";
        out << "" << c->requestParserOffset() << "/" << c->requestBufferSize() << "	";

        const x0::HttpRequest* r = c->request();
        if (r && c->state() != x0::HttpConnection::KeepAliveRead) {
            out << "" << r->bytesTransmitted() << "	";
            out << "" << sanitize(r->hostname) << "	";
            out << "" << sanitize(r->method) << "	";
            out << "" << sanitize(r->unparsedUri) << "	";
            out << "" << r->statusStr(r->status) << "	";
        } else {
            out << "" << "	";
        }

        if (debug) {
            out << "";
            out << "refcount:" << c->refCount() << ", ";

            x0::Buffer inspectBuffer;
            c->socket()->inspect(inspectBuffer);
            if (c->request()) {
                c->request()->inspect(inspectBuffer);
            }

            out << inspectBuffer.replaceAll("\n", "
");

            out << "	";
        }

        out << "
\n";
    }

    template<typename T>
    std::string sanitize(const T& value)
    {
        std::string out;
        char buf[32];
        for (auto i: value) {
            switch (i) {
                case '<':
                case '>':
                case '&':
                    snprintf(buf, sizeof(buf), "&#%d;", static_cast<unsigned>(i) & 0xFF);
                    out += buf;
                    break;
                default:
                    out += i;
                    break;
            }
        }
        return !out.empty() ? out : "(null)";
    }
};

X0_EXPORT_PLUGIN_CLASS(StatusPlugin)
<|endoftext|>


<commit_before>#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <nav_msgs/OccupancyGrid.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl_ros/point_cloud.h>
#include <tf/transform_listener.h>
#include <pcl_ros/transforms.h>
#include <cmath>

typedef pcl::PointCloud<pcl::PointXYZ> PointCloudXYZ;

nav_msgs::OccupancyGrid nav_map;
ros::Publisher map_pub;
tf::TransformListener* tf_listener;
bool map_changed;

int x_origin, y_origin;

inline int round (const float a) { return int (a + 0.5); }

void setCellOccupancy(const int x, const int y, const int value)
{
	// only include this point if it falls within the map grid
	if(x < nav_map.info.width && x >= 0 &&
	   y < nav_map.info.height && y >= 0)
	{
		nav_map.data[y*nav_map.info.width + x] = value;
	}
}

void setCellOccupancy(const float x, const float y, const int value)
{
	int x_coord = (int)(x/nav_map.info.resolution) + x_origin;
	int y_coord = (int)(y/nav_map.info.resolution) + y_origin;

	setCellOccupancy(x_coord, y_coord, value);
}

int getCellOccupancy(const int x, const int y)
{
	if(x < nav_map.info.width && x >= 0 &&
	   y < nav_map.info.height && y >= 0)
	{
		return nav_map.data[y*nav_map.info.width + x];
	}

	return 50;
}

int getCellOccupancy(const float x, const float y)
{
	int x_coord = (int)(x/nav_map.info.resolution) + x_origin;
	int y_coord = (int)(y/nav_map.info.resolution) + y_origin;

	return getCellOccupancy(x_coord, y_coord);
}

void beamDDA(const tf::Vector3& beamStart, const tf::Vector3& beamEnd)
{
	int dx = (int)((beamEnd.getX() - beamStart.getX())/nav_map.info.resolution);
	int dy = (int)((beamEnd.getY() - beamStart.getY())/nav_map.info.resolution);
	int steps;

	float xIncrement, yIncrement;
	float x = beamStart.getX()/nav_map.info.resolution;
	float y = beamStart.getY()/nav_map.info.resolution;

	if(fabs(dx) > fabs(dy))
		steps = fabs(dx);
	else
		steps = fabs(dy);
	xIncrement = ((float)dx)/((float)steps);
	yIncrement = ((float)dy)/((float)steps);

/*	ROS_INFO("%3.3f, %3.3f, %3.3f, %3.3f, %d, %d, %d, %3.3f, %3.3f",
			 beamStart.getX(), beamStart.getY(),
			 beamEnd.getX(), beamEnd.getY(),
			 dx, dy, steps,
			 xIncrement, yIncrement);*/

	setCellOccupancy(x, y, 0);
	for(int k=0; k<steps-1; k++)
	{
		x += xIncrement;
		y += yIncrement;
		int new_value = getCellOccupancy(round(x)+x_origin, round(y)+y_origin) - 5;
		if(new_value < 0)
			new_value = 0;
		setCellOccupancy(round(x)+x_origin, round(y)+y_origin, new_value);
	}
}

void laserCallback(const sensor_msgs::PointCloud2::ConstPtr& ros_cloud)
{
	sensor_msgs::PointCloud2 map_ros_cloud;
	if(!tf_listener->waitForTransform(
		   "/map",
		   "/base_link",
		   ros_cloud->header.stamp,
		   ros::Duration(1.0)))
	{
		ROS_INFO("Transform from /map to /base_link failed");
		return; // if a transform isn't found within one second give up
	}
	pcl_ros::transformPointCloud("/map", *ros_cloud, map_ros_cloud, *tf_listener);
	// get the start point of the laser beam
	tf::StampedTransform laser_transform;
	if(!tf_listener->waitForTransform(
		   "/map",
		   "/head_hokuyo_frame",
		   ros_cloud->header.stamp,
		   ros::Duration(1.0)))
	{
		ROS_INFO("Transform from /map to /head_hokuyo_frame failed");
		return;
	}
	tf_listener->lookupTransform("/map", "/head_hokuyo_frame", ros_cloud->header.stamp, laser_transform);
	tf::Vector3 beam_start = laser_transform.getOrigin();
	beam_start.setZ(0);

	// convert to the PCL format
	PointCloudXYZ pcl_cloud;
	pcl::fromROSMsg(map_ros_cloud, pcl_cloud);

	for(pcl::PointCloud<pcl::PointXYZ>::const_iterator point = pcl_cloud.points.begin();
		point != pcl_cloud.points.end();
		point++)
	{
		if(point->z > 1.2 || point->z < .8)
			continue; // ignore points outside of the scan plane

		// get the vector in the map plane
		tf::Vector3 beam_end(point->x, point->y, 0);

		beamDDA(beam_start, beam_end);

		double beam_length = beam_end.distance(beam_start);
		// ignore the max distance results and results that come from
		// within the bounding cylinder of the robot
		if(fabs(beam_length - 30.0) < .03 || fabs(beam_length) < .5)
			continue; // skip points that are near the max sensor range

		int new_value = getCellOccupancy(point->x, point->y) + 15;
		if(new_value > 100)
			new_value = 100;
		setCellOccupancy(point->x, point->y, new_value);
	}

	map_changed = true;
}

int main(int argc, char** argv)
{
	ros::init(argc, argv, "map_creator_2d");
	ros::NodeHandle nh;

	ros::NodeHandle priv_nh("~");
	double map_resolution, x_offset, y_offset;
	int map_width, map_height;
	priv_nh.param<double>("resolution", map_resolution, 0.05);
	priv_nh.param<int>("width", map_width, 1000);
	priv_nh.param<int>("height", map_height, 1000);
	priv_nh.param<double>("x_offset", x_offset, .5);
	priv_nh.param<double>("y_offset", y_offset, .5);

	if(x_offset < 0 || x_offset > 1)
	{
		ROS_ERROR("Invalid x_offset. Must be in range [0, 1.0]");
		return 1;
	}

	if(y_offset < 0 || y_offset > 1)
	{
		ROS_ERROR("Invalid y_offset. Must be in range [0, 1.0]");
		return 1;
	}

	tf_listener = new tf::TransformListener();

	// set the nav_map frame
	nav_map.header.frame_id = "/map";

	// fill out the meta-data
	nav_map.info.map_load_time = ros::Time::now();
	nav_map.info.resolution = map_resolution;
	nav_map.info.width = map_width;
	nav_map.info.height = map_height;
	nav_map.info.origin.position.x = -nav_map.info.resolution*nav_map.info.width*x_offset;
	nav_map.info.origin.position.y = -nav_map.info.resolution*nav_map.info.height*y_offset;
	nav_map.info.origin.orientation.w = 1.0;

	// initialize to all unknown
	for(int i=0; i<nav_map.info.height; i++)
	{
		for(int j=0; j<nav_map.info.width; j++)
		{
			nav_map.data.push_back(30);
		}
	}

	// get the origin coordinates in the map grid
	x_origin = (int)(nav_map.info.width*x_offset);
	y_origin = (int)(nav_map.info.height*y_offset);

	// now that the datastructure is initialized set up the ROS stuff
	map_pub = nh.advertise<nav_msgs::OccupancyGrid>("/map", 1, true);
	ros::Subscriber laser_sub = nh.subscribe<sensor_msgs::PointCloud2>("laser/points", 10, &laserCallback);

	ros::Rate loop_rate(10);
	while(ros::ok())
	{
		if(map_changed)
		{
			nav_map.header.stamp = ros::Time::now();
			map_pub.publish(nav_map);
			map_changed = false;
		}

		ros::spinOnce();
		loop_rate.sleep();
	}

	delete tf_listener;
}
<commit_msg>Changed the mapping prior and the probability of being occupied<commit_after>#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <nav_msgs/OccupancyGrid.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl_ros/point_cloud.h>
#include <tf/transform_listener.h>
#include <pcl_ros/transforms.h>
#include <cmath>

typedef pcl::PointCloud<pcl::PointXYZ> PointCloudXYZ;

nav_msgs::OccupancyGrid nav_map;
ros::Publisher map_pub;
tf::TransformListener* tf_listener;
bool map_changed;

int x_origin, y_origin;

inline int round (const float a) { return int (a + 0.5); }

void setCellOccupancy(const int x, const int y, const int value)
{
	// only include this point if it falls within the map grid
	if(x < nav_map.info.width && x >= 0 &&
	   y < nav_map.info.height && y >= 0)
	{
		nav_map.data[y*nav_map.info.width + x] = value;
	}
}

void setCellOccupancy(const float x, const float y, const int value)
{
	int x_coord = (int)(x/nav_map.info.resolution) + x_origin;
	int y_coord = (int)(y/nav_map.info.resolution) + y_origin;

	setCellOccupancy(x_coord, y_coord, value);
}

int getCellOccupancy(const int x, const int y)
{
	if(x < nav_map.info.width && x >= 0 &&
	   y < nav_map.info.height && y >= 0)
	{
		return nav_map.data[y*nav_map.info.width + x];
	}

	return 50;
}

int getCellOccupancy(const float x, const float y)
{
	int x_coord = (int)(x/nav_map.info.resolution) + x_origin;
	int y_coord = (int)(y/nav_map.info.resolution) + y_origin;

	return getCellOccupancy(x_coord, y_coord);
}

void beamDDA(const tf::Vector3& beamStart, const tf::Vector3& beamEnd)
{
	int dx = (int)((beamEnd.getX() - beamStart.getX())/nav_map.info.resolution);
	int dy = (int)((beamEnd.getY() - beamStart.getY())/nav_map.info.resolution);
	int steps;

	float xIncrement, yIncrement;
	float x = beamStart.getX()/nav_map.info.resolution;
	float y = beamStart.getY()/nav_map.info.resolution;

	if(fabs(dx) > fabs(dy))
		steps = fabs(dx);
	else
		steps = fabs(dy);
	xIncrement = ((float)dx)/((float)steps);
	yIncrement = ((float)dy)/((float)steps);

/*	ROS_INFO("%3.3f, %3.3f, %3.3f, %3.3f, %d, %d, %d, %3.3f, %3.3f",
			 beamStart.getX(), beamStart.getY(),
			 beamEnd.getX(), beamEnd.getY(),
			 dx, dy, steps,
			 xIncrement, yIncrement);*/

	setCellOccupancy(x, y, 0);
	for(int k=0; k<steps-1; k++)
	{
		x += xIncrement;
		y += yIncrement;
		int new_value = getCellOccupancy(round(x)+x_origin, round(y)+y_origin) - 5;
		if(new_value < 0)
			new_value = 0;
		setCellOccupancy(round(x)+x_origin, round(y)+y_origin, new_value);
	}
}

void laserCallback(const sensor_msgs::PointCloud2::ConstPtr& ros_cloud)
{
	sensor_msgs::PointCloud2 map_ros_cloud;
	if(!tf_listener->waitForTransform(
		   "/map",
		   "/base_link",
		   ros_cloud->header.stamp,
		   ros::Duration(1.0)))
	{
		ROS_INFO("Transform from /map to /base_link failed");
		return; // if a transform isn't found within one second give up
	}
	pcl_ros::transformPointCloud("/map", *ros_cloud, map_ros_cloud, *tf_listener);
	// get the start point of the laser beam
	tf::StampedTransform laser_transform;
	if(!tf_listener->waitForTransform(
		   "/map",
		   "/head_hokuyo_frame",
		   ros_cloud->header.stamp,
		   ros::Duration(1.0)))
	{
		ROS_INFO("Transform from /map to /head_hokuyo_frame failed");
		return;
	}
	tf_listener->lookupTransform("/map", "/head_hokuyo_frame", ros_cloud->header.stamp, laser_transform);
	tf::Vector3 beam_start = laser_transform.getOrigin();
	beam_start.setZ(0);

	// convert to the PCL format
	PointCloudXYZ pcl_cloud;
	pcl::fromROSMsg(map_ros_cloud, pcl_cloud);

	for(pcl::PointCloud<pcl::PointXYZ>::const_iterator point = pcl_cloud.points.begin();
		point != pcl_cloud.points.end();
		point++)
	{
		if(point->z > 1.2 || point->z < .8)
			continue; // ignore points outside of the scan plane

		// get the vector in the map plane
		tf::Vector3 beam_end(point->x, point->y, 0);

		beamDDA(beam_start, beam_end);

		double beam_length = beam_end.distance(beam_start);
		// ignore the max distance results and results that come from
		// within the bounding cylinder of the robot
		if(fabs(beam_length - 30.0) < .03 || fabs(beam_length) < .5)
			continue; // skip points that are near the max sensor range

		int new_value = getCellOccupancy(point->x, point->y) + 25;
		if(new_value > 100)
			new_value = 100;
		setCellOccupancy(point->x, point->y, new_value);
	}

	map_changed = true;
}

int main(int argc, char** argv)
{
	ros::init(argc, argv, "map_creator_2d");
	ros::NodeHandle nh;

	ros::NodeHandle priv_nh("~");
	double map_resolution, x_offset, y_offset;
	int map_width, map_height;
	priv_nh.param<double>("resolution", map_resolution, 0.05);
	priv_nh.param<int>("width", map_width, 1000);
	priv_nh.param<int>("height", map_height, 1000);
	priv_nh.param<double>("x_offset", x_offset, .5);
	priv_nh.param<double>("y_offset", y_offset, .5);

	if(x_offset < 0 || x_offset > 1)
	{
		ROS_ERROR("Invalid x_offset. Must be in range [0, 1.0]");
		return 1;
	}

	if(y_offset < 0 || y_offset > 1)
	{
		ROS_ERROR("Invalid y_offset. Must be in range [0, 1.0]");
		return 1;
	}

	tf_listener = new tf::TransformListener();

	// set the nav_map frame
	nav_map.header.frame_id = "/map";

	// fill out the meta-data
	nav_map.info.map_load_time = ros::Time::now();
	nav_map.info.resolution = map_resolution;
	nav_map.info.width = map_width;
	nav_map.info.height = map_height;
	nav_map.info.origin.position.x = -nav_map.info.resolution*nav_map.info.width*x_offset;
	nav_map.info.origin.position.y = -nav_map.info.resolution*nav_map.info.height*y_offset;
	nav_map.info.origin.orientation.w = 1.0;

	// initialize to all unknown
	for(int i=0; i<nav_map.info.height; i++)
	{
		for(int j=0; j<nav_map.info.width; j++)
		{
			nav_map.data.push_back(50);
		}
	}

	// get the origin coordinates in the map grid
	x_origin = (int)(nav_map.info.width*x_offset);
	y_origin = (int)(nav_map.info.height*y_offset);

	// now that the datastructure is initialized set up the ROS stuff
	map_pub = nh.advertise<nav_msgs::OccupancyGrid>("/map", 1, true);
	ros::Subscriber laser_sub = nh.subscribe<sensor_msgs::PointCloud2>("laser/points", 10, &laserCallback);

	ros::Rate loop_rate(10);
	while(ros::ok())
	{
		if(map_changed)
		{
			nav_map.header.stamp = ros::Time::now();
			map_pub.publish(nav_map);
			map_changed = false;
		}

		ros::spinOnce();
		loop_rate.sleep();
	}

	delete tf_listener;
}
<|endoftext|>


<commit_before>/*
 * Copyright 2009-2015 The VOTCA Development Team (http://www.votca.org)
 *
 * 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.
 *
 */

#include <algorithm>
#include <votca/csg/topology.h>
#include <votca/csg/exclusionlist.h>

namespace votca { namespace csg {

void ExclusionList::Clear(void)
{
    list< exclusion_t *>::iterator iter;

    for(iter=_exclusions.begin();iter!=_exclusions.end();++iter)
        delete *iter;
    _exclusions.clear();
}

void ExclusionList::CreateExclusions(Topology *top) {
	InteractionContainer &ic = top->BondedInteractions();
	InteractionContainer::iterator ia;

	for (ia = ic.begin(); ia != ic.end(); ++ia) {
		int beads_in_int = (*ia)->BeadCount();
		list<Bead *> l;

		for (int ibead = 0; ibead < beads_in_int; ibead ++) {
			int ii = (*ia)->getBeadId(ibead);
			l.push_back(top->getBead(ii));
		}
		ExcludeList(l);
	}
}

bool ExclusionList::IsExcluded(Bead *bead1, Bead *bead2) {
    exclusion_t *excl;
    if(bead1->getMolecule() != bead2->getMolecule()) return false;
    if (bead2->getId() < bead1->getId()) swap(bead1, bead2);
    if ((excl = GetExclusions(bead1))) {
        if(find(excl->_exclude.begin(), excl->_exclude.end(), bead2)
                != excl->_exclude.end()) return true;
    }
    return false;
}

bool compareAtomIdiExclusionList(const ExclusionList::exclusion_t *a, const ExclusionList::exclusion_t *b){
    return a->_atom->getId() < b->_atom->getId();
}

bool compareAtomIdBeadList(const Bead *a, const Bead *b){
    return a->getId() < b->getId();
}

std::ostream &operator<<(std::ostream &out, ExclusionList& exl)
{
    exl._exclusions.sort(compareAtomIdiExclusionList);

    list<ExclusionList::exclusion_t*>::iterator ex;
    for(ex=exl._exclusions.begin();ex!=exl._exclusions.end();++ex) {
        (*ex)->_exclude.sort(compareAtomIdBeadList);
        list<Bead *>::iterator i;
        out << (int)((*ex)->_atom->getId()) + 1;
        for(i=(*ex)->_exclude.begin(); i!=(*ex)->_exclude.end(); ++i) {
            out << " " << ((*i)->getId()+1);
        }
        out << endl;
    }
    return out;
}

}}
<commit_msg>Added using namespace std to source file exclusionlist.cc<commit_after>/*
 * Copyright 2009-2015 The VOTCA Development Team (http://www.votca.org)
 *
 * 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.
 *
 */

#include <algorithm>
#include <votca/csg/topology.h>
#include <votca/csg/exclusionlist.h>

namespace votca { namespace csg {

using namespace std;

void ExclusionList::Clear(void)
{
    list< exclusion_t *>::iterator iter;

    for(iter=_exclusions.begin();iter!=_exclusions.end();++iter)
        delete *iter;
    _exclusions.clear();
}

void ExclusionList::CreateExclusions(Topology *top) {
	InteractionContainer &ic = top->BondedInteractions();
	InteractionContainer::iterator ia;

	for (ia = ic.begin(); ia != ic.end(); ++ia) {
		int beads_in_int = (*ia)->BeadCount();
		list<Bead *> l;

		for (int ibead = 0; ibead < beads_in_int; ibead ++) {
			int ii = (*ia)->getBeadId(ibead);
			l.push_back(top->getBead(ii));
		}
		ExcludeList(l);
	}
}

bool ExclusionList::IsExcluded(Bead *bead1, Bead *bead2) {
    exclusion_t *excl;
    if(bead1->getMolecule() != bead2->getMolecule()) return false;
    if (bead2->getId() < bead1->getId()) swap(bead1, bead2);
    if ((excl = GetExclusions(bead1))) {
        if(find(excl->_exclude.begin(), excl->_exclude.end(), bead2)
                != excl->_exclude.end()) return true;
    }
    return false;
}

bool compareAtomIdiExclusionList(const ExclusionList::exclusion_t *a, const ExclusionList::exclusion_t *b){
    return a->_atom->getId() < b->_atom->getId();
}

bool compareAtomIdBeadList(const Bead *a, const Bead *b){
    return a->getId() < b->getId();
}

std::ostream &operator<<(std::ostream &out, ExclusionList& exl)
{
    exl._exclusions.sort(compareAtomIdiExclusionList);

    list<ExclusionList::exclusion_t*>::iterator ex;
    for(ex=exl._exclusions.begin();ex!=exl._exclusions.end();++ex) {
        (*ex)->_exclude.sort(compareAtomIdBeadList);
        list<Bead *>::iterator i;
        out << (int)((*ex)->_atom->getId()) + 1;
        for(i=(*ex)->_exclude.begin(); i!=(*ex)->_exclude.end(); ++i) {
            out << " " << ((*i)->getId()+1);
        }
        out << endl;
    }
    return out;
}

}}
<|endoftext|>


<commit_before>/*
 * SSL and TLS filter.
 *
 * author: Max Kellermann <mk@cm4all.com>
 */

#include "ssl_filter.hxx"
#include "ssl_factory.hxx"
#include "ssl_config.hxx"
#include "ssl_quark.hxx"
#include "Unique.hxx"
#include "Name.hxx"
#include "Error.hxx"
#include "thread_socket_filter.hxx"
#include "pool.hxx"
#include "gerrno.h"
#include "fb_pool.hxx"
#include "SliceFifoBuffer.hxx"
#include "util/AllocatedString.hxx"

#include <openssl/ssl.h>
#include <openssl/err.h>

#include <assert.h>
#include <string.h>

/**
 * Throttle if a #BIO grows larger than this number of bytes.
 */
static constexpr int SSL_THROTTLE_THRESHOLD = 16384;

struct SslFilter final : ThreadSocketFilterHandler {
    /**
     * Buffers which can be accessed from within the thread without
     * holding locks.  These will be copied to/from the according
     * #thread_socket_filter buffers.
     */
    SliceFifoBuffer decrypted_input, plain_output;

    /**
     * Memory BIOs for passing data to/from OpenSSL.
     */
    BIO *const encrypted_input, *const encrypted_output;

    const UniqueSSL ssl;

    bool handshaking = true;

    bool eof;

    AllocatedString<> peer_subject = nullptr, peer_issuer_subject = nullptr;

    SslFilter(UniqueSSL &&_ssl)
        :encrypted_input(BIO_new(BIO_s_mem())),
         encrypted_output(BIO_new(BIO_s_mem())),
         ssl(std::move(_ssl)) {
        decrypted_input.Allocate(fb_pool_get());
        SSL_set_bio(ssl.get(), encrypted_input, encrypted_output);
    }

    ~SslFilter() {
        decrypted_input.Free(fb_pool_get());
        plain_output.FreeIfDefined(fb_pool_get());
    }

    bool Encrypt(GError **error_r);

    /* virtual methods from class ThreadSocketFilterHandler */
    bool Run(ThreadSocketFilter &f, GError **error_r) override;

    void CycleBuffers(ThreadSocketFilter &f) override {
        if (f.IsIdle())
            decrypted_input.CycleIfEmpty(fb_pool_get());
    }

    void Destroy(ThreadSocketFilter &) override {
        this->~SslFilter();
    }
};

static void
ssl_set_error(GError **error_r)
{
    if (error_r == nullptr)
        return;

    unsigned long error = ERR_get_error();
    char buffer[120];
    g_set_error(error_r, ssl_quark(), 0, "%s",
                ERR_error_string(error, buffer));
}

/**
 * Is the #BIO full, i.e. above the #SSL_THROTTLE_THRESHOLD?
 */
gcc_pure
static bool
IsFull(BIO *bio)
{
    return BIO_pending(bio) >= SSL_THROTTLE_THRESHOLD;
}

/**
 * Move data from #src to #dest.
 */
static void
Move(BIO *dest, ForeignFifoBuffer<uint8_t> &src)
{
    auto r = src.Read();
    if (r.IsEmpty())
        return;

    if (IsFull(dest))
        /* throttle */
        return;

    int nbytes = BIO_write(dest, r.data, r.size);
    if (nbytes > 0)
        src.Consume(nbytes);
}

static void
Move(ForeignFifoBuffer<uint8_t> &dest, BIO *src)
{
    while (true) {
        auto w = dest.Write();
        if (w.IsEmpty())
            return;

        int nbytes = BIO_read(src, w.data, w.size);
        if (nbytes <= 0)
            return;

        dest.Append(nbytes);
    }
}

static AllocatedString<>
format_subject_name(X509 *cert)
{
    return ToString(X509_get_subject_name(cert));
}

static AllocatedString<>
format_issuer_subject_name(X509 *cert)
{
    return ToString(X509_get_issuer_name(cert));
}

gcc_pure
static bool
is_ssl_error(SSL *ssl, int ret)
{
    if (ret == 0)
        /* this is always an error according to the documentation of
           SSL_read(), SSL_write() and SSL_do_handshake() */
        return true;

    switch (SSL_get_error(ssl, ret)) {
    case SSL_ERROR_NONE:
    case SSL_ERROR_WANT_READ:
    case SSL_ERROR_WANT_WRITE:
    case SSL_ERROR_WANT_CONNECT:
    case SSL_ERROR_WANT_ACCEPT:
        return false;

    default:
        return true;
    }
}

static bool
check_ssl_error(SSL *ssl, int result, GError **error_r)
{
    if (is_ssl_error(ssl, result)) {
        ssl_set_error(error_r);
        return false;
    } else
        return true;
}

enum class SslDecryptResult {
    SUCCESS,
    ERROR,
    CLOSE_NOTIFY_ALERT,
};

static SslDecryptResult
ssl_decrypt(SSL *ssl, ForeignFifoBuffer<uint8_t> &buffer, GError **error_r)
{
    /* SSL_read() must be called repeatedly until there is no more
       data (or until the buffer is full) */

    while (true) {
        auto w = buffer.Write();
        if (w.IsEmpty())
            return SslDecryptResult::SUCCESS;

        int result = SSL_read(ssl, w.data, w.size);
        if (result <= 0) {
            if (SSL_get_error(ssl, result) == SSL_ERROR_ZERO_RETURN)
                /* got a "close notify" alert from the peer */
                return SslDecryptResult::CLOSE_NOTIFY_ALERT;

            return check_ssl_error(ssl, result, error_r)
                ? SslDecryptResult::SUCCESS
                : SslDecryptResult::ERROR;
        }

        buffer.Append(result);
    }
}

/**
 * @return false on error
 */
static bool
ssl_encrypt(SSL *ssl, ForeignFifoBuffer<uint8_t> &buffer, GError **error_r)
{
    auto r = buffer.Read();
    if (r.IsEmpty())
        return true;

    int result = SSL_write(ssl, r.data, r.size);
    if (result <= 0)
        return check_ssl_error(ssl, result, error_r);

    buffer.Consume(result);
    return true;
}

inline bool
SslFilter::Encrypt(GError **error_r)
{
    return IsFull(encrypted_output) || /* throttle? */
        ssl_encrypt(ssl.get(), plain_output, error_r);
}

/*
 * thread_socket_filter_handler
 *
 */

inline bool
SslFilter::Run(ThreadSocketFilter &f, GError **error_r)
{
    /* copy input (and output to make room for more output) */

    {
        std::unique_lock<std::mutex> lock(f.mutex);

        f.decrypted_input.MoveFrom(decrypted_input);
        plain_output.MoveFromAllowNull(f.plain_output);
        Move(encrypted_input, f.encrypted_input);
        Move(f.encrypted_output, encrypted_output);
    }

    /* let OpenSSL work */

    ERR_clear_error();

    if (gcc_unlikely(handshaking)) {
        int result = SSL_do_handshake(ssl.get());
        if (result == 1) {
            handshaking = false;

            UniqueX509 cert(SSL_get_peer_certificate(ssl.get()));
            if (cert != nullptr) {
                peer_subject = format_subject_name(cert.get());
                peer_issuer_subject = format_issuer_subject_name(cert.get());
            }
        } else if (!check_ssl_error(ssl.get(), result, error_r))
            return false;
    }

    if (gcc_likely(!handshaking)) {
        if (!Encrypt(error_r))
            return false;

        switch (ssl_decrypt(ssl.get(), decrypted_input, error_r)) {
        case SslDecryptResult::SUCCESS:
            break;

        case SslDecryptResult::ERROR:
            return false;

        case SslDecryptResult::CLOSE_NOTIFY_ALERT:
            {
                std::unique_lock<std::mutex> lock(f.mutex);
                f.input_eof = true;
            }
            break;
        }
    }

    /* copy output */

    {
        std::unique_lock<std::mutex> lock(f.mutex);

        f.decrypted_input.MoveFrom(decrypted_input);

        /* let the main thread free our plain_output buffer */
        plain_output.SwapIfNull(f.plain_output);

        Move(f.encrypted_output, encrypted_output);
        f.drained = plain_output.IsEmpty() &&
            BIO_eof(encrypted_output);
    }

    return true;
}

/*
 * constructor
 *
 */

SslFilter *
ssl_filter_new(struct pool &pool, UniqueSSL &&ssl)
{
    return NewFromPool<SslFilter>(pool, std::move(ssl));
}

SslFilter *
ssl_filter_new(struct pool *pool, SslFactory &factory,
               GError **error_r)
{
    assert(pool != nullptr);

    try {
        return NewFromPool<SslFilter>(*pool, ssl_factory_make(factory));
    } catch (const SslError &e) {
        g_set_error(error_r, ssl_quark(), 0, "SSL_new() failed: %s",
                    e.what());
        return nullptr;
    }
}

ThreadSocketFilterHandler &
ssl_filter_get_handler(SslFilter &ssl)
{
    return ssl;
}

const char *
ssl_filter_get_peer_subject(SslFilter *ssl)
{
    assert(ssl != nullptr);

    return ssl->peer_subject.c_str();
}

const char *
ssl_filter_get_peer_issuer_subject(SslFilter *ssl)
{
    assert(ssl != nullptr);

    return ssl->peer_issuer_subject.c_str();
}
<commit_msg>ssl/filter: move plain_output disposal to PostRun()<commit_after>/*
 * SSL and TLS filter.
 *
 * author: Max Kellermann <mk@cm4all.com>
 */

#include "ssl_filter.hxx"
#include "ssl_factory.hxx"
#include "ssl_config.hxx"
#include "ssl_quark.hxx"
#include "Unique.hxx"
#include "Name.hxx"
#include "Error.hxx"
#include "thread_socket_filter.hxx"
#include "pool.hxx"
#include "gerrno.h"
#include "fb_pool.hxx"
#include "SliceFifoBuffer.hxx"
#include "util/AllocatedString.hxx"

#include <openssl/ssl.h>
#include <openssl/err.h>

#include <assert.h>
#include <string.h>

/**
 * Throttle if a #BIO grows larger than this number of bytes.
 */
static constexpr int SSL_THROTTLE_THRESHOLD = 16384;

struct SslFilter final : ThreadSocketFilterHandler {
    /**
     * Buffers which can be accessed from within the thread without
     * holding locks.  These will be copied to/from the according
     * #thread_socket_filter buffers.
     */
    SliceFifoBuffer decrypted_input, plain_output;

    /**
     * Memory BIOs for passing data to/from OpenSSL.
     */
    BIO *const encrypted_input, *const encrypted_output;

    const UniqueSSL ssl;

    bool handshaking = true;

    bool eof;

    AllocatedString<> peer_subject = nullptr, peer_issuer_subject = nullptr;

    SslFilter(UniqueSSL &&_ssl)
        :encrypted_input(BIO_new(BIO_s_mem())),
         encrypted_output(BIO_new(BIO_s_mem())),
         ssl(std::move(_ssl)) {
        decrypted_input.Allocate(fb_pool_get());
        SSL_set_bio(ssl.get(), encrypted_input, encrypted_output);
    }

    ~SslFilter() {
        decrypted_input.Free(fb_pool_get());
        plain_output.FreeIfDefined(fb_pool_get());
    }

    bool Encrypt(GError **error_r);

    /* virtual methods from class ThreadSocketFilterHandler */
    bool Run(ThreadSocketFilter &f, GError **error_r) override;
    void PostRun(ThreadSocketFilter &f) override;

    void CycleBuffers(ThreadSocketFilter &f) override {
        if (f.IsIdle())
            decrypted_input.CycleIfEmpty(fb_pool_get());
    }

    void Destroy(ThreadSocketFilter &) override {
        this->~SslFilter();
    }
};

static void
ssl_set_error(GError **error_r)
{
    if (error_r == nullptr)
        return;

    unsigned long error = ERR_get_error();
    char buffer[120];
    g_set_error(error_r, ssl_quark(), 0, "%s",
                ERR_error_string(error, buffer));
}

/**
 * Is the #BIO full, i.e. above the #SSL_THROTTLE_THRESHOLD?
 */
gcc_pure
static bool
IsFull(BIO *bio)
{
    return BIO_pending(bio) >= SSL_THROTTLE_THRESHOLD;
}

/**
 * Move data from #src to #dest.
 */
static void
Move(BIO *dest, ForeignFifoBuffer<uint8_t> &src)
{
    auto r = src.Read();
    if (r.IsEmpty())
        return;

    if (IsFull(dest))
        /* throttle */
        return;

    int nbytes = BIO_write(dest, r.data, r.size);
    if (nbytes > 0)
        src.Consume(nbytes);
}

static void
Move(ForeignFifoBuffer<uint8_t> &dest, BIO *src)
{
    while (true) {
        auto w = dest.Write();
        if (w.IsEmpty())
            return;

        int nbytes = BIO_read(src, w.data, w.size);
        if (nbytes <= 0)
            return;

        dest.Append(nbytes);
    }
}

static AllocatedString<>
format_subject_name(X509 *cert)
{
    return ToString(X509_get_subject_name(cert));
}

static AllocatedString<>
format_issuer_subject_name(X509 *cert)
{
    return ToString(X509_get_issuer_name(cert));
}

gcc_pure
static bool
is_ssl_error(SSL *ssl, int ret)
{
    if (ret == 0)
        /* this is always an error according to the documentation of
           SSL_read(), SSL_write() and SSL_do_handshake() */
        return true;

    switch (SSL_get_error(ssl, ret)) {
    case SSL_ERROR_NONE:
    case SSL_ERROR_WANT_READ:
    case SSL_ERROR_WANT_WRITE:
    case SSL_ERROR_WANT_CONNECT:
    case SSL_ERROR_WANT_ACCEPT:
        return false;

    default:
        return true;
    }
}

static bool
check_ssl_error(SSL *ssl, int result, GError **error_r)
{
    if (is_ssl_error(ssl, result)) {
        ssl_set_error(error_r);
        return false;
    } else
        return true;
}

enum class SslDecryptResult {
    SUCCESS,
    ERROR,
    CLOSE_NOTIFY_ALERT,
};

static SslDecryptResult
ssl_decrypt(SSL *ssl, ForeignFifoBuffer<uint8_t> &buffer, GError **error_r)
{
    /* SSL_read() must be called repeatedly until there is no more
       data (or until the buffer is full) */

    while (true) {
        auto w = buffer.Write();
        if (w.IsEmpty())
            return SslDecryptResult::SUCCESS;

        int result = SSL_read(ssl, w.data, w.size);
        if (result <= 0) {
            if (SSL_get_error(ssl, result) == SSL_ERROR_ZERO_RETURN)
                /* got a "close notify" alert from the peer */
                return SslDecryptResult::CLOSE_NOTIFY_ALERT;

            return check_ssl_error(ssl, result, error_r)
                ? SslDecryptResult::SUCCESS
                : SslDecryptResult::ERROR;
        }

        buffer.Append(result);
    }
}

/**
 * @return false on error
 */
static bool
ssl_encrypt(SSL *ssl, ForeignFifoBuffer<uint8_t> &buffer, GError **error_r)
{
    auto r = buffer.Read();
    if (r.IsEmpty())
        return true;

    int result = SSL_write(ssl, r.data, r.size);
    if (result <= 0)
        return check_ssl_error(ssl, result, error_r);

    buffer.Consume(result);
    return true;
}

inline bool
SslFilter::Encrypt(GError **error_r)
{
    return IsFull(encrypted_output) || /* throttle? */
        ssl_encrypt(ssl.get(), plain_output, error_r);
}

/*
 * thread_socket_filter_handler
 *
 */

inline bool
SslFilter::Run(ThreadSocketFilter &f, GError **error_r)
{
    /* copy input (and output to make room for more output) */

    {
        std::unique_lock<std::mutex> lock(f.mutex);

        f.decrypted_input.MoveFrom(decrypted_input);
        plain_output.MoveFromAllowNull(f.plain_output);
        Move(encrypted_input, f.encrypted_input);
        Move(f.encrypted_output, encrypted_output);
    }

    /* let OpenSSL work */

    ERR_clear_error();

    if (gcc_unlikely(handshaking)) {
        int result = SSL_do_handshake(ssl.get());
        if (result == 1) {
            handshaking = false;

            UniqueX509 cert(SSL_get_peer_certificate(ssl.get()));
            if (cert != nullptr) {
                peer_subject = format_subject_name(cert.get());
                peer_issuer_subject = format_issuer_subject_name(cert.get());
            }
        } else if (!check_ssl_error(ssl.get(), result, error_r))
            return false;
    }

    if (gcc_likely(!handshaking)) {
        if (!Encrypt(error_r))
            return false;

        switch (ssl_decrypt(ssl.get(), decrypted_input, error_r)) {
        case SslDecryptResult::SUCCESS:
            break;

        case SslDecryptResult::ERROR:
            return false;

        case SslDecryptResult::CLOSE_NOTIFY_ALERT:
            {
                std::unique_lock<std::mutex> lock(f.mutex);
                f.input_eof = true;
            }
            break;
        }
    }

    /* copy output */

    {
        std::unique_lock<std::mutex> lock(f.mutex);

        f.decrypted_input.MoveFrom(decrypted_input);

        Move(f.encrypted_output, encrypted_output);
        f.drained = plain_output.IsEmpty() &&
            BIO_eof(encrypted_output);
    }

    return true;
}

void
SslFilter::PostRun(ThreadSocketFilter &f)
{
    if (f.IsIdle()) {
        plain_output.FreeIfEmpty(fb_pool_get());
    }
}

/*
 * constructor
 *
 */

SslFilter *
ssl_filter_new(struct pool &pool, UniqueSSL &&ssl)
{
    return NewFromPool<SslFilter>(pool, std::move(ssl));
}

SslFilter *
ssl_filter_new(struct pool *pool, SslFactory &factory,
               GError **error_r)
{
    assert(pool != nullptr);

    try {
        return NewFromPool<SslFilter>(*pool, ssl_factory_make(factory));
    } catch (const SslError &e) {
        g_set_error(error_r, ssl_quark(), 0, "SSL_new() failed: %s",
                    e.what());
        return nullptr;
    }
}

ThreadSocketFilterHandler &
ssl_filter_get_handler(SslFilter &ssl)
{
    return ssl;
}

const char *
ssl_filter_get_peer_subject(SslFilter *ssl)
{
    assert(ssl != nullptr);

    return ssl->peer_subject.c_str();
}

const char *
ssl_filter_get_peer_issuer_subject(SslFilter *ssl)
{
    assert(ssl != nullptr);

    return ssl->peer_issuer_subject.c_str();
}
<|endoftext|>


<commit_before>#include "xchainer/cuda/cudnn.h"

#include <algorithm>
#include <memory>
#include <utility>
#include <vector>

#include <cudnn.h>
#include <nonstd/optional.hpp>

#include "xchainer/array.h"
#include "xchainer/cuda/cuda_device.h"
#include "xchainer/device.h"
#include "xchainer/dtype.h"
#include "xchainer/error.h"
#include "xchainer/routines/creation.h"
#include "xchainer/stack_vector.h"

namespace xchainer {
namespace cuda {

CudnnError::CudnnError(cudnnStatus_t status) : XchainerError{cudnnGetErrorString(status)}, status_{status} {}

void CheckCudnnError(cudnnStatus_t status) {
    if (status != CUDNN_STATUS_SUCCESS) {
        throw CudnnError{status};
    }
}

namespace {

cudnnDataType_t GetCudnnDataType(Dtype dtype) {
    switch (dtype) {
        case Dtype::kFloat32:
            return CUDNN_DATA_FLOAT;
        case Dtype::kFloat64:
            return CUDNN_DATA_DOUBLE;
        // TODO(sonots): Support float16 if it becomes avaialable
        // case Dtype::kFloat16:
        //    return CUDNN_DATA_HALF;
        default:
            throw DtypeError{"Dtype ", dtype, " is not supported in cuDNN"};
    }
}

void SetTensorDescriptor(cudnnTensorDescriptor_t desc, const Array& arr, cudnnTensorFormat_t format) {
    if (!arr.IsContiguous()) {
        throw XchainerError{"CuDNN supports only c-contiguous arrays"};
    }
    cudnnDataType_t cudnn_dtype = GetCudnnDataType(arr.dtype());
    if (arr.shape().ndim() == 4) {
        int n = static_cast<int>(arr.shape()[0]);
        int c = static_cast<int>(arr.shape()[1]);
        int h = static_cast<int>(arr.shape()[2]);
        int w = static_cast<int>(arr.shape()[3]);
        CheckCudnnError(cudnnSetTensor4dDescriptor(desc, format, cudnn_dtype, n, c, h, w));
    } else {
        int64_t item_size = arr.item_size();
        StackVector<int, kMaxNdim> int_shape;
        StackVector<int, kMaxNdim> int_strides;
        for (int64_t v : arr.strides()) {
            int_strides.emplace_back(static_cast<int>(v / item_size));
        }
        for (int64_t v : arr.shape()) {
            int_shape.emplace_back(static_cast<int>(v));
        }
        CheckCudnnError(cudnnSetTensorNdDescriptor(desc, cudnn_dtype, arr.ndim(), &int_shape[0], &int_strides[0]));
    }
}

void SetFilterDescriptor(cudnnFilterDescriptor_t desc, const Array& arr, cudnnTensorFormat_t format) {
    if (!arr.IsContiguous()) {
        throw XchainerError{"CuDNN supports only c-contiguous arrays"};
    }
    cudnnDataType_t cudnn_dtype = GetCudnnDataType(arr.dtype());
    if (arr.shape().ndim() == 4) {
        int n = static_cast<int>(arr.shape()[0]);
        int c = static_cast<int>(arr.shape()[1]);
        int h = static_cast<int>(arr.shape()[2]);
        int w = static_cast<int>(arr.shape()[3]);
        CheckCudnnError(cudnnSetFilter4dDescriptor(desc, cudnn_dtype, format, n, c, h, w));
    } else {
        StackVector<int, kMaxNdim> int_shape;
        for (int64_t v : arr.shape()) {
            int_shape.emplace_back(static_cast<int>(v));
        }
        CheckCudnnError(cudnnSetFilterNdDescriptor(desc, cudnn_dtype, format, arr.ndim(), &int_shape[0]));
    }
}

void SetConvolutionDescriptor(
        cudnnConvolutionDescriptor_t desc,
        const StackVector<int64_t, kMaxNdim>& pad,
        const StackVector<int64_t, kMaxNdim>& stride,
        const nonstd::optional<StackVector<int64_t, kMaxNdim>>& dilation,
        int groups,
        Dtype dtype,
        cudnnConvolutionMode_t mode) {
    size_t ndim = pad.size();
    if (ndim != stride.size()) {
        throw DimensionError{"pad and stride must be of same length"};
    }
    if (dilation && ndim != dilation->size()) {
        throw DimensionError{"pad and dilation must be of same length"};
    }

    StackVector<int, kMaxNdim> int_stride;
    StackVector<int, kMaxNdim> int_pad;
    for (int64_t v : stride) {
        int_stride.emplace_back(static_cast<int>(v));
    }
    for (int64_t v : pad) {
        int_pad.emplace_back(static_cast<int>(v));
    }

    cudnnDataType_t compute_type = GetCudnnDataType(dtype);

    if (ndim != 2) {
        StackVector<int, kMaxNdim> int_dilation;
        if (!dilation) {
            // TODO(sonots): Use assign(ndim, 1) if it becomes available
            for (decltype(ndim) i = 0; i < ndim; ++i) {
                int_dilation.emplace_back(1);
            }
        } else {
            for (int64_t v : *dilation) {
                int_dilation.emplace_back(static_cast<int>(v));
            }
        }
        CheckCudnnError(cudnnSetConvolutionNdDescriptor(desc, ndim, &int_pad[0], &int_stride[0], &int_dilation[0], mode, compute_type));
    } else {
        int d0 = 0, d1 = 0;
        if (!dilation) {
            d0 = d1 = 1;
        } else {
            d0 = static_cast<int>((*dilation)[0]);
            d1 = static_cast<int>((*dilation)[1]);
        }
        int p0 = static_cast<int>(pad[0]);
        int p1 = static_cast<int>(pad[1]);
        int s0 = static_cast<int>(stride[0]);
        int s1 = static_cast<int>(stride[1]);

        CheckCudnnError(cudnnSetConvolution2dDescriptor(desc, p0, p1, s0, s1, d0, d1, mode, compute_type));
    }
    if (groups > 1) {
        CheckCudnnError(cudnnSetConvolutionGroupCount(desc, groups));
    }
}

std::shared_ptr<cudnnTensorStruct> CreateTensorDescriptor(const Array& arr, cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW) {
    cudnnTensorDescriptor_t desc{};
    CheckCudnnError(cudnnCreateTensorDescriptor(&desc));
    auto shared_desc = std::shared_ptr<cudnnTensorStruct>{
            desc, [](cudnnTensorDescriptor_t desc) { CheckCudnnError(cudnnDestroyTensorDescriptor(desc)); }};
    SetTensorDescriptor(desc, arr, format);
    return shared_desc;
}

std::shared_ptr<cudnnFilterStruct> CreateFilterDescriptor(const Array& arr, cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW) {
    cudnnFilterDescriptor_t desc{};
    CheckCudnnError(cudnnCreateFilterDescriptor(&desc));
    auto shared_desc = std::shared_ptr<cudnnFilterStruct>{
            desc, [](cudnnFilterDescriptor_t desc) { CheckCudnnError(cudnnDestroyFilterDescriptor(desc)); }};
    SetFilterDescriptor(desc, arr, format);
    return shared_desc;
}

std::shared_ptr<cudnnConvolutionStruct> CreateConvolutionDescriptor(
        const StackVector<int64_t, kMaxNdim>& pad,
        const StackVector<int64_t, kMaxNdim>& stride,
        Dtype dtype,
        cudnnConvolutionMode_t mode = CUDNN_CROSS_CORRELATION,
        const nonstd::optional<StackVector<int64_t, kMaxNdim>>& dilation = nonstd::nullopt,
        int groups = 1) {
    cudnnConvolutionDescriptor_t desc{};
    CheckCudnnError(cudnnCreateConvolutionDescriptor(&desc));
    auto shared_desc = std::shared_ptr<cudnnConvolutionStruct>{
            desc, [](cudnnConvolutionDescriptor_t desc) { CheckCudnnError(cudnnDestroyConvolutionDescriptor(desc)); }};
    SetConvolutionDescriptor(desc, pad, stride, dilation, groups, dtype, mode);
    return shared_desc;
}

std::pair<cudnnConvolutionFwdAlgo_t, size_t> FindConvolutionForwardAlgorithm(
        cudnnHandle_t handle,
        const std::shared_ptr<cudnnTensorStruct>& x_desc,
        const Array& x,
        const std::shared_ptr<cudnnFilterStruct>& filter_desc,
        const Array& w,
        const std::shared_ptr<cudnnConvolutionStruct>& conv_desc,
        const std::shared_ptr<cudnnTensorStruct>& y_desc,
        const Array& y,
        size_t max_workspace_size) {
    std::shared_ptr<void> workspace = y.device().Allocate(max_workspace_size);

    int requested_algo_count = 1;
    std::vector<cudnnConvolutionFwdAlgoPerf_t> perf_results(requested_algo_count);
    int returned_algo_count;

    CheckCudnnError(cudnnFindConvolutionForwardAlgorithmEx(
            handle,
            x_desc.get(),
            x.data().get(),
            filter_desc.get(),
            w.data().get(),
            conv_desc.get(),
            y_desc.get(),
            y.data().get(),
            requested_algo_count,
            &returned_algo_count,
            &perf_results[0],
            workspace.get(),
            max_workspace_size));

    return {perf_results[0].algo, perf_results[0].memory};
}

}  // namespace

namespace internal {

// TODO(sonots): Support tensor core
void ConvolutionForward(
        const Array& x,
        const Array& w,
        const nonstd::optional<Array>& b,
        Array& y,
        const StackVector<int64_t, kMaxNdim>& pad,
        const StackVector<int64_t, kMaxNdim>& stride,
        const nonstd::optional<StackVector<int64_t, kMaxNdim>>& dilation,
        int groups) {
    CudaDevice& device = *dynamic_cast<CudaDevice*>(&x.device());
    CudaBackend& backend = *dynamic_cast<CudaBackend*>(&device.backend());

    if (&y.device() != &device) {
        throw XchainerError{"The output array device: ", y.device().name(), " must be same with the input array device: ", x.device().name()};
    }
    if (&w.device() != &device) {
        throw XchainerError{"The filter (kernel) array device: ", w.device().name(), " must be same with the input array device: ", x.device().name()};
    }
    // TODO(sonots): Support float16
    if (x.dtype() != Dtype::kFloat32 && x.dtype() != Dtype::kFloat64) {
        throw XchainerError{"XChainer cuDNN supports only float32 or float64 arrays, but the input array dtype is: ", x.dtype()};
    }
    if (y.dtype() != x.dtype()) {
        throw XchainerError{"The output array dtype: ", y.dtype(), " must be same with the input array dtype: ", x.dtype()};
    }
    if (w.dtype() != x.dtype()) {
        throw XchainerError{"The filter (kernel) array dtype: ", w.dtype(), " must be same with the input array dtype: ", x.dtype()};
    }

    float float_zero = 0, float_one = 1;
    double double_zero = 0, double_one = 1;
    void* zero;
    void* one;
    if (x.dtype() == Dtype::kFloat64) {
        zero = &double_zero;
        one = &double_one;
    } else {
        zero = &float_zero;
        one = &float_one;
    }

    cudnnHandle_t handle = device.cudnn_handle();
    Array x_cont = AsContiguousArray(x);
    Array w_cont = AsContiguousArray(w);
    assert(y.IsContiguous());

    std::shared_ptr<cudnnTensorStruct> x_desc = CreateTensorDescriptor(x_cont);
    std::shared_ptr<cudnnTensorStruct> y_desc = CreateTensorDescriptor(y);
    std::shared_ptr<cudnnFilterStruct> filter_desc = CreateFilterDescriptor(w_cont, CUDNN_TENSOR_NCHW);
    std::shared_ptr<cudnnConvolutionStruct> conv_desc =
            CreateConvolutionDescriptor(pad, stride, x.dtype(), CUDNN_CROSS_CORRELATION, dilation, groups);
    size_t max_workspace_size = backend.max_workspace_size();

    // auto tune
    std::pair<cudnnConvolutionFwdAlgo_t, size_t> algo_workspace_size =
            FindConvolutionForwardAlgorithm(handle, x_desc, x_cont, filter_desc, w_cont, conv_desc, y_desc, y, max_workspace_size);
    cudnnConvolutionFwdAlgo_t algo = std::get<0>(algo_workspace_size);
    size_t workspace_size = std::max(max_workspace_size, std::get<1>(algo_workspace_size));
    std::shared_ptr<void> workspace = device.Allocate(workspace_size);

    CheckCudnnError(cudnnConvolutionForward(
            handle,
            one,
            x_desc.get(),
            x_cont.data().get(),
            filter_desc.get(),
            w_cont.data().get(),
            conv_desc.get(),
            algo,
            workspace.get(),
            workspace_size,
            zero,
            y_desc.get(),
            y.data().get()));

    if (b) {
        if (&b->device() != &device) {
            throw XchainerError{"The bias array device: ", b->device().name(), " must be same with the input array device: ", x.device().name()};
        }
        if (b->dtype() != x.dtype()) {
            throw XchainerError{"The bias array dtype: ", b->dtype(), " must be same with the input array dtype: ", x.dtype()};
        }

        int8_t ndim = x.ndim() - 2;
        assert(ndim > 0);

        Shape new_shape;
        new_shape.emplace_back(1);
        new_shape.emplace_back(b->GetTotalSize());
        // TODO(sonots): Use assign(ndim, 1) if it becomes available
        for (int8_t i = 0; i < ndim; ++i) {
            new_shape.emplace_back(1);
        }
        Array b_cont = AsContiguousArray(*b).Reshape(new_shape);
        std::shared_ptr<cudnnTensorStruct> b_desc = CreateTensorDescriptor(b_cont);
        CheckCudnnError(cudnnAddTensor(handle, one, b_desc.get(), b_cont.data().get(), one, y_desc.get(), y.data().get()));
    }
}

}  // namespace internal
}  // namespace cuda
}  // namespace xchainer
<commit_msg>throw if value is larger than INT_MAX<commit_after>#include "xchainer/cuda/cudnn.h"

#include <algorithm>
#include <limits>
#include <memory>
#include <utility>
#include <vector>

#include <cudnn.h>
#include <nonstd/optional.hpp>

#include "xchainer/array.h"
#include "xchainer/cuda/cuda_device.h"
#include "xchainer/device.h"
#include "xchainer/dtype.h"
#include "xchainer/error.h"
#include "xchainer/routines/creation.h"
#include "xchainer/stack_vector.h"

namespace xchainer {
namespace cuda {

CudnnError::CudnnError(cudnnStatus_t status) : XchainerError{cudnnGetErrorString(status)}, status_{status} {}

void CheckCudnnError(cudnnStatus_t status) {
    if (status != CUDNN_STATUS_SUCCESS) {
        throw CudnnError{status};
    }
}

namespace {

cudnnDataType_t GetCudnnDataType(Dtype dtype) {
    switch (dtype) {
        case Dtype::kFloat32:
            return CUDNN_DATA_FLOAT;
        case Dtype::kFloat64:
            return CUDNN_DATA_DOUBLE;
        // TODO(sonots): Support float16 if it becomes avaialable
        // case Dtype::kFloat16:
        //    return CUDNN_DATA_HALF;
        default:
            throw DtypeError{"Dtype ", dtype, " is not supported in cuDNN"};
    }
}

void SetTensorDescriptor(cudnnTensorDescriptor_t desc, const Array& arr, cudnnTensorFormat_t format) {
    if (!arr.IsContiguous()) {
        throw XchainerError{"CuDNN supports only c-contiguous arrays"};
    }
    constexpr int int_max = std::numeric_limits<int>::max();
    for (int8_t i = 0; i < arr.shape().ndim(); ++i) {
        if (arr.shape()[i] > int_max) {
            throw XchainerError{"The shape size: ", arr.shape()[i], " of dimension: ", i, " is larger than ", int_max};
        }
    }

    cudnnDataType_t cudnn_dtype = GetCudnnDataType(arr.dtype());
    if (arr.shape().ndim() == 4) {
        int n = static_cast<int>(arr.shape()[0]);
        int c = static_cast<int>(arr.shape()[1]);
        int h = static_cast<int>(arr.shape()[2]);
        int w = static_cast<int>(arr.shape()[3]);
        CheckCudnnError(cudnnSetTensor4dDescriptor(desc, format, cudnn_dtype, n, c, h, w));
    } else {
        StackVector<int, kMaxNdim> int_strides;
        for (int8_t i = 0; i < arr.strides().ndim(); ++i) {
            int64_t v = arr.strides()[i] / arr.item_size();
            if (v > int_max) {
                throw XchainerError{"The array stride: ", v, " (in number of items) of dimension: ", i, " is larger than ", int_max};
            }
            int_strides.emplace_back(static_cast<int>(v));
        }
        StackVector<int, kMaxNdim> int_shape;
        for (int64_t v : arr.shape()) {
            int_shape.emplace_back(static_cast<int>(v));
        }
        CheckCudnnError(cudnnSetTensorNdDescriptor(desc, cudnn_dtype, arr.ndim(), &int_shape[0], &int_strides[0]));
    }
}

void SetFilterDescriptor(cudnnFilterDescriptor_t desc, const Array& arr, cudnnTensorFormat_t format) {
    if (!arr.IsContiguous()) {
        throw XchainerError{"CuDNN supports only c-contiguous arrays"};
    }
    constexpr int int_max = std::numeric_limits<int>::max();
    for (int8_t i = 0; i < arr.shape().ndim(); ++i) {
        if (arr.shape()[i] > int_max) {
            throw XchainerError{"The shape size: ", arr.shape()[i], " of dimension: ", i, " is larger than ", int_max};
        }
    }

    cudnnDataType_t cudnn_dtype = GetCudnnDataType(arr.dtype());
    if (arr.shape().ndim() == 4) {
        int n = static_cast<int>(arr.shape()[0]);
        int c = static_cast<int>(arr.shape()[1]);
        int h = static_cast<int>(arr.shape()[2]);
        int w = static_cast<int>(arr.shape()[3]);
        CheckCudnnError(cudnnSetFilter4dDescriptor(desc, cudnn_dtype, format, n, c, h, w));
    } else {
        StackVector<int, kMaxNdim> int_shape;
        for (int64_t v : arr.shape()) {
            int_shape.emplace_back(static_cast<int>(v));
        }
        CheckCudnnError(cudnnSetFilterNdDescriptor(desc, cudnn_dtype, format, arr.ndim(), &int_shape[0]));
    }
}

void SetConvolutionDescriptor(
        cudnnConvolutionDescriptor_t desc,
        const StackVector<int64_t, kMaxNdim>& pad,
        const StackVector<int64_t, kMaxNdim>& stride,
        const nonstd::optional<StackVector<int64_t, kMaxNdim>>& dilation,
        int groups,
        Dtype dtype,
        cudnnConvolutionMode_t mode) {
    size_t ndim = pad.size();
    if (ndim != stride.size()) {
        throw DimensionError{"pad and stride must be of same length"};
    }
    if (dilation && ndim != dilation->size()) {
        throw DimensionError{"pad and dilation must be of same length"};
    }

    constexpr int int_max = std::numeric_limits<int>::max();
    StackVector<int, kMaxNdim> int_stride;
    for (size_t i = 0; i < stride.size(); ++i) {
        if (stride[i] > int_max) {
            throw XchainerError{"The stride: ", stride[i], " of dimension: ", i, " is larger than ", int_max};
        }
        int_stride.emplace_back(static_cast<int>(stride[i]));
    }
    StackVector<int, kMaxNdim> int_pad;
    for (size_t i = 0; i < pad.size(); ++i) {
        if (pad[i] > int_max) {
            throw XchainerError{"The pad: ", pad[i], " of dimension: ", i, " is larger than ", int_max};
        }
        int_pad.emplace_back(static_cast<int>(pad[i]));
    }
    if (dilation) {
        for (size_t i = 0; i < dilation->size(); ++i) {
            if ((*dilation)[i] > int_max) {
                throw XchainerError{"The dilation, ", (*dilation)[i], " of dimension: ", i, " is larger than ", int_max};
            }
        }
    }

    cudnnDataType_t compute_type = GetCudnnDataType(dtype);

    if (ndim != 2) {
        StackVector<int, kMaxNdim> int_dilation;
        if (!dilation) {
            // TODO(sonots): Use assign(ndim, 1) if it becomes available
            for (size_t i = 0; i < ndim; ++i) {
                int_dilation.emplace_back(1);
            }
        } else {
            for (int64_t v : *dilation) {
                int_dilation.emplace_back(static_cast<int>(v));
            }
        }
        CheckCudnnError(cudnnSetConvolutionNdDescriptor(desc, ndim, &int_pad[0], &int_stride[0], &int_dilation[0], mode, compute_type));
    } else {
        int d0 = 0, d1 = 0;
        if (!dilation) {
            d0 = d1 = 1;
        } else {
            d0 = static_cast<int>((*dilation)[0]);
            d1 = static_cast<int>((*dilation)[1]);
        }
        int p0 = static_cast<int>(pad[0]);
        int p1 = static_cast<int>(pad[1]);
        int s0 = static_cast<int>(stride[0]);
        int s1 = static_cast<int>(stride[1]);

        CheckCudnnError(cudnnSetConvolution2dDescriptor(desc, p0, p1, s0, s1, d0, d1, mode, compute_type));
    }
    if (groups > 1) {
        CheckCudnnError(cudnnSetConvolutionGroupCount(desc, groups));
    }
}

std::shared_ptr<cudnnTensorStruct> CreateTensorDescriptor(const Array& arr, cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW) {
    cudnnTensorDescriptor_t desc{};
    CheckCudnnError(cudnnCreateTensorDescriptor(&desc));
    auto shared_desc = std::shared_ptr<cudnnTensorStruct>{
            desc, [](cudnnTensorDescriptor_t desc) { CheckCudnnError(cudnnDestroyTensorDescriptor(desc)); }};
    SetTensorDescriptor(desc, arr, format);
    return shared_desc;
}

std::shared_ptr<cudnnFilterStruct> CreateFilterDescriptor(const Array& arr, cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW) {
    cudnnFilterDescriptor_t desc{};
    CheckCudnnError(cudnnCreateFilterDescriptor(&desc));
    auto shared_desc = std::shared_ptr<cudnnFilterStruct>{
            desc, [](cudnnFilterDescriptor_t desc) { CheckCudnnError(cudnnDestroyFilterDescriptor(desc)); }};
    SetFilterDescriptor(desc, arr, format);
    return shared_desc;
}

std::shared_ptr<cudnnConvolutionStruct> CreateConvolutionDescriptor(
        const StackVector<int64_t, kMaxNdim>& pad,
        const StackVector<int64_t, kMaxNdim>& stride,
        Dtype dtype,
        cudnnConvolutionMode_t mode = CUDNN_CROSS_CORRELATION,
        const nonstd::optional<StackVector<int64_t, kMaxNdim>>& dilation = nonstd::nullopt,
        int groups = 1) {
    cudnnConvolutionDescriptor_t desc{};
    CheckCudnnError(cudnnCreateConvolutionDescriptor(&desc));
    auto shared_desc = std::shared_ptr<cudnnConvolutionStruct>{
            desc, [](cudnnConvolutionDescriptor_t desc) { CheckCudnnError(cudnnDestroyConvolutionDescriptor(desc)); }};
    SetConvolutionDescriptor(desc, pad, stride, dilation, groups, dtype, mode);
    return shared_desc;
}

std::pair<cudnnConvolutionFwdAlgo_t, size_t> FindConvolutionForwardAlgorithm(
        cudnnHandle_t handle,
        const std::shared_ptr<cudnnTensorStruct>& x_desc,
        const Array& x,
        const std::shared_ptr<cudnnFilterStruct>& filter_desc,
        const Array& w,
        const std::shared_ptr<cudnnConvolutionStruct>& conv_desc,
        const std::shared_ptr<cudnnTensorStruct>& y_desc,
        const Array& y,
        size_t max_workspace_size) {
    std::shared_ptr<void> workspace = y.device().Allocate(max_workspace_size);

    int requested_algo_count = 1;
    std::vector<cudnnConvolutionFwdAlgoPerf_t> perf_results(requested_algo_count);
    int returned_algo_count;

    CheckCudnnError(cudnnFindConvolutionForwardAlgorithmEx(
            handle,
            x_desc.get(),
            x.data().get(),
            filter_desc.get(),
            w.data().get(),
            conv_desc.get(),
            y_desc.get(),
            y.data().get(),
            requested_algo_count,
            &returned_algo_count,
            &perf_results[0],
            workspace.get(),
            max_workspace_size));

    return {perf_results[0].algo, perf_results[0].memory};
}

}  // namespace

namespace internal {

// TODO(sonots): Support tensor core
void ConvolutionForward(
        const Array& x,
        const Array& w,
        const nonstd::optional<Array>& b,
        Array& y,
        const StackVector<int64_t, kMaxNdim>& pad,
        const StackVector<int64_t, kMaxNdim>& stride,
        const nonstd::optional<StackVector<int64_t, kMaxNdim>>& dilation,
        int groups) {
    CudaDevice& device = *dynamic_cast<CudaDevice*>(&x.device());
    CudaBackend& backend = *dynamic_cast<CudaBackend*>(&device.backend());

    if (&y.device() != &device) {
        throw XchainerError{
                "The output array device: ", y.device().name(), " must be same with the input array device: ", x.device().name()};
    }
    if (&w.device() != &device) {
        throw XchainerError{
                "The filter (kernel) array device: ", w.device().name(), " must be same with the input array device: ", x.device().name()};
    }
    // TODO(sonots): Support float16
    if (x.dtype() != Dtype::kFloat32 && x.dtype() != Dtype::kFloat64) {
        throw XchainerError{"XChainer cuDNN supports only float32 or float64 arrays, but the input array dtype is: ", x.dtype()};
    }
    if (y.dtype() != x.dtype()) {
        throw XchainerError{"The output array dtype: ", y.dtype(), " must be same with the input array dtype: ", x.dtype()};
    }
    if (w.dtype() != x.dtype()) {
        throw XchainerError{"The filter (kernel) array dtype: ", w.dtype(), " must be same with the input array dtype: ", x.dtype()};
    }

    float float_zero = 0, float_one = 1;
    double double_zero = 0, double_one = 1;
    void* zero;
    void* one;
    if (x.dtype() == Dtype::kFloat64) {
        zero = &double_zero;
        one = &double_one;
    } else {
        zero = &float_zero;
        one = &float_one;
    }

    cudnnHandle_t handle = device.cudnn_handle();
    Array x_cont = AsContiguousArray(x);
    Array w_cont = AsContiguousArray(w);
    assert(y.IsContiguous());

    std::shared_ptr<cudnnTensorStruct> x_desc = CreateTensorDescriptor(x_cont);
    std::shared_ptr<cudnnTensorStruct> y_desc = CreateTensorDescriptor(y);
    std::shared_ptr<cudnnFilterStruct> filter_desc = CreateFilterDescriptor(w_cont, CUDNN_TENSOR_NCHW);
    std::shared_ptr<cudnnConvolutionStruct> conv_desc =
            CreateConvolutionDescriptor(pad, stride, x.dtype(), CUDNN_CROSS_CORRELATION, dilation, groups);
    size_t max_workspace_size = backend.max_workspace_size();

    // auto tune
    std::pair<cudnnConvolutionFwdAlgo_t, size_t> algo_workspace_size =
            FindConvolutionForwardAlgorithm(handle, x_desc, x_cont, filter_desc, w_cont, conv_desc, y_desc, y, max_workspace_size);
    cudnnConvolutionFwdAlgo_t algo = std::get<0>(algo_workspace_size);
    size_t workspace_size = std::max(max_workspace_size, std::get<1>(algo_workspace_size));
    std::shared_ptr<void> workspace = device.Allocate(workspace_size);

    CheckCudnnError(cudnnConvolutionForward(
            handle,
            one,
            x_desc.get(),
            x_cont.data().get(),
            filter_desc.get(),
            w_cont.data().get(),
            conv_desc.get(),
            algo,
            workspace.get(),
            workspace_size,
            zero,
            y_desc.get(),
            y.data().get()));

    if (b) {
        if (&b->device() != &device) {
            throw XchainerError{
                    "The bias array device: ", b->device().name(), " must be same with the input array device: ", x.device().name()};
        }
        if (b->dtype() != x.dtype()) {
            throw XchainerError{"The bias array dtype: ", b->dtype(), " must be same with the input array dtype: ", x.dtype()};
        }

        int8_t ndim = x.ndim() - 2;
        assert(ndim > 0);

        Shape new_shape;
        new_shape.emplace_back(1);
        new_shape.emplace_back(b->GetTotalSize());
        // TODO(sonots): Use assign(ndim, 1) if it becomes available
        for (int8_t i = 0; i < ndim; ++i) {
            new_shape.emplace_back(1);
        }
        Array b_cont = AsContiguousArray(*b).Reshape(new_shape);
        std::shared_ptr<cudnnTensorStruct> b_desc = CreateTensorDescriptor(b_cont);
        CheckCudnnError(cudnnAddTensor(handle, one, b_desc.get(), b_cont.data().get(), one, y_desc.get(), y.data().get()));
    }
}

}  // namespace internal
}  // namespace cuda
}  // namespace xchainer
<|endoftext|>