trifid/tfclient/src/socketclient.rs

use std::error::Error;
use std::io::{BufRead, BufReader, Write};
use std::net::{IpAddr, SocketAddr, TcpStream};
use log::{error, info};
use crate::config::TFClientConfig;
use crate::socketworker::{ctob, DisconnectReason, JSON_API_VERSION, JsonMessage};

pub fn enroll(code: &str, config: &TFClientConfig) -> Result<(), Box<dyn Error>> {
    info!("Connecting to local command socket...");
    let mut stream = TcpStream::connect(SocketAddr::new(IpAddr::from([127, 0, 0, 1]), config.listen_port))?;
    let stream2 = stream.try_clone()?;
    let mut reader = BufReader::new(&stream2);

    info!("Sending Hello...");
    stream.write_all(&ctob(JsonMessage::Hello {
        version: JSON_API_VERSION,
    }))?;
    info!("Waiting for hello...");
    let msg = read_msg(&mut reader)?;
    match msg {
        JsonMessage::Hello { .. } => {
            info!("Server sent hello, connection established")
        }
        JsonMessage::Goodbye { reason } => {
            error!("Disconnected by server. Reason: {:?}", reason);
            return Err("Disconnected by server".into());
        }
        _ => {
            error!("Server returned unexpected message: {:?}", msg);
            error!("Sending goodbye and exiting");
            stream.write_all(&ctob(JsonMessage::Goodbye {
                reason: DisconnectReason::UnexpectedMessageType,
            }))?;
            return Err("Unexpected message type by server".into());
        }
    }

    info!("Sending enroll request...");
    stream.write_all(&ctob(JsonMessage::Enroll {
        code: code.to_string(),
    }))?;

    info!("Sending disconnect...");
    stream.write_all(&ctob(JsonMessage::Goodbye {
        reason: DisconnectReason::Done,
    }))?;

    info!("Sent enroll request to tfclient daemon. Check logs to see if the enrollment was successful.");

    Ok(())
}

fn read_msg(reader: &mut BufReader<&TcpStream>) -> Result<JsonMessage, Box<dyn Error>> {
    let mut str = String::new();
    reader.read_line(&mut str)?;
    let msg: JsonMessage = serde_json::from_str(&str)?;
    Ok(msg)
}