444 lines
16 KiB
Python
444 lines
16 KiB
Python
"""
|
|
TCP NAT — chip-initiated outbound connections to the host network.
|
|
|
|
The implementation follows RFC 793 §3.4-§3.9 closely enough to handle
|
|
real-world MicroPython workloads:
|
|
|
|
- Three-way handshake (SYN → SYN+ACK → ACK)
|
|
- Bidirectional data flow with proper seq/ack accounting
|
|
- Half-close handling (FIN from either side)
|
|
- RST as the cheap escape hatch on protocol errors
|
|
- MSS option negotiation (we advertise TCP_MSS = MTU - 40)
|
|
- Window clamped to TCP_WINDOW (no window scaling)
|
|
|
|
Per-connection state lives in a TcpConnection object keyed by
|
|
(chip_port, dst_ip, dst_port). Each connection owns an asyncio
|
|
StreamReader/StreamWriter to the real host endpoint.
|
|
|
|
States we transition through, simplified to chip-initiated only:
|
|
|
|
CLOSED
|
|
│ chip SYN
|
|
▼
|
|
SYN_RCVD ── send SYN+ACK back to chip
|
|
│ chip ACK
|
|
▼
|
|
ESTABLISHED ── pump bytes both ways
|
|
│
|
|
├── chip FIN ──► CLOSE_WAIT ── after host close: LAST_ACK ──► CLOSED
|
|
└── host EOF ──► FIN_WAIT_1 ── after chip ACK: FIN_WAIT_2 ──► CLOSED
|
|
|
|
We deliberately don't implement TIME_WAIT — the chip does, we just GC
|
|
once both sides have FIN'd. This is the same simplification slirp uses.
|
|
|
|
Sequence numbers wrap at 2³² — every comparison goes through
|
|
``_seq_lt`` / ``_seq_geq`` which use modular arithmetic.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import random
|
|
import struct
|
|
from dataclasses import dataclass
|
|
from typing import Awaitable, Callable, Dict, Optional, Tuple
|
|
|
|
from .consts import (
|
|
GATEWAY_MAC,
|
|
IPPROTO_TCP,
|
|
TCP_ACK,
|
|
TCP_FIN,
|
|
TCP_MSS,
|
|
TCP_PSH,
|
|
TCP_RST,
|
|
TCP_SYN,
|
|
TCP_WINDOW,
|
|
bytes_to_ip,
|
|
)
|
|
from .protocols import IPv4, TCP, make_frame_ipv4, parse_tcp_options
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
InjectFn = Callable[[bytes], Awaitable[None]]
|
|
|
|
|
|
# ─── Sequence number arithmetic (modular 32-bit) ─────────────────────
|
|
|
|
def _seq_add(a: int, b: int) -> int:
|
|
return (a + b) & 0xffffffff
|
|
|
|
|
|
def _seq_lt(a: int, b: int) -> bool:
|
|
"""RFC 1323-style: a < b modulo 2^32."""
|
|
return ((a - b) & 0xffffffff) >= 0x80000000
|
|
|
|
|
|
def _seq_leq(a: int, b: int) -> bool:
|
|
return a == b or _seq_lt(a, b)
|
|
|
|
|
|
def _seq_diff(a: int, b: int) -> int:
|
|
"""Distance a - b modulo 2^32, signed."""
|
|
d = (a - b) & 0xffffffff
|
|
if d & 0x80000000:
|
|
d -= 0x100000000
|
|
return d
|
|
|
|
|
|
# ─── Per-connection state ────────────────────────────────────────────
|
|
|
|
class _State:
|
|
SYN_RCVD = 'SYN_RCVD'
|
|
ESTABLISHED = 'ESTABLISHED'
|
|
FIN_WAIT_1 = 'FIN_WAIT_1' # we (host side) sent FIN, waiting for chip ACK
|
|
FIN_WAIT_2 = 'FIN_WAIT_2' # chip ACKed our FIN
|
|
CLOSE_WAIT = 'CLOSE_WAIT' # chip sent FIN, host still has more to send
|
|
LAST_ACK = 'LAST_ACK' # both sides FIN'd, waiting for last ACK
|
|
CLOSED = 'CLOSED'
|
|
|
|
|
|
@dataclass
|
|
class TcpConnection:
|
|
chip_ip: bytes
|
|
chip_port: int
|
|
dst_ip: bytes
|
|
dst_port: int
|
|
chip_mac: bytes
|
|
state: str = _State.CLOSED
|
|
chip_isn: int = 0 # initial chip seq we observed
|
|
our_isn: int = 0 # initial seq we picked
|
|
our_seq: int = 0 # next seq we'll put on the wire chipward
|
|
chip_seq: int = 0 # next seq we expect from chip
|
|
chip_window: int = 0
|
|
mss: int = TCP_MSS
|
|
host_reader: Optional[asyncio.StreamReader] = None
|
|
host_writer: Optional[asyncio.StreamWriter] = None
|
|
host_pump_task: Optional[asyncio.Task] = None
|
|
last_activity: float = 0.0
|
|
|
|
def key(self) -> Tuple[bytes, int, bytes, int]:
|
|
return (self.chip_ip, self.chip_port, self.dst_ip, self.dst_port)
|
|
|
|
|
|
# ─── NAT manager ─────────────────────────────────────────────────────
|
|
|
|
class TcpNat:
|
|
"""
|
|
Manages every chip-initiated TCP connection. ``inject`` is the
|
|
callback the bridge gives us to push Ethernet frames back to the
|
|
chip; calls into the manager are made from the bridge whenever an
|
|
IP-with-protocol-TCP frame arrives from the chip.
|
|
"""
|
|
|
|
def __init__(self, inject: InjectFn) -> None:
|
|
self._inject = inject
|
|
self._conns: Dict[Tuple[bytes, int, bytes, int], TcpConnection] = {}
|
|
|
|
# ── Entry point from the bridge ────────────────────────────────
|
|
|
|
async def handle_chip_segment(
|
|
self, chip_mac: bytes, ip: IPv4, tcp: TCP,
|
|
) -> None:
|
|
key = (bytes(ip.src), tcp.src_port, bytes(ip.dst), tcp.dst_port)
|
|
conn = self._conns.get(key)
|
|
|
|
if tcp.flags & TCP_RST:
|
|
# Chip aborted — tear down silently.
|
|
if conn:
|
|
await self._close(conn, send_rst=False)
|
|
return
|
|
|
|
if conn is None:
|
|
if tcp.flags & TCP_SYN and not (tcp.flags & TCP_ACK):
|
|
await self._on_passive_syn(chip_mac, ip, tcp)
|
|
else:
|
|
# Stray segment with no connection: respond with RST.
|
|
await self._send_rst(chip_mac, ip, tcp)
|
|
return
|
|
|
|
# Update bookkeeping that's the same in every state.
|
|
conn.chip_window = tcp.window
|
|
|
|
if conn.state == _State.SYN_RCVD:
|
|
await self._on_handshake_complete(conn, tcp)
|
|
elif conn.state == _State.ESTABLISHED:
|
|
await self._on_data(conn, tcp)
|
|
elif conn.state == _State.FIN_WAIT_1:
|
|
await self._on_fin_wait_1(conn, tcp)
|
|
elif conn.state == _State.FIN_WAIT_2:
|
|
await self._on_fin_wait_2(conn, tcp)
|
|
elif conn.state == _State.CLOSE_WAIT:
|
|
# Chip should be quiet; ignore unless RST/FIN retransmit.
|
|
pass
|
|
elif conn.state == _State.LAST_ACK:
|
|
if tcp.flags & TCP_ACK and _seq_geq_or_eq(tcp.ack, _seq_add(conn.our_seq, 0)):
|
|
await self._close(conn, send_rst=False)
|
|
|
|
# ── State handlers ─────────────────────────────────────────────
|
|
|
|
async def _on_passive_syn(self, chip_mac: bytes, ip: IPv4, tcp: TCP) -> None:
|
|
"""Chip is opening a new connection — we play the server."""
|
|
opts = parse_tcp_options(tcp.options)
|
|
mss = opts.get('mss', TCP_MSS)
|
|
if mss > TCP_MSS:
|
|
mss = TCP_MSS
|
|
|
|
our_isn = random.randint(0, 0xffffffff)
|
|
conn = TcpConnection(
|
|
chip_ip=bytes(ip.src),
|
|
chip_port=tcp.src_port,
|
|
dst_ip=bytes(ip.dst),
|
|
dst_port=tcp.dst_port,
|
|
chip_mac=chip_mac,
|
|
state=_State.CLOSED,
|
|
chip_isn=tcp.seq,
|
|
our_isn=our_isn,
|
|
our_seq=_seq_add(our_isn, 1), # SYN counts as 1 byte
|
|
chip_seq=_seq_add(tcp.seq, 1),
|
|
chip_window=tcp.window,
|
|
mss=mss,
|
|
last_activity=asyncio.get_event_loop().time(),
|
|
)
|
|
|
|
# Try to establish the host-side connection. If it fails we
|
|
# send RST to the chip and never store the connection.
|
|
try:
|
|
reader, writer = await asyncio.wait_for(
|
|
asyncio.open_connection(bytes_to_ip(conn.dst_ip), conn.dst_port),
|
|
timeout=10.0,
|
|
)
|
|
except (OSError, asyncio.TimeoutError) as e:
|
|
logger.info(
|
|
'[picow-tcp] connect %s:%d failed: %s',
|
|
bytes_to_ip(conn.dst_ip), conn.dst_port, e,
|
|
)
|
|
await self._send_rst(chip_mac, ip, tcp)
|
|
return
|
|
|
|
conn.host_reader = reader
|
|
conn.host_writer = writer
|
|
conn.state = _State.SYN_RCVD
|
|
self._conns[conn.key()] = conn
|
|
|
|
# Send SYN+ACK back. Advertise our MSS option.
|
|
await self._send(conn, flags=TCP_SYN | TCP_ACK,
|
|
seq=conn.our_isn, ack=conn.chip_seq,
|
|
options=_mss_option(conn.mss))
|
|
|
|
# Start the host → chip pump. It will block until handshake completes.
|
|
conn.host_pump_task = asyncio.create_task(self._pump_host_to_chip(conn))
|
|
|
|
async def _on_handshake_complete(self, conn: TcpConnection, tcp: TCP) -> None:
|
|
"""We're SYN_RCVD; this should be the chip's ACK of our SYN+ACK."""
|
|
if not (tcp.flags & TCP_ACK):
|
|
return
|
|
if tcp.ack != conn.our_seq:
|
|
# Stale or duplicate; ignore.
|
|
return
|
|
conn.state = _State.ESTABLISHED
|
|
# Some clients piggyback data on the final handshake ACK.
|
|
if tcp.payload:
|
|
await self._on_data(conn, tcp)
|
|
|
|
async def _on_data(self, conn: TcpConnection, tcp: TCP) -> None:
|
|
# Reject out-of-order. The chip will retransmit.
|
|
if tcp.payload:
|
|
if tcp.seq != conn.chip_seq:
|
|
# Re-ACK what we have (forces retransmit).
|
|
await self._ack_only(conn)
|
|
return
|
|
assert conn.host_writer is not None
|
|
try:
|
|
conn.host_writer.write(tcp.payload)
|
|
await conn.host_writer.drain()
|
|
except (ConnectionError, OSError):
|
|
await self._close(conn, send_rst=True)
|
|
return
|
|
conn.chip_seq = _seq_add(conn.chip_seq, len(tcp.payload))
|
|
await self._ack_only(conn)
|
|
elif (tcp.flags & TCP_ACK) and tcp.ack and tcp.seq == conn.chip_seq:
|
|
# Pure ACK or keep-alive — nothing to do.
|
|
pass
|
|
|
|
if tcp.flags & TCP_FIN:
|
|
conn.chip_seq = _seq_add(conn.chip_seq, 1)
|
|
conn.state = _State.CLOSE_WAIT
|
|
# Tell the host side we're done sending.
|
|
if conn.host_writer is not None:
|
|
try:
|
|
conn.host_writer.write_eof()
|
|
except (OSError, ConnectionError):
|
|
pass
|
|
await self._ack_only(conn)
|
|
# Stay in CLOSE_WAIT until the host pump finishes draining
|
|
# whatever's still inbound, then it transitions to LAST_ACK.
|
|
|
|
async def _on_fin_wait_1(self, conn: TcpConnection, tcp: TCP) -> None:
|
|
# Waiting for the chip to ACK our FIN.
|
|
if (tcp.flags & TCP_ACK) and tcp.ack == conn.our_seq:
|
|
conn.state = _State.FIN_WAIT_2
|
|
if tcp.flags & TCP_FIN:
|
|
conn.chip_seq = _seq_add(conn.chip_seq, 1)
|
|
await self._ack_only(conn)
|
|
await self._close(conn, send_rst=False)
|
|
|
|
async def _on_fin_wait_2(self, conn: TcpConnection, tcp: TCP) -> None:
|
|
if tcp.flags & TCP_FIN:
|
|
conn.chip_seq = _seq_add(conn.chip_seq, 1)
|
|
await self._ack_only(conn)
|
|
await self._close(conn, send_rst=False)
|
|
|
|
# ── Host → chip pump ───────────────────────────────────────────
|
|
|
|
async def _pump_host_to_chip(self, conn: TcpConnection) -> None:
|
|
"""Read bytes from the real host socket and segment them to the chip."""
|
|
try:
|
|
assert conn.host_reader is not None
|
|
# Wait for handshake to complete before pushing.
|
|
while conn.state == _State.SYN_RCVD:
|
|
await asyncio.sleep(0.005)
|
|
|
|
while conn.state in (_State.ESTABLISHED, _State.CLOSE_WAIT):
|
|
chunk = await conn.host_reader.read(conn.mss)
|
|
if not chunk:
|
|
break
|
|
# Segment if needed (read() should already cap at mss).
|
|
while chunk:
|
|
seg = chunk[:conn.mss]
|
|
chunk = chunk[conn.mss:]
|
|
await self._send(
|
|
conn, flags=TCP_ACK | TCP_PSH,
|
|
seq=conn.our_seq, ack=conn.chip_seq,
|
|
payload=seg,
|
|
)
|
|
conn.our_seq = _seq_add(conn.our_seq, len(seg))
|
|
conn.last_activity = asyncio.get_event_loop().time()
|
|
|
|
# Host side EOF — send FIN.
|
|
if conn.state == _State.ESTABLISHED:
|
|
conn.state = _State.FIN_WAIT_1
|
|
await self._send(conn, flags=TCP_ACK | TCP_FIN,
|
|
seq=conn.our_seq, ack=conn.chip_seq)
|
|
conn.our_seq = _seq_add(conn.our_seq, 1)
|
|
elif conn.state == _State.CLOSE_WAIT:
|
|
conn.state = _State.LAST_ACK
|
|
await self._send(conn, flags=TCP_ACK | TCP_FIN,
|
|
seq=conn.our_seq, ack=conn.chip_seq)
|
|
conn.our_seq = _seq_add(conn.our_seq, 1)
|
|
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
logger.exception('[picow-tcp] pump crashed')
|
|
await self._close(conn, send_rst=True)
|
|
|
|
# ── Frame emission ─────────────────────────────────────────────
|
|
|
|
async def _send(
|
|
self,
|
|
conn: TcpConnection,
|
|
flags: int,
|
|
seq: int,
|
|
ack: int,
|
|
options: bytes = b'',
|
|
payload: bytes = b'',
|
|
) -> None:
|
|
tcp = TCP(
|
|
src_port=conn.dst_port, # chip's "remote" = our destination
|
|
dst_port=conn.chip_port,
|
|
seq=seq & 0xffffffff,
|
|
ack=ack & 0xffffffff,
|
|
flags=flags,
|
|
window=TCP_WINDOW,
|
|
options=options,
|
|
payload=payload,
|
|
)
|
|
# Note: we swap src/dst here because we're emitting the chip's
|
|
# peer's segment — what would have come back from the host.
|
|
ipv4_payload = tcp.to_bytes(conn.dst_ip, conn.chip_ip)
|
|
frame = make_frame_ipv4(
|
|
dst_mac=conn.chip_mac,
|
|
src_mac=GATEWAY_MAC,
|
|
src_ip=conn.dst_ip,
|
|
dst_ip=conn.chip_ip,
|
|
protocol=IPPROTO_TCP,
|
|
l4_payload=ipv4_payload,
|
|
)
|
|
await self._inject(frame)
|
|
|
|
async def _ack_only(self, conn: TcpConnection) -> None:
|
|
await self._send(
|
|
conn, flags=TCP_ACK,
|
|
seq=conn.our_seq, ack=conn.chip_seq,
|
|
)
|
|
|
|
async def _send_rst(self, chip_mac: bytes, ip: IPv4, tcp: TCP) -> None:
|
|
rst = TCP(
|
|
src_port=tcp.dst_port,
|
|
dst_port=tcp.src_port,
|
|
seq=tcp.ack if (tcp.flags & TCP_ACK) else 0,
|
|
ack=_seq_add(tcp.seq, 1 if (tcp.flags & TCP_SYN) else len(tcp.payload)),
|
|
flags=TCP_RST | TCP_ACK,
|
|
window=0,
|
|
)
|
|
ipv4_payload = rst.to_bytes(bytes(ip.dst), bytes(ip.src))
|
|
frame = make_frame_ipv4(
|
|
dst_mac=chip_mac,
|
|
src_mac=GATEWAY_MAC,
|
|
src_ip=bytes(ip.dst),
|
|
dst_ip=bytes(ip.src),
|
|
protocol=IPPROTO_TCP,
|
|
l4_payload=ipv4_payload,
|
|
)
|
|
await self._inject(frame)
|
|
|
|
# ── Teardown ───────────────────────────────────────────────────
|
|
|
|
async def _close(self, conn: TcpConnection, send_rst: bool) -> None:
|
|
if conn.state == _State.CLOSED:
|
|
return
|
|
conn.state = _State.CLOSED
|
|
if send_rst:
|
|
try:
|
|
rst = TCP(
|
|
src_port=conn.dst_port,
|
|
dst_port=conn.chip_port,
|
|
seq=conn.our_seq,
|
|
ack=conn.chip_seq,
|
|
flags=TCP_RST,
|
|
)
|
|
ipv4_payload = rst.to_bytes(conn.dst_ip, conn.chip_ip)
|
|
await self._inject(make_frame_ipv4(
|
|
dst_mac=conn.chip_mac,
|
|
src_mac=GATEWAY_MAC,
|
|
src_ip=conn.dst_ip,
|
|
dst_ip=conn.chip_ip,
|
|
protocol=IPPROTO_TCP,
|
|
l4_payload=ipv4_payload,
|
|
))
|
|
except Exception:
|
|
pass
|
|
if conn.host_writer is not None:
|
|
try:
|
|
conn.host_writer.close()
|
|
except Exception:
|
|
pass
|
|
if conn.host_pump_task is not None and not conn.host_pump_task.done():
|
|
conn.host_pump_task.cancel()
|
|
self._conns.pop(conn.key(), None)
|
|
|
|
async def shutdown(self) -> None:
|
|
for conn in list(self._conns.values()):
|
|
await self._close(conn, send_rst=True)
|
|
|
|
|
|
# ─── helpers ────────────────────────────────────────────────────────
|
|
|
|
def _mss_option(mss: int) -> bytes:
|
|
return b'\x02\x04' + struct.pack('!H', mss)
|
|
|
|
|
|
def _seq_geq_or_eq(a: int, b: int) -> bool:
|
|
return a == b or not _seq_lt(a, b)
|