Set headers on stream proxy (#518)

Closes #401

source/as-stream.js

@@ -9,6 +9,8 @@ module.exports = options => {
 	const input = new PassThrough();
 	const output = new PassThrough();
 	const proxy = duplexer3(input, output);
+	const piped = new Set();
+	let isFinished = false;
 
 	options.gotRetry.retries = () => 0;
 
@@ -52,13 +54,31 @@ module.exports = options => {
 			proxy.emit('error', new ReadError(error, options));
 		});
 
-		response.pipe(output);
-
 		if (options.throwHttpErrors && statusCode !== 304 && (statusCode < 200 || statusCode > 299)) {
 			proxy.emit('error', new HTTPError(statusCode, response.statusMessage, response.headers, options), null, response);
 			return;
 		}
 
+		isFinished = true;
+
+		response.pipe(output);
+
+		for (const destination of piped) {
+			if (destination.headersSent) {
+				continue;
+			}
+
+			for (const [key, value] of Object.entries(response.headers)) {
+				// Got gives *uncompressed* data. Overriding `content-encoding` header would result in an error.
+				// It's not possible to decompress uncompressed data, is it?
+				if (key.toLowerCase() !== 'content-encoding') {
+					destination.setHeader(key, value);
+				}
+			}
+
+			destination.statusCode = response.statusCode;
+		}
+
 		proxy.emit('response', response);
 	});
 
@@ -69,5 +89,25 @@ module.exports = options => {
 		'downloadProgress'
 	].forEach(event => emitter.on(event, (...args) => proxy.emit(event, ...args)));
 
+	const pipe = proxy.pipe.bind(proxy);
+	const unpipe = proxy.unpipe.bind(proxy);
+	proxy.pipe = (destination, options) => {
+		if (isFinished) {
+			throw new Error('Failed to pipe. The response has been emitted already.');
+		}
+
+		const result = pipe(destination, options);
+
+		if (Reflect.has(destination, 'setHeader')) {
+			piped.add(destination);
+		}
+
+		return result;
+	};
+	proxy.unpipe = stream => {
+		piped.delete(stream);
+		return unpipe(stream);
+	};
+
 	return proxy;
 };

test/stream.js

@@ -12,6 +12,9 @@ test.before('setup', async () => {
 	s = await createServer();
 
 	s.on('/', (req, res) => {
+		res.writeHead(200, {
+			unicorn: 'rainbow'
+		});
 		res.end('ok');
 	});
 
@@ -111,6 +114,21 @@ test('piping works', async t => {
 	t.is(await getStream(got.stream(`${s.url}/`).on('error', () => {})), 'ok');
 });
 
+test('proxying headers works', async t => {
+	const server = await createServer();
+
+	server.on('/', (req, res) => {
+		got.stream(s.url).pipe(res);
+	});
+
+	await server.listen(server.port);
+
+	const {headers} = await got(server.url);
+	t.is(headers.unicorn, 'rainbow');
+
+	await server.close();
+});
+
 test.after('cleanup', async () => {
 	await s.close();
 });