<commit_before>/*
   (c) Copyright 2000  convergence integrated media GmbH.
   All rights reserved.

   Written by Denis Oliver Kropp <dok@convergence.de>,
              Andreas Hundt <andi@convergence.de> and
              Sven Neumann <sven@convergence.de>.

   This library is free software; you can redistribute it and/or
   modify it under the terms of the GNU Lesser General Public
   License as published by the Free Software Foundation; either
   version 2 of the License, or (at your option) any later version.

   This library is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public
   License along with this library; if not, write to the
   Free Software Foundation, Inc., 59 Temple Place - Suite 330,
   Boston, MA 02111-1307, USA.
*/

extern "C" {
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>

#include <malloc.h>

#include <directfb.h>

#include <misc/util.h>

#include <core/core.h>
#include <core/coredefs.h>
#include <core/layers.h>
#include <core/gfxcard.h>

#include <display/idirectfbsurface.h>
}

#include <aviplay.h>

/*
 * private data struct of IDirectFBVideoProvider
 */
typedef struct {
     int               ref;       /* reference counter */
     IAviPlayer       *player;

     IDirectFBSurface *destination;
     DFBRectangle      dest_rect;
     DVFrameCallback   callback;
     void             *ctx;

     CardState         state;
     CoreSurface       source;
} IDirectFBVideoProvider_AviFile_data;

static void IDirectFBVideoProvider_AviFile_Destruct(
                                                  IDirectFBVideoProvider *thiz )
{
     IDirectFBVideoProvider_AviFile_data *data;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (data->player->IsPlaying())
          data->player->Stop();

     reactor_free( data->source.reactor );

     delete thiz->priv;
     thiz->priv = NULL;

#ifndef DFB_DEBUG
     free( thiz );
#endif
}

static DFBResult IDirectFBVideoProvider_AviFile_AddRef(
                                                  IDirectFBVideoProvider *thiz )
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     data->ref++;

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_Release(
                                                  IDirectFBVideoProvider *thiz )
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
          return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     if (--data->ref == 0) {
          IDirectFBVideoProvider_AviFile_Destruct( thiz );
     }

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_GetSurfaceDescription(
                                                  IDirectFBVideoProvider *thiz,
                                                  DFBSurfaceDescription  *desc )
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz || !desc)
          return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     desc->flags = (DFBSurfaceDescriptionFlags)
                            (DSDESC_WIDTH | DSDESC_HEIGHT | DSDESC_PIXELFORMAT);
     desc->width  = data->player->GetWidth();
     desc->height = data->player->GetHeight();
     desc->pixelformat = layers->surface->format;

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_PlayTo(
                                           IDirectFBVideoProvider *thiz,
                                           IDirectFBSurface       *destination,
                                           DFBRectangle           *dstrect,
                                           DVFrameCallback         callback,
                                           void                   *ctx )
{
     DFBRectangle                         rect;
     IDirectFBVideoProvider_AviFile_data *data;
     IDirectFBSurface_data               *dst_data;

     if (!thiz || !destination)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;
     dst_data = (IDirectFBSurface_data*)destination->priv;

     if (!data || !dst_data)
          return DFB_DEAD;

//     thiz->Stop( thiz );

     /* URGENT: keep in sync with DrawCallback */

     /* build the destination rectangle */
     if (dstrect) {
          if (dstrect->w < 1  ||  dstrect->h < 1)
               return DFB_INVARG;

          rect = *dstrect;

          rect.x += dst_data->req_rect.x;
          rect.y += dst_data->req_rect.y;
     }
     else
          rect = dst_data->req_rect;

     /* save for later blitting operation */
     data->dest_rect = rect;

     /* build the clip rectangle */
     if (!rectangle_intersect( &rect, &dst_data->clip_rect ))
          return DFB_INVARG;

     /* put the destination clip into the state */
     data->state.clip.x1 = rect.x;
     data->state.clip.y1 = rect.y;
     data->state.clip.x2 = rect.x + rect.w - 1;
     data->state.clip.y2 = rect.y + rect.h - 1;
     data->state.destination = dst_data->surface;
     data->state.modified = (StateModificationFlags)
                            (data->state.modified | SMF_CLIP | SMF_DESTINATION);

     if (data->destination) {
          data->destination->Release( data->destination );
          data->destination = NULL;     /* FIXME: remove listener */
     }

     destination->AddRef( destination );
     data->destination = destination;   /* FIXME: install listener */

     data->callback = callback;
     data->ctx = ctx;

     switch (data->player->GetState( NULL )) {
          case IAviPlayer::Playing:
               break;
          case IAviPlayer::Paused:
               data->player->Pause(false);
               break;
          default:
               data->player->Start();
     }

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_Stop(
                                                  IDirectFBVideoProvider *thiz )
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     if (data->player->IsPlaying())
          data->player->Pause(true);

     if (data->destination) {
          data->destination->Release( data->destination );
          data->destination = NULL;     /* FIXME: remove listener */
     }

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_SeekTo(
                                               IDirectFBVideoProvider *thiz,
                                               double                  seconds)
{
     IDirectFBVideoProvider_AviFile_data *data;
     double curpos;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     curpos = data->player->GetPos();
     data->player->Reseek( seconds );
     /* seeking forward for some small amount may actually bring us
      * _back_ to the last key frame -> compensate via PageUp()
      */
     if (seconds > curpos && curpos > data->player->GetPos())
       data->player->PageUp();

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_GetPos(
     IDirectFBVideoProvider *thiz,
     double                 *seconds)
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     *seconds = data->player->GetPos();

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_GetLength(
     IDirectFBVideoProvider *thiz,
     double                 *seconds)
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     *seconds = data->player->GetVideoLength();

     return DFB_OK;
}

static void AviFile_KillCallback( int bogus, void *p )
{
     IDirectFBVideoProvider_AviFile_data *data =
          (IDirectFBVideoProvider_AviFile_data*)p;

     /* AviFile_KillCallback gets called when AviFile->Stop is called */
     /* At the moment we do nothing here...                           */
}

static void AviFile_DrawCallback( const CImage *image, void *p )
{
     IDirectFBVideoProvider_AviFile_data *data =
          (IDirectFBVideoProvider_AviFile_data*)p;

     data->source.front_buffer->system.addr   = (void*)(image->Data());
     data->source.front_buffer->system.pitch  = image->Bpl();

     DFBRectangle rect = { 0, 0, image->Width(), image->Height() };

     if (rect.w == data->dest_rect.w  &&  rect.h == data->dest_rect.h) {
          gfxcard_blit( &rect, data->dest_rect.x,
                        data->dest_rect.y, &data->state );
     }
     else {
          DFBRectangle drect = data->dest_rect;

          gfxcard_stretchblit( &rect, &drect, &data->state );
     }

     if (data->callback)
          data->callback( data->ctx );
}

/* exported symbols */

extern "C" {

const char *get_type()
{
     return "IDirectFBVideoProvider";
}

const char *get_implementation()
{
     return "AviFile";
}

DFBResult Probe( const char *filename )
{
     if (strstr( filename, ".avi" ) ||
         strstr( filename, ".AVI" ) ||
         strstr( filename, ".asf" ) ||
         strstr( filename, ".ASF" ))
          return DFB_OK;

     return DFB_UNSUPPORTED;
}

DFBResult Construct( IDirectFBVideoProvider *thiz, const char *filename )
{
     IDirectFBVideoProvider_AviFile_data *data;

     data = (IDirectFBVideoProvider_AviFile_data*)
          calloc( 1, sizeof(IDirectFBVideoProvider_AviFile_data) );

     thiz->priv = data;

     data->ref = 1;

     try {
          data->player = CreateAviPlayer( filename, 16 /* FIXME */ );

          data->player->SetDrawCallback2( AviFile_DrawCallback, data );
          data->player->SetKillHandler( AviFile_KillCallback, data );
     }
     catch (FatalError e) {
          ERRORMSG( "DirectFB/AviFile: CreateAviPlayer failed: %s\n",
                    e.GetDesc() );
          free( data );
          return DFB_FAILURE;
     }

     data->source.width  = data->player->GetWidth();
     data->source.height = data->player->GetHeight();
     data->source.format = DSPF_RGB16;

     data->source.front_buffer =
          (SurfaceBuffer*) calloc( 1, sizeof(SurfaceBuffer) );

     data->source.front_buffer->policy = CSP_SYSTEMONLY;
     data->source.front_buffer->system.health = CSH_STORED;

     data->source.back_buffer = data->source.front_buffer;

     pthread_mutex_init( &data->source.front_lock, NULL );
     pthread_mutex_init( &data->source.back_lock, NULL );

     data->source.reactor = reactor_new();

     data->state.source   = &data->source;
     data->state.modified = SMF_ALL;

     thiz->AddRef = IDirectFBVideoProvider_AviFile_AddRef;
     thiz->Release = IDirectFBVideoProvider_AviFile_Release;
     thiz->GetSurfaceDescription =
          IDirectFBVideoProvider_AviFile_GetSurfaceDescription;
     thiz->PlayTo = IDirectFBVideoProvider_AviFile_PlayTo;
     thiz->Stop = IDirectFBVideoProvider_AviFile_Stop;
     thiz->SeekTo = IDirectFBVideoProvider_AviFile_SeekTo;
     thiz->GetPos = IDirectFBVideoProvider_AviFile_GetPos;
     thiz->GetLength = IDirectFBVideoProvider_AviFile_GetLength;

     return DFB_OK;
}

}

<commit_msg>fixed a braino and some compiler warnings<commit_after>/*
   (c) Copyright 2000  convergence integrated media GmbH.
   All rights reserved.

   Written by Denis Oliver Kropp <dok@convergence.de>,
              Andreas Hundt <andi@convergence.de> and
              Sven Neumann <sven@convergence.de>.

   This library is free software; you can redistribute it and/or
   modify it under the terms of the GNU Lesser General Public
   License as published by the Free Software Foundation; either
   version 2 of the License, or (at your option) any later version.

   This library is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public
   License along with this library; if not, write to the
   Free Software Foundation, Inc., 59 Temple Place - Suite 330,
   Boston, MA 02111-1307, USA.
*/

extern "C" {
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>

#include <malloc.h>

#include <directfb.h>

#include <misc/util.h>

#include <core/core.h>
#include <core/coredefs.h>
#include <core/layers.h>
#include <core/gfxcard.h>

#include <display/idirectfbsurface.h>
}

#include <aviplay.h>

/*
 * private data struct of IDirectFBVideoProvider
 */
typedef struct {
     int               ref;       /* reference counter */
     IAviPlayer       *player;

     IDirectFBSurface *destination;
     DFBRectangle      dest_rect;
     DVFrameCallback   callback;
     void             *ctx;

     CardState         state;
     CoreSurface       source;
} IDirectFBVideoProvider_AviFile_data;

static void IDirectFBVideoProvider_AviFile_Destruct(
                                                  IDirectFBVideoProvider *thiz )
{
     IDirectFBVideoProvider_AviFile_data *data;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (data->player->IsPlaying())
          data->player->Stop();

     delete data->player;

     reactor_free( data->source.reactor );

     free (thiz->priv);
     thiz->priv = NULL;

#ifndef DFB_DEBUG
     free( thiz );
#endif
}

static DFBResult IDirectFBVideoProvider_AviFile_AddRef(
                                                  IDirectFBVideoProvider *thiz )
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     data->ref++;

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_Release(
                                                  IDirectFBVideoProvider *thiz )
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
          return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     if (--data->ref == 0) {
          IDirectFBVideoProvider_AviFile_Destruct( thiz );
     }

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_GetSurfaceDescription(
                                                  IDirectFBVideoProvider *thiz,
                                                  DFBSurfaceDescription  *desc )
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz || !desc)
          return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     desc->flags = (DFBSurfaceDescriptionFlags)
                            (DSDESC_WIDTH | DSDESC_HEIGHT | DSDESC_PIXELFORMAT);
     desc->width  = data->player->GetWidth();
     desc->height = data->player->GetHeight();
     desc->pixelformat = layers->surface->format;

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_PlayTo(
                                           IDirectFBVideoProvider *thiz,
                                           IDirectFBSurface       *destination,
                                           DFBRectangle           *dstrect,
                                           DVFrameCallback         callback,
                                           void                   *ctx )
{
     DFBRectangle                         rect;
     IDirectFBVideoProvider_AviFile_data *data;
     IDirectFBSurface_data               *dst_data;

     if (!thiz || !destination)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;
     dst_data = (IDirectFBSurface_data*)destination->priv;

     if (!data || !dst_data)
          return DFB_DEAD;

//     thiz->Stop( thiz );

     /* URGENT: keep in sync with DrawCallback */

     /* build the destination rectangle */
     if (dstrect) {
          if (dstrect->w < 1  ||  dstrect->h < 1)
               return DFB_INVARG;

          rect = *dstrect;

          rect.x += dst_data->req_rect.x;
          rect.y += dst_data->req_rect.y;
     }
     else
          rect = dst_data->req_rect;

     /* save for later blitting operation */
     data->dest_rect = rect;

     /* build the clip rectangle */
     if (!rectangle_intersect( &rect, &dst_data->clip_rect ))
          return DFB_INVARG;

     /* put the destination clip into the state */
     data->state.clip.x1 = rect.x;
     data->state.clip.y1 = rect.y;
     data->state.clip.x2 = rect.x + rect.w - 1;
     data->state.clip.y2 = rect.y + rect.h - 1;
     data->state.destination = dst_data->surface;
     data->state.modified = (StateModificationFlags)
                            (data->state.modified | SMF_CLIP | SMF_DESTINATION);

     if (data->destination) {
          data->destination->Release( data->destination );
          data->destination = NULL;     /* FIXME: remove listener */
     }

     destination->AddRef( destination );
     data->destination = destination;   /* FIXME: install listener */

     data->callback = callback;
     data->ctx = ctx;

     switch (data->player->GetState( NULL )) {
          case IAviPlayer::Playing:
               break;
          case IAviPlayer::Paused:
               data->player->Pause(false);
               break;
          default:
               data->player->Start();
     }

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_Stop(
                                                  IDirectFBVideoProvider *thiz )
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     if (data->player->IsPlaying())
          data->player->Pause(true);

     if (data->destination) {
          data->destination->Release( data->destination );
          data->destination = NULL;     /* FIXME: remove listener */
     }

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_SeekTo(
                                               IDirectFBVideoProvider *thiz,
                                               double                  seconds)
{
     IDirectFBVideoProvider_AviFile_data *data;
     double curpos;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     curpos = data->player->GetPos();
     data->player->Reseek( seconds );
     /* seeking forward for some small amount may actually bring us
      * _back_ to the last key frame -> compensate via PageUp()
      */
     if (seconds > curpos && curpos > data->player->GetPos())
       data->player->PageUp();

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_GetPos(
     IDirectFBVideoProvider *thiz,
     double                 *seconds)
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     *seconds = data->player->GetPos();

     return DFB_OK;
}

static DFBResult IDirectFBVideoProvider_AviFile_GetLength(
     IDirectFBVideoProvider *thiz,
     double                 *seconds)
{
     IDirectFBVideoProvider_AviFile_data *data;

     if (!thiz)
        return DFB_INVARG;

     data = (IDirectFBVideoProvider_AviFile_data*)thiz->priv;

     if (!data)
          return DFB_DEAD;

     *seconds = data->player->GetVideoLength();

     return DFB_OK;
}

static void AviFile_KillCallback( int bogus, void *p )
{
  /* AviFile_KillCallback gets called when AviFile->Stop is called.
     At the moment we do nothing here...

     IDirectFBVideoProvider_AviFile_data *data =
          (IDirectFBVideoProvider_AviFile_data*)p;
   */
}

static void AviFile_DrawCallback( const CImage *image, void *p )
{
     IDirectFBVideoProvider_AviFile_data *data =
          (IDirectFBVideoProvider_AviFile_data*)p;

     data->source.front_buffer->system.addr   = (void*)(image->Data());
     data->source.front_buffer->system.pitch  = image->Bpl();

     DFBRectangle rect = { 0, 0, image->Width(), image->Height() };

     if (rect.w == data->dest_rect.w  &&  rect.h == data->dest_rect.h) {
          gfxcard_blit( &rect, data->dest_rect.x,
                        data->dest_rect.y, &data->state );
     }
     else {
          DFBRectangle drect = data->dest_rect;

          gfxcard_stretchblit( &rect, &drect, &data->state );
     }

     if (data->callback)
          data->callback( data->ctx );
}

/* exported symbols */

extern "C" {

const char *get_type()
{
     return "IDirectFBVideoProvider";
}

const char *get_implementation()
{
     return "AviFile";
}

DFBResult Probe( const char *filename )
{
     if (strstr( filename, ".avi" ) ||
         strstr( filename, ".AVI" ) ||
         strstr( filename, ".asf" ) ||
         strstr( filename, ".ASF" ))
          return DFB_OK;

     return DFB_UNSUPPORTED;
}

DFBResult Construct( IDirectFBVideoProvider *thiz, const char *filename )
{
     IDirectFBVideoProvider_AviFile_data *data;

     data = (IDirectFBVideoProvider_AviFile_data*)
          calloc( 1, sizeof(IDirectFBVideoProvider_AviFile_data) );

     thiz->priv = data;

     data->ref = 1;

     try {
          data->player = CreateAviPlayer( filename, 16 /* FIXME */ );

          data->player->SetDrawCallback2( AviFile_DrawCallback, data );
          data->player->SetKillHandler( AviFile_KillCallback, data );
     }
     catch (FatalError e) {
          ERRORMSG( "DirectFB/AviFile: CreateAviPlayer failed: %s\n",
                    e.GetDesc() );
          free( data );
          return DFB_FAILURE;
     }

     data->source.width  = data->player->GetWidth();
     data->source.height = data->player->GetHeight();
     data->source.format = DSPF_RGB16;

     data->source.front_buffer =
          (SurfaceBuffer*) calloc( 1, sizeof(SurfaceBuffer) );

     data->source.front_buffer->policy = CSP_SYSTEMONLY;
     data->source.front_buffer->system.health = CSH_STORED;

     data->source.back_buffer = data->source.front_buffer;

     pthread_mutex_init( &data->source.front_lock, NULL );
     pthread_mutex_init( &data->source.back_lock, NULL );

     data->source.reactor = reactor_new();

     data->state.source   = &data->source;
     data->state.modified = SMF_ALL;

     thiz->AddRef = IDirectFBVideoProvider_AviFile_AddRef;
     thiz->Release = IDirectFBVideoProvider_AviFile_Release;
     thiz->GetSurfaceDescription =
          IDirectFBVideoProvider_AviFile_GetSurfaceDescription;
     thiz->PlayTo = IDirectFBVideoProvider_AviFile_PlayTo;
     thiz->Stop = IDirectFBVideoProvider_AviFile_Stop;
     thiz->SeekTo = IDirectFBVideoProvider_AviFile_SeekTo;
     thiz->GetPos = IDirectFBVideoProvider_AviFile_GetPos;
     thiz->GetLength = IDirectFBVideoProvider_AviFile_GetLength;

     return DFB_OK;
}

}

<|endoftext|>


<commit_before>#include "core.h"

static bool isNullSessionKey(uint8_t *key) {
	for (int i = 0; i < SESSION_KEY_SIZE; i++)
		if (key[i] != 0)
			return false;

	return true;
}

Client::Client(NetCore *_netCore) : SocketRWCommon(_netCore) {
	authState = AS_LOGIN_WAIT;
	memset(sessionKey, 0, sizeof(sessionKey));
	readBufPosition = 0;

	nextPacketID = 1;
	lastReceivedPacketID = 0;
}
Client::~Client() {
	std::list<Packet *>::iterator
		i = packetCache.begin(),
		  e = packetCache.end();

	for (; i != e; ++i)
		delete *i;
}

void Client::startService(int _sock, bool withTls) {
	close();

	sock = _sock;

	if (!setSocketNonBlocking(sock)) {
		perror("[Client::startService] Could not set non-blocking");
		close();
		return;
	}

#ifdef USE_GNUTLS
	if (withTls) {
		int initRet = gnutls_init(&tls, GNUTLS_SERVER);
		if (initRet != GNUTLS_E_SUCCESS) {
			printf("[Client::startService] gnutls_init borked\n");
			gnutls_perror(initRet);
			close();
			return;
		}

		// TODO: error check this
		int ret;
		const char *errPos;

		ret = gnutls_priority_set_direct(tls, "PERFORMANCE:%SERVER_PRECEDENCE", &errPos);
		if (ret != GNUTLS_E_SUCCESS) {
			printf("gnutls_priority_set_direct failure: %s\n", gnutls_strerror(ret));
			close();
			return;
		}

		ret = gnutls_credentials_set(tls, GNUTLS_CRD_CERTIFICATE, g_clientCreds);
		if (ret != GNUTLS_E_SUCCESS) {
			printf("gnutls_credentials_set failure: %s\n", gnutls_strerror(ret));
			close();
			return;
		}

		gnutls_certificate_server_set_request(tls, GNUTLS_CERT_IGNORE);

		gnutls_transport_set_int(tls, sock);

		tlsActive = true;

		state = CS_TLS_HANDSHAKE;

		printf("[fd=%d] preparing for TLS handshake\n", sock);
	} else
#endif
	{
		state = CS_CONNECTED;
	}
}

void Client::close() {
	SocketRWCommon::close();

	if (authState == AS_AUTHED)
		deadTime = time(NULL) + SESSION_KEEPALIVE;
	else
		deadTime = time(NULL) - 1; // kill instantly
}

void Client::generateSessionKey() {
	time_t now = time(NULL);

	while (true) {
		for (int i = 0; i < SESSION_KEY_SIZE; i++) {
			if (i < sizeof(time_t))
				sessionKey[i] = ((uint8_t*)&now)[i];
			else
				sessionKey[i] = rand() & 255;
		}

		// Is any other client already using this key?
		// It's ridiculously unlikely, but... probably best
		// to check just in case!
		bool foundMatch = false;

		for (int i = 0; i < netCore->clientCount; i++) {
			if (netCore->clients[i] != this) {
				if (!memcmp(netCore->clients[i]->sessionKey, sessionKey, SESSION_KEY_SIZE))
					foundMatch = true;
			}
		}

		// If there's none, we can safely leave!
		if (!foundMatch)
			break;
	}
}

void Client::clearCachedPackets(int maxID) {
	packetCache.remove_if([maxID](Packet *&pkt) {
			return (pkt->id <= maxID);
			});
}

void Client::handlePacket(Packet::Type type, char *data, int size) {
	Buffer pkt;
	pkt.useExistingBuffer(data, size);

	printf("[fd=%d] Packet : type %d, size %d\n", sock, type, size);

	if (authState == AS_LOGIN_WAIT) {
		if (type == Packet::C2B_OOB_LOGIN) {
			int error = 0;

			uint32_t protocolVersion = pkt.readU32();
			if (protocolVersion != PROTOCOL_VERSION)
				error = 1;

			uint32_t lastReceivedByClient = pkt.readU32();

			if (!pkt.readRemains(SESSION_KEY_SIZE))
				error = 2;

			// Authentication goes here at some point, too

			if (error != 0) {
				// Send an error...
				Buffer pkt;
				pkt.writeU32(error);
				sendPacket(Packet::B2C_OOB_LOGIN_FAILED, pkt, /*allowUnauthed=*/true);

				// Would close() now but this means the login failed packet never gets sent
				// need to figure out a fix for this. TODO FIXME etc etc.

			} else {
				// or log us in!
				uint8_t reqKey[SESSION_KEY_SIZE];
				pkt.read((char *)reqKey, SESSION_KEY_SIZE);

				printf("[fd=%d] Client authenticating\n", sock);

				if (!isNullSessionKey(reqKey)) {
					printf("[fd=%d] Trying to resume session...", sock);
					printf("(last they received = %d)\n", lastReceivedByClient);

					Client *other = netCore->findClientWithSessionKey(reqKey);
					printf("[fd=%d] Got client %p\n", sock, other);

					if (other && other->authState == AS_AUTHED) {
						printf("Valid: last packet we sent = %d\n", other->nextPacketID - 1);
						// Yep, we can go!
						other->resumeSession(this, lastReceivedByClient);
						return;
					}
				}

				// If we got here, it means we couldn't resume the session.
				// Start over.
				printf("[fd=%d] Creating new session\n", sock);

				generateSessionKey();
				authState = AS_AUTHED;

				Buffer pkt;
				pkt.append((char *)sessionKey, SESSION_KEY_SIZE);
				sendPacket(Packet::B2C_OOB_LOGIN_SUCCESS, pkt);

				sessionStartEvent();
			}

		} else {
			printf("[fd=%d] Unrecognised packet in AS_LOGIN_WAIT authstate: type %d, size %d\n",
					sock, type, size);
		}
	} else if (authState == AS_AUTHED) {
		packetReceivedEvent(type, pkt);
	}
}

void Client::processReadBuffer() {
	// Try to process as many packets as we have in inputBuf

	// Basic header is 8 bytes
	// Extended (non-OOB) header is 16 bytes
	inputBuf.readSeek(0);
	readBufPosition = 0;

	while (inputBuf.readRemains(8)) {
		// We have 8 bytes, so we can try to read a basic header
		Packet::Type type = (Packet::Type)inputBuf.readU16();
		int reserved = inputBuf.readU16();
		uint32_t packetSize = inputBuf.readU32();

		// Do we now have the whole packet in memory...?
		int extHeaderSize = (type & Packet::T_OUT_OF_BAND_FLAG) ? 0 : 8;

		if (!inputBuf.readRemains(packetSize + extHeaderSize))
			break;

		if (!(type & Packet::T_OUT_OF_BAND_FLAG)) {
			// Handle packet system things for non-OOB packets
			uint32_t packetID = inputBuf.readU32();
			uint32_t lastReceivedByClient = inputBuf.readU32();

			lastReceivedPacketID = packetID;
			clearCachedPackets(lastReceivedByClient);
		}

		// Yep, we can process it!

		// Save the position of the next packet
		readBufPosition = inputBuf.readTell() + packetSize;
		handlePacket(type, &inputBuf.data()[inputBuf.readTell()], packetSize);

		inputBuf.readSeek(readBufPosition);
	}

	// If we managed to handle anything, lop it off the buffer
	inputBuf.trimFromStart(readBufPosition);
	readBufPosition = 0;
}

void Client::resumeSession(Client *other, int lastReceivedByClient) {
	close();

	inputBuf.clear();
	inputBuf.append(
			&other->inputBuf.data()[other->readBufPosition],
			other->inputBuf.size() - other->readBufPosition);

	// Not sure if we need to copy the outputbuf but it can't hurt
	outputBuf.clear();
	outputBuf.append(other->outputBuf.data(), other->outputBuf.size());

	sock = other->sock;
	state = other->state;
#ifdef USE_GNUTLS
	tls = other->tls;
	tlsActive = other->tlsActive;
#endif

	other->sock = -1;
	other->state = CS_DISCONNECTED;
#ifdef USE_GNUTLS
	other->tls = 0;
	other->tlsActive = false;
#endif

	other->close();

	// Now send them everything we've got!
	Buffer pkt;
	pkt.writeU32(lastReceivedPacketID);
	sendPacket(Packet::B2C_OOB_SESSION_RESUMED, pkt);

	clearCachedPackets(lastReceivedByClient);

	std::list<Packet*>::iterator
		i = packetCache.begin(),
		  e = packetCache.end();

	for (; i != e; ++i)
		sendPacketOverWire(*i);
}

void Client::sendPacket(Packet::Type type, const Buffer &data, bool allowUnauthed) {
	Packet *packet = new Packet;
	packet->type = type;
	packet->data.append(data);

	if (type & Packet::T_OUT_OF_BAND_FLAG) {
		packet->id = 0;
	} else {
		packet->id = nextPacketID;
		nextPacketID++;
	}

	if (state == CS_CONNECTED)
		if (authState == AS_AUTHED || allowUnauthed)
			sendPacketOverWire(packet);

	if (type & Packet::T_OUT_OF_BAND_FLAG)
		delete packet;
	else
		packetCache.push_back(packet);
}

void Client::sendPacketOverWire(const Packet *packet) {
	Buffer header;
	header.writeU16(packet->type);
	header.writeU16(0);
	header.writeU32(packet->data.size());

	if (!(packet->type & Packet::T_OUT_OF_BAND_FLAG)) {
		header.writeU32(packet->id);
		header.writeU32(lastReceivedPacketID);
	}

	outputBuf.append(header);
	outputBuf.append(packet->data);
}
<commit_msg>bouncer: silently ignore packets from the client that we've already seen<commit_after>#include "core.h"

static bool isNullSessionKey(uint8_t *key) {
	for (int i = 0; i < SESSION_KEY_SIZE; i++)
		if (key[i] != 0)
			return false;

	return true;
}

Client::Client(NetCore *_netCore) : SocketRWCommon(_netCore) {
	authState = AS_LOGIN_WAIT;
	memset(sessionKey, 0, sizeof(sessionKey));
	readBufPosition = 0;

	nextPacketID = 1;
	lastReceivedPacketID = 0;
}
Client::~Client() {
	std::list<Packet *>::iterator
		i = packetCache.begin(),
		  e = packetCache.end();

	for (; i != e; ++i)
		delete *i;
}

void Client::startService(int _sock, bool withTls) {
	close();

	sock = _sock;

	if (!setSocketNonBlocking(sock)) {
		perror("[Client::startService] Could not set non-blocking");
		close();
		return;
	}

#ifdef USE_GNUTLS
	if (withTls) {
		int initRet = gnutls_init(&tls, GNUTLS_SERVER);
		if (initRet != GNUTLS_E_SUCCESS) {
			printf("[Client::startService] gnutls_init borked\n");
			gnutls_perror(initRet);
			close();
			return;
		}

		// TODO: error check this
		int ret;
		const char *errPos;

		ret = gnutls_priority_set_direct(tls, "PERFORMANCE:%SERVER_PRECEDENCE", &errPos);
		if (ret != GNUTLS_E_SUCCESS) {
			printf("gnutls_priority_set_direct failure: %s\n", gnutls_strerror(ret));
			close();
			return;
		}

		ret = gnutls_credentials_set(tls, GNUTLS_CRD_CERTIFICATE, g_clientCreds);
		if (ret != GNUTLS_E_SUCCESS) {
			printf("gnutls_credentials_set failure: %s\n", gnutls_strerror(ret));
			close();
			return;
		}

		gnutls_certificate_server_set_request(tls, GNUTLS_CERT_IGNORE);

		gnutls_transport_set_int(tls, sock);

		tlsActive = true;

		state = CS_TLS_HANDSHAKE;

		printf("[fd=%d] preparing for TLS handshake\n", sock);
	} else
#endif
	{
		state = CS_CONNECTED;
	}
}

void Client::close() {
	SocketRWCommon::close();

	if (authState == AS_AUTHED)
		deadTime = time(NULL) + SESSION_KEEPALIVE;
	else
		deadTime = time(NULL) - 1; // kill instantly
}

void Client::generateSessionKey() {
	time_t now = time(NULL);

	while (true) {
		for (int i = 0; i < SESSION_KEY_SIZE; i++) {
			if (i < sizeof(time_t))
				sessionKey[i] = ((uint8_t*)&now)[i];
			else
				sessionKey[i] = rand() & 255;
		}

		// Is any other client already using this key?
		// It's ridiculously unlikely, but... probably best
		// to check just in case!
		bool foundMatch = false;

		for (int i = 0; i < netCore->clientCount; i++) {
			if (netCore->clients[i] != this) {
				if (!memcmp(netCore->clients[i]->sessionKey, sessionKey, SESSION_KEY_SIZE))
					foundMatch = true;
			}
		}

		// If there's none, we can safely leave!
		if (!foundMatch)
			break;
	}
}

void Client::clearCachedPackets(int maxID) {
	packetCache.remove_if([maxID](Packet *&pkt) {
			return (pkt->id <= maxID);
			});
}

void Client::handlePacket(Packet::Type type, char *data, int size) {
	Buffer pkt;
	pkt.useExistingBuffer(data, size);

	printf("[fd=%d] Packet : type %d, size %d\n", sock, type, size);

	if (authState == AS_LOGIN_WAIT) {
		if (type == Packet::C2B_OOB_LOGIN) {
			int error = 0;

			uint32_t protocolVersion = pkt.readU32();
			if (protocolVersion != PROTOCOL_VERSION)
				error = 1;

			uint32_t lastReceivedByClient = pkt.readU32();

			if (!pkt.readRemains(SESSION_KEY_SIZE))
				error = 2;

			// Authentication goes here at some point, too

			if (error != 0) {
				// Send an error...
				Buffer pkt;
				pkt.writeU32(error);
				sendPacket(Packet::B2C_OOB_LOGIN_FAILED, pkt, /*allowUnauthed=*/true);

				// Would close() now but this means the login failed packet never gets sent
				// need to figure out a fix for this. TODO FIXME etc etc.

			} else {
				// or log us in!
				uint8_t reqKey[SESSION_KEY_SIZE];
				pkt.read((char *)reqKey, SESSION_KEY_SIZE);

				printf("[fd=%d] Client authenticating\n", sock);

				if (!isNullSessionKey(reqKey)) {
					printf("[fd=%d] Trying to resume session...", sock);
					printf("(last they received = %d)\n", lastReceivedByClient);

					Client *other = netCore->findClientWithSessionKey(reqKey);
					printf("[fd=%d] Got client %p\n", sock, other);

					if (other && other->authState == AS_AUTHED) {
						printf("Valid: last packet we sent = %d\n", other->nextPacketID - 1);
						// Yep, we can go!
						other->resumeSession(this, lastReceivedByClient);
						return;
					}
				}

				// If we got here, it means we couldn't resume the session.
				// Start over.
				printf("[fd=%d] Creating new session\n", sock);

				generateSessionKey();
				authState = AS_AUTHED;

				Buffer pkt;
				pkt.append((char *)sessionKey, SESSION_KEY_SIZE);
				sendPacket(Packet::B2C_OOB_LOGIN_SUCCESS, pkt);

				sessionStartEvent();
			}

		} else {
			printf("[fd=%d] Unrecognised packet in AS_LOGIN_WAIT authstate: type %d, size %d\n",
					sock, type, size);
		}
	} else if (authState == AS_AUTHED) {
		packetReceivedEvent(type, pkt);
	}
}

void Client::processReadBuffer() {
	// Try to process as many packets as we have in inputBuf

	// Basic header is 8 bytes
	// Extended (non-OOB) header is 16 bytes
	inputBuf.readSeek(0);
	readBufPosition = 0;

	while (inputBuf.readRemains(8)) {
		// We have 8 bytes, so we can try to read a basic header
		Packet::Type type = (Packet::Type)inputBuf.readU16();
		int reserved = inputBuf.readU16();
		uint32_t packetSize = inputBuf.readU32();
		bool silentlyIgnore = false;

		// Do we now have the whole packet in memory...?
		int extHeaderSize = (type & Packet::T_OUT_OF_BAND_FLAG) ? 0 : 8;

		if (!inputBuf.readRemains(packetSize + extHeaderSize))
			break;

		if (!(type & Packet::T_OUT_OF_BAND_FLAG)) {
			// Handle packet system things for non-OOB packets
			uint32_t packetID = inputBuf.readU32();
			uint32_t lastReceivedByClient = inputBuf.readU32();

			if (packetID > lastReceivedPacketID) {
				// This is a new packet
				lastReceivedPacketID = packetID;
			} else {
				// We've already seen this packet, silently ignore it!
				silentlyIgnore = true;
			}

			clearCachedPackets(lastReceivedByClient);
		}

		// Yep, we can process it!

		// Save the position of the next packet
		readBufPosition = inputBuf.readTell() + packetSize;

		if (!silentlyIgnore)
			handlePacket(type, &inputBuf.data()[inputBuf.readTell()], packetSize);

		inputBuf.readSeek(readBufPosition);
	}

	// If we managed to handle anything, lop it off the buffer
	inputBuf.trimFromStart(readBufPosition);
	readBufPosition = 0;
}

void Client::resumeSession(Client *other, int lastReceivedByClient) {
	close();

	inputBuf.clear();
	inputBuf.append(
			&other->inputBuf.data()[other->readBufPosition],
			other->inputBuf.size() - other->readBufPosition);

	// Not sure if we need to copy the outputbuf but it can't hurt
	outputBuf.clear();
	outputBuf.append(other->outputBuf.data(), other->outputBuf.size());

	sock = other->sock;
	state = other->state;
#ifdef USE_GNUTLS
	tls = other->tls;
	tlsActive = other->tlsActive;
#endif

	other->sock = -1;
	other->state = CS_DISCONNECTED;
#ifdef USE_GNUTLS
	other->tls = 0;
	other->tlsActive = false;
#endif

	other->close();

	// Now send them everything we've got!
	Buffer pkt;
	pkt.writeU32(lastReceivedPacketID);
	sendPacket(Packet::B2C_OOB_SESSION_RESUMED, pkt);

	clearCachedPackets(lastReceivedByClient);

	std::list<Packet*>::iterator
		i = packetCache.begin(),
		  e = packetCache.end();

	for (; i != e; ++i)
		sendPacketOverWire(*i);
}

void Client::sendPacket(Packet::Type type, const Buffer &data, bool allowUnauthed) {
	Packet *packet = new Packet;
	packet->type = type;
	packet->data.append(data);

	if (type & Packet::T_OUT_OF_BAND_FLAG) {
		packet->id = 0;
	} else {
		packet->id = nextPacketID;
		nextPacketID++;
	}

	if (state == CS_CONNECTED)
		if (authState == AS_AUTHED || allowUnauthed)
			sendPacketOverWire(packet);

	if (type & Packet::T_OUT_OF_BAND_FLAG)
		delete packet;
	else
		packetCache.push_back(packet);
}

void Client::sendPacketOverWire(const Packet *packet) {
	Buffer header;
	header.writeU16(packet->type);
	header.writeU16(0);
	header.writeU32(packet->data.size());

	if (!(packet->type & Packet::T_OUT_OF_BAND_FLAG)) {
		header.writeU32(packet->id);
		header.writeU32(lastReceivedPacketID);
	}

	outputBuf.append(header);
	outputBuf.append(packet->data);
}
<|endoftext|>


<commit_before>#include <cmath>
#include <vector>
#include <fstream>
#include <thread>
#include <unordered_map>
#include <ctime>

#include "tasks.hpp"
#include "sequence.h"
#include "io.h"
#include "alignment.h"
#include "algorithms.h"

// This utilitu function returns the bytes remaining until the end of the file
// is reached. It should be probably made publicly availale in a separate util
// section of the library (e.g. util/io.h)
size_t bytesRemainingToTheEnd(ifstream& ifs) {
  std::streampos tmp = ifs.tellg();
  ifs.seekg(0, ifs.end);
  size_t totalLength = ifs.tellg();
  ifs.seekg(tmp);
  return (totalLength - ifs.tellg());
}

void alignSmithWaterman(std::vector<Read>* reads, const Reference* ref,
			std::vector<ScoredPosition<int,int> >* aligns, int indexOffset) {
  int nReads = reads->size();
  string refBases((char*)ref->getSequence());
  std::cout << "Aligning " << nReads << " reads" << std::endl;
  for (int i = 0; i < nReads; ++i) {
    SmithWatermanDP sw((*reads)[i].getBases(), refBases);
    sw.computeMatrix();
    MatrixPoint2D maxP = sw.getGlobalBest();
    aligns->push_back(ScoredPosition<int, int>((i + indexOffset),maxP.j, sw.getScoreAt(maxP)));
  }
}

std::vector<ScoredPosition<int, int> > alignFastqReadsSimpleSW(const string& readsPath, const string& referencePath,
							       std::ostream& output, uint64_t nThreads, size_t nReads) {

  std::cout << "-- Smith Waterman alignment --" << std::endl;
  std::cout.flush();

  std::vector<ScoredPosition<int,int> > aligns;
						   // some constants
  uint64_t T = (nThreads > 0) ? nThreads : 1;
  uint64_t M = (nReads > 0) ? nReads : 0;

  std::cout << " ************* " << T << "  " <<  M << "**************"  << std::endl;
  std::cout.flush();

  // open files
  std::ifstream readsIn(readsPath, std::ifstream::in);
  //  std::ifstream refIn(referencePath, std::ifstream::in);

  // list of all threads (use later for joining)
  std::vector<std::thread> threads;
  std::vector<std::vector<ScoredPosition<int,int> >*> threadAligns;

  // read input reference...
  std::cout << "    Loading reference..." << std::endl;
  std::cout.flush();
  FastFormat fast;
  fast.loadFromFile(referencePath);
  Reference ref = (Reference)fast;

  // ...and reads
  std::cout << "    Loading reads..." << std::endl;
  std::cout.flush();
  if (nReads >= 0) {
    // CASE 1: A specific total number of reads has been specified or
    uint64_t readsPerThread = (size_t)ceil( ((double)(M)) / ((double)(T)) );
    uint64_t actualReads = 0;
    // repeat for each thread
    for (uint64_t t = 0; t < T; ++t) {
      std::vector<Read>* reads = new std::vector<Read>();
      actualReads = 0;
      while( (!readsIn.eof()) &&  (actualReads < readsPerThread) ) {
	FastqRead r;
	readsIn >> r;
	reads->push_back(r);
      	actualReads++;
      }
      // create the aligns vector for the next starting thread
      std::vector<ScoredPosition<int,int> >* alignsVector = new std::vector<ScoredPosition< int,int > >();
      threadAligns.push_back(alignsVector);
      threads.push_back(std::thread(alignSmithWaterman, reads, &ref, alignsVector, t * readsPerThread));
    }

  } else {
    // CASE 2: we need to estimate the amount of reads to be assigned to each thread
    double fraction = 0;
    double D = 1.0 / ((double)T);
    // repeat for each thread
    for (uint64_t t = 0; t < T; ++t) {
      while( (!readsIn.eof()) && (fraction < D)) {
      	fraction += 0.1;
      }
      //      threads.push_back(std::thread(alignSmithWaterman,reads,&ref));
    }

  }

  // synchronize all threads
  for (auto& thr : threads) {
    thr.join();
  }
  aligns.clear();
  for (std::vector<ScoredPosition<int,int> >* v : threadAligns) {
    std::cout << "++++ " << v->size() <<  std::endl;
    aligns.insert(aligns.end(), v->begin(), v->end());
    delete v;
  }

  output << "*** Alignments: " << std::endl;
  for (ScoredPosition< int,int > p : aligns) {
    output << p.getSequenceId() << "\t" << p.getPosition() << " (" << p.getScore() << ")" << std::endl;
  }

  return aligns;
}

/**************************** K-SPECTRUM FUNCTIONS ****************************/
void taskComputeKSpectrum(size_t k, const string& referenceFile) {
  FastFormat fast;
  fast.loadFromFile(referenceFile);
  Reference ref = (Reference) fast;
  std::unordered_map< uint64_t, uint64_t > index = spectrumAsIntMap(ref, k);
  for (std::pair< uint64_t, uint64_t > p : index) {
    std::cout << NumericKMer(p.first, k) << " " << p.second << std::endl;
  }
}

void taskMapReadsKmers(const string& reference, const string& reads, size_t k, const string& out) {
  std::cout << "-------------------- Reads Mapping --------------------" << std::endl;
  // open files
  std::ofstream outFileStream;
  if (!out.empty()) {
    outFileStream.open(out, std::ios::out);
  }
  std::ostream& outStream = (out.empty()) ? std::cout : outFileStream;
  FastFormat refFast;
  std::cout << "Loading reference (" << reference << ")...";
  std::cout.flush();
  refFast.loadFromFile(reference);
  std::cout << " Ok!" << std::endl;
  Reference ref = refFast.toReference();

  std::ifstream readsStream(reads, std::ios::in);

  // compute index for the reference
  std::cout << "Index creation...";
  std::cout.flush();
  NumericKmerIndex index = kmersMapping(ref, k);
  std::cout << " Ok!" << std::endl;
  // scan reads and find mappings
  size_t read_index = 0;
  size_t kmer_index = 0;
  std::cout << "Reads mapping...";
  std::cout.flush();
  while(!readsStream.eof()) {
    FastqRead r;
    readsStream >> r;
    list< KMer > kmers = r.getKMerList(k);
    kmer_index = 0;
    for (KMer kmer : kmers) {
      NumericKMer nkmer(kmer);
      std::unordered_map< uint64_t, std::list< size_t > >::const_iterator it = index.find((uint64_t)nkmer);
      outStream << read_index << ":" << kmer_index << " "  << ((it == index.end()) ? (int64_t)(-1) :  (int64_t)it->second.front() ) << std::endl;
      kmer_index++;
    }
    read_index++;
  }
  std::cout << "Ok!" << std::endl << "-------------------- DONE --------------------" << std::endl;
}

// Computes for each reads of the input set the kmerscore and the numero fo errors
// (see algorithms/kmerscore) against the reference sequence
void taskKmerScoreReads(const string& reference, const string& reads, size_t k, const string& out) {

  // define output (either a file or the standard out)
  std::ofstream outFileStream;
  if (!out.empty()) {
    outFileStream.open(out, std::ios::out);
  }
  std::ostream& outStream = (out.empty()) ? std::cout : outFileStream;
  time_t beginTime, endTime;

  std::cout << "-------------------- Reads Mapping --------------------" << std::endl;
  // open files
  std::cout << "Loading reference (" << reference << ")...";
  std::cout.flush(); // in case of stuck code we see at which point
  FastFormat refFast(reference);
  Reference ref = refFast.toReference();
  std::cout << " Done (" << ref.getSequenceLength() << " bases)" << std::endl;

  // open a stream for reading reads file
  std::ifstream readsStream(reads, std::ios::in);

  // compute the index for the reference
  std::cout << "Index creation...";
  std::cout.flush();
  time(&beginTime);
  NumericKmerIndex index = kmersMapping(ref, k);
  time(&endTime);
  std::cout << " Ok! (" << difftime(endTime, beginTime) << " sec)" << std::endl;
}
<commit_msg>little advance on the 'kscore' task<commit_after>#include <cmath>
#include <vector>
#include <fstream>
#include <thread>
#include <unordered_map>
#include <ctime>

#include "tasks.hpp"
#include "sequence.h"
#include "io.h"
#include "alignment.h"
#include "algorithms.h"

// This utilitu function returns the bytes remaining until the end of the file
// is reached. It should be probably made publicly availale in a separate util
// section of the library (e.g. util/io.h)
size_t bytesRemainingToTheEnd(ifstream& ifs) {
  std::streampos tmp = ifs.tellg();
  ifs.seekg(0, ifs.end);
  size_t totalLength = ifs.tellg();
  ifs.seekg(tmp);
  return (totalLength - ifs.tellg());
}

void alignSmithWaterman(std::vector<Read>* reads, const Reference* ref,
			std::vector<ScoredPosition<int,int> >* aligns, int indexOffset) {
  int nReads = reads->size();
  string refBases((char*)ref->getSequence());
  std::cout << "Aligning " << nReads << " reads" << std::endl;
  for (int i = 0; i < nReads; ++i) {
    SmithWatermanDP sw((*reads)[i].getBases(), refBases);
    sw.computeMatrix();
    MatrixPoint2D maxP = sw.getGlobalBest();
    aligns->push_back(ScoredPosition<int, int>((i + indexOffset),maxP.j, sw.getScoreAt(maxP)));
  }
}

std::vector<ScoredPosition<int, int> > alignFastqReadsSimpleSW(const string& readsPath, const string& referencePath,
							       std::ostream& output, uint64_t nThreads, size_t nReads) {

  std::cout << "-- Smith Waterman alignment --" << std::endl;
  std::cout.flush();

  std::vector<ScoredPosition<int,int> > aligns;
						   // some constants
  uint64_t T = (nThreads > 0) ? nThreads : 1;
  uint64_t M = (nReads > 0) ? nReads : 0;

  std::cout << " ************* " << T << "  " <<  M << "**************"  << std::endl;
  std::cout.flush();

  // open files
  std::ifstream readsIn(readsPath, std::ifstream::in);
  //  std::ifstream refIn(referencePath, std::ifstream::in);

  // list of all threads (use later for joining)
  std::vector<std::thread> threads;
  std::vector<std::vector<ScoredPosition<int,int> >*> threadAligns;

  // read input reference...
  std::cout << "    Loading reference..." << std::endl;
  std::cout.flush();
  FastFormat fast;
  fast.loadFromFile(referencePath);
  Reference ref = (Reference)fast;

  // ...and reads
  std::cout << "    Loading reads..." << std::endl;
  std::cout.flush();
  if (nReads >= 0) {
    // CASE 1: A specific total number of reads has been specified or
    uint64_t readsPerThread = (size_t)ceil( ((double)(M)) / ((double)(T)) );
    uint64_t actualReads = 0;
    // repeat for each thread
    for (uint64_t t = 0; t < T; ++t) {
      std::vector<Read>* reads = new std::vector<Read>();
      actualReads = 0;
      while( (!readsIn.eof()) &&  (actualReads < readsPerThread) ) {
	FastqRead r;
	readsIn >> r;
	reads->push_back(r);
      	actualReads++;
      }
      // create the aligns vector for the next starting thread
      std::vector<ScoredPosition<int,int> >* alignsVector = new std::vector<ScoredPosition< int,int > >();
      threadAligns.push_back(alignsVector);
      threads.push_back(std::thread(alignSmithWaterman, reads, &ref, alignsVector, t * readsPerThread));
    }

  } else {
    // CASE 2: we need to estimate the amount of reads to be assigned to each thread
    double fraction = 0;
    double D = 1.0 / ((double)T);
    // repeat for each thread
    for (uint64_t t = 0; t < T; ++t) {
      while( (!readsIn.eof()) && (fraction < D)) {
      	fraction += 0.1;
      }
      //      threads.push_back(std::thread(alignSmithWaterman,reads,&ref));
    }

  }

  // synchronize all threads
  for (auto& thr : threads) {
    thr.join();
  }
  aligns.clear();
  for (std::vector<ScoredPosition<int,int> >* v : threadAligns) {
    std::cout << "++++ " << v->size() <<  std::endl;
    aligns.insert(aligns.end(), v->begin(), v->end());
    delete v;
  }

  output << "*** Alignments: " << std::endl;
  for (ScoredPosition< int,int > p : aligns) {
    output << p.getSequenceId() << "\t" << p.getPosition() << " (" << p.getScore() << ")" << std::endl;
  }

  return aligns;
}

/**************************** K-SPECTRUM FUNCTIONS ****************************/
void taskComputeKSpectrum(size_t k, const string& referenceFile) {
  FastFormat fast;
  fast.loadFromFile(referenceFile);
  Reference ref = (Reference) fast;
  std::unordered_map< uint64_t, uint64_t > index = spectrumAsIntMap(ref, k);
  for (std::pair< uint64_t, uint64_t > p : index) {
    std::cout << NumericKMer(p.first, k) << " " << p.second << std::endl;
  }
}

void taskMapReadsKmers(const string& reference, const string& reads, size_t k, const string& out) {
  std::cout << "-------------------- Reads Mapping --------------------" << std::endl;
  // open files
  std::ofstream outFileStream;
  if (!out.empty()) {
    outFileStream.open(out, std::ios::out);
  }
  std::ostream& outStream = (out.empty()) ? std::cout : outFileStream;
  FastFormat refFast;
  std::cout << "Loading reference (" << reference << ")...";
  std::cout.flush();
  refFast.loadFromFile(reference);
  std::cout << " Ok!" << std::endl;
  Reference ref = refFast.toReference();

  std::ifstream readsStream(reads, std::ios::in);

  // compute index for the reference
  std::cout << "Index creation...";
  std::cout.flush();
  NumericKmerIndex index = kmersMapping(ref, k);
  std::cout << " Ok!" << std::endl;
  // scan reads and find mappings
  size_t read_index = 0;
  size_t kmer_index = 0;
  std::cout << "Reads mapping...";
  std::cout.flush();
  while(!readsStream.eof()) {
    FastqRead r;
    readsStream >> r;
    list< KMer > kmers = r.getKMerList(k);
    kmer_index = 0;
    for (KMer kmer : kmers) {
      NumericKMer nkmer(kmer);
      std::unordered_map< uint64_t, std::list< size_t > >::const_iterator it = index.find((uint64_t)nkmer);
      outStream << read_index << ":" << kmer_index << " "  << ((it == index.end()) ? (int64_t)(-1) :  (int64_t)it->second.front() ) << std::endl;
      kmer_index++;
    }
    read_index++;
  }
  std::cout << "Ok!" << std::endl << "-------------------- DONE --------------------" << std::endl;
}

// Computes for each reads of the input set the kmerscore and the numero fo errors
// (see algorithms/kmerscore) against the reference sequence
void taskKmerScoreReads(const string& reference, const string& reads, size_t k, const string& out) {

  // define output (either a file or the standard out)
  std::ofstream outFileStream;
  if (!out.empty()) {
    outFileStream.open(out, std::ios::out);
  }
  std::ostream& outStream = (out.empty()) ? std::cout : outFileStream;
  time_t beginTime, endTime;

  std::cout << "-------------------- Reads Mapping --------------------" << std::endl;
  // open files
  std::cout << "Loading reference (" << reference << ")...";
  std::cout.flush(); // in case of stuck code we see at which point
  FastFormat refFast(reference);
  Reference ref = refFast.toReference();
  std::cout << " Done (" << ref.getSequenceLength() << " bases)" << std::endl;

  // open a stream for reading reads file
  std::ifstream readsStream(reads, std::ios::in);

  // compute the index for the reference
  std::cout << "Index creation...";
  std::cout.flush();
  time(&beginTime);
  NumericKmerIndex index = kmersMapping(ref, k);
  time(&endTime);
  std::cout << " Ok! (" << difftime(endTime, beginTime) << " sec)" << std::endl;
  std::cout << "Calculating scores...";
  std::cout.flush();
  size_t M = 0;
  time(&beginTime);
  while(!readsStream.eof()) {
    FastqRead r;
    readsStream >> r;
    if (r.getSequenceLength() == 0) {
      continue;
    }
    //std::cout << r <<  "  (" << M << ")" << std::endl; std::cout.flush(); // just for debug
    KmersMap map = extractKmersMapPosition(r, index, k);
    M++;
  }
  readsStream.close();
  time(&endTime);
  double elapsed = difftime(endTime, beginTime);
  double rate = M / elapsed;
  std::cout << " Done (" << M << " reads in " << elapsed  << " sec, " << rate << " reads/sec)"  << std::endl;
}
<|endoftext|>


<commit_before>#include "hasher.hxx"

using namespace std;
using namespace Py;

SequenceHasher::SequenceHasher(){

    _ref = NULL;
    _initialized = false;

}

SequenceHasher::~SequenceHasher(){
    if( _ref ){
        delete _ref;
    }

    list<ObservationSet*>::iterator obs_itr;

    for( obs_itr = _obs.begin(); obs_itr != _obs.end(); obs_itr++ ){
        delete (*obs_itr);
    }
    _obs.clear();
}

priority_queue<pair<double, string> > SequenceHasher::hash(string sequence){

    if( !_initialized ){
        initialize();
    }

    ObservationSet os(sequence);
    list<ObservationSet*>::iterator l_itr;

    priority_queue<pair<double, string> > rpq;
    for(l_itr = _obs.begin(); l_itr != _obs.end(); l_itr++){

        double likelihood = (*l_itr)->likelihood(&os, _ref);
        string nm = (*l_itr)->getName();
        rpq.push(pair<double, string>(likelihood, nm));
    }

    return rpq;
}

void SequenceHasher::addReference(string name, string sequence){

    //TODO Ensure that arguments are in a consistent order
    ObservationSet* ob = new ObservationSet(sequence, name);
    _obs.push_back(ob);

    _initialized = false;
    return;
}

void SequenceHasher::initialize(){
    _initialized = true;
    if( !_ref ){ delete _ref; }

    _ref = new ReferenceSet(_obs);
    list<ObservationSet*>::iterator obs_itr;

    return;
}

Object SequenceHasher::py_addReference( const Tuple &args ){
    args.verify_length(2);
    String name = args[0];
    String sequence = args[1];

    addReference(name, sequence);

    return args[0];
}

Object SequenceHasher::py_hash( const Tuple &args ){
    args.verify_length(1);
    String sequence = args[0];
    priority_queue<pair<double, string> > results = hash(sequence);

    Dict d;

    while( !results.empty() ){
        double likelihood = results.top().first;
        string name = results.top().second;
        d[String(name)] = Float(floor(likelihood * 10000));
        results.pop();
    }

    return d;
}

Object SequenceHasher::getattr( const char *name ){
    return getattr_methods(name);
}

Object SequenceHasher::repr(){
    return Py::String("A Hashed based Naive Bayes Classifier for sequence classification");
}

void SequenceHasher::init_type(){
    behaviors().name("SequenceHasher");
    behaviors().doc("SequenceHasher objects: nil");
    behaviors().supportGetattr();
    behaviors().supportRepr();

    add_varargs_method("addReference",  &SequenceHasher::py_addReference, "addReference(name, sequence)");
    add_varargs_method("hash",          &SequenceHasher::py_hash,          "hash(sequence)");
    add_varargs_method("reference_count", &SequenceHasher::reference_count);
}

class hasher_module : public Py::ExtensionModule<hasher_module>
{
public:
    hasher_module()
    : Py::ExtensionModule<hasher_module>("hasher")
    {
        SequenceHasher::init_type();
        add_varargs_method("SequenceHasher", &hasher_module::new_hasher);
        initialize( "docs" );
    }

    virtual ~hasher_module()
    {}

private:
    Object new_hasher(const Py::Tuple& args){
        return asObject(new SequenceHasher());
    }
};

extern "C" void inithasher()
{
#if defined(PY_WIN32_DELAYLOAD_PYTHON_DLL)
    Py::InitialisePythonIndirectPy::Interface();
#endif
    static hasher_module* hasher = new hasher_module;
}
<commit_msg>We have a Naive Bayes Classifier for sequence analysis. It works with python. It produces correct results.<commit_after>#include "hasher.hxx"

using namespace std;
using namespace Py;

SequenceHasher::SequenceHasher(){

    _ref = NULL;
    _initialized = false;

}

SequenceHasher::~SequenceHasher(){
    if( _ref ){
        delete _ref;
    }

    list<ObservationSet*>::iterator obs_itr;

    for( obs_itr = _obs.begin(); obs_itr != _obs.end(); obs_itr++ ){
        delete (*obs_itr);
    }
    _obs.clear();
}

priority_queue<pair<double, string> > SequenceHasher::hash(string sequence){

    if( !_initialized ){
        initialize();
    }

    ObservationSet os(sequence);
    list<ObservationSet*>::iterator l_itr;

    priority_queue<pair<double, string> > rpq;
    for(l_itr = _obs.begin(); l_itr != _obs.end(); l_itr++){

        double likelihood = (*l_itr)->likelihood(&os, _ref);
        string nm = (*l_itr)->getName();
        rpq.push(pair<double, string>(likelihood, nm));
    }

    return rpq;
}

void SequenceHasher::addReference(string name, string sequence){

    //TODO Ensure that arguments are in a consistent order
    ObservationSet* ob = new ObservationSet(sequence, name);
    _obs.push_back(ob);

    _initialized = false;
    return;
}

void SequenceHasher::initialize(){
    _initialized = true;
    if( !_ref ){ delete _ref; }

    _ref = new ReferenceSet(_obs);
    list<ObservationSet*>::iterator obs_itr;

    return;
}

Object SequenceHasher::py_addReference( const Tuple &args ){
    args.verify_length(2);
    String name = args[0];
    String sequence = args[1];

    addReference(name, sequence);

    return args[0];
}

Object SequenceHasher::py_hash( const Tuple &args ){
    args.verify_length(1);
    String sequence = args[0];
    priority_queue<pair<double, string> > results = hash(sequence);

    List d;

    while( !results.empty() ){
        string name = results.top().second;
        d.append(String(name));
        results.pop();
    }

    return d;
}

Object SequenceHasher::getattr( const char *name ){
    return getattr_methods(name);
}

Object SequenceHasher::repr(){
    return Py::String("A Hashed based Naive Bayes Classifier for sequence classification");
}

void SequenceHasher::init_type(){
    behaviors().name("SequenceHasher");
    behaviors().doc("SequenceHasher objects: nil");
    behaviors().supportGetattr();
    behaviors().supportRepr();

    add_varargs_method("addReference",  &SequenceHasher::py_addReference, "addReference(name, sequence)");
    add_varargs_method("hash",          &SequenceHasher::py_hash,          "hash(sequence)");
    add_varargs_method("reference_count", &SequenceHasher::reference_count);
}

class hasher_module : public Py::ExtensionModule<hasher_module>
{
public:
    hasher_module()
    : Py::ExtensionModule<hasher_module>("hasher")
    {
        SequenceHasher::init_type();
        add_varargs_method("SequenceHasher", &hasher_module::new_hasher);
        initialize( "docs" );
    }

    virtual ~hasher_module()
    {}

private:
    Object new_hasher(const Py::Tuple& args){
        return asObject(new SequenceHasher());
    }
};

extern "C" void inithasher()
{
#if defined(PY_WIN32_DELAYLOAD_PYTHON_DLL)
    Py::InitialisePythonIndirectPy::Interface();
#endif
    static hasher_module* hasher = new hasher_module;
}
<|endoftext|>


<commit_before>/*
 * Copyright 2007-2017 Content Management AG
 * All rights reserved.
 *
 * author: Max Kellermann <mk@cm4all.com>
 *
 * 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.
 *
 * 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
 * FOUNDATION 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.
 */

#include "ForwardHttpRequest.hxx"
#include "HttpConnection.hxx"
#include "Cluster.hxx"
#include "ClusterConfig.hxx"
#include "ListenerConfig.hxx"
#include "Instance.hxx"
#include "Session.hxx"
#include "Cookie.hxx"
#include "JvmRoute.hxx"
#include "Headers.hxx"
#include "ssl/Filter.hxx"
#include "address_sticky.hxx"
#include "address_string.hxx"
#include "http_server/http_server.hxx"
#include "http/IncomingRequest.hxx"
#include "http_server/Handler.hxx"
#include "http_client.hxx"
#include "fs/Stock.hxx"
#include "fs/Balancer.hxx"
#include "HttpResponseHandler.hxx"
#include "http/Headers.hxx"
#include "stock/GetHandler.hxx"
#include "stock/Item.hxx"
#include "lease.hxx"
#include "strmap.hxx"
#include "pool/pool.hxx"
#include "net/IPv4Address.hxx"
#include "net/IPv6Address.hxx"
#include "net/SocketAddress.hxx"
#include "net/SocketDescriptor.hxx"
#include "net/FailureManager.hxx"
#include "net/FailureRef.hxx"
#include "istream/istream.hxx"
#include "istream/UnusedHoldPtr.hxx"
#include "util/Cancellable.hxx"
#include "util/LeakDetector.hxx"
#include "util/FNVHash.hxx"
#include "AllocatorPtr.hxx"
#include "stopwatch.hxx"

static constexpr Event::Duration LB_HTTP_CONNECT_TIMEOUT =
    std::chrono::seconds(20);

class LbRequest final
    : LeakDetector, Cancellable, StockGetHandler, Lease, HttpResponseHandler {

    struct pool &pool;

    LbHttpConnection &connection;
    LbCluster &cluster;
    const LbClusterConfig &cluster_config;

    FilteredSocketBalancer &balancer;

    IncomingHttpRequest &request;

    /**
     * The request body.
     */
    UnusedHoldIstreamPtr body;

    CancellablePointer cancel_ptr;

    /**
     * The cluster member we're currently connecting to (waiting for
     * #StockGetHandler to be called).  Only used for Zeroconf to be
     * able to call failure_set().
     */
    LbCluster::MemberPtr current_member;

    StockItem *stock_item = nullptr;
    FailurePtr failure;

    /**
     * The number of remaining connection attempts.  We give up when
     * we get an error and this attribute is already zero.
     */
    unsigned retries;

    unsigned new_cookie = 0;

    bool response_sent = false;

public:
    LbRequest(LbHttpConnection &_connection, LbCluster &_cluster,
              FilteredSocketBalancer &_balancer,
              IncomingHttpRequest &_request,
              CancellablePointer &_cancel_ptr) noexcept
        :pool(_request.pool), connection(_connection), cluster(_cluster),
         cluster_config(cluster.GetConfig()),
         balancer(_balancer),
         request(_request),
         body(pool, std::move(request.body)) {
        _cancel_ptr = *this;

        if (cluster_config.HasZeroConf())
            retries = CalculateRetries(cluster.GetZeroconfCount());
    }

    EventLoop &GetEventLoop() const noexcept {
        return connection.instance.event_loop;
    }

    FailureManager &GetFailureManager() noexcept {
        return balancer.GetFailureManager();
    }

    void Start() noexcept;

private:
    /* code copied from generic_balancer.hxx */
    static constexpr unsigned CalculateRetries(size_t size) noexcept {
        if (size <= 1)
            return 0;
        else if (size == 2)
            return 1;
        else if (size == 3)
            return 2;
        else
            return 3;
    }

    void Destroy() noexcept {
        assert(stock_item == nullptr);

        this->~LbRequest();
    }

    void SetForwardedTo() noexcept {
        // TODO: optimize this operation
        connection.per_request.forwarded_to =
            address_to_string(pool,
                              GetFailureManager().GetAddress(*failure));
    }

    void ResponseSent() noexcept {
        assert(!response_sent);
        response_sent = true;

        if (stock_item == nullptr)
            Destroy();
    }

    const char *GetCanonicalHost() const noexcept {
        return connection.per_request.GetCanonicalHost();
    }

    sticky_hash_t GetStickyHash() noexcept;
    sticky_hash_t GetHostHash() const noexcept;
    sticky_hash_t GetXHostHash() const noexcept;
    sticky_hash_t MakeCookieHash() noexcept;

    SocketAddress MakeBindAddress() const noexcept;

    /* virtual methods from class Cancellable */
    void Cancel() noexcept override {
        assert(!response_sent);

        cancel_ptr.Cancel();
        Destroy();
    }

    /* virtual methods from class StockGetHandler */
    void OnStockItemReady(StockItem &item) noexcept override;
    void OnStockItemError(std::exception_ptr ep) noexcept override;

    /* virtual methods from class Lease */
    void ReleaseLease(bool reuse) noexcept override;

    /* virtual methods from class HttpResponseHandler */
    void OnHttpResponse(http_status_t status, StringMap &&headers,
                        UnusedIstreamPtr body) noexcept override;
    void OnHttpError(std::exception_ptr ep) noexcept override;
};

static void
SendResponse(IncomingHttpRequest &request,
             const LbSimpleHttpResponse &response) noexcept
{
    assert(response.IsDefined());

    request.SendSimpleResponse(response.status,
                               response.location.empty() ? nullptr : response.location.c_str(),
                               response.message.empty() ? nullptr : response.message.c_str());
}

static bool
send_fallback(IncomingHttpRequest &request,
              const LbSimpleHttpResponse &fallback) noexcept
{
    if (!fallback.IsDefined())
        return false;

    SendResponse(request, fallback);
    return true;
}

inline sticky_hash_t
LbRequest::GetHostHash() const noexcept
{
    const char *host = GetCanonicalHost();
    if (host == nullptr)
        return 0;

    return FNV1aHash32(host);
}

inline sticky_hash_t
LbRequest::GetXHostHash() const noexcept
{
    const char *host = request.headers.Get("x-cm4all-host");
    if (host == nullptr)
        return 0;

    return FNV1aHash32(host);
}

/**
 * Generate a cookie for sticky worker selection.  Return only worker
 * numbers that are not known to be failing.  Returns 0 on total
 * failure.
 */
static unsigned
GenerateCookie(const FailureManager &failure_manager, Expiry now,
               const AddressList &list) noexcept
{
    assert(list.GetSize() >= 2);

    const unsigned first = lb_cookie_generate(list.GetSize());

    unsigned i = first;
    do {
        assert(i >= 1 && i <= list.GetSize());
        const SocketAddress address = list.addresses[i % list.GetSize()];
        if (failure_manager.Check(now, address))
            return i;

        i = lb_cookie_next(list.GetSize(), i);
    } while (i != first);

    /* all nodes have failed */
    return first;
}

sticky_hash_t
LbRequest::MakeCookieHash() noexcept
{
    unsigned hash = lb_cookie_get(request.headers);
    if (hash == 0)
        new_cookie = hash = GenerateCookie(GetFailureManager(),
                                           GetEventLoop().SteadyNow(),
                                           cluster_config.address_list);

    return hash;
}

sticky_hash_t
LbRequest::GetStickyHash() noexcept
{
    switch (cluster_config.sticky_mode) {
    case StickyMode::NONE:
    case StickyMode::FAILOVER:
        /* these modes require no preparation; they are handled
           completely by balancer_get() */
        return 0;

    case StickyMode::SOURCE_IP:
        /* calculate session_sticky from remote address */
        return socket_address_sticky(request.remote_address);

    case StickyMode::HOST:
        /* calculate session_sticky from "Host" request header */
        return GetHostHash();

    case StickyMode::XHOST:
        /* calculate session_sticky from "X-CM4all-Host" request
           header */
        return GetXHostHash();

    case StickyMode::SESSION_MODULO:
        /* calculate session_sticky from beng-proxy session id */
        return lb_session_get(request.headers,
                              cluster_config.session_cookie.c_str());

    case StickyMode::COOKIE:
        /* calculate session_sticky from beng-lb cookie */
        return MakeCookieHash();

    case StickyMode::JVM_ROUTE:
        /* calculate session_sticky from JSESSIONID cookie suffix */
        return lb_jvm_route_get(request.headers, cluster_config);
    }

    assert(false);
    gcc_unreachable();
}

/*
 * HTTP response handler
 *
 */

void
LbRequest::OnHttpResponse(http_status_t status, StringMap &&_headers,
                          UnusedIstreamPtr response_body) noexcept
{
    assert(!response_sent);

    failure->UnsetProtocol();

    SetForwardedTo();

    HttpHeaders headers(std::move(_headers));

    if (request.method == HTTP_METHOD_HEAD)
        /* pass Content-Length, even though there is no response body
           (RFC 2616 14.13) */
        headers.MoveToBuffer("content-length");

    if (new_cookie != 0) {
        char buffer[64];
        /* "Discard" must be last, to work around an Android bug*/
        snprintf(buffer, sizeof(buffer),
                 "beng_lb_node=0-%x; HttpOnly; Path=/; Version=1; Discard",
                 new_cookie);

        headers.Write("cookie2", "$Version=\"1\"");
        headers.Write("set-cookie", buffer);
    }

    request.SendResponse(status, std::move(headers), std::move(response_body));
    ResponseSent();
}

void
LbRequest::OnHttpError(std::exception_ptr ep) noexcept
{
    assert(!response_sent);

    if (IsHttpClientServerFailure(ep))
        failure->SetProtocol(GetEventLoop().SteadyNow(),
                             std::chrono::seconds(20));

    SetForwardedTo();

    connection.logger(2, ep);

    if (!send_fallback(request, cluster_config.fallback))
        connection.SendError(request, ep);

    ResponseSent();
}

/*
 * stock callback
 *
 */

void
LbRequest::OnStockItemReady(StockItem &item) noexcept
{
    assert(stock_item == nullptr);
    assert(!response_sent);

    stock_item = &item;

    if (cluster_config.HasZeroConf()) {
        failure = current_member->GetFailureRef();

        /* without the fs_balancer, we have to roll our own failure
           updates */
        failure->UnsetConnect();
    } else
        failure = GetFailureManager().Make(fs_stock_item_get_address(*stock_item));

    const char *peer_subject = connection.ssl_filter != nullptr
        ? ssl_filter_get_peer_subject(connection.ssl_filter)
        : nullptr;
    const char *peer_issuer_subject = connection.ssl_filter != nullptr
        ? ssl_filter_get_peer_issuer_subject(connection.ssl_filter)
        : nullptr;

    auto &headers = request.headers;
    lb_forward_request_headers(pool, headers,
                               request.local_host_and_port,
                               request.remote_host,
                               connection.IsEncrypted(),
                               peer_subject, peer_issuer_subject,
                               cluster_config.mangle_via);

    http_client_request(pool, nullptr,
                        fs_stock_item_get(item),
                        *this,
                        item.GetStockName(),
                        request.method, request.uri,
                        HttpHeaders(std::move(headers)),
                        std::move(body), true,
                        *this, cancel_ptr);
}

void
LbRequest::OnStockItemError(std::exception_ptr ep) noexcept
{
    assert(stock_item == nullptr);
    assert(!response_sent);

    connection.logger(2, "Connect error: ", ep);

    if (cluster_config.HasZeroConf()) {
        /* without the tcp_balancer, we have to roll our own failure
           updates and retries */
        current_member->GetFailureInfo().SetConnect(GetEventLoop().SteadyNow(),
                                                    std::chrono::seconds(20));

        if (retries-- > 0) {
            /* try the next Zeroconf member */
            Start();
            return;
        }
    }

    body.Clear();

    if (!send_fallback(request, cluster_config.fallback))
        connection.SendError(request, ep);

    ResponseSent();
}

/*
 * Lease
 *
 */

void
LbRequest::ReleaseLease(bool reuse) noexcept
{
    assert(stock_item != nullptr);

    stock_item->Put(!reuse);
    stock_item = nullptr;

    if (response_sent) {
        Destroy();
    }
}

/*
 * constructor
 *
 */

inline SocketAddress
LbRequest::MakeBindAddress() const noexcept
{
    if (cluster_config.transparent_source) {
        SocketAddress bind_address = request.remote_address;

        /* reset the port to 0 to allow the kernel to choose one */
        if (bind_address.GetFamily() == AF_INET) {
            auto &address = *NewFromPool<IPv4Address>(pool,
                                                      IPv4Address(bind_address));
            address.SetPort(0);
            return address;
        } else if (bind_address.GetFamily() == AF_INET6) {
            auto &address = *NewFromPool<IPv6Address>(pool,
                                                      IPv6Address(bind_address));
            address.SetPort(0);
            return address;
        }

        return bind_address;
    } else
        return SocketAddress::Null();
}

inline void
LbRequest::Start() noexcept
{
    const auto bind_address = MakeBindAddress();

    if (cluster_config.HasZeroConf()) {
        auto *member = cluster.Pick(GetEventLoop().SteadyNow(),
                                    GetStickyHash());
        if (member == nullptr) {
            auto &_request = request;
            Destroy();
            _request.SendMessage(HTTP_STATUS_INTERNAL_SERVER_ERROR,
                                 "Zeroconf cluster is empty");
            return;
        }

        current_member = *member;

        connection.instance.fs_stock->Get(pool,
                                          nullptr,
                                          member->GetLogName(),
                                          cluster_config.transparent_source,
                                          bind_address,
                                          member->GetAddress(),
                                          LB_HTTP_CONNECT_TIMEOUT,
                                          nullptr,
                                          *this, cancel_ptr);

        return;
    }

    balancer.Get(pool, nullptr,
                 cluster_config.transparent_source,
                 bind_address,
                 GetStickyHash(),
                 cluster_config.address_list,
                 LB_HTTP_CONNECT_TIMEOUT,
                 nullptr,
                 *this, cancel_ptr);
}

void
ForwardHttpRequest(LbHttpConnection &connection,
                   IncomingHttpRequest &request,
                   LbCluster &cluster,
                   CancellablePointer &cancel_ptr) noexcept
{
    const auto request2 =
        NewFromPool<LbRequest>(request.pool,
                               connection, cluster,
                               *connection.instance.fs_balancer,
                               request, cancel_ptr);
    request2->Start();
}
<commit_msg>lb/ForwardHttpRequest: destroy LbRequest before sending response<commit_after>/*
 * Copyright 2007-2017 Content Management AG
 * All rights reserved.
 *
 * author: Max Kellermann <mk@cm4all.com>
 *
 * 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.
 *
 * 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
 * FOUNDATION 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.
 */

#include "ForwardHttpRequest.hxx"
#include "HttpConnection.hxx"
#include "Cluster.hxx"
#include "ClusterConfig.hxx"
#include "ListenerConfig.hxx"
#include "Instance.hxx"
#include "Session.hxx"
#include "Cookie.hxx"
#include "JvmRoute.hxx"
#include "Headers.hxx"
#include "ssl/Filter.hxx"
#include "address_sticky.hxx"
#include "address_string.hxx"
#include "http_server/http_server.hxx"
#include "http/IncomingRequest.hxx"
#include "http_server/Handler.hxx"
#include "http_client.hxx"
#include "fs/Stock.hxx"
#include "fs/Balancer.hxx"
#include "HttpResponseHandler.hxx"
#include "http/Headers.hxx"
#include "stock/GetHandler.hxx"
#include "stock/Item.hxx"
#include "lease.hxx"
#include "strmap.hxx"
#include "pool/pool.hxx"
#include "net/IPv4Address.hxx"
#include "net/IPv6Address.hxx"
#include "net/SocketAddress.hxx"
#include "net/SocketDescriptor.hxx"
#include "net/FailureManager.hxx"
#include "net/FailureRef.hxx"
#include "istream/istream.hxx"
#include "istream/UnusedHoldPtr.hxx"
#include "util/Cancellable.hxx"
#include "util/LeakDetector.hxx"
#include "util/FNVHash.hxx"
#include "AllocatorPtr.hxx"
#include "stopwatch.hxx"

static constexpr Event::Duration LB_HTTP_CONNECT_TIMEOUT =
    std::chrono::seconds(20);

class LbRequest final
    : LeakDetector, Cancellable, StockGetHandler, Lease, HttpResponseHandler {

    struct pool &pool;

    LbHttpConnection &connection;
    LbCluster &cluster;
    const LbClusterConfig &cluster_config;

    FilteredSocketBalancer &balancer;

    IncomingHttpRequest &request;

    /**
     * The request body.
     */
    UnusedHoldIstreamPtr body;

    CancellablePointer cancel_ptr;

    /**
     * The cluster member we're currently connecting to (waiting for
     * #StockGetHandler to be called).  Only used for Zeroconf to be
     * able to call failure_set().
     */
    LbCluster::MemberPtr current_member;

    StockItem *stock_item = nullptr;
    FailurePtr failure;

    /**
     * The number of remaining connection attempts.  We give up when
     * we get an error and this attribute is already zero.
     */
    unsigned retries;

    unsigned new_cookie = 0;

    bool response_sent = false;

public:
    LbRequest(LbHttpConnection &_connection, LbCluster &_cluster,
              FilteredSocketBalancer &_balancer,
              IncomingHttpRequest &_request,
              CancellablePointer &_cancel_ptr) noexcept
        :pool(_request.pool), connection(_connection), cluster(_cluster),
         cluster_config(cluster.GetConfig()),
         balancer(_balancer),
         request(_request),
         body(pool, std::move(request.body)) {
        _cancel_ptr = *this;

        if (cluster_config.HasZeroConf())
            retries = CalculateRetries(cluster.GetZeroconfCount());
    }

    EventLoop &GetEventLoop() const noexcept {
        return connection.instance.event_loop;
    }

    FailureManager &GetFailureManager() noexcept {
        return balancer.GetFailureManager();
    }

    void Start() noexcept;

private:
    /* code copied from generic_balancer.hxx */
    static constexpr unsigned CalculateRetries(size_t size) noexcept {
        if (size <= 1)
            return 0;
        else if (size == 2)
            return 1;
        else if (size == 3)
            return 2;
        else
            return 3;
    }

    void Destroy() noexcept {
        assert(stock_item == nullptr);

        this->~LbRequest();
    }

    void SetForwardedTo() noexcept {
        // TODO: optimize this operation
        connection.per_request.forwarded_to =
            address_to_string(pool,
                              GetFailureManager().GetAddress(*failure));
    }

    void ResponseSent() noexcept {
        assert(!response_sent);
        response_sent = true;

        if (stock_item == nullptr)
            Destroy();
    }

    const char *GetCanonicalHost() const noexcept {
        return connection.per_request.GetCanonicalHost();
    }

    sticky_hash_t GetStickyHash() noexcept;
    sticky_hash_t GetHostHash() const noexcept;
    sticky_hash_t GetXHostHash() const noexcept;
    sticky_hash_t MakeCookieHash() noexcept;

    SocketAddress MakeBindAddress() const noexcept;

    /* virtual methods from class Cancellable */
    void Cancel() noexcept override {
        assert(!response_sent);

        cancel_ptr.Cancel();
        Destroy();
    }

    /* virtual methods from class StockGetHandler */
    void OnStockItemReady(StockItem &item) noexcept override;
    void OnStockItemError(std::exception_ptr ep) noexcept override;

    /* virtual methods from class Lease */
    void ReleaseLease(bool reuse) noexcept override;

    /* virtual methods from class HttpResponseHandler */
    void OnHttpResponse(http_status_t status, StringMap &&headers,
                        UnusedIstreamPtr body) noexcept override;
    void OnHttpError(std::exception_ptr ep) noexcept override;
};

static void
SendResponse(IncomingHttpRequest &request,
             const LbSimpleHttpResponse &response) noexcept
{
    assert(response.IsDefined());

    request.SendSimpleResponse(response.status,
                               response.location.empty() ? nullptr : response.location.c_str(),
                               response.message.empty() ? nullptr : response.message.c_str());
}

static bool
send_fallback(IncomingHttpRequest &request,
              const LbSimpleHttpResponse &fallback) noexcept
{
    if (!fallback.IsDefined())
        return false;

    SendResponse(request, fallback);
    return true;
}

inline sticky_hash_t
LbRequest::GetHostHash() const noexcept
{
    const char *host = GetCanonicalHost();
    if (host == nullptr)
        return 0;

    return FNV1aHash32(host);
}

inline sticky_hash_t
LbRequest::GetXHostHash() const noexcept
{
    const char *host = request.headers.Get("x-cm4all-host");
    if (host == nullptr)
        return 0;

    return FNV1aHash32(host);
}

/**
 * Generate a cookie for sticky worker selection.  Return only worker
 * numbers that are not known to be failing.  Returns 0 on total
 * failure.
 */
static unsigned
GenerateCookie(const FailureManager &failure_manager, Expiry now,
               const AddressList &list) noexcept
{
    assert(list.GetSize() >= 2);

    const unsigned first = lb_cookie_generate(list.GetSize());

    unsigned i = first;
    do {
        assert(i >= 1 && i <= list.GetSize());
        const SocketAddress address = list.addresses[i % list.GetSize()];
        if (failure_manager.Check(now, address))
            return i;

        i = lb_cookie_next(list.GetSize(), i);
    } while (i != first);

    /* all nodes have failed */
    return first;
}

sticky_hash_t
LbRequest::MakeCookieHash() noexcept
{
    unsigned hash = lb_cookie_get(request.headers);
    if (hash == 0)
        new_cookie = hash = GenerateCookie(GetFailureManager(),
                                           GetEventLoop().SteadyNow(),
                                           cluster_config.address_list);

    return hash;
}

sticky_hash_t
LbRequest::GetStickyHash() noexcept
{
    switch (cluster_config.sticky_mode) {
    case StickyMode::NONE:
    case StickyMode::FAILOVER:
        /* these modes require no preparation; they are handled
           completely by balancer_get() */
        return 0;

    case StickyMode::SOURCE_IP:
        /* calculate session_sticky from remote address */
        return socket_address_sticky(request.remote_address);

    case StickyMode::HOST:
        /* calculate session_sticky from "Host" request header */
        return GetHostHash();

    case StickyMode::XHOST:
        /* calculate session_sticky from "X-CM4all-Host" request
           header */
        return GetXHostHash();

    case StickyMode::SESSION_MODULO:
        /* calculate session_sticky from beng-proxy session id */
        return lb_session_get(request.headers,
                              cluster_config.session_cookie.c_str());

    case StickyMode::COOKIE:
        /* calculate session_sticky from beng-lb cookie */
        return MakeCookieHash();

    case StickyMode::JVM_ROUTE:
        /* calculate session_sticky from JSESSIONID cookie suffix */
        return lb_jvm_route_get(request.headers, cluster_config);
    }

    assert(false);
    gcc_unreachable();
}

/*
 * HTTP response handler
 *
 */

void
LbRequest::OnHttpResponse(http_status_t status, StringMap &&_headers,
                          UnusedIstreamPtr response_body) noexcept
{
    assert(!response_sent);

    failure->UnsetProtocol();

    SetForwardedTo();

    HttpHeaders headers(std::move(_headers));

    if (request.method == HTTP_METHOD_HEAD)
        /* pass Content-Length, even though there is no response body
           (RFC 2616 14.13) */
        headers.MoveToBuffer("content-length");

    if (new_cookie != 0) {
        char buffer[64];
        /* "Discard" must be last, to work around an Android bug*/
        snprintf(buffer, sizeof(buffer),
                 "beng_lb_node=0-%x; HttpOnly; Path=/; Version=1; Discard",
                 new_cookie);

        headers.Write("cookie2", "$Version=\"1\"");
        headers.Write("set-cookie", buffer);
    }

    auto &_request = request;
    ResponseSent();

    _request.SendResponse(status, std::move(headers),
                          std::move(response_body));
}

void
LbRequest::OnHttpError(std::exception_ptr ep) noexcept
{
    assert(!response_sent);

    if (IsHttpClientServerFailure(ep))
        failure->SetProtocol(GetEventLoop().SteadyNow(),
                             std::chrono::seconds(20));

    SetForwardedTo();

    connection.logger(2, ep);

    auto &_connection = connection;
    auto &_request = request;
    const auto &fallback = cluster_config.fallback;

    ResponseSent();

    if (!send_fallback(_request, fallback))
        _connection.SendError(_request, ep);
}

/*
 * stock callback
 *
 */

void
LbRequest::OnStockItemReady(StockItem &item) noexcept
{
    assert(stock_item == nullptr);
    assert(!response_sent);

    stock_item = &item;

    if (cluster_config.HasZeroConf()) {
        failure = current_member->GetFailureRef();

        /* without the fs_balancer, we have to roll our own failure
           updates */
        failure->UnsetConnect();
    } else
        failure = GetFailureManager().Make(fs_stock_item_get_address(*stock_item));

    const char *peer_subject = connection.ssl_filter != nullptr
        ? ssl_filter_get_peer_subject(connection.ssl_filter)
        : nullptr;
    const char *peer_issuer_subject = connection.ssl_filter != nullptr
        ? ssl_filter_get_peer_issuer_subject(connection.ssl_filter)
        : nullptr;

    auto &headers = request.headers;
    lb_forward_request_headers(pool, headers,
                               request.local_host_and_port,
                               request.remote_host,
                               connection.IsEncrypted(),
                               peer_subject, peer_issuer_subject,
                               cluster_config.mangle_via);

    http_client_request(pool, nullptr,
                        fs_stock_item_get(item),
                        *this,
                        item.GetStockName(),
                        request.method, request.uri,
                        HttpHeaders(std::move(headers)),
                        std::move(body), true,
                        *this, cancel_ptr);
}

void
LbRequest::OnStockItemError(std::exception_ptr ep) noexcept
{
    assert(stock_item == nullptr);
    assert(!response_sent);

    connection.logger(2, "Connect error: ", ep);

    if (cluster_config.HasZeroConf()) {
        /* without the tcp_balancer, we have to roll our own failure
           updates and retries */
        current_member->GetFailureInfo().SetConnect(GetEventLoop().SteadyNow(),
                                                    std::chrono::seconds(20));

        if (retries-- > 0) {
            /* try the next Zeroconf member */
            Start();
            return;
        }
    }

    body.Clear();

    auto &_connection = connection;
    auto &_request = request;
    const auto &fallback = cluster_config.fallback;

    ResponseSent();

    if (!send_fallback(_request, fallback))
        _connection.SendError(_request, ep);
}

/*
 * Lease
 *
 */

void
LbRequest::ReleaseLease(bool reuse) noexcept
{
    assert(stock_item != nullptr);

    stock_item->Put(!reuse);
    stock_item = nullptr;

    if (response_sent) {
        Destroy();
    }
}

/*
 * constructor
 *
 */

inline SocketAddress
LbRequest::MakeBindAddress() const noexcept
{
    if (cluster_config.transparent_source) {
        SocketAddress bind_address = request.remote_address;

        /* reset the port to 0 to allow the kernel to choose one */
        if (bind_address.GetFamily() == AF_INET) {
            auto &address = *NewFromPool<IPv4Address>(pool,
                                                      IPv4Address(bind_address));
            address.SetPort(0);
            return address;
        } else if (bind_address.GetFamily() == AF_INET6) {
            auto &address = *NewFromPool<IPv6Address>(pool,
                                                      IPv6Address(bind_address));
            address.SetPort(0);
            return address;
        }

        return bind_address;
    } else
        return SocketAddress::Null();
}

inline void
LbRequest::Start() noexcept
{
    const auto bind_address = MakeBindAddress();

    if (cluster_config.HasZeroConf()) {
        auto *member = cluster.Pick(GetEventLoop().SteadyNow(),
                                    GetStickyHash());
        if (member == nullptr) {
            auto &_request = request;
            Destroy();
            _request.SendMessage(HTTP_STATUS_INTERNAL_SERVER_ERROR,
                                 "Zeroconf cluster is empty");
            return;
        }

        current_member = *member;

        connection.instance.fs_stock->Get(pool,
                                          nullptr,
                                          member->GetLogName(),
                                          cluster_config.transparent_source,
                                          bind_address,
                                          member->GetAddress(),
                                          LB_HTTP_CONNECT_TIMEOUT,
                                          nullptr,
                                          *this, cancel_ptr);

        return;
    }

    balancer.Get(pool, nullptr,
                 cluster_config.transparent_source,
                 bind_address,
                 GetStickyHash(),
                 cluster_config.address_list,
                 LB_HTTP_CONNECT_TIMEOUT,
                 nullptr,
                 *this, cancel_ptr);
}

void
ForwardHttpRequest(LbHttpConnection &connection,
                   IncomingHttpRequest &request,
                   LbCluster &cluster,
                   CancellablePointer &cancel_ptr) noexcept
{
    const auto request2 =
        NewFromPool<LbRequest>(request.pool,
                               connection, cluster,
                               *connection.instance.fs_balancer,
                               request, cancel_ptr);
    request2->Start();
}
<|endoftext|>