From be57fda5423eae49133fa0fdecaa30e4d9f57026 Mon Sep 17 00:00:00 2001
From: Joshua Watt <JPEWhacker@gmail.com>
Date: Thu, 22 Feb 2024 15:31:48 -0700
Subject: bitbake: asyncrpc: Add support for server headers

Adds support for asyncrpc servers to send connection headers to clients
on connection. Since this is a breaking protocol change, clients must
opt-in to expect headers from the server, corresponding to a version
bump in the client protocol.

(Bitbake rev: 1cb2b8be6cc5269553f549285592e47b7d29db03)

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
Signed-off-by: Richard Purdie <richard.purdie@linuxfoundation.org>
---
 bitbake/lib/bb/asyncrpc/client.py | 33 ++++++++++++++++++++++++++++++++-
 bitbake/lib/bb/asyncrpc/serv.py   | 16 ++++++++++++++--
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/bitbake/lib/bb/asyncrpc/client.py b/bitbake/lib/bb/asyncrpc/client.py
index a6228bb0ba..29a5ab76aa 100644
--- a/bitbake/lib/bb/asyncrpc/client.py
+++ b/bitbake/lib/bb/asyncrpc/client.py
@@ -17,13 +17,24 @@ from .exceptions import ConnectionClosedError, InvokeError
 
 
 class AsyncClient(object):
-    def __init__(self, proto_name, proto_version, logger, timeout=30):
+    def __init__(
+        self,
+        proto_name,
+        proto_version,
+        logger,
+        timeout=30,
+        server_headers=False,
+        headers={},
+    ):
         self.socket = None
         self.max_chunk = DEFAULT_MAX_CHUNK
         self.proto_name = proto_name
         self.proto_version = proto_version
         self.logger = logger
         self.timeout = timeout
+        self.needs_server_headers = server_headers
+        self.server_headers = {}
+        self.headers = headers
 
     async def connect_tcp(self, address, port):
         async def connect_sock():
@@ -61,9 +72,29 @@ class AsyncClient(object):
     async def setup_connection(self):
         # Send headers
         await self.socket.send("%s %s" % (self.proto_name, self.proto_version))
+        await self.socket.send(
+            "needs-headers: %s" % ("true" if self.needs_server_headers else "false")
+        )
+        for k, v in self.headers.items():
+            await self.socket.send("%s: %s" % (k, v))
+
         # End of headers
         await self.socket.send("")
 
+        self.server_headers = {}
+        if self.needs_server_headers:
+            while True:
+                line = await self.socket.recv()
+                if not line:
+                    # End headers
+                    break
+                tag, value = line.split(":", 1)
+                self.server_headers[tag.lower()] = value.strip()
+
+    async def get_header(self, tag, default):
+        await self.connect()
+        return self.server_headers.get(tag, default)
+
     async def connect(self):
         if self.socket is None:
             self.socket = await self._connect_sock()
diff --git a/bitbake/lib/bb/asyncrpc/serv.py b/bitbake/lib/bb/asyncrpc/serv.py
index f0be9a6cdb..a66117acad 100644
--- a/bitbake/lib/bb/asyncrpc/serv.py
+++ b/bitbake/lib/bb/asyncrpc/serv.py
@@ -39,10 +39,14 @@ class AsyncServerConnection(object):
                 "address": socket.address,
             },
         )
+        self.client_headers = {}
 
     async def close(self):
         await self.socket.close()
 
+    async def handle_headers(self, headers):
+        return {}
+
     async def process_requests(self):
         try:
             self.logger.info("Client %r connected" % (self.socket.address,))
@@ -64,12 +68,20 @@ class AsyncServerConnection(object):
                 )
                 return
 
-            # Read headers. Currently, no headers are implemented, so look for
-            # an empty line to signal the end of the headers
+            # Read headers
+            self.client_headers = {}
             while True:
                 header = await self.socket.recv()
                 if not header:
+                    # Empty line. End of headers
                     break
+                tag, value = header.split(":", 1)
+                self.client_headers[tag.lower()] = value.strip()
+
+            if self.client_headers.get("needs-headers", "false") == "true":
+                for k, v in (await self.handle_headers(self.client_headers)).items():
+                    await self.socket.send("%s: %s" % (k, v))
+                await self.socket.send("")
 
             # Handle messages
             while True:
-- 
cgit v1.2.3-54-g00ecf