From 46a206535715a1a759ce336bd713790d61087501 Mon Sep 17 00:00:00 2001
From: Jean-Paul Calderone <exarkun@twistedmatrix.com>
Date: Fri, 31 May 2019 13:40:51 -0400
Subject: [PATCH] Offer a pass-through IStorageServer

---
 src/allmydata/storage_client.py | 134 +++++++++++++++++++++++++++++++-
 1 file changed, 133 insertions(+), 1 deletion(-)

diff --git a/src/allmydata/storage_client.py b/src/allmydata/storage_client.py
index 832f03858..cd0f2e4ad 100644
--- a/src/allmydata/storage_client.py
+++ b/src/allmydata/storage_client.py
@@ -30,12 +30,18 @@ the foolscap-based server implemented in src/allmydata/storage/*.py .
 
 
 import re, time, hashlib
+import attr
 from zope.interface import implementer
 from twisted.internet import defer
 from twisted.application import service
 
 from foolscap.api import eventually
-from allmydata.interfaces import IStorageBroker, IDisplayableServer, IServer
+from allmydata.interfaces import (
+    IStorageBroker,
+    IDisplayableServer,
+    IServer,
+    IStorageServer,
+)
 from allmydata.util import log, base32, connection_status
 from allmydata.util.assertutil import precondition
 from allmydata.util.observer import ObserverList
@@ -421,6 +427,18 @@ class NativeStorageServer(service.MultiService):
     def get_rref(self):
         return self._rref
 
+    def get_storage_server(self):
+        if self._rref is None:
+            # Somewhat questionable to create a third state (never connected,
+            # connected, disconnected) but it mirrors the get_rref behavior so
+            # we'll keep it for now.
+            return None
+        # Pass in an accessor for our _rref attribute.  The value of the
+        # attribute may change over time as connections are lost and
+        # re-established.  The _StorageServer should always be able to get the
+        # most up-to-date value.
+        return _StorageServer(get_rref=lambda: self._rref)
+
     def _lost(self):
         log.msg(format="lost connection to %(name)s", name=self.get_name(),
                 facility="tahoe.storage_broker", umid="zbRllw")
@@ -443,3 +461,117 @@ class NativeStorageServer(service.MultiService):
 
 class UnknownServerTypeError(Exception):
     pass
+
+
+@implementer(IStorageServer)
+@attr.s
+class _StorageServer(object):
+    """
+    ``_StorageServer`` is a direct pass-through to an ``RIStorageServer`` via
+    a ``RemoteReference``.
+    """
+    _get_rref = attr.ib()
+
+    @property
+    def _rref(self):
+        return self._get_rref()
+
+    def get_version(self):
+        return self._rref.callRemote(
+            "get_version",
+        )
+
+    def allocate_buckets(
+            self,
+            storage_index,
+            renew_secret,
+            cancel_secret,
+            sharenums,
+            allocated_size,
+            canary,
+    ):
+        return self._rref.callRemote(
+            "allocate_buckets",
+            storage_index,
+            renew_secret,
+            cancel_secret,
+            sharenums,
+            allocated_size,
+            canary,
+        )
+
+    def add_lease(
+            self,
+            storage_index,
+            renew_secret,
+            cancel_secret,
+    ):
+        return self._rref.callRemote(
+            "add_lease",
+            storage_index,
+            renew_secret,
+            cancel_secret,
+        )
+
+    def renew_lease(
+            self,
+            storage_index,
+            renew_secret,
+    ):
+        return self._rref.callRemote(
+            "renew_lease",
+            storage_index,
+            renew_secret,
+        )
+
+    def get_buckets(
+            self,
+            storage_index,
+    ):
+        return self._rref.callRemote(
+            "get_buckets",
+            storage_index,
+        )
+
+    def slot_readv(
+            self,
+            storage_index,
+            shares,
+            readv,
+    ):
+        return self._rref.callRemote(
+            "slot_readv",
+            storage_index,
+            shares,
+            readv,
+        )
+
+    def slot_testv_and_readv_and_writev(
+            self,
+            storage_index,
+            secrets,
+            tw_vectors,
+            r_vector,
+    ):
+        return self._rref.callRemote(
+            "slot_testv_and_readv_and_writev",
+            storage_index,
+            secrets,
+            tw_vectors,
+            r_vector,
+        )
+
+    def advise_corrupt_share(
+            self,
+            share_type,
+            storage_index,
+            shnum,
+            reason,
+    ):
+        return self._rref.callRemoteOnly(
+            "advise_corrupt_share",
+            share_type,
+            storage_index,
+            shnum,
+            reason,
+        )