Skip to content

Commit

Permalink
Feat: fifo control
Browse files Browse the repository at this point in the history
  • Loading branch information
zijiren233 committed Mar 9, 2024
1 parent d9a06af commit 16e7a60
Show file tree
Hide file tree
Showing 6 changed files with 251 additions and 398 deletions.
52 changes: 52 additions & 0 deletions cmd/start/control.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
//go:build !windows
// +build !windows

package start

import (
"fmt"
"os"
"syscall"
"time"

pty "github.com/MCSManager/pty/console"
)

func runControl(fifo string, con pty.Console) error {
err := os.Remove(fifo)
if err != nil {
if !os.IsNotExist(err) {
return fmt.Errorf("remove fifo error: %w", err)
}
}
if err := syscall.Mkfifo(fifo, 0666); err != nil {
return fmt.Errorf("create fifo error: %w", err)
}

if testFifoResize {
go func() {
time.Sleep(time.Second * 5)
_ = testUnixResize(fifo)
}()
}

for {
f, err := os.OpenFile(fifo, os.O_RDONLY, os.ModeNamedPipe)
if err != nil {
return fmt.Errorf("open fifo error: %w", err)
}
defer f.Close()
u := newConnUtils(f, f)
_ = handleConn(u, con)
}
}

func testUnixResize(fifo string) error {
n, err := os.OpenFile(fifo, os.O_WRONLY, os.ModeNamedPipe)
if err != nil {
return fmt.Errorf("open fifo error: %w", err)
}
defer n.Close()
u := newConnUtils(n, n)
return testResize(u)
}
86 changes: 86 additions & 0 deletions cmd/start/control_common.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
package start

import (
"encoding/json"
"fmt"
"io"

pty "github.com/MCSManager/pty/console"
"github.com/zijiren233/stream"
)

type connUtils struct {
r *stream.Reader
w *stream.Writer
}

func newConnUtils(r io.Reader, w io.Writer) *connUtils {
return &connUtils{
r: stream.NewReader(r, stream.BigEndian),
w: stream.NewWriter(w, stream.BigEndian),
}
}

func (cu *connUtils) ReadMessage() (uint8, []byte, error) {
var (
length uint16
msgType uint8
)
data, err := cu.r.U8(&msgType).U16(&length).ReadBytes(int(length))
return msgType, data, err
}

func (cu *connUtils) SendMessage(msgType uint8, data any) error {
b, err := json.Marshal(data)
if err != nil {
return err
}
return cu.w.U8(msgType).U16(uint16(len(b))).Bytes(b).Error()
}

func handleConn(u *connUtils, con pty.Console) error {
for {
t, msg, err := u.ReadMessage()
if err != nil {
return fmt.Errorf("read message error: %w", err)
}
switch t {
case RESIZE:
resize := resizeMsg{}
err := json.Unmarshal(msg, &resize)
if err != nil {
_ = u.SendMessage(
ERROR,
&errorMsg{
Msg: fmt.Sprintf("unmarshal resize message error: %s", err),
},
)
continue
}
err = con.SetSize(resize.Width, resize.Height)
if err != nil {
_ = u.SendMessage(
ERROR,
&errorMsg{
Msg: fmt.Sprintf("resize error: %s", err),
},
)
continue
}
}
}
}

func testResize(u *connUtils) error {
err := u.SendMessage(
RESIZE,
&resizeMsg{
Width: 20,
Height: 20,
},
)
if err != nil {
return fmt.Errorf("send resize message error: %w", err)
}
return nil
}
47 changes: 47 additions & 0 deletions cmd/start/control_windows.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
package start

import (
"fmt"
"time"

pty "github.com/MCSManager/pty/console"

winio "github.com/Microsoft/go-winio"
)

// \\.\pipe\mypipe
func runControl(fifo string, con pty.Console) error {
n, err := winio.ListenPipe(fifo, &winio.PipeConfig{})
if err != nil {
return fmt.Errorf("open fifo error: %w", err)
}
defer n.Close()

if testFifoResize {
go func() {
time.Sleep(time.Second * 5)
_ = testWinResize(fifo)
}()
}

for {
conn, err := n.Accept()
if err != nil {
return fmt.Errorf("accept fifo error: %w", err)
}
go func() {
defer conn.Close()
u := newConnUtils(conn, conn)
_ = handleConn(u, con)
}()
}
}

func testWinResize(fifo string) error {
n, err := winio.DialPipe(fifo, nil)
if err != nil {
return fmt.Errorf("open fifo error: %w", err)
}
u := newConnUtils(n, n)
return testResize(u)
}
71 changes: 51 additions & 20 deletions cmd/start/start.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ import (
var (
dir, cmd, coder, ptySize string
cmds []string
fifo string
testFifoResize bool
)

type PtyInfo struct {
Expand All @@ -33,54 +35,83 @@ func init() {
flag.StringVar(&coder, "coder", "auto", "Coder")
flag.StringVar(&dir, "dir", ".", "command work path")
flag.StringVar(&ptySize, "size", "80,50", "Initialize pty size, stdin will be forwarded directly")
flag.StringVar(&fifo, "fifo", "", "control FIFO name")
flag.BoolVar(&testFifoResize, "test-fifo-resize", false, "test fifo resize")
}

func Main() {
flag.Parse()
runPTY()
}

func runPTY() {
if err := json.Unmarshal([]byte(cmd), &cmds); err != nil {
fmt.Println("[MCSMANAGER-PTY] Unmarshal command error: ", err)
con, err := newPTY()
if err != nil {
fmt.Printf("[MCSMANAGER-PTY] New pty error: %v\n", err)
return
}
con := pty.New(utils.CoderToType(coder))
if err := con.ResizeWithString(ptySize); err != nil {
fmt.Printf("[MCSMANAGER-PTY] PTY Resize error: %v\n", err)
err = con.Start(dir, cmds)
if err != nil {
fmt.Printf("[MCSMANAGER-PTY] Process start error: %v\n", err)
return
}
err := con.Start(dir, cmds)
info, _ := json.Marshal(&PtyInfo{
Pid: con.Pid(),
})
fmt.Println(string(info))
if err != nil {
fmt.Printf("[MCSMANAGER-PTY] Process start error: %v\n", err)
return
}
defer con.Close()
handleStdIO(con)
if fifo != "" {
go func() {
err := runControl(fifo, con)
if err != nil {
fmt.Println("[MCSMANAGER-PTY] Control error: ", err)
}
}()
}
if err = handleStdIO(con); err != nil {
fmt.Println("[MCSMANAGER-PTY] Handle stdio error: ", err)
}
_, _ = con.Wait()
}

func handleStdIO(c pty.Console) {
func newPTY() (pty.Console, error) {
if err := json.Unmarshal([]byte(cmd), &cmds); err != nil {
return nil, fmt.Errorf("unmarshal command error: %w", err)
}
con := pty.New(utils.CoderToType(coder))
if err := con.ResizeWithString(ptySize); err != nil {
return nil, fmt.Errorf("pty resize error: %w", err)
}
return con, nil
}

func handleStdIO(c pty.Console) error {
if colorable.IsReaderTerminal(os.Stdin) {
oldState, err := term.MakeRaw(int(os.Stdin.Fd()))
if err != nil {
panic(err)
return fmt.Errorf("make raw error: %w", err)
}
defer func() { _ = term.Restore(int(os.Stdin.Fd()), oldState) }()
go func() { _, _ = io.Copy(c.StdIn(), os.Stdin) }()
} else {
go func() { _, _ = io.Copy(c.StdIn(), os.Stdin) }()
}
go func() { _, _ = io.Copy(c.StdIn(), os.Stdin) }()
if runtime.GOOS == "windows" && c.StdErr() != nil {
go func() { _, _ = io.Copy(colorable.NewColorableStderr(), c.StdErr()) }()
}
handleStdOut(c)
return nil
}

func handleStdOut(c pty.Console) {
_, _ = io.Copy(colorable.NewColorableStdout(), c.StdOut())
}

const (
ERROR uint8 = iota + 2
PING
RESIZE
)

type errorMsg struct {
Msg string `json:"msg"`
}

type resizeMsg struct {
Width uint `json:"width"`
Height uint `json:"height"`
}
30 changes: 5 additions & 25 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,40 +3,20 @@ module github.com/MCSManager/pty
go 1.18

require (
github.com/Microsoft/go-winio v0.6.1
github.com/creack/pty v1.1.21
github.com/juju/mutex/v2 v2.0.0
github.com/klauspost/compress v1.16.5
github.com/mholt/archiver/v4 v4.0.0-alpha.8
github.com/shirou/gopsutil/v3 v3.23.4
github.com/zijiren233/go-colorable v0.0.0-20230930131441-997304c961cb
github.com/zijiren233/stream v0.5.2
golang.org/x/term v0.18.0
golang.org/x/text v0.14.0
)

require (
github.com/andybalholm/brotli v1.0.5 // indirect
github.com/bodgit/plumbing v1.3.0 // indirect
github.com/bodgit/sevenzip v1.4.1 // indirect
github.com/bodgit/windows v1.0.1 // indirect
github.com/connesc/cipherio v0.2.1 // indirect
github.com/dsnet/compress v0.0.1 // indirect
github.com/go-ole/go-ole v1.2.6 // indirect
github.com/golang/snappy v0.0.4 // indirect
github.com/hashicorp/errwrap v1.1.0 // indirect
github.com/hashicorp/go-multierror v1.1.1 // indirect
github.com/juju/errors v1.0.0 // indirect
github.com/klauspost/pgzip v1.2.6 // indirect
github.com/lufia/plan9stats v0.0.0-20230326075908-cb1d2100619a // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/nwaples/rardecode/v2 v2.0.0-beta.2 // indirect
github.com/pierrec/lz4/v4 v4.1.17 // indirect
github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b // indirect
github.com/shoenig/go-m1cpu v0.1.6 // indirect
github.com/therootcompany/xz v1.0.1 // indirect
github.com/tklauser/go-sysconf v0.3.11 // indirect
github.com/tklauser/numcpus v0.6.0 // indirect
github.com/ulikunitz/xz v0.5.11 // indirect
github.com/yusufpapurcu/wmi v1.2.3 // indirect
go4.org v0.0.0-20230225012048-214862532bf5 // indirect
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 // indirect
golang.org/x/mod v0.10.0 // indirect
golang.org/x/sys v0.18.0 // indirect
golang.org/x/tools v0.9.1 // indirect
)
Loading

0 comments on commit 16e7a60

Please sign in to comment.