opendal/services/memcached/
binary.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use tokio::io::AsyncReadExt;
19use tokio::io::AsyncWriteExt;
20use tokio::io::BufReader;
21use tokio::io::{self};
22use tokio::net::TcpStream;
23
24use crate::raw::*;
25use crate::*;
26
27pub(super) mod constants {
28    pub const OK_STATUS: u16 = 0x0;
29    pub const KEY_NOT_FOUND: u16 = 0x1;
30}
31
32pub enum Opcode {
33    Get = 0x00,
34    Set = 0x01,
35    Delete = 0x04,
36    Version = 0x0b,
37    StartAuth = 0x21,
38}
39
40pub enum Magic {
41    Request = 0x80,
42}
43
44#[derive(Debug)]
45pub struct StoreExtras {
46    pub flags: u32,
47    pub expiration: u32,
48}
49
50#[derive(Debug, Default)]
51pub struct PacketHeader {
52    pub magic: u8,
53    pub opcode: u8,
54    pub key_length: u16,
55    pub extras_length: u8,
56    pub data_type: u8,
57    pub vbucket_id_or_status: u16,
58    pub total_body_length: u32,
59    pub opaque: u32,
60    pub cas: u64,
61}
62
63impl PacketHeader {
64    pub async fn write(self, writer: &mut TcpStream) -> io::Result<()> {
65        writer.write_u8(self.magic).await?;
66        writer.write_u8(self.opcode).await?;
67        writer.write_u16(self.key_length).await?;
68        writer.write_u8(self.extras_length).await?;
69        writer.write_u8(self.data_type).await?;
70        writer.write_u16(self.vbucket_id_or_status).await?;
71        writer.write_u32(self.total_body_length).await?;
72        writer.write_u32(self.opaque).await?;
73        writer.write_u64(self.cas).await?;
74        Ok(())
75    }
76
77    pub async fn read(reader: &mut TcpStream) -> Result<PacketHeader, io::Error> {
78        let header = PacketHeader {
79            magic: reader.read_u8().await?,
80            opcode: reader.read_u8().await?,
81            key_length: reader.read_u16().await?,
82            extras_length: reader.read_u8().await?,
83            data_type: reader.read_u8().await?,
84            vbucket_id_or_status: reader.read_u16().await?,
85            total_body_length: reader.read_u32().await?,
86            opaque: reader.read_u32().await?,
87            cas: reader.read_u64().await?,
88        };
89        Ok(header)
90    }
91}
92
93pub struct Response {
94    header: PacketHeader,
95    _key: Vec<u8>,
96    _extras: Vec<u8>,
97    value: Vec<u8>,
98}
99
100pub struct Connection {
101    io: BufReader<TcpStream>,
102}
103
104impl Connection {
105    pub fn new(io: TcpStream) -> Self {
106        Self {
107            io: BufReader::new(io),
108        }
109    }
110
111    pub async fn auth(&mut self, username: &str, password: &str) -> Result<()> {
112        let writer = self.io.get_mut();
113        let key = "PLAIN";
114        let request_header = PacketHeader {
115            magic: Magic::Request as u8,
116            opcode: Opcode::StartAuth as u8,
117            key_length: key.len() as u16,
118            total_body_length: (key.len() + username.len() + password.len() + 2) as u32,
119            ..Default::default()
120        };
121        request_header
122            .write(writer)
123            .await
124            .map_err(new_std_io_error)?;
125        writer
126            .write_all(key.as_bytes())
127            .await
128            .map_err(new_std_io_error)?;
129        writer
130            .write_all(format!("\x00{}\x00{}", username, password).as_bytes())
131            .await
132            .map_err(new_std_io_error)?;
133        writer.flush().await.map_err(new_std_io_error)?;
134        parse_response(writer).await?;
135        Ok(())
136    }
137
138    pub async fn version(&mut self) -> Result<String> {
139        let writer = self.io.get_mut();
140        let request_header = PacketHeader {
141            magic: Magic::Request as u8,
142            opcode: Opcode::Version as u8,
143            ..Default::default()
144        };
145        request_header
146            .write(writer)
147            .await
148            .map_err(new_std_io_error)?;
149        writer.flush().await.map_err(new_std_io_error)?;
150        let response = parse_response(writer).await?;
151        let version = String::from_utf8(response.value);
152        match version {
153            Ok(version) => Ok(version),
154            Err(e) => {
155                Err(Error::new(ErrorKind::Unexpected, "unexpected data received").set_source(e))
156            }
157        }
158    }
159
160    pub async fn get(&mut self, key: &str) -> Result<Option<Vec<u8>>> {
161        let writer = self.io.get_mut();
162        let request_header = PacketHeader {
163            magic: Magic::Request as u8,
164            opcode: Opcode::Get as u8,
165            key_length: key.len() as u16,
166            total_body_length: key.len() as u32,
167            ..Default::default()
168        };
169        request_header
170            .write(writer)
171            .await
172            .map_err(new_std_io_error)?;
173        writer
174            .write_all(key.as_bytes())
175            .await
176            .map_err(new_std_io_error)?;
177        writer.flush().await.map_err(new_std_io_error)?;
178        match parse_response(writer).await {
179            Ok(response) => {
180                if response.header.vbucket_id_or_status == 0x1 {
181                    return Ok(None);
182                }
183                Ok(Some(response.value))
184            }
185            Err(e) => Err(e),
186        }
187    }
188
189    pub async fn set(&mut self, key: &str, val: &[u8], expiration: u32) -> Result<()> {
190        let writer = self.io.get_mut();
191        let request_header = PacketHeader {
192            magic: Magic::Request as u8,
193            opcode: Opcode::Set as u8,
194            key_length: key.len() as u16,
195            extras_length: 8,
196            total_body_length: (8 + key.len() + val.len()) as u32,
197            ..Default::default()
198        };
199        let extras = StoreExtras {
200            flags: 0,
201            expiration,
202        };
203        request_header
204            .write(writer)
205            .await
206            .map_err(new_std_io_error)?;
207        writer
208            .write_u32(extras.flags)
209            .await
210            .map_err(new_std_io_error)?;
211        writer
212            .write_u32(extras.expiration)
213            .await
214            .map_err(new_std_io_error)?;
215        writer
216            .write_all(key.as_bytes())
217            .await
218            .map_err(new_std_io_error)?;
219        writer.write_all(val).await.map_err(new_std_io_error)?;
220        writer.flush().await.map_err(new_std_io_error)?;
221
222        parse_response(writer).await?;
223        Ok(())
224    }
225
226    pub async fn delete(&mut self, key: &str) -> Result<()> {
227        let writer = self.io.get_mut();
228        let request_header = PacketHeader {
229            magic: Magic::Request as u8,
230            opcode: Opcode::Delete as u8,
231            key_length: key.len() as u16,
232            total_body_length: key.len() as u32,
233            ..Default::default()
234        };
235        request_header
236            .write(writer)
237            .await
238            .map_err(new_std_io_error)?;
239        writer
240            .write_all(key.as_bytes())
241            .await
242            .map_err(new_std_io_error)?;
243        writer.flush().await.map_err(new_std_io_error)?;
244        parse_response(writer).await?;
245        Ok(())
246    }
247}
248
249pub async fn parse_response(reader: &mut TcpStream) -> Result<Response> {
250    let header = PacketHeader::read(reader).await.map_err(new_std_io_error)?;
251
252    if header.vbucket_id_or_status != constants::OK_STATUS
253        && header.vbucket_id_or_status != constants::KEY_NOT_FOUND
254    {
255        return Err(
256            Error::new(ErrorKind::Unexpected, "unexpected status received")
257                .with_context("message", format!("{}", header.vbucket_id_or_status)),
258        );
259    }
260
261    let mut extras = vec![0x0; header.extras_length as usize];
262    reader
263        .read_exact(extras.as_mut_slice())
264        .await
265        .map_err(new_std_io_error)?;
266
267    let mut key = vec![0x0; header.key_length as usize];
268    reader
269        .read_exact(key.as_mut_slice())
270        .await
271        .map_err(new_std_io_error)?;
272
273    let mut value = vec![
274        0x0;
275        (header.total_body_length - u32::from(header.key_length) - u32::from(header.extras_length))
276            as usize
277    ];
278    reader
279        .read_exact(value.as_mut_slice())
280        .await
281        .map_err(new_std_io_error)?;
282
283    Ok(Response {
284        header,
285        _key: key,
286        _extras: extras,
287        value,
288    })
289}
opendal/services/memcached/binary.rs

opendal/services/memcached/
binary.rs
