1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
//! A [`tokio_util::codec`] Codec that is used to encode and decode the
//! blather protocol.

use std::fmt;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::{cmp, collections::HashMap, mem};

use bytes::{BufMut, Bytes, BytesMut};

use tokio::io;

use tokio_util::codec::Decoder;
use tokio_util::codec::Encoder;

use crate::err::Error;
use crate::{KVLines, Params, Telegram};


/// Current state of decoder
/// Controls what, if anything, will be returned to the application.
#[derive(Clone, Debug, PartialEq)]
enum CodecState {
  Telegram,
  Params,
  KVLines,
  Chunks,
  Buf,
  File,
  Writer,
  Skip
}

/// Data returned to the application when the Codec's Decode iterator is
/// called and the decoder has a complete entity to return.
pub enum Input {
  Telegram(Telegram),
  KVLines(KVLines),
  Params(Params),
  Chunk(BytesMut, usize),
  Buf(BytesMut),
  File(PathBuf),
  WriteDone,
  SkipDone
}


/// The Codec (exposed externally as ClntIfCodec) is used to keep track of the
/// state of the inbound and outbound communication.
pub struct Codec {
  next_line_index: usize,
  max_line_length: usize,
  tg: Telegram,
  params: Params,
  kvlines: KVLines,
  state: CodecState,
  bin_remain: usize,
  pathname: Option<PathBuf>,
  writer: Option<Box<dyn Write + Send + Sync>>,
  buf: BytesMut
}

impl fmt::Debug for Codec {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    f.debug_struct("Codec").field("state", &self.state).finish()
  }
}

impl Default for Codec {
  fn default() -> Self {
    Codec::new()
  }
}


/// A Codec used to encode and decode the blather protocol.
///
/// # Notes
/// Normally the Codec object is hidden inside a [`Framed`] object. In order to
/// call methods in the codec it must be accessed through the Framed object:
///
/// ```compile_fail
/// let mut conn = Framed::new(socket, Codec::new());
/// // ...
/// conn.codec_mut().expect_chunks(len);
/// ```
///
/// [`Framed`]: https://docs.rs/tokio-util/0.3/tokio_util/codec/struct.Framed.html
impl Codec {
  pub fn new() -> Codec {
    Codec {
      next_line_index: 0,
      max_line_length: usize::MAX,
      tg: Telegram::new(),
      params: Params::new(),
      kvlines: KVLines::new(),
      state: CodecState::Telegram,
      bin_remain: 0,
      pathname: None,
      writer: None,
      buf: BytesMut::new()
    }
  }

  pub fn new_with_max_length(max_line_length: usize) -> Self {
    Codec {
      max_line_length,
      ..Codec::new()
    }
  }

  pub fn max_line_length(&self) -> usize {
    self.max_line_length
  }


  /// Determine how far into the buffer we'll search for a newline. If
  /// there's no max_length set, we'll read to the end of the buffer.
  fn find_newline(&self, buf: &BytesMut) -> (usize, Option<usize>) {
    let read_to = cmp::min(self.max_line_length.saturating_add(1), buf.len());
    let newline_offset = buf[self.next_line_index..read_to]
      .iter()
      .position(|b| *b == b'\n');

    (read_to, newline_offset)
  }


  /// `decode_telegram_lines` has encountered an eol, determined that the
  /// string is longer than zero characters, and thus passed the line to this
  /// function to process it.
  ///
  /// The first line received is a telegram topic.  This is a required line.
  /// Following lines are parameter lines, which are a single space character
  /// separated key/value pairs.
  fn decode_telegram_line(&mut self, line: &str) -> Result<(), Error> {
    if self.tg.get_topic().is_none() {
      self.tg.set_topic(line)?;
    } else {
      let idx = line.find(' ');
      if let Some(idx) = idx {
        let (k, v) = line.split_at(idx);
        let v = &v[1..v.len()];
        self.tg.add_param(k, v)?;
      }
    }
    Ok(())
  }

  /*
  fn getline_owned(
    &mut self,
    buf: &mut BytesMut
  ) -> Result<Option<String>, Error> {
    let (read_to, newline_offset) = self.find_newline(&buf);
    match newline_offset {
      Some(offset) => {
        // Found an eol
        let newline_index = offset + self.next_line_index;
        self.next_line_index = 0;
        let line = buf.split_to(newline_index + 1);
        let line = &line[..line.len() - 1];
        let line = utf8(without_carriage_return(line))?;

        Ok(Some(line.to_owned()))
      }
      None if buf.len() > self.max_line_length => Err(Error::BadFormat(
        "Exceeded maximum line length.".to_string()
      )),
      None => {
        // We didn't find a line or reach the length limit, so the next
        // call will resume searching at the current offset.
        self.next_line_index = read_to;

        // Returning Ok(None) instructs the FramedRead that more data is
        // needed.
        Ok(None)
      }
    }
  }
  */

  fn get_eol_idx(&mut self, buf: &BytesMut) -> Result<Option<usize>, Error> {
    let (read_to, newline_offset) = self.find_newline(&buf);
    match newline_offset {
      Some(offset) => {
        // Found an eol
        let newline_index = offset + self.next_line_index;
        self.next_line_index = 0;
        Ok(Some(newline_index + 1))
      }
      None if buf.len() > self.max_line_length => Err(Error::BadFormat(
        "Exceeded maximum line length.".to_string()
      )),
      None => {
        // Didn't find a line or reach the length limit, so the next
        // call will resume searching at the current offset.
        self.next_line_index = read_to;

        // Returning Ok(None) instructs the FramedRead that more data is
        // needed.
        Ok(None)
      }
    }
  }

  /// (New) data is available in the input buffer.
  /// Try to parse lines until an empty line as been encountered, at which
  /// point the buffer is parsed and returned in an [`Telegram`] buffer.
  ///
  /// If the buffer doesn't contain enough data to finalize a complete telegram
  /// buffer return `Ok(None)` to inform the calling FramedRead that more data
  /// is needed.
  ///
  /// [`Telegram`]: blather::Telegram
  fn decode_telegram_lines(
    &mut self,
    buf: &mut BytesMut
  ) -> Result<Option<Telegram>, Error> {
    loop {
      if let Some(idx) = self.get_eol_idx(buf)? {
        let line = buf.split_to(idx);
        let line = &line[..line.len() - 1];
        let line = utf8(without_carriage_return(line))?;

        // Empty line marks end of Telegram
        if line.is_empty() {
          // mem::take() can replace a member of a struct.
          // (This requires Default to be implemented for the object being
          // taken).
          return Ok(Some(mem::take(&mut self.tg)));
        } else {
          self.decode_telegram_line(&line)?;
        }
      } else {
        // Returning Ok(None) instructs the FramedRead that more data is
        // needed.
        return Ok(None);
      }
    }
  }


  /// Read buffer line-by-line, split each line at the first space character
  /// and store the left part as a key and the right part as a value in a
  /// Params structure.
  fn decode_params_lines(
    &mut self,
    buf: &mut BytesMut
  ) -> Result<Option<Params>, Error> {
    loop {
      if let Some(idx) = self.get_eol_idx(buf)? {
        // Found an eol
        let line = buf.split_to(idx);
        let line = &line[..line.len() - 1];
        let line = utf8(without_carriage_return(line))?;

        // Empty line marks end of Params
        if line.is_empty() {
          // Revert to expecting a telegram once a Params has been completed.
          // The application can override this when needed.
          self.state = CodecState::Telegram;

          // mem::take() can replace a member of a struct.
          // (This requires Default to be implemented for the object being
          // taken).
          return Ok(Some(mem::take(&mut self.params)));
        } else {
          let idx = line.find(' ');
          if let Some(idx) = idx {
            let (k, v) = line.split_at(idx);
            let v = &v[1..v.len()];
            self.params.add_param(k, v)?;
          }
        }
      } else {
        // Need more data
        return Ok(None);
      }
    }
  }

  fn decode_kvlines(
    &mut self,
    buf: &mut BytesMut
  ) -> Result<Option<KVLines>, Error> {
    loop {
      if let Some(idx) = self.get_eol_idx(buf)? {
        // Found an eol
        let line = buf.split_to(idx);
        let line = &line[..line.len() - 1];
        let line = utf8(without_carriage_return(line))?;

        // Empty line marks end of Params
        if line.is_empty() {
          // Revert to expecting a telegram once a KVLines  has been
          // completed.
          // The application can override this when needed.
          self.state = CodecState::Telegram;

          // mem::take() can replace a member of a struct.
          // (This requires Default to be implemented for the object being
          // taken).
          return Ok(Some(mem::take(&mut self.kvlines)));
        } else {
          let idx = line.find(' ');
          if let Some(idx) = idx {
            let (k, v) = line.split_at(idx);
            let v = &v[1..v.len()];
            self.kvlines.append(k, v);
          }
        }
      } else {
        // Need more data
        return Ok(None);
      }
    }
  }


  /// Set the decoder to treat the next `size` bytes as raw bytes to be
  /// received in chunks.
  ///
  /// # Decoder behavior
  /// The decoder will return an `Input::Chunk(buf, remain)` to the application
  /// each time a new chunk has been received. In addition to the actual
  /// chunk number of bytes remaining will be returned.  The remaining bytes
  /// value is adjusted to subtract the currently returned chunk, which means
  /// that the application can detect the end of the buffer by checking if
  /// the remaining value is zero.
  pub fn expect_chunks(&mut self, size: usize) {
    //println!("Expecting bin {}", size);
    self.state = CodecState::Chunks;
    self.bin_remain = size;
  }

  /// Expect a buffer of a certain size to be received.
  /// The returned buffer will be stored in process memory.
  ///
  /// # Decoder behavior
  /// One a complete buffer has been successfully reaceived the `Decoder` will
  /// return an `Input::Buf(b)` where `b` is a `bytes::BytesMut` containing the
  /// entire buffer.
  ///
  /// Once the entire buffer has been received by the `Decoder` it will revert
  /// to expect an `Input::Telegram`.
  pub fn expect_buf(&mut self, size: usize) -> Result<(), Error> {
    if size == 0 {
      return Err(Error::InvalidSize("The size must not be zero".to_string()));
    }
    self.state = CodecState::Buf;
    self.bin_remain = size;
    self.buf = BytesMut::with_capacity(size);
    Ok(())
  }

  /// Expects a certain amount of bytes of data to arrive from the peer, and
  /// that data should be stored to a file.
  ///
  /// # Decoder behavior
  /// On successful completion the Decoder will return an Input::File(pathname)
  /// once the entire file length has successfully been received, where the
  /// pathname is a PathBuf which matches the pathname parameter passed to
  /// this function.
  pub fn expect_file<P: Into<PathBuf>>(
    &mut self,
    pathname: P,
    size: usize
  ) -> Result<(), Error> {
    if size == 0 {
      return Err(Error::InvalidSize("The size must not be zero".to_string()));
    }
    self.state = CodecState::File;
    let pathname = pathname.into();
    self.writer = Some(Box::new(File::create(&pathname)?));
    self.pathname = Some(pathname);

    self.bin_remain = size;

    Ok(())
  }

  /// Called from an application to request that data should be written to a
  /// supplied writer.
  ///
  /// The writer's ownership will be transferred to the `Decoder` and will
  /// automatically be dropped once the entire buffer has been written.
  ///
  /// # Decoder behavior
  /// On successful completion the Decoder will return an Input::WriteDone to
  /// signal that the entire buffer has been received and written to the
  /// `Writer`.
  ///
  /// Once the complete `Params` buffer has been received the Decoder will
  /// revert back to waiting for a `Telegram`.
  pub fn expect_writer<W: 'static + Write + Send + Sync>(
    &mut self,
    writer: W,
    size: usize
  ) -> Result<(), Error> {
    if size == 0 {
      return Err(Error::InvalidSize("The size must not be zero".to_string()));
    }
    self.state = CodecState::Writer;
    self.writer = Some(Box::new(writer));
    self.bin_remain = size;
    Ok(())
  }

  /// Tell the Decoder to expect lines of key/value pairs.
  ///
  /// # Decoder behavior
  /// On successful completion the Framed StreamExt next() will return an
  /// Input::Params(params) once a complete `Params` buffer has been received.
  ///
  /// Once the complete `Params` buffer has been received the Decoder will
  /// revert back to waiting for a `Telegram`.
  pub fn expect_params(&mut self) {
    self.state = CodecState::Params;
  }

  /// Tell the Decoder to expect lines ordered key/value pairs.
  ///
  /// # Decoder behavior
  /// On successful completion the Framed StreamExt next() will return an
  /// Input::KVLines(kvlines) once a complete `KVLines` buffer has been
  /// received.
  ///
  /// Once the complete `KVLines` buffer has been received the Decoder will
  /// revert back to waiting for a `Telegram`.
  pub fn expect_kvlines(&mut self) {
    self.state = CodecState::KVLines;
  }

  /// Skip bytes.
  ///
  /// # Decoder behavior
  /// Simply ignore the number of specified bytes, then revert back to waiting
  /// for a Telegram.
  pub fn skip(&mut self, size: usize) -> Result<(), Error> {
    if size == 0 {
      return Err(Error::InvalidSize("The size must not be zero".to_string()));
    }
    self.state = CodecState::Skip;
    self.bin_remain = size;
    Ok(())
  }
}

fn utf8(buf: &[u8]) -> Result<&str, io::Error> {
  std::str::from_utf8(buf).map_err(|_| {
    io::Error::new(
      io::ErrorKind::InvalidData,
      "Unable to decode input as UTF8"
    )
  })
}

fn without_carriage_return(s: &[u8]) -> &[u8] {
  if let Some(&b'\r') = s.last() {
    &s[..s.len() - 1]
  } else {
    s
  }
}


/// A Decoder implementation that is used to assist in decoding data arriving
/// over a DDM client interface.
///
/// The default behavior for the Decoder is to wait for a Telegram buffer.  It
/// will, on success, return an `Input::Telegram(tg)`, where `tg` is a
/// `blather::Telegram` object.
impl Decoder for Codec {
  type Item = Input;
  type Error = crate::err::Error;

  fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Input>, Error> {
    // The codec's internal decoder state denotes whether lines or binary data
    // is currently being expected.
    match self.state {
      CodecState::Telegram => {
        // If decode_telegram_lines returns Some(value) it means that a
        // complete buffer has been received.
        let tg = self.decode_telegram_lines(buf)?;
        if let Some(tg) = tg {
          // A complete Telegram was received
          return Ok(Some(Input::Telegram(tg)));
        }

        // Returning Ok(None) tells the caller that we need more data
        Ok(None)
      }
      CodecState::Params => {
        // If decode_telegram_lines returns Some(value) it means that a
        // complete buffer has been received.
        let params = self.decode_params_lines(buf)?;
        if let Some(params) = params {
          // A complete Params buffer was received
          return Ok(Some(Input::Params(params)));
        }

        // Returning Ok(None) tells the caller that we need more data
        Ok(None)
      }
      CodecState::KVLines => {
        // If decode_telegram_lines returns Some(value) it means that a
        // complete buffer has been received.
        let kvlines = self.decode_kvlines(buf)?;
        if let Some(kvlines) = kvlines {
          // A complete Params buffer was received
          return Ok(Some(Input::KVLines(kvlines)));
        }

        // Returning Ok(None) tells the caller that we need more data
        Ok(None)
      }

      CodecState::Chunks => {
        if buf.is_empty() {
          // Need more data
          return Ok(None);
        }

        let read_to = cmp::min(self.bin_remain, buf.len());
        self.bin_remain -= read_to;

        if self.bin_remain == 0 {
          // When no more data is expected for this binary part, revert to
          // expecting Telegram lines
          self.state = CodecState::Telegram;
        }

        // Return a buffer and the amount of data remaining, this buffer
        // included.  The application can check if remain is 0 to determine
        // if it has received all the expected binary data.
        Ok(Some(Input::Chunk(buf.split_to(read_to), self.bin_remain)))
      }
      CodecState::Buf => {
        if buf.is_empty() {
          // Need more data
          return Ok(None);
        }
        let read_to = cmp::min(self.bin_remain, buf.len());

        // Transfer data from input to output buffer
        self.buf.put(buf.split_to(read_to));

        self.bin_remain -= read_to;
        if self.bin_remain != 0 {
          // Need more data
          return Ok(None);
        }

        // When no more data is expected for this binary part, revert to
        // expecting Telegram lines
        self.state = CodecState::Telegram;

        // Return a buffer and the amount of data remaining, this buffer
        // included.  The application can check if remain is 0 to determine
        // if it has received all the expected binary data.
        Ok(Some(Input::Buf(mem::take(&mut self.buf))))
      }
      CodecState::File | CodecState::Writer => {
        if buf.is_empty() {
          return Ok(None); // Need more data
        }

        // Read as much data as available or requested and write it to our
        // output.
        let read_to = cmp::min(self.bin_remain, buf.len());
        if let Some(ref mut f) = self.writer {
          f.write_all(&buf.split_to(read_to))?;
        }

        self.bin_remain -= read_to;
        if self.bin_remain != 0 {
          return Ok(None); // Need more data
        }

        // At this point the entire expected buffer has been received

        // Close file
        self.writer = None;

        // Return a buffer and the amount of data remaining, this buffer
        // included.  The application can check if remain is 0 to determine
        // if it has received all the expected binary data.
        let ret = if self.state == CodecState::File {
          let pathname = if let Some(ref fname) = self.pathname {
            fname.clone()
          } else {
            return Err(Error::BadState("Missing pathname".to_string()));
          };

          // Reset the pathname
          self.pathname = None;

          Input::File(pathname)
        } else {
          Input::WriteDone
        };

        // Revert to the default of expecting a telegram.
        self.state = CodecState::Telegram;

        Ok(Some(ret))
      } // CodecState::{File|Writer}
      CodecState::Skip => {
        if buf.is_empty() {
          return Ok(None); // Need more data
        }

        // Read as much data as available or requested and write it to our
        // output.
        let read_to = cmp::min(self.bin_remain, buf.len());
        let _ = buf.split_to(read_to);

        self.bin_remain -= read_to;
        if self.bin_remain != 0 {
          return Ok(None); // Need more data
        }

        // Revert to the default of expecting a telegram.
        self.state = CodecState::Telegram;

        Ok(Some(Input::SkipDone))
      } // CodecState::Skip
    } // match self.state
  }
}


impl Encoder<&Telegram> for Codec {
  type Error = crate::err::Error;

  fn encode(
    &mut self,
    tg: &Telegram,
    buf: &mut BytesMut
  ) -> Result<(), Error> {
    tg.encoder_write(buf)?;
    Ok(())
  }
}


impl Encoder<&Params> for Codec {
  type Error = crate::err::Error;

  fn encode(
    &mut self,
    params: &Params,
    buf: &mut BytesMut
  ) -> Result<(), Error> {
    params.encoder_write(buf)?;
    Ok(())
  }
}


impl Encoder<&HashMap<String, String>> for Codec {
  type Error = crate::err::Error;

  fn encode(
    &mut self,
    data: &HashMap<String, String>,
    buf: &mut BytesMut
  ) -> Result<(), Error> {
    // Calculate the amount of space required
    let mut sz = 0;
    for (k, v) in data.iter() {
      // key space + whitespace + value space + eol
      sz += k.len() + 1 + v.len() + 1;
    }

    // Terminating empty line
    sz += 1;

    //println!("Writing {} bin data", data.len());
    buf.reserve(sz);

    for (k, v) in data.iter() {
      buf.put(k.as_bytes());
      buf.put_u8(b' ');
      buf.put(v.as_bytes());
      buf.put_u8(b'\n');
    }
    buf.put_u8(b'\n');

    Ok(())
  }
}


impl Encoder<&KVLines> for Codec {
  type Error = crate::err::Error;

  fn encode(
    &mut self,
    kvlines: &KVLines,
    buf: &mut BytesMut
  ) -> Result<(), Error> {
    kvlines.encoder_write(buf)?;
    Ok(())
  }
}


impl Encoder<Bytes> for Codec {
  type Error = crate::err::Error;

  fn encode(
    &mut self,
    data: Bytes,
    buf: &mut BytesMut
  ) -> Result<(), crate::err::Error> {
    buf.reserve(data.len());
    buf.put(data);
    Ok(())
  }
}


impl Encoder<&[u8]> for Codec {
  type Error = crate::err::Error;

  fn encode(
    &mut self,
    data: &[u8],
    buf: &mut BytesMut
  ) -> Result<(), crate::err::Error> {
    buf.reserve(data.len());
    buf.put(data);
    Ok(())
  }
}

// vim: set ft=rust et sw=2 ts=2 sts=2 cinoptions=2 tw=79 :