Introduce Rust driver token-based authentication (#749)

## Usage and product changes
We update the protocol version and introduce token-based authentication
for all the available drivers. Now, instead of sending usernames and
passwords for authentication purposes, authorization tokens are
implicitly added to every network request to a TypeDB server. It
enhances the authentication speed and security.

These tokens are acquired as a result of driver instantiation and are
renewed automatically. This feature does not require any user-side
changes.

Additionally, as a part of the HTTP client introduction work, we rename
Concept Documents key style from snake_case to camelCase, which is a
common convention for JSONs (it affects only value types: `value_type`
-> `valueType`).

## Implementation
Every network call adds tokens instead of user credentials and is
repeated after a token renewal request if it receives an authentication
error.

Author: farost

Cargo.lock

@@ -25,7 +25,7 @@ def typedb_artifact():
         artifact_name = "typedb-all-{platform}-{version}.{ext}",
         tag_source = deployment["artifact"]["release"]["download"],
         commit_source = deployment["artifact"]["snapshot"]["download"],
-        commit = "588d742a36b24423b40eaf05d8013c0cd188425f"
+        commit = "f98dba89712c7bc2a3418e8d65326af75c3bf150"
     )
 
 #def typedb_cloud_artifact():

dependencies/typedb/artifacts.bzl

@@ -28,12 +28,12 @@ def typedb_protocol():
     git_repository(
         name = "typedb_protocol",
         remote = "https://github.com/typedb/typedb-protocol",
-        tag = "3.1.0",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_protocol
+        commit = "9e46e089a005d6ca9f017ffb482337b9d5718695",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_protocol
     )
 
 def typedb_behaviour():
     git_repository(
         name = "typedb_behaviour",
         remote = "https://github.com/typedb/typedb-behaviour",
-        commit = "7c2df22e4378b32e53440df4a7b20d59f2ea0e28",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_behaviour
+        commit = "57fc7216f4b0df05ce29044a1ba621e44bdd87e1",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_behaviour
     )

dependencies/typedb/repositories.bzl

@@ -31,13 +31,13 @@
 		version = "1.13.0"
 		default-features = false
 
-	[dev-dependencies.steps]
-		path = "../rust/tests/behaviour/steps"
+	[dev-dependencies.config]
+		path = "../rust/tests/behaviour/config"
 		features = []
 		default-features = false
 
-	[dev-dependencies.config]
-		path = "../rust/tests/behaviour/config"
+	[dev-dependencies.steps]
+		path = "../rust/tests/behaviour/steps"
 		features = []
 		default-features = false
 
@@ -60,7 +60,7 @@
 
 	[dependencies.typedb-protocol]
 		features = []
-		rev = "7c7e24c43a59ac3e3b2c4bb17566eb9478099c21"
+		rev = "9e46e089a005d6ca9f017ffb482337b9d5718695"
 		git = "https://github.com/typedb/typedb-protocol"
 		default-features = false
 
@@ -79,16 +79,16 @@
 		version = "0.3.31"
 		default-features = false
 
-	[dependencies.uuid]
-		features = ["default", "fast-rng", "rng", "serde", "std", "v4"]
-		version = "1.15.1"
-		default-features = false
-
 	[dependencies.itertools]
 		features = ["default", "use_alloc", "use_std"]
 		version = "0.10.5"
 		default-features = false
 
+	[dependencies.uuid]
+		features = ["default", "fast-rng", "rng", "serde", "std", "v4"]
+		version = "1.15.1"
+		default-features = false
+
 	[dependencies.prost]
 		features = ["default", "derive", "prost-derive", "std"]
 		version = "0.13.5"

rust/Cargo.toml

@@ -116,7 +116,7 @@ impl Leaf {
 
 const KIND: Cow<'static, str> = Cow::Borrowed("kind");
 const LABEL: Cow<'static, str> = Cow::Borrowed("label");
-const VALUE_TYPE: Cow<'static, str> = Cow::Borrowed("value_type");
+const VALUE_TYPE: Cow<'static, str> = Cow::Borrowed("valueType");
 
 fn json_type(kind: Kind, label: Cow<'static, str>) -> JSON {
     JSON::Object([(KIND, json_kind(kind)), (LABEL, JSON::String(label))].into())

rust/src/answer/concept_document.rs

@@ -21,7 +21,7 @@ use std::{collections::HashSet, error::Error as StdError, fmt};
 
 use itertools::Itertools;
 use tonic::{Code, Status};
-use tonic_types::StatusExt;
+use tonic_types::{ErrorDetails, ErrorInfo, StatusExt};
 
 use super::{address::Address, RequestID};
 
@@ -150,7 +150,7 @@ error_messages! { ConnectionError
         15: "The replica is not the primary replica.",
     ClusterAllNodesFailed { errors: String } =
         16: "Attempted connecting to all TypeDB Cluster servers, but the following errors occurred: \n{errors}.",
-    ClusterTokenCredentialInvalid =
+    TokenCredentialInvalid =
         17: "Invalid token credentials.",
     EncryptionSettingsMismatch =
         18: "Unable to connect to TypeDB: possible encryption settings mismatch.",
@@ -275,6 +275,16 @@ impl Error {
         }
     }
 
+    fn try_extracting_connection_error(status: &Status, code: &str) -> Option<ConnectionError> {
+        // TODO: We should probably catch more connection errors instead of wrapping them into
+        // ServerErrors. However, the most valuable information even for connection is inside
+        // stacktraces now.
+        match code {
+            "AUT2" | "AUT3" => Some(ConnectionError::TokenCredentialInvalid {}),
+            _ => None,
+        }
+    }
+
     fn from_message(message: &str) -> Self {
         // TODO: Consider converting some of the messages to connection errors
         Self::Other(message.to_owned())
@@ -352,9 +362,13 @@ impl From<Status> for Error {
                 })
             } else if let Some(error_info) = details.error_info() {
                 let code = error_info.reason.clone();
+                if let Some(connection_error) = Self::try_extracting_connection_error(&status, &code) {
+                    return Self::Connection(connection_error);
+                }
                 let domain = error_info.domain.clone();
                 let stack_trace =
                     if let Some(debug_info) = details.debug_info() { debug_info.stack_entries.clone() } else { vec![] };
+
                 Self::Server(ServerError::new(code, domain, status.message().to_owned(), stack_trace))
             } else {
                 Self::from_message(status.message())
@@ -364,7 +378,6 @@ impl From<Status> for Error {
                 Self::parse_unavailable(status.message())
             } else if status.code() == Code::Unknown
                 || is_rst_stream(&status)
-                || status.code() == Code::InvalidArgument
                 || status.code() == Code::FailedPrecondition
                 || status.code() == Code::AlreadyExists
             {

rust/src/common/error.rs

@@ -35,12 +35,12 @@ use crate::{
     error::ServerError,
     info::UserInfo,
     user::User,
-    Options, TransactionType,
+    Credentials, Options, TransactionType,
 };
 
 #[derive(Debug)]
 pub(super) enum Request {
-    ConnectionOpen { driver_lang: String, driver_version: String },
+    ConnectionOpen { driver_lang: String, driver_version: String, credentials: Credentials },
 
     ServersAll,
 

rust/src/connection/message.rs

@@ -17,11 +17,12 @@
  * under the License.
  */
 
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
 use tonic::{
     body::BoxBody,
     client::GrpcService,
+    metadata::MetadataValue,
     service::{
         interceptor::{InterceptedService, ResponseFuture as InterceptorResponseFuture},
         Interceptor,
@@ -65,20 +66,33 @@ pub(super) fn open_callcred_channel(
 #[derive(Debug)]
 pub(super) struct CallCredentials {
     credentials: Credentials,
+    token: RwLock<Option<String>>,
 }
 
 impl CallCredentials {
     pub(super) fn new(credentials: Credentials) -> Self {
-        Self { credentials }
+        Self { credentials, token: RwLock::new(None) }
     }
 
-    pub(super) fn username(&self) -> &str {
-        self.credentials.username()
+    pub(super) fn credentials(&self) -> &Credentials {
+        &self.credentials
+    }
+
+    pub(super) fn set_token(&self, token: String) {
+        *self.token.write().expect("Expected token write lock acquisition on set") = Some(token);
+    }
+
+    pub(super) fn reset_token(&self) {
+        *self.token.write().expect("Expected token write lock acquisition on reset") = None;
     }
 
     pub(super) fn inject(&self, mut request: Request<()>) -> Request<()> {
-        request.metadata_mut().insert("username", self.credentials.username().try_into().unwrap());
-        request.metadata_mut().insert("password", self.credentials.password().try_into().unwrap());
+        if let Some(token) = &*self.token.read().expect("Expected token read lock acquisition on inject") {
+            request.metadata_mut().insert(
+                "authorization",
+                format!("Bearer {}", token).try_into().expect("Expected authorization header formatting"),
+            );
+        }
         request
     }
 }

rust/src/connection/network/channel.rs

@@ -19,8 +19,8 @@
 
 use itertools::Itertools;
 use typedb_protocol::{
-    connection, database, database_manager, query::initial_res::Res, server_manager, transaction, user, user_manager,
-    Version::Version,
+    authentication, connection, database, database_manager, query::initial_res::Res, server_manager, transaction, user,
+    user_manager, Version::Version,
 };
 use uuid::Uuid;
 
@@ -32,14 +32,18 @@ use crate::{
     error::{ConnectionError, InternalError, ServerError},
     info::UserInfo,
     user::User,
+    Credentials,
 };
 
 impl TryIntoProto<connection::open::Req> for Request {
     fn try_into_proto(self) -> Result<connection::open::Req> {
         match self {
-            Self::ConnectionOpen { driver_lang, driver_version } => {
-                Ok(connection::open::Req { version: Version.into(), driver_lang, driver_version })
-            }
+            Self::ConnectionOpen { driver_lang, driver_version, credentials } => Ok(connection::open::Req {
+                version: Version.into(),
+                driver_lang,
+                driver_version,
+                authentication: Some(credentials.try_into_proto()?),
+            }),
             other => Err(InternalError::UnexpectedRequestType { request_type: format!("{other:?}") }.into()),
         }
     }
@@ -225,14 +229,28 @@ impl TryIntoProto<user::delete::Req> for Request {
     }
 }
 
+impl TryIntoProto<authentication::token::create::Req> for Credentials {
+    fn try_into_proto(self) -> Result<authentication::token::create::Req> {
+        Ok(authentication::token::create::Req {
+            credentials: Some(authentication::token::create::req::Credentials::Password(
+                authentication::token::create::req::Password {
+                    username: self.username().to_owned(),
+                    password: self.password().to_owned(),
+                },
+            )),
+        })
+    }
+}
+
 impl TryFromProto<connection::open::Res> for Response {
     fn try_from_proto(proto: connection::open::Res) -> Result<Self> {
         let mut database_infos = Vec::new();
-        for database_info_proto in proto.databases_all.unwrap().databases {
+        for database_info_proto in proto.databases_all.expect("Expected databases data").databases {
             database_infos.push(DatabaseInfo::try_from_proto(database_info_proto)?);
         }
         Ok(Self::ConnectionOpen {
-            connection_id: Uuid::from_slice(proto.connection_id.unwrap().id.as_slice()).unwrap(),
+            connection_id: Uuid::from_slice(proto.connection_id.expect("Expected connection id").id.as_slice())
+                .unwrap(),
             server_duration_millis: proto.server_duration_millis,
             databases: database_infos,
         })

rust/src/connection/network/proto/message.rs

@@ -25,12 +25,15 @@ use tokio::sync::mpsc::{unbounded_channel as unbounded_async, UnboundedSender};
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tonic::{Response, Status, Streaming};
 use typedb_protocol::{
-    connection, database, database_manager, server_manager, transaction, type_db_client::TypeDbClient as GRPC, user,
-    user_manager,
+    authentication, connection, database, database_manager, server_manager, transaction,
+    type_db_client::TypeDbClient as GRPC, user, user_manager,
 };
 
 use super::channel::{CallCredentials, GRPCChannel};
-use crate::common::{error::ConnectionError, Error, Result, StdResult};
+use crate::{
+    common::{error::ConnectionError, Error, Result, StdResult},
+    connection::network::proto::TryIntoProto,
+};
 
 type TonicResult<T> = StdResult<Response<T>, Status>;
 
@@ -45,15 +48,41 @@ impl<Channel: GRPCChannel> RPCStub<Channel> {
         Self { grpc: GRPC::new(channel), call_credentials }
     }
 
-    async fn call<F, R>(&mut self, call: F) -> Result<R>
+    async fn call_with_auto_renew_token<F, R>(&mut self, call: F) -> Result<R>
     where
         for<'a> F: Fn(&'a mut Self) -> BoxFuture<'a, Result<R>>,
     {
-        call(self).await
+        match call(self).await {
+            Err(Error::Connection(ConnectionError::TokenCredentialInvalid)) => {
+                debug!("Request rejected because token credential was invalid. Renewing token and trying again...");
+                self.renew_token().await?;
+                call(self).await
+            }
+            res => res,
+        }
+    }
+
+    async fn renew_token(&mut self) -> Result {
+        if let Some(call_credentials) = &self.call_credentials {
+            trace!("Renewing token...");
+            call_credentials.reset_token();
+            let request = call_credentials.credentials().clone().try_into_proto()?;
+            let token = self.grpc.authentication_token_create(request).await?.into_inner().token;
+            call_credentials.set_token(token);
+            trace!("Token renewed");
+        }
+        Ok(())
     }
 
     pub(super) async fn connection_open(&mut self, req: connection::open::Req) -> Result<connection::open::Res> {
-        self.single(|this| Box::pin(this.grpc.connection_open(req.clone()))).await
+        let result = self.single(|this| Box::pin(this.grpc.connection_open(req.clone()))).await;
+        if let Ok(response) = &result {
+            if let Some(call_credentials) = &self.call_credentials {
+                call_credentials
+                    .set_token(response.authentication.as_ref().expect("Expected authentication token").token.clone());
+            }
+        }
+        result
     }
 
     pub(super) async fn servers_all(&mut self, req: server_manager::all::Req) -> Result<server_manager::all::Res> {
@@ -107,7 +136,7 @@ impl<Channel: GRPCChannel> RPCStub<Channel> {
         &mut self,
         open_req: transaction::Req,
     ) -> Result<(UnboundedSender<transaction::Client>, Streaming<transaction::Server>)> {
-        self.call(|this| {
+        self.call_with_auto_renew_token(|this| {
             let transaction_req = transaction::Client { reqs: vec![open_req.clone()] };
             Box::pin(async {
                 let (sender, receiver) = unbounded_async();
@@ -154,6 +183,6 @@ impl<Channel: GRPCChannel> RPCStub<Channel> {
         for<'a> F: Fn(&'a mut Self) -> BoxFuture<'a, TonicResult<R>> + Send + Sync,
         R: 'static,
     {
-        self.call(|this| Box::pin(call(this).map(|r| Ok(r?.into_inner())))).await
+        self.call_with_auto_renew_token(|this| Box::pin(call(this).map(|r| Ok(r?.into_inner())))).await
     }
 }