ping_client and ping_server samples

TomerHeber · Feb 23, 2020 · 370001b · 370001b
1 parent 4023a80
commit 370001b
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 11 deletions.
diff --git a/samples/ping_pong_ssl/ping_client.cc b/samples/ping_pong_ssl/ping_client.cc
@@ -6,17 +6,27 @@
 #include <ael/event_loop.h>
 #include <ael/stream_buffer.h>
 
+#include <ael/openssl/ssl_stream_buffer_filter.h>
+
 #include "elapsed.h"
 
 using namespace std;
 using namespace ael;
 
 static Elapsed elapsed;
 
+#define TCP_CONNECT 2
+#define SSL_FIRST_CONNECT 1
+#define SSL_SECOND_CONNECT 0
+
 class PingClient : public StreamBufferHandler {
 public:
-    PingClient(shared_ptr<promise<void>> done_promise) : done_promise_(done_promise) {}
-    virtual ~PingClient() {}
+    PingClient(shared_ptr<promise<void>> done_promise) : done_promise_(done_promise), state_(TCP_CONNECT) {
+       ssl_ctx_ = SSL_CTX_new(SSLv23_client_method());
+    }
+    virtual ~PingClient() {
+        SSL_CTX_free(ssl_ctx_);
+    }
 
     void HandleData(std::shared_ptr<StreamBuffer> stream_buffer, const std::shared_ptr<const DataView> &data_view) override {
         string data;
@@ -29,9 +39,31 @@ class PingClient : public StreamBufferHandler {
     }
 
 	void HandleConnected(std::shared_ptr<StreamBuffer> stream_buffer) override { 
-        // After connected, write "ping" to server.
-        cout << elapsed << "connected, sending ping" << endl;
-        stream_buffer->Write(string("ping"));
+        SSL *ssl;
+        shared_ptr<SSLStreamBufferFilter> ssl_filter;
+
+        switch(state_) {
+        case TCP_CONNECT:
+            cout << elapsed << "connected TCP - upgrading to SSL" << endl;
+            state_ = SSL_FIRST_CONNECT;
+            ssl = SSL_new(ssl_ctx_);
+			ssl_filter = std::make_shared<SSLStreamBufferFilter>(stream_buffer, ssl);
+			stream_buffer->AddStreamBufferFilter(ssl_filter);    
+            break;        
+        case SSL_FIRST_CONNECT:        
+            cout << elapsed << "connected SSL - upgrading to SSL over SSL" << endl; 
+            state_ = SSL_SECOND_CONNECT;
+            ssl = SSL_new(ssl_ctx_);
+            ssl_filter = std::make_shared<SSLStreamBufferFilter>(stream_buffer, ssl);
+            stream_buffer->AddStreamBufferFilter(ssl_filter);    
+            break;           
+        case SSL_SECOND_CONNECT:
+            cout << elapsed << "connected SSL over SSL, sending ping" << endl;
+            stream_buffer->Write(string("ping"));
+            break;
+        default:
+            throw "unexpected case";
+        }    
     }
 
 	void HandleEOF(std::shared_ptr<StreamBuffer> stream_buffer) override {
@@ -41,6 +73,8 @@ class PingClient : public StreamBufferHandler {
     }
 private:
     shared_ptr<promise<void>> done_promise_;
+    SSL_CTX *ssl_ctx_;
+    int state_;
 };
 
 int main() 

diff --git a/samples/ping_pong_ssl/ping_server.cc b/samples/ping_pong_ssl/ping_server.cc
@@ -1,7 +1,7 @@
 #include <iostream>
 #include <chrono>
 #include <memory>
-#include <unordered_set>
+#include <unordered_map>
 #include <future>
 
 #include <ael/event_loop.h>
@@ -17,15 +17,32 @@ using namespace ael;
 
 static Elapsed elapsed;
 
+#define TCP_CONNECT 2
+#define SSL_FIRST_CONNECT 1
+#define SSL_SECOND_CONNECT 0
+
 class PingServer : public NewConnectionHandler, public StreamBufferHandler, public enable_shared_from_this<StreamBufferHandler> {
 public:
-    PingServer(shared_ptr<EventLoop> event_loop) : event_loop_(event_loop) {}
-    virtual ~PingServer() {}
+    PingServer(shared_ptr<EventLoop> event_loop) : event_loop_(event_loop) {
+        ssl_ctx_ = SSL_CTX_new(SSLv23_server_method());
+
+	    if (SSL_CTX_use_certificate_file(ssl_ctx_, "fake1.crt", SSL_FILETYPE_PEM) <= 0) {
+	    	throw "failed to load certificate";
+	    }
+
+	    if (SSL_CTX_use_PrivateKey_file(ssl_ctx_, "fake1.key", SSL_FILETYPE_PEM) <= 0 ) {
+	    	throw "failed to load key";
+	    }
+    }
+
+    virtual ~PingServer() {
+        SSL_CTX_free(ssl_ctx_);
+    }
 
     void HandleNewConnection(Handle handle) override {
         cout << elapsed << "new connection" << endl;
         auto stream_buffer = StreamBuffer::CreateForServer(shared_from_this(), handle);
-        streams_buffers_.insert(stream_buffer);
+        streams_buffers_[stream_buffer] = TCP_CONNECT;
         event_loop_->Attach(stream_buffer);
     }
 
@@ -38,15 +55,42 @@ class PingServer : public NewConnectionHandler, public StreamBufferHandler, publ
         }
     }
 
-	void HandleConnected(std::shared_ptr<StreamBuffer>) override { /* required for filters (check libael_openssl) */ }
+	void HandleConnected(std::shared_ptr<StreamBuffer> stream_buffer) override { 
+        int &state = streams_buffers_[stream_buffer];
+        SSL *ssl;
+        shared_ptr<SSLStreamBufferFilter> ssl_filter;
+
+        switch(state) {
+        case TCP_CONNECT:
+            cout << elapsed << "connected TCP - upgrading to SSL" << endl;
+            state = SSL_FIRST_CONNECT;
+            ssl = SSL_new(ssl_ctx_);
+			ssl_filter = std::make_shared<SSLStreamBufferFilter>(stream_buffer, ssl);
+			stream_buffer->AddStreamBufferFilter(ssl_filter);    
+            break;        
+        case SSL_FIRST_CONNECT:        
+            cout << elapsed << "connected SSL - upgrading to SSL over SSL" << endl; 
+            state = SSL_SECOND_CONNECT;
+            ssl = SSL_new(ssl_ctx_);
+            ssl_filter = std::make_shared<SSLStreamBufferFilter>(stream_buffer, ssl);
+            stream_buffer->AddStreamBufferFilter(ssl_filter);    
+            break;           
+        case SSL_SECOND_CONNECT:
+            cout << elapsed << "connected SSL over SSL" << endl;     
+            break;
+        default:
+            throw "unexpected case";
+        }
+    }
 
 	void HandleEOF(std::shared_ptr<StreamBuffer> stream_buffer) override {
         cout << elapsed << "connection closed" << endl;
         streams_buffers_.erase(stream_buffer);
     }
 private:
     shared_ptr<EventLoop> event_loop_;
-    unordered_set<std::shared_ptr<StreamBuffer>> streams_buffers_;
+    SSL_CTX *ssl_ctx_;
+    unordered_map<std::shared_ptr<StreamBuffer>, int> streams_buffers_;
 };
 
 int main()