filter instagram-server requests from hitting fbcode codenav server

Summary: Here we introduce a method to fail fast and immediately respond with a DaemonQueryFailure instance in cases where a specified regex directory pattern is matched against.

Reviewed By: kinto0

Differential Revision: D47655034

fbshipit-source-id: c4f992bf28c8221c368b85f0514b180420743e3a

client/commands/code_navigation.py

@@ -17,7 +17,7 @@
 import logging
 import traceback
 
-from typing import Optional
+from typing import Callable, Optional
 
 from .. import (
     backend_arguments,
@@ -39,6 +39,8 @@
     subscription,
 )
 
+from .daemon_query_failer import AbstractDaemonQueryFailer
+
 
 LOG: logging.Logger = logging.getLogger(__name__)
 
@@ -200,6 +202,9 @@ async def async_run_code_navigation_client(
     server_options_reader: pyre_server_options.PyreServerOptionsReader,
     remote_logging: Optional[backend_arguments.RemoteLogging],
     index: remote_index.AbstractRemoteIndex,
+    daemon_query_failer_provider: Callable[
+        [pyre_server_options.PyreServerOptions], AbstractDaemonQueryFailer
+    ],
 ) -> int:
     initial_server_options = launch_and_subscribe_handler.PyreDaemonLaunchAndSubscribeHandler.read_server_options(
         server_options_reader, remote_logging
@@ -236,9 +241,16 @@ async def async_run_code_navigation_client(
         client_capabilities=client_capabilities,
         server_options=initial_server_options,
     )
-    codenav_querier = daemon_querier.CodeNavigationDaemonQuerier(
-        server_state=server_state
+
+    daemon_query_failer = daemon_query_failer_provider(initial_server_options)
+
+    codenav_querier = daemon_querier.FailableDaemonQuerier(
+        base_querier=daemon_querier.CodeNavigationDaemonQuerier(
+            server_state=server_state
+        ),
+        daemon_query_failer=daemon_query_failer,
     )
+
     querier = daemon_querier.RemoteIndexBackedQuerier(codenav_querier, index)
     client_type_error_handler = type_error_handler.ClientTypeErrorHandler(
         stdout, server_state, remote_logging
@@ -273,6 +285,9 @@ def run(
     server_options_reader: pyre_server_options.PyreServerOptionsReader,
     remote_logging: Optional[backend_arguments.RemoteLogging],
     index: remote_index.AbstractRemoteIndex,
+    daemon_query_failer_provider: Callable[
+        [pyre_server_options.PyreServerOptions], AbstractDaemonQueryFailer
+    ],
 ) -> int:
     command_timer = timer.Timer()
     error_message: Optional[str] = None
@@ -282,6 +297,7 @@ def run(
                 server_options_reader,
                 remote_logging,
                 index,
+                daemon_query_failer_provider,
             )
         )
     except Exception:

client/commands/daemon_querier.py

@@ -19,7 +19,7 @@
 import logging
 import os
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Callable, List, Optional, Union
 
 from .. import dataclasses_json_extensions as json_mixins, error
 
@@ -31,6 +31,7 @@
     remote_index,
 )
 from . import daemon_query, expression_level_coverage, server_state as state
+from .daemon_query_failer import AbstractDaemonQueryFailer
 
 LOG: logging.Logger = logging.getLogger(__name__)
 
@@ -458,6 +459,173 @@ def _get_overlay_id(self, path: Path) -> Optional[str]:
         return f"{path}, pid_{os.getpid()}" if unsaved_changes_enabled else None
 
 
+class FailableDaemonQuerier(AbstractDaemonQuerier):
+    """
+    We may need to fail fast and return a DaemonQueryFailure instead
+    of querying the Pyre server backend - by passing a AbstractDaemonQueryFailer
+    to the constructor of the `FailableDaemonQuerier` we have the option to do this
+    """
+
+    def __init__(
+        self,
+        base_querier: AbstractDaemonQuerier,
+        daemon_query_failer: AbstractDaemonQueryFailer,
+    ) -> None:
+        super().__init__(base_querier.server_state)
+        self.base_querier: AbstractDaemonQuerier = base_querier
+        self.get_query_failure: Callable[
+            [str], Optional[daemon_query.DaemonQueryFailure]
+        ] = daemon_query_failer.query_failure
+        self.get_connection_failure: Callable[
+            [str], Optional[daemon_connection.DaemonConnectionFailure]
+        ] = daemon_query_failer.query_connection_failure
+
+    async def get_hover(
+        self,
+        path: Path,
+        position: lsp.PyrePosition,
+    ) -> Union[daemon_query.DaemonQueryFailure, GetHoverResponse]:
+        failure = self.get_query_failure(str(path))
+        return (
+            await self.base_querier.get_hover(path, position)
+            if failure is None
+            else failure
+        )
+
+    async def get_type_errors(
+        self,
+        path: Path,
+    ) -> Union[daemon_query.DaemonQueryFailure, List[error.Error]]:
+        failure = self.get_query_failure(str(path))
+        return (
+            await self.base_querier.get_type_errors(path)
+            if failure is None
+            else failure
+        )
+
+    async def get_definition_locations(
+        self,
+        path: Path,
+        position: lsp.PyrePosition,
+    ) -> Union[daemon_query.DaemonQueryFailure, GetDefinitionLocationsResponse]:
+        failure = self.get_query_failure(str(path))
+        return (
+            await self.base_querier.get_definition_locations(path, position)
+            if failure is None
+            else failure
+        )
+
+    async def get_completions(
+        self,
+        path: Path,
+        position: lsp.PyrePosition,
+    ) -> Union[daemon_query.DaemonQueryFailure, List[lsp.CompletionItem]]:
+        failure = self.get_query_failure(str(path))
+        return (
+            await self.base_querier.get_completions(path, position)
+            if failure is None
+            else failure
+        )
+
+    async def get_reference_locations(
+        self,
+        path: Path,
+        position: lsp.PyrePosition,
+    ) -> Union[daemon_query.DaemonQueryFailure, List[lsp.LspLocation]]:
+        failure = self.get_query_failure(str(path))
+        return (
+            await self.base_querier.get_reference_locations(path, position)
+            if failure is None
+            else failure
+        )
+
+    async def get_init_call_hierarchy(
+        self,
+        path: Path,
+        position: lsp.PyrePosition,
+        relation_direction: lsp.PyreCallHierarchyRelationDirection,
+    ) -> Union[daemon_query.DaemonQueryFailure, List[lsp.CallHierarchyItem]]:
+        failure = self.get_query_failure(str(path))
+        return (
+            await self.base_querier.get_init_call_hierarchy(
+                path, position, relation_direction
+            )
+            if failure is None
+            else failure
+        )
+
+    async def get_call_hierarchy_from_item(
+        self,
+        path: Path,
+        call_hierarchy_item: lsp.CallHierarchyItem,
+        relation_direction: lsp.PyreCallHierarchyRelationDirection,
+    ) -> Union[daemon_query.DaemonQueryFailure, List[lsp.CallHierarchyItem]]:
+        failure = self.get_query_failure(str(path))
+        return (
+            await self.base_querier.get_call_hierarchy_from_item(
+                path, call_hierarchy_item, relation_direction
+            )
+            if failure is None
+            else failure
+        )
+
+    async def get_type_coverage(
+        self,
+        path: Path,
+    ) -> Union[daemon_query.DaemonQueryFailure, Optional[lsp.TypeCoverageResponse]]:
+        failure = self.get_query_failure(str(path))
+        return (
+            await self.base_querier.get_type_coverage(path)
+            if failure is None
+            else failure
+        )
+
+    async def handle_file_opened(
+        self,
+        path: Path,
+        code: str,
+    ) -> Union[daemon_connection.DaemonConnectionFailure, str]:
+        failure = self.get_connection_failure(str(path))
+        return (
+            await self.base_querier.handle_file_opened(path, code)
+            if failure is None
+            else failure
+        )
+
+    async def handle_file_closed(
+        self,
+        path: Path,
+    ) -> Union[daemon_connection.DaemonConnectionFailure, str]:
+        failure = self.get_connection_failure(str(path))
+        return (
+            await self.base_querier.handle_file_closed(path)
+            if failure is None
+            else failure
+        )
+
+    async def update_overlay(
+        self,
+        path: Path,
+        code: str,
+    ) -> Union[daemon_connection.DaemonConnectionFailure, str]:
+        failure = self.get_connection_failure(str(path))
+        return (
+            await self.base_querier.update_overlay(path, code)
+            if failure is None
+            else failure
+        )
+
+    async def handle_register_client(
+        self,
+    ) -> Union[daemon_connection.DaemonConnectionFailure, str]:
+        return await self.base_querier.handle_register_client()
+
+    async def handle_dispose_client(
+        self,
+    ) -> Union[daemon_connection.DaemonConnectionFailure, str]:
+        return await self.base_querier.handle_dispose_client()
+
+
 class CodeNavigationDaemonQuerier(AbstractDaemonQuerier):
     async def get_type_errors(
         self,

client/commands/daemon_query_failer.py

@@ -0,0 +1,62 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import abc
+import re
+
+from re import Pattern
+from typing import Optional
+
+from ..language_server import daemon_connection
+
+from .daemon_query import DaemonQueryFailure
+
+
+class AbstractDaemonQueryFailer(abc.ABC):
+    @abc.abstractmethod
+    def query_failure(
+        self,
+        path: str,
+    ) -> Optional[DaemonQueryFailure]:
+        """A result of None indicates that failure should not take place"""
+        raise NotImplementedError()
+
+    def query_connection_failure(
+        self, path: str
+    ) -> Optional[daemon_connection.DaemonConnectionFailure]:
+        """A result of None indicates that failure should not take place"""
+        raise NotImplementedError()
+
+
+class DaemonQueryNoOpFailer(AbstractDaemonQueryFailer):
+    def query_failure(self, path: str) -> Optional[DaemonQueryFailure]:
+        return None
+
+    def query_connection_failure(
+        self, path: str
+    ) -> Optional[daemon_connection.DaemonConnectionFailure]:
+        return None
+
+
+class RegexDaemonQueryFailer(AbstractDaemonQueryFailer):
+    """Fails daemon queries matching a specified regex pattern"""
+
+    def __init__(self, reject_regex: str) -> None:
+        self.reject_regex = reject_regex
+        self.compiled_reject_regex: Pattern[str] = re.compile(reject_regex)
+
+    def _matches_regex(self, path: str) -> Optional[str]:
+        if self.compiled_reject_regex.match(path):
+            return f"Not querying daemon for path: {path} as matches regex: {self.reject_regex}"
+
+    def query_failure(self, path: str) -> Optional[DaemonQueryFailure]:
+        if (fail_message := self._matches_regex(path)) is not None:
+            return DaemonQueryFailure(fail_message)
+
+    def query_connection_failure(
+        self, path: str
+    ) -> Optional[daemon_connection.DaemonConnectionFailure]:
+        if (fail_message := self._matches_regex(path)) is not None:
+            return daemon_connection.DaemonConnectionFailure(fail_message)