<commit_before>package pipe

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
)

// Pipe functions implement arbitrary functionality that may be
// plugged within pipe scripts and pipe lines. Pipe functions must
// not block reading or writing to the state streams. These
// operations must be run within state tasks (see State.Task).
type Pipe func(*State) error

// State defines the environment for Pipe functions to run on.
// Create a new State via the NewState function.
type State struct {

	// Stdin, Stdout, and Stderr represent the respective data
	// streams that the Pipe may choose to manipulate when running.
	// Reading from and writing to these streams must be done from
	// within a task created via the Task method.
	// The three streams are initialized by NewState and must
	// never be set to nil.
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer

	// Dir represents the directory in which all filesystem-related
	// operations performed by the Pipe must be run on. It defaults
	// to the current directory, and may be changed by Pipe functions,
	// although changes must not be done from tasks or other goroutines.
	Dir string

	// Env is the process environment in which all executions performed
	// by the Pipe must be run on. It defaults to a copy of the
	// environmnet from the current process, and may be changed by Pipe
	// functions, although changes must not be done from tasks or other
	// goroutines.
	Env []string

	tasks []*task
}

// NewState returns a new state for running pipes with.
func NewState(stdout, stderr io.Writer) *State {
	if stdout == nil {
		stdout = ioutil.Discard
	}
	if stderr == nil {
		stderr = ioutil.Discard
	}
	return &State{
		Stdin:  strings.NewReader(""),
		Stdout: stdout,
		Stderr: stderr,
		Env:    os.Environ(),
	}
}

// EnvVar returns the value for the named environment variable in s.
func (s *State) EnvVar(name string) string {
	prefix := name + "="
	for _, kv := range s.Env {
		if strings.HasPrefix(kv, prefix) {
			return kv[len(prefix):]
		}
	}
	return ""
}

// SetEnvVar sets the named environment variable to the given value in s.
func (s *State) SetEnvVar(name, value string) {
	prefix := name + "="
	for i, kv := range s.Env {
		if strings.HasPrefix(kv, prefix) {
			s.Env[i] = prefix + value
			return
		}
	}
	s.Env = append(s.Env, prefix+value)
}

type task struct {
	s    State
	f    func(*State) error
	done func() error
	kill func() error
}

func dummy() error { return nil }

// Task creates a new task that can block reading from and/or
// writing to the state streams.
// Tasks are all run concurrently, and must only be run if
// all the involved Pipe functions have returned no errors.
func (s *State) Task(f func(*State) error) {
	t := &task{*s, f, dummy, dummy}
	t.s.Env = append([]string(nil), s.Env...)
	s.tasks = append(s.tasks, t)
}

// TaskKill registers a function that must be called if interrupting
// the previous task registered in s is necessary. Functions that are
// blocked simply reading from and/or writing to the state streams
// do not need a kill function, as they will be unblocked by the
// closing of the streams used.
func (s *State) TaskKill(f func() error) {
	task := s.tasks[len(s.tasks)-1]
	task.kill = chain(task.kill, f)
}

// TaskDone registers a function to be executed after the
// previous task created in s finishes running.
func (s *State) TaskDone(f func() error) {
	task := s.tasks[len(s.tasks)-1]
	task.done = chain(task.done, f)
}

func chain(f1, f2 func() error) func() error {
	return func() error { return firstErr(f1(), f2()) }
}

func firstErr(err1, err2 error) error {
	if err1 != nil {
		return err1
	}
	return err2
}

// RunTasks runs all tasks that were registered in the pipe state.
func (s *State) RunTasks() error {
	done := make(chan error, len(s.tasks))
	for i := range s.tasks {
		task := s.tasks[i]
		go func() {
			done <- firstErr(task.f(&task.s), task.done())
		}()
	}
	var first error
	for _ = range s.tasks {
		err := <-done
		if err != nil && first == nil {
			first = err
			for i := range s.tasks {
				s.tasks[i].kill()
			}
		}
	}
	s.tasks = nil
	return first
}

func Run(p Pipe) error {
	s := NewState(nil, nil)
	err := p(s)
	if err == nil {
		err = s.RunTasks()
	}
	return err
}

func Output(p Pipe) ([]byte, error) {
	outb := &OutputBuffer{}
	s := NewState(outb, nil)
	err := p(s)
	if err == nil {
		err = s.RunTasks()
	}
	return outb.Bytes(), err
}

func CombinedOutput(p Pipe) ([]byte, error) {
	outb := &OutputBuffer{}
	s := NewState(outb, outb)
	err := p(s)
	if err == nil {
		err = s.RunTasks()
	}
	return outb.Bytes(), err
}

func DisjointOutput(p Pipe) (stdout []byte, stderr []byte, err error) {
	outb := &OutputBuffer{}
	errb := &OutputBuffer{}
	s := NewState(outb, errb)
	err = p(s)
	if err == nil {
		err = s.RunTasks()
	}
	return outb.Bytes(), errb.Bytes(), err
}

type OutputBuffer struct {
	m   sync.Mutex
	buf []byte
}

func (out *OutputBuffer) Write(p []byte) (n int, err error) {
	out.m.Lock()
	out.buf = append(out.buf, p...)
	out.m.Unlock()
	return len(p), nil
}

func (out *OutputBuffer) Bytes() []byte {
	out.m.Lock()
	buf := out.buf
	out.m.Unlock()
	return buf
}

func Exec(name string, args ...string) Pipe {
	return func(s *State) error {
		ch := make(chan *os.Process, 1)
		s.Task(func(s *State) error {
			cmd := exec.Command(name, args...)
			cmd.Dir = s.Dir
			cmd.Env = s.Env
			cmd.Stdin = s.Stdin
			cmd.Stdout = s.Stdout
			cmd.Stderr = s.Stderr
			err := cmd.Start()
			ch <- cmd.Process
			if err != nil {
				return err
			}
			if err := cmd.Wait(); err != nil {
				return fmt.Errorf("command %q: %v", name, err)
			}
			return nil
		})
		s.TaskKill(func() error {
			if p, ok := <-ch; ok {
				p.Kill()
			}
			return nil
		})
		return nil
	}
}

func System(command string) Pipe {
	return Exec("/bin/sh", "-c", command)
}

func ChDir(dir string) Pipe {
	return func(s *State) error {
		s.Dir = filepath.Join(s.Dir, dir)
		return nil
	}
}

func MkDir(dir string, perm os.FileMode) Pipe {
	return func(s *State) error {
		return os.Mkdir(filepath.Join(s.Dir, dir), perm)
	}
}

func SetEnvVar(name string, value string) Pipe {
	return func(s *State) error {
		s.SetEnvVar(name, value)
		return nil
	}
}

func CombineToErr() Pipe {
	return func(s *State) error {
		s.Stdout = s.Stderr
		return nil
	}
}

func CombineToOut() Pipe {
	return func(s *State) error {
		s.Stderr = s.Stdout
		return nil
	}
}

func Line(p ...Pipe) Pipe {
	return func(s *State) error {
		end := len(p) - 1
		endStdout := s.Stdout
		var r *io.PipeReader
		var w *io.PipeWriter
		for i, p := range p {
			closer := r
			if i == end {
				r, w = nil, nil
				s.Stdout = endStdout
			} else {
				r, w = io.Pipe()
				s.Stdout = w
			}
			closew := w
			closef := func() error {
				if closew != nil {
					closew.Close()
				}
				if closer != nil {
					//io.Copy(ioutil.Discard, closer)
					closer.Close()
				}
				return nil
			}
			if err := p(s); err != nil {
				closef()
				return err
			}
			s.TaskDone(closef)
			if i < end {
				s.Stdin = r
			}
		}
		return nil
	}
}

func Script(p ...Pipe) Pipe {
	return func(s *State) error {
		env := append([]string(nil), s.Env...)
		for _, p := range p {
			if err := p(s); err != nil {
				return err
			}
			if err := s.RunTasks(); err != nil {
				return err
			}
		}
		s.Env = env
		return nil
	}
}

func Echo(str string) Pipe {
	return func(s *State) error {
		s.Task(func(s *State) error {
			_, err := s.Stdout.Write([]byte(str))
			return err
		})
		return nil
	}
}

func Read(r io.Reader) Pipe {
	return func(s *State) error {
		s.Task(func(s *State) error {
			_, err := io.Copy(s.Stdout, r)
			return err
		})
		return nil
	}
}

func Write(w io.Writer) Pipe {
	return func(s *State) error {
		s.Task(func(s *State) error {
			_, err := io.Copy(w, s.Stdin)
			return err
		})
		return nil
	}
}

func Discard() Pipe {
	return Write(ioutil.Discard)
}

func Tee(w io.Writer) Pipe {
	return func(s *State) error {
		s.Task(func(s *State) error {
			_, err := io.Copy(w, io.TeeReader(s.Stdin, s.Stdout))
			return err
		})
		return nil
	}
}

func ReadFile(path string) Pipe {
	return func(s *State) error {
		s.Task(func(s *State) error {
			file, err := os.Open(filepath.Join(s.Dir, path))
			if err != nil {
				return err
			}
			_, err = io.Copy(s.Stdout, file)
			file.Close()
			return err
		})
		return nil
	}
}

func WriteFile(path string, perm os.FileMode) Pipe {
	return func(s *State) error {
		s.Task(func(s *State) error {
			path = filepath.Join(s.Dir, path)
			file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
			if err != nil {
				return err
			}
			_, err = io.Copy(file, s.Stdin)
			return firstErr(err, file.Close())
		})
		return nil
	}
}

func TeeFile(path string, perm os.FileMode) Pipe {
	return func(s *State) error {
		s.Task(func(s *State) error {
			// BUG: filepath.Join missing; test it first.
			file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
			if err != nil {
				return err
			}
			_, err = io.Copy(file, io.TeeReader(s.Stdin, s.Stdout))
			return firstErr(err, file.Close())
		})
		return nil
	}
}

//pipe.If(pipe.IsFile("/foo"), func(p pipe.Line) (Line, error) {
<commit_msg>Significant improvements to the internal organization.<commit_after>package pipe

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
)

// Pipe functions implement arbitrary functionality that may be
// plugged within pipe scripts and pipe lines. Pipe functions must
// not block reading or writing to the state streams. These
// operations must be run within state tasks (see State.Task).

// State defines the environment for Pipe functions to run on.
// Create a new State via the NewState function.
type State struct {

	// Stdin, Stdout, and Stderr represent the respective data
	// streams that the Pipe may choose to manipulate when running.
	// Reading from and writing to these streams must be done from
	// within a task created via the Task method.
	// The three streams are initialized by NewState and must
	// never be set to nil.
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer

	// Dir represents the directory in which all filesystem-related
	// operations performed by the Pipe must be run on. It defaults
	// to the current directory, and may be changed by Pipe functions,
	// although changes must not be done from tasks or other goroutines.
	Dir string

	// Env is the process environment in which all executions performed
	// by the Pipe must be run on. It defaults to a copy of the
	// environmnet from the current process, and may be changed by Pipe
	// functions, although changes must not be done from tasks or other
	// goroutines.
	Env []string

	pendingFlushes []*pendingFlush
}

type pendingFlush struct {
	s State
	f Flusher
	c []io.Closer
}

func (pf *pendingFlush) closeWhenDone(c io.Closer) {
	pf.c = append(pf.c, c)
}

func (pf *pendingFlush) closeAll() {
	for _, c := range pf.c {
		c.Close()
	}
}

type Pipe func(s *State) error

type Flusher interface {
	Flush(s *State) error
	Kill()
}

func (s *State) AddFlusher(f Flusher) error {
	pf := &pendingFlush{*s, f, nil}
	pf.s.Env = append([]string(nil), s.Env...)
	s.pendingFlushes = append(s.pendingFlushes, pf)
	return nil
}

func (s *State) FlushAll() error {
	done := make(chan error, len(s.pendingFlushes))
	for _, f := range s.pendingFlushes {
		go func(pf *pendingFlush) {
			err := pf.f.Flush(&pf.s)
			pf.closeAll()
			done <- err
		}(f)
	}
	var first error
	for _ = range s.pendingFlushes {
		err := <-done
		if err != nil && first == nil {
			first = err
			for _, pf := range s.pendingFlushes {
				pf.f.Kill()
			}
		}
	}
	s.pendingFlushes = nil
	return first
}

type flusherFunc func(s *State) error

func (f flusherFunc) Flush(s *State) error { return f(s) }
func (f flusherFunc) Kill()                {}

func FlusherFunc(f func(s *State) error) Pipe {
	return func(s *State) error {
		s.AddFlusher(flusherFunc(f))
		return nil
	}
}

// NewState returns a new state for running pipes with.
func NewState(stdout, stderr io.Writer) *State {
	if stdout == nil {
		stdout = ioutil.Discard
	}
	if stderr == nil {
		stderr = ioutil.Discard
	}
	return &State{
		Stdin:  strings.NewReader(""),
		Stdout: stdout,
		Stderr: stderr,
		Env:    os.Environ(),
	}
}

// EnvVar returns the value for the named environment variable in s.
func (s *State) EnvVar(name string) string {
	prefix := name + "="
	for _, kv := range s.Env {
		if strings.HasPrefix(kv, prefix) {
			return kv[len(prefix):]
		}
	}
	return ""
}

// SetEnvVar sets the named environment variable to the given value in s.
func (s *State) SetEnvVar(name, value string) {
	prefix := name + "="
	for i, kv := range s.Env {
		if strings.HasPrefix(kv, prefix) {
			s.Env[i] = prefix + value
			return
		}
	}
	s.Env = append(s.Env, prefix+value)
}

func firstErr(err1, err2 error) error {
	if err1 != nil {
		return err1
	}
	return err2
}

func Run(p Pipe) error {
	s := NewState(nil, nil)
	err := p(s)
	if err == nil {
		err = s.FlushAll()
	}
	return err
}

func Output(p Pipe) ([]byte, error) {
	outb := &OutputBuffer{}
	s := NewState(outb, nil)
	err := p(s)
	if err == nil {
		err = s.FlushAll()
	}
	return outb.Bytes(), err
}

func CombinedOutput(p Pipe) ([]byte, error) {
	outb := &OutputBuffer{}
	s := NewState(outb, outb)
	err := p(s)
	if err == nil {
		err = s.FlushAll()
	}
	return outb.Bytes(), err
}

func DisjointOutput(p Pipe) (stdout []byte, stderr []byte, err error) {
	outb := &OutputBuffer{}
	errb := &OutputBuffer{}
	s := NewState(outb, errb)
	err = p(s)
	if err == nil {
		err = s.FlushAll()
	}
	return outb.Bytes(), errb.Bytes(), err
}

type OutputBuffer struct {
	m   sync.Mutex
	buf []byte
}

func (out *OutputBuffer) Write(p []byte) (n int, err error) {
	out.m.Lock()
	out.buf = append(out.buf, p...)
	out.m.Unlock()
	return len(p), nil
}

func (out *OutputBuffer) Bytes() []byte {
	out.m.Lock()
	buf := out.buf
	out.m.Unlock()
	return buf
}

func Exec(name string, args ...string) Pipe {
	return func(s *State) error {
		s.AddFlusher(&execFlusher{name, args, make(chan *os.Process, 1)})
		return nil
	}
}

func System(command string) Pipe {
	return Exec("/bin/sh", "-c", command)
}

type execFlusher struct {
	name string
	args []string
	ch   chan *os.Process
}

func (f *execFlusher) Flush(s *State) error {
	cmd := exec.Command(f.name, f.args...)
	cmd.Dir = s.Dir
	cmd.Env = s.Env
	cmd.Stdin = s.Stdin
	cmd.Stdout = s.Stdout
	cmd.Stderr = s.Stderr
	err := cmd.Start()
	f.ch <- cmd.Process
	if err != nil {
		return err
	}
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("command %q: %v", f.name, err)
	}
	return nil
}

func (f *execFlusher) Kill() {
	if p, ok := <-f.ch; ok {
		p.Kill()
	}
}

func ChDir(dir string) Pipe {
	return func(s *State) error {
		s.Dir = filepath.Join(s.Dir, dir)
		return nil
	}
}

func MkDir(dir string, perm os.FileMode) Pipe {
	return func(s *State) error {
		return os.Mkdir(filepath.Join(s.Dir, dir), perm)
	}
}

func SetEnvVar(name string, value string) Pipe {
	return func(s *State) error {
		s.SetEnvVar(name, value)
		return nil
	}
}

func CombineToErr() Pipe {
	return func(s *State) error {
		s.Stdout = s.Stderr
		return nil
	}
}

func CombineToOut() Pipe {
	return func(s *State) error {
		s.Stderr = s.Stdout
		return nil
	}
}

func Line(p ...Pipe) Pipe {
	return func(s *State) error {
		end := len(p) - 1
		endStdout := s.Stdout
		var r *io.PipeReader
		var w *io.PipeWriter
		for i, p := range p {
			closeIn := r
			if i == end {
				r, w = nil, nil
				s.Stdout = endStdout
			} else {
				r, w = io.Pipe()
				s.Stdout = w
			}
			closeOut := w

			oldLen := len(s.pendingFlushes)
			if err := p(s); err != nil {
				if closeIn != nil {
					closeIn.Close()
				}
				return err
			}
			newLen := len(s.pendingFlushes)

			// Close the created ends that were put in place for this
			// specific Pipe after the last flusher that was registered
			// as a consequence of running the given Pipe ends running.
			if newLen == oldLen {
				if closeIn != nil {
					closeIn.Close()
				}
				if closeOut != nil {
					closeOut.Close()
				}
			} else {
				if closeIn != nil {
					for fi := oldLen; fi < newLen+1; fi++ {
						if fi == newLen || s.pendingFlushes[fi].s.Stdin != closeIn {
							s.pendingFlushes[fi-1].closeWhenDone(closeIn)
						}
					}
				}
				if closeOut != nil {
					for fi := newLen - 1; fi >= oldLen; fi-- {
						if fi == oldLen || (s.pendingFlushes[fi].s.Stdout == closeOut || s.pendingFlushes[fi].s.Stderr == closeOut) {
							s.pendingFlushes[fi].closeWhenDone(closeOut)
						}
					}
				}
			}

			if i < end {
				s.Stdin = r
			}
		}
		return nil
	}
}

func Script(p ...Pipe) Pipe {
	return func(s *State) error {
		env := append([]string(nil), s.Env...)
		for _, p := range p {
			if err := p(s); err != nil {
				return err
			}
			if err := s.FlushAll(); err != nil {
				return err
			}
		}
		s.Env = env
		return nil
	}
}

func Echo(str string) Pipe {
	return FlusherFunc(func(s *State) error {
		_, err := s.Stdout.Write([]byte(str))
		return err
	})
}

func Read(r io.Reader) Pipe {
	return FlusherFunc(func(s *State) error {
		_, err := io.Copy(s.Stdout, r)
		return err
	})
}

func Write(w io.Writer) Pipe {
	return FlusherFunc(func(s *State) error {
		_, err := io.Copy(w, s.Stdin)
		return err
	})
}

func Discard() Pipe {
	return Write(ioutil.Discard)
}

func Tee(w io.Writer) Pipe {
	return FlusherFunc(func(s *State) error {
		_, err := io.Copy(w, io.TeeReader(s.Stdin, s.Stdout))
		return err
	})
}

func ReadFile(path string) Pipe {
	return FlusherFunc(func(s *State) error {
		file, err := os.Open(filepath.Join(s.Dir, path))
		if err != nil {
			return err
		}
		_, err = io.Copy(s.Stdout, file)
		file.Close()
		return err
	})
}

func WriteFile(path string, perm os.FileMode) Pipe {
	return FlusherFunc(func(s *State) error {
		path = filepath.Join(s.Dir, path)
		file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
		if err != nil {
			return err
		}
		_, err = io.Copy(file, s.Stdin)
		return firstErr(err, file.Close())
	})
}

func TeeFile(path string, perm os.FileMode) Pipe {
	return FlusherFunc(func(s *State) error {
		// BUG: filepath.Join missing; test it first.
		file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
		if err != nil {
			return err
		}
		_, err = io.Copy(file, io.TeeReader(s.Stdin, s.Stdout))
		return firstErr(err, file.Close())
	})
}

//pipe.If(pipe.IsFile("/foo"), func(p pipe.Line) (Line, error) {
<|endoftext|>


<commit_before>package sh

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"io"
	"os"
	"strings"
)

// unmarshal shell output to decode json
func (s *Session) UnmarshalJSON(data interface{}) (err error) {
	bufrw := bytes.NewBuffer(nil)
	s.Stdout = bufrw
	if err = s.Run(); err != nil {
		return
	}
	return json.NewDecoder(bufrw).Decode(data)
}

// unmarshal command output into xml
func (s *Session) UnmarshalXML(data interface{}) (err error) {
	bufrw := bytes.NewBuffer(nil)
	s.Stdout = bufrw
	if err = s.Run(); err != nil {
		return
	}
	return xml.NewDecoder(bufrw).Decode(data)
}

// start command
func (s *Session) Start() (err error) {
	s.started = true
	var rd *io.PipeReader
	var wr *io.PipeWriter
	var length = len(s.cmds)
	if s.ShowCMD {
		var cmds = make([]string, 0, 4)
		for _, cmd := range s.cmds {
			cmds = append(cmds, strings.Join(cmd.Args, " "))
		}
		s.writePrompt(strings.Join(cmds, " | "))
	}
	for index, cmd := range s.cmds {
		if index != 0 {
			cmd.Stdin = rd
		}
		if index != length {
			rd, wr = io.Pipe() // create pipe
			cmd.Stdout = wr
			cmd.Stderr = os.Stderr
		}
		if index == length-1 {
			cmd.Stdout = s.Stdout
			cmd.Stderr = s.Stderr
		}
		err = cmd.Start()
		if err != nil {
			return
		}
	}
	return
}

// Should be call after Start()
func (s *Session) Wait() (err error) {
	for _, cmd := range s.cmds {
		cmd.Wait()
		wr, ok := cmd.Stdout.(*io.PipeWriter)
		if ok {
			wr.Close()
		}
	}
	return
}

func (s *Session) Run() (err error) {
	if err = s.Start(); err != nil {
		return
	}
	return s.Wait()
}

func (s *Session) Output() (out []byte, err error) {
	oldout := s.Stdout
	defer func() {
		s.Stdout = oldout
	}()
	stdout := bytes.NewBuffer(nil)
	s.Stdout = stdout
	err = s.Run()
	out = stdout.Bytes()
	return
}
<commit_msg>fix go-sh error return bug<commit_after>package sh

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"io"
	"os"
	"strings"
)

// unmarshal shell output to decode json
func (s *Session) UnmarshalJSON(data interface{}) (err error) {
	bufrw := bytes.NewBuffer(nil)
	s.Stdout = bufrw
	if err = s.Run(); err != nil {
		return
	}
	return json.NewDecoder(bufrw).Decode(data)
}

// unmarshal command output into xml
func (s *Session) UnmarshalXML(data interface{}) (err error) {
	bufrw := bytes.NewBuffer(nil)
	s.Stdout = bufrw
	if err = s.Run(); err != nil {
		return
	}
	return xml.NewDecoder(bufrw).Decode(data)
}

// start command
func (s *Session) Start() (err error) {
	s.started = true
	var rd *io.PipeReader
	var wr *io.PipeWriter
	var length = len(s.cmds)
	if s.ShowCMD {
		var cmds = make([]string, 0, 4)
		for _, cmd := range s.cmds {
			cmds = append(cmds, strings.Join(cmd.Args, " "))
		}
		s.writePrompt(strings.Join(cmds, " | "))
	}
	for index, cmd := range s.cmds {
		if index != 0 {
			cmd.Stdin = rd
		}
		if index != length {
			rd, wr = io.Pipe() // create pipe
			cmd.Stdout = wr
			cmd.Stderr = os.Stderr
		}
		if index == length-1 {
			cmd.Stdout = s.Stdout
			cmd.Stderr = s.Stderr
		}
		err = cmd.Start()
		if err != nil {
			return
		}
	}
	return
}

// Should be call after Start()
// only catch the last command error
func (s *Session) Wait() (err error) {
	for _, cmd := range s.cmds {
		err = cmd.Wait()
		wr, ok := cmd.Stdout.(*io.PipeWriter)
		if ok {
			wr.Close()
		}
	}
	return err
}

func (s *Session) Run() (err error) {
	if err = s.Start(); err != nil {
		return
	}
	return s.Wait()
}

func (s *Session) Output() (out []byte, err error) {
	oldout := s.Stdout
	defer func() {
		s.Stdout = oldout
	}()
	stdout := bytes.NewBuffer(nil)
	s.Stdout = stdout
	err = s.Run()
	out = stdout.Bytes()
	return
}
<|endoftext|>


<commit_before>package piecedownloader

import (
	"crypto/sha1" // nolint: gosec
	"errors"
	"fmt"

	"github.com/cenkalti/rain/torrent/internal/peer"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
	"github.com/cenkalti/rain/torrent/internal/pieceio"
)

const maxQueuedBlocks = 10

// PieceDownloader downloads all blocks of a piece from a peer.
type PieceDownloader struct {
	Piece          *pieceio.Piece
	Peer           *peer.Peer
	buffer         []byte
	nextBlockIndex uint32
	requested      map[uint32]struct{}
	done           map[uint32]struct{}
	PieceC         chan Piece
	RejectC        chan *pieceio.Block
	ChokeC         chan struct{}
	UnchokeC       chan struct{}
	resultC        chan Result
	closeC         chan struct{}
	doneC          chan struct{}
}

type Result struct {
	Peer  *peer.Peer
	Piece *pieceio.Piece
	Error error
}

func New(pi *pieceio.Piece, pe *peer.Peer, resultC chan Result) *PieceDownloader {
	return &PieceDownloader{
		Piece:     pi,
		Peer:      pe,
		buffer:    make([]byte, pi.Length),
		requested: make(map[uint32]struct{}),
		done:      make(map[uint32]struct{}),
		PieceC:    make(chan Piece),
		RejectC:   make(chan *pieceio.Block),
		ChokeC:    make(chan struct{}),
		UnchokeC:  make(chan struct{}),
		resultC:   resultC,
		closeC:    make(chan struct{}),
		doneC:     make(chan struct{}),
	}
}

func (d *PieceDownloader) Close() {
	close(d.closeC)
}

func (d *PieceDownloader) Done() <-chan struct{} {
	return d.doneC
}

func (d *PieceDownloader) requestBlocks() {
	for ; d.nextBlockIndex < uint32(len(d.Piece.Blocks)) && len(d.requested) < maxQueuedBlocks; d.nextBlockIndex++ {
		b := d.Piece.Blocks[d.nextBlockIndex]
		if _, ok := d.done[d.nextBlockIndex]; ok {
			continue
		}
		d.requested[d.nextBlockIndex] = struct{}{}
		msg := peerprotocol.RequestMessage{Index: d.Piece.Index, Begin: b.Begin, Length: b.Length}
		d.Peer.SendMessage(msg)
	}
}

func (d *PieceDownloader) Run() {
	defer close(d.doneC)

	result := Result{
		Peer:  d.Peer,
		Piece: d.Piece,
	}
	defer func() {
		select {
		case d.resultC <- result:
		case <-d.closeC:
		}
	}()

	d.requestBlocks()
	for {
		select {
		case p := <-d.PieceC:
			if _, ok := d.requested[p.Block.Index]; !ok {
				result.Error = fmt.Errorf("peer sent unrequested piece block: %q", p.Block)
				return
			}
			delete(d.requested, p.Block.Index)
			_, result.Error = p.Piece.Reader.Read(d.buffer[p.Piece.Begin : p.Piece.Begin+p.Piece.Length])
			close(p.Piece.Done)
			if result.Error != nil {
				return
			}
			d.done[p.Block.Index] = struct{}{}
			if d.allDone() {
				ok := d.Piece.VerifyHash(d.buffer, sha1.New()) // nolint: gosec
				if !ok {
					result.Error = errors.New("received corrupt piece")
					break
				}
				_, result.Error = d.Piece.Data.Write(d.buffer)
				return
			}
			d.requestBlocks()
		case blk := <-d.RejectC:
			if _, ok := d.requested[blk.Index]; !ok {
				result.Error = fmt.Errorf("peer sent reject to unrequested block: %q", blk)
				return
			}
			delete(d.requested, blk.Index)
			d.nextBlockIndex = 0
		case <-d.ChokeC:
			d.requested = make(map[uint32]struct{})
			d.nextBlockIndex = 0
		case <-d.UnchokeC:
			d.requestBlocks()
		case <-d.closeC:
			return
		}
	}
}

func (d *PieceDownloader) allDone() bool {
	return len(d.done) == len(d.Piece.Blocks)
}
<commit_msg>read full piece from peer<commit_after>package piecedownloader

import (
	"crypto/sha1" // nolint: gosec
	"errors"
	"fmt"
	"io"

	"github.com/cenkalti/rain/torrent/internal/peer"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
	"github.com/cenkalti/rain/torrent/internal/pieceio"
)

const maxQueuedBlocks = 10

// PieceDownloader downloads all blocks of a piece from a peer.
type PieceDownloader struct {
	Piece          *pieceio.Piece
	Peer           *peer.Peer
	buffer         []byte
	nextBlockIndex uint32
	requested      map[uint32]struct{}
	done           map[uint32]struct{}
	PieceC         chan Piece
	RejectC        chan *pieceio.Block
	ChokeC         chan struct{}
	UnchokeC       chan struct{}
	resultC        chan Result
	closeC         chan struct{}
	doneC          chan struct{}
}

type Result struct {
	Peer  *peer.Peer
	Piece *pieceio.Piece
	Error error
}

func New(pi *pieceio.Piece, pe *peer.Peer, resultC chan Result) *PieceDownloader {
	return &PieceDownloader{
		Piece:     pi,
		Peer:      pe,
		buffer:    make([]byte, pi.Length),
		requested: make(map[uint32]struct{}),
		done:      make(map[uint32]struct{}),
		PieceC:    make(chan Piece),
		RejectC:   make(chan *pieceio.Block),
		ChokeC:    make(chan struct{}),
		UnchokeC:  make(chan struct{}),
		resultC:   resultC,
		closeC:    make(chan struct{}),
		doneC:     make(chan struct{}),
	}
}

func (d *PieceDownloader) Close() {
	close(d.closeC)
}

func (d *PieceDownloader) Done() <-chan struct{} {
	return d.doneC
}

func (d *PieceDownloader) requestBlocks() {
	for ; d.nextBlockIndex < uint32(len(d.Piece.Blocks)) && len(d.requested) < maxQueuedBlocks; d.nextBlockIndex++ {
		b := d.Piece.Blocks[d.nextBlockIndex]
		if _, ok := d.done[d.nextBlockIndex]; ok {
			continue
		}
		d.requested[d.nextBlockIndex] = struct{}{}
		msg := peerprotocol.RequestMessage{Index: d.Piece.Index, Begin: b.Begin, Length: b.Length}
		d.Peer.SendMessage(msg)
	}
}

func (d *PieceDownloader) Run() {
	defer close(d.doneC)

	result := Result{
		Peer:  d.Peer,
		Piece: d.Piece,
	}
	defer func() {
		select {
		case d.resultC <- result:
		case <-d.closeC:
		}
	}()

	d.requestBlocks()
	for {
		select {
		case p := <-d.PieceC:
			if _, ok := d.requested[p.Block.Index]; !ok {
				result.Error = fmt.Errorf("peer sent unrequested piece block: %q", p.Block)
				return
			}
			delete(d.requested, p.Block.Index)
			_, result.Error = io.ReadFull(p.Piece.Reader, d.buffer[p.Piece.Begin:p.Piece.Begin+p.Piece.Length])
			close(p.Piece.Done)
			if result.Error != nil {
				return
			}
			d.done[p.Block.Index] = struct{}{}
			if d.allDone() {
				ok := d.Piece.VerifyHash(d.buffer, sha1.New()) // nolint: gosec
				if !ok {
					result.Error = errors.New("received corrupt piece")
					break
				}
				_, result.Error = d.Piece.Data.Write(d.buffer)
				return
			}
			d.requestBlocks()
		case blk := <-d.RejectC:
			if _, ok := d.requested[blk.Index]; !ok {
				result.Error = fmt.Errorf("peer sent reject to unrequested block: %q", blk)
				return
			}
			delete(d.requested, blk.Index)
			d.nextBlockIndex = 0
		case <-d.ChokeC:
			d.requested = make(map[uint32]struct{})
			d.nextBlockIndex = 0
		case <-d.UnchokeC:
			d.requestBlocks()
		case <-d.closeC:
			return
		}
	}
}

func (d *PieceDownloader) allDone() bool {
	return len(d.done) == len(d.Piece.Blocks)
}
<|endoftext|>


<commit_before>package stripe

import (
	"strconv"

	"github.com/stripe/stripe-go/form"
)

// PlanInterval is the list of allowed values for a plan's interval.
type PlanInterval string

const (
	PlanIntervalDay   PlanInterval = "day"
	PlanIntervalWeek  PlanInterval = "week"
	PlanIntervalMonth PlanInterval = "month"
	PlanIntervalYear  PlanInterval = "year"
)

// PlanBillingScheme is the list of allowed values for a plan's billing scheme.
type PlanBillingScheme string

const (
	PlanBillingSchemePerUnit PlanBillingScheme = "per_unit"
	PlanBillingSchemeTiered  PlanBillingScheme = "tiered"
)

// PlanUsageType is the list of allowed values for a plan's usage type.
type PlanUsageType string

const (
	PlanUsageTypeLicensed PlanUsageType = "licensed"
	PlanUsageTypeMetered  PlanUsageType = "metered"
)

// PlanTiersMode is the list of allowed values for a plan's tiers mode.
type PlanTiersMode string

const (
	PlanTiersModeGraduated PlanTiersMode = "graduated"
	PlanTiersModeVolume    PlanTiersMode = "volume"
)

// PlanTransformUsageRound is the list of allowed values for a plan's transform usage round logic.
type PlanTransformUsageRound string

const (
	PlanTransformUsageRoundDown PlanTransformUsageRound = "down"
	PlanTransformUsageRoundUp   PlanTransformUsageRound = "up"
)

// PlanAggregateUsage is the list of allowed values for a plan's aggregate usage.
type PlanAggregateUsage string

const (
	PlanAggregateUsageLastDuringPeriod PlanAggregateUsage = "last_during_period"
	PlanAggregateUsageLastEver         PlanAggregateUsage = "last_ever"
	PlanAggregateUsageMax              PlanAggregateUsage = "max"
	PlanAggregateUsageSum              PlanAggregateUsage = "sum"
)

// Plan is the resource representing a Stripe plan.
// For more details see https://stripe.com/docs/api#plans.
type Plan struct {
	Active          bool                `json:"active"`
	AggregateUsage  string              `json:"aggregate_usage"`
	Amount          int64               `json:"amount"`
	BillingScheme   PlanBillingScheme   `json:"billing_scheme"`
	Created         int64               `json:"created"`
	Currency        Currency            `json:"currency"`
	Deleted         bool                `json:"deleted"`
	ID              string              `json:"id"`
	Interval        PlanInterval        `json:"interval"`
	IntervalCount   int64               `json:"interval_count"`
	Livemode        bool                `json:"livemode"`
	Metadata        map[string]string   `json:"metadata"`
	Nickname        string              `json:"nickname"`
	Product         string              `json:"product"`
	Tiers           []*PlanTier         `json:"tiers"`
	TiersMode       string              `json:"tiers_mode"`
	TransformUsage  *PlanTransformUsage `json:"transform_usage"`
	TrialPeriodDays int64               `json:"trial_period_days"`
	UsageType       PlanUsageType       `json:"usage_type"`
}

// PlanList is a list of plans as returned from a list endpoint.
type PlanList struct {
	ListMeta
	Data []*Plan `json:"data"`
}

// PlanListParams is the set of parameters that can be used when listing plans.
// For more details see https://stripe.com/docs/api#list_plans.
type PlanListParams struct {
	ListParams   `form:"*"`
	Active       *bool             `form:"active"`
	Created      *int64            `form:"created"`
	CreatedRange *RangeQueryParams `form:"created"`
}

// PlanParams is the set of parameters that can be used when creating or updating a plan.
// For more details see https://stripe.com/docs/api#create_plan and https://stripe.com/docs/api#update_plan.
type PlanParams struct {
	Params          `form:"*"`
	Active          *bool                     `form:"active"`
	AggregateUsage  *string                   `form:"aggregate_usage"`
	Amount          *int64                    `form:"amount"`
	BillingScheme   *string                   `form:"billing_scheme"`
	Currency        *string                   `form:"currency"`
	ID              *string                   `form:"id"`
	Interval        *string                   `form:"interval"`
	IntervalCount   *int64                    `form:"interval_count"`
	Nickname        *string                   `form:"nickname"`
	Product         *PlanProductParams        `form:"product"`
	ProductID       *string                   `form:"product"`
	Tiers           []*PlanTierParams         `form:"tiers,indexed"`
	TiersMode       *string                   `form:"tiers_mode"`
	TransformUsage  *PlanTransformUsageParams `form:"transform_usage"`
	TrialPeriodDays *int64                    `form:"trial_period_days"`
	UsageType       *string                   `form:"usage_type"`
}

// PlanTier configures tiered pricing
type PlanTier struct {
	Amount int64 `json:"amount"`
	UpTo   int64 `json:"up_to"`
}

// PlanTransformUsage represents the bucket billing configuration.
type PlanTransformUsage struct {
	DivideBy int64                   `json:"divide_by"`
	Round    PlanTransformUsageRound `json:"round"`
}

// PlanTransformUsageParams represents the bucket billing configuration.
type PlanTransformUsageParams struct {
	DivideBy *int64  `form:"divide_by"`
	Round    *string `form:"round"`
}

// PlanTierParams configures tiered pricing
type PlanTierParams struct {
	Params  `form:"*"`
	Amount  *int64 `form:"amount"`
	UpTo    *int64 `form:"-"` // handled in custom AppendTo
	UpToInf *bool  `form:"-"` // handled in custom AppendTo
}

// AppendTo implements custom up_to serialisation logic for tiers configuration
func (p *PlanTierParams) AppendTo(body *form.Values, keyParts []string) {
	if BoolValue(p.UpToInf) {
		body.Add(form.FormatKey(append(keyParts, "up_to")), "inf")
	} else {
		body.Add(form.FormatKey(append(keyParts, "up_to")), strconv.FormatInt(Int64Value(p.UpTo), 10))
	}
}

// PlanProductParams is the set of parameters that can be used when creating a product inside a plan
// This can only be used on plan creation and won't work on plan update.
// For more details see https://stripe.com/docs/api#create_plan-product and https://stripe.com/docs/api#update_plan-product
type PlanProductParams struct {
	ID                  *string           `form:"id"`
	Name                *string           `form:"name"`
	Metadata            map[string]string `form:"metadata"`
	StatementDescriptor *string           `form:"statement_descriptor"`
}
<commit_msg>Add support for listing plans by product<commit_after>package stripe

import (
	"strconv"

	"github.com/stripe/stripe-go/form"
)

// PlanInterval is the list of allowed values for a plan's interval.
type PlanInterval string

const (
	PlanIntervalDay   PlanInterval = "day"
	PlanIntervalWeek  PlanInterval = "week"
	PlanIntervalMonth PlanInterval = "month"
	PlanIntervalYear  PlanInterval = "year"
)

// PlanBillingScheme is the list of allowed values for a plan's billing scheme.
type PlanBillingScheme string

const (
	PlanBillingSchemePerUnit PlanBillingScheme = "per_unit"
	PlanBillingSchemeTiered  PlanBillingScheme = "tiered"
)

// PlanUsageType is the list of allowed values for a plan's usage type.
type PlanUsageType string

const (
	PlanUsageTypeLicensed PlanUsageType = "licensed"
	PlanUsageTypeMetered  PlanUsageType = "metered"
)

// PlanTiersMode is the list of allowed values for a plan's tiers mode.
type PlanTiersMode string

const (
	PlanTiersModeGraduated PlanTiersMode = "graduated"
	PlanTiersModeVolume    PlanTiersMode = "volume"
)

// PlanTransformUsageRound is the list of allowed values for a plan's transform usage round logic.
type PlanTransformUsageRound string

const (
	PlanTransformUsageRoundDown PlanTransformUsageRound = "down"
	PlanTransformUsageRoundUp   PlanTransformUsageRound = "up"
)

// PlanAggregateUsage is the list of allowed values for a plan's aggregate usage.
type PlanAggregateUsage string

const (
	PlanAggregateUsageLastDuringPeriod PlanAggregateUsage = "last_during_period"
	PlanAggregateUsageLastEver         PlanAggregateUsage = "last_ever"
	PlanAggregateUsageMax              PlanAggregateUsage = "max"
	PlanAggregateUsageSum              PlanAggregateUsage = "sum"
)

// Plan is the resource representing a Stripe plan.
// For more details see https://stripe.com/docs/api#plans.
type Plan struct {
	Active          bool                `json:"active"`
	AggregateUsage  string              `json:"aggregate_usage"`
	Amount          int64               `json:"amount"`
	BillingScheme   PlanBillingScheme   `json:"billing_scheme"`
	Created         int64               `json:"created"`
	Currency        Currency            `json:"currency"`
	Deleted         bool                `json:"deleted"`
	ID              string              `json:"id"`
	Interval        PlanInterval        `json:"interval"`
	IntervalCount   int64               `json:"interval_count"`
	Livemode        bool                `json:"livemode"`
	Metadata        map[string]string   `json:"metadata"`
	Nickname        string              `json:"nickname"`
	Product         string              `json:"product"`
	Tiers           []*PlanTier         `json:"tiers"`
	TiersMode       string              `json:"tiers_mode"`
	TransformUsage  *PlanTransformUsage `json:"transform_usage"`
	TrialPeriodDays int64               `json:"trial_period_days"`
	UsageType       PlanUsageType       `json:"usage_type"`
}

// PlanList is a list of plans as returned from a list endpoint.
type PlanList struct {
	ListMeta
	Data []*Plan `json:"data"`
}

// PlanListParams is the set of parameters that can be used when listing plans.
// For more details see https://stripe.com/docs/api#list_plans.
type PlanListParams struct {
	ListParams   `form:"*"`
	Active       *bool             `form:"active"`
	Created      *int64            `form:"created"`
	CreatedRange *RangeQueryParams `form:"created"`
	Product      *string           `form:"product"`
}

// PlanParams is the set of parameters that can be used when creating or updating a plan.
// For more details see https://stripe.com/docs/api#create_plan and https://stripe.com/docs/api#update_plan.
type PlanParams struct {
	Params          `form:"*"`
	Active          *bool                     `form:"active"`
	AggregateUsage  *string                   `form:"aggregate_usage"`
	Amount          *int64                    `form:"amount"`
	BillingScheme   *string                   `form:"billing_scheme"`
	Currency        *string                   `form:"currency"`
	ID              *string                   `form:"id"`
	Interval        *string                   `form:"interval"`
	IntervalCount   *int64                    `form:"interval_count"`
	Nickname        *string                   `form:"nickname"`
	Product         *PlanProductParams        `form:"product"`
	ProductID       *string                   `form:"product"`
	Tiers           []*PlanTierParams         `form:"tiers,indexed"`
	TiersMode       *string                   `form:"tiers_mode"`
	TransformUsage  *PlanTransformUsageParams `form:"transform_usage"`
	TrialPeriodDays *int64                    `form:"trial_period_days"`
	UsageType       *string                   `form:"usage_type"`
}

// PlanTier configures tiered pricing
type PlanTier struct {
	Amount int64 `json:"amount"`
	UpTo   int64 `json:"up_to"`
}

// PlanTransformUsage represents the bucket billing configuration.
type PlanTransformUsage struct {
	DivideBy int64                   `json:"divide_by"`
	Round    PlanTransformUsageRound `json:"round"`
}

// PlanTransformUsageParams represents the bucket billing configuration.
type PlanTransformUsageParams struct {
	DivideBy *int64  `form:"divide_by"`
	Round    *string `form:"round"`
}

// PlanTierParams configures tiered pricing
type PlanTierParams struct {
	Params  `form:"*"`
	Amount  *int64 `form:"amount"`
	UpTo    *int64 `form:"-"` // handled in custom AppendTo
	UpToInf *bool  `form:"-"` // handled in custom AppendTo
}

// AppendTo implements custom up_to serialisation logic for tiers configuration
func (p *PlanTierParams) AppendTo(body *form.Values, keyParts []string) {
	if BoolValue(p.UpToInf) {
		body.Add(form.FormatKey(append(keyParts, "up_to")), "inf")
	} else {
		body.Add(form.FormatKey(append(keyParts, "up_to")), strconv.FormatInt(Int64Value(p.UpTo), 10))
	}
}

// PlanProductParams is the set of parameters that can be used when creating a product inside a plan
// This can only be used on plan creation and won't work on plan update.
// For more details see https://stripe.com/docs/api#create_plan-product and https://stripe.com/docs/api#update_plan-product
type PlanProductParams struct {
	ID                  *string           `form:"id"`
	Name                *string           `form:"name"`
	Metadata            map[string]string `form:"metadata"`
	StatementDescriptor *string           `form:"statement_descriptor"`
}
<|endoftext|>


<commit_before>package scaleway

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnflag"
	"github.com/docker/machine/libmachine/ssh"
	"github.com/docker/machine/libmachine/state"
	"github.com/scaleway/scaleway-cli/pkg/api"
	"github.com/scaleway/scaleway-cli/pkg/config"
)

const (
	// VERSION represents the semver version of the package
	VERSION           = "v1.0.2+dev"
	defaultImage      = "ubuntu-xenial"
	defaultBootscript = "docker"
)

var scwAPI *api.ScalewayAPI

// Driver represents the docker driver interface
type Driver struct {
	*drivers.BaseDriver
	ServerID       string
	Organization   string
	IPID           string
	Token          string
	commercialType string
	name           string
	image          string
	stopping       bool
	// size         string
	// userDataFile string
	// ipv6         bool
}

// DriverName returns the name of the driver
func (d *Driver) DriverName() string {
	return "scaleway"
}

func (d *Driver) getClient() (cl *api.ScalewayAPI, err error) {
	if scwAPI == nil {
		scwAPI, err = api.NewScalewayAPI(d.Organization, d.Token, "docker-machine-driver-scaleway/%v"+VERSION)
	}
	cl = scwAPI
	return
}

// SetConfigFromFlags sets the flags
func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) (err error) {
	if flags.Bool("scaleway-debug") {
		logrus.SetOutput(os.Stderr)
		logrus.SetLevel(logrus.DebugLevel)
	}

	d.Token, d.Organization = flags.String("scaleway-token"), flags.String("scaleway-organization")
	if d.Token == "" || d.Organization == "" {
		config, cfgErr := config.GetConfig()
		if cfgErr == nil {
			if d.Token == "" {
				d.Token = config.Token
			}
			if d.Organization == "" {
				d.Organization = config.Organization
			}
		} else {
			return fmt.Errorf("You must provide organization and token")
		}
	}
	d.commercialType = flags.String("scaleway-commercial-type")
	d.name = flags.String("scaleway-name")
	d.image = flags.String("scaleway-image")
	return
}

// NewDriver returns a new driver
func NewDriver(hostName, storePath string) *Driver {
	return &Driver{
		BaseDriver: &drivers.BaseDriver{},
	}
}

// GetCreateFlags registers the flags
func (d *Driver) GetCreateFlags() []mcnflag.Flag {
	return []mcnflag.Flag{
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_TOKEN",
			Name:   "scaleway-token",
			Usage:  "Scaleway token",
		},
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_ORGANIZATION",
			Name:   "scaleway-organization",
			Usage:  "Scaleway organization",
		},
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_NAME",
			Name:   "scaleway-name",
			Usage:  "Assign a name",
		},
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_COMMERCIAL_TYPE",
			Name:   "scaleway-commercial-type",
			Usage:  "Specifies the commercial type",
			Value:  "VC1S",
		},
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_IMAGE",
			Name:   "scaleway-image",
			Usage:  "Specifies the image",
			Value:  defaultImage,
		},
		mcnflag.BoolFlag{
			EnvVar: "SCALEWAY_DEBUG",
			Name:   "scaleway-debug",
			Usage:  "Enables Scaleway client debugging",
		},
		// mcnflag.StringFlag{
		//     EnvVar: "SCALEWAY_USERDATA",
		//     Name:   "scaleway-userdata",
		//     Usage:  "Path to file with user-data",
		// },
		// mcnflag.BoolFlag{
		// 	EnvVar: "SCALEWAY_IPV6",
		// 	Name:   "scaleway-ipv6",
		// 	Usage:  "Enable ipv6",
		// },
	}
}

// Create configures and starts a scaleway server
func (d *Driver) Create() (err error) {
	var publicKey []byte
	var cl *api.ScalewayAPI
	var ip *api.ScalewayGetIP

	log.Infof("Creating SSH key...")
	if err = ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
		return err
	}
	publicKey, err = ioutil.ReadFile(d.GetSSHKeyPath() + ".pub")
	if err != nil {
		return
	}
	log.Infof("Creating server...")
	cl, err = d.getClient()
	if err != nil {
		return
	}
	ip, err = cl.NewIP()
	if err != nil {
		return
	}
	d.IPAddress = ip.IP.Address
	d.IPID = ip.IP.ID
	d.ServerID, err = api.CreateServer(cl, &api.ConfigCreateServer{
		ImageName:      d.image,
		CommercialType: d.commercialType,
		Name:           d.name,
		Bootscript:     defaultBootscript,
		IP:             ip.IP.ID,
		Env: strings.Join([]string{"AUTHORIZED_KEY",
			strings.Replace(string(publicKey[:len(publicKey)-1]), " ", "_", -1)}, "="),
	})
	if err != nil {
		return
	}
	log.Infof("Starting server...")
	err = api.StartServer(cl, d.ServerID, false)
	return
}

// GetSSHHostname returns the IP of the server
func (d *Driver) GetSSHHostname() (string, error) {
	return d.IPAddress, nil
}

// GetState returns the state of the server
func (d *Driver) GetState() (st state.State, err error) {
	var server *api.ScalewayServer
	var cl *api.ScalewayAPI

	st = state.Error
	cl, err = d.getClient()
	if err != nil {
		return
	}
	server, err = cl.GetServer(d.ServerID)
	if err != nil {
		return
	}
	st = state.None
	switch server.State {
	case "starting":
		st = state.Starting
	case "running":
		st = state.Running
	case "stopping":
		st = state.Stopping
	case "stopped":
		st = state.Stopped
	}
	if d.stopping {
		time.Sleep(5 * time.Second)
	}
	return
}

// GetURL returns IP + docker port
func (d *Driver) GetURL() (string, error) {
	if err := drivers.MustBeRunning(d); err != nil {
		return "", err
	}
	return fmt.Sprintf("tcp://%s", net.JoinHostPort(d.IPAddress, "2376")), nil
}

func (d *Driver) postAction(action string) (err error) {
	var cl *api.ScalewayAPI

	cl, err = d.getClient()
	if err != nil {
		return
	}
	err = cl.PostServerAction(d.ServerID, action)
	return
}

// Kill does nothing
func (d *Driver) Kill() error {
	return errors.New("scaleway driver does not support kill")
}

// Remove shutdowns the server and removes the IP
func (d *Driver) Remove() (err error) {
	var cl *api.ScalewayAPI

	cl, err = d.getClient()
	if err != nil {
		return
	}
	err = cl.PostServerAction(d.ServerID, "terminate")
	if err != nil {
		return
	}
	for {
		_, err = cl.GetServer(d.ServerID)
		if err != nil {
			break
		}
	}
	err = cl.DeleteIP(d.IPID)
	return
}

// Restart reboots the server
func (d *Driver) Restart() error {
	return d.postAction("reboot")
}

// Start starts the server
func (d *Driver) Start() error {
	return d.postAction("poweron")
}

// Stop stops the server
func (d *Driver) Stop() error {
	d.stopping = true
	return d.postAction("poweroff")
}
<commit_msg>Wait after the boot that all services are done<commit_after>package scaleway

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnflag"
	"github.com/docker/machine/libmachine/ssh"
	"github.com/docker/machine/libmachine/state"
	"github.com/scaleway/scaleway-cli/pkg/api"
	"github.com/scaleway/scaleway-cli/pkg/config"
)

const (
	// VERSION represents the semver version of the package
	VERSION           = "v1.0.2+dev"
	defaultImage      = "ubuntu-xenial"
	defaultBootscript = "docker"
)

var scwAPI *api.ScalewayAPI

// Driver represents the docker driver interface
type Driver struct {
	*drivers.BaseDriver
	ServerID       string
	Organization   string
	IPID           string
	Token          string
	commercialType string
	name           string
	image          string
	stopping       bool
	created        bool
	// size         string
	// userDataFile string
	// ipv6         bool
}

// DriverName returns the name of the driver
func (d *Driver) DriverName() string {
	return "scaleway"
}

func (d *Driver) getClient() (cl *api.ScalewayAPI, err error) {
	if scwAPI == nil {
		scwAPI, err = api.NewScalewayAPI(d.Organization, d.Token, "docker-machine-driver-scaleway/%v"+VERSION)
	}
	cl = scwAPI
	return
}

// SetConfigFromFlags sets the flags
func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) (err error) {
	if flags.Bool("scaleway-debug") {
		logrus.SetOutput(os.Stderr)
		logrus.SetLevel(logrus.DebugLevel)
	}

	d.Token, d.Organization = flags.String("scaleway-token"), flags.String("scaleway-organization")
	if d.Token == "" || d.Organization == "" {
		config, cfgErr := config.GetConfig()
		if cfgErr == nil {
			if d.Token == "" {
				d.Token = config.Token
			}
			if d.Organization == "" {
				d.Organization = config.Organization
			}
		} else {
			return fmt.Errorf("You must provide organization and token")
		}
	}
	d.commercialType = flags.String("scaleway-commercial-type")
	d.name = flags.String("scaleway-name")
	d.image = flags.String("scaleway-image")
	return
}

// NewDriver returns a new driver
func NewDriver(hostName, storePath string) *Driver {
	return &Driver{
		BaseDriver: &drivers.BaseDriver{},
	}
}

// GetCreateFlags registers the flags
func (d *Driver) GetCreateFlags() []mcnflag.Flag {
	return []mcnflag.Flag{
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_TOKEN",
			Name:   "scaleway-token",
			Usage:  "Scaleway token",
		},
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_ORGANIZATION",
			Name:   "scaleway-organization",
			Usage:  "Scaleway organization",
		},
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_NAME",
			Name:   "scaleway-name",
			Usage:  "Assign a name",
		},
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_COMMERCIAL_TYPE",
			Name:   "scaleway-commercial-type",
			Usage:  "Specifies the commercial type",
			Value:  "VC1S",
		},
		mcnflag.StringFlag{
			EnvVar: "SCALEWAY_IMAGE",
			Name:   "scaleway-image",
			Usage:  "Specifies the image",
			Value:  defaultImage,
		},
		mcnflag.BoolFlag{
			EnvVar: "SCALEWAY_DEBUG",
			Name:   "scaleway-debug",
			Usage:  "Enables Scaleway client debugging",
		},
		// mcnflag.StringFlag{
		//     EnvVar: "SCALEWAY_USERDATA",
		//     Name:   "scaleway-userdata",
		//     Usage:  "Path to file with user-data",
		// },
		// mcnflag.BoolFlag{
		// 	EnvVar: "SCALEWAY_IPV6",
		// 	Name:   "scaleway-ipv6",
		// 	Usage:  "Enable ipv6",
		// },
	}
}

// Create configures and starts a scaleway server
func (d *Driver) Create() (err error) {
	var publicKey []byte
	var cl *api.ScalewayAPI
	var ip *api.ScalewayGetIP

	log.Infof("Creating SSH key...")
	if err = ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
		return err
	}
	publicKey, err = ioutil.ReadFile(d.GetSSHKeyPath() + ".pub")
	if err != nil {
		return
	}
	log.Infof("Creating server...")
	cl, err = d.getClient()
	if err != nil {
		return
	}
	ip, err = cl.NewIP()
	if err != nil {
		return
	}
	d.IPAddress = ip.IP.Address
	d.IPID = ip.IP.ID
	d.ServerID, err = api.CreateServer(cl, &api.ConfigCreateServer{
		ImageName:      d.image,
		CommercialType: d.commercialType,
		Name:           d.name,
		Bootscript:     defaultBootscript,
		IP:             ip.IP.ID,
		Env: strings.Join([]string{"AUTHORIZED_KEY",
			strings.Replace(string(publicKey[:len(publicKey)-1]), " ", "_", -1)}, "="),
	})
	if err != nil {
		return
	}
	log.Infof("Starting server...")
	err = api.StartServer(cl, d.ServerID, false)
	d.created = true
	return
}

// GetSSHHostname returns the IP of the server
func (d *Driver) GetSSHHostname() (string, error) {
	return d.IPAddress, nil
}

// GetState returns the state of the server
func (d *Driver) GetState() (st state.State, err error) {
	var server *api.ScalewayServer
	var cl *api.ScalewayAPI

	st = state.Error
	cl, err = d.getClient()
	if err != nil {
		return
	}
	server, err = cl.GetServer(d.ServerID)
	if err != nil {
		return
	}
	st = state.None
	switch server.State {
	case "starting":
		st = state.Starting
	case "running":
		if d.created {
			time.Sleep(60 * time.Second)
			d.created = false
		}
		st = state.Running
	case "stopping":
		st = state.Stopping
	case "stopped":
		st = state.Stopped
	}
	if d.stopping {
		time.Sleep(5 * time.Second)
	}
	return
}

// GetURL returns IP + docker port
func (d *Driver) GetURL() (string, error) {
	if err := drivers.MustBeRunning(d); err != nil {
		return "", err
	}
	return fmt.Sprintf("tcp://%s", net.JoinHostPort(d.IPAddress, "2376")), nil
}

func (d *Driver) postAction(action string) (err error) {
	var cl *api.ScalewayAPI

	cl, err = d.getClient()
	if err != nil {
		return
	}
	err = cl.PostServerAction(d.ServerID, action)
	return
}

// Kill does nothing
func (d *Driver) Kill() error {
	return errors.New("scaleway driver does not support kill")
}

// Remove shutdowns the server and removes the IP
func (d *Driver) Remove() (err error) {
	var cl *api.ScalewayAPI

	cl, err = d.getClient()
	if err != nil {
		return
	}
	err = cl.PostServerAction(d.ServerID, "terminate")
	if err != nil {
		return
	}
	for {
		_, err = cl.GetServer(d.ServerID)
		if err != nil {
			break
		}
	}
	err = cl.DeleteIP(d.IPID)
	return
}

// Restart reboots the server
func (d *Driver) Restart() error {
	return d.postAction("reboot")
}

// Start starts the server
func (d *Driver) Start() error {
	return d.postAction("poweron")
}

// Stop stops the server
func (d *Driver) Stop() error {
	d.stopping = true
	return d.postAction("poweroff")
}
<|endoftext|>


<commit_before>package client

import (
	"errors"
	"time"
)

var (
	TimeOutError = errors.New("Timeout")
)

type reply struct {
	data interface{}
	err  error
}

type TimeoutExt struct {
	client  Client
	timeout time.Duration
	fail    chan struct{}
	reply   chan reply
}

func NewTimeoutExt(timeout time.Duration) ClientExtension {
	return func(c Client) Client {
		return &TimeoutExt{
			client:  c,
			timeout: timeout,
			fail:    make(chan struct{}, 1),
			reply:   make(chan reply, 1),
		}
	}
}

func (e *TimeoutExt) Call(method string, params ...interface{}) (interface{}, error) {
	go func() {
		data, err := e.client.Call(method, params...)
		e.reply <- reply{data, err}
	}()
	go func() {
		time.Sleep(e.timeout)
		e.fail <- struct{}{}
	}()

	select {
	case <-e.fail:
		return nil, TimeOutError
	case res := <-e.reply:
		return res.data, res.err
	}
}
<commit_msg>Fix timeout extension<commit_after>package client

import (
	"errors"
	"time"
)

var (
	TimeOutError = errors.New("Timeout")
)

type reply struct {
	data interface{}
	err  error
}

type TimeoutExt struct {
	client  Client
	timeout time.Duration
	reply   chan reply
}

func NewTimeoutExt(timeout time.Duration) ClientExtension {
	return func(c Client) Client {
		return &TimeoutExt{
			client:  c,
			timeout: timeout,
			reply:   make(chan reply, 1),
		}
	}
}

func (e *TimeoutExt) Call(method string, params ...interface{}) (interface{}, error) {
	go func() {
		data, err := e.client.Call(method, params...)
		e.reply <- reply{data, err}
	}()

	select {
	case <-time.After(e.timeout * time.Second):
		return nil, TimeOutError
	case res := <-e.reply:
		return res.data, res.err
	}
}
<|endoftext|>


<commit_before>package main

import (
	"strconv"
	"time"
)

func getNextTuesday(t time.Time) time.Time {
	for int(t.Weekday()) != 2 { //todo figure out how to use the constant instead of special number 2! t.Weekday.Tuesday {
		t = t.AddDate(0, 0, 1)
	}

	return t
}

func AgendaDate(t time.Time) string {
	month := strconv.Itoa(int(t.Month()))
	day := strconv.Itoa(t.Day())
	year := strconv.Itoa(t.Year())

	return month + "." + day + "." + year
}

func DateWithSlashes(t time.Time) string {
	month := strconv.Itoa(int(t.Month()))
	day := strconv.Itoa(t.Day())
	year := strconv.Itoa(t.Year())

	return month + "/" + day + "/" + year
}

func AgendaMonthDayYear(t time.Time) string {
	day := strconv.Itoa(t.Day())
	month := t.Month().String()
	year := strconv.Itoa(t.Year())

	return month + " " + day + ", " + year
}

func prettyPrintTime(curTime time.Time, minToAdd int) (time.Time, string) {

	nextTime := time.Minute * time.Duration(minToAdd)
	curTime = curTime.Add(nextTime)
	hour, min, _ := curTime.Clock()

	return curTime, strconv.Itoa(hour) + ":" + strconv.Itoa(min)
}
<commit_msg>use formula for the next weekday instead of a loop.<commit_after>package main

import (
	"strconv"
	"time"
)

func getNextTuesday(t time.Time) time.Time {
	const TUESDAY = 2
	t = t.AddDate(0, 0, (TUESDAY+(7-int(t.Weekday())))%7)

	return t
}

func AgendaDate(t time.Time) string {
	month := strconv.Itoa(int(t.Month()))
	day := strconv.Itoa(t.Day())
	year := strconv.Itoa(t.Year())

	return month + "." + day + "." + year
}

func DateWithSlashes(t time.Time) string {
	month := strconv.Itoa(int(t.Month()))
	day := strconv.Itoa(t.Day())
	year := strconv.Itoa(t.Year())

	return month + "/" + day + "/" + year
}

func AgendaMonthDayYear(t time.Time) string {
	day := strconv.Itoa(t.Day())
	month := t.Month().String()
	year := strconv.Itoa(t.Year())

	return month + " " + day + ", " + year
}

func prettyPrintTime(curTime time.Time, minToAdd int) (time.Time, string) {

	nextTime := time.Minute * time.Duration(minToAdd)
	curTime = curTime.Add(nextTime)
	hour, min, _ := curTime.Clock()

	return curTime, strconv.Itoa(hour) + ":" + strconv.Itoa(min)
}
<|endoftext|>


<commit_before>package quic

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/internal/wire"
)

var (
	clientMuxerOnce sync.Once
	clientMuxer     multiplexer
)

type multiplexer interface {
	AddConn(net.PacketConn, int) (packetHandlerManager, error)
	AddHandler(net.PacketConn, protocol.ConnectionID, packetHandler) error
}

type connManager struct {
	connIDLen int
	manager   packetHandlerManager
}

// The clientMultiplexer listens on multiple net.PacketConns and dispatches
// incoming packets to the session handler.
type clientMultiplexer struct {
	mutex sync.Mutex

	conns                   map[net.PacketConn]connManager
	newPacketHandlerManager func() packetHandlerManager // so it can be replaced in the tests

	logger utils.Logger
}

var _ multiplexer = &clientMultiplexer{}

func getClientMultiplexer() multiplexer {
	clientMuxerOnce.Do(func() {
		clientMuxer = &clientMultiplexer{
			conns:                   make(map[net.PacketConn]connManager),
			logger:                  utils.DefaultLogger.WithPrefix("client muxer"),
			newPacketHandlerManager: newPacketHandlerMap,
		}
	})
	return clientMuxer
}

func (m *clientMultiplexer) AddConn(c net.PacketConn, connIDLen int) (packetHandlerManager, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	p, ok := m.conns[c]
	if !ok {
		manager := m.newPacketHandlerManager()
		p = connManager{connIDLen: connIDLen, manager: manager}
		m.conns[c] = p
		// If we didn't know this packet conn before, listen for incoming packets
		// and dispatch them to the right sessions.
		go m.listen(c, &p)
	}
	if p.connIDLen != connIDLen {
		return nil, fmt.Errorf("cannot use %d byte connection IDs on a connection that is already using %d byte connction IDs", connIDLen, p.connIDLen)
	}
	return p.manager, nil
}

func (m *clientMultiplexer) AddHandler(c net.PacketConn, connID protocol.ConnectionID, handler packetHandler) error {
	p, ok := m.conns[c]
	if !ok {
		return errors.New("unknown packet conn %s")
	}
	p.manager.Add(connID, handler)
	return nil
}

func (m *clientMultiplexer) listen(c net.PacketConn, p *connManager) {
	for {
		data := *getPacketBuffer()
		data = data[:protocol.MaxReceivePacketSize]
		// The packet size should not exceed protocol.MaxReceivePacketSize bytes
		// If it does, we only read a truncated packet, which will then end up undecryptable
		n, addr, err := c.ReadFrom(data)
		if err != nil {
			if !strings.HasSuffix(err.Error(), "use of closed network connection") {
				p.manager.Close(err)
			}
			return
		}
		data = data[:n]

		if err := m.handlePacket(addr, data, p); err != nil {
			m.logger.Debugf("error handling packet from %s: %s", addr, err)
		}
	}
}

func (m *clientMultiplexer) handlePacket(addr net.Addr, data []byte, p *connManager) error {
	rcvTime := time.Now()

	r := bytes.NewReader(data)
	iHdr, err := wire.ParseInvariantHeader(r, p.connIDLen)
	// drop the packet if we can't parse the header
	if err != nil {
		return fmt.Errorf("error parsing invariant header: %s", err)
	}
	client, ok := p.manager.Get(iHdr.DestConnectionID)
	if !ok {
		return fmt.Errorf("received a packet with an unexpected connection ID %s", iHdr.DestConnectionID)
	}
	if client == nil {
		// Late packet for closed session
		return nil
	}
	hdr, err := iHdr.Parse(r, protocol.PerspectiveServer, client.GetVersion())
	if err != nil {
		return fmt.Errorf("error parsing header: %s", err)
	}
	hdr.Raw = data[:len(data)-r.Len()]
	packetData := data[len(data)-r.Len():]

	client.handlePacket(&receivedPacket{
		remoteAddr: addr,
		header:     hdr,
		data:       packetData,
		rcvTime:    rcvTime,
	})
	return nil
}
<commit_msg>fix race condition when adding packet handlers to the multiplexer<commit_after>package quic

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/internal/wire"
)

var (
	clientMuxerOnce sync.Once
	clientMuxer     multiplexer
)

type multiplexer interface {
	AddConn(net.PacketConn, int) (packetHandlerManager, error)
	AddHandler(net.PacketConn, protocol.ConnectionID, packetHandler) error
}

type connManager struct {
	connIDLen int
	manager   packetHandlerManager
}

// The clientMultiplexer listens on multiple net.PacketConns and dispatches
// incoming packets to the session handler.
type clientMultiplexer struct {
	mutex sync.Mutex

	conns                   map[net.PacketConn]connManager
	newPacketHandlerManager func() packetHandlerManager // so it can be replaced in the tests

	logger utils.Logger
}

var _ multiplexer = &clientMultiplexer{}

func getClientMultiplexer() multiplexer {
	clientMuxerOnce.Do(func() {
		clientMuxer = &clientMultiplexer{
			conns:                   make(map[net.PacketConn]connManager),
			logger:                  utils.DefaultLogger.WithPrefix("client muxer"),
			newPacketHandlerManager: newPacketHandlerMap,
		}
	})
	return clientMuxer
}

func (m *clientMultiplexer) AddConn(c net.PacketConn, connIDLen int) (packetHandlerManager, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	p, ok := m.conns[c]
	if !ok {
		manager := m.newPacketHandlerManager()
		p = connManager{connIDLen: connIDLen, manager: manager}
		m.conns[c] = p
		// If we didn't know this packet conn before, listen for incoming packets
		// and dispatch them to the right sessions.
		go m.listen(c, &p)
	}
	if p.connIDLen != connIDLen {
		return nil, fmt.Errorf("cannot use %d byte connection IDs on a connection that is already using %d byte connction IDs", connIDLen, p.connIDLen)
	}
	return p.manager, nil
}

func (m *clientMultiplexer) AddHandler(c net.PacketConn, connID protocol.ConnectionID, handler packetHandler) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	p, ok := m.conns[c]
	if !ok {
		return errors.New("unknown packet conn %s")
	}
	p.manager.Add(connID, handler)
	return nil
}

func (m *clientMultiplexer) listen(c net.PacketConn, p *connManager) {
	for {
		data := *getPacketBuffer()
		data = data[:protocol.MaxReceivePacketSize]
		// The packet size should not exceed protocol.MaxReceivePacketSize bytes
		// If it does, we only read a truncated packet, which will then end up undecryptable
		n, addr, err := c.ReadFrom(data)
		if err != nil {
			if !strings.HasSuffix(err.Error(), "use of closed network connection") {
				p.manager.Close(err)
			}
			return
		}
		data = data[:n]

		if err := m.handlePacket(addr, data, p); err != nil {
			m.logger.Debugf("error handling packet from %s: %s", addr, err)
		}
	}
}

func (m *clientMultiplexer) handlePacket(addr net.Addr, data []byte, p *connManager) error {
	rcvTime := time.Now()

	r := bytes.NewReader(data)
	iHdr, err := wire.ParseInvariantHeader(r, p.connIDLen)
	// drop the packet if we can't parse the header
	if err != nil {
		return fmt.Errorf("error parsing invariant header: %s", err)
	}
	client, ok := p.manager.Get(iHdr.DestConnectionID)
	if !ok {
		return fmt.Errorf("received a packet with an unexpected connection ID %s", iHdr.DestConnectionID)
	}
	if client == nil {
		// Late packet for closed session
		return nil
	}
	hdr, err := iHdr.Parse(r, protocol.PerspectiveServer, client.GetVersion())
	if err != nil {
		return fmt.Errorf("error parsing header: %s", err)
	}
	hdr.Raw = data[:len(data)-r.Len()]
	packetData := data[len(data)-r.Len():]

	client.handlePacket(&receivedPacket{
		remoteAddr: addr,
		header:     hdr,
		data:       packetData,
		rcvTime:    rcvTime,
	})
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"io/ioutil"
	"log"
	"time"
	"github.com/bradfitz/gomemcache/memcache"
	"github.com/BurntSushi/toml"
)

type Config struct {
	Upstream string
	ReqFanFactor int
	TimeoutMS int
	MemcacheHosts []string
	MemcacheSeconds int
}

func snd (c chan string, s string) {
	go func () { c <- s }()
}

// start starts Charm up and returns a done channel for the done message
func start(confPath string) (chan string) {
	done := make(chan string)

	// print a welcome message
	log.Print("Charm is starting up.")
	// read the config file
	log.Printf(".   . Reading %s", confPath)
	tomlData, err := ioutil.ReadFile(confPath)
	if err != nil {
		snd(done, fmt.Sprintf("Could not read file at %s", confPath))
		return done
	}
	// populate the config struct
	log.Print(".   . Loading config")
	var conf Config
	_, err = toml.Decode(string(tomlData), &conf)
	if err != nil {
		snd(done, "Could not decode config")
		return done
	}
	// report on the configuration
	log.Print("Charm is configured!")
	log.Printf(".   . Stabilizing %v", conf.Upstream)
	log.Printf(".   . with %v duplicate requests", conf.ReqFanFactor)
	log.Printf(".   . and a %v milisecond timeout.", conf.TimeoutMS)
	log.Printf(
		".   . memcached at %v for %v seconds.",
		conf.MemcacheHosts,
		conf.MemcacheSeconds,
	)
	go run(conf, done)
	return done
}

type stableTransport struct {
	wrappedTransport http.RoundTripper
	reqFanFactor int
	cacheResponse chan *http.Response
}

// stableTransport.RoundTrip makes many round trips and returns the first
// response
func (t *stableTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	c := make(chan *http.Response)
	if t.wrappedTransport == nil {
		t.wrappedTransport = http.DefaultTransport
	}
	// fan out requests, send responses to the channel, log errors, don't
	// wait very long for someone to recieve our response
	for i := 0; i < t.reqFanFactor; i++ {
		go func () {
			resp, err := t.wrappedTransport.RoundTrip(r)
			if err != nil {
				log.Printf("transport-error: %v", err)
			} else {
				select {
				case c <- resp:
					// they were still waiting for the first
					// response and they recieved it from c
					return
				case <-time.After(1 * time.Millisecond):
					// no one was waiting to recieve from c
					// so this is not the first response
					return
				}
			}
		}()
	}

	// wait to reviece the first response
	first := <-c
	// copy the first respnse for caching
	cacheCopy := new(http.Response)
	*cacheCopy = *first
	if first.Body != nil {
		bodyBytes, err := ioutil.ReadAll(first.Body)
		if err != nil {
			log.Fatal("error reading response body:", err)
		}
		cacheBytes := make([]byte, len(bodyBytes))
		copy(cacheBytes, bodyBytes)
		first.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
		cacheCopy.Body = ioutil.NopCloser(bytes.NewBuffer(cacheBytes))
	}
	// send the copy to be cached if they want it
	go func () {
		select {
		case t.cacheResponse <- cacheCopy:
			return
		case <-time.After(15 * time.Millisecond):
			return
		}
	}()
	// return the first response to the handler
	return first, nil
}

// cacheKey returns a string to be used as the cache key for a request
func cacheKey(r *http.Request) (string, error) {
	// We need to be careful here.
	// There is serious potential to accidently ignore permissions if we
	// cache requests too broadly. For example, if our cache key is the path
	// and a super-admin cache-misses on /some/restricted/path then a
	// restricted user could be given the cached result from that super
	//  admin request.

	keyStr := ""
	keyStr += r.Method
	keyStr += r.URL.Host
	keyStr += r.URL.Path
	keyStr += r.URL.RawQuery

	// TODO: extract which headers to cache on into a config option
	keyStr += r.Header["X-Forwarded-Email"][0]

	key := sha256.Sum224([]byte(keyStr))
	return hex.EncodeToString(key[:sha256.Size224]), nil
}

// Conf.ServeHTTP checks memcache then proxies/caches with a stable transport
func (conf Config) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// check memcache
	mc := memcache.New(conf.MemcacheHosts...)
	key, err := cacheKey(r)
	if err != nil {
		log.Println("cache key error for Request:", r)
	} else {
		item, err := mc.Get(key)
		if err == nil { //cache hit
			log.Println(
				"INFO: cache hit!",
				r.Method,
				r.URL.Host,
				r.URL.Path,
				r.URL.RawQuery,
				r.Header["X-Forwarded-Email"],
			)
			// get the cached response
			response, err := http.ReadResponse(
				bufio.NewReader(bytes.NewReader(item.Value)),
				r,
			)
			if err == nil {
				// write the response to the response writer
				response.Write(w)
				return
			}
		}
	}
	// cache miss
	log.Println(
		"INFO: cache miss!",
		r.Method,
		r.URL.Host,
		r.URL.Path,
		r.URL.RawQuery,
		r.Header["X-Forwarded-Email"],
	)
	upstreamURL, err := url.Parse(conf.Upstream)
	if err != nil {
		log.Fatal("error parsing Upstream URL", conf.Upstream)
	}
	responseChan := make(chan *http.Response)
	proxy := httputil.NewSingleHostReverseProxy(upstreamURL)
	proxy.Transport = &stableTransport{
		proxy.Transport,
		conf.ReqFanFactor,
		responseChan,
	}
	proxy.ServeHTTP(w, r)

	// if the transport has a response waiting on the channel, cache it if
	// we have a cache
	if cacheKey != nil {
		select {
		case resp := <- responseChan:
			log.Println("got response to cache", resp)
			dump, err := httputil.DumpResponse(resp, true)
			if err != nil {
				log.Println(
					"ERROR: couldn't dump response:",
					err,
				)
				return
			}
			item := &memcache.Item{
				Key: key,
				Value: dump,
				Expiration: int32(conf.MemcacheSeconds),
			}
			err = mc.Set(item)
			if err != nil {
				log.Println("ERROR: memcached set error:", err)
				return
			}
		case <-time.After(1 * time.Millisecond):
			return
		}
	}

}

// run
func run(conf Config, done chan string) {
	// serve the config under a timeout
	timeout := time.Duration(conf.TimeoutMS) * time.Millisecond
	log.Fatal(http.ListenAndServe(
		":8000",
		http.TimeoutHandler(conf, timeout, "upstream timeout"),
	))
}

func main() {
	// start Charm,
	done := start("/secret/charm.conf")
	// when Charm is done, log the message and quit.
        log.Print(<-done)
}
<commit_msg>remove another log<commit_after>package main

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"io/ioutil"
	"log"
	"time"
	"github.com/bradfitz/gomemcache/memcache"
	"github.com/BurntSushi/toml"
)

type Config struct {
	Upstream string
	ReqFanFactor int
	TimeoutMS int
	MemcacheHosts []string
	MemcacheSeconds int
}

func snd (c chan string, s string) {
	go func () { c <- s }()
}

// start starts Charm up and returns a done channel for the done message
func start(confPath string) (chan string) {
	done := make(chan string)

	// print a welcome message
	log.Print("Charm is starting up.")
	// read the config file
	log.Printf(".   . Reading %s", confPath)
	tomlData, err := ioutil.ReadFile(confPath)
	if err != nil {
		snd(done, fmt.Sprintf("Could not read file at %s", confPath))
		return done
	}
	// populate the config struct
	log.Print(".   . Loading config")
	var conf Config
	_, err = toml.Decode(string(tomlData), &conf)
	if err != nil {
		snd(done, "Could not decode config")
		return done
	}
	// report on the configuration
	log.Print("Charm is configured!")
	log.Printf(".   . Stabilizing %v", conf.Upstream)
	log.Printf(".   . with %v duplicate requests", conf.ReqFanFactor)
	log.Printf(".   . and a %v milisecond timeout.", conf.TimeoutMS)
	log.Printf(
		".   . memcached at %v for %v seconds.",
		conf.MemcacheHosts,
		conf.MemcacheSeconds,
	)
	go run(conf, done)
	return done
}

type stableTransport struct {
	wrappedTransport http.RoundTripper
	reqFanFactor int
	cacheResponse chan *http.Response
}

// stableTransport.RoundTrip makes many round trips and returns the first
// response
func (t *stableTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	c := make(chan *http.Response)
	if t.wrappedTransport == nil {
		t.wrappedTransport = http.DefaultTransport
	}
	// fan out requests, send responses to the channel, log errors, don't
	// wait very long for someone to recieve our response
	for i := 0; i < t.reqFanFactor; i++ {
		go func () {
			resp, err := t.wrappedTransport.RoundTrip(r)
			if err != nil {
				log.Printf("transport-error: %v", err)
			} else {
				select {
				case c <- resp:
					// they were still waiting for the first
					// response and they recieved it from c
					return
				case <-time.After(1 * time.Millisecond):
					// no one was waiting to recieve from c
					// so this is not the first response
					return
				}
			}
		}()
	}

	// wait to reviece the first response
	first := <-c
	// copy the first respnse for caching
	cacheCopy := new(http.Response)
	*cacheCopy = *first
	if first.Body != nil {
		bodyBytes, err := ioutil.ReadAll(first.Body)
		if err != nil {
			log.Fatal("error reading response body:", err)
		}
		cacheBytes := make([]byte, len(bodyBytes))
		copy(cacheBytes, bodyBytes)
		first.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
		cacheCopy.Body = ioutil.NopCloser(bytes.NewBuffer(cacheBytes))
	}
	// send the copy to be cached if they want it
	go func () {
		select {
		case t.cacheResponse <- cacheCopy:
			return
		case <-time.After(15 * time.Millisecond):
			return
		}
	}()
	// return the first response to the handler
	return first, nil
}

// cacheKey returns a string to be used as the cache key for a request
func cacheKey(r *http.Request) (string, error) {
	// We need to be careful here.
	// There is serious potential to accidently ignore permissions if we
	// cache requests too broadly. For example, if our cache key is the path
	// and a super-admin cache-misses on /some/restricted/path then a
	// restricted user could be given the cached result from that super
	//  admin request.

	keyStr := ""
	keyStr += r.Method
	keyStr += r.URL.Host
	keyStr += r.URL.Path
	keyStr += r.URL.RawQuery

	// TODO: extract which headers to cache on into a config option
	keyStr += r.Header["X-Forwarded-Email"][0]

	key := sha256.Sum224([]byte(keyStr))
	return hex.EncodeToString(key[:sha256.Size224]), nil
}

// Conf.ServeHTTP checks memcache then proxies/caches with a stable transport
func (conf Config) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// check memcache
	mc := memcache.New(conf.MemcacheHosts...)
	key, err := cacheKey(r)
	if err != nil {
		log.Println("cache key error for Request:", r)
	} else {
		item, err := mc.Get(key)
		if err == nil { //cache hit
			log.Println(
				"INFO: cache hit!",
				r.Method,
				r.URL.Host,
				r.URL.Path,
				r.URL.RawQuery,
				r.Header["X-Forwarded-Email"],
			)
			// get the cached response
			response, err := http.ReadResponse(
				bufio.NewReader(bytes.NewReader(item.Value)),
				r,
			)
			if err == nil {
				// write the response to the response writer
				response.Write(w)
				return
			}
		}
	}
	// cache miss
	log.Println(
		"INFO: cache miss!",
		r.Method,
		r.URL.Host,
		r.URL.Path,
		r.URL.RawQuery,
		r.Header["X-Forwarded-Email"],
	)
	upstreamURL, err := url.Parse(conf.Upstream)
	if err != nil {
		log.Fatal("error parsing Upstream URL", conf.Upstream)
	}
	responseChan := make(chan *http.Response)
	proxy := httputil.NewSingleHostReverseProxy(upstreamURL)
	proxy.Transport = &stableTransport{
		proxy.Transport,
		conf.ReqFanFactor,
		responseChan,
	}
	proxy.ServeHTTP(w, r)

	// if the transport has a response waiting on the channel, cache it if
	// we have a cache
	if cacheKey != nil {
		select {
		case resp := <- responseChan:
			dump, err := httputil.DumpResponse(resp, true)
			if err != nil {
				log.Println(
					"ERROR: couldn't dump response:",
					err,
				)
				return
			}
			item := &memcache.Item{
				Key: key,
				Value: dump,
				Expiration: int32(conf.MemcacheSeconds),
			}
			err = mc.Set(item)
			if err != nil {
				log.Println("ERROR: memcached set error:", err)
				return
			}
		case <-time.After(1 * time.Millisecond):
			return
		}
	}

}

// run
func run(conf Config, done chan string) {
	// serve the config under a timeout
	timeout := time.Duration(conf.TimeoutMS) * time.Millisecond
	log.Fatal(http.ListenAndServe(
		":8000",
		http.TimeoutHandler(conf, timeout, "upstream timeout"),
	))
}

func main() {
	// start Charm,
	done := start("/secret/charm.conf")
	// when Charm is done, log the message and quit.
        log.Print(<-done)
}
<|endoftext|>


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

package datastore

import (
	"encoding/json"
	"fmt"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
	. "go.chromium.org/luci/common/testing/assertions"
)

func ShouldEqualKey(actual interface{}, expected ...interface{}) string {
	if len(expected) != 1 {
		return fmt.Sprintf("Assertion requires 1 expected value, got %d", len(expected))
	}
	if actual.(*Key).Equal(expected[0].(*Key)) {
		return ""
	}
	return fmt.Sprintf("Expected: %q\nActual: %q", actual, expected[0])
}

func TestKeyEncode(t *testing.T) {
	t.Parallel()

	kc := MkKeyContext("appid", "ns")
	keys := []*Key{
		kc.MakeKey("kind", 1),
		kc.MakeKey("nerd", "moo"),
		kc.MakeKey("parent", 10, "renerd", "moo"),
	}

	Convey("Key Round trip", t, func() {
		for _, k := range keys {
			k := k
			Convey(k.String(), func() {
				enc := k.Encode()
				dec, err := NewKeyEncoded(enc)
				So(err, ShouldBeNil)
				So(dec, ShouldNotBeNil)
				So(dec, ShouldEqualKey, k)

				dec2, err := NewKeyEncoded(enc)
				So(err, ShouldBeNil)
				So(dec2, ShouldEqualKey, dec)
				So(dec2, ShouldEqualKey, k)
			})

			Convey(k.String()+" (json)", func() {
				data, err := k.MarshalJSON()
				So(err, ShouldBeNil)

				dec := &Key{}
				So(dec.UnmarshalJSON(data), ShouldBeNil)
				So(dec, ShouldEqualKey, k)
			})
		}
	})

	Convey("NewKey", t, func() {
		Convey("single", func() {
			k := MkKeyContext("appid", "ns").NewKey("kind", "", 1, nil)
			So(k, ShouldEqualKey, keys[0])
		})

		Convey("empty", func() {
			So(MkKeyContext("appid", "ns").NewKeyToks(nil), ShouldBeNil)
		})

		Convey("nest", func() {
			kc := MkKeyContext("appid", "ns")
			k := kc.NewKey("renerd", "moo", 0, kc.NewKey("parent", "", 10, nil))
			So(k, ShouldEqualKey, keys[2])
		})
	})

	Convey("Key bad encoding", t, func() {
		Convey("extra junk before", func() {
			enc := keys[2].Encode()
			_, err := NewKeyEncoded("/" + enc)
			So(err, ShouldErrLike, "illegal base64")
		})

		Convey("extra junk after", func() {
			enc := keys[2].Encode()
			_, err := NewKeyEncoded(enc[:len(enc)-1])
			So(err, ShouldErrLike, "EOF")
		})

		Convey("json encoding includes quotes", func() {
			data, err := keys[0].MarshalJSON()
			So(err, ShouldBeNil)

			dec := &Key{}
			err = dec.UnmarshalJSON(append(data, '!'))
			So(err, ShouldErrLike, "bad JSON key")
		})
	})
}

func TestKeyValidity(t *testing.T) {
	t.Parallel()

	Convey("keys validity", t, func() {
		kc := MkKeyContext("aid", "ns")

		Convey("incomplete", func() {
			So(kc.MakeKey("kind", 1).IsIncomplete(), ShouldBeFalse)
			So(kc.MakeKey("kind", 0).IsIncomplete(), ShouldBeTrue)
		})

		Convey("invalid", func() {
			So(kc.MakeKey("hat", "face", "__kind__", 1).Valid(true, kc), ShouldBeTrue)

			bads := []*Key{
				MkKeyContext("aid", "ns").NewKeyToks([]KeyTok{{"Kind", 1, "1"}}),
				MkKeyContext("", "ns").MakeKey("", "ns", "hat", "face"),
				kc.MakeKey("base", 1, "", "id"),
				kc.MakeKey("hat", "face", "__kind__", 1),
				kc.MakeKey("hat", 0, "kind", 1),
			}
			for _, k := range bads {
				Convey(k.String(), func() {
					So(k.Valid(false, kc), ShouldBeFalse)
				})
			}
		})

		Convey("partially valid", func() {
			So(kc.MakeKey("kind", "").PartialValid(kc), ShouldBeTrue)
			So(kc.MakeKey("kind", "", "child", "").PartialValid(kc), ShouldBeFalse)
		})
	})
}

func TestMiscKey(t *testing.T) {
	t.Parallel()

	Convey("KeyRoot", t, func() {
		kc := MkKeyContext("appid", "ns")

		k := kc.MakeKey("parent", 10, "renerd", "moo")
		r := kc.MakeKey("parent", 10)
		So(k.Root(), ShouldEqualKey, r)
	})

	Convey("KeysEqual", t, func() {
		kc := MkKeyContext("a", "n")

		k1 := kc.MakeKey("knd", 1)
		k2 := kc.MakeKey("knd", 1)
		So(k1.Equal(k2), ShouldBeTrue)
		k3 := kc.MakeKey("knd", 2)
		So(k1.Equal(k3), ShouldBeFalse)
	})

	Convey("KeyString", t, func() {
		kc := MkKeyContext("a", "n")

		k1 := kc.MakeKey("knd", 1, "other", "wat")
		So(k1.String(), ShouldEqual, "a:n:/knd,1/other,\"wat\"")
	})

	Convey("HasAncestor", t, func() {
		kc := MkKeyContext("a", "n")

		k1 := kc.MakeKey("kind", 1)
		k2 := kc.MakeKey("kind", 1, "other", "wat")
		k3 := kc.MakeKey("kind", 1, "other", "wat", "extra", "data")
		k4 := MkKeyContext("something", "n").MakeKey("kind", 1)
		k5 := kc.MakeKey("kind", 1, "other", "meep")

		So(k1.HasAncestor(k1), ShouldBeTrue)
		So(k1.HasAncestor(k2), ShouldBeFalse)
		So(k2.HasAncestor(k5), ShouldBeFalse)
		So(k5.HasAncestor(k2), ShouldBeFalse)
		So(k2.HasAncestor(k1), ShouldBeTrue)
		So(k3.HasAncestor(k2), ShouldBeTrue)
		So(k3.HasAncestor(k1), ShouldBeTrue)
		So(k3.HasAncestor(k4), ShouldBeFalse)
	})

	Convey("*GenericKey supports json encoding", t, func() {
		type TestStruct struct {
			Key *Key
		}
		t := &TestStruct{
			MkKeyContext("aid", "ns").NewKey("kind", "id", 0,
				MkKeyContext("aid", "ns").NewKey("parent", "", 1, nil),
			)}
		d, err := json.Marshal(t)
		So(err, ShouldBeNil)
		t2 := &TestStruct{}
		err = json.Unmarshal(d, t2)
		So(err, ShouldBeNil)
		So(t.Key, ShouldEqualKey, t2.Key)
	})
}

func shouldBeLess(actual interface{}, expected ...interface{}) string {
	a, b := actual.(*Key), expected[0].(*Key)
	if a.Less(b) {
		return ""
	}
	return fmt.Sprintf("Expected %q < %q", a.String(), b.String())
}

func shouldNotBeEqual(actual interface{}, expected ...interface{}) string {
	a, b := actual.(*Key), expected[0].(*Key)
	if !a.Equal(b) {
		return ""
	}
	return fmt.Sprintf("Expected %q != %q", a.String(), b.String())
}

func shouldNotBeLess(actual interface{}, expected ...interface{}) string {
	a, b := actual.(*Key), expected[0].(*Key)
	if !a.Less(b) {
		return ""
	}
	return fmt.Sprintf("Expected !(%q < %q)", a.String(), b.String())
}

func TestKeySort(t *testing.T) {
	t.Parallel()

	Convey("KeyTok.Less() works", t, func() {
		So((KeyTok{"a", 0, "1"}).Less(KeyTok{"b", 0, "2"}), ShouldBeTrue)
		So((KeyTok{"b", 0, "1"}).Less(KeyTok{"a", 0, "2"}), ShouldBeFalse)
		So((KeyTok{"kind", 0, "1"}).Less(KeyTok{"kind", 0, "2"}), ShouldBeTrue)
		So((KeyTok{"kind", 1, ""}).Less(KeyTok{"kind", 2, ""}), ShouldBeTrue)
		So((KeyTok{"kind", 1, ""}).Less(KeyTok{"kind", 0, "1"}), ShouldBeTrue)
	})

	Convey("Key comparison works", t, func() {
		s := []*Key{
			MkKeyContext("A", "").MakeKey("kind", 1),
			MkKeyContext("A", "n").MakeKey("kind", 1),
			MkKeyContext("A", "n").MakeKey("kind", 1, "something", "else"),
			MkKeyContext("A", "n").MakeKey("kind", "1"),
			MkKeyContext("A", "n").MakeKey("kind", "1", "something", "else"),
			MkKeyContext("A", "n").MakeKey("other", 1, "something", "else"),
			MkKeyContext("a", "").MakeKey("kind", 1),
			MkKeyContext("a", "n").MakeKey("kind", 1),
			MkKeyContext("a", "n").MakeKey("kind", 2),
			MkKeyContext("a", "p").MakeKey("aleph", 1),
			MkKeyContext("b", "n").MakeKey("kind", 2),
		}

		for i := 1; i < len(s); i++ {
			So(s[i-1], shouldBeLess, s[i])
			So(s[i-1], shouldNotBeEqual, s[i])
			So(s[i], shouldNotBeEqual, s[i-1])
			So(s[i], shouldNotBeLess, s[i-1])
		}
	})
}
<commit_msg>[datastore] Fix test.<commit_after>// Copyright 2015 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package datastore

import (
	"encoding/json"
	"fmt"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
	. "go.chromium.org/luci/common/testing/assertions"
)

func ShouldEqualKey(actual interface{}, expected ...interface{}) string {
	if len(expected) != 1 {
		return fmt.Sprintf("Assertion requires 1 expected value, got %d", len(expected))
	}
	if actual.(*Key).Equal(expected[0].(*Key)) {
		return ""
	}
	return fmt.Sprintf("Expected: %q\nActual: %q", actual, expected[0])
}

func TestKeyEncode(t *testing.T) {
	t.Parallel()

	kc := MkKeyContext("appid", "ns")
	keys := []*Key{
		kc.MakeKey("kind", 1),
		kc.MakeKey("nerd", "moo"),
		kc.MakeKey("parent", 10, "renerd", "moo"),
	}

	Convey("Key Round trip", t, func() {
		for _, k := range keys {
			k := k
			Convey(k.String(), func() {
				enc := k.Encode()
				dec, err := NewKeyEncoded(enc)
				So(err, ShouldBeNil)
				So(dec, ShouldNotBeNil)
				So(dec, ShouldEqualKey, k)

				dec2, err := NewKeyEncoded(enc)
				So(err, ShouldBeNil)
				So(dec2, ShouldEqualKey, dec)
				So(dec2, ShouldEqualKey, k)
			})

			Convey(k.String()+" (json)", func() {
				data, err := k.MarshalJSON()
				So(err, ShouldBeNil)

				dec := &Key{}
				So(dec.UnmarshalJSON(data), ShouldBeNil)
				So(dec, ShouldEqualKey, k)
			})
		}
	})

	Convey("NewKey", t, func() {
		Convey("single", func() {
			k := MkKeyContext("appid", "ns").NewKey("kind", "", 1, nil)
			So(k, ShouldEqualKey, keys[0])
		})

		Convey("empty", func() {
			So(MkKeyContext("appid", "ns").NewKeyToks(nil), ShouldBeNil)
		})

		Convey("nest", func() {
			kc := MkKeyContext("appid", "ns")
			k := kc.NewKey("renerd", "moo", 0, kc.NewKey("parent", "", 10, nil))
			So(k, ShouldEqualKey, keys[2])
		})
	})

	Convey("Key bad encoding", t, func() {
		Convey("extra junk before", func() {
			enc := keys[2].Encode()
			_, err := NewKeyEncoded("/" + enc)
			So(err, ShouldErrLike, "illegal base64")
		})

		Convey("extra junk after", func() {
			enc := keys[2].Encode()
			_, err := NewKeyEncoded(enc[:len(enc)-1])
			So(err, ShouldNotBeNil)
		})

		Convey("json encoding includes quotes", func() {
			data, err := keys[0].MarshalJSON()
			So(err, ShouldBeNil)

			dec := &Key{}
			err = dec.UnmarshalJSON(append(data, '!'))
			So(err, ShouldErrLike, "bad JSON key")
		})
	})
}

func TestKeyValidity(t *testing.T) {
	t.Parallel()

	Convey("keys validity", t, func() {
		kc := MkKeyContext("aid", "ns")

		Convey("incomplete", func() {
			So(kc.MakeKey("kind", 1).IsIncomplete(), ShouldBeFalse)
			So(kc.MakeKey("kind", 0).IsIncomplete(), ShouldBeTrue)
		})

		Convey("invalid", func() {
			So(kc.MakeKey("hat", "face", "__kind__", 1).Valid(true, kc), ShouldBeTrue)

			bads := []*Key{
				MkKeyContext("aid", "ns").NewKeyToks([]KeyTok{{"Kind", 1, "1"}}),
				MkKeyContext("", "ns").MakeKey("", "ns", "hat", "face"),
				kc.MakeKey("base", 1, "", "id"),
				kc.MakeKey("hat", "face", "__kind__", 1),
				kc.MakeKey("hat", 0, "kind", 1),
			}
			for _, k := range bads {
				Convey(k.String(), func() {
					So(k.Valid(false, kc), ShouldBeFalse)
				})
			}
		})

		Convey("partially valid", func() {
			So(kc.MakeKey("kind", "").PartialValid(kc), ShouldBeTrue)
			So(kc.MakeKey("kind", "", "child", "").PartialValid(kc), ShouldBeFalse)
		})
	})
}

func TestMiscKey(t *testing.T) {
	t.Parallel()

	Convey("KeyRoot", t, func() {
		kc := MkKeyContext("appid", "ns")

		k := kc.MakeKey("parent", 10, "renerd", "moo")
		r := kc.MakeKey("parent", 10)
		So(k.Root(), ShouldEqualKey, r)
	})

	Convey("KeysEqual", t, func() {
		kc := MkKeyContext("a", "n")

		k1 := kc.MakeKey("knd", 1)
		k2 := kc.MakeKey("knd", 1)
		So(k1.Equal(k2), ShouldBeTrue)
		k3 := kc.MakeKey("knd", 2)
		So(k1.Equal(k3), ShouldBeFalse)
	})

	Convey("KeyString", t, func() {
		kc := MkKeyContext("a", "n")

		k1 := kc.MakeKey("knd", 1, "other", "wat")
		So(k1.String(), ShouldEqual, "a:n:/knd,1/other,\"wat\"")
	})

	Convey("HasAncestor", t, func() {
		kc := MkKeyContext("a", "n")

		k1 := kc.MakeKey("kind", 1)
		k2 := kc.MakeKey("kind", 1, "other", "wat")
		k3 := kc.MakeKey("kind", 1, "other", "wat", "extra", "data")
		k4 := MkKeyContext("something", "n").MakeKey("kind", 1)
		k5 := kc.MakeKey("kind", 1, "other", "meep")

		So(k1.HasAncestor(k1), ShouldBeTrue)
		So(k1.HasAncestor(k2), ShouldBeFalse)
		So(k2.HasAncestor(k5), ShouldBeFalse)
		So(k5.HasAncestor(k2), ShouldBeFalse)
		So(k2.HasAncestor(k1), ShouldBeTrue)
		So(k3.HasAncestor(k2), ShouldBeTrue)
		So(k3.HasAncestor(k1), ShouldBeTrue)
		So(k3.HasAncestor(k4), ShouldBeFalse)
	})

	Convey("*GenericKey supports json encoding", t, func() {
		type TestStruct struct {
			Key *Key
		}
		t := &TestStruct{
			MkKeyContext("aid", "ns").NewKey("kind", "id", 0,
				MkKeyContext("aid", "ns").NewKey("parent", "", 1, nil),
			)}
		d, err := json.Marshal(t)
		So(err, ShouldBeNil)
		t2 := &TestStruct{}
		err = json.Unmarshal(d, t2)
		So(err, ShouldBeNil)
		So(t.Key, ShouldEqualKey, t2.Key)
	})
}

func shouldBeLess(actual interface{}, expected ...interface{}) string {
	a, b := actual.(*Key), expected[0].(*Key)
	if a.Less(b) {
		return ""
	}
	return fmt.Sprintf("Expected %q < %q", a.String(), b.String())
}

func shouldNotBeEqual(actual interface{}, expected ...interface{}) string {
	a, b := actual.(*Key), expected[0].(*Key)
	if !a.Equal(b) {
		return ""
	}
	return fmt.Sprintf("Expected %q != %q", a.String(), b.String())
}

func shouldNotBeLess(actual interface{}, expected ...interface{}) string {
	a, b := actual.(*Key), expected[0].(*Key)
	if !a.Less(b) {
		return ""
	}
	return fmt.Sprintf("Expected !(%q < %q)", a.String(), b.String())
}

func TestKeySort(t *testing.T) {
	t.Parallel()

	Convey("KeyTok.Less() works", t, func() {
		So((KeyTok{"a", 0, "1"}).Less(KeyTok{"b", 0, "2"}), ShouldBeTrue)
		So((KeyTok{"b", 0, "1"}).Less(KeyTok{"a", 0, "2"}), ShouldBeFalse)
		So((KeyTok{"kind", 0, "1"}).Less(KeyTok{"kind", 0, "2"}), ShouldBeTrue)
		So((KeyTok{"kind", 1, ""}).Less(KeyTok{"kind", 2, ""}), ShouldBeTrue)
		So((KeyTok{"kind", 1, ""}).Less(KeyTok{"kind", 0, "1"}), ShouldBeTrue)
	})

	Convey("Key comparison works", t, func() {
		s := []*Key{
			MkKeyContext("A", "").MakeKey("kind", 1),
			MkKeyContext("A", "n").MakeKey("kind", 1),
			MkKeyContext("A", "n").MakeKey("kind", 1, "something", "else"),
			MkKeyContext("A", "n").MakeKey("kind", "1"),
			MkKeyContext("A", "n").MakeKey("kind", "1", "something", "else"),
			MkKeyContext("A", "n").MakeKey("other", 1, "something", "else"),
			MkKeyContext("a", "").MakeKey("kind", 1),
			MkKeyContext("a", "n").MakeKey("kind", 1),
			MkKeyContext("a", "n").MakeKey("kind", 2),
			MkKeyContext("a", "p").MakeKey("aleph", 1),
			MkKeyContext("b", "n").MakeKey("kind", 2),
		}

		for i := 1; i < len(s); i++ {
			So(s[i-1], shouldBeLess, s[i])
			So(s[i-1], shouldNotBeEqual, s[i])
			So(s[i], shouldNotBeEqual, s[i-1])
			So(s[i], shouldNotBeLess, s[i-1])
		}
	})
}
<|endoftext|>