Skip to content

Fixes try_execute_command message parsing bug #560

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Aug 24, 2023
Merged

Fixes try_execute_command message parsing bug #560

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
66f2430
Fixes try_execute_command message parsing bug
zainkabani Aug 24, 2023
9601486
Fix initial segment logic
zainkabani Aug 24, 2023
fa21928
Add test
zainkabani Aug 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions src/query_router.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,9 @@ use crate::plugins::{Intercept, Plugin, PluginOutput, QueryLogger, TableAccess};
use crate::pool::PoolSettings;
use crate::sharding::Sharder;

use std::cmp;
use std::collections::BTreeSet;
use std::io::Cursor;
use std::{cmp, mem};

/// Regexes used to parse custom commands.
const CUSTOM_SQL_REGEXES: [&str; 7] = [
Expand Down Expand Up @@ -141,6 +141,7 @@ impl QueryRouter {
let mut message_cursor = Cursor::new(message_buffer);

let code = message_cursor.get_u8() as char;
let len = message_cursor.get_i32() as usize;

// Check for any sharding regex matches in any queries
match code as char {
Expand All @@ -150,9 +151,13 @@ impl QueryRouter {
|| self.pool_settings.sharding_key_regex.is_some()
{
// Check only the first block of bytes configured by the pool settings
let len = message_cursor.get_i32() as usize;
let seg = cmp::min(len - 5, self.pool_settings.regex_search_limit);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this variable unused now?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's used in the seg calculation. It needs to be read off the cursor to progress the index that it's pointing to

let initial_segment = String::from_utf8_lossy(&message_buffer[0..seg]);

let query_start_index = mem::size_of::<u8>() + mem::size_of::<i32>();

let initial_segment = String::from_utf8_lossy(
&message_buffer[query_start_index..query_start_index + seg],
);

// Check for a shard_id included in the query
if let Some(shard_id_regex) = &self.pool_settings.shard_id_regex {
Expand Down Expand Up @@ -192,7 +197,6 @@ impl QueryRouter {
return None;
}

let _len = message_cursor.get_i32() as usize;
let query = message_cursor.read_string().unwrap();

let regex_set = match CUSTOM_SQL_REGEX_SET.get() {
Expand Down Expand Up @@ -1291,6 +1295,11 @@ mod test {
// Shard should start out unset
assert_eq!(qr.active_shard, None);

// Don't panic when short query eg. ; is sent
let q0 = simple_query(";");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice! Happy to see test coverage here.

assert!(qr.try_execute_command(&q0) == None);
assert_eq!(qr.active_shard, None);

// Make sure setting it works
let q1 = simple_query("/* shard_id: 1 */ select 1 from foo;");
assert!(qr.try_execute_command(&q1) == None);
Expand Down