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
140
141
142
143
144

// Copyright (c) 2016 The Rouille developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::collections::HashMap;
use std::sync::Mutex;
use rand;
use rand::Rng;

use Request;
use Response;
use input;

/// Manages all active user sessions in memory.
///
/// # Example
///
/// ```no_run
/// #[derive(Debug, Clone)]
/// struct SessionData {
///     user_id: i32
/// }
///
/// let sessions = rouille::SessionsManager::<SessionData>::new("SID", 3600);
/// 
/// rouille::start_server("localhost:80", move |request| {
///     let session = sessions.start(&request);
///     // rest of the handler
/// # let response: rouille::Response = unsafe { ::std::mem::uninitialized() };
///     session.apply(response)
/// })
/// ```
///
pub struct SessionsManager<T> where T: Clone {
    // TODO: eventually replace the key with `[u8; 64]` or something similar
    sessions: Mutex<HashMap<String, T>>,
    cookie_name: String,
    timeout_s: u64,
}

impl<T> SessionsManager<T> where T: Clone {
    /// Initializes the sessions manager.
    ///
    /// # Parameters
    ///
    /// - `cookie_name`: The name of the cookie to use. Usually `SID`.
    /// - `timeout_s`: The duration of the session, in seconds. Usually 3600.
    ///
    pub fn new<S>(cookie_name: S, timeout_s: u64) -> SessionsManager<T> where S: Into<String> {
        SessionsManager {
            sessions: Mutex::new(HashMap::new()),
            cookie_name: cookie_name.into(),
            timeout_s: timeout_s,
        }
    }

    /// Tries to load an existing session from the request, or creates one if there
    /// is no session yet.
    pub fn start(&self, request: &Request) -> Session<T> {
        let mut cookie = input::get_cookies(request).into_iter();
        let cookie = cookie.find(|&(ref k, _)| k == &self.cookie_name);
        let cookie = cookie.map(|(k, v)| v);

        if let Some(cookie) = cookie {
            Session {
                manager: self,
                key: cookie,
            }
        } else {
            Session {
                manager: self,
                key: generate_session_id(),
            }
        }
    }
}

/// Represents an entry in the sessions manager.
pub struct Session<'a, T> where T: Clone + 'a {
    manager: &'a SessionsManager<T>,
    key: String,
}

impl<'a, T> Session<'a, T> where T: Clone {
    /// Load the session infos from the manager. Returns `None` if there is no data yet.
    ///
    /// Note that calling `get` twice in a row can produce different results. That can happen
    /// if two requests are processed in parallel and access the same session.
    pub fn get(&self) -> Option<T> {
        let session = self.manager.sessions.lock().unwrap();
        session.get(&self.key).map(|d| d.clone())
    }

    /// Returns true if there is session data.
    #[inline]
    pub fn has_data(&self) -> bool {
        let session = self.manager.sessions.lock().unwrap();
        session.get(&self.key).is_some()
    }

    /// Stores the session infos in the manager.
    pub fn set(&self, value: T) {
        let mut session = self.manager.sessions.lock().unwrap();
        session.insert(self.key.clone(), value);
    }

    /// Removes the session infos from the manager.
    pub fn clear(&self) {
        let mut session = self.manager.sessions.lock().unwrap();
        session.remove(&self.key);
    }

    /// Applies the session on the `Response`. If you don't do that, the session won't be
    /// maintained on further connections.
    pub fn apply(&self, mut response: Response) -> Response {
        if !self.has_data() {
            return response;
        }

        // FIXME: correct interactions with existing headers
        let header_value = format!("{}={}; Max-Age={}", self.manager.cookie_name, self.key,
                                                        self.manager.timeout_s);
        response.headers.push(("Set-Cookie".to_owned(), header_value));
        response
    }
}

/// Generates a string suitable for a session ID.
///
/// The output string doesn't contain any punctuation or character such as quotes or brackets
/// that could need to be escaped.
pub fn generate_session_id() -> String {
    // 5e+114 possibilities is reasonable
    rand::OsRng::new().unwrap()     // TODO: <- how to handle that?
                      .gen_ascii_chars()
                      .filter(|&c| (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
                                   (c >= '0' && c <= '9'))
                      .take(64).collect::<String>()
}