1use super::EncodeError;
15use bytes::BytesMut;
16
17pub struct PgEncoder;
20
21impl PgEncoder {
22 pub const FORMAT_TEXT: i16 = 0;
24 pub const FORMAT_BINARY: i16 = 1;
26
27 #[inline(always)]
28 fn validate_format_code(format: i16) -> Result<(), EncodeError> {
29 match format {
30 Self::FORMAT_TEXT | Self::FORMAT_BINARY => Ok(()),
31 other => Err(EncodeError::InvalidAst(format!(
32 "invalid PostgreSQL format code {other}; expected 0 text or 1 binary"
33 ))),
34 }
35 }
36
37 #[inline(always)]
38 fn validate_format_codes(param_format: i16, result_format: i16) -> Result<(), EncodeError> {
39 Self::validate_format_code(param_format)?;
40 Self::validate_format_code(result_format)?;
41 Ok(())
42 }
43
44 #[inline(always)]
45 fn param_format_wire_len(param_format: i16) -> usize {
46 if param_format == Self::FORMAT_TEXT {
47 2 } else {
49 4 }
51 }
52
53 #[inline(always)]
54 fn encode_param_formats_vec(content: &mut Vec<u8>, param_format: i16) {
55 if param_format == Self::FORMAT_TEXT {
56 content.extend_from_slice(&0i16.to_be_bytes());
57 } else {
58 content.extend_from_slice(&1i16.to_be_bytes());
59 content.extend_from_slice(¶m_format.to_be_bytes());
60 }
61 }
62
63 #[inline(always)]
64 fn encode_param_formats_bytesmut(buf: &mut BytesMut, param_format: i16) {
65 if param_format == Self::FORMAT_TEXT {
66 buf.extend_from_slice(&0i16.to_be_bytes());
67 } else {
68 buf.extend_from_slice(&1i16.to_be_bytes());
69 buf.extend_from_slice(¶m_format.to_be_bytes());
70 }
71 }
72
73 #[inline(always)]
74 fn result_format_wire_len(result_format: i16) -> usize {
75 if result_format == Self::FORMAT_TEXT {
76 2 } else {
78 4 }
80 }
81
82 #[inline(always)]
83 fn params_wire_len(params: &[Option<Vec<u8>>]) -> Result<usize, EncodeError> {
84 params.iter().try_fold(0usize, |acc, p| {
85 let field_size = 4usize
86 .checked_add(p.as_ref().map_or(0usize, |v| v.len()))
87 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
88 acc.checked_add(field_size)
89 .ok_or(EncodeError::MessageTooLarge(usize::MAX))
90 })
91 }
92
93 #[inline(always)]
94 fn encode_result_formats_vec(content: &mut Vec<u8>, result_format: i16) {
95 if result_format == Self::FORMAT_TEXT {
96 content.extend_from_slice(&0i16.to_be_bytes());
97 } else {
98 content.extend_from_slice(&1i16.to_be_bytes());
99 content.extend_from_slice(&result_format.to_be_bytes());
100 }
101 }
102
103 #[inline(always)]
104 fn encode_result_formats_bytesmut(buf: &mut BytesMut, result_format: i16) {
105 if result_format == Self::FORMAT_TEXT {
106 buf.extend_from_slice(&0i16.to_be_bytes());
107 } else {
108 buf.extend_from_slice(&1i16.to_be_bytes());
109 buf.extend_from_slice(&result_format.to_be_bytes());
110 }
111 }
112
113 #[inline(always)]
114 fn content_len_to_wire_len(content_len: usize) -> Result<i32, EncodeError> {
115 let total = content_len
116 .checked_add(4)
117 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
118 i32::try_from(total).map_err(|_| EncodeError::MessageTooLarge(total))
119 }
120
121 #[inline(always)]
122 fn usize_to_i16(n: usize) -> Result<i16, EncodeError> {
123 i16::try_from(n).map_err(|_| EncodeError::TooManyParameters(n))
124 }
125
126 #[inline(always)]
127 fn usize_to_i32(n: usize) -> Result<i32, EncodeError> {
128 i32::try_from(n).map_err(|_| EncodeError::MessageTooLarge(n))
129 }
130
131 #[inline(always)]
132 fn has_nul(s: &str) -> bool {
133 s.as_bytes().contains(&0)
134 }
135
136 #[inline]
138 pub fn bind_wire_len_with_formats(
139 statement: &str,
140 params: &[Option<Vec<u8>>],
141 param_format: i16,
142 result_format: i16,
143 ) -> Result<usize, EncodeError> {
144 Self::validate_format_codes(param_format, result_format)?;
145 if Self::has_nul(statement) {
146 return Err(EncodeError::NullByte);
147 }
148 if params.len() > i16::MAX as usize {
149 return Err(EncodeError::TooManyParameters(params.len()));
150 }
151
152 let params_size = Self::params_wire_len(params)?;
153 let param_formats_size = Self::param_format_wire_len(param_format);
154 let result_formats_size = Self::result_format_wire_len(result_format);
155 let content_len = 1usize
156 .checked_add(statement.len())
157 .and_then(|v| v.checked_add(1))
158 .and_then(|v| v.checked_add(2))
159 .and_then(|v| v.checked_add(param_formats_size))
160 .and_then(|v| v.checked_add(params_size))
161 .and_then(|v| v.checked_add(result_formats_size))
162 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
163 let wire_len = Self::content_len_to_wire_len(content_len)? as usize;
164 1usize
165 .checked_add(wire_len)
166 .ok_or(EncodeError::MessageTooLarge(usize::MAX))
167 }
168
169 #[inline]
171 pub fn bind_execute_wire_len_with_formats(
172 statement: &str,
173 params: &[Option<Vec<u8>>],
174 param_format: i16,
175 result_format: i16,
176 ) -> Result<usize, EncodeError> {
177 Self::bind_wire_len_with_formats(statement, params, param_format, result_format)?
178 .checked_add(10)
179 .ok_or(EncodeError::MessageTooLarge(usize::MAX))
180 }
181
182 #[inline]
184 pub fn bind_execute_sync_wire_len_with_formats(
185 statement: &str,
186 params: &[Option<Vec<u8>>],
187 param_format: i16,
188 result_format: i16,
189 ) -> Result<usize, EncodeError> {
190 Self::bind_execute_wire_len_with_formats(statement, params, param_format, result_format)?
191 .checked_add(5)
192 .ok_or(EncodeError::MessageTooLarge(usize::MAX))
193 }
194
195 pub fn try_encode_query_string(sql: &str) -> Result<BytesMut, EncodeError> {
197 if Self::has_nul(sql) {
198 return Err(EncodeError::NullByte);
199 }
200
201 let mut buf = BytesMut::new();
202 let content_len = sql.len() + 1; let total_len = Self::content_len_to_wire_len(content_len)?;
204
205 buf.extend_from_slice(b"Q");
206 buf.extend_from_slice(&total_len.to_be_bytes());
207 buf.extend_from_slice(sql.as_bytes());
208 buf.extend_from_slice(&[0]);
209 Ok(buf)
210 }
211
212 pub fn encode_terminate() -> BytesMut {
214 let mut buf = BytesMut::new();
215 buf.extend_from_slice(&[b'X', 0, 0, 0, 4]);
216 buf
217 }
218
219 pub fn encode_sync() -> BytesMut {
221 let mut buf = BytesMut::new();
222 buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
223 buf
224 }
225
226 pub fn try_encode_parse(
230 name: &str,
231 sql: &str,
232 param_types: &[u32],
233 ) -> Result<BytesMut, EncodeError> {
234 if Self::has_nul(name) || Self::has_nul(sql) {
235 return Err(EncodeError::NullByte);
236 }
237 if param_types.len() > i16::MAX as usize {
238 return Err(EncodeError::TooManyParameters(param_types.len()));
239 }
240
241 let mut buf = BytesMut::new();
242 buf.extend_from_slice(b"P");
243
244 let mut content = Vec::new();
245 content.extend_from_slice(name.as_bytes());
246 content.push(0);
247 content.extend_from_slice(sql.as_bytes());
248 content.push(0);
249 let param_count = Self::usize_to_i16(param_types.len())?;
250 content.extend_from_slice(¶m_count.to_be_bytes());
251 for &oid in param_types {
252 content.extend_from_slice(&oid.to_be_bytes());
253 }
254
255 let len = Self::content_len_to_wire_len(content.len())?;
256 buf.extend_from_slice(&len.to_be_bytes());
257 buf.extend_from_slice(&content);
258 Ok(buf)
259 }
260
261 pub fn try_encode_parse_to(
263 buf: &mut BytesMut,
264 name: &str,
265 sql: &str,
266 param_types: &[u32],
267 ) -> Result<(), EncodeError> {
268 if Self::has_nul(name) || Self::has_nul(sql) {
269 return Err(EncodeError::NullByte);
270 }
271 if param_types.len() > i16::MAX as usize {
272 return Err(EncodeError::TooManyParameters(param_types.len()));
273 }
274
275 let content_len = name
276 .len()
277 .checked_add(1)
278 .and_then(|v| v.checked_add(sql.len()))
279 .and_then(|v| v.checked_add(1))
280 .and_then(|v| v.checked_add(2))
281 .and_then(|v| v.checked_add(param_types.len().checked_mul(4)?))
282 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
283 let wire_len = Self::content_len_to_wire_len(content_len)?;
284
285 buf.reserve(1 + 4 + content_len);
286 buf.extend_from_slice(b"P");
287 buf.extend_from_slice(&wire_len.to_be_bytes());
288 buf.extend_from_slice(name.as_bytes());
289 buf.extend_from_slice(&[0]);
290 buf.extend_from_slice(sql.as_bytes());
291 buf.extend_from_slice(&[0]);
292 let param_count = Self::usize_to_i16(param_types.len())?;
293 buf.extend_from_slice(¶m_count.to_be_bytes());
294 for &oid in param_types {
295 buf.extend_from_slice(&oid.to_be_bytes());
296 }
297 Ok(())
298 }
299
300 pub fn encode_bind(
317 portal: &str,
318 statement: &str,
319 params: &[Option<Vec<u8>>],
320 ) -> Result<BytesMut, EncodeError> {
321 Self::encode_bind_with_result_format(portal, statement, params, Self::FORMAT_TEXT)
322 }
323
324 pub fn encode_bind_with_result_format(
330 portal: &str,
331 statement: &str,
332 params: &[Option<Vec<u8>>],
333 result_format: i16,
334 ) -> Result<BytesMut, EncodeError> {
335 Self::encode_bind_with_formats(portal, statement, params, Self::FORMAT_TEXT, result_format)
336 }
337
338 pub fn encode_bind_with_formats(
347 portal: &str,
348 statement: &str,
349 params: &[Option<Vec<u8>>],
350 param_format: i16,
351 result_format: i16,
352 ) -> Result<BytesMut, EncodeError> {
353 Self::validate_format_codes(param_format, result_format)?;
354 if Self::has_nul(portal) || Self::has_nul(statement) {
355 return Err(EncodeError::NullByte);
356 }
357 if params.len() > i16::MAX as usize {
358 return Err(EncodeError::TooManyParameters(params.len()));
359 }
360
361 let mut buf = BytesMut::new();
362
363 buf.extend_from_slice(b"B");
365
366 let mut content = Vec::new();
367
368 content.extend_from_slice(portal.as_bytes());
370 content.push(0);
371
372 content.extend_from_slice(statement.as_bytes());
374 content.push(0);
375
376 Self::encode_param_formats_vec(&mut content, param_format);
378
379 let param_count = Self::usize_to_i16(params.len())?;
381 content.extend_from_slice(¶m_count.to_be_bytes());
382
383 for param in params {
385 match param {
386 None => {
387 content.extend_from_slice(&(-1i32).to_be_bytes());
389 }
390 Some(data) => {
391 let data_len = Self::usize_to_i32(data.len())?;
392 content.extend_from_slice(&data_len.to_be_bytes());
393 content.extend_from_slice(data);
394 }
395 }
396 }
397
398 Self::encode_result_formats_vec(&mut content, result_format);
400
401 let len = Self::content_len_to_wire_len(content.len())?;
403 buf.extend_from_slice(&len.to_be_bytes());
404 buf.extend_from_slice(&content);
405
406 Ok(buf)
407 }
408
409 pub fn try_encode_execute(portal: &str, max_rows: i32) -> Result<BytesMut, EncodeError> {
411 if Self::has_nul(portal) {
412 return Err(EncodeError::NullByte);
413 }
414 if max_rows < 0 {
415 return Err(EncodeError::InvalidMaxRows(max_rows));
416 }
417
418 let mut buf = BytesMut::new();
419 buf.extend_from_slice(b"E");
420
421 let mut content = Vec::new();
422 content.extend_from_slice(portal.as_bytes());
423 content.push(0);
424 content.extend_from_slice(&max_rows.to_be_bytes());
425
426 let len = Self::content_len_to_wire_len(content.len())?;
427 buf.extend_from_slice(&len.to_be_bytes());
428 buf.extend_from_slice(&content);
429 Ok(buf)
430 }
431
432 pub fn try_encode_describe(is_portal: bool, name: &str) -> Result<BytesMut, EncodeError> {
434 if Self::has_nul(name) {
435 return Err(EncodeError::NullByte);
436 }
437
438 let mut buf = BytesMut::new();
439 buf.extend_from_slice(b"D");
440
441 let mut content = Vec::new();
442 content.push(if is_portal { b'P' } else { b'S' });
443 content.extend_from_slice(name.as_bytes());
444 content.push(0);
445
446 let len = Self::content_len_to_wire_len(content.len())?;
447 buf.extend_from_slice(&len.to_be_bytes());
448 buf.extend_from_slice(&content);
449 Ok(buf)
450 }
451
452 pub fn encode_extended_query(
456 sql: &str,
457 params: &[Option<Vec<u8>>],
458 ) -> Result<BytesMut, EncodeError> {
459 Self::encode_extended_query_with_result_format(sql, params, Self::FORMAT_TEXT)
460 }
461
462 pub fn encode_extended_query_with_result_format(
466 sql: &str,
467 params: &[Option<Vec<u8>>],
468 result_format: i16,
469 ) -> Result<BytesMut, EncodeError> {
470 Self::encode_extended_query_with_formats(sql, params, Self::FORMAT_TEXT, result_format)
471 }
472
473 pub fn encode_extended_query_with_formats(
478 sql: &str,
479 params: &[Option<Vec<u8>>],
480 param_format: i16,
481 result_format: i16,
482 ) -> Result<BytesMut, EncodeError> {
483 Self::validate_format_codes(param_format, result_format)?;
484 if Self::has_nul(sql) {
485 return Err(EncodeError::NullByte);
486 }
487 if params.len() > i16::MAX as usize {
488 return Err(EncodeError::TooManyParameters(params.len()));
489 }
490
491 let params_size = params.iter().try_fold(0usize, |acc, p| {
496 let field_size = 4usize
497 .checked_add(p.as_ref().map_or(0usize, |v| v.len()))
498 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
499 acc.checked_add(field_size)
500 .ok_or(EncodeError::MessageTooLarge(usize::MAX))
501 })?;
502 let param_formats_size = Self::param_format_wire_len(param_format);
503 let result_formats_size = Self::result_format_wire_len(result_format);
504 let total_size = 9usize
505 .checked_add(sql.len())
506 .and_then(|v| v.checked_add(9))
507 .and_then(|v| v.checked_add(params_size))
508 .and_then(|v| v.checked_add(param_formats_size))
509 .and_then(|v| v.checked_add(result_formats_size))
510 .and_then(|v| v.checked_add(10))
511 .and_then(|v| v.checked_add(5))
512 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
513
514 let mut buf = BytesMut::with_capacity(total_size);
515
516 buf.extend_from_slice(b"P");
518 let parse_content_len = 1usize
519 .checked_add(sql.len())
520 .and_then(|v| v.checked_add(1))
521 .and_then(|v| v.checked_add(2))
522 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
523 let parse_len = Self::content_len_to_wire_len(parse_content_len)?;
524 buf.extend_from_slice(&parse_len.to_be_bytes());
525 buf.extend_from_slice(&[0]); buf.extend_from_slice(sql.as_bytes());
527 buf.extend_from_slice(&[0]); buf.extend_from_slice(&0i16.to_be_bytes()); buf.extend_from_slice(b"B");
532 let bind_content_len = 1usize
533 .checked_add(1)
534 .and_then(|v| v.checked_add(2))
535 .and_then(|v| v.checked_add(param_formats_size))
536 .and_then(|v| v.checked_add(params_size))
537 .and_then(|v| v.checked_add(result_formats_size))
538 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
539 let bind_len = Self::content_len_to_wire_len(bind_content_len)?;
540 buf.extend_from_slice(&bind_len.to_be_bytes());
541 buf.extend_from_slice(&[0]); buf.extend_from_slice(&[0]); Self::encode_param_formats_bytesmut(&mut buf, param_format);
544 let param_count = Self::usize_to_i16(params.len())?;
545 buf.extend_from_slice(¶m_count.to_be_bytes());
546 for param in params {
547 match param {
548 None => buf.extend_from_slice(&(-1i32).to_be_bytes()),
549 Some(data) => {
550 let data_len = Self::usize_to_i32(data.len())?;
551 buf.extend_from_slice(&data_len.to_be_bytes());
552 buf.extend_from_slice(data);
553 }
554 }
555 }
556 Self::encode_result_formats_bytesmut(&mut buf, result_format);
557
558 buf.extend_from_slice(b"E");
560 buf.extend_from_slice(&9i32.to_be_bytes()); buf.extend_from_slice(&[0]); buf.extend_from_slice(&0i32.to_be_bytes()); buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
566
567 Ok(buf)
568 }
569
570 pub fn try_encode_copy_fail(reason: &str) -> Result<BytesMut, EncodeError> {
572 if Self::has_nul(reason) {
573 return Err(EncodeError::NullByte);
574 }
575
576 let mut buf = BytesMut::new();
577 buf.extend_from_slice(b"f");
578 let content_len = reason.len() + 1; let len = Self::content_len_to_wire_len(content_len)?;
580 buf.extend_from_slice(&len.to_be_bytes());
581 buf.extend_from_slice(reason.as_bytes());
582 buf.extend_from_slice(&[0]);
583 Ok(buf)
584 }
585
586 pub fn try_encode_close(is_portal: bool, name: &str) -> Result<BytesMut, EncodeError> {
588 if Self::has_nul(name) {
589 return Err(EncodeError::NullByte);
590 }
591
592 let mut buf = BytesMut::new();
593 buf.extend_from_slice(b"C");
594 let content_len = 1 + name.len() + 1; let len = Self::content_len_to_wire_len(content_len)?;
596 buf.extend_from_slice(&len.to_be_bytes());
597 buf.extend_from_slice(&[if is_portal { b'P' } else { b'S' }]);
598 buf.extend_from_slice(name.as_bytes());
599 buf.extend_from_slice(&[0]);
600 Ok(buf)
601 }
602}
603
604use bytes::BufMut;
613
614pub enum Param<'a> {
617 Null,
619 Bytes(&'a [u8]),
621}
622
623impl PgEncoder {
624 #[inline(always)]
627 fn put_i32_be(buf: &mut BytesMut, v: i32) {
628 buf.put_i32(v);
629 }
630
631 #[inline(always)]
632 fn put_i16_be(buf: &mut BytesMut, v: i16) {
633 buf.put_i16(v);
634 }
635
636 #[inline]
641 pub fn encode_bind_ultra<'a>(
642 buf: &mut BytesMut,
643 statement: &str,
644 params: &[Param<'a>],
645 ) -> Result<(), EncodeError> {
646 Self::encode_bind_ultra_with_result_format(buf, statement, params, Self::FORMAT_TEXT)
647 }
648
649 #[inline]
651 pub fn encode_bind_ultra_with_result_format<'a>(
652 buf: &mut BytesMut,
653 statement: &str,
654 params: &[Param<'a>],
655 result_format: i16,
656 ) -> Result<(), EncodeError> {
657 Self::encode_bind_ultra_with_formats(
658 buf,
659 statement,
660 params,
661 Self::FORMAT_TEXT,
662 result_format,
663 )
664 }
665
666 #[inline]
668 pub fn encode_bind_ultra_with_formats<'a>(
669 buf: &mut BytesMut,
670 statement: &str,
671 params: &[Param<'a>],
672 param_format: i16,
673 result_format: i16,
674 ) -> Result<(), EncodeError> {
675 Self::validate_format_codes(param_format, result_format)?;
676 if Self::has_nul(statement) {
677 return Err(EncodeError::NullByte);
678 }
679 if params.len() > i16::MAX as usize {
680 return Err(EncodeError::TooManyParameters(params.len()));
681 }
682
683 let params_size = params.iter().try_fold(0usize, |acc, p| {
685 let field_size = match p {
686 Param::Null => 4usize,
687 Param::Bytes(b) => 4usize
688 .checked_add(b.len())
689 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?,
690 };
691 acc.checked_add(field_size)
692 .ok_or(EncodeError::MessageTooLarge(usize::MAX))
693 })?;
694 let param_formats_size = Self::param_format_wire_len(param_format);
695 let result_formats_size = Self::result_format_wire_len(result_format);
696 let content_len = 1usize
697 .checked_add(statement.len())
698 .and_then(|v| v.checked_add(1))
699 .and_then(|v| v.checked_add(2))
700 .and_then(|v| v.checked_add(param_formats_size))
701 .and_then(|v| v.checked_add(params_size))
702 .and_then(|v| v.checked_add(result_formats_size))
703 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
704 let wire_len = Self::content_len_to_wire_len(content_len)?;
705
706 buf.reserve(1 + 4 + content_len);
708
709 buf.put_u8(b'B');
711
712 Self::put_i32_be(buf, wire_len);
714
715 buf.put_u8(0);
717
718 buf.extend_from_slice(statement.as_bytes());
720 buf.put_u8(0);
721
722 Self::encode_param_formats_bytesmut(buf, param_format);
724
725 let param_count = Self::usize_to_i16(params.len())?;
727 Self::put_i16_be(buf, param_count);
728
729 for param in params {
731 match param {
732 Param::Null => Self::put_i32_be(buf, -1),
733 Param::Bytes(data) => {
734 let data_len = Self::usize_to_i32(data.len())?;
735 Self::put_i32_be(buf, data_len);
736 buf.extend_from_slice(data);
737 }
738 }
739 }
740
741 Self::encode_result_formats_bytesmut(buf, result_format);
743 Ok(())
744 }
745
746 #[inline(always)]
748 pub fn encode_execute_ultra(buf: &mut BytesMut) {
749 buf.extend_from_slice(&[b'E', 0, 0, 0, 9, 0, 0, 0, 0, 0]);
752 }
753
754 #[inline(always)]
756 pub fn encode_sync_ultra(buf: &mut BytesMut) {
757 buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
758 }
759
760 #[inline]
763 pub fn encode_bind_to(
764 buf: &mut BytesMut,
765 statement: &str,
766 params: &[Option<Vec<u8>>],
767 ) -> Result<(), EncodeError> {
768 Self::encode_bind_to_with_result_format(buf, statement, params, Self::FORMAT_TEXT)
769 }
770
771 #[inline]
773 pub fn encode_bind_to_with_result_format(
774 buf: &mut BytesMut,
775 statement: &str,
776 params: &[Option<Vec<u8>>],
777 result_format: i16,
778 ) -> Result<(), EncodeError> {
779 Self::encode_bind_to_with_formats(buf, statement, params, Self::FORMAT_TEXT, result_format)
780 }
781
782 #[inline]
784 pub fn encode_bind_to_with_formats(
785 buf: &mut BytesMut,
786 statement: &str,
787 params: &[Option<Vec<u8>>],
788 param_format: i16,
789 result_format: i16,
790 ) -> Result<(), EncodeError> {
791 Self::validate_format_codes(param_format, result_format)?;
792 if Self::has_nul(statement) {
793 return Err(EncodeError::NullByte);
794 }
795 if params.len() > i16::MAX as usize {
796 return Err(EncodeError::TooManyParameters(params.len()));
797 }
798
799 let params_size = Self::params_wire_len(params)?;
803 let param_formats_size = Self::param_format_wire_len(param_format);
804 let result_formats_size = Self::result_format_wire_len(result_format);
805 let content_len = 1usize
806 .checked_add(statement.len())
807 .and_then(|v| v.checked_add(1))
808 .and_then(|v| v.checked_add(2))
809 .and_then(|v| v.checked_add(param_formats_size))
810 .and_then(|v| v.checked_add(params_size))
811 .and_then(|v| v.checked_add(result_formats_size))
812 .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
813 let wire_len = Self::content_len_to_wire_len(content_len)?;
814
815 buf.reserve(1 + 4 + content_len);
816
817 buf.put_u8(b'B');
819
820 Self::put_i32_be(buf, wire_len);
822
823 buf.put_u8(0);
825
826 buf.extend_from_slice(statement.as_bytes());
828 buf.put_u8(0);
829
830 Self::encode_param_formats_bytesmut(buf, param_format);
832
833 let param_count = Self::usize_to_i16(params.len())?;
835 Self::put_i16_be(buf, param_count);
836
837 for param in params {
839 match param {
840 None => Self::put_i32_be(buf, -1),
841 Some(data) => {
842 let data_len = Self::usize_to_i32(data.len())?;
843 Self::put_i32_be(buf, data_len);
844 buf.extend_from_slice(data);
845 }
846 }
847 }
848
849 Self::encode_result_formats_bytesmut(buf, result_format);
851 Ok(())
852 }
853
854 #[inline]
856 pub fn encode_execute_to(buf: &mut BytesMut) {
857 buf.extend_from_slice(&[b'E', 0, 0, 0, 9, 0, 0, 0, 0, 0]);
859 }
860
861 #[inline]
863 pub fn encode_sync_to(buf: &mut BytesMut) {
864 buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
865 }
866}
867
868#[cfg(test)]
869mod tests {
870 use super::*;
871
872 #[test]
874 fn test_encode_query_string() {
875 let sql = "SELECT 1";
876 let bytes = PgEncoder::try_encode_query_string(sql).unwrap();
877
878 assert_eq!(bytes[0], b'Q');
880
881 let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
883 assert_eq!(len, 13);
884
885 assert_eq!(&bytes[5..13], b"SELECT 1");
887
888 assert_eq!(bytes[13], 0);
890 }
891
892 #[test]
893 fn test_encode_terminate() {
894 let bytes = PgEncoder::encode_terminate();
895 assert_eq!(bytes.as_ref(), &[b'X', 0, 0, 0, 4]);
896 }
897
898 #[test]
899 fn test_encode_sync() {
900 let bytes = PgEncoder::encode_sync();
901 assert_eq!(bytes.as_ref(), &[b'S', 0, 0, 0, 4]);
902 }
903
904 #[test]
905 fn test_encode_parse() {
906 let bytes = PgEncoder::try_encode_parse("", "SELECT $1", &[]).unwrap();
907
908 assert_eq!(bytes[0], b'P');
910
911 let content = String::from_utf8_lossy(&bytes[5..]);
913 assert!(content.contains("SELECT $1"));
914 }
915
916 #[test]
917 fn test_encode_parse_to_matches_allocate_variant() {
918 let expected = PgEncoder::try_encode_parse("stmt", "SELECT $1::int4", &[23]).unwrap();
919
920 let mut buf = BytesMut::from(&b"prefix"[..]);
921 PgEncoder::try_encode_parse_to(&mut buf, "stmt", "SELECT $1::int4", &[23]).unwrap();
922
923 assert_eq!(&buf[6..], expected.as_ref());
924 }
925
926 #[test]
927 fn test_encode_bind() {
928 let params = vec![
929 Some(b"42".to_vec()),
930 None, ];
932 let bytes = PgEncoder::encode_bind("", "", ¶ms).unwrap();
933
934 assert_eq!(bytes[0], b'B');
936
937 let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
939 assert!(len > 4); }
941
942 #[test]
943 fn test_encode_bind_binary_result_format() {
944 let bytes =
945 PgEncoder::encode_bind_with_result_format("", "", &[], PgEncoder::FORMAT_BINARY)
946 .unwrap();
947
948 assert_eq!(&bytes[11..15], &[0, 1, 0, 1]);
951 }
952
953 #[test]
954 fn test_encode_bind_binary_param_and_result_format() {
955 let bytes = PgEncoder::encode_bind_with_formats(
956 "",
957 "",
958 &[],
959 PgEncoder::FORMAT_BINARY,
960 PgEncoder::FORMAT_BINARY,
961 )
962 .unwrap();
963
964 assert_eq!(&bytes[7..11], &[0, 1, 0, 1]);
966 assert_eq!(&bytes[11..13], &[0, 0]);
967 assert_eq!(&bytes[13..17], &[0, 1, 0, 1]);
968 }
969
970 #[test]
971 fn test_encode_bind_rejects_invalid_format_codes() {
972 let err = PgEncoder::encode_bind_with_result_format("", "", &[], 2)
973 .expect_err("invalid result format must fail");
974 assert!(
975 matches!(err, EncodeError::InvalidAst(ref message) if message.contains("format code 2")),
976 "{err}"
977 );
978
979 let err = PgEncoder::encode_bind_with_formats("", "", &[], -1, PgEncoder::FORMAT_TEXT)
980 .expect_err("invalid parameter format must fail");
981 assert!(
982 matches!(err, EncodeError::InvalidAst(ref message) if message.contains("format code -1")),
983 "{err}"
984 );
985 }
986
987 #[test]
988 fn test_encode_execute() {
989 let bytes = PgEncoder::try_encode_execute("", 0).unwrap();
990
991 assert_eq!(bytes[0], b'E');
993
994 let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
996 assert_eq!(len, 9);
997 }
998
999 #[test]
1000 fn test_encode_execute_negative_max_rows_returns_error() {
1001 let err = PgEncoder::try_encode_execute("", -1).expect_err("must reject negative max_rows");
1002 assert_eq!(err, EncodeError::InvalidMaxRows(-1));
1003 }
1004
1005 #[test]
1006 fn test_encode_extended_query() {
1007 let params = vec![Some(b"hello".to_vec())];
1008 let bytes = PgEncoder::encode_extended_query("SELECT $1", ¶ms).unwrap();
1009
1010 assert!(bytes.windows(1).any(|w| w == b"P"));
1012 assert!(bytes.windows(1).any(|w| w == b"B"));
1013 assert!(bytes.windows(1).any(|w| w == b"E"));
1014 assert!(bytes.windows(1).any(|w| w == b"S"));
1015 }
1016
1017 #[test]
1018 fn test_encode_extended_query_binary_result_format() {
1019 let bytes = PgEncoder::encode_extended_query_with_result_format(
1020 "SELECT 1",
1021 &[],
1022 PgEncoder::FORMAT_BINARY,
1023 )
1024 .unwrap();
1025
1026 let parse_len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize;
1027 let bind_start = 1 + parse_len;
1028 assert_eq!(bytes[bind_start], b'B');
1029
1030 let bind_len = i32::from_be_bytes([
1031 bytes[bind_start + 1],
1032 bytes[bind_start + 2],
1033 bytes[bind_start + 3],
1034 bytes[bind_start + 4],
1035 ]);
1036 assert_eq!(bind_len, 14);
1037
1038 let bind_content = &bytes[bind_start + 5..bind_start + 1 + bind_len as usize];
1039 assert_eq!(&bind_content[6..10], &[0, 1, 0, 1]);
1040 }
1041
1042 #[test]
1043 fn test_encode_extended_query_binary_param_and_result_format() {
1044 let bytes = PgEncoder::encode_extended_query_with_formats(
1045 "SELECT 1",
1046 &[],
1047 PgEncoder::FORMAT_BINARY,
1048 PgEncoder::FORMAT_BINARY,
1049 )
1050 .unwrap();
1051
1052 let parse_len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize;
1053 let bind_start = 1 + parse_len;
1054 let bind_len = i32::from_be_bytes([
1055 bytes[bind_start + 1],
1056 bytes[bind_start + 2],
1057 bytes[bind_start + 3],
1058 bytes[bind_start + 4],
1059 ]);
1060 assert_eq!(bind_len, 16);
1061
1062 let bind_content = &bytes[bind_start + 5..bind_start + 1 + bind_len as usize];
1063 assert_eq!(&bind_content[2..6], &[0, 1, 0, 1]);
1064 assert_eq!(&bind_content[6..8], &[0, 0]);
1065 assert_eq!(&bind_content[8..12], &[0, 1, 0, 1]);
1066 }
1067
1068 #[test]
1069 fn test_encode_extended_query_rejects_invalid_format_codes() {
1070 let err = PgEncoder::encode_extended_query_with_result_format("SELECT 1", &[], 2)
1071 .expect_err("invalid result format must fail");
1072 assert!(
1073 matches!(err, EncodeError::InvalidAst(ref message) if message.contains("format code 2")),
1074 "{err}"
1075 );
1076
1077 let err = PgEncoder::encode_extended_query_with_formats(
1078 "SELECT 1",
1079 &[],
1080 PgEncoder::FORMAT_TEXT,
1081 -1,
1082 )
1083 .expect_err("invalid result format must fail");
1084 assert!(
1085 matches!(err, EncodeError::InvalidAst(ref message) if message.contains("format code -1")),
1086 "{err}"
1087 );
1088 }
1089
1090 #[test]
1091 fn test_encode_copy_fail() {
1092 let bytes = PgEncoder::try_encode_copy_fail("bad data").unwrap();
1093 assert_eq!(bytes[0], b'f');
1094 let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
1095 assert_eq!(len as usize, 4 + "bad data".len() + 1);
1096 assert_eq!(&bytes[5..13], b"bad data");
1097 assert_eq!(bytes[13], 0);
1098 }
1099
1100 #[test]
1101 fn test_encode_close_statement() {
1102 let bytes = PgEncoder::try_encode_close(false, "my_stmt").unwrap();
1103 assert_eq!(bytes[0], b'C');
1104 assert_eq!(bytes[5], b'S'); assert_eq!(&bytes[6..13], b"my_stmt");
1106 assert_eq!(bytes[13], 0);
1107 }
1108
1109 #[test]
1110 fn test_encode_close_portal() {
1111 let bytes = PgEncoder::try_encode_close(true, "").unwrap();
1112 assert_eq!(bytes[0], b'C');
1113 assert_eq!(bytes[5], b'P'); assert_eq!(bytes[6], 0); }
1116
1117 #[test]
1118 fn test_encode_parse_too_many_param_types_returns_error() {
1119 let param_types = vec![0u32; (i16::MAX as usize) + 1];
1120 let err =
1121 PgEncoder::try_encode_parse("s", "SELECT 1", ¶m_types).expect_err("must reject");
1122 assert_eq!(err, EncodeError::TooManyParameters(param_types.len()));
1123 }
1124
1125 #[test]
1126 fn test_encode_parse_to_with_nul_rejected() {
1127 let mut buf = BytesMut::new();
1128 let err =
1129 PgEncoder::try_encode_parse_to(&mut buf, "s", "SELECT 1\0", &[]).expect_err("reject");
1130 assert_eq!(err, EncodeError::NullByte);
1131 }
1132
1133 #[test]
1134 fn test_encode_bind_to_binary_result_format() {
1135 let mut buf = BytesMut::new();
1136 PgEncoder::encode_bind_to_with_result_format(&mut buf, "", &[], PgEncoder::FORMAT_BINARY)
1137 .unwrap();
1138
1139 assert_eq!(&buf[11..15], &[0, 1, 0, 1]);
1140 }
1141
1142 #[test]
1143 fn test_encode_bind_to_binary_param_and_result_format() {
1144 let mut buf = BytesMut::new();
1145 PgEncoder::encode_bind_to_with_formats(
1146 &mut buf,
1147 "",
1148 &[],
1149 PgEncoder::FORMAT_BINARY,
1150 PgEncoder::FORMAT_BINARY,
1151 )
1152 .unwrap();
1153
1154 assert_eq!(&buf[7..11], &[0, 1, 0, 1]);
1155 assert_eq!(&buf[11..13], &[0, 0]);
1156 assert_eq!(&buf[13..17], &[0, 1, 0, 1]);
1157 }
1158
1159 #[test]
1160 fn test_encode_bind_to_rejects_invalid_format_codes() {
1161 let mut buf = BytesMut::new();
1162 let err = PgEncoder::encode_bind_to_with_result_format(&mut buf, "", &[], 2)
1163 .expect_err("invalid result format must fail");
1164 assert!(
1165 matches!(err, EncodeError::InvalidAst(ref message) if message.contains("format code 2")),
1166 "{err}"
1167 );
1168 assert!(buf.is_empty());
1169 }
1170
1171 #[test]
1172 fn test_bind_execute_sync_wire_len_matches_encoded_bytes() {
1173 let params = vec![Some(b"abc".to_vec()), None, Some(b"defghi".to_vec())];
1174 let mut buf = BytesMut::new();
1175 PgEncoder::encode_bind_to_with_result_format(
1176 &mut buf,
1177 "stmt",
1178 ¶ms,
1179 PgEncoder::FORMAT_TEXT,
1180 )
1181 .unwrap();
1182 PgEncoder::encode_execute_to(&mut buf);
1183 PgEncoder::encode_sync_to(&mut buf);
1184
1185 let expected = PgEncoder::bind_execute_sync_wire_len_with_formats(
1186 "stmt",
1187 ¶ms,
1188 PgEncoder::FORMAT_TEXT,
1189 PgEncoder::FORMAT_TEXT,
1190 )
1191 .unwrap();
1192 assert_eq!(buf.len(), expected);
1193 }
1194
1195 #[test]
1196 fn test_bind_execute_wire_len_matches_encoded_bytes_binary_formats() {
1197 let params = vec![Some(vec![1, 2, 3, 4]), Some(vec![5, 6])];
1198 let mut buf = BytesMut::new();
1199 PgEncoder::encode_bind_to_with_formats(
1200 &mut buf,
1201 "stmt",
1202 ¶ms,
1203 PgEncoder::FORMAT_BINARY,
1204 PgEncoder::FORMAT_BINARY,
1205 )
1206 .unwrap();
1207 PgEncoder::encode_execute_to(&mut buf);
1208
1209 let expected = PgEncoder::bind_execute_wire_len_with_formats(
1210 "stmt",
1211 ¶ms,
1212 PgEncoder::FORMAT_BINARY,
1213 PgEncoder::FORMAT_BINARY,
1214 )
1215 .unwrap();
1216 assert_eq!(buf.len(), expected);
1217 }
1218
1219 #[test]
1220 fn test_encode_bind_ultra_binary_result_format() {
1221 let mut buf = BytesMut::new();
1222 PgEncoder::encode_bind_ultra_with_result_format(
1223 &mut buf,
1224 "",
1225 &[],
1226 PgEncoder::FORMAT_BINARY,
1227 )
1228 .unwrap();
1229
1230 assert_eq!(&buf[11..15], &[0, 1, 0, 1]);
1231 }
1232
1233 #[test]
1234 fn test_encode_bind_ultra_binary_param_and_result_format() {
1235 let mut buf = BytesMut::new();
1236 PgEncoder::encode_bind_ultra_with_formats(
1237 &mut buf,
1238 "",
1239 &[],
1240 PgEncoder::FORMAT_BINARY,
1241 PgEncoder::FORMAT_BINARY,
1242 )
1243 .unwrap();
1244
1245 assert_eq!(&buf[7..11], &[0, 1, 0, 1]);
1246 assert_eq!(&buf[11..13], &[0, 0]);
1247 assert_eq!(&buf[13..17], &[0, 1, 0, 1]);
1248 }
1249
1250 #[test]
1251 fn test_encode_bind_ultra_rejects_invalid_format_codes() {
1252 let mut buf = BytesMut::new();
1253 let err = PgEncoder::encode_bind_ultra_with_formats(&mut buf, "", &[], 2, 1)
1254 .expect_err("invalid parameter format must fail");
1255 assert!(
1256 matches!(err, EncodeError::InvalidAst(ref message) if message.contains("format code 2")),
1257 "{err}"
1258 );
1259 assert!(buf.is_empty());
1260 }
1261
1262 #[test]
1263 fn test_encode_query_string_with_nul_returns_empty() {
1264 let err =
1265 PgEncoder::try_encode_query_string("select 1\0select 2").expect_err("must reject NUL");
1266 assert_eq!(err, EncodeError::NullByte);
1267 }
1268
1269 #[test]
1270 fn test_encode_parse_with_nul_returns_empty() {
1271 let err = PgEncoder::try_encode_parse("s", "SELECT 1\0", &[]).expect_err("must reject");
1272 assert_eq!(err, EncodeError::NullByte);
1273 }
1274
1275 #[test]
1276 fn test_encode_bind_with_nul_rejected() {
1277 let err = PgEncoder::encode_bind_with_result_format("\0", "", &[], PgEncoder::FORMAT_TEXT)
1278 .expect_err("bind with NUL portal must fail");
1279 assert_eq!(err, EncodeError::NullByte);
1280 }
1281
1282 #[test]
1283 fn test_encode_extended_query_with_nul_rejected() {
1284 let err = PgEncoder::encode_extended_query_with_result_format(
1285 "SELECT 1\0UNION SELECT 2",
1286 &[],
1287 PgEncoder::FORMAT_TEXT,
1288 )
1289 .expect_err("extended query with NUL SQL must fail");
1290 assert_eq!(err, EncodeError::NullByte);
1291 }
1292
1293 #[test]
1294 fn test_encode_copy_fail_with_nul_returns_empty() {
1295 let err = PgEncoder::try_encode_copy_fail("bad\0data").expect_err("must reject");
1296 assert_eq!(err, EncodeError::NullByte);
1297 }
1298}