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

#include <boot/net/IP.h>

#include <stdio.h>
#include <KernelExport.h>

#include <boot/net/ARP.h>
#include <boot/net/ChainBuffer.h>


//#define TRACE_IP
#ifdef TRACE_IP
#	define TRACE(x) dprintf x
#else
#	define TRACE(x) ;
#endif


// #pragma mark - IPSubService

// constructor
IPSubService::IPSubService(const char *serviceName)
	: NetService(serviceName)
{
}

// destructor
IPSubService::~IPSubService()
{
}


// #pragma mark - IPService

// constructor
IPService::IPService(EthernetService *ethernet, ARPService *arpService)
	: EthernetSubService(kIPServiceName),
		fEthernet(ethernet),
		fARPService(arpService)
{
}

// destructor
IPService::~IPService()
{
	if (fEthernet)
		fEthernet->UnregisterEthernetSubService(this);
}

// Init
status_t
IPService::Init()
{
	if (!fEthernet)
		return B_BAD_VALUE;
	if (!fEthernet->RegisterEthernetSubService(this))
		return B_NO_MEMORY;
	return B_OK;
}

// IPAddress
ip_addr_t
IPService::IPAddress() const
{
	return (fEthernet ? fEthernet->IPAddress() : INADDR_ANY);
}

// EthernetProtocol
uint16
IPService::EthernetProtocol() const
{
	return ETHERTYPE_IP;
}

// HandleEthernetPacket
void
IPService::HandleEthernetPacket(EthernetService *ethernet,
	const mac_addr_t &targetAddress, const void *data, size_t size)
{
	TRACE(("IPService::HandleEthernetPacket(): %lu - %lu bytes\n", size,
		sizeof(ip_header)));

	if (!data || size < sizeof(ip_header))
		return;

	// check header
	const ip_header *header = (const ip_header*)data;
	// header length OK?
	int headerLength = header->header_length * 4;
	if (headerLength < 20 || headerLength > (int)size
		// IP V4?
		|| header->version != IP_PROTOCOL_VERSION_4
		// length OK?
		|| ntohs(header->total_length) > size
		// broadcast or our IP?
		|| (header->destination != htonl(INADDR_BROADCAST)
			&& (fEthernet->IPAddress() == INADDR_ANY
				|| header->destination != htonl(fEthernet->IPAddress())))
		// checksum OK?
		|| _Checksum(*header) != 0) {
		return;
	}

	// find a service handling this kind of packet
	int serviceCount = fServices.Count();
	for (int i = 0; i < serviceCount; i++) {
		IPSubService *service = fServices.ElementAt(i);
		if (service->IPProtocol() == header->protocol) {
			service->HandleIPPacket(this, ntohl(header->source),
				ntohl(header->destination),
				(uint8*)data + headerLength,
				ntohs(header->total_length) - headerLength);
			break;
		}
	}
}

// Send
status_t
IPService::Send(ip_addr_t destination, uint8 protocol, ChainBuffer *buffer)
{
	TRACE(("IPService::Send(to: %08lx, proto: %lu, %lu bytes)\n", destination,
		(uint32)protocol, (buffer ? buffer->TotalSize() : 0)));

	if (!buffer)
		return B_BAD_VALUE;

	if (!fEthernet || !fARPService)
		return B_NO_INIT;

	// prepare header
	ip_header header;
	ChainBuffer headerBuffer(&header, sizeof(header), buffer);
	header.header_length = 5;	// 5 32 bit words, no options
	header.version = IP_PROTOCOL_VERSION_4;
	header.type_of_service = 0;
	header.total_length = htons(headerBuffer.TotalSize());
	header.identifier = 0;
	header.fragment_offset = htons(IP_DONT_FRAGMENT);
	header.time_to_live = IP_DEFAULT_TIME_TO_LIVE;
	header.protocol = protocol;
	header.checksum = 0;
	header.source = htonl(fEthernet->IPAddress());
	header.destination = htonl(destination);

	// compute check sum
	header.checksum = htons(_Checksum(header));

	// get target MAC address
	mac_addr_t targetMAC;
	status_t error = fARPService->GetMACForIP(destination, targetMAC);
	if (error != B_OK)
		return error;

	// send the packet
	return fEthernet->Send(targetMAC, ETHERTYPE_IP, &headerBuffer);
}

// ProcessIncomingPackets
void
IPService::ProcessIncomingPackets()
{
	if (fEthernet)
		fEthernet->ProcessIncomingPackets();
}

// RegisterIPSubService
bool
IPService::RegisterIPSubService(IPSubService *service)
{
	return (service && fServices.Add(service) == B_OK);
}

// UnregisterIPSubService
bool
IPService::UnregisterIPSubService(IPSubService *service)
{
	return (service && fServices.Remove(service) >= 0);
}

// CountSubNetServices
int
IPService::CountSubNetServices() const
{
	return fServices.Count();
}

// SubNetServiceAt
NetService *
IPService::SubNetServiceAt(int index) const
{
	return fServices.ElementAt(index);
}

// _Checksum
uint16
IPService::_Checksum(const ip_header &header)
{
	ChainBuffer buffer((void*)&header, header.header_length * 4);
	return ip_checksum(&buffer);
}


// #pragma mark -

// ip_checksum
uint16
ip_checksum(ChainBuffer *buffer)
{
	// ChainBuffer iterator returning a stream of uint16 (big endian).
	struct Iterator {
		Iterator(ChainBuffer *buffer)
			: fBuffer(buffer),
			  fOffset(-1)
		{
			_Next();
		}

		bool HasNext() const
		{
			return fBuffer;
		}

		uint16 Next()
		{
			uint16 byte = _NextByte();
			return (byte << 8) | _NextByte();
		}

	private:
		void _Next()
		{
			while (fBuffer) {
				fOffset++;
				if (fOffset < (int)fBuffer->Size())
					break;

				fOffset = -1;
				fBuffer = fBuffer->Next();
			}
		}

		uint8 _NextByte()
		{
			uint8 byte = (fBuffer ? ((uint8*)fBuffer->Data())[fOffset] : 0);
			_Next();
			return byte;
		}

		ChainBuffer	*fBuffer;
		int			fOffset;
	};

	Iterator it(buffer);

	uint32 checksum = 0;
	while (it.HasNext()) {
		checksum += it.Next();
		while (checksum >> 16)
			checksum = (checksum & 0xffff) + (checksum >> 16);
	}

	return ~checksum;
}


ip_addr_t
ip_parse_address(const char *string)
{
	ip_addr_t address = 0;
	int components = 0;

	// TODO: Handles only IPv4 addresses for now.
	while (components < 4) {
		address |= strtol(string, NULL, 0) << ((4 - components - 1) * 8);

		const char *dot = strchr(string, '.');
		if (dot == NULL)
			break;

		string = dot + 1;
		components++;
	}

	return address;
}