From f71ea05603199cf33b19b4d2df053f3f8e869d10 Mon Sep 17 00:00:00 2001
From: gwenn <gtreguier@gmail.com>
Date: Wed, 25 Mar 2020 19:20:05 +0100
Subject: [PATCH] Handle text with internal nuls

Fix insertion and selection (#657).
---
 src/lib.rs       | 20 +++++++-------------
 src/statement.rs | 38 ++++++++++++++++++++++++++++++++++----
 2 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 354a8ee..b17691d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -237,25 +237,19 @@ fn str_to_cstring(s: &str) -> Result<CString> {
 
 /// Returns `Ok((string ptr, len as c_int, SQLITE_STATIC | SQLITE_TRANSIENT))`
 /// normally.
-/// Returns errors if the string has embedded nuls or is too large for sqlite.
+/// Returns error if the string is too large for sqlite.
 /// The `sqlite3_destructor_type` item is always `SQLITE_TRANSIENT` unless
 /// the string was empty (in which case it's `SQLITE_STATIC`, and the ptr is
 /// static).
 fn str_for_sqlite(s: &[u8]) -> Result<(*const c_char, c_int, ffi::sqlite3_destructor_type)> {
     let len = len_as_c_int(s.len())?;
-    if memchr::memchr(0, s).is_none() {
-        let (ptr, dtor_info) = if len != 0 {
-            (s.as_ptr() as *const c_char, ffi::SQLITE_TRANSIENT())
-        } else {
-            // Return a pointer guaranteed to live forever
-            ("".as_ptr() as *const c_char, ffi::SQLITE_STATIC())
-        };
-        Ok((ptr, len, dtor_info))
+    let (ptr, dtor_info) = if len != 0 {
+        (s.as_ptr() as *const c_char, ffi::SQLITE_TRANSIENT())
     } else {
-        // There's an embedded nul, so we fabricate a NulError.
-        let e = CString::new(s);
-        Err(Error::NulError(e.unwrap_err()))
-    }
+        // Return a pointer guaranteed to live forever
+        ("".as_ptr() as *const c_char, ffi::SQLITE_STATIC())
+    };
+    Ok((ptr, len, dtor_info))
 }
 
 // Helper to cast to c_int safely, returning the correct error type if the cast
diff --git a/src/statement.rs b/src/statement.rs
index 2769d6f..30a35c4 100644
--- a/src/statement.rs
+++ b/src/statement.rs
@@ -1,6 +1,5 @@
-use std::ffi::CStr;
 use std::iter::IntoIterator;
-use std::os::raw::{c_char, c_int, c_void};
+use std::os::raw::{c_int, c_void};
 #[cfg(feature = "array")]
 use std::rc::Rc;
 use std::slice::from_raw_parts;
@@ -659,15 +658,18 @@ impl Statement<'_> {
             }
             ffi::SQLITE_TEXT => {
                 let s = unsafe {
+                    // Quoting from "Using SQLite" book:
+                    // To avoid problems, an application should first extract the desired type using a sqlite3_column_xxx() function,
+                    // and then call the appropriate sqlite3_column_bytes() function.
                     let text = ffi::sqlite3_column_text(raw, col as c_int);
+                    let len = ffi::sqlite3_column_bytes(raw, col as c_int);
                     assert!(
                         !text.is_null(),
                         "unexpected SQLITE_TEXT column type with NULL data"
                     );
-                    CStr::from_ptr(text as *const c_char)
+                    from_raw_parts(text as *const u8, len as usize)
                 };
 
-                let s = s.to_bytes();
                 ValueRef::Text(s)
             }
             ffi::SQLITE_BLOB => {
@@ -1089,4 +1091,32 @@ mod test {
         let stmt = conn.prepare(";").unwrap();
         assert_eq!(0, stmt.column_count());
     }
+
+    #[test]
+    fn test_utf16_conversion() {
+        let db = Connection::open_in_memory().unwrap();
+        db.pragma_update(None, "encoding", &"UTF-16le").unwrap();
+        let encoding: String = db
+            .pragma_query_value(None, "encoding", |row| row.get(0))
+            .unwrap();
+        assert_eq!("UTF-16le", encoding);
+        db.execute_batch("CREATE TABLE foo(x TEXT)").unwrap();
+        let expected = "テスト";
+        db.execute("INSERT INTO foo(x) VALUES (?)", &[&expected])
+            .unwrap();
+        let actual: String = db
+            .query_row("SELECT x FROM foo", NO_PARAMS, |row| row.get(0))
+            .unwrap();
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn test_nul_byte() {
+        let db = Connection::open_in_memory().unwrap();
+        let expected = "a\x00b";
+        let actual: String = db
+            .query_row("SELECT ?", &[&expected], |row| row.get(0))
+            .unwrap();
+        assert_eq!(expected, actual);
+    }
 }