shorty/
lib.rs

1// Copyright 2019 Federico Fissore
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//   http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! # Shorty
16//!
17//! `shorty` is a URL shortener: it assigns a short ID to a URL of any length, and when people will
18//! access the URL with that short ID, they will be redirected to the original URL.
19//!
20#[macro_use]
21extern crate serde_derive;
22
23use core::fmt;
24use std::error::Error;
25use std::fmt::{Display, Formatter};
26
27use redis::{ErrorKind, RedisError};
28use url::Url;
29
30#[cfg(test)]
31use tests::StubRedisFacade as RedisFacade;
32
33#[cfg(not(test))]
34use crate::redis_facade::RedisFacade;
35
36#[cfg(not(test))]
37pub mod redis_facade;
38
39#[derive(Debug)]
40pub struct ShortenerError {
41    message: &'static str,
42    cause: Option<Box<dyn Error>>,
43}
44
45impl ShortenerError {
46    fn new(message: &'static str) -> ShortenerError {
47        ShortenerError {
48            message,
49            cause: None,
50        }
51    }
52    fn new_with_cause(message: &'static str, error: Box<dyn Error>) -> ShortenerError {
53        ShortenerError {
54            message,
55            cause: Some(error),
56        }
57    }
58}
59
60impl Display for ShortenerError {
61    fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
62        f.write_str(&self.message)?;
63
64        if let Some(err) = &self.cause {
65            return f.write_str(&format!(" - {}", &err));
66        }
67
68        Ok(())
69    }
70}
71
72impl Error for ShortenerError {}
73
74/// `Shortener` is the struct exposing methods `lookup` and `shorten`.
75///
76/// `lookup` attempts to resolve an ID to a URL. If no URL is found or an error occurs, it returns
77/// `None`, otherwise it returns `Some(url)`.
78///
79/// `shorten` takes an optional API key and a URL to shorten. If the API key is present, it will
80/// validate it and shorten the URL only if validation passes. Otherwise, it will just shorten the
81/// URL.
82///
83/// `Shortener` interacts with a `RedisFacade`, which makes it easier to work with the `redis` crate
84/// and simplifies testing.
85pub struct Shortener {
86    id_length: usize,
87    id_alphabet: Vec<char>,
88    id_generation_max_attempts: u8,
89    redis: RedisFacade,
90    rate_limit_period: usize,
91    rate_limit: i64,
92}
93
94/// A struct with the successful result of a URL shortening. It holds the original `url` and the
95/// resulting `id`
96#[derive(Serialize)]
97pub struct ShortenerResult {
98    id: String,
99    url: String,
100}
101
102impl Shortener {
103    /// Creates a new Shortener
104    ///
105    /// `id_length` is the length of the generated ID.
106    ///
107    /// `id_alphabet` is the alphabet used in the ID: a decent one is `a-zA-Z0-9` as each entry has
108    /// 62 possible values and is ASCII
109    ///
110    /// `id_generation_max_attempts` is the number of attempts to generate an unique ID when a
111    /// conflict is detected.
112    ///
113    /// `redis` is a `RedisFacade` instance.
114    ///
115    /// `rate_limit_period` is the amount of seconds during which calls to `shorten` will be counted.
116    ///
117    /// `rate_limit` is the max number of calls that can be made to `shorten` in a period.
118    pub fn new(
119        id_length: usize,
120        id_alphabet: Vec<char>,
121        id_generation_max_attempts: u8,
122        redis: RedisFacade,
123        rate_limit_period: usize,
124        rate_limit: i64,
125    ) -> Shortener {
126        Shortener {
127            id_length,
128            id_alphabet,
129            id_generation_max_attempts,
130            redis,
131            rate_limit_period,
132            rate_limit,
133        }
134    }
135
136    /// Looks up a URL by the given ID. If no URL is found or an error occurs, it returns `None`,
137    /// otherwise it returns `Some(url)`.
138    pub fn lookup(&self, id: &str) -> Option<String> {
139        match self.redis.get_string(id) {
140            Ok(url) => Some(url),
141            Err(_) => None,
142        }
143    }
144
145    fn verify_api_key(&self, api_key: &str) -> Result<(), ShortenerError> {
146        let api_key = format!("API_KEY_{}", api_key);
147        log::trace!("verifying api key '{}'", api_key);
148
149        let verify_and_increment = self.redis.get_bool(&api_key).and_then(|valid| {
150            if !valid {
151                return Err(RedisError::from((
152                    ErrorKind::ExtensionError,
153                    "API key expired",
154                )));
155            }
156
157            if self.rate_limit <= 0 {
158                return Ok(-1);
159            }
160
161            let rate_key = format!("RATE_{}", api_key);
162            log::trace!("verifying rate key '{}'", rate_key);
163
164            self.redis.exists(&rate_key).and_then(|exists| {
165                log::trace!("rate key exists {}", exists);
166
167                self.redis.increment(&rate_key).and_then(|number_of_calls| {
168                    log::trace!("rate key {} number of calls {}", rate_key, number_of_calls);
169
170                    if !exists {
171                        self.redis
172                            .expire(&rate_key, self.rate_limit_period)
173                            .unwrap();
174                    }
175
176                    Ok(number_of_calls)
177                })
178            })
179        });
180
181        match verify_and_increment {
182            Ok(call_rate) if self.rate_limit > 0 && call_rate > self.rate_limit => {
183                Err(ShortenerError::new("Rate limit exceeded"))
184            }
185            Ok(_) => Ok(()),
186            Err(err) => Err(ShortenerError::new_with_cause(
187                "Invalid API key",
188                Box::new(err),
189            )),
190        }
191    }
192
193    fn generate_id(&self) -> Result<String, ShortenerError> {
194        for _ in 1..=self.id_generation_max_attempts {
195            let id = nanoid::custom(self.id_length, &self.id_alphabet);
196
197            let exists = self.redis.exists(&id).unwrap_or(false);
198
199            if !exists {
200                return Ok(id);
201            }
202        }
203
204        Err(ShortenerError::new(
205            "Failed to generate an ID: too many attempts. Consider using a longer ID",
206        ))
207    }
208
209    /// Shortens an URL, returning a `ShortenerResult` holding the provided URL and the generated ID.
210    ///
211    /// If the optional API key is present, it will validate it and shorten the URL only if
212    /// validation passes.
213    ///
214    /// If the optional host is present, it will ensure that the url to shorten is not a url from
215    /// the same host that's running shorty (which would create a link loop)
216    ///
217    /// Otherwise, it will just shorten the URL.
218    pub fn shorten(
219        &self,
220        api_key: &Option<&str>,
221        host: Option<&str>,
222        url: &str,
223    ) -> Result<ShortenerResult, ShortenerError> {
224        let verify_result = api_key
225            .as_ref()
226            .map(|api_key| self.verify_api_key(api_key))
227            .unwrap_or(Ok(()));
228
229        verify_result
230            .and_then(|_| self.generate_id())
231            .and_then(|id| {
232                let mut url = url.to_owned();
233                if !url.to_lowercase().starts_with("http") {
234                    url = format!("http://{}", url);
235                }
236                Url::parse(&url)
237                    .and_then(|parsed_url| Ok((id, url, parsed_url)))
238                    .map_err(|parse_err| {
239                        ShortenerError::new_with_cause("Unable to parse url", Box::new(parse_err))
240                    })
241            })
242            .and_then(|(id, url, parsed_url)| {
243                if host.is_none() {
244                    return Ok((id, url));
245                }
246
247                if parsed_url.host_str().unwrap().eq(host.unwrap()) {
248                    return Err(ShortenerError::new("Link loop is not allowed"));
249                }
250
251                Ok((id, url))
252            })
253            .and_then(|(id, url)| {
254                self.redis
255                    .set(&id, url.as_str())
256                    .map(|_| ShortenerResult { id, url })
257                    .map_err(|err| ShortenerError::new_with_cause("Redis error", Box::new(err)))
258            })
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    use std::cell::RefCell;
265
266    use redis::RedisResult;
267
268    use super::*;
269
270    pub struct StubRedisFacade {
271        get_string_answers: RefCell<Vec<RedisResult<String>>>,
272        get_bool_answers: RefCell<Vec<RedisResult<bool>>>,
273        exists_answers: RefCell<Vec<RedisResult<bool>>>,
274        set_answers: RefCell<Vec<RedisResult<()>>>,
275        incr_answers: RefCell<Vec<RedisResult<i64>>>,
276        expire_answers: RefCell<Vec<RedisResult<()>>>,
277    }
278
279    impl StubRedisFacade {
280        fn new() -> Self {
281            StubRedisFacade {
282                get_string_answers: RefCell::new(vec![]),
283                get_bool_answers: RefCell::new(vec![]),
284                exists_answers: RefCell::new(vec![]),
285                set_answers: RefCell::new(vec![]),
286                incr_answers: RefCell::new(vec![]),
287                expire_answers: RefCell::new(vec![]),
288            }
289        }
290
291        pub fn get_string(&self, _key: &str) -> RedisResult<String> {
292            if self.get_string_answers.borrow().len() > 0 {
293                return self.get_string_answers.borrow_mut().remove(0);
294            }
295            panic!("unexpected get_string call");
296        }
297
298        pub fn get_bool(&self, _key: &str) -> RedisResult<bool> {
299            if self.get_bool_answers.borrow().len() > 0 {
300                return self.get_bool_answers.borrow_mut().remove(0);
301            }
302            panic!("unexpected get_bool call");
303        }
304
305        pub fn exists(&self, _key: &str) -> RedisResult<bool> {
306            if self.exists_answers.borrow().len() > 0 {
307                return self.exists_answers.borrow_mut().remove(0);
308            }
309            panic!("unexpected exists call");
310        }
311
312        pub fn set(&self, _key: &str, _value: &str) -> RedisResult<()> {
313            if self.set_answers.borrow().len() > 0 {
314                return self.set_answers.borrow_mut().remove(0);
315            }
316            panic!("unexpected set call");
317        }
318
319        pub fn increment(&self, _key: &str) -> RedisResult<i64> {
320            if self.incr_answers.borrow().len() > 0 {
321                return self.incr_answers.borrow_mut().remove(0);
322            }
323            panic!("unexpected increment call");
324        }
325
326        pub fn expire(&self, _key: &str, _seconds: usize) -> RedisResult<()> {
327            if self.expire_answers.borrow().len() > 0 {
328                return self.expire_answers.borrow_mut().remove(0);
329            }
330            panic!("unexpected expire call");
331        }
332    }
333
334    #[test]
335    fn test_lookup() {
336        let redis = StubRedisFacade::new();
337        &redis
338            .get_string_answers
339            .borrow_mut()
340            .push(Ok(String::from("test url")));
341
342        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, 10);
343        assert_eq!(shortener.lookup("id").unwrap(), "test url");
344    }
345
346    #[test]
347    fn test_shorten_happy_path_first_call() {
348        let redis = StubRedisFacade::new();
349        // api key verification
350        &redis.get_bool_answers.borrow_mut().push(Ok(true));
351        &redis.exists_answers.borrow_mut().push(Ok(false));
352        &redis.incr_answers.borrow_mut().push(Ok(1));
353        &redis.expire_answers.borrow_mut().push(Ok(()));
354
355        // id generation
356        &redis.exists_answers.borrow_mut().push(Ok(false));
357
358        // shortened url storage
359        &redis.set_answers.borrow_mut().push(Ok(()));
360
361        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, 10);
362        let shorten_result = shortener
363            .shorten(&Some("api key"), Some("with.lv"), "example.com")
364            .unwrap();
365        assert_eq!(10, shorten_result.id.len());
366        assert_eq!("http://example.com", shorten_result.url);
367    }
368
369    #[test]
370    fn test_shorten_happy_path_no_rate_limit() {
371        let redis = StubRedisFacade::new();
372        // api key verification
373        &redis.get_bool_answers.borrow_mut().push(Ok(true));
374
375        // id generation
376        &redis.exists_answers.borrow_mut().push(Ok(false));
377
378        // shortened url storage
379        &redis.set_answers.borrow_mut().push(Ok(()));
380
381        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, -1);
382        let shorten_result = shortener
383            .shorten(&Some("api key"), Some("with.lv"), "example.com")
384            .unwrap();
385        assert_eq!(10, shorten_result.id.len());
386        assert_eq!("http://example.com", shorten_result.url);
387    }
388
389    #[test]
390    fn test_shorten_happy_path_second_call() {
391        let redis = StubRedisFacade::new();
392        // api key verification
393        &redis.get_bool_answers.borrow_mut().push(Ok(true));
394        &redis.exists_answers.borrow_mut().push(Ok(true));
395        &redis.incr_answers.borrow_mut().push(Ok(2));
396
397        // id generation
398        &redis.exists_answers.borrow_mut().push(Ok(false));
399
400        // shortened url storage
401        &redis.set_answers.borrow_mut().push(Ok(()));
402
403        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, 10);
404        let shorten_result = shortener
405            .shorten(&Some("api key"), Some("with.lv"), "example.com")
406            .unwrap();
407        assert_eq!(10, shorten_result.id.len());
408        assert_eq!("http://example.com", shorten_result.url);
409    }
410
411    #[test]
412    fn test_shorten_happy_path_no_api_key() {
413        let redis = StubRedisFacade::new();
414        // id generation
415        &redis.exists_answers.borrow_mut().push(Ok(false));
416
417        // shortened url storage
418        &redis.set_answers.borrow_mut().push(Ok(()));
419
420        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, 10);
421        let shorten_result = shortener
422            .shorten(&None, Some("with.lv"), "example.com")
423            .unwrap();
424        assert_eq!(10, shorten_result.id.len());
425        assert_eq!("http://example.com", shorten_result.url);
426    }
427
428    #[test]
429    fn test_shorten_unhappy_path_rate_limit_exceeded() {
430        let rate_limit = 10;
431        let redis = StubRedisFacade::new();
432        // api key verification
433        &redis.get_bool_answers.borrow_mut().push(Ok(true));
434        &redis.exists_answers.borrow_mut().push(Ok(true));
435        &redis.incr_answers.borrow_mut().push(Ok(rate_limit + 1));
436
437        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, rate_limit);
438        let shorten_result_err = shortener
439            .shorten(&Some("api key"), Some("with.lv"), "example.com")
440            .err()
441            .unwrap();
442        assert_eq!("Rate limit exceeded", shorten_result_err.message);
443    }
444
445    #[test]
446    fn test_shorten_unhappy_path_bad_url() {
447        let redis = StubRedisFacade::new();
448        // id generation
449        &redis.exists_answers.borrow_mut().push(Ok(false));
450
451        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, -1);
452        let shorten_result_err = shortener
453            .shorten(&None, Some("with.lv"), "wrong domain.com")
454            .err()
455            .unwrap();
456        assert_eq!("Unable to parse url", shorten_result_err.message);
457    }
458
459    #[test]
460    fn test_shorten_unhappy_path_same_domain() {
461        let redis = StubRedisFacade::new();
462        // id generation
463        &redis.exists_answers.borrow_mut().push(Ok(false));
464
465        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, -1);
466        let shorten_result_err = shortener
467            .shorten(&None, Some("example.com"), "example.com")
468            .err()
469            .unwrap();
470        assert_eq!("Link loop is not allowed", shorten_result_err.message);
471    }
472
473    #[test]
474    fn test_shorten_happy_path_rate_limit_expired() {
475        let redis = StubRedisFacade::new();
476
477        // api key verification
478        &redis.get_bool_answers.borrow_mut().push(Ok(true));
479        &redis.exists_answers.borrow_mut().push(Ok(true));
480        &redis.incr_answers.borrow_mut().push(Ok(1));
481
482        // id generation
483        &redis.exists_answers.borrow_mut().push(Ok(false));
484
485        // shortened url storage
486        &redis.set_answers.borrow_mut().push(Ok(()));
487
488        // api key verification
489        &redis.get_bool_answers.borrow_mut().push(Ok(true));
490        &redis.exists_answers.borrow_mut().push(Ok(false));
491        &redis.incr_answers.borrow_mut().push(Ok(1));
492        &redis.expire_answers.borrow_mut().push(Ok(()));
493
494        // id generation
495        &redis.exists_answers.borrow_mut().push(Ok(false));
496
497        // shortened url storage
498        &redis.set_answers.borrow_mut().push(Ok(()));
499
500        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, 10);
501
502        let shorten_result = shortener
503            .shorten(&Some("api key"), Some("with.lv"), "example.com")
504            .unwrap();
505        assert_eq!(10, shorten_result.id.len());
506        assert_eq!("http://example.com", shorten_result.url);
507
508        let shorten_result = shortener
509            .shorten(&Some("api key"), Some("with.lv"), "www.wikipedia.org")
510            .unwrap();
511        assert_eq!(10, shorten_result.id.len());
512        assert_eq!("http://www.wikipedia.org", shorten_result.url);
513    }
514
515    #[test]
516    fn test_shorten_unhappy_path_invalid_api_key() {
517        let redis = StubRedisFacade::new();
518
519        // api key verification
520        &redis.get_bool_answers.borrow_mut().push(Ok(false));
521
522        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 10, redis, 600, 10);
523        let shorten_result_err = shortener
524            .shorten(&Some("api key"), Some("with.lv"), "example.com")
525            .err()
526            .unwrap();
527        assert_eq!("Invalid API key", shorten_result_err.message);
528    }
529
530    #[test]
531    fn test_shorten_unhappy_path_too_many_attempts_generating_id() {
532        let redis = StubRedisFacade::new();
533
534        // id generation attempts
535        &redis.exists_answers.borrow_mut().push(Ok(true));
536        &redis.exists_answers.borrow_mut().push(Ok(true));
537
538        let shortener = Shortener::new(10, vec!['a', 'b', 'c'], 2, redis, 600, 10);
539        let shorten_result_err = shortener
540            .shorten(&None, Some("with.lv"), "example.com")
541            .err()
542            .unwrap();
543        assert_eq!(
544            "Failed to generate an ID: too many attempts. Consider using a longer ID",
545            shorten_result_err.message
546        );
547    }
548}
shorty/lib.rs

shorty/
lib.rs
