<commit_before>#include "iotools.h"
#include "zlib.h"
#include <stdio.h>

// Compression functions
int compress_onefeat(char * in, char * comp, int fsize) {
    // Struct init
    z_stream defstream;
    defstream.zalloc = Z_NULL;
    defstream.zfree = Z_NULL;
    defstream.opaque = Z_NULL;
    defstream.avail_in = (uInt)fsize; // size of input
    defstream.next_in = (Bytef *)in; // input char array
    defstream.avail_out = (uInt)fsize; // size of output
    defstream.next_out = (Bytef *)comp; // output char array

    // the actual compression work.
    deflateInit(&defstream, Z_BEST_COMPRESSION);
    deflate(&defstream, Z_FINISH);
    deflateEnd(&defstream);
    return defstream.total_out;
}

int decompress_onefeat(char * in, char * comp, int compsize, int fsize) {
    // Struct init
    z_stream infstream;
    infstream.zalloc = Z_NULL;
    infstream.zfree = Z_NULL;
    infstream.opaque = Z_NULL;
    infstream.avail_in = (uInt)compsize; // size of input
    infstream.next_in = (Bytef *)in; // input char array
    infstream.avail_out = (uInt)fsize; // size of output
    infstream.next_out = (Bytef *)comp; // output char array

    // the actual de-compression work.
    inflateInit(&infstream);
    inflate(&infstream, Z_NO_FLUSH);
    inflateEnd(&infstream);
    return infstream.total_out;
}

// File reading functions
std::ifstream::pos_type filesize(std::string filename)
{
    std::ifstream in(filename, std::ios::ate | std::ios::binary);
    if (!in.fail()) {
        return in.tellg();
    } else {
        return 0;
    }
}

int get_file_pos(int * accum, int nb_files, int query, int & res)
{
    int file_id = 0;
    while (query >= accum[file_id] && file_id<nb_files)
    {
        file_id++;
    }
    if (file_id==nb_files) {
        res=-1;
        std::cout << "Could not find feature "  << query << ", maximum id is " << accum[file_id-1] << std::endl;
        return -1;
    }

    if (!file_id)
        res = query;
    else
        res = query-accum[file_id-1];
    return file_id;
}

int get_onefeatcomp(int query_id, size_t read_size, int* accum, vector<ifstream*>& read_in_compfeatures, vector<ifstream*>& read_in_compidx, char* feature_cp) {
    int new_pos = 0;
    int file_id = 0;
    size_t idx_size = sizeof(unsigned long long int);
    unsigned long long int start_feat,end_feat;
    char* comp_feature = new char[read_size];
    file_id = get_file_pos(accum, (int)read_in_compidx.size(), query_id, new_pos);
    if (file_id==-1)
        return -1;
    std::cout << "Feature found in file "  << file_id << " at pos " << new_pos << std::endl;
    read_in_compidx[file_id]->seekg((unsigned long long int)(new_pos)*idx_size);
    read_in_compidx[file_id]->read((char*)&start_feat, idx_size);
    read_in_compidx[file_id]->read((char*)&end_feat, idx_size);
    read_in_compfeatures[file_id]->seekg(start_feat);
    read_in_compfeatures[file_id]->read(comp_feature, end_feat-start_feat);
    decompress_onefeat(comp_feature, feature_cp, (int)end_feat-start_feat, read_size);
    delete[] comp_feature;
    return 0;
}

int get_onefeat(int query_id, size_t read_size, int* accum, vector<ifstream*>& read_in_features, char* feature_cp) {
    int new_pos = 0;
    int file_id = 0;
    file_id = get_file_pos(accum, (int)read_in_features.size(), query_id, new_pos);
    if (file_id==-1)
        return -1;
    //std::cout << "Feature found in file "  << file_id << " at pos " << new_pos << std::endl;
    read_in_features[file_id]->seekg((unsigned long long int)(new_pos)*read_size);
    read_in_features[file_id]->read(feature_cp, read_size);
    return 0;
}

unsigned long long int fill_data_nums(vector<string>& update_hash_files, vector<unsigned long long int>& data_nums, int bit_num) {
    unsigned long long int data_num = 0;
    for (int i=0;i<update_hash_files.size();i++)
    {
        data_nums.push_back((unsigned long long int)filesize(update_hash_files[i])*8/bit_num);
        data_num += data_nums[i];
        std::cout << "We have a " << data_nums[i] << " features in file " << update_hash_files[i] << std::endl;
    }
    std::cout << "We have a total of " << data_num << " features." << std::endl;
    return data_num;
}

int fill_vector_files(vector<ifstream*>& read_in, vector<string>& update_files){
    for (int i=0;i<update_files.size();i++)
    {
    read_in.push_back(new ifstream(update_files[i],ios::in|ios::binary));
    /*read_in.push_back(new ifstream);
    read_in[i]->open(update_files[i],ios::in|ios::binary);*/
    if (!read_in[i]->is_open())
        {
            std::cout << "Cannot load the file " << update_files[i] << std::endl;
            return -1;
        }
    }
    return 0;
}

void fill_accum(vector<unsigned long long int>& data_nums, int * accum) {
    accum[0]=data_nums[0];
    for (int i=1;i<data_nums.size();i++)
    {
        accum[i]=accum[i-1]+data_nums[i];
    }
}

int get_n_features(string update_fn, int* query_ids, int query_num, int norm, int bit_num, size_t read_size, char* feature_cp) {
    // Some ugly hard coded string initialization...
    string bit_string = to_string((long long)bit_num);
    string str_norm = "";
    if (norm)
        str_norm = "_norm";
    string itq_name = "itq" + str_norm + "_" + bit_string;
    string update_compfeature_suffix = "_comp" + str_norm;
    string update_compidx_suffix = "_compidx" + str_norm;
    string update_hash_suffix = "";
    if (norm)
    {
        update_hash_suffix = "_" + itq_name;
    }

    // Config update
    string line;
    vector<string> update_hash_files;
    vector<string> update_compfeature_files;
    vector<string> update_compidx_files;
    ifstream fu(update_fn,ios::in);
    if (!fu.is_open())
    {
        std::cout << "No update! Was looking for " << update_fn << std::endl;
    return -1;
    }
    else
    {
        while (getline(fu, line)) {
            std::cout << "Loading update: " << line << std::endl;
            update_hash_files.push_back(update_hash_prefix+line+update_hash_suffix);
            update_compfeature_files.push_back(update_compfeature_prefix+line+update_compfeature_suffix);
            update_compidx_files.push_back(update_compidx_prefix+line+update_compidx_suffix);
        }
    }

    // get compressed features files pointers and their corresponding indices
    vector<ifstream*> read_in_compfeatures;
    vector<ifstream*> read_in_compidx;
    int status = 0;
    status = fill_vector_files(read_in_compfeatures,update_compfeature_files);
    if (status==-1) {
        std::cout << "Could not load compressed features properly. Exiting." << std::endl;
        // TODO: We should clean here
        return -1;
    }
    status = fill_vector_files(read_in_compidx,update_compidx_files);
    if (status==-1) {
        std::cout << "Could not load compressed indices properly. Exiting." << std::endl;
        // TODO: We should clean here
        return -1;
    }

    // read in all files size to know where to look for features...
    vector<unsigned long long int> data_nums;
    unsigned long long int data_num = fill_data_nums(update_hash_files,data_nums,bit_num);
    int * accum = new int[data_nums.size()];
    fill_accum(data_nums,accum);

    // Get query feature(s)
    for (int i=0;i<query_num;i++)
    {
        std::cout << "Looking for feature #" << query_ids[i] << std::endl;
        // BEWARE: we consider here ids are python/db, so in C they are ids+1...
        // TODO: maybe define a flag python id or not
        status = get_onefeatcomp(query_ids[i]-1,read_size,accum,read_in_compfeatures,read_in_compidx,feature_cp);
        if (status==-1) {
            std::cout << "Could not load compressed feature " << query_ids[i]-1 << ". Exiting." << std::endl;
            // TODO: We should clean here
            return -1;
        }
        feature_cp +=read_size;
    }

    // clean exit
    delete[] accum;
    for (int i = 1; i<data_nums.size();i++)
    {
        read_in_compfeatures[i]->close();
        read_in_compidx[i]->close();
        delete read_in_compfeatures[i];
        delete read_in_compidx[i];
    }
    return 0;
}
<commit_msg>trying to debug get precomp<commit_after>#include "iotools.h"
#include "zlib.h"
#include <stdio.h>

// Compression functions
int compress_onefeat(char * in, char * comp, int fsize) {
    // Struct init
    z_stream defstream;
    defstream.zalloc = Z_NULL;
    defstream.zfree = Z_NULL;
    defstream.opaque = Z_NULL;
    defstream.avail_in = (uInt)fsize; // size of input
    defstream.next_in = (Bytef *)in; // input char array
    defstream.avail_out = (uInt)fsize; // size of output
    defstream.next_out = (Bytef *)comp; // output char array

    // the actual compression work.
    deflateInit(&defstream, Z_BEST_COMPRESSION);
    deflate(&defstream, Z_FINISH);
    deflateEnd(&defstream);
    return defstream.total_out;
}

int decompress_onefeat(char * in, char * comp, int compsize, int fsize) {
    // Struct init
    z_stream infstream;
    infstream.zalloc = Z_NULL;
    infstream.zfree = Z_NULL;
    infstream.opaque = Z_NULL;
    infstream.avail_in = (uInt)compsize; // size of input
    infstream.next_in = (Bytef *)in; // input char array
    infstream.avail_out = (uInt)fsize; // size of output
    infstream.next_out = (Bytef *)comp; // output char array

    // the actual de-compression work.
    inflateInit(&infstream);
    inflate(&infstream, Z_NO_FLUSH);
    inflateEnd(&infstream);
    return infstream.total_out;
}

// File reading functions
std::ifstream::pos_type filesize(std::string filename)
{
    std::ifstream in(filename, std::ios::ate | std::ios::binary);
    if (!in.fail()) {
        return in.tellg();
    } else {
        return 0;
    }
}

int get_file_pos(int * accum, int nb_files, int query, int & res)
{
    int file_id = 0;
    while (query >= accum[file_id] && file_id<nb_files)
    {
        file_id++;
    }
    if (file_id==nb_files) {
        res=-1;
        std::cout << "Could not find feature "  << query << ", maximum id is " << accum[file_id-1] << std::endl;
        return -1;
    }

    if (!file_id)
        res = query;
    else
        res = query-accum[file_id-1];
    return file_id;
}

int get_onefeatcomp(int query_id, size_t read_size, int* accum, vector<ifstream*>& read_in_compfeatures, vector<ifstream*>& read_in_compidx, char* feature_cp) {
    int new_pos = 0;
    int file_id = 0;
    size_t idx_size = sizeof(unsigned long long int);
    unsigned long long int start_feat,end_feat;
    char* comp_feature = new char[read_size];
    file_id = get_file_pos(accum, (int)read_in_compidx.size(), query_id, new_pos);
    if (file_id==-1)
        return -1;
    std::cout << "Feature found in file "  << file_id << " at pos " << new_pos << std::endl;
    read_in_compidx[file_id]->seekg((unsigned long long int)(new_pos)*idx_size);
    read_in_compidx[file_id]->read((char*)&start_feat, idx_size);
    read_in_compidx[file_id]->read((char*)&end_feat, idx_size);
    read_in_compfeatures[file_id]->seekg(start_feat);
    read_in_compfeatures[file_id]->read(comp_feature, end_feat-start_feat);
    std::cout << "Reading compressed feature from "  << start_feat << " to " << end_feat << std::endl;
    uLong total_out = decompress_onefeat(comp_feature, feature_cp, (int)end_feat-start_feat, read_size);
    std::cout << "Decompressed size is "  << total_out << std::endl;
    delete[] comp_feature;
    return 0;
}

int get_onefeat(int query_id, size_t read_size, int* accum, vector<ifstream*>& read_in_features, char* feature_cp) {
    int new_pos = 0;
    int file_id = 0;
    file_id = get_file_pos(accum, (int)read_in_features.size(), query_id, new_pos);
    if (file_id==-1)
        return -1;
    //std::cout << "Feature found in file "  << file_id << " at pos " << new_pos << std::endl;
    read_in_features[file_id]->seekg((unsigned long long int)(new_pos)*read_size);
    read_in_features[file_id]->read(feature_cp, read_size);
    return 0;
}

unsigned long long int fill_data_nums(vector<string>& update_hash_files, vector<unsigned long long int>& data_nums, int bit_num) {
    unsigned long long int data_num = 0;
    for (int i=0;i<update_hash_files.size();i++)
    {
        data_nums.push_back((unsigned long long int)filesize(update_hash_files[i])*8/bit_num);
        data_num += data_nums[i];
        std::cout << "We have a " << data_nums[i] << " features in file " << update_hash_files[i] << std::endl;
    }
    std::cout << "We have a total of " << data_num << " features." << std::endl;
    return data_num;
}

int fill_vector_files(vector<ifstream*>& read_in, vector<string>& update_files){
    for (int i=0;i<update_files.size();i++)
    {
    read_in.push_back(new ifstream(update_files[i],ios::in|ios::binary));
    /*read_in.push_back(new ifstream);
    read_in[i]->open(update_files[i],ios::in|ios::binary);*/
    if (!read_in[i]->is_open())
        {
            std::cout << "Cannot load the file " << update_files[i] << std::endl;
            return -1;
        }
    }
    return 0;
}

void fill_accum(vector<unsigned long long int>& data_nums, int * accum) {
    accum[0]=data_nums[0];
    for (int i=1;i<data_nums.size();i++)
    {
        accum[i]=accum[i-1]+data_nums[i];
    }
}

int get_n_features(string update_fn, int* query_ids, int query_num, int norm, int bit_num, size_t read_size, char* feature_cp) {
    // Some ugly hard coded string initialization...
    string bit_string = to_string((long long)bit_num);
    string str_norm = "";
    if (norm)
        str_norm = "_norm";
    string itq_name = "itq" + str_norm + "_" + bit_string;
    string update_compfeature_suffix = "_comp" + str_norm;
    string update_compidx_suffix = "_compidx" + str_norm;
    string update_hash_suffix = "";
    if (norm)
    {
        update_hash_suffix = "_" + itq_name;
    }

    // Config update
    string line;
    vector<string> update_hash_files;
    vector<string> update_compfeature_files;
    vector<string> update_compidx_files;
    ifstream fu(update_fn,ios::in);
    if (!fu.is_open())
    {
        std::cout << "No update! Was looking for " << update_fn << std::endl;
    return -1;
    }
    else
    {
        while (getline(fu, line)) {
            std::cout << "Loading update: " << line << std::endl;
            update_hash_files.push_back(update_hash_prefix+line+update_hash_suffix);
            update_compfeature_files.push_back(update_compfeature_prefix+line+update_compfeature_suffix);
            update_compidx_files.push_back(update_compidx_prefix+line+update_compidx_suffix);
        }
    }

    // get compressed features files pointers and their corresponding indices
    vector<ifstream*> read_in_compfeatures;
    vector<ifstream*> read_in_compidx;
    int status = 0;
    status = fill_vector_files(read_in_compfeatures,update_compfeature_files);
    if (status==-1) {
        std::cout << "Could not load compressed features properly. Exiting." << std::endl;
        // TODO: We should clean here
        return -1;
    }
    status = fill_vector_files(read_in_compidx,update_compidx_files);
    if (status==-1) {
        std::cout << "Could not load compressed indices properly. Exiting." << std::endl;
        // TODO: We should clean here
        return -1;
    }

    // read in all files size to know where to look for features...
    vector<unsigned long long int> data_nums;
    unsigned long long int data_num = fill_data_nums(update_hash_files,data_nums,bit_num);
    int * accum = new int[data_nums.size()];
    fill_accum(data_nums,accum);

    // Get query feature(s)
    for (int i=0;i<query_num;i++)
    {
        std::cout << "Looking for feature #" << query_ids[i] << std::endl;
        // BEWARE: we consider here ids are python/db, so in C they are ids+1...
        // TODO: maybe define a flag python id or not
        status = get_onefeatcomp(query_ids[i]-1,read_size,accum,read_in_compfeatures,read_in_compidx,feature_cp);
        if (status==-1) {
            std::cout << "Could not load compressed feature " << query_ids[i]-1 << ". Exiting." << std::endl;
            // TODO: We should clean here
            return -1;
        }
        feature_cp +=read_size;
    }

    // clean exit
    delete[] accum;
    for (int i = 1; i<data_nums.size();i++)
    {
        read_in_compfeatures[i]->close();
        read_in_compidx[i]->close();
        delete read_in_compfeatures[i];
        delete read_in_compidx[i];
    }
    return 0;
}
<|endoftext|>


<commit_before>/*=========================================================================

Program:   Medical Imaging & Interaction Toolkit
Module:    $RCSfile$
Language:  C++
Date:      $Date$
Version:   $Revision$

Copyright (c) German Cancer Research Center, Division of Medical and
Biological Informatics. All rights reserved.
See MITKCopyright.txt or http://www.mitk.org/copyright.html for details.

This software is distributed WITHOUT ANY WARRANTY; without even
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

#if(_MSC_VER==1200)
#include <itkFixedCenterOfRotationAffineTransform.h>
#endif

#include "mitkHistogramGenerator.h"

#include "mitkImageAccessByItk.h"

#include "itkScalarImageToHistogramGenerator.h"

mitk::HistogramGenerator::HistogramGenerator() : m_Image(NULL), m_Size(256), m_Histogram(NULL)
{

}

mitk::HistogramGenerator::~HistogramGenerator()
{
}

template < typename TPixel, unsigned int VImageDimension >
void InternalCompute(itk::Image< TPixel, VImageDimension >* itkImage, const mitk::HistogramGenerator* mitkHistoGenerator, mitk::HistogramGenerator::HistogramType& histogram)
{

  typedef itk::Statistics::ScalarImageToHistogramGenerator<
                                              itk::Image< TPixel, VImageDimension >
                                                          >   HistogramGeneratorType;

  typename HistogramGeneratorType::Pointer histogramGenerator = HistogramGeneratorType::New();

  histogramGenerator->SetInput( itkImage );

  histogramGenerator->SetNumberOfBins( mitkHistoGenerator->GetSize() );
  histogramGenerator->SetMarginalScale( 10.0 );
  histogramGenerator->Compute();

  const HistogramGeneratorType::HistogramType* generatedHistogram = histogramGenerator->GetOutput();
  HistogramGeneratorType::HistogramType::SizeType size = generatedHistogram->GetSize();
  histogram.Initialize(size);
  unsigned int i, j;
  for ( i = 0 ; i < HistogramGeneratorType::HistogramType::MeasurementVectorSize ; i++)
  {
    for (j = 0; j < (size[i] - 1) ; ++j)
    {
      histogram.SetBinMin(0, j, generatedHistogram->GetBinMin(0, j));
      histogram.SetBinMax(0, j, generatedHistogram->GetBinMax(0, j));
    }
    for (j = 0; j < (size[i] - 1) ; ++j)
    {
      histogram.SetFrequency(j, generatedHistogram->GetFrequency(j));
    }
  }
}

void mitk::HistogramGenerator::ComputeHistogram()
{
  if((m_Histogram.IsNull()) || (m_Histogram->GetMTime() < m_Image->GetMTime()))
  {
    if(m_Histogram.IsNull())
      m_Histogram = HistogramType::New();
    const_cast<mitk::Image*>(m_Image.GetPointer())->SetRequestedRegionToLargestPossibleRegion(); //@todo without this, Image::GetScalarMin does not work for dim==3 (including sliceselector!)
    const_cast<mitk::Image*>(m_Image.GetPointer())->Update();
    AccessByItk_2(m_Image, InternalCompute, this, *m_Histogram);
    m_Histogram->Modified();
  }

// debug code
  /*
    std::cout << "Histogram modfied 1" << m_Histogram->GetMTime() << std::endl;
    m_Histogram->Modified();
    std::cout << "Histogram modfied 2" << m_Histogram->GetMTime() << std::endl;
    std::cout << "Image modfied" << m_Image->GetMTime() << std::endl;
  const unsigned int histogramSize = m_Histogram->Size();

  std::cout << "Histogram size " << histogramSize << std::endl;

  HistogramType::ConstIterator itr = GetHistogram()->Begin();
  HistogramType::ConstIterator end = GetHistogram()->End();

  int bin = 0;
  while( itr != end )
    {
    std::cout << "bin = " << GetHistogram()->GetBinMin(0,bin) << "--" << GetHistogram()->GetBinMax(0,bin) << " frequency = ";
    std::cout << itr.GetFrequency() << std::endl;
    ++itr;
    ++bin;
    }
    */
}

float mitk::HistogramGenerator::GetMaximumFrequency() const {
  return CalculateMaximumFrequency(this->m_Histogram);
};

float mitk::HistogramGenerator::CalculateMaximumFrequency(const HistogramType* histogram)
{
  HistogramType::ConstIterator itr = histogram->Begin();
  HistogramType::ConstIterator end = histogram->End();

  float maxFreq = 0;
  while( itr != end )
    {
    maxFreq = vnl_math_max(maxFreq,itr.GetFrequency());
    ++itr;
    }
  return maxFreq;
};
<commit_msg>FIX: gcc compile<commit_after>/*=========================================================================

Program:   Medical Imaging & Interaction Toolkit
Module:    $RCSfile$
Language:  C++
Date:      $Date$
Version:   $Revision$

Copyright (c) German Cancer Research Center, Division of Medical and
Biological Informatics. All rights reserved.
See MITKCopyright.txt or http://www.mitk.org/copyright.html for details.

This software is distributed WITHOUT ANY WARRANTY; without even
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

#if(_MSC_VER==1200)
#include <itkFixedCenterOfRotationAffineTransform.h>
#endif

#include "mitkHistogramGenerator.h"

#include "mitkImageAccessByItk.h"

#include "itkScalarImageToHistogramGenerator.h"

mitk::HistogramGenerator::HistogramGenerator() : m_Image(NULL), m_Size(256), m_Histogram(NULL)
{

}

mitk::HistogramGenerator::~HistogramGenerator()
{
}

template < typename TPixel, unsigned int VImageDimension >
void InternalCompute(itk::Image< TPixel, VImageDimension >* itkImage, const mitk::HistogramGenerator* mitkHistoGenerator, mitk::HistogramGenerator::HistogramType& histogram)
{

  typedef itk::Statistics::ScalarImageToHistogramGenerator<
                                              itk::Image< TPixel, VImageDimension >
                                                          >   HistogramGeneratorType;

  typename HistogramGeneratorType::Pointer histogramGenerator = HistogramGeneratorType::New();

  histogramGenerator->SetInput( itkImage );

  histogramGenerator->SetNumberOfBins( mitkHistoGenerator->GetSize() );
  histogramGenerator->SetMarginalScale( 10.0 );
  histogramGenerator->Compute();

  const typename HistogramGeneratorType::HistogramType* generatedHistogram = histogramGenerator->GetOutput();
  typename HistogramGeneratorType::HistogramType::SizeType size = generatedHistogram->GetSize();
  histogram.Initialize(size);
  unsigned int i, j;
  for ( i = 0 ; i < HistogramGeneratorType::HistogramType::MeasurementVectorSize ; i++)
  {
    for (j = 0; j < (size[i] - 1) ; ++j)
    {
      histogram.SetBinMin(0, j, generatedHistogram->GetBinMin(0, j));
      histogram.SetBinMax(0, j, generatedHistogram->GetBinMax(0, j));
    }
    for (j = 0; j < (size[i] - 1) ; ++j)
    {
      histogram.SetFrequency(j, generatedHistogram->GetFrequency(j));
    }
  }
}

void mitk::HistogramGenerator::ComputeHistogram()
{
  if((m_Histogram.IsNull()) || (m_Histogram->GetMTime() < m_Image->GetMTime()))
  {
    if(m_Histogram.IsNull())
      m_Histogram = HistogramType::New();
    const_cast<mitk::Image*>(m_Image.GetPointer())->SetRequestedRegionToLargestPossibleRegion(); //@todo without this, Image::GetScalarMin does not work for dim==3 (including sliceselector!)
    const_cast<mitk::Image*>(m_Image.GetPointer())->Update();
    AccessByItk_2(m_Image, InternalCompute, this, *m_Histogram);
    m_Histogram->Modified();
  }

// debug code
  /*
    std::cout << "Histogram modfied 1" << m_Histogram->GetMTime() << std::endl;
    m_Histogram->Modified();
    std::cout << "Histogram modfied 2" << m_Histogram->GetMTime() << std::endl;
    std::cout << "Image modfied" << m_Image->GetMTime() << std::endl;
  const unsigned int histogramSize = m_Histogram->Size();

  std::cout << "Histogram size " << histogramSize << std::endl;

  HistogramType::ConstIterator itr = GetHistogram()->Begin();
  HistogramType::ConstIterator end = GetHistogram()->End();

  int bin = 0;
  while( itr != end )
    {
    std::cout << "bin = " << GetHistogram()->GetBinMin(0,bin) << "--" << GetHistogram()->GetBinMax(0,bin) << " frequency = ";
    std::cout << itr.GetFrequency() << std::endl;
    ++itr;
    ++bin;
    }
    */
}

float mitk::HistogramGenerator::GetMaximumFrequency() const {
  return CalculateMaximumFrequency(this->m_Histogram);
};

float mitk::HistogramGenerator::CalculateMaximumFrequency(const HistogramType* histogram)
{
  HistogramType::ConstIterator itr = histogram->Begin();
  HistogramType::ConstIterator end = histogram->End();

  float maxFreq = 0;
  while( itr != end )
    {
    maxFreq = vnl_math_max(maxFreq,itr.GetFrequency());
    ++itr;
    }
  return maxFreq;
};
<|endoftext|>


<commit_before>//===- llvm/unittest/DebugInfo/CodeView/RandomAccessVisitorTest.cpp -------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "llvm/ADT/SmallBitVector.h"
#include "llvm/DebugInfo/CodeView/CVTypeVisitor.h"
#include "llvm/DebugInfo/CodeView/LazyRandomTypeCollection.h"
#include "llvm/DebugInfo/CodeView/TypeRecord.h"
#include "llvm/DebugInfo/CodeView/TypeRecordMapping.h"
#include "llvm/DebugInfo/CodeView/TypeSerializer.h"
#include "llvm/DebugInfo/CodeView/TypeServerHandler.h"
#include "llvm/DebugInfo/CodeView/TypeTableBuilder.h"
#include "llvm/DebugInfo/CodeView/TypeVisitorCallbacks.h"
#include "llvm/DebugInfo/PDB/Native/RawTypes.h"
#include "llvm/Support/Allocator.h"
#include "llvm/Support/BinaryItemStream.h"
#include "llvm/Support/Error.h"
#include "llvm/Testing/Support/Error.h"

#include "gtest/gtest.h"

using namespace llvm;
using namespace llvm::codeview;
using namespace llvm::pdb;

namespace llvm {
namespace codeview {
inline bool operator==(const ArrayRecord &R1, const ArrayRecord &R2) {
  if (R1.ElementType != R2.ElementType)
    return false;
  if (R1.IndexType != R2.IndexType)
    return false;
  if (R1.Name != R2.Name)
    return false;
  if (R1.Size != R2.Size)
    return false;
  return true;
}
inline bool operator!=(const ArrayRecord &R1, const ArrayRecord &R2) {
  return !(R1 == R2);
}

inline bool operator==(const CVType &R1, const CVType &R2) {
  if (R1.Type != R2.Type)
    return false;
  if (R1.RecordData != R2.RecordData)
    return false;
  return true;
}
inline bool operator!=(const CVType &R1, const CVType &R2) {
  return !(R1 == R2);
}
}
}

namespace llvm {
template <> struct BinaryItemTraits<CVType> {
  static size_t length(const CVType &Item) { return Item.length(); }
  static ArrayRef<uint8_t> bytes(const CVType &Item) { return Item.data(); }
};
}

namespace {

class MockCallbacks : public TypeVisitorCallbacks {
public:
  virtual Error visitTypeBegin(CVType &CVR, TypeIndex Index) {
    Indices.push_back(Index);
    return Error::success();
  }
  virtual Error visitKnownRecord(CVType &CVR, ArrayRecord &AR) {
    VisitedRecords.push_back(AR);
    RawRecords.push_back(CVR);
    return Error::success();
  }

  uint32_t count() const {
    assert(Indices.size() == RawRecords.size());
    assert(Indices.size() == VisitedRecords.size());
    return Indices.size();
  }
  std::vector<TypeIndex> Indices;
  std::vector<CVType> RawRecords;
  std::vector<ArrayRecord> VisitedRecords;
};

class RandomAccessVisitorTest : public testing::Test {
public:
  RandomAccessVisitorTest() {}

  static void SetUpTestCase() {
    GlobalState = llvm::make_unique<GlobalTestState>();

    TypeTableBuilder Builder(GlobalState->Allocator);

    uint32_t Offset = 0;
    for (int I = 0; I < 11; ++I) {
      ArrayRecord AR(TypeRecordKind::Array);
      AR.ElementType = TypeIndex::Int32();
      AR.IndexType = TypeIndex::UInt32();
      AR.Size = I;
      std::string Name;
      raw_string_ostream Stream(Name);
      Stream << "Array [" << I << "]";
      AR.Name = GlobalState->Strings.save(Stream.str());
      GlobalState->Records.push_back(AR);
      GlobalState->Indices.push_back(Builder.writeKnownType(AR));

      CVType Type(TypeLeafKind::LF_ARRAY, Builder.records().back());
      GlobalState->TypeVector.push_back(Type);

      GlobalState->AllOffsets.push_back(
          {GlobalState->Indices.back(), ulittle32_t(Offset)});
      Offset += Type.length();
    }

    GlobalState->ItemStream.setItems(GlobalState->TypeVector);
    GlobalState->TypeArray = VarStreamArray<CVType>(GlobalState->ItemStream);
  }

  static void TearDownTestCase() { GlobalState.reset(); }

  void SetUp() override {
    TestState = llvm::make_unique<PerTestState>();
  }

  void TearDown() override { TestState.reset(); }

protected:
  bool ValidateDatabaseRecord(LazyRandomTypeCollection &Types, uint32_t Index) {
    TypeIndex TI = TypeIndex::fromArrayIndex(Index);
    if (!Types.contains(TI))
      return false;
    if (GlobalState->TypeVector[Index] != Types.getType(TI))
      return false;
    return true;
  }

  bool ValidateVisitedRecord(uint32_t VisitationOrder,
                             uint32_t GlobalArrayIndex) {
    TypeIndex TI = TypeIndex::fromArrayIndex(GlobalArrayIndex);
    if (TI != TestState->Callbacks.Indices[VisitationOrder])
      return false;

    if (GlobalState->TypeVector[TI.toArrayIndex()] !=
        TestState->Callbacks.RawRecords[VisitationOrder])
      return false;

    if (GlobalState->Records[TI.toArrayIndex()] !=
        TestState->Callbacks.VisitedRecords[VisitationOrder])
      return false;

    return true;
  }

  struct GlobalTestState {
    GlobalTestState() : Strings(Allocator), ItemStream(llvm::support::little) {}

    BumpPtrAllocator Allocator;
    StringSaver Strings;

    std::vector<ArrayRecord> Records;
    std::vector<TypeIndex> Indices;
    std::vector<TypeIndexOffset> AllOffsets;
    std::vector<CVType> TypeVector;
    BinaryItemStream<CVType> ItemStream;
    VarStreamArray<CVType> TypeArray;

    MutableBinaryByteStream Stream;
  };

  struct PerTestState {
    FixedStreamArray<TypeIndexOffset> Offsets;

    MockCallbacks Callbacks;
  };

  FixedStreamArray<TypeIndexOffset>
  createPartialOffsets(MutableBinaryByteStream &Storage,
                       std::initializer_list<uint32_t> Indices) {

    uint32_t Count = Indices.size();
    uint32_t Size = Count * sizeof(TypeIndexOffset);
    uint8_t *Buffer = GlobalState->Allocator.Allocate<uint8_t>(Size);
    MutableArrayRef<uint8_t> Bytes(Buffer, Size);
    Storage = MutableBinaryByteStream(Bytes, support::little);
    BinaryStreamWriter Writer(Storage);
    for (const auto I : Indices)
      consumeError(Writer.writeObject(GlobalState->AllOffsets[I]));

    BinaryStreamReader Reader(Storage);
    FixedStreamArray<TypeIndexOffset> Result;
    consumeError(Reader.readArray(Result, Count));
    return Result;
  }

  static std::unique_ptr<GlobalTestState> GlobalState;
  std::unique_ptr<PerTestState> TestState;
};

std::unique_ptr<RandomAccessVisitorTest::GlobalTestState>
    RandomAccessVisitorTest::GlobalState;
}

TEST_F(RandomAccessVisitorTest, MultipleVisits) {
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 8});
  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);

  std::vector<uint32_t> IndicesToVisit = {5, 5, 5};

  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [0,8) should be present
  EXPECT_EQ(8u, Types.size());
  for (uint32_t I = 0; I < 8; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // 5, 5, 5
  EXPECT_EQ(3u, TestState->Callbacks.count());
  for (auto I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, DescendingWithinChunk) {
  // Visit multiple items from the same "chunk" in reverse order.  In this
  // example, it's 7 then 4 then 2.  At the end, all records from 0 to 7 should
  // be known by the database, but only 2, 4, and 7 should have been visited.
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 8});

  std::vector<uint32_t> IndicesToVisit = {7, 4, 2};

  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);
  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [0, 7]
  EXPECT_EQ(8u, Types.size());
  for (uint32_t I = 0; I < 8; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // 2, 4, 7
  EXPECT_EQ(3u, TestState->Callbacks.count());
  for (auto I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, AscendingWithinChunk) {
  // * Visit multiple items from the same chunk in ascending order, ensuring
  //   that intermediate items are not visited.  In the below example, it's
  //   5 -> 6 -> 7 which come from the [4,8) chunk.
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 8});

  std::vector<uint32_t> IndicesToVisit = {2, 4, 7};

  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);
  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [0, 7]
  EXPECT_EQ(8u, Types.size());
  for (uint32_t I = 0; I < 8; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // 2, 4, 7
  EXPECT_EQ(3u, TestState->Callbacks.count());
  for (auto &I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, StopPrematurelyInChunk) {
  // * Don't visit the last item in one chunk, ensuring that visitation stops
  //   at the record you specify, and the chunk is only partially visited.
  //   In the below example, this is tested by visiting 0 and 1 but not 2,
  //   all from the [0,3) chunk.
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 8});

  std::vector<uint32_t> IndicesToVisit = {0, 1, 2};

  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);

  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [0, 8) should be visited.
  EXPECT_EQ(8u, Types.size());
  for (uint32_t I = 0; I < 8; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // [0, 2]
  EXPECT_EQ(3u, TestState->Callbacks.count());
  for (auto I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, InnerChunk) {
  // Test that when a request comes from a chunk in the middle of the partial
  // offsets array, that items from surrounding chunks are not visited or
  // added to the database.
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 4, 9});

  std::vector<uint32_t> IndicesToVisit = {5, 7};

  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);

  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [4, 9)
  EXPECT_EQ(5u, Types.size());
  for (uint32_t I = 4; I < 9; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // 5, 7
  EXPECT_EQ(2u, TestState->Callbacks.count());
  for (auto &I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, CrossChunkName) {
  TypeTableBuilder Builder(GlobalState->Allocator);

  // TypeIndex 0
  ClassRecord Class(TypeRecordKind::Class);
  Class.Name = "FooClass";
  Class.Options = ClassOptions::None;
  Class.MemberCount = 0;
  Class.DerivationList = TypeIndex::fromArrayIndex(0);
  Class.FieldList = TypeIndex::fromArrayIndex(0);
  Class.VTableShape = TypeIndex::fromArrayIndex(0);
  TypeIndex IndexZero = Builder.writeKnownType(Class);

  // TypeIndex 1 refers to type index 0.
  ModifierRecord Modifier(TypeRecordKind::Modifier);
  Modifier.ModifiedType = TypeIndex::fromArrayIndex(0);
  Modifier.Modifiers = ModifierOptions::Const;
  TypeIndex IndexOne = Builder.writeKnownType(Modifier);

  // set up a type stream that refers to the above two serialized records.
  std::vector<CVType> TypeArray;
  TypeArray.push_back(
      CVType(static_cast<TypeLeafKind>(Class.Kind), Builder.records()[0]));
  TypeArray.push_back(
      CVType(static_cast<TypeLeafKind>(Modifier.Kind), Builder.records()[1]));
  BinaryItemStream<CVType> ItemStream(llvm::support::little);
  ItemStream.setItems(TypeArray);
  VarStreamArray<CVType> TypeStream(ItemStream);

  // Figure out the byte offset of the second item.
  auto ItemOneIter = TypeStream.begin();
  ++ItemOneIter;

  // Set up a partial offsets buffer that contains the first and second items
  // in separate chunks.
  std::vector<TypeIndexOffset> TIO;
  TIO.push_back({IndexZero, ulittle32_t(0u)});
  TIO.push_back({IndexOne, ulittle32_t(ItemOneIter.offset())});
  ArrayRef<uint8_t> Buffer(reinterpret_cast<const uint8_t *>(TIO.data()),
                           TIO.size() * sizeof(TypeIndexOffset));

  BinaryStreamReader Reader(Buffer, llvm::support::little);
  FixedStreamArray<TypeIndexOffset> PartialOffsets;
  ASSERT_THAT_ERROR(Reader.readArray(PartialOffsets, 2), Succeeded());

  LazyRandomTypeCollection Types(TypeStream, 2, PartialOffsets);

  StringRef Name = Types.getTypeName(IndexOne);
  EXPECT_EQ("const FooClass", Name);
}<commit_msg>Try to fix uninitialized read in unit test.<commit_after>//===- llvm/unittest/DebugInfo/CodeView/RandomAccessVisitorTest.cpp -------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "llvm/ADT/SmallBitVector.h"
#include "llvm/DebugInfo/CodeView/CVTypeVisitor.h"
#include "llvm/DebugInfo/CodeView/LazyRandomTypeCollection.h"
#include "llvm/DebugInfo/CodeView/TypeRecord.h"
#include "llvm/DebugInfo/CodeView/TypeRecordMapping.h"
#include "llvm/DebugInfo/CodeView/TypeSerializer.h"
#include "llvm/DebugInfo/CodeView/TypeServerHandler.h"
#include "llvm/DebugInfo/CodeView/TypeTableBuilder.h"
#include "llvm/DebugInfo/CodeView/TypeVisitorCallbacks.h"
#include "llvm/DebugInfo/PDB/Native/RawTypes.h"
#include "llvm/Support/Allocator.h"
#include "llvm/Support/BinaryItemStream.h"
#include "llvm/Support/Error.h"
#include "llvm/Testing/Support/Error.h"

#include "gtest/gtest.h"

using namespace llvm;
using namespace llvm::codeview;
using namespace llvm::pdb;

namespace llvm {
namespace codeview {
inline bool operator==(const ArrayRecord &R1, const ArrayRecord &R2) {
  if (R1.ElementType != R2.ElementType)
    return false;
  if (R1.IndexType != R2.IndexType)
    return false;
  if (R1.Name != R2.Name)
    return false;
  if (R1.Size != R2.Size)
    return false;
  return true;
}
inline bool operator!=(const ArrayRecord &R1, const ArrayRecord &R2) {
  return !(R1 == R2);
}

inline bool operator==(const CVType &R1, const CVType &R2) {
  if (R1.Type != R2.Type)
    return false;
  if (R1.RecordData != R2.RecordData)
    return false;
  return true;
}
inline bool operator!=(const CVType &R1, const CVType &R2) {
  return !(R1 == R2);
}
}
}

namespace llvm {
template <> struct BinaryItemTraits<CVType> {
  static size_t length(const CVType &Item) { return Item.length(); }
  static ArrayRef<uint8_t> bytes(const CVType &Item) { return Item.data(); }
};
}

namespace {

class MockCallbacks : public TypeVisitorCallbacks {
public:
  virtual Error visitTypeBegin(CVType &CVR, TypeIndex Index) {
    Indices.push_back(Index);
    return Error::success();
  }
  virtual Error visitKnownRecord(CVType &CVR, ArrayRecord &AR) {
    VisitedRecords.push_back(AR);
    RawRecords.push_back(CVR);
    return Error::success();
  }

  uint32_t count() const {
    assert(Indices.size() == RawRecords.size());
    assert(Indices.size() == VisitedRecords.size());
    return Indices.size();
  }
  std::vector<TypeIndex> Indices;
  std::vector<CVType> RawRecords;
  std::vector<ArrayRecord> VisitedRecords;
};

class RandomAccessVisitorTest : public testing::Test {
public:
  RandomAccessVisitorTest() {}

  static void SetUpTestCase() {
    GlobalState = llvm::make_unique<GlobalTestState>();

    TypeTableBuilder Builder(GlobalState->Allocator);

    uint32_t Offset = 0;
    for (int I = 0; I < 11; ++I) {
      ArrayRecord AR(TypeRecordKind::Array);
      AR.ElementType = TypeIndex::Int32();
      AR.IndexType = TypeIndex::UInt32();
      AR.Size = I;
      std::string Name;
      raw_string_ostream Stream(Name);
      Stream << "Array [" << I << "]";
      AR.Name = GlobalState->Strings.save(Stream.str());
      GlobalState->Records.push_back(AR);
      GlobalState->Indices.push_back(Builder.writeKnownType(AR));

      CVType Type(TypeLeafKind::LF_ARRAY, Builder.records().back());
      GlobalState->TypeVector.push_back(Type);

      GlobalState->AllOffsets.push_back(
          {GlobalState->Indices.back(), ulittle32_t(Offset)});
      Offset += Type.length();
    }

    GlobalState->ItemStream.setItems(GlobalState->TypeVector);
    GlobalState->TypeArray = VarStreamArray<CVType>(GlobalState->ItemStream);
  }

  static void TearDownTestCase() { GlobalState.reset(); }

  void SetUp() override {
    TestState = llvm::make_unique<PerTestState>();
  }

  void TearDown() override { TestState.reset(); }

protected:
  bool ValidateDatabaseRecord(LazyRandomTypeCollection &Types, uint32_t Index) {
    TypeIndex TI = TypeIndex::fromArrayIndex(Index);
    if (!Types.contains(TI))
      return false;
    if (GlobalState->TypeVector[Index] != Types.getType(TI))
      return false;
    return true;
  }

  bool ValidateVisitedRecord(uint32_t VisitationOrder,
                             uint32_t GlobalArrayIndex) {
    TypeIndex TI = TypeIndex::fromArrayIndex(GlobalArrayIndex);
    if (TI != TestState->Callbacks.Indices[VisitationOrder])
      return false;

    if (GlobalState->TypeVector[TI.toArrayIndex()] !=
        TestState->Callbacks.RawRecords[VisitationOrder])
      return false;

    if (GlobalState->Records[TI.toArrayIndex()] !=
        TestState->Callbacks.VisitedRecords[VisitationOrder])
      return false;

    return true;
  }

  struct GlobalTestState {
    GlobalTestState() : Strings(Allocator), ItemStream(llvm::support::little) {}

    BumpPtrAllocator Allocator;
    StringSaver Strings;

    std::vector<ArrayRecord> Records;
    std::vector<TypeIndex> Indices;
    std::vector<TypeIndexOffset> AllOffsets;
    std::vector<CVType> TypeVector;
    BinaryItemStream<CVType> ItemStream;
    VarStreamArray<CVType> TypeArray;

    MutableBinaryByteStream Stream;
  };

  struct PerTestState {
    FixedStreamArray<TypeIndexOffset> Offsets;

    MockCallbacks Callbacks;
  };

  FixedStreamArray<TypeIndexOffset>
  createPartialOffsets(MutableBinaryByteStream &Storage,
                       std::initializer_list<uint32_t> Indices) {

    uint32_t Count = Indices.size();
    uint32_t Size = Count * sizeof(TypeIndexOffset);
    uint8_t *Buffer = GlobalState->Allocator.Allocate<uint8_t>(Size);
    MutableArrayRef<uint8_t> Bytes(Buffer, Size);
    Storage = MutableBinaryByteStream(Bytes, support::little);
    BinaryStreamWriter Writer(Storage);
    for (const auto I : Indices)
      consumeError(Writer.writeObject(GlobalState->AllOffsets[I]));

    BinaryStreamReader Reader(Storage);
    FixedStreamArray<TypeIndexOffset> Result;
    consumeError(Reader.readArray(Result, Count));
    return Result;
  }

  static std::unique_ptr<GlobalTestState> GlobalState;
  std::unique_ptr<PerTestState> TestState;
};

std::unique_ptr<RandomAccessVisitorTest::GlobalTestState>
    RandomAccessVisitorTest::GlobalState;
}

TEST_F(RandomAccessVisitorTest, MultipleVisits) {
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 8});
  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);

  std::vector<uint32_t> IndicesToVisit = {5, 5, 5};

  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [0,8) should be present
  EXPECT_EQ(8u, Types.size());
  for (uint32_t I = 0; I < 8; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // 5, 5, 5
  EXPECT_EQ(3u, TestState->Callbacks.count());
  for (auto I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, DescendingWithinChunk) {
  // Visit multiple items from the same "chunk" in reverse order.  In this
  // example, it's 7 then 4 then 2.  At the end, all records from 0 to 7 should
  // be known by the database, but only 2, 4, and 7 should have been visited.
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 8});

  std::vector<uint32_t> IndicesToVisit = {7, 4, 2};

  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);
  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [0, 7]
  EXPECT_EQ(8u, Types.size());
  for (uint32_t I = 0; I < 8; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // 2, 4, 7
  EXPECT_EQ(3u, TestState->Callbacks.count());
  for (auto I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, AscendingWithinChunk) {
  // * Visit multiple items from the same chunk in ascending order, ensuring
  //   that intermediate items are not visited.  In the below example, it's
  //   5 -> 6 -> 7 which come from the [4,8) chunk.
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 8});

  std::vector<uint32_t> IndicesToVisit = {2, 4, 7};

  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);
  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [0, 7]
  EXPECT_EQ(8u, Types.size());
  for (uint32_t I = 0; I < 8; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // 2, 4, 7
  EXPECT_EQ(3u, TestState->Callbacks.count());
  for (auto &I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, StopPrematurelyInChunk) {
  // * Don't visit the last item in one chunk, ensuring that visitation stops
  //   at the record you specify, and the chunk is only partially visited.
  //   In the below example, this is tested by visiting 0 and 1 but not 2,
  //   all from the [0,3) chunk.
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 8});

  std::vector<uint32_t> IndicesToVisit = {0, 1, 2};

  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);

  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [0, 8) should be visited.
  EXPECT_EQ(8u, Types.size());
  for (uint32_t I = 0; I < 8; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // [0, 2]
  EXPECT_EQ(3u, TestState->Callbacks.count());
  for (auto I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, InnerChunk) {
  // Test that when a request comes from a chunk in the middle of the partial
  // offsets array, that items from surrounding chunks are not visited or
  // added to the database.
  TestState->Offsets = createPartialOffsets(GlobalState->Stream, {0, 4, 9});

  std::vector<uint32_t> IndicesToVisit = {5, 7};

  LazyRandomTypeCollection Types(GlobalState->TypeArray,
                                 GlobalState->TypeVector.size(),
                                 TestState->Offsets);

  for (uint32_t I : IndicesToVisit) {
    TypeIndex TI = TypeIndex::fromArrayIndex(I);
    CVType T = Types.getType(TI);
    EXPECT_THAT_ERROR(codeview::visitTypeRecord(T, TI, TestState->Callbacks),
                      Succeeded());
  }

  // [4, 9)
  EXPECT_EQ(5u, Types.size());
  for (uint32_t I = 4; I < 9; ++I)
    EXPECT_TRUE(ValidateDatabaseRecord(Types, I));

  // 5, 7
  EXPECT_EQ(2u, TestState->Callbacks.count());
  for (auto &I : enumerate(IndicesToVisit))
    EXPECT_TRUE(ValidateVisitedRecord(I.index(), I.value()));
}

TEST_F(RandomAccessVisitorTest, CrossChunkName) {
  TypeTableBuilder Builder(GlobalState->Allocator);

  // TypeIndex 0
  ClassRecord Class(TypeRecordKind::Class);
  Class.Name = "FooClass";
  Class.Options = ClassOptions::None;
  Class.MemberCount = 0;
  Class.Size = 4U;
  Class.DerivationList = TypeIndex::fromArrayIndex(0);
  Class.FieldList = TypeIndex::fromArrayIndex(0);
  Class.VTableShape = TypeIndex::fromArrayIndex(0);
  TypeIndex IndexZero = Builder.writeKnownType(Class);

  // TypeIndex 1 refers to type index 0.
  ModifierRecord Modifier(TypeRecordKind::Modifier);
  Modifier.ModifiedType = TypeIndex::fromArrayIndex(0);
  Modifier.Modifiers = ModifierOptions::Const;
  TypeIndex IndexOne = Builder.writeKnownType(Modifier);

  // set up a type stream that refers to the above two serialized records.
  std::vector<CVType> TypeArray;
  TypeArray.push_back(
      CVType(static_cast<TypeLeafKind>(Class.Kind), Builder.records()[0]));
  TypeArray.push_back(
      CVType(static_cast<TypeLeafKind>(Modifier.Kind), Builder.records()[1]));
  BinaryItemStream<CVType> ItemStream(llvm::support::little);
  ItemStream.setItems(TypeArray);
  VarStreamArray<CVType> TypeStream(ItemStream);

  // Figure out the byte offset of the second item.
  auto ItemOneIter = TypeStream.begin();
  ++ItemOneIter;

  // Set up a partial offsets buffer that contains the first and second items
  // in separate chunks.
  std::vector<TypeIndexOffset> TIO;
  TIO.push_back({IndexZero, ulittle32_t(0u)});
  TIO.push_back({IndexOne, ulittle32_t(ItemOneIter.offset())});
  ArrayRef<uint8_t> Buffer(reinterpret_cast<const uint8_t *>(TIO.data()),
                           TIO.size() * sizeof(TypeIndexOffset));

  BinaryStreamReader Reader(Buffer, llvm::support::little);
  FixedStreamArray<TypeIndexOffset> PartialOffsets;
  ASSERT_THAT_ERROR(Reader.readArray(PartialOffsets, 2), Succeeded());

  LazyRandomTypeCollection Types(TypeStream, 2, PartialOffsets);

  StringRef Name = Types.getTypeName(IndexOne);
  EXPECT_EQ("const FooClass", Name);
}<|endoftext|>


<commit_before>#pragma once

#include <type_traits>
#include <tuple>
#include <ostream>

#include "blackhole/detail/traits/attributes/pack/feeder.hpp"

#include "format/message/extraction.hpp"
#include "format/message/insitu.hpp"
#include "logger.hpp"
#include "utils/format.hpp"

namespace blackhole {

namespace aux {

class feeder_t {
protected:
    log::record_t& record;

public:
    feeder_t(log::record_t& record) :
        record(record)
    {}

    // Using initializer lists is problematic, since GCC 4.6 has an extension which can
    // decude { { "value", "42" } } expressions in function call to the initializer list
    // without looking at the overload resolution. This violates C++11 standard, which
    // allows such deducing only for `auto` keyword (including range-based `for` loop).
    // Please, specify `attribute::list` type explicitly when using initializer lists to
    // provide attributes.
    feeder_t operator()(std::initializer_list<std::pair<std::string, log::attribute_value_t>>&& args) {
        for (auto it = args.begin(); it != args.end(); ++it) {
            record.attributes.insert(std::make_pair(it->first, log::attribute_t(it->second)));
        }

        return feeder_t { record };
    }

    template<typename... Args>
    feeder_t operator()(Args&&... args) {
        static_assert((is_keyword_pack<Args...>::value || is_emplace_pack<Args...>::value),
                      "parameter pack must be either attribute pack or emplace pack");
        typedef typename pack_determiner<Args...>::type pack_type;
        pack_feeder<pack_type>::feed(record, std::forward<Args>(args)...);

        return feeder_t { record };
    }
};

template<typename Log>
class scoped_pump : public feeder_t {
    const Log& log;

public:
    template<typename... Args>
    scoped_pump(const Log& log, log::record_t& record, Args&&... args) :
        feeder_t(record),
        log(log)
    {
        record.attributes.insert(keyword::message() = aux::format(record.attributes, std::forward<Args>(args)...));
    }

    ~scoped_pump() {
        log.push(std::move(record));
    }
};

template<typename Log, typename... Args>
scoped_pump<Log> make_scoped_pump(Log& log, log::record_t& record, Args&&... args) {
    return scoped_pump<Log>(log, record, std::forward<Args>(args)...);
}

} // namespace aux

} // namespace blackhole

#define BH_LOG(__log__, level, ...) \
    for (blackhole::log::record_t record = __log__.open_record(level); record.valid(); record.attributes.clear()) \
        blackhole::aux::make_scoped_pump(__log__, record, __VA_ARGS__)
<commit_msg>[Refactoring] No need for defining separate type for attributes feeding.<commit_after>#pragma once

#include <type_traits>
#include <tuple>
#include <ostream>

#include "blackhole/detail/traits/attributes/pack/feeder.hpp"

#include "format/message/extraction.hpp"
#include "format/message/insitu.hpp"
#include "logger.hpp"
#include "utils/format.hpp"

namespace blackhole {

namespace aux {

template<typename Log>
class scoped_pump {
    const Log& log;
    log::record_t& record;

public:
    template<typename... Args>
    scoped_pump(const Log& log, log::record_t& record, Args&&... args) :
        log(log),
        record(record)
    {
        record.attributes.insert(keyword::message() = aux::format(record.attributes, std::forward<Args>(args)...));
    }

    ~scoped_pump() {
        log.push(std::move(record));
    }

    // Using initializer lists is problematic, since GCC 4.6 has an extension which can
    // decude { { "value", "42" } } expressions in function call to the initializer list
    // without looking at the overload resolution. This violates C++11 standard, which
    // allows such deducing only for `auto` keyword (including range-based `for` loop).
    // Please, specify `attribute::list` type explicitly when using initializer lists to
    // provide attributes.
    scoped_pump& operator()(std::initializer_list<std::pair<std::string, log::attribute_value_t>>&& args) {
        for (auto it = args.begin(); it != args.end(); ++it) {
            record.attributes.insert(std::make_pair(it->first, log::attribute_t(it->second)));
        }

        return *this;
    }

    template<typename... Args>
    scoped_pump& operator()(Args&&... args) {
        static_assert((is_keyword_pack<Args...>::value || is_emplace_pack<Args...>::value),
                      "parameter pack must be either attribute pack or emplace pack");
        typedef typename pack_determiner<Args...>::type pack_type;
        pack_feeder<pack_type>::feed(record, std::forward<Args>(args)...);

        return *this;
    }
};

template<typename Log, typename... Args>
scoped_pump<Log> make_scoped_pump(Log& log, log::record_t& record, Args&&... args) {
    return scoped_pump<Log>(log, record, std::forward<Args>(args)...);
}

} // namespace aux

} // namespace blackhole

#define BH_LOG(__log__, level, ...) \
    for (blackhole::log::record_t record = __log__.open_record(level); record.valid(); record.attributes.clear()) \
        blackhole::aux::make_scoped_pump(__log__, record, __VA_ARGS__)
<|endoftext|>


<commit_before>/*
* TLS Server Proxy
* (C) 2014,2015 Jack Lloyd
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include "apps.h"

#if defined(BOTAN_HAS_TLS)

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

#define _GLIBCXX_HAVE_GTHR_DEFAULT
#include <boost/asio.hpp>
#include <boost/bind.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/enable_shared_from_this.hpp>

#include <botan/tls_server.h>
#include <botan/x509cert.h>
#include <botan/pkcs8.h>
#include <botan/auto_rng.h>

#if defined(BOTAN_HAS_TLS_SQLITE3_SESSION_MANAGER)
  #include <botan/tls_session_manager_sqlite.h>
#endif

#include "credentials.h"

using boost::asio::ip::tcp;

namespace Botan {

namespace {

inline void log_exception(const char* where, const std::exception& e)
   {
   std::cout << where << ' ' << e.what() << std::endl;
   }

inline void log_error(const char* where, const boost::system::error_code& error)
   {
   std::cout << where << ' ' << error.message() << std::endl;
   }

inline void log_binary_message(const char* where, const byte buf[], size_t buf_len)
   {
   //std::cout << where << ' ' << hex_encode(buf, buf_len) << std::endl;
   }

void log_text_message(const char* where,  const byte buf[], size_t buf_len)
   {
   //const char* c = reinterpret_cast<const char*>(buf);
   //std::cout << where << ' ' << std::string(c, c + buf_len)  << std::endl;
   }

class tls_proxy_session : public boost::enable_shared_from_this<tls_proxy_session>
   {
   public:
      enum { readbuf_size = 4 * 1024 };

      typedef boost::shared_ptr<tls_proxy_session> pointer;

      static pointer create(boost::asio::io_service& io,
                            TLS::Session_Manager& session_manager,
                            Credentials_Manager& credentials,
                            TLS::Policy& policy,
                            RandomNumberGenerator& rng,
                            tcp::resolver::iterator endpoints)
         {
         return pointer(
            new tls_proxy_session(
               io,
               session_manager,
               credentials,
               policy,
               rng,
               endpoints)
            );
         }

      tcp::socket& client_socket() { return m_client_socket; }

      void start()
         {
         m_c2p.resize(readbuf_size);

         client_read(boost::system::error_code(), 0); // start read loop
         }

      void stop() { m_client_socket.close(); }

   private:
      tls_proxy_session(boost::asio::io_service& io,
                        TLS::Session_Manager& session_manager,
                        Credentials_Manager& credentials,
                        TLS::Policy& policy,
                        RandomNumberGenerator& rng,
                        tcp::resolver::iterator endpoints) :
         m_strand(io),
         m_server_endpoints(endpoints),
         m_client_socket(io),
         m_server_socket(io),
         m_tls(boost::bind(&tls_proxy_session::tls_proxy_write_to_client, this, _1, _2),
               boost::bind(&tls_proxy_session::tls_client_write_to_proxy, this, _1, _2),
               boost::bind(&tls_proxy_session::tls_alert_cb, this, _1, _2, _3),
               boost::bind(&tls_proxy_session::tls_handshake_complete, this, _1),
               session_manager,
               credentials,
               policy,
               rng)
         {
         }

      void client_read(const boost::system::error_code& error,
                       size_t bytes_transferred)
         {
         if(error)
            {
            log_error("Read failed", error);
            stop();
            return;
            }

         try
            {
            if(!m_tls.is_active())
               log_binary_message("From client", &m_c2p[0], bytes_transferred);
            m_tls.received_data(&m_c2p[0], bytes_transferred);
            }
         catch(std::exception& e)
            {
            log_exception("TLS connection failed", e);
            stop();
            return;
            }

         m_client_socket.async_read_some(
            boost::asio::buffer(&m_c2p[0], m_c2p.size()),
            m_strand.wrap(boost::bind(&tls_proxy_session::client_read, shared_from_this(),
                                      boost::asio::placeholders::error,
                                      boost::asio::placeholders::bytes_transferred)));
         }

      void handle_client_write_completion(const boost::system::error_code& error)
         {
         if(error)
            {
            log_error("Client write", error);
            stop();
            return;
            }

         m_p2c.clear();
         tls_proxy_write_to_client(nullptr, 0); // initiate another write if needed
         }

      void handle_server_write_completion(const boost::system::error_code& error)
         {
         if(error)
            {
            log_error("Server write", error);
            stop();
            return;
            }

         m_p2s.clear();
         proxy_write_to_server(nullptr, 0); // initiate another write if needed
         }

      void tls_client_write_to_proxy(const byte buf[], size_t buf_len)
         {
         // Immediately bounce message to server
         proxy_write_to_server(buf, buf_len);
         }

      void tls_proxy_write_to_client(const byte buf[], size_t buf_len)
         {
         if(buf_len > 0)
            m_p2c_pending.insert(m_p2c_pending.end(), buf, buf + buf_len);

         // no write now active and we still have output pending
         if(m_p2c.empty() && !m_p2c_pending.empty())
            {
            std::swap(m_p2c_pending, m_p2c);

            //log_binary_message("To Client", &m_p2c[0], m_p2c.size());

            boost::asio::async_write(
               m_client_socket,
               boost::asio::buffer(&m_p2c[0], m_p2c.size()),
               m_strand.wrap(boost::bind(
                                &tls_proxy_session::handle_client_write_completion,
                                shared_from_this(),
                                boost::asio::placeholders::error)));
            }
         }

      void proxy_write_to_server(const byte buf[], size_t buf_len)
         {
         if(buf_len > 0)
            m_p2s_pending.insert(m_p2s_pending.end(), buf, buf + buf_len);

         // no write now active and we still have output pending
         if(m_p2s.empty() && !m_p2s_pending.empty())
            {
            std::swap(m_p2s_pending, m_p2s);

            log_text_message("To Server", &m_p2s[0], m_p2s.size());

            boost::asio::async_write(
               m_server_socket,
               boost::asio::buffer(&m_p2s[0], m_p2s.size()),
               m_strand.wrap(boost::bind(
                                &tls_proxy_session::handle_server_write_completion,
                                shared_from_this(),
                                boost::asio::placeholders::error)));
            }
         }

      void server_read(const boost::system::error_code& error,
                       size_t bytes_transferred)
         {
         if(error)
            {
            log_error("Server read failed", error);
            stop();
            return;
            }

         try
            {
            if(bytes_transferred)
               {
               log_text_message("Server to client", &m_s2p[0], m_s2p.size());
               log_binary_message("Server to client", &m_s2p[0], m_s2p.size());
               m_tls.send(&m_s2p[0], bytes_transferred);
               }
            }
         catch(std::exception& e)
            {
            log_exception("TLS connection failed", e);
            stop();
            return;
            }

         m_s2p.resize(readbuf_size);

         m_server_socket.async_read_some(
            boost::asio::buffer(&m_s2p[0], m_s2p.size()),
            m_strand.wrap(boost::bind(&tls_proxy_session::server_read, shared_from_this(),
                                      boost::asio::placeholders::error,
                                      boost::asio::placeholders::bytes_transferred)));
         }

      bool tls_handshake_complete(const TLS::Session& session)
         {
         //std::cout << "Handshake from client complete\n";

         m_hostname = session.server_info().hostname();

         if(m_hostname != "")
            std::cout << "Client requested hostname '" << m_hostname << "'\n";

         async_connect(m_server_socket, m_server_endpoints,
                       [this](boost::system::error_code ec, tcp::resolver::iterator endpoint)
                       {
                       if(ec)
                          {
                          log_error("Server connection", ec);
                          return;
                          }

                       server_read(boost::system::error_code(), 0); // start read loop
                       proxy_write_to_server(nullptr, 0);
                       });
         return true;
         }

      void tls_alert_cb(TLS::Alert alert, const byte[], size_t)
         {
         if(alert.type() == TLS::Alert::CLOSE_NOTIFY)
            {
            m_tls.close();
            return;
            }
         else
            std::cout << "Alert " << alert.type_string() << "\n";
         }

      boost::asio::io_service::strand m_strand;

      tcp::resolver::iterator m_server_endpoints;

      tcp::socket m_client_socket;
      tcp::socket m_server_socket;

      TLS::Server m_tls;
      std::string m_hostname;

      std::vector<byte> m_c2p;
      std::vector<byte> m_p2c;
      std::vector<byte> m_p2c_pending;

      std::vector<byte> m_s2p;
      std::vector<byte> m_p2s;
      std::vector<byte> m_p2s_pending;
   };

class tls_proxy_server
   {
   public:
      typedef tls_proxy_session session;

      tls_proxy_server(boost::asio::io_service& io, unsigned short port,
                       tcp::resolver::iterator endpoints,
                       Credentials_Manager& creds,
                       TLS::Policy& policy,
                       TLS::Session_Manager& session_mgr,
                       RandomNumberGenerator& rng) :
         m_acceptor(io, tcp::endpoint(tcp::v4(), port)),
         m_server_endpoints(endpoints),
         m_creds(creds),
         m_policy(policy),
         m_session_manager(session_mgr),
         m_rng(rng)
         {
         session::pointer new_session = make_session();

         m_acceptor.async_accept(
            new_session->client_socket(),
            boost::bind(
               &tls_proxy_server::handle_accept,
               this,
               new_session,
               boost::asio::placeholders::error)
            );
         }

   private:
      session::pointer make_session()
         {
         return session::create(
            m_acceptor.get_io_service(),
            m_session_manager,
            m_creds,
            m_policy,
            m_rng,
            m_server_endpoints
            );
         }

      void handle_accept(session::pointer new_session,
                         const boost::system::error_code& error)
         {
         if (!error)
            {
            new_session->start();

            new_session = make_session();

            m_acceptor.async_accept(
               new_session->client_socket(),
               boost::bind(
                  &tls_proxy_server::handle_accept,
                  this,
                  new_session,
                  boost::asio::placeholders::error)
               );
            }
         }

      tcp::acceptor m_acceptor;
      tcp::resolver::iterator m_server_endpoints;

      Credentials_Manager& m_creds;
      TLS::Policy& m_policy;
      TLS::Session_Manager& m_session_manager;
      RandomNumberGenerator& m_rng;
   };

size_t choose_thread_count()
   {
   size_t result = std::thread::hardware_concurrency();

   if(result)
      return result;

   return 2;
   }

int tls_proxy(int argc, char* argv[])
   {
   if(argc != 6)
      {
      std::cout << "Usage: " << argv[0] << " listen_port target_host target_port server_cert server_key\n";
      return 1;
      }

   const size_t listen_port = to_u32bit(argv[1]);
   const std::string target = argv[2];
   const std::string target_port = argv[3];

   const std::string server_crt = argv[4];
   const std::string server_key = argv[5];

   const size_t num_threads = choose_thread_count(); // make configurable

   AutoSeeded_RNG rng;
   Basic_Credentials_Manager creds(rng, server_crt, server_key);

   TLS::Policy policy; // TODO: Read policy from text file

   try
      {
      boost::asio::io_service io;

      tcp::resolver resolver(io);
      auto server_endpoint_iterator = resolver.resolve({ target, target_port });

#if defined(BOTAN_HAS_TLS_SQLITE3_SESSION_MANAGER)
      // Todo: make configurable
      const std::string sessions_passphrase = "correct horse battery staple";
      const std::string sessions_db = "sessions.db";
      TLS::Session_Manager_SQLite sessions(sessions_passphrase, rng, sessions_db);
#else
      TLS::Session_Manager_In_Memory sessions(rng);
#endif

      tls_proxy_server server(io, listen_port, server_endpoint_iterator, creds, policy, sessions, rng);

      std::vector<std::shared_ptr<std::thread>> threads;

      for(size_t i = 2; i <= num_threads; ++i)
         threads.push_back(std::make_shared<std::thread>([&io]() { io.run(); }));

      io.run();

      for (size_t i = 0; i < threads.size(); ++i)
         threads[i]->join();
      }
   catch (std::exception& e)
      {
      std::cerr << e.what() << std::endl;
      }

   return 0;
   }

}

}

REGISTER_APP(tls_proxy);

#endif
<commit_msg>Proxy tweaks<commit_after>/*
* TLS Server Proxy
* (C) 2014,2015 Jack Lloyd
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include "apps.h"

#if defined(BOTAN_HAS_TLS)

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

#define _GLIBCXX_HAVE_GTHR_DEFAULT
#include <boost/asio.hpp>
#include <boost/bind.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/enable_shared_from_this.hpp>

#include <botan/tls_server.h>
#include <botan/x509cert.h>
#include <botan/pkcs8.h>
#include <botan/auto_rng.h>

#if defined(BOTAN_HAS_SYSTEM_RNG)
  #include <botan/system_rng.h>
#endif

#if defined(BOTAN_HAS_TLS_SQLITE3_SESSION_MANAGER)
  #include <botan/tls_session_manager_sqlite.h>
#endif

#include "credentials.h"

using boost::asio::ip::tcp;

namespace Botan {

namespace {

inline void log_exception(const char* where, const std::exception& e)
   {
   std::cout << where << ' ' << e.what() << std::endl;
   }

inline void log_error(const char* where, const boost::system::error_code& error)
   {
   //std::cout << where << ' ' << error.message() << std::endl;
   }

inline void log_binary_message(const char* where, const byte buf[], size_t buf_len)
   {
   //std::cout << where << ' ' << hex_encode(buf, buf_len) << std::endl;
   }

void log_text_message(const char* where,  const byte buf[], size_t buf_len)
   {
   //const char* c = reinterpret_cast<const char*>(buf);
   //std::cout << where << ' ' << std::string(c, c + buf_len)  << std::endl;
   }

class tls_proxy_session : public boost::enable_shared_from_this<tls_proxy_session>
   {
   public:
      enum { readbuf_size = 4 * 1024 };

      typedef boost::shared_ptr<tls_proxy_session> pointer;

      static pointer create(boost::asio::io_service& io,
                            TLS::Session_Manager& session_manager,
                            Credentials_Manager& credentials,
                            TLS::Policy& policy,
                            RandomNumberGenerator& rng,
                            tcp::resolver::iterator endpoints)
         {
         return pointer(
            new tls_proxy_session(
               io,
               session_manager,
               credentials,
               policy,
               rng,
               endpoints)
            );
         }

      tcp::socket& client_socket() { return m_client_socket; }

      void start()
         {
         m_c2p.resize(readbuf_size);

         client_read(boost::system::error_code(), 0); // start read loop
         }

      void stop()
         {
         m_tls.close();
         m_client_socket.close();
         m_server_socket.close();
         }

   private:
      tls_proxy_session(boost::asio::io_service& io,
                        TLS::Session_Manager& session_manager,
                        Credentials_Manager& credentials,
                        TLS::Policy& policy,
                        RandomNumberGenerator& rng,
                        tcp::resolver::iterator endpoints) :
         m_strand(io),
         m_server_endpoints(endpoints),
         m_client_socket(io),
         m_server_socket(io),
         m_tls(boost::bind(&tls_proxy_session::tls_proxy_write_to_client, this, _1, _2),
               boost::bind(&tls_proxy_session::tls_client_write_to_proxy, this, _1, _2),
               boost::bind(&tls_proxy_session::tls_alert_cb, this, _1, _2, _3),
               boost::bind(&tls_proxy_session::tls_handshake_complete, this, _1),
               session_manager,
               credentials,
               policy,
               rng)
         {
         }

      void client_read(const boost::system::error_code& error,
                       size_t bytes_transferred)
         {
         if(error)
            {
            log_error("Read failed", error);
            stop();
            return;
            }

         try
            {
            if(!m_tls.is_active())
               log_binary_message("From client", &m_c2p[0], bytes_transferred);
            m_tls.received_data(&m_c2p[0], bytes_transferred);
            }
         catch(std::exception& e)
            {
            log_exception("TLS connection failed", e);
            stop();
            return;
            }

         m_client_socket.async_read_some(
            boost::asio::buffer(&m_c2p[0], m_c2p.size()),
            m_strand.wrap(boost::bind(&tls_proxy_session::client_read, shared_from_this(),
                                      boost::asio::placeholders::error,
                                      boost::asio::placeholders::bytes_transferred)));
         }

      void handle_client_write_completion(const boost::system::error_code& error)
         {
         if(error)
            {
            log_error("Client write", error);
            stop();
            return;
            }

         m_p2c.clear();
         tls_proxy_write_to_client(nullptr, 0); // initiate another write if needed
         }

      void handle_server_write_completion(const boost::system::error_code& error)
         {
         if(error)
            {
            log_error("Server write", error);
            stop();
            return;
            }

         m_p2s.clear();
         proxy_write_to_server(nullptr, 0); // initiate another write if needed
         }

      void tls_client_write_to_proxy(const byte buf[], size_t buf_len)
         {
         // Immediately bounce message to server
         proxy_write_to_server(buf, buf_len);
         }

      void tls_proxy_write_to_client(const byte buf[], size_t buf_len)
         {
         if(buf_len > 0)
            m_p2c_pending.insert(m_p2c_pending.end(), buf, buf + buf_len);

         // no write now active and we still have output pending
         if(m_p2c.empty() && !m_p2c_pending.empty())
            {
            std::swap(m_p2c_pending, m_p2c);

            //log_binary_message("To Client", &m_p2c[0], m_p2c.size());

            boost::asio::async_write(
               m_client_socket,
               boost::asio::buffer(&m_p2c[0], m_p2c.size()),
               m_strand.wrap(boost::bind(
                                &tls_proxy_session::handle_client_write_completion,
                                shared_from_this(),
                                boost::asio::placeholders::error)));
            }
         }

      void proxy_write_to_server(const byte buf[], size_t buf_len)
         {
         if(buf_len > 0)
            m_p2s_pending.insert(m_p2s_pending.end(), buf, buf + buf_len);

         // no write now active and we still have output pending
         if(m_p2s.empty() && !m_p2s_pending.empty())
            {
            std::swap(m_p2s_pending, m_p2s);

            log_text_message("To Server", &m_p2s[0], m_p2s.size());

            boost::asio::async_write(
               m_server_socket,
               boost::asio::buffer(&m_p2s[0], m_p2s.size()),
               m_strand.wrap(boost::bind(
                                &tls_proxy_session::handle_server_write_completion,
                                shared_from_this(),
                                boost::asio::placeholders::error)));
            }
         }

      void server_read(const boost::system::error_code& error,
                       size_t bytes_transferred)
         {
         if(error)
            {
            log_error("Server read failed", error);
            stop();
            return;
            }

         try
            {
            if(bytes_transferred)
               {
               log_text_message("Server to client", &m_s2p[0], m_s2p.size());
               log_binary_message("Server to client", &m_s2p[0], m_s2p.size());
               m_tls.send(&m_s2p[0], bytes_transferred);
               }
            }
         catch(std::exception& e)
            {
            log_exception("TLS connection failed", e);
            stop();
            return;
            }

         m_s2p.resize(readbuf_size);

         m_server_socket.async_read_some(
            boost::asio::buffer(&m_s2p[0], m_s2p.size()),
            m_strand.wrap(boost::bind(&tls_proxy_session::server_read, shared_from_this(),
                                      boost::asio::placeholders::error,
                                      boost::asio::placeholders::bytes_transferred)));
         }

      bool tls_handshake_complete(const TLS::Session& session)
         {
         //std::cout << "Handshake from client complete\n";

         m_hostname = session.server_info().hostname();

         if(m_hostname != "")
            std::cout << "Client requested hostname '" << m_hostname << "'\n";

         async_connect(m_server_socket, m_server_endpoints,
                       [this](boost::system::error_code ec, tcp::resolver::iterator endpoint)
                       {
                       if(ec)
                          {
                          log_error("Server connection", ec);
                          return;
                          }

                       server_read(boost::system::error_code(), 0); // start read loop
                       proxy_write_to_server(nullptr, 0);
                       });
         return true;
         }

      void tls_alert_cb(TLS::Alert alert, const byte[], size_t)
         {
         if(alert.type() == TLS::Alert::CLOSE_NOTIFY)
            {
            m_tls.close();
            return;
            }
         else
            std::cout << "Alert " << alert.type_string() << "\n";
         }

      boost::asio::io_service::strand m_strand;

      tcp::resolver::iterator m_server_endpoints;

      tcp::socket m_client_socket;
      tcp::socket m_server_socket;

      TLS::Server m_tls;
      std::string m_hostname;

      std::vector<byte> m_c2p;
      std::vector<byte> m_p2c;
      std::vector<byte> m_p2c_pending;

      std::vector<byte> m_s2p;
      std::vector<byte> m_p2s;
      std::vector<byte> m_p2s_pending;
   };

class tls_proxy_server
   {
   public:
      typedef tls_proxy_session session;

      tls_proxy_server(boost::asio::io_service& io, unsigned short port,
                       tcp::resolver::iterator endpoints,
                       Credentials_Manager& creds,
                       TLS::Policy& policy,
                       TLS::Session_Manager& session_mgr,
                       RandomNumberGenerator& rng) :
         m_acceptor(io, tcp::endpoint(tcp::v4(), port)),
         m_server_endpoints(endpoints),
         m_creds(creds),
         m_policy(policy),
         m_session_manager(session_mgr),
         m_rng(rng)
         {
         session::pointer new_session = make_session();

         m_acceptor.async_accept(
            new_session->client_socket(),
            boost::bind(
               &tls_proxy_server::handle_accept,
               this,
               new_session,
               boost::asio::placeholders::error)
            );
         }

   private:
      session::pointer make_session()
         {
         return session::create(
            m_acceptor.get_io_service(),
            m_session_manager,
            m_creds,
            m_policy,
            m_rng,
            m_server_endpoints
            );
         }

      void handle_accept(session::pointer new_session,
                         const boost::system::error_code& error)
         {
         if (!error)
            {
            new_session->start();

            new_session = make_session();

            m_acceptor.async_accept(
               new_session->client_socket(),
               boost::bind(
                  &tls_proxy_server::handle_accept,
                  this,
                  new_session,
                  boost::asio::placeholders::error)
               );
            }
         }

      tcp::acceptor m_acceptor;
      tcp::resolver::iterator m_server_endpoints;

      Credentials_Manager& m_creds;
      TLS::Policy& m_policy;
      TLS::Session_Manager& m_session_manager;
      RandomNumberGenerator& m_rng;
   };

size_t choose_thread_count()
   {
   size_t result = std::thread::hardware_concurrency();

   if(result)
      return result;

   return 2;
   }

int tls_proxy(int argc, char* argv[])
   {
   if(argc != 6)
      {
      std::cout << "Usage: " << argv[0] << " listen_port target_host target_port server_cert server_key\n";
      return 1;
      }

   const size_t listen_port = to_u32bit(argv[1]);
   const std::string target = argv[2];
   const std::string target_port = argv[3];

   const std::string server_crt = argv[4];
   const std::string server_key = argv[5];

   const size_t num_threads = choose_thread_count(); // make configurable

#if defined(BOTAN_HAS_SYSTEM_RNG)
   RandomNumberGenerator& rng = system_rng();
#else
   Serialized_RNG rng;
#endif
   Basic_Credentials_Manager creds(rng, server_crt, server_key);

   TLS::Policy policy; // TODO: Read policy from text file

   try
      {
      boost::asio::io_service io;

      tcp::resolver resolver(io);
      auto server_endpoint_iterator = resolver.resolve({ target, target_port });

#if defined(BOTAN_HAS_TLS_SQLITE3_SESSION_MANAGER)
      // Todo: make configurable
      const std::string sessions_passphrase = "correct horse battery staple";
      const std::string sessions_db = "sessions.db";
      TLS::Session_Manager_SQLite sessions(sessions_passphrase, rng, sessions_db);
#else
      TLS::Session_Manager_In_Memory sessions(rng);
#endif

      tls_proxy_server server(io, listen_port, server_endpoint_iterator, creds, policy, sessions, rng);

      std::vector<std::shared_ptr<std::thread>> threads;

      for(size_t i = 2; i <= num_threads; ++i)
         threads.push_back(std::make_shared<std::thread>([&io]() { io.run(); }));

      io.run();

      for (size_t i = 0; i < threads.size(); ++i)
         threads[i]->join();
      }
   catch (std::exception& e)
      {
      std::cerr << e.what() << std::endl;
      }

   return 0;
   }

}

}

REGISTER_APP(tls_proxy);

#endif
<|endoftext|>


<commit_before>// Include this #define to use SystemTraceControlGuid in Evntrace.h.
#define INITGUID

#include <windows.h>

#include <evntcons.h>
#include <evntrace.h>
#include <stdio.h>

static void err(char *s) {
  // Retrieve the system error message for the last-error code
  LPVOID msg;
  auto e = GetLastError();
  FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM |
                    FORMAT_MESSAGE_IGNORE_INSERTS,
                NULL, e, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                (LPTSTR)&msg, 0, NULL);

  // Display the error message and exit the process
  printf("%s failed with error %d: %s\n", s, e, msg);
  ExitProcess(e);
}

static void WINAPI EventRecordCallback(EVENT_RECORD *EventRecord) {
  EVENT_HEADER &Header = EventRecord->EventHeader;

  UCHAR ProcessorNumber = EventRecord->BufferContext.ProcessorNumber;
  ULONG ThreadID = Header.ThreadId;

  // Process event here.
}

static TRACEHANDLE trace;

static DWORD WINAPI processThread(LPVOID Parameter) {
  ProcessTrace(&trace, 1, 0, 0);
  return (0);
}

int main(int argc, char **argv) {
  auto BufferSize = sizeof(EVENT_TRACE_PROPERTIES) + sizeof(KERNEL_LOGGER_NAME);
  auto properties = (EVENT_TRACE_PROPERTIES *)malloc(BufferSize);

  ZeroMemory(properties, BufferSize);
  properties->Wnode.BufferSize = BufferSize;
  properties->Wnode.Flags = WNODE_FLAG_TRACED_GUID;
  properties->Wnode.ClientContext = 3;
  properties->Wnode.Guid = SystemTraceControlGuid;
  properties->EnableFlags = EVENT_TRACE_FLAG_NETWORK_TCPIP;
  properties->LogFileMode = EVENT_TRACE_REAL_TIME_MODE;
  properties->LoggerNameOffset = sizeof(EVENT_TRACE_PROPERTIES);
  strcpy((char *)properties + properties->LoggerNameOffset, KERNEL_LOGGER_NAME);

  ControlTrace(0, KERNEL_LOGGER_NAME, properties, EVENT_TRACE_CONTROL_STOP);

  auto status =
      StartTrace((PTRACEHANDLE)&trace, KERNEL_LOGGER_NAME, properties);
  if (ERROR_SUCCESS != status)
    err("StartTrace");

  EVENT_TRACE_LOGFILE LogFile = {0};
  LogFile.LoggerName = KERNEL_LOGGER_NAME;
  LogFile.ProcessTraceMode =
      (PROCESS_TRACE_MODE_REAL_TIME | PROCESS_TRACE_MODE_EVENT_RECORD |
       PROCESS_TRACE_MODE_RAW_TIMESTAMP);
  LogFile.EventRecordCallback = EventRecordCallback;
  trace = OpenTrace(&LogFile);
  CreateThread(0, 0, processThread, 0, 0, 0);

  wprintf(L"Press any key to end trace session ");
  getchar();

  ControlTrace(0, KERNEL_LOGGER_NAME, properties, EVENT_TRACE_CONTROL_STOP);
  return 0;
}
<commit_msg>refactor<commit_after>// Include this #define to use SystemTraceControlGuid in Evntrace.h.
#define INITGUID

#include <windows.h>

#include <evntcons.h>
#include <evntrace.h>
#include <stdio.h>

static void err(char *s) {
  // Retrieve the system error message for the last-error code
  char *msg;
  auto e = GetLastError();
  FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM |
                    FORMAT_MESSAGE_IGNORE_INSERTS,
                0, e, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (char *)&msg,
                0, 0);

  // Display the error message and exit the process
  printf("%s failed with error %d: %s\n", s, e, msg);
  ExitProcess(e);
}

static void WINAPI EventRecordCallback(EVENT_RECORD *EventRecord) {
  EVENT_HEADER &Header = EventRecord->EventHeader;

  UCHAR ProcessorNumber = EventRecord->BufferContext.ProcessorNumber;
  ULONG ThreadID = Header.ThreadId;

  // Process event here.
}

static TRACEHANDLE trace;

static DWORD WINAPI processThread(void *) {
  ProcessTrace(&trace, 1, 0, 0);
  return 0;
}

int main(int argc, char **argv) {
  auto BufferSize = sizeof(EVENT_TRACE_PROPERTIES) + sizeof(KERNEL_LOGGER_NAME);
  auto properties = (EVENT_TRACE_PROPERTIES *)malloc(BufferSize);

  ZeroMemory(properties, BufferSize);
  properties->Wnode.BufferSize = BufferSize;
  properties->Wnode.Flags = WNODE_FLAG_TRACED_GUID;
  properties->Wnode.ClientContext = 3;
  properties->Wnode.Guid = SystemTraceControlGuid;
  properties->EnableFlags = EVENT_TRACE_FLAG_NETWORK_TCPIP;
  properties->LogFileMode = EVENT_TRACE_REAL_TIME_MODE;
  properties->LoggerNameOffset = sizeof(EVENT_TRACE_PROPERTIES);
  strcpy((char *)properties + properties->LoggerNameOffset, KERNEL_LOGGER_NAME);

  ControlTrace(0, KERNEL_LOGGER_NAME, properties, EVENT_TRACE_CONTROL_STOP);

  auto status =
      StartTrace((PTRACEHANDLE)&trace, KERNEL_LOGGER_NAME, properties);
  if (ERROR_SUCCESS != status)
    err("StartTrace");

  EVENT_TRACE_LOGFILE LogFile = {0};
  LogFile.LoggerName = KERNEL_LOGGER_NAME;
  LogFile.ProcessTraceMode =
      (PROCESS_TRACE_MODE_REAL_TIME | PROCESS_TRACE_MODE_EVENT_RECORD |
       PROCESS_TRACE_MODE_RAW_TIMESTAMP);
  LogFile.EventRecordCallback = EventRecordCallback;
  trace = OpenTrace(&LogFile);
  CreateThread(0, 0, processThread, 0, 0, 0);

  wprintf(L"Press any key to end trace session ");
  getchar();

  ControlTrace(0, KERNEL_LOGGER_NAME, properties, EVENT_TRACE_CONTROL_STOP);
  return 0;
}
<|endoftext|>