Merge pull request #1406 from dxbjavid/function-text-embedded-nul

preserve embedded NUL bytes in custom function text values

callback.go

@@ -18,7 +18,7 @@ package sqlite3
 #endif
 #include <stdlib.h>
 
-void _sqlite3_result_text(sqlite3_context* ctx, const char* s);
+void _sqlite3_result_text(sqlite3_context* ctx, const char* s, int n);
 void _sqlite3_result_blob(sqlite3_context* ctx, const void* b, int l);
 */
 import "C"
@@ -209,7 +209,8 @@ func callbackArgString(v *C.sqlite3_value) (reflect.Value, error) {
 		return reflect.ValueOf(C.GoStringN(p, l)), nil
 	case C.SQLITE_TEXT:
 		c := (*C.char)(unsafe.Pointer(C.sqlite3_value_text(v)))
-		return reflect.ValueOf(C.GoString(c)), nil
+		l := C.sqlite3_value_bytes(v)
+		return reflect.ValueOf(C.GoStringN(c, l)), nil
 	default:
 		return reflect.Value{}, fmt.Errorf("argument must be BLOB or TEXT")
 	}
@@ -349,8 +350,13 @@ func callbackRetText(ctx *C.sqlite3_context, v reflect.Value) error {
 	if v.Type().Kind() != reflect.String {
 		return fmt.Errorf("cannot convert %s to TEXT", v.Type())
 	}
-	cstr := C.CString(v.Interface().(string))
-	C._sqlite3_result_text(ctx, cstr)
+	s := v.Interface().(string)
+	if i64 && len(s) > math.MaxInt32 {
+		C.sqlite3_result_error_toobig(ctx)
+		return nil
+	}
+	cstr := C.CString(s)
+	C._sqlite3_result_text(ctx, cstr, C.int(len(s)))
 	return nil
 }
 

sqlite3.go

@@ -220,8 +220,8 @@ _sqlite3_prepare_v2_internal(sqlite3 *db, const char *zSql, int nBytes, sqlite3_
 }
 #endif
 
-void _sqlite3_result_text(sqlite3_context* ctx, const char* s) {
-  sqlite3_result_text(ctx, s, -1, &free);
+void _sqlite3_result_text(sqlite3_context* ctx, const char* s, int n) {
+  sqlite3_result_text(ctx, s, n, &free);
 }
 
 void _sqlite3_result_blob(sqlite3_context* ctx, const void* b, int l) {

sqlite3_test.go

@@ -1439,6 +1439,42 @@ func TestFunctionRegistration(t *testing.T) {
 	}
 }
 
+func TestFunctionArgStringContainingZero(t *testing.T) {
+	sql.Register("sqlite3_FunctionArgZero", &SQLiteDriver{
+		ConnectHook: func(conn *SQLiteConn) error {
+			// arglen reports how many bytes of the text argument reached the
+			// Go side; echo returns the string result verbatim.
+			if err := conn.RegisterFunc("arglen", func(s string) int64 { return int64(len(s)) }, true); err != nil {
+				return err
+			}
+			return conn.RegisterFunc("echo", func(s string) string { return s }, true)
+		},
+	})
+	db, err := sql.Open("sqlite3_FunctionArgZero", ":memory:")
+	if err != nil {
+		t.Fatal("Failed to open database:", err)
+	}
+	defer db.Close()
+
+	const text = "foo\x00bar"
+
+	var n int64
+	if err := db.QueryRow("SELECT arglen(?)", text).Scan(&n); err != nil {
+		t.Fatal("Failed to call db.QueryRow:", err)
+	}
+	if n != int64(len(text)) {
+		t.Errorf("text argument truncated at embedded NUL: got len %d, want %d", n, len(text))
+	}
+
+	var got string
+	if err := db.QueryRow("SELECT echo(?)", text).Scan(&got); err != nil {
+		t.Fatal("Failed to call db.QueryRow:", err)
+	}
+	if got != text {
+		t.Errorf("text result truncated at embedded NUL: got %q (len %d), want %q (len %d)", got, len(got), text, len(text))
+	}
+}
+
 type sumAggregator int64
 
 func (s *sumAggregator) Step(x int64) {