large-ish commit. added map in global omq instance to track client connections, implementation of OMQ request handling underway, re-writing subrequest and onion-request handling to process omq requests

Author: dr7ana

sogs/model/clientmanager.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import oxenmq
+
 from .. import crypto, db, config
 from ..db import query
 from ..web import app
@@ -92,7 +94,7 @@ def bqempty(self):
         return not self.bqueue._empty()
 
 
-    def register_client(self, cid, authlevel, bot: bool = False, priority: int = None):
+    def register_client(self, conn_id, cid, authlevel, bot, priority):
         if not bot:
             # add client to self.clients
             return
@@ -107,7 +109,7 @@ def register_client(self, cid, authlevel, bot: bool = False, priority: int = Non
         self.bqueue._put(PriorityTuple(priority, bot))
 
 
-    def deregister_client(self, cid, bot: bool = False):
+    def deregister_client(self, cid, bot):
         if not bot:
             # remove client from clients list
             return

sogs/mule.py

@@ -86,17 +86,21 @@ def setup_omq(self, omq: OMQ):
         self._omq.add_timer(cleanup.cleanup, timedelta(seconds=cleanup.INTERVAL))
 
         # Commands other workers can send to us, e.g. for notifications of activity for us to know about
-        worker = self._omq.add_category("worker", access_level=oxenmq.AuthLevel.admin)
+        worker = omq._omq.add_category("worker", access_level=oxenmq.AuthLevel.admin)
         worker.add_command("message_posted", self.message_posted)
         worker.add_command("messages_deleted", self.messages_deleted)
         worker.add_command("message_edited", self.message_edited)
 
-        # new client code
-        # TOFIX: use add_request_command to handle a response value
-        handler = self._omq.add_category("handler", access_level=oxenmq.AuthLevel.admin)
-        handler.add_command("register_client", omq.register_client)
-        handler.add_command("deregister_client", omq.deregister_client)
-        handler.add_command("send_to_handler", omq.manager.receive_message)
+        # client code
+        handler = omq._omq.add_category("handler", access_level=oxenmq.AuthLevel.admin)
+        handler.add_request_command("register_client", omq.register_client)
+        handler.add_request_command("deregister_client", omq.deregister_client)
+        handler.add_request_command("send_to_handler", omq.manager.receive_message)
+
+        # proxy handler for subrequest queue
+        internal = omq._omq.add_category("internal", access_level=oxenmq.AuthLevel.admin)
+        internal.add_request_command("get_next_request", omq.get_next_request)
+        internal.add_request_command("subreq_response", omq.subreq_response)
 
         app.logger.debug("Mule starting omq")
         self._omq.start()

sogs/omq.py

@@ -1,16 +1,19 @@
 # Common oxenmq object; this is used by workers and the oxenmq mule.  We create, but do not start,
 # this pre-forking.
 
-import oxenmq
-from oxenc import bt_serialize
+import oxenmq, queue
+from oxenc import bt_serialize, bt_deserialize
 
+from mule import log_exceptions
 from routes import omq_auth
 from . import crypto, config
 from .postfork import postfork
 from .model.clientmanager import ClientManager
 
 
 omq_global = None
+global blueprints_global
+blueprints_global = {}
 
 
 class OMQ:
@@ -25,13 +28,16 @@ def __init__(self):
             privkey=crypto._privkey.encode(), pubkey=crypto.server_pubkey.encode()
         )
         self._omq.ephemeral_routing_id = True
-
+        self.client_map = {}
         self.manager = ClientManager()
         self.test_suite = False
+        self.subreq_queue = queue.SimpleQueue()
 
         if uwsgi.mule_id() != 0:
             uwsgi.opt['mule'].setup_omq(self)
             return
+        
+        uwsgi.register_signal(123, 'internal', self.handle_proxied_omq_req)
 
         from .web import app  # Imported here to avoid circular import
 
@@ -46,15 +52,47 @@ def __init__(self):
         global omq_global
         omq_global = self
 
+    
+    @log_exceptions
+    def subreq_response(self):
+        pass
+
+
+    @log_exceptions
+    def handle_proxied_omq_req(self):
+        id, subreq_body = self.send_mule(
+            command='get_next_request',
+            prefix='internal'
+        )
+
+        '''
+            
+            Handle omq subrequest
+
+        '''
+
+        return 
+
+    @log_exceptions
+    def get_next_request(self):
+        subreq_body = self.subreq_queue.get()
+        id = list(subreq_body.keys())[0]
+        return id, subreq_body[id]
+
 
-    def register_client(self, cid, authlevel, bot: bool = False, priority: int = None):
-        self.manager.register_client(cid, authlevel, bot, priority)
-        # TODO: add omq logic
+    @log_exceptions
+    def register_client(self, msg: oxenmq.Message):
+        cid, authlevel, bot, priority = bt_deserialize(msg.data()[0])
+        conn_id = msg.conn()
+        self.client_map[conn_id] = cid
+        self.manager.register_client(msg)
 
 
-    def deregister_client(self, cid, bot: bool = False):
-        self.manager.register_client()
-        # TODO: add omq logic
+    @log_exceptions
+    def deregister_client(self, msg: oxenmq.Message):
+        cid, bot = bt_deserialize(msg.data()[0])
+        self.client_map.pop(cid)
+        self.manager.deregister_client(cid, bot)
 
 
     def send_mule(self, command, *args, prefix="worker."):

sogs/routes/clients.py

@@ -2,7 +2,7 @@
 from ..model import room as mroom
 from ..model.user import User
 from ..web import app
-from ..omq import omq_global
+from ..omq import omq_global, blueprints_global
 from . import omq_auth
 
 from flask import abort, jsonify, g, Blueprint, request
@@ -15,8 +15,8 @@
 
 """
 
-
 clients = Blueprint('clients', __name__)
+blueprints_global['clients'] = clients
 
 
 @omq_auth.first_request
@@ -28,7 +28,7 @@ def register(cid):
 
     ## URL Parameters
 
-        - 'cid': the client ID of the given client to be registered with the SOGS instance
+        - 'cid': the client ID (session ID) of the given client to be registered with the SOGS instance
 
     ## Query Parameters
 
@@ -61,7 +61,7 @@ def register_bot(cid, authlevel, priority):
 
     ## URL Parameters
 
-        - 'cid': the client ID of the given client to be registered with the SOGS instance
+        - 'cid': the client ID (session ID) of the given client to be registered with the SOGS instance
 
     ## Body Parameters
 
@@ -73,11 +73,12 @@ def register_bot(cid, authlevel, priority):
     """
 
     client = omq_global.send_mule(
-        command="register_client",
+        command='register_client',
         cid=cid,
         authlevel=authlevel,
-        bot=True,
-        priority=priority
+        bot=1,
+        priority=priority,
+        prefix='handler'
     )
 
     return client
@@ -90,7 +91,7 @@ def register_client(cid, authlevel):
 
     ## URL Parameters
 
-        - 'cid': the client ID of the given client to be registered with the SOGS instance
+        - 'cid': the client ID (session ID) of the given client to be registered with the SOGS instance
 
     ## Body Parameters
 
@@ -100,11 +101,12 @@ def register_client(cid, authlevel):
     """
 
     client = omq_global.send_mule(
-        command="register_client",
+        command='register_client',
         cid=cid,
         authlevel=authlevel,
-        bot=False,
-        priority=None
+        bot=0,
+        priority=None,
+        prefix='handler'
     )
 
     return client
@@ -117,7 +119,7 @@ def unregister(cid):
 
     ## URL Parameters
 
-        - 'cid': the client ID of the given client to be registered with the SOGS instance
+        - 'cid': the client ID (session ID) of the given client to be registered with the SOGS instance
 
     ## Query Parameters
 
@@ -136,9 +138,10 @@ def unregister(cid):
 def unregister_client(cid):
 
     client = omq_global.send_mule(
-        command="register_client",
+        command='unregister_client',
         cid=cid,
-        bot=False
+        bot=0,
+        prefix='handler'
     )
 
     return client
@@ -149,9 +152,10 @@ def unregister_client(cid):
 def unregister_bot(cid):
 
     client = omq_global.send_mule(
-        command="register_client",
+        command='unregister_bot',
         cid=cid,
-        bot=True
+        bot=1,
+        prefix='handler'
     )
 
     return client

sogs/routes/messages.py

@@ -1,15 +1,14 @@
 from .. import http, utils
 from . import auth
 from model.room import Room
-from ..omq import omq_global
+from ..omq import omq_global, blueprints_global
 
 from flask import abort, jsonify, g, Blueprint, request
 
 # Room message retrieving/submitting endpoints
 
-
 messages = Blueprint('messages', __name__)
-
+blueprints_global['messages'] = messages
 
 def qs_reactors():
     return utils.get_int_param('reactors', 4, min=0, max=20, truncate=True)

sogs/routes/omq_auth.py

@@ -3,10 +3,11 @@
 from ..web import app
 from ..db import query
 from .. import config, crypto, http, utils
+from ..omq import omq_global, blueprints_global
 from ..model.user import User
 from ..hashing import blake2b
 
-from flask import request, abort, Response, g
+from flask import request, abort, Response, g, Blueprint
 import time
 import nacl
 from nacl.signing import VerifyKey
@@ -17,6 +18,30 @@
 
 # Authentication for handling OMQ requests
 
+omq = Blueprint('endpoint', __name__)
+
+def endpoint(f):
+    """ 
+    Default endpoint for omq routes to pass requests to; constructs flask HTTP request
+
+    Message (request) components:
+
+    "blueprint" - the flask blueprint
+    "query" - the request query 
+    "pubkey" - pk of client making request
+    "params" - a json value to dump as the the query parameters
+
+        Example:
+            full request: `@omq.endpoint("messages", "room.<Room>.messages.since.<seqno>", {'Room:room', 'int:seqno'})`
+                blueprint: 'messages'
+                query: 'room.<Room>.messages.since.<seqno>'
+                params: {'Room:room', 'int:seqno'}
+    """
+
+    @wraps(f)
+    def endpoint_wrapper(*args, blueprint, query, pubkey, params, **kwargs):
+        bp = blueprints_global['messages']
+
 
 def abort_request(code, msg, warn=True):
     if warn:
@@ -27,16 +52,20 @@ def abort_request(code, msg, warn=True):
 
 
 def require_client():
-    """ Requires that an authenticated client was found in the OMQ instance; aborts with 
-    UNAUTHORIZED if the request has no client """
+    """ 
+    Requires that an authenticated client was found in the OMQ instance; aborts with 
+    UNAUTHORIZED if the request has no client 
+    """
     if g.client_id is None:
         abort_request(http.UNAUTHORIZED, 'OMQ client authentication required')
 
 
 def client_required(f):
-    """ Decorator for an endpoint that requires a client; this calls require_client() at the 
+    """ 
+    Decorator for an endpoint that requires a client; this calls require_client() at the 
     beginning of the request to abort the request as UNAUTHORIZED if the client has not been 
-    previously authenticated"""
+    previously authenticated 
+    """
 
     @wraps(f)
     def required_client_wrapper(*args, **kwargs):
@@ -126,14 +155,10 @@ def verify_omq_auth():
         return
 
 
-
 """
     TOFIX:
         - add some type of dict in omq_global to map conn_ID (onenmq conn ID) to session_ID/other info
             - do not persist:
                 - room specific access: check every time it makes a request because it can change
                 - values that admin level can change
-
-    
-
 """