//----------------------------------------------------------------------------------
// <copyright file="boot.c" company="LiquidRobotics">
//	Copyright (c) Liquid Robotics Corporation.  All rights reserved.
// </copyright>
//
// <summary>
// 	This module imports a binary image of the boot loader and places it into the
//  boot section to be correctly located.  It relies on PIBBOOT.BIN being available as
//  well.  There are also some routines to check the validity of the loaded image.
// </summary>
//
// <owner>Mike Cookson</owner>
//----------------------------------------------------------------------------------

#include "includes.h"
#include <string.h>
#include "avr/interrupt.h"
#include "avr/pgmspace.h"
#include "avr/sleep.h"
#include "avr/wdt.h"
#include "util/atomic.h"
#include "boot.h"
#include "BootEEPROM.h"
#include "MD5.h"
#include "avr\boot.h"

#ifndef DIM
#define DIM(a) (sizeof(a)/sizeof((a)[0]))
#endif

static BOOL corrupt = TRUE;				// is the load of the boot vars corrupt?
static BOOT_EEPROM_VARS bootVars  __attribute__((section(".xdata")));	// the boot vars loaded

// ---------------------------------------------
//
// Performs sanity check of program space and
// configuration at startup
//
// ---------------------------------------------

void BootCheckProgramImage()
{
#ifndef OMIT_BOOT_CHECKING

    // pull a character from the buffer atomically (if the buffer is not empty)
	OS_CPU_SR cpu_sr;
	OS_ENTER_CRITICAL();

	//uint8_t lowFuse = boot_lock_fuse_bits_get(GET_LOW_FUSE_BITS);
	uint8_t highFuse = boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS);
	//uint8_t extFuse = boot_lock_fuse_bits_get(GET_EXTENDED_FUSE_BITS);
	//uint8_t lockBits = boot_lock_fuse_bits_get(GET_LOCK_BITS);

	OS_EXIT_CRITICAL();


	// skip this if the boot loader is disabled, otherwise, we will just
	// keep rebooting
	if (highFuse & 1)
		return;


	// read the image from the EEPROM
	corrupt = !ReadBootEEPROM(&bootVars);

	BOOL success = FALSE;
	uint8_t failStep = 0;

	if (++failStep, !corrupt)
	if (++failStep, bootVars.selectedImage.imageLength)
	if (++failStep, bootVars.selectedImage.imageLength <= 0x3E000L)
	{
		++failStep;
		success = BootCheckProgramMD5(bootVars.selectedImage.imageLength, bootVars.selectedImage.imageMD5,
						bootVars.fallbackImage.imageMD5);
		WriteBootEEPROM(&bootVars);
	}

	if (!success)
		BootRestart(TRUE);
#endif
}

// ---------------------------------------------
//
// Marks the program image as no longer
// tentative
//
// ---------------------------------------------

void BootClearTentativeImage()
{
	if (BootImageIsTentative())
	{
		bootVars.tentativeImage = 0;
		WriteBootEEPROM(&bootVars);
	}
}

// ---------------------------------------------
//
// Determines if the boot image is tentative,
// meaning that it has not been vetted
//
// ---------------------------------------------

BOOL BootImageIsTentative()
{
	return !corrupt && bootVars.tentativeImage && !bootVars.useFallback;
}

// ---------------------------------------------
//
// Tests the configuration's restartedByWatchdog
// flag: the boot loader would have already
// cleared the hardware flag by the time the
// program loads
//
// ---------------------------------------------

BOOL BootCausedByWatchdog()
{
	return corrupt || bootVars.restartedByWatchdog;
}

// ---------------------------------------------
//
// Checks the program image vs the expected CRC
// and MD5
//
// ---------------------------------------------

BOOL BootCheckProgramMD5(uint32_t imageLength, uint8_t* pExpectedMD5, uint8_t* pCalculatedMD5)
{
	MD5_CTX md5ctx;						// the context used for incrementally computing MD5 digest
	uint32_t addr = 0;					// program address to read from
	uint8_t buffer[100];				// buffer to assemble for MD5

	// get the MD5 of the program space
	MD5Init(&md5ctx);

	while(addr < imageLength)
	{
		uint8_t* ptr = buffer;
		uint16_t blockSize = ((imageLength-addr) > DIM(buffer)) ? DIM(buffer) : ((uint16_t) (imageLength-addr));
		uint16_t remaining = blockSize;

		while(remaining--)
			*ptr++ = __ELPM(addr++);

		MD5Update(&md5ctx, (unsigned char*) buffer, blockSize);
	}
	
	// finalizes the MD5 calculation
	MD5Final(&md5ctx);

	// save a copy of the MD5 to the output location
	memcpy(pCalculatedMD5, md5ctx.digest, sizeof(md5ctx.digest));

	// fail in the case of invalid CRC
	BOOL success = FALSE;

	if (!memcmp(pExpectedMD5, pCalculatedMD5, sizeof(md5ctx.digest)))
		success = TRUE;

	// true if the digest matches, false if it's different
	return success;
}

// ---------------------------------------------
//
// Restarts the system, indicating if this is an
// intentional restart
//
// ---------------------------------------------

void BootRestart(BOOL intentional)
{
	// write the EEPROM flag to indicate if this is an intentional restart
	if (!corrupt)
	{
		bootVars.intentionalRestart = intentional ? 1 : 0;
		WriteBootEEPROM(&bootVars);
	}

	cli();						// turn off interrupts to prevent watchdog retriggers
	wdt_enable(WDTO_60MS);		// use a short timeout
	
	// just wait here until the watchdog restarts the system
	while(1)
		sleep_mode();
}

// ---------------------------------------------
//
// Updates the boot configuration to request
// a new program load on restart (call BootRestart
// to actually restart)
//
// ---------------------------------------------

void BootSelectProgram(BOOL newSelectedImage)
{
	if (!corrupt)
	{
		memset(bootVars.selectedImage.fileName, '\0', sizeof(bootVars.selectedImage.fileName));

		if (newSelectedImage)
			strcpy_P(bootVars.selectedImage.fileName, PSTR("$NEWPROG"));
		else
			strcpy_P(bootVars.selectedImage.fileName, PSTR("$DEFAULT"));

		WriteBootEEPROM(&bootVars);
	}
}

// this stub ensures that the boot loader is also included in the build
void __attribute__ ((section (".boot"))) __attribute__((naked)) boot_loader(void)
{
	asm (".incbin \"../pibboot.bin\"\n\t");
}