<commit_before>package ssh

import (
	"bufio"
	"bytes"
	"code.google.com/p/go.crypto/ssh"
	"errors"
	"fmt"
	"github.com/mitchellh/packer/packer"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type comm struct {
	client *ssh.ClientConn
	config *Config
	conn   net.Conn
}

// Config is the structure used to configure the SSH communicator.
type Config struct {
	// The configuration of the Go SSH connection
	SSHConfig *ssh.ClientConfig

	// Connection returns a new connection. The current connection
	// in use will be closed as part of the Close method, or in the
	// case an error occurs.
	Connection func() (net.Conn, error)

	// NoPty, if true, will not request a pty from the remote end.
	NoPty bool
}

// Creates a new packer.Communicator implementation over SSH. This takes
// an already existing TCP connection and SSH configuration.
func New(config *Config) (result *comm, err error) {
	// Establish an initial connection and connect
	result = &comm{
		config: config,
	}

	if err = result.reconnect(); err != nil {
		result = nil
		return
	}

	return
}

func (c *comm) Start(cmd *packer.RemoteCmd) (err error) {
	session, err := c.newSession()
	if err != nil {
		return
	}

	// Setup our session
	session.Stdin = cmd.Stdin
	session.Stdout = cmd.Stdout
	session.Stderr = cmd.Stderr

	if !c.config.NoPty {
		// Request a PTY
		termModes := ssh.TerminalModes{
			ssh.ECHO:          0,     // do not echo
			ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
			ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
		}

		if err = session.RequestPty("xterm", 80, 40, termModes); err != nil {
			return
		}
	}

	log.Printf("starting remote command: %s", cmd.Command)
	err = session.Start(cmd.Command + "\n")
	if err != nil {
		return
	}

	// A channel to keep track of our done state
	doneCh := make(chan struct{})
	sessionLock := new(sync.Mutex)
	timedOut := false

	// Start a goroutine to wait for the session to end and set the
	// exit boolean and status.
	go func() {
		defer session.Close()

		err := session.Wait()
		exitStatus := 0
		if err != nil {
			exitErr, ok := err.(*ssh.ExitError)
			if ok {
				exitStatus = exitErr.ExitStatus()
			}
		}

		sessionLock.Lock()
		defer sessionLock.Unlock()

		if timedOut {
			// We timed out, so set the exit status to -1
			exitStatus = -1
		}

		log.Printf("remote command exited with '%d': %s", exitStatus, cmd.Command)
		cmd.SetExited(exitStatus)
		close(doneCh)
	}()

	go func() {
		failures := 0
		for {
			dummy, err := c.config.Connection()
			if err == nil {
				failures = 0
				dummy.Close()
			}

			select {
			case <-doneCh:
				return
			default:
			}

			if err != nil {
				log.Printf("background SSH connection checker failure: %s", err)
				failures += 1
			}

			if failures < 5 {
				time.Sleep(5 * time.Second)
				continue
			}

			// Acquire a lock in order to modify session state
			sessionLock.Lock()
			defer sessionLock.Unlock()

			// Kill the connection and mark that we timed out.
			log.Printf("Too many SSH connection failures. Killing it!")
			c.conn.Close()
			timedOut = true

			return
		}
	}()

	return
}

func (c *comm) Upload(path string, input io.Reader) error {
	// The target directory and file for talking the SCP protocol
	target_dir := filepath.Dir(path)
	target_file := filepath.Base(path)

	// On windows, filepath.Dir uses backslash seperators (ie. "\tmp").
	// This does not work when the target host is unix.  Switch to forward slash
	// which works for unix and windows
	target_dir = filepath.ToSlash(target_dir)

	scpFunc := func(w io.Writer, stdoutR *bufio.Reader) error {
		return scpUploadFile(target_file, input, w, stdoutR)
	}

	return c.scpSession("scp -vt "+target_dir, scpFunc)
}

func (c *comm) UploadDir(dst string, src string, excl []string) error {
	log.Printf("Upload dir '%s' to '%s'", src, dst)
	scpFunc := func(w io.Writer, r *bufio.Reader) error {
		uploadEntries := func() error {
			f, err := os.Open(src)
			if err != nil {
				return err
			}
			defer f.Close()

			entries, err := f.Readdir(-1)
			if err != nil {
				return err
			}

			return scpUploadDir(src, entries, w, r)
		}

		if src[len(src)-1] != '/' {
			log.Printf("No trailing slash, creating the source directory name")
			return scpUploadDirProtocol(filepath.Base(src), w, r, uploadEntries)
		} else {
			// Trailing slash, so only upload the contents
			return uploadEntries()
		}
	}

	return c.scpSession("scp -rvt "+dst, scpFunc)
}

func (c *comm) Download(string, io.Writer) error {
	panic("not implemented yet")
}

func (c *comm) newSession() (session *ssh.Session, err error) {
	log.Println("opening new ssh session")
	if c.client == nil {
		err = errors.New("client not available")
	} else {
		session, err = c.client.NewSession()
	}

	if err != nil {
		log.Printf("ssh session open error: '%s', attempting reconnect", err)
		if err := c.reconnect(); err != nil {
			return nil, err
		}

		return c.client.NewSession()
	}

	return session, nil
}

func (c *comm) reconnect() (err error) {
	if c.conn != nil {
		c.conn.Close()
	}

	// Set the conn and client to nil since we'll recreate it
	c.conn = nil
	c.client = nil

	log.Printf("reconnecting to TCP connection for SSH")
	c.conn, err = c.config.Connection()
	if err != nil {
		log.Printf("reconnection error: %s", err)
		return
	}

	log.Printf("handshaking with SSH")
	c.client, err = ssh.Client(c.conn, c.config.SSHConfig)
	if err != nil {
		log.Printf("handshake error: %s", err)
	}

	return
}

func (c *comm) scpSession(scpCommand string, f func(io.Writer, *bufio.Reader) error) error {
	session, err := c.newSession()
	if err != nil {
		return err
	}
	defer session.Close()

	// Get a pipe to stdin so that we can send data down
	stdinW, err := session.StdinPipe()
	if err != nil {
		return err
	}

	// We only want to close once, so we nil w after we close it,
	// and only close in the defer if it hasn't been closed already.
	defer func() {
		if stdinW != nil {
			stdinW.Close()
		}
	}()

	// Get a pipe to stdout so that we can get responses back
	stdoutPipe, err := session.StdoutPipe()
	if err != nil {
		return err
	}
	stdoutR := bufio.NewReader(stdoutPipe)

	// Set stderr to a bytes buffer
	stderr := new(bytes.Buffer)
	session.Stderr = stderr

	// Start the sink mode on the other side
	// TODO(mitchellh): There are probably issues with shell escaping the path
	log.Println("Starting remote scp process: ", scpCommand)
	if err := session.Start(scpCommand); err != nil {
		return err
	}

	// Call our callback that executes in the context of SCP. We ignore
	// EOF errors if they occur because it usually means that SCP prematurely
	// ended on the other side.
	log.Println("Started SCP session, beginning transfers...")
	if err := f(stdinW, stdoutR); err != nil && err != io.EOF {
		return err
	}

	// Close the stdin, which sends an EOF, and then set w to nil so that
	// our defer func doesn't close it again since that is unsafe with
	// the Go SSH package.
	log.Println("SCP session complete, closing stdin pipe.")
	stdinW.Close()
	stdinW = nil

	// Wait for the SCP connection to close, meaning it has consumed all
	// our data and has completed. Or has errored.
	log.Println("Waiting for SSH session to complete.")
	err = session.Wait()
	if err != nil {
		if exitErr, ok := err.(*ssh.ExitError); ok {
			// Otherwise, we have an ExitErorr, meaning we can just read
			// the exit status
			log.Printf("non-zero exit status: %d", exitErr.ExitStatus())

			// If we exited with status 127, it means SCP isn't available.
			// Return a more descriptive error for that.
			if exitErr.ExitStatus() == 127 {
				return errors.New(
					"SCP failed to start. This usually means that SCP is not\n" +
						"properly installed on the remote system.")
			}
		}

		return err
	}

	log.Printf("scp stderr (length %d): %s", stderr.Len(), stderr.String())
	return nil
}

// checkSCPStatus checks that a prior command sent to SCP completed
// successfully. If it did not complete successfully, an error will
// be returned.
func checkSCPStatus(r *bufio.Reader) error {
	code, err := r.ReadByte()
	if err != nil {
		return err
	}

	if code != 0 {
		// Treat any non-zero (really 1 and 2) as fatal errors
		message, _, err := r.ReadLine()
		if err != nil {
			return fmt.Errorf("Error reading error message: %s", err)
		}

		return errors.New(string(message))
	}

	return nil
}

func scpUploadFile(dst string, src io.Reader, w io.Writer, r *bufio.Reader) error {
	// Determine the length of the upload content by copying it
	// into an in-memory buffer. Note that this means what we upload
	// must fit into memory.
	log.Println("Copying input data into in-memory buffer so we can get the length")
	inputBuf := new(bytes.Buffer)
	if _, err := io.Copy(inputBuf, src); err != nil {
		return err
	}

	// Start the protocol
	log.Println("Beginning file upload...")
	fmt.Fprintln(w, "C0644", inputBuf.Len(), dst)
	err := checkSCPStatus(r)
	if err != nil {
		return err
	}

	if _, err := io.Copy(w, inputBuf); err != nil {
		return err
	}

	fmt.Fprint(w, "\x00")
	err = checkSCPStatus(r)
	if err != nil {
		return err
	}

	return nil
}

func scpUploadDirProtocol(name string, w io.Writer, r *bufio.Reader, f func() error) error {
	log.Printf("SCP: starting directory upload: %s", name)
	fmt.Fprintln(w, "D0755 0", name)
	err := checkSCPStatus(r)
	if err != nil {
		return err
	}

	if err := f(); err != nil {
		return err
	}

	fmt.Fprintln(w, "E")
	if err != nil {
		return err
	}

	return nil
}

func scpUploadDir(root string, fs []os.FileInfo, w io.Writer, r *bufio.Reader) error {
	for _, fi := range fs {
		realPath := filepath.Join(root, fi.Name())

		// Track if this is actually a symlink to a directory. If it is
		// a symlink to a file we don't do any special behavior because uploading
		// a file just works. If it is a directory, we need to know so we
		// treat it as such.
		isSymlinkToDir := false
		if fi.Mode() & os.ModeSymlink == os.ModeSymlink {
			symPath, err := filepath.EvalSymlinks(realPath)
			if err != nil {
				return err
			}

			symFi, err := os.Lstat(symPath)
			if err != nil {
				return err
			}

			isSymlinkToDir = symFi.IsDir()
		}

		if !fi.IsDir() && !isSymlinkToDir {
			// It is a regular file (or symlink to a file), just upload it
			f, err := os.Open(realPath)
			if err != nil {
				return err
			}

			err = func() error {
				defer f.Close()
				return scpUploadFile(fi.Name(), f, w, r)
			}()

			if err != nil {
				return err
			}

			continue
		}

		// It is a directory, recursively upload
		err := scpUploadDirProtocol(fi.Name(), w, r, func() error {
			f, err := os.Open(realPath)
			if err != nil {
				return err
			}
			defer f.Close()

			entries, err := f.Readdir(-1)
			if err != nil {
				return err
			}

			return scpUploadDir(realPath, entries, w, r)
		})
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>go fmt<commit_after>package ssh

import (
	"bufio"
	"bytes"
	"code.google.com/p/go.crypto/ssh"
	"errors"
	"fmt"
	"github.com/mitchellh/packer/packer"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type comm struct {
	client *ssh.ClientConn
	config *Config
	conn   net.Conn
}

// Config is the structure used to configure the SSH communicator.
type Config struct {
	// The configuration of the Go SSH connection
	SSHConfig *ssh.ClientConfig

	// Connection returns a new connection. The current connection
	// in use will be closed as part of the Close method, or in the
	// case an error occurs.
	Connection func() (net.Conn, error)

	// NoPty, if true, will not request a pty from the remote end.
	NoPty bool
}

// Creates a new packer.Communicator implementation over SSH. This takes
// an already existing TCP connection and SSH configuration.
func New(config *Config) (result *comm, err error) {
	// Establish an initial connection and connect
	result = &comm{
		config: config,
	}

	if err = result.reconnect(); err != nil {
		result = nil
		return
	}

	return
}

func (c *comm) Start(cmd *packer.RemoteCmd) (err error) {
	session, err := c.newSession()
	if err != nil {
		return
	}

	// Setup our session
	session.Stdin = cmd.Stdin
	session.Stdout = cmd.Stdout
	session.Stderr = cmd.Stderr

	if !c.config.NoPty {
		// Request a PTY
		termModes := ssh.TerminalModes{
			ssh.ECHO:          0,     // do not echo
			ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
			ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
		}

		if err = session.RequestPty("xterm", 80, 40, termModes); err != nil {
			return
		}
	}

	log.Printf("starting remote command: %s", cmd.Command)
	err = session.Start(cmd.Command + "\n")
	if err != nil {
		return
	}

	// A channel to keep track of our done state
	doneCh := make(chan struct{})
	sessionLock := new(sync.Mutex)
	timedOut := false

	// Start a goroutine to wait for the session to end and set the
	// exit boolean and status.
	go func() {
		defer session.Close()

		err := session.Wait()
		exitStatus := 0
		if err != nil {
			exitErr, ok := err.(*ssh.ExitError)
			if ok {
				exitStatus = exitErr.ExitStatus()
			}
		}

		sessionLock.Lock()
		defer sessionLock.Unlock()

		if timedOut {
			// We timed out, so set the exit status to -1
			exitStatus = -1
		}

		log.Printf("remote command exited with '%d': %s", exitStatus, cmd.Command)
		cmd.SetExited(exitStatus)
		close(doneCh)
	}()

	go func() {
		failures := 0
		for {
			dummy, err := c.config.Connection()
			if err == nil {
				failures = 0
				dummy.Close()
			}

			select {
			case <-doneCh:
				return
			default:
			}

			if err != nil {
				log.Printf("background SSH connection checker failure: %s", err)
				failures += 1
			}

			if failures < 5 {
				time.Sleep(5 * time.Second)
				continue
			}

			// Acquire a lock in order to modify session state
			sessionLock.Lock()
			defer sessionLock.Unlock()

			// Kill the connection and mark that we timed out.
			log.Printf("Too many SSH connection failures. Killing it!")
			c.conn.Close()
			timedOut = true

			return
		}
	}()

	return
}

func (c *comm) Upload(path string, input io.Reader) error {
	// The target directory and file for talking the SCP protocol
	target_dir := filepath.Dir(path)
	target_file := filepath.Base(path)

	// On windows, filepath.Dir uses backslash seperators (ie. "\tmp").
	// This does not work when the target host is unix.  Switch to forward slash
	// which works for unix and windows
	target_dir = filepath.ToSlash(target_dir)

	scpFunc := func(w io.Writer, stdoutR *bufio.Reader) error {
		return scpUploadFile(target_file, input, w, stdoutR)
	}

	return c.scpSession("scp -vt "+target_dir, scpFunc)
}

func (c *comm) UploadDir(dst string, src string, excl []string) error {
	log.Printf("Upload dir '%s' to '%s'", src, dst)
	scpFunc := func(w io.Writer, r *bufio.Reader) error {
		uploadEntries := func() error {
			f, err := os.Open(src)
			if err != nil {
				return err
			}
			defer f.Close()

			entries, err := f.Readdir(-1)
			if err != nil {
				return err
			}

			return scpUploadDir(src, entries, w, r)
		}

		if src[len(src)-1] != '/' {
			log.Printf("No trailing slash, creating the source directory name")
			return scpUploadDirProtocol(filepath.Base(src), w, r, uploadEntries)
		} else {
			// Trailing slash, so only upload the contents
			return uploadEntries()
		}
	}

	return c.scpSession("scp -rvt "+dst, scpFunc)
}

func (c *comm) Download(string, io.Writer) error {
	panic("not implemented yet")
}

func (c *comm) newSession() (session *ssh.Session, err error) {
	log.Println("opening new ssh session")
	if c.client == nil {
		err = errors.New("client not available")
	} else {
		session, err = c.client.NewSession()
	}

	if err != nil {
		log.Printf("ssh session open error: '%s', attempting reconnect", err)
		if err := c.reconnect(); err != nil {
			return nil, err
		}

		return c.client.NewSession()
	}

	return session, nil
}

func (c *comm) reconnect() (err error) {
	if c.conn != nil {
		c.conn.Close()
	}

	// Set the conn and client to nil since we'll recreate it
	c.conn = nil
	c.client = nil

	log.Printf("reconnecting to TCP connection for SSH")
	c.conn, err = c.config.Connection()
	if err != nil {
		log.Printf("reconnection error: %s", err)
		return
	}

	log.Printf("handshaking with SSH")
	c.client, err = ssh.Client(c.conn, c.config.SSHConfig)
	if err != nil {
		log.Printf("handshake error: %s", err)
	}

	return
}

func (c *comm) scpSession(scpCommand string, f func(io.Writer, *bufio.Reader) error) error {
	session, err := c.newSession()
	if err != nil {
		return err
	}
	defer session.Close()

	// Get a pipe to stdin so that we can send data down
	stdinW, err := session.StdinPipe()
	if err != nil {
		return err
	}

	// We only want to close once, so we nil w after we close it,
	// and only close in the defer if it hasn't been closed already.
	defer func() {
		if stdinW != nil {
			stdinW.Close()
		}
	}()

	// Get a pipe to stdout so that we can get responses back
	stdoutPipe, err := session.StdoutPipe()
	if err != nil {
		return err
	}
	stdoutR := bufio.NewReader(stdoutPipe)

	// Set stderr to a bytes buffer
	stderr := new(bytes.Buffer)
	session.Stderr = stderr

	// Start the sink mode on the other side
	// TODO(mitchellh): There are probably issues with shell escaping the path
	log.Println("Starting remote scp process: ", scpCommand)
	if err := session.Start(scpCommand); err != nil {
		return err
	}

	// Call our callback that executes in the context of SCP. We ignore
	// EOF errors if they occur because it usually means that SCP prematurely
	// ended on the other side.
	log.Println("Started SCP session, beginning transfers...")
	if err := f(stdinW, stdoutR); err != nil && err != io.EOF {
		return err
	}

	// Close the stdin, which sends an EOF, and then set w to nil so that
	// our defer func doesn't close it again since that is unsafe with
	// the Go SSH package.
	log.Println("SCP session complete, closing stdin pipe.")
	stdinW.Close()
	stdinW = nil

	// Wait for the SCP connection to close, meaning it has consumed all
	// our data and has completed. Or has errored.
	log.Println("Waiting for SSH session to complete.")
	err = session.Wait()
	if err != nil {
		if exitErr, ok := err.(*ssh.ExitError); ok {
			// Otherwise, we have an ExitErorr, meaning we can just read
			// the exit status
			log.Printf("non-zero exit status: %d", exitErr.ExitStatus())

			// If we exited with status 127, it means SCP isn't available.
			// Return a more descriptive error for that.
			if exitErr.ExitStatus() == 127 {
				return errors.New(
					"SCP failed to start. This usually means that SCP is not\n" +
						"properly installed on the remote system.")
			}
		}

		return err
	}

	log.Printf("scp stderr (length %d): %s", stderr.Len(), stderr.String())
	return nil
}

// checkSCPStatus checks that a prior command sent to SCP completed
// successfully. If it did not complete successfully, an error will
// be returned.
func checkSCPStatus(r *bufio.Reader) error {
	code, err := r.ReadByte()
	if err != nil {
		return err
	}

	if code != 0 {
		// Treat any non-zero (really 1 and 2) as fatal errors
		message, _, err := r.ReadLine()
		if err != nil {
			return fmt.Errorf("Error reading error message: %s", err)
		}

		return errors.New(string(message))
	}

	return nil
}

func scpUploadFile(dst string, src io.Reader, w io.Writer, r *bufio.Reader) error {
	// Determine the length of the upload content by copying it
	// into an in-memory buffer. Note that this means what we upload
	// must fit into memory.
	log.Println("Copying input data into in-memory buffer so we can get the length")
	inputBuf := new(bytes.Buffer)
	if _, err := io.Copy(inputBuf, src); err != nil {
		return err
	}

	// Start the protocol
	log.Println("Beginning file upload...")
	fmt.Fprintln(w, "C0644", inputBuf.Len(), dst)
	err := checkSCPStatus(r)
	if err != nil {
		return err
	}

	if _, err := io.Copy(w, inputBuf); err != nil {
		return err
	}

	fmt.Fprint(w, "\x00")
	err = checkSCPStatus(r)
	if err != nil {
		return err
	}

	return nil
}

func scpUploadDirProtocol(name string, w io.Writer, r *bufio.Reader, f func() error) error {
	log.Printf("SCP: starting directory upload: %s", name)
	fmt.Fprintln(w, "D0755 0", name)
	err := checkSCPStatus(r)
	if err != nil {
		return err
	}

	if err := f(); err != nil {
		return err
	}

	fmt.Fprintln(w, "E")
	if err != nil {
		return err
	}

	return nil
}

func scpUploadDir(root string, fs []os.FileInfo, w io.Writer, r *bufio.Reader) error {
	for _, fi := range fs {
		realPath := filepath.Join(root, fi.Name())

		// Track if this is actually a symlink to a directory. If it is
		// a symlink to a file we don't do any special behavior because uploading
		// a file just works. If it is a directory, we need to know so we
		// treat it as such.
		isSymlinkToDir := false
		if fi.Mode()&os.ModeSymlink == os.ModeSymlink {
			symPath, err := filepath.EvalSymlinks(realPath)
			if err != nil {
				return err
			}

			symFi, err := os.Lstat(symPath)
			if err != nil {
				return err
			}

			isSymlinkToDir = symFi.IsDir()
		}

		if !fi.IsDir() && !isSymlinkToDir {
			// It is a regular file (or symlink to a file), just upload it
			f, err := os.Open(realPath)
			if err != nil {
				return err
			}

			err = func() error {
				defer f.Close()
				return scpUploadFile(fi.Name(), f, w, r)
			}()

			if err != nil {
				return err
			}

			continue
		}

		// It is a directory, recursively upload
		err := scpUploadDirProtocol(fi.Name(), w, r, func() error {
			f, err := os.Open(realPath)
			if err != nil {
				return err
			}
			defer f.Close()

			entries, err := f.Readdir(-1)
			if err != nil {
				return err
			}

			return scpUploadDir(realPath, entries, w, r)
		})
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package compress

import (
	"errors"

	"github.com/bkaradzic/go-lz4"
	"github.com/golang/snappy"
)

var (
	// ErrBadAlgo is returned on a unsupported/unknown algorithm.
	ErrBadAlgo = errors.New("Invalid algorithm type")
)

// Algorithm is the common interface for all supported algorithms.
type Algorithm interface {
	Encode([]byte) ([]byte, error)
	Decode([]byte) ([]byte, error)
}

type noneAlgo struct{}
type snappyAlgo struct{}
type lz4Algo struct{}

var (
	algoMap = map[AlgorithmType]Algorithm{
		AlgoNone:   noneAlgo{},
		AlgoSnappy: snappyAlgo{},
		AlgoLZ4:    lz4Algo{},
	}

	algoToString = map[AlgorithmType]string{
		AlgoNone:   "none",
		AlgoSnappy: "snappy",
		AlgoLZ4:    "lz4",
	}

	stringToAlgo = map[string]AlgorithmType{
		"none":   AlgoNone,
		"snappy": AlgoSnappy,
		"lz4":    AlgoLZ4,
	}
)

// AlgoNone
func (a noneAlgo) Encode(src []byte) ([]byte, error) {
	return src, nil
}

func (a noneAlgo) Decode(src []byte) ([]byte, error) {
	return src, nil
}

// AlgoSnappy
func (a snappyAlgo) Encode(src []byte) ([]byte, error) {
	return snappy.Encode(nil, src), nil

}

func (a snappyAlgo) Decode(src []byte) ([]byte, error) {
	return snappy.Decode(nil, src)
}

// AlgoLZ4
func (a lz4Algo) Encode(src []byte) ([]byte, error) {
	return lz4.Encode(nil, src)
}

func (a lz4Algo) Decode(src []byte) ([]byte, error) {
	return lz4.Decode(nil, src)
}

// AlgorithmFromType returns a interface to the given AlgorithmType.
func AlgorithmFromType(a AlgorithmType) (Algorithm, error) {
	if algo, ok := algoMap[a]; ok {
		return algo, nil
	}
	return nil, ErrBadAlgo
}

// AlgoToString converts a algorithm type to a string.
func AlgoToString(a AlgorithmType) string {
	algo, ok := algoToString[a]
	if !ok {
		return "unknown algorithm"
	}
	return algo
}

// AlgoFromString tries to convert a string to AlgorithmType
func AlgoFromString(s string) (AlgorithmType, error) {
	algoType, ok := stringToAlgo[s]
	if !ok {
		return 0, errors.New("Invalid algorithm name")
	}
	return algoType, nil
}
<commit_msg>catfs: mio: fix test build error<commit_after>package compress

import (
	"errors"

	"github.com/bkaradzic/go-lz4"
	"github.com/golang/snappy"
)

var (
	// ErrBadAlgo is returned on a unsupported/unknown algorithm.
	ErrBadAlgo = errors.New("Invalid algorithm type")
)

// Algorithm is the common interface for all supported algorithms.
type Algorithm interface {
	Encode([]byte) ([]byte, error)
	Decode([]byte) ([]byte, error)
}

type noneAlgo struct{}
type snappyAlgo struct{}
type lz4Algo struct{}

var (
	// AlgoMap is a map of available algorithms.
	AlgoMap = map[AlgorithmType]Algorithm{
		AlgoNone:   noneAlgo{},
		AlgoSnappy: snappyAlgo{},
		AlgoLZ4:    lz4Algo{},
	}

	algoToString = map[AlgorithmType]string{
		AlgoNone:   "none",
		AlgoSnappy: "snappy",
		AlgoLZ4:    "lz4",
	}

	stringToAlgo = map[string]AlgorithmType{
		"none":   AlgoNone,
		"snappy": AlgoSnappy,
		"lz4":    AlgoLZ4,
	}
)

// AlgoNone
func (a noneAlgo) Encode(src []byte) ([]byte, error) {
	return src, nil
}

func (a noneAlgo) Decode(src []byte) ([]byte, error) {
	return src, nil
}

// AlgoSnappy
func (a snappyAlgo) Encode(src []byte) ([]byte, error) {
	return snappy.Encode(nil, src), nil

}

func (a snappyAlgo) Decode(src []byte) ([]byte, error) {
	return snappy.Decode(nil, src)
}

// AlgoLZ4
func (a lz4Algo) Encode(src []byte) ([]byte, error) {
	return lz4.Encode(nil, src)
}

func (a lz4Algo) Decode(src []byte) ([]byte, error) {
	return lz4.Decode(nil, src)
}

// AlgorithmFromType returns a interface to the given AlgorithmType.
func AlgorithmFromType(a AlgorithmType) (Algorithm, error) {
	if algo, ok := AlgoMap[a]; ok {
		return algo, nil
	}
	return nil, ErrBadAlgo
}

// AlgoToString converts a algorithm type to a string.
func AlgoToString(a AlgorithmType) string {
	algo, ok := algoToString[a]
	if !ok {
		return "unknown algorithm"
	}
	return algo
}

// AlgoFromString tries to convert a string to AlgorithmType
func AlgoFromString(s string) (AlgorithmType, error) {
	algoType, ok := stringToAlgo[s]
	if !ok {
		return 0, errors.New("Invalid algorithm name")
	}
	return algoType, nil
}
<|endoftext|>


<commit_before>package user

import (
	"github.com/cloudfoundry/cli/cf/api"
	"github.com/cloudfoundry/cli/cf/api/spaces"
	"github.com/cloudfoundry/cli/cf/command_registry"
	"github.com/cloudfoundry/cli/cf/configuration/core_config"
	. "github.com/cloudfoundry/cli/cf/i18n"
	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/cli/cf/requirements"
	"github.com/cloudfoundry/cli/cf/terminal"
	"github.com/cloudfoundry/cli/flags"
	"github.com/cloudfoundry/cli/plugin/models"
)

type SpaceUsers struct {
	ui          terminal.UI
	config      core_config.Reader
	spaceRepo   spaces.SpaceRepository
	userRepo    api.UserRepository
	orgReq      requirements.OrganizationRequirement
	pluginModel *[]plugin_models.GetSpaceUsers_Model
	pluginCall  bool
}

func init() {
	command_registry.Register(&SpaceUsers{})
}

func (cmd *SpaceUsers) MetaData() command_registry.CommandMetadata {
	return command_registry.CommandMetadata{
		Name:        "space-users",
		Description: T("Show space users by role"),
		Usage:       T("CF_NAME space-users ORG SPACE"),
	}
}

func (cmd *SpaceUsers) Requirements(requirementsFactory requirements.Factory, fc flags.FlagContext) (reqs []requirements.Requirement, err error) {
	if len(fc.Args()) != 2 {
		cmd.ui.Failed(T("Incorrect Usage. Requires arguments\n\n") + command_registry.Commands.CommandUsage("space-users"))
	}

	cmd.orgReq = requirementsFactory.NewOrganizationRequirement(fc.Args()[0])

	reqs = []requirements.Requirement{
		requirementsFactory.NewLoginRequirement(),
		cmd.orgReq,
	}
	return
}

func (cmd *SpaceUsers) SetDependency(deps command_registry.Dependency, pluginCall bool) command_registry.Command {
	cmd.ui = deps.Ui
	cmd.config = deps.Config
	cmd.userRepo = deps.RepoLocator.GetUserRepository()
	cmd.spaceRepo = deps.RepoLocator.GetSpaceRepository()
	cmd.pluginCall = pluginCall
	cmd.pluginModel = deps.PluginModels.SpaceUsers

	return cmd
}

func (cmd *SpaceUsers) Execute(c flags.FlagContext) {
	spaceName := c.Args()[1]
	org := cmd.orgReq.GetOrganization()

	space, apiErr := cmd.spaceRepo.FindByNameInOrg(spaceName, org.Guid)
	if apiErr != nil {
		cmd.ui.Failed(apiErr.Error())
	}

	cmd.ui.Say(T("Getting users in org {{.TargetOrg}} / space {{.TargetSpace}} as {{.CurrentUser}}",
		map[string]interface{}{
			"TargetOrg":   terminal.EntityNameColor(org.Name),
			"TargetSpace": terminal.EntityNameColor(space.Name),
			"CurrentUser": terminal.EntityNameColor(cmd.config.Username()),
		}))

	var spaceRoleToDisplayName = map[string]string{
		models.SPACE_MANAGER:   T("SPACE MANAGER"),
		models.SPACE_DEVELOPER: T("SPACE DEVELOPER"),
		models.SPACE_AUDITOR:   T("SPACE AUDITOR"),
	}

	var usersMap = make(map[string]plugin_models.GetSpaceUsers_Model)

	listUsers := cmd.getUserLister()

	var users []models.UserFields
	for role, displayName := range spaceRoleToDisplayName {
		users, apiErr = listUsers(space.Guid, role)

		cmd.ui.Say("")
		cmd.ui.Say("%s", terminal.HeaderColor(displayName))

		if len(users) == 0 {
			cmd.ui.Say("none")
		} else {
			for _, user := range users {
				cmd.ui.Say("  %s", user.Username)

				if cmd.pluginCall {
					u, found := usersMap[user.Username]
					if !found {
						u = plugin_models.GetSpaceUsers_Model{}
						u.Username = user.Username
						u.Guid = user.Guid
						u.IsAdmin = user.IsAdmin
						u.Roles = make([]string, 1)
						u.Roles[0] = role
						usersMap[user.Username] = u
					} else {
						u.Roles = append(u.Roles, role)
						usersMap[user.Username] = u
					}
				}
			}
		}

		if apiErr != nil {
			cmd.ui.Failed(T("Failed fetching space-users for role {{.SpaceRoleToDisplayName}}.\n{{.Error}}",
				map[string]interface{}{
					"Error":                  apiErr.Error(),
					"SpaceRoleToDisplayName": displayName,
				}))
			return
		}
	}

	if cmd.pluginCall {
		for _, v := range usersMap {
			*(cmd.pluginModel) = append(*(cmd.pluginModel), v)
		}
	}
}

func (cmd *SpaceUsers) getUserLister() func(spaceGuid string, role string) ([]models.UserFields, error) {
	if cmd.config.IsMinApiVersion("2.21.0") {
		return cmd.userRepo.ListUsersInSpaceForRoleWithNoUAA
	}
	return cmd.userRepo.ListUsersInSpaceForRole
}
<commit_msg>Fail early when fetching users<commit_after>package user

import (
	"github.com/cloudfoundry/cli/cf/api"
	"github.com/cloudfoundry/cli/cf/api/spaces"
	"github.com/cloudfoundry/cli/cf/command_registry"
	"github.com/cloudfoundry/cli/cf/configuration/core_config"
	. "github.com/cloudfoundry/cli/cf/i18n"
	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/cli/cf/requirements"
	"github.com/cloudfoundry/cli/cf/terminal"
	"github.com/cloudfoundry/cli/flags"
	"github.com/cloudfoundry/cli/plugin/models"
)

type SpaceUsers struct {
	ui          terminal.UI
	config      core_config.Reader
	spaceRepo   spaces.SpaceRepository
	userRepo    api.UserRepository
	orgReq      requirements.OrganizationRequirement
	pluginModel *[]plugin_models.GetSpaceUsers_Model
	pluginCall  bool
}

func init() {
	command_registry.Register(&SpaceUsers{})
}

func (cmd *SpaceUsers) MetaData() command_registry.CommandMetadata {
	return command_registry.CommandMetadata{
		Name:        "space-users",
		Description: T("Show space users by role"),
		Usage:       T("CF_NAME space-users ORG SPACE"),
	}
}

func (cmd *SpaceUsers) Requirements(requirementsFactory requirements.Factory, fc flags.FlagContext) (reqs []requirements.Requirement, err error) {
	if len(fc.Args()) != 2 {
		cmd.ui.Failed(T("Incorrect Usage. Requires arguments\n\n") + command_registry.Commands.CommandUsage("space-users"))
	}

	cmd.orgReq = requirementsFactory.NewOrganizationRequirement(fc.Args()[0])

	reqs = []requirements.Requirement{
		requirementsFactory.NewLoginRequirement(),
		cmd.orgReq,
	}
	return
}

func (cmd *SpaceUsers) SetDependency(deps command_registry.Dependency, pluginCall bool) command_registry.Command {
	cmd.ui = deps.Ui
	cmd.config = deps.Config
	cmd.userRepo = deps.RepoLocator.GetUserRepository()
	cmd.spaceRepo = deps.RepoLocator.GetSpaceRepository()
	cmd.pluginCall = pluginCall
	cmd.pluginModel = deps.PluginModels.SpaceUsers

	return cmd
}

func (cmd *SpaceUsers) Execute(c flags.FlagContext) {
	spaceName := c.Args()[1]
	org := cmd.orgReq.GetOrganization()

	space, err := cmd.spaceRepo.FindByNameInOrg(spaceName, org.Guid)
	if err != nil {
		cmd.ui.Failed(err.Error())
	}

	cmd.ui.Say(T("Getting users in org {{.TargetOrg}} / space {{.TargetSpace}} as {{.CurrentUser}}",
		map[string]interface{}{
			"TargetOrg":   terminal.EntityNameColor(org.Name),
			"TargetSpace": terminal.EntityNameColor(space.Name),
			"CurrentUser": terminal.EntityNameColor(cmd.config.Username()),
		}))

	var spaceRoleToDisplayName = map[string]string{
		models.SPACE_MANAGER:   T("SPACE MANAGER"),
		models.SPACE_DEVELOPER: T("SPACE DEVELOPER"),
		models.SPACE_AUDITOR:   T("SPACE AUDITOR"),
	}

	var usersMap = make(map[string]plugin_models.GetSpaceUsers_Model)

	listUsers := cmd.getUserLister()

	for role, displayName := range spaceRoleToDisplayName {
		users, err := listUsers(space.Guid, role)
		if err != nil {
			cmd.ui.Failed(T("Failed fetching space-users for role {{.SpaceRoleToDisplayName}}.\n{{.Error}}",
				map[string]interface{}{
					"Error":                  err.Error(),
					"SpaceRoleToDisplayName": displayName,
				}))
			return
		}

		cmd.ui.Say("")
		cmd.ui.Say("%s", terminal.HeaderColor(displayName))

		if len(users) == 0 {
			cmd.ui.Say("none")
		} else {
			for _, user := range users {
				if cmd.pluginCall {
					u, found := usersMap[user.Username]
					if !found {
						u = plugin_models.GetSpaceUsers_Model{}
						u.Username = user.Username
						u.Guid = user.Guid
						u.IsAdmin = user.IsAdmin
						u.Roles = make([]string, 1)
						u.Roles[0] = role
						usersMap[user.Username] = u
					} else {
						u.Roles = append(u.Roles, role)
						usersMap[user.Username] = u
					}
				} else {
					cmd.ui.Say("  %s", user.Username)
				}
			}
		}
	}

	if cmd.pluginCall {
		for _, v := range usersMap {
			*(cmd.pluginModel) = append(*(cmd.pluginModel), v)
		}
	}
}

func (cmd *SpaceUsers) getUserLister() func(spaceGuid string, role string) ([]models.UserFields, error) {
	if cmd.config.IsMinApiVersion("2.21.0") {
		return cmd.userRepo.ListUsersInSpaceForRoleWithNoUAA
	}
	return cmd.userRepo.ListUsersInSpaceForRole
}
<|endoftext|>


<commit_before>package mock

import (
	"github.com/mackerelio/checkers"
)

type Executer struct {
	CommandName string
	CommandArgs []string
	CommandResult string
	Status checkers.Status
	ParseResult string
}

func (e *Executer) MakeCommandName() string {
	return e.CommandName
}

func (e *Executer) MakeCommandArgs() []string {
	return e.CommandArgs
}

func (e *Executer) Parse(result string) (checkers.Status, string) {
	e.CommandResult = result
	return e.Status, e.ParseResult
}
<commit_msg>write go doc for mock executer<commit_after>package mock

import (
	"github.com/mackerelio/checkers"
)

// Executer is dummy executer for testing
type Executer struct {
	CommandName string
	CommandArgs []string
	CommandResult string
	Status checkers.Status
	ParseResult string
}

// MakeCommandName is dummy command name maker for testing
func (e *Executer) MakeCommandName() string {
	return e.CommandName
}

// MakeCommandArgs is dummy command arguments maker for testing
func (e *Executer) MakeCommandArgs() []string {
	return e.CommandArgs
}

// Parse is dummy parse method for testing
func (e *Executer) Parse(result string) (checkers.Status, string) {
	e.CommandResult = result
	return e.Status, e.ParseResult
}
<|endoftext|>


<commit_before>package fm_test

import (
	"go/ast"
	"testing"

	fm "github.com/enocom/fm/lib"
)

// TestGenerateReturnsSliceOfSpyDecls ensures the generator produces
// two declarations for a single interface with a single method:
// 1) a struct with fields to store the result of a function call, and
// 2) a spy implementation of the interface's single method.
func TestGenerateReturnsSliceOfSpyDecls(t *testing.T) {
	gen := &fm.SpyGenerator{
		Converter:   &fm.SpyStructConverter{},
		Implementer: &fm.SpyFuncImplementer{},
	}
	interfaceDecls := buildInterfaceAST()
	spyDecls := gen.Generate(interfaceDecls)

	want := 2
	got := len(spyDecls)

	if want != got {
		t.Errorf("want %v, got %v", want, got)
	}
}

func TestGenerateSkipsAnythingButInterfaceTypes(t *testing.T) {
	gen := &fm.SpyGenerator{
		Converter:   &fm.SpyStructConverter{},
		Implementer: &fm.SpyFuncImplementer{},
	}
	decls := buildASTWithoutInterfaces()
	spyDecls := gen.Generate(decls)

	want := 0
	got := len(spyDecls)

	if want != got {
		t.Errorf("want %v, got %v", want, got)
	}
}

func TestGenerateReturnsEmptySliceForNoInput(t *testing.T) {
	gen := &fm.SpyGenerator{}

	result := gen.Generate(make([]ast.Decl, 0))

	want := 0
	got := len(result)

	if want != got {
		t.Errorf("want %v, got %v", want, got)
	}
}

// buildTestAST generates as AST for the following interface:
//
// type Tester interface {
//     Test()
// }
func buildInterfaceAST() []ast.Decl {
	fields := make([]*ast.Field, 0)
	fields = append(fields, &ast.Field{
		Names: []*ast.Ident{ast.NewIdent("Test")},
		Type: &ast.FuncType{
			Params:  &ast.FieldList{List: nil},
			Results: &ast.FieldList{List: nil},
		},
	})

	decls := make([]ast.Decl, 0)
	decls = append(decls, &ast.GenDecl{
		Specs: []ast.Spec{
			&ast.TypeSpec{
				Name: ast.NewIdent("Tester"),
				Type: &ast.InterfaceType{
					Methods: &ast.FieldList{
						List: fields,
					},
				},
			},
		},
	})

	return decls
}

// buildASTWithoutInterfaces generates as AST of the following code:
//
// type Tester struct {}
func buildASTWithoutInterfaces() []ast.Decl {
	decls := make([]ast.Decl, 0)
	decls = append(decls, &ast.GenDecl{
		Specs: []ast.Spec{
			&ast.TypeSpec{
				Name: ast.NewIdent("Tester"),
				Type: &ast.StructType{
					Fields: &ast.FieldList{
						List: make([]*ast.Field, 0),
					},
				},
			},
		},
	})

	return decls
}
<commit_msg>Add tests covering all type casts in generator<commit_after>package fm_test

import (
	"go/ast"
	"go/token"
	"testing"

	fm "github.com/enocom/fm/lib"
)

// TestGenerateReturnsSliceOfSpyDecls ensures the generator produces
// two declarations for a single interface with a single method:
// 1) a struct with fields to store the result of a function call, and
// 2) a spy implementation of the interface's single method.
func TestGenerateReturnsSliceOfSpyDecls(t *testing.T) {
	gen := &fm.SpyGenerator{
		Converter:   &fm.SpyStructConverter{},
		Implementer: &fm.SpyFuncImplementer{},
	}
	interfaceDecls := buildInterfaceAST()
	spyDecls := gen.Generate(interfaceDecls)

	want := 2
	got := len(spyDecls)

	if want != got {
		t.Errorf("want %v, got %v", want, got)
	}
}

// e.g., ast.FuncDecl
func TestGenerateSkipsDeclsThatAreNotGenDecls(t *testing.T) {
	gen := &fm.SpyGenerator{
		Converter:   &fm.SpyStructConverter{},
		Implementer: &fm.SpyFuncImplementer{},
	}
	decls := buildFuncDeclAST()
	spyDecls := gen.Generate(decls)

	want := 0
	got := len(spyDecls)

	if want != got {
		t.Errorf("want %v, got %v", want, got)
	}
}

// e.g., ast.ValueSpec
func TestGenerateSkipsSpecsThatAreNotTypeSpecs(t *testing.T) {
	gen := &fm.SpyGenerator{
		Converter:   &fm.SpyStructConverter{},
		Implementer: &fm.SpyFuncImplementer{},
	}
	decls := buildValueSpecAST()
	spyDecls := gen.Generate(decls)

	want := 0
	got := len(spyDecls)

	if want != got {
		t.Errorf("want %v, got %v", want, got)
	}
}

// e.g., ast.StructType
func TestGenerateSkipsTypeSpecsThatAreNotInterfaceTypes(t *testing.T) {
	gen := &fm.SpyGenerator{
		Converter:   &fm.SpyStructConverter{},
		Implementer: &fm.SpyFuncImplementer{},
	}
	decls := buildStructAST()
	spyDecls := gen.Generate(decls)

	want := 0
	got := len(spyDecls)

	if want != got {
		t.Errorf("want %v, got %v", want, got)
	}
}

func TestGenerateReturnsEmptySliceForNoInput(t *testing.T) {
	gen := &fm.SpyGenerator{}

	result := gen.Generate(make([]ast.Decl, 0))

	want := 0
	got := len(result)

	if want != got {
		t.Errorf("want %v, got %v", want, got)
	}
}

// buildTestAST generates as AST of the following code:
//
// type Tester interface {
//     Test()
// }
func buildInterfaceAST() []ast.Decl {
	var fields []*ast.Field
	fields = append(fields, &ast.Field{
		Names: []*ast.Ident{ast.NewIdent("Test")},
		Type: &ast.FuncType{
			Params:  &ast.FieldList{List: nil},
			Results: &ast.FieldList{List: nil},
		},
	})

	var decls []ast.Decl
	decls = append(decls, &ast.GenDecl{
		Specs: []ast.Spec{
			&ast.TypeSpec{
				Name: ast.NewIdent("Tester"),
				Type: &ast.InterfaceType{
					Methods: &ast.FieldList{
						List: fields,
					},
				},
			},
		},
	})

	return decls
}

// buildStructAST generates as AST of the following code:
//
// type Tester struct {}
func buildStructAST() []ast.Decl {
	var decls []ast.Decl
	decls = append(decls, &ast.GenDecl{
		Specs: []ast.Spec{
			&ast.TypeSpec{
				Name: ast.NewIdent("Tester"),
				Type: &ast.StructType{
					Fields: &ast.FieldList{
						List: make([]*ast.Field, 0),
					},
				},
			},
		},
	})

	return decls
}

// buildStructAST generates an AST of the followign code:
//
// func foobar() {}
func buildFuncDeclAST() []ast.Decl {
	var decls []ast.Decl
	decls = append(decls, &ast.FuncDecl{
		Recv: nil,
		Name: ast.NewIdent("foobar"),
		Type: &ast.FuncType{
			Params: &ast.FieldList{
				List: make([]*ast.Field, 0),
			},
			Results: nil,
		},
		Body: &ast.BlockStmt{
			List: nil,
		},
	})

	return decls
}

// buildValueSpecAST generates an AST of the following code:
//
// var foobar string
func buildValueSpecAST() []ast.Decl {
	var decls []ast.Decl
	decls = append(decls, &ast.GenDecl{
		Tok: token.VAR,
		Specs: []ast.Spec{
			&ast.ValueSpec{
				Names: []*ast.Ident{
					ast.NewIdent("foobar"),
					ast.NewIdent("string"),
				},
			},
		},
	})

	return decls
}
<|endoftext|>


<commit_before>package cluster

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"time"

	"github.com/kballard/go-shellquote"
	"github.com/pkg/errors"
	"golang.org/x/crypto/scrypt"

	"github.com/lxc/lxd/lxd/config"
	"github.com/lxc/lxd/lxd/db"
)

// Config holds cluster-wide configuration values.
type Config struct {
	tx *db.ClusterTx // DB transaction the values in this config are bound to.
	m  config.Map    // Low-level map holding the config values.
}

// ConfigLoad loads a new Config object with the current cluster configuration
// values fetched from the database.
func ConfigLoad(tx *db.ClusterTx) (*Config, error) {
	// Load current raw values from the database, any error is fatal.
	values, err := tx.Config()
	if err != nil {
		return nil, fmt.Errorf("cannot fetch node config from database: %v", err)
	}

	m, err := config.SafeLoad(ConfigSchema, values)
	if err != nil {
		return nil, fmt.Errorf("failed to load node config: %v", err)
	}

	return &Config{tx: tx, m: m}, nil
}

// HTTPSAllowedHeaders returns the relevant CORS setting.
func (c *Config) HTTPSAllowedHeaders() string {
	return c.m.GetString("core.https_allowed_headers")
}

// HTTPSAllowedMethods returns the relevant CORS setting.
func (c *Config) HTTPSAllowedMethods() string {
	return c.m.GetString("core.https_allowed_methods")
}

// HTTPSAllowedOrigin returns the relevant CORS setting.
func (c *Config) HTTPSAllowedOrigin() string {
	return c.m.GetString("core.https_allowed_origin")
}

// HTTPSAllowedCredentials returns the relevant CORS setting.
func (c *Config) HTTPSAllowedCredentials() bool {
	return c.m.GetBool("core.https_allowed_credentials")
}

// TrustPassword returns the LXD trust password for authenticating clients.
func (c *Config) TrustPassword() string {
	return c.m.GetString("core.trust_password")
}

// TrustCACertificates returns whether client certificates are checked
// against a CA.
func (c *Config) TrustCACertificates() bool {
	return c.m.GetBool("core.trust_ca_certificates")
}

// CandidServer returns all the Candid settings needed to connect to a server.
func (c *Config) CandidServer() (string, string, int64, string) {
	return c.m.GetString("candid.api.url"),
		c.m.GetString("candid.api.key"),
		c.m.GetInt64("candid.expiry"),
		c.m.GetString("candid.domains")
}

// RBACServer returns all the Candid settings needed to connect to a server.
func (c *Config) RBACServer() (string, string, int64, string, string, string, string) {
	return c.m.GetString("rbac.api.url"),
		c.m.GetString("rbac.api.key"),
		c.m.GetInt64("rbac.expiry"),
		c.m.GetString("rbac.agent.url"),
		c.m.GetString("rbac.agent.username"),
		c.m.GetString("rbac.agent.private_key"),
		c.m.GetString("rbac.agent.public_key")
}

// AutoUpdateInterval returns the configured images auto update interval.
func (c *Config) AutoUpdateInterval() time.Duration {
	n := c.m.GetInt64("images.auto_update_interval")
	return time.Duration(n) * time.Hour
}

// RemoteCacheExpiry returns the configured expiration value for remote images
// expiration.
func (c *Config) RemoteCacheExpiry() int64 {
	return c.m.GetInt64("images.remote_cache_expiry")
}

// ProxyHTTPS returns the configured HTTPS proxy, if any.
func (c *Config) ProxyHTTPS() string {
	return c.m.GetString("core.proxy_https")
}

// ProxyHTTP returns the configured HTTP proxy, if any.
func (c *Config) ProxyHTTP() string {
	return c.m.GetString("core.proxy_http")
}

// ProxyIgnoreHosts returns the configured ignore-hosts proxy setting, if any.
func (c *Config) ProxyIgnoreHosts() string {
	return c.m.GetString("core.proxy_ignore_hosts")
}

// MAASController the configured MAAS url and key, if any.
func (c *Config) MAASController() (string, string) {
	url := c.m.GetString("maas.api.url")
	key := c.m.GetString("maas.api.key")
	return url, key
}

// OfflineThreshold returns the configured heartbeat threshold, i.e. the
// number of seconds before after which an unresponsive node is considered
// offline..
func (c *Config) OfflineThreshold() time.Duration {
	n := c.m.GetInt64("cluster.offline_threshold")
	return time.Duration(n) * time.Second
}

// ImagesMinimalReplica returns the numbers of nodes for cluster images replication
func (c *Config) ImagesMinimalReplica() int64 {
	return c.m.GetInt64("cluster.images_minimal_replica")
}

// MaxVoters returns the maximum number of members in a cluster that will be
// assigned the voter role.
func (c *Config) MaxVoters() int64 {
	return c.m.GetInt64("cluster.max_voters")
}

// MaxStandBy returns the maximum number of standby members in a cluster that
// will be assigned the stand-by role.
func (c *Config) MaxStandBy() int64 {
	return c.m.GetInt64("cluster.max_standby")
}

// Dump current configuration keys and their values. Keys with values matching
// their defaults are omitted.
func (c *Config) Dump() map[string]interface{} {
	return c.m.Dump()
}

// Replace the current configuration with the given values.
//
// Return what has actually changed.
func (c *Config) Replace(values map[string]interface{}) (map[string]string, error) {
	return c.update(values)
}

// Patch changes only the configuration keys in the given map.
//
// Return what has actually changed.
func (c *Config) Patch(patch map[string]interface{}) (map[string]string, error) {
	values := c.Dump() // Use current values as defaults
	for name, value := range patch {
		values[name] = value
	}
	return c.update(values)
}

func (c *Config) update(values map[string]interface{}) (map[string]string, error) {
	changed, err := c.m.Change(values)
	if err != nil {
		return nil, err
	}

	err = c.tx.UpdateConfig(changed)
	if err != nil {
		return nil, errors.Wrap(err, "cannot persist configuration changes: %v")
	}

	return changed, nil
}

// ConfigGetString is a convenience for loading the cluster configuration and
// returning the value of a particular key.
//
// It's a deprecated API meant to be used by call sites that are not
// interacting with the database in a transactional way.
func ConfigGetString(cluster *db.Cluster, key string) (string, error) {
	config, err := configGet(cluster)
	if err != nil {
		return "", err
	}
	return config.m.GetString(key), nil
}

// ConfigGetBool is a convenience for loading the cluster configuration and
// returning the value of a particular boolean key.
//
// It's a deprecated API meant to be used by call sites that are not
// interacting with the database in a transactional way.
func ConfigGetBool(cluster *db.Cluster, key string) (bool, error) {
	config, err := configGet(cluster)
	if err != nil {
		return false, err
	}
	return config.m.GetBool(key), nil
}

// ConfigGetInt64 is a convenience for loading the cluster configuration and
// returning the value of a particular key.
//
// It's a deprecated API meant to be used by call sites that are not
// interacting with the database in a transactional way.
func ConfigGetInt64(cluster *db.Cluster, key string) (int64, error) {
	config, err := configGet(cluster)
	if err != nil {
		return 0, err
	}
	return config.m.GetInt64(key), nil
}

func configGet(cluster *db.Cluster) (*Config, error) {
	var config *Config
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		config, err = ConfigLoad(tx)
		return err
	})
	return config, err
}

// ConfigSchema defines available server configuration keys.
var ConfigSchema = config.Schema{
	"backups.compression_algorithm":  {Default: "gzip", Validator: validateCompression},
	"cluster.offline_threshold":      {Type: config.Int64, Default: offlineThresholdDefault(), Validator: offlineThresholdValidator},
	"cluster.images_minimal_replica": {Type: config.Int64, Default: "3", Validator: imageMinimalReplicaValidator},
	"cluster.max_voters":             {Type: config.Int64, Default: "3", Validator: maxVotersValidator},
	"cluster.max_standby":            {Type: config.Int64, Default: "2", Validator: maxStandByValidator},
	"core.https_allowed_headers":     {},
	"core.https_allowed_methods":     {},
	"core.https_allowed_origin":      {},
	"core.https_allowed_credentials": {Type: config.Bool},
	"core.proxy_http":                {},
	"core.proxy_https":               {},
	"core.proxy_ignore_hosts":        {},
	"core.trust_password":            {Hidden: true, Setter: passwordSetter},
	"core.trust_ca_certificates":     {Type: config.Bool},
	"candid.api.key":                 {},
	"candid.api.url":                 {},
	"candid.domains":                 {},
	"candid.expiry":                  {Type: config.Int64, Default: "3600"},
	"images.auto_update_cached":      {Type: config.Bool, Default: "true"},
	"images.auto_update_interval":    {Type: config.Int64, Default: "6"},
	"images.compression_algorithm":   {Default: "gzip", Validator: validateCompression},
	"images.remote_cache_expiry":     {Type: config.Int64, Default: "10"},
	"maas.api.key":                   {},
	"maas.api.url":                   {},
	"rbac.agent.url":                 {},
	"rbac.agent.username":            {},
	"rbac.agent.private_key":         {},
	"rbac.agent.public_key":          {},
	"rbac.api.expiry":                {Type: config.Int64, Default: "3600"},
	"rbac.api.key":                   {},
	"rbac.api.url":                   {},
	"rbac.expiry":                    {Type: config.Int64, Default: "3600"},

	// Keys deprecated since the implementation of the storage api.
	"storage.lvm_fstype":           {Setter: deprecatedStorage, Default: "ext4"},
	"storage.lvm_mount_options":    {Setter: deprecatedStorage, Default: "discard"},
	"storage.lvm_thinpool_name":    {Setter: deprecatedStorage, Default: "LXDThinPool"},
	"storage.lvm_vg_name":          {Setter: deprecatedStorage},
	"storage.lvm_volume_size":      {Setter: deprecatedStorage, Default: "10GiB"},
	"storage.zfs_pool_name":        {Setter: deprecatedStorage},
	"storage.zfs_remove_snapshots": {Setter: deprecatedStorage, Type: config.Bool},
	"storage.zfs_use_refquota":     {Setter: deprecatedStorage, Type: config.Bool},

	// OVN networking global keys.
	"network.ovn.integration_bridge":    {Default: "br-int"},
	"network.ovn.northbound_connection": {Default: "unix:/var/run/ovn/ovnnb_db.sock"},
}

func offlineThresholdDefault() string {
	return strconv.Itoa(db.DefaultOfflineThreshold)
}

func offlineThresholdValidator(value string) error {
	// Ensure that the given value is greater than the heartbeat interval,
	// which is the lower bound granularity of the offline check.
	threshold, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("Offline threshold is not a number")
	}

	if threshold <= heartbeatInterval {
		return fmt.Errorf("Value must be greater than '%d'", heartbeatInterval)
	}

	return nil
}

func imageMinimalReplicaValidator(value string) error {
	count, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("Minimal image replica count is not a number")
	}

	if count < 1 && count != -1 {
		return fmt.Errorf("Invalid value for image replica count")
	}

	return nil
}

func maxVotersValidator(value string) error {
	n, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("Value is not a number")
	}

	if n < 3 || n%2 != 1 {
		return fmt.Errorf("Value must be an odd number equal to or higher than 3")
	}

	return nil
}

func maxStandByValidator(value string) error {
	n, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("Value is not a number")
	}

	if n < 0 || n > 5 {
		return fmt.Errorf("Value must be between 0 and 5")
	}

	return nil
}

func passwordSetter(value string) (string, error) {
	// Nothing to do on unset
	if value == "" {
		return value, nil
	}

	// Hash the password
	buf := make([]byte, 32)
	_, err := io.ReadFull(rand.Reader, buf)
	if err != nil {
		return "", err
	}

	hash, err := scrypt.Key([]byte(value), buf, 1<<14, 8, 1, 64)
	if err != nil {
		return "", err
	}

	buf = append(buf, hash...)
	value = hex.EncodeToString(buf)

	return value, nil
}

func validateCompression(value string) error {
	if value == "none" {
		return nil
	}

	// Going to look up tar2sqfs executable binary
	if value == "squashfs" {
		value = "tar2sqfs"
	}

	// Parse the command.
	fields, err := shellquote.Split(value)
	if err != nil {
		return err
	}

	_, err = exec.LookPath(fields[0])
	return err
}

func deprecatedStorage(value string) (string, error) {
	if value == "" {
		return "", nil
	}
	return "", fmt.Errorf("deprecated: use storage pool configuration")
}
<commit_msg>lxd/cluster: Update compression validation<commit_after>package cluster

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"strconv"
	"time"

	"github.com/pkg/errors"
	"golang.org/x/crypto/scrypt"

	"github.com/lxc/lxd/lxd/config"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/shared/validate"
)

// Config holds cluster-wide configuration values.
type Config struct {
	tx *db.ClusterTx // DB transaction the values in this config are bound to.
	m  config.Map    // Low-level map holding the config values.
}

// ConfigLoad loads a new Config object with the current cluster configuration
// values fetched from the database.
func ConfigLoad(tx *db.ClusterTx) (*Config, error) {
	// Load current raw values from the database, any error is fatal.
	values, err := tx.Config()
	if err != nil {
		return nil, fmt.Errorf("cannot fetch node config from database: %v", err)
	}

	m, err := config.SafeLoad(ConfigSchema, values)
	if err != nil {
		return nil, fmt.Errorf("failed to load node config: %v", err)
	}

	return &Config{tx: tx, m: m}, nil
}

// HTTPSAllowedHeaders returns the relevant CORS setting.
func (c *Config) HTTPSAllowedHeaders() string {
	return c.m.GetString("core.https_allowed_headers")
}

// HTTPSAllowedMethods returns the relevant CORS setting.
func (c *Config) HTTPSAllowedMethods() string {
	return c.m.GetString("core.https_allowed_methods")
}

// HTTPSAllowedOrigin returns the relevant CORS setting.
func (c *Config) HTTPSAllowedOrigin() string {
	return c.m.GetString("core.https_allowed_origin")
}

// HTTPSAllowedCredentials returns the relevant CORS setting.
func (c *Config) HTTPSAllowedCredentials() bool {
	return c.m.GetBool("core.https_allowed_credentials")
}

// TrustPassword returns the LXD trust password for authenticating clients.
func (c *Config) TrustPassword() string {
	return c.m.GetString("core.trust_password")
}

// TrustCACertificates returns whether client certificates are checked
// against a CA.
func (c *Config) TrustCACertificates() bool {
	return c.m.GetBool("core.trust_ca_certificates")
}

// CandidServer returns all the Candid settings needed to connect to a server.
func (c *Config) CandidServer() (string, string, int64, string) {
	return c.m.GetString("candid.api.url"),
		c.m.GetString("candid.api.key"),
		c.m.GetInt64("candid.expiry"),
		c.m.GetString("candid.domains")
}

// RBACServer returns all the Candid settings needed to connect to a server.
func (c *Config) RBACServer() (string, string, int64, string, string, string, string) {
	return c.m.GetString("rbac.api.url"),
		c.m.GetString("rbac.api.key"),
		c.m.GetInt64("rbac.expiry"),
		c.m.GetString("rbac.agent.url"),
		c.m.GetString("rbac.agent.username"),
		c.m.GetString("rbac.agent.private_key"),
		c.m.GetString("rbac.agent.public_key")
}

// AutoUpdateInterval returns the configured images auto update interval.
func (c *Config) AutoUpdateInterval() time.Duration {
	n := c.m.GetInt64("images.auto_update_interval")
	return time.Duration(n) * time.Hour
}

// RemoteCacheExpiry returns the configured expiration value for remote images
// expiration.
func (c *Config) RemoteCacheExpiry() int64 {
	return c.m.GetInt64("images.remote_cache_expiry")
}

// ProxyHTTPS returns the configured HTTPS proxy, if any.
func (c *Config) ProxyHTTPS() string {
	return c.m.GetString("core.proxy_https")
}

// ProxyHTTP returns the configured HTTP proxy, if any.
func (c *Config) ProxyHTTP() string {
	return c.m.GetString("core.proxy_http")
}

// ProxyIgnoreHosts returns the configured ignore-hosts proxy setting, if any.
func (c *Config) ProxyIgnoreHosts() string {
	return c.m.GetString("core.proxy_ignore_hosts")
}

// MAASController the configured MAAS url and key, if any.
func (c *Config) MAASController() (string, string) {
	url := c.m.GetString("maas.api.url")
	key := c.m.GetString("maas.api.key")
	return url, key
}

// OfflineThreshold returns the configured heartbeat threshold, i.e. the
// number of seconds before after which an unresponsive node is considered
// offline..
func (c *Config) OfflineThreshold() time.Duration {
	n := c.m.GetInt64("cluster.offline_threshold")
	return time.Duration(n) * time.Second
}

// ImagesMinimalReplica returns the numbers of nodes for cluster images replication
func (c *Config) ImagesMinimalReplica() int64 {
	return c.m.GetInt64("cluster.images_minimal_replica")
}

// MaxVoters returns the maximum number of members in a cluster that will be
// assigned the voter role.
func (c *Config) MaxVoters() int64 {
	return c.m.GetInt64("cluster.max_voters")
}

// MaxStandBy returns the maximum number of standby members in a cluster that
// will be assigned the stand-by role.
func (c *Config) MaxStandBy() int64 {
	return c.m.GetInt64("cluster.max_standby")
}

// Dump current configuration keys and their values. Keys with values matching
// their defaults are omitted.
func (c *Config) Dump() map[string]interface{} {
	return c.m.Dump()
}

// Replace the current configuration with the given values.
//
// Return what has actually changed.
func (c *Config) Replace(values map[string]interface{}) (map[string]string, error) {
	return c.update(values)
}

// Patch changes only the configuration keys in the given map.
//
// Return what has actually changed.
func (c *Config) Patch(patch map[string]interface{}) (map[string]string, error) {
	values := c.Dump() // Use current values as defaults
	for name, value := range patch {
		values[name] = value
	}
	return c.update(values)
}

func (c *Config) update(values map[string]interface{}) (map[string]string, error) {
	changed, err := c.m.Change(values)
	if err != nil {
		return nil, err
	}

	err = c.tx.UpdateConfig(changed)
	if err != nil {
		return nil, errors.Wrap(err, "cannot persist configuration changes: %v")
	}

	return changed, nil
}

// ConfigGetString is a convenience for loading the cluster configuration and
// returning the value of a particular key.
//
// It's a deprecated API meant to be used by call sites that are not
// interacting with the database in a transactional way.
func ConfigGetString(cluster *db.Cluster, key string) (string, error) {
	config, err := configGet(cluster)
	if err != nil {
		return "", err
	}
	return config.m.GetString(key), nil
}

// ConfigGetBool is a convenience for loading the cluster configuration and
// returning the value of a particular boolean key.
//
// It's a deprecated API meant to be used by call sites that are not
// interacting with the database in a transactional way.
func ConfigGetBool(cluster *db.Cluster, key string) (bool, error) {
	config, err := configGet(cluster)
	if err != nil {
		return false, err
	}
	return config.m.GetBool(key), nil
}

// ConfigGetInt64 is a convenience for loading the cluster configuration and
// returning the value of a particular key.
//
// It's a deprecated API meant to be used by call sites that are not
// interacting with the database in a transactional way.
func ConfigGetInt64(cluster *db.Cluster, key string) (int64, error) {
	config, err := configGet(cluster)
	if err != nil {
		return 0, err
	}
	return config.m.GetInt64(key), nil
}

func configGet(cluster *db.Cluster) (*Config, error) {
	var config *Config
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		config, err = ConfigLoad(tx)
		return err
	})
	return config, err
}

// ConfigSchema defines available server configuration keys.
var ConfigSchema = config.Schema{
	"backups.compression_algorithm":  {Default: "gzip", Validator: validate.IsCompressionAlgorithm},
	"cluster.offline_threshold":      {Type: config.Int64, Default: offlineThresholdDefault(), Validator: offlineThresholdValidator},
	"cluster.images_minimal_replica": {Type: config.Int64, Default: "3", Validator: imageMinimalReplicaValidator},
	"cluster.max_voters":             {Type: config.Int64, Default: "3", Validator: maxVotersValidator},
	"cluster.max_standby":            {Type: config.Int64, Default: "2", Validator: maxStandByValidator},
	"core.https_allowed_headers":     {},
	"core.https_allowed_methods":     {},
	"core.https_allowed_origin":      {},
	"core.https_allowed_credentials": {Type: config.Bool},
	"core.proxy_http":                {},
	"core.proxy_https":               {},
	"core.proxy_ignore_hosts":        {},
	"core.trust_password":            {Hidden: true, Setter: passwordSetter},
	"core.trust_ca_certificates":     {Type: config.Bool},
	"candid.api.key":                 {},
	"candid.api.url":                 {},
	"candid.domains":                 {},
	"candid.expiry":                  {Type: config.Int64, Default: "3600"},
	"images.auto_update_cached":      {Type: config.Bool, Default: "true"},
	"images.auto_update_interval":    {Type: config.Int64, Default: "6"},
	"images.compression_algorithm":   {Default: "gzip", Validator: validate.IsCompressionAlgorithm},
	"images.remote_cache_expiry":     {Type: config.Int64, Default: "10"},
	"maas.api.key":                   {},
	"maas.api.url":                   {},
	"rbac.agent.url":                 {},
	"rbac.agent.username":            {},
	"rbac.agent.private_key":         {},
	"rbac.agent.public_key":          {},
	"rbac.api.expiry":                {Type: config.Int64, Default: "3600"},
	"rbac.api.key":                   {},
	"rbac.api.url":                   {},
	"rbac.expiry":                    {Type: config.Int64, Default: "3600"},

	// Keys deprecated since the implementation of the storage api.
	"storage.lvm_fstype":           {Setter: deprecatedStorage, Default: "ext4"},
	"storage.lvm_mount_options":    {Setter: deprecatedStorage, Default: "discard"},
	"storage.lvm_thinpool_name":    {Setter: deprecatedStorage, Default: "LXDThinPool"},
	"storage.lvm_vg_name":          {Setter: deprecatedStorage},
	"storage.lvm_volume_size":      {Setter: deprecatedStorage, Default: "10GiB"},
	"storage.zfs_pool_name":        {Setter: deprecatedStorage},
	"storage.zfs_remove_snapshots": {Setter: deprecatedStorage, Type: config.Bool},
	"storage.zfs_use_refquota":     {Setter: deprecatedStorage, Type: config.Bool},

	// OVN networking global keys.
	"network.ovn.integration_bridge":    {Default: "br-int"},
	"network.ovn.northbound_connection": {Default: "unix:/var/run/ovn/ovnnb_db.sock"},
}

func offlineThresholdDefault() string {
	return strconv.Itoa(db.DefaultOfflineThreshold)
}

func offlineThresholdValidator(value string) error {
	// Ensure that the given value is greater than the heartbeat interval,
	// which is the lower bound granularity of the offline check.
	threshold, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("Offline threshold is not a number")
	}

	if threshold <= heartbeatInterval {
		return fmt.Errorf("Value must be greater than '%d'", heartbeatInterval)
	}

	return nil
}

func imageMinimalReplicaValidator(value string) error {
	count, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("Minimal image replica count is not a number")
	}

	if count < 1 && count != -1 {
		return fmt.Errorf("Invalid value for image replica count")
	}

	return nil
}

func maxVotersValidator(value string) error {
	n, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("Value is not a number")
	}

	if n < 3 || n%2 != 1 {
		return fmt.Errorf("Value must be an odd number equal to or higher than 3")
	}

	return nil
}

func maxStandByValidator(value string) error {
	n, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("Value is not a number")
	}

	if n < 0 || n > 5 {
		return fmt.Errorf("Value must be between 0 and 5")
	}

	return nil
}

func passwordSetter(value string) (string, error) {
	// Nothing to do on unset
	if value == "" {
		return value, nil
	}

	// Hash the password
	buf := make([]byte, 32)
	_, err := io.ReadFull(rand.Reader, buf)
	if err != nil {
		return "", err
	}

	hash, err := scrypt.Key([]byte(value), buf, 1<<14, 8, 1, 64)
	if err != nil {
		return "", err
	}

	buf = append(buf, hash...)
	value = hex.EncodeToString(buf)

	return value, nil
}

func deprecatedStorage(value string) (string, error) {
	if value == "" {
		return "", nil
	}
	return "", fmt.Errorf("deprecated: use storage pool configuration")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"source.datanerd.us/site-engineering/go_nagios"
)

const (
	API_VERSION = "v1.10"
)

// A struct representing CLI opts that will be passed at runtime
type CliOpts struct {
	BaseUrl        string
	CritDataSpace  int
	WarnDataSpace  int
	CritMetaSpace  int
	WarnMetaSpace  int
	ImageId        string
	GhostsStatus   int
}

// Information describing the status of a Docker host
type DockerInfo struct {
	Containers     float64
	DriverStatus   [][]string
	DataSpaceUsed  float64
	DataSpaceTotal float64
	MetaSpaceUsed  float64
	MetaSpaceTotal float64
	ImageIsRunning bool
	GhostCount     int
}

// Used internally to build lists of checks to run
type checkArgs struct {
	tag                string
	value              string
	healthy            bool
	appendErrorMessage string
	statusVal          nagios.NagiosStatusVal
}

// Describes one container
type Container struct {
	Image  string
	Status string
}

// An interface to request things from the Web
type HttpResponseFetcher interface {
	Fetch(url string) ([]byte, error)
}

type Fetcher struct{}

// Properly format a Float64 as a string
func float64String(num float64) string {
	return strconv.FormatFloat(num, 'f', 0, 64)
}

// Return a float from a Docker info string for megabytes
func megabytesFloat64(value string) (float64, error) {
	numberStr := strings.Fields(value)[0]
	number, err := strconv.ParseFloat(numberStr, 64)

	if err != nil {
		return 0.00, err
	}

	return number, nil
}

// Look through a list of driveStatus slices and find the one that matches
func findDriverStatus(key string, driverStatus [][]string) string {
	for _, entry := range driverStatus {
		if entry[0] == key {
			return entry[1]
		}
	}

	return ""
}

// Connect to a Docker URL and return the contents as a []byte
func (Fetcher) Fetch(url string) ([]byte, error) {
	response, err := http.Get(url)

	if err != nil {
		return nil, err
	}

	defer response.Body.Close()

	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	return contents, nil
}

// Parses JSON and populates a DockerInfo
func populateInfo(contents []byte, info *DockerInfo) error {
	err := json.Unmarshal(contents, info)
	if err != nil {
		return err
	}

	fields := map[string]*float64{
		"Data Space Used":      &info.DataSpaceUsed,
		"Data Space Total":     &info.DataSpaceTotal,
		"Metadata Space Used":  &info.MetaSpaceUsed,
		"Metadata Space Total": &info.MetaSpaceTotal,
	}

	for key, val := range fields {
		entry := findDriverStatus(key, info.DriverStatus)
		if entry == "" {
			return errors.New("Error parsing response from API! Can't find key: " + key)
		}
		*val, err = megabytesFloat64(findDriverStatus(key, info.DriverStatus))
		if err != nil {
			return errors.New("Error parsing response from API! " + err.Error())
		}
	}

	return nil
}

// checkRunningContainers looks to see if a container is currently running from a given
// Image Id.
func checkRunningContainers(contents []byte, opts *CliOpts) (bool, int, error) {
	var containers []Container

	err := json.Unmarshal(contents, &containers)
	if err != nil {
		return false, 0, err
	}

	isRunning := false
	ghostCount := 0
	for _, container := range containers {
		if strings.HasPrefix(container.Image, opts.ImageId+":") && strings.HasPrefix(container.Status, "Up") {
			isRunning = true
		} else if strings.Contains(container.Status, "Ghost") {
			ghostCount += 1
		}
	}
	return isRunning, ghostCount, nil
}

// fetchInfo retrieves JSON from a Docker host and fills in a DockerInfo
func fetchInfo(fetcher HttpResponseFetcher, opts CliOpts, info *DockerInfo) error {
	var wg sync.WaitGroup
	wg.Add(2)

	var err, err2 error
	var imageFound bool
	var ghostCount int

	go func() {
		var infoResult []byte
		infoResult, err = fetcher.Fetch(opts.BaseUrl + "/" + API_VERSION + "/info")
		if err == nil {
			err = populateInfo(infoResult, info)
		}
		wg.Done()
	}()

	go func() {
		var containersResult []byte
		containersResult, err2 = fetcher.Fetch(opts.BaseUrl + "/" + API_VERSION + "/containers/json")
		if err2 == nil {
			imageFound, ghostCount, err2 = checkRunningContainers(containersResult, &opts)
		}
		wg.Done()
	}()

	wg.Wait()

	if err != nil || err2 != nil {
		return err
	}

	info.ImageIsRunning = imageFound
	info.GhostCount = ghostCount

	return nil
}

// defineChecks returns a list of checks we should run based on CLI flags
func defineChecks(info *DockerInfo, opts *CliOpts) []checkArgs {
	checks := []checkArgs{
		checkArgs{"Meta Space Used",
			float64String(info.MetaSpaceUsed / info.MetaSpaceTotal * 100),
			info.MetaSpaceUsed/info.MetaSpaceTotal*100 < float64(opts.CritMetaSpace),
			"%",
			nagios.NAGIOS_CRITICAL,
		},
		checkArgs{"Data Space Used",
			float64String(info.DataSpaceUsed / info.DataSpaceTotal * 100),
			info.DataSpaceUsed/info.DataSpaceTotal*100 < float64(opts.CritDataSpace),
			"%",
			nagios.NAGIOS_CRITICAL,
		},
		checkArgs{"Meta Space Used",
			float64String(info.MetaSpaceUsed / info.MetaSpaceTotal * 100),
			info.MetaSpaceUsed/info.MetaSpaceTotal*100 < float64(opts.WarnMetaSpace),
			"%",
			nagios.NAGIOS_WARNING,
		},
		checkArgs{"Data Space Used",
			float64String(info.DataSpaceUsed / info.DataSpaceTotal * 100),
			info.DataSpaceUsed/info.DataSpaceTotal*100 < float64(opts.WarnDataSpace),
			"%",
			nagios.NAGIOS_WARNING,
		},
		checkArgs{"Ghost Containers",
			strconv.Itoa(info.GhostCount),
			info.GhostCount == 0,
			"",
			nagios.NagiosStatusVal(opts.GhostsStatus),
		},
	}

	if opts.ImageId != "" {
		checks = append(checks,
			checkArgs{"Running Image",
				opts.ImageId,
				info.ImageIsRunning,
				" is not running!",
				nagios.NAGIOS_CRITICAL,
			},
		)
	}

	return checks
}

// Runs a set of checkes and returns an array of statuses
func mapAlertStatuses(info *DockerInfo, opts *CliOpts) []*nagios.NagiosStatus {
	var statuses []*nagios.NagiosStatus

	var check = func(args checkArgs) *nagios.NagiosStatus {
		if !args.healthy {
			return &nagios.NagiosStatus{args.tag + ": " + args.value + args.appendErrorMessage, args.statusVal}
		}
		return nil
	}

	checks := defineChecks(info, opts)

	for _, entry := range checks {
		result := check(entry)
		if result != nil {
			statuses = append(statuses, check(entry))
		}
	}

	return statuses
}

// parseCommandLine parses the flags passed on the CLI
func parseCommandLine() *CliOpts {
	var opts CliOpts

	flag.StringVar(&opts.BaseUrl, "base-url", "http://chi-staging-pool-1:4243/", "The Base URL for the Docker server")
	flag.IntVar(&opts.WarnMetaSpace, "warn-meta-space", 100, "Warning threshold for Metadata Space")
	flag.IntVar(&opts.WarnDataSpace, "warn-data-space", 100, "Warning threshold for Data Space")
	flag.IntVar(&opts.CritMetaSpace, "crit-meta-space", 100, "Critical threshold for Metadata Space")
	flag.IntVar(&opts.CritDataSpace, "crit-data-space", 100, "Critical threshold for Data Space")
	flag.StringVar(&opts.ImageId, "image-id", "", "An image ID that must be running on the Docker server")
	flag.IntVar(&opts.GhostsStatus, "ghosts-status", 1, "If ghosts are present, treat as this status")

	flag.Parse()

	return &opts
}

func main() {
	opts := parseCommandLine()

	var fetcher Fetcher
	var info DockerInfo

	err := fetchInfo(fetcher, *opts, &info)
	if err != nil {
		nagios.Critical(err)
	}

	statuses := mapAlertStatuses(&info, opts)
	baseStatus := nagios.NagiosStatus{float64String(info.Containers) + " containers", 0}

	baseStatus.Aggregate(statuses)
	nagios.ExitWithStatus(&baseStatus)
}
<commit_msg>Fix bug where the wrong error could be returned.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"source.datanerd.us/site-engineering/go_nagios"
)

const (
	API_VERSION = "v1.10"
)

// A struct representing CLI opts that will be passed at runtime
type CliOpts struct {
	BaseUrl        string
	CritDataSpace  int
	WarnDataSpace  int
	CritMetaSpace  int
	WarnMetaSpace  int
	ImageId        string
	GhostsStatus   int
}

// Information describing the status of a Docker host
type DockerInfo struct {
	Containers     float64
	DriverStatus   [][]string
	DataSpaceUsed  float64
	DataSpaceTotal float64
	MetaSpaceUsed  float64
	MetaSpaceTotal float64
	ImageIsRunning bool
	GhostCount     int
}

// Used internally to build lists of checks to run
type checkArgs struct {
	tag                string
	value              string
	healthy            bool
	appendErrorMessage string
	statusVal          nagios.NagiosStatusVal
}

// Describes one container
type Container struct {
	Image  string
	Status string
}

// An interface to request things from the Web
type HttpResponseFetcher interface {
	Fetch(url string) ([]byte, error)
}

type Fetcher struct{}

// Properly format a Float64 as a string
func float64String(num float64) string {
	return strconv.FormatFloat(num, 'f', 0, 64)
}

// Return a float from a Docker info string for megabytes
func megabytesFloat64(value string) (float64, error) {
	numberStr := strings.Fields(value)[0]
	number, err := strconv.ParseFloat(numberStr, 64)

	if err != nil {
		return 0.00, err
	}

	return number, nil
}

// Look through a list of driveStatus slices and find the one that matches
func findDriverStatus(key string, driverStatus [][]string) string {
	for _, entry := range driverStatus {
		if entry[0] == key {
			return entry[1]
		}
	}

	return ""
}

// Connect to a Docker URL and return the contents as a []byte
func (Fetcher) Fetch(url string) ([]byte, error) {
	response, err := http.Get(url)

	if err != nil {
		return nil, err
	}

	defer response.Body.Close()

	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	return contents, nil
}

// Parses JSON and populates a DockerInfo
func populateInfo(contents []byte, info *DockerInfo) error {
	err := json.Unmarshal(contents, info)
	if err != nil {
		return err
	}

	fields := map[string]*float64{
		"Data Space Used":      &info.DataSpaceUsed,
		"Data Space Total":     &info.DataSpaceTotal,
		"Metadata Space Used":  &info.MetaSpaceUsed,
		"Metadata Space Total": &info.MetaSpaceTotal,
	}

	for key, val := range fields {
		entry := findDriverStatus(key, info.DriverStatus)
		if entry == "" {
			return errors.New("Error parsing response from API! Can't find key: " + key)
		}
		*val, err = megabytesFloat64(findDriverStatus(key, info.DriverStatus))
		if err != nil {
			return errors.New("Error parsing response from API! " + err.Error())
		}
	}

	return nil
}

// checkRunningContainers looks to see if a container is currently running from a given
// Image Id.
func checkRunningContainers(contents []byte, opts *CliOpts) (bool, int, error) {
	var containers []Container

	err := json.Unmarshal(contents, &containers)
	if err != nil {
		return false, 0, err
	}

	isRunning := false
	ghostCount := 0
	for _, container := range containers {
		if strings.HasPrefix(container.Image, opts.ImageId+":") && strings.HasPrefix(container.Status, "Up") {
			isRunning = true
		} else if strings.Contains(container.Status, "Ghost") {
			ghostCount += 1
		}
	}
	return isRunning, ghostCount, nil
}

// fetchInfo retrieves JSON from a Docker host and fills in a DockerInfo
func fetchInfo(fetcher HttpResponseFetcher, opts CliOpts, info *DockerInfo) error {
	var wg sync.WaitGroup
	wg.Add(2)

	var err, err2 error
	var imageFound bool
	var ghostCount int

	go func() {
		var infoResult []byte
		infoResult, err = fetcher.Fetch(opts.BaseUrl + "/" + API_VERSION + "/info")
		if err == nil {
			err = populateInfo(infoResult, info)
		}
		wg.Done()
	}()

	go func() {
		var containersResult []byte
		containersResult, err2 = fetcher.Fetch(opts.BaseUrl + "/" + API_VERSION + "/containers/json")
		if err2 == nil {
			imageFound, ghostCount, err2 = checkRunningContainers(containersResult, &opts)
		}
		wg.Done()
	}()

	wg.Wait()

	if err != nil {
		return err
	}

	if err2 != nil {
		return err2
	}

	info.ImageIsRunning = imageFound
	info.GhostCount = ghostCount

	return nil
}

// defineChecks returns a list of checks we should run based on CLI flags
func defineChecks(info *DockerInfo, opts *CliOpts) []checkArgs {
	checks := []checkArgs{
		checkArgs{"Meta Space Used",
			float64String(info.MetaSpaceUsed / info.MetaSpaceTotal * 100),
			info.MetaSpaceUsed/info.MetaSpaceTotal*100 < float64(opts.CritMetaSpace),
			"%",
			nagios.NAGIOS_CRITICAL,
		},
		checkArgs{"Data Space Used",
			float64String(info.DataSpaceUsed / info.DataSpaceTotal * 100),
			info.DataSpaceUsed/info.DataSpaceTotal*100 < float64(opts.CritDataSpace),
			"%",
			nagios.NAGIOS_CRITICAL,
		},
		checkArgs{"Meta Space Used",
			float64String(info.MetaSpaceUsed / info.MetaSpaceTotal * 100),
			info.MetaSpaceUsed/info.MetaSpaceTotal*100 < float64(opts.WarnMetaSpace),
			"%",
			nagios.NAGIOS_WARNING,
		},
		checkArgs{"Data Space Used",
			float64String(info.DataSpaceUsed / info.DataSpaceTotal * 100),
			info.DataSpaceUsed/info.DataSpaceTotal*100 < float64(opts.WarnDataSpace),
			"%",
			nagios.NAGIOS_WARNING,
		},
		checkArgs{"Ghost Containers",
			strconv.Itoa(info.GhostCount),
			info.GhostCount == 0,
			"",
			nagios.NagiosStatusVal(opts.GhostsStatus),
		},
	}

	if opts.ImageId != "" {
		checks = append(checks,
			checkArgs{"Running Image",
				opts.ImageId,
				info.ImageIsRunning,
				" is not running!",
				nagios.NAGIOS_CRITICAL,
			},
		)
	}

	return checks
}

// Runs a set of checkes and returns an array of statuses
func mapAlertStatuses(info *DockerInfo, opts *CliOpts) []*nagios.NagiosStatus {
	var statuses []*nagios.NagiosStatus

	var check = func(args checkArgs) *nagios.NagiosStatus {
		if !args.healthy {
			return &nagios.NagiosStatus{args.tag + ": " + args.value + args.appendErrorMessage, args.statusVal}
		}
		return nil
	}

	checks := defineChecks(info, opts)

	for _, entry := range checks {
		result := check(entry)
		if result != nil {
			statuses = append(statuses, check(entry))
		}
	}

	return statuses
}

// parseCommandLine parses the flags passed on the CLI
func parseCommandLine() *CliOpts {
	var opts CliOpts

	flag.StringVar(&opts.BaseUrl, "base-url", "http://chi-staging-pool-1:4243/", "The Base URL for the Docker server")
	flag.IntVar(&opts.WarnMetaSpace, "warn-meta-space", 100, "Warning threshold for Metadata Space")
	flag.IntVar(&opts.WarnDataSpace, "warn-data-space", 100, "Warning threshold for Data Space")
	flag.IntVar(&opts.CritMetaSpace, "crit-meta-space", 100, "Critical threshold for Metadata Space")
	flag.IntVar(&opts.CritDataSpace, "crit-data-space", 100, "Critical threshold for Data Space")
	flag.StringVar(&opts.ImageId, "image-id", "", "An image ID that must be running on the Docker server")
	flag.IntVar(&opts.GhostsStatus, "ghosts-status", 1, "If ghosts are present, treat as this status")

	flag.Parse()

	return &opts
}

func main() {
	opts := parseCommandLine()

	var fetcher Fetcher
	var info DockerInfo

	err := fetchInfo(fetcher, *opts, &info)
	if err != nil {
		nagios.Critical(err)
	}

	statuses := mapAlertStatuses(&info, opts)
	baseStatus := nagios.NagiosStatus{float64String(info.Containers) + " containers", 0}

	baseStatus.Aggregate(statuses)
	nagios.ExitWithStatus(&baseStatus)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/samuel/go-zookeeper/zk"
	"github.com/zenazn/goji"
	lumberjack "github.com/natefinch/lumberjack"

	"github.com/QubitProducts/bamboo/api"
	"github.com/QubitProducts/bamboo/configuration"
	"github.com/QubitProducts/bamboo/qzk"
	"github.com/QubitProducts/bamboo/services/event_bus"
)

/*
	Commandline arguments
*/
var configFilePath string
var logPath string

func init() {
	flag.StringVar(&configFilePath, "config", "config/development.json", "Full path of the configuration JSON file")
	flag.StringVar(&logPath, "log", "", "Log path to a file. Default logs to stdout")
}

func main() {
	flag.Parse()
	configureLog()

	// Load configuration
	conf, err := configuration.FromFile(configFilePath)
	if err != nil {
		log.Fatal(err)
	}

	eventBus := event_bus.New()

	// Create StatsD client
	conf.StatsD.CreateClient()

	// Create Zookeeper connection
	zkConn := listenToZookeeper(conf, eventBus)

	// Register handlers
	handlers := event_bus.Handlers{ Conf: &conf, Zookeeper: zkConn }
	eventBus.Register(handlers.MarathonEventHandler)
	eventBus.Register(handlers.DomainEventHandler)

	// Start server
	initServer(&conf, zkConn, eventBus)
}

func initServer(conf *configuration.Configuration, conn *zk.Conn, eventBus *event_bus.EventBus) {
	stateAPI := api.State{Config: conf, Zookeeper: conn}
	domainAPI := api.Domain{Config: conf, Zookeeper: conn}
	eventSubAPI := api.EventSubscriptions{Conf: conf, EventBus: eventBus}

	conf.StatsD.Increment(1.0, "restart", 1)
	// Status live information
	goji.Get("/status", api.HandleStatus)

	// State API
	goji.Get("/api/state", stateAPI.Get)

	// Domains API
	goji.Get("/api/state/domains", domainAPI.All)
	goji.Post("/api/state/domains", domainAPI.Create)
	goji.Delete("/api/state/domains/:id", domainAPI.Delete)
	goji.Put("/api/state/domains/:id", domainAPI.Put)

	goji.Post("/api/marathon/event_callback", eventSubAPI.Callback)

	// Static pages
	goji.Get("/*", http.FileServer(http.Dir("./webapp")))

	registerMarathonEvent(conf)

	goji.Serve()
}

func registerMarathonEvent(conf *configuration.Configuration) {
	url := conf.Marathon.Endpoint + "/v2/eventSubscriptions?callbackUrl=" + conf.Bamboo.Host + "/api/marathon/event_callback"
	client := &http.Client{}
	req, _ := http.NewRequest("POST", url, nil)
	req.Header.Add("Content-Type", "application/json")
	client.Do(req)
}

func createAndListen(conf configuration.Zookeeper) (chan zk.Event, *zk.Conn) {
	conn, _, err := zk.Connect(conf.ConnectionString(), time.Second*10)

	if err != nil { log.Panic(err) }

	ch, _ := qzk.ListenToConn(conn, conf.Path, true, conf.Delay())
	return ch, conn
}

func listenToZookeeper(conf configuration.Configuration, eventBus *event_bus.EventBus) *zk.Conn {
	domainCh, domainConn := createAndListen(conf.DomainMapping.Zookeeper)

	go func() {
		for {
			select {
			case _ = <-domainCh:
				eventBus.Publish(&event_bus.DomainEvent{ EventType: "change" })
			}
		}
	}()
	return domainConn
}

func configureLog() {
	if len(logPath) > 0 {
		log.SetOutput(io.MultiWriter(&lumberjack.Logger{
			Filename:   logPath,
			// megabytes
			MaxSize:    100,
			MaxBackups: 3,
			//days
			MaxAge:     28,
		}, os.Stdout))
	}
}
<commit_msg>DomainEvent should not be a reference<commit_after>package main

import (
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/samuel/go-zookeeper/zk"
	"github.com/zenazn/goji"
	lumberjack "github.com/natefinch/lumberjack"

	"github.com/QubitProducts/bamboo/api"
	"github.com/QubitProducts/bamboo/configuration"
	"github.com/QubitProducts/bamboo/qzk"
	"github.com/QubitProducts/bamboo/services/event_bus"
)

/*
	Commandline arguments
*/
var configFilePath string
var logPath string

func init() {
	flag.StringVar(&configFilePath, "config", "config/development.json", "Full path of the configuration JSON file")
	flag.StringVar(&logPath, "log", "", "Log path to a file. Default logs to stdout")
}

func main() {
	flag.Parse()
	configureLog()

	// Load configuration
	conf, err := configuration.FromFile(configFilePath)
	if err != nil {
		log.Fatal(err)
	}

	eventBus := event_bus.New()

	// Create StatsD client
	conf.StatsD.CreateClient()

	// Create Zookeeper connection
	zkConn := listenToZookeeper(conf, eventBus)

	// Register handlers
	handlers := event_bus.Handlers{ Conf: &conf, Zookeeper: zkConn }
	eventBus.Register(handlers.MarathonEventHandler)
	eventBus.Register(handlers.DomainEventHandler)

	// Start server
	initServer(&conf, zkConn, eventBus)
}

func initServer(conf *configuration.Configuration, conn *zk.Conn, eventBus *event_bus.EventBus) {
	stateAPI := api.State{Config: conf, Zookeeper: conn}
	domainAPI := api.Domain{Config: conf, Zookeeper: conn}
	eventSubAPI := api.EventSubscriptions{Conf: conf, EventBus: eventBus}

	conf.StatsD.Increment(1.0, "restart", 1)
	// Status live information
	goji.Get("/status", api.HandleStatus)

	// State API
	goji.Get("/api/state", stateAPI.Get)

	// Domains API
	goji.Get("/api/state/domains", domainAPI.All)
	goji.Post("/api/state/domains", domainAPI.Create)
	goji.Delete("/api/state/domains/:id", domainAPI.Delete)
	goji.Put("/api/state/domains/:id", domainAPI.Put)

	goji.Post("/api/marathon/event_callback", eventSubAPI.Callback)

	// Static pages
	goji.Get("/*", http.FileServer(http.Dir("./webapp")))

	registerMarathonEvent(conf)

	goji.Serve()
}

func registerMarathonEvent(conf *configuration.Configuration) {
	url := conf.Marathon.Endpoint + "/v2/eventSubscriptions?callbackUrl=" + conf.Bamboo.Host + "/api/marathon/event_callback"
	client := &http.Client{}
	req, _ := http.NewRequest("POST", url, nil)
	req.Header.Add("Content-Type", "application/json")
	client.Do(req)
}

func createAndListen(conf configuration.Zookeeper) (chan zk.Event, *zk.Conn) {
	conn, _, err := zk.Connect(conf.ConnectionString(), time.Second*10)

	if err != nil { log.Panic(err) }

	ch, _ := qzk.ListenToConn(conn, conf.Path, true, conf.Delay())
	return ch, conn
}

func listenToZookeeper(conf configuration.Configuration, eventBus *event_bus.EventBus) *zk.Conn {
	domainCh, domainConn := createAndListen(conf.DomainMapping.Zookeeper)

	go func() {
		for {
			select {
			case _ = <-domainCh:
				eventBus.Publish(event_bus.DomainEvent{ EventType: "change" })
			}
		}
	}()
	return domainConn
}

func configureLog() {
	if len(logPath) > 0 {
		log.SetOutput(io.MultiWriter(&lumberjack.Logger{
			Filename:   logPath,
			// megabytes
			MaxSize:    100,
			MaxBackups: 3,
			//days
			MaxAge:     28,
		}, os.Stdout))
	}
}
<|endoftext|>


<commit_before>package reception

import (
	"errors"
	"fmt"
	"net"
	"net/http"

	"github.com/efritz/glock"
)

type checkServer struct {
	host      string
	port      int
	addr      string
	listeners []chan error
	err       error
	mutex.Mutex
	logger Logger
	clock  glock.Clock
}

// ErrZkDisconnect occurs when a health check server is created with an illegal host value.
var ErrIllegalHost = errors.New("illegal host")

func newCheckServer(host string, port int, logger Logger, clock glock.Clock) *checkServer {
	return &checkServer{
		host:      host,
		port:      port,
		listeners: []chan error{},
		logger:    logger,
		clock:     clock,
	}
}

func (s *checkServer) start() {
	if s.host == "" {
		s.signal(ErrIllegalHost)
		return
	}

	listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", s.port))
	if err != nil {
		s.signal(err)
		return
	}

	s.addr = fmt.Sprintf(
		"http://%s:%d/health",
		s.host,
		listener.Addr().(*net.TCPAddr).Port,
	)

	mux := http.NewServeMux()
	mux.HandleFunc("/health", s.handler)

	s.logger.Printf("Running health check at %s\n", s.addr)

	go func() {
		s.signal(http.Serve(listener, mux))
	}()
}

func (s *checkServer) handler(w http.ResponseWriter, r *http.Request) {
	s.logger.Printf("Consul performing health check\n")

	s.signal(nil)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(`{"alive": true}`))
}

func (s *checkServer) register() <-chan error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	ch := make(chan error, 1)

	if s.err != nil {
		ch <- s.err
		close(ch)
		return ch
	}

	s.listeners = append(s.listeners, ch)
	return ch
}

func (s *checkServer) signal(err error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	for _, ch := range s.listeners {
		ch <- err
	}

	if err == nil {
		return
	}

	for _, ch := range s.listeners {
		close(ch)
	}

	s.err = err
	s.listeners = nil
}
<commit_msg>Fix typo in last commit.<commit_after>package reception

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

	"github.com/efritz/glock"
)

type checkServer struct {
	host      string
	port      int
	addr      string
	listeners []chan error
	err       error
	mutex     sync.Mutex
	logger    Logger
	clock     glock.Clock
}

// ErrZkDisconnect occurs when a health check server is created with an illegal host value.
var ErrIllegalHost = errors.New("illegal host")

func newCheckServer(host string, port int, logger Logger, clock glock.Clock) *checkServer {
	return &checkServer{
		host:      host,
		port:      port,
		listeners: []chan error{},
		logger:    logger,
		clock:     clock,
	}
}

func (s *checkServer) start() {
	if s.host == "" {
		s.signal(ErrIllegalHost)
		return
	}

	listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", s.port))
	if err != nil {
		s.signal(err)
		return
	}

	s.addr = fmt.Sprintf(
		"http://%s:%d/health",
		s.host,
		listener.Addr().(*net.TCPAddr).Port,
	)

	mux := http.NewServeMux()
	mux.HandleFunc("/health", s.handler)

	s.logger.Printf("Running health check at %s\n", s.addr)

	go func() {
		s.signal(http.Serve(listener, mux))
	}()
}

func (s *checkServer) handler(w http.ResponseWriter, r *http.Request) {
	s.logger.Printf("Consul performing health check\n")

	s.signal(nil)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(`{"alive": true}`))
}

func (s *checkServer) register() <-chan error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	ch := make(chan error, 1)

	if s.err != nil {
		ch <- s.err
		close(ch)
		return ch
	}

	s.listeners = append(s.listeners, ch)
	return ch
}

func (s *checkServer) signal(err error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	for _, ch := range s.listeners {
		ch <- err
	}

	if err == nil {
		return
	}

	for _, ch := range s.listeners {
		close(ch)
	}

	s.err = err
	s.listeners = nil
}
<|endoftext|>