diff --git a/flashstub/lmi.c b/flashstub/lmi.c
new file mode 100644
index 0000000..27f8875
--- /dev/null
+++ b/flashstub/lmi.c
@@ -0,0 +1,49 @@
+/*
+ * This file is part of the Black Magic Debug project.
+ *
+ * Copyright (C) 2015  Black Sphere Technologies Ltd.
+ * Written by Gareth McMullin <gareth@blacksphere.co.nz>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+#include <stdint.h>
+#include "stub.h"
+
+#define LMI_FLASH_BASE       ((volatile uint32_t *)0x400FD000)
+#define LMI_FLASH_FMA        LMI_FLASH_BASE[0]
+#define LMI_FLASH_FMD        LMI_FLASH_BASE[1]
+#define LMI_FLASH_FMC        LMI_FLASH_BASE[2]
+
+#define LMI_FLASH_FMC_WRITE  (1 << 0)
+#define LMI_FLASH_FMC_ERASE  (1 << 1)
+#define LMI_FLASH_FMC_MERASE (1 << 2)
+#define LMI_FLASH_FMC_COMT   (1 << 3)
+#define LMI_FLASH_FMC_WRKEY  0xA4420000
+
+void __attribute__((naked))
+stm32f1_flash_write_stub(uint32_t *dest, uint32_t *src, uint32_t size)
+{
+	size /= 4;
+	for (int i; i < size; i++) {
+		LMI_FLASH_FMA = (uint32_t)&dest[i];
+		LMI_FLASH_FMD = src[i];
+		LMI_FLASH_FMC = LMI_FLASH_FMC_WRKEY | LMI_FLASH_FMC_WRITE;
+		while (LMI_FLASH_FMC & LMI_FLASH_FMC_WRITE)
+			;
+	}
+
+	stub_exit(0);
+}
+
+
diff --git a/flashstub/lmi.s b/flashstub/lmi.s
deleted file mode 100644
index b3b206e..0000000
--- a/flashstub/lmi.s
+++ /dev/null
@@ -1,42 +0,0 @@
-
-_start:
-	ldr r0, _flashbase
-	ldr r1, _addr
-	mov r2, pc
-	add r2, #(_data - . - 2)
-	ldr r3, _size
-	ldr r5, _flash_write_cmd
-_next:
-	cbz r3, _done
-	@ Write address to FMA
-	str r1, [r0]
-	@ Write data to FMD
-	ldr r4, [r2]
-	str r4, [r0, #4]
-	@ Write WRITE bit to FMC
-	str r5, [r0, #8]
-_wait:	@ Wait for WRITE bit to clear
-	ldr r4, [r0, #8]
-	mov r6, #1
-	tst r4, r6
-	bne _wait
-
-	sub r3, #1
-	add r1, #4
-	add r2, #4
-	b _next
-_done:
-	bkpt
-
-@.align 4
-.org 0x28
-_flashbase:
-	.word 0x400FD000
-_flash_write_cmd:
-	.word 0xA4420001
-_addr:
-	.word 0
-_size:
-	.word 4
-_data:
-	.string "Hello World!\n\0\0\0"
diff --git a/flashstub/lmi.stub b/flashstub/lmi.stub
index afaf939..5f2c696 100644
--- a/flashstub/lmi.stub
+++ b/flashstub/lmi.stub
@@ -1 +1 @@
-0x4809, 0x490B, 0x467A, 0x3230, 0x4B0A, 0x4D08, 0xB15B, 0x6001, 0x6814, 0x6044, 0x6085, 0x6884, 0x2601, 0x4234, 0xD1FB, 0x3B01, 0x3104, 0x3204, 0xE7F2, 0xBE00, 0xD000, 0x400F, 0x0001, 0xA442, 0x0000, 0x0000, 0x0004, 0x0000, 0x6548, 0x6C6C, 0x206F, 0x6F57, 0x6C72, 0x2164, 0x000A, 0x0000, 0x0000, 
+0x0892, 0x2300, 0x4293, 0xD20F, 0x4C08, 0x6020, 0xF851, 0x5023, 0x3408, 0xF844, 0x5C04, 0x4D06, 0x6025, 0x6825, 0xF015, 0x0F01, 0xD1FB, 0x3301, 0x3004, 0xE7ED, 0xBE00, 0xBF00, 0xD000, 0x400F, 0x0001, 0xA442, 
\ No newline at end of file
diff --git a/src/lmi.c b/src/lmi.c
index 6bf2240..e2d1809 100644
--- a/src/lmi.c
+++ b/src/lmi.c
@@ -29,7 +29,7 @@
 #include "cortexm.h"
 
 #define SRAM_BASE            0x20000000
-#define STUB_BUFFER_BASE     (SRAM_BASE + 0x30)
+#define STUB_BUFFER_BASE     ALIGN(SRAM_BASE + sizeof(lmi_flash_write_stub), 4)
 
 #define BLOCK_SIZE           0x400
 
@@ -109,15 +109,9 @@ int lmi_flash_write(struct target_flash *f,
 {
 	target  *t = f->t;
 
-	/* FIXME rewrite stub to use register args */
-	uint32_t data[(len>>2)+2];
-	data[0] = dest;
-	data[1] = len >> 2;
-	memcpy(&data[2], src, len);
-
 	target_mem_write(t, SRAM_BASE, lmi_flash_write_stub,
 	                 sizeof(lmi_flash_write_stub));
-	target_mem_write(t, STUB_BUFFER_BASE, data, len + 8);
-	return cortexm_run_stub(t, SRAM_BASE, 0, 0, 0, 0);
+	target_mem_write(t, STUB_BUFFER_BASE, src, len);
+	return cortexm_run_stub(t, SRAM_BASE, dest, STUB_BUFFER_BASE, len, 0);
 }