mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
516 lines
18 KiB
Python
516 lines
18 KiB
Python
import argparse
|
|
import logging
|
|
import os
|
|
import struct
|
|
|
|
from cryptography.hazmat.primitives.asymmetric import (
|
|
x25519,
|
|
)
|
|
import multiaddr
|
|
import trio
|
|
|
|
from libp2p import (
|
|
generate_new_rsa_identity,
|
|
new_host,
|
|
)
|
|
from libp2p.custom_types import (
|
|
TProtocol,
|
|
)
|
|
from libp2p.network.stream.net_stream import (
|
|
INetStream,
|
|
)
|
|
from libp2p.security.noise.transport import Transport as NoiseTransport
|
|
from libp2p.stream_muxer.yamux.yamux import (
|
|
Yamux,
|
|
YamuxStream,
|
|
)
|
|
from libp2p.stream_muxer.yamux.yamux import PROTOCOL_ID as YAMUX_PROTOCOL_ID
|
|
|
|
# Configure detailed logging
|
|
logging.basicConfig(
|
|
level=logging.DEBUG,
|
|
format="%(asctime)s - %(levelname)s - %(message)s",
|
|
handlers=[
|
|
logging.StreamHandler(),
|
|
logging.FileHandler("ping_debug.log", mode="w", encoding="utf-8"),
|
|
],
|
|
)
|
|
|
|
# Protocol constants - must match rust-libp2p exactly
|
|
PING_PROTOCOL_ID = TProtocol("/ipfs/ping/1.0.0")
|
|
PING_LENGTH = 32
|
|
RESP_TIMEOUT = 30
|
|
MAX_FRAME_SIZE = 1024 * 1024 # 1MB max frame size
|
|
|
|
|
|
class InteropYamux(Yamux):
|
|
"""Enhanced Yamux with proper rust-libp2p interoperability"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
logging.info("InteropYamux.__init__ called")
|
|
super().__init__(*args, **kwargs)
|
|
self.frame_count = 0
|
|
self.debug_frames = True
|
|
|
|
async def _read_exact_bytes(self, n):
|
|
"""Read exactly n bytes from the connection with proper error handling"""
|
|
if n == 0:
|
|
return b""
|
|
|
|
if n > MAX_FRAME_SIZE:
|
|
logging.error(f"Requested read size {n} exceeds maximum {MAX_FRAME_SIZE}")
|
|
return None
|
|
|
|
data = b""
|
|
while len(data) < n:
|
|
try:
|
|
remaining = n - len(data)
|
|
chunk = await self.secured_conn.read(remaining)
|
|
except (trio.ClosedResourceError, trio.BrokenResourceError):
|
|
logging.debug(
|
|
f"Connection closed while reading {n}"
|
|
f"bytes (got {len(data)}) for peer {self.peer_id}"
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
logging.error(f"Error reading {n} bytes: {e}")
|
|
return None
|
|
|
|
if not chunk:
|
|
logging.debug(
|
|
f"Connection closed while reading {n}"
|
|
f"bytes (got {len(data)}) for peer {self.peer_id}"
|
|
)
|
|
return None
|
|
data += chunk
|
|
return data
|
|
|
|
async def handle_incoming(self):
|
|
"""Enhanced incoming frame handler with better error recovery"""
|
|
logging.info(f"Starting Yamux for {self.peer_id}")
|
|
|
|
consecutive_errors = 0
|
|
max_consecutive_errors = 3
|
|
|
|
while not self.event_shutting_down.is_set():
|
|
try:
|
|
# Read frame header (12 bytes)
|
|
header_data = await self._read_exact_bytes(12)
|
|
if header_data is None:
|
|
logging.debug(
|
|
f"Connection closed or incomplete"
|
|
f"header for peer {self.peer_id}"
|
|
)
|
|
break
|
|
|
|
# Quick sanity check for protocol data leakage
|
|
if (
|
|
b"/ipfs" in header_data
|
|
or b"multi" in header_data
|
|
or b"noise" in header_data
|
|
):
|
|
logging.error(
|
|
f"Protocol data in header position: {header_data.hex()}"
|
|
)
|
|
break
|
|
|
|
try:
|
|
# Unpack header: version, type, flags, stream_id, length
|
|
version, msg_type, flags, stream_id, length = struct.unpack(
|
|
">BBHII", header_data
|
|
)
|
|
|
|
# Validate header values strictly
|
|
if version != 0:
|
|
logging.error(f"Invalid yamux version {version}, expected 0")
|
|
break
|
|
|
|
if msg_type not in [0, 1, 2, 3]:
|
|
logging.error(f"Invalid message type {msg_type}, expected 0-3")
|
|
break
|
|
|
|
if length > MAX_FRAME_SIZE:
|
|
logging.error(f"Frame too large: {length} > {MAX_FRAME_SIZE}")
|
|
break
|
|
|
|
# Additional validation for ping frames
|
|
if msg_type == 2 and length != 4:
|
|
logging.error(
|
|
f"Invalid ping frame length: {length}, expected 4"
|
|
)
|
|
break
|
|
|
|
# Log frame details
|
|
logging.debug(
|
|
f"Received header for peer {self.peer_id}"
|
|
f": type={msg_type}, flags={flags},"
|
|
f"stream_id={stream_id}, length={length}"
|
|
)
|
|
|
|
consecutive_errors = 0 # Reset error counter on successful parse
|
|
|
|
except struct.error as e:
|
|
consecutive_errors += 1
|
|
logging.error(
|
|
f"Header parse error #{consecutive_errors}"
|
|
f": {e}, data: {header_data.hex()}"
|
|
)
|
|
if consecutive_errors >= max_consecutive_errors:
|
|
logging.error("Too many consecutive header parse errors")
|
|
break
|
|
continue
|
|
|
|
# Read payload if present
|
|
payload = b""
|
|
if length > 0:
|
|
payload = await self._read_exact_bytes(length)
|
|
if payload is None:
|
|
logging.debug(
|
|
f"Failed to read payload of"
|
|
f"{length} bytes for peer {self.peer_id}"
|
|
)
|
|
break
|
|
if len(payload) != length:
|
|
logging.error(
|
|
f"Payload length mismatch:"
|
|
f"got {len(payload)}, expected {length}"
|
|
)
|
|
break
|
|
|
|
# Process frame by type
|
|
if msg_type == 0: # Data frame
|
|
await self._handle_data_frame(stream_id, flags, payload)
|
|
|
|
elif msg_type == 1: # Window update
|
|
await self._handle_window_update(stream_id, payload)
|
|
|
|
elif msg_type == 2: # Ping frame
|
|
await self._handle_ping_frame(stream_id, flags, payload)
|
|
|
|
elif msg_type == 3: # GoAway frame
|
|
await self._handle_goaway_frame(payload)
|
|
break
|
|
|
|
except (trio.ClosedResourceError, trio.BrokenResourceError):
|
|
logging.debug(
|
|
f"Connection closed during frame processing for peer {self.peer_id}"
|
|
)
|
|
break
|
|
except Exception as e:
|
|
consecutive_errors += 1
|
|
logging.error(f"Frame processing error #{consecutive_errors}: {e}")
|
|
if consecutive_errors >= max_consecutive_errors:
|
|
logging.error("Too many consecutive frame processing errors")
|
|
break
|
|
|
|
await self._cleanup_on_error()
|
|
|
|
async def _handle_data_frame(self, stream_id, flags, payload):
|
|
"""Handle data frames with proper stream lifecycle"""
|
|
if stream_id == 0:
|
|
logging.warning("Received data frame for stream 0 (control stream)")
|
|
return
|
|
|
|
# Handle SYN flag - new stream creation
|
|
if flags & 0x1: # SYN flag
|
|
if stream_id in self.streams:
|
|
logging.warning(f"SYN received for existing stream {stream_id}")
|
|
else:
|
|
logging.debug(
|
|
f"Creating new stream {stream_id} for peer {self.peer_id}"
|
|
)
|
|
stream = YamuxStream(self, stream_id, is_outbound=False)
|
|
async with self.streams_lock:
|
|
self.streams[stream_id] = stream
|
|
|
|
# Send the new stream to the handler
|
|
await self.new_stream_send_channel.send(stream)
|
|
logging.debug(f"Sent stream {stream_id} to handler")
|
|
|
|
# Add data to stream buffer if stream exists
|
|
if stream_id in self.streams:
|
|
stream = self.streams[stream_id]
|
|
if payload:
|
|
# Add to stream's receive buffer
|
|
async with self.streams_lock:
|
|
if not hasattr(stream, "_receive_buffer"):
|
|
stream._receive_buffer = bytearray()
|
|
if not hasattr(stream, "_receive_event"):
|
|
stream._receive_event = trio.Event()
|
|
stream._receive_buffer.extend(payload)
|
|
stream._receive_event.set()
|
|
|
|
# Handle stream closure flags
|
|
if flags & 0x2: # FIN flag
|
|
stream.recv_closed = True
|
|
logging.debug(f"Stream {stream_id} received FIN")
|
|
if flags & 0x4: # RST flag
|
|
stream.reset_received = True
|
|
logging.debug(f"Stream {stream_id} received RST")
|
|
else:
|
|
if payload:
|
|
logging.warning(f"Received data for unknown stream {stream_id}")
|
|
|
|
async def _handle_window_update(self, stream_id, payload):
|
|
"""Handle window update frames"""
|
|
if len(payload) != 4:
|
|
logging.warning(f"Invalid window update payload length: {len(payload)}")
|
|
return
|
|
|
|
delta = struct.unpack(">I", payload)[0]
|
|
logging.debug(f"Window update: stream={stream_id}, delta={delta}")
|
|
|
|
async with self.streams_lock:
|
|
if stream_id in self.streams:
|
|
if not hasattr(self.streams[stream_id], "send_window"):
|
|
self.streams[stream_id].send_window = 256 * 1024 # Default window
|
|
self.streams[stream_id].send_window += delta
|
|
|
|
async def _handle_ping_frame(self, stream_id, flags, payload):
|
|
"""Handle ping/pong frames with proper validation"""
|
|
if len(payload) != 4:
|
|
logging.warning(f"Invalid ping payload length: {len(payload)} (expected 4)")
|
|
return
|
|
|
|
ping_value = struct.unpack(">I", payload)[0]
|
|
|
|
if flags & 0x1: # SYN flag - ping request
|
|
logging.debug(
|
|
f"Received ping request with value {ping_value} for peer {self.peer_id}"
|
|
)
|
|
# Send pong response (ACK flag = 0x2)
|
|
try:
|
|
pong_header = struct.pack(
|
|
">BBHII", 0, 2, 0x2, 0, 4
|
|
) # Version=0, Type=2, Flags=ACK, StreamID=0, Length=4
|
|
pong_payload = struct.pack(">I", ping_value)
|
|
await self.secured_conn.write(pong_header + pong_payload)
|
|
logging.debug(f"Sent pong response with value {ping_value}")
|
|
except Exception as e:
|
|
logging.error(f"Failed to send pong response: {e}")
|
|
else:
|
|
# Pong response
|
|
logging.debug(f"Received pong response with value {ping_value}")
|
|
|
|
async def _handle_goaway_frame(self, payload):
|
|
"""Handle GoAway frames"""
|
|
if len(payload) != 4:
|
|
logging.warning(f"Invalid GoAway payload length: {len(payload)}")
|
|
return
|
|
|
|
code = struct.unpack(">I", payload)[0]
|
|
logging.info(f"Received GoAway frame with code {code}")
|
|
self.event_shutting_down.set()
|
|
|
|
|
|
async def handle_ping(stream: INetStream) -> None:
|
|
peer_id = stream.muxed_conn.peer_id
|
|
logging.info(f"Handling ping stream from {peer_id}")
|
|
|
|
try:
|
|
with trio.fail_after(RESP_TIMEOUT):
|
|
# Read initial protocol negotiation
|
|
initial_data = await stream.read(1024)
|
|
logging.debug(
|
|
f"Received initial stream data from {peer_id}"
|
|
f": {initial_data.hex()} (length={len(initial_data)})"
|
|
)
|
|
if initial_data == b"/ipfs/ping/1.0.0\n":
|
|
logging.debug(
|
|
f"Confirmed /ipfs/ping/1.0.0 protocol negotiation from {peer_id}"
|
|
)
|
|
else:
|
|
logging.warning(f"Unexpected initial data: {initial_data!r}")
|
|
|
|
# Read ping payload
|
|
payload = await stream.read(PING_LENGTH)
|
|
if not payload:
|
|
logging.info(f"Stream closed by {peer_id}")
|
|
return
|
|
if len(payload) != PING_LENGTH:
|
|
logging.warning(
|
|
f"Unexpected payload length"
|
|
f" {len(payload)} from {peer_id}: {payload.hex()}"
|
|
)
|
|
return
|
|
logging.info(
|
|
f"Received ping from {peer_id}:"
|
|
f" {payload[:8].hex()}... (length={len(payload)})"
|
|
)
|
|
await stream.write(payload)
|
|
logging.info(f"Sent pong to {peer_id}: {payload[:8].hex()}...")
|
|
|
|
except trio.TooSlowError:
|
|
logging.warning(f"Ping timeout with {peer_id}")
|
|
except trio.BrokenResourceError:
|
|
logging.info(f"Connection broken with {peer_id}")
|
|
except Exception as e:
|
|
logging.error(f"Error handling ping from {peer_id}: {e}")
|
|
finally:
|
|
try:
|
|
await stream.close()
|
|
logging.debug(f"Closed ping stream with {peer_id}")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def send_ping(stream: INetStream) -> None:
|
|
peer_id = stream.muxed_conn.peer_id
|
|
try:
|
|
payload = os.urandom(PING_LENGTH)
|
|
logging.info(f"Sending ping to {peer_id}: {payload[:8].hex()}...")
|
|
with trio.fail_after(RESP_TIMEOUT):
|
|
await stream.write(payload)
|
|
logging.debug(f"Ping sent to {peer_id}")
|
|
response = await stream.read(PING_LENGTH)
|
|
if not response:
|
|
logging.error(f"No pong response from {peer_id}")
|
|
return
|
|
if len(response) != PING_LENGTH:
|
|
logging.warning(
|
|
f"Pong length mismatch: got {len(response)}, expected {PING_LENGTH}"
|
|
)
|
|
if response == payload:
|
|
logging.info(f"Ping successful! Pong matches from {peer_id}")
|
|
else:
|
|
logging.warning(f"Pong mismatch from {peer_id}")
|
|
except trio.TooSlowError:
|
|
logging.error(f"Ping timeout to {peer_id}")
|
|
except Exception as e:
|
|
logging.error(f"Error sending ping to {peer_id}: {e}")
|
|
finally:
|
|
try:
|
|
await stream.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def create_noise_keypair():
|
|
try:
|
|
x25519_private_key = x25519.X25519PrivateKey.generate()
|
|
|
|
class NoisePrivateKey:
|
|
def __init__(self, key):
|
|
self._key = key
|
|
|
|
def to_bytes(self):
|
|
return self._key.private_bytes_raw()
|
|
|
|
def public_key(self):
|
|
return NoisePublicKey(self._key.public_key())
|
|
|
|
def get_public_key(self):
|
|
return NoisePublicKey(self._key.public_key())
|
|
|
|
class NoisePublicKey:
|
|
def __init__(self, key):
|
|
self._key = key
|
|
|
|
def to_bytes(self):
|
|
return self._key.public_bytes_raw()
|
|
|
|
return NoisePrivateKey(x25519_private_key)
|
|
except Exception as e:
|
|
logging.error(f"Failed to create Noise keypair: {e}")
|
|
return None
|
|
|
|
|
|
def info_from_p2p_addr(addr):
|
|
"""Extract peer info from multiaddr - you'll need to implement this"""
|
|
# This is a placeholder - you need to implement the actual parsing
|
|
# based on your libp2p implementation
|
|
|
|
|
|
async def run(port: int, destination: str) -> None:
|
|
listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
|
|
|
|
try:
|
|
key_pair = generate_new_rsa_identity()
|
|
logging.debug("Generated RSA keypair")
|
|
|
|
noise_privkey = create_noise_keypair()
|
|
logging.debug("Generated Noise keypair")
|
|
except Exception as e:
|
|
logging.error(f"Key generation failed: {e}")
|
|
raise
|
|
|
|
noise_transport = NoiseTransport(key_pair, noise_privkey=noise_privkey)
|
|
logging.debug(f"Noise transport initialized: {noise_transport}")
|
|
sec_opt = {TProtocol("/noise"): noise_transport}
|
|
muxer_opt = {TProtocol(YAMUX_PROTOCOL_ID): InteropYamux}
|
|
|
|
logging.info(f"Using muxer: {muxer_opt}")
|
|
|
|
host = new_host(key_pair=key_pair, sec_opt=sec_opt, muxer_opt=muxer_opt)
|
|
|
|
peer_id = host.get_id().pretty()
|
|
logging.info(f"Host peer ID: {peer_id}")
|
|
|
|
async with host.run(listen_addrs=[listen_addr]), trio.open_nursery():
|
|
if not destination:
|
|
host.set_stream_handler(PING_PROTOCOL_ID, handle_ping)
|
|
|
|
logging.info(f"Server listening on {listen_addr}")
|
|
logging.info(f"Full address: {listen_addr}/p2p/{peer_id}")
|
|
logging.info("Waiting for connections...")
|
|
|
|
await trio.sleep_forever()
|
|
else:
|
|
maddr = multiaddr.Multiaddr(destination)
|
|
info = info_from_p2p_addr(maddr)
|
|
|
|
logging.info(f"Connecting to {info.peer_id}")
|
|
|
|
try:
|
|
with trio.fail_after(30):
|
|
await host.connect(info)
|
|
logging.info(f"Connected to {info.peer_id}")
|
|
|
|
await trio.sleep(2.0)
|
|
|
|
logging.info(f"Opening ping stream to {info.peer_id}")
|
|
stream = await host.new_stream(info.peer_id, [PING_PROTOCOL_ID])
|
|
logging.info(f"Opened ping stream to {info.peer_id}")
|
|
|
|
await trio.sleep(0.5)
|
|
|
|
await send_ping(stream)
|
|
|
|
logging.info("Ping completed successfully")
|
|
|
|
logging.info("Keeping connection alive for 5 seconds...")
|
|
await trio.sleep(5.0)
|
|
except trio.TooSlowError:
|
|
logging.error(f"Connection timeout to {info.peer_id}")
|
|
raise
|
|
except Exception as e:
|
|
logging.error(f"Connection failed to {info.peer_id}: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="libp2p ping with Rust interoperability"
|
|
)
|
|
parser.add_argument(
|
|
"-p", "--port", default=8000, type=int, help="Port to listen on"
|
|
)
|
|
parser.add_argument("-d", "--destination", type=str, help="Destination multiaddr")
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
trio.run(run, args.port, args.destination)
|
|
except KeyboardInterrupt:
|
|
logging.info("Terminated by user")
|
|
except Exception as e:
|
|
logging.error(f"Fatal error: {e}")
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|