1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

use std::{
    env,
    fs::{self, File},
    io::{BufReader, Read, Write},
    path::{Path, PathBuf},
};

use assert_cmd::prelude::OutputOkExt;
use getset::{Getters, MutGetters, Setters};
use rslua::lexer::Lexer;
use tempfile::NamedTempFile;

use crate::{Headers, Result, WrkError};

const LUA_DEFAULT_DONE_FUNCTION: &str = r#"
-- The done() function is called at the end of wrk execution
-- and allows us to produce a well formed JSON output, prefixed
-- by the string "JSON" which allows us to parse the wrk output
-- easily.
done = function(summary, latency, requests)
    local errors = summary.errors.connect
        + summary.errors.read
        + summary.errors.write
        + summary.errors.status
        + summary.errors.timeout
    io.write("JSON")
    io.write(string.format(
        [[{
    "requests": %d,
    "errors": %d,
    "successes": %d,
    "requests_sec": %d,
    "avg_latency_ms": %.2f,
    "min_latency_ms": %.2f,
    "max_latency_ms": %.2f,
    "stdev_latency_ms": %.2f,
    "transfer_mb": %.d,
    "errors_connect": %d,
    "errors_read": %d,
    "errors_write": %d,
    "errors_status": %d,
    "errors_timeout": %d
}
]],
        summary.requests,
        errors,
        summary.requests - errors,
        summary.requests / (summary.duration / 1000000),
        (latency.mean / 1000),
        (latency.min / 1000),
        (latency.max / 1000),
        (latency.stdev / 1000),
        (summary.bytes / 1048576),
        summary.errors.connect,
        summary.errors.read,
        summary.errors.write,
        summary.errors.status,
        summary.errors.timeout
    ))
end
"#;

#[derive(Debug)]
pub struct LuaScript {}

impl LuaScript {
    fn lua_script_from_config(&mut self, uri: &str, method: &str, headers: &Headers, body: &str) -> Result<String> {
        let request = format!(
            r#"
-- The request() function is called by wrk on all requests
-- and allow us to configure things like headers, method, body, etc..
request = function()
    wrk.method = "{}"
    wrk.body = "{}"
    {}
    return wrk.format("{}", "{}")
end
        "#,
            method,
            body,
            self.lua_headers(headers)?,
            method,
            uri
        );
        let buffer = request + LUA_DEFAULT_DONE_FUNCTION;
        Ok(buffer)
    }

    fn lua_script_from_user(&mut self, lua_script: &Path) -> Result<String> {
        let file = File::open(lua_script)?;
        let mut reader = BufReader::new(file);
        let mut buffer = String::new();
        reader.read_to_string(&mut buffer)?;
        let mut lexer = Lexer::new();
        let tokens = lexer.run(&buffer).map_err(|e| WrkError::Lua(format!("{:?}", e)))?;
        let buffer = buffer + LUA_DEFAULT_DONE_FUNCTION;
        Ok(buffer)
    }

    fn lua_headers(&self, headers: &Headers) -> Result<String> {
        let mut result = String::new();
        for (k, v) in headers {
            result += &format!(r#"wrk.headers["{}"] = "{}"\n"#, k, v);
        }
        Ok(result)
    }

    pub fn render(
        script_file: &mut NamedTempFile,
        user_script: Option<&PathBuf>,
        uri: &str,
        method: &str,
        headers: &Headers,
        body: &str,
    ) -> Result<()> {
        let mut this = Self {};
        let script = match user_script {
            Some(lua_script) => {
                if !lua_script.exists() {
                    error!(
                        "Wrk Lua file {} not found in {}",
                        env::current_dir().expect("unable to get current directory").display(),
                        lua_script.display()
                    );
                    return Err(WrkError::Lua("Wrk Lua file not found".to_string()));
                } else {
                    this.lua_script_from_user(&lua_script)?
                }
            }
            None => this.lua_script_from_config(uri, method, headers, body)?,
        };
        script_file.write_all(script.as_bytes())?;
        Ok(())
    }
}