plugin_pdo.rs

// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::{log_exception, Plugin};
use crate::{
    component::COMPONENT_PHP_PDO_ID,
    context::RequestContext,
    execute::{get_this_mut, validate_num_args, AfterExecuteHook, BeforeExecuteHook},
    tag::{TAG_DB_STATEMENT, TAG_DB_TYPE},
};
use anyhow::Context;
use dashmap::DashMap;
use once_cell::sync::Lazy;
use phper::{
    arrays::ZArr,
    classes::ClassEntry,
    objects::ZObj,
    sys,
    values::{ExecuteData, ZVal},
};
use skywalking::{
    proto::v3::SpanLayer,
    trace::span::{HandleSpanObject, Span},
};
use std::{any::Any, str::FromStr};
use tracing::{debug, warn};

static DSN_MAP: Lazy<DashMap<u32, Dsn>> = Lazy::new(Default::default);
static DTOR_MAP: Lazy<DashMap<u32, sys::zend_object_dtor_obj_t>> = Lazy::new(Default::default);

#[derive(Default, Clone)]
pub struct PdoPlugin;

impl Plugin for PdoPlugin {
    fn class_names(&self) -> Option<&'static [&'static str]> {
        Some(&["PDO", "PDOStatement"])
    }

    fn function_name_prefix(&self) -> Option<&'static str> {
        None
    }

    fn hook(
        &self, class_name: Option<&str>, function_name: &str,
    ) -> Option<(
        Box<crate::execute::BeforeExecuteHook>,
        Box<crate::execute::AfterExecuteHook>,
    )> {
        match (class_name, function_name) {
            (Some("PDO"), "__construct") => Some(self.hook_pdo_construct()),
            (Some("PDO"), f)
                if [
                    "exec",
                    "query",
                    "prepare",
                    "commit",
                    "begintransaction",
                    "rollback",
                ]
                .contains(&f) =>
            {
                Some(self.hook_pdo_methods(function_name))
            }
            (Some("PDOStatement"), f)
                if ["execute", "fetch", "fetchAll", "fetchColumn", "fetchObject"].contains(&f) =>
            {
                Some(self.hook_pdo_statement_methods(function_name))
            }
            _ => None,
        }
    }
}

impl PdoPlugin {
    fn hook_pdo_construct(&self) -> (Box<BeforeExecuteHook>, Box<AfterExecuteHook>) {
        (
            Box::new(|request_id, execute_data| {
                validate_num_args(execute_data, 1)?;

                let this = get_this_mut(execute_data)?;
                let handle = this.handle();
                hack_dtor(this, Some(pdo_dtor));

                let dsn = execute_data.get_parameter(0);
                let dsn = dsn.as_z_str().context("dsn isn't str")?.to_str()?;
                debug!(dsn, handle, "construct PDO");

                let dsn: Dsn = dsn.parse()?;
                debug!(?dsn, "parse PDO dsn");

                let span = create_exit_span_with_dsn(request_id, "PDO", "__construct", &dsn)?;

                DSN_MAP.insert(handle, dsn);

                Ok(Box::new(span))
            }),
            Box::new(move |_, span, _, _| {
                let mut span = span.downcast::<Span>().unwrap();
                log_exception(&mut *span);
                Ok(())
            }),
        )
    }

    fn hook_pdo_methods(
        &self, function_name: &str,
    ) -> (Box<BeforeExecuteHook>, Box<AfterExecuteHook>) {
        let function_name = function_name.to_owned();
        (
            Box::new(move |request_id, execute_data| {
                let handle = get_this_mut(execute_data)?.handle();

                debug!(handle, function_name, "call PDO method");

                let mut span = with_dsn(handle, |dsn| {
                    create_exit_span_with_dsn(request_id, "PDO", &function_name, dsn)
                })?;

                if execute_data.num_args() >= 1 {
                    if let Some(statement) = execute_data.get_parameter(0).as_z_str() {
                        span.add_tag(TAG_DB_STATEMENT, statement.to_str()?);
                    }
                }

                Ok(Box::new(span) as _)
            }),
            Box::new(after_hook),
        )
    }

    fn hook_pdo_statement_methods(
        &self, function_name: &str,
    ) -> (Box<BeforeExecuteHook>, Box<AfterExecuteHook>) {
        let function_name = function_name.to_owned();
        (
            Box::new(move |request_id, execute_data| {
                let this = get_this_mut(execute_data)?;
                let handle = this.handle();

                debug!(handle, function_name, "call PDOStatement method");

                let mut span = with_dsn(handle, |dsn| {
                    create_exit_span_with_dsn(request_id, "PDOStatement", &function_name, dsn)
                })?;

                if let Some(query) = this.get_property("queryString").as_z_str() {
                    span.add_tag(TAG_DB_STATEMENT, query.to_str()?);
                } else {
                    warn!("PDOStatement queryString is empty");
                }

                Ok(Box::new(span) as _)
            }),
            Box::new(after_hook),
        )
    }
}

fn hack_dtor(this: &mut ZObj, new_dtor: sys::zend_object_dtor_obj_t) {
    let handle = this.handle();

    unsafe {
        let ori_dtor = (*(*this.as_mut_ptr()).handlers).dtor_obj;
        DTOR_MAP.insert(handle, ori_dtor);
        (*((*this.as_mut_ptr()).handlers as *mut sys::zend_object_handlers)).dtor_obj = new_dtor;
    }
}

unsafe extern "C" fn pdo_dtor(object: *mut sys::zend_object) {
    debug!("call PDO dtor");
    dtor(object);
}

unsafe extern "C" fn pdo_statement_dtor(object: *mut sys::zend_object) {
    debug!("call PDOStatement dtor");
    dtor(object);
}

unsafe extern "C" fn dtor(object: *mut sys::zend_object) {
    let handle = ZObj::from_ptr(object).handle();

    DSN_MAP.remove(&handle);
    if let Some((_, Some(dtor))) = DTOR_MAP.remove(&handle) {
        dtor(object);
    }
}

fn after_hook(
    _: Option<i64>, span: Box<dyn Any>, execute_data: &mut ExecuteData, return_value: &mut ZVal,
) -> crate::Result<()> {
    let mut span = span.downcast::<Span>().unwrap();

    if log_exception(&mut *span).is_some() {
        return Ok(());
    }

    if let Some(b) = return_value.as_bool() {
        if !b {
            return after_hook_when_false(get_this_mut(execute_data)?, &mut span);
        }
    } else if let Some(obj) = return_value.as_mut_z_obj() {
        let cls = obj.get_class();
        let pdo_cls = ClassEntry::from_globals("PDOStatement").unwrap();
        if cls.is_instance_of(pdo_cls) {
            return after_hook_when_pdo_statement(get_this_mut(execute_data)?, obj);
        } else {
            let cls = cls.get_name().to_str()?;
            debug!(cls, "not a subclass of PDOStatement");
        }
    }

    Ok(())
}

fn after_hook_when_false(this: &mut ZObj, span: &mut Span) -> crate::Result<()> {
    let info = this.call("errorInfo", [])?;
    let info = info.as_z_arr().context("errorInfo isn't array")?;

    let state = get_error_info_item(info, 0)?.expect_z_str()?.to_str()?;
    let code = {
        let code = get_error_info_item(info, 1)?;
        // PDOStatement::fetch
        // In all cases, false is returned on failure or if there are no more rows.
        if code.get_type_info().is_null() {
            return Ok(());
        }

        &code.expect_long()?.to_string()
    };
    let error = get_error_info_item(info, 2)?.expect_z_str()?.to_str()?;

    let span_object = span.span_object_mut();
    span_object.is_error = true;
    span_object.add_log([("SQLSTATE", state), ("Error Code", code), ("Error", error)]);

    Ok(())
}

fn after_hook_when_pdo_statement(pdo: &mut ZObj, pdo_statement: &mut ZObj) -> crate::Result<()> {
    let dsn = DSN_MAP
        .get(&pdo.handle())
        .map(|r| r.value().clone())
        .context("DSN not found")?;
    let handle = pdo_statement.handle();
    debug!(?dsn, handle, "Hook PDOStatement class");
    DSN_MAP.insert(handle, dsn);
    hack_dtor(pdo_statement, Some(pdo_statement_dtor));
    Ok(())
}

fn get_error_info_item(info: &ZArr, i: u64) -> anyhow::Result<&ZVal> {
    info.get(i)
        .with_context(|| format!("errorInfo[{}] not exists", i))
}

fn create_exit_span_with_dsn(
    request_id: Option<i64>, class_name: &str, function_name: &str, dsn: &Dsn,
) -> anyhow::Result<Span> {
    RequestContext::try_with_global_ctx(request_id, |ctx| {
        let mut span =
            ctx.create_exit_span(&format!("{}->{}", class_name, function_name), &dsn.peer);

        let span_object = span.span_object_mut();
        span_object.set_span_layer(SpanLayer::Database);
        span_object.component_id = COMPONENT_PHP_PDO_ID;
        span_object.add_tag(TAG_DB_TYPE, &dsn.db_type);
        span_object.add_tag("db.data_source", &dsn.data_source);

        Ok(span)
    })
}

fn with_dsn<T>(handle: u32, f: impl FnOnce(&Dsn) -> anyhow::Result<T>) -> anyhow::Result<T> {
    DSN_MAP
        .get(&handle)
        .map(|r| f(r.value()))
        .context("dsn not exists")?
}

#[derive(Debug, Clone)]
struct Dsn {
    db_type: String,
    data_source: String,
    peer: String,
}

impl FromStr for Dsn {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut ss = s.splitn(2, ':');
        let db_type = ss.next().context("unknown db type")?.to_owned();
        let data_source = ss.next().context("unknown datasource")?.to_owned();

        let mut host = "unknown";
        let mut port = match &*db_type {
            "mysql" => "3306",
            "oci" => "1521", // Oracle
            "sqlsrv" => "1433",
            "pgsql" => "5432",
            _ => "0",
        };

        let ss = data_source.split(';');
        for s in ss {
            if s.is_empty() {
                continue;
            }

            let mut kv = s.splitn(2, '=');
            let k = kv.next().context("unknown key")?;
            let v = kv.next().context("unknown value")?;

            // TODO compact the fields rather than mysql.
            match k {
                "host" => {
                    host = v;
                }
                "port" => {
                    port = v;
                }
                _ => {}
            }
        }

        let peer = if host.contains(':') {
            host.to_string()
        } else {
            host.to_string() + ":" + port
        };

        Ok(Dsn {
            db_type,
            data_source,
            peer,
        })
    }
}