<commit_before>package menu

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"sort"
	"strings"

	"github.com/fatih/color"
	"github.com/miquella/ask"
	"golang.org/x/crypto/ssh"

	"github.com/miquella/vaulted/lib"
)

type SSHKeyMenu struct {
	*Menu
}

func (m *SSHKeyMenu) Help() {
	menuColor.Set()
	defer color.Unset()

	fmt.Println("a,add      - Add")
	fmt.Println("D,delete   - Delete")
	fmt.Println("g,generate - Generate Key")
	fmt.Println("v          - HashiCorp Vault Signing URL")
	fmt.Println("u,users    - HashiCorp Vault User Principals")
	fmt.Println("E          - Expose External SSH Agent")
	fmt.Println("?,help     - Help")
	fmt.Println("b,back     - Back")
	fmt.Println("q,quit     - Quit")
}

func (m *SSHKeyMenu) Handler() error {
	for {
		var err error
		m.Printer()
		input, err := interaction.ReadMenu("Edit ssh keys: [a,D,g,v,u,E,b]: ")
		if err != nil {
			return err
		}
		switch input {
		case "a", "add", "key", "keys":
			err = m.AddSSHKey()
		case "D", "delete", "remove":
			var key string
			key, err = interaction.ReadValue("Key: ")
			if err == nil {
				if _, exists := m.Vault.SSHKeys[key]; exists {
					delete(m.Vault.SSHKeys, key)
				} else {
					color.Red("Key '%s' not found", key)
				}
			}
		case "g", "generate":
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.GenerateRSAKey = !m.Vault.SSHOptions.GenerateRSAKey
		case "v":
			signingUrl, err := interaction.ReadValue("HashiCorp Vault signing URL: ")
			if err != nil {
				return err
			}
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.VaultSigningUrl = signingUrl

			if signingUrl != "" && !m.Vault.SSHOptions.GenerateRSAKey {
				generateKey, _ := interaction.ReadValue("Would you like to enable RSA key generation (y/n): ")
				if generateKey == "y" {
					m.Vault.SSHOptions.GenerateRSAKey = true
				}
			}
		case "u", "users":
			userPrincipals, err := interaction.ReadValue("HashiCorp Vault user principals (comma separated): ")
			if err != nil {
				return err
			}
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			if userPrincipals != "" {
				m.Vault.SSHOptions.ValidPrincipals = strings.Split(userPrincipals, ",")
			} else {
				m.Vault.SSHOptions.ValidPrincipals = []string{}
			}
		case "E":
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.DisableProxy = !m.Vault.SSHOptions.DisableProxy
		case "b", "back":
			return nil
		case "q", "quit", "exit":
			var confirm string
			confirm, err = interaction.ReadValue("Are you sure you wish to save and exit the vault? (y/n): ")
			if err == nil {
				if confirm == "y" {
					return ErrSaveAndExit
				}
			}
		case "?", "help":
			m.Help()
		default:
			color.Red("Command not recognized")
		}

		if err != nil {
			return err
		}
	}
}

func (m *SSHKeyMenu) AddSSHKey() error {
	var err error

	homeDir := ""
	user, err := user.Current()
	if err == nil {
		homeDir = user.HomeDir
	} else {
		homeDir = os.Getenv("HOME")
	}

	defaultFilename := ""
	filename := ""
	if homeDir != "" {
		defaultFilename = filepath.Join(homeDir, ".ssh", "id_rsa")
		filename, err = interaction.ReadValue(fmt.Sprintf("Key file (default: %s): ", defaultFilename))
		if err != nil {
			return err
		}
		if filename == "" {
			filename = defaultFilename
		}
		if !filepath.IsAbs(filename) {
			filename = filepath.Join(filepath.Join(homeDir, ".ssh"), filename)
		}
	} else {
		filename, err = interaction.ReadValue("Key file: ")
		if err != nil {
			return err
		}
	}

	decryptedBlock, err := loadAndDecryptKey(filename)
	if err != nil {
		color.Red("%v", err)
		return nil
	}

	comment := loadPublicKeyComment(filename + ".pub")
	var name string
	if comment != "" {
		name, err = interaction.ReadValue(fmt.Sprintf("Name (default: %s): ", comment))
		if err != nil {
			return err
		}
		if name == "" {
			name = comment
		}
	} else {
		name, err = interaction.ReadValue("Name: ")
		if err != nil {
			return err
		}
		if name == "" {
			name = filename
		}
	}

	if m.Vault.SSHKeys == nil {
		m.Vault.SSHKeys = make(map[string]string)
	}
	m.Vault.SSHKeys[name] = string(pem.EncodeToMemory(decryptedBlock))

	return nil
}

func loadAndDecryptKey(filename string) (*pem.Block, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(data)
	if block == nil {
		return nil, err
	}

	if x509.IsEncryptedPEMBlock(block) {
		var passphrase string
		var decryptedBytes []byte
		for i := 0; i < 3; i++ {
			passphrase, err = ask.HiddenAsk("Passphrase: ")
			if err != nil {
				return nil, err
			}

			decryptedBytes, err = x509.DecryptPEMBlock(block, []byte(passphrase))
			if err == nil {
				break
			}
			if err != x509.IncorrectPasswordError {
				return nil, err
			}
		}

		if err != nil {
			return nil, err
		}

		return &pem.Block{
			Type:  block.Type,
			Bytes: decryptedBytes,
		}, nil
	}
	return block, nil
}

func loadPublicKeyComment(filename string) string {
	f, err := os.Open(filename)
	if err != nil {
		return ""
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	if err != nil {
		return ""
	}

	_, comment, _, _, err := ssh.ParseAuthorizedKey(data)
	if err != nil {
		return ""
	}
	return comment
}

func (m *SSHKeyMenu) Printer() {
	color.Cyan("\nSSH Agent:")
	color.Cyan("  Keys:")
	if len(m.Vault.SSHKeys) > 0 || m.Vault.SSHOptions != nil && m.Vault.SSHOptions.GenerateRSAKey {
		keys := []string{}
		for key := range m.Vault.SSHKeys {
			keys = append(keys, key)
		}
		sort.Strings(keys)

		for _, key := range keys {
			green.Printf("    %s\n", key)
		}

		if m.Vault.SSHOptions != nil && m.Vault.SSHOptions.GenerateRSAKey {
			faintColor.Print("    <generated RSA key>\n")
		}
	} else {
		fmt.Println("    [Empty]")
	}

	if m.Vault.SSHOptions != nil {
		if m.Vault.SSHOptions.VaultSigningUrl != "" || len(m.Vault.SSHOptions.ValidPrincipals) > 0 {
			color.Cyan("\n  Signing (HashiCorp Vault):")
			if m.Vault.SSHOptions.VaultSigningUrl != "" {
				green.Printf("    URL: ")
				fmt.Printf("%s\n", m.Vault.SSHOptions.VaultSigningUrl)
			}

			if len(m.Vault.SSHOptions.ValidPrincipals) > 0 {
				green.Printf("    User(s): ")
				fmt.Printf("%s\n", strings.Join(m.Vault.SSHOptions.ValidPrincipals, ", "))
			}
		}

		cyan.Print("\n  Expose external SSH agent: ")
		fmt.Printf("%t\n", !m.Vault.SSHOptions.DisableProxy)
	}
}
<commit_msg>remove key generate prompt if keys exist<commit_after>package menu

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"sort"
	"strings"

	"github.com/fatih/color"
	"github.com/miquella/ask"
	"golang.org/x/crypto/ssh"

	"github.com/miquella/vaulted/lib"
)

type SSHKeyMenu struct {
	*Menu
}

func (m *SSHKeyMenu) Help() {
	menuColor.Set()
	defer color.Unset()

	fmt.Println("a,add      - Add")
	fmt.Println("D,delete   - Delete")
	fmt.Println("g,generate - Generate Key")
	fmt.Println("v          - HashiCorp Vault Signing URL")
	fmt.Println("u,users    - HashiCorp Vault User Principals")
	fmt.Println("E          - Expose External SSH Agent")
	fmt.Println("?,help     - Help")
	fmt.Println("b,back     - Back")
	fmt.Println("q,quit     - Quit")
}

func (m *SSHKeyMenu) Handler() error {
	for {
		var err error
		m.Printer()
		input, err := interaction.ReadMenu("Edit ssh keys: [a,D,g,v,u,E,b]: ")
		if err != nil {
			return err
		}
		switch input {
		case "a", "add", "key", "keys":
			err = m.AddSSHKey()
		case "D", "delete", "remove":
			var key string
			key, err = interaction.ReadValue("Key: ")
			if err == nil {
				if _, exists := m.Vault.SSHKeys[key]; exists {
					delete(m.Vault.SSHKeys, key)
				} else {
					color.Red("Key '%s' not found", key)
				}
			}
		case "g", "generate":
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.GenerateRSAKey = !m.Vault.SSHOptions.GenerateRSAKey
		case "v":
			signingUrl, err := interaction.ReadValue("HashiCorp Vault signing URL: ")
			if err != nil {
				return err
			}
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.VaultSigningUrl = signingUrl

			if signingUrl != "" && !m.Vault.SSHOptions.GenerateRSAKey && len(m.Vault.SSHKeys) == 0 {
				generateKey, _ := interaction.ReadValue("Would you like to enable RSA key generation (y/n): ")
				if generateKey == "y" {
					m.Vault.SSHOptions.GenerateRSAKey = true
				}
			}
		case "u", "users":
			userPrincipals, err := interaction.ReadValue("HashiCorp Vault user principals (comma separated): ")
			if err != nil {
				return err
			}
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			if userPrincipals != "" {
				m.Vault.SSHOptions.ValidPrincipals = strings.Split(userPrincipals, ",")
			} else {
				m.Vault.SSHOptions.ValidPrincipals = []string{}
			}
		case "E":
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.DisableProxy = !m.Vault.SSHOptions.DisableProxy
		case "b", "back":
			return nil
		case "q", "quit", "exit":
			var confirm string
			confirm, err = interaction.ReadValue("Are you sure you wish to save and exit the vault? (y/n): ")
			if err == nil {
				if confirm == "y" {
					return ErrSaveAndExit
				}
			}
		case "?", "help":
			m.Help()
		default:
			color.Red("Command not recognized")
		}

		if err != nil {
			return err
		}
	}
}

func (m *SSHKeyMenu) AddSSHKey() error {
	var err error

	homeDir := ""
	user, err := user.Current()
	if err == nil {
		homeDir = user.HomeDir
	} else {
		homeDir = os.Getenv("HOME")
	}

	defaultFilename := ""
	filename := ""
	if homeDir != "" {
		defaultFilename = filepath.Join(homeDir, ".ssh", "id_rsa")
		filename, err = interaction.ReadValue(fmt.Sprintf("Key file (default: %s): ", defaultFilename))
		if err != nil {
			return err
		}
		if filename == "" {
			filename = defaultFilename
		}
		if !filepath.IsAbs(filename) {
			filename = filepath.Join(filepath.Join(homeDir, ".ssh"), filename)
		}
	} else {
		filename, err = interaction.ReadValue("Key file: ")
		if err != nil {
			return err
		}
	}

	decryptedBlock, err := loadAndDecryptKey(filename)
	if err != nil {
		color.Red("%v", err)
		return nil
	}

	comment := loadPublicKeyComment(filename + ".pub")
	var name string
	if comment != "" {
		name, err = interaction.ReadValue(fmt.Sprintf("Name (default: %s): ", comment))
		if err != nil {
			return err
		}
		if name == "" {
			name = comment
		}
	} else {
		name, err = interaction.ReadValue("Name: ")
		if err != nil {
			return err
		}
		if name == "" {
			name = filename
		}
	}

	if m.Vault.SSHKeys == nil {
		m.Vault.SSHKeys = make(map[string]string)
	}
	m.Vault.SSHKeys[name] = string(pem.EncodeToMemory(decryptedBlock))

	return nil
}

func loadAndDecryptKey(filename string) (*pem.Block, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(data)
	if block == nil {
		return nil, err
	}

	if x509.IsEncryptedPEMBlock(block) {
		var passphrase string
		var decryptedBytes []byte
		for i := 0; i < 3; i++ {
			passphrase, err = ask.HiddenAsk("Passphrase: ")
			if err != nil {
				return nil, err
			}

			decryptedBytes, err = x509.DecryptPEMBlock(block, []byte(passphrase))
			if err == nil {
				break
			}
			if err != x509.IncorrectPasswordError {
				return nil, err
			}
		}

		if err != nil {
			return nil, err
		}

		return &pem.Block{
			Type:  block.Type,
			Bytes: decryptedBytes,
		}, nil
	}
	return block, nil
}

func loadPublicKeyComment(filename string) string {
	f, err := os.Open(filename)
	if err != nil {
		return ""
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	if err != nil {
		return ""
	}

	_, comment, _, _, err := ssh.ParseAuthorizedKey(data)
	if err != nil {
		return ""
	}
	return comment
}

func (m *SSHKeyMenu) Printer() {
	color.Cyan("\nSSH Agent:")
	color.Cyan("  Keys:")
	if len(m.Vault.SSHKeys) > 0 || m.Vault.SSHOptions != nil && m.Vault.SSHOptions.GenerateRSAKey {
		keys := []string{}
		for key := range m.Vault.SSHKeys {
			keys = append(keys, key)
		}
		sort.Strings(keys)

		for _, key := range keys {
			green.Printf("    %s\n", key)
		}

		if m.Vault.SSHOptions != nil && m.Vault.SSHOptions.GenerateRSAKey {
			faintColor.Print("    <generated RSA key>\n")
		}
	} else {
		fmt.Println("    [Empty]")
	}

	if m.Vault.SSHOptions != nil {
		if m.Vault.SSHOptions.VaultSigningUrl != "" || len(m.Vault.SSHOptions.ValidPrincipals) > 0 {
			color.Cyan("\n  Signing (HashiCorp Vault):")
			if m.Vault.SSHOptions.VaultSigningUrl != "" {
				green.Printf("    URL: ")
				fmt.Printf("%s\n", m.Vault.SSHOptions.VaultSigningUrl)
			}

			if len(m.Vault.SSHOptions.ValidPrincipals) > 0 {
				green.Printf("    User(s): ")
				fmt.Printf("%s\n", strings.Join(m.Vault.SSHOptions.ValidPrincipals, ", "))
			}
		}

		cyan.Print("\n  Expose external SSH agent: ")
		fmt.Printf("%t\n", !m.Vault.SSHOptions.DisableProxy)
	}
}
<|endoftext|>


<commit_before>package merkledag

import (
	"fmt"

	"gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"

	key "github.com/ipfs/go-ipfs/blocks/key"
	mh "gx/ipfs/QmYf7ng2hG5XBtJA3tN34DQ2GUN5HNksEw1rLDkmr6vGku/go-multihash"
)

var ErrLinkNotFound = fmt.Errorf("no link by that name")

// Node represents a node in the IPFS Merkle DAG.
// nodes have opaque data and a set of navigable links.
type Node struct {
	Links []*Link
	Data  []byte

	// cache encoded/marshaled value
	encoded []byte

	cached mh.Multihash
}

// NodeStat is a statistics object for a Node. Mostly sizes.
type NodeStat struct {
	Hash           string
	NumLinks       int // number of links in link table
	BlockSize      int // size of the raw, encoded data
	LinksSize      int // size of the links segment
	DataSize       int // size of the data segment
	CumulativeSize int // cumulative size of object and its references
}

func (ns NodeStat) String() string {
	f := "NodeStat{NumLinks: %d, BlockSize: %d, LinksSize: %d, DataSize: %d, CumulativeSize: %d}"
	return fmt.Sprintf(f, ns.NumLinks, ns.BlockSize, ns.LinksSize, ns.DataSize, ns.CumulativeSize)
}

// Link represents an IPFS Merkle DAG Link between Nodes.
type Link struct {
	// utf string name. should be unique per object
	Name string // utf8

	// cumulative size of target object
	Size uint64

	// multihash of the target object
	Hash mh.Multihash

	// a ptr to the actual node for graph manipulation
	node *Node
}

type LinkSlice []*Link

func (ls LinkSlice) Len() int           { return len(ls) }
func (ls LinkSlice) Swap(a, b int)      { ls[a], ls[b] = ls[b], ls[a] }
func (ls LinkSlice) Less(a, b int) bool { return ls[a].Name < ls[b].Name }

// MakeLink creates a link to the given node
func MakeLink(n *Node) (*Link, error) {
	s, err := n.Size()
	if err != nil {
		return nil, err
	}

	h, err := n.Multihash()
	if err != nil {
		return nil, err
	}
	return &Link{
		Size: s,
		Hash: h,
	}, nil
}

// GetCachedNode returns the MDAG Node that was cached, or nil
func (l *Link) GetCachedNode() *Node {
	return l.node
}

// GetNode returns the MDAG Node that this link points to
func (l *Link) GetNode(ctx context.Context, serv DAGService) (*Node, error) {
	if l.node != nil {
		return l.node, nil
	}

	return serv.Get(ctx, key.Key(l.Hash))
}

// AddNodeLink adds a link to another node.
func (n *Node) AddNodeLink(name string, that *Node) error {
	n.encoded = nil

	lnk, err := MakeLink(that)

	lnk.Name = name
	lnk.node = that
	if err != nil {
		return err
	}

	n.AddRawLink(name, lnk)

	return nil
}

// AddNodeLinkClean adds a link to another node. without keeping a reference to
// the child node
func (n *Node) AddNodeLinkClean(name string, that *Node) error {
	n.encoded = nil
	lnk, err := MakeLink(that)
	if err != nil {
		return err
	}
	n.AddRawLink(name, lnk)

	return nil
}

// AddRawLink adds a copy of a link to this node
func (n *Node) AddRawLink(name string, l *Link) error {
	n.encoded = nil
	n.Links = append(n.Links, &Link{
		Name: name,
		Size: l.Size,
		Hash: l.Hash,
		node: l.node,
	})

	return nil
}

// Remove a link on this node by the given name
func (n *Node) RemoveNodeLink(name string) error {
	n.encoded = nil
	good := make([]*Link, 0, len(n.Links))
	var found bool

	for _, l := range n.Links {
		if l.Name != name {
			good = append(good, l)
		} else {
			found = true
		}
	}
	n.Links = good

	if !found {
		return ErrNotFound
	}

	return nil
}

// Return a copy of the link with given name
func (n *Node) GetNodeLink(name string) (*Link, error) {
	for _, l := range n.Links {
		if l.Name == name {
			return &Link{
				Name: l.Name,
				Size: l.Size,
				Hash: l.Hash,
				node: l.node,
			}, nil
		}
	}
	return nil, ErrLinkNotFound
}

func (n *Node) GetLinkedNode(ctx context.Context, ds DAGService, name string) (*Node, error) {
	lnk, err := n.GetNodeLink(name)
	if err != nil {
		return nil, err
	}

	return lnk.GetNode(ctx, ds)
}

// Copy returns a copy of the node.
// NOTE: does not make copies of Node objects in the links.
func (n *Node) Copy() *Node {
	nnode := new(Node)
	if len(n.Data) > 0 {
		nnode.Data = make([]byte, len(n.Data))
		copy(nnode.Data, n.Data)
	}

	if len(n.Links) > 0 {
		nnode.Links = make([]*Link, len(n.Links))
		copy(nnode.Links, n.Links)
	}
	return nnode
}

// UpdateNodeLink return a copy of the node with the link name set to point to
// that. If a link of the same name existed, it is removed.
func (n *Node) UpdateNodeLink(name string, that *Node) (*Node, error) {
	newnode := n.Copy()
	err := newnode.RemoveNodeLink(name)
	err = nil // ignore error
	err = newnode.AddNodeLink(name, that)
	return newnode, err
}

// Size returns the total size of the data addressed by node,
// including the total sizes of references.
func (n *Node) Size() (uint64, error) {
	b, err := n.Encoded(false)
	if err != nil {
		return 0, err
	}

	s := uint64(len(b))
	for _, l := range n.Links {
		s += l.Size
	}
	return s, nil
}

// Stat returns statistics on the node.
func (n *Node) Stat() (*NodeStat, error) {
	enc, err := n.Encoded(false)
	if err != nil {
		return nil, err
	}

	cumSize, err := n.Size()
	if err != nil {
		return nil, err
	}

	key, err := n.Key()
	if err != nil {
		return nil, err
	}

	return &NodeStat{
		Hash:           key.B58String(),
		NumLinks:       len(n.Links),
		BlockSize:      len(enc),
		LinksSize:      len(enc) - len(n.Data), // includes framing.
		DataSize:       len(n.Data),
		CumulativeSize: int(cumSize),
	}, nil
}

// Multihash hashes the encoded data of this node.
func (n *Node) Multihash() (mh.Multihash, error) {
	// Note: Encoded generates the hash and puts it in n.cached.
	_, err := n.Encoded(false)
	if err != nil {
		return nil, err
	}

	return n.cached, nil
}

// Key returns the Multihash as a key, for maps.
func (n *Node) Key() (key.Key, error) {
	h, err := n.Multihash()
	return key.Key(h), err
}
<commit_msg>Remove GetCachedNode()<commit_after>package merkledag

import (
	"fmt"

	"gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"

	key "github.com/ipfs/go-ipfs/blocks/key"
	mh "gx/ipfs/QmYf7ng2hG5XBtJA3tN34DQ2GUN5HNksEw1rLDkmr6vGku/go-multihash"
)

var ErrLinkNotFound = fmt.Errorf("no link by that name")

// Node represents a node in the IPFS Merkle DAG.
// nodes have opaque data and a set of navigable links.
type Node struct {
	Links []*Link
	Data  []byte

	// cache encoded/marshaled value
	encoded []byte

	cached mh.Multihash
}

// NodeStat is a statistics object for a Node. Mostly sizes.
type NodeStat struct {
	Hash           string
	NumLinks       int // number of links in link table
	BlockSize      int // size of the raw, encoded data
	LinksSize      int // size of the links segment
	DataSize       int // size of the data segment
	CumulativeSize int // cumulative size of object and its references
}

func (ns NodeStat) String() string {
	f := "NodeStat{NumLinks: %d, BlockSize: %d, LinksSize: %d, DataSize: %d, CumulativeSize: %d}"
	return fmt.Sprintf(f, ns.NumLinks, ns.BlockSize, ns.LinksSize, ns.DataSize, ns.CumulativeSize)
}

// Link represents an IPFS Merkle DAG Link between Nodes.
type Link struct {
	// utf string name. should be unique per object
	Name string // utf8

	// cumulative size of target object
	Size uint64

	// multihash of the target object
	Hash mh.Multihash

	// a ptr to the actual node for graph manipulation
	node *Node
}

type LinkSlice []*Link

func (ls LinkSlice) Len() int           { return len(ls) }
func (ls LinkSlice) Swap(a, b int)      { ls[a], ls[b] = ls[b], ls[a] }
func (ls LinkSlice) Less(a, b int) bool { return ls[a].Name < ls[b].Name }

// MakeLink creates a link to the given node
func MakeLink(n *Node) (*Link, error) {
	s, err := n.Size()
	if err != nil {
		return nil, err
	}

	h, err := n.Multihash()
	if err != nil {
		return nil, err
	}
	return &Link{
		Size: s,
		Hash: h,
	}, nil
}

// GetNode returns the MDAG Node that this link points to
func (l *Link) GetNode(ctx context.Context, serv DAGService) (*Node, error) {
	if l.node != nil {
		return l.node, nil
	}

	return serv.Get(ctx, key.Key(l.Hash))
}

// AddNodeLink adds a link to another node.
func (n *Node) AddNodeLink(name string, that *Node) error {
	n.encoded = nil

	lnk, err := MakeLink(that)

	lnk.Name = name
	lnk.node = that
	if err != nil {
		return err
	}

	n.AddRawLink(name, lnk)

	return nil
}

// AddNodeLinkClean adds a link to another node. without keeping a reference to
// the child node
func (n *Node) AddNodeLinkClean(name string, that *Node) error {
	n.encoded = nil
	lnk, err := MakeLink(that)
	if err != nil {
		return err
	}
	n.AddRawLink(name, lnk)

	return nil
}

// AddRawLink adds a copy of a link to this node
func (n *Node) AddRawLink(name string, l *Link) error {
	n.encoded = nil
	n.Links = append(n.Links, &Link{
		Name: name,
		Size: l.Size,
		Hash: l.Hash,
		node: l.node,
	})

	return nil
}

// Remove a link on this node by the given name
func (n *Node) RemoveNodeLink(name string) error {
	n.encoded = nil
	good := make([]*Link, 0, len(n.Links))
	var found bool

	for _, l := range n.Links {
		if l.Name != name {
			good = append(good, l)
		} else {
			found = true
		}
	}
	n.Links = good

	if !found {
		return ErrNotFound
	}

	return nil
}

// Return a copy of the link with given name
func (n *Node) GetNodeLink(name string) (*Link, error) {
	for _, l := range n.Links {
		if l.Name == name {
			return &Link{
				Name: l.Name,
				Size: l.Size,
				Hash: l.Hash,
				node: l.node,
			}, nil
		}
	}
	return nil, ErrLinkNotFound
}

func (n *Node) GetLinkedNode(ctx context.Context, ds DAGService, name string) (*Node, error) {
	lnk, err := n.GetNodeLink(name)
	if err != nil {
		return nil, err
	}

	return lnk.GetNode(ctx, ds)
}

// Copy returns a copy of the node.
// NOTE: does not make copies of Node objects in the links.
func (n *Node) Copy() *Node {
	nnode := new(Node)
	if len(n.Data) > 0 {
		nnode.Data = make([]byte, len(n.Data))
		copy(nnode.Data, n.Data)
	}

	if len(n.Links) > 0 {
		nnode.Links = make([]*Link, len(n.Links))
		copy(nnode.Links, n.Links)
	}
	return nnode
}

// UpdateNodeLink return a copy of the node with the link name set to point to
// that. If a link of the same name existed, it is removed.
func (n *Node) UpdateNodeLink(name string, that *Node) (*Node, error) {
	newnode := n.Copy()
	err := newnode.RemoveNodeLink(name)
	err = nil // ignore error
	err = newnode.AddNodeLink(name, that)
	return newnode, err
}

// Size returns the total size of the data addressed by node,
// including the total sizes of references.
func (n *Node) Size() (uint64, error) {
	b, err := n.Encoded(false)
	if err != nil {
		return 0, err
	}

	s := uint64(len(b))
	for _, l := range n.Links {
		s += l.Size
	}
	return s, nil
}

// Stat returns statistics on the node.
func (n *Node) Stat() (*NodeStat, error) {
	enc, err := n.Encoded(false)
	if err != nil {
		return nil, err
	}

	cumSize, err := n.Size()
	if err != nil {
		return nil, err
	}

	key, err := n.Key()
	if err != nil {
		return nil, err
	}

	return &NodeStat{
		Hash:           key.B58String(),
		NumLinks:       len(n.Links),
		BlockSize:      len(enc),
		LinksSize:      len(enc) - len(n.Data), // includes framing.
		DataSize:       len(n.Data),
		CumulativeSize: int(cumSize),
	}, nil
}

// Multihash hashes the encoded data of this node.
func (n *Node) Multihash() (mh.Multihash, error) {
	// Note: Encoded generates the hash and puts it in n.cached.
	_, err := n.Encoded(false)
	if err != nil {
		return nil, err
	}

	return n.cached, nil
}

// Key returns the Multihash as a key, for maps.
func (n *Node) Key() (key.Key, error) {
	h, err := n.Multihash()
	return key.Key(h), err
}
<|endoftext|>


<commit_before>// Package bulk has functions to load and save image metadata to a CSV file as
// needed by the 'Bulk Add CSV' Shimmie2 extension. The CSV file is assumed to
// have the following format:
//
//  "/path/to/image.jpg","spaced tags","source","rating s/q/e","/path/thumbnail.jpg"
//
// The last record (thumbnail) is left empty as thumbnails can easily be
// generated by the server.
//
// The package assumes that all images and the CSV file are under a certain
// directory path that is used as input in many package functions.
package bulk

import (
	"encoding/csv"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// Image holds the metadata of each image from the CSV file.
type Image struct {
	ID     int
	Name   string
	Tags   []string
	Source string
	Rating string
}

var supportedExt = []string{"gif", "jpeg", "jpg", "png", "swf"}

func isSupportedType(name string) bool {
	fname := strings.ToLower(name)
	for _, ext := range supportedExt {
		// The only possible returned error is ErrBadPattern, when pattern is
		// malformed. Patterns like *.jpg are never malformed so we ignore the
		// error.
		matches, _ := filepath.Match("*."+ext, fname)
		if matches {
			return true
		}
	}
	return false
}

// LoadImages expects a slice of directory entries (os.FileInfo) which is the
// result of a read directory like ioutil.ReadDir. It loops through the slice,
// ignoring any directory and keeps only the files with one of the following
// extensions: "gif", "jpeg", "jpg", "png", "swf"
//
// It returns a slice of images without metadata, using the filename as Name
// and the order the files were found as an increasing ID starting from 0.
//
// In case of a CSV file, the image metadata should be read using LoadCSV and
// then combined with the images (discovered by LoadImages) using Combine.
func LoadImages(files []os.FileInfo) []Image {
	images := []Image{}

	id := 0
	for _, f := range files {
		if !f.IsDir() {
			if isSupportedType(f.Name()) {
				img := Image{ID: id, Name: f.Name()}
				images = append(images, img)
				id++
			}
		}
	}
	return images
}

// LoadCSV loads the image metadata from a CSV file that is open for reading.
// The metadata are returned as slice of images and should be combined with the
// slice of images discovered by LoadImages by calling Combine.
func LoadCSV(file io.Reader) ([]Image, error) {
	images := []Image{}

	r := csv.NewReader(file)
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		if len(record) != 5 {
			return nil, fmt.Errorf("invalid csv file format")
		}
		// Image filepath (first column) should exist otherwise we cannot match
		// the metadata with the images found under the directory.
		if record[0] != "" {
			img := Image{
				Name:   filepath.Base(record[0]),
				Tags:   strings.Split(record[1], " "),
				Source: record[2],
				Rating: record[3],
			}
			images = append(images, img)
		}
	}
	return images, nil
}

// Combine takes the metadata of imagesWithInfo and copies them to images
// returning the combined result.
func Combine(images, imagesWithInfo []Image) []Image {
	for _, info := range imagesWithInfo {
		if info.Name == "" {
			continue
		}
		img := findByName(images, info.Name)
		if img != nil {
			img.Source = info.Source
			img.Rating = info.Rating
			img.Tags = info.Tags
		}
	}
	return images
}

type byName []Image

func (img byName) Len() int           { return len(img) }
func (img byName) Swap(i, j int)      { img[i], img[j] = img[j], img[i] }
func (img byName) Less(i, j int) bool { return img[i].Name < img[j].Name }

func findByName(image []Image, name string) *Image {
	sort.Sort(byName(image))
	i := sort.Search(len(image), func(i int) bool { return image[i].Name >= name })
	if i < len(image) && image[i].Name == name {
		return &image[i]
	}
	return nil
}

type byID []Image

func (img byID) Len() int           { return len(img) }
func (img byID) Swap(i, j int)      { img[i], img[j] = img[j], img[i] }
func (img byID) Less(i, j int) bool { return img[i].ID < img[j].ID }

// FindByID takes a slice of images, sorts them by ID and then finds then one
// with ID id.
func FindByID(image []Image, id int) *Image {
	sort.Sort(byID(image))
	i := sort.Search(len(image), func(i int) bool { return image[i].ID >= id })
	if i < len(image) && image[i].ID == id {
		return &image[i]
	}
	return nil
}

// CurrentPrefix reads from an open CSV file, reads the first line and uses the
// base directory of the provided dir path to find the path prefix used in the
// CSV file.
//
// As an example if the provided dir path is '/localpath/pics' and the first line
// has '/serverpath/pics/pic1' then the returned current prefix will be
// '/serverpath'.
func CurrentPrefix(workingDir string, file io.Reader) (string, error) {
	r := csv.NewReader(file)
	record, err := r.Read()
	if err == io.EOF {
		return "", nil
	}
	if err != nil {
		return "", err
	}

	serverDir := record[0]
	picFolder := filepath.Base(workingDir)
	sep := fmt.Sprintf("%c", filepath.Separator)
	if !strings.Contains(serverDir, picFolder) {
		return sep, nil
	}
	for {
		if filepath.Base(serverDir) == picFolder {
			break
		} else {
			serverDir, _ = filepath.Split(serverDir)
			serverDir = filepath.Dir(serverDir)
		}
	}
	return filepath.Dir(serverDir), nil
}

// Save will write the image metadata to an open for writing file. It will
// keep the base of the dir path and replace the prefix with the provided one.
func Save(file io.Writer, images []Image, dir, prefix string, useLinuxSep bool) error {
	w := csv.NewWriter(file)
	w.WriteAll(toRecords(images, dir, prefix, useLinuxSep))

	if err := w.Error(); err != nil {
		return fmt.Errorf("error writing csv: %v", err)
	}
	return nil
}

func toRecords(images []Image, dir, prefix string, useLinuxSep bool) [][]string {
	var records [][]string
	for _, img := range images {
		record := toRecord(img, dir, prefix, useLinuxSep)
		records = append(records, record)
	}
	return records
}

func toRecord(img Image, dir, prefix string, useLinuxSep bool) []string {
	var record []string
	p := filepath.Join(prefix, filepath.Base(dir), img.Name)
	if useLinuxSep {
		p = strings.Replace(p, "\\", "/", -1)
	}
	record = append(record, p)
	record = append(record, strings.Join(img.Tags, " "))
	record = append(record, img.Source)
	record = append(record, img.Rating)
	record = append(record, "")
	return record
}
<commit_msg>Rename record to firstLine in CurrentPrefix<commit_after>// Package bulk has functions to load and save image metadata to a CSV file as
// needed by the 'Bulk Add CSV' Shimmie2 extension. The CSV file is assumed to
// have the following format:
//
//  "/path/to/image.jpg","spaced tags","source","rating s/q/e","/path/thumbnail.jpg"
//
// The last record (thumbnail) is left empty as thumbnails can easily be
// generated by the server.
//
// The package assumes that all images and the CSV file are under a certain
// directory path that is used as input in many package functions.
package bulk

import (
	"encoding/csv"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// Image holds the metadata of each image from the CSV file.
type Image struct {
	ID     int
	Name   string
	Tags   []string
	Source string
	Rating string
}

var supportedExt = []string{"gif", "jpeg", "jpg", "png", "swf"}

func isSupportedType(name string) bool {
	fname := strings.ToLower(name)
	for _, ext := range supportedExt {
		// The only possible returned error is ErrBadPattern, when pattern is
		// malformed. Patterns like *.jpg are never malformed so we ignore the
		// error.
		matches, _ := filepath.Match("*."+ext, fname)
		if matches {
			return true
		}
	}
	return false
}

// LoadImages expects a slice of directory entries (os.FileInfo) which is the
// result of a read directory like ioutil.ReadDir. It loops through the slice,
// ignoring any directory and keeps only the files with one of the following
// extensions: "gif", "jpeg", "jpg", "png", "swf"
//
// It returns a slice of images without metadata, using the filename as Name
// and the order the files were found as an increasing ID starting from 0.
//
// In case of a CSV file, the image metadata should be read using LoadCSV and
// then combined with the images (discovered by LoadImages) using Combine.
func LoadImages(files []os.FileInfo) []Image {
	images := []Image{}

	id := 0
	for _, f := range files {
		if !f.IsDir() {
			if isSupportedType(f.Name()) {
				img := Image{ID: id, Name: f.Name()}
				images = append(images, img)
				id++
			}
		}
	}
	return images
}

// LoadCSV loads the image metadata from a CSV file that is open for reading.
// The metadata are returned as slice of images and should be combined with the
// slice of images discovered by LoadImages by calling Combine.
func LoadCSV(file io.Reader) ([]Image, error) {
	images := []Image{}

	r := csv.NewReader(file)
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		if len(record) != 5 {
			return nil, fmt.Errorf("invalid csv file format")
		}
		// Image filepath (first column) should exist otherwise we cannot match
		// the metadata with the images found under the directory.
		if record[0] != "" {
			img := Image{
				Name:   filepath.Base(record[0]),
				Tags:   strings.Split(record[1], " "),
				Source: record[2],
				Rating: record[3],
			}
			images = append(images, img)
		}
	}
	return images, nil
}

// Combine takes the metadata of imagesWithInfo and copies them to images
// returning the combined result.
func Combine(images, imagesWithInfo []Image) []Image {
	for _, info := range imagesWithInfo {
		if info.Name == "" {
			continue
		}
		img := findByName(images, info.Name)
		if img != nil {
			img.Source = info.Source
			img.Rating = info.Rating
			img.Tags = info.Tags
		}
	}
	return images
}

type byName []Image

func (img byName) Len() int           { return len(img) }
func (img byName) Swap(i, j int)      { img[i], img[j] = img[j], img[i] }
func (img byName) Less(i, j int) bool { return img[i].Name < img[j].Name }

func findByName(image []Image, name string) *Image {
	sort.Sort(byName(image))
	i := sort.Search(len(image), func(i int) bool { return image[i].Name >= name })
	if i < len(image) && image[i].Name == name {
		return &image[i]
	}
	return nil
}

type byID []Image

func (img byID) Len() int           { return len(img) }
func (img byID) Swap(i, j int)      { img[i], img[j] = img[j], img[i] }
func (img byID) Less(i, j int) bool { return img[i].ID < img[j].ID }

// FindByID takes a slice of images, sorts them by ID and then finds then one
// with ID id.
func FindByID(image []Image, id int) *Image {
	sort.Sort(byID(image))
	i := sort.Search(len(image), func(i int) bool { return image[i].ID >= id })
	if i < len(image) && image[i].ID == id {
		return &image[i]
	}
	return nil
}

// CurrentPrefix reads from an open CSV file, reads the first line and uses the
// base directory of the provided dir path to find the path prefix used in the
// CSV file.
//
// As an example if the provided dir path is '/localpath/pics' and the first line
// has '/serverpath/pics/pic1' then the returned current prefix will be
// '/serverpath'.
func CurrentPrefix(workingDir string, file io.Reader) (string, error) {
	r := csv.NewReader(file)
	firstLine, err := r.Read()
	if err == io.EOF {
		return "", nil
	}
	if err != nil {
		return "", err
	}

	serverDir := firstLine[0]
	picFolder := filepath.Base(workingDir)
	sep := fmt.Sprintf("%c", filepath.Separator)
	if !strings.Contains(serverDir, picFolder) {
		return sep, nil
	}
	for {
		if filepath.Base(serverDir) == picFolder {
			break
		} else {
			serverDir, _ = filepath.Split(serverDir)
			serverDir = filepath.Dir(serverDir)
		}
	}
	return filepath.Dir(serverDir), nil
}

// Save will write the image metadata to an open for writing file. It will
// keep the base of the dir path and replace the prefix with the provided one.
func Save(file io.Writer, images []Image, dir, prefix string, useLinuxSep bool) error {
	w := csv.NewWriter(file)
	w.WriteAll(toRecords(images, dir, prefix, useLinuxSep))

	if err := w.Error(); err != nil {
		return fmt.Errorf("error writing csv: %v", err)
	}
	return nil
}

func toRecords(images []Image, dir, prefix string, useLinuxSep bool) [][]string {
	var records [][]string
	for _, img := range images {
		record := toRecord(img, dir, prefix, useLinuxSep)
		records = append(records, record)
	}
	return records
}

func toRecord(img Image, dir, prefix string, useLinuxSep bool) []string {
	var record []string
	p := filepath.Join(prefix, filepath.Base(dir), img.Name)
	if useLinuxSep {
		p = strings.Replace(p, "\\", "/", -1)
	}
	record = append(record, p)
	record = append(record, strings.Join(img.Tags, " "))
	record = append(record, img.Source)
	record = append(record, img.Rating)
	record = append(record, "")
	return record
}
<|endoftext|>


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

package addresser_test

import (
	stdtesting "testing"
	"time"

	jc "github.com/juju/testing/checkers"
	"github.com/juju/utils"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/juju/testing"
	"github.com/juju/juju/network"
	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/worker"
	"github.com/juju/juju/worker/addresser"
)

func TestPackage(t *stdtesting.T) {
	coretesting.MgoTestPackage(t)
}

var _ = gc.Suite(&workerSuite{})
var shortAttempt = utils.AttemptStrategy{
	Total: 5 * time.Second,
	Delay: 200 * time.Millisecond,
}

type workerSuite struct {
	testing.JujuConnSuite
}

func (s *workerSuite) SetUpTest(c *gc.C) {
	s.AssertConfigParameterUpdated(c, "broken", []string{})
}

func (s *workerSuite) createAddresses(c *gc.C) {
	addresses := [][]string{
		{"0.1.2.3", "wibble"},
		{"0.1.2.4", "wibble"},
		{"0.1.2.5", "wobble"},
		{"0.1.2.6", "wobble"},
	}
	for i, details := range addresses {
		addr := network.NewScopedAddress(details[0], network.ScopePublic)
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Assert(err, jc.ErrorIsNil)
		err = ipAddr.AllocateTo(details[1], "wobble")
		c.Assert(err, jc.ErrorIsNil)
		if i%2 == 1 {
			// two of the addresses start out Dead
			err = ipAddr.EnsureDead()
			c.Assert(err, jc.ErrorIsNil)
		}
	}
}

func (s *workerSuite) TestWorker(c *gc.C) {
	s.createAddresses(c)
	s.State.StartSync()
	w, err := addresser.NewWorker(s.State)
	c.Assert(err, jc.ErrorIsNil)
	defer func() {
		c.Assert(worker.Stop(w), gc.IsNil)
	}()

}
<commit_msg>Add fail<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package addresser_test

import (
	stdtesting "testing"
	"time"

	jc "github.com/juju/testing/checkers"
	"github.com/juju/utils"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/juju/testing"
	"github.com/juju/juju/network"
	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/worker"
	"github.com/juju/juju/worker/addresser"
)

func TestPackage(t *stdtesting.T) {
	coretesting.MgoTestPackage(t)
}

var _ = gc.Suite(&workerSuite{})
var shortAttempt = utils.AttemptStrategy{
	Total: 5 * time.Second,
	Delay: 200 * time.Millisecond,
}

type workerSuite struct {
	testing.JujuConnSuite
}

func (s *workerSuite) SetUpTest(c *gc.C) {
	s.AssertConfigParameterUpdated(c, "broken", []string{})
}

func (s *workerSuite) createAddresses(c *gc.C) {
	addresses := [][]string{
		{"0.1.2.3", "wibble"},
		{"0.1.2.4", "wibble"},
		{"0.1.2.5", "wobble"},
		{"0.1.2.6", "wobble"},
	}
	for i, details := range addresses {
		addr := network.NewScopedAddress(details[0], network.ScopePublic)
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Assert(err, jc.ErrorIsNil)
		err = ipAddr.AllocateTo(details[1], "wobble")
		c.Assert(err, jc.ErrorIsNil)
		if i%2 == 1 {
			// two of the addresses start out Dead
			err = ipAddr.EnsureDead()
			c.Assert(err, jc.ErrorIsNil)
		}
	}
}

func (s *workerSuite) TestWorker(c *gc.C) {
	s.createAddresses(c)
	s.State.StartSync()
	w, err := addresser.NewWorker(s.State)
	c.Assert(err, jc.ErrorIsNil)
	defer func() {
		c.Assert(worker.Stop(w), gc.IsNil)
	}()

	for a := shortAttempt.Start(); a.Next(); {
		dead, err := s.State.DeadIPAddresses()
		c.Assert(err, jc.ErrorIsNil)
		if len(dead) == 0 {
			break
		}
		if !a.HasNext() {
			c.Fail()
		}
	}
}
<|endoftext|>


<commit_before>package firewaller

import (
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/watcher"
	"launchpad.net/tomb"
)

// Firewaller watches the state for ports opened or closed
// and reflects those changes onto the backing environment.
type Firewaller struct {
	st                  *state.State
	tomb                tomb.Tomb
	machinesWatcher     *state.MachinesWatcher
	machineDatas        map[int]*machineData
	machineUnitsChanges chan *machineUnitsChange
	unitDatas           map[string]*unitData
	unitPortsChanges    chan *unitPortsChange
	serviceDatas        map[string]*serviceData
}

// NewFirewaller returns a new Firewaller.
func NewFirewaller(st *state.State) (*Firewaller, error) {
	fw := &Firewaller{
		st:                  st,
		machinesWatcher:     st.WatchMachines(),
		machineDatas:        make(map[int]*machineData),
		machineUnitsChanges: make(chan *machineUnitsChange),
		unitDatas:           make(map[string]*unitData),
		unitPortsChanges:    make(chan *unitPortsChange),
		serviceDatas:        make(map[string]*serviceData),
	}
	go fw.loop()
	return fw, nil
}

func (fw *Firewaller) loop() {
	defer fw.finish()
	for {
		select {
		case <-fw.tomb.Dying():
			return
		case change, ok := <-fw.machinesWatcher.Changes():
			if !ok {
				return
			}
			for _, machine := range change.Removed {
				machined, ok := fw.machineDatas[machine.Id()]
				if !ok {
					panic("trying to remove machine that wasn't added")
				}
				delete(fw.machineDatas, machine.Id())
				if err := machined.stopWatch(); err != nil {
					log.Printf("machine data %d returned error when stopping: %v", machine.Id(), err)
				}
				log.Debugf("firewaller: stopped watching machine %d", machine.Id())
			}
			for _, machine := range change.Added {
				machined := newMachineData(machine, fw)
				fw.machineDatas[machine.Id()] = machined
				log.Debugf("firewaller: started watching machine %d", machine.Id())
			}
		case change := <-fw.machineUnitsChanges:
			for _, unit := range change.Removed {
				unitd, ok := fw.unitDatas[unit.Name()]
				if !ok {
					panic("trying to remove unit that wasn't added")
				}
				delete(fw.unitDatas, unit.Name())
				// TODO(mue) Close ports.
				if err := unitd.stopWatch(); err != nil {
					log.Printf("unit data %s returned error when stopping: %v", unit.Name(), err)
				}
				log.Debugf("firewaller: stopped watching unit %s", unit.Name())
			}
			for _, unit := range change.Added {
				unitd := newUnitData(unit, fw)
				fw.unitDatas[unit.Name()] = unitd
				if fw.serviceDatas[unit.ServiceName()] == nil {
					// TODO(mue) Add service watcher.
				}
				log.Debugf("firewaller: started watching unit %s", unit.Name())
			}
		case <-fw.unitPortsChanges:
			// TODO(mue) Handle changes of ports.
		}
	}
}

// finishes cleans up when the firewaller is stopping.
func (fw *Firewaller) finish() {
	watcher.Stop(fw.machinesWatcher, &fw.tomb)
	for _, unitd := range fw.unitDatas {
		fw.tomb.Kill(unitd.stopWatch())
	}
	for _, machined := range fw.machineDatas {
		fw.tomb.Kill(machined.stopWatch())
	}
	fw.tomb.Done()
}

// Wait waits for the Firewaller to exit.
func (fw *Firewaller) Wait() error {
	return fw.tomb.Wait()
}

// Stop stops the Firewaller and returns any error encountered while stopping.
func (fw *Firewaller) Stop() error {
	fw.tomb.Kill(nil)
	return fw.tomb.Wait()
}

// machineUnitsChange contains the changed units for one specific machine.
type machineUnitsChange struct {
	machined *machineData
	*state.MachineUnitsChange
}

// machineData watches the unit changes of a machine and passes them
// to the firewaller for handling.
type machineData struct {
	tomb       tomb.Tomb
	firewaller *Firewaller
	machine    *state.Machine
	watcher    *state.MachineUnitsWatcher
}

// newMachineData starts the watching of the passed machine.
func newMachineData(machine *state.Machine, fw *Firewaller) *machineData {
	md := &machineData{
		firewaller: fw,
		machine:    machine,
		watcher:    machine.WatchUnits(),
	}
	go md.watchLoop()
	return md
}

// watchLoop is the backend watching for machine units changes.
func (md *machineData) watchLoop() {
	defer md.tomb.Done()
	defer md.watcher.Stop()
	for {
		select {
		case <-md.tomb.Dying():
			return
		case change, ok := <-md.watcher.Changes():
			if !ok {
				md.firewaller.tomb.Kill(watcher.MustErr(md.watcher))
				return
			}
			select {
			case md.firewaller.machineUnitsChanges <- &machineUnitsChange{md, change}:
			case <-md.tomb.Dying():
				return
			}
		}
	}
}

// stopWatch stops the machine watching.
func (md *machineData) stopWatch() error {
	md.tomb.Kill(nil)
	return md.tomb.Wait()
}

// unitPortsChange contains the changed ports for one specific unit.
type unitPortsChange struct {
	unitd *unitData
	ports []state.Port
}

// unitData watches the port changes of a unit and passes them
// to the firewaller for handling.
type unitData struct {
	tomb       tomb.Tomb
	firewaller *Firewaller
	unit       *state.Unit
	watcher    *state.PortsWatcher
	service    *serviceData
	ports      []state.Port
}

// newMachineData starts the watching of the passed unit.
func newUnitData(unit *state.Unit, fw *Firewaller) *unitData {
	ud := &unitData{
		firewaller: fw,
		unit:       unit,
		watcher:    unit.WatchPorts(),
		ports:      make([]state.Port, 0),
	}
	go ud.watchLoop()
	return ud
}

func (ud *unitData) watchLoop() {
	defer ud.tomb.Done()
	defer ud.watcher.Stop()
	for {
		select {
		case <-ud.tomb.Dying():
			return
		case change, ok := <-ud.watcher.Changes():
			if !ok {
				ud.firewaller.tomb.Kill(watcher.MustErr(ud.watcher))
				return
			}
			select {
			case ud.firewaller.unitPortsChanges <- &unitPortsChange{ud, change}:
			case <-ud.tomb.Dying():
				return
			}
		}
	}
}

// stopWatch stops the unit watching.
func (ud *unitData) stopWatch() error {
	ud.tomb.Kill(nil)
	return ud.tomb.Wait()
}

// serviceData watches the exposed flag changes of a service and passes them
// to the firewaller for handling.
type serviceData struct {
	// TODO(mue) Fill with life.
	service *state.Service
	exposed bool
}
<commit_msg>firewaller: smaller naming changes<commit_after>package firewaller

import (
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/watcher"
	"launchpad.net/tomb"
)

// Firewaller watches the state for ports opened or closed
// and reflects those changes onto the backing environment.
type Firewaller struct {
	st              *state.State
	tomb            tomb.Tomb
	machinesWatcher *state.MachinesWatcher
	machineds       map[int]*machineData
	unitsChange     chan *machineUnitsChange
	unitds          map[string]*unitData
	portsChange     chan *unitPortsChange
	serviceds       map[string]*serviceData
}

// NewFirewaller returns a new Firewaller.
func NewFirewaller(st *state.State) (*Firewaller, error) {
	fw := &Firewaller{
		st:              st,
		machinesWatcher: st.WatchMachines(),
		machineds:       make(map[int]*machineData),
		unitsChange:     make(chan *machineUnitsChange),
		unitds:          make(map[string]*unitData),
		portsChange:     make(chan *unitPortsChange),
		serviceds:       make(map[string]*serviceData),
	}
	go fw.loop()
	return fw, nil
}

func (fw *Firewaller) loop() {
	defer fw.finish()
	for {
		select {
		case <-fw.tomb.Dying():
			return
		case change, ok := <-fw.machinesWatcher.Changes():
			if !ok {
				return
			}
			for _, machine := range change.Removed {
				machined, ok := fw.machineds[machine.Id()]
				if !ok {
					panic("trying to remove machine that wasn't added")
				}
				delete(fw.machineds, machine.Id())
				if err := machined.stopWatch(); err != nil {
					log.Printf("machine data %d returned error when stopping: %v", machine.Id(), err)
				}
				log.Debugf("firewaller: stopped watching machine %d", machine.Id())
			}
			for _, machine := range change.Added {
				machined := newMachineData(machine, fw)
				fw.machineds[machine.Id()] = machined
				log.Debugf("firewaller: started watching machine %d", machine.Id())
			}
		case change := <-fw.unitsChange:
			for _, unit := range change.Removed {
				unitd, ok := fw.unitds[unit.Name()]
				if !ok {
					panic("trying to remove unit that wasn't added")
				}
				delete(fw.unitds, unit.Name())
				// TODO(mue) Close ports.
				if err := unitd.stopWatch(); err != nil {
					log.Printf("unit watcher %q returned error when stopping: %v", unit.Name(), err)
				}
				log.Debugf("firewaller: stopped watching unit %s", unit.Name())
			}
			for _, unit := range change.Added {
				unitd := newUnitData(unit, fw)
				fw.unitds[unit.Name()] = unitd
				if fw.serviceds[unit.ServiceName()] == nil {
					// TODO(mue) Add service watcher.
				}
				log.Debugf("firewaller: started watching unit %s", unit.Name())
			}
		case <-fw.portsChange:
			// TODO(mue) Handle changes of ports.
		}
	}
}

// finishes cleans up when the firewaller is stopping.
func (fw *Firewaller) finish() {
	watcher.Stop(fw.machinesWatcher, &fw.tomb)
	for _, unitd := range fw.unitds {
		fw.tomb.Kill(unitd.stopWatch())
	}
	for _, machined := range fw.machineds {
		fw.tomb.Kill(machined.stopWatch())
	}
	fw.tomb.Done()
}

// Wait waits for the Firewaller to exit.
func (fw *Firewaller) Wait() error {
	return fw.tomb.Wait()
}

// Stop stops the Firewaller and returns any error encountered while stopping.
func (fw *Firewaller) Stop() error {
	fw.tomb.Kill(nil)
	return fw.tomb.Wait()
}

// machineUnitsChange contains the changed units for one specific machine.
type machineUnitsChange struct {
	machined *machineData
	*state.MachineUnitsChange
}

// machineData watches the unit changes of a machine and passes them
// to the firewaller for handling.
type machineData struct {
	tomb       tomb.Tomb
	firewaller *Firewaller
	machine    *state.Machine
	watcher    *state.MachineUnitsWatcher
}

// newMachineData starts the watching of the passed machine.
func newMachineData(machine *state.Machine, fw *Firewaller) *machineData {
	md := &machineData{
		firewaller: fw,
		machine:    machine,
		watcher:    machine.WatchUnits(),
	}
	go md.watchLoop()
	return md
}

// watchLoop is the backend watching for machine units changes.
func (md *machineData) watchLoop() {
	defer md.tomb.Done()
	defer md.watcher.Stop()
	for {
		select {
		case <-md.tomb.Dying():
			return
		case change, ok := <-md.watcher.Changes():
			if !ok {
				md.firewaller.tomb.Kill(watcher.MustErr(md.watcher))
				return
			}
			select {
			case md.firewaller.unitsChange <- &machineUnitsChange{md, change}:
			case <-md.tomb.Dying():
				return
			}
		}
	}
}

// stopWatch stops the machine watching.
func (md *machineData) stopWatch() error {
	md.tomb.Kill(nil)
	return md.tomb.Wait()
}

// unitPortsChange contains the changed ports for one specific unit.
type unitPortsChange struct {
	unitd *unitData
	ports []state.Port
}

// unitData watches the port changes of a unit and passes them
// to the firewaller for handling.
type unitData struct {
	tomb       tomb.Tomb
	firewaller *Firewaller
	unit       *state.Unit
	watcher    *state.PortsWatcher
	service    *serviceData
	ports      []state.Port
}

// newMachineData starts the watching of the passed unit.
func newUnitData(unit *state.Unit, fw *Firewaller) *unitData {
	ud := &unitData{
		firewaller: fw,
		unit:       unit,
		watcher:    unit.WatchPorts(),
		ports:      make([]state.Port, 0),
	}
	go ud.watchLoop()
	return ud
}

func (ud *unitData) watchLoop() {
	defer ud.tomb.Done()
	defer ud.watcher.Stop()
	for {
		select {
		case <-ud.tomb.Dying():
			return
		case change, ok := <-ud.watcher.Changes():
			if !ok {
				ud.firewaller.tomb.Kill(watcher.MustErr(ud.watcher))
				return
			}
			select {
			case ud.firewaller.portsChange <- &unitPortsChange{ud, change}:
			case <-ud.tomb.Dying():
				return
			}
		}
	}
}

// stopWatch stops the unit watching.
func (ud *unitData) stopWatch() error {
	ud.tomb.Kill(nil)
	return ud.tomb.Wait()
}

// serviceData watches the exposed flag changes of a service and passes them
// to the firewaller for handling.
type serviceData struct {
	// TODO(mue) Fill with life.
	service *state.Service
	exposed bool
}
<|endoftext|>


<commit_before>package workertest

import (
	"encoding/json"
	"fmt"
	"net/http/httptest"
	"os"
	"path"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/taskcluster/slugid-go/slugid"
	tcclient "github.com/taskcluster/taskcluster-client-go"
	"github.com/taskcluster/taskcluster-client-go/queue"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/atomics"
	"github.com/taskcluster/taskcluster-worker/runtime/util"
	"github.com/taskcluster/taskcluster-worker/worker"
	"github.com/taskcluster/taskcluster-worker/worker/workertest/fakequeue"
)

const defaultTestCaseTimeout = 10 * time.Minute

// Case is a worker test case
type Case struct {
	Engine            string        // Engine to be used
	EngineConfig      string        // Engine configuration as JSON
	PluginConfig      string        // Configuration of plugins, see plugins.PluginManagerConfigSchema()
	Tasks             []Task        // Tasks to create and associated assertions
	Concurrency       int           // Worker concurrency, if zero defaulted to 1 and tasks will sequantially dependent
	StoppedGracefully bool          // True, if worker is expected to stop gracefully
	StoppedNow        bool          // True, if worker is expected to stop now
	Timeout           time.Duration // Test timeout, defaults to 10 min
}

// A Task to be included in a worker test case
type Task struct {
	Title           string                  // Optional title (for debugging)
	Payload         string                  // Task payload as JSON
	Success         bool                    // True, if task should be successfully
	Exception       runtime.ExceptionReason // Reason, if exception is expected
	Artifacts       ArtifactAssertions      // Mapping from artifact name to assertion
	AllowAdditional bool                    // True, if additional artifacts is allowed
	Status          StatusAssertion         // Optional, custom assertion on status and queue
}

// A StatusAssertion is a function that can make an assertion on a task status
type StatusAssertion func(t *testing.T, q *queue.Queue, status queue.TaskStatusStructure)

// An ArtifactAssertions is a mapping from artifact name to assertion for the
// artifact. If mapping to nil value, any artifact will be permitted.
type ArtifactAssertions map[string]func(t *testing.T, a Artifact)

// Artifact contains artifact meta-data.
type Artifact struct {
	ContentType     string
	Expires         time.Time
	Name            string
	StorageType     string
	Data            []byte
	ContentEncoding string
}

// provisionerId/workerType for test cases, access granted by role:
//   assume:project:taskcluster:worker-test-scopes
var dummyProvisionerID = "test-dummy-provisioner"

func dummyWorkerType() string {
	return "dummy-worker-" + slugid.V4()[:9]
}

// TestWithFakeQueue runs integration tests against FakeQueue
func (c Case) TestWithFakeQueue(t *testing.T) {
	// Create FakeQueue
	fq := fakequeue.New()
	s := httptest.NewServer(fq)
	defer s.Close()

	// Create listener
	l := fakequeue.NewFakeQueueListener(fq)

	// Create queue client
	q := queue.New(&tcclient.Credentials{
		// Long enough to pass schema validation
		ClientID:    "dummy-test-client-id",
		AccessToken: "non-secret-dummy-test-access-token",
	})
	q.BaseURL = s.URL

	c.testWithQueue(t, q, l)
}

// TestWithRealQueue runs integration tests against production queue
func (c Case) TestWithRealQueue(t *testing.T) {
	u := os.Getenv("PULSE_USERNAME")
	p := os.Getenv("PULSE_PASSWORD")
	if u == "" || p == "" {
		t.Skip("Skipping integration tests, because PULSE_USERNAME and PULSE_PASSWORD are not specified")
	}

	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	// Create listener
	l, err := fakequeue.NewPulseListener(u, p)
	require.NoError(t, err, "Failed to create PulseListener")

	// Create queue client
	q := queue.New(&tcclient.Credentials{
		ClientID:    os.Getenv("TASKCLUSTER_CLIENT_ID"),
		AccessToken: os.Getenv("TASKCLUSTER_ACCESS_TOKEN"),
		Certificate: os.Getenv("TASKCLUSTER_CERTIFICATE"),
	})
	if os.Getenv("QUEUE_BASE_URL") != "" {
		q.BaseURL = os.Getenv("QUEUE_BASE_URL")
	}

	c.testWithQueue(t, q, l)
}

// Test runs the test case
func (c Case) Test(t *testing.T) {
	passedFake := t.Run("FakeQueue", c.TestWithFakeQueue)
	// We don't run real integration tests if the FakeQueue tests fails
	// This is aimed to avoid polluting the queue and reducing feedback cycle.
	// You can manually call TestWithRealQueue(t), if you want to debug it.
	if passedFake {
		t.Run("RealQueue", c.TestWithRealQueue)
	} else {
		t.Run("RealQueue", func(t *testing.T) {
			t.Skip("Skipping integration tests, because FakeQueue tests failed")
		})
	}
}

func mustUnmarshalJSON(data string) interface{} {
	var v interface{}
	err := json.Unmarshal([]byte(data), &v)
	if err != nil {
		panic(fmt.Sprintf("Failed to parse JSON, error: %s, json: '%s'", err, data))
	}
	return v
}

func (c Case) testWithQueue(t *testing.T, q *queue.Queue, l fakequeue.Listener) {
	// Create config
	tempFolder := path.Join(os.TempDir(), slugid.Nice())
	defer os.RemoveAll(tempFolder)
	concurrency := c.Concurrency
	if concurrency == 0 {
		concurrency = 1
	}
	creds := map[string]interface{}{
		"clientId":    q.Credentials.ClientID,
		"accessToken": q.Credentials.AccessToken,
	}
	if q.Credentials.Certificate != "" {
		creds["certificate"] = q.Credentials.Certificate
	}
	workerID := "dummy-worker-" + slugid.V4()[:9]
	workerType := dummyWorkerType()
	config := map[string]interface{}{
		"credentials": creds,
		"engine":      c.Engine,
		"engines": map[string]interface{}{
			c.Engine: mustUnmarshalJSON(c.EngineConfig),
		},
		"minimumDiskSpace": 0,
		"minimumMemory":    0,
		"monitor":          mustUnmarshalJSON(`{"panicOnError": false, "type": "mock"}`),
		"plugins":          mustUnmarshalJSON(c.PluginConfig),
		"queueBaseUrl":     q.BaseURL,
		"temporaryFolder":  tempFolder,
		"webHookServer":    mustUnmarshalJSON(`{"provider": "localhost"}`),
		"worker": map[string]interface{}{
			"concurrency":         concurrency,
			"minimumReclaimDelay": 30,
			"pollingInterval":     1,
			"reclaimOffset":       30,
			"workerGroup":         "test-dummy-workers",
			"workerId":            workerID,
			"provisionerId":       dummyProvisionerID,
			"workerType":          workerType,
		},
	}
	err := worker.ConfigSchema().Validate(config)
	require.NoError(t, err, "Failed to validate worker config against schema")

	// Create worker
	w, err := worker.New(config)
	require.NoError(t, err, "Failed to create worker")

	// Create taskIDs
	taskIDs := make([]string, len(c.Tasks))
	for i := range taskIDs {
		taskIDs[i] = slugid.Nice()
	}

	// Setup event listeners
	events := make([]<-chan error, len(c.Tasks))
	util.Spawn(len(c.Tasks), func(i int) {
		events[i] = l.WaitForTask(taskIDs[i])
	})

	// Wait for tasks to be resolved
	var tasksResolved atomics.Once
	go tasksResolved.Do(func() {
		// Wait for events
		debug("Waiting for tasks to be resolved")
		util.Spawn(len(c.Tasks), func(i int) {
			err := <-events[i]
			assert.NoError(t, err, "Failed to listen for task %d", i)
			debug("Finished waiting for %s", taskIDs[i])
			if err != nil {
				debug("Error '%s' waiting for %s", taskIDs[i], err)
			}
		})
	})

	// Create tasks
	for i, task := range c.Tasks {
		tdef := queue.TaskDefinitionRequest{
			ProvisionerID: dummyProvisionerID,
			WorkerType:    workerType,
			Created:       tcclient.Time(time.Now()),
			Deadline:      tcclient.Time(time.Now().Add(60 * time.Minute)),
			Expires:       tcclient.Time(time.Now().Add(31 * 24 * 60 * time.Minute)),
			Payload:       json.RawMessage(task.Payload),
		}
		// If tasks are to run sequantially, we'll make them dependent
		if c.Concurrency == 0 && i > 0 {
			tdef.Dependencies = []string{taskIDs[i-1]}
			tdef.Requires = "all-resolved"
		}
		title := task.Title
		if title == "" {
			title = fmt.Sprintf("Task %d", i)
		}
		tdef.Metadata.Name = title
		tdef.Metadata.Description = "Task from taskcluster-worker integration tests"
		tdef.Metadata.Source = "https://github.com/taskcluster/taskcluster-worker/tree/master/worker/workertest/workertest.go"
		tdef.Metadata.Owner = "jonasfj@mozilla.com"
		debug("creating task '%s' as taskId: %s", title, taskIDs[i])
		_, err := q.CreateTask(taskIDs[i], &tdef)
		require.NoError(t, err, "Failed to create task: %s", title)
	}

	// Start worker
	var serr error
	var stopped atomics.Once
	go stopped.Do(func() {
		debug("starting worker with workerType: %s and workerID: %s", workerType, workerID)
		serr = w.Start()
		debug("worker stopped")
	})

	// Wait for events to have been handled
	timeout := c.Timeout
	if timeout == 0 {
		timeout = defaultTestCaseTimeout
	}
	select {
	case <-tasksResolved.Done():
	case <-time.After(timeout):
		assert.Fail(t, "Test case timed out, see workertest.Case.Timeout Property!")
		debug("worker.StopNow() because of test case timeout")
		w.StopNow()
		// We give it 30s to stop now, otherwise we end the test-case
		select {
		case <-stopped.Done():
		case <-time.After(30 * time.Second):
			debug("worker.StopNow() didn't stop after 30s")
		}
		return
	}

	// if we expect the worker to stop then we don't want to stop it here
	if !c.StoppedGracefully && !c.StoppedNow {
		// Stop worker
		debug("gracefully stopping worker (since test-case isn't stopping the worker)")
		w.StopGracefully()
	}

	// Wait for the worker to stop
	select {
	case <-stopped.Done():
	case <-time.After(30 * time.Second):
		assert.Fail(t, "Expected worker to stop")
	}

	// Verify assertions
	// We must do this after the worker has stopped, since tasks resolved
	// with exception can have artifacts added after resolution.
	debug("Verifying task assertions")
	// We could run these in parallel, but debugging is easier if we don't...
	for i, task := range c.Tasks {
		title := task.Title
		if title == "" {
			title = fmt.Sprintf("Task %d", i)
		}
		t.Run(title, func(t *testing.T) {
			verifyAssertions(t, title, taskIDs[i], task, q)
		})
	}

	// Check the stopping condition
	if c.StoppedNow {
		assert.Exactly(t, worker.ErrWorkerStoppedNow, serr, "Expected StoppedNow!")
	} else {
		assert.NoError(t, serr, "Expected worker to stop gracefully")
	}
}
<commit_msg>Add support for declaring task.scopes in workertest integration tests<commit_after>package workertest

import (
	"encoding/json"
	"fmt"
	"net/http/httptest"
	"os"
	"path"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/taskcluster/slugid-go/slugid"
	tcclient "github.com/taskcluster/taskcluster-client-go"
	"github.com/taskcluster/taskcluster-client-go/queue"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/atomics"
	"github.com/taskcluster/taskcluster-worker/runtime/util"
	"github.com/taskcluster/taskcluster-worker/worker"
	"github.com/taskcluster/taskcluster-worker/worker/workertest/fakequeue"
)

const defaultTestCaseTimeout = 10 * time.Minute

// Case is a worker test case
type Case struct {
	Engine            string        // Engine to be used
	EngineConfig      string        // Engine configuration as JSON
	PluginConfig      string        // Configuration of plugins, see plugins.PluginManagerConfigSchema()
	Tasks             []Task        // Tasks to create and associated assertions
	Concurrency       int           // Worker concurrency, if zero defaulted to 1 and tasks will sequantially dependent
	StoppedGracefully bool          // True, if worker is expected to stop gracefully
	StoppedNow        bool          // True, if worker is expected to stop now
	Timeout           time.Duration // Test timeout, defaults to 10 min
}

// A Task to be included in a worker test case
type Task struct {
	Title           string                  // Optional title (for debugging)
	Scopes          []string                // Task scopes
	Payload         string                  // Task payload as JSON
	Success         bool                    // True, if task should be successfully
	Exception       runtime.ExceptionReason // Reason, if exception is expected
	Artifacts       ArtifactAssertions      // Mapping from artifact name to assertion
	AllowAdditional bool                    // True, if additional artifacts is allowed
	Status          StatusAssertion         // Optional, custom assertion on status and queue
}

// A StatusAssertion is a function that can make an assertion on a task status
type StatusAssertion func(t *testing.T, q *queue.Queue, status queue.TaskStatusStructure)

// An ArtifactAssertions is a mapping from artifact name to assertion for the
// artifact. If mapping to nil value, any artifact will be permitted.
type ArtifactAssertions map[string]func(t *testing.T, a Artifact)

// Artifact contains artifact meta-data.
type Artifact struct {
	ContentType     string
	Expires         time.Time
	Name            string
	StorageType     string
	Data            []byte
	ContentEncoding string
}

// provisionerId/workerType for test cases, access granted by role:
//   assume:project:taskcluster:worker-test-scopes
var dummyProvisionerID = "test-dummy-provisioner"

func dummyWorkerType() string {
	return "dummy-worker-" + slugid.V4()[:9]
}

// TestWithFakeQueue runs integration tests against FakeQueue
func (c Case) TestWithFakeQueue(t *testing.T) {
	// Create FakeQueue
	fq := fakequeue.New()
	s := httptest.NewServer(fq)
	defer s.Close()

	// Create listener
	l := fakequeue.NewFakeQueueListener(fq)

	// Create queue client
	q := queue.New(&tcclient.Credentials{
		// Long enough to pass schema validation
		ClientID:    "dummy-test-client-id",
		AccessToken: "non-secret-dummy-test-access-token",
	})
	q.BaseURL = s.URL

	c.testWithQueue(t, q, l)
}

// TestWithRealQueue runs integration tests against production queue
func (c Case) TestWithRealQueue(t *testing.T) {
	u := os.Getenv("PULSE_USERNAME")
	p := os.Getenv("PULSE_PASSWORD")
	if u == "" || p == "" {
		t.Skip("Skipping integration tests, because PULSE_USERNAME and PULSE_PASSWORD are not specified")
	}

	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	// Create listener
	l, err := fakequeue.NewPulseListener(u, p)
	require.NoError(t, err, "Failed to create PulseListener")

	// Create queue client
	q := queue.New(&tcclient.Credentials{
		ClientID:    os.Getenv("TASKCLUSTER_CLIENT_ID"),
		AccessToken: os.Getenv("TASKCLUSTER_ACCESS_TOKEN"),
		Certificate: os.Getenv("TASKCLUSTER_CERTIFICATE"),
	})
	if os.Getenv("QUEUE_BASE_URL") != "" {
		q.BaseURL = os.Getenv("QUEUE_BASE_URL")
	}

	c.testWithQueue(t, q, l)
}

// Test runs the test case
func (c Case) Test(t *testing.T) {
	passedFake := t.Run("FakeQueue", c.TestWithFakeQueue)
	// We don't run real integration tests if the FakeQueue tests fails
	// This is aimed to avoid polluting the queue and reducing feedback cycle.
	// You can manually call TestWithRealQueue(t), if you want to debug it.
	if passedFake {
		t.Run("RealQueue", c.TestWithRealQueue)
	} else {
		t.Run("RealQueue", func(t *testing.T) {
			t.Skip("Skipping integration tests, because FakeQueue tests failed")
		})
	}
}

func mustUnmarshalJSON(data string) interface{} {
	var v interface{}
	err := json.Unmarshal([]byte(data), &v)
	if err != nil {
		panic(fmt.Sprintf("Failed to parse JSON, error: %s, json: '%s'", err, data))
	}
	return v
}

func (c Case) testWithQueue(t *testing.T, q *queue.Queue, l fakequeue.Listener) {
	// Create config
	tempFolder := path.Join(os.TempDir(), slugid.Nice())
	defer os.RemoveAll(tempFolder)
	concurrency := c.Concurrency
	if concurrency == 0 {
		concurrency = 1
	}
	creds := map[string]interface{}{
		"clientId":    q.Credentials.ClientID,
		"accessToken": q.Credentials.AccessToken,
	}
	if q.Credentials.Certificate != "" {
		creds["certificate"] = q.Credentials.Certificate
	}
	workerID := "dummy-worker-" + slugid.V4()[:9]
	workerType := dummyWorkerType()
	config := map[string]interface{}{
		"credentials": creds,
		"engine":      c.Engine,
		"engines": map[string]interface{}{
			c.Engine: mustUnmarshalJSON(c.EngineConfig),
		},
		"minimumDiskSpace": 0,
		"minimumMemory":    0,
		"monitor":          mustUnmarshalJSON(`{"panicOnError": false, "type": "mock"}`),
		"plugins":          mustUnmarshalJSON(c.PluginConfig),
		"queueBaseUrl":     q.BaseURL,
		"temporaryFolder":  tempFolder,
		"webHookServer":    mustUnmarshalJSON(`{"provider": "localhost"}`),
		"worker": map[string]interface{}{
			"concurrency":         concurrency,
			"minimumReclaimDelay": 30,
			"pollingInterval":     1,
			"reclaimOffset":       30,
			"workerGroup":         "test-dummy-workers",
			"workerId":            workerID,
			"provisionerId":       dummyProvisionerID,
			"workerType":          workerType,
		},
	}
	err := worker.ConfigSchema().Validate(config)
	require.NoError(t, err, "Failed to validate worker config against schema")

	// Create worker
	w, err := worker.New(config)
	require.NoError(t, err, "Failed to create worker")

	// Create taskIDs
	taskIDs := make([]string, len(c.Tasks))
	for i := range taskIDs {
		taskIDs[i] = slugid.Nice()
	}

	// Setup event listeners
	events := make([]<-chan error, len(c.Tasks))
	util.Spawn(len(c.Tasks), func(i int) {
		events[i] = l.WaitForTask(taskIDs[i])
	})

	// Wait for tasks to be resolved
	var tasksResolved atomics.Once
	go tasksResolved.Do(func() {
		// Wait for events
		debug("Waiting for tasks to be resolved")
		util.Spawn(len(c.Tasks), func(i int) {
			err := <-events[i]
			assert.NoError(t, err, "Failed to listen for task %d", i)
			debug("Finished waiting for %s", taskIDs[i])
			if err != nil {
				debug("Error '%s' waiting for %s", taskIDs[i], err)
			}
		})
	})

	// Create tasks
	for i, task := range c.Tasks {
		tdef := queue.TaskDefinitionRequest{
			ProvisionerID: dummyProvisionerID,
			WorkerType:    workerType,
			Created:       tcclient.Time(time.Now()),
			Deadline:      tcclient.Time(time.Now().Add(60 * time.Minute)),
			Expires:       tcclient.Time(time.Now().Add(31 * 24 * 60 * time.Minute)),
			Payload:       json.RawMessage(task.Payload),
		}
		// If tasks are to run sequantially, we'll make them dependent
		if c.Concurrency == 0 && i > 0 {
			tdef.Dependencies = []string{taskIDs[i-1]}
			tdef.Requires = "all-resolved"
		}
		title := task.Title
		if title == "" {
			title = fmt.Sprintf("Task %d", i)
		}
		tdef.Scopes = task.Scopes
		tdef.Metadata.Name = title
		tdef.Metadata.Description = "Task from taskcluster-worker integration tests"
		tdef.Metadata.Source = "https://github.com/taskcluster/taskcluster-worker/tree/master/worker/workertest/workertest.go"
		tdef.Metadata.Owner = "jonasfj@mozilla.com"
		debug("creating task '%s' as taskId: %s", title, taskIDs[i])
		_, err := q.CreateTask(taskIDs[i], &tdef)
		require.NoError(t, err, "Failed to create task: %s", title)
	}

	// Start worker
	var serr error
	var stopped atomics.Once
	go stopped.Do(func() {
		debug("starting worker with workerType: %s and workerID: %s", workerType, workerID)
		serr = w.Start()
		debug("worker stopped")
	})

	// Wait for events to have been handled
	timeout := c.Timeout
	if timeout == 0 {
		timeout = defaultTestCaseTimeout
	}
	select {
	case <-tasksResolved.Done():
	case <-time.After(timeout):
		assert.Fail(t, "Test case timed out, see workertest.Case.Timeout Property!")
		debug("worker.StopNow() because of test case timeout")
		w.StopNow()
		// We give it 30s to stop now, otherwise we end the test-case
		select {
		case <-stopped.Done():
		case <-time.After(30 * time.Second):
			debug("worker.StopNow() didn't stop after 30s")
		}
		return
	}

	// if we expect the worker to stop then we don't want to stop it here
	if !c.StoppedGracefully && !c.StoppedNow {
		// Stop worker
		debug("gracefully stopping worker (since test-case isn't stopping the worker)")
		w.StopGracefully()
	}

	// Wait for the worker to stop
	select {
	case <-stopped.Done():
	case <-time.After(30 * time.Second):
		assert.Fail(t, "Expected worker to stop")
	}

	// Verify assertions
	// We must do this after the worker has stopped, since tasks resolved
	// with exception can have artifacts added after resolution.
	debug("Verifying task assertions")
	// We could run these in parallel, but debugging is easier if we don't...
	for i, task := range c.Tasks {
		title := task.Title
		if title == "" {
			title = fmt.Sprintf("Task %d", i)
		}
		t.Run(title, func(t *testing.T) {
			verifyAssertions(t, title, taskIDs[i], task, q)
		})
	}

	// Check the stopping condition
	if c.StoppedNow {
		assert.Exactly(t, worker.ErrWorkerStoppedNow, serr, "Expected StoppedNow!")
	} else {
		assert.NoError(t, serr, "Expected worker to stop gracefully")
	}
}
<|endoftext|>


<commit_before>package workflowhelpers

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/commandstarter"
	"github.com/cloudfoundry-incubator/cf-test-helpers/internal"
	workflowhelpersinternal "github.com/cloudfoundry-incubator/cf-test-helpers/workflowhelpers/internal"
	"github.com/onsi/ginkgo"
	ginkgoconfig "github.com/onsi/ginkgo/config"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

type userValues interface {
	Username() string
	Password() string
}

type spaceValues interface {
	OrganizationName() string
	SpaceName() string
}

type UserContext struct {
	ApiUrl    string
	TestSpace spaceValues
	TestUser  userValues

	SkipSSLValidation bool
	CommandStarter    internal.Starter
	Timeout           time.Duration

	// the followings are left around for CATS to use
	Username string
	Password string
	Org      string
	Space    string
}

func cliErrorMessage(session *Session) string {
	var command string

	if strings.EqualFold(session.Command.Args[1], "auth") {
		command = strings.Join(session.Command.Args[:2], " ")
	} else {
		command = strings.Join(session.Command.Args, " ")
	}

	return fmt.Sprintf("\n>>> [ %s ] exited with an error \n", command)
}

func NewUserContext(apiUrl string, testUser userValues, testSpace spaceValues, skipSSLValidation bool, timeout time.Duration) UserContext {
	var org, space string
	if testSpace != nil {
		org = testSpace.OrganizationName()
		space = testSpace.SpaceName()
	}

	return UserContext{
		ApiUrl:            apiUrl,
		Username:          testUser.Username(),
		Password:          testUser.Password(),
		TestSpace:         testSpace,
		TestUser:          testUser,
		Org:               org,
		Space:             space,
		SkipSSLValidation: skipSSLValidation,
		CommandStarter:    commandstarter.NewCommandStarter(),
		Timeout:           timeout,
	}
}

func (uc UserContext) Login() {
	args := []string{"api", uc.ApiUrl}
	if uc.SkipSSLValidation {
		args = append(args, "--skip-ssl-validation")
	}

	session := internal.Cf(uc.CommandStarter, args...)
	EventuallyWithOffset(1, session, uc.Timeout).Should(Exit(0), cliErrorMessage(session))

	redactor := internal.NewRedactor(uc.TestUser.Password())
	redactingReporter := internal.NewRedactingReporter(ginkgo.GinkgoWriter, redactor)

	session = workflowhelpersinternal.CfAuth(uc.CommandStarter, redactingReporter, uc.TestUser.Username(), uc.TestUser.Password())
	EventuallyWithOffset(1, session, uc.Timeout).Should(Exit(0), cliErrorMessage(session))
}

func (uc UserContext) SetCfHomeDir() (string, string) {
	originalCfHomeDir := os.Getenv("CF_HOME")
	currentCfHomeDir, err := ioutil.TempDir("", fmt.Sprintf("cf_home_%d", ginkgoconfig.GinkgoConfig.ParallelNode))
	if err != nil {
		panic("Error: could not create temporary home directory: " + err.Error())
	}

	os.Setenv("CF_HOME", currentCfHomeDir)
	return originalCfHomeDir, currentCfHomeDir
}

func (uc UserContext) TargetSpace() {
	if uc.TestSpace != nil && uc.TestSpace.OrganizationName() != "" {
		var session *Session
		session = internal.Cf(uc.CommandStarter, "target", "-o", uc.TestSpace.OrganizationName(), "-s", uc.TestSpace.SpaceName())
		EventuallyWithOffset(1, session, uc.Timeout).Should(Exit(0), cliErrorMessage(session))
	}
}

func (uc UserContext) AddUserToSpace() {
	username := uc.TestUser.Username()
	orgName := uc.TestSpace.OrganizationName()
	spaceName := uc.TestSpace.SpaceName()

	spaceManager := internal.Cf(uc.CommandStarter, "set-space-role", username, orgName, spaceName, "SpaceManager")
	EventuallyWithOffset(1, spaceManager, uc.Timeout).Should(Exit())
	if spaceManager.ExitCode() != 0 {
		ExpectWithOffset(1, spaceManager.Out).Should(gbytes.Say("not authorized"))
	}

	spaceDeveloper := internal.Cf(uc.CommandStarter, "set-space-role", username, orgName, spaceName, "SpaceDeveloper")
	EventuallyWithOffset(1, spaceDeveloper, uc.Timeout).Should(Exit())
	if spaceDeveloper.ExitCode() != 0 {
		ExpectWithOffset(1, spaceDeveloper.Out).Should(gbytes.Say("not authorized"))
	}

	spaceAuditor := internal.Cf(uc.CommandStarter, "set-space-role", username, orgName, spaceName, "SpaceAuditor")
	EventuallyWithOffset(1, spaceAuditor, uc.Timeout).Should(Exit())
	if spaceAuditor.ExitCode() != 0 {
		ExpectWithOffset(1, spaceAuditor.Out).Should(gbytes.Say("not authorized"))
	}
}

func (uc UserContext) Logout() {
	session := internal.Cf(uc.CommandStarter, "logout")
	EventuallyWithOffset(1, session, uc.Timeout).Should(Exit(0), cliErrorMessage(session))
}

func (uc UserContext) UnsetCfHomeDir(originalCfHomeDir, currentCfHomeDir string) {
	os.Setenv("CF_HOME", originalCfHomeDir)
	os.RemoveAll(currentCfHomeDir)
}
<commit_msg>Add cf command error output to test error message [#158405539]<commit_after>package workflowhelpers

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/commandstarter"
	"github.com/cloudfoundry-incubator/cf-test-helpers/internal"
	workflowhelpersinternal "github.com/cloudfoundry-incubator/cf-test-helpers/workflowhelpers/internal"
	"github.com/onsi/ginkgo"
	ginkgoconfig "github.com/onsi/ginkgo/config"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

type userValues interface {
	Username() string
	Password() string
}

type spaceValues interface {
	OrganizationName() string
	SpaceName() string
}

type UserContext struct {
	ApiUrl    string
	TestSpace spaceValues
	TestUser  userValues

	SkipSSLValidation bool
	CommandStarter    internal.Starter
	Timeout           time.Duration

	// the followings are left around for CATS to use
	Username string
	Password string
	Org      string
	Space    string
}

func cliErrorMessage(session *Session) string {
	var command string

	if strings.EqualFold(session.Command.Args[1], "auth") {
		command = strings.Join(session.Command.Args[:2], " ")
	} else {
		command = strings.Join(session.Command.Args, " ")
	}

	return fmt.Sprintf("\n>>> [ %s ] exited with an error \n", command)
}

func apiErrorMessage(session *Session) string {
	apiEndpoint := strings.Join(session.Command.Args, " ")
	stdError := string(session.Err.Contents())

	return fmt.Sprintf("\n>>> [ %s ] exited with an error \n\n%s\n", apiEndpoint, stdError)
}

func NewUserContext(apiUrl string, testUser userValues, testSpace spaceValues, skipSSLValidation bool, timeout time.Duration) UserContext {
	var org, space string
	if testSpace != nil {
		org = testSpace.OrganizationName()
		space = testSpace.SpaceName()
	}

	return UserContext{
		ApiUrl:            apiUrl,
		Username:          testUser.Username(),
		Password:          testUser.Password(),
		TestSpace:         testSpace,
		TestUser:          testUser,
		Org:               org,
		Space:             space,
		SkipSSLValidation: skipSSLValidation,
		CommandStarter:    commandstarter.NewCommandStarter(),
		Timeout:           timeout,
	}
}

func (uc UserContext) Login() {
	args := []string{"api", uc.ApiUrl}
	if uc.SkipSSLValidation {
		args = append(args, "--skip-ssl-validation")
	}

	session := internal.Cf(uc.CommandStarter, args...).Wait(uc.Timeout)
	ExpectWithOffset(1, session).Should(Exit(0), apiErrorMessage(session))

	redactor := internal.NewRedactor(uc.TestUser.Password())
	redactingReporter := internal.NewRedactingReporter(ginkgo.GinkgoWriter, redactor)

	session = workflowhelpersinternal.CfAuth(uc.CommandStarter, redactingReporter, uc.TestUser.Username(), uc.TestUser.Password())
	EventuallyWithOffset(1, session, uc.Timeout).Should(Exit(0), cliErrorMessage(session))
}

func (uc UserContext) SetCfHomeDir() (string, string) {
	originalCfHomeDir := os.Getenv("CF_HOME")
	currentCfHomeDir, err := ioutil.TempDir("", fmt.Sprintf("cf_home_%d", ginkgoconfig.GinkgoConfig.ParallelNode))
	if err != nil {
		panic("Error: could not create temporary home directory: " + err.Error())
	}

	os.Setenv("CF_HOME", currentCfHomeDir)
	return originalCfHomeDir, currentCfHomeDir
}

func (uc UserContext) TargetSpace() {
	if uc.TestSpace != nil && uc.TestSpace.OrganizationName() != "" {
		var session *Session
		session = internal.Cf(uc.CommandStarter, "target", "-o", uc.TestSpace.OrganizationName(), "-s", uc.TestSpace.SpaceName())
		EventuallyWithOffset(1, session, uc.Timeout).Should(Exit(0), cliErrorMessage(session))
	}
}

func (uc UserContext) AddUserToSpace() {
	username := uc.TestUser.Username()
	orgName := uc.TestSpace.OrganizationName()
	spaceName := uc.TestSpace.SpaceName()

	spaceManager := internal.Cf(uc.CommandStarter, "set-space-role", username, orgName, spaceName, "SpaceManager")
	EventuallyWithOffset(1, spaceManager, uc.Timeout).Should(Exit())
	if spaceManager.ExitCode() != 0 {
		ExpectWithOffset(1, spaceManager.Out).Should(gbytes.Say("not authorized"))
	}

	spaceDeveloper := internal.Cf(uc.CommandStarter, "set-space-role", username, orgName, spaceName, "SpaceDeveloper")
	EventuallyWithOffset(1, spaceDeveloper, uc.Timeout).Should(Exit())
	if spaceDeveloper.ExitCode() != 0 {
		ExpectWithOffset(1, spaceDeveloper.Out).Should(gbytes.Say("not authorized"))
	}

	spaceAuditor := internal.Cf(uc.CommandStarter, "set-space-role", username, orgName, spaceName, "SpaceAuditor")
	EventuallyWithOffset(1, spaceAuditor, uc.Timeout).Should(Exit())
	if spaceAuditor.ExitCode() != 0 {
		ExpectWithOffset(1, spaceAuditor.Out).Should(gbytes.Say("not authorized"))
	}
}

func (uc UserContext) Logout() {
	session := internal.Cf(uc.CommandStarter, "logout")
	EventuallyWithOffset(1, session, uc.Timeout).Should(Exit(0), cliErrorMessage(session))
}

func (uc UserContext) UnsetCfHomeDir(originalCfHomeDir, currentCfHomeDir string) {
	os.Setenv("CF_HOME", originalCfHomeDir)
	os.RemoveAll(currentCfHomeDir)
}
<|endoftext|>


<commit_before>package container

import (
	"time"

	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/go-connections/nat"
)

// HealthConfig holds configuration settings for the HEALTHCHECK feature.
type HealthConfig struct {
	// Test is the test to perform to check that the container is healthy.
	// An empty slice means to inherit the default.
	// The options are:
	// {} : inherit healthcheck
	// {"NONE"} : disable healthcheck
	// {"CMD", args...} : exec arguments directly
	// {"CMD-SHELL", command} : run command with system's default shell
	Test []string `json:",omitempty"`

	// Zero means to inherit. Durations are expressed as integer nanoseconds.
	Interval time.Duration `json:",omitempty"` // Interval is the time to wait between checks.
	Timeout  time.Duration `json:",omitempty"` // Timeout is the time to wait before considering the check to have hung.

	// Retries is the number of consecutive failures needed to consider a container as unhealthy.
	// Zero means inherit.
	Retries int `json:",omitempty"`
}

// Config contains the configuration data about a container.
// It should hold only portable information about the container.
// Here, "portable" means "independent from the host we are running on".
// Non-portable information *should* appear in HostConfig.
// All fields added to this struct must be marked `omitempty` to keep getting
// predictable hashes from the old `v1Compatibility` configuration.
type Config struct {
	Hostname        string                // Hostname
	Domainname      string                // Domainname
	User            string                // User that will run the command(s) inside the container, also support user:group
	AttachStdin     bool                  // Attach the standard input, makes possible user interaction
	AttachStdout    bool                  // Attach the standard output
	AttachStderr    bool                  // Attach the standard error
	ExposedPorts    map[nat.Port]struct{} `json:",omitempty"` // List of exposed ports
	Tty             bool                  // Attach standard streams to a tty, including stdin if it is not closed.
	OpenStdin       bool                  // Open stdin
	StdinOnce       bool                  // If true, close stdin after the 1 attached client disconnects.
	Env             []string              // List of environment variable to set in the container
	Cmd             strslice.StrSlice     // Command to run when starting the container
	Healthcheck     *HealthConfig         `json:",omitempty"` // Healthcheck describes how to check the container is healthy
	ArgsEscaped     bool                  `json:",omitempty"` // True if command is already escaped (Windows specific)
	Image           string                // Name of the image as it was passed by the operator (e.g. could be symbolic)
	Volumes         map[string]struct{}   // List of volumes (mounts) used for the container
	WorkingDir      string                // Current directory (PWD) in the command will be launched
	Entrypoint      strslice.StrSlice     // Entrypoint to run when starting the container
	NetworkDisabled bool                  `json:",omitempty"` // Is network disabled
	MacAddress      string                `json:",omitempty"` // Mac Address of the container
	OnBuild         []string              // ONBUILD metadata that were defined on the image Dockerfile
	Labels          map[string]string     // List of labels set to this container
	StopSignal      string                `json:",omitempty"` // Signal to stop a container
	StopTimeout     *int                  `json:",omitempty"` // Timeout (in seconds) to stop a container
	Shell           strslice.StrSlice     `json:",omitempty"` // Shell for shell-form of RUN, CMD, ENTRYPOINT
}
<commit_msg>The type of 'ExportdPorts' should be 'nat.PortSet'<commit_after>package container

import (
	"time"

	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/go-connections/nat"
)

// HealthConfig holds configuration settings for the HEALTHCHECK feature.
type HealthConfig struct {
	// Test is the test to perform to check that the container is healthy.
	// An empty slice means to inherit the default.
	// The options are:
	// {} : inherit healthcheck
	// {"NONE"} : disable healthcheck
	// {"CMD", args...} : exec arguments directly
	// {"CMD-SHELL", command} : run command with system's default shell
	Test []string `json:",omitempty"`

	// Zero means to inherit. Durations are expressed as integer nanoseconds.
	Interval time.Duration `json:",omitempty"` // Interval is the time to wait between checks.
	Timeout  time.Duration `json:",omitempty"` // Timeout is the time to wait before considering the check to have hung.

	// Retries is the number of consecutive failures needed to consider a container as unhealthy.
	// Zero means inherit.
	Retries int `json:",omitempty"`
}

// Config contains the configuration data about a container.
// It should hold only portable information about the container.
// Here, "portable" means "independent from the host we are running on".
// Non-portable information *should* appear in HostConfig.
// All fields added to this struct must be marked `omitempty` to keep getting
// predictable hashes from the old `v1Compatibility` configuration.
type Config struct {
	Hostname        string                // Hostname
	Domainname      string                // Domainname
	User            string                // User that will run the command(s) inside the container, also support user:group
	AttachStdin     bool                  // Attach the standard input, makes possible user interaction
	AttachStdout    bool                  // Attach the standard output
	AttachStderr    bool                  // Attach the standard error
	ExposedPorts    nat.PortSet           `json:",omitempty"` // List of exposed ports
	Tty             bool                  // Attach standard streams to a tty, including stdin if it is not closed.
	OpenStdin       bool                  // Open stdin
	StdinOnce       bool                  // If true, close stdin after the 1 attached client disconnects.
	Env             []string              // List of environment variable to set in the container
	Cmd             strslice.StrSlice     // Command to run when starting the container
	Healthcheck     *HealthConfig         `json:",omitempty"` // Healthcheck describes how to check the container is healthy
	ArgsEscaped     bool                  `json:",omitempty"` // True if command is already escaped (Windows specific)
	Image           string                // Name of the image as it was passed by the operator (e.g. could be symbolic)
	Volumes         map[string]struct{}   // List of volumes (mounts) used for the container
	WorkingDir      string                // Current directory (PWD) in the command will be launched
	Entrypoint      strslice.StrSlice     // Entrypoint to run when starting the container
	NetworkDisabled bool                  `json:",omitempty"` // Is network disabled
	MacAddress      string                `json:",omitempty"` // Mac Address of the container
	OnBuild         []string              // ONBUILD metadata that were defined on the image Dockerfile
	Labels          map[string]string     // List of labels set to this container
	StopSignal      string                `json:",omitempty"` // Signal to stop a container
	StopTimeout     *int                  `json:",omitempty"` // Timeout (in seconds) to stop a container
	Shell           strslice.StrSlice     `json:",omitempty"` // Shell for shell-form of RUN, CMD, ENTRYPOINT
}
<|endoftext|>


<commit_before>/*
SPDX-License-Identifier: MIT

MIT License

Copyright (c) 2017 Thanh Ha

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"unicode"
)

var VERSION = "0.1.0"

func check(e error) {
	if e != nil {
		fmt.Println(e)
		os.Exit(1)
	}
}

func stripSpaces(str string) string {
	return strings.Map(func(r rune) rune {
		if unicode.IsSpace(r) {
			return -1
		}
		return r
	}, str)
}

func usage() {
	fmt.Printf("Usage: %s [OPTIONS] [PATTERN]...\n", os.Args[0])
	fmt.Printf("Scans a directory for files matching PATTERN and compares them with an expected license header.\n")
	fmt.Printf("\nPATTERN is a space separated list of regex patterns to search for files.\n")
	fmt.Printf("\nOptions:\n")
	flag.PrintDefaults()
}

func main() {
	licensePtr := flag.String("license", "license.txt", "Comma-separated list of license files to compare against.")
	versionPtr := flag.Bool("version", false, "Print version")
	// directoryPtr := flag.String("directory", ".", "Directory to search for files.")

	flag.Usage = usage
	flag.Parse()

	if *versionPtr {
		fmt.Println("License Checker version", VERSION)
		os.Exit(0)
	}

	fmt.Println("Search Patterns:", flag.Args())

	buf, err := ioutil.ReadFile(*licensePtr)
	check(err)
	licenseText := stripSpaces(string(buf))
	fmt.Println("License Text")
	fmt.Println(licenseText)

	file, err := os.Open("lhc.go")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer file.Close()

	headerText := ""
	scanner := bufio.NewScanner(file)
	commentSection := false
	for scanner.Scan() {
		s := scanner.Text()

		if strings.HasPrefix(s, "/*") {
			commentSection = true
		} else if commentSection && strings.Contains(s, "*/") {
			commentSection = false
			// TODO: Ignore Copyright lines in license header text
			// } else if strings.Contains(s, "Copyright") {
			// 	continue
		} else if strings.Contains(s, "SPDX-License-Identifier") {
			continue
		}

		if !commentSection &&
			!strings.HasPrefix(s, "#") &&
			!strings.HasPrefix(s, "//") {
			break
		}

		s = strings.TrimPrefix(s, "#")
		s = strings.TrimPrefix(s, "//")
		s = strings.TrimPrefix(s, "/*")
		s = strings.Split(s, "*/")[0]
		headerText += stripSpaces(s)
	}

	fmt.Println("Header Text")
	fmt.Println(headerText)
	if licenseText != headerText {
		fmt.Println("WARNING: License header does not match.", "lhc.go")
	}
}
<commit_msg>Refactor code into fetchLicense()<commit_after>/*
SPDX-License-Identifier: MIT

MIT License

Copyright (c) 2017 Thanh Ha

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"unicode"
)

var VERSION = "0.1.0"

func check(e error) {
	if e != nil {
		fmt.Println(e)
		os.Exit(1)
	}
}

func fetchLicense(filename string) string {
    file, err := os.Open(filename)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer file.Close()

    commentSection := false
	licenseText := ""
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		s := scanner.Text()

		if strings.HasPrefix(s, "/*") {
			commentSection = true
		} else if commentSection && strings.Contains(s, "*/") {
			commentSection = false
			// TODO: Ignore Copyright lines in license header text
			// } else if strings.Contains(s, "Copyright") {
			// 	continue
		} else if strings.Contains(s, "SPDX-License-Identifier") {
			continue
		}

		if !commentSection &&
			!strings.HasPrefix(s, "#") &&
			!strings.HasPrefix(s, "//") {
			break
		}

		s = strings.TrimPrefix(s, "#")
		s = strings.TrimPrefix(s, "//")
		s = strings.TrimPrefix(s, "/*")
		s = strings.Split(s, "*/")[0]
		licenseText += stripSpaces(s)
	}

    return licenseText
}

func stripSpaces(str string) string {
	return strings.Map(func(r rune) rune {
		if unicode.IsSpace(r) {
			return -1
		}
		return r
	}, str)
}

func usage() {
	fmt.Printf("Usage: %s [OPTIONS] [PATTERN]...\n", os.Args[0])
	fmt.Printf("Scans a directory for files matching PATTERN and compares them with an expected license header.\n")
	fmt.Printf("\nPATTERN is a space separated list of regex patterns to search for files.\n")
	fmt.Printf("\nOptions:\n")
	flag.PrintDefaults()
}

func main() {
	licensePtr := flag.String("license", "license.txt", "Comma-separated list of license files to compare against.")
	versionPtr := flag.Bool("version", false, "Print version")
	// directoryPtr := flag.String("directory", ".", "Directory to search for files.")

	flag.Usage = usage
	flag.Parse()

	if *versionPtr {
		fmt.Println("License Checker version", VERSION)
		os.Exit(0)
	}

	fmt.Println("Search Patterns:", flag.Args())

	buf, err := ioutil.ReadFile(*licensePtr)
	check(err)
	licenseText := stripSpaces(string(buf))
	fmt.Println("License Text")
	fmt.Println(licenseText)

	headerText := fetchLicense("lhc.go")

	fmt.Println("Header Text")
	fmt.Println(headerText)
	if licenseText != headerText {
		fmt.Println("WARNING: License header does not match.", "lhc.go")
	}
}
<|endoftext|>