mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-19 16:16:45 +00:00
feat(plugin-sdk): support framed binary transfers
This commit is contained in:
t8y2
@@ -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.
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user