1 files changed, 81 insertions, 26 deletions

diff --git a/drivers/usb/host/xhci-mem.c b/drivers/usb/host/xhci-mem.c
index be7b7b6b5517..edfb31ad5949 100644
--- a/drivers/usb/host/xhci-mem.c
+++ b/drivers/usb/host/xhci-mem.c
@@ -180,53 +180,98 @@ static void xhci_link_rings(struct xhci_hcd *xhci, struct xhci_ring *ring,
  * extended systems (where the DMA address can be bigger than 32-bits),
  * if we allow the PCI dma mask to be bigger than 32-bits.  So don't do that.
  */
-static int xhci_update_stream_mapping(struct xhci_ring *ring, gfp_t mem_flags)
+static int xhci_insert_segment_mapping(struct radix_tree_root *trb_address_map,
+		struct xhci_ring *ring,
+		struct xhci_segment *seg,
+		gfp_t mem_flags)
 {
-	struct xhci_segment *seg;
 	unsigned long key;
 	int ret;
 
-	if (WARN_ON_ONCE(ring->trb_address_map == NULL))
+	key = (unsigned long)(seg->dma >> TRB_SEGMENT_SHIFT);
+	/* Skip any segments that were already added. */
+	if (radix_tree_lookup(trb_address_map, key))
 		return 0;
 
-	seg = ring->first_seg;
-	do {
-		key = (unsigned long)(seg->dma >> TRB_SEGMENT_SHIFT);
-		/* Skip any segments that were already added. */
-		if (radix_tree_lookup(ring->trb_address_map, key))
-			continue;
+	ret = radix_tree_maybe_preload(mem_flags);
+	if (ret)
+		return ret;
+	ret = radix_tree_insert(trb_address_map,
+			key, ring);
+	radix_tree_preload_end();
+	return ret;
+}
 
-		ret = radix_tree_maybe_preload(mem_flags);
-		if (ret)
-			return ret;
-		ret = radix_tree_insert(ring->trb_address_map,
-				key, ring);
-		radix_tree_preload_end();
+static void xhci_remove_segment_mapping(struct radix_tree_root *trb_address_map,
+		struct xhci_segment *seg)
+{
+	unsigned long key;
+
+	key = (unsigned long)(seg->dma >> TRB_SEGMENT_SHIFT);
+	if (radix_tree_lookup(trb_address_map, key))
+		radix_tree_delete(trb_address_map, key);
+}
+
+static int xhci_update_stream_segment_mapping(
+		struct radix_tree_root *trb_address_map,
+		struct xhci_ring *ring,
+		struct xhci_segment *first_seg,
+		struct xhci_segment *last_seg,
+		gfp_t mem_flags)
+{
+	struct xhci_segment *seg;
+	struct xhci_segment *failed_seg;
+	int ret;
+
+	if (WARN_ON_ONCE(trb_address_map == NULL))
+		return 0;
+
+	seg = first_seg;
+	do {
+		ret = xhci_insert_segment_mapping(trb_address_map,
+				ring, seg, mem_flags);
 		if (ret)
-			return ret;
+			goto remove_streams;
+		if (seg == last_seg)
+			return 0;
 		seg = seg->next;
-	} while (seg != ring->first_seg);
+	} while (seg != first_seg);
 
 	return 0;
+
+remove_streams:
+	failed_seg = seg;
+	seg = first_seg;
+	do {
+		xhci_remove_segment_mapping(trb_address_map, seg);
+		if (seg == failed_seg)
+			return ret;
+		seg = seg->next;
+	} while (seg != first_seg);
+
+	return ret;
 }
 
 static void xhci_remove_stream_mapping(struct xhci_ring *ring)
 {
 	struct xhci_segment *seg;
-	unsigned long key;
 
 	if (WARN_ON_ONCE(ring->trb_address_map == NULL))
 		return;
 
 	seg = ring->first_seg;
 	do {
-		key = (unsigned long)(seg->dma >> TRB_SEGMENT_SHIFT);
-		if (radix_tree_lookup(ring->trb_address_map, key))
-			radix_tree_delete(ring->trb_address_map, key);
+		xhci_remove_segment_mapping(ring->trb_address_map, seg);
 		seg = seg->next;
 	} while (seg != ring->first_seg);
 }
 
+static int xhci_update_stream_mapping(struct xhci_ring *ring, gfp_t mem_flags)
+{
+	return xhci_update_stream_segment_mapping(ring->trb_address_map, ring,
+			ring->first_seg, ring->last_seg, mem_flags);
+}
+
 /* XXX: Do we need the hcd structure in all these functions? */
 void xhci_ring_free(struct xhci_hcd *xhci, struct xhci_ring *ring)
 {
@@ -430,16 +475,26 @@ int xhci_ring_expansion(struct xhci_hcd *xhci, struct xhci_ring *ring,
 	if (ret)
 		return -ENOMEM;
 
+	if (ring->type == TYPE_STREAM)
+		ret = xhci_update_stream_segment_mapping(ring->trb_address_map,
+						ring, first, last, flags);
+	if (ret) {
+		struct xhci_segment *next;
+		do {
+			next = first->next;
+			xhci_segment_free(xhci, first);
+			if (first == last)
+				break;
+			first = next;
+		} while (true);
+		return ret;
+	}
+
 	xhci_link_rings(xhci, ring, first, last, num_segs);
 	xhci_dbg_trace(xhci, trace_xhci_dbg_ring_expansion,
 			"ring expansion succeed, now has %d segments",
 			ring->num_segs);
 
-	if (ring->type == TYPE_STREAM) {
-		ret = xhci_update_stream_mapping(ring, flags);
-		WARN_ON(ret); /* FIXME */
-	}
-
 	return 0;
 }