From 00843990c0d449ac7fa1139877642a76e366974a Mon Sep 17 00:00:00 2001
From: Ali Moazenzadeh <aliveravb@gmail.com>
Date: Sun, 16 Jul 2023 13:07:44 +0330
Subject: [PATCH] Implement custom request API (#29)

---
 Sources/NetShears/Models/RequestModel.swift | 11 +++---
 Sources/NetShears/NetShears.swift           | 41 +++++++++++++++++----
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/Sources/NetShears/Models/RequestModel.swift b/Sources/NetShears/Models/RequestModel.swift
index 4b15dc5..092f189 100644
--- a/Sources/NetShears/Models/RequestModel.swift
+++ b/Sources/NetShears/Models/RequestModel.swift
@@ -99,18 +99,19 @@ public final class NetShearsRequestModel: Codable {
          success: Bool,
          statusCode: Int,
          duration: Double?,
-         HPACKHeadersRequest: [String: String]?,
-         HPACKHeadersResponse: [String: String]?,
+         scheme: String,
+         requestHeaders: [String: String]?,
+         responseHeaders: [String: String]?,
          isFinished: Bool = true) {
         self.id = UUID().uuidString
         self.method = method
-        self.scheme = "gRPC"
+        self.scheme = scheme
         self.url = url
         self.host = host
         self.httpBody = requestObject
         self.code = statusCode
-        self.responseHeaders = HPACKHeadersResponse
-        self.headers = HPACKHeadersRequest ?? [:]
+        self.responseHeaders = responseHeaders
+        self.headers = requestHeaders ?? [:]
         self.dataResponse = responseObject
         self.date = Date()
         self.port = nil
diff --git a/Sources/NetShears/NetShears.swift b/Sources/NetShears/NetShears.swift
index 1f009f3..bc5d1df 100644
--- a/Sources/NetShears/NetShears.swift
+++ b/Sources/NetShears/NetShears.swift
@@ -91,6 +91,39 @@ public final class NetShears: NSObject {
         }
     }
 
+    public func addCustomRequest(url: String,
+                                 host: String,
+                                 method: String,
+                                 requestObject: Data?,
+                                 responseObject: Data?,
+                                 success: Bool,
+                                 statusCode: Int,
+                                 statusMessage: String?,
+                                 duration: Double?,
+                                 scheme: String,
+                                 requestHeaders: [String: String]?,
+                                 responseHeaders: [String: String]?) {
+        let request = NetShearsRequestModel(url: url,
+                                            host: host,
+                                            method: method,
+                                            requestObject: requestObject,
+                                            responseObject: responseObject,
+                                            success: success,
+                                            statusCode: statusCode,
+                                            duration: duration,
+                                            scheme: scheme,
+                                            requestHeaders: requestHeaders,
+                                            responseHeaders: responseHeaders,
+                                            isFinished: true)
+        if loggerEnable {
+            RequestStorage.shared.newRequestArrived(request)
+        }
+
+        if listenerEnable {
+            RequestBroadcast.shared.newRequestArrived(request)
+        }
+    }
+
     public func addGRPC(url: String,
                         host: String,
                         method: String,
@@ -102,12 +135,6 @@ public final class NetShears: NSObject {
                         duration: Double?,
                         HPACKHeadersRequest: [String: String]?,
                         HPACKHeadersResponse: [String: String]?){
-        let request = NetShearsRequestModel(url: url, host: host, method: method, requestObject: requestObject, responseObject: responseObject, success: success, statusCode: statusCode, duration: duration, HPACKHeadersRequest: HPACKHeadersRequest, HPACKHeadersResponse: HPACKHeadersResponse, isFinished: true)
-        if loggerEnable {
-            RequestStorage.shared.newRequestArrived(request)
-        }
-        if listenerEnable {
-            RequestBroadcast.shared.newRequestArrived(request)
-        }
+       addCustomRequest(url: url, host: host, method: method, requestObject: requestObject, responseObject: responseObject, success: success, statusCode: statusCode, statusMessage: statusMessage, duration: duration, scheme: "gRPC", requestHeaders: HPACKHeadersRequest, responseHeaders: HPACKHeadersResponse)
     }
 }