struct onmessage_work_context {
        struct work_struct work;
-       struct hv_message msg;
+       struct {
+               struct hv_message_header header;
+               u8 payload[];
+       } msg;
 };
 
 static void vmbus_onmessage_work(struct work_struct *work)
                goto msg_handled;
 
        if (entry->handler_type == VMHT_BLOCKING) {
-               ctx = kmalloc(sizeof(*ctx), GFP_ATOMIC);
+               ctx = kmalloc(sizeof(*ctx) + msg->header.payload_size,
+                             GFP_ATOMIC);
                if (ctx == NULL)
                        return;
 
        WARN_ON(!is_hvsock_channel(channel));
 
        /*
-        * sizeof(*ctx) is small and the allocation should really not fail,
+        * Allocation size is small and the allocation should really not fail,
         * otherwise the state of the hv_sock connections ends up in limbo.
         */
-       ctx = kzalloc(sizeof(*ctx), GFP_KERNEL | __GFP_NOFAIL);
+       ctx = kzalloc(sizeof(*ctx) + sizeof(*rescind),
+                     GFP_KERNEL | __GFP_NOFAIL);
 
        /*
         * So far, these are not really used by Linux. Just set them to the
        ctx->msg.header.payload_size = sizeof(*rescind);
 
        /* These values are actually used by Linux. */
-       rescind = (struct vmbus_channel_rescind_offer *)ctx->msg.u.payload;
+       rescind = (struct vmbus_channel_rescind_offer *)ctx->msg.payload;
        rescind->header.msgtype = CHANNELMSG_RESCIND_CHANNELOFFER;
        rescind->child_relid = channel->offermsg.child_relid;