arrow_cast/
base64.rs

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

arrow_cast/
base64.rs
