S061 · Networking I — TCP/IP, DNS, Sockets

#!/usr/bin/env python3
"""echo.py — minimal TCP echo server + client.
 
Run server:   python3 echo.py server
Run client:   python3 echo.py client hello world
Framed mode:  python3 echo.py framed-client hello world
"""
import socket
import struct
import sys
import threading
 
HOST, PORT = "127.0.0.1", 9999
 
 
# ---------- server ----------
def server() -> None:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # SO_REUSEADDR: reclaim the port even if a previous socket is in TIME_WAIT.
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((HOST, PORT))
    s.listen(5)
    print(f"[server] listening on {HOST}:{PORT}")
    while True:
        conn, addr = s.accept()
        print(f"[server] connected by {addr}")
        threading.Thread(target=handle, args=(conn,), daemon=True).start()
 
 
def handle(conn: socket.socket) -> None:
    """Naive handler — proves the byte-stream problem exists."""
    with conn:
        while True:
            data = conn.recv(1024)      # blocks until at least 1 byte arrives
            if not data:                 # empty = peer sent FIN
                print("[server] client hung up")
                return
            print(f"[server] recv {len(data):4d} bytes: {data!r}")
            conn.sendall(b"echo: " + data)
 
 
# ---------- naive client ----------
def client(msg: str) -> None:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, PORT))
        # Two sends back-to-back — watch what the server sees.
        s.sendall(msg.encode())
        s.sendall(b" [second-write]")
        reply = s.recv(4096)
        print(f"[client] got {reply!r}")
 
 
# ---------- framed client (the fix) ----------
def framed_client(msg: str) -> None:
    """Length-prefix each message so the receiver knows where it ends."""
    payload = msg.encode()
    header = struct.pack("!I", len(payload))   # 4-byte big-endian length
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, PORT))
        s.sendall(header + payload)
        # Read the length header first, then exactly that many bytes.
        raw_len = recv_exact(s, 4)
        (msg_len,) = struct.unpack("!I", raw_len)
        body = recv_exact(s, msg_len - len(b"echo: ") + len(b"echo: "))
        print(f"[framed-client] got {body!r}")
 
 
def recv_exact(s: socket.socket, n: int) -> bytes:
    """Read exactly n bytes or die trying. This is the pattern real protocols use."""
    buf = bytearray()
    while len(buf) < n:
        chunk = s.recv(n - len(buf))
        if not chunk:
            raise ConnectionError(f"peer closed after {len(buf)}/{n} bytes")
        buf.extend(chunk)
    return bytes(buf)
 
 
# ---------- entrypoint ----------
if __name__ == "__main__":
    mode = sys.argv[1] if len(sys.argv) > 1 else "server"
    if mode == "server":
        server()
    elif mode == "client":
        client(" ".join(sys.argv[2:]) or "hello")
    elif mode == "framed-client":
        framed_client(" ".join(sys.argv[2:]) or "hello")
    else:
        sys.exit(f"unknown mode: {mode}")

# terminal 1
python3 echo.py server
 
# terminal 2 — naive client (two back-to-back sends)
python3 echo.py client "hello world"
 
# what the server prints will often look like:
# [server] recv   25 bytes: b'hello world [second-write]'
# NOT two separate recvs — that's the byte stream merging your two sends.

# Linux: list all TCP sockets with process info
ss -tnp | grep 9999
 
# You'll see two rows:
# ESTAB  0  0  127.0.0.1:9999    127.0.0.1:54321   users:(("python3",pid=1234,fd=4))
# LISTEN 0  5  127.0.0.1:9999    0.0.0.0:*         users:(("python3",pid=1234,fd=3))
 
# And see the actual packets on the loopback interface:
sudo tcpdump -i lo -X port 9999
# You'll literally see SYN, SYN-ACK, ACK, PSH, FIN in the output.