<commit_before>/*
 * Use Gtk+ to show live video from the dc1394 camera
 *
 * Written by John Stowers <john.stowers@gmail.com>
 * Modified for STEREO VISION PROCESSING BY ANDREW BARRY <abarry@csail.mit.edu>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * 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 GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

#include "opencv-calibrate.hpp"

bool left_camera_mode = false;
bool right_camera_mode = false;
bool stereo_mode = true;
int force_brightness = -1, force_exposure = -1;

int camera_custom_setup(dc1394camera_t *camera);

int main(int argc, char *argv[])
{
    // parse configuration file
    // get input arguments

    OpenCvStereoConfig stereo_config;

    string config_file = "";

    ConciseArgs parser(argc, argv);
    parser.add(config_file, "c", "config", "Configuration file containing camera GUIDs, etc.", true);
    parser.add(left_camera_mode, "l", "left-camera", "Calibrate just the left camera.");
    parser.add(right_camera_mode, "r", "right-camera", "Calibrate just the right camera.");
    parser.parse();

    // parse the config file
    if (ParseConfigFile(config_file, &stereo_config) != true)
    {
        fprintf(stderr, "Failed to parse configuration file, quitting.\n");
        return -1;
    }

    if (left_camera_mode || right_camera_mode)
    {
        stereo_mode = false;
    }

    uint64 guid = stereo_config.guidLeft;
    uint64 guid2 = stereo_config.guidRight;

    dc1394_t        *d;
    dc1394camera_t  *camera;
    dc1394error_t   err;

    Mat frame_array_left[MAX_FRAMES];
    Mat frame_array_right[MAX_FRAMES];

    int numFrames = 0;

    // ----- cam 2 -----
    dc1394_t        *d2;
    dc1394camera_t  *camera2;
    dc1394error_t   err2;

    d = dc1394_new ();
    if (!d)
        g_critical("Could not create dc1394 context");

    d2 = dc1394_new ();
    if (!d2)
        g_critical("Could not create dc1394 context for camera 2");

    camera = dc1394_camera_new (d, guid);
    if (!camera)
        g_critical("Could not create dc1394 camera");

    camera2 = dc1394_camera_new (d2, guid2);
    if (!camera2)
        g_critical("Could not create dc1394 camera for camera 2");

    // setup
    err = setup_gray_capture(camera, DC1394_VIDEO_MODE_FORMAT7_1);
    DC1394_ERR_CLN_RTN(err, cleanup_and_exit(camera), "Could not setup camera");

    err2 = setup_gray_capture(camera2, DC1394_VIDEO_MODE_FORMAT7_1);
    DC1394_ERR_CLN_RTN(err2, cleanup_and_exit(camera2), "Could not setup camera number 2");

    // enable auto-exposure
    // turn on the auto exposure feature
    err = dc1394_feature_set_power(camera, DC1394_FEATURE_EXPOSURE, DC1394_ON);
    DC1394_ERR_RTN(err,"Could not turn on the exposure feature");

    err = dc1394_feature_set_mode(camera, DC1394_FEATURE_EXPOSURE, DC1394_FEATURE_MODE_ONE_PUSH_AUTO);
    DC1394_ERR_RTN(err,"Could not turn on Auto-exposure");

    // enable auto-exposure
    // turn on the auto exposure feature
    err = dc1394_feature_set_power(camera2, DC1394_FEATURE_EXPOSURE, DC1394_ON);
    DC1394_ERR_RTN(err,"Could not turn on the exposure feature for cam2");

    err = dc1394_feature_set_mode(camera2, DC1394_FEATURE_EXPOSURE, DC1394_FEATURE_MODE_ONE_PUSH_AUTO);
    DC1394_ERR_RTN(err,"Could not turn on Auto-exposure for cam2");

    // enable camera
    err = dc1394_video_set_transmission(camera, DC1394_ON);
    DC1394_ERR_CLN_RTN(err, cleanup_and_exit(camera), "Could not start camera iso transmission");
    err2 = dc1394_video_set_transmission(camera2, DC1394_ON);
    DC1394_ERR_CLN_RTN(err2, cleanup_and_exit(camera2), "Could not start camera iso transmission for camera number 2");

    if (left_camera_mode || stereo_mode)
    {
        InitBrightnessSettings(camera, camera2);
        MatchBrightnessSettings(camera, camera2, true, force_brightness, force_exposure);
    } else {
        // use the right camera as the master for brightness
        // since we're calibrating that one
        InitBrightnessSettings(camera2, camera);
        MatchBrightnessSettings(camera2, camera, true);
    }

    // make opencv windows
    if (left_camera_mode || stereo_mode)
    {
    	namedWindow("Input Left", CV_WINDOW_AUTOSIZE);
    }

    if (right_camera_mode || stereo_mode)
    {
    	namedWindow("Input Right", CV_WINDOW_AUTOSIZE);
    	moveWindow("Input Right", 500, 100);
    }

    CvSize size;

    Mat cornersL, cornersR;

    int i;
    while (numFrames < MAX_FRAMES) {

        Mat chessL, chessR;

        // each loop dump a bunch of frames to clear the buffer
        MatchBrightnessSettings(camera, camera2);
        for (i=0;i<10;i++)
        {
            if (left_camera_mode || stereo_mode)
            {
                chessL = GetFrameFormat7(camera);
            }

            if (right_camera_mode || stereo_mode)
            {
                chessR = GetFrameFormat7(camera2);
            }
        }

        // copy the images for drawing/display
        size = chessL.size();
        Mat chessLc;
        chessLc.create(size, CV_32FC3);
        Mat chessRc;
        chessRc.create(size, CV_32FC3);

        // attempt checkerboard matching
        bool foundPattern = true; // set to true so we can do an OR
                                  // later if we're only using one
                                  // camera

        if (left_camera_mode || stereo_mode)
        {
            foundPattern = findChessboardCorners(chessL, Size(CHESS_X, CHESS_Y), cornersL);
        }

        if (right_camera_mode || stereo_mode)
        {
            foundPattern = foundPattern & findChessboardCorners(chessR, Size(CHESS_X, CHESS_Y), cornersR);
        }

        if (left_camera_mode || stereo_mode)
        {
            cvtColor( chessL, chessLc, CV_GRAY2BGR );
            drawChessboardCorners(chessLc, Size(CHESS_X, CHESS_Y), cornersL, foundPattern);
            imshow("Input Left", chessLc);
        }

        if (right_camera_mode || stereo_mode)
        {
            cvtColor(chessR, chessRc, CV_GRAY2BGR);
            drawChessboardCorners(chessRc, Size(CHESS_X, CHESS_Y), cornersR, foundPattern);

            imshow("Input Right", chessRc);
        }

		// key codes:
		// page up: 654365
		// page down: 65366
		// b: 98

		char key = waitKey();
		//printf("%d\n", (int)key);
		if (key == 98)
		{
		    break;
		} else if (key == 86){
		    if (foundPattern)
		    {
		        // this was a good one -- save it
		        frame_array_left[numFrames] = chessL;
                frame_array_right[numFrames] = chessR;

                // give the user some guidence on the number
                // of frames they should be using
                if (stereo_mode)
                {
                    printf("Saved frame %d / about 10\n", numFrames);
                } else {
                    printf("Saved frame %d / about 20-30\n", numFrames);
                }

                numFrames ++;
            } else {
                printf("Not saving frame since did not find a checkboard.\n");
            }
		} else if (key == 'b') {
            force_brightness ++;
            cout << "Brightness: " << force_brightness;
        } else if (key == 'B') {
            force_brightness --;
            cout << "Brightness: " << force_brightness;
        } else if (key == 'e') {
            force_exposure ++;
            cout << "Exposure: " << force_exposure;
        } else if (key == 'E') {
            force_exposure --;
            cout << "Exposure: " << force_exposure;
        }

	}

    printf("\n\n");

    // clear out the calibration directory
    printf("Deleting old images...\nrm calibrationImages/*.ppm\n");
    system("rm calibrationImages/*.ppm");
    printf("done.\n");

    char filename[1000];

    for (i=0;i<numFrames;i++)
    {
        if (left_camera_mode || stereo_mode)
        {
            sprintf(filename, "calibrationImages/cam1-%05d.ppm", i+1);
            imwrite(filename, frame_array_left[i]);
        }

        if (right_camera_mode || stereo_mode)
        {
            sprintf(filename, "calibrationImages/cam2-%05d.ppm", i+1);
            imwrite(filename, frame_array_right[i]);
        }

        printf("Writing frame %d\n", i);
    }

    printf("\n\n");

    destroyWindow("Input Left");
    destroyWindow("Input Right");

    // stop data transmission
    err = dc1394_video_set_transmission(camera, DC1394_OFF);
    DC1394_ERR_CLN_RTN(err,cleanup_and_exit(camera),"Could not stop the camera");

    err2 = dc1394_video_set_transmission(camera2, DC1394_OFF);
    DC1394_ERR_CLN_RTN(err2,cleanup_and_exit(camera2),"Could not stop the camera 2");

    // close camera
    cleanup_and_exit(camera);
    cleanup_and_exit(camera2);
    dc1394_free (d);
    dc1394_free (d2);

    return 0;
}

<commit_msg>can pass in brightness and exposure as params<commit_after>/*
 * Use Gtk+ to show live video from the dc1394 camera
 *
 * Written by John Stowers <john.stowers@gmail.com>
 * Modified for STEREO VISION PROCESSING BY ANDREW BARRY <abarry@csail.mit.edu>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * 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 GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

#include "opencv-calibrate.hpp"

bool left_camera_mode = false;
bool right_camera_mode = false;
bool stereo_mode = true;
int force_brightness = -1, force_exposure = -1;

int camera_custom_setup(dc1394camera_t *camera);

int main(int argc, char *argv[])
{
    // parse configuration file
    // get input arguments

    OpenCvStereoConfig stereo_config;
    string config_file = "";

    ConciseArgs parser(argc, argv);
    parser.add(config_file, "c", "config", "Configuration file containing camera GUIDs, etc.", true);
    parser.add(left_camera_mode, "l", "left-camera", "Calibrate just the left camera.");
    parser.add(right_camera_mode, "r", "right-camera", "Calibrate just the right camera.");
    parser.add(force_brightness, "b", "brightness", "set brightness to this level");
    parser.add(force_exposure, "e", "exposure", "set exposure to this level");
    parser.parse();

    // parse the config file
    if (ParseConfigFile(config_file, &stereo_config) != true)
    {
        fprintf(stderr, "Failed to parse configuration file, quitting.\n");
        return -1;
    }

    if (left_camera_mode || right_camera_mode)
    {
        stereo_mode = false;
    }

    uint64 guid = stereo_config.guidLeft;
    uint64 guid2 = stereo_config.guidRight;

    dc1394_t        *d;
    dc1394camera_t  *camera;
    dc1394error_t   err;

    Mat frame_array_left[MAX_FRAMES];
    Mat frame_array_right[MAX_FRAMES];

    int numFrames = 0;

    // ----- cam 2 -----
    dc1394_t        *d2;
    dc1394camera_t  *camera2;
    dc1394error_t   err2;

    d = dc1394_new ();
    if (!d)
        g_critical("Could not create dc1394 context");

    d2 = dc1394_new ();
    if (!d2)
        g_critical("Could not create dc1394 context for camera 2");

    camera = dc1394_camera_new (d, guid);
    if (!camera)
        g_critical("Could not create dc1394 camera");

    camera2 = dc1394_camera_new (d2, guid2);
    if (!camera2)
        g_critical("Could not create dc1394 camera for camera 2");

    // setup
    err = setup_gray_capture(camera, DC1394_VIDEO_MODE_FORMAT7_1);
    DC1394_ERR_CLN_RTN(err, cleanup_and_exit(camera), "Could not setup camera");

    err2 = setup_gray_capture(camera2, DC1394_VIDEO_MODE_FORMAT7_1);
    DC1394_ERR_CLN_RTN(err2, cleanup_and_exit(camera2), "Could not setup camera number 2");

    // enable auto-exposure
    // turn on the auto exposure feature
    err = dc1394_feature_set_power(camera, DC1394_FEATURE_EXPOSURE, DC1394_ON);
    DC1394_ERR_RTN(err,"Could not turn on the exposure feature");

    err = dc1394_feature_set_mode(camera, DC1394_FEATURE_EXPOSURE, DC1394_FEATURE_MODE_ONE_PUSH_AUTO);
    DC1394_ERR_RTN(err,"Could not turn on Auto-exposure");

    // enable auto-exposure
    // turn on the auto exposure feature
    err = dc1394_feature_set_power(camera2, DC1394_FEATURE_EXPOSURE, DC1394_ON);
    DC1394_ERR_RTN(err,"Could not turn on the exposure feature for cam2");

    err = dc1394_feature_set_mode(camera2, DC1394_FEATURE_EXPOSURE, DC1394_FEATURE_MODE_ONE_PUSH_AUTO);
    DC1394_ERR_RTN(err,"Could not turn on Auto-exposure for cam2");

    // enable camera
    err = dc1394_video_set_transmission(camera, DC1394_ON);
    DC1394_ERR_CLN_RTN(err, cleanup_and_exit(camera), "Could not start camera iso transmission");
    err2 = dc1394_video_set_transmission(camera2, DC1394_ON);
    DC1394_ERR_CLN_RTN(err2, cleanup_and_exit(camera2), "Could not start camera iso transmission for camera number 2");

    if (left_camera_mode || stereo_mode)
    {
        InitBrightnessSettings(camera, camera2);
        MatchBrightnessSettings(camera, camera2, true, force_brightness, force_exposure);
    } else {
        // use the right camera as the master for brightness
        // since we're calibrating that one
        InitBrightnessSettings(camera2, camera);
        MatchBrightnessSettings(camera2, camera, true);
    }

    // make opencv windows
    if (left_camera_mode || stereo_mode)
    {
    	namedWindow("Input Left", CV_WINDOW_AUTOSIZE);
    }

    if (right_camera_mode || stereo_mode)
    {
    	namedWindow("Input Right", CV_WINDOW_AUTOSIZE);
    	moveWindow("Input Right", 500, 100);
    }

    CvSize size;

    Mat cornersL, cornersR;

    int i;

    while (numFrames < MAX_FRAMES) {

        Mat chessL, chessR;

        // each loop dump a bunch of frames to clear the buffer
        MatchBrightnessSettings(camera, camera2, true, force_brightness, force_exposure);
        for (i=0;i<10;i++)
        {
            if (left_camera_mode || stereo_mode)
            {
                chessL = GetFrameFormat7(camera);
            }

            if (right_camera_mode || stereo_mode)
            {
                chessR = GetFrameFormat7(camera2);
            }
        }

        // copy the images for drawing/display
        size = chessL.size();
        Mat chessLc;
        chessLc.create(size, CV_32FC3);
        Mat chessRc;
        chessRc.create(size, CV_32FC3);

        // attempt checkerboard matching
        bool foundPattern = true; // set to true so we can do an OR
                                  // later if we're only using one
                                  // camera

        if (left_camera_mode || stereo_mode)
        {
            foundPattern = findChessboardCorners(chessL, Size(CHESS_X, CHESS_Y), cornersL);
        }

        if (right_camera_mode || stereo_mode)
        {
            foundPattern = foundPattern & findChessboardCorners(chessR, Size(CHESS_X, CHESS_Y), cornersR);
        }

        if (left_camera_mode || stereo_mode)
        {
            cvtColor( chessL, chessLc, CV_GRAY2BGR );
            drawChessboardCorners(chessLc, Size(CHESS_X, CHESS_Y), cornersL, foundPattern);
            imshow("Input Left", chessLc);
        }

        if (right_camera_mode || stereo_mode)
        {
            cvtColor(chessR, chessRc, CV_GRAY2BGR);
            drawChessboardCorners(chessRc, Size(CHESS_X, CHESS_Y), cornersR, foundPattern);

            imshow("Input Right", chessRc);
        }

		// key codes:
		// page up: 654365
		// page down: 65366
		// b: 98
		char key = waitKey();
		//printf("%d\n", (int)key);
		if (key == 98)
		{
		    break;
		} else if (key == 86){
		    if (foundPattern)
		    {
		        // this was a good one -- save it
		        frame_array_left[numFrames] = chessL;
                frame_array_right[numFrames] = chessR;

                // give the user some guidence on the number
                // of frames they should be using
                if (stereo_mode)
                {
                    printf("Saved frame %d / about 10\n", numFrames);
                } else {
                    printf("Saved frame %d / about 20-30\n", numFrames);
                }

                numFrames ++;
            } else {
                printf("Not saving frame since did not find a checkboard.\n");
            }
		} else if (key == 'W') {
            force_brightness +=20;
            cout << "Brightness: " << force_brightness << "\n";
        } else if (key == 'w') {
            force_brightness -=20;
            cout << "Brightness: " << force_brightness << "\n";
        } else if (key == 'E') {
            force_exposure +=20;
            cout << "Exposure: " << force_exposure << "\n";
        } else if (key == 'e') {
            force_exposure -=20;
            cout << "Exposure: " << force_exposure << "\n";
        }
	}

    printf("\n\n");

    // clear out the calibration directory
    printf("Deleting old images...\nrm calibrationImages/*.ppm\n");
    system("rm calibrationImages/*.ppm");
    printf("done.\n");

    char filename[1000];

    for (i=0;i<numFrames;i++)
    {
        if (left_camera_mode || stereo_mode)
        {
            sprintf(filename, "calibrationImages/cam1-%05d.ppm", i+1);
            imwrite(filename, frame_array_left[i]);
        }

        if (right_camera_mode || stereo_mode)
        {
            sprintf(filename, "calibrationImages/cam2-%05d.ppm", i+1);
            imwrite(filename, frame_array_right[i]);
        }

        printf("Writing frame %d\n", i);
    }

    printf("\n\n");

    destroyWindow("Input Left");
    destroyWindow("Input Right");

    // stop data transmission
    err = dc1394_video_set_transmission(camera, DC1394_OFF);
    DC1394_ERR_CLN_RTN(err,cleanup_and_exit(camera),"Could not stop the camera");

    err2 = dc1394_video_set_transmission(camera2, DC1394_OFF);
    DC1394_ERR_CLN_RTN(err2,cleanup_and_exit(camera2),"Could not stop the camera 2");

    // close camera
    cleanup_and_exit(camera);
    cleanup_and_exit(camera2);
    dc1394_free (d);
    dc1394_free (d2);

    return 0;
}

<|endoftext|>


<commit_before>/*
 * Copyright (C) 2014  Daniel Vrátil <dvratil@redhat.com>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * 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 GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */

#include "partstreamer.h"
#include "parthelper.h"
#include "parttypehelper.h"
#include "selectquerybuilder.h"
#include "dbconfig.h"
#include "connection.h"
#include "capabilities_p.h"
#include "imapstreamparser.h"
#include "response.h"

#include <libs/imapparser_p.h>
#include <libs/protocol_p.h>

#include <unistd.h>

#include <QFile>
#include <QFileInfo>

using namespace Akonadi;
using namespace Akonadi::Server;

PartStreamer::PartStreamer(Connection *connection, ImapStreamParser *parser,
                           const PimItem &pimItem, QObject *parent)
    : QObject(parent)
    , mConnection(connection)
    , mStreamParser(parser)
    , mItem(pimItem)
{
}

PartStreamer::~PartStreamer()
{
}

QByteArray PartStreamer::error() const
{
    return mError;
}

bool PartStreamer::streamNonliteral(Part& part, qint64& partSize, QByteArray& value)
{
    value = mStreamParser->readString();
    if (part.partType().ns() == QLatin1String("PLD")) {
        partSize = value.size();
    }
    if (!mDataChanged) {
        mDataChanged = (value.size() != part.datasize());
    }

    // only relevant for non-literals or non-external literals
    // only fallback to data comparision if part already exists and sizes match
    if (!mDataChanged) {
        mDataChanged = (value != PartHelper::translateData(part));
    }

    if (mDataChanged) {
        if (part.isValid()) {
            PartHelper::update(&part, value, value.size());
        } else {
//           akDebug() << "insert from Store::handleLine: " << value.left(100);
            part.setData(value);
            part.setDatasize(value.size());
            if (!PartHelper::insert(&part)) {
                mError = "Unable to add item part";
                return false;
            }
        }
    }

    return true;
}

bool PartStreamer::streamLiteral(Part& part, qint64& partSize, QByteArray& value)
{
    const qint64 dataSize = mStreamParser->remainingLiteralSize();
    if (part.partType().ns() == QLatin1String("PLD")) {
        partSize = dataSize;
        // Shortcut: if sizes differ, we don't need to compare data later no in order
        // to detect whether the part has changed
        if (!mDataChanged) {
            mDataChanged = (partSize != part.datasize());
        }
    }

    //actual case when streaming storage is used: external payload is enabled, data is big enough in a literal
    const bool storeInFile = dataSize > DbConfig::configuredDatabase()->sizeThreshold();
    if (storeInFile) {
        return streamLiteralToFile(dataSize, part, value);
    } else {
        mStreamParser->sendContinuationResponse(dataSize);
        //don't write in streaming way as the data goes to the database
        while (!mStreamParser->atLiteralEnd()) {
            value += mStreamParser->readLiteralPart();
        }
        PartHelper::update(&part, value, value.size());
    }

    if (part.isValid()) {
        if (!part.update()) {
            mError = "Failed to update part in database";
            return false;
        }
    } else {
        if (!part.insert()) {
            mError = "Failed to insert part to database";
            return false;
        }
    }

    return true;
}

bool PartStreamer::streamLiteralToFile(qint64 dataSize, Part &part, QByteArray &value)
{
    const bool directStreaming = mConnection->capabilities().directStreaming();

    QByteArray origData;
    if (!mDataChanged && mCheckChanged) {
        origData = PartHelper::translateData(part);
    }

    if (directStreaming) {
        bool ok = streamLiteralToFileDirectly(dataSize, part);
        if (!ok) {
            return false;
        }
    } else {
        mStreamParser->sendContinuationResponse(dataSize);

        // use first part as value for the initial insert into / update to the database.
        // this will give us a proper filename to stream the rest of the parts contents into
        // NOTE: we have to set the correct size (== dataSize) directly
        value = mStreamParser->readLiteralPart();
        // akDebug() << Q_FUNC_INFO << "VALUE in STORE: " << value << value.size() << dataSize;

        if (part.isValid()) {
            PartHelper::update(&part, value, dataSize);
        } else {
            part.setData( value );
            part.setDatasize( dataSize );
            if ( !PartHelper::insert( &part ) ) {
                mError = "Unable to add item part";
                return false;
            }
        }

        //the actual streaming code for the remaining parts:
        // reads from the parser, writes immediately to the file
        QFile partFile(PartHelper::resolveAbsolutePath(part.data()));
        try {
            PartHelper::streamToFile(mStreamParser, partFile, QIODevice::WriteOnly | QIODevice::Append);
        } catch (const PartHelperException &e) {
            mError = e.what();
            return false;
        }
    }

    if (mCheckChanged && !mDataChanged) {
        // This is invoked only when part already exists, data sizes match and
        // caller wants to know whether parts really differ
        mDataChanged = (origData != PartHelper::translateData(part));
    }

    return true;
}

bool PartStreamer::streamLiteralToFileDirectly(qint64 dataSize, Part &part)
{
    QString filename;
    if (part.isValid()) {
        if (part.external()) {
            // Part was external and is still external
            filename = QString::fromLatin1(part.data());
        } else {
            // Part wasn't external, but is now
            filename = PartHelper::fileNameForPart(&part);
        }
        filename = PartHelper::updateFileNameRevision(filename);
    }

    QFileInfo finfo(filename);
    if (finfo.isAbsolute()) {
        filename = finfo.fileName();
    }

    part.setExternal(true);
    part.setDatasize(dataSize);
    part.setData(filename.toLatin1());

    if (part.isValid()) {
        if (!part.update()) {
            mError = "Failed to update part in database";
            return false;
        }
    } else {
        if (!part.insert()) {
            mError = "Failed to insert part into database";
            return false;
        }

        filename = PartHelper::updateFileNameRevision(PartHelper::fileNameForPart(&part));
        part.setData(filename.toLatin1());
        part.update();
    }

    Response response;
    response.setContinuation();
    response.setString("STREAM [FILE " + part.data() + "]");
    Q_EMIT responseAvailable(response);

    const QByteArray reply = mStreamParser->peekString();
    if (reply.startsWith("NO")) {
        akError() << "Client failed to store payload into file";
        mError = "Client failed to store payload into file";
        return false;
    }

    QFile file(PartHelper::resolveAbsolutePath(filename.toLatin1()), this);
    if (!file.exists()) {
        akError() << "External payload file does not exist";
        mError = "External payload file does not exist";
        return false;
    }

    if (file.size() != dataSize) {
        akError() << "Size mismatch! Client advertised" << dataSize << "bytes, but the file is" << file.size() << "bytes!";
        mError = "Payload size mismatch";
        return false;
    }

    return true;
}

bool PartStreamer::stream(const QByteArray &command, bool checkExists,
                          QByteArray &partName, qint64 &partSize, bool *changed)
{
    int partVersion = 0;
    partSize = 0;
    mError.clear();
    mDataChanged = false;
    mCheckChanged = (changed != 0);
    if (changed != 0) {
        *changed = false;
    }

    ImapParser::splitVersionedKey(command, partName, partVersion);
    const PartType partType = PartTypeHelper::fromFqName( partName );

    Part part;

    if (checkExists || mCheckChanged) {
        SelectQueryBuilder<Part> qb;
        qb.addValueCondition(Part::pimItemIdColumn(), Query::Equals, mItem.id());
        qb.addValueCondition(Part::partTypeIdColumn(), Query::Equals, partType.id());
        if (!qb.exec()) {
            mError = "Unable to check item part existence";
            return false;
        }

        Part::List result = qb.result();
        if (!result.isEmpty()) {
            part = result.first();
        }
    }

    // Shortcut: newly created parts are always "changed"
    if (!part.isValid()) {
        mDataChanged = true;
    }

    part.setPartType(partType);
    part.setVersion(partVersion);
    part.setPimItemId(mItem.id());

    QByteArray value;
    bool ok;
    if (mStreamParser->hasLiteral(false)) {
        ok = streamLiteral(part, partSize, value);
    } else {
        ok = streamNonliteral(part, partSize, value);
    }

    if (ok && mCheckChanged) {
        *changed = mDataChanged;
    }

    return ok;
}

<commit_msg>PartStreamer: Make sure the external part storage directory exists<commit_after>/*
 * Copyright (C) 2014  Daniel Vrátil <dvratil@redhat.com>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * 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 GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */

#include "partstreamer.h"
#include "parthelper.h"
#include "parttypehelper.h"
#include "selectquerybuilder.h"
#include "dbconfig.h"
#include "connection.h"
#include "capabilities_p.h"
#include "imapstreamparser.h"
#include "response.h"

#include <libs/imapparser_p.h>
#include <libs/protocol_p.h>
#include <shared/akstandarddirs.h>

#include <unistd.h>

#include <QFile>
#include <QFileInfo>

using namespace Akonadi;
using namespace Akonadi::Server;

PartStreamer::PartStreamer(Connection *connection, ImapStreamParser *parser,
                           const PimItem &pimItem, QObject *parent)
    : QObject(parent)
    , mConnection(connection)
    , mStreamParser(parser)
    , mItem(pimItem)
{
    // Make sure the file_db_data path exists
    AkStandardDirs::saveDir( "data", QLatin1String( "file_db_data" ) );
}

PartStreamer::~PartStreamer()
{
}

QByteArray PartStreamer::error() const
{
    return mError;
}

bool PartStreamer::streamNonliteral(Part& part, qint64& partSize, QByteArray& value)
{
    value = mStreamParser->readString();
    if (part.partType().ns() == QLatin1String("PLD")) {
        partSize = value.size();
    }
    if (!mDataChanged) {
        mDataChanged = (value.size() != part.datasize());
    }

    // only relevant for non-literals or non-external literals
    // only fallback to data comparision if part already exists and sizes match
    if (!mDataChanged) {
        mDataChanged = (value != PartHelper::translateData(part));
    }

    if (mDataChanged) {
        if (part.isValid()) {
            PartHelper::update(&part, value, value.size());
        } else {
//           akDebug() << "insert from Store::handleLine: " << value.left(100);
            part.setData(value);
            part.setDatasize(value.size());
            if (!PartHelper::insert(&part)) {
                mError = "Unable to add item part";
                return false;
            }
        }
    }

    return true;
}

bool PartStreamer::streamLiteral(Part& part, qint64& partSize, QByteArray& value)
{
    const qint64 dataSize = mStreamParser->remainingLiteralSize();
    if (part.partType().ns() == QLatin1String("PLD")) {
        partSize = dataSize;
        // Shortcut: if sizes differ, we don't need to compare data later no in order
        // to detect whether the part has changed
        if (!mDataChanged) {
            mDataChanged = (partSize != part.datasize());
        }
    }

    //actual case when streaming storage is used: external payload is enabled, data is big enough in a literal
    const bool storeInFile = dataSize > DbConfig::configuredDatabase()->sizeThreshold();
    if (storeInFile) {
        return streamLiteralToFile(dataSize, part, value);
    } else {
        mStreamParser->sendContinuationResponse(dataSize);
        //don't write in streaming way as the data goes to the database
        while (!mStreamParser->atLiteralEnd()) {
            value += mStreamParser->readLiteralPart();
        }
        PartHelper::update(&part, value, value.size());
    }

    if (part.isValid()) {
        if (!part.update()) {
            mError = "Failed to update part in database";
            return false;
        }
    } else {
        if (!part.insert()) {
            mError = "Failed to insert part to database";
            return false;
        }
    }

    return true;
}

bool PartStreamer::streamLiteralToFile(qint64 dataSize, Part &part, QByteArray &value)
{
    const bool directStreaming = mConnection->capabilities().directStreaming();

    QByteArray origData;
    if (!mDataChanged && mCheckChanged) {
        origData = PartHelper::translateData(part);
    }

    if (directStreaming) {
        bool ok = streamLiteralToFileDirectly(dataSize, part);
        if (!ok) {
            return false;
        }
    } else {
        mStreamParser->sendContinuationResponse(dataSize);

        // use first part as value for the initial insert into / update to the database.
        // this will give us a proper filename to stream the rest of the parts contents into
        // NOTE: we have to set the correct size (== dataSize) directly
        value = mStreamParser->readLiteralPart();
        // akDebug() << Q_FUNC_INFO << "VALUE in STORE: " << value << value.size() << dataSize;

        if (part.isValid()) {
            PartHelper::update(&part, value, dataSize);
        } else {
            part.setData( value );
            part.setDatasize( dataSize );
            if ( !PartHelper::insert( &part ) ) {
                mError = "Unable to add item part";
                return false;
            }
        }

        //the actual streaming code for the remaining parts:
        // reads from the parser, writes immediately to the file
        QFile partFile(PartHelper::resolveAbsolutePath(part.data()));
        try {
            PartHelper::streamToFile(mStreamParser, partFile, QIODevice::WriteOnly | QIODevice::Append);
        } catch (const PartHelperException &e) {
            mError = e.what();
            return false;
        }
    }

    if (mCheckChanged && !mDataChanged) {
        // This is invoked only when part already exists, data sizes match and
        // caller wants to know whether parts really differ
        mDataChanged = (origData != PartHelper::translateData(part));
    }

    return true;
}

bool PartStreamer::streamLiteralToFileDirectly(qint64 dataSize, Part &part)
{
    QString filename;
    if (part.isValid()) {
        if (part.external()) {
            // Part was external and is still external
            filename = QString::fromLatin1(part.data());
        } else {
            // Part wasn't external, but is now
            filename = PartHelper::fileNameForPart(&part);
        }
        filename = PartHelper::updateFileNameRevision(filename);
    }

    QFileInfo finfo(filename);
    if (finfo.isAbsolute()) {
        filename = finfo.fileName();
    }

    part.setExternal(true);
    part.setDatasize(dataSize);
    part.setData(filename.toLatin1());

    if (part.isValid()) {
        if (!part.update()) {
            mError = "Failed to update part in database";
            return false;
        }
    } else {
        if (!part.insert()) {
            mError = "Failed to insert part into database";
            return false;
        }

        filename = PartHelper::updateFileNameRevision(PartHelper::fileNameForPart(&part));
        part.setData(filename.toLatin1());
        part.update();
    }

    Response response;
    response.setContinuation();
    response.setString("STREAM [FILE " + part.data() + "]");
    Q_EMIT responseAvailable(response);

    const QByteArray reply = mStreamParser->peekString();
    if (reply.startsWith("NO")) {
        akError() << "Client failed to store payload into file";
        mError = "Client failed to store payload into file";
        return false;
    }

    QFile file(PartHelper::resolveAbsolutePath(filename.toLatin1()), this);
    if (!file.exists()) {
        akError() << "External payload file does not exist";
        mError = "External payload file does not exist";
        return false;
    }

    if (file.size() != dataSize) {
        akError() << "Size mismatch! Client advertised" << dataSize << "bytes, but the file is" << file.size() << "bytes!";
        mError = "Payload size mismatch";
        return false;
    }

    return true;
}

bool PartStreamer::stream(const QByteArray &command, bool checkExists,
                          QByteArray &partName, qint64 &partSize, bool *changed)
{
    int partVersion = 0;
    partSize = 0;
    mError.clear();
    mDataChanged = false;
    mCheckChanged = (changed != 0);
    if (changed != 0) {
        *changed = false;
    }

    ImapParser::splitVersionedKey(command, partName, partVersion);
    const PartType partType = PartTypeHelper::fromFqName( partName );

    Part part;

    if (checkExists || mCheckChanged) {
        SelectQueryBuilder<Part> qb;
        qb.addValueCondition(Part::pimItemIdColumn(), Query::Equals, mItem.id());
        qb.addValueCondition(Part::partTypeIdColumn(), Query::Equals, partType.id());
        if (!qb.exec()) {
            mError = "Unable to check item part existence";
            return false;
        }

        Part::List result = qb.result();
        if (!result.isEmpty()) {
            part = result.first();
        }
    }

    // Shortcut: newly created parts are always "changed"
    if (!part.isValid()) {
        mDataChanged = true;
    }

    part.setPartType(partType);
    part.setVersion(partVersion);
    part.setPimItemId(mItem.id());

    QByteArray value;
    bool ok;
    if (mStreamParser->hasLiteral(false)) {
        ok = streamLiteral(part, partSize, value);
    } else {
        ok = streamNonliteral(part, partSize, value);
    }

    if (ok && mCheckChanged) {
        *changed = mDataChanged;
    }

    return ok;
}

<|endoftext|>


<commit_before>
#ifndef ROPUFU_SETTLERS_ONLINE_PREFIX_DATABASE_HPP_INCLUDED
#define ROPUFU_SETTLERS_ONLINE_PREFIX_DATABASE_HPP_INCLUDED

#include <aftermath/not_an_error.hpp>
#include <experimental/filesystem>
#include <nlohmann/json.hpp>

#include "char_string.hpp"
#include "logger.hpp"
#include "prefix_tree.hpp"

#include <cstddef> // std::size_t
#include <fstream> // std::ifstream
#include <map> // std::map
#include <set> // std::set
#include <string> // std::string
#include <type_traits> // std::is_same
#include <vector> // std::vector

namespace ropufu
{
    namespace settlers_online
    {
        template <typename t_key_type, typename t_skeleton_type = t_key_type>
        struct lookup
        {
            using type = lookup<t_key_type, t_skeleton_type>;
            using key_type = t_key_type;
            using skeleton_type = t_skeleton_type;

        private:
            std::map<key_type, std::set<skeleton_type>> m_skeleton = { }; // For eack key, keeps track of all associated relaxed names.
            std::map<skeleton_type, std::set<key_type>> m_inverse_skeleton = { }; // For reach relaxed name, keeps track of all associated keys.

        public:
            void clear() noexcept
            {
                this->m_skeleton.clear();
                this->m_inverse_skeleton.clear();
            } // clear(...)

            void update(const key_type& key, const std::set<skeleton_type>& values) noexcept
            {
                // Record the values associated with the given key.
                this->m_skeleton.emplace(key, values);

                // Now the reverse: find all keys associated with values.
                for (const skeleton_type& weak : values)
                {
                    auto search = this->m_inverse_skeleton.find(weak);
                    if (search != this->m_inverse_skeleton.end()) search->second.insert(key);
                    else
                    {
                        std::set<key_type> matches { key }; // This is the only key that matches <weak>, since otherwise it would've been present in the inverse skeleton.
                        this->m_inverse_skeleton.emplace(weak, matches);
                    }
                }
            } // update(...)

            template <typename t_predicate_type>
            std::size_t try_find(const skeleton_type& query, key_type& key, const t_predicate_type& filter) const noexcept
            {
                auto search = this->m_inverse_skeleton.find(query);
                if (search != this->m_inverse_skeleton.end())
                {
                    const std::set<key_type>& matches = search->second;
                    std::size_t count = 0;
                    for (const key_type& maybe : matches)
                    {
                        if (filter(maybe)) { key = maybe; count++; }
                    }
                    return count;
                }
                return 0;
            } // try_find(...)
        }; // struct lookup

        /** @brief Auxiliary class for \c prefix_database specializations. */
        template <typename t_value_type>
        struct prefix_builder
        {
            using value_type = t_value_type;
            using key_type = std::string;

            static key_type build_key(const value_type& unit) noexcept
            {
                static_assert(false, "<prefix_builder> not specialized.");
            } // build_key(...)

            static key_type build_primary_name(const value_type& unit) noexcept
            {
                static_assert(false, "<prefix_builder> not specialized.");
            } // build_primary_name(...)

            /** @brief Collection of names associated with the \p unit.
             *  @remark Return type has to support range-based loops.
             */
            static std::vector<key_type> names(const value_type& unit) noexcept
            {
                static_assert(false, "<prefix_builder> not specialized.");
            } // names(...)
        }; // struct prefix_builder

        /** @brief Class for accessing known units.
         *  @remark This is an abstract class (CRTP). For more information on CRTP see https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern
         *  @remark The class has to implement the following functions:
         *          static key_type build_key(const value_type& unit) const noexcept
         */
        template <typename t_derived_type, typename t_value_type>
        struct prefix_database
        {
            using type = prefix_database<t_derived_type, t_value_type>;
            using derived_type = t_derived_type;
            using value_type = t_value_type;
            using key_type = std::string;

        private:
            using builder_type = prefix_builder<value_type>;

            value_type m_invalid = { }; // Used to indicates invalid return result.
            std::map<key_type, value_type> m_database = { }; // Primary key: shortest name.
            std::set<key_type> m_primary_names = { }; // Keep track of primary names to allow for fast prefix search.
            char_tree m_primary_name_tree = { }; // Fast prefix search by primary name.
            lookup<key_type, key_type> m_lowercase_lookup = { }; // Lookup for lowercase names.
            lookup<key_type, key_type> m_misspelled_lookup = { }; // Lookup for misspelled names.

            /** Relaxed lookup stage 1. */
            static key_type relax_to_lowercase(const key_type& query) noexcept
            {
                key_type relaxed = query;
                char_string::to_lower(relaxed);
                return relaxed;
            } // relax_to_lowercase(...)

            /** Relaxed lookup stage 2. Assuming stage 1 has already been applied. */
            static key_type relax_spelling(const key_type& query) noexcept
            {
                key_type relaxed = query;
                // Plural: man -> men.
                char_string::replace(relaxed, "men", "man");
                // Plural: ...es or ...s.
                if (char_string::ends_with(relaxed, "es")) relaxed = relaxed.substr(0, relaxed.length() - 2);
                else if (char_string::ends_with(relaxed, "s")) relaxed = relaxed.substr(0, relaxed.length() - 1);

                // Collapse all repeated letters for "reasonably" long words.
                if (relaxed.length() > 4)
                {
                    char previous = relaxed[0];
                    for (std::size_t i = 1; i < relaxed.size(); ++i)
                    {
                        if (relaxed[i] == previous) relaxed.erase(i, 1);
                        else previous = relaxed[i];
                    }
                }

                return relaxed;
            } // relax_spelling(...)

            /** Assuming the unit names have already been subject to \c deep_trim_copy. */
            void update_lookup(const key_type& key, const value_type& unit) noexcept
            {
                std::set<key_type> relaxed_names { };
                std::set<key_type> misspelled_names { };
                for (const key_type& name : builder_type::names(unit))
                {
                    key_type stage1 = type::relax_to_lowercase(name);
                    key_type stage2 = type::relax_spelling(stage1);
                    relaxed_names.insert(stage1);
                    misspelled_names.insert(stage2);
                }

                this->m_lowercase_lookup.update(key, relaxed_names);
                this->m_misspelled_lookup.update(key, misspelled_names);
            } // update_lookup(...)

        protected:
            /** @brief Happens when \c clear() is called. */
            void on_clear() noexcept
            {
                constexpr bool is_overwritten = std::is_same<
                    decltype(&derived_type::on_clear),
                    decltype(&type::on_clear)>::value;
                if (!is_overwritten) return; // Do nothing if the function has not been overridden.

                derived_type* that = static_cast<derived_type*>(this);
                that->on_clear();
            } // on_clear(...)

            /** @brief Happens when a unit is about to be loaded. */
            template <typename t_logger_type>
            void on_loading(const value_type& unit, bool& do_cancel, t_logger_type& logger) const noexcept
            {
                do_cancel = false;
                constexpr bool is_overwritten = std::is_same<
                    decltype(&derived_type::template on_loading<t_logger_type>),
                    decltype(&type::template on_loading<t_logger_type>)>::value;
                if (!is_overwritten) return; // Do nothing if the function has not been overridden.

                const derived_type* that = static_cast<const derived_type*>(this);
                that->on_loading(unit, do_cancel, logger);
            } // on_loading(...)

            /** @brief Happens when a unit has been loaded. */
            void on_loaded(const value_type& unit) noexcept
            {
                constexpr bool is_overwritten = std::is_same<
                    decltype(&derived_type::on_loaded),
                    decltype(&type::on_loaded)>::value;
                if (!is_overwritten) return; // Do nothing if the function has not been overridden.

                derived_type* that = static_cast<derived_type*>(this);
                that->on_loaded(unit);
            } // on_loaded(...)

        public:
            prefix_database() noexcept { }

            /** Clears the contents of the database. */
            void clear() noexcept
            {
                this->m_database.clear();
                this->m_primary_names.clear();
                this->m_primary_name_tree.clear();
                this->m_lowercase_lookup.clear();
                this->m_misspelled_lookup.clear();
                this->on_clear();
            } // clear(...)

            const std::map<key_type, value_type>& data() const noexcept { return this->m_database; }

            /** @brief Access elements by key.
             *  @exception not_an_error::out_of_range This error is pushed to \c quiet_error if specified \p key is not in the database.
             */
            const value_type& at(const key_type& key) const noexcept
            {
                auto search = this->m_database.find(key);
                if (search != this->m_database.end()) return search->second;

                aftermath::quiet_error::instance().push(
                    aftermath::not_an_error::out_of_range,
                    aftermath::severity_level::fatal,
                    "<key> is not present in the database.", __FUNCTION__, __LINE__);
                return this->m_invalid;
            } // at(...)

            template <typename t_predicate_type>
            bool try_find(const key_type& query, value_type& unit, const t_predicate_type& filter) const noexcept
            {
                detail::no_logger logger { };
                return this->try_find(query, unit, filter, logger);
            } // try_find(...)

            template <typename t_predicate_type, typename t_logger_type>
            bool try_find(const key_type& query, value_type& unit, const t_predicate_type& filter, t_logger_type& logger) const noexcept
            {
                auto search = this->m_database.find(query);
                if (search != this->m_database.end())
                {
                    unit = search->second;
                    if (filter(unit)) return true;
                }

                // Primary search failed. Secondary search: all lowercase!
                key_type key { };
                key_type lowercase = unit_database::relax_to_lowercase(query);
                key_type misspelled = unit_database::relax_spelling(lowercase);
                std::size_t count_matches = 0;

                // Stage 0: prefix tree search.
                bool is_single = false;
                key_type first_prefix_match = this->m_primary_name_tree.first(lowercase, is_single);
                //count_matches = this->m_primary_name_tree.count(lowercase);
                // Only one terminus matches the prefix.
                if (is_single) lowercase = first_prefix_match; // Overwrite the search query.

                // Stage 1: lowercase lookup.
                count_matches = this->m_lowercase_lookup.try_find(lowercase, key, [&] (const key_type& maybe) { return filter(this->m_database.at(maybe)); });
                if (count_matches >= 1)
                {
                    unit = this->m_database.at(key);
                    if (count_matches == 1) return true;
                    logger.write(std::string("Multiple units match the specified query: ") + lowercase + std::string("."));
                    return false;
                }

                // Stage 2: misspelled lookup.
                count_matches = this->m_misspelled_lookup.try_find(misspelled, key, [&] (const key_type& maybe) { return filter(this->m_database.at(maybe)); });
                if (count_matches >= 1)
                {
                    unit = this->m_database.at(key);
                    if (count_matches == 1) return true;
                    logger.write(std::string("Multiple units match the specified query: ") + misspelled + std::string("."));
                    return false;
                }
                return false;
            } // try_find(...)

            /** Load units from .json files in a specified folder. */
            template <typename t_logger_type>
            bool add(const value_type& unit, t_logger_type& logger) noexcept
            {
                bool do_cancel = false;
                this->on_loading(unit, do_cancel, logger);
                if (do_cancel) return false;

                key_type key = builder_type::build_key(unit);
                key_type primary_name = builder_type::build_primary_name(unit);

                if (this->m_database.count(key) != 0)
                {
                    logger.write(std::string("Unit with name ") + key + std::string(" already exists."));
                    return false;
                }
                else if (this->m_primary_names.count(primary_name) != 0)
                {
                    logger.write(std::string("Unit with a similar primary name ") + primary_name + std::string(" already exists."));
                    return false;
                }
                else
                {
                    this->m_database.emplace(key, unit);
                    this->m_primary_names.insert(primary_name);
                    this->m_primary_name_tree.add(primary_name);
                    this->update_lookup(key, unit);
                    this->on_loaded(unit);
                    return true;
                }
            } // add(...)
        }; // struct prefix_database
    } // namespace settlers_online
} // namespace ropufu

#endif // ROPUFU_SETTLERS_ONLINE_PREFIX_DATABASE_HPP_INCLUDED
<commit_msg>Now with gcc.<commit_after>
#ifndef ROPUFU_SETTLERS_ONLINE_PREFIX_DATABASE_HPP_INCLUDED
#define ROPUFU_SETTLERS_ONLINE_PREFIX_DATABASE_HPP_INCLUDED

#include <aftermath/not_an_error.hpp>
#include <experimental/filesystem>
#include <nlohmann/json.hpp>

#include "char_string.hpp"
#include "logger.hpp"
#include "prefix_tree.hpp"

#include <cstddef> // std::size_t
#include <fstream> // std::ifstream
#include <map> // std::map
#include <set> // std::set
#include <string> // std::string
#include <type_traits> // std::is_same
#include <vector> // std::vector

namespace ropufu
{
    namespace settlers_online
    {
        template <typename t_key_type, typename t_skeleton_type = t_key_type>
        struct lookup
        {
            using type = lookup<t_key_type, t_skeleton_type>;
            using key_type = t_key_type;
            using skeleton_type = t_skeleton_type;

        private:
            std::map<key_type, std::set<skeleton_type>> m_skeleton = { }; // For eack key, keeps track of all associated relaxed names.
            std::map<skeleton_type, std::set<key_type>> m_inverse_skeleton = { }; // For reach relaxed name, keeps track of all associated keys.

        public:
            void clear() noexcept
            {
                this->m_skeleton.clear();
                this->m_inverse_skeleton.clear();
            } // clear(...)

            void update(const key_type& key, const std::set<skeleton_type>& values) noexcept
            {
                // Record the values associated with the given key.
                this->m_skeleton.emplace(key, values);

                // Now the reverse: find all keys associated with values.
                for (const skeleton_type& weak : values)
                {
                    auto search = this->m_inverse_skeleton.find(weak);
                    if (search != this->m_inverse_skeleton.end()) search->second.insert(key);
                    else
                    {
                        std::set<key_type> matches { key }; // This is the only key that matches <weak>, since otherwise it would've been present in the inverse skeleton.
                        this->m_inverse_skeleton.emplace(weak, matches);
                    }
                }
            } // update(...)

            template <typename t_predicate_type>
            std::size_t try_find(const skeleton_type& query, key_type& key, const t_predicate_type& filter) const noexcept
            {
                auto search = this->m_inverse_skeleton.find(query);
                if (search != this->m_inverse_skeleton.end())
                {
                    const std::set<key_type>& matches = search->second;
                    std::size_t count = 0;
                    for (const key_type& maybe : matches)
                    {
                        if (filter(maybe)) { key = maybe; count++; }
                    }
                    return count;
                }
                return 0;
            } // try_find(...)
        }; // struct lookup

        /** @brief Auxiliary class for \c prefix_database specializations. */
        template <typename t_value_type>
        struct prefix_builder
        {
            using value_type = t_value_type;
            using key_type = std::string;

            static key_type build_key(const value_type& unit) noexcept;

            static key_type build_primary_name(const value_type& unit) noexcept;

            /** @brief Collection of names associated with the \p unit.
             *  @remark Return type has to support range-based loops.
             */
            static std::vector<key_type> names(const value_type& unit) noexcept;
        }; // struct prefix_builder

        /** @brief Class for accessing known units.
         *  @remark This is an abstract class (CRTP). For more information on CRTP see https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern
         *  @remark The class has to implement the following functions:
         *          static key_type build_key(const value_type& unit) const noexcept
         */
        template <typename t_derived_type, typename t_value_type>
        struct prefix_database
        {
            using type = prefix_database<t_derived_type, t_value_type>;
            using derived_type = t_derived_type;
            using value_type = t_value_type;
            using key_type = std::string;

        private:
            using builder_type = prefix_builder<value_type>;

            value_type m_invalid = { }; // Used to indicates invalid return result.
            std::map<key_type, value_type> m_database = { }; // Primary key: shortest name.
            std::set<key_type> m_primary_names = { }; // Keep track of primary names to allow for fast prefix search.
            char_tree m_primary_name_tree = { }; // Fast prefix search by primary name.
            lookup<key_type, key_type> m_lowercase_lookup = { }; // Lookup for lowercase names.
            lookup<key_type, key_type> m_misspelled_lookup = { }; // Lookup for misspelled names.

            /** Relaxed lookup stage 1. */
            static key_type relax_to_lowercase(const key_type& query) noexcept
            {
                key_type relaxed = query;
                char_string::to_lower(relaxed);
                return relaxed;
            } // relax_to_lowercase(...)

            /** Relaxed lookup stage 2. Assuming stage 1 has already been applied. */
            static key_type relax_spelling(const key_type& query) noexcept
            {
                key_type relaxed = query;
                // Plural: man -> men.
                char_string::replace(relaxed, "men", "man");
                // Plural: ...es or ...s.
                if (char_string::ends_with(relaxed, "es")) relaxed = relaxed.substr(0, relaxed.length() - 2);
                else if (char_string::ends_with(relaxed, "s")) relaxed = relaxed.substr(0, relaxed.length() - 1);

                // Collapse all repeated letters for "reasonably" long words.
                if (relaxed.length() > 4)
                {
                    char previous = relaxed[0];
                    for (std::size_t i = 1; i < relaxed.size(); ++i)
                    {
                        if (relaxed[i] == previous) relaxed.erase(i, 1);
                        else previous = relaxed[i];
                    }
                }

                return relaxed;
            } // relax_spelling(...)

            /** Assuming the unit names have already been subject to \c deep_trim_copy. */
            void update_lookup(const key_type& key, const value_type& unit) noexcept
            {
                std::set<key_type> relaxed_names { };
                std::set<key_type> misspelled_names { };
                for (const key_type& name : builder_type::names(unit))
                {
                    key_type stage1 = type::relax_to_lowercase(name);
                    key_type stage2 = type::relax_spelling(stage1);
                    relaxed_names.insert(stage1);
                    misspelled_names.insert(stage2);
                }

                this->m_lowercase_lookup.update(key, relaxed_names);
                this->m_misspelled_lookup.update(key, misspelled_names);
            } // update_lookup(...)

        protected:
            /** @brief Happens when \c clear() is called. */
            void on_clear() noexcept
            {
                constexpr bool is_overwritten = std::is_same<
                    decltype(&derived_type::on_clear),
                    decltype(&type::on_clear)>::value;
                if (!is_overwritten) return; // Do nothing if the function has not been overridden.

                derived_type* that = static_cast<derived_type*>(this);
                that->on_clear();
            } // on_clear(...)

            /** @brief Happens when a unit is about to be loaded. */
            template <typename t_logger_type>
            void on_loading(const value_type& unit, bool& do_cancel, t_logger_type& logger) const noexcept
            {
                do_cancel = false;
                constexpr bool is_overwritten = std::is_same<
                    decltype(&derived_type::template on_loading<t_logger_type>),
                    decltype(&type::template on_loading<t_logger_type>)>::value;
                if (!is_overwritten) return; // Do nothing if the function has not been overridden.

                const derived_type* that = static_cast<const derived_type*>(this);
                that->on_loading(unit, do_cancel, logger);
            } // on_loading(...)

            /** @brief Happens when a unit has been loaded. */
            void on_loaded(const value_type& unit) noexcept
            {
                constexpr bool is_overwritten = std::is_same<
                    decltype(&derived_type::on_loaded),
                    decltype(&type::on_loaded)>::value;
                if (!is_overwritten) return; // Do nothing if the function has not been overridden.

                derived_type* that = static_cast<derived_type*>(this);
                that->on_loaded(unit);
            } // on_loaded(...)

        public:
            prefix_database() noexcept { }

            /** Clears the contents of the database. */
            void clear() noexcept
            {
                this->m_database.clear();
                this->m_primary_names.clear();
                this->m_primary_name_tree.clear();
                this->m_lowercase_lookup.clear();
                this->m_misspelled_lookup.clear();
                this->on_clear();
            } // clear(...)

            const std::map<key_type, value_type>& data() const noexcept { return this->m_database; }

            /** @brief Access elements by key.
             *  @exception not_an_error::out_of_range This error is pushed to \c quiet_error if specified \p key is not in the database.
             */
            const value_type& at(const key_type& key) const noexcept
            {
                auto search = this->m_database.find(key);
                if (search != this->m_database.end()) return search->second;

                aftermath::quiet_error::instance().push(
                    aftermath::not_an_error::out_of_range,
                    aftermath::severity_level::fatal,
                    "<key> is not present in the database.", __FUNCTION__, __LINE__);
                return this->m_invalid;
            } // at(...)

            template <typename t_predicate_type>
            bool try_find(const key_type& query, value_type& unit, const t_predicate_type& filter) const noexcept
            {
                detail::no_logger logger { };
                return this->try_find(query, unit, filter, logger);
            } // try_find(...)

            template <typename t_predicate_type, typename t_logger_type>
            bool try_find(const key_type& query, value_type& unit, const t_predicate_type& filter, t_logger_type& logger) const noexcept
            {
                auto search = this->m_database.find(query);
                if (search != this->m_database.end())
                {
                    unit = search->second;
                    if (filter(unit)) return true;
                }

                // Primary search failed. Secondary search: all lowercase!
                key_type key { };
                key_type lowercase = type::relax_to_lowercase(query);
                key_type misspelled = type::relax_spelling(lowercase);
                std::size_t count_matches = 0;

                // Stage 0: prefix tree search.
                bool is_single = false;
                key_type first_prefix_match = this->m_primary_name_tree.first(lowercase, is_single);
                //count_matches = this->m_primary_name_tree.count(lowercase);
                // Only one terminus matches the prefix.
                if (is_single) lowercase = first_prefix_match; // Overwrite the search query.

                // Stage 1: lowercase lookup.
                count_matches = this->m_lowercase_lookup.try_find(lowercase, key, [&] (const key_type& maybe) { return filter(this->m_database.at(maybe)); });
                if (count_matches >= 1)
                {
                    unit = this->m_database.at(key);
                    if (count_matches == 1) return true;
                    logger.write(std::string("Multiple units match the specified query: ") + lowercase + std::string("."));
                    return false;
                }

                // Stage 2: misspelled lookup.
                count_matches = this->m_misspelled_lookup.try_find(misspelled, key, [&] (const key_type& maybe) { return filter(this->m_database.at(maybe)); });
                if (count_matches >= 1)
                {
                    unit = this->m_database.at(key);
                    if (count_matches == 1) return true;
                    logger.write(std::string("Multiple units match the specified query: ") + misspelled + std::string("."));
                    return false;
                }
                return false;
            } // try_find(...)

            /** Load units from .json files in a specified folder. */
            template <typename t_logger_type>
            bool add(const value_type& unit, t_logger_type& logger) noexcept
            {
                bool do_cancel = false;
                this->on_loading(unit, do_cancel, logger);
                if (do_cancel) return false;

                key_type key = builder_type::build_key(unit);
                key_type primary_name = builder_type::build_primary_name(unit);

                if (this->m_database.count(key) != 0)
                {
                    logger.write(std::string("Unit with name ") + key + std::string(" already exists."));
                    return false;
                }
                else if (this->m_primary_names.count(primary_name) != 0)
                {
                    logger.write(std::string("Unit with a similar primary name ") + primary_name + std::string(" already exists."));
                    return false;
                }
                else
                {
                    this->m_database.emplace(key, unit);
                    this->m_primary_names.insert(primary_name);
                    this->m_primary_name_tree.add(primary_name);
                    this->update_lookup(key, unit);
                    this->on_loaded(unit);
                    return true;
                }
            } // add(...)
        }; // struct prefix_database
    } // namespace settlers_online
} // namespace ropufu

#endif // ROPUFU_SETTLERS_ONLINE_PREFIX_DATABASE_HPP_INCLUDED
<|endoftext|>


<commit_before>/*************************************************************************
 *
 *  OpenOffice.org - a multi-platform office productivity suite
 *
 *  $RCSfile: sfxurlrelocator.hxx,v $
 *
 *  $Revision: 1.3 $
 *
 *  last change: $Author: rt $ $Date: 2005-09-07 19:11:35 $
 *
 *  The Contents of this file are made available subject to
 *  the terms of GNU Lesser General Public License Version 2.1.
 *
 *
 *    GNU Lesser General Public License Version 2.1
 *    =============================================
 *    Copyright 2005 by Sun Microsystems, Inc.
 *    901 San Antonio Road, Palo Alto, CA 94303, USA
 *
 *    This library is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU Lesser General Public
 *    License version 2.1, as published by the Free Software Foundation.
 *
 *    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 GNU
 *    Lesser General Public License for more details.
 *
 *    You should have received a copy of the GNU Lesser General Public
 *    License along with this library; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 *    MA  02111-1307  USA
 *
 ************************************************************************/

#ifndef _SFX_SFXURLRELOCATOR_HXX_
#define _SFX_SFXURLRELOCATOR_HXX_

#ifndef _COM_SUN_STAR_LANG_XMULTISERVICEFACTORY_HPP_
#include <com/sun/star/lang/XMultiServiceFactory.hpp>
#endif

#ifndef  _COM_SUN_STAR_UTIL_XOFFICEINSTALLATIONDIRECTORIES_HPP_
#include <com/sun/star/util/XOfficeInstallationDirectories.hpp>
#endif

#include <rtl/ustring.hxx>
#include <osl/mutex.hxx>

class SfxURLRelocator_Impl
{
    ::osl::Mutex maMutex;
    ::com::sun::star::uno::Reference< ::com::sun::star::lang::XMultiServiceFactory >           mxFactory;
    ::com::sun::star::uno::Reference< ::com::sun::star::util::XOfficeInstallationDirectories > mxOfficeInstDirs;

public:
    static bool                 propertyCanContainOfficeDir( const rtl::OUString & rPropName );
    void                        initOfficeInstDirs();
    void                        makeRelocatableURL( rtl::OUString & rURL );
    void                        makeAbsoluteURL( rtl::OUString & rURL );

    SfxURLRelocator_Impl( ::com::sun::star::uno::Reference< ::com::sun::star::lang::XMultiServiceFactory > xFactory );
    ~SfxURLRelocator_Impl();
};

#endif

<commit_msg>INTEGRATION: CWS changefileheader (1.3.732); FILE MERGED 2008/04/01 12:40:51 thb 1.3.732.2: #i85898# Stripping all external header guards 2008/03/31 13:38:35 rt 1.3.732.1: #i87441# Change license header to LPGL v3.<commit_after>/*************************************************************************
 *
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * Copyright 2008 by Sun Microsystems, Inc.
 *
 * OpenOffice.org - a multi-platform office productivity suite
 *
 * $RCSfile: sfxurlrelocator.hxx,v $
 * $Revision: 1.4 $
 *
 * This file is part of OpenOffice.org.
 *
 * OpenOffice.org is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License version 3
 * only, as published by the Free Software Foundation.
 *
 * OpenOffice.org 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 Lesser General Public License version 3 for more details
 * (a copy is included in the LICENSE file that accompanied this code).
 *
 * You should have received a copy of the GNU Lesser General Public License
 * version 3 along with OpenOffice.org.  If not, see
 * <http://www.openoffice.org/license.html>
 * for a copy of the LGPLv3 License.
 *
 ************************************************************************/

#ifndef _SFX_SFXURLRELOCATOR_HXX_
#define _SFX_SFXURLRELOCATOR_HXX_

#include <com/sun/star/lang/XMultiServiceFactory.hpp>
#include <com/sun/star/util/XOfficeInstallationDirectories.hpp>

#include <rtl/ustring.hxx>
#include <osl/mutex.hxx>

class SfxURLRelocator_Impl
{
    ::osl::Mutex maMutex;
    ::com::sun::star::uno::Reference< ::com::sun::star::lang::XMultiServiceFactory >           mxFactory;
    ::com::sun::star::uno::Reference< ::com::sun::star::util::XOfficeInstallationDirectories > mxOfficeInstDirs;

public:
    static bool                 propertyCanContainOfficeDir( const rtl::OUString & rPropName );
    void                        initOfficeInstDirs();
    void                        makeRelocatableURL( rtl::OUString & rURL );
    void                        makeAbsoluteURL( rtl::OUString & rURL );

    SfxURLRelocator_Impl( ::com::sun::star::uno::Reference< ::com::sun::star::lang::XMultiServiceFactory > xFactory );
    ~SfxURLRelocator_Impl();
};

#endif

<|endoftext|>


<commit_before>/// Copyright 2015 Google Inc. All rights reserved.
///
/// 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 "SantaDecisionManager.h"

#define super OSObject
OSDefineMetaClassAndStructors(SantaDecisionManager, OSObject);

#pragma mark Object Lifecycle

bool SantaDecisionManager::init() {
  sdm_lock_grp_ = lck_grp_alloc_init("santa-locks", lck_grp_attr_alloc_init());
  dataqueue_lock_ = lck_mtx_alloc_init(sdm_lock_grp_, lck_attr_alloc_init());
  cached_decisions_lock_ = lck_rw_alloc_init(sdm_lock_grp_, lck_attr_alloc_init());

  cached_decisions_ = OSDictionary::withCapacity(1000);

  return kIOReturnSuccess;
}

void SantaDecisionManager::free() {
  if (cached_decisions_) {
    cached_decisions_->release();
    cached_decisions_ = NULL;
  }

  if (cached_decisions_lock_) {
    lck_rw_free(cached_decisions_lock_, sdm_lock_grp_);
    cached_decisions_lock_ = NULL;
  }

  if (dataqueue_lock_ ) {
    lck_mtx_free(dataqueue_lock_, sdm_lock_grp_);
    dataqueue_lock_ = NULL;
  }

  if (sdm_lock_grp_) {
    lck_grp_free(sdm_lock_grp_);
    sdm_lock_grp_ = NULL;
  }

  super::free();
}

#pragma mark Client Management

void SantaDecisionManager::ConnectClient(IOSharedDataQueue *queue, pid_t pid) {
  if (!pid) return;
  if (!queue) return;

  // Any decisions made while the daemon wasn't
  // connected should be cleared
  ClearCache();

  lck_mtx_lock(dataqueue_lock_);
  dataqueue_ = queue;
  dataqueue_->retain();
  lck_mtx_unlock(dataqueue_lock_);

  owning_pid_ = pid;
  owning_proc_ = proc_find(pid);
}

void SantaDecisionManager::DisconnectClient() {
  owning_pid_ = -1;

  // Ask santad to shutdown, in case it's running.
  santa_message_t message;
  message.action = ACTION_REQUEST_SHUTDOWN;
  message.userId = 0;
  message.pid = 0;
  message.ppid = 0;
  message.vnode_id = 0;
  PostToQueue(message);

  lck_mtx_lock(dataqueue_lock_);
  dataqueue_->release();
  dataqueue_ = NULL;
  lck_mtx_unlock(dataqueue_lock_);

  proc_rele(owning_proc_);
  owning_proc_ = NULL;
}

bool SantaDecisionManager::ClientConnected() {
  return owning_pid_ > 0;
}

# pragma mark Listener Control

kern_return_t SantaDecisionManager::StartListener() {
  vnode_listener_ = kauth_listen_scope(KAUTH_SCOPE_VNODE,
                                       vnode_scope_callback,
                                       reinterpret_cast<void *>(this));
  if (!vnode_listener_) return kIOReturnInternalError;

  LOGD("Vnode listener started.");

  return kIOReturnSuccess;
}

kern_return_t SantaDecisionManager::StopListener() {
  kauth_unlisten_scope(vnode_listener_);
  vnode_listener_ = NULL;

  // Wait for any active invocations to finish before returning
  do {
    IOSleep(5);
  } while (listener_invocations_);

  // Delete any cached decisions
  ClearCache();

  LOGD("Vnode listener stopped.");

  return kIOReturnSuccess;
}

#pragma mark Cache Management

void SantaDecisionManager::AddToCache(
    const char *identifier, santa_action_t decision, uint64_t microsecs) {
  lck_rw_lock_exclusive(cached_decisions_lock_);

  if (cached_decisions_->getCount() > kMaxCacheSize) {
    // This could be made a _lot_ smarter, say only removing entries older
    // than a certain time period. However, with a kMaxCacheSize set
    // sufficiently large and a kMaxAllowCacheTimeMilliseconds set
    // sufficiently low, this should only ever occur if someone is purposefully
    // trying to make the cache grow.
    LOGD("Cache too large, flushing.");
    cached_decisions_->flushCollection();
  }

  if (decision == ACTION_REQUEST_CHECKBW) {
    SantaMessage *pending = new SantaMessage();
    pending->setAction(ACTION_REQUEST_CHECKBW, 0);
    cached_decisions_->setObject(identifier, pending);
    pending->release();  // it was retained when added to the dictionary
  } else {
    SantaMessage *pending = OSDynamicCast(
        SantaMessage, cached_decisions_->getObject(identifier));
    if (pending) {
      pending->setAction(decision, microsecs);
    }
  }

  lck_rw_unlock_exclusive(cached_decisions_lock_);
}

void SantaDecisionManager::CacheCheck(const char *identifier) {
  lck_rw_lock_shared(cached_decisions_lock_);
  bool shouldInvalidate = (cached_decisions_->getObject(identifier) != NULL);
  if (shouldInvalidate) {
    lck_rw_lock_shared_to_exclusive(cached_decisions_lock_);
    cached_decisions_->removeObject(identifier);
    lck_rw_unlock_exclusive(cached_decisions_lock_);
  } else {
    lck_rw_unlock_shared(cached_decisions_lock_);
  }
}

uint64_t SantaDecisionManager::CacheCount() {
  return cached_decisions_->getCount();
}

void SantaDecisionManager::ClearCache() {
  lck_rw_lock_exclusive(cached_decisions_lock_);
  cached_decisions_->flushCollection();
  lck_rw_unlock_exclusive(cached_decisions_lock_);
}

santa_action_t SantaDecisionManager::GetFromCache(const char *identifier) {
  santa_action_t result = ACTION_UNSET;
  uint64_t decision_time = 0;

  lck_rw_lock_shared(cached_decisions_lock_);
  SantaMessage *cached_decision = OSDynamicCast(
      SantaMessage, cached_decisions_->getObject(identifier));
  if (cached_decision) {
    result = cached_decision->getAction();
    decision_time = cached_decision->getMicrosecs();
  }
  lck_rw_unlock_shared(cached_decisions_lock_);

  if (CHECKBW_RESPONSE_VALID(result)) {
    uint64_t diff_time = GetCurrentUptime();

    if (result == ACTION_RESPOND_CHECKBW_ALLOW) {
      if ((kMaxAllowCacheTimeMilliseconds * 1000) > diff_time) {
        diff_time = 0;
      } else {
        diff_time -= (kMaxAllowCacheTimeMilliseconds * 1000);
      }
    } else if (result == ACTION_RESPOND_CHECKBW_DENY) {
      if ((kMaxDenyCacheTimeMilliseconds * 1000) > diff_time) {
        diff_time = 0;
      } else {
        diff_time -= (kMaxDenyCacheTimeMilliseconds * 1000);
      }
    }

    if (decision_time < diff_time) {
      lck_rw_lock_exclusive(cached_decisions_lock_);
      cached_decisions_->removeObject(identifier);
      lck_rw_unlock_exclusive(cached_decisions_lock_);
      return ACTION_UNSET;
    }
  }

  return result;
}

# pragma mark Queue Management

bool SantaDecisionManager::PostToQueue(santa_message_t message) {
  lck_mtx_lock(dataqueue_lock_);
  bool kr = false;
  if (dataqueue_) {
    kr = dataqueue_->enqueue(&message, sizeof(message));
  }
  lck_mtx_unlock(dataqueue_lock_);
  return kr;
}

santa_action_t SantaDecisionManager::FetchDecision(
    const kauth_cred_t credential,
    const vfs_context_t vfs_context,
    const vnode_t vnode) {
  santa_action_t return_action = ACTION_UNSET;

  // Fetch Vnode ID & string
  uint64_t vnode_id = GetVnodeIDForVnode(vfs_context, vnode);
  char vnode_id_str[MAX_VNODE_ID_STR];
  snprintf(vnode_id_str, MAX_VNODE_ID_STR, "%llu", vnode_id);

  // Check to see if item is in cache
  return_action = GetFromCache(vnode_id_str);

  // If item wasn't in cache, fetch decision from daemon.
  if (!CHECKBW_RESPONSE_VALID(return_action)) {
    // Add pending request to cache
    AddToCache(vnode_id_str, ACTION_REQUEST_CHECKBW, 0);

    // Get path
    char path[MAX_PATH_LEN];
    int name_len = MAX_PATH_LEN;
    if (vn_getpath(vnode, path, &name_len) != 0) {
      path[0] = '\0';
    }

    if (!ClientConnected()) {
      LOGI("Execution request without daemon running: %s", path);
      AddToCache(vnode_id_str,
                 ACTION_RESPOND_CHECKBW_ALLOW,
                 GetCurrentUptime());
      return ACTION_RESPOND_CHECKBW_ALLOW;
    }

    // Prepare to send message to daemon
    santa_message_t message;
    strncpy(message.path, path, MAX_PATH_LEN);
    message.userId = kauth_cred_getuid(credential);
    message.pid = proc_selfpid();
    message.ppid = proc_selfppid();
    message.action = ACTION_REQUEST_CHECKBW;
    message.vnode_id = vnode_id;

    // Wait for the daemon to respond or die.
    do {
      // Send request to daemon...
      if (!PostToQueue(message)) {
        LOGE("Failed to queue request for %s.", path);
        CacheCheck(vnode_id_str);
        return ACTION_ERROR;
      }

      // ... and wait for it to respond. If after kRequestLoopSleepMilliseconds
      // * kMaxRequestLoops it still hasn't responded, send request again.
      for (int i = 0; i < kMaxRequestLoops; ++i) {
        IOSleep(kRequestLoopSleepMilliseconds);
        return_action = GetFromCache(vnode_id_str);
        if (CHECKBW_RESPONSE_VALID(return_action)) break;
      }
    } while (!CHECKBW_RESPONSE_VALID(return_action) &&
             proc_exiting(owning_proc_) == 0);

    // If response is still not valid, the daemon exited
    if (!CHECKBW_RESPONSE_VALID(return_action)) {
      LOGE("Daemon process did not respond correctly. Allowing executions "
           "until it comes back.");
      CacheCheck(vnode_id_str);
      return ACTION_ERROR;
    }
  }

  return return_action;
}

# pragma mark Misc

uint64_t SantaDecisionManager::GetVnodeIDForVnode(const vfs_context_t context,
                                                  const vnode_t vp) {
  struct vnode_attr vap;
  VATTR_INIT(&vap);
  VATTR_WANTED(&vap, va_fileid);
  vnode_getattr(vp, &vap, context);
  return vap.va_fileid;
}

uint64_t SantaDecisionManager::GetCurrentUptime() {
  clock_sec_t sec;
  clock_usec_t usec;
  clock_get_system_microtime(&sec, &usec);
  return (uint64_t)((sec * 1000000) + usec);
}

# pragma mark Invocation Tracking & PID comparison

void SantaDecisionManager::IncrementListenerInvocations() {
  OSIncrementAtomic(&listener_invocations_);
}

void SantaDecisionManager::DecrementListenerInvocations() {
  OSDecrementAtomic(&listener_invocations_);
}

#undef super

#pragma mark Kauth Callback

extern "C" int vnode_scope_callback(
    kauth_cred_t credential, void *idata, kauth_action_t action,
    uintptr_t arg0, uintptr_t arg1, uintptr_t arg2, uintptr_t arg3) {
  // The default action is to defer
  int returnResult = KAUTH_RESULT_DEFER;

  // Cast arguments to correct types
  if (idata == NULL) {
    LOGE("Vnode callback established without valid decision manager.");
    return returnResult;
  }
  SantaDecisionManager *sdm = OSDynamicCast(
      SantaDecisionManager, reinterpret_cast<OSObject *>(idata));
  vfs_context_t vfs_context = reinterpret_cast<vfs_context_t>(arg0);
  vnode_t vnode = reinterpret_cast<vnode_t>(arg1);

  // Only operate on regular files (not directories, symlinks, etc.)
  vtype vt = vnode_vtype(vnode);
  if (vt != VREG) return returnResult;

  // Don't operate on ACCESS events, as they're advisory
  if (action & KAUTH_VNODE_ACCESS) return returnResult;

  // Filter for only writes
  if (action & KAUTH_VNODE_WRITE_DATA ||
      action & KAUTH_VNODE_APPEND_DATA ||
      action & KAUTH_VNODE_DELETE) {
    char vnode_id_str[MAX_VNODE_ID_STR];
    snprintf(vnode_id_str, MAX_VNODE_ID_STR, "%llu",
             sdm->GetVnodeIDForVnode(vfs_context, vnode));

    // If an execution request is pending, deny write
    if (sdm->GetFromCache(vnode_id_str) == ACTION_REQUEST_CHECKBW) {
      LOGD("Denying write due to pending execution: %s", vnode_id_str);
      *(reinterpret_cast<int *>(arg3)) = EACCES;
      return KAUTH_RESULT_DENY;
    }

    // Otherwise remove from cache
    sdm->CacheCheck(vnode_id_str);

    return returnResult;
  }

  // Filter for only EXECUTE actions
  if (action & KAUTH_VNODE_EXECUTE) {
    sdm->IncrementListenerInvocations();

    // Fetch decision
    santa_action_t returnedAction = sdm->FetchDecision(
        credential, vfs_context, vnode);

    switch (returnedAction) {
      case ACTION_RESPOND_CHECKBW_ALLOW:
        returnResult = KAUTH_RESULT_ALLOW;
        break;
      case ACTION_RESPOND_CHECKBW_DENY:
        *(reinterpret_cast<int *>(arg3)) = EACCES;
        returnResult = KAUTH_RESULT_DENY;
        break;
      default:
        // NOTE: Any unknown response or error condition causes us to fail open.
        // Whilst from a security perspective this is bad, it's important that
        // we don't break user's machines.
        break;
    }

    sdm->DecrementListenerInvocations();

    return returnResult;
  }

  return returnResult;
}
<commit_msg>santa-driver: DisconnectClient should return straight away if client is not properly connected<commit_after>/// Copyright 2015 Google Inc. All rights reserved.
///
/// 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 "SantaDecisionManager.h"

#define super OSObject
OSDefineMetaClassAndStructors(SantaDecisionManager, OSObject);

#pragma mark Object Lifecycle

bool SantaDecisionManager::init() {
  sdm_lock_grp_ = lck_grp_alloc_init("santa-locks", lck_grp_attr_alloc_init());
  dataqueue_lock_ = lck_mtx_alloc_init(sdm_lock_grp_, lck_attr_alloc_init());
  cached_decisions_lock_ = lck_rw_alloc_init(sdm_lock_grp_, lck_attr_alloc_init());

  cached_decisions_ = OSDictionary::withCapacity(1000);

  return kIOReturnSuccess;
}

void SantaDecisionManager::free() {
  if (cached_decisions_) {
    cached_decisions_->release();
    cached_decisions_ = NULL;
  }

  if (cached_decisions_lock_) {
    lck_rw_free(cached_decisions_lock_, sdm_lock_grp_);
    cached_decisions_lock_ = NULL;
  }

  if (dataqueue_lock_ ) {
    lck_mtx_free(dataqueue_lock_, sdm_lock_grp_);
    dataqueue_lock_ = NULL;
  }

  if (sdm_lock_grp_) {
    lck_grp_free(sdm_lock_grp_);
    sdm_lock_grp_ = NULL;
  }

  super::free();
}

#pragma mark Client Management

void SantaDecisionManager::ConnectClient(IOSharedDataQueue *queue, pid_t pid) {
  if (!pid) return;
  if (!queue) return;

  // Any decisions made while the daemon wasn't
  // connected should be cleared
  ClearCache();

  lck_mtx_lock(dataqueue_lock_);
  dataqueue_ = queue;
  dataqueue_->retain();
  lck_mtx_unlock(dataqueue_lock_);

  owning_pid_ = pid;
  owning_proc_ = proc_find(pid);
}

void SantaDecisionManager::DisconnectClient() {
  if (owning_pid_ < 1) return;

  owning_pid_ = -1;

  // Ask santad to shutdown, in case it's running.
  santa_message_t message;
  message.action = ACTION_REQUEST_SHUTDOWN;
  message.userId = 0;
  message.pid = 0;
  message.ppid = 0;
  message.vnode_id = 0;
  PostToQueue(message);

  lck_mtx_lock(dataqueue_lock_);
  dataqueue_->release();
  dataqueue_ = NULL;
  lck_mtx_unlock(dataqueue_lock_);

  proc_rele(owning_proc_);
  owning_proc_ = NULL;
}

bool SantaDecisionManager::ClientConnected() {
  return owning_pid_ > 0;
}

# pragma mark Listener Control

kern_return_t SantaDecisionManager::StartListener() {
  vnode_listener_ = kauth_listen_scope(KAUTH_SCOPE_VNODE,
                                       vnode_scope_callback,
                                       reinterpret_cast<void *>(this));
  if (!vnode_listener_) return kIOReturnInternalError;

  LOGD("Vnode listener started.");

  return kIOReturnSuccess;
}

kern_return_t SantaDecisionManager::StopListener() {
  kauth_unlisten_scope(vnode_listener_);
  vnode_listener_ = NULL;

  // Wait for any active invocations to finish before returning
  do {
    IOSleep(5);
  } while (listener_invocations_);

  // Delete any cached decisions
  ClearCache();

  LOGD("Vnode listener stopped.");

  return kIOReturnSuccess;
}

#pragma mark Cache Management

void SantaDecisionManager::AddToCache(
    const char *identifier, santa_action_t decision, uint64_t microsecs) {
  lck_rw_lock_exclusive(cached_decisions_lock_);

  if (cached_decisions_->getCount() > kMaxCacheSize) {
    // This could be made a _lot_ smarter, say only removing entries older
    // than a certain time period. However, with a kMaxCacheSize set
    // sufficiently large and a kMaxAllowCacheTimeMilliseconds set
    // sufficiently low, this should only ever occur if someone is purposefully
    // trying to make the cache grow.
    LOGD("Cache too large, flushing.");
    cached_decisions_->flushCollection();
  }

  if (decision == ACTION_REQUEST_CHECKBW) {
    SantaMessage *pending = new SantaMessage();
    pending->setAction(ACTION_REQUEST_CHECKBW, 0);
    cached_decisions_->setObject(identifier, pending);
    pending->release();  // it was retained when added to the dictionary
  } else {
    SantaMessage *pending = OSDynamicCast(
        SantaMessage, cached_decisions_->getObject(identifier));
    if (pending) {
      pending->setAction(decision, microsecs);
    }
  }

  lck_rw_unlock_exclusive(cached_decisions_lock_);
}

void SantaDecisionManager::CacheCheck(const char *identifier) {
  lck_rw_lock_shared(cached_decisions_lock_);
  bool shouldInvalidate = (cached_decisions_->getObject(identifier) != NULL);
  if (shouldInvalidate) {
    lck_rw_lock_shared_to_exclusive(cached_decisions_lock_);
    cached_decisions_->removeObject(identifier);
    lck_rw_unlock_exclusive(cached_decisions_lock_);
  } else {
    lck_rw_unlock_shared(cached_decisions_lock_);
  }
}

uint64_t SantaDecisionManager::CacheCount() {
  return cached_decisions_->getCount();
}

void SantaDecisionManager::ClearCache() {
  lck_rw_lock_exclusive(cached_decisions_lock_);
  cached_decisions_->flushCollection();
  lck_rw_unlock_exclusive(cached_decisions_lock_);
}

santa_action_t SantaDecisionManager::GetFromCache(const char *identifier) {
  santa_action_t result = ACTION_UNSET;
  uint64_t decision_time = 0;

  lck_rw_lock_shared(cached_decisions_lock_);
  SantaMessage *cached_decision = OSDynamicCast(
      SantaMessage, cached_decisions_->getObject(identifier));
  if (cached_decision) {
    result = cached_decision->getAction();
    decision_time = cached_decision->getMicrosecs();
  }
  lck_rw_unlock_shared(cached_decisions_lock_);

  if (CHECKBW_RESPONSE_VALID(result)) {
    uint64_t diff_time = GetCurrentUptime();

    if (result == ACTION_RESPOND_CHECKBW_ALLOW) {
      if ((kMaxAllowCacheTimeMilliseconds * 1000) > diff_time) {
        diff_time = 0;
      } else {
        diff_time -= (kMaxAllowCacheTimeMilliseconds * 1000);
      }
    } else if (result == ACTION_RESPOND_CHECKBW_DENY) {
      if ((kMaxDenyCacheTimeMilliseconds * 1000) > diff_time) {
        diff_time = 0;
      } else {
        diff_time -= (kMaxDenyCacheTimeMilliseconds * 1000);
      }
    }

    if (decision_time < diff_time) {
      lck_rw_lock_exclusive(cached_decisions_lock_);
      cached_decisions_->removeObject(identifier);
      lck_rw_unlock_exclusive(cached_decisions_lock_);
      return ACTION_UNSET;
    }
  }

  return result;
}

# pragma mark Queue Management

bool SantaDecisionManager::PostToQueue(santa_message_t message) {
  lck_mtx_lock(dataqueue_lock_);
  bool kr = false;
  if (dataqueue_) {
    kr = dataqueue_->enqueue(&message, sizeof(message));
  }
  lck_mtx_unlock(dataqueue_lock_);
  return kr;
}

santa_action_t SantaDecisionManager::FetchDecision(
    const kauth_cred_t credential,
    const vfs_context_t vfs_context,
    const vnode_t vnode) {
  santa_action_t return_action = ACTION_UNSET;

  // Fetch Vnode ID & string
  uint64_t vnode_id = GetVnodeIDForVnode(vfs_context, vnode);
  char vnode_id_str[MAX_VNODE_ID_STR];
  snprintf(vnode_id_str, MAX_VNODE_ID_STR, "%llu", vnode_id);

  // Check to see if item is in cache
  return_action = GetFromCache(vnode_id_str);

  // If item wasn't in cache, fetch decision from daemon.
  if (!CHECKBW_RESPONSE_VALID(return_action)) {
    // Add pending request to cache
    AddToCache(vnode_id_str, ACTION_REQUEST_CHECKBW, 0);

    // Get path
    char path[MAX_PATH_LEN];
    int name_len = MAX_PATH_LEN;
    if (vn_getpath(vnode, path, &name_len) != 0) {
      path[0] = '\0';
    }

    if (!ClientConnected()) {
      LOGI("Execution request without daemon running: %s", path);
      AddToCache(vnode_id_str,
                 ACTION_RESPOND_CHECKBW_ALLOW,
                 GetCurrentUptime());
      return ACTION_RESPOND_CHECKBW_ALLOW;
    }

    // Prepare to send message to daemon
    santa_message_t message;
    strncpy(message.path, path, MAX_PATH_LEN);
    message.userId = kauth_cred_getuid(credential);
    message.pid = proc_selfpid();
    message.ppid = proc_selfppid();
    message.action = ACTION_REQUEST_CHECKBW;
    message.vnode_id = vnode_id;

    // Wait for the daemon to respond or die.
    do {
      // Send request to daemon...
      if (!PostToQueue(message)) {
        LOGE("Failed to queue request for %s.", path);
        CacheCheck(vnode_id_str);
        return ACTION_ERROR;
      }

      // ... and wait for it to respond. If after kRequestLoopSleepMilliseconds
      // * kMaxRequestLoops it still hasn't responded, send request again.
      for (int i = 0; i < kMaxRequestLoops; ++i) {
        IOSleep(kRequestLoopSleepMilliseconds);
        return_action = GetFromCache(vnode_id_str);
        if (CHECKBW_RESPONSE_VALID(return_action)) break;
      }
    } while (!CHECKBW_RESPONSE_VALID(return_action) &&
             proc_exiting(owning_proc_) == 0);

    // If response is still not valid, the daemon exited
    if (!CHECKBW_RESPONSE_VALID(return_action)) {
      LOGE("Daemon process did not respond correctly. Allowing executions "
           "until it comes back.");
      CacheCheck(vnode_id_str);
      return ACTION_ERROR;
    }
  }

  return return_action;
}

# pragma mark Misc

uint64_t SantaDecisionManager::GetVnodeIDForVnode(const vfs_context_t context,
                                                  const vnode_t vp) {
  struct vnode_attr vap;
  VATTR_INIT(&vap);
  VATTR_WANTED(&vap, va_fileid);
  vnode_getattr(vp, &vap, context);
  return vap.va_fileid;
}

uint64_t SantaDecisionManager::GetCurrentUptime() {
  clock_sec_t sec;
  clock_usec_t usec;
  clock_get_system_microtime(&sec, &usec);
  return (uint64_t)((sec * 1000000) + usec);
}

# pragma mark Invocation Tracking & PID comparison

void SantaDecisionManager::IncrementListenerInvocations() {
  OSIncrementAtomic(&listener_invocations_);
}

void SantaDecisionManager::DecrementListenerInvocations() {
  OSDecrementAtomic(&listener_invocations_);
}

#undef super

#pragma mark Kauth Callback

extern "C" int vnode_scope_callback(
    kauth_cred_t credential, void *idata, kauth_action_t action,
    uintptr_t arg0, uintptr_t arg1, uintptr_t arg2, uintptr_t arg3) {
  // The default action is to defer
  int returnResult = KAUTH_RESULT_DEFER;

  // Cast arguments to correct types
  if (idata == NULL) {
    LOGE("Vnode callback established without valid decision manager.");
    return returnResult;
  }
  SantaDecisionManager *sdm = OSDynamicCast(
      SantaDecisionManager, reinterpret_cast<OSObject *>(idata));
  vfs_context_t vfs_context = reinterpret_cast<vfs_context_t>(arg0);
  vnode_t vnode = reinterpret_cast<vnode_t>(arg1);

  // Only operate on regular files (not directories, symlinks, etc.)
  vtype vt = vnode_vtype(vnode);
  if (vt != VREG) return returnResult;

  // Don't operate on ACCESS events, as they're advisory
  if (action & KAUTH_VNODE_ACCESS) return returnResult;

  // Filter for only writes
  if (action & KAUTH_VNODE_WRITE_DATA ||
      action & KAUTH_VNODE_APPEND_DATA ||
      action & KAUTH_VNODE_DELETE) {
    char vnode_id_str[MAX_VNODE_ID_STR];
    snprintf(vnode_id_str, MAX_VNODE_ID_STR, "%llu",
             sdm->GetVnodeIDForVnode(vfs_context, vnode));

    // If an execution request is pending, deny write
    if (sdm->GetFromCache(vnode_id_str) == ACTION_REQUEST_CHECKBW) {
      LOGD("Denying write due to pending execution: %s", vnode_id_str);
      *(reinterpret_cast<int *>(arg3)) = EACCES;
      return KAUTH_RESULT_DENY;
    }

    // Otherwise remove from cache
    sdm->CacheCheck(vnode_id_str);

    return returnResult;
  }

  // Filter for only EXECUTE actions
  if (action & KAUTH_VNODE_EXECUTE) {
    sdm->IncrementListenerInvocations();

    // Fetch decision
    santa_action_t returnedAction = sdm->FetchDecision(
        credential, vfs_context, vnode);

    switch (returnedAction) {
      case ACTION_RESPOND_CHECKBW_ALLOW:
        returnResult = KAUTH_RESULT_ALLOW;
        break;
      case ACTION_RESPOND_CHECKBW_DENY:
        *(reinterpret_cast<int *>(arg3)) = EACCES;
        returnResult = KAUTH_RESULT_DENY;
        break;
      default:
        // NOTE: Any unknown response or error condition causes us to fail open.
        // Whilst from a security perspective this is bad, it's important that
        // we don't break user's machines.
        break;
    }

    sdm->DecrementListenerInvocations();

    return returnResult;
  }

  return returnResult;
}
<|endoftext|>