support for proxy #2

Sreyas-Sreelal · Apr 15, 2020 · 2c65797 · 2c65797
1 parent ca6a664
commit 2c65797
Show file tree

Hide file tree

Showing 5 changed files with 73 additions and 20 deletions.
diff --git a/include/tgconnector.inc b/include/tgconnector.inc
@@ -30,8 +30,8 @@ enum TGUserStatus {
 };
 
 //basic
-native TGBot:TGConnect(const token[],thread_limit=3);
-native TGBot:TGConnectFromEnv(const variable[],thread_limit=3);
+native TGBot:TGConnect(const token[],thread_limit=3,proxy[]="");
+native TGBot:TGConnectFromEnv(const variable[],thread_limit=3,proxy[]="");
 native TGSendMessage(TGBot:bot,const TGChatId:chatid[],const text[],TGMessage:reply_id=INVALID_MESSAGE_ID,TGParseMode:parse_mode=TGParseMode:-1,callback[]="");
 native TGDeleteMessage(TGBot:bot,TGChatId:chatid[],TGMessage:messageid);
 native TGEditMessage(TGBot:bot,TGChatId:chatid[],TGMessage:messageid,text[],TGParseMode:parse_mode=TGParseMode:-1);

diff --git a/src/api.rs b/src/api.rs
@@ -17,10 +17,11 @@ pub struct BOT {
     pub send_message_reciever: Option<Receiver<(Message, String)>>,
     pub send_message_sender: Option<Sender<(Message, String)>>,
     pub pool: ThreadPool,
+    pub proxy_url: Option<String>,
 }
 
 impl BOT {
-    pub fn new(bot_token: String, thread_count: i32) -> Self {
+    pub fn new(bot_token: String, thread_count: i32, proxy_url: Option<String>) -> Self {
         let (update_sender, update_reciever) = channel();
         let (send_message_sender, send_message_reciever) = channel();
 
@@ -32,14 +33,16 @@ impl BOT {
             send_message_reciever: Some(send_message_reciever),
             send_message_sender: Some(send_message_sender),
             pool: ThreadPool::new(thread_count as usize),
+            proxy_url,
         }
     }
 
-    pub fn connect(&mut self) -> bool {
+    pub fn connect(&mut self, proxy_url: Option<String>) -> bool {
         let request = HttpRequest {
             url: format!("{}/getme", self.api_request_link),
             method: HttpMethod::Get,
             body: None,
+            proxy_url,
         };
 
         match request.make_request() {
@@ -65,11 +68,12 @@ impl BOT {
     fn get_updates(&self) {
         let update_move = self.update_sender.clone();
         let api_link = self.api_request_link.clone();
+        let proxy_url = self.proxy_url.clone();
         let mut getupdate = GetUpdates { offset: -2 };
 
         self.pool.execute(move || loop {
             let update: Result<APIResponse<VecDeque<Update>>, String> =
-                telegram_request("getUpdates", &api_link, &getupdate);
+                telegram_request("getUpdates", &api_link, &getupdate, &proxy_url);
             match update {
                 Ok(update) => {
                     let mut check_result: VecDeque<Update> = match update.body {
@@ -104,10 +108,11 @@ impl BOT {
     pub fn send_message(&self, send_message_obj: SendMessage, callback: Option<String>) {
         let send_message_move = self.send_message_sender.clone();
         let api_link = self.api_request_link.clone();
+        let proxy_url = self.proxy_url.clone();
 
         self.pool.execute(move || {
             let response: Result<APIResponse<Message>, String> =
-                telegram_request("sendmessage", &api_link, &send_message_obj);
+                telegram_request("sendmessage", &api_link, &send_message_obj, &proxy_url);
             match response {
                 Ok(response) => {
                     if !response.ok {
@@ -128,10 +133,11 @@ impl BOT {
 
     pub fn delete_message(&self, delete_message_obj: DeleteMessage) {
         let api_link = self.api_request_link.clone();
+        let proxy_url = self.proxy_url.clone();
 
         self.pool.execute(move || {
             let response: Result<APIResponse<bool>, String> =
-                telegram_request("deletemessage", &api_link, &delete_message_obj);
+                telegram_request("deletemessage", &api_link, &delete_message_obj, &proxy_url);
 
             match response {
                 Ok(response) => {
@@ -152,10 +158,11 @@ impl BOT {
 
     pub fn edit_message(&self, edit_message_obj: EditMessageText) {
         let api_link = self.api_request_link.clone();
+        let proxy_url = self.proxy_url.clone();
 
         self.pool.execute(move || {
             let response: Result<APIResponse<Message>, String> =
-                telegram_request("editmessagetext", &api_link, &edit_message_obj);
+                telegram_request("editmessagetext", &api_link, &edit_message_obj, &proxy_url);
             match response {
                 Ok(response) => {
                     if !response.ok {
@@ -174,8 +181,12 @@ impl BOT {
     }
 
     pub fn get_chat_member(&self, getchatmember: GetChatMember) -> Option<ChatMember> {
-        let response: Result<APIResponse<ChatMember>, String> =
-            telegram_request("getchatmember", &self.api_request_link, &getchatmember);
+        let response: Result<APIResponse<ChatMember>, String> = telegram_request(
+            "getchatmember",
+            &self.api_request_link,
+            &getchatmember,
+            &self.proxy_url,
+        );
 
         match response {
             Ok(response) => {
@@ -199,6 +210,7 @@ impl BOT {
             "getchatmemberscount",
             &self.api_request_link,
             &getchatmemberscount,
+            &self.proxy_url,
         );
 
         match response {
@@ -220,7 +232,7 @@ impl BOT {
 
     pub fn get_chat(&self, getchat: GetChat) -> Option<Chat> {
         let response: Result<APIResponse<Chat>, String> =
-            telegram_request("getchat", &self.api_request_link, &getchat);
+            telegram_request("getchat", &self.api_request_link, &getchat, &self.proxy_url);
 
         match response {
             Ok(response) => {
@@ -244,11 +256,13 @@ fn telegram_request<T: DeserializeOwned, B: Serialize>(
     endpoint: &str,
     api_link: &str,
     body: B,
+    proxy_url: &Option<String>,
 ) -> Result<APIResponse<T>, String> {
     let request = HttpRequest {
         url: format!("{}/{}", api_link, endpoint),
         method: HttpMethod::Post,
         body: Some(to_string(&body).unwrap()),
+        proxy_url: proxy_url.clone(),
     };
 
     match request.make_request() {

diff --git a/src/http.rs b/src/http.rs
@@ -10,6 +10,7 @@ pub struct HttpRequest {
     pub url: String,
     pub method: HttpMethod,
     pub body: Option<String>,
+    pub proxy_url: Option<String>,
 }
 
 impl HttpRequest {
@@ -22,18 +23,37 @@ impl HttpRequest {
             }
         };
         let method = match self.method {
-            HttpMethod::Get => requests_obj.get(),
+            HttpMethod::Get => {
+                if let Some(proxy_url) = &self.proxy_url {
+                    match requests_obj.proxy(&proxy_url) {
+                        Ok(method) => method.get(),
+                        Err(err) => {
+                            return Err(format!("Error connecting to proxy server \n{:?}", err));
+                        }
+                    }
+                } else {
+                    requests_obj.get()
+                }
+            }
 
             HttpMethod::Post => {
                 let body = &self.body.clone().unwrap();
                 requests_obj.body_str(&body);
                 let mut headers = HashMap::new();
                 headers.insert("Content-Type".to_string(), "application/json".to_string());
                 requests_obj.headers(headers);
-                requests_obj.post()
+                if let Some(proxy_url) = &self.proxy_url {
+                    match requests_obj.proxy(&proxy_url) {
+                        Ok(method) => method.post(),
+                        Err(err) => {
+                            return Err(format!("Error connecting to proxy server \n{:?}", err));
+                        }
+                    }
+                } else {
+                    requests_obj.post()
+                }
             }
         };
-
         match method.send() {
             Ok(data) => Ok(data.text()),
 

diff --git a/src/internals.rs b/src/internals.rs
@@ -138,8 +138,12 @@ pub fn on_send_message_process(plugin: &mut super::TgConnector) {
     }
 }
 
-pub fn create_bot(plugin: &mut super::TgConnector, mut api: BOT) -> AmxResult<i32> {
-    if api.connect() {
+pub fn create_bot(
+    plugin: &mut super::TgConnector,
+    mut api: BOT,
+    proxy_url: Option<String>,
+) -> AmxResult<i32> {
+    if api.connect(proxy_url) {
         plugin.bots.insert(plugin.bot_context_id, api);
         plugin.bot_context_id += 1;
         Ok(plugin.bot_context_id as i32 - 1)

diff --git a/src/natives.rs b/src/natives.rs
@@ -13,30 +13,45 @@ impl super::TgConnector {
         _amx: &Amx,
         token: AmxString,
         thread_count: i32,
+        proxy_url: AmxString,
     ) -> AmxResult<i32> {
-        let api = BOT::new(token.to_string(), thread_count);
-        create_bot(self, api)
+        let proxy_url = proxy_url.to_string();
+        let proxy_url = if proxy_url.is_empty() {
+            None
+        } else {
+            Some(proxy_url)
+        };
+
+        let api = BOT::new(token.to_string(), thread_count, proxy_url.clone());
+        create_bot(self, api, proxy_url)
     }
 
     #[native(name = "TGConnectFromEnv")]
     pub fn bot_connect_from_env(
         &mut self,
         _amx: &Amx,
         variable: AmxString,
+        proxy_url: AmxString,
         thread_count: i32,
     ) -> AmxResult<i32> {
         let variable = variable.to_string();
         let token = std::env::var_os(&variable);
+        let proxy_url = proxy_url.to_string();
+        let proxy_url = if proxy_url.is_empty() {
+            None
+        } else {
+            Some(proxy_url)
+        };
 
         if token == None {
             error!("Environment variable {:?} is not set", variable);
             return Ok(-1);
         }
 
         let token = token.unwrap().into_string().unwrap();
-        let api = BOT::new(token, thread_count);
+        let api = BOT::new(token, thread_count, proxy_url.clone());
 
-        create_bot(self, api)
+        create_bot(self, api, proxy_url)
     }
 
     #[native(name = "TGSendMessage")]