<commit_before>/*
* Copyright (C) 2010 Toni Gundogdu.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "config.h"

#include <fstream>
#include <stdexcept>
#include <sstream>
#include <iomanip>

#include <boost/format.hpp>
#include <boost/foreach.hpp>
#include <boost/filesystem.hpp>

#ifndef foreach
#define foreach BOOST_FOREACH
#endif

#include <curl/curl.h>

#include "quvicpp/quvicpp.h"

#include "cclive/options.h"
#include "cclive/progressbar.h"
#include "cclive/re.h"
#include "cclive/error.h"
#include "cclive/log.h"
#include "cclive/file.h"

namespace cclive {

file::file ()
    : _initial_length(0)
    { }

file::file (
    const quvicpp::video& video,
    const quvicpp::link& link,
    const int n,
    const options& opts)
    : _initial_length(0)
    { _init (video, link, n, opts); }

file::file (const file& f)
    : _initial_length(0)
    { _swap(f); }

file&
file::operator=(const file& f) {
    if (this != &f) _swap(f);
    return *this;
}

void
file::_swap (const file& f) {
    _name           = f._name;
    _path           = f._path;
    _initial_length = f._initial_length;
}

#define E "server response code %ld, expected 200 or 206 (conn_code=%ld)"

static std::string
format_unexpected_http_error (const long resp_code, const long conn_code)
    { return (boost::format (E) % resp_code % conn_code).str (); }

#undef E

#define E "%s (curl_code=%ld, resp_code=%ld, conn_code=%ld)"

static std::string
format_error (
    const CURLcode curl_code,
    const long resp_code,
    const long conn_code)
{
    const std::string e = curl_easy_strerror (curl_code);
    return (boost::format (E) % e % curl_code % resp_code % conn_code).str ();
}

#undef E

static size_t
write_cb (void *data, size_t size, size_t nmemb, void *ptr) {
    std::ofstream *o   = reinterpret_cast<std::ofstream*>(ptr);
    const size_t rsize = size*nmemb;
    o->write (static_cast<char*>(data), rsize);
    o->flush ();
    return rsize;
}

static int
progress_cb (void *ptr, double, double now, double, double) {
    reinterpret_cast<progressbar*>(ptr)->update (now);
    return 0;
}

namespace po = boost::program_options;

bool
file::write (
    const quvicpp::query& q,
    const quvicpp::link& l,
    const options& opts) const
{
    CURL *curl = q.curlHandle();

    curl_easy_setopt (curl, CURLOPT_URL, l.url().c_str());
    curl_easy_setopt (curl, CURLOPT_WRITEFUNCTION, write_cb);

    const po::variables_map map  = opts.map();
    std::ios_base::openmode mode = std::ofstream::binary;

    if ( map.count ("overwrite") )
        mode |= std::ofstream::trunc;
    else {
        if ( _should_continue() )
            mode |= std::ofstream::app;
    }

    std::ofstream out( _path.c_str(), mode );

    if ( out.fail() ) {

        std::string s = _path + ": ";

        if (errno)
            s += cclive::perror ();
        else
            s += "unknown file open error";

        throw std::runtime_error (s);
    }

    curl_easy_setopt (curl, CURLOPT_WRITEDATA, &out);

    curl_easy_setopt (curl, CURLOPT_ENCODING, "identity");
    curl_easy_setopt (curl, CURLOPT_HEADER, 0L);

    progressbar pb (*this, l, opts);
    curl_easy_setopt (curl, CURLOPT_PROGRESSDATA, &pb);
    curl_easy_setopt (curl, CURLOPT_NOPROGRESS, 0L);
    curl_easy_setopt (curl, CURLOPT_PROGRESSFUNCTION, progress_cb);

    curl_off_t resume_from = static_cast<curl_off_t>(_initial_length);
    curl_easy_setopt (curl, CURLOPT_RESUME_FROM_LARGE, resume_from);

    curl_easy_setopt (curl, CURLOPT_CONNECTTIMEOUT,
        map["connect-timeout"].as<int>());

    curl_easy_setopt (curl, CURLOPT_TIMEOUT,
        map["transfer-timeout"].as<int>());

    curl_easy_setopt (curl, CURLOPT_DNS_CACHE_TIMEOUT,
        map["dns-cache-timeout"].as<int>());

    curl_off_t throttle = map["throttle"].as<int>()*1024;
    curl_easy_setopt (curl, CURLOPT_MAX_RECV_SPEED_LARGE, throttle);

    const CURLcode rc = curl_easy_perform (curl);

    // Restore curl settings.

    curl_easy_setopt (curl, CURLOPT_HEADER, 1L);
    curl_easy_setopt (curl, CURLOPT_NOPROGRESS, 1L);
    curl_easy_setopt (curl, CURLOPT_RESUME_FROM_LARGE, 0L);
    curl_easy_setopt (curl, CURLOPT_MAX_RECV_SPEED_LARGE,
        static_cast<curl_off_t>(0L));

    out.flush();
    out.close();

    long resp_code = 0;
    long conn_code = 0;

    curl_easy_getinfo (curl,
        CURLINFO_RESPONSE_CODE, &resp_code);

    curl_easy_getinfo (curl,
        CURLINFO_HTTP_CONNECTCODE, &conn_code);

    std::string error;

    if (CURLE_OK == rc) {
        if (resp_code != 200 && resp_code != 206)
            error = format_unexpected_http_error (resp_code, conn_code);
    }
    else {
        error = format_error (rc, resp_code, conn_code);
    }

    if ( !error.empty() ) {

        cclive::log << std::endl;

        if (resp_code >= 400 && resp_code <= 500)
            throw std::runtime_error (error);
        else
            std::clog << "error: " << error << std::endl;

        return false; // Retry.
    }

    pb.finish();

    cclive::log << std::endl;

    return true;
}

double file::initial_length   () const { return _initial_length; }
const std::string& file::path () const { return _path; }
bool file::_should_continue   () const { return _initial_length > 0; }

static double
to_mb (const double bytes) { return bytes/(1024*1024); }

std::string
file::to_s (const quvicpp::link& link) const {

    const double length = to_mb (link.length ());

    boost::format fmt = boost::format("%s  %.2fM  [%s]")
        % _name % length % link.content_type ();

    return fmt.str();
}

namespace fs = boost::filesystem;

void
file::_init (
    const quvicpp::video& video,
    const quvicpp::link& link,
    const int n,
    const options& opts)
{
    const po::variables_map map = opts.map();

    if (map.count("output-file")) {

        const std::string output_file =
            map["output-file"].as<std::string>();

        const fs::path p =
            fs::system_complete (fs::path(output_file));

        _name           = p.filename();
        _path           = p.string();
        _initial_length = file::exists (_path);

        if ( _initial_length >= link.length() && !map.count("overwrite") )
            throw std::runtime_error ("nothing todo");
    }

    else {

        std::string title  = video.title();

        // Apply --regexp to title.

        cclive::re::match (map["regexp"].as<std::string>(), title);

        cclive::re::trim (title);

        // --filename-format

        std::string fname_format =
            map["filename-format"].as<std::string>();

        boost::format fmt;

        fmt = boost::format("s{%%i}{%1%}g") % video.id();
        cclive::re::subst (fmt.str(), fname_format);

        fmt = boost::format("s{%%t}{%1%}g") % title;
        cclive::re::subst (fmt.str(), fname_format);

        fmt = boost::format("s{%%s}{%1%}g") % link.suffix();
        cclive::re::subst (fmt.str(), fname_format);

        fmt = boost::format("s{%%h}{%1%}g") % video.host();
        cclive::re::subst (fmt.str(), fname_format);

        if (map.count("subst")) {

            std::istringstream iss(map["subst"].as<std::string>());
            std::vector<std::string> v;

            std::copy(
                std::istream_iterator<std::string >(iss),
                std::istream_iterator<std::string >(),
                std::back_inserter<std::vector<std::string> >(v)
            );

            foreach (std::string s, v)
                cclive::re::subst (s, fname_format);
        }

        std::stringstream b;

        b << fname_format;

        // A multi-segment video.

        if (n > 1) b << "_" << n;

        // Path, name.

        fs::path p( fs::system_complete (b.str()) );

        _name = p.filename();
        _path = p.string();

        if ( !map.count ("overwrite") ) {

            for (int i=0; i<INT_MAX; ++i) {

                _initial_length = file::exists (_path);

                if (!_initial_length)
                    break;

                else if (_initial_length >= link.length())
                    throw std::runtime_error ("nothing todo");

                else {
                    if (map.count("continue"))
                        break;

                }

                boost::format fmt =
                    boost::format ("%1%.%2%") % b.str() % i;

                p     = fs::system_complete( fmt.str() );
                _name = p.filename();
                _path = p.string();
            }
        }
    }

    if ( map.count("overwrite") )
        _initial_length = 0;
}

double
file::exists (const std::string& path) {

    fs::path p( fs::system_complete(path) );

    double size = 0;

    if (fs::exists(p))
        size = static_cast<double>(fs::file_size(p));

    return size;
}

} // End namespace.

<commit_msg>change "nothing todo" to more understandable error.<commit_after>/*
* Copyright (C) 2010 Toni Gundogdu.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "config.h"

#include <fstream>
#include <stdexcept>
#include <sstream>
#include <iomanip>

#include <boost/format.hpp>
#include <boost/foreach.hpp>
#include <boost/filesystem.hpp>

#ifndef foreach
#define foreach BOOST_FOREACH
#endif

#include <curl/curl.h>

#include "quvicpp/quvicpp.h"

#include "cclive/options.h"
#include "cclive/progressbar.h"
#include "cclive/re.h"
#include "cclive/error.h"
#include "cclive/log.h"
#include "cclive/file.h"

namespace cclive {

file::file ()
    : _initial_length(0)
    { }

file::file (
    const quvicpp::video& video,
    const quvicpp::link& link,
    const int n,
    const options& opts)
    : _initial_length(0)
    { _init (video, link, n, opts); }

file::file (const file& f)
    : _initial_length(0)
    { _swap(f); }

file&
file::operator=(const file& f) {
    if (this != &f) _swap(f);
    return *this;
}

void
file::_swap (const file& f) {
    _name           = f._name;
    _path           = f._path;
    _initial_length = f._initial_length;
}

#define E "server response code %ld, expected 200 or 206 (conn_code=%ld)"

static std::string
format_unexpected_http_error (const long resp_code, const long conn_code)
    { return (boost::format (E) % resp_code % conn_code).str (); }

#undef E

#define E "%s (curl_code=%ld, resp_code=%ld, conn_code=%ld)"

static std::string
format_error (
    const CURLcode curl_code,
    const long resp_code,
    const long conn_code)
{
    const std::string e = curl_easy_strerror (curl_code);
    return (boost::format (E) % e % curl_code % resp_code % conn_code).str ();
}

#undef E

static size_t
write_cb (void *data, size_t size, size_t nmemb, void *ptr) {
    std::ofstream *o   = reinterpret_cast<std::ofstream*>(ptr);
    const size_t rsize = size*nmemb;
    o->write (static_cast<char*>(data), rsize);
    o->flush ();
    return rsize;
}

static int
progress_cb (void *ptr, double, double now, double, double) {
    reinterpret_cast<progressbar*>(ptr)->update (now);
    return 0;
}

namespace po = boost::program_options;

bool
file::write (
    const quvicpp::query& q,
    const quvicpp::link& l,
    const options& opts) const
{
    CURL *curl = q.curlHandle();

    curl_easy_setopt (curl, CURLOPT_URL, l.url().c_str());
    curl_easy_setopt (curl, CURLOPT_WRITEFUNCTION, write_cb);

    const po::variables_map map  = opts.map();
    std::ios_base::openmode mode = std::ofstream::binary;

    if ( map.count ("overwrite") )
        mode |= std::ofstream::trunc;
    else {
        if ( _should_continue() )
            mode |= std::ofstream::app;
    }

    std::ofstream out( _path.c_str(), mode );

    if ( out.fail() ) {

        std::string s = _path + ": ";

        if (errno)
            s += cclive::perror ();
        else
            s += "unknown file open error";

        throw std::runtime_error (s);
    }

    curl_easy_setopt (curl, CURLOPT_WRITEDATA, &out);

    curl_easy_setopt (curl, CURLOPT_ENCODING, "identity");
    curl_easy_setopt (curl, CURLOPT_HEADER, 0L);

    progressbar pb (*this, l, opts);
    curl_easy_setopt (curl, CURLOPT_PROGRESSDATA, &pb);
    curl_easy_setopt (curl, CURLOPT_NOPROGRESS, 0L);
    curl_easy_setopt (curl, CURLOPT_PROGRESSFUNCTION, progress_cb);

    curl_off_t resume_from = static_cast<curl_off_t>(_initial_length);
    curl_easy_setopt (curl, CURLOPT_RESUME_FROM_LARGE, resume_from);

    curl_easy_setopt (curl, CURLOPT_CONNECTTIMEOUT,
        map["connect-timeout"].as<int>());

    curl_easy_setopt (curl, CURLOPT_TIMEOUT,
        map["transfer-timeout"].as<int>());

    curl_easy_setopt (curl, CURLOPT_DNS_CACHE_TIMEOUT,
        map["dns-cache-timeout"].as<int>());

    curl_off_t throttle = map["throttle"].as<int>()*1024;
    curl_easy_setopt (curl, CURLOPT_MAX_RECV_SPEED_LARGE, throttle);

    const CURLcode rc = curl_easy_perform (curl);

    // Restore curl settings.

    curl_easy_setopt (curl, CURLOPT_HEADER, 1L);
    curl_easy_setopt (curl, CURLOPT_NOPROGRESS, 1L);
    curl_easy_setopt (curl, CURLOPT_RESUME_FROM_LARGE, 0L);
    curl_easy_setopt (curl, CURLOPT_MAX_RECV_SPEED_LARGE,
        static_cast<curl_off_t>(0L));

    out.flush();
    out.close();

    long resp_code = 0;
    long conn_code = 0;

    curl_easy_getinfo (curl,
        CURLINFO_RESPONSE_CODE, &resp_code);

    curl_easy_getinfo (curl,
        CURLINFO_HTTP_CONNECTCODE, &conn_code);

    std::string error;

    if (CURLE_OK == rc) {
        if (resp_code != 200 && resp_code != 206)
            error = format_unexpected_http_error (resp_code, conn_code);
    }
    else {
        error = format_error (rc, resp_code, conn_code);
    }

    if ( !error.empty() ) {

        cclive::log << std::endl;

        if (resp_code >= 400 && resp_code <= 500)
            throw std::runtime_error (error);
        else
            std::clog << "error: " << error << std::endl;

        return false; // Retry.
    }

    pb.finish();

    cclive::log << std::endl;

    return true;
}

double file::initial_length   () const { return _initial_length; }
const std::string& file::path () const { return _path; }
bool file::_should_continue   () const { return _initial_length > 0; }

static double
to_mb (const double bytes) { return bytes/(1024*1024); }

std::string
file::to_s (const quvicpp::link& link) const {

    const double length = to_mb (link.length ());

    boost::format fmt = boost::format("%s  %.2fM  [%s]")
        % _name % length % link.content_type ();

    return fmt.str();
}

#define E_NOTHINGTODO "video completely retrieved already"

namespace fs = boost::filesystem;

void
file::_init (
    const quvicpp::video& video,
    const quvicpp::link& link,
    const int n,
    const options& opts)
{
    const po::variables_map map = opts.map();

    if (map.count("output-file")) {

        const std::string output_file =
            map["output-file"].as<std::string>();

        const fs::path p =
            fs::system_complete (fs::path(output_file));

        _name           = p.filename();
        _path           = p.string();
        _initial_length = file::exists (_path);

        if ( _initial_length >= link.length() && !map.count("overwrite") )
            throw std::runtime_error (E_NOTHINGTODO);
    }

    else {

        std::string title  = video.title();

        // Apply --regexp to title.

        cclive::re::match (map["regexp"].as<std::string>(), title);

        cclive::re::trim (title);

        // --filename-format

        std::string fname_format =
            map["filename-format"].as<std::string>();

        boost::format fmt;

        fmt = boost::format("s{%%i}{%1%}g") % video.id();
        cclive::re::subst (fmt.str(), fname_format);

        fmt = boost::format("s{%%t}{%1%}g") % title;
        cclive::re::subst (fmt.str(), fname_format);

        fmt = boost::format("s{%%s}{%1%}g") % link.suffix();
        cclive::re::subst (fmt.str(), fname_format);

        fmt = boost::format("s{%%h}{%1%}g") % video.host();
        cclive::re::subst (fmt.str(), fname_format);

        if (map.count("subst")) {

            std::istringstream iss(map["subst"].as<std::string>());
            std::vector<std::string> v;

            std::copy(
                std::istream_iterator<std::string >(iss),
                std::istream_iterator<std::string >(),
                std::back_inserter<std::vector<std::string> >(v)
            );

            foreach (std::string s, v)
                cclive::re::subst (s, fname_format);
        }

        std::stringstream b;

        b << fname_format;

        // A multi-segment video.

        if (n > 1) b << "_" << n;

        // Path, name.

        fs::path p( fs::system_complete (b.str()) );

        _name = p.filename();
        _path = p.string();

        if ( !map.count ("overwrite") ) {

            for (int i=0; i<INT_MAX; ++i) {

                _initial_length = file::exists (_path);

                if (!_initial_length)
                    break;

                else if (_initial_length >= link.length())
                    throw std::runtime_error (E_NOTHINGTODO);

                else {
                    if (map.count("continue"))
                        break;

                }

                boost::format fmt =
                    boost::format ("%1%.%2%") % b.str() % i;

                p     = fs::system_complete( fmt.str() );
                _name = p.filename();
                _path = p.string();
            }
        }
    }

    if ( map.count("overwrite") )
        _initial_length = 0;
}

#undef E_NOTHINGTODO

double
file::exists (const std::string& path) {

    fs::path p( fs::system_complete(path) );

    double size = 0;

    if (fs::exists(p))
        size = static_cast<double>(fs::file_size(p));

    return size;
}

} // End namespace.

<|endoftext|>


<commit_before>/* -*-c++-*- OpenSceneGraph - Copyright (C) 1998-2008 Robert Osfield
 *
 * This library is open source and may be redistributed and/or modified under
 * the terms of the OpenSceneGraph Public License (OSGPL) version 0.0 or
 * (at your option) any later version.  The full license is in LICENSE file
 * included with this distribution, and on the openscenegraph.org website.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * OpenSceneGraph Public License for more details.
*/

#include <osgDB/FileUtils>
#include <osgDB/FileNameUtils>
#include <osgDB/ReadFile>
#include <osgDB/WriteFile>
#include <osgDB/Registry>

#include <iostream>
#include <sstream>
#include <fstream>
#include <iterator>

#include <zlib.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
//
//  ReaderWriterGZ
//
class ReaderWriterGZ : public osgDB::ReaderWriter
{
    public:

        enum ObjectType
        {
            OBJECT,
            ARCHIVE,
            IMAGE,
            HEIGHTFIELD,
            NODE
        };

        ReaderWriterGZ();

        ~ReaderWriterGZ();

        virtual const char* className() const { return "HTTP Protocol Model Reader"; }

        virtual ReadResult openArchive(const std::string& fileName,ArchiveStatus status, unsigned int , const Options* options) const
        {
            if (status!=READ) return ReadResult(ReadResult::FILE_NOT_HANDLED);
            else return readFile(ARCHIVE,fileName,options);
        }

        virtual ReadResult readObject(const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return readFile(OBJECT,fileName,options);
        }

        virtual ReadResult readImage(const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return readFile(IMAGE,fileName,options);
        }

        virtual ReadResult readHeightField(const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return readFile(HEIGHTFIELD,fileName,options);
        }

        virtual ReadResult readNode(const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return readFile(NODE,fileName,options);
        }

        ReadResult readFile(ObjectType objectType, osgDB::ReaderWriter* rw, std::istream& fin, const osgDB::ReaderWriter::Options* options) const;

        ReadResult readFile(ObjectType objectType, const std::string& fullFileName, const osgDB::ReaderWriter::Options* options) const;

        virtual WriteResult writeObject(const osg::Object& obj, const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return writeFile(OBJECT, &obj, fileName, options);
        }

        virtual WriteResult writeImage(const osg::Image& image, const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return writeFile(IMAGE, &image, fileName, options);
        }

        virtual WriteResult writeHeightField(const osg::HeightField& hf, const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return writeFile(HEIGHTFIELD, &hf, fileName, options);
        }

        virtual WriteResult writeNode(const osg::Node& node, const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return writeFile(NODE, &node, fileName,options);
        }

        WriteResult writeFile(ObjectType objectType, const osg::Object* object, osgDB::ReaderWriter* rw, std::ostream& fin, const osgDB::ReaderWriter::Options* options) const;

        WriteResult writeFile(ObjectType objectType, const osg::Object* object, const std::string& fullFileName, const osgDB::ReaderWriter::Options* options) const;

        bool read(std::istream& fin, std::string& destination) const;
        bool write(std::ostream& fout, const std::string& source) const;

};

ReaderWriterGZ::ReaderWriterGZ()
{
    // OSG_NOTICE<<"ReaderWriterGZ::ReaderWriterGZ()"<<std::endl;

    supportsExtension("osgz","Compressed .osg file extension.");
    supportsExtension("ivez","Compressed .ive file extension.");
    supportsExtension("gz","Compressed file extension.");
}

ReaderWriterGZ::~ReaderWriterGZ()
{
    // OSG_NOTICE<<"ReaderWriterGZ::~ReaderWriterGZ()"<<std::endl;
}

osgDB::ReaderWriter::ReadResult ReaderWriterGZ::readFile(ObjectType objectType, osgDB::ReaderWriter* rw, std::istream& fin, const osgDB::ReaderWriter::Options *options) const
{
    switch(objectType)
    {
        case(OBJECT): return rw->readObject(fin,options);
        case(ARCHIVE): return rw->openArchive(fin,options);
        case(IMAGE): return rw->readImage(fin,options);
        case(HEIGHTFIELD): return rw->readHeightField(fin,options);
        case(NODE): return rw->readNode(fin,options);
        default: break;
    }
    return ReadResult::FILE_NOT_HANDLED;
}

osgDB::ReaderWriter::ReadResult ReaderWriterGZ::readFile(ObjectType objectType, const std::string& fullFileName, const osgDB::ReaderWriter::Options *options) const
{
    std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
    if (!acceptsExtension(ext)) return ReadResult::FILE_NOT_HANDLED;

    if (osgDB::containsServerAddress(fullFileName)) return ReadResult::FILE_NOT_HANDLED;

    osgDB::ReaderWriter* rw = 0;

    if (osgDB::equalCaseInsensitive(ext,"osgz"))
    {
        rw = osgDB::Registry::instance()->getReaderWriterForExtension("osg");
        OSG_INFO<<"osgz ReaderWriter "<<rw<<std::endl;
    }
    else if (osgDB::equalCaseInsensitive(ext,"ivez"))
    {
        rw = osgDB::Registry::instance()->getReaderWriterForExtension("ive");
        OSG_INFO<<"ivez ReaderWriter "<<rw<<std::endl;
    }
    else
    {
        std::string baseFileName = osgDB::getNameLessExtension(fullFileName);
        std::string baseExt = osgDB::getLowerCaseFileExtension(baseFileName);
        rw = osgDB::Registry::instance()->getReaderWriterForExtension(baseExt);
        OSG_INFO<<baseExt<<" ReaderWriter "<<rw<<std::endl;
    }

    std::string fileName = osgDB::findDataFile( fullFileName, options );
    if (fileName.empty()) return ReadResult::FILE_NOT_FOUND;

    // code for setting up the database path so that internally referenced file are searched for on relative paths.
    osg::ref_ptr<Options> local_opt = options ? static_cast<Options*>(options->clone(osg::CopyOp::SHALLOW_COPY)) : new Options;
    local_opt->getDatabasePathList().push_front(osgDB::getFilePath(fileName));

    osgDB::ifstream fin(fileName.c_str(), std::ios::binary|std::ios::in);
    if (!fin) return ReadResult::ERROR_IN_READING_FILE;

    std::string dest;
    read(fin, dest);

    std::stringstream strstream(dest);

    return readFile(objectType, rw, strstream, local_opt.get());
}

osgDB::ReaderWriter::WriteResult ReaderWriterGZ::writeFile(ObjectType objectType, const osg::Object* object, osgDB::ReaderWriter* rw, std::ostream& fout, const osgDB::ReaderWriter::Options *options) const
{
    switch(objectType)
    {
        case(OBJECT): return rw->writeObject(*object, fout, options);
        case(IMAGE): return rw->writeImage(*(dynamic_cast<const osg::Image*>(object)), fout, options);
        case(HEIGHTFIELD): return rw->writeHeightField(*(dynamic_cast<const osg::HeightField*>(object)), fout, options);
        case(NODE): return rw->writeNode(*(dynamic_cast<const osg::Node*>(object)), fout,options);
        default: break;
    }
    return WriteResult::FILE_NOT_HANDLED;
}

osgDB::ReaderWriter::WriteResult ReaderWriterGZ::writeFile(ObjectType objectType, const osg::Object* object, const std::string& fullFileName, const osgDB::ReaderWriter::Options *options) const
{
    std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
    if (!acceptsExtension(ext)) return WriteResult::FILE_NOT_HANDLED;

    osgDB::ReaderWriter* rw = 0;

    if (osgDB::equalCaseInsensitive(ext,"osgz"))
    {
        rw = osgDB::Registry::instance()->getReaderWriterForExtension("osg");
        OSG_NOTICE<<"osgz ReaderWriter "<<rw<<std::endl;
    }
    else if (osgDB::equalCaseInsensitive(ext,"ivez"))
    {
        rw = osgDB::Registry::instance()->getReaderWriterForExtension("ive");
        OSG_NOTICE<<"ivez ReaderWriter "<<rw<<std::endl;
    }
    else
    {
        std::string baseFileName = osgDB::getNameLessExtension(fullFileName);
        std::string baseExt = osgDB::getLowerCaseFileExtension(baseFileName);
        rw = osgDB::Registry::instance()->getReaderWriterForExtension(baseExt);
        OSG_NOTICE<<baseExt<<" ReaderWriter "<<rw<<std::endl;
    }

    std::stringstream strstream;
    osgDB::ReaderWriter::WriteResult writeResult = writeFile(objectType, object, rw, strstream, options);

    osgDB::ofstream fout(fullFileName.c_str(), std::ios::binary|std::ios::out);

    write(fout,strstream.str());

    return writeResult;
}

#define CHUNK 16384

bool ReaderWriterGZ::read(std::istream& fin, std::string& destination) const
{
    int ret;
    unsigned have;
    z_stream strm;
    unsigned char in[CHUNK];
    unsigned char out[CHUNK];

    /* allocate inflate state */
    strm.zalloc = Z_NULL;
    strm.zfree = Z_NULL;
    strm.opaque = Z_NULL;
    strm.avail_in = 0;
    strm.next_in = Z_NULL;
    ret = inflateInit2(&strm,
                       15 + 32 // autodetected zlib or gzip header
                       );
    if (ret != Z_OK)
        return false;

    /* decompress until deflate stream ends or end of file */
    do {

        fin.read((char*)in, CHUNK);
        strm.avail_in = fin.gcount();

        if (fin.bad())
        {
            (void)inflateEnd(&strm);
            return false;
        }
        if (strm.avail_in == 0)
            break;
        strm.next_in = in;

        /* run inflate() on input until output buffer not full */
        do {
            strm.avail_out = CHUNK;
            strm.next_out = out;
            ret = inflate(&strm, Z_NO_FLUSH);

            switch (ret) {
            case Z_NEED_DICT:
            case Z_DATA_ERROR:
            case Z_MEM_ERROR:
                (void)inflateEnd(&strm);
                return false;
            }
            have = CHUNK - strm.avail_out;

            destination.append((char*)out, have);

        } while (strm.avail_out == 0);

        /* done when inflate() says it's done */
    } while (ret != Z_STREAM_END);

    /* clean up and return */
    (void)inflateEnd(&strm);
    return ret == Z_STREAM_END ? true : false;
}

bool ReaderWriterGZ::write(std::ostream& fout, const std::string& source) const
{
    int ret, flush = Z_FINISH;
    unsigned have;
    z_stream strm;
    unsigned char out[CHUNK];

    int level = 6;
    int stategy = Z_DEFAULT_STRATEGY; // looks to be the best for .osg/.ive files
    //int stategy = Z_FILTERED;
    //int stategy = Z_HUFFMAN_ONLY;
    //int stategy = Z_RLE;

    /* allocate deflate state */
    strm.zalloc = Z_NULL;
    strm.zfree = Z_NULL;
    strm.opaque = Z_NULL;
    ret = deflateInit2(&strm,
                       level,
                       Z_DEFLATED,
                       15+16, // +16 to use gzip encoding
                       8, // default
                       stategy);
    if (ret != Z_OK)
    return false;

    strm.avail_in = source.size();
    strm.next_in = (Bytef*)(&(*source.begin()));

    /* run deflate() on input until output buffer not full, finish
       compression if all of source has been read in */
    do {
        strm.avail_out = CHUNK;
        strm.next_out = out;
        ret = deflate(&strm, flush);    /* no bad return value */

        if (ret == Z_STREAM_ERROR)
        {
            OSG_NOTICE<<"Z_STREAM_ERROR"<<std::endl;
            return false;
        }

        have = CHUNK - strm.avail_out;

        if (have>0) fout.write((const char*)out, have);

        if (fout.fail())
        {
            (void)deflateEnd(&strm);
            return false;
        }
    } while (strm.avail_out == 0);

    /* clean up and return */
    (void)deflateEnd(&strm);
    return true;
}

// now register with Registry to instantiate the above
// reader/writer.
REGISTER_OSGPLUGIN(GZ, ReaderWriterGZ)
<commit_msg>add funcs to read Object from a GZ stream<commit_after>/* -*-c++-*- OpenSceneGraph - Copyright (C) 1998-2008 Robert Osfield
 *
 * This library is open source and may be redistributed and/or modified under
 * the terms of the OpenSceneGraph Public License (OSGPL) version 0.0 or
 * (at your option) any later version.  The full license is in LICENSE file
 * included with this distribution, and on the openscenegraph.org website.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * OpenSceneGraph Public License for more details.
*/

#include <osgDB/FileUtils>
#include <osgDB/FileNameUtils>
#include <osgDB/ReadFile>
#include <osgDB/WriteFile>
#include <osgDB/Registry>

#include <iostream>
#include <sstream>
#include <fstream>
#include <iterator>

#include <zlib.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
//
//  ReaderWriterGZ
//
class ReaderWriterGZ : public osgDB::ReaderWriter
{
    public:

        enum ObjectType
        {
            OBJECT,
            ARCHIVE,
            IMAGE,
            HEIGHTFIELD,
            NODE
        };

        ReaderWriterGZ();

        ~ReaderWriterGZ();

        virtual const char* className() const { return "GZ Archive Reader/Writer"; }

        virtual ReadResult openArchive(const std::string& fileName,ArchiveStatus status, unsigned int , const Options* options) const
        {
            if (status!=READ) return ReadResult(ReadResult::FILE_NOT_HANDLED);
            else return readFile(ARCHIVE,fileName,options);
        }

        virtual ReadResult readObject(const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return readFile(OBJECT,fileName,options);
        }

        virtual ReadResult readImage(const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return readFile(IMAGE,fileName,options);
        }

        virtual ReadResult readHeightField(const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return readFile(HEIGHTFIELD,fileName,options);
        }

        virtual ReadResult readNode(const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return readFile(NODE,fileName,options);
        }

        ReadResult readFile(ObjectType objectType, osgDB::ReaderWriter* rw, std::istream& fin, const osgDB::ReaderWriter::Options* options) const;

        ReadResult readFile(ObjectType objectType, const std::string& fullFileName, const osgDB::ReaderWriter::Options* options) const;

        ///return the reader for fullFileName and a outuncompessed stream fetching this file in the archive
        osgDB::ReaderWriter *getStreamAndReader(std::stringstream& outuncompessed, std::istream& fin, const std::string& fullFileName) const {
            std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
            osgDB::ReaderWriter* rw = 0;
            rw = osgDB::Registry::instance()->getReaderWriterForExtension(ext);
            std::string baseFileName = osgDB::getNameLessExtension(fullFileName);
            std::string baseExt = osgDB::getLowerCaseFileExtension(baseFileName);
            rw = osgDB::Registry::instance()->getReaderWriterForExtension(baseExt);
            OSG_INFO<<classname()<<"::getStreamAndReader:"<<baseExt<<" ReaderWriter "<<rw<<std::endl;
            read(fin, outuncompessed);
            return rw;
        }

        virtual ReadResult readImage(std::istream& fin,const osgDB::ReaderWriter::Options* local_opt=NULL) const {
            std::string fullFileName = local_opt->getPluginStringData("STREAM_FILENAME");
            std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
            if (!acceptsExtension(ext)) return ReadResult::FILE_NOT_HANDLED;
            if (osgDB::containsServerAddress(fullFileName)) return ReadResult::FILE_NOT_HANDLED;
            std::stringstream decstream;
            osgDB::ReaderWriter* rw = getStreamAndReader(decstream, fin, fullFileName);
            return readFile(IMAGE, rw, decstream, local_opt);
        }

        virtual ReadResult readHeightField(std::istream& fin,const osgDB::ReaderWriter::Options* local_opt=NULL) const {
            std::string fullFileName = local_opt->getPluginStringData("STREAM_FILENAME");
            std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
            if (!acceptsExtension(ext)) return ReadResult::FILE_NOT_HANDLED;
            if (osgDB::containsServerAddress(fullFileName)) return ReadResult::FILE_NOT_HANDLED;
            std::stringstream decstream;
            osgDB::ReaderWriter* rw = getStreamAndReader(decstream, fin, fullFileName);
            return readFile(HEIGHTFIELD, rw, decstream, local_opt);
        }

        virtual ReadResult readNode(std::istream& fin,const osgDB::ReaderWriter::Options* local_opt=NULL) const {
            std::string fullFileName = local_opt->getPluginStringData("STREAM_FILENAME");
            std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
            if (!acceptsExtension(ext)) return ReadResult::FILE_NOT_HANDLED;
            if (osgDB::containsServerAddress(fullFileName)) return ReadResult::FILE_NOT_HANDLED;
            std::stringstream decstream;
            osgDB::ReaderWriter* rw = getStreamAndReader(decstream, fin, fullFileName);
            return readFile(NODE, rw, decstream, local_opt);
        }

        virtual ReadResult readObject(std::istream& fin,const osgDB::ReaderWriter::Options* local_opt=NULL) const {
            std::string fullFileName = local_opt->getPluginStringData("STREAM_FILENAME");
            std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
            if (!acceptsExtension(ext)) return ReadResult::FILE_NOT_HANDLED;
            if (osgDB::containsServerAddress(fullFileName)) return ReadResult::FILE_NOT_HANDLED;
            std::stringstream decstream;
            osgDB::ReaderWriter* rw = getStreamAndReader(decstream, fin, fullFileName);
            return readFile(OBJECT, rw, decstream, local_opt);
        }
        virtual ReadResult readArchive(std::istream& fin,const osgDB::ReaderWriter::Options* local_opt=NULL) const {
            std::string fullFileName = local_opt->getPluginStringData("STREAM_FILENAME");
            std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
            if (!acceptsExtension(ext)) return ReadResult::FILE_NOT_HANDLED;
            if (osgDB::containsServerAddress(fullFileName)) return ReadResult::FILE_NOT_HANDLED;
            std::stringstream decstream;
            osgDB::ReaderWriter* rw = getStreamAndReader(decstream, fin, fullFileName);
            return readFile(ARCHIVE, rw, decstream, local_opt);
        }

        virtual WriteResult writeObject(const osg::Object& obj, const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return writeFile(OBJECT, &obj, fileName, options);
        }

        virtual WriteResult writeImage(const osg::Image& image, const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return writeFile(IMAGE, &image, fileName, options);
        }

        virtual WriteResult writeHeightField(const osg::HeightField& hf, const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return writeFile(HEIGHTFIELD, &hf, fileName, options);
        }

        virtual WriteResult writeNode(const osg::Node& node, const std::string& fileName, const osgDB::ReaderWriter::Options* options) const
        {
            return writeFile(NODE, &node, fileName,options);
        }

        WriteResult writeFile(ObjectType objectType, const osg::Object* object, osgDB::ReaderWriter* rw, std::ostream& fin, const osgDB::ReaderWriter::Options* options) const;

        WriteResult writeFile(ObjectType objectType, const osg::Object* object, const std::string& fullFileName, const osgDB::ReaderWriter::Options* options) const;

        bool read(std::istream& fin, std::string& destination) const;
        bool write(std::ostream& fout, const std::string& source) const;

};

ReaderWriterGZ::ReaderWriterGZ()
{
    // OSG_NOTICE<<"ReaderWriterGZ::ReaderWriterGZ()"<<std::endl;

    supportsExtension("osgz","Compressed .osg file extension.");
    supportsExtension("ivez","Compressed .ive file extension.");
    supportsExtension("gz","Compressed file extension.");
}

ReaderWriterGZ::~ReaderWriterGZ()
{
    // OSG_NOTICE<<"ReaderWriterGZ::~ReaderWriterGZ()"<<std::endl;
}

osgDB::ReaderWriter::ReadResult ReaderWriterGZ::readFile(ObjectType objectType, osgDB::ReaderWriter* rw, std::istream& fin, const osgDB::ReaderWriter::Options *options) const
{
    switch(objectType)
    {
        case(OBJECT): return rw->readObject(fin,options);
        case(ARCHIVE): return rw->openArchive(fin,options);
        case(IMAGE): return rw->readImage(fin,options);
        case(HEIGHTFIELD): return rw->readHeightField(fin,options);
        case(NODE): return rw->readNode(fin,options);
        default: break;
    }
    return ReadResult::FILE_NOT_HANDLED;
}

osgDB::ReaderWriter::ReadResult ReaderWriterGZ::readFile(ObjectType objectType, const std::string& fullFileName, const osgDB::ReaderWriter::Options *options) const
{
    std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
    if (!acceptsExtension(ext)) return ReadResult::FILE_NOT_HANDLED;

    if (osgDB::containsServerAddress(fullFileName)) return ReadResult::FILE_NOT_HANDLED;

    osgDB::ReaderWriter* rw = 0;

    if (osgDB::equalCaseInsensitive(ext,"osgz"))
    {
        rw = osgDB::Registry::instance()->getReaderWriterForExtension("osg");
        OSG_INFO<<"osgz ReaderWriter "<<rw<<std::endl;
    }
    else if (osgDB::equalCaseInsensitive(ext,"ivez"))
    {
        rw = osgDB::Registry::instance()->getReaderWriterForExtension("ive");
        OSG_INFO<<"ivez ReaderWriter "<<rw<<std::endl;
    }
    else
    {
        std::string baseFileName = osgDB::getNameLessExtension(fullFileName);
        std::string baseExt = osgDB::getLowerCaseFileExtension(baseFileName);
        rw = osgDB::Registry::instance()->getReaderWriterForExtension(baseExt);
        OSG_INFO<<baseExt<<" ReaderWriter "<<rw<<std::endl;
    }

    std::string fileName = osgDB::findDataFile( fullFileName, options );
    if (fileName.empty()) return ReadResult::FILE_NOT_FOUND;

    // code for setting up the database path so that internally referenced file are searched for on relative paths.
    osg::ref_ptr<Options> local_opt = options ? static_cast<Options*>(options->clone(osg::CopyOp::SHALLOW_COPY)) : new Options;
    local_opt->getDatabasePathList().push_front(osgDB::getFilePath(fileName));

    osgDB::ifstream fin(fileName.c_str(), std::ios::binary|std::ios::in);
    if (!fin) return ReadResult::ERROR_IN_READING_FILE;

    std::string dest;
    read(fin, dest);

    std::stringstream strstream(dest);

    return readFile(objectType, rw, strstream, local_opt.get());
}

osgDB::ReaderWriter::WriteResult ReaderWriterGZ::writeFile(ObjectType objectType, const osg::Object* object, osgDB::ReaderWriter* rw, std::ostream& fout, const osgDB::ReaderWriter::Options *options) const
{
    switch(objectType)
    {
        case(OBJECT): return rw->writeObject(*object, fout, options);
        case(IMAGE): return rw->writeImage(*(dynamic_cast<const osg::Image*>(object)), fout, options);
        case(HEIGHTFIELD): return rw->writeHeightField(*(dynamic_cast<const osg::HeightField*>(object)), fout, options);
        case(NODE): return rw->writeNode(*(dynamic_cast<const osg::Node*>(object)), fout,options);
        default: break;
    }
    return WriteResult::FILE_NOT_HANDLED;
}

osgDB::ReaderWriter::WriteResult ReaderWriterGZ::writeFile(ObjectType objectType, const osg::Object* object, const std::string& fullFileName, const osgDB::ReaderWriter::Options *options) const
{
    std::string ext = osgDB::getLowerCaseFileExtension(fullFileName);
    if (!acceptsExtension(ext)) return WriteResult::FILE_NOT_HANDLED;

    osgDB::ReaderWriter* rw = 0;

    if (osgDB::equalCaseInsensitive(ext,"osgz"))
    {
        rw = osgDB::Registry::instance()->getReaderWriterForExtension("osg");
        OSG_NOTICE<<"osgz ReaderWriter "<<rw<<std::endl;
    }
    else if (osgDB::equalCaseInsensitive(ext,"ivez"))
    {
        rw = osgDB::Registry::instance()->getReaderWriterForExtension("ive");
        OSG_NOTICE<<"ivez ReaderWriter "<<rw<<std::endl;
    }
    else
    {
        std::string baseFileName = osgDB::getNameLessExtension(fullFileName);
        std::string baseExt = osgDB::getLowerCaseFileExtension(baseFileName);
        rw = osgDB::Registry::instance()->getReaderWriterForExtension(baseExt);
        OSG_NOTICE<<baseExt<<" ReaderWriter "<<rw<<std::endl;
    }

    std::stringstream strstream;
    osgDB::ReaderWriter::WriteResult writeResult = writeFile(objectType, object, rw, strstream, options);

    osgDB::ofstream fout(fullFileName.c_str(), std::ios::binary|std::ios::out);

    write(fout,strstream.str());

    return writeResult;
}

#define CHUNK 16384

bool ReaderWriterGZ::read(std::istream& fin, std::string& destination) const
{
    int ret;
    unsigned have;
    z_stream strm;
    unsigned char in[CHUNK];
    unsigned char out[CHUNK];

    /* allocate inflate state */
    strm.zalloc = Z_NULL;
    strm.zfree = Z_NULL;
    strm.opaque = Z_NULL;
    strm.avail_in = 0;
    strm.next_in = Z_NULL;
    ret = inflateInit2(&strm,
                       15 + 32 // autodetected zlib or gzip header
                       );
    if (ret != Z_OK)
        return false;

    /* decompress until deflate stream ends or end of file */
    do {

        fin.read((char*)in, CHUNK);
        strm.avail_in = fin.gcount();

        if (fin.bad())
        {
            (void)inflateEnd(&strm);
            return false;
        }
        if (strm.avail_in == 0)
            break;
        strm.next_in = in;

        /* run inflate() on input until output buffer not full */
        do {
            strm.avail_out = CHUNK;
            strm.next_out = out;
            ret = inflate(&strm, Z_NO_FLUSH);

            switch (ret) {
            case Z_NEED_DICT:
            case Z_DATA_ERROR:
            case Z_MEM_ERROR:
                (void)inflateEnd(&strm);
                return false;
            }
            have = CHUNK - strm.avail_out;

            destination.append((char*)out, have);

        } while (strm.avail_out == 0);

        /* done when inflate() says it's done */
    } while (ret != Z_STREAM_END);

    /* clean up and return */
    (void)inflateEnd(&strm);
    return ret == Z_STREAM_END ? true : false;
}

bool ReaderWriterGZ::write(std::ostream& fout, const std::string& source) const
{
    int ret, flush = Z_FINISH;
    unsigned have;
    z_stream strm;
    unsigned char out[CHUNK];

    int level = 6;
    int stategy = Z_DEFAULT_STRATEGY; // looks to be the best for .osg/.ive files
    //int stategy = Z_FILTERED;
    //int stategy = Z_HUFFMAN_ONLY;
    //int stategy = Z_RLE;

    /* allocate deflate state */
    strm.zalloc = Z_NULL;
    strm.zfree = Z_NULL;
    strm.opaque = Z_NULL;
    ret = deflateInit2(&strm,
                       level,
                       Z_DEFLATED,
                       15+16, // +16 to use gzip encoding
                       8, // default
                       stategy);
    if (ret != Z_OK)
    return false;

    strm.avail_in = source.size();
    strm.next_in = (Bytef*)(&(*source.begin()));

    /* run deflate() on input until output buffer not full, finish
       compression if all of source has been read in */
    do {
        strm.avail_out = CHUNK;
        strm.next_out = out;
        ret = deflate(&strm, flush);    /* no bad return value */

        if (ret == Z_STREAM_ERROR)
        {
            OSG_NOTICE<<"Z_STREAM_ERROR"<<std::endl;
            return false;
        }

        have = CHUNK - strm.avail_out;

        if (have>0) fout.write((const char*)out, have);

        if (fout.fail())
        {
            (void)deflateEnd(&strm);
            return false;
        }
    } while (strm.avail_out == 0);

    /* clean up and return */
    (void)deflateEnd(&strm);
    return true;
}

// now register with Registry to instantiate the above
// reader/writer.
REGISTER_OSGPLUGIN(GZ, ReaderWriterGZ)
<|endoftext|>


<commit_before>#include "TrPageArchive.h"
#include "TrPageParser.h"

#include <osg/Group>
#include <osg/Image>
#include <osg/Texture>
#include <osg/Material>
#include <osg/TexEnv>
#include <osg/Transparency>
#include <osg/CullFace>
#include <osg/Light>
#include <osg/StateSet>
#include <osg/Notify>
#include <osgDB/FileUtils>
#include <iostream>

#include <osgDB/ReadFile>
#include <osgDB/FileNameUtils>

#include "trpage_geom.h"
#include "trpage_read.h"
#include "trpage_write.h"
#include "trpage_scene.h"

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

using namespace txp;
using namespace osg;

TrPageArchive::TrPageArchive()
: trpgr_Archive()
, parse(new TrPageParser())
{
}

TrPageArchive::~TrPageArchive()
{
}

bool TrPageArchive::OpenFile(const char* file)
{
    m_alternate_path = osgDB::getFilePath(file);
    std::string name = osgDB::getSimpleFileName(file);

    if(m_alternate_path.empty())
        SetDirectory(".");
    else
    {
        osgDB::setFilePath(m_alternate_path.c_str());
        SetDirectory(m_alternate_path.c_str());
    }

    if (!trpgr_Archive::OpenFile(name.c_str()))
    {
        notify(WARN) << "TrPageArchive::OpenFile() error: "
            << "couldn't open archive: " << file << std::endl;
        return false;
    }
    if (!ReadHeader())
    {
        notify(WARN) << "TrPageArchive::OpenFile() error: "
            << "couldn't read header for archive: " << file
            << std::endl;
        return false;
    }

    // Set the max group ID here because it won't change
    const trpgHeader *head = GetHeader();
    int maxID;
    head->GetMaxGroupID(maxID);
    parse->SetMaxGroupID(maxID);

    return true;
}

// load textures and materials
// TODO : multitexturing
void TrPageArchive::LoadMaterials()
{
    int n_textures;
    texTable.GetNumTextures(n_textures);

    m_textures.resize(n_textures);

    // these extra braces are workaroud for annoying bug in MSVC
    // for( int i = ....) and i is visible outside the loop
    {
        for (int i=0; i < n_textures ; i++)
        {
            const trpgTexture *tex;
            tex = texTable.GetTextureRef(i);
            char texName[1024];  texName[0] = 0;
            tex->GetName(texName,1023);
            // Create a texture by name.
            Texture *osg_texture = new Texture();

            // Load Texture and Create Texture State
            std::string filename = osgDB::getSimpleFileName(texName);
            std::string path(getDir());
            std::string theFile = path + '/'+ filename ;
            ref_ptr<Image> image = osgDB::readImageFile(theFile);
            if (image.valid())
            {
                osg_texture->setImage(image.get());
            }
            m_textures[i] = osg_texture;
        }
    }

    // create materials
    int n_materials;
    materialTable.GetNumMaterial(n_materials);

    {
        for (int i = 0; i < n_materials; i++)
        {
            StateSet* osg_state_set = new StateSet;

            const trpgMaterial *mat;
            mat = materialTable.GetMaterialRef(0,i);
            // Set texture
            int numMatTex;
            mat->GetNumTexture(numMatTex);

            // TODO : multitextuting
            // also note that multitexturing in terrapage can came from two sides
            // - multiple textures per material, and multiple materials per geometry
            if(numMatTex)
            {
                int texId;
                trpgTextureEnv texEnv;
                mat->GetTexture(0,texId,texEnv);

                // Set up texture environment
                TexEnv       *osg_texenv       = new TexEnv();
                int32 te_mode;
                texEnv.GetEnvMode(te_mode);
                switch( te_mode )
                {
                case trpgTextureEnv::Alpha :
                    osg_texenv->setMode(TexEnv::REPLACE);
                    break;
                case trpgTextureEnv::Decal:
                    osg_texenv->setMode(TexEnv::DECAL);
                    break;
                case trpgTextureEnv::Blend :
                    osg_texenv->setMode(TexEnv::BLEND);
                    break;
                case trpgTextureEnv::Modulate :
                    osg_texenv->setMode(TexEnv::MODULATE);
                    break;
                }

                osg_state_set->setAttribute(osg_texenv);

                Material     *osg_material     = new Material;

                float64 alpha;
                mat->GetAlpha(alpha);

                trpgColor color;
                mat->GetAmbient(color);
                osg_material->setAmbient( Material::FRONT_AND_BACK ,
                    Vec4(color.red, color.green, color.blue, alpha));
                mat->GetDiffuse(color);
                osg_material->setDiffuse(Material::FRONT_AND_BACK ,
                    Vec4(color.red, color.green, color.blue, alpha));

                mat->GetSpecular(color);
                osg_material->setSpecular(Material::FRONT_AND_BACK ,
                    Vec4(color.red, color.green, color.blue, alpha));
                mat->GetEmission(color);
                osg_material->setEmission(Material::FRONT_AND_BACK ,
                    Vec4(color.red, color.green, color.blue, alpha));

                float64 shinines;
                mat->GetShininess(shinines);
                osg_material->setShininess(Material::FRONT_AND_BACK , (float)shinines/128.0);

                osg_material->setAlpha(Material::FRONT_AND_BACK ,(float)alpha);
                osg_state_set->setAttributeAndModes(osg_material, StateAttribute::ON);

                if( alpha < 1.0f )
                {
                    osg_state_set->setMode(GL_BLEND,StateAttribute::ON);
                    osg_state_set->setRenderingHint(StateSet::TRANSPARENT_BIN);
                }

                /*
                int alphaFunc;
                mat->GetAlphaFunc(alphaFunc);
                */

                // transparency
                Transparency *osg_transparency = new Transparency;
                osg_transparency->setFunction(Transparency::SRC_ALPHA, Transparency::ONE_MINUS_SRC_ALPHA);
                osg_state_set->setAttributeAndModes(osg_transparency, StateAttribute::ON);

                int wrap_s, wrap_t;
                texEnv.GetWrap(wrap_s, wrap_t);

                Texture* osg_texture = m_textures[texId].get();
                osg_texture->setWrap(Texture::WRAP_S, wrap_s == trpgTextureEnv::Repeat ? Texture::REPEAT: Texture::CLAMP );
                osg_texture->setWrap(Texture::WRAP_T, wrap_t == trpgTextureEnv::Repeat ? Texture::REPEAT: Texture::CLAMP );
                osg_state_set->setAttributeAndModes(osg_texture, StateAttribute::ON);

                if(osg_texture->getImage())
                {
                    switch (osg_texture->getImage()->pixelFormat())
                    {
                    case GL_LUMINANCE_ALPHA:
                    case GL_RGBA:
                        osg_state_set->setMode(GL_BLEND,StateAttribute::ON);
                        osg_state_set->setRenderingHint(StateSet::TRANSPARENT_BIN);
                    }
                }

                int cullMode;
                mat->GetCullMode(cullMode);

                // Culling mode in txp means opposite from osg i.e. Front-> show front face
                if( cullMode != trpgMaterial::FrontAndBack)
                {
                    CullFace* cull_face = new CullFace;
                    switch (cullMode)
                    {
                    case trpgMaterial::Front:
                        cull_face->setMode(CullFace::BACK);
                        break;
                    case trpgMaterial::Back:
                        cull_face->setMode(CullFace::FRONT);
                        break;
                    }
                    osg_state_set->setAttributeAndModes(cull_face, StateAttribute::ON);
                }
          }
          m_gstates.push_back(osg_state_set);
        }
    }
}

bool TrPageArchive::LoadModels()
{
   int numModel;
   modelTable.GetNumModels(numModel);

   // Iterate over the models
   for (int i=0; i< numModel; i++)
   {
      trpgModel *mod = modelTable.GetModelRef(i);
      int type;
      mod->GetType(type);

      // Only dealing with external models currently
      if (type == trpgModel::External)
      {
         char name[1024];
         mod->GetName(name,1023);

         // Load the model.  It's probably not TerraPage
         Node *osg_model = osgDB::readNodeFile(name);
         if (!osg_model)
         {
            notify(WARN) << "TrPageArchive::LoadModels() error: "
                              << "failed to load model: "
                              << name << std::endl;
         }
         // Do this even if it's NULL
         m_models.push_back(osg_model);
      }
/*
      else
      {
          trpgMemReadBuffer buf(GetEndian());
          mod->Read(buf);
          Group *osg_model = parse->ParseScene(buf, m_gstates , m_models);
          m_models.push_back(osg_model);
      }
*/
  }
   return true;
}

Group* TrPageArchive::LoadTile(int x,int y,int lod,int &parentID)
{
   trpgMemReadBuffer buf(GetEndian());

   // Read the tile data in, but don't parse it
   if (!ReadTile(x,y,lod,buf))
      return NULL;

   // Call the parser
   Group *tile = parse->ParseScene(buf, m_gstates , m_models);
   if (tile)
   {
      parentID = parse->GetParentID();
      // This is where you would page in textures and models
   }

   // That's it
   return tile;
}

//----------------------------------------------------------------------------
Group* TrPageArchive::LoadAllTiles()
{
   // Size information comes out of the header
   const trpgHeader *head = GetHeader();
   // Create one group as the top
   Group *topGroup = new Group;

   int32 numLod;
   head->GetNumLods(numLod);
   //  Iterate over the LODs.  Lower res LODs must be loaded
   //  first, otherwise there's nothing to hook the higher res
   //  LODs into.
   trpg2iPoint tileSize;

   // The group list is used to map parent IDs to pfGroup nodes
   std::vector<Group *> *groupList = parse->GetGroupList();

   for (int nl=0;nl<numLod;nl++)
   {
      // The number of tiles in X and Y changes per LOD
      head->GetLodSize(nl,tileSize);
      for (int x=0; x < tileSize.x; x++)
      {
         for (int y=0; y < tileSize.y; y++)
         {
            int parentID;
            Group *tile = LoadTile(x,y,nl,parentID);
            if (!tile)
            {
                notify(WARN)<< "TrPageArchive::LoadAllTiles error: "
                                 << "failed to load tile ("
                                 << x << "," << y << "," << nl << ")"
                                 << std::endl;
            }
            else
            {
               if (parentID == -1)
               {
                  // Get added to the top level node
                  topGroup->addChild(tile);
               }
               else
               {
                  // Added below some other node
                  (*groupList)[parentID]->addChild(tile);
               }
            }
         }
      }
   }
   return topGroup;
}
<commit_msg>Fix to TXP plugin to handle to new osg::Image method naming.<commit_after>#include "TrPageArchive.h"
#include "TrPageParser.h"

#include <osg/Group>
#include <osg/Image>
#include <osg/Texture>
#include <osg/Material>
#include <osg/TexEnv>
#include <osg/Transparency>
#include <osg/CullFace>
#include <osg/Light>
#include <osg/StateSet>
#include <osg/Notify>
#include <osgDB/FileUtils>
#include <iostream>

#include <osgDB/ReadFile>
#include <osgDB/FileNameUtils>

#include "trpage_geom.h"
#include "trpage_read.h"
#include "trpage_write.h"
#include "trpage_scene.h"

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

using namespace txp;
using namespace osg;

TrPageArchive::TrPageArchive()
: trpgr_Archive()
, parse(new TrPageParser())
{
}

TrPageArchive::~TrPageArchive()
{
}

bool TrPageArchive::OpenFile(const char* file)
{
    m_alternate_path = osgDB::getFilePath(file);
    std::string name = osgDB::getSimpleFileName(file);

    if(m_alternate_path.empty())
        SetDirectory(".");
    else
    {
        osgDB::setFilePath(m_alternate_path.c_str());
        SetDirectory(m_alternate_path.c_str());
    }

    if (!trpgr_Archive::OpenFile(name.c_str()))
    {
        notify(WARN) << "TrPageArchive::OpenFile() error: "
            << "couldn't open archive: " << file << std::endl;
        return false;
    }
    if (!ReadHeader())
    {
        notify(WARN) << "TrPageArchive::OpenFile() error: "
            << "couldn't read header for archive: " << file
            << std::endl;
        return false;
    }

    // Set the max group ID here because it won't change
    const trpgHeader *head = GetHeader();
    int maxID;
    head->GetMaxGroupID(maxID);
    parse->SetMaxGroupID(maxID);

    return true;
}

// load textures and materials
// TODO : multitexturing
void TrPageArchive::LoadMaterials()
{
    int n_textures;
    texTable.GetNumTextures(n_textures);

    m_textures.resize(n_textures);

    // these extra braces are workaroud for annoying bug in MSVC
    // for( int i = ....) and i is visible outside the loop
    {
        for (int i=0; i < n_textures ; i++)
        {
            const trpgTexture *tex;
            tex = texTable.GetTextureRef(i);
            char texName[1024];  texName[0] = 0;
            tex->GetName(texName,1023);
            // Create a texture by name.
            Texture *osg_texture = new Texture();

            // Load Texture and Create Texture State
            std::string filename = osgDB::getSimpleFileName(texName);
            std::string path(getDir());
            std::string theFile = path + '/'+ filename ;
            ref_ptr<Image> image = osgDB::readImageFile(theFile);
            if (image.valid())
            {
                osg_texture->setImage(image.get());
            }
            m_textures[i] = osg_texture;
        }
    }

    // create materials
    int n_materials;
    materialTable.GetNumMaterial(n_materials);

    {
        for (int i = 0; i < n_materials; i++)
        {
            StateSet* osg_state_set = new StateSet;

            const trpgMaterial *mat;
            mat = materialTable.GetMaterialRef(0,i);
            // Set texture
            int numMatTex;
            mat->GetNumTexture(numMatTex);

            // TODO : multitextuting
            // also note that multitexturing in terrapage can came from two sides
            // - multiple textures per material, and multiple materials per geometry
            if(numMatTex)
            {
                int texId;
                trpgTextureEnv texEnv;
                mat->GetTexture(0,texId,texEnv);

                // Set up texture environment
                TexEnv       *osg_texenv       = new TexEnv();
                int32 te_mode;
                texEnv.GetEnvMode(te_mode);
                switch( te_mode )
                {
                case trpgTextureEnv::Alpha :
                    osg_texenv->setMode(TexEnv::REPLACE);
                    break;
                case trpgTextureEnv::Decal:
                    osg_texenv->setMode(TexEnv::DECAL);
                    break;
                case trpgTextureEnv::Blend :
                    osg_texenv->setMode(TexEnv::BLEND);
                    break;
                case trpgTextureEnv::Modulate :
                    osg_texenv->setMode(TexEnv::MODULATE);
                    break;
                }

                osg_state_set->setAttribute(osg_texenv);

                Material     *osg_material     = new Material;

                float64 alpha;
                mat->GetAlpha(alpha);

                trpgColor color;
                mat->GetAmbient(color);
                osg_material->setAmbient( Material::FRONT_AND_BACK ,
                    Vec4(color.red, color.green, color.blue, alpha));
                mat->GetDiffuse(color);
                osg_material->setDiffuse(Material::FRONT_AND_BACK ,
                    Vec4(color.red, color.green, color.blue, alpha));

                mat->GetSpecular(color);
                osg_material->setSpecular(Material::FRONT_AND_BACK ,
                    Vec4(color.red, color.green, color.blue, alpha));
                mat->GetEmission(color);
                osg_material->setEmission(Material::FRONT_AND_BACK ,
                    Vec4(color.red, color.green, color.blue, alpha));

                float64 shinines;
                mat->GetShininess(shinines);
                osg_material->setShininess(Material::FRONT_AND_BACK , (float)shinines/128.0);

                osg_material->setAlpha(Material::FRONT_AND_BACK ,(float)alpha);
                osg_state_set->setAttributeAndModes(osg_material, StateAttribute::ON);

                if( alpha < 1.0f )
                {
                    osg_state_set->setMode(GL_BLEND,StateAttribute::ON);
                    osg_state_set->setRenderingHint(StateSet::TRANSPARENT_BIN);
                }

                /*
                int alphaFunc;
                mat->GetAlphaFunc(alphaFunc);
                */

                // transparency
                Transparency *osg_transparency = new Transparency;
                osg_transparency->setFunction(Transparency::SRC_ALPHA, Transparency::ONE_MINUS_SRC_ALPHA);
                osg_state_set->setAttributeAndModes(osg_transparency, StateAttribute::ON);

                int wrap_s, wrap_t;
                texEnv.GetWrap(wrap_s, wrap_t);

                Texture* osg_texture = m_textures[texId].get();
                osg_texture->setWrap(Texture::WRAP_S, wrap_s == trpgTextureEnv::Repeat ? Texture::REPEAT: Texture::CLAMP );
                osg_texture->setWrap(Texture::WRAP_T, wrap_t == trpgTextureEnv::Repeat ? Texture::REPEAT: Texture::CLAMP );
                osg_state_set->setAttributeAndModes(osg_texture, StateAttribute::ON);

                if(osg_texture->getImage())
                {
                    switch (osg_texture->getImage()->getPixelFormat())
                    {
                    case GL_LUMINANCE_ALPHA:
                    case GL_RGBA:
                        osg_state_set->setMode(GL_BLEND,StateAttribute::ON);
                        osg_state_set->setRenderingHint(StateSet::TRANSPARENT_BIN);
                    }
                }

                int cullMode;
                mat->GetCullMode(cullMode);

                // Culling mode in txp means opposite from osg i.e. Front-> show front face
                if( cullMode != trpgMaterial::FrontAndBack)
                {
                    CullFace* cull_face = new CullFace;
                    switch (cullMode)
                    {
                    case trpgMaterial::Front:
                        cull_face->setMode(CullFace::BACK);
                        break;
                    case trpgMaterial::Back:
                        cull_face->setMode(CullFace::FRONT);
                        break;
                    }
                    osg_state_set->setAttributeAndModes(cull_face, StateAttribute::ON);
                }
          }
          m_gstates.push_back(osg_state_set);
        }
    }
}

bool TrPageArchive::LoadModels()
{
   int numModel;
   modelTable.GetNumModels(numModel);

   // Iterate over the models
   for (int i=0; i< numModel; i++)
   {
      trpgModel *mod = modelTable.GetModelRef(i);
      int type;
      mod->GetType(type);

      // Only dealing with external models currently
      if (type == trpgModel::External)
      {
         char name[1024];
         mod->GetName(name,1023);

         // Load the model.  It's probably not TerraPage
         Node *osg_model = osgDB::readNodeFile(name);
         if (!osg_model)
         {
            notify(WARN) << "TrPageArchive::LoadModels() error: "
                              << "failed to load model: "
                              << name << std::endl;
         }
         // Do this even if it's NULL
         m_models.push_back(osg_model);
      }
/*
      else
      {
          trpgMemReadBuffer buf(GetEndian());
          mod->Read(buf);
          Group *osg_model = parse->ParseScene(buf, m_gstates , m_models);
          m_models.push_back(osg_model);
      }
*/
  }
   return true;
}

Group* TrPageArchive::LoadTile(int x,int y,int lod,int &parentID)
{
   trpgMemReadBuffer buf(GetEndian());

   // Read the tile data in, but don't parse it
   if (!ReadTile(x,y,lod,buf))
      return NULL;

   // Call the parser
   Group *tile = parse->ParseScene(buf, m_gstates , m_models);
   if (tile)
   {
      parentID = parse->GetParentID();
      // This is where you would page in textures and models
   }

   // That's it
   return tile;
}

//----------------------------------------------------------------------------
Group* TrPageArchive::LoadAllTiles()
{
   // Size information comes out of the header
   const trpgHeader *head = GetHeader();
   // Create one group as the top
   Group *topGroup = new Group;

   int32 numLod;
   head->GetNumLods(numLod);
   //  Iterate over the LODs.  Lower res LODs must be loaded
   //  first, otherwise there's nothing to hook the higher res
   //  LODs into.
   trpg2iPoint tileSize;

   // The group list is used to map parent IDs to pfGroup nodes
   std::vector<Group *> *groupList = parse->GetGroupList();

   for (int nl=0;nl<numLod;nl++)
   {
      // The number of tiles in X and Y changes per LOD
      head->GetLodSize(nl,tileSize);
      for (int x=0; x < tileSize.x; x++)
      {
         for (int y=0; y < tileSize.y; y++)
         {
            int parentID;
            Group *tile = LoadTile(x,y,nl,parentID);
            if (!tile)
            {
                notify(WARN)<< "TrPageArchive::LoadAllTiles error: "
                                 << "failed to load tile ("
                                 << x << "," << y << "," << nl << ")"
                                 << std::endl;
            }
            else
            {
               if (parentID == -1)
               {
                  // Get added to the top level node
                  topGroup->addChild(tile);
               }
               else
               {
                  // Added below some other node
                  (*groupList)[parentID]->addChild(tile);
               }
            }
         }
      }
   }
   return topGroup;
}
<|endoftext|>


<commit_before>/*****************************************************************************
 * Licensed to Qualys, Inc. (QUALYS) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * QUALYS licenses this file to You 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.
 ****************************************************************************/

/**
 * @file
 * @brief IronBee --- Trusted Proxy Module
 */

// Include all of IronBee++
#include <ironbeepp/all.hpp>
#include <ironbee/ipset.h>
#include <ironbee/string.h>

#include <boost/algorithm/string/split.hpp>
#include <boost/algorithm/string/trim.hpp>
#include <boost/bind.hpp>

#include <vector>

using namespace std;

namespace {

/**
 * Configuration for the Trusted Proxy Module
 */
class TrustedProxyConfig
{
public:
    //! Constructor.
    TrustedProxyConfig();

    //! Is X-Forwarded-For handling enabled?
    bool is_xff_enabled() const;

    //! Enable or disable X-Forwarded-For handling.
    void set_xff_enabled(bool enabled);

    //! Clear the trusted and untrusted network lists.
    void clear_networks();

    //! Add a network to the trusted list.
    void add_trusted_network(const char* network);

    //! Add a network to the untrusted list.
    void add_untrusted_network(const char* network);

    /**
     * Check if the ip address is trusted.
     *
     * The default is to trust all addresses.
     *
     * @param[in] ipstr IP address to check.
     *
     * @returns True if the address is trusted.
     */
    bool is_trusted(const char* ipstr) const;

    //! Finalize the configuration when the context is closed.
    void context_close(IronBee::Engine& ib);

private:
    //! X-Forwarding-For handling enabled?
    bool m_xff_enabled;

    //! List of trusted networks
    vector<ib_ipset4_entry_t> m_trusted_net_list;

    //! List of untrusted networks
    vector<ib_ipset4_entry_t> m_untrusted_net_list;

    //! IP set of the trusted and untrusted networks.
    ib_ipset4_t m_trusted_networks;
};

/**
 * Module to handle X-Forwarded-For headers from trusted Proxies.
 */
class TrustedProxyModule : public IronBee::ModuleDelegate
{
public:
    //! Constructor.
    explicit TrustedProxyModule(IronBee::Module module);

private:
    //! X-Forward-For lookup target.
    IronBee::ConstVarTarget m_xff_target;

    //! Source for recording the remote address
    IronBee::VarSource m_remote_addr_source;

    /**
     * Handle the TrustedProxyUseXFFHeader directive.
     *
     * @param[in] cp Configuration parser.
     * @param[in] enabled Should XFF handling be enabled?
     */
    void enable_xff_directive(
        IronBee::ConfigurationParser cp,
        bool                         enabled);

    /**
     * @param[in] cp Configuration parser.
     * @param[in] ip_list List of IPs or CIDR blocks.
     */
    void trusted_ips_directive(
        IronBee::ConfigurationParser    cp,
        IronBee::ConstList<const char*> ip_list);

    void on_context_close(IronBee::Engine ib, IronBee::Context ctx);

    /**
     * Update the transaction effective IP based.
     *
     * @param[in] ib Ironbee engine.
     * @param[in,out] tx Transaction to update.
     */
    void set_effective_ip(
        IronBee::Engine      ib,
        IronBee::Transaction tx);
};

} // Anonymous namespace

IBPP_BOOTSTRAP_MODULE_DELEGATE("TrustedProxyModule", TrustedProxyModule)

namespace {

void make_ipset_entry(const char* cidr_or_ip, ib_ipset4_entry_t& entry)
{
    if (strchr(cidr_or_ip, '/') != NULL) {
        // Has / assume CIDR
        IronBee::throw_if_error(
            ib_ip4_str_to_net(cidr_or_ip, &(entry.network)),
            "Invalid CIDR block");
    }
    else {
        // IP make /32
        IronBee::throw_if_error(
            ib_ip4_str_to_ip(cidr_or_ip, &(entry.network.ip)),
            "Invalid IP address");
        entry.network.size=32;
    }
}

TrustedProxyConfig::TrustedProxyConfig()
    : m_xff_enabled(true)
{
    ib_ipset4_init(&m_trusted_networks, NULL, 0, NULL, 0);
}

bool TrustedProxyConfig::is_xff_enabled() const
{
    return m_xff_enabled;
}

void TrustedProxyConfig::set_xff_enabled(bool enabled)
{
    m_xff_enabled = enabled;
}

void TrustedProxyConfig::clear_networks()
{
    m_trusted_net_list.clear();
    m_untrusted_net_list.clear();
}

void TrustedProxyConfig::add_trusted_network(const char* cidr_or_ip)
{
    ib_ipset4_entry_t net_entry;
    make_ipset_entry(cidr_or_ip, net_entry);
    m_trusted_net_list.push_back(net_entry);
}

void TrustedProxyConfig::add_untrusted_network(const char* cidr_or_ip)
{
    ib_ipset4_entry_t net_entry;
    make_ipset_entry(cidr_or_ip, net_entry);
    m_untrusted_net_list.push_back(net_entry);
}

void TrustedProxyConfig::context_close(IronBee::Engine& ib)
{
    IronBee::throw_if_error(
        ib_ipset4_init(
            &m_trusted_networks,
            m_untrusted_net_list.data(),
            m_untrusted_net_list.size(),
            m_trusted_net_list.data(),
            m_trusted_net_list.size()),
        "Failed to initialize IPv4 set.");
}

bool TrustedProxyConfig::is_trusted(const char* ipstr) const
{
    ib_ip4_t ip;

    IronBee::throw_if_error(
        ib_ip4_str_to_ip(ipstr, &ip),
        "Invalid remote IP address");
    ib_status_t rc =
        ib_ipset4_query(&m_trusted_networks, ip, NULL, NULL, NULL);
    return rc == IB_OK;
}

TrustedProxyModule::TrustedProxyModule(IronBee::Module module) :
    IronBee::ModuleDelegate(module)
{
    module.set_configuration_data<TrustedProxyConfig>();

    m_xff_target = IronBee::VarTarget::acquire_from_string(
        module.engine().main_memory_pool(),
        module.engine().var_config(),
        "request_headers:X-Forwarded-For");

    m_remote_addr_source = IronBee::VarSource::register_(
        module.engine().var_config(),
        IB_S2SL("remote_addr"),
        IB_PHASE_REQUEST_HEADER,
        IB_PHASE_REQUEST_HEADER);

    module.engine().register_configuration_directives()
        .on_off(
            "TrustedProxyUseXFFHeader",
            boost::bind(&TrustedProxyModule::enable_xff_directive,
                        this, _1, _3))
        .list(
            "TrustedProxyIPs",
            boost::bind(&TrustedProxyModule::trusted_ips_directive,
                        this, _1, _3));

    module.engine().register_hooks()
        .context_close(
            boost::bind(&TrustedProxyModule::on_context_close,
                        this, _1, _2))
        .request_header_finished(
            boost::bind(&TrustedProxyModule::set_effective_ip,
                        this, _1, _2));
}

void TrustedProxyModule::enable_xff_directive(
    IronBee::ConfigurationParser cp,
    bool                         enabled
)
{
    TrustedProxyConfig& config =
        module().configuration_data<TrustedProxyConfig>(cp.current_context());
    config.set_xff_enabled(enabled);
}

void TrustedProxyModule::trusted_ips_directive(
    IronBee::ConfigurationParser    cp,
    IronBee::ConstList<const char*> ip_list
)
{
    TrustedProxyConfig& config =
        module().configuration_data<TrustedProxyConfig>(cp.current_context());

    const char* first_arg = ip_list.front();
    if (*first_arg != '+' and *first_arg != '-') {
        config.clear_networks();
    }

    BOOST_FOREACH(const char* arg, ip_list) {
        if (*arg == '+') {
            config.add_trusted_network(arg+1);
        }
        else if (*arg == '-') {
            config.add_untrusted_network(arg+1);
        }
        else {
            config.add_trusted_network(arg);
        }
    }
}

void TrustedProxyModule::on_context_close(
    IronBee::Engine ib,
    IronBee::Context ctx
)
{
    TrustedProxyConfig& config =
        module().configuration_data<TrustedProxyConfig>(ctx);

    config.context_close(ib);
}

void TrustedProxyModule::set_effective_ip(
    IronBee::Engine      ib,
    IronBee::Transaction tx
)
{
    ib_status_t rc;
    IronBee::Context ctx = tx.context();
    TrustedProxyConfig& config =
        module().configuration_data<TrustedProxyConfig>(ctx);

    ib_log_debug_tx(tx.ib(), "checking: %s",
                   tx.connection().remote_ip_string());
    // check actual remote ip against trusted ips
    if (! config.is_trusted(tx.connection().remote_ip_string())) {
        ib_log_debug_tx(tx.ib(), "Remote address '%s' not a trusted proxy.",
                       tx.connection().remote_ip_string());
        return;
    }
    // Remote address is trusted, get the X-Forwarded-For value
    IronBee::ConstList<IronBee::ConstField> xfflist =
        m_xff_target.get(tx.memory_pool(), tx.var_store());

    if (xfflist.size() == 0) {
        // No X-Forwarded-For header found. Nothing to do.
        return;
    }
    // Use the final IP address in X-Forwarded-For
    string forwarded = xfflist.back().to_s();
    list<string> forwarded_list;
    boost::algorithm::split(forwarded_list,
                            forwarded,
                            boost::algorithm::is_any_of(","));

    string remote_ip = boost::algorithm::trim_copy(forwarded_list.back());

    /* Verify that it looks like a valid IP address, ignore it if not */
    rc = ib_ip_validate(remote_ip.c_str());
    if (rc != IB_OK) {
        ib_log_error_tx(tx.ib(),
                        "X-Forwarded-For \"%s\" is not a valid IP address",
                        remote_ip.c_str());
        return;
    }
    char* buf = static_cast<char*>(tx.memory_pool().alloc(remote_ip.length()+1));
    strcpy(buf, remote_ip.c_str());

    /* This will lose the pointer to the original address
     * buffer, but it should be cleaned up with the rest
     * of the memory pool. */
    tx.ib()->remote_ipstr = buf;

    ib_log_debug_tx(tx.ib(), "Remote address changed to \"%s\"", buf);
    IronBee::ByteString remote_addr_bs = IronBee::ByteString::create_alias(
        tx.memory_pool(),
        buf,
        remote_ip.length());
    IronBee::Field remote_addr_field = IronBee::Field::create_byte_string(
        tx.memory_pool(),
        "", 0,
        remote_addr_bs);
    m_remote_addr_source.set(tx.var_store(), remote_addr_field);
}

} // Anonymous namespace
<commit_msg>modules/trusted_proxy.cpp: Indentation cleanup.<commit_after>/*****************************************************************************
 * Licensed to Qualys, Inc. (QUALYS) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * QUALYS licenses this file to You 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.
 ****************************************************************************/

/**
 * @file
 * @brief IronBee --- Trusted Proxy Module
 */

// Include all of IronBee++
#include <ironbeepp/all.hpp>
#include <ironbee/ipset.h>
#include <ironbee/string.h>

#include <boost/algorithm/string/split.hpp>
#include <boost/algorithm/string/trim.hpp>
#include <boost/bind.hpp>

#include <vector>

using namespace std;

namespace {

/**
 * Configuration for the Trusted Proxy Module
 */
class TrustedProxyConfig
{
public:
    //! Constructor.
    TrustedProxyConfig();

    //! Is X-Forwarded-For handling enabled?
    bool is_xff_enabled() const;

    //! Enable or disable X-Forwarded-For handling.
    void set_xff_enabled(bool enabled);

    //! Clear the trusted and untrusted network lists.
    void clear_networks();

    //! Add a network to the trusted list.
    void add_trusted_network(const char* network);

    //! Add a network to the untrusted list.
    void add_untrusted_network(const char* network);

    /**
     * Check if the ip address is trusted.
     *
     * The default is to trust all addresses.
     *
     * @param[in] ipstr IP address to check.
     *
     * @returns True if the address is trusted.
     */
    bool is_trusted(const char* ipstr) const;

    //! Finalize the configuration when the context is closed.
    void context_close(IronBee::Engine& ib);

private:
    //! X-Forwarding-For handling enabled?
    bool m_xff_enabled;

    //! List of trusted networks
    vector<ib_ipset4_entry_t> m_trusted_net_list;

    //! List of untrusted networks
    vector<ib_ipset4_entry_t> m_untrusted_net_list;

    //! IP set of the trusted and untrusted networks.
    ib_ipset4_t m_trusted_networks;
};

/**
 * Module to handle X-Forwarded-For headers from trusted Proxies.
 */
class TrustedProxyModule : public IronBee::ModuleDelegate
{
public:
    //! Constructor.
    explicit TrustedProxyModule(IronBee::Module module);

private:
    //! X-Forward-For lookup target.
    IronBee::ConstVarTarget m_xff_target;

    //! Source for recording the remote address
    IronBee::VarSource m_remote_addr_source;

    /**
     * Handle the TrustedProxyUseXFFHeader directive.
     *
     * @param[in] cp Configuration parser.
     * @param[in] enabled Should XFF handling be enabled?
     */
    void enable_xff_directive(
        IronBee::ConfigurationParser cp,
        bool                         enabled);

    /**
     * @param[in] cp Configuration parser.
     * @param[in] ip_list List of IPs or CIDR blocks.
     */
    void trusted_ips_directive(
        IronBee::ConfigurationParser    cp,
        IronBee::ConstList<const char*> ip_list);

    void on_context_close(IronBee::Engine ib, IronBee::Context ctx);

    /**
     * Update the transaction effective IP based.
     *
     * @param[in] ib Ironbee engine.
     * @param[in,out] tx Transaction to update.
     */
    void set_effective_ip(
        IronBee::Engine      ib,
        IronBee::Transaction tx);
};

} // Anonymous namespace

IBPP_BOOTSTRAP_MODULE_DELEGATE("TrustedProxyModule", TrustedProxyModule)

namespace {

void make_ipset_entry(const char* cidr_or_ip, ib_ipset4_entry_t& entry)
{
    if (strchr(cidr_or_ip, '/') != NULL) {
        // Has / assume CIDR
        IronBee::throw_if_error(
            ib_ip4_str_to_net(cidr_or_ip, &(entry.network)),
            "Invalid CIDR block");
    }
    else {
        // IP make /32
        IronBee::throw_if_error(
            ib_ip4_str_to_ip(cidr_or_ip, &(entry.network.ip)),
            "Invalid IP address");
        entry.network.size=32;
    }
}

TrustedProxyConfig::TrustedProxyConfig()
    : m_xff_enabled(true)
{
    ib_ipset4_init(&m_trusted_networks, NULL, 0, NULL, 0);
}

bool TrustedProxyConfig::is_xff_enabled() const
{
    return m_xff_enabled;
}

void TrustedProxyConfig::set_xff_enabled(bool enabled)
{
    m_xff_enabled = enabled;
}

void TrustedProxyConfig::clear_networks()
{
    m_trusted_net_list.clear();
    m_untrusted_net_list.clear();
}

void TrustedProxyConfig::add_trusted_network(const char* cidr_or_ip)
{
    ib_ipset4_entry_t net_entry;
    make_ipset_entry(cidr_or_ip, net_entry);
    m_trusted_net_list.push_back(net_entry);
}

void TrustedProxyConfig::add_untrusted_network(const char* cidr_or_ip)
{
    ib_ipset4_entry_t net_entry;
    make_ipset_entry(cidr_or_ip, net_entry);
    m_untrusted_net_list.push_back(net_entry);
}

void TrustedProxyConfig::context_close(IronBee::Engine& ib)
{
    IronBee::throw_if_error(
        ib_ipset4_init(
            &m_trusted_networks,
            m_untrusted_net_list.data(),
            m_untrusted_net_list.size(),
            m_trusted_net_list.data(),
            m_trusted_net_list.size()),
        "Failed to initialize IPv4 set.");
}

bool TrustedProxyConfig::is_trusted(const char* ipstr) const
{
    ib_ip4_t ip;

    IronBee::throw_if_error(
        ib_ip4_str_to_ip(ipstr, &ip),
        "Invalid remote IP address");
    ib_status_t rc =
        ib_ipset4_query(&m_trusted_networks, ip, NULL, NULL, NULL);
    return rc == IB_OK;
}

TrustedProxyModule::TrustedProxyModule(IronBee::Module module) :
    IronBee::ModuleDelegate(module)
{
    module.set_configuration_data<TrustedProxyConfig>();

    m_xff_target = IronBee::VarTarget::acquire_from_string(
        module.engine().main_memory_pool(),
        module.engine().var_config(),
        "request_headers:X-Forwarded-For");

    m_remote_addr_source = IronBee::VarSource::register_(
        module.engine().var_config(),
        IB_S2SL("remote_addr"),
        IB_PHASE_REQUEST_HEADER,
        IB_PHASE_REQUEST_HEADER);

    module.engine().register_configuration_directives()
        .on_off(
            "TrustedProxyUseXFFHeader",
            boost::bind(&TrustedProxyModule::enable_xff_directive,
                        this, _1, _3))
        .list(
            "TrustedProxyIPs",
            boost::bind(&TrustedProxyModule::trusted_ips_directive,
                        this, _1, _3));

    module.engine().register_hooks()
        .context_close(
            boost::bind(&TrustedProxyModule::on_context_close,
                        this, _1, _2))
        .request_header_finished(
            boost::bind(&TrustedProxyModule::set_effective_ip,
                        this, _1, _2));
}

void TrustedProxyModule::enable_xff_directive(
    IronBee::ConfigurationParser cp,
    bool                         enabled
)
{
    TrustedProxyConfig& config =
        module().configuration_data<TrustedProxyConfig>(cp.current_context());
    config.set_xff_enabled(enabled);
}

void TrustedProxyModule::trusted_ips_directive(
    IronBee::ConfigurationParser    cp,
    IronBee::ConstList<const char*> ip_list
)
{
    TrustedProxyConfig& config =
        module().configuration_data<TrustedProxyConfig>(cp.current_context());

    const char* first_arg = ip_list.front();
    if (*first_arg != '+' and *first_arg != '-') {
        config.clear_networks();
    }

    BOOST_FOREACH(const char* arg, ip_list) {
        if (*arg == '+') {
            config.add_trusted_network(arg+1);
        }
        else if (*arg == '-') {
            config.add_untrusted_network(arg+1);
        }
        else {
            config.add_trusted_network(arg);
        }
    }
}

void TrustedProxyModule::on_context_close(
    IronBee::Engine ib,
    IronBee::Context ctx
)
{
    TrustedProxyConfig& config =
        module().configuration_data<TrustedProxyConfig>(ctx);

    config.context_close(ib);
}

void TrustedProxyModule::set_effective_ip(
    IronBee::Engine      ib,
    IronBee::Transaction tx
)
{
    ib_status_t rc;
    IronBee::Context ctx = tx.context();
    TrustedProxyConfig& config =
        module().configuration_data<TrustedProxyConfig>(ctx);

    ib_log_debug_tx(tx.ib(), "checking: %s",
                    tx.connection().remote_ip_string());
    // check actual remote ip against trusted ips
    if (! config.is_trusted(tx.connection().remote_ip_string())) {
        ib_log_debug_tx(tx.ib(), "Remote address '%s' not a trusted proxy.",
                        tx.connection().remote_ip_string());
        return;
    }
    // Remote address is trusted, get the X-Forwarded-For value
    IronBee::ConstList<IronBee::ConstField> xfflist =
        m_xff_target.get(tx.memory_pool(), tx.var_store());

    if (xfflist.size() == 0) {
        // No X-Forwarded-For header found. Nothing to do.
        return;
    }
    // Use the final IP address in X-Forwarded-For
    string forwarded = xfflist.back().to_s();
    list<string> forwarded_list;
    boost::algorithm::split(forwarded_list,
                            forwarded,
                            boost::algorithm::is_any_of(","));

    string remote_ip = boost::algorithm::trim_copy(forwarded_list.back());

    /* Verify that it looks like a valid IP address, ignore it if not */
    rc = ib_ip_validate(remote_ip.c_str());
    if (rc != IB_OK) {
        ib_log_error_tx(tx.ib(),
                        "X-Forwarded-For \"%s\" is not a valid IP address",
                        remote_ip.c_str());
        return;
    }
    char* buf = static_cast<char*>(tx.memory_pool().alloc(remote_ip.length()+1));
    strcpy(buf, remote_ip.c_str());

    /* This will lose the pointer to the original address
     * buffer, but it should be cleaned up with the rest
     * of the memory pool. */
    tx.ib()->remote_ipstr = buf;

    ib_log_debug_tx(tx.ib(), "Remote address changed to \"%s\"", buf);
    IronBee::ByteString remote_addr_bs = IronBee::ByteString::create_alias(
        tx.memory_pool(),
        buf,
        remote_ip.length());
    IronBee::Field remote_addr_field = IronBee::Field::create_byte_string(
        tx.memory_pool(),
        "", 0,
        remote_addr_bs);
    m_remote_addr_source.set(tx.var_store(), remote_addr_field);
}

} // Anonymous namespace
<|endoftext|>