Make followRedirect option accept a function (#2306)

Co-authored-by: Sindre Sorhus <[email protected]>

Author: VanTigranyan

documentation/2-options.md

@@ -681,10 +681,12 @@ Only useful when the `cookieJar` option has been set.
 
 ### `followRedirect`
 
-**Type: `boolean`**\
+**Type: `boolean | (response: PlainResponse) => boolean`**\
 **Default: `true`**
 
-Defines if redirect responses should be followed automatically.
+Whether redirect responses should be followed automatically.
+
+Optionally, pass a function to dynamically decide based on the response object.
 
 #### **Note:**
 > - If a `303` is sent by the server in response to any request type (POST, DELETE, etc.), Got will request the resource pointed to in the location header via GET.\

source/core/index.ts

@@ -336,7 +336,7 @@ export default class Request extends Duplex implements RequestEvents<Request> {
 		void (async () => {
 			// Node.js parser is really weird.
 			// It emits post-request Parse Errors on the same instance as previous request. WTF.
-			// Therefore we need to check if it has been destroyed as well.
+			// Therefore, we need to check if it has been destroyed as well.
 			//
 			// Furthermore, Node.js 16 `response.destroy()` doesn't immediately destroy the socket,
 			// but makes the response unreadable. So we additionally need to check `response.readable`.
@@ -723,95 +723,98 @@ export default class Request extends Duplex implements RequestEvents<Request> {
 			return;
 		}
 
-		if (options.followRedirect && response.headers.location && redirectCodes.has(statusCode)) {
+		if (response.headers.location && redirectCodes.has(statusCode)) {
 			// We're being redirected, we don't care about the response.
 			// It'd be best to abort the request, but we can't because
 			// we would have to sacrifice the TCP connection. We don't want that.
-			response.resume();
-
-			this._cancelTimeouts();
-			this._unproxyEvents();
-
-			if (this.redirectUrls.length >= options.maxRedirects) {
-				this._beforeError(new MaxRedirectsError(this));
-				return;
-			}
+			const shouldFollow = typeof options.followRedirect === 'function' ? options.followRedirect(typedResponse) : options.followRedirect;
+			if (shouldFollow) {
+				response.resume();
 
-			this._request = undefined;
+				this._cancelTimeouts();
+				this._unproxyEvents();
 
-			const updatedOptions = new Options(undefined, undefined, this.options);
+				if (this.redirectUrls.length >= options.maxRedirects) {
+					this._beforeError(new MaxRedirectsError(this));
+					return;
+				}
 
-			const serverRequestedGet = statusCode === 303 && updatedOptions.method !== 'GET' && updatedOptions.method !== 'HEAD';
-			const canRewrite = statusCode !== 307 && statusCode !== 308;
-			const userRequestedGet = updatedOptions.methodRewriting && canRewrite;
+				this._request = undefined;
 
-			if (serverRequestedGet || userRequestedGet) {
-				updatedOptions.method = 'GET';
+				const updatedOptions = new Options(undefined, undefined, this.options);
 
-				updatedOptions.body = undefined;
-				updatedOptions.json = undefined;
-				updatedOptions.form = undefined;
+				const serverRequestedGet = statusCode === 303 && updatedOptions.method !== 'GET' && updatedOptions.method !== 'HEAD';
+				const canRewrite = statusCode !== 307 && statusCode !== 308;
+				const userRequestedGet = updatedOptions.methodRewriting && canRewrite;
 
-				delete updatedOptions.headers['content-length'];
-			}
+				if (serverRequestedGet || userRequestedGet) {
+					updatedOptions.method = 'GET';
 
-			try {
-				// We need this in order to support UTF-8
-				const redirectBuffer = Buffer.from(response.headers.location, 'binary').toString();
-				const redirectUrl = new URL(redirectBuffer, url);
+					updatedOptions.body = undefined;
+					updatedOptions.json = undefined;
+					updatedOptions.form = undefined;
 
-				if (!isUnixSocketURL(url as URL) && isUnixSocketURL(redirectUrl)) {
-					this._beforeError(new RequestError('Cannot redirect to UNIX socket', {}, this));
-					return;
+					delete updatedOptions.headers['content-length'];
 				}
 
-				// Redirecting to a different site, clear sensitive data.
-				if (redirectUrl.hostname !== (url as URL).hostname || redirectUrl.port !== (url as URL).port) {
-					if ('host' in updatedOptions.headers) {
-						delete updatedOptions.headers.host;
-					}
+				try {
+					// We need this in order to support UTF-8
+					const redirectBuffer = Buffer.from(response.headers.location, 'binary').toString();
+					const redirectUrl = new URL(redirectBuffer, url);
 
-					if ('cookie' in updatedOptions.headers) {
-						delete updatedOptions.headers.cookie;
+					if (!isUnixSocketURL(url as URL) && isUnixSocketURL(redirectUrl)) {
+						this._beforeError(new RequestError('Cannot redirect to UNIX socket', {}, this));
+						return;
 					}
 
-					if ('authorization' in updatedOptions.headers) {
-						delete updatedOptions.headers.authorization;
-					}
+					// Redirecting to a different site, clear sensitive data.
+					if (redirectUrl.hostname !== (url as URL).hostname || redirectUrl.port !== (url as URL).port) {
+						if ('host' in updatedOptions.headers) {
+							delete updatedOptions.headers.host;
+						}
+
+						if ('cookie' in updatedOptions.headers) {
+							delete updatedOptions.headers.cookie;
+						}
 
-					if (updatedOptions.username || updatedOptions.password) {
-						updatedOptions.username = '';
-						updatedOptions.password = '';
+						if ('authorization' in updatedOptions.headers) {
+							delete updatedOptions.headers.authorization;
+						}
+
+						if (updatedOptions.username || updatedOptions.password) {
+							updatedOptions.username = '';
+							updatedOptions.password = '';
+						}
+					} else {
+						redirectUrl.username = updatedOptions.username;
+						redirectUrl.password = updatedOptions.password;
 					}
-				} else {
-					redirectUrl.username = updatedOptions.username;
-					redirectUrl.password = updatedOptions.password;
-				}
 
-				this.redirectUrls.push(redirectUrl);
-				updatedOptions.prefixUrl = '';
-				updatedOptions.url = redirectUrl;
+					this.redirectUrls.push(redirectUrl);
+					updatedOptions.prefixUrl = '';
+					updatedOptions.url = redirectUrl;
 
-				for (const hook of updatedOptions.hooks.beforeRedirect) {
-					// eslint-disable-next-line no-await-in-loop
-					await hook(updatedOptions, typedResponse);
-				}
+					for (const hook of updatedOptions.hooks.beforeRedirect) {
+						// eslint-disable-next-line no-await-in-loop
+						await hook(updatedOptions, typedResponse);
+					}
 
-				this.emit('redirect', updatedOptions, typedResponse);
+					this.emit('redirect', updatedOptions, typedResponse);
 
-				this.options = updatedOptions;
+					this.options = updatedOptions;
+
+					await this._makeRequest();
+				} catch (error: any) {
+					this._beforeError(error);
+					return;
+				}
 
-				await this._makeRequest();
-			} catch (error: any) {
-				this._beforeError(error);
 				return;
 			}
-
-			return;
 		}
 
 		// `HTTPError`s always have `error.response.body` defined.
-		// Therefore we cannot retry if `options.throwHttpErrors` is false.
+		// Therefore, we cannot retry if `options.throwHttpErrors` is false.
 		// On the last retry, if `options.throwHttpErrors` is false, we would need to return the body,
 		// but that wouldn't be possible since the body would be already read in `error.response.body`.
 		if (options.isStream && options.throwHttpErrors && !isResponseOk(typedResponse)) {

source/core/options.ts

@@ -1756,19 +1756,21 @@ export default class Options {
 	}
 
 	/**
-	Defines if redirect responses should be followed automatically.
+	Whether redirect responses should be followed automatically.
+
+ 	Optionally, pass a function to dynamically decide based on the response object.
 
 	Note that if a `303` is sent by the server in response to any request type (`POST`, `DELETE`, etc.), Got will automatically request the resource pointed to in the location header via `GET`.
 	This is in accordance with [the spec](https://tools.ietf.org/html/rfc7231#section-6.4.4). You can optionally turn on this behavior also for other redirect codes - see `methodRewriting`.
 
 	@default true
 	*/
-	get followRedirect(): boolean {
+	get followRedirect(): boolean | ((response: PlainResponse) => boolean) {
 		return this._internals.followRedirect;
 	}
 
-	set followRedirect(value: boolean) {
-		assert.boolean(value);
+	set followRedirect(value: boolean | ((response: PlainResponse) => boolean)) {
+		assert.any([is.boolean, is.function_], value);
 
 		this._internals.followRedirect = value;
 	}

source/core/response.ts

@@ -114,7 +114,9 @@ export type Response<T = unknown> = {
 
 export const isResponseOk = (response: PlainResponse): boolean => {
 	const {statusCode} = response;
-	const limitStatusCode = response.request.options.followRedirect ? 299 : 399;
+	const {followRedirect} = response.request.options;
+	const shouldFollow = typeof followRedirect === 'function' ? followRedirect(response) : followRedirect;
+	const limitStatusCode = shouldFollow ? 299 : 399;
 
 	return (statusCode >= 200 && statusCode <= limitStatusCode) || statusCode === 304;
 };

test/redirects.ts

@@ -89,6 +89,39 @@ test('follows redirect', withServer, async (t, server, got) => {
 	t.deepEqual(redirectUrls.map(String), [`${server.url}/`]);
 });
 
+test('does not follow redirect when followRedirect is a function and returns false', withServer, async (t, server, got) => {
+	server.get('/', reachedHandler);
+	server.get('/finite', finiteHandler);
+
+	const {body, statusCode} = await got('finite', {followRedirect: () => false});
+	t.not(body, 'reached');
+	t.is(statusCode, 302);
+});
+
+test('follows redirect when followRedirect is a function and returns true', withServer, async (t, server, got) => {
+	server.get('/', reachedHandler);
+	server.get('/finite', finiteHandler);
+
+	const {body, redirectUrls} = await got('finite', {followRedirect: () => true});
+	t.is(body, 'reached');
+	t.deepEqual(redirectUrls.map(String), [`${server.url}/`]);
+});
+
+test('followRedirect gets plainResponse and does not follow', withServer, async (t, server, got) => {
+	server.get('/temporary', (_request, response) => {
+		response.writeHead(307, {
+			location: '/redirect',
+		});
+		response.end();
+	});
+
+	const {statusCode} = await got('temporary', {followRedirect(response) {
+		t.is(response.headers.location, '/redirect');
+		return false;
+	}});
+	t.is(statusCode, 307);
+});
+
 test('follows 307, 308 redirect', withServer, async (t, server, got) => {
 	server.get('/', reachedHandler);