From c4a0e85d32fd9a422e262b2dd918fbb5d7a3711b Mon Sep 17 00:00:00 2001 From: Mike Sawka Date: Mon, 12 Aug 2024 10:58:39 -0700 Subject: [PATCH] update wsh code for easier creation of client servers (for readfile/readdir/fileinfo) (#218) --- Taskfile.yml | 12 ++ cmd/wsh/cmd/wshcmd-root.go | 16 +- cmd/wsh/cmd/wshcmd-server.go | 236 +++++++++++++++++++++++++ frontend/app/store/wos.ts | 1 - frontend/app/store/wshserver.ts | 15 ++ frontend/types/gotypes.d.ts | 14 +- pkg/service/fileservice/fileservice.go | 27 +-- pkg/wshrpc/wshclient/wshclient.go | 17 ++ pkg/wshrpc/wshclient/wshclientutil.go | 6 + pkg/wshrpc/wshrpctypes.go | 31 ++++ pkg/wshrpc/wshserver/wshserver.go | 37 ++++ pkg/wshrpc/wshserver/wshserverutil.go | 76 +------- pkg/wshutil/wshadapter.go | 132 ++++++++++++++ pkg/wshutil/wshrpc.go | 90 +++++++--- pkg/wshutil/wshutil.go | 12 +- 15 files changed, 590 insertions(+), 132 deletions(-) create mode 100644 cmd/wsh/cmd/wshcmd-server.go create mode 100644 pkg/wshutil/wshadapter.go diff --git a/Taskfile.yml b/Taskfile.yml index cf5f81e4..c02c02d0 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -120,6 +120,18 @@ tasks: GOOS: windows GOARCH: arm64 + dev:installwsh: + desc: quick shortcut to rebuild wsh and install for macos arm64 + requires: + vars: + - VERSION + cmds: + - task: build:wsh:internal + vars: + GOOS: darwin + GOARCH: arm64 + - cp dist/bin/wsh-{{.VERSION}}-darwin.arm64 ~/.w2-dev/bin/wsh + build:wsh:internal: vars: EXT: diff --git a/cmd/wsh/cmd/wshcmd-root.go b/cmd/wsh/cmd/wshcmd-root.go index ccd34894..93c14b8c 100644 --- a/cmd/wsh/cmd/wshcmd-root.go +++ b/cmd/wsh/cmd/wshcmd-root.go @@ -61,19 +61,19 @@ func WriteStdout(fmtStr string, args ...interface{}) { } // returns the wrapped stdin and a new rpc client (that wraps the stdin input and stdout output) -func setupRpcClient(handlerFn wshutil.CommandHandlerFnType) error { +func setupRpcClient(serverImpl wshutil.ServerImpl) error { jwtToken := os.Getenv("WAVETERM_JWT") if jwtToken == "" { wshutil.SetTermRawModeAndInstallShutdownHandlers(true) UsingTermWshMode = true - RpcClient, WrappedStdin = wshutil.SetupTerminalRpcClient(handlerFn) + RpcClient, WrappedStdin = wshutil.SetupTerminalRpcClient(serverImpl) return nil } sockName, err := wshutil.ExtractUnverifiedSocketName(jwtToken) if err != nil { return fmt.Errorf("error extracting socket name from WAVETERM_JWT: %v", err) } - RpcClient, err = wshutil.SetupDomainSocketRpcClient(sockName, handlerFn) + RpcClient, err = wshutil.SetupDomainSocketRpcClient(sockName, serverImpl) if err != nil { return fmt.Errorf("error setting up domain socket rpc client: %v", err) } @@ -149,7 +149,15 @@ func resolveSimpleId(id string) (*waveobj.ORef, error) { // Execute executes the root command. func Execute() { - defer wshutil.DoShutdown("", 0, false) + defer func() { + r := recover() + if r != nil { + WriteStderr("[panic] %v\n", r) + wshutil.DoShutdown("", 1, true) + } else { + wshutil.DoShutdown("", 0, false) + } + }() err := setupRpcClient(nil) if err != nil { log.Printf("[error] %v\n", err) diff --git a/cmd/wsh/cmd/wshcmd-server.go b/cmd/wsh/cmd/wshcmd-server.go new file mode 100644 index 00000000..d0e01ccb --- /dev/null +++ b/cmd/wsh/cmd/wshcmd-server.go @@ -0,0 +1,236 @@ +// Copyright 2024, Command Line Inc. +// SPDX-License-Identifier: Apache-2.0 + +package cmd + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "io" + "os" + "path/filepath" + + "github.com/spf13/cobra" + "github.com/wavetermdev/thenextwave/pkg/util/utilfn" + "github.com/wavetermdev/thenextwave/pkg/wavebase" + "github.com/wavetermdev/thenextwave/pkg/wshrpc" + "github.com/wavetermdev/thenextwave/pkg/wshrpc/wshclient" +) + +const MaxFileSize = 50 * 1024 * 1024 // 10M + +var serverCmd = &cobra.Command{ + Use: "server", + Short: "remote server to power wave blocks", + Args: cobra.NoArgs, + Run: serverRun, +} + +type ServerImpl struct{} + +func (*ServerImpl) WshServerImpl() {} + +func (*ServerImpl) MessageCommand(ctx context.Context, data wshrpc.CommandMessageData) error { + WriteStderr("[message] %q\n", data.Message) + return nil +} + +func respErr(err error) wshrpc.RespOrErrorUnion[wshrpc.CommandRemoteStreamFileRtnData] { + return wshrpc.RespOrErrorUnion[wshrpc.CommandRemoteStreamFileRtnData]{Error: err} +} + +type ByteRangeType struct { + All bool + Start int64 + End int64 +} + +func parseByteRange(rangeStr string) (ByteRangeType, error) { + if rangeStr == "" { + return ByteRangeType{All: true}, nil + } + var start, end int64 + _, err := fmt.Sscanf(rangeStr, "%d-%d", &start, &end) + if err != nil { + return ByteRangeType{}, errors.New("invalid byte range") + } + if start < 0 || end < 0 || start > end { + return ByteRangeType{}, errors.New("invalid byte range") + } + return ByteRangeType{Start: start, End: end}, nil +} + +func (impl *ServerImpl) remoteStreamFileDir(ctx context.Context, path string, byteRange ByteRangeType, dataCallback func(fileInfo *wshrpc.FileInfo, data []byte)) error { + innerFilesEntries, err := os.ReadDir(path) + if err != nil { + return fmt.Errorf("cannot open dir %q: %w", path, err) + } + if byteRange.All { + if len(innerFilesEntries) > 1000 { + innerFilesEntries = innerFilesEntries[:1000] + } + } else { + if byteRange.Start >= int64(len(innerFilesEntries)) { + return nil + } + realEnd := byteRange.End + if realEnd > int64(len(innerFilesEntries)) { + realEnd = int64(len(innerFilesEntries)) + } + innerFilesEntries = innerFilesEntries[byteRange.Start:realEnd] + } + parent := filepath.Dir(path) + parentFileInfo, err := impl.RemoteFileInfoCommand(ctx, parent) + if err == nil && parent != path { + parentFileInfo.Name = ".." + parentFileInfo.Size = -1 + dataCallback(parentFileInfo, nil) + } + for _, innerFileEntry := range innerFilesEntries { + if ctx.Err() != nil { + return ctx.Err() + } + innerFileInfoInt, err := innerFileEntry.Info() + if err != nil { + continue + } + mimeType := utilfn.DetectMimeType(filepath.Join(path, innerFileInfoInt.Name())) + var fileSize int64 + if mimeType == "directory" { + fileSize = -1 + } else { + fileSize = innerFileInfoInt.Size() + } + innerFileInfo := wshrpc.FileInfo{ + Path: filepath.Join(path, innerFileInfoInt.Name()), + Name: innerFileInfoInt.Name(), + Size: fileSize, + Mode: innerFileInfoInt.Mode(), + ModeStr: innerFileInfoInt.Mode().String(), + ModTime: innerFileInfoInt.ModTime().UnixMilli(), + IsDir: innerFileInfoInt.IsDir(), + MimeType: mimeType, + } + dataCallback(&innerFileInfo, nil) + } + return nil +} + +func (impl *ServerImpl) remoteStreamFileRegular(ctx context.Context, path string, byteRange ByteRangeType, dataCallback func(fileInfo *wshrpc.FileInfo, data []byte)) error { + fd, err := os.Open(path) + if err != nil { + return fmt.Errorf("cannot open file %q: %w", path, err) + } + defer fd.Close() + var filePos int64 + if !byteRange.All && byteRange.Start > 0 { + _, err := fd.Seek(byteRange.Start, io.SeekStart) + if err != nil { + return fmt.Errorf("seeking file %q: %w", path, err) + } + filePos = byteRange.Start + } + buf := make([]byte, 4096) + for { + if ctx.Err() != nil { + return ctx.Err() + } + n, err := fd.Read(buf) + if n > 0 { + if !byteRange.All && filePos+int64(n) > byteRange.End { + n = int(byteRange.End - filePos) + } + filePos += int64(n) + dataCallback(nil, buf[:n]) + } + if filePos >= byteRange.End { + break + } + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return fmt.Errorf("reading file %q: %w", path, err) + } + } + return nil +} + +func (impl *ServerImpl) remoteStreamFileInternal(ctx context.Context, data wshrpc.CommandRemoteStreamFileData, dataCallback func(fileInfo *wshrpc.FileInfo, data []byte)) error { + byteRange, err := parseByteRange(data.ByteRange) + if err != nil { + return err + } + path := data.Path + path = wavebase.ExpandHomeDir(path) + finfo, err := impl.RemoteFileInfoCommand(ctx, path) + if err != nil { + return fmt.Errorf("cannot stat file %q: %w", path, err) + } + dataCallback(finfo, nil) + if finfo.NotFound { + return nil + } + if finfo.Size > MaxFileSize { + return fmt.Errorf("file %q is too large to read, use /wave/stream-file", path) + } + if finfo.IsDir { + return impl.remoteStreamFileDir(ctx, path, byteRange, dataCallback) + } else { + return impl.remoteStreamFileRegular(ctx, path, byteRange, dataCallback) + } +} + +func (impl *ServerImpl) RemoteStreamFileCommand(ctx context.Context, data wshrpc.CommandRemoteStreamFileData) chan wshrpc.RespOrErrorUnion[wshrpc.CommandRemoteStreamFileRtnData] { + ch := make(chan wshrpc.RespOrErrorUnion[wshrpc.CommandRemoteStreamFileRtnData], 16) + defer close(ch) + err := impl.remoteStreamFileInternal(ctx, data, func(fileInfo *wshrpc.FileInfo, data []byte) { + resp := wshrpc.CommandRemoteStreamFileRtnData{} + resp.FileInfo = fileInfo + if len(data) > 0 { + resp.Data64 = base64.RawStdEncoding.EncodeToString(data) + } + ch <- wshrpc.RespOrErrorUnion[wshrpc.CommandRemoteStreamFileRtnData]{Response: resp} + }) + if err != nil { + ch <- respErr(err) + } + return ch +} + +func (*ServerImpl) RemoteFileInfoCommand(ctx context.Context, path string) (*wshrpc.FileInfo, error) { + cleanedPath := filepath.Clean(wavebase.ExpandHomeDir(path)) + finfo, err := os.Stat(cleanedPath) + if os.IsNotExist(err) { + return &wshrpc.FileInfo{Path: wavebase.ReplaceHomeDir(path), NotFound: true}, nil + } + if err != nil { + return nil, fmt.Errorf("cannot stat file %q: %w", path, err) + } + mimeType := utilfn.DetectMimeType(cleanedPath) + return &wshrpc.FileInfo{ + Path: cleanedPath, + Name: finfo.Name(), + Size: finfo.Size(), + Mode: finfo.Mode(), + ModeStr: finfo.Mode().String(), + ModTime: finfo.ModTime().UnixMilli(), + IsDir: finfo.IsDir(), + MimeType: mimeType, + }, nil +} + +func init() { + rootCmd.AddCommand(serverCmd) +} + +func serverRun(cmd *cobra.Command, args []string) { + WriteStdout("running wsh server\n") + RpcClient.SetServerImpl(&ServerImpl{}) + err := wshclient.TestCommand(RpcClient, "hello", nil) + WriteStdout("got test rtn: %v\n", err) + + select {} // run forever +} diff --git a/frontend/app/store/wos.ts b/frontend/app/store/wos.ts index bbc8bcb9..2620ffd9 100644 --- a/frontend/app/store/wos.ts +++ b/frontend/app/store/wos.ts @@ -168,7 +168,6 @@ function createWaveValueObject(oref: string, shouldFetch: boo const localPromise = GetObject(oref); wov.pendingPromise = localPromise; localPromise.then((val) => { - console.log("GetObject resolved", oref, val); if (wov.pendingPromise != localPromise) { return; } diff --git a/frontend/app/store/wshserver.ts b/frontend/app/store/wshserver.ts index d464f28f..6e51946e 100644 --- a/frontend/app/store/wshserver.ts +++ b/frontend/app/store/wshserver.ts @@ -87,6 +87,16 @@ class WshServerType { return WOS.wshServerRpcHelper_call("message", data, opts); } + // command "remotefileinfo" [call] + RemoteFileInfoCommand(data: string, opts?: WshRpcCommandOpts): Promise { + return WOS.wshServerRpcHelper_call("remotefileinfo", data, opts); + } + + // command "remotestreamfile" [responsestream] + RemoteStreamFileCommand(data: CommandRemoteStreamFileData, opts?: WshRpcCommandOpts): AsyncGenerator { + return WOS.wshServerRpcHelper_responsestream("remotestreamfile", data, opts); + } + // command "resolveids" [call] ResolveIdsCommand(data: CommandResolveIdsData, opts?: WshRpcCommandOpts): Promise { return WOS.wshServerRpcHelper_call("resolveids", data, opts); @@ -117,6 +127,11 @@ class WshServerType { return WOS.wshServerRpcHelper_responsestream("streamwaveai", data, opts); } + // command "test" [call] + TestCommand(data: string, opts?: WshRpcCommandOpts): Promise { + return WOS.wshServerRpcHelper_call("test", data, opts); + } + } export const WshServer = new WshServerType(); diff --git a/frontend/types/gotypes.d.ts b/frontend/types/gotypes.d.ts index c7799d01..1992843a 100644 --- a/frontend/types/gotypes.d.ts +++ b/frontend/types/gotypes.d.ts @@ -116,6 +116,18 @@ declare global { message: string; }; + // wshrpc.CommandRemoteStreamFileData + type CommandRemoteStreamFileData = { + path: string; + byterange?: string; + }; + + // wshrpc.CommandRemoteStreamFileRtnData + type CommandRemoteStreamFileRtnData = { + fileinfo?: FileInfo; + data64?: string; + }; + // wshrpc.CommandResolveIdsData type CommandResolveIdsData = { ids: string[]; @@ -153,7 +165,7 @@ declare global { meta?: {[key: string]: any}; }; - // fileservice.FileInfo + // wshrpc.FileInfo type FileInfo = { path: string; name: string; diff --git a/pkg/service/fileservice/fileservice.go b/pkg/service/fileservice/fileservice.go index e77d671c..a9656644 100644 --- a/pkg/service/fileservice/fileservice.go +++ b/pkg/service/fileservice/fileservice.go @@ -14,6 +14,7 @@ import ( "github.com/wavetermdev/thenextwave/pkg/util/utilfn" "github.com/wavetermdev/thenextwave/pkg/wavebase" "github.com/wavetermdev/thenextwave/pkg/wconfig" + "github.com/wavetermdev/thenextwave/pkg/wshrpc" ) const MaxFileSize = 10 * 1024 * 1024 // 10M @@ -21,21 +22,9 @@ const DefaultTimeout = 2 * time.Second type FileService struct{} -type FileInfo struct { - Path string `json:"path"` // cleaned path - Name string `json:"name"` - NotFound bool `json:"notfound,omitempty"` - Size int64 `json:"size"` - Mode os.FileMode `json:"mode"` - ModeStr string `json:"modestr"` - ModTime int64 `json:"modtime"` - IsDir bool `json:"isdir,omitempty"` - MimeType string `json:"mimetype,omitempty"` -} - type FullFile struct { - Info *FileInfo `json:"info"` - Data64 string `json:"data64"` // base64 encoded + Info *wshrpc.FileInfo `json:"info"` + Data64 string `json:"data64"` // base64 encoded } func (fs *FileService) SaveFile(path string, data64 string) error { @@ -51,17 +40,17 @@ func (fs *FileService) SaveFile(path string, data64 string) error { return nil } -func (fs *FileService) StatFile(path string) (*FileInfo, error) { +func (fs *FileService) StatFile(path string) (*wshrpc.FileInfo, error) { cleanedPath := filepath.Clean(wavebase.ExpandHomeDir(path)) finfo, err := os.Stat(cleanedPath) if os.IsNotExist(err) { - return &FileInfo{Path: wavebase.ReplaceHomeDir(path), NotFound: true}, nil + return &wshrpc.FileInfo{Path: wavebase.ReplaceHomeDir(path), NotFound: true}, nil } if err != nil { return nil, fmt.Errorf("cannot stat file %q: %w", path, err) } mimeType := utilfn.DetectMimeType(cleanedPath) - return &FileInfo{ + return &wshrpc.FileInfo{ Path: cleanedPath, Name: finfo.Name(), Size: finfo.Size(), @@ -92,7 +81,7 @@ func (fs *FileService) ReadFile(path string) (*FullFile, error) { if len(innerFilesEntries) > 1000 { innerFilesEntries = innerFilesEntries[:1000] } - var innerFilesInfo []FileInfo + var innerFilesInfo []wshrpc.FileInfo parent := filepath.Dir(finfo.Path) parentFileInfo, err := fs.StatFile(parent) if err == nil && parent != finfo.Path { @@ -114,7 +103,7 @@ func (fs *FileService) ReadFile(path string) (*FullFile, error) { } else { fileSize = innerFileInfoInt.Size() } - innerFileInfo := FileInfo{ + innerFileInfo := wshrpc.FileInfo{ Path: filepath.Join(finfo.Path, innerFileInfoInt.Name()), Name: innerFileInfoInt.Name(), Size: fileSize, diff --git a/pkg/wshrpc/wshclient/wshclient.go b/pkg/wshrpc/wshclient/wshclient.go index 677563d4..8481e048 100644 --- a/pkg/wshrpc/wshclient/wshclient.go +++ b/pkg/wshrpc/wshclient/wshclient.go @@ -107,6 +107,17 @@ func MessageCommand(w *wshutil.WshRpc, data wshrpc.CommandMessageData, opts *wsh return err } +// command "remotefileinfo", wshserver.RemoteFileInfoCommand +func RemoteFileInfoCommand(w *wshutil.WshRpc, data string, opts *wshrpc.WshRpcCommandOpts) (*wshrpc.FileInfo, error) { + resp, err := sendRpcRequestCallHelper[*wshrpc.FileInfo](w, "remotefileinfo", data, opts) + return resp, err +} + +// command "remotestreamfile", wshserver.RemoteStreamFileCommand +func RemoteStreamFileCommand(w *wshutil.WshRpc, data wshrpc.CommandRemoteStreamFileData, opts *wshrpc.WshRpcCommandOpts) chan wshrpc.RespOrErrorUnion[wshrpc.CommandRemoteStreamFileRtnData] { + return sendRpcRequestResponseStreamHelper[wshrpc.CommandRemoteStreamFileRtnData](w, "remotestreamfile", data, opts) +} + // command "resolveids", wshserver.ResolveIdsCommand func ResolveIdsCommand(w *wshutil.WshRpc, data wshrpc.CommandResolveIdsData, opts *wshrpc.WshRpcCommandOpts) (wshrpc.CommandResolveIdsRtnData, error) { resp, err := sendRpcRequestCallHelper[wshrpc.CommandResolveIdsRtnData](w, "resolveids", data, opts) @@ -140,4 +151,10 @@ func StreamWaveAiCommand(w *wshutil.WshRpc, data wshrpc.OpenAiStreamRequest, opt return sendRpcRequestResponseStreamHelper[wshrpc.OpenAIPacketType](w, "streamwaveai", data, opts) } +// command "test", wshserver.TestCommand +func TestCommand(w *wshutil.WshRpc, data string, opts *wshrpc.WshRpcCommandOpts) error { + _, err := sendRpcRequestCallHelper[any](w, "test", data, opts) + return err +} + diff --git a/pkg/wshrpc/wshclient/wshclientutil.go b/pkg/wshrpc/wshclient/wshclientutil.go index ad74adec..eba36e7f 100644 --- a/pkg/wshrpc/wshclient/wshclientutil.go +++ b/pkg/wshrpc/wshclient/wshclientutil.go @@ -10,6 +10,9 @@ import ( ) func sendRpcRequestCallHelper[T any](w *wshutil.WshRpc, command string, data interface{}, opts *wshrpc.WshRpcCommandOpts) (T, error) { + if opts == nil { + opts = &wshrpc.WshRpcCommandOpts{} + } var respData T if opts.NoResponse { err := w.SendCommand(command, data) @@ -30,6 +33,9 @@ func sendRpcRequestCallHelper[T any](w *wshutil.WshRpc, command string, data int } func sendRpcRequestResponseStreamHelper[T any](w *wshutil.WshRpc, command string, data interface{}, opts *wshrpc.WshRpcCommandOpts) chan wshrpc.RespOrErrorUnion[T] { + if opts == nil { + opts = &wshrpc.WshRpcCommandOpts{} + } respChan := make(chan wshrpc.RespOrErrorUnion[T]) reqHandler, err := w.SendComplexRequest(command, data, true, opts.Timeout) if err != nil { diff --git a/pkg/wshrpc/wshrpctypes.go b/pkg/wshrpc/wshrpctypes.go index cde23d94..ade1325b 100644 --- a/pkg/wshrpc/wshrpctypes.go +++ b/pkg/wshrpc/wshrpctypes.go @@ -6,6 +6,7 @@ package wshrpc import ( "context" + "os" "reflect" "github.com/wavetermdev/thenextwave/pkg/ijson" @@ -44,6 +45,9 @@ const ( Command_StreamTest = "streamtest" Command_StreamWaveAi = "streamwaveai" Command_StreamCpuData = "streamcpudata" + Command_Test = "test" + Command_RemoteStreamFile = "remotestreamfile" + Command_RemoteFileInfo = "remotefileinfo" ) type RespOrErrorUnion[T any] struct { @@ -74,6 +78,11 @@ type WshRpcInterface interface { StreamTestCommand(ctx context.Context) chan RespOrErrorUnion[int] StreamWaveAiCommand(ctx context.Context, request OpenAiStreamRequest) chan RespOrErrorUnion[OpenAIPacketType] StreamCpuDataCommand(ctx context.Context, request CpuDataRequest) chan RespOrErrorUnion[CpuDataType] + TestCommand(ctx context.Context, data string) error + + // remotes + RemoteStreamFileCommand(ctx context.Context, data CommandRemoteStreamFileData) chan RespOrErrorUnion[CommandRemoteStreamFileRtnData] + RemoteFileInfoCommand(ctx context.Context, path string) (*FileInfo, error) } // for frontend @@ -243,3 +252,25 @@ type CpuDataType struct { Time int64 `json:"time"` Value float64 `json:"value"` } + +type FileInfo struct { + Path string `json:"path"` // cleaned path + Name string `json:"name"` + NotFound bool `json:"notfound,omitempty"` + Size int64 `json:"size"` + Mode os.FileMode `json:"mode"` + ModeStr string `json:"modestr"` + ModTime int64 `json:"modtime"` + IsDir bool `json:"isdir,omitempty"` + MimeType string `json:"mimetype,omitempty"` +} + +type CommandRemoteStreamFileData struct { + Path string `json:"path"` + ByteRange string `json:"byterange,omitempty"` +} + +type CommandRemoteStreamFileRtnData struct { + FileInfo *FileInfo `json:"fileinfo,omitempty"` + Data64 string `json:"data64,omitempty"` +} diff --git a/pkg/wshrpc/wshserver/wshserver.go b/pkg/wshrpc/wshserver/wshserver.go index 6a75d261..9dfd8cdc 100644 --- a/pkg/wshrpc/wshserver/wshserver.go +++ b/pkg/wshrpc/wshserver/wshserver.go @@ -23,12 +23,49 @@ import ( "github.com/wavetermdev/thenextwave/pkg/waveobj" "github.com/wavetermdev/thenextwave/pkg/wps" "github.com/wavetermdev/thenextwave/pkg/wshrpc" + "github.com/wavetermdev/thenextwave/pkg/wshrpc/wshclient" "github.com/wavetermdev/thenextwave/pkg/wshutil" "github.com/wavetermdev/thenextwave/pkg/wstore" ) const SimpleId_This = "this" +func (ws *WshServer) TestCommand(ctx context.Context, data string) error { + defer func() { + if r := recover(); r != nil { + log.Printf("panic in TestCommand: %v", r) + } + }() + rpc := wshutil.GetWshRpcFromContext(ctx) + if rpc == nil { + return nil + } + go func() { + wshclient.MessageCommand(rpc, wshrpc.CommandMessageData{Message: "test message"}, &wshrpc.WshRpcCommandOpts{NoResponse: true}) + resp, err := wshclient.RemoteFileInfoCommand(rpc, "~/work/wails/thenextwave/README.md", nil) + if err != nil { + log.Printf("error getting remote file info: %v", err) + return + } + log.Printf("remote file info: %#v\n", resp) + + rch := wshclient.RemoteStreamFileCommand(rpc, wshrpc.CommandRemoteStreamFileData{Path: "~/work/wails/thenextwave/README.md"}, nil) + for msg := range rch { + if msg.Error != nil { + log.Printf("error in stream: %v", msg.Error) + break + } + if msg.Response.FileInfo != nil { + log.Printf("stream resp (fileinfo): %v\n", msg.Response.FileInfo) + } + if msg.Response.Data64 != "" { + log.Printf("stream resp (data): %v\n", len(msg.Response.Data64)) + } + } + }() + return nil +} + func (ws *WshServer) AuthenticateCommand(ctx context.Context, data string) error { w := wshutil.GetWshRpcFromContext(ctx) if w == nil { diff --git a/pkg/wshrpc/wshserver/wshserverutil.go b/pkg/wshrpc/wshserver/wshserverutil.go index 20824e86..565ab639 100644 --- a/pkg/wshrpc/wshserver/wshserverutil.go +++ b/pkg/wshrpc/wshserver/wshserverutil.go @@ -10,7 +10,6 @@ import ( "net" "reflect" - "github.com/wavetermdev/thenextwave/pkg/util/utilfn" "github.com/wavetermdev/thenextwave/pkg/wshrpc" "github.com/wavetermdev/thenextwave/pkg/wshutil" ) @@ -24,6 +23,8 @@ const ( type WshServer struct{} +func (*WshServer) WshServerImpl() {} + type WshServerMethodDecl struct { Command string CommandType string @@ -88,75 +89,6 @@ func decodeRtnVals(rtnVals []reflect.Value) (any, error) { } } -func mainWshServerHandler(handler *wshutil.RpcResponseHandler) bool { - command := handler.GetCommand() - methodDecl := wshCommandDeclMap[command] - if methodDecl == nil { - handler.SendResponseError(fmt.Errorf("command %q not found", command)) - return true - } - var callParams []reflect.Value - callParams = append(callParams, reflect.ValueOf(handler.Context())) - if methodDecl.CommandDataType != nil { - commandData := reflect.New(methodDecl.CommandDataType).Interface() - err := utilfn.ReUnmarshal(commandData, handler.GetCommandRawData()) - if err != nil { - handler.SendResponseError(fmt.Errorf("error re-marshalling command data: %w", err)) - return true - } - wshrpc.HackRpcContextIntoData(commandData, handler.GetRpcContext()) - callParams = append(callParams, reflect.ValueOf(commandData).Elem()) - } - implVal := reflect.ValueOf(&WshServerImpl) - implMethod := implVal.MethodByName(methodDecl.MethodName) - if !implMethod.IsValid() { - if !handler.NeedsResponse() { - // we also send an out of band message here since this is likely unexpected and will require debugging - handler.SendMessage(fmt.Sprintf("command %q method %q not found", handler.GetCommand(), methodDecl.MethodName)) - } - handler.SendResponseError(fmt.Errorf("method %q not found", methodDecl.MethodName)) - return true - } - if methodDecl.CommandType == wshrpc.RpcType_Call { - rtnVals := implMethod.Call(callParams) - rtnData, rtnErr := decodeRtnVals(rtnVals) - if rtnErr != nil { - handler.SendResponseError(rtnErr) - return true - } - handler.SendResponse(rtnData, true) - return true - } else if methodDecl.CommandType == wshrpc.RpcType_ResponseStream { - rtnVals := implMethod.Call(callParams) - rtnChVal := rtnVals[0] - if rtnChVal.IsNil() { - handler.SendResponse(nil, true) - return true - } - go func() { - defer handler.Finalize() - // must use reflection here because we don't know the generic type of RespOrErrorUnion - for { - respVal, ok := rtnChVal.Recv() - if !ok { - break - } - errorVal := respVal.FieldByName("Error") - if !errorVal.IsNil() { - handler.SendResponseError(errorVal.Interface().(error)) - break - } - respData := respVal.FieldByName("Response").Interface() - handler.SendResponse(respData, false) - } - }() - return false - } else { - handler.SendResponseError(fmt.Errorf("unsupported command type %q", methodDecl.CommandType)) - return true - } -} - func RunWshRpcOverListener(listener net.Listener) { defer log.Printf("domain socket listener shutting down\n") for { @@ -167,10 +99,10 @@ func RunWshRpcOverListener(listener net.Listener) { } log.Print("got domain socket connection\n") // TODO deal with closing connection - go wshutil.SetupConnRpcClient(conn, mainWshServerHandler) + go wshutil.SetupConnRpcClient(conn, &WshServerImpl) } } func MakeWshServer(inputCh chan []byte, outputCh chan []byte, initialCtx wshrpc.RpcContext) { - wshutil.MakeWshRpc(inputCh, outputCh, initialCtx, mainWshServerHandler) + wshutil.MakeWshRpc(inputCh, outputCh, initialCtx, &WshServerImpl) } diff --git a/pkg/wshutil/wshadapter.go b/pkg/wshutil/wshadapter.go new file mode 100644 index 00000000..b99ef2c1 --- /dev/null +++ b/pkg/wshutil/wshadapter.go @@ -0,0 +1,132 @@ +// Copyright 2024, Command Line Inc. +// SPDX-License-Identifier: Apache-2.0 + +package wshutil + +import ( + "fmt" + "reflect" + "strings" + + "github.com/wavetermdev/thenextwave/pkg/util/utilfn" + "github.com/wavetermdev/thenextwave/pkg/wshrpc" +) + +var WshCommandDeclMap = wshrpc.GenerateWshCommandDeclMap() + +func findCmdMethod(impl any, cmd string) *reflect.Method { + rtype := reflect.TypeOf(impl) + methodName := cmd + "command" + for i := 0; i < rtype.NumMethod(); i++ { + method := rtype.Method(i) + if strings.ToLower(method.Name) == methodName { + return &method + } + } + return nil +} + +func decodeRtnVals(rtnVals []reflect.Value) (any, error) { + switch len(rtnVals) { + case 0: + return nil, nil + case 1: + errIf := rtnVals[0].Interface() + if errIf == nil { + return nil, nil + } + return nil, errIf.(error) + case 2: + errIf := rtnVals[1].Interface() + if errIf == nil { + return rtnVals[0].Interface(), nil + } + return rtnVals[0].Interface(), errIf.(error) + default: + return nil, fmt.Errorf("too many return values: %d", len(rtnVals)) + } +} + +func noImplHandler(handler *RpcResponseHandler) bool { + handler.SendResponseError(fmt.Errorf("command %q not implemented", handler.GetCommand())) + return true +} + +func serverImplAdapter(impl any) func(*RpcResponseHandler) bool { + if impl == nil { + return noImplHandler + } + rtype := reflect.TypeOf(impl) + if rtype.Kind() != reflect.Ptr && rtype.Elem().Kind() != reflect.Struct { + panic(fmt.Sprintf("expected struct pointer, got %s", rtype)) + } + // returns isAsync + return func(handler *RpcResponseHandler) bool { + cmd := handler.GetCommand() + methodDecl := WshCommandDeclMap[cmd] + if methodDecl == nil { + handler.SendResponseError(fmt.Errorf("command %q not found", cmd)) + return true + } + rmethod := findCmdMethod(impl, cmd) + if rmethod == nil { + if !handler.NeedsResponse() { + // we also send an out of band message here since this is likely unexpected and will require debugging + handler.SendMessage(fmt.Sprintf("command %q method %q not found", handler.GetCommand(), methodDecl.MethodName)) + } + handler.SendResponseError(fmt.Errorf("command not implemented %q", cmd)) + return true + } + implMethod := reflect.ValueOf(impl).MethodByName(rmethod.Name) + var callParams []reflect.Value + callParams = append(callParams, reflect.ValueOf(handler.Context())) + if methodDecl.CommandDataType != nil { + commandData := reflect.New(methodDecl.CommandDataType).Interface() + err := utilfn.ReUnmarshal(commandData, handler.GetCommandRawData()) + if err != nil { + handler.SendResponseError(fmt.Errorf("error re-marshalling command data: %w", err)) + return true + } + wshrpc.HackRpcContextIntoData(commandData, handler.GetRpcContext()) + callParams = append(callParams, reflect.ValueOf(commandData).Elem()) + } + if methodDecl.CommandType == wshrpc.RpcType_Call { + rtnVals := implMethod.Call(callParams) + rtnData, rtnErr := decodeRtnVals(rtnVals) + if rtnErr != nil { + handler.SendResponseError(rtnErr) + return true + } + handler.SendResponse(rtnData, true) + return true + } else if methodDecl.CommandType == wshrpc.RpcType_ResponseStream { + rtnVals := implMethod.Call(callParams) + rtnChVal := rtnVals[0] + if rtnChVal.IsNil() { + handler.SendResponse(nil, true) + return true + } + go func() { + defer handler.Finalize() + // must use reflection here because we don't know the generic type of RespOrErrorUnion + for { + respVal, ok := rtnChVal.Recv() + if !ok { + break + } + errorVal := respVal.FieldByName("Error") + if !errorVal.IsNil() { + handler.SendResponseError(errorVal.Interface().(error)) + break + } + respData := respVal.FieldByName("Response").Interface() + handler.SendResponse(respData, false) + } + }() + return false + } else { + handler.SendResponseError(fmt.Errorf("unsupported command type %q", methodDecl.CommandType)) + return true + } + } +} diff --git a/pkg/wshutil/wshrpc.go b/pkg/wshutil/wshrpc.go index 8b52677a..2479f047 100644 --- a/pkg/wshutil/wshrpc.go +++ b/pkg/wshutil/wshrpc.go @@ -9,6 +9,7 @@ import ( "errors" "fmt" "log" + "reflect" "runtime/debug" "sync" "sync/atomic" @@ -27,6 +28,22 @@ type ResponseFnType = func(any) error // returns true if handler is complete, false for an async handler type CommandHandlerFnType = func(*RpcResponseHandler) bool +type ServerImpl interface { + WshServerImpl() +} + +type WshRpc struct { + Lock *sync.Mutex + clientId string + InputCh chan []byte + OutputCh chan []byte + RpcContext *atomic.Pointer[wshrpc.RpcContext] + RpcMap map[string]*rpcData + ServerImpl ServerImpl + + ResponseHandlerMap map[string]*RpcResponseHandler // reqId => handler +} + type wshRpcContextKey struct{} func withWshRpcContext(ctx context.Context, wshRpc *WshRpc) context.Context { @@ -109,26 +126,25 @@ func (r *RpcMessage) Validate() error { return fmt.Errorf("invalid packet: must have command, reqid, or resid set") } -type WshRpc struct { - Lock *sync.Mutex - clientId string - InputCh chan []byte - OutputCh chan []byte - RpcContext *atomic.Pointer[wshrpc.RpcContext] - RpcMap map[string]*rpcData - HandlerFn CommandHandlerFnType - - ResponseHandlerMap map[string]*RpcResponseHandler // reqId => handler -} - type rpcData struct { ResCh chan *RpcMessage Ctx context.Context } +func validateServerImpl(serverImpl ServerImpl) { + if serverImpl == nil { + return + } + serverType := reflect.TypeOf(serverImpl) + if serverType.Kind() != reflect.Pointer && serverType.Elem().Kind() != reflect.Struct { + panic(fmt.Sprintf("serverImpl must be a pointer to struct, got %v", serverType)) + } +} + // oscEsc is the OSC escape sequence to use for *sending* messages // closes outputCh when inputCh is closed/done -func MakeWshRpc(inputCh chan []byte, outputCh chan []byte, rpcCtx wshrpc.RpcContext, commandHandlerFn CommandHandlerFnType) *WshRpc { +func MakeWshRpc(inputCh chan []byte, outputCh chan []byte, rpcCtx wshrpc.RpcContext, serverImpl ServerImpl) *WshRpc { + validateServerImpl(serverImpl) rtn := &WshRpc{ Lock: &sync.Mutex{}, clientId: uuid.New().String(), @@ -136,7 +152,7 @@ func MakeWshRpc(inputCh chan []byte, outputCh chan []byte, rpcCtx wshrpc.RpcCont OutputCh: outputCh, RpcMap: make(map[string]*rpcData), RpcContext: &atomic.Pointer[wshrpc.RpcContext]{}, - HandlerFn: commandHandlerFn, + ServerImpl: serverImpl, ResponseHandlerMap: make(map[string]*RpcResponseHandler), } rtn.RpcContext.Store(&rpcCtx) @@ -243,9 +259,8 @@ func (w *WshRpc) handleRequest(req *RpcMessage) { respHandler.Finalize() } }() - if w.HandlerFn != nil { - isAsync = !w.HandlerFn(respHandler) - } + handlerFn := serverImplAdapter(w.ServerImpl) + isAsync = !handlerFn(respHandler) } func (w *WshRpc) runServer() { @@ -291,10 +306,11 @@ func (w *WshRpc) getResponseCh(resId string) chan *RpcMessage { return rd.ResCh } -func (w *WshRpc) SetHandler(handler CommandHandlerFnType) { +func (w *WshRpc) SetServerImpl(serverImpl ServerImpl) { + validateServerImpl(serverImpl) w.Lock.Lock() defer w.Lock.Unlock() - w.HandlerFn = handler + w.ServerImpl = serverImpl } func (w *WshRpc) registerRpc(ctx context.Context, reqId string) chan *RpcMessage { @@ -356,6 +372,7 @@ type RpcRequestHandler struct { ctxCancelFn *atomic.Pointer[context.CancelFunc] reqId string respCh chan *RpcMessage + cachedResp *RpcMessage } func (handler *RpcRequestHandler) Context() context.Context { @@ -379,16 +396,32 @@ func (handler *RpcRequestHandler) SendCancel() { } func (handler *RpcRequestHandler) ResponseDone() bool { + if handler.cachedResp != nil { + return false + } select { - case _, more := <-handler.respCh: - return !more + case msg, more := <-handler.respCh: + if !more { + return true + } + handler.cachedResp = msg + return false default: return false } } func (handler *RpcRequestHandler) NextResponse() (any, error) { - resp := <-handler.respCh + var resp *RpcMessage + if handler.cachedResp != nil { + resp = handler.cachedResp + handler.cachedResp = nil + } else { + resp = <-handler.respCh + } + if resp == nil { + return nil, errors.New("response channel closed") + } if resp.Error != "" { return nil, errors.New(resp.Error) } @@ -527,6 +560,9 @@ func (handler *RpcResponseHandler) IsDone() bool { } func (w *WshRpc) SendComplexRequest(command string, data any, expectsResponse bool, timeoutMs int) (rtnHandler *RpcRequestHandler, rtnErr error) { + if timeoutMs <= 0 { + timeoutMs = DefaultTimeoutMs + } defer func() { if r := recover(); r != nil { log.Printf("panic in SendComplexRequest: %v\n", r) @@ -540,13 +576,9 @@ func (w *WshRpc) SendComplexRequest(command string, data any, expectsResponse bo w: w, ctxCancelFn: &atomic.Pointer[context.CancelFunc]{}, } - if timeoutMs < 0 { - handler.ctx = context.Background() - } else { - var cancelFn context.CancelFunc - handler.ctx, cancelFn = context.WithTimeout(context.Background(), time.Duration(timeoutMs)*time.Millisecond) - handler.ctxCancelFn.Store(&cancelFn) - } + var cancelFn context.CancelFunc + handler.ctx, cancelFn = context.WithTimeout(context.Background(), time.Duration(timeoutMs)*time.Millisecond) + handler.ctxCancelFn.Store(&cancelFn) if expectsResponse { handler.reqId = uuid.New().String() } diff --git a/pkg/wshutil/wshutil.go b/pkg/wshutil/wshutil.go index 7a3d4b29..f2fd882d 100644 --- a/pkg/wshutil/wshutil.go +++ b/pkg/wshutil/wshutil.go @@ -186,11 +186,11 @@ func RestoreTermState() { } // returns (wshRpc, wrappedStdin) -func SetupTerminalRpcClient(handlerFn func(*RpcResponseHandler) bool) (*WshRpc, io.Reader) { +func SetupTerminalRpcClient(serverImpl ServerImpl) (*WshRpc, io.Reader) { messageCh := make(chan []byte, DefaultInputChSize) outputCh := make(chan []byte, DefaultOutputChSize) ptyBuf := MakePtyBuffer(WaveServerOSCPrefix, os.Stdin, messageCh) - rpcClient := MakeWshRpc(messageCh, outputCh, wshrpc.RpcContext{}, handlerFn) + rpcClient := MakeWshRpc(messageCh, outputCh, wshrpc.RpcContext{}, serverImpl) go func() { for msg := range outputCh { barr := EncodeWaveOSCBytes(WaveOSC, msg) @@ -200,7 +200,7 @@ func SetupTerminalRpcClient(handlerFn func(*RpcResponseHandler) bool) (*WshRpc, return rpcClient, ptyBuf } -func SetupConnRpcClient(conn net.Conn, handlerFn func(*RpcResponseHandler) bool) (*WshRpc, chan error, error) { +func SetupConnRpcClient(conn net.Conn, serverImpl ServerImpl) (*WshRpc, chan error, error) { inputCh := make(chan []byte, DefaultInputChSize) outputCh := make(chan []byte, DefaultOutputChSize) writeErrCh := make(chan error, 1) @@ -216,16 +216,16 @@ func SetupConnRpcClient(conn net.Conn, handlerFn func(*RpcResponseHandler) bool) defer conn.Close() AdaptStreamToMsgCh(conn, inputCh) }() - rtn := MakeWshRpc(inputCh, outputCh, wshrpc.RpcContext{}, handlerFn) + rtn := MakeWshRpc(inputCh, outputCh, wshrpc.RpcContext{}, serverImpl) return rtn, writeErrCh, nil } -func SetupDomainSocketRpcClient(sockName string, handlerFn func(*RpcResponseHandler) bool) (*WshRpc, error) { +func SetupDomainSocketRpcClient(sockName string, serverImpl ServerImpl) (*WshRpc, error) { conn, err := net.Dial("unix", sockName) if err != nil { return nil, fmt.Errorf("failed to connect to Unix domain socket: %w", err) } - rtn, errCh, err := SetupConnRpcClient(conn, handlerFn) + rtn, errCh, err := SetupConnRpcClient(conn, serverImpl) go func() { defer conn.Close() err := <-errCh