<commit_before>package c2

import (
	"runtime"
	"bufio"
	"strconv"
	"bytes"
)

type HttpError interface {
	error
	Stack() string
	HttpCode() int
}

type codedErrorString struct {
	s string
	c int
	stack string
}

func (e *codedErrorString) Error() string {
	return e.s
}

// 404 - Not Found
// 500 - Internal error
// 501 - Not implemented
// 400 - Bad request (user error)
// 409 - Conflict - Like bad request, but based on existing state of data/system.
func (e *codedErrorString) HttpCode() int {
	return e.c
}

func (e *codedErrorString) Stack() string {
	return e.stack
}

func NewErr(msg string) error {
	return &codedErrorString{
		s : msg,
		c : 500,
		stack: DumpStack(),
	}
}

func NewErrC(msg string, code int) error {
	return &codedErrorString{
		s : msg,
		c : code,
		stack: DumpStack(),
	}
}

func trim(s string, max int) string {
	if len(s) > max {
		return "..." + s[len(s) - (max + 3):]
	}
	return s
}

func DumpStack() string {
	var buff bytes.Buffer
	w := bufio.NewWriter(&buff)
	var stack [25]uintptr
	len := runtime.Callers(2, stack[:])
	for i := 1; i < len; i++ {
		f := runtime.FuncForPC(stack[i])
		w.WriteRune(' ')
		w.WriteString(f.Name())
		w.WriteRune(' ')
		file, lineno := f.FileLine(stack[i - 1])
		w.WriteString(trim(file, 20))
		w.WriteRune(':')
		w.WriteString(strconv.Itoa(lineno))
		w.WriteString("\n")
	}
	w.Flush()
	return buff.String()
}<commit_msg>add is 404 error helper<commit_after>package c2

import (
	"runtime"
	"bufio"
	"strconv"
	"bytes"
)

type HttpError interface {
	error
	Stack() string
	HttpCode() int
}

func IsNotFoundErr(err error) bool {
	if herr, isHErr := err.(HttpError); isHErr {
		return herr.HttpCode() == 404
	}
	return false
}

type codedErrorString struct {
	s string
	c int
	stack string
}

func (e *codedErrorString) Error() string {
	return e.s
}

// 404 - Not Found
// 500 - Internal error
// 501 - Not implemented
// 400 - Bad request (user error)
// 409 - Conflict - Like bad request, but based on existing state of data/system.
func (e *codedErrorString) HttpCode() int {
	return e.c
}

func (e *codedErrorString) Stack() string {
	return e.stack
}

func NewErr(msg string) error {
	return &codedErrorString{
		s : msg,
		c : 500,
		stack: DumpStack(),
	}
}

func NewErrC(msg string, code int) error {
	return &codedErrorString{
		s : msg,
		c : code,
		stack: DumpStack(),
	}
}

func trim(s string, max int) string {
	if len(s) > max {
		return "..." + s[len(s) - (max + 3):]
	}
	return s
}

func DumpStack() string {
	var buff bytes.Buffer
	w := bufio.NewWriter(&buff)
	var stack [25]uintptr
	len := runtime.Callers(2, stack[:])
	for i := 1; i < len; i++ {
		f := runtime.FuncForPC(stack[i])
		w.WriteRune(' ')
		w.WriteString(f.Name())
		w.WriteRune(' ')
		file, lineno := f.FileLine(stack[i - 1])
		w.WriteString(trim(file, 20))
		w.WriteRune(':')
		w.WriteString(strconv.Itoa(lineno))
		w.WriteString("\n")
	}
	w.Flush()
	return buff.String()
}<|endoftext|>


<commit_before>// Copyright 2011 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.

package http_test

import (
	"io/ioutil"
	"net/http"
	"path/filepath"
	"testing"
)

func checker(t *testing.T) func(string, error) {
	return func(call string, err error) {
		if err == nil {
			return
		}
		t.Fatalf("%s: %v", call, err)
	}
}

func TestFileTransport(t *testing.T) {
	check := checker(t)

	dname, err := ioutil.TempDir("", "")
	check("TempDir", err)
	fname := filepath.Join(dname, "foo.txt")
	err = ioutil.WriteFile(fname, []byte("Bar"), 0644)
	check("WriteFile", err)

	tr := &http.Transport{}
	tr.RegisterProtocol("file", http.NewFileTransport(http.Dir(dname)))
	c := &http.Client{Transport: tr}

	fooURLs := []string{"file:///foo.txt", "file://../foo.txt"}
	for _, urlstr := range fooURLs {
		res, err := c.Get(urlstr)
		check("Get "+urlstr, err)
		if res.StatusCode != 200 {
			t.Errorf("for %s, StatusCode = %d, want 200", urlstr, res.StatusCode)
		}
		if res.ContentLength != -1 {
			t.Errorf("for %s, ContentLength = %d, want -1", urlstr, res.ContentLength)
		}
		if res.Body == nil {
			t.Fatalf("for %s, nil Body", urlstr)
		}
		slurp, err := ioutil.ReadAll(res.Body)
		check("ReadAll "+urlstr, err)
		if string(slurp) != "Bar" {
			t.Errorf("for %s, got content %q, want %q", urlstr, string(slurp), "Bar")
		}
	}

	const badURL = "file://../no-exist.txt"
	res, err := c.Get(badURL)
	check("Get "+badURL, err)
	if res.StatusCode != 404 {
		t.Errorf("for %s, StatusCode = %d, want 404", badURL, res.StatusCode)
	}
}
<commit_msg>net/http: delete temporary files.<commit_after>// Copyright 2011 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.

package http_test

import (
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"testing"
)

func checker(t *testing.T) func(string, error) {
	return func(call string, err error) {
		if err == nil {
			return
		}
		t.Fatalf("%s: %v", call, err)
	}
}

func TestFileTransport(t *testing.T) {
	check := checker(t)

	dname, err := ioutil.TempDir("", "")
	check("TempDir", err)
	fname := filepath.Join(dname, "foo.txt")
	err = ioutil.WriteFile(fname, []byte("Bar"), 0644)
	check("WriteFile", err)
	defer os.Remove(dname)
	defer os.Remove(fname)

	tr := &http.Transport{}
	tr.RegisterProtocol("file", http.NewFileTransport(http.Dir(dname)))
	c := &http.Client{Transport: tr}

	fooURLs := []string{"file:///foo.txt", "file://../foo.txt"}
	for _, urlstr := range fooURLs {
		res, err := c.Get(urlstr)
		check("Get "+urlstr, err)
		if res.StatusCode != 200 {
			t.Errorf("for %s, StatusCode = %d, want 200", urlstr, res.StatusCode)
		}
		if res.ContentLength != -1 {
			t.Errorf("for %s, ContentLength = %d, want -1", urlstr, res.ContentLength)
		}
		if res.Body == nil {
			t.Fatalf("for %s, nil Body", urlstr)
		}
		slurp, err := ioutil.ReadAll(res.Body)
		check("ReadAll "+urlstr, err)
		if string(slurp) != "Bar" {
			t.Errorf("for %s, got content %q, want %q", urlstr, string(slurp), "Bar")
		}
	}

	const badURL = "file://../no-exist.txt"
	res, err := c.Get(badURL)
	check("Get "+badURL, err)
	if res.StatusCode != 404 {
		t.Errorf("for %s, StatusCode = %d, want 404", badURL, res.StatusCode)
	}
}
<|endoftext|>


<commit_before>package provision

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

	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/swarm"
)

type SwarmCommandContext struct {
	ContainerName string
	Env           []string
	DockerDir     string
	DockerPort    int
	IP            string
	Port          string
	AuthOptions   auth.Options
	SwarmOptions  swarm.Options
	SwarmImage    string
}

// Wrapper function to generate a docker run swarm command (manage or join)
// from a template/context and execute it.
func runSwarmCommandFromTemplate(p Provisioner, cmdTmpl string, swarmCmdContext SwarmCommandContext) error {
	var (
		executedCmdTmpl bytes.Buffer
	)

	parsedMasterCmdTemplate, err := template.New("swarmMasterCmd").Parse(cmdTmpl)
	if err != nil {
		return err
	}

	parsedMasterCmdTemplate.Execute(&executedCmdTmpl, swarmCmdContext)

	log.Debugf("The swarm command being run is: %s", executedCmdTmpl.String())

	if _, err := p.SSHCommand(executedCmdTmpl.String()); err != nil {
		return err
	}

	return nil
}

func configureSwarm(p Provisioner, swarmOptions swarm.Options, authOptions auth.Options) error {
	if !swarmOptions.IsSwarm {
		return nil
	}

	log.Info("Configuring swarm...")

	ip, err := p.GetDriver().GetIP()
	if err != nil {
		return err
	}

	u, err := url.Parse(swarmOptions.Host)
	if err != nil {
		return err
	}

	parts := strings.Split(u.Host, ":")
	port := parts[1]

	dockerDir := p.GetDockerOptionsDir()

	swarmCmdContext := SwarmCommandContext{
		ContainerName: "",
		Env:           swarmOptions.Env,
		DockerDir:     dockerDir,
		DockerPort:    2376,
		IP:            ip,
		Port:          port,
		AuthOptions:   authOptions,
		SwarmOptions:  swarmOptions,
		SwarmImage:    swarmOptions.Image,
	}

	// First things first, get the swarm image.
	if _, err := p.SSHCommand(fmt.Sprintf("sudo docker pull %s", swarmOptions.Image)); err != nil {
		return err
	}

	swarmMasterCmdTemplate := `sudo docker run -d \
--restart=always \
{{range .Env}} -e {{.}}{{end}} \
--name swarm-agent-master \
-p {{.Port}}:{{.Port}} \
-v {{.DockerDir}}:{{.DockerDir}} \
{{.SwarmImage}} \
manage \
--tlsverify \
--tlscacert={{.AuthOptions.CaCertRemotePath}} \
--tlscert={{.AuthOptions.ServerCertRemotePath}} \
--tlskey={{.AuthOptions.ServerKeyRemotePath}} \
-H {{.SwarmOptions.Host}} \
--strategy {{.SwarmOptions.Strategy}} {{range .SwarmOptions.ArbitraryFlags}} --{{.}}{{end}} {{.SwarmOptions.Discovery}}
`

	swarmWorkerCmdTemplate := `sudo docker run -d \
--restart=always \
{{range .Env}} -e {{.}}{{end}} \
--name swarm-agent \
{{.SwarmImage}} \
join --advertise {{.IP}}:{{.DockerPort}} {{.SwarmOptions.Discovery}}
`

	if swarmOptions.Master {
		log.Debug("Launching swarm master")
		if err := runSwarmCommandFromTemplate(p, swarmMasterCmdTemplate, swarmCmdContext); err != nil {
			return err
		}
	}

	log.Debug("Launch swarm worker")
	if err := runSwarmCommandFromTemplate(p, swarmWorkerCmdTemplate, swarmCmdContext); err != nil {
		return err
	}

	return nil
}
<commit_msg>added advertise flag to swarm manager<commit_after>package provision

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

	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/swarm"
)

type SwarmCommandContext struct {
	ContainerName string
	Env           []string
	DockerDir     string
	DockerPort    int
	IP            string
	Port          string
	AuthOptions   auth.Options
	SwarmOptions  swarm.Options
	SwarmImage    string
}

// Wrapper function to generate a docker run swarm command (manage or join)
// from a template/context and execute it.
func runSwarmCommandFromTemplate(p Provisioner, cmdTmpl string, swarmCmdContext SwarmCommandContext) error {
	var (
		executedCmdTmpl bytes.Buffer
	)

	parsedMasterCmdTemplate, err := template.New("swarmMasterCmd").Parse(cmdTmpl)
	if err != nil {
		return err
	}

	parsedMasterCmdTemplate.Execute(&executedCmdTmpl, swarmCmdContext)

	log.Debugf("The swarm command being run is: %s", executedCmdTmpl.String())

	if _, err := p.SSHCommand(executedCmdTmpl.String()); err != nil {
		return err
	}

	return nil
}

func configureSwarm(p Provisioner, swarmOptions swarm.Options, authOptions auth.Options) error {
	if !swarmOptions.IsSwarm {
		return nil
	}

	log.Info("Configuring swarm...")

	ip, err := p.GetDriver().GetIP()
	if err != nil {
		return err
	}

	u, err := url.Parse(swarmOptions.Host)
	if err != nil {
		return err
	}

	parts := strings.Split(u.Host, ":")
	port := parts[1]

	dockerDir := p.GetDockerOptionsDir()

	swarmCmdContext := SwarmCommandContext{
		ContainerName: "",
		Env:           swarmOptions.Env,
		DockerDir:     dockerDir,
		DockerPort:    2376,
		IP:            ip,
		Port:          port,
		AuthOptions:   authOptions,
		SwarmOptions:  swarmOptions,
		SwarmImage:    swarmOptions.Image,
	}

	// First things first, get the swarm image.
	if _, err := p.SSHCommand(fmt.Sprintf("sudo docker pull %s", swarmOptions.Image)); err != nil {
		return err
	}

	swarmMasterCmdTemplate := `sudo docker run -d \
--restart=always \
{{range .Env}} -e {{.}}{{end}} \
--name swarm-agent-master \
-p {{.Port}}:{{.Port}} \
-v {{.DockerDir}}:{{.DockerDir}} \
{{.SwarmImage}} \
manage \
--tlsverify \
--tlscacert={{.AuthOptions.CaCertRemotePath}} \
--tlscert={{.AuthOptions.ServerCertRemotePath}} \
--tlskey={{.AuthOptions.ServerKeyRemotePath}} \
-H {{.SwarmOptions.Host}} \
--strategy {{.SwarmOptions.Strategy}} --advertise {{.IP}}:{{.DockerPort}} {{range .SwarmOptions.ArbitraryFlags}} --{{.}}{{end}} {{.SwarmOptions.Discovery}}
`

	swarmWorkerCmdTemplate := `sudo docker run -d \
--restart=always \
{{range .Env}} -e {{.}}{{end}} \
--name swarm-agent \
{{.SwarmImage}} \
join --advertise {{.IP}}:{{.DockerPort}} {{.SwarmOptions.Discovery}}
`

	if swarmOptions.Master {
		log.Debug("Launching swarm master")
		if err := runSwarmCommandFromTemplate(p, swarmMasterCmdTemplate, swarmCmdContext); err != nil {
			return err
		}
	}

	log.Debug("Launch swarm worker")
	if err := runSwarmCommandFromTemplate(p, swarmWorkerCmdTemplate, swarmCmdContext); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package tlsutil

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"
	"time"
)

// Wrapper is a function that is used to wrap a non-TLS connection
// and returns an appropriate TLS connection or error
type Wrapper func(net.Conn) (net.Conn, error)

// Config used to create tls.Config
type Config struct {
	// VerifyIncoming is used to verify the authenticity of incoming connections.
	// This means that TCP requests are forbidden, only allowing for TLS. TLS connections
	// must match a provided certificate authority. This can be used to force client auth.
	VerifyIncoming bool

	// VerifyOutgoing is used to verify the authenticity of outgoing connections.
	// This means that TLS requests are used, and TCP requests are not made. TLS connections
	// must match a provided certificate authority. This is used to verify authenticity of
	// server nodes.
	VerifyOutgoing bool

	// CAFile is a path to a certificate authority file. This is used with VerifyIncoming
	// or VerifyOutgoing to verify the TLS connection.
	CAFile string

	// CertFile is used to provide a TLS certificate that is used for serving TLS connections.
	// Must be provided to serve TLS connections.
	CertFile string

	// KeyFile is used to provide a TLS key that is used for serving TLS connections.
	// Must be provided to serve TLS connections.
	KeyFile string

	// Node name is the name we use to advertise. Defaults to hostname.
	NodeName string

	// ServerName is used with the TLS certificate to ensure the name we
	// provide matches the certificate
	ServerName string
}

// AppendCA opens and parses the CA file and adds the certificates to
// the provided CertPool.
func (c *Config) AppendCA(pool *x509.CertPool) error {
	if c.CAFile == "" {
		return nil
	}

	// Read the file
	data, err := ioutil.ReadFile(c.CAFile)
	if err != nil {
		return fmt.Errorf("Failed to read CA file: %v", err)
	}

	if !pool.AppendCertsFromPEM(data) {
		return fmt.Errorf("Failed to parse any CA certificates")
	}

	return nil
}

// KeyPair is used to open and parse a certificate and key file
func (c *Config) KeyPair() (*tls.Certificate, error) {
	if c.CertFile == "" || c.KeyFile == "" {
		return nil, nil
	}
	cert, err := tls.LoadX509KeyPair(c.CertFile, c.KeyFile)
	if err != nil {
		return nil, fmt.Errorf("Failed to load cert/key pair: %v", err)
	}
	return &cert, err
}

// OutgoingTLSConfig generates a TLS configuration for outgoing
// requests. It will return a nil config if this configuration should
// not use TLS for outgoing connections.
func (c *Config) OutgoingTLSConfig() (*tls.Config, error) {
	if !c.VerifyOutgoing {
		return nil, nil
	}
	// Create the tlsConfig
	tlsConfig := &tls.Config{
		RootCAs:            x509.NewCertPool(),
		InsecureSkipVerify: true,
	}
	if c.ServerName != "" {
		tlsConfig.ServerName = c.ServerName
		tlsConfig.InsecureSkipVerify = false
	}

	// Ensure we have a CA if VerifyOutgoing is set
	if c.VerifyOutgoing && c.CAFile == "" {
		return nil, fmt.Errorf("VerifyOutgoing set, and no CA certificate provided!")
	}

	// Parse the CA cert if any
	err := c.AppendCA(tlsConfig.RootCAs)
	if err != nil {
		return nil, err
	}

	// Add cert/key
	cert, err := c.KeyPair()
	if err != nil {
		return nil, err
	} else if cert != nil {
		tlsConfig.Certificates = []tls.Certificate{*cert}
	}

	return tlsConfig, nil
}

// Wrap a net.Conn into a client tls connection, performing any
// additional verification as needed.
//
// As of go 1.3, crypto/tls only supports either doing no certificate
// verification, or doing full verification including of the peer's
// DNS name. For consul, we want to validate that the certificate is
// signed by a known CA, but because consul doesn't use DNS names for
// node names, we don't verify the certificate DNS names. Since go 1.3
// no longer supports this mode of operation, we have to do it
// manually.
func WrapTLSClient(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) {
	var err error
	var tlsConn *tls.Conn

	tlsConn = tls.Client(conn, tlsConfig)

	// If crypto/tls is doing verification, there's no need to do
	// our own.
	if tlsConfig.InsecureSkipVerify == false {
		return tlsConn, nil
	}

	if err = tlsConn.Handshake(); err != nil {
		tlsConn.Close()
		return nil, err
	}

	// The following is lightly-modified from the doFullHandshake
	// method in crypto/tls's handshake_client.go.
	opts := x509.VerifyOptions{
		Roots:         tlsConfig.RootCAs,
		CurrentTime:   time.Now(),
		DNSName:       "",
		Intermediates: x509.NewCertPool(),
	}

	certs := tlsConn.ConnectionState().PeerCertificates
	for i, cert := range certs {
		if i == 0 {
			continue
		}
		opts.Intermediates.AddCert(cert)
	}

	_, err = certs[0].Verify(opts)
	if err != nil {
		tlsConn.Close()
		return nil, err
	}

	return tlsConn, err
}

// IncomingTLSConfig generates a TLS configuration for incoming requests
func (c *Config) IncomingTLSConfig() (*tls.Config, error) {
	// Create the tlsConfig
	tlsConfig := &tls.Config{
		ServerName: c.ServerName,
		ClientCAs:  x509.NewCertPool(),
		ClientAuth: tls.NoClientCert,
	}
	if tlsConfig.ServerName == "" {
		tlsConfig.ServerName = c.NodeName
	}

	// Parse the CA cert if any
	err := c.AppendCA(tlsConfig.ClientCAs)
	if err != nil {
		return nil, err
	}

	// Add cert/key
	cert, err := c.KeyPair()
	if err != nil {
		return nil, err
	} else if cert != nil {
		tlsConfig.Certificates = []tls.Certificate{*cert}
	}

	// Check if we require verification
	if c.VerifyIncoming {
		tlsConfig.ClientAuth = tls.RequireAndVerifyClientCert
		if c.CAFile == "" {
			return nil, fmt.Errorf("VerifyIncoming set, and no CA certificate provided!")
		}
		if cert == nil {
			return nil, fmt.Errorf("VerifyIncoming set, and no Cert/Key pair provided!")
		}
	}
	return tlsConfig, nil
}
<commit_msg>tlsutil: Adding wrappers for hostname verification<commit_after>package tlsutil

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"
	"time"
)

// DCWrapper is a function that is used to wrap a non-TLS connection
// and returns an appropriate TLS connection or error. This takes
// a datacenter as an argument.
type DCWrapper func(dc string, conn net.Conn) (net.Conn, error)

// Wrapper is a variant of DCWrapper, where the DC is provided as
// a constant value. This is usually done by currying DCWrapper.
type Wrapper func(conn net.Conn) (net.Conn, error)

// Config used to create tls.Config
type Config struct {
	// VerifyIncoming is used to verify the authenticity of incoming connections.
	// This means that TCP requests are forbidden, only allowing for TLS. TLS connections
	// must match a provided certificate authority. This can be used to force client auth.
	VerifyIncoming bool

	// VerifyOutgoing is used to verify the authenticity of outgoing connections.
	// This means that TLS requests are used, and TCP requests are not made. TLS connections
	// must match a provided certificate authority. This is used to verify authenticity of
	// server nodes.
	VerifyOutgoing bool

	// VerifyServerHostname is used to enable hostname verification of servers. This
	// ensures that the certificate presented is valid for server.<datacenter>.<domain>.
	// This prevents a compromised client from being restarted as a server, and then
	// intercepting request traffic as well as being added as a raft peer. This should be
	// enabled by default with VerifyOutgoing, but for legacy reasons we cannot break
	// existing clients.
	VerifyServerHostname bool

	// CAFile is a path to a certificate authority file. This is used with VerifyIncoming
	// or VerifyOutgoing to verify the TLS connection.
	CAFile string

	// CertFile is used to provide a TLS certificate that is used for serving TLS connections.
	// Must be provided to serve TLS connections.
	CertFile string

	// KeyFile is used to provide a TLS key that is used for serving TLS connections.
	// Must be provided to serve TLS connections.
	KeyFile string

	// Node name is the name we use to advertise. Defaults to hostname.
	NodeName string

	// ServerName is used with the TLS certificate to ensure the name we
	// provide matches the certificate
	ServerName string

	// Domain is the Consul TLD being used. Defaults to "consul."
	Domain string
}

// AppendCA opens and parses the CA file and adds the certificates to
// the provided CertPool.
func (c *Config) AppendCA(pool *x509.CertPool) error {
	if c.CAFile == "" {
		return nil
	}

	// Read the file
	data, err := ioutil.ReadFile(c.CAFile)
	if err != nil {
		return fmt.Errorf("Failed to read CA file: %v", err)
	}

	if !pool.AppendCertsFromPEM(data) {
		return fmt.Errorf("Failed to parse any CA certificates")
	}

	return nil
}

// KeyPair is used to open and parse a certificate and key file
func (c *Config) KeyPair() (*tls.Certificate, error) {
	if c.CertFile == "" || c.KeyFile == "" {
		return nil, nil
	}
	cert, err := tls.LoadX509KeyPair(c.CertFile, c.KeyFile)
	if err != nil {
		return nil, fmt.Errorf("Failed to load cert/key pair: %v", err)
	}
	return &cert, err
}

// OutgoingTLSConfig generates a TLS configuration for outgoing
// requests. It will return a nil config if this configuration should
// not use TLS for outgoing connections.
func (c *Config) OutgoingTLSConfig() (*tls.Config, error) {
	if !c.VerifyOutgoing {
		return nil, nil
	}
	// Create the tlsConfig
	tlsConfig := &tls.Config{
		RootCAs:            x509.NewCertPool(),
		InsecureSkipVerify: true,
	}
	if c.ServerName != "" {
		tlsConfig.ServerName = c.ServerName
		tlsConfig.InsecureSkipVerify = false
	}
	if c.VerifyServerHostname {
		// ServerName is filled in dynamically based on the target DC
		tlsConfig.ServerName = "VerifyServerHostname"
		tlsConfig.InsecureSkipVerify = false
	}

	// Ensure we have a CA if VerifyOutgoing is set
	if c.VerifyOutgoing && c.CAFile == "" {
		return nil, fmt.Errorf("VerifyOutgoing set, and no CA certificate provided!")
	}

	// Parse the CA cert if any
	err := c.AppendCA(tlsConfig.RootCAs)
	if err != nil {
		return nil, err
	}

	// Add cert/key
	cert, err := c.KeyPair()
	if err != nil {
		return nil, err
	} else if cert != nil {
		tlsConfig.Certificates = []tls.Certificate{*cert}
	}

	return tlsConfig, nil
}

// OutgoingTLSWrapper returns a a DCWrapper based on the OutgoingTLS
// configuration. If hostname verification is on, the wrapper
// will properly generate the dynamic server name for verification.
func (c *Config) OutgoingTLSWrapper() (DCWrapper, error) {
	// Get the TLS config
	tlsConfig, err := c.OutgoingTLSConfig()
	if err != nil {
		return nil, err
	}

	// Check if TLS is not enabled
	if tlsConfig == nil {
		return nil, nil
	}

	// Generate the wrapper based on hostname verification
	if c.VerifyServerHostname {
		wrapper := func(dc string, conn net.Conn) (net.Conn, error) {
			conf := *tlsConfig
			conf.ServerName = "server." + dc + "." + c.Domain
			return WrapTLSClient(conn, &conf)
		}
		return wrapper, nil
	} else {
		wrapper := func(dc string, c net.Conn) (net.Conn, error) {
			return WrapTLSClient(c, tlsConfig)
		}
		return wrapper, nil
	}
}

// SpecificDC is used to invoke a static datacenter
// and turns a DCWrapper into a Wrapper type.
func SpecificDC(dc string, tlsWrap DCWrapper) Wrapper {
	if tlsWrap == nil {
		return nil
	}
	return func(conn net.Conn) (net.Conn, error) {
		return tlsWrap(dc, conn)
	}
}

// Wrap a net.Conn into a client tls connection, performing any
// additional verification as needed.
//
// As of go 1.3, crypto/tls only supports either doing no certificate
// verification, or doing full verification including of the peer's
// DNS name. For consul, we want to validate that the certificate is
// signed by a known CA, but because consul doesn't use DNS names for
// node names, we don't verify the certificate DNS names. Since go 1.3
// no longer supports this mode of operation, we have to do it
// manually.
func WrapTLSClient(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) {
	var err error
	var tlsConn *tls.Conn

	tlsConn = tls.Client(conn, tlsConfig)

	// If crypto/tls is doing verification, there's no need to do
	// our own.
	if tlsConfig.InsecureSkipVerify == false {
		return tlsConn, nil
	}

	if err = tlsConn.Handshake(); err != nil {
		tlsConn.Close()
		return nil, err
	}

	// The following is lightly-modified from the doFullHandshake
	// method in crypto/tls's handshake_client.go.
	opts := x509.VerifyOptions{
		Roots:         tlsConfig.RootCAs,
		CurrentTime:   time.Now(),
		DNSName:       "",
		Intermediates: x509.NewCertPool(),
	}

	certs := tlsConn.ConnectionState().PeerCertificates
	for i, cert := range certs {
		if i == 0 {
			continue
		}
		opts.Intermediates.AddCert(cert)
	}

	_, err = certs[0].Verify(opts)
	if err != nil {
		tlsConn.Close()
		return nil, err
	}

	return tlsConn, err
}

// IncomingTLSConfig generates a TLS configuration for incoming requests
func (c *Config) IncomingTLSConfig() (*tls.Config, error) {
	// Create the tlsConfig
	tlsConfig := &tls.Config{
		ServerName: c.ServerName,
		ClientCAs:  x509.NewCertPool(),
		ClientAuth: tls.NoClientCert,
	}
	if tlsConfig.ServerName == "" {
		tlsConfig.ServerName = c.NodeName
	}

	// Parse the CA cert if any
	err := c.AppendCA(tlsConfig.ClientCAs)
	if err != nil {
		return nil, err
	}

	// Add cert/key
	cert, err := c.KeyPair()
	if err != nil {
		return nil, err
	} else if cert != nil {
		tlsConfig.Certificates = []tls.Certificate{*cert}
	}

	// Check if we require verification
	if c.VerifyIncoming {
		tlsConfig.ClientAuth = tls.RequireAndVerifyClientCert
		if c.CAFile == "" {
			return nil, fmt.Errorf("VerifyIncoming set, and no CA certificate provided!")
		}
		if cert == nil {
			return nil, fmt.Errorf("VerifyIncoming set, and no Cert/Key pair provided!")
		}
	}
	return tlsConfig, nil
}
<|endoftext|>


<commit_before>package scamp

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"io"

	"strings"
	"sync"
	"sync/atomic"
)

type incomingMsgNo uint64
type outgoingMsgNo uint64

// Connection a scamp connection
type Connection struct {
	conn           *tls.Conn
	Fingerprint    string
	readWriter     *bufio.ReadWriter
	readWriterLock sync.Mutex
	incomingmsgno  incomingMsgNo
	outgoingmsgno  outgoingMsgNo
	pktToMsg       map[incomingMsgNo](*Message)
	msgs           chan *Message
	client         *Client
	isClosed       bool
	closedMutex    sync.Mutex
	scampDebugger  *scampDebugger
}

// DialConnection Used by Client to establish a secure connection to the remote service.
// TODO: You must use the *connection.Fingerprint to verify the
// remote host
func DialConnection(connspec string) (conn *Connection, err error) {
	// Trace.Printf("Dialing connection to `%s`", connspec)
	config := &tls.Config{
		InsecureSkipVerify: true,
	}
	config.BuildNameToCertificate()

	tlsConn, err := tls.Dial("tcp", connspec, config)
	if err != nil {
		return
	}
	// Trace.Printf("Past TLS")
	conn = NewConnection(tlsConn, "client")
	return
}

// NewConnection Used by Service
func NewConnection(tlsConn *tls.Conn, connType string) (conn *Connection) {
	conn = new(Connection)
	conn.conn = tlsConn

	// TODO get the end entity certificate instead
	peerCerts := conn.conn.ConnectionState().PeerCertificates
	if len(peerCerts) == 1 {
		peerCert := peerCerts[0]
		conn.Fingerprint = sha1FingerPrint(peerCert)
	}

	var reader io.Reader = conn.conn
	var writer io.Writer = conn.conn
	if enableWriteTee {
		var err error
		conn.scampDebugger, err = newScampDebugger(conn.conn, connType)
		if err != nil {
			panic(fmt.Sprintf("could not create debugger: %s", err))
		}
		// reader = conn.scampDebugger.WrapReader(reader)
		writer = io.MultiWriter(writer, conn.scampDebugger)
		debuggerReaderWriter := scampDebuggerReader{
			wraps: conn.scampDebugger,
		}
		reader = io.TeeReader(reader, &debuggerReaderWriter)
		// nothing
	}

	conn.readWriter = bufio.NewReadWriter(bufio.NewReader(reader), bufio.NewWriter(writer))
	conn.incomingmsgno = 0
	conn.outgoingmsgno = 0

	conn.pktToMsg = make(map[incomingMsgNo](*Message))
	conn.msgs = make(chan *Message)

	conn.isClosed = false
	go conn.packetReader()

	return
}

// SetClient sets the client for a *Connection
func (conn *Connection) SetClient(client *Client) {
	conn.client = client
}

func (conn *Connection) packetReader() (err error) {
	if conn == nil {
		return
	}
	// I think we only need to lock on writes, packetReader is only running
	// from one spot.
	// conn.readWriterLock.Lock()
	// defer conn.readWriterLock.Unlock()
	var pkt *Packet

PacketReaderLoop:
	for {
		// Trace.Printf("reading packet...")

		pkt, err = ReadPacket(conn.readWriter)
		if err != nil {
			// Warning.Printf("Client %v, packet reader go routine %v ReadPacket error %s\n", conn.client.ID, prNum, err)
			if strings.Contains(err.Error(), "readline error: EOF") {
				// Trace.Printf("%s", err)
			} else if strings.Contains(err.Error(), "use of closed network connection") {
				// Trace.Printf("%s", err)
			} else if strings.Contains(err.Error(), "connection reset by peer") {
				// Trace.Printf("%s", err)
			} else {
				// Trace.Printf("%s", err)
				Error.Printf("err: %s", err)
			}
			break PacketReaderLoop
		}

		err = conn.routePacket(pkt)
		if err != nil {
			// Trace.Printf("breaking PacketReaderLoop")
			break PacketReaderLoop
		}
	}

	close(conn.msgs)
	return
}

func (conn *Connection) routePacket(pkt *Packet) (err error) {
	var msg *Message
	// Trace.Printf("routing packet...")
	switch {
	case pkt.packetType == HEADER:
		// Trace.Printf("HEADER")

		incomingmsgno := atomic.LoadUint64((*uint64)(&conn.incomingmsgno))
		if pkt.msgNo != incomingmsgno {
			err = fmt.Errorf("out of sequence msgno: expected %d but got %d", incomingmsgno, pkt.msgNo)
			Error.Printf("%s", err)
			return err
		}

		msg = conn.pktToMsg[incomingMsgNo(pkt.msgNo)]
		if msg != nil {
			err = fmt.Errorf("Bad HEADER; already tracking msgno %d", pkt.msgNo)
			Error.Printf("%s", err)
			return err
		}

		// Allocate message and copy over header values so we don't have to track them
		// We copy out the packetHeader values and then we can discard it
		msg = NewMessage()
		msg.SetAction(pkt.packetHeader.Action)
		msg.SetEnvelope(pkt.packetHeader.Envelope)
		msg.SetVersion(pkt.packetHeader.Version)
		msg.SetMessageType(pkt.packetHeader.MessageType)
		msg.SetRequestID(pkt.packetHeader.RequestID)
		msg.SetError(pkt.packetHeader.Error)
		msg.SetErrorCode(pkt.packetHeader.ErrorCode)
		msg.SetTicket(pkt.packetHeader.Ticket)
		// TODO: Do we need the requestId?

		conn.pktToMsg[incomingMsgNo(pkt.msgNo)] = msg
		// This is for sending out data
		// conn.incomingNotifiers[pktMsgNo] = &make((chan *Message),1)

		atomic.AddUint64((*uint64)(&conn.incomingmsgno), 1)
	case pkt.packetType == DATA:
		// Trace.Printf("DATA")
		// Append data
		// Verify we are tracking that message
		msg = conn.pktToMsg[incomingMsgNo(pkt.msgNo)]
		if msg == nil {
			return fmt.Errorf("not tracking message number %d", pkt.msgNo)
		}

		msg.Write(pkt.body)
		conn.ackBytes(incomingMsgNo(pkt.msgNo), msg.BytesWritten())

	case pkt.packetType == EOF:
		// Trace.Printf("EOF")
		// Deliver message
		msg = conn.pktToMsg[incomingMsgNo(pkt.msgNo)]
		if msg == nil {
			err = fmt.Errorf("cannot process EOF for unknown msgno %d", pkt.msgNo)
			Error.Printf("err: `%s`", err)
			return
		}

		delete(conn.pktToMsg, incomingMsgNo(pkt.msgNo))
		// Trace.Printf("Delivering message number %d up the stack", pkt.msgNo)
		// Trace.Printf("Adding message to channel:")
		conn.msgs <- msg

	case pkt.packetType == TXERR:
		msg = conn.pktToMsg[incomingMsgNo(pkt.msgNo)]
		if msg == nil {
			err = fmt.Errorf("cannot process EOF for unknown msgno %d", pkt.msgNo)
			Error.Printf("err: `%s`", err)
			return
		}
		//get the error
		if len(pkt.body) > 0 {
			// Trace.Printf("getting error from packet body: %s", pkt.body)
			errMessage := string(pkt.body)
			msg.Error = errMessage
		} else {
			msg.Error = "There was an unkown error with the connection"
		}
		msg.Write(pkt.body)
		conn.ackBytes(incomingMsgNo(pkt.msgNo), msg.BytesWritten())

		delete(conn.pktToMsg, incomingMsgNo(pkt.msgNo))
		conn.msgs <- msg

	case pkt.packetType == ACK:
		// Trace.Printf("ACK `%v` for msgno %v", len(pkt.body), pkt.msgNo)
		// panic("Xavier needs to support this")
		// TODO: Add bytes to message stream tally
	}

	return
}

const RetryLimit = 50

// Send sends a scamp message using the current *Connection
func (conn *Connection) Send(msg *Message) (err error) {\
	if conn == nil {
		return fmt.Errorf("cannot send on nil connection")
	}
	if conn.isClosed {
		err = fmt.Errorf("connection already closed")
	}

	conn.readWriterLock.Lock()
	defer conn.readWriterLock.Unlock()
	if msg.RequestID == 0 {
		err = fmt.Errorf("must specify `ReqestId` on msg before sending")
		return
	}

	outgoingmsgno := atomic.LoadUint64((*uint64)(&conn.outgoingmsgno))
	atomic.AddUint64((*uint64)(&conn.outgoingmsgno), 1)

	// Trace.Printf("sending msgno %d", outgoingmsgno)

	for _, pkt := range msg.toPackets(outgoingmsgno) {
		// Trace.Printf("sending pkt %d", i)

		retries := 0
		if enableWriteTee {
			writer := io.MultiWriter(conn.readWriter, conn.scampDebugger)
			_, err := pkt.Write(writer)
			conn.scampDebugger.file.Write([]byte("\n"))
			if err != nil {
				Error.Printf("error writing packet: `%s`", err)
				return err
			}
		} else {
			for {
				_, err := pkt.Write(conn.readWriter)
				// TODO: should we actually blacklist this error?
				if err != nil {
					//temprarily
					if strings.Contains(err.Error(), "use of closed connection") {
						err = fmt.Errorf("connection closed")
						break
					}

					if retries > RetryLimit {
						return fmt.Errorf("Retried too many times: %s", err)
					}

					Error.Printf("error writing packet: `%s` (retrying)", err)
					retries += 1
					continue
				}
				break
			}
		}

	}
	conn.readWriter.Flush()
	// Trace.Printf("done sending msg")

	return
}

func (conn *Connection) ackBytes(msgno incomingMsgNo, unackedByteCount uint64) (err error) {
	// Trace.Printf("ACKing msg %v, unacked bytes = %v", msgno, unackedByteCount)
	conn.readWriterLock.Lock()
	defer conn.readWriterLock.Unlock()

	ackPacket := Packet{
		packetType: ACK,
		msgNo:      uint64(msgno),
		body:       []byte(fmt.Sprintf("%d", unackedByteCount)),
	}

	var thisWriter io.Writer
	if enableWriteTee {
		thisWriter = io.MultiWriter(conn.readWriter, conn.scampDebugger)
	} else {
		thisWriter = conn.readWriter
	}

	_, err = ackPacket.Write(thisWriter)
	if err != nil {
		return err
	}

	conn.readWriter.Flush()

	return
}

// Close closes the current *Connection
func (conn *Connection) Close() {
	conn.closedMutex.Lock()
	if conn.isClosed {
		// Trace.Printf("connection already closed. skipping shutdown.")
		conn.closedMutex.Unlock()
		return
	}

	// Trace.Printf("connection is closing")

	conn.conn.Close()
	// conn.conn = nil

	// conn.readWriterLock.Lock()
	// conn.readWriter.Flush()
	// conn.readWriterLock.Unlock()

	conn.isClosed = true
	conn.closedMutex.Unlock()
}
<commit_msg>removed illegal character (typo)<commit_after>package scamp

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"io"

	"strings"
	"sync"
	"sync/atomic"
)

type incomingMsgNo uint64
type outgoingMsgNo uint64

// Connection a scamp connection
type Connection struct {
	conn           *tls.Conn
	Fingerprint    string
	readWriter     *bufio.ReadWriter
	readWriterLock sync.Mutex
	incomingmsgno  incomingMsgNo
	outgoingmsgno  outgoingMsgNo
	pktToMsg       map[incomingMsgNo](*Message)
	msgs           chan *Message
	client         *Client
	isClosed       bool
	closedMutex    sync.Mutex
	scampDebugger  *scampDebugger
}

// DialConnection Used by Client to establish a secure connection to the remote service.
// TODO: You must use the *connection.Fingerprint to verify the
// remote host
func DialConnection(connspec string) (conn *Connection, err error) {
	// Trace.Printf("Dialing connection to `%s`", connspec)
	config := &tls.Config{
		InsecureSkipVerify: true,
	}
	config.BuildNameToCertificate()

	tlsConn, err := tls.Dial("tcp", connspec, config)
	if err != nil {
		return
	}
	// Trace.Printf("Past TLS")
	conn = NewConnection(tlsConn, "client")
	return
}

// NewConnection Used by Service
func NewConnection(tlsConn *tls.Conn, connType string) (conn *Connection) {
	conn = new(Connection)
	conn.conn = tlsConn

	// TODO get the end entity certificate instead
	peerCerts := conn.conn.ConnectionState().PeerCertificates
	if len(peerCerts) == 1 {
		peerCert := peerCerts[0]
		conn.Fingerprint = sha1FingerPrint(peerCert)
	}

	var reader io.Reader = conn.conn
	var writer io.Writer = conn.conn
	if enableWriteTee {
		var err error
		conn.scampDebugger, err = newScampDebugger(conn.conn, connType)
		if err != nil {
			panic(fmt.Sprintf("could not create debugger: %s", err))
		}
		// reader = conn.scampDebugger.WrapReader(reader)
		writer = io.MultiWriter(writer, conn.scampDebugger)
		debuggerReaderWriter := scampDebuggerReader{
			wraps: conn.scampDebugger,
		}
		reader = io.TeeReader(reader, &debuggerReaderWriter)
		// nothing
	}

	conn.readWriter = bufio.NewReadWriter(bufio.NewReader(reader), bufio.NewWriter(writer))
	conn.incomingmsgno = 0
	conn.outgoingmsgno = 0

	conn.pktToMsg = make(map[incomingMsgNo](*Message))
	conn.msgs = make(chan *Message)

	conn.isClosed = false
	go conn.packetReader()

	return
}

// SetClient sets the client for a *Connection
func (conn *Connection) SetClient(client *Client) {
	conn.client = client
}

func (conn *Connection) packetReader() (err error) {
	if conn == nil {
		return
	}
	// I think we only need to lock on writes, packetReader is only running
	// from one spot.
	// conn.readWriterLock.Lock()
	// defer conn.readWriterLock.Unlock()
	var pkt *Packet

PacketReaderLoop:
	for {
		// Trace.Printf("reading packet...")

		pkt, err = ReadPacket(conn.readWriter)
		if err != nil {
			// Warning.Printf("Client %v, packet reader go routine %v ReadPacket error %s\n", conn.client.ID, prNum, err)
			if strings.Contains(err.Error(), "readline error: EOF") {
				// Trace.Printf("%s", err)
			} else if strings.Contains(err.Error(), "use of closed network connection") {
				// Trace.Printf("%s", err)
			} else if strings.Contains(err.Error(), "connection reset by peer") {
				// Trace.Printf("%s", err)
			} else {
				// Trace.Printf("%s", err)
				Error.Printf("err: %s", err)
			}
			break PacketReaderLoop
		}

		err = conn.routePacket(pkt)
		if err != nil {
			// Trace.Printf("breaking PacketReaderLoop")
			break PacketReaderLoop
		}
	}

	close(conn.msgs)
	return
}

func (conn *Connection) routePacket(pkt *Packet) (err error) {
	var msg *Message
	// Trace.Printf("routing packet...")
	switch {
	case pkt.packetType == HEADER:
		// Trace.Printf("HEADER")

		incomingmsgno := atomic.LoadUint64((*uint64)(&conn.incomingmsgno))
		if pkt.msgNo != incomingmsgno {
			err = fmt.Errorf("out of sequence msgno: expected %d but got %d", incomingmsgno, pkt.msgNo)
			Error.Printf("%s", err)
			return err
		}

		msg = conn.pktToMsg[incomingMsgNo(pkt.msgNo)]
		if msg != nil {
			err = fmt.Errorf("Bad HEADER; already tracking msgno %d", pkt.msgNo)
			Error.Printf("%s", err)
			return err
		}

		// Allocate message and copy over header values so we don't have to track them
		// We copy out the packetHeader values and then we can discard it
		msg = NewMessage()
		msg.SetAction(pkt.packetHeader.Action)
		msg.SetEnvelope(pkt.packetHeader.Envelope)
		msg.SetVersion(pkt.packetHeader.Version)
		msg.SetMessageType(pkt.packetHeader.MessageType)
		msg.SetRequestID(pkt.packetHeader.RequestID)
		msg.SetError(pkt.packetHeader.Error)
		msg.SetErrorCode(pkt.packetHeader.ErrorCode)
		msg.SetTicket(pkt.packetHeader.Ticket)
		// TODO: Do we need the requestId?

		conn.pktToMsg[incomingMsgNo(pkt.msgNo)] = msg
		// This is for sending out data
		// conn.incomingNotifiers[pktMsgNo] = &make((chan *Message),1)

		atomic.AddUint64((*uint64)(&conn.incomingmsgno), 1)
	case pkt.packetType == DATA:
		// Trace.Printf("DATA")
		// Append data
		// Verify we are tracking that message
		msg = conn.pktToMsg[incomingMsgNo(pkt.msgNo)]
		if msg == nil {
			return fmt.Errorf("not tracking message number %d", pkt.msgNo)
		}

		msg.Write(pkt.body)
		conn.ackBytes(incomingMsgNo(pkt.msgNo), msg.BytesWritten())

	case pkt.packetType == EOF:
		// Trace.Printf("EOF")
		// Deliver message
		msg = conn.pktToMsg[incomingMsgNo(pkt.msgNo)]
		if msg == nil {
			err = fmt.Errorf("cannot process EOF for unknown msgno %d", pkt.msgNo)
			Error.Printf("err: `%s`", err)
			return
		}

		delete(conn.pktToMsg, incomingMsgNo(pkt.msgNo))
		// Trace.Printf("Delivering message number %d up the stack", pkt.msgNo)
		// Trace.Printf("Adding message to channel:")
		conn.msgs <- msg

	case pkt.packetType == TXERR:
		msg = conn.pktToMsg[incomingMsgNo(pkt.msgNo)]
		if msg == nil {
			err = fmt.Errorf("cannot process EOF for unknown msgno %d", pkt.msgNo)
			Error.Printf("err: `%s`", err)
			return
		}
		//get the error
		if len(pkt.body) > 0 {
			// Trace.Printf("getting error from packet body: %s", pkt.body)
			errMessage := string(pkt.body)
			msg.Error = errMessage
		} else {
			msg.Error = "There was an unkown error with the connection"
		}
		msg.Write(pkt.body)
		conn.ackBytes(incomingMsgNo(pkt.msgNo), msg.BytesWritten())

		delete(conn.pktToMsg, incomingMsgNo(pkt.msgNo))
		conn.msgs <- msg

	case pkt.packetType == ACK:
		// Trace.Printf("ACK `%v` for msgno %v", len(pkt.body), pkt.msgNo)
		// panic("Xavier needs to support this")
		// TODO: Add bytes to message stream tally
	}

	return
}

const RetryLimit = 50

// Send sends a scamp message using the current *Connection
func (conn *Connection) Send(msg *Message) (err error) {
	if conn == nil {
		return fmt.Errorf("cannot send on nil connection")
	}
	if conn.isClosed {
		err = fmt.Errorf("connection already closed")
	}

	conn.readWriterLock.Lock()
	defer conn.readWriterLock.Unlock()
	if msg.RequestID == 0 {
		err = fmt.Errorf("must specify `ReqestId` on msg before sending")
		return
	}

	outgoingmsgno := atomic.LoadUint64((*uint64)(&conn.outgoingmsgno))
	atomic.AddUint64((*uint64)(&conn.outgoingmsgno), 1)

	// Trace.Printf("sending msgno %d", outgoingmsgno)

	for _, pkt := range msg.toPackets(outgoingmsgno) {
		// Trace.Printf("sending pkt %d", i)

		retries := 0
		if enableWriteTee {
			writer := io.MultiWriter(conn.readWriter, conn.scampDebugger)
			_, err := pkt.Write(writer)
			conn.scampDebugger.file.Write([]byte("\n"))
			if err != nil {
				Error.Printf("error writing packet: `%s`", err)
				return err
			}
		} else {
			for {
				_, err := pkt.Write(conn.readWriter)
				// TODO: should we actually blacklist this error?
				if err != nil {
					//temprarily
					if strings.Contains(err.Error(), "use of closed connection") {
						err = fmt.Errorf("connection closed")
						break
					}

					if retries > RetryLimit {
						return fmt.Errorf("Retried too many times: %s", err)
					}

					Error.Printf("error writing packet: `%s` (retrying)", err)
					retries += 1
					continue
				}
				break
			}
		}

	}
	conn.readWriter.Flush()
	// Trace.Printf("done sending msg")

	return
}

func (conn *Connection) ackBytes(msgno incomingMsgNo, unackedByteCount uint64) (err error) {
	// Trace.Printf("ACKing msg %v, unacked bytes = %v", msgno, unackedByteCount)
	conn.readWriterLock.Lock()
	defer conn.readWriterLock.Unlock()

	ackPacket := Packet{
		packetType: ACK,
		msgNo:      uint64(msgno),
		body:       []byte(fmt.Sprintf("%d", unackedByteCount)),
	}

	var thisWriter io.Writer
	if enableWriteTee {
		thisWriter = io.MultiWriter(conn.readWriter, conn.scampDebugger)
	} else {
		thisWriter = conn.readWriter
	}

	_, err = ackPacket.Write(thisWriter)
	if err != nil {
		return err
	}

	conn.readWriter.Flush()

	return
}

// Close closes the current *Connection
func (conn *Connection) Close() {
	conn.closedMutex.Lock()
	if conn.isClosed {
		// Trace.Printf("connection already closed. skipping shutdown.")
		conn.closedMutex.Unlock()
		return
	}

	// Trace.Printf("connection is closing")

	conn.conn.Close()
	// conn.conn = nil

	// conn.readWriterLock.Lock()
	// conn.readWriter.Flush()
	// conn.readWriterLock.Unlock()

	conn.isClosed = true
	conn.closedMutex.Unlock()
}
<|endoftext|>


<commit_before>package scheduler

import (
	"sync/atomic"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/cache"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/mesos"
	"github.com/mesosphere/kubernetes-mesos/queue"
)

const (
	offerListenerMaxAge      = 5 // max number of times we'll attempt to fit an offer to a listener before requireing them to re-register themselves
	deferredDeclineTtlFactor = 2 // this factor, multiplied by the offer ttl, determines how long to wait before attempting to decline previously claimed offers that were subsequently deleted, then released. see offerStorage.Delete
)

type OfferFilter func(*mesos.Offer) bool

type OfferRegistry interface {
	// Initialize the instance, spawning necessary housekeeping go routines.
	Init()
	Add([]*mesos.Offer)

	// Listen for arriving offers that are acceptable to the filter, sending
	// a signal on (by closing) the returned channel. A listener will only
	// ever be notified once, if at all.
	Listen(id string, f OfferFilter) <-chan empty

	// invoked when offers are rescinded or expired
	Delete(string)

	Get(offerId string) (PerishableOffer, bool)

	Walk(Walker) error

	// invalidate one or all (when offerId="") offers; offers are not declined,
	// but are simply flagged as expired in the offer history
	Invalidate(offerId string)
}

// callback that is invoked during a walk through a series of live offers,
// returning with stop=true (or err != nil) if the walk should stop permaturely.
type Walker func(offer PerishableOffer) (stop bool, err error)

type OfferRegistryConfig struct {
	declineOffer  func(offerId string) error
	ttl           time.Duration // determines a perishable offer's expiration deadline: now+ttl
	lingerTtl     time.Duration // if zero, offers will not linger in the FIFO past their expiration deadline
	listenerDelay time.Duration // specifies the sleep time between offer listener notifications
}

type offerStorage struct {
	OfferRegistryConfig
	offers    *cache.FIFO       // collection of PerishableOffer, both live and expired
	listeners *cache.FIFO       // collection of *offerListener
	delayed   *queue.DelayQueue // deadline-oriented offer-event queue
}

type liveOffer struct {
	*mesos.Offer
	expiration time.Time
	acquired   int32 // 1 = acquired, 0 = free
}

type expiredOffer struct {
	id       string
	deadline time.Time
}

type PerishableOffer interface {
	queue.Delayed
	// returns true if this offer has expired
	HasExpired() bool
	// if not yet expired, return mesos offer details; otherwise nil
	Details() *mesos.Offer
	// mark this offer as acquired, returning true if it was previously unacquired. thread-safe.
	Acquire() bool
	// mark this offer as un-acquired. thread-safe.
	Release()
	// expire or delete this offer from storage
	age(s *offerStorage)
}

func (e *expiredOffer) HasExpired() bool {
	return true
}

func (e *expiredOffer) Details() *mesos.Offer {
	return nil
}

func (e *expiredOffer) Acquire() bool {
	return false
}

func (e *expiredOffer) Release() {}

func (e *expiredOffer) age(s *offerStorage) {
	log.V(3).Infof("Delete lingering offer: %v", e.id)
	s.offers.Delete(e.id)
}

// return the time left to linger
func (e *expiredOffer) GetDelay() time.Duration {
	return e.deadline.Sub(time.Now())
}

func (to *liveOffer) HasExpired() bool {
	return time.Now().After(to.expiration)
}

func (to *liveOffer) Details() *mesos.Offer {
	return to.Offer
}

func (to *liveOffer) Acquire() bool {
	return atomic.CompareAndSwapInt32(&to.acquired, 0, 1)
}

func (to *liveOffer) Release() {
	atomic.CompareAndSwapInt32(&to.acquired, 1, 0)
}

func (to *liveOffer) age(s *offerStorage) {
	s.Delete(to.Offer.Id.GetValue())
}

// return the time remaining before the offer expires
func (to *liveOffer) GetDelay() time.Duration {
	return to.expiration.Sub(time.Now())
}

func CreateOfferRegistry(c OfferRegistryConfig) OfferRegistry {
	return &offerStorage{c, cache.NewFIFO(), cache.NewFIFO(), queue.NewDelayQueue()}
}

func (s *offerStorage) Add(offers []*mesos.Offer) {
	now := time.Now()
	for _, offer := range offers {
		offerId := offer.Id.GetValue()
		log.V(3).Infof("Receiving offer %v", offerId)
		timed := &liveOffer{offer, now.Add(s.ttl), 0}
		s.offers.Add(offerId, timed)
		s.delayed.Add(timed)
	}
}

// delete an offer from storage, meaning that we expire it
func (s *offerStorage) Delete(offerId string) {
	if offer, ok := s.Get(offerId); ok {
		log.V(3).Infof("Deleting offer %v", offerId)
		// attempt to block others from consuming the offer. if it's already been
		// claimed and is not yet lingering then don't decline it - just mark it as
		// expired in the history: allow a prior claimant to attempt to launch with it
		myoffer := offer.Acquire()
		if offer.Details() != nil {
			if myoffer {
				log.V(3).Infof("Declining offer %v", offerId)
				if err := s.declineOffer(offerId); err != nil {
					log.Warningf("Failed to decline offer %v: %v", offerId, err)
				}
			} else {
				// some pod has acquired this and may attempt to launch a task with it
				// failed schedule/launch attempts are requried to Release() any claims on the offer
				go func() {
					// TODO(jdef): not sure what a good value is here. the goal is to provide a
					// launchTasks (driver) operation enough time to complete so that we don't end
					// up delining an offer that we're actually attempting to use.
					time.Sleep(deferredDeclineTtlFactor * s.ttl)

					// at this point the offer is in one of five states:
					// a) permanently deleted: expired due to timeout
					// b) permanently deleted: expired due to having been rescinded
					// c) lingering: expired due to timeout
					// d) lingering: expired due to having been rescinded
					// e) claimed: task launched and it using resources from this offer
					// we want to **avoid** declining an offer that's claimed: attempt to acquire
					if offer.Acquire() {
						// previously claimed offer was released, perhaps due to a launch
						// failure, so we should attempt to decline
						if err := s.declineOffer(offerId); err != nil {
							log.Warningf("Failed to decline (previously claimed) offer %v: %v", offerId, err)
						}
					}
				}()
			}
		}
		s.expireOffer(offer)
	} // else, ignore offers not in the history
}

// expire all known, live offers
func (s *offerStorage) Invalidate(offerId string) {
	if offerId != "" {
		s.invalidateOne(offerId)
		return
	}
	obj := s.offers.List()
	for _, o := range obj {
		offer, ok := o.(PerishableOffer)
		if !ok {
			log.Errorf("Expected perishable offer, not %v", o)
			continue
		}
		offer.Acquire() // attempt to block others from using it
		s.expireOffer(offer)
		// don't reject it from the controller, we already know that it's an invalid offer
	}
}

func (s *offerStorage) invalidateOne(offerId string) {
	if offer, ok := s.Get(offerId); ok {
		offer.Acquire() // attempt to block others from using it
		s.expireOffer(offer)
		// don't reject it from the controller, we already know that it's an invalid offer
	}
}

// Walk the collection of offers. The walk stops either as indicated by the
// Walker or when the end of the offer list is reached. Expired offers are
// never passed to a Walker.
func (s *offerStorage) Walk(w Walker) error {
	for offerId := range s.offers.Contains() {
		offer, ok := s.Get(offerId)
		if !ok {
			// offer disappeared...
			continue
		}
		if offer.HasExpired() {
			// never pass expired offers to walkers
			continue
		}
		if stop, err := w(offer); err != nil {
			return err
		} else if stop {
			return nil
		}
	}
	return nil
}

func (s *offerStorage) expireOffer(offer PerishableOffer) {
	// the offer may or may not be expired due to TTL so check for details
	// since that's a more reliable determinant of lingering status
	if details := offer.Details(); details != nil {
		// recently expired, should linger
		offerId := details.Id.GetValue()
		log.V(3).Infof("Expiring offer %v", offerId)
		if s.lingerTtl > 0 {
			log.V(3).Infof("offer will linger: %v", offerId)
			expired := &expiredOffer{offerId, time.Now().Add(s.lingerTtl)}
			s.offers.Update(offerId, expired)
			s.delayed.Add(expired)
		} else {
			log.V(3).Infof("Permanently deleting offer %v", offerId)
			s.offers.Delete(offerId)
		}
	} // else, it's still lingering...
}

func (s *offerStorage) Get(id string) (PerishableOffer, bool) {
	if obj, ok := s.offers.Get(id); !ok {
		return nil, false
	} else {
		to, ok := obj.(PerishableOffer)
		if !ok {
			log.Errorf("invalid offer object in fifo '%v'", obj)
		}
		return to, ok
	}
}

type offerListener struct {
	id      string
	accepts OfferFilter
	notify  chan<- empty
	age     int
}

// register a listener for new offers, whom we'll notify upon receiving such.
// notification is delivered in the form of closing the channel, nothing is ever sent.
func (s *offerStorage) Listen(id string, f OfferFilter) <-chan empty {
	if f == nil {
		return nil
	}
	ch := make(chan empty)
	listen := &offerListener{
		id:      id,
		accepts: f,
		notify:  ch,
		age:     0,
	}
	log.V(3).Infof("Registering offer listener %s", listen.id)
	s.listeners.Add(id, listen)
	return ch
}

func (s *offerStorage) ageOffers() {
	offer, ok := s.delayed.Pop().(PerishableOffer)
	if !ok {
		log.Errorf("Expected PerishableOffer, not %v", offer)
		return
	}
	if details := offer.Details(); details != nil && !offer.HasExpired() {
		// live offer has not expired yet: timed out early
		// FWIW: early timeouts are more frequent when GOMAXPROCS is > 1
		s.delayed.Add(offer)
	} else {
		offer.age(s)
	}
}

func (s *offerStorage) notifyListeners() {
	var listen *offerListener
	var ok bool
	// get the next offer listener
	for {
		obj := s.listeners.Pop()
		if listen, ok = obj.(*offerListener); ok {
			break
		}
		log.Warningf("unexpected listener object %v", obj)
	}
	// notify if we find an acceptable offer
	for id := range s.offers.Contains() {
		var offer PerishableOffer
		if offer, ok = s.Get(id); !ok || offer.HasExpired() {
			continue
		}
		if listen.accepts(offer.Details()) {
			log.V(3).Infof("Notifying offer listener %s", listen.id)
			close(listen.notify)
			return
		}
	}
	// no interesting offers found, re-queue the listener
	listen.age++
	if listen.age < offerListenerMaxAge {
		// if the same listener has re-registered in the meantime we don't want to
		// destroy the newer listener channel. this is racy since a newer listener
		// can register between the Get() and Update(), but the consequences aren't
		// very dire - the listener merely has to wait their full backoff period.
		if _, ok := s.listeners.Get(listen.id); !ok {
			log.V(3).Infof("Re-registering offer listener %s", listen.id)
			s.listeners.Update(listen.id, listen)
		}
	} // else, you're gc'd
}

func (s *offerStorage) Init() {
	go util.Forever(s.ageOffers, 0)

	// to avoid a rush on offers, we add a short delay between each registered
	// listener, so as to allow the most recently notified listener a bit of time
	// to act on the offer.
	go util.Forever(s.notifyListeners, s.listenerDelay)
}
<commit_msg>updated comments<commit_after>package scheduler

import (
	"sync/atomic"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/cache"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/mesos"
	"github.com/mesosphere/kubernetes-mesos/queue"
)

const (
	offerListenerMaxAge      = 5 // max number of times we'll attempt to fit an offer to a listener before requiring them to re-register themselves
	deferredDeclineTtlFactor = 2 // this factor, multiplied by the offer ttl, determines how long to wait before attempting to decline previously claimed offers that were subsequently deleted, then released. see offerStorage.Delete
)

type OfferFilter func(*mesos.Offer) bool

type OfferRegistry interface {
	// Initialize the instance, spawning necessary housekeeping go routines.
	Init()
	Add([]*mesos.Offer)

	// Listen for arriving offers that are acceptable to the filter, sending
	// a signal on (by closing) the returned channel. A listener will only
	// ever be notified once, if at all.
	Listen(id string, f OfferFilter) <-chan empty

	// invoked when offers are rescinded or expired
	Delete(string)

	Get(offerId string) (PerishableOffer, bool)

	Walk(Walker) error

	// invalidate one or all (when offerId="") offers; offers are not declined,
	// but are simply flagged as expired in the offer history
	Invalidate(offerId string)
}

// callback that is invoked during a walk through a series of live offers,
// returning with stop=true (or err != nil) if the walk should stop permaturely.
type Walker func(offer PerishableOffer) (stop bool, err error)

type OfferRegistryConfig struct {
	declineOffer  func(offerId string) error
	ttl           time.Duration // determines a perishable offer's expiration deadline: now+ttl
	lingerTtl     time.Duration // if zero, offers will not linger in the FIFO past their expiration deadline
	listenerDelay time.Duration // specifies the sleep time between offer listener notifications
}

type offerStorage struct {
	OfferRegistryConfig
	offers    *cache.FIFO       // collection of PerishableOffer, both live and expired
	listeners *cache.FIFO       // collection of *offerListener
	delayed   *queue.DelayQueue // deadline-oriented offer-event queue
}

type liveOffer struct {
	*mesos.Offer
	expiration time.Time
	acquired   int32 // 1 = acquired, 0 = free
}

type expiredOffer struct {
	id       string
	deadline time.Time
}

type PerishableOffer interface {
	queue.Delayed
	// returns true if this offer has expired
	HasExpired() bool
	// if not yet expired, return mesos offer details; otherwise nil
	Details() *mesos.Offer
	// mark this offer as acquired, returning true if it was previously unacquired. thread-safe.
	Acquire() bool
	// mark this offer as un-acquired. thread-safe.
	Release()
	// expire or delete this offer from storage
	age(s *offerStorage)
}

func (e *expiredOffer) HasExpired() bool {
	return true
}

func (e *expiredOffer) Details() *mesos.Offer {
	return nil
}

func (e *expiredOffer) Acquire() bool {
	return false
}

func (e *expiredOffer) Release() {}

func (e *expiredOffer) age(s *offerStorage) {
	log.V(3).Infof("Delete lingering offer: %v", e.id)
	s.offers.Delete(e.id)
}

// return the time left to linger
func (e *expiredOffer) GetDelay() time.Duration {
	return e.deadline.Sub(time.Now())
}

func (to *liveOffer) HasExpired() bool {
	return time.Now().After(to.expiration)
}

func (to *liveOffer) Details() *mesos.Offer {
	return to.Offer
}

func (to *liveOffer) Acquire() bool {
	return atomic.CompareAndSwapInt32(&to.acquired, 0, 1)
}

func (to *liveOffer) Release() {
	atomic.CompareAndSwapInt32(&to.acquired, 1, 0)
}

func (to *liveOffer) age(s *offerStorage) {
	s.Delete(to.Offer.Id.GetValue())
}

// return the time remaining before the offer expires
func (to *liveOffer) GetDelay() time.Duration {
	return to.expiration.Sub(time.Now())
}

func CreateOfferRegistry(c OfferRegistryConfig) OfferRegistry {
	return &offerStorage{c, cache.NewFIFO(), cache.NewFIFO(), queue.NewDelayQueue()}
}

func (s *offerStorage) Add(offers []*mesos.Offer) {
	now := time.Now()
	for _, offer := range offers {
		offerId := offer.Id.GetValue()
		log.V(3).Infof("Receiving offer %v", offerId)
		timed := &liveOffer{offer, now.Add(s.ttl), 0}
		s.offers.Add(offerId, timed)
		s.delayed.Add(timed)
	}
}

// delete an offer from storage, meaning that we expire it
func (s *offerStorage) Delete(offerId string) {
	if offer, ok := s.Get(offerId); ok {
		log.V(3).Infof("Deleting offer %v", offerId)
		// attempt to block others from consuming the offer. if it's already been
		// claimed and is not yet lingering then don't decline it - just mark it as
		// expired in the history: allow a prior claimant to attempt to launch with it
		myoffer := offer.Acquire()
		if offer.Details() != nil {
			if myoffer {
				log.V(3).Infof("Declining offer %v", offerId)
				if err := s.declineOffer(offerId); err != nil {
					log.Warningf("Failed to decline offer %v: %v", offerId, err)
				}
			} else {
				// some pod has acquired this and may attempt to launch a task with it
				// failed schedule/launch attempts are requried to Release() any claims on the offer
				go func() {
					// TODO(jdef): not sure what a good value is here. the goal is to provide a
					// launchTasks (driver) operation enough time to complete so that we don't end
					// up declining an offer that we're actually attempting to use.
					time.Sleep(deferredDeclineTtlFactor * s.ttl)

					// at this point the offer is in one of five states:
					// a) permanently deleted: expired due to timeout
					// b) permanently deleted: expired due to having been rescinded
					// c) lingering: expired due to timeout
					// d) lingering: expired due to having been rescinded
					// e) claimed: task launched and it using resources from this offer
					// we want to **avoid** declining an offer that's claimed: attempt to acquire
					if offer.Acquire() {
						// previously claimed offer was released, perhaps due to a launch
						// failure, so we should attempt to decline
						if err := s.declineOffer(offerId); err != nil {
							log.Warningf("Failed to decline (previously claimed) offer %v: %v", offerId, err)
						}
					}
				}()
			}
		}
		s.expireOffer(offer)
	} // else, ignore offers not in the history
}

// expire all known, live offers
func (s *offerStorage) Invalidate(offerId string) {
	if offerId != "" {
		s.invalidateOne(offerId)
		return
	}
	obj := s.offers.List()
	for _, o := range obj {
		offer, ok := o.(PerishableOffer)
		if !ok {
			log.Errorf("Expected perishable offer, not %v", o)
			continue
		}
		offer.Acquire() // attempt to block others from using it
		s.expireOffer(offer)
		// don't decline, we already know that it's an invalid offer
	}
}

func (s *offerStorage) invalidateOne(offerId string) {
	if offer, ok := s.Get(offerId); ok {
		offer.Acquire() // attempt to block others from using it
		s.expireOffer(offer)
		// don't decline, we already know that it's an invalid offer
	}
}

// Walk the collection of offers. The walk stops either as indicated by the
// Walker or when the end of the offer list is reached. Expired offers are
// never passed to a Walker.
func (s *offerStorage) Walk(w Walker) error {
	for offerId := range s.offers.Contains() {
		offer, ok := s.Get(offerId)
		if !ok {
			// offer disappeared...
			continue
		}
		if offer.HasExpired() {
			// never pass expired offers to walkers
			continue
		}
		if stop, err := w(offer); err != nil {
			return err
		} else if stop {
			return nil
		}
	}
	return nil
}

func (s *offerStorage) expireOffer(offer PerishableOffer) {
	// the offer may or may not be expired due to TTL so check for details
	// since that's a more reliable determinant of lingering status
	if details := offer.Details(); details != nil {
		// recently expired, should linger
		offerId := details.Id.GetValue()
		log.V(3).Infof("Expiring offer %v", offerId)
		if s.lingerTtl > 0 {
			log.V(3).Infof("offer will linger: %v", offerId)
			expired := &expiredOffer{offerId, time.Now().Add(s.lingerTtl)}
			s.offers.Update(offerId, expired)
			s.delayed.Add(expired)
		} else {
			log.V(3).Infof("Permanently deleting offer %v", offerId)
			s.offers.Delete(offerId)
		}
	} // else, it's still lingering...
}

func (s *offerStorage) Get(id string) (PerishableOffer, bool) {
	if obj, ok := s.offers.Get(id); !ok {
		return nil, false
	} else {
		to, ok := obj.(PerishableOffer)
		if !ok {
			log.Errorf("invalid offer object in fifo '%v'", obj)
		}
		return to, ok
	}
}

type offerListener struct {
	id      string
	accepts OfferFilter
	notify  chan<- empty
	age     int
}

// register a listener for new offers, whom we'll notify upon receiving such.
// notification is delivered in the form of closing the channel, nothing is ever sent.
func (s *offerStorage) Listen(id string, f OfferFilter) <-chan empty {
	if f == nil {
		return nil
	}
	ch := make(chan empty)
	listen := &offerListener{
		id:      id,
		accepts: f,
		notify:  ch,
		age:     0,
	}
	log.V(3).Infof("Registering offer listener %s", listen.id)
	s.listeners.Add(id, listen)
	return ch
}

func (s *offerStorage) ageOffers() {
	offer, ok := s.delayed.Pop().(PerishableOffer)
	if !ok {
		log.Errorf("Expected PerishableOffer, not %v", offer)
		return
	}
	if details := offer.Details(); details != nil && !offer.HasExpired() {
		// live offer has not expired yet: timed out early
		// FWIW: early timeouts are more frequent when GOMAXPROCS is > 1
		s.delayed.Add(offer)
	} else {
		offer.age(s)
	}
}

func (s *offerStorage) notifyListeners() {
	var listen *offerListener
	var ok bool
	// get the next offer listener
	for {
		obj := s.listeners.Pop()
		if listen, ok = obj.(*offerListener); ok {
			break
		}
		log.Warningf("unexpected listener object %v", obj)
	}
	// notify if we find an acceptable offer
	for id := range s.offers.Contains() {
		var offer PerishableOffer
		if offer, ok = s.Get(id); !ok || offer.HasExpired() {
			continue
		}
		if listen.accepts(offer.Details()) {
			log.V(3).Infof("Notifying offer listener %s", listen.id)
			close(listen.notify)
			return
		}
	}
	// no interesting offers found, re-queue the listener
	listen.age++
	if listen.age < offerListenerMaxAge {
		// if the same listener has re-registered in the meantime we don't want to
		// destroy the newer listener channel. this is racy since a newer listener
		// can register between the Get() and Update(), but the consequences aren't
		// very dire - the listener merely has to wait their full backoff period.
		if _, ok := s.listeners.Get(listen.id); !ok {
			log.V(3).Infof("Re-registering offer listener %s", listen.id)
			s.listeners.Update(listen.id, listen)
		}
	} // else, you're gc'd
}

func (s *offerStorage) Init() {
	go util.Forever(s.ageOffers, 0)

	// to avoid a rush on offers, we add a short delay between each registered
	// listener, so as to allow the most recently notified listener a bit of time
	// to act on the offer.
	go util.Forever(s.notifyListeners, s.listenerDelay)
}
<|endoftext|>


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

package auth

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/user"
	"path/filepath"

	"github.com/coreos/mantle/Godeps/_workspace/src/golang.org/x/oauth2"
)

// client registered under 'marineam-tools'
var conf = oauth2.Config{
	ClientID:     "937427706989-nbndmfkp0knqardoagk6lbcamrsh828i.apps.googleusercontent.com",
	ClientSecret: "F6Xs5wGHZzGw-QFXl3aylLUT",
	Endpoint: oauth2.Endpoint{
		AuthURL:  "https://accounts.google.com/o/oauth2/auth",
		TokenURL: "https://accounts.google.com/o/oauth2/token",
	},
	RedirectURL: "urn:ietf:wg:oauth:2.0:oob",
	Scopes: []string{"https://www.googleapis.com/auth/devstorage.full_control",
		"https://www.googleapis.com/auth/compute"},
}

func writeCache(cachePath string, tok *oauth2.Token) error {
	file, err := os.OpenFile(cachePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer file.Close()

	if err := json.NewEncoder(file).Encode(tok); err != nil {
		return err
	}
	return nil
}

func readCache(cachePath string) (*oauth2.Token, error) {
	file, err := os.Open(cachePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	tok := &oauth2.Token{}
	if err := json.NewDecoder(file).Decode(tok); err != nil {
		return nil, err
	}
	return tok, nil
}

func getToken() (*oauth2.Token, error) {
	userInfo, err := user.Current()
	if err != nil {
		return nil, err
	}

	cachePath := filepath.Join(userInfo.HomeDir, ".mantle-cache-google.json")
	tok, err := readCache(cachePath)
	if err != nil {
		log.Printf("Error reading google token cache file: %v", err)
	}
	if tok == nil || !tok.Valid() {
		url := conf.AuthCodeURL("state", oauth2.AccessTypeOffline)
		fmt.Printf("Visit the URL for the auth dialog: %v\n", url)
		fmt.Print("Enter token: ")

		var code string
		if _, err := fmt.Scan(&code); err != nil {
			return nil, err
		}
		tok, err = conf.Exchange(oauth2.NoContext, code)
		if err != nil {
			return nil, err
		}
		err = writeCache(cachePath, tok)
		if err != nil {
			log.Printf("Error writing google token cache file: %v", err)
		}
	}
	return tok, nil
}

func GoogleClient() (*http.Client, error) {
	tok, err := getToken()
	if err != nil {
		return nil, err
	}
	return conf.Client(oauth2.NoContext, tok), nil
}

func GoogleTokenSource() (oauth2.TokenSource, error) {
	tok, err := getToken()
	if err != nil {
		return nil, err
	}
	return conf.TokenSource(oauth2.NoContext, tok), nil
}
<commit_msg>auth: allow cached tokens to be refreshed<commit_after>// Copyright 2014 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package auth

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/user"
	"path/filepath"

	"github.com/coreos/mantle/Godeps/_workspace/src/golang.org/x/oauth2"
)

// client registered under 'marineam-tools'
var conf = oauth2.Config{
	ClientID:     "937427706989-nbndmfkp0knqardoagk6lbcamrsh828i.apps.googleusercontent.com",
	ClientSecret: "F6Xs5wGHZzGw-QFXl3aylLUT",
	Endpoint: oauth2.Endpoint{
		AuthURL:  "https://accounts.google.com/o/oauth2/auth",
		TokenURL: "https://accounts.google.com/o/oauth2/token",
	},
	RedirectURL: "urn:ietf:wg:oauth:2.0:oob",
	Scopes: []string{"https://www.googleapis.com/auth/devstorage.full_control",
		"https://www.googleapis.com/auth/compute"},
}

func writeCache(cachePath string, tok *oauth2.Token) error {
	file, err := os.OpenFile(cachePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer file.Close()

	if err := json.NewEncoder(file).Encode(tok); err != nil {
		return err
	}
	return nil
}

func readCache(cachePath string) (*oauth2.Token, error) {
	file, err := os.Open(cachePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	tok := &oauth2.Token{}
	if err := json.NewDecoder(file).Decode(tok); err != nil {
		return nil, err
	}

	// make sure token is refreshable
	if tok != nil && !tok.Valid() {
		ts := conf.TokenSource(oauth2.NoContext, tok)
		tok, err = ts.Token()
		if err != nil || !tok.Valid() {
			fmt.Printf("Could not refresh cached token: %v\n", err)
			return nil, nil
		}
	}
	return tok, nil
}

func getToken() (*oauth2.Token, error) {
	userInfo, err := user.Current()
	if err != nil {
		return nil, err
	}

	cachePath := filepath.Join(userInfo.HomeDir, ".mantle-cache-google.json")
	tok, err := readCache(cachePath)
	if err != nil {
		log.Printf("Error reading google token cache file: %v", err)
	}
	if tok == nil {
		url := conf.AuthCodeURL("state", oauth2.AccessTypeOffline)
		fmt.Printf("Visit the URL for the auth dialog: %v\n", url)
		fmt.Print("Enter token: ")

		var code string
		if _, err := fmt.Scan(&code); err != nil {
			return nil, err
		}
		tok, err = conf.Exchange(oauth2.NoContext, code)
		if err != nil {
			return nil, err
		}
		err = writeCache(cachePath, tok)
		if err != nil {
			log.Printf("Error writing google token cache file: %v", err)
		}
	}
	return tok, nil
}

func GoogleClient() (*http.Client, error) {
	tok, err := getToken()
	if err != nil {
		return nil, err
	}
	return conf.Client(oauth2.NoContext, tok), nil
}

func GoogleTokenSource() (oauth2.TokenSource, error) {
	tok, err := getToken()
	if err != nil {
		return nil, err
	}
	return conf.TokenSource(oauth2.NoContext, tok), nil
}
<|endoftext|>


<commit_before>package clover

import (
	"encoding/base64"
	"net/http"
	"strings"
)

type AuthServer struct {
	Config    *AuthConfig
	authCtrl  *authController
	tokenCtrl *tokenController
}

func NewAuthServer(config *AuthConfig) *AuthServer {
	a := &AuthServer{
		Config: config,
	}

	tokenResp := a.getTokenRespType()
	authResp := newCodeRespType(a.Config)

	a.tokenCtrl = newTokenController(a.Config, tokenResp)
	a.authCtrl = newAuthController(a.Config, authResp, tokenResp)

	return a
}

func (a *AuthServer) Authorize(w http.ResponseWriter, r *http.Request, isAuthorized bool) *Response {
	ar := a.parseAuthRequest(r)
	resp := a.authCtrl.handleAuthorize(ar, isAuthorized)

	return resp
}

func (a *AuthServer) ValidateAuthorize(w http.ResponseWriter, r *http.Request) (Client, []string, *Response) {
	ar := a.parseAuthRequest(r)

	resp := a.authCtrl.validateAuthRequest(ar)
	if resp != nil {
		return nil, nil, resp
	}

	return ar.client, ar.scopeArr, newResp()
}

func (a *AuthServer) Token(w http.ResponseWriter, r *http.Request) *Response {
	if strings.ToLower(r.Method) != "post" {
		return errMustBePostMetthod
	}

	tr, resp := a.parseTokenRequest(r)
	if resp != nil {
		return resp
	}

	return a.tokenCtrl.handleAccessToken(tr)
}

func (a *AuthServer) parseTokenRequest(r *http.Request) (*TokenRequest, *Response) {
	clientID, clientSecret, resp := a.getCredentialsFromHttp(r, a.Config)
	if resp != nil {
		return nil, resp
	}

	tr := &TokenRequest{
		ClientID:     clientID,
		ClientSecret: clientSecret,
		Scope:        r.FormValue("scope"),
		GrantType:    r.FormValue("grant_type"),
		Username:     r.FormValue("username"),
		Password:     r.FormValue("password"),
		Code:         r.FormValue("code"),
		RedirectURI:  r.FormValue("redirect_uri"),
		RefreshToken: r.FormValue("refresh_token"),
		Assertion:    r.FormValue("assertion"),
	}

	return tr, nil
}

func (a *AuthServer) parseAuthRequest(r *http.Request) *authorizeRequest {
	ar := &authorizeRequest{
		state:        r.FormValue("state"),
		redirectURI:  r.FormValue("redirect_uri"),
		responseType: strings.ToLower(r.FormValue("response_type")),
		clientID:     r.FormValue("client_id"),
		scope:        r.FormValue("scope"),
	}

	return ar
}

func (a *AuthServer) getCredentialsFromHttp(r *http.Request, config *AuthConfig) (string, string, *Response) {
	headerAuth := r.Header.Get("Authorization")

	switch {
	case headerAuth != "":
		s := strings.SplitN(headerAuth, " ", 2)
		if len(s) != 2 || s[0] != "Basic" {
			return "", "", errInvalidAuthHeader
		}

		b, err := base64.StdEncoding.DecodeString(s[1])
		if err != nil {
			return "", "", errInternal(err.Error())
		}

		pair := strings.SplitN(string(b), ":", 2)
		if len(pair) != 2 {
			return "", "", errInvalidAuthMSG
		}

		return pair[0], pair[1], nil
	case config.AllowCredentialsBody:
		if r.PostForm == nil {
			r.ParseForm()
		}
		if r.PostForm.Get(`client_id`) == "" || r.PostForm.Get(`client_secret`) == "" {
			return "", "", errCredentailsNotInBody
		}

		return r.PostForm.Get("client_id"), r.PostForm.Get("client_secret"), nil
	}

	return "", "", errCredentailsRequired
}

func (a *AuthServer) getTokenRespType() ResponseType {
	if a.Config.PublicKeyStore != nil {
		return newJWTResponseType(a.Config)
	}

	return newTokenRespType(a.Config)
}
<commit_msg>change config in authserver to unexport<commit_after>package clover

import (
	"encoding/base64"
	"net/http"
	"strings"
)

type AuthServer struct {
	config    *AuthConfig
	authCtrl  *authController
	tokenCtrl *tokenController
}

func NewAuthServer(config *AuthConfig) *AuthServer {
	a := &AuthServer{
		config: config,
	}

	tokenResp := a.getTokenRespType()
	authResp := newCodeRespType(a.config)

	a.tokenCtrl = newTokenController(a.config, tokenResp)
	a.authCtrl = newAuthController(a.config, authResp, tokenResp)

	return a
}

func (a *AuthServer) Authorize(w http.ResponseWriter, r *http.Request, isAuthorized bool) *Response {
	ar := a.parseAuthRequest(r)
	resp := a.authCtrl.handleAuthorize(ar, isAuthorized)

	return resp
}

func (a *AuthServer) ValidateAuthorize(w http.ResponseWriter, r *http.Request) (Client, []string, *Response) {
	ar := a.parseAuthRequest(r)

	resp := a.authCtrl.validateAuthRequest(ar)
	if resp != nil {
		return nil, nil, resp
	}

	return ar.client, ar.scopeArr, newResp()
}

func (a *AuthServer) Token(w http.ResponseWriter, r *http.Request) *Response {
	if strings.ToLower(r.Method) != "post" {
		return errMustBePostMetthod
	}

	tr, resp := a.parseTokenRequest(r)
	if resp != nil {
		return resp
	}

	return a.tokenCtrl.handleAccessToken(tr)
}

func (a *AuthServer) parseTokenRequest(r *http.Request) (*TokenRequest, *Response) {
	clientID, clientSecret, resp := a.getCredentialsFromHttp(r, a.config)
	if resp != nil {
		return nil, resp
	}

	tr := &TokenRequest{
		ClientID:     clientID,
		ClientSecret: clientSecret,
		Scope:        r.FormValue("scope"),
		GrantType:    r.FormValue("grant_type"),
		Username:     r.FormValue("username"),
		Password:     r.FormValue("password"),
		Code:         r.FormValue("code"),
		RedirectURI:  r.FormValue("redirect_uri"),
		RefreshToken: r.FormValue("refresh_token"),
		Assertion:    r.FormValue("assertion"),
	}

	return tr, nil
}

func (a *AuthServer) parseAuthRequest(r *http.Request) *authorizeRequest {
	ar := &authorizeRequest{
		state:        r.FormValue("state"),
		redirectURI:  r.FormValue("redirect_uri"),
		responseType: strings.ToLower(r.FormValue("response_type")),
		clientID:     r.FormValue("client_id"),
		scope:        r.FormValue("scope"),
	}

	return ar
}

func (a *AuthServer) getCredentialsFromHttp(r *http.Request, config *AuthConfig) (string, string, *Response) {
	headerAuth := r.Header.Get("Authorization")

	switch {
	case headerAuth != "":
		s := strings.SplitN(headerAuth, " ", 2)
		if len(s) != 2 || s[0] != "Basic" {
			return "", "", errInvalidAuthHeader
		}

		b, err := base64.StdEncoding.DecodeString(s[1])
		if err != nil {
			return "", "", errInternal(err.Error())
		}

		pair := strings.SplitN(string(b), ":", 2)
		if len(pair) != 2 {
			return "", "", errInvalidAuthMSG
		}

		return pair[0], pair[1], nil
	case config.AllowCredentialsBody:
		if r.PostForm == nil {
			r.ParseForm()
		}
		if r.PostForm.Get(`client_id`) == "" || r.PostForm.Get(`client_secret`) == "" {
			return "", "", errCredentailsNotInBody
		}

		return r.PostForm.Get("client_id"), r.PostForm.Get("client_secret"), nil
	}

	return "", "", errCredentailsRequired
}

func (a *AuthServer) getTokenRespType() ResponseType {
	if a.config.PublicKeyStore != nil {
		return newJWTResponseType(a.config)
	}

	return newTokenRespType(a.config)
}
<|endoftext|>


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

// Test package for mysql.  Requires a mysql server to be running locally with
// a user `root` with a blank password and a database called `test`.
package mysql_test

import "testing"
import "mysql"
import "os"

func defaultConn(t *testing.T) *mysql.Conn {
	conn := mysql.NewConn();
	err := conn.Connect(&mysql.ConnInfo{
		"127.0.0.1", 3306,
		"root", "",
		"test"
	});
	if err != nil {
		t.Log("Couldn't connect to root:@127.0.0.1:3306:test\n%s",
			err);
		return nil;
	}
	return conn
}

func error(t *testing.T, err os.Error, msg string) {
	if err == nil {
		t.Error(msg)
	}
	else {
		t.Errorf("%s: %s\n", msg, err.String());
	}
}

func TestCursorErrors(t *testing.T) {
	// Cursor on an unconnected object returns nil
	conn := mysql.NewConn();
	if conn.Cursor() != nil {
		error(t, nil, "Unconnected Cursor should be nil")
	}

	conn = defaultConn(t);
	if conn == nil { return }

	cur := conn.Cursor();

	// Fetch called before Execute should return nil with an error
	res, err := cur.FetchOne();
	if res != nil || err == nil {
		error(t, nil, "FetchOne before Execute should error")
	}

	resm, err := cur.FetchMany(10);
	if resm != nil || err == nil {
		error(t, nil, "FetchMany(10) before Execute should error")
	}

	resm, err = cur.FetchAll();
	if resm != nil || err == nil {
		error(t, nil, "FetchAll before Execute should error")
	}

	// Invalid statements return errors
	err = cur.Execute("1");
	if res != nil || err == nil {
		error(t, nil, "Invalid statement should return errors")
	}

	// No result statements should not error, but should not have results.
	err = cur.Execute("# No results");
	if err != nil {
		error(t, err, "No-result SQL statement should not error")
	}

	res, err = cur.FetchOne();
	if res != nil || err == nil {
		error(t, nil, "FetchOne on no-result statement should error")
	}

	cols := cur.Description();
	if len(cols) > 0 {
		t.Error("Description should return no 0-length columns")
	}

	conn.Close();
	cur.Close();
}

func TestCursor(t *testing.T) {
	conn := defaultConn(t);
	if conn == nil { return }
	cur := conn.Cursor();

	err := cur.Execute("CREATE TEMPORARY TABLE t (i INT, s VARCHAR(20));");
	if err != nil {
		error(t, err, "Couldn't create temporary table test.t")
	}

	/*
	TODO: Do proper unicode escaping.
	values := []string{
		"道可道，非常道。", "名可名，非常名。",
		"無名天地之始；", "有名萬物之母。",
		"故常無欲以觀其妙；", "常有欲以觀其徼。",
		"此兩者同出而異名，", "同謂之玄。",
		"玄之又玄，眾妙之門。"
	};
	*/
	values := []string{
		"lorem", "ipsum", "dolor", "sit", "amet", "consectetur",
		"adipisicing", "elit", "sed"
	};
	for i, s := range values {
		err = cur.Execute("INSERT INTO t (i, s) VALUES (%d, %q)",
			i, s);
		if err != nil {
			error(t, err, "Couldn't insert into temporary table test.t")
		}
		if count := cur.RowCount(); count != 0 {
			t.Error("Returned rows for INSERT statement.")
		}
	}

	err = cur.Execute("SELECT i AS pos, s AS phrase FROM t ORDER BY pos ASC");
	if err != nil {
		error(t, err, "Couldn't select from temporary table test.t")
	}
	if count := cur.RowCount(); int(count) != len(values) {
		t.Error("Result count doesn't match inserted count.")
	}

	cols := cur.Description();
	if len(cols) != 2 { t.Error("Description should return 2 columns") }
	if cols[0].Name != "pos" { t.Error("Description()[0] != 'pos'") }
	if cols[1].Name != "phrase" { t.Error("Description()[0] != 'phrase'") }

	i := 0;
	var row []interface {};
	row, err = cur.FetchOne();
	for row != nil {
		if err != nil { error(t, err, "Couldn't FetchOne()") }
		if v, ok := row[1].(string); !ok || values[i] != v {
			if ok { t.Errorf("Mismatch %q != %q", values[i], v) }
			else { t.Errorf("Couldn't convert %v to string.", row[1]) }
		}
		i += 1;
		row, err = cur.FetchOne();
	}

	// Test FetchMany
	err = cur.Execute("SELECT i AS pos, s AS phrase FROM t ORDER BY pos ASC");
	if err != nil {
		error(t, err, "Couldn't select from temporary table test.t")
	}

	var results [][]interface {};
	results, err = cur.FetchMany(3);
	if err != nil { error(t, err, "Error") }
	if len(results) != 3 {
		t.Errorf("Result count mismatch 3 != %d", len(results))
	}
	for i, v := range results {
		if v, ok := v[1].(string); !ok || values[i] != v {
			if ok { t.Errorf("Mismatch %q != %q", values[i], v) }
			else { t.Errorf("Couldn't convert %v to string.", row[1]) }
		}
	}

	// Test FetchAll
	err = cur.Execute("SELECT i AS pos, s AS phrase FROM t ORDER BY pos ASC");
	if err != nil {
		error(t, err, "Couldn't select from temporary table test.t")
	}

	results, err = cur.FetchAll();
	if err != nil { error(t, err, "Error") }

	if len(results) != len(values) {
		t.Errorf("Result count mismatch %d != %d", len(values), len(results))
	}
	for i, v := range results {
		if v, ok := v[1].(string); !ok || values[i] != v {
			if ok { t.Errorf("Mismatch %q != %q", values[i], v) }
			else { t.Errorf("Couldn't convert %v to string.", row[1]) }
		}
	}

	cur.Close();
	conn.Close();
}
<commit_msg>Add Fetch{Many,All} tests when no results are returned<commit_after>// Copyright 2009 Eden Li. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Test package for mysql.  Requires a mysql server to be running locally with
// a user `root` with a blank password and a database called `test`.
package mysql_test

import "testing"
import "mysql"
import "os"

func defaultConn(t *testing.T) *mysql.Conn {
	conn := mysql.NewConn();
	err := conn.Connect(&mysql.ConnInfo{
		"127.0.0.1", 3306,
		"root", "",
		"test"
	});
	if err != nil {
		t.Log("Couldn't connect to root:@127.0.0.1:3306:test\n%s",
			err);
		return nil;
	}
	return conn
}

func error(t *testing.T, err os.Error, msg string) {
	if err == nil {
		t.Error(msg)
	}
	else {
		t.Errorf("%s: %s\n", msg, err.String());
	}
}

func TestCursorErrors(t *testing.T) {
	// Cursor on an unconnected object returns nil
	conn := mysql.NewConn();
	if conn.Cursor() != nil {
		error(t, nil, "Unconnected Cursor should be nil")
	}

	conn = defaultConn(t);
	if conn == nil { return }

	cur := conn.Cursor();

	// Fetch called before Execute should return nil with an error
	res, err := cur.FetchOne();
	if res != nil || err == nil {
		error(t, nil, "FetchOne before Execute should error")
	}

	resm, err := cur.FetchMany(10);
	if resm != nil || err == nil {
		error(t, nil, "FetchMany(10) before Execute should error")
	}

	resm, err = cur.FetchAll();
	if resm != nil || err == nil {
		error(t, nil, "FetchAll before Execute should error")
	}

	// Invalid statements return errors
	err = cur.Execute("1");
	if res != nil || err == nil {
		error(t, nil, "Invalid statement should return errors")
	}

	// No result statements should not error, but should not have results.
	err = cur.Execute("# No results");
	if err != nil {
		error(t, err, "No-result SQL statement should not error")
	}

	res, err = cur.FetchOne();
	if res != nil || err == nil {
		error(t, nil, "FetchOne on no-result statement should error")
	}

	resm, err = cur.FetchMany(10);
	if res != nil || err == nil {
		error(t, nil, "FetchMany on no-result statement should error")
	}

	resm, err = cur.FetchAll();
	if res != nil || err == nil {
		error(t, nil, "FetchAll on no-result statement should error")
	}

	cols := cur.Description();
	if len(cols) > 0 {
		t.Error("Description should return no 0-length columns")
	}

	conn.Close();
	cur.Close();
}

func TestCursor(t *testing.T) {
	conn := defaultConn(t);
	if conn == nil { return }
	cur := conn.Cursor();

	err := cur.Execute("CREATE TEMPORARY TABLE t (i INT, s VARCHAR(20));");
	if err != nil {
		error(t, err, "Couldn't create temporary table test.t")
	}

	/*
	TODO: Do proper unicode escaping.
	values := []string{
		"道可道，非常道。", "名可名，非常名。",
		"無名天地之始；", "有名萬物之母。",
		"故常無欲以觀其妙；", "常有欲以觀其徼。",
		"此兩者同出而異名，", "同謂之玄。",
		"玄之又玄，眾妙之門。"
	};
	*/
	values := []string{
		"lorem", "ipsum", "dolor", "sit", "amet", "consectetur",
		"adipisicing", "elit", "sed"
	};
	for i, s := range values {
		err = cur.Execute("INSERT INTO t (i, s) VALUES (%d, %q)",
			i, s);
		if err != nil {
			error(t, err, "Couldn't insert into temporary table test.t")
		}
		if count := cur.RowCount(); count != 0 {
			t.Error("Returned rows for INSERT statement.")
		}
	}

	err = cur.Execute("SELECT i AS pos, s AS phrase FROM t ORDER BY pos ASC");
	if err != nil {
		error(t, err, "Couldn't select from temporary table test.t")
	}
	if count := cur.RowCount(); int(count) != len(values) {
		t.Error("Result count doesn't match inserted count.")
	}

	cols := cur.Description();
	if len(cols) != 2 { t.Error("Description should return 2 columns") }
	if cols[0].Name != "pos" { t.Error("Description()[0] != 'pos'") }
	if cols[1].Name != "phrase" { t.Error("Description()[0] != 'phrase'") }

	i := 0;
	var row []interface {};
	row, err = cur.FetchOne();
	for row != nil {
		if err != nil { error(t, err, "Couldn't FetchOne()") }
		if v, ok := row[1].(string); !ok || values[i] != v {
			if ok { t.Errorf("Mismatch %q != %q", values[i], v) }
			else { t.Errorf("Couldn't convert %v to string.", row[1]) }
		}
		i += 1;
		row, err = cur.FetchOne();
	}

	// Test FetchMany
	err = cur.Execute("SELECT i AS pos, s AS phrase FROM t ORDER BY pos ASC");
	if err != nil {
		error(t, err, "Couldn't select from temporary table test.t")
	}

	var results [][]interface {};
	results, err = cur.FetchMany(3);
	if err != nil { error(t, err, "Error") }
	if len(results) != 3 {
		t.Errorf("Result count mismatch 3 != %d", len(results))
	}
	for i, v := range results {
		if v, ok := v[1].(string); !ok || values[i] != v {
			if ok { t.Errorf("Mismatch %q != %q", values[i], v) }
			else { t.Errorf("Couldn't convert %v to string.", row[1]) }
		}
	}

	// Test FetchAll
	err = cur.Execute("SELECT i AS pos, s AS phrase FROM t ORDER BY pos ASC");
	if err != nil {
		error(t, err, "Couldn't select from temporary table test.t")
	}

	results, err = cur.FetchAll();
	if err != nil { error(t, err, "Error") }

	if len(results) != len(values) {
		t.Errorf("Result count mismatch %d != %d", len(values), len(results))
	}
	for i, v := range results {
		if v, ok := v[1].(string); !ok || values[i] != v {
			if ok { t.Errorf("Mismatch %q != %q", values[i], v) }
			else { t.Errorf("Couldn't convert %v to string.", row[1]) }
		}
	}

	cur.Close();
	conn.Close();
}
<|endoftext|>


<commit_before>package main

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

	"github.com/inge4pres/gmq/communication"
)

var action, qname, payload, protocol, server, logfile, output, user, password string
var logger *log.Logger

func main() {

	flag.StringVar(&user, "user", "", "Username to authenticate with GMQ")
	flag.StringVar(&password, "password", "", "Password to authenticate to GMQ")
	flag.StringVar(&action, "a", "", "Action to be performed: P = Publish, S = Subscribe, L = List")
	flag.StringVar(&qname, "q", "", "Queue name")
	flag.StringVar(&payload, "m", "", "Base64 encoded payload to use (only needed with \"-a P\")")
	flag.StringVar(&protocol, "p", "tcp", "Protocol type: tcp4 or tcp6 (default to system settings)")
	flag.StringVar(&logfile, "l", "", "Log destination: if not set, defaults to STDOUT")
	flag.StringVar(&server, "s", "localhost:4567", "Server address: (DNS|IP)[:PORT] form ")
	flag.StringVar(&output, "o", "", "Output file, if empty defaults to STDOUT")
	flag.Parse()

	logger = initLog()

	out := configureOutput()
	mex := configureCall()
	resp := callServer(mex)
	if _, err := out.Write(resp); err != nil {
		logger.Printf("There was an error printing response from server to %s\n", out.Name())
		logger.Println("Response:\n%s", string(resp))
	}

}

func initLog() *log.Logger {
	file, err := os.OpenFile(logfile, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0660)
	if err != nil {
		return log.New(os.Stdout, "[GMQ Client] ", log.LstdFlags)

	}
	return log.New(file, "[GMQ Client] ", log.LstdFlags)
}

func configureCall() *gmqnet.Message {
	if user == "" {
		logger.Fatalln("User cannot be null! Set -user")
	}
	if password == "" {
		logger.Fatalln("Password cannot be null! Set -password")
	}
	if action == "" {
		logger.Fatalln("Action cannot be null! Set -a")
	}
	if qname == "" {
		logger.Fatalln("Queue name cannot be null! Set -q")
	}

	return &gmqnet.Message{
		Auth: gmqnet.AuthToken{
			UserTok: gmqnet.GenToken(user),
			PwdTok:  gmqnet.GenToken(password),
		},
		Operation: action,
		Queue:     qname,
		Payload:   payload,
	}
}

func configureOutput() *os.File {
	if output == "" {
		return os.Stdout
	}
	out, err := os.OpenFile(output, os.O_RDWR|os.O_CREATE, 0660)
	if err != nil {
		logger.Fatalln("Error opening output file!")
	}
	return out
}

func callServer(mex *gmqnet.Message) []byte {
	conn, err := net.Dial(protocol, server)
	defer conn.Close()
	if err != nil {
		logger.Fatalf("Error connecting to the server:\n%T\n%s\n", err, err.Error())
	}
	if _, err := conn.Write(gmqnet.WriteMessage(mex)); err != nil {
		logger.Fatalf("Error posting the message to server:\n%T\n%s\n", err, err.Error())
	}
	resp, err := ioutil.ReadAll(conn)
	if err != nil {
		logger.Fatalf("Error reading response from server:\n%T\n%s\n", err, err.Error())
	}
	return resp
}
<commit_msg>Updated imports<commit_after>package main

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

	"github.com/inge4pres/gmq/communication"
)

var action, qname, payload, protocol, server, logfile, output, user, password string
var logger *log.Logger

func main() {

	flag.StringVar(&user, "user", "", "Username to authenticate with GMQ")
	flag.StringVar(&password, "password", "", "Password to authenticate to GMQ")
	flag.StringVar(&action, "a", "", "Action to be performed: P = Publish, S = Subscribe, L = List")
	flag.StringVar(&qname, "q", "", "Queue name")
	flag.StringVar(&payload, "m", "", "Base64 encoded payload to use (only needed with \"-a P\")")
	flag.StringVar(&protocol, "p", "tcp", "Protocol type: tcp4 or tcp6 (default to system settings)")
	flag.StringVar(&logfile, "l", "", "Log destination: if not set, defaults to STDOUT")
	flag.StringVar(&server, "s", "localhost:4567", "Server address: (DNS|IP)[:PORT] form ")
	flag.StringVar(&output, "o", "", "Output file, if empty defaults to STDOUT")
	flag.Parse()

	logger = initLog()

	out := configureOutput()
	mex := configureCall()
	resp := callServer(mex)
	if _, err := out.Write(resp); err != nil {
		logger.Printf("There was an error printing response from server to %s\n", out.Name())
		logger.Println("Response:\n%s", string(resp))
	}

}

func initLog() *log.Logger {
	file, err := os.OpenFile(logfile, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0660)
	if err != nil {
		return log.New(os.Stdout, "[GMQ Client] ", log.LstdFlags)

	}
	return log.New(file, "[GMQ Client] ", log.LstdFlags)
}

func configureCall() *gmqnet.Message {
	if user == "" {
		logger.Fatalln("User cannot be null! Set -user")
	}
	if password == "" {
		logger.Fatalln("Password cannot be null! Set -password")
	}
	if action == "" {
		logger.Fatalln("Action cannot be null! Set -a")
	}
	if qname == "" {
		logger.Fatalln("Queue name cannot be null! Set -q")
	}

	return &gmqnet.Message{
		Auth: gmqnet.AuthToken{
			UserTok: gmqnet.GenSha512Token(user),
			PwdTok:  gmqnet.GenSha512Token(password),
		},
		Operation: action,
		Queue:     qname,
		Payload:   payload,
	}
}

func configureOutput() *os.File {
	if output == "" {
		return os.Stdout
	}
	out, err := os.OpenFile(output, os.O_RDWR|os.O_CREATE, 0660)
	if err != nil {
		logger.Fatalln("Error opening output file!")
	}
	return out
}

func callServer(mex *gmqnet.Message) []byte {
	conn, err := net.Dial(protocol, server)
	defer conn.Close()
	if err != nil {
		logger.Fatalf("Error connecting to the server:\n%T\n%s\n", err, err.Error())
	}
	if _, err := conn.Write(gmqnet.WriteMessage(mex)); err != nil {
		logger.Fatalf("Error posting the message to server:\n%T\n%s\n", err, err.Error())
	}
	resp, err := ioutil.ReadAll(conn)
	if err != nil {
		logger.Fatalf("Error reading response from server:\n%T\n%s\n", err, err.Error())
	}
	return resp
}
<|endoftext|>