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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
use std::env;

use crate::error::RuntimeError;
use log::*;

/// The name of the environment variable in the Lambda execution
/// environment for the Runtime APIs endpoint. The value of this
/// variable is read once as the runtime starts.
pub const RUNTIME_ENDPOINT_VAR: &str = "AWS_LAMBDA_RUNTIME_API";

/// Clone-able generic function settings object. The data is loaded
/// from environment variables during the init process. The data
/// for the object is cloned in the `Context` for each invocation.
#[derive(Clone)]
pub struct FunctionSettings {
    pub function_name: String,
    pub memory_size: i32,
    pub version: String,
    pub log_stream: String,
    pub log_group: String,
}

/// Trait used by the `RustRuntime` module to retrieve configuration information
/// about the environement. This is implemented by the `EnvConfigProvider` using
/// the environment variables. We also have a mock implementation for the unit tests
pub trait ConfigProvider {
    /// Loads the function settings such as name, arn, memory amount, version, etc.
    ///
    /// # Return
    /// A `Result` of `FunctionSettings` or a `RuntimeError`. The runtime
    /// fails the init process if this function returns an error.
    fn get_function_settings(&self) -> Result<FunctionSettings, RuntimeError>;

    /// Returns the endpoint (hostname:port) for the Runtime API endpoint
    fn get_runtime_api_endpoint(&self) -> Result<String, RuntimeError>;
}

/// Implementation of the `ConfigProvider` trait that reads the settings from
/// environment variables in the Lambda execution environment. This is the config
/// used by the `start()` method of this module.
pub struct EnvConfigProvider;

impl std::default::Default for EnvConfigProvider {
    fn default() -> Self {
        EnvConfigProvider
    }
}

impl ConfigProvider for EnvConfigProvider {
    /// Loads the function settings from the Lambda environment variables:
    /// https://docs.aws.amazon.com/lambda/latest/dg/current-supported-versions.html
    fn get_function_settings(&self) -> Result<FunctionSettings, RuntimeError> {
        let function_name = env::var("AWS_LAMBDA_FUNCTION_NAME")?;
        let version = env::var("AWS_LAMBDA_FUNCTION_VERSION")?;
        let log_stream = env::var("AWS_LAMBDA_LOG_STREAM_NAME")?;
        let log_group = env::var("AWS_LAMBDA_LOG_GROUP_NAME")?;
        let memory_str = env::var("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")?;
        let parsed_memory_str = memory_str.parse::<i32>();
        let memory_size: i32;
        match parsed_memory_str {
            Ok(int_value) => memory_size = int_value,
            Err(_parse_err) => {
                error!(
                    "Memory value from environment is not i32: {}",
                    env::var("AWS_LAMBDA_FUNCTION_MEMORY_SIZE").unwrap()
                );
                return Err(RuntimeError::unrecoverable(&format!(
                    "Could not parse memory value: {}",
                    memory_str
                )));
            }
        };

        Ok(FunctionSettings {
            function_name,
            memory_size,
            version,
            log_stream,
            log_group,
        })
    }

    /// Loads the endpoint from Lambda's default environment variable: AWS_LAMBDA_RUNTIME_API
    fn get_runtime_api_endpoint(&self) -> Result<String, RuntimeError> {
        let endpoint = env::var(RUNTIME_ENDPOINT_VAR)?;
        Ok(endpoint)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use crate::{env::*, error};
    use std::{env, error::Error};

    pub(crate) struct MockConfigProvider {
        pub(crate) error: bool,
    }

    impl ConfigProvider for MockConfigProvider {
        fn get_function_settings(&self) -> Result<FunctionSettings, error::RuntimeError> {
            if self.error {
                return Err(error::RuntimeError::unrecoverable("Mock error"));
            }

            Ok(FunctionSettings {
                function_name: String::from("MockFunction"),
                memory_size: 128,
                version: String::from("$LATEST"),
                log_stream: String::from("LogStream"),
                log_group: String::from("LogGroup"),
            })
        }

        fn get_runtime_api_endpoint(&self) -> Result<String, error::RuntimeError> {
            if self.error {
                return Err(error::RuntimeError::unrecoverable("Mock error"));
            }

            Ok(String::from("http://localhost:8080"))
        }
    }

    fn set_endpoint_env_var() {
        env::set_var(RUNTIME_ENDPOINT_VAR, "localhost:8080");
    }

    fn set_lambda_env_vars() {
        env::set_var("AWS_LAMBDA_FUNCTION_NAME", "test_func");
        env::set_var("AWS_LAMBDA_FUNCTION_VERSION", "$LATEST");
        env::set_var("AWS_LAMBDA_LOG_STREAM_NAME", "LogStreamName");
        env::set_var("AWS_LAMBDA_LOG_GROUP_NAME", "LogGroup2");
        env::set_var("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "128");
    }

    fn unset_env_vars() {
        env::remove_var(RUNTIME_ENDPOINT_VAR);
        env::remove_var("AWS_LAMBDA_FUNCTION_NAME");
        env::remove_var("AWS_LAMBDA_FUNCTION_VERSION");
        env::remove_var("AWS_LAMBDA_LOG_STREAM_NAME");
        env::remove_var("AWS_LAMBDA_LOG_GROUP_NAME");
        env::remove_var("AWS_LAMBDA_FUNCTION_MEMORY_SIZE");
    }

    #[test]
    fn function_config_from_env_vars() {
        unset_env_vars();
        set_endpoint_env_var();
        set_lambda_env_vars();
        let config_provider: &dyn ConfigProvider = &EnvConfigProvider {};
        let env_settings = config_provider.get_function_settings();
        assert_eq!(
            env_settings.is_err(),
            false,
            "Env settings returned an error: {}",
            env_settings.err().unwrap().description()
        );
        let settings = env_settings.unwrap();
        assert_eq!(
            settings.memory_size, 128,
            "Invalid memory size: {}",
            settings.memory_size
        );
        let endpoint = config_provider.get_runtime_api_endpoint();
        assert_eq!(
            endpoint.is_err(),
            false,
            "Env endpoint returned an error: {}",
            endpoint.err().unwrap().description()
        );

        unset_env_vars();
        let err_env_settings = config_provider.get_function_settings();
        assert!(
            err_env_settings.is_err(),
            "Env config did not return error without variables"
        );
        let err_endpoint = config_provider.get_runtime_api_endpoint();
        assert!(
            err_endpoint.is_err(),
            "Env endpoint did not return error without variables"
        );
    }
}