4 files changed, 14 insertions, 26 deletions

diff --git a/drivers/thunderbolt/icm.c b/drivers/thunderbolt/icm.c
index 8b7f9131e9d1..7c923e16a7d8 100644
--- a/drivers/thunderbolt/icm.c
+++ b/drivers/thunderbolt/icm.c
@@ -798,9 +798,11 @@ icm_fr_xdomain_connected(struct tb *tb, const struct icm_pkg_header *hdr)
          * connected another host to the same port, remove the switch
          * first.
          */
-        sw = get_switch_at_route(tb->root_switch, route);
-        if (sw)
+        sw = tb_switch_find_by_route(tb, route);
+        if (sw) {
                 remove_switch(sw);
+                tb_switch_put(sw);
+        }
 
         sw = tb_switch_find_by_link_depth(tb, link, depth);
         if (!sw) {
@@ -1143,9 +1145,11 @@ icm_tr_xdomain_connected(struct tb *tb, const struct icm_pkg_header *hdr)
          * connected another host to the same port, remove the switch
          * first.
          */
-        sw = get_switch_at_route(tb->root_switch, route);
-        if (sw)
+        sw = tb_switch_find_by_route(tb, route);
+        if (sw) {
                 remove_switch(sw);
+                tb_switch_put(sw);
+        }
 
         sw = tb_switch_find_by_route(tb, get_parent_route(route));
         if (!sw) {
diff --git a/drivers/thunderbolt/switch.c b/drivers/thunderbolt/switch.c
index 504365d46827..5c2c0201ae7f 100644
--- a/drivers/thunderbolt/switch.c
+++ b/drivers/thunderbolt/switch.c
@@ -644,24 +644,6 @@ int tb_switch_reset(struct tb *tb, u64 route)
         return res.err;
 }
 
-struct tb_switch *get_switch_at_route(struct tb_switch *sw, u64 route)
-{
-        u8 next_port = route; /*
-                               * Routes use a stride of 8 bits,
-                               * eventhough a port index has 6 bits at most.
-                               * */
-        if (route == 0)
-                return sw;
-        if (next_port > sw->config.max_port_number)
-                return NULL;
-        if (tb_is_upstream_port(&sw->ports[next_port]))
-                return NULL;
-        if (!sw->ports[next_port].remote)
-                return NULL;
-        return get_switch_at_route(sw->ports[next_port].remote->sw,
-                                   route >> TB_ROUTE_SHIFT);
-}
-
 /**
  * tb_plug_events_active() - enable/disable plug events on a switch
  *
diff --git a/drivers/thunderbolt/tb.c b/drivers/thunderbolt/tb.c
index 30e02c716f6c..d8f4ed0f2ef8 100644
--- a/drivers/thunderbolt/tb.c
+++ b/drivers/thunderbolt/tb.c
@@ -258,7 +258,7 @@ static void tb_handle_hotplug(struct work_struct *work)
         if (!tcm->hotplug_active)
                 goto out; /* during init, suspend or shutdown */
 
-        sw = get_switch_at_route(tb->root_switch, ev->route);
+        sw = tb_switch_find_by_route(tb, ev->route);
         if (!sw) {
                 tb_warn(tb,
                         "hotplug event from non existent switch %llx:%x (unplug: %d)\n",
@@ -269,14 +269,14 @@ static void tb_handle_hotplug(struct work_struct *work)
                 tb_warn(tb,
                         "hotplug event from non existent port %llx:%x (unplug: %d)\n",
                         ev->route, ev->port, ev->unplug);
-                goto out;
+                goto put_sw;
         }
         port = &sw->ports[ev->port];
         if (tb_is_upstream_port(port)) {
                 tb_warn(tb,
                         "hotplug event for upstream port %llx:%x (unplug: %d)\n",
                         ev->route, ev->port, ev->unplug);
-                goto out;
+                goto put_sw;
         }
         if (ev->unplug) {
                 if (port->remote) {
@@ -306,6 +306,9 @@ static void tb_handle_hotplug(struct work_struct *work)
                         tb_activate_pcie_devices(tb);
                 }
         }
+
+put_sw:
+        tb_switch_put(sw);
 out:
         mutex_unlock(&tb->lock);
         kfree(ev);
diff --git a/drivers/thunderbolt/tb.h b/drivers/thunderbolt/tb.h
index 8058ea02d572..aea668c40d27 100644
--- a/drivers/thunderbolt/tb.h
+++ b/drivers/thunderbolt/tb.h
@@ -399,7 +399,6 @@ void tb_switch_suspend(struct tb_switch *sw);
 int tb_switch_resume(struct tb_switch *sw);
 int tb_switch_reset(struct tb *tb, u64 route);
 void tb_sw_set_unplugged(struct tb_switch *sw);
-struct tb_switch *get_switch_at_route(struct tb_switch *sw, u64 route);
 struct tb_switch *tb_switch_find_by_link_depth(struct tb *tb, u8 link,
                                                u8 depth);
 struct tb_switch *tb_switch_find_by_uuid(struct tb *tb, const uuid_t *uuid);