<commit_before>package metadata

import (
	"strings"

	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/scollector/util"
)

func init() {
	metafuncs = append(metafuncs, metaWindowsVersion, metaWindowsIfaces)
}

func metaWindowsVersion() {
	util.ReadCommand(func(line string) error {
		fields := strings.Fields(line)
		if len(fields) == 0 {
			return nil
		}
		AddMeta("", nil, "version", strings.Join(fields, " "), true)
		return nil
	}, "cmd", "/c", "ver")
}

func metaWindowsIfaces() {
	var iface string
	util.ReadCommand(func(line string) error {
		fields := strings.Fields(line)
		sp := strings.Split(line, ":")
		if len(fields) == 0 || len(sp) != 2 {
			return nil
		}
		if line[0] != ' ' {
			iface = sp[0]
		} else if strings.HasPrefix(line, "   IPv4 Address") {
			AddMeta("", opentsdb.TagSet{"iface": iface}, "addr", strings.TrimSpace(sp[1]), true)
		}
		return nil
	}, "ipconfig")
}
<commit_msg>cmd/scollector: Merge pull request #91 from StackExchange/winver<commit_after>package metadata

import (
	"strings"

	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/scollector/util"
	"github.com/StackExchange/slog"
	"github.com/StackExchange/wmi"
)

func init() {
	metafuncs = append(metafuncs, metaWindowsVersion, metaWindowsIfaces)
}

func metaWindowsVersion() {
	var dst []Win32_OperatingSystem
	q := wmi.CreateQuery(&dst, "")
	err := wmi.Query(q, &dst)
	if err != nil {
		slog.Error(err)
		return
	}
	for _, v := range dst {
		AddMeta("", nil, "version", v.Version, true)
		AddMeta("", nil, "versionCaption", v.Caption, true)
	}
}

type Win32_OperatingSystem struct {
	Caption string
	Version string
}

func metaWindowsIfaces() {
	var iface string
	util.ReadCommand(func(line string) error {
		fields := strings.Fields(line)
		sp := strings.Split(line, ":")
		if len(fields) == 0 || len(sp) != 2 {
			return nil
		}
		if line[0] != ' ' {
			iface = sp[0]
		} else if strings.HasPrefix(line, "   IPv4 Address") {
			AddMeta("", opentsdb.TagSet{"iface": iface}, "addr", strings.TrimSpace(sp[1]), true)
		}
		return nil
	}, "ipconfig")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/kr/s3/s3util" //Ref: http://godoc.org/github.com/kr/s3/s3util
	"io"
	"os"
	"path/filepath"
	"strings"
)

var (
	cfg         Config
	locationPtr int //contains the index of the current location being processed
)

func main() {
	cfg = getConfigs()

	if len(cfg.Locations.Source) != len(cfg.Locations.Destination) {
		fmt.Println("Error: The number of sources and destinations should be the same.")
		os.Exit(1)
	}

	if cfg.Auth.AccessKey != "" && cfg.Auth.SecretKey != "" {
		s3util.DefaultConfig.AccessKey = cfg.Auth.AccessKey
		s3util.DefaultConfig.SecretKey = cfg.Auth.SecretKey
	} else {
		fmt.Println("Error: AccessKey and/or SecretKey not set.")
		os.Exit(1)
	}

	//loop through all the source paths to check if they exist before processing
	for _, src := range cfg.Locations.Source {
		if _, err := os.Stat(src); os.IsNotExist(err) {
			fmt.Println(src + " does not exist.")
			os.Exit(1)
		}
	}

	for i, src := range cfg.Locations.Source {
		locationPtr = i
		_ = filepath.Walk(src, upload)
	}
}

func upload(path string, f os.FileInfo, err error) error {
	if !f.IsDir() { //if path is a file
		uploadFile(path)
	}

	return nil
}

func uploadFile(path string) {
	dest := strings.Replace(path, cfg.Locations.Source[locationPtr], "", 1)
	dest = strings.Trim(dest, "\\")             //in case its a Windows path
	dest = strings.Replace(dest, "\\", "/", -1) //in case its a Windows path

	fmt.Print(path)

	r, rerr := os.Open(path)
	if rerr != nil {
		fmt.Println("....." + rerr.Error())
		return
	}
	defer r.Close()

	w, werr := s3util.Create(cfg.Locations.Destination[locationPtr]+dest, nil, nil)
	if werr != nil {
		fmt.Println("....." + werr.Error())
		return
	}
	defer w.Close()

	io.Copy(w, r)

	fmt.Println(".....Done")
}

func checkFile(file, dest string) bool {
	r, err := s3util.Open(dest+file, nil)
	if err != nil {
		return false
	}
	defer r.Close()

	return true
}
<commit_msg>Unify the way paths are combined.<commit_after>package main

import (
	"fmt"
	"github.com/kr/s3/s3util" //Ref: http://godoc.org/github.com/kr/s3/s3util
	"io"
	"os"
	"path/filepath"
	"strings"
)

var (
	cfg         Config
	locationPtr int //contains the index of the current location being processed
)

func main() {
	cfg = getConfigs()

	if len(cfg.Locations.Source) != len(cfg.Locations.Destination) {
		fmt.Println("Error: The number of sources and destinations should be the same.")
		os.Exit(1)
	}

	if cfg.Auth.AccessKey != "" && cfg.Auth.SecretKey != "" {
		s3util.DefaultConfig.AccessKey = cfg.Auth.AccessKey
		s3util.DefaultConfig.SecretKey = cfg.Auth.SecretKey
	} else {
		fmt.Println("Error: AccessKey and/or SecretKey not set.")
		os.Exit(1)
	}

	for i := 0; i < len(cfg.Locations.Destination); i++ {
		//remove trailing slashes to unify the way paths are combined
		//ie: wether they are entered in the config file or not it still works
		cfg.Locations.Destination[i] = strings.Trim(cfg.Locations.Destination[i], "/")
	}

	//loop through all the source paths to check if they exist before processing
	for _, src := range cfg.Locations.Source {
		if _, err := os.Stat(src); os.IsNotExist(err) {
			fmt.Println(src + " does not exist.")
			os.Exit(1)
		}
	}

	for i, src := range cfg.Locations.Source {
		locationPtr = i
		_ = filepath.Walk(src, upload)
	}
}

func upload(path string, f os.FileInfo, err error) error {
	if !f.IsDir() { //if path is a file
		uploadFile(path)
	}

	return nil
}

func uploadFile(path string) {
	dest := strings.Replace(path, cfg.Locations.Source[locationPtr], "", 1)
	dest = strings.Trim(dest, "\\/")            //remove *nix and Windows dir separators
	dest = strings.Replace(dest, "\\", "/", -1) //in case its a Windows path

	fmt.Print(dest)

	r, rerr := os.Open(path)
	if rerr != nil {
		fmt.Println("....." + rerr.Error())
		return
	}
	defer r.Close()

	w, werr := s3util.Create(cfg.Locations.Destination[locationPtr]+"/"+dest, nil, nil)
	if werr != nil {
		fmt.Println("....." + werr.Error())
		return
	}
	defer w.Close()

	io.Copy(w, r)

	fmt.Println(".....Done")
}

func checkFile(file, dest string) bool {
	r, err := s3util.Open(dest+file, nil)
	if err != nil {
		return false
	}
	defer r.Close()

	return true
}
<|endoftext|>


<commit_before>package commons

import (
	"math/rand"
	"strconv"
	"time"
	"net/http"
	"encoding/json"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func GetUID() string {
	s := 6
	return randomString(s)
}

func randomString(l int) string {
	r := strconv.Itoa(rand.Intn(10000))

	bytes := make([]byte, l)
	for i := 0; i < l; i++ {
		bytes[i] = byte(randInt(64, 90))
	}
	return string(bytes) + r
}

func randInt(min int, max int) int {
	return min + rand.Intn(max-min)
}<commit_msg>[Fix] remove unused import<commit_after>package commons

import (
	"math/rand"
	"strconv"
	"time"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func GetUID() string {
	s := 6
	return randomString(s)
}

func randomString(l int) string {
	r := strconv.Itoa(rand.Intn(10000))

	bytes := make([]byte, l)
	for i := 0; i < l; i++ {
		bytes[i] = byte(randInt(64, 90))
	}
	return string(bytes) + r
}

func randInt(min int, max int) int {
	return min + rand.Intn(max-min)
}<|endoftext|>


<commit_before>package shell

import (
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/types"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
	"io"
	"path/filepath"
	"sync"
	"time"

	"github.com/chrislusf/seaweedfs/weed/storage/needle"
)

func init() {
	Commands = append(Commands, &commandVolumeTierMove{})
}

type volumeTierMoveJob struct {
	src pb.ServerAddress
	vid needle.VolumeId
}

type commandVolumeTierMove struct {
	activeServers sync.Map
	queues        map[pb.ServerAddress]chan volumeTierMoveJob
	//activeServers     map[pb.ServerAddress]struct{}
	//activeServersLock sync.Mutex
	//activeServersCond *sync.Cond
}

func (c *commandVolumeTierMove) Name() string {
	return "volume.tier.move"
}

func (c *commandVolumeTierMove) Help() string {
	return `change a volume from one disk type to another

	volume.tier.move -fromDiskType=hdd -toDiskType=ssd [-collectionPattern=""] [-fullPercent=95] [-quietFor=1h] [-parallelLimit=4]

	Even if the volume is replicated, only one replica will be changed and the rest replicas will be dropped.
	So "volume.fix.replication" and "volume.balance" should be followed.

`
}

func (c *commandVolumeTierMove) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	tierCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	collectionPattern := tierCommand.String("collectionPattern", "", "match with wildcard characters '*' and '?'")
	fullPercentage := tierCommand.Float64("fullPercent", 95, "the volume reaches the percentage of max volume size")
	quietPeriod := tierCommand.Duration("quietFor", 24*time.Hour, "select volumes without no writes for this period")
	source := tierCommand.String("fromDiskType", "", "the source disk type")
	target := tierCommand.String("toDiskType", "", "the target disk type")
	parallelLimit := tierCommand.Int("parallelLimit", 0, "limit the number of parallel copying jobs")
	applyChange := tierCommand.Bool("force", false, "actually apply the changes")
	if err = tierCommand.Parse(args); err != nil {
		return nil
	}

	if err = commandEnv.confirmIsLocked(args); err != nil {
		return
	}

	fromDiskType := types.ToDiskType(*source)
	toDiskType := types.ToDiskType(*target)

	if fromDiskType == toDiskType {
		return fmt.Errorf("source tier %s is the same as target tier %s", fromDiskType, toDiskType)
	}

	// collect topology information
	topologyInfo, volumeSizeLimitMb, err := collectTopologyInfo(commandEnv, 0)
	if err != nil {
		return err
	}

	// collect all volumes that should change
	volumeIds, err := collectVolumeIdsForTierChange(commandEnv, topologyInfo, volumeSizeLimitMb, fromDiskType, *collectionPattern, *fullPercentage, *quietPeriod)
	if err != nil {
		return err
	}
	fmt.Printf("tier move volumes: %v\n", volumeIds)

	_, allLocations := collectVolumeReplicaLocations(topologyInfo)
	allLocations = filterLocationsByDiskType(allLocations, toDiskType)
	keepDataNodesSorted(allLocations, toDiskType)

	if len(allLocations) > 0 && *parallelLimit > 0 && *parallelLimit < len(allLocations) {
		allLocations = allLocations[:*parallelLimit]
	}

	wg := sync.WaitGroup{}
	bufferLen := len(allLocations)
	c.queues = make(map[pb.ServerAddress]chan volumeTierMoveJob)

	for _, dst := range allLocations {
		destServerAddress := pb.NewServerAddressFromDataNode(dst.dataNode)
		c.queues[destServerAddress] = make(chan volumeTierMoveJob, bufferLen)

		wg.Add(1)
		go func(dst location, jobs <-chan volumeTierMoveJob, applyChanges bool) {
			defer wg.Done()
			for job := range jobs {
				fmt.Fprintf(writer, "moving volume %d from %s to %s with disk type %s ...\n", job.vid, job.src, dst.dataNode.Id, toDiskType.ReadableString())

				locations, found := commandEnv.MasterClient.GetLocations(uint32(job.vid))
				if !found {
					fmt.Printf("volume %d not found", job.vid)
					continue
				}

				unlock := c.Lock(job.src)

				if err := c.doMoveOneVolume(commandEnv, writer, job.vid, toDiskType, locations, job.src, dst, applyChanges); err != nil {
					fmt.Fprintf(writer, "move volume %d %s => %s: %v\n", job.vid, job.src, dst.dataNode.Id, err)
				}
				unlock()
			}
		}(dst, c.queues[destServerAddress], *applyChange)
	}

	for _, vid := range volumeIds {
		if err = c.doVolumeTierMove(commandEnv, writer, vid, toDiskType, allLocations); err != nil {
			fmt.Printf("tier move volume %d: %v\n", vid, err)
		}
		allLocations = rotateDataNodes(allLocations)
	}
	for key, _ := range c.queues {
		close(c.queues[key])
	}

	wg.Wait()

	return nil
}

func (c *commandVolumeTierMove) Lock(key pb.ServerAddress) func() {
	value, _ := c.activeServers.LoadOrStore(key, &sync.Mutex{})
	mtx := value.(*sync.Mutex)
	mtx.Lock()

	return func() { mtx.Unlock() }
}

func filterLocationsByDiskType(dataNodes []location, diskType types.DiskType) (ret []location) {
	for _, loc := range dataNodes {
		_, found := loc.dataNode.DiskInfos[string(diskType)]
		if found {
			ret = append(ret, loc)
		}
	}
	return
}

func rotateDataNodes(dataNodes []location) []location {
	if len(dataNodes) > 0 {
		return append(dataNodes[1:], dataNodes[0])
	} else {
		return dataNodes
	}
}

func isOneOf(server string, locations []wdclient.Location) bool {
	for _, loc := range locations {
		if server == loc.Url {
			return true
		}
	}
	return false
}

func (c *commandVolumeTierMove) doVolumeTierMove(commandEnv *CommandEnv, writer io.Writer, vid needle.VolumeId, toDiskType types.DiskType, allLocations []location) (err error) {
	// find volume location
	locations, found := commandEnv.MasterClient.GetLocations(uint32(vid))
	if !found {
		return fmt.Errorf("volume %d not found", vid)
	}

	// find one server with the most empty volume slots with target disk type
	hasFoundTarget := false
	fn := capacityByFreeVolumeCount(toDiskType)
	for _, dst := range allLocations {
		if fn(dst.dataNode) > 0 && !hasFoundTarget {
			// ask the volume server to replicate the volume
			if isOneOf(dst.dataNode.Id, locations) {
				continue
			}
			var sourceVolumeServer pb.ServerAddress
			for _, loc := range locations {
				if loc.Url != dst.dataNode.Id {
					sourceVolumeServer = loc.ServerAddress()
				}
			}
			if sourceVolumeServer == "" {
				continue
			}
			hasFoundTarget = true

			// adjust volume count
			dst.dataNode.DiskInfos[string(toDiskType)].VolumeCount++

			destServerAddress := pb.NewServerAddressFromDataNode(dst.dataNode)
			c.queues[destServerAddress] <- volumeTierMoveJob{sourceVolumeServer, vid}
		}
	}

	if !hasFoundTarget {
		fmt.Fprintf(writer, "can not find disk type %s for volume %d\n", toDiskType.ReadableString(), vid)
	}

	return nil
}

func (c *commandVolumeTierMove) doMoveOneVolume(commandEnv *CommandEnv, writer io.Writer, vid needle.VolumeId, toDiskType types.DiskType, locations []wdclient.Location, sourceVolumeServer pb.ServerAddress, dst location, applyChanges bool) (err error) {

	// mark all replicas as read only
	if applyChanges {
		if err = markVolumeReplicasWritable(commandEnv.option.GrpcDialOption, vid, locations, false); err != nil {
			return fmt.Errorf("mark volume %d as readonly on %s: %v", vid, locations[0].Url, err)
		}
		if err = LiveMoveVolume(commandEnv.option.GrpcDialOption, writer, vid, sourceVolumeServer, pb.NewServerAddressFromDataNode(dst.dataNode), 5*time.Second, toDiskType.ReadableString(), true); err != nil {

			// mark all replicas as writable
			if err = markVolumeReplicasWritable(commandEnv.option.GrpcDialOption, vid, locations, true); err != nil {
				glog.Errorf("mark volume %d as writable on %s: %v", vid, locations[0].Url, err)
			}

			return fmt.Errorf("move volume %d %s => %s : %v", vid, locations[0].Url, dst.dataNode.Id, err)
		}
	}

	// adjust volume count
	dst.dataNode.DiskInfos[string(toDiskType)].VolumeCount++

	// remove the remaining replicas
	for _, loc := range locations {
		if loc.Url != dst.dataNode.Id && loc.ServerAddress() != sourceVolumeServer {
			if applyChanges {
				if err = deleteVolume(commandEnv.option.GrpcDialOption, vid, loc.ServerAddress()); err != nil {
					fmt.Fprintf(writer, "failed to delete volume %d on %s: %v\n", vid, loc.Url, err)
				}
			}
			// reduce volume count? Not really necessary since they are "more" full and will not be a candidate to move to
		}
	}
	return nil
}

func collectVolumeIdsForTierChange(commandEnv *CommandEnv, topologyInfo *master_pb.TopologyInfo, volumeSizeLimitMb uint64, sourceTier types.DiskType, collectionPattern string, fullPercentage float64, quietPeriod time.Duration) (vids []needle.VolumeId, err error) {

	quietSeconds := int64(quietPeriod / time.Second)
	nowUnixSeconds := time.Now().Unix()

	fmt.Printf("collect %s volumes quiet for: %d seconds\n", sourceTier, quietSeconds)

	vidMap := make(map[uint32]bool)
	eachDataNode(topologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		for _, diskInfo := range dn.DiskInfos {
			for _, v := range diskInfo.VolumeInfos {
				// check collection name pattern
				if collectionPattern != "" {
					matched, err := filepath.Match(collectionPattern, v.Collection)
					if err != nil {
						return
					}
					if !matched {
						continue
					}
				}

				if v.ModifiedAtSecond+quietSeconds < nowUnixSeconds && types.ToDiskType(v.DiskType) == sourceTier {
					if float64(v.Size) > fullPercentage/100*float64(volumeSizeLimitMb)*1024*1024 {
						vidMap[v.Id] = true
					}
				}
			}
		}
	})

	for vid := range vidMap {
		vids = append(vids, needle.VolumeId(vid))
	}

	return
}
<commit_msg>volume.tier.move: avoid double counting<commit_after>package shell

import (
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/types"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
	"io"
	"path/filepath"
	"sync"
	"time"

	"github.com/chrislusf/seaweedfs/weed/storage/needle"
)

func init() {
	Commands = append(Commands, &commandVolumeTierMove{})
}

type volumeTierMoveJob struct {
	src pb.ServerAddress
	vid needle.VolumeId
}

type commandVolumeTierMove struct {
	activeServers sync.Map
	queues        map[pb.ServerAddress]chan volumeTierMoveJob
	//activeServers     map[pb.ServerAddress]struct{}
	//activeServersLock sync.Mutex
	//activeServersCond *sync.Cond
}

func (c *commandVolumeTierMove) Name() string {
	return "volume.tier.move"
}

func (c *commandVolumeTierMove) Help() string {
	return `change a volume from one disk type to another

	volume.tier.move -fromDiskType=hdd -toDiskType=ssd [-collectionPattern=""] [-fullPercent=95] [-quietFor=1h] [-parallelLimit=4]

	Even if the volume is replicated, only one replica will be changed and the rest replicas will be dropped.
	So "volume.fix.replication" and "volume.balance" should be followed.

`
}

func (c *commandVolumeTierMove) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	tierCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	collectionPattern := tierCommand.String("collectionPattern", "", "match with wildcard characters '*' and '?'")
	fullPercentage := tierCommand.Float64("fullPercent", 95, "the volume reaches the percentage of max volume size")
	quietPeriod := tierCommand.Duration("quietFor", 24*time.Hour, "select volumes without no writes for this period")
	source := tierCommand.String("fromDiskType", "", "the source disk type")
	target := tierCommand.String("toDiskType", "", "the target disk type")
	parallelLimit := tierCommand.Int("parallelLimit", 0, "limit the number of parallel copying jobs")
	applyChange := tierCommand.Bool("force", false, "actually apply the changes")
	if err = tierCommand.Parse(args); err != nil {
		return nil
	}

	if err = commandEnv.confirmIsLocked(args); err != nil {
		return
	}

	fromDiskType := types.ToDiskType(*source)
	toDiskType := types.ToDiskType(*target)

	if fromDiskType == toDiskType {
		return fmt.Errorf("source tier %s is the same as target tier %s", fromDiskType, toDiskType)
	}

	// collect topology information
	topologyInfo, volumeSizeLimitMb, err := collectTopologyInfo(commandEnv, 0)
	if err != nil {
		return err
	}

	// collect all volumes that should change
	volumeIds, err := collectVolumeIdsForTierChange(commandEnv, topologyInfo, volumeSizeLimitMb, fromDiskType, *collectionPattern, *fullPercentage, *quietPeriod)
	if err != nil {
		return err
	}
	fmt.Printf("tier move volumes: %v\n", volumeIds)

	_, allLocations := collectVolumeReplicaLocations(topologyInfo)
	allLocations = filterLocationsByDiskType(allLocations, toDiskType)
	keepDataNodesSorted(allLocations, toDiskType)

	if len(allLocations) > 0 && *parallelLimit > 0 && *parallelLimit < len(allLocations) {
		allLocations = allLocations[:*parallelLimit]
	}

	wg := sync.WaitGroup{}
	bufferLen := len(allLocations)
	c.queues = make(map[pb.ServerAddress]chan volumeTierMoveJob)

	for _, dst := range allLocations {
		destServerAddress := pb.NewServerAddressFromDataNode(dst.dataNode)
		c.queues[destServerAddress] = make(chan volumeTierMoveJob, bufferLen)

		wg.Add(1)
		go func(dst location, jobs <-chan volumeTierMoveJob, applyChanges bool) {
			defer wg.Done()
			for job := range jobs {
				fmt.Fprintf(writer, "moving volume %d from %s to %s with disk type %s ...\n", job.vid, job.src, dst.dataNode.Id, toDiskType.ReadableString())

				locations, found := commandEnv.MasterClient.GetLocations(uint32(job.vid))
				if !found {
					fmt.Printf("volume %d not found", job.vid)
					continue
				}

				unlock := c.Lock(job.src)

				if applyChanges {
					if err := c.doMoveOneVolume(commandEnv, writer, job.vid, toDiskType, locations, job.src, dst); err != nil {
						fmt.Fprintf(writer, "move volume %d %s => %s: %v\n", job.vid, job.src, dst.dataNode.Id, err)
					}
				}
				unlock()
			}
		}(dst, c.queues[destServerAddress], *applyChange)
	}

	for _, vid := range volumeIds {
		if err = c.doVolumeTierMove(commandEnv, writer, vid, toDiskType, allLocations); err != nil {
			fmt.Printf("tier move volume %d: %v\n", vid, err)
		}
		allLocations = rotateDataNodes(allLocations)
	}
	for key, _ := range c.queues {
		close(c.queues[key])
	}

	wg.Wait()

	return nil
}

func (c *commandVolumeTierMove) Lock(key pb.ServerAddress) func() {
	value, _ := c.activeServers.LoadOrStore(key, &sync.Mutex{})
	mtx := value.(*sync.Mutex)
	mtx.Lock()

	return func() { mtx.Unlock() }
}

func filterLocationsByDiskType(dataNodes []location, diskType types.DiskType) (ret []location) {
	for _, loc := range dataNodes {
		_, found := loc.dataNode.DiskInfos[string(diskType)]
		if found {
			ret = append(ret, loc)
		}
	}
	return
}

func rotateDataNodes(dataNodes []location) []location {
	if len(dataNodes) > 0 {
		return append(dataNodes[1:], dataNodes[0])
	} else {
		return dataNodes
	}
}

func isOneOf(server string, locations []wdclient.Location) bool {
	for _, loc := range locations {
		if server == loc.Url {
			return true
		}
	}
	return false
}

func (c *commandVolumeTierMove) doVolumeTierMove(commandEnv *CommandEnv, writer io.Writer, vid needle.VolumeId, toDiskType types.DiskType, allLocations []location) (err error) {
	// find volume location
	locations, found := commandEnv.MasterClient.GetLocations(uint32(vid))
	if !found {
		return fmt.Errorf("volume %d not found", vid)
	}

	// find one server with the most empty volume slots with target disk type
	hasFoundTarget := false
	fn := capacityByFreeVolumeCount(toDiskType)
	for _, dst := range allLocations {
		if fn(dst.dataNode) > 0 && !hasFoundTarget {
			// ask the volume server to replicate the volume
			if isOneOf(dst.dataNode.Id, locations) {
				continue
			}
			var sourceVolumeServer pb.ServerAddress
			for _, loc := range locations {
				if loc.Url != dst.dataNode.Id {
					sourceVolumeServer = loc.ServerAddress()
				}
			}
			if sourceVolumeServer == "" {
				continue
			}
			hasFoundTarget = true

			// adjust volume count
			dst.dataNode.DiskInfos[string(toDiskType)].VolumeCount++

			destServerAddress := pb.NewServerAddressFromDataNode(dst.dataNode)
			c.queues[destServerAddress] <- volumeTierMoveJob{sourceVolumeServer, vid}
		}
	}

	if !hasFoundTarget {
		fmt.Fprintf(writer, "can not find disk type %s for volume %d\n", toDiskType.ReadableString(), vid)
	}

	return nil
}

func (c *commandVolumeTierMove) doMoveOneVolume(commandEnv *CommandEnv, writer io.Writer, vid needle.VolumeId, toDiskType types.DiskType, locations []wdclient.Location, sourceVolumeServer pb.ServerAddress, dst location) (err error) {

	// mark all replicas as read only
	if err = markVolumeReplicasWritable(commandEnv.option.GrpcDialOption, vid, locations, false); err != nil {
		return fmt.Errorf("mark volume %d as readonly on %s: %v", vid, locations[0].Url, err)
	}
	if err = LiveMoveVolume(commandEnv.option.GrpcDialOption, writer, vid, sourceVolumeServer, pb.NewServerAddressFromDataNode(dst.dataNode), 5*time.Second, toDiskType.ReadableString(), true); err != nil {

		// mark all replicas as writable
		if err = markVolumeReplicasWritable(commandEnv.option.GrpcDialOption, vid, locations, true); err != nil {
			glog.Errorf("mark volume %d as writable on %s: %v", vid, locations[0].Url, err)
		}

		return fmt.Errorf("move volume %d %s => %s : %v", vid, locations[0].Url, dst.dataNode.Id, err)
	}

	// remove the remaining replicas
	for _, loc := range locations {
		if loc.Url != dst.dataNode.Id && loc.ServerAddress() != sourceVolumeServer {
			if err = deleteVolume(commandEnv.option.GrpcDialOption, vid, loc.ServerAddress()); err != nil {
				fmt.Fprintf(writer, "failed to delete volume %d on %s: %v\n", vid, loc.Url, err)
			}
			// reduce volume count? Not really necessary since they are "more" full and will not be a candidate to move to
		}
	}
	return nil
}

func collectVolumeIdsForTierChange(commandEnv *CommandEnv, topologyInfo *master_pb.TopologyInfo, volumeSizeLimitMb uint64, sourceTier types.DiskType, collectionPattern string, fullPercentage float64, quietPeriod time.Duration) (vids []needle.VolumeId, err error) {

	quietSeconds := int64(quietPeriod / time.Second)
	nowUnixSeconds := time.Now().Unix()

	fmt.Printf("collect %s volumes quiet for: %d seconds\n", sourceTier, quietSeconds)

	vidMap := make(map[uint32]bool)
	eachDataNode(topologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		for _, diskInfo := range dn.DiskInfos {
			for _, v := range diskInfo.VolumeInfos {
				// check collection name pattern
				if collectionPattern != "" {
					matched, err := filepath.Match(collectionPattern, v.Collection)
					if err != nil {
						return
					}
					if !matched {
						continue
					}
				}

				if v.ModifiedAtSecond+quietSeconds < nowUnixSeconds && types.ToDiskType(v.DiskType) == sourceTier {
					if float64(v.Size) > fullPercentage/100*float64(volumeSizeLimitMb)*1024*1024 {
						vidMap[v.Id] = true
					}
				}
			}
		}
	})

	for vid := range vidMap {
		vids = append(vids, needle.VolumeId(vid))
	}

	return
}
<|endoftext|>


<commit_before>package consumer

import (
	"fmt"
	"gollum/shared"
	"math/rand"
	"sync"
	"time"
)

type Profiler struct {
	standardConsumer
	profileRuns int
	batches     int
	length      int
}

func init() {
	shared.Plugin.Register(Profiler{})
}

func (cons Profiler) Create(conf shared.PluginConfig, pool *shared.BytePool) (shared.Consumer, error) {
	err := cons.configureStandardConsumer(conf, pool)

	cons.profileRuns = conf.GetInt("Runs", 10000)
	cons.batches = conf.GetInt("Batches", 10)
	cons.length = conf.GetInt("Length", 256)

	return cons, err
}

var stringBase = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890 _.!?/&%$§'")

func (cons Profiler) profile() {

	randString := make([]rune, cons.length)
	for i := 0; i < cons.length; i++ {
		randString[i] = stringBase[rand.Intn(len(stringBase))]
	}

	var msg string
	for b := 0; b < cons.batches; b++ {

		start := time.Now()
		for i := 0; i < cons.profileRuns; i++ {
			msg = fmt.Sprintf("%d/%d %s", i, cons.profileRuns, string(randString))
			cons.postMessage(msg)
		}
		runTime := time.Since(start)

		shared.Log.Note(fmt.Sprintf(
			"Profile run #%d: %.4f sec = %4.f msg/sec",
			b, runTime.Seconds(),
			float64(cons.profileRuns)/runTime.Seconds()))
	}
}

func (cons Profiler) Consume(threads *sync.WaitGroup) {
	go cons.profile()

	// Wait for control statements

	for {
		command := <-cons.control
		if command == shared.ConsumerControlStop {
			return // ### return ###
		}
	}
}
<commit_msg>Additional profiling stats<commit_after>package consumer

import (
	"fmt"
	"gollum/shared"
	"math/rand"
	"sync"
	"time"
)

type Profiler struct {
	standardConsumer
	profileRuns int
	batches     int
	length      int
}

func init() {
	shared.Plugin.Register(Profiler{})
}

func (cons Profiler) Create(conf shared.PluginConfig, pool *shared.BytePool) (shared.Consumer, error) {
	err := cons.configureStandardConsumer(conf, pool)

	cons.profileRuns = conf.GetInt("Runs", 10000)
	cons.batches = conf.GetInt("Batches", 10)
	cons.length = conf.GetInt("Length", 256)

	return cons, err
}

var stringBase = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890 _.!?/&%$§'")

func (cons Profiler) profile() {

	randString := make([]rune, cons.length)
	for i := 0; i < cons.length; i++ {
		randString[i] = stringBase[rand.Intn(len(stringBase))]
	}

	testStart := time.Now()

	var msg string
	for b := 0; b < cons.batches; b++ {

		start := time.Now()
		for i := 0; i < cons.profileRuns; i++ {
			msg = fmt.Sprintf("%d/%d %s", i, cons.profileRuns, string(randString))
			cons.postMessage(msg)
		}

		runTime := time.Since(start)
		shared.Log.Note(fmt.Sprintf(
			"Profile run #%d: %.4f sec = %4.f msg/sec",
			b, runTime.Seconds(),
			float64(cons.profileRuns)/runTime.Seconds()))
	}

	runTime := time.Since(testStart)
	shared.Log.Note(fmt.Sprintf(
		"Total: %.4f sec = %4.f msg/sec",
		runTime.Seconds(),
		float64(cons.profileRuns*cons.batches)/runTime.Seconds()))
}

func (cons Profiler) Consume(threads *sync.WaitGroup) {
	go cons.profile()

	// Wait for control statements

	for {
		command := <-cons.control
		if command == shared.ConsumerControlStop {
			return // ### return ###
		}
	}
}
<|endoftext|>


<commit_before>package modata

import (
    "fmt"
    "sync"
    "web"
)

type ReplicationServer struct {
    mu sync.Mutex
    name string
    identifier string
    server *web.Server
}

// Requests to know the node
func (rs *ReplicationServer) WhoHasNode (c *web.Context, node string) {
    c.WriteString("Don't know who owns " + node)
}

func (rs *ReplicationServer) Identifier (c *web.Context) {
    c.WriteString(rs.identifier)
}

func StartReplicationServer(name string) *ReplicationServer{
    rs := new(ReplicationServer)

    // Node identifier for chord
    rs.identifier = MakeGUID()
    rs.name = name

    rs.server = web.NewServer()

    go func() {
        // Identifier for this node
        rs.server.Get("/id", func (c *web.Context) {
            rs.Identifier(c)
        })
        rs.server.Get("/whohas/(.*)", func (c *web.Context, node string) {
            rs.WhoHasNode(c, node)
        })

        rs.server.Run(name)
    }();

    return rs
}

<commit_msg>unbreaking build<commit_after>package modata

import (
    "fmt"
    "sync"
    "web"
)

type ReplicationServer struct {
    mu sync.Mutex
    name string
    identifier string
    server *web.Server
}

// Requests to know the node
func (rs *ReplicationServer) WhoHasNode (c *web.Context, node string) {
    c.WriteString("Don't know who owns " + node)
}

func (rs *ReplicationServer) Identifier (c *web.Context) {
    c.WriteString(rs.identifier)
}

func StartReplicationServer(name string) *ReplicationServer{
    rs := new(ReplicationServer)

    // Node identifier for chord
    rs.identifier = MakeGUID()
    rs.name = name

    rs.server = web.NewServer()

    go func() {
        // Identifier for this node
        rs.server.Get("/id", func (c *web.Context) {
            rs.Identifier(c)
        })
        rs.server.Get("/whohas/(.*)", func (c *web.Context, node string) {
            rs.WhoHasNode(c, node)
        })

        fmt.Printf("Listening on %v\n", name)
        rs.server.Run(name)
    }();

    return rs
}

<|endoftext|>


<commit_before>package chatblast

import (
	"log"
	"sync"
)

type Room struct {
	sync.RWMutex
	Name        string
	Id          string
	Private     bool
	Owner       *User
	Subscribers map[string]*User
	Channel     chan *Message
}

func (r *Room) GetSubscriber(userId string) (*User, bool) {
	r.RLock()
	defer r.RUnlock()
	user, ok := r.Subscribers[userId]
	return user, ok
}

func (r *Room) SetSubscriber(userId string, u *User) {
	r.Lock()
	defer r.Unlock()
	r.Subscribers[userId] = u
}

func (r *Room) RemoveSubscriber(userId string) {
	r.Lock()
	defer r.Unlock()
	delete(r.Subscribers, userId)
}

func NewRoom(name string, owner *User) *Room {
	newRoom := &Room{
		Name:        name,
		Id:          GUID(),
		Owner:       owner,
		Subscribers: map[string]*User{},
		Channel:     make(chan *Message),
	}
	go newRoom.Listen()
	log.Println("New room", name, "created")
	return newRoom
}

func (r *Room) Listen() {
	for incoming := range r.Channel {
		log.Println("Message sent in room", r.Name)
		for _, user := range r.Subscribers {
			user.Tell(incoming)
		}
	}
}

func (r *Room) Broadcast(msg *Message) {
	r.Channel <- msg
}

func (r *Room) Whisper(msg *Message, userId string) {
	whisperee, ok := r.Subscribers[userId]
	if ok {
		whisperee.Tell(msg)
	}
}

func (r *Room) RequestInvite(u *User) {
	msg := &Message{
		Cmd:    "inviterequest",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Owner.Channel <- msg
}

func (r *Room) Join(u *User) {
	r.SetSubscriber(u.Id, u)
	msg := &Message{
		Cmd:    "join",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Broadcast(msg)
}

func (r *Room) Leave(u *User) {
	r.RemoveSubscriber(u.Id)
	msg := &Message{
		Cmd:    "leave",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Broadcast(msg)
}

func (r *Room) CloseRoom(u *User) (closed bool, err error) {
	closed = false
	err = nil
	if r.Owner != nil && u.Id == r.Owner.Id {
		closed = true
		msg := &Message{
			Cmd:    "closing",
			RoomId: r.Id,
		}
		r.Broadcast(msg)
	}
	return closed, err

}
<commit_msg>Add room JSON serialization<commit_after>package chatblast

import (
	"log"
	"sync"
)

type Room struct {
	sync.RWMutex
	Name        string           `json:"name,omitempty"`
	Id          string           `json:"id,omitempty"`
	Private     bool             `json:"private,omitempty"`
	Owner       *User            `json:"owner,omitempty"`
	Subscribers map[string]*User `json:"subscribers,omitempty"`
	Channel     chan *Message    `json:"-"`
}

func (r *Room) GetSubscriber(userId string) (*User, bool) {
	r.RLock()
	defer r.RUnlock()
	user, ok := r.Subscribers[userId]
	return user, ok
}

func (r *Room) SetSubscriber(userId string, u *User) {
	r.Lock()
	defer r.Unlock()
	r.Subscribers[userId] = u
}

func (r *Room) RemoveSubscriber(userId string) {
	r.Lock()
	defer r.Unlock()
	delete(r.Subscribers, userId)
}

func NewRoom(name string, owner *User) *Room {
	newRoom := &Room{
		Name:        name,
		Id:          GUID(),
		Owner:       owner,
		Subscribers: map[string]*User{},
		Channel:     make(chan *Message),
	}
	go newRoom.Listen()
	log.Println("New room", name, "created")
	return newRoom
}

func (r *Room) Listen() {
	for incoming := range r.Channel {
		log.Println("Message sent in room", r.Name)
		for _, user := range r.Subscribers {
			user.Tell(incoming)
		}
	}
}

func (r *Room) Broadcast(msg *Message) {
	r.Channel <- msg
}

func (r *Room) Whisper(msg *Message, userId string) {
	if whisperee, ok := r.Subscribers[userId]; ok {
		whisperee.Tell(msg)
	}
}

func (r *Room) RequestInvite(u *User) {
	msg := &Message{
		Cmd:    "inviterequest",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Owner.Channel <- msg
}

func (r *Room) Join(u *User) {
	// TODO Add re-join check
	r.SetSubscriber(u.Id, u)
	msg := &Message{
		Cmd:    "join",
		Text:   u.Name,
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Broadcast(msg)
}

func (r *Room) Leave(u *User) {
	// TODO Add bad leave request check (i.e. user tries to leave room they
	// aren't subscribed to)
	r.RemoveSubscriber(u.Id)
	msg := &Message{
		Cmd:    "leave",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Broadcast(msg)
}

func (r *Room) CloseRoom(u *User) (closed bool, err error) {
	closed = false
	err = nil
	if r.Owner != nil && u.Id == r.Owner.Id {
		closed = true
		msg := &Message{
			Cmd:    "closing",
			RoomId: r.Id,
		}
		r.Broadcast(msg)
	}
	return closed, err

}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssh

import (
	"math/big"
	"math/rand"
	"reflect"
	"testing"
	"testing/quick"
)

var intLengthTests = []struct {
	val, length int
}{
	{0, 4 + 0},
	{1, 4 + 1},
	{127, 4 + 1},
	{128, 4 + 2},
	{-1, 4 + 1},
}

func TestIntLength(t *testing.T) {
	for _, test := range intLengthTests {
		v := new(big.Int).SetInt64(int64(test.val))
		length := intLength(v)
		if length != test.length {
			t.Errorf("For %d, got length %d but expected %d", test.val, length, test.length)
		}
	}
}

var messageTypes = []interface{}{
	&kexInitMsg{},
	&kexDHInitMsg{},
	&serviceRequestMsg{},
	&serviceAcceptMsg{},
	&userAuthRequestMsg{},
	&channelOpenMsg{},
	&channelOpenConfirmMsg{},
	&channelRequestMsg{},
	&channelRequestSuccessMsg{},
}

func TestMarshalUnmarshal(t *testing.T) {
	rand := rand.New(rand.NewSource(0))
	for i, iface := range messageTypes {
		ty := reflect.ValueOf(iface).Type()

		n := 100
		if testing.Short() {
			n = 5
		}
		for j := 0; j < n; j++ {
			v, ok := quick.Value(ty, rand)
			if !ok {
				t.Errorf("#%d: failed to create value", i)
				break
			}

			m1 := v.Elem().Interface()
			m2 := iface

			marshaled := marshal(msgIgnore, m1)
			if err := unmarshal(m2, marshaled, msgIgnore); err != nil {
				t.Errorf("#%d failed to unmarshal %#v: %s", i, m1, err)
				break
			}

			if !reflect.DeepEqual(v.Interface(), m2) {
				t.Errorf("#%d\ngot: %#v\nwant:%#v\n%x", i, m2, m1, marshaled)
				break
			}
		}
	}
}

func randomBytes(out []byte, rand *rand.Rand) {
	for i := 0; i < len(out); i++ {
		out[i] = byte(rand.Int31())
	}
}

func randomNameList(rand *rand.Rand) []string {
	ret := make([]string, rand.Int31()&15)
	for i := range ret {
		s := make([]byte, 1+(rand.Int31()&15))
		for j := range s {
			s[j] = 'a' + uint8(rand.Int31()&15)
		}
		ret[i] = string(s)
	}
	return ret
}

func randomInt(rand *rand.Rand) *big.Int {
	return new(big.Int).SetInt64(int64(int32(rand.Uint32())))
}

func (*kexInitMsg) Generate(rand *rand.Rand, size int) reflect.Value {
	ki := &kexInitMsg{}
	randomBytes(ki.Cookie[:], rand)
	ki.KexAlgos = randomNameList(rand)
	ki.ServerHostKeyAlgos = randomNameList(rand)
	ki.CiphersClientServer = randomNameList(rand)
	ki.CiphersServerClient = randomNameList(rand)
	ki.MACsClientServer = randomNameList(rand)
	ki.MACsServerClient = randomNameList(rand)
	ki.CompressionClientServer = randomNameList(rand)
	ki.CompressionServerClient = randomNameList(rand)
	ki.LanguagesClientServer = randomNameList(rand)
	ki.LanguagesServerClient = randomNameList(rand)
	if rand.Int31()&1 == 1 {
		ki.FirstKexFollows = true
	}
	return reflect.ValueOf(ki)
}

func (*kexDHInitMsg) Generate(rand *rand.Rand, size int) reflect.Value {
	dhi := &kexDHInitMsg{}
	dhi.X = randomInt(rand)
	return reflect.ValueOf(dhi)
}
<commit_msg>go.crypto/ssh: add benchmarks for marshal and unmarshal<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssh

import (
	"math/big"
	"math/rand"
	"reflect"
	"testing"
	"testing/quick"
)

var intLengthTests = []struct {
	val, length int
}{
	{0, 4 + 0},
	{1, 4 + 1},
	{127, 4 + 1},
	{128, 4 + 2},
	{-1, 4 + 1},
}

func TestIntLength(t *testing.T) {
	for _, test := range intLengthTests {
		v := new(big.Int).SetInt64(int64(test.val))
		length := intLength(v)
		if length != test.length {
			t.Errorf("For %d, got length %d but expected %d", test.val, length, test.length)
		}
	}
}

var messageTypes = []interface{}{
	&kexInitMsg{},
	&kexDHInitMsg{},
	&serviceRequestMsg{},
	&serviceAcceptMsg{},
	&userAuthRequestMsg{},
	&channelOpenMsg{},
	&channelOpenConfirmMsg{},
	&channelRequestMsg{},
	&channelRequestSuccessMsg{},
}

func TestMarshalUnmarshal(t *testing.T) {
	rand := rand.New(rand.NewSource(0))
	for i, iface := range messageTypes {
		ty := reflect.ValueOf(iface).Type()

		n := 100
		if testing.Short() {
			n = 5
		}
		for j := 0; j < n; j++ {
			v, ok := quick.Value(ty, rand)
			if !ok {
				t.Errorf("#%d: failed to create value", i)
				break
			}

			m1 := v.Elem().Interface()
			m2 := iface

			marshaled := marshal(msgIgnore, m1)
			if err := unmarshal(m2, marshaled, msgIgnore); err != nil {
				t.Errorf("#%d failed to unmarshal %#v: %s", i, m1, err)
				break
			}

			if !reflect.DeepEqual(v.Interface(), m2) {
				t.Errorf("#%d\ngot: %#v\nwant:%#v\n%x", i, m2, m1, marshaled)
				break
			}
		}
	}
}

func randomBytes(out []byte, rand *rand.Rand) {
	for i := 0; i < len(out); i++ {
		out[i] = byte(rand.Int31())
	}
}

func randomNameList(rand *rand.Rand) []string {
	ret := make([]string, rand.Int31()&15)
	for i := range ret {
		s := make([]byte, 1+(rand.Int31()&15))
		for j := range s {
			s[j] = 'a' + uint8(rand.Int31()&15)
		}
		ret[i] = string(s)
	}
	return ret
}

func randomInt(rand *rand.Rand) *big.Int {
	return new(big.Int).SetInt64(int64(int32(rand.Uint32())))
}

func (*kexInitMsg) Generate(rand *rand.Rand, size int) reflect.Value {
	ki := &kexInitMsg{}
	randomBytes(ki.Cookie[:], rand)
	ki.KexAlgos = randomNameList(rand)
	ki.ServerHostKeyAlgos = randomNameList(rand)
	ki.CiphersClientServer = randomNameList(rand)
	ki.CiphersServerClient = randomNameList(rand)
	ki.MACsClientServer = randomNameList(rand)
	ki.MACsServerClient = randomNameList(rand)
	ki.CompressionClientServer = randomNameList(rand)
	ki.CompressionServerClient = randomNameList(rand)
	ki.LanguagesClientServer = randomNameList(rand)
	ki.LanguagesServerClient = randomNameList(rand)
	if rand.Int31()&1 == 1 {
		ki.FirstKexFollows = true
	}
	return reflect.ValueOf(ki)
}

func (*kexDHInitMsg) Generate(rand *rand.Rand, size int) reflect.Value {
	dhi := &kexDHInitMsg{}
	dhi.X = randomInt(rand)
	return reflect.ValueOf(dhi)
}

var (
	_kexInitMsg   = new(kexInitMsg).Generate(rand.New(rand.NewSource(0)), 10).Elem().Interface()
	_kexDHInitMsg = new(kexDHInitMsg).Generate(rand.New(rand.NewSource(0)), 10).Elem().Interface()

	_kexInit   = marshal(msgKexInit, _kexInitMsg)
	_kexDHInit = marshal(msgKexDHInit, _kexDHInitMsg)
)

func BenchmarkMarshalKexInitMsg(b *testing.B) {
	for i := 0; i < b.N; i++ {
		marshal(msgKexInit, _kexInitMsg)
	}
}

func BenchmarkUnmarshalKexInitMsg(b *testing.B) {
	m := new(kexInitMsg)
	for i := 0; i < b.N; i++ {
		unmarshal(m, _kexInit, msgKexInit)
	}
}

func BenchmarkMarshalKexDHInitMsg(b *testing.B) {
	for i := 0; i < b.N; i++ {
		marshal(msgKexDHInit, _kexDHInitMsg)
	}
}

func BenchmarkUnmarshalKexDHInitMsg(b *testing.B) {
	m := new(kexDHInitMsg)
	for i := 0; i < b.N; i++ {
		unmarshal(m, _kexDHInit, msgKexDHInit)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
)

const (
	// colors (linux)
	BLUE   = "\033[36m"
	RED    = "\033[31m"
	ENDC   = "\033[0m"
	BEGIN  = BLUE + "# " + ENDC
	CMDOUT = RED + "> " + ENDC
)

type SassGit struct {
	file string
	msg  string
	push bool
}

// Parses arguments to create struct
func ParseArgs() (*SassGit, error) {
	s := new(SassGit)
	if len(os.Args) < 2 {
		return s, errors.New("not enough args")
	} else {
		for i := 1; i < len(os.Args); i++ {
			if os.Args[i][0] == '-' {
				// flag
				if os.Args[i] == "-p" {
					s.push = true
				}
				if os.Args[i] == "-m" {
					// begin quoted string
					i++
					s.msg = os.Args[i]
				}
			} else {
				s.file = os.Args[i]
			}
		}
		if s.msg == "" {
			s.msg = fmt.Sprintf("Compiled %[1]s.scss to %[1]s.css", s.file)
		}
		return s, nil
	}
}

func (s *SassGit) CmdExec() error {
	// sass compile to file
	PrintMsg(fmt.Sprintf("Compiling %[1]s.scss %[1]s.css\n", s.file))
	cmd := exec.Command("sass", s.file+".scss", s.file+".css")
	str, err := cmd.Output()
	if err != nil {
		return err
	} else if string(str) != "" {
		PrintExec(str)
	}

	// git add
	PrintMsg(fmt.Sprintf("Adding %[1]s.scss %[1]s.css\n", s.file))
	str, err = exec.Command("git", "add", s.file+".scss", s.file+".css").Output()
	if err != nil {
		return err
	} else if string(str) != "" {
		PrintExec(str)
	}

	// git commit
	PrintMsg(fmt.Sprintf("Committing \"%s\"\n", s.msg))
	str, err = exec.Command("git", "commit", "-m", s.msg).Output()
	if err != nil {
		return err
	} else if string(str) != "" {
		PrintExec(str)
	}

	if s.push == true {
		// git push
		PrintMsg("Pushing\n")
		str, err = exec.Command("git", "push").Output()
		if err != nil {
			return err
		} else if string(str) != "" {
			PrintExec(str)
		}
	}
	return nil
}

// used to print error messages
func PrintError(err error) {
	PrintMsg(fmt.Sprintf("Sorry, %s ("+RED+"!"+ENDC+")\n", err))
}

// used to print normal messages
func PrintMsg(s string) {
	fmt.Printf(BEGIN+"%s", s)
}

// used to print messages from executed code
func PrintExec(b []byte) {
	s := string(b)
	sub := ""
	for i := 0; i < len(s); i++ {
		if s[i] != '\n' {
			sub += string(s[i])
		} else {
			fmt.Printf(CMDOUT+"%s\n", sub)
			sub = ""
		}
	}
}

// calls functions
func main() {
	s, err := ParseArgs()
	if err != nil {
		PrintError(err)
	} else {
		err = s.CmdExec()
		if err != nil {
			PrintError(err)
		}
	}
}
<commit_msg>fixed it -> else if<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
)

const (
	// colors (linux)
	BLUE   = "\033[36m"
	RED    = "\033[31m"
	ENDC   = "\033[0m"
	BEGIN  = BLUE + "# " + ENDC
	CMDOUT = RED + "> " + ENDC
)

type SassGit struct {
	file string
	msg  string
	push bool
}

// Parses arguments to create struct
func ParseArgs() (*SassGit, error) {
	s := new(SassGit)
	if len(os.Args) < 2 {
		return s, errors.New("not enough args")
	} else {
		for i := 1; i < len(os.Args); i++ {
			if os.Args[i][0] == '-' {
				// flag
				if os.Args[i] == "-p" {
					s.push = true
				} else if os.Args[i] == "-m" {
					// begin quoted string
					i++
					s.msg = os.Args[i]
				}
			} else {
				s.file = os.Args[i]
			}
		}
		if s.msg == "" {
			s.msg = fmt.Sprintf("Compiled %[1]s.scss to %[1]s.css", s.file)
		}
		return s, nil
	}
}

func (s *SassGit) CmdExec() error {
	// sass compile to file
	PrintMsg(fmt.Sprintf("Compiling %[1]s.scss %[1]s.css\n", s.file))
	cmd := exec.Command("sass", s.file+".scss", s.file+".css")
	str, err := cmd.Output()
	if err != nil {
		return err
	} else if string(str) != "" {
		PrintExec(str)
	}

	// git add
	PrintMsg(fmt.Sprintf("Adding %[1]s.scss %[1]s.css\n", s.file))
	str, err = exec.Command("git", "add", s.file+".scss", s.file+".css").Output()
	if err != nil {
		return err
	} else if string(str) != "" {
		PrintExec(str)
	}

	// git commit
	PrintMsg(fmt.Sprintf("Committing \"%s\"\n", s.msg))
	str, err = exec.Command("git", "commit", "-m", s.msg).Output()
	if err != nil {
		return err
	} else if string(str) != "" {
		PrintExec(str)
	}

	if s.push == true {
		// git push
		PrintMsg("Pushing\n")
		str, err = exec.Command("git", "push").Output()
		if err != nil {
			return err
		} else if string(str) != "" {
			PrintExec(str)
		}
	}
	return nil
}

// used to print error messages
func PrintError(err error) {
	PrintMsg(fmt.Sprintf("Sorry, %s ("+RED+"!"+ENDC+")\n", err))
}

// used to print normal messages
func PrintMsg(s string) {
	fmt.Printf(BEGIN+"%s", s)
}

// used to print messages from executed code
func PrintExec(b []byte) {
	s := string(b)
	sub := ""
	for i := 0; i < len(s); i++ {
		if s[i] != '\n' {
			sub += string(s[i])
		} else {
			fmt.Printf(CMDOUT+"%s\n", sub)
			sub = ""
		}
	}
}

// calls functions
func main() {
	s, err := ParseArgs()
	if err != nil {
		PrintError(err)
	} else {
		err = s.CmdExec()
		if err != nil {
			PrintError(err)
		}
	}
}
<|endoftext|>


<commit_before>package victorops

import (
	"encoding/json"
	"time"
)

// Schedule current schedule
type Schedule struct {
	Oncall         string    `json:"oncall,omitempty"`
	OverrideOncall string    `json:"overrideoncall,omitempty"`
	PolicyType     string    `json:"policyType,omitempty"`
	RotationName   string    `json:"rotationName,omitempty"`
	ShiftName      string    `json:"shiftName,omitempty"`
	ShiftRoll      time.Time `json:"shiftRoll,omitempty"`
	Rolls          []Roll    `json:"rolls,omitempty"`
}

// Roll is shcedule roll
type Roll struct {
	Start  time.Time `json:"change,omitempty"`
	End    time.Time `json:"until,omitempty"`
	Oncall string    `json:"oncall,omitempty"`
	IsRoll bool      `json:"isroll,omitempty"`
}

// Override describes oncall overrides
type Override struct {
	Orig  string    `json:"orig,omitempty"`
	Over  string    `json:"over,omitempty"`
	Start time.Time `json:"start,omitempty"`
	End   time.Time `json:"end,omitempty"`
}

// TeamSchedule is api response struct
type TeamSchedule struct {
	Team      string     `json:"team"`
	Schedules []Schedule `json:"schedule"`
	Overrides []Override `json:"overrides"`
}

// TeamOverride is a pyaload for creating override request
type teamOverride struct {
	FromUser string `json:"fromUser"`
	ToUser   string `json:"toUser"`
}

// ResponseResult is result of an api request
type ResponseResult struct {
	Result string `json:"result"`
}

//UserOncallSchedule returns the on-call schedule for a user for all teams, including on-call overrides
func (client *Client) UserOncallSchedule(nick string) ([]TeamSchedule, error) {
	var schd []TeamSchedule
	err := client.sendRequest("GET", "api-public/v1/user/"+nick+"/oncall/schedule", nil, &schd)
	if err != nil {
		return schd, err
	}
	return schd, nil
}

// TeamOncallSchedule Get the on-call schedule for a user for all teams, including on-call overrides
func (client *Client) TeamOncallSchedule(team string) (TeamSchedule, error) {
	var schd TeamSchedule
	err := client.sendRequest("GET", "api-public/v1/team/"+team+"/oncall/schedule", nil, &schd)
	if err != nil {
		return schd, err
	}
	return schd, nil
}

// CreateOncallOverride replaces a currently on-call user on the team with another
func (client *Client) CreateOncallOverride(fromUser, toUser, team string) (ResponseResult, error) {
	var resp ResponseResult
	values := teamOverride{fromUser, toUser}
	jsonData, err := json.Marshal(values)
	if err != nil {
		return resp, err
	}
	err = client.sendRequest("PATCH", "api-public/v1/team/"+team+"/oncall/user", jsonData, &resp)
	return resp, err
}
<commit_msg>small coe optimization<commit_after>package victorops

import (
	"encoding/json"
	"time"
)

// Schedule current schedule
type Schedule struct {
	Oncall         string    `json:"oncall,omitempty"`
	OverrideOncall string    `json:"overrideoncall,omitempty"`
	PolicyType     string    `json:"policyType,omitempty"`
	RotationName   string    `json:"rotationName,omitempty"`
	ShiftName      string    `json:"shiftName,omitempty"`
	ShiftRoll      time.Time `json:"shiftRoll,omitempty"`
	Rolls          []Roll    `json:"rolls,omitempty"`
}

// Roll is shcedule roll
type Roll struct {
	Start  time.Time `json:"change,omitempty"`
	End    time.Time `json:"until,omitempty"`
	Oncall string    `json:"oncall,omitempty"`
	IsRoll bool      `json:"isroll,omitempty"`
}

// Override describes oncall overrides
type Override struct {
	Orig  string    `json:"orig,omitempty"`
	Over  string    `json:"over,omitempty"`
	Start time.Time `json:"start,omitempty"`
	End   time.Time `json:"end,omitempty"`
}

// TeamSchedule is api response struct
type TeamSchedule struct {
	Team      string     `json:"team"`
	Schedules []Schedule `json:"schedule"`
	Overrides []Override `json:"overrides"`
}

// TeamOverride is a pyaload for creating override request
type teamOverride struct {
	FromUser string `json:"fromUser"`
	ToUser   string `json:"toUser"`
}

// ResponseResult is result of an api request
type ResponseResult struct {
	Result string `json:"result"`
}

//UserOncallSchedule returns the on-call schedule for a user for all teams, including on-call overrides
func (client *Client) UserOncallSchedule(nick string) ([]TeamSchedule, error) {
	var schd []TeamSchedule
	err := client.sendRequest("GET", "api-public/v1/user/"+nick+"/oncall/schedule", nil, &schd)
	return schd, err
}

// TeamOncallSchedule Get the on-call schedule for a user for all teams, including on-call overrides
func (client *Client) TeamOncallSchedule(team string) (TeamSchedule, error) {
	var schd TeamSchedule
	err := client.sendRequest("GET", "api-public/v1/team/"+team+"/oncall/schedule", nil, &schd)
	return schd, err
}

// CreateOncallOverride replaces a currently on-call user on the team with another
func (client *Client) CreateOncallOverride(fromUser, toUser, team string) (ResponseResult, error) {
	var resp ResponseResult
	values := teamOverride{fromUser, toUser}
	jsonData, err := json.Marshal(values)
	if err != nil {
		return resp, err
	}
	err = client.sendRequest("PATCH", "api-public/v1/team/"+team+"/oncall/user", jsonData, &resp)
	return resp, err
}
<|endoftext|>


<commit_before>package libkbfs

import (
	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/engine"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/maxtaco/go-framed-msgpack-rpc/rpc2"
	"golang.org/x/net/context"
)

// KeybaseDaemonRPC implements the KeybaseDaemon interface using RPC
// calls.
type KeybaseDaemonRPC struct {
	identify keybase1.IdentifyInterface
	session  keybase1.SessionInterface
	favorite keybase1.FavoriteInterface
	user     keybase1.UserInterface
	log      logger.Logger
}

var _ KeybaseDaemon = KeybaseDaemonRPC{}

// NewKeybaseDaemonRPC makes a new KeybaseDaemonRPC that makes RPC
// calls using the socket of the given context.
func NewKeybaseDaemonRPC(ctx *libkb.GlobalContext, log logger.Logger) (KeybaseDaemonRPC, error) {
	_, xp, err := ctx.GetSocket()
	if err != nil {
		return KeybaseDaemonRPC{}, err
	}

	srv := rpc2.NewServer(xp, libkb.WrapError)

	protocols := []rpc2.Protocol{
		client.NewLogUIProtocol(),
		client.NewIdentifyUIProtocol(),
	}

	for _, p := range protocols {
		if err := srv.Register(p); err != nil {
			if _, ok := err.(rpc2.AlreadyRegisteredError); !ok {
				return KeybaseDaemonRPC{}, err
			}
		}
	}

	client := rpc2.NewClient(xp, libkb.UnwrapError)
	identifyClient := keybase1.IdentifyClient{Cli: client}
	sessionClient := keybase1.SessionClient{Cli: client}
	favoriteClient := keybase1.FavoriteClient{Cli: client}
	userClient := keybase1.UserClient{Cli: client}
	return newKeybaseDaemonRPCWithInterfaces(
		identifyClient, sessionClient, favoriteClient, userClient, log), nil
}

// For testing.
func newKeybaseDaemonRPCWithInterfaces(
	identify keybase1.IdentifyInterface,
	session keybase1.SessionInterface,
	favorite keybase1.FavoriteInterface,
	user keybase1.UserInterface,
	log logger.Logger,
) KeybaseDaemonRPC {
	return KeybaseDaemonRPC{
		identify: identify,
		session:  session,
		favorite: favorite,
		user:     user,
		log:      log,
	}
}

func (k KeybaseDaemonRPC) filterKeys(ctx context.Context, uid keybase1.UID, keys []keybase1.PublicKey) ([]VerifyingKey, []CryptPublicKey, error) {
	var verifyingKeys []VerifyingKey
	var cryptPublicKeys []CryptPublicKey
	for _, publicKey := range keys {
		if len(publicKey.PGPFingerprint) > 0 {
			continue
		}
		// Import the KID to validate it.
		key, err := libkb.ImportKeypairFromKID(publicKey.KID)
		if err != nil {
			return nil, nil, err
		}
		if publicKey.IsSibkey {
			k.log.CDebugf(
				ctx, "got verifying key %s for user %s",
				key.VerboseDescription(), uid)
			verifyingKeys = append(
				verifyingKeys, VerifyingKey{key.GetKID()})
		} else {
			k.log.CDebugf(
				ctx, "got crypt public key %s for user %s",
				key.VerboseDescription(), uid)
			cryptPublicKeys = append(
				cryptPublicKeys, CryptPublicKey{key.GetKID()})
		}
	}
	return verifyingKeys, cryptPublicKeys, nil
}

// Identify implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) Identify(ctx context.Context, assertion string) (
	UserInfo, error) {
	arg := engine.IDEngineArg{UserAssertion: assertion}.Export()
	var res keybase1.IdentifyRes
	f := func() error {
		var err error
		res, err = k.identify.Identify(arg)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return UserInfo{}, err
	}

	name := libkb.NewNormalizedUsername(res.User.Username)
	uid := keybase1.UID(res.User.Uid)

	verifyingKeys, cryptPublicKeys, err := k.filterKeys(ctx, uid, res.PublicKeys)
	if err != nil {
		return UserInfo{}, err
	}

	return UserInfo{
		Name:            name,
		UID:             uid,
		VerifyingKeys:   verifyingKeys,
		CryptPublicKeys: cryptPublicKeys,
	}, nil
}

// LoadUserPlusKeys implements the KeybaseDaemon interface for KeybaseDaemonRPC.
// If you have the UID for a user and don't require Identify, use
// this to get UserInfo.
func (k KeybaseDaemonRPC) LoadUserPlusKeys(ctx context.Context, uid keybase1.UID) (
	UserInfo, error) {
	arg := keybase1.LoadUserPlusKeysArg{Uid: uid, CacheOK: true}
	var res keybase1.UserPlusKeys
	f := func() error {
		var err error
		res, err = k.user.LoadUserPlusKeys(arg)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return UserInfo{}, err
	}

	verifyingKeys, cryptPublicKeys, err := k.filterKeys(ctx, uid, res.DeviceKeys)
	if err != nil {
		return UserInfo{}, err
	}

	return UserInfo{
		Name:            libkb.NormalizedUsername(res.Username),
		UID:             res.Uid,
		VerifyingKeys:   verifyingKeys,
		CryptPublicKeys: cryptPublicKeys,
	}, nil
}

// CurrentUID implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) CurrentUID(ctx context.Context, sessionID int) (
	keybase1.UID, error) {
	var currentUID keybase1.UID
	f := func() error {
		var err error
		currentUID, err = k.session.CurrentUID(sessionID)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return keybase1.UID(""), err
	}
	return currentUID, nil
}

// CurrentSession implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) CurrentSession(ctx context.Context, sessionID int) (
	SessionInfo, error) {
	var res keybase1.Session
	f := func() error {
		var err error
		res, err = k.session.CurrentSession(sessionID)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return SessionInfo{}, err
	}
	// Import the KID to validate it.
	deviceSubkey, err := libkb.ImportKeypairFromKID(res.DeviceSubkeyKid)
	if err != nil {
		return SessionInfo{}, err
	}
	k.log.CDebugf(ctx, "got device subkey %s",
		deviceSubkey.GetKID().ToShortIDString())
	cryptPublicKey := CryptPublicKey{deviceSubkey.GetKID()}
	return SessionInfo{
		UID:            keybase1.UID(res.Uid),
		Token:          res.Token,
		CryptPublicKey: cryptPublicKey,
	}, nil
}

// FavoriteAdd implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) FavoriteAdd(ctx context.Context, folder keybase1.Folder) error {
	f := func() error {
		return k.favorite.FavoriteAdd(keybase1.FavoriteAddArg{Folder: folder})
	}
	return runUnlessCanceled(ctx, f)
}

// FavoriteDelete implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) FavoriteDelete(ctx context.Context, folder keybase1.Folder) error {
	f := func() error {
		return k.favorite.FavoriteDelete(keybase1.FavoriteDeleteArg{Folder: folder})
	}
	return runUnlessCanceled(ctx, f)
}

// FavoriteList implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) FavoriteList(ctx context.Context, sessionID int) ([]keybase1.Folder, error) {
	var folders []keybase1.Folder
	f := func() error {
		var err error
		folders, err = k.favorite.FavoriteList(sessionID)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return nil, err
	}
	return folders, nil
}

// Shutdown implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) Shutdown() {
	// Nothing to do.
}
<commit_msg>Cleanup<commit_after>package libkbfs

import (
	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/engine"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/maxtaco/go-framed-msgpack-rpc/rpc2"
	"golang.org/x/net/context"
)

// KeybaseDaemonRPC implements the KeybaseDaemon interface using RPC
// calls.
type KeybaseDaemonRPC struct {
	identify keybase1.IdentifyInterface
	session  keybase1.SessionInterface
	favorite keybase1.FavoriteInterface
	user     keybase1.UserInterface
	log      logger.Logger
}

var _ KeybaseDaemon = KeybaseDaemonRPC{}

// NewKeybaseDaemonRPC makes a new KeybaseDaemonRPC that makes RPC
// calls using the socket of the given context.
func NewKeybaseDaemonRPC(ctx *libkb.GlobalContext, log logger.Logger) (KeybaseDaemonRPC, error) {
	_, xp, err := ctx.GetSocket()
	if err != nil {
		return KeybaseDaemonRPC{}, err
	}

	srv := rpc2.NewServer(xp, libkb.WrapError)

	protocols := []rpc2.Protocol{
		client.NewLogUIProtocol(),
		client.NewIdentifyUIProtocol(),
	}

	for _, p := range protocols {
		if err := srv.Register(p); err != nil {
			if _, ok := err.(rpc2.AlreadyRegisteredError); !ok {
				return KeybaseDaemonRPC{}, err
			}
		}
	}

	client := rpc2.NewClient(xp, libkb.UnwrapError)
	identifyClient := keybase1.IdentifyClient{Cli: client}
	sessionClient := keybase1.SessionClient{Cli: client}
	favoriteClient := keybase1.FavoriteClient{Cli: client}
	userClient := keybase1.UserClient{Cli: client}
	return newKeybaseDaemonRPCWithInterfaces(
		identifyClient, sessionClient, favoriteClient, userClient, log), nil
}

// For testing.
func newKeybaseDaemonRPCWithInterfaces(
	identify keybase1.IdentifyInterface,
	session keybase1.SessionInterface,
	favorite keybase1.FavoriteInterface,
	user keybase1.UserInterface,
	log logger.Logger,
) KeybaseDaemonRPC {
	return KeybaseDaemonRPC{
		identify: identify,
		session:  session,
		favorite: favorite,
		user:     user,
		log:      log,
	}
}

func (k KeybaseDaemonRPC) filterKeys(ctx context.Context, uid keybase1.UID, keys []keybase1.PublicKey) ([]VerifyingKey, []CryptPublicKey, error) {
	var verifyingKeys []VerifyingKey
	var cryptPublicKeys []CryptPublicKey
	for _, publicKey := range keys {
		if len(publicKey.PGPFingerprint) > 0 {
			continue
		}
		// Import the KID to validate it.
		key, err := libkb.ImportKeypairFromKID(publicKey.KID)
		if err != nil {
			return nil, nil, err
		}
		if publicKey.IsSibkey {
			k.log.CDebugf(
				ctx, "got verifying key %s for user %s",
				key.VerboseDescription(), uid)
			verifyingKeys = append(
				verifyingKeys, VerifyingKey{key.GetKID()})
		} else {
			k.log.CDebugf(
				ctx, "got crypt public key %s for user %s",
				key.VerboseDescription(), uid)
			cryptPublicKeys = append(
				cryptPublicKeys, CryptPublicKey{key.GetKID()})
		}
	}
	return verifyingKeys, cryptPublicKeys, nil
}

// Identify implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) Identify(ctx context.Context, assertion string) (
	UserInfo, error) {
	arg := engine.IDEngineArg{UserAssertion: assertion}.Export()
	var res keybase1.IdentifyRes
	f := func() error {
		var err error
		res, err = k.identify.Identify(arg)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return UserInfo{}, err
	}

	name := libkb.NewNormalizedUsername(res.User.Username)
	uid := keybase1.UID(res.User.Uid)

	verifyingKeys, cryptPublicKeys, err := k.filterKeys(ctx, uid, res.PublicKeys)
	if err != nil {
		return UserInfo{}, err
	}

	return UserInfo{
		Name:            name,
		UID:             uid,
		VerifyingKeys:   verifyingKeys,
		CryptPublicKeys: cryptPublicKeys,
	}, nil
}

// LoadUserPlusKeys implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) LoadUserPlusKeys(ctx context.Context, uid keybase1.UID) (
	UserInfo, error) {
	arg := keybase1.LoadUserPlusKeysArg{Uid: uid, CacheOK: true}
	var res keybase1.UserPlusKeys
	f := func() error {
		var err error
		res, err = k.user.LoadUserPlusKeys(arg)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return UserInfo{}, err
	}

	verifyingKeys, cryptPublicKeys, err := k.filterKeys(ctx, uid, res.DeviceKeys)
	if err != nil {
		return UserInfo{}, err
	}

	return UserInfo{
		Name:            libkb.NormalizedUsername(res.Username),
		UID:             res.Uid,
		VerifyingKeys:   verifyingKeys,
		CryptPublicKeys: cryptPublicKeys,
	}, nil
}

// CurrentUID implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) CurrentUID(ctx context.Context, sessionID int) (
	keybase1.UID, error) {
	var currentUID keybase1.UID
	f := func() error {
		var err error
		currentUID, err = k.session.CurrentUID(sessionID)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return keybase1.UID(""), err
	}
	return currentUID, nil
}

// CurrentSession implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) CurrentSession(ctx context.Context, sessionID int) (
	SessionInfo, error) {
	var res keybase1.Session
	f := func() error {
		var err error
		res, err = k.session.CurrentSession(sessionID)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return SessionInfo{}, err
	}
	// Import the KID to validate it.
	deviceSubkey, err := libkb.ImportKeypairFromKID(res.DeviceSubkeyKid)
	if err != nil {
		return SessionInfo{}, err
	}
	k.log.CDebugf(ctx, "got device subkey %s",
		deviceSubkey.GetKID().ToShortIDString())
	cryptPublicKey := CryptPublicKey{deviceSubkey.GetKID()}
	return SessionInfo{
		UID:            keybase1.UID(res.Uid),
		Token:          res.Token,
		CryptPublicKey: cryptPublicKey,
	}, nil
}

// FavoriteAdd implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) FavoriteAdd(ctx context.Context, folder keybase1.Folder) error {
	f := func() error {
		return k.favorite.FavoriteAdd(keybase1.FavoriteAddArg{Folder: folder})
	}
	return runUnlessCanceled(ctx, f)
}

// FavoriteDelete implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) FavoriteDelete(ctx context.Context, folder keybase1.Folder) error {
	f := func() error {
		return k.favorite.FavoriteDelete(keybase1.FavoriteDeleteArg{Folder: folder})
	}
	return runUnlessCanceled(ctx, f)
}

// FavoriteList implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) FavoriteList(ctx context.Context, sessionID int) ([]keybase1.Folder, error) {
	var folders []keybase1.Folder
	f := func() error {
		var err error
		folders, err = k.favorite.FavoriteList(sessionID)
		return err
	}
	if err := runUnlessCanceled(ctx, f); err != nil {
		return nil, err
	}
	return folders, nil
}

// Shutdown implements the KeybaseDaemon interface for KeybaseDaemonRPC.
func (k KeybaseDaemonRPC) Shutdown() {
	// Nothing to do.
}
<|endoftext|>


<commit_before>package libproxy

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
	"sync"

	"github.com/Sirupsen/logrus"
)

type udpListener interface {
  ReadFromUDP(b []byte) (int, *net.UDPAddr, error)
	WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
	Close() error
}

type udpEncapsulator struct {
	conn     *net.Conn
	listener net.Listener
	m        *sync.Mutex
	r        *sync.Mutex
	w        *sync.Mutex
}

func (u *udpEncapsulator) getConn() (net.Conn, error) {
	u.m.Lock()
	defer u.m.Unlock()
	if u.conn != nil {
		return *u.conn, nil
	}
	conn, err := u.listener.Accept()
	if err != nil {
		logrus.Printf("Failed to accept connection: %#v", err)
		return nil, err
	}
	u.conn = &conn
	return conn, nil
}

func (u *udpEncapsulator) ReadFromUDP(b []byte) (int, *net.UDPAddr, error) {
	conn, err := u.getConn()
	if err != nil {
		return 0, nil, err
	}
	u.r.Lock()
	defer u.r.Unlock()
	datagram := &udpDatagram{payload: b}
	err = datagram.Unmarshal(conn)
	if err != nil {
		return 0, nil, err
	}
	return len(datagram.payload), &net.UDPAddr{IP: *datagram.IP, Port: int(datagram.Port), Zone: datagram.Zone}, nil
}

func (u *udpEncapsulator) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) {
	conn, err := u.getConn()
	if err != nil {
		return 0, err
	}
	u.w.Lock()
	defer u.w.Unlock()
	datagram := &udpDatagram{payload: b, IP: &addr.IP, Port: uint16(addr.Port), Zone: addr.Zone}
	return len(b), datagram.Marshal(conn)
}

func (u *udpEncapsulator) Close() error {
	if u.conn != nil {
		conn := *u.conn
		conn.Close()
	}
	u.listener.Close()
	return nil
}

func NewUDPListener(listener net.Listener) udpListener {
	var m sync.Mutex;
	return &udpEncapsulator{
		conn: nil,
		listener: listener,
		m: &m,
	}
}

type udpDatagram struct {
	payload []byte
	IP      *net.IP
	Port    uint16
	Zone    string
}

func (u *udpDatagram) Marshal(conn net.Conn) error {
	var length uint16
	length = uint16(len(*u.IP))
	if err := binary.Write(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	if err := binary.Write(conn, binary.LittleEndian, &u.IP); err != nil {
		return err
	}
	if err := binary.Write(conn, binary.LittleEndian, &u.Port); err != nil {
		return err
	}
	length = uint16(len(u.Zone))
	if err := binary.Write(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	if err := binary.Write(conn, binary.LittleEndian, &u.Zone); err != nil {
		return nil
	}
	length = uint16(len(u.payload))
	if err := binary.Write(conn, binary.LittleEndian, &length); err != nil {
		return nil
	}
	payload := bytes.NewBuffer(u.payload)
	_, err := io.Copy(conn, payload)
	if err != nil {
		return err
	}
	return nil
}

func (u *udpDatagram) Unmarshal(conn net.Conn) error {
	var length uint16
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	var IP net.IP
	IP = make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &IP); err != nil {
		return err
	}
	u.IP = &IP
	if err := binary.Read(conn, binary.LittleEndian, &u.Port); err != nil {
		return err
	}
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	Zone := make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &Zone); err != nil {
		return err
	}
	u.Zone = string(Zone)
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	u.payload = make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &u.payload); err != nil {
		return err
	}
	return nil
}
<commit_msg>proxy: add an initial length to the marshalled frames<commit_after>package libproxy

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
	"sync"

	"github.com/Sirupsen/logrus"
)

type udpListener interface {
  ReadFromUDP(b []byte) (int, *net.UDPAddr, error)
	WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
	Close() error
}

type udpEncapsulator struct {
	conn     *net.Conn
	listener net.Listener
	m        *sync.Mutex
	r        *sync.Mutex
	w        *sync.Mutex
}

func (u *udpEncapsulator) getConn() (net.Conn, error) {
	u.m.Lock()
	defer u.m.Unlock()
	if u.conn != nil {
		return *u.conn, nil
	}
	conn, err := u.listener.Accept()
	if err != nil {
		logrus.Printf("Failed to accept connection: %#v", err)
		return nil, err
	}
	u.conn = &conn
	return conn, nil
}

func (u *udpEncapsulator) ReadFromUDP(b []byte) (int, *net.UDPAddr, error) {
	conn, err := u.getConn()
	if err != nil {
		return 0, nil, err
	}
	u.r.Lock()
	defer u.r.Unlock()
	datagram := &udpDatagram{payload: b}
	err = datagram.Unmarshal(conn)
	if err != nil {
		return 0, nil, err
	}
	return len(datagram.payload), &net.UDPAddr{IP: *datagram.IP, Port: int(datagram.Port), Zone: datagram.Zone}, nil
}

func (u *udpEncapsulator) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) {
	conn, err := u.getConn()
	if err != nil {
		return 0, err
	}
	u.w.Lock()
	defer u.w.Unlock()
	datagram := &udpDatagram{payload: b, IP: &addr.IP, Port: uint16(addr.Port), Zone: addr.Zone}
	return len(b), datagram.Marshal(conn)
}

func (u *udpEncapsulator) Close() error {
	if u.conn != nil {
		conn := *u.conn
		conn.Close()
	}
	u.listener.Close()
	return nil
}

func NewUDPListener(listener net.Listener) udpListener {
	var m sync.Mutex;
	return &udpEncapsulator{
		conn: nil,
		listener: listener,
		m: &m,
	}
}

type udpDatagram struct {
	payload []byte
	IP      *net.IP
	Port    uint16
	Zone    string
}

func (u *udpDatagram) Marshal(conn net.Conn) error {
	// marshal the variable length header to a temporary buffer
	var header bytes.Buffer
	var length uint16
	length = uint16(len(*u.IP))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return err
	}
	if err := binary.Write(&header, binary.LittleEndian, &u.IP); err != nil {
		return err
	}
	if err := binary.Write(&header, binary.LittleEndian, &u.Port); err != nil {
		return err
	}
	length = uint16(len(u.Zone))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return err
	}
	if err := binary.Write(&header, binary.LittleEndian, &u.Zone); err != nil {
		return nil
	}
	length = uint16(len(u.payload))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return nil
	}
	length = uint16(header.Len() + len(u.payload))
	if err := binary.Write(conn, binary.LittleEndian, &length); err != nil {
		return nil
	}
	_, err := io.Copy(conn, &header)
	if err != nil {
		return err
	}
	payload := bytes.NewBuffer(u.payload)
	_, err = io.Copy(conn, payload)
	if err != nil {
		return err
	}
	return nil
}

func (u *udpDatagram) Unmarshal(conn net.Conn) error {
	var length uint16
	// frame length
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	var IP net.IP
	IP = make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &IP); err != nil {
		return err
	}
	u.IP = &IP
	if err := binary.Read(conn, binary.LittleEndian, &u.Port); err != nil {
		return err
	}
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	Zone := make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &Zone); err != nil {
		return err
	}
	u.Zone = string(Zone)
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return err
	}
	u.payload = make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &u.payload); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 The Author - Unknown. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Help linls:
// https://codereview.appspot.com/7305072
package gocompresser

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"log"
	"os"
	"path"
)

// main functions shows how to TarGz a directory and
// UnTarGz a file
//func main() {
//	targetFilePath := "testdata.tar.gz"
//	srcDirPath := "testdata"
//	TarGz(srcDirPath, targetFilePath)
//	UnTarGz(targetFilePath, srcDirPath+"_temp")
//}

// Gzip and tar from source directory to destination file
// you need check file exist before you call this function
func TarGz(srcDirPath string, destFilePath string) {
	fw, err := os.Create(destFilePath)
	handleError(err)
	defer fw.Close()

	// Gzip writer
	gw := gzip.NewWriter(fw)
	defer gw.Close()

	// Tar writer
	tw := tar.NewWriter(gw)
	defer tw.Close()

	// Check if it's a file a directory
	f, err := os.Open(srcDirPath)
	handleError(err)
	fi, err := f.Stat()
	handleError(err)
	switch {
	case fi.IsDir():
		// handle source directory
		fmt.Println("Cerating tar.gz from directory...")
		tarGzDir(srcDirPath, path.Base(srcDirPath), tw)
		fallthrough
	case 0 == (os.ModeType & fi.Mode()):
		// handle file directly
		fmt.Println("Cerating tar.gz from file...")
		tarGzFile(srcDirPath, "", tw, fi)
	}
	fmt.Println("Well done!")
}

// Deal with directories
// if find files, handle them with HandleFile
// Every recurrence append the base path to the recPath
// recPath is the path inside of tar.gz
func tarGzDir(srcDirPath string, recPath string, tw *tar.Writer) {
	// Open source diretory
	dir, err := os.Open(srcDirPath)
	handleError(err)
	defer dir.Close()

	// Get file info slice
	fis, err := dir.Readdir(0)
	handleError(err)
	for _, fi := range fis {
		// Append path
		curPath := srcDirPath + "/" + fi.Name()
		// Check it is directory or file
		if fi.IsDir() {
			// Directory
			// (Directory won't add unitl all subfiles are added)
			fmt.Printf("Adding path...%s\n", curPath)
			tarGzDir(curPath, recPath+"/"+fi.Name(), tw)
		} else {
			// File
			fmt.Printf("Adding file...%s\n", curPath)
		}

		tarGzFile(curPath, recPath+"/"+fi.Name(), tw, fi)
	}
}

// Deal with files
func tarGzFile(srcFile string, recPath string, tw *tar.Writer, fi os.FileInfo) {
	if fi.IsDir() {
		// Create tar header
		hdr := new(tar.Header)
		// if last character of header name is '/' it also can be directory
		// but if you don't set Typeflag, error will occur when you untargz
		hdr.Name = recPath + "/"
		hdr.Typeflag = tar.TypeDir
		hdr.Size = 0
		//hdr.Mode = 0755 | c_ISDIR
		hdr.Mode = int64(fi.Mode())
		hdr.ModTime = fi.ModTime()

		// Write hander
		err := tw.WriteHeader(hdr)
		handleError(err)
	} else {
		// File reader
		fr, err := os.Open(srcFile)
		handleError(err)
		defer fr.Close()

		// Create tar header
		hdr := new(tar.Header)
		hdr.Name = recPath
		hdr.Size = fi.Size()
		hdr.Mode = int64(fi.Mode())
		hdr.ModTime = fi.ModTime()

		// Write hander
		err = tw.WriteHeader(hdr)
		handleError(err)

		// Write file data
		_, err = io.Copy(tw, fr)
		handleError(err)
	}
}

// Ungzip and untar from source file to destination directory
// you need check file exist before you call this function
func UnTarGz(srcFilePath string, destDirPath string) {
	fmt.Println("UnTarGzing " + srcFilePath + "...")
	// Create destination directory
	os.Mkdir(destDirPath, os.ModePerm)

	fr, err := os.Open(srcFilePath)
	handleError(err)
	defer fr.Close()

	// Gzip reader
	gr, err := gzip.NewReader(fr)

	// Tar reader
	tr := tar.NewReader(gr)

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// End of tar archive
			break
		}
		//handleError(err)
		fmt.Println("UnTarGzing file..." + hdr.Name)
		// Check if it is diretory or file
		if hdr.Typeflag != tar.TypeDir {
			// Get files from archive
			// Create diretory before create file
			os.MkdirAll(destDirPath+"/"+path.Dir(hdr.Name), os.ModePerm)
			// Write data to file
			fw, _ := os.Create(destDirPath + "/" + hdr.Name)
			handleError(err)
			_, err = io.Copy(fw, tr)
			handleError(err)
		}
	}
}

func handleError(e error) {
	if e != nil {
		log.Fatal(e)
	}

}
<commit_msg>go-tar.gz:bugs fix<commit_after>// Copyright 2013 The Author - Unknown. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Help linls:
// https://codereview.appspot.com/7305072
package gocompresser

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"log"
	"os"
	"path"
)

// main functions shows how to TarGz a directory and
// UnTarGz a file
//func main() {
//	targetFilePath := "testdata.tar.gz"
//	srcDirPath := "testdata"
//	TarGz(srcDirPath, targetFilePath)
//	UnTarGz(targetFilePath, srcDirPath+"_temp")
//}

// Gzip and tar from source directory to destination file
// you need check file exist before you call this function
func TarGz(srcDirPath string, destFilePath string) {
	fw, err := os.Create(destFilePath)
	handleError(err)
	defer fw.Close()

	// Gzip writer
	gw := gzip.NewWriter(fw)
	defer gw.Close()

	// Tar writer
	tw := tar.NewWriter(gw)
	defer tw.Close()

	// Check if it's a file a directory
	f, err := os.Open(srcDirPath)
	handleError(err)
	fi, err := f.Stat()
	handleError(err)
	switch {
	case fi.IsDir():
		// handle source directory
		fmt.Println("Cerating tar.gz from directory...")
		tarGzDir(srcDirPath, path.Base(srcDirPath), tw)
		fallthrough
	case 0 == (os.ModeType & fi.Mode()):
		// handle file directly
		fmt.Println("Cerating tar.gz from " + fi.Name() + "...")
		tarGzFile(srcDirPath, fi.Name(), tw, fi)
	}
	fmt.Println("Well done!")
}

// Deal with directories
// if find files, handle them with HandleFile
// Every recurrence append the base path to the recPath
// recPath is the path inside of tar.gz
func tarGzDir(srcDirPath string, recPath string, tw *tar.Writer) {
	// Open source diretory
	dir, err := os.Open(srcDirPath)
	handleError(err)
	defer dir.Close()

	// Get file info slice
	fis, err := dir.Readdir(0)
	handleError(err)
	for _, fi := range fis {
		// Append path
		curPath := srcDirPath + "/" + fi.Name()
		// Check it is directory or file
		if fi.IsDir() {
			// Directory
			// (Directory won't add unitl all subfiles are added)
			fmt.Printf("Adding path...%s\n", curPath)
			tarGzDir(curPath, recPath+"/"+fi.Name(), tw)
		} else {
			// File
			fmt.Printf("Adding file...%s\n", curPath)
		}

		tarGzFile(curPath, recPath+"/"+fi.Name(), tw, fi)
	}
}

// Deal with files
func tarGzFile(srcFile string, recPath string, tw *tar.Writer, fi os.FileInfo) {
	if fi.IsDir() {
		// Create tar header
		hdr := new(tar.Header)
		// if last character of header name is '/' it also can be directory
		// but if you don't set Typeflag, error will occur when you untargz
		hdr.Name = recPath + "/"
		hdr.Typeflag = tar.TypeDir
		hdr.Size = 0
		//hdr.Mode = 0755 | c_ISDIR
		hdr.Mode = int64(fi.Mode())
		hdr.ModTime = fi.ModTime()

		// Write hander
		err := tw.WriteHeader(hdr)
		handleError(err)
	} else {
		// File reader
		fr, err := os.Open(srcFile)
		handleError(err)
		defer fr.Close()

		// Create tar header
		hdr := new(tar.Header)
		hdr.Name = recPath
		hdr.Size = fi.Size()
		hdr.Mode = int64(fi.Mode())
		hdr.ModTime = fi.ModTime()

		// Write hander
		err = tw.WriteHeader(hdr)
		handleError(err)

		// Write file data
		_, err = io.Copy(tw, fr)
		handleError(err)
	}
}

// Ungzip and untar from source file to destination directory
// you need check file exist before you call this function
func UnTarGz(srcFilePath string, destDirPath string) {
	fmt.Println("UnTarGzing " + srcFilePath + "...")
	// Create destination directory
	os.Mkdir(destDirPath, os.ModePerm)

	fr, err := os.Open(srcFilePath)
	handleError(err)
	defer fr.Close()

	// Gzip reader
	gr, err := gzip.NewReader(fr)

	// Tar reader
	tr := tar.NewReader(gr)

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// End of tar archive
			break
		}
		//handleError(err)
		fmt.Println("UnTarGzing file..." + hdr.Name)
		// Check if it is diretory or file
		if hdr.Typeflag != tar.TypeDir {
			// Get files from archive
			// Create diretory before create file
			os.MkdirAll(destDirPath+"/"+path.Dir(hdr.Name), os.ModePerm)
			// Write data to file
			fw, _ := os.Create(destDirPath + "/" + hdr.Name)
			handleError(err)
			_, err = io.Copy(fw, tr)
			handleError(err)
		}
	}
}

func handleError(e error) {
	if e != nil {
		log.Fatal(e)
	}

}
<|endoftext|>


<commit_before>/*
   Copyright 2014 Outbrain Inc.

   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 db

import (
	"crypto/tls"
	"fmt"
	"strings"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/openark/golib/log"
	"github.com/openark/golib/sqlutils"
	"github.com/patrickmn/go-cache"
	"github.com/rcrowley/go-metrics"

	"github.com/github/orchestrator/go/config"
	"github.com/github/orchestrator/go/ssl"
)

const Error3159 = "Error 3159:"
const Error1045 = "Access denied for user"

// Track if a TLS has already been configured for topology
var topologyTLSConfigured bool = false

// Track if a TLS has already been configured for Orchestrator
var orchestratorTLSConfigured bool = false

var requireTLSCache *cache.Cache = cache.New(time.Duration(config.Config.TLSCacheTTLFactor*config.Config.InstancePollSeconds)*time.Second, time.Second)

var readInstanceTLSCounter = metrics.NewCounter()
var writeInstanceTLSCounter = metrics.NewCounter()
var readInstanceTLSCacheCounter = metrics.NewCounter()
var writeInstanceTLSCacheCounter = metrics.NewCounter()

func init() {
	metrics.Register("instance_tls.read", readInstanceTLSCounter)
	metrics.Register("instance_tls.write", writeInstanceTLSCounter)
	metrics.Register("instance_tls.read_cache", readInstanceTLSCacheCounter)
	metrics.Register("instance_tls.write_cache", writeInstanceTLSCacheCounter)
}

func requiresTLS(host string, port int, mysql_uri string) bool {
	cacheKey := fmt.Sprintf("%s:%d", host, port)

	if value, found := requireTLSCache.Get(cacheKey); found {
		readInstanceTLSCacheCounter.Inc(1)
		return value.(bool)
	}

	required := false
	db, _, _ := sqlutils.GetDB(mysql_uri)
	if err := db.Ping(); err != nil && (strings.Contains(err.Error(), Error3159) || strings.Contains(err.Error(), Error1045)) {
		required = true
	}

	query := `
			insert into
				database_instance_tls (
					hostname, port, required
				) values (
					?, ?, ?
				)
				on duplicate key update
					required=values(required)
				`
	if _, err := ExecOrchestrator(query, host, port, required); err != nil {
		log.Errore(err)
	}
	writeInstanceTLSCounter.Inc(1)

	requireTLSCache.Set(cacheKey, required, cache.DefaultExpiration)
	writeInstanceTLSCacheCounter.Inc(1)

	return required
}

// Create a TLS configuration from the config supplied CA, Certificate, and Private key.
// Register the TLS config with the mysql drivers as the "topology" config
// Modify the supplied URI to call the TLS config
func SetupMySQLTopologyTLS(uri string) (string, error) {
	if !topologyTLSConfigured {
		tlsConfig, err := ssl.NewTLSConfig(config.Config.MySQLTopologySSLCAFile, !config.Config.MySQLTopologySSLSkipVerify)
		// Drop to TLS 1.0 for talking to MySQL
		tlsConfig.MinVersion = tls.VersionTLS10
		if err != nil {
			return "", log.Errorf("Can't create TLS configuration for Topology connection %s: %s", uri, err)
		}
		tlsConfig.InsecureSkipVerify = config.Config.MySQLTopologySSLSkipVerify

		if config.Config.MySQLTopologyUseMutualTLS ||
			config.Config.MySQLTopologySSLCertFile != "" ||
			config.Config.MySQLTopologySSLPrivateKeyFile != "" {
			if err = ssl.AppendKeyPair(tlsConfig, config.Config.MySQLTopologySSLCertFile, config.Config.MySQLTopologySSLPrivateKeyFile); err != nil {
				return "", log.Errorf("Can't setup TLS key pairs for %s: %s", uri, err)
			}
		}
		if err = mysql.RegisterTLSConfig("topology", tlsConfig); err != nil {
			return "", log.Errorf("Can't register mysql TLS config for topology: %s", err)
		}
		topologyTLSConfigured = true
	}
	return fmt.Sprintf("%s&tls=topology", uri), nil
}

// Create a TLS configuration from the config supplied CA, Certificate, and Private key.
// Register the TLS config with the mysql drivers as the "orchestrator" config
// Modify the supplied URI to call the TLS config
func SetupMySQLOrchestratorTLS(uri string) (string, error) {
	if !orchestratorTLSConfigured {
		tlsConfig, err := ssl.NewTLSConfig(config.Config.MySQLOrchestratorSSLCAFile, true)
		// Drop to TLS 1.0 for talking to MySQL
		tlsConfig.MinVersion = tls.VersionTLS10
		if err != nil {
			return "", log.Fatalf("Can't create TLS configuration for Orchestrator connection %s: %s", uri, err)
		}
		tlsConfig.InsecureSkipVerify = config.Config.MySQLOrchestratorSSLSkipVerify
		if err = ssl.AppendKeyPair(tlsConfig, config.Config.MySQLOrchestratorSSLCertFile, config.Config.MySQLOrchestratorSSLPrivateKeyFile); err != nil {
			return "", log.Fatalf("Can't setup TLS key pairs for %s: %s", uri, err)
		}
		if err = mysql.RegisterTLSConfig("orchestrator", tlsConfig); err != nil {
			return "", log.Fatalf("Can't register mysql TLS config for orchestrator: %s", err)
		}
		orchestratorTLSConfigured = true
	}
	return fmt.Sprintf("%s&tls=orchestrator", uri), nil
}
<commit_msg>MySQLOrchestratorSSLSkipVerify to apply on backend TLS config<commit_after>/*
   Copyright 2014 Outbrain Inc.

   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 db

import (
	"crypto/tls"
	"fmt"
	"strings"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/openark/golib/log"
	"github.com/openark/golib/sqlutils"
	"github.com/patrickmn/go-cache"
	"github.com/rcrowley/go-metrics"

	"github.com/github/orchestrator/go/config"
	"github.com/github/orchestrator/go/ssl"
)

const Error3159 = "Error 3159:"
const Error1045 = "Access denied for user"

// Track if a TLS has already been configured for topology
var topologyTLSConfigured bool = false

// Track if a TLS has already been configured for Orchestrator
var orchestratorTLSConfigured bool = false

var requireTLSCache *cache.Cache = cache.New(time.Duration(config.Config.TLSCacheTTLFactor*config.Config.InstancePollSeconds)*time.Second, time.Second)

var readInstanceTLSCounter = metrics.NewCounter()
var writeInstanceTLSCounter = metrics.NewCounter()
var readInstanceTLSCacheCounter = metrics.NewCounter()
var writeInstanceTLSCacheCounter = metrics.NewCounter()

func init() {
	metrics.Register("instance_tls.read", readInstanceTLSCounter)
	metrics.Register("instance_tls.write", writeInstanceTLSCounter)
	metrics.Register("instance_tls.read_cache", readInstanceTLSCacheCounter)
	metrics.Register("instance_tls.write_cache", writeInstanceTLSCacheCounter)
}

func requiresTLS(host string, port int, mysql_uri string) bool {
	cacheKey := fmt.Sprintf("%s:%d", host, port)

	if value, found := requireTLSCache.Get(cacheKey); found {
		readInstanceTLSCacheCounter.Inc(1)
		return value.(bool)
	}

	required := false
	db, _, _ := sqlutils.GetDB(mysql_uri)
	if err := db.Ping(); err != nil && (strings.Contains(err.Error(), Error3159) || strings.Contains(err.Error(), Error1045)) {
		required = true
	}

	query := `
			insert into
				database_instance_tls (
					hostname, port, required
				) values (
					?, ?, ?
				)
				on duplicate key update
					required=values(required)
				`
	if _, err := ExecOrchestrator(query, host, port, required); err != nil {
		log.Errore(err)
	}
	writeInstanceTLSCounter.Inc(1)

	requireTLSCache.Set(cacheKey, required, cache.DefaultExpiration)
	writeInstanceTLSCacheCounter.Inc(1)

	return required
}

// Create a TLS configuration from the config supplied CA, Certificate, and Private key.
// Register the TLS config with the mysql drivers as the "topology" config
// Modify the supplied URI to call the TLS config
func SetupMySQLTopologyTLS(uri string) (string, error) {
	if !topologyTLSConfigured {
		tlsConfig, err := ssl.NewTLSConfig(config.Config.MySQLTopologySSLCAFile, !config.Config.MySQLTopologySSLSkipVerify)
		// Drop to TLS 1.0 for talking to MySQL
		tlsConfig.MinVersion = tls.VersionTLS10
		if err != nil {
			return "", log.Errorf("Can't create TLS configuration for Topology connection %s: %s", uri, err)
		}
		tlsConfig.InsecureSkipVerify = config.Config.MySQLTopologySSLSkipVerify

		if config.Config.MySQLTopologyUseMutualTLS ||
			config.Config.MySQLTopologySSLCertFile != "" ||
			config.Config.MySQLTopologySSLPrivateKeyFile != "" {
			if err = ssl.AppendKeyPair(tlsConfig, config.Config.MySQLTopologySSLCertFile, config.Config.MySQLTopologySSLPrivateKeyFile); err != nil {
				return "", log.Errorf("Can't setup TLS key pairs for %s: %s", uri, err)
			}
		}
		if err = mysql.RegisterTLSConfig("topology", tlsConfig); err != nil {
			return "", log.Errorf("Can't register mysql TLS config for topology: %s", err)
		}
		topologyTLSConfigured = true
	}
	return fmt.Sprintf("%s&tls=topology", uri), nil
}

// Create a TLS configuration from the config supplied CA, Certificate, and Private key.
// Register the TLS config with the mysql drivers as the "orchestrator" config
// Modify the supplied URI to call the TLS config
func SetupMySQLOrchestratorTLS(uri string) (string, error) {
	if !orchestratorTLSConfigured {
		tlsConfig, err := ssl.NewTLSConfig(config.Config.MySQLOrchestratorSSLCAFile, !config.Config.MySQLOrchestratorSSLSkipVerify)
		// Drop to TLS 1.0 for talking to MySQL
		tlsConfig.MinVersion = tls.VersionTLS10
		if err != nil {
			return "", log.Fatalf("Can't create TLS configuration for Orchestrator connection %s: %s", uri, err)
		}
		tlsConfig.InsecureSkipVerify = config.Config.MySQLOrchestratorSSLSkipVerify
		if err = ssl.AppendKeyPair(tlsConfig, config.Config.MySQLOrchestratorSSLCertFile, config.Config.MySQLOrchestratorSSLPrivateKeyFile); err != nil {
			return "", log.Fatalf("Can't setup TLS key pairs for %s: %s", uri, err)
		}
		if err = mysql.RegisterTLSConfig("orchestrator", tlsConfig); err != nil {
			return "", log.Fatalf("Can't register mysql TLS config for orchestrator: %s", err)
		}
		orchestratorTLSConfigured = true
	}
	return fmt.Sprintf("%s&tls=orchestrator", uri), nil
}
<|endoftext|>