Repository: libsql/libsql
File: libsql-server/src/hrana/ws/conn.rs (Rust)

Code before:
        tungstenite::Message::Pong(_) => Ok(true),
        tungstenite::Message::Close(_) => Ok(false),
        tungstenite::Message::Frame(_) => panic!("Received a tungstenite::Message::Frame"),
    }
}

async fn handle_client_msg(conn: &mut Conn, client_msg: proto::ClientMsg) -> Result<bool> {
    tracing::trace!("Received client msg: {:?}", client_msg);
    match client_msg {
        proto::ClientMsg::None => bail!(ProtocolError::NoneClientMsg),
        proto::ClientMsg::Hello(msg) => handle_hello_msg(conn, msg.jwt).await,
        proto::ClientMsg::Request(msg) => match msg.request {
            Some(request) => handle_request_msg(conn, msg.request_id, request).await,
            None => bail!(ProtocolError::NoneRequest),
        },
    }
}

async fn handle_hello_msg(conn: &mut Conn, jwt: Option<String>) -> Result<bool> {
    let auth = session::handle_hello(&conn.server, jwt, conn.namespace.clone()).await;

    let hello_res = match conn.session.as_mut() {
        None => auth
            .map(|auth| session::Session::new(auth, conn.version))
            .map(|s| conn.session = Some(s)),
        Some(session) => auth.map(|a| session.update_auth(a)).and_then(|op| op),
    };

    match hello_res {
        Ok(_) => {
            send_msg(conn, &proto::ServerMsg::HelloOk(proto::HelloOkMsg {})).await?;
            Ok(true)
        }
        Err(err) => match downcast_error(err) {
            Ok(error) => {
                send_msg(
                    conn,
                    &proto::ServerMsg::HelloError(proto::HelloErrorMsg { error }),
                )
                .await?;
                Ok(false)
            }
            Err(err) => Err(err),
        },
    }
}

async fn handle_request_msg(
    conn: &mut Conn,
    request_id: i32,
    request: proto::Request,

Review comment (security):
I'm wondering if the code wouldn't be a bit easier to follow if we structure it like this:
```
let hello_res = match auth {
    Ok(a) => match conn.session.as_mut() {
        None => {
            conn.session = Some(session::Session::new(a, conn.version));
            Ok(())
        }
        Some(session) => session.update_auth(a)
    }
    Err(e) => Err(e)
}
```

Code after:
        tungstenite::Message::Pong(_) => Ok(true),
        tungstenite::Message::Close(_) => Ok(false),
        tungstenite::Message::Frame(_) => panic!("Received a tungstenite::Message::Frame"),
    }
}

async fn handle_client_msg(conn: &mut Conn, client_msg: proto::ClientMsg) -> Result<bool> {
    tracing::trace!("Received client msg: {:?}", client_msg);
    match client_msg {
        proto::ClientMsg::None => bail!(ProtocolError::NoneClientMsg),
        proto::ClientMsg::Hello(msg) => handle_hello_msg(conn, msg.jwt).await,
        proto::ClientMsg::Request(msg) => match msg.request {
            Some(request) => handle_request_msg(conn, msg.request_id, request).await,
            None => bail!(ProtocolError::NoneRequest),
        },
    }
}

async fn handle_hello_msg(conn: &mut Conn, jwt: Option<String>) -> Result<bool> {
    let auth = session::handle_hello(&conn.server, jwt, conn.namespace.clone()).await;

    let hello_res = auth
        .map(|a| {
            if let Some(sess) = conn.session.as_mut() {
                sess.update_auth(a)
            } else {
                conn.session = Some(session::Session::new(a, conn.version));
                Ok(())
            }
        })
        .and_then(|o| o);

    match hello_res {
        Ok(_) => {
            send_msg(conn, &proto::ServerMsg::HelloOk(proto::HelloOkMsg {})).await?;
            Ok(true)
        }
        Err(err) => match downcast_error(err) {
            Ok(error) => {
                send_msg(
                    conn,
                    &proto::ServerMsg::HelloError(proto::HelloErrorMsg { error }),
                )
                .await?;
                Ok(false)
            }
            Err(err) => Err(err),
        },
    }
}