Add starlark reader implementation
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful

This commit is contained in:
2023-06-22 14:47:06 -07:00
parent 2cdea934eb
commit 228bf59c85
5 changed files with 298 additions and 40 deletions

View File

@@ -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

202
internal/builtins/reader.go Normal file
View File

@@ -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
}