ronantakizawa/codereview-bench huggingface dataset:
Repository: libsql/libsql
File: libsql-server/src/stats.rs (Rust)

Code before:

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct QueriesStats {
    #[serde(default)]
    id: Option<Uuid>,

    #[serde(default)]
    stats_threshold: AtomicU64,
    #[serde(default)]
    stats: HashMap<String, QueryStats>,
}

impl QueriesStats {
    pub fn new() -> Arc<RwLock<Self>> {
        let mut this = QueriesStats::default();
        this.id = Some(Uuid::new_v4());
        Arc::new(RwLock::new(this))
    }

    pub fn register_query(&mut self, sql: &String, stat: QueryStats) {
        let (aggregated, new) = match self.stats.get(sql) {
            Some(stat) => (stat.clone(), false),
            None => (QueryStats::empty(), true),
        };

        let aggregated = aggregated.merge(&stat);
        if aggregated.elapsed_ms < self.stats_threshold.load(Ordering::Relaxed) {
            return;
        }

        self.stats.insert(sql.clone(), aggregated);

        if !new || self.stats.len() <= 30 {
            return;
        }

        let mut vec = self.stats.clone().into_iter().collect_vec();
        vec.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        let len = vec.len();
        if len <= 30 {
            return;
        }

        for i in 0..len - 30 {
            self.stats.remove(&vec[i].0);
        }

        self.stats_threshold
            .store(vec[len - 30].1.elapsed_ms, Ordering::Relaxed);
    }


Review comment (suggestion):
I think there is a way here if you embed the merge within the borrow of the get that you don't need to clone the stat coming from the hashmap. You might do something like

```rust
let aggregated = if let Some(aggregated) = self.stats.get(&sql) {
    aggregated.merge(&stat)
} else {
    stat
};
```

Code after:
    stats: HashMap<String, QueryStats>,
    elapsed: Duration,
    stats_threshold: u64,
    hist: Histogram<u32>,
    expires_at: Option<Instant>,
    created_at: DateTime<Utc>,
}

impl QueriesStats {
    fn new() -> Self {
        Self {
            id: Uuid::new_v4(),
            stats: HashMap::new(),
            elapsed: Duration::default(),
            stats_threshold: 0,
            hist: Histogram::<u32>::new(3).unwrap(),
            expires_at: None,
            created_at: Utc::now(),
        }
    }

    fn with_expiration(expires_at: Instant) -> Self {
        let mut this = Self::new();
        this.expires_at = Some(expires_at);
        this
    }

    fn register_query(&mut self, sql: &String, stat: QueryStats) {
        self.elapsed += stat.elapsed;
        let _ = self.hist.record(stat.elapsed.as_millis() as u64);

        let (aggregated, new) = match self.stats.get(sql) {
            Some(aggregated) => (aggregated.merge(&stat), false),
            None => (stat, true),
        };

        if (aggregated.elapsed.as_micros() as u64) < self.stats_threshold {
            return;
        }

        self.stats.insert(sql.clone(), aggregated);

        if !new || self.stats.len() <= 30 {
            return;
        }

        while self.stats.len() > 30 {
            self.pop()
        }

        self.update_threshold();
Repository: libsql/libsql
File: libsql-server/src/auth/user_auth_strategies/mod.rs (Rust)

Code before:
pub mod disabled;
pub mod http_basic;
pub mod jwt;

use anyhow::Context;
pub use disabled::*;
pub use http_basic::*;
pub use jwt::*;

use super::{AuthError, Authenticated};

pub struct UserAuthContext {
    pub scheme: Option<String>,
    pub token: Option<String>, // token might not be required in some cases
}

impl TryFrom<&str> for UserAuthContext {
    type Error = anyhow::Error;

    fn try_from(auth_string: &str) -> Result<Self, Self::Error> {
        let (scheme, token) = auth_string.split_once(' ').context("malformed auth string`")?;
        Ok(UserAuthContext{scheme: Some(scheme.into()), token: Some(token.into())})    
    }
}
pub trait UserAuthStrategy: Sync + Send {
    fn authenticate(&self, context: UserAuthContext) -> Result<Authenticated, AuthError>;
}


Review comment (suggestion):
I think we should create some errors via `thiserror` and use them in this module, we could do a sub error enum that then embeds in the overall server error enum. That way instead of using context + strings we qualify it via a typed enum.

Code after:
pub mod disabled;
pub mod http_basic;
pub mod jwt;

pub use disabled::Disabled;
pub use http_basic::HttpBasic;
pub use jwt::Jwt;

use super::{AuthError, Authenticated};

#[derive(Debug)]
pub struct UserAuthContext {
    scheme: Option<String>,
    token: Option<String>,
}

impl UserAuthContext {
    pub fn scheme(&self) -> &Option<String> {
        &self.scheme
    }

    pub fn token(&self) -> &Option<String> {
        &self.token
    }

    pub fn empty() -> UserAuthContext {
        UserAuthContext {
            scheme: None,
            token: None,
        }
    }

    pub fn basic(creds: &str) -> UserAuthContext {
        UserAuthContext {
            scheme: Some("Basic".into()),
            token: Some(creds.into()),
        }
    }

    pub fn bearer(token: &str) -> UserAuthContext {
        UserAuthContext {
            scheme: Some("Bearer".into()),
            token: Some(token.into()),
Repository: libsql/libsql
File: libsql-server/src/auth/user_auth_strategies/mod.rs (Rust)

Code before:
pub mod disabled;
pub mod http_basic;
pub mod jwt;

use anyhow::Context;
pub use disabled::*;
pub use http_basic::*;
pub use jwt::*;

use super::{AuthError, Authenticated};

pub struct UserAuthContext {
    pub scheme: Option<String>,
    pub token: Option<String>, // token might not be required in some cases
}

impl TryFrom<&str> for UserAuthContext {
    type Error = anyhow::Error;

    fn try_from(auth_string: &str) -> Result<Self, Self::Error> {
        let (scheme, token) = auth_string.split_once(' ').context("malformed auth string`")?;
        Ok(UserAuthContext{scheme: Some(scheme.into()), token: Some(token.into())})    
    }
}
pub trait UserAuthStrategy: Sync + Send {
    fn authenticate(&self, context: UserAuthContext) -> Result<Authenticated, AuthError>;
}


Review comment (refactor):
Since we are converting things via TryFrom now it would be a good time to turn these into private fields and using setters/gettings. It will make refactoring easier in the future. We should in general avoid pub fields, something that we have a lot of now but plan on fixing. Especially, if we are going to release this as a crate.

Code after:
pub mod disabled;
pub mod http_basic;
pub mod jwt;

pub use disabled::Disabled;
pub use http_basic::HttpBasic;
pub use jwt::Jwt;

use super::{AuthError, Authenticated};

#[derive(Debug)]
pub struct UserAuthContext {
    scheme: Option<String>,
    token: Option<String>,
}

impl UserAuthContext {
    pub fn scheme(&self) -> &Option<String> {
        &self.scheme
    }

    pub fn token(&self) -> &Option<String> {
        &self.token
    }

    pub fn empty() -> UserAuthContext {
        UserAuthContext {
            scheme: None,
            token: None,
        }
    }

    pub fn basic(creds: &str) -> UserAuthContext {
        UserAuthContext {
            scheme: Some("Basic".into()),
            token: Some(creds.into()),
        }
    }
Repository: libsql/libsql
File: libsql-server/src/hrana/ws/session.rs (Rust)

Code before:
    #[error("Cursor {cursor_id} has failed to open")]
    CursorNotOpen { cursor_id: i32 },
    #[error("The server already stores {count} SQL texts, it cannot store more")]
    SqlTooMany { count: usize },
    #[error(transparent)]
    Stmt(stmt::StmtError),
    #[error(transparent)]
    Batch(batch::BatchError),
}

pub(super) async fn handle_initial_hello(
    server: &Server,
    version: Version,
    jwt: Option<String>,
    namespace: NamespaceName,
) -> Result<Session> {
    let namespace_jwt_key = server
        .namespaces
        .with(namespace.clone(), |ns| ns.jwt_key())
        .await??;

    let auth = namespace_jwt_key
        .map(Jwt::new)
        .map(Auth::new)
        .unwrap_or(server.user_auth_strategy.clone())
        .authenticate(UserAuthContext { scheme: Some("Bearer".into()), token: jwt })
        .map_err(|err| anyhow!(ResponseError::Auth { source: err }))?;

    Ok(Session {
        auth,
        version,
        streams: HashMap::new(),
        sqls: HashMap::new(),
        cursors: HashMap::new(),
    })
}

pub(super) async fn handle_repeated_hello(
    server: &Server,
    session: &mut Session,
    jwt: Option<String>,
    namespace: NamespaceName,
) -> Result<()> {
    if session.version < Version::Hrana2 {
        bail!(ProtocolError::NotSupported {
            what: "Repeated hello message",
            min_version: Version::Hrana2,
        })
    }
    let namespace_jwt_key = server
        .namespaces

Review comment (security):
Related to my comment above this would read better if we had a constructor like `UserAuthContext::from_scheme("Bearer", jwt)`.

Code after:
    #[error("Cursor {cursor_id} has failed to open")]
    CursorNotOpen { cursor_id: i32 },
    #[error("The server already stores {count} SQL texts, it cannot store more")]
    SqlTooMany { count: usize },
    #[error(transparent)]
    Stmt(stmt::StmtError),
    #[error(transparent)]
    Batch(batch::BatchError),
}

pub(super) async fn handle_initial_hello(
    server: &Server,
    version: Version,
    jwt: Option<String>,
    namespace: NamespaceName,
) -> Result<Session> {
    // todo dupe #auth
    let namespace_jwt_key = server
        .namespaces
        .with(namespace.clone(), |ns| ns.jwt_key())
        .await??;

    let auth = namespace_jwt_key
        .map(Jwt::new)
        .map(Auth::new)
        .unwrap_or(server.user_auth_strategy.clone())
        .authenticate(Ok(UserAuthContext::bearer_opt(jwt)))
        .map_err(|err| anyhow!(ResponseError::Auth { source: err }))?;

    Ok(Session {
        auth,
        version,
        streams: HashMap::new(),
        sqls: HashMap::new(),
        cursors: HashMap::new(),
    })
}

pub(super) async fn handle_repeated_hello(
    server: &Server,
    session: &mut Session,
    jwt: Option<String>,
    namespace: NamespaceName,
) -> Result<()> {
    if session.version < Version::Hrana2 {
        bail!(ProtocolError::NotSupported {
            what: "Repeated hello message",
            min_version: Version::Hrana2,
        })
    }
    // todo dupe #auth
Repository: libsql/libsql
File: libsql-server/src/http/user/extract.rs (Rust)

Code before:

use super::{db_factory, AppState};

#[async_trait::async_trait]
impl FromRequestParts<AppState> for RequestContext {
    type Rejection = crate::error::Error;

    async fn from_request_parts(
        parts: &mut axum::http::request::Parts,
        state: &AppState,
    ) -> std::result::Result<Self, Self::Rejection> {

        // start todo this block is same as the one in mod.rs
        let namespace = db_factory::namespace_from_headers(
            &parts.headers,
            state.disable_default_namespace,
            state.disable_namespaces,
        )?;

        let namespace_jwt_key = state
            .namespaces
            .with(namespace.clone(), |ns| ns.jwt_key())
            .await??;

        let context = parts.headers
        .get(hyper::header::AUTHORIZATION).context("auth header not found") // todo this context is swallowed for now, gotta fix that but not with panicking
        .and_then(|h| h.to_str().context("non ascii auth token"))
        .and_then(|t| t.try_into())
        .unwrap_or(UserAuthContext{scheme: None, token: None});


        let authenticated = namespace_jwt_key
        .map(Jwt::new)
        .map(Auth::new)
        .unwrap_or_else(|| state.user_auth_strategy.clone())
        .authenticate(context)?;            

        Ok(Self::new(
            authenticated,
            namespace,
            state.namespaces.meta_store().clone(),
        ))
    }
}


Review comment (suggestion):
Agreed, on the removing the panic 👍 I think moving to thiserror based on my comment above would be good here.

Code after:
use super::{db_factory, AppState};

#[async_trait::async_trait]
impl FromRequestParts<AppState> for RequestContext {
    type Rejection = crate::error::Error;

    async fn from_request_parts(
        parts: &mut axum::http::request::Parts,
        state: &AppState,
    ) -> std::result::Result<Self, Self::Rejection> {
        let namespace = db_factory::namespace_from_headers(
            &parts.headers,
            state.disable_default_namespace,
            state.disable_namespaces,
        )?;
        // todo dupe #auth
        let namespace_jwt_key = state
            .namespaces
            .with(namespace.clone(), |ns| ns.jwt_key())
            .await??;

        let context = parts
            .headers
            .get(hyper::header::AUTHORIZATION)
            .ok_or(AuthError::AuthHeaderNotFound)
            .and_then(|h| h.to_str().map_err(|_| AuthError::AuthHeaderNonAscii))
            .and_then(|t| UserAuthContext::from_auth_str(t));

        let authenticated = namespace_jwt_key
            .map(Jwt::new)
            .map(Auth::new)
            .unwrap_or_else(|| state.user_auth_strategy.clone())
            .authenticate(context)?;

        Ok(Self::new(
            authenticated,
            namespace,
            state.namespaces.meta_store().clone(),
        ))
    }
}
Repository: libsql/libsql
File: libsql-server/src/auth/user_auth_strategies/mod.rs (Rust)

Code before:
    }

    pub fn bearer(token: &str) -> UserAuthContext {
        UserAuthContext {
            scheme: Some("Bearer".into()),
            token: Some(token.into()),
        }
    }

    pub fn bearer_opt(token: Option<String>) -> UserAuthContext {
        UserAuthContext {
            scheme: Some("Bearer".into()),
            token: token,
        }
    }

    pub fn new(scheme: &str, token: &str) -> UserAuthContext {
        UserAuthContext {
            scheme: Some(scheme.into()),
            token: Some(token.into()),
        }
    }

    pub fn from_auth_str(auth_string: &str) -> Result<Self, AuthError> {
        let (scheme, token) = auth_string
            .split_once(' ')
            .ok_or(AuthError::AuthStringMalformed)?;
        Ok(UserAuthContext::new(scheme, token))
    }
}

pub trait UserAuthStrategy: Sync + Send {
    fn authenticate(&self, context: UserAuthContext) -> Result<Authenticated, AuthError>;
}


Review comment (bug):
Related to my other comment I see now what the problem is. I think we should support an empty auth string mostly because this is the workflow we have for a ton of tools outside of libsql and in tests. The problem is that now that the parsing of the header is separate from the authentication method we don't know that auth is disabled, because we should throw this error if jwt is enabled and we can't split. 

I think the solution here would be to not return an error here but somehow represent the state within the auth context of an empty auth header and let the auth strategy decide what to do with that. This would allow us to revert all the `dummy_token` stuff and fix the encryption test, as well as ensure that all our other tools continue to work locally.

Code after:
    }

    pub fn bearer(token: &str) -> UserAuthContext {
        UserAuthContext {
            scheme: Some("Bearer".into()),
            token: Some(token.into()),
        }
    }

    pub fn bearer_opt(token: Option<String>) -> UserAuthContext {
        UserAuthContext {
            scheme: Some("Bearer".into()),
            token: token,
        }
    }

    pub fn new(scheme: &str, token: &str) -> UserAuthContext {
        UserAuthContext {
            scheme: Some(scheme.into()),
            token: Some(token.into()),
        }
    }

    pub fn from_auth_str(auth_string: &str) -> Result<Self, AuthError> {
        let (scheme, token) = auth_string
            .split_once(' ')
            .ok_or(AuthError::AuthStringMalformed)?;
        Ok(UserAuthContext::new(scheme, token))
    }
}

pub trait UserAuthStrategy: Sync + Send {
    fn authenticate(
        &self,
        context: Result<UserAuthContext, AuthError>,
    ) -> Result<Authenticated, AuthError>;
}
Repository: libsql/libsql
File: libsql-server/src/hrana/ws/session.rs (Rust)

Code before:
use std::collections::HashMap;
use std::sync::Arc;

use anyhow::{anyhow, bail, Result};
use futures::future::BoxFuture;
use tokio::sync::{mpsc, oneshot};

use super::super::{batch, cursor, stmt, ProtocolError, Version};
use super::{proto, Server};
use crate::auth::user_auth_strategies::UserAuthContext;
use crate::auth::{Auth, AuthError, Authenticated, Jwt};
use crate::connection::{Connection as _, RequestContext};
use crate::database::Connection;
use crate::namespace::NamespaceName;

/// Session-level state of an authenticated Hrana connection.
pub struct Session {
    pub auth: Authenticated,
    pub version: Version,
    streams: HashMap<i32, StreamHandle>,
    sqls: HashMap<i32, String>,
    cursors: HashMap<i32, i32>,
}

impl Session {
    pub fn new(auth: Authenticated, version: Version) -> Self {
        Self { 
            auth, 
            version, 
            streams: HashMap::new(), 
            sqls: HashMap::new(), 
            cursors: HashMap::new(),
        }
    }
}


struct StreamHandle {
    job_tx: mpsc::Sender<StreamJob>,
    cursor_id: Option<i32>,
}

/// An arbitrary job that is executed on a [`Stream`].

Review comment (suggestion):
I think you could avoid leaking those details by having `handle_initial_hello`, `handle_repeated_hello` that both call private `handle_hello`.

Code after:
use std::collections::HashMap;
use std::sync::Arc;

use anyhow::{anyhow, bail, Error, Result};
use futures::future::BoxFuture;
use tokio::sync::{mpsc, oneshot};

use super::super::{batch, cursor, stmt, ProtocolError, Version};
use super::{proto, Server};
use crate::auth::user_auth_strategies::UserAuthContext;
use crate::auth::{Auth, AuthError, Authenticated, Jwt};
use crate::connection::{Connection as _, RequestContext};
use crate::database::Connection;
use crate::namespace::NamespaceName;

/// Session-level state of an authenticated Hrana connection.
pub struct Session {
    auth: Authenticated,
    version: Version,
    streams: HashMap<i32, StreamHandle>,
    sqls: HashMap<i32, String>,
    cursors: HashMap<i32, i32>,
}

impl Session {
    pub fn new(auth: Authenticated, version: Version) -> Self {
        Self {
            auth,
            version,
            streams: HashMap::new(),
            sqls: HashMap::new(),
            cursors: HashMap::new(),
        }
    }

    pub fn update_auth(&mut self, auth: Authenticated) -> Result<(), Error> {
        if self.version < Version::Hrana2 {
            bail!(ProtocolError::NotSupported {
                what: "Repeated hello message",
                min_version: Version::Hrana2,
            })
        }
        self.auth = auth;
Repository: libsql/libsql
File: libsql-server/src/auth/user_auth_strategies/jwt.rs (Rust)

Code before:
use chrono::{DateTime, Utc};

use crate::{
    auth::{authenticated::LegacyAuth, AuthError, Authenticated, Authorized, Permission},
    namespace::NamespaceName,
};

use super::{UserAuthContext, UserAuthStrategy};

pub enum DecodingKeyContainer {
    Single(jsonwebtoken::DecodingKey),
    Multiple(Vec<jsonwebtoken::DecodingKey>),
}

pub struct Jwt {
    key: DecodingKeyContainer,
}

impl UserAuthStrategy for Jwt {
    fn authenticate(
        &self,
        context: Result<UserAuthContext, AuthError>,
    ) -> Result<Authenticated, AuthError> {
        tracing::trace!("executing jwt auth");

        let ctx = context?;

        let UserAuthContext {
            scheme: Some(scheme),
            token: Some(token),
        } = ctx
        else {
            return Err(AuthError::HttpAuthHeaderInvalid);
        };


Review comment (performance):
It won't be long until we forget why we implemented this single/multiple distinction
it is a seed of complexity -> it requires forked logic to handle
any chance we could have one more round of conversations about how much real overhead a Vec<> introduces?

Code after:
use chrono::{DateTime, Utc};

use crate::{
    auth::{authenticated::LegacyAuth, AuthError, Authenticated, Authorized, Permission},
    namespace::NamespaceName,
};

use super::{UserAuthContext, UserAuthStrategy};

pub struct Jwt {
    keys: Vec<jsonwebtoken::DecodingKey>,
}

impl UserAuthStrategy for Jwt {
    fn authenticate(
        &self,
        context: Result<UserAuthContext, AuthError>,
    ) -> Result<Authenticated, AuthError> {
        tracing::trace!("executing jwt auth");

        let ctx = context?;

        let UserAuthContext {
            scheme: Some(scheme),
            token: Some(token),
        } = ctx
        else {
            return Err(AuthError::HttpAuthHeaderInvalid);
        };

        if !scheme.eq_ignore_ascii_case("bearer") {
            return Err(AuthError::HttpAuthHeaderUnsupportedScheme);
        }

        validate_any_jwt(&self.keys, &token)
Repository: libsql/libsql
File: libsql-server/src/http/user/listen.rs (Rust)

Code before:
) {
    let mut msg = msg;
    if let (Some(requested), Some(received)) = (action, msg.action.as_ref()) {
        if !received.eq(requested) {
            return;
        }
        // hide it from the reponse, in case user requested a specific action
        msg.action = None;
    }
    let msg = serde_json::to_string(&msg).unwrap() + "\n";
    if let Err(error) = writer.write_all(msg.as_bytes()) {
        tracing::warn!("failed to write listen error: {}", error);
    }
}

pub enum ListenError {
    Error(RecvError),
}

impl Serialize for ListenError {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let msg = match self {
            ListenError::Error(RecvError::Closed) => "listen channel was closd".to_string(),
            ListenError::Error(RecvError::Lagged(n)) => format!("lost {} messages", n),
        };
        let mut state = serializer.serialize_struct("foo", 1)?;
        state.serialize_field("error", &msg)?;
        state.end()
    }
}


Review comment (suggestion):
```suggestion
            ListenError::Error(RecvError::Closed) => "listen channel was closed".to_string(),
```

Code after:
}

async fn listen_stream(
    store: NamespaceStore,
    namespace: NamespaceName,
    table: String,
    actions: Vec<Action>,
) -> impl Stream<Item = crate::Result<AggregatorEvent>> {
    async_stream::try_stream! {
        let _sub = Subscription {
            store: store.clone(),
            namespace: namespace.clone(),
            table: table.clone(),
        };

        let mut stream = store.subscribe(namespace.clone(), table.clone());

        while let Some(item) = stream.next().await  {
            match item {
                Ok(msg) => if filter_actions(&msg, &actions) {
                    LISTEN_EVENTS_SENT.increment(1);
                    yield AggregatorEvent::Changes(msg);
                },
                Err(BroadcastStreamRecvError::Lagged(n)) => {
                    LISTEN_EVENTS_DROPPED.increment(n as u64);
                    yield AggregatorEvent::Error(&LAGGED_MSG);
                },
            }
        }
    }
}

fn filter_actions(msg: &BroadcastMsg, actions: &Vec<Action>) -> bool {
    for action in actions {
        let count = match action {
            Action::DELETE => msg.delete,
            Action::INSERT => msg.insert,
            Action::UPDATE => msg.update,
            Action::UNKNOWN => msg.unknown,
        };
        if count > 0 {
            return true;
        }
    }
    actions.is_empty()
}