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(
 72                    self._stream_user_updates()
 73                )
 74            except RpcError:
 75                self.synced = False
 76                self.set_status(NodeStatus.unhealthy)
 77                if self._streaming_task:
 78                    self._streaming_task.cancel()
 79            else:
 80                self.set_status(NodeStatus.healthy)
 81                logger.info("Connected to node %i", self.id)
 82
 83            await self._channel.wait_for_state_change(state)
 84
 85    async def _stream_user_updates(self):
 86        logger.debug("opened the stream")
 87        stream = self._stub.SyncUsers()
 88        while True:
 89            user_update = await self._updates_queue.get()
 90            logger.debug("got something from queue")
 91            user = user_update["user"]
 92            try:
 93                await stream.write(
 94                    UserData(
 95                        user=User(
 96                            id=user.id, username=user.username, key=user.key
 97                        ),
 98                        inbounds=[
 99                            Inbound(tag=t) for t in user_update["inbounds"]
100                        ],
101                    )
102                )
103            except RpcError:
104                self.synced = False
105                self.set_status(NodeStatus.unhealthy)
106                return
107
108    async def update_user(self, user, inbounds: set[str] | None = None):
109        if inbounds is None:
110            inbounds = set()
111
112        await self._updates_queue.put({"user": user, "inbounds": inbounds})
113
114    async def _repopulate_users(self, users_data: list[dict]) -> None:
115        updates = [
116            UserData(
117                user=User(id=u["id"], username=u["username"], key=u["key"]),
118                inbounds=[Inbound(tag=t) for t in u["inbounds"]],
119            )
120            for u in users_data
121        ]
122        await self._stub.RepopulateUsers(UsersData(users_data=updates))
123
124    async def fetch_users_stats(self):
125        response = await self._stub.FetchUsersStats(Empty())
126        return response.users_stats
127
128    async def _fetch_inbounds(self) -> list:
129        response = await self._stub.FetchInbounds(Empty())
130        return response.inbounds
131
132    async def _sync(self):
133        inbounds = await self._fetch_inbounds()
134        self.store_inbounds(inbounds)
135        users = self.list_users()
136        await self._repopulate_users(users)
137        self.synced = True
138
139    async def get_logs(self, include_buffer=True):
140        async for response in self._stub.StreamXrayLogs(
141            XrayLogsRequest(include_buffer=include_buffer)
142        ):
143            yield response.line
144
145    async def restart_xray(self, config: str):
146        try:
147            await self._stub.RestartXray(XrayConfig(configuration=config))
148            await self._sync()
149        except RpcError:
150            self.synced = False
151            self.set_status(NodeStatus.unhealthy)
152            raise
153        else:
154            self.set_status(NodeStatus.healthy)
155
156    async def get_xray_config(self):
157        response = await self._stub.FetchXrayConfig(Empty())
158        return response.configuration