A Python port of the Invisible Internet Project (I2P)
1"""Stream connection state machine."""
2
3import enum
4from i2p_streaming.options import StreamOptions
5
6
7class ConnectionState(enum.Enum):
8 NEW = 0
9 SYN_SENT = 1
10 SYN_RECEIVED = 2
11 ESTABLISHED = 3
12 CLOSE_WAIT = 4
13 LAST_ACK = 5
14 FIN_WAIT = 6
15 CLOSING = 7
16 TIME_WAIT = 8
17 CLOSED = 9
18 RESET = 10
19
20
21class StreamConnection:
22 """TCP-like connection state machine for I2P streaming."""
23
24 def __init__(self, options: StreamOptions | None = None):
25 self._options = options or StreamOptions()
26 self.state = ConnectionState.NEW
27 self.next_send_seq = 0
28 self.next_recv_seq = 0
29 self.send_window = self._options.initial_window_size
30 self.recv_window = self._options.initial_window_size
31 self.destination: bytes = b""
32
33 def _require_state(self, *states: ConnectionState):
34 if self.state not in states:
35 raise RuntimeError(f"Invalid transition from {self.state}")
36
37 def send_syn(self):
38 self._require_state(ConnectionState.NEW)
39 self.state = ConnectionState.SYN_SENT
40
41 def receive_syn(self):
42 self._require_state(ConnectionState.NEW)
43 self.state = ConnectionState.SYN_RECEIVED
44
45 def receive_syn_ack(self):
46 self._require_state(ConnectionState.SYN_SENT)
47 self.state = ConnectionState.ESTABLISHED
48
49 def send_syn_ack(self):
50 self._require_state(ConnectionState.SYN_RECEIVED)
51 self.state = ConnectionState.ESTABLISHED
52
53 def send_close(self):
54 self._require_state(ConnectionState.ESTABLISHED, ConnectionState.CLOSE_WAIT)
55 if self.state == ConnectionState.ESTABLISHED:
56 self.state = ConnectionState.FIN_WAIT
57 elif self.state == ConnectionState.CLOSE_WAIT:
58 self.state = ConnectionState.LAST_ACK
59
60 def receive_close(self):
61 self._require_state(ConnectionState.ESTABLISHED, ConnectionState.FIN_WAIT)
62 if self.state == ConnectionState.ESTABLISHED:
63 self.state = ConnectionState.CLOSE_WAIT
64 elif self.state == ConnectionState.FIN_WAIT:
65 self.state = ConnectionState.CLOSING
66
67 def receive_ack(self):
68 self._require_state(ConnectionState.LAST_ACK)
69 self.state = ConnectionState.CLOSED
70
71 def receive_close_ack(self):
72 self._require_state(ConnectionState.FIN_WAIT)
73 self.state = ConnectionState.TIME_WAIT
74
75 def timeout(self):
76 self._require_state(ConnectionState.TIME_WAIT)
77 self.state = ConnectionState.CLOSED
78
79 def reset(self):
80 self.state = ConnectionState.RESET
81
82 def increment_send_seq(self):
83 self.next_send_seq += 1
84
85 def increment_recv_seq(self):
86 self.next_recv_seq += 1