viam-sdk 0.45.2__py3-none-win_amd64.whl

viam/robot/service.py

@@ -0,0 +1,69 @@
+from typing import Any, Dict, List, Set
+
+from grpclib.server import Stream
+
+from viam import logging
+from viam.components.movement_sensor import MovementSensor
+from viam.components.sensor import Sensor
+from viam.errors import ViamGRPCError
+from viam.proto.common import ResourceName
+from viam.proto.robot import (
+    ResourceNamesRequest,
+    ResourceNamesResponse,
+    StopAllRequest,
+    StopAllResponse,
+    UnimplementedRobotServiceBase,
+)
+from viam.resource.rpc_service_base import ResourceRPCServiceBase
+from viam.utils import resource_names_for_resource, struct_to_dict
+
+LOGGER = logging.getLogger(__name__)
+
+
+class RobotService(UnimplementedRobotServiceBase, ResourceRPCServiceBase):
+    def _generate_metadata(self) -> List[ResourceName]:
+        md: Set[ResourceName] = set()
+
+        for resource in self.manager.resources.values():
+            # If the resource is a MovementSensor, DO NOT include Sensor as well (it will get added via MovementSensor)
+            if resource.API == Sensor.API and MovementSensor.get_resource_name(resource.name) in self.manager.resources:
+                continue
+
+            md.update(resource_names_for_resource(resource))
+
+        return list(md)
+
+    async def ResourceNames(self, stream: Stream[ResourceNamesRequest, ResourceNamesResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        metadata = self._generate_metadata()
+        response = ResourceNamesResponse(resources=metadata)
+        await stream.send_message(response)
+
+    async def StopAll(self, stream: Stream[StopAllRequest, StopAllResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+
+        extra: Dict[ResourceName, Dict[str, Any]] = {}
+        for ex in request.extra:
+            extra[ex.name] = struct_to_dict(ex.params)
+
+        errors: List[str] = []
+        for component in self.manager.resources.values():
+            if callable(getattr(component, "stop", None)):
+                try:
+                    rn = component.get_resource_name(component.name)
+                    if rn in extra:
+                        try:
+                            await component.stop(extra=extra[rn])  # type: ignore
+                        except TypeError:
+                            await component.stop()  # type: ignore
+                    else:
+                        await component.stop()  # type: ignore
+                except Exception:
+                    LOGGER.exception(f"Failed to stop component named {component.name}")
+                    errors.append(component.name)
+
+        if errors:
+            raise ViamGRPCError(f'Failed to stop components named {", ".join(errors)}')
+        await stream.send_message(StopAllResponse())

viam/rpc/dial.py

@@ -0,0 +1,420 @@
+import ctypes
+import pathlib
+import re
+import socket
+import ssl
+import sys
+import uuid
+import warnings
+from dataclasses import dataclass
+from typing import Callable, Literal, Optional, Tuple, Type, Union
+
+from grpclib.client import Channel, Stream
+from grpclib.const import Cardinality
+from grpclib.events import SendRequest, listen
+from grpclib.metadata import Deadline, _MetadataLike
+from grpclib.protocol import H2Protocol
+from grpclib.stream import _RecvType, _SendType
+from typing_extensions import Self
+
+from viam import logging
+from viam.errors import InsecureConnectionError, ViamError
+from viam.proto.rpc.auth import AuthenticateRequest, AuthServiceStub
+from viam.proto.rpc.auth import Credentials as PBCredentials
+from viam.utils import to_thread
+from viam.version_metadata import API_VERSION, SDK_VERSION
+
+LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class Credentials:
+    """Credentials to connect to the robot and the Viam app."""
+
+    type: Union[Literal["robot-location-secret"], Literal["robot-secret"], Literal["api-key"]]
+    """The type of credential
+    """
+
+    payload: str
+    """The credential
+    """
+
+
+class DialOptions:
+    disable_webrtc: bool
+    """Bypass Web RTC and connect directly to the robot.
+    """
+
+    auth_entity: Optional[str]
+    """The URL to authenticate against. Should be used if the address passed in and FQDN of the server do not match.
+    """
+
+    credentials: Optional[Credentials]
+    """Credentials for connecting to the robot
+    """
+
+    insecure: bool = False
+    """Determine if the RPC connection is TLS based. Must be provided to
+    establish an insecure connection. Otherwise, a TLS based connection
+    will be assumed."""
+
+    allow_insecure_downgrade: bool = False
+    """Allow the RPC connection to be downgraded to an insecure connection
+    if detected. This is only used when credentials are not present."""
+
+    allow_insecure_with_creds_downgrade: bool = False
+    """Allow the RPC connection to be downgraded to an insecure connection
+    if detected, even with credentials present. This is generally
+    unsafe to use, but can be requested."""
+
+    max_reconnect_attempts: int = 3
+    """Max number of times the client attempts to reconnect when connection is lost"""
+
+    initial_connection_attempts: int = 3
+    """Max number of times the client will attempt to establish an initial connection
+    If set to a non-positive integer, then there will be no limit to initial connection attempts"""
+
+    initial_connection_attempt_timeout: float
+    """Number of seconds before dial connection times out on initial connection attempts
+    Defaults to whatever value is set in the `timeout` field"""
+
+    timeout: float = 20
+    """Number of seconds before the dial connection times out
+    Set to 20sec to match _defaultOfferDeadline in goutils/rpc/wrtc_call_queue.go"""
+
+    def __init__(
+        self,
+        *,
+        disable_webrtc: bool = False,
+        auth_entity: Optional[str] = None,
+        credentials: Optional[Credentials] = None,
+        insecure: bool = False,
+        allow_insecure_downgrade: bool = False,
+        allow_insecure_with_creds_downgrade: bool = False,
+        max_reconnect_attempts: int = 3,
+        timeout: float = 20,
+        initial_connection_attempts: int = 3,
+        initial_connection_attempt_timeout: Optional[float] = None,
+    ) -> None:
+        self.disable_webrtc = disable_webrtc
+        self.auth_entity = auth_entity
+        self.credentials = credentials
+        self.insecure = insecure
+        self.allow_insecure_downgrade = allow_insecure_downgrade
+        self.allow_insecure_with_creds_downgrade = allow_insecure_with_creds_downgrade
+        self.max_reconnect_attempts = max_reconnect_attempts
+        self.timeout = timeout
+        self.initial_connection_attempts = initial_connection_attempts
+        self.initial_connection_attempt_timeout = initial_connection_attempt_timeout if initial_connection_attempt_timeout else timeout
+
+    @classmethod
+    def with_api_key(cls, api_key: str, api_key_id: str) -> Self:
+        """Create DialOptions with an API key for credentials and default values for other arguments.
+
+        Args:
+            api_key (str): your API key
+            api_key_id (str): your API key ID. Must be a valid UUID
+
+        Raises:
+            ValueError: Raised if the api_key_id is not a valid UUID
+
+        Returns:
+            Self: the DialOptions
+        """
+        try:
+            uuid.UUID(api_key_id)
+        except ValueError:
+            raise ValueError(f"{api_key_id} is not a valid UUID")
+
+        credentials = Credentials(type="api-key", payload=api_key)
+        return cls(credentials=credentials, auth_entity=api_key_id)
+
+
+def _host_port_from_url(url) -> Tuple[Optional[str], Optional[int]]:
+    query = "(?:.*://)?(?P<host>[^:/ ]+).?(?P<port>[0-9]*).*"
+    match = re.search(query, url)
+
+    if not match:
+        return (None, None)
+
+    host = match.group("host")
+    try:
+        port = int(match.group("port"))
+    except ValueError:
+        port = None
+    return (host, port)
+
+
+async def _get_access_token(channel: Channel, address: str, opts: DialOptions) -> str:
+    entity = opts.auth_entity if opts.auth_entity else re.sub(r"^(.*:\/\/)/", "", address)
+    creds = PBCredentials(
+        type=opts.credentials.type if opts.credentials else "", payload=opts.credentials.payload if opts.credentials else ""
+    )
+    request = AuthenticateRequest(entity=entity, credentials=creds)
+
+    auth_service = AuthServiceStub(channel=channel)
+    response = await auth_service.Authenticate(request)
+    return response.access_token
+
+
+class AuthenticatedChannel(Channel):
+    _metadata: _MetadataLike
+
+    def __init__(
+        self,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        *,
+        ssl: Union[None, bool, ssl.SSLContext] = None,
+        server_hostname: Optional[str] = None,
+    ):
+        super().__init__(host, port, ssl=ssl)
+        self._server_hostname = server_hostname
+
+    async def _create_connection(self) -> H2Protocol:
+        _, protocol = await self._loop.create_connection(
+            self._protocol_factory,
+            self._host,
+            self._port,
+            ssl=self._ssl,
+            server_hostname=self._server_hostname,
+        )
+        return protocol
+
+    def request(
+        self,
+        name: str,
+        cardinality: Cardinality,
+        request_type: Type[_SendType],
+        reply_type: Type[_RecvType],
+        *,
+        timeout: Optional[float] = None,
+        deadline: Optional[Deadline] = None,
+        metadata: Optional[_MetadataLike] = None,
+    ) -> Stream[_SendType, _RecvType]:
+        if not metadata and hasattr(self, "_metadata"):
+            metadata = self._metadata
+        return super().request(name, cardinality, request_type, reply_type, timeout=timeout, deadline=deadline, metadata=metadata)
+
+
+@dataclass
+class ViamChannel:
+    channel: Channel
+    release: Callable[[], None]
+    _closed: bool = False
+
+    def close(self):
+        if not self._closed:
+            try:
+                self.channel.close()
+            except RuntimeError as e:
+                # ignore event loop is closed errors - robot is getting shutdown
+                if len(e.args) > 0 and e.args[0] == "Event loop is closed":
+                    LOGGER.debug("ViamChannel might not have shut down cleanly - Event loop was closed")
+                    return
+                raise
+            finally:
+                self.release()
+                self._closed = True
+
+    def __del__(self):
+        self.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+
+class _Runtime:
+    _lib: ctypes.CDLL
+    _ptr: ctypes.c_void_p
+
+    def __init__(self) -> None:
+        suffix = "dylib" if sys.platform == "darwin" else "so" if "linux" in sys.platform else "dll"
+        LOGGER.debug("Creating new viam-rust-utils runtime")
+        libname = pathlib.Path(__file__).parent.absolute() / f"libviam_rust_utils.{suffix}"
+        self._lib = ctypes.CDLL(libname.__str__())
+        self._lib.init_rust_runtime.argtypes = ()
+        self._lib.init_rust_runtime.restype = ctypes.c_void_p
+
+        self._lib.dial.argtypes = (
+            ctypes.c_char_p,
+            ctypes.c_char_p,
+            ctypes.c_char_p,
+            ctypes.c_char_p,
+            ctypes.c_bool,
+            ctypes.c_float,
+            ctypes.c_void_p,
+        )
+        self._lib.dial.restype = ctypes.c_void_p
+
+        self._lib.free_rust_runtime.argtypes = (ctypes.c_void_p,)
+        self._lib.free_rust_runtime.restype = None
+
+        self._lib.free_string.argtypes = (ctypes.c_void_p,)
+        self._lib.free_string.restype = None
+
+        self._ptr = self._lib.init_rust_runtime()
+
+    async def dial(self, address: str, options: DialOptions) -> Tuple[Optional[str], ctypes.c_void_p]:
+        type = options.credentials.type if options.credentials else ""
+        payload = options.credentials.payload if options.credentials else ""
+        insecure = (
+            options.insecure
+            or options.allow_insecure_with_creds_downgrade
+            or (not type and not payload and options.allow_insecure_downgrade)
+        )
+
+        LOGGER.debug(f"Dialing {address} using viam-rust-utils library")
+        path_ptr = await to_thread(
+            self._lib.dial,
+            address.encode("utf-8"),
+            options.auth_entity.encode("utf-8") if options.auth_entity else None,
+            type.encode("utf-8") if type else None,
+            payload.encode("utf-8") if payload else None,
+            insecure,
+            ctypes.c_float(options.timeout),
+            self._ptr,
+        )
+        path = ctypes.cast(path_ptr, ctypes.c_char_p).value
+        path = path.decode("utf-8") if path else ""
+        return (path, path_ptr)
+
+    def release(self):
+        LOGGER.debug("Freeing viam-rust-utils runtime")
+        self._lib.free_rust_runtime(self._ptr)
+
+    def free_str(self, ptr: ctypes.c_void_p):
+        LOGGER.debug("Freeing socket string")
+        self._lib.free_string(ptr)
+
+
+async def dial(address: str, options: Optional[DialOptions] = None) -> ViamChannel:
+    options = options if options else DialOptions()
+    timeout = options.timeout
+    options.timeout = options.initial_connection_attempt_timeout
+    if options.initial_connection_attempts == 0:
+        options.initial_connection_attempts = -1
+    attempt_countdown = options.initial_connection_attempts
+    exception: Exception
+    while attempt_countdown != 0:
+        try:
+            chan = await _dial_inner(address, options)
+            options.timeout = timeout
+            return chan
+        except Exception as e:
+            exception = e
+            attempt_countdown -= 1
+    # the only way we could get here is if we failed at least once which means we've set the
+    # exception, so typechecker concerns about a possibly unbounded variable are unfounded
+    raise exception  # type: ignore
+
+
+def _create_chan(path: str) -> Channel:
+    if sys.platform == "win32" or sys.platform == "cygwin":
+        # we have to use a TCP connection, so we want a host and port for our channel.
+        host, port = _host_port_from_url(path)
+        return Channel(host=host, port=port, ssl=None)
+    # we're not on windows and so can use a UDS
+    return Channel(path=path, ssl=None)
+
+
+async def _dial_inner(address: str, options: Optional[DialOptions] = None) -> ViamChannel:
+    async def send_request(event: SendRequest):
+        event.metadata["viam-client"] = f"python;v{SDK_VERSION};v{API_VERSION}"
+
+    opts = options if options else DialOptions()
+    if opts.disable_webrtc:
+        channel = await _dial_direct(address, options)
+        listen(channel, SendRequest, send_request)
+        return ViamChannel(channel, lambda: None)
+    runtime = _Runtime()
+    path, path_ptr = await runtime.dial(address, opts)
+    if path:
+        LOGGER.info(f"Connecting to socket: {path}")
+        chan = _create_chan(path)
+        listen(chan, SendRequest, send_request)
+
+        def release():
+            runtime.free_str(path_ptr)
+            runtime.release()
+
+        channel = ViamChannel(chan, release)
+        return channel
+
+    runtime.release()
+    raise ViamError(f"Unable to establish a connection to {address}")
+
+
+async def _dial_direct(address: str, options: Optional[DialOptions] = None) -> Channel:
+    opts = options if options else DialOptions()
+    insecure = opts.insecure
+
+    if pathlib.Path(address).is_socket():
+        return Channel(path=address)
+
+    host, port = _host_port_from_url(address)
+    if not port:
+        port = 80 if insecure else 443
+    server_hostname = host
+
+    if insecure:
+        ctx = None
+    else:
+        is_local_host = host is not None and (host.startswith("localhost") or host.startswith("0.0.0.0") or host.startswith("127."))
+        if is_local_host:
+            ctx = ssl._create_unverified_context(purpose=ssl.Purpose.SERVER_AUTH)
+        else:
+            ctx = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
+        ctx.minimum_version = ssl.TLSVersion.TLSv1_2
+        ctx.set_ciphers("ECDHE+AESGCM:ECDHE+CHACHA20:DHE+AESGCM:DHE+CHACHA20")
+        ctx.set_alpn_protocols(["h2"])
+
+        if (
+            options is not None
+            and options.auth_entity
+            and host != options.auth_entity
+            and options.credentials is not None
+            and options.credentials.type != "api-key"
+        ):
+            server_hostname = options.auth_entity
+
+        # Test if downgrade is required.
+        downgrade = False
+        with socket.create_connection((host, port), timeout=opts.timeout) as sock:
+            try:
+                with ctx.wrap_socket(sock, server_hostname=server_hostname) as ssock:
+                    _ = ssock.version()
+            except ssl.SSLError as e:
+                if e.reason != "WRONG_VERSION_NUMBER":
+                    raise e
+                downgrade = True
+
+        if downgrade:
+            if opts.credentials:
+                if not opts.allow_insecure_with_creds_downgrade:
+                    raise InsecureConnectionError(address, authenticated=True)
+            elif not opts.allow_insecure_downgrade:
+                raise InsecureConnectionError(address)
+            ctx = None
+
+    if opts.credentials:
+        channel = AuthenticatedChannel(host, port, ssl=ctx, server_hostname=server_hostname)
+        access_token = await _get_access_token(channel, address, opts)
+        metadata = {"authorization": f"Bearer {access_token}"}
+        channel._metadata = metadata
+    else:
+        channel = Channel(host, port, ssl=ctx)
+
+    return channel
+
+
+async def dial_direct(address: str, options: Optional[DialOptions] = None) -> Channel:
+    warnings.warn("dial_direct is deprecated. Use rpc.dial.dial instead.", DeprecationWarning, stacklevel=2)
+    return await _dial_direct(address, options)
+
+
+async def _dial_app(app_url: str) -> Channel:
+    return await _dial_direct(app_url)

viam/rpc/server.py

@@ -0,0 +1,201 @@
+from typing import TYPE_CHECKING, Callable, List, Optional
+
+from grpclib import GRPCError, Status
+from grpclib._typing import IServable
+from grpclib.const import Handler
+from grpclib.events import RecvRequest, listen
+from grpclib.reflection.service import ServerReflection
+from grpclib.server import Server as GRPCServer
+from grpclib.server import Stream
+from grpclib.utils import graceful_exit
+
+from viam import logging
+from viam.errors import ViamGRPCError
+from viam.resource.base import ResourceBase
+from viam.resource.manager import ResourceManager
+from viam.resource.registry import Registry
+from viam.resource.rpc_service_base import ResourceRPCServiceBase
+from viam.robot.service import RobotService
+
+from .signaling import SignalingService
+
+if TYPE_CHECKING:
+    from viam.module.service import ModuleRPCService
+
+LOGGER = logging.getLogger(__name__)
+
+
+class Server(ResourceManager):
+    """
+    gRPC Server
+    """
+
+    def __init__(self, resources: List[ResourceBase], *, module_service: Optional["ModuleRPCService"] = None):
+        """
+        Initialize the Server with a list of resources
+        to be managed.
+
+        Args:
+            resources (List[ComponentBase]): List of resources to be managed
+        """
+        super().__init__(resources)
+
+        services = [SignalingService(), RobotService(manager=self)]
+        for registration in Registry.REGISTERED_APIS().values():
+            if issubclass(registration.rpc_service, ResourceRPCServiceBase):
+                services.append(registration.rpc_service(manager=self))
+            else:
+                services.append(registration.rpc_service())
+
+        if module_service is not None:
+            services.append(module_service)
+        services = ServerReflection.extend(services)
+        services = _patch_mappings(services)
+
+        self._server = GRPCServer(services)
+
+    async def _grpc_recvrequest_handler(self, event: RecvRequest):
+        host = None
+        port = None
+        address = event.peer.addr()
+        if address:
+            host = address[0]
+            port = address[1]
+        method_func = event.method_func
+
+        async def log_resource_name(stream: Stream):
+            recv_msg = stream.recv_message
+
+            async def rcv_and_log_msg():
+                msg = await recv_msg()
+                log_msg = f"[gRPC] Received message from {host or 'xxxx'}:{port or 'xxxx'} - {event.method_name}"
+                if msg and hasattr(msg, "name"):
+                    log_msg += f" for resource named: {msg.name}"
+                LOGGER.debug(log_msg)
+                return msg
+
+            stream.recv_message = rcv_and_log_msg
+            try:
+                return await method_func(stream)
+            finally:
+                LOGGER.debug(f"[gRPC] Finished call from {host or 'xxxx'}:{port or 'xxxx'} - {event.method_name}")
+
+        event.method_func = log_resource_name
+
+    async def serve(
+        self,
+        host: Optional[str] = "localhost",
+        port: Optional[int] = 9090,
+        log_level: Optional[int] = logging.INFO,
+        *,
+        path: Optional[str] = None,
+    ):
+        """
+        Server the gRPC server on the provided host and port
+
+        Args:
+            host (Optional[str], optional): Desired hostname of the server. Defaults to "localhost".
+            port (Optional[int], optional): Desired port of the server. Defaults to 9090.
+            log_level (Optional[int], optional): The minimum log level. To not receive any logs, set to None. Defaults to logging.INFO.
+            path (Optional[str], optional): UNIX socket path. Takes precedence over `host` and `port` if set. Defaults to None.
+        """
+        if log_level is None:
+            logging.silence()
+        else:
+            logging.setLevel(log_level)
+        listen(self._server, RecvRequest, self._grpc_recvrequest_handler)
+
+        with graceful_exit([self._server]):
+            if path:
+                await self._server.start(path=path)
+                LOGGER.info(f"Serving on {path}")
+            else:
+                await self._server.start(host, port)
+                LOGGER.info(f"Serving on {host}:{port}")
+            await self._server.wait_closed()
+            await self.close()
+            LOGGER.debug("gRPC server closed")
+
+    @classmethod
+    async def create_and_serve(
+        cls,
+        components: List[ResourceBase],
+        host: Optional[str] = "localhost",
+        port: Optional[int] = 9090,
+        log_level: int = logging.INFO,
+        *,
+        path: Optional[str] = None,
+    ):
+        """
+        Convenience method to create and start the server.
+
+        Args:
+            components (List[ComponentBase]): List of components to manage
+            host (str, optional): Desired hostname. Defaults to "localhost".
+            port (int, optional): Desired port. Defaults to 9090.
+            log_level(int, optional): The minimum log level.
+                To not receive any logs, set to None.
+                Defaults to logging.INFO
+            path (Optional[str], optional): UNIX socket path. Takes precedence over `host` and `port` if set. Defaults to None.
+        """
+        server = cls(components)
+        await server.serve(host, port, log_level, path=path)
+
+
+def _grpc_error_wrapper(func: Callable):
+    """
+    Wrap a function so that any exceptions get raised as GRPCErrors to the client.
+
+    Args:
+        func (Callable): The function that should be wrapped
+
+    Returns:
+        The method that is now wrapped to raise GRPCErrors
+    """
+
+    async def interceptor(*args, **kwargs):
+        try:
+            new_func = await func(*args, **kwargs)
+            return new_func
+        except GRPCError:
+            raise
+        except ViamGRPCError as e:
+            raise e.grpc_error
+        except Exception as e:
+            tb = e.__traceback__
+            file_name = None
+            func_name = None
+            line_num = None
+            # only print the last entry in the stacktrace - not perfect but gives users a starting point
+            while tb is not None:
+                file_name = tb.tb_frame.f_code.co_filename
+                func_name = tb.tb_frame.f_code.co_name
+                line_num = tb.tb_lineno
+                tb = tb.tb_next
+            raise GRPCError(Status.UNKNOWN, f"{e.__class__.__name__} - {e} - {file_name=} {func_name=} {line_num=}")
+
+    return interceptor
+
+
+def _patch_mappings(services: List[IServable]) -> List[IServable]:
+    """Replace the methods of all given services with a wrapped method that has error handling
+
+    Args:
+        services (List[IServable]): The services that should be patched
+
+    Returns:
+        services (List[IServable]): The patched services with new mapping functions
+    """
+    for service in services:
+
+        def patch_mapping():
+            mapping = service.__mapping__()
+            new_mapping = {}
+            for method, handler in mapping.items():
+                new_method = _grpc_error_wrapper(handler[0])
+                new_mapping[method] = Handler(new_method, *handler[1:])
+
+            return lambda: new_mapping
+
+        service.__mapping__ = patch_mapping()
+    return services