feat(ipc): support exec command in restful api (#52)

Signed-off-by: Kevin Cui <[email protected]>

Author: BlackHole1

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/pkg/errors v0.9.1
 	github.com/prashantgupta24/mac-sleep-notifier v1.0.1
 	github.com/shirou/gopsutil/v3 v3.23.12
+	golang.org/x/crypto v0.18.0
 	golang.org/x/net v0.20.0
 	golang.org/x/sync v0.5.0
 	inet.af/tcpproxy v0.0.0-20221017015627-91f861402626
@@ -20,7 +21,7 @@ require (
 	github.com/Microsoft/go-winio v0.6.1 // indirect
 	github.com/apparentlymart/go-cidr v1.1.0 // indirect
 	github.com/go-ole/go-ole v1.3.0 // indirect
-	github.com/google/btree v1.0.1 // indirect
+	github.com/google/btree v1.1.2 // indirect
 	github.com/google/gopacket v1.1.19 // indirect
 	github.com/insomniacslk/dhcp v0.0.0-20220504074936-1ca156eafb9f // indirect
 	github.com/linuxkit/virtsock v0.0.0-20220523201153-1a23e78aa7a2 // indirect
@@ -35,10 +36,11 @@ require (
 	github.com/tklauser/numcpus v0.7.0 // indirect
 	github.com/u-root/uio v0.0.0-20210528114334-82958018845c // indirect
 	github.com/yusufpapurcu/wmi v1.2.3 // indirect
-	golang.org/x/crypto v0.18.0 // indirect
 	golang.org/x/mod v0.13.0 // indirect
 	golang.org/x/sys v0.16.0 // indirect
 	golang.org/x/time v0.5.0 // indirect
-	golang.org/x/tools v0.13.0 // indirect
+	golang.org/x/tools v0.14.0 // indirect
 	gvisor.dev/gvisor v0.0.0-20230715022000-fd277b20b8db // indirect
 )
+
+replace gvisor.dev/gvisor v0.0.0-20230715022000-fd277b20b8db => gvisor.dev/gvisor v0.0.0-20231023213702-2691a8f9b1cf

go.sum

@@ -20,8 +20,8 @@ github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4
 github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=
 github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE=
 github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78=
-github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=
-github.com/google/btree v1.0.1/go.mod h1:xXMiIv4Fb/0kKde4SpL7qlzvu5cMJDRkFDxJfI9uaxA=
+github.com/google/btree v1.1.2 h1:xf4v41cLI2Z6FxbKm+8Bu+m8ifhj15JuZ9sa0jZCMUU=
+github.com/google/btree v1.1.2/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
@@ -166,8 +166,8 @@ golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190328211700-ab21143f2384/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
-golang.org/x/tools v0.13.0 h1:Iey4qkscZuv0VvIt8E0neZjtPVQFSc870HQ448QgEmQ=
-golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
+golang.org/x/tools v0.14.0 h1:jvNa2pY0M4r62jkRQ6RwEZZyPcymeL9XZMLBbV7U2nc=
+golang.org/x/tools v0.14.0/go.mod h1:uYBEerGOWcJyEORxN+Ek8+TT266gXkNlHdJBwexUsBg=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
@@ -176,7 +176,7 @@ gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWD
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gvisor.dev/gvisor v0.0.0-20230715022000-fd277b20b8db h1:WZSmkyu/hep9YhWIlBZefwGVBrnGE5yW8JPD56YRsXs=
-gvisor.dev/gvisor v0.0.0-20230715022000-fd277b20b8db/go.mod h1:sQuqOkxbfJq/GS2uSnqHphtXclHyk/ZrAGhZBxxsq6g=
+gvisor.dev/gvisor v0.0.0-20231023213702-2691a8f9b1cf h1:0A28IFBR6VcMacM0m6Rn5/nr8pk8xa2TyIkjSaFAOPc=
+gvisor.dev/gvisor v0.0.0-20231023213702-2691a8f9b1cf/go.mod h1:8hmigyCdYtw5xJGfQDJzSH5Ju8XEIDBnpyi8+O6GRt8=
 inet.af/tcpproxy v0.0.0-20221017015627-91f861402626 h1:2dMP3Ox/Wh5BiItwOt4jxRsfzkgyBrHzx2nW28Yg6nc=
 inet.af/tcpproxy v0.0.0-20221017015627-91f861402626/go.mod h1:Tojt5kmHpDIR2jMojxzZK2w2ZR7OILODmUo2gaSwjrk=

pkg/cli/setup.go

@@ -14,6 +14,7 @@ import (
 	"strings"
 
 	"github.com/oomol-lab/ovm/pkg/utils"
+	"golang.org/x/crypto/ssh"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -38,6 +39,7 @@ type Context struct {
 	SSHPublicKeyPath  string
 	SSHPrivateKey     string
 	SSHPublicKey      string
+	SSHSigner         ssh.Signer
 
 	ForwardSocketPath     string
 	SocketNetworkPath     string
@@ -224,6 +226,9 @@ func (c *Context) ssh() error {
 		}
 
 		c.SSHPrivateKey = strings.TrimSpace(string(b))
+		if c.SSHSigner, err = ssh.ParsePrivateKey(b); err != nil {
+			return fmt.Errorf("parse private key error: %w", err)
+		}
 	}
 
 	return nil

pkg/ipc/restful/restful.go

@@ -7,13 +7,19 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
+	"strings"
+	"sync"
+	"time"
 
+	"github.com/Code-Hex/go-infinity-channel"
 	"github.com/Code-Hex/vz/v3"
 	"github.com/crc-org/vfkit/pkg/config"
 	"github.com/oomol-lab/ovm/pkg/cli"
 	"github.com/oomol-lab/ovm/pkg/logger"
+	"golang.org/x/crypto/ssh"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -56,6 +62,10 @@ type powerSaveModeBody struct {
 	Enable bool `json:"enable"`
 }
 
+type execBody struct {
+	Command string `json:"command"`
+}
+
 func (s *Restful) mux() *http.ServeMux {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/info", func(w http.ResponseWriter, r *http.Request) {
@@ -129,6 +139,71 @@ func (s *Restful) mux() *http.ServeMux {
 
 		s.powerSaveMode(body.Enable)
 	})
+	mux.HandleFunc("/exec", func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			http.Error(w, "post only", http.StatusBadRequest)
+			return
+		}
+
+		var body execBody
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			s.log.Warnf("Failed to decode request body: %v", err)
+			http.Error(w, "failed to decode request body", http.StatusBadRequest)
+			return
+		}
+
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.Header().Set("Cache-Control", "no-cache")
+		w.Header().Set("Connection", "keep-alive")
+
+		if _, ok := w.(http.Flusher); !ok {
+			s.log.Warnf("Bowser does not support server-sent events")
+			return
+		}
+
+		outCh := infinity.NewChannel[string]()
+		errCh := make(chan string)
+		doneCh := make(chan struct{})
+
+		go func() {
+			if err := s.exec(body.Command, outCh, errCh); err != nil {
+				s.log.Warnf("Failed to execute command: %v", err)
+			}
+
+			_, _ = fmt.Fprintf(w, "event: done\n")
+			_, _ = fmt.Fprintf(w, "data: done\n\n")
+			w.(http.Flusher).Flush()
+
+			doneCh <- struct{}{}
+			outCh.Close()
+			close(errCh)
+		}()
+
+		for {
+			select {
+			case <-doneCh:
+				s.log.Warnf("Command execution finished")
+				return
+			case err := <-errCh:
+				_, _ = fmt.Fprintf(w, "event: error\n")
+				_, _ = fmt.Fprintf(w, "data: %s\n\n", encodeSSE(err))
+				w.(http.Flusher).Flush()
+				continue
+			case out := <-outCh.Out():
+				_, _ = fmt.Fprintf(w, "event: out\n")
+				_, _ = fmt.Fprintf(w, "data: %s\n\n", encodeSSE(out))
+				w.(http.Flusher).Flush()
+				continue
+			case <-r.Context().Done():
+				s.log.Warnf("Client closed connection")
+				return
+			case <-time.After(3 * time.Second):
+				_, _ = fmt.Fprintf(w, ": ping\n\n")
+				w.(http.Flusher).Flush()
+				continue
+			}
+		}
+	})
 
 	return mux
 }
@@ -219,3 +294,84 @@ func (s *Restful) powerSaveMode(enable bool) {
 	s.log.Info("request /powerSaveMode")
 	s.opt.PowerSaveMode = enable
 }
+
+func (s *Restful) exec(command string, outCh *infinity.Channel[string], errCh chan string) error {
+	s.log.Info("request /exec")
+
+	conf := &ssh.ClientConfig{
+		User:            "root",
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		Auth: []ssh.AuthMethod{
+			ssh.PublicKeys(s.opt.SSHSigner),
+		},
+	}
+
+	conn, err := ssh.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", s.opt.SSHPort), conf)
+	if err != nil {
+		errCh <- fmt.Sprintf("dial ssh error: %v", err)
+		return fmt.Errorf("dial ssh error: %w", err)
+	}
+	defer conn.Close()
+
+	session, err := conn.NewSession()
+	if err != nil {
+		errCh <- fmt.Sprintf("new ssh session error: %v", err)
+		return fmt.Errorf("new ssh session error: %w", err)
+	}
+	defer session.Close()
+
+	w := ch2Writer(outCh)
+	session.Stdout = w
+	stderr := recordWriter(w)
+	session.Stderr = stderr
+
+	if err := session.Run(command); err != nil {
+		newErr := fmt.Errorf("%s\n%s", stderr.LastRecord(), err)
+		errCh <- fmt.Sprintf(newErr.Error())
+		return fmt.Errorf("run ssh command error: %w", newErr)
+	}
+
+	return nil
+}
+
+type chWriter struct {
+	ch *infinity.Channel[string]
+	mu sync.Mutex
+}
+
+func (w *chWriter) Write(p []byte) (n int, err error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	w.ch.In() <- string(p)
+	return len(p), nil
+}
+
+func ch2Writer(ch *infinity.Channel[string]) io.Writer {
+	return &chWriter{
+		ch: ch,
+	}
+}
+
+type writer struct {
+	w    io.Writer
+	last []byte
+}
+
+func (w *writer) Write(p []byte) (n int, err error) {
+	w.last = p
+	return w.w.Write(p)
+}
+
+func (w *writer) LastRecord() string {
+	return string(w.last)
+}
+
+func recordWriter(w io.Writer) *writer {
+	return &writer{
+		w: w,
+	}
+}
+
+func encodeSSE(str string) string {
+	return strings.ReplaceAll(strings.TrimSpace(str), "\n", "\ndata: ")
+}