From c057da88ad7650f2625d9d4696744d4ac3228bcc Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Wed, 5 Jul 2023 15:01:29 -0700
Subject: [PATCH] [yellowstone] Add registration to yellowstone stub.

---
 lib/mammoth/include/mammoth/endpoint_client.h |  2 +-
 lib/mammoth/include/mammoth/port_client.h     |  5 ++++-
 sys/denali/denali.cpp                         | 15 +--------------
 sys/yellowstone/include/yellowstone_stub.h    |  5 +++++
 sys/yellowstone/stub/yellowstone_stub.cpp     | 16 ++++++++++++++++
 5 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/lib/mammoth/include/mammoth/endpoint_client.h b/lib/mammoth/include/mammoth/endpoint_client.h
index 86a16c3..83baae9 100644
--- a/lib/mammoth/include/mammoth/endpoint_client.h
+++ b/lib/mammoth/include/mammoth/endpoint_client.h
@@ -20,7 +20,7 @@ class EndpointClient {
   template <typename Req, typename Resp>
   glcr::ErrorOr<Resp> CallEndpoint(const Req& req);
 
-  z_cap_t GetCap() { return cap_; }
+  z_cap_t GetCap() const { return cap_; }
 
  private:
   EndpointClient(uint64_t cap) : cap_(cap) {}
diff --git a/lib/mammoth/include/mammoth/port_client.h b/lib/mammoth/include/mammoth/port_client.h
index 5f33db6..95683b6 100644
--- a/lib/mammoth/include/mammoth/port_client.h
+++ b/lib/mammoth/include/mammoth/port_client.h
@@ -7,6 +7,7 @@
 
 class PortClient {
  public:
+  PortClient() {}
   static PortClient AdoptPort(z_cap_t port_cap);
 
   template <typename T>
@@ -16,8 +17,10 @@ class PortClient {
 
   z_cap_t cap() { return port_cap_; }
 
+  bool empty() { return port_cap_ == 0; }
+
  private:
-  z_cap_t port_cap_;
+  z_cap_t port_cap_ = 0;
 
   PortClient(z_cap_t port_cap);
 };
diff --git a/sys/denali/denali.cpp b/sys/denali/denali.cpp
index 106020d..382ec07 100644
--- a/sys/denali/denali.cpp
+++ b/sys/denali/denali.cpp
@@ -19,24 +19,11 @@ uint64_t main(uint64_t init_port_cap) {
   ASSIGN_OR_RETURN(MappedMemoryRegion ahci_region, stub.GetAhciConfig());
   ASSIGN_OR_RETURN(auto driver, AhciDriver::Init(ahci_region));
 
-  YellowstoneGetReq req{
-      .type = kYellowstoneGetRegistration,
-  };
-  auto resp_cap_or =
-      yellowstone->CallEndpointGetCap<YellowstoneGetReq,
-                                      YellowstoneGetRegistrationResp>(req);
-  if (!resp_cap_or.ok()) {
-    dbgln("Bad call");
-    check(resp_cap_or.error());
-  }
-  auto resp_cap = resp_cap_or.value();
-  PortClient notify = PortClient::AdoptPort(resp_cap.second());
-
   ASSIGN_OR_RETURN(glcr::UniquePtr<EndpointServer> endpoint,
                    EndpointServer::Create());
   ASSIGN_OR_RETURN(glcr::UniquePtr<EndpointClient> client,
                    endpoint->CreateClient());
-  notify.WriteMessage("denali", client->GetCap());
+  check(stub.Register("denali", *client));
 
   DenaliServer server(glcr::Move(endpoint), *driver);
   RET_ERR(server.RunServer());
diff --git a/sys/yellowstone/include/yellowstone_stub.h b/sys/yellowstone/include/yellowstone_stub.h
index b665c22..9199bc0 100644
--- a/sys/yellowstone/include/yellowstone_stub.h
+++ b/sys/yellowstone/include/yellowstone_stub.h
@@ -2,6 +2,7 @@
 
 #include <mammoth/endpoint_client.h>
 #include <mammoth/memory_region.h>
+#include <mammoth/port_client.h>
 #include <ztypes.h>
 
 class YellowstoneStub {
@@ -10,6 +11,10 @@ class YellowstoneStub {
 
   glcr::ErrorOr<MappedMemoryRegion> GetAhciConfig();
 
+  [[nodiscard]] glcr::ErrorCode Register(glcr::String name,
+                                         const EndpointClient& client);
+
  private:
   glcr::UniquePtr<EndpointClient> yellowstone_stub_;
+  PortClient register_port_;
 };
diff --git a/sys/yellowstone/stub/yellowstone_stub.cpp b/sys/yellowstone/stub/yellowstone_stub.cpp
index d05ecb7..993b73b 100644
--- a/sys/yellowstone/stub/yellowstone_stub.cpp
+++ b/sys/yellowstone/stub/yellowstone_stub.cpp
@@ -21,3 +21,19 @@ glcr::ErrorOr<MappedMemoryRegion> YellowstoneStub::GetAhciConfig() {
            ->CallEndpoint<YellowstoneGetReq, YellowstoneGetAhciResp>(req)));
   return MappedMemoryRegion::DirectPhysical(resp.ahci_phys_offset, kPciSize);
 }
+
+glcr::ErrorCode YellowstoneStub::Register(glcr::String name,
+                                          const EndpointClient& client) {
+  if (register_port_.empty()) {
+    YellowstoneGetReq req{
+        .type = kYellowstoneGetRegistration,
+    };
+    ASSIGN_OR_RETURN(
+        auto resp_cap,
+        (yellowstone_stub_->CallEndpointGetCap<
+            YellowstoneGetReq, YellowstoneGetRegistrationResp>(req)));
+    register_port_ = PortClient::AdoptPort(resp_cap.second());
+  }
+
+  return register_port_.WriteString(name, client.GetCap());
+}