Setting up paging for real mode to 64-bit long mode switch

I am trying to write a bootloader using NASM and as a result, I have found OSDEV to be extremely helpful. However, somewhere in the process of setting up paging, loading my GDT, or transitioning (im coming directly from real mode), there is an error that causes the machine to reboot. I've based my code on the Long mode OSDEV article. This is what I have that's important to the question:

GDT

gdt_start:

.gdt_null: equ $-gdt_start      ; mandatory null descriptor
  dw 0
  dw 0
  db 0                         ; define double word
  db 0
  db 0
  db 0

.gdt_code: equ $-gdt_start      ; code segment
                               ; base = 0x0, limif = 0xffff
                               ; 1st flags: (present)1 (privilege)00 (descriptor type)1 -> 1001
                               ; type flags: (code)1 (conforming)0 (readable)1 (accessed)0 -> 1010
                               ; 2nd flags: (granularity)1 (32 bit default)1 (64 bit seg)0 (AVL)0 -> 1100
  dw 0                         ; limit (0-15)
  dw 0                         ; base (0-15)
  db 0                         ; base (16-23)
  db 10011010b                 ; 1st/type flags
  db 00100000b                 ; 2nd flags, limit (16-19)
  db 0                         ; base (bits 24-31)

.gdt_data: equ $-gdt_start      ; data segment descriptor
                ; type flags (code)0 (expand down)0 (writable)1 (accessed)0 -> 0010
  dw 0                         ; limit
  dw 0                         ; base
  db 0                         ; base
  db 10010010b                 ; 1st/type flags
  db 00000000b                 ; 2nd flags, limit
  db 0                         ; base

gdt_end:

gdt_descriptor:
  dw gdt_end - gdt_start - 1   ; size of GDT

  dq gdt_start

CODE_SEG equ gdt_start.gdt_code
DATA_SEG equ gdt_start.gdt_data

Paging and switch:

%include "gdt.ns"

;test/enable A20 line
call test_a20
fin:

cmp ax, 1
je enabled

call enable_A20

enabled:

;switch
call switch_to_lm
jmp $

switch_to_lm:
;
; SET UP PAGING!!!!!
;

;no previous paging defined so the below code is unnecessary
;mov eax, cr0
;and eax, 01111111111111111111111111111111b
;mov cr0, eax

;clear tables
mov edi, 0x1000
mov cr3, edi
xor eax, eax
mov ecx, 4096
rep stosd
mov edi, cr3

;set up new tables
mov DWORD [edi], 0x2003
add edi, 0x1000
mov DWORD [edi], 0x3003
add edi, 0x1000
mov DWORD [edi], 0x4003
add edi, 0x1000

mov ebx, 0x00000003
mov ecx, 512

.setEntry:
  mov DWORD [edi], ebx
  add ebx, 0x1000
  add edi, 8
  loop .setEntry

;enable PAE bit in CR4
mov eax, cr4
or eax, 1<<5
mov cr4, eax

;switch from REAL MODE
;set long mode bit
mov ecx, 0xc0000080
rdmsr
or eax, 1<<8
wrmsr

;enable paging
mov eax, cr0
or eax, 1<<31
mov cr0, eax

lgdt [gdt_descriptor]
jmp CODE_SEG:init_lm

[bits 64]

init_lm:

  cli
  mov ax, DATA_SEG
  mov ds, ax
  mov es, ax
  mov fs, ax
  mov gs, ax
  mov ss, ax

  mov ebp, 0x90000
  mov esp, ebp

  call BEGIN_LM

My code for testing for A20:

test_a20:
  pushf
  push ds
  push es
  push di
  push si

  cli

  xor ax, ax
  mov es, ax

  not ax
  mov ds, ax

  mov di, 0x0500
  mov si, 0x0510

  mov al, byte [es:di]
  push ax

  mov byte [es:di], 0x00
  mov byte [ds:si], 0xff

  cmp byte [es:di], 0xff

  pop ax
  mov byte [ds:si], al

  pop ax
  mov byte [es:di], al

  mov ax, 0
  je test_exit

  mov ax, 1

test_exit:
  pop si
  pop di
  pop es
  pop ds
  popf

  jmp fin

In your code, there is a problem with your test_a20 function. In particular you have this code:

mov al, byte [es:di]
push ax

mov byte [es:di], 0x00
mov byte [ds:si], 0xff

cmp byte [es:di], 0xff

pop ax
mov byte [ds:si], al

pop ax
mov byte [es:di], al

You appear to be pushing one value of AX on the stack, popping 2 off after. This will mess up the stack, registers will be restored incorrectly including DS and the flags. It appears you may have been trying to circumvent this bug by not using ret to return. Instead you used jmp fin to jump to a point after the call test_a20 instruction.

It appears you were trying to use the A20 test code from OSDEV Wiki. You'll note you are missing these lines:

mov al, byte [ds:si]
push ax

If you modify your test_a20 function to add the missing line and use ret it should look like:

test_a20:
  pushf
  push ds
  push es
  push di
  push si

  cli

  xor ax, ax
  mov es, ax

  not ax
  mov ds, ax

  mov di, 0x0500
  mov si, 0x0510

  mov al, byte [es:di]
  push ax

  mov al, byte [ds:si]
  push ax

  mov byte [es:di], 0x00
  mov byte [ds:si], 0xff

  cmp byte [es:di], 0xff

  pop ax
  mov byte [ds:si], al

  pop ax
  mov byte [es:di], al

  mov ax, 0
  je test_exit

  mov ax, 1

test_exit:
  pop si
  pop di
  pop es
  pop ds
  popf
  ret

This change should fix the problems with the DS register being destroyed and subsquent memory access in the page writing code to work incorrectly. You also have to amend your page enabling code to also enable protected mode. This code:

;enable paging
mov eax, cr0
or eax, 1<<31
mov cr0, eax

Should be:

;enable paging
mov eax, cr0
or eax, (1<<31) | (1<<0)
mov cr0, eax

With these changes you should be able to get into 64-bit long mode.