server/wd-hub.js

const Promise = require('bluebird');
const path = require('path');
const fs = require('fs-extra');
const outputFile = Promise.promisify(fs.outputFile);
const router = require('express').Router();

const config = require('../lib/config');
const { logger } = require('../lib/log');
const rawBodyParser = require('../lib/raw-body-parser');
const proxyRequest = require('../lib/proxy-request');
const logUtils = require('../lib/selenium-log-utils');
const uid = require('../lib/unique-id');
const createSeleniumError = require ('../lib/create-selenium-error');
const getSessionIdFromReq = require('../lib/get-session-id-from-req');
const { isEndingSession } = require('../lib/identify-command-from-req');
const doesRequestTriggerScreenshot = require('../lib/does-request-trigger-screenshot');
const getLogEntryFromReq = require('../lib/get-log-entry-from-req');
const takeScreenshot = require('../lib/take-screenshot');


router.all('*', rawBodyParser, function(req, res) {
  const targetEndpoint = config.get('targetEndpoint');
  const sessionId = getSessionIdFromReq(req);
  const commandId = uid();
  const startTime = Date.now();

  // TODO: We may need to insert our own loggingPrefs into their capabilities when they `POST /wd/hub/session`
  // {"desiredCapabilities":{"browserName":"chrome","loggingPrefs":{"browser":"ALL","client":"ALL","driver":"ALL","performance":"ALL","server":"ALL"}}}

  //logger.log('debug', 'req', req);
  Promise.resolve()
    .then(() => {
      if(!targetEndpoint) {
        throw new Error('No target endpoint provided. You probably need to set the `GITLAB_TARGET_SELENIUM_REMOTE_URL` environment variable');
      }
    })
    .then(() => {
      // If the session is ending, grab the logs
      let fetchLogsMiddlewarePromise = Promise.resolve();
      if(isEndingSession(req).result) {
        fetchLogsMiddlewarePromise = logUtils.fetchAllLogs(targetEndpoint, sessionId)
          .then((logMap) => {
            const seleniumLogPath = path.join(config.get('logDir'), `./${sessionId}/selenium-logs.json`);
            return outputFile(seleniumLogPath, JSON.stringify(logMap, null, 2));
          });
      }

      return fetchLogsMiddlewarePromise;
    })
    .then(() => {
      // Forward on the request to the target
      return proxyRequest(targetEndpoint, req);
    })
    .then((proxyRes) => {
      const endTime = Date.now();
      // Add an entry for our custom GitLab log we use in CI views
      logUtils.updateLog(sessionId, 'gitlab', getLogEntryFromReq(req, commandId, startTime, endTime));

      let screenshotMiddlewarePromise = Promise.resolve();
      if(sessionId && doesRequestTriggerScreenshot(req)) {
        screenshotMiddlewarePromise = takeScreenshot(targetEndpoint, sessionId, commandId);
      }

      return screenshotMiddlewarePromise
        .then(() => proxyRes);
    })
    .tap((proxyRes) => {
      // Return the proxied response back to our originating user
      Object.keys(proxyRes.headers).forEach((headerName) => {
        const headerValue = proxyRes.headers[headerName];
        res.set(headerName, headerValue);
      });
      res.status(proxyRes.statusCode);
      //logger.log('debug', '->', req.method, req.url, proxyRes.headers, proxyRes.body);
      res.send(proxyRes.body);
    })
    .tap((proxyRes) => {
      // TODO: If someone uses `POST /wd/hub/session/:sessionId/log` add it to our own log
      // because the log buffer is cleared on any request (need to test this)
      if(sessionId && req.method === 'POST' && (/session\/[a-f\d-]+\/log\/?$/i).test(req.path)) {
        const type = req.body.type;
        const logValue = proxyRes.body.value;
        logUtils.updateLog(sessionId, type, logValue);
      }
    })
    .catch((err) => {
      logger.error('Error', err.message, err.stack);
      res.status(500).json(createSeleniumError(err));
    });
});

module.exports = router;