<commit_before>package koding

import (
	"fmt"
	"koding/db/mongodb"
	"koding/kites/kloud/klient"
	"strconv"
	"strings"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	"github.com/koding/kite"
	aws "github.com/koding/kloud/api/amazon"
	"github.com/koding/kloud/machinestate"
	"github.com/koding/kloud/protocol"
	"github.com/koding/kloud/provider/amazon"
	"github.com/koding/logging"
	"github.com/mitchellh/goamz/ec2"
)

// Checker checks various aspects of a machine. It used to limit certain
// aspects of a machine, such a the total machine or total storage.
type Checker interface {
	// Total checks whether the user has reached the current plan's limit of
	// having a total number numbers of machines. It returns an error if the
	// limit is reached or an unexplained error happaned.
	Total() error

	// AlwaysOn checks whether the given machine has reached the current plans
	// always on limit
	AlwaysOn() error

	// Timeout checks whether the user has reached the current plan's
	// inactivity timeout.
	Timeout() error

	// Storage checks whether the user has reached the current plan's limit
	// total storage with the supplied wantStorage information. It returns an
	// error if the limit is reached or an unexplained error happaned.
	Storage(wantStorage int) error

	// AllowedInstances checks whether the given machine has the permisison to
	// create the given instance type
	AllowedInstances(wantInstance InstanceType) error
}

type PlanChecker struct {
	api      *amazon.AmazonClient
	db       *mongodb.MongoDB
	machine  *protocol.Machine
	provider *Provider
	kite     *kite.Kite
	username string
	log      logging.Logger
}

// PlanChecker creates and returns a new Checker interface that is responsible
// of checking various pieces of informations based on a Plan
func (p *Provider) PlanChecker(opts *protocol.Machine) (Checker, error) {
	a, err := p.NewClient(opts)
	if err != nil {
		return nil, err
	}

	checker := &PlanChecker{
		api:      a,
		provider: p,
		db:       p.Session,
		kite:     p.Kite,
		username: opts.Builder["username"].(string),
		log:      p.Log,
		machine:  opts,
	}

	return checker, nil
}

// Plan returns user's current plan
func (p *PlanChecker) Plan() (Plan, error) {
	var account = struct {
		GlobalFlags []string `bson:"globalFlags"`
	}{}
	if err := p.db.Run("jAccounts", func(c *mgo.Collection) error {
		return c.Find(bson.M{"profile.nickname": p.username}).One(&account)
	}); err != nil {
		return 0, err
	}

	if len(account.GlobalFlags) == 0 {
		p.log.Warning("[%s] retrieving plan failed, no flag defined. Using free plan",
			p.machine.MachineId)
		return Free, nil
	}

	splitted := strings.Split(account.GlobalFlags[0], "-")
	if len(splitted) != 2 {
		p.log.Warning("[%s] retrieving plan failed, flag '%v' malformed. Using free plan",
			p.machine.MachineId, account.GlobalFlags)
		return Free, nil
	}

	plan, ok := plans[strings.Title(splitted[1])]
	if !ok {
		p.log.Warning("[%s] retrieving plan failed, flag plan '%v' does not exist. Using free plan",
			p.machine.MachineId, splitted[1])
		return Free, nil
	}

	return plan, nil
}

func (p *PlanChecker) AllowedInstances(wantInstance InstanceType) error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	allowedInstances := plan.Limits().AllowedInstances

	p.log.Info("[%s] checking instance type. want: %s (plan: %s)",
		p.machine.MachineId, wantInstance, plan)

	if _, ok := allowedInstances[wantInstance]; ok {
		return nil
	}

	return fmt.Errorf("not allowed to create instance type: %s", wantInstance)
}

func (p *PlanChecker) AlwaysOn() error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	machineData, ok := p.machine.CurrentData.(*Machine)
	if !ok {
		return fmt.Errorf("current data is malformed: %v", p.machine.CurrentData)
	}

	alwaysOnLimit := plan.Limits().AlwaysOn

	user := machineData.Users[0]

	// get all alwaysOn machines that belongs to this user
	alwaysOnMachines := 0
	err = p.db.Run("jMachines", func(c *mgo.Collection) error {
		alwaysOnMachines, err = c.Find(bson.M{
			"users.id":      user.Id,
			"meta.alwaysOn": true,
		}).Count()

		return err
	})

	// if it's something else just return an error, needs to be fixed
	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	p.log.Info("[%s] checking alwaysOn limit. current alwaysOn count: %d (plan limit: %d, plan: %s)",
		p.machine.MachineId, alwaysOnMachines, alwaysOnLimit, plan)
	// the user has still not reached the limit
	if alwaysOnMachines <= alwaysOnLimit {
		p.log.Info("[%s] allowing user '%s'. current alwaysOn count: %d (plan limit: %d, plan: %s)",
			p.machine.MachineId, p.username, alwaysOnMachines, alwaysOnLimit, plan)
		return nil
	}

	p.log.Info("[%s] denying user '%s'. current alwaysOn count: %d (plan limit: %d, plan: %s)",
		p.machine.MachineId, p.username, alwaysOnMachines, alwaysOnLimit, plan)
	return fmt.Errorf("total alwaysOn limit has been reached")
}

func (p *PlanChecker) Timeout() error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	// get the timeout from the plan in which the user belongs to
	planTimeout := plan.Limits().Timeout

	machineData, ok := p.machine.CurrentData.(*Machine)
	if !ok {
		return fmt.Errorf("current data is malformed: %v", p.machine.CurrentData)
	}

	// connect and get real time data directly from the machines klient
	klient, err := klient.New(p.kite, machineData.QueryString)
	if err != nil {
		return err
	}
	defer klient.Close()

	// get the usage directly from the klient, which is the most predictable source
	usg, err := klient.Usage()
	if err != nil {
		return err
	}

	p.log.Info("[%s] machine [%s] is inactive for %s (plan limit: %s, plan: %s).",
		machineData.Id.Hex(), machineData.IpAddress, usg.InactiveDuration, planTimeout, plan)

	// It still have plenty of time to work, do not stop it
	if usg.InactiveDuration <= planTimeout {
		return nil
	}

	p.log.Info("[%s] machine [%s] has reached current plan limit of %s (plan: %s). Shutting down...",
		machineData.Id.Hex(), machineData.IpAddress, usg.InactiveDuration, planTimeout, plan)

	// mark our state as stopping so others know what we are doing
	p.provider.UpdateState(machineData.Id.Hex(), machinestate.Stopping)

	// replace with the real and authenticated username
	p.machine.Builder["username"] = klient.Username

	// Hasta la vista, baby!
	err = p.provider.Stop(p.machine)
	if err != nil {
		return err
	}

	// update to final state too
	return p.provider.UpdateState(machineData.Id.Hex(), machinestate.Stopped)
}

func (p *PlanChecker) Total() error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	allowedMachines := plan.Limits().Total

	instances, err := p.userInstances()

	// no match, allow to create instance
	if err == aws.ErrNoInstances {
		p.log.Info("[%s] allowing user '%s'. current machine count: %d (plan limit: %d, plan: %s)",
			p.machine.MachineId, p.username, len(instances), allowedMachines, plan)
		return nil
	}

	// if it's something else don't allow it until it's solved
	if err != nil {
		return err
	}

	if len(instances) >= allowedMachines {
		p.log.Info("[%s] denying user '%s'. current machine count: %d (plan limit: %d, plan: %s)",
			p.machine.MachineId, p.username, len(instances), allowedMachines, plan)

		return fmt.Errorf("total limit of %d machines has been reached", allowedMachines)
	}

	p.log.Info("[%s] allowing user '%s'. current machine count: %d (plan limit: %d, plan: %s)",
		p.machine.MachineId, p.username, len(instances), allowedMachines, plan)

	return nil
}

func (p *PlanChecker) Storage(wantStorage int) error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	totalStorage := plan.Limits().Storage

	instances, err := p.userInstances()

	// i hate for loops too, but unfortunaly the responses are always in form
	// of slices
	currentStorage := 0
	for _, instance := range instances {
		for _, blockDevice := range instance.BlockDevices {
			volumes, err := p.api.Client.Volumes([]string{blockDevice.VolumeId}, ec2.NewFilter())
			if err != nil {
				return err
			}

			for _, volume := range volumes.Volumes {
				volumeStorage, err := strconv.Atoi(volume.Size)
				if err != nil {
					return err
				}

				currentStorage += volumeStorage
			}
		}
	}

	p.log.Info("[%s] Checking storage. Current: %dGB. Want: %dGB (plan limit: %dGB, plan: %s)",
		p.machine.MachineId, currentStorage, wantStorage, totalStorage, plan)

	if currentStorage+wantStorage > totalStorage {
		return fmt.Errorf("total storage limit has been reached. Can use %dGB of %dGB (plan: %s)",
			totalStorage-currentStorage, totalStorage, plan)
	}

	p.log.Info("[%s] Allowing user '%s'. Current: %dGB. Want: %dGB (plan limit: %dGB, plan: %s)",
		p.machine.MachineId, p.username, currentStorage, wantStorage, totalStorage, plan)

	// allow to create storage
	return nil
}

func (p *PlanChecker) userInstances() ([]ec2.Instance, error) {
	filter := ec2.NewFilter()
	// instances in Amazon have a `koding-user` tag with the username as the
	// value. We can easily find them acording to this tag
	filter.Add("tag:koding-user", p.username)
	filter.Add("tag:koding-env", p.kite.Config.Environment)

	// Anything except "terminated" and "shutting-down"
	filter.Add("instance-state-name", "pending", "running", "stopping", "stopped")

	return p.api.InstancesByFilter(filter)

}
<commit_msg>koding/checker: do not fail for mongodb lookups<commit_after>package koding

import (
	"fmt"
	"koding/db/mongodb"
	"koding/kites/kloud/klient"
	"strconv"
	"strings"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	"github.com/koding/kite"
	aws "github.com/koding/kloud/api/amazon"
	"github.com/koding/kloud/machinestate"
	"github.com/koding/kloud/protocol"
	"github.com/koding/kloud/provider/amazon"
	"github.com/koding/logging"
	"github.com/mitchellh/goamz/ec2"
)

// Checker checks various aspects of a machine. It used to limit certain
// aspects of a machine, such a the total machine or total storage.
type Checker interface {
	// Total checks whether the user has reached the current plan's limit of
	// having a total number numbers of machines. It returns an error if the
	// limit is reached or an unexplained error happaned.
	Total() error

	// AlwaysOn checks whether the given machine has reached the current plans
	// always on limit
	AlwaysOn() error

	// Timeout checks whether the user has reached the current plan's
	// inactivity timeout.
	Timeout() error

	// Storage checks whether the user has reached the current plan's limit
	// total storage with the supplied wantStorage information. It returns an
	// error if the limit is reached or an unexplained error happaned.
	Storage(wantStorage int) error

	// AllowedInstances checks whether the given machine has the permisison to
	// create the given instance type
	AllowedInstances(wantInstance InstanceType) error
}

type PlanChecker struct {
	api      *amazon.AmazonClient
	db       *mongodb.MongoDB
	machine  *protocol.Machine
	provider *Provider
	kite     *kite.Kite
	username string
	log      logging.Logger
}

// PlanChecker creates and returns a new Checker interface that is responsible
// of checking various pieces of informations based on a Plan
func (p *Provider) PlanChecker(opts *protocol.Machine) (Checker, error) {
	a, err := p.NewClient(opts)
	if err != nil {
		return nil, err
	}

	checker := &PlanChecker{
		api:      a,
		provider: p,
		db:       p.Session,
		kite:     p.Kite,
		username: opts.Builder["username"].(string),
		log:      p.Log,
		machine:  opts,
	}

	return checker, nil
}

// Plan returns user's current plan
func (p *PlanChecker) Plan() (Plan, error) {
	var account = struct {
		GlobalFlags []string `bson:"globalFlags"`
	}{}
	if err := p.db.Run("jAccounts", func(c *mgo.Collection) error {
		return c.Find(bson.M{"profile.nickname": p.username}).One(&account)
	}); err != nil {
		p.log.Warning("[%s] retrieving plan failed, mongodb lookup err: %s. Using free plan",
			p.machine.MachineId, err.Error())
		return Free, nil
	}

	if len(account.GlobalFlags) == 0 {
		p.log.Warning("[%s] retrieving plan failed, no flag defined. Using free plan",
			p.machine.MachineId)
		return Free, nil
	}

	splitted := strings.Split(account.GlobalFlags[0], "-")
	if len(splitted) != 2 {
		p.log.Warning("[%s] retrieving plan failed, flag '%v' malformed. Using free plan",
			p.machine.MachineId, account.GlobalFlags)
		return Free, nil
	}

	plan, ok := plans[strings.Title(splitted[1])]
	if !ok {
		p.log.Warning("[%s] retrieving plan failed, flag plan '%v' does not exist. Using free plan",
			p.machine.MachineId, splitted[1])
		return Free, nil
	}

	return plan, nil
}

func (p *PlanChecker) AllowedInstances(wantInstance InstanceType) error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	allowedInstances := plan.Limits().AllowedInstances

	p.log.Info("[%s] checking instance type. want: %s (plan: %s)",
		p.machine.MachineId, wantInstance, plan)

	if _, ok := allowedInstances[wantInstance]; ok {
		return nil
	}

	return fmt.Errorf("not allowed to create instance type: %s", wantInstance)
}

func (p *PlanChecker) AlwaysOn() error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	machineData, ok := p.machine.CurrentData.(*Machine)
	if !ok {
		return fmt.Errorf("current data is malformed: %v", p.machine.CurrentData)
	}

	alwaysOnLimit := plan.Limits().AlwaysOn

	user := machineData.Users[0]

	// get all alwaysOn machines that belongs to this user
	alwaysOnMachines := 0
	err = p.db.Run("jMachines", func(c *mgo.Collection) error {
		alwaysOnMachines, err = c.Find(bson.M{
			"users.id":      user.Id,
			"meta.alwaysOn": true,
		}).Count()

		return err
	})

	// if it's something else just return an error, needs to be fixed
	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	p.log.Info("[%s] checking alwaysOn limit. current alwaysOn count: %d (plan limit: %d, plan: %s)",
		p.machine.MachineId, alwaysOnMachines, alwaysOnLimit, plan)
	// the user has still not reached the limit
	if alwaysOnMachines <= alwaysOnLimit {
		p.log.Info("[%s] allowing user '%s'. current alwaysOn count: %d (plan limit: %d, plan: %s)",
			p.machine.MachineId, p.username, alwaysOnMachines, alwaysOnLimit, plan)
		return nil
	}

	p.log.Info("[%s] denying user '%s'. current alwaysOn count: %d (plan limit: %d, plan: %s)",
		p.machine.MachineId, p.username, alwaysOnMachines, alwaysOnLimit, plan)
	return fmt.Errorf("total alwaysOn limit has been reached")
}

func (p *PlanChecker) Timeout() error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	// get the timeout from the plan in which the user belongs to
	planTimeout := plan.Limits().Timeout

	machineData, ok := p.machine.CurrentData.(*Machine)
	if !ok {
		return fmt.Errorf("current data is malformed: %v", p.machine.CurrentData)
	}

	// connect and get real time data directly from the machines klient
	klient, err := klient.New(p.kite, machineData.QueryString)
	if err != nil {
		return err
	}
	defer klient.Close()

	// get the usage directly from the klient, which is the most predictable source
	usg, err := klient.Usage()
	if err != nil {
		return err
	}

	p.log.Info("[%s] machine [%s] is inactive for %s (plan limit: %s, plan: %s).",
		machineData.Id.Hex(), machineData.IpAddress, usg.InactiveDuration, planTimeout, plan)

	// It still have plenty of time to work, do not stop it
	if usg.InactiveDuration <= planTimeout {
		return nil
	}

	p.log.Info("[%s] machine [%s] has reached current plan limit of %s (plan: %s). Shutting down...",
		machineData.Id.Hex(), machineData.IpAddress, usg.InactiveDuration, planTimeout, plan)

	// mark our state as stopping so others know what we are doing
	p.provider.UpdateState(machineData.Id.Hex(), machinestate.Stopping)

	// replace with the real and authenticated username
	p.machine.Builder["username"] = klient.Username

	// Hasta la vista, baby!
	err = p.provider.Stop(p.machine)
	if err != nil {
		return err
	}

	// update to final state too
	return p.provider.UpdateState(machineData.Id.Hex(), machinestate.Stopped)
}

func (p *PlanChecker) Total() error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	allowedMachines := plan.Limits().Total

	instances, err := p.userInstances()

	// no match, allow to create instance
	if err == aws.ErrNoInstances {
		p.log.Info("[%s] allowing user '%s'. current machine count: %d (plan limit: %d, plan: %s)",
			p.machine.MachineId, p.username, len(instances), allowedMachines, plan)
		return nil
	}

	// if it's something else don't allow it until it's solved
	if err != nil {
		return err
	}

	if len(instances) >= allowedMachines {
		p.log.Info("[%s] denying user '%s'. current machine count: %d (plan limit: %d, plan: %s)",
			p.machine.MachineId, p.username, len(instances), allowedMachines, plan)

		return fmt.Errorf("total limit of %d machines has been reached", allowedMachines)
	}

	p.log.Info("[%s] allowing user '%s'. current machine count: %d (plan limit: %d, plan: %s)",
		p.machine.MachineId, p.username, len(instances), allowedMachines, plan)

	return nil
}

func (p *PlanChecker) Storage(wantStorage int) error {
	plan, err := p.Plan()
	if err != nil {
		return err
	}

	totalStorage := plan.Limits().Storage

	instances, err := p.userInstances()

	// i hate for loops too, but unfortunaly the responses are always in form
	// of slices
	currentStorage := 0
	for _, instance := range instances {
		for _, blockDevice := range instance.BlockDevices {
			volumes, err := p.api.Client.Volumes([]string{blockDevice.VolumeId}, ec2.NewFilter())
			if err != nil {
				return err
			}

			for _, volume := range volumes.Volumes {
				volumeStorage, err := strconv.Atoi(volume.Size)
				if err != nil {
					return err
				}

				currentStorage += volumeStorage
			}
		}
	}

	p.log.Info("[%s] Checking storage. Current: %dGB. Want: %dGB (plan limit: %dGB, plan: %s)",
		p.machine.MachineId, currentStorage, wantStorage, totalStorage, plan)

	if currentStorage+wantStorage > totalStorage {
		return fmt.Errorf("total storage limit has been reached. Can use %dGB of %dGB (plan: %s)",
			totalStorage-currentStorage, totalStorage, plan)
	}

	p.log.Info("[%s] Allowing user '%s'. Current: %dGB. Want: %dGB (plan limit: %dGB, plan: %s)",
		p.machine.MachineId, p.username, currentStorage, wantStorage, totalStorage, plan)

	// allow to create storage
	return nil
}

func (p *PlanChecker) userInstances() ([]ec2.Instance, error) {
	filter := ec2.NewFilter()
	// instances in Amazon have a `koding-user` tag with the username as the
	// value. We can easily find them acording to this tag
	filter.Add("tag:koding-user", p.username)
	filter.Add("tag:koding-env", p.kite.Config.Environment)

	// Anything except "terminated" and "shutting-down"
	filter.Add("instance-state-name", "pending", "running", "stopping", "stopped")

	return p.api.InstancesByFilter(filter)

}
<|endoftext|>


<commit_before>package index

import (
	"encoding/json"
	"hash/crc32"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"

	"github.com/djherbis/times"
)

// Version stores current version of index.
const Version = 1

// Entry represents a single file registered to index.
type Entry struct {
	Name  string      `json:"name"`  // The relative name of the file.
	CTime int64       `json:"ctime"` // Metadata change time since EPOCH.
	MTime int64       `json:"mtime"` // File data change time since EPOCH.
	Mode  os.FileMode `json:"mode"`  // File mode and permission bits.
	Size  int64       `json:"size"`  // Size of the file.
	SHA1  []byte      `json:"sha1"`  // SHA-1 hash of file content.
}

// NewEntryFile creates new Entry from a file stored under path argument.
// Info is optional and, if given, should store LStat result on the given file.
func NewEntryFile(root, path string, info os.FileInfo) (e *Entry, err error) {
	if info == nil {
		if info, err = os.Lstat(path); err != nil {
			return nil, err
		}
	}

	// Get relative file name.
	name, err := filepath.Rel(root, path)
	if err != nil {
		return nil, err
	}

	// Compute file's SHA-1 sum.
	var sum []byte
	if !info.IsDir() {
		if sum, err = readCRC32(path); err != nil {
			return nil, err
		}
	}

	return &Entry{
		Name:  filepath.ToSlash(name),
		CTime: ctime(info),
		MTime: info.ModTime().UnixNano(),
		Mode:  info.Mode(),
		Size:  info.Size(),
		SHA1:  sum,
	}, nil
}

var copyBufPool = sync.Pool{
	New: func() interface{} {
		b := make([]byte, 64*1024)
		return &b
	},
}

// readCRC32 computes CRC-32 checksum of a given file content.
func readCRC32(path string) ([]byte, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	hash := crc32.NewIEEE()

	bufp := copyBufPool.Get().(*[]byte)
	defer copyBufPool.Put(bufp)

	if _, err := io.CopyBuffer(hash, file, *bufp); err != nil {
		return nil, err
	}

	return hash.Sum(nil), nil
}

// ChangeMeta indicates what change has been done on a given file.
type ChangeMeta uint32

const (
	ChangeMetaUpdate ChangeMeta = 1 << iota // File was updated.
	ChangeMetaRemove                        // File was removed.
	ChangeMetaAdd                           // File was added.

	ChangeMetaLarge ChangeMeta = 1 << (8 + iota) // File size is above 4GB.
)

// Change describes single file change.
type Change struct {
	Name      string     `json:"name"`      // The relative name of the file.
	Size      uint32     `json:"size"`      // Size of the file truncated to 32 bits.
	Meta      ChangeMeta `json:"meta"`      // The type of operation made on file entry.
	CreatedAt int64      `json:"createdAt"` // Change creation time since EPOCH.
}

// ChangeSlice stores multiple changes.
type ChangeSlice []Change

func (cs ChangeSlice) Len() int           { return len(cs) }
func (cs ChangeSlice) Swap(i, j int)      { cs[i], cs[j] = cs[j], cs[i] }
func (cs ChangeSlice) Less(i, j int) bool { return cs[i].Name < cs[j].Name }

// Index stores a virtual working tree state. It recursively records objects in
// a given root path and allows to efficiently detect changes on it.
type Index struct {
	limitC chan struct{}

	mu      sync.RWMutex
	entries map[string]*Entry
}

var (
	_ json.Marshaler   = (*Index)(nil)
	_ json.Unmarshaler = (*Index)(nil)
)

// NewIndex creates the empty index object.
func NewIndex() *Index {
	return &Index{
		limitC:  make(chan struct{}, 2*runtime.NumCPU()),
		entries: make(map[string]*Entry, 0),
	}
}

// NewIndexFiles walks the given file tree roted at root and records file
// states to resulting Index object.
func NewIndexFiles(root string) (*Index, error) {
	idx := NewIndex()

	// In order to get as much entries as we can we ignore errors.
	var wg sync.WaitGroup
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		// Skip root path.
		if name, err := filepath.Rel(root, path); err != nil || name == "." {
			return nil
		}

		wg.Add(1)
		idx.addEntry(&wg, root, path, info)
		return nil
	}

	if err := filepath.Walk(root, walkFn); err != nil {
		return nil, err
	}

	wg.Wait()
	return idx, nil
}

// addEntry asynchronously adds new entry to index. This function has limits
// to a number of started go-routines. Errors are ignored.
func (idx *Index) addEntry(wg *sync.WaitGroup, root, path string, info os.FileInfo) {
	idx.limitC <- struct{}{}

	go func() {
		defer func() {
			<-idx.limitC
			wg.Done()
		}()

		entry, err := NewEntryFile(root, path, info)
		if err != nil {
			return
		}

		idx.mu.Lock()
		idx.entries[entry.Name] = entry
		idx.mu.Unlock()
	}()
}

// Count returns the number of entries stored in index. Only items which size is
// below provided value are counted. If provided argument is negative, this
// function will return the number of all entries.
func (idx *Index) Count(maxsize int64) (count int) {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	if maxsize < 0 {
		return len(idx.entries)
	} else if maxsize == 0 {
		return 0
	}

	for _, entry := range idx.entries {
		if entry != nil && entry.Size <= maxsize {
			count++
		}
	}

	return count
}

// DiskSize tells how much disk space would be used by entries stored in index.
// Only items which size is below provided value are counted. If provided
// argument is negative, this function will count disk size of all items.
func (idx *Index) DiskSize(maxsize int64) (size int64) {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	if maxsize == 0 {
		return 0
	}

	for _, entry := range idx.entries {
		if entry != nil && (maxsize < 0 || entry.Size <= maxsize) {
			size += entry.Size
		}
	}

	return size
}

// Compare rereads the given file tree roted at root and compares its entries
// to previous state of the index. All detected changes will be stored in
// returned Change slice.
func (idx *Index) Compare(root string) (cs ChangeSlice) {
	visited := make(map[string]struct{})

	// Walk over current root path and check it files.
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		name, err := filepath.Rel(root, path)
		if err != nil || name == "." {
			return nil
		}
		name = filepath.ToSlash(name)

		idx.mu.RLock()
		entry, ok := idx.entries[name]
		idx.mu.RUnlock()

		// Not found in current index - file was added.
		if !ok {
			cs = append(cs, Change{
				Name:      name,
				Size:      safeTruncate(info.Size()),
				Meta:      ChangeMetaAdd | markLargeMeta(info.Size()),
				CreatedAt: time.Now().UnixNano(),
			})
			return nil
		}

		// Entry is read only-now. Check for changes.
		visited[name] = struct{}{}
		if entry.MTime != info.ModTime().UnixNano() ||
			entry.CTime != ctime(info) ||
			entry.Size != info.Size() {
			cs = append(cs, Change{
				Name:      name,
				Size:      safeTruncate(info.Size()),
				Meta:      ChangeMetaUpdate | markLargeMeta(info.Size()),
				CreatedAt: time.Now().UnixNano(),
			})
		}

		return nil
	}

	if err := filepath.Walk(root, walkFn); err != nil {
		return nil
	}

	// Check for removes.
	idx.mu.RLock()
	for name := range idx.entries {
		if _, ok := visited[name]; !ok {
			path := filepath.Join(root, filepath.FromSlash(name))
			if _, err := os.Lstat(path); os.IsNotExist(err) {
				cs = append(cs, Change{
					Name:      name,
					Meta:      ChangeMetaRemove | markLargeMeta(idx.entries[name].Size),
					CreatedAt: time.Now().UnixNano(),
				})
			}
		}
	}
	idx.mu.RUnlock()

	return cs
}

// safeTruncate converts signed integer to unsigned one returning 0 for negative
// values of provided argument.
func safeTruncate(n int64) uint32 {
	if n < 0 {
		return 0
	}

	return uint32(n)
}

// markLargeMeta adds large file flag for files which size is over 4GiB.
func markLargeMeta(n int64) ChangeMeta {
	if n < 0 || (n>>32) == 0 {
		return 0
	}

	return ChangeMetaLarge
}

// ctime gets file's change time in UNIX Nano format.
func ctime(fi os.FileInfo) int64 {
	if tspec := times.Get(fi); tspec.HasChangeTime() {
		return tspec.ChangeTime().UnixNano()
	}

	return 0
}

// Apply modifies index according to provided changes. This function doesn't
// guarantee that changes from Compare function applied to the index will
// result in actual directory state.
func (idx *Index) Apply(root string, cs ChangeSlice) {
	var wg sync.WaitGroup

	for i := range cs {
		switch {
		case cs[i].Meta&(ChangeMetaUpdate|ChangeMetaAdd) != 0:
			// Check if the event is still valid or if it was replaced by newer
			// change.
			idx.mu.RLock()
			entry, ok := idx.entries[cs[i].Name]
			idx.mu.RUnlock()

			// Entry was updated/added after the event was created.
			if ok && entry.MTime > cs[i].CreatedAt {
				continue
			}
			fallthrough
		case cs[i].Meta&ChangeMetaRemove != 0:
			// Check if the file still exists, since it could be removed before
			// Apply was called. If the file exists, create new entry from it
			// and replace its value inside index map.
			path := filepath.Join(root, filepath.FromSlash(cs[i].Name))
			info, err := os.Lstat(path)
			if os.IsNotExist(err) {
				idx.mu.Lock()
				delete(idx.entries, cs[i].Name)
				idx.mu.Unlock()
				continue
			}

			wg.Add(1)
			idx.addEntry(&wg, root, path, info)
		}
	}

	wg.Wait()
}

// MarshalJSON satisfies json.Marshaler interface. It safely marshals index
// private data to JSON format.
func (idx *Index) MarshalJSON() ([]byte, error) {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	return json.Marshal(idx.entries)
}

// UnmarshalJSON satisfies json.Unmarshaler interface. It is used to unmarshal
// data into private index fields.
func (idx *Index) UnmarshalJSON(data []byte) error {
	idx.mu.Lock()
	defer idx.mu.Unlock()

	return json.Unmarshal(data, &idx.entries)
}
<commit_msg>klient/machine: use workers to add index entries<commit_after>package index

import (
	"encoding/json"
	"hash/crc32"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"

	"github.com/djherbis/times"
)

// Version stores current version of index.
const Version = 1

// Entry represents a single file registered to index.
type Entry struct {
	Name  string      `json:"name"`  // The relative name of the file.
	CTime int64       `json:"ctime"` // Metadata change time since EPOCH.
	MTime int64       `json:"mtime"` // File data change time since EPOCH.
	Mode  os.FileMode `json:"mode"`  // File mode and permission bits.
	Size  int64       `json:"size"`  // Size of the file.
	SHA1  []byte      `json:"sha1"`  // SHA-1 hash of file content.
}

// NewEntryFile creates new Entry from a file stored under path argument.
// Info is optional and, if given, should store LStat result on the given file.
func NewEntryFile(root, path string, info os.FileInfo) (e *Entry, err error) {
	if info == nil {
		if info, err = os.Lstat(path); err != nil {
			return nil, err
		}
	}

	// Get relative file name.
	name, err := filepath.Rel(root, path)
	if err != nil {
		return nil, err
	}

	// Compute file's SHA-1 sum.
	var sum []byte
	if !info.IsDir() {
		if sum, err = readCRC32(path); err != nil {
			return nil, err
		}
	}

	return &Entry{
		Name:  filepath.ToSlash(name),
		CTime: ctime(info),
		MTime: info.ModTime().UnixNano(),
		Mode:  info.Mode(),
		Size:  info.Size(),
		SHA1:  sum,
	}, nil
}

var copyBufPool = sync.Pool{
	New: func() interface{} {
		b := make([]byte, 64*1024)
		return &b
	},
}

// readCRC32 computes CRC-32 checksum of a given file content.
func readCRC32(path string) ([]byte, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	hash := crc32.NewIEEE()

	bufp := copyBufPool.Get().(*[]byte)
	defer copyBufPool.Put(bufp)

	if _, err := io.CopyBuffer(hash, file, *bufp); err != nil {
		return nil, err
	}

	return hash.Sum(nil), nil
}

// ChangeMeta indicates what change has been done on a given file.
type ChangeMeta uint32

const (
	ChangeMetaUpdate ChangeMeta = 1 << iota // File was updated.
	ChangeMetaRemove                        // File was removed.
	ChangeMetaAdd                           // File was added.

	ChangeMetaLarge ChangeMeta = 1 << (8 + iota) // File size is above 4GB.
)

// Change describes single file change.
type Change struct {
	Name      string     `json:"name"`      // The relative name of the file.
	Size      uint32     `json:"size"`      // Size of the file truncated to 32 bits.
	Meta      ChangeMeta `json:"meta"`      // The type of operation made on file entry.
	CreatedAt int64      `json:"createdAt"` // Change creation time since EPOCH.
}

// ChangeSlice stores multiple changes.
type ChangeSlice []Change

func (cs ChangeSlice) Len() int           { return len(cs) }
func (cs ChangeSlice) Swap(i, j int)      { cs[i], cs[j] = cs[j], cs[i] }
func (cs ChangeSlice) Less(i, j int) bool { return cs[i].Name < cs[j].Name }

// Index stores a virtual working tree state. It recursively records objects in
// a given root path and allows to efficiently detect changes on it.
type Index struct {
	mu      sync.RWMutex
	entries map[string]*Entry
}

var (
	_ json.Marshaler   = (*Index)(nil)
	_ json.Unmarshaler = (*Index)(nil)
)

// NewIndex creates the empty index object.
func NewIndex() *Index {
	return &Index{
		entries: make(map[string]*Entry, 0),
	}
}

type fileDesc struct {
	path string      // relative path to the file.
	info os.FileInfo // file LStat result.
}

// NewIndexFiles walks the given file tree roted at root and records file
// states to resulting Index object.
func NewIndexFiles(root string) (*Index, error) {
	idx := NewIndex()

	// Start worker pool.
	var wg sync.WaitGroup
	fC := make(chan *fileDesc)
	for i := 0; i < 2*runtime.NumCPU(); i++ {
		wg.Add(1)
		go idx.addEntryWorker(root, &wg, fC)
	}
	defer func() {
		close(fC)
		wg.Wait()
	}()

	// In order to get as much entries as we can we ignore errors.
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		// Skip root path.
		if name, err := filepath.Rel(root, path); err != nil || name == "." {
			return nil
		}

		fC <- &fileDesc{path: path, info: info}
		return nil
	}

	if err := filepath.Walk(root, walkFn); err != nil {
		return nil, err
	}

	return idx, nil
}

// addEntryWorker asynchronously adds new entry to index. Errors are ignored.
func (idx *Index) addEntryWorker(root string, wg *sync.WaitGroup, fC <-chan *fileDesc) {
	defer wg.Done()

	for f := range fC {
		entry, err := NewEntryFile(root, f.path, f.info)
		if err != nil {
			continue
		}

		idx.mu.Lock()
		idx.entries[entry.Name] = entry
		idx.mu.Unlock()
	}
}

// Count returns the number of entries stored in index. Only items which size is
// below provided value are counted. If provided argument is negative, this
// function will return the number of all entries.
func (idx *Index) Count(maxsize int64) (count int) {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	if maxsize < 0 {
		return len(idx.entries)
	} else if maxsize == 0 {
		return 0
	}

	for _, entry := range idx.entries {
		if entry != nil && entry.Size <= maxsize {
			count++
		}
	}

	return count
}

// DiskSize tells how much disk space would be used by entries stored in index.
// Only items which size is below provided value are counted. If provided
// argument is negative, this function will count disk size of all items.
func (idx *Index) DiskSize(maxsize int64) (size int64) {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	if maxsize == 0 {
		return 0
	}

	for _, entry := range idx.entries {
		if entry != nil && (maxsize < 0 || entry.Size <= maxsize) {
			size += entry.Size
		}
	}

	return size
}

// Compare rereads the given file tree roted at root and compares its entries
// to previous state of the index. All detected changes will be stored in
// returned Change slice.
func (idx *Index) Compare(root string) (cs ChangeSlice) {
	visited := make(map[string]struct{})

	// Walk over current root path and check it files.
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		name, err := filepath.Rel(root, path)
		if err != nil || name == "." {
			return nil
		}
		name = filepath.ToSlash(name)

		idx.mu.RLock()
		entry, ok := idx.entries[name]
		idx.mu.RUnlock()

		// Not found in current index - file was added.
		if !ok {
			cs = append(cs, Change{
				Name:      name,
				Size:      safeTruncate(info.Size()),
				Meta:      ChangeMetaAdd | markLargeMeta(info.Size()),
				CreatedAt: time.Now().UnixNano(),
			})
			return nil
		}

		// Entry is read only-now. Check for changes.
		visited[name] = struct{}{}
		if entry.MTime != info.ModTime().UnixNano() ||
			entry.CTime != ctime(info) ||
			entry.Size != info.Size() {
			cs = append(cs, Change{
				Name:      name,
				Size:      safeTruncate(info.Size()),
				Meta:      ChangeMetaUpdate | markLargeMeta(info.Size()),
				CreatedAt: time.Now().UnixNano(),
			})
		}

		return nil
	}

	if err := filepath.Walk(root, walkFn); err != nil {
		return nil
	}

	// Check for removes.
	idx.mu.RLock()
	for name := range idx.entries {
		if _, ok := visited[name]; !ok {
			path := filepath.Join(root, filepath.FromSlash(name))
			if _, err := os.Lstat(path); os.IsNotExist(err) {
				cs = append(cs, Change{
					Name:      name,
					Meta:      ChangeMetaRemove | markLargeMeta(idx.entries[name].Size),
					CreatedAt: time.Now().UnixNano(),
				})
			}
		}
	}
	idx.mu.RUnlock()

	return cs
}

// safeTruncate converts signed integer to unsigned one returning 0 for negative
// values of provided argument.
func safeTruncate(n int64) uint32 {
	if n < 0 {
		return 0
	}

	return uint32(n)
}

// markLargeMeta adds large file flag for files which size is over 4GiB.
func markLargeMeta(n int64) ChangeMeta {
	if n < 0 || (n>>32) == 0 {
		return 0
	}

	return ChangeMetaLarge
}

// ctime gets file's change time in UNIX Nano format.
func ctime(fi os.FileInfo) int64 {
	if tspec := times.Get(fi); tspec.HasChangeTime() {
		return tspec.ChangeTime().UnixNano()
	}

	return 0
}

// Apply modifies index according to provided changes. This function doesn't
// guarantee that changes from Compare function applied to the index will
// result in actual directory state.
func (idx *Index) Apply(root string, cs ChangeSlice) {
	// Start worker pool.
	var wg sync.WaitGroup
	fC := make(chan *fileDesc)
	for i := 0; i < naturalMin(2*runtime.NumCPU(), len(cs)); i++ {
		wg.Add(1)
		go idx.addEntryWorker(root, &wg, fC)
	}

	for i := range cs {
		switch {
		case cs[i].Meta&(ChangeMetaUpdate|ChangeMetaAdd) != 0:
			// Check if the event is still valid or if it was replaced by newer
			// change.
			idx.mu.RLock()
			entry, ok := idx.entries[cs[i].Name]
			idx.mu.RUnlock()

			// Entry was updated/added after the event was created.
			if ok && entry.MTime > cs[i].CreatedAt {
				continue
			}
			fallthrough
		case cs[i].Meta&ChangeMetaRemove != 0:
			// Check if the file still exists, since it could be removed before
			// Apply was called. If the file exists, create new entry from it
			// and replace its value inside index map.
			path := filepath.Join(root, filepath.FromSlash(cs[i].Name))
			info, err := os.Lstat(path)
			if os.IsNotExist(err) {
				idx.mu.Lock()
				delete(idx.entries, cs[i].Name)
				idx.mu.Unlock()
				continue
			}

			fC <- &fileDesc{path: path, info: info}
		}
	}

	close(fC)
	wg.Wait()
}

// naturalMin returns the minimal value of provided arguments but not less than
// one.
func naturalMin(a, b int) (n int) {
	if n = b; a < b {
		n = a
	}

	if n < 1 {
		return 1
	}

	return n
}

// MarshalJSON satisfies json.Marshaler interface. It safely marshals index
// private data to JSON format.
func (idx *Index) MarshalJSON() ([]byte, error) {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	return json.Marshal(idx.entries)
}

// UnmarshalJSON satisfies json.Unmarshaler interface. It is used to unmarshal
// data into private index fields.
func (idx *Index) UnmarshalJSON(data []byte) error {
	idx.mu.Lock()
	defer idx.mu.Unlock()

	return json.Unmarshal(data, &idx.entries)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"koding/db/models"
	"koding/db/mongodb"
	"koding/tools/config"
	"math"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/gorilla/mux"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type ApiWorker struct {
	Name               string    `json:"name"`
	Version            int       `json:"version"`
	Environment        string    `json:"environment"`
	Hostname           string    `json:"hostname"`
	ServiceGenericName string    `json:"serviceGenericName"`
	ServiceUniqueName  string    `json:"serviceUniqueName"`
	Uuid               string    `json:"uuid"`
	Timestamp          time.Time `json:"timestamp"`
	Pid                int       `json:"pid"`
	State              string    `json:"state"`
	Uptime             int       `json:"uptime"`
	Port               int       `json:"port"`
}

type Workers []ApiWorker

var StatusCode = map[models.WorkerStatus]string{
	models.Started: "started",
	models.Waiting: "waiting",
	models.Killed:  "dead",
	models.Dead:    "dead",
}

const (
	WorkersCollection = "jKontrolWorkers"
	WorkersDB         = "kontrol"
)

var (
	kontrolDB *mongodb.MongoDB

	// used for loadbalance modes, like roundrobin or random
	index AtomicUint32
)

func GetWorkerURL(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	workerName := vars["workername"]

	queries, _ := url.ParseQuery(req.URL.RawQuery)
	worker := models.Worker{}
	workers := make(Workers, 0)

	query := func(c *mgo.Collection) error {
		iter := c.Find(bson.M{"name": workerName}).Iter()
		for iter.Next(&worker) {
			apiWorker := &ApiWorker{
				Name:               worker.Name,
				Version:            worker.Version,
				Environment:        worker.Environment,
				Hostname:           worker.Hostname,
				ServiceGenericName: worker.ServiceGenericName,
				ServiceUniqueName:  worker.ServiceUniqueName,
				Uuid:               worker.Uuid,
				Timestamp:          worker.Timestamp,
				Pid:                worker.Pid,
				State:              StatusCode[worker.Status],
				Uptime:             worker.Monitor.Uptime,
				Port:               worker.Port,
			}

			workers = append(workers, *apiWorker)
		}

		return nil
	}

	kontrolDB.RunOnDatabase(WorkersDB, WorkersCollection, query)

	// use http for all workers because they don't have ssl certs
	protocolScheme := "http:"

	brokerConf := config.Broker{}
	switch workerName {
	case "brokerKite":
		brokerConf = conf.BrokerKite
	case "broker":
		brokerConf = conf.Broker
	default:
	}

	// broker has ssl cert and a custom url scheme, look what it's it
	if workerName == "broker" || workerName == "brokerKite" {
		if brokerConf.WebProtocol != "" {
			protocolScheme = brokerConf.WebProtocol
		} else {
			protocolScheme = "https:" // fallback
		}
	}

	hostnames := make([]string, len(workers))
	for i, worker := range workers {
		hostnames[i] = fmt.Sprintf("%s//%s:%d", protocolScheme, worker.Hostname, worker.Port)
	}

	var data []byte
	var err error

	_, ok := queries["all"]
	if ok {
		// return all hostnames back
		data, err = json.MarshalIndent(hostnames, "", "  ")
		if err != nil {
			io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
			return
		}
	} else {
		// quit early when there is no workers available
		if len(workers) == 0 {
			writer.Write([]byte(string("")))
			return
		}

		// return only one hostname back, roundrobin
		oldIndex := index.Get() // gives 0 for first time

		N := float64(len(hostnames))
		newIndex := int(math.Mod(float64(oldIndex+1), N))
		hostname := hostnames[newIndex]

		index.CompareAndSwap(oldIndex, uint32(newIndex))

		data = []byte(fmt.Sprintf("\"%s\"", hostname))
	}

	writer.Write(data)
}

type AtomicUint32 uint32

func (i *AtomicUint32) Add(n uint32) uint32 {
	return atomic.AddUint32((*uint32)(i), n)
}

func (i *AtomicUint32) Set(n uint32) {
	atomic.StoreUint32((*uint32)(i), n)
}

func (i *AtomicUint32) Get() uint32 {
	return atomic.LoadUint32((*uint32)(i))
}

func (i *AtomicUint32) CompareAndSwap(oldval, newval uint32) (swapped bool) {
	return atomic.CompareAndSwapUint32((*uint32)(i), oldval, newval)
}

func GetWorkers(writer http.ResponseWriter, req *http.Request) {
	queries, _ := url.ParseQuery(req.URL.RawQuery)

	var latestVersion bool
	var sortFields []string // not initialized means do not sort
	query := bson.M{}
	for key, value := range queries {
		switch key {
		case "version", "pid":
			if value[0] == "latest" {
				latestVersion = true
				break
			}
			v, _ := strconv.Atoi(value[0])
			query[key] = v
		case "state":
			for status, state := range StatusCode {
				if value[0] == state {
					query["status"] = status
				}
			}
		case "sort":
			sortFields = []string{value[0]}
			// override "state" with status, they are not the same in db
			if value[0] == "state" {
				sortFields = []string{"status"}
			}
		default:
			if key == "name" {
				name := value[0]
				if counts := strings.Count(value[0], "-"); counts > 0 {
					s := strings.Split(value[0], "-")
					name = s[0]
				}
				// if searched for social-1, social-2, then return all workers
				// that begins with social
				query[key] = bson.RegEx{Pattern: name, Options: "i"}
			} else {
				query[key] = value[0]
			}
		}
	}

	matchedWorkers := queryResult(query, latestVersion, sortFields)
	data, err := json.MarshalIndent(matchedWorkers, "", "  ")
	if err != nil {
		io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
		return
	}
	writer.Write(data)

}

func queryResult(query bson.M, latestVersion bool, sortFields []string) Workers {
	workers := make(Workers, 0)
	worker := models.Worker{}

	queryFunc := func(c *mgo.Collection) error {
		// sorting is no-op when sortFields is empty
		iter := c.Find(query).Sort(sortFields...).Iter()
		for iter.Next(&worker) {
			apiWorker := &ApiWorker{
				Name:               worker.Name,
				Version:            worker.Version,
				Environment:        worker.Environment,
				Hostname:           worker.Hostname,
				ServiceGenericName: worker.ServiceGenericName,
				ServiceUniqueName:  worker.ServiceUniqueName,
				Uuid:               worker.Uuid,
				Timestamp:          worker.Timestamp,
				Pid:                worker.Pid,
				State:              StatusCode[worker.Status],
				Uptime:             worker.Monitor.Uptime,
				Port:               worker.Port,
			}

			workers = append(workers, *apiWorker)
		}
		return nil
	}

	kontrolDB.RunOnDatabase(WorkersDB, WorkersCollection, queryFunc)

	// finding the largest number of a field in mongo is kinda problematic.
	// therefore we are doing it on our side
	if latestVersion {
		versions := make([]int, len(workers))

		if len(workers) == 0 {
			return workers
		}

		for i, val := range workers {
			versions[i] = val.Version
		}

		sort.Ints(versions)
		maxVersion := versions[len(versions)-1] // get largest version number

		filteredWorkers := make(Workers, 0)
		for _, val := range workers {
			if maxVersion == val.Version {
				filteredWorkers = append(filteredWorkers, val)
			}
		}

		return filteredWorkers
	}

	return workers
}

func GetWorker(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	uuid := vars["uuid"]

	query := bson.M{"uuid": uuid}
	matchedWorkers := queryResult(query, false, nil)
	data, err := json.MarshalIndent(matchedWorkers, "", "  ")
	if err != nil {
		io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
		return
	}
	writer.Write(data)
}
<commit_msg>kontrolapi: add premiumBroker and premiumBrokerKite<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"koding/db/models"
	"koding/db/mongodb"
	"koding/tools/config"
	"math"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/gorilla/mux"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type ApiWorker struct {
	Name               string    `json:"name"`
	Version            int       `json:"version"`
	Environment        string    `json:"environment"`
	Hostname           string    `json:"hostname"`
	ServiceGenericName string    `json:"serviceGenericName"`
	ServiceUniqueName  string    `json:"serviceUniqueName"`
	Uuid               string    `json:"uuid"`
	Timestamp          time.Time `json:"timestamp"`
	Pid                int       `json:"pid"`
	State              string    `json:"state"`
	Uptime             int       `json:"uptime"`
	Port               int       `json:"port"`
}

type Workers []ApiWorker

var StatusCode = map[models.WorkerStatus]string{
	models.Started: "started",
	models.Waiting: "waiting",
	models.Killed:  "dead",
	models.Dead:    "dead",
}

const (
	WorkersCollection = "jKontrolWorkers"
	WorkersDB         = "kontrol"
)

var (
	kontrolDB *mongodb.MongoDB

	// used for loadbalance modes, like roundrobin or random
	index AtomicUint32
)

func GetWorkerURL(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	workerName := vars["workername"]

	queries, _ := url.ParseQuery(req.URL.RawQuery)
	worker := models.Worker{}
	workers := make(Workers, 0)

	query := func(c *mgo.Collection) error {
		iter := c.Find(bson.M{"name": workerName}).Iter()
		for iter.Next(&worker) {
			apiWorker := &ApiWorker{
				Name:               worker.Name,
				Version:            worker.Version,
				Environment:        worker.Environment,
				Hostname:           worker.Hostname,
				ServiceGenericName: worker.ServiceGenericName,
				ServiceUniqueName:  worker.ServiceUniqueName,
				Uuid:               worker.Uuid,
				Timestamp:          worker.Timestamp,
				Pid:                worker.Pid,
				State:              StatusCode[worker.Status],
				Uptime:             worker.Monitor.Uptime,
				Port:               worker.Port,
			}

			workers = append(workers, *apiWorker)
		}

		return nil
	}

	kontrolDB.RunOnDatabase(WorkersDB, WorkersCollection, query)

	// use http for all workers because they don't have ssl certs
	protocolScheme := "http:"

	brokerConf := config.Broker{}
	switch workerName {
	case "brokerKite":
		brokerConf = conf.BrokerKite
	case "broker":
		brokerConf = conf.Broker
	case "premiumBrokerKite":
		brokerConf = conf.PremiumBrokerKite
	case "premiumBroker":
		brokerConf = conf.PremiumBroker
	default:
	}

	// broker has ssl cert and a custom url scheme, look what it's it
	// TODO: this is ugly, fix some time
	if workerName == "broker" ||
		workerName == "brokerKite" ||
		workerName == "premiumBroker" ||
		workerName == "premiumBrokerKite" {

		protocolScheme = "https:"
		if brokerConf.WebProtocol != "" {
			protocolScheme = brokerConf.WebProtocol
		}
	}

	hostnames := make([]string, len(workers))
	for i, worker := range workers {
		hostnames[i] = fmt.Sprintf("%s//%s:%d", protocolScheme, worker.Hostname, worker.Port)
	}

	var data []byte
	var err error

	_, ok := queries["all"]
	if ok {
		// return all hostnames back
		data, err = json.MarshalIndent(hostnames, "", "  ")
		if err != nil {
			io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
			return
		}
	} else {
		// quit early when there is no workers available
		if len(workers) == 0 {
			writer.Write([]byte(string("")))
			return
		}

		// return only one hostname back, roundrobin
		oldIndex := index.Get() // gives 0 for first time

		N := float64(len(hostnames))
		newIndex := int(math.Mod(float64(oldIndex+1), N))
		hostname := hostnames[newIndex]

		index.CompareAndSwap(oldIndex, uint32(newIndex))

		data = []byte(fmt.Sprintf("\"%s\"", hostname))
	}

	writer.Write(data)
}

type AtomicUint32 uint32

func (i *AtomicUint32) Add(n uint32) uint32 {
	return atomic.AddUint32((*uint32)(i), n)
}

func (i *AtomicUint32) Set(n uint32) {
	atomic.StoreUint32((*uint32)(i), n)
}

func (i *AtomicUint32) Get() uint32 {
	return atomic.LoadUint32((*uint32)(i))
}

func (i *AtomicUint32) CompareAndSwap(oldval, newval uint32) (swapped bool) {
	return atomic.CompareAndSwapUint32((*uint32)(i), oldval, newval)
}

func GetWorkers(writer http.ResponseWriter, req *http.Request) {
	queries, _ := url.ParseQuery(req.URL.RawQuery)

	var latestVersion bool
	var sortFields []string // not initialized means do not sort
	query := bson.M{}
	for key, value := range queries {
		switch key {
		case "version", "pid":
			if value[0] == "latest" {
				latestVersion = true
				break
			}
			v, _ := strconv.Atoi(value[0])
			query[key] = v
		case "state":
			for status, state := range StatusCode {
				if value[0] == state {
					query["status"] = status
				}
			}
		case "sort":
			sortFields = []string{value[0]}
			// override "state" with status, they are not the same in db
			if value[0] == "state" {
				sortFields = []string{"status"}
			}
		default:
			if key == "name" {
				name := value[0]
				if counts := strings.Count(value[0], "-"); counts > 0 {
					s := strings.Split(value[0], "-")
					name = s[0]
				}
				// if searched for social-1, social-2, then return all workers
				// that begins with social
				query[key] = bson.RegEx{Pattern: name, Options: "i"}
			} else {
				query[key] = value[0]
			}
		}
	}

	matchedWorkers := queryResult(query, latestVersion, sortFields)
	data, err := json.MarshalIndent(matchedWorkers, "", "  ")
	if err != nil {
		io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
		return
	}
	writer.Write(data)

}

func queryResult(query bson.M, latestVersion bool, sortFields []string) Workers {
	workers := make(Workers, 0)
	worker := models.Worker{}

	queryFunc := func(c *mgo.Collection) error {
		// sorting is no-op when sortFields is empty
		iter := c.Find(query).Sort(sortFields...).Iter()
		for iter.Next(&worker) {
			apiWorker := &ApiWorker{
				Name:               worker.Name,
				Version:            worker.Version,
				Environment:        worker.Environment,
				Hostname:           worker.Hostname,
				ServiceGenericName: worker.ServiceGenericName,
				ServiceUniqueName:  worker.ServiceUniqueName,
				Uuid:               worker.Uuid,
				Timestamp:          worker.Timestamp,
				Pid:                worker.Pid,
				State:              StatusCode[worker.Status],
				Uptime:             worker.Monitor.Uptime,
				Port:               worker.Port,
			}

			workers = append(workers, *apiWorker)
		}
		return nil
	}

	kontrolDB.RunOnDatabase(WorkersDB, WorkersCollection, queryFunc)

	// finding the largest number of a field in mongo is kinda problematic.
	// therefore we are doing it on our side
	if latestVersion {
		versions := make([]int, len(workers))

		if len(workers) == 0 {
			return workers
		}

		for i, val := range workers {
			versions[i] = val.Version
		}

		sort.Ints(versions)
		maxVersion := versions[len(versions)-1] // get largest version number

		filteredWorkers := make(Workers, 0)
		for _, val := range workers {
			if maxVersion == val.Version {
				filteredWorkers = append(filteredWorkers, val)
			}
		}

		return filteredWorkers
	}

	return workers
}

func GetWorker(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	uuid := vars["uuid"]

	query := bson.M{"uuid": uuid}
	matchedWorkers := queryResult(query, false, nil)
	data, err := json.MarshalIndent(matchedWorkers, "", "  ")
	if err != nil {
		io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
		return
	}
	writer.Write(data)
}
<|endoftext|>


<commit_before>package api

import (
	"net/http"
	"net/url"

	"socialapi/workers/common/response"
	"socialapi/workers/payment"
)

func InitCheckers() error {
	err := payment.InitCheckers()
	return err
}

//----------------------------------------------------------
// Subscribe
//----------------------------------------------------------

func AccountSubscribe(u *url.URL, h http.Header, req *payment.AccountSubscribeRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Do(),
	)
}

func GroupSubscribe(u *url.URL, h http.Header, req *payment.GroupSubscribeRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Do(),
	)
}

func AccountSubscriptionRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	subscriptionRequest := &payment.AccountRequest{
		AccountId: u.Query().Get("account_id"),
	}

	return response.HandleResultAndClientError(
		subscriptionRequest.Subscriptions(),
	)
}

func AccountCancelSubscriptionRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	subscriptionRequest := &payment.AccountRequest{
		AccountId: u.Query().Get("account_id"),
	}

	return response.HandleResultAndClientError(
		subscriptionRequest.CancelSubscription(),
	)
}

func GroupSubscriptionRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	subscriptionRequest := &payment.GroupRequest{
		GroupId: u.Query().Get("group_id"),
	}

	return response.HandleResultAndClientError(
		subscriptionRequest.Subscriptions(),
	)
}

func GetCustomersRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.AccountRequest{}

	return response.HandleResultAndClientError(
		req.ActiveUsernames(),
	)
}

func ExpireCustomerRequest(u *url.URL, h http.Header, req *payment.AccountRequest) (int, http.Header, interface{}, error) {
	req.AccountId = u.Query().Get("accountId")

	return response.HandleResultAndClientError(
		req.Expire(),
	)
}

func DeleteCustomerRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.AccountRequest{
		AccountId: u.Query().Get("accountId"),
	}

	return response.HandleResultAndClientError(
		req.Delete(),
	)
}

func InvoiceRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.AccountRequest{
		AccountId: u.Query().Get("accountId"),
	}

	return response.HandleResultAndClientError(
		req.Invoices(),
	)
}

func CreditCardRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.AccountRequest{
		AccountId: u.Query().Get("accountId"),
	}

	return response.HandleResultAndClientError(
		req.GetCreditCard(),
	)
}

func AccountUpdateCreditCardRequest(u *url.URL, h http.Header, req *payment.AccountUpdateCreditCardRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Do(),
	)
}

func GroupUpdateCreditCardRequest(u *url.URL, h http.Header, req *payment.GroupUpdateCreditCardRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Do(),
	)
}

//----------------------------------------------------------
// Paypal
//----------------------------------------------------------

func PaypalGetToken(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.PaypalGetTokenRequest{
		PlanTitle:    u.Query().Get("planTitle"),
		PlanInterval: u.Query().Get("planInterval"),
	}

	return response.HandleResultAndClientError(
		req.Do(),
	)
}

func PaypalSuccess(u *url.URL, h http.Header, req *payment.PaypalRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Success(),
	)
}

func PaypalCancel(u *url.URL, h http.Header, req *payment.PaypalRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Cancel(),
	)
}
<commit_msg>fixed PUT method for subscription cancelation<commit_after>package api

import (
	"net/http"
	"net/url"

	"socialapi/workers/common/response"
	"socialapi/workers/payment"
)

func InitCheckers() error {
	err := payment.InitCheckers()
	return err
}

//----------------------------------------------------------
// Subscribe
//----------------------------------------------------------

func AccountSubscribe(u *url.URL, h http.Header, req *payment.AccountSubscribeRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Do(),
	)
}

func GroupSubscribe(u *url.URL, h http.Header, req *payment.GroupSubscribeRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Do(),
	)
}

func AccountSubscriptionRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	subscriptionRequest := &payment.AccountRequest{
		AccountId: u.Query().Get("account_id"),
	}

	return response.HandleResultAndClientError(
		subscriptionRequest.Subscriptions(),
	)
}

func AccountCancelSubscriptionRequest(u *url.URL, h http.Header, req *payment.AccountRequest) (int, http.Header, interface{}, error) {
	req.AccountId = u.Query().Get("account_id")

	return response.HandleResultAndClientError(
		req.CancelSubscription(),
	)
}

func GroupSubscriptionRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	subscriptionRequest := &payment.GroupRequest{
		GroupId: u.Query().Get("group_id"),
	}

	return response.HandleResultAndClientError(
		subscriptionRequest.Subscriptions(),
	)
}

func GetCustomersRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.AccountRequest{}

	return response.HandleResultAndClientError(
		req.ActiveUsernames(),
	)
}

func ExpireCustomerRequest(u *url.URL, h http.Header, req *payment.AccountRequest) (int, http.Header, interface{}, error) {
	req.AccountId = u.Query().Get("accountId")

	return response.HandleResultAndClientError(
		req.Expire(),
	)
}

func DeleteCustomerRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.AccountRequest{
		AccountId: u.Query().Get("accountId"),
	}

	return response.HandleResultAndClientError(
		req.Delete(),
	)
}

func InvoiceRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.AccountRequest{
		AccountId: u.Query().Get("accountId"),
	}

	return response.HandleResultAndClientError(
		req.Invoices(),
	)
}

func CreditCardRequest(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.AccountRequest{
		AccountId: u.Query().Get("accountId"),
	}

	return response.HandleResultAndClientError(
		req.GetCreditCard(),
	)
}

func AccountUpdateCreditCardRequest(u *url.URL, h http.Header, req *payment.AccountUpdateCreditCardRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Do(),
	)
}

func GroupUpdateCreditCardRequest(u *url.URL, h http.Header, req *payment.GroupUpdateCreditCardRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Do(),
	)
}

//----------------------------------------------------------
// Paypal
//----------------------------------------------------------

func PaypalGetToken(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	req := &payment.PaypalGetTokenRequest{
		PlanTitle:    u.Query().Get("planTitle"),
		PlanInterval: u.Query().Get("planInterval"),
	}

	return response.HandleResultAndClientError(
		req.Do(),
	)
}

func PaypalSuccess(u *url.URL, h http.Header, req *payment.PaypalRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Success(),
	)
}

func PaypalCancel(u *url.URL, h http.Header, req *payment.PaypalRequest) (int, http.Header, interface{}, error) {
	return response.HandleResultAndClientError(
		req.Cancel(),
	)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"
	"errors"
	"io/ioutil"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
)

type S3CreateBucketImpl struct{}

func (dt S3CreateBucketImpl) CreateBucket(ctx context.Context,
	params *s3.CreateBucketInput,
	optFns ...func(*s3.Options)) (*s3.CreateBucketOutput, error) {

	output := &s3.CreateBucketOutput{
		Location: aws.String("us-west-2"),
	}

	return output, nil
}

type Config struct {
	BucketName string `json:"BucketName"`
}

var configFileName = "config.json"

var globalConfig Config

func populateConfiguration(t *testing.T) error {
	content, err := ioutil.ReadFile(configFileName)
	if err != nil {
		return err
	}

	text := string(content)

	err = json.Unmarshal([]byte(text), &globalConfig)
	if err != nil {
		return err
	}

	if globalConfig.BucketName == "" {
		msg := "You must supply a BucketName value in " + configFileName
		return errors.New(msg)
	}

	t.Log("BucketName: " + globalConfig.BucketName)

	return nil
}

func TestCreateBucket(t *testing.T) {
	thisTime := time.Now()
	nowString := thisTime.Format("2006-01-02 15:04:05 Monday")
	t.Log("Starting unit test at " + nowString)

	err := populateConfiguration(t)
	if err != nil {
		t.Fatal(err)
	}

	input := s3.CreateBucketInput{
		Bucket: &globalConfig.BucketName,
	}

	api := &S3CreateBucketImpl{}

	resp, err := MakeBucket(context.Background(), *api, &input)
	if err != nil {
		t.Log("Got an error ...:")
		t.Log(err)
		return
	}

	t.Log("Created bucket " + globalConfig.BucketName + " in " + *resp.Location)
}
<commit_msg>Go CI test that should fail<commit_after>package main

import (
	"context"
	"encoding/json"
	"errors"
	"io/ioutil"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
)

type S3CreateBucketImpl struct{}

func (dt S3CreateBucketImpl) CreateBucket(ctx context.Context,
	params *s3.CreateBucketInput,
	optFns ...func(*s3.Options)) (*s3.CreateBucketOutput, error) {

	output := &s3.CreateBucketOutput{
		Location: aws.String("us-west-2"),
	}

	//return output, nil
	return output, errors.New("error")
}

type Config struct {
	BucketName string `json:"BucketName"`
}

var configFileName = "config.json"

var globalConfig Config

func populateConfiguration(t *testing.T) error {
	content, err := ioutil.ReadFile(configFileName)
	if err != nil {
		return err
	}

	text := string(content)

	err = json.Unmarshal([]byte(text), &globalConfig)
	if err != nil {
		return err
	}

	if globalConfig.BucketName == "" {
		msg := "You must supply a BucketName value in " + configFileName
		return errors.New(msg)
	}

	t.Log("BucketName: " + globalConfig.BucketName)

	return nil
}

func TestCreateBucket(t *testing.T) {
	thisTime := time.Now()
	nowString := thisTime.Format("2006-01-02 15:04:05 Monday")
	t.Log("Starting unit test at " + nowString)

	err := populateConfiguration(t)
	if err != nil {
		t.Fatal(err)
	}

	input := s3.CreateBucketInput{
		Bucket: &globalConfig.BucketName,
	}

	api := &S3CreateBucketImpl{}

	resp, err := MakeBucket(context.Background(), *api, &input)
	if err != nil {
		t.Log("Got an error ...:")
		t.Log(err)
		return
	}

	t.Log("Created bucket " + globalConfig.BucketName + " in " + *resp.Location)
}
<|endoftext|>


<commit_before>package rest

import (
	"strings"
	"time"

	"github.com/lfq7413/tomato/authdatamanager"
	"github.com/lfq7413/tomato/config"
	"github.com/lfq7413/tomato/orm"
	"github.com/lfq7413/tomato/utils"
)

// Write ...
type Write struct {
	auth         *Auth
	className    string
	query        map[string]interface{}
	data         map[string]interface{}
	originalData map[string]interface{}
	storage      map[string]interface{}
	runOptions   map[string]interface{}
	response     map[string]interface{}
	updatedAt    time.Time
}

// NewWrite 可用于 create 和 update ， create 时 query 为 nil
func NewWrite(
	auth *Auth,
	className string,
	query map[string]interface{},
	data map[string]interface{},
	originalData map[string]interface{},
) *Write {
	if query == nil && data["objectId"] != nil {
		// TODO objectId 无效
	}
	write := &Write{
		auth:         auth,
		className:    className,
		query:        query,
		data:         data,
		originalData: originalData,
		storage:      map[string]interface{}{},
		runOptions:   map[string]interface{}{},
		response:     nil,
		updatedAt:    time.Now().UTC(),
	}
	return write
}

// Execute ...
func (w *Write) Execute() map[string]interface{} {
	w.getUserAndRoleACL()
	w.validateClientClassCreation()
	w.validateSchema()
	w.handleInstallation()
	w.handleSession()
	w.validateAuthData()
	w.runBeforeTrigger()
	w.setRequiredFieldsIfNeeded()
	w.transformUser()
	w.expandFilesForExistingObjects()
	w.runDatabaseOperation()
	w.handleFollowup()
	w.runAfterTrigger()
	return w.response
}

func (w *Write) getUserAndRoleACL() error {
	if w.auth.IsMaster {
		return nil
	}
	w.runOptions["acl"] = []string{"*"}
	if w.auth.User != nil {
		roles := w.auth.GetUserRoles()
		roles = append(roles, w.auth.User.ID)
		if v, ok := w.runOptions["acl"].([]string); ok {
			v = utils.AppendString(v, roles)
		}
	}
	return nil
}

func (w *Write) validateClientClassCreation() error {
	sysClass := []string{"_User", "_Installation", "_Role", "_Session", "_Product"}
	if config.TConfig.AllowClientClassCreation {
		return nil
	}
	if w.auth.IsMaster {
		return nil
	}
	for _, v := range sysClass {
		if v == w.className {
			return nil
		}
	}
	if orm.CollectionExists(w.className) {
		return nil
	}
	// TODO 无法操作不存在的表
	return nil
}

func (w *Write) validateSchema() error {
	// TODO 判断是否可以进行操作
	orm.ValidateObject(w.className, w.data, w.query, w.runOptions)
	return nil
}

func (w *Write) handleInstallation() error {
	if w.response != nil || w.className != "_Installation" {
		return nil
	}

	if w.query == nil && w.data["deviceToken"] == nil && w.data["installationId"] == nil {
		// TODO 设备 id 不能为空
		return nil
	}

	if w.query == nil && w.data["deviceType"] == nil {
		// TODO 设备类型不能为空
		return nil
	}

	if w.data["deviceToken"] != nil && len(utils.String(w.data["deviceToken"])) == 64 {
		w.data["deviceToken"] = strings.ToLower(utils.String(w.data["deviceToken"]))
	}

	if w.data["installationId"] != nil {
		w.data["installationId"] = strings.ToLower(utils.String(w.data["installationId"]))
	}

	var idMatch map[string]interface{}
	var deviceTokenMatches []interface{}

	if w.query != nil && w.query["objectId"] != nil {
		results := orm.Find("_Installation", map[string]interface{}{"objectId": w.query["objectId"]}, map[string]interface{}{})
		if results == nil || len(results) == 0 {
			// TODO 更新对象未找到
			return nil
		}
		idMatch = utils.MapInterface(results[0])
		if w.data["installationId"] != nil && idMatch["installationId"] != nil &&
			w.data["installationId"] != idMatch["installationId"] {
			//TODO installationId 不能修改
			return nil
		}
		if w.data["deviceToken"] != nil && idMatch["deviceToken"] != nil &&
			w.data["deviceToken"] != idMatch["deviceToken"] &&
			w.data["installationId"] == nil && idMatch["installationId"] == nil {
			//TODO deviceToken 不能修改
			return nil
		}
		if w.data["deviceType"] != nil && idMatch["deviceType"] != nil &&
			w.data["deviceType"] != idMatch["deviceType"] {
			//TODO deviceType 不能修改
			return nil
		}
	}

	idMatch = nil
	if w.data["installationId"] != nil {
		results := orm.Find("_Installation", map[string]interface{}{"installationId": w.data["installationId"]}, map[string]interface{}{})
		if results != nil && len(results) > 0 {
			idMatch = utils.MapInterface(results[0])
		}
	}
	if w.data["deviceToken"] != nil {
		results := orm.Find("_Installation", map[string]interface{}{"deviceToken": w.data["deviceToken"]}, map[string]interface{}{})
		if results != nil {
			deviceTokenMatches = results
		}
	}

	var objID string
	if idMatch == nil {
		if deviceTokenMatches == nil || len(deviceTokenMatches) == 0 {
			objID = ""
		} else if len(deviceTokenMatches) == 1 &&
			(utils.MapInterface(deviceTokenMatches[0])["installationId"] == nil || w.data["installationId"] == nil) {
			objID = utils.String(utils.MapInterface(deviceTokenMatches[0])["objectId"])
		} else if w.data["installationId"] == nil {
			// TODO 当有多个 deviceToken 时，必须指定 installationId
			return nil
		} else {
			// 清理多余数据
			installationID := map[string]interface{}{
				"$ne": w.data["installationId"],
			}
			delQuery := map[string]interface{}{
				"deviceToken":    w.data["deviceToken"],
				"installationId": installationID,
			}
			if w.data["appIdentifier"] != nil {
				delQuery["appIdentifier"] = w.data["appIdentifier"]
			}
			orm.Destroy("_Installation", delQuery)
			objID = ""
		}
	} else {
		if deviceTokenMatches != nil && len(deviceTokenMatches) == 1 &&
			utils.MapInterface(deviceTokenMatches[0])["installationId"] == nil {
			// 合并
			delQuery := map[string]interface{}{
				"objectId": idMatch["objectId"],
			}
			orm.Destroy("_Installation", delQuery)
			objID = utils.String(utils.MapInterface(deviceTokenMatches[0])["objectId"])
		} else {
			if w.data["deviceToken"] != nil && idMatch["deviceToken"] != w.data["deviceToken"] {
				// 清理多余数据
				installationID := map[string]interface{}{
					"$ne": w.data["installationId"],
				}
				delQuery := map[string]interface{}{
					"deviceToken":    w.data["deviceToken"],
					"installationId": installationID,
				}
				if w.data["appIdentifier"] != nil {
					delQuery["appIdentifier"] = w.data["appIdentifier"]
				}
				orm.Destroy("_Installation", delQuery)
			}
			objID = utils.String(idMatch["objectId"])
		}
	}
	if objID != "" {
		w.query = map[string]interface{}{
			"objectId": objID,
		}
		delete(w.data, "objectId")
		delete(w.data, "createdAt")
	}
	// TODO Validate ops (add/remove on channels, $inc on badge, etc.)

	return nil
}

func (w *Write) handleSession() error {
	if w.response != nil || w.className != "_Session" {
		return nil
	}

	if w.auth.User == nil && w.auth.IsMaster == false {
		// TODO 需要 Session token
		return nil
	}

	if w.data["ACL"] != nil {
		// TODO Session 不能设置 ACL
		return nil
	}

	if w.query == nil && w.auth.IsMaster == false {
		token := "r:" + utils.CreateToken()
		expiresAt := time.Now().UTC()
		expiresAt.AddDate(1, 0, 0)
		user := map[string]interface{}{
			"__type":    "Pointer",
			"className": "_User",
			"objectId":  w.auth.User.ID,
		}
		createdWith := map[string]interface{}{
			"action": "create",
		}
		sessionData := map[string]interface{}{
			"sessionToken": token,
			"user":         user,
			"createdWith":  createdWith,
			"restricted":   true,
			"expiresAt":    expiresAt,
		}
		for k, v := range w.data {
			if k == "objectId" {
				continue
			}
			sessionData[k] = v
		}
		results := NewWrite(Master(), "_Session", nil, sessionData, map[string]interface{}{}).Execute()
		if results["response"] == nil {
			// TODO 创建 Session 失败
			return nil
		}
		sessionData["objectId"] = utils.MapInterface(results["response"])["objectId"]
		w.response = map[string]interface{}{
			"status":   201,
			"location": results["location"],
			"response": sessionData,
		}
	}

	return nil
}

func (w *Write) validateAuthData() error {
	if w.className != "_User" {
		return nil
	}

	if w.query == nil && w.data["authData"] == nil {
		if utils.String(w.data["username"]) == "" {
			// TODO 没有设置 username
			return nil
		}
		if utils.String(w.data["password"]) == "" {
			// TODO 没有设置 password
			return nil
		}
	}

	if w.data["authData"] == nil || len(utils.MapInterface(w.data["authData"])) == 0 {
		return nil
	}

	authData := utils.MapInterface(w.data["authData"])
	canHandleAuthData := true

	for _, v := range authData {
		providerAuthData := utils.MapInterface(v)
		hasToken := (providerAuthData != nil && providerAuthData["id"] != nil)
		canHandleAuthData = (canHandleAuthData && (hasToken || providerAuthData == nil))
	}
	if canHandleAuthData {
		return w.handleAuthData(authData)
	}
	// TODO 这个 authentication 不支持
	return nil
}

func (w *Write) handleAuthData(authData map[string]interface{}) error {
	w.handleAuthDataValidation(authData)
	results := w.findUsersWithAuthData(authData)
	if results != nil && len(results) > 1 {
		// TODO auth 已经被使用
		return nil
	}

	keys := []string{}
	for k := range authData {
		keys = append(keys, k)
	}
	w.storage["authProvider"] = strings.Join(keys, ",")

	if results == nil || len(results) == 0 {
		w.data["username"] = utils.CreateToken()
	} else if w.query == nil {
		// 登录
		user := utils.MapInterface(results[0])
		delete(user, "password")
		w.response = map[string]interface{}{
			"response": user,
			"location": w.location(),
		}
		w.data["objectId"] = user["objectId"]
	} else if w.query != nil && w.query["objectId"] != nil {
		// 更新
		user := utils.MapInterface(results[0])
		if utils.String(user["objectId"]) != utils.String(w.query["objectId"]) {
			// TODO auth 已经被使用
			return nil
		}
	}

	return nil
}

func (w *Write) handleAuthDataValidation(authData map[string]interface{}) error {
	for k, v := range authData {
		if v == nil {
			continue
		}
		result := authdatamanager.ValidateAuthData(k, utils.MapInterface(v))
		if result != nil {
			// 验证出现问题
			return nil
		}
	}

	return nil
}

func (w *Write) findUsersWithAuthData(authData map[string]interface{}) []interface{} {
	query := []interface{}{}
	for k, v := range authData {
		if v == nil {
			continue
		}
		key := "authData." + k + ".id"
		provider := utils.MapInterface(v)
		q := map[string]interface{}{
			key: provider["id"],
		}
		query = append(query, q)
	}

	findPromise := []interface{}{}
	if len(query) > 0 {
		where := map[string]interface{}{
			"$or": query,
		}
		findPromise = orm.Find(w.className, where, map[string]interface{}{})
	}

	return findPromise
}

func (w *Write) runBeforeTrigger() error {
	if w.response != nil {
		return nil
	}
	if TriggerExists(TypeBeforeSave, w.className) == false {
		return nil
	}

	updatedObject := map[string]interface{}{}
	if w.query != nil && w.query["objectId"] != nil {
		// 如果是更新，则把原始数据添加进来
		for k, v := range w.originalData {
			updatedObject[k] = v
		}
	}
	// 把需要更新的数据添加进来
	for k, v := range w.data {
		updatedObject[k] = v
	}

	response := RunTrigger(TypeBeforeSave, w.className, w.auth, updatedObject, w.originalData)
	if response != nil && response["object"] != nil {
		w.data = utils.MapInterface(response["object"])
		w.storage["changedByTrigger"] = true
		if w.query != nil && w.query["objectId"] != nil {
			delete(w.data, "objectId")
		}
	}

	return nil
}

func (w *Write) setRequiredFieldsIfNeeded() error {
	return nil
}

func (w *Write) transformUser() error {
	return nil
}

func (w *Write) expandFilesForExistingObjects() error {
	return nil
}

func (w *Write) runDatabaseOperation() error {
	return nil
}

func (w *Write) handleFollowup() error {
	return nil
}

func (w *Write) runAfterTrigger() error {
	return nil
}

func (w *Write) location() string {
	var middle string
	if w.className == "_User" {
		middle = "/users/"
	} else {
		middle = "/classes/" + w.className + "/"
	}
	return config.TConfig.ServerURL + middle + utils.String(w.data["objectId"])
}
<commit_msg>完成 setRequiredFieldsIfNeeded<commit_after>package rest

import (
	"strings"
	"time"

	"github.com/lfq7413/tomato/authdatamanager"
	"github.com/lfq7413/tomato/config"
	"github.com/lfq7413/tomato/orm"
	"github.com/lfq7413/tomato/utils"
)

// Write ...
type Write struct {
	auth         *Auth
	className    string
	query        map[string]interface{}
	data         map[string]interface{}
	originalData map[string]interface{}
	storage      map[string]interface{}
	runOptions   map[string]interface{}
	response     map[string]interface{}
	updatedAt    time.Time
}

// NewWrite 可用于 create 和 update ， create 时 query 为 nil
func NewWrite(
	auth *Auth,
	className string,
	query map[string]interface{},
	data map[string]interface{},
	originalData map[string]interface{},
) *Write {
	if query == nil && data["objectId"] != nil {
		// TODO objectId 无效
	}
	write := &Write{
		auth:         auth,
		className:    className,
		query:        query,
		data:         data,
		originalData: originalData,
		storage:      map[string]interface{}{},
		runOptions:   map[string]interface{}{},
		response:     nil,
		updatedAt:    time.Now().UTC(),
	}
	return write
}

// Execute ...
func (w *Write) Execute() map[string]interface{} {
	w.getUserAndRoleACL()
	w.validateClientClassCreation()
	w.validateSchema()
	w.handleInstallation()
	w.handleSession()
	w.validateAuthData()
	w.runBeforeTrigger()
	w.setRequiredFieldsIfNeeded()
	w.transformUser()
	w.expandFilesForExistingObjects()
	w.runDatabaseOperation()
	w.handleFollowup()
	w.runAfterTrigger()
	return w.response
}

func (w *Write) getUserAndRoleACL() error {
	if w.auth.IsMaster {
		return nil
	}
	w.runOptions["acl"] = []string{"*"}
	if w.auth.User != nil {
		roles := w.auth.GetUserRoles()
		roles = append(roles, w.auth.User.ID)
		if v, ok := w.runOptions["acl"].([]string); ok {
			v = utils.AppendString(v, roles)
		}
	}
	return nil
}

func (w *Write) validateClientClassCreation() error {
	sysClass := []string{"_User", "_Installation", "_Role", "_Session", "_Product"}
	if config.TConfig.AllowClientClassCreation {
		return nil
	}
	if w.auth.IsMaster {
		return nil
	}
	for _, v := range sysClass {
		if v == w.className {
			return nil
		}
	}
	if orm.CollectionExists(w.className) {
		return nil
	}
	// TODO 无法操作不存在的表
	return nil
}

func (w *Write) validateSchema() error {
	// TODO 判断是否可以进行操作
	orm.ValidateObject(w.className, w.data, w.query, w.runOptions)
	return nil
}

func (w *Write) handleInstallation() error {
	if w.response != nil || w.className != "_Installation" {
		return nil
	}

	if w.query == nil && w.data["deviceToken"] == nil && w.data["installationId"] == nil {
		// TODO 设备 id 不能为空
		return nil
	}

	if w.query == nil && w.data["deviceType"] == nil {
		// TODO 设备类型不能为空
		return nil
	}

	if w.data["deviceToken"] != nil && len(utils.String(w.data["deviceToken"])) == 64 {
		w.data["deviceToken"] = strings.ToLower(utils.String(w.data["deviceToken"]))
	}

	if w.data["installationId"] != nil {
		w.data["installationId"] = strings.ToLower(utils.String(w.data["installationId"]))
	}

	var idMatch map[string]interface{}
	var deviceTokenMatches []interface{}

	if w.query != nil && w.query["objectId"] != nil {
		results := orm.Find("_Installation", map[string]interface{}{"objectId": w.query["objectId"]}, map[string]interface{}{})
		if results == nil || len(results) == 0 {
			// TODO 更新对象未找到
			return nil
		}
		idMatch = utils.MapInterface(results[0])
		if w.data["installationId"] != nil && idMatch["installationId"] != nil &&
			w.data["installationId"] != idMatch["installationId"] {
			//TODO installationId 不能修改
			return nil
		}
		if w.data["deviceToken"] != nil && idMatch["deviceToken"] != nil &&
			w.data["deviceToken"] != idMatch["deviceToken"] &&
			w.data["installationId"] == nil && idMatch["installationId"] == nil {
			//TODO deviceToken 不能修改
			return nil
		}
		if w.data["deviceType"] != nil && idMatch["deviceType"] != nil &&
			w.data["deviceType"] != idMatch["deviceType"] {
			//TODO deviceType 不能修改
			return nil
		}
	}

	idMatch = nil
	if w.data["installationId"] != nil {
		results := orm.Find("_Installation", map[string]interface{}{"installationId": w.data["installationId"]}, map[string]interface{}{})
		if results != nil && len(results) > 0 {
			idMatch = utils.MapInterface(results[0])
		}
	}
	if w.data["deviceToken"] != nil {
		results := orm.Find("_Installation", map[string]interface{}{"deviceToken": w.data["deviceToken"]}, map[string]interface{}{})
		if results != nil {
			deviceTokenMatches = results
		}
	}

	var objID string
	if idMatch == nil {
		if deviceTokenMatches == nil || len(deviceTokenMatches) == 0 {
			objID = ""
		} else if len(deviceTokenMatches) == 1 &&
			(utils.MapInterface(deviceTokenMatches[0])["installationId"] == nil || w.data["installationId"] == nil) {
			objID = utils.String(utils.MapInterface(deviceTokenMatches[0])["objectId"])
		} else if w.data["installationId"] == nil {
			// TODO 当有多个 deviceToken 时，必须指定 installationId
			return nil
		} else {
			// 清理多余数据
			installationID := map[string]interface{}{
				"$ne": w.data["installationId"],
			}
			delQuery := map[string]interface{}{
				"deviceToken":    w.data["deviceToken"],
				"installationId": installationID,
			}
			if w.data["appIdentifier"] != nil {
				delQuery["appIdentifier"] = w.data["appIdentifier"]
			}
			orm.Destroy("_Installation", delQuery)
			objID = ""
		}
	} else {
		if deviceTokenMatches != nil && len(deviceTokenMatches) == 1 &&
			utils.MapInterface(deviceTokenMatches[0])["installationId"] == nil {
			// 合并
			delQuery := map[string]interface{}{
				"objectId": idMatch["objectId"],
			}
			orm.Destroy("_Installation", delQuery)
			objID = utils.String(utils.MapInterface(deviceTokenMatches[0])["objectId"])
		} else {
			if w.data["deviceToken"] != nil && idMatch["deviceToken"] != w.data["deviceToken"] {
				// 清理多余数据
				installationID := map[string]interface{}{
					"$ne": w.data["installationId"],
				}
				delQuery := map[string]interface{}{
					"deviceToken":    w.data["deviceToken"],
					"installationId": installationID,
				}
				if w.data["appIdentifier"] != nil {
					delQuery["appIdentifier"] = w.data["appIdentifier"]
				}
				orm.Destroy("_Installation", delQuery)
			}
			objID = utils.String(idMatch["objectId"])
		}
	}
	if objID != "" {
		w.query = map[string]interface{}{
			"objectId": objID,
		}
		delete(w.data, "objectId")
		delete(w.data, "createdAt")
	}
	// TODO Validate ops (add/remove on channels, $inc on badge, etc.)

	return nil
}

func (w *Write) handleSession() error {
	if w.response != nil || w.className != "_Session" {
		return nil
	}

	if w.auth.User == nil && w.auth.IsMaster == false {
		// TODO 需要 Session token
		return nil
	}

	if w.data["ACL"] != nil {
		// TODO Session 不能设置 ACL
		return nil
	}

	if w.query == nil && w.auth.IsMaster == false {
		token := "r:" + utils.CreateToken()
		expiresAt := time.Now().UTC()
		expiresAt.AddDate(1, 0, 0)
		user := map[string]interface{}{
			"__type":    "Pointer",
			"className": "_User",
			"objectId":  w.auth.User.ID,
		}
		createdWith := map[string]interface{}{
			"action": "create",
		}
		sessionData := map[string]interface{}{
			"sessionToken": token,
			"user":         user,
			"createdWith":  createdWith,
			"restricted":   true,
			"expiresAt":    expiresAt,
		}
		for k, v := range w.data {
			if k == "objectId" {
				continue
			}
			sessionData[k] = v
		}
		results := NewWrite(Master(), "_Session", nil, sessionData, map[string]interface{}{}).Execute()
		if results["response"] == nil {
			// TODO 创建 Session 失败
			return nil
		}
		sessionData["objectId"] = utils.MapInterface(results["response"])["objectId"]
		w.response = map[string]interface{}{
			"status":   201,
			"location": results["location"],
			"response": sessionData,
		}
	}

	return nil
}

func (w *Write) validateAuthData() error {
	if w.className != "_User" {
		return nil
	}

	if w.query == nil && w.data["authData"] == nil {
		if utils.String(w.data["username"]) == "" {
			// TODO 没有设置 username
			return nil
		}
		if utils.String(w.data["password"]) == "" {
			// TODO 没有设置 password
			return nil
		}
	}

	if w.data["authData"] == nil || len(utils.MapInterface(w.data["authData"])) == 0 {
		return nil
	}

	authData := utils.MapInterface(w.data["authData"])
	canHandleAuthData := true

	for _, v := range authData {
		providerAuthData := utils.MapInterface(v)
		hasToken := (providerAuthData != nil && providerAuthData["id"] != nil)
		canHandleAuthData = (canHandleAuthData && (hasToken || providerAuthData == nil))
	}
	if canHandleAuthData {
		return w.handleAuthData(authData)
	}
	// TODO 这个 authentication 不支持
	return nil
}

func (w *Write) handleAuthData(authData map[string]interface{}) error {
	w.handleAuthDataValidation(authData)
	results := w.findUsersWithAuthData(authData)
	if results != nil && len(results) > 1 {
		// TODO auth 已经被使用
		return nil
	}

	keys := []string{}
	for k := range authData {
		keys = append(keys, k)
	}
	w.storage["authProvider"] = strings.Join(keys, ",")

	if results == nil || len(results) == 0 {
		w.data["username"] = utils.CreateToken()
	} else if w.query == nil {
		// 登录
		user := utils.MapInterface(results[0])
		delete(user, "password")
		w.response = map[string]interface{}{
			"response": user,
			"location": w.location(),
		}
		w.data["objectId"] = user["objectId"]
	} else if w.query != nil && w.query["objectId"] != nil {
		// 更新
		user := utils.MapInterface(results[0])
		if utils.String(user["objectId"]) != utils.String(w.query["objectId"]) {
			// TODO auth 已经被使用
			return nil
		}
	}

	return nil
}

func (w *Write) handleAuthDataValidation(authData map[string]interface{}) error {
	for k, v := range authData {
		if v == nil {
			continue
		}
		result := authdatamanager.ValidateAuthData(k, utils.MapInterface(v))
		if result != nil {
			// 验证出现问题
			return nil
		}
	}

	return nil
}

func (w *Write) findUsersWithAuthData(authData map[string]interface{}) []interface{} {
	query := []interface{}{}
	for k, v := range authData {
		if v == nil {
			continue
		}
		key := "authData." + k + ".id"
		provider := utils.MapInterface(v)
		q := map[string]interface{}{
			key: provider["id"],
		}
		query = append(query, q)
	}

	findPromise := []interface{}{}
	if len(query) > 0 {
		where := map[string]interface{}{
			"$or": query,
		}
		findPromise = orm.Find(w.className, where, map[string]interface{}{})
	}

	return findPromise
}

func (w *Write) runBeforeTrigger() error {
	if w.response != nil {
		return nil
	}
	if TriggerExists(TypeBeforeSave, w.className) == false {
		return nil
	}

	updatedObject := map[string]interface{}{}
	if w.query != nil && w.query["objectId"] != nil {
		// 如果是更新，则把原始数据添加进来
		for k, v := range w.originalData {
			updatedObject[k] = v
		}
	}
	// 把需要更新的数据添加进来
	for k, v := range w.data {
		updatedObject[k] = v
	}

	response := RunTrigger(TypeBeforeSave, w.className, w.auth, updatedObject, w.originalData)
	if response != nil && response["object"] != nil {
		w.data = utils.MapInterface(response["object"])
		w.storage["changedByTrigger"] = true
		if w.query != nil && w.query["objectId"] != nil {
			delete(w.data, "objectId")
		}
	}

	return nil
}

func (w *Write) setRequiredFieldsIfNeeded() error {
	if w.data != nil {
		w.data["updatedAt"] = w.updatedAt
		if w.query == nil {
			w.data["createdAt"] = w.updatedAt

			if w.data["objectId"] == nil {
				w.data["objectId"] = utils.CreateObjectID()
			}
		}
	}

	return nil
}

func (w *Write) transformUser() error {
	return nil
}

func (w *Write) expandFilesForExistingObjects() error {
	return nil
}

func (w *Write) runDatabaseOperation() error {
	return nil
}

func (w *Write) handleFollowup() error {
	return nil
}

func (w *Write) runAfterTrigger() error {
	return nil
}

func (w *Write) location() string {
	var middle string
	if w.className == "_User" {
		middle = "/users/"
	} else {
		middle = "/classes/" + w.className + "/"
	}
	return config.TConfig.ServerURL + middle + utils.String(w.data["objectId"])
}
<|endoftext|>


<commit_before>package resources

import (
    "fmt"
    "strings"
    "html/template"
)

type PageModel struct {
    Raw string `json:"raw"`
}

func (c *PageModel) HTML() (res template.HTML) {
    lines := strings.Split(c.Raw, "\n")
    var code bool
    var lineCnt int
    for _, line := range(lines) {
        if code {
            lineCnt++
            res += template.HTML(fmt.Sprintf("%d.\t%s\n", lineCnt, line))
            continue
        }
        if line == "" {continue}
        switch line[0] {
        case '!':
            res += template.HTML(fmt.Sprintf("
%s
\n", line[1:]))
        case '[':
            res += template.HTML(fmt.Sprintf(`<img src="/image/%s"/>`, line[1:]))
        case '*':
            res += template.HTML(fmt.Sprintf("
%s
\n", line[1:]))
        case '{':
            code = true
            res += template.HTML("
")
        default:
            res += template.HTML(fmt.Sprintf("
%s
\n", line))
        }
    }
    if code {
        res += template.HTML("
")
    }
    return
}
<commit_msg>add hyperlink in page<commit_after>package resources

import (
    "fmt"
    "strings"
    "html/template"
)

type PageModel struct {
    Raw string `json:"raw"`
}

func (c *PageModel) HTML() (res template.HTML) {
    lines := strings.Split(c.Raw, "\n")
    var code bool
    var lineCnt int
    for _, line := range(lines) {
        if code {
            lineCnt++
            res += template.HTML(fmt.Sprintf("%d.\t%s\n", lineCnt, line))
            continue
        }
        if line == "" {continue}
        switch line[0] {
        case '!':
            res += template.HTML(fmt.Sprintf("
%s
\n", line[1:]))
        case '[':
            res += template.HTML(fmt.Sprintf(`<img src="/image/%s"/>`, line[1:]))
        case '*':
            res += template.HTML(fmt.Sprintf("
%s
\n", line[1:]))
        case '{':
            code = true
            res += template.HTML("
")
        case '@':
            res += template.HTML(fmt.Sprintf(`<a href="%s" target=_blank><h3>%s</h3></a>`, line[1:], line[1:]))
        default:
            res += template.HTML(fmt.Sprintf("
%s
\n", line))
        }
    }
    if code {
        res += template.HTML("
")
    }
    return
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcdv3

import (
	"crypto/tls"
	"crypto/x509"

	"time"

	"os"
	"strings"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/pkg/tlsutil"
)

// Config represents a part of etcd configuration that can be
// loaded from a file. The Config might be afterwards transformed into
// clientv3.Config using ConfigToClienv3 function.
type Config struct {
	Endpoints             []string      `json:"endpoints"`
	DialTimeout           time.Duration `json:"dial-timeout"`
	InsecureTransport     bool          `json:"insecure-transport"`
	InsecureSkipTLSVerify bool          `json:"insecure-skip-tls-verify"`
	Certfile              string        `json:"cert-file"`
	Keyfile               string        `json:"key-file"`
	CAfile                string        `json:"ca-file"`
}

// default timeout for connecting to etcd.
const defaultTimeout = 1 * time.Second

// ConfigToClientv3 transforms the configuration modelled by yaml structure
// into clientv3.Config. If the endpoints are not specified the function tries to load endpoints
// ETCDV3_ENDPOINTS environment variable.
func ConfigToClientv3(yc *Config) (*clientv3.Config, error) {

	timeout := defaultTimeout

	if yc.DialTimeout != 0 {
		timeout = yc.DialTimeout
	}

	cfg := &clientv3.Config{
		Endpoints:   yc.Endpoints,
		DialTimeout: timeout,
	}

	if len(cfg.Endpoints) == 0 {
		if ep := os.Getenv("ETCDV3_ENDPOINTS"); ep != "" {
			cfg.Endpoints = strings.Split(ep, ",")
		} else {
			cfg.Endpoints = []string{"127.0.0.1:2379"}
		}
	}

	if yc.InsecureTransport {
		cfg.TLS = nil
		return cfg, nil
	}

	var (
		cert *tls.Certificate
		cp   *x509.CertPool
		err  error
	)

	if yc.Certfile != "" && yc.Keyfile != "" {
		cert, err = tlsutil.NewCert(yc.Certfile, yc.Keyfile, nil)
		if err != nil {
			return nil, err
		}
	}

	if yc.CAfile != "" {
		cp, err = tlsutil.NewCertPool([]string{yc.CAfile})
		if err != nil {
			return nil, err
		}
	}

	tlscfg := &tls.Config{
		MinVersion:         tls.VersionTLS10,
		InsecureSkipVerify: yc.InsecureSkipTLSVerify,
		RootCAs:            cp,
	}
	if cert != nil {
		tlscfg.Certificates = []tls.Certificate{*cert}
	}
	cfg.TLS = tlscfg

	return cfg, nil
}
<commit_msg>db/keyval/etcd: apply TLS config only if certFile, keyFile or CAfile specified<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcdv3

import (
	"crypto/tls"
	"crypto/x509"

	"time"

	"os"
	"strings"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/pkg/tlsutil"
)

// Config represents a part of etcd configuration that can be
// loaded from a file. The Config might be afterwards transformed into
// clientv3.Config using ConfigToClienv3 function.
type Config struct {
	Endpoints             []string      `json:"endpoints"`
	DialTimeout           time.Duration `json:"dial-timeout"`
	InsecureTransport     bool          `json:"insecure-transport"`
	InsecureSkipTLSVerify bool          `json:"insecure-skip-tls-verify"`
	Certfile              string        `json:"cert-file"`
	Keyfile               string        `json:"key-file"`
	CAfile                string        `json:"ca-file"`
}

// default timeout for connecting to etcd.
const defaultTimeout = 1 * time.Second

// ConfigToClientv3 transforms the configuration modelled by yaml structure
// into clientv3.Config. If the endpoints are not specified the function tries to load endpoints
// ETCDV3_ENDPOINTS environment variable.
func ConfigToClientv3(yc *Config) (*clientv3.Config, error) {

	timeout := defaultTimeout

	if yc.DialTimeout != 0 {
		timeout = yc.DialTimeout
	}

	cfg := &clientv3.Config{
		Endpoints:   yc.Endpoints,
		DialTimeout: timeout,
	}

	if len(cfg.Endpoints) == 0 {
		if ep := os.Getenv("ETCDV3_ENDPOINTS"); ep != "" {
			cfg.Endpoints = strings.Split(ep, ",")
		} else {
			cfg.Endpoints = []string{"127.0.0.1:2379"}
		}
	}

	if yc.InsecureTransport {
		cfg.TLS = nil
		return cfg, nil
	}

	var (
		cert *tls.Certificate
		cp   *x509.CertPool
		err  error
	)

	if yc.Certfile != "" && yc.Keyfile != "" {
		cert, err = tlsutil.NewCert(yc.Certfile, yc.Keyfile, nil)
		if err != nil {
			return nil, err
		}
	}

	if yc.CAfile != "" {
		cp, err = tlsutil.NewCertPool([]string{yc.CAfile})
		if err != nil {
			return nil, err
		}
	}

	tlscfg := &tls.Config{
		MinVersion:         tls.VersionTLS10,
		InsecureSkipVerify: yc.InsecureSkipTLSVerify,
		RootCAs:            cp,
	}
	if cert != nil {
		tlscfg.Certificates = []tls.Certificate{*cert}
	}
	if yc.Certfile != "" || yc.Keyfile != "" || yc.CAfile != "" {
		cfg.TLS = tlscfg
	}

	return cfg, nil
}
<|endoftext|>