inn_network/
server.rs

1//-------------------------------------------------------------------
2// MIT License
3// Copyright (c) 2022 black-mongo
4// @author CameronYang
5// @doc
6//
7// @end
8// Created : 2022-04-15T12:32:45+08:00
9//-------------------------------------------------------------------
10
11use crate::codec::AuthChoice;
12use crate::messages::*;
13use actix::prelude::*;
14use actix::{Actor, Context, Handler};
15use log::debug;
16use rand::rngs::ThreadRng;
17use rand::Rng;
18use std::collections::HashMap;
19// Enable it socks5 must auth
20const MUST_AUTH: bool = true;
21#[derive(Clone)]
22pub struct ProxyServer {
23    sessions: HashMap<usize, Recipient<ToSession>>,
24    rng: ThreadRng,
25}
26impl Default for ProxyServer {
27    fn default() -> Self {
28        ProxyServer {
29            sessions: HashMap::new(),
30            rng: rand::thread_rng(),
31        }
32    }
33}
34impl ProxyServer {
35    pub fn auth_choice(auths: &[u8]) -> AuthChoice {
36        for auth in auths {
37            if *auth == 0x02 {
38                return AuthChoice::UserNamePwd;
39            }
40            if *auth == 0x00 && !MUST_AUTH {
41                return AuthChoice::NoAuth;
42            }
43            if *auth == 0x00 && MUST_AUTH {
44                return AuthChoice::UserNamePwd;
45            }
46        }
47        AuthChoice::NoAcceptable
48    }
49}
50impl Actor for ProxyServer {
51    type Context = Context<Self>;
52    fn started(&mut self, _ctx: &mut Self::Context) {}
53}
54impl Handler<ToProxyServer> for ProxyServer {
55    type Result = MessageResult<ToProxyServer>;
56    fn handle(&mut self, command: ToProxyServer, _ctx: &mut Context<Self>) -> Self::Result {
57        match command {
58            ToProxyServer::Connect(recipient) => {
59                debug!("connect");
60                let id = self.rng.gen::<usize>();
61                self.sessions.insert(id, recipient.clone());
62                recipient.do_send(ToSession::ProxyServerReply(ProxyServerReply::Id(id)));
63                MessageResult(ProxyServerReply::Id(id))
64            }
65            ToProxyServer::DisConnect(id) => {
66                debug!("disconnect");
67                self.sessions.remove(&id);
68                MessageResult(ProxyServerReply::Ok)
69            }
70            ToProxyServer::OnlineCounter(id) => {
71                let len = self.sessions.len();
72                if let Some(session) = self.sessions.get(&id) {
73                    session.do_send(ToSession::ProxyServerReply(
74                        ProxyServerReply::OnlineCounter(len),
75                    ));
76                }
77                MessageResult(ProxyServerReply::OnlineCounter(len))
78            }
79            ToProxyServer::Cli(id, cli) => {
80                let cmd: Vec<String> = cli.into();
81                debug!("server recv cmd = {:?}", cmd);
82                match (self.sessions.get(&id), cmd.as_slice()) {
83                    (Some(session), [c]) if c == "online_counter" => {
84                        let len = self.sessions.len();
85                        session.do_send(ToSession::ProxyServerReply(ProxyServerReply::Cli(
86                            inn_common::cli::Cli::Integers(len as i64),
87                        )));
88                    }
89                    (Some(session), [c]) if c == "stop" => {
90                        session.do_send(ToSession::Stop);
91                    }
92                    _ => {}
93                }
94                MessageResult(ProxyServerReply::Ok)
95            }
96            ToProxyServer::HttpReq(http_req) => {
97                debug!("{:?}", http_req);
98                let msg = ToSession::ProxyServerReply(ProxyServerReply::HttpReq(http_req));
99                for s in self.sessions.values() {
100                    s.do_send(msg.clone());
101                }
102                MessageResult(ProxyServerReply::Ok)
103            }
104        }
105    }
106}
107#[cfg(test)]
108mod test {
109    use super::*;
110    pub struct Session;
111    impl Actor for Session {
112        type Context = Context<Self>;
113        fn started(&mut self, _: &mut Context<Self>) {}
114    }
115    impl Handler<ToSession> for Session {
116        type Result = MessageResult<ToSession>;
117        fn handle(&mut self, _: ToSession, _: &mut Context<Self>) -> Self::Result {
118            MessageResult(SessionReply::Ok)
119        }
120    }
121    #[actix_rt::test]
122    async fn new_session() {
123        let addr = ProxyServer::default().start();
124        let session = Session.start();
125        let counter = addr.send(ToProxyServer::OnlineCounter(0)).await.unwrap();
126        assert_eq!(counter, ProxyServerReply::OnlineCounter(0));
127        let id = addr
128            .send(ToProxyServer::Connect(session.recipient()))
129            .await
130            .unwrap();
131        match id {
132            ProxyServerReply::Id(id) => {
133                let counter = addr.send(ToProxyServer::OnlineCounter(0)).await.unwrap();
134                assert_eq!(counter, ProxyServerReply::OnlineCounter(1));
135                let rs = addr.send(ToProxyServer::DisConnect(id)).await.unwrap();
136                assert_eq!(rs, ProxyServerReply::Ok);
137                let counter = addr.send(ToProxyServer::OnlineCounter(0)).await.unwrap();
138                assert_eq!(counter, ProxyServerReply::OnlineCounter(0));
139            }
140            e => {
141                panic!("{:?}", e);
142            }
143        }
144    }
145}
inn_network/server.rs

inn_network/
server.rs
