diff --git a/inc/hf/cpu.h b/inc/hf/cpu.h
index af7391f..7c93ad2 100644
--- a/inc/hf/cpu.h
+++ b/inc/hf/cpu.h
@@ -39,5 +39,5 @@
 bool cpu_on(struct cpu *c, ipaddr_t entry, uintreg_t arg);
 void cpu_off(struct cpu *c);
 struct cpu *cpu_find(cpu_id_t id);
-uint8_t *cpu_get_buffer(cpu_id_t cpu_id);
-uint32_t cpu_get_buffer_size(cpu_id_t cpu_id);
+uint8_t *cpu_get_buffer(struct cpu *c);
+uint32_t cpu_get_buffer_size(struct cpu *c);
diff --git a/src/api.c b/src/api.c
index 5170a41..062a495 100644
--- a/src/api.c
+++ b/src/api.c
@@ -1467,8 +1467,8 @@
 	 * sender can't change it underneath us.
 	 */
 	memory_region =
-		(struct spci_memory_region *)cpu_get_buffer(current->cpu->id);
-	message_buffer_size = cpu_get_buffer_size(current->cpu->id);
+		(struct spci_memory_region *)cpu_get_buffer(current->cpu);
+	message_buffer_size = cpu_get_buffer_size(current->cpu);
 	if (length > HF_MAILBOX_SIZE || length > message_buffer_size) {
 		return spci_error(SPCI_INVALID_PARAMETERS);
 	}
diff --git a/src/cpu.c b/src/cpu.c
index f8beed6..02c312f 100644
--- a/src/cpu.c
+++ b/src/cpu.c
@@ -52,18 +52,22 @@
  */
 alignas(PAGE_SIZE) static uint8_t cpu_message_buffer[MAX_CPUS][PAGE_SIZE];
 
-uint8_t *cpu_get_buffer(cpu_id_t cpu_id)
+uint8_t *cpu_get_buffer(struct cpu *c)
 {
-	CHECK(cpu_id < MAX_CPUS);
+	size_t cpu_indx = cpu_index(c);
 
-	return cpu_message_buffer[cpu_id];
+	CHECK(cpu_indx < MAX_CPUS);
+
+	return cpu_message_buffer[cpu_indx];
 }
 
-uint32_t cpu_get_buffer_size(cpu_id_t cpu_id)
+uint32_t cpu_get_buffer_size(struct cpu *c)
 {
-	CHECK(cpu_id < MAX_CPUS);
+	size_t cpu_indx = cpu_index(c);
 
-	return sizeof(cpu_message_buffer[cpu_id]);
+	CHECK(cpu_indx < MAX_CPUS);
+
+	return sizeof(cpu_message_buffer[cpu_indx]);
 }
 
 /* State of all supported CPUs. The stack of the first one is initialized. */