diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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, []