arrow_cast/
base64.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Functions for converting data in [`GenericBinaryArray`] such as [`StringArray`] to/from base64 encoded strings
19//!
20//! [`StringArray`]: arrow_array::StringArray
21
22use arrow_array::{Array, GenericBinaryArray, GenericStringArray, OffsetSizeTrait};
23use arrow_buffer::{Buffer, OffsetBuffer};
24use arrow_schema::ArrowError;
25use base64::encoded_len;
26use base64::engine::Config;
27
28pub use base64::prelude::*;
29
30/// Base64 encode each element of `array` with the provided [`Engine`]
31///
32/// # Panics
33///
34/// Panics if the `Engine` emits output that is not valid UTF-8. A correct
35/// `Engine` never does, but it is a safe trait so a misbehaving impl could;
36/// validating keeps the returned [`GenericStringArray`] sound (#10284).
37pub fn b64_encode<E: Engine, O: OffsetSizeTrait>(
38    engine: &E,
39    array: &GenericBinaryArray<O>,
40) -> GenericStringArray<O> {
41    let lengths = array.offsets().windows(2).map(|w| {
42        let len = w[1].as_usize() - w[0].as_usize();
43        encoded_len(len, engine.config().encode_padding()).unwrap()
44    });
45    let offsets = OffsetBuffer::<O>::from_lengths(lengths);
46    let buffer_len = offsets.last().as_usize();
47    let mut buffer = vec![0_u8; buffer_len];
48    let mut offset = 0;
49
50    for i in 0..array.len() {
51        let len = engine
52            .encode_slice(array.value(i), &mut buffer[offset..])
53            .unwrap();
54        offset += len;
55    }
56    assert_eq!(offset, buffer_len);
57
58    // `try_new` validates UTF-8 instead of trusting the (safe-trait) Engine.
59    GenericStringArray::try_new(offsets, Buffer::from_vec(buffer), array.nulls().cloned())
60        .expect("Engine produced invalid UTF-8")
61}
62
63/// Base64 decode each element of `array` with the provided [`Engine`]
64///
65/// # Errors
66///
67/// Returns an error if a value is not valid base64 for `engine`.
68pub fn b64_decode<E: Engine, O: OffsetSizeTrait>(
69    engine: &E,
70    array: &GenericBinaryArray<O>,
71) -> Result<GenericBinaryArray<O>, ArrowError> {
72    let estimated_len = array.values().len(); // This is an overestimate
73    let mut buffer = vec![0; estimated_len];
74
75    let mut offsets = Vec::with_capacity(array.len() + 1);
76    offsets.push(O::usize_as(0));
77    let mut offset = 0;
78
79    for v in array {
80        if let Some(v) = v {
81            let len = engine
82                .decode_slice(v, &mut buffer[offset..])
83                .map_err(|err| {
84                    ArrowError::InvalidArgumentError(format!("Failed to decode base64: {err}"))
85                })?;
86            // This cannot overflow as `len` is less than `v.len()` and `a` is valid
87            offset += len;
88        }
89        offsets.push(O::usize_as(offset));
90    }
91
92    // Safety: offsets monotonically increasing by construction
93    let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
94
95    GenericBinaryArray::try_new(offsets, Buffer::from_vec(buffer), array.nulls().cloned())
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use arrow_array::BinaryArray;
102    use rand::{RngExt, rng};
103
104    fn test_engine<E: Engine>(e: &E, a: &BinaryArray) {
105        let encoded = b64_encode(e, a);
106        encoded.to_data().validate_full().unwrap();
107
108        let to_decode = encoded.into();
109        let decoded = b64_decode(e, &to_decode).unwrap();
110        decoded.to_data().validate_full().unwrap();
111
112        assert_eq!(&decoded, a);
113    }
114
115    #[test]
116    #[cfg_attr(miri, ignore)] // Takes too long
117    fn test_b64() {
118        let mut rng = rng();
119        let len = rng.random_range(1024..1050);
120        let data: BinaryArray = (0..len)
121            .map(|_| {
122                let len = rng.random_range(0..16);
123                Some((0..len).map(|_| rng.random::<u8>()).collect::<Vec<u8>>())
124            })
125            .collect();
126
127        test_engine(&BASE64_STANDARD, &data);
128        test_engine(&BASE64_STANDARD_NO_PAD, &data);
129    }
130
131    #[test]
132    fn test_b64_decode_invalid_input() {
133        let data: BinaryArray = vec![Some(b"!!!not base64!!!".to_vec())]
134            .into_iter()
135            .collect();
136        let err = b64_decode(&BASE64_STANDARD, &data).unwrap_err().to_string();
137        assert!(err.contains("Failed to decode base64"), "{err}");
138    }
139
140    /// Safe-Rust `Engine` that writes invalid UTF-8 into the encode buffer
141    /// (#10284). `b64_encode` must reject it rather than build an unsound
142    /// `StringArray`.
143    struct EvilEngine;
144
145    impl Engine for EvilEngine {
146        type Config = <base64::engine::GeneralPurpose as Engine>::Config;
147        type DecodeEstimate = <base64::engine::GeneralPurpose as Engine>::DecodeEstimate;
148
149        fn internal_encode(&self, input: &[u8], output: &mut [u8]) -> usize {
150            BASE64_STANDARD.internal_encode(input, output)
151        }
152        fn internal_decoded_len_estimate(&self, input_len: usize) -> Self::DecodeEstimate {
153            BASE64_STANDARD.internal_decoded_len_estimate(input_len)
154        }
155        fn internal_decode(
156            &self,
157            input: &[u8],
158            output: &mut [u8],
159            estimate: Self::DecodeEstimate,
160        ) -> Result<base64::engine::DecodeMetadata, base64::DecodeSliceError> {
161            BASE64_STANDARD.internal_decode(input, output, estimate)
162        }
163        fn config(&self) -> &Self::Config {
164            BASE64_STANDARD.config()
165        }
166        fn encode_slice<T: AsRef<[u8]>>(
167            &self,
168            input: T,
169            output_buf: &mut [u8],
170        ) -> Result<usize, base64::EncodeSliceError> {
171            let len = BASE64_STANDARD.encode_slice(input, output_buf)?;
172            for b in &mut output_buf[..len] {
173                *b = 0xFF; // invalid UTF-8, but correct length
174            }
175            Ok(len)
176        }
177
178        fn padding(&self) -> base64::alphabet::Symbol {
179            base64::alphabet::Symbol::new(b'=').unwrap()
180        }
181    }
182
183    #[test]
184    #[should_panic(expected = "produced invalid UTF-8")]
185    fn test_b64_encode_rejects_invalid_utf8() {
186        let data: BinaryArray = vec![Some(b"hello".to_vec())].into_iter().collect();
187        let _ = b64_encode(&EvilEngine, &data);
188    }
189}
arrow_cast/base64.rs

arrow_cast/
base64.rs
