1use crate::Commands;
8use crate::conn::Conn;
9use crate::message::Part;
10use crate::shard::Shard;
11use kevy_resp::ArgvView;
12
13#[derive(Debug, Clone, PartialEq, Eq)]
16pub enum ClientKillFilter {
17 Addr(Vec<u8>),
19 Id(u64),
21}
22
23impl ClientKillFilter {
24 pub fn parse<A: ArgvView + ?Sized>(args: &A) -> Option<(Self, bool)> {
30 match args.len() {
31 3 => {
32 let a = args.get(2)?;
33 a.contains(&b':').then(|| (Self::Addr(a.to_vec()), true))
34 }
35 4 => {
36 let kind = args.get(2)?.to_ascii_uppercase();
37 let val = args.get(3)?;
38 match kind.as_slice() {
39 b"ID" => {
40 std::str::from_utf8(val).ok()?.parse().ok().map(|id| (Self::Id(id), false))
41 }
42 b"ADDR" => Some((Self::Addr(val.to_vec()), false)),
43 _ => None,
44 }
45 }
46 _ => None,
47 }
48 }
49}
50
51pub(crate) fn client_row(id: u64, conn: &Conn, out: &mut Vec<u8>) {
56 use std::fmt::Write as _;
57 let mut s = String::with_capacity(224);
58 let _ = writeln!(
59 s,
60 "id={id} addr={}:{} laddr=0.0.0.0:0 fd={} name={} age={} idle=0 \
61 flags=N db=0 sub={} psub={} ssub=0 multi={} watch={} qbuf={} \
62 qbuf-free=0 argv-mem=0 multi-mem=0 tot-mem=0 rbs=0 rbp=0 obl={} \
63 oll=0 omem=0 events=r cmd=NULL user=default redir=-1 resp={} \
64 lib-name= lib-ver=",
65 conn.peer.0,
66 conn.peer.1,
67 conn.sock.raw(),
68 String::from_utf8_lossy(&conn.client_name),
69 conn.created.elapsed().as_secs(),
70 conn.sub.len(),
71 conn.psub.len(),
72 conn.multi.as_ref().map_or(-1, |q| q.len() as i64),
73 conn.watched.len(),
74 conn.input.len(),
75 conn.output.len().saturating_sub(conn.write_pos),
76 match conn.proto {
77 kevy_resp::RespVersion::V2 => 2,
78 kevy_resp::RespVersion::V3 => 3,
79 },
80 );
81 out.extend_from_slice(s.as_bytes());
82}
83
84impl<C: Commands> Shard<C> {
85 pub(crate) fn exec_client_list(&mut self) -> Part {
88 let mut text = Vec::with_capacity(self.conns.len() * 192);
89 for (id, conn) in &self.conns {
90 if conn.cluster {
91 continue;
92 }
93 client_row(*id, conn, &mut text);
94 }
95 Part::ExtensionChunk(text)
96 }
97
98 pub(crate) fn exec_client_kill(&mut self, filter: &ClientKillFilter) -> Part {
104 let mut victims: Vec<u64> = Vec::new();
105 for (id, conn) in &self.conns {
106 if conn.cluster || conn.closing {
107 continue;
108 }
109 let hit = match filter {
110 ClientKillFilter::Id(want) => *id == *want,
111 ClientKillFilter::Addr(addr) => {
112 format!("{}:{}", conn.peer.0, conn.peer.1).as_bytes() == addr.as_slice()
113 }
114 };
115 if hit {
116 victims.push(*id);
117 }
118 }
119 for id in &victims {
120 if let Some(conn) = self.conns.get_mut(id) {
121 conn.closing = true;
122 }
123 self.dirty.push(*id);
124 self.closing_uring_conns.push(*id);
125 self.blocked.drop_for_conn(*id);
132 self.cancel_xshard_on_close(*id);
133 }
134 Part::Int(victims.len() as i64)
135 }
136}
137
138#[cfg(test)]
139mod tests {
140 use super::ClientKillFilter;
141 use kevy_resp::Argv;
142
143 fn argv(parts: &[&[u8]]) -> Argv {
144 let mut a = Argv::default();
145 for p in parts {
146 a.push(p);
147 }
148 a
149 }
150
151 #[test]
152 fn parse_legacy_addr_form() {
153 let a = argv(&[b"CLIENT", b"KILL", b"127.0.0.1:50123"]);
154 assert_eq!(
155 ClientKillFilter::parse(&a),
156 Some((ClientKillFilter::Addr(b"127.0.0.1:50123".to_vec()), true))
157 );
158 }
159
160 #[test]
161 fn parse_id_and_addr_filters() {
162 let a = argv(&[b"CLIENT", b"KILL", b"ID", b"42"]);
163 assert_eq!(ClientKillFilter::parse(&a), Some((ClientKillFilter::Id(42), false)));
164 let a = argv(&[b"CLIENT", b"KILL", b"addr", b"10.0.0.1:1"]);
165 assert_eq!(
166 ClientKillFilter::parse(&a),
167 Some((ClientKillFilter::Addr(b"10.0.0.1:1".to_vec()), false))
168 );
169 }
170
171 #[test]
172 fn parse_rejects_unsupported_shapes() {
173 assert_eq!(ClientKillFilter::parse(&argv(&[b"CLIENT", b"KILL"])), None);
174 assert_eq!(ClientKillFilter::parse(&argv(&[b"CLIENT", b"KILL", b"noport"])), None);
175 assert_eq!(
176 ClientKillFilter::parse(&argv(&[b"CLIENT", b"KILL", b"LADDR", b"1.2.3.4:5"])),
177 None
178 );
179 assert_eq!(ClientKillFilter::parse(&argv(&[b"CLIENT", b"KILL", b"ID", b"notanum"])), None);
180 }
181}