Allow subdomains to be used to specify branch instead of query param

Automattic · Jan 25, 2016 · 48cac8b · 48cac8b
1 parent 1e9fd47
commit 48cac8b
Showing 1 changed file with 62 additions and 24 deletions.
diff --git a/lib/server.js b/lib/server.js
@@ -92,7 +92,7 @@ module.exports = function(config) {
 		});
 	}
 
-	function proxy(branchName, req, res) {
+	function proxyRequest(branchName, req, res) {
 		var proxy = proxies[branchName];
 		if(!proxy) return res.send('proxy stopped');
 		proxy.web(req, res, function(err) {
@@ -118,12 +118,38 @@ module.exports = function(config) {
 		return false;
 	}
 
-	app.use(function(req, res, next) {
-		var branchName = req.query.branch || req.session.branch || 'master';
-		if(branchName !== req.session.branch) {
-			debug('Using branch', branchName);
-			req.session.branch = branchName;
+	function escapeRegExp(text) {
+		return text.replace(/([.*+?^=!:${}()|\[\]\/\\])/g, "\\$1");
+	}
+
+	var subdomainRegexp = null;
+	if(process.env.HOST) {
+		subdomainRegexp = new RegExp("(.+?)" + escapeRegExp('.'+process.env.HOST));
+	}
+	function getBranchName(req) {
+		var branchName;
+		if(subdomainRegexp) {
+			var matches = req.hostname.match(subdomainRegexp);
+			if(matches && matches[1]) {
+				branchName = matches[1].replace('.', '/');
+			}
 		}
+		if(!branchName) {
+			branchName = req.query.branch || req.session.branch || 'master';
+			if(branchName !== req.session.branch) {
+				debug('Using branch', branchName);
+				req.session.branch = branchName;
+			}
+		}
+		return normalizeBranchName(branchName);
+	}
+
+	function normalizeBranchName(branchName) {
+		return branchName.toLowerCase();
+	}
+
+	app.use(function(req, res, next) {
+		var branchName = getBranchName(req);
 		if(!workers[branchName] && !proxies[branchName]) {
 			serveBranch(branchName, function(err) {
 				if(err) {
@@ -158,31 +184,43 @@ module.exports = function(config) {
 		checkUpdated(branchName, function(err) {
 			if(err) return next(err);
 			if(!res.headersSent) {
-				proxy(branchName, req, res);
+				proxyRequest(branchName, req, res);
 			}
 		});
 	});
 
 	var server = http.createServer(app);
+	var initMiddleware = require('express/lib/middleware/init').init(app);
+	var queryMiddleware = require('express/lib/middleware/query')(app.get('query parser fn'));
 	// Proxy WebSockets as well proxy.ws.bind( proxy )
 	server.on('upgrade', function(req, socket, head) {
-		sessionMiddleware(req, { headers: {} } , function(err) {
-			if(err) {
-				socket.end('Error opening session: '+(err.message || err.toString()));
-				return;
-			}
-			var branchName = req.session.branch || 'master';
-			var proxy = proxies[branchName];
-			if(!branchName || !proxies[branchName]) {
-				socket.end('Session not found');
-				return;
-			}
-			if(!proxy) {
-				socket.end('Proxy not found');
-				return;
-			}
-			proxy.ws(req, socket, head, function(err) {
-				debug(err);
+		function handleError(err) {
+			socket.end('Error opening session: '+(err.message || err.toString()));
+		}
+		var res = { headers: {}, setHeader: function(key, value) { this.headers[key] = value; } };
+		// add express.js req methods
+		initMiddleware(req, res, function(err) {
+			if(err) return handleError(err);
+			// parse query string
+			queryMiddleware(req, res, function(err) {
+				if(err) return handleError(err);
+				// load session from cookie header
+				sessionMiddleware(req, res, function(err) {
+					if(err) return handleError(err);
+					var branchName = getBranchName(req);
+					var proxy = proxies[branchName];
+					if(!branchName || !proxies[branchName]) {
+						socket.end('Session not found');
+						return;
+					}
+					if(!proxy) {
+						socket.end('Proxy not found');
+						return;
+					}
+					proxy.ws(req, socket, head, function(err) {
+						debug(err);
+					});
+				});
 			});
 		});
 	});