<commit_before>package doc

import (
	"bytes"
	"text/template"

	"strings"

	"fmt"

	"github.com/ninedraft/boxofstuff/str"
	"github.com/octago/sflags"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

const (
	ListFormat       = "{{.Path}}"
	MarkdownTemplate = `#### <a name="{{.Link}}">{{.Path}}</a>` + "\n\n" +
		"**Description**:\n\n{{.Description}}\n\n" +
		"**Example**:\n\n{{.Example}}\n\n" +
		"**Flags**:\n\n" +
		"| Short | Name | Usage | Default value |\n" +
		"| ----- | ---- | ----- | ------------- |\n" +
		"{{range .Flags}}" +
		"| {{if .Short}}-{{.Short}}{{end}} " +
		"| --{{.Name}} " +
		"| {{.Usage}} " +
		"| {{if ne .DefValue \"[]\"}}{{.DefValue}}{{end}} " +
		"|\n" +
		"{{end}}\n\n" +
		"**Subcommands**:\n\n" +
		"{{range .Subcommands}}" +
		"* **[{{.Name}}](#{{.Link}})** {{.ShortDescription}}\n" +
		"{{end}}\n\n"
	TextTamplate = "Command: {{.Path}}\n" +
		"Description:\n{{.Description}}\n" +
		"Example:\n{{.Example}}\n" +
		"Flags:\n" +
		"{{range .Flags}}" +
		"{{if .Short}}-{{.Short}}{{else}}  {{end}} " +
		"--{{.Name}} " +
		"{{.Usage}} " +
		"{{if ne .DefValue \"[]\"}}{{.DefValue}}{{end}} " +
		"\n" +
		"{{end}}\n"
)

type Command struct {
	cobra.Command
}

type Doc struct {
	Link        string        `json:"link"`
	Path        string        `json:"path"`
	Name        string        `json:"name"`
	Description string        `json:"help"`
	Example     string        `json:"example"`
	Flags       []sflags.Flag `json:"flags"`
	Subcommands []SubCommand  `json:"subcommands,omitempty"`
}

type SubCommand struct {
	Link             string `json:"link"`
	Name             string `json:"name"`
	ShortDescription string `json:"short_description"`
}

func (cmd Command) String() string {
	var str, err = cmd.Format(TextTamplate)
	if err != nil {
		panic(fmt.Errorf("Command.String: %v", err))
	}
	return str
}

func (cmd Command) Doc() Doc {
	var cmdFlags []sflags.Flag
	cmd.LocalFlags().VisitAll(func(flag *pflag.Flag) {
		if !flag.Hidden {
			cmdFlags = append(cmdFlags, sflags.Flag{
				Name:     flag.Name,
				Short:    flag.Shorthand,
				Usage:    strings.Replace(flag.Usage, "\n", " ", -1),
				DefValue: flag.DefValue,
			})
		}
	})
	var subc = make([]SubCommand, 0, len(cmd.Commands()))
	for _, sc := range cmd.Commands() {
		subc = append(subc, SubCommand{
			Link: str.Vector{sc.CommandPath()}.
				Map(str.TrimPrefix("chkit ")).
				Map(strings.NewReplacer(" ", "_").Replace).
				FirstNonEmpty(),
			Name:             strings.TrimPrefix(sc.CommandPath(), "chkit "),
			ShortDescription: sc.Short,
		})
	}
	var cmdPath = strings.TrimPrefix(cmd.CommandPath(), "chkit ")
	return Doc{
		Link: strings.Replace(cmdPath, " ", "_", -1),
		Path: cmdPath,
		Name: cmd.Name(),
		Description: strings.Replace(str.Vector{
			cmd.Long,
			cmd.Short,
			cmd.Example,
		}.FirstNonEmpty(), "\n", " ", -1),
		Example:     cmd.Example,
		Flags:       cmdFlags,
		Subcommands: subc,
	}
}

func (cmd Command) Format(format string) (string, error) {
	var templ, err = template.New(cmd.Name()).Parse(format)
	if err != nil {
		return "", err
	}
	var buf = &bytes.Buffer{}
	err = templ.Execute(buf, cmd.Doc())
	return buf.String(), err
}

func (cmd Command) Markdown() string {
	var str, err = cmd.Format(MarkdownTemplate)
	if err != nil {
		panic(fmt.Errorf("Command.Markdown: %v", err))
	}
	return str
}
<commit_msg>remove newline replacing<commit_after>package doc

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"

	"github.com/ninedraft/boxofstuff/str"
	"github.com/octago/sflags"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

const (
	ListFormat       = "{{.Path}}"
	MarkdownTemplate = `#### <a name="{{.Link}}">{{.Path}}</a>` + "\n\n" +
		"**Description**:\n\n{{.Description}}\n\n" +
		"**Example**:\n\n{{.Example}}\n\n" +
		"**Flags**:\n\n" +
		"| Short | Name | Usage | Default value |\n" +
		"| ----- | ---- | ----- | ------------- |\n" +
		"{{range .Flags}}" +
		"| {{if .Short}}-{{.Short}}{{end}} " +
		"| --{{.Name}} " +
		"| {{.Usage}} " +
		"| {{if ne .DefValue \"[]\"}}{{.DefValue}}{{end}} " +
		"|\n" +
		"{{end}}\n\n" +
		"**Subcommands**:\n\n" +
		"{{range .Subcommands}}" +
		"* **[{{.Name}}](#{{.Link}})** {{.ShortDescription}}\n" +
		"{{end}}\n\n"
	TextTamplate = "Command: {{.Path}}\n" +
		"Description:\n{{.Description}}\n" +
		"Example:\n{{.Example}}\n" +
		"Flags:\n" +
		"{{range .Flags}}" +
		"{{if .Short}}-{{.Short}}{{else}}  {{end}} " +
		"--{{.Name}} " +
		"{{.Usage}} " +
		"{{if ne .DefValue \"[]\"}}{{.DefValue}}{{end}} " +
		"\n" +
		"{{end}}\n"
)

type Command struct {
	cobra.Command
}

type Doc struct {
	Link        string        `json:"link"`
	Path        string        `json:"path"`
	Name        string        `json:"name"`
	Description string        `json:"help"`
	Example     string        `json:"example"`
	Flags       []sflags.Flag `json:"flags"`
	Subcommands []SubCommand  `json:"subcommands,omitempty"`
}

type SubCommand struct {
	Link             string `json:"link"`
	Name             string `json:"name"`
	ShortDescription string `json:"short_description"`
}

func (cmd Command) String() string {
	var str, err = cmd.Format(TextTamplate)
	if err != nil {
		panic(fmt.Errorf("Command.String: %v", err))
	}
	return str
}

func (cmd Command) Doc() Doc {
	var cmdFlags []sflags.Flag
	cmd.LocalFlags().VisitAll(func(flag *pflag.Flag) {
		if !flag.Hidden {
			cmdFlags = append(cmdFlags, sflags.Flag{
				Name:     flag.Name,
				Short:    flag.Shorthand,
				Usage:    strings.Replace(flag.Usage, "\n", " ", -1),
				DefValue: flag.DefValue,
			})
		}
	})
	var subc = make([]SubCommand, 0, len(cmd.Commands()))
	for _, sc := range cmd.Commands() {
		subc = append(subc, SubCommand{
			Link: str.Vector{sc.CommandPath()}.
				Map(str.TrimPrefix("chkit ")).
				Map(strings.NewReplacer(" ", "_").Replace).
				FirstNonEmpty(),
			Name:             strings.TrimPrefix(sc.CommandPath(), "chkit "),
			ShortDescription: sc.Short,
		})
	}
	var cmdPath = strings.TrimPrefix(cmd.CommandPath(), "chkit ")
	return Doc{
		Link: strings.Replace(cmdPath, " ", "_", -1),
		Path: cmdPath,
		Name: cmd.Name(),
		Description: str.Vector{
			cmd.Long,
			cmd.Short,
			cmd.Example,
		}.FirstNonEmpty(),
		Example:     cmd.Example,
		Flags:       cmdFlags,
		Subcommands: subc,
	}
}

func (cmd Command) Format(format string) (string, error) {
	var templ, err = template.New(cmd.Name()).Parse(format)
	if err != nil {
		return "", err
	}
	var buf = &bytes.Buffer{}
	err = templ.Execute(buf, cmd.Doc())
	return buf.String(), err
}

func (cmd Command) Markdown() string {
	var str, err = cmd.Format(MarkdownTemplate)
	if err != nil {
		panic(fmt.Errorf("Command.Markdown: %v", err))
	}
	return str
}
<|endoftext|>


<commit_before>package server

import (
	"os"
	"net"

	"junta/util"
	"junta/paxos"
	"junta/proto"
	"junta/store"
	"strconv"
	"time"
)

type ReadFromWriteToer interface {
	ReadFrom([]byte) (int, net.Addr, os.Error)
	WriteTo([]byte, net.Addr) (int, os.Error)
	LocalAddr() net.Addr
}

const packetSize = 3000

var ErrBadPrefix = os.NewError("bad prefix in path")

type conn struct {
	net.Conn
	s *Server
}

type Manager interface {
	PutFrom(string, paxos.Msg)
	Propose(string) (uint64, string, os.Error)
	Alpha() int
}

type Server struct {
	Addr string
	St *store.Store
	Mg Manager
	Self, Prefix string
}

func (sv *Server) ListenAndServe() os.Error {
	logger := util.NewLogger("server %s", sv.Addr)

	logger.Log("binding")
	l, err := net.Listen("tcp", sv.Addr)
	if err != nil {
		logger.Log(err)
		return err
	}
	defer l.Close()
	logger.Log("listening")

	err = sv.Serve(l)
	if err != nil {
		logger.Logf("%s: %s", l, err)
}
	return err
}

func (sv *Server) ListenAndServeUdp(outs chan paxos.Packet) os.Error {
	logger := util.NewLogger("udp server %s", sv.Addr)

	logger.Log("binding")
	u, err := net.ListenPacket("udp", sv.Addr)
	if err != nil {
		logger.Log(err)
		return err
	}
	defer u.Close()
	logger.Log("listening")

	err = sv.ServeUdp(u, outs)
	if err != nil {
		logger.Logf("%s: %s", u, err)
	}
	return err
}

func (sv *Server) ServeUdp(u ReadFromWriteToer, outs chan paxos.Packet) os.Error {
	recvd := make(chan paxos.Packet)
	sent := make(chan paxos.Packet)

	logger := util.NewLogger("udp server %s", u.LocalAddr())
	go func() {
		logger.Log("reading messages...")
		for {
			msg, addr, err := paxos.ReadMsg(u, packetSize)
			if err != nil {
				logger.Log(err)
				continue
			}
			logger.Logf("read %v from %s", msg, addr)
			recvd <- paxos.Packet{msg, addr}
			sv.Mg.PutFrom(addr, msg)
		}
	}()

	go func() {
		logger.Log("sending messages...")
		for pk := range outs {
			logger.Logf("sending %v", pk)
			udpAddr, err := net.ResolveUDPAddr(pk.Addr)
			if err != nil {
				logger.Log(err)
				continue
			}

			_, err = u.WriteTo(pk.Msg.WireBytes(), udpAddr)
			if err != nil {
				logger.Log(err)
				continue
			}
			sent <- paxos.Packet{pk.Msg, pk.Addr}
		}
	}()

	needsAck := make(map[string]bool)
	resend := make(chan paxos.Packet)
	for {
		select {
		case pk := <-recvd:
			if pk.Msg.HasFlags(paxos.Ack) {
				logger.Logf("got ack %s %v", pk.Addr, pk.Msg)
				needsAck[pk.Id()] = false
			} else {
				logger.Logf("sending ack %s %v", pk.Addr, pk.Msg)
				udpAddr, err := net.ResolveUDPAddr(pk.Addr)
				if err != nil {
					break
				}
				ack := pk.Msg.Dup().SetFlags(paxos.Ack)
				u.WriteTo(ack.WireBytes(), udpAddr)
			}
		case pk := <-sent:
			needsAck[pk.Id()] = true
			logger.Logf("needs ack %s %v", pk.Addr, pk.Msg)
			go func() {
				time.Sleep(100000000) // ns == 0.1s
				resend <- pk
			}()
		case pk := <-resend:
			if needsAck[pk.Id()] {
				logger.Logf("resending %s %v", pk.Addr, pk.Msg)
				go func() {
					outs <- pk
				}()
			} else {
				needsAck[pk.Id()] = false, false
			}
		}
	}

	panic("not reached")
}

func (s *Server) Serve(l net.Listener) os.Error {
	for {
		rw, e := l.Accept()
		if e != nil {
			return e
		}
		c := &conn{rw, s}
		go c.serve()
	}

	panic("not reached")
}

func (sv *Server) leader() string {
	parts, cas := sv.St.Lookup("/junta/leader")
	if cas == store.Dir && cas == store.Missing {
		return ""
	}
	return parts[0]
}

func (sv *Server) addrFor(id string) string {
	parts, cas := sv.St.Lookup("/junta/members/"+id)
	if cas == store.Dir && cas == store.Missing {
		return ""
	}
	return parts[0]
}

// Checks that path begins with the proper prefix and returns the short path
// without the prefix.
func (sv *Server) checkPath(path string) (string, os.Error) {
	logger := util.NewLogger("checkPath")
	if path[0:len(sv.Prefix)] != sv.Prefix {
		logger.Logf("prefix %q not in %q", sv.Prefix, path)
		return "", ErrBadPrefix
	}
	return path[len(sv.Prefix):], nil
}

func (sv *Server) setOnce(path, body, cas string) (uint64, os.Error) {
	shortPath, err := sv.checkPath(path)
	if err != nil {
		return 0, err
	}

	mut, err := store.EncodeSet(shortPath, body, cas)
	if err != nil {
		return 0, err
	}

	seqn, v, err := sv.Mg.Propose(mut)
	if err != nil {
		return 0, err
	}

	// We failed, but only because of a competing proposal. The client should
	// retry.
	if v != mut {
		return 0, os.EAGAIN
	}

	return seqn, nil
}

func (sv *Server) Set(path, body, cas string) (seqn uint64, err os.Error) {
	err = os.EAGAIN
	for err == os.EAGAIN {
		seqn, err = sv.setOnce(path, body, cas)
	}
	return
}

func (sv *Server) delOnce(path, cas string) (uint64, os.Error) {
	shortPath, err := sv.checkPath(path)
	if err != nil {
		return 0, err
	}

	mut, err := store.EncodeDel(shortPath, cas)
	if err != nil {
		return 0, err
	}

	seqn, v, err := sv.Mg.Propose(mut)
	if err != nil {
		return 0, err
	}

	// We failed, but only because of a competing proposal. The client should
	// retry.
	if v != mut {
		return 0, os.EAGAIN
	}

	return seqn, nil
}

func (sv *Server) Del(path, cas string) (seqn uint64, err os.Error) {
	err = os.EAGAIN
	for err == os.EAGAIN {
		seqn, err = sv.delOnce(path, cas)
	}
	return
}

func (sv *Server) WaitForPathSet(path string) (body string, err os.Error) {
	evs := make(chan store.Event)
	defer close(evs)
	sv.St.Watch(path, evs)

	parts, cas := sv.St.Lookup(path)
	if cas != store.Dir && cas != store.Missing {
		return parts[0], nil
	}

	for ev := range evs {
		if ev.IsSet() {
			return ev.Body, nil
		}
	}

	panic("not reached")
}

// Repeatedly propose nop values until a successful read from `done`.
func (sv *Server) AdvanceUntil(done chan int) {
	for _, ok := <-done; !ok; _, ok = <-done {
		sv.Mg.Propose(store.Nop)
	}
}

func (c *conn) serve() {
	pc := proto.NewConn(c)
	logger := util.NewLogger("%v", c.RemoteAddr())
	logger.Log("accepted connection")
	for {
		rid, parts, err := pc.ReadRequest()
		if err != nil {
			if err == os.EOF {
				logger.Log("connection closed by peer")
			} else {
				logger.Log(err)
			}
			return
		}

		rlogger := util.NewLogger("%v - req [%d]", c.RemoteAddr(), rid)
		rlogger.Logf("received <%v>", parts)

		if len(parts) == 0 {
			rlogger.Log("zero parts supplied")
			pc.SendError(rid, proto.InvalidCommand + ": no command")
			continue
		}

		switch parts[0] {
		default:
			rlogger.Logf("unknown command <%s>", parts[0])
			pc.SendError(rid, proto.InvalidCommand + " " + parts[0])
		case "set":
			if len(parts) != 4 {
				rlogger.Logf("invalid set command: %#v", parts)
				pc.SendError(rid, "wrong number of parts")
				break
			}

			leader := c.s.leader()
			if c.s.Self != leader {
				addr := c.s.addrFor(leader)
				if addr == "" {
					rlogger.Logf("unknown address for leader: %s", leader)
					pc.SendError(rid, "unknown address for leader")
					break
				}

				rlogger.Logf("redirect to %s", addr)
				pc.SendRedirect(rid, addr)
				break
			}

			rlogger.Logf("set %q=%q (cas %q)", parts[1], parts[2], parts[3])
			seqn, err := c.s.Set(parts[1], parts[2], parts[3])
			if err != nil {
				rlogger.Logf("bad: %s", err)
				pc.SendError(rid, err.String())
			} else {
				rlogger.Logf("good")
				pc.SendResponse(rid, strconv.Uitoa64(seqn))
			}
		case "del":
			if len(parts) != 3 {
				rlogger.Logf("invalid del command: %v", parts)
				pc.SendError(rid, "wrong number of parts")
				break
			}

			leader := c.s.leader()
			if c.s.Self != leader {
				rlogger.Logf("redirect to %s", leader)
				pc.SendRedirect(rid, leader)
				break
			}

			rlogger.Logf("del %q (cas %q)", parts[1], parts[2])
			_, err := c.s.Del(parts[1], parts[2])
			if err != nil {
				rlogger.Logf("bad: %s", err)
				pc.SendError(rid, err.String())
			} else {
				rlogger.Logf("good")
				pc.SendResponse(rid, "true")
			}
		case "wait-for-path-set": // TODO this is for demo purposes only
			if len(parts) != 2 {
				rlogger.Logf("invalid wait-for-path-set command: %v", parts)
				pc.SendError(rid, "wrong number of parts")
				break
			}
			rlogger.Logf("wait-for-path-set %q", parts[1])
			body, err := c.s.WaitForPathSet(parts[1])
			if err != nil {
				rlogger.Logf("bad: %s", err)
				pc.SendError(rid, err.String())
			} else {
				rlogger.Logf("good %q", body)
				pc.SendResponse(rid, body)
			}
		case "join":
			// join abc123 1.2.3.4:999
			if len(parts) != 3 {
				rlogger.Logf("invalid join command: %v", parts)
				pc.SendError(rid, "wrong number of parts")
				break
			}

			leader := c.s.leader()
			if c.s.Self != leader {
				rlogger.Logf("redirect to %s", leader)
				pc.SendRedirect(rid, leader)
				break
			}

			who, addr := parts[1], parts[2]
			rlogger.Logf("membership requested for %s at %s", who, addr)

			key := c.s.Prefix + "/junta/members/" + who

			seqn, err := c.s.Set(key, addr, store.Missing)
			if err != nil {
				rlogger.Logf("bad: %s", err)
				pc.SendError(rid, err.String())
			} else {
				rlogger.Logf("good")
				done := make(chan int)
				go c.s.AdvanceUntil(done)
				c.s.St.Sync(seqn + uint64(c.s.Mg.Alpha()))
				close(done)
				seqn, snap := c.s.St.Snapshot()
				pc.SendResponse(rid, strconv.Uitoa64(seqn), snap)
			}
		}
	}
}
<commit_msg>properly check path prefix<commit_after>package server

import (
	"os"
	"net"

	"junta/util"
	"junta/paxos"
	"junta/proto"
	"junta/store"
	"strconv"
	"strings"
	"time"
)

type ReadFromWriteToer interface {
	ReadFrom([]byte) (int, net.Addr, os.Error)
	WriteTo([]byte, net.Addr) (int, os.Error)
	LocalAddr() net.Addr
}

const packetSize = 3000

var ErrBadPrefix = os.NewError("bad prefix in path")

type conn struct {
	net.Conn
	s *Server
}

type Manager interface {
	PutFrom(string, paxos.Msg)
	Propose(string) (uint64, string, os.Error)
	Alpha() int
}

type Server struct {
	Addr string
	St *store.Store
	Mg Manager
	Self, Prefix string
}

func (sv *Server) ListenAndServe() os.Error {
	logger := util.NewLogger("server %s", sv.Addr)

	logger.Log("binding")
	l, err := net.Listen("tcp", sv.Addr)
	if err != nil {
		logger.Log(err)
		return err
	}
	defer l.Close()
	logger.Log("listening")

	err = sv.Serve(l)
	if err != nil {
		logger.Logf("%s: %s", l, err)
}
	return err
}

func (sv *Server) ListenAndServeUdp(outs chan paxos.Packet) os.Error {
	logger := util.NewLogger("udp server %s", sv.Addr)

	logger.Log("binding")
	u, err := net.ListenPacket("udp", sv.Addr)
	if err != nil {
		logger.Log(err)
		return err
	}
	defer u.Close()
	logger.Log("listening")

	err = sv.ServeUdp(u, outs)
	if err != nil {
		logger.Logf("%s: %s", u, err)
	}
	return err
}

func (sv *Server) ServeUdp(u ReadFromWriteToer, outs chan paxos.Packet) os.Error {
	recvd := make(chan paxos.Packet)
	sent := make(chan paxos.Packet)

	logger := util.NewLogger("udp server %s", u.LocalAddr())
	go func() {
		logger.Log("reading messages...")
		for {
			msg, addr, err := paxos.ReadMsg(u, packetSize)
			if err != nil {
				logger.Log(err)
				continue
			}
			logger.Logf("read %v from %s", msg, addr)
			recvd <- paxos.Packet{msg, addr}
			sv.Mg.PutFrom(addr, msg)
		}
	}()

	go func() {
		logger.Log("sending messages...")
		for pk := range outs {
			logger.Logf("sending %v", pk)
			udpAddr, err := net.ResolveUDPAddr(pk.Addr)
			if err != nil {
				logger.Log(err)
				continue
			}

			_, err = u.WriteTo(pk.Msg.WireBytes(), udpAddr)
			if err != nil {
				logger.Log(err)
				continue
			}
			sent <- paxos.Packet{pk.Msg, pk.Addr}
		}
	}()

	needsAck := make(map[string]bool)
	resend := make(chan paxos.Packet)
	for {
		select {
		case pk := <-recvd:
			if pk.Msg.HasFlags(paxos.Ack) {
				logger.Logf("got ack %s %v", pk.Addr, pk.Msg)
				needsAck[pk.Id()] = false
			} else {
				logger.Logf("sending ack %s %v", pk.Addr, pk.Msg)
				udpAddr, err := net.ResolveUDPAddr(pk.Addr)
				if err != nil {
					break
				}
				ack := pk.Msg.Dup().SetFlags(paxos.Ack)
				u.WriteTo(ack.WireBytes(), udpAddr)
			}
		case pk := <-sent:
			needsAck[pk.Id()] = true
			logger.Logf("needs ack %s %v", pk.Addr, pk.Msg)
			go func() {
				time.Sleep(100000000) // ns == 0.1s
				resend <- pk
			}()
		case pk := <-resend:
			if needsAck[pk.Id()] {
				logger.Logf("resending %s %v", pk.Addr, pk.Msg)
				go func() {
					outs <- pk
				}()
			} else {
				needsAck[pk.Id()] = false, false
			}
		}
	}

	panic("not reached")
}

func (s *Server) Serve(l net.Listener) os.Error {
	for {
		rw, e := l.Accept()
		if e != nil {
			return e
		}
		c := &conn{rw, s}
		go c.serve()
	}

	panic("not reached")
}

func (sv *Server) leader() string {
	parts, cas := sv.St.Lookup("/junta/leader")
	if cas == store.Dir && cas == store.Missing {
		return ""
	}
	return parts[0]
}

func (sv *Server) addrFor(id string) string {
	parts, cas := sv.St.Lookup("/junta/members/"+id)
	if cas == store.Dir && cas == store.Missing {
		return ""
	}
	return parts[0]
}

// Checks that path begins with the proper prefix and returns the short path
// without the prefix.
func (sv *Server) checkPath(path string) (string, os.Error) {
	logger := util.NewLogger("checkPath")
	if !strings.HasPrefix(path, sv.Prefix+"/") {
		logger.Logf("prefix %q not in %q", sv.Prefix+"/", path)
		return "", ErrBadPrefix
	}
	return path[len(sv.Prefix):], nil
}

func (sv *Server) setOnce(path, body, cas string) (uint64, os.Error) {
	shortPath, err := sv.checkPath(path)
	if err != nil {
		return 0, err
	}

	mut, err := store.EncodeSet(shortPath, body, cas)
	if err != nil {
		return 0, err
	}

	seqn, v, err := sv.Mg.Propose(mut)
	if err != nil {
		return 0, err
	}

	// We failed, but only because of a competing proposal. The client should
	// retry.
	if v != mut {
		return 0, os.EAGAIN
	}

	return seqn, nil
}

func (sv *Server) Set(path, body, cas string) (seqn uint64, err os.Error) {
	err = os.EAGAIN
	for err == os.EAGAIN {
		seqn, err = sv.setOnce(path, body, cas)
	}
	return
}

func (sv *Server) delOnce(path, cas string) (uint64, os.Error) {
	shortPath, err := sv.checkPath(path)
	if err != nil {
		return 0, err
	}

	mut, err := store.EncodeDel(shortPath, cas)
	if err != nil {
		return 0, err
	}

	seqn, v, err := sv.Mg.Propose(mut)
	if err != nil {
		return 0, err
	}

	// We failed, but only because of a competing proposal. The client should
	// retry.
	if v != mut {
		return 0, os.EAGAIN
	}

	return seqn, nil
}

func (sv *Server) Del(path, cas string) (seqn uint64, err os.Error) {
	err = os.EAGAIN
	for err == os.EAGAIN {
		seqn, err = sv.delOnce(path, cas)
	}
	return
}

func (sv *Server) WaitForPathSet(path string) (body string, err os.Error) {
	evs := make(chan store.Event)
	defer close(evs)
	sv.St.Watch(path, evs)

	parts, cas := sv.St.Lookup(path)
	if cas != store.Dir && cas != store.Missing {
		return parts[0], nil
	}

	for ev := range evs {
		if ev.IsSet() {
			return ev.Body, nil
		}
	}

	panic("not reached")
}

// Repeatedly propose nop values until a successful read from `done`.
func (sv *Server) AdvanceUntil(done chan int) {
	for _, ok := <-done; !ok; _, ok = <-done {
		sv.Mg.Propose(store.Nop)
	}
}

func (c *conn) serve() {
	pc := proto.NewConn(c)
	logger := util.NewLogger("%v", c.RemoteAddr())
	logger.Log("accepted connection")
	for {
		rid, parts, err := pc.ReadRequest()
		if err != nil {
			if err == os.EOF {
				logger.Log("connection closed by peer")
			} else {
				logger.Log(err)
			}
			return
		}

		rlogger := util.NewLogger("%v - req [%d]", c.RemoteAddr(), rid)
		rlogger.Logf("received <%v>", parts)

		if len(parts) == 0 {
			rlogger.Log("zero parts supplied")
			pc.SendError(rid, proto.InvalidCommand + ": no command")
			continue
		}

		switch parts[0] {
		default:
			rlogger.Logf("unknown command <%s>", parts[0])
			pc.SendError(rid, proto.InvalidCommand + " " + parts[0])
		case "set":
			if len(parts) != 4 {
				rlogger.Logf("invalid set command: %#v", parts)
				pc.SendError(rid, "wrong number of parts")
				break
			}

			leader := c.s.leader()
			if c.s.Self != leader {
				addr := c.s.addrFor(leader)
				if addr == "" {
					rlogger.Logf("unknown address for leader: %s", leader)
					pc.SendError(rid, "unknown address for leader")
					break
				}

				rlogger.Logf("redirect to %s", addr)
				pc.SendRedirect(rid, addr)
				break
			}

			rlogger.Logf("set %q=%q (cas %q)", parts[1], parts[2], parts[3])
			seqn, err := c.s.Set(parts[1], parts[2], parts[3])
			if err != nil {
				rlogger.Logf("bad: %s", err)
				pc.SendError(rid, err.String())
			} else {
				rlogger.Logf("good")
				pc.SendResponse(rid, strconv.Uitoa64(seqn))
			}
		case "del":
			if len(parts) != 3 {
				rlogger.Logf("invalid del command: %v", parts)
				pc.SendError(rid, "wrong number of parts")
				break
			}

			leader := c.s.leader()
			if c.s.Self != leader {
				rlogger.Logf("redirect to %s", leader)
				pc.SendRedirect(rid, leader)
				break
			}

			rlogger.Logf("del %q (cas %q)", parts[1], parts[2])
			_, err := c.s.Del(parts[1], parts[2])
			if err != nil {
				rlogger.Logf("bad: %s", err)
				pc.SendError(rid, err.String())
			} else {
				rlogger.Logf("good")
				pc.SendResponse(rid, "true")
			}
		case "wait-for-path-set": // TODO this is for demo purposes only
			if len(parts) != 2 {
				rlogger.Logf("invalid wait-for-path-set command: %v", parts)
				pc.SendError(rid, "wrong number of parts")
				break
			}
			rlogger.Logf("wait-for-path-set %q", parts[1])
			body, err := c.s.WaitForPathSet(parts[1])
			if err != nil {
				rlogger.Logf("bad: %s", err)
				pc.SendError(rid, err.String())
			} else {
				rlogger.Logf("good %q", body)
				pc.SendResponse(rid, body)
			}
		case "join":
			// join abc123 1.2.3.4:999
			if len(parts) != 3 {
				rlogger.Logf("invalid join command: %v", parts)
				pc.SendError(rid, "wrong number of parts")
				break
			}

			leader := c.s.leader()
			if c.s.Self != leader {
				rlogger.Logf("redirect to %s", leader)
				pc.SendRedirect(rid, leader)
				break
			}

			who, addr := parts[1], parts[2]
			rlogger.Logf("membership requested for %s at %s", who, addr)

			key := c.s.Prefix + "/junta/members/" + who

			seqn, err := c.s.Set(key, addr, store.Missing)
			if err != nil {
				rlogger.Logf("bad: %s", err)
				pc.SendError(rid, err.String())
			} else {
				rlogger.Logf("good")
				done := make(chan int)
				go c.s.AdvanceUntil(done)
				c.s.St.Sync(seqn + uint64(c.s.Mg.Alpha()))
				close(done)
				seqn, snap := c.s.St.Snapshot()
				pc.SendResponse(rid, strconv.Uitoa64(seqn), snap)
			}
		}
	}
}
<|endoftext|>


<commit_before>package slinga

import (
	"fmt"
	"github.com/julienschmidt/httprouter"
	"net/http"
)

func Index(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	fmt.Fprint(w, "Welcome!\n")
}

func Hello(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	fmt.Fprintf(w, "hello, %s!\n", ps.ByName("name"))
}

func Serve(host string, port int) {
	router := httprouter.New()
	router.GET("/", Index)
	router.GET("/hello/:name", Hello)
	router.ServeFiles("/static/*filepath", http.Dir("static/"))

	fmt.Println("Serving")
	// todo handle error returned from ListenAndServe (path to Fatal??)
	http.ListenAndServe(fmt.Sprintf("%s:%d", host, port), router)
}
<commit_msg>Fix server to serve files from public<commit_after>package slinga

import (
	"fmt"
	"github.com/julienschmidt/httprouter"
	"net/http"
)

func Index(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	fmt.Fprint(w, "Welcome!\n")
}

func Hello(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	fmt.Fprintf(w, "hello, %s!\n", ps.ByName("name"))
}

func Serve(host string, port int) {
	router := httprouter.New()
	router.GET("/", Index)
	router.GET("/hello/:name", Hello)
	router.ServeFiles("/static/*filepath", http.Dir("public/static"))

	fmt.Println("Serving")
	// todo handle error returned from ListenAndServe (path to Fatal??)
	http.ListenAndServe(fmt.Sprintf("%s:%d", host, port), router)
}
<|endoftext|>


<commit_before>package store

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	level "github.com/go-kit/kit/log/experimental_level"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/oklog/prototype/pkg/cluster"
	"github.com/oklog/prototype/pkg/ingest"
)

// Consumer reads segments from the ingesters, and replicates merged segments to
// the rest of the cluster. It's implemented as a state machine: gather
// segments, replicate, commit, and repeat. All failures invalidate the entire
// batch.
type Consumer struct {
	peer               *cluster.Peer
	client             *http.Client
	segmentTargetSize  int64
	replicationFactor  int
	gatherErrors       int                 // heuristic to move out of gather state
	pending            map[string][]string // ingester: segment IDs
	active             *bytes.Buffer       // merged pending segments
	activeSince        time.Time           // active segment has been "open" since this time
	stop               chan chan struct{}
	consumedSegments   prometheus.Counter
	consumedBytes      prometheus.Counter
	replicatedSegments prometheus.Counter
	replicatedBytes    prometheus.Counter
	logger             log.Logger
}

// NewConsumer creates a consumer.
// Don't forget to Run it.
func NewConsumer(
	peer *cluster.Peer,
	client *http.Client,
	segmentTargetSize int64,
	replicationFactor int,
	consumedSegments, consumedBytes prometheus.Counter,
	replicatedSegments, replicatedBytes prometheus.Counter,
	logger log.Logger,
) *Consumer {
	return &Consumer{
		peer:               peer,
		client:             client,
		segmentTargetSize:  segmentTargetSize,
		replicationFactor:  replicationFactor,
		gatherErrors:       0,
		pending:            map[string][]string{},
		active:             &bytes.Buffer{},
		activeSince:        time.Time{},
		stop:               make(chan chan struct{}),
		consumedSegments:   consumedSegments,
		consumedBytes:      consumedBytes,
		replicatedSegments: replicatedSegments,
		replicatedBytes:    replicatedBytes,
		logger:             logger,
	}
}

// Run consumes segments from ingest nodes, and replicates them to the cluster.
// Run returns when Stop is invoked.
func (c *Consumer) Run() {
	step := time.Tick(100 * time.Millisecond)
	state := c.gather
	for {
		select {
		case <-step:
			state = state()

		case q := <-c.stop:
			c.fail() // any outstanding segments
			close(q)
			return
		}
	}
}

// Stop the consumer from consuming.
func (c *Consumer) Stop() {
	q := make(chan struct{})
	c.stop <- q
	<-q
}

type stateFn func() stateFn

func (c *Consumer) gather() stateFn {
	var (
		base = log.NewContext(c.logger).With("state", "gather")
		warn = level.Warn(base)
	)

	// A naïve way to break out of the gather loop in atypical conditions.
	// TODO(pb): this obviously needs more thought and consideration
	instances := c.peer.Current(cluster.PeerTypeIngest)
	if c.gatherErrors > 0 && c.gatherErrors > 2*len(instances) {
		if c.active.Len() <= 0 {
			// We didn't successfully consume any segments.
			// Nothing to do but reset and try again.
			c.gatherErrors = 0
			return c.gather
		}
		// We consumed some segment, at least.
		// Press forward to persistence.
		return c.replicate
	}
	if len(instances) == 0 {
		return c.gather // maybe some will come back later
	}
	if want, have := c.replicationFactor, len(c.peer.Current(cluster.PeerTypeStore)); want < have {
		// Don't gather if we can't replicate.
		// Better to queue up on the ingesters.
		warn.Log("replication_factor", want, "available_peers", have, "err", "replication currently impossible")
		time.Sleep(time.Second)
		c.gatherErrors++
		return c.gather
	}

	// More typical exit clauses.
	const maxAge = time.Second // TODO(pb): parameterize?
	var (
		tooBig = int64(c.active.Len()) > c.segmentTargetSize
		tooOld = !c.activeSince.IsZero() && time.Now().Sub(c.activeSince) > maxAge
	)
	if tooBig || tooOld {
		return c.replicate
	}

	// Get the oldest segment ID from a random ingester.
	instance := instances[rand.Intn(len(instances))]
	nextResp, err := c.client.Get(fmt.Sprintf("http://%s/ingest%s", instance, ingest.APIPathNext))
	if err != nil {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "err", err)
		c.gatherErrors++
		return c.gather
	}
	defer nextResp.Body.Close()
	nextRespBody, err := ioutil.ReadAll(nextResp.Body)
	if err != nil {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "err", err)
		c.gatherErrors++
		return c.gather
	}
	nextID := strings.TrimSpace(string(nextRespBody))
	if nextResp.StatusCode == http.StatusNotFound {
		// Normal, when the ingester has no more segments to give right now.
		c.gatherErrors++ // after enough of these errors, we should replicate
		return c.gather
	}
	if nextResp.StatusCode != http.StatusOK {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "returned", nextResp.Status)
		c.gatherErrors++
		return c.gather
	}

	// Mark the segment ID as pending.
	// From this point forward, we must either commit or fail the segment.
	// If we do neither, it will eventually time out, but we should be nice.
	c.pending[instance] = append(c.pending[instance], nextID)

	// Read the segment.
	readResp, err := c.client.Get(fmt.Sprintf("http://%s/ingest%s?id=%s", instance, ingest.APIPathRead, nextID))
	if err != nil {
		// Reading failed, so we can't possibly commit the segment.
		// The simplest thing to do now is to fail everything.
		// TODO(pb): this could be improved i.e. made more granular
		warn.Log("ingester", instance, "during", ingest.APIPathRead, "err", err)
		c.gatherErrors++
		return c.fail // fail everything
	}
	defer readResp.Body.Close()
	if readResp.StatusCode != http.StatusOK {
		warn.Log("ingester", instance, "during", ingest.APIPathRead, "returned", readResp.Status)
		c.gatherErrors++
		return c.fail // fail everything, same as above
	}

	// Merge the segment into our active segment.
	var cw countingWriter
	if _, _, _, err := mergeRecords(c.active, io.TeeReader(readResp.Body, &cw)); err != nil {
		warn.Log("ingester", instance, "during", "mergeRecords", "err", err)
		c.gatherErrors++
		return c.fail // fail everything, same as above
	}
	if c.activeSince.IsZero() {
		c.activeSince = time.Now()
	}

	// Repeat!
	c.consumedSegments.Inc()
	c.consumedBytes.Add(float64(cw.n))
	return c.gather
}

func (c *Consumer) replicate() stateFn {
	var (
		base = log.NewContext(c.logger).With("state", "replicate")
		warn = level.Warn(base)
	)

	// Replicate the segment to the cluster.
	var (
		peers      = c.peer.Current(cluster.PeerTypeStore)
		indices    = rand.Perm(len(peers))
		replicated = 0
	)
	if want, have := c.replicationFactor, len(peers); want < have {
		warn.Log("replication_factor", want, "available_peers", have, "err", "replication currently impossible")
		return c.fail // can't do anything here
	}
	for i := 0; i < len(indices) && replicated < c.replicationFactor; i++ {
		var (
			index    = indices[i]
			target   = peers[index]
			uri      = fmt.Sprintf("http://%s/store%s", target, APIPathReplicate)
			bodyType = "application/binary"
			body     = bytes.NewReader(c.active.Bytes())
		)
		resp, err := c.client.Post(uri, bodyType, body)
		if err != nil {
			warn.Log("target", target, "during", APIPathReplicate, "err", err)
			continue // we'll try another one
		}
		resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			warn.Log("target", target, "during", APIPathReplicate, "got", resp.Status)
			continue // we'll try another one
		}
		replicated++
	}
	if replicated < c.replicationFactor {
		warn.Log("err", "failed to fully replicate", "want", c.replicationFactor, "have", replicated)
		return c.fail // harsh, but OK
	}

	// All good!
	c.replicatedSegments.Inc()
	c.replicatedBytes.Add(float64(c.active.Len()))
	return c.commit
}

func (c *Consumer) commit() stateFn {
	return c.resetVia("commit")
}

func (c *Consumer) fail() stateFn {
	return c.resetVia("failed")
}

func (c *Consumer) resetVia(commitOrFailed string) stateFn {
	var (
		base = log.NewContext(c.logger).With("state", commitOrFailed)
		warn = level.Warn(base)
	)

	// If commits fail, the segment may be re-replicated; that's OK.
	// If fails fail, the segment will eventually time-out; that's also OK.
	// So we have best-effort semantics, just log the error and move on.
	var wg sync.WaitGroup
	for instance, ids := range c.pending {
		wg.Add(len(ids))
		for _, id := range ids {
			go func(id string) {
				defer wg.Done()
				resp, err := c.client.Post(fmt.Sprintf("http://%s/ingest/%s?id=%s", instance, commitOrFailed, id), "text/plain", nil)
				if err != nil {
					warn.Log("instance", instance, "during", "POST", "err", err)
					return
				}
				resp.Body.Close()
				if resp.StatusCode != http.StatusOK {
					warn.Log("instance", instance, "during", "POST", "status", resp.Status)
					return
				}
			}(id)
		}
	}
	wg.Wait()

	// Reset various pending things.
	c.gatherErrors = 0
	c.pending = map[string][]string{}
	c.active.Reset()
	c.activeSince = time.Time{}

	// Back to the beginning.
	return c.gather
}

type countingWriter struct{ n int64 }

func (cw *countingWriter) Write(p []byte) (int, error) {
	cw.n += int64(len(p))
	return len(p), nil
}
<commit_msg>Big oops!<commit_after>package store

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	level "github.com/go-kit/kit/log/experimental_level"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/oklog/prototype/pkg/cluster"
	"github.com/oklog/prototype/pkg/ingest"
)

// Consumer reads segments from the ingesters, and replicates merged segments to
// the rest of the cluster. It's implemented as a state machine: gather
// segments, replicate, commit, and repeat. All failures invalidate the entire
// batch.
type Consumer struct {
	peer               *cluster.Peer
	client             *http.Client
	segmentTargetSize  int64
	replicationFactor  int
	gatherErrors       int                 // heuristic to move out of gather state
	pending            map[string][]string // ingester: segment IDs
	active             *bytes.Buffer       // merged pending segments
	activeSince        time.Time           // active segment has been "open" since this time
	stop               chan chan struct{}
	consumedSegments   prometheus.Counter
	consumedBytes      prometheus.Counter
	replicatedSegments prometheus.Counter
	replicatedBytes    prometheus.Counter
	logger             log.Logger
}

// NewConsumer creates a consumer.
// Don't forget to Run it.
func NewConsumer(
	peer *cluster.Peer,
	client *http.Client,
	segmentTargetSize int64,
	replicationFactor int,
	consumedSegments, consumedBytes prometheus.Counter,
	replicatedSegments, replicatedBytes prometheus.Counter,
	logger log.Logger,
) *Consumer {
	return &Consumer{
		peer:               peer,
		client:             client,
		segmentTargetSize:  segmentTargetSize,
		replicationFactor:  replicationFactor,
		gatherErrors:       0,
		pending:            map[string][]string{},
		active:             &bytes.Buffer{},
		activeSince:        time.Time{},
		stop:               make(chan chan struct{}),
		consumedSegments:   consumedSegments,
		consumedBytes:      consumedBytes,
		replicatedSegments: replicatedSegments,
		replicatedBytes:    replicatedBytes,
		logger:             logger,
	}
}

// Run consumes segments from ingest nodes, and replicates them to the cluster.
// Run returns when Stop is invoked.
func (c *Consumer) Run() {
	step := time.Tick(100 * time.Millisecond)
	state := c.gather
	for {
		select {
		case <-step:
			state = state()

		case q := <-c.stop:
			c.fail() // any outstanding segments
			close(q)
			return
		}
	}
}

// Stop the consumer from consuming.
func (c *Consumer) Stop() {
	q := make(chan struct{})
	c.stop <- q
	<-q
}

type stateFn func() stateFn

func (c *Consumer) gather() stateFn {
	var (
		base = log.NewContext(c.logger).With("state", "gather")
		warn = level.Warn(base)
	)

	// A naïve way to break out of the gather loop in atypical conditions.
	// TODO(pb): this obviously needs more thought and consideration
	instances := c.peer.Current(cluster.PeerTypeIngest)
	if c.gatherErrors > 0 && c.gatherErrors > 2*len(instances) {
		if c.active.Len() <= 0 {
			// We didn't successfully consume any segments.
			// Nothing to do but reset and try again.
			c.gatherErrors = 0
			return c.gather
		}
		// We consumed some segment, at least.
		// Press forward to persistence.
		return c.replicate
	}
	if len(instances) == 0 {
		return c.gather // maybe some will come back later
	}
	if want, have := c.replicationFactor, len(c.peer.Current(cluster.PeerTypeStore)); have < want {
		// Don't gather if we can't replicate.
		// Better to queue up on the ingesters.
		warn.Log("replication_factor", want, "available_peers", have, "err", "replication currently impossible")
		time.Sleep(time.Second)
		c.gatherErrors++
		return c.gather
	}

	// More typical exit clauses.
	const maxAge = time.Second // TODO(pb): parameterize?
	var (
		tooBig = int64(c.active.Len()) > c.segmentTargetSize
		tooOld = !c.activeSince.IsZero() && time.Now().Sub(c.activeSince) > maxAge
	)
	if tooBig || tooOld {
		return c.replicate
	}

	// Get the oldest segment ID from a random ingester.
	instance := instances[rand.Intn(len(instances))]
	nextResp, err := c.client.Get(fmt.Sprintf("http://%s/ingest%s", instance, ingest.APIPathNext))
	if err != nil {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "err", err)
		c.gatherErrors++
		return c.gather
	}
	defer nextResp.Body.Close()
	nextRespBody, err := ioutil.ReadAll(nextResp.Body)
	if err != nil {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "err", err)
		c.gatherErrors++
		return c.gather
	}
	nextID := strings.TrimSpace(string(nextRespBody))
	if nextResp.StatusCode == http.StatusNotFound {
		// Normal, when the ingester has no more segments to give right now.
		c.gatherErrors++ // after enough of these errors, we should replicate
		return c.gather
	}
	if nextResp.StatusCode != http.StatusOK {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "returned", nextResp.Status)
		c.gatherErrors++
		return c.gather
	}

	// Mark the segment ID as pending.
	// From this point forward, we must either commit or fail the segment.
	// If we do neither, it will eventually time out, but we should be nice.
	c.pending[instance] = append(c.pending[instance], nextID)

	// Read the segment.
	readResp, err := c.client.Get(fmt.Sprintf("http://%s/ingest%s?id=%s", instance, ingest.APIPathRead, nextID))
	if err != nil {
		// Reading failed, so we can't possibly commit the segment.
		// The simplest thing to do now is to fail everything.
		// TODO(pb): this could be improved i.e. made more granular
		warn.Log("ingester", instance, "during", ingest.APIPathRead, "err", err)
		c.gatherErrors++
		return c.fail // fail everything
	}
	defer readResp.Body.Close()
	if readResp.StatusCode != http.StatusOK {
		warn.Log("ingester", instance, "during", ingest.APIPathRead, "returned", readResp.Status)
		c.gatherErrors++
		return c.fail // fail everything, same as above
	}

	// Merge the segment into our active segment.
	var cw countingWriter
	if _, _, _, err := mergeRecords(c.active, io.TeeReader(readResp.Body, &cw)); err != nil {
		warn.Log("ingester", instance, "during", "mergeRecords", "err", err)
		c.gatherErrors++
		return c.fail // fail everything, same as above
	}
	if c.activeSince.IsZero() {
		c.activeSince = time.Now()
	}

	// Repeat!
	c.consumedSegments.Inc()
	c.consumedBytes.Add(float64(cw.n))
	return c.gather
}

func (c *Consumer) replicate() stateFn {
	var (
		base = log.NewContext(c.logger).With("state", "replicate")
		warn = level.Warn(base)
	)

	// Replicate the segment to the cluster.
	var (
		peers      = c.peer.Current(cluster.PeerTypeStore)
		indices    = rand.Perm(len(peers))
		replicated = 0
	)
	if want, have := c.replicationFactor, len(peers); have < want {
		warn.Log("replication_factor", want, "available_peers", have, "err", "replication currently impossible")
		return c.fail // can't do anything here
	}
	for i := 0; i < len(indices) && replicated < c.replicationFactor; i++ {
		var (
			index    = indices[i]
			target   = peers[index]
			uri      = fmt.Sprintf("http://%s/store%s", target, APIPathReplicate)
			bodyType = "application/binary"
			body     = bytes.NewReader(c.active.Bytes())
		)
		resp, err := c.client.Post(uri, bodyType, body)
		if err != nil {
			warn.Log("target", target, "during", APIPathReplicate, "err", err)
			continue // we'll try another one
		}
		resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			warn.Log("target", target, "during", APIPathReplicate, "got", resp.Status)
			continue // we'll try another one
		}
		replicated++
	}
	if replicated < c.replicationFactor {
		warn.Log("err", "failed to fully replicate", "want", c.replicationFactor, "have", replicated)
		return c.fail // harsh, but OK
	}

	// All good!
	c.replicatedSegments.Inc()
	c.replicatedBytes.Add(float64(c.active.Len()))
	return c.commit
}

func (c *Consumer) commit() stateFn {
	return c.resetVia("commit")
}

func (c *Consumer) fail() stateFn {
	return c.resetVia("failed")
}

func (c *Consumer) resetVia(commitOrFailed string) stateFn {
	var (
		base = log.NewContext(c.logger).With("state", commitOrFailed)
		warn = level.Warn(base)
	)

	// If commits fail, the segment may be re-replicated; that's OK.
	// If fails fail, the segment will eventually time-out; that's also OK.
	// So we have best-effort semantics, just log the error and move on.
	var wg sync.WaitGroup
	for instance, ids := range c.pending {
		wg.Add(len(ids))
		for _, id := range ids {
			go func(id string) {
				defer wg.Done()
				resp, err := c.client.Post(fmt.Sprintf("http://%s/ingest/%s?id=%s", instance, commitOrFailed, id), "text/plain", nil)
				if err != nil {
					warn.Log("instance", instance, "during", "POST", "err", err)
					return
				}
				resp.Body.Close()
				if resp.StatusCode != http.StatusOK {
					warn.Log("instance", instance, "during", "POST", "status", resp.Status)
					return
				}
			}(id)
		}
	}
	wg.Wait()

	// Reset various pending things.
	c.gatherErrors = 0
	c.pending = map[string][]string{}
	c.active.Reset()
	c.activeSince = time.Time{}

	// Back to the beginning.
	return c.gather
}

type countingWriter struct{ n int64 }

func (cw *countingWriter) Write(p []byte) (int, error) {
	cw.n += int64(len(p))
	return len(p), nil
}
<|endoftext|>


<commit_before>package tsdb

import (
	"fmt"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/models"
)

var (
	defaultRes         int64 = 1500
	defaultMinInterval       = time.Millisecond * 1
	year                     = time.Hour * 24 * 365
	day                      = time.Hour * 24
)

type Interval struct {
	Text  string
	Value time.Duration
}

type intervalCalculator struct {
	minInterval time.Duration
}

type IntervalCalculator interface {
	Calculate(timeRange *TimeRange, minInterval time.Duration) Interval
}

type IntervalOptions struct {
	MinInterval time.Duration
}

func NewIntervalCalculator(opt *IntervalOptions) *intervalCalculator {
	if opt == nil {
		opt = &IntervalOptions{}
	}

	calc := &intervalCalculator{}

	if opt.MinInterval == 0 {
		calc.minInterval = defaultMinInterval
	} else {
		calc.minInterval = opt.MinInterval
	}

	return calc
}

func (i *Interval) Milliseconds() int64 {
	return i.Value.Nanoseconds() / int64(time.Millisecond)
}

func (ic *intervalCalculator) Calculate(timerange *TimeRange, minInterval time.Duration) Interval {
	to := timerange.MustGetTo().UnixNano()
	from := timerange.MustGetFrom().UnixNano()
	interval := time.Duration((to - from) / defaultRes)

	if interval < minInterval {
		return Interval{Text: formatDuration(minInterval), Value: minInterval}
	}

	rounded := roundInterval(interval)
	return Interval{Text: formatDuration(rounded), Value: rounded}
}

func GetIntervalFrom(dsInfo *models.DataSource, queryModel *simplejson.Json, defaultInterval time.Duration) (time.Duration, error) {
	interval := queryModel.Get("interval").MustString("")

	if interval == "" && dsInfo.JsonData != nil {
		dsInterval := dsInfo.JsonData.Get("timeInterval").MustString("")
		if dsInterval != "" {
			interval = dsInterval
		}
	}

	if interval == "" {
		return defaultInterval, nil
	}

	interval = strings.Replace(strings.Replace(interval, "<", "", 1), ">", "", 1)
	parsedInterval, err := time.ParseDuration(interval)
	if err != nil {
		return time.Duration(0), err
	}

	return parsedInterval, nil
}

func formatDuration(inter time.Duration) string {
	if inter >= year {
		return fmt.Sprintf("%dy", inter/year)
	}

	if inter >= day {
		return fmt.Sprintf("%dd", inter/day)
	}

	if inter >= time.Hour {
		return fmt.Sprintf("%dh", inter/time.Hour)
	}

	if inter >= time.Minute {
		return fmt.Sprintf("%dm", inter/time.Minute)
	}

	if inter >= time.Second {
		return fmt.Sprintf("%ds", inter/time.Second)
	}

	if inter >= time.Millisecond {
		return fmt.Sprintf("%dms", inter/time.Millisecond)
	}

	return "1ms"
}

func roundInterval(interval time.Duration) time.Duration {
	switch true {
	// 0.015s
	case interval <= 15*time.Millisecond:
		return time.Millisecond * 10 // 0.01s
	// 0.035s
	case interval <= 35*time.Millisecond:
		return time.Millisecond * 20 // 0.02s
	// 0.075s
	case interval <= 75*time.Millisecond:
		return time.Millisecond * 50 // 0.05s
	// 0.15s
	case interval <= 150*time.Millisecond:
		return time.Millisecond * 100 // 0.1s
	// 0.35s
	case interval <= 350*time.Millisecond:
		return time.Millisecond * 200 // 0.2s
	// 0.75s
	case interval <= 750*time.Millisecond:
		return time.Millisecond * 500 // 0.5s
	// 1.5s
	case interval <= 1500*time.Millisecond:
		return time.Millisecond * 1000 // 1s
	// 3.5s
	case interval <= 3500*time.Millisecond:
		return time.Millisecond * 2000 // 2s
	// 7.5s
	case interval <= 7500*time.Millisecond:
		return time.Millisecond * 5000 // 5s
	// 12.5s
	case interval <= 12500*time.Millisecond:
		return time.Millisecond * 10000 // 10s
	// 17.5s
	case interval <= 17500*time.Millisecond:
		return time.Millisecond * 15000 // 15s
	// 25s
	case interval <= 25000*time.Millisecond:
		return time.Millisecond * 20000 // 20s
	// 45s
	case interval <= 45000*time.Millisecond:
		return time.Millisecond * 30000 // 30s
	// 1.5m
	case interval <= 90000*time.Millisecond:
		return time.Millisecond * 60000 // 1m
	// 3.5m
	case interval <= 210000*time.Millisecond:
		return time.Millisecond * 120000 // 2m
	// 7.5m
	case interval <= 450000*time.Millisecond:
		return time.Millisecond * 300000 // 5m
	// 12.5m
	case interval <= 750000*time.Millisecond:
		return time.Millisecond * 600000 // 10m
	// 12.5m
	case interval <= 1050000*time.Millisecond:
		return time.Millisecond * 900000 // 15m
	// 25m
	case interval <= 1500000*time.Millisecond:
		return time.Millisecond * 1200000 // 20m
	// 45m
	case interval <= 2700000*time.Millisecond:
		return time.Millisecond * 1800000 // 30m
	// 1.5h
	case interval <= 5400000*time.Millisecond:
		return time.Millisecond * 3600000 // 1h
	// 2.5h
	case interval <= 9000000*time.Millisecond:
		return time.Millisecond * 7200000 // 2h
	// 4.5h
	case interval <= 16200000*time.Millisecond:
		return time.Millisecond * 10800000 // 3h
	// 9h
	case interval <= 32400000*time.Millisecond:
		return time.Millisecond * 21600000 // 6h
	// 24h
	case interval <= 86400000*time.Millisecond:
		return time.Millisecond * 43200000 // 12h
	// 48h
	case interval <= 172800000*time.Millisecond:
		return time.Millisecond * 86400000 // 24h
	// 1w
	case interval <= 604800000*time.Millisecond:
		return time.Millisecond * 86400000 // 24h
	// 3w
	case interval <= 1814400000*time.Millisecond:
		return time.Millisecond * 604800000 // 1w
	// 2y
	case interval < 3628800000*time.Millisecond:
		return time.Millisecond * 2592000000 // 30d
	default:
		return time.Millisecond * 31536000000 // 1y
	}
}
<commit_msg>interval: make the FormatDuration function public<commit_after>package tsdb

import (
	"fmt"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/models"
)

var (
	defaultRes         int64 = 1500
	defaultMinInterval       = time.Millisecond * 1
	year                     = time.Hour * 24 * 365
	day                      = time.Hour * 24
)

type Interval struct {
	Text  string
	Value time.Duration
}

type intervalCalculator struct {
	minInterval time.Duration
}

type IntervalCalculator interface {
	Calculate(timeRange *TimeRange, minInterval time.Duration) Interval
}

type IntervalOptions struct {
	MinInterval time.Duration
}

func NewIntervalCalculator(opt *IntervalOptions) *intervalCalculator {
	if opt == nil {
		opt = &IntervalOptions{}
	}

	calc := &intervalCalculator{}

	if opt.MinInterval == 0 {
		calc.minInterval = defaultMinInterval
	} else {
		calc.minInterval = opt.MinInterval
	}

	return calc
}

func (i *Interval) Milliseconds() int64 {
	return i.Value.Nanoseconds() / int64(time.Millisecond)
}

func (ic *intervalCalculator) Calculate(timerange *TimeRange, minInterval time.Duration) Interval {
	to := timerange.MustGetTo().UnixNano()
	from := timerange.MustGetFrom().UnixNano()
	interval := time.Duration((to - from) / defaultRes)

	if interval < minInterval {
		return Interval{Text: FormatDuration(minInterval), Value: minInterval}
	}

	rounded := roundInterval(interval)
	return Interval{Text: FormatDuration(rounded), Value: rounded}
}

func GetIntervalFrom(dsInfo *models.DataSource, queryModel *simplejson.Json, defaultInterval time.Duration) (time.Duration, error) {
	interval := queryModel.Get("interval").MustString("")

	if interval == "" && dsInfo.JsonData != nil {
		dsInterval := dsInfo.JsonData.Get("timeInterval").MustString("")
		if dsInterval != "" {
			interval = dsInterval
		}
	}

	if interval == "" {
		return defaultInterval, nil
	}

	interval = strings.Replace(strings.Replace(interval, "<", "", 1), ">", "", 1)
	parsedInterval, err := time.ParseDuration(interval)
	if err != nil {
		return time.Duration(0), err
	}

	return parsedInterval, nil
}

// FormatDuration converts a duration into the kbn format e.g. 1m 2h or 3d
func FormatDuration(inter time.Duration) string {
	if inter >= year {
		return fmt.Sprintf("%dy", inter/year)
	}

	if inter >= day {
		return fmt.Sprintf("%dd", inter/day)
	}

	if inter >= time.Hour {
		return fmt.Sprintf("%dh", inter/time.Hour)
	}

	if inter >= time.Minute {
		return fmt.Sprintf("%dm", inter/time.Minute)
	}

	if inter >= time.Second {
		return fmt.Sprintf("%ds", inter/time.Second)
	}

	if inter >= time.Millisecond {
		return fmt.Sprintf("%dms", inter/time.Millisecond)
	}

	return "1ms"
}

func roundInterval(interval time.Duration) time.Duration {
	switch true {
	// 0.015s
	case interval <= 15*time.Millisecond:
		return time.Millisecond * 10 // 0.01s
	// 0.035s
	case interval <= 35*time.Millisecond:
		return time.Millisecond * 20 // 0.02s
	// 0.075s
	case interval <= 75*time.Millisecond:
		return time.Millisecond * 50 // 0.05s
	// 0.15s
	case interval <= 150*time.Millisecond:
		return time.Millisecond * 100 // 0.1s
	// 0.35s
	case interval <= 350*time.Millisecond:
		return time.Millisecond * 200 // 0.2s
	// 0.75s
	case interval <= 750*time.Millisecond:
		return time.Millisecond * 500 // 0.5s
	// 1.5s
	case interval <= 1500*time.Millisecond:
		return time.Millisecond * 1000 // 1s
	// 3.5s
	case interval <= 3500*time.Millisecond:
		return time.Millisecond * 2000 // 2s
	// 7.5s
	case interval <= 7500*time.Millisecond:
		return time.Millisecond * 5000 // 5s
	// 12.5s
	case interval <= 12500*time.Millisecond:
		return time.Millisecond * 10000 // 10s
	// 17.5s
	case interval <= 17500*time.Millisecond:
		return time.Millisecond * 15000 // 15s
	// 25s
	case interval <= 25000*time.Millisecond:
		return time.Millisecond * 20000 // 20s
	// 45s
	case interval <= 45000*time.Millisecond:
		return time.Millisecond * 30000 // 30s
	// 1.5m
	case interval <= 90000*time.Millisecond:
		return time.Millisecond * 60000 // 1m
	// 3.5m
	case interval <= 210000*time.Millisecond:
		return time.Millisecond * 120000 // 2m
	// 7.5m
	case interval <= 450000*time.Millisecond:
		return time.Millisecond * 300000 // 5m
	// 12.5m
	case interval <= 750000*time.Millisecond:
		return time.Millisecond * 600000 // 10m
	// 12.5m
	case interval <= 1050000*time.Millisecond:
		return time.Millisecond * 900000 // 15m
	// 25m
	case interval <= 1500000*time.Millisecond:
		return time.Millisecond * 1200000 // 20m
	// 45m
	case interval <= 2700000*time.Millisecond:
		return time.Millisecond * 1800000 // 30m
	// 1.5h
	case interval <= 5400000*time.Millisecond:
		return time.Millisecond * 3600000 // 1h
	// 2.5h
	case interval <= 9000000*time.Millisecond:
		return time.Millisecond * 7200000 // 2h
	// 4.5h
	case interval <= 16200000*time.Millisecond:
		return time.Millisecond * 10800000 // 3h
	// 9h
	case interval <= 32400000*time.Millisecond:
		return time.Millisecond * 21600000 // 6h
	// 24h
	case interval <= 86400000*time.Millisecond:
		return time.Millisecond * 43200000 // 12h
	// 48h
	case interval <= 172800000*time.Millisecond:
		return time.Millisecond * 86400000 // 24h
	// 1w
	case interval <= 604800000*time.Millisecond:
		return time.Millisecond * 86400000 // 24h
	// 3w
	case interval <= 1814400000*time.Millisecond:
		return time.Millisecond * 604800000 // 1w
	// 2y
	case interval < 3628800000*time.Millisecond:
		return time.Millisecond * 2592000000 // 30d
	default:
		return time.Millisecond * 31536000000 // 1y
	}
}
<|endoftext|>


<commit_before>// package sqrl provides a fluent SQL generator.
//
// See https://github.com/elgris/sqrl for examples.
package sqrl

import (
	"context"
	"database/sql"
	"fmt"
)

// Sqlizer is the interface that wraps the ToSql method.
//
// ToSql returns a SQL representation of the Sqlizer, along with a slice of args
// as passed to e.g. database/sql.Exec. It can also return an error.
type Sqlizer interface {
	ToSql() (string, []interface{}, error)
}

// Execer is the interface that wraps the Exec method.
//
// Exec executes the given query as implemented by database/sql.Exec.
type Execer interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
}

// ExecerContext is the interface that wraps the Exec method.
//
// ExecContext executes the given query using given context as implemented by database/sql.ExecContext.
type ExecerContext interface {
	ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
}

// Queryer is the interface that wraps the Query method.
//
// Query executes the given query as implemented by database/sql.Query.
type Queryer interface {
	Query(query string, args ...interface{}) (*sql.Rows, error)
}

// QueryerContext is the interface that wraps the Query method.
//
// QueryerContext executes the given query using given context as implemented by database/sql.QueryContext.
type QueryerContext interface {
	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
}

// QueryRower is the interface that wraps the QueryRow method.
//
// QueryRow executes the given query as implemented by database/sql.QueryRow.
type QueryRower interface {
	QueryRow(query string, args ...interface{}) RowScanner
}

// QueryRowerContext is the interface that wraps the QueryRow method.
//
// QueryRowContext executes the given query using given context as implemented by database/sql.QueryRowContext.
type QueryRowerContext interface {
	QueryRowContext(ctx context.Context, query string, args ...interface{}) RowScanner
}

// BaseRunner groups the Execer and Queryer interfaces.
type BaseRunner interface {
	Execer
	ExecerContext
	Queryer
	QueryerContext
}

// Runner groups the Execer, Queryer, and QueryRower interfaces.
type Runner interface {
	Execer
	ExecerContext
	Queryer
	QueryerContext
	QueryRower
	QueryRowerContext
}

// ErrRunnerNotSet is returned by methods that need a Runner if it isn't set.
var ErrRunnerNotSet = fmt.Errorf("cannot run; no Runner set (RunWith)")

// ErrRunnerNotQueryRunner is returned by QueryRow if the RunWith value doesn't implement QueryRower.
var ErrRunnerNotQueryRunner = fmt.Errorf("cannot QueryRow; Runner is not a QueryRower")

// ErrRunnerNotQueryRunnerContext is returned by QueryRowContext if the RunWith value doesn't implement QueryRowerContext.
var ErrRunnerNotQueryRunnerContext = fmt.Errorf("cannot QueryRow; Runner is not a QueryRowerContext")

// ExecWith Execs the SQL returned by s with db.
func ExecWith(db Execer, s Sqlizer) (res sql.Result, err error) {
	query, args, err := s.ToSql()
	if err != nil {
		return
	}
	return db.Exec(query, args...)
}

// ExecWithContext Execs the SQL returned by s with db.
func ExecWithContext(ctx context.Context, db ExecerContext, s Sqlizer) (res sql.Result, err error) {
	query, args, err := s.ToSql()
	if err != nil {
		return
	}
	return db.ExecContext(ctx, query, args...)
}

// QueryWith Querys the SQL returned by s with db.
func QueryWith(db Queryer, s Sqlizer) (rows *sql.Rows, err error) {
	query, args, err := s.ToSql()
	if err != nil {
		return
	}
	return db.Query(query, args...)
}

// QueryWithContext Querys the SQL returned by s with db.
func QueryWithContext(ctx context.Context, db QueryerContext, s Sqlizer) (rows *sql.Rows, err error) {
	query, args, err := s.ToSql()
	if err != nil {
		return
	}
	return db.QueryContext(ctx, query, args...)
}

// QueryRowWith QueryRows the SQL returned by s with db.
func QueryRowWith(db QueryRower, s Sqlizer) RowScanner {
	query, args, err := s.ToSql()
	return &Row{RowScanner: db.QueryRow(query, args...), err: err}
}

// QueryRowWithContext QueryRows the SQL returned by s with db.
func QueryRowWithContext(ctx context.Context, db QueryRowerContext, s Sqlizer) RowScanner {
	query, args, err := s.ToSql()
	return &Row{RowScanner: db.QueryRowContext(ctx, query, args...), err: err}
}

// DBRunner wraps sql.DB to implement Runner.
type dbRunner struct {
	*sql.DB
}

func (r *dbRunner) QueryRow(query string, args ...interface{}) RowScanner {
	return r.DB.QueryRow(query, args...)
}

func (r *dbRunner) QueryRowContext(ctx context.Context, query string, args ...interface{}) RowScanner {
	return r.DB.QueryRowContext(ctx, query, args...)
}

// TxRunner wraps sql.Tx to implement Runner.
type txRunner struct {
	*sql.Tx
}

func (r *txRunner) QueryRow(query string, args ...interface{}) RowScanner {
	return r.Tx.QueryRow(query, args...)
}

func (r *txRunner) QueryRowContext(ctx context.Context, query string, args ...interface{}) RowScanner {
	return r.Tx.QueryRowContext(ctx, query, args...)
}

// WrapRunner returns Runner that implements QueryRower and QueryRowerContext.
func wrapRunner(baseRunner BaseRunner) (runner Runner) {
	switch r := baseRunner.(type) {
	case Runner:
		runner = r
	case *sql.DB:
		runner = &dbRunner{r}
	case *sql.Tx:
		runner = &txRunner{r}
	}
	return
}
<commit_msg>Update wrapRunner to return BaseRunner for wider compatibility<commit_after>// package sqrl provides a fluent SQL generator.
//
// See https://github.com/elgris/sqrl for examples.
package sqrl

import (
	"context"
	"database/sql"
	"fmt"
)

// Sqlizer is the interface that wraps the ToSql method.
//
// ToSql returns a SQL representation of the Sqlizer, along with a slice of args
// as passed to e.g. database/sql.Exec. It can also return an error.
type Sqlizer interface {
	ToSql() (string, []interface{}, error)
}

// Execer is the interface that wraps the Exec method.
//
// Exec executes the given query as implemented by database/sql.Exec.
type Execer interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
}

// ExecerContext is the interface that wraps the Exec method.
//
// ExecContext executes the given query using given context as implemented by database/sql.ExecContext.
type ExecerContext interface {
	ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
}

// Queryer is the interface that wraps the Query method.
//
// Query executes the given query as implemented by database/sql.Query.
type Queryer interface {
	Query(query string, args ...interface{}) (*sql.Rows, error)
}

// QueryerContext is the interface that wraps the Query method.
//
// QueryerContext executes the given query using given context as implemented by database/sql.QueryContext.
type QueryerContext interface {
	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
}

// QueryRower is the interface that wraps the QueryRow method.
//
// QueryRow executes the given query as implemented by database/sql.QueryRow.
type QueryRower interface {
	QueryRow(query string, args ...interface{}) RowScanner
}

// QueryRowerContext is the interface that wraps the QueryRow method.
//
// QueryRowContext executes the given query using given context as implemented by database/sql.QueryRowContext.
type QueryRowerContext interface {
	QueryRowContext(ctx context.Context, query string, args ...interface{}) RowScanner
}

// BaseRunner groups the Execer and Queryer interfaces.
type BaseRunner interface {
	Execer
	ExecerContext
	Queryer
	QueryerContext
}

// Runner groups the Execer, Queryer, and QueryRower interfaces.
type Runner interface {
	Execer
	ExecerContext
	Queryer
	QueryerContext
	QueryRower
	QueryRowerContext
}

// ErrRunnerNotSet is returned by methods that need a Runner if it isn't set.
var ErrRunnerNotSet = fmt.Errorf("cannot run; no Runner set (RunWith)")

// ErrRunnerNotQueryRunner is returned by QueryRow if the RunWith value doesn't implement QueryRower.
var ErrRunnerNotQueryRunner = fmt.Errorf("cannot QueryRow; Runner is not a QueryRower")

// ErrRunnerNotQueryRunnerContext is returned by QueryRowContext if the RunWith value doesn't implement QueryRowerContext.
var ErrRunnerNotQueryRunnerContext = fmt.Errorf("cannot QueryRow; Runner is not a QueryRowerContext")

// ExecWith Execs the SQL returned by s with db.
func ExecWith(db Execer, s Sqlizer) (res sql.Result, err error) {
	query, args, err := s.ToSql()
	if err != nil {
		return
	}
	return db.Exec(query, args...)
}

// ExecWithContext Execs the SQL returned by s with db.
func ExecWithContext(ctx context.Context, db ExecerContext, s Sqlizer) (res sql.Result, err error) {
	query, args, err := s.ToSql()
	if err != nil {
		return
	}
	return db.ExecContext(ctx, query, args...)
}

// QueryWith Querys the SQL returned by s with db.
func QueryWith(db Queryer, s Sqlizer) (rows *sql.Rows, err error) {
	query, args, err := s.ToSql()
	if err != nil {
		return
	}
	return db.Query(query, args...)
}

// QueryWithContext Querys the SQL returned by s with db.
func QueryWithContext(ctx context.Context, db QueryerContext, s Sqlizer) (rows *sql.Rows, err error) {
	query, args, err := s.ToSql()
	if err != nil {
		return
	}
	return db.QueryContext(ctx, query, args...)
}

// QueryRowWith QueryRows the SQL returned by s with db.
func QueryRowWith(db QueryRower, s Sqlizer) RowScanner {
	query, args, err := s.ToSql()
	return &Row{RowScanner: db.QueryRow(query, args...), err: err}
}

// QueryRowWithContext QueryRows the SQL returned by s with db.
func QueryRowWithContext(ctx context.Context, db QueryRowerContext, s Sqlizer) RowScanner {
	query, args, err := s.ToSql()
	return &Row{RowScanner: db.QueryRowContext(ctx, query, args...), err: err}
}

// DBRunner wraps sql.DB to implement Runner.
type dbRunner struct {
	*sql.DB
}

func (r *dbRunner) QueryRow(query string, args ...interface{}) RowScanner {
	return r.DB.QueryRow(query, args...)
}

func (r *dbRunner) QueryRowContext(ctx context.Context, query string, args ...interface{}) RowScanner {
	return r.DB.QueryRowContext(ctx, query, args...)
}

// TxRunner wraps sql.Tx to implement Runner.
type txRunner struct {
	*sql.Tx
}

func (r *txRunner) QueryRow(query string, args ...interface{}) RowScanner {
	return r.Tx.QueryRow(query, args...)
}

func (r *txRunner) QueryRowContext(ctx context.Context, query string, args ...interface{}) RowScanner {
	return r.Tx.QueryRowContext(ctx, query, args...)
}

// WrapRunner returns Runner for sql.DB and sql.Tx, or BaseRunner otherwise.
func wrapRunner(baseRunner BaseRunner) (runner BaseRunner) {
	switch r := baseRunner.(type) {
	case *sql.DB:
		runner = &dbRunner{r}
	case *sql.Tx:
		runner = &txRunner{r}
	default:
		runner = r
	}
	return
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"github.com/hashicorp/serf/serf"
	"io/ioutil"
	"net"
	"testing"
)

const eventScript = `#!/bin/sh
RESULT_FILE="%s"
echo $SERF_SELF_NAME $SERF_SELF_ROLE >>${RESULT_FILE}
echo $SERF_EVENT $SERF_USER_EVENT "$@" >>${RESULT_FILE}
while read line; do
	printf "${line}\n" >>${RESULT_FILE}
done
`

const userEventScript = `#!/bin/sh
RESULT_FILE="%s"
echo $SERF_SELF_NAME $SERF_SELF_ROLE >>${RESULT_FILE}
echo $SERF_EVENT $SERF_USER_EVENT "$@" >>${RESULT_FILE}
echo $SERF_EVENT $SERF_USER_LTIME "$@" >>${RESULT_FILE}
while read line; do
	printf "${line}\n" >>${RESULT_FILE}
done
`

// testEventScript creates an event script that can be used with the
// agent. It returns the path to the event script itself and a path to
// the file that will contain the events that that script receives.
func testEventScript(t *testing.T, script string) (string, string) {
	scriptFile, err := ioutil.TempFile("", "serf")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer scriptFile.Close()

	if err := scriptFile.Chmod(0755); err != nil {
		t.Fatalf("err: %s", err)
	}

	resultFile, err := ioutil.TempFile("", "serf-result")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer resultFile.Close()

	_, err = scriptFile.Write([]byte(
		fmt.Sprintf(script, resultFile.Name())))
	if err != nil {
		t.Fatalf("err: %s")
	}

	return scriptFile.Name(), resultFile.Name()
}

func TestScriptEventHandler(t *testing.T) {
	script, results := testEventScript(t, eventScript)

	h := &ScriptEventHandler{
		Self: serf.Member{
			Name: "ourname",
			Role: "ourrole",
		},
		Scripts: []EventScript{
			{
				EventFilter: EventFilter{
					Event: "*",
				},
				Script: script,
			},
		},
	}

	event := serf.MemberEvent{
		Type: serf.EventMemberJoin,
		Members: []serf.Member{
			{
				Name: "foo",
				Addr: net.ParseIP("1.2.3.4"),
				Role: "bar",
			},
		},
	}

	h.HandleEvent(event)

	result, err := ioutil.ReadFile(results)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := "ourname ourrole\nmember-join\nfoo\t1.2.3.4\tbar\n"
	if string(result) != expected {
		t.Fatalf("bad: %#v. Expected: %#v", string(result), expected)
	}
}

func TestScriptUserEventHandler(t *testing.T) {
	script, results := testEventScript(t, userEventScript)

	h := &ScriptEventHandler{
		Self: serf.Member{
			Name: "ourname",
			Role: "ourrole",
		},
		Scripts: []EventScript{
			{
				EventFilter: EventFilter{
					Event: "*",
				},
				Script: script,
			},
		},
	}

	userEvent := serf.UserEvent{
		LTime:    1,
		Name:     "baz",
		Payload:  []byte("foobar"),
		Coalesce: true,
	}

	h.HandleEvent(userEvent)

	result, err := ioutil.ReadFile(results)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := "ourname ourrole\nuser baz\nuser 1\n"
	if string(result) != expected {
		t.Fatalf("bad: %#v. Expected: %#v", string(result), expected)
	}
}

func TestEventScriptInvoke(t *testing.T) {
	testCases := []struct {
		script EventScript
		event  serf.Event
		invoke bool
	}{
		{
			EventScript{EventFilter{"*", ""}, "script.sh"},
			serf.MemberEvent{},
			true,
		},
		{
			EventScript{EventFilter{"user", ""}, "script.sh"},
			serf.MemberEvent{},
			false,
		},
		{
			EventScript{EventFilter{"user", "deploy"}, "script.sh"},
			serf.UserEvent{Name: "deploy"},
			true,
		},
		{
			EventScript{EventFilter{"user", "deploy"}, "script.sh"},
			serf.UserEvent{Name: "restart"},
			false,
		},
		{
			EventScript{EventFilter{"member-join", ""}, "script.sh"},
			serf.MemberEvent{Type: serf.EventMemberJoin},
			true,
		},
		{
			EventScript{EventFilter{"member-join", ""}, "script.sh"},
			serf.MemberEvent{Type: serf.EventMemberLeave},
			false,
		},
	}

	for _, tc := range testCases {
		result := tc.script.Invoke(tc.event)
		if result != tc.invoke {
			t.Errorf("bad: %#v", tc)
		}
	}
}

func TestEventScriptValid(t *testing.T) {
	testCases := []struct {
		Event string
		Valid bool
	}{
		{"member-join", true},
		{"member-leave", true},
		{"member-failed", true},
		{"user", true},
		{"User", false},
		{"member", false},
		{"*", true},
	}

	for _, tc := range testCases {
		script := EventScript{EventFilter: EventFilter{Event: tc.Event}}
		if script.Valid() != tc.Valid {
			t.Errorf("bad: %#v", tc)
		}
	}
}

func TestParseEventScript(t *testing.T) {
	testCases := []struct {
		v       string
		err     bool
		results []EventScript
	}{
		{
			"script.sh",
			false,
			[]EventScript{{EventFilter{"*", ""}, "script.sh"}},
		},

		{
			"member-join=script.sh",
			false,
			[]EventScript{{EventFilter{"member-join", ""}, "script.sh"}},
		},

		{
			"foo,bar=script.sh",
			false,
			[]EventScript{
				{EventFilter{"foo", ""}, "script.sh"},
				{EventFilter{"bar", ""}, "script.sh"},
			},
		},

		{
			"user:deploy=script.sh",
			false,
			[]EventScript{{EventFilter{"user", "deploy"}, "script.sh"}},
		},

		{
			"foo,user:blah,bar=script.sh",
			false,
			[]EventScript{
				{EventFilter{"foo", ""}, "script.sh"},
				{EventFilter{"user", "blah"}, "script.sh"},
				{EventFilter{"bar", ""}, "script.sh"},
			},
		},
	}

	for _, tc := range testCases {
		results := ParseEventScript(tc.v)
		if results == nil {
			t.Errorf("result should not be nil")
			continue
		}

		if len(results) != len(tc.results) {
			t.Errorf("bad: %#v", results)
			continue
		}

		for i, r := range results {
			expected := tc.results[i]

			if r.Event != expected.Event {
				t.Errorf("Events not equal: %s %s", r.Event, expected.Event)
			}

			if r.UserEvent != expected.UserEvent {
				t.Errorf("User events not equal: %s %s", r.UserEvent, expected.UserEvent)
			}

			if r.Script != expected.Script {
				t.Errorf("Scripts not equal: %s %s", r.Script, expected.Script)
			}
		}
	}
}
<commit_msg>agent: Adding tests for ParseEventFilter<commit_after>package agent

import (
	"fmt"
	"github.com/hashicorp/serf/serf"
	"io/ioutil"
	"net"
	"testing"
)

const eventScript = `#!/bin/sh
RESULT_FILE="%s"
echo $SERF_SELF_NAME $SERF_SELF_ROLE >>${RESULT_FILE}
echo $SERF_EVENT $SERF_USER_EVENT "$@" >>${RESULT_FILE}
while read line; do
	printf "${line}\n" >>${RESULT_FILE}
done
`

const userEventScript = `#!/bin/sh
RESULT_FILE="%s"
echo $SERF_SELF_NAME $SERF_SELF_ROLE >>${RESULT_FILE}
echo $SERF_EVENT $SERF_USER_EVENT "$@" >>${RESULT_FILE}
echo $SERF_EVENT $SERF_USER_LTIME "$@" >>${RESULT_FILE}
while read line; do
	printf "${line}\n" >>${RESULT_FILE}
done
`

// testEventScript creates an event script that can be used with the
// agent. It returns the path to the event script itself and a path to
// the file that will contain the events that that script receives.
func testEventScript(t *testing.T, script string) (string, string) {
	scriptFile, err := ioutil.TempFile("", "serf")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer scriptFile.Close()

	if err := scriptFile.Chmod(0755); err != nil {
		t.Fatalf("err: %s", err)
	}

	resultFile, err := ioutil.TempFile("", "serf-result")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer resultFile.Close()

	_, err = scriptFile.Write([]byte(
		fmt.Sprintf(script, resultFile.Name())))
	if err != nil {
		t.Fatalf("err: %s")
	}

	return scriptFile.Name(), resultFile.Name()
}

func TestScriptEventHandler(t *testing.T) {
	script, results := testEventScript(t, eventScript)

	h := &ScriptEventHandler{
		Self: serf.Member{
			Name: "ourname",
			Role: "ourrole",
		},
		Scripts: []EventScript{
			{
				EventFilter: EventFilter{
					Event: "*",
				},
				Script: script,
			},
		},
	}

	event := serf.MemberEvent{
		Type: serf.EventMemberJoin,
		Members: []serf.Member{
			{
				Name: "foo",
				Addr: net.ParseIP("1.2.3.4"),
				Role: "bar",
			},
		},
	}

	h.HandleEvent(event)

	result, err := ioutil.ReadFile(results)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := "ourname ourrole\nmember-join\nfoo\t1.2.3.4\tbar\n"
	if string(result) != expected {
		t.Fatalf("bad: %#v. Expected: %#v", string(result), expected)
	}
}

func TestScriptUserEventHandler(t *testing.T) {
	script, results := testEventScript(t, userEventScript)

	h := &ScriptEventHandler{
		Self: serf.Member{
			Name: "ourname",
			Role: "ourrole",
		},
		Scripts: []EventScript{
			{
				EventFilter: EventFilter{
					Event: "*",
				},
				Script: script,
			},
		},
	}

	userEvent := serf.UserEvent{
		LTime:    1,
		Name:     "baz",
		Payload:  []byte("foobar"),
		Coalesce: true,
	}

	h.HandleEvent(userEvent)

	result, err := ioutil.ReadFile(results)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := "ourname ourrole\nuser baz\nuser 1\n"
	if string(result) != expected {
		t.Fatalf("bad: %#v. Expected: %#v", string(result), expected)
	}
}

func TestEventScriptInvoke(t *testing.T) {
	testCases := []struct {
		script EventScript
		event  serf.Event
		invoke bool
	}{
		{
			EventScript{EventFilter{"*", ""}, "script.sh"},
			serf.MemberEvent{},
			true,
		},
		{
			EventScript{EventFilter{"user", ""}, "script.sh"},
			serf.MemberEvent{},
			false,
		},
		{
			EventScript{EventFilter{"user", "deploy"}, "script.sh"},
			serf.UserEvent{Name: "deploy"},
			true,
		},
		{
			EventScript{EventFilter{"user", "deploy"}, "script.sh"},
			serf.UserEvent{Name: "restart"},
			false,
		},
		{
			EventScript{EventFilter{"member-join", ""}, "script.sh"},
			serf.MemberEvent{Type: serf.EventMemberJoin},
			true,
		},
		{
			EventScript{EventFilter{"member-join", ""}, "script.sh"},
			serf.MemberEvent{Type: serf.EventMemberLeave},
			false,
		},
	}

	for _, tc := range testCases {
		result := tc.script.Invoke(tc.event)
		if result != tc.invoke {
			t.Errorf("bad: %#v", tc)
		}
	}
}

func TestEventScriptValid(t *testing.T) {
	testCases := []struct {
		Event string
		Valid bool
	}{
		{"member-join", true},
		{"member-leave", true},
		{"member-failed", true},
		{"user", true},
		{"User", false},
		{"member", false},
		{"*", true},
	}

	for _, tc := range testCases {
		script := EventScript{EventFilter: EventFilter{Event: tc.Event}}
		if script.Valid() != tc.Valid {
			t.Errorf("bad: %#v", tc)
		}
	}
}

func TestParseEventScript(t *testing.T) {
	testCases := []struct {
		v       string
		err     bool
		results []EventScript
	}{
		{
			"script.sh",
			false,
			[]EventScript{{EventFilter{"*", ""}, "script.sh"}},
		},

		{
			"member-join=script.sh",
			false,
			[]EventScript{{EventFilter{"member-join", ""}, "script.sh"}},
		},

		{
			"foo,bar=script.sh",
			false,
			[]EventScript{
				{EventFilter{"foo", ""}, "script.sh"},
				{EventFilter{"bar", ""}, "script.sh"},
			},
		},

		{
			"user:deploy=script.sh",
			false,
			[]EventScript{{EventFilter{"user", "deploy"}, "script.sh"}},
		},

		{
			"foo,user:blah,bar=script.sh",
			false,
			[]EventScript{
				{EventFilter{"foo", ""}, "script.sh"},
				{EventFilter{"user", "blah"}, "script.sh"},
				{EventFilter{"bar", ""}, "script.sh"},
			},
		},
	}

	for _, tc := range testCases {
		results := ParseEventScript(tc.v)
		if results == nil {
			t.Errorf("result should not be nil")
			continue
		}

		if len(results) != len(tc.results) {
			t.Errorf("bad: %#v", results)
			continue
		}

		for i, r := range results {
			expected := tc.results[i]

			if r.Event != expected.Event {
				t.Errorf("Events not equal: %s %s", r.Event, expected.Event)
			}

			if r.UserEvent != expected.UserEvent {
				t.Errorf("User events not equal: %s %s", r.UserEvent, expected.UserEvent)
			}

			if r.Script != expected.Script {
				t.Errorf("Scripts not equal: %s %s", r.Script, expected.Script)
			}
		}
	}
}

func TestParseEventFilter(t *testing.T) {
	testCases := []struct {
		v       string
		results []EventFilter
	}{
		{
			"",
			[]EventFilter{EventFilter{"*", ""}},
		},

		{
			"member-join",
			[]EventFilter{EventFilter{"member-join", ""}},
		},

		{
			"foo,bar",
			[]EventFilter{
				EventFilter{"foo", ""},
				EventFilter{"bar", ""},
			},
		},

		{
			"user:deploy",
			[]EventFilter{EventFilter{"user", "deploy"}},
		},

		{
			"foo,user:blah,bar",
			[]EventFilter{
				EventFilter{"foo", ""},
				EventFilter{"user", "blah"},
				EventFilter{"bar", ""},
			},
		},
	}

	for _, tc := range testCases {
		results := ParseEventFilter(tc.v)
		if results == nil {
			t.Errorf("result should not be nil")
			continue
		}

		if len(results) != len(tc.results) {
			t.Errorf("bad: %#v", results)
			continue
		}

		for i, r := range results {
			expected := tc.results[i]

			if r.Event != expected.Event {
				t.Errorf("Events not equal: %s %s", r.Event, expected.Event)
			}

			if r.UserEvent != expected.UserEvent {
				t.Errorf("User events not equal: %s %s", r.UserEvent, expected.UserEvent)
			}
		}
	}
}
<|endoftext|>


<commit_before>package signature

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

const testTs = "1544544948"
const testQp = "abc=foo&def=bar"
const testBody = `{"a key":"some value"}`
const testSignature = "orb0adPhRCYND1WCAvPBr+qjm4STGtyvNDIDNBZ4Ir4="
const testKey = "other-secret"

func TestCalculateSignature(t *testing.T) {
	var cases = []struct {
		sKey string
		ts   string
		qp   string
		b    string
		es   string
		e    bool
	}{
		{
			sKey: testKey,
			ts:   testTs,
			qp:   testQp,
			b:    testBody,
			es:   testSignature,
			e:    true,
		},
		{
			sKey: testKey,
			ts:   testTs,
			qp:   testQp,
			b:    testBody,
			es:   "LISw4Je7n0/MkYDgVSzTJm8dW6BkytKTXMZZk1IElMs=",
			e:    false,
		},
		{
			sKey: "secret",
			ts:   testTs,
			qp:   "",
			b:    "",
			es:   "LISw4Je7n0/MkYDgVSzTJm8dW6BkytKTXMZZk1IElMs=",
			e:    true,
		},
		{
			sKey: "secret",
			ts:   testTs,
			qp:   "",
			b:    testBody,
			es:   "p2e20OtAg39DEmz1ORHpjQ556U4o1ZaH4NWbM9Q8Qjk=",
			e:    true,
		},
		{
			sKey: "secret",
			ts:   testTs,
			qp:   testQp,
			b:    "",
			es:   "Tfn+nRUBsn6lQgf6IpxBMS1j9lm7XsGjt5xh47M3jCk=",
			e:    true,
		},
	}
	for i, tt := range cases {
		v := NewValidator(tt.sKey, nil)
		s, err := v.CalculateSignature(tt.ts, tt.qp, []byte(tt.b))
		if err != nil {
			t.Errorf("Error calculating signature: %s, expected: %s", s, tt.es)
		}
		drs, _ := base64.StdEncoding.DecodeString(tt.es)
		e := bool(bytes.Compare(s, drs) == 0)
		if e != tt.e {
			t.Errorf("Unexpected signature: %s, test case: %d", s, i)
		}
	}
}
func TestValidTimestamp(t *testing.T) {
	var p float64 = 2
	now := time.Now()
	nowts := fmt.Sprintf("%d", now.Unix())
	var cases = []struct {
		ts string
		p  ValidityPeriod
		e  bool
	}{
		{
			ts: nowts,
			p:  nil,
			e:  true,
		},
		{
			ts: "",
			p:  nil,
			e:  false,
		},
		{
			ts: "wrongTs",
			p:  nil,
			e:  false,
		},
		{
			ts: nowts,
			p:  &p,
			e:  true,
		},
		{
			ts: fmt.Sprintf("%d", now.AddDate(0, 0, 1).Unix()),
			p:  &p,
			e:  false,
		},
		{
			ts: fmt.Sprintf("%d", now.AddDate(0, 0, -1).Unix()),
			p:  &p,
			e:  false,
		},
	}

	for i, tt := range cases {
		v := NewValidator(testKey, tt.p)
		r := v.ValidTimestamp(tt.ts)
		if r != tt.e {
			t.Errorf("Unexpected error validating ts: %s, test case: %d", tt.ts, i)
		}
	}
}

func TestValidSignature(t *testing.T) {
	var cases = []struct {
		ts string
		qp string
		b  string
		s  string
		e  bool
	}{
		{
			ts: testTs,
			qp: testQp,
			b:  testBody,
			s:  testSignature,
			e:  true,
		},
		{
			ts: testTs,
			qp: "def=bar&abc=foo",
			b:  testBody,
			s:  testSignature,
			e:  true,
		},
		{
			ts: testTs,
			qp: testQp,
			b:  testBody,
			s:  "wrong signature",
			e:  false,
		},
	}

	for i, tt := range cases {
		v := NewValidator(testKey, nil)
		r := v.ValidSignature(tt.ts, tt.qp, []byte(tt.b), tt.s)
		if r != tt.e {
			t.Errorf("Unexpected error validating signature: %s, test case: %d", tt.s, i)
		}
	}
}
func TestValidate(t *testing.T) {
	var cases = []struct {
		k   string
		ts  string
		s   string
		sh  string
		tsh string
		e   int
	}{
		{
			k:   testKey,
			ts:  testTs,
			s:   testSignature,
			sh:  sHeader,
			tsh: tsHeader,
			e:   http.StatusOK,
		},
		{
			k:   "",
			ts:  testTs,
			s:   testSignature,
			sh:  sHeader,
			tsh: tsHeader,
			e:   http.StatusUnauthorized,
		},
		{
			k:   testKey,
			ts:  "",
			s:   testSignature,
			sh:  sHeader,
			tsh: tsHeader,
			e:   http.StatusUnauthorized,
		},
		{
			k:   testKey,
			ts:  testTs,
			s:   "",
			sh:  sHeader,
			tsh: tsHeader,
			e:   http.StatusUnauthorized,
		},
		{
			k:   testKey,
			ts:  testTs,
			s:   testSignature,
			sh:  "wrong-header",
			tsh: tsHeader,
			e:   http.StatusUnauthorized,
		},
		{
			k:   testKey,
			ts:  testTs,
			s:   testSignature,
			sh:  sHeader,
			tsh: "wrong-header",
			e:   http.StatusUnauthorized,
		},
	}

	for i, tt := range cases {
		v := NewValidator(tt.k, nil)
		ts := httptest.NewServer(v.Validate(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
		})))
		defer ts.Close()

		client := &http.Client{}
		req, _ := http.NewRequest("GET", ts.URL+"?"+testQp, strings.NewReader(testBody))
		req.Header.Set(tt.sh, tt.s)
		req.Header.Set(tt.tsh, tt.ts)
		res, _ := client.Do(req)
		if res.StatusCode != tt.e {
			t.Errorf("Unexpected response code: %s, test case: %d", res.Status, i)
		}
	}

}
<commit_msg>Style: Add name to test cases<commit_after>package signature

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

const testTs = "1544544948"
const testQp = "abc=foo&def=bar"
const testBody = `{"a key":"some value"}`
const testSignature = "orb0adPhRCYND1WCAvPBr+qjm4STGtyvNDIDNBZ4Ir4="
const testKey = "other-secret"

func TestCalculateSignature(t *testing.T) {
	var cases = []struct {
		name string
		sKey string
		ts   string
		qp   string
		b    string
		es   string
		e    bool
	}{
		{
			name: "Succesful",
			sKey: testKey,
			ts:   testTs,
			qp:   testQp,
			b:    testBody,
			es:   testSignature,
			e:    true,
		},
		{
			name: "Wrong signature",
			sKey: testKey,
			ts:   testTs,
			qp:   testQp,
			b:    testBody,
			es:   "LISw4Je7n0/MkYDgVSzTJm8dW6BkytKTXMZZk1IElMs=",
			e:    false,
		},
		{
			name: "Empty query params and body",
			sKey: "secret",
			ts:   testTs,
			qp:   "",
			b:    "",
			es:   "LISw4Je7n0/MkYDgVSzTJm8dW6BkytKTXMZZk1IElMs=",
			e:    true,
		},
		{
			name: "Empty query params",
			sKey: "secret",
			ts:   testTs,
			qp:   "",
			b:    testBody,
			es:   "p2e20OtAg39DEmz1ORHpjQ556U4o1ZaH4NWbM9Q8Qjk=",
			e:    true,
		},
		{
			name: "Empty body",
			sKey: "secret",
			ts:   testTs,
			qp:   testQp,
			b:    "",
			es:   "Tfn+nRUBsn6lQgf6IpxBMS1j9lm7XsGjt5xh47M3jCk=",
			e:    true,
		},
	}
	for _, tt := range cases {
		v := NewValidator(tt.sKey, nil)
		s, err := v.CalculateSignature(tt.ts, tt.qp, []byte(tt.b))
		if err != nil {
			t.Errorf("Error calculating signature: %s, expected: %s", s, tt.es)
		}
		drs, _ := base64.StdEncoding.DecodeString(tt.es)
		e := bool(bytes.Compare(s, drs) == 0)
		if e != tt.e {
			t.Errorf("Unexpected signature: %s, test case: %s", s, tt.name)
		}
	}
}
func TestValidTimestamp(t *testing.T) {
	var p float64 = 2
	now := time.Now()
	nowts := fmt.Sprintf("%d", now.Unix())
	var cases = []struct {
		name string
		ts   string
		p    ValidityPeriod
		e    bool
	}{
		{
			name: "Succesful",
			ts:   nowts,
			p:    nil,
			e:    true,
		},
		{
			name: "Empty time stamp",
			ts:   "",
			p:    nil,
			e:    false,
		},
		{
			name: "Invalid time stamp",
			ts:   "wrongTs",
			p:    nil,
			e:    false,
		},
		{
			name: "Succesful",
			ts:   nowts,
			p:    &p,
			e:    true,
		},
		{
			name: "Time stamp 24 hours in the futute",
			ts:   fmt.Sprintf("%d", now.AddDate(0, 0, 1).Unix()),
			p:    &p,
			e:    false,
		},
		{
			name: "Time stamp 24 hours in the past",
			ts:   fmt.Sprintf("%d", now.AddDate(0, 0, -1).Unix()),
			p:    &p,
			e:    false,
		},
	}

	for _, tt := range cases {
		v := NewValidator(testKey, tt.p)
		r := v.ValidTimestamp(tt.ts)
		if r != tt.e {
			t.Errorf("Unexpected error validating ts: %s, test case: %s", tt.ts, tt.name)
		}
	}
}

func TestValidSignature(t *testing.T) {
	var cases = []struct {
		name string
		ts   string
		qp   string
		b    string
		s    string
		e    bool
	}{
		{
			name: "succesful",
			ts:   testTs,
			qp:   testQp,
			b:    testBody,
			s:    testSignature,
			e:    true,
		},
		{
			name: "Unorganized query params",
			ts:   testTs,
			qp:   "def=bar&abc=foo",
			b:    testBody,
			s:    testSignature,
			e:    true,
		},
		{
			name: "Wrong signature received",
			ts:   testTs,
			qp:   testQp,
			b:    testBody,
			s:    "wrong signature",
			e:    false,
		},
	}

	for _, tt := range cases {
		v := NewValidator(testKey, nil)
		r := v.ValidSignature(tt.ts, tt.qp, []byte(tt.b), tt.s)
		if r != tt.e {
			t.Errorf("Unexpected error validating signature: %s, test case: %s", tt.s, tt.name)
		}
	}
}
func TestValidate(t *testing.T) {
	var cases = []struct {
		name string
		k    string
		ts   string
		s    string
		sh   string
		tsh  string
		e    int
	}{
		{
			name: "Succesful",
			k:    testKey,
			ts:   testTs,
			s:    testSignature,
			sh:   sHeader,
			tsh:  tsHeader,
			e:    http.StatusOK,
		},
		{
			name: "NO Access key",
			k:    "",
			ts:   testTs,
			s:    testSignature,
			sh:   sHeader,
			tsh:  tsHeader,
			e:    http.StatusUnauthorized,
		},
		{
			name: "Request with empty time stamp",
			k:    testKey,
			ts:   "",
			s:    testSignature,
			sh:   sHeader,
			tsh:  tsHeader,
			e:    http.StatusUnauthorized,
		},
		{
			name: "Request with empty signature",
			k:    testKey,
			ts:   testTs,
			s:    "",
			sh:   sHeader,
			tsh:  tsHeader,
			e:    http.StatusUnauthorized,
		},
		{
			name: "Request with wrong signature header",
			k:    testKey,
			ts:   testTs,
			s:    testSignature,
			sh:   "wrong-header",
			tsh:  tsHeader,
			e:    http.StatusUnauthorized,
		},
		{
			name: "Request with wrong timestamp header",
			k:    testKey,
			ts:   testTs,
			s:    testSignature,
			sh:   sHeader,
			tsh:  "wrong-header",
			e:    http.StatusUnauthorized,
		},
	}

	for _, tt := range cases {
		v := NewValidator(tt.k, nil)
		ts := httptest.NewServer(v.Validate(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
		})))
		defer ts.Close()

		client := &http.Client{}
		req, _ := http.NewRequest("GET", ts.URL+"?"+testQp, strings.NewReader(testBody))
		req.Header.Set(tt.sh, tt.s)
		req.Header.Set(tt.tsh, tt.ts)
		res, _ := client.Do(req)
		if res.StatusCode != tt.e {
			t.Errorf("Unexpected response code: %s, test case: %s", res.Status, tt.name)
		}
	}

}
<|endoftext|>