AkurAI-Build / src/server/search_api.rs

Latest change e8aa34a9be50d7d7ea72f66991ff82dfca7d23e1 - OLI-276: Git over SSH, forks and mirrors, code search through akurai-index, repository insights by Ólafur Búi Ólafsson

//! Cross-repository code search (OLI-276), answered by akurai-index and
//! filtered to the repositories the caller may read.

use std::process::Command;

use axum::{
    Json,
    extract::{Query, State},
    http::HeaderMap,
    response::{Html, IntoResponse, Redirect, Response},
};
use minijinja::context;
use serde::{Deserialize, Serialize};

use super::access_api::{Identity, identify, role_on};
use super::{ApiError, AppState, Envelope, encode_code_segment, ok, valid_session};
use crate::access::Role;

/// Where akurai-index answers and the bearer token it expects.
#[derive(Clone, Debug)]
pub struct IndexConfig {
    pub url: String,
    pub token: String,
}

/// How hits are fetched: the real one runs `curl`; tests inject their own.
pub trait Searcher: Send + Sync {
    fn search(&self, query: &str, repo: Option<&str>, k: usize) -> anyhow::Result<Vec<RawHit>>;
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct RawHit {
    pub repo: String,
    #[serde(default)]
    pub branch: String,
    pub path: String,
    #[serde(default)]
    pub symbol: Option<String>,
    #[serde(default)]
    pub kind: String,
    #[serde(default)]
    pub start_line: usize,
    #[serde(default)]
    pub end_line: usize,
    #[serde(default)]
    pub score: f64,
    #[serde(default)]
    pub snippet: String,
}

#[derive(Clone, Debug, Serialize)]
pub struct Hit {
    pub repo: String,
    pub branch: String,
    pub path: String,
    pub symbol: Option<String>,
    pub kind: String,
    pub start_line: usize,
    pub end_line: usize,
    pub score: f64,
    pub snippet: String,
    /// The forge's own file view for this hit.
    pub href: String,
}

#[derive(Clone, Debug, Serialize)]
pub struct SearchResult {
    pub query: String,
    pub hits: Vec<Hit>,
}

pub struct CurlSearcher(pub IndexConfig);

fn percent_encode(text: &str) -> String {
    let mut out = String::with_capacity(text.len());
    for byte in text.bytes() {
        if byte.is_ascii_alphanumeric() || b"-_.~".contains(&byte) {
            out.push(byte as char);
        } else {
            out.push_str(&format!("%{byte:02X}"));
        }
    }
    out
}

impl Searcher for CurlSearcher {
    fn search(&self, query: &str, repo: Option<&str>, k: usize) -> anyhow::Result<Vec<RawHit>> {
        let mut url = format!(
            "{}/search?q={}&k={k}",
            self.0.url.trim_end_matches('/'),
            percent_encode(query)
        );
        if let Some(repo) = repo {
            url.push_str(&format!("&repo={}", percent_encode(repo)));
        }
        let output = Command::new("curl")
            .args(["--silent", "--show-error", "--fail", "--max-time", "20"])
            .arg(&url)
            .arg("-H")
            .arg(format!("Authorization: Bearer {}", self.0.token))
            .output()?;
        anyhow::ensure!(
            output.status.success(),
            "akurai-index did not answer: {}",
            String::from_utf8_lossy(&output.stderr).trim()
        );
        let body: serde_json::Value = serde_json::from_slice(&output.stdout)?;
        if let Some(error) = body.get("error").and_then(|e| e.as_str()) {
            anyhow::bail!("akurai-index: {error}");
        }
        Ok(serde_json::from_value(
            body.get("hits")
                .cloned()
                .unwrap_or(serde_json::Value::Array(vec![])),
        )?)
    }
}

#[derive(Deserialize)]
pub(super) struct SearchQuery {
    #[serde(default)]
    q: String,
    repo: Option<String>,
    k: Option<usize>,
}

const MAX_QUERY: usize = 300;
const MAX_K: usize = 50;

/// Hits for `identity`: only repositories it may read (public ones included).
async fn run(
    state: &AppState,
    identity: &Identity,
    query: &SearchQuery,
) -> Result<SearchResult, ApiError> {
    let Some(searcher) = state.searcher.clone() else {
        return Err(ApiError::not_found(
            "code search is not configured on this forge (AKURAI_BUILD_INDEX_URL)",
        ));
    };
    let q = query.q.trim().to_string();
    if q.is_empty() || q.len() > MAX_QUERY {
        return Err(ApiError::bad(format!("q must be 1 to {MAX_QUERY} bytes")));
    }
    let k = query.k.unwrap_or(20).clamp(1, MAX_K);
    let repo = query.repo.clone();
    let wanted = q.clone();
    let raw = tokio::task::spawn_blocking(move || searcher.search(&wanted, repo.as_deref(), k * 2))
        .await
        .map_err(ApiError::internal)?
        .map_err(ApiError::internal)?;
    let mut hits = Vec::new();
    let mut allowed: std::collections::BTreeMap<String, bool> = Default::default();
    for hit in raw {
        let visible = match allowed.get(&hit.repo) {
            Some(known) => *known,
            None => {
                let known = match state.database.repository(&hit.repo) {
                    Ok(repository) if repository.visibility == "public" => true,
                    Ok(_) => {
                        matches!(role_on(identity, state, &hit.repo)?, Some(role) if role.allows(Role::Read))
                    }
                    Err(_) => false,
                };
                allowed.insert(hit.repo.clone(), known);
                known
            }
        };
        if !visible {
            continue;
        }
        let branch = if hit.branch.is_empty() {
            "HEAD".to_string()
        } else {
            hit.branch.clone()
        };
        let href = format!(
            "/app/repositories/{}/tree/{}/{}#L{}",
            encode_code_segment(&hit.repo),
            encode_code_segment(&branch),
            hit.path
                .split('/')
                .map(encode_code_segment)
                .collect::<Vec<_>>()
                .join("/"),
            hit.start_line.max(1)
        );
        hits.push(Hit {
            repo: hit.repo,
            branch,
            path: hit.path,
            symbol: hit.symbol,
            kind: hit.kind,
            start_line: hit.start_line,
            end_line: hit.end_line,
            score: hit.score,
            snippet: hit.snippet,
            href,
        });
        if hits.len() == k {
            break;
        }
    }
    Ok(SearchResult { query: q, hits })
}

pub(super) async fn api(
    State(state): State<AppState>,
    headers: HeaderMap,
    Query(query): Query<SearchQuery>,
) -> Result<Json<Envelope<SearchResult>>, ApiError> {
    let identity = identify(&headers, &state)?;
    Ok(ok(run(&state, &identity, &query).await?))
}

pub(super) async fn page(
    State(state): State<AppState>,
    headers: HeaderMap,
    Query(query): Query<SearchQuery>,
) -> Result<Response, ApiError> {
    if !valid_session(&headers, &state) {
        return Ok(Redirect::to("/login").into_response());
    }
    let identity = identify(&headers, &state)?;
    let (result, error) = if query.q.trim().is_empty() {
        (None, None)
    } else {
        match run(&state, &identity, &query).await {
            Ok(result) => (Some(result), None),
            Err(error) => (None, Some(error.message.clone())),
        }
    };
    let repositories: Vec<String> = state
        .database
        .repositories()
        .map_err(ApiError::internal)?
        .into_iter()
        .map(|r| r.name)
        .collect();
    let rendered = state
        .templates
        .get_template("search.html")
        .map_err(ApiError::internal)?
        .render(context! {
            version => env!("CARGO_PKG_VERSION"),
            q => query.q, repo => query.repo, result => result, error => error,
            repositories => repositories, enabled => state.searcher.is_some(),
        })
        .map_err(ApiError::internal)?;
    Ok(Html(rendered).into_response())
}

#[cfg(test)]
pub(crate) mod testing {
    use super::*;

    /// Answers every query with the given hits.
    pub struct FixedSearcher(pub Vec<RawHit>);

    impl Searcher for FixedSearcher {
        fn search(
            &self,
            _query: &str,
            repo: Option<&str>,
            _k: usize,
        ) -> anyhow::Result<Vec<RawHit>> {
            Ok(self
                .0
                .iter()
                .filter(|h| repo.is_none_or(|r| r == h.repo))
                .cloned()
                .collect())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{
        Router,
        body::{Body, to_bytes},
        http::{Request, StatusCode, header},
        routing::get,
    };
    use std::sync::Arc;
    use tower::ServiceExt;

    fn hit(repo: &str, path: &str, line: usize) -> RawHit {
        RawHit {
            repo: repo.into(),
            branch: "main".into(),
            path: path.into(),
            symbol: None,
            kind: "chunk".into(),
            start_line: line,
            end_line: line + 3,
            score: 0.5,
            snippet: "fn x() {}".into(),
        }
    }

    #[tokio::test]
    async fn search_filters_hits_to_readable_repositories_and_links_into_the_forge() {
        let mut state = crate::server::tests::test_state("api-token-123", "webhook-secret");
        state
            .database
            .add_repository("open", "https://example.com/open.git", "main")
            .expect("add");
        state
            .database
            .set_repository_visibility("open", "public")
            .expect("public");
        state
            .database
            .add_repository("secret", "https://example.com/secret.git", "main")
            .expect("add");
        state.searcher = Some(Arc::new(testing::FixedSearcher(vec![
            hit("secret", "src/lib.rs", 10),
            hit("open", "src/a b.rs", 3),
            hit("gone", "x.rs", 1),
        ])));
        let now = super::super::unix_time().ok().expect("now") as i64;
        let (_, bob) = crate::access::create_token(
            &state.database,
            &crate::access::NewToken {
                kind: crate::access::Kind::Pat,
                name: "t",
                subject: Some("bob@example"),
                repository: None,
                role: Role::Admin,
                ttl_secs: None,
            },
            now,
        )
        .expect("pat");
        let app = Router::new()
            .route("/api/search", get(api))
            .with_state(state.clone());
        let get_json = |uri: &str, auth: Option<&str>| {
            let app = app.clone();
            let uri = uri.to_string();
            let auth = auth.map(str::to_string);
            async move {
                let mut request = Request::builder().uri(uri);
                if let Some(auth) = auth {
                    request = request.header(header::AUTHORIZATION, auth);
                }
                let response = app
                    .oneshot(request.body(Body::empty()).expect("request"))
                    .await
                    .expect("response");
                let status = response.status();
                let bytes = to_bytes(response.into_body(), usize::MAX)
                    .await
                    .expect("body");
                (
                    status,
                    serde_json::from_slice::<serde_json::Value>(&bytes)
                        .unwrap_or(serde_json::Value::Null),
                )
            }
        };

        let (status, _) = get_json("/api/search?q=x", None).await;
        assert_eq!(status, StatusCode::UNAUTHORIZED);
        let (status, body) = get_json("/api/search?q=", Some("Bearer api-token-123")).await;
        assert_eq!(status, StatusCode::BAD_REQUEST, "{body}");
        let (status, body) = get_json("/api/search?q=fn+x", Some("Bearer api-token-123")).await;
        assert_eq!(status, StatusCode::OK, "{body}");
        let hits = body["data"]["hits"].as_array().expect("hits");
        assert_eq!(
            hits.len(),
            2,
            "controller sees both known repos, never unknown ones: {body}"
        );
        assert_eq!(
            hits[0]["href"],
            "/app/repositories/secret/tree/main/src/lib.rs#L10"
        );
        assert_eq!(
            hits[1]["href"],
            "/app/repositories/open/tree/main/src/a%20b.rs#L3"
        );

        let (status, body) = get_json("/api/search?q=fn+x", Some(&format!("Bearer {bob}"))).await;
        assert_eq!(status, StatusCode::OK, "{body}");
        let hits = body["data"]["hits"].as_array().expect("hits");
        assert_eq!(hits.len(), 1, "bob only sees the public repo: {body}");
        assert_eq!(hits[0]["repo"], "open");
        let (status, body) = get_json(
            "/api/search?q=fn+x&repo=secret",
            Some(&format!("Bearer {bob}")),
        )
        .await;
        assert_eq!(status, StatusCode::OK);
        assert_eq!(body["data"]["hits"].as_array().map(Vec::len), Some(0));

        state.searcher = None;
        let app = Router::new()
            .route("/api/search", get(api))
            .with_state(state);
        let response = app
            .oneshot(
                Request::builder()
                    .uri("/api/search?q=x")
                    .header(header::AUTHORIZATION, "Bearer api-token-123")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");
        assert_eq!(response.status(), StatusCode::NOT_FOUND, "search off");
    }

    #[test]
    fn percent_encoding_keeps_unreserved_bytes() {
        assert_eq!(percent_encode("fn x/y.rs~"), "fn%20x%2Fy.rs~");
    }
}