rust-cache-server 0.1.0

High-performance Redis-like cache server in Rust - open source alternative without Redis licensing concerns
Documentation
use std::sync::Arc;
use std::time::Duration;

use crate::protocol::RespFrame;
use crate::store::Store;

pub async fn handle_command(store: Arc<Store>, frame: RespFrame) -> RespFrame {
    let RespFrame::Array(Some(items)) = frame else {
        return RespFrame::Error("ERR protocol error".into());
    };

    if items.is_empty() {
        return RespFrame::Error("ERR empty command".into());
    }

    let cmd_name = match &items[0] {
        RespFrame::BulkString(Some(b)) => String::from_utf8_lossy(b).to_ascii_uppercase(),
        RespFrame::SimpleString(s) => s.to_ascii_uppercase(),
        _ => return RespFrame::Error("ERR invalid command".into()),
    };

    match cmd_name.as_str() {
        "PING" => RespFrame::SimpleString("PONG".into()),

        "INFO" => {
            // Minimal INFO response so that ioredis is happy.
            // Redis INFO is a bulk string with lines "key:value\r\n".
            let info = "redis_version:rust-cache-0.1.0\r\nrole:master\r\n";
            RespFrame::BulkString(Some(info.as_bytes().to_vec()))
        }

        "GET" => {
            if items.len() != 2 {
                return RespFrame::Error("ERR wrong number of arguments for 'get' command".into());
            }
            let key = bulk_to_string(&items[1]);
            match key {
                None => RespFrame::Error("ERR invalid key".into()),
                Some(k) => match store.get(&k) {
                    Some(v) => RespFrame::BulkString(Some(v)),
                    None => RespFrame::BulkString(None),
                },
            }
        }

        "SET" => {
            if items.len() < 3 {
                return RespFrame::Error("ERR wrong number of arguments for 'set' command".into());
            }
            let key = match bulk_to_string(&items[1]) {
                Some(k) => k,
                None => return RespFrame::Error("ERR invalid key".into()),
            };
            let val = match &items[2] {
                RespFrame::BulkString(Some(b)) => b.clone(),
                RespFrame::SimpleString(s) => s.clone().into_bytes(),
                _ => return RespFrame::Error("ERR invalid value".into()),
            };

            // Optional EX seconds
            let mut ttl: Option<Duration> = None;
            if items.len() >= 5 {
                if let RespFrame::BulkString(Some(flag)) = &items[3] {
                    if flag.eq_ignore_ascii_case(b"EX") {
                        if let RespFrame::BulkString(Some(sec_bytes)) = &items[4] {
                            if let Ok(sec_str) = std::str::from_utf8(sec_bytes) {
                                if let Ok(sec) = sec_str.parse::<u64>() {
                                    ttl = Some(Duration::from_secs(sec));
                                }
                            }
                        }
                    }
                }
            }

            match store.set(key, val, ttl) {
                Ok(()) => RespFrame::SimpleString("OK".into()),
                Err(()) => RespFrame::Error("ERR memory limit reached".into()),
            }
        }

        "DEL" => {
            if items.len() < 2 {
                return RespFrame::Error("ERR wrong number of arguments for 'del' command".into());
            }
            let mut keys = Vec::with_capacity(items.len() - 1);
            for it in &items[1..] {
                if let Some(k) = bulk_to_string(it) {
                    keys.push(k);
                }
            }
            let removed = store.del(&keys);
            RespFrame::Integer(removed as i64)
        }

        "EXPIRE" => {
            if items.len() != 3 {
                return RespFrame::Error(
                    "ERR wrong number of arguments for 'expire' command".into(),
                );
            }
            let key = match bulk_to_string(&items[1]) {
                Some(k) => k,
                None => return RespFrame::Error("ERR invalid key".into()),
            };
            let seconds = match &items[2] {
                RespFrame::BulkString(Some(b)) => match std::str::from_utf8(b) {
                    Ok(s) => match s.parse::<u64>() {
                        Ok(v) => v,
                        Err(_) => return RespFrame::Error("ERR value is not an integer".into()),
                    },
                    Err(_) => return RespFrame::Error("ERR value is not an integer".into()),
                },
                _ => return RespFrame::Error("ERR value is not an integer".into()),
            };

            let ok = store.expire(&key, Duration::from_secs(seconds));
            RespFrame::Integer(if ok { 1 } else { 0 })
        }

        "TTL" => {
            if items.len() != 2 {
                return RespFrame::Error("ERR wrong number of arguments for 'ttl' command".into());
            }
            let key = match bulk_to_string(&items[1]) {
                Some(k) => k,
                None => return RespFrame::Error("ERR invalid key".into()),
            };
            match store.ttl(&key) {
                Some(d) => RespFrame::Integer(d.as_secs() as i64),
                None => RespFrame::Integer(-2),
            }
        }

        "KEYS" => {
            if items.len() != 2 {
                return RespFrame::Error("ERR wrong number of arguments for 'keys' command".into());
            }
            let pattern = match bulk_to_string(&items[1]) {
                Some(p) => p,
                None => return RespFrame::Error("ERR invalid pattern".into()),
            };

            let all_keys = store.keys();
            let matched: Vec<RespFrame> = all_keys
                .into_iter()
                .filter(|k| pattern_match(&pattern, k))
                .map(|k| RespFrame::BulkString(Some(k.into_bytes())))
                .collect();

            RespFrame::Array(Some(matched))
        }

        "FLUSHALL" | "CLEARALL" => {
            if items.len() != 1 {
                return RespFrame::Error(
                    "ERR wrong number of arguments for 'flushall' command".into(),
                );
            }
            store.clear_all();
            RespFrame::SimpleString("OK".into())
        }

        "INCR" => {
            if items.len() != 2 {
                return RespFrame::Error("ERR wrong number of arguments for 'incr' command".into());
            }
            let key = match bulk_to_string(&items[1]) {
                Some(k) => k,
                None => return RespFrame::Error("ERR invalid key".into()),
            };
            match store.incr_by(&key, 1, true) {
                Some(val) => RespFrame::Integer(val),
                None => RespFrame::Error("ERR value is not an integer or out of range".into()),
            }
        }

        "DECR" => {
            if items.len() != 2 {
                return RespFrame::Error("ERR wrong number of arguments for 'decr' command".into());
            }
            let key = match bulk_to_string(&items[1]) {
                Some(k) => k,
                None => return RespFrame::Error("ERR invalid key".into()),
            };
            match store.incr_by(&key, -1, true) {
                Some(val) => RespFrame::Integer(val),
                None => RespFrame::Error("ERR value is not an integer or out of range".into()),
            }
        }

        "INCRBY" => {
            if items.len() != 3 {
                return RespFrame::Error(
                    "ERR wrong number of arguments for 'incrby' command".into(),
                );
            }
            let key = match bulk_to_string(&items[1]) {
                Some(k) => k,
                None => return RespFrame::Error("ERR invalid key".into()),
            };
            let increment = match &items[2] {
                RespFrame::BulkString(Some(b)) => match std::str::from_utf8(b) {
                    Ok(s) => match s.parse::<i64>() {
                        Ok(v) => v,
                        Err(_) => {
                            return RespFrame::Error("ERR value is not an integer".into());
                        }
                    },
                    Err(_) => return RespFrame::Error("ERR value is not an integer".into()),
                },
                _ => return RespFrame::Error("ERR value is not an integer".into()),
            };
            match store.incr_by(&key, increment, true) {
                Some(val) => RespFrame::Integer(val),
                None => RespFrame::Error("ERR value is not an integer or out of range".into()),
            }
        }

        "DECRBY" => {
            if items.len() != 3 {
                return RespFrame::Error(
                    "ERR wrong number of arguments for 'decrby' command".into(),
                );
            }
            let key = match bulk_to_string(&items[1]) {
                Some(k) => k,
                None => return RespFrame::Error("ERR invalid key".into()),
            };
            let decrement = match &items[2] {
                RespFrame::BulkString(Some(b)) => match std::str::from_utf8(b) {
                    Ok(s) => match s.parse::<i64>() {
                        Ok(v) => v,
                        Err(_) => {
                            return RespFrame::Error("ERR value is not an integer".into());
                        }
                    },
                    Err(_) => return RespFrame::Error("ERR value is not an integer".into()),
                },
                _ => return RespFrame::Error("ERR value is not an integer".into()),
            };
            match store.incr_by(&key, -decrement, true) {
                Some(val) => RespFrame::Integer(val),
                None => RespFrame::Error("ERR value is not an integer or out of range".into()),
            }
        }

        _ => RespFrame::Error(format!("ERR unknown command '{cmd_name}'")),
    }
}

fn bulk_to_string(frame: &RespFrame) -> Option<String> {
    match frame {
        RespFrame::BulkString(Some(b)) => String::from_utf8(b.clone()).ok(),
        RespFrame::SimpleString(s) => Some(s.clone()),
        _ => None,
    }
}

fn pattern_match(pattern: &str, text: &str) -> bool {
    // Very small glob subset: only '*' wildcard, like Redis KEYS.
    if pattern == "*" {
        return true;
    }

    let mut p_iter = pattern.chars().peekable();
    let mut t_iter = text.chars().peekable();

    let mut star: Option<(usize, usize)> = None;
    let pattern_chars: Vec<char> = pattern.chars().collect();
    let text_chars: Vec<char> = text.chars().collect();
    let (mut i, mut j) = (0usize, 0usize);

    while j < text_chars.len() {
        if i < pattern_chars.len()
            && (pattern_chars[i] == '?' || pattern_chars[i] == text_chars[j])
        {
            i += 1;
            j += 1;
        } else if i < pattern_chars.len() && pattern_chars[i] == '*' {
            star = Some((i, j));
            i += 1;
        } else if let Some((si, sj)) = star {
            i = si + 1;
            j = sj + 1;
            star = Some((si, j));
        } else {
            return false;
        }
    }

    while i < pattern_chars.len() && pattern_chars[i] == '*' {
        i += 1;
    }

    i == pattern_chars.len()
}