summaryrefslogtreecommitdiffstats
path: root/packages/meshbay-node/src/meshbay_node/replication.py
blob: c2f5cd33b13fb15f301013f3c9875bd5e563e996 (plain) (blame)
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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
"""
MeshBay Node — content replication (node-to-node, admin-authorized).

A replication node downloads files from a source node and stores them
locally, then registers itself as an additional swarm source in the hub.
This provides redundancy and improves availability for public content.

Only public content is replicated (no GEK needed).
Private content replication requires the GEK and is admin-controlled.

Usage:
    replicator = ContentReplicator(
        hub_url=..., access_token=..., source_endpoint=...,
        local_dir=Path("/data/replicated"), node_pk_b64=...,
    )
    await replicator.replicate_file(file_id, file_name, file_size)
"""

import asyncio
import hashlib
import logging
from pathlib import Path

import blake3
import httpx

log = logging.getLogger(__name__)

CHUNK_SIZE = 1024 * 1024   # 1 MB


class ContentReplicator:
    """
    Downloads public files from a source node and registers as swarm source.
    """

    def __init__(
        self,
        hub_url: str,
        access_token: str,
        source_endpoint: str,    # "http://ip:port" of source node HTTP API
        local_dir: Path,
        node_pk_b64: str,
    ):
        self._hub_url        = hub_url.rstrip("/")
        self._access_token   = access_token
        self._source_endpoint = source_endpoint.rstrip("/")
        self._local_dir      = local_dir
        self._node_pk_b64    = node_pk_b64
        self._local_dir.mkdir(parents=True, exist_ok=True)

    @property
    def _auth_headers(self) -> dict:
        return {"Authorization": f"Bearer {self._access_token}"}

    async def fetch_index(self) -> list[dict]:
        """Fetch the public Mesh Group Index from the source node."""
        async with httpx.AsyncClient(timeout=30) as c:
            r = await c.get(f"{self._source_endpoint}/index")
            r.raise_for_status()
            return r.json()["entries"]

    async def replicate_file(
        self,
        file_id: str,
        file_name: str,
        file_size: int,
        progress_cb=None,
    ) -> Path:
        """
        Download a public file from the source node, verify integrity,
        save locally, and register as swarm source in the hub.
        Returns the local file path.
        """
        local_path = self._local_dir / file_name
        if local_path.exists():
            # Verify hash
            existing_hash = blake3.blake3(local_path.read_bytes()).hexdigest()
            if existing_hash == file_id:
                log.info("Already have %s, skipping", file_name)
                await self._register_swarm(file_id, local_path)
                return local_path

        log.info("Replicating %s (%d bytes) from %s", file_name, file_size, self._source_endpoint)

        # Stream download chunk by chunk
        n_chunks = max(1, (file_size + CHUNK_SIZE - 1) // CHUNK_SIZE)
        with open(local_path, "wb") as f:
            async with httpx.AsyncClient(timeout=60) as c:
                for chunk_idx in range(n_chunks):
                    # Download full file (simpler for public content)
                    if chunk_idx == 0:
                        r = await c.get(
                            f"{self._source_endpoint}/file/{file_id}",
                            headers=self._auth_headers,
                        )
                        r.raise_for_status()
                        f.write(r.content)
                        if progress_cb:
                            progress_cb(len(r.content), file_size)
                        break   # full file downloaded in one request

        # Verify hash
        actual_hash = blake3.blake3(local_path.read_bytes()).hexdigest()
        if actual_hash != file_id:
            local_path.unlink(missing_ok=True)
            raise ValueError(f"Hash mismatch: expected {file_id[:16]}, got {actual_hash[:16]}")

        log.info("Replicated %s (hash OK)", file_name)
        await self._register_swarm(file_id, local_path)
        return local_path

    async def _register_swarm(self, content_hash: str, local_path: Path) -> None:
        """Register this node as a swarm source for the content hash in the hub."""
        try:
            async with httpx.AsyncClient(timeout=10) as c:
                r = await c.post(
                    f"{self._hub_url}/v1/swarm/register",
                    json={"content_hash": content_hash, "endpoint": self._source_endpoint},
                    headers=self._auth_headers,
                )
                r.raise_for_status()
                log.debug("Registered as swarm source for %s", content_hash[:16])
        except Exception as e:
            log.warning("Failed to register swarm source: %s", e)

    async def replicate_all(self, progress_cb=None) -> list[Path]:
        """Replicate all public files from the source node."""
        entries = await self.fetch_index()
        results = []
        for entry in entries:
            try:
                path = await self.replicate_file(
                    file_id=entry["id"],
                    file_name=entry["name"],
                    file_size=entry["size"],
                    progress_cb=progress_cb,
                )
                results.append(path)
            except Exception as e:
                log.error("Failed to replicate %s: %s", entry["name"], e)
        return results