Merge pull request #323 from cipherstash/go-json-struct-tests

fix: simple protocol type error for JSONB encrypted types

Author: calvinbrewer

packages/cipherstash-proxy/src/postgresql/backend.rs

@@ -578,7 +578,7 @@ where
 
         debug!(target: PROTOCOL, client_id = self.context.client_id, RowDescription = ?description);
 
-        if let Some(statement) = self.context.get_statement_from_describe() {
+        if let Some(statement) = self.context.get_statement_for_row_decription() {
             let projection_types = statement
                 .projection_columns
                 .iter()

packages/cipherstash-proxy/src/postgresql/context/mod.rs

@@ -248,6 +248,19 @@ impl Context {
         }
     }
 
+    pub fn get_statement_for_row_decription(&self) -> Option<Arc<Statement>> {
+        if let Some(statement) = self.get_statement_from_describe() {
+            return Some(statement.clone());
+        }
+
+        if let Some(Portal::Encrypted { statement, .. }) = self.get_portal_from_execute().as_deref()
+        {
+            return Some(statement.clone());
+        };
+
+        None
+    }
+
     pub fn get_statement_from_describe(&self) -> Option<Arc<Statement>> {
         let queue = self.describe.read().ok()?;
         let describe = queue.next()?;

tests/integration/golang/pgx_test.go

@@ -220,57 +220,88 @@ func TestPgxInsertEncryptedWithDomainTypeAndReturning(t *testing.T) {
 	}
 }
 
-/*
-func TestPgxEncryptedMapDate(t *testing.T) {
+// EncryptedRow represents a row with id and encrypted_text fields
+type EncryptedRow struct {
+	ID            int    `db:"id"`
+	EncryptedText string `db:"encrypted_text"`
+}
+
+// Scan implements the sql.Scanner interface for EncryptedRow
+func (er *EncryptedRow) Scan(src interface{}) error {
+	return fmt.Errorf("Scan method not implemented for single value")
+}
+
+func TestPgxInsertEncryptedWithStructScan(t *testing.T) {
+	t.Parallel()
 	conn := setupPgxConnection(t)
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	column := "encrypted_date"
-	dates := []time.Time{
-		time.Date(1, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(1000, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(1600, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(1700, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(1800, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(1900, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(1990, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(1999, 12, 31, 0, 0, 0, 0, time.UTC),
-		time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(2001, 1, 2, 0, 0, 0, 0, time.UTC),
-		time.Date(2004, 2, 29, 0, 0, 0, 0, time.UTC),
-		time.Date(2013, 7, 4, 0, 0, 0, 0, time.UTC),
-		time.Date(2013, 12, 25, 0, 0, 0, 0, time.UTC),
-		time.Date(2029, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(2081, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(2096, 2, 29, 0, 0, 0, 0, time.UTC),
-		time.Date(2550, 1, 1, 0, 0, 0, 0, time.UTC),
-		time.Date(9999, 12, 31, 0, 0, 0, 0, time.UTC),
+	tx, err := conn.Begin(ctx)
+	require.NoError(t, err)
+	defer tx.Rollback(ctx)
+
+	encryptedTextValue := "test encrypted content"
+	insertStmt := `INSERT INTO encrypted (id, encrypted_text) VALUES ($1, $2) RETURNING id, encrypted_text`
+
+	for _, mode := range modes {
+		id := rand.Int()
+
+		t.Run(mode.String(), func(t *testing.T) {
+			t.Run("insert_with_struct_scan", func(t *testing.T) {
+				var row EncryptedRow
+				err := conn.QueryRow(context.Background(), insertStmt, mode, id, encryptedTextValue).Scan(&row.ID, &row.EncryptedText)
+				require.NoError(t, err)
+				require.Equal(t, id, row.ID)
+				require.Equal(t, encryptedTextValue, row.EncryptedText)
+			})
+		})
 	}
-	insertStmt := fmt.Sprintf(`INSERT INTO encrypted (id, %s) VALUES ($1, $2)`, column)
-	selectStmt := fmt.Sprintf(`SELECT id, %s FROM encrypted WHERE id=$1`, column)
+}
 
-	for _, value := range dates {
-		t.Run(value.String(), func(t *testing.T) {
-			for _, mode := range modes {
-				id := rand.Int()
-				t.Run(mode.String(), func(t *testing.T) {
-					t.Run("insert", func(t *testing.T) {
-						_, err := conn.Exec(ctx, insertStmt, mode, id, value)
-						require.NoError(t, err)
-					})
+// EncryptedRowWithJsonb represents a row with id, encrypted_text and encrypted_jsonb fields
+type EncryptedRowWithJsonb struct {
+	ID              int                    `db:"id"`
+	EncryptedText   string                 `db:"encrypted_text"`
+	EncryptedJsonb  map[string]interface{} `db:"encrypted_jsonb"`
+}
 
-					t.Run("select", func(t *testing.T) {
-						var rid int
-						var rv time.Time
-						err := conn.QueryRow(context.Background(), selectStmt, mode, id).Scan(&rid, &rv)
-						require.NoError(t, err)
-						require.Equal(t, id, rid)
-						require.Equal(t, value, rv)
-					})
-				})
-			}
+// Scan implements the sql.Scanner interface for EncryptedRowWithJsonb
+func (er *EncryptedRowWithJsonb) Scan(src interface{}) error {
+	return fmt.Errorf("Scan method not implemented for single value")
+}
+
+func TestPgxInsertEncryptedWithJsonbStructScan(t *testing.T) {
+	t.Parallel()
+	conn := setupPgxConnection(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	tx, err := conn.Begin(ctx)
+	require.NoError(t, err)
+	defer tx.Rollback(ctx)
+
+	encryptedTextValue := "test encrypted content"
+	encryptedJsonbValue := `{"key":"value","number":42}`
+	expectedJsonbValue := map[string]interface{}{
+		"key":    "value",
+		"number": float64(42),
+	}
+	insertStmt := `INSERT INTO encrypted (id, encrypted_text, encrypted_jsonb) VALUES ($1, $2, $3) RETURNING id, encrypted_text, encrypted_jsonb`
+
+	for _, mode := range modes {
+		id := rand.Int()
+
+		t.Run(mode.String(), func(t *testing.T) {
+			t.Run("insert_with_jsonb_struct_scan", func(t *testing.T) {
+				var row EncryptedRowWithJsonb
+				err := conn.QueryRow(context.Background(), insertStmt, mode, id, encryptedTextValue, encryptedJsonbValue).Scan(&row.ID, &row.EncryptedText, &row.EncryptedJsonb)
+				require.NoError(t, err)
+				require.Equal(t, id, row.ID)
+				require.Equal(t, encryptedTextValue, row.EncryptedText)
+				require.Equal(t, expectedJsonbValue, row.EncryptedJsonb)
+			})
 		})
 	}
 }
-*/
+