use std::collections::BTreeMap;
use std::io::{BufRead, Write};
use onetaskgraph_plugin_api::{SecretResolver, SourceError, SourceName, TaskSource};
use secrecy::SecretString;
use serde::Deserialize;
use serde_json::{Value, json};
use super::connection::{Line, MAX_LINE, read_line};
use super::wire::{
DependencyParams, HandshakePluginKind, IdParams, InitializeParams, InitializeResult,
LabelParams, PROTOCOL_VERSION, ProjectQueryParams, ProjectWriteParams, Request, Response,
TaskQueryParams, TaskWriteParams,
};
use crate::registry::PluginKind;
#[derive(Debug, Clone, Deserialize)]
struct HostedSettings {
kind: PluginKind,
#[serde(default)]
config: Value,
}
pub async fn serve(input: impl BufRead, output: impl Write) -> std::io::Result<()> {
serve_kind(input, output, None).await
}
pub async fn serve_plugin(
input: impl BufRead,
output: impl Write,
kind: PluginKind,
) -> std::io::Result<()> {
serve_kind(input, output, Some(kind)).await
}
async fn serve_kind(
mut input: impl BufRead,
mut output: impl Write,
kind: Option<PluginKind>,
) -> std::io::Result<()> {
let mut source: Option<Box<dyn TaskSource>> = None;
loop {
let line = match read_line(&mut input) {
Line::Read(line) => line,
Line::Ended => return Ok(()),
Line::Failed(error) => return Err(error),
Line::TooLong => {
eprintln!(
"onetaskgraph-source: a request ran past {MAX_LINE} bytes without \
ending its line; closing the connection"
);
return Ok(());
}
};
if line.trim().is_empty() {
continue;
}
let Some(id) = addressed(&line) else {
eprintln!("onetaskgraph-source: ignoring a line with no request id: {line}");
continue;
};
let response = match serde_json::from_str::<Request>(&line) {
Ok(request) => answer(&mut source, request, kind).await,
Err(error) => Response::failed(
id,
SourceError::Malformed {
message: format!("that is not a request envelope: {error}"),
},
),
};
let finished = ended_the_connection(&response);
writeln!(
output,
"{}",
serde_json::to_string(&response).expect("a response is plain data")
)?;
output.flush()?;
if finished {
return Ok(());
}
}
}
fn addressed(line: &str) -> Option<String> {
serde_json::from_str::<Value>(line)
.ok()?
.get("id")?
.as_str()
.map(str::to_owned)
}
fn ended_the_connection(response: &Response) -> bool {
matches!(
response.error.as_ref(),
Some(SourceError::Config { message }) if message.starts_with(VERSION_REFUSAL)
)
}
const VERSION_REFUSAL: &str = "protocol version ";
async fn answer(
source: &mut Option<Box<dyn TaskSource>>,
request: Request,
kind: Option<PluginKind>,
) -> Response {
let Request { id, method, params } = request;
if method == "initialize" {
return match source {
Some(_) => Response::failed(
id,
SourceError::Malformed {
message: "this connection was already initialized".to_owned(),
},
),
None => initialize(source, id, params, kind),
};
}
let Some(built) = source.as_deref() else {
return Response::failed(
id,
SourceError::Malformed {
message: format!("{method} arrived before the handshake"),
},
);
};
match dispatch(built, &method, params).await {
Ok(result) => Response::ok(id, result),
Err(error) => Response::failed(id, error),
}
}
fn initialize(
source: &mut Option<Box<dyn TaskSource>>,
id: String,
params: Value,
kind: Option<PluginKind>,
) -> Response {
let params: InitializeParams = match serde_json::from_value(params) {
Ok(params) => params,
Err(error) => {
return Response::failed(
id,
SourceError::Config {
message: format!("that is not an initialize request: {error}"),
},
);
}
};
if params.protocol_version != PROTOCOL_VERSION {
return Response::failed(
id,
SourceError::Config {
message: format!(
"{VERSION_REFUSAL}{} is not supported by this plugin; it speaks \
version {PROTOCOL_VERSION}",
params.protocol_version
),
},
);
}
match build(¶ms, kind) {
Ok(built) => {
let kind = match HandshakePluginKind::new(built.kind()) {
Ok(kind) => kind,
Err(error) => {
return Response::failed(
id,
SourceError::Malformed {
message: format!("the hosted plugin reported an invalid kind: {error}"),
},
);
}
};
let result = InitializeResult {
protocol_version: Some(PROTOCOL_VERSION),
kind,
capabilities: built.capabilities(),
writes: Some(built.writes()),
};
*source = Some(built);
Response::ok(
id,
serde_json::to_value(&result).expect("a result is plain data"),
)
}
Err(error) => Response::failed(id, error),
}
}
fn build(
params: &InitializeParams,
selected: Option<PluginKind>,
) -> Result<Box<dyn TaskSource>, SourceError> {
let (kind, config) = match selected {
Some(kind) => (kind, ¶ms.config),
None => {
let settings: HostedSettings = serde_json::from_value(params.config.clone()).map_err(
|error| SourceError::Config {
message: format!(
"this host serves a plugin of this build, and its settings must name one \
as {{\"kind\": …, \"config\": …}}: {error}"
),
},
)?;
return build_plugin(params, settings.kind, &settings.config);
}
};
build_plugin(params, kind, config)
}
fn build_plugin(
params: &InitializeParams,
kind: PluginKind,
config: &Value,
) -> Result<Box<dyn TaskSource>, SourceError> {
let name = SourceName::new(params.source_name.clone())?;
kind.plugin()
.build(&name, config, &Handshake(¶ms.secrets))
}
struct Handshake<'a>(&'a BTreeMap<String, String>);
impl SecretResolver for Handshake<'_> {
fn get(&self, var: &str) -> Option<SecretString> {
self.0
.get(var)
.map(|value| SecretString::from(value.clone()))
}
}
async fn dispatch(
source: &dyn TaskSource,
method: &str,
params: Value,
) -> Result<Value, SourceError> {
match method {
"health" => encode(source.health().await?),
"get_task" => {
let params: IdParams = decode(method, params)?;
encode(json!({ "task": source.get_task(¶ms.id).await? }))
}
"get_project" => {
let params: IdParams = decode(method, params)?;
encode(json!({ "project": source.get_project(¶ms.id).await? }))
}
"query_tasks" => {
let params: TaskQueryParams = decode(method, params)?;
encode(source.query_tasks(¶ms.query, ¶ms.page).await?)
}
"query_projects" => {
let params: ProjectQueryParams = decode(method, params)?;
encode(source.query_projects(¶ms.query, ¶ms.page).await?)
}
"labels" => {
let params: LabelParams = decode(method, params)?;
encode(source.labels(¶ms.page).await?)
}
"task_dependencies" => {
let params: DependencyParams = decode(method, params)?;
encode(
source
.task_dependencies(¶ms.id, params.direction, ¶ms.page)
.await?,
)
}
"project_dependencies" => {
let params: DependencyParams = decode(method, params)?;
encode(
source
.project_dependencies(¶ms.id, params.direction, ¶ms.page)
.await?,
)
}
"write_task" => {
let params: TaskWriteParams = decode(method, params)?;
encode(json!({ "id": source.write_task(¶ms.write).await? }))
}
"write_project" => {
let params: ProjectWriteParams = decode(method, params)?;
encode(json!({ "id": source.write_project(¶ms.write).await? }))
}
other => Err(SourceError::Malformed {
message: format!("protocol version {PROTOCOL_VERSION} has no method called {other:?}"),
}),
}
}
fn decode<T: for<'de> Deserialize<'de>>(method: &str, params: Value) -> Result<T, SourceError> {
serde_json::from_value(params).map_err(|error| SourceError::Malformed {
message: format!("the parameters of {method} are not the shape it takes: {error}"),
})
}
fn encode<T: serde::Serialize>(value: T) -> Result<Value, SourceError> {
serde_json::to_value(value).map_err(|error| SourceError::Malformed {
message: format!("this source returned data that will not serialize: {error}"),
})
}