From 6bfe9e48e4e16e41e21f4784c6c6c58514873acf Mon Sep 17 00:00:00 2001
From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com>
Date: Sat, 24 Feb 2024 20:05:26 +0530
Subject: [PATCH] added app_message endpoint

---
 packages/grid/veilid/server/main.py        | 11 ++++++
 packages/grid/veilid/server/veilid_core.py | 42 ++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/packages/grid/veilid/server/main.py b/packages/grid/veilid/server/main.py
index 2e847ec8a8b..6ea3132d8e5 100644
--- a/packages/grid/veilid/server/main.py
+++ b/packages/grid/veilid/server/main.py
@@ -1,8 +1,12 @@
 # third party
+from fastapi import Body
 from fastapi import FastAPI
+from fastapi import Request
+from typing_extensions import Annotated
 
 # relative
 from .veilid_core import VeilidConnectionSingleton
+from .veilid_core import app_message
 from .veilid_core import generate_dht_key
 from .veilid_core import get_veilid_conn
 from .veilid_core import retrieve_dht_key
@@ -36,6 +40,13 @@ async def retrieve_dht_key_endpoint() -> dict[str, str]:
     return await retrieve_dht_key()
 
 
+@app.post("/app_message")
+async def app_message_endpoint(
+    request: Request, dht_key: Annotated[str, Body()], message: Annotated[bytes, Body()]
+) -> dict[str, str]:
+    return await app_message(dht_key=dht_key, message=message)
+
+
 @app.on_event("startup")
 async def startup_event() -> None:
     try:
diff --git a/packages/grid/veilid/server/veilid_core.py b/packages/grid/veilid/server/veilid_core.py
index 9cdd6fef5b4..0f09e2a521f 100644
--- a/packages/grid/veilid/server/veilid_core.py
+++ b/packages/grid/veilid/server/veilid_core.py
@@ -1,11 +1,15 @@
 # stdlib
 from typing import Callable
 from typing import Optional
+from typing import Union
 
 # third party
 import veilid
 from veilid import KeyPair
+from veilid import TypedKey
+from veilid import ValueData
 from veilid import VeilidUpdate
+from veilid.json_api import _JsonRoutingContext
 from veilid.json_api import _JsonVeilidAPI
 
 # relative
@@ -87,3 +91,41 @@ async def retrieve_dht_key() -> dict[str, str]:
     if dht_key is None:
         return {"message": "DHT Key does not exist"}
     return {"message": str(dht_key)}
+
+
+async def get_dht_value(
+    router: _JsonRoutingContext,
+    dht_key: TypedKey,
+    subkey: int,
+    force_refresh: bool = True,
+) -> Union[dict[str, str], ValueData]:
+    try:
+        await router.open_dht_record(key=dht_key, writer=None)
+    except Exception:
+        return {"message": f"DHT Key:{dht_key} does not exist in the veilid network"}
+
+    try:
+        return await router.get_dht_value(
+            key=dht_key, subkey=subkey, force_refresh=force_refresh
+        )
+    except Exception:
+        return {"message": f"Subkey:{subkey} does not exist in the DHT Key:{dht_key}"}
+
+
+async def app_message(dht_key: str, message: bytes) -> dict[str, str]:
+    conn = await get_veilid_conn()
+    router = await (await conn.new_routing_context()).with_default_safety()
+
+    dht_key = veilid.TypedKey(dht_key)
+    dht_value = await get_dht_value(router, dht_key, 0)
+
+    if isinstance(dht_value, dict):
+        return dht_value
+
+    # Private Router to peer
+    prr_peer = await conn.import_remote_private_route(dht_value.data)
+
+    # Send message to peer
+    await router.app_message(prr_peer, message)
+
+    return {"message": "Message sent successfully"}