/*
 * Copyright 2005-2007, Ingo Weinhold <bonefish@cs.tu-berlin.de>.
 * All rights reserved. Distributed under the terms of the MIT License.
 */

#include "RemoteDisk.h"

#include <new>

#include <endian.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>

#include <KernelExport.h>
#include <OS.h>

#include <kernel.h>		// for IS_USER_ADDRESS()


//#define TRACE_REMOTE_DISK
#ifdef TRACE_REMOTE_DISK
#	define TRACE(x) dprintf x
#else
#	define TRACE(x) do {} while (false)
#endif


static const int kMaxRequestResendCount = 5;
static const bigtime_t kReceiveTimeout = 800000LL;
static const bigtime_t kRequestTimeout = 5000000LL;

#if __BYTE_ORDER == __LITTLE_ENDIAN

static inline
uint64_t swap_uint64(uint64_t data)
{
	return ((data & 0xff) << 56)
		| ((data & 0xff00) << 40)
		| ((data & 0xff0000) << 24)
		| ((data & 0xff000000) << 8)
		| ((data >> 8) & 0xff000000)
		| ((data >> 24) & 0xff0000)
		| ((data >> 40) & 0xff00)
		| ((data >> 56) & 0xff);
}

#define host_to_net64(data)	swap_uint64(data)
#define net_to_host64(data)	swap_uint64(data)

#endif

#if __BYTE_ORDER == __BIG_ENDIAN
#define host_to_net64(data)	(data)
#define net_to_host64(data)	(data)
#endif

#undef htonll
#undef ntohll
#define htonll(data)	host_to_net64(data)
#define ntohll(data)	net_to_host64(data)


enum {
	BUFFER_SIZE	= 2048
};


// constructor
RemoteDisk::RemoteDisk()
	: fImageSize(0),
	  fRequestID(0),
	  fSocket(-1),
	  fPacket(NULL),
	  fPacketSize(0)
{
}


// destructor
RemoteDisk::~RemoteDisk()
{
	if (fSocket >= 0)
		close(fSocket);

	free(fPacket);
}


// Init
status_t
RemoteDisk::Init(uint32 serverAddress, uint16 serverPort, off_t imageSize)
{
	status_t error = _Init();
	if (error != B_OK)
		return error;

	fServerAddress.sin_family = AF_INET;
	fServerAddress.sin_port = htons(serverPort);
	fServerAddress.sin_addr.s_addr = htonl(serverAddress);
	fServerAddress.sin_len = sizeof(sockaddr_in);

	fImageSize = imageSize;

	return B_OK;
}


// FindAnyRemoteDisk
status_t
RemoteDisk::FindAnyRemoteDisk()
{
	status_t error = _Init();
	if (error != B_OK)
		return error;

	// prepare request
	remote_disk_header request;
	request.command = REMOTE_DISK_HELLO_REQUEST;

	// init server address to broadcast
	fServerAddress.sin_family = AF_INET;
	fServerAddress.sin_port = htons(REMOTE_DISK_SERVER_PORT);
	fServerAddress.sin_addr.s_addr = htonl(INADDR_BROADCAST);
	fServerAddress.sin_len = sizeof(sockaddr_in);

	// set SO_BROADCAST on socket
	int soBroadcastValue = 1;
	if (setsockopt(fSocket, SOL_SOCKET, SO_BROADCAST, &soBroadcastValue,
			sizeof(soBroadcastValue)) < 0) {
		dprintf("RemoteDisk::Init(): Failed to set SO_BROADCAST on socket: "
			"%s\n", strerror(errno));
	}

	// send request
	sockaddr_in serverAddress;
	error = _SendRequest(&request, sizeof(request), REMOTE_DISK_HELLO_REPLY,
		&serverAddress);
	if (error != B_OK) {
		dprintf("RemoteDisk::FindAnyRemoteDisk(): Got no server reply: %s\n",
			strerror(error));
		return error;
	}
	remote_disk_header* reply = (remote_disk_header*)fPacket;

	// unset SO_BROADCAST on socket
	soBroadcastValue = 0;
	if (setsockopt(fSocket, SOL_SOCKET, SO_BROADCAST, &soBroadcastValue,
			sizeof(soBroadcastValue)) < 0) {
		dprintf("RemoteDisk::Init(): Failed to unset SO_BROADCAST on socket: "
			"%s\n", strerror(errno));
	}

	// init server address and size
	fServerAddress = serverAddress;
	fServerAddress.sin_port = reply->port;

	fImageSize = ntohll(reply->offset);

	return B_OK;
}


// ReadAt
ssize_t
RemoteDisk::ReadAt(off_t pos, void *_buffer, size_t bufferSize)
{
	if (fSocket < 0)
		return B_NO_INIT;

	uint8 *buffer = (uint8*)_buffer;
	if (!buffer || pos < 0)
		return B_BAD_VALUE;

	if (bufferSize == 0)
		return 0;

	// Check whether the current packet already contains the beginning of the
	// data to read.
	ssize_t bytesRead = _ReadFromPacket(pos, buffer, bufferSize);
	if (bytesRead < 0)
		return bytesRead;

	// If there still remains something to be read, we need to get it from the
	// server.
	status_t error = B_OK;
	while (bufferSize > 0) {
		// prepare request
		remote_disk_header request;
		request.offset = htonll(pos);
		uint32 toRead = min_c(bufferSize, REMOTE_DISK_BLOCK_SIZE);
		request.size = htons(toRead);
		request.command = REMOTE_DISK_READ_REQUEST;

		// send request
		error = _SendRequest(&request, sizeof(request), REMOTE_DISK_READ_REPLY);
		if (error != B_OK)
			break;

		// check for errors
		int16 packetSize = ntohs(((remote_disk_header*)fPacket)->size);
		if (packetSize < 0) {
			if (packetSize == REMOTE_DISK_IO_ERROR)
				error = B_IO_ERROR;
			else if (packetSize == REMOTE_DISK_BAD_REQUEST)
				error = B_BAD_VALUE;
			fPacketSize = 0;
			break;
		}

		// read from the packet
		size_t packetBytesRead = _ReadFromPacket(pos, buffer, bufferSize);
		if (packetBytesRead <= 0) {
			if (packetBytesRead < 0)
				error = packetBytesRead;
			break;
		}
		bytesRead += packetBytesRead;
	}

	// only return an error, when we were not able to read anything at all
	return (bytesRead == 0 ? error : bytesRead);
}


// WriteAt
ssize_t
RemoteDisk::WriteAt(off_t pos, const void *_buffer, size_t bufferSize)
{
	if (fSocket < 0)
		return B_NO_INIT;

	const uint8 *buffer = (const uint8*)_buffer;
	if (!buffer || pos < 0)
		return B_BAD_VALUE;

	if (bufferSize == 0)
		return 0;

	status_t error = B_OK;
	size_t bytesWritten = 0;
	while (bufferSize > 0) {
		// prepare request
		remote_disk_header* request = (remote_disk_header*)fPacket;
		request->offset = htonll(pos);
		uint32 toWrite = min_c(bufferSize, REMOTE_DISK_BLOCK_SIZE);
		request->size = htons(toWrite);
		request->command = REMOTE_DISK_WRITE_REQUEST;

		// copy to packet buffer
		if (IS_USER_ADDRESS(buffer)) {
			status_t error = user_memcpy(request->data, buffer, toWrite);
			if (error != B_OK)
				return error;
		} else
			memcpy(request->data, buffer, toWrite);

		// send request
		size_t requestSize = request->data + toWrite - (uint8_t*)request;
		remote_disk_header reply;
		int32 replySize;
		error = _SendRequest(request, requestSize, REMOTE_DISK_WRITE_REPLY,
			NULL, &reply, sizeof(reply), &replySize);
		if (error != B_OK)
			break;

		// check for errors
		int16 packetSize = ntohs(reply.size);
		if (packetSize < 0) {
			if (packetSize == REMOTE_DISK_IO_ERROR)
				error = B_IO_ERROR;
			else if (packetSize == REMOTE_DISK_BAD_REQUEST)
				error = B_BAD_VALUE;
			break;
		}

		bytesWritten += toWrite;
		pos += toWrite;
		buffer += toWrite;
		bufferSize -= toWrite;
	}

	// only return an error, when we were not able to write anything at all
	return (bytesWritten == 0 ? error : bytesWritten);
}


// _Init
status_t
RemoteDisk::_Init()
{
	// open a control socket for playing with the stack
	fSocket = socket(AF_INET, SOCK_DGRAM, 0);
	if (fSocket < 0) {
		dprintf("RemoteDisk::Init(): Failed to open socket: %s\n",
			strerror(errno));
		return errno;
	}

	// bind socket
	fSocketAddress.sin_family = AF_INET;
	fSocketAddress.sin_port = 0;
	fSocketAddress.sin_addr.s_addr = INADDR_ANY;
	fSocketAddress.sin_len = sizeof(sockaddr_in);
	if (bind(fSocket, (sockaddr*)&fSocketAddress, sizeof(fSocketAddress)) < 0) {
		dprintf("RemoteDisk::Init(): Failed to bind socket: %s\n",
			strerror(errno));
		return errno;
	}

   // get the port
    socklen_t addrSize = sizeof(fSocketAddress);
    if (getsockname(fSocket, (sockaddr*)&fSocketAddress, &addrSize) < 0) {
		dprintf("RemoteDisk::Init(): Failed to get socket address: %s\n",
			strerror(errno));
        return errno;
    }

	// set receive timeout
	timeval timeout;
	timeout.tv_sec = time_t(kReceiveTimeout / 1000000LL);
	timeout.tv_usec = suseconds_t(kReceiveTimeout % 1000000LL);
	if (setsockopt(fSocket, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout))
			< 0) {
		dprintf("RemoteDisk::Init(): Failed to set socket receive timeout: "
			"%s\n", strerror(errno));
        return errno;
	}

	// allocate buffer
	fPacket = malloc(BUFFER_SIZE);
	if (!fPacket)
		return B_NO_MEMORY;

	return B_OK;
}


// _ReadFromPacket
ssize_t
RemoteDisk::_ReadFromPacket(off_t& pos, uint8*& buffer, size_t& bufferSize)
{
	if (fPacketSize == 0)
		return 0;

	// check whether the cached packet is indeed a read reply
	remote_disk_header* header = (remote_disk_header*)fPacket;
	if (header->command != REMOTE_DISK_READ_REPLY)
		return 0;

	uint64 packetOffset = ntohll(header->offset);
	uint32 packetSize = ntohs(header->size);
	if (packetOffset > (uint64)pos || packetOffset + packetSize <= (uint64)pos)
		return 0;

	// we have something to copy
	size_t toCopy = size_t(packetOffset + packetSize - (uint64)pos);
	if (toCopy > bufferSize)
		toCopy = bufferSize;

	if (IS_USER_ADDRESS(buffer)) {
		status_t error = user_memcpy(buffer,
			header->data + (pos - packetOffset), toCopy);
		if (error != B_OK)
			return error;
	} else
		memcpy(buffer, header->data + (pos - packetOffset), toCopy);

	pos += toCopy;
	buffer += toCopy;
	bufferSize -= toCopy;
	return toCopy;
}


// _SendRequest
status_t
RemoteDisk::_SendRequest(remote_disk_header* request, size_t size,
	uint8 expectedReply, sockaddr_in* peerAddress)
{
	return _SendRequest(request, size, expectedReply, peerAddress, fPacket,
		BUFFER_SIZE, &fPacketSize);
}


status_t
RemoteDisk::_SendRequest(remote_disk_header *request, size_t size,
	uint8 expectedReply, sockaddr_in* peerAddress, void* receiveBuffer,
	size_t receiveBufferSize, int32* _bytesReceived)
{
	request->request_id = fRequestID++;
	request->port = fSocketAddress.sin_port;

	// try sending the request kMaxRequestResendCount times at most
	for (int i = 0; i < kMaxRequestResendCount; i++) {
		// send request
		ssize_t bytesSent;
		do {
			bytesSent = sendto(fSocket, request, size, 0,
				(sockaddr*)&fServerAddress, sizeof(fServerAddress));
		} while (bytesSent < 0 && errno == B_INTERRUPTED);

		if (bytesSent < 0) {
			dprintf("RemoteDisk::_SendRequest(): failed to send packet: %s\n",
				strerror(errno));
			return errno;
		}
		if (bytesSent != (ssize_t)size) {
			dprintf("RemoteDisk::_SendRequest(): sent less bytes than "
				"desired\n");
			return B_ERROR;
		}

		// receive reply
		bigtime_t timeout = system_time() + kRequestTimeout;
		do {
			*_bytesReceived = 0;
			socklen_t addrSize = sizeof(sockaddr_in);
			ssize_t bytesReceived = recvfrom(fSocket, receiveBuffer,
				receiveBufferSize, 0, (sockaddr*)peerAddress,
				(peerAddress ? &addrSize : NULL));
			if (bytesReceived < 0) {
				status_t error = errno;
				if (error != B_TIMED_OUT && error != B_WOULD_BLOCK
						&& error != B_INTERRUPTED) {
					dprintf("RemoteDisk::_SendRequest(): recvfrom() failed: "
						"%s\n", strerror(error));
					return error;
				}
				continue;
			}

			// got something; check, if it is looks good
			if (bytesReceived >= (ssize_t)sizeof(remote_disk_header)) {
				remote_disk_header* reply = (remote_disk_header*)receiveBuffer;
				if (reply->request_id == request->request_id
					&& reply->command == expectedReply) {
					*_bytesReceived = bytesReceived;
					return B_OK;
				}
			}
		} while (timeout > system_time());
	}

	// no reply
	return B_TIMED_OUT;
}