bootloader/tsb/tsb_asm.cpp

// TinySafeBoot on libavr - tier 3: full feature parity in the 512-byte boot
// section, in C++ except where the C ABI itself is the cost.
//
// The complete TinySafeBoot feature set - watchdog-reset bail, one-wire
// half-duplex UART, a config-page activation timeout, the password gate,
// emergency erase, and config/flash/EEPROM read-write - inside the 512-byte
// BOOTSZ=11 section the hand-written oracle occupies (oracle/README.md holds
// what each tier measures, in one table rather than four). The
// body is the tricks tier's C++ (same register protocol, same structure - see
// tsb_tricks.cpp, including the global-register miscompile rules) with exactly
// two routines kept in assembly, the two whose remaining cost *is* the calling
// convention:
//
//   rx     the bounded receive: C++ must re-floor the timeout window on every
//          call (the global-register-store miscompile) and split it across
//          call-saved registers; the asm keeps the oracle's X-register nested
//          countdown.
//   store  the page-store loop: C++ cannot hold the receive byte pair and the
//          walked Z pointer across the rx calls without call-saved staging
//          (push/pop + a Y->Z copy per word); the asm calls rx knowing exactly
//          which registers it touches and walks Z live across the whole page.
//
// Everything else - bring-up, activation, password gate, emergency erase,
// dispatch, every SPM/EEPROM/flash primitive, every geometry/baud/info
// constant - is C++ on libavr, and the two asm routines splice into the same
// global-register protocol the C++ uses (g_addr in Y, g_cnt in r16, g_window
// in r7, g_receiving in r6), so calls cross the boundary with no marshalling.
//
// The wire protocol is strict request/response, which is what makes the shared
// line safe: the device drives it only between a received command and its
// reply, and releases it (RXEN0 only) whenever it waits.

#include <libavr/libavr.hpp>

#include <avr/io.h> // SPMCSR's I/O address for the page-store routine

using namespace avr::literals;
namespace spm = avr::spm;
namespace ee = avr::eeprom;
namespace hw = avr::hw;

namespace tsb {
namespace {

// The loader is purely polled - it never enables interrupts - so every SPM and
// EEPROM lock folds to nothing under this posture.
constexpr auto off = avr::irq::guard_policy::unused;

// Strict request/response: every SPM operation is waited out before the next
// byte moves, so no flash operation is ever in flight at an EEPROM access -
// the write procedure's step 2 has nothing to guard, the omission the
// datasheet grants (DS40002061B section 8.6.3).
constexpr auto no_spm = ee::spm_interlock::omitted;

constexpr std::uint8_t confirm = '!';
constexpr std::uint8_t request = '?';
constexpr std::uint8_t knock = '@';

constexpr auto clock = 16_MHz;

// Boot geometry for the 512 B boot section (BOOTSZ=11); the page size and the
// flash/EEPROM extents are the chip database's to know. app_end is the config
// page (TSB's LASTPAGE), one page below the boot section.
constexpr std::uint16_t page = spm::page_bytes;
constexpr std::uint16_t boot_bytes = 512;
constexpr std::uint16_t app_end = spm::flash_bytes - boot_bytes - page;
constexpr std::uint16_t eeprom_end = avr::hw::db.mem.eeprom_size - 1;

// Lockout-proof floor for the activation window: the oracle's F_CPU/1MHz, so
// it follows the clock rather than restating it.
constexpr auto act_min = static_cast<std::uint8_t>(clock.hz / 1'000'000);
// Post-activation window: the host gets seconds, not milliseconds, mid-session.
constexpr std::uint8_t comm_window = 200;

constexpr std::uint16_t build_date = 26 * 512 + 7 * 32 + 20;

// Fixed 115200 8N1; the library solves UBRR + U2X from clock and baud.
constexpr auto baud = avr::uart::solve_baud(clock, 115200_Bd, 8, avr::uart::parity::none);

// One bit time on the wire: the turn-around a shared-line peer needs to stop
// driving before this one starts. Derived from the solved rate, so it follows
// the link rather than a count measured against one.
constexpr auto guard_cycles = static_cast<std::uint32_t>(clock.hz / baud.actual);

// The 16-byte device-info block, streamed out on activation.
// clang-format off
[[gnu::progmem]] constexpr auto info = std::to_array<std::uint8_t>({
    'T', 'S', 'B',
    build_date & 0xFF, build_date >> 8,
    0xF3,             // status: native-UART fixed-baud lineage
    avr::hw::db.signature[0], avr::hw::db.signature[1], avr::hw::db.signature[2],
    page / 2,         // page size in words
    (app_end / 2) & 0xFF, (app_end / 2) >> 8,
    eeprom_end & 0xFF, eeprom_end >> 8,
    0xAA, 0xAA,
});
// clang-format on

register std::uint16_t g_addr asm("r28");
register std::uint8_t g_cnt asm("r16");
register std::uint8_t g_window asm("r7");
register std::uint8_t g_receiving asm("r6");

const std::uint8_t *flash_ptr(std::uint16_t addr)
{
	return reinterpret_cast<const std::uint8_t *>(addr);
}

// Bounded byte receive (asm 1 of 2): release the one-wire line on a direction
// change, poll RXC0 under the oracle's nested X-register countdown seeded from
// g_window (floored against lockout), byte or 0-on-silence in r24. Z survives
// - the property the store's word loop rides on.
[[gnu::noinline, gnu::noclone]] std::uint8_t rx()
{
	std::uint8_t byte;
	asm volatile("	tst	%[dir]			\n\t" // already receiving? keep the line released
	             "	brne	1f			\n\t"
	             "	ldi	%[b], %[rxen]		\n\t" // the receiver alone: release the line and listen
	             "	sts	%[ucsr0b], %[b]		\n\t"
	             "	ser	%[b]			\n\t"
	             "	mov	%[dir], %[b]		\n\t"
	             "1:	mov	r27, %[to]		\n\t" // outer countdown high byte = window
	             "	ori	r27, %[actmin]		\n\t" // lockout-proof floor
	             "	clr	r26			\n\t"
	             "2:	ser	%[b]			\n\t"
	             "3:	lds	%[b], %[ucsr0a]		\n\t"
	             "	sbrc	%[b], %[rxc]		\n\t"
	             "	rjmp	4f			\n\t"
	             "	dec	%[b]			\n\t"
	             "	brne	3b			\n\t"
	             "	sbiw	r26, 1			\n\t"
	             "	brcc	2b			\n\t"
	             "	clr	%[b]			\n\t" // silence -> 0, which no compare accepts
	             "	rjmp	5f			\n\t"
	             "4:	lds	%[b], %[udr0]		\n\t"
	             "5:					\n\t"
	             : [b] "=&d"(byte), [dir] "+r"(g_receiving)
	             : [to] "r"(g_window), [actmin] "M"(act_min), [rxen] "M"(hw::ucsr0b::rxen0(1).value),
	               [rxc] "I"(std::countr_zero(hw::ucsr0a::rxc0.mask)), [ucsr0a] "n"(hw::ucsr0a::addr),
	               [ucsr0b] "n"(hw::ucsr0b::addr), [udr0] "n"(hw::udr0::addr)
	             : "r26", "r27", "cc");
	return byte;
}

// One-wire transmit: take the line (TXEN0 alone) on a direction change with a
// turn-around guard, put the byte out, hold the line until the whole frame is
// out (TXC0, not UDRE0), W1C TXC0 by storing the sampled status back (keeps
// U2X0). Plain C++ - it compiles *smaller* than the oracle's routine.
[[gnu::noinline, gnu::noclone]] void tx(std::uint8_t byte)
{
	if (g_receiving) {
		g_receiving = 0;
		hw::ucsr0b::write(hw::ucsr0b::txen0(1));
		avr::delay::cycles<guard_cycles>();
	}
	hw::udr0::write(byte);
	std::uint8_t status;
	do {
		status = hw::ucsr0a::read();
	} while (!(status & hw::ucsr0a::txc0(1).value));
	hw::ucsr0a::write(status);
}

// '?', then hand back the host's reply for the callers' one-byte compare.
[[gnu::noinline, gnu::noclone]] std::uint8_t rcnf()
{
	tx(request);
	return rx();
}

// One flash byte <- [g_addr++] (the advance right before ret - the
// global-register rule, see tsb_tricks.cpp).
[[gnu::noinline, gnu::noclone]] std::uint8_t sflash()
{
	const auto byte = avr::flash_load(flash_ptr(g_addr));
	++g_addr;
	return byte;
}

// One EEPROM byte <- [g_addr++].
[[gnu::noinline, gnu::noclone]] std::uint8_t eerd()
{
	const auto byte = ee::read<no_spm>(g_addr);
	++g_addr;
	return byte;
}

// One EEPROM byte -> [g_addr++].
[[gnu::noinline, gnu::noclone]] void eewr(std::uint8_t byte)
{
	ee::write<off, no_spm>(g_addr, byte);
	++g_addr;
}

// Stream g_cnt flash bytes from g_addr to the host.
[[gnu::noinline, gnu::noclone]] void sendf()
{
	do {
		tx(sflash());
	} while (--g_cnt);
}

// Wait out a running SPM op, then re-open the RWW section - after every page
// op and before handing over, as the oracle does.
[[gnu::noinline, gnu::noclone]] void settle()
{
	spm::wait();
	spm::rww_enable<off>();
}

extern "C" [[noreturn]] void tsb_app(); // the application's reset vector: --defsym=tsb_app=0

[[noreturn]] void appjump()
{
	settle();
	tsb_app();
}

// Step g_addr one page down and erase that page (the decrement lives here -
// the global-register rule).
[[gnu::noinline, gnu::noclone]] void erase_below()
{
	g_addr -= page;
	spm::command<off>(spm::op::erase, g_addr);
	settle();
}

// Erase the whole application, top-down like the oracle: the loop bound is a
// compare with zero, and g_addr = 0 is handed back for free.
[[gnu::noinline, gnu::noclone]] void erase_application()
{
	g_addr = app_end;
	do {
		erase_below();
	} while (g_addr != 0);
}

// Stream one host page into the erased flash page at g_addr (asm 2 of 2): the
// word pair stages in r0:r1 straight from rx (whose register set is known -
// the cross-call liveness C++ cannot express), Z walks the page and PGWRT
// programs it. g_addr is left at the next page base.
[[gnu::noinline, gnu::noclone]] void store_flash()
{
	asm volatile("	movw	r30, r28		\n\t" // Z = page base; rx leaves Z live
	             "	ldi	r20, %[words]		\n\t"
	             "1:	rcall	%x[rx]			\n\t"
	             "	mov	r0, r24			\n\t" // word low byte
	             "	rcall	%x[rx]			\n\t"
	             "	mov	r1, r24			\n\t"     // word high byte
	             "	ldi	r24, %[fill]		\n\t" // buffer the word at Z
	             "	out	%[spmcsr], r24		\n\t"
	             "	spm				\n\t"
	             "	clr	r1			\n\t"
	             "	adiw	r30, 2			\n\t"
	             "	dec	r20			\n\t"
	             "	brne	1b			\n\t"
	             "	movw	%[base], r30		\n\t" // g_addr = the next page base
	             "	subi	r30, %[pagelo]		\n\t" // Z back to this page's base
	             "	sbci	r31, %[pagehi]		\n\t"
	             "	ldi	r24, %[write]		\n\t" // program the page
	             "	out	%[spmcsr], r24		\n\t"
	             "	spm				\n\t"
	             : [base] "+r"(g_addr)
	             : [rx] "i"(&rx), [spmcsr] "I"(_SFR_IO_ADDR(SPMCSR)), [fill] "M"(hw::spmcsr::spmen(1).value),
	               [write] "M"(hw::spmcsr::pgwrt(1).value | hw::spmcsr::spmen(1).value), [words] "M"(page / 2),
	               [pagelo] "M"(page & 0xff), [pagehi] "M"(page >> 8)
	             : "r0", "r1", "r20", "r24", "r26", "r27", "r30", "r31", "cc", "memory");
	settle();
}

[[noreturn, gnu::noinline]] void run()
{
	// A watchdog reset hands straight back to the application, as the
	// reference loader does, rather than re-entering the bootloader.
	if (hw::mcusr::wdrf.test()) {
		appjump();
	}

	// Lean bring-up from reset state: UCSR0C already reads 8N1, UBRR0H reads
	// 0, and rx()/tx() raise RXEN0/TXEN0 on first use - only the divisor low
	// byte and U2X0 need a store. The library still does the datasheet work.
	static_assert(baud.u2x && baud.ubrr < 256, "lean bring-up writes UBRR0L only, with U2X0");
	hw::ubrr0::write(static_cast<std::uint8_t>(baud.ubrr));
	hw::ucsr0a::write(hw::ucsr0a::u2x0(1));
	// General-purpose registers are undefined at power-on (no crt zeroes them);
	// the direction latch must start "not receiving" so the first rx() enables
	// the receiver. The reference loader clears its shadow register for the
	// same reason.
	g_receiving = 0;

	// Activation: 3x'@', each inside the config page's timeout window (rx
	// floors it so a corrupt page cannot lock the loader out); anything else -
	// including silence - hands over.
	g_window = avr::flash_load(flash_ptr(app_end + 2));
	for (std::uint8_t k = 3; k; --k) {
		if (rx() != knock) {
			appjump();
		}
	}
	g_window = comm_window;

	// Password gate (config page from app_end+3, 0xff-terminated; a blank
	// page is no password). A wrong byte blanks the comparison and drains the
	// line forever, so a wrong password can never fall through; a 0 requests
	// emergency erase behind two confirms. On pass the info block goes out;
	// the emergency path skips it and drops into the command loop.
	g_addr = app_end + 3;
	std::uint8_t mask = 0xff;
	for (;;) {
		std::uint8_t expected = avr::flash_load(flash_ptr(g_addr)) & mask;
		++g_addr;
		if (expected == 0xff) {
			g_addr = reinterpret_cast<std::uint16_t>(info.data());
			g_cnt = sizeof(info);
			sendf();
			break;
		}
		const auto got = rx();
		if (got == 0) {
			if (mask == 0) {
				continue;
			}
			if (rcnf() != confirm || rcnf() != confirm) {
				appjump();
			}
			erase_application(); // leaves g_addr = 0 for the EEPROM walk
			do {
				eewr(0xff);
			} while (g_addr <= eeprom_end);
			g_addr = app_end + page;
			erase_below();
			break;
		}
		if (got != expected) {
			mask = 0;
		}
	}

	for (;;) {
		tx(confirm); // Mainloop ready
		g_addr = 0;
		switch (rx()) {
		case 'f': // read application flash, one page per host '!'
			for (;;) {
				if (rx() != confirm) {
					break;
				}
				g_cnt = page;
				sendf();
				if (g_addr >= app_end) {
					break;
				}
			}
			break;
		case 'F':                // erase the application, then take pages behind '?'
			erase_application(); // leaves g_addr = 0, the write start
			while (rcnf() == confirm) {
				store_flash();
			}
			break;
		case 'e': // read EEPROM, one page per host '!', until the host stops
			for (;;) {
				if (rx() != confirm) {
					break;
				}
				g_cnt = page;
				do {
					tx(eerd());
				} while (--g_cnt);
			}
			break;
		case 'E': // take EEPROM pages behind '?'
			while (rcnf() == confirm) {
				g_cnt = page;
				do {
					eewr(rx());
				} while (--g_cnt);
			}
			break;
		case 'c': // read the config page
		read_config:
			g_addr = app_end;
			g_cnt = page;
			sendf();
			break;
		case 'C': // replace the config page, then echo it back to verify
			if (rcnf() != confirm) {
				break;
			}
			g_addr = app_end + page;
			erase_below(); // leaves g_addr = app_end, the store target
			store_flash();
			goto read_config;
		default: // 'q' or any other byte runs the application
			appjump();
		}
	}
}

} // namespace
} // namespace tsb

// Reset lands at the boot section base (BOOTRST): the entry stub in .vectors
// is laid first and does the one line of crt a crt-less image needs.
template struct avr::startup::entry<tsb::run, avr::startup::stack::hardware>;