1
1
mirror of https://github.com/openziti/zrok.git synced 2025-01-11 16:38:22 +01:00

rough proctree implementation ()

This commit is contained in:
Michael Quigley 2024-09-12 12:44:54 -04:00
parent 82d8f4ba2e
commit 42c3ec48b0
No known key found for this signature in database
GPG Key ID: 9B60314A9DD20A62
3 changed files with 204 additions and 0 deletions

View 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
View 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
View 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]
}
}