Skip to content

Commit a8c753e

Browse files
Amogh-BharadwajAmogh-Bharadwaj
committed
fix(pg-pg): better tls wiring for pgdump command
1 parent 04a87e8 commit a8c753e

2 files changed

Lines changed: 299 additions & 21 deletions

File tree

‎flow/connectors/postgres/pgdump_schema.go‎

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -281,29 +281,39 @@ func buildPsqlArgs(config *protos.PostgresConfig) []string {
281281
}
282282

283283
func appendTLSEnv(ctx context.Context, cmd *exec.Cmd, config *protos.PostgresConfig) {
284-
if internal.PGMustUseTlsConnection(config) {
285-
sslMode := "verify-ca"
286-
if config.SkipCertVerification {
287-
sslMode = "require"
288-
}
289-
cmd.Env = append(cmd.Env, "PGSSLMODE="+sslMode)
284+
hasRootCA := config.RootCa != nil && *config.RootCa != ""
285+
// Match the regular connector: enable TLS when explicitly required OR when
286+
// a root CA is provided (which implies the user wants a verified connection).
287+
if !internal.PGMustUseTlsConnection(config) && !hasRootCA {
288+
return
289+
}
290290

291-
if config.RootCa != nil && *config.RootCa != "" {
292-
// write root CA to a temp file
293-
tmpFile, err := os.CreateTemp("", "peerdb-root-ca-*.pem")
294-
if err != nil {
295-
slog.WarnContext(ctx, "failed to create temp file for root CA, skipping sslrootcert", slog.Any("error", err))
296-
return
297-
}
298-
if _, err := tmpFile.WriteString(*config.RootCa); err != nil {
299-
slog.WarnContext(ctx, "failed to write root CA to temp file", slog.Any("error", err))
300-
tmpFile.Close()
301-
os.Remove(tmpFile.Name())
302-
return
303-
}
291+
switch {
292+
case config.SkipCertVerification:
293+
cmd.Env = append(cmd.Env, "PGSSLMODE=require")
294+
case hasRootCA:
295+
cmd.Env = append(cmd.Env, "PGSSLMODE=verify-ca")
296+
default:
297+
// TLS required but no CA provided — encrypt the connection without
298+
// attempting certificate verification (matches pgx "require" behaviour).
299+
cmd.Env = append(cmd.Env, "PGSSLMODE=require")
300+
}
301+
302+
if hasRootCA {
303+
// write root CA to a temp file
304+
tmpFile, err := os.CreateTemp("", "peerdb-root-ca-*.pem")
305+
if err != nil {
306+
slog.WarnContext(ctx, "failed to create temp file for root CA, skipping sslrootcert", slog.Any("error", err))
307+
return
308+
}
309+
if _, err := tmpFile.WriteString(*config.RootCa); err != nil {
310+
slog.WarnContext(ctx, "failed to write root CA to temp file", slog.Any("error", err))
304311
tmpFile.Close()
305-
cmd.Env = append(cmd.Env, "PGSSLROOTCERT="+tmpFile.Name())
306-
// note: temp file is cleaned up when the process exits
312+
os.Remove(tmpFile.Name())
313+
return
307314
}
315+
tmpFile.Close()
316+
cmd.Env = append(cmd.Env, "PGSSLROOTCERT="+tmpFile.Name())
317+
// note: temp file is cleaned up when the process exits
308318
}
309319
}

‎flow/connectors/postgres/pgdump_schema_test.go‎

Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,15 @@ import (
44
"bytes"
55
"context"
66
"errors"
7+
"os"
78
"os/exec"
89
"runtime"
10+
"slices"
911
"strings"
1012
"testing"
1113
"time"
14+
15+
"github.com/PeerDB-io/peerdb/flow/generated/protos"
1216
)
1317

1418
// requireUnix skips the test on platforms without the shell utilities used here.
@@ -268,3 +272,267 @@ func TestRunPipeline_FilterStripsLines(t *testing.T) {
268272
t.Fatalf("filtered output = %q, want %q", got, want)
269273
}
270274
}
275+
276+
func strPtr(s string) *string { return &s }
277+
func boolPtr(b bool) *bool { return &b }
278+
279+
func TestBuildPgDumpArgs(t *testing.T) {
280+
tests := []struct {
281+
name string
282+
config *protos.PostgresConfig
283+
wantArgs []string
284+
}{
285+
{
286+
name: "basic config",
287+
config: &protos.PostgresConfig{
288+
Host: "src-host",
289+
Port: 5432,
290+
Database: "mydb",
291+
User: "admin",
292+
},
293+
wantArgs: []string{
294+
"--schema-only", "--no-owner", "--no-privileges",
295+
"-h", "src-host", "-p", "5432", "-d", "mydb", "-U", "admin",
296+
},
297+
},
298+
{
299+
name: "default port when zero",
300+
config: &protos.PostgresConfig{
301+
Host: "localhost",
302+
Port: 0,
303+
Database: "test",
304+
},
305+
wantArgs: []string{
306+
"--schema-only", "--no-owner", "--no-privileges",
307+
"-h", "localhost", "-p", "5432", "-d", "test",
308+
},
309+
},
310+
{
311+
name: "custom port no user",
312+
config: &protos.PostgresConfig{
313+
Host: "remote",
314+
Port: 6543,
315+
Database: "prod",
316+
},
317+
wantArgs: []string{
318+
"--schema-only", "--no-owner", "--no-privileges",
319+
"-h", "remote", "-p", "6543", "-d", "prod",
320+
},
321+
},
322+
}
323+
324+
for _, tt := range tests {
325+
t.Run(tt.name, func(t *testing.T) {
326+
got := buildPgDumpArgs(tt.config)
327+
if !slices.Equal(got, tt.wantArgs) {
328+
t.Fatalf("buildPgDumpArgs = %v, want %v", got, tt.wantArgs)
329+
}
330+
})
331+
}
332+
}
333+
334+
func TestBuildPsqlArgs(t *testing.T) {
335+
tests := []struct {
336+
name string
337+
config *protos.PostgresConfig
338+
wantArgs []string
339+
}{
340+
{
341+
name: "with user",
342+
config: &protos.PostgresConfig{
343+
Host: "dst-host",
344+
Port: 5433,
345+
Database: "target",
346+
User: "writer",
347+
},
348+
wantArgs: []string{
349+
"-h", "dst-host", "-p", "5433", "-d", "target",
350+
"--single-transaction",
351+
"-v", "ON_ERROR_STOP=1",
352+
"--quiet",
353+
"-U", "writer",
354+
},
355+
},
356+
{
357+
name: "default port no user",
358+
config: &protos.PostgresConfig{
359+
Host: "localhost",
360+
Port: 0,
361+
Database: "db",
362+
},
363+
wantArgs: []string{
364+
"-h", "localhost", "-p", "5432", "-d", "db",
365+
"--single-transaction",
366+
"-v", "ON_ERROR_STOP=1",
367+
"--quiet",
368+
},
369+
},
370+
}
371+
372+
for _, tt := range tests {
373+
t.Run(tt.name, func(t *testing.T) {
374+
got := buildPsqlArgs(tt.config)
375+
if !slices.Equal(got, tt.wantArgs) {
376+
t.Fatalf("buildPsqlArgs = %v, want %v", got, tt.wantArgs)
377+
}
378+
})
379+
}
380+
}
381+
382+
// envValue returns the last value for key in cmd.Env, or "" if absent.
383+
func envValue(cmd *exec.Cmd, key string) string {
384+
prefix := key + "="
385+
for i := len(cmd.Env) - 1; i >= 0; i-- {
386+
if strings.HasPrefix(cmd.Env[i], prefix) {
387+
return strings.TrimPrefix(cmd.Env[i], prefix)
388+
}
389+
}
390+
return ""
391+
}
392+
393+
func TestAppendTLSEnv(t *testing.T) {
394+
ctx := t.Context()
395+
396+
tests := []struct {
397+
name string
398+
config *protos.PostgresConfig
399+
wantSSLMode string
400+
wantRootCertSet bool
401+
}{
402+
{
403+
name: "no TLS when not required and no root CA",
404+
config: &protos.PostgresConfig{
405+
Host: "h", Port: 5432, Database: "d",
406+
},
407+
wantSSLMode: "",
408+
},
409+
{
410+
name: "require TLS via RequireTls flag",
411+
config: &protos.PostgresConfig{
412+
Host: "h", Port: 5432, Database: "d",
413+
RequireTls: true,
414+
},
415+
wantSSLMode: "require",
416+
},
417+
{
418+
name: "require TLS via DisableTls=false",
419+
config: &protos.PostgresConfig{
420+
Host: "h", Port: 5432, Database: "d",
421+
DisableTls: boolPtr(false),
422+
},
423+
wantSSLMode: "require",
424+
},
425+
{
426+
name: "no TLS when DisableTls=true",
427+
config: &protos.PostgresConfig{
428+
Host: "h", Port: 5432, Database: "d",
429+
DisableTls: boolPtr(true),
430+
},
431+
wantSSLMode: "",
432+
},
433+
{
434+
name: "require with skip cert verification",
435+
config: &protos.PostgresConfig{
436+
Host: "h", Port: 5432, Database: "d",
437+
RequireTls: true,
438+
SkipCertVerification: true,
439+
},
440+
wantSSLMode: "require",
441+
},
442+
{
443+
name: "verify-ca when root CA provided with RequireTls",
444+
config: &protos.PostgresConfig{
445+
Host: "h", Port: 5432, Database: "d",
446+
RequireTls: true,
447+
RootCa: strPtr("-----BEGIN CERTIFICATE-----\ntest\n-----END CERTIFICATE-----"),
448+
},
449+
wantSSLMode: "verify-ca",
450+
wantRootCertSet: true,
451+
},
452+
{
453+
name: "root CA alone triggers TLS with verify-ca",
454+
config: &protos.PostgresConfig{
455+
Host: "h", Port: 5432, Database: "d",
456+
RootCa: strPtr("-----BEGIN CERTIFICATE-----\ntest\n-----END CERTIFICATE-----"),
457+
},
458+
wantSSLMode: "verify-ca",
459+
wantRootCertSet: true,
460+
},
461+
{
462+
name: "skip cert verification with root CA uses require",
463+
config: &protos.PostgresConfig{
464+
Host: "h", Port: 5432, Database: "d",
465+
RequireTls: true,
466+
SkipCertVerification: true,
467+
RootCa: strPtr("-----BEGIN CERTIFICATE-----\ntest\n-----END CERTIFICATE-----"),
468+
},
469+
wantSSLMode: "require",
470+
wantRootCertSet: false, // skip_cert_verification takes precedence, no sslrootcert needed
471+
},
472+
{
473+
name: "empty root CA string is treated as absent",
474+
config: &protos.PostgresConfig{
475+
Host: "h", Port: 5432, Database: "d",
476+
RequireTls: true,
477+
RootCa: strPtr(""),
478+
},
479+
wantSSLMode: "require",
480+
wantRootCertSet: false,
481+
},
482+
}
483+
484+
for _, tt := range tests {
485+
t.Run(tt.name, func(t *testing.T) {
486+
cmd := exec.CommandContext(ctx, "echo") //nolint:gosec
487+
cmd.Env = os.Environ()
488+
489+
appendTLSEnv(ctx, cmd, tt.config)
490+
491+
gotSSLMode := envValue(cmd, "PGSSLMODE")
492+
if gotSSLMode != tt.wantSSLMode {
493+
t.Errorf("PGSSLMODE = %q, want %q", gotSSLMode, tt.wantSSLMode)
494+
}
495+
496+
gotRootCert := envValue(cmd, "PGSSLROOTCERT")
497+
if tt.wantRootCertSet {
498+
if gotRootCert == "" {
499+
t.Error("expected PGSSLROOTCERT to be set, but it was empty")
500+
} else {
501+
// Verify the temp file was actually written with the CA content.
502+
content, err := os.ReadFile(gotRootCert)
503+
if err != nil {
504+
t.Fatalf("failed to read PGSSLROOTCERT file %q: %v", gotRootCert, err)
505+
}
506+
if !strings.Contains(string(content), "BEGIN CERTIFICATE") {
507+
t.Errorf("PGSSLROOTCERT file content = %q, expected PEM data", string(content))
508+
}
509+
os.Remove(gotRootCert)
510+
}
511+
} else if gotRootCert != "" {
512+
t.Errorf("expected PGSSLROOTCERT to be empty, got %q", gotRootCert)
513+
}
514+
})
515+
}
516+
}
517+
518+
func TestIncompatibleLineRegex(t *testing.T) {
519+
tests := []struct {
520+
line string
521+
match bool
522+
}{
523+
{"SET transaction_timeout = 0;\n", true},
524+
{"SET transaction_timeout=0;\n", true},
525+
{"SET statement_timeout = 0;\n", false},
526+
{"\\restrict abc123\n", true},
527+
{"\\unrestrict abc123\n", true},
528+
{"\\restrict\n", true},
529+
{"CREATE TABLE t(id int);\n", false},
530+
{"SELECT 1;\n", false},
531+
}
532+
for _, tt := range tests {
533+
got := incompatibleLineRE.MatchString(tt.line)
534+
if got != tt.match {
535+
t.Errorf("incompatibleLineRE.Match(%q) = %v, want %v", tt.line, got, tt.match)
536+
}
537+
}
538+
}

0 commit comments

Comments
 (0)