Repository: GreenmaskIO/greenmask
Pull request #3 (2024-01-30 12:22:37)

Greenmask v0.1.2 release
This PR introduces a bug fixes:

* Fixed bug when raw COPY lines were parsed incorrectly #2 
* Fixed `--version` parameter behaviour
* Fixed `--dbname` parameter - now it correctly works with PostgreSQL connection string in URI format `postgresql:///`


Patch:
diff --git a/Makefile b/Makefile
index b9268fe..b61f2d5 100644
--- a/Makefile
+++ b/Makefile
@@ -4,7 +4,7 @@ CMD_FILES = $(wildcard *.go)
 TEST_FILES = $(wildcard *.go)
 COVERAGE_FILE := coverage.out
 VERSION ?= $(shell git tag --points-at HEAD)
-LDFLAGS ?= -X main.version=$(VERSION)
+LDFLAGS ?= -X github.com/greenmaskio/greenmask/cmd/greenmask/cmd.Version=$(VERSION)
 
 .PHONY: build
 
diff --git a/cmd/greenmask/main.go b/cmd/greenmask/main.go
index 4afb899..21fb64b 100644
--- a/cmd/greenmask/main.go
+++ b/cmd/greenmask/main.go
@@ -20,10 +20,7 @@ import (
 	"github.com/greenmaskio/greenmask/cmd/greenmask/cmd"
 )
 
-var version string
-
 func main() {
-	cmd.Version = version
 	if err := cmd.Execute(); err != nil {
 		log.Fatal().Err(err).Msg("")
 	}
diff --git a/internal/db/postgres/pgcopy/row.go b/internal/db/postgres/pgcopy/row.go
index 1143717..be27ce1 100644
--- a/internal/db/postgres/pgcopy/row.go
+++ b/internal/db/postgres/pgcopy/row.go
@@ -73,7 +73,7 @@ func (r *Row) Decode(raw []byte) error {
 
 	// Building column position slice
 	idx := 0
-	for colStartPos < len(raw) {
+	for colStartPos <= len(raw) {
 
 		colEndPos = slices.Index(raw[colStartPos:], DefaultCopyDelimiter)
 		if colEndPos == -1 {
diff --git a/internal/db/postgres/pgdump/pgdump.go b/internal/db/postgres/pgdump/pgdump.go
index 817dd9e..d511c3e 100644
--- a/internal/db/postgres/pgdump/pgdump.go
+++ b/internal/db/postgres/pgdump/pgdump.go
@@ -104,10 +104,11 @@ type Options struct {
 }
 
 func (o *Options) GetPgDSN() (string, error) {
-	//return "host=localhost port=5432 user=postgres dbname=postgres", nil
-	if strings.Contains(o.DbName, "=") {
+	// URI or Standard format
+	if strings.HasPrefix(o.DbName, "postgresql://") || strings.Contains(o.DbName, "=") {
 		return o.DbName, nil
 	}
+
 	return fmt.Sprintf("host=%s port=%d user=%s dbname=%s", o.Host, o.Port, o.UserName, o.DbName), nil
 }
 


Test patch:
diff --git a/internal/db/postgres/pgcopy/row_test.go b/internal/db/postgres/pgcopy/row_test.go
index 7872f1f..9f651a8 100644
--- a/internal/db/postgres/pgcopy/row_test.go
+++ b/internal/db/postgres/pgcopy/row_test.go
@@ -23,6 +23,80 @@ import (
 	"github.com/greenmaskio/greenmask/pkg/toolkit"
 )
 
+func TestDecode_positions_panic_error_regression(t *testing.T) {
+	// This case was found in https://github.com/GreenmaskIO/greenmask/issues/2
+	// Fixed bug:
+	// 	The position array Row.columnPos contained wrong values for the last column if the value is empty string
+	// 	This caused panic or corrupt dump due to the wrong positions that were not overridden after previous COPY line
+	//	parsing
+
+	type params struct {
+		record []byte
+	}
+
+	type expected struct {
+		positions []*columnPos
+		values    []string
+	}
+
+	type test struct {
+		name     string
+		params   params
+		expected expected
+	}
+
+	length := 4
+
+	tests := []test{
+		{
+			name: "not empty",
+			params: params{
+				record: []byte("27\t1213\t\tasda"),
+			},
+			expected: expected{
+				positions: []*columnPos{
+					{0, 2}, {3, 7}, {8, 8}, {9, 13},
+				},
+				values: []string{"27", "1213", "", "asda"},
+			},
+		},
+		{
+			name: "empty last two col",
+			params: params{
+				record: []byte("27\t12\t\t"),
+			},
+			expected: expected{
+				positions: []*columnPos{
+					{0, 2}, {3, 5}, {6, 6}, {7, 7},
+				},
+				values: []string{"27", "12", "", ""},
+			},
+		},
+	}
+
+	row := NewRow(length)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := row.Decode(tt.params.record)
+			require.NoError(t, err)
+			require.Equal(t, len(tt.expected.positions), len(row.columnPos))
+
+			for idx, _ := range row.columnPos {
+				assert.Equalf(t, tt.expected.positions[idx].start, row.columnPos[idx].start, "wrong start value for idx %d", idx)
+				assert.Equalf(t, tt.expected.positions[idx].end, row.columnPos[idx].end, "wrong end value for idx %d", idx)
+			}
+
+			for idx, _ := range tt.expected.values {
+				v, err := row.GetColumn(idx)
+				require.NoErrorf(t, err, "unexpected error for idx %d", idx)
+				require.Falsef(t, v.IsNull, "unexpected NULL value for idx %d", idx)
+				require.Equalf(t, tt.expected.values[idx], string(v.Data), "unexpected value for idx %d", idx)
+			}
+		})
+	}
+}
+
 func TestDecode(t *testing.T) {
 
 	type result struct {