Comprender la concurrencia y el paralelismo
Uno de los fundadores de Golang tiene una explicación muy reveladora e intuitiva: la concurrencia es la capacidad de procesar muchas cosas al mismo tiempo, y el paralelismo es un medio para ejecutar muchas cosas al mismo tiempo.
Procesamos las cosas que necesitamos hacer en múltiples subprocesos o múltiples tareas asincrónicas, esta es la capacidad de concurrencia. Ejecutar estos subprocesos o tareas asincrónicas simultáneamente en una máquina de múltiples núcleos y múltiples CPU es un método paralelo. Se puede decir que la concurrencia potencia el paralelismo. Cuando tenemos la capacidad de concurrencia, el paralelismo es algo natural.
En el proceso de lidiar con la concurrencia, la dificultad no es cómo crear múltiples subprocesos para distribuir el trabajo, sino cómo sincronizar estas tareas concurrentes . Echemos un vistazo a varios modos de trabajo comunes en el estado concurrente: modo de libre competencia, modo de mapa/reducción y modo DAG: el modo de
mapa/reducción divide el trabajo y, una vez completado el mismo procesamiento, los resultados se organizan en un cierto orden Levántese; modo DAG, divida el trabajo en subtareas independientes y dependientes, y luego ejecútelas al mismo tiempo de acuerdo con las dependencias. (Aquí puede consultar cómo se maneja la concurrencia en la programación concurrente de C++)
Detrás de estos modos de concurrencia, ¿qué primitivas de concurrencia podemos usar? Estas dos conferencias se centrarán en explicar y profundizar cinco conceptos: modelos Atomic, Mutex, Condvar, Channel y Actor. . Hoy hablaremos de los dos primeros Atomic y Mutex.
Atomic es la base de todas las primitivas de concurrencia y sienta una base sólida para la sincronización de tareas concurrentes. Detrás está el principio CAS:
la garantía más básica es: puede leer una dirección de memoria a través de una instrucción para determinar si su valor es igual a un determinado valor de prefijo, si es igual, modifíquelo a un nuevo valor . Esta es la operación de comparar e intercambiar, o CAS para abreviar. Es la piedra angular de casi todas las primitivas de concurrencia del sistema operativo, lo que nos permite implementar un bloqueo que funcione correctamente. compare_exchange es una operación CAS proporcionada por Rust, que se compilará en la instrucción CAS correspondiente de la CPU.
Pero para 3 y 4 relacionados con la optimización automática del compilador/CPU, todavía necesitamos algún procesamiento adicional. Estos son los dos parámetros extraños adicionales relacionados con el pedido en esta función. Esto también puede estar relacionado con la clasificación de declaraciones en C++.
pub enum Ordering {
Relaxed,
Release,
Acquire,
AcqRel,
SeqCst,
}
Lo que más uso personalmente es crear varias estructuras de datos sin bloqueos. Por ejemplo, se necesita un generador de ID global. Por supuesto, puede usar un módulo como UUID para generar una ID única, pero si también necesitamos ordenar la ID, entonces AtomicUsize es la mejor opción.
Mutex (mutex y bloqueo de giro) también puede contactar mutex y variables de condición para lograr el mecanismo de sincronización
Atomic. Aunque Atomic puede manejar la demanda de bloqueo en modo de libre competencia**, después de todo no es tan conveniente de usar. Necesitamos un nivel superior La primitiva de concurrencia** garantiza que el sistema de software controle el acceso de múltiples subprocesos al mismo recurso compartido, de modo que cada subproceso pueda tener acceso exclusivo o mutuo exclusivo al acceder al recurso compartido.
SpinLock, como sugiere el nombre, es un subproceso que espera un bloqueo disponible en una sección crítica a través de la CPU inactiva (giro, al igual que el bucle while anterior) y espera ocupada (espera ocupada). Sin embargo, esta implementación de exclusión mutua a través de SpinLock tiene limitaciones en los escenarios de uso: si la sección crítica protegida es demasiado grande, el rendimiento general disminuirá drásticamente, la CPU estará ocupada esperando, desperdiciando recursos y sin realizar trabajos prácticos , y no es adecuada. como método de procesamiento común.
Bloqueo Mutex: con el bloqueo Mutex, el hilo se programará cuando se espera el bloqueo y luego se programará nuevamente cuando el bloqueo esté disponible.
Parece que SpinLock es muy ineficiente, pero no lo es. Depende del tamaño de la sección crítica de la cerradura. Si hay muy poco código para ejecutar en la sección crítica, SpinLock vale la pena en comparación con el cambio de contexto causado por el bloqueo Mutex. En el Kernel de Linux muchas veces sólo podemos usar SpinLock.
Atomic/Mutex resuelve el problema de sincronización de tareas concurrentes en modo de libre competencia, y también puede resolver bien el problema de sincronización en modo mapa/reducción, porque la sincronización solo ocurre en el mapa y en las etapas de reducción.
Sin embargo, no resuelven un problema de nivel superior que es el patrón DAG: ¿Qué debemos hacer si este tipo de acceso necesita realizarse en un orden determinado o hay dependencias antes y después?
El escenario típico de este problema es el modelo productor-consumidor : después de que el productor produce el contenido, debe haber un mecanismo para notificar al consumidor que puede consumirlo. Por ejemplo, cuando hay datos en el socket, se notifica al subproceso de procesamiento para que procese los datos y, una vez completado el procesamiento, se notifica al subproceso de envío y recepción del socket para que envíe los datos.
Condvar debería ser similar a la variable de condición en C++. Tenga en cuenta la comparación,
por lo que el sistema operativo también proporciona Condvar. Condvar tiene dos estados: esperar: el hilo espera en la cola hasta que se cumpla una determinada condición. Notificar: cuando se cumplen las condiciones de condvar, el hilo actual notifica a otros hilos en espera que se pueden despertar. Las notificaciones pueden ser una sola notificación, múltiples notificaciones o incluso transmitidas (notificando a todos). En la práctica, Condvar se usa a menudo junto con Mutex: Mutex se usa para garantizar que las condiciones sean mutuamente excluyentes durante la lectura y la escritura, y Condvar se usa para controlar la espera y la activación de subprocesos. a nosotros
Channel
, pero será más difícil usar Mutex y Condvar para manejar el complejo modo de concurrencia DAG. Por lo tanto, Rust también proporciona una variedad de canales para manejar la comunicación entre tareas simultáneas. El canal encapsula el bloqueo en un área pequeña para la escritura y lectura en cola, y luego separa completamente a los lectores y escritores, lo que permite a los lectores leer datos y a los escritores escribir datos. Para los desarrolladores, además del posible cambio de contexto, Además, no tiene nada que que ver con las cerraduras, al igual que acceder a una cola local.
En comparación con Mutex, Channel tiene el nivel más alto de abstracción, la interfaz más intuitiva y la carga psicológica de usarlo no es tan grande. Al utilizar Mutex, se debe tener mucho cuidado para evitar puntos muertos, controlar el tamaño de las secciones críticas y prevenir posibles accidentes.
Al implementar Channel, se seleccionarán diferentes herramientas según los diferentes escenarios de uso. Rust proporciona los siguientes cuatro canales:
oneshot: este es probablemente el canal más simple: el escritor solo envía datos una vez y el lector solo los lee una vez. Esta sincronización única entre múltiples subprocesos se puede lograr utilizando un canal de una sola vez. Debido al propósito especial de oneshot, el intercambio atómico se puede utilizar directamente para implementarlo.
acotado: el canal acotado tiene una cola, pero la cola tiene un límite superior. Una vez que la cola está llena, el escritor también debe suspenderse y esperar. Cuando se produce el bloqueo, una vez que el lector lee los datos, el canal utilizará internamente notify_one de Condvar para notificar al escritor y despertarlo para que pueda continuar escribiendo.
ilimitado: la cola no tiene límite superior, si está llena, se expandirá automáticamente. Sabemos que muchas de las estructuras de datos de Rust, como Vec y VecDeque, se expanden automáticamente. En comparación con acotado, excepto que no bloquea a los escritores, otras implementaciones son muy similares.
Para todos estos tipos de canales, las ideas de implementación de síncrono y asíncrono son similares. La principal diferencia radica en los objetos que se suspenden/despertan. En el mundo sincrónico, los objetos que se suspenden/despertan son hilos; en el mundo asincrónico, son tareas con muy pequeña granularidad.
Práctica de etapa (4): cree un procesamiento de red de servidor KV simple
(para el análisis de protobuf, también puede comunicarse con el proyecto C ++).
Hemos estado usando una misteriosa biblioteca async-prost antes y mágicamente completamos la paquetización y descompresión de TCP. marcos. . La idea principal es agregar un encabezado para proporcionar la longitud del marco al serializar datos. Al deserializar, primero lea el encabezado, obtenga la longitud y luego lea los datos correspondientes. Nuestro desafío hoy es intentar manejar la lógica de paquetización y descompresión por nosotros mismos sin depender de async-prost
basado en el servidor KV que completamos la última vez . Si domina esta habilidad y coopera con protobuf, puede diseñar cualquier protocolo que pueda llevar a cabo negocios reales .
protobuf nos ayuda a resolver el problema de cómo definir mensajes de protocolo, pero cómo distinguir un mensaje de otro es un dolor de cabeza. Necesitamos definir delimitadores apropiados. El delimitador + datos del mensaje es un marco
(muchos protocolos basados en TCP usan \r\n como delimitador, como FTP; algunos usan la longitud del mensaje como delimitador, como gRPC; y algunos usan una combinación de los dos, como Redis RESP; para otros más complejos como HTTP, \r\n se usa para separar encabezados, \r\n\r\n se usa entre encabezado/cuerpo, y la longitud del cuerpo se proporcionará en el encabezado. "\r\ n" así El delimitador es adecuado para mensajes de protocolo que son datos ASCII; y la separación por longitud es adecuada para mensajes de protocolo que son datos binarios. El protobuf que lleva nuestro servidor KV es binario, por lo que ponemos una longitud antes de la carga útil como un separador de cuadros. .)
Tokio tiene una biblioteca tokio-util, que nos ha ayudado a lidiar con los requisitos principales para el desempaquetado de paquetes relacionados con cuadros, incluyendo LinesDelimited (procesamiento de delimitadores \r\n) y LongitudDelimitado (procesamiento de delimitadores de longitud) para permitir el flujo mut. =
Enmarcado ::new(flujo, LongitudDelimitedCodec::new());
(¿ Por qué necesita diseñarlo usted mismo? Porque las necesidades reales son modificables. No es solo el delimitador el que determina la longitud. Por ejemplo, ¿también puede personalizar si se necesita compresión? ¿Se requiere otro procesamiento especial? El uso de la biblioteca El código es limitado porque su interfaz proporciona La función es fija )
Para estar más cerca de la realidad, tomamos el bit más alto de la longitud de 4 bytes como una señal de si comprimir o no. Si está configurado, significa que el la carga útil posterior es un protobuf comprimido con gzip; de lo contrario, es directamente protobuf: según
la convención, primero definamos el rasgo que maneja esta lógica:
pub trait FrameCoder
where
Self: Message + Sized + Default,
{
/// 把一个 Message encode 成一个 frame
fn encode_frame(&self, buf: &mut BytesMut) -> Result<(), KvError>;
/// 把一个完整的 frame decode 成一个 Message
fn decode_frame(buf: &mut BytesMut) -> Result<Self, KvError>;
}
Implementar rasgos
use std::io::{
Read, Write};
use crate::{
CommandRequest, CommandResponse, KvError};
use bytes::{
Buf, BufMut, BytesMut};
use flate2::{
read::GzDecoder, write::GzEncoder, Compression};
use prost::Message;
use tokio::io::{
AsyncRead, AsyncReadExt};
use tracing::debug;
/// 长度整个占用 4 个字节
pub const LEN_LEN: usize = 4;
/// 长度占 31 bit，所以最大的 frame 是 2G
const MAX_FRAME: usize = 2 * 1024 * 1024 * 1024;
/// 如果 payload 超过了 1436 字节，就做压缩
const COMPRESSION_LIMIT: usize = 1436;
/// 代表压缩的 bit（整个长度 4 字节的最高位）
const COMPRESSION_BIT: usize = 1 << 31;
/// 处理 Frame 的 encode/decode
pub trait FrameCoder
where
Self: Message + Sized + Default,
{
/// 把一个 Message encode 成一个 frame
fn encode_frame(&self, buf: &mut BytesMut) -> Result<(), KvError> {
let size = self.encoded_len();
if size >= MAX_FRAME {
return Err(KvError::FrameError);
}
// 我们先写入长度，如果需要压缩，再重写压缩后的长度
buf.put_u32(size as _);
if size > COMPRESSION_LIMIT {
let mut buf1 = Vec::with_capacity(size);
self.encode(&mut buf1)?;
// BytesMut 支持逻辑上的 split（之后还能 unsplit）
// 所以我们先把长度这 4 字节拿走，清除
let payload = buf.split_off(LEN_LEN);
buf.clear();
// 处理 gzip 压缩，具体可以参考 flate2 文档
let mut encoder = GzEncoder::new(payload.writer(), Compression::default());
encoder.write_all(&buf1[..])?;
// 压缩完成后，从 gzip encoder 中把 BytesMut 再拿回来
let payload = encoder.finish()?.into_inner();
debug!("Encode a frame: size {}({})", size, payload.len());
// 写入压缩后的长度
buf.put_u32((payload.len() | COMPRESSION_BIT) as _);
// 把 BytesMut 再合并回来
buf.unsplit(payload);
Ok(())
} else {
self.encode(buf)?;
Ok(())
}
}
/// 把一个完整的 frame decode 成一个 Message
fn decode_frame(buf: &mut BytesMut) -> Result<Self, KvError> {
// 先取 4 字节，从中拿出长度和 compression bit
let header = buf.get_u32() as usize;
let (len, compressed) = decode_header(header);
debug!("Got a frame: msg len {}, compressed {}", len, compressed);
if compressed {
// 解压缩
let mut decoder = GzDecoder::new(&buf[..len]);
let mut buf1 = Vec::with_capacity(len * 2);
decoder.read_to_end(&mut buf1)?;
buf.advance(len);
// decode 成相应的消息
Ok(Self::decode(&buf1[..buf1.len()])?)
} else {
let msg = Self::decode(&buf[..len])?;
buf.advance(len);
Ok(msg)
}
}
}
impl FrameCoder for CommandRequest {
}
impl FrameCoder for CommandResponse {
}
fn decode_header(header: usize) -> (usize, bool) {
let len = header & !COMPRESSION_BIT;
let compressed = header & COMPRESSION_BIT == COMPRESSION_BIT;
(len, compressed)
}
Si se pregunta por qué COMPRESSION_LIMIT está configurado en 1436.
Esto se debe a que la MTU de Ethernet es 1500. Después de excluir los 20 bytes del encabezado IP y los 20 bytes del encabezado TCP, todavía quedan 1460. Generalmente, los paquetes TCP contendrán algunas opciones (como la marca de tiempo) y los paquetes IP. también puede contenerlos, por lo que reservamos 20 Bytes; menos la longitud de 4 bytes, es 1436, la longitud máxima del mensaje sin fragmentación. Si es más grande que esto, es probable que cause fragmentación, por lo que simplemente lo comprimiremos.
Actualmente, este código no toca nada relacionado con el socket IO, es solo pura lógica, luego necesitamos conectarlo con el TcpStream que usamos para manejar el cliente del servidor. Hay algo de procesamiento en el medio que permite que la transmisión procese el cuadro, del cual no hablaré todavía.
El objetivo principal es permitir que la transmisión lea el fotograma completo, lo que implica algunas funciones de la biblioteca, por lo que no entraré en detalles.
stream.read_exact(&mut buf[LEN_LEN…]).await?;
A continuación, debemos pensar en cómo encapsular el servidor y el cliente.
En el lado del servidor, use el proceso para encapsular
#[tokio::main]
async fn main() -> Result<()> {
tracing_subscriber::fmt::init();
let addr = "127.0.0.1:9527";
let service: Service = ServiceInner::new(MemTable::new()).into();
let listener = TcpListener::bind(addr).await?;
info!("Start listening on {}", addr);
loop {
let (stream, addr) = listener.accept().await?;
info!("Client {:?} connected", addr);
let stream = ProstServerStream::new(stream, service.clone());
tokio::spawn(async move {
stream.process().await });
}
}
Este método Process() es en realidad una encapsulación del bucle while en tokio::spawn en ejemplos/server.rs:
while let Some(Ok(cmd)) = stream.next().await {
info!("Got a new command: {:?}", cmd);
let res = svc.execute(cmd);
stream.send(res).await.unwrap();
}
Para el cliente, también esperamos poder ejecutar directamente() un comando y obtener el resultado:
#[tokio::main]
async fn main() -> Result<()> {
tracing_subscriber::fmt::init();
let addr = "127.0.0.1:9527";
// 连接服务器
let stream = TcpStream::connect(addr).await?;
let mut client = ProstClientStream::new(stream);
// 生成一个 HSET 命令
let cmd = CommandRequest::new_hset("table1", "hello", "world".to_string().into());
// 发送 HSET 命令
let data = client.execute(cmd).await?;
info!("Got response {:?}", data);
Ok(())
}
Esta ejecución() es en realidad una encapsulación del código de envío y recepción en ejemplos/client.rs:
client.send(cmd).await?;
if let Some(Ok(data)) = client.next().await {
info!("Got response {:?}", data);
}
Bien, primero veamos la estructura de datos del servidor que procesa un TcpStream. Debe contener TcpStream, así como el Servicio que creamos antes para procesar los comandos del cliente. Por lo tanto, la estructura para que el servidor procese TcpStream contiene estas dos partes:
pub struct ProstServerStream<S> {
inner: S,
service: Service,
}
La estructura del cliente para procesar TcpStream solo necesita contener TcpStream:
pub struct ProstClientStream<S> {
inner: S,
}
Aquí todavía se utiliza el parámetro genérico S. **En el futuro, si queremos admitir WebSocket o TLS sobre TCP, nos ahorrará la necesidad de cambiar esta capa de código. **Esto también refleja los beneficios de los parámetros genéricos, al igual que el rasgo de la tienda anterior.
El siguiente paso es implementar el proceso y ejecutarlo específicamente.
mod frame;
use bytes::BytesMut;
pub use frame::{
read_frame, FrameCoder};
use tokio::io::{
AsyncRead, AsyncWrite, AsyncWriteExt};
use tracing::info;
use crate::{
CommandRequest, CommandResponse, KvError, Service};
/// 处理服务器端的某个 accept 下来的 socket 的读写
pub struct ProstServerStream<S> {
inner: S,
service: Service,
}
/// 处理客户端 socket 的读写
pub struct ProstClientStream<S> {
inner: S,
}
impl<S> ProstServerStream<S>
where
S: AsyncRead + AsyncWrite + Unpin + Send,
{
pub fn new(stream: S, service: Service) -> Self {
Self {
inner: stream,
service,
}
}
pub async fn process(mut self) -> Result<(), KvError> {
while let Ok(cmd) = self.recv().await {
info!("Got a new command: {:?}", cmd);
let res = self.service.execute(cmd);
self.send(res).await?;
}
// info!("Client {:?} disconnected", self.addr);
Ok(())
}
async fn send(&mut self, msg: CommandResponse) -> Result<(), KvError> {
let mut buf = BytesMut::new();
msg.encode_frame(&mut buf)?;
let encoded = buf.freeze();
self.inner.write_all(&encoded[..]).await?;
Ok(())
}
async fn recv(&mut self) -> Result<CommandRequest, KvError> {
let mut buf = BytesMut::new();
let stream = &mut self.inner;
read_frame(stream, &mut buf).await?;
CommandRequest::decode_frame(&mut buf)
}
}
impl<S> ProstClientStream<S>
where
S: AsyncRead + AsyncWrite + Unpin + Send,
{
pub fn new(stream: S) -> Self {
Self {
inner: stream }
}
pub async fn execute(&mut self, cmd: CommandRequest) -> Result<CommandResponse, KvError> {
self.send(cmd).await?;
Ok(self.recv().await?)
}
async fn send(&mut self, msg: CommandRequest) -> Result<(), KvError> {
let mut buf = BytesMut::new();
msg.encode_frame(&mut buf)?;
let encoded = buf.freeze();
self.inner.write_all(&encoded[..]).await?;
Ok(())
}
async fn recv(&mut self) -> Result<CommandResponse, KvError> {
let mut buf = BytesMut::new();
let stream = &mut self.inner;
read_frame(stream, &mut buf).await?;
CommandResponse::decode_frame(&mut buf)
}
}
Después de escribirlo, descubrí que los códigos del servidor y del cliente son más concisos: el proceso se usa para reemplazar el proceso de procesamiento del servidor y la ejecución se usa para reemplazar el proceso de comando de ejecución del cliente. Y también se utiliza un método de procesamiento de cuadros personalizado . Esta es la mejora de esta sección. La secuencia personalizada utiliza un parámetro genérico S, que elimina la necesidad de modificar el código al agregar nuevos tipos de protocolo en el futuro.
Etapa de operación práctica (5): construir un servidor KV simple: seguridad de red
Entonces, cuando la arquitectura de nuestra aplicación se basa en TCP, ¿cómo usar TLS para garantizar la seguridad entre el cliente y el servidor?
Para usar TLS, primero necesitamos un certificado x509. TLS requiere un certificado x509 para que el cliente verifique que el servidor es un servidor confiable, e incluso para que el servidor verifique al cliente para confirmar que la otra parte es un cliente confiable.
Para facilitar las pruebas, debemos tener la capacidad de generar nuestro propio certificado de CA, certificado de servidor e incluso certificado de cliente. Los detalles de la generación de certificados no se presentarán en detalle hoy. Anteriormente creé una biblioteca llamada certificar, que se puede usar para generar varios certificados. Podemos agregar esta biblioteca a Cargo.toml:
[dependencias de desarrollo]
…
certificar = “0.3”
…
Luego cree un directorio de accesorios en el directorio raíz para almacenar el certificado, cree el archivo ejemplos/gen_cert.rs y agregue el siguiente código:
use anyhow::Result;
use certify::{
generate_ca, generate_cert, load_ca, CertType, CA};
use tokio::fs;
struct CertPem {
cert_type: CertType,
cert: String,
key: String,
}
#[tokio::main]
async fn main() -> Result<()> {
let pem = create_ca()?;
gen_files(&pem).await?;
let ca = load_ca(&pem.cert, &pem.key)?;
let pem = create_cert(&ca, &["kvserver.acme.inc"], "Acme KV server", false)?;
gen_files(&pem).await?;
let pem = create_cert(&ca, &[], "awesome-device-id", true)?;
gen_files(&pem).await?;
Ok(())
}
fn create_ca() -> Result<CertPem> {
let (cert, key) = generate_ca(
&["acme.inc"],
"CN",
"Acme Inc.",
"Acme CA",
None,
Some(10 * 365),
)?;
Ok(CertPem {
cert_type: CertType::CA,
cert,
key,
})
}
fn create_cert(ca: &CA, domains: &[&str], cn: &str, is_client: bool) -> Result<CertPem> {
let (days, cert_type) = if is_client {
(Some(365), CertType::Client)
} else {
(Some(5 * 365), CertType::Server)
};
let (cert, key) = generate_cert(ca, domains, "CN", "Acme Inc.", cn, None, is_client, days)?;
Ok(CertPem {
cert_type,
cert,
key,
})
}
async fn gen_files(pem: &CertPem) -> Result<()> {
let name = match pem.cert_type {
CertType::Client => "client",
CertType::Server => "server",
CertType::CA => "ca",
};
fs::write(format!("fixtures/{}.cert", name), pem.cert.as_bytes()).await?;
fs::write(format!("fixtures/{}.key", name), pem.key.as_bytes()).await?;
Ok(())
}
Este código es muy simple: primero genera un certificado CA y luego genera certificados de servidor y cliente, todos los cuales se almacenan en el directorio de accesorios recién creado . Necesita cargar ejecutar --examples gen_cert para ejecutar este comando. Usaremos estos certificados y claves en la prueba más adelante.
No se detallarán los detalles específicos sobre TLS.
Para el servidor KV, después de usar TLS, la encapsulación de datos de todo el protocolo se muestra en la siguiente figura:
Se estima que muchas personas se adormecen cuando escuchan TLS o SSL, porque han tenido muchas malas experiencias con openssl antes. La base del código de openssl es demasiado compleja, la API no es amigable y la compilación y vinculación son muy difíciles. Sin embargo, la experiencia de usar TLS en Rust sigue siendo muy buena: Rust tiene una muy buena encapsulación de openssl y también hay Rustls escritos en Rust que no dependen de openssl. Tokio además proporciona soporte TLS en línea con el ecosistema de Tokio, con versiones openssl y Rustls disponibles.
Hoy usaremos tokio-rustls para escribir soporte TLS. Creo que podrá ver durante el proceso de implementación lo fácil que es agregar el protocolo TLS a la aplicación para proteger la capa de red.
Primero agrega tokio-rustls en Cargo.toml:
luego crea src/network/tls.rs y escribe el siguiente código (recuerda introducir este archivo en src/network/mod.rs):
use std::io::Cursor;
use std::sync::Arc;
use tokio::io::{
AsyncRead, AsyncWrite};
use tokio_rustls::rustls::{
internal::pemfile, Certificate, ClientConfig, ServerConfig};
use tokio_rustls::rustls::{
AllowAnyAuthenticatedClient, NoClientAuth, PrivateKey, RootCertStore};
use tokio_rustls::webpki::DNSNameRef;
use tokio_rustls::TlsConnector;
use tokio_rustls::{
client::TlsStream as ClientTlsStream, server::TlsStream as ServerTlsStream, TlsAcceptor,
};
use crate::KvError;
/// KV Server 自己的 ALPN (Application-Layer Protocol Negotiation)
const ALPN_KV: &str = "kv";
/// 存放 TLS ServerConfig 并提供方法 accept 把底层的协议转换成 TLS
#[derive(Clone)]
pub struct TlsServerAcceptor {
inner: Arc<ServerConfig>,
}
/// 存放 TLS Client 并提供方法 connect 把底层的协议转换成 TLS
#[derive(Clone)]
pub struct TlsClientConnector {
pub config: Arc<ClientConfig>,
pub domain: Arc<String>,
}
impl TlsClientConnector {
/// 加载 client cert / CA cert，生成 ClientConfig
pub fn new(
domain: impl Into<String>,
identity: Option<(&str, &str)>,
server_ca: Option<&str>,
) -> Result<Self, KvError> {
let mut config = ClientConfig::new();
// 如果有客户端证书，加载之
if let Some((cert, key)) = identity {
let certs = load_certs(cert)?;
let key = load_key(key)?;
config.set_single_client_cert(certs, key)?;
}
// 加载本地信任的根证书链
config.root_store = match rustls_native_certs::load_native_certs() {
Ok(store) | Err((Some(store), _)) => store,
Err((None, error)) => return Err(error.into()),
};
// 如果有签署服务器的 CA 证书，则加载它，这样服务器证书不在根证书链
// 但是这个 CA 证书能验证它，也可以
if let Some(cert) = server_ca {
let mut buf = Cursor::new(cert);
config.root_store.add_pem_file(&mut buf).unwrap();
}
Ok(Self {
config: Arc::new(config),
domain: Arc::new(domain.into()),
})
}
/// 触发 TLS 协议，把底层的 stream 转换成 TLS stream
pub async fn connect<S>(&self, stream: S) -> Result<ClientTlsStream<S>, KvError>
where
S: AsyncRead + AsyncWrite + Unpin + Send,
{
let dns = DNSNameRef::try_from_ascii_str(self.domain.as_str())
.map_err(|_| KvError::Internal("Invalid DNS name".into()))?;
let stream = TlsConnector::from(self.config.clone())
.connect(dns, stream)
.await?;
Ok(stream)
}
}
impl TlsServerAcceptor {
/// 加载 server cert / CA cert，生成 ServerConfig
pub fn new(cert: &str, key: &str, client_ca: Option<&str>) -> Result<Self, KvError> {
let certs = load_certs(cert)?;
let key = load_key(key)?;
let mut config = match client_ca {
None => ServerConfig::new(NoClientAuth::new()),
Some(cert) => {
// 如果客户端证书是某个 CA 证书签发的，则把这个 CA 证书加载到信任链中
let mut cert = Cursor::new(cert);
let mut client_root_cert_store = RootCertStore::empty();
client_root_cert_store
.add_pem_file(&mut cert)
.map_err(|_| KvError::CertifcateParseError("CA", "cert"))?;
let client_auth = AllowAnyAuthenticatedClient::new(client_root_cert_store);
ServerConfig::new(client_auth)
}
};
// 加载服务器证书
config
.set_single_cert(certs, key)
.map_err(|_| KvError::CertifcateParseError("server", "cert"))?;
config.set_protocols(&[Vec::from(&ALPN_KV[..])]);
Ok(Self {
inner: Arc::new(config),
})
}
/// 触发 TLS 协议，把底层的 stream 转换成 TLS stream
pub async fn accept<S>(&self, stream: S) -> Result<ServerTlsStream<S>, KvError>
where
S: AsyncRead + AsyncWrite + Unpin + Send,
{
let acceptor = TlsAcceptor::from(self.inner.clone());
Ok(acceptor.accept(stream).await?)
}
}
fn load_certs(cert: &str) -> Result<Vec<Certificate>, KvError> {
let mut cert = Cursor::new(cert);
pemfile::certs(&mut cert).map_err(|_| KvError::CertifcateParseError("server", "cert"))
}
fn load_key(key: &str) -> Result<PrivateKey, KvError> {
let mut cursor = Cursor::new(key);
// 先尝试用 PKCS8 加载私钥
if let Ok(mut keys) = pemfile::pkcs8_private_keys(&mut cursor) {
if !keys.is_empty() {
return Ok(keys.remove(0));
}
}
// 再尝试加载 RSA key
cursor.set_position(0);
if let Ok(mut keys) = pemfile::rsa_private_keys(&mut cursor) {
if !keys.is_empty() {
return Ok(keys.remove(0));
}
}
// 不支持的私钥类型
Err(KvError::CertifcateParseError("private", "key"))
}
Aunque tiene más de 100 líneas, el trabajo principal es en realidad generar el ServerConfig/ClientConfig requerido por tokio-tls en función del certificado proporcionado. Después de procesar la configuración, la lógica central de este código es en realidad el método connect () del cliente y el método aceptar () del servidor. Ambos aceptan una secuencia que satisface AsyncRead + AsyncWrite + Unpin + Send. De manera similar a la conferencia anterior, no queremos que el código TLS solo acepte TcpStream, por lo que aquí se proporciona un parámetro genérico S:
Después de usar TlsConnector o TlsAcceptor para procesar conectar/aceptar, obtenemos un TlsStream, que también satisface AsyncRead + AsyncWrite + Desanclar + Enviar, se pueden completar operaciones posteriores en él.
Debido a nuestro buen diseño de interfaz a lo largo del camino, especialmente ProstClientStream / ProstServerStream, todos aceptan parámetros genéricos, el código TLS se puede incrustar sin problemas. Por ejemplo, cliente:
// 新加的代码
let connector = TlsClientConnector::new("kvserver.acme.inc", None, Some(ca_cert))?;
let stream = TcpStream::connect(addr).await?;
// 新加的代码
let stream = connector.connect(stream).await?;
let mut client = ProstClientStream::new(stream);
Simplemente cambie la transmisión pasada a ProstClientStream de TcpStream al TlsStream generado para admitir TLS sin problemas.
Lado del servidor completo
use anyhow::Result;
use kv3::{
MemTable, ProstServerStream, Service, ServiceInner, TlsServerAcceptor};
use tokio::net::TcpListener;
use tracing::info;
#[tokio::main]
async fn main() -> Result<()> {
tracing_subscriber::fmt::init();
let addr = "127.0.0.1:9527";
// 以后从配置文件取
let server_cert = include_str!("../fixtures/server.cert");
let server_key = include_str!("../fixtures/server.key");
let acceptor = TlsServerAcceptor::new(server_cert, server_key, None)?;
let service: Service = ServiceInner::new(MemTable::new()).into();
let listener = TcpListener::bind(addr).await?;
info!("Start listening on {}", addr);
loop {
let tls = acceptor.clone();
let (stream, addr) = listener.accept().await?;
info!("Client {:?} connected", addr);
let stream = tls.accept(stream).await?;
let stream = ProstServerStream::new(stream, service.clone());
tokio::spawn(async move {
stream.process().await });
}
}
cliente
use anyhow::Result;
use kv3::{
CommandRequest, ProstClientStream, TlsClientConnector};
use tokio::net::TcpStream;
use tracing::info;
#[tokio::main]
async fn main() -> Result<()> {
tracing_subscriber::fmt::init();
// 以后用配置替换
let ca_cert = include_str!("../fixtures/ca.cert");
let addr = "127.0.0.1:9527";
// 连接服务器
let connector = TlsClientConnector::new("kvserver.acme.inc", None, Some(ca_cert))?;
let stream = TcpStream::connect(addr).await?;
let stream = connector.connect(stream).await?;
let mut client = ProstClientStream::new(stream);
// 生成一个 HSET 命令
let cmd = CommandRequest::new_hset("table1", "hello", "world".to_string().into());
// 发送 HSET 命令
let data = client.execute(cmd).await?;
info!("Got response {:?}", data);
Ok(())
}
En comparación con el proyecto de código de la conferencia anterior, los códigos de cliente y servidor actualizados solo tienen una línea más cada uno, encapsulando TcpStream en TlsStream . Este es el gran poder de usar rasgos para la programación orientada a interfaces. Varios componentes del sistema pueden provenir de diferentes cajas, pero siempre y cuando sus interfaces sean consistentes (o creemos un adaptador para que sus interfaces sean consistentes), se pueden integrar sin problemas. insertado.