From 13f51260b2671c5e97f109a34dbf7c2864d93b14 Mon Sep 17 00:00:00 2001
From: John Jones <jmjatlanta@gmail.com>
Date: Thu, 23 Nov 2017 07:44:42 -0500
Subject: [PATCH] Minor fixes to yamux and identify protocols

---
 identify/identify.c           | 2 +-
 include/libp2p/net/protocol.h | 6 ++++++
 include/libp2p/yamux/yamux.h  | 4 ++++
 net/protocol.c                | 9 +++++++++
 yamux/yamux.c                 | 8 +++++++-
 5 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/identify/identify.c b/identify/identify.c
index f56a2ea..afe68fc 100644
--- a/identify/identify.c
+++ b/identify/identify.c
@@ -329,7 +329,7 @@ int libp2p_identify_shutdown(void* protocol_context) {
 struct Libp2pProtocolHandler* libp2p_identify_build_protocol_handler(struct Libp2pVector* handlers) {
 	struct Libp2pProtocolHandler* handler = libp2p_protocol_handler_new();
 	if (handler != NULL) {
-		handler->context = handler;
+		handler->context = NULL;
 		handler->CanHandle = libp2p_identify_can_handle;
 		handler->HandleMessage = libp2p_identify_handle_message;
 		handler->Shutdown = libp2p_identify_shutdown;
diff --git a/include/libp2p/net/protocol.h b/include/libp2p/net/protocol.h
index 801636f..ab1b222 100644
--- a/include/libp2p/net/protocol.h
+++ b/include/libp2p/net/protocol.h
@@ -42,6 +42,12 @@ struct Libp2pProtocolHandler {
  */
 struct Libp2pProtocolHandler* libp2p_protocol_handler_new();
 
+/***
+ * Release resources of a protocol handler
+ * @param handler the handler to free
+ */
+void libp2p_protocol_handler_free(struct Libp2pProtocolHandler* handler);
+
 /***
  * Handle an incoming message
  * @param message the incoming message
diff --git a/include/libp2p/yamux/yamux.h b/include/libp2p/yamux/yamux.h
index 3a191a1..f57c20f 100644
--- a/include/libp2p/yamux/yamux.h
+++ b/include/libp2p/yamux/yamux.h
@@ -13,6 +13,10 @@ static const int yamux_default_timeout = 10;
 static const char YAMUX_CONTEXT = 'Y';
 static const char YAMUX_CHANNEL_CONTEXT = 'C';
 
+struct YamuxProtocolContext {
+	struct Libp2pVector* protocol_handlers;
+};
+
 /***
  * Context struct for Yamux
  */
diff --git a/net/protocol.c b/net/protocol.c
index 84bb887..82fc6b8 100644
--- a/net/protocol.c
+++ b/net/protocol.c
@@ -41,6 +41,15 @@ struct Libp2pProtocolHandler* libp2p_protocol_handler_new() {
 	return h;
 }
 
+/***
+ * Release resources of a protocol handler
+ * @param handler the handler to free
+ */
+void libp2p_protocol_handler_free(struct Libp2pProtocolHandler* handler) {
+	if (handler != NULL)
+		free(handler);
+}
+
 /***
  * Handle an incoming message
  * @param message the incoming message
diff --git a/yamux/yamux.c b/yamux/yamux.c
index f6b5117..4e838bc 100644
--- a/yamux/yamux.c
+++ b/yamux/yamux.c
@@ -163,7 +163,13 @@ int yamux_shutdown(void* protocol_context) {
 struct Libp2pProtocolHandler* libp2p_yamux_build_protocol_handler(struct Libp2pVector* handlers) {
 	struct Libp2pProtocolHandler* handler = libp2p_protocol_handler_new();
 	if (handler != NULL) {
-		handler->context = handlers;
+		struct YamuxProtocolContext* ctx = (struct YamuxProtocolContext*) malloc(sizeof(struct YamuxProtocolContext));
+		if (ctx == NULL) {
+			libp2p_protocol_handler_free(handler);
+			return NULL;
+		}
+		ctx->protocol_handlers = handlers;
+		handler->context = ctx;
 		handler->CanHandle = yamux_can_handle;
 		handler->HandleMessage = yamux_handle_message;
 		handler->Shutdown = yamux_shutdown;