From ece771a27548f0425ea7d703eb4f78f8b1b5802a Mon Sep 17 00:00:00 2001
From: Markus Sattler <help.markus+git@gmail.com>
Date: Thu, 14 Jan 2016 17:28:00 +0100
Subject: [PATCH] add support for setting:  - Access-Control-Allow-Origin (#25)
  - Sec-WebSocket-Protocol ()

add _server->close(); for ESP
---
 src/WebSocketsServer.cpp | 26 ++++++++++++++++++++++----
 src/WebSocketsServer.h   |  4 +++-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/src/WebSocketsServer.cpp b/src/WebSocketsServer.cpp
index 5ec4eca..1dfcdcd 100644
--- a/src/WebSocketsServer.cpp
+++ b/src/WebSocketsServer.cpp
@@ -25,8 +25,11 @@
 #include "WebSockets.h"
 #include "WebSocketsServer.h"
 
-WebSocketsServer::WebSocketsServer(uint16_t port) {
+WebSocketsServer::WebSocketsServer(uint16_t port, String origin, String protocol) {
     _port = port;
+    _origin = origin;
+    _protocol = protocol;
+
     _server = new WEBSOCKETS_NETWORK_SERVER_CLASS(port);
 
     _cbEvent = NULL;
@@ -37,7 +40,12 @@ WebSocketsServer::~WebSocketsServer() {
     // disconnect all clients
     disconnect();
 
+#if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266)
+    _server->close();
+#else
     // TODO how to close server?
+#endif
+
 }
 
 /**
@@ -540,11 +548,21 @@ void WebSocketsServer::handleHeader(WSclient_t * client) {
                     "Sec-WebSocket-Version: 13\r\n"
                     "Sec-WebSocket-Accept: ");
             client->tcp->write(sKey.c_str(), sKey.length());
-            client->tcp->write("\r\n");
+
+            if(_origin.length() > 0) {
+                String origin = "\r\nAccess-Control-Allow-Origin: ";
+                origin += _origin;
+                origin += "\r\n";
+                client->tcp->write(origin.c_str(), origin.length());
+            }
 
             if(client->cProtocol.length() > 0) {
-                // TODO add api to set Protocol of Server
-                client->tcp->write("Sec-WebSocket-Protocol: arduino\r\n");
+                String protocol = "\r\nSec-WebSocket-Protocol: ";
+                protocol += _protocol;
+                protocol += "\r\n";
+                client->tcp->write(protocol.c_str(), protocol.length());
+            } else {
+                client->tcp->write("\r\n");
             }
 
             // header end
diff --git a/src/WebSocketsServer.h b/src/WebSocketsServer.h
index e7c00e6..a4155da 100644
--- a/src/WebSocketsServer.h
+++ b/src/WebSocketsServer.h
@@ -38,7 +38,7 @@ public:
 
         typedef std::function<void (uint8_t num, WStype_t type, uint8_t * payload, size_t length)> WebSocketServerEvent;
 
-        WebSocketsServer(uint16_t port);
+        WebSocketsServer(uint16_t port, String origin = "", String protocol = "arduino");
         ~WebSocketsServer(void);
 
         void begin(void);
@@ -74,6 +74,8 @@ public:
 
 protected:
         uint16_t _port;
+        String _origin;
+        String _protocol;
 
         WEBSOCKETS_NETWORK_SERVER_CLASS * _server;