Skip to content

Commit

Permalink
Handle consumed buffer legth in returned result/value
Browse files Browse the repository at this point in the history 
Signed-off-by: Jean-Christophe Dubois <[email protected]>
  • Loading branch information
@jcdubois
jcdubois committed Oct 23, 2023
1 parent 59aa87f commit c7ba61f
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 50 deletions.
40 changes: 17 additions & 23 deletions src/decoder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,15 +40,13 @@ pub fn decode<'a>(src: &[u8], dst: &'a mut [u8]) -> Result<&'a [u8], HSError> {
let mut dec = HeatshrinkDecoder::new();

while total_input_size < src.len() {
let mut segment_input_size = 0;

// Fill the input buffer from the src buffer
match dec.sink(&src[total_input_size..], &mut segment_input_size) {
HSsinkRes::HSRSinkOK => {
match dec.sink(&src[total_input_size..]) {
(HSsinkRes::HSRSinkOK, segment_input_size) => {
total_input_size += segment_input_size;
}
HSsinkRes::HSRSinkFull => {}
HSsinkRes::HSRSinkErrorMisuse => {
(HSsinkRes::HSRSinkFull, _) => {}
(HSsinkRes::HSRSinkErrorMisuse, _) => {
return Err(HSError::HSErrorInternal);
}
}
Expand All @@ -58,17 +56,15 @@ pub fn decode<'a>(src: &[u8], dst: &'a mut [u8]) -> Result<&'a [u8], HSError> {
} else {
// process the current input buffer
loop {
let mut segment_output_size = 0;

match dec.poll(&mut dst[total_output_size..], &mut segment_output_size) {
HSpollRes::HSRPollMore => {
match dec.poll(&mut dst[total_output_size..]) {
(HSpollRes::HSRPollMore, _) => {
return Err(HSError::HSErrorOutputFull);
}
HSpollRes::HSRPollEmpty => {
(HSpollRes::HSRPollEmpty, segment_output_size) => {
total_output_size += segment_output_size;
break;
}
HSpollRes::HSRPollErrorMisuse => {
(HSpollRes::HSRPollErrorMisuse, _) => {
return Err(HSError::HSErrorInternal);
}
}
Expand Down Expand Up @@ -122,12 +118,11 @@ impl HeatshrinkDecoder {
}

/// Add an input buffer to be processed/uncompressed
pub fn sink(&mut self, input_buffer: &[u8], input_size: &mut usize) -> HSsinkRes {
pub fn sink(&mut self, input_buffer: &[u8]) -> (HSsinkRes, usize) {
let remaining_size = self.input_buffer.len() - self.input_size as usize;

if remaining_size == 0 {
*input_size = 0;
return HSsinkRes::HSRSinkFull;
return (HSsinkRes::HSRSinkFull, 0);
}

let copy_size = if remaining_size < input_buffer.len() {
Expand All @@ -140,20 +135,19 @@ impl HeatshrinkDecoder {
self.input_buffer[self.input_size as usize..(self.input_size as usize + copy_size)]
.copy_from_slice(&input_buffer[0..copy_size]);
self.input_size += copy_size as u16;
*input_size = copy_size;

return HSsinkRes::HSRSinkOK;
return (HSsinkRes::HSRSinkOK, copy_size);
}

/// function to process the input/internal buffer and put the uncompressed
/// stream in the provided buffer.
pub fn poll(&mut self, output_buffer: &mut [u8], output_size: &mut usize) -> HSpollRes {
pub fn poll(&mut self, output_buffer: &mut [u8]) -> (HSpollRes, usize) {
if output_buffer.len() == 0 {
return HSpollRes::HSRPollErrorMisuse;
return (HSpollRes::HSRPollErrorMisuse, 0);
} else {
*output_size = 0;
let mut output_size: usize = 0;

let mut output_info = OutputInfo::new(output_buffer, output_size);
let mut output_info = OutputInfo::new(output_buffer, &mut output_size);

loop {
let in_state = self.state;
Expand Down Expand Up @@ -183,9 +177,9 @@ impl HeatshrinkDecoder {
// output buffer are exhausted.
if self.state == in_state {
if output_info.can_take_byte() {
return HSpollRes::HSRPollEmpty;
return (HSpollRes::HSRPollEmpty, output_size);
} else {
return HSpollRes::HSRPollMore;
return (HSpollRes::HSRPollMore, output_size);
}
}
}
Expand Down
47 changes: 20 additions & 27 deletions src/encoder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,15 +66,13 @@ pub fn encode<'a>(src: &[u8], dst: &'a mut [u8]) -> Result<&'a [u8], HSError> {
let mut enc = HeatshrinkEncoder::new();

while total_input_size < src.len() {
let mut segment_input_size = 0;

// Fill the input buffer from the src buffer
match enc.sink(&src[total_input_size..], &mut segment_input_size) {
HSsinkRes::HSRSinkOK => {
match enc.sink(&src[total_input_size..]) {
(HSsinkRes::HSRSinkOK, segment_input_size) => {
total_input_size += segment_input_size;
}
HSsinkRes::HSRSinkFull => {}
HSsinkRes::HSRSinkErrorMisuse => {
(HSsinkRes::HSRSinkFull, _) => {}
(HSsinkRes::HSRSinkErrorMisuse, _) => {
return Err(HSError::HSErrorInternal);
}
}
Expand All @@ -92,17 +90,15 @@ pub fn encode<'a>(src: &[u8], dst: &'a mut [u8]) -> Result<&'a [u8], HSError> {
} else {
// process the current input buffer
loop {
let mut segment_output_size = 0;

match enc.poll(&mut dst[total_output_size..], &mut segment_output_size) {
HSpollRes::HSRPollMore => {
match enc.poll(&mut dst[total_output_size..]) {
(HSpollRes::HSRPollMore, _) => {
return Err(HSError::HSErrorOutputFull);
}
HSpollRes::HSRPollEmpty => {
(HSpollRes::HSRPollEmpty, segment_output_size) => {
total_output_size += segment_output_size;
break;
}
HSpollRes::HSRPollErrorMisuse => {
(HSpollRes::HSRPollErrorMisuse, _) => {
return Err(HSError::HSErrorInternal);
}
}
Expand Down Expand Up @@ -174,23 +170,22 @@ impl HeatshrinkEncoder {
}

/// Add an input buffer to be processed/compressed
pub fn sink(&mut self, input_buffer: &[u8], input_size: &mut usize) -> HSsinkRes {
pub fn sink(&mut self, input_buffer: &[u8]) -> (HSsinkRes, usize) {
/* Sinking more content after saying the content is done, tsk tsk */
if self.is_finishing() {
return HSsinkRes::HSRSinkErrorMisuse;
return (HSsinkRes::HSRSinkErrorMisuse, 0);
}

/* Sinking more content before processing is done */
if self.state != HSEstate::HSESNotFull {
return HSsinkRes::HSRSinkErrorMisuse;
return (HSsinkRes::HSRSinkErrorMisuse, 0);
}

let write_offset: usize = (self.get_input_offset() + self.input_size).into();
let remaining_size: usize = (self.get_input_buffer_size() - self.input_size).into();

if remaining_size == 0 {
*input_size = 0;
return HSsinkRes::HSRSinkFull;
return (HSsinkRes::HSRSinkFull, 0);
}

let copy_size = if remaining_size < input_buffer.len() {
Expand All @@ -203,31 +198,29 @@ impl HeatshrinkEncoder {
self.input_buffer[write_offset..write_offset + copy_size]
.copy_from_slice(&input_buffer[0..copy_size]);
self.input_size += copy_size as u16;
*input_size = copy_size;

if self.input_size == self.get_input_buffer_size() {
self.state = HSEstate::HSESFilled;
}

return HSsinkRes::HSRSinkOK;
return (HSsinkRes::HSRSinkOK, copy_size);
}

/// function to process the input/internal buffer and put the compressed
/// stream in the provided buffer.
pub fn poll(&mut self, output_buffer: &mut [u8], output_size: &mut usize) -> HSpollRes {
*output_size = 0;

pub fn poll(&mut self, output_buffer: &mut [u8]) -> (HSpollRes, usize) {
if output_buffer.len() == 0 {
return HSpollRes::HSRPollMore;
return (HSpollRes::HSRPollMore, 0);
} else {
let mut output_info = OutputInfo::new(output_buffer, output_size);
let mut output_size: usize = 0;
let mut output_info = OutputInfo::new(output_buffer, &mut output_size);

loop {
let in_state = self.state;

match in_state {
HSEstate::HSESNotFull => {
return HSpollRes::HSRPollEmpty;
return (HSpollRes::HSRPollEmpty, output_size);
}
HSEstate::HSESFilled => {
self.do_indexing();
Expand Down Expand Up @@ -255,15 +248,15 @@ impl HeatshrinkEncoder {
self.state = self.st_flush_bit_buffer(&mut output_info);
}
HSEstate::HSESDone => {
return HSpollRes::HSRPollEmpty;
return (HSpollRes::HSRPollEmpty, output_size);
}
}

// If the current state cannot advance, check if output
// buffer is exhausted.
if self.state == in_state {
if !output_info.can_take_byte() {
return HSpollRes::HSRPollMore;
return (HSpollRes::HSRPollMore, output_size);
}
}
}
Expand Down

0 comments on commit c7ba61f

Please sign in to comment.