apim-tools/lib/middleware/content-rewriter.js

APIM Tools

/**
 * @file 内容改写中间件
 * @author sparklewhy@gmail.com
 */

'use strict';

/* eslint-disable prefer-rest-params */
const helper = require('../util');

/**
 * 初始化响应对象
 *
 * @param {Object} context 请求上下文对象
 * @param {Object} context.req 请求对象
 * @param {Object} context.res 响应对象
 * @param {Object} options 选项
 * @return {Object}
 */
function initResponse(context, options) {
    const res = context.res;
    const rawSetHeader = res.setHeader;
    const rawWriteHead = res.writeHead;
    const rawWrite = res.write;
    const rawEnd = res.end;

    const WritableStream = require('stream-buffers').WritableStreamBuffer;
    const tasks = options.tasks || [];
    const fakeRes = {

        /**
         * 注入处理
         *
         * @param {Buffer} data 原始数据
         * @param {string=} encoding 编码
         * @return {Buffer}
         */
        inject: function (data, encoding) {
            let content = data.toString(encoding);
            for (let i = 0, len = tasks.length; i < len; i++) {
                let item = tasks[i];
                if (this.checkTaskInject(item) && item.process) {
                    content = item.process(content, item.options, context);
                }
            }

            return new Buffer(content, encoding);
        },

        /**
         * 检查是否需要注入处理
         *
         * @param {Object} task 注入的任务
         * @return {boolean}
         */
        checkTaskInject: function (task) {
            if (this._enableInject) {
                return !task.when || task.when(context);
            }

            return false;
        },

        resetContentLenInfo() {
            let contentLenInfo = this.contentLenInfo;
            if (contentLenInfo) {
                rawSetHeader.call(this, contentLenInfo.name, contentLenInfo.value);
            }
        },

        /**
         * 检查是否需要注入处理
         *
         * @return {boolean}
         */
        checkInject: function () {
            if (this._enableInject != null) {
                return this._enableInject;
            }

            let result = false;
            if (this._contentTypeInfo && options.whenContentType) {
                result = !!options.whenContentType(this._contentTypeInfo, context);
            }
            else {
                result = !!tasks.some(function (item) {
                    return !item.when || item.when(context);
                });
            }

            if (!result) {
                // 不需要注入，恢复内容长度信息设置
                this.resetContentLenInfo();

                this.setHeader = rawSetHeader;
                this.writeHead = rawWriteHead;
                this.write = rawWrite;
                this.end = rawEnd;
            }

            this._enableInject = result;

            return result;
        },

        /**
         * @override
         */
        setHeader: function (name, value) {
            name = name.toLowerCase();
            if (name === 'content-length') {
                this.contentLenInfo = {
                    name,
                    value
                };
                return;
            }

            if (name === 'content-type') {
                this._contentTypeInfo = value;
            }

            return rawSetHeader.apply(this, arguments);
        },

        /**
         * @override
         */
        writeHead: function (status, statusMessage, headers) {
            let headerInfo = headers || statusMessage;
            if (headerInfo && typeof headerInfo !== 'string') {
                Object.keys(headerInfo).forEach(h => {
                    let value = headerInfo[h];
                    this.setHeader(h, value);
                });
            }

            return rawWriteHead.call(
                this, status,
                typeof statusMessage === 'string' ? statusMessage : undefined
            );
        },

        /**
         * @override
         */
        write: function (chunk, encoding) {
            if (this.checkInject()) {
                if (!this.injectBuffer) {
                    this.injectBuffer = new WritableStream();
                }

                return this.injectBuffer.write(chunk, encoding);
            }

            return rawWrite.apply(this, arguments);
        },

        /**
         * @override
         */
        end: function (data, encoding) {
            if (!this.checkInject()) {
                return rawEnd.apply(this, arguments);
            }

            if (data) {
                this.write(data, encoding);
            }

            this.write = rawWrite;
            if (this._hasBody && this.injectBuffer) {
                data = this.injectBuffer.getContents();
                this.injectBuffer = null;

                let contentEncoding = res.getHeader('content-encoding');
                let result = helper.unzip(data, contentEncoding);
                data = result.data;

                // 对于无法处理的编码不做注入处理
                if (!result.encoding) {
                    data = this.inject(data, encoding);
                }
                else {
                    this.resetContentLenInfo();
                }

                data = helper.zip(data, contentEncoding).data;
            }

            return rawEnd.call(this, data, encoding);
        }
    };

    return Object.assign(res, fakeRes);
}

/**
 * 注入中间件
 *
 * @param {Object} options 注入选项
 * @param {function(string):boolean=} options.whenContentType 判断内容类型是否允许注入处理
 * @param {Array.<{process: Function, when: Function, options: Object}>} options.tasks
 *        注入处理的任务定义
 * @return {Function}
 */
module.exports = exports = function (options) {
    return function (context, next) {
        initResponse(context, options);
        next();
    };
};