You are almost correct. But stateful inspection allows you to have more control above data streams going through the firewall.
A Cybrary article says:
Stateless firewalls watch network traffic and restrict or block
packets based on source and destination addresses or other static
values. They’re not ‘aware’ of traffic patterns or data flows. A
stateless firewall uses simple rule-sets that do not account for the
possibility that a packet might be received by the firewall
‘pretending’ to be something you asked for.
A stateless firewall filter, also known as an access control list
(ACL), does not statefully inspect traffic. Instead, it evaluates
packet contents statically and does not keep track of the state of
network connections.
Stateful firewalls can watch traffic streams from end to end. They are
aware of communication paths and can implement various IP Security
(IPsec) functions such as tunnels and encryption. In technical terms,
this means that stateful firewalls can tell what stage a TCP
connection is in (open, open sent, synchronized, synchronization
acknowledge or established). It can tell if the MTU has changed and
whether packets have fragmented. etc.
They are named as such because they are really doing what the name means.
Stateless: don't remember any state, just filter the current packet passing through according to the rule list.
Stateful: remember the state of the data stream passed through (what happened in all previous packets) and filter based on the same list as in case of stateless, but filter also based on state information.
The state information stored in the firewall memory allows you to do exactly what you asked before and much more.
You can add a rule that allows passing of the new incoming TCP connection inside (TCP packet marked with a SYN flag) and all established / related traffic inside / outside.
In case of TCP, the client (SRC IP/SRC PORT) sends SYN packet to the server (DST IP/DST PORT) to open the TCP connection. Firewall has a rule that this packet can pass (only new TCP connections to that server). Firewall now remembers the request to open the connection was sent to server (it remembers among others src IP, src port, dst IP, dst port and state of the connection - in this case CONNECT) and now expects only SYN/ACK packet from server side. No other packets of this concrete TCP connection can pass through it (in both directions). Server then sends SYN/ACK, firewall remembers HALF-OPEN state and expects only ACK flagged packet from client. Once the client sends ACK, the state of the connection is changed to CONNECTED and only packets with flag ACK (data) FIN/ACK (close connection) can pass through now. The same is happening during the connection closing procedure - so it is possible to have complete control above the protocol passing through and filter out invalid packets.
For TCP it's a bit hard to explain and understand. ICMP is much easier. There are two basic ICMP packets - ECHO and REPLY.
In case of stateful firewall you can ensure that REPLY can pass the firewall only in case the ECHO was previously sent.
In case of stateless it is possible only to control if ECHO and REPLY can pass through so the server can send REPLY anytime even if it was not asked for it.