vyn_core/
relay_storage.rs

1use std::path::Path;
2use std::process::{Command, Stdio};
3use std::sync::Arc;
4
5use tonic::metadata::MetadataValue;
6use tonic::Request;
7use tonic::transport::{Channel, Endpoint};
8
9use crate::storage::{StorageError, StorageProvider, StorageResult};
10
11#[derive(Clone)]
12pub struct RelayStorageProvider {
13    endpoint: String,
14    pub token: Arc<tokio::sync::RwLock<Option<String>>>,
15}
16
17impl RelayStorageProvider {
18    pub fn new(endpoint: impl Into<String>) -> Self {
19        Self {
20            endpoint: endpoint.into(),
21            token: Arc::new(tokio::sync::RwLock::new(None)),
22        }
23    }
24
25    async fn connect(
26        &self,
27    ) -> StorageResult<vyn_relay::proto::vyn_relay_client::VynRelayClient<Channel>> {
28        let endpoint = Endpoint::from_shared(self.endpoint.clone())
29            .map_err(|err| StorageError::Transport(err.to_string()))?;
30        vyn_relay::proto::vyn_relay_client::VynRelayClient::connect(endpoint)
31            .await
32            .map_err(|err| StorageError::Transport(err.to_string()))
33    }
34
35    /// Perform the two-step challenge-response auth and store the session token.
36    pub async fn authenticate(
37        &self,
38        user_id: &str,
39        sign_fn: impl Fn(&[u8]) -> StorageResult<String>,
40    ) -> StorageResult<()> {
41        let mut client = self.connect().await?;
42
43        // Step 1: request challenge
44        let challenge_resp = client
45            .authenticate(Request::new(vyn_relay::proto::AuthRequest {
46                user_id: user_id.to_string(),
47                nonce: Vec::new(),
48                signature: String::new(),
49            }))
50            .await
51            .map_err(|e| StorageError::Transport(e.to_string()))?
52            .into_inner();
53
54        let nonce = challenge_resp.challenge_nonce;
55
56        // Step 2: sign nonce and send back
57        let signature = sign_fn(&nonce)?;
58        let auth_resp = client
59            .authenticate(Request::new(vyn_relay::proto::AuthRequest {
60                user_id: user_id.to_string(),
61                nonce: nonce.clone(),
62                signature,
63            }))
64            .await
65            .map_err(|e| StorageError::Transport(e.to_string()))?
66            .into_inner();
67
68        if !auth_resp.authenticated {
69            return Err(StorageError::Transport(
70                "relay authentication failed".to_string(),
71            ));
72        }
73
74        *self.token.write().await = Some(auth_resp.token);
75        Ok(())
76    }
77
78    /// Register identity on the relay (idempotent) then authenticate.
79    /// `vault_dir` is the `.vyn/` directory (contains `identity.toml`).
80    pub async fn authenticate_with_identity(&self, vault_dir: &Path) -> StorageResult<()> {
81        let identity = load_identity(vault_dir)?;
82        let private_key_path = identity.ssh_private_key.clone();
83        let public_key_path = identity.ssh_public_key.clone();
84        let user_id = identity.github_username.clone();
85
86        let public_key = std::fs::read_to_string(&public_key_path)
87            .map_err(|e| StorageError::Transport(format!("failed to read public key: {e}")))?;
88        let public_key = public_key.trim().to_string();
89
90        // Register identity on the relay (no-op if already registered with same key)
91        self.ensure_identity_registered(&user_id, &public_key, &private_key_path)
92            .await?;
93
94        let private_key_path2 = private_key_path.clone();
95        self.authenticate(&user_id, move |nonce| {
96            sign_nonce_with_ssh_key(nonce, Path::new(&private_key_path2))
97        })
98        .await
99    }
100
101    async fn ensure_identity_registered(
102        &self,
103        user_id: &str,
104        public_key: &str,
105        private_key_path: &str,
106    ) -> StorageResult<()> {
107        let registration_payload = format!("vyn-register:{user_id}:{public_key}");
108        let signature = sign_nonce_with_ssh_key(
109            registration_payload.as_bytes(),
110            Path::new(private_key_path),
111        )?;
112
113        let mut client = self.connect().await?;
114        client
115            .register_identity(Request::new(vyn_relay::proto::RegisterRequest {
116                user_id: user_id.to_string(),
117                public_key: public_key.to_string(),
118                signature,
119            }))
120            .await
121            .map_err(|e| StorageError::Transport(e.to_string()))?;
122
123        Ok(())
124    }
125
126    async fn inject_token<T>(&self, mut request: Request<T>) -> StorageResult<Request<T>> {
127        if let Some(ref tok) = *self.token.read().await {
128            let val = MetadataValue::try_from(tok.as_str())
129                .map_err(|e| StorageError::Transport(e.to_string()))?;
130            request.metadata_mut().insert("x-vyn-token", val);
131        }
132        Ok(request)
133    }
134}
135
136impl StorageProvider for RelayStorageProvider {
137    async fn get_manifest(&self, project_id: &str) -> StorageResult<Option<Vec<u8>>> {
138        let mut client = self.connect().await?;
139        let response = client
140            .get_manifest(Request::new(vyn_relay::proto::GetManifestRequest {
141                project_id: project_id.to_string(),
142            }))
143            .await
144            .map_err(|err| StorageError::Transport(err.to_string()))?
145            .into_inner();
146
147        if response.found {
148            Ok(Some(response.payload))
149        } else {
150            Ok(None)
151        }
152    }
153
154    async fn put_manifest(&self, project_id: &str, manifest_payload: &[u8]) -> StorageResult<()> {
155        let mut client = self.connect().await?;
156        let req = self
157            .inject_token(Request::new(vyn_relay::proto::PutManifestRequest {
158                project_id: project_id.to_string(),
159                payload: manifest_payload.to_vec(),
160            }))
161            .await?;
162        client
163            .put_manifest(req)
164            .await
165            .map_err(|err| StorageError::Transport(err.to_string()))?;
166        Ok(())
167    }
168
169    async fn upload_blob(&self, hash: &str, data: Vec<u8>) -> StorageResult<()> {
170        let mut client = self.connect().await?;
171        let message = vyn_relay::proto::UploadBlobChunk {
172            hash: hash.to_string(),
173            chunk: data,
174        };
175
176        let mut req = Request::new(tokio_stream::iter(vec![message]));
177        if let Some(ref tok) = *self.token.read().await {
178            let val = MetadataValue::try_from(tok.as_str())
179                .map_err(|e| StorageError::Transport(e.to_string()))?;
180            req.metadata_mut().insert("x-vyn-token", val);
181        }
182        client
183            .upload_blob(req)
184            .await
185            .map_err(|err| StorageError::Transport(err.to_string()))?;
186        Ok(())
187    }
188
189    async fn download_blob(&self, hash: &str) -> StorageResult<Option<Vec<u8>>> {
190        let mut client = self.connect().await?;
191        let stream = client
192            .download_blob(Request::new(vyn_relay::proto::DownloadBlobRequest {
193                hash: hash.to_string(),
194            }))
195            .await
196            .map_err(|err| {
197                if err.code() == tonic::Code::NotFound {
198                    StorageError::NotFound
199                } else {
200                    StorageError::Transport(err.to_string())
201                }
202            });
203
204        if let Err(StorageError::NotFound) = stream {
205            return Ok(None);
206        }
207
208        let mut stream = stream?.into_inner();
209
210        let mut out = Vec::new();
211        while let Some(chunk) = stream
212            .message()
213            .await
214            .map_err(|err| StorageError::Transport(err.to_string()))?
215        {
216            out.extend_from_slice(&chunk.chunk);
217        }
218
219        Ok(Some(out))
220    }
221
222    async fn create_invite(
223        &self,
224        user_id: &str,
225        vault_id: &str,
226        payload: Vec<u8>,
227    ) -> StorageResult<()> {
228        let mut client = self.connect().await?;
229        let req = self
230            .inject_token(Request::new(vyn_relay::proto::CreateInviteRequest {
231                user_id: user_id.to_string(),
232                vault_id: vault_id.to_string(),
233                payload,
234            }))
235            .await?;
236        client
237            .create_invite(req)
238            .await
239            .map_err(|err| StorageError::Transport(err.to_string()))?;
240        Ok(())
241    }
242
243    async fn get_invites(&self, user_id: &str, vault_id: &str) -> StorageResult<Vec<Vec<u8>>> {
244        let mut client = self.connect().await?;
245        let response = client
246            .get_invites(Request::new(vyn_relay::proto::GetInvitesRequest {
247                user_id: user_id.to_string(),
248                vault_id: vault_id.to_string(),
249            }))
250            .await
251            .map_err(|err| StorageError::Transport(err.to_string()))?
252            .into_inner();
253        Ok(response.payloads)
254    }
255}
256
257struct IdentityConfig {
258    github_username: String,
259    ssh_private_key: String,
260    ssh_public_key: String,
261}
262
263fn load_identity(vault_dir: &Path) -> StorageResult<IdentityConfig> {
264    let path = vault_dir.join("identity.toml");
265    let text = std::fs::read_to_string(&path)
266        .map_err(|e| StorageError::Transport(format!("failed to read identity.toml: {e}")))?;
267    let github_username = parse_toml_string(&text, "github_username")
268        .ok_or_else(|| StorageError::Transport("missing github_username in identity.toml".into()))?;
269    let ssh_private_key = parse_toml_string(&text, "ssh_private_key")
270        .ok_or_else(|| StorageError::Transport("missing ssh_private_key in identity.toml".into()))?;
271    let ssh_public_key = parse_toml_string(&text, "ssh_public_key")
272        .ok_or_else(|| StorageError::Transport("missing ssh_public_key in identity.toml".into()))?;
273    Ok(IdentityConfig { github_username, ssh_private_key, ssh_public_key })
274}
275
276fn parse_toml_string(text: &str, key: &str) -> Option<String> {
277    for line in text.lines() {
278        let line = line.trim();
279        if let Some(rest) = line.strip_prefix(key) {
280            let rest = rest.trim().strip_prefix('=')?;
281            let val = rest.trim().trim_matches('"');
282            return Some(val.to_string());
283        }
284    }
285    None
286}
287
288fn sign_nonce_with_ssh_key(nonce: &[u8], private_key: &Path) -> StorageResult<String> {
289    let tmp = tempfile::TempDir::new()
290        .map_err(|e| StorageError::Transport(format!("failed to create temp dir: {e}")))?;
291    let nonce_file = tmp.path().join("nonce");
292    std::fs::write(&nonce_file, nonce)
293        .map_err(|e| StorageError::Transport(format!("failed to write nonce: {e}")))?;
294
295    let status = Command::new("ssh-keygen")
296        .args([
297            "-Y", "sign",
298            "-f", private_key.to_str().unwrap_or(""),
299            "-n", "vyn",
300            nonce_file.to_str().unwrap_or(""),
301        ])
302        .stdout(Stdio::null())
303        .stderr(Stdio::null())
304        .status()
305        .map_err(|e| StorageError::Transport(format!("failed to run ssh-keygen: {e}")))?;
306
307    if !status.success() {
308        return Err(StorageError::Transport("ssh-keygen signing failed".into()));
309    }
310
311    let sig_file = tmp.path().join("nonce.sig");
312    let sig = std::fs::read_to_string(&sig_file)
313        .map_err(|e| StorageError::Transport(format!("failed to read signature: {e}")))?;
314    Ok(sig)
315}
316
317#[cfg(test)]
318mod tests {
319    use crate::crypto::secret_bytes;
320    use crate::manifest::Manifest;
321    use crate::storage::{StorageProvider, decrypt_manifest, encrypt_manifest};
322
323    use super::RelayStorageProvider;
324
325    #[tokio::test]
326    async fn relay_roundtrip() {
327        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
328            .await
329            .expect("bind test port");
330        let port = listener.local_addr().expect("read local addr").port();
331
332        let data_dir = std::env::temp_dir().join(format!("vyn-relay-it-{}", uuid::Uuid::new_v4()));
333        let data_dir_string = data_dir.to_string_lossy().to_string();
334
335        let handle = tokio::spawn(async move {
336            let _ = vyn_relay::serve_with_listener(listener, data_dir_string).await;
337        });
338
339        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
340
341        let provider = RelayStorageProvider::new(format!("http://127.0.0.1:{port}"));
342        // Inject the test bypass token so write RPCs pass auth in test builds.
343        *provider.token.write().await = Some("test-bypass".to_string());
344
345        let mut uploaded = false;
346        for _ in 0..20 {
347            if provider
348                .upload_blob("blob123", b"hello-relay".to_vec())
349                .await
350                .is_ok()
351            {
352                uploaded = true;
353                break;
354            }
355            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
356        }
357        assert!(uploaded, "blob upload should succeed once relay is ready");
358        let blob = provider
359            .download_blob("blob123")
360            .await
361            .expect("blob download should succeed")
362            .expect("blob should exist");
363        assert_eq!(blob, b"hello-relay");
364
365        let key = secret_bytes(vec![3u8; 32]);
366        let manifest = Manifest {
367            version: 1,
368            files: vec![],
369        };
370        let payload = encrypt_manifest(&manifest, &key).expect("encrypt manifest");
371        provider
372            .put_manifest("vault-it", &payload)
373            .await
374            .expect("put manifest should succeed");
375        let pulled = provider
376            .get_manifest("vault-it")
377            .await
378            .expect("get manifest should succeed")
379            .expect("manifest should exist");
380        let restored = decrypt_manifest(&pulled, &key).expect("decrypt manifest");
381        assert_eq!(restored.version, 1);
382
383        provider
384            .create_invite("alice", "vault-it", b"invite-data".to_vec())
385            .await
386            .expect("create invite should succeed");
387        let invites = provider
388            .get_invites("alice", "vault-it")
389            .await
390            .expect("get invites should succeed");
391        assert_eq!(invites.len(), 1);
392        assert_eq!(invites[0], b"invite-data");
393
394        handle.abort();
395        let _ = std::fs::remove_dir_all(data_dir);
396    }
397}
vyn_core/relay_storage.rs

vyn_core/
relay_storage.rs
