🐛 Fix bug preventing non-utf8 characters being saved in some databases (#87)

postgres/postgres.go

@@ -5,9 +5,9 @@ import (
 	"errors"
 	"fmt"
 	"net/url"
+	"strings"
 	"time"
 
-	"github.com/gofiber/utils"
 	_ "github.com/lib/pq"
 )
 
@@ -25,15 +25,19 @@ type Storage struct {
 }
 
 var (
+	checkSchemaMsg = "The `v` row has an incorrect data type. " +
+		"It should be BYTEA but is instead %s. This will cause encoding-related panics if the DB is not migrated (see https://github.com/gofiber/storage/blob/main/MIGRATE.md)."
 	dropQuery = `DROP TABLE IF EXISTS %s;`
 	initQuery = []string{
 		`CREATE TABLE IF NOT EXISTS %s (
 			k  VARCHAR(64) PRIMARY KEY NOT NULL DEFAULT '',
-			v  TEXT NOT NULL,
+			v  BYTEA NOT NULL,
 			e  BIGINT NOT NULL DEFAULT '0'
 		);`,
 		`CREATE INDEX IF NOT EXISTS e ON %s (e);`,
 	}
+	checkSchemaQuery = `SELECT DATA_TYPE FROM INFORMATION_SCHEMA.COLUMNS
+		WHERE table_name = '%s' AND COLUMN_NAME = 'v';`
 )
 
 // New creates a new storage
@@ -103,6 +107,8 @@ func New(config ...Config) *Storage {
 		sqlGC:      fmt.Sprintf("DELETE FROM %s WHERE e <= $1 AND e != 0", cfg.Table),
 	}
 
+	store.checkSchema(cfg.Table)
+
 	// Start garbage collector
 	go store.gcTicker()
 
@@ -137,7 +143,6 @@ func (s *Storage) Get(key string) ([]byte, error) {
 	return data, nil
 }
 
-// Set key with value
 // Set key with value
 func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 	// Ain't Nobody Got Time For That
@@ -148,8 +153,7 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 	if exp != 0 {
 		expSeconds = time.Now().Add(exp).Unix()
 	}
-	valStr := utils.UnsafeString(val)
-	_, err := s.db.Exec(s.sqlInsert, key, valStr, expSeconds, valStr, expSeconds)
+	_, err := s.db.Exec(s.sqlInsert, key, val, expSeconds, val, expSeconds)
 	return err
 }
 
@@ -193,3 +197,16 @@ func (s *Storage) gcTicker() {
 func (s *Storage) gc(t time.Time) {
 	_, _ = s.db.Exec(s.sqlGC, t.Unix())
 }
+
+func (s *Storage) checkSchema(tableName string) {
+	var data []byte
+
+	row := s.db.QueryRow(fmt.Sprintf(checkSchemaQuery, tableName))
+	if err := row.Scan(&data); err != nil {
+		panic(err)
+	}
+
+	if strings.ToLower(string(data)) != "bytea" {
+		fmt.Printf(checkSchemaMsg, string(data))
+	}
+}

postgres/postgres_test.go

@@ -148,6 +148,17 @@ func Test_Postgres_GC(t *testing.T) {
 
 }
 
+func Test_Postgres_Non_UTF8(t *testing.T) {
+	val := []byte("0xF5")
+
+	err := testStore.Set("0xF6", val, 0)
+	utils.AssertEqual(t, nil, err)
+
+	result, err := testStore.Get("0xF6")
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, val, result)
+}
+
 func Test_Postgres_Close(t *testing.T) {
 	utils.AssertEqual(t, nil, testStore.Close())
 }