mirror of
https://github.com/openziti/zrok.git
synced 2025-01-11 16:38:22 +01:00
rough proctree implementation (#748)
This commit is contained in:
parent
82d8f4ba2e
commit
42c3ec48b0
agent/proctree
59
agent/proctree/impl_posix.go
Normal file
59
agent/proctree/impl_posix.go
Normal file
@ -0,0 +1,59 @@
|
||||
//go:build !windows
|
||||
|
||||
package proctree
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
"sync"
|
||||
)
|
||||
|
||||
func Init(_ string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func StartChild(tail TailFunction, args ...string) (*Child, error) {
|
||||
cmd := exec.Command(args[0], args[1:]...)
|
||||
|
||||
cld := &Child{
|
||||
TailFunction: tail,
|
||||
cmd: cmd,
|
||||
outStream: make(chan []byte),
|
||||
errStream: make(chan []byte),
|
||||
wg: new(sync.WaitGroup),
|
||||
}
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stderr, err := cmd.StderrPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cld.wg.Add(3)
|
||||
go reader(stdout, cld.outStream, cld.wg)
|
||||
go reader(stderr, cld.errStream, cld.wg)
|
||||
go cld.combiner(cld.wg)
|
||||
|
||||
return cld, nil
|
||||
}
|
||||
|
||||
func WaitChild(c *Child) error {
|
||||
c.wg.Wait()
|
||||
if err := c.cmd.Wait(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func StopChild(c *Child) error {
|
||||
if err := c.cmd.Process.Kill(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
79
agent/proctree/impl_windows.go
Executable file
79
agent/proctree/impl_windows.go
Executable file
@ -0,0 +1,79 @@
|
||||
//go:build windows
|
||||
|
||||
package proctree
|
||||
|
||||
import (
|
||||
"github.com/kolesnikovae/go-winjob"
|
||||
"golang.org/x/sys/windows"
|
||||
"os/exec"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var job *winjob.JobObject
|
||||
|
||||
func Init(name string) error {
|
||||
var err error
|
||||
if job == nil {
|
||||
job, err = winjob.Create(name, winjob.LimitKillOnJobClose, winjob.LimitBreakawayOK)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func StartChild(tail TailFunction, args ...string) (*Child, error) {
|
||||
cmd := exec.Command(args[0], args[1:]...)
|
||||
cmd.SysProcAttr = &windows.SysProcAttr{CreationFlags: windows.CREATE_SUSPENDED}
|
||||
|
||||
cld := &Child{
|
||||
TailFunction: tail,
|
||||
cmd: cmd,
|
||||
outStream: make(chan []byte),
|
||||
errStream: make(chan []byte),
|
||||
wg: new(sync.WaitGroup),
|
||||
}
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stderr, err := cmd.StderrPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := job.Assign(cmd.Process); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := winjob.ResumeProcess(cmd.Process.Pid); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cld.wg.Add(3)
|
||||
go reader(stdout, cld.outStream, cld.wg)
|
||||
go reader(stderr, cld.errStream, cld.wg)
|
||||
go cld.combiner(cld.wg)
|
||||
|
||||
return cld, nil
|
||||
}
|
||||
|
||||
func WaitChild(c *Child) error {
|
||||
c.wg.Wait()
|
||||
if err := c.cmd.Wait(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func StopChild(c *Child) error {
|
||||
if err := c.cmd.Process.Kill(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
66
agent/proctree/proctree.go
Executable file
66
agent/proctree/proctree.go
Executable file
@ -0,0 +1,66 @@
|
||||
package proctree
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os/exec"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type Child struct {
|
||||
TailFunction TailFunction
|
||||
cmd *exec.Cmd
|
||||
outStream chan []byte
|
||||
errStream chan []byte
|
||||
wg *sync.WaitGroup
|
||||
}
|
||||
|
||||
type TailFunction func(data []byte)
|
||||
|
||||
func (c *Child) combiner(wg *sync.WaitGroup) {
|
||||
defer wg.Done()
|
||||
|
||||
outDone := false
|
||||
errDone := false
|
||||
for {
|
||||
select {
|
||||
case data := <-c.outStream:
|
||||
if data != nil {
|
||||
if c.TailFunction != nil {
|
||||
c.TailFunction(data)
|
||||
}
|
||||
} else {
|
||||
outDone = true
|
||||
}
|
||||
case data := <-c.errStream:
|
||||
if data != nil {
|
||||
if c.TailFunction != nil {
|
||||
c.TailFunction(data)
|
||||
}
|
||||
} else {
|
||||
errDone = true
|
||||
}
|
||||
}
|
||||
if outDone && errDone {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func reader(r io.ReadCloser, o chan []byte, wg *sync.WaitGroup) {
|
||||
defer close(o)
|
||||
defer wg.Done()
|
||||
|
||||
buf := make([]byte, 64*1024)
|
||||
for {
|
||||
n, err := r.Read(buf)
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
return
|
||||
}
|
||||
fmt.Printf("error reading: %v", err)
|
||||
return
|
||||
}
|
||||
o <- buf[:n]
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue
Block a user