featherbit/consumers/

mod.rs

//! Consumer identities and credentials.
//!
//! The featherbit analogue of APISIX's consumer concept: named API clients
//! declared in `gateway.yaml` (or via the Admin API), each carrying
//! per-auth-plugin credentials. Auth plugins configured with
//! `use_consumers: true` resolve the presented credential against the
//! [`ConsumerStore`] and, on a match, attach the consumer's identity to the
//! request via [`attach_consumer`] so downstream nodes (consumer-restriction,
//! loggers, rate limits keyed on `$consumer_name`) can act on it.
//!
//! The store is immutable once built; reloads build a fresh store and swap
//! it atomically (`ArcSwap` in `PluginResources`), so lookups on the request
//! path are lock-free.

use std::collections::HashMap;
use std::sync::Arc;

use serde::{Deserialize, Serialize};

use crate::context::Context;

/// A consumer as declared in `gateway.yaml` under `consumers:`.
///
/// ```yaml
/// consumers:
///   - name: alice
///     group: partners
///     labels: { custom_id: "42", tier: gold }
///     credentials:
///       key-auth: { key: "s3cret" }
///       basic-auth: { username: alice, password: pw }
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConsumerConfig {
    /// Unique consumer name (the identity attached to matching requests).
    pub name: String,
    /// Optional consumer group, exposed as `consumer.group` / `$consumer_group_id`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub group: Option<String>,
    /// Free-form labels; `custom_id` is special-cased into the
    /// `X-Consumer-Custom-ID` header (matching APISIX).
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub labels: HashMap<String, String>,
    /// Per-auth-plugin credentials, keyed by plugin type
    /// (e.g. `key-auth: {key}`, `basic-auth: {username, password}`).
    #[serde(default)]
    pub credentials: HashMap<String, serde_json::Value>,
}

/// A resolved consumer, shared read-only across requests.
#[derive(Debug)]
pub struct Consumer {
    pub name: String,
    pub group: Option<String>,
    pub labels: HashMap<String, String>,
    pub credentials: HashMap<String, serde_json::Value>,
}

/// Lock-free lookup structure over the declared consumers.
///
/// Indexes each auth type's primary credential value → consumer, so an auth
/// plugin resolves a presented credential in one hash lookup. The primary
/// credential field per auth type: `key-auth` → `key`, `basic-auth` →
/// `username`, `jwt-auth` → `key` (falling back to the consumer name). Auth
/// plugins that need more than the primary field (e.g. basic-auth's password
/// check) read the matched consumer's `credentials`.
#[derive(Default)]
pub struct ConsumerStore {
    by_name: HashMap<String, Arc<Consumer>>,
    /// auth type → credential value → consumer
    by_credential: HashMap<String, HashMap<String, Arc<Consumer>>>,
}

/// The credential field indexed per auth type.
fn primary_credential_field(auth_type: &str) -> &'static str {
    match auth_type {
        "basic-auth" => "username",
        "hmac-auth" => "access_key",
        _ => "key",
    }
}

impl ConsumerStore {
    /// Builds the store, failing fast on duplicate consumer names or on two
    /// consumers sharing the same credential value for the same auth type.
    pub fn from_config(consumers: &[ConsumerConfig]) -> Result<Self, String> {
        let mut store = Self::default();

        // Consumer declarations arrive with `${VAR}` placeholders intact
        // (gateway config is loaded raw so the Admin API never serves
        // resolved secrets); the store is where they resolve, so credential
        // lookups index the actual values.
        let resolve = |s: &str| -> String {
            if s.contains("${") {
                crate::config::interpolate_env(s)
            } else {
                s.to_string()
            }
        };
        for config in consumers {
            if config.name.trim().is_empty() {
                return Err("consumer with empty name".to_string());
            }
            let credentials = config
                .credentials
                .iter()
                .map(|(auth_type, credential)| {
                    let mut credential = credential.clone();
                    crate::config::interpolate_env_json(&mut credential);
                    (auth_type.clone(), credential)
                })
                .collect();
            let consumer = Arc::new(Consumer {
                name: resolve(&config.name),
                group: config.group.as_deref().map(resolve),
                labels: config
                    .labels
                    .iter()
                    .map(|(k, v)| (k.clone(), resolve(v)))
                    .collect(),
                credentials,
            });

            if store
                .by_name
                .insert(consumer.name.clone(), consumer.clone())
                .is_some()
            {
                return Err(format!("duplicate consumer name '{}'", consumer.name));
            }

            for (auth_type, credential) in &consumer.credentials {
                let field = primary_credential_field(auth_type);
                let value = credential
                    .get(field)
                    .and_then(|v| v.as_str())
                    .ok_or_else(|| {
                        format!(
                            "consumer '{}': credential '{}' is missing string field '{}'",
                            consumer.name, auth_type, field
                        )
                    })?;

                let index = store.by_credential.entry(auth_type.clone()).or_default();
                if let Some(existing) = index.insert(value.to_string(), consumer.clone()) {
                    return Err(format!(
                        "consumers '{}' and '{}' share the same {} credential",
                        existing.name, consumer.name, auth_type
                    ));
                }
            }
        }

        Ok(store)
    }

    /// Resolves a presented credential value for an auth type.
    pub fn find_by_credential(&self, auth_type: &str, value: &str) -> Option<Arc<Consumer>> {
        self.by_credential.get(auth_type)?.get(value).cloned()
    }

    /// Looks a consumer up by name (used for `anonymous_consumer` fallback).
    pub fn get(&self, name: &str) -> Option<Arc<Consumer>> {
        self.by_name.get(name).cloned()
    }

    /// Number of declared consumers.
    #[allow(dead_code)] // store accessors; kept for completeness
    pub fn len(&self) -> usize {
        self.by_name.len()
    }

    /// True when no consumers are declared.
    #[allow(dead_code)]
    pub fn is_empty(&self) -> bool {
        self.by_name.is_empty()
    }
}

/// Attaches a matched consumer's identity to the request.
///
/// Writes the `context.message` keys downstream nodes read —
/// `consumer.name`, `consumer.group` (when set), `consumer.labels`,
/// `consumer.custom_id` (from `labels.custom_id`), `consumer.auth_type` —
/// and sets the upstream request headers `X-Consumer-Username` and
/// `X-Consumer-Custom-ID` (APISIX's upstream contract).
pub fn attach_consumer(ctx: &mut Context, consumer: &Consumer, auth_type: &str) {
    ctx.message.insert(
        "consumer.name".to_string(),
        serde_json::json!(consumer.name),
    );
    if let Some(ref group) = consumer.group {
        ctx.message
            .insert("consumer.group".to_string(), serde_json::json!(group));
    }
    if !consumer.labels.is_empty() {
        ctx.message.insert(
            "consumer.labels".to_string(),
            serde_json::json!(consumer.labels),
        );
    }
    ctx.message.insert(
        "consumer.auth_type".to_string(),
        serde_json::json!(auth_type),
    );

    ctx.request.headers.insert(
        "x-consumer-username".to_string(),
        vec![consumer.name.clone()],
    );
    if let Some(custom_id) = consumer.labels.get("custom_id") {
        ctx.message.insert(
            "consumer.custom_id".to_string(),
            serde_json::json!(custom_id),
        );
        ctx.request
            .headers
            .insert("x-consumer-custom-id".to_string(), vec![custom_id.clone()]);
    }
}

/// Returns a copy of `c` with credential secrets replaced by `"<masked>"`.
///
/// Every scalar leaf under `credentials` is masked except the identifying
/// halves of a credential pair (`username`, `access_key`), so a read-only
/// viewer can still tell *which* credential exists without learning it.
pub fn mask_credentials(c: &ConsumerConfig) -> ConsumerConfig {
    fn mask(v: &serde_json::Value, key: Option<&str>) -> serde_json::Value {
        match v {
            serde_json::Value::Object(map) => serde_json::Value::Object(
                map.iter()
                    .map(|(k, v)| (k.clone(), mask(v, Some(k))))
                    .collect(),
            ),
            serde_json::Value::Array(items) => {
                serde_json::Value::Array(items.iter().map(|v| mask(v, key)).collect())
            }
            serde_json::Value::Null => serde_json::Value::Null,
            _ if matches!(key, Some("username") | Some("access_key")) => v.clone(),
            _ => serde_json::Value::String("<masked>".into()),
        }
    }
    let mut out = c.clone();
    out.credentials = c
        .credentials
        .iter()
        .map(|(plugin, v)| (plugin.clone(), mask(v, None)))
        .collect();
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    fn config(json: serde_json::Value) -> Vec<ConsumerConfig> {
        serde_json::from_value(json).expect("valid consumer config")
    }

    #[test]
    fn mask_credentials_keeps_identifiers_only() {
        let c: ConsumerConfig = serde_yaml::from_str(
            "name: alice\ncredentials:\n  key-auth:\n    key: s3cret\n  basic-auth:\n    username: alice\n    password: pw\n  hmac-auth:\n    access_key: ak\n    secret_key: sk\n    nested: {inner: x}\n    tags: [tag1, tag2]\n",
        )
        .unwrap();
        let m = mask_credentials(&c);
        assert_eq!(m.name, "alice");
        assert_eq!(m.credentials["key-auth"]["key"], "<masked>");
        assert_eq!(m.credentials["basic-auth"]["username"], "alice");
        assert_eq!(m.credentials["basic-auth"]["password"], "<masked>");
        assert_eq!(m.credentials["hmac-auth"]["access_key"], "ak");
        assert_eq!(m.credentials["hmac-auth"]["secret_key"], "<masked>");
        assert_eq!(m.credentials["hmac-auth"]["nested"]["inner"], "<masked>");
        assert_eq!(m.credentials["hmac-auth"]["tags"][0], "<masked>");
        assert_eq!(m.credentials["hmac-auth"]["tags"][1], "<masked>");
        // The original is untouched.
        assert_eq!(c.credentials["key-auth"]["key"], "s3cret");
    }

    #[test]
    fn test_consumer_credentials_resolve_env_placeholders() {
        // gateway.yaml is loaded with `${VAR}` placeholders intact (so the
        // Admin API never serves resolved secrets); the store must resolve
        // them when it builds, or env-provided credentials stop matching.
        std::env::set_var("TEST_CONSUMER_KEY", "k-resolved-123");
        let store = ConsumerStore::from_config(&config(serde_json::json!([
            {
                "name": "alice",
                "labels": { "tier": "${TEST_CONSUMER_TIER:-gold}" },
                "credentials": { "key-auth": { "key": "${TEST_CONSUMER_KEY}" } }
            }
        ])))
        .unwrap();

        let alice = store
            .find_by_credential("key-auth", "k-resolved-123")
            .expect("credential indexed under the resolved value");
        assert_eq!(alice.name, "alice");
        assert_eq!(
            alice.credentials["key-auth"]["key"],
            serde_json::json!("k-resolved-123")
        );
        assert_eq!(alice.labels["tier"], "gold");
        std::env::remove_var("TEST_CONSUMER_KEY");
    }

    #[test]
    fn test_store_lookup_and_attach() {
        let store = ConsumerStore::from_config(&config(serde_json::json!([
            {
                "name": "alice",
                "group": "partners",
                "labels": { "custom_id": "42", "tier": "gold" },
                "credentials": {
                    "key-auth": { "key": "alice-key" },
                    "basic-auth": { "username": "alice", "password": "pw" }
                }
            },
            { "name": "bob", "credentials": { "key-auth": { "key": "bob-key" } } }
        ])))
        .unwrap();

        assert_eq!(store.len(), 2);
        let alice = store.find_by_credential("key-auth", "alice-key").unwrap();
        assert_eq!(alice.name, "alice");
        assert!(store.find_by_credential("key-auth", "nope").is_none());
        assert!(store.find_by_credential("jwt-auth", "alice-key").is_none());
        assert_eq!(
            store
                .find_by_credential("basic-auth", "alice")
                .unwrap()
                .name,
            "alice"
        );
        assert_eq!(store.get("bob").unwrap().name, "bob");

        // attach writes message keys + headers
        let mut ctx = crate::context::Context::new(crate::context::GatewayRequest {
            method: "GET".into(),
            path: "/".into(),
            host: "h".into(),
            scheme: "http".into(),
            headers: HashMap::new(),
            query_params: HashMap::new(),
            body: bytes::Bytes::new(),
            remote_addr: "1.2.3.4:5".into(),
            protocol: crate::context::Protocol::Http1,
        });
        attach_consumer(&mut ctx, &alice, "key-auth");
        assert_eq!(
            ctx.message.get("consumer.name"),
            Some(&serde_json::json!("alice"))
        );
        assert_eq!(
            ctx.message.get("consumer.group"),
            Some(&serde_json::json!("partners"))
        );
        assert_eq!(
            ctx.message.get("consumer.custom_id"),
            Some(&serde_json::json!("42"))
        );
        assert_eq!(
            ctx.message.get("consumer.auth_type"),
            Some(&serde_json::json!("key-auth"))
        );
        assert_eq!(
            ctx.request.headers.get("x-consumer-username"),
            Some(&vec!["alice".to_string()])
        );
        assert_eq!(
            ctx.request.headers.get("x-consumer-custom-id"),
            Some(&vec!["42".to_string()])
        );
    }

    #[test]
    fn test_store_rejects_duplicates_and_malformed() {
        // duplicate name
        assert!(ConsumerStore::from_config(&config(serde_json::json!([
            { "name": "a" }, { "name": "a" }
        ])))
        .is_err());

        // shared credential value
        assert!(ConsumerStore::from_config(&config(serde_json::json!([
            { "name": "a", "credentials": { "key-auth": { "key": "same" } } },
            { "name": "b", "credentials": { "key-auth": { "key": "same" } } }
        ])))
        .is_err());

        // credential missing its primary field
        assert!(ConsumerStore::from_config(&config(serde_json::json!([
            { "name": "a", "credentials": { "key-auth": { "wrong": "x" } } }
        ])))
        .is_err());
    }
}