<commit_before>// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "media/filters/ffmpeg_audio_decoder.h"

#include "media/base/callback.h"
#include "media/base/data_buffer.h"
#include "media/base/limits.h"
#include "media/ffmpeg/ffmpeg_common.h"
#include "media/filters/ffmpeg_demuxer.h"

namespace media {

// Size of the decoded audio buffer.
const size_t FFmpegAudioDecoder::kOutputBufferSize =
    AVCODEC_MAX_AUDIO_FRAME_SIZE;

FFmpegAudioDecoder::FFmpegAudioDecoder()
    : codec_context_(NULL) {
}

FFmpegAudioDecoder::~FFmpegAudioDecoder() {
}

// static
bool FFmpegAudioDecoder::IsMediaFormatSupported(const MediaFormat& format) {
  std::string mime_type;
  return format.GetAsString(MediaFormat::kMimeType, &mime_type) &&
      mime_type::kFFmpegAudio == mime_type;
}

void FFmpegAudioDecoder::DoInitialize(DemuxerStream* demuxer_stream,
                                      bool* success,
                                      Task* done_cb) {
  AutoTaskRunner done_runner(done_cb);
  *success = false;

  // Get the AVStream by querying for the provider interface.
  AVStreamProvider* av_stream_provider;
  if (!demuxer_stream->QueryInterface(&av_stream_provider)) {
    return;
  }
  AVStream* av_stream = av_stream_provider->GetAVStream();

  // Grab the AVStream's codec context and make sure we have sensible values.
  codec_context_ = av_stream->codec;
  int bps = av_get_bits_per_sample_format(codec_context_->sample_fmt);
  DCHECK_GT(codec_context_->channels, 0);
  DCHECK_GT(bps, 0);
  DCHECK_GT(codec_context_->sample_rate, 0);
  if (codec_context_->channels == 0 ||
      static_cast<size_t>(codec_context_->channels) > Limits::kMaxChannels ||
      bps == 0 ||
      static_cast<size_t>(bps) > Limits::kMaxBPS ||
      codec_context_->sample_rate == 0 ||
      (static_cast<size_t>(codec_context_->sample_rate) >
       Limits::kMaxSampleRate)) {
    return;
  }

  // Serialize calls to avcodec_open().
  AVCodec* codec = avcodec_find_decoder(codec_context_->codec_id);
  if (!codec || avcodec_open(codec_context_, codec) < 0) {
    return;
  }

  // When calling avcodec_find_decoder(), |codec_context_| might be altered by
  // the decoder to give more accurate values of the output format of the
  // decoder. So set the media format after a decoder is allocated.
  // TODO(hclam): Reuse the information provided by the demuxer for now, we may
  // need to wait until the first buffer is decoded to know the correct
  // information.
  media_format_.SetAsInteger(MediaFormat::kChannels, codec_context_->channels);
  media_format_.SetAsInteger(MediaFormat::kSampleBits,
      av_get_bits_per_sample_format(codec_context_->sample_fmt));
  media_format_.SetAsInteger(MediaFormat::kSampleRate,
      codec_context_->sample_rate);
  media_format_.SetAsString(MediaFormat::kMimeType,
      mime_type::kUncompressedAudio);

  // Prepare the output buffer.
  output_buffer_.reset(static_cast<uint8*>(av_malloc(kOutputBufferSize)));
  if (!output_buffer_.get()) {
    host()->SetError(PIPELINE_ERROR_OUT_OF_MEMORY);
    return;
  }
  *success = true;
}

void FFmpegAudioDecoder::DoSeek(base::TimeDelta time, Task* done_cb) {
  avcodec_flush_buffers(codec_context_);
  estimated_next_timestamp_ = StreamSample::kInvalidTimestamp;
  done_cb->Run();
  delete done_cb;
}

void FFmpegAudioDecoder::DoDecode(Buffer* input, Task* done_cb) {
  AutoTaskRunner done_runner(done_cb);

  // Due to FFmpeg API changes we no longer have const read-only pointers.
  AVPacket packet;
  av_init_packet(&packet);
  packet.data = const_cast<uint8*>(input->GetData());
  packet.size = input->GetDataSize();

  int16_t* output_buffer = reinterpret_cast<int16_t*>(output_buffer_.get());
  int output_buffer_size = kOutputBufferSize;
  int result = avcodec_decode_audio3(codec_context_,
                                     output_buffer,
                                     &output_buffer_size,
                                     &packet);

  // TODO(ajwong): Consider if kOutputBufferSize should just be an int instead
  // of a size_t.
  if (result < 0 ||
      output_buffer_size < 0 ||
      static_cast<size_t>(output_buffer_size) > kOutputBufferSize) {
    LOG(INFO) << "Error decoding an audio frame with timestamp: "
              << input->GetTimestamp().InMicroseconds() << " us"
              << " , duration: "
              << input->GetDuration().InMicroseconds() << " us"
              << " , packet size: "
              << input->GetDataSize() << " bytes";
    return;
  }

  // If we have decoded something, enqueue the result.
  if (output_buffer_size) {
    DataBuffer* result_buffer = new DataBuffer(output_buffer_size);
    result_buffer->SetDataSize(output_buffer_size);
    uint8* data = result_buffer->GetWritableData();
    memcpy(data, output_buffer, output_buffer_size);

    // Determine the duration if the demuxer couldn't figure it out, otherwise
    // copy it over.
    if (input->GetDuration().ToInternalValue() == 0) {
      result_buffer->SetDuration(CalculateDuration(output_buffer_size));
    } else {
      DCHECK(input->GetDuration() != StreamSample::kInvalidTimestamp);
      result_buffer->SetDuration(input->GetDuration());
    }

    // Use our estimate for the timestamp if |input| does not have one.
    // Otherwise, copy over the timestamp.
    if (input->GetTimestamp() == StreamSample::kInvalidTimestamp) {
      result_buffer->SetTimestamp(estimated_next_timestamp_);
    } else {
      result_buffer->SetTimestamp(input->GetTimestamp());
    }

    // Only use the timestamp of |result_buffer| to estimate the next timestamp
    // if it is valid (i.e. != StreamSample::kInvalidTimestamp).  Otherwise the
    // error will stack together and we will get a series of incorrect
    // timestamps.  In this case, this will maintain a series of zero
    // timestamps.
    if (result_buffer->GetTimestamp() != StreamSample::kInvalidTimestamp) {
      // Update our estimated timestamp for the next packet.
      estimated_next_timestamp_ = result_buffer->GetTimestamp() +
          result_buffer->GetDuration();
    }

    EnqueueResult(result_buffer);
    return;
  }

  // We can get a positive result but no decoded data.  This is ok because this
  // this can be a marker packet that only contains timestamp.  In this case we
  // save the timestamp for later use.
  if (result && !input->IsEndOfStream() &&
      input->GetTimestamp() != StreamSample::kInvalidTimestamp &&
      input->GetDuration() != StreamSample::kInvalidTimestamp) {
    estimated_next_timestamp_ = input->GetTimestamp() + input->GetDuration();
    return;
  }

  // Three conditions to meet to declare end of stream for this decoder:
  // 1. FFmpeg didn't read anything.
  // 2. FFmpeg didn't output anything.
  // 3. An end of stream buffer is received.
  if (result == 0 && output_buffer_size == 0 && input->IsEndOfStream()) {
    DataBuffer* result_buffer = new DataBuffer(0);
    result_buffer->SetTimestamp(input->GetTimestamp());
    result_buffer->SetDuration(input->GetDuration());
    EnqueueResult(result_buffer);
  }
}

base::TimeDelta FFmpegAudioDecoder::CalculateDuration(size_t size) {
  int64 denominator = codec_context_->channels *
      av_get_bits_per_sample_format(codec_context_->sample_fmt) / 8 *
      codec_context_->sample_rate;
  double microseconds = size /
      (denominator / static_cast<double>(base::Time::kMicrosecondsPerSecond));
  return base::TimeDelta::FromMicroseconds(static_cast<int64>(microseconds));
}

}  // namespace
<commit_msg>Clean up the audio timestamp and duration code.<commit_after>// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "media/filters/ffmpeg_audio_decoder.h"

#include "media/base/callback.h"
#include "media/base/data_buffer.h"
#include "media/base/limits.h"
#include "media/ffmpeg/ffmpeg_common.h"
#include "media/filters/ffmpeg_demuxer.h"

namespace media {

// Size of the decoded audio buffer.
const size_t FFmpegAudioDecoder::kOutputBufferSize =
    AVCODEC_MAX_AUDIO_FRAME_SIZE;

FFmpegAudioDecoder::FFmpegAudioDecoder()
    : codec_context_(NULL),
      estimated_next_timestamp_(StreamSample::kInvalidTimestamp) {
}

FFmpegAudioDecoder::~FFmpegAudioDecoder() {
}

// static
bool FFmpegAudioDecoder::IsMediaFormatSupported(const MediaFormat& format) {
  std::string mime_type;
  return format.GetAsString(MediaFormat::kMimeType, &mime_type) &&
      mime_type::kFFmpegAudio == mime_type;
}

void FFmpegAudioDecoder::DoInitialize(DemuxerStream* demuxer_stream,
                                      bool* success,
                                      Task* done_cb) {
  AutoTaskRunner done_runner(done_cb);
  *success = false;

  // Get the AVStream by querying for the provider interface.
  AVStreamProvider* av_stream_provider;
  if (!demuxer_stream->QueryInterface(&av_stream_provider)) {
    return;
  }
  AVStream* av_stream = av_stream_provider->GetAVStream();

  // Grab the AVStream's codec context and make sure we have sensible values.
  codec_context_ = av_stream->codec;
  int bps = av_get_bits_per_sample_format(codec_context_->sample_fmt);
  DCHECK_GT(codec_context_->channels, 0);
  DCHECK_GT(bps, 0);
  DCHECK_GT(codec_context_->sample_rate, 0);
  if (codec_context_->channels == 0 ||
      static_cast<size_t>(codec_context_->channels) > Limits::kMaxChannels ||
      bps == 0 ||
      static_cast<size_t>(bps) > Limits::kMaxBPS ||
      codec_context_->sample_rate == 0 ||
      (static_cast<size_t>(codec_context_->sample_rate) >
       Limits::kMaxSampleRate)) {
    return;
  }

  // Serialize calls to avcodec_open().
  AVCodec* codec = avcodec_find_decoder(codec_context_->codec_id);
  if (!codec || avcodec_open(codec_context_, codec) < 0) {
    return;
  }

  // When calling avcodec_find_decoder(), |codec_context_| might be altered by
  // the decoder to give more accurate values of the output format of the
  // decoder. So set the media format after a decoder is allocated.
  // TODO(hclam): Reuse the information provided by the demuxer for now, we may
  // need to wait until the first buffer is decoded to know the correct
  // information.
  media_format_.SetAsInteger(MediaFormat::kChannels, codec_context_->channels);
  media_format_.SetAsInteger(MediaFormat::kSampleBits,
      av_get_bits_per_sample_format(codec_context_->sample_fmt));
  media_format_.SetAsInteger(MediaFormat::kSampleRate,
      codec_context_->sample_rate);
  media_format_.SetAsString(MediaFormat::kMimeType,
      mime_type::kUncompressedAudio);

  // Prepare the output buffer.
  output_buffer_.reset(static_cast<uint8*>(av_malloc(kOutputBufferSize)));
  if (!output_buffer_.get()) {
    host()->SetError(PIPELINE_ERROR_OUT_OF_MEMORY);
    return;
  }
  *success = true;
}

void FFmpegAudioDecoder::DoSeek(base::TimeDelta time, Task* done_cb) {
  avcodec_flush_buffers(codec_context_);
  estimated_next_timestamp_ = StreamSample::kInvalidTimestamp;
  done_cb->Run();
  delete done_cb;
}

void FFmpegAudioDecoder::DoDecode(Buffer* input, Task* done_cb) {
  AutoTaskRunner done_runner(done_cb);

  // Due to FFmpeg API changes we no longer have const read-only pointers.
  AVPacket packet;
  av_init_packet(&packet);
  packet.data = const_cast<uint8*>(input->GetData());
  packet.size = input->GetDataSize();

  int16_t* output_buffer = reinterpret_cast<int16_t*>(output_buffer_.get());
  int output_buffer_size = kOutputBufferSize;
  int result = avcodec_decode_audio3(codec_context_,
                                     output_buffer,
                                     &output_buffer_size,
                                     &packet);

  // TODO(ajwong): Consider if kOutputBufferSize should just be an int instead
  // of a size_t.
  if (result < 0 ||
      output_buffer_size < 0 ||
      static_cast<size_t>(output_buffer_size) > kOutputBufferSize) {
    LOG(INFO) << "Error decoding an audio frame with timestamp: "
              << input->GetTimestamp().InMicroseconds() << " us"
              << " , duration: "
              << input->GetDuration().InMicroseconds() << " us"
              << " , packet size: "
              << input->GetDataSize() << " bytes";
    return;
  }

  // If we have decoded something, enqueue the result.
  if (output_buffer_size) {
    DataBuffer* result_buffer = new DataBuffer(output_buffer_size);
    result_buffer->SetDataSize(output_buffer_size);
    uint8* data = result_buffer->GetWritableData();
    memcpy(data, output_buffer, output_buffer_size);

    // We don't trust the demuxer, so always calculate the duration based on
    // the actual number of samples decoded.
    base::TimeDelta duration = CalculateDuration(output_buffer_size);
    result_buffer->SetDuration(duration);

    // Use an estimated timestamp unless the incoming buffer has a valid one.
    if (input->GetTimestamp() == StreamSample::kInvalidTimestamp) {
      result_buffer->SetTimestamp(estimated_next_timestamp_);

      // Keep the estimated timestamp invalid until we get an incoming buffer
      // with a valid timestamp.  This can happen during seeks, etc...
      if (estimated_next_timestamp_ != StreamSample::kInvalidTimestamp) {
        estimated_next_timestamp_ += duration;
      }
    } else {
      result_buffer->SetTimestamp(input->GetTimestamp());
      estimated_next_timestamp_ = input->GetTimestamp() + duration;
    }

    EnqueueResult(result_buffer);
    return;
  }

  // We can get a positive result but no decoded data.  This is ok because this
  // this can be a marker packet that only contains timestamp.  In this case we
  // save the timestamp for later use.
  if (result && !input->IsEndOfStream() &&
      input->GetTimestamp() != StreamSample::kInvalidTimestamp &&
      input->GetDuration() != StreamSample::kInvalidTimestamp) {
    estimated_next_timestamp_ = input->GetTimestamp() + input->GetDuration();
    return;
  }

  // Three conditions to meet to declare end of stream for this decoder:
  // 1. FFmpeg didn't read anything.
  // 2. FFmpeg didn't output anything.
  // 3. An end of stream buffer is received.
  if (result == 0 && output_buffer_size == 0 && input->IsEndOfStream()) {
    DataBuffer* result_buffer = new DataBuffer(0);
    result_buffer->SetTimestamp(input->GetTimestamp());
    result_buffer->SetDuration(input->GetDuration());
    EnqueueResult(result_buffer);
  }
}

base::TimeDelta FFmpegAudioDecoder::CalculateDuration(size_t size) {
  int64 denominator = codec_context_->channels *
      av_get_bits_per_sample_format(codec_context_->sample_fmt) / 8 *
      codec_context_->sample_rate;
  double microseconds = size /
      (denominator / static_cast<double>(base::Time::kMicrosecondsPerSecond));
  return base::TimeDelta::FromMicroseconds(static_cast<int64>(microseconds));
}

}  // namespace
<|endoftext|>


<commit_before>//////////////////////////////////////////////////////////////////////////
//
//  Copyright (c) 2011, Image Engine Design Inc. All rights reserved.
//  Copyright (c) 2012, John Haddon. All rights reserved.
//
//  Redistribution and use in source and binary forms, with or without
//  modification, are permitted provided that the following conditions are
//  met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
//     * Neither the name of Image Engine Design nor the names of any
//       other contributors to this software may be used to endorse or
//       promote products derived from this software without specific prior
//       written permission.
//
//  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
//  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
//  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
//  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
//  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
//  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
//  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
//  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
//  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
//  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
//  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
//////////////////////////////////////////////////////////////////////////

// This must come before the Cortex includes, because on OSX headers included
// by TBB define macros which conflict with the inline functions in ai_types.h.
#include "ai.h"

#include "IECore/CompoundParameter.h"
#include "IECore/SimpleTypedData.h"
#include "IECore/MessageHandler.h"
#include "IECore/DespatchTypedData.h"

#include "IECoreArnold/ToArnoldConverter.h"

using namespace IECore;
using namespace IECoreArnold;

IE_CORE_DEFINERUNTIMETYPED( ToArnoldConverter );

ToArnoldConverter::ToArnoldConverter( const std::string &description, IECore::TypeId supportedType )
	:	IECore::FromCoreConverter( description, supportedType )
{
}

ToArnoldConverter::~ToArnoldConverter()
{
}

AtNode *ToArnoldConverter::convert() const
{
	IECore::ConstCompoundObjectPtr operands = parameters()->getTypedValidatedValue<IECore::CompoundObject>();
	return doConversion( srcParameter()->getValidatedValue(), operands );
}

void ToArnoldConverter::setParameter( AtNode *node, const AtParamEntry *parameter, const IECore::Data *value )
{
	bool isArray = false;
	int type = AiParamGetType( parameter );
	if( type == AI_TYPE_ARRAY )
	{
		type = AiParamGetDefault( parameter )->ARRAY->type;
		isArray = true;
	}

	setParameterInternal( node, AiParamGetName( parameter ), type, isArray, value );
}

void ToArnoldConverter::setParameter( AtNode *node, const char *name, const IECore::Data *value )
{
	const AtParamEntry *parameter = AiNodeEntryLookUpParameter( AiNodeGetNodeEntry( node ), name );
	if( parameter )
	{
		setParameter( node, parameter, value );
	}
	else
	{
		bool array = false;
		int type = parameterType( value->typeId(), array );
		if( type != AI_TYPE_NONE )
		{
			std::string typeString = "constant ";
			if( array )
			{
				typeString += "ARRAY ";
			}
			typeString += AiParamGetTypeName( type );
			AiNodeDeclare( node, name, typeString.c_str() );
			setParameterInternal( node, name, type, array, value );
		}
		else
		{
			msg(
				Msg::Warning,
				"ToArnoldConverter::setParameter",
				boost::format( "Unsupported data type \"%s\" for name \"%s\"" ) % value->typeName() % name
			);
		}
	}
}

void ToArnoldConverter::setParameters( AtNode *node, const IECore::CompoundDataMap &values )
{
	for( CompoundDataMap::const_iterator it=values.begin(); it!=values.end(); it++ )
	{
		setParameter( node, it->first.value().c_str(), it->second );
	}
}

IECore::DataPtr ToArnoldConverter::getParameter( AtNode *node, const AtParamEntry *parameter )
{
	return getParameterInternal( node, AiParamGetName( parameter ), AiParamGetType( parameter ) );
}

IECore::DataPtr ToArnoldConverter::getParameter( AtNode *node, const AtUserParamEntry *parameter )
{
	return getParameterInternal( node, AiUserParamGetName( parameter ), AiUserParamGetType( parameter ) );
}

IECore::DataPtr ToArnoldConverter::getParameter( AtNode *node, const char *name )
{
	const AtParamEntry *parameter = AiNodeEntryLookUpParameter( AiNodeGetNodeEntry( node ), name );
	if( parameter )
	{
		return getParameter( node, parameter );
	}
	else
	{
		const AtUserParamEntry *userParameter = AiNodeLookUpUserParameter( node, name );
		if( userParameter )
		{
			return getParameter( node, userParameter );
		}
	}

	return 0;
}

void ToArnoldConverter::getParameters( AtNode *node, IECore::CompoundDataMap &values )
{
	/// \todo Non-user parameters

	AtUserParamIterator *it = AiNodeGetUserParamIterator( node );
	while( const AtUserParamEntry *param = AiUserParamIteratorGetNext( it ) )
	{
		DataPtr d = getParameter( node, param );
		if( d )
		{
			values[AiUserParamGetName( param )] = d;
		}
		else
		{
			msg(
				Msg::Warning,
				"ToArnoldConverter::getParameters",
				boost::format( "Unable to convert user parameter \"%s\"" ) % AiUserParamGetName( param )
			);
		}
	}
	AiUserParamIteratorDestroy( it );
}

int ToArnoldConverter::parameterType( IECore::TypeId dataType, bool &array )
{
	switch( dataType )
	{
		// non-array types

		case IntDataTypeId :
			array = false;
			return AI_TYPE_INT;
		case FloatDataTypeId :
			array = false;
			return AI_TYPE_FLOAT;
		case StringDataTypeId :
			array = false;
			return AI_TYPE_STRING;
		case Color3fDataTypeId :
			array = false;
			return AI_TYPE_RGB;
		case BoolDataTypeId :
			array = false;
			return AI_TYPE_BOOLEAN;

		// array types

		case IntVectorDataTypeId :
			array = true;
			return AI_TYPE_INT;
		case FloatVectorDataTypeId :
			array = true;
			return AI_TYPE_FLOAT;
		case StringVectorDataTypeId :
			array = true;
			return AI_TYPE_STRING;
		case Color3fVectorDataTypeId :
			array = true;
			return AI_TYPE_RGB;
		case BoolVectorDataTypeId :
			array = true;
			return AI_TYPE_BOOLEAN;
		default :
			return AI_TYPE_NONE;
	}
}

AtArray *ToArnoldConverter::dataToArray( const IECore::Data *data )
{
	bool isArray = false;
	int type = parameterType( data->typeId(), isArray );
	if( type == AI_TYPE_NONE || !isArray )
	{
		return 0;
	}

	const void *dataAddress = despatchTypedData<TypedDataAddress, TypeTraits::IsTypedData, DespatchTypedDataIgnoreError>( DataPtr( const_cast<Data *>( data ) ) );
	size_t dataSize = despatchTypedData<TypedDataSize, TypeTraits::IsTypedData, DespatchTypedDataIgnoreError>( DataPtr( const_cast<Data *>( data ) ) );
	return AiArrayConvert( dataSize, 1, type, dataAddress );
}

template<typename T>
static inline const T *dataCast( const char *name, const IECore::Data *data )
{
	const T *result = runTimeCast<const T>( data );
	if( result )
	{
		return result;
	}
	msg( Msg::Warning, "ToArnoldConverter::setParameter", boost::format( "Unsupported value type \"%s\" for parameter \"%s\" (expected %s)." ) % data->typeName() % name % T::staticTypeName() );
	return 0;
}

void ToArnoldConverter::setParameterInternal( AtNode *node, const char *name, int parameterType, bool array, const IECore::Data *value )
{
	if( array )
	{
		AtArray *a = dataToArray( value );
		if( !a )
		{
			msg( Msg::Warning, "ToArnoldConverter::setParameter", boost::format( "Unable to create array from data of type \"%s\" for parameter \"%s\"" ) % value->typeName() % name );
			return;
		}
		if( a->type != parameterType )
		{
			msg( Msg::Warning, "ToArnoldConverter::setParameter", boost::format( "Unable to create array of type %s from data of type \"%s\" for parameter \"%s\"" ) % AiParamGetTypeName( parameterType ) % value->typeName() % name );
			return;
		}
		AiNodeSetArray( node, name, a );
	}
	else
	{
		switch( parameterType )
		{
			case AI_TYPE_INT :
				if( const IntData *data = dataCast<IntData>( name, value ) )
				{
					AiNodeSetInt( node, name, data->readable() );
				}
				break;
			case AI_TYPE_FLOAT :
				if( const FloatData *data = dataCast<FloatData>( name, value ) )
				{
					AiNodeSetFlt( node, name, data->readable() );
				}
				break;
			case AI_TYPE_STRING :
				if( const StringData *data = dataCast<StringData>( name, value ) )
				{
					AiNodeSetStr( node, name, data->readable().c_str() );
				}
				break;
			case AI_TYPE_RGB :
				if( const Color3fData *data = dataCast<Color3fData>( name, value ) )
				{
					const Imath::Color3f &c = data->readable();
					AiNodeSetRGB( node, name, c[0], c[1], c[2] );
				}
				break;
			case AI_TYPE_ENUM :
				if( const StringData *data = dataCast<StringData>( name, value ) )
				{
					AiNodeSetStr( node, name, data->readable().c_str() );
				}
				break;
			case AI_TYPE_BOOLEAN :
			{
				if( const BoolData *data = dataCast<BoolData>( name, value ) )
				{
					AiNodeSetBool( node, name, data->readable() );
				}
				break;
			}
			default :
			{
				msg( Msg::Warning, "ToArnoldConverter::setParameter", boost::format( "Arnold parameter \"%s\" has unsupported type \"%s\"." ) % name % AiParamGetTypeName( parameterType ) );
			}
		}
	}
}

IECore::DataPtr ToArnoldConverter::getParameterInternal( AtNode *node, const char *name, int parameterType )
{
	switch( parameterType )
	{
		case AI_TYPE_BOOLEAN :
			return new BoolData( AiNodeGetBool( node, name ) );
		case AI_TYPE_INT :
			return new IntData( AiNodeGetInt( node, name ) );
		case AI_TYPE_FLOAT :
			return new FloatData( AiNodeGetFlt( node, name ) );
		case AI_TYPE_STRING :
			return new StringData( AiNodeGetStr( node, name ) );
	}
	return 0;
}

ToArnoldConverterPtr ToArnoldConverter::create( IECore::ObjectPtr object )
{
	const CreatorMap &m = creators();
	CreatorMap::const_iterator it = m.find( object->typeId() );
	if( it != m.end() )
	{
		return it->second( object );
	}
	return 0;
}

ToArnoldConverter::CreatorMap &ToArnoldConverter::creators()
{
	static CreatorMap m;
	return m;
}

<commit_msg>added AI_TYPE_BYTE case to ToArnoldConverter.cpp<commit_after>//////////////////////////////////////////////////////////////////////////
//
//  Copyright (c) 2011, Image Engine Design Inc. All rights reserved.
//  Copyright (c) 2012, John Haddon. All rights reserved.
//
//  Redistribution and use in source and binary forms, with or without
//  modification, are permitted provided that the following conditions are
//  met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
//     * Neither the name of Image Engine Design nor the names of any
//       other contributors to this software may be used to endorse or
//       promote products derived from this software without specific prior
//       written permission.
//
//  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
//  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
//  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
//  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
//  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
//  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
//  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
//  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
//  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
//  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
//  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
//////////////////////////////////////////////////////////////////////////

// This must come before the Cortex includes, because on OSX headers included
// by TBB define macros which conflict with the inline functions in ai_types.h.
#include "ai.h"

#include "IECore/CompoundParameter.h"
#include "IECore/SimpleTypedData.h"
#include "IECore/MessageHandler.h"
#include "IECore/DespatchTypedData.h"

#include "IECoreArnold/ToArnoldConverter.h"

using namespace IECore;
using namespace IECoreArnold;

IE_CORE_DEFINERUNTIMETYPED( ToArnoldConverter );

ToArnoldConverter::ToArnoldConverter( const std::string &description, IECore::TypeId supportedType )
	:	IECore::FromCoreConverter( description, supportedType )
{
}

ToArnoldConverter::~ToArnoldConverter()
{
}

AtNode *ToArnoldConverter::convert() const
{
	IECore::ConstCompoundObjectPtr operands = parameters()->getTypedValidatedValue<IECore::CompoundObject>();
	return doConversion( srcParameter()->getValidatedValue(), operands );
}

void ToArnoldConverter::setParameter( AtNode *node, const AtParamEntry *parameter, const IECore::Data *value )
{
	bool isArray = false;
	int type = AiParamGetType( parameter );
	if( type == AI_TYPE_ARRAY )
	{
		type = AiParamGetDefault( parameter )->ARRAY->type;
		isArray = true;
	}

	setParameterInternal( node, AiParamGetName( parameter ), type, isArray, value );
}

void ToArnoldConverter::setParameter( AtNode *node, const char *name, const IECore::Data *value )
{
	const AtParamEntry *parameter = AiNodeEntryLookUpParameter( AiNodeGetNodeEntry( node ), name );
	if( parameter )
	{
		setParameter( node, parameter, value );
	}
	else
	{
		bool array = false;
		int type = parameterType( value->typeId(), array );
		if( type != AI_TYPE_NONE )
		{
			std::string typeString = "constant ";
			if( array )
			{
				typeString += "ARRAY ";
			}
			typeString += AiParamGetTypeName( type );
			AiNodeDeclare( node, name, typeString.c_str() );
			setParameterInternal( node, name, type, array, value );
		}
		else
		{
			msg(
				Msg::Warning,
				"ToArnoldConverter::setParameter",
				boost::format( "Unsupported data type \"%s\" for name \"%s\"" ) % value->typeName() % name
			);
		}
	}
}

void ToArnoldConverter::setParameters( AtNode *node, const IECore::CompoundDataMap &values )
{
	for( CompoundDataMap::const_iterator it=values.begin(); it!=values.end(); it++ )
	{
		setParameter( node, it->first.value().c_str(), it->second );
	}
}

IECore::DataPtr ToArnoldConverter::getParameter( AtNode *node, const AtParamEntry *parameter )
{
	return getParameterInternal( node, AiParamGetName( parameter ), AiParamGetType( parameter ) );
}

IECore::DataPtr ToArnoldConverter::getParameter( AtNode *node, const AtUserParamEntry *parameter )
{
	return getParameterInternal( node, AiUserParamGetName( parameter ), AiUserParamGetType( parameter ) );
}

IECore::DataPtr ToArnoldConverter::getParameter( AtNode *node, const char *name )
{
	const AtParamEntry *parameter = AiNodeEntryLookUpParameter( AiNodeGetNodeEntry( node ), name );
	if( parameter )
	{
		return getParameter( node, parameter );
	}
	else
	{
		const AtUserParamEntry *userParameter = AiNodeLookUpUserParameter( node, name );
		if( userParameter )
		{
			return getParameter( node, userParameter );
		}
	}

	return 0;
}

void ToArnoldConverter::getParameters( AtNode *node, IECore::CompoundDataMap &values )
{
	/// \todo Non-user parameters

	AtUserParamIterator *it = AiNodeGetUserParamIterator( node );
	while( const AtUserParamEntry *param = AiUserParamIteratorGetNext( it ) )
	{
		DataPtr d = getParameter( node, param );
		if( d )
		{
			values[AiUserParamGetName( param )] = d;
		}
		else
		{
			msg(
				Msg::Warning,
				"ToArnoldConverter::getParameters",
				boost::format( "Unable to convert user parameter \"%s\"" ) % AiUserParamGetName( param )
			);
		}
	}
	AiUserParamIteratorDestroy( it );
}

int ToArnoldConverter::parameterType( IECore::TypeId dataType, bool &array )
{
	switch( dataType )
	{
		// non-array types

		case IntDataTypeId :
			array = false;
			return AI_TYPE_INT;
		case FloatDataTypeId :
			array = false;
			return AI_TYPE_FLOAT;
		case StringDataTypeId :
			array = false;
			return AI_TYPE_STRING;
		case Color3fDataTypeId :
			array = false;
			return AI_TYPE_RGB;
		case BoolDataTypeId :
			array = false;
			return AI_TYPE_BOOLEAN;

		// array types

		case IntVectorDataTypeId :
			array = true;
			return AI_TYPE_INT;
		case FloatVectorDataTypeId :
			array = true;
			return AI_TYPE_FLOAT;
		case StringVectorDataTypeId :
			array = true;
			return AI_TYPE_STRING;
		case Color3fVectorDataTypeId :
			array = true;
			return AI_TYPE_RGB;
		case BoolVectorDataTypeId :
			array = true;
			return AI_TYPE_BOOLEAN;
		default :
			return AI_TYPE_NONE;
	}
}

AtArray *ToArnoldConverter::dataToArray( const IECore::Data *data )
{
	bool isArray = false;
	int type = parameterType( data->typeId(), isArray );
	if( type == AI_TYPE_NONE || !isArray )
	{
		return 0;
	}

	const void *dataAddress = despatchTypedData<TypedDataAddress, TypeTraits::IsTypedData, DespatchTypedDataIgnoreError>( DataPtr( const_cast<Data *>( data ) ) );
	size_t dataSize = despatchTypedData<TypedDataSize, TypeTraits::IsTypedData, DespatchTypedDataIgnoreError>( DataPtr( const_cast<Data *>( data ) ) );
	return AiArrayConvert( dataSize, 1, type, dataAddress );
}

template<typename T>
static inline const T *dataCast( const char *name, const IECore::Data *data )
{
	const T *result = runTimeCast<const T>( data );
	if( result )
	{
		return result;
	}
	msg( Msg::Warning, "ToArnoldConverter::setParameter", boost::format( "Unsupported value type \"%s\" for parameter \"%s\" (expected %s)." ) % data->typeName() % name % T::staticTypeName() );
	return 0;
}

void ToArnoldConverter::setParameterInternal( AtNode *node, const char *name, int parameterType, bool array, const IECore::Data *value )
{
	if( array )
	{
		AtArray *a = dataToArray( value );
		if( !a )
		{
			msg( Msg::Warning, "ToArnoldConverter::setParameter", boost::format( "Unable to create array from data of type \"%s\" for parameter \"%s\"" ) % value->typeName() % name );
			return;
		}
		if( a->type != parameterType )
		{
			msg( Msg::Warning, "ToArnoldConverter::setParameter", boost::format( "Unable to create array of type %s from data of type \"%s\" for parameter \"%s\"" ) % AiParamGetTypeName( parameterType ) % value->typeName() % name );
			return;
		}
		AiNodeSetArray( node, name, a );
	}
	else
	{
		switch( parameterType )
		{
			case AI_TYPE_INT :
				if( const IntData *data = dataCast<IntData>( name, value ) )
				{
					AiNodeSetInt( node, name, data->readable() );
				}
				break;
			case AI_TYPE_BYTE :
				if( const IntData *data = dataCast<IntData>( name, value ) )
				{
					AiNodeSetByte( node, name, data->readable() );
				}
				break;
			case AI_TYPE_FLOAT :
				if( const FloatData *data = dataCast<FloatData>( name, value ) )
				{
					AiNodeSetFlt( node, name, data->readable() );
				}
				break;
			case AI_TYPE_STRING :
				if( const StringData *data = dataCast<StringData>( name, value ) )
				{
					AiNodeSetStr( node, name, data->readable().c_str() );
				}
				break;
			case AI_TYPE_RGB :
				if( const Color3fData *data = dataCast<Color3fData>( name, value ) )
				{
					const Imath::Color3f &c = data->readable();
					AiNodeSetRGB( node, name, c[0], c[1], c[2] );
				}
				break;
			case AI_TYPE_ENUM :
				if( const StringData *data = dataCast<StringData>( name, value ) )
				{
					AiNodeSetStr( node, name, data->readable().c_str() );
				}
				break;
			case AI_TYPE_BOOLEAN :
			{
				if( const BoolData *data = dataCast<BoolData>( name, value ) )
				{
					AiNodeSetBool( node, name, data->readable() );
				}
				break;
			}
			default :
			{
				msg( Msg::Warning, "ToArnoldConverter::setParameter", boost::format( "Arnold parameter \"%s\" has unsupported type \"%s\"." ) % name % AiParamGetTypeName( parameterType ) );
			}
		}
	}
}

IECore::DataPtr ToArnoldConverter::getParameterInternal( AtNode *node, const char *name, int parameterType )
{
	switch( parameterType )
	{
		case AI_TYPE_BOOLEAN :
			return new BoolData( AiNodeGetBool( node, name ) );
		case AI_TYPE_INT :
			return new IntData( AiNodeGetInt( node, name ) );
		case AI_TYPE_FLOAT :
			return new FloatData( AiNodeGetFlt( node, name ) );
		case AI_TYPE_STRING :
			return new StringData( AiNodeGetStr( node, name ) );
	}
	return 0;
}

ToArnoldConverterPtr ToArnoldConverter::create( IECore::ObjectPtr object )
{
	const CreatorMap &m = creators();
	CreatorMap::const_iterator it = m.find( object->typeId() );
	if( it != m.end() )
	{
		return it->second( object );
	}
	return 0;
}

ToArnoldConverter::CreatorMap &ToArnoldConverter::creators()
{
	static CreatorMap m;
	return m;
}

<|endoftext|>


<commit_before>#include <iostream>

#include <CL/cl.h>

#include "cl_util.h"

const char *source =
    "__kernel void saxpy(float a, __global float *x, __global float "
    "*y)"
    "{"
    "   uint id = get_global_id(0);"
    "   y[id] = a * x[id] + y[id];"
    "}";

int main()
{
    // input data
    const size_t inputSize = 100000000;
    float a = 7;
    std::vector<float> x(inputSize, 4);
    std::vector<float> y(inputSize, 5);

    const auto device =
        cl::chooseDevice(cl::allDevices(), CL_DEVICE_MAX_COMPUTE_UNITS);
    cl::printDeviceInfo(device);
    cl_int error = 0;

    cl_context context =
        clCreateContext(nullptr, 1, &device, nullptr, nullptr, &error);
    cl_command_queue queue = clCreateCommandQueue(context, device, 0, &error);
    cl_program program =
        clCreateProgramWithSource(context, 1, &source, nullptr, &error);
    clBuildProgram(program, 1, &device, nullptr, nullptr, nullptr);

    size_t logSize;
    clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr,
                          &logSize);
    std::vector<char> infoLog(logSize);
    clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, logSize,
                          infoLog.data(), nullptr);
    std::cout << std::string{infoLog.begin(), infoLog.end()} << std::endl;

    cl_kernel kernel = clCreateKernel(program, "saxpy", nullptr);
    cl_mem bufferX =
        clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                       sizeof(float) * x.size(), x.data(), &error);
    cl_mem bufferY =
        clCreateBuffer(context, CL_MEM_READ_WRITE | CL_MEM_COPY_HOST_PTR,
                       sizeof(float) * y.size(), y.data(), &error);

    clSetKernelArg(kernel, 0, sizeof(a), static_cast<void *>(&a));
    clSetKernelArg(kernel, 1, sizeof(bufferX), static_cast<void *>(&bufferX));
    clSetKernelArg(kernel, 2, sizeof(bufferY), static_cast<void *>(&bufferY));

    clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, &inputSize, nullptr, 0,
                           nullptr, nullptr);
    clFinish(queue);

    auto data = static_cast<cl_float *>(clEnqueueMapBuffer(
        queue, bufferY, CL_TRUE, CL_MAP_READ, 0, sizeof(float) * inputSize, 0,
        nullptr, nullptr, &error));

    for (int i = 0; i < 20; ++i)
    {
        std::cout << data[i] << std::endl;
    }
}
<commit_msg>Add context error callback<commit_after>#include <iostream>

#include <CL/cl.h>

#include "cl_util.h"

const char *source =
    "__kernel void saxpy(float a, __global float *x, __global float *y)"
    "{"
    "   uint id = get_global_id(0);"
    "   y[id] = a * x[id] + y[id];"
    "}";

void errorCallback(const char *errinfo, const void *private_info, size_t cb,
                   void *user_data)
{
    std::cout << "ContextError: " << errinfo << std::endl;
}

int main()
{
    // input data
    const size_t inputSize = 100000000;
    float a = 7;
    std::vector<float> x(inputSize, 4);
    std::vector<float> y(inputSize, 5);

    const auto device =
        cl::chooseDevice(cl::allDevices(), CL_DEVICE_MAX_COMPUTE_UNITS);
    cl::printDeviceInfo(device);
    cl_int error = 0;

    cl_context context =
        clCreateContext(nullptr, 1, &device, errorCallback, nullptr, &error);
    cl_command_queue queue = clCreateCommandQueue(context, device, 0, &error);
    cl_program program =
        clCreateProgramWithSource(context, 1, &source, nullptr, &error);
    clBuildProgram(program, 1, &device, nullptr, nullptr, nullptr);

    size_t logSize;
    clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr,
                          &logSize);
    std::vector<char> infoLog(logSize);
    clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, logSize,
                          infoLog.data(), nullptr);
    std::cout << std::string{infoLog.begin(), infoLog.end()} << std::endl;

    cl_kernel kernel = clCreateKernel(program, "saxpy", nullptr);
    cl_mem bufferX =
        clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                       sizeof(float) * x.size(), x.data(), &error);
    cl_mem bufferY =
        clCreateBuffer(context, CL_MEM_READ_WRITE | CL_MEM_COPY_HOST_PTR,
                       sizeof(float) * y.size(), y.data(), &error);

    clSetKernelArg(kernel, 0, sizeof(a), static_cast<void *>(&a));
    clSetKernelArg(kernel, 1, sizeof(bufferX), static_cast<void *>(&bufferX));
    clSetKernelArg(kernel, 2, sizeof(bufferY), static_cast<void *>(&bufferY));

    clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, &inputSize, nullptr, 0,
                           nullptr, nullptr);
    clFinish(queue);

    auto data = static_cast<cl_float *>(clEnqueueMapBuffer(
        queue, bufferY, CL_TRUE, CL_MAP_READ, 0, sizeof(float) * inputSize, 0,
        nullptr, nullptr, &error));

    for (int i = 0; i < 20; ++i)
    {
        std::cout << data[i] << std::endl;
    }
}
<|endoftext|>


<commit_before>#include <vector>
#include <iostream>
#include <stdlib.h>
#include <string.h>
#include <fstream>
#include "overlapgraph.h"

#define VERTEX_INACTIVE 0
#define VERTEX_ACTIVE 1
#define VERTEX_ELIMINATED -1
#define FUZZ 10

using namespace std;

OverlapGraph::OverlapGraph(std::vector<std::string> readvector, std::vector<overlap> overlapvector)
{
	reads = readvector;
	overlaps = overlapvector;
}

void OverlapGraph::initialize()
{
	graph.resize(reads.size());
	reducedGraph.resize(reads.size());
	nonContainedReads.resize(reads.size());
	nonContainedReads.assign(reads.size(), 1);

	for(int i=0;i<overlaps.size();i++)
	{
		overlap o = overlaps[i];

		//removing contained reads
		if(o.ahg > 0 & o.bhg > 0)
			graph[o.read1].push_back(o);
		else if (o.ahg < 0 & o.bhg < 0)
			graph[o.read2].push_back(o);
		else if (o.ahg >= 0 & o.bhg <= 0)
			nonContainedReads[o.read2] = 0;
		else if (o.ahg <= 0 & o.bhg>=0)
			nonContainedReads[o.read1] = 0;
	}
}

void OverlapGraph::runUnitigging()
{
	int N = reads.size();

	initialize();

	int *vertexstatus = (int *)malloc(N*sizeof(int));
	memset(vertexstatus, VERTEX_INACTIVE, N*sizeof(int));
	vector< vector<bool> > reduceflags(N);

	for (int i = 0; i < N; i++)
		for (int j = 0; j < graph[i].size(); j++)
			reduceflags[i].push_back(false);

	for (int I = 0; I < N; I++)
	{
		if (graph[I].size() == 0)
			continue;

		for (int j = 0; j < graph[I].size(); j++)
			vertexstatus[graph[I][j].read2] = VERTEX_ACTIVE;

		int longest = 0;
		for (int j = 0; j < graph[I].size(); j++)
		{
			if (graph[I][j].bhg > longest)
				longest = graph[I][j].bhg;
		}

		longest += FUZZ;
		cout << longest << endl;

		for (int j = 0; j < graph[I].size(); j++)
		{
			if (vertexstatus[graph[I][j].read2] == VERTEX_ACTIVE)
			{
				int x = graph[I][j].read2;
				for (int k = 0; k < graph[x].size(); k++)
				{
					if (graph[I][j].bhg + graph[x][k].bhg <= longest)
					{
						if (vertexstatus[graph[x][k].read2] == VERTEX_ACTIVE)
							vertexstatus[graph[x][k].read2] = VERTEX_ELIMINATED;
					}
				}
			}
		}

		for (int j = 0; j < graph[I].size(); j++)
		{
			int x = graph[I][j].read2;
			for (int k = 0; k < graph[x].size(); k++)
			{
				if (graph[x][k].bhg < FUZZ)
				{
					if (vertexstatus[graph[x][k].read2] == VERTEX_ACTIVE)
						vertexstatus[graph[x][k].read2] = VERTEX_ELIMINATED;
				}
			}
		}

		for (int j = 0; j < graph[I].size(); j++)
		{
			if (vertexstatus[graph[I][j].read2] == VERTEX_ELIMINATED)
			{
				reduceflags[I][j] = true;
			}
			vertexstatus[graph[I][j].read2] == VERTEX_INACTIVE;
		}
	}

	for(int i=0;i<graph.size();i++)
    {
        for(int j=0;j<graph[i].size();j++)
        {
			if (!reduceflags[i][j] & nonContainedReads[graph[i][j].read1] & nonContainedReads[graph[i][j].read2]){
				// Add edge to the reduced graph
				reducedGraph[i].push_back(graph[i][j]);
				cout << graph[i][j].read1 << " -> " << graph[i][j].read2 << endl;
			}
        }
    }
	uniqueJoinCollapsing();

	free(vertexstatus);
	return;
}

void OverlapGraph::unitigsPrinting()
{
	ofstream output("unitigsPrinting.afg");

	int counter = 1;

	output << "Unitigs{" << endl;

	for (Chunk c : collapsedReducedGraph)
	{
		output << "\t" << counter << "{" << endl;

		for (int a : c.members)
		{
		}

		output << c.members[1] << endl;
		output << "\t" << "}" << endl;
	}

	output << "}" << endl;
}

void OverlapGraph::printLayouts()
{
	ofstream output("layouts.afg");
	ofstream pretty("pretty.afg");
	int *offsets = (int *)malloc(reads.size()*sizeof(int));
	memset(offsets, 0, reads.size()*sizeof(int));

	for (int i = 0; i < reads.size(); i++)
	{
		for (int j = 0; j < graph[i].size(); j++)
		{
			overlap o = graph[i][j];
			if (nonContainedReads[o.read1] & nonContainedReads[o.read2]) {
				offsets[o.read2] = offsets[o.read1] + o.ahg;
			}
		}
	}

	output << "{LAY" << endl;
	for (int i = 0; i < reads.size(); i++)
	{
		if (nonContainedReads[i])
		{
			output << "{TLE" << endl;
			output << "clr:0," << reads[i].length() << endl;
			output << "off:" << offsets[i] << endl;
			output << "src:" << i + 1 << endl;
			output << "}" << endl;

			pretty << string(offsets[i], ' ') + reads[i] << endl;
			//pretty <<reads[i] << endl;

		}
	}

	output << "}" << endl;
	output.close();

	free(offsets);
}

void OverlapGraph::uniqueJoinCollapsing(){

	reducedGraphVertexReferenceCounter.assign(reducedGraph.size(), 0);

	// Calculate references to vertexes
	for (int i = 0; i < reducedGraph.size(); i++){
		for (int j = 0; j < reducedGraph[i].size(); j++){
			reducedGraphVertexReferenceCounter[reducedGraph[i][j].read2]++;
		}
	}

	// Find first "juncture" vertex
	int startingVertex = 0;
	bool startingVertexFound = false;
	while (startingVertex < reducedGraph.size()){
		if (reducedGraph[startingVertex].size() > 0){
			startingVertexFound = true;
			break;
		}
		startingVertex++;
	}
	if (startingVertexFound){
		calculateChunks(startingVertex);
	}
}

void OverlapGraph::calculateChunks(int chunkStartingVertexId){

	int currentVertexId = chunkStartingVertexId;

	if (vertexAlreadyInChunk(currentVertexId)){
		return;
	}

	// create new chunk
	Chunk chunk;

	while (true){
		if (reducedGraphVertexReferenceCounter[currentVertexId] > 1 && currentVertexId != chunkStartingVertexId){
			// special case: juncture vertex (not part of our chunk)
			calculateChunks(currentVertexId);
			break;
		}

		chunk.members.push_back(currentVertexId);

		if (reducedGraph[currentVertexId].size() < 1){
			// dead end
			break;
		}

		if (reducedGraph[currentVertexId].size() > 1){
			// juncture point - each path vertex represents a new chunk start
			for (int i = 0; i < reducedGraph[currentVertexId].size(); i++){
				calculateChunks(reducedGraph[currentVertexId][i].read2);
			}
			break;
		}

		currentVertexId = reducedGraph[currentVertexId][0].read2;
	}
	collapsedReducedGraph.push_back(chunk);

	return;
}

bool OverlapGraph::vertexAlreadyInChunk(int vertexId){
	for (int i = 0; i < collapsedReducedGraph.size(); ++i){
		for (int j = 0; j < collapsedReducedGraph[i].members.size(); ++j){
			if (vertexId == collapsedReducedGraph[i].members[j]){
				return true;
			}
		}
	}
	return false;
}<commit_msg>Unitigs printing function - implementing<commit_after>#include <vector>
#include <iostream>
#include <stdlib.h>
#include <string.h>
#include <fstream>
#include "overlapgraph.h"

#define VERTEX_INACTIVE 0
#define VERTEX_ACTIVE 1
#define VERTEX_ELIMINATED -1
#define FUZZ 10

using namespace std;

OverlapGraph::OverlapGraph(std::vector<std::string> readvector, std::vector<overlap> overlapvector)
{
	reads = readvector;
	overlaps = overlapvector;
}

void OverlapGraph::initialize()
{
	graph.resize(reads.size());
	reducedGraph.resize(reads.size());
	nonContainedReads.resize(reads.size());
	nonContainedReads.assign(reads.size(), 1);

	for(int i=0;i<overlaps.size();i++)
	{
		overlap o = overlaps[i];

		//removing contained reads
		if(o.ahg > 0 & o.bhg > 0)
			graph[o.read1].push_back(o);
		else if (o.ahg < 0 & o.bhg < 0)
			graph[o.read2].push_back(o);
		else if (o.ahg >= 0 & o.bhg <= 0)
			nonContainedReads[o.read2] = 0;
		else if (o.ahg <= 0 & o.bhg>=0)
			nonContainedReads[o.read1] = 0;
	}
}

void OverlapGraph::runUnitigging()
{
	int N = reads.size();

	initialize();

	int *vertexstatus = (int *)malloc(N*sizeof(int));
	memset(vertexstatus, VERTEX_INACTIVE, N*sizeof(int));
	vector< vector<bool> > reduceflags(N);

	for (int i = 0; i < N; i++)
		for (int j = 0; j < graph[i].size(); j++)
			reduceflags[i].push_back(false);

	for (int I = 0; I < N; I++)
	{
		if (graph[I].size() == 0)
			continue;

		for (int j = 0; j < graph[I].size(); j++)
			vertexstatus[graph[I][j].read2] = VERTEX_ACTIVE;

		int longest = 0;
		for (int j = 0; j < graph[I].size(); j++)
		{
			if (graph[I][j].bhg > longest)
				longest = graph[I][j].bhg;
		}

		longest += FUZZ;
		cout << longest << endl;

		for (int j = 0; j < graph[I].size(); j++)
		{
			if (vertexstatus[graph[I][j].read2] == VERTEX_ACTIVE)
			{
				int x = graph[I][j].read2;
				for (int k = 0; k < graph[x].size(); k++)
				{
					if (graph[I][j].bhg + graph[x][k].bhg <= longest)
					{
						if (vertexstatus[graph[x][k].read2] == VERTEX_ACTIVE)
							vertexstatus[graph[x][k].read2] = VERTEX_ELIMINATED;
					}
				}
			}
		}

		for (int j = 0; j < graph[I].size(); j++)
		{
			int x = graph[I][j].read2;
			for (int k = 0; k < graph[x].size(); k++)
			{
				if (graph[x][k].bhg < FUZZ)
				{
					if (vertexstatus[graph[x][k].read2] == VERTEX_ACTIVE)
						vertexstatus[graph[x][k].read2] = VERTEX_ELIMINATED;
				}
			}
		}

		for (int j = 0; j < graph[I].size(); j++)
		{
			if (vertexstatus[graph[I][j].read2] == VERTEX_ELIMINATED)
			{
				reduceflags[I][j] = true;
			}
			vertexstatus[graph[I][j].read2] == VERTEX_INACTIVE;
		}
	}

	for(int i=0;i<graph.size();i++)
    {
        for(int j=0;j<graph[i].size();j++)
        {
			if (!reduceflags[i][j] & nonContainedReads[graph[i][j].read1] & nonContainedReads[graph[i][j].read2]){
				// Add edge to the reduced graph
				reducedGraph[i].push_back(graph[i][j]);
				cout << graph[i][j].read1 << " -> " << graph[i][j].read2 << endl;
			}
        }
    }
	uniqueJoinCollapsing();

	free(vertexstatus);
	return;
}

void OverlapGraph::unitigsPrinting()
{
	ofstream output("unitigsPrinting.afg");

	int counter = 1;

	output << "Unitigs{" << endl;

	for (Chunk c : collapsedReducedGraph)
	{
		output << "\t" << counter << "{" << endl;
		output << "\t\t members: ";
		for (int a : c.members)
		{
			output << a << " ";
		}

		output << endl;
		output << "\t" << "}" << endl;
	}

	output << "}" << endl;
}

void OverlapGraph::printLayouts()
{
	ofstream output("layouts.afg");
	ofstream pretty("pretty.afg");
	int *offsets = (int *)malloc(reads.size()*sizeof(int));
	memset(offsets, 0, reads.size()*sizeof(int));

	for (int i = 0; i < reads.size(); i++)
	{
		for (int j = 0; j < graph[i].size(); j++)
		{
			overlap o = graph[i][j];
			if (nonContainedReads[o.read1] & nonContainedReads[o.read2]) {
				offsets[o.read2] = offsets[o.read1] + o.ahg;
			}
		}
	}

	output << "{LAY" << endl;
	for (int i = 0; i < reads.size(); i++)
	{
		if (nonContainedReads[i])
		{
			output << "{TLE" << endl;
			output << "clr:0," << reads[i].length() << endl;
			output << "off:" << offsets[i] << endl;
			output << "src:" << i + 1 << endl;
			output << "}" << endl;

			pretty << string(offsets[i], ' ') + reads[i] << endl;
			//pretty <<reads[i] << endl;

		}
	}

	output << "}" << endl;
	output.close();

	free(offsets);
}

void OverlapGraph::uniqueJoinCollapsing(){

	reducedGraphVertexReferenceCounter.assign(reducedGraph.size(), 0);

	// Calculate references to vertexes
	for (int i = 0; i < reducedGraph.size(); i++){
		for (int j = 0; j < reducedGraph[i].size(); j++){
			reducedGraphVertexReferenceCounter[reducedGraph[i][j].read2]++;
		}
	}

	// Find first "juncture" vertex
	int startingVertex = 0;
	bool startingVertexFound = false;
	while (startingVertex < reducedGraph.size()){
		if (reducedGraph[startingVertex].size() > 0){
			startingVertexFound = true;
			break;
		}
		startingVertex++;
	}
	if (startingVertexFound){
		calculateChunks(startingVertex);
	}
}

void OverlapGraph::calculateChunks(int chunkStartingVertexId){

	int currentVertexId = chunkStartingVertexId;

	if (vertexAlreadyInChunk(currentVertexId)){
		return;
	}

	// create new chunk
	Chunk chunk;

	while (true){
		if (reducedGraphVertexReferenceCounter[currentVertexId] > 1 && currentVertexId != chunkStartingVertexId){
			// special case: juncture vertex (not part of our chunk)
			calculateChunks(currentVertexId);
			break;
		}

		chunk.members.push_back(currentVertexId);

		if (reducedGraph[currentVertexId].size() < 1){
			// dead end
			break;
		}

		if (reducedGraph[currentVertexId].size() > 1){
			// juncture point - each path vertex represents a new chunk start
			for (int i = 0; i < reducedGraph[currentVertexId].size(); i++){
				calculateChunks(reducedGraph[currentVertexId][i].read2);
			}
			break;
		}

		currentVertexId = reducedGraph[currentVertexId][0].read2;
	}
	collapsedReducedGraph.push_back(chunk);

	return;
}

bool OverlapGraph::vertexAlreadyInChunk(int vertexId){
	for (int i = 0; i < collapsedReducedGraph.size(); ++i){
		for (int j = 0; j < collapsedReducedGraph[i].members.size(); ++j){
			if (vertexId == collapsedReducedGraph[i].members[j]){
				return true;
			}
		}
	}
	return false;
}<|endoftext|>


<commit_before>#include "ResourcePool.h"

#include <osgDB/ReadFile>

using namespace troen;

ResourcePool::ResourcePool() {}

ResourcePool::~ResourcePool() {}

void ResourcePool::readData()
{
	m_dataAlreadyRead = true;

	std::string folder = "data/models/cycle/";

	std::string textureFileNames[m_textureCount] = {
		"MG_MovieCycle_Body_SPEC.tga",
		"MG_MovieCycle_BodyHeadLight_EMSS.tga",
		"MG_MovieCycle_Body_NORM.tga",
		"MG_Player_Body_SPEC.tga",
		"MG_Player_Body_EMSS.tga",
		"MG_Player_Body_NORM.tga",
		"MG_MovieCycle_Tire_DIFF.tga",
		"MG_MovieCycle_Tire_EMSS.tga",
		"MG_MovieCycle_Tire_NORM.tga",
		"MG_Player_Helmet_SPEC.tga",
		"MG_Player_Helmet_EMSS.tga",
		"MG_Player_Helmet_NORM.tga",
		"MG_Player_Disc_SPEC.tga",
		"MG_Player_Disc_EMSS.tga",
		"MG_Player_Disc_NORM.tga",
		"Glass.tga",
		"MG_MovieCycle_Engine_EMSS.tga",
		"MG_MovieCycle_Engine_NORM.tga",
		"MG_Player_Baton_SPEC.tga",
		"MG_Player_Baton_EMSS.tga",
		"MG_Player_Baton_NORM.tga"
	};

	for (int i = 0; i < m_textureCount; ++i)
	{
		const std::string filePath = folder + textureFileNames[i];
		//m_images[i] = osgDB::readImageFile(filePath);

		osg::ref_ptr<osg::Image> img = osgDB::readImageFile(filePath);
		m_images[i] = img.release();

	}

	std::string objFileNames[m_objectCount] = {
		"MG_MovieCycle_Body_MI.ive",
		"MG_MovieCycle_PlayerBody_MI.ive",
		"MG_MovieCycle_Tire_MI.ive",
		"MG_MovieCycle_PlayerHelmet_MI.ive",
		"MG_MovieCycle_PlayerDisc_MI.ive",
		"MG_MovieCycle_Glass_MI.ive",
		"MG_MovieCycle_Engine_MI.ive",
		"MG_MovieCycle_PlayerBaton_MI.ive"
	};

	for (int i = 0; i < m_objectCount; ++i)
	{
		const std::string filePath = folder + objFileNames[i];
		std::cout << "[TroenGame::bikeView] Loading Model \"" << filePath << "\"" << std::endl;
		m_objects[i] = osgDB::readNodeFile(filePath);
	}

}

osg::Image* ResourcePool::getImage(TextureResource texture)
{
	if (!m_dataAlreadyRead)
		readData();

	if (!m_images[texture]->valid() || !m_images[texture]) {
		std::cout << "rereading data for texture: " << texture << std::endl;
		readData();
	}

	return m_images[texture];
}

osg::Node* ResourcePool::getNode(ModelResource model)
{
	if (!m_dataAlreadyRead)
		readData();

	const osg::CopyOp copyOp(osg::CopyOp::SHALLOW_COPY);
	return dynamic_cast<osg::Node*>(m_objects[model]->clone(copyOp));
}<commit_msg>yet another check for valid resources; hopefully fixes #73<commit_after>#include "ResourcePool.h"

#include <osgDB/ReadFile>

using namespace troen;

ResourcePool::ResourcePool() {}

ResourcePool::~ResourcePool() {}

void ResourcePool::readData()
{
	m_dataAlreadyRead = true;

	std::string folder = "data/models/cycle/";

	std::string textureFileNames[m_textureCount] = {
		"MG_MovieCycle_Body_SPEC.tga",
		"MG_MovieCycle_BodyHeadLight_EMSS.tga",
		"MG_MovieCycle_Body_NORM.tga",
		"MG_Player_Body_SPEC.tga",
		"MG_Player_Body_EMSS.tga",
		"MG_Player_Body_NORM.tga",
		"MG_MovieCycle_Tire_DIFF.tga",
		"MG_MovieCycle_Tire_EMSS.tga",
		"MG_MovieCycle_Tire_NORM.tga",
		"MG_Player_Helmet_SPEC.tga",
		"MG_Player_Helmet_EMSS.tga",
		"MG_Player_Helmet_NORM.tga",
		"MG_Player_Disc_SPEC.tga",
		"MG_Player_Disc_EMSS.tga",
		"MG_Player_Disc_NORM.tga",
		"Glass.tga",
		"MG_MovieCycle_Engine_EMSS.tga",
		"MG_MovieCycle_Engine_NORM.tga",
		"MG_Player_Baton_SPEC.tga",
		"MG_Player_Baton_EMSS.tga",
		"MG_Player_Baton_NORM.tga"
	};

	for (int i = 0; i < m_textureCount; ++i)
	{
		const std::string filePath = folder + textureFileNames[i];
		//m_images[i] = osgDB::readImageFile(filePath);

		osg::ref_ptr<osg::Image> img = osgDB::readImageFile(filePath);
		m_images[i] = img.release();

	}

	std::string objFileNames[m_objectCount] = {
		"MG_MovieCycle_Body_MI.ive",
		"MG_MovieCycle_PlayerBody_MI.ive",
		"MG_MovieCycle_Tire_MI.ive",
		"MG_MovieCycle_PlayerHelmet_MI.ive",
		"MG_MovieCycle_PlayerDisc_MI.ive",
		"MG_MovieCycle_Glass_MI.ive",
		"MG_MovieCycle_Engine_MI.ive",
		"MG_MovieCycle_PlayerBaton_MI.ive"
	};

	for (int i = 0; i < m_objectCount; ++i)
	{
		const std::string filePath = folder + objFileNames[i];
		std::cout << "[TroenGame::bikeView] Loading Model \"" << filePath << "\"" << std::endl;
		m_objects[i] = osgDB::readNodeFile(filePath);
	}

}

osg::Image* ResourcePool::getImage(TextureResource texture)
{
	if (!m_dataAlreadyRead)
		readData();

	if (!m_images[texture]->valid() || !m_images[texture] || m_images[texture]->getImageSizeInBytes() <= 0) {
		std::cout << "rereading data for texture: " << texture << std::endl;
		readData();
	}

	return m_images[texture];
}

osg::Node* ResourcePool::getNode(ModelResource model)
{
	if (!m_dataAlreadyRead)
		readData();

	const osg::CopyOp copyOp(osg::CopyOp::SHALLOW_COPY);
	return dynamic_cast<osg::Node*>(m_objects[model]->clone(copyOp));
}<|endoftext|>


<commit_before>/*!
    \file ws.cpp
    \brief WebSocket C++ Library implementation
    \author Ivan Shynkarenka
    \date 22.05.2019
    \copyright MIT License
*/

#include "server/ws/ws.h"

#include "string/encoding.h"
#include "string/format.h"

#include <algorithm>
#include <openssl/sha.h>

namespace CppServer {
namespace WS {

bool WebSocket::PerformClientUpgrade(const HTTP::HTTPResponse& response, const CppCommon::UUID& id)
{
    if (response.status() != 101)
        return false;

    bool error = false;
    bool accept = false;
    bool connection = false;
    bool upgrade = false;

    // Validate WebSocket handshake headers
    for (size_t i = 0; i < response.headers(); ++i)
    {
        auto header = response.header(i);
        auto key = std::get<0>(header);
        auto value = std::get<1>(header);

        if (key == "Connection")
        {
            if (value != "Upgrade")
            {
                error = true;
                onWSError("Invalid WebSocket handshaked response: 'Connection' header value must be 'Upgrade'");
                break;
            }

            connection = true;
        }
        else if (key == "Upgrade")
        {
            if (value != "websocket")
            {
                error = true;
                onWSError("Invalid WebSocket handshaked response: 'Upgrade' header value must be 'websocket'");
                break;
            }

            upgrade = true;
        }
        else if (key == "Sec-WebSocket-Accept")
        {
            // Calculate the original WebSocket hash
            std::string wskey = CppCommon::Encoding::Base64Encode(id.string()) + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
            char wshash[SHA_DIGEST_LENGTH];
            SHA1((const unsigned char*)wskey.data(), wskey.size(), (unsigned char*)wshash);

            // Get the received WebSocket hash
            wskey = CppCommon::Encoding::Base64Decode(value);

            // Compare original and received hashes
            if (std::strncmp(wskey.data(), wshash, std::min(wskey.size(), sizeof(wshash))) != 0)
            {
                error = true;
                onWSError("Invalid WebSocket handshaked response: 'Sec-WebSocket-Accept' value validation failed");
                break;
            }

            accept = true;
        }
    }

    // Failed to perfrom WebSocket handshake
    if (!accept || !connection || !upgrade)
    {
        if (!error)
            onWSError("Invalid WebSocket response");
        return false;
    }

    // WebSocket successfully handshaked!
    _ws_handshaked = true;
    *((uint32_t*)_ws_send_mask) = rand();
    onWSConnected(response);

    return true;
}

bool WebSocket::PerformServerUpgrade(const HTTP::HTTPRequest& request, HTTP::HTTPResponse& response)
{
    if (request.method() != "GET")
        return false;

    bool error = false;
    bool connection = false;
    bool upgrade = false;
    bool ws_key = false;
    bool ws_version = false;

    std::string accept;

    // Validate WebSocket handshake headers
    for (size_t i = 0; i < request.headers(); ++i)
    {
        auto header = request.header(i);
        auto key = std::get<0>(header);
        auto value = std::get<1>(header);

        if (key == "Connection")
        {
            if ((value != "Upgrade") && (value != "keep-alive, Upgrade"))
            {
                error = true;
                response.MakeErrorResponse("Invalid WebSocket handshaked request: 'Connection' header value must be 'Upgrade' or 'keep-alive, Upgrade'", 400);
                break;
            }

            connection = true;
        }
        else if (key == "Upgrade")
        {
            if (value != "websocket")
            {
                error = true;
                response.MakeErrorResponse("Invalid WebSocket handshaked request: 'Upgrade' header value must be 'websocket'", 400);
                break;
            }

            upgrade = true;
        }
        else if (key == "Sec-WebSocket-Key")
        {
            if (value.empty())
            {
                error = true;
                response.MakeErrorResponse("Invalid WebSocket handshaked request: 'Sec-WebSocket-Key' header value must be non empty", 400);
                break;
            }

            // Calculate WebSocket accept value
            std::string wskey = std::string(value) + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
            char wshash[SHA_DIGEST_LENGTH];
            SHA1((const unsigned char*)wskey.data(), wskey.size(), (unsigned char*)wshash);

            accept = CppCommon::Encoding::Base64Encode(std::string(wshash, sizeof(wshash)));

            ws_key = true;
        }
        else if (key == "Sec-WebSocket-Version")
        {
            if (value != "13")
            {
                error = true;
                response.MakeErrorResponse("Invalid WebSocket handshaked request: 'Sec-WebSocket-Version' header value must be '13'", 400);
                break;
            }

            ws_version = true;
        }
    }

    // Filter out non WebSocket handshake requests
    if (!connection && !upgrade && !ws_key && !ws_version)
        return false;

    // Failed to perfrom WebSocket handshake
    if (!connection || !upgrade || !ws_key || !ws_version)
    {
        if (!error)
            response.MakeErrorResponse("Invalid WebSocket response", 400);
        SendResponse(response);
        return false;
    }

    // Prepare WebSocket upgrade success response
    response.Clear();
    response.SetBegin(101);
    response.SetHeader("Connection", "Upgrade");
    response.SetHeader("Upgrade", "websocket");
    response.SetHeader("Sec-WebSocket-Accept", accept);
    response.SetBody();

    // Validate WebSocket upgrade request and response
    if (!onWSConnecting(request, response))
        return false;

    // Send WebSocket upgrade response
    SendResponse(response);

    // WebSocket successfully handshaked!
    _ws_handshaked = true;
    *((uint32_t*)_ws_send_mask) = 0;
    onWSConnected(request);

    return true;
}

void WebSocket::PrepareSendFrame(uint8_t opcode, bool mask, const void* buffer, size_t size, int status)
{
    // Clear the previous WebSocket send buffer
    _ws_send_buffer.clear();

    // Append WebSocket frame opcode
    _ws_send_buffer.push_back(opcode);

    // Append WebSocket frame size
    if (size <= 125)
        _ws_send_buffer.push_back((size & 0xFF) | (mask ? 0x80 : 0));
    else if (size <= 65535)
    {
        _ws_send_buffer.push_back(126 | (mask ? 0x80 : 0));
        _ws_send_buffer.push_back((size >> 8) & 0xFF);
        _ws_send_buffer.push_back(size & 0xFF);
    }
    else
    {
        _ws_send_buffer.push_back(127 | (mask ? 0x80 : 0));
        for (int i = 7; i >= 0; --i)
            _ws_send_buffer.push_back((size >> (8 * i)) & 0xFF);
    }

    if (mask)
    {
        // Append WebSocket frame mask
        _ws_send_buffer.push_back(_ws_send_mask[0]);
        _ws_send_buffer.push_back(_ws_send_mask[1]);
        _ws_send_buffer.push_back(_ws_send_mask[2]);
        _ws_send_buffer.push_back(_ws_send_mask[3]);
    }

    // Resize WebSocket frame buffer
    size_t offset = _ws_send_buffer.size();
    _ws_send_buffer.resize(offset + size);

    // Mask WebSocket frame content
    const uint8_t* data = (const uint8_t*)buffer;
    for (size_t i = 0; i < size; ++i)
        _ws_send_buffer[offset + i] = data[i] ^ _ws_send_mask[i % 4];
}

void WebSocket::PrepareReceiveFrame(const void* buffer, size_t size)
{
    const uint8_t* data = (const uint8_t*)buffer;

    // Clear received data after WebSocket frame was processed
    if (_ws_received)
    {
        _ws_received = false;
        _ws_header_size = 0;
        _ws_payload_size = 0;
        _ws_receive_buffer.clear();
        *((uint32_t*)_ws_receive_mask) = 0;
    }

    while (size > 0)
    {
        // Clear received data after WebSocket frame was processed
        if (_ws_received)
        {
            _ws_received = false;
            _ws_header_size = 0;
            _ws_payload_size = 0;
            _ws_receive_buffer.clear();
            *((uint32_t*)_ws_receive_mask) = 0;
        }

        // Prepare WebSocket frame opcode and mask flag
        if (_ws_receive_buffer.size() < 2)
        {
            for (size_t i = 0; i < 2; ++i, ++data, --size)
            {
                if (size == 0)
                    return;
                _ws_receive_buffer.push_back(*data);
            }
        }

        uint8_t opcode = _ws_receive_buffer[0] & 0x0F;
        [[maybe_unused]] bool fin = ((_ws_receive_buffer[0] >> 7) & 0x01) != 0;
        bool mask = ((_ws_receive_buffer[1] >> 7) & 0x01) != 0;
        size_t payload = _ws_receive_buffer[1] & (~0x80);

        // Prepare WebSocket frame size
        if (payload <= 125)
        {
            _ws_header_size = 2 + (mask ? 4 : 0);
            _ws_payload_size = payload;
            _ws_receive_buffer.reserve(_ws_header_size + _ws_payload_size);
        }
        else if (payload == 126)
        {
            if (_ws_receive_buffer.size() < 4)
            {
                for (size_t i = 0; i < 2; ++i, ++data, --size)
                {
                    if (size == 0)
                        return;
                    _ws_receive_buffer.push_back(*data);
                }
            }

            payload = (((size_t)_ws_receive_buffer[2] << 8) | ((size_t)_ws_receive_buffer[3] << 0));
            _ws_header_size = 4 + (mask ? 4 : 0);
            _ws_payload_size = payload;
            _ws_receive_buffer.reserve(_ws_header_size + _ws_payload_size);
        }
        else if (payload == 127)
        {
            if (_ws_receive_buffer.size() < 10)
            {
                for (size_t i = 0; i < 8; ++i, ++data, --size)
                {
                    if (size == 0)
                        return;
                    _ws_receive_buffer.push_back(*data);
                }
            }

            payload = (((size_t)_ws_receive_buffer[2] << 56) | ((size_t)_ws_receive_buffer[3] << 48) | ((size_t)_ws_receive_buffer[4] << 40) | ((size_t)_ws_receive_buffer[5] << 32) | ((size_t)_ws_receive_buffer[6] << 24) | ((size_t)_ws_receive_buffer[7] << 16) | ((size_t)_ws_receive_buffer[8] << 8) | ((size_t)_ws_receive_buffer[9] << 0));
            _ws_header_size = 10 + (mask ? 4 : 0);
            _ws_payload_size = payload;
            _ws_receive_buffer.reserve(_ws_header_size + _ws_payload_size);
        }

        // Prepare WebSocket frame mask
        if (mask)
        {
            if (_ws_receive_buffer.size() < _ws_header_size)
            {
                for (size_t i = 0; i < 4; ++i, ++data, --size)
                {
                    if (size == 0)
                        return;
                    _ws_receive_buffer.push_back(*data);
                    _ws_receive_mask[i] = *data;
                }
            }
        }

        size_t total = _ws_header_size + _ws_payload_size;
        size_t length = std::min(total - _ws_receive_buffer.size(), size);

        // Prepare WebSocket frame payload
        _ws_receive_buffer.insert(_ws_receive_buffer.end(), data, data + length);
        data += length;
        size -= length;

        // Process WebSocket frame
        if (_ws_receive_buffer.size() == total)
        {
            size_t offset = _ws_header_size;

            // Unmask WebSocket frame content
            if (mask)
                for (size_t i = 0; i < _ws_payload_size; ++i)
                    _ws_receive_buffer[offset + i] ^= _ws_receive_mask[i % 4];

            _ws_received = true;

            if ((opcode & WS_PING) == WS_PING)
            {
                // Call the WebSocket ping handler
                onWSPing(_ws_receive_buffer.data() + offset, _ws_payload_size);
            }
            else if ((opcode & WS_PONG) == WS_PONG)
            {
                // Call the WebSocket pong handler
                onWSPong(_ws_receive_buffer.data() + offset, _ws_payload_size);
            }
            else if ((opcode & WS_CLOSE) == WS_CLOSE)
            {
                // Call the WebSocket close handler
                onWSClose(_ws_receive_buffer.data() + offset, _ws_payload_size);
            }
            else if (((opcode & WS_TEXT) == WS_TEXT) || ((opcode & WS_BINARY) == WS_BINARY))
            {
                // Call the WebSocket received handler
                onWSReceived(_ws_receive_buffer.data() + offset, _ws_payload_size);
            }
        }
    }
}

size_t WebSocket::RequiredReceiveFrameSize()
{
    if (_ws_received)
        return 0;

    // Required WebSocket frame opcode and mask flag
    if (_ws_receive_buffer.size() < 2)
        return 2 - _ws_receive_buffer.size();

    bool mask = ((_ws_receive_buffer[1] >> 7) & 0x01) != 0;
    size_t payload = _ws_receive_buffer[1] & (~0x80);

    // Required WebSocket frame size
    if ((payload == 126) && (_ws_receive_buffer.size() < 4))
        return 4 - _ws_receive_buffer.size();
    if ((payload == 127) && (_ws_receive_buffer.size() < 10))
        return 10 - _ws_receive_buffer.size();

    // Required WebSocket frame mask
    if ((mask) && (_ws_receive_buffer.size() < _ws_header_size))
        return _ws_header_size - _ws_receive_buffer.size();

    // Required WebSocket frame payload
    return _ws_header_size + _ws_payload_size - _ws_receive_buffer.size();
}

void WebSocket::ClearWSBuffers()
{
    _ws_received = false;
    _ws_header_size = 0;
    _ws_payload_size = 0;
    _ws_receive_buffer.clear();
    *((uint32_t*)_ws_receive_mask) = 0;

    std::scoped_lock locker(_ws_send_lock);

    _ws_send_buffer.clear();
    *((uint32_t*)_ws_send_mask) = 0;
}

} // namespace WS
} // namespace CppServer
<commit_msg>Fixed WebSockets for Firefox<commit_after>/*!
    \file ws.cpp
    \brief WebSocket C++ Library implementation
    \author Ivan Shynkarenka
    \date 22.05.2019
    \copyright MIT License
*/

#include "server/ws/ws.h"

#include "string/encoding.h"
#include "string/format.h"

#include <algorithm>
#include <openssl/sha.h>

namespace CppServer {
namespace WS {

bool WebSocket::PerformClientUpgrade(const HTTP::HTTPResponse& response, const CppCommon::UUID& id)
{
    if (response.status() != 101)
        return false;

    bool error = false;
    bool accept = false;
    bool connection = false;
    bool upgrade = false;

    // Validate WebSocket handshake headers
    for (size_t i = 0; i < response.headers(); ++i)
    {
        auto header = response.header(i);
        auto key = std::get<0>(header);
        auto value = std::get<1>(header);

        if (key == "Connection")
        {
            if (value != "Upgrade")
            {
                error = true;
                onWSError("Invalid WebSocket handshaked response: 'Connection' header value must be 'Upgrade'");
                break;
            }

            connection = true;
        }
        else if (key == "Upgrade")
        {
            if (value != "websocket")
            {
                error = true;
                onWSError("Invalid WebSocket handshaked response: 'Upgrade' header value must be 'websocket'");
                break;
            }

            upgrade = true;
        }
        else if (key == "Sec-WebSocket-Accept")
        {
            // Calculate the original WebSocket hash
            std::string wskey = CppCommon::Encoding::Base64Encode(id.string()) + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
            char wshash[SHA_DIGEST_LENGTH];
            SHA1((const unsigned char*)wskey.data(), wskey.size(), (unsigned char*)wshash);

            // Get the received WebSocket hash
            wskey = CppCommon::Encoding::Base64Decode(value);

            // Compare original and received hashes
            if (std::strncmp(wskey.data(), wshash, std::min(wskey.size(), sizeof(wshash))) != 0)
            {
                error = true;
                onWSError("Invalid WebSocket handshaked response: 'Sec-WebSocket-Accept' value validation failed");
                break;
            }

            accept = true;
        }
    }

    // Failed to perfrom WebSocket handshake
    if (!accept || !connection || !upgrade)
    {
        if (!error)
            onWSError("Invalid WebSocket response");
        return false;
    }

    // WebSocket successfully handshaked!
    _ws_handshaked = true;
    *((uint32_t*)_ws_send_mask) = rand();
    onWSConnected(response);

    return true;
}

bool WebSocket::PerformServerUpgrade(const HTTP::HTTPRequest& request, HTTP::HTTPResponse& response)
{
    if (request.method() != "GET")
        return false;

    bool error = false;
    bool connection = false;
    bool upgrade = false;
    bool ws_key = false;
    bool ws_version = false;

    std::string accept;

    // Validate WebSocket handshake headers
    for (size_t i = 0; i < request.headers(); ++i)
    {
        auto header = request.header(i);
        auto key = std::get<0>(header);
        auto value = std::get<1>(header);

        if (key == "Connection")
        {
            if ((value != "Upgrade") && (value != "keep-alive, Upgrade"))
            {
                error = true;
                response.MakeErrorResponse("Invalid WebSocket handshaked request: 'Connection' header value must be 'Upgrade' or 'keep-alive, Upgrade'", 400);
                break;
            }

            connection = true;
        }
        else if (key == "Upgrade")
        {
            if (value != "websocket")
            {
                error = true;
                response.MakeErrorResponse("Invalid WebSocket handshaked request: 'Upgrade' header value must be 'websocket'", 400);
                break;
            }

            upgrade = true;
        }
        else if (key == "Sec-WebSocket-Key")
        {
            if (value.empty())
            {
                error = true;
                response.MakeErrorResponse("Invalid WebSocket handshaked request: 'Sec-WebSocket-Key' header value must be non empty", 400);
                break;
            }

            // Calculate WebSocket accept value
            std::string wskey = std::string(value) + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
            char wshash[SHA_DIGEST_LENGTH];
            SHA1((const unsigned char*)wskey.data(), wskey.size(), (unsigned char*)wshash);

            accept = CppCommon::Encoding::Base64Encode(std::string(wshash, sizeof(wshash)));

            ws_key = true;
        }
        else if (key == "Sec-WebSocket-Version")
        {
            if (value != "13")
            {
                error = true;
                response.MakeErrorResponse("Invalid WebSocket handshaked request: 'Sec-WebSocket-Version' header value must be '13'", 400);
                break;
            }

            ws_version = true;
        }
    }

    // Filter out non WebSocket handshake requests
    if (!connection && !upgrade && !ws_key && !ws_version)
        return false;

    // Failed to perfrom WebSocket handshake
    if (!connection || !upgrade || !ws_key || !ws_version)
    {
        if (!error)
            response.MakeErrorResponse("Invalid WebSocket response", 400);
        SendResponse(response);
        return false;
    }

    // Prepare WebSocket upgrade success response
    response.Clear();
    response.SetBegin(101, "HTTP/1.1");
    response.SetHeader("Connection", "Upgrade");
    response.SetHeader("Upgrade", "websocket");
    response.SetHeader("Sec-WebSocket-Accept", accept);
    response.SetBody();

    // Validate WebSocket upgrade request and response
    if (!onWSConnecting(request, response))
        return false;

    // Send WebSocket upgrade response
    SendResponse(response);

    // WebSocket successfully handshaked!
    _ws_handshaked = true;
    *((uint32_t*)_ws_send_mask) = 0;
    onWSConnected(request);

    return true;
}

void WebSocket::PrepareSendFrame(uint8_t opcode, bool mask, const void* buffer, size_t size, int status)
{
    // Clear the previous WebSocket send buffer
    _ws_send_buffer.clear();

    // Append WebSocket frame opcode
    _ws_send_buffer.push_back(opcode);

    // Append WebSocket frame size
    if (size <= 125)
        _ws_send_buffer.push_back((size & 0xFF) | (mask ? 0x80 : 0));
    else if (size <= 65535)
    {
        _ws_send_buffer.push_back(126 | (mask ? 0x80 : 0));
        _ws_send_buffer.push_back((size >> 8) & 0xFF);
        _ws_send_buffer.push_back(size & 0xFF);
    }
    else
    {
        _ws_send_buffer.push_back(127 | (mask ? 0x80 : 0));
        for (int i = 7; i >= 0; --i)
            _ws_send_buffer.push_back((size >> (8 * i)) & 0xFF);
    }

    if (mask)
    {
        // Append WebSocket frame mask
        _ws_send_buffer.push_back(_ws_send_mask[0]);
        _ws_send_buffer.push_back(_ws_send_mask[1]);
        _ws_send_buffer.push_back(_ws_send_mask[2]);
        _ws_send_buffer.push_back(_ws_send_mask[3]);
    }

    // Resize WebSocket frame buffer
    size_t offset = _ws_send_buffer.size();
    _ws_send_buffer.resize(offset + size);

    // Mask WebSocket frame content
    const uint8_t* data = (const uint8_t*)buffer;
    for (size_t i = 0; i < size; ++i)
        _ws_send_buffer[offset + i] = data[i] ^ _ws_send_mask[i % 4];
}

void WebSocket::PrepareReceiveFrame(const void* buffer, size_t size)
{
    const uint8_t* data = (const uint8_t*)buffer;

    // Clear received data after WebSocket frame was processed
    if (_ws_received)
    {
        _ws_received = false;
        _ws_header_size = 0;
        _ws_payload_size = 0;
        _ws_receive_buffer.clear();
        *((uint32_t*)_ws_receive_mask) = 0;
    }

    while (size > 0)
    {
        // Clear received data after WebSocket frame was processed
        if (_ws_received)
        {
            _ws_received = false;
            _ws_header_size = 0;
            _ws_payload_size = 0;
            _ws_receive_buffer.clear();
            *((uint32_t*)_ws_receive_mask) = 0;
        }

        // Prepare WebSocket frame opcode and mask flag
        if (_ws_receive_buffer.size() < 2)
        {
            for (size_t i = 0; i < 2; ++i, ++data, --size)
            {
                if (size == 0)
                    return;
                _ws_receive_buffer.push_back(*data);
            }
        }

        uint8_t opcode = _ws_receive_buffer[0] & 0x0F;
        [[maybe_unused]] bool fin = ((_ws_receive_buffer[0] >> 7) & 0x01) != 0;
        bool mask = ((_ws_receive_buffer[1] >> 7) & 0x01) != 0;
        size_t payload = _ws_receive_buffer[1] & (~0x80);

        // Prepare WebSocket frame size
        if (payload <= 125)
        {
            _ws_header_size = 2 + (mask ? 4 : 0);
            _ws_payload_size = payload;
            _ws_receive_buffer.reserve(_ws_header_size + _ws_payload_size);
        }
        else if (payload == 126)
        {
            if (_ws_receive_buffer.size() < 4)
            {
                for (size_t i = 0; i < 2; ++i, ++data, --size)
                {
                    if (size == 0)
                        return;
                    _ws_receive_buffer.push_back(*data);
                }
            }

            payload = (((size_t)_ws_receive_buffer[2] << 8) | ((size_t)_ws_receive_buffer[3] << 0));
            _ws_header_size = 4 + (mask ? 4 : 0);
            _ws_payload_size = payload;
            _ws_receive_buffer.reserve(_ws_header_size + _ws_payload_size);
        }
        else if (payload == 127)
        {
            if (_ws_receive_buffer.size() < 10)
            {
                for (size_t i = 0; i < 8; ++i, ++data, --size)
                {
                    if (size == 0)
                        return;
                    _ws_receive_buffer.push_back(*data);
                }
            }

            payload = (((size_t)_ws_receive_buffer[2] << 56) | ((size_t)_ws_receive_buffer[3] << 48) | ((size_t)_ws_receive_buffer[4] << 40) | ((size_t)_ws_receive_buffer[5] << 32) | ((size_t)_ws_receive_buffer[6] << 24) | ((size_t)_ws_receive_buffer[7] << 16) | ((size_t)_ws_receive_buffer[8] << 8) | ((size_t)_ws_receive_buffer[9] << 0));
            _ws_header_size = 10 + (mask ? 4 : 0);
            _ws_payload_size = payload;
            _ws_receive_buffer.reserve(_ws_header_size + _ws_payload_size);
        }

        // Prepare WebSocket frame mask
        if (mask)
        {
            if (_ws_receive_buffer.size() < _ws_header_size)
            {
                for (size_t i = 0; i < 4; ++i, ++data, --size)
                {
                    if (size == 0)
                        return;
                    _ws_receive_buffer.push_back(*data);
                    _ws_receive_mask[i] = *data;
                }
            }
        }

        size_t total = _ws_header_size + _ws_payload_size;
        size_t length = std::min(total - _ws_receive_buffer.size(), size);

        // Prepare WebSocket frame payload
        _ws_receive_buffer.insert(_ws_receive_buffer.end(), data, data + length);
        data += length;
        size -= length;

        // Process WebSocket frame
        if (_ws_receive_buffer.size() == total)
        {
            size_t offset = _ws_header_size;

            // Unmask WebSocket frame content
            if (mask)
                for (size_t i = 0; i < _ws_payload_size; ++i)
                    _ws_receive_buffer[offset + i] ^= _ws_receive_mask[i % 4];

            _ws_received = true;

            if ((opcode & WS_PING) == WS_PING)
            {
                // Call the WebSocket ping handler
                onWSPing(_ws_receive_buffer.data() + offset, _ws_payload_size);
            }
            else if ((opcode & WS_PONG) == WS_PONG)
            {
                // Call the WebSocket pong handler
                onWSPong(_ws_receive_buffer.data() + offset, _ws_payload_size);
            }
            else if ((opcode & WS_CLOSE) == WS_CLOSE)
            {
                // Call the WebSocket close handler
                onWSClose(_ws_receive_buffer.data() + offset, _ws_payload_size);
            }
            else if (((opcode & WS_TEXT) == WS_TEXT) || ((opcode & WS_BINARY) == WS_BINARY))
            {
                // Call the WebSocket received handler
                onWSReceived(_ws_receive_buffer.data() + offset, _ws_payload_size);
            }
        }
    }
}

size_t WebSocket::RequiredReceiveFrameSize()
{
    if (_ws_received)
        return 0;

    // Required WebSocket frame opcode and mask flag
    if (_ws_receive_buffer.size() < 2)
        return 2 - _ws_receive_buffer.size();

    bool mask = ((_ws_receive_buffer[1] >> 7) & 0x01) != 0;
    size_t payload = _ws_receive_buffer[1] & (~0x80);

    // Required WebSocket frame size
    if ((payload == 126) && (_ws_receive_buffer.size() < 4))
        return 4 - _ws_receive_buffer.size();
    if ((payload == 127) && (_ws_receive_buffer.size() < 10))
        return 10 - _ws_receive_buffer.size();

    // Required WebSocket frame mask
    if ((mask) && (_ws_receive_buffer.size() < _ws_header_size))
        return _ws_header_size - _ws_receive_buffer.size();

    // Required WebSocket frame payload
    return _ws_header_size + _ws_payload_size - _ws_receive_buffer.size();
}

void WebSocket::ClearWSBuffers()
{
    _ws_received = false;
    _ws_header_size = 0;
    _ws_payload_size = 0;
    _ws_receive_buffer.clear();
    *((uint32_t*)_ws_receive_mask) = 0;

    std::scoped_lock locker(_ws_send_lock);

    _ws_send_buffer.clear();
    *((uint32_t*)_ws_send_mask) = 0;
}

} // namespace WS
} // namespace CppServer
<|endoftext|>


<commit_before>#include "map.h"

Map::Map(Vector2i mapSize, Vector2f blockSize) {
	this->size = mapSize;

	this->renderTex.create(this->size.x * blockSize.x, this->size.y * blockSize.y);

	for (int row = 0; row < mapSize.x; row++) {
		this->mapData.push_back(std::vector<int>());

		for (int col = 0; col < mapSize.y; col++) {
			this->mapData[row].push_back(255);
		}
	}

	block.setSize(blockSize);
}

Map::~Map() {}

void Map::Save(std::string fileName) {
	// SAVE BLOCK HEIGHTS
	this->renderTex.clear();

	RectangleShape block;

	for (unsigned int row = 0; row < this->size.x; row++) {
		for (unsigned int col = 0; col < this->size.y; col++) {
			block.setFillColor(Color(this->mapData[row][col], this->mapData[row][col], this->mapData[row][col], 255));
			block.setPosition(Vector2f(row * block.getSize().x, col * block.getSize().y));
			renderTex.draw(block);
		}
	}

	this->renderTex.display();

	sf::Image image;
	image.create(size.x, size.y, sf::Color::Black);

	for (unsigned int row = 0; row < this->size.x; row++) {
		for (unsigned int col = 0; col < this->size.y; col++) {
			image.setPixel(row, col, Color(mapData[row][col], mapData[row][col], mapData[row][col], 255));
		}
	}

	image.saveToFile("maps/" + fileName + ".png");

	// SAVE OBJECT DATA
	std::ofstream out;

	out.open("maps/" + fileName + ".toobj");

	if (out.is_open()) {
		for (size_t i = 0; i < this->enemySpawners.size(); i++) {
			out << enemySpawners[i].position.x << " " << enemySpawners[i].position.y << " " << enemySpawners[i].eCount << " " << enemySpawners[i].eHP << "\n";
		}
	} else
		std::cout << "ERROR OPENING FILE " << (fileName) << "\n";

	out.close();
}

void Map::Draw(RenderTarget& target) {
	this->renderTex.clear(Color::White);

	for (unsigned int row = 0; row < this->size.x; row++) {
		for (unsigned int col = 0; col < this->size.y; col++) {
			block.setFillColor(Color(this->mapData[row][col], this->mapData[row][col], this->mapData[row][col], 255));
			block.setPosition(Vector2f(row * block.getSize().x, col * block.getSize().y));
			renderTex.draw(block);

			/*std::cout << block.getPosition().x << " " <<  block.getPosition().y << "\n";*/
		}
	}

	this->renderTex.display();

	this->renderSprite.setTexture(this->renderTex.getTexture());

	target.draw(this->renderSprite);
}
<commit_msg>Fixed int compare thingy<commit_after>#include "map.h"

Map::Map(Vector2i mapSize, Vector2f blockSize) {
	this->size = mapSize;

	this->renderTex.create(this->size.x * blockSize.x, this->size.y * blockSize.y);

	for (int row = 0; row < mapSize.x; row++) {
		this->mapData.push_back(std::vector<int>());

		for (int col = 0; col < mapSize.y; col++) {
			this->mapData[row].push_back(255);
		}
	}

	block.setSize(blockSize);
}

Map::~Map() {}

void Map::Save(std::string fileName) {
	// SAVE BLOCK HEIGHTS
	this->renderTex.clear();

	RectangleShape block;

	for (int row = 0; row < this->size.x; row++) {
		for (int col = 0; col < this->size.y; col++) {
			block.setFillColor(Color(this->mapData[row][col], this->mapData[row][col], this->mapData[row][col], 255));
			block.setPosition(Vector2f(row * block.getSize().x, col * block.getSize().y));
			renderTex.draw(block);
		}
	}

	this->renderTex.display();

	sf::Image image;
	image.create(size.x, size.y, sf::Color::Black);

	for (int row = 0; row < this->size.x; row++) {
		for (int col = 0; col < this->size.y; col++) {
			image.setPixel(row, col, Color(mapData[row][col], mapData[row][col], mapData[row][col], 255));
		}
	}

	image.saveToFile("maps/" + fileName + ".png");

	// SAVE OBJECT DATA
	std::ofstream out;

	out.open("maps/" + fileName + ".toobj");

	if (out.is_open()) {
		for (size_t i = 0; i < this->enemySpawners.size(); i++) {
			out << enemySpawners[i].position.x << " " << enemySpawners[i].position.y << " " << enemySpawners[i].eCount << " " << enemySpawners[i].eHP << "\n";
		}
	} else
		std::cout << "ERROR OPENING FILE " << (fileName) << "\n";

	out.close();
}

void Map::Draw(RenderTarget& target) {
	this->renderTex.clear(Color::White);

	for (int row = 0; row < this->size.x; row++) {
		for (int col = 0; col < this->size.y; col++) {
			block.setFillColor(Color(this->mapData[row][col], this->mapData[row][col], this->mapData[row][col], 255));
			block.setPosition(Vector2f(row * block.getSize().x, col * block.getSize().y));
			renderTex.draw(block);

			/*std::cout << block.getPosition().x << " " <<  block.getPosition().y << "\n";*/
		}
	}

	this->renderTex.display();

	this->renderSprite.setTexture(this->renderTex.getTexture());

	target.draw(this->renderSprite);
}
<|endoftext|>


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

#include "main.h"
#include "primitives/transaction.h"
#include "consensus/validation.h"

TEST(checktransaction_tests, check_vpub_not_both_nonzero) {
    CMutableTransaction tx;
    tx.nVersion = 2;

    {
        // Ensure that values within the pour are well-formed.
        CMutableTransaction newTx(tx);
        CValidationState state;
        state.SetPerformPourVerification(false); // don't verify the snark

        newTx.vpour.push_back(CPourTx());

        CPourTx *pourtx = &newTx.vpour[0];
        pourtx->vpub_old = 1;
        pourtx->vpub_new = 1;

        EXPECT_FALSE(CheckTransaction(newTx, state));
        EXPECT_EQ(state.GetRejectReason(), "bad-txns-vpubs-both-nonzero");
    }
}

class MockCValidationState : public CValidationState {
public:
    MOCK_METHOD5(DoS, bool(int level, bool ret,
             unsigned char chRejectCodeIn, std::string strRejectReasonIn,
             bool corruptionIn));
    MOCK_METHOD1(SetPerformPourVerification, bool(bool pourVerifyIn));
    MOCK_METHOD0(PerformPourVerification, bool());
    MOCK_METHOD3(Invalid, bool(bool ret,
                 unsigned char _chRejectCode, std::string _strRejectReason));
    MOCK_METHOD1(Error, bool(std::string strRejectReasonIn));
    MOCK_CONST_METHOD0(IsValid, bool());
    MOCK_CONST_METHOD0(IsInvalid, bool());
    MOCK_CONST_METHOD0(IsError, bool());
    MOCK_CONST_METHOD1(IsInvalid, bool(int &nDoSOut));
    MOCK_CONST_METHOD0(CorruptionPossible, bool());
    MOCK_CONST_METHOD0(GetRejectCode, unsigned char());
    MOCK_CONST_METHOD0(GetRejectReason, std::string());
};

TEST(checktransaction_tests, mock_proof_of_concept) {
    CTransaction tx;
    MockCValidationState state;
    EXPECT_CALL(state, DoS(10, false, REJECT_INVALID, "bad-txns-vin-empty", false)).Times(1);
    CheckTransaction(tx, state);
}
<commit_msg>WIP: Add mock test coverage of CheckTransaction<commit_after>#include <gtest/gtest.h>
#include <gmock/gmock.h>

#include "main.h"
#include "primitives/transaction.h"
#include "consensus/validation.h"

TEST(checktransaction_tests, check_vpub_not_both_nonzero) {
    CMutableTransaction tx;
    tx.nVersion = 2;

    {
        // Ensure that values within the pour are well-formed.
        CMutableTransaction newTx(tx);
        CValidationState state;
        state.SetPerformPourVerification(false); // don't verify the snark

        newTx.vpour.push_back(CPourTx());

        CPourTx *pourtx = &newTx.vpour[0];
        pourtx->vpub_old = 1;
        pourtx->vpub_new = 1;

        EXPECT_FALSE(CheckTransaction(newTx, state));
        EXPECT_EQ(state.GetRejectReason(), "bad-txns-vpubs-both-nonzero");
    }
}

class MockCValidationState : public CValidationState {
public:
    MOCK_METHOD5(DoS, bool(int level, bool ret,
             unsigned char chRejectCodeIn, std::string strRejectReasonIn,
             bool corruptionIn));
    MOCK_METHOD1(SetPerformPourVerification, bool(bool pourVerifyIn));
    MOCK_METHOD0(PerformPourVerification, bool());
    MOCK_METHOD3(Invalid, bool(bool ret,
                 unsigned char _chRejectCode, std::string _strRejectReason));
    MOCK_METHOD1(Error, bool(std::string strRejectReasonIn));
    MOCK_CONST_METHOD0(IsValid, bool());
    MOCK_CONST_METHOD0(IsInvalid, bool());
    MOCK_CONST_METHOD0(IsError, bool());
    MOCK_CONST_METHOD1(IsInvalid, bool(int &nDoSOut));
    MOCK_CONST_METHOD0(CorruptionPossible, bool());
    MOCK_CONST_METHOD0(GetRejectCode, unsigned char());
    MOCK_CONST_METHOD0(GetRejectReason, std::string());
};

CMutableTransaction GetValidTransaction() {
    CMutableTransaction mtx;
    mtx.vin.resize(2);
    mtx.vin[0].prevout.hash = uint256S("0000000000000000000000000000000000000000000000000000000000000000");
    mtx.vin[0].prevout.n = 0;
    mtx.vin[1].prevout.hash = uint256S("0000000000000000000000000000000000000000000000000000000000000001");
    mtx.vin[1].prevout.n = 0;
    mtx.vout.resize(2);
    // mtx.vout[0].scriptPubKey =
    mtx.vout[0].nValue = 0;
    mtx.vout[1].nValue = 0;
    mtx.vpour.resize(2);
    mtx.vpour[0].serials.at(0) = uint256S("0000000000000000000000000000000000000000000000000000000000000000");
    mtx.vpour[0].serials.at(1) = uint256S("0000000000000000000000000000000000000000000000000000000000000001");
    mtx.vpour[1].serials.at(0) = uint256S("0000000000000000000000000000000000000000000000000000000000000002");
    mtx.vpour[1].serials.at(1) = uint256S("0000000000000000000000000000000000000000000000000000000000000003");
    return mtx;
}

TEST(checktransaction_tests, valid_transaction) {
    CMutableTransaction mtx = GetValidTransaction();
    CTransaction tx(mtx);
    MockCValidationState state;
    EXPECT_TRUE(CheckTransaction(tx, state));
}

TEST(checktransaction_tests, bad_txns_vin_empty) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vpour.resize(0);
    mtx.vin.resize(0);

    CTransaction tx(mtx);
    MockCValidationState state;
    EXPECT_CALL(state, DoS(10, false, REJECT_INVALID, "bad-txns-vin-empty", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_vout_empty) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vpour.resize(0);
    mtx.vout.resize(0);

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(10, false, REJECT_INVALID, "bad-txns-vout-empty", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_oversize) {
    CMutableTransaction mtx = GetValidTransaction();

    mtx.vin[0].scriptSig = CScript();
    // 18 * (520char + DROP) + OP_1 = 9433 bytes
    std::vector<unsigned char> vchData(520);
    for (unsigned int i = 0; i < 2000; ++i)
        mtx.vin[0].scriptSig << vchData << OP_DROP;
    mtx.vin[0].scriptSig << OP_1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-oversize", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_vout_negative) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vout[0].nValue = -1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-vout-negative", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_vout_toolarge) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vout[0].nValue = MAX_MONEY + 1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-vout-toolarge", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_txouttotal_toolarge_outputs) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vout[0].nValue = MAX_MONEY;
    mtx.vout[1].nValue = 1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-txouttotal-toolarge", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_txouttotal_toolarge_joinsplit) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vout[0].nValue = 1;
    mtx.vpour[0].vpub_new = MAX_MONEY;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-txouttotal-toolarge", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_vpub_old_negative) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vpour[0].vpub_old = -1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-vpub_old-negative", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_vpub_new_negative) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vpour[0].vpub_new = -1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-vpub_new-negative", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_vpub_old_toolarge) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vpour[0].vpub_old = MAX_MONEY + 1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-vpub_old-toolarge", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_vpub_new_toolarge) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vpour[0].vpub_new = MAX_MONEY + 1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-vpub_new-toolarge", false)).Times(1);
    CheckTransaction(tx, state);
}

TEST(checktransaction_tests, bad_txns_vpubs_both_nonzero) {
    CMutableTransaction mtx = GetValidTransaction();
    mtx.vpour[0].vpub_old = 1;
    mtx.vpour[0].vpub_new = 1;

    CTransaction tx(mtx);

    MockCValidationState state;
    EXPECT_CALL(state, DoS(100, false, REJECT_INVALID, "bad-txns-vpubs-both-nonzero", false)).Times(1);
    CheckTransaction(tx, state);
}
<|endoftext|>