2 files changed, 292 insertions, 0 deletions
diff --git a/internal/net/server.go b/internal/net/server.go
index 9d99acc..05a98d1 100644
--- a/internal/net/server.go
+++ b/internal/net/server.go
@@ -168,6 +168,14 @@ type Server struct {
 	// stopping flips when Shutdown begins so listener goroutines treat
 	// their subsequent Accept/serve errors as expected, not fatal.
 	stopping atomic.Bool
+	// conns counts live sessions across every transport. handleSession
+	// admits against maxConnections so a pre-auth connection flood cannot
+	// hold an unbounded number of reader goroutines, buffers, and outboxes.
+	conns atomic.Int64
+	// idleTimeout returns the read timeout armed before each message. It is a
+	// field so tests can drive the idle-drop path without waiting out the
+	// production windows; nil means sessionIdleTimeout.
+	idleTimeout func(*Session) time.Duration
 }
 
 type Hub struct {
@@ -187,6 +195,64 @@ const (
 	closeFlushTimeout = 2 * time.Second
 )
 
+// Connection admission and idle-read limits. The listener admits at most
+// maxConnections live sessions; each read arms an idle deadline (short before
+// authentication, generous once the session has an account) that every
+// message resets, so a silent client cannot pin a reader goroutine forever
+// while legitimate long sessions keep working.
+const (
+	maxConnections   = 1000
+	loginReadTimeout = 2 * time.Minute
+	idleReadTimeout  = 30 * time.Minute
+)
+
+// admitConnection reserves a connection slot, returning false when the server
+// is at capacity. Every admitted connection must call releaseConnection.
+func (s *Server) admitConnection() bool {
+	if s.conns.Add(1) > maxConnections {
+		s.conns.Add(-1)
+		return false
+	}
+	return true
+}
+
+func (s *Server) releaseConnection() {
+	s.conns.Add(-1)
+}
+
+// sessionIdleTimeout returns the read timeout to arm for this session.
+func (s *Server) sessionIdleTimeout(sess *Session) time.Duration {
+	if s.idleTimeout != nil {
+		return s.idleTimeout(sess)
+	}
+	return sessionIdleTimeout(sess)
+}
+
+// sessionIdleTimeout returns the read timeout for a session: the account-name
+// and password prompts get a short window (pre-auth sockets are the ones an
+// anonymous flood targets), everything else gets a generous idle timeout.
+func sessionIdleTimeout(sess *Session) time.Duration {
+	switch sess.State() {
+	case StateAccountName, StatePassword:
+		return loginReadTimeout
+	default:
+		return idleReadTimeout
+	}
+}
+
+// armReadDeadline sets an idle read deadline on the session's transport when
+// it supports one. tcpConn and wsConn (both in this package) expose their
+// underlying connection, so telnet, telnet-TLS, and websocket are all covered.
+func armReadDeadline(sess *Session, d time.Duration) {
+	deadline := time.Now().Add(d)
+	switch c := sess.Conn.(type) {
+	case *tcpConn:
+		_ = c.conn.SetReadDeadline(deadline)
+	case *wsConn:
+		_ = c.conn.SetReadDeadline(deadline)
+	}
+}
+
 // StartWriter puts the session into queued mode: writes are enqueued and a
 // dedicated goroutine drains the queue to the network with per-message write
 // deadlines. The goroutine exits after Close() — draining whatever was
@@ -705,12 +771,24 @@ func (s *Server) handleSession(sess *Session, handler func(*Session, string)) {
 		sess.Close()
 	}()
 
+	// Admit before reading anything: an over-capacity connection (including
+	// one arriving through the websocket transport, whose handler shares this
+	// path) is told so and dropped without holding a read loop.
+	if !s.admitConnection() {
+		sess.Write("Server is at full capacity. Please try again later.\r\n")
+		return
+	}
+	defer s.releaseConnection()
+
 	sess.Write("\033[2J\033[H")
 
 	sess.Write(WelcomeBanner())
 	sess.Write("What's your account name? ")
 
 	for {
+		// Re-armed every iteration so any message from the client resets the
+		// idle countdown.
+		armReadDeadline(sess, s.sessionIdleTimeout(sess))
 		line, err := sess.Conn.ReadMessage()
 		if err != nil {
 			log.Printf("session read error: %v", err)
diff --git a/internal/net/server_test.go b/internal/net/server_test.go
index c91c701..da8084a 100644
--- a/internal/net/server_test.go
+++ b/internal/net/server_test.go
@@ -1,6 +1,11 @@
 package net
 
 import (
+	"bytes"
+	"errors"
+	"io"
+	"net"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -46,6 +51,215 @@ func (c *testConn) SetWriteDeadline(time.Time) error { return nil }
 func (c *testConn) Close() error       { return nil }
 func (c *testConn) SetEcho(bool) error { return nil }
 
+// fakeNetConn is a net.Conn stub for the session read loop: Read reports the
+// peer hung up so the read loop ends immediately, Write captures what the
+// server sent, and read deadlines are recorded so a test can assert the
+// transport was armed.
+type fakeNetConn struct {
+	mu       sync.Mutex
+	written  bytes.Buffer
+	deadline time.Time
+}
+
+func newFakeNetConn() *fakeNetConn { return &fakeNetConn{} }
+
+func (c *fakeNetConn) Read([]byte) (int, error) { return 0, io.EOF }
+
+func (c *fakeNetConn) Write(b []byte) (int, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.written.Write(b)
+}
+
+func (c *fakeNetConn) Close() error { return nil }
+
+func (c *fakeNetConn) LocalAddr() net.Addr  { return nil }
+func (c *fakeNetConn) RemoteAddr() net.Addr { return nil }
+
+func (c *fakeNetConn) SetDeadline(t time.Time) error { return c.SetReadDeadline(t) }
+
+func (c *fakeNetConn) SetReadDeadline(t time.Time) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.deadline = t
+	return nil
+}
+
+func (c *fakeNetConn) SetWriteDeadline(time.Time) error { return nil }
+
+func (c *fakeNetConn) output() string {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.written.String()
+}
+
+func (c *fakeNetConn) readDeadline() time.Time {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.deadline
+}
+
+// TestConnectionCapRefusalKeepsTheServerDrained: at maxConnections a new
+// connection is refused without holding a slot (the counter stays at the cap),
+// and the refused peer is told why. Once the admitted sessions end, the live
+// count is back to zero.
+func TestConnectionCapRefusalKeepsTheServerDrained(t *testing.T) {
+	s := &Server{hub: NewHub()}
+	s.handler = func(*Session, string) {}
+
+	for i := 0; i < maxConnections; i++ {
+		if !s.admitConnection() {
+			t.Fatalf("connection %d of %d refused below the cap", i+1, maxConnections)
+		}
+	}
+
+	refused := newFakeNetConn()
+	s.handleSession(&Session{Conn: newTCPConn(refused)}, s.handler)
+
+	if got := s.conns.Load(); got != maxConnections {
+		t.Fatalf("live connection count after a refusal = %d, want %d (the refused connection held a slot)", got, maxConnections)
+	}
+	if !strings.Contains(refused.output(), "full capacity") {
+		t.Fatalf("refused connection was not told the server is full: %q", refused.output())
+	}
+
+	for i := 0; i < maxConnections; i++ {
+		s.releaseConnection()
+	}
+	if got := s.conns.Load(); got != 0 {
+		t.Fatalf("live connection count after every session ended = %d, want 0", got)
+	}
+}
+
+// TestAdmittedSessionReleasesItsSlot: an admitted session whose client hangs
+// up returns its slot, so a client cannot leak capacity by disconnecting.
+func TestAdmittedSessionReleasesItsSlot(t *testing.T) {
+	s := &Server{hub: NewHub()}
+	s.handler = func(*Session, string) {}
+
+	s.handleSession(&Session{Conn: newTCPConn(newFakeNetConn())}, s.handler)
+
+	if got := s.conns.Load(); got != 0 {
+		t.Fatalf("live connection count after the session ended = %d, want 0", got)
+	}
+}
+
+// TestArmReadDeadlineArmsTheTransport: the idle window is applied to the real
+// transport (not just computed), and the window itself depends on the protocol
+// state: pre-auth prompts get the short one.
+func TestArmReadDeadlineArmsTheTransport(t *testing.T) {
+	rec := newFakeNetConn()
+	sess := &Session{Conn: newTCPConn(rec)}
+
+	before := time.Now()
+	armReadDeadline(sess, loginReadTimeout)
+	got := rec.readDeadline()
+	if got.Before(before.Add(loginReadTimeout)) || got.After(time.Now().Add(loginReadTimeout)) {
+		t.Fatalf("armed read deadline = %v, want about %v", got, before.Add(loginReadTimeout))
+	}
+
+	sess.SetState(StateAccountName)
+	if got := sessionIdleTimeout(sess); got != loginReadTimeout {
+		t.Fatalf("account-name window = %v, want %v", got, loginReadTimeout)
+	}
+	sess.SetState(StatePassword)
+	if got := sessionIdleTimeout(sess); got != loginReadTimeout {
+		t.Fatalf("password window = %v, want %v", got, loginReadTimeout)
+	}
+	sess.SetState(StateGame)
+	if got := sessionIdleTimeout(sess); got != idleReadTimeout {
+		t.Fatalf("in-game window = %v, want %v", got, idleReadTimeout)
+	}
+}
+
+// TestIdleSessionDroppedWithoutAffectingActiveOne: an idle socket is dropped
+// once its armed read deadline expires, while a client that keeps sending
+// messages survives well past that window because every read re-arms the
+// deadline.
+func TestIdleSessionDroppedWithoutAffectingActiveOne(t *testing.T) {
+	s := &Server{hub: NewHub()}
+	var mu sync.Mutex
+	served := 0
+	s.handler = func(*Session, string) {
+		mu.Lock()
+		served++
+		mu.Unlock()
+	}
+	const window = 200 * time.Millisecond
+	s.idleTimeout = func(*Session) time.Duration { return window }
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			sess := &Session{Conn: newTCPConn(conn)}
+			sess.SetState(StateAccountName)
+			s.hub.Add(sess)
+			go s.handleSession(sess, s.handler)
+		}
+	}()
+
+	idle, err := net.Dial("tcp", ln.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer idle.Close()
+	active, err := net.Dial("tcp", ln.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer active.Close()
+
+	// The active client talks for five windows' worth of messages: it can only
+	// be served to the end if each message re-armed its read deadline.
+	sent := 10
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for i := 0; i < sent; i++ {
+			if _, err := active.Write([]byte("look\r\n")); err != nil {
+				return
+			}
+			time.Sleep(window / 4)
+		}
+	}()
+
+	// The silent client is dropped: its socket is closed by the server as soon
+	// as the armed deadline expires.
+	if err := idle.SetReadDeadline(time.Now().Add(10 * window)); err != nil {
+		t.Fatal(err)
+	}
+	buf := make([]byte, 4096)
+	for {
+		n, err := idle.Read(buf)
+		if n > 0 {
+			continue
+		}
+		if err == nil {
+			continue
+		}
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		t.Fatalf("idle client was not dropped after its read deadline: %v", err)
+	}
+
+	<-done
+	mu.Lock()
+	got := served
+	mu.Unlock()
+	if got < sent-2 {
+		t.Fatalf("the active client was served %d of %d messages; its read deadline was not re-armed per message", got, sent)
+	}
+}
+
 func TestSessionWriteConcurrency(t *testing.T) {
 	conn := &testConn{}
 	sess := &Session{Conn: conn}