laburnum 1.17.1

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  crate::protocol::jsonrpc::{
    Notification,
    Request,
    Response,
  },
  serde::{
    Deserialize,
    Serialize,
  },
  std::io::{
    self,
    BufRead,
    Write,
  },
};

fn invalid_data(
  error: impl Into<Box<dyn std::error::Error + Send + Sync>>,
) -> io::Error {
  io::Error::new(io::ErrorKind::InvalidData, error)
}

macro_rules! invalid_data {
  ($($tt:tt)*) => (invalid_data(format!($($tt)*)))
}

#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
#[serde(untagged)]
pub enum Message {
  Response(Response),
  Request(Request),
  Notification(Notification),
}

impl Message {
  pub fn read(r: &mut impl BufRead) -> io::Result<Option<Message>> {
    Message::_read(r)
  }

  fn _read(r: &mut dyn BufRead) -> io::Result<Option<Message>> {
    let text = match read_msg_text(r)? {
      | None => return Ok(None),
      | Some(text) => text,
    };

    let msg = match serde_json::from_str(&text) {
      | Ok(msg) => msg,
      | Err(e) => {
        return Err(invalid_data!("malformed JSON-RPC payload: {:?}", e));
      },
    };

    Ok(Some(msg))
  }

  pub fn write(&self, w: &mut impl Write) -> io::Result<()> {
    self._write(w)
  }

  fn _write(&self, w: &mut dyn Write) -> io::Result<()> {
    let text = serde_json::to_string(self)?;
    write_msg_text(w, &text)
  }
}

impl From<Request> for Message {
  fn from(request: Request) -> Message {
    Message::Request(request)
  }
}

impl From<Response> for Message {
  fn from(response: Response) -> Message {
    Message::Response(response)
  }
}

impl From<Notification> for Message {
  fn from(notification: Notification) -> Message {
    Message::Notification(notification)
  }
}

fn read_msg_text(inp: &mut dyn BufRead) -> io::Result<Option<String>> {
  let mut size = None;
  let mut buf = String::new();
  loop {
    buf.clear();
    if inp.read_line(&mut buf)? == 0 {
      return Ok(None);
    }
    if !buf.ends_with("\r\n") {
      return Err(invalid_data!("malformed header: {:?}", buf));
    }
    let buf = &buf[..buf.len() - 2];
    if buf.is_empty() {
      break;
    }
    let mut parts = buf.splitn(2, ": ");
    // splitn always returns at least one element, but we handle None for safety
    let Some(header_name) = parts.next() else {
      return Err(invalid_data!("malformed header: {:?}", buf));
    };
    let header_value = parts
      .next()
      .ok_or_else(|| invalid_data!("malformed header: {:?}", buf))?;
    if header_name.eq_ignore_ascii_case("Content-Length") {
      size = Some(header_value.parse::<usize>().map_err(invalid_data)?);
    }
  }
  let size: usize = size.ok_or_else(|| invalid_data!("no Content-Length"))?;
  let mut buf = buf.into_bytes();
  buf.resize(size, 0);
  inp.read_exact(&mut buf)?;
  let buf = String::from_utf8(buf).map_err(invalid_data)?;

  Ok(Some(buf))
}

fn write_msg_text(out: &mut dyn Write, msg: &str) -> io::Result<()> {
  write!(out, "Content-Length: {}\r\n\r\n", msg.len())?;
  out.write_all(msg.as_bytes())?;
  out.flush()?;
  Ok(())
}