diameter_interface/transport/
client.rs

1use crate::errors::DiameterResult;
2use crate::errors::Error::ClientError;
3use crate::modeling::diameter::DiameterMessage;
4use crate::modeling::message::dictionary::Dictionary;
5use std::io::Write;
6use std::net::{Shutdown, TcpStream};
7use std::sync::Arc;
8
9pub struct DiameterClient {
10    address: &'static str,
11    stream: Option<TcpStream>,
12}
13
14impl DiameterClient {
15    pub fn new(address: &'static str) -> Self {
16        Self {
17            address,
18            stream: None,
19        }
20    }
21
22    pub fn connect(&mut self) -> DiameterResult<()> {
23        let stream = TcpStream::connect(self.address)?;
24        self.stream = Some(stream);
25        Ok(())
26    }
27
28    pub fn close(&mut self) -> DiameterResult<()> {
29        if let Some(ref mut stream) = self.stream {
30            stream.shutdown(Shutdown::Both)?;
31            Ok(())
32        } else {
33            Err(ClientError("Connection not established yet!"))
34        }
35    }
36
37    pub fn send_message(
38        &mut self,
39        message: &mut DiameterMessage,
40        dict: Arc<Dictionary>,
41    ) -> DiameterResult<DiameterMessage> {
42        if let Some(ref mut stream) = self.stream {
43            let mut buffer = vec![];
44            message.encode_to(&mut buffer)?;
45            stream.write_all(&buffer)?;
46            let answer = DiameterMessage::decode_from(stream, Arc::clone(&dict))?;
47            Ok(answer)
48        } else {
49            Err(ClientError("Connection not established yet!"))
50        }
51    }
52}