11package pgdump
22
33import (
4- "fmt"
4+ "bytes"
5+ "os"
6+ "os/exec"
57 "path/filepath"
8+ "strings"
69
710 "github.com/sourcegraph/sourcegraph/lib/errors"
811)
@@ -30,24 +33,114 @@ type Target struct {
3033 Password string `yaml:"password"`
3134}
3235
33- // RestoreCommand generates a psql command that can be used for migrations.
34- func RestoreCommand (t Target ) string {
35- dump := fmt .Sprintf ("psql --username=%s --dbname=%s 1>/dev/null" ,
36- t .Username , t .DBName )
37- if t .Password == "" {
38- return dump
36+ // Command represents a command to run as an argv array rather than a shell
37+ // string. Values that originate from user-controlled targets are always passed
38+ // as discrete arguments (or safely shell-quoted when they must be embedded in a
39+ // nested shell command), never interpolated into a string that is handed to
40+ // "bash -c". This prevents shell injection via malicious targets files.
41+ type Command struct {
42+ // Env holds additional environment variables in "KEY=value" form.
43+ Env []string
44+ // Args is the argv to execute; Args[0] is the executable.
45+ Args []string
46+ // InputFile, if set, is opened and connected to the command's stdin.
47+ InputFile string
48+ // OutputFile, if set, receives the command's stdout. When empty, stdout is
49+ // discarded (matching the original "1>/dev/null" behaviour for restores).
50+ OutputFile string
51+ }
52+
53+ // Run executes the command safely (without a shell), wiring up any input/output
54+ // file redirection, and returns the combined stderr (and stdout when it is not
55+ // redirected to a file) output.
56+ func (c Command ) Run () ([]byte , error ) {
57+ if len (c .Args ) == 0 {
58+ return nil , errors .New ("no command to run" )
59+ }
60+
61+ cmd := exec .Command (c .Args [0 ], c .Args [1 :]... )
62+ cmd .Env = append (os .Environ (), c .Env ... )
63+
64+ var combined bytes.Buffer
65+ cmd .Stderr = & combined
66+
67+ if c .InputFile != "" {
68+ f , err := os .Open (c .InputFile )
69+ if err != nil {
70+ return nil , errors .Wrapf (err , "opening input file %q" , c .InputFile )
71+ }
72+ defer f .Close ()
73+ cmd .Stdin = f
74+ }
75+
76+ if c .OutputFile != "" {
77+ f , err := os .Create (c .OutputFile )
78+ if err != nil {
79+ return nil , errors .Wrapf (err , "creating output file %q" , c .OutputFile )
80+ }
81+ defer f .Close ()
82+ cmd .Stdout = f
83+ } else {
84+ // No output file means stdout is not of interest (e.g. restores), so
85+ // capture it alongside stderr for diagnostics.
86+ cmd .Stdout = & combined
87+ }
88+
89+ err := cmd .Run ()
90+ return combined .Bytes (), err
91+ }
92+
93+ // String renders the command as a copy-pasteable, safely shell-quoted string.
94+ // It is intended for display only; execution always goes through Run.
95+ func (c Command ) String () string {
96+ var parts []string
97+ for _ , e := range c .Env {
98+ parts = append (parts , shellQuoteEnv (e ))
99+ }
100+ for _ , a := range c .Args {
101+ parts = append (parts , shellQuote (a ))
102+ }
103+ s := strings .Join (parts , " " )
104+ if c .OutputFile != "" {
105+ s += " > " + shellQuote (c .OutputFile )
39106 }
40- return fmt .Sprintf ("PGPASSWORD=%s %s" , t .Password , dump )
107+ if c .InputFile != "" {
108+ s += " < " + shellQuote (c .InputFile )
109+ }
110+ return s
41111}
42112
43- // DumpCommand generates a pg_dump command that can be used for on-prem-to-Cloud migrations.
44- func DumpCommand (t Target ) string {
45- dump := fmt .Sprintf ("pg_dump --clean --format=plain --if-exists --no-acl --no-owner --quote-all-identifiers --username=%s --dbname=%s" ,
46- t .Username , t .DBName )
47- if t .Password == "" {
48- return dump
113+ // RestoreCommand generates a psql invocation that can be used for migrations.
114+ func RestoreCommand (t Target ) (args []string , env []string ) {
115+ args = []string {
116+ "psql" ,
117+ "--username=" + t .Username ,
118+ "--dbname=" + t .DBName ,
119+ }
120+ if t .Password != "" {
121+ env = append (env , "PGPASSWORD=" + t .Password )
49122 }
50- return fmt .Sprintf ("PGPASSWORD=%s %s" , t .Password , dump )
123+ return args , env
124+ }
125+
126+ // DumpCommand generates a pg_dump invocation that can be used for
127+ // on-prem-to-Cloud migrations.
128+ func DumpCommand (t Target ) (args []string , env []string ) {
129+ args = []string {
130+ "pg_dump" ,
131+ "--clean" ,
132+ "--format=plain" ,
133+ "--if-exists" ,
134+ "--no-acl" ,
135+ "--no-owner" ,
136+ "--quote-all-identifiers" ,
137+ "--username=" + t .Username ,
138+ "--dbname=" + t .DBName ,
139+ }
140+ if t .Password != "" {
141+ env = append (env , "PGPASSWORD=" + t .Password )
142+ }
143+ return args , env
51144}
52145
53146type Output struct {
@@ -70,30 +163,37 @@ func Outputs(dir string, targets Targets) []Output {
70163 }}
71164}
72165
73- type CommandBuilder func (Target ) (string , error )
74- type PGCommand func (Target ) string
166+ type CommandBuilder func (Target ) (Command , error )
167+
168+ // PGCommand generates the argv and environment for a base Postgres command
169+ // (pg_dump or psql) targeting the given Target.
170+ type PGCommand func (Target ) (args []string , env []string )
75171
76172// Builder generates the CommandBuilder and targetKey for a given builder and PGCommand
77173func Builder (builder string , command PGCommand ) (commandBuilder CommandBuilder , targetKey string ) {
78174 switch builder {
79175 case "pg_dump" , "" :
80176 targetKey = "local"
81- commandBuilder = func (t Target ) (string , error ) {
82- cmd := command (t )
177+ commandBuilder = func (t Target ) (Command , error ) {
178+ args , env := command (t )
83179 if t .Target != "" {
84- return fmt . Sprintf ( "%s --host=%s" , cmd , t .Target ), nil
180+ args = append ( args , " --host=" + t .Target )
85181 }
86- return cmd , nil
182+ return Command { Env : env , Args : args } , nil
87183 }
88184 case "docker" :
89185 targetKey = "docker"
90- commandBuilder = func (t Target ) (string , error ) {
91- return fmt .Sprintf ("docker exec -i %s sh -c '%s'" , t .Target , command (t )), nil
186+ commandBuilder = func (t Target ) (Command , error ) {
187+ args , env := command (t )
188+ inner := shellCommand (env , args )
189+ return Command {Args : []string {"docker" , "exec" , "-i" , t .Target , "sh" , "-c" , inner }}, nil
92190 }
93191 case "kubectl" :
94192 targetKey = "k8s"
95- commandBuilder = func (t Target ) (string , error ) {
96- return fmt .Sprintf ("kubectl exec -i %s -- bash -c '%s'" , t .Target , command (t )), nil
193+ commandBuilder = func (t Target ) (Command , error ) {
194+ args , env := command (t )
195+ inner := shellCommand (env , args )
196+ return Command {Args : []string {"kubectl" , "exec" , "-i" , t .Target , "--" , "bash" , "-c" , inner }}, nil
97197 }
98198 default :
99199 return commandBuilder , targetKey
@@ -103,21 +203,70 @@ func Builder(builder string, command PGCommand) (commandBuilder CommandBuilder,
103203
104204// BuildCommands generates commands that output Postgres dumps and sends them to predefined
105205// files for each target database.
106- func BuildCommands (outDir string , commandBuilder CommandBuilder , targets Targets , dump bool ) ([]string , error ) {
107- var commands []string
206+ func BuildCommands (outDir string , commandBuilder CommandBuilder , targets Targets , dump bool ) ([]Command , error ) {
207+ var commands []Command
108208 for _ , t := range Outputs (outDir , targets ) {
109209 c , err := commandBuilder (t .Target )
110210 if err != nil {
111211 return nil , errors .Wrapf (err , "generating command for %q" , t .Output )
112212 }
113213
114214 if dump {
115- // When dumping use output redirection to dump command stdout to target file
116- commands = append ( commands , fmt . Sprintf ( "%s > %s" , c , t .Output ))
215+ // When dumping, redirect command stdout to the target file.
216+ c . OutputFile = t .Output
117217 } else {
118- // When restoring use input redirection to pass target file to command stdin
119- commands = append ( commands , fmt . Sprintf ( "%s < %s" , c , t .Output ))
218+ // When restoring, feed the target file to the command's stdin.
219+ c . InputFile = t .Output
120220 }
221+ commands = append (commands , c )
121222 }
122223 return commands , nil
123224}
225+
226+ // shellCommand renders env and args as a single, safely shell-quoted string for
227+ // use as the argument to a nested "sh -c"/"bash -c" (e.g. inside a container).
228+ func shellCommand (env , args []string ) string {
229+ var parts []string
230+ for _ , e := range env {
231+ parts = append (parts , shellQuoteEnv (e ))
232+ }
233+ for _ , a := range args {
234+ parts = append (parts , shellQuote (a ))
235+ }
236+ return strings .Join (parts , " " )
237+ }
238+
239+ // shellQuoteEnv quotes a "KEY=value" assignment, leaving the KEY= prefix
240+ // unquoted (so the shell still treats it as an assignment) and quoting only the
241+ // value.
242+ func shellQuoteEnv (assignment string ) string {
243+ i := strings .IndexByte (assignment , '=' )
244+ if i < 0 {
245+ return shellQuote (assignment )
246+ }
247+ return assignment [:i + 1 ] + shellQuote (assignment [i + 1 :])
248+ }
249+
250+ // shellQuote returns s quoted such that a POSIX shell will treat it as a single
251+ // literal argument.
252+ func shellQuote (s string ) string {
253+ if s == "" {
254+ return "''"
255+ }
256+ safe := true
257+ for _ , r := range s {
258+ if (r >= 'a' && r <= 'z' ) || (r >= 'A' && r <= 'Z' ) || (r >= '0' && r <= '9' ) {
259+ continue
260+ }
261+ switch r {
262+ case '-' , '_' , '.' , '/' , ':' , '=' , '@' , '%' , '+' , ',' :
263+ continue
264+ }
265+ safe = false
266+ break
267+ }
268+ if safe {
269+ return s
270+ }
271+ return "'" + strings .ReplaceAll (s , "'" , `'\''` ) + "'"
272+ }
0 commit comments