From d6f49a003b10b27fae2493c87a13349b9e81afe4 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 29 Mar 2018 13:52:19 +0545
Subject: [PATCH] issue #106: ansible: beginnings of FileService.

---
 ansible_mitogen/services.py | 70 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 69 insertions(+), 1 deletion(-)

diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py
index 0a576351..98a231d4 100644
--- a/ansible_mitogen/services.py
+++ b/ansible_mitogen/services.py
@@ -27,9 +27,16 @@
 # POSSIBILITY OF SUCH DAMAGE.
 
 from __future__ import absolute_import
+import logging
+import zlib
+
+import mitogen
 import mitogen.service
 
 
+LOG = logging.getLogger(__name__)
+
+
 class ContextService(mitogen.service.DeduplicatingService):
     """
     Used by worker processes connecting back into the top-level process to
@@ -40,7 +47,7 @@ class ContextService(mitogen.service.DeduplicatingService):
         https://mitogen.readthedocs.io/en/latest/api.html#context-factories
 
     This concentrates all SSH connections in the top-level process, which may
-    become a bottleneck. There are multiple ways to fix that: 
+    become a bottleneck. There are multiple ways to fix that:
         * creating one .local() child context per CPU and sharding connections
           between them, using the master process to route messages, or
         * as above, but having each child create a unique UNIX listener and
@@ -68,3 +75,64 @@ class ContextService(mitogen.service.DeduplicatingService):
         args.pop('discriminator', None)
         method = getattr(self.router, args.pop('method'))
         return method(**args)
+
+
+class FileService(mitogen.service.Service):
+    """
+    Primitive latency-inducing file server for old-style incantations of the
+    module runner. This is to be replaced later with a scheme that forwards
+    files known to be missing without the target having to ask for them,
+    avoiding a corresponding roundtrip per file.
+
+    Paths must be explicitly added to the service by a trusted context before
+    they will be served to an untrusted context.
+
+    :param tuple args:
+        Tuple of `(cmd, path)`, where:
+            - cmd: one of "register", "fetch", where:
+                - register: register a file that may be fetched
+                - fetch: fetch a file that was previously registered
+            - path: key of the file to fetch or register
+
+    :returns:
+        Returns ``None` for "register", or the file data for "fetch".
+
+    :raises mitogen.core.CallError:
+        Security violation occurred, either path not registered, or attempt to
+        register path from unprivileged context.
+    """
+    handle = 501
+    max_message_size = 1000
+
+    unprivileged_msg = 'Cannot register from unprivileged context.'
+    unregistered_msg = 'Path is not registered with FileService.'
+
+    def __init__(self, router):
+        super(FileService, self).__init__(router)
+        self._paths = {}
+
+    def validate_args(self, args):
+        return (
+            isinstance(args, tuple) and
+            len(args) == 2 and
+            args[0] in ('register', 'fetch') and
+            isinstance(args[1], str)
+        )
+
+    def dispatch(self, args, msg):
+        cmd, path = msg
+        return getattr(self, cmd)(path, msg)
+
+    def register(self, path, msg):
+        if msg.auth_id not in mitogen.parent_ids:
+            raise mitogen.core.CallError(self.unprivileged_msg)
+
+        with open(path, 'rb') as fp:
+            self._paths[path] = zlib.compress(fp.read())
+
+    def fetch(self, path, msg):
+        if path not in self._paths:
+            raise mitogen.core.CallError(self.unregistered_msg)
+
+        LOG.debug('Serving %r to context %r', path, msg.src_id)
+        return self._paths[path]