From 0bbb2484c7e5c8e94ff6bbc74ef6dab6769a68d6 Mon Sep 17 00:00:00 2001
From: Michael Brown
Date: Tue, 10 Jul 2012 22:08:38 +0100
Subject: [pxe] Exit back to the PXE stack instead of rebooting

Signed-off-by: Michael Brown <mcb30@ipxe.org>
---
 defs.h |  3 +++
 exit.S | 31 +++++++++++++++++++++++++++++++
 head.S |  6 ++++++
 pxe.S  | 12 ++++++++++++
 4 files changed, 52 insertions(+)

diff --git a/defs.h b/defs.h
index 3891699..06bb2dd 100644
--- a/defs.h
+++ b/defs.h
@@ -24,3 +24,6 @@
 #define KERNEL_DS	0x18			/* 32 bit segment adrs for data */
 #define REAL_CS		0x20			/* 16 bit segment adrs for code */
 #define REAL_DS		0x28			/* 16 bit segment adrs for data */
+
+/* Magic marker to determine if exit stack is intact */
+#define EXIT_MAGIC	( 'E' + ( 'X' << 8 ) + ( 'I' << 16 ) + ( 'T' << 24 ) )
diff --git a/exit.S b/exit.S
index 23858da..a118c13 100644
--- a/exit.S
+++ b/exit.S
@@ -12,6 +12,15 @@ exit:
 	/* Exit status to %eax */
 	movl	4(%esp), %eax
 
+	/* Calculate real-mode exit %ss:esp in %dx:%edi */
+	xorw	%dx, %dx
+	xorl	%edi, %edi
+	movl	exit_esp, %ebp
+	cmpl	$EXIT_MAGIC, 0(%ebp)
+	jne	1f
+	movl	8(%ebp), %edi
+	movw	6(%ebp), %dx
+1:
 	/* Load IDT and GDT and switch to 16-bit code segment */
 	cli
 	lidt	idt_descr
@@ -45,6 +54,24 @@ exit:
 	/* Reenable interrupts */
 	sti
 
+	/* If we have a real-mode exit stack, restore registers
+	 * (except exit status) and return
+	 */
+	testl	%edi, %edi
+	jz	reset
+	movw	%dx, %ss
+	movl	%edi, %esp
+	movw	%sp, %bp
+	movl	%eax, 36(%bp)
+	popw	%ds
+	popw	%es
+	popw	%fs
+	popw	%gs
+	popal
+	popfl
+	lret
+
+reset:
 	/* Perform a warm reset */
 	movw	$0x1234, %ax
 	movw	%ax, 0x472
@@ -79,3 +106,7 @@ gdt_descr:
 idt_descr:
 	.word	0xffff
 	.long	0
+
+.globl exit_esp
+exit_esp:
+	.long 0
diff --git a/head.S b/head.S
index e8f64a8..2464791 100644
--- a/head.S
+++ b/head.S
@@ -35,6 +35,12 @@ startup_32:
 0:	popl	%ebx
 	addl	$_GLOBAL_OFFSET_TABLE_+[.-0b], %ebx
 
+	/* Store the exit stack address */
+	cmpl	$0, exit_esp@GOTOFF(%ebx)
+	jnz	0f
+	movl	%esp, exit_esp@GOTOFF(%ebx)
+0:
+
 	/* Pick the appropriate boot_stack address */
 	leal	boot_stack_top@GOTOFF(%ebx), %esp
 
diff --git a/pxe.S b/pxe.S
index 98d8b98..932757c 100644
--- a/pxe.S
+++ b/pxe.S
@@ -23,6 +23,18 @@ _main:
 	ljmp	$BOOTSEG, $pxe_start
 
 pxe_start:
+	/* Preserve registers, %ss:%esp, and magic marker on PXE stack */
+	pushfl
+	pushal
+	pushw	%gs
+	pushw	%fs
+	pushw	%es
+	pushw	%ds
+	pushl	%esp
+	pushw	%ss
+	pushw	%ax	/* Padding */
+	pushl	$EXIT_MAGIC
+
 	/* Store PXENV+ entry point */
 	movl	%es:PXENV_ENTRY(%bx), %eax
 	movl	%eax, %cs:pxenv_vector
-- 
cgit v1.2.3-55-g7522