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
// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{Factory, Result, Swarm, Version};
use anyhow::format_err;
use rand::rngs::StdRng;
use std::{env, fs::File, io::Read, num::NonZeroUsize, path::PathBuf};
use tokio::runtime::Runtime;

mod cluster_helper;
mod node;
mod swarm;

pub use cluster_helper::*;
pub use node::K8sNode;
pub use swarm::*;

use diem_sdk::crypto::ed25519::ED25519_PRIVATE_KEY_LENGTH;
use diem_secure_storage::{CryptoStorage, KVStorage, VaultStorage};

const DEFAULT_TESTNET_IMAGE_TAG: &str = "devnet";

pub struct K8sFactory {
    root_key: [u8; ED25519_PRIVATE_KEY_LENGTH],
    treasury_compliance_key: [u8; ED25519_PRIVATE_KEY_LENGTH],
    cluster_name: String,
    helm_repo: String,
    image_tag: String,
    base_image_tag: String,
}

impl K8sFactory {
    pub fn new(
        cluster_name: String,
        helm_repo: String,
        image_tag: String,
        base_image_tag: String,
    ) -> Result<K8sFactory> {
        let vault_addr = env::var("VAULT_ADDR")
            .map_err(|_| format_err!("Expected environment variable VAULT_ADDR"))?;
        let vault_cacert = env::var("VAULT_CACERT")
            .map_err(|_| format_err!("Expected environment variable VAULT_CACERT"))?;
        let vault_token = env::var("VAULT_TOKEN")
            .map_err(|_| format_err!("Expected environment variable VAULT_TOKEN"))?;

        let vault_cacert_path = PathBuf::from(vault_cacert.clone());

        let mut vault_cacert_file = File::open(vault_cacert_path)
            .map_err(|_| format_err!("Failed to open VAULT_CACERT file at {}", &vault_cacert))?;
        let mut vault_cacert_contents = String::new();
        vault_cacert_file
            .read_to_string(&mut vault_cacert_contents)
            .map_err(|_| format_err!("Failed to read VAULT_CACERT file at {}", &vault_cacert))?;

        let vault = VaultStorage::new(
            vault_addr,
            vault_token,
            Some(vault_cacert_contents),
            None,
            false,
            None,
            None,
        );
        vault.available()?;
        let root_key = vault
            .export_private_key("diem__diem_root")
            .unwrap()
            .to_bytes();
        let treasury_compliance_key = vault
            .export_private_key("diem__treasury_compliance")
            .unwrap()
            .to_bytes();

        Ok(Self {
            root_key,
            treasury_compliance_key,
            cluster_name,
            helm_repo,
            image_tag,
            base_image_tag,
        })
    }
}

impl Factory for K8sFactory {
    fn versions<'a>(&'a self) -> Box<dyn Iterator<Item = Version> + 'a> {
        let version = vec![
            Version::new(0, self.base_image_tag.clone()),
            Version::new(1, self.image_tag.clone()),
        ];
        Box::new(version.into_iter())
    }

    fn launch_swarm(
        &self,
        _rng: &mut StdRng,
        node_num: NonZeroUsize,
        version: &Version,
    ) -> Result<Box<dyn Swarm>> {
        set_eks_nodegroup_size(self.cluster_name.clone(), node_num.get(), true)?;
        uninstall_from_k8s_cluster()?;
        clean_k8s_cluster(
            self.helm_repo.clone(),
            node_num.get(),
            format!("{}", version),
            DEFAULT_TESTNET_IMAGE_TAG.to_string(),
            true,
        )?;
        let rt = Runtime::new().unwrap();
        let swarm = rt
            .block_on(K8sSwarm::new(
                &self.root_key,
                &self.treasury_compliance_key,
                &self.cluster_name,
                &self.helm_repo,
                &self.image_tag,
                &self.base_image_tag,
            ))
            .unwrap();
        Ok(Box::new(swarm))
    }
}