/*
 * Copyright 2013 Google Inc.
 *
 * See file CREDITS for list of people who contributed to this
 * project.
 *
 * 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 <assert.h>
#include <endian.h>
#include <stdint.h>

#include "base/time.h"
#include "base/xalloc.h"
#include "drivers/net/net.h"
#include "net/connection.h"
#include "net/ipv4/ipv4.h"
#include "net/ipv4/uip/arp.h"
#include "net/ipv4/uip/uip.h"
#include "net/netboot/dhcp.h"

typedef enum DhcpOpcode
{
	DhcpRequestOp = 1,
	DhcpReplyOp = 2
} DhcpOpcode;

typedef enum DhcpHardwareType
{
	DhcpEthernet = 1
} DhcpHardwareType;

static const uint16_t DhcpServerPort = 67;
static const uint16_t DhcpClientPort = 68;

// Wait for a response for 3 seconds before resending a request.
static const uint64_t DhcpRespTimeoutUs = 3 * 1000 * 1000;

typedef struct __attribute__((packed)) DhcpPacket
{
	uint8_t opcode;
	uint8_t hw_type;
	uint8_t hw_length;
	uint8_t hops;
	uint32_t transaction_id;
	uint16_t seconds;
	uint16_t flags;
	uint32_t client_ip;
	uint32_t your_ip;
	uint32_t server_ip;
	uint32_t gateway_ip;
	uint8_t client_hw_addr[16];
	uint8_t server_name[64];
	uint8_t bootfile_name[128];
	uint32_t cookie;
	uint8_t options[308];
} DhcpPacket;

static const unsigned int DhcpMinPacketSize = 576;
static const unsigned int DhcpMaxPacketSize =
	sizeof(DhcpPacket) + UIP_IPH_LEN + UIP_UDPH_LEN;

static const uint8_t DhcpCookie[] = { 99, 130, 83, 99 };

typedef enum DhcpTags {
	// "Vendor extensions".
	DhcpTagPadding = 0,
	DhcpTagSubnetMask = 1,
	DhcpTagTimeOffset = 2,
	DhcpTagDefaultRouter = 3,
	DhcpTagTimeServer = 4,
	DhcpTagDnsServer = 6,
	DhcpTagLogServer = 7,
	DhcpTagCookieServer = 8,
	DhcpTagPrintServer = 9,
	DhcpTagImpressServer = 10,
	DhcpTagResourceLocationServer = 11,
	DhcpTagHostName = 12,
	DhcpTagBootFileSize = 13,
	DhcpTagMeritDumpFile = 14,
	DhcpTagDomainName = 15,
	DhcpTagSwapServer = 16,
	DhcpTagRootPath = 17,
	DhcpTagExtensionsPath = 18,

	// IP layer paramters per host.
	DhcpTagIpForwarding = 19,
	DhcpTagNonLocalSourceRouting = 20,
	DhcpTagPolicyFilter = 21,
	DhcpTagMaximumDatagramSize = 22,
	DhcpTagDefaultTimeToLive = 23,
	DhcpTagPathMtuAgingTimeout = 24,
	DhcpTagPathMtuPlateauTable = 25,

	// IP later parameters per interface.
	DhcpTagInterfaceMtu = 26,
	DhcpTagAllSubnetsAreLocal = 27,
	DhcpTagBroadcastAddress = 28,
	DhcpTagPerformMaskDiscovery = 29,
	DhcpTagMaskSupplier = 30,
	DhcpTagPerformRouterDiscovery = 31,
	DhcpTagRouterSolicitationAddress = 32,
	DhcpTagStaticRouteOption = 33,

	// Link layer paramters per interface.
	DhcpTagTrailerEncapsulation = 34,
	DhcpTagArpCacheTimeout = 35,
	DhcpTagEthernetEncapsulation = 36,

	// Tcp parameters.
	DhcpTagTcpDefaultTtl = 37,
	DhcpTagTcpKeepaliveInterval = 38,
	DhcpTagTcpKeepaliveGarbage = 39,

	// Application and service parameters.
	DhcpTagNetworkInformationServiceDomain = 40,
	DhcpTagNetworkInformationServers = 41,
	DhcpTagNetworkTimeProtocolServers = 42,
	DhcpTagVendorSpecificInformation = 43,
	DhcpTagNetbiosOverTcpIpNameServer = 44,
	DhcpTagNetbiosOverTcpIpDatagramDistributionServer = 45,
	DhcpTagNetbiosOverTcpIpNodeType = 46,
	DhcpTagNetbiosOverTcpIpScope = 47,
	DhcpTagXWindowSystemFontServer = 48,
	DhcpTagXWindowSystemDisplayManager = 49,

	// Dhcp extensions.
	DhcpTagRequestedIpAddress = 50,
	DhcpTagIpAddressLeaseTime = 51,
	DhcpTagOptionOverload = 52,
	DhcpTagMessageType = 53,
	DhcpTagServerIdentifier = 54,
	DhcpTagParameterRequestList = 55,
	DhcpTagMessage = 56,
	DhcpTagMaximumDhcpMessageSize = 57,
	DhcpTagRenewalTimeValue = 58,
	DhcpTagRebindTimeValue = 59,
	DhcpTagClassIdentifier = 60,
	DhcpTagClientIdentifier = 61,

	DhcpTagEndOfList = 255
} DhcpTags;

typedef enum DhcpState
{
	DhcpInit,
	DhcpRequesting,
	DhcpBound
} DhcpState;

typedef enum DhcpMessageType
{
	DhcpNoMessageType = 0,
	DhcpDiscover = 1,
	DhcpOffer = 2,
	DhcpRequest = 3,
	DhcpDecline = 4,
	DhcpAck = 5,
	DhcpNak = 6,
	DhcpRelease = 7
} DhcpMessageType;

enum {
	OptionOverloadNone = 0,
	OptionOverloadFile = 1,
	OptionOverloadSname = 2,
	OptionOverloadBoth = 3
};

static void uip_ipaddr_from_int(uip_ipaddr_t *ip, uint32_t num)
{
	uip_ipaddr(ip, num >> 0, num >> 8, num >> 16, num >> 24);
}

static void uip_ipaddr_from_array(uip_ipaddr_t *ip, uint8_t *array)
{
	uip_ipaddr(ip, array[0], array[1], array[2], array[3]);
}

static uint32_t uip_ipaddr_to_int(uip_ipaddr_t ip)
{
	return (uip_ipaddr1(&ip) << 0) | (uip_ipaddr2(&ip) << 8) |
	       (uip_ipaddr2(&ip) << 16) | (uip_ipaddr3(&ip) << 24);
}

typedef int (*DhcpOptionFunc)(uint8_t tag, uint8_t length, uint8_t *value,
			      void *data);

static int dhcp_process_options(DhcpPacket *packet, int overload,
				DhcpOptionFunc func, void *data)
{
	uint8_t *options;
	int size;
	switch (overload) {
	case OptionOverloadNone:
		options = packet->options;
		size = sizeof(packet->options);
		break;
	case OptionOverloadFile:
		options = packet->bootfile_name;
		size = sizeof(packet->bootfile_name);
		break;
	case OptionOverloadSname:
		options = packet->server_name;
		size = sizeof(packet->server_name);
		break;
	case OptionOverloadBoth:
		options = packet->server_name;
		size = sizeof(packet->server_name) +
			sizeof(packet->bootfile_name);
		break;
	default:
		return 1;
	}

	int pos = 0;

	int new_overload = 0;

	while (pos < size) {
		uint8_t tag = options[pos++];
		if (pos >= size)
			return 1;
		uint8_t length = options[pos++];
		if (pos >= size)
			return 1;
		uint8_t *value = &options[pos];
		pos += length;
		if (pos >= size)
			return 1;

		// Handle "end of list" and "option overload" options.
		if (tag == DhcpTagEndOfList)
			break;
		if (tag == DhcpTagOptionOverload) {
			if (length < 1)
				return 1;
			new_overload = *value;
			continue;
		}

		// Otherwise, use the provided callback.
		assert(func);
		if (func(tag, length, value, data))
			return 1;
	}

	if (!overload && new_overload)
		return dhcp_process_options(packet, new_overload, func, data);

	return 0;
}

static int dhcp_get_type(uint8_t tag, uint8_t length, uint8_t *value,
			 void *data)
{
	if (tag != DhcpTagMessageType)
		return 0;
	if (length < 1)
		return 1;

	DhcpMessageType *type = (DhcpMessageType *)data;
	*type = (DhcpMessageType)*value;

	return 0;
}

static int dhcp_get_server(uint8_t tag, uint8_t length, uint8_t *value,
			   void *data)
{
	if (tag != DhcpTagServerIdentifier)
		return 0;
	if (length < sizeof(uint32_t))
		return 1;

	memcpy(data, value, sizeof(uint32_t));

	return 0;
}

static int dhcp_apply_options(uint8_t tag, uint8_t length, uint8_t *value,
			      void *data)
{
	switch (tag) {
	case DhcpTagSubnetMask:
		if (length < 4) {
			printf("Truncated subnet mask.\n");
			return 1;
		} else {
			uip_ipaddr_t netmask;
			uip_ipaddr_from_array(&netmask, value);
			uip_setnetmask(&netmask);
			return 0;
		}
	case DhcpTagDefaultRouter:
		if (length < 4) {
			printf("Truncated routers.\n");
			return 1;
		} else {
			uip_ipaddr_t router;
			uip_ipaddr_from_array(&router, value);
			uip_setdraddr(&router);
			return 0;
		}
	}
	return 0;
}

int dhcp_check_for_reply(NetConOps *con, DhcpState state,
			 DhcpPacket *in, DhcpPacket *out)
{
	size_t incoming;
	if (netcon_incoming(con, &incoming) || !incoming)
		return 0;

	// If the packet is too big, ignore it.
	if (incoming > sizeof(*in)) {
		netcon_receive(con, NULL, &incoming, incoming);
		return 0;
	}

	// Receive the packet.
	memset(in, 0, sizeof(*in));
	if (netcon_receive(con, in, &incoming, sizeof(*in)))
		return 0;

	// Check for problems in the reply.
	if (in->opcode != DhcpReplyOp ||
		in->hw_length != out->hw_length ||
		memcmp(in->client_hw_addr,
		       out->client_hw_addr,
		       out->hw_length) ||
		(in->transaction_id != out->transaction_id)) {
		return 0;
	}

	if (memcmp(&in->cookie, DhcpCookie, sizeof(DhcpCookie)))
		return 0;

	DhcpMessageType type = DhcpNoMessageType;
	if (dhcp_process_options(in, OptionOverloadNone, &dhcp_get_type, &type))
		return 0;

	switch (state) {
	case DhcpInit:
		if (type != DhcpOffer)
			return 0;
		break;
	case DhcpRequesting:
		if (type != DhcpAck && type != DhcpNak)
			return 0;
		break;
	case DhcpBound:
		// We shouldn't get any more packets once we're bound.
		break;
	}

	// Everything checks out. We have a valid reply.
	return 1;
}

static int dhcp_send_packet(NetConOps *con, DhcpState state, const char *name,
			    DhcpPacket *out, DhcpPacket *in)
{
	// Send the outbound packet.
	printf("Sending %s... ", name);
	if (netcon_send(con, out, sizeof(*out)))
		return 1;
	printf("done.\n");

	// Prepare for the reply.
	printf("Waiting for reply... ");
	uint64_t start = time_us(0);
	while (!dhcp_check_for_reply(con, state, in, out)) {
		if (time_us(start) > DhcpRespTimeoutUs) {
			if (netcon_send(con, out, sizeof(*out)))
				return 1;
			start = time_us(0);
		}
	}
	printf("done.\n");

	return 0;
}

static void dhcp_prep_packet(DhcpPacket *packet, uint32_t transaction_id)
{
	memset(packet, 0, sizeof(*packet));
	packet->opcode = DhcpRequestOp;
	packet->hw_type = DhcpEthernet;
	packet->hw_length = sizeof(uip_ethaddr);
	packet->hops = 0;
	packet->transaction_id = transaction_id;
	packet->seconds = htonw(100);
	assert(sizeof(uip_ethaddr) < sizeof(packet->client_hw_addr));
	memcpy(packet->client_hw_addr, &uip_ethaddr, sizeof(uip_ethaddr));
	memcpy(&packet->cookie, DhcpCookie, sizeof(DhcpCookie));
}

static void dhcp_add_option(uint8_t **options, uint8_t tag, void *value,
			    uint8_t length, int *remaining)
{
	assert(*remaining >= length + 2);
	(*options)[0] = tag;
	(*options)[1] = length;
	if (length) {
		assert(value);
		memcpy(*options + 2, value, length);
	}
	*remaining -= length + 2;
	*options += length + 2;
}

static int dhcp_request_work(NetConOps *con, uip_ipaddr_t *next_ip,
			     uip_ipaddr_t *server_ip, const char **bootfile)
{
	DhcpPacket out, in;
	uint8_t byte;
	uint8_t *options;
	int remaining;
	uint8_t requested[] = { DhcpTagSubnetMask, DhcpTagDefaultRouter };
	assert(DhcpMaxPacketSize >= DhcpMinPacketSize);
	uint16_t max_size = htonw(DhcpMaxPacketSize);
	uint8_t client_id[1 + sizeof(uip_ethaddr)];
	client_id[0] = DhcpEthernet;
	memcpy(client_id + 1, &uip_ethaddr, sizeof(uip_ethaddr));

	// Send a DHCP discover packet.
	dhcp_prep_packet(&out, rand());
	options = out.options;
	remaining = sizeof(out.options);
	byte = DhcpDiscover;
	dhcp_add_option(&options, DhcpTagMessageType, &byte, sizeof(byte),
			&remaining);
	dhcp_add_option(&options, DhcpTagClientIdentifier, client_id,
			sizeof(client_id), &remaining);
	dhcp_add_option(&options, DhcpTagParameterRequestList, requested,
			sizeof(requested), &remaining);
	dhcp_add_option(&options, DhcpTagMaximumDhcpMessageSize,
			&max_size, sizeof(max_size), &remaining);
	dhcp_add_option(&options, DhcpTagEndOfList, NULL, 0, &remaining);
	if (dhcp_send_packet(con, DhcpInit, "DHCP discover", &out, &in))
		return 1;

	// Extract the DHCP server id.
	uint32_t server_id;
	if (dhcp_process_options(&in, OptionOverloadNone, &dhcp_get_server,
				 &server_id)) {
		printf("Failed to extract server id.\n");
		return 1;
	}

	// We got an offer. Request it.
	dhcp_prep_packet(&out, rand());
	options = out.options;
	remaining = sizeof(out.options);
	byte = DhcpRequest;
	dhcp_add_option(&options, DhcpTagMessageType, &byte, sizeof(byte),
			&remaining);
	dhcp_add_option(&options, DhcpTagClientIdentifier, client_id,
			sizeof(client_id), &remaining);
	dhcp_add_option(&options, DhcpTagRequestedIpAddress, &in.your_ip,
			sizeof(in.your_ip), &remaining);
	dhcp_add_option(&options, DhcpTagParameterRequestList, requested,
			sizeof(requested), &remaining);
	dhcp_add_option(&options, DhcpTagMaximumDhcpMessageSize,
			&max_size, sizeof(max_size), &remaining);
	dhcp_add_option(&options, DhcpTagServerIdentifier,
			&server_id, sizeof(server_id), &remaining);
	dhcp_add_option(&options, DhcpTagEndOfList, NULL, 0, &remaining);
	if (dhcp_send_packet(con, DhcpRequesting, "DHCP request", &out, &in))
		return 1;

	DhcpMessageType type;
	if (dhcp_process_options(&in, OptionOverloadNone, &dhcp_get_type,
				 &type)) {
		printf("Failed to extract message type.\n");
		return 1;
	}
	if (type == DhcpNak) {
		printf("DHCP request nak-ed by the server.\n");
		return 1;
	}

	// The server acked, completing the transaction.

	// Apply the settings.
	if (dhcp_process_options(&in, OptionOverloadNone,
				 &dhcp_apply_options, NULL)) {
		return 1;
	}

	int bootfile_size = sizeof(in.bootfile_name) + 1;
	char *file = xmalloc(bootfile_size);
	file[bootfile_size - 1] = 0;
	memcpy(file, in.bootfile_name, sizeof(in.bootfile_name));
	*bootfile = file;
	uip_ipaddr_from_int(next_ip, in.server_ip);
	uip_ipaddr_from_int(server_ip, server_id);

	uip_ipaddr_t my_ip;
	uip_ipaddr_from_int(&my_ip, in.your_ip);
	uip_sethostaddr(&my_ip);

	return 0;
}

int dhcp_request(NetDevice *dev, uip_ipaddr_t *next_ip,
		 uip_ipaddr_t *server_ip, const char **bootfile)
{
	// Set up the UDP connection.
	uip_ipaddr_t addr;
	uip_ipaddr(&addr, 255,255,255,255);
	Ipv4UdpCon *con = new_ipv4_udp_con(dev, addr, DhcpServerPort);
	int ret = netcon_bind(&con->ops, DhcpClientPort);
	// We shouldn't attempt to DHCP request if we failed to bind.
	ret = ret || dhcp_request_work(&con->ops, next_ip, server_ip, bootfile);
	// Failure should not prevent closing the connection.
	ret = netcon_close(&con->ops) || ret;
	free(con);
	return ret;
}

static int dhcp_release_work(NetConOps *con)
{
	DhcpPacket release;

	// Prepare the DHCP release packet.
	dhcp_prep_packet(&release, rand());
	uip_ipaddr_t my_ip;
	uip_gethostaddr(&my_ip);
	release.client_ip = uip_ipaddr_to_int(my_ip);

	uint8_t *options = release.options;
	int remaining = sizeof(release.options);
	uint8_t byte = DhcpRelease;
	dhcp_add_option(&options, DhcpTagMessageType, &byte, sizeof(byte),
			&remaining);
	dhcp_add_option(&options, DhcpTagEndOfList, NULL, 0, &remaining);

	return netcon_send(con, &release, sizeof(release));
}

int dhcp_release(NetDevice *dev, uip_ipaddr_t server_ip)
{
	// Set up the UDP connection.
	Ipv4UdpCon *con = new_ipv4_udp_con(dev, server_ip, DhcpServerPort);
	int ret = netcon_bind(&con->ops, DhcpClientPort);
	// We shouldn't attempt to DHCP release if we failed to bind.
	ret = ret || dhcp_release_work(&con->ops);
	// Failure should not prevent closing the connection.
	ret = netcon_close(&con->ops) || ret;
	free(con);
	return ret;
}