<commit_before>#include "SoloFeature.h"
#include "streamFuns.h"
#include "TimeFunctions.h"
#include "SequenceFuns.h"
#include "SoloCommon.h"
#include <unordered_map>
#include <bitset>
#include "serviceFuns.cpp"
#include <math.h>       /* log */

void SoloFeature::quantTranscript()
{//velocyto counting gets info from Gene counting
    time_t rawTime;

    if (pSolo.clusterCBfile=="-")
        return;//no transcript quantification w/o cluster file

    std::unordered_map<uint32,uint32> clusterCBind; //for each CB - cluster index of detected CB
    std::set<uint32> clusterInd; //cluster index for each cluster - the integer from clusterCBfile
    {//load cluster information
        ifstream &clusterStream = ifstrOpen(pSolo.clusterCBfile, ERROR_OUT, "SOLUTION: check the path and permissions of the cluster CB file: " + pSolo.clusterCBfile, P);
        //std::set<uint32> clusterInd; //all cluster indexes
        string seq1;
        while (clusterStream >> seq1) {
            uint32 icl1;
            clusterStream >> icl1;
            uint64 cb1;
            if (convertNuclStrToInt64(seq1,cb1)) {//convert to 2-bit format
                auto cb1it=std::equal_range(pSolo.cbWL.begin(), pSolo.cbWL.end(), cb1);
                uint32 cb1ind=(uint32) (cb1it.first-pSolo.cbWL.begin());
                if (cb1ind < pSolo.cbWL.size()) {
                    clusterCBind[cb1ind]=icl1;
                    clusterInd.emplace(icl1);
                } else {
                    P.inOut->logMain << "WARNING: cluster CB sequence not present in whitelist and is ignored: " << seq1 <<endl;
                };
            } else {
                P.inOut->logMain << "WARNING: cluster CB sequence contains non-ACGT base and is ignored: " << seq1 <<endl;
            };
        };
    };

    auto &trDistCount=readFeatSum->transcriptDistCount;
    vector<double> trDistFun(trDistCount.size(),0.0);
    vector<double> trDistFunTrFactor(Trans.nTr,0.0);

    {//process distance distribution function
        //running average
        int32 runAverN=50, runAverStart=0;
        //for (uint32 ii=0; ii<runAverN; ii++)
        //    trDistFun[0] += trDistCount[ii];
        for (int32 ii=0; ii<runAverStart; ii++)
            trDistFun[ii] = (double) trDistCount[ii];
        for (int32 ii=runAverStart; ii<(int32)trDistCount.size()-runAverN-1; ii++) {
            trDistFun[ii] = (double) std::accumulate(trDistCount.begin()+max(runAverStart,ii-runAverN), trDistCount.begin()+ii+runAverN+1, 0) / min(2*runAverN+1, ii-runAverStart+runAverN);
        };

        //cut when becomes non-monotonic
        uint32 imax=1000;
        while (trDistFun[imax+1]>trDistFun[imax])
            imax++;
        P.inOut->logMain << "SoloQuant: distance distribution past maximum = " << imax <<endl;

        while (trDistFun[imax+1]<trDistFun[imax])
            imax++;
        P.inOut->logMain << "SoloQuant: distance distribution cutoff = " << imax <<endl;

        trDistFun.resize(imax);

        //normalize
        double norm1 = std::accumulate(trDistFun.begin(), trDistFun.end(), 0.0);

        ofstream *streamTrDistFun = &ofstrOpen(outputPrefix+"transcriptEndDistanceDistribution.txt",ERROR_OUT, P);
        for (auto & ff : trDistFun) {
            ff = ff/norm1;
            *streamTrDistFun << ff <<'\n';
        };
        streamTrDistFun->close();

        //normalization factors for all transcripts
        vector<double> trDistFunCum(trDistFun);
        std::partial_sum(trDistFun.begin(), trDistFun.end(), trDistFunCum.begin());
        for (uint32 ii=0; ii<Trans.nTr; ii++)
            if (Trans.trLen[ii]<trDistFunCum.size())
                trDistFunTrFactor[ii]=-std::log(trDistFunCum[Trans.trLen[ii]-1]);

        for (auto & ff : trDistFun)
            ff = std::log(ff);

    };

    typedef struct {
        uint32 tr;
        double d;
    } transcriptDistProbStruct;

    map<uint32, unordered_map<uint64,vector<transcriptDistProbStruct>>> mapTrDist;

    //////////// input records
    for (int iThread=0; iThread<P.runThreadN; iThread++) {//TODO: this can be parallelized
        fstream *streamReads = readFeatAll[iThread]->streamReads;
        streamReads->flush();
        streamReads->seekg(0,ios::beg);

        string line1;
        while (std::getline(*streamReads, line1)) {//until the end of file
            stringstream lineStream(line1);
            uint32 cb, cbCl, nTr;
            uint64 umi;
            lineStream >> cb >> umi >> nTr;

            if (clusterCBind.count(cb)==0) //this cb is not in the clusters
                continue;

            umi += (((uint64)cb)<<32); //to make UMI from different CB distinguishable

            cbCl=clusterCBind[cb];

            vector<transcriptDistProbStruct> tD;
            tD.reserve(nTr);
            for (uint32 ii=0; ii<nTr; ii++) {
                uint32 tr1,d1;
                lineStream >> tr1 >> d1;
                if (d1>=trDistFun.size())
                    continue; //do not record such outlier

                tD.push_back( {tr1, trDistFun[d1] + trDistFunTrFactor[tr1]} );
            };

            if (tD.size()==0)
                continue;

            std::sort(tD.begin(), tD.end(), [](const transcriptDistProbStruct &t1, const transcriptDistProbStruct &t2) {
                                                    return (t1.tr<t2.tr);
                                              }
                     );

            if (mapTrDist[cbCl].count(umi)==0) {//1st entry for this umi
                mapTrDist[cbCl][umi]=tD;
                continue;
            };

            uint32 inew=0;
            vector<transcriptDistProbStruct> tD1;
            tD1.reserve(mapTrDist[cbCl][umi].size());

            for (uint32 iold=0; iold<mapTrDist[cbCl][umi].size(); iold++) {//intersection of old with new
                while (inew < tD.size() && mapTrDist[cbCl][umi][iold].tr>tD[inew].tr) //move through the sorted lists
                    ++inew;

                if (inew == tD.size() ) //end of tD reached
                    break;

                if (mapTrDist[cbCl][umi][iold].tr == tD[inew].tr) {//add new log probability to the
                    tD1.push_back({  tD[inew].tr, mapTrDist[cbCl][umi][iold].d + tD[inew].d  });
                };
            };
            mapTrDist[cbCl][umi]=tD1;//replace with intersection
        };
    };

    time(&rawTime);
    P.inOut->logMain << timeMonthDayTime(rawTime) << " ... Transcript3p counting: finished input" <<endl;

    map<uint32, vector<double>> clusterExpression; //per cluster, relative abundance
    //TODO parallelize this loop
    for (auto & mapTrDist1 : mapTrDist) {
        auto &clTrDist=mapTrDist1.second;

        vector<double> trUnique(Trans.nTr,0), trInitial(Trans.nTr,0);//counts of unique read for each transcripts (i.e. reads that map uniquely only to this transcript)
        uint64 nUMItot=0, nUMI0=0, nUMI1=0;

        {//pre-process mapTrDist:
            auto clTrDist1=clTrDist.begin();
            while ( clTrDist1!=clTrDist.end() ) {
                auto &trDist=clTrDist1->second;

                if (trDist.size()==0) {
                    clTrDist1=clTrDist.erase(clTrDist1);
                    nUMI0++;
                    continue;
                } else if (trDist.size()==1) {
                    trUnique[trDist[0].tr]++;
                    trInitial[trDist[0].tr] += 1.0;
                    clTrDist1=clTrDist.erase(clTrDist1);
                    nUMI1++;
                    nUMItot++;
                    continue;
                };

                double max1= std::max_element(trDist.begin(), trDist.end(), [](const transcriptDistProbStruct &t1, const transcriptDistProbStruct &t2) {
                                                                                return (t1.d<t2.d);
                                                                                }) -> d;
                for (auto & tt : trDist) {
                    trInitial[tt.tr] += 1.0/trDist.size();//for initialization, split each count between multimappers evenly
                    tt.d =std::exp(tt.d-max1);
                };

                nUMItot++;
                clTrDist1++;
            };
            P.inOut->logMain << timeMonthDayTime(rawTime) << " ... Transcript3p counting: cluster " << mapTrDist1.first <<" nUMItot="<<nUMItot<<" nUMI0="<<nUMI0<<" nUMI1="<<nUMI1 <<endl;
        };

//         double trExpressed=0;
//         for (auto & tt : trInitial)
//             trExpressed += tt;
//
//         for (auto & tt : trInitial)
//             tt *= nUMItot/trExpressed;

        vector<double> thOldNew[2];
        thOldNew[0]=trInitial;
        thOldNew[1].resize(Trans.nTr,0);

        auto *thOldP = &thOldNew[0];
        auto *thNewP = &thOldNew[1];

        vector<bool> trConverged(trInitial.size(), false);

        for (uint32 iteration=0; iteration<10000; iteration++) {//main EM loop
            auto &thNew = *thNewP;
            auto &thOld = *thOldP;

            //always start with trUnique counts
            std::copy(trUnique.begin(), trUnique.end(), thNew.begin());

            //calculate thNew
            for (auto & clTrDist1: clTrDist) {//loop over all UMIs
                auto &trDist=clTrDist1.second;

                double denom1=0.0;
                for  (auto & td : trDist) {
                    denom1 += td.d * thOld[td.tr];
                };

                for  (auto & td : trDist) {
                    if (!trConverged[td.tr])
                        thNew[td.tr] += td.d * thOld[td.tr] / denom1;
                };
            };

            //convergence check
            double diffThresholdMax=1e-6;
            double diffThresholdOne=1e-9;
            double exprThreshold=0*nUMItot;
            double diffMax=0, diffSum=0, aboveThrN=0, aboveThrExprSum=0, aboveThrOneN=0 ;
            for (uint32 itr=0; itr<thNew.size(); itr++) {
                if (trConverged[itr] || thOld[itr]==0)
                    continue;

                double diff1 = std::abs(thNew[itr]-thOld[itr])/thOld[itr];
                diffSum += diff1;
                diffMax = max(diffMax,diff1);
                if (diff1 > diffThresholdMax) {
                    aboveThrN++;
                    aboveThrExprSum += thNew[itr];
                };

                if (thNew[itr]<exprThreshold) {
                    trConverged[itr]=true;
                    trUnique[itr]=thNew[itr];//0 here could make more sense, but convergence suffers
                };
                if (diff1<diffThresholdOne) {
                    trConverged[itr]=true;
                    trUnique[itr]=thNew[itr]; //this is the final value for this transcript
                } else {
                    aboveThrOneN++;
                };
            };

            cout <<iteration <<" "<<diffMax<<" "<<diffSum<<" "<<aboveThrN<<" "<<aboveThrExprSum<<" "<<aboveThrOneN<<endl;

            if (diffMax<diffThresholdMax)
                break;
            swap(thNewP,thOldP); //swap for the next iteration
        };

        auto &thOut=*thNewP;
        {//renormalize into TPMs
            double norm1=0;
            for (uint32 itr=0; itr<thOut.size(); itr++) {
                 thOut[itr] *= std::exp(trDistFunTrFactor[itr]);
                 norm1 += thOut[itr];
            };
            norm1=nUMItot/norm1;
            for (uint32 itr=0; itr<thOut.size(); itr++) {
                 thOut[itr] *= norm1;
            };
        };

        clusterExpression[mapTrDist1.first]=thOut;

        time(&rawTime);
        P.inOut->logMain << timeMonthDayTime(rawTime) << " ... Transcript3p counting: finished cluster" << mapTrDist1.first <<endl;
    };

    {//output counting matrix
        string matrixFileName=outputPrefix+pSolo.outFileNames[3];
        ofstream &countMatrixStream=ofstrOpen(matrixFileName,ERROR_OUT, P);
        countMatrixStream <<"%%MatrixMarket matrix coordinate real general\n%\n";

        uint32 nCellGeneEntries = 0;
        for (auto & ctpm : clusterExpression)
            for (auto & tt : ctpm.second)
                if (tt>0)
                    nCellGeneEntries++;

        countMatrixStream << Trans.nTr <<' '<< *std::max_element(clusterInd.begin(),clusterInd.end()) <<' '<< nCellGeneEntries << '\n';

        for (auto & ctpm : clusterExpression)
            for (auto itt=ctpm.second.begin(); itt!=ctpm.second.end(); itt++)
                if (*itt>0)
                    countMatrixStream << itt-ctpm.second.begin()+1  <<' '<<  ctpm.first  <<' '<< *itt << '\n';

        countMatrixStream.flush();
    };

    ofstream &outStr=ofstrOpen(outputPrefix+"/features.tsv",ERROR_OUT, P);
    for (uint32 ii=0; ii<Trans.nTr; ii++)
        outStr << Trans.trID[ii] <<"\t"<< Trans.trLen[ii] <<"\t"<< Trans.geName[Trans.trGene[ii]] << '\n';
    outStr.close();

    time(&rawTime);
    P.inOut->logMain << timeMonthDayTime(rawTime) << " ... Transcript3p counting: finished transcript quantification" <<endl;
};
<commit_msg>Minor changes to EM algorithm.<commit_after>#include "SoloFeature.h"
#include "streamFuns.h"
#include "TimeFunctions.h"
#include "SequenceFuns.h"
#include "SoloCommon.h"
#include <unordered_map>
#include <bitset>
#include "serviceFuns.cpp"
#include <math.h>       /* log */

void SoloFeature::quantTranscript()
{//velocyto counting gets info from Gene counting
    time_t rawTime;

    if (pSolo.clusterCBfile=="-")
        return;//no transcript quantification w/o cluster file

    std::unordered_map<uint32,uint32> clusterCBind; //for each CB - cluster index of detected CB
    std::set<uint32> clusterInd; //cluster index for each cluster - the integer from clusterCBfile
    {//load cluster information
        ifstream &clusterStream = ifstrOpen(pSolo.clusterCBfile, ERROR_OUT, "SOLUTION: check the path and permissions of the cluster CB file: " + pSolo.clusterCBfile, P);
        //std::set<uint32> clusterInd; //all cluster indexes
        string seq1;
        while (clusterStream >> seq1) {
            uint32 icl1;
            clusterStream >> icl1;
            uint64 cb1;
            if (convertNuclStrToInt64(seq1,cb1)) {//convert to 2-bit format
                auto cb1it=std::equal_range(pSolo.cbWL.begin(), pSolo.cbWL.end(), cb1);
                uint32 cb1ind=(uint32) (cb1it.first-pSolo.cbWL.begin());
                if (cb1ind < pSolo.cbWL.size()) {
                    clusterCBind[cb1ind]=icl1;
                    clusterInd.emplace(icl1);
                } else {
                    P.inOut->logMain << "WARNING: cluster CB sequence not present in whitelist and is ignored: " << seq1 <<endl;
                };
            } else {
                P.inOut->logMain << "WARNING: cluster CB sequence contains non-ACGT base and is ignored: " << seq1 <<endl;
            };
        };
    };

    auto &trDistCount=readFeatSum->transcriptDistCount;
    vector<double> trDistFun(trDistCount.size(),0.0);
    vector<double> trDistFunTrFactor(Trans.nTr,0.0);

    {//process distance distribution function
        //running average
        int32 runAverN=50, runAverStart=0;
        //for (uint32 ii=0; ii<runAverN; ii++)
        //    trDistFun[0] += trDistCount[ii];
        for (int32 ii=0; ii<runAverStart; ii++)
            trDistFun[ii] = (double) trDistCount[ii];
        for (int32 ii=runAverStart; ii<(int32)trDistCount.size()-runAverN-1; ii++) {
            trDistFun[ii] = (double) std::accumulate(trDistCount.begin()+max(runAverStart,ii-runAverN), trDistCount.begin()+ii+runAverN+1, 0) / min(2*runAverN+1, ii-runAverStart+runAverN);
        };

        //cut when becomes non-monotonic
        uint32 imax=1000;
        while (trDistFun[imax+1]>trDistFun[imax])
            imax++;
        P.inOut->logMain << "SoloQuant: distance distribution past maximum = " << imax <<endl;

        while (trDistFun[imax+1]<trDistFun[imax])
            imax++;
        P.inOut->logMain << "SoloQuant: distance distribution cutoff = " << imax <<endl;

        trDistFun.resize(imax);

        //normalize
        double norm1 = std::accumulate(trDistFun.begin(), trDistFun.end(), 0.0);

        ofstream *streamTrDistFun = &ofstrOpen(outputPrefix+"transcriptEndDistanceDistribution.txt",ERROR_OUT, P);
        for (auto & ff : trDistFun) {
            ff = ff/norm1;
            *streamTrDistFun << ff <<'\n';
        };
        streamTrDistFun->close();

        //normalization factors for all transcripts
        vector<double> trDistFunCum(trDistFun);
        std::partial_sum(trDistFun.begin(), trDistFun.end(), trDistFunCum.begin());
        for (uint32 ii=0; ii<Trans.nTr; ii++)
            if (Trans.trLen[ii]<trDistFunCum.size())
                trDistFunTrFactor[ii]=-std::log(trDistFunCum[Trans.trLen[ii]-1]);

        for (auto & ff : trDistFun)
            ff = std::log(ff);

    };

    typedef struct {
        uint32 tr;
        double d;
    } transcriptDistProbStruct;

    map<uint32, unordered_map<uint64,vector<transcriptDistProbStruct>>> mapTrDist;

    //////////// input records
    for (int iThread=0; iThread<P.runThreadN; iThread++) {//TODO: this can be parallelized
        fstream *streamReads = readFeatAll[iThread]->streamReads;
        streamReads->flush();
        streamReads->seekg(0,ios::beg);

        string line1;
        while (std::getline(*streamReads, line1)) {//until the end of file
            stringstream lineStream(line1);
            uint32 cb, cbCl, nTr;
            uint64 umi;
            lineStream >> cb >> umi >> nTr;

            if (clusterCBind.count(cb)==0) //this cb is not in the clusters
                continue;

            umi += (((uint64)cb)<<32); //to make UMI from different CB distinguishable

            cbCl=clusterCBind[cb];

            vector<transcriptDistProbStruct> tD;
            tD.reserve(nTr);
            for (uint32 ii=0; ii<nTr; ii++) {
                uint32 tr1,d1;
                lineStream >> tr1 >> d1;
                if (d1>=trDistFun.size())
                    continue; //do not record such outlier

                tD.push_back( {tr1, trDistFun[d1] + trDistFunTrFactor[tr1]} );
            };

            if (tD.size()==0)
                continue;

            std::sort(tD.begin(), tD.end(), [](const transcriptDistProbStruct &t1, const transcriptDistProbStruct &t2) {
                                                    return (t1.tr<t2.tr);
                                              }
                     );

            if (mapTrDist[cbCl].count(umi)==0) {//1st entry for this umi
                mapTrDist[cbCl][umi]=tD;
                continue;
            };

            uint32 inew=0;
            vector<transcriptDistProbStruct> tD1;
            tD1.reserve(mapTrDist[cbCl][umi].size());

            for (uint32 iold=0; iold<mapTrDist[cbCl][umi].size(); iold++) {//intersection of old with new
                while (inew < tD.size() && mapTrDist[cbCl][umi][iold].tr>tD[inew].tr) //move through the sorted lists
                    ++inew;

                if (inew == tD.size() ) //end of tD reached
                    break;

                if (mapTrDist[cbCl][umi][iold].tr == tD[inew].tr) {//add new log probability to the
                    tD1.push_back({  tD[inew].tr, mapTrDist[cbCl][umi][iold].d + tD[inew].d  });
                };
            };
            mapTrDist[cbCl][umi]=tD1;//replace with intersection
        };
    };

    time(&rawTime);
    P.inOut->logMain << timeMonthDayTime(rawTime) << " ... Transcript3p counting: finished input" <<endl;

    map<uint32, vector<double>> clusterExpression; //per cluster, relative abundance
    //TODO parallelize this loop
    for (auto & mapTrDist1 : mapTrDist) {
        auto &clTrDist=mapTrDist1.second;

        vector<double> trUnique(Trans.nTr,0), trInitial(Trans.nTr,0);//counts of unique read for each transcripts (i.e. reads that map uniquely only to this transcript)
        uint64 nUMItot=0, nUMI0=0, nUMI1=0;

        {//pre-process mapTrDist:
            auto clTrDist1=clTrDist.begin();
            while ( clTrDist1!=clTrDist.end() ) {
                auto &trDist=clTrDist1->second;

                if (trDist.size()==0) {
                    clTrDist1=clTrDist.erase(clTrDist1);
                    nUMI0++;
                    continue;
                } else if (trDist.size()==1) {
                    trUnique[trDist[0].tr]++;
                    trInitial[trDist[0].tr] += 1.0;
                    clTrDist1=clTrDist.erase(clTrDist1);
                    nUMI1++;
                    nUMItot++;
                    continue;
                };

                double max1= std::max_element(trDist.begin(), trDist.end(), [](const transcriptDistProbStruct &t1, const transcriptDistProbStruct &t2) {
                                                                                return (t1.d<t2.d);
                                                                                }) -> d;
                for (auto & tt : trDist) {
                    trInitial[tt.tr] += 1.0/trDist.size();//for initialization, split each count between multimappers evenly
                    tt.d =std::exp(tt.d-max1);
                };

                nUMItot++;
                clTrDist1++;
            };
            P.inOut->logMain << timeMonthDayTime(rawTime) << " ... Transcript3p counting: cluster " << mapTrDist1.first <<" nUMItot="<<nUMItot<<" nUMI0="<<nUMI0<<" nUMI1="<<nUMI1 <<endl;
        };

//         double trExpressed=0;
//         for (auto & tt : trInitial)
//             trExpressed += tt;
//
//         for (auto & tt : trInitial)
//             tt *= nUMItot/trExpressed;

        vector<double> thOldNew[2];
        thOldNew[0]=trInitial;
        thOldNew[1].resize(Trans.nTr,0);

        auto *thOldP = &thOldNew[0];
        auto *thNewP = &thOldNew[1];

        vector<bool> trConverged(trInitial.size(), false);

        for (uint32 iteration=0; iteration<10000; iteration++) {//main EM loop
            auto &thNew = *thNewP;
            auto &thOld = *thOldP;

            //always start with trUnique counts
            std::copy(trUnique.begin(), trUnique.end(), thNew.begin());

            //calculate thNew
            for (auto & clTrDist1: clTrDist) {//loop over all UMIs
                auto &trDist=clTrDist1.second;

                double denom1=0.0;
                for  (auto & td : trDist) {
                    denom1 += td.d * thOld[td.tr];
                };

                for  (auto & td : trDist) {
                    if (!trConverged[td.tr])
                        thNew[td.tr] += td.d * thOld[td.tr] / denom1;
                };
            };

            //convergence check
            double diffThresholdMax=1e-5;
            double diffThresholdOne=diffThresholdMax*0.1;
            double exprThreshold=1e-8*nUMItot;
            double diffMax=0, diffSum=0, aboveThrN=0, aboveThrExprSum=0, aboveThrOneN=0 ;
            for (uint32 itr=0; itr<thNew.size(); itr++) {
                if (trConverged[itr] || thOld[itr]==0)
                    continue;

                double diff1 = std::abs(thNew[itr]-thOld[itr])/thOld[itr];
                diffSum += diff1;
                diffMax = max(diffMax,diff1);
                if (diff1 > diffThresholdMax) {
                    aboveThrN++;
                    aboveThrExprSum += thNew[itr];
                };

                if (thNew[itr]<exprThreshold) {
                    trConverged[itr]=true;
                    //trUnique[itr]=thNew[itr];//0 here could make more sense, but convergence suffers
                    trUnique[itr]=0;
                };
                if (diff1<diffThresholdOne) {
                    trConverged[itr]=true;
                    trUnique[itr]=thNew[itr]; //this is the final value for this transcript
                } else {
                    aboveThrOneN++;
                };
            };

            cout <<iteration <<" "<<diffMax<<" "<<diffSum<<" "<<aboveThrN<<" "<<aboveThrExprSum<<" "<<aboveThrOneN<<endl;

            if (diffMax<diffThresholdMax)
                break;
            swap(thNewP,thOldP); //swap for the next iteration
        };

        auto &thOut=*thNewP;
        {//renormalize into TPMs
            double norm1=0;
            for (uint32 itr=0; itr<thOut.size(); itr++) {
                 thOut[itr] *= std::exp(trDistFunTrFactor[itr]);
                 norm1 += thOut[itr];
            };
            norm1=nUMItot/norm1;
            for (uint32 itr=0; itr<thOut.size(); itr++) {
                 thOut[itr] *= norm1;
            };
        };

        clusterExpression[mapTrDist1.first]=thOut;

        time(&rawTime);
        P.inOut->logMain << timeMonthDayTime(rawTime) << " ... Transcript3p counting: finished cluster" << mapTrDist1.first <<endl;
    };

    {//output counting matrix
        string matrixFileName=outputPrefix+pSolo.outFileNames[3];
        ofstream &countMatrixStream=ofstrOpen(matrixFileName,ERROR_OUT, P);
        countMatrixStream <<"%%MatrixMarket matrix coordinate real general\n%\n";

        uint32 nCellGeneEntries = 0;
        for (auto & ctpm : clusterExpression)
            for (auto & tt : ctpm.second)
                if (tt>0)
                    nCellGeneEntries++;

        countMatrixStream << Trans.nTr <<' '<< *std::max_element(clusterInd.begin(),clusterInd.end()) <<' '<< nCellGeneEntries << '\n';

        for (auto & ctpm : clusterExpression)
            for (auto itt=ctpm.second.begin(); itt!=ctpm.second.end(); itt++)
                if (*itt>0)
                    countMatrixStream << itt-ctpm.second.begin()+1  <<' '<<  ctpm.first  <<' '<< *itt << '\n';

        countMatrixStream.flush();
    };

    ofstream &outStr=ofstrOpen(outputPrefix+"/features.tsv",ERROR_OUT, P);
    for (uint32 ii=0; ii<Trans.nTr; ii++)
        outStr << Trans.trID[ii] <<"\t"<< Trans.trLen[ii] <<"\t"<< Trans.geName[Trans.trGene[ii]] << '\n';
    outStr.close();

    time(&rawTime);
    P.inOut->logMain << timeMonthDayTime(rawTime) << " ... Transcript3p counting: finished transcript quantification" <<endl;
};
<|endoftext|>


<commit_before>//
// Created by lf-z on 12/28/16.
//
#include <fstream>
#include "engy/energy_mgmt.hh"
#include "debug/EnergyMgmt.hh"

EnergyMgmt::EnergyMgmt(const Params *p)
        : SimObject(p),
          time_unit(p->energy_time_unit),
          energy_remained(0),
          event_poweroff(this),
          event_poweron(this),
          event_energy_harvest(this),
          _path_energy_profile(p->path_energy_profile)
{

}

EnergyMgmt::~EnergyMgmt()
{

}

void EnergyMgmt::init()
{
    /* Read energy profile */
    energy_harvest_data = readEnergyProfile();
    /* Reset energy remained to 0. */
    energy_remained = 0;

    DPRINTF(EnergyMgmt, "Energy Management module initialized!\n");
    DPRINTF(EnergyMgmt, "Energy profile: %s (Time unit: %d ticks)\n",
            _path_energy_profile.c_str(), time_unit);

    /* Trigger first energy harvest event here */
    energyHarvest();

}

int EnergyMgmt::consumeEnergy(double val)
{
    /* Todo: Pass the module which consumed the energy to this function. (Or DPRINTF in the module which consumes energy) */
    /* Consume energy if val > 0, and harvest energy if val < 0 */
    energy_remained -= val;
    if (val > 0)
        DPRINTF(EnergyMgmt, "Energy %lf is consumed by xxx. Energy remained: %lf\n", val, energy_remained);
    else
        DPRINTF(EnergyMgmt, "Energy %lf is harvested. Energy remained: %lf\n", -val, energy_remained);
    return 1;
}

void EnergyMgmt::broadcastPowerOff()
{
    EnergyMsg msg;
    msg.type = POWEROFF;
    _meport.broadcastMsg(msg);
    DPRINTF(EnergyMgmt, "Insufficient energy, system power off.\n");
}

void EnergyMgmt::broadcastPowerOn()
{
    EnergyMsg msg;
    msg.type = POWERON;
    _meport.broadcastMsg(msg);
    DPRINTF(EnergyMgmt, "Sufficient energy, system power on.\n");
}

std::vector<double> EnergyMgmt::readEnergyProfile()
{
    std::vector<double> data;
    double temp;
    std::ifstream fin;
    fin.open(_path_energy_profile.c_str());
    /* Read energy profile and store the data into vector. */
    data.resize(0);
    while (fin>>temp) {
        data.push_back(temp);
    }
    /* Reverse the energy harvest queue so that the first energy unit pops first */
    reverse(data.begin(), data.end());
    fin.close();
    return data;
}

void EnergyMgmt::energyHarvest()
{
    /* Todo: process energy harvest event and trigger next one */
    /* Add harvested energy into capacity. */
    double energy_val = energy_harvest_data.back();
    consumeEnergy(-energy_val);
    energy_harvest_data.pop_back();

    /* Trigger the next harvest function. */
    schedule(event_energy_harvest, curTick() + time_unit);

    // DPRINTF(EnergyMgmt, "Energy %lf harvested.\n", energy_val);
}

EnergyMgmt *
EnergyMgmtParams::create()
{
    return new EnergyMgmt(this);
}
<commit_msg>engy: add wrappers for harvesting events<commit_after>//
// Created by lf-z on 12/28/16.
//
#include <fstream>
#include "engy/energy_mgmt.hh"
#include "debug/EnergyMgmt.hh"

EnergyMgmt::EnergyMgmt(const Params *p)
        : SimObject(p),
          time_unit(p->energy_time_unit),
          energy_remained(0),
          event_poweroff(this),
          event_poweron(this),
          event_energy_harvest(this),
          _path_energy_profile(p->path_energy_profile)
{

}

EnergyMgmt::~EnergyMgmt()
{

}

void EnergyMgmt::init()
{
    /* Read energy profile */
    energy_harvest_data = readEnergyProfile();
    /* Reset energy remained to 0. */
    energy_remained = 0;

    DPRINTF(EnergyMgmt, "Energy Management module initialized!\n");
    DPRINTF(EnergyMgmt, "Energy profile: %s (Time unit: %d ticks)\n",
            _path_energy_profile.c_str(), time_unit);

    /* Trigger first energy harvest event here */
    energyHarvest();

}

int EnergyMgmt::consumeEnergy(double val)
{
    /* Todo: Pass the module which consumed the energy to this function. (Or DPRINTF in the module which consumes energy) */
    /* Consume energy if val > 0, and harvest energy if val < 0 */
    energy_remained -= val;
    if (val > 0)
        DPRINTF(EnergyMgmt, "Energy %lf is consumed by xxx. Energy remained: %lf\n", val, energy_remained);
    else
        DPRINTF(EnergyMgmt, "Energy %lf is harvested. Energy remained: %lf\n", -val, energy_remained);
    return 1;
}

void EnergyMgmt::broadcastPowerOff()
{
    EnergyMsg msg;
    msg.type = POWEROFF;
    _meport.broadcastMsg(msg);
    DPRINTF(EnergyMgmt, "Insufficient energy, system power off.\n");
}

void EnergyMgmt::broadcastPowerOn()
{
    EnergyMsg msg;
    msg.type = POWERON;
    _meport.broadcastMsg(msg);
    DPRINTF(EnergyMgmt, "Sufficient energy, system power on.\n");
}

std::vector<double> EnergyMgmt::readEnergyProfile()
{
    std::vector<double> data;
    double temp;
    std::ifstream fin;
    fin.open(_path_energy_profile.c_str());
    /* Read energy profile and store the data into vector. */
    data.resize(0);
    while (fin>>temp) {
        data.push_back(temp);
    }
    /* Reverse the energy harvest queue so that the first energy unit pops first */
    reverse(data.begin(), data.end());
    fin.close();
    return data;
}

void EnergyMgmt::energyHarvest()
{
    /* Add harvested energy into capacity. */
    double energy_val = energy_harvest_data.back();
    consumeEnergy(-energy_val);
    energy_harvest_data.pop_back();

    /* Trigger the next harvest function. */
    schedule(event_energy_harvest, curTick() + time_unit);

    // DPRINTF(EnergyMgmt, "Energy %lf harvested.\n", energy_val);
}

EnergyMgmt *
EnergyMgmtParams::create()
{
    return new EnergyMgmt(this);
}
<|endoftext|>


<commit_before>/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: t; c-basic-offset: 4 -*- vim:set ts=4 sw=4 sts=4 noet: */
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "exec/task-runner.h"

#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <future>

#define BOOST_FILESYSTEM_NO_DEPRECATED
#include <boost/filesystem.hpp>

#include "bc/index.h"
#include "cas/dimension.h"
#include "compiler.h"
#include "db/driver.h"
#include "db/read-only-driver.h"
#include "db/statement-driver.h"
#include "exec.h"
#include "exec/job-runner.h"
#include "exec/parameter.h"
#include "exec/progress.h"
#include "flint/ls.h"
#include "filter.h"
#include "job.h"
#include "layout.h"
#include "lo/layout.h"
#include "lo/layout_loader.h"
#include "load.h"
#include "task.h"

namespace flint {
namespace exec {

namespace {

bool CreateSpec(int id, sqlite3 *db)
{
	char spec_file[64]; // large enough
	std::sprintf(spec_file, "%d/spec.txt", id);
	FILE *fp = fopen(spec_file, "w");
	if (!fp) {
		perror(spec_file);
		return false;
	}
	bool r = task::Spec(id, db, fp);
	fclose(fp);
	return r;
}

int CountParameterSamples(sqlite3 *db)
{
	db::StatementDriver sd(db, "SELECT COUNT(*) FROM parameter_samples");
	int e = sqlite3_step(sd.stmt());
	if (e != SQLITE_ROW) {
		std::cerr << "failed to step statement: " << e << std::endl;
		return -1;
	}
	return sqlite3_column_int(sd.stmt(), 0);
}

const int kFilenameLength = 64;

}

TaskRunner::TaskRunner(int id, char *path)
	: id_(id)
	, path_(path)
	, dir_(new char[kFilenameLength])
	, layout_(new char[kFilenameLength])
	, generated_layout_(new char[kFilenameLength])
{
	std::sprintf(dir_.get(), "%d", id);
	std::sprintf(layout_.get(), "%d/layout", id);
	std::sprintf(generated_layout_.get(), "%d/generated-layout", id);
}

TaskRunner::~TaskRunner() = default;

task::Task *TaskRunner::GetTask()
{
	return task_.get();
}

sqlite3 *TaskRunner::GetDatabase()
{
	return db_driver_->db();
}

sqlite3 *TaskRunner::GetModelDatabase()
{
	return modeldb_driver_->db();
}

const cas::DimensionAnalyzer *TaskRunner::GetDimensionAnalyzer() const
{
	return dimension_analyzer_.get();
}

void *TaskRunner::GetProgressAddress(int job_id)
{
	assert(progress_region_);
	assert(static_cast<size_t>(job_id) < progress_region_->get_size());
	char *addr = static_cast<char *>(progress_region_->get_address());
	return addr + job_id;
}

bool TaskRunner::Setup(int id, const char *path, std::vector<double> *data)
{
	task_.reset(load::Load(path, load::kExec, id, data));
	if (!task_)
		return false;
	{
		db::Driver driver("x.db");
		auto db = driver.db();
		if (!db)
			return false;
		if (!task::Config(id, db))
			return false;
	}
	db::ReadOnlyDriver driver("x.db");
	auto db = driver.db();
	if (!db)
		return false;
	return CreateSpec(id, db);
}

bool TaskRunner::Run()
{
	if (!Setup(id_, path_.get(), &data_))
		return false;

	{
		char canceled_file[kFilenameLength];
		std::sprintf(canceled_file, "%s/canceled", dir_.get());
		if (boost::filesystem::exists(canceled_file)) {
			// exit early if file "canceled" exists
			return true;
		}
	}

	char modeldb_file[kFilenameLength];
	std::sprintf(modeldb_file, "%s/model.db", dir_.get());
	modeldb_driver_.reset(new db::ReadOnlyDriver(modeldb_file));

	char spec_file[kFilenameLength];
	std::sprintf(spec_file, "%s/spec.txt", dir_.get());
	char filter_file[kFilenameLength];
	std::sprintf(filter_file, "%s/filter", dir_.get());
	if (!filter::Create(modeldb_driver_->db(), spec_file, layout_.get(), filter_file))
		return false;
	char track_file[kFilenameLength];
	std::sprintf(track_file, "%s/track", dir_.get());
	if (!filter::Track(filter_file, track_file))
		return false;
	char isdh_file[kFilenameLength];
	std::sprintf(isdh_file, "%s/isdh", dir_.get());
	if (!filter::Isdh(filter_file, isdh_file))
		return false;
	reader_.reset(new task::ConfigReader(modeldb_driver_->db()));
	if (!reader_->Read())
		return false;
	{
		std::unique_ptr<Layout> layout(new Layout);
		LayoutLoader loader(layout_.get());
		if (!loader.Load(layout.get()))
			return false;
		size_t layer_size = layout->Calculate();
		assert(layer_size > kOffsetBase);
		task_->layout.swap(layout);
		task_->layer_size = layer_size;

		if (reader_->GetDpsPath())
			task_->ls_config = ls::CreateConfiguration(reader_->GetDpsPath(), *task_->layout);
	}
	if (reader_->GetMethod() != compiler::Method::kArk) {
		cas::DimensionAnalyzer da;
		if (!da.Load(modeldb_driver_->db()))
			return false;
		compiler::Compiler c(&da);
		task_->reinit_bc.reset(c.Compile(modeldb_driver_->db(), "dependent_ivs", compiler::Method::kAssign));
		if (!task_->reinit_bc)
			return false;
		task_->bc.reset(c.Compile(modeldb_driver_->db(), "input_eqs", reader_->GetMethod()));
		if (!task_->bc)
			return false;
	}

	char db_file[kFilenameLength];
	std::sprintf(db_file, "%s/task.db", dir_.get());
	db_driver_.reset(new db::Driver(db_file));
	if (!exec::SaveParameters(id_, db_driver_->db()))
		return false;
	int n = CountParameterSamples(db_driver_->db());
	if (n <= 0)
		return false;
	std::unique_ptr<boost::interprocess::file_mapping> fm(exec::CreateProgressFile(n, dir_.get()));
	if (!fm)
		return false;
	progress_region_.reset(new boost::interprocess::mapped_region(*fm, boost::interprocess::read_write));
	if (!task::Form(db_driver_->db()))
		return false;
	dimension_analyzer_.reset(new cas::DimensionAnalyzer);
	if (!dimension_analyzer_->Load(db_driver_->db()))
		return false;
	if (!layout::Generate(db_driver_->db(), generated_layout_.get()))
		return false;
	std::vector<std::future<bool> > v;
	for (;;) {
		int job_id;
		if (!job::Generate(db_driver_->db(), dir_.get(), &job_id))
			return false;
		if (job_id == 0)
			break; // all jobs has been generated

		auto lmbd = [](TaskRunner *tr, int id){
			JobRunner runner(tr, id);
			return runner.Run();
		};
		v.emplace_back(std::async(std::launch::async, lmbd, this, job_id));
	}
	assert(static_cast<size_t>(n) == v.size());
	return MonitorTaskProgress(v, progress_region_.get());
}

}
}
<commit_msg>Return false immediately when failing configuration<commit_after>/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: t; c-basic-offset: 4 -*- vim:set ts=4 sw=4 sts=4 noet: */
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "exec/task-runner.h"

#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <future>

#define BOOST_FILESYSTEM_NO_DEPRECATED
#include <boost/filesystem.hpp>

#include "bc/index.h"
#include "cas/dimension.h"
#include "compiler.h"
#include "db/driver.h"
#include "db/read-only-driver.h"
#include "db/statement-driver.h"
#include "exec.h"
#include "exec/job-runner.h"
#include "exec/parameter.h"
#include "exec/progress.h"
#include "flint/ls.h"
#include "filter.h"
#include "job.h"
#include "layout.h"
#include "lo/layout.h"
#include "lo/layout_loader.h"
#include "load.h"
#include "task.h"

namespace flint {
namespace exec {

namespace {

bool CreateSpec(int id, sqlite3 *db)
{
	char spec_file[64]; // large enough
	std::sprintf(spec_file, "%d/spec.txt", id);
	FILE *fp = fopen(spec_file, "w");
	if (!fp) {
		perror(spec_file);
		return false;
	}
	bool r = task::Spec(id, db, fp);
	fclose(fp);
	return r;
}

int CountParameterSamples(sqlite3 *db)
{
	db::StatementDriver sd(db, "SELECT COUNT(*) FROM parameter_samples");
	int e = sqlite3_step(sd.stmt());
	if (e != SQLITE_ROW) {
		std::cerr << "failed to step statement: " << e << std::endl;
		return -1;
	}
	return sqlite3_column_int(sd.stmt(), 0);
}

const int kFilenameLength = 64;

}

TaskRunner::TaskRunner(int id, char *path)
	: id_(id)
	, path_(path)
	, dir_(new char[kFilenameLength])
	, layout_(new char[kFilenameLength])
	, generated_layout_(new char[kFilenameLength])
{
	std::sprintf(dir_.get(), "%d", id);
	std::sprintf(layout_.get(), "%d/layout", id);
	std::sprintf(generated_layout_.get(), "%d/generated-layout", id);
}

TaskRunner::~TaskRunner() = default;

task::Task *TaskRunner::GetTask()
{
	return task_.get();
}

sqlite3 *TaskRunner::GetDatabase()
{
	return db_driver_->db();
}

sqlite3 *TaskRunner::GetModelDatabase()
{
	return modeldb_driver_->db();
}

const cas::DimensionAnalyzer *TaskRunner::GetDimensionAnalyzer() const
{
	return dimension_analyzer_.get();
}

void *TaskRunner::GetProgressAddress(int job_id)
{
	assert(progress_region_);
	assert(static_cast<size_t>(job_id) < progress_region_->get_size());
	char *addr = static_cast<char *>(progress_region_->get_address());
	return addr + job_id;
}

bool TaskRunner::Setup(int id, const char *path, std::vector<double> *data)
{
	task_.reset(load::Load(path, load::kExec, id, data));
	if (!task_)
		return false;
	{
		db::Driver driver("x.db");
		auto db = driver.db();
		if (!db)
			return false;
		if (!task::Config(id, db))
			return false;
	}
	db::ReadOnlyDriver driver("x.db");
	auto db = driver.db();
	if (!db)
		return false;
	return CreateSpec(id, db);
}

bool TaskRunner::Run()
{
	if (!Setup(id_, path_.get(), &data_))
		return false;

	{
		char canceled_file[kFilenameLength];
		std::sprintf(canceled_file, "%s/canceled", dir_.get());
		if (boost::filesystem::exists(canceled_file)) {
			// exit early if file "canceled" exists
			return true;
		}
	}

	char modeldb_file[kFilenameLength];
	std::sprintf(modeldb_file, "%s/model.db", dir_.get());
	modeldb_driver_.reset(new db::ReadOnlyDriver(modeldb_file));

	char spec_file[kFilenameLength];
	std::sprintf(spec_file, "%s/spec.txt", dir_.get());
	char filter_file[kFilenameLength];
	std::sprintf(filter_file, "%s/filter", dir_.get());
	if (!filter::Create(modeldb_driver_->db(), spec_file, layout_.get(), filter_file))
		return false;
	char track_file[kFilenameLength];
	std::sprintf(track_file, "%s/track", dir_.get());
	if (!filter::Track(filter_file, track_file))
		return false;
	char isdh_file[kFilenameLength];
	std::sprintf(isdh_file, "%s/isdh", dir_.get());
	if (!filter::Isdh(filter_file, isdh_file))
		return false;
	reader_.reset(new task::ConfigReader(modeldb_driver_->db()));
	if (!reader_->Read())
		return false;
	{
		std::unique_ptr<Layout> layout(new Layout);
		LayoutLoader loader(layout_.get());
		if (!loader.Load(layout.get()))
			return false;
		size_t layer_size = layout->Calculate();
		assert(layer_size > kOffsetBase);
		task_->layout.swap(layout);
		task_->layer_size = layer_size;

		if (reader_->GetDpsPath()) {
			auto ls_config = ls::CreateConfiguration(reader_->GetDpsPath(), *task_->layout);
			if (!ls_config)
				return false;
			task_->ls_config.swap(ls_config);
		}
	}
	if (reader_->GetMethod() != compiler::Method::kArk) {
		cas::DimensionAnalyzer da;
		if (!da.Load(modeldb_driver_->db()))
			return false;
		compiler::Compiler c(&da);
		task_->reinit_bc.reset(c.Compile(modeldb_driver_->db(), "dependent_ivs", compiler::Method::kAssign));
		if (!task_->reinit_bc)
			return false;
		task_->bc.reset(c.Compile(modeldb_driver_->db(), "input_eqs", reader_->GetMethod()));
		if (!task_->bc)
			return false;
	}

	char db_file[kFilenameLength];
	std::sprintf(db_file, "%s/task.db", dir_.get());
	db_driver_.reset(new db::Driver(db_file));
	if (!exec::SaveParameters(id_, db_driver_->db()))
		return false;
	int n = CountParameterSamples(db_driver_->db());
	if (n <= 0)
		return false;
	std::unique_ptr<boost::interprocess::file_mapping> fm(exec::CreateProgressFile(n, dir_.get()));
	if (!fm)
		return false;
	progress_region_.reset(new boost::interprocess::mapped_region(*fm, boost::interprocess::read_write));
	if (!task::Form(db_driver_->db()))
		return false;
	dimension_analyzer_.reset(new cas::DimensionAnalyzer);
	if (!dimension_analyzer_->Load(db_driver_->db()))
		return false;
	if (!layout::Generate(db_driver_->db(), generated_layout_.get()))
		return false;
	std::vector<std::future<bool> > v;
	for (;;) {
		int job_id;
		if (!job::Generate(db_driver_->db(), dir_.get(), &job_id))
			return false;
		if (job_id == 0)
			break; // all jobs has been generated

		auto lmbd = [](TaskRunner *tr, int id){
			JobRunner runner(tr, id);
			return runner.Run();
		};
		v.emplace_back(std::async(std::launch::async, lmbd, this, job_id));
	}
	assert(static_cast<size_t>(n) == v.size());
	return MonitorTaskProgress(v, progress_region_.get());
}

}
}
<|endoftext|>


<commit_before>// Copyright (c) 2009-2012 The Bitcoin developers
// Distributed under the MIT/X11 software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <map>

#include <openssl/ecdsa.h>
#include <openssl/obj_mac.h>

#include "headers.h"
#include "key.h"

// Generate a private key from just the secret parameter
int EC_KEY_regenerate_key(EC_KEY *eckey, BIGNUM *priv_key)
{
    int ok = 0;
    BN_CTX *ctx = NULL;
    EC_POINT *pub_key = NULL;

    if (!eckey) return 0;

    const EC_GROUP *group = EC_KEY_get0_group(eckey);

    if ((ctx = BN_CTX_new()) == NULL)
        goto err;

    pub_key = EC_POINT_new(group);

    if (pub_key == NULL)
        goto err;

    if (!EC_POINT_mul(group, pub_key, priv_key, NULL, NULL, ctx))
        goto err;

    EC_KEY_set_private_key(eckey,priv_key);
    EC_KEY_set_public_key(eckey,pub_key);

    ok = 1;

err:

    if (pub_key)
        EC_POINT_free(pub_key);
    if (ctx != NULL)
        BN_CTX_free(ctx);

    return(ok);
}

// Perform ECDSA key recovery (see SEC1 4.1.6) for curves over (mod p)-fields
// recid selects which key is recovered
// if check is non-zero, additional checks are performed
int ECDSA_SIG_recover_key_GFp(EC_KEY *eckey, ECDSA_SIG *ecsig, const unsigned char *msg, int msglen, int recid, int check)
{
    if (!eckey) return 0;

    int ret = 0;
    BN_CTX *ctx = NULL;

    BIGNUM *x = NULL;
    BIGNUM *e = NULL;
    BIGNUM *order = NULL;
    BIGNUM *sor = NULL;
    BIGNUM *eor = NULL;
    BIGNUM *field = NULL;
    EC_POINT *R = NULL;
    EC_POINT *O = NULL;
    EC_POINT *Q = NULL;
    BIGNUM *rr = NULL;
    BIGNUM *zero = NULL;
    int n = 0;
    int i = recid / 2;

    const EC_GROUP *group = EC_KEY_get0_group(eckey);
    if ((ctx = BN_CTX_new()) == NULL) { ret = -1; goto err; }
    BN_CTX_start(ctx);
    order = BN_CTX_get(ctx);
    if (!EC_GROUP_get_order(group, order, ctx)) { ret = -2; goto err; }
    x = BN_CTX_get(ctx);
    if (!BN_copy(x, order)) { ret=-1; goto err; }
    if (!BN_mul_word(x, i)) { ret=-1; goto err; }
    if (!BN_add(x, x, ecsig->r)) { ret=-1; goto err; }
    field = BN_CTX_get(ctx);
    if (!EC_GROUP_get_curve_GFp(group, field, NULL, NULL, ctx)) { ret=-2; goto err; }
    if (BN_cmp(x, field) >= 0) { ret=0; goto err; }
    if ((R = EC_POINT_new(group)) == NULL) { ret = -2; goto err; }
    if (!EC_POINT_set_compressed_coordinates_GFp(group, R, x, recid % 2, ctx)) { ret=0; goto err; }
    if (check)
    {
        if ((O = EC_POINT_new(group)) == NULL) { ret = -2; goto err; }
        if (!EC_POINT_mul(group, O, NULL, R, order, ctx)) { ret=-2; goto err; }
        if (!EC_POINT_is_at_infinity(group, O)) { ret = 0; goto err; }
    }
    if ((Q = EC_POINT_new(group)) == NULL) { ret = -2; goto err; }
    n = EC_GROUP_get_degree(group);
    e = BN_CTX_get(ctx);
    if (!BN_bin2bn(msg, msglen, e)) { ret=-1; goto err; }
    if (8*msglen > n) BN_rshift(e, e, 8-(n & 7));
    zero = BN_CTX_get(ctx);
    if (!BN_zero(zero)) { ret=-1; goto err; }
    if (!BN_mod_sub(e, zero, e, order, ctx)) { ret=-1; goto err; }
    rr = BN_CTX_get(ctx);
    if (!BN_mod_inverse(rr, ecsig->r, order, ctx)) { ret=-1; goto err; }
    sor = BN_CTX_get(ctx);
    if (!BN_mod_mul(sor, ecsig->s, rr, order, ctx)) { ret=-1; goto err; }
    eor = BN_CTX_get(ctx);
    if (!BN_mod_mul(eor, e, rr, order, ctx)) { ret=-1; goto err; }
    if (!EC_POINT_mul(group, Q, eor, R, sor, ctx)) { ret=-2; goto err; }
    if (!EC_KEY_set_public_key(eckey, Q)) { ret=-2; goto err; }

    ret = 1;

err:
    if (ctx) {
        BN_CTX_end(ctx);
        BN_CTX_free(ctx);
    }
    if (R != NULL) EC_POINT_free(R);
    if (O != NULL) EC_POINT_free(O);
    if (Q != NULL) EC_POINT_free(Q);
    return ret;
}

void CKey::SetCompressedPubKey(bool fCompressed)
{
    EC_KEY_set_conv_form(pkey, fCompressed ? POINT_CONVERSION_COMPRESSED : POINT_CONVERSION_UNCOMPRESSED);
    fCompressedPubKey = true;
}

// create a compact signature (65 bytes), which allows reconstructing the used public key
// The format is one header byte, followed by two times 32 bytes for the serialized r and s values.
// The header byte: 0x1B = first key with even y, 0x1C = first key with odd y,
//                  0x1D = second key with even y, 0x1E = second key with odd y
bool CKey::SignCompact(uint256 hash, std::vector<unsigned char>& vchSig)
{
    bool fOk = false;
    ECDSA_SIG *sig = ECDSA_do_sign((unsigned char*)&hash, sizeof(hash), pkey);
    if (sig==NULL)
        return false;
    vchSig.clear();
    vchSig.resize(65,0);
    int nBitsR = BN_num_bits(sig->r);
    int nBitsS = BN_num_bits(sig->s);
    if (nBitsR <= 256 && nBitsS <= 256)
    {
        int nRecId = -1;
        for (int i=0; i<4; i++)
        {
            CKey keyRec;
            keyRec.fSet = true;
            if (fCompressedPubKey)
                keyRec.SetCompressedPubKey();
            if (ECDSA_SIG_recover_key_GFp(keyRec.pkey, sig, (unsigned char*)&hash, sizeof(hash), i, 1) == 1)
                if (keyRec.GetPubKey() == this->GetPubKey())
                {
                    nRecId = i;
                    break;
                }
        }

        if (nRecId == -1)
            throw key_error("CKey::SignCompact() : unable to construct recoverable key");

        vchSig[0] = nRecId+27+(fCompressedPubKey ? 4 : 0);
        BN_bn2bin(sig->r,&vchSig[33-(nBitsR+7)/8]);
        BN_bn2bin(sig->s,&vchSig[65-(nBitsS+7)/8]);
        fOk = true;
    }
    ECDSA_SIG_free(sig);
    return fOk;
}

// reconstruct public key from a compact signature
// This is only slightly more CPU intensive than just verifying it.
// If this function succeeds, the recovered public key is guaranteed to be valid
// (the signature is a valid signature of the given data for that key)
bool CKey::SetCompactSignature(uint256 hash, const std::vector<unsigned char>& vchSig)
{
    if (vchSig.size() != 65)
        return false;
    int nV = vchSig[0];
    if (nV<27 || nV>=35)
        return false;
    ECDSA_SIG *sig = ECDSA_SIG_new();
    BN_bin2bn(&vchSig[1],32,sig->r);
    BN_bin2bn(&vchSig[33],32,sig->s);

    EC_KEY_free(pkey);
    pkey = EC_KEY_new_by_curve_name(NID_secp256k1);
    if (nV >= 31)
    {
        SetCompressedPubKey();
        nV -= 4;
    }
    if (ECDSA_SIG_recover_key_GFp(pkey, sig, (unsigned char*)&hash, sizeof(hash), nV - 27, 0) == 1)
    {
        fSet = true;
        ECDSA_SIG_free(sig);
        return true;
    }
    return false;
}

bool CKey::VerifyCompact(uint256 hash, const std::vector<unsigned char>& vchSig)
{
    CKey key;
    if (!key.SetCompactSignature(hash, vchSig))
        return false;
    if (GetPubKey() != key.GetPubKey())
        return false;

    return true;
}

bool CKey::SetSecret(const CSecret& vchSecret, bool fCompressed)
{
    EC_KEY_free(pkey);
    pkey = EC_KEY_new_by_curve_name(NID_secp256k1);
    if (pkey == NULL)
        throw key_error("CKey::SetSecret() : EC_KEY_new_by_curve_name failed");
    if (vchSecret.size() != 32)
        throw key_error("CKey::SetSecret() : secret must be 32 bytes");
    BIGNUM *bn = BN_bin2bn(&vchSecret[0],32,BN_new());
    if (bn == NULL)
        throw key_error("CKey::SetSecret() : BN_bin2bn failed");
    if (!EC_KEY_regenerate_key(pkey,bn))
    {
        BN_clear_free(bn);
        throw key_error("CKey::SetSecret() : EC_KEY_regenerate_key failed");
    }
    BN_clear_free(bn);
    fSet = true;
    if (fCompressed || fCompressedPubKey)
        SetCompressedPubKey();
    return true;
}

CSecret32 CKey::GetSecret(bool &fCompressed) const
{
    CSecret vchRet;
    vchRet.resize(32);
    const BIGNUM *bn = EC_KEY_get0_private_key(pkey);
    int nBytes = BN_num_bytes(bn);
    if (bn == NULL)
        throw key_error("CKey::GetSecret() : EC_KEY_get0_private_key failed");
    int n=BN_bn2bin(bn,&vchRet[32 - nBytes]);
    if (n != nBytes)
        throw key_error("CKey::GetSecret(): BN_bn2bin failed");
    fCompressed = fCompressedPubKey;
    return vchRet;
}
<commit_msg>OpenSSL Fix<commit_after>// Copyright (c) 2009-2012 The Bitcoin developers
// Distributed under the MIT/X11 software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <map>

#include <openssl/ecdsa.h>
#include <openssl/obj_mac.h>

#include "headers.h"
#include "key.h"

// Generate a private key from just the secret parameter
int EC_KEY_regenerate_key(EC_KEY *eckey, BIGNUM *priv_key)
{
    int ok = 0;
    BN_CTX *ctx = NULL;
    EC_POINT *pub_key = NULL;

    if (!eckey) return 0;

    const EC_GROUP *group = EC_KEY_get0_group(eckey);

    if ((ctx = BN_CTX_new()) == NULL)
        goto err;

    pub_key = EC_POINT_new(group);

    if (pub_key == NULL)
        goto err;

    if (!EC_POINT_mul(group, pub_key, priv_key, NULL, NULL, ctx))
        goto err;

    EC_KEY_set_private_key(eckey,priv_key);
    EC_KEY_set_public_key(eckey,pub_key);

    ok = 1;

err:

    if (pub_key)
        EC_POINT_free(pub_key);
    if (ctx != NULL)
        BN_CTX_free(ctx);

    return(ok);
}

// Perform ECDSA key recovery (see SEC1 4.1.6) for curves over (mod p)-fields
// recid selects which key is recovered
// if check is non-zero, additional checks are performed
int ECDSA_SIG_recover_key_GFp(EC_KEY *eckey, ECDSA_SIG *ecsig, const unsigned char *msg, int msglen, int recid, int check)
{
    if (!eckey) return 0;

    int ret = 0;
    BN_CTX *ctx = NULL;

    BIGNUM *x = NULL;
    BIGNUM *e = NULL;
    BIGNUM *order = NULL;
    BIGNUM *sor = NULL;
    BIGNUM *eor = NULL;
    BIGNUM *field = NULL;
    EC_POINT *R = NULL;
    EC_POINT *O = NULL;
    EC_POINT *Q = NULL;
    BIGNUM *rr = NULL;
    BIGNUM *zero = NULL;
    int n = 0;
    int i = recid / 2;

    const EC_GROUP *group = EC_KEY_get0_group(eckey);
    if ((ctx = BN_CTX_new()) == NULL) { ret = -1; goto err; }
    BN_CTX_start(ctx);
    order = BN_CTX_get(ctx);
    if (!EC_GROUP_get_order(group, order, ctx)) { ret = -2; goto err; }
    x = BN_CTX_get(ctx);
    if (!BN_copy(x, order)) { ret=-1; goto err; }
    if (!BN_mul_word(x, i)) { ret=-1; goto err; }
    if (!BN_add(x, x, ecsig->r)) { ret=-1; goto err; }
    field = BN_CTX_get(ctx);
    if (!EC_GROUP_get_curve_GFp(group, field, NULL, NULL, ctx)) { ret=-2; goto err; }
    if (BN_cmp(x, field) >= 0) { ret=0; goto err; }
    if ((R = EC_POINT_new(group)) == NULL) { ret = -2; goto err; }
    if (!EC_POINT_set_compressed_coordinates_GFp(group, R, x, recid % 2, ctx)) { ret=0; goto err; }
    if (check)
    {
        if ((O = EC_POINT_new(group)) == NULL) { ret = -2; goto err; }
        if (!EC_POINT_mul(group, O, NULL, R, order, ctx)) { ret=-2; goto err; }
        if (!EC_POINT_is_at_infinity(group, O)) { ret = 0; goto err; }
    }
    if ((Q = EC_POINT_new(group)) == NULL) { ret = -2; goto err; }
    n = EC_GROUP_get_degree(group);
    e = BN_CTX_get(ctx);
    if (!BN_bin2bn(msg, msglen, e)) { ret=-1; goto err; }
    if (8*msglen > n) BN_rshift(e, e, 8-(n & 7));
    zero = BN_CTX_get(ctx);
    if (!BN_zero(zero)) { ret=-1; goto err; }
    if (!BN_mod_sub(e, zero, e, order, ctx)) { ret=-1; goto err; }
    rr = BN_CTX_get(ctx);
    if (!BN_mod_inverse(rr, ecsig->r, order, ctx)) { ret=-1; goto err; }
    sor = BN_CTX_get(ctx);
    if (!BN_mod_mul(sor, ecsig->s, rr, order, ctx)) { ret=-1; goto err; }
    eor = BN_CTX_get(ctx);
    if (!BN_mod_mul(eor, e, rr, order, ctx)) { ret=-1; goto err; }
    if (!EC_POINT_mul(group, Q, eor, R, sor, ctx)) { ret=-2; goto err; }
    if (!EC_KEY_set_public_key(eckey, Q)) { ret=-2; goto err; }

    ret = 1;

err:
    if (ctx) {
        BN_CTX_end(ctx);
        BN_CTX_free(ctx);
    }
    if (R != NULL) EC_POINT_free(R);
    if (O != NULL) EC_POINT_free(O);
    if (Q != NULL) EC_POINT_free(Q);
    return ret;
}

void CKey::SetCompressedPubKey(bool fCompressed)
{
    EC_KEY_set_conv_form(pkey, fCompressed ? POINT_CONVERSION_COMPRESSED : POINT_CONVERSION_UNCOMPRESSED);
    fCompressedPubKey = true;
}

// create a compact signature (65 bytes), which allows reconstructing the used public key
// The format is one header byte, followed by two times 32 bytes for the serialized r and s values.
// The header byte: 0x1B = first key with even y, 0x1C = first key with odd y,
//                  0x1D = second key with even y, 0x1E = second key with odd y
bool CKey::SignCompact(uint256 hash, std::vector<unsigned char>& vchSig)
{
    bool fOk = false;
    ECDSA_SIG *sig = ECDSA_do_sign((unsigned char*)&hash, sizeof(hash), pkey);
    if (sig==NULL)
        return false;
    vchSig.clear();
    vchSig.resize(65,0);
    int nBitsR = BN_num_bits(sig->r);
    int nBitsS = BN_num_bits(sig->s);
    if (nBitsR <= 256 && nBitsS <= 256)
    {
        int nRecId = -1;
        for (int i=0; i<4; i++)
        {
            CKey keyRec;
            keyRec.fSet = true;
            if (fCompressedPubKey)
                keyRec.SetCompressedPubKey();
            if (ECDSA_SIG_recover_key_GFp(keyRec.pkey, sig, (unsigned char*)&hash, sizeof(hash), i, 1) == 1)
                if (keyRec.GetPubKey() == this->GetPubKey())
                {
                    nRecId = i;
                    break;
                }
        }

        if (nRecId == -1)
            throw key_error("CKey::SignCompact() : unable to construct recoverable key");

        vchSig[0] = nRecId+27+(fCompressedPubKey ? 4 : 0);
        BN_bn2bin(sig->r,&vchSig[33-(nBitsR+7)/8]);
        BN_bn2bin(sig->s,&vchSig[65-(nBitsS+7)/8]);
        fOk = true;
    }
    ECDSA_SIG_free(sig);
    return fOk;
}

// reconstruct public key from a compact signature
// This is only slightly more CPU intensive than just verifying it.
// If this function succeeds, the recovered public key is guaranteed to be valid
// (the signature is a valid signature of the given data for that key)
bool CKey::SetCompactSignature(uint256 hash, const std::vector<unsigned char>& vchSig)
{
    if (vchSig.size() != 65)
        return false;
    int nV = vchSig[0];
    if (nV<27 || nV>=35)
        return false;
    ECDSA_SIG *sig = ECDSA_SIG_new();
    BN_bin2bn(&vchSig[1],32,sig->r);
    BN_bin2bn(&vchSig[33],32,sig->s);

    EC_KEY_free(pkey);
    pkey = EC_KEY_new_by_curve_name(NID_secp256k1);
    if (nV >= 31)
    {
        SetCompressedPubKey();
        nV -= 4;
    }
    if (ECDSA_SIG_recover_key_GFp(pkey, sig, (unsigned char*)&hash, sizeof(hash), nV - 27, 0) == 1)
    {
        fSet = true;
        ECDSA_SIG_free(sig);
        return true;
    }
    return false;
}

bool CKey::VerifyCompact(uint256 hash, const std::vector<unsigned char>& vchSig)
{
    if (vchSig.empty())
        return false;

    // New versions of OpenSSL will reject non-canonical DER signatures. de/re-serialize first.
    unsigned char *norm_der = NULL;
    ECDSA_SIG *norm_sig = ECDSA_SIG_new();
    const unsigned char* sigptr = &vchSig[0];
    assert(norm_sig);
    if (d2i_ECDSA_SIG(&norm_sig, &sigptr, vchSig.size()) == NULL)
    {
        /* As of OpenSSL 1.0.0p d2i_ECDSA_SIG frees and nulls the pointer on
         * error. But OpenSSL's own use of this function redundantly frees the
         * result. As ECDSA_SIG_free(NULL) is a no-op, and in the absence of a
         * clear contract for the function behaving the same way is more
         * conservative.
         */
         ECDSA_SIG_free(norm_sig);
         return false;
    }
    int derlen = i2d_ECDSA_SIG(norm_sig, &norm_der);
    ECDSA_SIG_free(norm_sig);
    if (derlen <= 0)
        return false;
    // -1 = error, 0 = bad sig, 1 = good
    bool ret = ECDSA_verify(0, (unsigned char*)&hash, sizeof(hash), norm_der, derlen, pkey) == 1;
    OPENSSL_free(norm_der);
    return ret;
}

bool CKey::SetSecret(const CSecret& vchSecret, bool fCompressed)
{
    EC_KEY_free(pkey);
    pkey = EC_KEY_new_by_curve_name(NID_secp256k1);
    if (pkey == NULL)
        throw key_error("CKey::SetSecret() : EC_KEY_new_by_curve_name failed");
    if (vchSecret.size() != 32)
        throw key_error("CKey::SetSecret() : secret must be 32 bytes");
    BIGNUM *bn = BN_bin2bn(&vchSecret[0],32,BN_new());
    if (bn == NULL)
        throw key_error("CKey::SetSecret() : BN_bin2bn failed");
    if (!EC_KEY_regenerate_key(pkey,bn))
    {
        BN_clear_free(bn);
        throw key_error("CKey::SetSecret() : EC_KEY_regenerate_key failed");
    }
    BN_clear_free(bn);
    fSet = true;
    if (fCompressed || fCompressedPubKey)
        SetCompressedPubKey();
    return true;
}

CSecret32 CKey::GetSecret(bool &fCompressed) const
{
    CSecret vchRet;
    vchRet.resize(32);
    const BIGNUM *bn = EC_KEY_get0_private_key(pkey);
    int nBytes = BN_num_bytes(bn);
    if (bn == NULL)
        throw key_error("CKey::GetSecret() : EC_KEY_get0_private_key failed");
    int n=BN_bn2bin(bn,&vchRet[32 - nBytes]);
    if (n != nBytes)
        throw key_error("CKey::GetSecret(): BN_bn2bin failed");
    fCompressed = fCompressedPubKey;
    return vchRet;
}
<|endoftext|>


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

#include "host/desktop_session_manager.h"

#include "base/location.h"
#include "base/logging.h"
#include "base/task_runner.h"
#include "desktop/desktop_frame.h"
#include "host/desktop_session_fake.h"
#include "host/desktop_session_ipc.h"
#include "host/desktop_session_process.h"
#include "host/desktop_session_proxy.h"
#include "ipc/ipc_channel.h"

namespace host {

DesktopSessionManager::DesktopSessionManager(
    std::shared_ptr<base::TaskRunner> task_runner, DesktopSession::Delegate* delegate)
    : task_runner_(task_runner),
      session_attach_timer_(task_runner),
      session_proxy_(std::make_shared<DesktopSessionProxy>()),
      delegate_(delegate)
{
    // Nothing
}

DesktopSessionManager::~DesktopSessionManager()
{
    state_ = State::STOPPING;
    dettachSession(FROM_HERE);
}

void DesktopSessionManager::attachSession(
    const base::Location& location, base::win::SessionId session_id)
{
    if (state_ == State::ATTACHED)
        return;

    LOG(LS_INFO) << "Attach session with ID: " << session_id
                 << " (from: " << location.toString() << ")";

    if (state_ == State::STOPPED)
    {
        session_attach_timer_.start(std::chrono::minutes(1), [this]()
        {
            LOG(LS_WARNING) << "Session attach timeout";
            onErrorOccurred();
        });
    }

    state_ = State::STARTING;

    std::u16string channel_id = ipc::Server::createUniqueId();

    server_ = std::make_unique<ipc::Server>();
    if (!server_->start(channel_id, this))
    {
        LOG(LS_ERROR) << "Failed to start IPC server";

        onErrorOccurred();
        return;
    }

    std::unique_ptr<DesktopSessionProcess> process =
        DesktopSessionProcess::create(session_id, channel_id);
    if (!process)
    {
        LOG(LS_ERROR) << "Failed to create session process";

        onErrorOccurred();
        return;
    }
}

void DesktopSessionManager::dettachSession(const base::Location& location)
{
    if (state_ == State::STOPPED || state_ == State::DETACHED)
        return;

    LOG(LS_INFO) << "Dettach session (from: " << location.toString() << ")";

    if (state_ != State::STOPPING)
        state_ = State::DETACHED;

    session_attach_timer_.stop();
    session_proxy_->dettach();
    session_.reset();

    LOG(LS_INFO) << "Session process is detached";

    if (state_ == State::STOPPING)
        return;

    session_attach_timer_.start(std::chrono::minutes(1), [this]()
    {
        LOG(LS_ERROR) << "Timeout while waiting for session";
        onErrorOccurred();
    });

    // The real session process has ended. We create a temporary fake session.
    session_ = std::make_unique<DesktopSessionFake>(this);
    session_proxy_->attach(session_.get());
    session_->start();
}

std::shared_ptr<DesktopSessionProxy> DesktopSessionManager::sessionProxy() const
{
    return session_proxy_;
}

void DesktopSessionManager::onNewConnection(std::unique_ptr<ipc::Channel> channel)
{
    if (DesktopSessionProcess::filePath() != channel->peerFilePath())
    {
        LOG(LS_ERROR) << "An attempt was made to connect from an unknown application";
        return;
    }

    LOG(LS_INFO) << "Session process successfully connected";

    session_attach_timer_.stop();
    task_runner_->deleteSoon(std::move(server_));

    session_ = std::make_unique<DesktopSessionIpc>(std::move(channel), this);
    session_proxy_->attach(session_.get());

    state_ = State::ATTACHED;
    session_->start();
}

void DesktopSessionManager::onErrorOccurred()
{
    if (state_ == State::STOPPED || state_ == State::STOPPING)
        return;

    state_ = State::STOPPING;
    dettachSession(FROM_HERE);
    state_ = State::STOPPED;
}

void DesktopSessionManager::onDesktopSessionStarted()
{
    delegate_->onDesktopSessionStarted();
}

void DesktopSessionManager::onDesktopSessionStopped()
{
    dettachSession(FROM_HERE);
}

void DesktopSessionManager::onScreenCaptured(const desktop::Frame& frame)
{
    delegate_->onScreenCaptured(frame);
}

void DesktopSessionManager::onCursorCaptured(std::shared_ptr<desktop::MouseCursor> mouse_cursor)
{
    delegate_->onCursorCaptured(std::move(mouse_cursor));
}

void DesktopSessionManager::onScreenListChanged(const proto::ScreenList& list)
{
    delegate_->onScreenListChanged(list);
}

void DesktopSessionManager::onClipboardEvent(const proto::ClipboardEvent& event)
{
    delegate_->onClipboardEvent(event);
}

} // namespace host
<commit_msg>- Delete session over message loop.<commit_after>//
// Aspia Project
// Copyright (C) 2020 Dmitry Chapyshev <dmitry@aspia.ru>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.
//

#include "host/desktop_session_manager.h"

#include "base/location.h"
#include "base/logging.h"
#include "base/task_runner.h"
#include "desktop/desktop_frame.h"
#include "host/desktop_session_fake.h"
#include "host/desktop_session_ipc.h"
#include "host/desktop_session_process.h"
#include "host/desktop_session_proxy.h"
#include "ipc/ipc_channel.h"

namespace host {

DesktopSessionManager::DesktopSessionManager(
    std::shared_ptr<base::TaskRunner> task_runner, DesktopSession::Delegate* delegate)
    : task_runner_(task_runner),
      session_attach_timer_(task_runner),
      session_proxy_(std::make_shared<DesktopSessionProxy>()),
      delegate_(delegate)
{
    // Nothing
}

DesktopSessionManager::~DesktopSessionManager()
{
    state_ = State::STOPPING;
    dettachSession(FROM_HERE);
}

void DesktopSessionManager::attachSession(
    const base::Location& location, base::win::SessionId session_id)
{
    if (state_ == State::ATTACHED)
        return;

    LOG(LS_INFO) << "Attach session with ID: " << session_id
                 << " (from: " << location.toString() << ")";

    if (state_ == State::STOPPED)
    {
        session_attach_timer_.start(std::chrono::minutes(1), [this]()
        {
            LOG(LS_WARNING) << "Session attach timeout";
            onErrorOccurred();
        });
    }

    state_ = State::STARTING;

    std::u16string channel_id = ipc::Server::createUniqueId();

    server_ = std::make_unique<ipc::Server>();
    if (!server_->start(channel_id, this))
    {
        LOG(LS_ERROR) << "Failed to start IPC server";

        onErrorOccurred();
        return;
    }

    std::unique_ptr<DesktopSessionProcess> process =
        DesktopSessionProcess::create(session_id, channel_id);
    if (!process)
    {
        LOG(LS_ERROR) << "Failed to create session process";

        onErrorOccurred();
        return;
    }
}

void DesktopSessionManager::dettachSession(const base::Location& location)
{
    if (state_ == State::STOPPED || state_ == State::DETACHED)
        return;

    LOG(LS_INFO) << "Dettach session (from: " << location.toString() << ")";

    if (state_ != State::STOPPING)
        state_ = State::DETACHED;

    session_attach_timer_.stop();
    session_proxy_->dettach();
    task_runner_->deleteSoon(std::move(session_));

    LOG(LS_INFO) << "Session process is detached";

    if (state_ == State::STOPPING)
        return;

    session_attach_timer_.start(std::chrono::minutes(1), [this]()
    {
        LOG(LS_ERROR) << "Timeout while waiting for session";
        onErrorOccurred();
    });

    // The real session process has ended. We create a temporary fake session.
    session_ = std::make_unique<DesktopSessionFake>(this);
    session_proxy_->attach(session_.get());
    session_->start();
}

std::shared_ptr<DesktopSessionProxy> DesktopSessionManager::sessionProxy() const
{
    return session_proxy_;
}

void DesktopSessionManager::onNewConnection(std::unique_ptr<ipc::Channel> channel)
{
    if (DesktopSessionProcess::filePath() != channel->peerFilePath())
    {
        LOG(LS_ERROR) << "An attempt was made to connect from an unknown application";
        return;
    }

    LOG(LS_INFO) << "Session process successfully connected";

    session_attach_timer_.stop();
    task_runner_->deleteSoon(std::move(server_));

    session_ = std::make_unique<DesktopSessionIpc>(std::move(channel), this);
    session_proxy_->attach(session_.get());

    state_ = State::ATTACHED;
    session_->start();
}

void DesktopSessionManager::onErrorOccurred()
{
    if (state_ == State::STOPPED || state_ == State::STOPPING)
        return;

    state_ = State::STOPPING;
    dettachSession(FROM_HERE);
    state_ = State::STOPPED;
}

void DesktopSessionManager::onDesktopSessionStarted()
{
    delegate_->onDesktopSessionStarted();
}

void DesktopSessionManager::onDesktopSessionStopped()
{
    dettachSession(FROM_HERE);
}

void DesktopSessionManager::onScreenCaptured(const desktop::Frame& frame)
{
    delegate_->onScreenCaptured(frame);
}

void DesktopSessionManager::onCursorCaptured(std::shared_ptr<desktop::MouseCursor> mouse_cursor)
{
    delegate_->onCursorCaptured(std::move(mouse_cursor));
}

void DesktopSessionManager::onScreenListChanged(const proto::ScreenList& list)
{
    delegate_->onScreenListChanged(list);
}

void DesktopSessionManager::onClipboardEvent(const proto::ClipboardEvent& event)
{
    delegate_->onClipboardEvent(event);
}

} // namespace host
<|endoftext|>


<commit_before>//
// PROJECT:         Aspia
// FILE:            host/host_notifier_window.cc
// LICENSE:         GNU General Public License 3
// PROGRAMMERS:     Dmitry Chapyshev (dmitry@aspia.ru)
//

#include "host/ui/host_notifier_window.h"

#include <QMenu>
#include <QMouseEvent>
#include <QScreen>
#include <QTranslator>

#include "host/host_settings.h"

namespace aspia {

namespace {

class SessionTreeItem : public QTreeWidgetItem
{
public:
    SessionTreeItem(const proto::notifier::Session& session)
        : session_(session)
    {
        switch (session_.session_type())
        {
            case proto::auth::SESSION_TYPE_DESKTOP_MANAGE:
                setIcon(0, QIcon(QStringLiteral(":/icon/monitor-keyboard.png")));
                break;

            case proto::auth::SESSION_TYPE_DESKTOP_VIEW:
                setIcon(0, QIcon(QStringLiteral(":/icon/monitor.png")));
                break;

            case proto::auth::SESSION_TYPE_FILE_TRANSFER:
                setIcon(0, QIcon(QStringLiteral(":/icon/folder-stand.png")));
                break;

            default:
                qFatal("Unexpected session type: %d", session_.session_type());
                return;
        }

        setText(0, QString("%1 (%2)")
                .arg(QString::fromStdString(session_.username()))
                .arg(QString::fromStdString(session_.remote_address())));
    }

    const proto::notifier::Session& session() const { return session_; }

private:
    proto::notifier::Session session_;
    Q_DISABLE_COPY(SessionTreeItem)
};

} // namespace

HostNotifierWindow::HostNotifierWindow(QWidget* parent)
    : QWidget(parent, Qt::FramelessWindowHint | Qt::Tool | Qt::WindowStaysOnTopHint)
{
    HostSettings settings;

    QString current_locale = settings.locale();

    if (!locale_loader_.contains(current_locale))
    {
        current_locale = HostSettings::defaultLocale();
        settings.setLocale(current_locale);
    }

    locale_loader_.installTranslators(current_locale);
    ui.setupUi(this);

    ui.label_title->installEventFilter(this);
    ui.label_connections->installEventFilter(this);

    connect(ui.button_show_hide, &QPushButton::pressed,
            this, &HostNotifierWindow::onShowHidePressed);

    connect(ui.button_disconnect_all, &QPushButton::pressed,
            this, &HostNotifierWindow::onDisconnectAllPressed);

    connect(ui.tree, &QTreeWidget::customContextMenuRequested,
            this, &HostNotifierWindow::onContextMenu);

    setAttribute(Qt::WA_TranslucentBackground);

    connect(QApplication::primaryScreen(), &QScreen::availableGeometryChanged,
            this, &HostNotifierWindow::updateWindowPosition);
}

void HostNotifierWindow::setChannelId(const QString& channel_id)
{
    channel_id_ = channel_id;
}

void HostNotifierWindow::sessionOpen(const proto::notifier::Session& session)
{
    ui.tree->addTopLevelItem(new SessionTreeItem(session));
    ui.button_disconnect_all->setEnabled(true);
}

void HostNotifierWindow::sessionClose(const proto::notifier::SessionClose& session_close)
{
    for (int i = 0; i < ui.tree->topLevelItemCount(); ++i)
    {
        SessionTreeItem* item = dynamic_cast<SessionTreeItem*>(ui.tree->topLevelItem(i));
        if (item && item->session().uuid() == session_close.uuid())
        {
            delete item;
            break;
        }
    }

    if (!ui.tree->topLevelItemCount())
        quit();
}

bool HostNotifierWindow::eventFilter(QObject* object, QEvent* event)
{
    if (object == ui.label_title || object == ui.label_connections)
    {
        switch (event->type())
        {
            case QEvent::MouseButtonPress:
            {
                QMouseEvent* mouse_event = dynamic_cast<QMouseEvent*>(event);
                if (mouse_event && mouse_event->button() == Qt::LeftButton)
                {
                    start_pos_ = mouse_event->pos();
                    return true;
                }
            }
            break;

            case QEvent::MouseMove:
            {
                QMouseEvent* mouse_event = dynamic_cast<QMouseEvent*>(event);
                if (mouse_event && mouse_event->buttons() & Qt::LeftButton && start_pos_.x() >= 0)
                {
                    QPoint diff = mouse_event->pos() - start_pos_;
                    move(pos() + diff);
                    return true;
                }
            }
            break;

            case QEvent::MouseButtonRelease:
            {
                start_pos_ = QPoint(-1, -1);
            }
            break;

            default:
                break;
        }
    }

    return QWidget::eventFilter(object, event);
}

void HostNotifierWindow::showEvent(QShowEvent* event)
{
    if (notifier_.isNull())
    {
        notifier_ = new HostNotifier(this);

        connect(notifier_, &HostNotifier::finished, this, &HostNotifierWindow::quit);
        connect(notifier_, &HostNotifier::sessionOpen, this, &HostNotifierWindow::sessionOpen);
        connect(notifier_, &HostNotifier::sessionClose, this, &HostNotifierWindow::sessionClose);
        connect(this, &HostNotifierWindow::killSession, notifier_, &HostNotifier::killSession);

        if (!notifier_->start(channel_id_))
        {
            quit();
            return;
        }
    }

    QWidget::showEvent(event);
}

void HostNotifierWindow::quit()
{
    close();
    QApplication::quit();
}

void HostNotifierWindow::onShowHidePressed()
{
    if (ui.content->isVisible())
        hideNotifier();
    else
        showNotifier();
}

void HostNotifierWindow::onDisconnectAllPressed()
{
    for (int i = 0; i < ui.tree->topLevelItemCount(); ++i)
    {
        SessionTreeItem* item = dynamic_cast<SessionTreeItem*>(ui.tree->topLevelItem(i));
        if (item)
            emit killSession(item->session().uuid());
    }
}

void HostNotifierWindow::onContextMenu(const QPoint& point)
{
    SessionTreeItem* item = dynamic_cast<SessionTreeItem*>(ui.tree->itemAt(point));
    if (!item)
        return;

    QAction disconnect_action(tr("Disconnect"));

    QMenu menu;
    menu.addAction(&disconnect_action);

    if (menu.exec(ui.tree->mapToGlobal(point)) == &disconnect_action)
        emit killSession(item->session().uuid());
}

void HostNotifierWindow::updateWindowPosition()
{
    showNotifier();

    QRect screen_rect = QApplication::primaryScreen()->availableGeometry();
    QSize window_size = frameSize();

    move(screen_rect.x() + screen_rect.width() - window_size.width(),
         screen_rect.y() + screen_rect.height() - window_size.height());
}

void HostNotifierWindow::showNotifier()
{
    if (ui.content->isHidden() && ui.title->isHidden())
    {
        ui.content->show();
        ui.title->show();

        move(window_rect_.topLeft());
        setFixedSize(window_rect_.size());

        ui.button_show_hide->setIcon(QIcon(QStringLiteral(":/icon/arrow-left-gray.png")));
    }
}

void HostNotifierWindow::hideNotifier()
{
    QRect screen_rect = QApplication::primaryScreen()->availableGeometry();
    QSize content_size = ui.content->frameSize();
    window_rect_ = frameGeometry();

    ui.content->hide();
    ui.title->hide();

    move(screen_rect.x() + screen_rect.width() - ui.button_show_hide->width(), pos().y());

    setFixedSize(window_rect_.width() - content_size.width(),
                 window_rect_.height());

    ui.button_show_hide->setIcon(QIcon(QStringLiteral(":/icon/arrow-right-gray.png")));
}

} // namespace aspia
<commit_msg>- Added system information session.<commit_after>//
// PROJECT:         Aspia
// FILE:            host/host_notifier_window.cc
// LICENSE:         GNU General Public License 3
// PROGRAMMERS:     Dmitry Chapyshev (dmitry@aspia.ru)
//

#include "host/ui/host_notifier_window.h"

#include <QMenu>
#include <QMouseEvent>
#include <QScreen>
#include <QTranslator>

#include "host/host_settings.h"

namespace aspia {

namespace {

class SessionTreeItem : public QTreeWidgetItem
{
public:
    SessionTreeItem(const proto::notifier::Session& session)
        : session_(session)
    {
        switch (session_.session_type())
        {
            case proto::auth::SESSION_TYPE_DESKTOP_MANAGE:
                setIcon(0, QIcon(QStringLiteral(":/icon/monitor-keyboard.png")));
                break;

            case proto::auth::SESSION_TYPE_DESKTOP_VIEW:
                setIcon(0, QIcon(QStringLiteral(":/icon/monitor.png")));
                break;

            case proto::auth::SESSION_TYPE_FILE_TRANSFER:
                setIcon(0, QIcon(QStringLiteral(":/icon/folder-stand.png")));
                break;

            case proto::auth::SESSION_TYPE_SYSTEM_INFO:
                setIcon(0, QIcon(QStringLiteral(":/icon/system-monitor.png")));
                break;

            default:
                qFatal("Unexpected session type: %d", session_.session_type());
                return;
        }

        setText(0, QString("%1 (%2)")
                .arg(QString::fromStdString(session_.username()))
                .arg(QString::fromStdString(session_.remote_address())));
    }

    const proto::notifier::Session& session() const { return session_; }

private:
    proto::notifier::Session session_;
    Q_DISABLE_COPY(SessionTreeItem)
};

} // namespace

HostNotifierWindow::HostNotifierWindow(QWidget* parent)
    : QWidget(parent, Qt::FramelessWindowHint | Qt::Tool | Qt::WindowStaysOnTopHint)
{
    HostSettings settings;

    QString current_locale = settings.locale();

    if (!locale_loader_.contains(current_locale))
    {
        current_locale = HostSettings::defaultLocale();
        settings.setLocale(current_locale);
    }

    locale_loader_.installTranslators(current_locale);
    ui.setupUi(this);

    ui.label_title->installEventFilter(this);
    ui.label_connections->installEventFilter(this);

    connect(ui.button_show_hide, &QPushButton::pressed,
            this, &HostNotifierWindow::onShowHidePressed);

    connect(ui.button_disconnect_all, &QPushButton::pressed,
            this, &HostNotifierWindow::onDisconnectAllPressed);

    connect(ui.tree, &QTreeWidget::customContextMenuRequested,
            this, &HostNotifierWindow::onContextMenu);

    setAttribute(Qt::WA_TranslucentBackground);

    connect(QApplication::primaryScreen(), &QScreen::availableGeometryChanged,
            this, &HostNotifierWindow::updateWindowPosition);
}

void HostNotifierWindow::setChannelId(const QString& channel_id)
{
    channel_id_ = channel_id;
}

void HostNotifierWindow::sessionOpen(const proto::notifier::Session& session)
{
    ui.tree->addTopLevelItem(new SessionTreeItem(session));
    ui.button_disconnect_all->setEnabled(true);
}

void HostNotifierWindow::sessionClose(const proto::notifier::SessionClose& session_close)
{
    for (int i = 0; i < ui.tree->topLevelItemCount(); ++i)
    {
        SessionTreeItem* item = dynamic_cast<SessionTreeItem*>(ui.tree->topLevelItem(i));
        if (item && item->session().uuid() == session_close.uuid())
        {
            delete item;
            break;
        }
    }

    if (!ui.tree->topLevelItemCount())
        quit();
}

bool HostNotifierWindow::eventFilter(QObject* object, QEvent* event)
{
    if (object == ui.label_title || object == ui.label_connections)
    {
        switch (event->type())
        {
            case QEvent::MouseButtonPress:
            {
                QMouseEvent* mouse_event = dynamic_cast<QMouseEvent*>(event);
                if (mouse_event && mouse_event->button() == Qt::LeftButton)
                {
                    start_pos_ = mouse_event->pos();
                    return true;
                }
            }
            break;

            case QEvent::MouseMove:
            {
                QMouseEvent* mouse_event = dynamic_cast<QMouseEvent*>(event);
                if (mouse_event && mouse_event->buttons() & Qt::LeftButton && start_pos_.x() >= 0)
                {
                    QPoint diff = mouse_event->pos() - start_pos_;
                    move(pos() + diff);
                    return true;
                }
            }
            break;

            case QEvent::MouseButtonRelease:
            {
                start_pos_ = QPoint(-1, -1);
            }
            break;

            default:
                break;
        }
    }

    return QWidget::eventFilter(object, event);
}

void HostNotifierWindow::showEvent(QShowEvent* event)
{
    if (notifier_.isNull())
    {
        notifier_ = new HostNotifier(this);

        connect(notifier_, &HostNotifier::finished, this, &HostNotifierWindow::quit);
        connect(notifier_, &HostNotifier::sessionOpen, this, &HostNotifierWindow::sessionOpen);
        connect(notifier_, &HostNotifier::sessionClose, this, &HostNotifierWindow::sessionClose);
        connect(this, &HostNotifierWindow::killSession, notifier_, &HostNotifier::killSession);

        if (!notifier_->start(channel_id_))
        {
            quit();
            return;
        }
    }

    QWidget::showEvent(event);
}

void HostNotifierWindow::quit()
{
    close();
    QApplication::quit();
}

void HostNotifierWindow::onShowHidePressed()
{
    if (ui.content->isVisible())
        hideNotifier();
    else
        showNotifier();
}

void HostNotifierWindow::onDisconnectAllPressed()
{
    for (int i = 0; i < ui.tree->topLevelItemCount(); ++i)
    {
        SessionTreeItem* item = dynamic_cast<SessionTreeItem*>(ui.tree->topLevelItem(i));
        if (item)
            emit killSession(item->session().uuid());
    }
}

void HostNotifierWindow::onContextMenu(const QPoint& point)
{
    SessionTreeItem* item = dynamic_cast<SessionTreeItem*>(ui.tree->itemAt(point));
    if (!item)
        return;

    QAction disconnect_action(tr("Disconnect"));

    QMenu menu;
    menu.addAction(&disconnect_action);

    if (menu.exec(ui.tree->mapToGlobal(point)) == &disconnect_action)
        emit killSession(item->session().uuid());
}

void HostNotifierWindow::updateWindowPosition()
{
    showNotifier();

    QRect screen_rect = QApplication::primaryScreen()->availableGeometry();
    QSize window_size = frameSize();

    move(screen_rect.x() + screen_rect.width() - window_size.width(),
         screen_rect.y() + screen_rect.height() - window_size.height());
}

void HostNotifierWindow::showNotifier()
{
    if (ui.content->isHidden() && ui.title->isHidden())
    {
        ui.content->show();
        ui.title->show();

        move(window_rect_.topLeft());
        setFixedSize(window_rect_.size());

        ui.button_show_hide->setIcon(QIcon(QStringLiteral(":/icon/arrow-left-gray.png")));
    }
}

void HostNotifierWindow::hideNotifier()
{
    QRect screen_rect = QApplication::primaryScreen()->availableGeometry();
    QSize content_size = ui.content->frameSize();
    window_rect_ = frameGeometry();

    ui.content->hide();
    ui.title->hide();

    move(screen_rect.x() + screen_rect.width() - ui.button_show_hide->width(), pos().y());

    setFixedSize(window_rect_.width() - content_size.width(),
                 window_rect_.height());

    ui.button_show_hide->setIcon(QIcon(QStringLiteral(":/icon/arrow-right-gray.png")));
}

} // namespace aspia
<|endoftext|>


<commit_before><commit_msg>disable test of pipelined writes because it works only with emulated mutexes at the moment<commit_after><|endoftext|>


<commit_before>#include <gtest/gtest.h>

#include "yall/logger.hpp"
#include "yall/mocks.hpp"

namespace {

struct YallLoggerShould: public ::testing::Test {
  YallLoggerShould():
    backendMock(std::make_shared<MockLoggerBackend>()),
    uut(backendMock) {
  }
  std::shared_ptr<MockLoggerBackend> backendMock;
  ::yall::Logger uut;
};

TEST_F(YallLoggerShould, ForwardToBackend) {
  ::yall::LoggerMessage msg;

  EXPECT_CALL(*backendMock, take(::testing::_))
    .Times(1).WillOnce(::testing::SaveArg<0>(&msg));

  uut.log("test");

  EXPECT_EQ(1, msg.sequence.size());
  EXPECT_EQ(msg.sequence[0].value, "test");
}

}
<commit_msg>Additional tests for logger<commit_after>#include <gtest/gtest.h>

#include "yall/logger.hpp"
#include "yall/mocks.hpp"

namespace {

struct YallLoggerShould: public ::testing::Test {
  YallLoggerShould():
    backendMock(std::make_shared<MockLoggerBackend>()),
    uut(backendMock) {
  }
  std::shared_ptr<MockLoggerBackend> backendMock;
  ::yall::Logger uut;

  ::yall::LoggerMessage msg;

  void SetUp() {
    EXPECT_CALL(*backendMock, take(::testing::_))
      .Times(1).WillOnce(::testing::SaveArg<0>(&msg));
  }

  void verifyCall2() {
    EXPECT_EQ(2, msg.sequence.size());
    EXPECT_EQ(msg.sequence[0].value, "test");
    EXPECT_EQ(msg.sequence[1].value, "1");
  }
};

TEST_F(YallLoggerShould, ForwardToBackend) {
  uut.log("test");

  EXPECT_EQ(1, msg.sequence.size());
  EXPECT_EQ(msg.sequence[0].value, "test");
}

TEST_F(YallLoggerShould, ForwardToBackend2) {
  uut.log("test", 1);
  verifyCall2();
}

TEST_F(YallLoggerShould, ForwardToBackend2WithOperator) {
  uut("test", 1);
  verifyCall2();
}

TEST_F(YallLoggerShould, ForwardToBackend2WithStream) {
  uut("test", 1);
  verifyCall2();
}

TEST_F(YallLoggerShould, ProvideMetaData) {
  uut("test");
  EXPECT_EQ(1, msg.meta.count("yall::TimeStamp"));
  EXPECT_EQ(1, msg.meta.count("yall::ThreadId"));
}

TEST_F(YallLoggerShould, HandleNumbers) {
  uut(short(1), int(1), long(1), (long long)(1),
    (unsigned short)(1), (unsigned int)1, (unsigned long)1, (unsigned long long)1,
    float(1), double(1));
  EXPECT_EQ(10, msg.sequence.size());
}

}
<|endoftext|>