summaryrefslogtreecommitdiff
path: root/src/session/manager.rs
blob: 128316430807b23c02baa16e886ff5cc36584bd3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
use crate::session::session::reverse_session_id;
use crate::session::session::Session;
use crate::session::session::SessionProto;
use crate::session::session::SessionState;
use chrono::Utc;
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

#[inline] // 内联
pub(crate) fn get_sessionmange() -> Rc<RefCell<SessionManager>> {
    // ex 本身是 Rc<RefCell<EventManager>>, 所以 clone 的是指向 实例 的引用
    crate::event::executor::EX.with(|ex| ex.session_manager.clone())
}

/******************************************************************************
 * Struct
 ******************************************************************************/

#[derive(Debug)]
pub struct SessionManager {
    sessions: HashMap<String, Rc<RefCell<Session>>>,
}

/******************************************************************************
 * API
 ******************************************************************************/

impl SessionManager {
    pub fn new(capacity: usize) -> SessionManager {
        SessionManager {
            sessions: HashMap::with_capacity(capacity),
        }
    }

    pub fn get_session(&self, session_id: &String) -> Option<Rc<RefCell<Session>>> {
        let result = self.sessions.get(session_id).map(|session| session.clone());
        if result.is_some() {
            return result;
        }

        let reversed_id = reverse_session_id(session_id);
        self.sessions
            .get(reversed_id.as_str())
            .map(|session| session.clone())
    }

    pub fn insert_session(&mut self, session_id: String, session: Rc<RefCell<Session>>) {
        self.sessions.insert(session_id, session);
    }

    pub fn remove_session(&mut self, session_id: &str) -> Option<Rc<RefCell<Session>>> {
        let result = self.sessions.remove(session_id);
        if result.is_some() {
            return result;
        }

        let reversed_id = reverse_session_id(&session_id.to_string());
        self.sessions.remove(reversed_id.as_str())
    }

    pub fn expire_sessions(&mut self) {
        let now = Utc::now().timestamp();
        let mut expired_sessions = Vec::with_capacity(1024);
        for (session_id, session) in &self.sessions {
            if session.borrow().get_session_expire_ts() < now {
                expired_sessions.push(session_id.clone());
            }
        }
        for session_id in expired_sessions {
            let option = self.remove_session(&session_id);
            if let Some(session) = option {
                println!("Session expired: {}", session_id);
                session
                    .borrow_mut()
                    .set_session_state(SessionState::Expired);
                session
                    .borrow_mut()
                    .set_session_end_ts(Utc::now().timestamp());

                match session.borrow().get_session_proto() {
                    SessionProto::TCP => {
                        //BuildInEvent::trigger_tcp_expire_event(event_mgr, Some(session));
                    }
                    SessionProto::UDP => {
                        //BuildInEvent::trigger_udp_expire_event(event_mgr, Some(session));
                    }
                }
            }
        }
    }
}

/******************************************************************************
 * TEST
 ******************************************************************************/

#[cfg(test)]
mod tests {
    use super::SessionManager;
    use crate::session::session::Session;
    use std::cell::RefCell;
    use std::rc::Rc;

    // use std::{thread, time};

    #[test]
    fn test_session_manager() {
        let session_id = "IP4->IP4;192.168.0.1->192.168.0.2;UDP->UDP;2345->80;".to_string();
        let reversed_id = "IP4->IP4;192.168.0.2->192.168.0.1;UDP->UDP;80->2345;".to_string();
        let session = Rc::new(RefCell::new(Session::new(session_id.clone())));

        // Create Session Manager
        let mut session_mgr = SessionManager::new(1024);

        // Search Session
        assert_eq!(session_mgr.get_session(&session_id).is_none(), true);

        // Insert Session
        session_mgr.insert_session(session_id.clone(), session);

        // Update Session
        assert_eq!(session_mgr.get_session(&session_id).is_some(), true);
        assert_eq!(session_mgr.get_session(&reversed_id).is_some(), true);

        // Expire Session
        // thread::sleep(time::Duration::from_secs(61));
        // session_mgr.expire_sessions();

        // Delete session
        assert_eq!(session_mgr.remove_session(&reversed_id).is_some(), true);

        // Research Session
        assert_eq!(session_mgr.get_session(&session_id).is_none(), true);
        assert_eq!(session_mgr.get_session(&reversed_id).is_none(), true);

        dbg!(session_mgr);
    }
}