Add starlark reader implementation

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

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

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