Add starlark reader implementation
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
This commit is contained in:
@@ -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
202
internal/builtins/reader.go
Normal 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
|
||||
}
|
||||
Reference in New Issue
Block a user