Share modified UTF-8 codec across classfile and data streams

Author: dlunch

classfile/src/constant_pool.rs

@@ -1,62 +1,20 @@
-use alloc::{collections::BTreeMap, string::String, sync::Arc, vec::Vec};
-use core::char::REPLACEMENT_CHARACTER;
+use alloc::{collections::BTreeMap, string::String, sync::Arc};
 
+use jvm_types::modified_utf8;
 use nom::{
     IResult, Parser,
     bytes::complete::take,
-    combinator::map_res,
+    combinator::map_opt,
     error::{Error, ErrorKind},
     number::complete::{be_f32, be_f64, be_i32, be_i64, be_u16, u8},
 };
 
 fn parse_utf8(data: &[u8]) -> IResult<&[u8], Arc<String>> {
     let (data, length) = be_u16(data)?;
-    map_res(take(length as usize), |utf8: &[u8]| decode_modified_utf8(utf8).map(Arc::new)).parse(data)
-}
-
-// CONSTANT_Utf8 is *modified* UTF-8 (JVMS §4.4.7), not UTF-8: NUL is the two bytes C0 80, and a
-// supplementary character is its UTF-16 surrogate pair with each half encoded as three bytes
-// (ED A0..AF xx ED B0..BF xx). `String::from_utf8` rejects both, so any class holding "\0" or a
-// supplementary character in a constant failed to load with ClassFormatError. Decoding goes
-// through UTF-16 units, which is what the Java string is anyway.
-//
-// Strictly wider than before: input that is valid UTF-8 takes the old path unchanged. On the slow
-// path the standard 4-byte form is accepted too (javac never emits it, but the old parser did
-// accept it, so rejecting it now would be a regression), and a lone surrogate — legal in a Java
-// string, unrepresentable in a Rust `String` — becomes U+FFFD instead of failing the whole class.
-fn decode_modified_utf8(bytes: &[u8]) -> Result<String, ()> {
-    if let Ok(x) = core::str::from_utf8(bytes) {
-        return Ok(x.into());
-    }
-
-    let continuation = |i: usize| match bytes.get(i) {
-        Some(&b) if b & 0xC0 == 0x80 => Ok((b & 0x3F) as u32),
-        _ => Err(()),
-    };
-
-    let mut units = Vec::with_capacity(bytes.len());
-    let mut i = 0;
-    while i < bytes.len() {
-        let b = bytes[i] as u32;
-        let (code_point, length) = match bytes[i] {
-            0x00..=0x7F => (b, 1),
-            0xC0..=0xDF => (((b & 0x1F) << 6) | continuation(i + 1)?, 2),
-            0xE0..=0xEF => (((b & 0x0F) << 12) | (continuation(i + 1)? << 6) | continuation(i + 2)?, 3),
-            0xF0..=0xF4 => (
-                ((b & 0x07) << 18) | (continuation(i + 1)? << 12) | (continuation(i + 2)? << 6) | continuation(i + 3)?,
-                4,
-            ),
-            _ => return Err(()),
-        };
-        match char::from_u32(code_point) {
-            Some(c) if code_point > 0xFFFF => units.extend_from_slice(c.encode_utf16(&mut [0; 2])),
-            None if code_point > 0x10FFFF => return Err(()),
-            _ => units.push(code_point as u16), // BMP, or one surrogate half (paired up below)
-        }
-        i += length;
-    }
-
-    Ok(char::decode_utf16(units).map(|x| x.unwrap_or(REPLACEMENT_CHARACTER)).collect())
+    map_opt(take(length as usize), |utf8: &[u8]| {
+        modified_utf8::decode(utf8).map(|units| Arc::new(String::from_utf16_lossy(&units)))
+    })
+    .parse(data)
 }
 
 #[derive(Debug)]
@@ -381,9 +339,8 @@ mod modified_utf8_tests {
     }
 
     #[test]
-    fn edges_of_the_wider_decoder() {
-        // standard 4-byte UTF-8 next to C0 80: the old parser accepted the first alone
-        assert_eq!(constant(b"\xF0\x9F\x90\x8D\xC0\x80").as_deref(), Some("\u{1F40D}\0"));
+    fn malformed_input_and_unpaired_surrogates() {
+        assert_eq!(constant(b"\xF0\x9F\x90\x8D"), None);
         // a lone surrogate half cannot live in a Rust String
         assert_eq!(constant(b"x\xED\xA0\xBD").as_deref(), Some("x\u{FFFD}"));
         // still malformed: truncated sequence, stray continuation byte, invalid lead byte

jvm-types/src/lib.rs

@@ -1,3 +1,9 @@
+#![no_std]
+
+extern crate alloc;
+
+pub mod modified_utf8;
+
 bitflags::bitflags! {
     #[derive(Eq, PartialEq, Default, Clone, Copy, Debug, Ord, PartialOrd)]
     pub struct FieldAccessFlags: u16 {

jvm-types/src/modified_utf8.rs

@@ -0,0 +1,45 @@
+//! Java modified UTF-8 over UTF-16 code units, without a length prefix.
+
+use alloc::vec::Vec;
+
+/// Encodes UTF-16 code units, preserving unpaired surrogates.
+pub fn encode(units: &[u16]) -> Vec<u8> {
+    let mut bytes = Vec::with_capacity(units.len());
+    for &value in units {
+        if (0x0001..=0x007f).contains(&value) {
+            bytes.push(value as u8);
+        } else if value <= 0x07ff {
+            bytes.push((0xc0 | (value >> 6)) as u8);
+            bytes.push((0x80 | (value & 0x3f)) as u8);
+        } else {
+            bytes.push((0xe0 | (value >> 12)) as u8);
+            bytes.push((0x80 | ((value >> 6) & 0x3f)) as u8);
+            bytes.push((0x80 | (value & 0x3f)) as u8);
+        }
+    }
+    bytes
+}
+
+/// Decodes using Java DataInput's one-, two-, and three-byte rules.
+/// Returns `None` for malformed input; unpaired surrogates are preserved.
+pub fn decode(bytes: &[u8]) -> Option<Vec<u16>> {
+    let continuation = |i: usize| match bytes.get(i) {
+        Some(&b) if b & 0xc0 == 0x80 => Some((b & 0x3f) as u16),
+        _ => None,
+    };
+
+    let mut units = Vec::with_capacity(bytes.len());
+    let mut i = 0;
+    while i < bytes.len() {
+        let b = bytes[i] as u16;
+        let (unit, length) = match bytes[i] {
+            0x00..=0x7f => (b, 1),
+            0xc0..=0xdf => (((b & 0x1f) << 6) | continuation(i + 1)?, 2),
+            0xe0..=0xef => (((b & 0x0f) << 12) | (continuation(i + 1)? << 6) | continuation(i + 2)?, 3),
+            _ => return None,
+        };
+        units.push(unit);
+        i += length;
+    }
+    Some(units)
+}

rustjava-runtime/src/classes/java/io/data_input_stream.rs

@@ -2,7 +2,7 @@ use alloc::{vec, vec::Vec};
 
 use jvm::{Array, ClassInstanceRef, JavaChar, Jvm, Result, runtime::JavaLangString};
 use jvm_class_proto::JavaMethodProto;
-use jvm_types::{ClassAccessFlags, MethodAccessFlags};
+use jvm_types::{ClassAccessFlags, MethodAccessFlags, modified_utf8};
 
 use crate::{
     RuntimeClassProto, RuntimeContext,
@@ -205,36 +205,9 @@ impl DataInputStream {
         let bytes: Vec<i8> = jvm.load_array(&java_array, 0, length as usize).await?;
         let bytes: Vec<u8> = bytes.into_iter().map(|value| value as u8).collect();
 
-        let mut chars = Vec::with_capacity(length as usize);
-        let mut index = 0;
-        while index < bytes.len() {
-            let first = bytes[index];
-            match first >> 4 {
-                0..=7 => {
-                    chars.push(first as JavaChar);
-                    index += 1;
-                }
-                12 | 13 => {
-                    if index + 1 >= bytes.len() || bytes[index + 1] & 0xc0 != 0x80 {
-                        return Err(jvm.exception("java/io/UTFDataFormatException", "malformed modified UTF-8").await);
-                    }
-                    chars.push((((first & 0x1f) as JavaChar) << 6) | ((bytes[index + 1] & 0x3f) as JavaChar));
-                    index += 2;
-                }
-                14 => {
-                    if index + 2 >= bytes.len() || bytes[index + 1] & 0xc0 != 0x80 || bytes[index + 2] & 0xc0 != 0x80 {
-                        return Err(jvm.exception("java/io/UTFDataFormatException", "malformed modified UTF-8").await);
-                    }
-                    chars.push(
-                        (((first & 0x0f) as JavaChar) << 12)
-                            | (((bytes[index + 1] & 0x3f) as JavaChar) << 6)
-                            | ((bytes[index + 2] & 0x3f) as JavaChar),
-                    );
-                    index += 3;
-                }
-                _ => return Err(jvm.exception("java/io/UTFDataFormatException", "malformed modified UTF-8").await),
-            }
-        }
+        let Some(chars) = modified_utf8::decode(&bytes) else {
+            return Err(jvm.exception("java/io/UTFDataFormatException", "malformed modified UTF-8").await);
+        };
 
         Ok(JavaLangString::from_utf16(jvm, chars).await?.into())
     }

rustjava-runtime/src/classes/java/io/data_output_stream.rs

@@ -4,7 +4,7 @@ use bytemuck::cast_vec;
 
 use jvm::{Array, ClassInstanceRef, JavaChar, Jvm, Result};
 use jvm_class_proto::JavaMethodProto;
-use jvm_types::{ClassAccessFlags, MethodAccessFlags};
+use jvm_types::{ClassAccessFlags, MethodAccessFlags, modified_utf8};
 
 use crate::{
     RuntimeClassProto, RuntimeContext,
@@ -225,19 +225,7 @@ impl DataOutputStream {
         let chars: ClassInstanceRef<Array<JavaChar>> = jvm.invoke_virtual(&s, "java/lang/String", "toCharArray", "()[C", ()).await?;
         let length = jvm.array_length(&chars).await?;
         let chars: Vec<JavaChar> = jvm.load_array(&chars, 0, length).await?;
-        let mut data = Vec::new();
-        for value in chars {
-            if (0x0001..=0x007f).contains(&value) {
-                data.push(value as i8);
-            } else if value <= 0x07ff {
-                data.push((0xc0 | ((value >> 6) & 0x1f)) as i8);
-                data.push((0x80 | (value & 0x3f)) as i8);
-            } else {
-                data.push((0xe0 | ((value >> 12) & 0x0f)) as i8);
-                data.push((0x80 | ((value >> 6) & 0x3f)) as i8);
-                data.push((0x80 | (value & 0x3f)) as i8);
-            }
-        }
+        let data: Vec<i8> = cast_vec(modified_utf8::encode(&chars));
 
         if data.len() > u16::MAX as usize {
             return Err(jvm.exception("java/io/UTFDataFormatException", "encoded string is too long").await);

rustjava-runtime/tests/classes/java/io/test_data_output_stream.rs

@@ -214,7 +214,8 @@ async fn test_data_stream_modified_utf_round_trip_and_malformed_input() -> Resul
     let data_output = jvm
         .new_class("java/io/DataOutputStream", "(Ljava/io/OutputStream;)V", (output.clone(),))
         .await?;
-    let expected = JavaLangString::from_rust_string(&jvm, "\0A😀한").await?;
+    let units = [0, 0x41, 0x7f, 0x80, 0x7ff, 0x800, 0xd83d, 0xde00, 0xd800, 0x41, 0xdc00, 0xffff];
+    let expected = JavaLangString::from_utf16(&jvm, units.to_vec()).await?;
     let _: () = jvm
         .invoke_virtual(
             &data_output,
@@ -229,26 +230,81 @@ async fn test_data_stream_modified_utf_round_trip_and_malformed_input() -> Resul
         .invoke_virtual(&output, &output.class_definition().name(), "toByteArray", "()[B", ())
         .await?;
     let encoded: Vec<i8> = jvm.load_array(&bytes, 0, jvm.array_length(&bytes).await?).await?;
-    assert_eq!(&encoded[2..4], &[-64, -128]);
+    assert_eq!(
+        encoded.iter().map(|&byte| byte as u8).collect::<Vec<_>>(),
+        [
+            0, 27, 0xc0, 0x80, 0x41, 0x7f, 0xc2, 0x80, 0xdf, 0xbf, 0xe0, 0xa0, 0x80, 0xed, 0xa0, 0xbd, 0xed, 0xb8, 0x80, 0xed, 0xa0, 0x80, 0x41,
+            0xed, 0xb0, 0x80, 0xef, 0xbf, 0xbf,
+        ]
+    );
 
     let input = jvm.new_class("java/io/ByteArrayInputStream", "([B)V", (bytes,)).await?;
     let data_input = jvm.new_class("java/io/DataInputStream", "(Ljava/io/InputStream;)V", (input,)).await?;
     let decoded: ClassInstanceRef<rustjava_runtime::classes::java::lang::String> = jvm
         .invoke_virtual(&data_input, &data_input.class_definition().name(), "readUTF", "()Ljava/lang/String;", ())
         .await?;
-    assert_eq!(JavaLangString::to_rust_string(&jvm, &decoded).await?, "\0A😀한");
+    assert_eq!(JavaLangString::to_utf16(&jvm, &decoded).await?, units);
 
-    let mut malformed = jvm.instantiate_array("B", 4).await?;
-    jvm.store_array(&mut malformed, 0, [0i8, 2, -62, 32]).await?;
-    let input = jvm.new_class("java/io/ByteArrayInputStream", "([B)V", (malformed,)).await?;
-    let data_input = jvm.new_class("java/io/DataInputStream", "(Ljava/io/InputStream;)V", (input,)).await?;
-    let result: Result<ClassInstanceRef<rustjava_runtime::classes::java::lang::String>> = jvm
-        .invoke_virtual(&data_input, &data_input.class_definition().name(), "readUTF", "()Ljava/lang/String;", ())
-        .await;
-    let Err(JavaError::JavaException(exception)) = result else {
-        panic!("malformed modified UTF-8 must throw UTFDataFormatException");
-    };
-    assert!(jvm.is_instance(&*exception, "java/io/UTFDataFormatException"));
+    for payload in [
+        &[0xc2, 0x20][..],
+        &[0xc0],
+        &[0xe0, 0xa0],
+        &[0xe0, 0x20, 0x80],
+        &[0xe0, 0xa0, 0x20],
+        &[0xf0, 0x9f, 0x98, 0x80],
+        &[0x80],
+        &[0xff],
+    ] {
+        let mut data = (payload.len() as u16).to_be_bytes().to_vec();
+        data.extend_from_slice(payload);
+        let mut malformed = jvm.instantiate_array("B", data.len()).await?;
+        jvm.store_array(&mut malformed, 0, data.into_iter().map(|byte| byte as i8).collect::<Vec<_>>())
+            .await?;
+        let input = jvm.new_class("java/io/ByteArrayInputStream", "([B)V", (malformed,)).await?;
+        let data_input = jvm.new_class("java/io/DataInputStream", "(Ljava/io/InputStream;)V", (input,)).await?;
+        let result: Result<ClassInstanceRef<rustjava_runtime::classes::java::lang::String>> = jvm
+            .invoke_virtual(&data_input, &data_input.class_definition().name(), "readUTF", "()Ljava/lang/String;", ())
+            .await;
+        let Err(JavaError::JavaException(exception)) = result else {
+            panic!("malformed modified UTF-8 must throw UTFDataFormatException: {payload:?}");
+        };
+        assert!(jvm.is_instance(&*exception, "java/io/UTFDataFormatException"));
+    }
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_write_utf_encoded_length_limit() -> Result<()> {
+    let jvm = test_jvm().await?;
+    let output = jvm.new_class("java/io/ByteArrayOutputStream", "()V", ()).await?;
+    let data_output = jvm
+        .new_class("java/io/DataOutputStream", "(Ljava/io/OutputStream;)V", (output.clone(),))
+        .await?;
+
+    for (length, succeeds) in [(0, true), (21845, true), (21846, false)] {
+        let _: () = jvm.invoke_virtual(&output, "java/io/ByteArrayOutputStream", "reset", "()V", ()).await?;
+        let value = JavaLangString::from_utf16(&jvm, vec![0x800; length]).await?;
+        let result: Result<()> = jvm
+            .invoke_virtual(&data_output, "java/io/DataOutputStream", "writeUTF", "(Ljava/lang/String;)V", (value,))
+            .await;
+        if succeeds {
+            result?;
+            let bytes: ClassInstanceRef<Array<i8>> = jvm
+                .invoke_virtual(&output, "java/io/ByteArrayOutputStream", "toByteArray", "()[B", ())
+                .await?;
+            assert_eq!(jvm.array_length(&bytes).await?, length * 3 + 2);
+            let header: Vec<i8> = jvm.load_array(&bytes, 0, 2).await?;
+            assert_eq!(header, if length == 0 { [0, 0] } else { [-1, -1] });
+        } else {
+            let Err(JavaError::JavaException(exception)) = result else {
+                panic!("oversized modified UTF-8 must throw UTFDataFormatException");
+            };
+            assert!(jvm.is_instance(&*exception, "java/io/UTFDataFormatException"));
+            let size: i32 = jvm.invoke_virtual(&output, "java/io/ByteArrayOutputStream", "size", "()I", ()).await?;
+            assert_eq!(size, 0);
+        }
+    }
 
     Ok(())
 }