1
// Copyright (C) Moondance Labs Ltd.
2
// This file is part of Tanssi.
3

            
4
// Tanssi is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Tanssi is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>.
16

            
17
use {
18
    cumulus_primitives_core::{relay_chain::HeadData, ParaId},
19
    dp_collator_assignment::AssignedCollators,
20
    dp_core::well_known_keys,
21
    frame_support::Hashable,
22
    parity_scale_codec::Encode,
23
    sp_runtime::traits::{BlakeTwo256, HashingFor},
24
    sp_state_machine::Backend,
25
    sp_trie::{PrefixedMemoryDB, StorageProof},
26
};
27

            
28
/// Enum representing how we want to insert the Header
29
#[derive(Clone)]
30
pub enum HeaderAs {
31
    AlreadyEncoded(Vec<u8>),
32
    NonEncoded(sp_runtime::generic::Header<u32, BlakeTwo256>),
33
}
34

            
35
/// Builds a sproof (portmanteau of 'spoof' and 'proof') of the relay chain state.
36
#[derive(Clone)]
37
pub struct ParaHeaderSproofBuilderItem {
38
    /// The para id of the current parachain.
39
    pub para_id: ParaId,
40

            
41
    /// The author_id, which represents a Header with a Aura Digest
42
    pub author_id: HeaderAs,
43
}
44

            
45
impl Default for ParaHeaderSproofBuilderItem {
46
    fn default() -> Self {
47
        Self {
48
            para_id: ParaId::from(200),
49
            author_id: HeaderAs::AlreadyEncoded(vec![]),
50
        }
51
    }
52
}
53

            
54
/// Builds a sproof (portmanteau of 'spoof' and 'proof') of the relay chain state.
55
/// Receives a vec of individual ParaHeaderSproofBuilderItem items of which
56
/// we need to insert the header
57
#[derive(Clone, Default)]
58
pub struct ParaHeaderSproofBuilder {
59
    pub items: Vec<ParaHeaderSproofBuilderItem>,
60
}
61

            
62
impl ParaHeaderSproofBuilder {
63
    pub fn into_state_root_and_proof(
64
        self,
65
    ) -> (
66
        cumulus_primitives_core::relay_chain::Hash,
67
        sp_state_machine::StorageProof,
68
    ) {
69
        let (db, root) =
70
            PrefixedMemoryDB::<HashingFor<cumulus_primitives_core::relay_chain::Block>>::default_with_root();
71
        let state_version = Default::default(); // for test using default.
72
        let mut backend = sp_state_machine::TrieBackendBuilder::new(db, root).build();
73

            
74
        let mut relevant_keys = Vec::new();
75
        {
76
            use parity_scale_codec::Encode as _;
77

            
78
            let mut insert = |key: Vec<u8>, value: Vec<u8>| {
79
                relevant_keys.push(key.clone());
80
                backend.insert(vec![(None, vec![(key, Some(value))])], state_version);
81
            };
82

            
83
            for item in self.items {
84
                let para_key = item.para_id.twox_64_concat();
85
                let key = [well_known_keys::PARAS_HEADS_INDEX, para_key.as_slice()].concat();
86

            
87
                let encoded = match item.author_id {
88
                    HeaderAs::AlreadyEncoded(encoded) => encoded,
89
                    HeaderAs::NonEncoded(header) => header.encode(),
90
                };
91

            
92
                let head_data: HeadData = encoded.into();
93
                insert(key, head_data.encode());
94
            }
95
        }
96

            
97
        let root = *backend.root();
98
        let proof = sp_state_machine::prove_read(backend, relevant_keys).expect("prove read");
99

            
100
        (root, proof)
101
    }
102

            
103
    pub fn relevant_keys(self) -> Vec<Vec<u8>> {
104
        let mut relevant_keys = Vec::new();
105
        {
106
            for item in self.items {
107
                let para_key = item.para_id.twox_64_concat();
108
                let key = [well_known_keys::PARAS_HEADS_INDEX, para_key.as_slice()].concat();
109

            
110
                relevant_keys.push(key.clone());
111
            }
112
        }
113
        relevant_keys
114
    }
115

            
116
    // Construct the proof from an existing state and proof
117
    pub fn from_existing_state(
118
        self,
119
        root: cumulus_primitives_core::relay_chain::Hash,
120
        state: StorageProof,
121
    ) -> (
122
        cumulus_primitives_core::relay_chain::Hash,
123
        sp_state_machine::StorageProof,
124
    ) {
125
        // Recover the db
126
        let db = state.into_memory_db::<HashingFor<cumulus_primitives_core::relay_chain::Block>>();
127

            
128
        // We assume this backend already has the keys injected, and we just need to fetch the proof
129
        let backend = sp_state_machine::TrieBackendBuilder::new(db, root).build();
130

            
131
        // Fetch all existing keys
132
        let mut relevant_keys = backend
133
            .keys(Default::default())
134
            .expect("we should have keys if entering this func")
135
            .map(|result| result.unwrap())
136
            .collect::<Vec<_>>();
137

            
138
        // Fetch relevant keys
139
        for item in self.items {
140
            let para_key = item.para_id.twox_64_concat();
141
            let key = [well_known_keys::PARAS_HEADS_INDEX, para_key.as_slice()].concat();
142
            relevant_keys.push(key.clone());
143
        }
144

            
145
        let proof = sp_state_machine::prove_read(backend, relevant_keys).expect("prove read");
146

            
147
        (root, proof)
148
    }
149

            
150
    // Construct the proof from an existing state and proof
151
    pub fn key_values(self) -> Vec<(Vec<u8>, Vec<u8>)> {
152
        // Fetch all existing keys
153
        let mut key_values = vec![];
154

            
155
        // Fetch relevant keys
156
        for item in self.items {
157
            let para_key = item.para_id.twox_64_concat();
158
            let key = [well_known_keys::PARAS_HEADS_INDEX, para_key.as_slice()].concat();
159

            
160
            let encoded = match item.author_id {
161
                HeaderAs::AlreadyEncoded(encoded) => encoded,
162
                HeaderAs::NonEncoded(header) => header.encode(),
163
            };
164

            
165
            let head_data: HeadData = encoded.into();
166
            key_values.push((key, head_data.encode()))
167
        }
168
        key_values
169
    }
170
}
171

            
172
/// Builds a sproof (portmanteau of 'spoof' and 'proof') of the orchestrator chain state.
173
#[derive(Clone, Encode, Default)]
174
pub struct AuthorityAssignmentSproofBuilder<T> {
175
    pub session_index: u32,
176
    pub authority_assignment: AssignedCollators<T>,
177
}
178

            
179
impl<T: Encode> AuthorityAssignmentSproofBuilder<T> {
180
    pub fn into_state_root_and_proof(
181
        self,
182
    ) -> (
183
        cumulus_primitives_core::relay_chain::Hash,
184
        sp_state_machine::StorageProof,
185
    ) {
186
        let (db, root) =
187
            PrefixedMemoryDB::<HashingFor<cumulus_primitives_core::relay_chain::Block>>::default_with_root();
188

            
189
        let state_version = Default::default();
190
        let mut backend = sp_state_machine::TrieBackendBuilder::new(db, root).build();
191
        let mut relevant_keys = Vec::new();
192

            
193
        let mut insert = |key: Vec<u8>, value: Vec<u8>| {
194
            relevant_keys.push(key.clone());
195
            backend.insert(vec![(None, vec![(key, Some(value))])], state_version);
196
        };
197

            
198
        insert(
199
            well_known_keys::SESSION_INDEX.to_vec(),
200
            self.session_index.encode(),
201
        );
202
        insert(
203
            well_known_keys::authority_assignment_for_session(self.session_index, None).to_vec(),
204
            self.authority_assignment.encode(),
205
        );
206

            
207
        let root = *backend.root();
208
        let proof = sp_state_machine::prove_read(backend, relevant_keys).expect("prove read");
209

            
210
        (root, proof)
211
    }
212

            
213
    pub fn into_state_root_and_proof_solochain(
214
        self,
215
    ) -> (
216
        cumulus_primitives_core::relay_chain::Hash,
217
        sp_state_machine::StorageProof,
218
    ) {
219
        let (db, root) =
220
            PrefixedMemoryDB::<HashingFor<cumulus_primitives_core::relay_chain::Block>>::default_with_root();
221

            
222
        let state_version = Default::default();
223
        let mut backend = sp_state_machine::TrieBackendBuilder::new(db, root).build();
224
        let mut relevant_keys = Vec::new();
225

            
226
        let mut insert = |key: Vec<u8>, value: Vec<u8>| {
227
            relevant_keys.push(key.clone());
228
            backend.insert(vec![(None, vec![(key, Some(value))])], state_version);
229
        };
230

            
231
        insert(
232
            well_known_keys::SESSION_INDEX.to_vec(),
233
            self.session_index.encode(),
234
        );
235
        insert(
236
            well_known_keys::authority_assignment_for_session(
237
                self.session_index,
238
                Some(well_known_keys::SOLOCHAIN_AUTHORITY_ASSIGNMENT_PREFIX),
239
            )
240
            .to_vec(),
241
            self.authority_assignment.encode(),
242
        );
243

            
244
        let root = *backend.root();
245
        let proof = sp_state_machine::prove_read(backend, relevant_keys).expect("prove read");
246

            
247
        (root, proof)
248
    }
249
}