signer_remote/storage/

client.rs

//! 存储客户端实现
use std::{io::SeekFrom, path::PathBuf};

use base64::Engine as _;
use futures::stream::{self, StreamExt, TryStreamExt};
use sha3::Digest;
use signer_core::{SignerKeys, SignerUser};
use tokio::io::{AsyncReadExt as _, AsyncSeekExt, AsyncWriteExt, BufReader};

use crate::{
    error::{RemoteError, RemoteResult},
    remote::{HttpClient, HttpClientConfig},
    storage::{
        StorageBlobsCheckExistenceRequest, StorageBlobsCheckExistenceResponse, StorageItemData,
        StorageItemResponse,
    },
};

pub struct SignerDaemonStorage {
    addr: String,
    keys: SignerKeys,
    user: SignerUser,
}

#[derive(Debug, Clone)]
struct ChunkDescriptor {
    hash: String,
    offset: u64,
    length: u64,
}

impl SignerDaemonStorage {
    pub fn new(addr: String, keys: SignerKeys, user: SignerUser) -> Self {
        Self { addr, keys, user }
    }

    pub async fn upload(&self, path: PathBuf) -> RemoteResult<StorageItemResponse> {
        let file = tokio::fs::File::open(&path)
            .await
            .map_err(RemoteError::Io)?;
        let mut reader = BufReader::new(file);

        // 1. Read file in chunks and compute hashes, but only store descriptors
        let mut chunk_descriptors = Vec::new();
        let mut current_offset = 0;

        loop {
            let mut buffer = vec![0; 16 * 1024 * 1024];
            let n = reader.read(&mut buffer).await.map_err(RemoteError::Io)?;
            if n == 0 {
                break;
            }
            buffer.truncate(n);

            let mut hasher = sha3::Sha3_256::new();
            hasher.update(&buffer);
            let hash = base64::prelude::BASE64_URL_SAFE.encode(hasher.finalize().to_vec());

            chunk_descriptors.push(ChunkDescriptor {
                hash,
                offset: current_offset,
                length: n as u64,
            });

            current_offset += n as u64;
        }

        let blob_hashes: Vec<String> = chunk_descriptors.iter().map(|d| d.hash.clone()).collect();

        // 2. Compute item hash
        let mut item_hasher = sha3::Sha3_256::new();
        for hash in &blob_hashes {
            item_hasher.update(hash.as_bytes());
        }
        let item_hash = base64::prelude::BASE64_URL_SAFE.encode(item_hasher.finalize().to_vec());

        let config = HttpClientConfig::new(self.keys.clone(), self.user.clone(), self.addr.clone());
        let client = HttpClient::new(config);

        // 3. Check if item exists for the user
        let head_response = client
            .head(&format!(
                "/api/storage/item/{}/{}",
                &self.user.pub_key, &item_hash
            ))
            .await;

        match head_response {
            Ok(response) => {
                if response.status().is_success() {
                    return Ok(StorageItemResponse {
                        pubkey: self.user.pub_key.clone(),
                        hash: item_hash,
                        file_name: path.file_name().unwrap().to_str().unwrap().to_string(),
                    });
                }
            }
            Err(_) => {
                // 如果 HEAD 请求失败，继续执行后续逻辑
            }
        }

        // 尝试创建项目，如果返回 421 状态码则表示项目已存在
        let create_response = client
            .post_raw(
                "/api/storage/item",
                &StorageItemData {
                    chunks: blob_hashes.clone(),
                    file_name: path.file_name().unwrap().to_str().unwrap().to_string(),
                },
            )
            .await;

        match create_response {
            Ok(response) => {
                if response.status() == reqwest::StatusCode::UNAVAILABLE_FOR_LEGAL_REASONS {
                    let item: StorageItemResponse = response
                        .json()
                        .await
                        .map_err(|e| RemoteError::Internal(e.to_string()))?;
                    return Ok(item);
                }
            }
            Err(_) => {
                // 继续执行后续逻辑
            }
        }

        // 4. Check for missing blobs
        let missing_blobs: StorageBlobsCheckExistenceResponse = client
            .post(
                "/api/storage/blobs/check-existence",
                &StorageBlobsCheckExistenceRequest {
                    hashes: blob_hashes.clone(),
                },
            )
            .await?;

        let missing_hashes: std::collections::HashSet<String> =
            missing_blobs.missing_hashes.into_iter().collect();

        // 5. Upload missing blobs by streaming from the file
        let descriptors_to_upload: Vec<ChunkDescriptor> = chunk_descriptors
            .into_iter()
            .filter(|d| missing_hashes.contains(&d.hash))
            .collect();

        let upload_futs = descriptors_to_upload.into_iter().map(|desc| {
            let keys = self.keys.clone();
            let user = self.user.clone();
            let addr = self.addr.clone();
            let path = path.clone();
            async move {
                let mut file = tokio::fs::File::open(&path)
                    .await
                    .map_err(RemoteError::Io)?;
                file.seek(SeekFrom::Start(desc.offset))
                    .await
                    .map_err(RemoteError::Io)?;
                let mut chunk_reader = file.take(desc.length);
                let mut chunk_data = Vec::with_capacity(desc.length as usize);
                chunk_reader
                    .read_to_end(&mut chunk_data)
                    .await
                    .map_err(RemoteError::Io)?;

                let config = HttpClientConfig::new(keys, user, addr.clone());
                let client = HttpClient::new(config);

                let file_data = reqwest::multipart::Part::bytes(chunk_data);
                let form_data = reqwest::multipart::Form::new().part("blob", file_data);

                client
                    .post_multipart("/api/storage/blob", form_data)
                    .await?
                    .error_for_status()
                    .map_err(RemoteError::Network)?;

                let result: RemoteResult<()> = Ok(());
                result
            }
        });

        stream::iter(upload_futs)
            .buffer_unordered(10)
            .try_collect::<Vec<()>>()
            .await?;

        // 6. Create the item
        let item: StorageItemResponse = client
            .post(
                "/api/storage/item",
                &StorageItemData {
                    chunks: blob_hashes,
                    file_name: path.file_name().unwrap().to_str().unwrap().to_string(),
                },
            )
            .await?;

        Ok(item)
    }

    pub async fn download(&self, key: StorageItemResponse, path: PathBuf) -> RemoteResult<()> {
        let config = HttpClientConfig::new(self.keys.clone(), self.user.clone(), self.addr.clone());
        let client = HttpClient::new(config);

        let r = client
            .get_raw(&format!("/api/storage/item/{}/{}", &key.pubkey, &key.hash))
            .await?
            .error_for_status()
            .map_err(RemoteError::Network)?;

        let mut f = tokio::fs::File::create(path)
            .await
            .map_err(RemoteError::Io)?;
        let mut response_stream = r.bytes_stream();
        while let Some(chunk) = response_stream
            .try_next()
            .await
            .map_err(RemoteError::Network)?
        {
            f.write_all(&chunk).await.map_err(RemoteError::Io)?;
        }

        Ok(())
    }
}