1 files changed, 24 insertions, 6 deletions
diff --git a/include/sibs/DirectConnection.hpp b/include/sibs/DirectConnection.hpp
index 8e3865f..9be55f1 100644
--- a/include/sibs/DirectConnection.hpp
+++ b/include/sibs/DirectConnection.hpp
@@ -11,6 +11,8 @@
 #include "IpAddress.hpp"
 #include "../types.hpp"
 #include "../utils.hpp"
+#include "Socket.hpp"
+#include "Message.hpp"
 
 namespace sibs
 {
@@ -29,15 +31,19 @@ namespace sibs
     struct DirectConnectionPeer;
     
     using PubSubConnectCallback = std::function<void(std::shared_ptr<DirectConnectionPeer> peer, PubSubResult result, const std::string &resultStr)>;
-    using PubSubReceiveDataCallback = std::function<void(std::shared_ptr<DirectConnectionPeer> peer, const void *data, const usize size)>;
+    using PubSubReceiveDataCallback = std::function<void(std::shared_ptr<DirectConnectionPeer> peer, MessageType messageType, const void *data, const usize size)>;
     using PubSubSendDataCallback = std::function<void(PubSubResult result, const std::string &resultStr)>;
     using PubSubOnRemoveDisconnectedPeerCallback = std::function<void(std::shared_ptr<DirectConnectionPeer> peer)>;
     
     struct DirectConnectionPeer
     {
-        int socket;
+        std::unique_ptr<Socket> socket;
         Ipv4 address;
         PubSubReceiveDataCallback receiveDataCallbackFunc;
+        int sharedKeys = 0;
+
+        bool operator == (const DirectConnectionPeer &other) const;
+        bool operator != (const DirectConnectionPeer &other) const;
     };
     
     class DirectConnections
@@ -52,18 +58,22 @@ namespace sibs
         void connectServer(const Ipv4 &address, PubSubConnectCallback connectCallbackFunc, PubSubReceiveDataCallback receiveDataCallbackFunc);
         // Throws ConnectionException on error
         void connect(const Ipv4 &address, PubSubConnectCallback connectCallbackFunc, PubSubReceiveDataCallback receiveDataCallbackFunc);
-
-        void send(const std::shared_ptr<DirectConnectionPeer> peer, std::shared_ptr<std::vector<u8>> data, PubSubSendDataCallback sendDataCallbackFunc = nullptr);
+        // Returns false if data is larger than 800kb
+        bool send(const std::shared_ptr<DirectConnectionPeer> peer, std::shared_ptr<Message> data, PubSubSendDataCallback sendDataCallbackFunc = nullptr);
 
         void onRemoveDisconnectedPeer(PubSubOnRemoveDisconnectedPeerCallback callbackFunc);
+        bool removePeer(int peerSocket);
+
+        std::vector<std::shared_ptr<DirectConnectionPeer>> getPeers();
+
+        std::shared_ptr<DirectConnectionPeer> getPeerByAddress(const Ipv4 &address) const;
     protected:
-        int createSocket(const Ipv4 &addressToBind, bool rendezvous, bool reuseAddr, bool bind = true);
+        std::unique_ptr<Socket> createSocket(const Ipv4 &addressToBind, bool rendezvous, bool reuseAddr, bool bind = true);
     private:
         void connect(const Ipv4 &address, bool rendezvous, bool reuseAddr, PubSubConnectCallback connectCallbackFunc, PubSubReceiveDataCallback receiveDataCallbackFunc, bool bind);
         void removeDisconnectedPeers();
         void receiveData();
         int receiveDataFromPeer(const int socket, char *output, usize *receivedTotalSize);
-        bool removePeer(int peerSocket);
     private:
         u16 port;
         int eid;
@@ -72,5 +82,13 @@ namespace sibs
         std::mutex peersMutex;
         bool alive;
         PubSubOnRemoveDisconnectedPeerCallback removeDisconnectedPeerCallback;
+        Ipv4Map<std::shared_ptr<DirectConnectionPeer>> peerByAddressMap;
+    };
+
+    struct DirectConnectionsUtils
+    {
+        static std::vector<u8> serializePeers(const std::vector<std::shared_ptr<DirectConnectionPeer>> &peers);
+        // Throws DeserializeException on error
+        static std::vector<std::shared_ptr<DirectConnectionPeer>> deserializePeers(const u8 *data, const usize size);
     };
 }