1 files changed, 94 insertions, 42 deletions

diff --git a/drivers/thunderbolt/cap.c b/drivers/thunderbolt/cap.c
index 19db6cdc5b70..6f571e912cf2 100644
--- a/drivers/thunderbolt/cap.c
+++ b/drivers/thunderbolt/cap.c
@@ -15,14 +15,6 @@
 #define VSE_CAP_OFFSET_MAX	0xffff
 #define TMU_ACCESS_EN		BIT(20)
 
-struct tb_cap_any {
-	union {
-		struct tb_cap_basic basic;
-		struct tb_cap_extended_short extended_short;
-		struct tb_cap_extended_long extended_long;
-	};
-} __packed;
-
 static int tb_port_enable_tmu(struct tb_port *port, bool enable)
 {
 	struct tb_switch *sw = port->sw;
@@ -67,23 +59,50 @@ static void tb_port_dummy_read(struct tb_port *port)
 	}
 }
 
+/**
+ * tb_port_next_cap() - Return next capability in the linked list
+ * @port: Port to find the capability for
+ * @offset: Previous capability offset (%0 for start)
+ *
+ * Returns dword offset of the next capability in port config space
+ * capability list and returns it. Passing %0 returns the first entry in
+ * the capability list. If no next capability is found returns %0. In case
+ * of failure returns negative errno.
+ */
+int tb_port_next_cap(struct tb_port *port, unsigned int offset)
+{
+	struct tb_cap_any header;
+	int ret;
+
+	if (!offset)
+		return port->config.first_cap_offset;
+
+	ret = tb_port_read(port, &header, TB_CFG_PORT, offset, 1);
+	if (ret)
+		return ret;
+
+	return header.basic.next;
+}
+
 static int __tb_port_find_cap(struct tb_port *port, enum tb_port_cap cap)
 {
-	u32 offset = 1;
+	int offset = 0;
 
 	do {
 		struct tb_cap_any header;
 		int ret;
 
+		offset = tb_port_next_cap(port, offset);
+		if (offset < 0)
+			return offset;
+
 		ret = tb_port_read(port, &header, TB_CFG_PORT, offset, 1);
 		if (ret)
 			return ret;
 
 		if (header.basic.cap == cap)
 			return offset;
-
-		offset = header.basic.next;
-	} while (offset);
+	} while (offset > 0);
 
 	return -ENOENT;
 }
@@ -114,6 +133,50 @@ int tb_port_find_cap(struct tb_port *port, enum tb_port_cap cap)
 }
 
 /**
+ * tb_switch_next_cap() - Return next capability in the linked list
+ * @sw: Switch to find the capability for
+ * @offset: Previous capability offset (%0 for start)
+ *
+ * Finds dword offset of the next capability in router config space
+ * capability list and returns it. Passing %0 returns the first entry in
+ * the capability list. If no next capability is found returns %0. In case
+ * of failure returns negative errno.
+ */
+int tb_switch_next_cap(struct tb_switch *sw, unsigned int offset)
+{
+	struct tb_cap_any header;
+	int ret;
+
+	if (!offset)
+		return sw->config.first_cap_offset;
+
+	ret = tb_sw_read(sw, &header, TB_CFG_SWITCH, offset, 2);
+	if (ret)
+		return ret;
+
+	switch (header.basic.cap) {
+	case TB_SWITCH_CAP_TMU:
+		ret = header.basic.next;
+		break;
+
+	case TB_SWITCH_CAP_VSE:
+		if (!header.extended_short.length)
+			ret = header.extended_long.next;
+		else
+			ret = header.extended_short.next;
+		break;
+
+	default:
+		tb_sw_dbg(sw, "unknown capability %#x at %#x\n",
+			  header.basic.cap, offset);
+		ret = -EINVAL;
+		break;
+	}
+
+	return ret >= VSE_CAP_OFFSET_MAX ? 0 : ret;
+}
+
+/**
  * tb_switch_find_cap() - Find switch capability
  * @sw Switch to find the capability for
  * @cap: Capability to look
@@ -124,21 +187,23 @@ int tb_port_find_cap(struct tb_port *port, enum tb_port_cap cap)
  */
 int tb_switch_find_cap(struct tb_switch *sw, enum tb_switch_cap cap)
 {
-	int offset = sw->config.first_cap_offset;
+	int offset = 0;
 
-	while (offset > 0 && offset < CAP_OFFSET_MAX) {
+	do {
 		struct tb_cap_any header;
 		int ret;
 
+		offset = tb_switch_next_cap(sw, offset);
+		if (offset < 0)
+			return offset;
+
 		ret = tb_sw_read(sw, &header, TB_CFG_SWITCH, offset, 1);
 		if (ret)
 			return ret;
 
 		if (header.basic.cap == cap)
 			return offset;
-
-		offset = header.basic.next;
-	}
+	} while (offset);
 
 	return -ENOENT;
 }
@@ -155,37 +220,24 @@ int tb_switch_find_cap(struct tb_switch *sw, enum tb_switch_cap cap)
  */
 int tb_switch_find_vse_cap(struct tb_switch *sw, enum tb_switch_vse_cap vsec)
 {
-	struct tb_cap_any header;
-	int offset;
+	int offset = 0;
 
-	offset = tb_switch_find_cap(sw, TB_SWITCH_CAP_VSE);
-	if (offset < 0)
-		return offset;
-
-	while (offset > 0 && offset < VSE_CAP_OFFSET_MAX) {
+	do {
+		struct tb_cap_any header;
 		int ret;
 
-		ret = tb_sw_read(sw, &header, TB_CFG_SWITCH, offset, 2);
+		offset = tb_switch_next_cap(sw, offset);
+		if (offset < 0)
+			return offset;
+
+		ret = tb_sw_read(sw, &header, TB_CFG_SWITCH, offset, 1);
 		if (ret)
 			return ret;
 
-		/*
-		 * Extended vendor specific capabilities come in two
-		 * flavors: short and long. The latter is used when
-		 * offset is over 0xff.
-		 */
-		if (offset >= CAP_OFFSET_MAX) {
-			if (header.extended_long.vsec_id == vsec)
-				return offset;
-			offset = header.extended_long.next;
-		} else {
-			if (header.extended_short.vsec_id == vsec)
-				return offset;
-			if (!header.extended_short.length)
-				return -ENOENT;
-			offset = header.extended_short.next;
-		}
-	}
+		if (header.extended_short.cap == TB_SWITCH_CAP_VSE &&
+		    header.extended_short.vsec_id == vsec)
+			return offset;
+	} while (offset);
 
 	return -ENOENT;
 }