From fc50e2e7c1e3a9daa7336afb5ed90a972707361a Mon Sep 17 00:00:00 2001
From: yuuzheng <yuuzheng@linkedin.com>
Date: Sun, 2 Jan 2022 14:11:28 +0800
Subject: [PATCH] Update socket functions

---
 Sources/Heze/Socket/HezeSocketClient.swift  | 17 +----------------
 Sources/Heze/Socket/HezeSocketHandler.swift | 14 +++++++++++---
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/Sources/Heze/Socket/HezeSocketClient.swift b/Sources/Heze/Socket/HezeSocketClient.swift
index e1a92b2..a0d7a08 100644
--- a/Sources/Heze/Socket/HezeSocketClient.swift
+++ b/Sources/Heze/Socket/HezeSocketClient.swift
@@ -22,23 +22,8 @@ open class HezeSocketClient: HezeSocketHandler {
         return nil
     }
 
-    public func sendMessage(_ msg: String) {
-        socket?.sendStringMessage(string: msg, final: true) { }
-    }
-
-    public func sendBytes(_ bytes: [UInt8]) {
-        socket?.sendBinaryMessage(bytes: bytes, final: true) { }
-    }
-
-    public func close() {
+    open func close() {
         socket?.close()
     }
 
-    open override func connected() {
-
-    }
-
-    open override func closed() {
-
-    }
 }
diff --git a/Sources/Heze/Socket/HezeSocketHandler.swift b/Sources/Heze/Socket/HezeSocketHandler.swift
index 0930694..d828a17 100644
--- a/Sources/Heze/Socket/HezeSocketHandler.swift
+++ b/Sources/Heze/Socket/HezeSocketHandler.swift
@@ -60,7 +60,7 @@ open class HezeSocketHandler: HezeHandler, WebSocketSessionHandler {
                 return
             }
             if let bytes = self.handleBytes(bytes, opcodeType: opcodeType, final: final, socket: socket)?.toBytes() {
-                socket.sendBinaryMessage(bytes: bytes, final: final) {
+                self.sendBytes(bytes) {
                     self.handleSession(request: req, socket: socket)
                 }
             } else {
@@ -85,11 +85,19 @@ open class HezeSocketHandler: HezeHandler, WebSocketSessionHandler {
         return nil
     }
 
-    public func connected() {
+    open func sendMessage(_ msg: String, completion: @escaping () -> Void) {
+        sendBytes([UInt8](msg.utf8), completion: completion)
+    }
+
+    open func sendBytes(_ bytes: [UInt8], completion: @escaping () -> Void) {
+        socket?.sendBinaryMessage(bytes: bytes, final: true, completion: completion)
+    }
+
+    open func connected() {
 
     }
 
-    public func closed() {
+    open func closed() {
 
     }
 }