gosh: make {Send,Await}Vars handle "\n" in JSON
Change-Id: Ib7f0778b93d1a97332cb62bde2ff0d2ac2dda66a
diff --git a/gosh/child.go b/gosh/child.go
index 81d0ba8..99562c9 100644
--- a/gosh/child.go
+++ b/gosh/child.go
@@ -14,18 +14,19 @@
"time"
)
-const varsPrefix = "# gosh "
+var (
+ varsPrefix = []byte("<goshVars")
+ varsSuffix = []byte("goshVars>")
+)
-// SendVars sends the given vars to the parent process. Writes a line of the
-// form "# gosh { ... JSON object ... }" to stderr.
+// SendVars sends the given vars to the parent process. Writes a string of the
+// form "<goshVars{ ... JSON-encoded vars ... }goshVars>\n" to stderr.
func SendVars(vars map[string]string) {
data, err := json.Marshal(vars)
if err != nil {
panic(err)
}
- // TODO(sadovsky): Handle the case where the JSON object contains a newline
- // character.
- fmt.Fprintf(os.Stderr, "%s%s\n", varsPrefix, data)
+ fmt.Fprintf(os.Stderr, "%s%s%s\n", varsPrefix, data, varsSuffix)
}
// watchParent periodically checks whether the parent process has exited and, if
diff --git a/gosh/cmd.go b/gosh/cmd.go
index 3d01ce3..e748670 100644
--- a/gosh/cmd.go
+++ b/gosh/cmd.go
@@ -292,40 +292,49 @@
return !c.exited
}
-// recvWriter listens for gosh messages from a child process.
+// recvWriter listens for gosh vars from a child process.
type recvWriter struct {
- c *Cmd
- buf bytes.Buffer
- readPrefix bool // if true, we've read len(varsPrefix) for the current line
- skipLine bool // if true, ignore bytes until next '\n'
+ c *Cmd
+ buf []byte
+ matchedPrefix int
+ matchedSuffix int
}
func (w *recvWriter) Write(p []byte) (n int, err error) {
- for _, b := range p {
- if b == '\n' {
- if w.readPrefix && !w.skipLine {
- vars := make(map[string]string)
- if err := json.Unmarshal(w.buf.Bytes(), &vars); err != nil {
- return 0, err
- }
- w.c.cond.L.Lock()
- w.c.recvVars = mergeMaps(w.c.recvVars, vars)
- w.c.cond.Signal()
- w.c.cond.L.Unlock()
+ for i, b := range p {
+ if w.matchedPrefix < len(varsPrefix) {
+ // Look for matching prefix.
+ if b != varsPrefix[w.matchedPrefix] {
+ w.matchedPrefix = 0
}
- // Reset state for next line.
- w.readPrefix, w.skipLine = false, false
- w.buf.Reset()
- } else if !w.skipLine {
- w.buf.WriteByte(b)
- if !w.readPrefix && w.buf.Len() == len(varsPrefix) {
- w.readPrefix = true
- prefix := string(w.buf.Next(len(varsPrefix)))
- if prefix != varsPrefix {
- w.skipLine = true
- }
+ if b == varsPrefix[w.matchedPrefix] {
+ w.matchedPrefix++
}
+ continue
}
+ w.buf = append(w.buf, b)
+ // Look for matching suffix.
+ if b != varsSuffix[w.matchedSuffix] {
+ w.matchedSuffix = 0
+ }
+ if b == varsSuffix[w.matchedSuffix] {
+ w.matchedSuffix++
+ }
+ if w.matchedSuffix != len(varsSuffix) {
+ continue
+ }
+ // Found matching suffix.
+ data := w.buf[:len(w.buf)-len(varsSuffix)]
+ w.buf = w.buf[:0]
+ w.matchedPrefix, w.matchedSuffix = 0, 0
+ vars := make(map[string]string)
+ if err := json.Unmarshal(data, &vars); err != nil {
+ return i, err
+ }
+ w.c.cond.L.Lock()
+ w.c.recvVars = mergeMaps(w.c.recvVars, vars)
+ w.c.cond.Signal()
+ w.c.cond.L.Unlock()
}
return len(p), nil
}
diff --git a/gosh/shell_test.go b/gosh/shell_test.go
index 56aa87e..c497498 100644
--- a/gosh/shell_test.go
+++ b/gosh/shell_test.go
@@ -238,6 +238,72 @@
eq(t, c.Stdout(), "Hello, world!\n")
}
+var (
+ sendVarsFunc = gosh.RegisterFunc("sendVarsFunc", func(vars map[string]string) {
+ gosh.SendVars(vars)
+ time.Sleep(time.Hour)
+ })
+ stderrFunc = gosh.RegisterFunc("stderrFunc", func(s string) {
+ fmt.Fprintf(os.Stderr, s)
+ time.Sleep(time.Hour)
+ })
+)
+
+// Tests that AwaitVars works under various conditions.
+func TestAwaitVars(t *testing.T) {
+ sh := gosh.NewShell(gosh.Opts{Fatalf: makeFatalf(t), Logf: t.Logf})
+ defer sh.Cleanup()
+
+ c := sh.FuncCmd(sendVarsFunc, map[string]string{"a": "1"})
+ c.Start()
+ eq(t, c.AwaitVars("a")["a"], "1")
+
+ c = sh.FuncCmd(stderrFunc, `<goshVars{"a":"1","b":"2"}goshVars>`)
+ c.Start()
+ vars := c.AwaitVars("a", "b")
+ eq(t, vars["a"], "1")
+ eq(t, vars["b"], "2")
+
+ c = sh.FuncCmd(stderrFunc, `<goshVars{"a":"1"}goshVars><gosh`)
+ c.Start()
+ eq(t, c.AwaitVars("a")["a"], "1")
+
+ c = sh.FuncCmd(stderrFunc, `<goshVars{"a":"1"}goshVars><goshVars{"b":"2"}goshVars>`)
+ c.Start()
+ vars = c.AwaitVars("a", "b")
+ eq(t, vars["a"], "1")
+ eq(t, vars["b"], "2")
+
+ c = sh.FuncCmd(stderrFunc, `<goshVars{"a":"1","b":"2"}goshVars>`)
+ c.Start()
+ vars = c.AwaitVars("a")
+ eq(t, vars["a"], "1")
+ eq(t, vars["b"], "")
+ vars = c.AwaitVars("b")
+ eq(t, vars["a"], "")
+ eq(t, vars["b"], "2")
+
+ c = sh.FuncCmd(stderrFunc, `<g<goshVars{"a":"goshVars"}goshVars>s><goshVars`)
+ c.Start()
+ eq(t, c.AwaitVars("a")["a"], "goshVars")
+
+ c = sh.FuncCmd(stderrFunc, `<<goshVars{"a":"1"}goshVars>><<goshVars{"b":"<goshVars"}goshVars>>`)
+ c.Start()
+ vars = c.AwaitVars("a", "b")
+ eq(t, vars["a"], "1")
+ eq(t, vars["b"], "<goshVars")
+}
+
+// Tests that AwaitVars returns immediately when the process exits.
+func TestAwaitProcessExit(t *testing.T) {
+ sh := gosh.NewShell(gosh.Opts{Fatalf: makeFatalf(t), Logf: t.Logf})
+ defer sh.Cleanup()
+
+ c := sh.FuncCmd(exitFunc, 0)
+ c.Start()
+ setsErr(t, sh, func() { c.AwaitVars("foo") })
+}
+
// Functions designed for TestRegistry.
var (
printIntsFunc = gosh.RegisterFunc("printIntsFunc", func(v ...int) {
@@ -256,16 +322,6 @@
})
)
-// Tests that AwaitVars returns immediately when the process exits.
-func TestAwaitProcessExit(t *testing.T) {
- sh := gosh.NewShell(gosh.Opts{Fatalf: makeFatalf(t), Logf: t.Logf})
- defer sh.Cleanup()
-
- c := sh.FuncCmd(exitFunc, 0)
- c.Start()
- setsErr(t, sh, func() { c.AwaitVars("foo") })
-}
-
// Tests function signature-checking and execution.
func TestRegistry(t *testing.T) {
sh := gosh.NewShell(gosh.Opts{Fatalf: makeFatalf(t), Logf: t.Logf})
