kos/entry/cli/context.go

package cli

import (
	"context"
	"fmt"
	"io"
	"math"
	"sync"
)

type Context struct {
	Id        int64
	seq       uint16
	ctx       context.Context
	wc        io.WriteCloser
	params    map[string]string
	locker    sync.RWMutex
	variables map[string]any
	args      []string
}

func (ctx *Context) reset(id int64, wc io.WriteCloser) {
	ctx.Id = id
	ctx.wc = wc
	ctx.seq = 0
	ctx.ctx = context.Background()
	ctx.args = make([]string, 0)
	ctx.params = make(map[string]string)
	ctx.variables = make(map[string]any)
}

func (ctx *Context) setArgs(args []string) {
	ctx.args = args
}

func (ctx *Context) setParam(ps map[string]string) {
	ctx.params = ps
}

func (ctx *Context) Bind(v any) (err error) {
	return
}

func (ctx *Context) setContext(c context.Context) {
	ctx.ctx = c
}

func (ctx *Context) Context() context.Context {
	return ctx.ctx
}

func (ctx *Context) Argument(index int) string {
	if index >= len(ctx.args) || index < 0 {
		return ""
	}
	return ctx.args[index]
}

func (ctx *Context) Param(s string) string {
	if v, ok := ctx.params[s]; ok {
		return v
	}
	return ""
}

func (ctx *Context) SetValue(name string, value any) {
	ctx.locker.Lock()
	if ctx.variables == nil {
		ctx.variables = make(map[string]any)
	}
	ctx.variables[name] = value
	ctx.locker.Unlock()
}

func (ctx *Context) GetValue(name string) (val any, ok bool) {
	ctx.locker.RLock()
	defer ctx.locker.RUnlock()
	val, ok = ctx.variables[name]
	return
}

func (ctx *Context) Success(v any) (err error) {
	return ctx.send(responsePayload{Type: PacketTypeCommand, Data: v})
}

func (ctx *Context) Error(code int, reason string) (err error) {
	return ctx.send(responsePayload{Type: PacketTypeCommand, Code: code, Reason: reason})
}

func (ctx *Context) Close() (err error) {
	return ctx.wc.Close()
}

func (ctx *Context) send(res responsePayload) (err error) {
	var (
		ok      bool
		buf     []byte
		marshal encoder
	)
	if res.Code > 0 {
		err = writeFrame(ctx.wc, &Frame{
			Feature: Feature,
			Type:    res.Type,
			Seq:     ctx.seq,
			Flag:    FlagComplete,
			Error:   fmt.Sprintf("ERROR(%d): %s", res.Code, res.Reason),
		})
		return
	}
	if res.Data == nil {
		buf = OK
		goto __END
	}
	if marshal, ok = res.Data.(encoder); ok {
		buf, err = marshal.Marshal()
		goto __END
	}
	buf, err = serialize(res.Data)
__END:
	if err != nil {
		return
	}
	offset := 0
	chunkSize := math.MaxInt16 - 1
	n := len(buf) / chunkSize
	for i := 0; i < n; i++ {
		if err = writeFrame(ctx.wc, newFrame(res.Type, FlagPortion, ctx.seq, 0, buf[offset:chunkSize+offset])); err != nil {
			return
		}
		offset += chunkSize
	}
	err = writeFrame(ctx.wc, newFrame(res.Type, FlagComplete, ctx.seq, 0, buf[offset:]))
	return
}

func newContext(id int64, wc io.WriteCloser) *Context {
	return &Context{
		Id: id,
		wc: wc,
	}
}