<commit_before>package main

import "math/rand"
import "time"

const DEFAULT_AD_CHANCE = 85

type Plug struct {
	ID             int
	S3ID           string
	Owner          string
	ViewsRemaining int
}

func (p Plug) IsDefault() bool {
	return p.ViewsRemaining >= 0
}

func ChoosePlug(plugs []Plug) Plug {
	rand.Seed(time.Now().Unix())
	// Split plugs into default and custom ads
	var defaults []Plug
	var customs []Plug
	for i := 0; i < len(plugs); i++ {
		if plugs[i].IsDefault() {
			defaults = append(defaults, plugs[i])
		} else {
			customs = append(customs, plugs[i])
		}
	}
	// Decide whether to chose default ad or user submitted ad
	var pickDefault int = rand.Intn(100)
	if pickDefault >= DEFAULT_AD_CHANCE && len(defaults) != 0 {
		return defaults[rand.Intn(len(defaults))]
	} else {
		return customs[rand.Intn(len(customs))]
	}
}
<commit_msg>Make Default Ads super rare<commit_after>package main

import "math/rand"
import "time"

const DEFAULT_AD_CHANCE = 95

type Plug struct {
	ID             int
	S3ID           string
	Owner          string
	ViewsRemaining int
}

func (p Plug) IsDefault() bool {
	return p.ViewsRemaining >= 0
}

func ChoosePlug(plugs []Plug) Plug {
	rand.Seed(time.Now().Unix())
	// Split plugs into default and custom ads
	var defaults []Plug
	var customs []Plug
	for i := 0; i < len(plugs); i++ {
		if plugs[i].IsDefault() {
			defaults = append(defaults, plugs[i])
		} else {
			customs = append(customs, plugs[i])
		}
	}
	// Decide whether to chose default ad or user submitted ad
	var pickDefault int = rand.Intn(100)
	if pickDefault >= DEFAULT_AD_CHANCE && len(defaults) != 0 {
		return defaults[rand.Intn(len(defaults))]
	} else {
		return customs[rand.Intn(len(customs))]
	}
}
<|endoftext|>


<commit_before>/**
 *  Copyright 2014 Paul Querna
 *
 *  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 generator

import (
	"bytes"
	"errors"
	"fmt"
	"go/format"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/pquerna/ffjson/shared"
)

const inceptionMainTemplate = `
// DO NOT EDIT!
// Code generated by ffjson <https://github.com/pquerna/ffjson>
// DO NOT EDIT!

package main

import (
	"github.com/pquerna/ffjson/inception"
	importedinceptionpackage "{{.ImportName}}"
)

func main() {
	i := ffjsoninception.NewInception("{{.InputPath}}", "{{.PackageName}}", "{{.OutputPath}}", {{.ResetFields}})
	i.AddMany(importedinceptionpackage.FFJSONExpose())
	i.Execute()
}
`

const ffjsonExposeTemplate = `
// Code generated by ffjson <https://github.com/pquerna/ffjson>
//
// This should be automatically deleted by running 'ffjson',
// if leftover, please delete it.

package {{.PackageName}}

import (
	ffjsonshared "github.com/pquerna/ffjson/shared"
)

func FFJSONExpose() []ffjsonshared.InceptionType {
	rv := make([]ffjsonshared.InceptionType, 0)
{{range .StructNames}}
	rv = append(rv, ffjsonshared.InceptionType{Obj: {{.Name}}{}, Options: ffjson{{printf "%#v" .Options}} } )
{{end}}
	return rv
}
`

type structName struct {
	Name    string
	Options shared.StructOptions
}

type templateCtx struct {
	StructNames []structName
	ImportName  string
	PackageName string
	InputPath   string
	OutputPath  string
	ResetFields bool
}

type InceptionMain struct {
	goCmd        string
	inputPath    string
	exposePath   string
	outputPath   string
	TempMainPath string
	tempDir      string
	tempMain     *os.File
	tempExpose   *os.File
	resetFields  bool
}

func NewInceptionMain(goCmd string, inputPath string, outputPath string, resetFields bool) *InceptionMain {
	exposePath := getExposePath(inputPath)
	return &InceptionMain{
		goCmd:       goCmd,
		inputPath:   inputPath,
		outputPath:  outputPath,
		exposePath:  exposePath,
		resetFields: resetFields,
	}
}

func getImportName(inputPath string) (string, error) {
	p, err := filepath.Abs(inputPath)
	if err != nil {
		return "", err
	}

	dir := filepath.Dir(p)
	gopaths := strings.Split(os.Getenv("GOPATH"), string(os.PathListSeparator))

	for _, path := range gopaths {
		gpath, err := filepath.Abs(path)
		if err != nil {
			continue
		}
		rel, err := filepath.Rel(filepath.ToSlash(gpath), dir)
		if err != nil {
			return "", err
		}

		if len(rel) < 4 || rel[:4] != "src"+string(os.PathSeparator) {
			continue
		}
		return rel[4:], nil
	}
	return "", errors.New(fmt.Sprintf("Could not find source directory: GOPATH=%q REL=%q", gopaths, dir))

}

func getExposePath(inputPath string) string {
	return inputPath[0:len(inputPath)-3] + "_ffjson_expose.go"
}

func (im *InceptionMain) renderTpl(f *os.File, t *template.Template, tc *templateCtx) error {
	buf := new(bytes.Buffer)
	err := t.Execute(buf, tc)
	if err != nil {
		return err
	}
	formatted, err := format.Source(buf.Bytes())
	if err != nil {
		return err
	}
	_, err = f.Write(formatted)
	return err
}

func (im *InceptionMain) Generate(packageName string, si []*StructInfo, importName string) error {
	var err error

	if importName == "" {
		importName, err = getImportName(im.inputPath)
		if err != nil {
			return err
		}
	}

	im.tempDir, err = ioutil.TempDir(filepath.Dir(im.inputPath), "ffjson-inception")
	if err != nil {
		return err
	}

	importName = filepath.ToSlash(importName)
	// for `go run` to work, we must have a file ending in ".go".
	im.tempMain, err = TempFileWithPostfix(im.tempDir, "ffjson-inception", ".go")
	if err != nil {
		return err
	}

	im.TempMainPath = im.tempMain.Name()
	sn := make([]structName, len(si))
	for i, st := range si {
		sn[i].Name = st.Name
		sn[i].Options = st.Options
	}

	tc := &templateCtx{
		ImportName:  importName,
		PackageName: packageName,
		StructNames: sn,
		InputPath:   im.inputPath,
		OutputPath:  im.outputPath,
		ResetFields: im.resetFields,
	}

	t := template.Must(template.New("inception.go").Parse(inceptionMainTemplate))

	err = im.renderTpl(im.tempMain, t, tc)
	if err != nil {
		return err
	}

	im.tempExpose, err = os.Create(im.exposePath)
	if err != nil {
		return err
	}

	t = template.Must(template.New("ffjson_expose.go").Parse(ffjsonExposeTemplate))

	err = im.renderTpl(im.tempExpose, t, tc)
	if err != nil {
		return err
	}

	return nil
}

func (im *InceptionMain) Run() error {
	var out bytes.Buffer
	var errOut bytes.Buffer

	cmd := exec.Command(im.goCmd, "run", "-a", im.TempMainPath)
	cmd.Stdout = &out
	cmd.Stderr = &errOut

	err := cmd.Run()

	if err != nil {
		return errors.New(
			fmt.Sprintf("Go Run Failed for: %s\nSTDOUT:\n%s\nSTDERR:\n%s\n",
				im.TempMainPath,
				string(out.Bytes()),
				string(errOut.Bytes())))
	}

	defer func() {
		if im.tempExpose != nil {
			im.tempExpose.Close()
		}

		if im.tempMain != nil {
			im.tempMain.Close()
		}

		os.Remove(im.TempMainPath)
		os.Remove(im.exposePath)
		os.Remove(im.tempDir)
	}()

	return nil
}
<commit_msg>Fix for go mod<commit_after>/**
 *  Copyright 2014 Paul Querna
 *
 *  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 generator

import (
	"bytes"
	"errors"
	"fmt"
	"go/format"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/pquerna/ffjson/shared"
)

const inceptionMainTemplate = `
// DO NOT EDIT!
// Code generated by ffjson <https://github.com/pquerna/ffjson>
// DO NOT EDIT!

package main

import (
	"github.com/pquerna/ffjson/inception"
	importedinceptionpackage "{{.ImportName}}"
)

func main() {
	i := ffjsoninception.NewInception("{{.InputPath}}", "{{.PackageName}}", "{{.OutputPath}}", {{.ResetFields}})
	i.AddMany(importedinceptionpackage.FFJSONExpose())
	i.Execute()
}
`

const ffjsonExposeTemplate = `
// Code generated by ffjson <https://github.com/pquerna/ffjson>
//
// This should be automatically deleted by running 'ffjson',
// if leftover, please delete it.

package {{.PackageName}}

import (
	ffjsonshared "github.com/pquerna/ffjson/shared"
)

func FFJSONExpose() []ffjsonshared.InceptionType {
	rv := make([]ffjsonshared.InceptionType, 0)
{{range .StructNames}}
	rv = append(rv, ffjsonshared.InceptionType{Obj: {{.Name}}{}, Options: ffjson{{printf "%#v" .Options}} } )
{{end}}
	return rv
}
`

type structName struct {
	Name    string
	Options shared.StructOptions
}

type templateCtx struct {
	StructNames []structName
	ImportName  string
	PackageName string
	InputPath   string
	OutputPath  string
	ResetFields bool
}

type InceptionMain struct {
	goCmd        string
	inputPath    string
	exposePath   string
	outputPath   string
	TempMainPath string
	tempDir      string
	tempMain     *os.File
	tempExpose   *os.File
	resetFields  bool
}

func NewInceptionMain(goCmd string, inputPath string, outputPath string, resetFields bool) *InceptionMain {
	exposePath := getExposePath(inputPath)
	return &InceptionMain{
		goCmd:       goCmd,
		inputPath:   inputPath,
		outputPath:  outputPath,
		exposePath:  exposePath,
		resetFields: resetFields,
	}
}

func getImportName(goCmd, inputPath string) (string, error) {
	p, err := filepath.Abs(inputPath)
	if err != nil {
		return "", err
	}
	dir := filepath.Dir(p)

	// `go list dir` gives back the module name
	// Should work for GOPATH as well as with modules
	// Errors if no go files are found
	cmd := exec.Command(goCmd, "list", dir)
	b, err := cmd.Output()
	if err == nil {
		return string(b[:len(b)-1]), nil
	}

	gopaths := strings.Split(os.Getenv("GOPATH"), string(os.PathListSeparator))

	for _, path := range gopaths {
		gpath, err := filepath.Abs(path)
		if err != nil {
			continue
		}
		rel, err := filepath.Rel(filepath.ToSlash(gpath), dir)
		if err != nil {
			return "", err
		}

		if len(rel) < 4 || rel[:4] != "src"+string(os.PathSeparator) {
			continue
		}
		return rel[4:], nil
	}
	return "", errors.New(fmt.Sprintf("Could not find source directory: GOPATH=%q REL=%q", gopaths, dir))

}

func getExposePath(inputPath string) string {
	return inputPath[0:len(inputPath)-3] + "_ffjson_expose.go"
}

func (im *InceptionMain) renderTpl(f *os.File, t *template.Template, tc *templateCtx) error {
	buf := new(bytes.Buffer)
	err := t.Execute(buf, tc)
	if err != nil {
		return err
	}
	formatted, err := format.Source(buf.Bytes())
	if err != nil {
		return err
	}
	_, err = f.Write(formatted)
	return err
}

func (im *InceptionMain) Generate(packageName string, si []*StructInfo, importName string) error {
	var err error
	if importName == "" {
		importName, err = getImportName(im.goCmd, im.inputPath)
		if err != nil {
			return err
		}
	}

	im.tempDir, err = ioutil.TempDir(filepath.Dir(im.inputPath), "ffjson-inception")
	if err != nil {
		return err
	}

	importName = filepath.ToSlash(importName)
	// for `go run` to work, we must have a file ending in ".go".
	im.tempMain, err = TempFileWithPostfix(im.tempDir, "ffjson-inception", ".go")
	if err != nil {
		return err
	}

	im.TempMainPath = im.tempMain.Name()
	sn := make([]structName, len(si))
	for i, st := range si {
		sn[i].Name = st.Name
		sn[i].Options = st.Options
	}

	tc := &templateCtx{
		ImportName:  importName,
		PackageName: packageName,
		StructNames: sn,
		InputPath:   im.inputPath,
		OutputPath:  im.outputPath,
		ResetFields: im.resetFields,
	}

	t := template.Must(template.New("inception.go").Parse(inceptionMainTemplate))

	err = im.renderTpl(im.tempMain, t, tc)
	if err != nil {
		return err
	}

	im.tempExpose, err = os.Create(im.exposePath)
	if err != nil {
		return err
	}

	t = template.Must(template.New("ffjson_expose.go").Parse(ffjsonExposeTemplate))

	err = im.renderTpl(im.tempExpose, t, tc)
	if err != nil {
		return err
	}

	return nil
}

func (im *InceptionMain) Run() error {
	var out bytes.Buffer
	var errOut bytes.Buffer

	cmd := exec.Command(im.goCmd, "run", "-a", im.TempMainPath)
	cmd.Stdout = &out
	cmd.Stderr = &errOut

	err := cmd.Run()

	if err != nil {
		return errors.New(
			fmt.Sprintf("Go Run Failed for: %s\nSTDOUT:\n%s\nSTDERR:\n%s\n",
				im.TempMainPath,
				string(out.Bytes()),
				string(errOut.Bytes())))
	}

	defer func() {
		if im.tempExpose != nil {
			im.tempExpose.Close()
		}

		if im.tempMain != nil {
			im.tempMain.Close()
		}

		os.Remove(im.TempMainPath)
		os.Remove(im.exposePath)
		os.Remove(im.tempDir)
	}()

	return nil
}
<|endoftext|>


<commit_before>// +build windows

package generic

import (
	"io"
	"net"
	"syscall"
)

func rawCopy(dst io.Writer, src *net.TCPConn, ctrl *CopyControl) (written int64, err error) {
	c, err := src.SyscallConn()
	if err != nil {
		return 0, err
	}

	buf := ctrl.Buffer
	for {
		var er error
		var nr int
		rr := c.Read(func(s uintptr) bool {
			ctrl.Lock()
			defer ctrl.Unlock()
			var read uint32
			var flags uint32
			var wsabuf syscall.WSABuf
			wsabuf.Buf = &buf[0]
			wsabuf.Len = uint32(len(buf))
			er = syscall.WSARecv(syscall.Handle(s), &wsabuf, 1, &read, &flags, nil, nil)
			nr = int(read)
			return true
		})

		// read EOF
		if nr == 0 && er == nil {
			break
		}

		if nr > 0 {
			ctrl.Lock()
			nw, ew := dst.Write(buf[0:nr])
			ctrl.Unlock()
			buf = nil
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				err = er
			}
			break
		}
		if rr != nil {
			if rr != io.EOF {
				err = rr
			}
			break
		}
	}

	return written, err
}
<commit_msg>fix error<commit_after>// +build windows

package generic

import (
	"io"
	"net"
	"syscall"
)

func rawCopy(dst io.Writer, src *net.TCPConn, ctrl *CopyControl) (written int64, err error) {
	c, err := src.SyscallConn()
	if err != nil {
		return 0, err
	}

	buf := ctrl.Buffer
	for {
		var er error
		var nr int
		rr := c.Read(func(s uintptr) bool {
			ctrl.Lock()
			defer ctrl.Unlock()
			var read uint32
			var flags uint32
			var wsabuf syscall.WSABuf
			wsabuf.Buf = &buf[0]
			wsabuf.Len = uint32(len(buf))
			er = syscall.WSARecv(syscall.Handle(s), &wsabuf, 1, &read, &flags, nil, nil)
			nr = int(read)
			return true
		})

		// read EOF
		if nr == 0 && er == nil {
			break
		}

		if nr > 0 {
			ctrl.Lock()
			nw, ew := dst.Write(buf[0:nr])
			ctrl.Unlock()
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				err = er
			}
			break
		}
		if rr != nil {
			if rr != io.EOF {
				err = rr
			}
			break
		}
	}

	return written, err
}
<|endoftext|>


<commit_before>package prosper

import "github.com/mtlynch/gofn-prosper/types"

// Accounter supports the Account interface for retrieving user account
// information.
type Accounter interface {
	Account() (types.AccountInformation, error)
}

// Accounts queries the Propser API for properties of the user's account,
// including balance information and note summaries. Accounts partially
// implements the REST API described at:
// https://developers.prosper.com/docs/investor/accounts-api/
func (c Client) Account() (types.AccountInformation, error) {
	rawResponse, err := c.rawClient.Accounts()
	if err != nil {
		return types.AccountInformation{}, err
	}
	return c.ap.Parse(rawResponse)
}
<commit_msg>Fixing comment on Account (#26)<commit_after>package prosper

import "github.com/mtlynch/gofn-prosper/types"

// Accounter supports the Account interface for retrieving user account
// information.
type Accounter interface {
	Account() (types.AccountInformation, error)
}

// Account queries the Propser API for properties of the user's account,
// including balance information and note summaries. Accounts partially
// implements the REST API described at:
// https://developers.prosper.com/docs/investor/accounts-api/
func (c Client) Account() (types.AccountInformation, error) {
	rawResponse, err := c.rawClient.Accounts()
	if err != nil {
		return types.AccountInformation{}, err
	}
	return c.ap.Parse(rawResponse)
}
<|endoftext|>


<commit_before>package provider

import (
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	fsnotify "gopkg.in/fsnotify.v1"

	"github.com/BurntSushi/toml"
	"github.com/Nitro/sidecar/catalog"
	"github.com/Nitro/sidecar/service"
	"github.com/containous/traefik/log"
	"github.com/containous/traefik/safe"
	"github.com/containous/traefik/types"
)

const (
	method               = "drr"
	sticky               = false
	maxConnAmount        = 300
	maxConnExtractorFunc = "client.ip"
	circuitBreaker       = "ResponseCodeRatio(500, 600, 0, 600) > 0.3"
	weight               = 0
)

var _ Provider = (*Sidecar)(nil)

// Sidecar holds configurations of the Sidecar provider
type Sidecar struct {
	BaseProvider      `mapstructure:",squash"`
	Endpoint          string `description:"Sidecar URL"`
	Frontend          string `description:"Configuration file for frontend"`
	configurationChan chan<- types.ConfigMessage
	RefreshConn       time.Duration `description:"How often to refresh the connection to Sidecar backend"`
}

type callback func(map[string][]*service.Service, error)

// Provide allows the provider to provide configurations to traefik
// using the given configuration channel.
func (provider *Sidecar) Provide(configurationChan chan<- types.ConfigMessage, pool *safe.Pool, constraints types.Constraints) error {
	provider.configurationChan = configurationChan
	if provider.Watch {
		safe.Go(func() { provider.sidecarWatcher() })

		watcher, err := fsnotify.NewWatcher()
		if err != nil {
			log.Errorln("Error creating file watcher", err)
			return err
		}

		file, err := os.Open(provider.Frontend)
		if err != nil {
			log.Errorln("Error opening file", err)
			return err
		}
		defer file.Close()
		pool.Go(func(stop chan bool) {
			defer watcher.Close()
			for {
				select {
				case <-stop:
					return
				case event := <-watcher.Events:
					if strings.Contains(event.Name, file.Name()) {
						log.Debug("Sidecar Frontend File event:", event)
						states, errState := provider.fetchState()
						if errState != nil {
							log.Errorln("Error reloading Sidecar config", errState)
						}
						provider.loadSidecarConfig(states.ByService())
					}
				case errWatcher := <-watcher.Errors:
					log.Errorln("Watcher event error", errWatcher)
				}
			}
		})
		err = watcher.Add(filepath.Dir(file.Name()))
		if err != nil {
			log.Error("Error adding file watcher", err)
			return err
		}
	}
	states, err := provider.fetchState()
	if err != nil {
		log.Fatalln("Error reloading Sidecar config", err)
	}
	err = provider.loadSidecarConfig(states.ByService())
	if err != nil {
		return err
	}
	return nil
}

func (provider *Sidecar) constructConfig(sidecarStates map[string][]*service.Service) (*types.Configuration, error) {
	sidecarConfig := types.Configuration{}
	log.Infoln("loading sidecar config")
	sidecarConfig.Backends = provider.makeBackends(sidecarStates)
	log.Infoln("loading frontend config from file: ", provider.Frontend)
	var err error
	sidecarConfig.Frontends, err = provider.makeFrontend()
	if err != nil {
		return nil, err
	}
	return &sidecarConfig, nil
}

func (provider *Sidecar) loadSidecarConfig(sidecarStates map[string][]*service.Service) error {
	conf, err := provider.constructConfig(sidecarStates)
	if err != nil {
		return err
	}
	provider.configurationChan <- types.ConfigMessage{
		ProviderName:  "sidecar",
		Configuration: conf,
	}
	return nil
}

func (provider *Sidecar) sidecarWatcher() error {
	//set timeout to be just a bot more than connection refresh interval
	client := &http.Client{Timeout: (provider.RefreshConn + 3) * time.Second}
	log.Infof("Using %s Sidecar connection refresh interval", provider.RefreshConn)
	provider.recycleConn(client)
	return nil
}

func (provider *Sidecar) recycleConn(client *http.Client) {
	var err error
	var resp *http.Response
	for { //use refresh interval to occasionally reconnect to Sidecar in case the stream connection is lost
		resp, err = client.Get(provider.Endpoint + "/watch")
		if err != nil {
			log.Errorf("Error connecting to Sidecar: %s, Error: %s", provider.Endpoint, err)
			time.Sleep(5 * time.Second)
			continue
		}
		go catalog.DecodeStream(resp.Body, provider.callbackLoader)
		time.Sleep(provider.RefreshConn * time.Second)
	}
}

func (provider *Sidecar) callbackLoader(sidecarStates map[string][]*service.Service, err error) error {
	if err != nil {
		return err
	}
	provider.loadSidecarConfig(sidecarStates)
	return nil
}

func (provider *Sidecar) makeFrontend() (map[string]*types.Frontend, error) {
	configuration := new(types.Configuration)
	if _, err := toml.DecodeFile(provider.Frontend, configuration); err != nil {
		log.Errorf("Error reading file: %s", err)
		return nil, err
	}
	return configuration.Frontends, nil
}

func (provider *Sidecar) makeBackends(sidecarStates map[string][]*service.Service) map[string]*types.Backend {
	sidecarBacks := make(map[string]*types.Backend)
	for serviceName, services := range sidecarStates {
		newServers := make(map[string]types.Server)
		newBackend := &types.Backend{LoadBalancer: &types.LoadBalancer{Method: method, Sticky: sticky},
			MaxConn:        &types.MaxConn{Amount: maxConnAmount, ExtractorFunc: maxConnExtractorFunc},
			CircuitBreaker: &types.CircuitBreaker{Expression: circuitBreaker},
			Servers:        newServers}
		for _, serv := range services {
			if serv.IsAlive() {
				for i := 0; i < len(serv.Ports); i++ {
					ipAddr, err := net.LookupIP(serv.Hostname)
					if err != nil {
						log.Errorln("Error resolving Ip address, ", err)
						newBackend.Servers[serv.Hostname] = types.Server{URL: "http://" + serv.Hostname + ":" + strconv.FormatInt(serv.Ports[i].Port, 10), Weight: weight}
					} else {
						newBackend.Servers[serv.Hostname] = types.Server{URL: "http://" + ipAddr[0].String() + ":" + strconv.FormatInt(serv.Ports[i].Port, 10), Weight: weight}
					}
				}
			}
		}
		sidecarBacks[serviceName] = newBackend
	}
	return sidecarBacks
}

func (provider *Sidecar) fetchState() (*catalog.ServicesState, error) {
	client := &http.Client{Timeout: 5 * time.Second}
	resp, err := client.Get(provider.Endpoint + "/state.json")
	if err != nil {
		return nil, err
	}

	bytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	state, err := catalog.Decode(bytes)
	if err != nil {
		return nil, err
	}
	return state, nil
}
<commit_msg>add timer external to sidecar watcher<commit_after>package provider

import (
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	fsnotify "gopkg.in/fsnotify.v1"

	"github.com/BurntSushi/toml"
	"github.com/Nitro/sidecar/catalog"
	"github.com/Nitro/sidecar/service"
	"github.com/containous/traefik/log"
	"github.com/containous/traefik/safe"
	"github.com/containous/traefik/types"
)

const (
	method               = "drr"
	sticky               = false
	maxConnAmount        = 300
	maxConnExtractorFunc = "client.ip"
	circuitBreaker       = "ResponseCodeRatio(500, 600, 0, 600) > 0.3"
	weight               = 0
)

var _ Provider = (*Sidecar)(nil)

// Sidecar holds configurations of the Sidecar provider
type Sidecar struct {
	BaseProvider      `mapstructure:",squash"`
	Endpoint          string `description:"Sidecar URL"`
	Frontend          string `description:"Configuration file for frontend"`
	configurationChan chan<- types.ConfigMessage
	RefreshConn       time.Duration `description:"How often to refresh the connection to Sidecar backend"`
	connTimer         *time.Timer
}

type callback func(map[string][]*service.Service, error)

// Provide allows the provider to provide configurations to traefik
// using the given configuration channel.
func (provider *Sidecar) Provide(configurationChan chan<- types.ConfigMessage, pool *safe.Pool, constraints types.Constraints) error {
	provider.configurationChan = configurationChan
	if provider.Watch {
		safe.Go(func() { provider.sidecarWatcher() })

		watcher, err := fsnotify.NewWatcher()
		if err != nil {
			log.Errorln("Error creating file watcher", err)
			return err
		}

		file, err := os.Open(provider.Frontend)
		if err != nil {
			log.Errorln("Error opening file", err)
			return err
		}
		defer file.Close()
		pool.Go(func(stop chan bool) {
			defer watcher.Close()
			for {
				select {
				case <-stop:
					return
				case event := <-watcher.Events:
					if strings.Contains(event.Name, file.Name()) {
						log.Debug("Sidecar Frontend File event:", event)
						states, errState := provider.fetchState()
						if errState != nil {
							log.Errorln("Error reloading Sidecar config", errState)
						}
						provider.loadSidecarConfig(states.ByService())
					}
				case errWatcher := <-watcher.Errors:
					log.Errorln("Watcher event error", errWatcher)
				}
			}
		})
		err = watcher.Add(filepath.Dir(file.Name()))
		if err != nil {
			log.Error("Error adding file watcher", err)
			return err
		}
	}
	states, err := provider.fetchState()
	if err != nil {
		log.Fatalln("Error reloading Sidecar config", err)
	}
	err = provider.loadSidecarConfig(states.ByService())
	if err != nil {
		return err
	}
	return nil
}

func (provider *Sidecar) constructConfig(sidecarStates map[string][]*service.Service) (*types.Configuration, error) {
	sidecarConfig := types.Configuration{}
	log.Infoln("loading sidecar config")
	sidecarConfig.Backends = provider.makeBackends(sidecarStates)
	log.Infoln("loading frontend config from file: ", provider.Frontend)
	var err error
	sidecarConfig.Frontends, err = provider.makeFrontend()
	if err != nil {
		return nil, err
	}
	return &sidecarConfig, nil
}

func (provider *Sidecar) loadSidecarConfig(sidecarStates map[string][]*service.Service) error {
	conf, err := provider.constructConfig(sidecarStates)
	if err != nil {
		return err
	}
	provider.configurationChan <- types.ConfigMessage{
		ProviderName:  "sidecar",
		Configuration: conf,
	}
	return nil
}

func (provider *Sidecar) sidecarWatcher() error {
	//set timeout to be just a bot more than connection refresh interval
	provider.connTimer = time.NewTimer(provider.RefreshConn * time.Second)
	tr := &http.Transport{ResponseHeaderTimeout: 0}
	client := &http.Client{
		Timeout:   0,
		Transport: tr}
	log.Infof("Using %s Sidecar connection refresh interval", provider.RefreshConn)
	provider.recycleConn(client, tr)
	return nil
}

func (provider *Sidecar) recycleConn(client *http.Client, tr *http.Transport) {
	var err error
	var resp *http.Response
	var req *http.Request
	for { //use refresh interval to occasionally reconnect to Sidecar in case the stream connection is lost
		req, err = http.NewRequest("GET", provider.Endpoint+"/watch", nil)
		if err != nil {
			log.Errorf("Error creating http request to Sidecar: %s, Error: %s", provider.Endpoint, err)
			continue
		}
		resp, err = client.Do(req)
		if err != nil {
			log.Errorf("Error connecting to Sidecar: %s, Error: %s", provider.Endpoint, err)
			time.Sleep(5 * time.Second)
			continue
		}
		go catalog.DecodeStream(resp.Body, provider.callbackLoader)

		//wait on refresh connection timer.  If this expires we haven't seen an update in a
		//while and should cancel the request, reset the time, and reconnect just in case
		<-provider.connTimer.C
		provider.connTimer.Reset(provider.RefreshConn * time.Second)
		tr.CancelRequest(req)
	}
}

func (provider *Sidecar) callbackLoader(sidecarStates map[string][]*service.Service, err error) error {
	if err != nil {
		return err
	}
	//reset refresh connection timer
	if !provider.connTimer.Stop() {
		<-provider.connTimer.C
	}
	provider.connTimer.Reset(provider.RefreshConn * time.Second)
	provider.loadSidecarConfig(sidecarStates)
	return nil
}

func (provider *Sidecar) makeFrontend() (map[string]*types.Frontend, error) {
	configuration := new(types.Configuration)
	if _, err := toml.DecodeFile(provider.Frontend, configuration); err != nil {
		log.Errorf("Error reading file: %s", err)
		return nil, err
	}
	return configuration.Frontends, nil
}

func (provider *Sidecar) makeBackends(sidecarStates map[string][]*service.Service) map[string]*types.Backend {
	sidecarBacks := make(map[string]*types.Backend)
	for serviceName, services := range sidecarStates {
		newServers := make(map[string]types.Server)
		newBackend := &types.Backend{LoadBalancer: &types.LoadBalancer{Method: method, Sticky: sticky},
			MaxConn:        &types.MaxConn{Amount: maxConnAmount, ExtractorFunc: maxConnExtractorFunc},
			CircuitBreaker: &types.CircuitBreaker{Expression: circuitBreaker},
			Servers:        newServers}
		for _, serv := range services {
			if serv.IsAlive() {
				for i := 0; i < len(serv.Ports); i++ {
					ipAddr, err := net.LookupIP(serv.Hostname)
					if err != nil {
						log.Errorln("Error resolving Ip address, ", err)
						newBackend.Servers[serv.Hostname] = types.Server{URL: "http://" + serv.Hostname + ":" + strconv.FormatInt(serv.Ports[i].Port, 10), Weight: weight}
					} else {
						newBackend.Servers[serv.Hostname] = types.Server{URL: "http://" + ipAddr[0].String() + ":" + strconv.FormatInt(serv.Ports[i].Port, 10), Weight: weight}
					}
				}
			}
		}
		sidecarBacks[serviceName] = newBackend
	}
	return sidecarBacks
}

func (provider *Sidecar) fetchState() (*catalog.ServicesState, error) {
	client := &http.Client{Timeout: 5 * time.Second}
	resp, err := client.Get(provider.Endpoint + "/state.json")
	if err != nil {
		return nil, err
	}

	bytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	state, err := catalog.Decode(bytes)
	if err != nil {
		return nil, err
	}
	return state, nil
}
<|endoftext|>


<commit_before>package providers

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strconv"
	"strings"
)

type GitHubProvider struct {
	*ProviderData
	Org  string
	Team string
}

func NewGitHubProvider(p *ProviderData) *GitHubProvider {
	p.ProviderName = "GitHub"
	if p.LoginURL == nil || p.LoginURL.String() == "" {
		p.LoginURL = &url.URL{
			Scheme: "https",
			Host:   "github.com",
			Path:   "/login/oauth/authorize",
		}
	}
	if p.RedeemURL == nil || p.RedeemURL.String() == "" {
		p.RedeemURL = &url.URL{
			Scheme: "https",
			Host:   "github.com",
			Path:   "/login/oauth/access_token",
		}
	}
	// ValidationURL is the API Base URL
	if p.ValidateURL == nil || p.ValidateURL.String() == "" {
		p.ValidateURL = &url.URL{
			Scheme: "https",
			Host:   "api.github.com",
			Path:   "/",
		}
	}
	if p.Scope == "" {
		p.Scope = "user:email"
	}
	return &GitHubProvider{ProviderData: p}
}
func (p *GitHubProvider) SetOrgTeam(org, team string) {
	p.Org = org
	p.Team = team
	if org != "" || team != "" {
		p.Scope += " read:org"
	}
}

func (p *GitHubProvider) hasOrg(accessToken string) (bool, error) {
	// https://developer.github.com/v3/orgs/#list-your-organizations

	var orgs []struct {
		Login string `json:"login"`
	}

	type orgsPage []struct {
		Login string `json:"login"`
	}

	pn := 1
	for {
		params := url.Values{
			"limit": {"100"},
			"page":  {strconv.Itoa(pn)},
		}

		endpoint := &url.URL{
			Scheme:   p.ValidateURL.Scheme,
			Host:     p.ValidateURL.Host,
			Path:     path.Join(p.ValidateURL.Path, "/user/orgs"),
			RawQuery: params.Encode(),
		}
		req, _ := http.NewRequest("GET", endpoint.String(), nil)
		req.Header.Set("Accept", "application/vnd.github.v3+json")
		req.Header.Set("Authorization", fmt.Sprintf("token %s", accessToken))
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return false, err
		}

		body, err := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			return false, err
		}
		if resp.StatusCode != 200 {
			return false, fmt.Errorf(
				"got %d from %q %s", resp.StatusCode, endpoint.String(), body)
		}

		var op orgsPage
		if err := json.Unmarshal(body, &op); err != nil {
			return false, err
		}
		if len(op) == 0 {
			break
		}

		orgs = append(orgs, op...)
		pn += 1
	}

	var presentOrgs []string
	for _, org := range orgs {
		if p.Org == org.Login {
			log.Printf("Found Github Organization: %q", org.Login)
			return true, nil
		}
		presentOrgs = append(presentOrgs, org.Login)
	}

	log.Printf("Missing Organization:%q in %v", p.Org, presentOrgs)
	return false, nil
}

func (p *GitHubProvider) hasOrgAndTeam(accessToken string) (bool, error) {
	// https://developer.github.com/v3/orgs/teams/#list-user-teams

	var teams []struct {
		Name string `json:"name"`
		Slug string `json:"slug"`
		Org  struct {
			Login string `json:"login"`
		} `json:"organization"`
	}

	params := url.Values{
		"limit": {"100"},
	}

	endpoint := &url.URL{
		Scheme:   p.ValidateURL.Scheme,
		Host:     p.ValidateURL.Host,
		Path:     path.Join(p.ValidateURL.Path, "/user/teams"),
		RawQuery: params.Encode(),
	}
	team_url := endpoint.String()

	pattern := regexp.MustCompile(`<([^>]+)>; rel="next"`)
	var hasOrg bool
	presentOrgs := make(map[string]bool)
	var presentTeams []string
	for {
		req, _ := http.NewRequest("GET", team_url, nil)
		req.Header.Set("Accept", "application/vnd.github.hellcat-preview+json")
		req.Header.Set("Authorization", fmt.Sprintf("token %s", accessToken))
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return false, err
		}

		body, err := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			return false, err
		}
		if resp.StatusCode != 200 {
			return false, fmt.Errorf(
				"got %d from %q %s", resp.StatusCode, endpoint.String(), body)
		}

		if err := json.Unmarshal(body, &teams); err != nil {
			return false, fmt.Errorf("%s unmarshaling %s", err, body)
		}

		for _, team := range teams {
			presentOrgs[team.Org.Login] = true
			if p.Org == team.Org.Login {
				hasOrg = true
				ts := strings.Split(p.Team, ",")
				for _, t := range ts {
					if t == team.Slug {
						log.Printf("Found Github Organization:%q Team:%q (Name:%q)",
							team.Org.Login, team.Slug, team.Name)
						return true, nil
					}
				}
				presentTeams = append(presentTeams, team.Slug)
			}
		}

		matches := pattern.FindStringSubmatch(resp.Header["Link"][0])
		if len(matches) == 0 {
			break
		}
		team_url = matches[1]
	}

	if hasOrg {
		log.Printf("Missing Team:%q from Org:%q in teams: %v", p.Team, p.Org, presentTeams)
	} else {
		var allOrgs []string
		for org, _ := range presentOrgs {
			allOrgs = append(allOrgs, org)
		}
		log.Printf("Missing Organization:%q in %#v", p.Org, allOrgs)
	}
	return false, nil
}

func (p *GitHubProvider) GetEmailAddress(s *SessionState) (string, error) {

	var emails []struct {
		Email   string `json:"email"`
		Primary bool   `json:"primary"`
	}

	// if we require an Org or Team, check that first
	if p.Org != "" {
		if p.Team != "" {
			if ok, err := p.hasOrgAndTeam(s.AccessToken); err != nil || !ok {
				return "", err
			}
		} else {
			if ok, err := p.hasOrg(s.AccessToken); err != nil || !ok {
				return "", err
			}
		}
	}

	endpoint := &url.URL{
		Scheme: p.ValidateURL.Scheme,
		Host:   p.ValidateURL.Host,
		Path:   path.Join(p.ValidateURL.Path, "/user/emails"),
	}
	req, _ := http.NewRequest("GET", endpoint.String(), nil)
	req.Header.Set("Authorization", fmt.Sprintf("token %s", s.AccessToken))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("got %d from %q %s",
			resp.StatusCode, endpoint.String(), body)
	}

	log.Printf("got %d from %q %s", resp.StatusCode, endpoint.String(), body)

	if err := json.Unmarshal(body, &emails); err != nil {
		return "", fmt.Errorf("%s unmarshaling %s", err, body)
	}

	for _, email := range emails {
		if email.Primary {
			return email.Email, nil
		}
	}

	return "", nil
}

func (p *GitHubProvider) GetUserName(s *SessionState) (string, error) {
	var user struct {
		Login string `json:"login"`
		Email string `json:"email"`
	}

	endpoint := &url.URL{
		Scheme: p.ValidateURL.Scheme,
		Host:   p.ValidateURL.Host,
		Path:   path.Join(p.ValidateURL.Path, "/user"),
	}

	req, err := http.NewRequest("GET", endpoint.String(), nil)
	if err != nil {
		return "", fmt.Errorf("could not create new GET request: %v", err)
	}

	req.Header.Set("Authorization", fmt.Sprintf("token %s", s.AccessToken))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	body, err := ioutil.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("got %d from %q %s",
			resp.StatusCode, endpoint.String(), body)
	}

	log.Printf("got %d from %q %s", resp.StatusCode, endpoint.String(), body)

	if err := json.Unmarshal(body, &user); err != nil {
		return "", fmt.Errorf("%s unmarshaling %s", err, body)
	}

	return user.Login, nil
}
<commit_msg>GitHub provider: limit to 10 pages of organizations/teams<commit_after>package providers

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strconv"
	"strings"
)

type GitHubProvider struct {
	*ProviderData
	Org  string
	Team string
}

func NewGitHubProvider(p *ProviderData) *GitHubProvider {
	p.ProviderName = "GitHub"
	if p.LoginURL == nil || p.LoginURL.String() == "" {
		p.LoginURL = &url.URL{
			Scheme: "https",
			Host:   "github.com",
			Path:   "/login/oauth/authorize",
		}
	}
	if p.RedeemURL == nil || p.RedeemURL.String() == "" {
		p.RedeemURL = &url.URL{
			Scheme: "https",
			Host:   "github.com",
			Path:   "/login/oauth/access_token",
		}
	}
	// ValidationURL is the API Base URL
	if p.ValidateURL == nil || p.ValidateURL.String() == "" {
		p.ValidateURL = &url.URL{
			Scheme: "https",
			Host:   "api.github.com",
			Path:   "/",
		}
	}
	if p.Scope == "" {
		p.Scope = "user:email"
	}
	return &GitHubProvider{ProviderData: p}
}
func (p *GitHubProvider) SetOrgTeam(org, team string) {
	p.Org = org
	p.Team = team
	if org != "" || team != "" {
		p.Scope += " read:org"
	}
}

func (p *GitHubProvider) hasOrg(accessToken string) (bool, error) {
	// https://developer.github.com/v3/orgs/#list-your-organizations
	var orgs []struct {
		Login string `json:"login"`
	}
	type orgsPage []struct {
		Login string `json:"login"`
	}

	for pn := 1; pn <= 10; pn++ {
		params := url.Values{
			"limit": {"100"},
			"page":  {strconv.Itoa(pn)},
		}
		endpoint := &url.URL{
			Scheme:   p.ValidateURL.Scheme,
			Host:     p.ValidateURL.Host,
			Path:     path.Join(p.ValidateURL.Path, "/user/orgs"),
			RawQuery: params.Encode(),
		}
		req, _ := http.NewRequest("GET", endpoint.String(), nil)
		req.Header.Set("Accept", "application/vnd.github.v3+json")
		req.Header.Set("Authorization", fmt.Sprintf("token %s", accessToken))
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return false, err
		}

		body, err := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			return false, err
		}
		if resp.StatusCode != 200 {
			return false, fmt.Errorf(
				"got %d from %q %s", resp.StatusCode, endpoint.String(), body)
		}

		var op orgsPage
		if err := json.Unmarshal(body, &op); err != nil {
			return false, err
		}
		if len(op) == 0 {
			break
		}

		orgs = append(orgs, op...)
	}

	var presentOrgs []string
	for _, org := range orgs {
		if p.Org == org.Login {
			log.Printf("Found Github Organization: %q", org.Login)
			return true, nil
		}
		presentOrgs = append(presentOrgs, org.Login)
	}

	log.Printf("Missing Organization:%q in %v", p.Org, presentOrgs)
	return false, nil
}

func (p *GitHubProvider) hasOrgAndTeam(accessToken string) (bool, error) {
	// https://developer.github.com/v3/orgs/teams/#list-user-teams

	var teams []struct {
		Name string `json:"name"`
		Slug string `json:"slug"`
		Org  struct {
			Login string `json:"login"`
		} `json:"organization"`
	}

	params := url.Values{
		"limit": {"100"},
	}
	endpoint := &url.URL{
		Scheme:   p.ValidateURL.Scheme,
		Host:     p.ValidateURL.Host,
		Path:     path.Join(p.ValidateURL.Path, "/user/teams"),
		RawQuery: params.Encode(),
	}
	team_url := endpoint.String()

	pattern := regexp.MustCompile(`<([^>]+)>; rel="next"`)
	var hasOrg bool
	presentOrgs := make(map[string]bool)
	var presentTeams []string

	for i := 0; i < 10; i++ {
		req, _ := http.NewRequest("GET", team_url, nil)
		req.Header.Set("Accept", "application/vnd.github.hellcat-preview+json")
		req.Header.Set("Authorization", fmt.Sprintf("token %s", accessToken))
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return false, err
		}

		body, err := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			return false, err
		}
		if resp.StatusCode != 200 {
			return false, fmt.Errorf(
				"got %d from %q %s", resp.StatusCode, endpoint.String(), body)
		}

		if err := json.Unmarshal(body, &teams); err != nil {
			return false, fmt.Errorf("%s unmarshaling %s", err, body)
		}

		for _, team := range teams {
			presentOrgs[team.Org.Login] = true
			if p.Org == team.Org.Login {
				hasOrg = true
				ts := strings.Split(p.Team, ",")
				for _, t := range ts {
					if t == team.Slug {
						log.Printf("Found Github Organization:%q Team:%q (Name:%q)",
							team.Org.Login, team.Slug, team.Name)
						return true, nil
					}
				}
				presentTeams = append(presentTeams, team.Slug)
			}
		}

		matches := pattern.FindStringSubmatch(resp.Header["Link"][0])
		if len(matches) == 0 {
			break
		}
		team_url = matches[1]
	}

	if hasOrg {
		log.Printf("Missing Team:%q from Org:%q in teams: %v", p.Team, p.Org, presentTeams)
	} else {
		var allOrgs []string
		for org, _ := range presentOrgs {
			allOrgs = append(allOrgs, org)
		}
		log.Printf("Missing Organization:%q in %#v", p.Org, allOrgs)
	}
	return false, nil
}

func (p *GitHubProvider) GetEmailAddress(s *SessionState) (string, error) {

	var emails []struct {
		Email   string `json:"email"`
		Primary bool   `json:"primary"`
	}

	// if we require an Org or Team, check that first
	if p.Org != "" {
		if p.Team != "" {
			if ok, err := p.hasOrgAndTeam(s.AccessToken); err != nil || !ok {
				return "", err
			}
		} else {
			if ok, err := p.hasOrg(s.AccessToken); err != nil || !ok {
				return "", err
			}
		}
	}

	endpoint := &url.URL{
		Scheme: p.ValidateURL.Scheme,
		Host:   p.ValidateURL.Host,
		Path:   path.Join(p.ValidateURL.Path, "/user/emails"),
	}
	req, _ := http.NewRequest("GET", endpoint.String(), nil)
	req.Header.Set("Authorization", fmt.Sprintf("token %s", s.AccessToken))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("got %d from %q %s",
			resp.StatusCode, endpoint.String(), body)
	}

	log.Printf("got %d from %q %s", resp.StatusCode, endpoint.String(), body)

	if err := json.Unmarshal(body, &emails); err != nil {
		return "", fmt.Errorf("%s unmarshaling %s", err, body)
	}

	for _, email := range emails {
		if email.Primary {
			return email.Email, nil
		}
	}

	return "", nil
}

func (p *GitHubProvider) GetUserName(s *SessionState) (string, error) {
	var user struct {
		Login string `json:"login"`
		Email string `json:"email"`
	}

	endpoint := &url.URL{
		Scheme: p.ValidateURL.Scheme,
		Host:   p.ValidateURL.Host,
		Path:   path.Join(p.ValidateURL.Path, "/user"),
	}

	req, err := http.NewRequest("GET", endpoint.String(), nil)
	if err != nil {
		return "", fmt.Errorf("could not create new GET request: %v", err)
	}

	req.Header.Set("Authorization", fmt.Sprintf("token %s", s.AccessToken))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	body, err := ioutil.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("got %d from %q %s",
			resp.StatusCode, endpoint.String(), body)
	}

	log.Printf("got %d from %q %s", resp.StatusCode, endpoint.String(), body)

	if err := json.Unmarshal(body, &user); err != nil {
		return "", fmt.Errorf("%s unmarshaling %s", err, body)
	}

	return user.Login, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/andrewchambers/cc/cpp"
	"github.com/andrewchambers/cc/emit"
	"github.com/andrewchambers/cc/parse"
	"io"
	"os"
)

func printVersion() {
	fmt.Println("cc version 0.01")
}

func printUsage() {
	printVersion()
	fmt.Println()
	fmt.Println("Usage:")
	fmt.Println("  cc [FLAGS] FILE.c")
	fmt.Println()
	fmt.Println("Environment variables:")
	fmt.Println("  CCDEBUG=true enables extended error messages for debugging the compiler.")
	fmt.Println()
	fmt.Println("Flags:")
	flag.PrintDefaults()
	fmt.Println()
	fmt.Println("Software by Andrew Chambers 2014-2015 - andrewchamberss@gmail.com")
}

func compileFile(path string, out io.Writer) error {
	f, err := os.Open(path)
	if err != nil {
		err = fmt.Errorf("Failed to open source file %s for parsing: %s\n", path, err)
		return err
	}
	lexer := cpp.Lex(path, f)
	pp := cpp.New(lexer, nil)
	toplevels, err := parse.Parse(pp)
	if err != nil {
		return err
	}
	return emit.Emit(toplevels, out)
}

func preprocessFile(sourceFile string, out io.Writer) error {
	f, err := os.Open(sourceFile)
	if err != nil {
		return fmt.Errorf("Failed to open source file %s for preprocessing: %s\n", sourceFile, err)
	}
	lexer := cpp.Lex(sourceFile, f)
	pp := cpp.New(lexer, nil)
	for {
		tok, err := pp.Next()
		if err != nil {
			return err
		}
		fmt.Fprintf(out, "%s:%s:%d:%d\n", tok.Kind, tok.Val, tok.Pos.Line, tok.Pos.Col)
		if tok.Kind == cpp.EOF {
			return nil
		}
	}
}

func tokenizeFile(sourceFile string, out io.Writer) error {
	f, err := os.Open(sourceFile)
	if err != nil {
		return fmt.Errorf("Failed to open source file %s for preprocessing: %s\n", sourceFile, err)
	}
	lexer := cpp.Lex(sourceFile, f)
	for {
		tok, err := lexer.Next()
		if err != nil {
			return err
		}
		fmt.Fprintf(out, "%s:%s:%d:%d\n", tok.Kind, tok.Val, tok.Pos.Line, tok.Pos.Col)
		if tok.Kind == cpp.EOF {
			return nil
		}
	}
}

func main() {
	flag.Usage = printUsage
	preprocessOnly := flag.Bool("P", false, "Print tokens after preprocessing (For debugging).")
	tokenizeOnly := flag.Bool("T", false, "Print tokens after lexing (For debugging).")
	version := flag.Bool("version", false, "Print version info and exit.")
	outputPath := flag.String("o", "-", "File to write output to, - for stdout.")
	flag.Parse()

	if *version {
		printVersion()
		return
	}
	if flag.NArg() == 0 {
		printUsage()
		os.Exit(1)
	}
	if flag.NArg() != 1 {
		fmt.Fprintf(os.Stderr, "Bad number of args, please specify a single source file.\n")
		os.Exit(1)
	}

	input := flag.Args()[0]
	var output io.WriteCloser
	var err error

	if *outputPath == "-" {
		output = os.Stdout
	} else {
		output, err = os.Create(*outputPath)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to open output file %s\n", err)
			os.Exit(1)
		}
	}

	if *preprocessOnly {
		err := preprocessFile(input, output)
		reportError(err)
	} else if *tokenizeOnly {
		err := tokenizeFile(input, output)
		reportError(err)
	} else {
		err := compileFile(input, output)
		reportError(err)
	}
}
<commit_msg>profile<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/andrewchambers/cc/cpp"
	"github.com/andrewchambers/cc/emit"
	"github.com/andrewchambers/cc/parse"
	"io"
	"os"
	"runtime/pprof"
)

func printVersion() {
	fmt.Println("cc version 0.01")
}

func printUsage() {
	printVersion()
	fmt.Println()
	fmt.Println("Usage:")
	fmt.Println("  cc [FLAGS] FILE.c")
	fmt.Println()
	fmt.Println("Environment variables:")
	fmt.Println("  CCDEBUG=true enables extended error messages for debugging the compiler.")
	fmt.Println()
	fmt.Println("Flags:")
	flag.PrintDefaults()
	fmt.Println()
	fmt.Println("Software by Andrew Chambers 2014-2015 - andrewchamberss@gmail.com")
}

func compileFile(path string, out io.Writer) error {
	f, err := os.Open(path)
	if err != nil {
		err = fmt.Errorf("Failed to open source file %s for parsing: %s\n", path, err)
		return err
	}
	lexer := cpp.Lex(path, f)
	pp := cpp.New(lexer, nil)
	toplevels, err := parse.Parse(pp)
	if err != nil {
		return err
	}
	return emit.Emit(toplevels, out)
}

func preprocessFile(sourceFile string, out io.Writer) error {
	f, err := os.Open(sourceFile)
	if err != nil {
		return fmt.Errorf("Failed to open source file %s for preprocessing: %s\n", sourceFile, err)
	}
	lexer := cpp.Lex(sourceFile, f)
	pp := cpp.New(lexer, nil)
	for {
		tok, err := pp.Next()
		if err != nil {
			return err
		}
		fmt.Fprintf(out, "%s:%s:%d:%d\n", tok.Kind, tok.Val, tok.Pos.Line, tok.Pos.Col)
		if tok.Kind == cpp.EOF {
			return nil
		}
	}
}

func tokenizeFile(sourceFile string, out io.Writer) error {
	f, err := os.Open(sourceFile)
	if err != nil {
		return fmt.Errorf("Failed to open source file %s for preprocessing: %s\n", sourceFile, err)
	}
	lexer := cpp.Lex(sourceFile, f)
	for {
		tok, err := lexer.Next()
		if err != nil {
			return err
		}
		fmt.Fprintf(out, "%s:%s:%d:%d\n", tok.Kind, tok.Val, tok.Pos.Line, tok.Pos.Col)
		if tok.Kind == cpp.EOF {
			return nil
		}
	}
}

func main() {
	flag.Usage = printUsage
	preprocessOnly := flag.Bool("P", false, "Print tokens after preprocessing (For debugging).")
	tokenizeOnly := flag.Bool("T", false, "Print tokens after lexing (For debugging).")
	version := flag.Bool("version", false, "Print version info and exit.")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
	outputPath := flag.String("o", "-", "File to write output to, - for stdout.")
	flag.Parse()

	if *version {
		printVersion()
		return
	}
	if flag.NArg() == 0 {
		printUsage()
		os.Exit(1)
	}
	if flag.NArg() != 1 {
		fmt.Fprintf(os.Stderr, "Bad number of args, please specify a single source file.\n")
		os.Exit(1)
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to open cpu profile file %s\n", err)
			os.Exit(1)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	input := flag.Args()[0]
	var output io.WriteCloser
	var err error

	if *outputPath == "-" {
		output = os.Stdout
	} else {
		output, err = os.Create(*outputPath)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to open output file %s\n", err)
			os.Exit(1)
		}
	}

	if *preprocessOnly {
		err := preprocessFile(input, output)
		reportError(err)
	} else if *tokenizeOnly {
		err := tokenizeFile(input, output)
		reportError(err)
	} else {
		err := compileFile(input, output)
		reportError(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build linux freebsd

package install

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"regexp"

	"github.com/keybase/client/go/libkb"
)

// Similar to the Brew install on OSX, the Unix install happens in two steps.
// First, the system package manager installs all the binaries as root. Second,
// an autostart file needs to be written to the user's home dir, so that
// Keybase launches when that user logs in. The second step is done the first
// time the user starts Keybase.
//
// ".desktop" files and the ~/.config/autostart directory are part of the
// freedesktop.org set of standards, which the popular desktop environments
// like Gnome and KDE all support. See
// http://standards.freedesktop.org/desktop-entry-spec/latest/.

const autostartFileText = `# This file is autogenerated when Keybase starts. Keybase will recreate it if you
# delete it, so to disable autostart you should uncomment the line below.

[Desktop Entry]
Name=Keybase
Comment=Keybase Filesystem Service and GUI
Type=Application
Exec=env KEYBASE_START_UI=hideWindow run_keybase

# Uncomment the following line to disable Keybase autostart:
#Hidden=true
`

func autostartDir(context Context) string {
	// strip off the "keybase" folder on the end of the config dir
	return path.Join(context.GetConfigDir(), "..", "autostart")
}

func autostartFilePath(context Context) string {
	return path.Join(autostartDir(context), "keybase_autostart.desktop")
}

// AutoInstall installs auto start on unix
func AutoInstall(context Context, _ string, _ bool, log Log) ( /* newProc */ bool, error) {
	// If the desktop file already exists and has been disabled by the user, short circuit.
	if file, err := os.Open(autostartFilePath(context)); err == nil {
		defer file.Close()
		// Check if the user has disabled autostart manually.
		// If the user inserts Hidden=true, our autostart will be ignored.
		// GNOME inserts X-GNOME-Autostart-enabled=false when startup apps are disabled via gnome-session-properties.
		autostartDisabledRegex := regexp.MustCompile(`(?m)^Hidden=true|^X-GNOME-Autostart-enabled=false`)
		if autostartDisabledRegex.MatchReader(bufio.NewReader(file)) {
			return false, nil
		}
	}

	// If the desktop file doesn't exist or hasn't been disabled by the user, replace it.
	err := os.MkdirAll(autostartDir(context), 0755)
	if err != nil {
		return false, err
	}
	return false, ioutil.WriteFile(autostartFilePath(context), []byte(autostartFileText), 0644)
}

// CheckIfValidLocation is not used on unix
func CheckIfValidLocation() error {
	return nil
}

// KBFSBinPath returns the path to the KBFS executable
func KBFSBinPath(runMode libkb.RunMode, binPath string) (string, error) {
	return kbfsBinPathDefault(runMode, binPath)
}

// RunAfterStartup is not used on unix
func RunAfterStartup(context Context, isService bool, log Log) error {
	return nil
}

// kbfsBinName returns the name for the KBFS executable
func kbfsBinName() string {
	return "kbfsfuse"
}

func updaterBinName() (string, error) {
	return "", fmt.Errorf("Updater isn't supported on unix")
}

// RunApp starts the app
func RunApp(context Context, log Log) error {
	// TODO: Start app, see run_keybase: /opt/keybase/Keybase &>> "$logdir/Keybase.app.log"
	return nil
}
<commit_msg>Don't clobber Linux autostart if disabled by commenting Exec<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build linux freebsd

package install

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"regexp"

	"github.com/keybase/client/go/libkb"
)

// Similar to the Brew install on OSX, the Unix install happens in two steps.
// First, the system package manager installs all the binaries as root. Second,
// an autostart file needs to be written to the user's home dir, so that
// Keybase launches when that user logs in. The second step is done the first
// time the user starts Keybase.
//
// ".desktop" files and the ~/.config/autostart directory are part of the
// freedesktop.org set of standards, which the popular desktop environments
// like Gnome and KDE all support. See
// http://standards.freedesktop.org/desktop-entry-spec/latest/.

const autostartFileText = `# This file is autogenerated when Keybase starts. Keybase will recreate it if you
# delete it, so to disable autostart you should uncomment the line below.

[Desktop Entry]
Name=Keybase
Comment=Keybase Filesystem Service and GUI
Type=Application
Exec=env KEYBASE_START_UI=hideWindow run_keybase

# Uncomment the following line to disable Keybase autostart:
#Hidden=true
`

func autostartDir(context Context) string {
	// strip off the "keybase" folder on the end of the config dir
	return path.Join(context.GetConfigDir(), "..", "autostart")
}

func autostartFilePath(context Context) string {
	return path.Join(autostartDir(context), "keybase_autostart.desktop")
}

// AutoInstall installs auto start on unix
func AutoInstall(context Context, _ string, _ bool, log Log) ( /* newProc */ bool, error) {
	// If the desktop file already exists and has been disabled by the user, short circuit.
	if file, err := os.Open(autostartFilePath(context)); err == nil {
		defer file.Close()
		// Check if the user has disabled autostart manually.
		// If the user inserts Hidden=true, our autostart will be ignored.
		// GNOME inserts X-GNOME-Autostart-enabled=false when startup apps are disabled via gnome-session-properties.
		// In past versions, users have been instructed to disable autostart by commenting out the "Exec" line.
		autostartDisabledRegex := regexp.MustCompile(`(?m)^Hidden=true|^X-GNOME-Autostart-enabled=false|^#\s*Exec`)
		if autostartDisabledRegex.MatchReader(bufio.NewReader(file)) {
			return false, nil
		}
	}

	// If the desktop file doesn't exist or hasn't been disabled by the user, replace it.
	err := os.MkdirAll(autostartDir(context), 0755)
	if err != nil {
		return false, err
	}
	return false, ioutil.WriteFile(autostartFilePath(context), []byte(autostartFileText), 0644)
}

// CheckIfValidLocation is not used on unix
func CheckIfValidLocation() error {
	return nil
}

// KBFSBinPath returns the path to the KBFS executable
func KBFSBinPath(runMode libkb.RunMode, binPath string) (string, error) {
	return kbfsBinPathDefault(runMode, binPath)
}

// RunAfterStartup is not used on unix
func RunAfterStartup(context Context, isService bool, log Log) error {
	return nil
}

// kbfsBinName returns the name for the KBFS executable
func kbfsBinName() string {
	return "kbfsfuse"
}

func updaterBinName() (string, error) {
	return "", fmt.Errorf("Updater isn't supported on unix")
}

// RunApp starts the app
func RunApp(context Context, log Log) error {
	// TODO: Start app, see run_keybase: /opt/keybase/Keybase &>> "$logdir/Keybase.app.log"
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package wrangler

import (
	"time"

	"code.google.com/p/vitess/go/relog"
	tm "code.google.com/p/vitess/go/vt/tabletmanager"
	"code.google.com/p/vitess/go/zk"
)

const (
	DefaultActionTimeout = 30 * time.Second
	DefaultLockTimeout   = 30 * time.Second
)

type Wrangler struct {
	zconn       zk.Conn
	ai          *tm.ActionInitiator
	deadline    time.Time
	lockTimeout time.Duration
}

// actionTimeout: how long should we wait for an action to complete?
// lockTimeout: how long should we wait for the initial lock to start a complex action?
//   This is distinct from actionTimeout because most of the time, we want to immediately
//   know that out action will fail. However, automated action will need some time to
//   arbitrate the locks.
func NewWrangler(zconn zk.Conn, actionTimeout, lockTimeout time.Duration) *Wrangler {
	return &Wrangler{zconn, tm.NewActionInitiator(zconn), time.Now().Add(actionTimeout), lockTimeout}
}

func (wr *Wrangler) actionTimeout() time.Duration {
	return wr.deadline.Sub(time.Now())
}

func (wr *Wrangler) readTablet(zkTabletPath string) (*tm.TabletInfo, error) {
	return tm.ReadTablet(wr.zconn, zkTabletPath)
}

func (wr *Wrangler) ZkConn() zk.Conn {
	return wr.zconn
}

func (wr *Wrangler) ActionInitiator() *tm.ActionInitiator {
	return wr.ai
}

// Change the type of tablet and recompute all necessary derived paths in the
// serving graph.
// force: Bypass the vtaction system and make the data change directly, and
// do not run the idle_server_check nor live_server_check hooks
func (wr *Wrangler) ChangeType(zkTabletPath string, dbType tm.TabletType, force bool) error {
	// Load tablet to find keyspace and shard assignment.
	// Don't load after the ChangeType which might have unassigned
	// the tablet.
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	if force {
		// with --force, we do not run any hook
		err = tm.ChangeType(wr.zconn, zkTabletPath, dbType)
	} else {
		// the remote action will run the hooks
		actionPath, err := wr.ai.ChangeType(zkTabletPath, dbType)
		// You don't have a choice - you must wait for
		// completion before rebuilding.
		if err == nil {
			err = wr.ai.WaitForCompletion(actionPath, DefaultActionTimeout)
		}
	}

	if err != nil {
		return err
	}

	// we rebuild if the tablet was serving, or if it is now
	var shardToRebuild string
	if rebuildRequired {
		shardToRebuild = ti.ShardPath()
	} else {
		// re-read the tablet, see if we become serving
		ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
		if err != nil {
			return err
		}
		if ti.Tablet.IsServingType() {
			rebuildRequired = true
			shardToRebuild = ti.ShardPath()
		}
	}

	if rebuildRequired {
		if _, err := wr.RebuildShardGraph(shardToRebuild); err != nil {
			return err
		}
	}
	return nil
}

// same as ChangeType, but assume we already have the shard lock,
// and do not have the option to force anything
// FIXME(alainjobart): doesn't rebuild the Keyspace, as that part has locks,
// so the local serving graphs will be wrong. To do that, I need to refactor
// some code, might be a bigger change.
// Mike says: Updating the shard should be good enough. I'm debating dropping the entire
// keyspace rollup, since I think that is adding complexity and feels like it might
// be a premature optimization.
func (wr *Wrangler) changeTypeInternal(zkTabletPath string, dbType tm.TabletType) error {
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	// change the type
	actionPath, err := wr.ai.ChangeType(ti.Path(), dbType)
	if err != nil {
		return err
	}
	err = wr.ai.WaitForCompletion(actionPath, wr.actionTimeout())
	if err != nil {
		return err
	}

	// rebuild if necessary
	if rebuildRequired {
		err = wr.rebuildShard(ti.ShardPath(), false)
		if err != nil {
			return err
		}
		// FIXME(alainjobart) We already have the lock on one shard, so this is not
		// possible. But maybe it's not necessary anyway.
		// We could pass in a shard path we already have the lock on, and skip it?
		//		err = wr.rebuildKeyspace(ti.KeyspacePath())
		//		if err != nil {
		//			return err
		//		}
	}
	return nil
}

// Cleanup an action node and write back status/error to zk.
// Only returns an error if something went wrong with zk.
func (wr *Wrangler) handleActionError(actionPath string, actionErr error) error {
	// re-read the action node
	data, _, err := wr.zconn.Get(actionPath)
	if err != nil {
		return err
	}
	var actionNode *tm.ActionNode
	actionNode, err = tm.ActionNodeFromJson(data, actionPath)
	if err != nil {
		return err
	}

	// write what happened to the action log
	err = tm.StoreActionResponse(wr.zconn, actionNode, actionPath, actionErr)
	if err != nil {
		return err
	}

	// no error, we can unblock the action queue
	if actionErr == nil {
		return zk.DeleteRecursive(wr.zconn, actionPath, -1)
	}
	return nil
}

// Scrap a tablet. If force is used, we write to ZK directly and don't
// remote-execute the command.
func (wr *Wrangler) Scrap(zkTabletPath string, force, skipRebuild bool) (actionPath string, err error) {
	// load the tablet, see if we'll need to rebuild
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return "", err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	if force {
		err = tm.Scrap(wr.zconn, zkTabletPath, force)
	} else {
		actionPath, err = wr.ai.Scrap(zkTabletPath)
	}

	if !rebuildRequired {
		relog.Info("Rebuild not required")
		return
	}
	if skipRebuild {
		relog.Warning("Rebuild required, but skipping it")
		return
	}

	// wait for the remote Scrap if necessary
	if actionPath != "" {
		err = wr.ai.WaitForCompletion(actionPath, wr.actionTimeout())
		if err != nil {
			return "", nil
		}
	}

	// and rebuild the original shard / keyspace
	return wr.RebuildShardGraph(ti.ShardPath())
}
<commit_msg>Fixing an error reporting bug.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package wrangler

import (
	"time"

	"code.google.com/p/vitess/go/relog"
	tm "code.google.com/p/vitess/go/vt/tabletmanager"
	"code.google.com/p/vitess/go/zk"
)

const (
	DefaultActionTimeout = 30 * time.Second
	DefaultLockTimeout   = 30 * time.Second
)

type Wrangler struct {
	zconn       zk.Conn
	ai          *tm.ActionInitiator
	deadline    time.Time
	lockTimeout time.Duration
}

// actionTimeout: how long should we wait for an action to complete?
// lockTimeout: how long should we wait for the initial lock to start a complex action?
//   This is distinct from actionTimeout because most of the time, we want to immediately
//   know that out action will fail. However, automated action will need some time to
//   arbitrate the locks.
func NewWrangler(zconn zk.Conn, actionTimeout, lockTimeout time.Duration) *Wrangler {
	return &Wrangler{zconn, tm.NewActionInitiator(zconn), time.Now().Add(actionTimeout), lockTimeout}
}

func (wr *Wrangler) actionTimeout() time.Duration {
	return wr.deadline.Sub(time.Now())
}

func (wr *Wrangler) readTablet(zkTabletPath string) (*tm.TabletInfo, error) {
	return tm.ReadTablet(wr.zconn, zkTabletPath)
}

func (wr *Wrangler) ZkConn() zk.Conn {
	return wr.zconn
}

func (wr *Wrangler) ActionInitiator() *tm.ActionInitiator {
	return wr.ai
}

// Change the type of tablet and recompute all necessary derived paths in the
// serving graph.
// force: Bypass the vtaction system and make the data change directly, and
// do not run the idle_server_check nor live_server_check hooks
func (wr *Wrangler) ChangeType(zkTabletPath string, dbType tm.TabletType, force bool) error {
	// Load tablet to find keyspace and shard assignment.
	// Don't load after the ChangeType which might have unassigned
	// the tablet.
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	if force {
		// with --force, we do not run any hook
		err = tm.ChangeType(wr.zconn, zkTabletPath, dbType)
	} else {
		// the remote action will run the hooks
		var actionPath string
		actionPath, err = wr.ai.ChangeType(zkTabletPath, dbType)
		// You don't have a choice - you must wait for
		// completion before rebuilding.
		if err == nil {
			err = wr.ai.WaitForCompletion(actionPath, DefaultActionTimeout)
		}
	}

	if err != nil {
		return err
	}

	// we rebuild if the tablet was serving, or if it is now
	var shardToRebuild string
	if rebuildRequired {
		shardToRebuild = ti.ShardPath()
	} else {
		// re-read the tablet, see if we become serving
		ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
		if err != nil {
			return err
		}
		if ti.Tablet.IsServingType() {
			rebuildRequired = true
			shardToRebuild = ti.ShardPath()
		}
	}

	if rebuildRequired {
		if _, err := wr.RebuildShardGraph(shardToRebuild); err != nil {
			return err
		}
	}
	return nil
}

// same as ChangeType, but assume we already have the shard lock,
// and do not have the option to force anything
// FIXME(alainjobart): doesn't rebuild the Keyspace, as that part has locks,
// so the local serving graphs will be wrong. To do that, I need to refactor
// some code, might be a bigger change.
// Mike says: Updating the shard should be good enough. I'm debating dropping the entire
// keyspace rollup, since I think that is adding complexity and feels like it might
// be a premature optimization.
func (wr *Wrangler) changeTypeInternal(zkTabletPath string, dbType tm.TabletType) error {
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	// change the type
	actionPath, err := wr.ai.ChangeType(ti.Path(), dbType)
	if err != nil {
		return err
	}
	err = wr.ai.WaitForCompletion(actionPath, wr.actionTimeout())
	if err != nil {
		return err
	}

	// rebuild if necessary
	if rebuildRequired {
		err = wr.rebuildShard(ti.ShardPath(), false)
		if err != nil {
			return err
		}
		// FIXME(alainjobart) We already have the lock on one shard, so this is not
		// possible. But maybe it's not necessary anyway.
		// We could pass in a shard path we already have the lock on, and skip it?
		//		err = wr.rebuildKeyspace(ti.KeyspacePath())
		//		if err != nil {
		//			return err
		//		}
	}
	return nil
}

// Cleanup an action node and write back status/error to zk.
// Only returns an error if something went wrong with zk.
func (wr *Wrangler) handleActionError(actionPath string, actionErr error) error {
	// re-read the action node
	data, _, err := wr.zconn.Get(actionPath)
	if err != nil {
		return err
	}
	var actionNode *tm.ActionNode
	actionNode, err = tm.ActionNodeFromJson(data, actionPath)
	if err != nil {
		return err
	}

	// write what happened to the action log
	err = tm.StoreActionResponse(wr.zconn, actionNode, actionPath, actionErr)
	if err != nil {
		return err
	}

	// no error, we can unblock the action queue
	if actionErr == nil {
		return zk.DeleteRecursive(wr.zconn, actionPath, -1)
	}
	return nil
}

// Scrap a tablet. If force is used, we write to ZK directly and don't
// remote-execute the command.
func (wr *Wrangler) Scrap(zkTabletPath string, force, skipRebuild bool) (actionPath string, err error) {
	// load the tablet, see if we'll need to rebuild
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return "", err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	if force {
		err = tm.Scrap(wr.zconn, zkTabletPath, force)
	} else {
		actionPath, err = wr.ai.Scrap(zkTabletPath)
	}

	if !rebuildRequired {
		relog.Info("Rebuild not required")
		return
	}
	if skipRebuild {
		relog.Warning("Rebuild required, but skipping it")
		return
	}

	// wait for the remote Scrap if necessary
	if actionPath != "" {
		err = wr.ai.WaitForCompletion(actionPath, wr.actionTimeout())
		if err != nil {
			return "", nil
		}
	}

	// and rebuild the original shard / keyspace
	return wr.RebuildShardGraph(ti.ShardPath())
}
<|endoftext|>


<commit_before>package gobusterdir

import (
	"bytes"
	"fmt"
	"log"

	"github.com/OJ/gobuster/libgobuster"
	uuid "github.com/satori/go.uuid"
)

// SetupDir is the setup implementation of gobusterdir
func SetupDir(g *libgobuster.Gobuster) error {
	_, _, err := g.GetRequest(g.Opts.URL)
	if err != nil {
		return fmt.Errorf("unable to connect to %s: %v", g.Opts.URL, err)
	}

	guid := uuid.Must(uuid.NewV4())
	url := fmt.Sprintf("%s%s", g.Opts.URL, guid)
	wildcardResp, _, err := g.GetRequest(url)

	if err != nil {
		return err
	}

	if g.Opts.StatusCodesParsed.Contains(*wildcardResp) {
		g.IsWildcard = true
		log.Printf("[-] Wildcard response found: %s => %d", url, *wildcardResp)
		if !g.Opts.WildcardForced {
			return fmt.Errorf("[-] To force processing of Wildcard responses, specify the '-fw' switch.")
		}
	}

	return nil
}

// ProcessDirEntry is the process implementation of gobusterdir
func ProcessDirEntry(g *libgobuster.Gobuster, word string) ([]libgobuster.Result, error) {
	suffix := ""
	if g.Opts.UseSlash {
		suffix = "/"
	}

	// Try the DIR first
	url := fmt.Sprintf("%s%s%s", g.Opts.URL, word, suffix)
	dirResp, dirSize, err := g.GetRequest(url)
	if err != nil {
		return nil, fmt.Errorf("can not call URL %s: %v", url, err)
	}
	var ret []libgobuster.Result
	if dirResp != nil {
		ret = append(ret, libgobuster.Result{
			Entity: fmt.Sprintf("%s%s", word, suffix),
			Status: *dirResp,
			Size:   dirSize,
		})
	}

	// Follow up with files using each ext.
	for _, ext := range g.Opts.ExtensionsParsed {
		file := fmt.Sprintf("%s.%s", word, ext)
		url = fmt.Sprintf("%s%s", g.Opts.URL, file)
		fileResp, fileSize, err := g.GetRequest(url)
		if err != nil {
			return nil, fmt.Errorf("can not call URL %s: %v", url, err)
		}

		if fileResp != nil {
			ret = append(ret, libgobuster.Result{
				Entity: file,
				Status: *fileResp,
				Size:   fileSize,
			})
		}
	}
	return ret, nil
}

// DirResultToString is the to string implementation of gobusterdir
func DirResultToString(g *libgobuster.Gobuster, r *libgobuster.Result) (*string, error) {
	buf := &bytes.Buffer{}

	// Prefix if we're in verbose mode
	if g.Opts.Verbose {
		if g.Opts.StatusCodesParsed.Contains(r.Status) {
			if _, err := fmt.Fprintf(buf, "Found: "); err != nil {
				return nil, err
			}
		} else {
			if _, err := fmt.Fprintf(buf, "Missed: "); err != nil {
				return nil, err
			}
		}
	}

	if g.Opts.StatusCodesParsed.Contains(r.Status) || g.Opts.Verbose {
		if g.Opts.Expanded {
			if _, err := fmt.Fprintf(buf, g.Opts.URL); err != nil {
				return nil, err
			}
		} else {
			if _, err := fmt.Fprintf(buf, "/"); err != nil {
				return nil, err
			}
		}
		if _, err := fmt.Fprintf(buf, r.Entity); err != nil {
			return nil, err
		}

		if !g.Opts.NoStatus {
			if _, err := fmt.Fprintf(buf, " (Status: %d)", r.Status); err != nil {
				return nil, err
			}
		}

		if r.Size != nil {
			if _, err := fmt.Fprintf(buf, " [Size: %d]", *r.Size); err != nil {
				return nil, err
			}
		}
		if _, err := fmt.Fprintf(buf, "\n"); err != nil {
			return nil, err
		}
	}
	s := buf.String()
	return &s, nil
}
<commit_msg>shorter error messages<commit_after>package gobusterdir

import (
	"bytes"
	"fmt"
	"log"

	"github.com/OJ/gobuster/libgobuster"
	uuid "github.com/satori/go.uuid"
)

// SetupDir is the setup implementation of gobusterdir
func SetupDir(g *libgobuster.Gobuster) error {
	_, _, err := g.GetRequest(g.Opts.URL)
	if err != nil {
		return fmt.Errorf("unable to connect to %s: %v", g.Opts.URL, err)
	}

	guid := uuid.Must(uuid.NewV4())
	url := fmt.Sprintf("%s%s", g.Opts.URL, guid)
	wildcardResp, _, err := g.GetRequest(url)

	if err != nil {
		return err
	}

	if g.Opts.StatusCodesParsed.Contains(*wildcardResp) {
		g.IsWildcard = true
		log.Printf("[-] Wildcard response found: %s => %d", url, *wildcardResp)
		if !g.Opts.WildcardForced {
			return fmt.Errorf("[-] To force processing of Wildcard responses, specify the '-fw' switch.")
		}
	}

	return nil
}

// ProcessDirEntry is the process implementation of gobusterdir
func ProcessDirEntry(g *libgobuster.Gobuster, word string) ([]libgobuster.Result, error) {
	suffix := ""
	if g.Opts.UseSlash {
		suffix = "/"
	}

	// Try the DIR first
	url := fmt.Sprintf("%s%s%s", g.Opts.URL, word, suffix)
	dirResp, dirSize, err := g.GetRequest(url)
	if err != nil {
		return nil, err
	}
	var ret []libgobuster.Result
	if dirResp != nil {
		ret = append(ret, libgobuster.Result{
			Entity: fmt.Sprintf("%s%s", word, suffix),
			Status: *dirResp,
			Size:   dirSize,
		})
	}

	// Follow up with files using each ext.
	for _, ext := range g.Opts.ExtensionsParsed {
		file := fmt.Sprintf("%s.%s", word, ext)
		url = fmt.Sprintf("%s%s", g.Opts.URL, file)
		fileResp, fileSize, err := g.GetRequest(url)
		if err != nil {
			return nil, err
		}

		if fileResp != nil {
			ret = append(ret, libgobuster.Result{
				Entity: file,
				Status: *fileResp,
				Size:   fileSize,
			})
		}
	}
	return ret, nil
}

// DirResultToString is the to string implementation of gobusterdir
func DirResultToString(g *libgobuster.Gobuster, r *libgobuster.Result) (*string, error) {
	buf := &bytes.Buffer{}

	// Prefix if we're in verbose mode
	if g.Opts.Verbose {
		if g.Opts.StatusCodesParsed.Contains(r.Status) {
			if _, err := fmt.Fprintf(buf, "Found: "); err != nil {
				return nil, err
			}
		} else {
			if _, err := fmt.Fprintf(buf, "Missed: "); err != nil {
				return nil, err
			}
		}
	}

	if g.Opts.StatusCodesParsed.Contains(r.Status) || g.Opts.Verbose {
		if g.Opts.Expanded {
			if _, err := fmt.Fprintf(buf, g.Opts.URL); err != nil {
				return nil, err
			}
		} else {
			if _, err := fmt.Fprintf(buf, "/"); err != nil {
				return nil, err
			}
		}
		if _, err := fmt.Fprintf(buf, r.Entity); err != nil {
			return nil, err
		}

		if !g.Opts.NoStatus {
			if _, err := fmt.Fprintf(buf, " (Status: %d)", r.Status); err != nil {
				return nil, err
			}
		}

		if r.Size != nil {
			if _, err := fmt.Fprintf(buf, " [Size: %d]", *r.Size); err != nil {
				return nil, err
			}
		}
		if _, err := fmt.Fprintf(buf, "\n"); err != nil {
			return nil, err
		}
	}
	s := buf.String()
	return &s, nil
}
<|endoftext|>


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

package cbft

import (
	"os"
	"testing"
)

// This test is kind of long, but it effectively writes some data,
// checks the size and then reads it back while verifying you can't do
// operations you didn't request in states you didn't allow them.
func TestFileLike(t *testing.T) {
	fn := ",file-like-thing"
	defer os.Remove(fn)

	fs := NewFileService(1)
	defer fs.Close()
	f, err := fs.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_EXCL)
	if err != nil {
		t.Fatalf("Error opening file: %v", err)
	}

	buf := make([]byte, 4096)
	copy(buf, []byte("first write"))

	n, err := f.WriteAt(buf, 8192)
	if err != nil {
		t.Fatalf("Error writing: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short write: %v", n)
	}

	copy(buf, []byte("second write"))

	n, err = f.WriteAt(buf, 32768)
	if err != nil {
		t.Fatalf("Error writing: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short write: %v", n)
	}

	//
	// Now try to read them back
	//

	// An immediate read should fail because we're write only.
	n, err = f.ReadAt(buf, 4096)
	if err != unReadable {
		t.Fatalf("Should've failed read with unReadable, "+
			"got %v and %v bytes instead", err, n)
	}

	// Reopen for reading.
	f.Close()
	f, err = fs.OpenFile(fn, os.O_RDONLY)
	if err != nil {
		t.Fatalf("Error reopening for read: %v", err)
	}

	// Let's throw in a stat
	fi, err := f.Stat()
	if err != nil {
		t.Errorf("Stat failed: %v", err)
	}
	if fi.Size() != 32768+4096 {
		t.Errorf("File size didn't meet our expectations: %v", fi.Size())
	}

	// And a write should fail
	n, err = f.WriteAt(buf, 4096)
	if err == nil {
		t.Fatalf("Should've failed write, wrote %v bytes instead", n)
	}

	n, err = f.ReadAt(buf, 32768)
	if err != nil {
		t.Fatalf("Error reading data: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short read: %v", n)
	}

	s := string(buf[:len("second write")])
	if s != "second write" {
		t.Fatalf("Misread:  %q", s)
	}

	n, err = f.ReadAt(buf, 8192)
	if err != nil {
		t.Fatalf("Error reading data: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short read: %v", n)
	}

	s = string(buf[:len("first write")])
	if s != "first write" {
		t.Fatalf("Misread:  %q", s)
	}
}

func TestFileLikeRW(t *testing.T) {
	fn := ",file-like-thing"
	defer os.Remove(fn)

	fs := NewFileService(1)
	defer fs.Close()
	f, err := fs.OpenFile(fn, os.O_CREATE|os.O_RDWR|os.O_EXCL)
	if err != nil {
		t.Fatalf("Error opening file: %v", err)
	}

	buf := make([]byte, 4096)
	copy(buf, []byte("a write"))

	n, err := f.WriteAt(buf, 8192)
	if err != nil {
		t.Fatalf("Error writing: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short write: %v", n)
	}

	buf[0] = 'x'

	n, err = f.ReadAt(buf, 8192)
	if err != nil {
		t.Fatalf("Error reading data: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short read: %v", n)
	}

	s := string(buf[:len("a write")])
	if s != "a write" {
		t.Fatalf("Misread:  %q", s)
	}
}
<commit_msg>test Truncate()<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package cbft

import (
	"os"
	"testing"
)

// This test is kind of long, but it effectively writes some data,
// checks the size and then reads it back while verifying you can't do
// operations you didn't request in states you didn't allow them.
func TestFileLike(t *testing.T) {
	fn := ",file-like-thing"
	defer os.Remove(fn)

	fs := NewFileService(1)
	defer fs.Close()
	f, err := fs.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_EXCL)
	if err != nil {
		t.Fatalf("Error opening file: %v", err)
	}

	buf := make([]byte, 4096)
	copy(buf, []byte("first write"))

	n, err := f.WriteAt(buf, 8192)
	if err != nil {
		t.Fatalf("Error writing: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short write: %v", n)
	}

	copy(buf, []byte("second write"))

	n, err = f.WriteAt(buf, 32768)
	if err != nil {
		t.Fatalf("Error writing: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short write: %v", n)
	}

	//
	// Now try to read them back
	//

	// An immediate read should fail because we're write only.
	n, err = f.ReadAt(buf, 4096)
	if err != unReadable {
		t.Fatalf("Should've failed read with unReadable, "+
			"got %v and %v bytes instead", err, n)
	}

	// Reopen for reading.
	f.Close()
	f, err = fs.OpenFile(fn, os.O_RDONLY)
	if err != nil {
		t.Fatalf("Error reopening for read: %v", err)
	}

	// Let's throw in a stat
	fi, err := f.Stat()
	if err != nil {
		t.Errorf("Stat failed: %v", err)
	}
	if fi.Size() != 32768+4096 {
		t.Errorf("File size didn't meet our expectations: %v", fi.Size())
	}

	// And a write should fail
	n, err = f.WriteAt(buf, 4096)
	if err == nil {
		t.Fatalf("Should've failed write, wrote %v bytes instead", n)
	}

	n, err = f.ReadAt(buf, 32768)
	if err != nil {
		t.Fatalf("Error reading data: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short read: %v", n)
	}

	s := string(buf[:len("second write")])
	if s != "second write" {
		t.Fatalf("Misread:  %q", s)
	}

	n, err = f.ReadAt(buf, 8192)
	if err != nil {
		t.Fatalf("Error reading data: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short read: %v", n)
	}

	s = string(buf[:len("first write")])
	if s != "first write" {
		t.Fatalf("Misread:  %q", s)
	}

	err = f.Truncate(10)
	if err == nil {
		t.Errorf("expected truncate to fail on read-only")
	}
}

func TestFileLikeRW(t *testing.T) {
	fn := ",file-like-thing"
	defer os.Remove(fn)

	fs := NewFileService(1)
	defer fs.Close()
	f, err := fs.OpenFile(fn, os.O_CREATE|os.O_RDWR|os.O_EXCL)
	if err != nil {
		t.Fatalf("Error opening file: %v", err)
	}

	buf := make([]byte, 4096)
	copy(buf, []byte("a write"))

	n, err := f.WriteAt(buf, 8192)
	if err != nil {
		t.Fatalf("Error writing: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short write: %v", n)
	}

	buf[0] = 'x'

	n, err = f.ReadAt(buf, 8192)
	if err != nil {
		t.Fatalf("Error reading data: %v", err)
	}
	if n != 4096 {
		t.Fatalf("Short read: %v", n)
	}

	s := string(buf[:len("a write")])
	if s != "a write" {
		t.Fatalf("Misread:  %q", s)
	}

	err = f.Truncate(10)
	if err != nil {
		t.Errorf("expected truncate to work")
	}
}
<|endoftext|>


<commit_before>package limio

import (
	"errors"
	"io"
	"time"
)

//A Manager enables consumers to treat a group of Limiters as a single Limiter,
//enabling hierarchies of limiters. For example, a network interface could have
//a global limit that is distributed across connections, each of which can
//manage their own distribution of the bandwidth they are allocated.
type Manager interface {
	Limiter
	Manage(Limiter) error
	Unmanage(Limiter)
}

//A SimpleManager is an implementation of the limio.Manager interface. It
//allows simple rate-based and arbitrary channel-based limits to be set.
//
//A SimpleManager is designed so that Limit and Manage may be called
//concurrently.
type SimpleManager struct {
	m map[Limiter]chan int

	newLimit chan *limit
	cls      chan struct{}

	newLimiter chan Limiter
	clsLimiter chan Limiter
}

//NewSimpleManager creates and initializes a SimpleManager.
func NewSimpleManager() *SimpleManager {
	lm := SimpleManager{
		m:          make(map[Limiter]chan int),
		newLimit:   make(chan *limit),
		cls:        make(chan struct{}),
		newLimiter: make(chan Limiter),
		clsLimiter: make(chan Limiter),
	}
	go lm.run()
	return &lm
}

//DefaultWindow is the window used to smooth SimpleLimit rates. That is,
//SimpleLimit distributes the given quantity evenly into buckets of size t.
//This is useful for avoiding tcp silly window syndrome and providing
//predictable resource usage.
const DefaultWindow = 10 * time.Millisecond

func (lm *SimpleManager) run() {
	limited := false
	cl := &limit{}
	ct := &time.Ticker{}

	er := rate{}

	for {
		select {
		case <-ct.C:
			lm.distribute(cl.rate.n)
		case tot := <-cl.lim:
			lm.distribute(tot)
		case newLim := <-lm.newLimit:
			go notify(cl.done, false)
			cl = &limit{}
			ct.Stop()

			if newLim != nil {
				limited = true
				cl = newLim

				for l := range lm.m {
					lm.limit(l)
				}

				if newLim.rate != er && cl.rate.n > 0 {
					cl.rate.n, cl.rate.t = Distribute(cl.rate.n, cl.rate.t, DefaultWindow)
					ct = time.NewTicker(cl.rate.t)
				}
			} else {
				limited = false
				for l := range lm.m {
					l.Unlimit()
				}
			}
		case l := <-lm.newLimiter:
			if limited {
				lm.limit(l)
			} else {
				l.Unlimit()
				lm.m[l] = nil
			}
		case toClose := <-lm.clsLimiter:
			close(lm.m[toClose])
			delete(lm.m, toClose)
		case <-lm.cls:
			for l := range lm.m {
				l.Unlimit()
			}
			go notify(cl.done, true)
			return
		}
	}
}

//NOTE must ONLY be used synchonously with the run() goroutine for concurrency
//safety.
//func distribute(int) takes a number and iterates over each channel in the map of
//managed Limiters, sending an evenly-distriuted limit to each "sublimiter".
//distribute takes a number to distribute and returns the number of bytes
//remaining
func (lm *SimpleManager) distribute(n int) int {
	if len(lm.m) > 0 {
		each := n / len(lm.m)
		for _, ch := range lm.m {
			if ch != nil {
				select {
				case ch <- each:
					n -= each
				default:
					//Skip if not ready
				}
			}
		}
	}
	return n
}

//NOTE must ONLY be used inside of run() for concurrency safety
//limit sets up a new channel for each limiter in the map. It then waits on the
//newly returned bool channel so that limiters can be removed when closed.
func (lm *SimpleManager) limit(l Limiter) {
	lm.m[l] = make(chan int)
	done := l.Limit(lm.m[l])
	go func() {
		//If `true` passed on channel, limiter is closed
		if <-done {
			lm.Unmanage(l)
		}
	}()
}

//NewReader takes an io.Reader and Limits it according to its limit
//policy/strategy
func (lm *SimpleManager) NewReader(r io.Reader) *Reader {
	lr := NewReader(r)
	lm.Manage(lr)
	return lr
}

//SimpleLimit takes an int and time.Duration that will be distributed evenly
//across all managed Limiters.
func (lm *SimpleManager) SimpleLimit(n int, t time.Duration) <-chan bool {
	done := make(chan bool, 1)
	lm.newLimit <- &limit{
		rate: rate{n, t},
		done: done,
	}
	return done
}

//Limit implements the limio.Limiter interface.
func (lm *SimpleManager) Limit(l chan int) <-chan bool {
	done := make(chan bool, 1)
	lm.newLimit <- &limit{
		lim:  l,
		done: done,
	}
	return done
}

//Unlimit implements the limio.Limiter interface.
func (lm *SimpleManager) Unlimit() {
	lm.newLimit <- nil
}

//Close allows the SimpleManager to free any resources it is using if the
//consumer has no further need for the SimpleManager.
func (lm *SimpleManager) Close() error {
	lm.cls <- struct{}{}
	return nil
}

//Unmanage allows consumers to remove a specific Limiter from its management
//strategy
func (lm *SimpleManager) Unmanage(l Limiter) {
	lm.clsLimiter <- l
}

//Manage takes a Limiter that will be adopted under the management policy of
//the SimpleManager.
func (lm *SimpleManager) Manage(l Limiter) error {
	if l == lm {
		return errors.New("A manager cannot manage itself.")
	}

	lm.newLimiter <- l
	return nil
}
<commit_msg>Much closer to having fixed the race condition<commit_after>package limio

import (
	"errors"
	"io"
	"time"
)

//A Manager enables consumers to treat a group of Limiters as a single Limiter,
//enabling hierarchies of limiters. For example, a network interface could have
//a global limit that is distributed across connections, each of which can
//manage their own distribution of the bandwidth they are allocated.
type Manager interface {
	Limiter
	Manage(Limiter) error
	Unmanage(Limiter)
}

//A SimpleManager is an implementation of the limio.Manager interface. It
//allows simple rate-based and arbitrary channel-based limits to be set.
//
//A SimpleManager is designed so that Limit and Manage may be called
//concurrently.
type SimpleManager struct {
	m map[Limiter]chan int

	newLimit chan *limit
	cls      chan struct{}

	newLimiter chan Limiter
	clsLimiter chan Limiter
}

//NewSimpleManager creates and initializes a SimpleManager.
func NewSimpleManager() *SimpleManager {
	lm := SimpleManager{
		m:          make(map[Limiter]chan int),
		newLimit:   make(chan *limit),
		cls:        make(chan struct{}),
		newLimiter: make(chan Limiter),
		clsLimiter: make(chan Limiter),
	}
	go lm.run()
	return &lm
}

//DefaultWindow is the window used to smooth SimpleLimit rates. That is,
//SimpleLimit distributes the given quantity evenly into buckets of size t.
//This is useful for avoiding tcp silly window syndrome and providing
//predictable resource usage.
const DefaultWindow = 10 * time.Millisecond

func (lm *SimpleManager) run() {
	limited := false
	cl := &limit{}
	ct := &time.Ticker{}

	er := rate{}

	for {
		select {
		case <-ct.C:
			lm.distribute(cl.rate.n)
		case tot := <-cl.lim:
			lm.distribute(tot)
		case newLim := <-lm.newLimit:
			go notify(cl.done, false)
			cl = &limit{}
			ct.Stop()

			if newLim != nil {
				limited = true
				cl = newLim

				for l := range lm.m {
					lm.limit(l)
				}

				if newLim.rate != er && cl.rate.n > 0 {
					cl.rate.n, cl.rate.t = Distribute(cl.rate.n, cl.rate.t, DefaultWindow)
					ct = time.NewTicker(cl.rate.t)
				}
			} else {
				limited = false
				for l := range lm.m {
					l.Unlimit()
				}
			}
		case l := <-lm.newLimiter:
			if limited {
				lm.limit(l)
			} else {
				l.Unlimit()
				lm.m[l] = nil
			}
		case toClose := <-lm.clsLimiter:
			close(lm.m[toClose])
			delete(lm.m, toClose)
		case <-lm.cls:
			for l := range lm.m {
				l.Unlimit()
			}
			go notify(cl.done, true)
			return
		}
	}
}

//NOTE must ONLY be used synchonously with the run() goroutine for concurrency
//safety.
//func distribute(int) takes a number and iterates over each channel in the map of
//managed Limiters, sending an evenly-distriuted limit to each "sublimiter".
//distribute takes a number to distribute and returns the number of bytes
//remaining
func (lm *SimpleManager) distribute(n int) int {
	if len(lm.m) > 0 {

		each := n / len(lm.m)

		cp := map[Limiter]chan int{}
		for k, v := range lm.m {
			cp[k] = v
		}

		for n >= each && len(cp) > 0 {
			for k, ch := range cp {
				if ch != nil {
					select {
					case ch <- each:
						n -= each
						delete(cp, k)
					default:
						//Skip if not ready; come back
					}
				}
			}
		}
	}
	return n
}

//NOTE must ONLY be used inside of run() for concurrency safety
//limit sets up a new channel for each limiter in the map. It then waits on the
//newly returned bool channel so that limiters can be removed when closed.
func (lm *SimpleManager) limit(l Limiter) {
	lm.m[l] = make(chan int)
	done := l.Limit(lm.m[l])
	go func() {
		//If `true` passed on channel, limiter is closed
		if <-done {
			lm.Unmanage(l)
		}
	}()
}

//NewReader takes an io.Reader and Limits it according to its limit
//policy/strategy
func (lm *SimpleManager) NewReader(r io.Reader) *Reader {
	lr := NewReader(r)
	lm.Manage(lr)
	return lr
}

//SimpleLimit takes an int and time.Duration that will be distributed evenly
//across all managed Limiters.
func (lm *SimpleManager) SimpleLimit(n int, t time.Duration) <-chan bool {
	done := make(chan bool, 1)
	lm.newLimit <- &limit{
		rate: rate{n, t},
		done: done,
	}
	return done
}

//Limit implements the limio.Limiter interface.
func (lm *SimpleManager) Limit(l chan int) <-chan bool {
	done := make(chan bool, 1)
	lm.newLimit <- &limit{
		lim:  l,
		done: done,
	}
	return done
}

//Unlimit implements the limio.Limiter interface.
func (lm *SimpleManager) Unlimit() {
	lm.newLimit <- nil
}

//Close allows the SimpleManager to free any resources it is using if the
//consumer has no further need for the SimpleManager.
func (lm *SimpleManager) Close() error {
	lm.cls <- struct{}{}
	return nil
}

//Unmanage allows consumers to remove a specific Limiter from its management
//strategy
func (lm *SimpleManager) Unmanage(l Limiter) {
	lm.clsLimiter <- l
}

//Manage takes a Limiter that will be adopted under the management policy of
//the SimpleManager.
func (lm *SimpleManager) Manage(l Limiter) error {
	if l == lm {
		return errors.New("A manager cannot manage itself.")
	}

	lm.newLimiter <- l
	return nil
}
<|endoftext|>