Merge pull request #36 from nicholasaleks/bugfix/subscriptions-and-audit

possible fix for subscriptions and audit logs
dolevf · May 20, 2022 · 50fc7ee · 50fc7ee
2 parents 54d2d68 + 94771a1
commit 50fc7ee
Show file tree

Hide file tree

Showing 5 changed files with 99 additions and 68 deletions.
diff --git a/core/view_override.py b/core/view_override.py
@@ -3,7 +3,15 @@
 from functools import partial
 from flask import Response, request
 from rx import AnonymousObservable
-
+from graphql_ws.gevent import GeventConnectionContext
+from graphql_ws.base_sync import BaseSyncSubscriptionServer
+from graphql_ws.base import (
+    ConnectionClosedException,
+)
+from core.models import (
+  Audit
+)
+import json
 
 def format_execution_result(execution_result, format_error,):
     status_code = 200
@@ -110,3 +118,28 @@ def dispatch_request(self):
                 headers=e.headers,
                 content_type='application/json'
             )
+
+import copy
+
+class GeventSubscriptionServerCustom(BaseSyncSubscriptionServer):
+    def handle(self, ws, request_context=None):
+        connection_context = GeventConnectionContext(ws, request_context)
+        self.on_open(connection_context)
+        while True:
+            try:
+                if connection_context.closed:
+                    raise ConnectionClosedException()
+                message = connection_context.receive()
+            except ConnectionClosedException:
+                self.on_close(connection_context)
+                return
+
+            if message:
+
+                msg = json.loads(message)
+
+                if msg.get('type', '') == 'start':
+                  Audit.create_audit_entry(msg['payload']['query'], operation_type='subscription')
+
+
+            self.on_message(connection_context, message)
diff --git a/core/views.py b/core/views.py
@@ -1,4 +1,3 @@
-import json
 import graphene
 
 from core.directives import *
@@ -16,7 +15,6 @@
 from flask_sockets import Sockets
 
 from graphql.backend import GraphQLCoreBackend
-from graphql_ws.gevent import  GeventSubscriptionServer
 
 from sqlalchemy.sql import text
 from graphene_sqlalchemy import (
@@ -31,7 +29,7 @@
   middleware
 )
 from sqlalchemy import event
-from core.view_override import OverriddenView
+from core.view_override import OverriddenView, GeventSubscriptionServerCustom
 
 from core.models import (
   Owner,
@@ -426,18 +424,16 @@ def set_difficulty():
 
 schema = graphene.Schema(query=Query, mutation=Mutations, subscription=Subscription, directives=[ShowNetworkDirective, SkipDirective, DeprecatedDirective])
 
-subscription_server = GeventSubscriptionServer(schema)
+subscription_server = GeventSubscriptionServerCustom(schema)
+
 
 sockets = Sockets(app)
 
 @sockets.route('/subscriptions')
 def echo_socket(ws):
-  msg = json.loads(ws.read_message())
-
-  if msg.get('type', '') == 'start':
-    Audit.create_audit_entry(msg['payload']['query'], operation_type='subscription')
-
+
   subscription_server.handle(ws)
+
   return []
 
 

diff --git a/templates/partials/base.html b/templates/partials/base.html
@@ -21,63 +21,6 @@
   <script src="/static/bootstrap/js/bootstrap.bundle.min.js"></script>
   <script src="/static/jquery/graphql.js"></script>
 
-  <script type="text/javascript">
-
-    const GQL = {
-      CONNECTION_INIT: 'connection_init',
-      CONNECTION_ACK: 'connection_ack',
-      CONNECTION_ERROR: 'connection_error',
-      CONNECTION_KEEP_ALIVE: 'ka',
-      START: 'start',
-      STOP: 'stop',
-      CONNECTION_TERMINATE: 'connection_terminate',
-      DATA: 'data',
-      ERROR: 'error',
-      COMPLETE: 'complete'
-    }
-
-    function subscribeToPastes() {
-
-        var exampleSocket = new WebSocket("ws://{{host}}:{{port}}/subscriptions");
-
-        exampleSocket.onopen = function () {
-            var query = 'subscription getPaste {paste { id title content ipAddr userAgent public owner {name} } }';
-
-            var graphqlMsg = {
-                type: GQL.START,
-                payload: {query}
-            };
-            exampleSocket.send(JSON.stringify(graphqlMsg));
-        };
-
-        exampleSocket.onmessage = function (event) {
-          data = JSON.parse(event.data)
-          paste = data.payload.data.paste
-
-          var pasteHTML = `<div class="card-header">
-              <i class="fa fa-paste"></i> &nbsp; ${paste.title}
-            </div>
-            <div class="card-body">
-              <p class="card-text">
-              <pre>${paste.content}</pre>
-              <br><hr />
-              <i class="fa fa-user"></i>
-              <i><small><b>${paste.owner.name}</b><br>- ${paste.ipAddr}<br>- (${paste.userAgent})</small></i></p>
-            </div>`
-
-          if(paste.public){
-            $(pasteHTML).hide().prependTo("#public_gallery").fadeIn(1000);
-          } else {
-            $(pasteHTML).hide().prependTo("#private_gallery").fadeIn(1000);
-          }
-
-
-        }
-    }
-
-    window.addEventListener("load", subscribeToPastes);
-  </script>
-
   {% block scripts %}
   {% endblock %}
 </body>

diff --git a/templates/paste.html b/templates/paste.html
@@ -26,6 +26,65 @@
 {% endblock %}
 
 {% block scripts %}
+  <script type="text/javascript">
+
+    const GQL = {
+      CONNECTION_INIT: 'connection_init',
+      CONNECTION_ACK: 'connection_ack',
+      CONNECTION_ERROR: 'connection_error',
+      CONNECTION_KEEP_ALIVE: 'ka',
+      START: 'start',
+      STOP: 'stop',
+      CONNECTION_TERMINATE: 'connection_terminate',
+      DATA: 'data',
+      ERROR: 'error',
+      COMPLETE: 'complete'
+    }
+
+    function subscribeToPastes() {
+
+        var exampleSocket = new WebSocket("ws://{{host}}:{{port}}/subscriptions");
+
+        exampleSocket.onopen = function () {
+            var query = 'subscription getPaste {paste { id title content ipAddr userAgent public owner {name} } }';
+
+            var graphqlMsg = {
+                type: GQL.START,
+                payload: {query}
+            };
+            exampleSocket.send(JSON.stringify(graphqlMsg));
+        };
+
+        var PASTE_UPDATES = {};
+
+        exampleSocket.onmessage = function (event) {
+          data = JSON.parse(event.data)
+          paste = data.payload.data.paste
+
+
+          var pasteHTML = `<div class="card-header">
+              <i class="fa fa-paste"></i> &nbsp; ${paste.title}
+            </div>
+            <div class="card-body">
+              <p class="card-text">
+              <pre>${paste.content}</pre>
+              <br><hr />
+              <i class="fa fa-user"></i>
+              <i><small><b>${paste.owner.name}</b><br>- ${paste.ipAddr}<br>- (${paste.userAgent})</small></i></p>
+            </div>`
+
+          if(paste.public){
+            $(pasteHTML).hide().prependTo("#public_gallery").fadeIn(1000);
+          } else {
+            $(pasteHTML).hide().prependTo("#private_gallery").fadeIn(1000);
+          }
+
+
+        }
+    }
+
+    window.addEventListener("load", subscribeToPastes);
+  </script>
   {% include 'partials/base_scripts.html' %}
   <script>
     function burnSelect() {

diff --git a/version.py b/version.py
@@ -1 +1 @@
-VERSION = '2.0.1'
+VERSION = '2.0.2'