SNF4CGP/SNF4CGP/JobPool.cpp

// SNF4CGP/JobPool.cpp
// Copyright (C) 2009 ARM Research Labs, LLC.
// See www.armresearch.com for more information.

#include "JobPool.hpp"
#include "Command.hpp"
#include "ScannerPool.hpp"
#include "OutputProcessor.hpp"

#include "../SNFMulti/SNFMulti.hpp"

#include "../CodeDweller/threading.hpp"
#include "../CodeDweller/faults.hpp"

#include <iostream>
#include <string>
#include <vector>

using namespace std;

//// Tuning Constants //////////////////////////////////////////////////////////

const int ReadBufferSize = snf_ScanHorizon;
const int StringReserveSize = 2048;

//// Job ///////////////////////////////////////////////////////////////////////

void Job::emitComment(const string& Comment) {
    ostringstream O;
    O << "* " << Comment << endl;
    OutputBuffer.append(O.str());
}

void Job::generateCommandCommentPrefix(ostringstream& O) {
    O << "* SNF4CGP[" << CurrentCommand.Number << "] ";
}

void Job::emitOK() {
    ostringstream O;
    O << CurrentCommand.Number << " OK" << endl;
    OutputBuffer.append(O.str());
}

const string INTFVersion = "1";

void Job::emitINTF() {
    ostringstream O;
    O << CurrentCommand.Number << " INTF " << INTFVersion << endl;
    OutputBuffer.append(O.str());
}

void Job::emitFAILURE() {
    ostringstream O;
    O << CurrentCommand.Number << " FAILURE" << endl;
    OutputBuffer.append(O.str());
}

string formatAsCGPString(const string& S) {
    string CGPString;
    CGPString.push_back('"');
    for(size_t i = 0; i < S.length(); i++) {
        char C = S[i];
        switch(C) {
            case '\r': break;
            case '\t': { CGPString.append("\\t"); break; }
            case '\n': { CGPString.append("\\e"); break; }
            case '\\': { CGPString.append("\\"); break; }
            case '"': { CGPString.append("\\\""); break; }
            default: CGPString.push_back(C);
        }
    }
    CGPString.push_back('"');
    return CGPString;
}

void Job::emitADDHEADER(const string& Headers) {
    ostringstream O;
    O << CurrentCommand.Number << " ADDHEADER " << formatAsCGPString(Headers) << endl;
    OutputBuffer.append(O.str());
}

void Job::emitERROR(const string& Report) {
    ostringstream O;
    O << CurrentCommand.Number << " ERROR " << formatAsCGPString(Report) << endl;
    OutputBuffer.append(O.str());
}

void Job::emitDISCARD() {
    ostringstream O;
    O << CurrentCommand.Number << " DISCARD" << endl;
    OutputBuffer.append(O.str());
}

void Job::emitREJECTED(const string& Report) {
    ostringstream O;
    O << CurrentCommand.Number << " REJECTED " << formatAsCGPString(Report) << endl;
    OutputBuffer.append(O.str());
}

void Job::finalize() {                                                          // Cleanup and report this job.
    closeWriter();                                                              // In case of exception let go of our
    closeReader();                                                              // Reader and Writer now. The rest
    Output.outputJob(*this);                                                    // can wait for the output processor.
}

void Job::doWakeUp() {
    emitComment("SNF4CGP Waking Up");
    emitComment(CurrentCommand.Data);
}

void Job::doINTF() {
    emitINTF();
}

void Job::doUNKNOWN() {
    ostringstream O;
    O << "* SNF4CGP[" << CurrentCommand.Number << "] Does not understand: "
      << formatAsCGPString(CurrentCommand.Data) << endl;
    OutputBuffer.append(O.str());
    emitFAILURE();
}

void Job::doFILE() {
    doInitialRead();
    doScan();
    doAction();
    closeReader();
}

RuntimeCheck CheckMessageReaderIsValid("Job::Reader() Check(0 != Reader_)");

ifstream& Job::Reader() {                                                       // Safe access to Reader_.
    CheckMessageReaderIsValid(0 != Reader_);
    return (*Reader_);
}

LogicFault FaultMessageReaderOpenedTwice("Job::openReader() Fault(0 != Reader_)");
RuntimeCheck CheckopenReaderWasSuccessful("Job::openReader() Check(Reader().good())");

void Job::openReader(string Path) {
    FaultMessageReaderOpenedTwice(0 != Reader_);
    Reader_ = new ifstream(Path.c_str(), ios::binary);
    CheckopenReaderWasSuccessful(Reader().good());
}

void Job::closeReader() {
    if(Reader_) {
        Reader().close();
        delete Reader_;
        Reader_ = 0;
    }
}

RuntimeCheck CheckJobWriterIsValid("Job::Writer() Check(0 != Writer_)");

ofstream& Job::Writer() {                                                       // Safe access to Writer_.
    CheckJobWriterIsValid(0 != Writer_);
    return (*Writer_);
}

LogicFault FaultMessageWriterOpenedTwice("Job::openWriter() Fault(0 != Writer_)");
RuntimeCheck CheckopenWriterWasSuccessful("Job::openWriter() Check(Writer().good())");

void Job::openWriter(string Path) {
    FaultMessageWriterOpenedTwice(0 != Writer_);
    Writer_ = new ofstream(Path.c_str(), ios::trunc | ios::binary);
    CheckopenWriterWasSuccessful(Writer().good());
}

void Job::closeWriter() {
    if(Writer_) {
        Writer().close();
        delete Writer_;
        Writer_ = 0;
    }
}

void Job::doInitialRead() {
    openReader(Job::CurrentCommand.Data);
    ReadBuffer.assign(ReadBufferSize, 0);
    Reader().read(reinterpret_cast<char*>(&ReadBuffer[0]), ReadBuffer.size());
}

string Job::ScanName() {                                                        // Scan name from FILE command
    ostringstream ScanNameFormatter;
    ScanNameFormatter << "[" << CurrentCommand.Number << "]" << CurrentCommand.Data;
    return ScanNameFormatter.str();
}

const int NoScanYet = -1;
LogicFault FaultIfSettingScanResultConfigBeforeScan("Job::setConfigurationFromScanResult() Fault(NoScanYet == ScanResultCode)");

void Job::setConfigurationFromScanResult(ScopeScanner& myScanner) {

    FaultIfSettingScanResultConfigBeforeScan(NoScanYet == ScanResultCode);
    ScanResultConfiguration = Scanners.ConfigurationForResultCode(ScanResultCode);

    /*//// Testing -- forcing scan results

    ScanResultConfiguration.Action = ResultConfiguration::Hold;
    ScanResultConfiguration.HoldPath = "C:\\M\\Projects\\MessageSniffer\\SNF4CGP_Work\\TestEnvironment\\hold\\";
    ScanResultConfiguration.RejectionReason = "I don't like the look of it -- so there!";
    ScanResultConfiguration.LogComment = "This is my happy little log comment.";
    ScanResultConfiguration.EmitXMLLog = true;
    ScanResultConfiguration.EmitClassicLog = false;
    ScanResultConfiguration.InjectHeaders = true;

    *///////////////////////////////////////////////////////////////////////

    if(ScanResultConfiguration.InjectHeaders) HeadersToInject = myScanner.Engine().getXHDRs();
    if(ScanResultConfiguration.EmitXMLLog) XMLLogData = myScanner.Engine().getXMLLog();
    if(ScanResultConfiguration.EmitClassicLog) ClassicLogData = myScanner.Engine().getClassicLog();
}

void Job::doScan() {
    ScopeScanner myScanner(Scanners);
    ScanResultCode =
      myScanner.Engine().scanMessage(
        &ReadBuffer[0], Reader().gcount(), ScanName(), JobTimer.getElapsedTime()
    );
    setConfigurationFromScanResult(myScanner);
}

void convertLogLinesToComments(ostringstream& O, string& Lines) {
    istringstream I(Lines);
    while(I) {
        string LogLine;
        getline(I, LogLine);
        if(0 < LogLine.length()) {
            O << "* " << LogLine << endl;
        }
    }
}

void Job::emitLogComment() {
    ostringstream O;
    generateCommandCommentPrefix(O);
    O << "Comment: " << ScanResultConfiguration.LogComment << endl;
    OutputBuffer.append(O.str());
}

void Job::emitXMLLogData() {
    ostringstream O;
    generateCommandCommentPrefix(O);
    O << "XML Scan log data: " << endl;
    convertLogLinesToComments(O, XMLLogData);
    OutputBuffer.append(O.str());
}

void Job::emitClassicLogData() {
    ostringstream O;
    generateCommandCommentPrefix(O);
    O << "Classic Scan log data: " << endl;
    convertLogLinesToComments(O, ClassicLogData);
    OutputBuffer.append(O.str());
}

LogicFault FaultUhandledScanResultAction("Job::doAction() switch(ScanResultConfiguration.Action) fell to default!");

void Job::doAction() {
    if(0 < ScanResultConfiguration.LogComment.length()) emitLogComment();
    if(ScanResultConfiguration.EmitXMLLog) emitXMLLogData();
    if(ScanResultConfiguration.EmitClassicLog) emitClassicLogData();
    switch(ScanResultConfiguration.Action) {
        case ResultConfiguration::Allow :  { doAllow(); break; }
        case ResultConfiguration::Bypass : { doBypass(); break; }
        case ResultConfiguration::Delete : { doDelete(); break; }
        case ResultConfiguration::Hold :   { doHold(); break; }
        case ResultConfiguration::Reject : { doReject(); break; }
        default: { throw FaultUhandledScanResultAction; }
    }
}

void Job::doBypass() {
    emitOK();
}

void Job::doAllow() {
    if(0 < HeadersToInject.length()) emitADDHEADER(HeadersToInject);
    else emitOK();
}

void Job::doReject() {
    emitREJECTED(ScanResultConfiguration.RejectionReason);
}

void Job::doDelete() {
    emitDISCARD();
}

void Job::emitMessageMovedTo(string& Destination) {
    ostringstream O;
    generateCommandCommentPrefix(O);
    O << "Moved message to " << Destination << endl;
    OutputBuffer.append(O.str());
}

void Job::identifyLocalLineEnd() {                                              // Identify the line endings in use:
    if(0 < LocalLineEnd.length()) return;                                       // Do this only once for any system.
    for(int i = 0; i < Reader().gcount(); i++) {                                // Look for the first line end.
        if('\r' == ReadBuffer[i] && '\n' == ReadBuffer[i+1]) {                  // It may be \r\n.
            LocalLineEnd = "\r\n"; return;
        } else
        if('\n' == ReadBuffer[i] && '\n' == ReadBuffer[i+1]) {                  // It may be \n. (find block end \n\n)
            LocalLineEnd = "\n"; return;
        }
    }
}

RuntimeCheck CheckLineEndsIdentified("Job::findHeaderInsertPoint() Check(0 < LocalLineEnd.length())");
RuntimeCheck CheckInsertCursorPlausable("Job::findHeaderInsertPoint() Check(isMinimumSafeDistanceFromEOF(InsertCursor))");

bool Job::isMinimumSafeDistanceFromEOF(const size_t Position) {
    string MinimalEmptyMessageBody = LocalLineEnd + "." + LocalLineEnd;
    size_t MinimalDistanceToEOF = MinimalEmptyMessageBody.length();
    size_t SafetyCalculation = Position + MinimalDistanceToEOF;
    size_t AvailableMessageLength = Reader().gcount();
    return (SafetyCalculation <= AvailableMessageLength);
}

void Job::moveInsertCursorForwardToBlockTerminus(size_t& InsertCursor, string& BlockTerminus) {
    size_t SafetyLimit = (Reader().gcount() - BlockTerminus.length());
    for(;InsertCursor < SafetyLimit; InsertCursor++) {                          // Scan forward safely until we find it.
        if(                                                                     // At each position compare the BlockTerminus
          0 == BlockTerminus.compare(                                           // with an equal length string in the buffer.
            0, BlockTerminus.length(),
            reinterpret_cast<char*>(&ReadBuffer[InsertCursor]), BlockTerminus.length())
          ) return;
    }
}

size_t Job::findHeaderInsertPoint() {                                           // How to find the insert point:
    identifyLocalLineEnd();
    CheckLineEndsIdentified(0 < LocalLineEnd.length());
    string BlockTerminus = LocalLineEnd + LocalLineEnd;
    size_t InsertCursor = 0;

    moveInsertCursorForwardToBlockTerminus(InsertCursor, BlockTerminus);        // Find the end of the CGP control block.
    InsertCursor += BlockTerminus.length();                                     // Move past this block terminus.
    moveInsertCursorForwardToBlockTerminus(InsertCursor, BlockTerminus);        // Find the end of the message headers.

    InsertCursor += LocalLineEnd.length();                                      // Split the headers BlockTerminus.

    CheckInsertCursorPlausable(isMinimumSafeDistanceFromEOF(InsertCursor));

    return InsertCursor;
}

RuntimeCheck CheckFirstSegmentWriterOk("Job::writeUpToInjectionPoint() Check(Writer().good())");

void Job::writeUpToInjectionPoint(size_t InjectionPoint) {
    Writer().write(reinterpret_cast<char*>(&ReadBuffer[0]), InjectionPoint);
    CheckFirstSegmentWriterOk(Writer().good());
}

RuntimeCheck CheckHeaderWriterOk("Job::writeHeaders() Check(Writer().good())");

void Job::writeHeaders() {                                                      // Write (inject) headers while converting
    for(size_t i = 0; i < HeadersToInject.length(); i++) {                      // line ends to the format observed in the
        switch(HeadersToInject.at(i)) {                                         // message file.
            case '\r' : { break; }                                              // Convert by ignoring <CR> and
            case '\n' : { Writer() << LocalLineEnd; break; }                    // triggering the LocalLineEnd on <LF>
            default : { Writer().put(HeadersToInject.at(i)); break; }           // All other bytes as seen in the message.
        }
    }
    CheckHeaderWriterOk(Writer().good());
}

RuntimeCheck CheckEndFirstBufferWriterOk("Job::writeRestOfFirstBuffer() Check(Writer().good())");

void Job::writeRestOfFirstBuffer(size_t InjectionPoint) {
    int LengthOfRestOfBuffer = Reader().gcount() - (InjectionPoint + 1);
    Writer().write(reinterpret_cast<char*>(&ReadBuffer[InjectionPoint]), LengthOfRestOfBuffer);
    CheckEndFirstBufferWriterOk(Writer().good());
}

void Job::copyFirstBufferAndInjectHeadersInMessageMove() {
    size_t HeaderInjectionPoint = findHeaderInsertPoint();
    writeUpToInjectionPoint(HeaderInjectionPoint);
    writeHeaders();
    writeRestOfFirstBuffer(HeaderInjectionPoint);
}

RuntimeCheck CheckFirstBufferWriterOk("Job::copyFirstBufferInMessageMove() Check(Writer().good())");

void Job::copyFirstBufferInMessageMove() {
    Writer().write(reinterpret_cast<char*>(&ReadBuffer[0]), Reader().gcount());
    CheckFirstBufferWriterOk(Writer().good());
}

RuntimeFault FaultLongCopyReaderBad("Job::readLongCopyBlock() Fault(Reader().bad())");

void Job::readLongCopyBlock() {
    Reader().read(reinterpret_cast<char*>(&ReadBuffer[0]), ReadBufferSize);
    FaultLongCopyReaderBad(Reader().bad());
}

RuntimeFault FaultLongCopyWriterBad("Job::writeLongCopyBlock() Fault(Writer().bad())");

void Job::writeLongCopyBlock() {
     if(Reader().gcount()) {
        Writer().write(reinterpret_cast<char*>(&ReadBuffer[0]), Reader().gcount());
        FaultLongCopyWriterBad(Writer().bad());
     }
}

void Job::copyRestOfLongMessage() {
    while(Reader().good()) {
        readLongCopyBlock();
        writeLongCopyBlock();
    }
}

void Job::moveMessageToHoldPath(string& Destination) {
    openWriter(Destination);
    if(ScanResultConfiguration.InjectHeaders) {
        copyFirstBufferAndInjectHeadersInMessageMove();
    } else {
        copyFirstBufferInMessageMove();
    }
    copyRestOfLongMessage();
    closeWriter();
}

const string PathSeparators = "\\/";

string FileNamePart(string Path) {                                              // How to get the file name part of a path:
    string::size_type Found = Path.find_last_of(PathSeparators);
    if(string::npos != Found) Path = Path.substr(Found + 1);                    // Return all after the last separator
    return Path;                                                                // or Path if no separators are found.
}

string Job::calculateDestinationFileName() {                                    // How to get the message move destination:
    string Source = CurrentCommand.Data;
    string Destination = ScanResultConfiguration.HoldPath;
    Destination.append(FileNamePart(Source));
    return Destination;
}

void Job::doHold() {
    string Destination = calculateDestinationFileName();
    moveMessageToHoldPath(Destination);
    emitMessageMovedTo(Destination);
    emitDISCARD();
}

Job::Job(ScannerPool& S, OutputProcessor& O) :                                  // Construct with important links.
  Scanners(S),
  Output(O),
  ScanResultCode(NoScanYet),
  Reader_(0),
  Writer_(0) {                                                                  // Minimize heap thrashing.
    OutputBuffer.reserve(StringReserveSize);
    HeadersToInject.reserve(StringReserveSize);
    XMLLogData.reserve(StringReserveSize);
    ClassicLogData.reserve(StringReserveSize);
    ReadBuffer.reserve(ReadBufferSize);
}

Job::~Job() {                                                                   // Cleanup when destructing.
    try { clear(); }
    catch(...) {}
}

void Job::clear() {                                                             // Cleanup for the next command.
    CurrentCommand.clear();
    ScanResultCode = NoScanYet;
    OutputBuffer.clear();
    HeadersToInject.clear();
    XMLLogData.clear();
    ClassicLogData.clear();
    ReadBuffer.clear();
    closeReader();
    closeWriter();
    JobTimer.clear();
}

LogicFault FaultIfQuitGetsHere("Job::setCommand() Fault(Command::QUIT == C.Type)");

void Job::setCommand(Command& C) {                                              // Assign a command for this job.
    JobTimer.start();
    FaultIfQuitGetsHere(Command::QUIT == C.Type);
    CurrentCommand = C;
}

void Job::executeCommand() {
    switch(CurrentCommand.Type) {
        case Command::WAKE: { doWakeUp(); break; }
        case Command::INTF: { doINTF(); break; }
        case Command::FILE: { doFILE(); break; }
        default: { doUNKNOWN(); break; }
    }
}

void Job::emitException(const string& What) {
    ostringstream O;
    generateCommandCommentPrefix(O);
    O << "Exception: " << What << endl;
    OutputBuffer.append(O.str());
}

void Job::emitUnknownException() {
    ostringstream O;
    generateCommandCommentPrefix(O);
    O << "Unknown Exception!" << endl;
    OutputBuffer.append(O.str());
}

void Job::doIt() {                                                              // Get the job done.
    try {                                                                       // emitFAILURE() on all exceptions.
        try { executeCommand(); }                                               // executeCommand() and report exceptions.
        catch(exception& e) { emitException(e.what()); throw; }
        catch(...) { emitUnknownException(); throw; }
    }
    catch(...) { emitFAILURE(); }
    finalize();
}

string& Job::Responses() {                                                      // Read the job's report.
    return OutputBuffer;
}

//// JobPool ///////////////////////////////////////////////////////////////////

JobPool::JobPool() :                                                            // Simple constructor.
  Output_(0),
  Scanners_(new ScannerPool()),
  AllocatedJobs(0),
  Started(false) {
}

JobPool::~JobPool() {                                                           // Clean up on the way out.
    try {
        stop();
        Output_ = 0;
        if(Scanners_) delete Scanners_;
        Scanners_ = 0;
    }
    catch(...) {}
}

RuntimeCheck CheckForValidOutputPool("JobPool::Output() Check(0 != Output_)");

OutputProcessor& JobPool::Output() {
    CheckForValidOutputPool(0 != Output_);
    return (*Output_);
}

RuntimeCheck CheckForValidScannerPool("JobPool::Scanners() Check(0 != Scanners_)");

ScannerPool& JobPool::Scanners() {
    CheckForValidScannerPool(0 != Scanners_);
    return (*Scanners_);
}

// Watch out -- Initializing the JobPool also means starting up SNFMulti and
// the ScannerPool.

void JobPool::init(
  string VersionInfo,
  string Configuration,
  OutputProcessor& O) {
    Output_ = &O;
    Scanners().init(VersionInfo, Configuration);
    ScopeMutex Busy(AllocationMutex);
    Job* FirstJob = makeJob();
    Jobs.give(FirstJob);
    Started = true;
}

Job* JobPool::makeJob() {                                                       // Create and count a Job.
    Job* NewJob = new Job(Scanners(), Output());
    ++AllocatedJobs;
    return NewJob;
}

LogicCheck CheckGrabJobsOnlyWhenStarted("JobPool::grab() Check(Started)");
RuntimeCheck CheckForValidGrabbedJob("JobPool::grab() Check(0 != GrabbedJob)");

Job& JobPool::grab() {                                                          // Grab a job (prefer from the pool).
    ScopeMutex Busy(AllocationMutex);
    CheckGrabJobsOnlyWhenStarted(Started);
    Job* GrabbedJob = 0;
    if(0 < Jobs.size()) GrabbedJob = Jobs.take();                               // Prefer jobs from the pool.
    else GrabbedJob = makeJob();                                                // Make new ones if needed.
    CheckForValidGrabbedJob(0 != GrabbedJob);
    return (*GrabbedJob);
}

void JobPool::drop(Job& J) {                                                    // Drop a job into the pool.
    Jobs.give(&J);
}

void JobPool::killJobFromPool() {                                               // Kill and count a Job from the pool.
    Job* JobToKill = 0;
    JobToKill = Jobs.take();
    delete JobToKill;
    --AllocatedJobs;
}

// Watch out -- Stopping the JobPool also means shutting down SNFMulti and
// the ScannerPool.

void JobPool::stop() {                                                          // Shut down the JobPool.
    ScopeMutex Busy(AllocationMutex);
    while(0 < AllocatedJobs) killJobFromPool();
    Scanners().stop();
    Started = false;
}