class RateLimiterTestCase(unittest.TestCase):
def setUp(self):
super().setUp()
self.rate_limiter = RateLimiter(reactor=self.clock)
def test_limiter(self):
key = 'test'
self.rate_limiter.set_limit(key, 2, 2)
# Hits limit
self.assertTrue(self.rate_limiter.add_hit(key))
self.assertTrue(self.rate_limiter.add_hit(key))
self.assertFalse(self.rate_limiter.add_hit(key))
# Advance 3 seconds to release limit
self.clock.advance(3)
# Add hits until limit again
self.assertTrue(self.rate_limiter.add_hit(key))
self.assertTrue(self.rate_limiter.add_hit(key))
self.assertFalse(self.rate_limiter.add_hit(key))
# Reset hits
self.rate_limiter.reset(key)
# Limit is free again
self.assertTrue(self.rate_limiter.add_hit(key))
# Get limit
self.assertEqual(self.rate_limiter.get_limit(key).max_hits, 2)
# Unset limit
self.rate_limiter.unset_limit(key)
self.assertIsNone(self.rate_limiter.get_limit(key))
class HathorAdminWebsocketFactory(WebSocketServerFactory):
""" Factory of the admin websocket protocol so we can subscribe to events and
send messages in the Admin page to clients when the events are published
"""
protocol = HathorAdminWebsocketProtocol
def buildProtocol(self, addr):
return self.protocol(self)
def __init__(self,
metrics: Optional[Metrics] = None,
wallet_index: Optional[WalletIndex] = None):
"""
:param metrics: If not given, a new one is created.
:type metrics: :py:class:`hathor.metrics.Metrics`
"""
# Opened websocket connections so I can broadcast messages later
self.connections: Set[HathorAdminWebsocketProtocol] = set()
# Websocket connection for each address
self.address_connections: DefaultDict[
str, Set[HathorAdminWebsocketProtocol]] = defaultdict(set)
super().__init__()
# Limit the send message rate for specific type of data
self.rate_limiter = RateLimiter(reactor=reactor)
# Stores the buffer of messages that exceeded the rate limit and will be sent
self.buffer_deques: Dict[str, Deque[Dict[str, Any]]] = {}
self.metrics = metrics
self.wallet_index = wallet_index
self.is_running = False
def start(self):
self.is_running = True
# Start limiter
self._setup_rate_limit()
# Start metric sender
self._schedule_and_send_metric()
def stop(self):
self.is_running = False
def _setup_rate_limit(self):
""" Set the limit of the RateLimiter and start the buffer deques with BUFFER_SIZE
"""
for control_type, config in CONTROLLED_TYPES.items():
self.rate_limiter.set_limit(control_type, config['max_hits'],
config['hits_window_seconds'])
self.buffer_deques[control_type] = deque(
maxlen=config['buffer_size'])
def _schedule_and_send_metric(self):
""" Send dashboard metric to websocket and schedule next message
"""
data = {
'transactions': self.metrics.transactions,
'blocks': self.metrics.blocks,
'best_block_height': self.metrics.best_block_height,
'hash_rate': self.metrics.hash_rate,
'block_hash_rate': self.metrics.block_hash_rate,
'tx_hash_rate': self.metrics.tx_hash_rate,
'network_hash_rate':
self.metrics.tx_hash_rate + self.metrics.block_hash_rate,
'peers': self.metrics.peers,
'type': 'dashboard:metrics',
'time': reactor.seconds(),
}
self.broadcast_message(data)
if self.is_running:
# Schedule next message
reactor.callLater(1, self._schedule_and_send_metric)
def subscribe(self, pubsub):
""" Subscribe to defined events for the pubsub received
"""
events = [
HathorEvents.NETWORK_NEW_TX_ACCEPTED,
HathorEvents.WALLET_OUTPUT_RECEIVED,
HathorEvents.WALLET_INPUT_SPENT,
HathorEvents.WALLET_BALANCE_UPDATED,
HathorEvents.WALLET_KEYS_GENERATED,
HathorEvents.WALLET_GAP_LIMIT,
HathorEvents.WALLET_HISTORY_UPDATED,
HathorEvents.WALLET_ADDRESS_HISTORY,
HathorEvents.WALLET_ELEMENT_WINNER,
HathorEvents.WALLET_ELEMENT_VOIDED,
]
for event in events:
pubsub.subscribe(event, self.handle_publish)
def handle_publish(self, key, args):
""" This method is called when pubsub publishes an event that we subscribed
Then we broadcast the data to all connected clients
"""
data = self.serialize_message_data(key, args)
data['type'] = key.value
self.send_or_enqueue(data)
def serialize_message_data(self, event, args):
""" Receives the event and the args from the pubsub
and serializes the data so it can be passed in the websocket
"""
# Ready events don't need extra serialization
ready_events = [
HathorEvents.WALLET_KEYS_GENERATED,
HathorEvents.WALLET_GAP_LIMIT,
HathorEvents.WALLET_HISTORY_UPDATED,
HathorEvents.WALLET_ADDRESS_HISTORY,
]
data = args.__dict__
if event in ready_events:
return data
elif event == HathorEvents.WALLET_OUTPUT_RECEIVED:
data['output'] = data['output'].to_dict()
return data
elif event == HathorEvents.WALLET_INPUT_SPENT:
data['output_spent'] = data['output_spent'].to_dict()
return data
elif event == HathorEvents.NETWORK_NEW_TX_ACCEPTED:
tx = data['tx']
data = tx.to_json_extended()
data['is_block'] = tx.is_block
return data
elif event == HathorEvents.WALLET_BALANCE_UPDATED:
data['balance'] = data['balance'][
settings.HATHOR_TOKEN_UID]._asdict()
return data
else:
raise ValueError(
'Should never have entered here! We dont know this event')
def execute_send(self, data: Dict[str, Any],
connections: Set[HathorAdminWebsocketProtocol]) -> None:
""" Send data in ws message for the connections
"""
payload = json.dumps(data).encode('utf-8')
for c in connections:
c.sendMessage(payload, False)
def broadcast_message(self, data: Dict[str, Any]) -> None:
""" Broadcast the update message to the connections
"""
self.execute_send(data, self.connections)
def send_message(self, data: Dict[str, Any]) -> None:
""" Check if should broadcast the message to all connections or send directly to some connections only
"""
if data['type'] in ADDRESS_EVENTS:
# This ws message will only be sent if the address was subscribed
if data['address'] in self.address_connections:
self.execute_send(data,
self.address_connections[data['address']])
else:
self.broadcast_message(data)
def send_or_enqueue(self, data):
""" Try to broadcast the message, or enqueue it when rate limit is exceeded and we've been throttled.
Enqueued messages are automatically sent after a while if they are not discarded first.
A message is discarded when new messages arrive and the queue buffer is full.
Rate limits change according to the message type, which is obtained from data['type'].
:param data: message to be sent
:type data: Dict[string, X] -> X can be different types, depending on the type of message
"""
if data['type'] in CONTROLLED_TYPES:
# This type is controlled, so I need to check the deque
if len(self.buffer_deques[data['type']]
) or not self.rate_limiter.add_hit(data['type']):
# If I am already with a buffer or if I hit the limit now, I enqueue for later
self.enqueue_for_later(data)
else:
data['throttled'] = False
self.send_message(data)
else:
self.send_message(data)
def enqueue_for_later(self, data):
""" Add this date to the correct deque to be processed later
If this deque is not programed to be called later yet, we call it
:param data: message to be sent
:type data: Dict[string, X] -> X can be different types, depending on the type of message
"""
# Add data to deque
# We always add the new messages in the end
# Adding parameter deque=True, so the admin can know this message was delayed
data['throttled'] = True
self.buffer_deques[data['type']].append(data)
if len(self.buffer_deques[data['type']]) == 1:
# If it's the first time we hit the limit (only one message in deque), we schedule process_deque
reactor.callLater(CONTROLLED_TYPES[data['type']]['time_buffering'],
self.process_deque,
data_type=data['type'])
def process_deque(self, data_type):
""" Process the deque and check if I have limit to send the messages now
:param data_type: Type of the message to be sent
:type data_type: string
"""
while len(self.buffer_deques[data_type]) > 0:
if self.rate_limiter.add_hit(data_type):
# We always process the older message first
data = self.buffer_deques[data_type].popleft()
if len(self.buffer_deques[data_type]) == 0:
data['throttled'] = False
self.send_message(data)
else:
reactor.callLater(
CONTROLLED_TYPES[data_type]['time_buffering'],
self.process_deque,
data_type=data_type)
break
def handle_message(self, connection: HathorAdminWebsocketProtocol,
data: bytes) -> None:
""" General message handler, detects type and deletages to specific handler."""
message = json.loads(data)
# we only handle ping messages for now
if message['type'] == 'ping':
self._handle_ping(connection, message)
elif message['type'] == 'subscribe_address':
self._handle_subscribe_address(connection, message)
elif message['type'] == 'unsubscribe_address':
self._handle_unsubscribe_address(connection, message)
def _handle_ping(self, connection: HathorAdminWebsocketProtocol,
message: Dict[Any, Any]) -> None:
""" Handler for ping message, should respond with a simple {"type": "pong"}"""
payload = json.dumps({'type': 'pong'}).encode('utf-8')
connection.sendMessage(payload, False)
def _handle_subscribe_address(self,
connection: HathorAdminWebsocketProtocol,
message: Dict[Any, Any]) -> None:
""" Handler for subscription to an address, consideirs subscription limits."""
addr: str = message['address']
subs: Set[str] = connection.subscribed_to
if len(subs) >= settings.WS_MAX_SUBS_ADDRS_CONN:
payload = json.dumps({
'message':
'Reached maximum number of subscribed '
f'addresses ({settings.WS_MAX_SUBS_ADDRS_CONN}).',
'type':
'subscribe_address',
'success':
False
}).encode('utf-8')
elif self.wallet_index and _count_empty(
subs, self.wallet_index) >= settings.WS_MAX_SUBS_ADDRS_EMPTY:
payload = json.dumps({
'message':
'Reached maximum number of subscribed '
f'addresses without output ({settings.WS_MAX_SUBS_ADDRS_EMPTY}).',
'type':
'subscribe_address',
'success':
False
}).encode('utf-8')
else:
self.address_connections[addr].add(connection)
connection.subscribed_to.add(addr)
payload = json.dumps({
'type': 'subscribe_address',
'success': True
}).encode('utf-8')
connection.sendMessage(payload, False)
def _handle_unsubscribe_address(self,
connection: HathorAdminWebsocketProtocol,
message: Dict[Any, Any]) -> None:
""" Handler for unsubscribing from an address, also removes address connection set if it ends up empty."""
addr = message['address']
if addr in self.address_connections and connection in self.address_connections[
addr]:
connection.subscribed_to.remove(addr)
self._remove_connection_from_address_dict(connection, addr)
# Reply back to the client
payload = json.dumps({
'type': 'unsubscribe_address',
'success': True
}).encode('utf-8')
connection.sendMessage(payload, False)
def _remove_connection_from_address_dict(
self, connection: HathorAdminWebsocketProtocol,
address: str) -> None:
""" Remove a connection from the address connections dict
If this was the last connection for this address, we remove the key
"""
self.address_connections[address].remove(connection)
# If this was the last connection for this address, we delete it from the dict
if len(self.address_connections[address]) == 0:
del self.address_connections[address]
def connection_closed(self,
connection: HathorAdminWebsocketProtocol) -> None:
""" Called when a ws connection is closed
We should remove it from self.connections and from self.address_connections set for each address
"""
self.connections.remove(connection)
for address in connection.subscribed_to:
self._remove_connection_from_address_dict(connection, address)
class HathorProtocol:
""" Implements Hathor Peer-to-Peer Protocol. An instance of this class is
created for each connection.
When the connection is established, the protocol waits for a
HELLO message, which will identify the application and give a
nonce value.
After receiving a HELLO message, the peer must reply with a PEER-ID
message, which will identity the peer through its id, public key,
and endpoints. There must be a signature of the nonce value which
will be checked against the public key.
After the PEER-ID message, the peer is ready to communicate.
The available states are listed in PeerState class.
The available commands are listed in the ProtocolMessages class.
"""
log = Logger()
class PeerState(Enum):
HELLO = HelloState
PEER_ID = PeerIdState
READY = ReadyState
class Metrics:
def __init__(self) -> None:
self.received_messages: int = 0
self.sent_messages: int = 0
self.received_bytes: int = 0
self.sent_bytes: int = 0
self.received_txs: int = 0
self.discarded_txs: int = 0
self.received_blocks: int = 0
self.discarded_blocks: int = 0
def format_bytes(self, value: int) -> str:
""" Format bytes in MB and kB.
"""
if value > 1024*1024:
return '{:11.2f} MB'.format(value / 1024 / 1024)
elif value > 1024:
return '{:11.2f} kB'.format(value / 1024)
else:
return '{} B'.format(value)
def print_stats(self, prefix: str = '') -> None:
""" Print a status of the metrics in stdout.
"""
print('----')
print('{}Received: {:8d} messages {}'.format(
prefix,
self.received_messages,
self.format_bytes(self.received_bytes))
)
print('{}Sent: {:8d} messages {}'.format(
prefix,
self.sent_messages,
self.format_bytes(self.sent_bytes))
)
print('{}Blocks: {:8d} received {:8d} discarded ({:2.0f}%)'.format(
prefix,
self.received_blocks,
self.discarded_blocks,
100.0 * self.discarded_blocks / (self.received_blocks + self.discarded_blocks)
))
print('{}Transactions: {:8d} received {:8d} discarded ({:2.0f}%)'.format(
prefix,
self.received_txs,
self.discarded_txs,
100.0 * self.discarded_txs / (self.received_txs + self.discarded_txs)
))
print('----')
class RateLimitKeys(str, Enum):
GLOBAL = 'global'
class WarningFlags(str, Enum):
NO_PEER_ID_URL = 'no_peer_id_url'
NO_ENTRYPOINTS = 'no_entrypoints'
network: str
my_peer: PeerId
connections: Optional['ConnectionsManager']
node: 'HathorManager'
app_version: str
last_message: float
peer: Optional[PeerId]
transport: ITransport
state: Optional[BaseState]
connection_time: float
_state_instances: Dict[PeerState, BaseState]
connection_string: Optional[str]
expected_peer_id: Optional[str]
warning_flags: Set[str]
connected: bool
initiated_connection: bool
def __init__(self, network: str, my_peer: PeerId, connections: Optional['ConnectionsManager'] = None, *,
node: 'HathorManager', use_ssl: bool) -> None:
self.network = network
self.my_peer = my_peer
self.connections = connections
self.node = node
self._state_instances = {}
self.app_version = 'Unknown'
# The peer on the other side of the connection.
self.peer = None
# The last time a message has been received from this peer.
self.last_message = 0
self.metrics = self.Metrics()
# The last time a request was send to this peer.
self.last_request = 0
# The time in which the connection was established.
self.connection_time = 0.0
# The current state of the connection.
self.state: Optional[BaseState] = None
# Default rate limit
self.ratelimit = RateLimiter()
# self.ratelimit.set_limit(self.RateLimitKeys.GLOBAL, 120, 60)
# Connection string of the peer
# Used to validate if entrypoints has this string
self.connection_string: Optional[str] = None
# Peer id sent in the connection url that is expected to connect (optional)
self.expected_peer_id: Optional[str] = None
# Set of warning flags that may be added during the connection process
self.warning_flags: Set[str] = set()
# If peer is connected
self.connected = False
self.use_ssl = use_ssl
# Set to true if this node initiated the connection
self.initiated_connection = False
def change_state(self, state_enum: PeerState) -> None:
if state_enum not in self._state_instances:
state_cls = state_enum.value
instance = state_cls(self)
instance.state_name = state_enum.name
self._state_instances[state_enum] = instance
new_state = self._state_instances[state_enum]
if new_state != self.state:
if self.state:
self.state.on_exit()
self.state = new_state
if self.state:
self.state.on_enter()
def on_connect(self) -> None:
""" Executed when the connection is established.
"""
remote = self.transport.getPeer()
self.log.info('HathorProtocol.connectionMade(): {remote}', remote=remote)
self.connection_time = time.time()
# The initial state is HELLO.
self.change_state(self.PeerState.HELLO)
self.connected = True
if self.connections:
self.connections.on_peer_connect(self)
def on_disconnect(self, reason: str) -> None:
""" Executed when the connection is lost.
"""
remote = self.transport.getPeer()
self.log.info('HathorProtocol.connectionLost(): {remote} {reason}', remote=remote, reason=reason)
self.connected = False
if self.state:
self.state.on_exit()
if self.connections:
self.connections.on_peer_disconnect(self)
def send_message(self, cmd: ProtocolMessages, payload: Optional[str] = None) -> None:
""" A generic message which must be implemented to send a message
to the peer. It depends on the underlying protocol in which
HathorProtocol is running.
"""
raise NotImplementedError
def recv_message(self, cmd: ProtocolMessages, payload: str) -> Optional[Generator[Any, Any, None]]:
""" Executed when a new message arrives.
"""
assert self.state is not None
self.last_message = self.node.reactor.seconds()
if not self.ratelimit.add_hit(self.RateLimitKeys.GLOBAL):
self.state.send_throttle(self.RateLimitKeys.GLOBAL)
return None
fn = self.state.cmd_map.get(cmd)
if fn is not None:
try:
return fn(payload)
except Exception as e:
self.log.warn('Unhandled Exception: {e!r}', e=e)
raise
else:
self.send_error('Invalid Command: {}'.format(cmd))
return None
def send_error(self, msg: str) -> None:
""" Send an error message to the peer.
"""
self.send_message(ProtocolMessages.ERROR, msg)
def send_error_and_close_connection(self, msg: str) -> None:
""" Send an ERROR message to the peer, and then closes the connection.
"""
self.send_error(msg)
# from twisted docs: "If a producer is being used with the transport, loseConnection will only close
# the connection once the producer is unregistered." We call on_exit to make sure any producers (like
# the one from node_sync) are unregistered
if self.state:
self.state.on_exit()
self.transport.loseConnection()
def handle_error(self, payload: str) -> None:
""" Executed when an ERROR command is received.
"""
self.log.warn('ERROR {payload}', payload=payload)