From 228bf59c8511a1b93f05684f1186b702fb7c710f Mon Sep 17 00:00:00 2001 From: Elara Musayelyan Date: Thu, 22 Jun 2023 14:47:06 -0700 Subject: [PATCH] Add starlark reader implementation --- go.mod | 2 + go.sum | 5 + internal/builtins/http.go | 42 +------- internal/builtins/reader.go | 202 ++++++++++++++++++++++++++++++++++++ internal/convert/convert.go | 87 ++++++++++++++++ 5 files changed, 298 insertions(+), 40 deletions(-) create mode 100644 internal/builtins/reader.go create mode 100644 internal/convert/convert.go diff --git a/go.mod b/go.mod index ae52875..1efcc90 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,7 @@ require ( github.com/go-git/go-git/v5 v5.7.0 github.com/pelletier/go-toml/v2 v2.0.8 github.com/spf13/pflag v1.0.5 + github.com/vmihailenco/msgpack/v5 v5.3.5 go.elara.ws/logger v0.0.0-20230421022458-e80700db2090 go.elara.ws/pcre v0.0.0-20230421030233-daf2d2e6973f go.etcd.io/bbolt v1.3.7 @@ -34,6 +35,7 @@ require ( github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect github.com/sergi/go-diff v1.1.0 // indirect github.com/skeema/knownhosts v1.1.1 // indirect + github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect github.com/xanzy/ssh-agent v0.3.3 // indirect github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778 // indirect golang.org/x/net v0.10.0 // indirect diff --git a/go.sum b/go.sum index 0246f36..9269ced 100644 --- a/go.sum +++ b/go.sum @@ -104,11 +104,16 @@ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSS github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.2/go.mod h1:R6va5+xMeoiuVRoj+gSkQ7d3FALtqAAGI1FQKckRals= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.3 h1:RP3t2pwF7cMEbC1dqtB6poj3niw/9gnV4Cjg5oW5gtY= github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/vmihailenco/msgpack/v5 v5.3.5 h1:5gO0H1iULLWGhs2H5tbAHIZTV8/cYafcFOr9znI5mJU= +github.com/vmihailenco/msgpack/v5 v5.3.5/go.mod h1:7xyJ9e+0+9SaZT0Wt1RGleJXzli6Q/V5KbhBonMG9jc= +github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g= +github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds= github.com/xanzy/ssh-agent v0.3.3 h1:+/15pJfg/RsTxqYcX6fHqOXZwwMP+2VyYWJeWM2qQFM= github.com/xanzy/ssh-agent v0.3.3/go.mod h1:6dzNDKs0J9rVPHPhaGCukekBHKqfl+L3KghI1Bc68Uw= github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778 h1:QldyIu/L63oPpyvQmHgvgickp1Yw510KJOqX7H24mg8= diff --git a/internal/builtins/http.go b/internal/builtins/http.go index baf9fa5..354d1f7 100644 --- a/internal/builtins/http.go +++ b/internal/builtins/http.go @@ -174,7 +174,7 @@ func starlarkResponse(res *http.Response) *starlarkstruct.Struct { return starlarkstruct.FromStringDict(starlarkstruct.Default, starlark.StringDict{ "code": starlark.MakeInt(res.StatusCode), "headers": starlarkStringSliceMap(res.Header), - "body": starlarkBody(res.Body), + "body": newStarlarkReader(res.Body), }) } @@ -184,7 +184,7 @@ func starlarkRequest(req *http.Request) *starlarkstruct.Struct { "remote_addr": starlark.String(req.RemoteAddr), "headers": starlarkStringSliceMap(req.Header), "query": starlarkStringSliceMap(req.URL.Query()), - "body": starlarkBody(req.Body), + "body": newStarlarkReader(req.Body), }) } @@ -200,44 +200,6 @@ func starlarkStringSliceMap(ssm map[string][]string) *starlark.Dict { return dict } -func starlarkBody(body io.ReadCloser) *starlarkstruct.Struct { - return starlarkstruct.FromStringDict(starlarkstruct.Default, starlark.StringDict{ - "string": bodyAsString(body), - "bytes": bodyAsBytes(body), - "close": bodyClose(body), - }) -} - -func bodyAsBytes(body io.ReadCloser) *starlark.Builtin { - return starlark.NewBuiltin("http.response.body.bytes", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { - data, err := io.ReadAll(io.LimitReader(body, maxBodySize)) - if err != nil { - return nil, err - } - return starlark.Bytes(data), nil - }) -} - -func bodyAsString(body io.ReadCloser) *starlark.Builtin { - return starlark.NewBuiltin("http.response.body.string", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { - data, err := io.ReadAll(io.LimitReader(body, maxBodySize)) - if err != nil { - return nil, err - } - return starlark.String(data), nil - }) -} - -func bodyClose(body io.ReadCloser) *starlark.Builtin { - return starlark.NewBuiltin("http.response.body.close", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { - err := body.Close() - if err != nil { - return nil, err - } - return starlark.None, nil - }) -} - func registerWebhook(mux *http.ServeMux, cfg *config.Config, pluginName string) *starlark.Builtin { return starlark.NewBuiltin("register_webhook", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var fn *starlark.Function diff --git a/internal/builtins/reader.go b/internal/builtins/reader.go new file mode 100644 index 0000000..aab96ba --- /dev/null +++ b/internal/builtins/reader.go @@ -0,0 +1,202 @@ +package builtins + +import ( + "bufio" + "encoding/json" + "io" + + "github.com/vmihailenco/msgpack/v5" + "go.elara.ws/lure-updater/internal/convert" + "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" +) + +type starlarkReader struct { + closeFunc func() error + br *bufio.Reader + *starlarkstruct.Struct +} + +func newStarlarkReader(r io.Reader) starlarkReader { + sr := starlarkReader{br: bufio.NewReader(r)} + + if rc, ok := r.(io.ReadCloser); ok { + sr.closeFunc = rc.Close + } + + sr.Struct = starlarkstruct.FromStringDict(starlark.String("regex"), starlark.StringDict{ + "read": starlark.NewBuiltin("reader.read", sr.read), + "peek": starlark.NewBuiltin("reader.peek", sr.peek), + "discard": starlark.NewBuiltin("reader.discard", sr.discard), + "read_string": starlark.NewBuiltin("reader.read_string", sr.readString), + "read_until": starlark.NewBuiltin("reader.read_until", sr.readUntil), + "read_string_until": starlark.NewBuiltin("reader.read_string_until", sr.readStringUntil), + "read_all": starlark.NewBuiltin("reader.read_all", sr.readAll), + "read_all_string": starlark.NewBuiltin("reader.read_all_string", sr.readAllString), + "read_json": starlark.NewBuiltin("reader.read_json", sr.readJSON), + "read_msgpack": starlark.NewBuiltin("read_msgpack", sr.readMsgpack), + "close": starlark.NewBuiltin("reader.close", sr.closeReader), + }) + + return sr +} + +func (sr starlarkReader) read(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var n int + err := starlark.UnpackArgs("reader.read", args, kwargs, "n", &n) + if err != nil { + return nil, err + } + + buf := make([]byte, n) + _, err = io.ReadFull(sr.br, buf) + if err != nil { + return nil, err + } + + return starlark.Bytes(buf), nil +} + +func (sr starlarkReader) readString(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var n int + err := starlark.UnpackArgs("reader.read_string", args, kwargs, "n", &n) + if err != nil { + return nil, err + } + + buf := make([]byte, n) + _, err = io.ReadFull(sr.br, buf) + if err != nil { + return nil, err + } + + return starlark.String(buf), nil +} + +func (sr starlarkReader) readUntil(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var delimiter string + err := starlark.UnpackArgs("reader.read_until", args, kwargs, "delimiter", &delimiter) + if err != nil { + return nil, err + } + + buf, err := sr.br.ReadBytes(delimiter[0]) + if err != nil { + return nil, err + } + + return starlark.Bytes(buf), nil +} + +func (sr starlarkReader) readStringUntil(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var delimiter string + err := starlark.UnpackArgs("reader.read_string_until", args, kwargs, "delimiter", &delimiter) + if err != nil { + return nil, err + } + + buf, err := sr.br.ReadString(delimiter[0]) + if err != nil { + return nil, err + } + + return starlark.String(buf), nil +} + +func (sr starlarkReader) peek(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var n int + err := starlark.UnpackArgs("reader.peek", args, kwargs, "n", &n) + if err != nil { + return nil, err + } + + buf, err := sr.br.Peek(n) + if err != nil { + return nil, err + } + + return starlark.Bytes(buf), nil +} + +func (sr starlarkReader) discard(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var n int + err := starlark.UnpackArgs("reader.discard", args, kwargs, "n", &n) + if err != nil { + return nil, err + } + + dn, err := sr.br.Discard(n) + if err != nil { + return nil, err + } + + return starlark.MakeInt(dn), nil +} + +func (sr starlarkReader) readAll(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var limit int64 = 102400 + err := starlark.UnpackArgs("reader.read_all", args, kwargs, "limit??", &limit) + if err != nil { + return nil, err + } + + var r io.Reader = sr.br + if limit > 0 { + r = io.LimitReader(sr.br, limit) + } + + buf, err := io.ReadAll(r) + if err != nil { + return nil, err + } + + return starlark.Bytes(buf), nil +} + +func (sr starlarkReader) readAllString(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var limit int64 = 102400 + err := starlark.UnpackArgs("reader.read_all_string", args, kwargs, "limit??", &limit) + if err != nil { + return nil, err + } + + var r io.Reader = sr.br + if limit > 0 { + r = io.LimitReader(sr.br, limit) + } + + buf, err := io.ReadAll(r) + if err != nil { + return nil, err + } + + return starlark.String(buf), nil +} + +func (sr starlarkReader) readJSON(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var v any + err := json.NewDecoder(sr.br).Decode(&v) + if err != nil { + return nil, err + } + return convert.Convert(v) +} + +func (sr starlarkReader) readMsgpack(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var v any + err := msgpack.NewDecoder(sr.br).Decode(&v) + if err != nil { + return nil, err + } + return convert.Convert(v) +} + +func (sr starlarkReader) closeReader(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + if sr.closeFunc != nil { + err := sr.closeFunc() + if err != nil { + return nil, err + } + } + return starlark.None, nil +} diff --git a/internal/convert/convert.go b/internal/convert/convert.go new file mode 100644 index 0000000..4a1c75a --- /dev/null +++ b/internal/convert/convert.go @@ -0,0 +1,87 @@ +package convert + +import ( + "errors" + "fmt" + "reflect" + + "go.starlark.net/starlark" +) + +var ErrInvalidType = errors.New("unknown type") + +func Convert(v any) (starlark.Value, error) { + if v == nil { + return starlark.None, nil + } + val := reflect.ValueOf(v) + kind := val.Kind() + for kind == reflect.Pointer || kind == reflect.Interface { + val = val.Elem() + } + return convert(val) +} + +func convert(val reflect.Value) (starlark.Value, error) { + switch val.Kind() { + case reflect.Interface: + return convert(val.Elem()) + case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int16, reflect.Int8: + return starlark.MakeInt64(val.Int()), nil + case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16, reflect.Uint8: + return starlark.MakeUint64(val.Uint()), nil + case reflect.Float64, reflect.Float32: + return starlark.Float(val.Float()), nil + case reflect.Bool: + return starlark.Bool(val.Bool()), nil + case reflect.String: + return starlark.String(val.String()), nil + case reflect.Slice, reflect.Array: + return convertSlice(val) + case reflect.Map: + return convertMap(val) + default: + return nil, fmt.Errorf("%w: %s", ErrInvalidType, val.Type()) + } +} + +func convertSlice(val reflect.Value) (starlark.Value, error) { + // Detect byte slice + if val.Type().Elem().Kind() == reflect.Uint8 { + return starlark.Bytes(val.Bytes()), nil + } + + elems := make([]starlark.Value, val.Len()) + + for i := 0; i < val.Len(); i++ { + elem, err := convert(val.Index(i)) + if err != nil { + return nil, err + } + elems[i] = elem + } + + return starlark.NewList(elems), nil +} + +func convertMap(val reflect.Value) (starlark.Value, error) { + dict := starlark.NewDict(val.Len()) + iter := val.MapRange() + for iter.Next() { + k, err := convert(iter.Key()) + if err != nil { + return nil, err + } + + v, err := convert(iter.Value()) + if err != nil { + return nil, err + } + + err = dict.SetKey(k, v) + if err != nil { + return nil, err + } + } + return dict, nil +}