1use std::collections::HashMap;
9
10use prost::encoding::{WireType, decode_key, decode_varint, encode_key, encode_varint};
11
12use super::wire::{decode_len, skip_wire_value};
13
14const FIELD_PRESENCE_EXPLICIT: i32 = 1;
16const FIELD_PRESENCE_LEGACY_REQUIRED: i32 = 3;
17
18const MESSAGE_ENCODING_DELIMITED: i32 = 2;
20
21const LABEL_OPTIONAL: i32 = 1;
23const LABEL_REQUIRED: i32 = 2;
24const LABEL_REPEATED: i32 = 3;
25
26const TYPE_MESSAGE: i32 = 11;
28const TYPE_GROUP: i32 = 10;
29
30mod file_tags {
32 pub const MESSAGE_TYPE: u32 = 4;
33 pub const EXTENSION: u32 = 7;
34 pub const OPTIONS: u32 = 8;
35 pub const SYNTAX: u32 = 12;
36 pub const EDITION: u32 = 14;
37}
38
39mod message_tags {
41 pub const FIELD: u32 = 2;
42 pub const NESTED_TYPE: u32 = 3;
43 pub const EXTENSION: u32 = 6;
44 pub const ONEOF_DECL: u32 = 8;
45}
46
47mod field_tags {
49 pub const NAME: u32 = 1;
50 pub const LABEL: u32 = 4;
51 pub const TYPE: u32 = 5;
52 pub const OPTIONS: u32 = 8;
53 pub const ONEOF_INDEX: u32 = 9;
54 pub const PROTO3_OPTIONAL: u32 = 17;
55}
56
57mod field_option_tags {
59 pub const FEATURES: u32 = 21;
60}
61
62mod feature_tags {
64 pub const FIELD_PRESENCE: u32 = 1;
65 pub const MESSAGE_ENCODING: u32 = 5;
66}
67
68mod option_tags {
70 pub const FEATURES: u32 = 50;
71}
72
73#[must_use]
77pub fn normalize_edition_descriptor_set(bytes: &[u8]) -> Vec<u8> {
78 let mut cursor = bytes;
79 let mut has_editions = false;
80
81 while !cursor.is_empty() {
83 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
84 return bytes.to_vec();
85 };
86 match (tag, wire_type) {
87 (1, WireType::LengthDelimited) => {
88 let Ok(len) = decode_len(&mut cursor) else {
89 return bytes.to_vec();
90 };
91 if cursor.len() < len {
92 return bytes.to_vec();
93 }
94 if file_has_editions_syntax(&cursor[..len]) {
95 has_editions = true;
96 break;
97 }
98 cursor = &cursor[len..];
99 }
100 _ => {
101 if skip_wire_value(&mut cursor, wire_type).is_err() {
102 return bytes.to_vec();
103 }
104 }
105 }
106 }
107
108 if !has_editions {
109 return bytes.to_vec();
110 }
111
112 let mut cursor = bytes;
114 let mut out = Vec::with_capacity(bytes.len());
115
116 while !cursor.is_empty() {
117 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
118 return bytes.to_vec();
119 };
120 if let (1, WireType::LengthDelimited) = (tag, wire_type) {
121 let Ok(len) = decode_len(&mut cursor) else {
122 return bytes.to_vec();
123 };
124 if cursor.len() < len {
125 return bytes.to_vec();
126 }
127 let file_bytes = &cursor[..len];
128 cursor = &cursor[len..];
129
130 let normalized = normalize_file_descriptor(file_bytes);
131 encode_key(1, WireType::LengthDelimited, &mut out);
132 encode_varint(normalized.len() as u64, &mut out);
133 out.extend_from_slice(&normalized);
134 } else {
135 let start = cursor;
136 if skip_wire_value(&mut cursor, wire_type).is_err() {
137 return bytes.to_vec();
138 }
139 encode_key(tag, wire_type, &mut out);
141 out.extend_from_slice(&start[..start.len() - cursor.len()]);
142 }
143 }
144
145 out
146}
147
148fn file_has_editions_syntax(bytes: &[u8]) -> bool {
150 let mut cursor = bytes;
151 while !cursor.is_empty() {
152 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
153 return false;
154 };
155 match (tag, wire_type) {
156 (file_tags::SYNTAX, WireType::LengthDelimited) => {
157 let Ok(len) = decode_len(&mut cursor) else {
158 return false;
159 };
160 if cursor.len() < len {
161 return false;
162 }
163 return &cursor[..len] == b"editions";
164 }
165 _ => {
166 if skip_wire_value(&mut cursor, wire_type).is_err() {
167 return false;
168 }
169 }
170 }
171 }
172 false
173}
174
175#[allow(clippy::cast_possible_truncation)] fn extract_feature_set_varint(bytes: &[u8], field_tag: u32) -> i32 {
178 let mut cursor = bytes;
179 while !cursor.is_empty() {
180 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
181 break;
182 };
183 match (tag, wire_type) {
184 (t, WireType::Varint) if t == field_tag => {
185 let Ok(v) = decode_varint(&mut cursor) else {
186 break;
187 };
188 return v as i32;
189 }
190 _ => {
191 if skip_wire_value(&mut cursor, wire_type).is_err() {
192 break;
193 }
194 }
195 }
196 }
197 0
198}
199
200fn extract_feature_varint(options_bytes: &[u8], features_tag: u32, field_tag: u32) -> i32 {
205 let mut cursor = options_bytes;
206 while !cursor.is_empty() {
207 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
208 break;
209 };
210 match (tag, wire_type) {
211 (t, WireType::LengthDelimited) if t == features_tag => {
212 let Ok(len) = decode_len(&mut cursor) else {
213 break;
214 };
215 if cursor.len() < len {
216 break;
217 }
218 let feature_set = &cursor[..len];
219 let val = extract_feature_set_varint(feature_set, field_tag);
220 if val != 0 {
221 return val;
222 }
223 cursor = &cursor[len..];
224 }
225 _ => {
226 if skip_wire_value(&mut cursor, wire_type).is_err() {
227 break;
228 }
229 }
230 }
231 }
232 0
233}
234
235fn extract_file_level_feature(bytes: &[u8], feature_field_tag: u32) -> i32 {
240 let mut cursor = bytes;
241 while !cursor.is_empty() {
242 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
243 break;
244 };
245 match (tag, wire_type) {
246 (file_tags::OPTIONS, WireType::LengthDelimited) => {
247 let Ok(len) = decode_len(&mut cursor) else {
248 break;
249 };
250 if cursor.len() < len {
251 break;
252 }
253 let options_bytes = &cursor[..len];
254 cursor = &cursor[len..];
255 let val =
256 extract_feature_varint(options_bytes, option_tags::FEATURES, feature_field_tag);
257 if val != 0 {
258 return val;
259 }
260 }
261 _ => {
262 if skip_wire_value(&mut cursor, wire_type).is_err() {
263 break;
264 }
265 }
266 }
267 }
268 0
269}
270
271fn normalize_file_descriptor(bytes: &[u8]) -> Vec<u8> {
274 if !file_has_editions_syntax(bytes) {
275 return bytes.to_vec();
276 }
277
278 let presence = extract_file_level_feature(bytes, feature_tags::FIELD_PRESENCE);
279 let file_default_presence = if presence != 0 {
280 presence
281 } else {
282 FIELD_PRESENCE_EXPLICIT
283 };
284 let file_default_encoding = extract_file_level_feature(bytes, feature_tags::MESSAGE_ENCODING);
285
286 let mut cursor = bytes;
287 let mut out = Vec::with_capacity(bytes.len());
288
289 while !cursor.is_empty() {
290 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
291 return bytes.to_vec();
292 };
293
294 match (tag, wire_type) {
295 (file_tags::SYNTAX, WireType::LengthDelimited) => {
297 let Ok(len) = decode_len(&mut cursor) else {
298 return bytes.to_vec();
299 };
300 if cursor.len() < len {
301 return bytes.to_vec();
302 }
303 cursor = &cursor[len..];
304 encode_key(file_tags::SYNTAX, WireType::LengthDelimited, &mut out);
306 encode_varint(6, &mut out); out.extend_from_slice(b"proto3");
308 }
309 (file_tags::EDITION, WireType::Varint) => {
311 let Ok(_) = decode_varint(&mut cursor) else {
312 return bytes.to_vec();
313 };
314 }
316 (file_tags::MESSAGE_TYPE, WireType::LengthDelimited) => {
318 let Ok(len) = decode_len(&mut cursor) else {
319 return bytes.to_vec();
320 };
321 if cursor.len() < len {
322 return bytes.to_vec();
323 }
324 let msg_bytes = &cursor[..len];
325 cursor = &cursor[len..];
326 let normalized = normalize_message_descriptor(
327 msg_bytes,
328 file_default_presence,
329 file_default_encoding,
330 );
331 encode_key(file_tags::MESSAGE_TYPE, WireType::LengthDelimited, &mut out);
332 encode_varint(normalized.len() as u64, &mut out);
333 out.extend_from_slice(&normalized);
334 }
335 (file_tags::EXTENSION, WireType::LengthDelimited) => {
337 let Ok(len) = decode_len(&mut cursor) else {
338 return bytes.to_vec();
339 };
340 if cursor.len() < len {
341 return bytes.to_vec();
342 }
343 let field_bytes = &cursor[..len];
344 cursor = &cursor[len..];
345 let normalized = normalize_field_descriptor(
346 field_bytes,
347 file_default_presence,
348 file_default_encoding,
349 );
350 encode_key(file_tags::EXTENSION, WireType::LengthDelimited, &mut out);
351 encode_varint(normalized.len() as u64, &mut out);
352 out.extend_from_slice(&normalized);
353 }
354 _ => {
356 let pre = cursor;
357 if skip_wire_value(&mut cursor, wire_type).is_err() {
358 return bytes.to_vec();
359 }
360 encode_key(tag, wire_type, &mut out);
361 out.extend_from_slice(&pre[..pre.len() - cursor.len()]);
362 }
363 }
364 }
365
366 out
367}
368
369#[allow(clippy::too_many_lines)] fn normalize_message_descriptor(
372 bytes: &[u8],
373 parent_presence: i32,
374 parent_encoding: i32,
375) -> Vec<u8> {
376 let msg_presence = extract_message_level_feature(bytes, feature_tags::FIELD_PRESENCE)
378 .unwrap_or(parent_presence);
379 let msg_encoding = extract_message_level_feature(bytes, feature_tags::MESSAGE_ENCODING)
380 .unwrap_or(parent_encoding);
381
382 let mut cursor = bytes;
383 let mut out = Vec::with_capacity(bytes.len());
384 let mut oneof_count = 0u32;
385
386 {
388 let mut scan = bytes;
389 while !scan.is_empty() {
390 let Ok((tag, wire_type)) = decode_key(&mut scan) else {
391 break;
392 };
393 if tag == message_tags::ONEOF_DECL && wire_type == WireType::LengthDelimited {
394 oneof_count += 1;
395 }
396 if skip_wire_value(&mut scan, wire_type).is_err() {
397 break;
398 }
399 }
400 }
401
402 let mut fields_needing_synthetic_oneof = Vec::new();
404 let mut field_index = 0u32;
405
406 {
408 let mut scan = bytes;
409 while !scan.is_empty() {
410 let Ok((tag, wire_type)) = decode_key(&mut scan) else {
411 break;
412 };
413 if tag == message_tags::FIELD && wire_type == WireType::LengthDelimited {
414 let Ok(len) = decode_len(&mut scan) else {
415 break;
416 };
417 if scan.len() < len {
418 break;
419 }
420 let field_bytes = &scan[..len];
421 scan = &scan[len..];
422 let info = analyze_field(field_bytes, msg_presence, msg_encoding);
423 if info.needs_proto3_optional && !info.has_oneof_index {
424 fields_needing_synthetic_oneof.push((field_index, info.name.clone()));
425 }
426 field_index += 1;
427 } else if skip_wire_value(&mut scan, wire_type).is_err() {
428 break;
429 }
430 }
431 }
432
433 let mut synthetic_oneof_map: HashMap<u32, u32> = HashMap::new();
435 for (i, (fi, _)) in fields_needing_synthetic_oneof.iter().enumerate() {
436 #[allow(clippy::cast_possible_truncation)] let idx = i as u32;
438 synthetic_oneof_map.insert(*fi, oneof_count + idx);
439 }
440
441 field_index = 0;
443 while !cursor.is_empty() {
444 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
445 return bytes.to_vec();
446 };
447
448 match (tag, wire_type) {
449 (message_tags::FIELD, WireType::LengthDelimited) => {
450 let Ok(len) = decode_len(&mut cursor) else {
451 return bytes.to_vec();
452 };
453 if cursor.len() < len {
454 return bytes.to_vec();
455 }
456 let field_bytes = &cursor[..len];
457 cursor = &cursor[len..];
458 let synthetic_oneof = synthetic_oneof_map.get(&field_index).copied();
459 let normalized = normalize_field_descriptor_with_oneof(
460 field_bytes,
461 msg_presence,
462 msg_encoding,
463 synthetic_oneof,
464 );
465 encode_key(message_tags::FIELD, WireType::LengthDelimited, &mut out);
466 encode_varint(normalized.len() as u64, &mut out);
467 out.extend_from_slice(&normalized);
468 field_index += 1;
469 }
470 (message_tags::NESTED_TYPE, WireType::LengthDelimited) => {
471 let Ok(len) = decode_len(&mut cursor) else {
472 return bytes.to_vec();
473 };
474 if cursor.len() < len {
475 return bytes.to_vec();
476 }
477 let nested_bytes = &cursor[..len];
478 cursor = &cursor[len..];
479 let normalized =
480 normalize_message_descriptor(nested_bytes, msg_presence, msg_encoding);
481 encode_key(
482 message_tags::NESTED_TYPE,
483 WireType::LengthDelimited,
484 &mut out,
485 );
486 encode_varint(normalized.len() as u64, &mut out);
487 out.extend_from_slice(&normalized);
488 }
489 (message_tags::EXTENSION, WireType::LengthDelimited) => {
490 let Ok(len) = decode_len(&mut cursor) else {
491 return bytes.to_vec();
492 };
493 if cursor.len() < len {
494 return bytes.to_vec();
495 }
496 let ext_bytes = &cursor[..len];
497 cursor = &cursor[len..];
498 let normalized = normalize_field_descriptor(ext_bytes, msg_presence, msg_encoding);
499 encode_key(message_tags::EXTENSION, WireType::LengthDelimited, &mut out);
500 encode_varint(normalized.len() as u64, &mut out);
501 out.extend_from_slice(&normalized);
502 }
503 _ => {
504 let pre = cursor;
505 if skip_wire_value(&mut cursor, wire_type).is_err() {
506 return bytes.to_vec();
507 }
508 encode_key(tag, wire_type, &mut out);
509 out.extend_from_slice(&pre[..pre.len() - cursor.len()]);
510 }
511 }
512 }
513
514 for (_, name) in &fields_needing_synthetic_oneof {
516 let oneof_name = format!("_{name}");
517 let mut oneof_bytes = Vec::new();
518 encode_key(1, WireType::LengthDelimited, &mut oneof_bytes);
520 encode_varint(oneof_name.len() as u64, &mut oneof_bytes);
521 oneof_bytes.extend_from_slice(oneof_name.as_bytes());
522
523 encode_key(
524 message_tags::ONEOF_DECL,
525 WireType::LengthDelimited,
526 &mut out,
527 );
528 encode_varint(oneof_bytes.len() as u64, &mut out);
529 out.extend_from_slice(&oneof_bytes);
530 }
531
532 out
533}
534
535fn extract_message_level_feature(bytes: &[u8], feature_field_tag: u32) -> Option<i32> {
540 let mut cursor = bytes;
541 while !cursor.is_empty() {
542 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
543 break;
544 };
545 match (tag, wire_type) {
546 (7, WireType::LengthDelimited) => {
548 let Ok(len) = decode_len(&mut cursor) else {
549 break;
550 };
551 if cursor.len() < len {
552 break;
553 }
554 let options_bytes = &cursor[..len];
555 cursor = &cursor[len..];
556 let val =
557 extract_feature_varint(options_bytes, option_tags::FEATURES, feature_field_tag);
558 if val != 0 {
559 return Some(val);
560 }
561 }
562 _ => {
563 if skip_wire_value(&mut cursor, wire_type).is_err() {
564 break;
565 }
566 }
567 }
568 }
569 None
570}
571
572#[allow(clippy::struct_excessive_bools)] struct FieldInfo {
574 name: String,
575 needs_proto3_optional: bool,
576 has_oneof_index: bool,
577 is_delimited: bool,
578 is_legacy_required: bool,
579}
580
581#[allow(clippy::too_many_lines, clippy::cast_possible_truncation)]
582fn analyze_field(bytes: &[u8], parent_presence: i32, parent_encoding: i32) -> FieldInfo {
584 let mut cursor = bytes;
585 let mut name = String::new();
586 let mut label = 0i32;
587 let mut field_type = 0i32;
588 let mut has_oneof_index = false;
589 let mut field_presence = 0i32;
590 let mut field_encoding = 0i32;
591 let mut has_proto3_optional = false;
592
593 while !cursor.is_empty() {
594 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
595 break;
596 };
597 match (tag, wire_type) {
598 (field_tags::NAME, WireType::LengthDelimited) => {
599 let Ok(len) = decode_len(&mut cursor) else {
600 break;
601 };
602 if cursor.len() < len {
603 break;
604 }
605 name = String::from_utf8_lossy(&cursor[..len]).to_string();
606 cursor = &cursor[len..];
607 }
608 (field_tags::LABEL, WireType::Varint) => {
609 let Ok(v) = decode_varint(&mut cursor) else {
610 break;
611 };
612 label = v as i32;
613 }
614 (field_tags::TYPE, WireType::Varint) => {
615 let Ok(v) = decode_varint(&mut cursor) else {
616 break;
617 };
618 field_type = v as i32;
619 }
620 (field_tags::ONEOF_INDEX, WireType::Varint) => {
621 let Ok(_) = decode_varint(&mut cursor) else {
622 break;
623 };
624 has_oneof_index = true;
625 }
626 (field_tags::PROTO3_OPTIONAL, WireType::Varint) => {
627 let Ok(v) = decode_varint(&mut cursor) else {
628 break;
629 };
630 has_proto3_optional = v != 0;
631 }
632 (field_tags::OPTIONS, WireType::LengthDelimited) => {
633 let Ok(len) = decode_len(&mut cursor) else {
634 break;
635 };
636 if cursor.len() < len {
637 break;
638 }
639 let options = &cursor[..len];
640 field_presence = extract_feature_varint(
641 options,
642 field_option_tags::FEATURES,
643 feature_tags::FIELD_PRESENCE,
644 );
645 field_encoding = extract_feature_varint(
646 options,
647 field_option_tags::FEATURES,
648 feature_tags::MESSAGE_ENCODING,
649 );
650 cursor = &cursor[len..];
651 }
652 _ => {
653 if skip_wire_value(&mut cursor, wire_type).is_err() {
654 break;
655 }
656 }
657 }
658 }
659
660 let effective_presence = if field_presence != 0 {
661 field_presence
662 } else {
663 parent_presence
664 };
665
666 let is_repeated = label == LABEL_REPEATED;
668 let is_message = field_type == TYPE_MESSAGE || field_type == TYPE_GROUP;
669 let needs_proto3_optional = !has_proto3_optional
670 && !is_repeated
671 && !has_oneof_index
672 && effective_presence == FIELD_PRESENCE_EXPLICIT
673 && !is_message;
674
675 let effective_encoding = if field_encoding != 0 {
677 field_encoding
678 } else {
679 parent_encoding
680 };
681 let is_delimited =
682 field_type == TYPE_MESSAGE && effective_encoding == MESSAGE_ENCODING_DELIMITED;
683 let is_legacy_required = effective_presence == FIELD_PRESENCE_LEGACY_REQUIRED;
684
685 FieldInfo {
686 name,
687 needs_proto3_optional,
688 has_oneof_index,
689 is_delimited,
690 is_legacy_required,
691 }
692}
693
694fn normalize_field_descriptor(bytes: &[u8], parent_presence: i32, parent_encoding: i32) -> Vec<u8> {
696 normalize_field_descriptor_with_oneof(bytes, parent_presence, parent_encoding, None)
697}
698
699#[allow(clippy::cast_possible_truncation)] fn normalize_field_descriptor_with_oneof(
701 bytes: &[u8],
702 parent_presence: i32,
703 parent_encoding: i32,
704 synthetic_oneof_index: Option<u32>,
705) -> Vec<u8> {
706 let info = analyze_field(bytes, parent_presence, parent_encoding);
707
708 let mut cursor = bytes;
709 let mut out = Vec::with_capacity(bytes.len() + 8);
710
711 while !cursor.is_empty() {
712 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
713 return bytes.to_vec();
714 };
715
716 match (tag, wire_type) {
717 (field_tags::LABEL, WireType::Varint) => {
719 let Ok(v) = decode_varint(&mut cursor) else {
720 return bytes.to_vec();
721 };
722 encode_key(field_tags::LABEL, WireType::Varint, &mut out);
723 if info.is_legacy_required && v as i32 == LABEL_OPTIONAL {
724 encode_varint(LABEL_REQUIRED as u64, &mut out);
725 } else {
726 encode_varint(v, &mut out);
727 }
728 }
729 (field_tags::TYPE, WireType::Varint) => {
731 let Ok(v) = decode_varint(&mut cursor) else {
732 return bytes.to_vec();
733 };
734 encode_key(field_tags::TYPE, WireType::Varint, &mut out);
735 if info.is_delimited && v as i32 == TYPE_MESSAGE {
736 encode_varint(TYPE_GROUP as u64, &mut out);
737 } else {
738 encode_varint(v, &mut out);
739 }
740 }
741 _ => {
743 let pre = cursor;
744 if skip_wire_value(&mut cursor, wire_type).is_err() {
745 return bytes.to_vec();
746 }
747 encode_key(tag, wire_type, &mut out);
748 out.extend_from_slice(&pre[..pre.len() - cursor.len()]);
749 }
750 }
751 }
752
753 if (info.needs_proto3_optional || synthetic_oneof_index.is_some())
755 && !has_field_tag(bytes, field_tags::PROTO3_OPTIONAL)
756 {
757 encode_key(field_tags::PROTO3_OPTIONAL, WireType::Varint, &mut out);
758 encode_varint(1, &mut out);
759 }
760
761 if let Some(idx) = synthetic_oneof_index {
763 if !info.has_oneof_index {
764 encode_key(field_tags::ONEOF_INDEX, WireType::Varint, &mut out);
765 encode_varint(u64::from(idx), &mut out);
766 }
767 }
768
769 out
770}
771
772fn has_field_tag(bytes: &[u8], target_tag: u32) -> bool {
774 let mut cursor = bytes;
775 while !cursor.is_empty() {
776 let Ok((tag, wire_type)) = decode_key(&mut cursor) else {
777 return false;
778 };
779 if tag == target_tag {
780 return true;
781 }
782 if skip_wire_value(&mut cursor, wire_type).is_err() {
783 return false;
784 }
785 }
786 false
787}
788
789#[cfg(test)]
790mod tests {
791 use super::normalize_edition_descriptor_set;
792 use proptest::collection::vec;
793 use proptest::prelude::*;
794
795 proptest! {
796 #[test]
797 fn normalization_is_idempotent_for_arbitrary_bytes(input in vec(any::<u8>(), 0..2048)) {
798 let once = normalize_edition_descriptor_set(&input);
799 let twice = normalize_edition_descriptor_set(&once);
800 prop_assert_eq!(twice, once);
801 }
802 }
803}