sparrow/services/registry/registry.go

package main

import (
	"errors"
	"flag"
	"fmt"
	"sparrow/pkg/generator"
	"sparrow/pkg/models"
	"sparrow/pkg/rpcs"

	opentracing "github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
)

const (
	flagAESKey = "aeskey"
)

var confAESKey = flag.String(flagAESKey, "", "use your own aes encryting key.")

var errDbActFailt = errors.New("数据库操作失败")

type Registry struct {
	keygen *generator.KeyGenerator
}

// NewRegistry create new registry
func NewRegistry() (*Registry, error) {
	gen, err := generator.NewKeyGenerator(*confAESKey)
	if err != nil {
		return nil, err
	}
	return &Registry{
		keygen: gen,
	}, nil
}

func setVendor(target *models.Vendor, src *models.Vendor) {
	target.ID = src.ID
	target.VendorName = src.VendorName
	target.VendorDescription = src.VendorDescription
	target.VendorKey = src.VendorKey
	target.CreatedAt = src.CreatedAt
	target.UpdatedAt = src.UpdatedAt
}

func setProduct(target *models.Product, src *models.Product) {
	target.ID = src.ID
	target.ProductName = src.ProductName
	target.ProductDescription = src.ProductDescription
	target.ProductKey = src.ProductKey
	target.ProductConfig = src.ProductConfig
	target.VendorID = src.VendorID
	target.CreatedAt = src.CreatedAt
	target.UpdatedAt = src.UpdatedAt
}

func setApplication(target *models.Application, src *models.Application) {
	target.ID = src.ID
	target.AppName = src.AppName
	target.AppDescription = src.AppDescription
	target.AppKey = src.AppKey
	target.ReportUrl = src.ReportUrl
	target.AppToken = src.AppToken
	target.AppDomain = src.AppDomain
	target.CreatedAt = src.CreatedAt
	target.UpdatedAt = src.UpdatedAt
}

func setDevice(target *models.Device, src *models.Device) {
	target.ID = src.ID
	target.ProductID = src.ProductID
	target.DeviceIdentifier = src.DeviceIdentifier
	target.DeviceSecret = src.DeviceIdentifier
	target.DeviceKey = src.DeviceKey
	target.DeviceName = src.DeviceName
	target.DeviceDescription = src.DeviceDescription
	target.DeviceVersion = src.DeviceVersion
	target.CreatedAt = src.CreatedAt
	target.UpdatedAt = src.UpdatedAt
}

// SaveVendor will create a vendor if the ID field is not initialized
// if ID field is initialized, it will update the conresponding vendor.
func (r *Registry) SaveVendor(vendor *models.Vendor, reply *models.Vendor) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	if vendor.ID == 0 {
		// if ID field is not initialized, will generate key first
		err = db.Save(vendor).Error
		if err != nil {
			return err
		}

		key, err := r.keygen.GenRandomKey(int64(vendor.ID))
		if err != nil {
			return err
		}

		vendor.VendorKey = key
	}

	err = db.Save(vendor).Error
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Vendor:%v", vendor.ID)
	if _, ok := cache.Get(cacheKey); ok {
		cache.Delete(cacheKey)
	}
	// write cache
	cache.Set(cacheKey, vendor)
	setVendor(reply, vendor)

	return nil
}

// SaveApplication will create a application if the ID field is not initialized
// if ID field is initialized, it will update the conresponding application.
func (r *Registry) SaveApplication(app *models.Application, reply *models.Application) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	if app.ID == 0 {
		err = db.Save(app).Error
		if err != nil {
			return err
		}

		key, err := r.keygen.GenRandomKey(int64(app.ID))
		if err != nil {
			return err
		}

		app.AppKey = key
	}

	err = db.Save(app).Error
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Application:%v", app.ID)
	if _, ok := cache.Get(cacheKey); ok {
		cache.Delete(cacheKey)
	}

	setApplication(reply, app)

	return nil
}

// ValidateApplication try to validate the given app key.
// if success, it will reply the corresponding application
func (r *Registry) ValidateApplication(key string, reply *models.Application) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	id, err := r.keygen.DecodeIDFromRandomKey(key)
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Application:%v", id)
	if cacheValue, ok := cache.Get(cacheKey); ok {
		app := cacheValue.(*models.Application)
		setApplication(reply, app)
	} else {
		err = db.First(reply, id).Error
		if err != nil {
			return err
		}
		var storage models.Application
		storage = *reply
		cache.Set(cacheKey, &storage)
	}

	if reply.AppKey != key {
		return errors.New("app key not match.")
	}

	return nil
}

// FindVendor will find product by specified ID
func (r *Registry) FindVendor(id int32, reply *models.Vendor) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Vendor:%v", id)
	if cacheValue, ok := cache.Get(cacheKey); ok {
		vendor := cacheValue.(*models.Vendor)
		setVendor(reply, vendor)
	} else {
		err = db.First(reply, id).Error
		if err != nil {
			return err
		}
		var storage models.Vendor
		storage = *reply
		cache.Set(cacheKey, &storage)
	}

	return nil
}

// GetVendors will get all vendors in the platform.
func (r *Registry) GetVendors(noarg int, reply *[]models.Vendor) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	return db.Find(reply).Error
}

// GetApplications will get all applications in the platform.
func (r *Registry) GetApplications(noarg int, reply *[]models.Application) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	return db.Find(reply).Error
}

// FindApplication will find product by specified ID
func (r *Registry) FindApplication(id int32, reply *models.Application) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Application:%v", id)
	if cacheValue, ok := cache.Get(cacheKey); ok {
		app := cacheValue.(*models.Application)
		setApplication(reply, app)
	} else {
		err = db.First(reply, id).Error
		if err != nil {
			return err
		}
		var storage models.Application
		storage = *reply
		cache.Set(cacheKey, &storage)
	}

	return nil
}

// RegisterDevice try to register a device to our platform.
// if the device has already been registered,
// the registration will success return the registered device before.
func (r *Registry) RegisterDevice(args *rpcs.ArgsDeviceRegister, reply *models.Device) error {
	spanCtx, _ := opentracing.GlobalTracer().Extract(opentracing.TextMap, opentracing.TextMapCarrier(args.SpanCarrier))
	span := opentracing.StartSpan("RegisterDevice", ext.RPCServerOption(spanCtx))
	defer span.Finish()
	span.SetTag("args", args)
	db, err := getDB()
	if err != nil {
		return err
	}

	product := &models.Product{}
	err = r.ValidateProduct(args.ProductKey, product)
	if err != nil {
		return err
	}

	identifier := genDeviceIdentifier(product.VendorID, product.ID, args.DeviceCode)
	if db.Where(&models.Device{DeviceIdentifier: identifier}).First(reply).RecordNotFound() {
		// device is not registered yet.
		reply.ProductID = int32(product.ID)
		reply.DeviceIdentifier = identifier
		reply.DeviceName = product.ProductName // product name as default device name.
		reply.DeviceDescription = product.ProductDescription
		reply.DeviceVersion = args.DeviceVersion
		reply.VendorID = uint(product.VendorID)
		reply.ModuleName = args.ModuleName
		err = db.Save(reply).Error
		if err != nil {
			return err
		}
		// generate a random device key with hex encoding.
		reply.DeviceKey, err = r.keygen.GenRandomKey(int64(reply.ID))
		if err != nil {
			return err
		}
		// generate a random password with base64 encoding.
		reply.DeviceSecret, err = generator.GenRandomPassword()
		if err != nil {
			return err
		}

		err = db.Save(reply).Error
		if err != nil {
			return err
		}
	} else {

		//delete cache
		cache := getCache()
		cacheKey := fmt.Sprintf("Device:%v", identifier)
		if _, ok := cache.Get(cacheKey); ok {
			cache.Delete(cacheKey)
		}

		// device has aleady been saved. just update version info.
		reply.DeviceVersion = args.DeviceVersion
		err = db.Save(reply).Error
		if err != nil {
			return err
		}
	}
	span.LogKV("device_identifier", reply.DeviceIdentifier)
	span.LogKV("device_key", reply.DeviceKey)
	span.LogKV("device_secret", reply.DeviceSecret)
	return nil
}

// FindDeviceByIdentifier will find the device by indentifier
func (r *Registry) FindDeviceByIdentifier(identifier string, reply *models.Device) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Device:%v", identifier)
	if cacheValue, ok := cache.Get(identifier); ok {
		device := cacheValue.(*models.Device)
		setDevice(reply, device)
	} else {
		err = db.Where(&models.Device{
			DeviceIdentifier: identifier,
		}).First(reply).Error
		if err != nil {
			return err
		}
		var storage models.Device
		storage = *reply
		cache.Set(cacheKey, &storage)
	}

	return nil
}

// FindDeviceById will find the device with given id
func (r *Registry) FindDeviceById(args *rpcs.ArgsDeviceAuth, reply *models.Device) error {
	spanCtx, _ := opentracing.GlobalTracer().Extract(opentracing.TextMap, opentracing.TextMapCarrier(args.SpanCarrier))
	span := opentracing.StartSpan("FindDeviceById", ext.RPCServerOption(spanCtx))
	defer span.Finish()
	span.SetTag("args", args)
	db, err := getDB()
	if err != nil {
		return err
	}
	d := &models.Device{}
	d.ID = uint(args.DeviceID)
	err = db.Where(d).First(reply).Error

	if err != nil {
		return err
	}
	span.LogKV("device", reply)
	return nil
}

// ValidateDevice will validate a device key and return the model if success.
func (r *Registry) ValidateDevice(key string, device *models.Device) error {
	id, err := r.keygen.DecodeIDFromRandomKey(key)
	if err != nil {
		return err
	}
	args := rpcs.ArgsDeviceAuth{
		DeviceID: id,
	}
	err = r.FindDeviceById(&args, device)
	if err != nil {
		return err
	}

	if device.DeviceKey != key {
		return errors.New("device key not match.")
	}

	return nil
}

// UpdateDeviceInfo will update a device info by identifier
func (r *Registry) UpdateDeviceInfo(args *rpcs.ArgsDeviceUpdate, reply *models.Device) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	err = r.FindDeviceByIdentifier(args.DeviceIdentifier, reply)
	if err != nil {
		return err
	}

	//delete cache
	cache := getCache()
	cacheKey := fmt.Sprintf("Device:%v", args.DeviceIdentifier)
	if _, ok := cache.Get(cacheKey); ok {
		cache.Delete(cacheKey)
	}

	reply.DeviceName = args.DeviceName
	reply.DeviceDescription = args.DeviceDescription

	err = db.Save(reply).Error
	if err != nil {
		return err
	}

	return nil
}

// CreateRule create a new rule with specified parameters.
func (r *Registry) CreateRule(args *models.Rule, reply *rpcs.ReplyEmptyResult) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	return db.Save(args).Error
}

// QueryRules queries rules by trigger and rule type.
func (r *Registry) QueryRules(args *models.Rule, reply *[]models.Rule) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	err = db.Where(args).Find(reply).Error
	if err != nil {
		return err
	}

	return nil
}