349 lines
12 KiB
Python
349 lines
12 KiB
Python
import socket
|
|
import json
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import logging
|
|
|
|
# 配置日志
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ServiceProvider:
|
|
def __init__(self, coordinator_addr, services):
|
|
"""
|
|
初始化服务提供者
|
|
:param coordinator_addr: 协调服务器地址 (IP, port)
|
|
:param services: 提供的服务列表 {服务名: 端口号}
|
|
"""
|
|
self.provider_id = f"provider-{uuid.uuid4().hex[:8]}"
|
|
self.coordinator_addr = coordinator_addr
|
|
self.services = services
|
|
|
|
# 创建UDP套接字用于协调通信
|
|
self.udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
self.udp_sock.bind(('0.0.0.0', 0))
|
|
self.udp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536)
|
|
|
|
# 创建线程池用于处理UDP消息
|
|
self.thread_pool = ThreadPoolExecutor(max_workers=10)
|
|
|
|
self.clients = {}
|
|
# 存储活动连接
|
|
self.active_connections = {}
|
|
self.running = True
|
|
|
|
# 心跳线程
|
|
self.heartbeat_thread = threading.Thread(target=self.send_heartbeat, daemon=True)
|
|
|
|
def send_heartbeat(self):
|
|
"""
|
|
发送心跳包到协调服务器
|
|
"""
|
|
while self.running:
|
|
try:
|
|
message = {
|
|
'action': 'heartbeat',
|
|
'provider_id': self.provider_id
|
|
}
|
|
self.udp_sock.sendto(json.dumps(message).encode(), self.coordinator_addr)
|
|
logger.debug(f"发送心跳包给协调服务器 {self.coordinator_addr}")
|
|
time.sleep(20) # 每20秒发送一次心跳包
|
|
except Exception as e:
|
|
logger.error(f"发送心跳包失败: {str(e)}")
|
|
|
|
def register_service(self):
|
|
"""
|
|
向协调服务器注册服务
|
|
:return: 注册是否成功
|
|
"""
|
|
message = {
|
|
'action': 'register',
|
|
'services': self.services,
|
|
'provider_id': self.provider_id
|
|
}
|
|
logger.info(f"向协调服务器 {self.coordinator_addr} 注册服务 '{self.services}'")
|
|
self.udp_sock.sendto(json.dumps(message).encode(), self.coordinator_addr)
|
|
|
|
# 等待响应
|
|
try:
|
|
data, _ = self.udp_sock.recvfrom(4096)
|
|
response = json.loads(data.decode())
|
|
if response['status'] == 'success':
|
|
logger.info(f"服务 '{self.services}' 注册成功")
|
|
return True
|
|
else:
|
|
logger.error(f"注册失败: {response['message']}")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"注册服务时发生错误: {str(e)}")
|
|
return False
|
|
|
|
def handle_punch(self, message, addr):
|
|
"""
|
|
处理打洞请求
|
|
:param message: 打洞请求消息
|
|
:param addr: 客户端地址
|
|
"""
|
|
self.udp_sock.sendto(json.dumps(
|
|
{'action': 'punch_response',
|
|
'client_id': message['client_id'],
|
|
'provider_id': self.provider_id
|
|
}).encode(), addr)
|
|
logger.debug(f"收到来自 {addr} 的打洞请求，已响应")
|
|
|
|
def handle_punch_response(self, _, addr):
|
|
"""
|
|
处理打洞响应
|
|
:param addr: 客户端地址
|
|
"""
|
|
logger.debug(f"收到来自 {addr} 的打洞响应")
|
|
|
|
def handle_connect_request(self, message, addr):
|
|
"""
|
|
处理连接请求
|
|
:param message: 连接请求消息
|
|
:param addr: 客户端地址
|
|
"""
|
|
conn_id = message['conn_id']
|
|
client_id = message['client_id']
|
|
service_name = message['service_name']
|
|
logger.debug(f"收到来自 {addr} 的连接请求")
|
|
threading.Thread(
|
|
target=self.handle_connection,
|
|
args=(conn_id, addr, client_id, service_name),
|
|
daemon=True
|
|
).start()
|
|
|
|
def handle_punch_request(self, message, _):
|
|
"""
|
|
处理打洞请求
|
|
:param message: 打洞请求消息
|
|
"""
|
|
for i in range(10):
|
|
try:
|
|
self.udp_sock.sendto(json.dumps({
|
|
'action': 'punch',
|
|
'client_id': message['client_id'],
|
|
'provider_id': self.provider_id,
|
|
}).encode(), tuple(message['client_addr']))
|
|
time.sleep(0.5)
|
|
except Exception as e:
|
|
logger.error(f"打洞失败: {str(e)}")
|
|
time.sleep(1)
|
|
|
|
def handle_data(self, message, addr):
|
|
"""
|
|
处理数据消息
|
|
:param message: 数据消息
|
|
:param addr: 客户端地址
|
|
"""
|
|
conn_id = message['conn_id']
|
|
data = bytes.fromhex(message['data'])
|
|
if conn_id in self.active_connections:
|
|
# 转发数据到本地服务
|
|
logger.debug(f"收到来自 {addr} 的数据，转发到本地服务")
|
|
self.active_connections[conn_id]['local_sock'].sendall(data)
|
|
else:
|
|
self.udp_sock.sendto(json.dumps({
|
|
'action': 'stop_conn',
|
|
'conn_id': conn_id
|
|
}).encode(), addr)
|
|
logger.debug(f"收到来自 {addr} 的数据，但未找到对应的连接")
|
|
|
|
def handle_stop_conn(self, message, _):
|
|
"""
|
|
处理停止连接请求
|
|
:param message: 停止连接请求消息
|
|
"""
|
|
conn_id = message['conn_id']
|
|
if conn_id in self.active_connections:
|
|
self.active_connections[conn_id]['local_sock'].close()
|
|
self.active_connections.pop(conn_id, None)
|
|
|
|
def handle_stop_client(self, message, addr):
|
|
"""
|
|
处理停止客户端请求
|
|
:param message: 停止客户端请求消息
|
|
:param addr: 客户端地址
|
|
"""
|
|
client_id = message['client_id']
|
|
for conn_id, conn_info in self.active_connections.items():
|
|
if conn_info['client_id'] == client_id:
|
|
conn_info['local_sock'].close()
|
|
self.active_connections.pop(conn_id, None)
|
|
|
|
def handle_stop_provider(self, message, _):
|
|
"""
|
|
处理停止服务提供者请求
|
|
:param message: 停止服务提供者请求消息
|
|
"""
|
|
logger.info("收到停止服务提供者请求，正在关闭所有连接...")
|
|
for conn_id, conn_info in self.active_connections.items():
|
|
conn_info['local_sock'].close()
|
|
self.udp_sock.sendto(json.dumps({
|
|
'action': 'stop_conn',
|
|
'conn_id': conn_id
|
|
}).encode(), conn_info['client_addr'])
|
|
self.active_connections.clear()
|
|
self.running = False
|
|
self.udp_sock.close()
|
|
self.thread_pool.shutdown(wait=True)
|
|
logger.info("服务提供者已停止")
|
|
|
|
def handle_connection(self, conn_id, client_addr, client_id, service_name):
|
|
"""
|
|
处理来自客户端的连接
|
|
:param conn_id: 连接ID
|
|
:param client_addr: 客户端地址
|
|
:param client_id: 客户端ID
|
|
:param service_name: 服务名称
|
|
"""
|
|
internal_port = self.services[service_name]
|
|
try:
|
|
# 接受本地服务连接
|
|
local_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
local_sock.connect(('127.0.0.1', internal_port))
|
|
if not local_sock:
|
|
logger.error("无法连接到本地服务")
|
|
return
|
|
|
|
# 创建与客户端的UDP隧道
|
|
logger.debug(f"建立连接 {conn_id} : {client_addr} -> {('127.0.0.1', internal_port)}")
|
|
|
|
# 存储连接
|
|
self.active_connections[conn_id] = {
|
|
'local_sock': local_sock,
|
|
'client_addr': client_addr,
|
|
'client_id': client_id
|
|
}
|
|
|
|
# 通知客户端连接就绪
|
|
self.udp_sock.sendto(json.dumps({
|
|
'action': 'connected',
|
|
'client_id': conn_id
|
|
}).encode(), client_addr)
|
|
|
|
# 启动数据转发
|
|
self.forward_data(conn_id, local_sock, client_addr)
|
|
except Exception as e:
|
|
logger.error(f"连接失败: {str(e)}")
|
|
self.udp_sock.sendto(json.dumps({
|
|
'action': 'connect_failed',
|
|
'client_id': conn_id,
|
|
'message': str(e)
|
|
}).encode(), client_addr)
|
|
|
|
def forward_data(self, conn_id, local_sock, client_addr):
|
|
"""
|
|
转发TCP数据到UDP隧道
|
|
:param conn_id: 连接ID
|
|
:param local_sock: 本地服务套接字
|
|
:param client_addr: 客户端地址
|
|
"""
|
|
try:
|
|
while True:
|
|
# 从本地服务读取数据
|
|
data = local_sock.recv(4096)
|
|
if not data:
|
|
break
|
|
|
|
# 通过UDP发送给客户端
|
|
self.udp_sock.sendto(json.dumps({
|
|
'action': 'data',
|
|
'conn_id': conn_id,
|
|
'data': data.hex() # 十六进制编码二进制数据
|
|
}).encode(), client_addr)
|
|
logger.debug(f"转发数据给客户端{client_addr}")
|
|
except Exception as e:
|
|
logger.error(f"转发数据失败: {str(e)}")
|
|
finally:
|
|
local_sock.close()
|
|
if conn_id in self.active_connections:
|
|
del self.active_connections[conn_id]
|
|
logger.debug(f"连接 {conn_id} 已关闭")
|
|
|
|
def udp_listener(self):
|
|
"""
|
|
监听UDP消息并处理
|
|
"""
|
|
data = None
|
|
while self.running:
|
|
try:
|
|
data, addr = self.udp_sock.recvfrom(4096)
|
|
logger.debug(f"收到来自 {addr} 的消息: {data}")
|
|
message = json.loads(data.decode())
|
|
|
|
# 使用字典映射处理不同消息类型
|
|
action_handlers = {
|
|
'punch': self.handle_punch,
|
|
'punch_check': self.handle_punch,
|
|
'punch_response': self.handle_punch_response,
|
|
'connect': self.handle_connect_request,
|
|
'punch_request': self.handle_punch_request,
|
|
'data': self.handle_data,
|
|
'stop_conn': self.handle_stop_conn,
|
|
'stop_client': self.handle_stop_client,
|
|
'stop_provider': self.handle_stop_provider,
|
|
}
|
|
|
|
# 提交任务到线程池
|
|
if message.get('action') in action_handlers:
|
|
self.thread_pool.submit(action_handlers[message['action']], message, addr)
|
|
elif message.get('status') == 'error':
|
|
logger.error(f"来自 {addr} 的错误消息: {message}")
|
|
elif message.get('status') == 'success':
|
|
logger.debug(f"来自 {addr} 的成功消息: {message}")
|
|
else:
|
|
logger.warning(f"收到未知消息: {message}")
|
|
except Exception as e:
|
|
logger.error(f"处理UDP消息时发生错误: {str(e)}")
|
|
if data:
|
|
logger.error(f"无法处理消息: {data}")
|
|
|
|
def run(self):
|
|
"""
|
|
运行服务提供端
|
|
"""
|
|
# 注册服务
|
|
if not self.register_service():
|
|
logger.error("服务注册失败，退出程序")
|
|
return
|
|
|
|
# 启动UDP监听线程
|
|
threading.Thread(target=self.udp_listener, daemon=True).start()
|
|
|
|
# 启动心跳线程
|
|
self.heartbeat_thread.start()
|
|
|
|
# 保持主线程运行
|
|
try:
|
|
while self.running:
|
|
time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
self.running = False
|
|
self.udp_sock.sendto(json.dumps({'action': 'stop_provider'}).encode(), self.coordinator_addr)
|
|
for conn_id, conn_info in self.active_connections.items():
|
|
self.udp_sock.sendto(json.dumps({
|
|
'action': 'stop_conn',
|
|
'conn_id': conn_id
|
|
}).encode(), conn_info['client_addr'])
|
|
self.udp_sock.close()
|
|
logger.info("服务提供端已停止")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# 配置信息
|
|
COORDINATOR_ADDR = ('www.awin-x.top', 5000) # 替换为公网服务器IP
|
|
SERVICES = {
|
|
'terraria-jk-2cxht5': 5001,
|
|
'minecraft-jk-ytsvb54u6': 5002,
|
|
'alist-jk-5shf43h6fdg': 5244,
|
|
'ssh-jk-54htrsd324n6': 22
|
|
}
|
|
|
|
provider = ServiceProvider(COORDINATOR_ADDR, SERVICES)
|
|
provider.run()
|