/* Socks Server 5
 * Copyright (C) 2002 - 2006 by Matteo Ricchetti - <matteo.ricchetti@libero.it>

 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 */

#include "SS5Main.h"
#include "SS5Mod_balance.h"


S5RetCode InitModule( struct _module *m )
{
  m->AddVip  = AddVip;
  m->AddConn = S5AddConn2Real;
  m->RemoveConn = S5RemoveConn2Real;
  m->LoadBalancing = LoadBalancing;
  m->Balancing = Balancing;
  m->FreeConnectionTable = FreeConnectionTable;
  m->FreeAffinity = FreeAffinity;

  return OK;
}

S5RetCode AddVip (char *real, unsigned int vid, unsigned int idx)
{
  S5AddReal2ConnectionTable(real, vid, idx);
 
  return OK;
}

S5RetCode FreeConnectionTable (struct _S5ConnectionEntry *ce)
{
  free(ce);
 
  return OK;
}

S5RetCode LoadBalancing( struct _SS5ClientInfo *ci, struct _SS5RequestInfo *ri )
{
  struct in_addr s;
  struct in_addr d;

  unsigned int vid;
  unsigned int ttl_status;

  /*
   *    If affinity enabled, looks for affinity between src ip and dst ip
   *    before using balancing
   */
  if( SS5SocksOpt.Sticky ) {
    IFLINUX( inet_aton(ci->SrcAddr,&s); )
    IFSOLARIS( inet_pton(AF_INET, ci->SrcAddr, &s); )

    ttl_status = OK;
    vid = S5GetRealVid(ri->DstAddr);

    LOCKMUTEXCA()
    if( (d.s_addr = S5GetAffinity(s.s_addr,&ttl_status,vid)) == ERR ) {
      if( ttl_status == ERR ) {
        /*
         *    If age expired, remove affinity between src ip and dst ip
         */
        S5RemoveAffinity(s.s_addr,vid);
      }

      /*
       *    Balances with least connections
       */
      S5LeastConnectionReal(ri->DstAddr);

      /*
       *    Setup affinity between src ip and dst ip
       */
      IFSOLARIS( inet_pton(AF_INET, ri->DstAddr, &d); )
      IFLINUX( inet_aton((const char *)ri->DstAddr,&d); )

      S5SetAffinity(s.s_addr,d.s_addr,vid);
    }
    else {
      strncpy(ri->DstAddr,inet_ntoa(d),sizeof(ri->DstAddr));
    }
    UNLOCKMUTEXCA()
  }
  else {
    /*
     *    Balances with least connections
     */
    S5LeastConnectionReal(ri->DstAddr);
  }
  
  return OK;
}

S5RetCode S5LeastConnectionReal(char *s5application)
{
  register unsigned int i;
  register unsigned int j;

  unsigned int id1;
  unsigned int vid;
  unsigned int conn;

  LOCKMUTEXCT()
  /*
   *    Search for one real with the smaller number of connections
   */
  for(i = 0; i < NReal; i++) {
    if( STREQ(S5ConnectionTable.S5ConnectionEntry[i]->Real,s5application,strlen(s5application)) ) {
      id1  = i;
      vid  = S5ConnectionTable.S5ConnectionEntry[i]->Vid; 
      conn = S5ConnectionTable.S5ConnectionEntry[i]->Connection;

      for(j = 0; j<NReal; j++)
        if( S5ConnectionTable.S5ConnectionEntry[j]->Vid == vid )
	  if(S5ConnectionTable.S5ConnectionEntry[j]->Connection < conn ) {
            conn = S5ConnectionTable.S5ConnectionEntry[j]->Connection;
	    id1  = j;
	  }
      strncpy(s5application,S5ConnectionTable.S5ConnectionEntry[id1]->Real,sizeof(S5ConnectionTable.S5ConnectionEntry[id1]->Real) - 1);
      UNLOCKMUTEXCT()
      return OK;
    }
  }
  UNLOCKMUTEXCT()
  return ERR;
}

S5RetCode S5AddReal2ConnectionTable(char *real, unsigned int vid, unsigned int idx)
{
  S5ConnectionTable._tmp_S5ConnectionEntry[idx] = (struct _S5ConnectionEntry *)calloc(1,sizeof(struct _S5ConnectionEntry));

  strncpy(S5ConnectionTable._tmp_S5ConnectionEntry[idx]->Real,real,strlen(real));
  S5ConnectionTable._tmp_S5ConnectionEntry[idx]->Vid = vid;
  S5ConnectionTable._tmp_S5ConnectionEntry[idx]->Connection = 0;

  return OK;
}

S5RetCode S5GetRealVid(char *real)
{
  register unsigned int idx;

  for(idx = 0; idx < NReal; idx++)
    if( STREQ(S5ConnectionTable.S5ConnectionEntry[idx]->Real,real,strlen(real)) ) {
      return S5ConnectionTable.S5ConnectionEntry[idx]->Vid;
    }
  return ERR;
}

S5RetCode S5AddConn2Real(char *real)
{
  register unsigned int idx;

  for(idx = 0; idx < NReal; idx++)
    if( STREQ(S5ConnectionTable.S5ConnectionEntry[idx]->Real,real,strlen(real)) ) {
      LOCKMUTEXCT()
      S5ConnectionTable.S5ConnectionEntry[idx]->Connection++;
      UNLOCKMUTEXCT()
      return OK;
    }
  return ERR;
}

S5RetCode S5RemoveConn2Real(char *real)
{
  register unsigned int idx;

  for(idx = 0; idx < NReal; idx++)
    if( STREQ(S5ConnectionTable.S5ConnectionEntry[idx]->Real,real,strlen(real)) ) {
      if( S5ConnectionTable.S5ConnectionEntry[idx]->Connection ) {
        LOCKMUTEXCT()
        S5ConnectionTable.S5ConnectionEntry[idx]->Connection--;
        UNLOCKMUTEXCT()
      }
      return OK;
    }
  return ERR;
}

inline S5Limit S5StickyHash( unsigned long int srcip )
{
  return (srcip % MAXSTICKYLIST);
}

unsigned long int S5GetAffinity(unsigned long int srcip, unsigned int *ttl_status, unsigned int vid)
{
  S5Limit idx;

  struct _S5StickyNode *node;

  idx = S5StickyHash( srcip );

  if( S5StickyList[idx] == NULL )
    return ERR;
  else {
    node = S5StickyList[idx];

    do {
     if( (node->srcip == srcip)  && (node->vid == vid)) {
       if( node->ttl > time(NULL) ) {
         return node->dstip;
       }
       else {
	 *ttl_status = ERR;
         return ERR;
       }
     }  
     node = node->next;
    } while(node != NULL );
  }
  return ERR;
}

S5RetCode S5SetAffinity(unsigned long int srcip, unsigned long int dstip, unsigned int vid )
{
  int idx;
  struct _S5StickyNode *node;
  struct in_addr s,d;

  s.s_addr = srcip;
  d.s_addr = dstip;

  idx=S5StickyHash( srcip );

  if( S5StickyList[idx]== NULL ) {
    S5StickyList[idx] = (struct _S5StickyNode *)calloc(1,sizeof(struct _S5StickyNode));
    S5StickyList[idx]->srcip = srcip;
    S5StickyList[idx]->dstip = dstip;
    S5StickyList[idx]->ttl = (time(NULL) + SS5SocksOpt.StickyAge);
    S5StickyList[idx]->vid = vid;
  }
  else {
    node=S5StickyList[idx];
    while( node->next != NULL ){
      node=node->next;
    }
    node->next = (struct _S5StickyNode *)calloc(1,sizeof(struct _S5StickyNode));
    node->next->srcip = srcip;
    node->next->dstip = dstip;
    node->next->ttl = (time(NULL) + SS5SocksOpt.StickyAge);
    node->next->vid = vid;
    node->next->next = NULL;
  }
  return OK;
}

S5RetCode S5RemoveAffinity(unsigned long int srcip, unsigned int vid)
{
  int idx;
  struct _S5StickyNode *node;
  struct _S5StickyNode **backnode;

  idx = S5StickyHash( srcip );

  node = S5StickyList[idx];
  backnode = &S5StickyList[idx];

  do {
   if( (node->srcip == srcip) && (node->vid == vid) ) {
     if( node->next != NULL )
       *backnode = node->next;
     else
       *backnode = NULL;

     free(node);
     break;
   }  
   backnode = &node->next;
   node = node->next;

  } while(node != NULL );

  return OK;
}

S5RetCode Balancing( struct _SS5ClientInfo *ci, struct _SS5Socks5Data *sd )
{

  register unsigned int idx;

  struct in_addr s;
  struct in_addr d;

  time_t currentAge;

  struct _S5StickyNode *node;

  char *buf,
       sa[16],
       da[16];

  buf = (char *)calloc(512,sizeof(char));

  if( STREQ(sd->MethodRequest,"GET /balancing HTTP/1.",sizeof("GET /balancing HTTP/1.") - 1) ) {
    /*
     *     Create response
     */

    for(idx = 0; idx < NReal; idx++) {
      snprintf(buf,512 - 1, "%s\n%u\n%u\n",
        S5ConnectionTable.S5ConnectionEntry[idx]->Real,
        S5ConnectionTable.S5ConnectionEntry[idx]->Vid,
        S5ConnectionTable.S5ConnectionEntry[idx]->Connection);

      /*
       *    Send http response
       */
      if( send(ci->Socket,buf,512,SS5_SEND_OPT) == -1) {
        free(buf);
        return ERR;
      }
    }

    fcntl(ci->Socket,F_SETFL,O_NONBLOCK);
    recv(ci->Socket,buf,strlen(buf),0);
    free(buf);
    return OK;
  }
  else if( STREQ(sd->MethodRequest,"GET /sticky HTTP/1.",sizeof("GET /sticky HTTP/1.") - 1) ) {
    /*
     *     Create response
     */

    for(idx = 0; idx < 997; idx++) {
      node = S5StickyList[idx];

      while( node != NULL ) {

        s.s_addr = node->srcip;
        d.s_addr = node->dstip;

        strncpy(sa,inet_ntoa(s),sizeof(sa));
        strncpy(da,inet_ntoa(d),sizeof(da));

        currentAge = time(NULL);

        snprintf(buf,512 - 1, "%s\n%u\n%s\n%lu\n%lu\n",sa,node->vid,da,node->ttl,currentAge);
        /*
         *    Send http response
         */
        if( send(ci->Socket,buf,512,SS5_SEND_OPT) == -1) {
          free(buf);
          return ERR;
        }

        node = node->next;
      }

    }
    fcntl(ci->Socket,F_SETFL,O_NONBLOCK);
    recv(ci->Socket,buf,strlen(buf),0);
    free(buf);
    return OK;
  }

  return ERR;
}

S5RetCode FreeAffinity( struct _S5StickyNode **node )
{
  struct _S5StickyNode *lnode;
  struct _S5StickyNode *lnode_prev=NULL;

  lnode=*node;

  if( lnode != NULL ) {
    do {
      while( lnode->next != NULL ) {
        lnode_prev=lnode;
        lnode=lnode->next;
      }
      free(lnode);
      if( lnode_prev != NULL ) {
        lnode_prev->next=NULL;
        lnode=lnode_prev;
        lnode_prev=NULL;
      }
      else
        lnode=NULL;
    } while( (lnode) != NULL );
  }
  *node=NULL;
  
  return OK;

}