/*
 * Copyright 2016 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 <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "base/xalloc.h"
#include "base/time.h"
#include "boot/boot.h"
#include "drivers/net/net.h"
#include "net/connection.h"
#include "net/gigaboot/gigaboot.h"
#include "net/ipv6/ipv6.h"

enum {
	GigabootMagic = 0xAA774217,
	GigabootAdvertise = 0x77777777,
};

enum {
	GigabootError = 0x80000000,
	GigabootErrorBadCmd = 0x80000001,
	GigabootErrorBadParam = 0x80000002,
	GigabootErrorTooLarge = 0x80000003,
	GigabootErrorBadFile = 0x80000004,
};

enum {
	GigabootVersion_1_0 = 0x0001000,
	GigabootVersion_1_1 = 0x0001010,
	GigabootVersion_Current = GigabootVersion_1_1,
};

enum {
	// Used to acknowledge receipt of a command.
	GigabootAck = 0,
	// Used to acknowledge the receipt of all a file's data.
	GigabootFileReceived = 0x70000001,
	// Doesn't do anything.
	GigabootCommandCommand = 1,
	// Start receiving a file.
	GigabootCommandSendFile = 2,
	// Receive some data for the file started with SendFile.
	GigabootCommandData = 3,
	// Boot using the data received so far.
	GigabootCommandBoot = 4,
	// Receive the last chunk of data for the file started with SendFile.
	GigabootCommandLastData = 11,
};

typedef struct {
	uint32_t magic;
	uint32_t cookie;
	uint32_t command;
	uint32_t arg;
} GigabootHeader;

typedef struct {
	GigabootHeader hdr;
	uint8_t data[0];
} GigabootMessage;


typedef struct {
	void *data;
	size_t size; // Max size of the buffer.
	size_t offset; // Write pointer.
} GigabootBuffer;

static GigabootBuffer gigaboot_kernel = {
	.data = (void *)(uintptr_t)CONFIG_KERNEL_START,
	.size = CONFIG_KERNEL_SIZE,
};

static char gigaboot_cmdline_buf[4096];
static GigabootBuffer gigaboot_cmdline = {
	.data = gigaboot_cmdline_buf,
	.size = sizeof(gigaboot_cmdline_buf),
};

static GigabootBuffer gigaboot_ramdisk;


typedef struct {
	NetConOps *con;

	GigabootHeader last;
	GigabootHeader ack;

	GigabootBuffer *buffer;

	GigabootMessage *message;
	size_t message_len;
} GigabootState;


/*
 * Receive messages and act on commands.
 */

static GigabootBuffer *gigaboot_prepare_buffer(const char *name, size_t size)
{
	if (!strcmp(name, "kernel.bin"))
		return &gigaboot_kernel;
	if (!strcmp(name, "cmdline"))
		return &gigaboot_cmdline;
	if (!strcmp(name, "ramdisk.bin")) {
		if (gigaboot_ramdisk.size > size)
			return &gigaboot_ramdisk;

		if (gigaboot_free_buffer(gigaboot_ramdisk.data,
					 gigaboot_ramdisk.size)) {
			return NULL;
		}
		gigaboot_ramdisk.data = NULL;
		gigaboot_ramdisk.size = 0;

		if (!size)
			size = 512 * 1024 * 1024;
		void *data = gigaboot_allocate_buffer(&size);
		if (!data)
			return NULL;
		gigaboot_ramdisk.data = data;
		gigaboot_ramdisk.size = size;

		return &gigaboot_ramdisk;
	}
	return NULL;
}

static void gigaboot_recv_message(GigabootState *state, GigabootMessage *msg,
				  size_t size)
{
	if (state->last.cookie == msg->hdr.cookie &&
	    state->last.command == msg->hdr.command &&
	    state->last.arg == msg->hdr.arg) {
		// Host must have missed the ack. Resend.
		netcon_send(state->con, &state->ack, sizeof(state->ack));
		return;
	}

	memcpy(&state->last, &msg->hdr, sizeof(GigabootHeader));

	size_t data_size = size - sizeof(GigabootHeader);

	state->ack.cookie = msg->hdr.cookie;
	state->ack.arg = msg->hdr.arg;

	int send_ack = 1;

	switch (msg->hdr.command) {
	case GigabootCommandCommand:
		if (data_size == 0)
			return;
		break;

	case GigabootCommandSendFile:
		if (data_size == 0)
			return;

		msg->data[data_size - 1] = 0;
		for (int i = 0; i < data_size - 1; i++) {
			if (msg->data[i] < ' ' || msg->data[i] > 127)
				msg->data[i] = '.';
		}
		state->buffer = gigaboot_prepare_buffer(
			(const char *)msg->data, msg->hdr.arg);
		if (state->buffer) {
			state->buffer->offset = 0;
			state->ack.arg = msg->hdr.arg;
			printf("gigaboot: Receive file '%s'...\n",
			       (char *)msg->data);
		} else {
			state->ack.command = GigabootErrorBadFile;
			printf("gigaboot: Rejected file '%s'...\n",
			       (char *)msg->data);
		}
		break;

	case GigabootCommandData:
	case GigabootCommandLastData:
		if (!state->buffer)
			return;
		if (msg->hdr.arg != state->buffer->offset) {
			printf("gigaboot: received chunk at offset %d but "
			       "current offset is %zu\n",
			       msg->hdr.arg, state->buffer->offset);
			state->ack.arg = state->buffer->offset;
			state->ack.command = GigabootAck;
		} else if (state->buffer->offset + data_size >
			   state->buffer->size) {
			state->ack.command = GigabootErrorTooLarge;
			state->ack.arg = msg->hdr.arg;
		} else {
			memcpy((uint8_t *)state->buffer->data +
				state->buffer->offset,
			       msg->data, data_size);
			state->buffer->offset += data_size;
			if (msg->hdr.command == GigabootCommandLastData) {
				state->ack.command = GigabootFileReceived;
			} else {
				state->ack.command = GigabootAck;
				send_ack = 0;
			}
		}
		break;

	case GigabootCommandBoot:
	{
		printf("netboot: Boot Kernel...\n");

		if (!gigaboot_kernel.offset) {
			state->ack.command = GigabootErrorBadParam;
			break;
		}

		// Send an ack here since we're (probably) not coming back.
		netcon_send(state->con, &state->ack, sizeof(state->ack));
		boot(gigaboot_kernel.data, gigaboot_cmdline.data, NULL,
		     NULL, gigaboot_ramdisk.data, gigaboot_ramdisk.offset);
		return;
	}

	default:
		state->ack.command = GigabootErrorBadCmd;
		state->ack.arg = 0;
	}

	if (send_ack)
		netcon_send(state->con, &state->ack, sizeof(state->ack));
}


/*
 * Poke the host to get it to start sending commands.
 */

static void gigaboot_advertise(GigabootState *state)
{
	static const char ad_data[] =
		"version\01.1\0"
		"serialno\0unknown\0"
		"board\0unknown\0";

	static struct {
		GigabootHeader header;
		uint8_t data[sizeof(ad_data)];
	} ad_message = {
		.header = {
			.magic = GigabootMagic,
			.cookie = 0,
			.command = GigabootAdvertise,
			.arg = GigabootVersion_Current,
		},
	};
	if (!ad_message.data[0])
		memcpy(ad_message.data, ad_data, sizeof(ad_data));

	netcon_send(state->con, &ad_message, sizeof(ad_message));
}


/*
 * Try running the gigaboot protocol on a given connection.
 */

static void gigaboot_attempt(NetConOps *con)
{
	GigabootState state;
	memset(&state, 0, sizeof(state));
	state.con = con;
	state.ack.magic = GigabootMagic;

	gigaboot_kernel.offset = 0;
	gigaboot_cmdline.offset = 0;
	gigaboot_ramdisk.offset = 0;

	gigaboot_advertise(&state);

	uint64_t start = time_us(0);
	while (time_us(start) < 1000000) {
		size_t incoming;
		if (netcon_incoming(con, &incoming)) {
			printf("Incoming failed.\n");
			return;
		}
		if (!incoming)
			continue;
		if (incoming < sizeof(GigabootHeader))
			return;

		if (incoming > state.message_len) {
			free(state.message);
			state.message = xmalloc(incoming);
			state.message_len = incoming;
		}

		netcon_receive(con, state.message, &incoming, incoming);
		gigaboot_recv_message(&state, (GigabootMessage *)state.message,
				      incoming);
		start = time_us(0);
	}

	free(state.message);
}


/*
 * Keep trying to gigaboot on each available network device in a loop.
 */

void gigaboot(void)
{
	NetDevice *dev = NULL;
	while (1) {
		dev = net_scan_for_link(dev);
		if (dev) {
			Ipv6UdpCon *con =
				new_ipv6_udp_con(dev, &Ipv6LlAllNodes,
						 33331);
			netcon_bind(&con->ops, 33330);
			gigaboot_attempt(&con->ops);
			netcon_close(&con->ops);
			free(con);
		} else {
			net_find_devices();
		}
	}
}