/* This file is part of wmaloader.*
 * Copyright 2004 Andrew Wild */

/* wmaloader 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. */

/* wmaloader 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 wmaloader; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

/*------------------------------------------------------------------------*/

/* system */
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <errno.h>
#include <assert.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

/* wmaloader components */
#include "common.h"

/* Globals */
int sock;            // Socket listening for connections
int client_sock = 0; // Socket of client
int source_fd = 0;   // File descriptor of source file
const char* source;  // Filename of source file

struct sockaddr_in address; // socket address IPV4

/*------------------------------------------------------------------------*/
int it_init(const char* interface, const char* source_file){
  /* Initialise socket. Return true for success, false for
	 failure. */

  int len;

  /* Save file name of source file */
  source = source_file;
  
  /* Create socket */
  sock = socket(PF_INET, SOCK_STREAM, 0);
  if (sock == -1){
	perror("it_init() socket");
	return FALSE;
  }

  /* Prepare local address */
  address.sin_family = AF_INET;
  address.sin_port = htons( 0 ); // Use any port
  address.sin_addr.s_addr = INADDR_ANY; // Use all interfaces
  if (interface != NULL){
	if (inet_aton(interface, &address.sin_addr) == 0){
	  message("%s is not a valid IP address\n", interface);
	  return FALSE;
	}
  }

  /* Bind socket */
  if (bind(sock, (struct sockaddr*)&address, sizeof(address)) == -1){
	perror("it_init() bind");
	return FALSE;
  }

  /* Fetch socket name */
  len = sizeof(address);
  if (getsockname(sock, (struct sockaddr*)&address, &len) == -1){
	perror("it_init() getsockname");
	return FALSE;
  }

  /* Listen */
  if (listen(sock, 5) == -1){
	perror("it_init() listen");
	return FALSE;
  }

  return TRUE;
}

/*------------------------------------------------------------------------*/
int it_get_port(){
  /* Return the port number in use. Can only be called after it_init() */  
  return ntohs(address.sin_port);
}

/*------------------------------------------------------------------------*/
int it_get_image_size(){
  /* Return the image size in bytes. If this fails zero is returned. */
  /* Can only be used after it_init() */
  struct stat st;
  if (stat(source, &st) == -1){
	perror("it_get_image_size() stat");
	return 0;
  }

  return st.st_size;
}

/*------------------------------------------------------------------------*/
void accept_new_connection(){
  /* Internal use only */
  
  struct sockaddr_in addr; // socket address IPV4
  int len = sizeof(struct sockaddr_in);
  int new_socket;
  int temp_fd;

  assert /* No current client */ (client_sock == 0);
  assert /* File is not open */ (source_fd == 0);

  new_socket = accept(sock, (struct sockaddr *)&addr, &len);
  if (new_socket == -1){
	perror("accept_new_connection() accept");
	return;
  }

  message("Accepted image request from %s:%d\n", 
		  inet_ntoa(addr.sin_addr), addr.sin_port);
  
  /* Open file to read data from */

  temp_fd = open(source, O_RDONLY);

  if (temp_fd == -1){
	/* Open failure */
	perror("accept_new_connection() open");
	close(new_socket);
	return;
  }

  client_sock = new_socket;
  source_fd = temp_fd;

  return;
}

/*------------------------------------------------------------------------*/
void transfer_data(){
  /* Internal use only */
  /* Read data from file and write to socket */

  const int block_size = 2048;
  unsigned char data[block_size];
  int ret;

  assert /* Block not too large */ (block_size <= SSIZE_MAX);
  assert /* File is open */ (source_fd != 0 && source_fd != -1);
  assert /* Client exists */ (client_sock != 0 && client_sock != -1);

  /* Read data */
  ret = read(source_fd, data, sizeof(data));
  switch (ret){
  case -1:
	/* Write error */
	perror("transfer_data() read");
	goto write_data_done;

  case 0:
	/* All data sent */
	//printf("Image transfer complete\n");
	goto write_data_done;

  default:
	/* Send data */
	if (send(client_sock, data, ret, 0) == -1){
	  perror("transfer_data() send");
	  goto write_data_done;
	}
  }
  return;

  /* An error occured or transmission is over.
   * Close client socket and close file */
  write_data_done:
  close(client_sock); client_sock = 0;
  close(source_fd); source_fd = 0;
  return;
}

/*------------------------------------------------------------------------*/
void* it_thread_entry(void *cookie){
  /* Thread entry routine. Can only be called after it_init() */

  int ret;
  int max_fd;

  fd_set temp_read, temp_write;

  while(1){
	
	FD_ZERO(&temp_read);
	FD_ZERO(&temp_write);
	if (sock)        FD_SET(sock, &temp_read);
	if (client_sock) FD_SET(client_sock, &temp_write);

	max_fd = sock > client_sock ? sock : client_sock;

	/* Wait for activity */
	ret = select(max_fd + 1, &temp_read, &temp_write, NULL, NULL);

	if (ret == -1){
	  perror("it_thread_entry() select");
	  return NULL;
	}

	/* Write data to client if socket read to send */
	if (FD_ISSET(client_sock, &temp_write)){
	  transfer_data();
	  continue;
	}

	/* Accept connection if one is available and not currently busy */
	if (FD_ISSET(sock, &temp_read) && client_sock == 0){
	  accept_new_connection();
	  continue;
	}	
  }

  return NULL;
}
