Week 3 - M4: Multi-object transfer + reachable objects

minigit/remote.py

@@ -9,12 +9,12 @@
 Build the `RemoteClient` class here, per the interface contract.
 """
 
-# from .objects import ObjectStore  (Module 1 & 3)
-# from .commits import CommitManager
 import os
 import socket
 
-from minigit.errors import NetworkProtocolError
+from minigit.commits import CommitManager
+from minigit.errors import NetworkProtocolError, ObjectCorruptError, ObjectNotFoundError
+from minigit.objects import ObjectStore
 
 
 def send_line(sock, text: str) -> None:
@@ -40,21 +40,41 @@ def receive_line(sock, buf: bytearray) -> str:
 
     line, _, rest = buf.partition(b"\n")
     buf[:] = rest
-    return line.decode()
+    try:
+        return line.decode("utf-8")
+    except UnicodeDecodeError as exc:
+        raise NetworkProtocolError("protocol line is not UTF-8") from exc
+
+
+def recv_exact(sock, buf: bytearray, size: int) -> bytes:
+    """Read exactly `size` bytes from `sock`, sharing `buf` with `receive_line`.
+
+    Same buffer contract: bytes past the `size`th belong to whatever message
+    comes next on this connection and are left in `buf` for that call to
+    consume, instead of being read (and discarded) here.
+    """
+
+    while len(buf) < size:
+        chunk = sock.recv(4096)
+        if not chunk:
+            raise NetworkProtocolError("connection closed mid-message")
+        buf.extend(chunk)
+
+    data = bytes(buf[:size])
+    del buf[:size]
+    return data
 
 
 class RemoteClient:
     """Push and pull commits between two minigit repos over a TCP connection."""
 
     def __init__(self, repo_path=".", store=None, commits=None):
-
         self.repo_path = repo_path
         self.config_path = os.path.join(self.repo_path, ".minigit", "config")
-        self.store = store
-        self.commits = commits
-        # Below is correct but need name of function within module 1 & 3
-        # self.store = store if store is not None else ObjectStore(self.repo_path)
-        # self.commits = commits if commits is not None else CommitManager(self.repo_path)
+        self.store = store if store is not None else ObjectStore(self.repo_path)
+        self.commits = (
+            commits if commits is not None else CommitManager(self.repo_path, store=self.store)
+        )
 
     def _parse_address(self, address: str) -> tuple[str, int]:
         """split a string address by host part(string) and the port part(integer)"""
@@ -135,14 +155,105 @@ def pull(self, remote_address: str, branch: str, token: str) -> None:
         finally:
             sock.close()
 
+    def fetch_objects(self, remote_address: str, hashes: list[str], token: str) -> None:
+        """Fetch each of `hashes` from the remote and write it into the local store.
+
+        Authenticates once, then sends one `WANT` per unique hash over the
+        same connection. Each reply's content is re-hashed and checked
+        against the hash that was requested before it is stored, so a
+        corrupted or mismatched reply never lands in the object store.
+        """
+
+        host, port = self._parse_address(remote_address)
+        if len(token) == 0:
+            raise NetworkProtocolError("fetch needs a token: pass --token")
+
+        try:
+            sock = socket.create_connection((host, port), timeout=5)
+        except OSError as exc:
+            raise NetworkProtocolError(f"could not connect to {host}:{port}: {exc}") from exc
+
+        try:
+            buf = bytearray()
+            send_line(sock, f"AUTH {token}")
+            reply = receive_line(sock, buf)
+            if reply != "OK":
+                raise NetworkProtocolError(f"auth failed: {reply}")
+
+            for obj_hash in dict.fromkeys(hashes):
+                send_line(sock, f"WANT {obj_hash}")
+                self._receive_object(sock, buf, obj_hash)
+
+            send_line(sock, "DONE")
+        except OSError as exc:
+            raise NetworkProtocolError(f"connection to {host}:{port} failed: {exc}") from exc
+        finally:
+            sock.close()
+
+    def _receive_object(self, sock, buf: bytearray, expected_hash: str) -> None:
+        """Read one `OBJ`/`ERR` reply for `expected_hash` and store it if it checks out."""
+
+        header = receive_line(sock, buf)
+        command, _, rest = header.partition(" ")
+
+        if command == "ERR":
+            raise NetworkProtocolError(f"remote could not provide {expected_hash}: {rest}")
+        if command != "OBJ":
+            raise NetworkProtocolError(f"expected OBJ, got {header!r}")
+
+        obj_type, _, length_text = rest.partition(" ")
+        if obj_type not in {"blob", "tree", "commit"} or not (
+            length_text.isascii() and length_text.isdigit()
+        ):
+            raise NetworkProtocolError(f"malformed OBJ header: {header!r}")
+
+        content = recv_exact(sock, buf, int(length_text))
+        if self.store.hash_object(content, obj_type) != expected_hash:
+            raise NetworkProtocolError(f"object {expected_hash} failed hash verification")
+
+        self.store.write_object(content, obj_type)
+
+    def collect_reachable(self, branch: str) -> set[str]:
+        """Return every commit, tree, and blob hash reachable from `branch`'s tip.
+
+        Local-only this week: this is what push will later diff against the
+        remote's advertised hash to find what's actually missing there.
+        """
+
+        tip = self.commits.read_ref(branch)
+        if tip is None:
+            return set()
+
+        reachable: set[str] = set()
+        for commit_hash in self.commits.walk_history(tip):
+            reachable.add(commit_hash)
+            commit = self.commits.read_commit(commit_hash)
+            self._collect_tree(commit.tree, reachable)
+
+        return reachable
+
+    def _collect_tree(self, tree_hash: str, reachable: set[str]) -> None:
+        """Add `tree_hash` and everything nested under it to `reachable`, once each."""
+
+        if tree_hash in reachable:
+            return
+        reachable.add(tree_hash)
+
+        for entry in self.store.read_tree(tree_hash):
+            if entry.type == "tree":
+                self._collect_tree(entry.hash, reachable)
+            else:
+                reachable.add(entry.hash)
+
 
 class RemoteServer:
     """Accepts a RemoteClient's AUTH + REF handshake over TCP, one client at a time."""
 
-    def __init__(self, repo_path=".", token="", host="127.0.0.1", port=0):
+    def __init__(self, repo_path=".", token="", host="127.0.0.1", port=0, store=None):
         self.repo_path = repo_path
         self.token = token
         self.host = host
+        self.store = store if store is not None else ObjectStore(repo_path)
 
         self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
@@ -180,7 +291,7 @@ def serve_forever(self) -> None:
                 conn.close()
 
     def _handle_client(self, conn) -> None:
-        """Run one client's AUTH + REF handshake."""
+        """Authenticate the connection, then answer REF / WANT requests until DONE."""
 
         buf = bytearray()
 
@@ -191,12 +302,30 @@ def _handle_client(self, conn) -> None:
             return
         send_line(conn, "OK")
 
-        line = receive_line(conn, buf)
-        command, _, branch = line.partition(" ")
-        if command != "REF":
-            send_line(conn, "ERR expected REF")
-            return
+        while True:
+            line = receive_line(conn, buf)
+            command, _, value = line.partition(" ")
 
+            if command == "DONE":
+                return
+            elif command == "REF":
+                self._send_ref(conn, value)
+            elif command == "WANT":
+                self._send_object(conn, value)
+            else:
+                send_line(conn, "ERR expected REF, WANT, or DONE")
+                return
+
+    def _send_ref(self, conn, branch: str) -> None:
+        """Reply with the commit hash `branch` currently points at, or `-` if it has none."""
+
+        if (
+            not branch
+            or any(part in {"", ".", ".."} for part in branch.split("/"))
+            or any(char in branch for char in "\0\\\r\n")
+        ):
+            send_line(conn, "ERR invalid branch")
+            return
         ref_path = os.path.join(self.repo_path, ".minigit", "refs", "heads", branch)
         if os.path.exists(ref_path):
             with open(ref_path) as f:
@@ -205,6 +334,25 @@ def _handle_client(self, conn) -> None:
             commit_hash = "-"
         send_line(conn, f"REF {branch} {commit_hash}")
 
+    def _send_object(self, conn, obj_hash: str) -> None:
+        """Reply with the requested object's bytes, or `ERR` if it isn't in the store."""
+
+        if len(obj_hash) != 40 or any(c not in "0123456789abcdef" for c in obj_hash):
+            send_line(conn, "ERR invalid object hash")
+            return
+        try:
+            obj_type, content = self.store.read_object(obj_hash)
+        except ObjectNotFoundError:
+            send_line(conn, f"ERR unknown object {obj_hash}")
+            return
+
+        except ObjectCorruptError:
+            send_line(conn, f"ERR corrupt object {obj_hash}")
+            return
+
+        send_line(conn, f"OBJ {obj_type} {len(content)}")
+        conn.sendall(content)
+
 
 # Wire protocol (draft only - Week 2 makes this real):
 # One message per line, UTF-8 encoded, terminated with "\n".