1 use std::io;
2 use std::io::Write;
3 use std::mem;
4 use std::mem::MaybeUninit;
5 use std::ptr;
6 use std::slice;
7 
8 use crate::misc::maybe_uninit_write;
9 use crate::misc::maybe_uninit_write_slice;
10 use crate::misc::vec_spare_capacity_mut;
11 use crate::varint;
12 use crate::wire_format;
13 use crate::zigzag::encode_zig_zag_32;
14 use crate::zigzag::encode_zig_zag_64;
15 use crate::Message;
16 use crate::ProtobufEnum;
17 use crate::ProtobufError;
18 use crate::ProtobufResult;
19 use crate::UnknownFields;
20 use crate::UnknownValueRef;
21 
22 /// Equal to the default buffer size of `BufWriter`, so when
23 /// `CodedOutputStream` wraps `BufWriter`, it often skips double buffering.
24 const OUTPUT_STREAM_BUFFER_SIZE: usize = 8 * 1024;
25 
26 #[doc(hidden)]
27 pub trait WithCodedOutputStream {
with_coded_output_stream<T, F>(self, cb: F) -> ProtobufResult<T> where F: FnOnce(&mut CodedOutputStream) -> ProtobufResult<T>28     fn with_coded_output_stream<T, F>(self, cb: F) -> ProtobufResult<T>
29     where
30         F: FnOnce(&mut CodedOutputStream) -> ProtobufResult<T>;
31 }
32 
33 impl<'a> WithCodedOutputStream for &'a mut (dyn Write + 'a) {
with_coded_output_stream<T, F>(self, cb: F) -> ProtobufResult<T> where F: FnOnce(&mut CodedOutputStream) -> ProtobufResult<T>,34     fn with_coded_output_stream<T, F>(self, cb: F) -> ProtobufResult<T>
35     where
36         F: FnOnce(&mut CodedOutputStream) -> ProtobufResult<T>,
37     {
38         let mut os = CodedOutputStream::new(self);
39         let r = cb(&mut os)?;
40         os.flush()?;
41         Ok(r)
42     }
43 }
44 
45 impl<'a> WithCodedOutputStream for &'a mut Vec<u8> {
with_coded_output_stream<T, F>(mut self, cb: F) -> ProtobufResult<T> where F: FnOnce(&mut CodedOutputStream) -> ProtobufResult<T>,46     fn with_coded_output_stream<T, F>(mut self, cb: F) -> ProtobufResult<T>
47     where
48         F: FnOnce(&mut CodedOutputStream) -> ProtobufResult<T>,
49     {
50         let mut os = CodedOutputStream::vec(&mut self);
51         let r = cb(&mut os)?;
52         os.flush()?;
53         Ok(r)
54     }
55 }
56 
57 #[doc(hidden)]
with_coded_output_stream_to_bytes<F>(cb: F) -> ProtobufResult<Vec<u8>> where F: FnOnce(&mut CodedOutputStream) -> ProtobufResult<()>,58 pub fn with_coded_output_stream_to_bytes<F>(cb: F) -> ProtobufResult<Vec<u8>>
59 where
60     F: FnOnce(&mut CodedOutputStream) -> ProtobufResult<()>,
61 {
62     let mut v = Vec::new();
63     v.with_coded_output_stream(cb)?;
64     Ok(v)
65 }
66 
67 /// Output buffer/writer for `CodedOutputStream`.
68 enum OutputTarget<'a> {
69     Write(&'a mut dyn Write, Vec<u8>),
70     Vec(&'a mut Vec<u8>),
71     /// The buffer is passed as `&[u8]` to `CodedOutputStream` constructor
72     /// and immediately converted to `buffer` field of `CodedOutputStream`,
73     /// it is not needed to be stored here.
74     /// Lifetime parameter of `CodedOutputStream` guarantees the buffer is valid
75     /// during the lifetime of `CodedOutputStream`.
76     Bytes,
77 }
78 
79 /// Buffered write with handy utilities
80 pub struct CodedOutputStream<'a> {
81     target: OutputTarget<'a>,
82     // Actual buffer is owned by `OutputTarget`,
83     // and here we alias the buffer so access to the buffer is branchless:
84     // access does not require switch by actual target type: `&[], `Vec`, `Write` etc.
85     // We don't access the actual buffer in `OutputTarget` except when
86     // we initialize `buffer` field here.
87     buffer: *mut [MaybeUninit<u8>],
88     // within buffer
89     position: usize,
90 }
91 
92 impl<'a> CodedOutputStream<'a> {
93     /// Construct from given `Write`.
94     ///
95     /// `CodedOutputStream` is buffered even if `Write` is not
new(writer: &'a mut dyn Write) -> CodedOutputStream<'a>96     pub fn new(writer: &'a mut dyn Write) -> CodedOutputStream<'a> {
97         let buffer_len = OUTPUT_STREAM_BUFFER_SIZE;
98 
99         let mut buffer_storage = Vec::with_capacity(buffer_len);
100 
101         // SAFETY: we are not using the `buffer_storage`
102         // except for initializing the `buffer` field.
103         // See `buffer` field documentation.
104         let buffer = vec_spare_capacity_mut(&mut buffer_storage);
105         let buffer: *mut [MaybeUninit<u8>] = buffer;
106 
107         CodedOutputStream {
108             target: OutputTarget::Write(writer, buffer_storage),
109             buffer,
110             position: 0,
111         }
112     }
113 
114     /// `CodedOutputStream` which writes directly to bytes.
115     ///
116     /// Attempt to write more than bytes capacity results in error.
bytes(bytes: &'a mut [u8]) -> CodedOutputStream<'a>117     pub fn bytes(bytes: &'a mut [u8]) -> CodedOutputStream<'a> {
118         // SAFETY: it is safe to cast from &mut [u8] to &mut [MaybeUninit<u8>].
119         let buffer =
120             ptr::slice_from_raw_parts_mut(bytes.as_mut_ptr() as *mut MaybeUninit<u8>, bytes.len());
121         CodedOutputStream {
122             target: OutputTarget::Bytes,
123             buffer,
124             position: 0,
125         }
126     }
127 
128     /// `CodedOutputStream` which writes directly to `Vec<u8>`.
129     ///
130     /// Caller should call `flush` at the end to guarantee vec contains
131     /// all written data.
vec(vec: &'a mut Vec<u8>) -> CodedOutputStream<'a>132     pub fn vec(vec: &'a mut Vec<u8>) -> CodedOutputStream<'a> {
133         let buffer: *mut [MaybeUninit<u8>] = &mut [];
134         CodedOutputStream {
135             target: OutputTarget::Vec(vec),
136             buffer,
137             position: 0,
138         }
139     }
140 
141     /// Check if EOF is reached.
142     ///
143     /// # Panics
144     ///
145     /// If underlying write has no EOF
check_eof(&self)146     pub fn check_eof(&self) {
147         match self.target {
148             OutputTarget::Bytes => {
149                 assert_eq!(self.buffer().len() as u64, self.position as u64);
150             }
151             OutputTarget::Write(..) | OutputTarget::Vec(..) => {
152                 panic!("must not be called with Writer or Vec");
153             }
154         }
155     }
156 
157     #[inline(always)]
buffer(&self) -> &[MaybeUninit<u8>]158     fn buffer(&self) -> &[MaybeUninit<u8>] {
159         // SAFETY: see the `buffer` field documentation about invariants.
160         unsafe { &*(self.buffer as *mut [MaybeUninit<u8>]) }
161     }
162 
163     #[inline(always)]
filled_buffer_impl<'s>(buffer: *mut [MaybeUninit<u8>], position: usize) -> &'s [u8]164     fn filled_buffer_impl<'s>(buffer: *mut [MaybeUninit<u8>], position: usize) -> &'s [u8] {
165         // SAFETY: this function is safe assuming `buffer` and `position`
166         //   are `self.buffer` and `safe.position`:
167         //   * `CodedOutputStream` has invariant that `position <= buffer.len()`.
168         //   * `buffer` is filled up to `position`.
169         unsafe { slice::from_raw_parts_mut(buffer as *mut u8, position) }
170     }
171 
refresh_buffer(&mut self) -> ProtobufResult<()>172     fn refresh_buffer(&mut self) -> ProtobufResult<()> {
173         match self.target {
174             OutputTarget::Write(ref mut write, _) => {
175                 write.write_all(Self::filled_buffer_impl(self.buffer, self.position))?;
176                 self.position = 0;
177             }
178             OutputTarget::Vec(ref mut vec) => unsafe {
179                 let vec_len = vec.len();
180                 assert!(vec_len + self.position <= vec.capacity());
181                 vec.set_len(vec_len + self.position);
182                 vec.reserve(1);
183                 self.buffer = vec_spare_capacity_mut(vec);
184                 self.position = 0;
185             },
186             OutputTarget::Bytes => {
187                 return Err(ProtobufError::IoError(io::Error::new(
188                     io::ErrorKind::Other,
189                     "given slice is too small to serialize the message",
190                 )));
191             }
192         }
193         Ok(())
194     }
195 
196     /// Flush the buffer to underlying write
flush(&mut self) -> ProtobufResult<()>197     pub fn flush(&mut self) -> ProtobufResult<()> {
198         match self.target {
199             OutputTarget::Bytes => Ok(()),
200             OutputTarget::Write(..) | OutputTarget::Vec(..) => {
201                 // TODO: must not reserve additional in Vec
202                 self.refresh_buffer()
203             }
204         }
205     }
206 
207     /// Write a byte
write_raw_byte(&mut self, byte: u8) -> ProtobufResult<()>208     pub fn write_raw_byte(&mut self, byte: u8) -> ProtobufResult<()> {
209         if self.position as usize == self.buffer().len() {
210             self.refresh_buffer()?;
211         }
212         unsafe { maybe_uninit_write(&mut (&mut *self.buffer)[self.position as usize], byte) };
213         self.position += 1;
214         Ok(())
215     }
216 
217     /// Write bytes
write_raw_bytes(&mut self, bytes: &[u8]) -> ProtobufResult<()>218     pub fn write_raw_bytes(&mut self, bytes: &[u8]) -> ProtobufResult<()> {
219         if bytes.len() <= self.buffer().len() - self.position {
220             let bottom = self.position as usize;
221             let top = bottom + (bytes.len() as usize);
222             // SAFETY: see the `buffer` field documentation about invariants.
223             let buffer = unsafe { &mut (&mut *self.buffer)[bottom..top] };
224             maybe_uninit_write_slice(buffer, bytes);
225             self.position += bytes.len();
226             return Ok(());
227         }
228 
229         self.refresh_buffer()?;
230 
231         assert!(self.position == 0);
232 
233         if self.position + bytes.len() < self.buffer().len() {
234             // SAFETY: see the `buffer` field documentation about invariants.
235             let buffer =
236                 unsafe { &mut (&mut *self.buffer)[self.position..self.position + bytes.len()] };
237             maybe_uninit_write_slice(buffer, bytes);
238             self.position += bytes.len();
239             return Ok(());
240         }
241 
242         match self.target {
243             OutputTarget::Bytes => {
244                 unreachable!();
245             }
246             OutputTarget::Write(ref mut write, _) => {
247                 write.write_all(bytes)?;
248             }
249             OutputTarget::Vec(ref mut vec) => {
250                 vec.extend(bytes);
251                 self.buffer = vec_spare_capacity_mut(vec)
252             }
253         }
254         Ok(())
255     }
256 
257     /// Write a tag
write_tag( &mut self, field_number: u32, wire_type: wire_format::WireType, ) -> ProtobufResult<()>258     pub fn write_tag(
259         &mut self,
260         field_number: u32,
261         wire_type: wire_format::WireType,
262     ) -> ProtobufResult<()> {
263         self.write_raw_varint32(wire_format::Tag::make(field_number, wire_type).value())
264     }
265 
266     /// Write varint
write_raw_varint32(&mut self, value: u32) -> ProtobufResult<()>267     pub fn write_raw_varint32(&mut self, value: u32) -> ProtobufResult<()> {
268         if self.buffer().len() - self.position >= 5 {
269             // fast path
270             let len = unsafe {
271                 varint::encode_varint32(value, &mut (&mut *self.buffer)[self.position..])
272             };
273             self.position += len;
274             Ok(())
275         } else {
276             // slow path
277             let buf = &mut [0u8; 5];
278             let len = varint::encode_varint32(value, unsafe {
279                 slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut MaybeUninit<u8>, buf.len())
280             });
281             self.write_raw_bytes(&buf[..len])
282         }
283     }
284 
285     /// Write varint
write_raw_varint64(&mut self, value: u64) -> ProtobufResult<()>286     pub fn write_raw_varint64(&mut self, value: u64) -> ProtobufResult<()> {
287         if self.buffer().len() - self.position >= 10 {
288             // fast path
289             let len = unsafe {
290                 varint::encode_varint64(value, &mut (&mut *self.buffer)[self.position..])
291             };
292             self.position += len;
293             Ok(())
294         } else {
295             // slow path
296             let buf = &mut [0u8; 10];
297             let len = varint::encode_varint64(value, unsafe {
298                 slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut MaybeUninit<u8>, buf.len())
299             });
300             self.write_raw_bytes(&buf[..len])
301         }
302     }
303 
304     /// Write 32-bit integer little endian
write_raw_little_endian32(&mut self, value: u32) -> ProtobufResult<()>305     pub fn write_raw_little_endian32(&mut self, value: u32) -> ProtobufResult<()> {
306         let bytes = unsafe { mem::transmute::<_, [u8; 4]>(value.to_le()) };
307         self.write_raw_bytes(&bytes)
308     }
309 
310     /// Write 64-bit integer little endian
write_raw_little_endian64(&mut self, value: u64) -> ProtobufResult<()>311     pub fn write_raw_little_endian64(&mut self, value: u64) -> ProtobufResult<()> {
312         let bytes = unsafe { mem::transmute::<_, [u8; 8]>(value.to_le()) };
313         self.write_raw_bytes(&bytes)
314     }
315 
316     /// Write `float`
write_float_no_tag(&mut self, value: f32) -> ProtobufResult<()>317     pub fn write_float_no_tag(&mut self, value: f32) -> ProtobufResult<()> {
318         let bits = unsafe { mem::transmute::<f32, u32>(value) };
319         self.write_raw_little_endian32(bits)
320     }
321 
322     /// Write `double`
write_double_no_tag(&mut self, value: f64) -> ProtobufResult<()>323     pub fn write_double_no_tag(&mut self, value: f64) -> ProtobufResult<()> {
324         let bits = unsafe { mem::transmute::<f64, u64>(value) };
325         self.write_raw_little_endian64(bits)
326     }
327 
328     /// Write `float` field
write_float(&mut self, field_number: u32, value: f32) -> ProtobufResult<()>329     pub fn write_float(&mut self, field_number: u32, value: f32) -> ProtobufResult<()> {
330         self.write_tag(field_number, wire_format::WireTypeFixed32)?;
331         self.write_float_no_tag(value)?;
332         Ok(())
333     }
334 
335     /// Write `double` field
write_double(&mut self, field_number: u32, value: f64) -> ProtobufResult<()>336     pub fn write_double(&mut self, field_number: u32, value: f64) -> ProtobufResult<()> {
337         self.write_tag(field_number, wire_format::WireTypeFixed64)?;
338         self.write_double_no_tag(value)?;
339         Ok(())
340     }
341 
342     /// Write varint
write_uint64_no_tag(&mut self, value: u64) -> ProtobufResult<()>343     pub fn write_uint64_no_tag(&mut self, value: u64) -> ProtobufResult<()> {
344         self.write_raw_varint64(value)
345     }
346 
347     /// Write varint
write_uint32_no_tag(&mut self, value: u32) -> ProtobufResult<()>348     pub fn write_uint32_no_tag(&mut self, value: u32) -> ProtobufResult<()> {
349         self.write_raw_varint32(value)
350     }
351 
352     /// Write varint
write_int64_no_tag(&mut self, value: i64) -> ProtobufResult<()>353     pub fn write_int64_no_tag(&mut self, value: i64) -> ProtobufResult<()> {
354         self.write_raw_varint64(value as u64)
355     }
356 
357     /// Write varint
write_int32_no_tag(&mut self, value: i32) -> ProtobufResult<()>358     pub fn write_int32_no_tag(&mut self, value: i32) -> ProtobufResult<()> {
359         self.write_raw_varint64(value as u64)
360     }
361 
362     /// Write zigzag varint
write_sint64_no_tag(&mut self, value: i64) -> ProtobufResult<()>363     pub fn write_sint64_no_tag(&mut self, value: i64) -> ProtobufResult<()> {
364         self.write_uint64_no_tag(encode_zig_zag_64(value))
365     }
366 
367     /// Write zigzag varint
write_sint32_no_tag(&mut self, value: i32) -> ProtobufResult<()>368     pub fn write_sint32_no_tag(&mut self, value: i32) -> ProtobufResult<()> {
369         self.write_uint32_no_tag(encode_zig_zag_32(value))
370     }
371 
372     /// Write `fixed64`
write_fixed64_no_tag(&mut self, value: u64) -> ProtobufResult<()>373     pub fn write_fixed64_no_tag(&mut self, value: u64) -> ProtobufResult<()> {
374         self.write_raw_little_endian64(value)
375     }
376 
377     /// Write `fixed32`
write_fixed32_no_tag(&mut self, value: u32) -> ProtobufResult<()>378     pub fn write_fixed32_no_tag(&mut self, value: u32) -> ProtobufResult<()> {
379         self.write_raw_little_endian32(value)
380     }
381 
382     /// Write `sfixed64`
write_sfixed64_no_tag(&mut self, value: i64) -> ProtobufResult<()>383     pub fn write_sfixed64_no_tag(&mut self, value: i64) -> ProtobufResult<()> {
384         self.write_raw_little_endian64(value as u64)
385     }
386 
387     /// Write `sfixed32`
write_sfixed32_no_tag(&mut self, value: i32) -> ProtobufResult<()>388     pub fn write_sfixed32_no_tag(&mut self, value: i32) -> ProtobufResult<()> {
389         self.write_raw_little_endian32(value as u32)
390     }
391 
392     /// Write `bool`
write_bool_no_tag(&mut self, value: bool) -> ProtobufResult<()>393     pub fn write_bool_no_tag(&mut self, value: bool) -> ProtobufResult<()> {
394         self.write_raw_varint32(if value { 1 } else { 0 })
395     }
396 
397     /// Write `enum`
write_enum_no_tag(&mut self, value: i32) -> ProtobufResult<()>398     pub fn write_enum_no_tag(&mut self, value: i32) -> ProtobufResult<()> {
399         self.write_int32_no_tag(value)
400     }
401 
402     /// Write `enum`
write_enum_obj_no_tag<E>(&mut self, value: E) -> ProtobufResult<()> where E: ProtobufEnum,403     pub fn write_enum_obj_no_tag<E>(&mut self, value: E) -> ProtobufResult<()>
404     where
405         E: ProtobufEnum,
406     {
407         self.write_enum_no_tag(value.value())
408     }
409 
410     /// Write unknown value
write_unknown_no_tag(&mut self, unknown: UnknownValueRef) -> ProtobufResult<()>411     pub fn write_unknown_no_tag(&mut self, unknown: UnknownValueRef) -> ProtobufResult<()> {
412         match unknown {
413             UnknownValueRef::Fixed64(fixed64) => self.write_raw_little_endian64(fixed64),
414             UnknownValueRef::Fixed32(fixed32) => self.write_raw_little_endian32(fixed32),
415             UnknownValueRef::Varint(varint) => self.write_raw_varint64(varint),
416             UnknownValueRef::LengthDelimited(bytes) => self.write_bytes_no_tag(bytes),
417         }
418     }
419 
420     /// Write `uint64` field
write_uint64(&mut self, field_number: u32, value: u64) -> ProtobufResult<()>421     pub fn write_uint64(&mut self, field_number: u32, value: u64) -> ProtobufResult<()> {
422         self.write_tag(field_number, wire_format::WireTypeVarint)?;
423         self.write_uint64_no_tag(value)?;
424         Ok(())
425     }
426 
427     /// Write `uint32` field
write_uint32(&mut self, field_number: u32, value: u32) -> ProtobufResult<()>428     pub fn write_uint32(&mut self, field_number: u32, value: u32) -> ProtobufResult<()> {
429         self.write_tag(field_number, wire_format::WireTypeVarint)?;
430         self.write_uint32_no_tag(value)?;
431         Ok(())
432     }
433 
434     /// Write `int64` field
write_int64(&mut self, field_number: u32, value: i64) -> ProtobufResult<()>435     pub fn write_int64(&mut self, field_number: u32, value: i64) -> ProtobufResult<()> {
436         self.write_tag(field_number, wire_format::WireTypeVarint)?;
437         self.write_int64_no_tag(value)?;
438         Ok(())
439     }
440 
441     /// Write `int32` field
write_int32(&mut self, field_number: u32, value: i32) -> ProtobufResult<()>442     pub fn write_int32(&mut self, field_number: u32, value: i32) -> ProtobufResult<()> {
443         self.write_tag(field_number, wire_format::WireTypeVarint)?;
444         self.write_int32_no_tag(value)?;
445         Ok(())
446     }
447 
448     /// Write `sint64` field
write_sint64(&mut self, field_number: u32, value: i64) -> ProtobufResult<()>449     pub fn write_sint64(&mut self, field_number: u32, value: i64) -> ProtobufResult<()> {
450         self.write_tag(field_number, wire_format::WireTypeVarint)?;
451         self.write_sint64_no_tag(value)?;
452         Ok(())
453     }
454 
455     /// Write `sint32` field
write_sint32(&mut self, field_number: u32, value: i32) -> ProtobufResult<()>456     pub fn write_sint32(&mut self, field_number: u32, value: i32) -> ProtobufResult<()> {
457         self.write_tag(field_number, wire_format::WireTypeVarint)?;
458         self.write_sint32_no_tag(value)?;
459         Ok(())
460     }
461 
462     /// Write `fixed64` field
write_fixed64(&mut self, field_number: u32, value: u64) -> ProtobufResult<()>463     pub fn write_fixed64(&mut self, field_number: u32, value: u64) -> ProtobufResult<()> {
464         self.write_tag(field_number, wire_format::WireTypeFixed64)?;
465         self.write_fixed64_no_tag(value)?;
466         Ok(())
467     }
468 
469     /// Write `fixed32` field
write_fixed32(&mut self, field_number: u32, value: u32) -> ProtobufResult<()>470     pub fn write_fixed32(&mut self, field_number: u32, value: u32) -> ProtobufResult<()> {
471         self.write_tag(field_number, wire_format::WireTypeFixed32)?;
472         self.write_fixed32_no_tag(value)?;
473         Ok(())
474     }
475 
476     /// Write `sfixed64` field
write_sfixed64(&mut self, field_number: u32, value: i64) -> ProtobufResult<()>477     pub fn write_sfixed64(&mut self, field_number: u32, value: i64) -> ProtobufResult<()> {
478         self.write_tag(field_number, wire_format::WireTypeFixed64)?;
479         self.write_sfixed64_no_tag(value)?;
480         Ok(())
481     }
482 
483     /// Write `sfixed32` field
write_sfixed32(&mut self, field_number: u32, value: i32) -> ProtobufResult<()>484     pub fn write_sfixed32(&mut self, field_number: u32, value: i32) -> ProtobufResult<()> {
485         self.write_tag(field_number, wire_format::WireTypeFixed32)?;
486         self.write_sfixed32_no_tag(value)?;
487         Ok(())
488     }
489 
490     /// Write `bool` field
write_bool(&mut self, field_number: u32, value: bool) -> ProtobufResult<()>491     pub fn write_bool(&mut self, field_number: u32, value: bool) -> ProtobufResult<()> {
492         self.write_tag(field_number, wire_format::WireTypeVarint)?;
493         self.write_bool_no_tag(value)?;
494         Ok(())
495     }
496 
497     /// Write `enum` field
write_enum(&mut self, field_number: u32, value: i32) -> ProtobufResult<()>498     pub fn write_enum(&mut self, field_number: u32, value: i32) -> ProtobufResult<()> {
499         self.write_tag(field_number, wire_format::WireTypeVarint)?;
500         self.write_enum_no_tag(value)?;
501         Ok(())
502     }
503 
504     /// Write `enum` field
write_enum_obj<E>(&mut self, field_number: u32, value: E) -> ProtobufResult<()> where E: ProtobufEnum,505     pub fn write_enum_obj<E>(&mut self, field_number: u32, value: E) -> ProtobufResult<()>
506     where
507         E: ProtobufEnum,
508     {
509         self.write_enum(field_number, value.value())
510     }
511 
512     /// Write unknown field
write_unknown( &mut self, field_number: u32, value: UnknownValueRef, ) -> ProtobufResult<()>513     pub fn write_unknown(
514         &mut self,
515         field_number: u32,
516         value: UnknownValueRef,
517     ) -> ProtobufResult<()> {
518         self.write_tag(field_number, value.wire_type())?;
519         self.write_unknown_no_tag(value)?;
520         Ok(())
521     }
522 
523     /// Write unknown fields
write_unknown_fields(&mut self, fields: &UnknownFields) -> ProtobufResult<()>524     pub fn write_unknown_fields(&mut self, fields: &UnknownFields) -> ProtobufResult<()> {
525         for (number, values) in fields {
526             for value in values {
527                 self.write_unknown(number, value)?;
528             }
529         }
530         Ok(())
531     }
532 
533     /// Write bytes
write_bytes_no_tag(&mut self, bytes: &[u8]) -> ProtobufResult<()>534     pub fn write_bytes_no_tag(&mut self, bytes: &[u8]) -> ProtobufResult<()> {
535         self.write_raw_varint32(bytes.len() as u32)?;
536         self.write_raw_bytes(bytes)?;
537         Ok(())
538     }
539 
540     /// Write string
write_string_no_tag(&mut self, s: &str) -> ProtobufResult<()>541     pub fn write_string_no_tag(&mut self, s: &str) -> ProtobufResult<()> {
542         self.write_bytes_no_tag(s.as_bytes())
543     }
544 
545     /// Write message
write_message_no_tag<M: Message>(&mut self, msg: &M) -> ProtobufResult<()>546     pub fn write_message_no_tag<M: Message>(&mut self, msg: &M) -> ProtobufResult<()> {
547         msg.write_length_delimited_to(self)
548     }
549 
550     /// Write `bytes` field
write_bytes(&mut self, field_number: u32, bytes: &[u8]) -> ProtobufResult<()>551     pub fn write_bytes(&mut self, field_number: u32, bytes: &[u8]) -> ProtobufResult<()> {
552         self.write_tag(field_number, wire_format::WireTypeLengthDelimited)?;
553         self.write_bytes_no_tag(bytes)?;
554         Ok(())
555     }
556 
557     /// Write `string` field
write_string(&mut self, field_number: u32, s: &str) -> ProtobufResult<()>558     pub fn write_string(&mut self, field_number: u32, s: &str) -> ProtobufResult<()> {
559         self.write_tag(field_number, wire_format::WireTypeLengthDelimited)?;
560         self.write_string_no_tag(s)?;
561         Ok(())
562     }
563 
564     /// Write `message` field
write_message<M: Message>(&mut self, field_number: u32, msg: &M) -> ProtobufResult<()>565     pub fn write_message<M: Message>(&mut self, field_number: u32, msg: &M) -> ProtobufResult<()> {
566         self.write_tag(field_number, wire_format::WireTypeLengthDelimited)?;
567         self.write_message_no_tag(msg)?;
568         Ok(())
569     }
570 }
571 
572 impl<'a> Write for CodedOutputStream<'a> {
write(&mut self, buf: &[u8]) -> io::Result<usize>573     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
574         self.write_raw_bytes(buf)?;
575         Ok(buf.len())
576     }
577 
flush(&mut self) -> io::Result<()>578     fn flush(&mut self) -> io::Result<()> {
579         CodedOutputStream::flush(self).map_err(Into::into)
580     }
581 }
582 
583 #[cfg(test)]
584 mod test {
585     use std::io::Write;
586     use std::iter;
587 
588     use crate::coded_output_stream::CodedOutputStream;
589     use crate::hex::decode_hex;
590     use crate::hex::encode_hex;
591     use crate::wire_format;
592     use crate::ProtobufResult;
593 
test_write<F>(expected: &str, mut gen: F) where F: FnMut(&mut CodedOutputStream) -> ProtobufResult<()>,594     fn test_write<F>(expected: &str, mut gen: F)
595     where
596         F: FnMut(&mut CodedOutputStream) -> ProtobufResult<()>,
597     {
598         let expected_bytes = decode_hex(expected);
599 
600         // write to Write
601         {
602             let mut v = Vec::new();
603             {
604                 let mut os = CodedOutputStream::new(&mut v as &mut dyn Write);
605                 gen(&mut os).unwrap();
606                 os.flush().unwrap();
607             }
608             assert_eq!(encode_hex(&expected_bytes), encode_hex(&v));
609         }
610 
611         // write to &[u8]
612         {
613             let mut r = Vec::with_capacity(expected_bytes.len());
614             r.resize(expected_bytes.len(), 0);
615             {
616                 let mut os = CodedOutputStream::bytes(&mut r);
617                 gen(&mut os).unwrap();
618                 os.check_eof();
619             }
620             assert_eq!(encode_hex(&expected_bytes), encode_hex(&r));
621         }
622 
623         // write to Vec<u8>
624         {
625             let mut r = Vec::new();
626             r.extend(&[11, 22, 33, 44, 55, 66, 77]);
627             {
628                 let mut os = CodedOutputStream::vec(&mut r);
629                 gen(&mut os).unwrap();
630                 os.flush().unwrap();
631             }
632 
633             r.drain(..7);
634             assert_eq!(encode_hex(&expected_bytes), encode_hex(&r));
635         }
636     }
637 
638     #[test]
test_output_stream_write_raw_byte()639     fn test_output_stream_write_raw_byte() {
640         test_write("a1", |os| os.write_raw_byte(0xa1));
641     }
642 
643     #[test]
test_output_stream_write_tag()644     fn test_output_stream_write_tag() {
645         test_write("08", |os| os.write_tag(1, wire_format::WireTypeVarint));
646     }
647 
648     #[test]
test_output_stream_write_raw_bytes()649     fn test_output_stream_write_raw_bytes() {
650         test_write("00 ab", |os| os.write_raw_bytes(&[0x00, 0xab]));
651 
652         let expected = iter::repeat("01 02 03 04")
653             .take(2048)
654             .collect::<Vec<_>>()
655             .join(" ");
656         test_write(&expected, |os| {
657             for _ in 0..2048 {
658                 os.write_raw_bytes(&[0x01, 0x02, 0x03, 0x04])?;
659             }
660 
661             Ok(())
662         });
663     }
664 
665     #[test]
test_output_stream_write_raw_varint32()666     fn test_output_stream_write_raw_varint32() {
667         test_write("96 01", |os| os.write_raw_varint32(150));
668         test_write("ff ff ff ff 0f", |os| os.write_raw_varint32(0xffffffff));
669     }
670 
671     #[test]
test_output_stream_write_raw_varint64()672     fn test_output_stream_write_raw_varint64() {
673         test_write("96 01", |os| os.write_raw_varint64(150));
674         test_write("ff ff ff ff ff ff ff ff ff 01", |os| {
675             os.write_raw_varint64(0xffffffffffffffff)
676         });
677     }
678 
679     #[test]
test_output_stream_write_int32_no_tag()680     fn test_output_stream_write_int32_no_tag() {
681         test_write("ff ff ff ff ff ff ff ff ff 01", |os| {
682             os.write_int32_no_tag(-1)
683         });
684     }
685 
686     #[test]
test_output_stream_write_int64_no_tag()687     fn test_output_stream_write_int64_no_tag() {
688         test_write("ff ff ff ff ff ff ff ff ff 01", |os| {
689             os.write_int64_no_tag(-1)
690         });
691     }
692 
693     #[test]
test_output_stream_write_raw_little_endian32()694     fn test_output_stream_write_raw_little_endian32() {
695         test_write("f1 e2 d3 c4", |os| os.write_raw_little_endian32(0xc4d3e2f1));
696     }
697 
698     #[test]
test_output_stream_write_float_no_tag()699     fn test_output_stream_write_float_no_tag() {
700         test_write("95 73 13 61", |os| os.write_float_no_tag(17e19));
701     }
702 
703     #[test]
test_output_stream_write_double_no_tag()704     fn test_output_stream_write_double_no_tag() {
705         test_write("40 d5 ab 68 b3 07 3d 46", |os| {
706             os.write_double_no_tag(23e29)
707         });
708     }
709 
710     #[test]
test_output_stream_write_raw_little_endian64()711     fn test_output_stream_write_raw_little_endian64() {
712         test_write("f1 e2 d3 c4 b5 a6 07 f8", |os| {
713             os.write_raw_little_endian64(0xf807a6b5c4d3e2f1)
714         });
715     }
716 
717     #[test]
test_output_stream_io_write()718     fn test_output_stream_io_write() {
719         let expected = [0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77];
720 
721         // write to Write
722         {
723             let mut v = Vec::new();
724             {
725                 let mut os = CodedOutputStream::new(&mut v as &mut dyn Write);
726                 Write::write(&mut os, &expected).expect("io::Write::write");
727                 Write::flush(&mut os).expect("io::Write::flush");
728             }
729             assert_eq!(expected, *v);
730         }
731 
732         // write to &[u8]
733         {
734             let mut v = Vec::with_capacity(expected.len());
735             v.resize(expected.len(), 0);
736             {
737                 let mut os = CodedOutputStream::bytes(&mut v);
738                 Write::write(&mut os, &expected).expect("io::Write::write");
739                 Write::flush(&mut os).expect("io::Write::flush");
740                 os.check_eof();
741             }
742             assert_eq!(expected, *v);
743         }
744 
745         // write to Vec<u8>
746         {
747             let mut v = Vec::new();
748             {
749                 let mut os = CodedOutputStream::vec(&mut v);
750                 Write::write(&mut os, &expected).expect("io::Write::write");
751                 Write::flush(&mut os).expect("io::Write::flush");
752             }
753             assert_eq!(expected, *v);
754         }
755     }
756 }
757