unsigned int size,
                                gfp_t gfp_mask)
 {
-       struct gbuf *gbuf;
-       int retval;
-
-       gbuf = kmem_cache_zalloc(gbuf_head_cache, gfp_mask);
-       if (!gbuf)
-               return NULL;
-
-       gbuf->hd = hd;
-       gbuf->dest_cport_id = dest_cport_id;
-       gbuf->status = -EBADR;  /* Initial value--means "never set" */
-
-       /* Host controller specific allocation for the actual buffer */
-       retval = hd->driver->alloc_gbuf_data(gbuf, size, gfp_mask);
-       if (retval) {
-               kmem_cache_free(gbuf_head_cache, gbuf);
-               return NULL;
-       }
-
-       return gbuf;
+       return kmem_cache_zalloc(gbuf_head_cache, gfp_mask);
 }
 EXPORT_SYMBOL_GPL(greybus_alloc_gbuf);
 
 void greybus_free_gbuf(struct gbuf *gbuf)
 {
-       gbuf->hd->driver->free_gbuf_data(gbuf);
-
        kmem_cache_free(gbuf_head_cache, gbuf);
 }
 EXPORT_SYMBOL_GPL(greybus_free_gbuf);
 
                        u8 *data, size_t length);
 
 struct gbuf *greybus_alloc_gbuf(struct greybus_host_device *hd,
-                               u16 dest_cport_id, unsigned int size,
+                               u16 dest_cport_id,
+                               unsigned int size,
                                gfp_t gfp_mask);
 void greybus_free_gbuf(struct gbuf *gbuf);
 
 
                                        bool request, bool data_out)
 {
        struct gb_connection *connection = operation->connection;
+       struct greybus_host_device *hd = connection->hd;
        struct gb_message *message;
        struct gb_operation_msg_hdr *header;
+       struct gbuf *gbuf;
        gfp_t gfp_flags = data_out ? GFP_KERNEL : GFP_ATOMIC;
        u16 dest_cport_id;
+       int ret;
 
        if (size > GB_OPERATION_MESSAGE_SIZE_MAX)
                return -E2BIG;
        if (message->gbuf)
                return -EALREADY;       /* Sanity check */
        size += sizeof(*header);
-       message->gbuf = greybus_alloc_gbuf(connection->hd, dest_cport_id,
-                                       size, gfp_flags);
-       if (!message->gbuf)
+       gbuf = greybus_alloc_gbuf(hd, dest_cport_id, size, gfp_flags);
+       if (!gbuf)
                return -ENOMEM;
+       gbuf->hd = hd;
+       gbuf->dest_cport_id = dest_cport_id;
+       gbuf->status = -EBADR;  /* Initial value--means "never set" */
+       ret = hd->driver->alloc_gbuf_data(gbuf, size, gfp_flags);
+       if (ret) {
+               greybus_free_gbuf(gbuf);
+               return ret;
+       }
 
        /* Fill in the header structure */
-       header = (struct gb_operation_msg_hdr *)message->gbuf->transfer_buffer;
+       header = (struct gb_operation_msg_hdr *)gbuf->transfer_buffer;
        header->size = cpu_to_le16(size);
        header->id = 0;         /* Filled in when submitted */
        header->type = type;
 
        message->payload = header + 1;
        message->operation = operation;
+       message->gbuf = gbuf;
 
        return 0;
 }
 {
        message->operation = NULL;
        message->payload = NULL;
+       message->gbuf->hd->driver->free_gbuf_data(message->gbuf);
        greybus_free_gbuf(message->gbuf);
        message->gbuf = NULL;
 }