x86/efi: Preserve segment registers in mixed mode

I was triggering a #GP(0) from userland when running with
CONFIG_EFI_MIXED and CONFIG_IA32_EMULATION, from what looked like
register corruption. Turns out that the mixed mode code was trashing the
contents of %ds, %es and %ss in __efi64_thunk().

Save and restore the contents of these segment registers across the call
to __efi64_thunk() so that we don't corrupt the CPU context.

Signed-off-by: Matt Fleming <matt.fleming@intel.com>
diff --git a/arch/x86/platform/efi/efi_stub_64.S b/arch/x86/platform/efi/efi_stub_64.S
index 65b787a..e0984ef 100644
--- a/arch/x86/platform/efi/efi_stub_64.S
+++ b/arch/x86/platform/efi/efi_stub_64.S
@@ -176,6 +176,13 @@
  * This function must be invoked with a 1:1 mapped stack.
  */
 ENTRY(__efi64_thunk)
+	movl	%ds, %eax
+	push	%rax
+	movl	%es, %eax
+	push	%rax
+	movl	%ss, %eax
+	push	%rax
+
 	subq	$32, %rsp
 	movl	%esi, 0x0(%rsp)
 	movl	%edx, 0x4(%rsp)
@@ -191,7 +198,7 @@
 	movq	%rbx, func_rt_ptr(%rip)
 
 	/* Switch to gdt with 32-bit segments */
-	movl	40(%rsp), %eax
+	movl	64(%rsp), %eax
 	lgdt	(%rax)
 
 	leaq	efi_enter32(%rip), %rax
@@ -203,6 +210,13 @@
 
 	lgdt	save_gdt(%rip)
 
+	pop	%rbx
+	movl	%ebx, %ss
+	pop	%rbx
+	movl	%ebx, %es
+	pop	%rbx
+	movl	%ebx, %ds
+
 	/*
 	 * Convert 32-bit status code into 64-bit.
 	 */
@@ -218,11 +232,6 @@
 ENDPROC(__efi64_thunk)
 
 ENTRY(efi_exit32)
-	xorq	%rax, %rax
-	movl	%eax, %ds
-	movl	%eax, %es
-	movl	%eax, %ss
-
 	movq	func_rt_ptr(%rip), %rax
 	push	%rax
 	mov	%rdi, %rax
@@ -267,7 +276,7 @@
 	 */
 	cli
 
-	movl	44(%esp), %eax
+	movl	68(%esp), %eax
 	movl	%eax, 2(%eax)
 	lgdtl	(%eax)
 
@@ -286,7 +295,7 @@
 	xorl	%eax, %eax
 	lldt	%ax
 
-	movl	48(%esp), %eax
+	movl	72(%esp), %eax
 	pushl	$__KERNEL_CS
 	pushl	%eax