<commit_before>package endpoints

import (
	"fmt"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	tomb "gopkg.in/tomb.v2"
)

// Config holds various configuration values that affect LXD endpoints
// initialization.
type Config struct {
	// The LXD var directory to create Unix sockets in.
	Dir string

	// UnixSocket is the path to the Unix socket to bind
	UnixSocket string

	// HTTP server handling requests for the LXD RESTful API.
	RestServer *http.Server

	// HTTP server for the internal /dev/lxd API exposed to containers.
	DevLxdServer *http.Server

	// The TLS keypair and optional CA to use for the network endpoint. It
	// must be always provided, since the pubblic key will be included in
	// the response of the /1.0 REST API as part of the server info.
	//
	// It can be updated after the endpoints are up using NetworkUpdateCert().
	Cert *shared.CertInfo

	// System group name to which the unix socket for the local endpoint should be
	// chgrp'ed when starting. The default is to use the process group. An empty
	// string means "use the default".
	LocalUnixSocketGroup string

	// NetworkSetAddress sets the address for the network endpoint. If not
	// set, the network endpoint won't be started (unless it's passed via
	// socket-based activation).
	//
	// It can be updated after the endpoints are up using NetworkUpdateAddress().
	NetworkAddress string

	// Optional dedicated network address for clustering traffic. If not
	// set, NetworkAddress will be used.
	//
	// It can be updated after the endpoints are up using ClusterUpdateAddress().
	ClusterAddress string

	// Address of the debug endpoint.
	//
	// It can be updated after the endpoints are up using PprofUpdateAddress().
	DebugAddress string
}

// Up brings up all applicable LXD endpoints and starts accepting HTTP
// requests.
//
// The endpoints will be activated in the following order and according to the
// following rules:
//
// local endpoint (unix socket)
// ----------------------------
//
// If socket-based activation is detected, look for a unix socket among the
// inherited file descriptors and use it for the local endpoint (or if no such
// file descriptor exists, don't bring up the local endpoint at all).
//
// If no socket-based activation is detected, create a unix socket using the
// default <lxd-var-dir>/unix.socket path. The file mode of this socket will be set
// to 660, the file owner will be set to the process' UID, and the file group
// will be set to the process GID, or to the GID of the system group name
// specified via config.LocalUnixSocketGroup.
//
// devlxd endpoint (unix socket)
// ----------------------------
//
// Created using <lxd-var-dir>/devlxd/sock, with file mode set to 666 (actual
// authorization will be performed by the HTTP server using the socket ucred
// struct).
//
// remote endpoint (TCP socket with TLS)
// -------------------------------------
//
// If socket-based activation is detected, look for a network socket among the
// inherited file descriptors and use it for the network endpoint.
//
// If a network address was set via config.NetworkAddress, then close any listener
// that was detected via socket-based activation and create a new network
// socket bound to the given address.
//
// The network endpoint socket will use TLS encryption, using the certificate
// keypair and CA passed via config.Cert.
//
// cluster endpoint (TCP socket with TLS)
// -------------------------------------
//
// If a network address was set via config.ClusterAddress, then attach
// config.RestServer to it.
func Up(config *Config) (*Endpoints, error) {
	if config.Dir == "" {
		return nil, fmt.Errorf("No directory configured")
	}
	if config.UnixSocket == "" {
		return nil, fmt.Errorf("No unix socket configured")
	}
	if config.RestServer == nil {
		return nil, fmt.Errorf("No REST server configured")
	}
	if config.DevLxdServer == nil {
		return nil, fmt.Errorf("No devlxd server configured")
	}
	if config.Cert == nil {
		return nil, fmt.Errorf("No TLS certificate configured")
	}

	endpoints := &Endpoints{
		systemdListenFDsStart: util.SystemdListenFDsStart,
	}

	err := endpoints.up(config)
	if err != nil {
		endpoints.Down()
		return nil, err
	}
	return endpoints, nil
}

// Endpoints are in charge of bringing up and down the HTTP endpoints for
// serving the LXD RESTful API.
//
// When LXD starts up, they start listen to the appropriate sockets and attach
// the relevant HTTP handlers to them. When LXD shuts down they close all
// sockets.
type Endpoints struct {
	tomb      *tomb.Tomb            // Controls the HTTP servers shutdown.
	mu        sync.RWMutex          // Serialize access to internal state.
	listeners map[kind]net.Listener // Activer listeners by endpoint type.
	servers   map[kind]*http.Server // HTTP servers by endpoint type.
	cert      *shared.CertInfo      // Keypair and CA to use for TLS.
	inherited map[kind]bool         // Store whether the listener came through socket activation

	systemdListenFDsStart int // First socket activation FD, for tests.
}

// Up brings up all configured endpoints and starts accepting HTTP requests.
func (e *Endpoints) up(config *Config) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	e.servers = map[kind]*http.Server{
		devlxd:  config.DevLxdServer,
		local:   config.RestServer,
		network: config.RestServer,
		cluster: config.RestServer,
		pprof:   pprofCreateServer(),
	}
	e.cert = config.Cert
	e.inherited = map[kind]bool{}

	var err error

	// Check for socket activation.
	systemdListeners := util.GetListeners(e.systemdListenFDsStart)
	if len(systemdListeners) > 0 {
		e.listeners = activatedListeners(systemdListeners, e.cert)
		for kind := range e.listeners {
			e.inherited[kind] = true
		}
	} else {
		e.listeners = map[kind]net.Listener{}

		e.listeners[local], err = localCreateListener(config.UnixSocket, config.LocalUnixSocketGroup)
		if err != nil {
			return fmt.Errorf("local endpoint: %v", err)
		}
	}

	// Start the devlxd listener
	e.listeners[devlxd], err = createDevLxdlListener(config.Dir)
	if err != nil {
		return err
	}

	if config.NetworkAddress != "" {
		listener, ok := e.listeners[network]
		if ok {
			logger.Infof("Replacing inherited TCP socket with configured one")
			listener.Close()
			e.inherited[network] = false
		}

		// Errors here are not fatal and are just logged (unless we're clustered, see below).
		var networkAddressErr error
		attempts := 0
	againHttps:
		e.listeners[network], networkAddressErr = networkCreateListener(config.NetworkAddress, e.cert)

		isCovered := util.IsAddressCovered(config.ClusterAddress, config.NetworkAddress)
		if config.ClusterAddress != "" {
			if isCovered {
				// In case of clustering we fail if we can't bind the network address.
				if networkAddressErr != nil {
					if attempts == 0 {
						logger.Infof("Unable to bind https address %q, re-trying for a minute", config.NetworkAddress)
					}

					attempts++
					if attempts < 60 {
						time.Sleep(1 * time.Second)
						goto againHttps
					}

					return networkAddressErr
				}

				logger.Infof("Starting cluster handler:")
				e.serve(cluster)
			}
		} else if networkAddressErr != nil {
			logger.Error("Cannot currently listen on https socket, re-trying once in 30s...", log.Ctx{"err": networkAddressErr})

			go func() {
				time.Sleep(30 * time.Second)
				err := e.NetworkUpdateAddress(config.NetworkAddress)
				if err != nil {
					logger.Error("Still unable to listen on https socket", log.Ctx{"err": err})
				}
			}()
		}
	}

	isCovered := util.IsAddressCovered(config.ClusterAddress, config.NetworkAddress)
	if config.ClusterAddress != "" && !isCovered {
		attempts := 0
	againCluster:
		e.listeners[cluster], err = networkCreateListener(config.ClusterAddress, e.cert)
		if err != nil {
			if attempts == 0 {
				logger.Infof("Unable to bind cluster address %q, re-trying for a minute", config.ClusterAddress)
			}

			attempts++
			if attempts < 60 {
				time.Sleep(1 * time.Second)
				goto againCluster
			}

			return err
		}

		logger.Infof("Starting cluster handler:")
		e.serve(cluster)
	}

	if config.DebugAddress != "" {
		e.listeners[pprof], err = pprofCreateListener(config.DebugAddress)
		if err != nil {
			return err
		}

		logger.Infof("Starting pprof handler:")
		e.serve(pprof)
	}

	logger.Infof("Starting /dev/lxd handler:")
	e.serve(devlxd)

	logger.Infof("REST API daemon:")
	e.serve(local)
	e.serve(network)

	return nil
}

// Down brings down all endpoints and stops serving HTTP requests.
func (e *Endpoints) Down() error {
	e.mu.Lock()
	defer e.mu.Unlock()

	if e.listeners[network] != nil || e.listeners[local] != nil {
		logger.Infof("Stopping REST API handler:")
		err := e.closeListener(network)
		if err != nil {
			return err
		}

		err = e.closeListener(local)
		if err != nil {
			return err
		}
	}

	if e.listeners[cluster] != nil {
		logger.Infof("Stopping cluster handler:")
		err := e.closeListener(cluster)
		if err != nil {
			return err
		}
	}

	if e.listeners[devlxd] != nil {
		logger.Infof("Stopping /dev/lxd handler:")
		err := e.closeListener(devlxd)
		if err != nil {
			return err
		}
	}

	if e.listeners[pprof] != nil {
		logger.Infof("Stopping pprof handler:")
		err := e.closeListener(pprof)
		if err != nil {
			return err
		}
	}

	if e.tomb != nil {
		e.tomb.Kill(nil)
		e.tomb.Wait()
	}

	return nil
}

// Start an HTTP server for the endpoint associated with the given code.
func (e *Endpoints) serve(kind kind) {
	listener := e.listeners[kind]

	if listener == nil {
		return
	}

	ctx := log.Ctx{"socket": listener.Addr()}
	if e.inherited[kind] {
		ctx["inherited"] = true
	}

	message := fmt.Sprintf(" - binding %s", descriptions[kind])
	logger.Info(message, ctx)

	server := e.servers[kind]

	// Defer the creation of the tomb, so Down() doesn't wait on it unless
	// we actually have spawned at least a server.
	if e.tomb == nil {
		e.tomb = &tomb.Tomb{}
	}

	e.tomb.Go(func() error {
		server.Serve(listener)
		return nil
	})
}

// Stop the HTTP server of the endpoint associated with the given code. The
// associated socket will be shutdown too.
func (e *Endpoints) closeListener(kind kind) error {
	listener := e.listeners[kind]
	if listener == nil {
		return nil
	}
	delete(e.listeners, kind)

	logger.Info(" - closing socket", log.Ctx{"socket": listener.Addr()})

	return listener.Close()
}

// Use the listeners associated with the file descriptors passed via
// socket-based activation.
func activatedListeners(systemdListeners []net.Listener, cert *shared.CertInfo) map[kind]net.Listener {
	listeners := map[kind]net.Listener{}
	for _, listener := range systemdListeners {
		var kind kind
		switch listener.(type) {
		case *net.UnixListener:
			kind = local
		case *net.TCPListener:
			kind = network
			listener = networkTLSListener(listener, cert)
		default:
			continue
		}
		listeners[kind] = listener
	}
	return listeners
}

// Numeric code identifying a specific API endpoint type.
type kind int

// Numeric codes identifying the various endpoints.
const (
	local kind = iota
	devlxd
	network
	pprof
	cluster
)

// Human-readable descriptions of the various kinds of endpoints.
var descriptions = map[kind]string{
	local:   "Unix socket",
	devlxd:  "devlxd socket",
	network: "TCP socket",
	pprof:   "pprof socket",
	cluster: "cluster socket",
}
<commit_msg>lxd/endpoints/endpoints: require set network listener before checking coverage<commit_after>package endpoints

import (
	"fmt"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	tomb "gopkg.in/tomb.v2"
)

// Config holds various configuration values that affect LXD endpoints
// initialization.
type Config struct {
	// The LXD var directory to create Unix sockets in.
	Dir string

	// UnixSocket is the path to the Unix socket to bind
	UnixSocket string

	// HTTP server handling requests for the LXD RESTful API.
	RestServer *http.Server

	// HTTP server for the internal /dev/lxd API exposed to containers.
	DevLxdServer *http.Server

	// The TLS keypair and optional CA to use for the network endpoint. It
	// must be always provided, since the pubblic key will be included in
	// the response of the /1.0 REST API as part of the server info.
	//
	// It can be updated after the endpoints are up using NetworkUpdateCert().
	Cert *shared.CertInfo

	// System group name to which the unix socket for the local endpoint should be
	// chgrp'ed when starting. The default is to use the process group. An empty
	// string means "use the default".
	LocalUnixSocketGroup string

	// NetworkSetAddress sets the address for the network endpoint. If not
	// set, the network endpoint won't be started (unless it's passed via
	// socket-based activation).
	//
	// It can be updated after the endpoints are up using NetworkUpdateAddress().
	NetworkAddress string

	// Optional dedicated network address for clustering traffic. If not
	// set, NetworkAddress will be used.
	//
	// It can be updated after the endpoints are up using ClusterUpdateAddress().
	ClusterAddress string

	// Address of the debug endpoint.
	//
	// It can be updated after the endpoints are up using PprofUpdateAddress().
	DebugAddress string
}

// Up brings up all applicable LXD endpoints and starts accepting HTTP
// requests.
//
// The endpoints will be activated in the following order and according to the
// following rules:
//
// local endpoint (unix socket)
// ----------------------------
//
// If socket-based activation is detected, look for a unix socket among the
// inherited file descriptors and use it for the local endpoint (or if no such
// file descriptor exists, don't bring up the local endpoint at all).
//
// If no socket-based activation is detected, create a unix socket using the
// default <lxd-var-dir>/unix.socket path. The file mode of this socket will be set
// to 660, the file owner will be set to the process' UID, and the file group
// will be set to the process GID, or to the GID of the system group name
// specified via config.LocalUnixSocketGroup.
//
// devlxd endpoint (unix socket)
// ----------------------------
//
// Created using <lxd-var-dir>/devlxd/sock, with file mode set to 666 (actual
// authorization will be performed by the HTTP server using the socket ucred
// struct).
//
// remote endpoint (TCP socket with TLS)
// -------------------------------------
//
// If socket-based activation is detected, look for a network socket among the
// inherited file descriptors and use it for the network endpoint.
//
// If a network address was set via config.NetworkAddress, then close any listener
// that was detected via socket-based activation and create a new network
// socket bound to the given address.
//
// The network endpoint socket will use TLS encryption, using the certificate
// keypair and CA passed via config.Cert.
//
// cluster endpoint (TCP socket with TLS)
// -------------------------------------
//
// If a network address was set via config.ClusterAddress, then attach
// config.RestServer to it.
func Up(config *Config) (*Endpoints, error) {
	if config.Dir == "" {
		return nil, fmt.Errorf("No directory configured")
	}
	if config.UnixSocket == "" {
		return nil, fmt.Errorf("No unix socket configured")
	}
	if config.RestServer == nil {
		return nil, fmt.Errorf("No REST server configured")
	}
	if config.DevLxdServer == nil {
		return nil, fmt.Errorf("No devlxd server configured")
	}
	if config.Cert == nil {
		return nil, fmt.Errorf("No TLS certificate configured")
	}

	endpoints := &Endpoints{
		systemdListenFDsStart: util.SystemdListenFDsStart,
	}

	err := endpoints.up(config)
	if err != nil {
		endpoints.Down()
		return nil, err
	}
	return endpoints, nil
}

// Endpoints are in charge of bringing up and down the HTTP endpoints for
// serving the LXD RESTful API.
//
// When LXD starts up, they start listen to the appropriate sockets and attach
// the relevant HTTP handlers to them. When LXD shuts down they close all
// sockets.
type Endpoints struct {
	tomb      *tomb.Tomb            // Controls the HTTP servers shutdown.
	mu        sync.RWMutex          // Serialize access to internal state.
	listeners map[kind]net.Listener // Activer listeners by endpoint type.
	servers   map[kind]*http.Server // HTTP servers by endpoint type.
	cert      *shared.CertInfo      // Keypair and CA to use for TLS.
	inherited map[kind]bool         // Store whether the listener came through socket activation

	systemdListenFDsStart int // First socket activation FD, for tests.
}

// Up brings up all configured endpoints and starts accepting HTTP requests.
func (e *Endpoints) up(config *Config) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	e.servers = map[kind]*http.Server{
		devlxd:  config.DevLxdServer,
		local:   config.RestServer,
		network: config.RestServer,
		cluster: config.RestServer,
		pprof:   pprofCreateServer(),
	}
	e.cert = config.Cert
	e.inherited = map[kind]bool{}

	var err error

	// Check for socket activation.
	systemdListeners := util.GetListeners(e.systemdListenFDsStart)
	if len(systemdListeners) > 0 {
		e.listeners = activatedListeners(systemdListeners, e.cert)
		for kind := range e.listeners {
			e.inherited[kind] = true
		}
	} else {
		e.listeners = map[kind]net.Listener{}

		e.listeners[local], err = localCreateListener(config.UnixSocket, config.LocalUnixSocketGroup)
		if err != nil {
			return fmt.Errorf("local endpoint: %v", err)
		}
	}

	// Start the devlxd listener
	e.listeners[devlxd], err = createDevLxdlListener(config.Dir)
	if err != nil {
		return err
	}

	if config.NetworkAddress != "" {
		listener, ok := e.listeners[network]
		if ok {
			logger.Infof("Replacing inherited TCP socket with configured one")
			listener.Close()
			e.inherited[network] = false
		}

		// Errors here are not fatal and are just logged (unless we're clustered, see below).
		var networkAddressErr error
		attempts := 0
	againHttps:
		e.listeners[network], networkAddressErr = networkCreateListener(config.NetworkAddress, e.cert)

		isCovered := util.IsAddressCovered(config.ClusterAddress, config.NetworkAddress)
		if config.ClusterAddress != "" {
			if isCovered {
				// In case of clustering we fail if we can't bind the network address.
				if networkAddressErr != nil {
					if attempts == 0 {
						logger.Infof("Unable to bind https address %q, re-trying for a minute", config.NetworkAddress)
					}

					attempts++
					if attempts < 60 {
						time.Sleep(1 * time.Second)
						goto againHttps
					}

					return networkAddressErr
				}

				logger.Infof("Starting cluster handler:")
				e.serve(cluster)
			}
		} else if networkAddressErr != nil {
			logger.Error("Cannot currently listen on https socket, re-trying once in 30s...", log.Ctx{"err": networkAddressErr})

			go func() {
				time.Sleep(30 * time.Second)
				err := e.NetworkUpdateAddress(config.NetworkAddress)
				if err != nil {
					logger.Error("Still unable to listen on https socket", log.Ctx{"err": err})
				}
			}()
		}
	}

	isCovered := false
	if config.NetworkAddress != "" {
		isCovered = util.IsAddressCovered(config.ClusterAddress, config.NetworkAddress)
	}

	if config.ClusterAddress != "" && !isCovered {
		attempts := 0
	againCluster:
		e.listeners[cluster], err = networkCreateListener(config.ClusterAddress, e.cert)
		if err != nil {
			if attempts == 0 {
				logger.Infof("Unable to bind cluster address %q, re-trying for a minute", config.ClusterAddress)
			}

			attempts++
			if attempts < 60 {
				time.Sleep(1 * time.Second)
				goto againCluster
			}

			return err
		}

		logger.Infof("Starting cluster handler:")
		e.serve(cluster)
	}

	if config.DebugAddress != "" {
		e.listeners[pprof], err = pprofCreateListener(config.DebugAddress)
		if err != nil {
			return err
		}

		logger.Infof("Starting pprof handler:")
		e.serve(pprof)
	}

	logger.Infof("Starting /dev/lxd handler:")
	e.serve(devlxd)

	logger.Infof("REST API daemon:")
	e.serve(local)
	e.serve(network)

	return nil
}

// Down brings down all endpoints and stops serving HTTP requests.
func (e *Endpoints) Down() error {
	e.mu.Lock()
	defer e.mu.Unlock()

	if e.listeners[network] != nil || e.listeners[local] != nil {
		logger.Infof("Stopping REST API handler:")
		err := e.closeListener(network)
		if err != nil {
			return err
		}

		err = e.closeListener(local)
		if err != nil {
			return err
		}
	}

	if e.listeners[cluster] != nil {
		logger.Infof("Stopping cluster handler:")
		err := e.closeListener(cluster)
		if err != nil {
			return err
		}
	}

	if e.listeners[devlxd] != nil {
		logger.Infof("Stopping /dev/lxd handler:")
		err := e.closeListener(devlxd)
		if err != nil {
			return err
		}
	}

	if e.listeners[pprof] != nil {
		logger.Infof("Stopping pprof handler:")
		err := e.closeListener(pprof)
		if err != nil {
			return err
		}
	}

	if e.tomb != nil {
		e.tomb.Kill(nil)
		e.tomb.Wait()
	}

	return nil
}

// Start an HTTP server for the endpoint associated with the given code.
func (e *Endpoints) serve(kind kind) {
	listener := e.listeners[kind]

	if listener == nil {
		return
	}

	ctx := log.Ctx{"socket": listener.Addr()}
	if e.inherited[kind] {
		ctx["inherited"] = true
	}

	message := fmt.Sprintf(" - binding %s", descriptions[kind])
	logger.Info(message, ctx)

	server := e.servers[kind]

	// Defer the creation of the tomb, so Down() doesn't wait on it unless
	// we actually have spawned at least a server.
	if e.tomb == nil {
		e.tomb = &tomb.Tomb{}
	}

	e.tomb.Go(func() error {
		server.Serve(listener)
		return nil
	})
}

// Stop the HTTP server of the endpoint associated with the given code. The
// associated socket will be shutdown too.
func (e *Endpoints) closeListener(kind kind) error {
	listener := e.listeners[kind]
	if listener == nil {
		return nil
	}
	delete(e.listeners, kind)

	logger.Info(" - closing socket", log.Ctx{"socket": listener.Addr()})

	return listener.Close()
}

// Use the listeners associated with the file descriptors passed via
// socket-based activation.
func activatedListeners(systemdListeners []net.Listener, cert *shared.CertInfo) map[kind]net.Listener {
	listeners := map[kind]net.Listener{}
	for _, listener := range systemdListeners {
		var kind kind
		switch listener.(type) {
		case *net.UnixListener:
			kind = local
		case *net.TCPListener:
			kind = network
			listener = networkTLSListener(listener, cert)
		default:
			continue
		}
		listeners[kind] = listener
	}
	return listeners
}

// Numeric code identifying a specific API endpoint type.
type kind int

// Numeric codes identifying the various endpoints.
const (
	local kind = iota
	devlxd
	network
	pprof
	cluster
)

// Human-readable descriptions of the various kinds of endpoints.
var descriptions = map[kind]string{
	local:   "Unix socket",
	devlxd:  "devlxd socket",
	network: "TCP socket",
	pprof:   "pprof socket",
	cluster: "cluster socket",
}
<|endoftext|>


<commit_before>package fakes

import (
	"fmt"
	"strings"
	"sync"
	"time"

	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type FakeCmdRunner struct {
	commandResults     map[string][]FakeCmdResult
	commandResultsLock sync.Mutex

	processes     map[string][]*FakeProcess
	processesLock sync.Mutex

	RunComplexCommands   []boshsys.Command
	RunCommands          [][]string
	RunCommandsWithInput [][]string

	CommandExistsValue bool
	AvailableCommands  map[string]bool
}

type FakeCmdResult struct {
	Stdout     string
	Stderr     string
	ExitStatus int
	Error      error

	Sticky bool // Set to true if this result should ALWAYS be returned for the given command
}

type FakeProcess struct {
	WaitCh chan boshsys.Result

	Waited     bool
	WaitResult boshsys.Result

	TerminatedNicely               bool
	TerminatedNicelyCallBack       func(*FakeProcess)
	TerminateNicelyKillGracePeriod time.Duration
	TerminateNicelyErr             error
}

func (p *FakeProcess) Wait() <-chan boshsys.Result {
	if p.Waited {
		panic("Cannot Wait() on process multiple times")
	}

	p.Waited = true
	p.WaitCh = make(chan boshsys.Result, 1)

	if p.TerminatedNicelyCallBack == nil {
		p.WaitCh <- p.WaitResult
	}
	return p.WaitCh
}

func (p *FakeProcess) TerminateNicely(killGracePeriod time.Duration) error {
	p.TerminateNicelyKillGracePeriod = killGracePeriod
	p.TerminatedNicely = true
	if p.TerminatedNicelyCallBack != nil {
		p.TerminatedNicelyCallBack(p)
	}
	return p.TerminateNicelyErr
}

func NewFakeCmdRunner() *FakeCmdRunner {
	return &FakeCmdRunner{
		AvailableCommands: map[string]bool{},
		commandResults:    map[string][]FakeCmdResult{},
		processes:         map[string][]*FakeProcess{},
	}
}

func (r *FakeCmdRunner) RunComplexCommand(cmd boshsys.Command) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	r.RunComplexCommands = append(r.RunComplexCommands, cmd)

	runCmd := append([]string{cmd.Name}, cmd.Args...)
	stdout, stderr, exitstatus, err := r.getOutputsForCmd(runCmd)

	if cmd.Stdout != nil {
		cmd.Stdout.Write([]byte(stdout))
	}

	if cmd.Stderr != nil {
		cmd.Stderr.Write([]byte(stderr))
	}

	return stdout, stderr, exitstatus, err
}

func (r *FakeCmdRunner) RunComplexCommandAsync(cmd boshsys.Command) (boshsys.Process, error) {
	r.processesLock.Lock()
	defer r.processesLock.Unlock()

	r.RunComplexCommands = append(r.RunComplexCommands, cmd)

	runCmd := append([]string{cmd.Name}, cmd.Args...)
	fullCmd := strings.Join(runCmd, " ")
	results, found := r.processes[fullCmd]
	if !found {
		panic(fmt.Sprintf("Failed to find process for %s", fullCmd))
	}

	return results[0], nil
}

func (r *FakeCmdRunner) RunCommand(cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{cmdName}, args...)
	r.RunCommands = append(r.RunCommands, runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) RunCommandWithInput(input, cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{input, cmdName}, args...)
	r.RunCommandsWithInput = append(r.RunCommandsWithInput, runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) CommandExists(cmdName string) bool {
	return r.CommandExistsValue || r.AvailableCommands[cmdName]
}

func (r *FakeCmdRunner) AddCmdResult(fullCmd string, result FakeCmdResult) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	results := r.commandResults[fullCmd]
	r.commandResults[fullCmd] = append(results, result)
}

func (r *FakeCmdRunner) AddProcess(fullCmd string, process *FakeProcess) {
	r.processesLock.Lock()
	defer r.processesLock.Unlock()

	processes := r.processes[fullCmd]
	r.processes[fullCmd] = append(processes, process)
}

func (r *FakeCmdRunner) getOutputsForCmd(runCmd []string) (string, string, int, error) {
	fullCmd := strings.Join(runCmd, " ")
	results, found := r.commandResults[fullCmd]
	if found {
		result := results[0]
		newResults := []FakeCmdResult{}
		if len(results) > 1 {
			newResults = results[1:]
		}

		if !result.Sticky {
			r.commandResults[fullCmd] = newResults
		}

		return result.Stdout, result.Stderr, result.ExitStatus, result.Error
	}

	return "", "", -1, nil
}
<commit_msg>Sync changes made in cloudfoundry/bosh-agent#44<commit_after>package fakes

import (
	"fmt"
	"strings"
	"sync"
	"time"

	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type FakeCmdRunner struct {
	commandResults     map[string][]FakeCmdResult
	commandResultsLock sync.Mutex

	processes     map[string][]*FakeProcess
	processesLock sync.Mutex

	RunComplexCommands   []boshsys.Command
	RunCommands          [][]string
	RunCommandsWithInput [][]string

	CommandExistsValue bool
	AvailableCommands  map[string]bool
}

type FakeCmdResult struct {
	Stdout     string
	Stderr     string
	ExitStatus int
	Error      error

	Sticky bool // Set to true if this result should ALWAYS be returned for the given command
}

type FakeProcess struct {
	WaitCh chan boshsys.Result

	Waited     bool
	WaitResult boshsys.Result

	TerminatedNicely               bool
	TerminatedNicelyCallBack       func(*FakeProcess)
	TerminateNicelyKillGracePeriod time.Duration
	TerminateNicelyErr             error
}

func (p *FakeProcess) Wait() <-chan boshsys.Result {
	if p.Waited {
		panic("Cannot Wait() on process multiple times")
	}

	p.Waited = true
	p.WaitCh = make(chan boshsys.Result, 1)

	if p.TerminatedNicelyCallBack == nil {
		p.WaitCh <- p.WaitResult
	}
	return p.WaitCh
}

func (p *FakeProcess) TerminateNicely(killGracePeriod time.Duration) error {
	p.TerminateNicelyKillGracePeriod = killGracePeriod
	p.TerminatedNicely = true
	if p.TerminatedNicelyCallBack != nil {
		p.TerminatedNicelyCallBack(p)
	}
	return p.TerminateNicelyErr
}

func NewFakeCmdRunner() *FakeCmdRunner {
	return &FakeCmdRunner{
		AvailableCommands: map[string]bool{},
		commandResults:    map[string][]FakeCmdResult{},
		processes:         map[string][]*FakeProcess{},
	}
}

func (r *FakeCmdRunner) RunComplexCommand(cmd boshsys.Command) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	r.RunComplexCommands = append(r.RunComplexCommands, cmd)

	runCmd := append([]string{cmd.Name}, cmd.Args...)
	stdout, stderr, exitstatus, err := r.getOutputsForCmd(runCmd)

	if cmd.Stdout != nil {
		cmd.Stdout.Write([]byte(stdout))
	}

	if cmd.Stderr != nil {
		cmd.Stderr.Write([]byte(stderr))
	}

	return stdout, stderr, exitstatus, err
}

func (r *FakeCmdRunner) RunComplexCommandAsync(cmd boshsys.Command) (boshsys.Process, error) {
	r.processesLock.Lock()
	defer r.processesLock.Unlock()

	r.RunComplexCommands = append(r.RunComplexCommands, cmd)

	runCmd := append([]string{cmd.Name}, cmd.Args...)
	fullCmd := strings.Join(runCmd, " ")
	results, found := r.processes[fullCmd]
	if !found {
		panic(fmt.Sprintf("Failed to find process for %s", fullCmd))
	}

	for _, proc := range results {
		if !proc.Waited {
			return proc, nil
		}
	}
	panic(fmt.Sprintf("Failed to find available process for %s", fullCmd))
}

func (r *FakeCmdRunner) RunCommand(cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{cmdName}, args...)
	r.RunCommands = append(r.RunCommands, runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) RunCommandWithInput(input, cmdName string, args ...string) (string, string, int, error) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	runCmd := append([]string{input, cmdName}, args...)
	r.RunCommandsWithInput = append(r.RunCommandsWithInput, runCmd)

	return r.getOutputsForCmd(runCmd)
}

func (r *FakeCmdRunner) CommandExists(cmdName string) bool {
	return r.CommandExistsValue || r.AvailableCommands[cmdName]
}

func (r *FakeCmdRunner) AddCmdResult(fullCmd string, result FakeCmdResult) {
	r.commandResultsLock.Lock()
	defer r.commandResultsLock.Unlock()

	results := r.commandResults[fullCmd]
	r.commandResults[fullCmd] = append(results, result)
}

func (r *FakeCmdRunner) AddProcess(fullCmd string, process *FakeProcess) {
	r.processesLock.Lock()
	defer r.processesLock.Unlock()

	processes := r.processes[fullCmd]
	r.processes[fullCmd] = append(processes, process)
}

func (r *FakeCmdRunner) getOutputsForCmd(runCmd []string) (string, string, int, error) {
	fullCmd := strings.Join(runCmd, " ")
	results, found := r.commandResults[fullCmd]
	if found {
		result := results[0]
		newResults := []FakeCmdResult{}
		if len(results) > 1 {
			newResults = results[1:]
		}

		if !result.Sticky {
			r.commandResults[fullCmd] = newResults
		}

		return result.Stdout, result.Stderr, result.ExitStatus, result.Error
	}

	return "", "", -1, nil
}
<|endoftext|>


<commit_before>/*
Package system runs system-level testing of rqlite. This includes testing of single nodes, and multi-node clusters.
*/
package system

import (
	"testing"
)

func Test_SingleNode(t *testing.T) {
	node := mustNewLeaderNode()
	defer node.Deprovision()

	tests := []struct {
		stmt     string
		expected string
		execute  bool
	}{
		{
			stmt:     `CREATE TABLE foo (id integer not null primary key, name text)`,
			expected: `{"results":[{}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT INTO foo(name) VALUES("fiona")`,
			expected: `{"results":[{"last_insert_id":1,"rows_affected":1}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT INTO bar(name) VALUES("fiona")`,
			expected: `{"results":[{"error":"no such table: bar"}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT blah blah`,
			expected: `{"results":[{"error":"near \"blah\": syntax error"}]}`,
			execute:  true,
		},
		{
			stmt:     `SELECT * FROM foo`,
			expected: `{"results":[{"columns":["id","name"],"types":["integer","text"],"values":[[1,"fiona"]]}]}`,
			execute:  false,
		},
		{
			stmt:     `DROP TABLE bar`,
			expected: `{"results":[{"error":"no such table: bar"}]}`,
			execute:  true,
		},
		{
			stmt:     `DROP TABLE foo`,
			expected: `{"results":[{"last_insert_id":1,"rows_affected":1}]}`,
			execute:  true,
		},
	}

	for i, tt := range tests {
		var r string
		var err error
		if tt.execute {
			r, err = node.Execute(tt.stmt)
		} else {
			r, err = node.Query(tt.stmt)
		}
		if err != nil {
			t.Fatalf(`test %d failed "%s": %s`, i, tt.stmt, err.Error())
		}
		if r != tt.expected {
			t.Fatalf(`test %d received wrong result "%s" got: %s exp: %s`, i, tt.stmt, r, tt.expected)
		}
	}
}

func Test_SingleNodeCoverage(t *testing.T) {
	node := mustNewLeaderNode()
	defer node.Deprovision()

	// Access endpoints to ensure the code is covered.
	var err error
	var str string

	str, err = node.Status()
	if err != nil {
		t.Fatalf("failed to access status endpoint: %s", err.Error())
	}
	if !isJSON(str) {
		t.Fatalf("output from status endpoint is not valid JSON: %s", str)
	}

	str, err = node.Expvar()
	if err != nil {
		t.Fatalf("failed to access expvar endpoint: %s", err.Error())
	}
	if !isJSON(str) {
		t.Fatalf("output from expvar endpoint is not valid JSON: %s", str)
	}
}
<commit_msg>Add multi-table SELECT unit test<commit_after>/*
Package system runs system-level testing of rqlite. This includes testing of single nodes, and multi-node clusters.
*/
package system

import (
	"testing"
)

func Test_SingleNode(t *testing.T) {
	node := mustNewLeaderNode()
	defer node.Deprovision()

	tests := []struct {
		stmt     string
		expected string
		execute  bool
	}{
		{
			stmt:     `CREATE TABLE foo (id integer not null primary key, name text)`,
			expected: `{"results":[{}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT INTO foo(name) VALUES("fiona")`,
			expected: `{"results":[{"last_insert_id":1,"rows_affected":1}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT INTO bar(name) VALUES("fiona")`,
			expected: `{"results":[{"error":"no such table: bar"}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT blah blah`,
			expected: `{"results":[{"error":"near \"blah\": syntax error"}]}`,
			execute:  true,
		},
		{
			stmt:     `SELECT * FROM foo`,
			expected: `{"results":[{"columns":["id","name"],"types":["integer","text"],"values":[[1,"fiona"]]}]}`,
			execute:  false,
		},
		{
			stmt:     `DROP TABLE bar`,
			expected: `{"results":[{"error":"no such table: bar"}]}`,
			execute:  true,
		},
		{
			stmt:     `DROP TABLE foo`,
			expected: `{"results":[{"last_insert_id":1,"rows_affected":1}]}`,
			execute:  true,
		},
	}

	for i, tt := range tests {
		var r string
		var err error
		if tt.execute {
			r, err = node.Execute(tt.stmt)
		} else {
			r, err = node.Query(tt.stmt)
		}
		if err != nil {
			t.Fatalf(`test %d failed "%s": %s`, i, tt.stmt, err.Error())
		}
		if r != tt.expected {
			t.Fatalf(`test %d received wrong result "%s" got: %s exp: %s`, i, tt.stmt, r, tt.expected)
		}
	}
}

func Test_SingleNodeMulti(t *testing.T) {
	node := mustNewLeaderNode()
	defer node.Deprovision()

	tests := []struct {
		stmt     string
		expected string
		execute  bool
	}{
		{
			stmt:     `CREATE TABLE foo (id integer not null primary key, name text)`,
			expected: `{"results":[{}]}`,
			execute:  true,
		},
		{
			stmt:     `CREATE TABLE bar (id integer not null primary key, sequence integer)`,
			expected: `{"results":[{}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT INTO foo(name) VALUES("fiona")`,
			expected: `{"results":[{"last_insert_id":1,"rows_affected":1}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT INTO foo(name) VALUES("declan")`,
			expected: `{"results":[{"last_insert_id":2,"rows_affected":1}]}`,
			execute:  true,
		},
		{
			stmt:     `INSERT INTO bar(sequence) VALUES(5)`,
			expected: `{"results":[{"last_insert_id":1,"rows_affected":1}]}`,
			execute:  true,
		},
	}

	for i, tt := range tests {
		var r string
		var err error
		if tt.execute {
			r, err = node.Execute(tt.stmt)
		} else {
			r, err = node.Query(tt.stmt)
		}
		if err != nil {
			t.Fatalf(`test %d failed "%s": %s`, i, tt.stmt, err.Error())
		}
		if r != tt.expected {
			t.Fatalf(`test %d received wrong result "%s" got: %s exp: %s`, i, tt.stmt, r, tt.expected)
		}
	}

	r, err := node.QueryMulti([]string{"SELECT * FROM foo", "SELECT * FROM bar"})
	if err != nil {
		t.Fatalf("failed to run multiple queries: %s", err.Error())
	}
	if r != `{"results":[{"columns":["id","name"],"types":["integer","text"],"values":[[1,"fiona"],[2,"declan"]]},{"columns":["id","sequence"],"types":["integer","integer"],"values":[[1,5]]}]}` {
		t.Fatalf("test received wrong result got %s", r)
	}
}

func Test_SingleNodeCoverage(t *testing.T) {
	node := mustNewLeaderNode()
	defer node.Deprovision()

	// Access endpoints to ensure the code is covered.
	var err error
	var str string

	str, err = node.Status()
	if err != nil {
		t.Fatalf("failed to access status endpoint: %s", err.Error())
	}
	if !isJSON(str) {
		t.Fatalf("output from status endpoint is not valid JSON: %s", str)
	}

	str, err = node.Expvar()
	if err != nil {
		t.Fatalf("failed to access expvar endpoint: %s", err.Error())
	}
	if !isJSON(str) {
		t.Fatalf("output from expvar endpoint is not valid JSON: %s", str)
	}
}
<|endoftext|>


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

// +build gccgo

package main

func cansse2() bool { return false }
<commit_msg>cmd/dist: detect sse2 even with gccgo<commit_after>// Copyright 2015 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build gccgo

package main

/*
int supports_sse2() {
#if defined(__i386__) || defined(__x86_64__)
	return __builtin_cpu_supports("sse2");
#else
	return 0;
#endif
}
*/
import "C"

func cansse2() bool { return C.supports_sse2() != 0 }
<|endoftext|>


<commit_before>/*
 * Copyright 2014 Google Inc. All rights reserved.
 *
 * 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 cli is a command line interface to shipshape.
// It (optionally) pulls a docker container, runs it,
// and runs the analysis service with the specified local
// files and configuration.
package cli

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"shipshape/service"
	"shipshape/util/docker"
	"shipshape/util/rpc/client"
	glog "third_party/go-glog"

	"github.com/golang/protobuf/proto"

	ctxpb "shipshape/proto/shipshape_context_proto"
	rpcpb "shipshape/proto/shipshape_rpc_proto"
)

const (
	workspace  = "/shipshape-workspace"
	logsDir    = "/shipshape-output"
	localLogs  = "/tmp"
	image      = "service"
	kytheImage = "kythe"
)

type Options struct {
	File                string
	ThirdPartyAnalyzers []string
	// TODO(ciera): make an enum
	Build          string
	TriggerCats    []string
	Dind           bool
	Event          string
	Repo           string
	StayUp         bool
	Tag            string
	LocalKythe     bool
	HandleResponse func(msg *rpcpb.ShipshapeResponse, directory string) error
}

type Invocation struct {
	options Options
}

func New(options Options) *Invocation {
	return &Invocation{options}
}

func (i *Invocation) Run() (int, error) {
	glog.Infof("Starting shipshape...")
	fs, err := os.Stat(i.options.File)
	if err != nil {
		return 0, fmt.Errorf("%s is not a valid file or directory\n", i.options.File)
	}

	origDir := i.options.File
	if !fs.IsDir() {
		origDir = filepath.Dir(i.options.File)
	}

	absRoot, err := filepath.Abs(origDir)
	if err != nil {
		return 0, fmt.Errorf("could not get absolute path for %s: %v\n", origDir, err)
	}

	if !docker.HasDocker() {
		return 0, fmt.Errorf("docker could not be found. Make sure you have docker installed.")
	}

	image := docker.FullImageName(i.options.Repo, image, i.options.Tag)
	glog.Infof("Starting shipshape using %s on %s", image, absRoot)

	// Create the request

	if len(i.options.TriggerCats) == 0 {
		glog.Infof("No categories provided. Will be using categories specified by the config file for the event %s", i.options.Event)
	}

	if len(i.options.ThirdPartyAnalyzers) == 0 {
		i.options.ThirdPartyAnalyzers, err = service.GlobalConfig(absRoot)
		if err != nil {
			glog.Infof("Could not get global config; using only the default analyzers: %v", err)
		}
	}

	// If we are not running in local mode, pull the latest copy
	// Notice this will use the local tag as a signal to not pull the
	// third-party analyzers either.
	if i.options.Tag != "local" {
		pull(image)
		pullAnalyzers(i.options.ThirdPartyAnalyzers)
	}

	// Put in this defer before calling run. Even if run fails, it can
	// still create the container.
	if !i.options.StayUp {
		// TODO(ciera): Rather than immediately sending a SIGKILL,
		// we should use the default 10 seconds and properly handle
		// SIGTERMs in the endpoint script.
		defer stop("shipping_container", 0)
		// Stop all the analyzers, even the ones that had trouble starting,
		// in case they did actually start
		for id, analyzerRepo := range i.options.ThirdPartyAnalyzers {
			container, _ := getContainerAndAddress(analyzerRepo, id)
			defer stop(container, 0)
		}
	}

	containers, errs := startAnalyzers(absRoot, i.options.ThirdPartyAnalyzers, i.options.Dind)
	for _, err := range errs {
		glog.Errorf("Could not start up third party analyzer: %v", err)
	}

	var c *client.Client
	var req *rpcpb.ShipshapeRequest
	var numNotes int

	// Run it on files
	relativeRoot := ""
	c, relativeRoot, err = startShipshapeService(image, absRoot, containers, i.options.Dind)
	if err != nil {
		return 0, fmt.Errorf("HTTP client did not become healthy: %v", err)
	}
	var files []string
	if !fs.IsDir() {
		files = []string{filepath.Base(i.options.File)}
	}
	req = createRequest(i.options.TriggerCats, files, i.options.Event, filepath.Join(workspace, relativeRoot), ctxpb.Stage_PRE_BUILD.Enum())
	glog.Infof("Calling with request %v", req)
	numNotes, err = analyze(c, req, origDir, i.options.HandleResponse)
	if err != nil {
		return numNotes, fmt.Errorf("error making service call: %v", err)
	}

	// If desired, generate compilation units with a kythe image
	if i.options.Build != "" {
		// TODO(ciera): Handle other build systems
		fullKytheImage := docker.FullImageName(i.options.Repo, kytheImage, i.options.Tag)
		if !i.options.LocalKythe {
			pull(fullKytheImage)
		}

		// TODO(emso): Add a check for an already running kythe container.
		// The below defer should stop the one started below but in case this
		// failed for some reason (or a kythe container was started in some other
		// way) the below run command will fail.
		defer stop("kythe", 10*time.Second)
		glog.Infof("Retrieving compilation units with %s", i.options.Build)

		result := docker.RunKythe(fullKytheImage, "kythe", absRoot, i.options.Build, i.options.Dind)
		if result.Err != nil {
			// kythe spews output, so only capture it if something went wrong.
			printStreams(result)
			return numNotes, fmt.Errorf("error from run: %v", result.Err)
		}
		glog.Infoln("CompilationUnits prepared")

		req.Stage = ctxpb.Stage_POST_BUILD.Enum()
		glog.Infof("Calling with request %v", req)
		numBuildNotes, err := analyze(c, req, origDir, i.options.HandleResponse)
		numNotes += numBuildNotes
		if err != nil {
			return numNotes, fmt.Errorf("error making service call: %v", err)
		}
	}

	glog.Infoln("End of Results.")
	return numNotes, nil
}

func numNotes(msg *rpcpb.ShipshapeResponse) int {
	numNotes := 0
	for _, analysis := range msg.AnalyzeResponse {
		numNotes += len(analysis.Note)
	}
	return numNotes
}

// startShipshapeService ensures that there is a service started with the given image and
// attached analyzers that can analyze the directory at absRoot (an absolute path). If a
// service is not started up that can do this, it will shut down the existing one and start
// a new one.
// The methods returns the (ready) client, the relative path from the docker container's mapped
// volume to the absRoot that we are analyzing, and any errors from attempting to run the service.
// TODO(ciera): This *should* check the analyzers that are connected, but does not yet
// do so.
func startShipshapeService(image, absRoot string, analyzers []string, dind bool) (*client.Client, string, error) {
	glog.Infof("Starting shipshape...")
	container := "shipping_container"
	// subPath is the relatve path from the mapped volume on shipping container
	// to the directory we are analyzing (absRoot)
	isMapped, subPath := docker.MappedVolume(absRoot, container)
	// Stop and restart the container if:
	// 1: The container is not using the latest image OR
	// 2: The container is not mapped to the right directory OR
	// 3: The container is not linked to the right analyzer containers
	// Otherwise, use the existing container
	if !docker.ImageMatches(image, container) || !isMapped || !docker.ContainsLinks(container, analyzers) {
		glog.Infof("Restarting container with %s", image)
		stop(container, 0)
		result := docker.RunService(image, container, absRoot, localLogs, analyzers, dind)
		subPath = ""
		printStreams(result)
		if result.Err != nil {
			return nil, "", result.Err
		}
	}
	glog.Infof("Image %s running in service mode", image)
	c := client.NewHTTPClient("localhost:10007")
	return c, subPath, c.WaitUntilReady(10 * time.Second)
}

func analyze(c *client.Client, req *rpcpb.ShipshapeRequest, originalDir string, handleResponse func(msg *rpcpb.ShipshapeResponse, directory string) error) (int, error) {
	var totalNotes = 0
	glog.Infof("Calling to the shipshape service with %v", req)
	rd := c.Stream("/ShipshapeService/Run", req)
	defer rd.Close()
	for {
		var msg rpcpb.ShipshapeResponse
		if err := rd.NextResult(&msg); err == io.EOF {
			break
		} else if err != nil {
			return 0, fmt.Errorf("received an error from calling run: %v", err.Error())
		}

		err := handleResponse(&msg, originalDir)
		if err != nil {
			return 0, fmt.Errorf("could not parse results: %v", err.Error())
		}
		totalNotes += numNotes(&msg)
	}
	return totalNotes, nil
}

func pull(image string) {
	if !docker.OutOfDate(image) {
		return
	}
	glog.Infof("Pulling image %s", image)
	result := docker.Pull(image)
	printStreams(result)
	if result.Err != nil {
		glog.Errorf("Error from pull: %v", result.Err)
		return
	}
	glog.Infoln("Pulling complete")
}

func stop(container string, timeWait time.Duration) {
	glog.Infof("Stopping and removing %s", container)
	result := docker.Stop(container, timeWait, true)
	printStreams(result)
	if result.Err != nil {
		glog.Infof("Could not stop %s: %v", container, result.Err)
	} else {
		glog.Infoln("Removed.")
	}
}

func pullAnalyzers(images []string) {
	var wg sync.WaitGroup
	for _, analyzerImage := range images {
		wg.Add(1)
		go func(image string) {
			pull(image)
			wg.Done()
		}(analyzerImage)
	}
	glog.Info("Pulling dockerized analyzers...")
	wg.Wait()
	glog.Info("Analyzers pulled")
}

func startAnalyzers(sourceDir string, images []string, dind bool) (containers []string, errs []error) {
	var wg sync.WaitGroup
	for id, fullImage := range images {
		wg.Add(1)
		go func(id int, image string) {
			analyzerContainer, port := getContainerAndAddress(image, id)
			if docker.ImageMatches(image, analyzerContainer) {
				glog.Infof("Reusing analyzer %v started at localhost:%d", image, port)
			} else {
				glog.Infof("Found no analyzer container (%v) to reuse for %v", analyzerContainer, image)
				// Analyzer is either running with the wrong image version, or not running
				// Stopping in case it's the first case
				result := docker.Stop(analyzerContainer, 0, true)
				if result.Err != nil {
					glog.Infof("Failed to stop %v (may not be running)", analyzerContainer)
				}
				result = docker.RunAnalyzer(image, analyzerContainer, sourceDir, localLogs, port, dind)
				if result.Err != nil {
					glog.Infof("Could not start %v at localhost:%d: %v, stderr: %v", image, port, result.Err.Error(), result.Stderr)
					errs = append(errs, result.Err)
				} else {
					glog.Infof("Analyzer %v started at localhost:%d", image, port)
					containers = append(containers, analyzerContainer)
				}
			}
			wg.Done()
		}(id, fullImage)
	}
	if len(images) > 0 {
		glog.Info("Waiting for dockerized analyzers to start up...")
		wg.Wait()
		glog.Info("Analyzers up")
	}
	return containers, errs
}

func printStreams(result docker.CommandResult) {
	out := strings.TrimSpace(result.Stdout)
	err := strings.TrimSpace(result.Stderr)
	if len(out) > 0 {
		glog.Infof("stdout:\n%s\n", strings.TrimSpace(result.Stdout))
	}
	if len(err) > 0 {
		glog.Errorf("stderr:\n%s\n", strings.TrimSpace(result.Stderr))
	}
}

func getContainerAndAddress(fullImage string, id int) (analyzerContainer string, port int) {
	// A docker image URI (location:port/path:tag) can have a host part
	// with a port number and a path part with a tag.  Both tag and port
	// are separated by colon, so we need to find out if the last colon is
	// the one that separates the tag from the path, or the port in the
	// location.
	end := strings.LastIndex(fullImage, ":")
	slash := strings.LastIndex(fullImage, "/")
	if end == -1 || end < slash {
		// No colon, or last colon is part of the location.
		end = len(fullImage)
	}
	image := fullImage[slash+1 : end]
	port = 10010 + id
	analyzerContainer = fmt.Sprintf("%s_%d", image, id)
	return analyzerContainer, port
}

func createRequest(triggerCats, files []string, event, repoRoot string, stage *ctxpb.Stage) *rpcpb.ShipshapeRequest {
	return &rpcpb.ShipshapeRequest{
		TriggeredCategory: triggerCats,
		ShipshapeContext: &ctxpb.ShipshapeContext{
			RepoRoot: proto.String(repoRoot),
			FilePath: files,
		},
		Event: proto.String(event),
		Stage: stage,
	}
}
<commit_msg>Added a comment.<commit_after>/*
 * Copyright 2014 Google Inc. All rights reserved.
 *
 * 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 cli is a command line interface to shipshape.
// It (optionally) pulls a docker container, runs it,
// and runs the analysis service with the specified local
// files and configuration.
package cli

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"shipshape/service"
	"shipshape/util/docker"
	"shipshape/util/rpc/client"
	glog "third_party/go-glog"

	"github.com/golang/protobuf/proto"

	ctxpb "shipshape/proto/shipshape_context_proto"
	rpcpb "shipshape/proto/shipshape_rpc_proto"
)

const (
	workspace  = "/shipshape-workspace"
	logsDir    = "/shipshape-output"
	localLogs  = "/tmp"
	image      = "service"
	kytheImage = "kythe"
)

type Options struct {
	File                string
	ThirdPartyAnalyzers []string
	// TODO(ciera): make an enum
	Build          string
	TriggerCats    []string
	Dind           bool
	Event          string
	Repo           string
	StayUp         bool
	Tag            string
	LocalKythe     bool
	// Directory has the path the analyzed file is in (msg.AnalyzeResponse.Note.Location.GetPath()
	// contains only the basename).
	HandleResponse func(msg *rpcpb.ShipshapeResponse, directory string) error
}

type Invocation struct {
	options Options
}

func New(options Options) *Invocation {
	return &Invocation{options}
}

func (i *Invocation) Run() (int, error) {
	glog.Infof("Starting shipshape...")
	fs, err := os.Stat(i.options.File)
	if err != nil {
		return 0, fmt.Errorf("%s is not a valid file or directory\n", i.options.File)
	}

	origDir := i.options.File
	if !fs.IsDir() {
		origDir = filepath.Dir(i.options.File)
	}

	absRoot, err := filepath.Abs(origDir)
	if err != nil {
		return 0, fmt.Errorf("could not get absolute path for %s: %v\n", origDir, err)
	}

	if !docker.HasDocker() {
		return 0, fmt.Errorf("docker could not be found. Make sure you have docker installed.")
	}

	image := docker.FullImageName(i.options.Repo, image, i.options.Tag)
	glog.Infof("Starting shipshape using %s on %s", image, absRoot)

	// Create the request

	if len(i.options.TriggerCats) == 0 {
		glog.Infof("No categories provided. Will be using categories specified by the config file for the event %s", i.options.Event)
	}

	if len(i.options.ThirdPartyAnalyzers) == 0 {
		i.options.ThirdPartyAnalyzers, err = service.GlobalConfig(absRoot)
		if err != nil {
			glog.Infof("Could not get global config; using only the default analyzers: %v", err)
		}
	}

	// If we are not running in local mode, pull the latest copy
	// Notice this will use the local tag as a signal to not pull the
	// third-party analyzers either.
	if i.options.Tag != "local" {
		pull(image)
		pullAnalyzers(i.options.ThirdPartyAnalyzers)
	}

	// Put in this defer before calling run. Even if run fails, it can
	// still create the container.
	if !i.options.StayUp {
		// TODO(ciera): Rather than immediately sending a SIGKILL,
		// we should use the default 10 seconds and properly handle
		// SIGTERMs in the endpoint script.
		defer stop("shipping_container", 0)
		// Stop all the analyzers, even the ones that had trouble starting,
		// in case they did actually start
		for id, analyzerRepo := range i.options.ThirdPartyAnalyzers {
			container, _ := getContainerAndAddress(analyzerRepo, id)
			defer stop(container, 0)
		}
	}

	containers, errs := startAnalyzers(absRoot, i.options.ThirdPartyAnalyzers, i.options.Dind)
	for _, err := range errs {
		glog.Errorf("Could not start up third party analyzer: %v", err)
	}

	var c *client.Client
	var req *rpcpb.ShipshapeRequest
	var numNotes int

	// Run it on files
	relativeRoot := ""
	c, relativeRoot, err = startShipshapeService(image, absRoot, containers, i.options.Dind)
	if err != nil {
		return 0, fmt.Errorf("HTTP client did not become healthy: %v", err)
	}
	var files []string
	if !fs.IsDir() {
		files = []string{filepath.Base(i.options.File)}
	}
	req = createRequest(i.options.TriggerCats, files, i.options.Event, filepath.Join(workspace, relativeRoot), ctxpb.Stage_PRE_BUILD.Enum())
	glog.Infof("Calling with request %v", req)
	numNotes, err = analyze(c, req, origDir, i.options.HandleResponse)
	if err != nil {
		return numNotes, fmt.Errorf("error making service call: %v", err)
	}

	// If desired, generate compilation units with a kythe image
	if i.options.Build != "" {
		// TODO(ciera): Handle other build systems
		fullKytheImage := docker.FullImageName(i.options.Repo, kytheImage, i.options.Tag)
		if !i.options.LocalKythe {
			pull(fullKytheImage)
		}

		// TODO(emso): Add a check for an already running kythe container.
		// The below defer should stop the one started below but in case this
		// failed for some reason (or a kythe container was started in some other
		// way) the below run command will fail.
		defer stop("kythe", 10*time.Second)
		glog.Infof("Retrieving compilation units with %s", i.options.Build)

		result := docker.RunKythe(fullKytheImage, "kythe", absRoot, i.options.Build, i.options.Dind)
		if result.Err != nil {
			// kythe spews output, so only capture it if something went wrong.
			printStreams(result)
			return numNotes, fmt.Errorf("error from run: %v", result.Err)
		}
		glog.Infoln("CompilationUnits prepared")

		req.Stage = ctxpb.Stage_POST_BUILD.Enum()
		glog.Infof("Calling with request %v", req)
		numBuildNotes, err := analyze(c, req, origDir, i.options.HandleResponse)
		numNotes += numBuildNotes
		if err != nil {
			return numNotes, fmt.Errorf("error making service call: %v", err)
		}
	}

	glog.Infoln("End of Results.")
	return numNotes, nil
}

func numNotes(msg *rpcpb.ShipshapeResponse) int {
	numNotes := 0
	for _, analysis := range msg.AnalyzeResponse {
		numNotes += len(analysis.Note)
	}
	return numNotes
}

// startShipshapeService ensures that there is a service started with the given image and
// attached analyzers that can analyze the directory at absRoot (an absolute path). If a
// service is not started up that can do this, it will shut down the existing one and start
// a new one.
// The methods returns the (ready) client, the relative path from the docker container's mapped
// volume to the absRoot that we are analyzing, and any errors from attempting to run the service.
// TODO(ciera): This *should* check the analyzers that are connected, but does not yet
// do so.
func startShipshapeService(image, absRoot string, analyzers []string, dind bool) (*client.Client, string, error) {
	glog.Infof("Starting shipshape...")
	container := "shipping_container"
	// subPath is the relatve path from the mapped volume on shipping container
	// to the directory we are analyzing (absRoot)
	isMapped, subPath := docker.MappedVolume(absRoot, container)
	// Stop and restart the container if:
	// 1: The container is not using the latest image OR
	// 2: The container is not mapped to the right directory OR
	// 3: The container is not linked to the right analyzer containers
	// Otherwise, use the existing container
	if !docker.ImageMatches(image, container) || !isMapped || !docker.ContainsLinks(container, analyzers) {
		glog.Infof("Restarting container with %s", image)
		stop(container, 0)
		result := docker.RunService(image, container, absRoot, localLogs, analyzers, dind)
		subPath = ""
		printStreams(result)
		if result.Err != nil {
			return nil, "", result.Err
		}
	}
	glog.Infof("Image %s running in service mode", image)
	c := client.NewHTTPClient("localhost:10007")
	return c, subPath, c.WaitUntilReady(10 * time.Second)
}

func analyze(c *client.Client, req *rpcpb.ShipshapeRequest, originalDir string, handleResponse func(msg *rpcpb.ShipshapeResponse, directory string) error) (int, error) {
	var totalNotes = 0
	glog.Infof("Calling to the shipshape service with %v", req)
	rd := c.Stream("/ShipshapeService/Run", req)
	defer rd.Close()
	for {
		var msg rpcpb.ShipshapeResponse
		if err := rd.NextResult(&msg); err == io.EOF {
			break
		} else if err != nil {
			return 0, fmt.Errorf("received an error from calling run: %v", err.Error())
		}

		err := handleResponse(&msg, originalDir)
		if err != nil {
			return 0, fmt.Errorf("could not parse results: %v", err.Error())
		}
		totalNotes += numNotes(&msg)
	}
	return totalNotes, nil
}

func pull(image string) {
	if !docker.OutOfDate(image) {
		return
	}
	glog.Infof("Pulling image %s", image)
	result := docker.Pull(image)
	printStreams(result)
	if result.Err != nil {
		glog.Errorf("Error from pull: %v", result.Err)
		return
	}
	glog.Infoln("Pulling complete")
}

func stop(container string, timeWait time.Duration) {
	glog.Infof("Stopping and removing %s", container)
	result := docker.Stop(container, timeWait, true)
	printStreams(result)
	if result.Err != nil {
		glog.Infof("Could not stop %s: %v", container, result.Err)
	} else {
		glog.Infoln("Removed.")
	}
}

func pullAnalyzers(images []string) {
	var wg sync.WaitGroup
	for _, analyzerImage := range images {
		wg.Add(1)
		go func(image string) {
			pull(image)
			wg.Done()
		}(analyzerImage)
	}
	glog.Info("Pulling dockerized analyzers...")
	wg.Wait()
	glog.Info("Analyzers pulled")
}

func startAnalyzers(sourceDir string, images []string, dind bool) (containers []string, errs []error) {
	var wg sync.WaitGroup
	for id, fullImage := range images {
		wg.Add(1)
		go func(id int, image string) {
			analyzerContainer, port := getContainerAndAddress(image, id)
			if docker.ImageMatches(image, analyzerContainer) {
				glog.Infof("Reusing analyzer %v started at localhost:%d", image, port)
			} else {
				glog.Infof("Found no analyzer container (%v) to reuse for %v", analyzerContainer, image)
				// Analyzer is either running with the wrong image version, or not running
				// Stopping in case it's the first case
				result := docker.Stop(analyzerContainer, 0, true)
				if result.Err != nil {
					glog.Infof("Failed to stop %v (may not be running)", analyzerContainer)
				}
				result = docker.RunAnalyzer(image, analyzerContainer, sourceDir, localLogs, port, dind)
				if result.Err != nil {
					glog.Infof("Could not start %v at localhost:%d: %v, stderr: %v", image, port, result.Err.Error(), result.Stderr)
					errs = append(errs, result.Err)
				} else {
					glog.Infof("Analyzer %v started at localhost:%d", image, port)
					containers = append(containers, analyzerContainer)
				}
			}
			wg.Done()
		}(id, fullImage)
	}
	if len(images) > 0 {
		glog.Info("Waiting for dockerized analyzers to start up...")
		wg.Wait()
		glog.Info("Analyzers up")
	}
	return containers, errs
}

func printStreams(result docker.CommandResult) {
	out := strings.TrimSpace(result.Stdout)
	err := strings.TrimSpace(result.Stderr)
	if len(out) > 0 {
		glog.Infof("stdout:\n%s\n", strings.TrimSpace(result.Stdout))
	}
	if len(err) > 0 {
		glog.Errorf("stderr:\n%s\n", strings.TrimSpace(result.Stderr))
	}
}

func getContainerAndAddress(fullImage string, id int) (analyzerContainer string, port int) {
	// A docker image URI (location:port/path:tag) can have a host part
	// with a port number and a path part with a tag.  Both tag and port
	// are separated by colon, so we need to find out if the last colon is
	// the one that separates the tag from the path, or the port in the
	// location.
	end := strings.LastIndex(fullImage, ":")
	slash := strings.LastIndex(fullImage, "/")
	if end == -1 || end < slash {
		// No colon, or last colon is part of the location.
		end = len(fullImage)
	}
	image := fullImage[slash+1 : end]
	port = 10010 + id
	analyzerContainer = fmt.Sprintf("%s_%d", image, id)
	return analyzerContainer, port
}

func createRequest(triggerCats, files []string, event, repoRoot string, stage *ctxpb.Stage) *rpcpb.ShipshapeRequest {
	return &rpcpb.ShipshapeRequest{
		TriggeredCategory: triggerCats,
		ShipshapeContext: &ctxpb.ShipshapeContext{
			RepoRoot: proto.String(repoRoot),
			FilePath: files,
		},
		Event: proto.String(event),
		Stage: stage,
	}
}
<|endoftext|>


<commit_before>package database_test

import (
	. "github.com/pivotal-cf-experimental/cf-mysql-quota-enforcer/database"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"errors"

	"database/sql"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/pivotal-golang/lager/lagertest"
)

var _ = Describe("Database", func() {

	const (
		dbName = "fake-db-name"
		dbUser = "fake-db-user"
	)

	var (
		logger                 *lagertest.TestLogger
		database               Database
		fakeDB                 *sql.DB
		flushPrivilegesPattern = "FLUSH PRIVILEGES"
	)

	BeforeEach(func() {
		var err error
		fakeDB, err = sqlmock.New()
		Expect(err).ToNot(HaveOccurred())

		logger = lagertest.NewTestLogger("Database test")
		database = New(dbName, dbUser, fakeDB, logger)
	})

	AfterEach(func() {
		err := fakeDB.Close()
		Expect(err).ToNot(HaveOccurred())
	})

	Describe("RevokePrivileges", func() {
		var (
			revokePrivilegesPattern = `REVOKE INSERT, UPDATE, CREATE ON \?.\* FROM '\?'@'\%'`
		)

		It("makes a sql query to revoke privileges on a database and then flushes privileges", func() {
			sqlmock.ExpectExec(revokePrivilegesPattern).
				WithArgs(dbName, dbUser).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			sqlmock.ExpectExec(flushPrivilegesPattern).
				WithArgs().
				WillReturnResult(sqlmock.NewResult(-1, 1))

			err := database.RevokePrivileges()
			Expect(err).ToNot(HaveOccurred())
		})

		Context("when the query fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(revokePrivilegesPattern).
					WithArgs(dbName, dbUser).
					WillReturnError(errors.New("fake-query-error"))
			})

			It("returns an error", func() {
				err := database.RevokePrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-query-error"))
				Expect(err.Error()).To(ContainSubstring(dbName))
				Expect(err.Error()).To(ContainSubstring(dbUser))
			})
		})

		Context("when getting the number of affected rows fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(revokePrivilegesPattern).
					WithArgs(dbName, dbUser).
					WillReturnResult(sqlmock.NewErrorResult(errors.New("fake-rows-affected-error")))
			})

			It("returns an error", func() {
				err := database.RevokePrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-rows-affected-error"))
				Expect(err.Error()).To(ContainSubstring("Getting rows affected"))
				Expect(err.Error()).To(ContainSubstring(dbName))
				Expect(err.Error()).To(ContainSubstring(dbUser))
			})
		})

		Context("when flushing privileges fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(revokePrivilegesPattern).
					WithArgs(dbName, dbUser).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				sqlmock.ExpectExec(flushPrivilegesPattern).
					WithArgs().
					WillReturnError(errors.New("fake-flush-error"))
			})

			It("returns an error", func() {
				err := database.RevokePrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-flush-error"))
			})
		})
	})

	Describe("GrantPrivileges", func() {
		var (
			grantPrivilegesPattern = `GRANT INSERT, UPDATE, CREATE ON \?.\* TO '\?'@'\%'`
		)

		It("grants privileges to the database", func() {
			sqlmock.ExpectExec(grantPrivilegesPattern).
				WithArgs(dbName, dbUser).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			sqlmock.ExpectExec(flushPrivilegesPattern).
				WithArgs().
				WillReturnResult(sqlmock.NewResult(-1, 1))

			err := database.GrantPrivileges()
			Expect(err).ToNot(HaveOccurred())
		})

		Context("when the query fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(grantPrivilegesPattern).
					WithArgs(dbName, dbUser).
					WillReturnError(errors.New("fake-query-error"))
			})

			It("returns an error", func() {
				err := database.GrantPrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-query-error"))
				Expect(err.Error()).To(ContainSubstring(dbName))
				Expect(err.Error()).To(ContainSubstring(dbUser))
			})
		})

		Context("when getting the number of affected rows fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(grantPrivilegesPattern).
					WithArgs(dbName, dbUser).
					WillReturnResult(sqlmock.NewErrorResult(errors.New("fake-rows-affected-error")))
			})

			It("returns an error", func() {
				err := database.GrantPrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-rows-affected-error"))
				Expect(err.Error()).To(ContainSubstring("Getting rows affected"))
				Expect(err.Error()).To(ContainSubstring(dbName))
				Expect(err.Error()).To(ContainSubstring(dbUser))
			})
		})

		Context("when flushing privileges fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(grantPrivilegesPattern).
					WithArgs(dbName, dbUser).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				sqlmock.ExpectExec(flushPrivilegesPattern).
					WithArgs().
					WillReturnError(errors.New("fake-flush-error"))
			})

			It("returns an error", func() {
				err := database.GrantPrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-flush-error"))
			})
		})

	})

	Describe("KillActiveConnections", func() {
		var (
			processListColumns    = []string{"ID"}
			processQueryPattern   = `SELECT ID FROM INFORMATION_SCHEMA.PROCESSLIST WHERE DB = \? AND USER <> 'root'`
			killConnectionPattern = "KILL CONNECTION \\?"
		)

		It("kills all active connections to DB", func() {
			sqlmock.ExpectQuery(processQueryPattern).
				WithArgs(dbName).
				WillReturnRows(sqlmock.NewRows(processListColumns).AddRow(1).AddRow(123))

			sqlmock.ExpectExec(killConnectionPattern).
				WithArgs(1).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			sqlmock.ExpectExec(killConnectionPattern).
				WithArgs(123).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			err := database.KillActiveConnections()
			Expect(err).ToNot(HaveOccurred())
		})

		Context("when there are no active connections to the database", func() {
			It("does not kill any connections", func() {
				sqlmock.ExpectQuery(processQueryPattern).
					WithArgs(dbName).
					WillReturnRows(sqlmock.NewRows(processListColumns))

				err := database.KillActiveConnections()
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("when there is only one active connections to the database", func() {
			It("kills the active connection", func() {
				sqlmock.ExpectQuery(processQueryPattern).
					WithArgs(dbName).
					WillReturnRows(sqlmock.NewRows(processListColumns).AddRow(123))

				sqlmock.ExpectExec(killConnectionPattern).
					WithArgs(123).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				err := database.KillActiveConnections()
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("when querying for active connections fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectQuery(processQueryPattern).
					WithArgs(dbName).
					WillReturnError(errors.New("fake-query-error"))
			})

			It("returns an error", func() {
				err := database.KillActiveConnections()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-query-error"))
				Expect(err.Error()).To(ContainSubstring(dbName))
			})
		})

		Context("when killing a connection fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectQuery(processQueryPattern).
					WithArgs(dbName).
					WillReturnRows(sqlmock.NewRows(processListColumns).AddRow(1).AddRow(2).AddRow(3))

				sqlmock.ExpectExec(killConnectionPattern).
					WithArgs(2).
					WillReturnError(errors.New("fake-exec-error"))
			})

			It("kills all other active connections", func() {
				sqlmock.ExpectExec(killConnectionPattern).
					WithArgs(1).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				sqlmock.ExpectExec(killConnectionPattern).
					WithArgs(3).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				err := database.KillActiveConnections()
				Expect(err).ToNot(HaveOccurred())
			})
		})
	})
})
<commit_msg>Grant/revoke queries are not using exec with args anymore<commit_after>package database_test

import (
	. "github.com/pivotal-cf-experimental/cf-mysql-quota-enforcer/database"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"errors"

	"database/sql"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/pivotal-golang/lager/lagertest"
)

var _ = Describe("Database", func() {

	const (
		dbName = "fake-db-name"
		dbUser = "fake-db-user"
	)

	var (
		logger                 *lagertest.TestLogger
		database               Database
		fakeDB                 *sql.DB
		flushPrivilegesPattern = "FLUSH PRIVILEGES"
	)

	BeforeEach(func() {
		var err error
		fakeDB, err = sqlmock.New()
		Expect(err).ToNot(HaveOccurred())

		logger = lagertest.NewTestLogger("Database test")
		database = New(dbName, dbUser, fakeDB, logger)
	})

	AfterEach(func() {
		err := fakeDB.Close()
		Expect(err).ToNot(HaveOccurred())
	})

	Describe("RevokePrivileges", func() {
		var (
			revokePrivilegesPattern = `REVOKE INSERT, UPDATE, CREATE ON fake-db-name.\* FROM 'fake-db-user'@'\%'`
		)

		It("makes a sql query to revoke privileges on a database and then flushes privileges", func() {
			sqlmock.ExpectExec(revokePrivilegesPattern).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			sqlmock.ExpectExec(flushPrivilegesPattern).
				WithArgs().
				WillReturnResult(sqlmock.NewResult(-1, 1))

			err := database.RevokePrivileges()
			Expect(err).ToNot(HaveOccurred())
		})

		Context("when the query fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(revokePrivilegesPattern).
					WillReturnError(errors.New("fake-query-error"))
			})

			It("returns an error", func() {
				err := database.RevokePrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-query-error"))
				Expect(err.Error()).To(ContainSubstring(dbName))
				Expect(err.Error()).To(ContainSubstring(dbUser))
			})
		})

		Context("when getting the number of affected rows fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(revokePrivilegesPattern).
					WillReturnResult(sqlmock.NewErrorResult(errors.New("fake-rows-affected-error")))
			})

			It("returns an error", func() {
				err := database.RevokePrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-rows-affected-error"))
				Expect(err.Error()).To(ContainSubstring("Getting rows affected"))
				Expect(err.Error()).To(ContainSubstring(dbName))
				Expect(err.Error()).To(ContainSubstring(dbUser))
			})
		})

		Context("when flushing privileges fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(revokePrivilegesPattern).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				sqlmock.ExpectExec(flushPrivilegesPattern).
					WithArgs().
					WillReturnError(errors.New("fake-flush-error"))
			})

			It("returns an error", func() {
				err := database.RevokePrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-flush-error"))
			})
		})
	})

	Describe("GrantPrivileges", func() {
		var (
			grantPrivilegesPattern = `GRANT INSERT, UPDATE, CREATE ON fake-db-name.\* TO 'fake-db-user'@'\%'`
		)

		It("grants privileges to the database", func() {
			sqlmock.ExpectExec(grantPrivilegesPattern).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			sqlmock.ExpectExec(flushPrivilegesPattern).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			err := database.GrantPrivileges()
			Expect(err).ToNot(HaveOccurred())
		})

		Context("when the query fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(grantPrivilegesPattern).
					WillReturnError(errors.New("fake-query-error"))
			})

			It("returns an error", func() {
				err := database.GrantPrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-query-error"))
				Expect(err.Error()).To(ContainSubstring(dbName))
				Expect(err.Error()).To(ContainSubstring(dbUser))
			})
		})

		Context("when getting the number of affected rows fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(grantPrivilegesPattern).
					WillReturnResult(sqlmock.NewErrorResult(errors.New("fake-rows-affected-error")))
			})

			It("returns an error", func() {
				err := database.GrantPrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-rows-affected-error"))
				Expect(err.Error()).To(ContainSubstring("Getting rows affected"))
				Expect(err.Error()).To(ContainSubstring(dbName))
				Expect(err.Error()).To(ContainSubstring(dbUser))
			})
		})

		Context("when flushing privileges fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectExec(grantPrivilegesPattern).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				sqlmock.ExpectExec(flushPrivilegesPattern).
					WithArgs().
					WillReturnError(errors.New("fake-flush-error"))
			})

			It("returns an error", func() {
				err := database.GrantPrivileges()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-flush-error"))
			})
		})

	})

	Describe("KillActiveConnections", func() {
		var (
			processListColumns    = []string{"ID"}
			processQueryPattern   = `SELECT ID FROM INFORMATION_SCHEMA.PROCESSLIST WHERE DB = \? AND USER <> 'root'`
			killConnectionPattern = "KILL CONNECTION \\?"
		)

		It("kills all active connections to DB", func() {
			sqlmock.ExpectQuery(processQueryPattern).
				WithArgs(dbName).
				WillReturnRows(sqlmock.NewRows(processListColumns).AddRow(1).AddRow(123))

			sqlmock.ExpectExec(killConnectionPattern).
				WithArgs(1).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			sqlmock.ExpectExec(killConnectionPattern).
				WithArgs(123).
				WillReturnResult(sqlmock.NewResult(-1, 1))

			err := database.KillActiveConnections()
			Expect(err).ToNot(HaveOccurred())
		})

		Context("when there are no active connections to the database", func() {
			It("does not kill any connections", func() {
				sqlmock.ExpectQuery(processQueryPattern).
					WithArgs(dbName).
					WillReturnRows(sqlmock.NewRows(processListColumns))

				err := database.KillActiveConnections()
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("when there is only one active connections to the database", func() {
			It("kills the active connection", func() {
				sqlmock.ExpectQuery(processQueryPattern).
					WithArgs(dbName).
					WillReturnRows(sqlmock.NewRows(processListColumns).AddRow(123))

				sqlmock.ExpectExec(killConnectionPattern).
					WithArgs(123).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				err := database.KillActiveConnections()
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("when querying for active connections fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectQuery(processQueryPattern).
					WithArgs(dbName).
					WillReturnError(errors.New("fake-query-error"))
			})

			It("returns an error", func() {
				err := database.KillActiveConnections()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("fake-query-error"))
				Expect(err.Error()).To(ContainSubstring(dbName))
			})
		})

		Context("when killing a connection fails", func() {
			BeforeEach(func() {
				sqlmock.ExpectQuery(processQueryPattern).
					WithArgs(dbName).
					WillReturnRows(sqlmock.NewRows(processListColumns).AddRow(1).AddRow(2).AddRow(3))

				sqlmock.ExpectExec(killConnectionPattern).
					WithArgs(2).
					WillReturnError(errors.New("fake-exec-error"))
			})

			It("kills all other active connections", func() {
				sqlmock.ExpectExec(killConnectionPattern).
					WithArgs(1).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				sqlmock.ExpectExec(killConnectionPattern).
					WithArgs(3).
					WillReturnResult(sqlmock.NewResult(-1, 1))

				err := database.KillActiveConnections()
				Expect(err).ToNot(HaveOccurred())
			})
		})
	})
})
<|endoftext|>


<commit_before>package database

import (
	"errors"
	"testing"
	"time"

	"context"

	"github.com/Bit-Nation/BITNATION-Panthalassa/message"
)

// TODO: avoid opening and closing the DB each times
// TODO: avoid filling the DB each times

var messages = []message.Message{
	message.Message{From: "<sample pubkey>", Previous: "", Seq: 1, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<content to retrieve>"}, Hash: "<sample hash 1>", Signature: "<sample sign>"},
	message.Message{From: "<sample pubkey>", Previous: "<sample hash 1>", Seq: 2, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<content to retrieve>"}, Hash: "<sample hash 2>", Signature: "<sample sign>"},
	message.Message{From: "<sample pubkey>", Previous: "<sample hash 2>", Seq: 3, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<sample content>"}, Hash: "<sample hash 3>", Signature: "<sample sign>"},
	message.Message{From: "<sample pubkey>", Previous: "<sample hash 3>", Seq: 4, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<sample content>"}, Hash: "<sample hash 4>", Signature: "<sample sign>"},

	message.Message{From: "<sample pubkey remove>", Previous: "", Seq: 1, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<sample content>"}, Hash: "<sample hash>", Signature: "<sample sign>"},
}

func TestOpenAndClose(t *testing.T) {
	db := DB{File: "/tmp/test.db"}

	err := db.Open()
	if err != nil {
		t.Error(err)
	}

	err = db.Close()
	if err != nil {
		t.Error(err)
	}
}
// TODO: check adding valid and invalid msg
func TestAdd(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	for _, msg := range messages {
		err := db.AddMessage(msg)
		if err != nil {
			t.Error(err)
		}
	}
}

// This test is stupid but who knows...
func TestGetLastMessage(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	msg, err := db.GetLastMessage(messages[3].From)
	if err != nil {
		t.Error(err)
	}

	if msg != messages[3] {
		t.Error(errors.New("previous message doesn't match expected one"))
	}
}

func TestRemove_ExistingMessage(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	err := db.Remove("<sample hash>")
	if err != nil {
		t.Error(err)
	}
}

// TODO: better tests
func TestGetFeed(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	doFeedTest(db, t, 3, "<sample pubkey>", "", -1, "", "")
	doFeedTest(db, t, 1, "", "", 2, "", "")
	doFeedTest(db, t, 1, "", "", -1, "", ".*retrieve.*")
	doFeedTest(db, t, 3, "<sample pubkey>", "", -1, "test", "")
}

func doFeedTest(db DB, t *testing.T, goal int, from string, previous string, seq int, msg_type string, data string) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	dst := make(chan message.Message)

	go db.GetFeed(ctx, dst, from, previous, seq, msg_type, data)

	counter := 0

	msg := messages[0] // Doesn't matter
	for (msg != message.Message{}) {
		msg = <-dst

		// An empty message means the GetFeed func finished its work
		if (msg != message.Message{}) {
			counter += 1
			t.Log(msg)
		}
	}

	if counter != goal {
		t.Error(errors.New("feed didn't send the correct number of messages"))
	}
}

func TestGetMessage(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	msg, err := db.GetMessage("<sample hash 2>")
	if err != nil {
		t.Error(err)
	}

	if msg != messages[0] {
		t.Error(errors.New("messages doesn't match"))
	}
}
<commit_msg>[database] Automatic cleanup<commit_after>package database

import (
	"errors"
	"testing"
	"time"
	"os"

	"context"

	"github.com/Bit-Nation/BITNATION-Panthalassa/message"
)

// TODO: avoid opening and closing the DB each times
// TODO: avoid filling the DB each times

var messages = []message.Message{
	message.Message{From: "<sample pubkey>", Previous: "", Seq: 1, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<content to retrieve>"}, Hash: "<sample hash 1>", Signature: "<sample sign>"},
	message.Message{From: "<sample pubkey>", Previous: "<sample hash 1>", Seq: 2, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<content to retrieve>"}, Hash: "<sample hash 2>", Signature: "<sample sign>"},
	message.Message{From: "<sample pubkey>", Previous: "<sample hash 2>", Seq: 3, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<sample content>"}, Hash: "<sample hash 3>", Signature: "<sample sign>"},
	message.Message{From: "<sample pubkey>", Previous: "<sample hash 3>", Seq: 4, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<sample content>"}, Hash: "<sample hash 4>", Signature: "<sample sign>"},

	message.Message{From: "<sample pubkey remove>", Previous: "", Seq: 1, Timestamp: time.Now(), Content: message.MessageContent{Type: "test", Data:"<sample content>"}, Hash: "<sample hash>", Signature: "<sample sign>"},
}

func TestOpenAndClose(t *testing.T) {
	os.RemoveAll("/tmp/test.db") // Delete previous attempts

	db := DB{File: "/tmp/test.db"}

	err := db.Open()
	if err != nil {
		t.Error(err)
	}

	err = db.Close()
	if err != nil {
		t.Error(err)
	}
}
// TODO: check adding valid and invalid msg
func TestAdd(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	for _, msg := range messages {
		err := db.AddMessage(msg)
		if err != nil {
			t.Error(err)
		}
	}
}

// This test is stupid but who knows...
func TestGetLastMessage(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	msg, err := db.GetLastMessage(messages[3].From)
	if err != nil {
		t.Error(err)
	}

	if msg != messages[3] {
		t.Error(errors.New("previous message doesn't match expected one"))
	}
}

func TestRemove_ExistingMessage(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	err := db.Remove("<sample hash>")
	if err != nil {
		t.Error(err)
	}
}

// TODO: better tests
func TestGetFeed(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	doFeedTest(db, t, 3, "<sample pubkey>", "", -1, "", "")
	doFeedTest(db, t, 1, "", "", 2, "", "")
	doFeedTest(db, t, 1, "", "", -1, "", ".*retrieve.*")
	doFeedTest(db, t, 3, "<sample pubkey>", "", -1, "test", "")
}

func doFeedTest(db DB, t *testing.T, goal int, from string, previous string, seq int, msg_type string, data string) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	dst := make(chan message.Message)

	go db.GetFeed(ctx, dst, from, previous, seq, msg_type, data)

	counter := 0

	msg := messages[0] // Doesn't matter
	for (msg != message.Message{}) {
		msg = <-dst

		// An empty message means the GetFeed func finished its work
		if (msg != message.Message{}) {
			counter += 1
			t.Log(msg)
		}
	}

	if counter != goal {
		t.Error(errors.New("feed didn't send the correct number of messages"))
	}
}

func TestGetMessage(t *testing.T) {
	db := DB{File: "/tmp/test.db"}
	db.Open()
	defer db.Close()

	msg, err := db.GetMessage("<sample hash 2>")
	if err != nil {
		t.Error(err)
	}

	if msg != messages[0] {
		t.Error(errors.New("messages doesn't match"))
	}
}
<|endoftext|>


<commit_before>package datastruct

import (
	"sync/atomic"
)

// RingBuffer is a rudimentary implementation of a fixed-size circular buffer.
// All buffered elements are strings.
type RingBuffer struct {
	size    int64
	counter int64
	buf     []string
}

// NewRingBuffer returns an initialised ring buffer.
func NewRingBuffer(size int64) *RingBuffer {
	if size < 1 {
		panic("NewRingBuffer: size must be greater than 0")
	}
	return &RingBuffer{
		size: size,
		buf:  make([]string, size),
	}
}

// Push places a new element into ring buffer.
func (r *RingBuffer) Push(elem string) {
	elemIndex := atomic.AddInt64(&r.counter, 1)
	r.buf[elemIndex%r.size] = elem
}

/*
Clear sets all buffered elements to empty string, consequently GetAll function will return an empty string array
indicating there's no element.
*/
func (r *RingBuffer) Clear() {
	r.buf = make([]string, r.size)
}

/*
IterateReverse traverses the ring buffer from the latest element to the oldest element.
The iterator function is fed an element value as sole parameter. If the function returns false, iteration is stopped
immediately. The total number of elements iterated is not predictable, and iteration loop always skips empty elements.
*/
func (r *RingBuffer) IterateReverse(fun func(string) bool) {
	currentIndex := r.counter % r.size
	for i := currentIndex; i >= 0; i-- {
		value := r.buf[i]
		if value != "" {
			if !fun(value) {
				return
			}
		}
	}
	for i := r.size - 1; i > currentIndex; i-- {
		value := r.buf[i]
		if value != "" {
			if !fun(value) {
				return
			}
		}
	}
}

// GetAll returns all elements (oldest to the latest) of the ring buffer in a string array.
func (r *RingBuffer) GetAll() (ret []string) {
	reversed := make([]string, 0, r.size)
	r.IterateReverse(func(elem string) bool {
		reversed = append(reversed, elem)
		return true
	})

	ret = make([]string, len(reversed))
	for i, s := range reversed {
		ret[len(ret)-1-i] = s
	}
	return
}
<commit_msg>fix unlikely data race in RingBuffer<commit_after>package datastruct

import (
	"sync"
)

// RingBuffer is a rudimentary implementation of a fixed-size circular buffer.
// All buffered elements are strings.
type RingBuffer struct {
	size    int64
	counter int64
	buf     []string
	mutex   *sync.Mutex
}

// NewRingBuffer returns an initialised ring buffer.
func NewRingBuffer(size int64) *RingBuffer {
	if size < 1 {
		panic("NewRingBuffer: size must be greater than 0")
	}
	return &RingBuffer{
		size:  size,
		buf:   make([]string, size),
		mutex: new(sync.Mutex),
	}
}

// Push places a new element into ring buffer.
func (r *RingBuffer) Push(elem string) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	r.counter++
	r.buf[r.counter%r.size] = elem
}

/*
Clear sets all buffered elements to empty string, consequently GetAll function will return an empty string array
indicating there's no element.
*/
func (r *RingBuffer) Clear() {
	r.buf = make([]string, r.size)
}

/*
IterateReverse traverses the ring buffer from the latest element to the oldest element.
The iterator function is fed an element value as sole parameter. If the function returns false, iteration is stopped
immediately. The total number of elements iterated is not predictable, and iteration loop always skips empty elements.
*/
func (r *RingBuffer) IterateReverse(fun func(string) bool) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	currentIndex := r.counter % r.size
	for i := currentIndex; i >= 0; i-- {
		value := r.buf[i]
		if value != "" {
			if !fun(value) {
				return
			}
		}
	}
	for i := r.size - 1; i > currentIndex; i-- {
		value := r.buf[i]
		if value != "" {
			if !fun(value) {
				return
			}
		}
	}
}

// GetAll returns all elements (oldest to the latest) of the ring buffer in a string array.
func (r *RingBuffer) GetAll() (ret []string) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	reversed := make([]string, 0, r.size)
	r.IterateReverse(func(elem string) bool {
		reversed = append(reversed, elem)
		return true
	})

	ret = make([]string, len(reversed))
	for i, s := range reversed {
		ret[len(ret)-1-i] = s
	}
	return
}
<|endoftext|>