A Python port of the Invisible Internet Project (I2P)
at main 86 lines 2.7 kB view raw
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