Feature/gateway auth (#186)

* Added auth module, just a simple token at this stage * Pass auth token GATEWAY_SECRET through * Auth token not mandatory, can be provided in env var
trustgraph-ai · Dec 2, 2024 · 1b9c6be · 1b9c6be
1 parent 6d200c7
commit 1b9c6be
Show file tree

Hide file tree

Showing 18 changed files with 126 additions and 33 deletions.
diff --git a/templates/components/trustgraph.jsonnet b/templates/components/trustgraph.jsonnet
@@ -15,6 +15,9 @@ local prompt = import "prompt-template.jsonnet";
 
         create:: function(engine)
 
+            local envSecrets = engine.envSecrets("gateway-secret")
+                .with_env_var("GATEWAY_SECRET", "gateway-secret");
+
             local port = $["api-gateway-port"];
 
             local container =
@@ -29,6 +32,7 @@ local prompt = import "prompt-template.jsonnet";
                         "--port",
                         std.toString(port),
                     ])
+                    .with_env_var_secrets(envSecrets)
                     .with_limits("0.5", "256M")
                     .with_reservations("0.1", "256M")
                     .with_port(8000, 8000, "metrics")
@@ -44,6 +48,7 @@ local prompt = import "prompt-template.jsonnet";
                 .with_port(port, port, "api");
 
             engine.resources([
+                envSecrets,
                 containerSet,
                 service,
             ])

diff --git a/trustgraph-flow/trustgraph/api/gateway/agent.py b/trustgraph-flow/trustgraph/api/gateway/agent.py
@@ -6,7 +6,7 @@
 from . endpoint import MultiResponseServiceEndpoint
 
 class AgentEndpoint(MultiResponseServiceEndpoint):
-    def __init__(self, pulsar_host, timeout):
+    def __init__(self, pulsar_host, timeout, auth):
 
         super(AgentEndpoint, self).__init__(
             pulsar_host=pulsar_host,
@@ -16,6 +16,7 @@ def __init__(self, pulsar_host, timeout):
             response_schema=AgentResponse,
             endpoint_path="/api/v1/agent",
             timeout=timeout,
+            auth=auth,
         )
 
     def to_request(self, body):

diff --git a/trustgraph-flow/trustgraph/api/gateway/auth.py b/trustgraph-flow/trustgraph/api/gateway/auth.py
@@ -0,0 +1,22 @@
+
+class Authenticator:
+
+    def __init__(self, token=None, allow_all=False):
+
+        if not allow_all and token is None:
+            raise RuntimeError("Need a token")
+
+        if not allow_all and token == "":
+            raise RuntimeError("Need a token")
+
+        self.token = token
+        self.allow_all = allow_all
+
+    def permitted(self, token, roles):
+
+        if self.allow_all: return True
+
+        if self.token != token: return False
+
+        return True
+
diff --git a/trustgraph-flow/trustgraph/api/gateway/dbpedia.py b/trustgraph-flow/trustgraph/api/gateway/dbpedia.py
@@ -6,7 +6,7 @@
 from . endpoint import ServiceEndpoint
 
 class DbpediaEndpoint(ServiceEndpoint):
-    def __init__(self, pulsar_host, timeout):
+    def __init__(self, pulsar_host, timeout, auth):
 
         super(DbpediaEndpoint, self).__init__(
             pulsar_host=pulsar_host,
@@ -16,6 +16,7 @@ def __init__(self, pulsar_host, timeout):
             response_schema=LookupResponse,
             endpoint_path="/api/v1/dbpedia",
             timeout=timeout,
+            auth=auth,
         )
 
     def to_request(self, body):

diff --git a/trustgraph-flow/trustgraph/api/gateway/embeddings.py b/trustgraph-flow/trustgraph/api/gateway/embeddings.py
@@ -6,7 +6,7 @@
 from . endpoint import ServiceEndpoint
 
 class EmbeddingsEndpoint(ServiceEndpoint):
-    def __init__(self, pulsar_host, timeout):
+    def __init__(self, pulsar_host, timeout, auth):
 
         super(EmbeddingsEndpoint, self).__init__(
             pulsar_host=pulsar_host,
@@ -16,6 +16,7 @@ def __init__(self, pulsar_host, timeout):
             response_schema=EmbeddingsResponse,
             endpoint_path="/api/v1/embeddings",
             timeout=timeout,
+            auth=auth,
         )
 
     def to_request(self, body):

diff --git a/trustgraph-flow/trustgraph/api/gateway/encyclopedia.py b/trustgraph-flow/trustgraph/api/gateway/encyclopedia.py
@@ -6,7 +6,7 @@
 from . endpoint import ServiceEndpoint
 
 class EncyclopediaEndpoint(ServiceEndpoint):
-    def __init__(self, pulsar_host, timeout):
+    def __init__(self, pulsar_host, timeout, auth):
 
         super(EncyclopediaEndpoint, self).__init__(
             pulsar_host=pulsar_host,
@@ -16,6 +16,7 @@ def __init__(self, pulsar_host, timeout):
             response_schema=LookupResponse,
             endpoint_path="/api/v1/encyclopedia",
             timeout=timeout,
+            auth=auth,
         )
 
     def to_request(self, body):

diff --git a/trustgraph-flow/trustgraph/api/gateway/endpoint.py b/trustgraph-flow/trustgraph/api/gateway/endpoint.py
@@ -19,6 +19,7 @@ def __init__(
             request_queue, request_schema,
             response_queue, response_schema,
             endpoint_path,
+            auth,
             subscription="api-gateway", consumer_name="api-gateway",
             timeout=600,
     ):
@@ -36,6 +37,9 @@ def __init__(
 
         self.path = endpoint_path
         self.timeout = timeout
+        self.auth = auth
+
+        self.operation = "service"
 
     async def start(self):
 
@@ -58,14 +62,24 @@ async def handle(self, request):
 
         id = str(uuid.uuid4())
 
+        try:
+            ht = request.headers["Authorization"]
+            tokens = ht.split(" ", 2)
+            if tokens[0] != "Bearer":
+                return web.HTTPUnauthorized()
+            token = tokens[1]
+        except:
+            token = ""
+
+        if not self.auth.permitted(token, self.operation):
+            return web.HTTPUnauthorized()
+
         try:
 
             data = await request.json()
 
             q = await self.sub.subscribe(id)
 
-            print(data)
-
             await self.pub.send(
                 id,
                 self.to_request(data),
@@ -76,8 +90,6 @@ async def handle(self, request):
             except:
                 raise RuntimeError("Timeout waiting for response")
 
-            print(resp)
-
             if resp.error:
                 return web.json_response(
                     { "error": resp.error.message }
@@ -110,8 +122,6 @@ async def handle(self, request):
 
             q = await self.sub.subscribe(id)
 
-            print(data)
-
             await self.pub.send(
                 id,
                 self.to_request(data),
@@ -126,8 +136,6 @@ async def handle(self, request):
                 except:
                     raise RuntimeError("Timeout waiting for response")
 
-                print(resp)
-
                 if resp.error:
                     return web.json_response(
                         { "error": resp.error.message }

diff --git a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py
@@ -14,10 +14,12 @@
 
 class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
 
-    def __init__(self, pulsar_host, path="/api/v1/load/graph-embeddings"):
+    def __init__(
+            self, pulsar_host, auth, path="/api/v1/load/graph-embeddings",
+    ):
 
         super(GraphEmbeddingsLoadEndpoint, self).__init__(
-            endpoint_path=path
+            endpoint_path=path, auth=auth,
         )
 
         self.pulsar_host=pulsar_host

diff --git a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py
@@ -12,10 +12,12 @@
 
 class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
 
-    def __init__(self, pulsar_host, path="/api/v1/stream/graph-embeddings"):
+    def __init__(
+            self, pulsar_host, auth, path="/api/v1/stream/graph-embeddings"
+    ):
 
         super(GraphEmbeddingsStreamEndpoint, self).__init__(
-            endpoint_path=path
+            endpoint_path=path, auth=auth,
         )
 
         self.pulsar_host=pulsar_host

diff --git a/trustgraph-flow/trustgraph/api/gateway/graph_rag.py b/trustgraph-flow/trustgraph/api/gateway/graph_rag.py
@@ -6,7 +6,7 @@
 from . endpoint import ServiceEndpoint
 
 class GraphRagEndpoint(ServiceEndpoint):
-    def __init__(self, pulsar_host, timeout):
+    def __init__(self, pulsar_host, timeout, auth):
 
         super(GraphRagEndpoint, self).__init__(
             pulsar_host=pulsar_host,
@@ -16,6 +16,7 @@ def __init__(self, pulsar_host, timeout):
             response_schema=GraphRagResponse,
             endpoint_path="/api/v1/graph-rag",
             timeout=timeout,
+            auth=auth,
         )
 
     def to_request(self, body):

diff --git a/trustgraph-flow/trustgraph/api/gateway/internet_search.py b/trustgraph-flow/trustgraph/api/gateway/internet_search.py
@@ -6,7 +6,7 @@
 from . endpoint import ServiceEndpoint
 
 class InternetSearchEndpoint(ServiceEndpoint):
-    def __init__(self, pulsar_host, timeout):
+    def __init__(self, pulsar_host, timeout, auth):
 
         super(InternetSearchEndpoint, self).__init__(
             pulsar_host=pulsar_host,
@@ -16,6 +16,7 @@ def __init__(self, pulsar_host, timeout):
             response_schema=LookupResponse,
             endpoint_path="/api/v1/internet-search",
             timeout=timeout,
+            auth=auth,
         )
 
     def to_request(self, body):

diff --git a/trustgraph-flow/trustgraph/api/gateway/prompt.py b/trustgraph-flow/trustgraph/api/gateway/prompt.py
@@ -8,7 +8,7 @@
 from . endpoint import ServiceEndpoint
 
 class PromptEndpoint(ServiceEndpoint):
-    def __init__(self, pulsar_host, timeout):
+    def __init__(self, pulsar_host, timeout, auth):
 
         super(PromptEndpoint, self).__init__(
             pulsar_host=pulsar_host,
@@ -18,6 +18,7 @@ def __init__(self, pulsar_host, timeout):
             response_schema=PromptResponse,
             endpoint_path="/api/v1/prompt",
             timeout=timeout,
+            auth=auth,
         )
 
     def to_request(self, body):

diff --git a/trustgraph-flow/trustgraph/api/gateway/service.py b/trustgraph-flow/trustgraph/api/gateway/service.py
@@ -45,13 +45,15 @@
 from . graph_embeddings_stream import GraphEmbeddingsStreamEndpoint
 from . triples_load import TriplesLoadEndpoint
 from . graph_embeddings_load import GraphEmbeddingsLoadEndpoint
+from . auth import Authenticator
 
 logger = logging.getLogger("api")
 logger.setLevel(logging.INFO)
 
 default_pulsar_host = os.getenv("PULSAR_HOST", "pulsar://pulsar:6650")
 default_timeout = 600
 default_port = 8088
+default_api_token = os.getenv("GATEWAY_SECRET", "")
 
 class Api:
 
@@ -66,45 +68,66 @@ def __init__(self, **config):
         self.timeout = int(config.get("timeout", default_timeout))
         self.pulsar_host = config.get("pulsar_host", default_pulsar_host)
 
+        api_token = config.get("api_token", default_api_token)
+
+        # Token not set, or token equal empty string means no auth
+        if api_token:
+            self.auth = Authenticator(token=api_token)
+        else:
+            self.auth = Authenticator(allow_all=True)
+
         self.endpoints = [
             TextCompletionEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             PromptEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             GraphRagEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             TriplesQueryEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             EmbeddingsEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             AgentEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             EncyclopediaEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             DbpediaEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             InternetSearchEndpoint(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
             ),
             TriplesStreamEndpoint(
-                pulsar_host=self.pulsar_host
+                pulsar_host=self.pulsar_host,
+                auth = self.auth,
             ),
             GraphEmbeddingsStreamEndpoint(
-                pulsar_host=self.pulsar_host
+                pulsar_host=self.pulsar_host,
+                auth = self.auth,
             ),
             TriplesLoadEndpoint(
-                pulsar_host=self.pulsar_host
+                pulsar_host=self.pulsar_host,
+                auth = self.auth,
             ),
             GraphEmbeddingsLoadEndpoint(
-                pulsar_host=self.pulsar_host
+                pulsar_host=self.pulsar_host,
+                auth = self.auth,
             ),
         ]
 
@@ -254,6 +277,12 @@ def run():
         help=f'API request timeout in seconds (default: {default_timeout})',
     )
 
+    parser.add_argument(
+        '--api-token',
+        default=default_api_token,
+        help=f'Secret API token (default: no auth)',
+    )
+
     parser.add_argument(
         '-l', '--log-level',
         type=LogLevel,