Repository: cloudflare/daphne
Pull request #155 (2022-10-28 19:27:20)

daphne_worker: Implement /internal/test/endpoint_for_task
Partially addresses #138.
Based on #154 (merge that first).

cc/ @divergentdave 


Patch:
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index d25c84b..b0e5e5c 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -13,7 +13,7 @@ use crate::{
         DurableConnector, BINDING_DAP_GARBAGE_COLLECTOR, BINDING_DAP_LEADER_BATCH_QUEUE,
         DURABLE_DELETE_ALL,
     },
-    int_err, InternalTestAddTask, InternalTestRole,
+    int_err, InternalTestAddTask, InternalTestEndpointForTask, InternalTestRole,
 };
 use daphne::{
     auth::BearerToken,
@@ -76,6 +76,12 @@ pub(crate) struct DaphneWorkerConfig<D> {
 
     /// Shard count, the number of report storage shards. This should be a power of 2.
     report_shard_count: u64,
+
+    /// Base URL of the Aggregator.
+    base_url: Url,
+
+    /// Indicates if DaphneWorker is used as the Leader.
+    is_leader: bool,
 }
 
 impl<D> DaphneWorkerConfig<D> {
@@ -86,6 +92,12 @@ impl<D> DaphneWorkerConfig<D> {
         )
         .map_err(|e| Error::RustError(format!("Failed to parse DAP_GLOBAL_CONFIG: {}", e)))?;
 
+        let base_url: Url = ctx
+            .var("DAP_BASE_URL")?
+            .to_string()
+            .parse()
+            .map_err(int_err)?;
+
         let report_shard_key = Seed::get_decoded(
             &hex::decode(ctx.secret("DAP_REPORT_SHARD_KEY")?.to_string()).map_err(int_err)?,
         )
@@ -147,6 +159,8 @@ impl<D> DaphneWorkerConfig<D> {
             deployment,
             report_shard_key,
             report_shard_count,
+            base_url,
+            is_leader,
         })
     }
 
@@ -382,7 +396,28 @@ impl<D> DaphneWorkerConfig<D> {
         }
     }
 
-    /// Configure Daphne-Worker a task.
+    /// Get the URL to use for this endpoint, as required by
+    /// draft-dcook-ppm-dap-interop-test-design-02.
+    pub(crate) async fn internal_endpoint_for_task(
+        &self,
+        cmd: InternalTestEndpointForTask,
+    ) -> Result<Response> {
+        if self.is_leader && !matches!(cmd.role, InternalTestRole::Leader)
+            || !self.is_leader && !matches!(cmd.role, InternalTestRole::Helper)
+        {
+            return Response::from_json(&serde_json::json!({
+                "status": "error",
+                "error": "role mismatch",
+            }));
+        }
+
+        Response::from_json(&serde_json::json!({
+            "status": "success",
+            "endpoint": self.base_url,
+        }))
+    }
+
+    /// Configure Daphne-Worker a task, as required by draft-dcook-ppm-dap-interop-test-design-02.
     pub(crate) async fn internal_add_task(&self, cmd: InternalTestAddTask) -> Result<()> {
         // Task ID.
         let task_id_data =
diff --git a/daphne_worker/src/lib.rs b/daphne_worker/src/lib.rs
index 3abbb35..4c7d34d 100644
--- a/daphne_worker/src/lib.rs
+++ b/daphne_worker/src/lib.rs
@@ -347,6 +347,14 @@ impl DaphneWorkerRouter {
                 .post_async("/internal/test/ready", |_req, _ctx| async move {
                     Response::from_json(&())
                 })
+                .post_async(
+                    "/internal/test/endpoint_for_task",
+                    |mut req, ctx| async move {
+                        let config = DaphneWorkerConfig::from_worker_context(ctx)?;
+                        let cmd: InternalTestEndpointForTask = req.json().await?;
+                        config.internal_endpoint_for_task(cmd).await
+                    },
+                )
                 .post_async("/internal/test/add_task", |mut req, ctx| async move {
                     let config = DaphneWorkerConfig::from_worker_context(ctx)?;
                     let cmd: InternalTestAddTask = req.json().await?;
@@ -407,6 +415,11 @@ pub(crate) enum InternalTestRole {
     Leader,
     Helper,
 }
+#[derive(Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub(crate) struct InternalTestEndpointForTask {
+    role: InternalTestRole,
+}
 
 #[derive(Deserialize)]
 pub(crate) struct InternalTestVdaf {


Test patch:
diff --git a/daphne_worker_test/tests/e2e.rs b/daphne_worker_test/tests/e2e.rs
index 4fe89a9..cf773a7 100644
--- a/daphne_worker_test/tests/e2e.rs
+++ b/daphne_worker_test/tests/e2e.rs
@@ -16,20 +16,67 @@ use daphne::{
 use daphne_worker::DaphneWorkerReportSelector;
 use prio::codec::{Decode, Encode};
 use rand::prelude::*;
+use serde::Deserialize;
+use serde_json::json;
 use test_runner::TestRunner;
 
+#[derive(Deserialize)]
+struct InternalTestEndpointForTaskResult {
+    status: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    error: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    endpoint: Option<String>,
+}
+
+#[tokio::test]
+#[cfg_attr(not(feature = "test_e2e"), ignore)]
+async fn e2e_helper_ready() {
+    let t = TestRunner::default().await;
+    t.helper_post_internal::<_, ()>("/internal/test/ready", &())
+        .await;
+}
+
 #[tokio::test]
 #[cfg_attr(not(feature = "test_e2e"), ignore)]
 async fn e2e_leader_ready() {
     let t = TestRunner::default().await;
-    t.leader_post_internal("/internal/test/ready", &()).await;
+    t.leader_post_internal::<_, ()>("/internal/test/ready", &())
+        .await;
+}
+
+#[tokio::test]
+#[cfg_attr(not(feature = "test_e2e"), ignore)]
+async fn e2e_leader_endpoint_for_task() {
+    let t = TestRunner::default().await;
+    let res: InternalTestEndpointForTaskResult = t
+        .leader_post_internal(
+            "/internal/test/endpoint_for_task",
+            &json!({
+                "task_id": "blah blah ignored",
+                "role": "leader",
+            }),
+        )
+        .await;
+    assert_eq!(res.status, "success");
+    assert_eq!(res.endpoint.unwrap(), "http://127.0.0.1:8787/v02/");
 }
 
 #[tokio::test]
 #[cfg_attr(not(feature = "test_e2e"), ignore)]
-async fn e2e_helper_ready() {
+async fn e2e_helper_endpoint_for_task() {
     let t = TestRunner::default().await;
-    t.helper_post_internal("/internal/test/ready", &()).await;
+    let res: InternalTestEndpointForTaskResult = t
+        .helper_post_internal(
+            "/internal/test/endpoint_for_task",
+            &json!({
+                "task_id": "blah blah ignored",
+                "role": "helper",
+            }),
+        )
+        .await;
+    assert_eq!(res.status, "success");
+    assert_eq!(res.endpoint.unwrap(), "http://127.0.0.1:8788/v02/");
 }
 
 #[tokio::test]
diff --git a/daphne_worker_test/tests/test_runner.rs b/daphne_worker_test/tests/test_runner.rs
index 8bfd7df..cb82721 100644
--- a/daphne_worker_test/tests/test_runner.rs
+++ b/daphne_worker_test/tests/test_runner.rs
@@ -16,7 +16,7 @@ use daphne_worker::DaphneWorkerReportSelector;
 use futures::channel::oneshot::Sender;
 use prio::codec::{Decode, Encode};
 use rand::prelude::*;
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 use serde_json::json;
 use std::time::SystemTime;
 #[cfg(feature = "test_janus")]
@@ -159,7 +159,7 @@ impl TestRunner {
             "collector_hpke_config": collector_hpke_config_base64url.clone(),
             "task_expiration": t.task_config.expiration,
         });
-        t.leader_post_internal("/internal/test/add_task", &leader_add_task_cmd)
+        t.leader_post_internal::<_, ()>("/internal/test/add_task", &leader_add_task_cmd)
             .await;
 
         // Configure the Helper with the task.
@@ -178,7 +178,7 @@ impl TestRunner {
             "collector_hpke_config": collector_hpke_config_base64url.clone(),
             "task_expiration": t.task_config.expiration,
         });
-        t.helper_post_internal("/internal/test/add_task", &helper_add_task_cmd)
+        t.helper_post_internal::<_, ()>("/internal/test/add_task", &helper_add_task_cmd)
             .await;
 
         t
@@ -344,7 +344,12 @@ impl TestRunner {
         resp.json().await.unwrap()
     }
 
-    async fn post_internal<I: Serialize>(&self, is_leader: bool, path: &str, data: &I) {
+    async fn post_internal<I: Serialize, O: for<'a> Deserialize<'a>>(
+        &self,
+        is_leader: bool,
+        path: &str,
+        data: &I,
+    ) -> O {
         let client = self.http_client();
         let mut url = if is_leader {
             self.leader_url.clone()
@@ -361,15 +366,24 @@ impl TestRunner {
         if resp.status() != 200 {
             panic!("request to {} failed: response: {:?}", url, resp);
         }
+        resp.json().await.expect("failed to parse result")
     }
 
     #[allow(dead_code)]
-    pub async fn leader_post_internal<I: Serialize>(&self, path: &str, data: &I) {
+    pub async fn leader_post_internal<I: Serialize, O: for<'a> Deserialize<'a>>(
+        &self,
+        path: &str,
+        data: &I,
+    ) -> O {
         self.post_internal(true /* is_leader */, path, data).await
     }
 
     #[allow(dead_code)]
-    pub async fn helper_post_internal<I: Serialize>(&self, path: &str, data: &I) {
+    pub async fn helper_post_internal<I: Serialize, O: for<'a> Deserialize<'a>>(
+        &self,
+        path: &str,
+        data: &I,
+    ) -> O {
         self.post_internal(false /* is_leader */, path, data).await
     }
 
diff --git a/daphne_worker_test/wrangler.toml b/daphne_worker_test/wrangler.toml
index 1324598..50e6271 100644
--- a/daphne_worker_test/wrangler.toml
+++ b/daphne_worker_test/wrangler.toml
@@ -21,6 +21,7 @@ fallthrough = false
 # https://developers.cloudflare.com/workers/wrangler/commands/#secret.
 DAP_DEPLOYMENT = "dev"
 DAP_AGGREGATOR_ROLE = "leader"
+DAP_BASE_URL = "http://127.0.0.1:8787/v02/"
 DAP_ISSUE73_DISABLE_AGG_JOB_QUEUE_GARBAGE_COLLECTION = "true"
 DAP_COLLECT_ID_KEY = "b416a85d280591d6da14e5b75a7d6e31" # SECRET
 DAP_REPORT_SHARD_KEY = "61cd9685547370cfea76c2eb8d156ad9" # SECRET
@@ -64,6 +65,7 @@ preview_id = "<ignored>" # TODO(cjpatton) Figure out how to pick this.
 # https://developers.cloudflare.com/workers/wrangler/commands/#secret.
 DAP_DEPLOYMENT = "dev"
 DAP_AGGREGATOR_ROLE = "helper"
+DAP_BASE_URL = "http://127.0.0.1:8788/v02/"
 DAP_ISSUE73_DISABLE_AGG_JOB_QUEUE_GARBAGE_COLLECTION = "true"
 DAP_REPORT_SHARD_KEY = "f79c352056982bae1737e34bdac24d63" # SECRET
 DAP_REPORT_SHARD_COUNT = "2"