<commit_before>// Command deployctl implements single-command operator's interface to manage
// deployments running under deploy-registry and deploy-agent
package main

import (
	"archive/tar"
	"compress/gzip"
	"crypto/sha256"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"

	"github.com/artyom/autoflags"
	"github.com/artyom/deploy-tools/internal/shared"
	"github.com/pkg/errors"
	"github.com/pkg/sftp"
)

// defaults can be pre-defined on build time using -ldflags="-X=..."
var defaultAddress, defaultFingerprint string

func main() {
	args := &runArgs{
		Addr: defaultAddress,
		Fp:   defaultFingerprint,
	}
	if val, ok := os.LookupEnv("DEPLOYCTL_ADDR"); ok {
		args.Addr = val
	}
	if val, ok := os.LookupEnv("DEPLOYCTL_FINGERPRINT"); ok {
		args.Fp = val
	}
	fs := flag.NewFlagSet("deployctl", flag.ExitOnError)
	fs.Usage = usageFunc(fs.PrintDefaults)
	autoflags.DefineFlagSet(fs, args)
	fs.Parse(os.Args[1:])
	if err := args.Validate(); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n\n", err)
		fs.Usage()
		os.Exit(2)
	}
	if len(fs.Args()) == 0 {
		fs.Usage()
		os.Exit(2)
	}
	if err := dispatch(args.Addr, args.Key, args.Fp, fs.Args()); err != nil {
		if err == errFlagParseError {
			os.Exit(2)
		}
		if _, ok := err.(*ssh.ExitError); !ok { // don't write "Process exited with status 1"
			fmt.Fprintln(os.Stderr, err)
		}
		os.Exit(1)
	}
}

func dispatch(addr, keyFile, fingerprint string, rawArgs []string) error {
	if len(rawArgs) == 0 {
		return errors.New("nothing to do")
	}
	cmd, args := rawArgs[0], rawArgs[1:]
	switch cmd {
	case "components", "configurations":
		return proxyCommand(addr, keyFile, fingerprint, rawArgs)
	case "addver":
		val := &shared.ArgsAddVersionByFile{}
		if err := parseArgs(cmd, val, os.Stderr, args); err != nil {
			return err
		}
		return uploadAndUpdate(addr, keyFile, fingerprint, val)
	}
	val, err := validatorForCommand(cmd)
	if err != nil {
		return err
	}
	if err := parseArgs(cmd, val, os.Stderr, args); err != nil {
		return err
	}
	return proxyCommand(addr, keyFile, fingerprint, rawArgs)
}

func uploadAndUpdate(addr, keyFile, fingerprint string, args *shared.ArgsAddVersionByFile) error {
	src, err := os.Open(args.File)
	if err != nil {
		return err
	}
	defer src.Close()
	client, cancel, err := dialSSH(addr, keyFile, fingerprint)
	if err != nil {
		return err
	}
	defer cancel()
	sftpconn, err := sftp.NewClient(client)
	if err != nil {
		return err
	}
	defer sftpconn.Close()
	dst, err := sftpconn.Create("upload")
	if err != nil {
		return err
	}
	defer dst.Close()
	pr, pw := io.Pipe()
	defer pw.Close()
	defer pr.Close()
	verifyErr := make(chan error)
	go func() {
		err := decodeArchive(pr)
		pr.CloseWithError(err)
		verifyErr <- err
	}()
	h := sha256.New()
	if _, err := io.Copy(io.MultiWriter(pw, dst, h), src); err != nil {
		return errors.WithMessage(err, "upload failure")
	}
	pw.Close()
	if err := dst.Close(); err != nil {
		return err
	}
	// decodeArchive buffers data, so it may not fail during io.Copy to
	// MultiWriter above if data is copied faster than decoded; ensure we
	// check decode result
	switch err := <-verifyErr; err {
	case io.EOF:
	default:
		return errors.WithMessage(err, "file validation failed")
	}
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	return session.Run(fmt.Sprintf("addver -name=%q -version=%q -hash=%x",
		args.Name, args.Version, h.Sum(nil)))
}

func dialSSH(addr, keyFile, fingerprint string) (client *ssh.Client, closeFunc func(), err error) {
	var signers []ssh.Signer
	switch {
	case keyFile != "":
		s, err := readKey(keyFile)
		if err != nil {
			return nil, nil, errors.Wrapf(err, "failed to read/parse key file %q", keyFile)
		}
		signers = append(signers, s)
	default:
		agentConn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
		if err != nil {
			return nil, nil, errors.WithMessage(err, "cannot connect to ssh-agent, check if SSH_AUTH_SOCK is set")
		}
		defer agentConn.Close()
		sshAgent := agent.NewClient(agentConn)
		signers, err = sshAgent.Signers()
		if err != nil {
			return nil, nil, err
		}
	}
	config := &ssh.ClientConfig{
		User:    os.Getenv("USER"),
		Auth:    []ssh.AuthMethod{ssh.PublicKeys(signers...)},
		Timeout: 30 * time.Second,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			if hostFp := ssh.FingerprintSHA256(key); hostFp != fingerprint {
				return errors.Errorf("host key fingerprint mismatch: %v", hostFp)
			}
			return nil
		},
	}
	client, err = ssh.Dial("tcp", addr, config)
	if err != nil {
		return nil, nil, err
	}
	closeFunc = func() { client.Close() }
	return client, closeFunc, nil
}

func proxyCommand(addr, keyFile, fingerprint string, args []string) error {
	client, cancel, err := dialSSH(addr, keyFile, fingerprint)
	if err != nil {
		return err
	}
	defer cancel()
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	return session.Run(strings.Join(args, " "))
}

type validator interface {
	Validate() error
}

func validatorForCommand(name string) (validator, error) {
	switch name {
	case "delver":
		return &shared.ArgsDelVersion{}, nil
	case "delcomp":
		return &shared.ArgsDelComponent{}, nil
	case "addconf":
		return &shared.ArgsAddConfiguration{}, nil
	case "delconf":
		return &shared.ArgsDelConfiguration{}, nil
	case "changeconf":
		return &shared.ArgsUpdateConfiguration{}, nil
	case "bumpconf":
		return &shared.ArgsBumpConfiguration{}, nil
	case "showconf":
		return &shared.ArgsShowConfiguration{}, nil
	case "showcomp":
		return &shared.ArgsShowComponent{}, nil
	}
	return nil, errors.Errorf("unknown command: %q", name)
}

// errFlagParseError is a sentinel error value used to determine whether error
// originates from flagset that already reported error to stderr so its
// reporting can be omitted
var errFlagParseError = errors.New("flag parse error")

// parseArgs defines new flag set with flags from argStruct that writes its
// errors to w, then calls flag set Parse method on provided raw arguments and
// calls Validate() method on provided argStruct. If parseArgs returns
// errFlagParseError, it means that flag set already reported error to w.
func parseArgs(command string, argStruct validator, w io.Writer, raw []string) error {
	fs := flag.NewFlagSet(command, flag.ContinueOnError)
	fs.SetOutput(w)
	autoflags.DefineFlagSet(fs, argStruct)
	if err := fs.Parse(raw); err != nil {
		return errFlagParseError
	}
	return argStruct.Validate()
}

type runArgs struct {
	Addr string `flag:"addr,$DEPLOYCTL_ADDR, registry host address (host:port)"`
	Fp   string `flag:"fp,$DEPLOYCTL_FINGERPRINT, sha256 host key fingerprint (sha256:...)"`
	Key  string `flag:"key,ssh private key to use; if not set, ssh-agent is used"`
}

func (a *runArgs) Validate() error {
	if a.Addr == "" || a.Fp == "" {
		return errors.New("both addr and fp should be set")
	}
	return nil
}

func usageFunc(printDefaults func()) func() {
	return func() {
		fmt.Fprintln(os.Stderr, "Usage: deployctl [flags] subcommand [subcommand flags]")
		printDefaults()
		fmt.Fprintln(os.Stderr, "\nSubcommands:")
		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, strings.TrimSpace(shared.CommandsListing))
	}
}

// decodeArchive reads and unpacks rd as tar.gz stream, discarding data and
// returning first error it encounters
func decodeArchive(rd io.Reader) error {
	gr, err := gzip.NewReader(rd)
	if err != nil {
		return err
	}
	defer gr.Close()
	tr := tar.NewReader(gr)
	for {
		if _, err := tr.Next(); err != nil {
			return err
		}
		if _, err := io.Copy(ioutil.Discard, tr); err != nil {
			return err
		}
	}
}

func readKey(name string) (ssh.Signer, error) {
	privateBytes, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return ssh.ParsePrivateKey(privateBytes)
}
<commit_msg>deployctl: print progress on file upload<commit_after>// Command deployctl implements single-command operator's interface to manage
// deployments running under deploy-registry and deploy-agent
package main

import (
	"archive/tar"
	"compress/gzip"
	"crypto/sha256"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"

	"github.com/artyom/autoflags"
	"github.com/artyom/deploy-tools/internal/shared"
	"github.com/artyom/isterm"
	"github.com/artyom/progress"
	"github.com/pkg/errors"
	"github.com/pkg/sftp"
)

// defaults can be pre-defined on build time using -ldflags="-X=..."
var defaultAddress, defaultFingerprint string

func main() {
	args := &runArgs{
		Addr: defaultAddress,
		Fp:   defaultFingerprint,
	}
	if val, ok := os.LookupEnv("DEPLOYCTL_ADDR"); ok {
		args.Addr = val
	}
	if val, ok := os.LookupEnv("DEPLOYCTL_FINGERPRINT"); ok {
		args.Fp = val
	}
	fs := flag.NewFlagSet("deployctl", flag.ExitOnError)
	fs.Usage = usageFunc(fs.PrintDefaults)
	autoflags.DefineFlagSet(fs, args)
	fs.Parse(os.Args[1:])
	if err := args.Validate(); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n\n", err)
		fs.Usage()
		os.Exit(2)
	}
	if len(fs.Args()) == 0 {
		fs.Usage()
		os.Exit(2)
	}
	if err := dispatch(args.Addr, args.Key, args.Fp, fs.Args()); err != nil {
		if err == errFlagParseError {
			os.Exit(2)
		}
		if _, ok := err.(*ssh.ExitError); !ok { // don't write "Process exited with status 1"
			fmt.Fprintln(os.Stderr, err)
		}
		os.Exit(1)
	}
}

func dispatch(addr, keyFile, fingerprint string, rawArgs []string) error {
	if len(rawArgs) == 0 {
		return errors.New("nothing to do")
	}
	cmd, args := rawArgs[0], rawArgs[1:]
	switch cmd {
	case "components", "configurations":
		return proxyCommand(addr, keyFile, fingerprint, rawArgs)
	case "addver":
		val := &shared.ArgsAddVersionByFile{}
		if err := parseArgs(cmd, val, os.Stderr, args); err != nil {
			return err
		}
		return uploadAndUpdate(addr, keyFile, fingerprint, val)
	}
	val, err := validatorForCommand(cmd)
	if err != nil {
		return err
	}
	if err := parseArgs(cmd, val, os.Stderr, args); err != nil {
		return err
	}
	return proxyCommand(addr, keyFile, fingerprint, rawArgs)
}

func uploadAndUpdate(addr, keyFile, fingerprint string, args *shared.ArgsAddVersionByFile) error {
	src, err := os.Open(args.File)
	if err != nil {
		return err
	}
	defer src.Close()
	var srcSize int64
	if fi, err := src.Stat(); err == nil {
		srcSize = fi.Size()
	}
	client, cancel, err := dialSSH(addr, keyFile, fingerprint)
	if err != nil {
		return err
	}
	defer cancel()
	sftpconn, err := sftp.NewClient(client)
	if err != nil {
		return err
	}
	defer sftpconn.Close()
	dst, err := sftpconn.Create("upload")
	if err != nil {
		return err
	}
	defer dst.Close()
	pr, pw := io.Pipe()
	defer pw.Close()
	defer pr.Close()
	verifyErr := make(chan error)
	go func() {
		err := decodeArchive(pr)
		pr.CloseWithError(err)
		verifyErr <- err
	}()
	h := sha256.New()
	writers := []io.Writer{pw, dst, h}
	withProgress := isterm.IsTerminal()
	if withProgress {
		writers = append(writers, progress.New(srcSize))
	}
	if _, err := io.Copy(io.MultiWriter(writers...), src); err != nil {
		return errors.WithMessage(err, "upload failure")
	}
	if withProgress {
		fmt.Println() // newline to leave progress above
	}
	pw.Close()
	if err := dst.Close(); err != nil {
		return err
	}
	// decodeArchive buffers data, so it may not fail during io.Copy to
	// MultiWriter above if data is copied faster than decoded; ensure we
	// check decode result
	switch err := <-verifyErr; err {
	case io.EOF:
	default:
		return errors.WithMessage(err, "file validation failed")
	}
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	return session.Run(fmt.Sprintf("addver -name=%q -version=%q -hash=%x",
		args.Name, args.Version, h.Sum(nil)))
}

func dialSSH(addr, keyFile, fingerprint string) (client *ssh.Client, closeFunc func(), err error) {
	var signers []ssh.Signer
	switch {
	case keyFile != "":
		s, err := readKey(keyFile)
		if err != nil {
			return nil, nil, errors.Wrapf(err, "failed to read/parse key file %q", keyFile)
		}
		signers = append(signers, s)
	default:
		agentConn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
		if err != nil {
			return nil, nil, errors.WithMessage(err, "cannot connect to ssh-agent, check if SSH_AUTH_SOCK is set")
		}
		defer agentConn.Close()
		sshAgent := agent.NewClient(agentConn)
		signers, err = sshAgent.Signers()
		if err != nil {
			return nil, nil, err
		}
	}
	config := &ssh.ClientConfig{
		User:    os.Getenv("USER"),
		Auth:    []ssh.AuthMethod{ssh.PublicKeys(signers...)},
		Timeout: 30 * time.Second,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			if hostFp := ssh.FingerprintSHA256(key); hostFp != fingerprint {
				return errors.Errorf("host key fingerprint mismatch: %v", hostFp)
			}
			return nil
		},
	}
	client, err = ssh.Dial("tcp", addr, config)
	if err != nil {
		return nil, nil, err
	}
	closeFunc = func() { client.Close() }
	return client, closeFunc, nil
}

func proxyCommand(addr, keyFile, fingerprint string, args []string) error {
	client, cancel, err := dialSSH(addr, keyFile, fingerprint)
	if err != nil {
		return err
	}
	defer cancel()
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	return session.Run(strings.Join(args, " "))
}

type validator interface {
	Validate() error
}

func validatorForCommand(name string) (validator, error) {
	switch name {
	case "delver":
		return &shared.ArgsDelVersion{}, nil
	case "delcomp":
		return &shared.ArgsDelComponent{}, nil
	case "addconf":
		return &shared.ArgsAddConfiguration{}, nil
	case "delconf":
		return &shared.ArgsDelConfiguration{}, nil
	case "changeconf":
		return &shared.ArgsUpdateConfiguration{}, nil
	case "bumpconf":
		return &shared.ArgsBumpConfiguration{}, nil
	case "showconf":
		return &shared.ArgsShowConfiguration{}, nil
	case "showcomp":
		return &shared.ArgsShowComponent{}, nil
	}
	return nil, errors.Errorf("unknown command: %q", name)
}

// errFlagParseError is a sentinel error value used to determine whether error
// originates from flagset that already reported error to stderr so its
// reporting can be omitted
var errFlagParseError = errors.New("flag parse error")

// parseArgs defines new flag set with flags from argStruct that writes its
// errors to w, then calls flag set Parse method on provided raw arguments and
// calls Validate() method on provided argStruct. If parseArgs returns
// errFlagParseError, it means that flag set already reported error to w.
func parseArgs(command string, argStruct validator, w io.Writer, raw []string) error {
	fs := flag.NewFlagSet(command, flag.ContinueOnError)
	fs.SetOutput(w)
	autoflags.DefineFlagSet(fs, argStruct)
	if err := fs.Parse(raw); err != nil {
		return errFlagParseError
	}
	return argStruct.Validate()
}

type runArgs struct {
	Addr string `flag:"addr,$DEPLOYCTL_ADDR, registry host address (host:port)"`
	Fp   string `flag:"fp,$DEPLOYCTL_FINGERPRINT, sha256 host key fingerprint (sha256:...)"`
	Key  string `flag:"key,ssh private key to use; if not set, ssh-agent is used"`
}

func (a *runArgs) Validate() error {
	if a.Addr == "" || a.Fp == "" {
		return errors.New("both addr and fp should be set")
	}
	return nil
}

func usageFunc(printDefaults func()) func() {
	return func() {
		fmt.Fprintln(os.Stderr, "Usage: deployctl [flags] subcommand [subcommand flags]")
		printDefaults()
		fmt.Fprintln(os.Stderr, "\nSubcommands:")
		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, strings.TrimSpace(shared.CommandsListing))
	}
}

// decodeArchive reads and unpacks rd as tar.gz stream, discarding data and
// returning first error it encounters
func decodeArchive(rd io.Reader) error {
	gr, err := gzip.NewReader(rd)
	if err != nil {
		return err
	}
	defer gr.Close()
	tr := tar.NewReader(gr)
	for {
		if _, err := tr.Next(); err != nil {
			return err
		}
		if _, err := io.Copy(ioutil.Discard, tr); err != nil {
			return err
		}
	}
}

func readKey(name string) (ssh.Signer, error) {
	privateBytes, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return ssh.ParsePrivateKey(privateBytes)
}
<|endoftext|>


<commit_before>package cafs

import (
	"bytes"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sync"

	"github.com/tchaik/tchaik/store"
)

// Index is an interface which contains methods for implementing
// a content addressable file system.
type Index interface {
	// Get returns the real path for the given filename, with true if and only
	// if the path exists in the index.
	Get(path string) (string, bool)

	// Add adds the path to the index, and returns the path to the file
	// and whether the path/content already exists.
	Add(path string, sum string) bool

	// Exists returns true of the sum is in the index.
	Exists(sum string) bool
}

type index struct {
	sync.RWMutex

	files map[string]string // path -> sha1
	index map[string]bool   // {sha1}
}

// NewIndex creates a new file system index.
func NewIndex() *index {
	return &index{
		files: make(map[string]string),
		index: make(map[string]bool),
	}
}

func (i *index) MarshalJSON() ([]byte, error) {
	i.RLock()
	defer i.RUnlock()

	exp := struct {
		Files map[string]string `json:"files"`
	}{
		Files: i.files,
	}
	return json.Marshal(exp)
}

func (i *index) UnmarshalJSON(b []byte) error {
	i.Lock()
	defer i.Unlock()

	var exp struct {
		Files map[string]string
	}
	err := json.Unmarshal(b, &exp)
	if err != nil {
		return err
	}

	i.files = make(map[string]string)
	i.index = make(map[string]bool)
	for k, v := range exp.Files {
		i.index[v] = true
		i.files[k] = v
	}
	return nil
}

// Get implements Index.
func (i *index) Get(path string) (string, bool) {
	i.RLock()
	defer i.RUnlock()

	x, ok := i.files[path]
	return x, ok
}

// Add implements Index.
func (i *index) Add(path, sum string) bool {
	i.Lock()
	defer i.Unlock()

	i.files[path] = sum
	old := i.index[sum]
	i.index[sum] = true
	return old
}

// Exists implements Index.
func (i *index) Exists(sum string) bool {
	i.RLock()
	defer i.RUnlock()

	return i.index[sum]
}

// Add the path+data to the index.  Returns the content sum and true if the content
// already existed in the index, false otherwise.
func AddContent(idx Index, path string, content []byte) (string, bool) {
	if x, ok := idx.Get(path); ok {
		return x, true
	}
	s := fmt.Sprintf("%x", sha1.Sum(content))
	return s, idx.Add(path, s)
}

// FileSystem is a type which defines a content addressable filesystem.
type FileSystem struct {
	idx Index

	fs store.RWFileSystem
}

// Open the file with the given path.  Uses the internal index to identify
// which file to open.  NB: Stat on the returned file will refer to the
// internal file.
func (s *FileSystem) Open(path string) (http.File, error) {
	realPath, ok := s.idx.Get(path)
	if !ok {
		return nil, fmt.Errorf("no such file: %v", path)
	}
	return s.open(realPath)
}

// Wait implements RWFileSystem.
func (s *FileSystem) Wait() error { return nil }

func (s *FileSystem) open(path string) (http.File, error) {
	return s.fs.Open(path)
}

func (s *FileSystem) create(path string) (io.WriteCloser, error) {
	return s.fs.Create(path)
}

type file struct {
	bytes.Buffer

	fs   *FileSystem
	path string
}

// Close acts as a signal that all the information has been written to
// the underlying buffer, and the file can be written to the RWFileSystem.
func (a *file) Close() error {
	_, ok := a.fs.idx.Get(a.path)
	if ok {
		return fmt.Errorf("file already exists: %v", a.path)
	}

	path, ok := AddContent(a.fs.idx, a.path, a.Bytes())
	if !ok {
		fmt.Println("creating", path)
		f, err := a.fs.create(path)
		if err != nil {
			return fmt.Errorf("error creating file: %v", err)
		}
		defer f.Close()
		_, err = io.Copy(f, a)
		if err != nil {
			return fmt.Errorf("error copying data into file '%v': %v", path, err)
		}
	}
	err := a.fs.persist()
	if err != nil {
		return err
	}
	return nil
}

// Create a new file with path. We buffer the contents written to the io.WriteCloser
// so that the content can be hashed and then written to the underlying RWFileSystem.
func (s *FileSystem) Create(path string) (io.WriteCloser, error) {
	_, ok := s.idx.Get(path)
	if ok {
		return nil, fmt.Errorf("file already exists for '%v'", path)
	}

	return &file{
		Buffer: bytes.Buffer{},
		path:   path,
		fs:     s,
	}, nil
}

// initIndex initialises the cafs index.
func (s *FileSystem) initIndex() error {
	f, err := s.open(".idx")
	if err != nil {
		// FIXME: Improve this
		// Can't guarantee that we will get an IsNotExist(err) here
		return nil
	}
	defer f.Close()

	b, err := ioutil.ReadAll(f)
	if err != nil {
		return fmt.Errorf("error reading index: %v", err)
	}
	err = json.Unmarshal(b, s.idx)
	if err != nil {
		return fmt.Errorf("error decoding index: %v", err)
	}
	fmt.Printf("Index initialised: %d files (%d paths)", len(s.idx.(*index).index), len(s.idx.(*index).files))
	return nil
}

func (s *FileSystem) persist() error {
	f, err := s.fs.Create(".idx")
	if err != nil {
		return fmt.Errorf("error creating index: %v", err)
	}
	defer f.Close()

	b, err := json.Marshal(s.idx)
	if err != nil {
		return fmt.Errorf("error encoding index: %v", err)
	}

	_, err = f.Write(b)
	if err != nil {
		return fmt.Errorf("error writing index: %v", err)
	}
	return nil
}

// New creates a new content addressable RWFileSystem.
func New(fs store.RWFileSystem) (*FileSystem, error) {
	s := &FileSystem{
		idx: NewIndex(),
		fs:  fs,
	}
	err := s.initIndex()
	if err != nil {
		return nil, err
	}
	return s, nil
}
<commit_msg>Refactor CAFS to better separate Index from the underlying RWFileSystem.<commit_after>package cafs

import (
	"bytes"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sync"

	"github.com/tchaik/tchaik/store"
)

// Index is an interface which contains methods for implementing
// a content addressable file system.
type Index interface {
	// Get returns the real path for the given filename, with true if and only
	// if the path exists in the index.
	Get(path string) (string, bool)

	// Add adds the path to the index, and returns the path to the file
	// and whether the path/content already exists.
	Add(path string, sum string) (bool, error)

	// Exists returns true of the sum is in the index.
	Exists(sum string) bool
}

type index struct {
	sync.RWMutex

	files map[string]string // path -> sha1
	index map[string]bool   // {sha1}

	fs store.RWFileSystem
}

// NewIndex creates a new file system index.
func NewIndex(fs store.RWFileSystem) (*index, error) {
	idx := &index{
		files: make(map[string]string),
		index: make(map[string]bool),
		fs:    fs,
	}

	f, err := fs.Open("index.json")
	if err != nil {
		// FIXME: Improve this
		// Can't guarantee that we will get an IsNotExist(err) here
		return idx, nil
	}
	defer f.Close()

	b, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, fmt.Errorf("error reading index: %v", err)
	}
	err = json.Unmarshal(b, idx)
	if err != nil {
		return nil, fmt.Errorf("error decoding index: %v", err)
	}
	fmt.Printf("Index initialised: %d files (%d paths)", len(idx.index), len(idx.files))
	return idx, nil
}

func (i *index) MarshalJSON() ([]byte, error) {
	i.RLock()
	defer i.RUnlock()

	exp := struct {
		Files map[string]string `json:"files"`
	}{
		Files: i.files,
	}
	return json.Marshal(exp)
}

func (i *index) UnmarshalJSON(b []byte) error {
	i.Lock()
	defer i.Unlock()

	var exp struct {
		Files map[string]string
	}
	err := json.Unmarshal(b, &exp)
	if err != nil {
		return err
	}

	i.files = make(map[string]string)
	i.index = make(map[string]bool)
	for k, v := range exp.Files {
		i.index[v] = true
		i.files[k] = v
	}
	return nil
}

// Get implements Index.
func (i *index) Get(path string) (string, bool) {
	i.RLock()
	defer i.RUnlock()

	x, ok := i.files[path]
	return x, ok
}

// Add implements Index.
func (i *index) Add(path, sum string) (bool, error) {
	i.Lock()
	defer i.Unlock()

	i.files[path] = sum
	old := i.index[sum]
	i.index[sum] = true
	return old, i.persist()
}

// Exists implements Index.
func (i *index) Exists(sum string) bool {
	i.RLock()
	defer i.RUnlock()

	return i.index[sum]
}

func (i *index) persist() error {
	f, err := i.fs.Create("index.json")
	if err != nil {
		return fmt.Errorf("error creating index: %v", err)
	}
	defer f.Close()

	b, err := json.Marshal(i)
	if err != nil {
		return fmt.Errorf("error encoding index: %v", err)
	}

	_, err = f.Write(b)
	if err != nil {
		return fmt.Errorf("error writing index: %v", err)
	}
	return nil
}

// Add the path+data to the index.  Returns the content sum and true if the content
// already existed in the index, false otherwise.
func AddContent(idx Index, path string, content []byte) (string, bool, error) {
	if x, ok := idx.Get(path); ok {
		return x, true, nil
	}
	s := fmt.Sprintf("%x", sha1.Sum(content))
	ok, err := idx.Add(path, s)
	return s, ok, err
}

// FileSystem is a type which defines a content addressable filesystem.
type FileSystem struct {
	idx Index

	fs store.RWFileSystem
}

// Open the file with the given path.  Uses the internal index to identify
// which file to open.  NB: Stat on the returned file will refer to the
// internal file.
func (s *FileSystem) Open(path string) (http.File, error) {
	realPath, ok := s.idx.Get(path)
	if !ok {
		return nil, fmt.Errorf("no such file: %v", path)
	}
	return s.open(realPath)
}

// Wait implements RWFileSystem.
func (s *FileSystem) Wait() error { return nil }

func (s *FileSystem) open(path string) (http.File, error) {
	return s.fs.Open("content/" + path)
}

func (s *FileSystem) create(path string) (io.WriteCloser, error) {
	return s.fs.Create("content/" + path)
}

type file struct {
	bytes.Buffer

	fs   *FileSystem
	path string
}

// Close acts as a signal that all the information has been written to
// the underlying buffer, and the file can be written to the RWFileSystem.
func (a *file) Close() error {
	_, ok := a.fs.idx.Get(a.path)
	if ok {
		return fmt.Errorf("file already exists: %v", a.path)
	}

	path, ok, err := AddContent(a.fs.idx, a.path, a.Bytes())
	if !ok {
		fmt.Println("creating", path)
		f, err := a.fs.create(path)
		if err != nil {
			return fmt.Errorf("error creating file: %v", err)
		}
		defer f.Close()
		_, err = io.Copy(f, a)
		if err != nil {
			return fmt.Errorf("error copying data into file '%v': %v", path, err)
		}
	}
	return err
}

// Create a new file with path. We buffer the contents written to the io.WriteCloser
// so that the content can be hashed and then written to the underlying RWFileSystem.
func (s *FileSystem) Create(path string) (io.WriteCloser, error) {
	_, ok := s.idx.Get(path)
	if ok {
		return nil, fmt.Errorf("file already exists for '%v'", path)
	}

	return &file{
		Buffer: bytes.Buffer{},
		path:   path,
		fs:     s,
	}, nil
}

// New creates a new content addressable RWFileSystem.
func New(fs store.RWFileSystem) (*FileSystem, error) {
	idx, err := NewIndex(fs)
	if err != nil {
		return nil, err
	}
	return &FileSystem{
		idx: idx,
		fs:  fs,
	}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/cert"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/agent"
	"launchpad.net/juju-core/environs/dummy"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/version"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

type BootstrapSuite struct {
	testing.LoggingSuite
	testing.MgoSuite
}

var _ = Suite(&BootstrapSuite{})

func (s *BootstrapSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
	s.MgoSuite.SetUpSuite(c)
}

func (s *BootstrapSuite) SetUpTest(c *C) {
	s.LoggingSuite.SetUpTest(c)
	s.MgoSuite.SetUpTest(c)
}

func (s *BootstrapSuite) TearDownSuite(c *C) {
	s.MgoSuite.TearDownSuite(c)
	s.LoggingSuite.TearDownSuite(c)
}

func (s *BootstrapSuite) TearDownTest(c *C) {
	s.MgoSuite.TearDownTest(c)
	s.LoggingSuite.TearDownTest(c)
	dummy.Reset()
}

func (*BootstrapSuite) TestWriteCertAndKeyToHome(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()

	cert := []byte("a cert")
	key := []byte("a key")
	err := writeCertAndKeyToHome("foo", cert, key)
	c.Assert(err, IsNil)

	// Check that the generated CA key has been written correctly.
	caCertPEM, err := ioutil.ReadFile(testing.HomePath(".juju", "foo-cert.pem"))
	c.Assert(err, IsNil)
	c.Assert(string(caCertPEM), Equals, "a cert")

	caKeyPEM, err := ioutil.ReadFile(testing.HomePath(".juju", "foo-private-key.pem"))
	c.Assert(err, IsNil)
	c.Assert(string(caKeyPEM), Equals, "a key")
}

func (*BootstrapSuite) TestCheckCertificateMissingKey(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	envName := "peckham"
	keyPath := filepath.Join(os.Getenv("HOME"), ".juju", envName+"-cert.pem")
	ioutil.WriteFile(keyPath, []byte(testing.CACert), 0600)

	env, err := environs.NewFromName(envName)
	c.Assert(err, IsNil)
	err = checkCertificate(env)

	c.Assert(err, ErrorMatches, "environment configuration with a certificate but no CA private key")
}

func (*BootstrapSuite) TestGenerateCertificate(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	envName := "peckham"
	env, err := environs.NewFromName(envName)
	c.Assert(err, IsNil)

	generateCertificate(env)

	// Check that the generated CA key has been written correctly.
	caCertPEM, err := ioutil.ReadFile(testing.HomePath(".juju", envName+"-cert.pem"))
	c.Assert(err, IsNil)
	caKeyPEM, err := ioutil.ReadFile(testing.HomePath(".juju", envName+"-private-key.pem"))
	c.Assert(err, IsNil)

	// Check that the cert and key have been set correctly in the configuration
	cfgCertPEM, cfgCertOK := env.Config().CACert()
	cfgKeyPEM, cfgKeyOK := env.Config().CAPrivateKey()
	c.Assert(cfgCertOK, Equals, true)
	c.Assert(cfgKeyOK, Equals, true)
	c.Assert(cfgCertPEM, DeepEquals, caCertPEM)
	c.Assert(cfgKeyPEM, DeepEquals, caKeyPEM)

	// Check the common name of the generated cert
	caCert, _, err := cert.ParseCertAndKey(cfgCertPEM, cfgKeyPEM)
	c.Assert(err, IsNil)
	c.Assert(caCert.Subject.CommonName, Equals, `juju-generated CA for environment peckham`)
}

func (*BootstrapSuite) TestRunGeneratesCertificate(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	envName := "peckham"
	_, err := testing.RunCommand(c, new(BootstrapCommand), nil)
	c.Assert(err, IsNil)

	// Check that the CA certificate and key have been automatically generated
	// for the environment.
	info, err := os.Stat(testing.HomePath(".juju", envName+"-cert.pem"))
	c.Assert(err, IsNil)
	c.Assert(info.Size() > 0, Equals, true)
	info, err = os.Stat(testing.HomePath(".juju", envName+"-private-key.pem"))
	c.Assert(err, IsNil)
	c.Assert(info.Size() > 0, Equals, true)

	// Check that the environment validates the cert and key.
	env, err := environs.NewFromName(envName)
	c.Assert(err, IsNil)
}

func (*BootstrapSuite) TestBootstrapCommandNoParams(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	opc, errc := runCommand(new(BootstrapCommand))
	c.Check(<-errc, IsNil)
	c.Check((<-opc).(dummy.OpBootstrap).Env, Equals, "peckham")
}

func (*BootstrapSuite) TestBootstrapCommandUploadTools(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	// bootstrap with tool uploading - checking that a file
	// is uploaded should be sufficient, as the detailed semantics
	// of UploadTools are tested in environs.
	opc, errc := runCommand(new(BootstrapCommand), "--upload-tools")
	c.Check(<-errc, IsNil)
	c.Check((<-opc).(dummy.OpPutFile).Env, Equals, "peckham")
	c.Check((<-opc).(dummy.OpBootstrap).Env, Equals, "peckham")

	envs, err := environs.ReadEnvirons("")
	c.Assert(err, IsNil)
	env, err := envs.Open("peckham")
	c.Assert(err, IsNil)

	tools, err := environs.FindTools(env, version.Current, environs.CompatVersion)
	c.Assert(err, IsNil)
	resp, err := http.Get(tools.URL)
	c.Assert(err, IsNil)
	defer resp.Body.Close()

	err = agent.UnpackTools(c.MkDir(), tools, resp.Body)
	c.Assert(err, IsNil)
}

func (*BootstrapSuite) TestBootstrapCommandBrokenEnvironment(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	opc, errc := runCommand(new(BootstrapCommand), "-e", "brokenenv")
	c.Check(<-errc, ErrorMatches, "dummy.Bootstrap is broken")
	c.Check(<-opc, IsNil)
}

func (*BootstrapSuite) TestMissingEnvironment(c *C) {
	defer testing.MakeFakeHomeNoEnvironments(c, "empty").Restore()
	// bootstrap without an environments.yaml
	ctx := testing.Context(c)
	code := cmd.Main(&BootstrapCommand{}, ctx, nil)
	c.Check(code, Equals, 1)
	errStr := ctx.Stderr.(*bytes.Buffer).String()
	strippedErr := strings.Replace(errStr, "\n", "", -1)
	c.Assert(strippedErr, Matches, ".*No juju environment configuration file exists.*")
}
<commit_msg>dumb errors<commit_after>package main

import (
	"bytes"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/cert"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/agent"
	"launchpad.net/juju-core/environs/dummy"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/version"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

type BootstrapSuite struct {
	testing.LoggingSuite
	testing.MgoSuite
}

var _ = Suite(&BootstrapSuite{})

func (s *BootstrapSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
	s.MgoSuite.SetUpSuite(c)
}

func (s *BootstrapSuite) SetUpTest(c *C) {
	s.LoggingSuite.SetUpTest(c)
	s.MgoSuite.SetUpTest(c)
}

func (s *BootstrapSuite) TearDownSuite(c *C) {
	s.MgoSuite.TearDownSuite(c)
	s.LoggingSuite.TearDownSuite(c)
}

func (s *BootstrapSuite) TearDownTest(c *C) {
	s.MgoSuite.TearDownTest(c)
	s.LoggingSuite.TearDownTest(c)
	dummy.Reset()
}

func (*BootstrapSuite) TestWriteCertAndKeyToHome(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()

	cert := []byte("a cert")
	key := []byte("a key")
	err := writeCertAndKeyToHome("foo", cert, key)
	c.Assert(err, IsNil)

	// Check that the generated CA key has been written correctly.
	caCertPEM, err := ioutil.ReadFile(testing.HomePath(".juju", "foo-cert.pem"))
	c.Assert(err, IsNil)
	c.Assert(string(caCertPEM), Equals, "a cert")

	caKeyPEM, err := ioutil.ReadFile(testing.HomePath(".juju", "foo-private-key.pem"))
	c.Assert(err, IsNil)
	c.Assert(string(caKeyPEM), Equals, "a key")
}

func (*BootstrapSuite) TestCheckCertificateMissingKey(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	envName := "peckham"
	keyPath := filepath.Join(os.Getenv("HOME"), ".juju", envName+"-cert.pem")
	ioutil.WriteFile(keyPath, []byte(testing.CACert), 0600)

	env, err := environs.NewFromName(envName)
	c.Assert(err, IsNil)
	err = checkCertificate(env)

	c.Assert(err, ErrorMatches, "environment configuration with a certificate but no CA private key")
}

func (*BootstrapSuite) TestGenerateCertificate(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	envName := "peckham"
	env, err := environs.NewFromName(envName)
	c.Assert(err, IsNil)

	generateCertificate(env)

	// Check that the generated CA key has been written correctly.
	caCertPEM, err := ioutil.ReadFile(testing.HomePath(".juju", envName+"-cert.pem"))
	c.Assert(err, IsNil)
	caKeyPEM, err := ioutil.ReadFile(testing.HomePath(".juju", envName+"-private-key.pem"))
	c.Assert(err, IsNil)

	// Check that the cert and key have been set correctly in the configuration
	cfgCertPEM, cfgCertOK := env.Config().CACert()
	cfgKeyPEM, cfgKeyOK := env.Config().CAPrivateKey()
	c.Assert(cfgCertOK, Equals, true)
	c.Assert(cfgKeyOK, Equals, true)
	c.Assert(cfgCertPEM, DeepEquals, caCertPEM)
	c.Assert(cfgKeyPEM, DeepEquals, caKeyPEM)

	// Check the common name of the generated cert
	caCert, _, err := cert.ParseCertAndKey(cfgCertPEM, cfgKeyPEM)
	c.Assert(err, IsNil)
	c.Assert(caCert.Subject.CommonName, Equals, `juju-generated CA for environment peckham`)
}

func (*BootstrapSuite) TestRunGeneratesCertificate(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	envName := "peckham"
	_, err := testing.RunCommand(c, new(BootstrapCommand), nil)
	c.Assert(err, IsNil)

	// Check that the CA certificate and key have been automatically generated
	// for the environment.
	info, err := os.Stat(testing.HomePath(".juju", envName+"-cert.pem"))
	c.Assert(err, IsNil)
	c.Assert(info.Size() > 0, Equals, true)
	info, err = os.Stat(testing.HomePath(".juju", envName+"-private-key.pem"))
	c.Assert(err, IsNil)
	c.Assert(info.Size() > 0, Equals, true)

	// Check that the environment validates the cert and key.
	_, err = environs.NewFromName(envName)
	c.Assert(err, IsNil)
}

func (*BootstrapSuite) TestBootstrapCommandNoParams(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	opc, errc := runCommand(new(BootstrapCommand))
	c.Check(<-errc, IsNil)
	c.Check((<-opc).(dummy.OpBootstrap).Env, Equals, "peckham")
}

func (*BootstrapSuite) TestBootstrapCommandUploadTools(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	// bootstrap with tool uploading - checking that a file
	// is uploaded should be sufficient, as the detailed semantics
	// of UploadTools are tested in environs.
	opc, errc := runCommand(new(BootstrapCommand), "--upload-tools")
	c.Check(<-errc, IsNil)
	c.Check((<-opc).(dummy.OpPutFile).Env, Equals, "peckham")
	c.Check((<-opc).(dummy.OpBootstrap).Env, Equals, "peckham")

	envs, err := environs.ReadEnvirons("")
	c.Assert(err, IsNil)
	env, err := envs.Open("peckham")
	c.Assert(err, IsNil)

	tools, err := environs.FindTools(env, version.Current, environs.CompatVersion)
	c.Assert(err, IsNil)
	resp, err := http.Get(tools.URL)
	c.Assert(err, IsNil)
	defer resp.Body.Close()

	err = agent.UnpackTools(c.MkDir(), tools, resp.Body)
	c.Assert(err, IsNil)
}

func (*BootstrapSuite) TestBootstrapCommandBrokenEnvironment(c *C) {
	defer testing.MakeFakeHome(c, envConfig).Restore()
	opc, errc := runCommand(new(BootstrapCommand), "-e", "brokenenv")
	c.Check(<-errc, ErrorMatches, "dummy.Bootstrap is broken")
	c.Check(<-opc, IsNil)
}

func (*BootstrapSuite) TestMissingEnvironment(c *C) {
	defer testing.MakeFakeHomeNoEnvironments(c, "empty").Restore()
	// bootstrap without an environments.yaml
	ctx := testing.Context(c)
	code := cmd.Main(&BootstrapCommand{}, ctx, nil)
	c.Check(code, Equals, 1)
	errStr := ctx.Stderr.(*bytes.Buffer).String()
	strippedErr := strings.Replace(errStr, "\n", "", -1)
	c.Assert(strippedErr, Matches, ".*No juju environment configuration file exists.*")
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"crypto/tls"
	"encoding/pem"
	"flag"
	"fmt"
	"os"
	"strings"
)

func main() {

	var server = flag.String("server", "", "The TLS host (name or IP) to connect to")
	var port = flag.Int("port", 443, "The port the TLS service is running on")
	var network = flag.String("net", "tcp", "Connect to this kind of network: tcp, udp, or unix")
	var ipv4only = flag.Bool("4", false, "Only connect via IPv4")
	var ipv6only = flag.Bool("6", false, "Only connect via IPv6")
	var verify = flag.Bool("verify", false, "Verify the provided certificate against trusted CAs")
	var fingerprint = flag.Bool("fingerprint", false, "Print the SHA-256 fingerprint instead of the certificate")
	flag.Parse()

	if *server == "" {
		fmt.Fprintln(os.Stderr, "Need to specify a server.")
		flag.Usage()
		os.Exit(2)
	}

	if *network != "tcp" && *network != "udp" && *network != "unix" {
		fmt.Fprintln(os.Stderr, "Unknown kind of network type! Try tcp.")
		flag.Usage()
		os.Exit(2)
	}

	var network_suffix string = ""
	if *ipv4only && *ipv6only {
		fmt.Println("Specifying both -4 & -6 is redundant")
	} else if *ipv4only {
		network_suffix = "4"
	} else if *ipv6only {
		network_suffix = "6"
	}

	target_host := fmt.Sprintf("%s:%d", *server, *port)
	network_type := fmt.Sprintf("%s%s", *network, network_suffix)

	// Merely connect to the service, do the full handshake and then immediately
	// close the connection. All the information we care about is given in the
	// handshake.
	conn, err := tls.Dial(network_type, target_host, &tls.Config{
		InsecureSkipVerify: !(*verify),
	})
	if err != nil {
		fmt.Fprintln(os.Stderr, "failed to connect: "+err.Error())
		os.Exit(1)
	}
	conn.Close()

	tls_state := conn.ConnectionState()
	remote_certs := tls_state.PeerCertificates
	for _, cert := range remote_certs {
		// skip all of the intermediary & root CA certificates provided by the
		// server
		if cert.BasicConstraintsValid {
			if cert.IsCA {
				continue
			}
		} else {
			continue
		}

		if *fingerprint {
			fmt.Println(strings.Replace(fmt.Sprintf("% X", sha256.Sum256(cert.Raw)),
				" ", ":", -1))
		} else {
			pem.Encode(os.Stdout, &pem.Block{
				Type:  "CERTIFICATE",
				Bytes: cert.Raw,
			})
		}
	}
}
<commit_msg>A bit of refactoring for idiomatic code<commit_after>package main

import (
	"crypto/sha256"
	"crypto/tls"
	"encoding/pem"
	"flag"
	"fmt"
	"os"
	"strings"
	"path"
	"log"
)

var (
	server      = flag.String("server", "", "The TLS host (name or IP) to connect to")
	port        = flag.Int("port", 443, "The port the TLS service is running on")
	network     = flag.String("net", "tcp", "Connect to this kind of network: tcp, udp, or unix")
	ipv4only    = flag.Bool("4", false, "Only connect via IPv4")
	ipv6only    = flag.Bool("6", false, "Only connect via IPv6")
	verify      = flag.Bool("verify", false, "Verify the provided certificate against trusted CAs")
	fingerprint = flag.Bool("fingerprint", false, "Print the SHA-256 fingerprint instead of the certificate")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(
			os.Stderr,
			"Usage: %s [-4|-6] [-net=tcp|udp|unix] [-verify] [-port PORT] -server HOST\n",
			path.Base(os.Args[0]),
		)
		fmt.Fprintf(os.Stderr, "A simple utility to grab just a server's public certificate or fingerprint.\n\n")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}

	flag.Parse()

	if *server == "" {
		flag.Usage()
		log.Fatal("Need to specify a server.")
	}

	if *network != "tcp" && *network != "udp" && *network != "unix" {
		flag.Usage()
		log.Fatal("Unknown kind of network type! Try tcp.")
	}

	var networkSuffix string = ""
	if *ipv4only && *ipv6only {
		log.Print("Specifying both -4 & -6 is redundant.")
	} else if *ipv4only {
		networkSuffix = "4"
	} else if *ipv6only {
		networkSuffix = "6"
	}

	targetHost := fmt.Sprintf("%s:%d", *server, *port)
	networkType := fmt.Sprintf("%s%s", *network, networkSuffix)

	// Merely connect to the service, do the full handshake and then immediately
	// close the connection. All the information we care about is given in the
	// handshake.
	conn, err := tls.Dial(networkType, targetHost, &tls.Config{
		InsecureSkipVerify: !(*verify),
	})
	if err != nil {
		log.Fatalf("failed to connect: %s", err.Error())
	}
	conn.Close()

	tlsState := conn.ConnectionState()
	remoteCerts := tlsState.PeerCertificates
	for _, cert := range remoteCerts {
		// skip all of the intermediary & root CA certificates provided by the
		// server
		if cert.BasicConstraintsValid {
			if cert.IsCA {
				continue
			}
		} else {
			continue
		}

		if *fingerprint {
			// A certificate's fingerprint is simply the hash of the raw ASN.1 encoded
			// certificate. Super simple.
			fmt.Println(strings.Replace(fmt.Sprintf("% X", sha256.Sum256(cert.Raw)),
				" ", ":", -1))
		} else {
			pem.Encode(os.Stdout, &pem.Block{
				Type:  "CERTIFICATE",
				Bytes: cert.Raw,
			})
		}
	}
}
<|endoftext|>


<commit_before>package drain

import (
	bosherr "bosh/errors"
	boshsys "bosh/system"
	"strconv"
	"strings"
)

type ConcreteDrainScript struct {
	fs              boshsys.FileSystem
	runner          boshsys.CmdRunner
	drainScriptPath string
}

func NewConcreteDrainScript(fs boshsys.FileSystem, runner boshsys.CmdRunner, drainScriptPath string) (script ConcreteDrainScript) {
	script = ConcreteDrainScript{
		fs:              fs,
		runner:          runner,
		drainScriptPath: drainScriptPath,
	}
	return
}

func (script ConcreteDrainScript) Exists() bool {
	return script.fs.FileExists(script.drainScriptPath)
}

func (script ConcreteDrainScript) Path() string {
	return script.drainScriptPath
}

func (script ConcreteDrainScript) Run(params DrainScriptParams) (value int, err error) {
	jobChange := params.JobChange()
	hashChange := params.HashChange()
	updatedPkgs := params.UpdatedPackages()

	command := boshsys.Command{
		Name: script.drainScriptPath,
		Env: map[string]string{
			"PATH": "/usr/sbin:/usr/bin:/sbin:/bin",
		},
	}
	command.Args = append(command.Args, jobChange, hashChange)
	command.Args = append(command.Args, updatedPkgs...)

	stdout, _, err := script.runner.RunComplexCommand(command)
	if err != nil {
		err = bosherr.WrapError(err, "Running drain script")
		return
	}

	value, err = strconv.Atoi(strings.TrimSpace(stdout))
	if err != nil {
		err = bosherr.WrapError(err, "Script did not return a signed integer")
	}

	return
}
<commit_msg>formatting<commit_after>package drain

import (
	bosherr "bosh/errors"
	boshsys "bosh/system"
	"strconv"
	"strings"
)

type ConcreteDrainScript struct {
	fs              boshsys.FileSystem
	runner          boshsys.CmdRunner
	drainScriptPath string
}

func NewConcreteDrainScript(
	fs boshsys.FileSystem,
	runner boshsys.CmdRunner,
	drainScriptPath string,
) (script ConcreteDrainScript) {
	script = ConcreteDrainScript{
		fs:              fs,
		runner:          runner,
		drainScriptPath: drainScriptPath,
	}
	return
}

func (script ConcreteDrainScript) Exists() bool {
	return script.fs.FileExists(script.drainScriptPath)
}

func (script ConcreteDrainScript) Path() string {
	return script.drainScriptPath
}

func (script ConcreteDrainScript) Run(params DrainScriptParams) (int, error) {
	jobChange := params.JobChange()
	hashChange := params.HashChange()
	updatedPkgs := params.UpdatedPackages()

	command := boshsys.Command{
		Name: script.drainScriptPath,
		Env: map[string]string{
			"PATH": "/usr/sbin:/usr/bin:/sbin:/bin",
		},
	}
	command.Args = append(command.Args, jobChange, hashChange)
	command.Args = append(command.Args, updatedPkgs...)

	stdout, _, err := script.runner.RunComplexCommand(command)
	if err != nil {
		return 0, bosherr.WrapError(err, "Running drain script")
	}

	value, err := strconv.Atoi(strings.TrimSpace(stdout))
	if err != nil {
		return 0, bosherr.WrapError(err, "Script did not return a signed integer")
	}

	return value, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package framework

import (
	"fmt"
	"net"
	"os"
	"path/filepath"

	"golang.org/x/crypto/ssh"
	v1 "k8s.io/api/core/v1"
	clientset "k8s.io/client-go/kubernetes"
	sshutil "k8s.io/kubernetes/pkg/ssh"
)

// GetSigner returns an ssh.Signer for the provider ("gce", etc.) that can be
// used to SSH to their nodes.
func GetSigner(provider string) (ssh.Signer, error) {
	// Select the key itself to use. When implementing more providers here,
	// please also add them to any SSH tests that are disabled because of signer
	// support.
	keyfile := ""
	switch provider {
	case "gce", "gke", "kubemark":
		keyfile = os.Getenv("GCE_SSH_KEY")
		if keyfile == "" {
			keyfile = "google_compute_engine"
		}
	case "aws", "eks":
		keyfile = os.Getenv("AWS_SSH_KEY")
		if keyfile == "" {
			keyfile = "kube_aws_rsa"
		}
	case "local", "vsphere":
		keyfile = os.Getenv("LOCAL_SSH_KEY")
		if keyfile == "" {
			keyfile = "id_rsa"
		}
	case "skeleton":
		keyfile = os.Getenv("KUBE_SSH_KEY")
		if keyfile == "" {
			keyfile = "id_rsa"
		}
	default:
		return nil, fmt.Errorf("GetSigner(...) not implemented for %s", provider)
	}

	// Respect absolute paths for keys given by user, fallback to assuming
	// relative paths are in ~/.ssh
	if !filepath.IsAbs(keyfile) {
		keydir := filepath.Join(os.Getenv("HOME"), ".ssh")
		keyfile = filepath.Join(keydir, keyfile)
	}

	return sshutil.MakePrivateKeySignerFromFile(keyfile)
}

// NodeSSHHosts returns SSH-able host names for all schedulable nodes - this
// excludes master node. If it can't find any external IPs, it falls back to
// looking for internal IPs. If it can't find an internal IP for every node it
// returns an error, though it still returns all hosts that it found in that
// case.
func NodeSSHHosts(c clientset.Interface) ([]string, error) {
	nodelist := waitListSchedulableNodesOrDie(c)

	hosts := NodeAddresses(nodelist, v1.NodeExternalIP)
	// If ExternalIPs aren't set, assume the test programs can reach the
	// InternalIP. Simplified exception logic here assumes that the hosts will
	// either all have ExternalIP or none will. Simplifies handling here and
	// should be adequate since the setting of the external IPs is provider
	// specific: they should either all have them or none of them will.
	if len(hosts) == 0 {
		Logf("No external IP address on nodes, falling back to internal IPs")
		hosts = NodeAddresses(nodelist, v1.NodeInternalIP)
	}

	// Error if any node didn't have an external/internal IP.
	if len(hosts) != len(nodelist.Items) {
		return hosts, fmt.Errorf(
			"only found %d IPs on nodes, but found %d nodes. Nodelist: %v",
			len(hosts), len(nodelist.Items), nodelist)
	}

	sshHosts := make([]string, 0, len(hosts))
	for _, h := range hosts {
		sshHosts = append(sshHosts, net.JoinHostPort(h, sshPort))
	}
	return sshHosts, nil
}

type SSHResult struct {
	User   string
	Host   string
	Cmd    string
	Stdout string
	Stderr string
	Code   int
}

// NodeExec execs the given cmd on node via SSH. Note that the nodeName is an sshable name,
// eg: the name returned by framework.GetMasterHost(). This is also not guaranteed to work across
// cloud providers since it involves ssh.
func NodeExec(nodeName, cmd string) (SSHResult, error) {
	return SSH(cmd, net.JoinHostPort(nodeName, sshPort), TestContext.Provider)
}

// SSH synchronously SSHs to a node running on provider and runs cmd. If there
// is no error performing the SSH, the stdout, stderr, and exit code are
// returned.
func SSH(cmd, host, provider string) (SSHResult, error) {
	result := SSHResult{Host: host, Cmd: cmd}

	// Get a signer for the provider.
	signer, err := GetSigner(provider)
	if err != nil {
		return result, fmt.Errorf("error getting signer for provider %s: '%v'", provider, err)
	}

	// RunSSHCommand will default to Getenv("USER") if user == "", but we're
	// defaulting here as well for logging clarity.
	result.User = os.Getenv("KUBE_SSH_USER")
	if result.User == "" {
		result.User = os.Getenv("USER")
	}

	stdout, stderr, code, err := sshutil.RunSSHCommand(cmd, result.User, host, signer)
	result.Stdout = stdout
	result.Stderr = stderr
	result.Code = code

	return result, err
}

func LogSSHResult(result SSHResult) {
	remote := fmt.Sprintf("%s@%s", result.User, result.Host)
	Logf("ssh %s: command:   %s", remote, result.Cmd)
	Logf("ssh %s: stdout:    %q", remote, result.Stdout)
	Logf("ssh %s: stderr:    %q", remote, result.Stderr)
	Logf("ssh %s: exit code: %d", remote, result.Code)
}

func IssueSSHCommandWithResult(cmd, provider string, node *v1.Node) (*SSHResult, error) {
	Logf("Getting external IP address for %s", node.Name)
	host := ""
	for _, a := range node.Status.Addresses {
		if a.Type == v1.NodeExternalIP {
			host = net.JoinHostPort(a.Address, sshPort)
			break
		}
	}

	if host == "" {
		// No external IPs were found, let's try to use internal as plan B
		for _, a := range node.Status.Addresses {
			if a.Type == v1.NodeInternalIP {
				host = net.JoinHostPort(a.Address, sshPort)
				break
			}
		}
	}

	if host == "" {
		return nil, fmt.Errorf("couldn't find any IP address for node %s", node.Name)
	}

	Logf("SSH %q on %s(%s)", cmd, node.Name, host)
	result, err := SSH(cmd, host, provider)
	LogSSHResult(result)

	if result.Code != 0 || err != nil {
		return nil, fmt.Errorf("failed running %q: %v (exit code %d, stderr %v)",
			cmd, err, result.Code, result.Stderr)
	}

	return &result, nil
}

func IssueSSHCommand(cmd, provider string, node *v1.Node) error {
	_, err := IssueSSHCommandWithResult(cmd, provider, node)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Allow a single SSH key to be used in all providers<commit_after>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package framework

import (
	"fmt"
	"net"
	"os"
	"path/filepath"

	"golang.org/x/crypto/ssh"
	v1 "k8s.io/api/core/v1"
	clientset "k8s.io/client-go/kubernetes"
	sshutil "k8s.io/kubernetes/pkg/ssh"
)

// GetSigner returns an ssh.Signer for the provider ("gce", etc.) that can be
// used to SSH to their nodes.
func GetSigner(provider string) (ssh.Signer, error) {
	// honor a consistent SSH key across all providers
	if path := os.Getenv("KUBE_SSH_KEY_PATH"); len(path) > 0 {
		return sshutil.MakePrivateKeySignerFromFile(path)
	}

	// Select the key itself to use. When implementing more providers here,
	// please also add them to any SSH tests that are disabled because of signer
	// support.
	keyfile := ""
	switch provider {
	case "gce", "gke", "kubemark":
		keyfile = os.Getenv("GCE_SSH_KEY")
		if keyfile == "" {
			keyfile = "google_compute_engine"
		}
	case "aws", "eks":
		keyfile = os.Getenv("AWS_SSH_KEY")
		if keyfile == "" {
			keyfile = "kube_aws_rsa"
		}
	case "local", "vsphere":
		keyfile = os.Getenv("LOCAL_SSH_KEY")
		if keyfile == "" {
			keyfile = "id_rsa"
		}
	case "skeleton":
		keyfile = os.Getenv("KUBE_SSH_KEY")
		if keyfile == "" {
			keyfile = "id_rsa"
		}
	default:
		return nil, fmt.Errorf("GetSigner(...) not implemented for %s", provider)
	}

	// Respect absolute paths for keys given by user, fallback to assuming
	// relative paths are in ~/.ssh
	if !filepath.IsAbs(keyfile) {
		keydir := filepath.Join(os.Getenv("HOME"), ".ssh")
		keyfile = filepath.Join(keydir, keyfile)
	}

	return sshutil.MakePrivateKeySignerFromFile(keyfile)
}

// NodeSSHHosts returns SSH-able host names for all schedulable nodes - this
// excludes master node. If it can't find any external IPs, it falls back to
// looking for internal IPs. If it can't find an internal IP for every node it
// returns an error, though it still returns all hosts that it found in that
// case.
func NodeSSHHosts(c clientset.Interface) ([]string, error) {
	nodelist := waitListSchedulableNodesOrDie(c)

	hosts := NodeAddresses(nodelist, v1.NodeExternalIP)
	// If ExternalIPs aren't set, assume the test programs can reach the
	// InternalIP. Simplified exception logic here assumes that the hosts will
	// either all have ExternalIP or none will. Simplifies handling here and
	// should be adequate since the setting of the external IPs is provider
	// specific: they should either all have them or none of them will.
	if len(hosts) == 0 {
		Logf("No external IP address on nodes, falling back to internal IPs")
		hosts = NodeAddresses(nodelist, v1.NodeInternalIP)
	}

	// Error if any node didn't have an external/internal IP.
	if len(hosts) != len(nodelist.Items) {
		return hosts, fmt.Errorf(
			"only found %d IPs on nodes, but found %d nodes. Nodelist: %v",
			len(hosts), len(nodelist.Items), nodelist)
	}

	sshHosts := make([]string, 0, len(hosts))
	for _, h := range hosts {
		sshHosts = append(sshHosts, net.JoinHostPort(h, sshPort))
	}
	return sshHosts, nil
}

type SSHResult struct {
	User   string
	Host   string
	Cmd    string
	Stdout string
	Stderr string
	Code   int
}

// NodeExec execs the given cmd on node via SSH. Note that the nodeName is an sshable name,
// eg: the name returned by framework.GetMasterHost(). This is also not guaranteed to work across
// cloud providers since it involves ssh.
func NodeExec(nodeName, cmd string) (SSHResult, error) {
	return SSH(cmd, net.JoinHostPort(nodeName, sshPort), TestContext.Provider)
}

// SSH synchronously SSHs to a node running on provider and runs cmd. If there
// is no error performing the SSH, the stdout, stderr, and exit code are
// returned.
func SSH(cmd, host, provider string) (SSHResult, error) {
	result := SSHResult{Host: host, Cmd: cmd}

	// Get a signer for the provider.
	signer, err := GetSigner(provider)
	if err != nil {
		return result, fmt.Errorf("error getting signer for provider %s: '%v'", provider, err)
	}

	// RunSSHCommand will default to Getenv("USER") if user == "", but we're
	// defaulting here as well for logging clarity.
	result.User = os.Getenv("KUBE_SSH_USER")
	if result.User == "" {
		result.User = os.Getenv("USER")
	}

	stdout, stderr, code, err := sshutil.RunSSHCommand(cmd, result.User, host, signer)
	result.Stdout = stdout
	result.Stderr = stderr
	result.Code = code

	return result, err
}

func LogSSHResult(result SSHResult) {
	remote := fmt.Sprintf("%s@%s", result.User, result.Host)
	Logf("ssh %s: command:   %s", remote, result.Cmd)
	Logf("ssh %s: stdout:    %q", remote, result.Stdout)
	Logf("ssh %s: stderr:    %q", remote, result.Stderr)
	Logf("ssh %s: exit code: %d", remote, result.Code)
}

func IssueSSHCommandWithResult(cmd, provider string, node *v1.Node) (*SSHResult, error) {
	Logf("Getting external IP address for %s", node.Name)
	host := ""
	for _, a := range node.Status.Addresses {
		if a.Type == v1.NodeExternalIP {
			host = net.JoinHostPort(a.Address, sshPort)
			break
		}
	}

	if host == "" {
		// No external IPs were found, let's try to use internal as plan B
		for _, a := range node.Status.Addresses {
			if a.Type == v1.NodeInternalIP {
				host = net.JoinHostPort(a.Address, sshPort)
				break
			}
		}
	}

	if host == "" {
		return nil, fmt.Errorf("couldn't find any IP address for node %s", node.Name)
	}

	Logf("SSH %q on %s(%s)", cmd, node.Name, host)
	result, err := SSH(cmd, host, provider)
	LogSSHResult(result)

	if result.Code != 0 || err != nil {
		return nil, fmt.Errorf("failed running %q: %v (exit code %d, stderr %v)",
			cmd, err, result.Code, result.Stderr)
	}

	return &result, nil
}

func IssueSSHCommand(cmd, provider string, node *v1.Node) error {
	_, err := IssueSSHCommandWithResult(cmd, provider, node)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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 e2e

import (
	"fmt"
	"strings"
	"testing"
	"time"

	habv1beta1 "github.com/habitat-sh/habitat-operator/pkg/apis/habitat/v1beta1"
	utils "github.com/habitat-sh/habitat-operator/test/e2e/framework"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	defaultWaitTime = 1 * time.Minute
	configMapName   = "peer-watch-file"
)

// TestBind tests that the operator correctly created two Habitat Services and bound them together.
func TestBind(t *testing.T) {
	// Get Habitat object from Habitat go example.
	web, err := utils.ConvertHabitat("resources/bind-config/webapp.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(web); err != nil {
		t.Fatal(err)
	}

	// Get Habitat object from Habitat db example.
	db, err := utils.ConvertHabitat("resources/bind-config/db.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(db); err != nil {
		t.Fatal(err)
	}

	// Get Service object from example file.
	svc, err := utils.ConvertService("resources/bind-config/service.yml")
	if err != nil {
		t.Fatal(err)
	}

	// Create Service.
	_, err = framework.KubeClient.CoreV1().Services(utils.TestNs).Create(svc)
	if err != nil {
		t.Fatal(err)
	}
	// Delete Service so it doesn't interfere with other tests.
	defer (func(name string) {
		if err := framework.DeleteService(name); err != nil {
			t.Fatal(err)
		}
	})(svc.Name)

	// Get Secret object from example file.
	sec, err := utils.ConvertSecret("resources/bind-config/secret.yml")
	if err != nil {
		t.Fatal(err)
	}

	// Create Secret.
	sec, err = framework.KubeClient.CoreV1().Secrets(utils.TestNs).Create(sec)
	if err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be ready.
	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, web.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}
	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, db.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Wait until endpoints are ready.
	if err := framework.WaitForEndpoints(svc.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}

	time.Sleep(defaultWaitTime)

	// Get response from Habitat Service.
	url := fmt.Sprintf("http://%s:30001/", framework.ExternalIP)

	body, err := utils.QueryService(url)
	if err != nil {
		t.Fatal(err)
	}

	// This msg is set in the config of the habitat/bindgo-hab Go Habitat Service.
	expectedMsg := "hello from port: 4444"
	actualMsg := body
	// actualMsg can contain whitespace and newlines or different formatting,
	// the only thing we need to check is it contains the expectedMsg.
	if !strings.Contains(actualMsg, expectedMsg) {
		t.Fatalf("Habitat Service msg does not match one in default.toml. Expected: \"%s\", got: \"%s\"", expectedMsg, actualMsg)
	}

	// Update secret.
	newPort := "port = 6333"

	sec.Data["user.toml"] = []byte(newPort)
	if _, err = framework.KubeClient.CoreV1().Secrets(utils.TestNs).Update(sec); err != nil {
		t.Fatalf("Could not update Secret: \"%s\"", err)
	}

	// Wait for SecretVolume to be updated.
	time.Sleep(defaultWaitTime)

	// Check that the port differs after the update.
	body, err = utils.QueryService(url)
	if err != nil {
		t.Fatal(err)
	}

	// Update the message set in the config of the habitat/bindgo-hab Go Habitat Service.
	expectedMsg = fmt.Sprintf("hello from port: %v", 6333)
	actualMsg = body
	// actualMsg can contain whitespace and newlines or different formatting,
	// the only thing we need to check is it contains the expectedMsg.
	if !strings.Contains(actualMsg, expectedMsg) {
		t.Fatalf("Configuration update did not go through. Expected: \"%s\", got: \"%s\"", expectedMsg, actualMsg)
	}
}

// TestHabitatDelete tests Habitat deletion.
func TestHabitatDelete(t *testing.T) {
	// Get Habitat object from Habitat go example.
	habitat, err := utils.ConvertHabitat("resources/standalone/habitat.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(habitat); err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be ready.
	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, habitat.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Delete Habitat.
	if err := framework.DeleteHabitat(habitat.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be deleted.
	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, habitat.ObjectMeta.Name, 0); err != nil {
		t.Fatal(err)
	}

	// Check if all the resources the operator creates are deleted.
	// We do not care about secrets being deleted, as the user needs to delete those manually.
	d, err := framework.KubeClient.AppsV1beta1().Deployments(utils.TestNs).Get(habitat.ObjectMeta.Name, metav1.GetOptions{})
	if err == nil && d != nil {
		t.Fatal("Deployment was not deleted.")
	}

	// The CM with the peer IP should still be alive, despite the Habitat being deleted as it was created outside of the scope of a Habitat.
	_, err = framework.KubeClient.CoreV1().ConfigMaps(utils.TestNs).Get(configMapName, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}
}

func TestHabitatPersistence(t *testing.T) {
	ephemeral, err := utils.ConvertHabitat("resources/standalone/habitat.yml")
	if err != nil {
		t.Fatal(err)
	}

	persisted, err := utils.ConvertHabitat("resources/persisted/habitat.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(ephemeral); err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(persisted); err != nil {
		t.Fatal(err)
	}

	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, persisted.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Test that persistence is only enabled if requested
	ephemeralSTS, err := framework.KubeClient.AppsV1beta1().StatefulSets(utils.TestNs).Get(ephemeral.Name, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}

	if len(ephemeralSTS.Spec.VolumeClaimTemplates) != 0 {
		t.Fatal("PersistentVolumeClaims created for ephemeral StatefulSet")
	}

	persistedSTS, err := framework.KubeClient.AppsV1beta1().StatefulSets(utils.TestNs).Get(persisted.Name, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}

	if len(persistedSTS.Spec.VolumeClaimTemplates) == 0 {
		t.Fatal("No PersistentVolumeClaims created for persistent StatefulSet")
	}
}
<commit_msg>Add "section" comment in E2E test<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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 e2e

import (
	"fmt"
	"strings"
	"testing"
	"time"

	habv1beta1 "github.com/habitat-sh/habitat-operator/pkg/apis/habitat/v1beta1"
	utils "github.com/habitat-sh/habitat-operator/test/e2e/framework"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	defaultWaitTime = 1 * time.Minute
	configMapName   = "peer-watch-file"
)

// TestBind tests that the operator correctly created two Habitat Services and bound them together.
func TestBind(t *testing.T) {
	// Get Habitat object from Habitat go example.
	web, err := utils.ConvertHabitat("resources/bind-config/webapp.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(web); err != nil {
		t.Fatal(err)
	}

	// Get Habitat object from Habitat db example.
	db, err := utils.ConvertHabitat("resources/bind-config/db.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(db); err != nil {
		t.Fatal(err)
	}

	// Get Service object from example file.
	svc, err := utils.ConvertService("resources/bind-config/service.yml")
	if err != nil {
		t.Fatal(err)
	}

	// Create Service.
	_, err = framework.KubeClient.CoreV1().Services(utils.TestNs).Create(svc)
	if err != nil {
		t.Fatal(err)
	}
	// Delete Service so it doesn't interfere with other tests.
	defer (func(name string) {
		if err := framework.DeleteService(name); err != nil {
			t.Fatal(err)
		}
	})(svc.Name)

	// Get Secret object from example file.
	sec, err := utils.ConvertSecret("resources/bind-config/secret.yml")
	if err != nil {
		t.Fatal(err)
	}

	// Create Secret.
	sec, err = framework.KubeClient.CoreV1().Secrets(utils.TestNs).Create(sec)
	if err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be ready.
	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, web.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}
	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, db.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Wait until endpoints are ready.
	if err := framework.WaitForEndpoints(svc.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}

	time.Sleep(defaultWaitTime)

	// Get response from Habitat Service.
	url := fmt.Sprintf("http://%s:30001/", framework.ExternalIP)

	body, err := utils.QueryService(url)
	if err != nil {
		t.Fatal(err)
	}

	// This msg is set in the config of the habitat/bindgo-hab Go Habitat Service.
	expectedMsg := "hello from port: 4444"
	actualMsg := body
	// actualMsg can contain whitespace and newlines or different formatting,
	// the only thing we need to check is it contains the expectedMsg.
	if !strings.Contains(actualMsg, expectedMsg) {
		t.Fatalf("Habitat Service msg does not match one in default.toml. Expected: \"%s\", got: \"%s\"", expectedMsg, actualMsg)
	}

	// Test `user.toml` updates.

	// Update secret.
	newPort := "port = 6333"

	sec.Data["user.toml"] = []byte(newPort)
	if _, err = framework.KubeClient.CoreV1().Secrets(utils.TestNs).Update(sec); err != nil {
		t.Fatalf("Could not update Secret: \"%s\"", err)
	}

	// Wait for SecretVolume to be updated.
	time.Sleep(defaultWaitTime)

	// Check that the port differs after the update.
	body, err = utils.QueryService(url)
	if err != nil {
		t.Fatal(err)
	}

	// Update the message set in the config of the habitat/bindgo-hab Go Habitat Service.
	expectedMsg = fmt.Sprintf("hello from port: %v", 6333)
	actualMsg = body
	// actualMsg can contain whitespace and newlines or different formatting,
	// the only thing we need to check is it contains the expectedMsg.
	if !strings.Contains(actualMsg, expectedMsg) {
		t.Fatalf("Configuration update did not go through. Expected: \"%s\", got: \"%s\"", expectedMsg, actualMsg)
	}
}

// TestHabitatDelete tests Habitat deletion.
func TestHabitatDelete(t *testing.T) {
	// Get Habitat object from Habitat go example.
	habitat, err := utils.ConvertHabitat("resources/standalone/habitat.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(habitat); err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be ready.
	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, habitat.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Delete Habitat.
	if err := framework.DeleteHabitat(habitat.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be deleted.
	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, habitat.ObjectMeta.Name, 0); err != nil {
		t.Fatal(err)
	}

	// Check if all the resources the operator creates are deleted.
	// We do not care about secrets being deleted, as the user needs to delete those manually.
	d, err := framework.KubeClient.AppsV1beta1().Deployments(utils.TestNs).Get(habitat.ObjectMeta.Name, metav1.GetOptions{})
	if err == nil && d != nil {
		t.Fatal("Deployment was not deleted.")
	}

	// The CM with the peer IP should still be alive, despite the Habitat being deleted as it was created outside of the scope of a Habitat.
	_, err = framework.KubeClient.CoreV1().ConfigMaps(utils.TestNs).Get(configMapName, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}
}

func TestHabitatPersistence(t *testing.T) {
	ephemeral, err := utils.ConvertHabitat("resources/standalone/habitat.yml")
	if err != nil {
		t.Fatal(err)
	}

	persisted, err := utils.ConvertHabitat("resources/persisted/habitat.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(ephemeral); err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(persisted); err != nil {
		t.Fatal(err)
	}

	if err := framework.WaitForResources(habv1beta1.HabitatNameLabel, persisted.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Test that persistence is only enabled if requested
	ephemeralSTS, err := framework.KubeClient.AppsV1beta1().StatefulSets(utils.TestNs).Get(ephemeral.Name, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}

	if len(ephemeralSTS.Spec.VolumeClaimTemplates) != 0 {
		t.Fatal("PersistentVolumeClaims created for ephemeral StatefulSet")
	}

	persistedSTS, err := framework.KubeClient.AppsV1beta1().StatefulSets(utils.TestNs).Get(persisted.Name, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}

	if len(persistedSTS.Spec.VolumeClaimTemplates) == 0 {
		t.Fatal("No PersistentVolumeClaims created for persistent StatefulSet")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"

	"github.com/fern4lvarez/piladb/config/vars"
	"github.com/fern4lvarez/piladb/pila"
	"github.com/gorilla/mux"
)

// These vars represent the command line flags.
// They are only used to initialize the Connection
// Config at pilad start-up.
var (
	maxStackSizeFlag                  int
	readTimeoutFlag, writeTimeoutFlag int
	portFlag                          int
)

func init() {
	flag.IntVar(&maxStackSizeFlag, "max-stack-size", -1, "Max size of Stacks")
	flag.IntVar(&readTimeoutFlag, "read-timeout", 30, "Read request timeout")
	flag.IntVar(&writeTimeoutFlag, "write-timeout", 45, "Write response timeout")
	flag.IntVar(&portFlag, "port", 1205, "Port number")
}

type flagKey struct {
	flag interface{}
	key  string
}

// buildConfig sets non-default config values to the Connection
// reading from environment variables and cli flags.
func (c *Conn) buildConfig() {
	flagKeys := []flagKey{
		flagKey{maxStackSizeFlag, vars.MaxStackSize},
		flagKey{readTimeoutFlag, vars.ReadTimeout},
		flagKey{writeTimeoutFlag, vars.WriteTimeout},
		flagKey{portFlag, vars.Port},
	}

	for _, fk := range flagKeys {
		if e := os.Getenv(vars.Env(fk.key)); e != "" {
			fmt.Println("DEBUG", fk.key)
			if i, err := strconv.Atoi(e); err != nil {
				fmt.Println("DEBUG", vars.DefaultInt(fk.key))
				c.Config.Set(fk.key, vars.DefaultInt(fk.key))
			} else {
				c.Config.Set(fk.key, i)

			}
			continue
		}
		c.Config.Set(fk.key, fk.flag)
	}
}

// stackHandlerFunc represents a Handler of a Stack.
type stackHandlerFunc func(w http.ResponseWriter, r *http.Request, stack *pila.Stack)

// configHandler handles a request to the Conn configuration.
func (c *Conn) configHandler(w http.ResponseWriter, r *http.Request) {
	res, err := c.Config.Values.StacksKV().ToJSON()
	if err != nil {
		log.Println(r.Method, r.URL, http.StatusBadRequest,
			"error on response serialization:", err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(res)
	log.Println(r.Method, r.URL, http.StatusOK)
}

// configKeyHandler handles a config value.
func (c *Conn) configKeyHandler(configKey string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)

		// we override the mux vars to be able to test
		// an arbitrary configKey
		if configKey != "" {
			vars = map[string]string{
				"key": configKey,
			}
		}
		value := c.Config.Get(vars["key"])
		if value == nil {
			c.goneHandler(w, r, fmt.Sprintf("%s is not set", vars["key"]))
			return
		}

		var element pila.Element
		if r.Method == "GET" {
			value := c.Config.Get(vars["key"])
			element.Value = value
		}
		if r.Method == "POST" {
			if r.Body == nil {
				log.Println(r.Method, r.URL, http.StatusBadRequest,
					"no element provided")
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			err := element.Decode(r.Body)
			if err != nil {
				log.Println(r.Method, r.URL, http.StatusBadRequest,
					"error on decoding element:", err)
				w.WriteHeader(http.StatusBadRequest)
				return
			}

			c.Config.Set(vars["key"], element.Value)
		}

		log.Println(r.Method, r.URL, http.StatusOK, element.Value)
		w.Header().Set("Content-Type", "application/json")

		b, err := element.ToJSON()
		if err != nil {
			log.Println(r.Method, r.URL, http.StatusBadRequest,
				"error on decoding element:", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		w.Write(b)
	})
}

// checkMaxStackSize checks config value for MaxStackSize and execute the
// wrapped handler if check is validated.
func (c *Conn) checkMaxStackSize(handler stackHandlerFunc) stackHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, stack *pila.Stack) {
		if s := c.Config.MaxStackSize(); stack.Size() >= s && s != -1 {
			log.Println(r.Method, r.URL, http.StatusNotAcceptable, vars.MaxStackSize, "value reached")
			w.WriteHeader(http.StatusNotAcceptable)
			return
		}

		handler(w, r, stack)
	}
}
<commit_msg>pilad: cleanup old debugging<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"

	"github.com/fern4lvarez/piladb/config/vars"
	"github.com/fern4lvarez/piladb/pila"
	"github.com/gorilla/mux"
)

// These vars represent the command line flags.
// They are only used to initialize the Connection
// Config at pilad start-up.
var (
	maxStackSizeFlag                  int
	readTimeoutFlag, writeTimeoutFlag int
	portFlag                          int
)

func init() {
	flag.IntVar(&maxStackSizeFlag, "max-stack-size", -1, "Max size of Stacks")
	flag.IntVar(&readTimeoutFlag, "read-timeout", 30, "Read request timeout")
	flag.IntVar(&writeTimeoutFlag, "write-timeout", 45, "Write response timeout")
	flag.IntVar(&portFlag, "port", 1205, "Port number")
}

type flagKey struct {
	flag interface{}
	key  string
}

// buildConfig sets non-default config values to the Connection
// reading from environment variables and cli flags.
func (c *Conn) buildConfig() {
	flagKeys := []flagKey{
		flagKey{maxStackSizeFlag, vars.MaxStackSize},
		flagKey{readTimeoutFlag, vars.ReadTimeout},
		flagKey{writeTimeoutFlag, vars.WriteTimeout},
		flagKey{portFlag, vars.Port},
	}

	for _, fk := range flagKeys {
		if e := os.Getenv(vars.Env(fk.key)); e != "" {
			if i, err := strconv.Atoi(e); err != nil {
				c.Config.Set(fk.key, vars.DefaultInt(fk.key))
			} else {
				c.Config.Set(fk.key, i)

			}
			continue
		}
		c.Config.Set(fk.key, fk.flag)
	}
}

// stackHandlerFunc represents a Handler of a Stack.
type stackHandlerFunc func(w http.ResponseWriter, r *http.Request, stack *pila.Stack)

// configHandler handles a request to the Conn configuration.
func (c *Conn) configHandler(w http.ResponseWriter, r *http.Request) {
	res, err := c.Config.Values.StacksKV().ToJSON()
	if err != nil {
		log.Println(r.Method, r.URL, http.StatusBadRequest,
			"error on response serialization:", err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(res)
	log.Println(r.Method, r.URL, http.StatusOK)
}

// configKeyHandler handles a config value.
func (c *Conn) configKeyHandler(configKey string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)

		// we override the mux vars to be able to test
		// an arbitrary configKey
		if configKey != "" {
			vars = map[string]string{
				"key": configKey,
			}
		}
		value := c.Config.Get(vars["key"])
		if value == nil {
			c.goneHandler(w, r, fmt.Sprintf("%s is not set", vars["key"]))
			return
		}

		var element pila.Element
		if r.Method == "GET" {
			value := c.Config.Get(vars["key"])
			element.Value = value
		}
		if r.Method == "POST" {
			if r.Body == nil {
				log.Println(r.Method, r.URL, http.StatusBadRequest,
					"no element provided")
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			err := element.Decode(r.Body)
			if err != nil {
				log.Println(r.Method, r.URL, http.StatusBadRequest,
					"error on decoding element:", err)
				w.WriteHeader(http.StatusBadRequest)
				return
			}

			c.Config.Set(vars["key"], element.Value)
		}

		log.Println(r.Method, r.URL, http.StatusOK, element.Value)
		w.Header().Set("Content-Type", "application/json")

		b, err := element.ToJSON()
		if err != nil {
			log.Println(r.Method, r.URL, http.StatusBadRequest,
				"error on decoding element:", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		w.Write(b)
	})
}

// checkMaxStackSize checks config value for MaxStackSize and execute the
// wrapped handler if check is validated.
func (c *Conn) checkMaxStackSize(handler stackHandlerFunc) stackHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, stack *pila.Stack) {
		if s := c.Config.MaxStackSize(); stack.Size() >= s && s != -1 {
			log.Println(r.Method, r.URL, http.StatusNotAcceptable, vars.MaxStackSize, "value reached")
			w.WriteHeader(http.StatusNotAcceptable)
			return
		}

		handler(w, r, stack)
	}
}
<|endoftext|>


<commit_before>/*
 */
package trayhost

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"sort"
)

/*
#cgo linux pkg-config: gtk+-2.0
#cgo linux CFLAGS: -DLINUX -I/usr/include/libappindicator-0.1/
#cgo linux LDFLAGS: -ldl
#cgo windows CFLAGS: -DWIN32 -DUNICODE
#cgo darwin CFLAGS: -DDARWIN -x objective-c
#cgo darwin LDFLAGS: -framework Cocoa
#include <stdlib.h>
#include "platform/platform.h"
*/
import "C"

const (
	ICON_PRIMARY     = iota
	ICON_ALTERNATIVE = iota
	ICON_ATTENTION   = iota
)

const (
	WINDOWS       = iota
	OSX           = iota
	GNOME         = iota
	UNITY         = iota
	LINUX_GENERIC = iota
)

type MenuItem struct {
	Title    string
	Disabled bool
	Handler  func()
}

type MenuItemUpdate struct {
	ItemId int
	Item   MenuItem
}

type MenuItems map[int]MenuItem

var isExiting bool = false
var menuItems MenuItems
var UpdateCh = make(chan MenuItemUpdate, 99)
var icons = map[int]string{}
var tmpFiles []string = make([]string, 0, 3)
var clickHandler func()
var Debug bool = false

// Run the host system's event loop
func Initialize(title string, imageData []byte, items MenuItems) (err error) {
	if !Debug {
		log.SetOutput(ioutil.Discard)
	}
	SetIconImage(ICON_PRIMARY, imageData)
	initialize(title)
	SetIcon(ICON_PRIMARY)
	setMenu(items)
	return
}

func SetIconImage(iconId int, imageData []byte) (err error) {
	iconPth, err := createTempFile(imageData)
	if err != nil {
		return
	}
	icons[iconId] = iconPth
	return
}

func EnterLoop() {
	log.Println("Entering native loop")
	go menuUpdateLoop()
	C.native_loop()
	// If reached, user clicked Exit
	isExiting = true
}

func Exit() {
	log.Println("Exiting native loop")
	close(UpdateCh)
	C.exit_loop()
	cleanup()
}

func SetIcon(iconId int) (err error) {

	iconPth, ok := icons[iconId]
	if !ok {
		err = fmt.Errorf("No icon with icon id %d", iconId)
		return
	}

	log.Printf("Setting icon %s (id: %d)", iconPth, iconId)
	setIcon(iconPth)
	return
}

func SetClickHandler(handler func()) {
	clickHandler = handler
}

func menuUpdateLoop() {
	runtime.UnlockOSThread()
	for update := range UpdateCh {
		updateMenuItem(update.ItemId, update.Item)
	}
}

func updateMenuItem(id int, item MenuItem) {
	menuItems[id] = item
	setMenuItem(id, item)
}

func setMenu(menu MenuItems) {
	menuItems = menu
	menuItemOrder := make([]int, 0, len(menuItems))
	for key, _ := range menuItems {
		menuItemOrder = append(menuItemOrder, key)
	}
	sort.Ints(menuItemOrder)

	for id := range menuItemOrder {
		item := menuItems[id]
		setMenuItem(id, item)
	}
}

func createTempFile(iconData []byte) (filename string, err error) {
	file, err := ioutil.TempFile(os.TempDir(), "trayhosticon")
	if err != nil {
		return
	}
	defer file.Close()
	_, err = file.Write(iconData)
	filename = file.Name()
	tmpFiles = append(tmpFiles, filename)
	log.Printf("Temp file created: %s\n", filename)
	return
}

func cleanup() {
	for _, file := range tmpFiles {
		err := os.Remove(file)
		if err != nil {
			log.Printf("Failed to remove tmp file %s: %v\n", file, err)
		} else {
			log.Printf("Tmp file %s removed\n", file)
		}
	}
}

func cbool(b bool) C.int {
	if b {
		return 1
	} else {
		return 0
	}
}
<commit_msg>Call c code to change icon only when really needed<commit_after>/*
 */
package trayhost

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"sort"
)

/*
#cgo linux pkg-config: gtk+-2.0
#cgo linux CFLAGS: -DLINUX -I/usr/include/libappindicator-0.1/
#cgo linux LDFLAGS: -ldl
#cgo windows CFLAGS: -DWIN32 -DUNICODE
#cgo darwin CFLAGS: -DDARWIN -x objective-c
#cgo darwin LDFLAGS: -framework Cocoa
#include <stdlib.h>
#include "platform/platform.h"
*/
import "C"

const (
	ICON_PRIMARY     = iota
	ICON_ALTERNATIVE = iota
	ICON_ATTENTION   = iota
)

const (
	WINDOWS       = iota
	OSX           = iota
	GNOME         = iota
	UNITY         = iota
	LINUX_GENERIC = iota
)

type MenuItem struct {
	Title    string
	Disabled bool
	Handler  func()
}

type MenuItemUpdate struct {
	ItemId int
	Item   MenuItem
}

type MenuItems map[int]MenuItem

var isExiting bool = false
var menuItems MenuItems
var UpdateCh = make(chan MenuItemUpdate, 99)
var icons = map[int]string{}
var tmpFiles []string = make([]string, 0, 3)
var clickHandler func()
var Debug bool = false
var curIconId int = -1

// Run the host system's event loop
func Initialize(title string, imageData []byte, items MenuItems) (err error) {
	if !Debug {
		log.SetOutput(ioutil.Discard)
	}
	SetIconImage(ICON_PRIMARY, imageData)
	initialize(title)
	SetIcon(ICON_PRIMARY)
	setMenu(items)
	return
}

func SetIconImage(iconId int, imageData []byte) (err error) {
	iconPth, err := createTempFile(imageData)
	if err != nil {
		return
	}
	icons[iconId] = iconPth
	return
}

func EnterLoop() {
	log.Println("Entering native loop")
	go menuUpdateLoop()
	C.native_loop()
	// If reached, user clicked Exit
	isExiting = true
}

func Exit() {
	log.Println("Exiting native loop")
	close(UpdateCh)
	C.exit_loop()
	cleanup()
}

func SetIcon(iconId int) (err error) {

	if iconId != curIconId {
		iconPth, ok := icons[iconId]
		if !ok {
			err = fmt.Errorf("No icon with icon id %d", iconId)
			return
		}

		log.Printf("Setting icon %s (id: %d)", iconPth, iconId)
		setIcon(iconPth)
		curIconId = iconId
	}
	return
}

func SetClickHandler(handler func()) {
	clickHandler = handler
}

func menuUpdateLoop() {
	runtime.UnlockOSThread()
	for update := range UpdateCh {
		updateMenuItem(update.ItemId, update.Item)
	}
}

func updateMenuItem(id int, item MenuItem) {
	menuItems[id] = item
	setMenuItem(id, item)
}

func setMenu(menu MenuItems) {
	menuItems = menu
	menuItemOrder := make([]int, 0, len(menuItems))
	for key, _ := range menuItems {
		menuItemOrder = append(menuItemOrder, key)
	}
	sort.Ints(menuItemOrder)

	for id := range menuItemOrder {
		item := menuItems[id]
		setMenuItem(id, item)
	}
}

func createTempFile(iconData []byte) (filename string, err error) {
	file, err := ioutil.TempFile(os.TempDir(), "trayhosticon")
	if err != nil {
		return
	}
	defer file.Close()
	_, err = file.Write(iconData)
	filename = file.Name()
	tmpFiles = append(tmpFiles, filename)
	log.Printf("Temp file created: %s\n", filename)
	return
}

func cleanup() {
	for _, file := range tmpFiles {
		err := os.Remove(file)
		if err != nil {
			log.Printf("Failed to remove tmp file %s: %v\n", file, err)
		} else {
			log.Printf("Tmp file %s removed\n", file)
		}
	}
}

func cbool(b bool) C.int {
	if b {
		return 1
	} else {
		return 0
	}
}
<|endoftext|>