diff --git a/plugins/sdk/go/dbx-plugin-sdk/README.md b/plugins/sdk/go/dbx-plugin-sdk/README.md index a24a95b3b3..5b984d14c8 100644 --- a/plugins/sdk/go/dbx-plugin-sdk/README.md +++ b/plugins/sdk/go/dbx-plugin-sdk/README.md @@ -1,6 +1,6 @@ # DBX Go plugin SDK -Go SDK for DBX native sidecar plugins using protocol v1 over JSON Lines stdin/stdout. +Go SDK for DBX native sidecar plugins using protocol v1 over JSON Lines or framed stdin/stdout. ```go metadata := dbxpluginsdk.Metadata{ @@ -15,3 +15,12 @@ if err := server.Serve(); err != nil { ``` Keep stdout reserved for protocol messages. Write diagnostics to stderr. + +Use framed transport for binary input and output: + +```go +server := dbxpluginsdk.NewServer(metadata, handler). + WithTransport(dbxpluginsdk.TransportFramed) +``` + +Handlers can implement `HandleBinary(channel string, data []byte, emitter *Emitter) *PluginError` for host-to-plugin frames. Use `Emitter.Binary` for plugin-to-host frames and application-level chunking for large transfers. diff --git a/plugins/sdk/go/dbx-plugin-sdk/sdk.go b/plugins/sdk/go/dbx-plugin-sdk/sdk.go index 7fc4d7b691..fd26517396 100644 --- a/plugins/sdk/go/dbx-plugin-sdk/sdk.go +++ b/plugins/sdk/go/dbx-plugin-sdk/sdk.go @@ -3,6 +3,7 @@ package dbxpluginsdk import ( "bufio" "bytes" + "encoding/binary" "encoding/json" "errors" "fmt" @@ -14,6 +15,15 @@ import ( const ProtocolVersion = 1 const maxJSONBytes = 8 * 1024 * 1024 +const maxBinaryBytes = 64 * 1024 * 1024 +const frameHeaderBytes = 5 + +type Transport int + +const ( + TransportJSONLines Transport = iota + TransportFramed +) type Metadata struct { ID string @@ -44,6 +54,10 @@ type Handler interface { Handle(context RequestContext, method string, params json.RawMessage, emitter *Emitter) (any, *PluginError) } +type BinaryHandler interface { + HandleBinary(channel string, data []byte, emitter *Emitter) *PluginError +} + type HandlerFunc func(context RequestContext, method string, params json.RawMessage, emitter *Emitter) (any, *PluginError) func (handler HandlerFunc) Handle( @@ -56,8 +70,9 @@ func (handler HandlerFunc) Handle( } type Emitter struct { - writer io.Writer - mutex *sync.Mutex + writer io.Writer + mutex *sync.Mutex + transport Transport } func (emitter *Emitter) Event(method string, params any) *PluginError { @@ -71,6 +86,24 @@ func (emitter *Emitter) Event(method string, params any) *PluginError { }) } +func (emitter *Emitter) Binary(channel string, data []byte) *PluginError { + if emitter.transport != TransportFramed { + return NewError(-32000, "Binary messages require framed transport") + } + if !validProtocolName(channel) { + return NewError(-32600, "Invalid binary channel") + } + channelBytes := []byte(channel) + if len(channelBytes) > 65535 || len(data) > maxBinaryBytes { + return NewError(-32600, "Binary message is too large") + } + payload := make([]byte, 2+len(channelBytes)+len(data)) + binary.BigEndian.PutUint16(payload[:2], uint16(len(channelBytes))) + copy(payload[2:], channelBytes) + copy(payload[2+len(channelBytes):], data) + return emitter.writeFrame(frameKindBinary, payload) +} + func (emitter *Emitter) respond(id json.RawMessage, result any, pluginError *PluginError) *PluginError { response := map[string]any{"jsonrpc": "2.0", "id": id} if pluginError != nil { @@ -89,6 +122,9 @@ func (emitter *Emitter) write(value any) *PluginError { if len(payload) > maxJSONBytes { return NewError(-32600, "JSON message is too large") } + if emitter.transport == TransportFramed { + return emitter.writeFrame(frameKindJSON, payload) + } emitter.mutex.Lock() defer emitter.mutex.Unlock() if _, err := emitter.writer.Write(append(payload, '\n')); err != nil { @@ -97,12 +133,36 @@ func (emitter *Emitter) write(value any) *PluginError { return nil } +func (emitter *Emitter) writeFrame(kind byte, payload []byte) *PluginError { + if len(payload) > int(^uint32(0)) { + return NewError(-32600, "Frame payload is too large") + } + header := make([]byte, frameHeaderBytes) + header[0] = kind + binary.BigEndian.PutUint32(header[1:], uint32(len(payload))) + emitter.mutex.Lock() + defer emitter.mutex.Unlock() + if _, err := emitter.writer.Write(header); err != nil { + return NewError(-32000, err.Error()) + } + if _, err := emitter.writer.Write(payload); err != nil { + return NewError(-32000, err.Error()) + } + if flusher, ok := emitter.writer.(interface{ Flush() error }); ok { + if err := flusher.Flush(); err != nil { + return NewError(-32000, err.Error()) + } + } + return nil +} + type Server struct { - metadata Metadata - handler Handler - input io.Reader - output io.Writer - errors io.Writer + metadata Metadata + handler Handler + transport Transport + input io.Reader + output io.Writer + errors io.Writer } func NewServer(metadata Metadata, handler Handler) *Server { @@ -122,6 +182,11 @@ func (server *Server) WithIO(input io.Reader, output io.Writer, errorsWriter io. return server } +func (server *Server) WithTransport(transport Transport) *Server { + server.transport = transport + return server +} + func (server *Server) Serve() error { if server.handler == nil { return errors.New("plugin handler is required") @@ -129,7 +194,10 @@ func (server *Server) Serve() error { if !validProtocolName(server.metadata.ID) { return errors.New("plugin id is invalid") } - emitter := &Emitter{writer: server.output, mutex: &sync.Mutex{}} + emitter := &Emitter{writer: server.output, mutex: &sync.Mutex{}, transport: server.transport} + if server.transport == TransportFramed { + return server.serveFramed(emitter) + } scanner := bufio.NewScanner(server.input) scanner.Buffer(make([]byte, 64*1024), maxJSONBytes) var workers sync.WaitGroup @@ -178,6 +246,100 @@ func (server *Server) Serve() error { return scanner.Err() } +func (server *Server) serveFramed(emitter *Emitter) error { + var workers sync.WaitGroup + for { + var header [frameHeaderBytes]byte + if _, err := io.ReadFull(server.input, header[:]); err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + workers.Wait() + return nil + } + return err + } + kind := header[0] + length := int(binary.BigEndian.Uint32(header[1:])) + maximum := maxJSONBytes + if kind == frameKindBinary { + maximum = maxBinaryBytes + 2 + 256 + } else if kind != frameKindJSON { + return errors.New("unknown plugin frame kind") + } + if length > maximum { + return errors.New("plugin frame is too large") + } + payload := make([]byte, length) + if _, err := io.ReadFull(server.input, payload); err != nil { + return err + } + if kind == frameKindBinary { + if pluginError := server.dispatchBinary(payload, emitter); pluginError != nil { + fmt.Fprintf(server.errors, "[dbx-plugin-sdk-go] %s\n", pluginError.Message) + } + continue + } + if err := server.dispatchJSON(payload, emitter, &workers); err != nil { + fmt.Fprintf(server.errors, "[dbx-plugin-sdk-go] %v\n", err) + } + } +} + +func (server *Server) dispatchJSON(payload []byte, emitter *Emitter, workers *sync.WaitGroup) error { + request, err := decodeRequest(payload) + if err != nil { + return err + } + if request.Method == "plugin/initialize" { + if len(request.ID) == 0 { + return errors.New("plugin/initialize must be a request") + } + result, pluginError := server.initialize(request.Params) + if writeError := emitter.respond(request.ID, result, pluginError); writeError != nil { + return errors.New(writeError.Message) + } + return nil + } + workers.Add(1) + go func(request protocolRequest) { + defer workers.Done() + result, pluginError := server.handler.Handle( + RequestContext{RequestID: request.ID, Driver: request.Driver}, + request.Method, + request.Params, + emitter, + ) + if len(request.ID) == 0 { + if pluginError != nil { + fmt.Fprintf(server.errors, "[dbx-plugin-sdk-go] %s\n", pluginError.Message) + } + return + } + if writeError := emitter.respond(request.ID, result, pluginError); writeError != nil { + fmt.Fprintf(server.errors, "[dbx-plugin-sdk-go] failed to write response: %s\n", writeError.Message) + } + }(request) + return nil +} + +func (server *Server) dispatchBinary(payload []byte, emitter *Emitter) *PluginError { + if len(payload) < 2 { + return NewError(-32600, "Invalid binary frame") + } + channelLength := int(binary.BigEndian.Uint16(payload[:2])) + if channelLength == 0 || len(payload) < 2+channelLength { + return NewError(-32600, "Invalid binary channel") + } + channel := string(payload[2 : 2+channelLength]) + if !validProtocolName(channel) { + return NewError(-32600, "Invalid binary channel") + } + handler, ok := server.handler.(BinaryHandler) + if !ok { + return NewError(-32601, "Binary input is not supported") + } + return handler.HandleBinary(channel, append([]byte(nil), payload[2+channelLength:]...), emitter) +} + func (server *Server) initialize(params json.RawMessage) (any, *PluginError) { var request struct { Host struct { @@ -242,3 +404,8 @@ func validProtocolName(value string) bool { } return true } + +const ( + frameKindJSON = 0 + frameKindBinary = 1 +) diff --git a/plugins/sdk/go/dbx-plugin-sdk/sdk_test.go b/plugins/sdk/go/dbx-plugin-sdk/sdk_test.go index b8aa97ca92..6d4a201f2c 100644 --- a/plugins/sdk/go/dbx-plugin-sdk/sdk_test.go +++ b/plugins/sdk/go/dbx-plugin-sdk/sdk_test.go @@ -2,11 +2,24 @@ package dbxpluginsdk import ( "bytes" + "encoding/binary" "encoding/json" "sync" "testing" ) +type binaryTestHandler struct { + HandlerFunc + channel string + data []byte +} + +func (handler *binaryTestHandler) HandleBinary(channel string, data []byte, _ *Emitter) *PluginError { + handler.channel = channel + handler.data = append(handler.data, data...) + return nil +} + func TestServerInitializesAndDispatches(t *testing.T) { input := bytes.NewBufferString( "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"plugin/initialize\",\"params\":{\"host\":{\"protocolVersions\":[1]}}}\n" + @@ -60,3 +73,36 @@ func TestEmitterWritesEvents(t *testing.T) { t.Fatalf("unexpected event: %#v", event) } } + +func TestFramedServerDispatchesBinaryInput(t *testing.T) { + input := &bytes.Buffer{} + writeFrame := func(kind byte, payload []byte) { + input.WriteByte(kind) + _ = binary.Write(input, binary.BigEndian, uint32(len(payload))) + input.Write(payload) + } + initialize, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": "plugin/initialize", "params": map[string]any{"host": map[string]any{"protocolVersions": []int{1}}}}) + writeFrame(frameKindJSON, initialize) + channel := []byte("sample.upload") + binaryPayload := make([]byte, 2+len(channel)+3) + binary.BigEndian.PutUint16(binaryPayload[:2], uint16(len(channel))) + copy(binaryPayload[2:], channel) + copy(binaryPayload[2+len(channel):], []byte("abc")) + writeFrame(frameKindBinary, binaryPayload) + var output bytes.Buffer + handler := &binaryTestHandler{HandlerFunc: HandlerFunc(func(_ RequestContext, _ string, _ json.RawMessage, _ *Emitter) (any, *PluginError) { + return nil, nil + })} + server := NewServer(Metadata{ID: "sample.plugin", Version: "1.0.0"}, handler). + WithTransport(TransportFramed). + WithIO(input, &output, &bytes.Buffer{}) + if err := server.Serve(); err != nil { + t.Fatal(err) + } + if handler.channel != "sample.upload" || string(handler.data) != "abc" { + t.Fatalf("unexpected binary callback: channel=%q data=%q", handler.channel, handler.data) + } + if output.Len() == 0 { + t.Fatal("expected framed initialize response") + } +}