@@ -17,7 +17,7 @@ HTTPConnection::HTTPConnection(ResourceResolver * resResolver):
 	_isKeepAlive = false;
 	_lastTransmissionTS = millis();
 	_shutdownTS = 0;
-	_websocket = nullptr;
+	_wsHandler = nullptr;
 }
 
 HTTPConnection::~HTTPConnection() {
@@ -125,6 +125,11 @@ void HTTPConnection::closeConnection() {
 		delete _httpHeaders;
 		_httpHeaders = NULL;
 	}
+
+	if (_wsHandler != nullptr) {
+		HTTPS_DLOG("[   ] Freeing WS Handler");
+		delete _wsHandler;
+	}
 }
 
 /**
@@ -417,48 +422,29 @@ void HTTPConnection::loop() {
 			{
 				HTTPS_DLOG("[   ] Resolving resource...");
 				ResolvedResource resolvedResource;
-				_resResolver->resolveNode(_httpMethod, _httpResource, resolvedResource);
+
+				// Check which kind of node we need (Websocket or regular)
+				bool websocketRequested = checkWebsocket();
+
+				_resResolver->resolveNode(_httpMethod, _httpResource, resolvedResource, websocketRequested ? WEBSOCKET : HANDLER_CALLBACK);
+
+				// Is there any match (may be the defaultNode, if it is configured)
 				if (resolvedResource.didMatch()) {
-					// Did the client request connection:keep-alive?
-					HTTPHeader * connectionHeader = _httpHeaders->get("Connection");
-					if (connectionHeader != NULL && std::string("keep-alive").compare(connectionHeader->_value)==0) {
-						HTTPS_DLOGHEX("[   ] Keep-Alive activated. fid=", _socket);
-						_isKeepAlive = true;
+					// Check for client's request to keep-alive if we have a handler function.
+					if (resolvedResource.getMatchingNode()->_nodeType == HANDLER_CALLBACK) {
+						// Did the client set connection:keep-alive?
+						HTTPHeader * connectionHeader = _httpHeaders->get("Connection");
+						if (connectionHeader != NULL && std::string("keep-alive").compare(connectionHeader->_value)==0) {
+							HTTPS_DLOGHEX("[   ] Keep-Alive activated. fid=", _socket);
+							_isKeepAlive = true;
+						} else {
+							HTTPS_DLOGHEX("[   ] Keep-Alive disabled. fid=", _socket);
+							_isKeepAlive = false;
+						}
 					} else {
-						HTTPS_DLOGHEX("[   ] Keep-Alive disabled. fid=", _socket);
 						_isKeepAlive = false;
 					}
 
-					// do we have a websocket connection?
-					if(checkWebsocket()) {
-						// Create response
-						HTTPResponse res = HTTPResponse(this);
-						// Add default headers to the response
-						auto allDefaultHeaders = _defaultHeaders->getAll();
-						for(std::vector<HTTPHeader*>::iterator header = allDefaultHeaders->begin(); header != allDefaultHeaders->end(); ++header) {
-							res.setHeader((*header)->_name, (*header)->_value);
-						}
-						// Set the response status
-						res.setStatusCode(101);
-						res.setStatusText("Switching Protocols");
-						res.setHeader("Upgrade", "websocket");
-						res.setHeader("Connection", "Upgrade");
-						res.setHeader("Sec-WebSocket-Accept", websocketKeyResponseHash(_httpHeaders->getValue("Sec-WebSocket-Key")));
-						res.print("");
-
-						// Callback for the actual resource
-						HTTPSCallbackFunction * resourceCallback = resolvedResource.getMatchingNode()->_callback;
-						// persistent request for websocket. To be destroyed when connection closes.
-						HTTPRequest *req  = new HTTPRequest(this, _httpHeaders, resolvedResource.getParams(), _httpResource, _httpMethod, "");
-						// bind function call to the actual resource
-						std::function<void()> next = std::function<void()>(std::bind(resourceCallback, req, nullptr));
-						_websocket = new Websocket(this);  // make websocket with this connection 
-						next();  // call callback
-						delete req;
-						_connectionState = STATE_WEBSOCKET;
-						break;
-					}
-
 					// Create request context
 					HTTPRequest req  = HTTPRequest(this, _httpHeaders, resolvedResource.getParams(), _httpResource, _httpMethod, resolvedResource.getMatchingNode()->_tag);
 					HTTPResponse res = HTTPResponse(this);
@@ -469,8 +455,15 @@ void HTTPConnection::loop() {
 						res.setHeader((*header)->_name, (*header)->_value);
 					}
 
-					// Callback for the actual resource
-					HTTPSCallbackFunction * resourceCallback = resolvedResource.getMatchingNode()->_callback;
+					// Find the request handler callback
+					HTTPSCallbackFunction * resourceCallback;
+					if (websocketRequested) {
+						// For the websocket, we use the handshake callback defined below
+						resourceCallback = &handleWebsocketHandshake;
+					} else {
+						// For resource nodes, we use the callback defined by the node itself
+						resourceCallback = ((ResourceNode*)resolvedResource.getMatchingNode())->_callback;
+					}
 
  					// Get the current middleware chain
 					auto vecMw = _resResolver->getMiddleware();
@@ -495,37 +488,45 @@ void HTTPConnection::loop() {
 						req.discardRequestBody();
 					}
 
-					// Handling the request is done
-					HTTPS_DLOG("[   ] Handler function done, request complete");
-
-					// Now we need to check if we can use keep-alive to reuse the SSL connection
-					// However, if the client did not set content-size or defined connection: close,
-					// we have no chance to do so.
-					if (!_isKeepAlive) {
-						// No KeepAlive -> We are done. Transition to next state.
-						if (!isClosed()) {
-							_connectionState = STATE_BODY_FINISHED;
-						}
+					// Finally, after the handshake is done, we create the WebsocketHandler and change the internal state.
+					if(websocketRequested) {
+						_wsHandler = ((WebsocketNode*)resolvedResource.getMatchingNode())->newHandler();
+						_wsHandler->initialize(this);  // make websocket with this connection 
+						_connectionState = STATE_WEBSOCKET;
 					} else {
-						if (res.isResponseBuffered()) {
-							// If the response could be buffered:
-							res.setHeader("Connection", "keep-alive");
-							res.finalize();
-							if (_clientState != CSTATE_CLOSED) {
-								// Refresh the timeout for the new request
-								refreshTimeout();
-								// Reset headers for the new connection
-								_httpHeaders->clearAll();
-								// Go back to initial state
-								_connectionState = STATE_INITIAL;
+						// Handling the request is done
+						HTTPS_DLOG("[   ] Handler function done, request complete");
+
+						// Now we need to check if we can use keep-alive to reuse the SSL connection
+						// However, if the client did not set content-size or defined connection: close,
+						// we have no chance to do so.
+						if (!_isKeepAlive) {
+							// No KeepAlive -> We are done. Transition to next state.
+							if (!isClosed()) {
+								_connectionState = STATE_BODY_FINISHED;
+							}
+						} else {
+							if (res.isResponseBuffered()) {
+								// If the response could be buffered:
+								res.setHeader("Connection", "keep-alive");
+								res.finalize();
+								if (_clientState != CSTATE_CLOSED) {
+									// Refresh the timeout for the new request
+									refreshTimeout();
+									// Reset headers for the new connection
+									_httpHeaders->clearAll();
+									// Go back to initial state
+									_connectionState = STATE_INITIAL;
+								}
+							}
+							// The response could not be buffered or the client has closed:
+							if (!isClosed() && _connectionState!=STATE_INITIAL) {
+								_connectionState = STATE_BODY_FINISHED;
 							}
-						}
-						// The response could not be buffered or the client has closed:
-						if (!isClosed() && _connectionState!=STATE_INITIAL) {
-							_connectionState = STATE_BODY_FINISHED;
 						}
 					}
 				} else {
+					// No match (no default route configured, nothing does match)
 					HTTPS_DLOG("[ERR] Could not find a matching resource. Server error.");
 					serverError();
 				}
@@ -542,14 +543,15 @@ void HTTPConnection::loop() {
 			refreshTimeout();  // don't timeout websocket connection
 			if(pendingBufferSize() > 0) {
 				HTTPS_DLOG("[   ] websocket handler");
-				if(_websocket->read() < 0) {
-					_websocket->close();
-					delete _websocket;
-					_httpHeaders->clearAll();
-					_connectionState = STATE_CLOSING;
-				}
+				_wsHandler->loop();
+			}
+			// If the handler has terminated the connection, clean up and close the socket too
+			if (_wsHandler->closed()) {
+				HTTPS_DLOG("[   ] WS Connection closed. Freeing WS Handler");
+				delete _wsHandler;
+				_wsHandler = nullptr;
+				_connectionState = STATE_CLOSING;
 			}
-			//readBuffer();
 			break;
 		default:;
 		}
@@ -579,7 +581,22 @@ bool HTTPConnection::checkWebsocket() {
 	   	return false;
 }
 
-std::string HTTPConnection::websocketKeyResponseHash(std::string key) {
+/**
+ * Handler function for the websocket handshake. Will be used by HTTPConnection if a websocket is detected
+ */
+void handleWebsocketHandshake(HTTPRequest * req, HTTPResponse * res) {
+	res->setStatusCode(101);
+	res->setStatusText("Switching Protocols");
+	res->setHeader("Upgrade", "websocket");
+	res->setHeader("Connection", "Upgrade");
+	res->setHeader("Sec-WebSocket-Accept", websocketKeyResponseHash(req->getHeader("Sec-WebSocket-Key")));
+	res->print("");
+}
+
+/**
+ * Function used to compute the value of the Sec-WebSocket-Accept during Websocket handshake
+ */
+std::string websocketKeyResponseHash(std::string key) {
 	std::string newKey = key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
 	uint8_t shaData[HTTPS_SHA1_LENGTH];
 	esp_sha(SHA1, (uint8_t*)newKey.data(), newKey.length(), shaData);
 
@@ -14,15 +14,20 @@
 #include "lwip/sockets.h"
 
 #include "HTTPSServerConstants.hpp"
+#include "ConnectionContext.hpp"
+
 #include "HTTPHeaders.hpp"
 #include "HTTPHeader.hpp"
+
 #include "ResourceResolver.hpp"
 #include "ResolvedResource.hpp"
+
 #include "ResourceNode.hpp"
 #include "HTTPRequest.hpp"
 #include "HTTPResponse.hpp"
-#include "ConnectionContext.hpp"
-#include "Websocket.hpp"
+
+#include "WebsocketHandler.hpp"
+#include "WebsocketNode.hpp"
 
 namespace httpsserver {
 
@@ -150,10 +155,14 @@ class HTTPConnection : private ConnectionContext {
 	bool _isKeepAlive;
 
 	//Websocket connection
-	Websocket * _websocket;
+	WebsocketHandler * _wsHandler;
 
 };
 
+void handleWebsocketHandshake(HTTPRequest * req, HTTPResponse * res);
+
+std::string websocketKeyResponseHash(std::string key);
+
 } /* namespace httpsserver */
 
 #endif /* SRC_HTTPCONNECTION_HPP_ */
@@ -0,0 +1,52 @@
+#include "HTTPNode.hpp"
+
+namespace httpsserver {
+
+  HTTPNode::HTTPNode(const std::string path, const HTTPNodeType nodeType, const std::string tag):
+    _path(std::move(path)),
+    _tag(std::move(tag)),
+    _nodeType(nodeType) {
+
+    // Count the parameters
+    _urlParamCount = 0;
+    size_t idx = 0;
+    while((idx = path.find("/*", idx)) != std::string::npos) {
+      _urlParamCount+=1;
+      // If we don't do this, the same index will be found again... and again... and again...
+      idx+=1;
+    };
+
+    // Check if there are parameters
+    if (_urlParamCount > 0) {
+      // If there are parameters, store their indices
+      _urlParamIdx = new size_t[_urlParamCount];
+      for(int i = 0; i < _urlParamCount; i++) {
+        _urlParamIdx[i] = path.find("/*", i==0 ? 0 : _urlParamIdx[i-1])+1;
+      }
+    } else {
+      _urlParamIdx = NULL;
+    }
+  }
+
+  HTTPNode::~HTTPNode() {
+    if (_urlParamIdx != NULL) {
+      delete[] _urlParamIdx;
+    }
+  }
+
+  bool HTTPNode::hasUrlParameter() {
+    return _urlParamCount > 0;
+  }
+
+  size_t HTTPNode::getParamIdx(uint8_t idx) {
+    if (idx<_urlParamCount) {
+      return _urlParamIdx[idx];
+    } else {
+      return -1;
+    }
+  }
+
+  uint8_t HTTPNode::getUrlParamCount() {
+    return _urlParamCount;
+  }
+}
@@ -0,0 +1,47 @@
+#ifndef SRC_HTTPNODE_HPP_
+#define SRC_HTTPNODE_HPP_
+
+#include <Arduino.h>
+#include <string>
+
+namespace httpsserver {
+
+enum HTTPNodeType {
+	/** Node with a handler callback function (class ResourceNode) */
+	HANDLER_CALLBACK,
+	/** Node with a websocket handler (class WebsocketNode) */
+	WEBSOCKET
+};
+
+class HTTPNode {
+public:
+	HTTPNode(const std::string path, const HTTPNodeType nodeType, const std::string tag = "");
+  virtual ~HTTPNode();
+
+	/**
+	 * The path under which this node will be available. Should start with a slash. Example:
+	 * "/myResource"
+	 */
+	const std::string _path;
+
+	/**
+	 * Stores a tag that can be used in middleware and handler functions to identify this
+	 * specific node, tag the node with a required permission, ...
+	 */
+	const std::string _tag;
+
+	/** Stores the type of the node (as we have not runtime type information by default) */
+	const HTTPNodeType _nodeType;
+
+	bool hasUrlParameter();
+	uint8_t getUrlParamCount();
+	size_t getParamIdx(uint8_t);
+
+private:
+	uint8_t _urlParamCount;
+	size_t * _urlParamIdx;
+};
+
+} // namespace httpserver
+
+#endif
@@ -18,11 +18,11 @@ bool ResolvedResource::didMatch() {
 	return _matchingNode != NULL;
 }
 
-ResourceNode * ResolvedResource::getMatchingNode() {
+HTTPNode * ResolvedResource::getMatchingNode() {
 	return _matchingNode;
 }
 
-void ResolvedResource::setMatchingNode(ResourceNode * node) {
+void ResolvedResource::setMatchingNode(HTTPNode * node) {
 	_matchingNode = node;
 }
 
 
@@ -11,14 +11,14 @@ class ResolvedResource {
 	ResolvedResource();
 	~ResolvedResource();
 
-	void setMatchingNode(ResourceNode * node);
-	ResourceNode * getMatchingNode();
+	void setMatchingNode(HTTPNode * node);
+	HTTPNode * getMatchingNode();
 	bool didMatch();
 	ResourceParameters * getParams();
 	void setParams(ResourceParameters * params);
 
 private:
-	ResourceNode * _matchingNode;
+	HTTPNode * _matchingNode;
 	ResourceParameters * _params;
 };
Original file line number	Diff line number	Diff line change
@@ -18,11 +18,11 @@ bool ResolvedResource::didMatch() {
18	18	return _matchingNode != NULL;
19	19	}
20	20
21		- ResourceNode * ResolvedResource::getMatchingNode() {
	21	+ HTTPNode * ResolvedResource::getMatchingNode() {
22	22	return _matchingNode;
23	23	}
24	24
25		- void ResolvedResource::setMatchingNode(ResourceNode * node) {
	25	+ void ResolvedResource::setMatchingNode(HTTPNode * node) {
26	26	_matchingNode = node;
27	27	}
28	28