<commit_before>// Package monkey lets you wrap bad behaviour around your http.Handlers, such as request delays, incorrect responses, bodies and returning garbage from configuration. By using monkey you can simulate the unpredictable nature of calling services over HTTP.
package monkey

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

// server wraps around a http.Handler and adds destructive behaviour (monkey business) based on the behaviours passed in
type server struct {
	delegate   http.Handler
	behaviours []behaviour
	randomiser randomiser
}

// NewServerFromYAML creates a http.Handler which wraps monkey business defined from YAML around it, to return a new http.Handler. If the YAML is invalid, it will return an error.
func NewServerFromYAML(server http.Handler, YAML []byte) (http.Handler, error) {
	behaviours, err := monkeyConfigFromYAML(YAML)

	if err != nil {
		return nil, errors.New(fmt.Sprintf("Problem occured when trying to parse the config file: %v", err))
	}

	log.Println("Monkey config loaded")
	for _, b := range behaviours {
		log.Println(b)
	}

	return newServerFromBehaviour(server, behaviours), nil
}

// NewServer creates a http.Handler which wraps it's monkey business around it, to return a new http.Handler. If no behaviours are defined in the config it will return the original handler, otherwise an error
func NewServer(server http.Handler, configPath string) (http.Handler, error) {
	if configPath == "" {
		return server, nil
	}

	data, err := ioutil.ReadFile(configPath)

	if err != nil {
		return nil, errors.New(fmt.Sprintf("Problem occured when trying to read the config file: %v", err))
	}

	return NewServerFromYAML(server, data)

}

func newServerFromBehaviour(degegate http.Handler, behaviours []behaviour) http.Handler {
	return &server{degegate, behaviours, new(defaultRandomiser)}
}

func (s *server) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	var responseWriter http.ResponseWriter
	if chosenBehaviour := getBehaviour(s.behaviours, s.randomiser); chosenBehaviour != nil {
		s.misbehave(*chosenBehaviour, w)
		responseWriter = monkeyWriter{w, []byte(chosenBehaviour.Body), chosenBehaviour.Garbage}
	} else {
		responseWriter = w
	}

	s.delegate.ServeHTTP(responseWriter, r)
}

func (s *server) misbehave(behaviour behaviour, w http.ResponseWriter) {
	time.Sleep(behaviour.Delay * time.Millisecond)
	if behaviour.Status != 0 {
		w.WriteHeader(behaviour.Status)
	}
}

type monkeyWriter struct {
	http.ResponseWriter
	newBody      []byte
	garbageCount int
}

func (w monkeyWriter) Write(data []byte) (int, error) {

	if w.garbageCount > 0 {
		content := []byte{}
		for i := 0; i < w.garbageCount; i++ {
			content = append(content, byte('a'))
		}
		return w.ResponseWriter.Write(content)
	}

	if len(w.newBody) > 0 {
		return w.ResponseWriter.Write(w.newBody)
	}
	return w.ResponseWriter.Write(data)
}
<commit_msg>Fixed multiple WriteHeader issue #7<commit_after>// Package monkey lets you wrap bad behaviour around your http.Handlers, such as request delays, incorrect responses, bodies and returning garbage from configuration. By using monkey you can simulate the unpredictable nature of calling services over HTTP.
package monkey

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

// server wraps around a http.Handler and adds destructive behaviour (monkey business) based on the behaviours passed in
type server struct {
	delegate   http.Handler
	behaviours []behaviour
	randomiser randomiser
}

// NewServerFromYAML creates a http.Handler which wraps monkey business defined from YAML around it, to return a new http.Handler. If the YAML is invalid, it will return an error.
func NewServerFromYAML(server http.Handler, YAML []byte) (http.Handler, error) {
	behaviours, err := monkeyConfigFromYAML(YAML)

	if err != nil {
		return nil, errors.New(fmt.Sprintf("Problem occured when trying to parse the config file: %v", err))
	}

	log.Println("Monkey config loaded")
	for _, b := range behaviours {
		log.Println(b)
	}

	return newServerFromBehaviour(server, behaviours), nil
}

// NewServer creates a http.Handler which wraps it's monkey business around it, to return a new http.Handler. If no behaviours are defined in the config it will return the original handler, otherwise an error
func NewServer(server http.Handler, configPath string) (http.Handler, error) {
	if configPath == "" {
		return server, nil
	}

	data, err := ioutil.ReadFile(configPath)

	if err != nil {
		return nil, errors.New(fmt.Sprintf("Problem occured when trying to read the config file: %v", err))
	}

	return NewServerFromYAML(server, data)

}

func newServerFromBehaviour(degegate http.Handler, behaviours []behaviour) http.Handler {
	return &server{degegate, behaviours, new(defaultRandomiser)}
}

func (s *server) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	var responseWriter http.ResponseWriter
	if chosenBehaviour := getBehaviour(s.behaviours, s.randomiser); chosenBehaviour != nil {
		responseWriter = monkeyWriter{w, chosenBehaviour}
	} else {
		responseWriter = w
	}

	s.delegate.ServeHTTP(responseWriter, r)
}

type monkeyWriter struct {
	http.ResponseWriter
	behaviour *behaviour
}

func (w monkeyWriter) Write(data []byte) (int, error) {

	time.Sleep(w.behaviour.Delay * time.Millisecond)

	if w.behaviour.Garbage > 0 {
		content := []byte{}
		for i := 0; i < w.behaviour.Garbage; i++ {
			content = append(content, byte('a'))
		}
		return w.ResponseWriter.Write(content)
	}

	if len(w.behaviour.Body) > 0 {
		return w.ResponseWriter.Write([]byte(w.behaviour.Body))
	}
	return w.ResponseWriter.Write(data)
}

func (w monkeyWriter) WriteHeader(code int) {
	if w.behaviour.Status != 0 {
		w.ResponseWriter.WriteHeader(w.behaviour.Status)
	} else {
		w.ResponseWriter.WriteHeader(code)
	}
}
<|endoftext|>


<commit_before>package msgbox

import "github.com/andlabs/ui"

// New creates a new Window and hides the parent window until the "Ok" button is pressed
// once clicked, the done channel is closed so that the calling function can continue
//
// WARNING: New can't be called by the goroutine that created the parent window
func New(p ui.Window, titel, msg string) {
	done := make(chan struct{})
	go ui.Do(func() {
		p.Hide()
		msgLabel := ui.NewLabel(msg)
		btn := ui.NewButton("Ok")
		stack := ui.NewVerticalStack(
			msgLabel,
			btn,
		)
		stack.SetStretchy(0)
		w := ui.NewWindow(titel, 500, 200, stack)
		btn.OnClicked(func() {
			close(done)
			w.Close()
			p.Show()
		})
		w.Show()
	})
	<-done
}
<commit_msg>msgbox: use a TextField for errors<commit_after>package msgbox

import "DST/Godeps/_workspace/src/github.com/andlabs/ui"

// New creates a new Window and hides the parent window until the "Ok" button is pressed
// once clicked, the done channel is closed so that the calling function can continue
//
// WARNING: New can't be called by the goroutine that created the parent window
func New(p ui.Window, titel, msg string) {
	done := make(chan struct{})
	go ui.Do(func() {
		p.Hide()
		msgField := ui.NewTextField()
		msgField.SetReadOnly(true)
		msgField.SetText(msg)
		btn := ui.NewButton("Ok")
		stack := ui.NewVerticalStack(
			msgField,
			btn,
		)
		stack.SetStretchy(0)
		w := ui.NewWindow(titel, 500, 200, stack)
		btn.OnClicked(func() {
			close(done)
			w.Close()
			p.Show()
		})
		w.Show()
	})
	<-done
}
<|endoftext|>


<commit_before>package mstats

import (
	"runtime"
	"strconv"
	"sync/atomic"
	"time"
)

// Var is an atomic variable satisfying expvar.Var
type Var struct {
	atomic.Value
}

func (a *Var) String() string {
	v := a.Load()
	if v == nil {
		return "0"
	}
	return strconv.FormatUint(v.(uint64), 10)
}

// PauseNS is the total number of nanoseconds the GC has paused the application
var PauseNS Var

// NumGC is the number of collections
var NumGC Var

// Alloc is the number of bytes allocated and not yet freed by the application
var Alloc Var

// Start polls runtime.ReadMemStats with interval d and updates the package level variables
func Start(d time.Duration) {
	for range time.Tick(d) {
		var m runtime.MemStats
		runtime.ReadMemStats(&m)
		PauseNS.Store(m.PauseTotalNs)
		Alloc.Store(m.Alloc)
		NumGC.Store(uint64(m.NumGC))
	}
}
<commit_msg>mstats: add TotalAlloc<commit_after>package mstats

import (
	"runtime"
	"strconv"
	"sync/atomic"
	"time"
)

// Var is an atomic variable satisfying expvar.Var
type Var struct {
	atomic.Value
}

func (a *Var) String() string {
	v := a.Load()
	if v == nil {
		return "0"
	}
	return strconv.FormatUint(v.(uint64), 10)
}

// PauseNS is the total number of nanoseconds the GC has paused the application
var PauseNS Var

// NumGC is the number of collections
var NumGC Var

// Alloc is the number of bytes allocated and not yet freed by the application
var Alloc Var

// TotalAlloc is the total number of bytes allocated by the application
var TotalAlloc Var

// Start polls runtime.ReadMemStats with interval d and updates the package level variables
func Start(d time.Duration) {
	for range time.Tick(d) {
		var m runtime.MemStats
		runtime.ReadMemStats(&m)
		PauseNS.Store(m.PauseTotalNs)
		Alloc.Store(m.Alloc)
		TotalAlloc.Store(m.TotalAlloc)
		NumGC.Store(uint64(m.NumGC))
	}
}
<|endoftext|>


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

// Package testserver provides helpers to run a cockroach binary within tests.
// It automatically downloads the latest cockroach binary for your platform
// (Linux-amd64 and Darwin-amd64 only for now), or attempts to run "cockroach"
// from your PATH.
//
// A normal invocation is (check err every time):
// ts, err := testserver.NewTestServer()
// err = ts.Start()
// defer ts.Stop()
// url := ts.PGURL()
//
// To use, run as follows:
//   import "github.com/cockroachdb/cockroach-go/testserver"
//   import "testing"
//   import "time"
//
//   func TestRunServer(t *testing.T) {
//      ts, err := testserver.NewTestServer()
//      if err != nil {
//        t.Fatal(err)
//      }
//      err := ts.Start()
//      if err != nil {
//        t.Fatal(err)
//      }
//      defer ts.Stop()
//
//      url := ts.PGURL()
//      if url != nil {
//        t.FatalF("url not found")
//      }
//      t.Logf("URL: %s", url.String())
//
//      db, err := sql.Open("postgres", url.String())
//      if err != nil {
//        t.Fatal(err)
//      }
//    }
package testserver

import (
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"testing"
	"time"
)

var (
	sqlURLRegexp = regexp.MustCompile("sql:\\s+(postgresql:.+)\n")
	customBinary = flag.String("cockroach-binary", "", "Use specified cockroach binary")
)

const (
	stateNew     = iota
	stateRunning = iota
	stateStopped = iota
	stateFailed  = iota

	socketPort     = 26257
	socketFileBase = ".s.PGSQL"
)

// TestServer is a helper to run a real cockroach node.
type TestServer struct {
	mu        sync.RWMutex
	state     int
	baseDir   string
	pgURL     *url.URL
	cmd       *exec.Cmd
	args      []string
	stdout    string
	stderr    string
	stdoutBuf logWriter
	stderrBuf logWriter
}

// NewDBForTest creates a new CockroachDB TestServer instance and
// opens a SQL database connection to it. Returns a sql *DB instance a
// shutdown function. The caller is responsible for executing the
// returned shutdown function on exit.
func NewDBForTest(t *testing.T) (*sql.DB, func()) {
	return NewDBForTestWithDatabase(t, "")
}

// NewDBForTestWithDatabase creates a new CockroachDB TestServer
// instance and opens a SQL database connection to it. If database is
// specified, the returned connection will explicitly connect to
// it. Returns a sql *DB instance a shutdown function. The caller is
// responsible for executing the returned shutdown function on exit.
func NewDBForTestWithDatabase(t *testing.T, database string) (*sql.DB, func()) {
	ts, err := NewTestServer()
	if err != nil {
		t.Fatal(err)
	}

	err = ts.Start()
	if err != nil {
		t.Fatal(err)
	}

	url := ts.PGURL()
	if url == nil {
		t.Fatalf("url not found")
	}
	if len(database) > 0 {
		url.Path = database
	}

	db, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	if err := ts.WaitForInit(db); err != nil {
		t.Fatal(err)
	}

	return db, func() {
		_ = db.Close()
		ts.Stop()
	}
}

// NewTestServer creates a new TestServer, but does not start it.
// The cockroach binary for your OS and ARCH is downloaded automatically.
// If the download fails, we attempt just call "cockroach", hoping it is
// found in your path.
func NewTestServer() (*TestServer, error) {
	var cockroachBinary string
	var err error
	if len(*customBinary) > 0 {
		cockroachBinary = *customBinary
		log.Printf("Using custom cockroach binary: %s", cockroachBinary)
	} else if cockroachBinary, err = downloadLatestBinary(); err != nil {
		log.Printf("Failed to fetch latest binary: %s, attempting to use cockroach binary from your PATH", err)
		cockroachBinary = "cockroach"
	} else {
		log.Printf("Using automatically-downloaded binary: %s", cockroachBinary)
	}

	// Force "/tmp/" so avoid OSX's really long temp directory names
	// which get us over the socket filename length limit.
	baseDir, err := ioutil.TempDir("/tmp", "cockroach-testserver")
	if err != nil {
		return nil, fmt.Errorf("could not create temp directory: %s", err)
	}

	logDir := filepath.Join(baseDir, "logs")
	if err := os.MkdirAll(logDir, 0755); err != nil {
		return nil, fmt.Errorf("could not create logs directory: %s: %s", logDir, err)
	}

	options := url.Values{
		"host": []string{baseDir},
	}
	pgurl := &url.URL{
		Scheme:   "postgres",
		User:     url.User("root"),
		Host:     fmt.Sprintf(":%d", socketPort),
		RawQuery: options.Encode(),
	}
	socketPath := filepath.Join(baseDir, fmt.Sprintf("%s.%d", socketFileBase, socketPort))

	args := []string{
		cockroachBinary,
		"start",
		"--logtostderr",
		"--insecure",
		"--port=0",
		"--http-port=0",
		"--socket=" + socketPath,
		"--store=" + baseDir,
	}

	ts := &TestServer{
		baseDir: baseDir,
		pgURL:   pgurl,
		args:    args,
		stdout:  filepath.Join(logDir, "cockroach.stdout"),
		stderr:  filepath.Join(logDir, "cockroach.stderr"),
	}
	return ts, nil
}

// Stdout returns the entire contents of the process' stdout.
func (ts *TestServer) Stdout() string {
	return ts.stdoutBuf.String()
}

// Stderr returns the entire contents of the process' stderr.
func (ts *TestServer) Stderr() string {
	return ts.stderrBuf.String()
}

// PGURL returns the postgres connection URL to reach the started
// cockroach node.
// It loops until the expected unix socket file exists.
// This does not timeout, relying instead on test timeouts.
func (ts *TestServer) PGURL() *url.URL {
	socketPath := filepath.Join(ts.baseDir, fmt.Sprintf("%s.%d", socketFileBase, socketPort))
	for {
		if _, err := os.Stat(socketPath); err == nil {
			return ts.pgURL
		}
		time.Sleep(time.Millisecond * 10)
	}
}

// WaitForInit retries until a connection is successfully established.
func (ts *TestServer) WaitForInit(db *sql.DB) error {
	var err error
	for i := 0; i < 50; i++ {
		if _, err = db.Query("SHOW DATABASES"); err == nil {
			return err
		}
		log.Printf("WaitForInit: %v", err)
		time.Sleep(time.Millisecond * 100)
	}
	return err
}

// Start runs the process, returning an error on any problems,
// including being unable to start, but not unexpected failure.
// It should only be called once in the lifetime of a TestServer object.
func (ts *TestServer) Start() error {
	ts.mu.Lock()
	if ts.state != stateNew {
		ts.mu.Unlock()
		return errors.New("Start() can only be called once")
	}
	ts.state = stateRunning
	ts.mu.Unlock()

	ts.cmd = exec.Command(ts.args[0], ts.args[1:]...)
	ts.cmd.Env = []string{"COCKROACH_MAX_OFFSET=1ns"}

	if len(ts.stdout) > 0 {
		wr, err := newFileLogWriter(ts.stdout)
		if err != nil {
			return fmt.Errorf("unable to open file %s: %s", ts.stdout, err)
		}
		ts.stdoutBuf = wr
	}
	ts.cmd.Stdout = ts.stdoutBuf

	if len(ts.stderr) > 0 {
		wr, err := newFileLogWriter(ts.stderr)
		if err != nil {
			return fmt.Errorf("unable to open file %s: %s", ts.stderr, err)
		}
		ts.stderrBuf = wr
	}
	ts.cmd.Stderr = ts.stderrBuf

	for k, v := range defaultEnv() {
		ts.cmd.Env = append(ts.cmd.Env, k+"="+v)
	}

	err := ts.cmd.Start()
	if ts.cmd.Process != nil {
		log.Printf("process %d started: %s", ts.cmd.Process.Pid, strings.Join(ts.args, " "))
	}
	if err != nil {
		log.Printf(err.Error())
		ts.stdoutBuf.Close()
		ts.stderrBuf.Close()

		ts.mu.Lock()
		ts.state = stateFailed
		ts.mu.Unlock()

		return fmt.Errorf("failure starting process: %s", err)
	}

	go func() {
		ts.cmd.Wait()

		ts.stdoutBuf.Close()
		ts.stderrBuf.Close()

		ps := ts.cmd.ProcessState
		sy := ps.Sys().(syscall.WaitStatus)

		log.Printf("Process %d exited with status %d", ps.Pid(), sy.ExitStatus())
		log.Printf(ps.String())

		ts.mu.Lock()
		if sy.ExitStatus() == 0 {
			ts.state = stateStopped
		} else {
			ts.state = stateFailed
		}
		ts.mu.Unlock()
	}()

	return nil
}

// Stop kills the process if it is still running and cleans its directory.
// It should only be called once in the lifetime of a TestServer object.
// Logs fatal if the process has already failed.
func (ts *TestServer) Stop() {
	ts.mu.RLock()
	defer ts.mu.RUnlock()

	if ts.state == stateNew {
		log.Fatal("Stop() called, but Start() was never called")
	}
	if ts.state == stateFailed {
		log.Fatalf("Stop() called, but process exited unexpectedly. Stdout:\n%s\nStderr:\n%s\n",
			ts.Stdout(), ts.Stderr())
		return
	}

	if ts.state != stateStopped {
		// Only call kill if not running. It could have exited properly.
		ts.cmd.Process.Kill()
	}

	// Only cleanup on intentional stops.
	_ = os.RemoveAll(ts.baseDir)
}

type logWriter interface {
	Write(p []byte) (n int, err error)
	String() string
	Len() int64
	Close()
}

type fileLogWriter struct {
	filename string
	file     *os.File
}

func newFileLogWriter(file string) (*fileLogWriter, error) {
	f, err := os.Create(file)
	if err != nil {
		return nil, err
	}

	return &fileLogWriter{
		filename: file,
		file:     f,
	}, nil
}

func (w fileLogWriter) Close() {
	w.file.Close()
}

func (w fileLogWriter) Write(p []byte) (n int, err error) {
	return w.file.Write(p)
}

func (w fileLogWriter) String() string {
	b, err := ioutil.ReadFile(w.filename)
	if err == nil {
		return string(b)
	}
	return ""
}

func (w fileLogWriter) Len() int64 {
	s, err := os.Stat(w.filename)
	if err == nil {
		return s.Size()
	}
	return 0
}

func defaultEnv() map[string]string {
	vars := map[string]string{}
	u, err := user.Current()
	if err == nil {
		if _, ok := vars["USER"]; !ok {
			vars["USER"] = u.Username
		}
		if _, ok := vars["UID"]; !ok {
			vars["UID"] = u.Uid
		}
		if _, ok := vars["GID"]; !ok {
			vars["GID"] = u.Gid
		}
		if _, ok := vars["HOME"]; !ok {
			vars["HOME"] = u.HomeDir
		}
	}
	if _, ok := vars["PATH"]; !ok {
		vars["PATH"] = os.Getenv("PATH")
	}
	return vars
}
<commit_msg>Expose TCP/IP URL instead of Socket URL for TestServers<commit_after>// Copyright 2016 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

// Package testserver provides helpers to run a cockroach binary within tests.
// It automatically downloads the latest cockroach binary for your platform
// (Linux-amd64 and Darwin-amd64 only for now), or attempts to run "cockroach"
// from your PATH.
//
// A normal invocation is (check err every time):
// ts, err := testserver.NewTestServer()
// err = ts.Start()
// defer ts.Stop()
// url := ts.PGURL()
//
// To use, run as follows:
//   import "github.com/cockroachdb/cockroach-go/testserver"
//   import "testing"
//   import "time"
//
//   func TestRunServer(t *testing.T) {
//      ts, err := testserver.NewTestServer()
//      if err != nil {
//        t.Fatal(err)
//      }
//      err := ts.Start()
//      if err != nil {
//        t.Fatal(err)
//      }
//      defer ts.Stop()
//
//      url := ts.PGURL()
//      if url != nil {
//        t.FatalF("url not found")
//      }
//      t.Logf("URL: %s", url.String())
//
//      db, err := sql.Open("postgres", url.String())
//      if err != nil {
//        t.Fatal(err)
//      }
//    }
package testserver

import (
	"bytes"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"testing"
	"time"
)

var (
	sqlURLRegexp = regexp.MustCompile(`sql:\s+(postgresql:.+)`)
	customBinary = flag.String("cockroach-binary", "", "Use specified cockroach binary")
)

const (
	stateNew = 1 + iota
	stateRunning
	stateStopped
	stateFailed
)

// TestServer is a helper to run a real cockroach node.
type TestServer struct {
	mu      sync.RWMutex
	state   int
	baseDir string
	pgURL   struct {
		set chan struct{}
		u   *url.URL
	}
	cmd        *exec.Cmd
	args       []string
	stdout     string
	stderr     string
	stdoutIntr *patternInterceptor
	stdoutBuf  logWriter
	stderrBuf  logWriter
}

// NewDBForTest creates a new CockroachDB TestServer instance and
// opens a SQL database connection to it. Returns a sql *DB instance a
// shutdown function. The caller is responsible for executing the
// returned shutdown function on exit.
func NewDBForTest(t *testing.T) (*sql.DB, func()) {
	return NewDBForTestWithDatabase(t, "")
}

// NewDBForTestWithDatabase creates a new CockroachDB TestServer
// instance and opens a SQL database connection to it. If database is
// specified, the returned connection will explicitly connect to
// it. Returns a sql *DB instance a shutdown function. The caller is
// responsible for executing the returned shutdown function on exit.
func NewDBForTestWithDatabase(t *testing.T, database string) (*sql.DB, func()) {
	ts, err := NewTestServer()
	if err != nil {
		t.Fatal(err)
	}

	err = ts.Start()
	if err != nil {
		t.Fatal(err)
	}

	url := ts.PGURL()
	if url == nil {
		t.Fatalf("url not found")
	}
	if len(database) > 0 {
		url.Path = database
	}

	db, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	if err := ts.WaitForInit(db); err != nil {
		t.Fatal(err)
	}

	return db, func() {
		_ = db.Close()
		ts.Stop()
	}
}

// NewTestServer creates a new TestServer, but does not start it.
// The cockroach binary for your OS and ARCH is downloaded automatically.
// If the download fails, we attempt just call "cockroach", hoping it is
// found in your path.
func NewTestServer() (*TestServer, error) {
	var cockroachBinary string
	var err error
	if len(*customBinary) > 0 {
		cockroachBinary = *customBinary
		log.Printf("Using custom cockroach binary: %s", cockroachBinary)
	} else if cockroachBinary, err = downloadLatestBinary(); err != nil {
		log.Printf("Failed to fetch latest binary: %s, attempting to use cockroach binary from your PATH", err)
		cockroachBinary = "cockroach"
	} else {
		log.Printf("Using automatically-downloaded binary: %s", cockroachBinary)
	}

	// Force "/tmp/" so avoid OSX's really long temp directory names
	// which get us over the socket filename length limit.
	baseDir, err := ioutil.TempDir("/tmp", "cockroach-testserver")
	if err != nil {
		return nil, fmt.Errorf("could not create temp directory: %s", err)
	}

	logDir := filepath.Join(baseDir, "logs")
	if err := os.MkdirAll(logDir, 0755); err != nil {
		return nil, fmt.Errorf("could not create logs directory: %s: %s", logDir, err)
	}

	args := []string{
		cockroachBinary,
		"start",
		"--logtostderr",
		"--insecure",
		"--host=localhost",
		"--port=0",
		"--http-port=0",
		"--store=" + baseDir,
	}

	ts := &TestServer{
		state:   stateNew,
		baseDir: baseDir,
		args:    args,
		stdout:  filepath.Join(logDir, "cockroach.stdout"),
		stderr:  filepath.Join(logDir, "cockroach.stderr"),
	}
	ts.pgURL.set = make(chan struct{})
	return ts, nil
}

// Stdout returns the entire contents of the process' stdout.
func (ts *TestServer) Stdout() string {
	return ts.stdoutBuf.String()
}

// Stderr returns the entire contents of the process' stderr.
func (ts *TestServer) Stderr() string {
	return ts.stderrBuf.String()
}

// PGURL returns the postgres connection URL to reach the started
// cockroach node.
//
// It blocks until the network URL is determined and does not timeout,
// relying instead on test timeouts.
func (ts *TestServer) PGURL() *url.URL {
	<-ts.pgURL.set
	return ts.pgURL.u
}

func (ts *TestServer) setPGURL(u *url.URL) {
	ts.pgURL.u = u
	close(ts.pgURL.set)
}

// WaitForInit retries until a connection is successfully established.
func (ts *TestServer) WaitForInit(db *sql.DB) error {
	var err error
	for i := 0; i < 50; i++ {
		if _, err = db.Query("SHOW DATABASES"); err == nil {
			return err
		}
		log.Printf("WaitForInit: %v", err)
		time.Sleep(time.Millisecond * 100)
	}
	return err
}

// Start runs the process, returning an error on any problems,
// including being unable to start, but not unexpected failure.
// It should only be called once in the lifetime of a TestServer object.
func (ts *TestServer) Start() error {
	ts.mu.Lock()
	if ts.state != stateNew {
		ts.mu.Unlock()
		return errors.New("Start() can only be called once")
	}
	ts.state = stateRunning
	ts.mu.Unlock()

	ts.cmd = exec.Command(ts.args[0], ts.args[1:]...)
	ts.cmd.Env = []string{"COCKROACH_MAX_OFFSET=1ns"}

	if len(ts.stdout) > 0 {
		wr, err := newFileLogWriter(ts.stdout)
		if err != nil {
			return fmt.Errorf("unable to open file %s: %s", ts.stdout, err)
		}
		ts.stdoutBuf = wr
	}

	pi := newPatternInterceptor(ts.stdoutBuf, sqlURLRegexp, func(match [][]byte) error {
		u, err := url.Parse(string(match[1]))
		if err != nil {
			return fmt.Errorf("failure to parse SQL URL: %v", err)
		}
		ts.setPGURL(u)
		ts.stdoutIntr.Disable()
		return nil
	})
	ts.stdoutIntr = pi
	ts.cmd.Stdout = pi

	if len(ts.stderr) > 0 {
		wr, err := newFileLogWriter(ts.stderr)
		if err != nil {
			return fmt.Errorf("unable to open file %s: %s", ts.stderr, err)
		}
		ts.stderrBuf = wr
	}
	ts.cmd.Stderr = ts.stderrBuf

	for k, v := range defaultEnv() {
		ts.cmd.Env = append(ts.cmd.Env, k+"="+v)
	}

	err := ts.cmd.Start()
	if ts.cmd.Process != nil {
		log.Printf("process %d started: %s", ts.cmd.Process.Pid, strings.Join(ts.args, " "))
	}
	if err != nil {
		log.Printf(err.Error())
		ts.stdoutBuf.Close()
		ts.stderrBuf.Close()

		ts.mu.Lock()
		ts.state = stateFailed
		ts.mu.Unlock()

		return fmt.Errorf("failure starting process: %s", err)
	}

	go func() {
		ts.cmd.Wait()

		ts.stdoutBuf.Close()
		ts.stderrBuf.Close()

		ps := ts.cmd.ProcessState
		sy := ps.Sys().(syscall.WaitStatus)

		log.Printf("Process %d exited with status %d", ps.Pid(), sy.ExitStatus())
		log.Printf(ps.String())

		ts.mu.Lock()
		if sy.ExitStatus() == 0 {
			ts.state = stateStopped
		} else {
			ts.state = stateFailed
		}
		ts.mu.Unlock()
	}()

	return nil
}

// Stop kills the process if it is still running and cleans its directory.
// It should only be called once in the lifetime of a TestServer object.
// Logs fatal if the process has already failed.
func (ts *TestServer) Stop() {
	ts.mu.RLock()
	defer ts.mu.RUnlock()

	if ts.state == stateNew {
		log.Fatal("Stop() called, but Start() was never called")
	}
	if ts.state == stateFailed {
		log.Fatalf("Stop() called, but process exited unexpectedly. Stdout:\n%s\nStderr:\n%s\n",
			ts.Stdout(), ts.Stderr())
		return
	}

	if ts.state != stateStopped {
		// Only call kill if not running. It could have exited properly.
		ts.cmd.Process.Kill()
	}

	// Only cleanup on intentional stops.
	_ = os.RemoveAll(ts.baseDir)
}

// patternInterceptor wraps an io.Writer and attempts to match the data stream
// to its regular expression pattern, calling the provided callback for all matches.
type patternInterceptor struct {
	w       io.Writer
	re      *regexp.Regexp
	onMatch func([][]byte) error

	buf      []byte
	disabled bool
}

func newPatternInterceptor(w io.Writer, re *regexp.Regexp, onMatch func([][]byte) error) *patternInterceptor {
	return &patternInterceptor{
		re:      re,
		onMatch: onMatch,
	}
}

func (pi *patternInterceptor) Write(p []byte) (n int, err error) {
	if !pi.disabled {
		// Search each full line in p for matches. Buffer partial lines.
		for sp := p; ; {
			i := bytes.IndexByte(sp, '\n')
			if i == -1 {
				pi.buf = append(pi.buf, sp...)
				break
			}

			l := sp[:i]
			if len(pi.buf) > 0 {
				l = append(pi.buf, l...)
			}

			if matches := pi.re.FindAllSubmatch(l, -1); matches != nil {
				for _, match := range matches {
					if err := pi.onMatch(match); err != nil {
						return 0, err
					}
				}
			}

			sp = sp[i+1:]
			pi.buf = pi.buf[:0]
		}
	}
	if pi.w == nil {
		return len(p), nil
	}
	return pi.w.Write(p)
}

// Disable disables the patternInterceptor from attempting to match the data
// stream to its pattern, allowing writes to pass through without inspection.
func (pi *patternInterceptor) Disable() {
	pi.disabled = true
	pi.buf = nil
}

type logWriter interface {
	Write(p []byte) (n int, err error)
	String() string
	Len() int64
	Close()
}

type fileLogWriter struct {
	filename string
	file     *os.File
}

func newFileLogWriter(file string) (*fileLogWriter, error) {
	f, err := os.Create(file)
	if err != nil {
		return nil, err
	}

	return &fileLogWriter{
		filename: file,
		file:     f,
	}, nil
}

func (w fileLogWriter) Close() {
	w.file.Close()
}

func (w fileLogWriter) Write(p []byte) (n int, err error) {
	return w.file.Write(p)
}

func (w fileLogWriter) String() string {
	b, err := ioutil.ReadFile(w.filename)
	if err == nil {
		return string(b)
	}
	return ""
}

func (w fileLogWriter) Len() int64 {
	s, err := os.Stat(w.filename)
	if err == nil {
		return s.Size()
	}
	return 0
}

func defaultEnv() map[string]string {
	vars := map[string]string{}
	u, err := user.Current()
	if err == nil {
		if _, ok := vars["USER"]; !ok {
			vars["USER"] = u.Username
		}
		if _, ok := vars["UID"]; !ok {
			vars["UID"] = u.Uid
		}
		if _, ok := vars["GID"]; !ok {
			vars["GID"] = u.Gid
		}
		if _, ok := vars["HOME"]; !ok {
			vars["HOME"] = u.HomeDir
		}
	}
	if _, ok := vars["PATH"]; !ok {
		vars["PATH"] = os.Getenv("PATH")
	}
	return vars
}
<|endoftext|>


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

// Package testserver provides helpers to run a cockroach binary within tests.
// It automatically downloads the latest cockroach binary for your platform
// (Linux-amd64 and Darwin-amd64 only for now), or attempts to run "cockroach"
// from your PATH.
//
// A normal invocation is (check err every time):
// ts, err := testserver.NewTestServer()
// err = ts.Start()
// defer ts.Stop()
// url := ts.PGURL()
//
// To use, run as follows:
//   import "github.com/cockroachdb/cockroach-go/testserver"
//   import "testing"
//   import "time"
//
//   func TestRunServer(t *testing.T) {
//      ts, err := testserver.NewTestServer()
//      if err != nil {
//        t.Fatal(err)
//      }
//      err := ts.Start()
//      if err != nil {
//        t.Fatal(err)
//      }
//      defer ts.Stop()
//
//      url := ts.PGURL()
//      if url != nil {
//        t.FatalF("url not found")
//      }
//      t.Logf("URL: %s", url.String())
//
//      db, err := sql.Open("postgres", url.String())
//      if err != nil {
//        t.Fatal(err)
//      }
//    }
package testserver

import (
	"bytes"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"testing"
	"time"

	// Import postgres driver.
	_ "github.com/lib/pq"
)

var (
	sqlURLRegexp = regexp.MustCompile(`sql:\s+(postgresql:.+)`)
	customBinary = flag.String("cockroach-binary", "", "Use specified cockroach binary")
)

const (
	stateNew = 1 + iota
	stateRunning
	stateStopped
	stateFailed
)

// TestServer is a helper to run a real cockroach node.
type TestServer struct {
	mu      sync.RWMutex
	state   int
	baseDir string
	pgURL   struct {
		set chan struct{}
		u   *url.URL
	}
	cmd        *exec.Cmd
	args       []string
	stdout     string
	stderr     string
	stdoutIntr *patternInterceptor
	stdoutBuf  logWriter
	stderrBuf  logWriter
}

// NewDBForTest creates a new CockroachDB TestServer instance and
// opens a SQL database connection to it. Returns a sql *DB instance a
// shutdown function. The caller is responsible for executing the
// returned shutdown function on exit.
func NewDBForTest(t *testing.T) (*sql.DB, func()) {
	return NewDBForTestWithDatabase(t, "")
}

// NewDBForTestWithDatabase creates a new CockroachDB TestServer
// instance and opens a SQL database connection to it. If database is
// specified, the returned connection will explicitly connect to
// it. Returns a sql *DB instance a shutdown function. The caller is
// responsible for executing the returned shutdown function on exit.
func NewDBForTestWithDatabase(t *testing.T, database string) (*sql.DB, func()) {
	ts, err := NewTestServer()
	if err != nil {
		t.Fatal(err)
	}

	err = ts.Start()
	if err != nil {
		t.Fatal(err)
	}

	url := ts.PGURL()
	if url == nil {
		t.Fatalf("url not found")
	}
	if len(database) > 0 {
		url.Path = database
	}

	db, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	if err := ts.WaitForInit(db); err != nil {
		t.Fatal(err)
	}

	return db, func() {
		_ = db.Close()
		ts.Stop()
	}
}

// NewTestServer creates a new TestServer, but does not start it.
// The cockroach binary for your OS and ARCH is downloaded automatically.
// If the download fails, we attempt just call "cockroach", hoping it is
// found in your path.
func NewTestServer() (*TestServer, error) {
	var cockroachBinary string
	var err error
	if len(*customBinary) > 0 {
		cockroachBinary = *customBinary
		log.Printf("Using custom cockroach binary: %s", cockroachBinary)
	} else if cockroachBinary, err = downloadLatestBinary(); err != nil {
		log.Printf("Failed to fetch latest binary: %s, attempting to use cockroach binary from your PATH", err)
		cockroachBinary = "cockroach"
	} else {
		log.Printf("Using automatically-downloaded binary: %s", cockroachBinary)
	}

	// Force "/tmp/" so avoid OSX's really long temp directory names
	// which get us over the socket filename length limit.
	baseDir, err := ioutil.TempDir("/tmp", "cockroach-testserver")
	if err != nil {
		return nil, fmt.Errorf("could not create temp directory: %s", err)
	}

	logDir := filepath.Join(baseDir, "logs")
	if err := os.MkdirAll(logDir, 0755); err != nil {
		return nil, fmt.Errorf("could not create logs directory: %s: %s", logDir, err)
	}

	args := []string{
		cockroachBinary,
		"start",
		"--logtostderr",
		"--insecure",
		"--host=localhost",
		"--port=0",
		"--http-port=0",
		"--store=" + baseDir,
	}

	ts := &TestServer{
		state:   stateNew,
		baseDir: baseDir,
		args:    args,
		stdout:  filepath.Join(logDir, "cockroach.stdout"),
		stderr:  filepath.Join(logDir, "cockroach.stderr"),
	}
	ts.pgURL.set = make(chan struct{})
	return ts, nil
}

// Stdout returns the entire contents of the process' stdout.
func (ts *TestServer) Stdout() string {
	return ts.stdoutBuf.String()
}

// Stderr returns the entire contents of the process' stderr.
func (ts *TestServer) Stderr() string {
	return ts.stderrBuf.String()
}

// PGURL returns the postgres connection URL to reach the started
// cockroach node.
//
// It blocks until the network URL is determined and does not timeout,
// relying instead on test timeouts.
func (ts *TestServer) PGURL() *url.URL {
	<-ts.pgURL.set
	return ts.pgURL.u
}

func (ts *TestServer) setPGURL(u *url.URL) {
	ts.pgURL.u = u
	close(ts.pgURL.set)
}

// WaitForInit retries until a connection is successfully established.
func (ts *TestServer) WaitForInit(db *sql.DB) error {
	var err error
	for i := 0; i < 50; i++ {
		if _, err = db.Query("SHOW DATABASES"); err == nil {
			return err
		}
		log.Printf("WaitForInit: %v", err)
		time.Sleep(time.Millisecond * 100)
	}
	return err
}

// Start runs the process, returning an error on any problems,
// including being unable to start, but not unexpected failure.
// It should only be called once in the lifetime of a TestServer object.
func (ts *TestServer) Start() error {
	ts.mu.Lock()
	if ts.state != stateNew {
		ts.mu.Unlock()
		return errors.New("Start() can only be called once")
	}
	ts.state = stateRunning
	ts.mu.Unlock()

	ts.cmd = exec.Command(ts.args[0], ts.args[1:]...)
	ts.cmd.Env = []string{"COCKROACH_MAX_OFFSET=1ns"}

	if len(ts.stdout) > 0 {
		wr, err := newFileLogWriter(ts.stdout)
		if err != nil {
			return fmt.Errorf("unable to open file %s: %s", ts.stdout, err)
		}
		ts.stdoutBuf = wr
	}

	pi := newPatternInterceptor(ts.stdoutBuf, sqlURLRegexp, func(match [][]byte) error {
		u, err := url.Parse(string(match[1]))
		if err != nil {
			return fmt.Errorf("failure to parse SQL URL: %v", err)
		}
		ts.setPGURL(u)
		ts.stdoutIntr.Disable()
		return nil
	})
	ts.stdoutIntr = pi
	ts.cmd.Stdout = pi

	if len(ts.stderr) > 0 {
		wr, err := newFileLogWriter(ts.stderr)
		if err != nil {
			return fmt.Errorf("unable to open file %s: %s", ts.stderr, err)
		}
		ts.stderrBuf = wr
	}
	ts.cmd.Stderr = ts.stderrBuf

	for k, v := range defaultEnv() {
		ts.cmd.Env = append(ts.cmd.Env, k+"="+v)
	}

	err := ts.cmd.Start()
	if ts.cmd.Process != nil {
		log.Printf("process %d started: %s", ts.cmd.Process.Pid, strings.Join(ts.args, " "))
	}
	if err != nil {
		log.Printf(err.Error())
		ts.stdoutBuf.Close()
		ts.stderrBuf.Close()

		ts.mu.Lock()
		ts.state = stateFailed
		ts.mu.Unlock()

		return fmt.Errorf("failure starting process: %s", err)
	}

	go func() {
		ts.cmd.Wait()

		ts.stdoutBuf.Close()
		ts.stderrBuf.Close()

		ps := ts.cmd.ProcessState
		sy := ps.Sys().(syscall.WaitStatus)

		log.Printf("Process %d exited with status %d", ps.Pid(), sy.ExitStatus())
		log.Printf(ps.String())

		ts.mu.Lock()
		if sy.ExitStatus() == 0 {
			ts.state = stateStopped
		} else {
			ts.state = stateFailed
		}
		ts.mu.Unlock()
	}()

	return nil
}

// Stop kills the process if it is still running and cleans its directory.
// It should only be called once in the lifetime of a TestServer object.
// Logs fatal if the process has already failed.
func (ts *TestServer) Stop() {
	ts.mu.RLock()
	defer ts.mu.RUnlock()

	if ts.state == stateNew {
		log.Fatal("Stop() called, but Start() was never called")
	}
	if ts.state == stateFailed {
		log.Fatalf("Stop() called, but process exited unexpectedly. Stdout:\n%s\nStderr:\n%s\n",
			ts.Stdout(), ts.Stderr())
		return
	}

	if ts.state != stateStopped {
		// Only call kill if not running. It could have exited properly.
		ts.cmd.Process.Kill()
	}

	// Only cleanup on intentional stops.
	_ = os.RemoveAll(ts.baseDir)
}

// patternInterceptor wraps an io.Writer and attempts to match the data stream
// to its regular expression pattern, calling the provided callback for all matches.
type patternInterceptor struct {
	w       io.Writer
	re      *regexp.Regexp
	onMatch func([][]byte) error

	buf      []byte
	disabled bool
}

func newPatternInterceptor(w io.Writer, re *regexp.Regexp, onMatch func([][]byte) error) *patternInterceptor {
	return &patternInterceptor{
		re:      re,
		onMatch: onMatch,
	}
}

func (pi *patternInterceptor) Write(p []byte) (n int, err error) {
	if !pi.disabled {
		// Search each full line in p for matches. Buffer partial lines.
		for sp := p; ; {
			i := bytes.IndexByte(sp, '\n')
			if i == -1 {
				pi.buf = append(pi.buf, sp...)
				break
			}

			l := sp[:i]
			if len(pi.buf) > 0 {
				l = append(pi.buf, l...)
			}

			if matches := pi.re.FindAllSubmatch(l, -1); matches != nil {
				for _, match := range matches {
					if err := pi.onMatch(match); err != nil {
						return 0, err
					}
				}
			}

			sp = sp[i+1:]
			pi.buf = pi.buf[:0]
		}
	}
	if pi.w == nil {
		return len(p), nil
	}
	return pi.w.Write(p)
}

// Disable disables the patternInterceptor from attempting to match the data
// stream to its pattern, allowing writes to pass through without inspection.
func (pi *patternInterceptor) Disable() {
	pi.disabled = true
	pi.buf = nil
}

type logWriter interface {
	Write(p []byte) (n int, err error)
	String() string
	Len() int64
	Close()
}

type fileLogWriter struct {
	filename string
	file     *os.File
}

func newFileLogWriter(file string) (*fileLogWriter, error) {
	f, err := os.Create(file)
	if err != nil {
		return nil, err
	}

	return &fileLogWriter{
		filename: file,
		file:     f,
	}, nil
}

func (w fileLogWriter) Close() {
	w.file.Close()
}

func (w fileLogWriter) Write(p []byte) (n int, err error) {
	return w.file.Write(p)
}

func (w fileLogWriter) String() string {
	b, err := ioutil.ReadFile(w.filename)
	if err == nil {
		return string(b)
	}
	return ""
}

func (w fileLogWriter) Len() int64 {
	s, err := os.Stat(w.filename)
	if err == nil {
		return s.Size()
	}
	return 0
}

func defaultEnv() map[string]string {
	vars := map[string]string{}
	u, err := user.Current()
	if err == nil {
		if _, ok := vars["USER"]; !ok {
			vars["USER"] = u.Username
		}
		if _, ok := vars["UID"]; !ok {
			vars["UID"] = u.Uid
		}
		if _, ok := vars["GID"]; !ok {
			vars["GID"] = u.Gid
		}
		if _, ok := vars["HOME"]; !ok {
			vars["HOME"] = u.HomeDir
		}
	}
	if _, ok := vars["PATH"]; !ok {
		vars["PATH"] = os.Getenv("PATH")
	}
	return vars
}
<commit_msg>testserver: use listening URL file to get SQL URL<commit_after>// Copyright 2016 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

// Package testserver provides helpers to run a cockroach binary within tests.
// It automatically downloads the latest cockroach binary for your platform
// (Linux-amd64 and Darwin-amd64 only for now), or attempts to run "cockroach"
// from your PATH.
//
// A normal invocation is (check err every time):
// ts, err := testserver.NewTestServer()
// err = ts.Start()
// defer ts.Stop()
// url := ts.PGURL()
//
// To use, run as follows:
//   import "github.com/cockroachdb/cockroach-go/testserver"
//   import "testing"
//   import "time"
//
//   func TestRunServer(t *testing.T) {
//      ts, err := testserver.NewTestServer()
//      if err != nil {
//        t.Fatal(err)
//      }
//      err := ts.Start()
//      if err != nil {
//        t.Fatal(err)
//      }
//      defer ts.Stop()
//
//      url := ts.PGURL()
//      if url != nil {
//        t.FatalF("url not found")
//      }
//      t.Logf("URL: %s", url.String())
//
//      db, err := sql.Open("postgres", url.String())
//      if err != nil {
//        t.Fatal(err)
//      }
//    }
package testserver

import (
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"testing"
	"time"

	// Import postgres driver.

	_ "github.com/lib/pq"
)

var customBinary = flag.String("cockroach-binary", "", "Use specified cockroach binary")

const (
	stateNew = 1 + iota
	stateRunning
	stateStopped
	stateFailed
)

// TestServer is a helper to run a real cockroach node.
type TestServer struct {
	mu      sync.RWMutex
	state   int
	baseDir string
	pgURL   struct {
		set chan struct{}
		u   *url.URL
	}
	cmd              *exec.Cmd
	args             []string
	stdout           string
	stderr           string
	stdoutBuf        logWriter
	stderrBuf        logWriter
	listeningURLFile string
}

// NewDBForTest creates a new CockroachDB TestServer instance and
// opens a SQL database connection to it. Returns a sql *DB instance a
// shutdown function. The caller is responsible for executing the
// returned shutdown function on exit.
func NewDBForTest(t *testing.T) (*sql.DB, func()) {
	return NewDBForTestWithDatabase(t, "")
}

// NewDBForTestWithDatabase creates a new CockroachDB TestServer
// instance and opens a SQL database connection to it. If database is
// specified, the returned connection will explicitly connect to
// it. Returns a sql *DB instance a shutdown function. The caller is
// responsible for executing the returned shutdown function on exit.
func NewDBForTestWithDatabase(t *testing.T, database string) (*sql.DB, func()) {
	ts, err := NewTestServer()
	if err != nil {
		t.Fatal(err)
	}

	err = ts.Start()
	if err != nil {
		t.Fatal(err)
	}

	url := ts.PGURL()
	if url == nil {
		t.Fatalf("url not found")
	}
	if len(database) > 0 {
		url.Path = database
	}

	db, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	if err := ts.WaitForInit(db); err != nil {
		t.Fatal(err)
	}

	return db, func() {
		_ = db.Close()
		ts.Stop()
	}
}

// NewTestServer creates a new TestServer, but does not start it.
// The cockroach binary for your OS and ARCH is downloaded automatically.
// If the download fails, we attempt just call "cockroach", hoping it is
// found in your path.
func NewTestServer() (*TestServer, error) {
	var cockroachBinary string
	var err error
	if len(*customBinary) > 0 {
		cockroachBinary = *customBinary
		log.Printf("Using custom cockroach binary: %s", cockroachBinary)
	} else if cockroachBinary, err = downloadLatestBinary(); err != nil {
		log.Printf("Failed to fetch latest binary: %s, attempting to use cockroach binary from your PATH", err)
		cockroachBinary = "cockroach"
	} else {
		log.Printf("Using automatically-downloaded binary: %s", cockroachBinary)
	}

	// Force "/tmp/" so avoid OSX's really long temp directory names
	// which get us over the socket filename length limit.
	baseDir, err := ioutil.TempDir("/tmp", "cockroach-testserver")
	if err != nil {
		return nil, fmt.Errorf("could not create temp directory: %s", err)
	}

	logDir := filepath.Join(baseDir, "logs")
	if err := os.MkdirAll(logDir, 0755); err != nil {
		return nil, fmt.Errorf("could not create logs directory: %s: %s", logDir, err)
	}

	listeningURLFile := filepath.Join(baseDir, "listen-url")

	args := []string{
		cockroachBinary,
		"start",
		"--logtostderr",
		"--insecure",
		"--host=localhost",
		"--port=0",
		"--http-port=0",
		"--store=" + baseDir,
		"--listening-url-file=" + listeningURLFile,
	}

	ts := &TestServer{
		state:            stateNew,
		baseDir:          baseDir,
		args:             args,
		stdout:           filepath.Join(logDir, "cockroach.stdout"),
		stderr:           filepath.Join(logDir, "cockroach.stderr"),
		listeningURLFile: listeningURLFile,
	}
	ts.pgURL.set = make(chan struct{})
	return ts, nil
}

// Stdout returns the entire contents of the process' stdout.
func (ts *TestServer) Stdout() string {
	return ts.stdoutBuf.String()
}

// Stderr returns the entire contents of the process' stderr.
func (ts *TestServer) Stderr() string {
	return ts.stderrBuf.String()
}

// PGURL returns the postgres connection URL to reach the started
// cockroach node.
//
// It blocks until the network URL is determined and does not timeout,
// relying instead on test timeouts.
func (ts *TestServer) PGURL() *url.URL {
	<-ts.pgURL.set
	return ts.pgURL.u
}

func (ts *TestServer) setPGURL(u *url.URL) {
	ts.pgURL.u = u
	close(ts.pgURL.set)
}

// WaitForInit retries until a connection is successfully established.
func (ts *TestServer) WaitForInit(db *sql.DB) error {
	var err error
	for i := 0; i < 50; i++ {
		if _, err = db.Query("SHOW DATABASES"); err == nil {
			return err
		}
		log.Printf("WaitForInit: %v", err)
		time.Sleep(time.Millisecond * 100)
	}
	return err
}

func (ts *TestServer) pollListeningURLFile() error {
	var data []byte
	for {
		ts.mu.Lock()
		state := ts.state
		ts.mu.Unlock()
		if state != stateRunning {
			return fmt.Errorf("server stopped or crashed before listening URL file was available")
		}

		var err error
		data, err = ioutil.ReadFile(ts.listeningURLFile)
		if err == nil {
			break
		} else if !os.IsNotExist(err) {
			return fmt.Errorf("unexpected error while reading listening URL file: %v", err)
		}
		time.Sleep(100 * time.Millisecond)
	}

	u, err := url.Parse(string(data))
	if err != nil {
		return fmt.Errorf("failed to parse SQL URL: %v", err)
	}
	ts.setPGURL(u)

	return nil
}

// Start runs the process, returning an error on any problems,
// including being unable to start, but not unexpected failure.
// It should only be called once in the lifetime of a TestServer object.
func (ts *TestServer) Start() error {
	ts.mu.Lock()
	if ts.state != stateNew {
		ts.mu.Unlock()
		return errors.New("Start() can only be called once")
	}
	ts.state = stateRunning
	ts.mu.Unlock()

	ts.cmd = exec.Command(ts.args[0], ts.args[1:]...)
	ts.cmd.Env = []string{"COCKROACH_MAX_OFFSET=1ns"}

	if len(ts.stdout) > 0 {
		wr, err := newFileLogWriter(ts.stdout)
		if err != nil {
			return fmt.Errorf("unable to open file %s: %s", ts.stdout, err)
		}
		ts.stdoutBuf = wr
	}
	ts.cmd.Stdout = ts.stdoutBuf

	if len(ts.stderr) > 0 {
		wr, err := newFileLogWriter(ts.stderr)
		if err != nil {
			return fmt.Errorf("unable to open file %s: %s", ts.stderr, err)
		}
		ts.stderrBuf = wr
	}
	ts.cmd.Stderr = ts.stderrBuf

	for k, v := range defaultEnv() {
		ts.cmd.Env = append(ts.cmd.Env, k+"="+v)
	}

	err := ts.cmd.Start()
	if ts.cmd.Process != nil {
		log.Printf("process %d started: %s", ts.cmd.Process.Pid, strings.Join(ts.args, " "))
	}
	if err != nil {
		log.Printf(err.Error())
		ts.stdoutBuf.Close()
		ts.stderrBuf.Close()

		ts.mu.Lock()
		ts.state = stateFailed
		ts.mu.Unlock()

		return fmt.Errorf("failure starting process: %s", err)
	}

	go func() {
		ts.cmd.Wait()

		ts.stdoutBuf.Close()
		ts.stderrBuf.Close()

		ps := ts.cmd.ProcessState
		sy := ps.Sys().(syscall.WaitStatus)

		log.Printf("Process %d exited with status %d", ps.Pid(), sy.ExitStatus())
		log.Printf(ps.String())

		ts.mu.Lock()
		if sy.ExitStatus() == 0 {
			ts.state = stateStopped
		} else {
			ts.state = stateFailed
		}
		ts.mu.Unlock()
	}()

	go func() {
		if err := ts.pollListeningURLFile(); err != nil {
			log.Printf("%v", err)
			close(ts.pgURL.set)
			ts.Stop()
		}
	}()

	return nil
}

// Stop kills the process if it is still running and cleans its directory.
// It should only be called once in the lifetime of a TestServer object.
// Logs fatal if the process has already failed.
func (ts *TestServer) Stop() {
	ts.mu.RLock()
	defer ts.mu.RUnlock()

	if ts.state == stateNew {
		log.Fatal("Stop() called, but Start() was never called")
	}
	if ts.state == stateFailed {
		log.Fatalf("Stop() called, but process exited unexpectedly. Stdout:\n%s\nStderr:\n%s\n",
			ts.Stdout(), ts.Stderr())
		return
	}

	if ts.state != stateStopped {
		// Only call kill if not running. It could have exited properly.
		ts.cmd.Process.Kill()
	}

	// Only cleanup on intentional stops.
	_ = os.RemoveAll(ts.baseDir)
}

type logWriter interface {
	Write(p []byte) (n int, err error)
	String() string
	Len() int64
	Close()
}

type fileLogWriter struct {
	filename string
	file     *os.File
}

func newFileLogWriter(file string) (*fileLogWriter, error) {
	f, err := os.Create(file)
	if err != nil {
		return nil, err
	}

	return &fileLogWriter{
		filename: file,
		file:     f,
	}, nil
}

func (w fileLogWriter) Close() {
	w.file.Close()
}

func (w fileLogWriter) Write(p []byte) (n int, err error) {
	return w.file.Write(p)
}

func (w fileLogWriter) String() string {
	b, err := ioutil.ReadFile(w.filename)
	if err == nil {
		return string(b)
	}
	return ""
}

func (w fileLogWriter) Len() int64 {
	s, err := os.Stat(w.filename)
	if err == nil {
		return s.Size()
	}
	return 0
}

func defaultEnv() map[string]string {
	vars := map[string]string{}
	u, err := user.Current()
	if err == nil {
		if _, ok := vars["USER"]; !ok {
			vars["USER"] = u.Username
		}
		if _, ok := vars["UID"]; !ok {
			vars["UID"] = u.Uid
		}
		if _, ok := vars["GID"]; !ok {
			vars["GID"] = u.Gid
		}
		if _, ok := vars["HOME"]; !ok {
			vars["HOME"] = u.HomeDir
		}
	}
	if _, ok := vars["PATH"]; !ok {
		vars["PATH"] = os.Getenv("PATH")
	}
	return vars
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"github.com/juju/errors"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/network"
	jujutxn "github.com/juju/txn"
)

// AddressState represents the states an IP address can be in. They are created
// in an unknown state and then either become allocated or unavailable if
// allocation fails.
type AddressState string

const (
	// AddressStateUnknown is the initial state an IP address is
	// created with.
	AddressStateUnknown AddressState = ""

	// AddressStateAllocated means that the IP address has
	// successfully been allocated by the provider and is now in use
	// by an interface on a machine.
	AddressStateAllocated AddressState = "allocated"

	// AddressStateUnavailable means that allocating the address with
	// the provider failed. We shouldn't use this address, nor should
	// we attempt to allocate it again in the future.
	AddressStateUnavailable AddressState = "unavailable"
)

// String implements fmt.Stringer.
func (s AddressState) String() string {
	if s == AddressStateUnknown {
		return "<unknown>"
	}
	return string(s)
}

// GoString implements fmt.GoStringer.
func (i *IPAddress) GoString() string {
	return i.String()
}

// IPAddress represents the state of an IP address.
type IPAddress struct {
	st  *State
	doc ipaddressDoc
}

type ipaddressDoc struct {
	DocID       string       `bson:"_id"`
	EnvUUID     string       `bson:"env-uuid"`
	Life        Life         `bson:"life"`
	SubnetId    string       `bson:"subnetid,omitempty"`
	MachineId   string       `bson:"machineid,omitempty"`
	InterfaceId string       `bson:"interfaceid,omitempty"`
	Value       string       `bson:"value"`
	Type        string       `bson:"type"`
	Scope       string       `bson:"networkscope,omitempty"`
	State       AddressState `bson:"state"`
}

// Life returns whether the IP address is Alive, Dying or Dead.
func (i *IPAddress) Life() Life {
	return i.doc.Life
}

// SubnetId returns the ID of the subnet the IP address is associated with. If
// the address is not associated with a subnet this returns "".
func (i *IPAddress) SubnetId() string {
	return i.doc.SubnetId
}

// MachineId returns the ID of the machine the IP address is associated with. If
// the address is not associated with a machine this returns "".
func (i *IPAddress) MachineId() string {
	return i.doc.MachineId
}

// InterfaceId returns the ID of the network interface the IP address is
// associated with. If the address is not associated with a netowrk interface
// this returns "".
func (i *IPAddress) InterfaceId() string {
	return i.doc.InterfaceId
}

// Value returns the IP address.
func (i *IPAddress) Value() string {
	return i.doc.Value
}

// Address returns the network.Address represent the IP address
func (i *IPAddress) Address() network.Address {
	return network.NewAddress(i.doc.Value, i.Scope())
}

// Type returns the type of the IP address. The IP address will have a type of
// IPv4, IPv6 or hostname.
func (i *IPAddress) Type() network.AddressType {
	return network.AddressType(i.doc.Type)
}

// Scope returns the scope of the IP address. If the scope is not set this
// returns "".
func (i *IPAddress) Scope() network.Scope {
	return network.Scope(i.doc.Scope)
}

// State returns the state of an IP address.
func (i *IPAddress) State() AddressState {
	return i.doc.State
}

// String implements fmt.Stringer.
func (i *IPAddress) String() string {
	return i.Address().String()
}

// EnsureDead sets the Life of the IP address to Dead, if it's Alive. It
// does nothing otherwise.
func (i *IPAddress) EnsureDead() (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot set address %q to dead", i)

	if i.doc.Life == Dead {
		return nil
	}

	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := i.Refresh(); err != nil {
				// Address is either gone or
				// another error occurred.
				return nil, err
			}
			if i.Life() == Dead {
				return nil, jujutxn.ErrNoOperations
			}
			return nil, errors.Errorf("unexpected life value: %s", i.Life().String())
		}
		return []txn.Op{{
			C:      ipaddressesC,
			Id:     i.doc.DocID,
			Update: bson.D{{"$set", bson.D{{"life", Dead}}}},
			Assert: isAliveDoc,
		}}, nil
	}

	err = i.st.run(buildTxn)
	if err != nil {
		return err
	}

	i.doc.Life = Dead
	return nil
}

// Remove removes an existing IP address. Trying to remove a missing
// address is not an error.
func (i *IPAddress) Remove() (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot remove IP address %q", i)

	if i.doc.Life != Dead {
		return errors.New("IP address is not dead")
	}

	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := i.Refresh(); errors.IsNotFound(err) {
				return nil, jujutxn.ErrNoOperations
			} else if err != nil {
				return nil, err
			}
			if i.Life() != Dead {
				return nil, errors.New("address is not dead")
			}
		}
		return []txn.Op{{
			C:      ipaddressesC,
			Id:     i.doc.DocID,
			Assert: isDeadDoc,
			Remove: true,
		}}, nil
	}

	return i.st.run(buildTxn)
}

// SetState sets the State of an IPAddress. Valid state transitions
// are Unknown to Allocated or Unavailable, as well as setting the
// same state more than once. Any other transition will result in
// returning an error satisfying errors.IsNotValid().
func (i *IPAddress) SetState(newState AddressState) (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot set IP address %q to state %q", i, newState)

	validStates := []AddressState{AddressStateUnknown, newState}
	unknownOrSame := bson.DocElem{"state", bson.D{{"$in", validStates}}}
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := i.Refresh(); errors.IsNotFound(err) {
				return nil, err
			} else if i.Life() == Dead {
				return nil, errors.New("address is dead")
			} else if err == txn.ErrAborted {
				return nil, errors.NotValidf("transition from %q", i.doc.State)
			} else if err != nil {
				return nil, err
			}

		}
		return []txn.Op{{
			C:      ipaddressesC,
			Id:     i.doc.DocID,
			Assert: append(isAliveDoc, unknownOrSame),
			Update: bson.D{{"$set", bson.D{{"state", string(newState)}}}},
		}}, nil
	}

	err = i.st.run(buildTxn)
	if err != nil {
		return err
	}

	i.doc.State = newState
	return nil
}

// AllocateTo sets the machine ID and interface ID of the IP address.
// It will fail if the state is not AddressStateUnknown. On success,
// the address state will also change to AddressStateAllocated.
func (i *IPAddress) AllocateTo(machineId, interfaceId string) (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot allocate IP address %q to machine %q, interface %q", i, machineId, interfaceId)

	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := i.Refresh(); errors.IsNotFound(err) {
				return nil, err
			} else if i.Life() == Dead {
				return nil, errors.New("address is dead")
			} else if err == txn.ErrAborted {
				return nil, errors.Errorf("already allocated or unavailable")
			} else if err != nil {
				return nil, err
			}

		}
		return []txn.Op{{
			C:      ipaddressesC,
			Id:     i.doc.DocID,
			Assert: append(isAliveDoc, bson.DocElem{"state", AddressStateUnknown}),
			Update: bson.D{{"$set", bson.D{
				{"machineid", machineId},
				{"interfaceid", interfaceId},
				{"state", string(AddressStateAllocated)},
			}}},
		}}, nil
	}

	err = i.st.run(buildTxn)
	if err != nil {
		return err
	}
	i.doc.MachineId = machineId
	i.doc.InterfaceId = interfaceId
	i.doc.State = AddressStateAllocated
	return nil
}

// Refresh refreshes the contents of the IPAddress from the underlying
// state. It an error that satisfies errors.IsNotFound if the Subnet has
// been removed.
func (i *IPAddress) Refresh() error {
	addresses, closer := i.st.getCollection(ipaddressesC)
	defer closer()

	err := addresses.FindId(i.doc.DocID).One(&i.doc)
	if err == mgo.ErrNotFound {
		return errors.NotFoundf("IP address %q", i)
	}
	if err != nil {
		return errors.Errorf("cannot refresh IP address %q: %v", i, err)
	}
	return nil
}
<commit_msg>Better error message<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"github.com/juju/errors"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/network"
	jujutxn "github.com/juju/txn"
)

// AddressState represents the states an IP address can be in. They are created
// in an unknown state and then either become allocated or unavailable if
// allocation fails.
type AddressState string

const (
	// AddressStateUnknown is the initial state an IP address is
	// created with.
	AddressStateUnknown AddressState = ""

	// AddressStateAllocated means that the IP address has
	// successfully been allocated by the provider and is now in use
	// by an interface on a machine.
	AddressStateAllocated AddressState = "allocated"

	// AddressStateUnavailable means that allocating the address with
	// the provider failed. We shouldn't use this address, nor should
	// we attempt to allocate it again in the future.
	AddressStateUnavailable AddressState = "unavailable"
)

// String implements fmt.Stringer.
func (s AddressState) String() string {
	if s == AddressStateUnknown {
		return "<unknown>"
	}
	return string(s)
}

// GoString implements fmt.GoStringer.
func (i *IPAddress) GoString() string {
	return i.String()
}

// IPAddress represents the state of an IP address.
type IPAddress struct {
	st  *State
	doc ipaddressDoc
}

type ipaddressDoc struct {
	DocID       string       `bson:"_id"`
	EnvUUID     string       `bson:"env-uuid"`
	Life        Life         `bson:"life"`
	SubnetId    string       `bson:"subnetid,omitempty"`
	MachineId   string       `bson:"machineid,omitempty"`
	InterfaceId string       `bson:"interfaceid,omitempty"`
	Value       string       `bson:"value"`
	Type        string       `bson:"type"`
	Scope       string       `bson:"networkscope,omitempty"`
	State       AddressState `bson:"state"`
}

// Life returns whether the IP address is Alive, Dying or Dead.
func (i *IPAddress) Life() Life {
	return i.doc.Life
}

// SubnetId returns the ID of the subnet the IP address is associated with. If
// the address is not associated with a subnet this returns "".
func (i *IPAddress) SubnetId() string {
	return i.doc.SubnetId
}

// MachineId returns the ID of the machine the IP address is associated with. If
// the address is not associated with a machine this returns "".
func (i *IPAddress) MachineId() string {
	return i.doc.MachineId
}

// InterfaceId returns the ID of the network interface the IP address is
// associated with. If the address is not associated with a netowrk interface
// this returns "".
func (i *IPAddress) InterfaceId() string {
	return i.doc.InterfaceId
}

// Value returns the IP address.
func (i *IPAddress) Value() string {
	return i.doc.Value
}

// Address returns the network.Address represent the IP address
func (i *IPAddress) Address() network.Address {
	return network.NewAddress(i.doc.Value, i.Scope())
}

// Type returns the type of the IP address. The IP address will have a type of
// IPv4, IPv6 or hostname.
func (i *IPAddress) Type() network.AddressType {
	return network.AddressType(i.doc.Type)
}

// Scope returns the scope of the IP address. If the scope is not set this
// returns "".
func (i *IPAddress) Scope() network.Scope {
	return network.Scope(i.doc.Scope)
}

// State returns the state of an IP address.
func (i *IPAddress) State() AddressState {
	return i.doc.State
}

// String implements fmt.Stringer.
func (i *IPAddress) String() string {
	return i.Address().String()
}

// EnsureDead sets the Life of the IP address to Dead, if it's Alive. It
// does nothing otherwise.
func (i *IPAddress) EnsureDead() (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot set address %q to dead", i)

	if i.doc.Life == Dead {
		return nil
	}

	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := i.Refresh(); err != nil {
				// Address is either gone or
				// another error occurred.
				return nil, err
			}
			if i.Life() == Dead {
				return nil, jujutxn.ErrNoOperations
			}
			return nil, errors.Errorf("unexpected life value: %s", i.Life().String())
		}
		return []txn.Op{{
			C:      ipaddressesC,
			Id:     i.doc.DocID,
			Update: bson.D{{"$set", bson.D{{"life", Dead}}}},
			Assert: isAliveDoc,
		}}, nil
	}

	err = i.st.run(buildTxn)
	if err != nil {
		return err
	}

	i.doc.Life = Dead
	return nil
}

// Remove removes an existing IP address. Trying to remove a missing
// address is not an error.
func (i *IPAddress) Remove() (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot remove IP address %q", i)

	if i.doc.Life != Dead {
		return errors.New("IP address is not dead")
	}

	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := i.Refresh(); errors.IsNotFound(err) {
				return nil, jujutxn.ErrNoOperations
			} else if err != nil {
				return nil, err
			}
			if i.Life() != Dead {
				return nil, errors.New("address is not dead")
			}
		}
		return []txn.Op{{
			C:      ipaddressesC,
			Id:     i.doc.DocID,
			Assert: isDeadDoc,
			Remove: true,
		}}, nil
	}

	return i.st.run(buildTxn)
}

// SetState sets the State of an IPAddress. Valid state transitions
// are Unknown to Allocated or Unavailable, as well as setting the
// same state more than once. Any other transition will result in
// returning an error satisfying errors.IsNotValid().
func (i *IPAddress) SetState(newState AddressState) (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot set IP address %q to state %q", i, newState)

	validStates := []AddressState{AddressStateUnknown, newState}
	unknownOrSame := bson.DocElem{"state", bson.D{{"$in", validStates}}}
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := i.Refresh(); errors.IsNotFound(err) {
				return nil, err
			} else if i.Life() == Dead {
				return nil, errors.New("address is dead")
			} else if err == txn.ErrAborted {
				return nil, errors.NotValidf("transition from %q", i.doc.State)
			} else if err != nil {
				return nil, err
			}

		}
		return []txn.Op{{
			C:      ipaddressesC,
			Id:     i.doc.DocID,
			Assert: append(isAliveDoc, unknownOrSame),
			Update: bson.D{{"$set", bson.D{{"state", string(newState)}}}},
		}}, nil
	}

	err = i.st.run(buildTxn)
	if err != nil {
		return err
	}

	i.doc.State = newState
	return nil
}

// AllocateTo sets the machine ID and interface ID of the IP address.
// It will fail if the state is not AddressStateUnknown. On success,
// the address state will also change to AddressStateAllocated.
func (i *IPAddress) AllocateTo(machineId, interfaceId string) (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot allocate IP address %q to machine %q, interface %q", i, machineId, interfaceId)

	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := i.Refresh(); errors.IsNotFound(err) {
				return nil, err
			} else if i.Life() == Dead {
				return nil, errors.New("address is dead")
			} else if err == txn.ErrAborted {
				return nil, errors.Errorf("already allocated or unavailable")
			} else if err != nil {
				return nil, err
			}

		}
		return []txn.Op{{
			C:      ipaddressesC,
			Id:     i.doc.DocID,
			Assert: append(isAliveDoc, bson.DocElem{"state", AddressStateUnknown}),
			Update: bson.D{{"$set", bson.D{
				{"machineid", machineId},
				{"interfaceid", interfaceId},
				{"state", string(AddressStateAllocated)},
			}}},
		}}, nil
	}

	err = i.st.run(buildTxn)
	if err != nil {
		return err
	}
	i.doc.MachineId = machineId
	i.doc.InterfaceId = interfaceId
	i.doc.State = AddressStateAllocated
	return nil
}

// Refresh refreshes the contents of the IPAddress from the underlying
// state. It an error that satisfies errors.IsNotFound if the Subnet has
// been removed.
func (i *IPAddress) Refresh() error {
	addresses, closer := i.st.getCollection(ipaddressesC)
	defer closer()

	err := addresses.FindId(i.doc.DocID).One(&i.doc)
	if err == mgo.ErrNotFound {
		return errors.NotFoundf("IP address %q", i)
	}
	if err != nil {
		return errors.Annotatef(err, "cannot refresh IP address %q", i)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	"fmt"
	"net"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/juju/state"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
)

type SpacesSuite struct {
	ConnSuite
}

var _ = gc.Suite(&SpacesSuite{})

func (s *SpacesSuite) addSubnets(c *gc.C, CIDRs []string) {
	for i, cidr := range CIDRs {
		ip, ipNet, err := net.ParseCIDR(cidr)
		c.Assert(err, jc.ErrorIsNil)

		// Generate the high IP address from the CIDR
		// First create a copy of the low IP address
		highIp := ip

		// By default we always get 16 bytes for each IP address. We want to
		// reduce this to 4 if we were provided an IPv4 address.
		if ip.To4() != nil {
			highIp = ip.To4()
			ip = ip.To4()
		}

		// To generate a high IP address we bitwise not each byte of the subnet
		// mask and OR it to the low IP address.
		for i, b := range ipNet.Mask {
			if i < len(ip) {
				highIp[i] |= ^b
			}
		}

		providerId := fmt.Sprintf("ProviderId%d", i)
		subnetInfo := state.SubnetInfo{
			ProviderId:        providerId,
			CIDR:              cidr,
			VLANTag:           79,
			AllocatableIPLow:  ip.String(),
			AllocatableIPHigh: highIp.String(),
			AvailabilityZone:  "AvailabilityZone",
		}
		_, err = s.State.AddSubnet(subnetInfo)
		c.Assert(err, jc.ErrorIsNil)
	}
}

func (s *SpacesSuite) assertNoSpace(c *gc.C, name string) {
	_, err := s.State.Space(name)
	c.Assert(err, gc.ErrorMatches, "space \""+name+"\" not found")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func assertSpace(c *gc.C, space *state.Space, name string, subnets []string, isPublic bool) {
	c.Assert(state.SpaceDoc(space).Name, gc.Equals, name)
	actualSubnets, err := space.Subnets()
	c.Assert(err, jc.ErrorIsNil)
	actualSubnetIds := make([]string, len(actualSubnets))
	for i, subnet := range actualSubnets {
		actualSubnetIds[i] = subnet.CIDR()
	}
	c.Assert(actualSubnetIds, jc.SameContents, subnets)
	c.Assert(state.SpaceDoc(space).IsPublic, gc.Equals, isPublic)

	c.Assert(space.Life(), gc.Equals, state.Alive)
	c.Assert(strings.HasSuffix(space.ID(), name), jc.IsTrue)
	c.Assert(space.String(), gc.Equals, name)
}

func (s *SpacesSuite) TestAddSpace(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	space, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)

	// We should get the same space back from the database
	id := space.ID()
	space, err = s.State.Space(name)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)
	c.Assert(id, gc.Equals, space.ID())
}

func (s *SpacesSuite) TestAddSpaceManySubnets(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24", "2.1.1.0/24", "3.1.1.0/24", "4.1.1.0/24", "5.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	space, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)

	// We should get the same space back from the database
	id := space.ID()
	space, err = s.State.Space(name)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)
	c.Assert(id, gc.Equals, space.ID())
}

func (s *SpacesSuite) TestAddSpaceSubnetsDoNotExist(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24"}
	isPublic := false

	_, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, gc.ErrorMatches, "cannot add space \"my-space\": subnet \"1.1.1.0/24\" not found")
	s.assertNoSpace(c, name)
}

func (s *SpacesSuite) TestAddSpaceDuplicateSpace(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	space, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)

	// We should get the same space back from the database
	id := space.ID()
	space, err = s.State.Space(name)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)
	c.Assert(id, gc.Equals, space.ID())

	// Trying to add the same space again should fail
	space, err = s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, gc.ErrorMatches, "cannot add space \"my-space\": space \"my-space\" already exists")

	// The space should still be there
	space, err = s.State.Space(name)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)
	c.Assert(id, gc.Equals, space.ID())
}

func (s *SpacesSuite) TestAddSpaceInvalidName(c *gc.C) {
	name := "-"
	subnets := []string{"1.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	_, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, gc.ErrorMatches, "cannot add space \"-\": invalid space name")
	s.assertNoSpace(c, name)
}

func (s *SpacesSuite) TestAddSpaceEmptyName(c *gc.C) {
	name := ""
	subnets := []string{"1.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	_, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, gc.ErrorMatches, "cannot add space \"\": invalid space name")
	s.assertNoSpace(c, name)
}
<commit_msg>Test Space.Subnets<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	"fmt"
	"net"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/juju/state"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
)

type SpacesSuite struct {
	ConnSuite
}

var _ = gc.Suite(&SpacesSuite{})

func (s *SpacesSuite) addSubnets(c *gc.C, CIDRs []string) {
	for i, cidr := range CIDRs {
		ip, ipNet, err := net.ParseCIDR(cidr)
		c.Assert(err, jc.ErrorIsNil)

		// Generate the high IP address from the CIDR
		// First create a copy of the low IP address
		highIp := ip

		// By default we always get 16 bytes for each IP address. We want to
		// reduce this to 4 if we were provided an IPv4 address.
		if ip.To4() != nil {
			highIp = ip.To4()
			ip = ip.To4()
		}

		// To generate a high IP address we bitwise not each byte of the subnet
		// mask and OR it to the low IP address.
		for i, b := range ipNet.Mask {
			if i < len(ip) {
				highIp[i] |= ^b
			}
		}

		providerId := fmt.Sprintf("ProviderId%d", i)
		subnetInfo := state.SubnetInfo{
			ProviderId:        providerId,
			CIDR:              cidr,
			VLANTag:           79,
			AllocatableIPLow:  ip.String(),
			AllocatableIPHigh: highIp.String(),
			AvailabilityZone:  "AvailabilityZone",
		}
		_, err = s.State.AddSubnet(subnetInfo)
		c.Assert(err, jc.ErrorIsNil)
	}
}

func (s *SpacesSuite) assertNoSpace(c *gc.C, name string) {
	_, err := s.State.Space(name)
	c.Assert(err, gc.ErrorMatches, "space \""+name+"\" not found")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func assertSpace(c *gc.C, space *state.Space, name string, subnets []string, isPublic bool) {
	c.Assert(state.SpaceDoc(space).Name, gc.Equals, name)
	actualSubnets, err := space.Subnets()
	c.Assert(err, jc.ErrorIsNil)
	actualSubnetIds := make([]string, len(actualSubnets))
	for i, subnet := range actualSubnets {
		actualSubnetIds[i] = subnet.CIDR()
	}
	c.Assert(actualSubnetIds, jc.SameContents, subnets)
	c.Assert(state.SpaceDoc(space).IsPublic, gc.Equals, isPublic)

	c.Assert(space.Life(), gc.Equals, state.Alive)
	c.Assert(strings.HasSuffix(space.ID(), name), jc.IsTrue)
	c.Assert(space.String(), gc.Equals, name)
}

func (s *SpacesSuite) TestAddSpace(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	space, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)

	// We should get the same space back from the database
	id := space.ID()
	space, err = s.State.Space(name)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)
	c.Assert(id, gc.Equals, space.ID())
}

func (s *SpacesSuite) TestAddSpaceManySubnets(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24", "2.1.1.0/24", "3.1.1.0/24", "4.1.1.0/24", "5.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	space, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)

	// We should get the same space back from the database
	id := space.ID()
	space, err = s.State.Space(name)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)
	c.Assert(id, gc.Equals, space.ID())
}

func (s *SpacesSuite) TestAddSpaceSubnetsDoNotExist(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24"}
	isPublic := false

	_, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, gc.ErrorMatches, "cannot add space \"my-space\": subnet \"1.1.1.0/24\" not found")
	s.assertNoSpace(c, name)
}

func (s *SpacesSuite) TestAddSpaceDuplicateSpace(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	space, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)

	// We should get the same space back from the database
	id := space.ID()
	space, err = s.State.Space(name)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)
	c.Assert(id, gc.Equals, space.ID())

	// Trying to add the same space again should fail
	space, err = s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, gc.ErrorMatches, "cannot add space \"my-space\": space \"my-space\" already exists")

	// The space should still be there
	space, err = s.State.Space(name)
	c.Assert(err, jc.ErrorIsNil)
	assertSpace(c, space, name, subnets, isPublic)
	c.Assert(id, gc.Equals, space.ID())
}

func (s *SpacesSuite) TestAddSpaceInvalidName(c *gc.C) {
	name := "-"
	subnets := []string{"1.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	_, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, gc.ErrorMatches, "cannot add space \"-\": invalid space name")
	s.assertNoSpace(c, name)
}

func (s *SpacesSuite) TestAddSpaceEmptyName(c *gc.C) {
	name := ""
	subnets := []string{"1.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	_, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, gc.ErrorMatches, "cannot add space \"\": invalid space name")
	s.assertNoSpace(c, name)
}

func (s *SpacesSuite) TestSpaceSubnets(c *gc.C) {
	name := "my-space"
	subnets := []string{"1.1.1.0/24", "2.1.1.0/24", "3.1.1.0/24", "4.1.1.0/24", "5.1.1.0/24"}
	isPublic := false
	s.addSubnets(c, subnets)

	space, err := s.State.AddSpace(name, subnets, isPublic)
	c.Assert(err, jc.ErrorIsNil)

	expected := []*state.Subnet{}
	for _, cidr := range subnets {
		subnet, err := s.State.Subnet(cidr)
		c.Assert(err, jc.ErrorIsNil)
		expected = append(expected, subnet)
	}
	actual, err := space.Subnets()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(actual, jc.DeepEquals, expected)
}
<|endoftext|>


<commit_before>package gateway

import (
	"net"
	"strconv"
	"testing"
	"time"

	"github.com/inconshreveable/muxado"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

func TestAddPeer(t *testing.T) {
	g := newTestingGateway("TestAddPeer", t)
	defer g.Close()
	id := g.mu.Lock()
	defer g.mu.Unlock(id)
	g.addPeer(&peer{addr: "foo", sess: muxado.Client(nil)})
	if len(g.peers) != 1 {
		t.Fatal("gateway did not add peer")
	}
}

func TestListen(t *testing.T) {
	g := newTestingGateway("TestListen", t)
	defer g.Close()

	// "compliant" connect
	conn, err := net.Dial("tcp", string(g.Address()))
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	addr := modules.NetAddress(conn.LocalAddr().String())
	// send version
	if err := encoding.WriteObject(conn, build.Version); err != nil {
		t.Fatal("couldn't write version")
	}
	// read ack
	var ack string
	if err := encoding.ReadObject(conn, &ack, maxAddrLength); err != nil {
		t.Fatal(err)
	} else if ack == "reject" {
		t.Fatal("gateway should have given ack")
	}

	// g should add the peer
	var ok bool
	for !ok {
		id := g.mu.RLock()
		_, ok = g.peers[addr]
		g.mu.RUnlock(id)
	}

	// a simple 'conn.Close' would not obey the muxado disconnect protocol
	muxado.Client(conn).Close()

	// g should remove the peer
	for ok {
		id := g.mu.RLock()
		_, ok = g.peers[addr]
		g.mu.RUnlock(id)
	}

	// "uncompliant" connect
	conn, err = net.Dial("tcp", string(g.Address()))
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	if _, err := conn.Write([]byte("missing length prefix")); err != nil {
		t.Fatal("couldn't write malformed header")
	}
	// g should have closed the connection
	if n, err := conn.Write([]byte("closed")); err != nil && n > 0 {
		t.Error("write succeeded after closed connection")
	}
}

func TestConnect(t *testing.T) {
	// create bootstrap peer
	bootstrap := newTestingGateway("TestConnect1", t)
	defer bootstrap.Close()

	// give it a node
	id := bootstrap.mu.Lock()
	bootstrap.addNode(dummyNode)
	bootstrap.mu.Unlock(id)

	// create peer who will connect to bootstrap
	g := newTestingGateway("TestConnect2", t)
	defer g.Close()

	// first simulate a "bad" connect, where bootstrap won't share its nodes
	bootstrap.RegisterRPC("ShareNodes", func(modules.PeerConn) error {
		return nil
	})
	// connect
	err := g.Connect(bootstrap.Address())
	if err != nil {
		t.Fatal(err)
	}
	// g should not have the node
	if g.removeNode(dummyNode) == nil {
		t.Fatal("bootstrapper should not have received dummyNode:", g.nodes)
	}

	// split 'em up
	g.Disconnect(bootstrap.Address())
	bootstrap.Disconnect(g.Address())

	// now restore the correct ShareNodes RPC and try again
	bootstrap.RegisterRPC("ShareNodes", bootstrap.shareNodes)
	err = g.Connect(bootstrap.Address())
	if err != nil {
		t.Fatal(err)
	}
	// g should have the node
	id = g.mu.RLock()
	defer g.mu.RUnlock(id)
	if _, ok := g.nodes[dummyNode]; !ok {
		t.Fatal("bootstrapper should have received dummyNode:", g.nodes)
	}
}

func TestDisconnect(t *testing.T) {
	g := newTestingGateway("TestDisconnect", t)
	defer g.Close()

	if err := g.Disconnect("bar"); err == nil {
		t.Fatal("disconnect removed unconnected peer")
	}

	// dummy listener to accept connection
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatal("couldn't start listener:", err)
	}
	go func() {
		conn, err := l.Accept()
		if err != nil {
			t.Fatal("accept failed:", err)
		}
		conn.Close()
	}()
	// skip standard connection protocol
	conn, err := net.Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	id := g.mu.Lock()
	g.addPeer(&peer{addr: "foo", sess: muxado.Client(conn)})
	g.mu.Unlock(id)
	if err := g.Disconnect("foo"); err != nil {
		t.Fatal("disconnect failed:", err)
	}
}

func TestMakeOutboundConnections(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	g1 := newTestingGateway("TestMakeOutboundConnections1", t)
	defer g1.Close()

	// first add 8 dummy peers
	id := g1.mu.Lock()
	for i := 0; i < 8; i++ {
		peerAddr := modules.NetAddress("foo" + strconv.Itoa(i))
		g1.peers[peerAddr] = &peer{addr: peerAddr, sess: nil}
	}
	g1.mu.Unlock(id)

	// makeOutboundConnections should now sleep for 5 seconds
	time.Sleep(1 * time.Second)

	// remove a peer while makeOutboundConnections is asleep, and add a new
	// connectable address to the node list
	id = g1.mu.Lock()
	delete(g1.peers, "foo1")
	g1.mu.Unlock(id)

	g2 := newTestingGateway("TestMakeOutboundConnections2", t)
	defer g2.Close()
	id = g1.mu.Lock()
	g1.nodes[g2.Address()] = struct{}{} // manual insertion to bypass addNode
	g1.mu.Unlock(id)

	// when makeOutboundConnections wakes up, it should connect to g2.
	time.Sleep(5 * time.Second)

	id = g1.mu.RLock()
	defer g1.mu.RUnlock(id)
	if len(g1.peers) != 8 {
		t.Fatal("gateway did not reach 8 peers:", g1.peers)
	}
	if g1.peers[g2.Address()] == nil {
		t.Fatal("gateway did not connect to g2")
	}
}
<commit_msg>add bad version test<commit_after>package gateway

import (
	"net"
	"strconv"
	"testing"
	"time"

	"github.com/inconshreveable/muxado"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

func TestAddPeer(t *testing.T) {
	g := newTestingGateway("TestAddPeer", t)
	defer g.Close()
	id := g.mu.Lock()
	defer g.mu.Unlock(id)
	g.addPeer(&peer{addr: "foo", sess: muxado.Client(nil)})
	if len(g.peers) != 1 {
		t.Fatal("gateway did not add peer")
	}
}

func TestListen(t *testing.T) {
	g := newTestingGateway("TestListen", t)
	defer g.Close()

	// compliant connect with old version
	conn, err := net.Dial("tcp", string(g.Address()))
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	addr := modules.NetAddress(conn.LocalAddr().String())
	// send version
	if err := encoding.WriteObject(conn, "0.1"); err != nil {
		t.Fatal("couldn't write version")
	}
	// read ack
	var ack string
	if err := encoding.ReadObject(conn, &ack, maxAddrLength); err != nil {
		t.Fatal(err)
	} else if ack != "reject" {
		t.Fatal("gateway should have rejected old version")
	}

	// a simple 'conn.Close' would not obey the muxado disconnect protocol
	muxado.Client(conn).Close()

	// compliant connect
	conn, err = net.Dial("tcp", string(g.Address()))
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	addr = modules.NetAddress(conn.LocalAddr().String())
	// send version
	if err := encoding.WriteObject(conn, build.Version); err != nil {
		t.Fatal("couldn't write version")
	}
	// read ack
	if err := encoding.ReadObject(conn, &ack, maxAddrLength); err != nil {
		t.Fatal(err)
	} else if ack == "reject" {
		t.Fatal("gateway should have given ack")
	}

	// g should add the peer
	var ok bool
	for !ok {
		id := g.mu.RLock()
		_, ok = g.peers[addr]
		g.mu.RUnlock(id)
	}

	muxado.Client(conn).Close()

	// g should remove the peer
	for ok {
		id := g.mu.RLock()
		_, ok = g.peers[addr]
		g.mu.RUnlock(id)
	}

	// uncompliant connect
	conn, err = net.Dial("tcp", string(g.Address()))
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	if _, err := conn.Write([]byte("missing length prefix")); err != nil {
		t.Fatal("couldn't write malformed header")
	}
	// g should have closed the connection
	if n, err := conn.Write([]byte("closed")); err != nil && n > 0 {
		t.Error("write succeeded after closed connection")
	}
}

func TestConnect(t *testing.T) {
	// create bootstrap peer
	bootstrap := newTestingGateway("TestConnect1", t)
	defer bootstrap.Close()

	// give it a node
	id := bootstrap.mu.Lock()
	bootstrap.addNode(dummyNode)
	bootstrap.mu.Unlock(id)

	// create peer who will connect to bootstrap
	g := newTestingGateway("TestConnect2", t)
	defer g.Close()

	// first simulate a "bad" connect, where bootstrap won't share its nodes
	bootstrap.RegisterRPC("ShareNodes", func(modules.PeerConn) error {
		return nil
	})
	// connect
	err := g.Connect(bootstrap.Address())
	if err != nil {
		t.Fatal(err)
	}
	// g should not have the node
	if g.removeNode(dummyNode) == nil {
		t.Fatal("bootstrapper should not have received dummyNode:", g.nodes)
	}

	// split 'em up
	g.Disconnect(bootstrap.Address())
	bootstrap.Disconnect(g.Address())

	// now restore the correct ShareNodes RPC and try again
	bootstrap.RegisterRPC("ShareNodes", bootstrap.shareNodes)
	err = g.Connect(bootstrap.Address())
	if err != nil {
		t.Fatal(err)
	}
	// g should have the node
	id = g.mu.RLock()
	defer g.mu.RUnlock(id)
	if _, ok := g.nodes[dummyNode]; !ok {
		t.Fatal("bootstrapper should have received dummyNode:", g.nodes)
	}
}

func TestDisconnect(t *testing.T) {
	g := newTestingGateway("TestDisconnect", t)
	defer g.Close()

	if err := g.Disconnect("bar"); err == nil {
		t.Fatal("disconnect removed unconnected peer")
	}

	// dummy listener to accept connection
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatal("couldn't start listener:", err)
	}
	go func() {
		conn, err := l.Accept()
		if err != nil {
			t.Fatal("accept failed:", err)
		}
		conn.Close()
	}()
	// skip standard connection protocol
	conn, err := net.Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	id := g.mu.Lock()
	g.addPeer(&peer{addr: "foo", sess: muxado.Client(conn)})
	g.mu.Unlock(id)
	if err := g.Disconnect("foo"); err != nil {
		t.Fatal("disconnect failed:", err)
	}
}

func TestMakeOutboundConnections(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	g1 := newTestingGateway("TestMakeOutboundConnections1", t)
	defer g1.Close()

	// first add 8 dummy peers
	id := g1.mu.Lock()
	for i := 0; i < 8; i++ {
		peerAddr := modules.NetAddress("foo" + strconv.Itoa(i))
		g1.peers[peerAddr] = &peer{addr: peerAddr, sess: nil}
	}
	g1.mu.Unlock(id)

	// makeOutboundConnections should now sleep for 5 seconds
	time.Sleep(1 * time.Second)

	// remove a peer while makeOutboundConnections is asleep, and add a new
	// connectable address to the node list
	id = g1.mu.Lock()
	delete(g1.peers, "foo1")
	g1.mu.Unlock(id)

	g2 := newTestingGateway("TestMakeOutboundConnections2", t)
	defer g2.Close()
	id = g1.mu.Lock()
	g1.nodes[g2.Address()] = struct{}{} // manual insertion to bypass addNode
	g1.mu.Unlock(id)

	// when makeOutboundConnections wakes up, it should connect to g2.
	time.Sleep(5 * time.Second)

	id = g1.mu.RLock()
	defer g1.mu.RUnlock(id)
	if len(g1.peers) != 8 {
		t.Fatal("gateway did not reach 8 peers:", g1.peers)
	}
	if g1.peers[g2.Address()] == nil {
		t.Fatal("gateway did not connect to g2")
	}
}
<|endoftext|>