review feedback

divviup · Sep 22, 2023 · 5fca9e0 · 5fca9e0
1 parent 0273c6d
commit 5fca9e0
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 1 deletion.
diff --git a/aggregator_core/src/datastore/tests.rs b/aggregator_core/src/datastore/tests.rs
@@ -203,6 +203,76 @@ async fn roundtrip_task(ephemeral_datastore: EphemeralDatastore) {
     assert_eq!(want_tasks, got_tasks);
 }
 
+#[rstest_reuse::apply(schema_versions_template)]
+#[tokio::test]
+async fn put_task_invalid_aggregator_auth_tokens(ephemeral_datastore: EphemeralDatastore) {
+    install_test_trace_subscriber();
+    let ds = ephemeral_datastore.datastore(MockClock::default()).await;
+
+    let task = TaskBuilder::new(
+        task::QueryType::TimeInterval,
+        VdafInstance::Prio3Count,
+        Role::Leader,
+    )
+    .build();
+
+    ds.put_task(&task).await.unwrap();
+
+    for (auth_token, auth_token_type) in [("NULL", "'BEARER'"), ("\\xDEADBEEF::bytea", "NULL")] {
+        ds.run_tx(|tx| {
+            Box::pin(async move {
+                tx.query_one(
+                    &format!(
+                        "UPDATE tasks SET aggregator_auth_token = {auth_token},
+                            aggregator_auth_token_type = {auth_token_type}"
+                    ),
+                    &[],
+                )
+                .await
+                .unwrap_err();
+                Ok(())
+            })
+        })
+        .await
+        .unwrap();
+    }
+}
+
+#[rstest_reuse::apply(schema_versions_template)]
+#[tokio::test]
+async fn put_task_invalid_collector_auth_tokens(ephemeral_datastore: EphemeralDatastore) {
+    install_test_trace_subscriber();
+    let ds = ephemeral_datastore.datastore(MockClock::default()).await;
+
+    let task = TaskBuilder::new(
+        task::QueryType::TimeInterval,
+        VdafInstance::Prio3Count,
+        Role::Leader,
+    )
+    .build();
+
+    ds.put_task(&task).await.unwrap();
+
+    for (auth_token, auth_token_type) in [("NULL", "'BEARER'"), ("\\xDEADBEEF::bytea", "NULL")] {
+        ds.run_tx(|tx| {
+            Box::pin(async move {
+                tx.query_one(
+                    &format!(
+                        "UPDATE tasks SET collector_auth_token = {auth_token},
+                            collector_auth_token_type = {auth_token_type}"
+                    ),
+                    &[],
+                )
+                .await
+                .unwrap_err();
+                Ok(())
+            })
+        })
+        .await
+        .unwrap();
+    }
+}
+
 #[rstest_reuse::apply(schema_versions_template)]
 #[tokio::test]
 async fn get_task_metrics(ephemeral_datastore: EphemeralDatastore) {

diff --git a/aggregator_core/src/task.rs b/aggregator_core/src/task.rs
@@ -132,7 +132,7 @@ pub struct Task {
     /// Token used to authenticate messages sent to or received from the other aggregator. Only set
     /// if the task was not created via taskprov.
     aggregator_auth_token: Option<AuthenticationToken>,
-    /// Token used to authenticate messages sent to received from the collector. Only set if this
+    /// Token used to authenticate messages sent to or received from the collector. Only set if this
     /// aggregator is the leader.
     collector_auth_token: Option<AuthenticationToken>,
     /// HPKE configurations & private keys used by this aggregator to decrypt client reports.