106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
import ipaddress
|
|
import json
|
|
import math
|
|
import threading
|
|
from enum import Enum
|
|
|
|
from colony.core.dispatcher import parse_command
|
|
from colony.core.network.relay import GATEWAY_ROLE
|
|
from colony.core.roles.role import Role, Command
|
|
from colony.settings import setting
|
|
from colony.transport.whisperer import ConnectFailed
|
|
|
|
|
|
def failure(code):
|
|
return {"ok": False, "code": code}
|
|
|
|
|
|
class Gateway(Role):
|
|
class COMMANDS(Enum):
|
|
FORWARD = 'Gateway.Forward'
|
|
|
|
def __init__(self, context):
|
|
commands = {
|
|
self.COMMANDS.FORWARD.value: Command(self.forward),
|
|
}
|
|
|
|
super().__init__(commands, context)
|
|
|
|
self.allow = set(setting("gateway.allow_commands"))
|
|
self.max_payload = setting("gateway.max_payload_bytes")
|
|
self.forward_timeout = setting("gateway.forward_timeout")
|
|
self.connect_timeout = setting("gateway.connect_timeout")
|
|
self._slots = threading.BoundedSemaphore(setting("gateway.max_concurrent"))
|
|
self.forwarded = 0
|
|
|
|
def forward(self, request):
|
|
result = self._forward(request)
|
|
if result["ok"]:
|
|
self.forwarded += 1
|
|
self.logger.info(f"Forward finished: {'ok' if result['ok'] else result['code']}")
|
|
return result
|
|
|
|
def _read(self, request):
|
|
if not isinstance(request, dict):
|
|
raise ValueError("Request must be an object")
|
|
|
|
target = request.get("target")
|
|
if not isinstance(target, list) or len(target) != 2:
|
|
raise ValueError("Target must be [ip, port]")
|
|
ip, port = target
|
|
if not isinstance(ip, str):
|
|
raise ValueError("Bad target ip")
|
|
ip = str(ipaddress.ip_address(ip))
|
|
if not isinstance(port, int) or isinstance(port, bool) or not 1 <= port <= 65535:
|
|
raise ValueError("Bad target port")
|
|
|
|
command = request.get("command")
|
|
parse_command({"command": command})
|
|
|
|
timeout = request.get("timeout", self.forward_timeout)
|
|
if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout) or timeout <= 0:
|
|
raise ValueError("Bad timeout")
|
|
return (ip, port), command, min(timeout, self.forward_timeout)
|
|
|
|
def _allowed(self, command):
|
|
name = f"{command['role']}.{command['name']}"
|
|
return command["role"] != GATEWAY_ROLE and ("*" in self.allow or name in self.allow)
|
|
|
|
def _is_member(self, target):
|
|
membership = self.context.membership()
|
|
return any(record.key == target for record in membership.alive_records() if record.key != membership.key)
|
|
|
|
def _forward(self, request):
|
|
try:
|
|
target, command, timeout = self._read(request)
|
|
except (ValueError, TypeError):
|
|
return failure("bad_request")
|
|
|
|
if not self._allowed(command):
|
|
return failure("denied")
|
|
if not self._is_member(target):
|
|
return failure("unknown_target")
|
|
|
|
payload = {"command": command}
|
|
if len(json.dumps(payload)) > self.max_payload:
|
|
return failure("too_large")
|
|
if not self._slots.acquire(blocking=False):
|
|
return failure("busy")
|
|
|
|
try:
|
|
response = self.context.send_direct(
|
|
target[0], target[1], payload, timeout=timeout, connect_timeout=min(self.connect_timeout, timeout)
|
|
)
|
|
except ConnectFailed:
|
|
return failure("unreachable")
|
|
except TimeoutError:
|
|
return failure("timeout")
|
|
except Exception:
|
|
return failure("error")
|
|
finally:
|
|
self._slots.release()
|
|
|
|
if len(json.dumps(response)) > self.max_payload:
|
|
return failure("too_large")
|
|
return {"ok": True, "response": response}
|