diff --git a/protocol.c b/protocol.c
index 080e424..7a5b491 100644
--- a/protocol.c
+++ b/protocol.c
@@ -130,6 +130,45 @@
 	return 0;
 }
 
+int protocol_v1_mark_dirty(struct mbox_context *context,
+			   struct protocol_mark_dirty *io)
+{
+	uint32_t offset = io->req.v1.offset;
+	uint32_t size = io->req.v1.size;
+	uint32_t off;
+
+	if (!(context->current && context->current_is_write)) {
+		MSG_ERR("Tried to call mark dirty without open write window\n");
+		return -EPERM;
+	}
+
+	/* For V1 offset given relative to flash - we want the window */
+	off = offset - ((context->current->flash_offset) >>
+			context->block_size_shift);
+	if (off > offset) { /* Underflow - before current window */
+		MSG_ERR("Tried to mark dirty before start of window\n");
+		MSG_ERR("requested offset: 0x%x window start: 0x%x\n",
+				offset << context->block_size_shift,
+				context->current->flash_offset);
+		return -EINVAL;
+	}
+	offset = off;
+	/*
+	 * We only track dirty at the block level.
+	 * For protocol V1 we can get away with just marking the whole
+	 * block dirty.
+	 */
+	size = align_up(size, 1 << context->block_size_shift);
+	size >>= context->block_size_shift;
+
+	MSG_INFO("Dirty window @ 0x%.8x for 0x%.8x\n",
+		 offset << context->block_size_shift,
+		 size << context->block_size_shift);
+
+	return window_set_bytemap(context, context->current, offset, size,
+				  WINDOW_DIRTY);
+}
+
 /*
  * get_suggested_timeout() - get the suggested timeout value in seconds
  * @context:	The mbox context pointer
@@ -210,11 +249,28 @@
 	return 0;
 }
 
+int protocol_v2_mark_dirty(struct mbox_context *context,
+			   struct protocol_mark_dirty *io)
+{
+	if (!(context->current && context->current_is_write)) {
+		MSG_ERR("Tried to call mark dirty without open write window\n");
+		return -EPERM;
+	}
+
+	MSG_INFO("Dirty window @ 0x%.8x for 0x%.8x\n",
+		 io->req.v2.offset << context->block_size_shift,
+		 io->req.v2.size << context->block_size_shift);
+
+	return window_set_bytemap(context, context->current, io->req.v2.offset,
+				  io->req.v2.size, WINDOW_DIRTY);
+}
+
 static const struct protocol_ops protocol_ops_v1 = {
 	.reset = protocol_v1_reset,
 	.get_info = protocol_v1_get_info,
 	.get_flash_info = protocol_v1_get_flash_info,
 	.create_window = protocol_v1_create_window,
+	.mark_dirty = protocol_v1_mark_dirty,
 };
 
 static const struct protocol_ops protocol_ops_v2 = {
@@ -222,6 +278,7 @@
 	.get_info = protocol_v2_get_info,
 	.get_flash_info = protocol_v2_get_flash_info,
 	.create_window = protocol_v2_create_window,
+	.mark_dirty = protocol_v2_mark_dirty,
 };
 
 static const struct protocol_ops *protocol_ops_map[] = {
diff --git a/protocol.h b/protocol.h
index 9b5fc09..cc6bc8f 100644
--- a/protocol.h
+++ b/protocol.h
@@ -58,6 +58,21 @@
 	} resp;
 };
 
+struct protocol_mark_dirty {
+	struct {
+		union {
+			struct {
+				uint16_t offset;
+				uint32_t size;
+			} v1;
+			struct {
+				uint16_t offset;
+				uint16_t size;
+			} v2;
+		};
+	} req;
+};
+
 struct protocol_ops {
 	int (*reset)(struct mbox_context *context);
 	int (*get_info)(struct mbox_context *context,
@@ -66,6 +81,8 @@
 			      struct protocol_get_flash_info *io);
 	int (*create_window)(struct mbox_context *context,
 			     struct protocol_create_window *io);
+	int (*mark_dirty)(struct mbox_context *context,
+			  struct protocol_mark_dirty *io);
 };
 
 int protocol_init(struct mbox_context *context);
@@ -81,6 +98,8 @@
 			       struct protocol_get_flash_info *io);
 int protocol_v1_create_window(struct mbox_context *context,
 			      struct protocol_create_window *io);
+int protocol_v1_mark_dirty(struct mbox_context *context,
+			   struct protocol_mark_dirty *io);
 
 /* Protocol v2 */
 int protocol_v2_get_info(struct mbox_context *context,
@@ -89,5 +108,7 @@
 			       struct protocol_get_flash_info *io);
 int protocol_v2_create_window(struct mbox_context *context,
 			      struct protocol_create_window *io);
+int protocol_v2_mark_dirty(struct mbox_context *context,
+			   struct protocol_mark_dirty *io);
 
 #endif /* PROTOCOL_H */
diff --git a/transport_mbox.c b/transport_mbox.c
index 8ea400c..3382d35 100644
--- a/transport_mbox.c
+++ b/transport_mbox.c
@@ -387,51 +387,20 @@
 int mbox_handle_dirty_window(struct mbox_context *context,
 				    union mbox_regs *req, struct mbox_msg *resp)
 {
-	uint32_t offset, size;
+	struct protocol_mark_dirty io;
 	int rc;
 
-	if (!(context->current && context->current_is_write)) {
-		MSG_ERR("Tried to call mark dirty without open write window\n");
-		return context->version >= API_VERSION_2 ? -MBOX_R_WINDOW_ERROR
-							 : -MBOX_R_PARAM_ERROR;
-	}
-
-	offset = get_u16(&req->msg.args[0]);
-
-	if (context->version >= API_VERSION_2) {
-		size = get_u16(&req->msg.args[2]);
+	if (context->version == API_VERSION_1) {
+		io.req.v1.offset = get_u16(&req->msg.args[0]);
+		io.req.v1.size = get_u32(&req->msg.args[2]);
 	} else {
-		uint32_t off;
-		/* For V1 offset given relative to flash - we want the window */
-		off = offset - ((context->current->flash_offset) >>
-				context->block_size_shift);
-		if (off > offset) { /* Underflow - before current window */
-			MSG_ERR("Tried to mark dirty before start of window\n");
-			MSG_ERR("requested offset: 0x%x window start: 0x%x\n",
-				offset << context->block_size_shift,
-				context->current->flash_offset);
-			return -MBOX_R_PARAM_ERROR;
-		}
-		offset = off;
-		size = get_u32(&req->msg.args[2]);
-		/*
-		 * We only track dirty at the block level.
-		 * For protocol V1 we can get away with just marking the whole
-		 * block dirty.
-		 */
-		size = align_up(size, 1 << context->block_size_shift);
-		size >>= context->block_size_shift;
+		io.req.v2.offset = get_u16(&req->msg.args[0]);
+		io.req.v2.size = get_u16(&req->msg.args[2]);
 	}
 
-	MSG_INFO("Dirty window @ 0x%.8x for 0x%.8x\n",
-		 offset << context->block_size_shift,
-		 size << context->block_size_shift);
-
-	rc = window_set_bytemap(context, context->current, offset, size,
-				  WINDOW_DIRTY);
+	rc = context->protocol->mark_dirty(context, &io);
 	if (rc < 0) {
-		return (rc == -EACCES) ? -MBOX_R_PARAM_ERROR
-				       : -MBOX_R_SYSTEM_ERROR;
+		return mbox_xlate_errno(context, rc);
 	}
 
 	return rc;