diff --git a/app/db/crud.py b/app/db/crud.py index 45f1669..a43f706 100644 --- a/app/db/crud.py +++ b/app/db/crud.py @@ -561,7 +561,8 @@ def get_nodes_usage(db: Session, start: datetime, end: datetime) -> List[NodeUsa def create_node(db: Session, node: NodeCreate): dbnode = Node(name=node.name, address=node.address, - port=node.port) + port=node.port, + connection_backend=node.connection_backend) db.add(dbnode) db.commit() @@ -600,14 +601,13 @@ def update_node(db: Session, dbnode: Node, modify: NodeModify): return dbnode -def update_node_status(db: Session, dbnode: Node, status: NodeStatus, message: str = None, version: str = None): - dbnode.status = status - dbnode.message = message - dbnode.xray_version = version - dbnode.last_status_change = datetime.utcnow() +def update_node_status(db: Session, node_id: int, status: NodeStatus, message: str = None, version: str = None): + db_node = db.query(Node).where(Node.id == node_id).first() + db_node.status = status + if message: + db_node.message = message + db_node.last_status_change = datetime.utcnow() db.commit() - db.refresh(dbnode) - return dbnode def create_notification_reminder( diff --git a/app/db/migrations/versions/20faa9f18c0a_init.py b/app/db/migrations/versions/20faa9f18c0a_init.py index 0bcbd56..4b8ed4c 100644 --- a/app/db/migrations/versions/20faa9f18c0a_init.py +++ b/app/db/migrations/versions/20faa9f18c0a_init.py @@ -43,6 +43,7 @@ def upgrade() -> None: sa.Column('name', sa.String(length=256), nullable=True), sa.Column('address', sa.String(length=256), nullable=True), sa.Column('port', sa.Integer(), nullable=True), + sa.Column('connection_backend', sa.String(), nullable=True), sa.Column('xray_version', sa.String(length=32), nullable=True), sa.Column('status', sa.Enum('healthy', 'unhealthy', 'disabled', name='nodestatus'), nullable=False), sa.Column('last_status_change', sa.DateTime(), nullable=True), diff --git a/app/db/models.py b/app/db/models.py index accd8ce..b2a32f1 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -198,6 +198,7 @@ class Node(Base): ) id = Column(Integer, primary_key=True) name = Column(String(256), unique=True) + connection_backend = Column(String()) address = Column(String(256)) port = Column(Integer) xray_version = Column(String(32)) diff --git a/app/marzneshin.py b/app/marzneshin.py index 0340bed..941754c 100644 --- a/app/marzneshin.py +++ b/app/marzneshin.py @@ -21,9 +21,9 @@ from config import DOCS, DEBUG, UVICORN_SSL_KEYFILE, UVICORN_SSL_CERTFILE, UVICO from config import HOME_PAGE_TEMPLATE from . import __version__, telegram from .dashboard import build_dir, build, base_dir -from .tasks import record_user_usages, review_users, reset_user_data_usage, nodes_health_check, \ - record_realtime_bandwidth, send_notifications, delete_expired_reminders, nodes_startup from .routes import api_router +from .tasks import record_user_usages, review_users, reset_user_data_usage, \ + record_realtime_bandwidth, send_notifications, delete_expired_reminders, nodes_startup logger = logging.getLogger(__name__) dashboard_path = "/dashboard/" @@ -56,7 +56,6 @@ scheduler = AsyncIOScheduler(timezone="UTC") scheduler.add_job(record_user_usages, 'interval', coalesce=True, seconds=30) scheduler.add_job(review_users, 'interval', seconds=600, coalesce=True, max_instances=1) scheduler.add_job(reset_user_data_usage, 'interval', coalesce=True, hours=1) -scheduler.add_job(nodes_health_check, 'interval', seconds=5, coalesce=True, max_instances=1) scheduler.add_job(record_realtime_bandwidth, 'interval', seconds=1, coalesce=True, max_instances=1) if config.WEBHOOK_ADDRESS: @@ -64,6 +63,11 @@ if config.WEBHOOK_ADDRESS: scheduler.add_job(delete_expired_reminders, "interval", hours=2, start_date=dt.utcnow() + td(minutes=1)) +@app.on_event("startup") +async def on_start(): + await nodes_startup() + + @app.on_event("shutdown") async def on_shutdown(): scheduler.shutdown() @@ -103,7 +107,6 @@ async def main(): ) scheduler.start() asyncio.create_task(telegram.start_bot()) - await nodes_startup() cfg = Config( app=app, host=UVICORN_HOST, diff --git a/app/marznode/__init__.py b/app/marznode/__init__.py index 5a3ebe4..aa57ae2 100644 --- a/app/marznode/__init__.py +++ b/app/marznode/__init__.py @@ -3,7 +3,8 @@ from typing import Dict from . import operations from .base import MarzNodeBase -from .grpc import MarzNodeGRPC +from .grpcio import MarzNodeGRPCIO +from .grpclib import MarzNodeGRPCLIB nodes: Dict[int, MarzNodeBase] = {} @@ -11,6 +12,7 @@ nodes: Dict[int, MarzNodeBase] = {} __all__ = [ "nodes", "operations", - "MarzNodeGRPC", + "MarzNodeGRPCIO", + "MarzNodeGRPCLIB", "MarzNodeBase" ] diff --git a/app/marznode/base.py b/app/marznode/base.py index dd38c83..961b85b 100644 --- a/app/marznode/base.py +++ b/app/marznode/base.py @@ -3,22 +3,9 @@ from abc import ABC class MarzNodeBase(ABC): - async def fetch_inbounds(self): - """returns a list of inbounds on the node""" - - async def add_user(self, user, inbounds: list[str] | None = None) -> None: - """adds a user to the node with inbounds specified""" - - async def remove_user(self, user, inbounds: list[str] | None = None) -> None: - """removes a user or some inbounds from a user on the node""" - - async def update_user_inbounds(self, user, - inbound_additions: list[str] | None = None, - inbound_reductions: list[str] | None = None) -> None: - """adds some inbounds and removes some inbounds for a user on the node""" - - async def repopulate_users(self, users) -> None: - """syncs users with the node""" + async def update_user(self, user, + inbounds: list[str] | None = None) -> None: + """updates a user on the node""" async def fetch_users_stats(self): """get user stats from the node""" diff --git a/app/marznode/database.py b/app/marznode/database.py new file mode 100644 index 0000000..5e364e1 --- /dev/null +++ b/app/marznode/database.py @@ -0,0 +1,23 @@ +from app.db import crud, GetDB +from app.models.node import NodeStatus +from app.models.user import UserStatus + + +class MarzNodeDB: + def list_users(self): + with GetDB() as db: + relations = crud.get_node_users(db, self.id, [UserStatus.active, UserStatus.on_hold]) + users = dict() + for rel in relations: + if not users.get(rel[0]): + users[rel[0]] = dict(username=rel[1], id=rel[0], key=rel[2], inbounds=[]) + users[rel[0]]["inbounds"].append(rel[3].tag) + return list(users.values()) + + def store_inbounds(self, inbounds): + with GetDB() as db: + crud.assure_node_inbounds(db, inbounds, self.id) + + def set_status(self, status: NodeStatus, message: str | None = None): + with GetDB() as db: + crud.update_node_status(db, self.id, status, message) diff --git a/app/marznode/grpc.py b/app/marznode/grpc.py deleted file mode 100644 index 0324a8d..0000000 --- a/app/marznode/grpc.py +++ /dev/null @@ -1,113 +0,0 @@ -import atexit -import ssl -import tempfile - -from grpclib import GRPCError -from grpclib.client import Channel -from grpclib.exceptions import StreamTerminatedError -from grpclib.health.v1.health_grpc import HealthStub -from grpclib.health.v1.health_pb2 import HealthCheckResponse, HealthCheckRequest - -from .base import MarzNodeBase -from .marznode_grpc import MarzServiceStub -from .marznode_pb2 import UserUpdate, UsersUpdate, Empty, User, Inbound - - -def string_to_temp_file(content: str): - file = tempfile.NamedTemporaryFile(mode='w+t') - file.write(content) - file.flush() - return file - - -class MarzNodeGRPC(MarzNodeBase): - def __init__(self, address: str, port: int, ssl_key: str, ssl_cert: str, usage_coefficient: int = 1): - self._address = address - self._port = port - - self._keyfile = string_to_temp_file(ssl_key) - self._certfile = string_to_temp_file(ssl_cert) - - ctx = ssl.create_default_context() # cadata=ssl_cert) - ctx.load_cert_chain(self._certfile.name, self._keyfile.name) - ctx.check_hostname = False - ctx.verify_mode = ssl.CERT_NONE - self._channel = Channel(self._address, self._port, ssl=ctx) - self._stub = MarzServiceStub(self._channel) - self._health = HealthStub(self._channel) - self.synced = False - - self.usage_coefficient = usage_coefficient - atexit.register(self._channel.close) - - @staticmethod - def catch_connection_problem(func): - async def catcher(self, *args, **kwargs): - try: - await func(self, *args, **kwargs) - except (OSError, ConnectionError, GRPCError): - self.synced = False - return catcher - - @catch_connection_problem - async def add_user(self, user, inbounds: list[str] | None = None) -> None: - if inbounds is None: - inbounds = [] - - async with self._stub.AddUser.open() as stm: - await stm.send_message(UserUpdate(user=User(id=user.id, username=user.username, key=user.key), - inbound_additions=[Inbound(tag=i) for i in inbounds])) - await stm.recv_message() - - @catch_connection_problem - async def remove_user(self, user, inbounds: list[str] | None = None) -> None: - if inbounds is None: - inbounds = [] - async with self._stub.RemoveUser.open() as stm: - await stm.send_message(UserUpdate(user=User(id=user.id, username=user.username, key=user.key), - inbound_reductions=[Inbound(tag=i) for i in inbounds])) - await stm.recv_message() - - @catch_connection_problem - async def update_user_inbounds(self, user, - inbound_additions: list[str] | None = None, - inbound_reductions: list[str] | None = None) -> None: - if inbound_reductions is None: - inbound_reductions = [] - if inbound_additions is None: - inbound_additions = [] - - async with self._stub.UpdateUserInbounds.open() as stm: - await stm.send_message(UserUpdate(user=User(id=user.id, username=user.username, key=user.key), - inbound_additions=[Inbound(tag=i) for i in inbound_additions], - inbound_reductions=[Inbound(tag=i) for i in inbound_reductions])) - await stm.recv_message() - - @catch_connection_problem - async def repopulate_users(self, user_updates: list[dict]) -> None: - async with self._stub.RepopulateUsers.open() as stream: - updates = [UserUpdate(user=User(id=u["id"], username=u["username"], key=u["key"]), - inbound_additions=[Inbound(tag=t) for t in u["inbounds"]]) for u in user_updates] - await stream.send_message(UsersUpdate(users_updates=updates)) - - async def fetch_users_stats(self): - async with self._stub.FetchUsersStats.open() as stream: - await stream.send_message(Empty()) - stats = await stream.recv_message() - return stats.users_stats - - async def fetch_inbounds(self) -> list: - async with self._stub.FetchInbounds.open() as stm: - await stm.send_message(Empty()) - inbounds = await stm.recv_message() - return inbounds.inbounds - - async def is_alive(self) -> bool: - try: - response = await self._health.Check(HealthCheckRequest()) - except (OSError, ConnectionError, GRPCError, StreamTerminatedError): - pass - else: - if response.status is HealthCheckResponse.SERVING: - return True - return False diff --git a/app/marznode/grpcio.py b/app/marznode/grpcio.py new file mode 100644 index 0000000..3a97c4d --- /dev/null +++ b/app/marznode/grpcio.py @@ -0,0 +1,104 @@ +import asyncio +import atexit +import logging +from _testcapi import INT_MAX + +from grpc import ChannelConnectivity +from grpc.aio import insecure_channel + +from .base import MarzNodeBase +from .database import MarzNodeDB +from .marznode_pb2_grpc import MarzServiceStub +from .marznode_pb2 import UserData, UsersData, Empty, User, Inbound +from ..models.node import NodeStatus + +logger = logging.getLogger(__name__) + +channel_options = [ + ("grpc.keepalive_time_ms", 8000), + ("grpc.keepalive_timeout_ms", 5000), + ("grpc.http2.max_pings_without_data", 0), + ("grpc.keepalive_permit_without_calls", 1), + ("grpc.max_connection_idle_ms", INT_MAX), + ("grpc.client_idle_timeout_ms", INT_MAX), + ("grpc.max_connection_age_ms", INT_MAX) +] + + +class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB): + def __init__(self, node_id: int, address: str, port: int, usage_coefficient: int = 1): + self.id = node_id + self._address = address + self._port = port + + self._channel = insecure_channel(f"{self._address}:{self._port}", channel_options) + self._stub = MarzServiceStub(self._channel) + asyncio.create_task(self._monitor_channel()) + self._streaming_task = None + + self._updates_queue = asyncio.Queue(1) + self.synced = False + self.usage_coefficient = usage_coefficient + atexit.register(self._close_channel) + + def _close_channel(self): + asyncio.run(self._channel.close()) + + async def _monitor_channel(self): + try: + await asyncio.wait_for(self._channel.channel_ready(), timeout=5) + except TimeoutError: + logger.info("timeout for node, id: %i", self.id) + self.set_status(NodeStatus.unhealthy, "timeout") + while state := self._channel.get_state(): + logger.debug("node %i state: %s", self.id, state.value) + if state == ChannelConnectivity.READY: + await self._sync() + self._streaming_task = asyncio.create_task(self._stream_user_updates()) + self.set_status(NodeStatus.healthy) + logger.info("Connected to node %i", self.id) + else: + self.synced = False + self.set_status(NodeStatus.unhealthy) + if self._streaming_task: + self._streaming_task.cancel() + + await self._channel.wait_for_state_change(state) + + async def _stream_user_updates(self): + logger.debug("opened the stream") + stream = self._stub.SyncUsers() + while True: + user_update = await self._updates_queue.get() + logger.debug("got something from queue") + user = user_update["user"] + await stream.write( + UserData(user=User(id=user.id, username=user.username, key=user.key), + inbounds=[Inbound(tag=t) for t in user_update["inbounds"]])) + + async def update_user(self, user, inbounds: set[str] | None = None): + if inbounds is None: + inbounds = set() + + if self.synced: + await self._updates_queue.put({"user": user, "inbounds": inbounds}) + + async def _repopulate_users(self, users_data: list[dict]) -> None: + updates = [UserData(user=User(id=u["id"], username=u["username"], key=u["key"]), + inbounds=[Inbound(tag=t) for t in u["inbounds"]]) for u in users_data] + await self._stub.RepopulateUsers(UsersData(users_data=updates)) + + async def fetch_users_stats(self): + response = await self._stub.FetchUsersStats(Empty()) + return response.users_stats + + async def _fetch_inbounds(self) -> list: + response = await self._stub.FetchInbounds(Empty()) + return response.inbounds + + async def _sync(self): + inbounds = await self._fetch_inbounds() + self.store_inbounds(inbounds) + users = self.list_users() + await self._repopulate_users(users) + self.synced = True diff --git a/app/marznode/grpclib.py b/app/marznode/grpclib.py new file mode 100644 index 0000000..b63e4f2 --- /dev/null +++ b/app/marznode/grpclib.py @@ -0,0 +1,116 @@ +import asyncio +import atexit +import logging +import ssl +import tempfile + +from grpclib import GRPCError +from grpclib.client import Channel +from grpclib.exceptions import StreamTerminatedError + +from .base import MarzNodeBase +from .database import MarzNodeDB +from .marznode_grpc import MarzServiceStub +from .marznode_pb2 import UserData, UsersData, Empty, User, Inbound +from ..models.node import NodeStatus + +logger = logging.getLogger(__name__) + + +def string_to_temp_file(content: str): + file = tempfile.NamedTemporaryFile(mode='w+t') + file.write(content) + file.flush() + return file + + +class MarzNodeGRPCLIB(MarzNodeBase, MarzNodeDB): + def __init__(self, node_id: int, address: str, port: int, ssl_key: str, ssl_cert: str, usage_coefficient: int = 1): + self.id = node_id + self._address = address + self._port = port + + self._key_file = string_to_temp_file(ssl_key) + self._cert_file = string_to_temp_file(ssl_cert) + + ctx = ssl.create_default_context() + ctx.load_cert_chain(self._cert_file.name, self._key_file.name) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + + self._channel = Channel(self._address, self._port, ssl=ctx) + self._stub = MarzServiceStub(self._channel) + asyncio.create_task(self._monitor_channel()) + self._streaming_task = None + + self._updates_queue = asyncio.Queue(1) + self.synced = False + self.usage_coefficient = usage_coefficient + atexit.register(self._channel.close) + + async def _monitor_channel(self): + while state := self._channel.state: + logger.debug("node %i channel state: %s", self.id, state.value) + try: + await asyncio.wait_for(self._channel.__connect__(), timeout=2) + except Exception: + logger.info("timeout for node, id: %i", self.id) + self.set_status(NodeStatus.unhealthy, "timeout") + self.synced = False + if self._streaming_task: + self._streaming_task.cancel() + await asyncio.sleep(10) + continue + else: + if not self.synced: + try: + await self._sync() + self._streaming_task = asyncio.create_task(self._stream_user_updates()) + self.set_status(NodeStatus.healthy) + logger.info("Connected to node %i", self.id) + except: + await asyncio.sleep(10) + continue + await asyncio.sleep(30) + + async def _stream_user_updates(self): + try: + async with self._stub.SyncUsers.open() as stream: + logger.info("opened the stream") + while True: + user_update = await self._updates_queue.get() + logger.info("got something from queue") + user = user_update["user"] + await stream.send_message( + UserData(user=User(id=user.id, username=user.username, key=user.key), + inbounds=[Inbound(tag=t) for t in user_update["inbounds"]])) + except (OSError, ConnectionError, GRPCError, StreamTerminatedError): + logger.info("node %i detached", self.id) + self.synced = False + + async def update_user(self, user, inbounds: set[str] | None = None): + if inbounds is None: + inbounds = set() + + if self.synced: + await self._updates_queue.put({"user": user, "inbounds": inbounds}) + + async def _repopulate_users(self, users_data: list[dict]) -> None: + updates = [UserData(user=User(id=u["id"], username=u["username"], key=u["key"]), + inbounds=[Inbound(tag=t) for t in u["inbounds"]]) for u in users_data] + await self._stub.RepopulateUsers(UsersData(users_data=updates)) + + async def fetch_users_stats(self): + response = await self._stub.FetchUsersStats(Empty()) + return response.users_stats + + async def _fetch_inbounds(self) -> list: + response = await self._stub.FetchInbounds(Empty()) + return response.inbounds + + async def _sync(self): + inbounds = await self._fetch_inbounds() + self.store_inbounds(inbounds) + users = self.list_users() + await self._repopulate_users(users) + self.synced = True diff --git a/app/marznode/marznode.proto b/app/marznode/marznode.proto index 815133b..8a41466 100644 --- a/app/marznode/marznode.proto +++ b/app/marznode/marznode.proto @@ -19,14 +19,13 @@ message User { string key = 3; } -message UserUpdate { +message UserData { User user = 1; - repeated Inbound inbound_additions = 2; - repeated Inbound inbound_reductions = 3; + repeated Inbound inbounds = 2; } -message UsersUpdate { - repeated UserUpdate users_updates = 1; +message UsersData { + repeated UserData users_data = 1; } message UsersStats { @@ -38,10 +37,8 @@ message UsersStats { } service MarzService { - rpc AddUser(UserUpdate) returns (Empty); - rpc RemoveUser(UserUpdate) returns (Empty); - rpc UpdateUserInbounds(UserUpdate) returns (Empty); - rpc RepopulateUsers(UsersUpdate) returns (Empty); + rpc SyncUsers(stream UserData) returns (Empty); + rpc RepopulateUsers(UsersData) returns (Empty); rpc FetchInbounds(Empty) returns (InboundsResponse); rpc FetchUsersStats(Empty) returns (UsersStats); } diff --git a/app/marznode/marznode_grpc.py b/app/marznode/marznode_grpc.py index ddc13c0..75e7675 100644 --- a/app/marznode/marznode_grpc.py +++ b/app/marznode/marznode_grpc.py @@ -15,19 +15,11 @@ import app.marznode.marznode_pb2 class MarzServiceBase(abc.ABC): @abc.abstractmethod - async def AddUser(self, stream: 'grpclib.server.Stream[app.marznode.marznode_pb2.UserUpdate, app.marznode.marznode_pb2.Empty]') -> None: + async def SyncUsers(self, stream: 'grpclib.server.Stream[app.marznode.marznode_pb2.UserData, app.marznode.marznode_pb2.Empty]') -> None: pass @abc.abstractmethod - async def RemoveUser(self, stream: 'grpclib.server.Stream[app.marznode.marznode_pb2.UserUpdate, app.marznode.marznode_pb2.Empty]') -> None: - pass - - @abc.abstractmethod - async def UpdateUserInbounds(self, stream: 'grpclib.server.Stream[app.marznode.marznode_pb2.UserUpdate, app.marznode.marznode_pb2.Empty]') -> None: - pass - - @abc.abstractmethod - async def RepopulateUsers(self, stream: 'grpclib.server.Stream[app.marznode.marznode_pb2.UsersUpdate, app.marznode.marznode_pb2.Empty]') -> None: + async def RepopulateUsers(self, stream: 'grpclib.server.Stream[app.marznode.marznode_pb2.UsersData, app.marznode.marznode_pb2.Empty]') -> None: pass @abc.abstractmethod @@ -40,28 +32,16 @@ class MarzServiceBase(abc.ABC): def __mapping__(self) -> typing.Dict[str, grpclib.const.Handler]: return { - '/marznode.MarzService/AddUser': grpclib.const.Handler( - self.AddUser, - grpclib.const.Cardinality.UNARY_UNARY, - app.marznode.marznode_pb2.UserUpdate, - app.marznode.marznode_pb2.Empty, - ), - '/marznode.MarzService/RemoveUser': grpclib.const.Handler( - self.RemoveUser, - grpclib.const.Cardinality.UNARY_UNARY, - app.marznode.marznode_pb2.UserUpdate, - app.marznode.marznode_pb2.Empty, - ), - '/marznode.MarzService/UpdateUserInbounds': grpclib.const.Handler( - self.UpdateUserInbounds, - grpclib.const.Cardinality.UNARY_UNARY, - app.marznode.marznode_pb2.UserUpdate, + '/marznode.MarzService/SyncUsers': grpclib.const.Handler( + self.SyncUsers, + grpclib.const.Cardinality.STREAM_UNARY, + app.marznode.marznode_pb2.UserData, app.marznode.marznode_pb2.Empty, ), '/marznode.MarzService/RepopulateUsers': grpclib.const.Handler( self.RepopulateUsers, grpclib.const.Cardinality.UNARY_UNARY, - app.marznode.marznode_pb2.UsersUpdate, + app.marznode.marznode_pb2.UsersData, app.marznode.marznode_pb2.Empty, ), '/marznode.MarzService/FetchInbounds': grpclib.const.Handler( @@ -82,28 +62,16 @@ class MarzServiceBase(abc.ABC): class MarzServiceStub: def __init__(self, channel: grpclib.client.Channel) -> None: - self.AddUser = grpclib.client.UnaryUnaryMethod( - channel, - '/marznode.MarzService/AddUser', - app.marznode.marznode_pb2.UserUpdate, - app.marznode.marznode_pb2.Empty, - ) - self.RemoveUser = grpclib.client.UnaryUnaryMethod( - channel, - '/marznode.MarzService/RemoveUser', - app.marznode.marznode_pb2.UserUpdate, - app.marznode.marznode_pb2.Empty, - ) - self.UpdateUserInbounds = grpclib.client.UnaryUnaryMethod( + self.SyncUsers = grpclib.client.StreamUnaryMethod( channel, - '/marznode.MarzService/UpdateUserInbounds', - app.marznode.marznode_pb2.UserUpdate, + '/marznode.MarzService/SyncUsers', + app.marznode.marznode_pb2.UserData, app.marznode.marznode_pb2.Empty, ) self.RepopulateUsers = grpclib.client.UnaryUnaryMethod( channel, '/marznode.MarzService/RepopulateUsers', - app.marznode.marznode_pb2.UsersUpdate, + app.marznode.marznode_pb2.UsersData, app.marznode.marznode_pb2.Empty, ) self.FetchInbounds = grpclib.client.UnaryUnaryMethod( diff --git a/app/marznode/marznode_pb2.py b/app/marznode/marznode_pb2.py index f67c7a0..9b2feb2 100644 --- a/app/marznode/marznode_pb2.py +++ b/app/marznode/marznode_pb2.py @@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default() -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x61pp/marznode/marznode.proto\x12\x08marznode\"\x07\n\x05\x45mpty\"6\n\x07Inbound\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12\x13\n\x06\x63onfig\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\t\n\x07_config\"7\n\x10InboundsResponse\x12#\n\x08inbounds\x18\x01 \x03(\x0b\x32\x11.marznode.Inbound\"1\n\x04User\x12\n\n\x02id\x18\x01 \x01(\r\x12\x10\n\x08username\x18\x02 \x01(\t\x12\x0b\n\x03key\x18\x03 \x01(\t\"\x87\x01\n\nUserUpdate\x12\x1c\n\x04user\x18\x01 \x01(\x0b\x32\x0e.marznode.User\x12,\n\x11inbound_additions\x18\x02 \x03(\x0b\x32\x11.marznode.Inbound\x12-\n\x12inbound_reductions\x18\x03 \x03(\x0b\x32\x11.marznode.Inbound\":\n\x0bUsersUpdate\x12+\n\rusers_updates\x18\x01 \x03(\x0b\x32\x14.marznode.UserUpdate\"j\n\nUsersStats\x12\x33\n\x0busers_stats\x18\x01 \x03(\x0b\x32\x1e.marznode.UsersStats.UserStats\x1a\'\n\tUserStats\x12\x0b\n\x03uid\x18\x01 \x01(\r\x12\r\n\x05usage\x18\x02 \x01(\x04\x32\xe4\x02\n\x0bMarzService\x12\x30\n\x07\x41\x64\x64User\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12\x33\n\nRemoveUser\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12;\n\x12UpdateUserInbounds\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12\x39\n\x0fRepopulateUsers\x12\x15.marznode.UsersUpdate\x1a\x0f.marznode.Empty\x12<\n\rFetchInbounds\x12\x0f.marznode.Empty\x1a\x1a.marznode.InboundsResponse\x12\x38\n\x0f\x46\x65tchUsersStats\x12\x0f.marznode.Empty\x1a\x14.marznode.UsersStatsb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x61pp/marznode/marznode.proto\x12\x08marznode\"\x07\n\x05\x45mpty\"6\n\x07Inbound\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12\x13\n\x06\x63onfig\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\t\n\x07_config\"7\n\x10InboundsResponse\x12#\n\x08inbounds\x18\x01 \x03(\x0b\x32\x11.marznode.Inbound\"1\n\x04User\x12\n\n\x02id\x18\x01 \x01(\r\x12\x10\n\x08username\x18\x02 \x01(\t\x12\x0b\n\x03key\x18\x03 \x01(\t\"M\n\x08UserData\x12\x1c\n\x04user\x18\x01 \x01(\x0b\x32\x0e.marznode.User\x12#\n\x08inbounds\x18\x02 \x03(\x0b\x32\x11.marznode.Inbound\"3\n\tUsersData\x12&\n\nusers_data\x18\x01 \x03(\x0b\x32\x12.marznode.UserData\"j\n\nUsersStats\x12\x33\n\x0busers_stats\x18\x01 \x03(\x0b\x32\x1e.marznode.UsersStats.UserStats\x1a\'\n\tUserStats\x12\x0b\n\x03uid\x18\x01 \x01(\r\x12\r\n\x05usage\x18\x02 \x01(\x04\x32\xf2\x01\n\x0bMarzService\x12\x32\n\tSyncUsers\x12\x12.marznode.UserData\x1a\x0f.marznode.Empty(\x01\x12\x37\n\x0fRepopulateUsers\x12\x13.marznode.UsersData\x1a\x0f.marznode.Empty\x12<\n\rFetchInbounds\x12\x0f.marznode.Empty\x1a\x1a.marznode.InboundsResponse\x12\x38\n\x0f\x46\x65tchUsersStats\x12\x0f.marznode.Empty\x1a\x14.marznode.UsersStatsb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -29,14 +29,14 @@ if _descriptor._USE_C_DESCRIPTORS == False: _globals['_INBOUNDSRESPONSE']._serialized_end=161 _globals['_USER']._serialized_start=163 _globals['_USER']._serialized_end=212 - _globals['_USERUPDATE']._serialized_start=215 - _globals['_USERUPDATE']._serialized_end=350 - _globals['_USERSUPDATE']._serialized_start=352 - _globals['_USERSUPDATE']._serialized_end=410 - _globals['_USERSSTATS']._serialized_start=412 - _globals['_USERSSTATS']._serialized_end=518 - _globals['_USERSSTATS_USERSTATS']._serialized_start=479 - _globals['_USERSSTATS_USERSTATS']._serialized_end=518 - _globals['_MARZSERVICE']._serialized_start=521 - _globals['_MARZSERVICE']._serialized_end=877 + _globals['_USERDATA']._serialized_start=214 + _globals['_USERDATA']._serialized_end=291 + _globals['_USERSDATA']._serialized_start=293 + _globals['_USERSDATA']._serialized_end=344 + _globals['_USERSSTATS']._serialized_start=346 + _globals['_USERSSTATS']._serialized_end=452 + _globals['_USERSSTATS_USERSTATS']._serialized_start=413 + _globals['_USERSSTATS_USERSTATS']._serialized_end=452 + _globals['_MARZSERVICE']._serialized_start=455 + _globals['_MARZSERVICE']._serialized_end=697 # @@protoc_insertion_point(module_scope) diff --git a/app/marznode/marznode_pb2.pyi b/app/marznode/marznode_pb2.pyi new file mode 100644 index 0000000..b882c83 --- /dev/null +++ b/app/marznode/marznode_pb2.pyi @@ -0,0 +1,61 @@ +from google.protobuf.internal import containers as _containers +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union + +DESCRIPTOR: _descriptor.FileDescriptor + +class Empty(_message.Message): + __slots__ = () + def __init__(self) -> None: ... + +class Inbound(_message.Message): + __slots__ = ("tag", "config") + TAG_FIELD_NUMBER: _ClassVar[int] + CONFIG_FIELD_NUMBER: _ClassVar[int] + tag: str + config: str + def __init__(self, tag: _Optional[str] = ..., config: _Optional[str] = ...) -> None: ... + +class InboundsResponse(_message.Message): + __slots__ = ("inbounds",) + INBOUNDS_FIELD_NUMBER: _ClassVar[int] + inbounds: _containers.RepeatedCompositeFieldContainer[Inbound] + def __init__(self, inbounds: _Optional[_Iterable[_Union[Inbound, _Mapping]]] = ...) -> None: ... + +class User(_message.Message): + __slots__ = ("id", "username", "key") + ID_FIELD_NUMBER: _ClassVar[int] + USERNAME_FIELD_NUMBER: _ClassVar[int] + KEY_FIELD_NUMBER: _ClassVar[int] + id: int + username: str + key: str + def __init__(self, id: _Optional[int] = ..., username: _Optional[str] = ..., key: _Optional[str] = ...) -> None: ... + +class UserData(_message.Message): + __slots__ = ("user", "inbounds") + USER_FIELD_NUMBER: _ClassVar[int] + INBOUNDS_FIELD_NUMBER: _ClassVar[int] + user: User + inbounds: _containers.RepeatedCompositeFieldContainer[Inbound] + def __init__(self, user: _Optional[_Union[User, _Mapping]] = ..., inbounds: _Optional[_Iterable[_Union[Inbound, _Mapping]]] = ...) -> None: ... + +class UsersData(_message.Message): + __slots__ = ("users_data",) + USERS_DATA_FIELD_NUMBER: _ClassVar[int] + users_data: _containers.RepeatedCompositeFieldContainer[UserData] + def __init__(self, users_data: _Optional[_Iterable[_Union[UserData, _Mapping]]] = ...) -> None: ... + +class UsersStats(_message.Message): + __slots__ = ("users_stats",) + class UserStats(_message.Message): + __slots__ = ("uid", "usage") + UID_FIELD_NUMBER: _ClassVar[int] + USAGE_FIELD_NUMBER: _ClassVar[int] + uid: int + usage: int + def __init__(self, uid: _Optional[int] = ..., usage: _Optional[int] = ...) -> None: ... + USERS_STATS_FIELD_NUMBER: _ClassVar[int] + users_stats: _containers.RepeatedCompositeFieldContainer[UsersStats.UserStats] + def __init__(self, users_stats: _Optional[_Iterable[_Union[UsersStats.UserStats, _Mapping]]] = ...) -> None: ... diff --git a/app/marznode/marznode_pb2_grpc.py b/app/marznode/marznode_pb2_grpc.py new file mode 100644 index 0000000..74c1b99 --- /dev/null +++ b/app/marznode/marznode_pb2_grpc.py @@ -0,0 +1,165 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +from app.marznode import marznode_pb2 as app_dot_marznode_dot_marznode__pb2 + + +class MarzServiceStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.SyncUsers = channel.stream_unary( + '/marznode.MarzService/SyncUsers', + request_serializer=app_dot_marznode_dot_marznode__pb2.UserData.SerializeToString, + response_deserializer=app_dot_marznode_dot_marznode__pb2.Empty.FromString, + ) + self.RepopulateUsers = channel.unary_unary( + '/marznode.MarzService/RepopulateUsers', + request_serializer=app_dot_marznode_dot_marznode__pb2.UsersData.SerializeToString, + response_deserializer=app_dot_marznode_dot_marznode__pb2.Empty.FromString, + ) + self.FetchInbounds = channel.unary_unary( + '/marznode.MarzService/FetchInbounds', + request_serializer=app_dot_marznode_dot_marznode__pb2.Empty.SerializeToString, + response_deserializer=app_dot_marznode_dot_marznode__pb2.InboundsResponse.FromString, + ) + self.FetchUsersStats = channel.unary_unary( + '/marznode.MarzService/FetchUsersStats', + request_serializer=app_dot_marznode_dot_marznode__pb2.Empty.SerializeToString, + response_deserializer=app_dot_marznode_dot_marznode__pb2.UsersStats.FromString, + ) + + +class MarzServiceServicer(object): + """Missing associated documentation comment in .proto file.""" + + def SyncUsers(self, request_iterator, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def RepopulateUsers(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def FetchInbounds(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def FetchUsersStats(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_MarzServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'SyncUsers': grpc.stream_unary_rpc_method_handler( + servicer.SyncUsers, + request_deserializer=app_dot_marznode_dot_marznode__pb2.UserData.FromString, + response_serializer=app_dot_marznode_dot_marznode__pb2.Empty.SerializeToString, + ), + 'RepopulateUsers': grpc.unary_unary_rpc_method_handler( + servicer.RepopulateUsers, + request_deserializer=app_dot_marznode_dot_marznode__pb2.UsersData.FromString, + response_serializer=app_dot_marznode_dot_marznode__pb2.Empty.SerializeToString, + ), + 'FetchInbounds': grpc.unary_unary_rpc_method_handler( + servicer.FetchInbounds, + request_deserializer=app_dot_marznode_dot_marznode__pb2.Empty.FromString, + response_serializer=app_dot_marznode_dot_marznode__pb2.InboundsResponse.SerializeToString, + ), + 'FetchUsersStats': grpc.unary_unary_rpc_method_handler( + servicer.FetchUsersStats, + request_deserializer=app_dot_marznode_dot_marznode__pb2.Empty.FromString, + response_serializer=app_dot_marznode_dot_marznode__pb2.UsersStats.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'marznode.MarzService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class MarzService(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def SyncUsers(request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.stream_unary(request_iterator, target, '/marznode.MarzService/SyncUsers', + app_dot_marznode_dot_marznode__pb2.UserData.SerializeToString, + app_dot_marznode_dot_marznode__pb2.Empty.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def RepopulateUsers(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/marznode.MarzService/RepopulateUsers', + app_dot_marznode_dot_marznode__pb2.UsersData.SerializeToString, + app_dot_marznode_dot_marznode__pb2.Empty.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def FetchInbounds(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/marznode.MarzService/FetchInbounds', + app_dot_marznode_dot_marznode__pb2.Empty.SerializeToString, + app_dot_marznode_dot_marznode__pb2.InboundsResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def FetchUsersStats(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/marznode.MarzService/FetchUsersStats', + app_dot_marznode_dot_marznode__pb2.Empty.SerializeToString, + app_dot_marznode_dot_marznode__pb2.UsersStats.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/app/marznode/operations.py b/app/marznode/operations.py index 7a4f7d1..667e66b 100644 --- a/app/marznode/operations.py +++ b/app/marznode/operations.py @@ -1,54 +1,35 @@ from collections import defaultdict from typing import TYPE_CHECKING -from .base import MarzNodeBase +from .grpclib import MarzNodeGRPCLIB +from .grpcio import MarzNodeGRPCIO from app import marznode +from ..models.node import NodeConnectionBackend +from ..models.user import UserStatus if TYPE_CHECKING: from app.db import User as DBUser -async def add_user(user: "DBUser"): - node_inbounds = defaultdict(list) - for inb in user.inbounds: - node_inbounds[inb.node_id].append(inb.tag) - - for node_id, tags in node_inbounds.items(): - await marznode.nodes[node_id].add_user(user=user, inbounds=tags) - - -async def remove_user(user: "DBUser"): - nodes_set = set() - - for inb in user.inbounds: - nodes_set.add(inb.node_id) +async def update_user(user: "DBUser", old_inbounds: set | None = None): + """updates a user on all nodes + even though this isn't efficient it is extremely precise""" + if old_inbounds is None: + old_inbounds = set() - for n in nodes_set: - await marznode.nodes[n].remove_user(user) - - -async def update_user_inbounds(user: "DBUser", new_inbounds: set, old_inbounds: set) -> None: - """ - updates user inbounds by finding out inbound additions and reductions - it then calculates all inbound additions and reductions for each node separately, - and then sends requests to affected nodes - :param user: the user - :param new_inbounds: a set of all the new inbounds - :param old_inbounds: a set of all the old inbounds - :return: nothing - """ - inbound_additions = new_inbounds - old_inbounds - inbound_reductions = old_inbounds - new_inbounds - - node_inbounds = defaultdict(lambda: tuple((list(), list()))) + node_inbounds = defaultdict(list) + if user.status in (UserStatus.on_hold, UserStatus.active): + for inb in user.inbounds: + node_inbounds[inb.node_id].append(inb.tag) + else: + for inb in user.inbounds: + node_inbounds[inb.node_id] - for inbound_addition in inbound_additions: - node_inbounds[inbound_addition[0]][0].append(inbound_addition[2]) - for inbound_reduction in inbound_reductions: - node_inbounds[inbound_reduction[0]][1].append(inbound_reduction[2]) + for inb in old_inbounds: + node_inbounds[inb[0]] - for node, inbound_change in node_inbounds.items(): - await marznode.nodes[node].update_user_inbounds(user, inbound_change[0], inbound_change[1]) + for node_id, tags in node_inbounds.items(): + await marznode.nodes[node_id].update_user(user=user, inbounds=tags) async def remove_node(node_id: int): @@ -56,15 +37,18 @@ async def remove_node(node_id: int): del marznode.nodes[node_id] -async def add_node(node_id: int, node: MarzNodeBase): - await remove_node(node_id) - marznode.nodes[node_id] = node +async def add_node(db_node, certificate): + await remove_node(db_node.id) + if db_node.connection_backend == NodeConnectionBackend.grpcio: + node = MarzNodeGRPCIO(db_node.id, db_node.address, db_node.port) + else: + node = MarzNodeGRPCLIB(db_node.id, db_node.address, db_node.port, certificate.key, + certificate.certificate) + marznode.nodes[db_node.id] = node __all__ = [ - "add_user", - "remove_user", - "update_user_inbounds", + "update_user", "add_node", "remove_node" ] diff --git a/app/models/node.py b/app/models/node.py index c7ad8be..e232b43 100644 --- a/app/models/node.py +++ b/app/models/node.py @@ -11,6 +11,11 @@ class NodeStatus(str, Enum): none = "none" +class NodeConnectionBackend(str, Enum): + grpcio = "grpcio" + grpclib = "grpclib" + + class NodeSettings(BaseModel): min_node_version: str = "v0.2.0" certificate: str @@ -21,6 +26,7 @@ class NodeBase(BaseModel): name: str address: str port: int = 62050 + connection_backend: NodeConnectionBackend = Field(default=NodeConnectionBackend.grpclib) usage_coefficient: float = Field(ge=0, default=1.0) model_config = ConfigDict(from_attributes=True) @@ -47,6 +53,7 @@ class NodeModify(Node): name: Optional[str] = Field(None) address: Optional[str] = Field(None) port: Optional[int] = Field(None) + connection_backend: Optional[NodeConnectionBackend] = Field(None) status: Optional[NodeStatus] = Field(None) usage_coefficient: Optional[float] = Field(None) model_config = ConfigDict(json_schema_extra={ diff --git a/app/models/user.py b/app/models/user.py index ab965ff..9e49c0b 100644 --- a/app/models/user.py +++ b/app/models/user.py @@ -7,7 +7,6 @@ from typing import List, Optional, Union, Annotated from pydantic import field_validator, ConfigDict, BaseModel, Field, validator, computed_field, ValidationInfo, constr, \ StringConstraints -from app.models.proxy import InboundBase from config import XRAY_SUBSCRIPTION_URL_PREFIX USERNAME_REGEXP = re.compile(r"^(?=\w{3,32}\b)[a-zA-Z0-9-_@.]+(?:_[a-zA-Z0-9-_@.]+)*$") @@ -96,6 +95,7 @@ class UserBase(BaseModel): from app.models.service import ServiceBase +from app.models.proxy import InboundBase class User(UserBase): diff --git a/app/routes/node.py b/app/routes/node.py index 44974fa..fce0434 100644 --- a/app/routes/node.py +++ b/app/routes/node.py @@ -9,7 +9,7 @@ from fastapi import HTTPException, WebSocket from app import marznode from app.db import crud, get_tls_certificate from app.dependencies import DBDep, SudoAdminDep, sudo_admin, EndDateDep, StartDateDep -from app.marznode import MarzNodeGRPC +from app.marznode import MarzNodeGRPCIO from app.models.admin import Admin from app.models.node import (NodeCreate, NodeModify, NodeResponse, NodeSettings, NodeStatus, NodesUsageResponse) @@ -39,10 +39,7 @@ async def add_node(new_node: NodeCreate, raise HTTPException(status_code=409, detail=f"Node \"{new_node.name}\" already exists") certificate = get_tls_certificate(db) - node = MarzNodeGRPC(db_node.address, db_node.port, - ssl_key=certificate.key, ssl_cert=certificate.certificate) - - await marznode.operations.add_node(db_node.id, node) + await marznode.operations.add_node(db_node, certificate) logger.info("New node `%s` added", db_node.name) return db_node @@ -115,9 +112,7 @@ async def modify_node(node_id: int, await marznode.operations.remove_node(db_node.id) if db_node.status != NodeStatus.disabled: certificate = get_tls_certificate(db) - node = MarzNodeGRPC(db_node.address, db_node.port, - ssl_key=certificate.key, ssl_cert=certificate.certificate) - await marznode.operations.add_node(db_node.id, node) + await marznode.operations.add_node(db_node, certificate) logger.info("Node `%s` modified", db_node.name) return db_node diff --git a/app/routes/user.py b/app/routes/user.py index 80c2061..3ae7004 100644 --- a/app/routes/user.py +++ b/app/routes/user.py @@ -40,7 +40,7 @@ async def add_user(new_user: UserCreate, raise HTTPException(status_code=409, detail="User already exists") user = UserResponse.model_validate(db_user) - await marznode.operations.add_user(user=db_user) + await marznode.operations.update_user(user=db_user) asyncio.create_task( report.user_created( user=user, @@ -119,7 +119,8 @@ async def remove_user(db_user: UserDep, """ Remove a user """ - await marznode.operations.remove_user(db_user) + db_user.inbounds = [] + await marznode.operations.update_user(db_user) crud.remove_user(db, db_user) db.flush() diff --git a/app/tasks/__init__.py b/app/tasks/__init__.py index 1396639..85bfd27 100644 --- a/app/tasks/__init__.py +++ b/app/tasks/__init__.py @@ -1,5 +1,5 @@ import logging -from .nodes import nodes_startup, nodes_health_check +from .nodes import nodes_startup from .record_usages import record_user_usages from .reset_user_data_usage import reset_user_data_usage from .review_users import review_users @@ -8,7 +8,6 @@ from .store_bandwidth import record_realtime_bandwidth __all__ = [ "nodes_startup", - "nodes_health_check", "record_user_usages", "reset_user_data_usage", "review_users", diff --git a/app/tasks/nodes.py b/app/tasks/nodes.py index a16b64d..cf75c68 100644 --- a/app/tasks/nodes.py +++ b/app/tasks/nodes.py @@ -1,38 +1,5 @@ -from collections import defaultdict - -from grpclib import GRPCError - from app import marznode from app.db import GetDB, crud, get_tls_certificate -from app.marznode import MarzNodeGRPC -from app.models.node import NodeStatus -from app.models.proxy import InboundBase -from app.models.user import UserStatus - - -async def nodes_health_check(): - with GetDB() as db: - for node_id, node in marznode.nodes.items(): - db_node = crud.get_node_by_id(db, node_id) - if await node.is_alive(): - crud.update_node_status(db, db_node, NodeStatus.healthy) - if not node.synced: - await sync_node(db, node, node_id) - else: - crud.update_node_status(db, db_node, NodeStatus.unhealthy) - - -async def sync_node(db, node, node_id): - inbounds = await node.fetch_inbounds() - crud.assure_node_inbounds(db, inbounds, node_id) - relations = crud.get_node_users(db, node_id, [UserStatus.active, UserStatus.on_hold]) - users = dict() - for rel in relations: - if not users.get(rel[0]): - users[rel[0]] = dict(username=rel[1], id=rel[0], key=rel[2], inbounds=[]) - users[rel[0]]["inbounds"].append(rel[3].tag) - await node.repopulate_users(list(users.values())) - node.synced = True async def nodes_startup(): @@ -40,13 +7,4 @@ async def nodes_startup(): certificate = get_tls_certificate(db) db_nodes = crud.get_nodes(db=db, enabled=True) for db_node in db_nodes: - node = MarzNodeGRPC(db_node.address, db_node.port, - ssl_key=certificate.key, ssl_cert=certificate.certificate) - if not await node.is_alive(): - node.synced = False - await marznode.operations.add_node(db_node.id, node) - continue - await sync_node(db, node, db_node.id) - await marznode.operations.add_node(db_node.id, node) - # crud.update_node_status(db, db_node, NodeStatus.connecting) - + await marznode.operations.add_node(db_node, certificate) diff --git a/app/tasks/record_usages.py b/app/tasks/record_usages.py index d15d67d..599a5cf 100644 --- a/app/tasks/record_usages.py +++ b/app/tasks/record_usages.py @@ -82,7 +82,7 @@ async def get_users_stats(node_id: int, node: MarzNodeBase) -> tuple[int, list[d if stat.usage: params.append({"uid": stat.uid, "value": stat.usage}) return node_id, params - except (GRPCError, ConnectionError, OSError): + except: return node_id, []