1import asyncio
  2import atexit
  3import logging
  4
  5from _testcapi import INT_MAX
  6from grpc import ChannelConnectivity, RpcError
  7from grpc.aio import insecure_channel
  8
  9from .base import MarzNodeBase
 10from .database import MarzNodeDB
 11from .marznode_pb2 import (
 12    UserData,
 13    UsersData,
 14    Empty,
 15    User,
 16    Inbound,
 17    XrayLogsRequest,
 18    XrayConfig,
 19)
 20from .marznode_pb2_grpc import MarzServiceStub
 21from ..models.node import NodeStatus
 22
 23logger = logging.getLogger(__name__)
 24
 25channel_options = [
 26    ("grpc.keepalive_time_ms", 8000),
 27    ("grpc.keepalive_timeout_ms", 5000),
 28    ("grpc.http2.max_pings_without_data", 0),
 29    ("grpc.keepalive_permit_without_calls", 1),
 30    ("grpc.max_connection_idle_ms", INT_MAX),
 31    ("grpc.client_idle_timeout_ms", INT_MAX),
 32    ("grpc.max_connection_age_ms", INT_MAX),
 33]
 34
 35
 36class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB):
 37    def __init__(
 38        self, node_id: int, address: str, port: int, usage_coefficient: int = 1
 39    ):
 40        self.id = node_id
 41        self._address = address
 42        self._port = port
 43
 44        self._channel = insecure_channel(
 45            f"{self._address}:{self._port}", channel_options
 46        )
 47        self._stub = MarzServiceStub(self._channel)
 48        asyncio.create_task(self._monitor_channel())
 49        self._streaming_task = None
 50
 51        self._updates_queue = asyncio.Queue(5)
 52        self.synced = False
 53        self.usage_coefficient = usage_coefficient
 54        atexit.register(self._close_channel)
 55
 56    def _close_channel(self):
 57        asyncio.run(self._channel.close())
 58
 59    async def _monitor_channel(self):
 60        try:
 61            await asyncio.wait_for(self._channel.channel_ready(), timeout=5)
 62        except TimeoutError:
 63            logger.info("timeout for node, id: %i", self.id)
 64            self.set_status(NodeStatus.unhealthy, "timeout")
 65        while state := self._channel.get_state():
 66            logger.debug("node %i state: %s", self.id, state.value)
 67            try:
 68                if state != ChannelConnectivity.READY:
 69                    raise RpcError
 70                await self._sync()
 71                self._streaming_task = asyncio.create_task(self._stream_user_updates())
 72            except RpcError:
 73                self.synced = False
 74                self.set_status(NodeStatus.unhealthy)
 75                if self._streaming_task:
 76                    self._streaming_task.cancel()
 77            else:
 78                self.set_status(NodeStatus.healthy)
 79                logger.info("Connected to node %i", self.id)
 80
 81            await self._channel.wait_for_state_change(state)
 82
 83    async def _stream_user_updates(self):
 84        logger.debug("opened the stream")
 85        stream = self._stub.SyncUsers()
 86        while True:
 87            user_update = await self._updates_queue.get()
 88            logger.debug("got something from queue")
 89            user = user_update["user"]
 90            try:
 91                await stream.write(
 92                    UserData(
 93                        user=User(id=user.id, username=user.username, key=user.key),
 94                        inbounds=[Inbound(tag=t) for t in user_update["inbounds"]],
 95                    )
 96                )
 97            except RpcError:
 98                self.synced = False
 99                self.set_status(NodeStatus.unhealthy)
100                return
101
102    async def update_user(self, user, inbounds: set[str] | None = None):
103        if inbounds is None:
104            inbounds = set()
105
106        if self.synced:
107            await self._updates_queue.put({"user": user, "inbounds": inbounds})
108
109    async def _repopulate_users(self, users_data: list[dict]) -> None:
110        updates = [
111            UserData(
112                user=User(id=u["id"], username=u["username"], key=u["key"]),
113                inbounds=[Inbound(tag=t) for t in u["inbounds"]],
114            )
115            for u in users_data
116        ]
117        await self._stub.RepopulateUsers(UsersData(users_data=updates))
118
119    async def fetch_users_stats(self):
120        response = await self._stub.FetchUsersStats(Empty())
121        return response.users_stats
122
123    async def _fetch_inbounds(self) -> list:
124        response = await self._stub.FetchInbounds(Empty())
125        return response.inbounds
126
127    async def _sync(self):
128        inbounds = await self._fetch_inbounds()
129        self.store_inbounds(inbounds)
130        users = self.list_users()
131        await self._repopulate_users(users)
132        self.synced = True
133
134    async def get_logs(self, include_buffer=True):
135        async for response in self._stub.StreamXrayLogs(
136            XrayLogsRequest(include_buffer=include_buffer)
137        ):
138            yield response.line
139
140    async def restart_xray(self, config: str):
141        try:
142            await self._stub.RestartXray(XrayConfig(configuration=config))
143            await self._sync()
144        except RpcError:
145            self.synced = False
146            self.set_status(NodeStatus.unhealthy)
147            raise
148        else:
149            self.set_status(NodeStatus.healthy)
150
151    async def get_xray_config(self):
152        response = await self._stub.FetchXrayConfig(Empty())
153        return response.configuration