gitscope
/
← marzneshin / app/marznode/grpclib.py / blame
Files Commits 446 Branches 1 Tags 43 Markers 4 README License
https://dev.kronoz.win/r/marzneshin.git
app/marznode/grpclib.py · v0.6.4
Source
1 f6a32ea Dawsh 2y ago import asyncio
2 f6a32ea Dawsh 2y ago import atexit
3 f6a32ea Dawsh 2y ago import logging
4 f6a32ea Dawsh 2y ago import ssl
5 f6a32ea Dawsh 2y ago import tempfile
6 f6a32ea Dawsh 2y ago
7 f6a32ea Dawsh 2y ago from grpclib import GRPCError
8 f6a32ea Dawsh 2y ago from grpclib.client import Channel
9 f6a32ea Dawsh 2y ago from grpclib.exceptions import StreamTerminatedError
10 f6a32ea Dawsh 2y ago
11 f6a32ea Dawsh 2y ago from .base import MarzNodeBase
12 f6a32ea Dawsh 2y ago from .database import MarzNodeDB
13 f6a32ea Dawsh 2y ago from .marznode_grpc import MarzServiceStub
14 01cdf2e Dawsh 2y ago from .marznode_pb2 import (
15 01cdf2e Dawsh 2y ago UserData,
16 01cdf2e Dawsh 2y ago UsersData,
17 01cdf2e Dawsh 2y ago Empty,
18 01cdf2e Dawsh 2y ago User,
19 01cdf2e Dawsh 2y ago Inbound,
20 f5c168d Mardin 2y ago BackendConfig,
21 f5c168d Mardin 2y ago BackendLogsRequest,
22 f5c168d Mardin 2y ago Backend,
23 f5c168d Mardin 2y ago RestartBackendRequest,
24 f5c168d Mardin 2y ago BackendStats,
25 01cdf2e Dawsh 2y ago )
26 f6a32ea Dawsh 2y ago from ..models.node import NodeStatus
27 f6a32ea Dawsh 2y ago
28 f6a32ea Dawsh 2y ago logger = logging.getLogger(__name__)
29 f6a32ea Dawsh 2y ago
30 f6a32ea Dawsh 2y ago
31 f6a32ea Dawsh 2y ago def string_to_temp_file(content: str):
32 01cdf2e Dawsh 2y ago file = tempfile.NamedTemporaryFile(mode="w+t")
33 f6a32ea Dawsh 2y ago file.write(content)
34 f6a32ea Dawsh 2y ago file.flush()
35 f6a32ea Dawsh 2y ago return file
36 f6a32ea Dawsh 2y ago
37 f6a32ea Dawsh 2y ago
38 f6a32ea Dawsh 2y ago class MarzNodeGRPCLIB(MarzNodeBase, MarzNodeDB):
39 01cdf2e Dawsh 2y ago def __init__(
40 01cdf2e Dawsh 2y ago self,
41 01cdf2e Dawsh 2y ago node_id: int,
42 01cdf2e Dawsh 2y ago address: str,
43 01cdf2e Dawsh 2y ago port: int,
44 01cdf2e Dawsh 2y ago ssl_key: str,
45 01cdf2e Dawsh 2y ago ssl_cert: str,
46 01cdf2e Dawsh 2y ago usage_coefficient: int = 1,
47 01cdf2e Dawsh 2y ago ):
48 f6a32ea Dawsh 2y ago self.id = node_id
49 f6a32ea Dawsh 2y ago self._address = address
50 f6a32ea Dawsh 2y ago self._port = port
51 f6a32ea Dawsh 2y ago
52 f6a32ea Dawsh 2y ago self._key_file = string_to_temp_file(ssl_key)
53 f6a32ea Dawsh 2y ago self._cert_file = string_to_temp_file(ssl_cert)
54 f6a32ea Dawsh 2y ago
55 f6a32ea Dawsh 2y ago ctx = ssl.create_default_context()
56 f6a32ea Dawsh 2y ago ctx.load_cert_chain(self._cert_file.name, self._key_file.name)
57 f6a32ea Dawsh 2y ago ctx.check_hostname = False
58 f6a32ea Dawsh 2y ago ctx.verify_mode = ssl.CERT_NONE
59 f6a32ea Dawsh 2y ago
60 f6a32ea Dawsh 2y ago self._channel = Channel(self._address, self._port, ssl=ctx)
61 f6a32ea Dawsh 2y ago self._stub = MarzServiceStub(self._channel)
62 961d06d Dawsh 2y ago self._monitor_task = asyncio.create_task(self._monitor_channel())
63 f6a32ea Dawsh 2y ago self._streaming_task = None
64 f6a32ea Dawsh 2y ago
65 f6a32ea Dawsh 2y ago self._updates_queue = asyncio.Queue(1)
66 f6a32ea Dawsh 2y ago self.synced = False
67 f6a32ea Dawsh 2y ago self.usage_coefficient = usage_coefficient
68 f6a32ea Dawsh 2y ago atexit.register(self._channel.close)
69 f6a32ea Dawsh 2y ago
70 961d06d Dawsh 2y ago async def stop(self):
71 961d06d Dawsh 2y ago self._channel.close()
72 961d06d Dawsh 2y ago self._monitor_task.cancel()
73 961d06d Dawsh 2y ago
74 f6a32ea Dawsh 2y ago async def _monitor_channel(self):
75 6504ac9 Dawsh 2y ago while state := self._channel._state:
76 f6a32ea Dawsh 2y ago logger.debug("node %i channel state: %s", self.id, state.value)
77 f6a32ea Dawsh 2y ago try:
78 f6a32ea Dawsh 2y ago await asyncio.wait_for(self._channel.__connect__(), timeout=2)
79 f6a32ea Dawsh 2y ago except Exception:
80 512d183 Mrdn 2y ago logger.debug("timeout for node, id: %i", self.id)
81 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.unhealthy, "timeout")
82 f6a32ea Dawsh 2y ago self.synced = False
83 f6a32ea Dawsh 2y ago if self._streaming_task:
84 f6a32ea Dawsh 2y ago self._streaming_task.cancel()
85 f6a32ea Dawsh 2y ago else:
86 f6a32ea Dawsh 2y ago if not self.synced:
87 4a2faf0 Dawsh 2y ago try:
88 4a2faf0 Dawsh 2y ago await self._sync()
89 4a2faf0 Dawsh 2y ago except:
90 4a2faf0 Dawsh 2y ago pass
91 4a2faf0 Dawsh 2y ago else:
92 4a2faf0 Dawsh 2y ago self._streaming_task = asyncio.create_task(
93 4a2faf0 Dawsh 2y ago self._stream_user_updates()
94 4a2faf0 Dawsh 2y ago )
95 4a2faf0 Dawsh 2y ago self.set_status(NodeStatus.healthy)
96 4a2faf0 Dawsh 2y ago logger.info("Connected to node %i", self.id)
97 512d183 Mrdn 2y ago await asyncio.sleep(10)
98 f6a32ea Dawsh 2y ago
99 f6a32ea Dawsh 2y ago async def _stream_user_updates(self):
100 f6a32ea Dawsh 2y ago try:
101 f6a32ea Dawsh 2y ago async with self._stub.SyncUsers.open() as stream:
102 4a2faf0 Dawsh 2y ago logger.debug("opened the stream")
103 f6a32ea Dawsh 2y ago while True:
104 f6a32ea Dawsh 2y ago user_update = await self._updates_queue.get()
105 4a2faf0 Dawsh 2y ago logger.debug("got something from queue")
106 f6a32ea Dawsh 2y ago user = user_update["user"]
107 f6a32ea Dawsh 2y ago await stream.send_message(
108 01cdf2e Dawsh 2y ago UserData(
109 f1ab16c Dawsh 2y ago user=User(
110 f1ab16c Dawsh 2y ago id=user.id,
111 f1ab16c Dawsh 2y ago username=user.username,
112 f1ab16c Dawsh 2y ago key=user.key,
113 f1ab16c Dawsh 2y ago ),
114 f1ab16c Dawsh 2y ago inbounds=[
115 f1ab16c Dawsh 2y ago Inbound(tag=t) for t in user_update["inbounds"]
116 f1ab16c Dawsh 2y ago ],
117 01cdf2e Dawsh 2y ago )
118 01cdf2e Dawsh 2y ago )
119 f6a32ea Dawsh 2y ago except (OSError, ConnectionError, GRPCError, StreamTerminatedError):
120 f6a32ea Dawsh 2y ago logger.info("node %i detached", self.id)
121 f6a32ea Dawsh 2y ago self.synced = False
122 f6a32ea Dawsh 2y ago
123 f6a32ea Dawsh 2y ago async def update_user(self, user, inbounds: set[str] | None = None):
124 f6a32ea Dawsh 2y ago if inbounds is None:
125 f6a32ea Dawsh 2y ago inbounds = set()
126 f6a32ea Dawsh 2y ago
127 4a2faf0 Dawsh 2y ago await self._updates_queue.put({"user": user, "inbounds": inbounds})
128 f6a32ea Dawsh 2y ago
129 f6a32ea Dawsh 2y ago async def _repopulate_users(self, users_data: list[dict]) -> None:
130 01cdf2e Dawsh 2y ago updates = [
131 01cdf2e Dawsh 2y ago UserData(
132 01cdf2e Dawsh 2y ago user=User(id=u["id"], username=u["username"], key=u["key"]),
133 01cdf2e Dawsh 2y ago inbounds=[Inbound(tag=t) for t in u["inbounds"]],
134 01cdf2e Dawsh 2y ago )
135 01cdf2e Dawsh 2y ago for u in users_data
136 01cdf2e Dawsh 2y ago ]
137 f6a32ea Dawsh 2y ago await self._stub.RepopulateUsers(UsersData(users_data=updates))
138 f6a32ea Dawsh 2y ago
139 f6a32ea Dawsh 2y ago async def fetch_users_stats(self):
140 f6a32ea Dawsh 2y ago response = await self._stub.FetchUsersStats(Empty())
141 f6a32ea Dawsh 2y ago return response.users_stats
142 f6a32ea Dawsh 2y ago
143 f5c168d Mardin 2y ago async def _fetch_backends(self) -> list:
144 f5c168d Mardin 2y ago response = await self._stub.FetchBackends(Empty())
145 f5c168d Mardin 2y ago return response.backends
146 f6a32ea Dawsh 2y ago
147 f6a32ea Dawsh 2y ago async def _sync(self):
148 f5c168d Mardin 2y ago backends = await self._fetch_backends()
149 f5c168d Mardin 2y ago self.store_backends(backends)
150 f6a32ea Dawsh 2y ago users = self.list_users()
151 f6a32ea Dawsh 2y ago await self._repopulate_users(users)
152 f6a32ea Dawsh 2y ago self.synced = True
153 512d183 Mrdn 2y ago
154 f5c168d Mardin 2y ago async def get_logs(self, name: str = "xray", include_buffer=True):
155 f5c168d Mardin 2y ago async with self._stub.StreamBackendLogs.open() as stm:
156 f1ab16c Dawsh 2y ago await stm.send_message(
157 f5c168d Mardin 2y ago BackendLogsRequest(
158 f5c168d Mardin 2y ago backend_name=name, include_buffer=include_buffer
159 f5c168d Mardin 2y ago )
160 f1ab16c Dawsh 2y ago )
161 512d183 Mrdn 2y ago while True:
162 512d183 Mrdn 2y ago response = await stm.recv_message()
163 512d183 Mrdn 2y ago yield response.line
164 512d183 Mrdn 2y ago
165 f5c168d Mardin 2y ago async def restart_backend(
166 f5c168d Mardin 2y ago self, name: str, config: str, config_format: int
167 f5c168d Mardin 2y ago ):
168 758c903 Dawsh 2y ago try:
169 f5c168d Mardin 2y ago await self._stub.RestartBackend(
170 f5c168d Mardin 2y ago RestartBackendRequest(
171 f5c168d Mardin 2y ago backend_name=name,
172 f5c168d Mardin 2y ago config=BackendConfig(
173 f5c168d Mardin 2y ago configuration=config, config_format=config_format
174 f5c168d Mardin 2y ago ),
175 f5c168d Mardin 2y ago )
176 f5c168d Mardin 2y ago )
177 758c903 Dawsh 2y ago await self._sync()
178 758c903 Dawsh 2y ago except:
179 758c903 Dawsh 2y ago self.synced = False
180 758c903 Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
181 758c903 Dawsh 2y ago raise
182 758c903 Dawsh 2y ago else:
183 758c903 Dawsh 2y ago self.set_status(NodeStatus.healthy)
184 512d183 Mrdn 2y ago
185 f5c168d Mardin 2y ago async def get_backend_config(self, name: str):
186 f5c168d Mardin 2y ago response: BackendConfig = await self._stub.FetchBackendConfig(
187 f5c168d Mardin 2y ago Backend(name=name)
188 f5c168d Mardin 2y ago )
189 f5c168d Mardin 2y ago return response.configuration, response.config_format
190 f5c168d Mardin 2y ago
191 f5c168d Mardin 2y ago async def get_backend_stats(self, name: str):
192 f5c168d Mardin 2y ago response: BackendStats = await self._stub.GetBackendStats(
193 f5c168d Mardin 2y ago Backend(name=name)
194 f5c168d Mardin 2y ago )
195 f5c168d Mardin 2y ago return response