<commit_before>package commands

import (
	"bufio"
	"fmt"
	"github.com/jingweno/gh/cmd"
	"github.com/jingweno/gh/git"
	"github.com/jingweno/gh/github"
	"github.com/jingweno/gh/utils"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strings"
)

var cmdPullRequest = &Command{
	Run:   pullRequest,
	Usage: "pull-request [-f] [-i ISSUE] [-b BASE] [-d HEAD] [TITLE]",
	Short: "Open a pull request on GitHub",
	Long: `Opens a pull request on GitHub for the project that the "origin" remote
points to. The default head of the pull request is the current branch.
Both base and head of the pull request can be explicitly given in one of
the following formats: "branch", "owner:branch", "owner/repo:branch".
This command will abort operation if it detects that the current topic
branch has local commits that are not yet pushed to its upstream branch
on the remote. To skip this check, use -f.

If TITLE is omitted, a text editor will open in which title and body of
the pull request can be entered in the same manner as git commit message.

If instead of normal TITLE an issue number is given with -i, the pull
request will be attached to an existing GitHub issue. Alternatively, instead
of title you can paste a full URL to an issue on GitHub.
`,
}

var flagPullRequestBase, flagPullRequestHead, flagPullRequestIssue string

func init() {
	cmdPullRequest.Flag.StringVar(&flagPullRequestBase, "b", "master", "BASE")
	cmdPullRequest.Flag.StringVar(&flagPullRequestHead, "d", "", "HEAD")
	cmdPullRequest.Flag.StringVar(&flagPullRequestIssue, "i", "", "ISSUE")
}

func pullRequest(cmd *Command, args *Args) {
	var title, body string
	if args.Size() == 1 {
		title = args.First()
	}

	gh := github.New()
	repo := gh.Project.LocalRepoWith(flagPullRequestBase, flagPullRequestHead)
	if title == "" && flagPullRequestIssue == "" {
		messageFile, err := git.PullReqMsgFile()
		utils.Check(err)

		defer removeFile(messageFile)

		err = writePullRequestChanges(repo, messageFile)
		utils.Check(err)

		editorPath, err := git.EditorPath()
		utils.Check(err)

		err = editTitleAndBody(editorPath, messageFile)
		utils.Check(err)

		title, body, err = readTitleAndBody(messageFile)
		utils.Check(err)
	}

	if title == "" && flagPullRequestIssue == "" {
		log.Fatal("Aborting due to empty pull request title")
	}

	var pullRequestURL string
	var err error
	if title != "" {
		pullRequestURL, err = gh.CreatePullRequest(repo.Base, repo.Head, title, body)
	}
	if flagPullRequestIssue != "" {
		pullRequestURL, err = gh.CreatePullRequestForIssue(repo.Base, repo.Head, flagPullRequestIssue)
	}

	utils.Check(err)

	fmt.Println(pullRequestURL)

	os.Exit(0)
}

func writePullRequestChanges(repo *github.Repo, messageFile string) error {
	message := `
# Requesting a pull to %s from %s
#
# Write a message for this pull request. The first block
# of the text is the title and the rest is description.%s
`
	startRegexp := regexp.MustCompilePOSIX("^")
	endRegexp := regexp.MustCompilePOSIX(" +$")

	commitLogs, _ := git.Log(repo.Base, repo.Head)
	var changesMsg string
	if len(commitLogs) > 0 {
		commitLogs = strings.TrimSpace(commitLogs)
		commitLogs = startRegexp.ReplaceAllString(commitLogs, "# ")
		commitLogs = endRegexp.ReplaceAllString(commitLogs, "")
		changesMsg = `
#
# Changes:
#
%s`
		changesMsg = fmt.Sprintf(changesMsg, commitLogs)
	}

	message = fmt.Sprintf(message, repo.FullBase(), repo.FullHead(), changesMsg)

	return ioutil.WriteFile(messageFile, []byte(message), 0644)
}

func editTitleAndBody(editorPath, messageFile string) error {
	editCmd := cmd.New(editorPath)
	r := regexp.MustCompile("[mg]?vi[m]$")
	if r.MatchString(editorPath) {
		editCmd.WithArg("-c")
		editCmd.WithArg("set ft=gitcommit tw=0 wrap lbr")
	}
	editCmd.WithArg(messageFile)

	return editCmd.Exec()
}

func readTitleAndBody(messageFile string) (title, body string, err error) {
	f, err := os.Open(messageFile)
	defer f.Close()
	if err != nil {
		return "", "", err
	}

	reader := bufio.NewReader(f)

	return readTitleAndBodyFrom(reader)
}

func readTitleAndBodyFrom(reader *bufio.Reader) (title, body string, err error) {
	r := regexp.MustCompile("\\S")
	var titleParts, bodyParts []string

	line, err := readln(reader)
	for err == nil {
		if strings.HasPrefix(line, "#") {
			break
		}
		if len(bodyParts) == 0 && r.MatchString(line) {
			titleParts = append(titleParts, line)
		} else {
			bodyParts = append(bodyParts, line)
		}
		line, err = readln(reader)
	}

	title = strings.Join(titleParts, " ")
	title = strings.TrimSpace(title)

	body = strings.Join(bodyParts, "\n")
	body = strings.TrimSpace(body)

	return title, body, nil
}

func readln(r *bufio.Reader) (string, error) {
	var (
		isPrefix = true
		err      error
		line, ln []byte
	)
	for isPrefix && err == nil {
		line, isPrefix, err = r.ReadLine()
		ln = append(ln, line...)
	}

	return string(ln), err
}

func removeFile(file string) {
	err := os.Remove(file)
	utils.Check(err)
}
<commit_msg>Direct remove instead of defer<commit_after>package commands

import (
	"bufio"
	"fmt"
	"github.com/jingweno/gh/cmd"
	"github.com/jingweno/gh/git"
	"github.com/jingweno/gh/github"
	"github.com/jingweno/gh/utils"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strings"
)

var cmdPullRequest = &Command{
	Run:   pullRequest,
	Usage: "pull-request [-f] [-i ISSUE] [-b BASE] [-d HEAD] [TITLE]",
	Short: "Open a pull request on GitHub",
	Long: `Opens a pull request on GitHub for the project that the "origin" remote
points to. The default head of the pull request is the current branch.
Both base and head of the pull request can be explicitly given in one of
the following formats: "branch", "owner:branch", "owner/repo:branch".
This command will abort operation if it detects that the current topic
branch has local commits that are not yet pushed to its upstream branch
on the remote. To skip this check, use -f.

If TITLE is omitted, a text editor will open in which title and body of
the pull request can be entered in the same manner as git commit message.

If instead of normal TITLE an issue number is given with -i, the pull
request will be attached to an existing GitHub issue. Alternatively, instead
of title you can paste a full URL to an issue on GitHub.
`,
}

var flagPullRequestBase, flagPullRequestHead, flagPullRequestIssue string

func init() {
	cmdPullRequest.Flag.StringVar(&flagPullRequestBase, "b", "master", "BASE")
	cmdPullRequest.Flag.StringVar(&flagPullRequestHead, "d", "", "HEAD")
	cmdPullRequest.Flag.StringVar(&flagPullRequestIssue, "i", "", "ISSUE")
}

func pullRequest(cmd *Command, args *Args) {
	var title, body string
	if args.Size() == 1 {
		title = args.First()
	}

	gh := github.New()
	repo := gh.Project.LocalRepoWith(flagPullRequestBase, flagPullRequestHead)
	if title == "" && flagPullRequestIssue == "" {
		messageFile, err := git.PullReqMsgFile()
		utils.Check(err)

		err = writePullRequestChanges(repo, messageFile)
		utils.Check(err)

		editorPath, err := git.EditorPath()
		utils.Check(err)

		err = editTitleAndBody(editorPath, messageFile)
		utils.Check(err)

		title, body, err = readTitleAndBody(messageFile)
		utils.Check(err)

		err = os.Remove(messageFile)
		utils.Check(err)
	}

	if title == "" && flagPullRequestIssue == "" {
		log.Fatal("Aborting due to empty pull request title")
	}

	var pullRequestURL string
	var err error
	if title != "" {
		pullRequestURL, err = gh.CreatePullRequest(repo.Base, repo.Head, title, body)
	}
	if flagPullRequestIssue != "" {
		pullRequestURL, err = gh.CreatePullRequestForIssue(repo.Base, repo.Head, flagPullRequestIssue)
	}

	utils.Check(err)

	fmt.Println(pullRequestURL)

	os.Exit(0)
}

func writePullRequestChanges(repo *github.Repo, messageFile string) error {
	message := `
# Requesting a pull to %s from %s
#
# Write a message for this pull request. The first block
# of the text is the title and the rest is description.%s
`
	startRegexp := regexp.MustCompilePOSIX("^")
	endRegexp := regexp.MustCompilePOSIX(" +$")

	commitLogs, _ := git.Log(repo.Base, repo.Head)
	var changesMsg string
	if len(commitLogs) > 0 {
		commitLogs = strings.TrimSpace(commitLogs)
		commitLogs = startRegexp.ReplaceAllString(commitLogs, "# ")
		commitLogs = endRegexp.ReplaceAllString(commitLogs, "")
		changesMsg = `
#
# Changes:
#
%s`
		changesMsg = fmt.Sprintf(changesMsg, commitLogs)
	}

	message = fmt.Sprintf(message, repo.FullBase(), repo.FullHead(), changesMsg)

	return ioutil.WriteFile(messageFile, []byte(message), 0644)
}

func editTitleAndBody(editorPath, messageFile string) error {
	editCmd := cmd.New(editorPath)
	r := regexp.MustCompile("[mg]?vi[m]$")
	if r.MatchString(editorPath) {
		editCmd.WithArg("-c")
		editCmd.WithArg("set ft=gitcommit tw=0 wrap lbr")
	}
	editCmd.WithArg(messageFile)

	return editCmd.Exec()
}

func readTitleAndBody(messageFile string) (title, body string, err error) {
	f, err := os.Open(messageFile)
	defer f.Close()
	if err != nil {
		return "", "", err
	}

	reader := bufio.NewReader(f)

	return readTitleAndBodyFrom(reader)
}

func readTitleAndBodyFrom(reader *bufio.Reader) (title, body string, err error) {
	r := regexp.MustCompile("\\S")
	var titleParts, bodyParts []string

	line, err := readln(reader)
	for err == nil {
		if strings.HasPrefix(line, "#") {
			break
		}
		if len(bodyParts) == 0 && r.MatchString(line) {
			titleParts = append(titleParts, line)
		} else {
			bodyParts = append(bodyParts, line)
		}
		line, err = readln(reader)
	}

	title = strings.Join(titleParts, " ")
	title = strings.TrimSpace(title)

	body = strings.Join(bodyParts, "\n")
	body = strings.TrimSpace(body)

	return title, body, nil
}

func readln(r *bufio.Reader) (string, error) {
	var (
		isPrefix = true
		err      error
		line, ln []byte
	)
	for isPrefix && err == nil {
		line, isPrefix, err = r.ReadLine()
		ln = append(ln, line...)
	}

	return string(ln), err
}
<|endoftext|>


<commit_before>// Copyright 2016 The etcd Authors
//
// 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 cmd

import (
	"fmt"
	"time"

	v3 "github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/pkg/report"

	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"gopkg.in/cheggaaa/pb.v1"
)

var leaseKeepaliveCmd = &cobra.Command{
	Use:   "lease-keepalive",
	Short: "Benchmark lease keepalive",

	Run: leaseKeepaliveFunc,
}

var (
	leaseKeepaliveTotal int
)

func init() {
	RootCmd.AddCommand(leaseKeepaliveCmd)
	leaseKeepaliveCmd.Flags().IntVar(&leaseKeepaliveTotal, "total", 10000, "Total number of lease keepalive requests")
}

func leaseKeepaliveFunc(cmd *cobra.Command, args []string) {
	requests := make(chan struct{})
	clients := mustCreateClients(totalClients, totalConns)

	bar = pb.New(leaseKeepaliveTotal)
	bar.Format("Bom !")
	bar.Start()

	r := newReport()
	for i := range clients {
		wg.Add(1)
		go func(c v3.Lease) {
			defer wg.Done()
			resp, err := c.Grant(context.Background(), 100)
			if err != nil {
				panic(err)
			}
			for range requests {
				st := time.Now()
				_, err := c.KeepAliveOnce(context.TODO(), resp.ID)
				r.Results() <- report.Result{Err: err, Start: st, End: time.Now()}
				bar.Increment()
			}
		}(clients[i])
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		for i := 0; i < leaseKeepaliveTotal; i++ {
			requests <- struct{}{}
		}
		close(requests)
	}()

	rc := r.Run()
	wg.Wait()
	close(r.Results())
	bar.Finish()
	fmt.Printf("%s", <-rc)
}
<commit_msg>benchmark: use new lease interface<commit_after>// Copyright 2016 The etcd Authors
//
// 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 cmd

import (
	"fmt"
	"time"

	v3 "github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/pkg/report"

	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"gopkg.in/cheggaaa/pb.v1"
)

var leaseKeepaliveCmd = &cobra.Command{
	Use:   "lease-keepalive",
	Short: "Benchmark lease keepalive",

	Run: leaseKeepaliveFunc,
}

var (
	leaseKeepaliveTotal int
)

func init() {
	RootCmd.AddCommand(leaseKeepaliveCmd)
	leaseKeepaliveCmd.Flags().IntVar(&leaseKeepaliveTotal, "total", 10000, "Total number of lease keepalive requests")
}

func leaseKeepaliveFunc(cmd *cobra.Command, args []string) {
	requests := make(chan struct{})
	clients := mustCreateClients(totalClients, totalConns)

	bar = pb.New(leaseKeepaliveTotal)
	bar.Format("Bom !")
	bar.Start()

	r := newReport()
	for i := range clients {
		wg.Add(1)
		go func(c v3.Lease) {
			defer wg.Done()
			resp, err := c.Grant(context.Background(), 100)
			if err != nil {
				panic(err)
			}
			for range requests {
				st := time.Now()
				ka := c.KeepAliveOnce(context.TODO(), resp.ID)
				r.Results() <- report.Result{Err: ka.Err, Start: st, End: time.Now()}
				bar.Increment()
			}
		}(clients[i])
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		for i := 0; i < leaseKeepaliveTotal; i++ {
			requests <- struct{}{}
		}
		close(requests)
	}()

	rc := r.Run()
	wg.Wait()
	close(r.Results())
	bar.Finish()
	fmt.Printf("%s", <-rc)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"time"
)

// boot2docker config
var B2D struct {
	VBM           string // VirtualBox management utility
	VM            string // boot2docker virtual machine name
	DIR           string // boot2docker directory
	ISO           string // boot2docker ISO image path
	DISK          string // boot2docker disk image path
	DISKSIZE      string // boot2docker disk image size (MB)
	MEMORY        string // boot2docker memory size (MB)
	SSH_HOST_PORT string // boot2docker host SSH port
	DOCKER_PORT   string // boot2docker docker port
}

// helper function to get env var with default values
func getenv(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func init() {
	u, err := user.Current()
	if err != nil {
		log.Fatalf("cannot get current user: %s", err)
	}
	B2D.VBM = getenv("BOOT2DOCKER_VBM", "VBoxManage")
	B2D.VM = getenv("BOOT2DOCKER_VM", "boot2docker-vm")
	B2D.DIR = getenv("BOOT2DOCKER_DIR", filepath.Join(u.HomeDir, ".boot2docker"))
	B2D.ISO = getenv("BOOT2DOCKER_ISO", filepath.Join(B2D.DIR, "boot2docker.iso"))
	B2D.DISK = getenv("BOOT2DOCKER_DISK", filepath.Join(B2D.DIR, "boot2docker.vmdk"))
	B2D.DISKSIZE = getenv("BOOT2DOCKER_DISKSIZE", "20000")
	B2D.MEMORY = getenv("BOOT2DOCKER_MEMORY", "1000")
	B2D.SSH_HOST_PORT = getenv("BOOT2DOCKER_SSH_HOST_PORT", "2022")
	B2D.DOCKER_PORT = getenv("BOOT2DOCKER_DOCKER_PORT", "4243")
}

type VM_STATE int

const (
	VM_UNKNOWN VM_STATE = iota
	VM_RUNNING
	VM_STOPPED
	VM_PAUSED
	VM_SUSPENDED
	VM_ABORTED
)

func main() {
	flag.Parse()

	vm := flag.Arg(1)
	if vm == "" {
		vm = B2D.VM
	}

	// TODO maybe use reflect here?
	switch flag.Arg(0) { // choose subcommand
	case "download":
		cmdDownload()
	case "init":
		cmdInit(vm)
	case "start":
		cmdStart(vm)
	case "up":
		cmdStart(vm)
	case "resume":
		cmdResume(vm)
	case "save":
		cmdSuspend(vm)
	case "pause":
		cmdSuspend(vm)
	case "suspend":
		cmdSuspend(vm)
	case "halt":
		cmdStop(vm)
	case "down":
		cmdStop(vm)
	case "stop":
		cmdStop(vm)
	case "restart":
		cmdRestart(vm)
	case "info":
		cmdInfo(vm)
	case "status":
		cmdStatus(vm)
	case "delete":
		cmdDelete(vm)
	default:
		help()
	}
}

func cmdStart(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
		return
	}
	state := status(vm)
	if state == VM_RUNNING {
		log.Printf("%v is already running.", vm)
		return
	}

	if state == VM_PAUSED {
		log.Printf("Resuming %s", vm)
		vbm("controlvm %s resume", vm)
		wait_vm()
		log.Printf("Resumed.")
	} else {
		log.Printf("Starting %s...", vm)
		vbm("startvm %s --type headless", vm)
		wait_vm()
		log.Printf("Started.")
	}
}

// ping boot2docker VM until it's started
func wait_vm() {
	addr := fmt.Sprintf("localhost:%s", B2D.SSH_HOST_PORT)
	for !ping(addr) {
		time.Sleep(1 * time.Second)
	}
}

func cmdResume(vm string) {
	if status(vm) == VM_SUSPENDED {
		vbm("controlvm %s resume", vm)
	} else {
		log.Printf("%s is not suspended.", vm)
	}
}

func cmdSuspend(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
	}
	if status(vm) == VM_RUNNING {
		log.Printf("Suspending %s", vm)
		vbm("controlvm %s savestate", vm)
	} else {
		log.Printf("%s is not running.", vm)
	}
}

func cmdStop(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
	}

	if status(vm) == VM_RUNNING {
		log.Printf("Shutting down %s...", vm)
		vbm("controlvm %s acpipowerbutton", vm)
		for status(vm) == VM_RUNNING {
			time.Sleep(1 * time.Second)
		}
	} else {
		log.Printf("%s is not running.", vm)
	}
}

func cmdRestart(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
	}

	state := status(vm)
	if state == VM_RUNNING {
		cmdStop(vm)
		time.Sleep(1 * time.Second)
		cmdStart(vm)
	} else {
		cmdStart(vm)
	}
}

func cmdInfo(vm string) {
	if installed(vm) {
		b, _ := vminfo(vm)
		fmt.Printf("%s", b)
	} else {
		fmt.Printf("%s does not exist.\n", vm)
	}
}

func cmdStatus(vm string) {
	switch status(vm) {
	case VM_RUNNING:
		fmt.Printf("%s is running.\n", vm)
	case VM_PAUSED:
		log.Fatalf("%s is paused.", vm)
	case VM_SUSPENDED:
		log.Fatalf("%s is suspended.", vm)
	case VM_STOPPED:
		log.Fatalf("%s is stopped.", vm)
	case VM_ABORTED:
		log.Fatalf("%s is aborted.")
	default:
		log.Fatalf("%s does not exist.", vm)
	}
}

func cmdInit(vm string) {
	if installed(vm) {
		log.Fatalf("%s already exists.\n")
	}

	if ping(fmt.Sprintf("localhost:%s", B2D.DOCKER_PORT)) {
		log.Fatalf("DOCKER_PORT=%s on localhost is occupied. Please choose another port.", B2D.DOCKER_PORT)
	}

	if ping(fmt.Sprintf("localhost:%s", B2D.SSH_HOST_PORT)) {
		log.Fatalf("SSH_HOST_PORT=%s on localhost is occupied. Please choose another port.", B2D.SSH_HOST_PORT)
	}

	log.Printf("Creating VM %s", vm)
	vbm("createvm --name %s --register", vm)
	if vbm(`modifyvm %s
		--ostype Linux26_64
		--cpus %d
		--memory %s
		--rtcuseutc on
		--acpi on
		--ioapic on
		--hpet on
		--hwvirtex on
		--vtxvpid on
		--largepages on
		--nestedpaging on
		--firmware bios
		--bioslogofadein off
		--bioslogofadeout off
		--bioslogodisplaytime 0
		--biosbootmenu disabled
		--boot1 dvd`, vm, runtime.NumCPU(), B2D.MEMORY) != nil {
		log.Fatal("An error occured, upgrade VirtualBox")
		cmdDelete(vm)
	}

	log.Printf("Setting VM networking")
	vbm("modifyvm %s --nic1 nat --nictype1 virtio --cableconnected1 on", vm)
	vbm(`modifyvm %s
		--natpf1 "ssh,tcp,127.0.0.1,%s,,22"
		--natpf1 "docker,tcp,127.0.0.1,%s,,4243"`, vm, B2D.SSH_HOST_PORT, B2D.DOCKER_PORT)

	if _, err := os.Stat(B2D.ISO); os.IsNotExist(err) {
		cmdDownload()
	}

	if _, err := os.Stat(B2D.DISK); os.IsNotExist(err) {
		makeDiskImage()
	}

	log.Printf("Setting VM disks")
	vbm(`storagectl %s --name "SATA" --add sata --hostiocache on`, vm)
	vbm(`storageattach %s --storagectl "SATA" --port 0 --device 0 --type dvddrive --medium %s`, vm, B2D.ISO)
	vbm(`storageattach %s --storagectl "SATA" --port 1 --device 0 --type hdd --medium %s`, vm, B2D.DISK)

	log.Printf("Done.")
	log.Printf("You can now type `boot2docker up` and wait for the VM to start.")
}

func cmdDownload() {
	log.Printf("downloading boot2docker ISO image...")
	tag, err := getLatestReleaseName()
	if err != nil {
		log.Fatalf("failed to get latest release: %s", err)
	}
	log.Printf("  %s", tag)
	err = download(tag, B2D.ISO)
	if err != nil {
		log.Fatalf("failed to download ISO image: %s", err)
	}
}

func cmdDelete(vm string) {
	state := status(vm)
	if state == VM_STOPPED || state == VM_ABORTED {
		vbm("unregistervm --delete %s", vm)
		return
	}
	log.Fatalf("%s needs to be stopped to delete it.", vm)
}

// convenient function to exec a command
func cmd(name string, args ...string) error {
	log.Println(name, args)
	//return nil
	cmd := exec.Command(name, args...)
	return cmd.Run()
}

// convenient function to launch VBoxManage
func vbm(arg string, x ...interface{}) error {
	return cmd(B2D.VBM, strings.Fields(fmt.Sprintf(arg, x...))...)
}

// get the latest boot2docker release tag e.g. v0.5.4
func getLatestReleaseName() (string, error) {
	rsp, err := http.Get("https://api.github.com/repos/steeve/boot2docker/releases")
	if err != nil {
		return "", err
	}
	defer rsp.Body.Close()

	var t [1]struct {
		TagName string `json:"tag_name"`
	}
	err = json.NewDecoder(rsp.Body).Decode(&t)
	if err != nil {
		return "", err
	}
	if len(t) == 0 {
		return "", fmt.Errorf("no releases found")
	}
	return t[0].TagName, nil
}

// download boot2docker ISO for the given tag and save it at dest
func download(tag, dest string) error {
	rsp, err := http.Get(fmt.Sprintf("https://github.com/steeve/boot2docker/releases/download/%s/boot2docker.iso", tag))
	if err != nil {
		return err
	}
	defer rsp.Body.Close()
	f, err := os.Create(dest)
	if err != nil {
		return err
	}
	_, err = io.Copy(f, rsp.Body)
	return err
}

// check if we already have the virtual machine in VirtualBox
func installed(vm string) bool {
	cmd := exec.Command(B2D.VBM, "list", "vms")
	stdout, _ := cmd.Output()
	matched, _ := regexp.Match(fmt.Sprintf(`(?m)^"%s"`, regexp.QuoteMeta(vm)), stdout)
	return matched
}

// get the state of a VM
func status(vm string) VM_STATE {
	b, err := vminfo(vm)
	if err != nil {
		return VM_UNKNOWN
	}
	re := regexp.MustCompile(`(?m)^State:\s+(\w+)`)
	groups := re.FindSubmatch(b)
	if len(groups) < 1 {
		return VM_UNKNOWN
	}
	switch string(groups[1]) {
	case "running":
		return VM_RUNNING
	case "paused":
		return VM_PAUSED
	case "saved":
		return VM_SUSPENDED
	case "powered": // it's actually "powered off"
		return VM_STOPPED
	case "aborted":
		return VM_ABORTED
	default:
		return VM_UNKNOWN
	}
}

// print help message
func help() {
	log.Fatalf("Usage %s {init|start|up|save|pause|stop|restart|resume|status|info|delete|download}\n", os.Args[0])
}

// get VM info
func vminfo(vm string) (stdout []byte, err error) {
	cmd := exec.Command(B2D.VBM, "showvminfo", vm)
	stdout, err = cmd.Output()
	return
}

// check if an addr can be successfully connected
func ping(addr string) bool {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return false
	}
	defer conn.Close()
	return true
}

func makeDiskImage() error {
	log.Printf("Creating %s Meg hard drive...", B2D.DISKSIZE)
	vbm("closemedium disk %s", B2D.DISK)
	vbm("createhd --format VMDK --filename %s --size %s", B2D.DISK, B2D.DISKSIZE)

	f, err := os.Create("format-flag.txt")
	if err != nil {
		return err
	}
	err = f.Truncate(5 * 1024 * 1024) // 5MB
	if err != nil {
		return err
	}
	_, err = f.WriteString("boot2docker, please format-me\n")
	if err != nil {
		return err
	}
	err = f.Close()
	if err != nil {
		return err
	}

	vbm("convertfromraw format-flag.txt format-flag.vmdk --format VMDK")
	vbm("clonehd format-flag.vmdk %s --existing", B2D.DISK)
	vbm("closemedium disk format-flag.vmdk")
	os.Remove("format-flag.txt")
	os.Remove("format-flag.vmdk")
	return nil
}
<commit_msg>add boot2docker ssh<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"time"
)

// boot2docker config
var B2D struct {
	VBM           string // VirtualBox management utility
	VM            string // boot2docker virtual machine name
	DIR           string // boot2docker directory
	ISO           string // boot2docker ISO image path
	DISK          string // boot2docker disk image path
	DISKSIZE      string // boot2docker disk image size (MB)
	MEMORY        string // boot2docker memory size (MB)
	SSH_HOST_PORT string // boot2docker host SSH port
	DOCKER_PORT   string // boot2docker docker port
	SSH	      string // ssh executable
}

// helper function to get env var with default values
func getenv(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func init() {
	u, err := user.Current()
	if err != nil {
		log.Fatalf("cannot get current user: %s", err)
	}
	B2D.VBM = getenv("BOOT2DOCKER_VBM", "VBoxManage")
	B2D.VM = getenv("BOOT2DOCKER_VM", "boot2docker-vm")
	B2D.DIR = getenv("BOOT2DOCKER_DIR", filepath.Join(u.HomeDir, ".boot2docker"))
	B2D.ISO = getenv("BOOT2DOCKER_ISO", filepath.Join(B2D.DIR, "boot2docker.iso"))
	B2D.DISK = getenv("BOOT2DOCKER_DISK", filepath.Join(B2D.DIR, "boot2docker.vmdk"))
	B2D.DISKSIZE = getenv("BOOT2DOCKER_DISKSIZE", "20000")
	B2D.MEMORY = getenv("BOOT2DOCKER_MEMORY", "1000")
	B2D.SSH_HOST_PORT = getenv("BOOT2DOCKER_SSH_HOST_PORT", "2022")
	B2D.DOCKER_PORT = getenv("BOOT2DOCKER_DOCKER_PORT", "4243")
	B2D.SSH = getenv("BOOT2DOCKER_DOCKER_SSH", "ssh")
}

type VM_STATE int

const (
	VM_UNKNOWN VM_STATE = iota
	VM_RUNNING
	VM_STOPPED
	VM_PAUSED
	VM_SUSPENDED
	VM_ABORTED
)

func main() {
	flag.Parse()

	vm := flag.Arg(1)
	if vm == "" {
		vm = B2D.VM
	}

	// TODO maybe use reflect here?
	switch flag.Arg(0) { // choose subcommand
	case "download":
		cmdDownload()
	case "init":
		cmdInit(vm)
	case "start":
		cmdStart(vm)
	case "up":
		cmdStart(vm)
	case "ssh":
		cmdSsh(vm)
	case "resume":
		cmdResume(vm)
	case "save":
		cmdSuspend(vm)
	case "pause":
		cmdSuspend(vm)
	case "suspend":
		cmdSuspend(vm)
	case "halt":
		cmdStop(vm)
	case "down":
		cmdStop(vm)
	case "stop":
		cmdStop(vm)
	case "restart":
		cmdRestart(vm)
	case "info":
		cmdInfo(vm)
	case "status":
		cmdStatus(vm)
	case "delete":
		cmdDelete(vm)
	default:
		help()
	}
}

func cmdSsh(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
		return
	}
	state := status(vm)
	if state != VM_RUNNING {
		log.Printf("%v is not running.", vm)
		return
	}
	cmd := exec.Command(B2D.SSH, "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null", "-p", B2D.SSH_HOST_PORT, "docker@localhost")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
}
func cmdStart(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
		return
	}
	state := status(vm)
	if state == VM_RUNNING {
		log.Printf("%v is already running.", vm)
		return
	}

	if state == VM_PAUSED {
		log.Printf("Resuming %s", vm)
		vbm("controlvm %s resume", vm)
		wait_vm()
		log.Printf("Resumed.")
	} else {
		log.Printf("Starting %s...", vm)
		vbm("startvm %s --type headless", vm)
		wait_vm()
		log.Printf("Started.")
	}
}

// ping boot2docker VM until it's started
func wait_vm() {
	addr := fmt.Sprintf("localhost:%s", B2D.SSH_HOST_PORT)
	for !ping(addr) {
		time.Sleep(1 * time.Second)
	}
}

func cmdResume(vm string) {
	if status(vm) == VM_SUSPENDED {
		vbm("controlvm %s resume", vm)
	} else {
		log.Printf("%s is not suspended.", vm)
	}
}

func cmdSuspend(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
	}
	if status(vm) == VM_RUNNING {
		log.Printf("Suspending %s", vm)
		vbm("controlvm %s savestate", vm)
	} else {
		log.Printf("%s is not running.", vm)
	}
}

func cmdStop(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
	}

	if status(vm) == VM_RUNNING {
		log.Printf("Shutting down %s...", vm)
		vbm("controlvm %s acpipowerbutton", vm)
		for status(vm) == VM_RUNNING {
			time.Sleep(1 * time.Second)
		}
	} else {
		log.Printf("%s is not running.", vm)
	}
}

func cmdRestart(vm string) {
	if !installed(vm) {
		cmdStatus(vm)
	}

	state := status(vm)
	if state == VM_RUNNING {
		cmdStop(vm)
		time.Sleep(1 * time.Second)
		cmdStart(vm)
	} else {
		cmdStart(vm)
	}
}

func cmdInfo(vm string) {
	if installed(vm) {
		b, _ := vminfo(vm)
		fmt.Printf("%s", b)
	} else {
		fmt.Printf("%s does not exist.\n", vm)
	}
}

func cmdStatus(vm string) {
	switch status(vm) {
	case VM_RUNNING:
		fmt.Printf("%s is running.\n", vm)
	case VM_PAUSED:
		log.Fatalf("%s is paused.", vm)
	case VM_SUSPENDED:
		log.Fatalf("%s is suspended.", vm)
	case VM_STOPPED:
		log.Fatalf("%s is stopped.", vm)
	case VM_ABORTED:
		log.Fatalf("%s is aborted.")
	default:
		log.Fatalf("%s does not exist.", vm)
	}
}

func cmdInit(vm string) {
	if installed(vm) {
		log.Fatalf("%s already exists.\n")
	}

	if ping(fmt.Sprintf("localhost:%s", B2D.DOCKER_PORT)) {
		log.Fatalf("DOCKER_PORT=%s on localhost is occupied. Please choose another port.", B2D.DOCKER_PORT)
	}

	if ping(fmt.Sprintf("localhost:%s", B2D.SSH_HOST_PORT)) {
		log.Fatalf("SSH_HOST_PORT=%s on localhost is occupied. Please choose another port.", B2D.SSH_HOST_PORT)
	}

	log.Printf("Creating VM %s", vm)
	vbm("createvm --name %s --register", vm)
	if vbm(`modifyvm %s
		--ostype Linux26_64
		--cpus %d
		--memory %s
		--rtcuseutc on
		--acpi on
		--ioapic on
		--hpet on
		--hwvirtex on
		--vtxvpid on
		--largepages on
		--nestedpaging on
		--firmware bios
		--bioslogofadein off
		--bioslogofadeout off
		--bioslogodisplaytime 0
		--biosbootmenu disabled
		--boot1 dvd`, vm, runtime.NumCPU(), B2D.MEMORY) != nil {
		log.Fatal("An error occured, upgrade VirtualBox")
		cmdDelete(vm)
	}

	log.Printf("Setting VM networking")
	vbm("modifyvm %s --nic1 nat --nictype1 virtio --cableconnected1 on", vm)
	vbm(`modifyvm %s
		--natpf1 "ssh,tcp,127.0.0.1,%s,,22"
		--natpf1 "docker,tcp,127.0.0.1,%s,,4243"`, vm, B2D.SSH_HOST_PORT, B2D.DOCKER_PORT)

	if _, err := os.Stat(B2D.ISO); os.IsNotExist(err) {
		cmdDownload()
	}

	if _, err := os.Stat(B2D.DISK); os.IsNotExist(err) {
		makeDiskImage()
	}

	log.Printf("Setting VM disks")
	vbm(`storagectl %s --name "SATA" --add sata --hostiocache on`, vm)
	vbm(`storageattach %s --storagectl "SATA" --port 0 --device 0 --type dvddrive --medium %s`, vm, B2D.ISO)
	vbm(`storageattach %s --storagectl "SATA" --port 1 --device 0 --type hdd --medium %s`, vm, B2D.DISK)

	log.Printf("Done.")
	log.Printf("You can now type `boot2docker up` and wait for the VM to start.")
}

func cmdDownload() {
	log.Printf("downloading boot2docker ISO image...")
	tag, err := getLatestReleaseName()
	if err != nil {
		log.Fatalf("failed to get latest release: %s", err)
	}
	log.Printf("  %s", tag)
	err = download(tag, B2D.ISO)
	if err != nil {
		log.Fatalf("failed to download ISO image: %s", err)
	}
}

func cmdDelete(vm string) {
	state := status(vm)
	if state == VM_STOPPED || state == VM_ABORTED {
		vbm("unregistervm --delete %s", vm)
		return
	}
	log.Fatalf("%s needs to be stopped to delete it.", vm)
}

// convenient function to exec a command
func cmd(name string, args ...string) error {
	log.Println(name, args)
	//return nil
	cmd := exec.Command(name, args...)
	return cmd.Run()
}

// convenient function to launch VBoxManage
func vbm(arg string, x ...interface{}) error {
	return cmd(B2D.VBM, strings.Fields(fmt.Sprintf(arg, x...))...)
}

// get the latest boot2docker release tag e.g. v0.5.4
func getLatestReleaseName() (string, error) {
	rsp, err := http.Get("https://api.github.com/repos/steeve/boot2docker/releases")
	if err != nil {
		return "", err
	}
	defer rsp.Body.Close()

	var t [1]struct {
		TagName string `json:"tag_name"`
	}
	err = json.NewDecoder(rsp.Body).Decode(&t)
	if err != nil {
		return "", err
	}
	if len(t) == 0 {
		return "", fmt.Errorf("no releases found")
	}
	return t[0].TagName, nil
}

// download boot2docker ISO for the given tag and save it at dest
func download(tag, dest string) error {
	rsp, err := http.Get(fmt.Sprintf("https://github.com/steeve/boot2docker/releases/download/%s/boot2docker.iso", tag))
	if err != nil {
		return err
	}
	defer rsp.Body.Close()
	f, err := os.Create(dest)
	if err != nil {
		return err
	}
	_, err = io.Copy(f, rsp.Body)
	return err
}

// check if we already have the virtual machine in VirtualBox
func installed(vm string) bool {
	cmd := exec.Command(B2D.VBM, "list", "vms")
	stdout, _ := cmd.Output()
	matched, _ := regexp.Match(fmt.Sprintf(`(?m)^"%s"`, regexp.QuoteMeta(vm)), stdout)
	return matched
}

// get the state of a VM
func status(vm string) VM_STATE {
	b, err := vminfo(vm)
	if err != nil {
		return VM_UNKNOWN
	}
	re := regexp.MustCompile(`(?m)^State:\s+(\w+)`)
	groups := re.FindSubmatch(b)
	if len(groups) < 1 {
		return VM_UNKNOWN
	}
	switch string(groups[1]) {
	case "running":
		return VM_RUNNING
	case "paused":
		return VM_PAUSED
	case "saved":
		return VM_SUSPENDED
	case "powered": // it's actually "powered off"
		return VM_STOPPED
	case "aborted":
		return VM_ABORTED
	default:
		return VM_UNKNOWN
	}
}

// print help message
func help() {
	log.Fatalf("Usage %s {init|start|up|save|pause|stop|restart|resume|status|info|delete|download}\n", os.Args[0])
}

// get VM info
func vminfo(vm string) (stdout []byte, err error) {
	cmd := exec.Command(B2D.VBM, "showvminfo", vm)
	stdout, err = cmd.Output()
	return
}

// check if an addr can be successfully connected
func ping(addr string) bool {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return false
	}
	defer conn.Close()
	return true
}

func makeDiskImage() error {
	log.Printf("Creating %s Meg hard drive...", B2D.DISKSIZE)
	vbm("closemedium disk %s", B2D.DISK)
	vbm("createhd --format VMDK --filename %s --size %s", B2D.DISK, B2D.DISKSIZE)

	f, err := os.Create("format-flag.txt")
	if err != nil {
		return err
	}
	err = f.Truncate(5 * 1024 * 1024) // 5MB
	if err != nil {
		return err
	}
	_, err = f.WriteString("boot2docker, please format-me\n")
	if err != nil {
		return err
	}
	err = f.Close()
	if err != nil {
		return err
	}

	vbm("convertfromraw format-flag.txt format-flag.vmdk --format VMDK")
	vbm("clonehd format-flag.vmdk %s --existing", B2D.DISK)
	vbm("closemedium disk format-flag.vmdk")
	os.Remove("format-flag.txt")
	os.Remove("format-flag.vmdk")
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/Sirupsen/logrus"
	"sync"
	"time"
)

var (
	certDomainsObtaining      = make(map[string]bool)
	certDomainsObtainingMutex = &sync.Mutex{}

	skipDomainMap      = make(map[string]time.Time)
	skipDomainMapMutex = &sync.Mutex{}
)

/*
Check for all of domains doesn't in obtaining process. If ok - lock all domains and return true.
If any of domains already in obtaining process - return false and lock nothing.
*/
func obtainDomainsLock(domains []string) bool {
	certDomainsObtainingMutex.Lock()
	defer certDomainsObtainingMutex.Unlock()

	alreadyObtaining := false
	for _, domain := range domains {
		alreadyObtaining = certDomainsObtaining[domain]
		if alreadyObtaining {
			return false
		}
	}

	for _, domain := range domains {
		certDomainsObtaining[domain] = true
	}
	return true
}

/*
Unlock all domains from obtaining lock without any check
*/
func obtainDomainsUnlock(domains []string) {
	certDomainsObtainingMutex.Lock()
	defer certDomainsObtainingMutex.Unlock()

	if logrus.GetLevel() >= logrus.DebugLevel {
		for _, domain := range domains {
			if !certDomainsObtaining[domain] {
				logrus.Debugf("Release from obtain cert not locked domain: %v", domain)
			}
		}
	}

	for _, domain := range domains {
		delete(certDomainsObtaining, domain)
	}
}

func skipDomainsAdd(domains []string) {
	skipDomainMapMutex.Lock()
	defer skipDomainMapMutex.Unlock()

	deadline := time.Now().Add(*blockBadDomainDuration)
	for _, domain := range domains {
		skipDomainMap[domain] = deadline
	}
	logrus.Debugf("Add domains to block '%v', dedline: '%v'", domains, deadline)
}

// return true if some of domains is bad domain
func skipDomainsCheck(domains []string) bool {
	skipDomainMapMutex.Lock()
	defer skipDomainMapMutex.Unlock()

	now := time.Now()
	for _, domain := range domains {
		if deadline, ok := skipDomainMap[domain]; ok && now.Before(deadline) {
			return true
		}
	}
	return false
}

func skipDomainsFlush() {
	logrus.Debug("Flush skip domains set")
	skipDomainMapMutex.Lock()
	skipDomainMap = make(map[string]time.Time)
	skipDomainMapMutex.Unlock()
}

func skipDomainsStartCleaner() {
	skipDomainMapMutex.Lock()
	defer skipDomainMapMutex.Unlock()

	now := time.Now()
	toClean := make([]string, 0, len(skipDomainMap))
	for domain, deadline := range skipDomainMap {
		if deadline.Before(now) {
			toClean = append(toClean, domain)
		}
	}

	for _, domain := range toClean {
		delete(skipDomainMap, domain)
	}

	logrus.Debugf("Clean blocked domains list, remove domains: '%v'", toClean)

	time.AfterFunc(*blockBadDomainDuration, skipDomainsStartCleaner)
}
<commit_msg>zero bad block domains list duration<commit_after>package main

import (
	"github.com/Sirupsen/logrus"
	"sync"
	"time"
)

var (
	certDomainsObtaining      = make(map[string]bool)
	certDomainsObtainingMutex = &sync.Mutex{}

	skipDomainMap      = make(map[string]time.Time)
	skipDomainMapMutex = &sync.Mutex{}
)

/*
Check for all of domains doesn't in obtaining process. If ok - lock all domains and return true.
If any of domains already in obtaining process - return false and lock nothing.
*/
func obtainDomainsLock(domains []string) bool {
	certDomainsObtainingMutex.Lock()
	defer certDomainsObtainingMutex.Unlock()

	alreadyObtaining := false
	for _, domain := range domains {
		alreadyObtaining = certDomainsObtaining[domain]
		if alreadyObtaining {
			return false
		}
	}

	for _, domain := range domains {
		certDomainsObtaining[domain] = true
	}
	return true
}

/*
Unlock all domains from obtaining lock without any check
*/
func obtainDomainsUnlock(domains []string) {
	certDomainsObtainingMutex.Lock()
	defer certDomainsObtainingMutex.Unlock()

	if logrus.GetLevel() >= logrus.DebugLevel {
		for _, domain := range domains {
			if !certDomainsObtaining[domain] {
				logrus.Debugf("Release from obtain cert not locked domain: %v", domain)
			}
		}
	}

	for _, domain := range domains {
		delete(certDomainsObtaining, domain)
	}
}

func skipDomainsAdd(domains []string) {
	if *blockBadDomainDuration == 0 {
		logrus.Debugf("Skip add domains to bad list, becouse bad list is disabled: '%v'", domains)
		return
	}

	skipDomainMapMutex.Lock()
	defer skipDomainMapMutex.Unlock()

	deadline := time.Now().Add(*blockBadDomainDuration)
	for _, domain := range domains {
		skipDomainMap[domain] = deadline
	}
	logrus.Debugf("Add domains to block '%v', dedline: '%v'", domains, deadline)
}

// return true if some of domains is bad domain
func skipDomainsCheck(domains []string) bool {
	skipDomainMapMutex.Lock()
	defer skipDomainMapMutex.Unlock()

	now := time.Now()
	for _, domain := range domains {
		if deadline, ok := skipDomainMap[domain]; ok && now.Before(deadline) {
			return true
		}
	}
	return false
}

func skipDomainsFlush() {
	logrus.Debug("Flush skip domains set")
	skipDomainMapMutex.Lock()
	skipDomainMap = make(map[string]time.Time)
	skipDomainMapMutex.Unlock()
}

func skipDomainsStartCleaner() {
	if *blockBadDomainDuration == 0 {
		logrus.Debugf("Bad domain cleaner doesn't start becouse block domains is disabled")
		return
	}

	skipDomainMapMutex.Lock()
	defer skipDomainMapMutex.Unlock()

	now := time.Now()
	toClean := make([]string, 0, len(skipDomainMap))
	for domain, deadline := range skipDomainMap {
		if deadline.Before(now) {
			toClean = append(toClean, domain)
		}
	}

	for _, domain := range toClean {
		delete(skipDomainMap, domain)
	}

	logrus.Debugf("Clean blocked domains list, remove domains: '%v'", toClean)

	time.AfterFunc(*blockBadDomainDuration, skipDomainsStartCleaner)
}
<|endoftext|>


<commit_before>package encrypted

import (
	"io/ioutil"
	"log"
	"os"

	"github.com/tinzenite/shared"
)

/*
handleLockMessage handles the logic upon receiving a LockMessage. Notably this
includes allowing or disallowing a lock for a specific time frame.
*/
func (c *chaninterface) handleLockMessage(address string, lm *shared.LockMessage) {
	switch lm.Action {
	case shared.LoRequest:
		if c.enc.setLock(address) {
			// if successful notify peer of success
			accept := shared.CreateLockMessage(shared.LoAccept)
			c.enc.channel.Send(address, accept.JSON())
		}
		log.Println("Encrypted: denying request from", address[:8])
		// if not successful we don't do anything, peer will retry
		return
	case shared.LoRelease:
		if c.enc.isLockedAddress(address) {
			c.enc.ClearLock()
			// TODO notify of clear?
			return
		}
		log.Println("handleLockMessage: WARNING: received release request from invalid peer!", address[:8])
	default:
		log.Println("handleLockMessage: Invalid action received!")
	}
}

/*
handleRequestMessage handles the logic upon receiving a RequestMessage. NOTE:
will only be actually handled if Encrypted is currently locked.
*/
func (c *chaninterface) handleRequestMessage(address string, rm *shared.RequestMessage) {
	var data []byte           // data to send
	var identification string // identification for writing temp file
	var err error
	// check file type and fetch data accordingly
	switch rm.ObjType {
	case shared.OtObject:
		// TODO differentiate ORGDIR from path? how? we don't have it... :( FIXME: add more special ID____ things
		// fetch data for normal objects from storage
		data, err = c.enc.storage.Retrieve(rm.Identification)
		identification = rm.Identification
	case shared.OtModel:
		// model is read from specially named file
		data, err = ioutil.ReadFile(c.enc.RootPath + "/" + shared.IDMODEL)
		identification = shared.IDMODEL
	default:
		log.Println("handleRequestMessage: Invalid ObjType requested!", rm.ObjType.String())
		return
	}
	// if error return
	if err != nil {
		log.Println("handleRequestMessage: retrieval of", rm.ObjType, "failed:", err)
		return
	}
	// path for temp file
	filePath := c.enc.RootPath + "/" + shared.SENDINGDIR + "/" + c.buildKey(address, identification)
	// write data to temp sending file
	err = ioutil.WriteFile(filePath, data, shared.FILEPERMISSIONMODE)
	if err != nil {
		log.Println("handleRequestMessage: failed to write data to SEDIR:", err)
		return
	}
	// send file
	err = c.enc.channel.SendFile(address, filePath, rm.Identification, func(success bool) {
		// if NOT success, log and keep file for debugging
		if !success {
			log.Println("handleRequestMessage: Failed to send file on request!", filePath)
			return
		}
		// remove file
		err := os.Remove(filePath)
		if err != nil {
			log.Println("handleRequestMessage: failed to remove temp file:", err)
			return
		}
	})
	// if error log
	if err != nil {
		log.Println("handleRequestMessage: SendFile returned error:", err)
	}
}

/*
handlePushMessage handles the logic upon receiving a PushMessage.
*/
func (c *chaninterface) handlePushMessage(address string, pm *shared.PushMessage) {
	// note that file transfer is allowed for when file is received
	var key string
	switch pm.ObjType {
	case shared.OtObject:
		key = c.buildKey(address, pm.Identification)
	case shared.OtModel:
		key = c.buildKey(address, shared.IDMODEL)
	default:
		log.Println("handlePushMessage: Invalid ObjType pushed!", pm.ObjType.String())
		return
	}
	// if we reach this, allow
	c.enc.allowedTransfers[key] = true
}

/*
buildKey is a helper function that builds the key used to identify transfers.
*/
func (c *chaninterface) buildKey(address, identification string) string {
	return address + ":" + identification
}
<commit_msg>applied NotifyMessage<commit_after>package encrypted

import (
	"io/ioutil"
	"log"
	"os"

	"github.com/tinzenite/shared"
)

/*
handleLockMessage handles the logic upon receiving a LockMessage. Notably this
includes allowing or disallowing a lock for a specific time frame.
*/
func (c *chaninterface) handleLockMessage(address string, lm *shared.LockMessage) {
	switch lm.Action {
	case shared.LoRequest:
		if c.enc.setLock(address) {
			// if successful notify peer of success
			accept := shared.CreateLockMessage(shared.LoAccept)
			c.enc.channel.Send(address, accept.JSON())
		}
		log.Println("Encrypted: denying request from", address[:8])
		// if not successful we don't do anything, peer will retry
		return
	case shared.LoRelease:
		if c.enc.isLockedAddress(address) {
			c.enc.ClearLock()
			// TODO notify of clear?
			return
		}
		log.Println("handleLockMessage: WARNING: received release request from invalid peer!", address[:8])
	default:
		log.Println("handleLockMessage: Invalid action received!")
	}
}

/*
handleRequestMessage handles the logic upon receiving a RequestMessage. NOTE:
will only be actually handled if Encrypted is currently locked.
*/
func (c *chaninterface) handleRequestMessage(address string, rm *shared.RequestMessage) {
	var data []byte           // data to send
	var identification string // identification for writing temp file
	var err error
	// check file type and fetch data accordingly
	switch rm.ObjType {
	case shared.OtObject:
		// TODO differentiate ORGDIR from path? how? we don't have it... :( FIXME: add more special ID____ things
		// fetch data for normal objects from storage
		data, err = c.enc.storage.Retrieve(rm.Identification)
		identification = rm.Identification
	case shared.OtModel:
		// model is read from specially named file
		data, err = ioutil.ReadFile(c.enc.RootPath + "/" + shared.IDMODEL)
		identification = shared.IDMODEL
	default:
		log.Println("handleRequestMessage: Invalid ObjType requested!", rm.ObjType.String())
		return
	}
	// if error return
	if err != nil {
		log.Println("handleRequestMessage: retrieval of", rm.ObjType, "failed:", err)
		// notify sender that it don't exist
		nm := shared.CreateNotifyMessage(shared.NoMissing, identification)
		c.enc.channel.Send(address, nm.JSON())
		return
	}
	// path for temp file
	filePath := c.enc.RootPath + "/" + shared.SENDINGDIR + "/" + c.buildKey(address, identification)
	// write data to temp sending file
	err = ioutil.WriteFile(filePath, data, shared.FILEPERMISSIONMODE)
	if err != nil {
		log.Println("handleRequestMessage: failed to write data to SEDIR:", err)
		return
	}
	// send file
	err = c.enc.channel.SendFile(address, filePath, rm.Identification, func(success bool) {
		// if NOT success, log and keep file for debugging
		if !success {
			log.Println("handleRequestMessage: Failed to send file on request!", filePath)
			return
		}
		// remove file
		err := os.Remove(filePath)
		if err != nil {
			log.Println("handleRequestMessage: failed to remove temp file:", err)
			return
		}
	})
	// if error log
	if err != nil {
		log.Println("handleRequestMessage: SendFile returned error:", err)
	}
}

/*
handlePushMessage handles the logic upon receiving a PushMessage.
*/
func (c *chaninterface) handlePushMessage(address string, pm *shared.PushMessage) {
	// note that file transfer is allowed for when file is received
	var key string
	switch pm.ObjType {
	case shared.OtObject:
		key = c.buildKey(address, pm.Identification)
	case shared.OtModel:
		key = c.buildKey(address, shared.IDMODEL)
	default:
		log.Println("handlePushMessage: Invalid ObjType pushed!", pm.ObjType.String())
		return
	}
	// if we reach this, allow
	c.enc.allowedTransfers[key] = true
}

/*
buildKey is a helper function that builds the key used to identify transfers.
*/
func (c *chaninterface) buildKey(address, identification string) string {
	return address + ":" + identification
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package util

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"
)

type buffer struct {
	b    []byte
	miss int
}

// BufferPool is a 'buffer pool'.
type BufferPool struct {
	pool       [6]chan []byte
	size       [5]uint32
	sizeMiss   [5]uint32
	sizeHalf   [5]uint32
	baseline   [4]int
	baselinex0 int
	baselinex1 int
	baseline0  int
	baseline1  int
	baseline2  int

	mu     sync.RWMutex
	closed bool
	closeC chan struct{}

	get     uint32
	put     uint32
	half    uint32
	less    uint32
	equal   uint32
	greater uint32
	miss    uint32
}

func (p *BufferPool) poolNum(n int) int {
	if n <= p.baseline0 && n > p.baseline0/2 {
		return 0
	}
	for i, x := range p.baseline {
		if n <= x {
			return i + 1
		}
	}
	return len(p.baseline) + 1
}

// Get returns buffer with length of n.
func (p *BufferPool) Get(n int) []byte {
	if p == nil {
		return make([]byte, n)
	}

	p.mu.RLock()
	defer p.mu.RUnlock()

	if p.closed {
		return make([]byte, n)
	}

	atomic.AddUint32(&p.get, 1)

	poolNum := p.poolNum(n)
	pool := p.pool[poolNum]
	if poolNum == 0 {
		// Fast path.
		select {
		case b := <-pool:
			switch {
			case cap(b) > n:
				if cap(b)-n >= n {
					atomic.AddUint32(&p.half, 1)
					select {
					case pool <- b:
					default:
					}
					return make([]byte, n)
				} else {
					atomic.AddUint32(&p.less, 1)
					return b[:n]
				}
			case cap(b) == n:
				atomic.AddUint32(&p.equal, 1)
				return b[:n]
			default:
				atomic.AddUint32(&p.greater, 1)
			}
		default:
			atomic.AddUint32(&p.miss, 1)
		}

		return make([]byte, n, p.baseline0)
	} else {
		sizePtr := &p.size[poolNum-1]

		select {
		case b := <-pool:
			switch {
			case cap(b) > n:
				if cap(b)-n >= n {
					atomic.AddUint32(&p.half, 1)
					sizeHalfPtr := &p.sizeHalf[poolNum-1]
					if atomic.AddUint32(sizeHalfPtr, 1) == 20 {
						atomic.StoreUint32(sizePtr, uint32(cap(b)/2))
						atomic.StoreUint32(sizeHalfPtr, 0)
					} else {
						select {
						case pool <- b:
						default:
						}
					}
					return make([]byte, n)
				} else {
					atomic.AddUint32(&p.less, 1)
					return b[:n]
				}
			case cap(b) == n:
				atomic.AddUint32(&p.equal, 1)
				return b[:n]
			default:
				atomic.AddUint32(&p.greater, 1)
				if uint32(cap(b)) >= atomic.LoadUint32(sizePtr) {
					select {
					case pool <- b:
					default:
					}
				}
			}
		default:
			atomic.AddUint32(&p.miss, 1)
		}

		if size := atomic.LoadUint32(sizePtr); uint32(n) > size {
			if size == 0 {
				atomic.CompareAndSwapUint32(sizePtr, 0, uint32(n))
			} else {
				sizeMissPtr := &p.sizeMiss[poolNum-1]
				if atomic.AddUint32(sizeMissPtr, 1) == 20 {
					atomic.StoreUint32(sizePtr, uint32(n))
					atomic.StoreUint32(sizeMissPtr, 0)
				}
			}
			return make([]byte, n)
		} else {
			return make([]byte, n, size)
		}
	}
}

// Put adds given buffer to the pool.
func (p *BufferPool) Put(b []byte) {
	if p == nil {
		return
	}

	p.mu.RLock()
	defer p.mu.RUnlock()

	if p.closed {
		return
	}

	atomic.AddUint32(&p.put, 1)

	pool := p.pool[p.poolNum(cap(b))]
	select {
	case pool <- b:
	default:
	}

}

func (p *BufferPool) Close() {
	if p == nil {
		return
	}

	p.mu.Lock()
	if !p.closed {
		p.closed = true
		p.closeC <- struct{}{}
	}
	p.mu.Unlock()
}

func (p *BufferPool) String() string {
	if p == nil {
		return "<nil>"
	}

	return fmt.Sprintf("BufferPool{B·%d Z·%v Zm·%v Zh·%v G·%d P·%d H·%d <·%d =·%d >·%d M·%d}",
		p.baseline0, p.size, p.sizeMiss, p.sizeHalf, p.get, p.put, p.half, p.less, p.equal, p.greater, p.miss)
}

func (p *BufferPool) drain() {
	ticker := time.NewTicker(2 * time.Second)
	for {
		select {
		case <-ticker.C:
			for _, ch := range p.pool {
				select {
				case <-ch:
				default:
				}
			}
		case <-p.closeC:
			close(p.closeC)
			for _, ch := range p.pool {
				close(ch)
			}
			return
		}
	}
}

// NewBufferPool creates a new initialized 'buffer pool'.
func NewBufferPool(baseline int) *BufferPool {
	if baseline <= 0 {
		panic("baseline can't be <= 0")
	}
	p := &BufferPool{
		baseline0: baseline,
		baseline:  [...]int{baseline / 4, baseline / 2, baseline * 2, baseline * 4},
		closeC:    make(chan struct{}, 1),
	}
	for i, cap := range []int{2, 2, 4, 4, 2, 1} {
		p.pool[i] = make(chan []byte, cap)
	}
	go p.drain()
	return p
}
<commit_msg>util: BufferPool: removes unused fields<commit_after>// Copyright (c) 2014, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package util

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"
)

type buffer struct {
	b    []byte
	miss int
}

// BufferPool is a 'buffer pool'.
type BufferPool struct {
	pool      [6]chan []byte
	size      [5]uint32
	sizeMiss  [5]uint32
	sizeHalf  [5]uint32
	baseline  [4]int
	baseline0 int

	mu     sync.RWMutex
	closed bool
	closeC chan struct{}

	get     uint32
	put     uint32
	half    uint32
	less    uint32
	equal   uint32
	greater uint32
	miss    uint32
}

func (p *BufferPool) poolNum(n int) int {
	if n <= p.baseline0 && n > p.baseline0/2 {
		return 0
	}
	for i, x := range p.baseline {
		if n <= x {
			return i + 1
		}
	}
	return len(p.baseline) + 1
}

// Get returns buffer with length of n.
func (p *BufferPool) Get(n int) []byte {
	if p == nil {
		return make([]byte, n)
	}

	p.mu.RLock()
	defer p.mu.RUnlock()

	if p.closed {
		return make([]byte, n)
	}

	atomic.AddUint32(&p.get, 1)

	poolNum := p.poolNum(n)
	pool := p.pool[poolNum]
	if poolNum == 0 {
		// Fast path.
		select {
		case b := <-pool:
			switch {
			case cap(b) > n:
				if cap(b)-n >= n {
					atomic.AddUint32(&p.half, 1)
					select {
					case pool <- b:
					default:
					}
					return make([]byte, n)
				} else {
					atomic.AddUint32(&p.less, 1)
					return b[:n]
				}
			case cap(b) == n:
				atomic.AddUint32(&p.equal, 1)
				return b[:n]
			default:
				atomic.AddUint32(&p.greater, 1)
			}
		default:
			atomic.AddUint32(&p.miss, 1)
		}

		return make([]byte, n, p.baseline0)
	} else {
		sizePtr := &p.size[poolNum-1]

		select {
		case b := <-pool:
			switch {
			case cap(b) > n:
				if cap(b)-n >= n {
					atomic.AddUint32(&p.half, 1)
					sizeHalfPtr := &p.sizeHalf[poolNum-1]
					if atomic.AddUint32(sizeHalfPtr, 1) == 20 {
						atomic.StoreUint32(sizePtr, uint32(cap(b)/2))
						atomic.StoreUint32(sizeHalfPtr, 0)
					} else {
						select {
						case pool <- b:
						default:
						}
					}
					return make([]byte, n)
				} else {
					atomic.AddUint32(&p.less, 1)
					return b[:n]
				}
			case cap(b) == n:
				atomic.AddUint32(&p.equal, 1)
				return b[:n]
			default:
				atomic.AddUint32(&p.greater, 1)
				if uint32(cap(b)) >= atomic.LoadUint32(sizePtr) {
					select {
					case pool <- b:
					default:
					}
				}
			}
		default:
			atomic.AddUint32(&p.miss, 1)
		}

		if size := atomic.LoadUint32(sizePtr); uint32(n) > size {
			if size == 0 {
				atomic.CompareAndSwapUint32(sizePtr, 0, uint32(n))
			} else {
				sizeMissPtr := &p.sizeMiss[poolNum-1]
				if atomic.AddUint32(sizeMissPtr, 1) == 20 {
					atomic.StoreUint32(sizePtr, uint32(n))
					atomic.StoreUint32(sizeMissPtr, 0)
				}
			}
			return make([]byte, n)
		} else {
			return make([]byte, n, size)
		}
	}
}

// Put adds given buffer to the pool.
func (p *BufferPool) Put(b []byte) {
	if p == nil {
		return
	}

	p.mu.RLock()
	defer p.mu.RUnlock()

	if p.closed {
		return
	}

	atomic.AddUint32(&p.put, 1)

	pool := p.pool[p.poolNum(cap(b))]
	select {
	case pool <- b:
	default:
	}

}

func (p *BufferPool) Close() {
	if p == nil {
		return
	}

	p.mu.Lock()
	if !p.closed {
		p.closed = true
		p.closeC <- struct{}{}
	}
	p.mu.Unlock()
}

func (p *BufferPool) String() string {
	if p == nil {
		return "<nil>"
	}

	return fmt.Sprintf("BufferPool{B·%d Z·%v Zm·%v Zh·%v G·%d P·%d H·%d <·%d =·%d >·%d M·%d}",
		p.baseline0, p.size, p.sizeMiss, p.sizeHalf, p.get, p.put, p.half, p.less, p.equal, p.greater, p.miss)
}

func (p *BufferPool) drain() {
	ticker := time.NewTicker(2 * time.Second)
	for {
		select {
		case <-ticker.C:
			for _, ch := range p.pool {
				select {
				case <-ch:
				default:
				}
			}
		case <-p.closeC:
			close(p.closeC)
			for _, ch := range p.pool {
				close(ch)
			}
			return
		}
	}
}

// NewBufferPool creates a new initialized 'buffer pool'.
func NewBufferPool(baseline int) *BufferPool {
	if baseline <= 0 {
		panic("baseline can't be <= 0")
	}
	p := &BufferPool{
		baseline0: baseline,
		baseline:  [...]int{baseline / 4, baseline / 2, baseline * 2, baseline * 4},
		closeC:    make(chan struct{}, 1),
	}
	for i, cap := range []int{2, 2, 4, 4, 2, 1} {
		p.pool[i] = make(chan []byte, cap)
	}
	go p.drain()
	return p
}
<|endoftext|>


<commit_before>//#excludeif true
//#-
package main

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"image/png"
	"log"
)

func main() {
	//#+
	//#local image="\""+binfile("./image.png","base64")+"\""
	var imageArray, err = base64.StdEncoding.DecodeString( /*$image+")"$*/ /*-*/ "")
	if err != nil {
		log.Fatal(err)
	}
	var imageConfig, errDecode = png.DecodeConfig(bytes.NewBuffer(imageArray))
	if errDecode != nil {
		log.Fatal(err)
	}

	fmt.Printf("Embedded image has size %dx%d\n", imageConfig.Width, imageConfig.Height)

	//#local text=str2java(evalfile("./text.txt"),false)
	fmt.Println( /*$"\""+text+"\""$*/ )
	//#-
}

//#+
<commit_msg>fixed variable name<commit_after>//#excludeif true
//#-
package main

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"image/png"
	"log"
)

func main() {
	//#+
	//#local image="\""+binfile("./image.png","base64")+"\""
	var imageArray, err = base64.StdEncoding.DecodeString( /*$image+")"$*/ /*-*/ "")
	if err != nil {
		log.Fatal(err)
	}
	var imageConfig, errDecode = png.DecodeConfig(bytes.NewBuffer(imageArray))
	if errDecode != nil {
		log.Fatal(errDecode)
	}

	fmt.Printf("Embedded image has size %dx%d\n", imageConfig.Width, imageConfig.Height)

	//#local text=str2java(evalfile("./text.txt"),false)
	fmt.Println( /*$"\""+text+"\""$*/ )
	//#-
}

//#+
<|endoftext|>


<commit_before>package iocap

import (
	"io"
	"time"
)

// RateOpts is used to encapsulate rate limiting options.
type RateOpts struct {
	D time.Duration
	N int
}

// Reader implements the io.Reader interface and limits the rate at which
// bytes come off of the underlying source reader.
type Reader struct {
	src    io.Reader
	bucket *bucket
}

// NewReader wraps src in a new rate limited reader.
func NewReader(src io.Reader, opts RateOpts) *Reader {
	return &Reader{
		src:    src,
		bucket: newBucket(opts),
	}
}

// Read reads bytes off of the underlying source reader onto p with rate
// limiting. Reads until EOF or until p is filled.
func (r *Reader) Read(p []byte) (n int, err error) {
	for n < len(p) {
		v := r.bucket.wait(len(p) - n)
		v, err = r.src.Read(p[n : n+v])
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return
		}
		n += v
	}
	return
}

// Writer implements the io.Writer interface and limits the rate at which
// bytes are written to the underlying writer.
type Writer struct {
	dst    io.Writer
	bucket *bucket
}

// NewWriter wraps dst in a new rate limited writer.
func NewWriter(dst io.Writer, opts RateOpts) *Writer {
	return &Writer{
		dst:    dst,
		bucket: newBucket(opts),
	}
}

// Write writes len(p) bytes onto the underlying io.Writer, respecting the
// configured rate limit options.
func (w *Writer) Write(p []byte) (n int, err error) {
	for n < len(p) {
		v := w.bucket.wait(len(p) - n)
		v, err = w.dst.Write(p[n : n+v])
		if err != nil {
			return
		}
		n += v
	}
	return
}

// PerMinute returns a RateOpts configured for the given rate per minute.
func PerMinute(n int) RateOpts {
	return RateOpts{time.Minute, n}
}

// PerSecond returns a RateOpts configured for the given rate per second.
func PerSecond(n int) RateOpts {
	return RateOpts{time.Second, n}
}

// bucket is a simple "leaky bucket" abstraction to provide a way to
// limit the number of operations (in this case, byte reads/writes)
// allowed within a given interval.
type bucket struct {
	tokenCh chan struct{}
	opts    RateOpts
	drained time.Time
}

// newBucket creates a new bucket to use for readers and writers.
func newBucket(opts RateOpts) *bucket {
	return &bucket{
		tokenCh: make(chan struct{}, opts.N),
		opts:    opts,
	}
}

// wait is used to wait for n tokens to fit into the bucket. The token
// insert is best-effort, and the actual number of tokens inserted is
// returned. This allows grabbing a bulk of tokens in a single pass.
// wait will block until at least one token is inserted.
func (b *bucket) wait(n int) (v int) {
	// Call a non-blocking drain up-front to make room for tokens.
	b.drain(false)

	for v < n {
		select {
		case b.tokenCh <- struct{}{}:
			v++
		default:
			if v == 0 {
				// Call a blocking drain, because the bucket cannot
				// make progress until the drain interval arrives.
				b.drain(true)
				continue
			}
			return
		}
	}
	return
}

// drain is used to drain the bucket of tokens. If wait is true, drain
// will wait until the next drain cycle and then continue. Otherwise,
// drain only drains the bucket if it is due.
func (b *bucket) drain(wait bool) {
	if wait {
		delay := b.drained.Add(b.opts.D).Sub(time.Now())
		time.Sleep(delay)
	}

	if time.Since(b.drained) >= b.opts.D {
		defer func() {
			b.drained = time.Now()
		}()
		for i := 0; i < b.opts.N; i++ {
			select {
			case <-b.tokenCh:
			default:
				return
			}
		}
	}
}
<commit_msg>Remove the useless PerMinute()<commit_after>package iocap

import (
	"io"
	"time"
)

// RateOpts is used to encapsulate rate limiting options.
type RateOpts struct {
	D time.Duration
	N int
}

// Reader implements the io.Reader interface and limits the rate at which
// bytes come off of the underlying source reader.
type Reader struct {
	src    io.Reader
	bucket *bucket
}

// NewReader wraps src in a new rate limited reader.
func NewReader(src io.Reader, opts RateOpts) *Reader {
	return &Reader{
		src:    src,
		bucket: newBucket(opts),
	}
}

// Read reads bytes off of the underlying source reader onto p with rate
// limiting. Reads until EOF or until p is filled.
func (r *Reader) Read(p []byte) (n int, err error) {
	for n < len(p) {
		v := r.bucket.wait(len(p) - n)
		v, err = r.src.Read(p[n : n+v])
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return
		}
		n += v
	}
	return
}

// Writer implements the io.Writer interface and limits the rate at which
// bytes are written to the underlying writer.
type Writer struct {
	dst    io.Writer
	bucket *bucket
}

// NewWriter wraps dst in a new rate limited writer.
func NewWriter(dst io.Writer, opts RateOpts) *Writer {
	return &Writer{
		dst:    dst,
		bucket: newBucket(opts),
	}
}

// Write writes len(p) bytes onto the underlying io.Writer, respecting the
// configured rate limit options.
func (w *Writer) Write(p []byte) (n int, err error) {
	for n < len(p) {
		v := w.bucket.wait(len(p) - n)
		v, err = w.dst.Write(p[n : n+v])
		if err != nil {
			return
		}
		n += v
	}
	return
}

// PerSecond returns a RateOpts configured to allow n bytes per second.
func PerSecond(n int) RateOpts {
	return RateOpts{time.Second, n}
}

// bucket is a simple "leaky bucket" abstraction to provide a way to
// limit the number of operations (in this case, byte reads/writes)
// allowed within a given interval.
type bucket struct {
	tokenCh chan struct{}
	opts    RateOpts
	drained time.Time
}

// newBucket creates a new bucket to use for readers and writers.
func newBucket(opts RateOpts) *bucket {
	return &bucket{
		tokenCh: make(chan struct{}, opts.N),
		opts:    opts,
	}
}

// wait is used to wait for n tokens to fit into the bucket. The token
// insert is best-effort, and the actual number of tokens inserted is
// returned. This allows grabbing a bulk of tokens in a single pass.
// wait will block until at least one token is inserted.
func (b *bucket) wait(n int) (v int) {
	// Call a non-blocking drain up-front to make room for tokens.
	b.drain(false)

	for v < n {
		select {
		case b.tokenCh <- struct{}{}:
			v++
		default:
			if v == 0 {
				// Call a blocking drain, because the bucket cannot
				// make progress until the drain interval arrives.
				b.drain(true)
				continue
			}
			return
		}
	}
	return
}

// drain is used to drain the bucket of tokens. If wait is true, drain
// will wait until the next drain cycle and then continue. Otherwise,
// drain only drains the bucket if it is due.
func (b *bucket) drain(wait bool) {
	if wait {
		delay := b.drained.Add(b.opts.D).Sub(time.Now())
		time.Sleep(delay)
	}

	if time.Since(b.drained) >= b.opts.D {
		defer func() {
			b.drained = time.Now()
		}()
		for i := 0; i < b.opts.N; i++ {
			select {
			case <-b.tokenCh:
			default:
				return
			}
		}
	}
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strings"

	log "github.com/Sirupsen/logrus"
	discord "github.com/bwmarrin/discordgo"

	"github.com/GinjaNinja32/DisGoIRC/format"
)

// DiscordConfig represents the required config to connect to Discord
type DiscordConfig struct {
	Token        string `json:"token"`
	UseNicknames bool   `json:"use_nicknames"`
}

var (
	dBotID      string
	dSession    *discord.Session
	dGuilds     = map[string]string{}
	dGuildChans = map[string]map[string]string{}

	dMsgQueue = make(chan func())
)

func dInit() {
	d, err := discord.New(fmt.Sprintf("Bot %s", conf.Discord.Token))
	dSession = d
	if err != nil {
		log.Fatalf("Failed to initialise Discord session: %s", err)
	}

	u, err := dSession.User("@me")
	if err != nil {
		log.Fatalf("Failed to get own Discord user: %s", err)
	}

	dBotID = u.ID

	guilds, err := dSession.UserGuilds(99, "", "")
	if err != nil {
		log.Fatalf("Failed to get guilds: %s", err)
	}

	for _, g := range guilds {
		var chans []*discord.Channel
		chans, err = dSession.GuildChannels(g.ID)
		if err != nil {
			log.Fatalf("Failed to get channels for %s: %s", g.Name, err)
		}

		dGuilds[g.Name] = g.ID
		dGuildChans[g.Name] = map[string]string{}
		for _, c := range chans {
			if c.Type == discord.ChannelTypeGuildText {
				dGuildChans[g.Name][c.Name] = c.ID
			}
		}
	}

	dSession.AddHandler(dMessageCreate)

	err = dSession.Open()
	if err != nil {
		log.Fatalf("Failed to connect to Discord: %s", err)
	}

	go func() {
		for f := range dMsgQueue {
			f()
		}
	}()

	log.Infof("Connected to Discord")
}

func dMessageCreate(s *discord.Session, m *discord.MessageCreate) {
	if m.Author.ID == dBotID {
		return
	}

	c, err := s.Channel(m.ChannelID)
	if err != nil {
		log.Errorf("Failed to get channel for incoming message with CID %s: %s", m.ChannelID, err)
		return
	}

	guildID := c.GuildID

	g, err := s.Guild(guildID)
	if err != nil {
		log.Errorf("Failed to get guild with ID %s: %s", guildID, err)
		return
	}

	channel := fmt.Sprintf("%s#%s", g.Name, c.Name)
	authorName := getDisplayNameForUser(m.Author, g.Members)

	if m.Content != "" {
		message := convertMentionsForIRC(g, m)

		dispatchMessageToIRC(authorName, channel, message)
	}
	for _, a := range m.Attachments {
		incomingDiscord(authorName, channel, a.ProxyURL)
	}
}
func convertMentionsForIRC(g *discord.Guild, m *discord.MessageCreate) string {
	message := m.Content

	// Channels
	for _, c := range g.Channels {
		if c.Type != discord.ChannelTypeGuildText {
			continue
		}
		find := fmt.Sprintf("<#%s>", c.ID)
		replace := fmt.Sprintf("#%s", c.Name)
		message = strings.Replace(message, find, replace, -1)
	}

	// Users
	for _, u := range g.Members {
		display := getDisplayNameForMember(u)
		if display == "" {
			log.Errorf("%s/%q/%q had an invalid display name", u.User.ID, u.User.Username, u.Nick)
			continue
		}
		find := fmt.Sprintf("<@%s>", u.User.ID)
		find2 := fmt.Sprintf("<@!%s>", u.User.ID)
		replace := fmt.Sprintf("@%s", display)
		message = strings.Replace(message, find, replace, -1)
		message = strings.Replace(message, find2, replace, -1)
	}

	// Roles
	for _, r := range g.Roles {
		find := fmt.Sprintf("<@&%s>", r.ID)
		replace := fmt.Sprintf("@%s", r.Name)
		message = strings.Replace(message, find, replace, -1)
	}

	// Emojis
	for _, e := range g.Emojis {
		find := fmt.Sprintf("<:%s:%s>", e.Name, e.ID)
		replace := fmt.Sprintf(":%s:", e.Name)
		message = strings.Replace(message, find, replace, -1)
	}

	return message
}

func dispatchMessageToIRC(authorName, channel, message string) {
	// Multiline
	lines := strings.Split(message, "\n")
	lines, forceClip := clipLinesForIRC(lines)
	if len(lines) > 3 || forceClip {
		url := uploadToPtpb(message)

		n := 2
		if len(lines) < 2 {
			n = len(lines)
		}

		for _, line := range lines[:n] {
			incomingDiscord(authorName, channel, line)
		}
		incomingDiscord("[SYSTEM]", channel, fmt.Sprintf("full message from %s: %s", iAddAntiPing(authorName), url))
	} else {
		for _, line := range lines {
			incomingDiscord(authorName, channel, line)
		}
	}
}

func clipLinesForIRC(s []string) ([]string, bool) {
	ret := []string{}
	anyLineForceClip := false

	for _, line := range s {
		if len(line) < 300 {
			ret = append(ret, line)
		} else {
			words := strings.Split(line, " ")
			for len(words) != 0 {
				l := words[0]
				words = words[1:]
				for len(words) != 0 && len(l)+len(words[0]) < 300 {
					l = l + " " + words[0]
					words = words[1:]
				}

				anyLineForceClip = anyLineForceClip || len(l) > 300
				ret = append(ret, l)
			}
		}
	}

	return ret, anyLineForceClip
}

func uploadToPtpb(s string) string {
	resp, err := http.PostForm("https://ptpb.pw/",
		url.Values{"c": {s}, "p": {"1"}})

	if err != nil {
		log.Errorf("Failed to upload to PTPB: %s", err)
		return "Failed to upload to PTPB"
	}
	defer func() {
		_ = resp.Body.Close()
	}()
	if resp.StatusCode == http.StatusOK {
		return resp.Header.Get("Location")
	}

	log.Errorf("Failed to upload to PTPB: HTTP %d", resp.StatusCode)
	ret, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Errorf("Failed to read body: %s", err)
	} else {
		log.Errorf(string(ret))
	}
	return fmt.Sprintf("Failed to upload to PTPB: HTTP %d", resp.StatusCode)

}

var discordEscaper = strings.NewReplacer(
	"\\", "\\\\",
	"*", "\\*",
	"_", "\\_",
)

// StringReplace represents a single string replacement
type StringReplace struct {
	Find    string
	Replace string
}

// StringReplaceGroup represents a group of string replacements to be performed longest-match-first
type StringReplaceGroup []StringReplace

// Add adds a find/replace pair to this group
func (s *StringReplaceGroup) Add(find, replace string) {
	*s = append(*s, StringReplace{find, replace})
}

// Replace performs the replacement represented by this group on the string `str`, returning the result.
// Replacements are performed in length order, longest first.
func (s *StringReplaceGroup) Replace(str string) string {
	sort.Sort(s)
	for _, r := range *s {
		str = strings.Replace(str, r.Find, r.Replace, -1)
	}
	return str
}

func (s StringReplaceGroup) Len() int      { return len(s) }
func (s StringReplaceGroup) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s StringReplaceGroup) Less(i, j int) bool {
	si := s[i].Find
	sj := s[j].Find

	if len(si) > len(sj) {
		return true
	}

	return si < sj
}

func dOutgoing(nick, channel string, messageParsed format.FormattedString) {
	chanParts := strings.Split(channel, "#")
	guildID := dGuilds[chanParts[0]]
	chanID := dGuildChans[chanParts[0]][chanParts[1]]

	g, err := dSession.Guild(guildID)
	if err != nil {
		log.Errorf("Failed to get guild with ID %s: %s", guildID, err)
		return
	}

	message := messageParsed.RenderDiscord()

	// Channels
	for _, c := range g.Channels {
		if c.Type != discord.ChannelTypeGuildText {
			continue
		}
		find := discordEscaper.Replace(fmt.Sprintf("#%s", c.Name))
		replace := fmt.Sprintf("<#%s>", c.ID)
		message = strings.Replace(message, find, replace, -1)
	}

	// Users
	var sr StringReplaceGroup
	for _, u := range g.Members {
		display := getDisplayNameForMember(u)
		if display == "" {
			log.Errorf("%s/%q/%q had an invalid display name", u.User.ID, u.User.Username, u.Nick)
			continue
		}
		find := discordEscaper.Replace(fmt.Sprintf("@%s", display))
		replace := fmt.Sprintf("<@%s>", u.User.ID)
		sr.Add(find, replace)
	}
	message = sr.Replace(message)

	// Roles
	for _, r := range g.Roles {
		find := discordEscaper.Replace(fmt.Sprintf("@%s", r.Name))
		replace := fmt.Sprintf("<@&%s>", r.ID)
		message = strings.Replace(message, find, replace, -1)
	}

	// Emojis
	for _, e := range g.Emojis {
		find := fmt.Sprintf(":%s:", e.Name)
		replace := fmt.Sprintf("<:%s:%s>", e.Name, e.ID)
		message = strings.Replace(message, find, replace, -1)
	}

	dMsgQueue <- func() {
		_, err := dSession.ChannelMessageSend(chanID, fmt.Sprintf("**<%s>** %s", nick, message))
		if err != nil {
			log.Errorf("Failed to send message to %s: <%s> %s", chanID, nick, message)
		}
	}
}

func getDisplayNameForMember(member *discord.Member) string {
	if conf.Discord.UseNicknames && member.Nick != "" {
		return member.Nick
	}

	return member.User.Username
}

func getDisplayNameForUser(user *discord.User, members []*discord.Member) string {
	if conf.Discord.UseNicknames {
		for _, m := range members {
			if m.User.ID == user.ID {
				return getDisplayNameForMember(m)
			}
		}
	}

	return user.Username
}
<commit_msg>fix sorting in StringReplaceGroup<commit_after>package bot

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strings"

	log "github.com/Sirupsen/logrus"
	discord "github.com/bwmarrin/discordgo"

	"github.com/GinjaNinja32/DisGoIRC/format"
)

// DiscordConfig represents the required config to connect to Discord
type DiscordConfig struct {
	Token        string `json:"token"`
	UseNicknames bool   `json:"use_nicknames"`
}

var (
	dBotID      string
	dSession    *discord.Session
	dGuilds     = map[string]string{}
	dGuildChans = map[string]map[string]string{}

	dMsgQueue = make(chan func())
)

func dInit() {
	d, err := discord.New(fmt.Sprintf("Bot %s", conf.Discord.Token))
	dSession = d
	if err != nil {
		log.Fatalf("Failed to initialise Discord session: %s", err)
	}

	u, err := dSession.User("@me")
	if err != nil {
		log.Fatalf("Failed to get own Discord user: %s", err)
	}

	dBotID = u.ID

	guilds, err := dSession.UserGuilds(99, "", "")
	if err != nil {
		log.Fatalf("Failed to get guilds: %s", err)
	}

	for _, g := range guilds {
		var chans []*discord.Channel
		chans, err = dSession.GuildChannels(g.ID)
		if err != nil {
			log.Fatalf("Failed to get channels for %s: %s", g.Name, err)
		}

		dGuilds[g.Name] = g.ID
		dGuildChans[g.Name] = map[string]string{}
		for _, c := range chans {
			if c.Type == discord.ChannelTypeGuildText {
				dGuildChans[g.Name][c.Name] = c.ID
			}
		}
	}

	dSession.AddHandler(dMessageCreate)

	err = dSession.Open()
	if err != nil {
		log.Fatalf("Failed to connect to Discord: %s", err)
	}

	go func() {
		for f := range dMsgQueue {
			f()
		}
	}()

	log.Infof("Connected to Discord")
}

func dMessageCreate(s *discord.Session, m *discord.MessageCreate) {
	if m.Author.ID == dBotID {
		return
	}

	c, err := s.Channel(m.ChannelID)
	if err != nil {
		log.Errorf("Failed to get channel for incoming message with CID %s: %s", m.ChannelID, err)
		return
	}

	guildID := c.GuildID

	g, err := s.Guild(guildID)
	if err != nil {
		log.Errorf("Failed to get guild with ID %s: %s", guildID, err)
		return
	}

	channel := fmt.Sprintf("%s#%s", g.Name, c.Name)
	authorName := getDisplayNameForUser(m.Author, g.Members)

	if m.Content != "" {
		message := convertMentionsForIRC(g, m)

		dispatchMessageToIRC(authorName, channel, message)
	}
	for _, a := range m.Attachments {
		incomingDiscord(authorName, channel, a.ProxyURL)
	}
}
func convertMentionsForIRC(g *discord.Guild, m *discord.MessageCreate) string {
	message := m.Content

	// Channels
	for _, c := range g.Channels {
		if c.Type != discord.ChannelTypeGuildText {
			continue
		}
		find := fmt.Sprintf("<#%s>", c.ID)
		replace := fmt.Sprintf("#%s", c.Name)
		message = strings.Replace(message, find, replace, -1)
	}

	// Users
	for _, u := range g.Members {
		display := getDisplayNameForMember(u)
		if display == "" {
			log.Errorf("%s/%q/%q had an invalid display name", u.User.ID, u.User.Username, u.Nick)
			continue
		}
		find := fmt.Sprintf("<@%s>", u.User.ID)
		find2 := fmt.Sprintf("<@!%s>", u.User.ID)
		replace := fmt.Sprintf("@%s", display)
		message = strings.Replace(message, find, replace, -1)
		message = strings.Replace(message, find2, replace, -1)
	}

	// Roles
	for _, r := range g.Roles {
		find := fmt.Sprintf("<@&%s>", r.ID)
		replace := fmt.Sprintf("@%s", r.Name)
		message = strings.Replace(message, find, replace, -1)
	}

	// Emojis
	for _, e := range g.Emojis {
		find := fmt.Sprintf("<:%s:%s>", e.Name, e.ID)
		replace := fmt.Sprintf(":%s:", e.Name)
		message = strings.Replace(message, find, replace, -1)
	}

	return message
}

func dispatchMessageToIRC(authorName, channel, message string) {
	// Multiline
	lines := strings.Split(message, "\n")
	lines, forceClip := clipLinesForIRC(lines)
	if len(lines) > 3 || forceClip {
		url := uploadToPtpb(message)

		n := 2
		if len(lines) < 2 {
			n = len(lines)
		}

		for _, line := range lines[:n] {
			incomingDiscord(authorName, channel, line)
		}
		incomingDiscord("[SYSTEM]", channel, fmt.Sprintf("full message from %s: %s", iAddAntiPing(authorName), url))
	} else {
		for _, line := range lines {
			incomingDiscord(authorName, channel, line)
		}
	}
}

func clipLinesForIRC(s []string) ([]string, bool) {
	ret := []string{}
	anyLineForceClip := false

	for _, line := range s {
		if len(line) < 300 {
			ret = append(ret, line)
		} else {
			words := strings.Split(line, " ")
			for len(words) != 0 {
				l := words[0]
				words = words[1:]
				for len(words) != 0 && len(l)+len(words[0]) < 300 {
					l = l + " " + words[0]
					words = words[1:]
				}

				anyLineForceClip = anyLineForceClip || len(l) > 300
				ret = append(ret, l)
			}
		}
	}

	return ret, anyLineForceClip
}

func uploadToPtpb(s string) string {
	resp, err := http.PostForm("https://ptpb.pw/",
		url.Values{"c": {s}, "p": {"1"}})

	if err != nil {
		log.Errorf("Failed to upload to PTPB: %s", err)
		return "Failed to upload to PTPB"
	}
	defer func() {
		_ = resp.Body.Close()
	}()
	if resp.StatusCode == http.StatusOK {
		return resp.Header.Get("Location")
	}

	log.Errorf("Failed to upload to PTPB: HTTP %d", resp.StatusCode)
	ret, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Errorf("Failed to read body: %s", err)
	} else {
		log.Errorf(string(ret))
	}
	return fmt.Sprintf("Failed to upload to PTPB: HTTP %d", resp.StatusCode)

}

var discordEscaper = strings.NewReplacer(
	"\\", "\\\\",
	"*", "\\*",
	"_", "\\_",
)

// StringReplace represents a single string replacement
type StringReplace struct {
	Find    string
	Replace string
}

// StringReplaceGroup represents a group of string replacements to be performed longest-match-first
type StringReplaceGroup []StringReplace

// Add adds a find/replace pair to this group
func (s *StringReplaceGroup) Add(find, replace string) {
	*s = append(*s, StringReplace{find, replace})
}

// Replace performs the replacement represented by this group on the string `str`, returning the result.
// Replacements are performed in length order, longest first.
func (s *StringReplaceGroup) Replace(str string) string {
	sort.Sort(s)
	for _, r := range *s {
		str = strings.Replace(str, r.Find, r.Replace, -1)
	}
	return str
}

func (s StringReplaceGroup) Len() int      { return len(s) }
func (s StringReplaceGroup) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s StringReplaceGroup) Less(i, j int) bool {
	si := s[i].Find
	sj := s[j].Find

	if len(si) != len(sj) {
		return len(si) > len(sj)
	}

	return si < sj
}

func dOutgoing(nick, channel string, messageParsed format.FormattedString) {
	chanParts := strings.Split(channel, "#")
	guildID := dGuilds[chanParts[0]]
	chanID := dGuildChans[chanParts[0]][chanParts[1]]

	g, err := dSession.Guild(guildID)
	if err != nil {
		log.Errorf("Failed to get guild with ID %s: %s", guildID, err)
		return
	}

	message := messageParsed.RenderDiscord()

	// Channels
	for _, c := range g.Channels {
		if c.Type != discord.ChannelTypeGuildText {
			continue
		}
		find := discordEscaper.Replace(fmt.Sprintf("#%s", c.Name))
		replace := fmt.Sprintf("<#%s>", c.ID)
		message = strings.Replace(message, find, replace, -1)
	}

	// Users
	var sr StringReplaceGroup
	for _, u := range g.Members {
		display := getDisplayNameForMember(u)
		if display == "" {
			log.Errorf("%s/%q/%q had an invalid display name", u.User.ID, u.User.Username, u.Nick)
			continue
		}
		find := discordEscaper.Replace(fmt.Sprintf("@%s", display))
		replace := fmt.Sprintf("<@%s>", u.User.ID)
		sr.Add(find, replace)
	}
	message = sr.Replace(message)

	// Roles
	for _, r := range g.Roles {
		find := discordEscaper.Replace(fmt.Sprintf("@%s", r.Name))
		replace := fmt.Sprintf("<@&%s>", r.ID)
		message = strings.Replace(message, find, replace, -1)
	}

	// Emojis
	for _, e := range g.Emojis {
		find := fmt.Sprintf(":%s:", e.Name)
		replace := fmt.Sprintf("<:%s:%s>", e.Name, e.ID)
		message = strings.Replace(message, find, replace, -1)
	}

	dMsgQueue <- func() {
		_, err := dSession.ChannelMessageSend(chanID, fmt.Sprintf("**<%s>** %s", nick, message))
		if err != nil {
			log.Errorf("Failed to send message to %s: <%s> %s", chanID, nick, message)
		}
	}
}

func getDisplayNameForMember(member *discord.Member) string {
	if conf.Discord.UseNicknames && member.Nick != "" {
		return member.Nick
	}

	return member.User.Username
}

func getDisplayNameForUser(user *discord.User, members []*discord.Member) string {
	if conf.Discord.UseNicknames {
		for _, m := range members {
			if m.User.ID == user.ID {
				return getDisplayNameForMember(m)
			}
		}
	}

	return user.Username
}
<|endoftext|>


<commit_before>package dallimin

import (
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"fmt"
	"hash/crc32"
	"io"
	"math"
	"net"
	"sort"
	"strconv"
	"strings"
)

const (
	pointsPerServer = 160 // MEMCACHED_POINTS_PER_SERVER_KETAMA
)

var (
	ErrNoServers = errors.New("memcache: no servers configured or available")
)

type Node struct {
	Label  string
	Addr   net.Addr
	Weight int
}

type Entry struct {
	Node  Node
	Point uint
}

type Ring struct {
	addrs []net.Addr
	rings entries
}

type entries []Entry

func (c entries) Less(i, j int) bool { return c[i].Point < c[j].Point }
func (c entries) Len() int           { return len(c) }
func (c entries) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }

func New(servers []string) (*Ring, error) {
	h := &Ring{}

	if len(servers) == 0 {
		return h, nil
	}

	if len(servers) == 1 {
		server := servers[0]

		addr, err := nodeAddr(server)
		if err != nil {
			return nil, err
		}

		h.addrs = append(h.addrs, addr)
		h.rings = append(h.rings, buildEntry(server, addr, 1, 0))

		return h, nil
	}

	totalWeight := len(servers)
	totalServers := len(servers)

	var rings entries
	var addrs []net.Addr

	for _, server := range servers {
		weight := 1
		count := entryCount(weight, totalServers, totalWeight)

		addr, err := nodeAddr(server)
		if err != nil {
			return nil, err
		}

		for i := 0; i < count; i++ {
			rings = append(rings, buildEntry(server, addr, weight, i))
		}

		addrs = append(addrs, addr)
	}

	sort.Sort(rings)

	return &Ring{addrs: addrs, rings: rings}, nil
}

func (h *Ring) Each(f func(net.Addr) error) error {
	for _, a := range h.addrs {
		if err := f(a); nil != err {
			return err
		}
	}
	return nil
}

func (h *Ring) PickServer(key string) (net.Addr, error) {
	if len(h.rings) == 0 {
		return nil, ErrNoServers
	}

	if len(h.rings) == 1 {
		return h.rings[0].Node.Addr, nil
	}

	x := hash(key)
	i := search(h.rings, x)

	return h.rings[i].Node.Addr, nil
}

func hash(key string) uint {
	return uint(crc32.ChecksumIEEE([]byte(key)))
}

func search(ring entries, h uint) uint {
	var maxp = uint(len(ring))
	var lowp = uint(0)
	var highp = maxp

	for {
		midp := (lowp + highp) / 2
		if midp >= maxp {
			if midp == maxp {
				midp = 1
			} else {
				midp = maxp
			}

			return midp - 1
		}

		midval := ring[midp].Point

		var midval1 uint

		if midp == 0 {
			midval1 = 0
		} else {
			midval1 = ring[midp-1].Point
		}

		if h <= midval && h > midval1 {
			return midp - 1
		}

		if midval < h {
			lowp = midp + 1
		} else {
			highp = midp - 1
		}

		if lowp > highp {
			return 0
		}
	}
}

func buildEntry(label string, addr net.Addr, weight int, index int) Entry {
	return Entry{
		Node:  Node{Addr: addr, Weight: weight},
		Point: serverPoint(label, index),
	}
}

func entryCount(weight, totalServers, totalWeight int) int {
	return (totalServers * 160 * weight) / int(math.Floor(float64(totalWeight)))
}

func serverPoint(server string, index int) uint {
	c := sha1.New()
	io.WriteString(c, fmt.Sprintf("%s:%d", server, index))

	hash := hex.EncodeToString(c.Sum(nil))
	hx := "0x" + hash[0:8]

	d, _ := strconv.ParseUint(hx, 0, 64)

	return uint(d)
}

func nodeAddr(node string) (net.Addr, error) {
	if strings.Contains(node, "/") {
		return net.ResolveUnixAddr("unix", node)
	}

	return net.ResolveTCPAddr("tcp", node)
}
<commit_msg>added ability to pass custom weights<commit_after>package dallimin

import (
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"fmt"
	"hash/crc32"
	"io"
	"math"
	"net"
	"sort"
	"strconv"
	"strings"
)

const (
	pointsPerServer = 160 // MEMCACHED_POINTS_PER_SERVER_KETAMA
)

var (
	ErrNoServers = errors.New("memcache: no servers configured or available")
)

type Node struct {
	Label  string
	Addr   net.Addr
	Weight int
}

type Entry struct {
	Node  Node
	Point uint
}

type Ring struct {
	addrs []net.Addr
	rings entries
}

type entries []Entry

func (c entries) Less(i, j int) bool { return c[i].Point < c[j].Point }
func (c entries) Len() int           { return len(c) }
func (c entries) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }

func New(servers []string) (*Ring, error) {
	if len(servers) == 0 {
		return &Ring{}, nil
	}

	sw := make([]int, len(servers))
	for i := range sw {
		sw[i] = 1
	}

	return newRingWeights(servers, sw)
}

func NewWithWeights(servers map[string]int) (*Ring, error) {
	if len(servers) == 0 {
		return &Ring{}, nil
	}

	ss, sw := extract(servers)
	return newRingWeights(ss, sw)
}

func (h *Ring) Each(f func(net.Addr) error) error {
	for _, a := range h.addrs {
		if err := f(a); nil != err {
			return err
		}
	}
	return nil
}

func (h *Ring) PickServer(key string) (net.Addr, error) {
	if len(h.rings) == 0 {
		return nil, ErrNoServers
	}

	if len(h.rings) == 1 {
		return h.rings[0].Node.Addr, nil
	}

	x := hash(key)
	i := search(h.rings, x)

	return h.rings[i].Node.Addr, nil
}

func newRingWeights(ss []string, sw []int) (*Ring, error) {
	h := &Ring{}

	if len(sw) == 1 {
		server := ss[0]
		weight := sw[0]

		addr, err := nodeAddr(server)
		if err != nil {
			return nil, err
		}

		h.addrs = append(h.addrs, addr)
		h.rings = append(h.rings, buildEntry(server, addr, weight, 0))

		return h, nil
	}

	totalWeight := 0
	totalServers := len(ss)

	for i := range sw {
		totalWeight += sw[i]
	}

	var rings entries
	var addrs []net.Addr

	for i, server := range ss {
		weight := sw[i]
		count := entryCount(weight, totalServers, totalWeight)

		addr, err := nodeAddr(server)
		if err != nil {
			return nil, err
		}

		for i := 0; i < count; i++ {
			rings = append(rings, buildEntry(server, addr, weight, i))
		}

		addrs = append(addrs, addr)
	}

	sort.Sort(rings)

	return &Ring{addrs: addrs, rings: rings}, nil
}

func extract(servers map[string]int) ([]string, []int) {
	var ss []string
	var sw []int

	for server, weight := range servers {
		ss = append(ss, server)
		sw = append(sw, weight)
	}

	return ss, sw
}

func hash(key string) uint {
	return uint(crc32.ChecksumIEEE([]byte(key)))
}

func search(ring entries, h uint) uint {
	var maxp = uint(len(ring))
	var lowp = uint(0)
	var highp = maxp

	for {
		midp := (lowp + highp) / 2
		if midp >= maxp {
			if midp == maxp {
				midp = 1
			} else {
				midp = maxp
			}

			return midp - 1
		}

		midval := ring[midp].Point

		var midval1 uint

		if midp == 0 {
			midval1 = 0
		} else {
			midval1 = ring[midp-1].Point
		}

		if h <= midval && h > midval1 {
			return midp - 1
		}

		if midval < h {
			lowp = midp + 1
		} else {
			highp = midp - 1
		}

		if lowp > highp {
			return 0
		}
	}
}

func buildEntry(label string, addr net.Addr, weight int, index int) Entry {
	return Entry{
		Node:  Node{Addr: addr, Weight: weight},
		Point: serverPoint(label, index),
	}
}

func entryCount(weight, totalServers, totalWeight int) int {
	return (totalServers * 160 * weight) / int(math.Floor(float64(totalWeight)))
}

func serverPoint(server string, index int) uint {
	c := sha1.New()
	io.WriteString(c, fmt.Sprintf("%s:%d", server, index))

	hash := hex.EncodeToString(c.Sum(nil))
	hx := "0x" + hash[0:8]

	d, _ := strconv.ParseUint(hx, 0, 64)

	return uint(d)
}

func nodeAddr(node string) (net.Addr, error) {
	if strings.Contains(node, "/") {
		return net.ResolveUnixAddr("unix", node)
	}

	return net.ResolveTCPAddr("tcp", node)
}
<|endoftext|>


<commit_before>package handlers

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/openfaas/faas/gateway/metrics"
	"github.com/prometheus/client_golang/prometheus"
)

// HTTPNotifier notify about HTTP request/response
type HTTPNotifier interface {
	Notify(method string, URL string, originalURL string, statusCode int, duration time.Duration)
}

// PrometheusServiceNotifier notifier for core service endpoints
type PrometheusServiceNotifier struct {
	ServiceMetrics *metrics.ServiceMetricOptions
}

// Notify about service metrics
func (psn PrometheusServiceNotifier) Notify(method string, URL string, originalURL string, statusCode int, duration time.Duration) {
	code := fmt.Sprintf("%d", statusCode)
	path := urlToLabel(URL)
	psn.ServiceMetrics.Counter.WithLabelValues(code, method, path).Inc()
	psn.ServiceMetrics.Histogram.WithLabelValues(method, path, code).Observe(duration.Seconds())
}

func urlToLabel(path string) string {
	if len(path) > 0 {
		path = strings.TrimRight(path, "/")
	}
	if path == "" {
		path = "/"
	}
	return path
}

// PrometheusFunctionNotifier records metrics to Prometheus
type PrometheusFunctionNotifier struct {
	Metrics *metrics.MetricOptions
}

// Notify records metrics in Prometheus
func (p PrometheusFunctionNotifier) Notify(method string, URL string, originalURL string, statusCode int, duration time.Duration) {
	seconds := duration.Seconds()
	serviceName := getServiceName(originalURL)

	p.Metrics.GatewayFunctionsHistogram.
		WithLabelValues(serviceName).
		Observe(seconds)

	code := strconv.Itoa(statusCode)

	p.Metrics.GatewayFunctionInvocation.
		With(prometheus.Labels{"function_name": serviceName, "code": code}).
		Inc()
}

func getServiceName(urlValue string) string {
	var serviceName string
	forward := "/function/"
	if strings.HasPrefix(urlValue, forward) {
		// With a path like `/function/xyz/rest/of/path?q=a`, the service
		// name we wish to locate is just the `xyz` portion.  With a postive
		// match on the regex below, it will return a three-element slice.
		// The item at index `0` is the same as `urlValue`, at `1`
		// will be the service name we need, and at `2` the rest of the path.
		matcher := functionMatcher.Copy()
		matches := matcher.FindStringSubmatch(urlValue)
		if len(matches) == hasPathCount {
			serviceName = matches[nameIndex]
		}
	}
	return strings.Trim(serviceName, "/")
}

// LoggingNotifier notifies a log about a request
type LoggingNotifier struct {
}

// Notify a log about a request
func (LoggingNotifier) Notify(method string, URL string, originalURL string, statusCode int, duration time.Duration) {
	log.Printf("Forwarded [%s] to %s - [%d] - %f seconds", method, originalURL, statusCode, duration.Seconds())
}
<commit_msg>Fix label order for http_requests_total<commit_after>package handlers

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/openfaas/faas/gateway/metrics"
	"github.com/prometheus/client_golang/prometheus"
)

// HTTPNotifier notify about HTTP request/response
type HTTPNotifier interface {
	Notify(method string, URL string, originalURL string, statusCode int, duration time.Duration)
}

// PrometheusServiceNotifier notifier for core service endpoints
type PrometheusServiceNotifier struct {
	ServiceMetrics *metrics.ServiceMetricOptions
}

// Notify about service metrics
func (psn PrometheusServiceNotifier) Notify(method string, URL string, originalURL string, statusCode int, duration time.Duration) {
	code := fmt.Sprintf("%d", statusCode)
	path := urlToLabel(URL)

	psn.ServiceMetrics.Counter.WithLabelValues(method, path, code).Inc()
	psn.ServiceMetrics.Histogram.WithLabelValues(method, path, code).Observe(duration.Seconds())
}

func urlToLabel(path string) string {
	if len(path) > 0 {
		path = strings.TrimRight(path, "/")
	}
	if path == "" {
		path = "/"
	}
	return path
}

// PrometheusFunctionNotifier records metrics to Prometheus
type PrometheusFunctionNotifier struct {
	Metrics *metrics.MetricOptions
}

// Notify records metrics in Prometheus
func (p PrometheusFunctionNotifier) Notify(method string, URL string, originalURL string, statusCode int, duration time.Duration) {
	seconds := duration.Seconds()
	serviceName := getServiceName(originalURL)

	p.Metrics.GatewayFunctionsHistogram.
		WithLabelValues(serviceName).
		Observe(seconds)

	code := strconv.Itoa(statusCode)

	p.Metrics.GatewayFunctionInvocation.
		With(prometheus.Labels{"function_name": serviceName, "code": code}).
		Inc()
}

func getServiceName(urlValue string) string {
	var serviceName string
	forward := "/function/"
	if strings.HasPrefix(urlValue, forward) {
		// With a path like `/function/xyz/rest/of/path?q=a`, the service
		// name we wish to locate is just the `xyz` portion.  With a postive
		// match on the regex below, it will return a three-element slice.
		// The item at index `0` is the same as `urlValue`, at `1`
		// will be the service name we need, and at `2` the rest of the path.
		matcher := functionMatcher.Copy()
		matches := matcher.FindStringSubmatch(urlValue)
		if len(matches) == hasPathCount {
			serviceName = matches[nameIndex]
		}
	}
	return strings.Trim(serviceName, "/")
}

// LoggingNotifier notifies a log about a request
type LoggingNotifier struct {
}

// Notify a log about a request
func (LoggingNotifier) Notify(method string, URL string, originalURL string, statusCode int, duration time.Duration) {
	log.Printf("Forwarded [%s] to %s - [%d] - %f seconds", method, originalURL, statusCode, duration.Seconds())
}
<|endoftext|>