Prevent partial custom component overlays (#21070)

* Prevent partial custom component overlays

* Fix tests

homeassistant/loader.py

@@ -15,7 +15,7 @@ import importlib
 import logging
 import sys
 from types import ModuleType
-from typing import Optional, Set, TYPE_CHECKING, Callable, Any, TypeVar  # noqa pylint: disable=unused-import
+from typing import Optional, Set, TYPE_CHECKING, Callable, Any, TypeVar, List  # noqa pylint: disable=unused-import
 
 from homeassistant.const import PLATFORM_FORMAT
 
@@ -34,8 +34,9 @@ _LOGGER = logging.getLogger(__name__)
 
 
 DATA_KEY = 'components'
-PATH_CUSTOM_COMPONENTS = 'custom_components'
-PACKAGE_COMPONENTS = 'homeassistant.components'
+PACKAGE_CUSTOM_COMPONENTS = 'custom_components'
+PACKAGE_BUILTIN = 'homeassistant.components'
+LOOKUP_PATHS = [PACKAGE_CUSTOM_COMPONENTS, PACKAGE_BUILTIN]
 
 
 class LoaderError(Exception):
@@ -76,23 +77,43 @@ def get_platform(hass,  # type: HomeAssistant
                  domain: str, platform_name: str) -> Optional[ModuleType]:
     """Try to load specified platform.
 
+    Example invocation: get_platform(hass, 'light', 'hue')
+
     Async friendly.
     """
-    platform = _load_file(hass, PLATFORM_FORMAT.format(
-        domain=domain, platform=platform_name))
+    # If the platform has a component, we will limit the platform loading path
+    # to be the same source (custom/built-in).
+    component = get_component(hass, platform_name)
+
+    # Until we have moved all platforms under their component/own folder, it
+    # can be that the component is None.
+    if component is not None:
+        base_paths = [component.__name__.rsplit('.', 1)[0]]
+    else:
+        base_paths = LOOKUP_PATHS
+
+    platform = _load_file(
+        hass, PLATFORM_FORMAT.format(domain=domain, platform=platform_name),
+        base_paths)
 
     if platform is not None:
         return platform
 
     # Legacy platform check: light/hue.py
-    platform = _load_file(hass, PLATFORM_FORMAT.format(
-        domain=platform_name, platform=domain))
+    platform = _load_file(
+        hass, PLATFORM_FORMAT.format(domain=platform_name, platform=domain),
+        base_paths)
 
     if platform is None:
-        _LOGGER.error("Unable to find platform %s", platform_name)
+        if component is None:
+            extra = ""
+        else:
+            extra = " Search path was limited to path of component: {}".format(
+                base_paths[0])
+        _LOGGER.error("Unable to find platform %s.%s", platform_name, extra)
         return None
 
-    if platform.__name__.startswith(PATH_CUSTOM_COMPONENTS):
+    if platform.__name__.startswith(PACKAGE_CUSTOM_COMPONENTS):
         _LOGGER.warning(
             "Integrations need to be in their own folder. Change %s/%s.py to "
             "%s/%s.py. This will stop working soon.",
@@ -107,7 +128,7 @@ def get_component(hass,  # type: HomeAssistant
 
     Async friendly.
     """
-    comp = _load_file(hass, comp_or_platform)
+    comp = _load_file(hass, comp_or_platform, LOOKUP_PATHS)
 
     if comp is None:
         _LOGGER.error("Unable to find component %s", comp_or_platform)
@@ -116,7 +137,8 @@ def get_component(hass,  # type: HomeAssistant
 
 
 def _load_file(hass,  # type: HomeAssistant
-               comp_or_platform: str) -> Optional[ModuleType]:
+               comp_or_platform: str,
+               base_paths: List[str]) -> Optional[ModuleType]:
     """Try to load specified file.
 
     Looks in config dir first, then built-in components.
@@ -138,11 +160,8 @@ def _load_file(hass,  # type: HomeAssistant
             sys.path.insert(0, hass.config.config_dir)
         cache = hass.data[DATA_KEY] = {}
 
-    # First check custom, then built-in
-    potential_paths = ['custom_components.{}'.format(comp_or_platform),
-                       'homeassistant.components.{}'.format(comp_or_platform)]
-
-    for index, path in enumerate(potential_paths):
+    for index, path in enumerate('{}.{}'.format(base, comp_or_platform)
+                                 for base in base_paths):
         try:
             module = importlib.import_module(path)
 

tests/common.py

@@ -454,6 +454,7 @@ class MockModule:
                  async_setup_entry=None, async_unload_entry=None,
                  async_migrate_entry=None):
         """Initialize the mock module."""
+        self.__name__ = 'homeassistant.components.{}'.format(domain)
         self.DOMAIN = domain
         self.DEPENDENCIES = dependencies or []
         self.REQUIREMENTS = requirements or []

tests/test_loader.py

@@ -135,3 +135,10 @@ async def test_get_platform(hass, caplog):
     legacy_platform = loader.get_platform(hass, 'switch', 'test')
     assert legacy_platform.__name__ == 'custom_components.switch.test'
     assert 'Integrations need to be in their own folder.' in caplog.text
+
+
+async def test_get_platform_enforces_component_path(hass, caplog):
+    """Test that existence of a component limits lookup path of platforms."""
+    assert loader.get_platform(hass, 'comp_path_test', 'hue') is None
+    assert ('Search path was limited to path of component: '
+            'homeassistant.components') in caplog.text

tests/testing_config/custom_components/hue/comp_path_test.py

@@ -0,0 +1 @@
+"""Custom platform for a built-in component, should not be allowed."""