betterdocs / includes / Mcp / MCPRateLimiter.php

MCPRateLimiter.php in BetterDocs – AI Documentation, Knowledge Base, MCP Server, Docs, Wikis, FAQ & Chatbot 4.9.3, at includes/Mcp/MCPRateLimiter.php

293 lines 7.9 KB

Raw Download Zip

1	<?php
2	/**
3	* MCP rate limiter — a per-IP lockout on FAILED token authentication.
4	*
5	* @package BetterDocs
6	* @since 4.9.0
7	*/
8
9	namespace WPDeveloper\BetterDocs\Mcp;
10
11	if ( ! defined( 'ABSPATH' ) ) {
12	exit; // Exit if accessed directly.
13	}
14
15	/**
16	* The pairing token is a 256-bit secret, so online brute-forcing is already
17	* infeasible. This limiter stops the cheaper abuse — a flood of bad-token
18	* requests burning CPU and filling logs — and gives a rotated token's stale
19	* clients a hard wall. Defence in depth, not the primary control.
20	*
21	* Model: count consecutive FAILED attempts per client IP in a rolling window
22	* (transient-backed). At or past the threshold the IP is locked out for the rest
23	* of that window; a successful authentication clears the counter immediately.
24	*
25	* "Failed" means a credential was presented and rejected. A request with no
26	* `Authorization` header never reaches the counter — that is the first step of
27	* the OAuth handshake (the client is asking for the RFC 9728 challenge), so
28	* counting it would lock out every OAuth client during ordinary discovery.
29	*
30	* The window is fixed from the first failure: recording a failure never extends
31	* it, so a stranded client that retries every few minutes cannot re-arm its own
32	* lockout forever. Waiting out one window always works.
33	*
34	* Threshold and window come from the `BETTERDOCS_MCP_MAX_FAILS` /
35	* `BETTERDOCS_MCP_LOCKOUT_SECONDS` constants, then the
36	* `betterdocs_mcp_rate_limit` filter.
37	*
38	* @since 4.9.0
39	*/
40	final class MCPRateLimiter {
41
42	/**
43	* Transient key prefix; the client-IP hash is appended.
44	*
45	* @since 4.9.0
46	*/
47	const PREFIX = 'betterdocs_mcp_rl_';
48
49	/**
50	* Default: lock out after this many failed attempts.
51	*
52	* @since 4.9.0
53	*/
54	const DEFAULT_MAX_FAILS = 10;
55
56	/**
57	* Default: lockout / rolling-window length, in seconds.
58	*
59	* @since 4.9.0
60	*/
61	const DEFAULT_LOCKOUT = 900;
62
63	/**
64	* Is the current client locked out?
65	*
66	* Call this before comparing the token, so a locked IP never reaches the
67	* constant-time compare at all.
68	*
69	* @since 4.9.0
70	*
71	* @return bool
72	*/
73	public static function is_locked() {
74	$limits = self::limits();
75
76	return self::attempts() >= $limits[0];
77	}
78
79	/**
80	* Record a failed attempt for the current client.
81	*
82	* @since 4.9.0
83	*
84	* @return bool True when this failure crossed into a lockout.
85	*/
86	public static function record_failure() {
87	$limits = self::limits();
88	$max = $limits[0];
89	$window = $limits[1];
90
91	$key = self::key();
92	$entry = get_transient( $key );
93
94	if ( is_array( $entry ) && isset( $entry['count'], $entry['until'] ) ) {
95	++$entry['count'];
96
97	// Preserve the ORIGINAL window end: the TTL is the remaining time
98	// only, never a fresh full window.
99	$remaining = max( 1, (int) $entry['until'] - time() );
100
101	set_transient( $key, $entry, $remaining );
102
103	return $entry['count'] >= $max;
104	}
105
106	// First failure in a window. This also migrates a legacy integer entry:
107	// a stale int simply restarts as a fresh window of one.
108	$entry = [
109	'count' => 1,
110	'until' => time() + $window
111	];
112
113	set_transient( $key, $entry, $window );
114
115	return 1 >= $max;
116	}
117
118	/**
119	* Clear the counter for the current client — call on a successful auth.
120	*
121	* @since 4.9.0
122	*
123	* @return void
124	*/
125	public static function clear() {
126	delete_transient( self::key() );
127	}
128
129	/**
130	* Seconds until the current client's window ends.
131	*
132	* Falls back to the full window length when no entry exists — honest, now
133	* that the window is fixed.
134	*
135	* @since 4.9.0
136	*
137	* @return int
138	*/
139	public static function retry_after() {
140	$entry = get_transient( self::key() );
141
142	if ( is_array( $entry ) && isset( $entry['until'] ) ) {
143	return max( 1, (int) $entry['until'] - time() );
144	}
145
146	$limits = self::limits();
147
148	return $limits[1];
149	}
150
151	/**
152	* How many clients are currently locked out, across all IPs.
153	*
154	* A diagnostic for the self-test: a healthy loopback plus a locked-out
155	* remote client is exactly the state a stranded connector produces — a
156	* rotated-away token, still retrying — and it is otherwise invisible.
157	*
158	* Returns null under a persistent object cache: transients do not live in
159	* the options table there, so the count is unknowable and claiming zero
160	* would be a lie.
161	*
162	* @since 4.9.0
163	*
164	* @return int\|null Locked-out client count, or null when unknowable.
165	*/
166	public static function active_lockouts() {
167	if ( wp_using_ext_object_cache() ) {
168	return null;
169	}
170
171	global $wpdb;
172
173	if ( ! is_object( $wpdb ) \|\| ! method_exists( $wpdb, 'get_col' ) ) {
174	return null;
175	}
176
177	$limits = self::limits();
178
179	// phpcs:ignore WordPress.DB.DirectDatabaseQuery.DirectQuery, WordPress.DB.DirectDatabaseQuery.NoCaching -- read-only diagnostic scan over transient rows; no core API enumerates them.
180	$rows = $wpdb->get_col(
181	$wpdb->prepare(
182	"SELECT option_value FROM {$wpdb->options} WHERE option_name LIKE %s",
183	$wpdb->esc_like( '_transient_' . self::PREFIX ) . '%'
184	)
185	);
186
187	$locked = 0;
188
189	foreach ( (array) $rows as $row ) {
190	$entry = maybe_unserialize( $row );
191	$count = is_array( $entry ) && isset( $entry['count'] )
192	? (int) $entry['count']
193	: ( is_numeric( $entry ) ? (int) $entry : 0 );
194
195	if ( $count >= $limits[0] ) {
196	++$locked;
197	}
198	}
199
200	return $locked;
201	}
202
203	/**
204	* Current failed-attempt count for this client (0 when none).
205	*
206	* @since 4.9.0
207	*
208	* @return int
209	*/
210	private static function attempts() {
211	$value = get_transient( self::key() );
212
213	if ( is_array( $value ) && isset( $value['count'] ) ) {
214	return (int) $value['count'];
215	}
216
217	// Legacy integer entries, from before the fixed-window format.
218	return is_numeric( $value ) ? (int) $value : 0;
219	}
220
221	/**
222	* Transient key bound to the hashed client IP.
223	*
224	* @since 4.9.0
225	*
226	* @return string
227	*/
228	private static function key() {
229	return self::PREFIX . md5( self::client_ip() );
230	}
231
232	/**
233	* Resolve `[ max_fails, lockout_seconds ]` from the constants, then filter.
234	*
235	* @since 4.9.0
236	*
237	* @return array
238	*/
239	private static function limits() {
240	$max = defined( 'BETTERDOCS_MCP_MAX_FAILS' ) ? (int) BETTERDOCS_MCP_MAX_FAILS : self::DEFAULT_MAX_FAILS;
241	$window = defined( 'BETTERDOCS_MCP_LOCKOUT_SECONDS' ) ? (int) BETTERDOCS_MCP_LOCKOUT_SECONDS : self::DEFAULT_LOCKOUT;
242
243	/**
244	* Filter the MCP failed-auth rate limit.
245	*
246	* @since 4.9.0
247	*
248	* @param array $limits `[ max_fails, lockout_seconds ]`.
249	*/
250	$limits = (array) apply_filters( 'betterdocs_mcp_rate_limit', [ $max, $window ] );
251
252	$max = isset( $limits[0] ) ? max( 1, (int) $limits[0] ) : self::DEFAULT_MAX_FAILS;
253	$window = isset( $limits[1] ) ? max( 1, (int) $limits[1] ) : self::DEFAULT_LOCKOUT;
254
255	return [ $max, $window ];
256	}
257
258	/**
259	* Best-effort client IP.
260	*
261	* `REMOTE_ADDR` only — `X-Forwarded-For` is spoofable, so trusting it would
262	* let an attacker dodge the limit or lock out a victim. Behind a reverse
263	* proxy every client shares one `REMOTE_ADDR` and therefore one bucket; such
264	* a site should set `REMOTE_ADDR` upstream, or opt in through the filter
265	* below, which is safe only when the proxy is known to overwrite the
266	* forwarded header.
267	*
268	* @since 4.9.0
269	*
270	* @return string
271	*/
272	private static function client_ip() {
273	// phpcs:ignore WordPress.Security.ValidatedSanitizedInput -- used only as a rate-limit bucket key (md5'd), never output or stored raw.
274	$ip = isset( $_SERVER['REMOTE_ADDR'] ) ? (string) wp_unslash( $_SERVER['REMOTE_ADDR'] ) : '';
275
276	/**
277	* Filter the IP used as the MCP rate-limit bucket key.
278	*
279	* An opt-in escape hatch for sites behind a trusted reverse proxy, where
280	* `REMOTE_ADDR` is the proxy and every client would otherwise collapse
281	* into a single bucket. Only return a forwarded header's value when the
282	* proxy is known to overwrite it.
283	*
284	* @since 4.9.0
285	*
286	* @param string $ip Resolved `REMOTE_ADDR` ('' when unavailable).
287	*/
288	$ip = (string) apply_filters( 'betterdocs_mcp_client_ip', $ip );
289
290	return '' !== $ip ? $ip : 'unknown';
291	}
292	}
293