Websocket provider

1 files changed, 141 insertions, 11 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index e7358613..c2529942 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -5200,6 +5200,13 @@ __attribute__ ((format (printf, 1, 2)))
 
 	class DownloadInstance: public CoreRefCountObject<BNDownloadInstance, BNNewDownloadInstanceReference, BNFreeDownloadInstance>
 	{
+	public:
+		struct Response
+		{
+			uint16_t statusCode;
+			std::unordered_map<std::string, std::string> headers;
+		};
+
 	protected:
 		DownloadInstance(DownloadProvider* provider);
 		DownloadInstance(BNDownloadInstance* instance);
@@ -5208,25 +5215,52 @@ __attribute__ ((format (printf, 1, 2)))
 		static int PerformRequestCallback(void* ctxt, const char* url);
 		static int PerformCustomRequestCallback(void* ctxt, const char* method, const char* url, uint64_t headerCount, const char* const* headerKeys, const char* const* headerValues, BNDownloadInstanceResponse** response);
 		static void PerformFreeResponse(void* ctxt, BNDownloadInstanceResponse* response);
-
+		/*!
+		    Cleanup any resources created by the instance
+		 */
 		virtual void DestroyInstance();
-
-	public:
-		struct Response
-		{
-			uint16_t statusCode;
-			std::unordered_map<std::string, std::string> headers;
-		};
-
+		/*!
+		    Virtual method to synchronously perform a GET request to a url, overridden by a subclass
+		    \param url Full url to request
+		    \return Zero or greater on successful request and HTTP 200 status
+		 */
 		virtual int PerformRequest(const std::string& url) = 0;
+		/*!
+		    Virtual method to synchronously perform a request to a url, overridden by a subclass
+		    \param method Request method e.g. GET
+		    \param url Full url to request
+		    \param headers HTTP headers as keys/values
+		    \param response Structure into which the response status code and headers should be stored
+		    \return Zero or greater on success
+		 */
 		virtual int PerformCustomRequest(const std::string& method, const std::string& url, const std::unordered_map<std::string, std::string>& headers, Response& response) = 0;
 
-		int PerformRequest(const std::string& url, BNDownloadInstanceOutputCallbacks* callbacks);
-		int PerformCustomRequest(const std::string& method, const std::string& url, const std::unordered_map<std::string, std::string>& headers, Response& response, BNDownloadInstanceInputOutputCallbacks* callbacks);
 		int64_t ReadDataCallback(uint8_t* data, uint64_t len);
 		uint64_t WriteDataCallback(uint8_t* data, uint64_t len);
 		bool NotifyProgressCallback(uint64_t progress, uint64_t total);
 		void SetError(const std::string& error);
+
+	public:
+		/*!
+		    Send a GET request to a url, synchronously
+		    \param url Full url to request
+		    \param callbacks Structure with callback functions for output data
+		    \return Zero or greater on successful request and HTTP 200 status
+		 */
+		int PerformRequest(const std::string& url, BNDownloadInstanceOutputCallbacks* callbacks);
+		/*!
+		    Send a request to a url, synchronously
+		    \param method Request method e.g. GET
+		    \param url Full url to request
+		    \param headers HTTP headers as keys/values
+		    \param response Structure into which the response status code and headers are stored
+		    \param callbacks Structure with callback functions for input and output data
+		    \return Zero or greater on success
+		 */
+		int PerformCustomRequest(const std::string& method, const std::string& url, const std::unordered_map<std::string, std::string>& headers, Response& response, BNDownloadInstanceInputOutputCallbacks* callbacks);
+		/*!
+		    Retrieve the error from the last request sent by this instance
+		 */
 		std::string GetError() const;
 	};
 
@@ -5265,6 +5299,102 @@ __attribute__ ((format (printf, 1, 2)))
 		virtual Ref<DownloadInstance> CreateNewInstance() override;
 	};
 
+	// WebsocketProvider
+	class WebsocketProvider;
+
+	class WebsocketClient: public CoreRefCountObject<BNWebsocketClient, BNNewWebsocketClientReference, BNFreeWebsocketClient>
+	{
+	protected:
+		WebsocketClient(WebsocketProvider* provider);
+		WebsocketClient(BNWebsocketClient* instance);
+
+		static void DestroyClientCallback(void* ctxt);
+		static bool ConnectCallback(void* ctxt, const char* host, uint64_t headerCount, const char* const* headerKeys, const char* const* headerValues);
+		static bool WriteCallback(const uint8_t* data, uint64_t len, void* ctxt);
+		static bool DisconnectCallback(void* ctxt);
+		static void ErrorCallback(const char* msg, void* ctxt);
+		bool ReadData(uint8_t* data, uint64_t len);
+
+		/*!
+		    Cleanup any resources created by the client
+		 */
+		virtual void DestroyClient();
+		/*!
+		    Virtual method for performing the connection, overridden by a subclass.
+		    \param host Full url with scheme, domain, optionally port, and path
+		    \param headers HTTP header keys and values
+		    \return True if the connection has started, but not necessarily if it succeeded
+		 */
+		virtual bool Connect(const std::string& host, const std::unordered_map<std::string, std::string>& headers) = 0;
+	public:
+		/*!
+		    Connect to a given url, asynchronously. The connection will be run in a separate thread managed by the websocket provider.
+
+		    Callbacks will be called **on the thread of the connection**, so be sure to ExecuteOnMainThread any long-running
+		    or gui operations in the callbacks.
+
+		    If the connection succeeds, connectedCallback will be called. On normal termination, disconnectedCallback will be called.
+		    If the connection succeeds, but later fails, disconnectedCallback will not be called, and errorCallback will be called instead.
+		    If the connection fails, neither connectedCallback nor disconnectedCallback will be called, and errorCallback will be called instead.
+
+		    If connectedCallback or readCallback return false, the connection will be aborted.
+
+		    \param host Full url with scheme, domain, optionally port, and path
+		    \param headers HTTP header keys and values
+		    \param callbacks Structure with callbacks for various websocket events
+		    \return True if the connection has started, but not necessarily if it succeeded
+		 */
+		bool Connect(const std::string& host, const std::unordered_map<std::string, std::string>& headers, BNWebsocketClientOutputCallbacks* callbacks);
+
+		/*!
+		    Write some data to the websocket
+		    \param data Data to write
+		    \return True if successful
+		 */
+		virtual bool Write(const std::vector<uint8_t>& data) = 0;
+		/*!
+		    Disconnect the websocket
+		    \return True if successful
+		 */
+		virtual bool Disconnect() = 0;
+	};
+
+	class CoreWebsocketClient: public WebsocketClient
+	{
+	public:
+		CoreWebsocketClient(BNWebsocketClient* instance);
+		virtual ~CoreWebsocketClient() {};
+
+		virtual bool Connect(const std::string& host, const std::unordered_map<std::string, std::string>& headers) override;
+		virtual bool Write(const std::vector<uint8_t>& data) override;
+		virtual bool Disconnect() override;
+	};
+
+	class WebsocketProvider: public StaticCoreRefCountObject<BNWebsocketProvider>
+	{
+		std::string m_nameForRegister;
+
+	protected:
+		WebsocketProvider(const std::string& name);
+		WebsocketProvider(BNWebsocketProvider* provider);
+
+		static BNWebsocketClient* CreateClientCallback(void* ctxt);
+
+	public:
+		virtual Ref<WebsocketClient> CreateNewClient() = 0;
+
+		static std::vector<Ref<WebsocketProvider>> GetList();
+		static Ref<WebsocketProvider> GetByName(const std::string& name);
+		static void Register(WebsocketProvider* provider);
+	};
+
+	class CoreWebsocketProvider: public WebsocketProvider
+	{
+	public:
+		CoreWebsocketProvider(BNWebsocketProvider* provider);
+		virtual Ref<WebsocketClient> CreateNewClient() override;
+	};
+
 	// Scripting Provider
 	class ScriptingOutputListener
 	{