waf-platform/EdgeNode/internal/js/lib_redis.go

// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build script
// +build script

package js

import (
	"context"
	"errors"
	"fmt"
	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
	"github.com/go-redis/redis/v8"
	_ "github.com/go-redis/redis/v8"
	"github.com/iwind/TeaGo/types"
	"time"
)

func init() {
	if !teaconst.IsMain {
		return
	}

	SharedLibraryManager.Register(&JSRedisLibrary{})
}

type JSRedisLibrary struct {
	JSBaseLibrary
}

func (this *JSRedisLibrary) JSNamespace() string {
	return "gojs.redis"
}

func (this *JSRedisLibrary) JSPrototype() string {
	return `gojs.Redis = class {
	#goObjectId

	constructor(options) {
		this.#goObjectId = $this.Connect(options)
	}

	set(key, value, ttl) {
		$this.Set(this.#goObjectId, key, value, ttl)
	}

	incr(key) {
		return $this.Incr(this.#goObjectId, key)
	}

	incrBy(key, increment) {
		return $this.IncrBy(this.#goObjectId, key, increment)
	}

	incrByFloat(key, increment) {
		return $this.IncrByFloat(this.#goObjectId, key, increment)
	}

	get(key) {
		return $this.Get(this.#goObjectId, key)
	}

	close() {
		$this.Close(this.#goObjectId)
	}
}
`
}

func (this *JSRedisLibrary) Connect(arguments *FunctionArguments) (any, error) {
	var optionAddr = "127.0.0.1:6379"
	var optionUsername = ""
	var optionPassword = ""
	var optionDB = 0

	arg, ok := arguments.ArgAt(0)
	if ok && arg.IsObject() {
		var options = arg.Object()
		if options.Has("addr") {
			addr, err := options.Get("addr")
			if err == nil {
				optionAddr = addr.String()
			}
		}
		if options.Has("username") {
			username, err := options.Get("username")
			if err == nil {
				optionUsername = username.String()
			}
		}
		if options.Has("password") {
			password, err := options.Get("password")
			if err == nil {
				optionPassword = password.String()
			}
		}
		if options.Has("db") {
			db, err := options.Get("db")
			if err == nil {
				if db.IsNumber() {
					optionDB = int(db.Int32())
				} else if db.IsString() {
					optionDB = types.Int(db.String())
				}
			}
		}
	}

	var client = redis.NewClient(&redis.Options{
		Network:  "tcp",
		Addr:     optionAddr,
		Username: optionUsername,
		Password: optionPassword,
		DB:       optionDB,
	})

	var objectId = arguments.Context().AddGoObject(client)

	return objectId, nil
}

func (this *JSRedisLibrary) Set(arguments *FunctionArguments) (any, error) {
	goObject, ok := arguments.GoObjectAt(0)
	if !ok {
		return "", nil
	}

	key, ok := arguments.FormatStringAt(1)
	if !ok {
		return "", errors.New("set(): invalid key")
	}

	value, ok := arguments.FormatStringAt(2)
	if !ok {
		return "", errors.New("set(): invalid value")
	}

	ttl, ok := arguments.IntAt(3)
	if !ok {
		ttl = 0
	}

	var client = goObject.(*redis.Client)
	var cmd = client.Set(context.Background(), key, value, time.Duration(ttl))
	if cmd.Err() != nil {
		return nil, fmt.Errorf("set(): %w", cmd.Err())
	}
	return nil, nil
}

func (this *JSRedisLibrary) Incr(arguments *FunctionArguments) (any, error) {
	goObject, ok := arguments.GoObjectAt(0)
	if !ok {
		return "", nil
	}

	key, ok := arguments.FormatStringAt(1)
	if !ok {
		return "", errors.New("incr(): invalid key")
	}

	var client = goObject.(*redis.Client)
	var cmd = client.Incr(context.Background(), key)
	if cmd.Err() != nil {
		return nil, fmt.Errorf("incr(): %w", cmd.Err())
	}
	result, err := cmd.Uint64()
	if err != nil {
		return nil, fmt.Errorf("incr(): %w", err)
	}
	return result, nil
}

func (this *JSRedisLibrary) IncrBy(arguments *FunctionArguments) (any, error) {
	goObject, ok := arguments.GoObjectAt(0)
	if !ok {
		return "", nil
	}

	key, ok := arguments.FormatStringAt(1)
	if !ok {
		return "", errors.New("incrBy(): invalid key")
	}

	increment, ok := arguments.Int64At(2)
	if !ok {
		return "", errors.New("incrBy(): invalid increment")
	}

	var client = goObject.(*redis.Client)
	var cmd = client.IncrBy(context.Background(), key, increment)
	if cmd.Err() != nil {
		return nil, fmt.Errorf("incrBy(): %w", cmd.Err())
	}
	result, err := cmd.Uint64()
	if err != nil {
		return nil, fmt.Errorf("incrBy(): %w", err)
	}
	return result, nil
}

func (this *JSRedisLibrary) IncrByFloat(arguments *FunctionArguments) (any, error) {
	goObject, ok := arguments.GoObjectAt(0)
	if !ok {
		return "", nil
	}

	key, ok := arguments.FormatStringAt(1)
	if !ok {
		return "", errors.New("incrByFloat(): invalid key")
	}

	increment, ok := arguments.Float64At(2)
	if !ok {
		return "", errors.New("incrByFloat(): invalid increment")
	}

	var client = goObject.(*redis.Client)
	var cmd = client.IncrByFloat(context.Background(), key, increment)
	if cmd.Err() != nil {
		return nil, fmt.Errorf("incrByFloat(): %w", cmd.Err())
	}
	return cmd.Val(), nil
}

func (this *JSRedisLibrary) Get(arguments *FunctionArguments) (string, error) {
	goObject, ok := arguments.GoObjectAt(0)
	if !ok {
		return "", nil
	}

	key, ok := arguments.FormatStringAt(1)
	if !ok {
		return "", errors.New("get(): invalid key")
	}

	var client = goObject.(*redis.Client)
	var cmd = client.Get(context.Background(), key)
	if cmd.Err() != nil {
		// the key not found
		if cmd.Err() == redis.Nil {
			return "", nil
		}
		return "", fmt.Errorf("get(): %w", cmd.Err())
	}
	return cmd.Result()
}

func (this *JSRedisLibrary) Close(arguments *FunctionArguments) (any, error) {
	goObject, ok := arguments.GoObjectAt(0)
	if !ok {
		return nil, nil
	}
	err := goObject.(*redis.Client).Close()
	return nil, err
}