firewire: Add ref-counting for sbp2 orbs (fix command abortion)

This handles the case where we get the status write before getting the
complete_transaction callback ("status write for unknown orb"). In
this case, we just assume that the initial orb pointer transaction
succeeded and finish the orb. To prevent the transaction callback
from touching freed memory, we ref-count the orb structures.

Signed-off-by: Kristian Høgsberg <[e-mail]>
Signed-off-by: Stefan Richter <[e-mail]>

drivers/firewire/fw-sbp2.c

···
 
 struct sbp2_orb {
 	struct fw_transaction t;
+	struct kref kref;
 	dma_addr_t request_bus;
 	int rcode;
 	struct sbp2_pointer pointer;
···
 };
 
 static void
+free_orb(struct kref *kref)
+{
+	struct sbp2_orb *orb = container_of(kref, struct sbp2_orb, kref);
+
+	kfree(orb);
+}
+
+static void
 sbp2_status_write(struct fw_card *card, struct fw_request *request,
 		  int tcode, int destination, int source,
 		  int generation, int speed,
···
 	spin_lock_irqsave(&card->lock, flags);
 	list_for_each_entry(orb, &sd->orb_list, link) {
 		if (STATUS_GET_ORB_HIGH(status) == 0 &&
-		    STATUS_GET_ORB_LOW(status) == orb->request_bus &&
-		    orb->rcode == RCODE_COMPLETE) {
+		    STATUS_GET_ORB_LOW(status) == orb->request_bus) {
+			orb->rcode = RCODE_COMPLETE;
 			list_del(&orb->link);
 			break;
 		}
···
 	else
 		fw_error("status write for unknown orb\n");
 
+	kref_put(&orb->kref, free_orb);
+
 	fw_send_response(card, request, RCODE_COMPLETE);
 }
 
···
 	struct sbp2_orb *orb = data;
 	unsigned long flags;
 
-	orb->rcode = rcode;
-	if (rcode != RCODE_COMPLETE) {
-		spin_lock_irqsave(&card->lock, flags);
+	/*
+	 * This is a little tricky.  We can get the status write for
+	 * the orb before we get this callback.  The status write
+	 * handler above will assume the orb pointer transaction was
+	 * successful and set the rcode to RCODE_COMPLETE for the orb.
+	 * So this callback only sets the rcode if it hasn't already
+	 * been set and only does the cleanup if the transaction
+	 * failed and we didn't already get a status write.
+	 */
+	spin_lock_irqsave(&card->lock, flags);
+
+	if (orb->rcode == -1)
+		orb->rcode = rcode;
+	if (orb->rcode != RCODE_COMPLETE) {
 		list_del(&orb->link);
-		spin_unlock_irqrestore(&card->lock, flags);
 		orb->callback(orb, NULL);
 	}
+
+	spin_unlock_irqrestore(&card->lock, flags);
+
+	kref_put(&orb->kref, free_orb);
 }
 
 static void
···
 	spin_lock_irqsave(&device->card->lock, flags);
 	list_add_tail(&orb->link, &sd->orb_list);
 	spin_unlock_irqrestore(&device->card->lock, flags);
+
+	/* Take a ref for the orb list and for the transaction callback. */
+	kref_get(&orb->kref);
+	kref_get(&orb->kref);
 
 	fw_send_request(device->card, &orb->t, TCODE_WRITE_BLOCK_REQUEST,
 			node_id, generation, device->max_speed, offset,
···
 	if (orb == NULL)
 		return -ENOMEM;
 
+	kref_init(&orb->base.kref);
 	orb->response_bus =
 		dma_map_single(device->card->device, &orb->response,
 			       sizeof(orb->response), DMA_FROM_DEVICE);
···
 	if (response)
 		fw_memcpy_from_be32(response,
 				    orb->response, sizeof(orb->response));
-	kfree(orb);
+	kref_put(&orb->base.kref, free_orb);
 
 	return retval;
 }
···
 
 	orb->cmd->result = result;
 	orb->done(orb->cmd);
-	kfree(orb);
 }
 
 static int sbp2_command_orb_map_scatterlist(struct sbp2_command_orb *orb)
···
 
 	/* Initialize rcode to something not RCODE_COMPLETE. */
 	orb->base.rcode = -1;
+	kref_init(&orb->base.kref);
 
 	orb->unit = unit;
 	orb->done = done;
···
 	sbp2_send_orb(&orb->base, unit, sd->node_id, sd->generation,
 		      sd->command_block_agent_address + SBP2_ORB_POINTER);
 
+	kref_put(&orb->base.kref, free_orb);
 	return 0;
 
  fail_mapping:
-	kfree(orb);
+	kref_put(&orb->base.kref, free_orb);
  fail_alloc:
 	return SCSI_MLQUEUE_HOST_BUSY;
 }