feat: implement Phase 2 Civitai Guard (v0.2)

Cargo.toml

@@ -11,6 +11,8 @@ tokio-stream = "0.1"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 anyhow = "1.0"
+notify = "6.1"
+uuid = { version = "1.0", features = ["v4"] }
 
 [build-dependencies]
 tonic-build = "0.11"

STRUCTURE.tree

@@ -1,6 +1,5 @@
 .
 ├── ARCHITECTURE.md
-├── .benchmarks
 ├── buf.gen.yaml
 ├── buf.yaml
 ├── build.rs
@@ -40,28 +39,22 @@
 │   └── vtuber_image
 │       └── v1
 │           └── image.proto
-├── .pytest_cache
-│   ├── CACHEDIR.TAG
-│   ├── .gitignore
-│   ├── README.md
-│   └── v
-│       └── cache
-│           └── nodeids
 ├── python
 │   ├── comfy_client.py
-│   ├── __pycache__
-│   │   └── comfy_client.cpython-314.pyc
 │   ├── requirements.txt
 │   └── test_comfy_client.py
 ├── README.md
 ├── ROADMAP.md
 ├── SECURITY.md
 ├── src
+│   ├── guard
+│   │   ├── cache.rs
+│   │   └── mod.rs
 │   └── main.rs
 ├── STRATEGY.md
 ├── STRUCTURE.tree
 ├── SUPPORT.md
 ├── TROUBLESHOOTING.md
 └── VISION.md
 
-15 directories, 50 files
+11 directories, 47 files

deploy/k8s/vtuber-image.yaml

@@ -0,0 +1,50 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: vtuber-image
+  namespace: vtuber-image
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: vtuber-image
+  template:
+    metadata:
+      labels:
+        app: vtuber-image
+    spec:
+      containers:
+        - name: vtuber-image
+          image: echo-layer/vtuber-image:latest
+          ports:
+            - containerPort: 8083
+          env:
+            - name: CONFIG_PATH
+              value: "/mnt/config"
+            - name: ALLOWLIST_FILENAME
+              value: "allowlist.json"
+            - name: S3_ENDPOINT_URL
+              value: "http://seaweedfs-s3:8333"
+            - name: COMFYUI_URL
+              value: "http://comfyui:8188"
+          volumeMounts:
+            - name: config-volume
+              mountPath: /mnt/config
+              readOnly: true
+
+        # git-sync Sidecar to sync vtuber-commons
+        - name: git-sync
+          image: registry.k8s.io/git-sync/git-sync:v4.2.3
+          args:
+            - "--repo=https://github.com/echo-layer/vtuber-commons"
+            - "--branch=main"
+            - "--root=/mnt/config"
+            - "--dest=vtuber-commons"
+            - "--wait=60"
+          volumeMounts:
+            - name: config-volume
+              mountPath: /mnt/config
+
+      volumes:
+        - name: config-volume
+          emptyDir: {}

python/comfy_client.py

@@ -80,6 +80,45 @@ def upload_result(self, local_filename, target_bucket, target_key):
         )
         return f"s3://{target_bucket}/{target_key}"
 
+    def verify_model(self, model_id, expected_hash, allow_nsfw):
+        print(f"Verifying model {model_id} on Civitai...", file=sys.stderr)
+        # Using a timeout to avoid hanging
+        try:
+            response = requests.get(f"https://civitai.com/api/v1/models/{model_id}", timeout=10)
+            if response.status_code != 200:
+                raise Exception(f"Failed to fetch metadata for model {model_id} from Civitai: {response.status_code}")
+
+            metadata = response.json()
+
+            # Check NSFW if restricted
+            if not allow_nsfw and metadata.get('nsfw', False):
+                raise Exception(f"Model {model_id} is marked as NSFW, but NSFW is not allowed.")
+
+            found_hash = False
+            versions = metadata.get('modelVersions', [])
+            if not versions:
+                raise Exception(f"No versions found for model {model_id}")
+
+            # We check all versions for the hash to be safe, though usually it's the latest
+            for version in versions:
+                for file in version.get('files', []):
+                    hashes = file.get('hashes', {})
+                    sha256 = hashes.get('SHA256')
+                    if sha256:
+                        if sha256.lower() == expected_hash.lower():
+                            found_hash = True
+                            break
+                if found_hash:
+                    break
+
+            if not found_hash:
+                raise Exception(f"SHA256 hash mismatch for model {model_id}. Expected {expected_hash}")
+
+            print(f"Model {model_id} verified successfully.", file=sys.stderr)
+            return True
+        except requests.exceptions.RequestException as e:
+            raise Exception(f"Network error verifying model {model_id}: {str(e)}")
+
 if __name__ == "__main__":
     client = ComfyClient()
 
@@ -97,6 +136,15 @@ def upload_result(self, local_filename, target_bucket, target_key):
         # 2. Inject overrides
         workflow = client.inject_overrides(workflow, req.get('overrides', {}))
 
+        # 2.5 Verify models
+        model_auth = req.get('model_auth', [])
+        for auth in model_auth:
+            client.verify_model(
+                auth['model_id'],
+                auth['expected_hash'],
+                auth.get('allow_nsfw', False)
+            )
+
         # 3. Queue prompt
         prompt_response = client.queue_prompt(workflow)
         prompt_id = prompt_response['prompt_id']

python/test_comfy_client.py

@@ -45,5 +45,64 @@ def test_upload_result(self, mock_get):
         self.assertEqual(kwargs['Key'], "test_key")
         self.assertEqual(kwargs['Body'], b"fake_image_data")
 
+    @patch('requests.get')
+    def test_verify_model_success(self, mock_get):
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "nsfw": False,
+            "modelVersions": [
+                {
+                    "files": [
+                        {
+                            "hashes": {
+                                "SHA256": "ABCDEF123456"
+                            }
+                        }
+                    ]
+                }
+            ]
+        }
+        mock_get.return_value = mock_response
+
+        result = self.client.verify_model("123", "abcdef123456", False)
+        self.assertTrue(result)
+
+    @patch('requests.get')
+    def test_verify_model_nsfw_rejected(self, mock_get):
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "nsfw": True,
+            "modelVersions": []
+        }
+        mock_get.return_value = mock_response
+
+        with self.assertRaisesRegex(Exception, "Model 123 is marked as NSFW"):
+            self.client.verify_model("123", "hash", False)
+
+    @patch('requests.get')
+    def test_verify_model_hash_mismatch(self, mock_get):
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "nsfw": False,
+            "modelVersions": [
+                {
+                    "files": [
+                        {
+                            "hashes": {
+                                "SHA256": "WRONGHASH"
+                            }
+                        }
+                    ]
+                }
+            ]
+        }
+        mock_get.return_value = mock_response
+
+        with self.assertRaisesRegex(Exception, "SHA256 hash mismatch"):
+            self.client.verify_model("123", "expectedhash", False)
+
 if __name__ == '__main__':
     unittest.main()

src/guard/cache.rs

@@ -0,0 +1,53 @@
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::fs;
+use std::path::Path;
+use std::sync::{Arc, RwLock};
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct ModelEntry {
+    pub model_id: String,
+    pub hash: String,
+    pub license: String,
+    pub allow_nsfw: bool,
+}
+
+#[derive(Debug, Clone)]
+pub struct GuardCache {
+    cache: Arc<RwLock<HashMap<String, ModelEntry>>>,
+}
+
+impl Default for GuardCache {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl GuardCache {
+    pub fn new() -> Self {
+        Self {
+            cache: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+
+    pub fn load_from_file<P: AsRef<Path>>(&self, path: P) -> anyhow::Result<()> {
+        let content = fs::read_to_string(path)?;
+        let entries: Vec<ModelEntry> = serde_json::from_str(&content)?;
+
+        let mut cache = self
+            .cache
+            .write()
+            .map_err(|_| anyhow::anyhow!("Failed to acquire write lock"))?;
+        cache.clear();
+        for entry in entries {
+            cache.insert(entry.model_id.clone(), entry);
+        }
+
+        Ok(())
+    }
+
+    pub fn get_model(&self, model_id: &str) -> Option<ModelEntry> {
+        let cache = self.cache.read().ok()?;
+        cache.get(model_id).cloned()
+    }
+}

src/guard/mod.rs

@@ -0,0 +1 @@
+pub mod cache;