server.ts

import { Socket } from 'node:net';
import { createServer, Server, IncomingMessage, ServerResponse } from 'node:http';
import { createServer as createHttpsServer } from 'node:https';
import { createBrotliCompress, createGzip, gzipSync, brotliCompressSync } from 'node:zlib';
import { createReadStream } from 'node:fs';
import busboy from 'busboy';
import iconv from 'iconv-lite';
import selfsigned from 'selfsigned';
import { readableToBytes, sleep } from '../utils';

const requestsPerSocket = Symbol('requestsPerSocket');

export async function startServer(options?: {
  keepAliveTimeout?: number;
  https?: boolean;
}): Promise<{ server: Server, url: string, closeServer: any }> {
  let server: Server;
  const requestHandler = async (req: IncomingMessage, res: ServerResponse) => {
    const startTime = Date.now();
    req.socket[requestsPerSocket] = (req.socket[requestsPerSocket] || 0) + 1;
    if (server.keepAliveTimeout) {
      res.setHeader('Keep-Alive', 'timeout=' + server.keepAliveTimeout / 1000);
    }
    const urlObject = new URL(req.url!, `http://${req.headers.host}`);
    const pathname = urlObject.pathname;
    res.setHeader('x-requests-persocket', req.socket[requestsPerSocket]);
    res.setHeader('x-requests-socket-port', req.socket.remotePort!);
    res.setHeader('X-Foo', 'bar');
    res.setHeader('x-href', urlObject.href);
    res.setHeader('x-method', req.method ?? '');
    res.setHeader('x-request-headers', JSON.stringify(req.headers));

    if (pathname === '/block') {
      return;
    }

    if (pathname === '/mock-status') {
      const statusText = urlObject.searchParams.get('status') ?? '200';
      const statusCode = parseInt(statusText);
      res.statusCode = statusCode;
      res.end(`Mock status ${statusCode}`);
      return;
    }

    const timeout = urlObject.searchParams.get('timeout');

    if (pathname === '/mock-bytes') {
      const size = urlObject.searchParams.get('size') ?? '1024';
      const bytes = Buffer.alloc(parseInt(size));
      if (timeout) {
        res.write(bytes);
        await sleep(parseInt(timeout));
        res.end();
      } else {
        res.end(bytes);
      }
      return;
    }

    if (timeout) {
      await sleep(parseInt(timeout));
    }

    if (pathname === '/auth') {
      const authorization = req.headers.authorization?.split(' ')[1] ?? '';
      const data = Buffer.from(authorization, 'base64');
      const auth = data.toString().split(':');
      return res.end(JSON.stringify({
        user: auth[0],
        password: auth[1],
      }));
    }

    if (pathname === '/hello/json') {
      return res.end(JSON.stringify({
        hello: 'urllib',
      }));
    }

    if (pathname === '/digestAuth') {
      const authorization = req.headers.authorization;
      if (!authorization) {
        res.setHeader('www-authenticate', 'Digest realm="testrealm@urllib.com", qop="auth,auth-int", nonce="dcd98b7102dd2f0e8b11d0f600bfb0c093", opaque="5ccc069c403ebaf9f0171e9517f40e41"');
        res.statusCode = 401;
        return res.end(JSON.stringify({
          error: 'authorization required',
        }));
      }
      if (!authorization.includes('Digest username="user"')) {
        res.setHeader('www-authenticate', 'Digest realm="testrealm@urllib.com", qop="auth,auth-int", nonce="dcd98b7102dd2f0e8b11d0f600bfb0c093", opaque="5ccc069c403ebaf9f0171e9517f40e41"');
        res.statusCode = 401;
        return res.end(JSON.stringify({
          error: 'authorization invaild',
        }));
      }
      return res.end(JSON.stringify({
        authorization,
      }));
    }

    if (pathname === '/digestAuth/multi') {
      const authorization = req.headers.authorization;
      if (!authorization) {
        res.setHeader('www-authenticate', [
          'Digest realm="testrealm@urllib.com", algorithm=MD5-sess, qop="auth", nonce="dcd98b7102dd2f0e8b11d0f600bfb0c093", opaque="5ccc069c403ebaf9f0171e9517f40e41"',
          'Digest realm="testrealm@urllib.com", qop="auth,auth-int", nonce="dcd98b7102dd2f0e8b11d0f600bfb0c093", opaque="5ccc069c403ebaf9f0171e9517f40e41"',
        ]);
        res.statusCode = 401;
        return res.end(JSON.stringify({
          error: 'authorization required',
        }));
      }
      if (!authorization.includes('Digest username="user"')) {
        res.setHeader('www-authenticate', 'Digest realm="testrealm@urllib.com", qop="auth,auth-int", nonce="dcd98b7102dd2f0e8b11d0f600bfb0c093", opaque="5ccc069c403ebaf9f0171e9517f40e41"');
        res.statusCode = 401;
        return res.end(JSON.stringify({
          error: 'authorization invaild',
        }));
      }
      return res.end(JSON.stringify({
        authorization,
      }));
    }

    if (pathname === '/wrongjson') {
      res.setHeader('content-type', 'application/json');
      return res.end(Buffer.from('{"foo":""'));
    }

    if (pathname === '/html') {
      res.setHeader('content-type', 'text/html');
      return res.end('<h1>hello</h1>');
    }

    if (pathname === '/redirect') {
      res.setHeader('Location', '/redirect-to-url');
      res.statusCode = 302;
      return res.end('Redirect to /redirect-to-url');
    }
    if (pathname === '/redirect-to-root') {
      res.setHeader('Location', '/');
      res.statusCode = 302;
      return res.end('Redirect to /');
    }
    if (pathname === '/redirect-deadlock') {
      res.setHeader('Location', '/redirect-deadlock');
      res.statusCode = 302;
      return res.end('Redirect to /redirect-deadlock');
    }
    if (pathname === '/redirect-301') {
      res.setHeader('Location', '/redirect-301-to-url');
      res.statusCode = 301;
      return res.end('Redirect to /redirect-301-to-url');
    }
    if (pathname === '/redirect-full') {
      const url = `http://${req.headers.host}/redirect-full-to-url`;
      res.setHeader('Location', url);
      res.statusCode = 302;
      return res.end(`Redirect to ${url}`);
    }
    if (pathname === '/redirect-full-301') {
      const url = `http://${req.headers.host}/redirect-full-301-to-url`;
      res.setHeader('Location', url);
      res.statusCode = 301;
      return res.end(`Redirect to ${url}`);
    }
    if (req.url === '/redirect-to-localhost') {
      res.statusCode = 302;
      res.setHeader('Location', 'http://localhost/');
      return res.end('Redirect to http://localhost/');
    }

    if (req.url === '/304-with-gzip') {
      res.setHeader('Content-Encoding', 'gzip');
      res.statusCode = 304;
      return res.end();
    }

    if (pathname === '/socket.end.error') {
      res.write('foo haha\n');
      await sleep(200);
      res.write('foo haha 2');
      await sleep(200);
      res.socket!.end('balabala');
      return;
    }

    if (pathname === '/wrongjson-gbk') {
      res.setHeader('content-type', 'application/json');
      createReadStream(__filename).pipe(res);
      return;
    }
    if (pathname === '/json_with_controls_unicode') {
      return res.end(Buffer.from('{"foo":"\b\f\n\r\tbar\u000e!1!\u0086!2!\u0000!3!\u001F!4!\u005C!5!end\u005C\\"}'));
    }
    if (pathname === '/json_with_t') {
      return res.end(Buffer.from('{"foo":"ba\tr\t\t"}'));
    }
    if (pathname === '/gbk/json') {
      res.setHeader('Content-Type', 'application/json;charset=gbk');
      const content = iconv.encode(JSON.stringify({ hello: '你好' }), 'gbk');
      return res.end(content);
    }
    if (pathname === '/gbk/text') {
      res.setHeader('Content-Type', 'text/plain;charset=gbk');
      const content = iconv.encode('你好', 'gbk');
      return res.end(content);
    }
    if (pathname === '/errorcharset') {
      res.setHeader('Content-Type', 'text/plain;charset=notfound');
      return res.end('你好');
    }

    if (pathname === '/brotli') {
      res.setHeader('Content-Encoding', 'br');
      createReadStream(__filename).pipe(createBrotliCompress()).pipe(res);
      return;
    }
    if (pathname === '/gzip') {
      res.setHeader('Content-Encoding', 'gzip');
      createReadStream(__filename).pipe(createGzip()).pipe(res);
      return;
    }
    if (pathname === '/error-gzip') {
      res.setHeader('Content-Encoding', 'gzip');
      createReadStream(__filename).pipe(res);
      return;
    }
    if (pathname === '/error-brotli') {
      res.setHeader('Content-Encoding', 'br');
      createReadStream(__filename).pipe(res);
      return;
    }
    if (pathname === '/error') {
      return res.destroy();
    }

    if (pathname === '/multipart' && (req.method === 'POST' || req.method === 'PUT')) {
      const bb = busboy({ headers: req.headers });
      const result = {
        method: req.method,
        url: req.url,
        href: urlObject.href,
        headers: req.headers,
        files: {},
        form: {},
      };
      bb.on('file', (name, file, info) => {
        const { filename, encoding, mimeType } = info;
        // console.log(`File [${name}]: info %j`, info);
        let size = 0;
        file.on('data', data => {
          // console.log(`File [${name}] got ${data.length} bytes`);
          size += data.length;
        }).on('close', () => {
          // console.log(`File [${name}] done`);
          result.files[name] = {
            filename,
            encoding,
            mimeType,
            size,
          };
        });
      });
      bb.on('field', (name, val) => {
        // console.log(`Field [${name}]: value length: %d, info: %j`, val.length, info);
        result.form[name] = val;
      });
      bb.on('close', () => {
        // console.log('Done parsing form!');
        res.writeHead(200, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify(result));
      });
      req.pipe(bb);
      return;
    }

    if (pathname === '/raw') {
      req.pipe(res);
      return;
    }

    let requestBody: any;
    let requestBytes: Buffer = Buffer.from('');
    if (req.method !== 'GET' && req.method !== 'HEAD') {
      requestBytes = await readableToBytes(req);
    }

    if (req.headers['content-type']?.startsWith('application/x-www-form-urlencoded')) {
      const searchParams = new URLSearchParams(requestBytes.toString());
      requestBody = {};
      for (const [ field, value ] of searchParams.entries()) {
        requestBody[field] = value;
      }
    } else if (req.headers['content-type']?.startsWith('application/json')) {
      requestBody = JSON.parse(requestBytes.toString());
    } else {
      requestBody = requestBytes.toString();
    }
    res.statusCode = 200;
    res.setHeader('Content-Type', 'application/json');
    res.setHeader('x-rt', `${Date.now() - startTime}`);
    const responseBody = JSON.stringify({
      method: req.method,
      url: req.url,
      href: urlObject.href,
      headers: req.headers,
      requestBody,
    });
    const contentEncoding = urlObject.searchParams.get('content-encoding');
    if (contentEncoding === 'gzip') {
      res.setHeader('content-encoding', contentEncoding);
      return res.end(gzipSync(responseBody));
    }
    if (contentEncoding === 'br') {
      res.setHeader('content-encoding', contentEncoding);
      return res.end(brotliCompressSync(responseBody));
    }
    res.end(responseBody);
  };

  if (options?.https) {
    const pems = selfsigned.generate();
    server = createHttpsServer({
      key: pems.private,
      cert: pems.cert,
    }, requestHandler);
  } else {
    server = createServer(requestHandler);
  }

  if (options?.keepAliveTimeout) {
    server.keepAliveTimeout = options.keepAliveTimeout;
  }

  // handle active connection on Node.js 16
  const hasCloseAllConnections = !!(server as any).closeAllConnections;
  const connections: Socket[] = [];
  if (!hasCloseAllConnections) {
    server.on('connection', connection => {
      connections.push(connection);
    });
  }
  const protocol = options?.https ? 'https' : 'http';

  return new Promise(resolve => {
    server.listen(0, () => {
      const address: any = server.address();
      resolve({
        url: `${protocol}://localhost:${address.port}/`,
        server,
        closeServer() {
          if (hasCloseAllConnections) {
            (server as any).closeAllConnections();
          } else {
            // console.log('Closing %d http connections', connections.length);
            for (const connection of connections) {
              connection.destroy();
            }
          }
          return new Promise(resolve => {
            server.close(resolve);
          });
        },
      });
    });
  });
}