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

// +build !appengine

// Package socket implements an WebSocket-based playground backend.
// Clients connect to a websocket handler and send run/kill commands, and
// the server sends the output and exit status of the running processes.
// Multiple clients running multiple processes may be served concurrently.
// The wire format is JSON and is described by the Message type.
//
// This will not run on App Engine as WebSockets are not supported there.
package socket

import (
	"bytes"
	"encoding/json"
	"errors"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"code.google.com/p/go.net/websocket"
)

// Handler implements a WebSocket handler for a client connection.
var Handler = websocket.Handler(socketHandler)

// Environ provides an environment when a binary, such as the go tool, is
// invoked.
var Environ func() []string = os.Environ

const (
	// The maximum number of messages to send per session (avoid flooding).
	msgLimit = 1000

	// Batch messages sent in this interval and send as a single message.
	msgDelay = 10 * time.Millisecond
)

// Message is the wire format for the websocket connection to the browser.
// It is used for both sending output messages and receiving commands, as
// distinguished by the Kind field.
type Message struct {
	Id      string // client-provided unique id for the process
	Kind    string // in: "run", "kill" out: "stdout", "stderr", "end"
	Body    string
	Options *Options `json:",omitempty"`
}

// Options specify additional message options.
type Options struct {
	Race bool // use -race flag when building code (for "run" only)
}

// socketHandler handles the websocket connection for a given present session.
// It handles transcoding Messages to and from JSON format, and starting
// and killing processes.
func socketHandler(c *websocket.Conn) {
	in, out := make(chan *Message), make(chan *Message)
	errc := make(chan error, 1)

	// Decode messages from client and send to the in channel.
	go func() {
		dec := json.NewDecoder(c)
		for {
			var m Message
			if err := dec.Decode(&m); err != nil {
				errc <- err
				return
			}
			in <- &m
		}
	}()

	// Receive messages from the out channel and encode to the client.
	go func() {
		enc := json.NewEncoder(c)
		for m := range out {
			if err := enc.Encode(m); err != nil {
				errc <- err
				return
			}
		}
	}()

	// Start and kill processes and handle errors.
	proc := make(map[string]*process)
	for {
		select {
		case m := <-in:
			switch m.Kind {
			case "run":
				proc[m.Id].Kill()
				lOut := limiter(in, out)
				proc[m.Id] = startProcess(m.Id, m.Body, lOut, m.Options)
			case "kill":
				proc[m.Id].Kill()
			}
		case err := <-errc:
			if err != io.EOF {
				// A encode or decode has failed; bail.
				log.Println(err)
			}
			// Shut down any running processes.
			for _, p := range proc {
				p.Kill()
			}
			return
		}
	}
}

// process represents a running process.
type process struct {
	id   string
	out  chan<- *Message
	done chan struct{} // closed when wait completes
	run  *exec.Cmd
	bin  string
}

// startProcess builds and runs the given program, sending its output
// and end event as Messages on the provided channel.
func startProcess(id, body string, out chan<- *Message, opt *Options) *process {
	p := &process{
		id:   id,
		out:  out,
		done: make(chan struct{}),
	}
	var err error
	if path, args := shebang(body); path != "" {
		err = p.startProcess(path, args, body)
	} else {
		err = p.start(body, opt)
	}
	if err != nil {
		p.end(err)
		return nil
	}
	go p.wait()
	return p
}

// Kill stops the process if it is running and waits for it to exit.
func (p *process) Kill() {
	if p == nil {
		return
	}
	p.run.Process.Kill()
	<-p.done // block until process exits
}

// shebang looks for a shebang ('#!') at the beginning of the passed string.
// If found, it returns the path and args after the shebang.
func shebang(body string) (path string, args []string) {
	body = strings.TrimSpace(body)
	if !strings.HasPrefix(body, "#!") {
		return "", nil
	}
	if i := strings.Index(body, "\n"); i >= 0 {
		body = body[:i]
	}
	fs := strings.Fields(body[2:])
	return fs[0], fs[1:]
}

// startProcess starts a given program given its path and passing the given body
// to the command standard input.
func (p *process) startProcess(path string, args []string, body string) error {
	cmd := &exec.Cmd{
		Path:   path,
		Args:   args,
		Stdin:  strings.NewReader(body),
		Stdout: &messageWriter{id: p.id, kind: "stdout", out: p.out},
		Stderr: &messageWriter{id: p.id, kind: "stderr", out: p.out},
	}
	if err := cmd.Start(); err != nil {
		return err
	}
	p.run = cmd
	return nil
}

// start builds and starts the given program, sending its output to p.out,
// and stores the running *exec.Cmd in the run field.
func (p *process) start(body string, opt *Options) error {
	// We "go build" and then exec the binary so that the
	// resultant *exec.Cmd is a handle to the user's program
	// (rather than the go tool process).
	// This makes Kill work.

	bin := filepath.Join(tmpdir, "compile"+strconv.Itoa(<-uniq))
	src := bin + ".go"
	if runtime.GOOS == "windows" {
		bin += ".exe"
	}

	// write body to x.go
	defer os.Remove(src)
	err := ioutil.WriteFile(src, []byte(body), 0666)
	if err != nil {
		return err
	}

	// build x.go, creating x
	p.bin = bin // to be removed by p.end
	dir, file := filepath.Split(src)
	args := []string{"go", "build", "-tags", "OMIT"}
	if opt != nil && opt.Race {
		p.out <- &Message{
			Id: p.id, Kind: "stderr",
			Body: "Running with race detector.\n",
		}
		args = append(args, "-race")
	}
	args = append(args, "-o", bin, file)
	cmd := p.cmd(dir, args...)
	cmd.Stdout = cmd.Stderr // send compiler output to stderr
	if err := cmd.Run(); err != nil {
		return err
	}

	// run x
	cmd = p.cmd("", bin)
	if opt != nil && opt.Race {
		cmd.Env = append(cmd.Env, "GOMAXPROCS=2")
	}
	if err := cmd.Start(); err != nil {
		// If we failed to exec, that might be because they built
		// a non-main package instead of an executable.
		// Check and report that.
		if name, err := packageName(body); err == nil && name != "main" {
			return errors.New(`executable programs must use "package main"`)
		}
		return err
	}
	p.run = cmd
	return nil
}

// wait waits for the running process to complete
// and sends its error state to the client.
func (p *process) wait() {
	p.end(p.run.Wait())
	close(p.done) // unblock waiting Kill calls
}

// end sends an "end" message to the client, containing the process id and the
// given error value. It also removes the binary.
func (p *process) end(err error) {
	if p.bin != "" {
		defer os.Remove(p.bin)
	}
	m := &Message{Id: p.id, Kind: "end"}
	if err != nil {
		m.Body = err.Error()
	}
	// Wait for any outstanding reads to finish (potential race here).
	time.AfterFunc(msgDelay, func() { p.out <- m })
}

// cmd builds an *exec.Cmd that writes its standard output and error to the
// process' output channel.
func (p *process) cmd(dir string, args ...string) *exec.Cmd {
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Dir = dir
	cmd.Env = Environ()
	cmd.Stdout = &messageWriter{id: p.id, kind: "stdout", out: p.out}
	cmd.Stderr = &messageWriter{id: p.id, kind: "stderr", out: p.out}
	return cmd
}

func packageName(body string) (string, error) {
	f, err := parser.ParseFile(token.NewFileSet(), "prog.go",
		strings.NewReader(body), parser.PackageClauseOnly)
	if err != nil {
		return "", err
	}
	return f.Name.String(), nil
}

// messageWriter is an io.Writer that converts all writes to Message sends on
// the out channel with the specified id and kind.
type messageWriter struct {
	id, kind string
	out      chan<- *Message

	mu   sync.Mutex
	buf  []byte
	send *time.Timer
}

func (w *messageWriter) Write(b []byte) (n int, err error) {
	// Buffer writes that occur in a short period to send as one Message.
	w.mu.Lock()
	w.buf = append(w.buf, b...)
	if w.send == nil {
		w.send = time.AfterFunc(msgDelay, w.sendNow)
	}
	w.mu.Unlock()
	return len(b), nil
}

func (w *messageWriter) sendNow() {
	w.mu.Lock()
	body := safeString(w.buf)
	w.buf, w.send = nil, nil
	w.mu.Unlock()
	w.out <- &Message{Id: w.id, Kind: w.kind, Body: body}
}

// safeString returns b as a valid UTF-8 string.
func safeString(b []byte) string {
	if utf8.Valid(b) {
		return string(b)
	}
	var buf bytes.Buffer
	for len(b) > 0 {
		r, size := utf8.DecodeRune(b)
		b = b[size:]
		buf.WriteRune(r)
	}
	return buf.String()
}

// limiter returns a channel that wraps dest. Messages sent to the channel are
// sent to dest. After msgLimit Messages have been passed on, a "kill" Message
// is sent to the kill channel, and only "end" messages are passed.
func limiter(kill chan<- *Message, dest chan<- *Message) chan<- *Message {
	ch := make(chan *Message)
	go func() {
		n := 0
		for m := range ch {
			switch {
			case n < msgLimit || m.Kind == "end":
				dest <- m
				if m.Kind == "end" {
					return
				}
			case n == msgLimit:
				// process produced too much output. Kill it.
				kill <- &Message{Id: m.Id, Kind: "kill"}
			}
			n++
		}
	}()
	return ch
}

var tmpdir string

func init() {
	// find real path to temporary directory
	var err error
	tmpdir, err = filepath.EvalSymlinks(os.TempDir())
	if err != nil {
		log.Fatal(err)
	}
}

var uniq = make(chan int) // a source of numbers for naming temporary files

func init() {
	go func() {
		for i := 0; ; i++ {
			uniq <- i
		}
	}()
}
<commit_msg>go.tools/playground: parse shebang correctly<commit_after>// Copyright 2012 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !appengine

// Package socket implements an WebSocket-based playground backend.
// Clients connect to a websocket handler and send run/kill commands, and
// the server sends the output and exit status of the running processes.
// Multiple clients running multiple processes may be served concurrently.
// The wire format is JSON and is described by the Message type.
//
// This will not run on App Engine as WebSockets are not supported there.
package socket

import (
	"bytes"
	"encoding/json"
	"errors"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"code.google.com/p/go.net/websocket"
)

// Handler implements a WebSocket handler for a client connection.
var Handler = websocket.Handler(socketHandler)

// Environ provides an environment when a binary, such as the go tool, is
// invoked.
var Environ func() []string = os.Environ

const (
	// The maximum number of messages to send per session (avoid flooding).
	msgLimit = 1000

	// Batch messages sent in this interval and send as a single message.
	msgDelay = 10 * time.Millisecond
)

// Message is the wire format for the websocket connection to the browser.
// It is used for both sending output messages and receiving commands, as
// distinguished by the Kind field.
type Message struct {
	Id      string // client-provided unique id for the process
	Kind    string // in: "run", "kill" out: "stdout", "stderr", "end"
	Body    string
	Options *Options `json:",omitempty"`
}

// Options specify additional message options.
type Options struct {
	Race bool // use -race flag when building code (for "run" only)
}

// socketHandler handles the websocket connection for a given present session.
// It handles transcoding Messages to and from JSON format, and starting
// and killing processes.
func socketHandler(c *websocket.Conn) {
	in, out := make(chan *Message), make(chan *Message)
	errc := make(chan error, 1)

	// Decode messages from client and send to the in channel.
	go func() {
		dec := json.NewDecoder(c)
		for {
			var m Message
			if err := dec.Decode(&m); err != nil {
				errc <- err
				return
			}
			in <- &m
		}
	}()

	// Receive messages from the out channel and encode to the client.
	go func() {
		enc := json.NewEncoder(c)
		for m := range out {
			if err := enc.Encode(m); err != nil {
				errc <- err
				return
			}
		}
	}()

	// Start and kill processes and handle errors.
	proc := make(map[string]*process)
	for {
		select {
		case m := <-in:
			switch m.Kind {
			case "run":
				proc[m.Id].Kill()
				lOut := limiter(in, out)
				proc[m.Id] = startProcess(m.Id, m.Body, lOut, m.Options)
			case "kill":
				proc[m.Id].Kill()
			}
		case err := <-errc:
			if err != io.EOF {
				// A encode or decode has failed; bail.
				log.Println(err)
			}
			// Shut down any running processes.
			for _, p := range proc {
				p.Kill()
			}
			return
		}
	}
}

// process represents a running process.
type process struct {
	id   string
	out  chan<- *Message
	done chan struct{} // closed when wait completes
	run  *exec.Cmd
	bin  string
}

// startProcess builds and runs the given program, sending its output
// and end event as Messages on the provided channel.
func startProcess(id, body string, out chan<- *Message, opt *Options) *process {
	p := &process{
		id:   id,
		out:  out,
		done: make(chan struct{}),
	}
	var err error
	if path, args := shebang(body); path != "" {
		err = p.startProcess(path, args, body)
	} else {
		err = p.start(body, opt)
	}
	if err != nil {
		p.end(err)
		return nil
	}
	go p.wait()
	return p
}

// Kill stops the process if it is running and waits for it to exit.
func (p *process) Kill() {
	if p == nil {
		return
	}
	p.run.Process.Kill()
	<-p.done // block until process exits
}

// shebang looks for a shebang ('#!') at the beginning of the passed string.
// If found, it returns the path and args after the shebang.
// args includes the command as args[0].
func shebang(body string) (path string, args []string) {
	body = strings.TrimSpace(body)
	if !strings.HasPrefix(body, "#!") {
		return "", nil
	}
	if i := strings.Index(body, "\n"); i >= 0 {
		body = body[:i]
	}
	fs := strings.Fields(body[2:])
	return fs[0], fs
}

// startProcess starts a given program given its path and passing the given body
// to the command standard input.
func (p *process) startProcess(path string, args []string, body string) error {
	cmd := &exec.Cmd{
		Path:   path,
		Args:   args,
		Stdin:  strings.NewReader(body),
		Stdout: &messageWriter{id: p.id, kind: "stdout", out: p.out},
		Stderr: &messageWriter{id: p.id, kind: "stderr", out: p.out},
	}
	if err := cmd.Start(); err != nil {
		return err
	}
	p.run = cmd
	return nil
}

// start builds and starts the given program, sending its output to p.out,
// and stores the running *exec.Cmd in the run field.
func (p *process) start(body string, opt *Options) error {
	// We "go build" and then exec the binary so that the
	// resultant *exec.Cmd is a handle to the user's program
	// (rather than the go tool process).
	// This makes Kill work.

	bin := filepath.Join(tmpdir, "compile"+strconv.Itoa(<-uniq))
	src := bin + ".go"
	if runtime.GOOS == "windows" {
		bin += ".exe"
	}

	// write body to x.go
	defer os.Remove(src)
	err := ioutil.WriteFile(src, []byte(body), 0666)
	if err != nil {
		return err
	}

	// build x.go, creating x
	p.bin = bin // to be removed by p.end
	dir, file := filepath.Split(src)
	args := []string{"go", "build", "-tags", "OMIT"}
	if opt != nil && opt.Race {
		p.out <- &Message{
			Id: p.id, Kind: "stderr",
			Body: "Running with race detector.\n",
		}
		args = append(args, "-race")
	}
	args = append(args, "-o", bin, file)
	cmd := p.cmd(dir, args...)
	cmd.Stdout = cmd.Stderr // send compiler output to stderr
	if err := cmd.Run(); err != nil {
		return err
	}

	// run x
	cmd = p.cmd("", bin)
	if opt != nil && opt.Race {
		cmd.Env = append(cmd.Env, "GOMAXPROCS=2")
	}
	if err := cmd.Start(); err != nil {
		// If we failed to exec, that might be because they built
		// a non-main package instead of an executable.
		// Check and report that.
		if name, err := packageName(body); err == nil && name != "main" {
			return errors.New(`executable programs must use "package main"`)
		}
		return err
	}
	p.run = cmd
	return nil
}

// wait waits for the running process to complete
// and sends its error state to the client.
func (p *process) wait() {
	p.end(p.run.Wait())
	close(p.done) // unblock waiting Kill calls
}

// end sends an "end" message to the client, containing the process id and the
// given error value. It also removes the binary.
func (p *process) end(err error) {
	if p.bin != "" {
		defer os.Remove(p.bin)
	}
	m := &Message{Id: p.id, Kind: "end"}
	if err != nil {
		m.Body = err.Error()
	}
	// Wait for any outstanding reads to finish (potential race here).
	time.AfterFunc(msgDelay, func() { p.out <- m })
}

// cmd builds an *exec.Cmd that writes its standard output and error to the
// process' output channel.
func (p *process) cmd(dir string, args ...string) *exec.Cmd {
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Dir = dir
	cmd.Env = Environ()
	cmd.Stdout = &messageWriter{id: p.id, kind: "stdout", out: p.out}
	cmd.Stderr = &messageWriter{id: p.id, kind: "stderr", out: p.out}
	return cmd
}

func packageName(body string) (string, error) {
	f, err := parser.ParseFile(token.NewFileSet(), "prog.go",
		strings.NewReader(body), parser.PackageClauseOnly)
	if err != nil {
		return "", err
	}
	return f.Name.String(), nil
}

// messageWriter is an io.Writer that converts all writes to Message sends on
// the out channel with the specified id and kind.
type messageWriter struct {
	id, kind string
	out      chan<- *Message

	mu   sync.Mutex
	buf  []byte
	send *time.Timer
}

func (w *messageWriter) Write(b []byte) (n int, err error) {
	// Buffer writes that occur in a short period to send as one Message.
	w.mu.Lock()
	w.buf = append(w.buf, b...)
	if w.send == nil {
		w.send = time.AfterFunc(msgDelay, w.sendNow)
	}
	w.mu.Unlock()
	return len(b), nil
}

func (w *messageWriter) sendNow() {
	w.mu.Lock()
	body := safeString(w.buf)
	w.buf, w.send = nil, nil
	w.mu.Unlock()
	w.out <- &Message{Id: w.id, Kind: w.kind, Body: body}
}

// safeString returns b as a valid UTF-8 string.
func safeString(b []byte) string {
	if utf8.Valid(b) {
		return string(b)
	}
	var buf bytes.Buffer
	for len(b) > 0 {
		r, size := utf8.DecodeRune(b)
		b = b[size:]
		buf.WriteRune(r)
	}
	return buf.String()
}

// limiter returns a channel that wraps dest. Messages sent to the channel are
// sent to dest. After msgLimit Messages have been passed on, a "kill" Message
// is sent to the kill channel, and only "end" messages are passed.
func limiter(kill chan<- *Message, dest chan<- *Message) chan<- *Message {
	ch := make(chan *Message)
	go func() {
		n := 0
		for m := range ch {
			switch {
			case n < msgLimit || m.Kind == "end":
				dest <- m
				if m.Kind == "end" {
					return
				}
			case n == msgLimit:
				// process produced too much output. Kill it.
				kill <- &Message{Id: m.Id, Kind: "kill"}
			}
			n++
		}
	}()
	return ch
}

var tmpdir string

func init() {
	// find real path to temporary directory
	var err error
	tmpdir, err = filepath.EvalSymlinks(os.TempDir())
	if err != nil {
		log.Fatal(err)
	}
}

var uniq = make(chan int) // a source of numbers for naming temporary files

func init() {
	go func() {
		for i := 0; ; i++ {
			uniq <- i
		}
	}()
}
<|endoftext|>


<commit_before>package http

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gansoi/gansoi/plugins"
)

func init() {
	plugins.RegisterAgent("http", HTTP{})
}

const (
	redirectsToFollow = 10

	httpsScheme = "https"
	httpScheme  = "http"
)

// HTTP will request a ressource from a HTTP server.
type HTTP struct {
	URL            string `json:"url" description:"The URL to request"`
	FollowRedirect bool   `json:"followRedirect" description:"Follow 30x redirects"`
	Insecure       bool   `json:"insecure" description:"Ignore SSL errors"`
	IncludeBody    bool   `json:"includeBody" description:"Include body in results"`
}

func getHostPort(URL *url.URL) (string, string) {
	if !strings.ContainsRune(URL.Host, ':') {
		switch URL.Scheme {
		case httpScheme:
			return URL.Host, "80"
		case httpsScheme:
			return URL.Host, "443"
		}
	}

	host, port, _ := net.SplitHostPort(URL.Host)

	return host, port
}

func camelCaseHeader(key string) string {
	fields := strings.FieldsFunc(key, func(r rune) bool {
		return !plugins.ValidateResultKeyRune(r)
	})

	result := "Header"

	for _, f := range fields {
		result += strings.Title(f)
	}

	return result
}

// Check implements plugins.Agent.
func (h *HTTP) Check(result plugins.AgentResult) error {
	URL, err := url.Parse(h.URL)
	if err != nil {
		return err
	}

	for try := 0; try < redirectsToFollow; try++ {
		if !(URL.Scheme == "http" || URL.Scheme == "https") {
			return http.ErrNotSupported
		}

		host, port := getHostPort(URL)

		t0 := time.Now()
		raddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(host, port))
		if err != nil {
			return err
		}

		t1 := time.Now()
		var conn net.Conn
		conn, err = net.DialTCP("tcp", nil, raddr)
		if err != nil {
			return err
		}
		defer conn.Close()

		t2 := time.Now()
		if URL.Scheme == httpsScheme {
			c := tls.Client(conn, &tls.Config{
				ServerName:         host,
				InsecureSkipVerify: h.Insecure,
			})

			err = c.Handshake()
			if err != nil {
				return err
			}

			state := c.ConnectionState()

			if len(state.PeerCertificates) > 0 {
				cert := state.PeerCertificates[0]
				notAfter := cert.NotAfter
				result.AddValue("SSLValidDays", notAfter.Sub(time.Now()).Hours()/24.0)
				result.AddValue("SSLCommonName", cert.Subject.CommonName)
			}

			conn = c
		}

		req, err := http.NewRequest("GET", URL.String(), nil)
		if err != nil {
			return err
		}

		t3 := time.Now()
		err = req.Write(conn)
		if err != nil {
			return err
		}

		t4 := time.Now()
		resp, err := http.ReadResponse(bufio.NewReader(conn), req)
		if err != nil {
			return err
		}

		if h.FollowRedirect && (resp.StatusCode == 301 || resp.StatusCode == 302) {
			URL, err = resp.Location()
			if err != nil {
				return err
			}

			result.AddValue(fmt.Sprintf("Redirect%d", try), URL.String())

			continue
		}

		t5 := time.Now()
		b := make([]byte, 1024)
		n, _ := resp.Body.Read(b)
		resp.Body.Close()
		t6 := time.Now()

		if h.IncludeBody {
			result.AddValue("Body", string(b[:n]))
		}

		result.AddValue("StatusCode", resp.StatusCode)
		result.AddValue("TimeDNS", ms(t1.Sub(t0)))
		result.AddValue("TimeConnect", ms(t2.Sub(t1)))
		result.AddValue("TimeTLS", ms(t3.Sub(t2)))
		result.AddValue("TimeRequest", ms(t4.Sub(t3)))
		result.AddValue("TimeReadHeaders", ms(t5.Sub(t4)))
		result.AddValue("TimeReadBody", ms(t6.Sub(t5)))
		result.AddValue("TimeAccumulated", ms(t6.Sub(t0)))

		for k, v := range resp.Header {
			result.AddValue(camelCaseHeader(k), strings.Join(v, " "))
		}

		break
	}

	return nil
}

// ms will convert a time.Duration to milliseconds.
func ms(d time.Duration) int64 {
	return (d.Nanoseconds() + 1000000/2) / 1000000
}
<commit_msg>Added an arbitrary host setting to the http check.<commit_after>package http

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gansoi/gansoi/plugins"
)

func init() {
	plugins.RegisterAgent("http", HTTP{})
}

const (
	redirectsToFollow = 10

	httpsScheme = "https"
	httpScheme  = "http"
)

// HTTP will request a ressource from a HTTP server.
type HTTP struct {
	URL            string `json:"url" description:"The URL to request"`
	FollowRedirect bool   `json:"followRedirect" description:"Follow 30x redirects"`
	Insecure       bool   `json:"insecure" description:"Ignore SSL errors"`
	IncludeBody    bool   `json:"includeBody" description:"Include body in results"`
	Host           string `json:"host" description:"Host to contact (leave empty to use host derived from URL)"`
}

func getHostPort(URL *url.URL) (string, string) {
	if !strings.ContainsRune(URL.Host, ':') {
		switch URL.Scheme {
		case httpScheme:
			return URL.Host, "80"
		case httpsScheme:
			return URL.Host, "443"
		}
	}

	host, port, _ := net.SplitHostPort(URL.Host)

	return host, port
}

func camelCaseHeader(key string) string {
	fields := strings.FieldsFunc(key, func(r rune) bool {
		return !plugins.ValidateResultKeyRune(r)
	})

	result := "Header"

	for _, f := range fields {
		result += strings.Title(f)
	}

	return result
}

// Check implements plugins.Agent.
func (h *HTTP) Check(result plugins.AgentResult) error {
	URL, err := url.Parse(h.URL)
	if err != nil {
		return err
	}

	for try := 0; try < redirectsToFollow; try++ {
		if !(URL.Scheme == "http" || URL.Scheme == "https") {
			return http.ErrNotSupported
		}

		host, port := getHostPort(URL)

		if h.Host != "" {
			host = h.Host
		}

		t0 := time.Now()
		raddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(host, port))
		if err != nil {
			return err
		}

		t1 := time.Now()
		var conn net.Conn
		conn, err = net.DialTCP("tcp", nil, raddr)
		if err != nil {
			return err
		}
		defer conn.Close()

		t2 := time.Now()
		if URL.Scheme == httpsScheme {
			c := tls.Client(conn, &tls.Config{
				ServerName:         host,
				InsecureSkipVerify: h.Insecure,
			})

			err = c.Handshake()
			if err != nil {
				return err
			}

			state := c.ConnectionState()

			if len(state.PeerCertificates) > 0 {
				cert := state.PeerCertificates[0]
				notAfter := cert.NotAfter
				result.AddValue("SSLValidDays", notAfter.Sub(time.Now()).Hours()/24.0)
				result.AddValue("SSLCommonName", cert.Subject.CommonName)
			}

			conn = c
		}

		req, err := http.NewRequest("GET", URL.String(), nil)
		if err != nil {
			return err
		}

		t3 := time.Now()
		err = req.Write(conn)
		if err != nil {
			return err
		}

		t4 := time.Now()
		resp, err := http.ReadResponse(bufio.NewReader(conn), req)
		if err != nil {
			return err
		}

		if h.FollowRedirect && (resp.StatusCode == 301 || resp.StatusCode == 302) {
			URL, err = resp.Location()
			if err != nil {
				return err
			}

			result.AddValue(fmt.Sprintf("Redirect%d", try), URL.String())

			continue
		}

		t5 := time.Now()
		b := make([]byte, 1024)
		n, _ := resp.Body.Read(b)
		resp.Body.Close()
		t6 := time.Now()

		if h.IncludeBody {
			result.AddValue("Body", string(b[:n]))
		}

		result.AddValue("StatusCode", resp.StatusCode)
		result.AddValue("TimeDNS", ms(t1.Sub(t0)))
		result.AddValue("TimeConnect", ms(t2.Sub(t1)))
		result.AddValue("TimeTLS", ms(t3.Sub(t2)))
		result.AddValue("TimeRequest", ms(t4.Sub(t3)))
		result.AddValue("TimeReadHeaders", ms(t5.Sub(t4)))
		result.AddValue("TimeReadBody", ms(t6.Sub(t5)))
		result.AddValue("TimeAccumulated", ms(t6.Sub(t0)))

		for k, v := range resp.Header {
			result.AddValue(camelCaseHeader(k), strings.Join(v, " "))
		}

		break
	}

	return nil
}

// ms will convert a time.Duration to milliseconds.
func ms(d time.Duration) int64 {
	return (d.Nanoseconds() + 1000000/2) / 1000000
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"net"
)

var Config struct {
	localAddr  string
	serverAddr string
}

func proxy(clientConn net.Conn) {
	defer clientConn.Close()

	// create socks5 server
	s := NewSocks5Server(clientConn)
	err := s.MethodSelect()
	if err != nil {
		return
	}
	dstAddr, err := s.ReceiveDstAddr()
	if err != nil {
		return
	}
	log.Printf("proxy_request: [%s] => [%s]\n",
		clientConn.RemoteAddr(), dstAddr)

	// create tl client
	serverConn, err := net.Dial("tcp4", Config.serverAddr)
	if err != nil {
		log.Printf("%v", err)
		return
	}
	c := NewTlClient(serverConn)

	err = c.Connect(dstAddr)
	if err != nil {
		return
	}

	err = s.NotifyConnectSuccess()
	if err != nil {
		return
	}

	go io.Copy(c, clientConn)
	io.Copy(clientConn, c)
}

func handleProxy() {
	l, err := net.Listen("tcp4", Config.localAddr)
	if err != nil {
		log.Printf("%v", err)
		return
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			log.Printf("%v", err)
		}
		go proxy(conn)
	}
}

func main() {
	Config.localAddr = "0.0.0.0:8000"
	Config.serverAddr = "127.0.0.1:8001"

	handleProxy()
}
<commit_msg>fix connection leak bug<commit_after>package main

import (
	"io"
	"log"
	"net"
)

var Config struct {
	localAddr  string
	serverAddr string
}

func proxy(clientConn net.Conn) {
	defer clientConn.Close()

	// create socks5 server
	s := NewSocks5Server(clientConn)
	err := s.MethodSelect()
	if err != nil {
		return
	}
	dstAddr, err := s.ReceiveDstAddr()
	if err != nil {
		return
	}
	log.Printf("proxy_request: [%s] => [%s]\n",
		clientConn.RemoteAddr(), dstAddr)

	// create tl client
	serverConn, err := net.Dial("tcp4", Config.serverAddr)
	if err != nil {
		log.Printf("%v", err)
		return
	}
	defer serverConn.Close()

	c := NewTlClient(serverConn)

	err = c.Connect(dstAddr)
	if err != nil {
		return
	}

	err = s.NotifyConnectSuccess()
	if err != nil {
		return
	}

	go io.Copy(c, clientConn)
	io.Copy(clientConn, c)
}

func handleProxy() {
	l, err := net.Listen("tcp4", Config.localAddr)
	if err != nil {
		log.Printf("%v", err)
		return
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			log.Printf("%v", err)
		}
		go proxy(conn)
	}
}

func main() {
	Config.localAddr = "0.0.0.0:8000"
	Config.serverAddr = "127.0.0.1:8001"

	handleProxy()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"testing"

	"github.com/stretchr/testify/assert"
)

var testClient *Client

func setup() {
	out, err := exec.Command("createdb", "-U", "postgres", "-h", "localhost", "booktown").CombinedOutput()

	if err != nil {
		fmt.Print("Database creation failed:", string(out))
		os.Exit(1)
	}

	out, err = exec.Command("psql", "-U", "postgres", "-h", "localhost", "-f", "./sql/booktown.sql", "booktown").CombinedOutput()

	if err != nil {
		fmt.Print("Database import failed:", string(out))
		os.Exit(1)
	}
}

func setupClient() {
	testClient, _ = NewClientFromUrl("postgres://postgres@localhost/booktown?sslmode=disable")
}

func teardownClient() {
	if testClient != nil {
		testClient.db.Close()
	}
}

func teardown() {
	out, err := exec.Command("dropdb", "-U", "postgres", "-h", "localhost", "booktown").CombinedOutput()

	if err != nil {
		fmt.Print(string(out))
	}
}

func test_NewClientFromUrl(t *testing.T) {
	url := "postgres://postgres@localhost/booktown?sslmode=disable"
	client, err := NewClientFromUrl(url)

	if err != nil {
		defer client.db.Close()
	}

	assert.Equal(t, nil, err)
	assert.Equal(t, url, client.connectionString)
}

func test_Test(t *testing.T) {
	assert.Equal(t, nil, testClient.Test())
}

func test_Info(t *testing.T) {
	res, err := testClient.Info()

	assert.Equal(t, nil, err)
	assert.NotEqual(t, nil, res)
}

func test_Databases(t *testing.T) {
	res, err := testClient.Databases()

	assert.Equal(t, nil, err)
	assert.Contains(t, res, "booktown")
	assert.Contains(t, res, "postgres")
}

func test_Tables(t *testing.T) {
	res, err := testClient.Tables()

	expected := []string{
		"alternate_stock",
		"authors",
		"book_backup",
		"book_queue",
		"books",
		"customers",
		"daily_inventory",
		"distinguished_authors",
		"editions",
		"employees",
		"favorite_authors",
		"favorite_books",
		"money_example",
		"my_list",
		"numeric_values",
		"publishers",
		"recent_shipments",
		"schedules",
		"shipments",
		"states",
		"stock",
		"stock_backup",
		"stock_view",
		"subjects",
		"text_sorting",
	}

	assert.Equal(t, nil, err)
	assert.Equal(t, expected, res)
}

func test_Table(t *testing.T) {
	res, err := testClient.Table("books")

	columns := []string{
		"column_name",
		"data_type",
		"is_nullable",
		"character_maximum_length",
		"character_set_catalog",
		"column_default",
	}

	assert.Equal(t, nil, err)
	assert.Equal(t, columns, res.Columns)
	assert.Equal(t, 4, len(res.Rows))
}

func test_TableRows(t *testing.T) {
	res, err := testClient.TableRows("books", RowsOptions{})

	assert.Equal(t, nil, err)
	assert.Equal(t, 4, len(res.Columns))
	assert.Equal(t, 15, len(res.Rows))
}

func test_TableInfo(t *testing.T) {
	res, err := testClient.TableInfo("books")

	assert.Equal(t, nil, err)
	assert.Equal(t, 4, len(res.Columns))
	assert.Equal(t, 1, len(res.Rows))
}

func test_TableIndexes(t *testing.T) {
	res, err := testClient.TableIndexes("books")

	assert.Equal(t, nil, err)
	assert.Equal(t, 2, len(res.Columns))
	assert.Equal(t, 2, len(res.Rows))
}

func test_Query(t *testing.T) {
	res, err := testClient.Query("SELECT * FROM books")

	assert.Equal(t, nil, err)
	assert.Equal(t, 4, len(res.Columns))
	assert.Equal(t, 15, len(res.Rows))
}

func test_QueryError(t *testing.T) {
	res, err := testClient.Query("SELCT * FROM books")

	assert.NotEqual(t, nil, err)
	assert.Equal(t, "pq: syntax error at or near \"SELCT\"", err.Error())
	assert.Equal(t, true, res == nil)
}

func test_QueryInvalidTable(t *testing.T) {
	res, err := testClient.Query("SELECT * FROM books2")

	assert.NotEqual(t, nil, err)
	assert.Equal(t, "pq: relation \"books2\" does not exist", err.Error())
	assert.Equal(t, true, res == nil)
}

func test_ResultCsv(t *testing.T) {
	res, _ := testClient.Query("SELECT * FROM books ORDER BY id ASC LIMIT 1")
	csv := res.CSV()

	expected := "id,title,author_id,subject_id\n156,The Tell-Tale Heart,115,9\n"

	assert.Equal(t, expected, string(csv))
}

func TestAll(t *testing.T) {
	teardown()
	setup()
	setupClient()

	test_NewClientFromUrl(t)
	test_Test(t)
	test_Info(t)
	test_Databases(t)
	test_Tables(t)
	test_Table(t)
	test_TableRows(t)
	test_TableInfo(t)
	test_TableIndexes(t)
	test_Query(t)
	test_QueryError(t)
	test_QueryInvalidTable(t)
	test_ResultCsv(t)

	teardownClient()
	teardown()
}
<commit_msg>Test history adding<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"testing"

	"github.com/stretchr/testify/assert"
)

var testClient *Client

func setup() {
	out, err := exec.Command("createdb", "-U", "postgres", "-h", "localhost", "booktown").CombinedOutput()

	if err != nil {
		fmt.Print("Database creation failed:", string(out))
		os.Exit(1)
	}

	out, err = exec.Command("psql", "-U", "postgres", "-h", "localhost", "-f", "./sql/booktown.sql", "booktown").CombinedOutput()

	if err != nil {
		fmt.Print("Database import failed:", string(out))
		os.Exit(1)
	}
}

func setupClient() {
	testClient, _ = NewClientFromUrl("postgres://postgres@localhost/booktown?sslmode=disable")
}

func teardownClient() {
	if testClient != nil {
		testClient.db.Close()
	}
}

func teardown() {
	out, err := exec.Command("dropdb", "-U", "postgres", "-h", "localhost", "booktown").CombinedOutput()

	if err != nil {
		fmt.Print(string(out))
	}
}

func test_NewClientFromUrl(t *testing.T) {
	url := "postgres://postgres@localhost/booktown?sslmode=disable"
	client, err := NewClientFromUrl(url)

	if err != nil {
		defer client.db.Close()
	}

	assert.Equal(t, nil, err)
	assert.Equal(t, url, client.connectionString)
}

func test_Test(t *testing.T) {
	assert.Equal(t, nil, testClient.Test())
}

func test_Info(t *testing.T) {
	res, err := testClient.Info()

	assert.Equal(t, nil, err)
	assert.NotEqual(t, nil, res)
}

func test_Databases(t *testing.T) {
	res, err := testClient.Databases()

	assert.Equal(t, nil, err)
	assert.Contains(t, res, "booktown")
	assert.Contains(t, res, "postgres")
}

func test_Tables(t *testing.T) {
	res, err := testClient.Tables()

	expected := []string{
		"alternate_stock",
		"authors",
		"book_backup",
		"book_queue",
		"books",
		"customers",
		"daily_inventory",
		"distinguished_authors",
		"editions",
		"employees",
		"favorite_authors",
		"favorite_books",
		"money_example",
		"my_list",
		"numeric_values",
		"publishers",
		"recent_shipments",
		"schedules",
		"shipments",
		"states",
		"stock",
		"stock_backup",
		"stock_view",
		"subjects",
		"text_sorting",
	}

	assert.Equal(t, nil, err)
	assert.Equal(t, expected, res)
}

func test_Table(t *testing.T) {
	res, err := testClient.Table("books")

	columns := []string{
		"column_name",
		"data_type",
		"is_nullable",
		"character_maximum_length",
		"character_set_catalog",
		"column_default",
	}

	assert.Equal(t, nil, err)
	assert.Equal(t, columns, res.Columns)
	assert.Equal(t, 4, len(res.Rows))
}

func test_TableRows(t *testing.T) {
	res, err := testClient.TableRows("books", RowsOptions{})

	assert.Equal(t, nil, err)
	assert.Equal(t, 4, len(res.Columns))
	assert.Equal(t, 15, len(res.Rows))
}

func test_TableInfo(t *testing.T) {
	res, err := testClient.TableInfo("books")

	assert.Equal(t, nil, err)
	assert.Equal(t, 4, len(res.Columns))
	assert.Equal(t, 1, len(res.Rows))
}

func test_TableIndexes(t *testing.T) {
	res, err := testClient.TableIndexes("books")

	assert.Equal(t, nil, err)
	assert.Equal(t, 2, len(res.Columns))
	assert.Equal(t, 2, len(res.Rows))
}

func test_Query(t *testing.T) {
	res, err := testClient.Query("SELECT * FROM books")

	assert.Equal(t, nil, err)
	assert.Equal(t, 4, len(res.Columns))
	assert.Equal(t, 15, len(res.Rows))
}

func test_QueryError(t *testing.T) {
	res, err := testClient.Query("SELCT * FROM books")

	assert.NotEqual(t, nil, err)
	assert.Equal(t, "pq: syntax error at or near \"SELCT\"", err.Error())
	assert.Equal(t, true, res == nil)
}

func test_QueryInvalidTable(t *testing.T) {
	res, err := testClient.Query("SELECT * FROM books2")

	assert.NotEqual(t, nil, err)
	assert.Equal(t, "pq: relation \"books2\" does not exist", err.Error())
	assert.Equal(t, true, res == nil)
}

func test_ResultCsv(t *testing.T) {
	res, _ := testClient.Query("SELECT * FROM books ORDER BY id ASC LIMIT 1")
	csv := res.CSV()

	expected := "id,title,author_id,subject_id\n156,The Tell-Tale Heart,115,9\n"

	assert.Equal(t, expected, string(csv))
}

func test_History(t *testing.T) {
	_, err := testClient.Query("SELECT * FROM books")
	query := testClient.history[len(testClient.history)-1].Query

	assert.Equal(t, nil, err)
	assert.Equal(t, "SELECT * FROM books", query)
}

func test_HistoryError(t *testing.T) {
	_, err := testClient.Query("SELECT * FROM books123")
	query := testClient.history[len(testClient.history)-1].Query

	assert.NotEqual(t, nil, err)
	assert.NotEqual(t, "SELECT * FROM books123", query)
}

func TestAll(t *testing.T) {
	teardown()
	setup()
	setupClient()

	test_NewClientFromUrl(t)
	test_Test(t)
	test_Info(t)
	test_Databases(t)
	test_Tables(t)
	test_Table(t)
	test_TableRows(t)
	test_TableInfo(t)
	test_TableIndexes(t)
	test_Query(t)
	test_QueryError(t)
	test_QueryInvalidTable(t)
	test_ResultCsv(t)
	test_History(t)
	test_HistoryError(t)

	teardownClient()
	teardown()
}
<|endoftext|>


<commit_before>// Package redisbroker implements a juggler broker using Redis
// as backend. RPC calls and results are stored in Redis lists
// and queried via the BRPOP command, while pub-sub events
// are handled using Redis' built-in pub-sub support.
//
// Call timeouts are handled by an expiring key associated
// with each call request, and in a similar way for results.
// Keys are named in such a way that the call request list
// and associated expiring keys are in the same hash slot,
// and the same is true for results and their expiring key,
// so that using a redis cluster is supported. The call
// requests are hashed on the call URI, and the results
// are hashed on the calling connection's UUID.
//
// If an RPC URI is much more sollicitated than others,
// it can be spread over multiple URIs using
// "RPC_URI_%d" where %d is e.g. a number from 1 to 100.
// Clients that need to call this function can use a random
// over that range to spread the load over different cluster
// nodes.
//
package redisbroker

import (
	"encoding/json"
	"expvar"
	"fmt"
	"log"
	"time"

	"github.com/PuerkitoBio/juggler/broker"
	"github.com/PuerkitoBio/juggler/message"
	"github.com/PuerkitoBio/redisc"
	"github.com/garyburd/redigo/redis"
	"github.com/pborman/uuid"
)

var (
	// static check that *Broker implements all the broker interfaces
	_ broker.CallerBroker = (*Broker)(nil)
	_ broker.CalleeBroker = (*Broker)(nil)
	_ broker.PubSubBroker = (*Broker)(nil)
)

// Pool defines the methods required for a redis pool that provides
// a method to get a connection and to release the pool's resources.
type Pool interface {
	// Get returns a redis connection.
	Get() redis.Conn

	// Close releases the resources used by the pool.
	Close() error
}

// Broker is a broker that provides the methods to
// interact with Redis using the juggler protocol.
type Broker struct {
	// Pool is the redis pool or redisc cluster to use to get
	// short-lived connections.
	Pool Pool

	// Dial is the function to call to get a non-pooled, long-lived
	// redis connection. Typically, it can be set to redis.Pool.Dial
	// or redisc.Cluster.Dial.
	Dial func() (redis.Conn, error)

	// BlockingTimeout is the time to wait for a value on calls to
	// BRPOP before trying again. The default of 0 means no timeout.
	BlockingTimeout time.Duration

	// LogFunc is the logging function to use. If nil, log.Printf
	// is used. It can be set to juggler.DiscardLog to disable logging.
	LogFunc func(string, ...interface{})

	// CallCap is the capacity of the CALL queue per URI. If it is
	// exceeded for a given URI, subsequent Broker.Call calls for that
	// URI will fail with an error.
	CallCap int

	// ResultCap is the capacity of the RES queue per connection UUID.
	// If it is exceeded for a given connection, Broker.Result calls
	// for that connection will fail with an error.
	ResultCap int

	// Vars can be set to an *expvar.Map to collect metrics about the
	// broker. It should be set before starting to make calls with the
	// broker.
	Vars *expvar.Map
}

// script to store the call request or call result along with
// its expiration information.
var callOrResScript = redis.NewScript(2, `
	redis.call("SET", KEYS[1], ARGV[1], "PX", tonumber(ARGV[1]))
	local res = redis.call("LPUSH", KEYS[2], ARGV[2])
	local limit = tonumber(ARGV[3])
	if res > limit and limit > 0 then
		local diff = res - limit
		redis.call("LTRIM", KEYS[2], diff, limit + diff)
		return redis.error_reply("list capacity exceeded")
	end
	return res
`)

const (
	// redis cluster-compliant keys, so that both keys are in the same slot
	callKey        = "juggler:calls:{%s}"            // 1: URI
	callTimeoutKey = "juggler:calls:timeout:{%s}:%s" // 1: URI, 2: mUUID

	// redis cluster-compliant keys, so that both keys are in the same slot
	resKey        = "juggler:results:{%s}"            // 1: cUUID
	resTimeoutKey = "juggler:results:timeout:{%s}:%s" // 1: cUUID, 2: mUUID
)

// Call registers a call request in the broker.
func (b *Broker) Call(cp *message.CallPayload, timeout time.Duration) error {
	k1 := fmt.Sprintf(callTimeoutKey, cp.URI, cp.MsgUUID)
	k2 := fmt.Sprintf(callKey, cp.URI)
	return registerCallOrRes(b.Pool, cp, timeout, b.CallCap, k1, k2)
}

// Result registers a call result in the broker.
func (b *Broker) Result(rp *message.ResPayload, timeout time.Duration) error {
	k1 := fmt.Sprintf(resTimeoutKey, rp.ConnUUID, rp.MsgUUID)
	k2 := fmt.Sprintf(resKey, rp.ConnUUID)
	return registerCallOrRes(b.Pool, rp, timeout, b.ResultCap, k1, k2)
}

func registerCallOrRes(pool Pool, pld interface{}, timeout time.Duration, cap int, k1, k2 string) error {
	p, err := json.Marshal(pld)
	if err != nil {
		return err
	}

	rc := pool.Get()
	defer rc.Close()

	// turn it into a cluster-aware RetryConn if running in a cluster
	rc = clusterifyConn(rc, k1, k2)

	to := int(timeout / time.Millisecond)
	if to == 0 {
		to = int(broker.DefaultCallTimeout / time.Millisecond)
	}

	_, err = callOrResScript.Do(rc,
		k1,  // key[1] : the SET key with expiration
		k2,  // key[2] : the LIST key
		to,  // argv[1] : the timeout in milliseconds
		p,   // argv[2] : the call payload
		cap, // argv[3] : the LIST capacity
	)
	return err
}

// Publish publishes an event to a channel.
func (b *Broker) Publish(channel string, pp *message.PubPayload) error {
	p, err := json.Marshal(pp)
	if err != nil {
		return err
	}

	rc := b.Pool.Get()
	defer rc.Close()

	// force selection of a random node (otherwise it would use
	// the node of the hash of the channel - which may hit the
	// same node over and over again if there are few channels).
	if bc, ok := rc.(binder); ok {
		// ignore the error, if it fails, use the connection as-is.
		// Bind without a key selects a random node.
		bc.Bind()
	}
	_, err = rc.Do("PUBLISH", channel, p)
	return err
}

// NewPubSubConn returns a new pub-sub connection that can be used
// to subscribe to and unsubscribe from channels, and to process
// incoming events.
func (b *Broker) NewPubSubConn() (broker.PubSubConn, error) {
	rc, err := b.Dial()
	if err != nil {
		return nil, err
	}
	return &pubSubConn{
		psc:   redis.PubSubConn{Conn: rc},
		logFn: b.LogFunc,
		vars:  b.Vars,
	}, nil
}

// NewCallsConn returns a new calls connection that can be used
// to process the call requests for the specified URIs.
func (b *Broker) NewCallsConn(uris ...string) (broker.CallsConn, error) {
	rc, err := b.Dial()
	if err != nil {
		return nil, err
	}
	return &callsConn{
		c:       rc,
		pool:    b.Pool,
		uris:    uris,
		vars:    b.Vars,
		timeout: b.BlockingTimeout,
		logFn:   b.LogFunc,
	}, nil
}

// NewResultsConn returns a new results connection that can be used
// to process the call results for the specified connection UUID.
func (b *Broker) NewResultsConn(connUUID uuid.UUID) (broker.ResultsConn, error) {
	rc, err := b.Dial()
	if err != nil {
		return nil, err
	}
	return &resultsConn{
		c:        rc,
		pool:     b.Pool,
		connUUID: connUUID,
		vars:     b.Vars,
		timeout:  b.BlockingTimeout,
		logFn:    b.LogFunc,
	}, nil
}

const (
	// TODO : maybe make that customizable, not super critical.
	clusterConnMaxAttempts   = 4
	clusterConnTryAgainDelay = 100 * time.Millisecond
)

type binder interface {
	Bind(...string) error
}

func clusterifyConn(rc redis.Conn, keys ...string) redis.Conn {
	// if it implements Bind, call it and make it a RetryConn so
	// that it follows redirections in a cluster.
	if bc, ok := rc.(binder); ok {
		// if Bind fails, go on with the call as usual, but if it
		// succeeds, try to turn it into a RetryConn.
		if err := bc.Bind(keys...); err == nil {
			retry, err := redisc.RetryConn(rc, clusterConnMaxAttempts, clusterConnTryAgainDelay)
			// again, if it fails, ignore and go on with the normal conn,
			// but if it succeds, replace the conn with this one.
			if err == nil {
				rc = retry
			}
		}
	}
	return rc
}

func logf(fn func(string, ...interface{}), f string, args ...interface{}) {
	if fn != nil {
		fn(f, args...)
	} else {
		log.Printf(f, args...)
	}
}
<commit_msg>juggler/broker/redisbroker: add the DiscardLog var<commit_after>// Package redisbroker implements a juggler broker using Redis
// as backend. RPC calls and results are stored in Redis lists
// and queried via the BRPOP command, while pub-sub events
// are handled using Redis' built-in pub-sub support.
//
// Call timeouts are handled by an expiring key associated
// with each call request, and in a similar way for results.
// Keys are named in such a way that the call request list
// and associated expiring keys are in the same hash slot,
// and the same is true for results and their expiring key,
// so that using a redis cluster is supported. The call
// requests are hashed on the call URI, and the results
// are hashed on the calling connection's UUID.
//
// If an RPC URI is much more sollicitated than others,
// it can be spread over multiple URIs using
// "RPC_URI_%d" where %d is e.g. a number from 1 to 100.
// Clients that need to call this function can use a random
// over that range to spread the load over different cluster
// nodes.
//
package redisbroker

import (
	"encoding/json"
	"expvar"
	"fmt"
	"log"
	"time"

	"github.com/PuerkitoBio/juggler/broker"
	"github.com/PuerkitoBio/juggler/message"
	"github.com/PuerkitoBio/redisc"
	"github.com/garyburd/redigo/redis"
	"github.com/pborman/uuid"
)

var (
	// static check that *Broker implements all the broker interfaces
	_ broker.CallerBroker = (*Broker)(nil)
	_ broker.CalleeBroker = (*Broker)(nil)
	_ broker.PubSubBroker = (*Broker)(nil)
)

// DiscardLog is a no-op logging function that can be used as Broker.LogFunc
// to disable logging.
var DiscardLog = func(_ string, _ ...interface{}) {}

// Pool defines the methods required for a redis pool that provides
// a method to get a connection and to release the pool's resources.
type Pool interface {
	// Get returns a redis connection.
	Get() redis.Conn

	// Close releases the resources used by the pool.
	Close() error
}

// Broker is a broker that provides the methods to
// interact with Redis using the juggler protocol.
type Broker struct {
	// Pool is the redis pool or redisc cluster to use to get
	// short-lived connections.
	Pool Pool

	// Dial is the function to call to get a non-pooled, long-lived
	// redis connection. Typically, it can be set to redis.Pool.Dial
	// or redisc.Cluster.Dial.
	Dial func() (redis.Conn, error)

	// BlockingTimeout is the time to wait for a value on calls to
	// BRPOP before trying again. The default of 0 means no timeout.
	BlockingTimeout time.Duration

	// LogFunc is the logging function to use. If nil, log.Printf
	// is used. It can be set to DiscardLog to disable logging.
	LogFunc func(string, ...interface{})

	// CallCap is the capacity of the CALL queue per URI. If it is
	// exceeded for a given URI, subsequent Broker.Call calls for that
	// URI will fail with an error.
	CallCap int

	// ResultCap is the capacity of the RES queue per connection UUID.
	// If it is exceeded for a given connection, Broker.Result calls
	// for that connection will fail with an error.
	ResultCap int

	// Vars can be set to an *expvar.Map to collect metrics about the
	// broker. It should be set before starting to make calls with the
	// broker.
	Vars *expvar.Map
}

// script to store the call request or call result along with
// its expiration information.
var callOrResScript = redis.NewScript(2, `
	redis.call("SET", KEYS[1], ARGV[1], "PX", tonumber(ARGV[1]))
	local res = redis.call("LPUSH", KEYS[2], ARGV[2])
	local limit = tonumber(ARGV[3])
	if res > limit and limit > 0 then
		local diff = res - limit
		redis.call("LTRIM", KEYS[2], diff, limit + diff)
		return redis.error_reply("list capacity exceeded")
	end
	return res
`)

const (
	// redis cluster-compliant keys, so that both keys are in the same slot
	callKey        = "juggler:calls:{%s}"            // 1: URI
	callTimeoutKey = "juggler:calls:timeout:{%s}:%s" // 1: URI, 2: mUUID

	// redis cluster-compliant keys, so that both keys are in the same slot
	resKey        = "juggler:results:{%s}"            // 1: cUUID
	resTimeoutKey = "juggler:results:timeout:{%s}:%s" // 1: cUUID, 2: mUUID
)

// Call registers a call request in the broker.
func (b *Broker) Call(cp *message.CallPayload, timeout time.Duration) error {
	k1 := fmt.Sprintf(callTimeoutKey, cp.URI, cp.MsgUUID)
	k2 := fmt.Sprintf(callKey, cp.URI)
	return registerCallOrRes(b.Pool, cp, timeout, b.CallCap, k1, k2)
}

// Result registers a call result in the broker.
func (b *Broker) Result(rp *message.ResPayload, timeout time.Duration) error {
	k1 := fmt.Sprintf(resTimeoutKey, rp.ConnUUID, rp.MsgUUID)
	k2 := fmt.Sprintf(resKey, rp.ConnUUID)
	return registerCallOrRes(b.Pool, rp, timeout, b.ResultCap, k1, k2)
}

func registerCallOrRes(pool Pool, pld interface{}, timeout time.Duration, cap int, k1, k2 string) error {
	p, err := json.Marshal(pld)
	if err != nil {
		return err
	}

	rc := pool.Get()
	defer rc.Close()

	// turn it into a cluster-aware RetryConn if running in a cluster
	rc = clusterifyConn(rc, k1, k2)

	to := int(timeout / time.Millisecond)
	if to == 0 {
		to = int(broker.DefaultCallTimeout / time.Millisecond)
	}

	_, err = callOrResScript.Do(rc,
		k1,  // key[1] : the SET key with expiration
		k2,  // key[2] : the LIST key
		to,  // argv[1] : the timeout in milliseconds
		p,   // argv[2] : the call payload
		cap, // argv[3] : the LIST capacity
	)
	return err
}

// Publish publishes an event to a channel.
func (b *Broker) Publish(channel string, pp *message.PubPayload) error {
	p, err := json.Marshal(pp)
	if err != nil {
		return err
	}

	rc := b.Pool.Get()
	defer rc.Close()

	// force selection of a random node (otherwise it would use
	// the node of the hash of the channel - which may hit the
	// same node over and over again if there are few channels).
	if bc, ok := rc.(binder); ok {
		// ignore the error, if it fails, use the connection as-is.
		// Bind without a key selects a random node.
		bc.Bind()
	}
	_, err = rc.Do("PUBLISH", channel, p)
	return err
}

// NewPubSubConn returns a new pub-sub connection that can be used
// to subscribe to and unsubscribe from channels, and to process
// incoming events.
func (b *Broker) NewPubSubConn() (broker.PubSubConn, error) {
	rc, err := b.Dial()
	if err != nil {
		return nil, err
	}
	return &pubSubConn{
		psc:   redis.PubSubConn{Conn: rc},
		logFn: b.LogFunc,
		vars:  b.Vars,
	}, nil
}

// NewCallsConn returns a new calls connection that can be used
// to process the call requests for the specified URIs.
func (b *Broker) NewCallsConn(uris ...string) (broker.CallsConn, error) {
	rc, err := b.Dial()
	if err != nil {
		return nil, err
	}
	return &callsConn{
		c:       rc,
		pool:    b.Pool,
		uris:    uris,
		vars:    b.Vars,
		timeout: b.BlockingTimeout,
		logFn:   b.LogFunc,
	}, nil
}

// NewResultsConn returns a new results connection that can be used
// to process the call results for the specified connection UUID.
func (b *Broker) NewResultsConn(connUUID uuid.UUID) (broker.ResultsConn, error) {
	rc, err := b.Dial()
	if err != nil {
		return nil, err
	}
	return &resultsConn{
		c:        rc,
		pool:     b.Pool,
		connUUID: connUUID,
		vars:     b.Vars,
		timeout:  b.BlockingTimeout,
		logFn:    b.LogFunc,
	}, nil
}

const (
	// TODO : maybe make that customizable, not super critical.
	clusterConnMaxAttempts   = 4
	clusterConnTryAgainDelay = 100 * time.Millisecond
)

type binder interface {
	Bind(...string) error
}

func clusterifyConn(rc redis.Conn, keys ...string) redis.Conn {
	// if it implements Bind, call it and make it a RetryConn so
	// that it follows redirections in a cluster.
	if bc, ok := rc.(binder); ok {
		// if Bind fails, go on with the call as usual, but if it
		// succeeds, try to turn it into a RetryConn.
		if err := bc.Bind(keys...); err == nil {
			retry, err := redisc.RetryConn(rc, clusterConnMaxAttempts, clusterConnTryAgainDelay)
			// again, if it fails, ignore and go on with the normal conn,
			// but if it succeds, replace the conn with this one.
			if err == nil {
				rc = retry
			}
		}
	}
	return rc
}

func logf(fn func(string, ...interface{}), f string, args ...interface{}) {
	if fn != nil {
		fn(f, args...)
	} else {
		log.Printf(f, args...)
	}
}
<|endoftext|>


<commit_before>package fbapi_test

import (
	"flag"
	"fmt"
	"net/http"
	"net/url"
	"testing"
	"time"

	"github.com/daaku/go.fbapi"
	"github.com/daaku/go.flagconfig"
	"github.com/daaku/go.httpcontrol"
)

var (
	defaultHttpTransport = &httpcontrol.Transport{
		MaxIdleConnsPerHost:   50,
		DialTimeout:           3 * time.Second,
		ResponseHeaderTimeout: 30 * time.Second,
		RequestTimeout:        time.Minute,
		Stats:                 logRequestHandler,
	}
	defaultFbClient = fbapi.ClientFlag("fbapi-test")

	logRequest = flag.Bool(
		"log-requests",
		false,
		"will trigger verbose logging of requests",
	)
)

func init() {
	flag.Usage = flagconfig.Usage
	flagconfig.Parse()
	if err := defaultHttpTransport.Start(); err != nil {
		panic(err)
	}
	defaultFbClient.Transport = defaultHttpTransport
}

func logRequestHandler(stats *httpcontrol.Stats) {
	if *logRequest {
		fmt.Println(stats.String())
		fmt.Println("Header", stats.Request.Header)
	}
}

func TestPublicGet(t *testing.T) {
	t.Parallel()
	user := struct {
		Username string `json:"username"`
	}{}
	res, err := defaultFbClient.Do(
		&http.Request{
			Method: "GET",
			URL: &url.URL{
				Path: "5526183",
			},
		},
		&user,
	)
	if err != nil {
		t.Fatal(err)
	}
	if res.StatusCode != 200 {
		t.Fatalf("was expecting status 200 but got %d", res.StatusCode)
	}
	if user.Username != "naitik" {
		t.Fatalf("was expecting naitik but got %s", user.Username)
	}
}

func TestInvalidGet(t *testing.T) {
	t.Parallel()
	res, err := defaultFbClient.Do(
		&http.Request{
			Method: "GET",
			URL: &url.URL{
				Path: "20aa2519-4745-4522-92a9-4522b8edf6e9",
			},
		},
		nil,
	)
	if err == nil {
		t.Fatal("was expecting error")
	}

	const expected = `GET ` +
		`https://graph.facebook.com/20aa2519-4745-4522-92a9-4522b8edf6e9 got ` +
		`404 Not Found failed with code 803 type OAuthException message (#803) ` +
		`Some of the aliases you requested do not exist: ` +
		`20aa2519-4745-4522-92a9-4522b8edf6e9`

	if err.Error() != expected {
		t.Fatalf(`expected "%s" got "%s"`, expected, err)
	}

	if res.StatusCode != 404 {
		t.Fatalf("was expecting status 404 but got %d", res.StatusCode)
	}
}
<commit_msg>test nil urls<commit_after>package fbapi_test

import (
	"flag"
	"fmt"
	"net/http"
	"net/url"
	"testing"
	"time"

	"github.com/daaku/go.fbapi"
	"github.com/daaku/go.flagconfig"
	"github.com/daaku/go.httpcontrol"
)

var (
	defaultHttpTransport = &httpcontrol.Transport{
		MaxIdleConnsPerHost:   50,
		DialTimeout:           3 * time.Second,
		ResponseHeaderTimeout: 30 * time.Second,
		RequestTimeout:        time.Minute,
		Stats:                 logRequestHandler,
	}
	defaultFbClient = fbapi.ClientFlag("fbapi-test")

	logRequest = flag.Bool(
		"log-requests",
		false,
		"will trigger verbose logging of requests",
	)
)

func init() {
	flag.Usage = flagconfig.Usage
	flagconfig.Parse()
	if err := defaultHttpTransport.Start(); err != nil {
		panic(err)
	}
	defaultFbClient.Transport = defaultHttpTransport
}

func logRequestHandler(stats *httpcontrol.Stats) {
	if *logRequest {
		fmt.Println(stats.String())
		fmt.Println("Header", stats.Request.Header)
	}
}

func TestPublicGet(t *testing.T) {
	t.Parallel()
	user := struct {
		Username string `json:"username"`
	}{}
	res, err := defaultFbClient.Do(
		&http.Request{
			Method: "GET",
			URL: &url.URL{
				Path: "5526183",
			},
		},
		&user,
	)
	if err != nil {
		t.Fatal(err)
	}
	if res.StatusCode != 200 {
		t.Fatalf("was expecting status 200 but got %d", res.StatusCode)
	}
	if user.Username != "naitik" {
		t.Fatalf("was expecting naitik but got %s", user.Username)
	}
}

func TestInvalidGet(t *testing.T) {
	t.Parallel()
	res, err := defaultFbClient.Do(
		&http.Request{
			Method: "GET",
			URL: &url.URL{
				Path: "20aa2519-4745-4522-92a9-4522b8edf6e9",
			},
		},
		nil,
	)
	if err == nil {
		t.Fatal("was expecting error")
	}

	const expected = `GET ` +
		`https://graph.facebook.com/20aa2519-4745-4522-92a9-4522b8edf6e9 got ` +
		`404 Not Found failed with code 803 type OAuthException message (#803) ` +
		`Some of the aliases you requested do not exist: ` +
		`20aa2519-4745-4522-92a9-4522b8edf6e9`

	if err.Error() != expected {
		t.Fatalf(`expected "%s" got "%s"`, expected, err)
	}

	if res.StatusCode != 404 {
		t.Fatalf("was expecting status 404 but got %d", res.StatusCode)
	}
}

func TestNilURLWithDefaultBaseURL(t *testing.T) {
	t.Parallel()
	res, err := defaultFbClient.Do(&http.Request{Method: "GET"}, nil)
	if err == nil {
		t.Fatal("was expecting error")
	}

	const expected = `GET https://graph.facebook.com/ got 400 Bad Request ` +
		`failed with code 100 type GraphMethodException message Unsupported get ` +
		`request`

	if err.Error() != expected {
		t.Fatalf(`expected "%s" got "%s"`, expected, err)
	}

	if res.StatusCode != 400 {
		t.Fatalf("was expecting status 400 but got %d", res.StatusCode)
	}
}

func TestNilURLWithBaseURL(t *testing.T) {
	t.Parallel()
	client := &fbapi.Client{
		BaseURL: &url.URL{
			Scheme: "https",
			Host:   "graph.facebook.com",
			Path:   "/20aa2519-4745-4522-92a9-4522b8edf6e9",
		},
	}
	res, err := client.Do(&http.Request{Method: "GET"}, nil)
	if err == nil {
		t.Fatal("was expecting error")
	}

	const expected = `GET ` +
		`https://graph.facebook.com/20aa2519-4745-4522-92a9-4522b8edf6e9 got ` +
		`404 Not Found failed with code 803 type OAuthException message (#803) ` +
		`Some of the aliases you requested do not exist: ` +
		`20aa2519-4745-4522-92a9-4522b8edf6e9`

	if err.Error() != expected {
		t.Fatalf(`expected "%s" got "%s"`, expected, err)
	}

	if res.StatusCode != 404 {
		t.Fatalf("was expecting status 404 but got %d", res.StatusCode)
	}
}
<|endoftext|>


<commit_before>package dns

import (
	"strconv"
	"testing"
	"time"
)

func TestClientSync(t *testing.T) {
	HandleFunc("miek.nl.", HelloServer)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeSOA)

	c := new(Client)
	r, _, err := c.Exchange(m, addrstr)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}
	if r != nil && r.Rcode != RcodeSuccess {
		t.Errorf("failed to get an valid answer\n%v", r)
	}
	// And now with plain Exchange().
	r, err = Exchange(m, addrstr)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}
	if r == nil || r.Rcode != RcodeSuccess {
		t.Errorf("failed to get an valid answer\n%v", r)
	}
}

func TestClientSyncBadId(t *testing.T) {
	HandleFunc("miek.nl.", HelloServerBadId)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeSOA)

	c := new(Client)
	if _, _, err := c.Exchange(m, addrstr); err != ErrId {
		t.Errorf("did not find a bad Id")
	}
	// And now with plain Exchange().
	if _, err := Exchange(m, addrstr); err != ErrId {
		t.Errorf("did not find a bad Id")
	}
}

func TestClientEDNS0(t *testing.T) {
	HandleFunc("miek.nl.", HelloServer)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeDNSKEY)

	m.SetEdns0(2048, true)

	c := new(Client)
	r, _, err := c.Exchange(m, addrstr)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}

	if r != nil && r.Rcode != RcodeSuccess {
		t.Errorf("failed to get an valid answer\n%v", r)
	}
}

// Validates the transmission and parsing of local EDNS0 options.
func TestClientEDNS0Local(t *testing.T) {
	optStr1 := "1979:0x0707"
	optStr2 := strconv.Itoa(EDNS0LOCALSTART) + ":0x0601"

	handler := func(w ResponseWriter, req *Msg) {
		m := new(Msg)
		m.SetReply(req)

		m.Extra = make([]RR, 1, 2)
		m.Extra[0] = &TXT{Hdr: RR_Header{Name: m.Question[0].Name, Rrtype: TypeTXT, Class: ClassINET, Ttl: 0}, Txt: []string{"Hello local edns"}}

		// If the local options are what we expect, then reflect them back.
		ec1 := req.Extra[0].(*OPT).Option[0].(*EDNS0_LOCAL).String()
		ec2 := req.Extra[0].(*OPT).Option[1].(*EDNS0_LOCAL).String()
		if ec1 == optStr1 && ec2 == optStr2 {
			m.Extra = append(m.Extra, req.Extra[0])
		}

		w.WriteMsg(m)
	}

	HandleFunc("miek.nl.", handler)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %s", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeTXT)

	// Add two local edns options to the query.
	ec1 := &EDNS0_LOCAL{Code: 1979, Data: []byte{7, 7}}
	ec2 := &EDNS0_LOCAL{Code: EDNS0LOCALSTART, Data: []byte{6, 1}}
	o := &OPT{Hdr: RR_Header{Name: ".", Rrtype: TypeOPT}, Option: []EDNS0{ec1, ec2}}
	m.Extra = append(m.Extra, o)

	c := new(Client)
	r, _, e := c.Exchange(m, addrstr)
	if e != nil {
		t.Logf("failed to exchange: %s", e.Error())
		t.Fail()
	}

	if r != nil && r.Rcode != RcodeSuccess {
		t.Log("failed to get a valid answer")
		t.Fail()
		t.Logf("%v\n", r)
	}

	txt := r.Extra[0].(*TXT).Txt[0]
	if txt != "Hello local edns" {
		t.Log("Unexpected result for miek.nl", txt, "!= Hello local edns")
		t.Fail()
	}

	// Validate the local options in the reply.
	got := r.Extra[1].(*OPT).Option[0].(*EDNS0_LOCAL).String()
	if got != optStr1 {
		t.Logf("failed to get local edns0 answer; got %s, expected %s", got, optStr1)
		t.Fail()
		t.Logf("%v\n", r)
	}

	got = r.Extra[1].(*OPT).Option[1].(*EDNS0_LOCAL).String()
	if got != optStr2 {
		t.Logf("failed to get local edns0 answer; got %s, expected %s", got, optStr2)
		t.Fail()
		t.Logf("%v\n", r)
	}
}

func TestSingleInflight(t *testing.T) {
	// Test is inherently racy, because queries might actually be returned before the test
	// is over, leading to multiple queries even with SingleInflight. This ofcourse then
	// leads to diff. rrts and the test fails. Number of tests is now 3, to lower the chance
	// for the race to hit.
	HandleFunc("miek.nl.", HelloServer)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeDNSKEY)

	c := new(Client)
	c.SingleInflight = true
	nr := 3
	ch := make(chan time.Duration)
	for i := 0; i < nr; i++ {
		go func() {
			_, rtt, _ := c.Exchange(m, addrstr)
			ch <- rtt
		}()
	}
	i := 0
	var first time.Duration
	// With inflight *all* rtt are identical, and by doing actual lookups
	// the chances that this is a coincidence is small.
Loop:
	for {
		select {
		case rtt := <-ch:
			if i == 0 {
				first = rtt
			} else {
				if first != rtt {
					t.Errorf("all rtts should be equal, got %d want %d", rtt, first)
				}
			}
			i++
			if i == nr {
				break Loop
			}
		}
	}
}

// ExampleUpdateLeaseTSIG shows how to update a lease signed with TSIG.
func ExampleUpdateLeaseTSIG(t *testing.T) {
	m := new(Msg)
	m.SetUpdate("t.local.ip6.io.")
	rr, _ := NewRR("t.local.ip6.io. 30 A 127.0.0.1")
	rrs := make([]RR, 1)
	rrs[0] = rr
	m.Insert(rrs)

	leaseRr := new(OPT)
	leaseRr.Hdr.Name = "."
	leaseRr.Hdr.Rrtype = TypeOPT
	e := new(EDNS0_UL)
	e.Code = EDNS0UL
	e.Lease = 120
	leaseRr.Option = append(leaseRr.Option, e)
	m.Extra = append(m.Extra, leaseRr)

	c := new(Client)
	m.SetTsig("polvi.", HmacMD5, 300, time.Now().Unix())
	c.TsigSecret = map[string]string{"polvi.": "pRZgBrBvI4NAHZYhxmhs/Q=="}

	_, _, err := c.Exchange(m, "127.0.0.1:53")
	if err != nil {
		t.Error(err)
	}
}

func TestClientConn(t *testing.T) {
	HandleFunc("miek.nl.", HelloServer)
	defer HandleRemove("miek.nl.")

	// This uses TCP just to make it slightly different than TestClientSync
	s, addrstr, err := RunLocalTCPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeSOA)

	cn, err := Dial("tcp", addrstr)
	if err != nil {
		t.Errorf("failed to dial %s: %v", addrstr, err)
	}

	err = cn.WriteMsg(m)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}
	r, err := cn.ReadMsg()
	if r == nil || r.Rcode != RcodeSuccess {
		t.Errorf("failed to get an valid answer\n%v", r)
	}

	err = cn.WriteMsg(m)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}
	h := new(Header)
	buf, err := cn.ReadMsgHeader(h)
	if buf == nil {
		t.Errorf("failed to get an valid answer\n%v", r)
	}
	if int(h.Bits&0xF) != RcodeSuccess {
		t.Errorf("failed to get an valid answer in ReadMsgHeader\n%v", r)
	}
	if h.Ancount != 0 || h.Qdcount != 1 || h.Nscount != 0 || h.Arcount != 1 {
		t.Errorf("expected to have question and additional in response; got something else: %+v", h)
	}
	if err = r.Unpack(buf); err != nil {
		t.Errorf("unable to unpack message fully: %v", err)
	}
}
<commit_msg>Remove TestSingleInflight as per #250<commit_after>package dns

import (
	"strconv"
	"testing"
	"time"
)

func TestClientSync(t *testing.T) {
	HandleFunc("miek.nl.", HelloServer)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeSOA)

	c := new(Client)
	r, _, err := c.Exchange(m, addrstr)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}
	if r != nil && r.Rcode != RcodeSuccess {
		t.Errorf("failed to get an valid answer\n%v", r)
	}
	// And now with plain Exchange().
	r, err = Exchange(m, addrstr)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}
	if r == nil || r.Rcode != RcodeSuccess {
		t.Errorf("failed to get an valid answer\n%v", r)
	}
}

func TestClientSyncBadId(t *testing.T) {
	HandleFunc("miek.nl.", HelloServerBadId)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeSOA)

	c := new(Client)
	if _, _, err := c.Exchange(m, addrstr); err != ErrId {
		t.Errorf("did not find a bad Id")
	}
	// And now with plain Exchange().
	if _, err := Exchange(m, addrstr); err != ErrId {
		t.Errorf("did not find a bad Id")
	}
}

func TestClientEDNS0(t *testing.T) {
	HandleFunc("miek.nl.", HelloServer)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeDNSKEY)

	m.SetEdns0(2048, true)

	c := new(Client)
	r, _, err := c.Exchange(m, addrstr)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}

	if r != nil && r.Rcode != RcodeSuccess {
		t.Errorf("failed to get an valid answer\n%v", r)
	}
}

// Validates the transmission and parsing of local EDNS0 options.
func TestClientEDNS0Local(t *testing.T) {
	optStr1 := "1979:0x0707"
	optStr2 := strconv.Itoa(EDNS0LOCALSTART) + ":0x0601"

	handler := func(w ResponseWriter, req *Msg) {
		m := new(Msg)
		m.SetReply(req)

		m.Extra = make([]RR, 1, 2)
		m.Extra[0] = &TXT{Hdr: RR_Header{Name: m.Question[0].Name, Rrtype: TypeTXT, Class: ClassINET, Ttl: 0}, Txt: []string{"Hello local edns"}}

		// If the local options are what we expect, then reflect them back.
		ec1 := req.Extra[0].(*OPT).Option[0].(*EDNS0_LOCAL).String()
		ec2 := req.Extra[0].(*OPT).Option[1].(*EDNS0_LOCAL).String()
		if ec1 == optStr1 && ec2 == optStr2 {
			m.Extra = append(m.Extra, req.Extra[0])
		}

		w.WriteMsg(m)
	}

	HandleFunc("miek.nl.", handler)
	defer HandleRemove("miek.nl.")

	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %s", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeTXT)

	// Add two local edns options to the query.
	ec1 := &EDNS0_LOCAL{Code: 1979, Data: []byte{7, 7}}
	ec2 := &EDNS0_LOCAL{Code: EDNS0LOCALSTART, Data: []byte{6, 1}}
	o := &OPT{Hdr: RR_Header{Name: ".", Rrtype: TypeOPT}, Option: []EDNS0{ec1, ec2}}
	m.Extra = append(m.Extra, o)

	c := new(Client)
	r, _, e := c.Exchange(m, addrstr)
	if e != nil {
		t.Logf("failed to exchange: %s", e.Error())
		t.Fail()
	}

	if r != nil && r.Rcode != RcodeSuccess {
		t.Log("failed to get a valid answer")
		t.Fail()
		t.Logf("%v\n", r)
	}

	txt := r.Extra[0].(*TXT).Txt[0]
	if txt != "Hello local edns" {
		t.Log("Unexpected result for miek.nl", txt, "!= Hello local edns")
		t.Fail()
	}

	// Validate the local options in the reply.
	got := r.Extra[1].(*OPT).Option[0].(*EDNS0_LOCAL).String()
	if got != optStr1 {
		t.Logf("failed to get local edns0 answer; got %s, expected %s", got, optStr1)
		t.Fail()
		t.Logf("%v\n", r)
	}

	got = r.Extra[1].(*OPT).Option[1].(*EDNS0_LOCAL).String()
	if got != optStr2 {
		t.Logf("failed to get local edns0 answer; got %s, expected %s", got, optStr2)
		t.Fail()
		t.Logf("%v\n", r)
	}
}

// ExampleUpdateLeaseTSIG shows how to update a lease signed with TSIG.
func ExampleUpdateLeaseTSIG(t *testing.T) {
	m := new(Msg)
	m.SetUpdate("t.local.ip6.io.")
	rr, _ := NewRR("t.local.ip6.io. 30 A 127.0.0.1")
	rrs := make([]RR, 1)
	rrs[0] = rr
	m.Insert(rrs)

	leaseRr := new(OPT)
	leaseRr.Hdr.Name = "."
	leaseRr.Hdr.Rrtype = TypeOPT
	e := new(EDNS0_UL)
	e.Code = EDNS0UL
	e.Lease = 120
	leaseRr.Option = append(leaseRr.Option, e)
	m.Extra = append(m.Extra, leaseRr)

	c := new(Client)
	m.SetTsig("polvi.", HmacMD5, 300, time.Now().Unix())
	c.TsigSecret = map[string]string{"polvi.": "pRZgBrBvI4NAHZYhxmhs/Q=="}

	_, _, err := c.Exchange(m, "127.0.0.1:53")
	if err != nil {
		t.Error(err)
	}
}

func TestClientConn(t *testing.T) {
	HandleFunc("miek.nl.", HelloServer)
	defer HandleRemove("miek.nl.")

	// This uses TCP just to make it slightly different than TestClientSync
	s, addrstr, err := RunLocalTCPServer("127.0.0.1:0")
	if err != nil {
		t.Fatalf("Unable to run test server: %v", err)
	}
	defer s.Shutdown()

	m := new(Msg)
	m.SetQuestion("miek.nl.", TypeSOA)

	cn, err := Dial("tcp", addrstr)
	if err != nil {
		t.Errorf("failed to dial %s: %v", addrstr, err)
	}

	err = cn.WriteMsg(m)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}
	r, err := cn.ReadMsg()
	if r == nil || r.Rcode != RcodeSuccess {
		t.Errorf("failed to get an valid answer\n%v", r)
	}

	err = cn.WriteMsg(m)
	if err != nil {
		t.Errorf("failed to exchange: %v", err)
	}
	h := new(Header)
	buf, err := cn.ReadMsgHeader(h)
	if buf == nil {
		t.Errorf("failed to get an valid answer\n%v", r)
	}
	if int(h.Bits&0xF) != RcodeSuccess {
		t.Errorf("failed to get an valid answer in ReadMsgHeader\n%v", r)
	}
	if h.Ancount != 0 || h.Qdcount != 1 || h.Nscount != 0 || h.Arcount != 1 {
		t.Errorf("expected to have question and additional in response; got something else: %+v", h)
	}
	if err = r.Unpack(buf); err != nil {
		t.Errorf("unable to unpack message fully: %v", err)
	}
}
<|endoftext|>


<commit_before>package baggageclaim_test

import (
	"net/http"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/ghttp"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/lager/lagertest"

	"github.com/concourse/baggageclaim"
	"github.com/concourse/baggageclaim/api"
	"github.com/concourse/baggageclaim/client"
	"github.com/concourse/baggageclaim/volume"
)

var _ = Describe("Baggage Claim Client", func() {
	Describe("getting the heartbeat interval from a TTL", func() {
		It("has an upper bound of 1 minute", func() {
			ttlInSeconds := uint(500)
			interval := client.IntervalForTTL(ttlInSeconds)

			Expect(interval).To(Equal(time.Minute))
		})

		Context("when the TTL is small", func() {
			It("Returns an interval that is half of the TTL", func() {
				ttlInSeconds := uint(5)
				interval := client.IntervalForTTL(ttlInSeconds)

				Expect(interval).To(Equal(2500 * time.Millisecond))
			})
		})
	})
	Describe("Interacting with the server", func() {
		var (
			bcServer *ghttp.Server
			logger   lager.Logger
			bcClient baggageclaim.Client
		)

		BeforeEach(func() {
			bcServer = ghttp.NewServer()
			logger = lagertest.NewTestLogger("client")
			bcClient = client.New(bcServer.URL())
		})

		AfterEach(func() {
			bcServer.Close()
		})

		Describe("Looking up a volume by handle", func() {
			It("heartbeats immediately to reset the TTL", func() {
				didHeartbeat := make(chan struct{})

				bcServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/volumes/some-handle"),
						ghttp.RespondWithJSONEncoded(200, volume.Volume{
							Handle:     "some-handle",
							Path:       "some-path",
							Properties: volume.Properties{},
							TTL:        volume.TTL(1),
							ExpiresAt:  time.Now().Add(time.Second),
						}),
					),
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("PUT", "/volumes/some-handle/ttl"),
						func(w http.ResponseWriter, r *http.Request) {
							close(didHeartbeat)
						},
						ghttp.RespondWith(http.StatusNoContent, ""),
					),
				)
				bcClient.LookupVolume(logger, "some-handle")
				Eventually(didHeartbeat, time.Second).Should(BeClosed())
			})
			Context("when the volume's TTL is 0", func() {
				It("does not heartbeat", func() {
					bcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/volumes/some-handle"),
							ghttp.RespondWithJSONEncoded(200, volume.Volume{
								Handle:     "some-handle",
								Path:       "some-path",
								Properties: volume.Properties{},
								TTL:        volume.TTL(0),
								ExpiresAt:  time.Now().Add(time.Second),
							}),
						),
					)
					_, err := bcClient.LookupVolume(logger, "some-handle")
					Expect(err).NotTo(HaveOccurred())
					time.Sleep(1) // wait to verify it is not heartbeating
				})
			})
			Context("when the intial heartbeat fails", func() {
				It("reports that the volume could not be found", func() {
					bcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/volumes/some-handle"),
							ghttp.RespondWithJSONEncoded(200, volume.Volume{
								Handle:     "some-handle",
								Path:       "some-path",
								Properties: volume.Properties{},
								TTL:        volume.TTL(1),
								ExpiresAt:  time.Now().Add(time.Second),
							}),
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("PUT", "/volumes/some-handle/ttl"),
							func(w http.ResponseWriter, r *http.Request) {
								api.RespondWithError(w, volume.ErrSetTTLFailed, http.StatusNotFound)
							},
						),
					)
					foundVolume, err := bcClient.LookupVolume(logger, "some-handle")
					Expect(foundVolume).To(BeNil())
					Expect(err).To(Equal(volume.ErrVolumeDoesNotExist))
				})
			})
		})
		Describe("Listing volumes", func() {
			Context("when the inital heartbeat fails for a volume", func() {
				It("it is omitted from the returned list of volumes", func() {
					bcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/volumes"),
							ghttp.RespondWithJSONEncoded(200, []volume.Volume{
								{
									Handle:     "some-handle",
									Path:       "some-path",
									Properties: volume.Properties{},
									TTL:        volume.TTL(1),
									ExpiresAt:  time.Now().Add(time.Second),
								},
								{
									Handle:     "another-handle",
									Path:       "some-path",
									Properties: volume.Properties{},
									TTL:        volume.TTL(1),
									ExpiresAt:  time.Now().Add(time.Second),
								},
							}),
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("PUT", "/volumes/some-handle/ttl"),
							func(w http.ResponseWriter, r *http.Request) {
								w.WriteHeader(http.StatusNoContent)
							},
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("PUT", "/volumes/another-handle/ttl"),
							func(w http.ResponseWriter, r *http.Request) {
								api.RespondWithError(w, volume.ErrSetTTLFailed, http.StatusNotFound)
							},
						),
					)
					volumes, err := bcClient.ListVolumes(logger, baggageclaim.VolumeProperties{})
					Expect(err).NotTo(HaveOccurred())
					Expect(len(volumes)).To(Equal(1))
					Expect(volumes[0].Handle()).To(Equal("some-handle"))
				})
			})
		})
		Describe("Creating volumes", func() {
			Context("when the inital heartbeat fails for the volume", func() {
				It("reports that the volume could not be found", func() {
					bcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("POST", "/volumes"),
							ghttp.RespondWithJSONEncoded(201, volume.Volume{
								Handle:     "some-handle",
								Path:       "some-path",
								Properties: volume.Properties{},
								TTL:        volume.TTL(1),
								ExpiresAt:  time.Now().Add(time.Second),
							}),
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("PUT", "/volumes/some-handle/ttl"),
							func(w http.ResponseWriter, r *http.Request) {
								api.RespondWithError(w, volume.ErrSetTTLFailed, http.StatusNotFound)
							},
						),
					)
					createdVolume, err := bcClient.CreateVolume(logger, baggageclaim.VolumeSpec{})
					Expect(createdVolume).To(BeNil())
					Expect(err).To(Equal(volume.ErrVolumeDoesNotExist))
				})
			})
		})
	})
})
<commit_msg>fix units which apparently exist<commit_after>package baggageclaim_test

import (
	"net/http"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/ghttp"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/lager/lagertest"

	"github.com/concourse/baggageclaim"
	"github.com/concourse/baggageclaim/api"
	"github.com/concourse/baggageclaim/client"
	"github.com/concourse/baggageclaim/volume"
)

var _ = Describe("Baggage Claim Client", func() {
	Describe("getting the heartbeat interval from a TTL", func() {
		It("has an upper bound of 1 minute", func() {
			interval := client.IntervalForTTL(500 * time.Second)

			Expect(interval).To(Equal(time.Minute))
		})

		Context("when the TTL is small", func() {
			It("Returns an interval that is half of the TTL", func() {
				interval := client.IntervalForTTL(5 * time.Second)

				Expect(interval).To(Equal(2500 * time.Millisecond))
			})
		})
	})
	Describe("Interacting with the server", func() {
		var (
			bcServer *ghttp.Server
			logger   lager.Logger
			bcClient baggageclaim.Client
		)

		BeforeEach(func() {
			bcServer = ghttp.NewServer()
			logger = lagertest.NewTestLogger("client")
			bcClient = client.New(bcServer.URL())
		})

		AfterEach(func() {
			bcServer.Close()
		})

		Describe("Looking up a volume by handle", func() {
			It("heartbeats immediately to reset the TTL", func() {
				didHeartbeat := make(chan struct{})

				bcServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/volumes/some-handle"),
						ghttp.RespondWithJSONEncoded(200, volume.Volume{
							Handle:     "some-handle",
							Path:       "some-path",
							Properties: volume.Properties{},
							TTL:        volume.TTL(1),
							ExpiresAt:  time.Now().Add(time.Second),
						}),
					),
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("PUT", "/volumes/some-handle/ttl"),
						func(w http.ResponseWriter, r *http.Request) {
							close(didHeartbeat)
						},
						ghttp.RespondWith(http.StatusNoContent, ""),
					),
				)
				bcClient.LookupVolume(logger, "some-handle")
				Eventually(didHeartbeat, time.Second).Should(BeClosed())
			})
			Context("when the volume's TTL is 0", func() {
				It("does not heartbeat", func() {
					bcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/volumes/some-handle"),
							ghttp.RespondWithJSONEncoded(200, volume.Volume{
								Handle:     "some-handle",
								Path:       "some-path",
								Properties: volume.Properties{},
								TTL:        volume.TTL(0),
								ExpiresAt:  time.Now().Add(time.Second),
							}),
						),
					)
					_, err := bcClient.LookupVolume(logger, "some-handle")
					Expect(err).NotTo(HaveOccurred())
					time.Sleep(1) // wait to verify it is not heartbeating
				})
			})
			Context("when the intial heartbeat fails", func() {
				It("reports that the volume could not be found", func() {
					bcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/volumes/some-handle"),
							ghttp.RespondWithJSONEncoded(200, volume.Volume{
								Handle:     "some-handle",
								Path:       "some-path",
								Properties: volume.Properties{},
								TTL:        volume.TTL(1),
								ExpiresAt:  time.Now().Add(time.Second),
							}),
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("PUT", "/volumes/some-handle/ttl"),
							func(w http.ResponseWriter, r *http.Request) {
								api.RespondWithError(w, volume.ErrSetTTLFailed, http.StatusNotFound)
							},
						),
					)
					foundVolume, err := bcClient.LookupVolume(logger, "some-handle")
					Expect(foundVolume).To(BeNil())
					Expect(err).To(Equal(volume.ErrVolumeDoesNotExist))
				})
			})
		})
		Describe("Listing volumes", func() {
			Context("when the inital heartbeat fails for a volume", func() {
				It("it is omitted from the returned list of volumes", func() {
					bcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/volumes"),
							ghttp.RespondWithJSONEncoded(200, []volume.Volume{
								{
									Handle:     "some-handle",
									Path:       "some-path",
									Properties: volume.Properties{},
									TTL:        volume.TTL(1),
									ExpiresAt:  time.Now().Add(time.Second),
								},
								{
									Handle:     "another-handle",
									Path:       "some-path",
									Properties: volume.Properties{},
									TTL:        volume.TTL(1),
									ExpiresAt:  time.Now().Add(time.Second),
								},
							}),
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("PUT", "/volumes/some-handle/ttl"),
							func(w http.ResponseWriter, r *http.Request) {
								w.WriteHeader(http.StatusNoContent)
							},
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("PUT", "/volumes/another-handle/ttl"),
							func(w http.ResponseWriter, r *http.Request) {
								api.RespondWithError(w, volume.ErrSetTTLFailed, http.StatusNotFound)
							},
						),
					)
					volumes, err := bcClient.ListVolumes(logger, baggageclaim.VolumeProperties{})
					Expect(err).NotTo(HaveOccurred())
					Expect(len(volumes)).To(Equal(1))
					Expect(volumes[0].Handle()).To(Equal("some-handle"))
				})
			})
		})
		Describe("Creating volumes", func() {
			Context("when the inital heartbeat fails for the volume", func() {
				It("reports that the volume could not be found", func() {
					bcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("POST", "/volumes"),
							ghttp.RespondWithJSONEncoded(201, volume.Volume{
								Handle:     "some-handle",
								Path:       "some-path",
								Properties: volume.Properties{},
								TTL:        volume.TTL(1),
								ExpiresAt:  time.Now().Add(time.Second),
							}),
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("PUT", "/volumes/some-handle/ttl"),
							func(w http.ResponseWriter, r *http.Request) {
								api.RespondWithError(w, volume.ErrSetTTLFailed, http.StatusNotFound)
							},
						),
					)
					createdVolume, err := bcClient.CreateVolume(logger, baggageclaim.VolumeSpec{})
					Expect(createdVolume).To(BeNil())
					Expect(err).To(Equal(volume.ErrVolumeDoesNotExist))
				})
			})
		})
	})
})
<|endoftext|>