}
 
 
-/* -- SSAM base-hub driver. ------------------------------------------------- */
+/* -- SSAM generic subsystem hub driver framework. -------------------------- */
 
-/*
- * Some devices (especially battery) may need a bit of time to be fully usable
- * after being (re-)connected. This delay has been determined via
- * experimentation.
- */
-#define SSAM_BASE_UPDATE_CONNECT_DELAY         msecs_to_jiffies(2500)
+enum ssam_hub_state {
+       SSAM_HUB_UNINITIALIZED,         /* Only set during initialization. */
+       SSAM_HUB_CONNECTED,
+       SSAM_HUB_DISCONNECTED,
+};
 
-enum ssam_base_hub_state {
-       SSAM_BASE_HUB_UNINITIALIZED,
-       SSAM_BASE_HUB_CONNECTED,
-       SSAM_BASE_HUB_DISCONNECTED,
+enum ssam_hub_flags {
+       SSAM_HUB_HOT_REMOVED,
 };
 
-struct ssam_base_hub {
+struct ssam_hub {
        struct ssam_device *sdev;
 
-       enum ssam_base_hub_state state;
+       enum ssam_hub_state state;
+       unsigned long flags;
+
        struct delayed_work update_work;
+       unsigned long connect_delay;
 
        struct ssam_event_notifier notif;
+
+       int (*get_state)(struct ssam_hub *hub, enum ssam_hub_state *state);
 };
 
+static void ssam_hub_update_workfn(struct work_struct *work)
+{
+       struct ssam_hub *hub = container_of(work, struct ssam_hub, update_work.work);
+       struct fwnode_handle *node = dev_fwnode(&hub->sdev->dev);
+       enum ssam_hub_state state;
+       int status = 0;
+
+       status = hub->get_state(hub, &state);
+       if (status)
+               return;
+
+       /*
+        * There is a small possibility that hub devices were hot-removed and
+        * re-added before we were able to remove them here. In that case, both
+        * the state returned by get_state() and the state of the hub will
+        * equal SSAM_HUB_CONNECTED and we would bail early below, which would
+        * leave child devices without proper (re-)initialization and the
+        * hot-remove flag set.
+        *
+        * Therefore, we check whether devices have been hot-removed via an
+        * additional flag on the hub and, in this case, override the returned
+        * hub state. In case of a missed disconnect (i.e. get_state returned
+        * "connected"), we further need to re-schedule this work (with the
+        * appropriate delay) as the actual connect work submission might have
+        * been merged with this one.
+        *
+        * This then leads to one of two cases: Either we submit an unnecessary
+        * work item (which will get ignored via either the queue or the state
+        * checks) or, in the unlikely case that the work is actually required,
+        * double the normal connect delay.
+        */
+       if (test_and_clear_bit(SSAM_HUB_HOT_REMOVED, &hub->flags)) {
+               if (state == SSAM_HUB_CONNECTED)
+                       schedule_delayed_work(&hub->update_work, hub->connect_delay);
+
+               state = SSAM_HUB_DISCONNECTED;
+       }
+
+       if (hub->state == state)
+               return;
+       hub->state = state;
+
+       if (hub->state == SSAM_HUB_CONNECTED)
+               status = ssam_hub_register_clients(&hub->sdev->dev, hub->sdev->ctrl, node);
+       else
+               ssam_remove_clients(&hub->sdev->dev);
+
+       if (status)
+               dev_err(&hub->sdev->dev, "failed to update hub child devices: %d\n", status);
+}
+
+static int ssam_hub_mark_hot_removed(struct device *dev, void *_data)
+{
+       struct ssam_device *sdev = to_ssam_device(dev);
+
+       if (is_ssam_device(dev))
+               ssam_device_mark_hot_removed(sdev);
+
+       return 0;
+}
+
+static void ssam_hub_update(struct ssam_hub *hub, bool connected)
+{
+       unsigned long delay;
+
+       /* Mark devices as hot-removed before we remove any. */
+       if (!connected) {
+               set_bit(SSAM_HUB_HOT_REMOVED, &hub->flags);
+               device_for_each_child_reverse(&hub->sdev->dev, NULL, ssam_hub_mark_hot_removed);
+       }
+
+       /*
+        * Delay update when the base/keyboard cover is being connected to give
+        * devices/EC some time to set up.
+        */
+       delay = connected ? hub->connect_delay : 0;
+
+       schedule_delayed_work(&hub->update_work, delay);
+}
+
+static int __maybe_unused ssam_hub_resume(struct device *dev)
+{
+       struct ssam_hub *hub = dev_get_drvdata(dev);
+
+       schedule_delayed_work(&hub->update_work, 0);
+       return 0;
+}
+static SIMPLE_DEV_PM_OPS(ssam_hub_pm_ops, NULL, ssam_hub_resume);
+
+static int ssam_hub_setup(struct ssam_device *sdev, struct ssam_hub *hub)
+{
+       int status;
+
+       hub->sdev = sdev;
+       hub->state = SSAM_HUB_UNINITIALIZED;
+
+       INIT_DELAYED_WORK(&hub->update_work, ssam_hub_update_workfn);
+
+       ssam_device_set_drvdata(sdev, hub);
+
+       status = ssam_device_notifier_register(sdev, &hub->notif);
+       if (status)
+               return status;
+
+       schedule_delayed_work(&hub->update_work, 0);
+       return 0;
+}
+
+static void ssam_hub_remove(struct ssam_device *sdev)
+{
+       struct ssam_hub *hub = ssam_device_get_drvdata(sdev);
+
+       ssam_device_notifier_unregister(sdev, &hub->notif);
+       cancel_delayed_work_sync(&hub->update_work);
+       ssam_remove_clients(&sdev->dev);
+}
+
+
+/* -- SSAM base-hub driver. ------------------------------------------------- */
+
+/*
+ * Some devices (especially battery) may need a bit of time to be fully usable
+ * after being (re-)connected. This delay has been determined via
+ * experimentation.
+ */
+#define SSAM_BASE_UPDATE_CONNECT_DELAY         msecs_to_jiffies(2500)
+
 SSAM_DEFINE_SYNC_REQUEST_R(ssam_bas_query_opmode, u8, {
        .target_category = SSAM_SSH_TC_BAS,
        .target_id       = 0x01,
 #define SSAM_BAS_OPMODE_TABLET         0x00
 #define SSAM_EVENT_BAS_CID_CONNECTION  0x0c
 
-static int ssam_base_hub_query_state(struct ssam_base_hub *hub, enum ssam_base_hub_state *state)
+static int ssam_base_hub_query_state(struct ssam_hub *hub, enum ssam_hub_state *state)
 {
        u8 opmode;
        int status;
        }
 
        if (opmode != SSAM_BAS_OPMODE_TABLET)
-               *state = SSAM_BASE_HUB_CONNECTED;
+               *state = SSAM_HUB_CONNECTED;
        else
-               *state = SSAM_BASE_HUB_DISCONNECTED;
+               *state = SSAM_HUB_DISCONNECTED;
 
        return 0;
 }
 
-static ssize_t ssam_base_hub_state_show(struct device *dev, struct device_attribute *attr,
-                                       char *buf)
-{
-       struct ssam_base_hub *hub = dev_get_drvdata(dev);
-       bool connected = hub->state == SSAM_BASE_HUB_CONNECTED;
-
-       return sysfs_emit(buf, "%d\n", connected);
-}
-
-static struct device_attribute ssam_base_hub_attr_state =
-       __ATTR(state, 0444, ssam_base_hub_state_show, NULL);
-
-static struct attribute *ssam_base_hub_attrs[] = {
-       &ssam_base_hub_attr_state.attr,
-       NULL,
-};
-
-static const struct attribute_group ssam_base_hub_group = {
-       .attrs = ssam_base_hub_attrs,
-};
-
-static void ssam_base_hub_update_workfn(struct work_struct *work)
-{
-       struct ssam_base_hub *hub = container_of(work, struct ssam_base_hub, update_work.work);
-       struct fwnode_handle *node = dev_fwnode(&hub->sdev->dev);
-       enum ssam_base_hub_state state;
-       int status = 0;
-
-       status = ssam_base_hub_query_state(hub, &state);
-       if (status)
-               return;
-
-       if (hub->state == state)
-               return;
-       hub->state = state;
-
-       if (hub->state == SSAM_BASE_HUB_CONNECTED)
-               status = ssam_hub_register_clients(&hub->sdev->dev, hub->sdev->ctrl, node);
-       else
-               ssam_remove_clients(&hub->sdev->dev);
-
-       if (status)
-               dev_err(&hub->sdev->dev, "failed to update base-hub devices: %d\n", status);
-}
-
 static u32 ssam_base_hub_notif(struct ssam_event_notifier *nf, const struct ssam_event *event)
 {
-       struct ssam_base_hub *hub = container_of(nf, struct ssam_base_hub, notif);
-       unsigned long delay;
+       struct ssam_hub *hub = container_of(nf, struct ssam_hub, notif);
 
        if (event->command_id != SSAM_EVENT_BAS_CID_CONNECTION)
                return 0;
                return 0;
        }
 
-       /*
-        * Delay update when the base is being connected to give devices/EC
-        * some time to set up.
-        */
-       delay = event->data[0] ? SSAM_BASE_UPDATE_CONNECT_DELAY : 0;
-
-       schedule_delayed_work(&hub->update_work, delay);
+       ssam_hub_update(hub, event->data[0]);
 
        /*
         * Do not return SSAM_NOTIF_HANDLED: The event should be picked up and
        return 0;
 }
 
-static int __maybe_unused ssam_base_hub_resume(struct device *dev)
-{
-       struct ssam_base_hub *hub = dev_get_drvdata(dev);
-
-       schedule_delayed_work(&hub->update_work, 0);
-       return 0;
-}
-static SIMPLE_DEV_PM_OPS(ssam_base_hub_pm_ops, NULL, ssam_base_hub_resume);
-
 static int ssam_base_hub_probe(struct ssam_device *sdev)
 {
-       struct ssam_base_hub *hub;
-       int status;
+       struct ssam_hub *hub;
 
        hub = devm_kzalloc(&sdev->dev, sizeof(*hub), GFP_KERNEL);
        if (!hub)
                return -ENOMEM;
 
-       hub->sdev = sdev;
-       hub->state = SSAM_BASE_HUB_UNINITIALIZED;
-
        hub->notif.base.priority = INT_MAX;  /* This notifier should run first. */
        hub->notif.base.fn = ssam_base_hub_notif;
        hub->notif.event.reg = SSAM_EVENT_REGISTRY_SAM;
        hub->notif.event.mask = SSAM_EVENT_MASK_NONE;
        hub->notif.event.flags = SSAM_EVENT_SEQUENCED;
 
-       INIT_DELAYED_WORK(&hub->update_work, ssam_base_hub_update_workfn);
-
-       ssam_device_set_drvdata(sdev, hub);
-
-       status = ssam_device_notifier_register(sdev, &hub->notif);
-       if (status)
-               return status;
-
-       status = sysfs_create_group(&sdev->dev.kobj, &ssam_base_hub_group);
-       if (status)
-               goto err;
-
-       schedule_delayed_work(&hub->update_work, 0);
-       return 0;
+       hub->connect_delay = SSAM_BASE_UPDATE_CONNECT_DELAY;
+       hub->get_state = ssam_base_hub_query_state;
 
-err:
-       ssam_device_notifier_unregister(sdev, &hub->notif);
-       cancel_delayed_work_sync(&hub->update_work);
-       ssam_remove_clients(&sdev->dev);
-       return status;
-}
-
-static void ssam_base_hub_remove(struct ssam_device *sdev)
-{
-       struct ssam_base_hub *hub = ssam_device_get_drvdata(sdev);
-
-       sysfs_remove_group(&sdev->dev.kobj, &ssam_base_hub_group);
-
-       ssam_device_notifier_unregister(sdev, &hub->notif);
-       cancel_delayed_work_sync(&hub->update_work);
-       ssam_remove_clients(&sdev->dev);
+       return ssam_hub_setup(sdev, hub);
 }
 
 static const struct ssam_device_id ssam_base_hub_match[] = {
 
 static struct ssam_device_driver ssam_base_hub_driver = {
        .probe = ssam_base_hub_probe,
-       .remove = ssam_base_hub_remove,
+       .remove = ssam_hub_remove,
        .match_table = ssam_base_hub_match,
        .driver = {
                .name = "surface_aggregator_base_hub",
                .probe_type = PROBE_PREFER_ASYNCHRONOUS,
-               .pm = &ssam_base_hub_pm_ops,
+               .pm = &ssam_hub_pm_ops,
        },
 };