Check data types are safe in returned query result.

internal/store/store.go

@@ -3,11 +3,30 @@ package store
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"github.com/jackc/pgx/v4"
 	"github.com/weaponry/pgscv/internal/log"
 	"github.com/weaponry/pgscv/internal/model"
 )
 
+const (
+	// Data types supported by parser of query results.
+	dataTypeBool uint32 = 16
+	// dataTypeChar uint32 = 18 is not supported - its conversion to sql.NullString lead to panic 'pgx' driver.
+	dataTypeName    uint32 = 19
+	dataTypeInt8    uint32 = 20
+	dataTypeInt2    uint32 = 21
+	dataTypeInt4    uint32 = 23
+	dataTypeText    uint32 = 25
+	dataTypeOid     uint32 = 26
+	dataTypeFloat4  uint32 = 700
+	dataTypeFloat8  uint32 = 701
+	dataTypeInet    uint32 = 869
+	dataTypeBpchar  uint32 = 1042
+	dataTypeVarchar uint32 = 1043
+	dataTypeNumeric uint32 = 1700
+)
+
 // DB is the database representation
 type DB struct {
 	conn *pgx.Conn // database connection object
@@ -69,6 +88,15 @@ func (db *DB) query(query string) (*model.PGResult, error) {
 		nrows    int
 	)
 
+	// Not the all data types could be safely converted into sql.NullString
+	// and conversion errors lead to panic.
+	// Check the data types are safe in returned result.
+	for _, c := range colnames {
+		if !isDataTypeSupported(c.DataTypeOID) {
+			return nil, fmt.Errorf("query '%s', unsupported data type OID: %d", query, c.DataTypeOID)
+		}
+	}
+
 	// Storage used for data extracted from rows.
 	// Scan operation supports only slice of interfaces, 'pointers' slice is the intermediate store where all values written.
 	// Next values from 'pointers' associated with type-strict slice - 'values'. When Scan is writing to the 'pointers' it
@@ -110,3 +138,16 @@ func (db *DB) close() {
 		log.Warnf("failed to close database connection: %s; ignore", err)
 	}
 }
+
+// isDataTypeSupported tests passed type OID is supported.
+func isDataTypeSupported(t uint32) bool {
+	switch t {
+	case dataTypeName, dataTypeBpchar, dataTypeVarchar, dataTypeText,
+		dataTypeInt2, dataTypeInt4, dataTypeInt8,
+		dataTypeOid, dataTypeFloat4, dataTypeFloat8, dataTypeNumeric,
+		dataTypeBool, dataTypeInet:
+		return true
+	default:
+		return false
+	}
+}

internal/store/store_test.go

@@ -2,6 +2,7 @@ package store
 
 import (
 	"database/sql"
+	"fmt"
 	"github.com/jackc/pgproto3/v2"
 	"github.com/jackc/pgx/v4"
 	"github.com/stretchr/testify/assert"
@@ -83,6 +84,11 @@ func TestDB_Query(t *testing.T) {
 			},
 			valid: true,
 		},
+		{
+			name:  "valid query, with unsupported data type",
+			query: "SELECT 'example' AS v1, 100 AS v2, '0x1111'::bytea AS v3",
+			valid: false,
+		},
 		{
 			name:  "invalid query",
 			query: "invalid",
@@ -112,3 +118,37 @@ func TestDB_Close(t *testing.T) {
 
 	db.Close()
 }
+
+func TestExample(t *testing.T) {
+	db := NewTest(t)
+	q := "select relkind::char as relkind from pg_class where relname in ('pg_class')"
+	_, err := db.query(q)
+	fmt.Println(err)
+	//fmt.Println(res.Rows)
+}
+
+func Test_isDataTypeSupported(t *testing.T) {
+	testcases := []struct {
+		t    uint32
+		want bool
+	}{
+		{t: 18, want: false}, // char
+		{t: dataTypeBool, want: true},
+		{t: dataTypeName, want: true},
+		{t: dataTypeInt8, want: true},
+		{t: dataTypeInt2, want: true},
+		{t: dataTypeInt4, want: true},
+		{t: dataTypeText, want: true},
+		{t: dataTypeOid, want: true},
+		{t: dataTypeFloat4, want: true},
+		{t: dataTypeFloat8, want: true},
+		{t: dataTypeInet, want: true},
+		{t: dataTypeVarchar, want: true},
+		{t: dataTypeNumeric, want: true},
+		{t: dataTypeBpchar, want: true},
+	}
+
+	for _, tc := range testcases {
+		assert.Equal(t, tc.want, isDataTypeSupported(tc.t))
+	}
+}