use super::args::{self, Args, is};
use super::clients::{self, Client};
use super::table::{self, Spec};
use super::{Flow, Reply, Server, Session};
use crate::proto::Proto;
use crate::reply::Out;
use core::fmt::Write;
use std::sync::atomic::Ordering::{Acquire, Relaxed};
use yo_common::{Code, Error, Result};
pub(super) fn execute(
server: &Server,
session: &mut Session,
_spec: &Spec,
args: Args<'_>,
out: &mut Out,
) -> Result<Flow> {
let sub = args.get(1);
if is(sub, b"ID") {
one(args, "id")?;
out.int(session.id as i64);
} else if is(sub, b"GETNAME") {
one(args, "getname")?;
if session.name.is_empty() {
out.nil();
} else {
out.bulk(&session.name);
}
} else if is(sub, b"SETNAME") {
two(args, "setname")?;
let name = args.get(2);
if !printable(name) {
return Err(Error::new(
Code::Invalid,
"Client names cannot contain spaces, newlines or special characters.",
));
}
session.set_name(name);
out.ok();
} else if is(sub, b"SETINFO") {
setinfo(session, args)?;
out.ok();
} else if is(sub, b"INFO") {
one(args, "info")?;
let text = report(server, session, out.proto(), out.len());
out.verbatim(b"txt", text.as_bytes());
} else if is(sub, b"LIST") {
list(server, args, out)?;
} else if is(sub, b"KILL") {
if kill(server, session, args, out)? {
return Ok(Flow::Close);
}
} else if is(sub, b"PAUSE") {
pause(server, args)?;
out.ok();
} else if is(sub, b"UNPAUSE") {
one(args, "unpause")?;
server.unpause();
out.ok();
} else if is(sub, b"REPLY") {
reply(session, args)?;
out.ok();
} else if is(sub, b"NO-EVICT") {
let on = on_off(args, "no-evict")?;
session.set_no_evict(on);
out.ok();
} else if is(sub, b"NO-TOUCH") {
let on = on_off(args, "no-touch")?;
session.set_no_touch(on);
out.ok();
} else if is(sub, b"HELP") {
super::server::help(out, CLIENT_HELP);
} else {
return Err(args::unknown_subcommand(sub, "CLIENT"));
}
Ok(Flow::Continue)
}
fn one(args: Args<'_>, sub: &str) -> Result<()> {
if args.len() == 2 {
Ok(())
} else {
Err(args::wrong_arity_sub("client", sub))
}
}
fn two(args: Args<'_>, sub: &str) -> Result<()> {
if args.len() == 3 {
Ok(())
} else {
Err(args::wrong_arity_sub("client", sub))
}
}
fn printable(value: &[u8]) -> bool {
value.iter().all(|b| (b'!'..=b'~').contains(b))
}
fn setinfo(session: &mut Session, args: Args<'_>) -> Result<()> {
if args.len() != 4 {
return Err(args::wrong_arity_sub("client", "setinfo"));
}
let what = args.get(2);
let value = args.get(3);
let name = if is(what, b"LIB-NAME") {
"lib-name"
} else if is(what, b"LIB-VER") {
"lib-ver"
} else {
return Err(Error::fmt(
Code::Invalid,
format_args!("Unrecognized option '{}'", String::from_utf8_lossy(what)),
));
};
if !printable(value) {
return Err(Error::fmt(
Code::Invalid,
format_args!("{name} cannot contain spaces, newlines or special characters."),
));
}
if name == "lib-name" {
session.set_lib_name(value);
} else {
session.set_lib_ver(value);
}
Ok(())
}
fn reply(session: &mut Session, args: Args<'_>) -> Result<()> {
if args.len() != 3 {
return Err(args::wrong_arity_sub("client", "reply"));
}
let mode = args.get(2);
session.reply = if is(mode, b"ON") {
Reply::On
} else if is(mode, b"OFF") {
Reply::Off
} else if is(mode, b"SKIP") {
Reply::SkipNext
} else {
return Err(args::syntax());
};
Ok(())
}
fn on_off(args: Args<'_>, sub: &str) -> Result<bool> {
if args.len() != 3 {
return Err(args::wrong_arity_sub("client", sub));
}
let arg = args.get(2);
if is(arg, b"ON") {
Ok(true)
} else if is(arg, b"OFF") {
Ok(false)
} else {
Err(args::syntax())
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Kind {
Any,
Normal,
Pubsub,
Link,
}
impl Kind {
fn parse(word: &[u8]) -> Result<Kind> {
if is(word, b"normal") {
Ok(Kind::Normal)
} else if is(word, b"pubsub") {
Ok(Kind::Pubsub)
} else if is(word, b"master") || is(word, b"replica") || is(word, b"slave") {
Ok(Kind::Link)
} else {
Err(Error::fmt(
Code::Invalid,
format_args!("Unknown client type '{}'", String::from_utf8_lossy(word)),
))
}
}
fn covers(self, row: &Client) -> bool {
match self {
Kind::Any => true,
Kind::Normal => !row.flag(clients::SUBSCRIBED),
Kind::Pubsub => row.flag(clients::SUBSCRIBED),
Kind::Link => false,
}
}
}
fn list(server: &Server, args: Args<'_>, out: &mut Out) -> Result<()> {
let mut want = Kind::Any;
let mut ids: Vec<u64> = Vec::new();
if args.len() == 4 && is(args.get(2), b"TYPE") {
want = Kind::parse(args.get(3))?;
} else if args.len() > 3 && is(args.get(2), b"ID") {
for at in 3..args.len() {
let id = args
.int(at)
.map_err(|_| Error::new(Code::Invalid, "Invalid client ID"))?;
yo_alloc::allow(|| ids.push(id as u64));
}
} else if args.len() != 2 {
return Err(args::syntax());
}
let now = server.now_ms();
let text = yo_alloc::allow(|| {
let mut text = String::new();
for row in server.client_rows() {
if !want.covers(&row) || (!ids.is_empty() && !ids.contains(&row.id)) {
continue;
}
let obl = row.obl.load(Relaxed);
let resp = row.resp.load(Relaxed);
line(&row, now, obl, resp, &mut text);
}
text
});
out.verbatim(b"txt", text.as_bytes());
Ok(())
}
#[derive(Default)]
struct Filter<'a> {
id: Option<u64>,
addr: Option<&'a [u8]>,
laddr: Option<&'a [u8]>,
kind: Option<Kind>,
maxage: Option<u64>,
skipme: bool,
}
impl Filter<'_> {
fn covers(&self, row: &Client, now: u64, me: u64) -> bool {
if self.skipme && row.id == me {
return false;
}
if self.id.is_some_and(|id| id != row.id) {
return false;
}
if self.kind.is_some_and(|kind| !kind.covers(row)) {
return false;
}
if let Some(want) = self.maxage {
let since = row.since_ms.load(Relaxed);
let age = if since == 0 {
0
} else {
now.saturating_sub(since) / 1000
};
if age < want {
return false;
}
}
if self.addr.is_none() && self.laddr.is_none() {
return true;
}
let text = row.text.lock();
self.addr.is_none_or(|want| want == text.peer.as_slice())
&& self.laddr.is_none_or(|want| want == text.local.as_slice())
}
}
fn kill(server: &Server, session: &mut Session, args: Args<'_>, out: &mut Out) -> Result<bool> {
let old = args.len() == 3;
let mut filter = Filter {
skipme: !old,
..Filter::default()
};
if old {
filter.addr = Some(args.get(2));
} else if args.len() > 3 {
let mut at = 2;
while at + 1 < args.len() {
let word = args.get(at);
let value = args.get(at + 1);
if is(word, b"ID") {
filter.id = Some(client_id(value)?);
} else if is(word, b"ADDR") {
filter.addr = Some(value);
} else if is(word, b"LADDR") {
filter.laddr = Some(value);
} else if is(word, b"TYPE") {
filter.kind = Some(Kind::parse(value)?);
} else if is(word, b"USER") {
if !is(value, b"default") {
return Err(Error::fmt(
Code::Invalid,
format_args!("No such user '{}'", String::from_utf8_lossy(value)),
));
}
} else if is(word, b"MAXAGE") {
filter.maxage = Some(maxage(value)?);
} else if is(word, b"SKIPME") {
filter.skipme = if is(value, b"yes") {
true
} else if is(value, b"no") {
false
} else {
return Err(args::syntax());
};
} else {
return Err(args::syntax());
}
at += 2;
}
if at != args.len() {
return Err(args::syntax());
}
} else {
return Err(args::wrong_arity_sub("client", "kill"));
}
let now = server.now_ms();
let me = session.id;
let mut killed = 0u64;
let mut myself = false;
let mut posted = 0;
for row in server.client_rows() {
if !filter.covers(&row, now, me) {
continue;
}
killed += 1;
if row.id == me {
myself = true;
} else if row.kill() {
posted += 1;
}
}
server.note_kills(posted);
if old {
if killed == 0 {
return Err(Error::new(Code::Invalid, "No such client"));
}
out.ok();
} else {
out.uint(killed);
}
Ok(myself)
}
fn client_id(value: &[u8]) -> Result<u64> {
let text = core::str::from_utf8(value).ok();
let id = text.and_then(|t| t.parse::<i64>().ok());
match id {
Some(id) if id > 0 => Ok(id as u64),
_ => Err(Error::new(
Code::Invalid,
"client-id should be greater than 0",
)),
}
}
fn maxage(value: &[u8]) -> Result<u64> {
let text = core::str::from_utf8(value).ok();
let Some(age) = text.and_then(|t| t.parse::<i64>().ok()) else {
return Err(Error::new(
Code::Invalid,
"maxage is not an integer or out of range",
));
};
if age <= 0 {
return Err(Error::new(Code::Invalid, "maxage should be greater than 0"));
}
Ok(age as u64)
}
fn pause(server: &Server, args: Args<'_>) -> Result<()> {
if args.len() < 3 {
return Err(args::wrong_arity_sub("client", "pause"));
}
if args.len() > 4 {
return Err(args::subcommand_syntax(args.get(1), "CLIENT"));
}
let text = core::str::from_utf8(args.get(2)).ok();
let Some(ms) = text.and_then(|t| t.parse::<i64>().ok()) else {
return Err(Error::new(
Code::Invalid,
"timeout is not an integer or out of range",
));
};
if ms < 0 {
return Err(Error::new(Code::Invalid, "timeout is negative"));
}
let all = if args.len() == 3 {
true
} else {
let mode = args.get(3);
if is(mode, b"ALL") {
true
} else if is(mode, b"WRITE") {
false
} else {
return Err(Error::new(
Code::Invalid,
"CLIENT PAUSE mode must be WRITE or ALL",
));
}
};
server.pause(server.now_ms().saturating_add(ms as u64), all);
Ok(())
}
fn line(row: &Client, now_ms: u64, obl: u64, resp: u32, into: &mut String) {
let since = row.since_ms.load(Relaxed);
let now = if since == 0 { 0 } else { now_ms };
let qbuf = row.qbuf.load(Relaxed);
let qbuf_free = row.qbuf_free.load(Relaxed);
let rbs = row.rbs.load(Relaxed);
let tot_mem = qbuf + qbuf_free + rbs;
let named = row.has_sub.load(Acquire) == 1;
let text = row.text.lock();
let _ = writeln!(
into,
"id={id} addr={addr} laddr={laddr} fd={fd} name={name} age={age} idle={idle} \
flags={flags} db={db} sub={sub} psub={psub} ssub={ssub} multi={multi} watch={watch} \
qbuf={qbuf} qbuf-free={qbuf_free} argv-mem={argv_mem} multi-mem={multi_mem} \
rbs={rbs} rbp={rbp} obl={obl} oll=0 omem=0 omem-shared=0 omem-unshared=0 \
tot-mem={tot_mem} events=r cmd={cmd} user=default redir=-1 resp={resp} \
lib-name={lib_name} lib-ver={lib_ver} io-thread={io_thread} tot-net-in={net_in} \
tot-net-out={net_out} tot-cmds={cmds} read-events={reads} \
avg-pipeline-len-sum={cmds} avg-pipeline-len-cnt={reads}",
id = row.id,
addr = String::from_utf8_lossy(&text.peer),
laddr = String::from_utf8_lossy(&text.local),
fd = row.fd.load(Relaxed),
name = String::from_utf8_lossy(&text.name),
age = (now.saturating_sub(since)) / 1000,
idle = (now.saturating_sub(row.last_ms.load(Relaxed))) / 1000,
flags = Flags(row.flags.load(Relaxed)),
db = row.db.load(Relaxed),
sub = row.sub.load(Relaxed),
psub = row.psub.load(Relaxed),
ssub = row.ssub.load(Relaxed),
multi = row.multi.load(Relaxed),
watch = row.watch.load(Relaxed),
argv_mem = row.argv_mem.load(Relaxed),
multi_mem = row.multi_mem.load(Relaxed),
rbp = row.rbp.load(Relaxed),
cmd = Named(row.spec.load(Relaxed), named.then_some(&text.sub)),
lib_name = String::from_utf8_lossy(&text.lib_name),
lib_ver = String::from_utf8_lossy(&text.lib_ver),
io_thread = row.thread.load(Relaxed),
net_in = row.net_in.load(Relaxed),
net_out = row.net_out.load(Relaxed),
cmds = row.cmds.load(Relaxed),
reads = row.reads.load(Relaxed),
);
}
fn report(server: &Server, session: &Session, proto: Proto, mark: usize) -> String {
yo_alloc::allow(|| {
let mut s = String::with_capacity(512);
line(
session.row(),
server.now_ms(),
mark as u64,
proto.version() as u32,
&mut s,
);
s
})
}
struct Flags(u32);
impl core::fmt::Display for Flags {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let letters = [
(clients::SUBSCRIBED, 'P'),
(clients::IN_MULTI, 'x'),
(clients::UNIX, 'U'),
(clients::NO_EVICT, 'e'),
(clients::NO_TOUCH, 'T'),
];
let mut wrote = false;
for (bit, letter) in letters {
if self.0 & bit != 0 {
f.write_char(letter)?;
wrote = true;
}
}
if wrote { Ok(()) } else { f.write_char('N') }
}
}
struct Named<'a>(u32, Option<&'a Vec<u8>>);
impl core::fmt::Display for Named<'_> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let Ok(at) = usize::try_from(self.0) else {
return f.write_str("NULL");
};
if at >= table::count() {
return f.write_str("NULL");
}
f.write_str(table::name_at(at))?;
if let Some(sub) = self.1.filter(|sub| !sub.is_empty()) {
f.write_str("|")?;
for b in sub {
f.write_char(b.to_ascii_lowercase() as char)?;
}
}
Ok(())
}
}
const CLIENT_HELP: &[&str] = &[
"CLIENT <subcommand> [<arg> [value] [opt] ...]. Subcommands are:",
"GETNAME",
" Return the name of the current connection.",
"ID",
" Return the ID of the current connection.",
"INFO",
" Return information about the current client connection.",
"KILL <ip:port>",
" Close the connection from the specified address and port.",
"KILL <option> <value> [<option> <value> [...]]",
" Kill connections. Options are:",
" * ADDR (<ip:port>|<unixsocket>:0)",
" Kill connections made from the specified address",
" * LADDR (<ip:port>|<unixsocket>:0)",
" Kill connections made to specified local address",
" * TYPE (NORMAL|PUBSUB|MASTER|REPLICA)",
" Kill connections by type.",
" * USER <username>",
" Kill connections authenticated by <username>.",
" * SKIPME (YES|NO)",
" Skip killing current connection (default: yes).",
" * ID <client-id>",
" Kill connections by client id.",
" * MAXAGE <maxage>",
" Kill connections older than the specified age.",
"LIST [options ...]",
" Return information about client connections. Options:",
" * TYPE (NORMAL|PUBSUB|MASTER|REPLICA)",
" Return clients of specified type.",
"NO-EVICT (ON|OFF)",
" Protect current client connection from eviction.",
"NO-TOUCH (ON|OFF)",
" Will not touch LRU/LFU stats when this mode is on.",
"UNPAUSE",
" Stop the current client pause, resuming traffic.",
"PAUSE <timeout> [WRITE|ALL]",
" Suspend all, or just write, clients for <timeout> milliseconds.",
"REPLY (ON|OFF|SKIP)",
" Control the replies sent to the current connection.",
"SETINFO <option> <value>",
" Set client meta attr. Options are:",
" * LIB-NAME: the client lib name.",
" * LIB-VER: the client lib version.",
"SETNAME <name>",
" Assign the name <name> to the current connection.",
"HELP",
" Print this help.",
];