//------------------------------------------------------------------------------
// File:          main.c
// Function:      Bootloader main program
// Supported chip(s):
//    - AT91SAM7XC128
//    - AT91SAM7XC256
// Supported toolchain(s):
//    - IAR Embedded Workbench
// Date created:  28 April 2006
// Created by:    JJo
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
// Includes
//------------------------------------------------------------------------------

#include "common.h"
#include "debug.h"
#include "media.h"
#include "trigger.h"
#include "memory.h"
#include "communication.h"
#include "timing.h"
#include "encryption.h"
#include "safety/partitioning/partitioning.h"

//------------------------------------------------------------------------------
// Functions
//------------------------------------------------------------------------------

/**
 * Name:      bootloader_load
 * Purpose:   Receive and upgrade a new firmware
 * Input:
 *  - Application starting address
 * Output:    Application ending address
 */
void * bootloader_load(void * vStart) {
  
  // Local variables
  char page[MEMORY_PAGE_SIZE];
  unsigned int bytes;
  char * pCurrent;

  // Initialize address and notify host
  pCurrent = (char *) vStart;
  communication_ready();
  
  debug_printf("Bootloader: Upgrading ");
  do {
    // Receive a page
    bytes = communication_receive(page);
    debug_printf(".");

#ifdef USE_MEMORY_PARTITIONING
    // Check for an overflow
    if (((unsigned int) pCurrent + bytes - (unsigned int) vStart) >
        (unsigned int) REGION_SIZE) {

      debug_printf("Partitioning: Firmware too big.\n");
      return vStart;
    }
#endif

    // End of firmware ?
    if (bytes > 0) {

      // Pad page with 0xFF and write it
      while (bytes < MEMORY_PAGE_SIZE) {
  
        page[bytes] = 0xFF;
        bytes++;
      }

      // Decrypt page (when USE_ENCRYPTION is defined)
      if (encryption_decrypt((unsigned char *) page, (unsigned char *) page, bytes) != OK) {

        debug_printf("Bootloader: Error while trying to decrypt page !\n");
      }

      memory_write(pCurrent, page);
  
      // Loop progression
      pCurrent += MEMORY_PAGE_SIZE;
    }
  } while (bytes > 0);

  debug_printf("\n");
  debug_printf("Bootloader: Loading done.\n");

  return pCurrent;
}

//------------------------------------------------------------------------------
// Main program
//
// Returns application starting address
//------------------------------------------------------------------------------
unsigned int main(void) {

#ifdef USE_MEMORY_PARTITIONING
  // Local variables
  regionInfo_t regionInfo;

  partitioning_readInfo(REGION_INFO_START, &regionInfo);
#endif
  
  // Initialize debug port and display start message
  debug_init();
  debug_printf("\n\n--------------------------------\n");
  debug_printf("Safe and Secure Bootloader\n");
  debug_printf("--------------------------------\n");

  /**
   * Encryption timing measurement
   */
//#define ENCRYPTION_TIMING
#ifdef ENCRYPTION_TIMING
  unsigned char page[MEMORY_PAGE_SIZE];

  timing_init();
  debug_printf("Timing: Starting decryption timing measurement ...\n");
  
  // 1k
  timing_start();

  for (unsigned int i=0; i < 100; i++) {

    encryption_init();
    for (unsigned int j=0; j < 131072; j += MEMORY_PAGE_SIZE) {
      
      encryption_decrypt(page, page, MEMORY_PAGE_SIZE);
    }
    encryption_cleanup();
  }

  timing_stop();

  debug_printf("Timing: Decrypting 128 kByte = %d 1/10ms", timing_get() / 100);

  while (1);
#endif

  /**
   * Bootloader initialization
   */
  debug_printf("Bootloader: Initializing ...\n");

  timing_init();
  trigger_init();
  memory_init();
  media_init();
  communication_init();
  encryption_init();

  debug_printf("Bootloader: Initialization done.\n");

  // Boot region lock
#ifdef USE_BOOT_REGION_LOCKING
  
  memory_unlock(APP_START_ADDRESS, NULL);
  debug_printf("Bootloader: Locking boot region ...\n");

  if (memory_lock(MEMORY_START_ADDRESS, APP_START_ADDRESS)) {
  
    debug_printf("Bootloader: Failed to lock boot region !\n");
  }
  else {

    debug_printf("Bootloader: Boot region locked.\n");
  }
#endif
  
  // Check trigger condition
  debug_printf("Bootloader: Checking trigger ...\n");

  if (trigger_poll()) {
      
    /**
     * Firmware loading
     */
#ifndef USE_MEMORY_PARTITIONING
    bootloader_load(APP_START_ADDRESS);
  }
#else
    // Erase old signature
    regionInfo.signatureB = 0;
    partitioning_writeInfo(REGION_INFO_START, &regionInfo);

    // Load firmware in region B and compute its length
    regionInfo.lengthB = (unsigned int) bootloader_load(REGION_B_START);
    regionInfo.lengthB -= (unsigned int) REGION_B_START;

    // Compute region B signature
    debug_printf("Partitioning: Calculating region B signature ...\n");
    regionInfo.signatureB = partitioning_sign(REGION_B_START, regionInfo.lengthB);
    partitioning_writeInfo(REGION_INFO_START, &regionInfo);
    debug_printf("Partitioning: Signing done.\n");
  }
  
  // Check if region A == region B
  debug_printf("Partitioning: Checking signatures ...\n");
  if (partitioning_check(regionInfo.signatureA, regionInfo.signatureB) != OK) {

    debug_printf("Partitioning: Signatures are different.\n");

    // Check that B is valid
    debug_printf("Partitioning: Checking that region B is valid ...\n");

    if (partitioning_isRegionValid(REGION_B_START,
                                   regionInfo.signatureB,
                                   regionInfo.lengthB) == OK) {

      debug_printf("Partitioning: Region B is valid.\n");

      // Copy B over A
      char * pCurrent = (char *) REGION_A_START;
      char * pData = (char *) REGION_B_START;
      char * pEnd = (char *) ((unsigned int) REGION_A_START + regionInfo.lengthB);

      debug_printf("Partitioning: Copying region B over region A ");

      while (pCurrent != pEnd) {

        flash_write(pCurrent, pData);
        pCurrent += MEMORY_PAGE_SIZE;
        pData += MEMORY_PAGE_SIZE;

        debug_printf(".");
      }

      regionInfo.signatureA = partitioning_sign(REGION_A_START, regionInfo.lengthB);
      regionInfo.lengthA = regionInfo.lengthB;
      partitioning_writeInfo(REGION_INFO_START, &regionInfo);      

      debug_printf("\nPartitioning: Copying done.\n");
    }
    else {

      debug_printf("Partitioning: Region B is not valid.\n");
    }
  }
  else {

    debug_printf("Partitioning: Signatures are identical.\n");
  }
  
  // Check if region A is valid
  debug_printf("Partitioning: Checking that region A is valid ...\n");

  if (partitioning_isRegionValid(REGION_A_START,
                                 regionInfo.signatureA,
                                 regionInfo.lengthA) != OK) {

    debug_printf("Partitioning: Region A is not valid.\n");
    debug_printf("Partitioning: No valid region to boot on.\n");
    while (1);
  }
  else {

    debug_printf("Partitioning: Region A is valid.\n");
  }
#endif

  /**
   * Bootloader cleanup
   */
  debug_printf("Bootloader: Cleaning up ...\n");
  
  media_cleanup();
  communication_cleanup();
  memory_cleanup();
  encryption_cleanup();
  trigger_cleanup();
  timing_cleanup();
  
  debug_printf("Bootloader: Cleanup done.\n");
 
  /**
   * Run application after return
   */
  debug_printf("Bootloader: Running application ...\n");

  // Cleanup debug port
  debug_cleanup();

  return (unsigned int) APP_START_ADDRESS;
}
