#![doc = include_str!("../README.md")]
#![deny(missing_docs)]
#[cfg(feature = "tokio")]
pub mod tokio;
#[cfg(feature = "async-std")]
pub mod async_std;
use std::{
fs::File,
io::{self, Error, Read, Write},
net::{TcpStream, ToSocketAddrs},
path::Path,
str::{self, Utf8Error},
};
#[cfg(unix)]
use std::os::unix::net::UnixStream;
pub type IoResult = Result<Vec<u8>, Error>;
pub type Utf8Result = Result<bool, Utf8Error>;
const DEFAULT_CHUNK_SIZE: usize = 4096;
const PING: &[u8; 6] = b"zPING\0";
const VERSION: &[u8; 9] = b"zVERSION\0";
const SHUTDOWN: &[u8; 10] = b"zSHUTDOWN\0";
const INSTREAM: &[u8; 10] = b"zINSTREAM\0";
const END_OF_STREAM: &[u8; 4] = &[0, 0, 0, 0];
pub const PONG: &[u8; 5] = b"PONG\0";
fn send_command<RW: Read + Write>(
mut stream: RW,
command: &[u8],
expected_response_length: Option<usize>,
) -> IoResult {
stream.write_all(command)?;
stream.flush()?;
let mut response = match expected_response_length {
Some(len) => Vec::with_capacity(len),
None => Vec::new(),
};
stream.read_to_end(&mut response)?;
Ok(response)
}
fn _ping<RW: Read + Write>(stream: RW) -> IoResult {
send_command(stream, PING, Some(PONG.len()))
}
fn _get_version<RW: Read + Write>(stream: RW) -> IoResult {
send_command(stream, VERSION, None)
}
fn scan<R: Read, RW: Read + Write>(
mut input: R,
chunk_size: Option<usize>,
mut stream: RW,
) -> IoResult {
stream.write_all(INSTREAM)?;
let chunk_size = chunk_size
.unwrap_or(DEFAULT_CHUNK_SIZE)
.min(u32::MAX as usize);
let mut buffer = vec![0; chunk_size];
loop {
let len = input.read(&mut buffer[..])?;
if len != 0 {
stream.write_all(&(len as u32).to_be_bytes())?;
stream.write_all(&buffer[..len])?;
} else {
stream.write_all(END_OF_STREAM)?;
stream.flush()?;
break;
}
}
let mut response = Vec::new();
stream.read_to_end(&mut response)?;
Ok(response)
}
pub fn clean(response: &[u8]) -> Utf8Result {
let response = str::from_utf8(response)?;
Ok(response.contains("OK") && !response.contains("FOUND"))
}
#[derive(Copy, Clone)]
pub struct Tcp<A: ToSocketAddrs> {
pub host_address: A,
}
#[derive(Copy, Clone)]
#[cfg(unix)]
pub struct Socket<P: AsRef<Path>> {
pub socket_path: P,
}
pub trait TransportProtocol {
type Stream: Read + Write;
fn connect(&self) -> io::Result<Self::Stream>;
}
impl<A: ToSocketAddrs> TransportProtocol for Tcp<A> {
type Stream = TcpStream;
fn connect(&self) -> io::Result<Self::Stream> {
TcpStream::connect(&self.host_address)
}
}
#[cfg(unix)]
impl<P: AsRef<Path>> TransportProtocol for Socket<P> {
type Stream = UnixStream;
fn connect(&self) -> io::Result<Self::Stream> {
UnixStream::connect(&self.socket_path)
}
}
pub fn ping<T: TransportProtocol>(connection: T) -> IoResult {
let stream = connection.connect()?;
_ping(stream)
}
pub fn get_version<T: TransportProtocol>(connection: T) -> IoResult {
let stream = connection.connect()?;
_get_version(stream)
}
pub fn scan_file<P: AsRef<Path>, T: TransportProtocol>(
file_path: P,
connection: T,
chunk_size: Option<usize>,
) -> IoResult {
let file = File::open(file_path)?;
let stream = connection.connect()?;
scan(file, chunk_size, stream)
}
pub fn scan_buffer<T: TransportProtocol>(
buffer: &[u8],
connection: T,
chunk_size: Option<usize>,
) -> IoResult {
let stream = connection.connect()?;
scan(buffer, chunk_size, stream)
}
pub fn shutdown<T: TransportProtocol>(connection: T) -> IoResult {
let stream = connection.connect()?;
send_command(stream, SHUTDOWN, None)
}