<commit_before>package debian

import (
	"bufio"
	"bytes"
	"code.google.com/p/go-uuid/uuid"
	"fmt"
	"github.com/smira/aptly/database"
	"github.com/smira/aptly/utils"
	"github.com/ugorji/go/codec"
	"log"
	"path/filepath"
	"strings"
	"time"
)

// PublishedRepo is a published for http/ftp representation of snapshot as Debian repository
type PublishedRepo struct {
	// Internal unique ID
	UUID string
	// Prefix & distribution should be unique across all published repositories
	Prefix       string
	Distribution string
	Component    string
	// Architectures is a list of all architectures published
	Architectures []string
	// Snapshot as a source of publishing
	SnapshotUUID string

	snapshot *Snapshot
}

// NewPublishedRepo creates new published repository
func NewPublishedRepo(prefix string, distribution string, component string, architectures []string, snapshot *Snapshot) (*PublishedRepo, error) {
	prefix = filepath.Clean(prefix)
	if strings.HasPrefix(prefix, "/") {
		prefix = prefix[1:]
	}
	if strings.HasSuffix(prefix, "/") {
		prefix = prefix[:len(prefix)-1]
	}
	prefix = filepath.Clean(prefix)

	for _, component := range strings.Split(prefix, "/") {
		if component == ".." || component == "dists" || component == "pool" {
			return nil, fmt.Errorf("invalid prefix %s", prefix)
		}
	}

	return &PublishedRepo{
		UUID:          uuid.New(),
		Prefix:        prefix,
		Distribution:  distribution,
		Component:     component,
		Architectures: architectures,
		SnapshotUUID:  snapshot.UUID,
		snapshot:      snapshot,
	}, nil
}

// String returns human-readable represenation of PublishedRepo
func (p *PublishedRepo) String() string {
	var archs string

	if len(p.Architectures) > 0 {
		archs = fmt.Sprintf(" [%s]", strings.Join(p.Architectures, ", "))
	}

	return fmt.Sprintf("%s/%s (%s)%s publishes %s", p.Prefix, p.Distribution, p.Component, archs, p.snapshot.String())
}

// Key returns unique key identifying PublishedRepo
func (p *PublishedRepo) Key() []byte {
	return []byte("U" + p.Prefix + ">>" + p.Distribution)
}

// Encode does msgpack encoding of PublishedRepo
func (p *PublishedRepo) Encode() []byte {
	var buf bytes.Buffer

	encoder := codec.NewEncoder(&buf, &codec.MsgpackHandle{})
	encoder.Encode(p)

	return buf.Bytes()
}

// Decode decodes msgpack representation into PublishedRepo
func (p *PublishedRepo) Decode(input []byte) error {
	decoder := codec.NewDecoderBytes(input, &codec.MsgpackHandle{})
	return decoder.Decode(p)
}

// Publish publishes snapshot (repository) contents, links package files, generates Packages & Release files, signs them
func (p *PublishedRepo) Publish(repo *Repository, packageCollection *PackageCollection, signer utils.Signer) error {
	err := repo.MkDir(filepath.Join(p.Prefix, "pool"))
	if err != nil {
		return err
	}
	basePath := filepath.Join(p.Prefix, "dists", p.Distribution)
	err = repo.MkDir(basePath)
	if err != nil {
		return err
	}

	// Load all packages
	list, err := NewPackageListFromRefList(p.snapshot.RefList(), packageCollection)
	if err != nil {
		return fmt.Errorf("unable to load packages: %s", err)
	}

	if list.Len() == 0 {
		return fmt.Errorf("snapshot is empty")
	}

	if len(p.Architectures) == 0 {
		p.Architectures = list.Architectures(true)
	}

	if len(p.Architectures) == 0 {
		return fmt.Errorf("unable to figure out list of architectures, please supply explicit list")
	}

	generatedFiles := map[string]utils.ChecksumInfo{}

	// For all architectures, generate release file
	for _, arch := range p.Architectures {
		relativePath := filepath.Join(p.Component, fmt.Sprintf("binary-%s", arch), "Packages")
		err = repo.MkDir(filepath.Dir(filepath.Join(basePath, relativePath)))
		if err != nil {
			return err
		}

		packagesFile, err := repo.CreateFile(filepath.Join(basePath, relativePath))
		if err != nil {
			return fmt.Errorf("unable to creates Packages file: %s", err)
		}

		bufWriter := bufio.NewWriter(packagesFile)

		err = list.ForEach(func(pkg *Package) error {
			if pkg.MatchesArchitecture(arch) {
				err = pkg.LinkFromPool(repo, p.Prefix, p.Component)
				if err != nil {
					return err
				}

				err = pkg.Stanza().WriteTo(bufWriter)
				if err != nil {
					return err
				}
				err = bufWriter.WriteByte('\n')
				if err != nil {
					return err
				}

			}

			return nil
		})

		if err != nil {
			return fmt.Errorf("unable to creates process packages: %s", err)
		}

		err = bufWriter.Flush()
		if err != nil {
			return fmt.Errorf("unable to write Packages file: %s", err)
		}

		err = utils.CompressFile(packagesFile)
		if err != nil {
			return fmt.Errorf("unable to compress Packages files: %s", err)
		}

		packagesFile.Close()

		checksumInfo, err := repo.ChecksumsForFile(filepath.Join(basePath, relativePath))
		if err != nil {
			return fmt.Errorf("unable to collect checksums: %s", err)
		}
		generatedFiles[relativePath] = checksumInfo

		checksumInfo, err = repo.ChecksumsForFile(filepath.Join(basePath, relativePath+".gz"))
		if err != nil {
			return fmt.Errorf("unable to collect checksums: %s", err)
		}
		generatedFiles[relativePath+".gz"] = checksumInfo

		checksumInfo, err = repo.ChecksumsForFile(filepath.Join(basePath, relativePath+".bz2"))
		if err != nil {
			return fmt.Errorf("unable to collect checksums: %s", err)
		}
		generatedFiles[relativePath+".bz2"] = checksumInfo

	}

	release := make(Stanza)
	release["Origin"] = p.Prefix + " " + p.Distribution
	release["Label"] = p.Prefix + " " + p.Distribution
	release["Codename"] = p.Distribution
	release["Date"] = time.Now().UTC().Format("Mon, 2 Jan 2006 15:04:05 MST")
	release["Components"] = p.Component
	release["Architectures"] = strings.Join(p.Architectures, " ")
	release["Description"] = " Generated by aptly\n"
	release["MD5Sum"] = "\n"
	release["SHA1"] = "\n"
	release["SHA256"] = "\n"

	for path, info := range generatedFiles {
		release["MD5Sum"] += fmt.Sprintf(" %s %8d %s\n", info.MD5, info.Size, path)
		release["SHA1"] += fmt.Sprintf(" %s %8d %s\n", info.SHA1, info.Size, path)
		release["SHA256"] += fmt.Sprintf(" %s %8d %s\n", info.SHA256, info.Size, path)
	}

	releaseFile, err := repo.CreateFile(filepath.Join(basePath, "Release"))
	if err != nil {
		return fmt.Errorf("unable to create Release file: %s", err)
	}

	bufWriter := bufio.NewWriter(releaseFile)

	err = release.WriteTo(bufWriter)
	if err != nil {
		return fmt.Errorf("unable to create Release file: %s", err)
	}

	err = bufWriter.Flush()
	if err != nil {
		return fmt.Errorf("unable to create Release file: %s", err)
	}

	releaseFilename := releaseFile.Name()
	releaseFile.Close()

	if signer != nil {
		err = signer.DetachedSign(releaseFilename, releaseFilename+".gpg")
		if err != nil {
			return fmt.Errorf("unable to sign Release file: %s", err)
		}

		err = signer.ClearSign(releaseFilename, filepath.Join(filepath.Dir(releaseFilename), "InRelease"))
		if err != nil {
			return fmt.Errorf("unable to sign Release file: %s", err)
		}
	}

	return nil
}

// RemoveFiles removes files that were created by Publish
//
// It can remove prefix fully, and part of pool (for specific component)
func (p *PublishedRepo) RemoveFiles(repo *Repository, removePrefix, removePoolComponent bool) error {
	if removePrefix {
		err := repo.RemoveDirs(filepath.Join(p.Prefix, "dists"))
		if err != nil {
			return err
		}

		return repo.RemoveDirs(filepath.Join(p.Prefix, "pool"))
	}

	err := repo.RemoveDirs(filepath.Join(p.Prefix, "dists", p.Distribution))
	if err != nil {
		return err
	}

	if removePoolComponent {
		err = repo.RemoveDirs(filepath.Join(p.Prefix, "pool", p.Component))
		if err != nil {
			return err
		}
	}
	return nil
}

// PublishedRepoCollection does listing, updating/adding/deleting of PublishedRepos
type PublishedRepoCollection struct {
	db   database.Storage
	list []*PublishedRepo
}

// NewPublishedRepoCollection loads PublishedRepos from DB and makes up collection
func NewPublishedRepoCollection(db database.Storage) *PublishedRepoCollection {
	result := &PublishedRepoCollection{
		db: db,
	}

	blobs := db.FetchByPrefix([]byte("U"))
	result.list = make([]*PublishedRepo, 0, len(blobs))

	for _, blob := range blobs {
		r := &PublishedRepo{}
		if err := r.Decode(blob); err != nil {
			log.Printf("Error decoding published repo: %s\n", err)
		} else {
			result.list = append(result.list, r)
		}
	}

	return result
}

// Add appends new repo to collection and saves it
func (collection *PublishedRepoCollection) Add(repo *PublishedRepo) error {
	if collection.CheckDuplicate(repo) != nil {
		return fmt.Errorf("published repo with prefix/distribution %s/%s already exists", repo.Prefix, repo.Distribution)
	}

	err := collection.Update(repo)
	if err != nil {
		return err
	}

	collection.list = append(collection.list, repo)
	return nil
}

// CheckDuplicate verifies that there's no published repo with the same name
func (collection *PublishedRepoCollection) CheckDuplicate(repo *PublishedRepo) *PublishedRepo {
	for _, r := range collection.list {
		if r.Prefix == repo.Prefix && r.Distribution == repo.Distribution {
			return r
		}
	}

	return nil
}

// Update stores updated information about repo in DB
func (collection *PublishedRepoCollection) Update(repo *PublishedRepo) error {
	err := collection.db.Put(repo.Key(), repo.Encode())
	if err != nil {
		return err
	}
	return nil
}

// LoadComplete loads additional information for remote repo
func (collection *PublishedRepoCollection) LoadComplete(repo *PublishedRepo, snapshotCollection *SnapshotCollection) error {
	snapshot, err := snapshotCollection.ByUUID(repo.SnapshotUUID)
	if err != nil {
		return err
	}

	repo.snapshot = snapshot
	return nil
}

// ByPrefixDistribution looks up repository by prefix & distribution
func (collection *PublishedRepoCollection) ByPrefixDistribution(prefix, distribution string) (*PublishedRepo, error) {
	for _, r := range collection.list {
		if r.Prefix == prefix && r.Distribution == distribution {
			return r, nil
		}
	}
	return nil, fmt.Errorf("published repo with prefix/distribution %s/%s not found", prefix, distribution)
}

// ByUUID looks up repository by uuid
func (collection *PublishedRepoCollection) ByUUID(uuid string) (*PublishedRepo, error) {
	for _, r := range collection.list {
		if r.UUID == uuid {
			return r, nil
		}
	}
	return nil, fmt.Errorf("published repo with uuid %s not found", uuid)
}

// BySnapshot looks up repository by snapshot source
func (collection *PublishedRepoCollection) BySnapshot(snapshot *Snapshot) []*PublishedRepo {
	result := make([]*PublishedRepo, 0)
	for _, r := range collection.list {
		if r.SnapshotUUID == snapshot.UUID {
			result = append(result, r)
		}
	}
	return result
}

// ForEach runs method for each repository
func (collection *PublishedRepoCollection) ForEach(handler func(*PublishedRepo) error) error {
	var err error
	for _, r := range collection.list {
		err = handler(r)
		if err != nil {
			return err
		}
	}
	return err
}

// Len returns number of remote repos
func (collection *PublishedRepoCollection) Len() int {
	return len(collection.list)
}

// Remove removes published repository, cleaning up directories, files
func (collection *PublishedRepoCollection) Remove(packageRepo *Repository, prefix, distribution string) error {
	repo, err := collection.ByPrefixDistribution(prefix, distribution)
	if err != nil {
		return err
	}

	removePrefix := true
	removePoolComponent := true
	repoPosition := -1

	for i, r := range collection.list {
		if r == repo {
			repoPosition = i
			continue
		}
		if r.Prefix == repo.Prefix {
			removePrefix = false
			if r.Component == repo.Component {
				removePoolComponent = false
			}
		}
	}

	err = repo.RemoveFiles(packageRepo, removePrefix, removePoolComponent)
	if err != nil {
		return err
	}

	collection.list[len(collection.list)-1], collection.list[repoPosition], collection.list =
		nil, collection.list[len(collection.list)-1], collection.list[:len(collection.list)-1]

	return collection.db.Delete(repo.Key())
}
<commit_msg>Publishing source packages index.<commit_after>package debian

import (
	"bufio"
	"bytes"
	"code.google.com/p/go-uuid/uuid"
	"fmt"
	"github.com/smira/aptly/database"
	"github.com/smira/aptly/utils"
	"github.com/ugorji/go/codec"
	"log"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

// PublishedRepo is a published for http/ftp representation of snapshot as Debian repository
type PublishedRepo struct {
	// Internal unique ID
	UUID string
	// Prefix & distribution should be unique across all published repositories
	Prefix       string
	Distribution string
	Component    string
	// Architectures is a list of all architectures published
	Architectures []string
	// Snapshot as a source of publishing
	SnapshotUUID string

	snapshot *Snapshot
}

// NewPublishedRepo creates new published repository
func NewPublishedRepo(prefix string, distribution string, component string, architectures []string, snapshot *Snapshot) (*PublishedRepo, error) {
	prefix = filepath.Clean(prefix)
	if strings.HasPrefix(prefix, "/") {
		prefix = prefix[1:]
	}
	if strings.HasSuffix(prefix, "/") {
		prefix = prefix[:len(prefix)-1]
	}
	prefix = filepath.Clean(prefix)

	for _, component := range strings.Split(prefix, "/") {
		if component == ".." || component == "dists" || component == "pool" {
			return nil, fmt.Errorf("invalid prefix %s", prefix)
		}
	}

	return &PublishedRepo{
		UUID:          uuid.New(),
		Prefix:        prefix,
		Distribution:  distribution,
		Component:     component,
		Architectures: architectures,
		SnapshotUUID:  snapshot.UUID,
		snapshot:      snapshot,
	}, nil
}

// String returns human-readable represenation of PublishedRepo
func (p *PublishedRepo) String() string {
	var archs string

	if len(p.Architectures) > 0 {
		archs = fmt.Sprintf(" [%s]", strings.Join(p.Architectures, ", "))
	}

	return fmt.Sprintf("%s/%s (%s)%s publishes %s", p.Prefix, p.Distribution, p.Component, archs, p.snapshot.String())
}

// Key returns unique key identifying PublishedRepo
func (p *PublishedRepo) Key() []byte {
	return []byte("U" + p.Prefix + ">>" + p.Distribution)
}

// Encode does msgpack encoding of PublishedRepo
func (p *PublishedRepo) Encode() []byte {
	var buf bytes.Buffer

	encoder := codec.NewEncoder(&buf, &codec.MsgpackHandle{})
	encoder.Encode(p)

	return buf.Bytes()
}

// Decode decodes msgpack representation into PublishedRepo
func (p *PublishedRepo) Decode(input []byte) error {
	decoder := codec.NewDecoderBytes(input, &codec.MsgpackHandle{})
	return decoder.Decode(p)
}

// Publish publishes snapshot (repository) contents, links package files, generates Packages & Release files, signs them
func (p *PublishedRepo) Publish(repo *Repository, packageCollection *PackageCollection, signer utils.Signer) error {
	err := repo.MkDir(filepath.Join(p.Prefix, "pool"))
	if err != nil {
		return err
	}
	basePath := filepath.Join(p.Prefix, "dists", p.Distribution)
	err = repo.MkDir(basePath)
	if err != nil {
		return err
	}

	// Load all packages
	list, err := NewPackageListFromRefList(p.snapshot.RefList(), packageCollection)
	if err != nil {
		return fmt.Errorf("unable to load packages: %s", err)
	}

	if list.Len() == 0 {
		return fmt.Errorf("snapshot is empty")
	}

	if len(p.Architectures) == 0 {
		p.Architectures = list.Architectures(true)
	}

	if len(p.Architectures) == 0 {
		return fmt.Errorf("unable to figure out list of architectures, please supply explicit list")
	}

	sort.Strings(p.Architectures)

	generatedFiles := map[string]utils.ChecksumInfo{}

	// For all architectures, generate release file
	for _, arch := range p.Architectures {
		var relativePath string
		if arch == "source" {
			relativePath = filepath.Join(p.Component, "source", "Sources")
		} else {
			relativePath = filepath.Join(p.Component, fmt.Sprintf("binary-%s", arch), "Packages")
		}
		err = repo.MkDir(filepath.Dir(filepath.Join(basePath, relativePath)))
		if err != nil {
			return err
		}

		packagesFile, err := repo.CreateFile(filepath.Join(basePath, relativePath))
		if err != nil {
			return fmt.Errorf("unable to creates Packages file: %s", err)
		}

		bufWriter := bufio.NewWriter(packagesFile)

		err = list.ForEach(func(pkg *Package) error {
			if pkg.MatchesArchitecture(arch) {
				err = pkg.LinkFromPool(repo, p.Prefix, p.Component)
				if err != nil {
					return err
				}

				err = pkg.Stanza().WriteTo(bufWriter)
				if err != nil {
					return err
				}
				err = bufWriter.WriteByte('\n')
				if err != nil {
					return err
				}

			}

			return nil
		})

		if err != nil {
			return fmt.Errorf("unable to process packages: %s", err)
		}

		err = bufWriter.Flush()
		if err != nil {
			return fmt.Errorf("unable to write Packages file: %s", err)
		}

		err = utils.CompressFile(packagesFile)
		if err != nil {
			return fmt.Errorf("unable to compress Packages files: %s", err)
		}

		packagesFile.Close()

		checksumInfo, err := repo.ChecksumsForFile(filepath.Join(basePath, relativePath))
		if err != nil {
			return fmt.Errorf("unable to collect checksums: %s", err)
		}
		generatedFiles[relativePath] = checksumInfo

		checksumInfo, err = repo.ChecksumsForFile(filepath.Join(basePath, relativePath+".gz"))
		if err != nil {
			return fmt.Errorf("unable to collect checksums: %s", err)
		}
		generatedFiles[relativePath+".gz"] = checksumInfo

		checksumInfo, err = repo.ChecksumsForFile(filepath.Join(basePath, relativePath+".bz2"))
		if err != nil {
			return fmt.Errorf("unable to collect checksums: %s", err)
		}
		generatedFiles[relativePath+".bz2"] = checksumInfo

	}

	release := make(Stanza)
	release["Origin"] = p.Prefix + " " + p.Distribution
	release["Label"] = p.Prefix + " " + p.Distribution
	release["Codename"] = p.Distribution
	release["Date"] = time.Now().UTC().Format("Mon, 2 Jan 2006 15:04:05 MST")
	release["Components"] = p.Component
	release["Architectures"] = strings.Join(utils.StrSlicesSubstract(p.Architectures, []string{"source"}), " ")
	release["Description"] = " Generated by aptly\n"
	release["MD5Sum"] = "\n"
	release["SHA1"] = "\n"
	release["SHA256"] = "\n"

	for path, info := range generatedFiles {
		release["MD5Sum"] += fmt.Sprintf(" %s %8d %s\n", info.MD5, info.Size, path)
		release["SHA1"] += fmt.Sprintf(" %s %8d %s\n", info.SHA1, info.Size, path)
		release["SHA256"] += fmt.Sprintf(" %s %8d %s\n", info.SHA256, info.Size, path)
	}

	releaseFile, err := repo.CreateFile(filepath.Join(basePath, "Release"))
	if err != nil {
		return fmt.Errorf("unable to create Release file: %s", err)
	}

	bufWriter := bufio.NewWriter(releaseFile)

	err = release.WriteTo(bufWriter)
	if err != nil {
		return fmt.Errorf("unable to create Release file: %s", err)
	}

	err = bufWriter.Flush()
	if err != nil {
		return fmt.Errorf("unable to create Release file: %s", err)
	}

	releaseFilename := releaseFile.Name()
	releaseFile.Close()

	if signer != nil {
		err = signer.DetachedSign(releaseFilename, releaseFilename+".gpg")
		if err != nil {
			return fmt.Errorf("unable to sign Release file: %s", err)
		}

		err = signer.ClearSign(releaseFilename, filepath.Join(filepath.Dir(releaseFilename), "InRelease"))
		if err != nil {
			return fmt.Errorf("unable to sign Release file: %s", err)
		}
	}

	return nil
}

// RemoveFiles removes files that were created by Publish
//
// It can remove prefix fully, and part of pool (for specific component)
func (p *PublishedRepo) RemoveFiles(repo *Repository, removePrefix, removePoolComponent bool) error {
	if removePrefix {
		err := repo.RemoveDirs(filepath.Join(p.Prefix, "dists"))
		if err != nil {
			return err
		}

		return repo.RemoveDirs(filepath.Join(p.Prefix, "pool"))
	}

	err := repo.RemoveDirs(filepath.Join(p.Prefix, "dists", p.Distribution))
	if err != nil {
		return err
	}

	if removePoolComponent {
		err = repo.RemoveDirs(filepath.Join(p.Prefix, "pool", p.Component))
		if err != nil {
			return err
		}
	}
	return nil
}

// PublishedRepoCollection does listing, updating/adding/deleting of PublishedRepos
type PublishedRepoCollection struct {
	db   database.Storage
	list []*PublishedRepo
}

// NewPublishedRepoCollection loads PublishedRepos from DB and makes up collection
func NewPublishedRepoCollection(db database.Storage) *PublishedRepoCollection {
	result := &PublishedRepoCollection{
		db: db,
	}

	blobs := db.FetchByPrefix([]byte("U"))
	result.list = make([]*PublishedRepo, 0, len(blobs))

	for _, blob := range blobs {
		r := &PublishedRepo{}
		if err := r.Decode(blob); err != nil {
			log.Printf("Error decoding published repo: %s\n", err)
		} else {
			result.list = append(result.list, r)
		}
	}

	return result
}

// Add appends new repo to collection and saves it
func (collection *PublishedRepoCollection) Add(repo *PublishedRepo) error {
	if collection.CheckDuplicate(repo) != nil {
		return fmt.Errorf("published repo with prefix/distribution %s/%s already exists", repo.Prefix, repo.Distribution)
	}

	err := collection.Update(repo)
	if err != nil {
		return err
	}

	collection.list = append(collection.list, repo)
	return nil
}

// CheckDuplicate verifies that there's no published repo with the same name
func (collection *PublishedRepoCollection) CheckDuplicate(repo *PublishedRepo) *PublishedRepo {
	for _, r := range collection.list {
		if r.Prefix == repo.Prefix && r.Distribution == repo.Distribution {
			return r
		}
	}

	return nil
}

// Update stores updated information about repo in DB
func (collection *PublishedRepoCollection) Update(repo *PublishedRepo) error {
	err := collection.db.Put(repo.Key(), repo.Encode())
	if err != nil {
		return err
	}
	return nil
}

// LoadComplete loads additional information for remote repo
func (collection *PublishedRepoCollection) LoadComplete(repo *PublishedRepo, snapshotCollection *SnapshotCollection) error {
	snapshot, err := snapshotCollection.ByUUID(repo.SnapshotUUID)
	if err != nil {
		return err
	}

	repo.snapshot = snapshot
	return nil
}

// ByPrefixDistribution looks up repository by prefix & distribution
func (collection *PublishedRepoCollection) ByPrefixDistribution(prefix, distribution string) (*PublishedRepo, error) {
	for _, r := range collection.list {
		if r.Prefix == prefix && r.Distribution == distribution {
			return r, nil
		}
	}
	return nil, fmt.Errorf("published repo with prefix/distribution %s/%s not found", prefix, distribution)
}

// ByUUID looks up repository by uuid
func (collection *PublishedRepoCollection) ByUUID(uuid string) (*PublishedRepo, error) {
	for _, r := range collection.list {
		if r.UUID == uuid {
			return r, nil
		}
	}
	return nil, fmt.Errorf("published repo with uuid %s not found", uuid)
}

// BySnapshot looks up repository by snapshot source
func (collection *PublishedRepoCollection) BySnapshot(snapshot *Snapshot) []*PublishedRepo {
	result := make([]*PublishedRepo, 0)
	for _, r := range collection.list {
		if r.SnapshotUUID == snapshot.UUID {
			result = append(result, r)
		}
	}
	return result
}

// ForEach runs method for each repository
func (collection *PublishedRepoCollection) ForEach(handler func(*PublishedRepo) error) error {
	var err error
	for _, r := range collection.list {
		err = handler(r)
		if err != nil {
			return err
		}
	}
	return err
}

// Len returns number of remote repos
func (collection *PublishedRepoCollection) Len() int {
	return len(collection.list)
}

// Remove removes published repository, cleaning up directories, files
func (collection *PublishedRepoCollection) Remove(packageRepo *Repository, prefix, distribution string) error {
	repo, err := collection.ByPrefixDistribution(prefix, distribution)
	if err != nil {
		return err
	}

	removePrefix := true
	removePoolComponent := true
	repoPosition := -1

	for i, r := range collection.list {
		if r == repo {
			repoPosition = i
			continue
		}
		if r.Prefix == repo.Prefix {
			removePrefix = false
			if r.Component == repo.Component {
				removePoolComponent = false
			}
		}
	}

	err = repo.RemoveFiles(packageRepo, removePrefix, removePoolComponent)
	if err != nil {
		return err
	}

	collection.list[len(collection.list)-1], collection.list[repoPosition], collection.list =
		nil, collection.list[len(collection.list)-1], collection.list[:len(collection.list)-1]

	return collection.db.Delete(repo.Key())
}
<|endoftext|>


<commit_before>package stripe

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"
)

// enable logging to print the request and reponses to stdout
var _log bool

// the API Key used to authenticate all Stripe API requests
var _key string

// the default URL for all Stripe API requests
var _url string = "https://api.stripe.com"

// SetUrl will override the default Stripe API URL. This is primarily used
// for unit testing.
func SetUrl(url string) {
	_url = url
}

// SetKey will set the default Stripe API key used to authenticate all Stripe
// API requests.
func SetKey(key string) {
	_key = key
}

// Available APIs
var (
	Charges       = new(ChargeClient)
	Coupons       = new(CouponClient)
	Customers     = new(CustomerClient)
	Invoices      = new(InvoiceClient)
	InvoiceItems  = new(InvoiceItemClient)
	Plans         = new(PlanClient)
	Subscriptions = new(SubscriptionClient)
	Tokens        = new(TokenClient)
)

// SetKeyEnv retrieves the Stripe API key using the STRIPE_API_KEY environment
// variable.
func SetKeyEnv() (err error) {
	_key = os.Getenv("STRIPE_API_KEY")
	if _key == "" {
		err = errors.New("STRIPE_API_KEY not found in environment")
	}
	return
}

// query submits an http.Request and parses the JSON-encoded http.Response,
// storing the result in the value pointed to by v.
func query(method, path string, values url.Values, v interface{}) error {
	// parse the stripe URL
	endpoint, err := url.Parse(_url)
	if err != nil {
		return err
	}

	// set the endpoint for the specific API
	endpoint.Path = path
	endpoint.User = url.User(_key)

	// if this is an http GET, add the url.Values to the endpoint
	if method == "GET" {
		endpoint.RawQuery = values.Encode()
	}

	// else if this is not a GET, encode the url.Values in the body.
	var reqBody io.Reader
	if method != "GET" && values != nil {
		reqBody = strings.NewReader(values.Encode())
	}

	// Log request if logging enabled
	if _log {
		fmt.Println("REQUEST: ", method, endpoint.String())
		fmt.Println(values.Encode())
	}

	// create the request
	req, err := http.NewRequest(method, endpoint.String(), reqBody)
	if err != nil {
		return err
	}

	// submit the http request
	r, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	// read the body of the http message into a byte array
	body, err := ioutil.ReadAll(r.Body)
	defer r.Body.Close()
	if err != nil {
		return err
	}

	// Log response if logging enabled
	if _log {
		fmt.Println("RESPONSE: ", r.StatusCode)
		fmt.Println(string(body))
	}

	// is this an error?
	if r.StatusCode != 200 {
		error := Error{}
		json.Unmarshal(body, &error)
		return &error
	}

	//parse the JSON response into the response object
	return json.Unmarshal(body, v)
}

// Error encapsulates an error returned by the Stripe REST API.
type Error struct {
	Code   int
	Detail struct {
		Code    string `json:"code"`
		Message string `json:"message"`
		Param   string `json:"param"`
		Type    string `json:"type"`
	} `json:"error"`
}

func (e *Error) Error() string {
	return e.Detail.Message
}

// Response to a Deletion request.
type DeleteResp struct {
	// ID of the Object that was deleted
	Id string `json:"id"`
	// Boolean value indicating object was successfully deleted.
	Deleted bool `json:"deleted"`
}
<commit_msg>Set Stripe-Version header to known good version<commit_after>package stripe

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"
)

// enable logging to print the request and reponses to stdout
var _log bool

// the API Key used to authenticate all Stripe API requests
var _key string

// the default URL for all Stripe API requests
var _url string = "https://api.stripe.com"

const apiVersion = "2013-02-13"

// SetUrl will override the default Stripe API URL. This is primarily used
// for unit testing.
func SetUrl(url string) {
	_url = url
}

// SetKey will set the default Stripe API key used to authenticate all Stripe
// API requests.
func SetKey(key string) {
	_key = key
}

// Available APIs
var (
	Charges       = new(ChargeClient)
	Coupons       = new(CouponClient)
	Customers     = new(CustomerClient)
	Invoices      = new(InvoiceClient)
	InvoiceItems  = new(InvoiceItemClient)
	Plans         = new(PlanClient)
	Subscriptions = new(SubscriptionClient)
	Tokens        = new(TokenClient)
)

// SetKeyEnv retrieves the Stripe API key using the STRIPE_API_KEY environment
// variable.
func SetKeyEnv() (err error) {
	_key = os.Getenv("STRIPE_API_KEY")
	if _key == "" {
		err = errors.New("STRIPE_API_KEY not found in environment")
	}
	return
}

// query submits an http.Request and parses the JSON-encoded http.Response,
// storing the result in the value pointed to by v.
func query(method, path string, values url.Values, v interface{}) error {
	// parse the stripe URL
	endpoint, err := url.Parse(_url)
	if err != nil {
		return err
	}

	// set the endpoint for the specific API
	endpoint.Path = path
	endpoint.User = url.User(_key)

	// if this is an http GET, add the url.Values to the endpoint
	if method == "GET" {
		endpoint.RawQuery = values.Encode()
	}

	// else if this is not a GET, encode the url.Values in the body.
	var reqBody io.Reader
	if method != "GET" && values != nil {
		reqBody = strings.NewReader(values.Encode())
	}

	// Log request if logging enabled
	if _log {
		fmt.Println("REQUEST: ", method, endpoint.String())
		fmt.Println(values.Encode())
	}

	// create the request
	req, err := http.NewRequest(method, endpoint.String(), reqBody)
	if err != nil {
		return err
	}

	req.Header.Set("Stripe-Version", apiVersion)

	// submit the http request
	r, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	// read the body of the http message into a byte array
	body, err := ioutil.ReadAll(r.Body)
	defer r.Body.Close()
	if err != nil {
		return err
	}

	// Log response if logging enabled
	if _log {
		fmt.Println("RESPONSE: ", r.StatusCode)
		fmt.Println(string(body))
	}

	// is this an error?
	if r.StatusCode != 200 {
		error := Error{}
		json.Unmarshal(body, &error)
		return &error
	}

	//parse the JSON response into the response object
	return json.Unmarshal(body, v)
}

// Error encapsulates an error returned by the Stripe REST API.
type Error struct {
	Code   int
	Detail struct {
		Code    string `json:"code"`
		Message string `json:"message"`
		Param   string `json:"param"`
		Type    string `json:"type"`
	} `json:"error"`
}

func (e *Error) Error() string {
	return e.Detail.Message
}

// Response to a Deletion request.
type DeleteResp struct {
	// ID of the Object that was deleted
	Id string `json:"id"`
	// Boolean value indicating object was successfully deleted.
	Deleted bool `json:"deleted"`
}
<|endoftext|>


<commit_before>package pipl

type PiplResponse struct {
	HttpStatusCode   int      `json:"@http_status_code"`
	VisibleSources   int      `json:"@visible_sources"`
	AvailableSources int      `json:"@available_sources"`
	SearchID         string   `json:"@search_id"`
	Person           Person   `json:"person"`
	PossiblePersons  []Person `json:"possible_persons"`
	Query            Person   `json:"query"`
	Warnings         []string `json:"warnings,omitempty"`
	Error            string   `json:"error,omitempty"`
}

type Person struct {
	Id             string         `json:"@id"`
	Search_pointer string         `json:"@search_pointer"`
	Match          float32        `json:"@match"`
	Dob            Dob            `json:"dob"`
	Gender         Gender         `json:"gender"`
	Ethnicity      Ethnicity      `json:"ethnicity"`
	Language       Language       `json:"language"`
	OriginCountry  OriginCountry  `json:"Country"`
	Addresses      []Address      `json:"addresses"`
	Relationships  []Relationship `json:"relationships"`
	Educations     []Education    `json:"educations"`
	Emails         []Email        `json:"emails"`
	Images         []Image        `json:"images"`
	Jobs           []Job          `json:"jobs"`
	Names          []Name         `json:"names"`
	Phones         []Phone        `json:"phones"`
	URLs           []URL          `json:"urls"`
	Usernames      []Username     `json:"usernames"`
	UserIDs        []UserID       `json:"user_ids"`
}

type DateRange struct {
	End   string `json:"end"`
	Start string `json:"start"`
}

type Dob struct {
	DateRange DateRange `json:"date_range"`
	Display   string    `json:"display"`
}

type Gender struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Content  string `json:"content"`
}

type OriginCountry struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Country  string `json:"content"`
}

type Ethnicity struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Content  string `json:"content"`
}

type Language struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Language string `json:"language"`
	Region   string `json:"region"`
	Display  string `json:"display"`
}

type Username struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Content  string `json:"content"`
}

type UserID struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Content  string `json:"content"`
}

type Relationship struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Type     string `json:"@type"`
	Subtype  string `json:"@subtype"`
	Person
}

type Address struct {
	Inferred  bool   `json:"@inferred,omitempty"`
	Type      string `json:"@type"`
	Country   string `json:"country"`
	State     string `json:"state"`
	City      string `json:"city"`
	Street    string `json:"street"`
	House     string `json:"house"`
	Apartment string `json:"apartment"`
	ZipCode   string `json:"zip_code"`
	POBox     string `json:"po_box"`
	Raw       string `json:"raw"`
	Display   string `json:"display"`
}

type Education struct {
	Inferred  bool      `json:"@inferred,omitempty"`
	DateRange DateRange `json:"date_range"`
	Degree    string    `json:"degree"`
	Display   string    `json:"display"`
	School    string    `json:"school"`
}

type Email struct {
	Inferred      bool   `json:"@inferred,omitempty"`
	Type          string `json:"@type"`
	Address       string `json:"address"`
	AddressMd5    string `json:"address_md5"`
	Disposable    string `json:"@disposable"`
	EmailProvider string `json:"@email_provider"`
}

type Image struct {
	Inferred       bool   `json:"@inferred,omitempty"`
	ThumbnailToken string `json:"thumbnail_token"`
	URL            string `json:"url"`
}

type Job struct {
	Inferred     bool      `json:"@inferred,omitempty"`
	Title        string    `json:"title"`
	Organization string    `json:"organization"`
	Industry     string    `json:"industry"`
	DateRange    DateRange `json:"date_range"`
	Display      string    `json:"display"`
}

type Name struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Type     string `json:"@type"`
	First    string `json:"first"`
	Middle   string `json:"middle"`
	Last     string `json:"last"`
	Prefix   string `json:"prefix"`
	Suffix   string `json:"suffix"`
	Raw      string `json:"raw"`
	Display  string `json:"display"`
}

type Tag struct {
	Classification string `json:"@classification"`
	Content        string `json:"content"`
}

type URL struct {
	Inferred bool   `json:"@inferred,omitempty"`
	SourceID string `json:"@source_id"`
	Name     string `json:"@name"`
	Category string `json:"@category"`
	Domain   string `json:"@domain"`
	URL      string `json:"url"`
}

type Phone struct {
	Inferred             bool   `json:"@inferred,omitempty"`
	Type                 string `json:"@type"`
	CountryCode          int    `json:"country_code"`
	Number               int    `json:"number"`
	Extension            string `json:"extension"`
	Raw                  string `json:"raw"`
	Display              string `json:"display"`
	DisplayInternational string `json:"display_international"`
}
<commit_msg>changed string to bool<commit_after>package pipl

type PiplResponse struct {
	HttpStatusCode   int      `json:"@http_status_code"`
	VisibleSources   int      `json:"@visible_sources"`
	AvailableSources int      `json:"@available_sources"`
	SearchID         string   `json:"@search_id"`
	Person           Person   `json:"person"`
	PossiblePersons  []Person `json:"possible_persons"`
	Query            Person   `json:"query"`
	Warnings         []string `json:"warnings,omitempty"`
	Error            string   `json:"error,omitempty"`
}

type Person struct {
	Id             string         `json:"@id"`
	Search_pointer string         `json:"@search_pointer"`
	Match          float32        `json:"@match"`
	Dob            Dob            `json:"dob"`
	Gender         Gender         `json:"gender"`
	Ethnicity      Ethnicity      `json:"ethnicity"`
	Language       Language       `json:"language"`
	OriginCountry  OriginCountry  `json:"Country"`
	Addresses      []Address      `json:"addresses"`
	Relationships  []Relationship `json:"relationships"`
	Educations     []Education    `json:"educations"`
	Emails         []Email        `json:"emails"`
	Images         []Image        `json:"images"`
	Jobs           []Job          `json:"jobs"`
	Names          []Name         `json:"names"`
	Phones         []Phone        `json:"phones"`
	URLs           []URL          `json:"urls"`
	Usernames      []Username     `json:"usernames"`
	UserIDs        []UserID       `json:"user_ids"`
}

type DateRange struct {
	End   string `json:"end"`
	Start string `json:"start"`
}

type Dob struct {
	DateRange DateRange `json:"date_range"`
	Display   string    `json:"display"`
}

type Gender struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Content  string `json:"content"`
}

type OriginCountry struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Country  string `json:"content"`
}

type Ethnicity struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Content  string `json:"content"`
}

type Language struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Language string `json:"language"`
	Region   string `json:"region"`
	Display  string `json:"display"`
}

type Username struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Content  string `json:"content"`
}

type UserID struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Content  string `json:"content"`
}

type Relationship struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Type     string `json:"@type"`
	Subtype  string `json:"@subtype"`
	Person
}

type Address struct {
	Inferred  bool   `json:"@inferred,omitempty"`
	Type      string `json:"@type"`
	Country   string `json:"country"`
	State     string `json:"state"`
	City      string `json:"city"`
	Street    string `json:"street"`
	House     string `json:"house"`
	Apartment string `json:"apartment"`
	ZipCode   string `json:"zip_code"`
	POBox     string `json:"po_box"`
	Raw       string `json:"raw"`
	Display   string `json:"display"`
}

type Education struct {
	Inferred  bool      `json:"@inferred,omitempty"`
	DateRange DateRange `json:"date_range"`
	Degree    string    `json:"degree"`
	Display   string    `json:"display"`
	School    string    `json:"school"`
}

type Email struct {
	Inferred      bool   `json:"@inferred,omitempty"`
	Type          string `json:"@type"`
	Address       string `json:"address"`
	AddressMd5    string `json:"address_md5"`
	Disposable    string `json:"@disposable"`
	EmailProvider bool   `json:"@email_provider"`
}

type Image struct {
	Inferred       bool   `json:"@inferred,omitempty"`
	ThumbnailToken string `json:"thumbnail_token"`
	URL            string `json:"url"`
}

type Job struct {
	Inferred     bool      `json:"@inferred,omitempty"`
	Title        string    `json:"title"`
	Organization string    `json:"organization"`
	Industry     string    `json:"industry"`
	DateRange    DateRange `json:"date_range"`
	Display      string    `json:"display"`
}

type Name struct {
	Inferred bool   `json:"@inferred,omitempty"`
	Type     string `json:"@type"`
	First    string `json:"first"`
	Middle   string `json:"middle"`
	Last     string `json:"last"`
	Prefix   string `json:"prefix"`
	Suffix   string `json:"suffix"`
	Raw      string `json:"raw"`
	Display  string `json:"display"`
}

type Tag struct {
	Classification string `json:"@classification"`
	Content        string `json:"content"`
}

type URL struct {
	Inferred bool   `json:"@inferred,omitempty"`
	SourceID string `json:"@source_id"`
	Name     string `json:"@name"`
	Category string `json:"@category"`
	Domain   string `json:"@domain"`
	URL      string `json:"url"`
}

type Phone struct {
	Inferred             bool   `json:"@inferred,omitempty"`
	Type                 string `json:"@type"`
	CountryCode          int    `json:"country_code"`
	Number               int    `json:"number"`
	Extension            string `json:"extension"`
	Raw                  string `json:"raw"`
	Display              string `json:"display"`
	DisplayInternational string `json:"display_international"`
}
<|endoftext|>


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

package command

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"go/types"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	"github.com/zchee/nvim-go/src/internal/guru"
	"golang.org/x/tools/cmd/guru/serial"
	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
)

// Definition parse definition from the current cursor.
func Definition(q *guru.Query) (*serial.Definition, error) {
	// First try the simple resolution done by parser.
	// It only works for intra-file references but it is very fast.
	// (Extending this approach to all the files of the package,
	// resolved using ast.NewPackage, was not worth the effort.)
	{
		qpos, err := fastQueryPos(q.Build, q.Pos)
		if err != nil {
			return nil, err
		}

		id, _ := qpos.path[0].(*ast.Ident)
		if id == nil {
			err := errors.New("no identifier here")
			return nil, err
		}

		// Did the parser resolve it to a local object?
		if obj := id.Obj; obj != nil && obj.Pos().IsValid() {
			return &serial.Definition{
				ObjPos: qpos.fset.Position(obj.Pos()).String(),
				Desc:   fmt.Sprintf("%s %s", obj.Kind, obj.Name),
			}, nil
		}

		// Qualified identifier?
		if pkg := guru.PackageForQualIdent(qpos.path, id); pkg != "" {
			srcdir := filepath.Dir(qpos.fset.File(qpos.start).Name())
			tok, pos, err := guru.FindPackageMember(q.Build, qpos.fset, srcdir, pkg, id.Name)
			if err != nil {
				return nil, err
			}
			return &serial.Definition{
				ObjPos: qpos.fset.Position(pos).String(),
				Desc:   fmt.Sprintf("%s %s.%s", tok, pkg, id.Name),
			}, nil
		}

		// Fall back on the type checker.
	}

	// Set loader.Config, same allowErrors() function result except CgoEnabled = false
	q.Build.CgoEnabled = true
	// Run the type checker.
	lconf := loader.Config{
		Build:       q.Build,
		AllowErrors: true,
		// AllErrors makes the parser always return an AST instead of
		// bailing out after 10 errors and returning an empty ast.File.
		ParserMode: parser.AllErrors,
		TypeChecker: types.Config{
			Error: func(err error) {},
		},
	}

	if _, err := importQueryPackage(q.Pos, &lconf); err != nil {
		return nil, err
	}

	// Load/parse/type-check the program.
	lprog, err := lconf.Load()
	if err != nil {
		return nil, err
	}

	qpos, err := parseQueryPos(lprog, q.Pos, false)
	if err != nil {
		return nil, err
	}

	id, _ := qpos.path[0].(*ast.Ident)
	if id == nil {
		err := errors.New("no identifier here")
		return nil, err
	}

	obj := qpos.info.ObjectOf(id)
	if obj == nil {
		// Happens for y in "switch y := x.(type)",
		// and the package declaration,
		// but I think that's all.
		err := errors.New("no object for identifier")
		return nil, err
	}

	if !obj.Pos().IsValid() {
		err := errors.Errorf("%s is built in", obj.Name())
		return nil, err
	}

	return &serial.Definition{
		ObjPos: qpos.fset.Position(obj.Pos()).String(),
		Desc:   qpos.ObjectString(obj),
	}, nil
}

// A QueryPos represents the position provided as input to a query:
// a textual extent in the program's source code, the AST node it
// corresponds to, and the package to which it belongs.
// Instances are created by parseQueryPos.
type queryPos struct {
	fset       *token.FileSet
	start, end token.Pos           // source extent of query
	path       []ast.Node          // AST path from query node to root of ast.File
	exact      bool                // 2nd result of PathEnclosingInterval
	info       *loader.PackageInfo // type info for the queried package (nil for fastQueryPos)
}

// TypeString prints type T relative to the query position.
func (qpos *queryPos) typeString(T types.Type) string {
	return types.TypeString(T, types.RelativeTo(qpos.info.Pkg))
}

// ObjectString prints object obj relative to the query position.
func (qpos *queryPos) ObjectString(obj types.Object) string {
	return types.ObjectString(obj, types.RelativeTo(qpos.info.Pkg))
}

// SelectionString prints selection sel relative to the query position.
func (qpos *queryPos) selectionString(sel *types.Selection) string {
	return types.SelectionString(sel, types.RelativeTo(qpos.info.Pkg))
}

// parseOctothorpDecimal returns the numeric value if s matches "#%d",
// otherwise -1.
func parseOctothorpDecimal(s string) int {
	if s != "" && s[0] == '#' {
		if s, err := strconv.ParseInt(s[1:], 10, 32); err == nil {
			return int(s)
		}
	}
	return -1
}

// parsePos parses a string of the form "file:pos" or
// file:start,end" where pos, start, end match #%d and represent byte
// offsets, and returns its components.
//
// (Numbers without a '#' prefix are reserved for future use,
// e.g. to indicate line/column positions.)
//
func parsePos(pos string) (filename string, startOffset, endOffset int, err error) {
	if pos == "" {
		err = fmt.Errorf("no source position specified")
		return
	}

	colon := strings.LastIndex(pos, ":")
	if colon < 0 {
		err = fmt.Errorf("bad position syntax %q", pos)
		return
	}
	filename, offset := pos[:colon], pos[colon+1:]
	startOffset = -1
	endOffset = -1
	if hyphen := strings.Index(offset, ","); hyphen < 0 {
		// e.g. "foo.go:#123"
		startOffset = parseOctothorpDecimal(offset)
		endOffset = startOffset
	} else {
		// e.g. "foo.go:#123,#456"
		startOffset = parseOctothorpDecimal(offset[:hyphen])
		endOffset = parseOctothorpDecimal(offset[hyphen+1:])
	}
	if startOffset < 0 || endOffset < 0 {
		err = fmt.Errorf("invalid offset %q in query position", offset)
		return
	}
	return
}

// fileOffsetToPos translates the specified file-relative byte offsets
// into token.Pos form.  It returns an error if the file was not found
// or the offsets were out of bounds.
//
func fileOffsetToPos(file *token.File, startOffset, endOffset int) (start, end token.Pos, err error) {
	// Range check [start..end], inclusive of both end-points.

	if 0 <= startOffset && startOffset <= file.Size() {
		start = file.Pos(int(startOffset))
	} else {
		err = fmt.Errorf("start position is beyond end of file")
		return
	}

	if 0 <= endOffset && endOffset <= file.Size() {
		end = file.Pos(int(endOffset))
	} else {
		err = fmt.Errorf("end position is beyond end of file")
		return
	}

	return
}

// sameFile returns true if x and y have the same basename and denote
// the same file.
//
func sameFile(x, y string) bool {
	if filepath.Base(x) == filepath.Base(y) { // (optimisation)
		if xi, err := os.Stat(x); err == nil {
			if yi, err := os.Stat(y); err == nil {
				return os.SameFile(xi, yi)
			}
		}
	}
	return false
}

// fastQueryPos parses the position string and returns a queryPos.
// It parses only a single file and does not run the type checker.
func fastQueryPos(ctxt *build.Context, pos string) (*queryPos, error) {
	filename, startOffset, endOffset, err := parsePos(pos)
	if err != nil {
		return nil, err
	}

	// Parse the file, opening it the file via the build.Context
	// so that we observe the effects of the -modified flag.
	fset := token.NewFileSet()
	fset.AddFile(filename, fset.Base(), startOffset)
	// cwd, _ := os.Getwd()
	cwd, _ := filepath.Split(filename)
	f, err := buildutil.ParseFile(fset, ctxt, nil, filepath.Clean(cwd), filename, parser.Mode(0))
	// ParseFile usually returns a partial file along with an error.
	// Only fail if there is no file.
	if f == nil {
		return nil, err
	}
	if !f.Pos().IsValid() {
		return nil, fmt.Errorf("%s is not a Go source file", filename)
	}

	start, end, err := fileOffsetToPos(fset.File(f.Pos()), startOffset, endOffset)
	if err != nil {
		return nil, err
	}

	path, exact := astutil.PathEnclosingInterval(f, start, end)
	if path == nil {
		return nil, fmt.Errorf("no syntax here")
	}

	return &queryPos{fset, start, end, path, exact, nil}, nil
}

// ParseQueryPos parses the source query position pos and returns the
// AST node of the loaded program lprog that it identifies.
// If needExact, it must identify a single AST subtree;
// this is appropriate for queries that allow fairly arbitrary syntax,
// e.g. "describe".
//
func parseQueryPos(lprog *loader.Program, pos string, needExact bool) (*queryPos, error) {
	filename, startOffset, endOffset, err := parsePos(pos)
	if err != nil {
		return nil, err
	}

	// Find the named file among those in the loaded program.
	var file *token.File
	lprog.Fset.Iterate(func(f *token.File) bool {
		if sameFile(filename, f.Name()) {
			file = f
			return false // done
		}
		return true // continue
	})
	if file == nil {
		return nil, fmt.Errorf("file %s not found in loaded program", filename)
	}

	start, end, err := fileOffsetToPos(file, startOffset, endOffset)
	if err != nil {
		return nil, err
	}
	info, path, exact := lprog.PathEnclosingInterval(start, end)
	if path == nil {
		return nil, fmt.Errorf("no syntax here")
	}
	if needExact && !exact {
		return nil, fmt.Errorf("ambiguous selection within %s", astutil.NodeDescription(path[0]))
	}
	return &queryPos{lprog.Fset, start, end, path, exact, info}, nil
}

// importQueryPackage finds the package P containing the
// query position and tells conf to import it.
// It returns the package's path.
func importQueryPackage(pos string, conf *loader.Config) (string, error) {
	fqpos, err := fastQueryPos(conf.Build, pos)
	if err != nil {
		return "", err // bad query
	}
	filename := fqpos.fset.File(fqpos.start).Name()

	_, importPath, err := guru.GuessImportPath(filename, conf.Build)
	if err != nil {
		return "", err // can't find GOPATH dir
	}

	// Check that it's possible to load the queried package.
	// (e.g. guru tests contain different 'package' decls in same dir.)
	// Keep consistent with logic in loader/util.go!
	cfg2 := *conf.Build
	cfg2.CgoEnabled = false
	bp, err := cfg2.Import(importPath, "", 0)
	if err != nil {
		return "", err // no files for package
	}

	switch pkgContainsFile(bp, filename) {
	case 'T':
		conf.ImportWithTests(importPath)
	case 'X':
		conf.ImportWithTests(importPath)
		importPath += "_test" // for TypeCheckFuncBodies
	case 'G':
		conf.Import(importPath)
	default:
		// This happens for ad-hoc packages like
		// $GOROOT/src/net/http/triv.go.
		return "", fmt.Errorf("package %q doesn't contain file %s",
			importPath, filename)
	}

	conf.TypeCheckFuncBodies = func(p string) bool { return p == importPath }

	return importPath, nil
}

// pkgContainsFile reports whether file was among the packages Go
// files, Test files, eXternal test files, or not found.
func pkgContainsFile(bp *build.Package, filename string) byte {
	for i, files := range [][]string{bp.GoFiles, bp.TestGoFiles, bp.XTestGoFiles} {
		for _, file := range files {
			if sameFile(filepath.Join(bp.Dir, file), filename) {
				return "GTX"[i]
			}
		}
	}
	return 0 // not found
}
<commit_msg>command/guru: fix get current working directory logic<commit_after>// Copyright 2016 The nvim-go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package command

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"go/types"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	"github.com/zchee/nvim-go/src/internal/guru"
	"golang.org/x/tools/cmd/guru/serial"
	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
)

// Definition parse definition from the current cursor.
func Definition(q *guru.Query) (*serial.Definition, error) {
	// First try the simple resolution done by parser.
	// It only works for intra-file references but it is very fast.
	// (Extending this approach to all the files of the package,
	// resolved using ast.NewPackage, was not worth the effort.)
	{
		qpos, err := fastQueryPos(q.Build, q.Pos)
		if err != nil {
			return nil, err
		}

		id, _ := qpos.path[0].(*ast.Ident)
		if id == nil {
			err := errors.New("no identifier here")
			return nil, err
		}

		// Did the parser resolve it to a local object?
		if obj := id.Obj; obj != nil && obj.Pos().IsValid() {
			return &serial.Definition{
				ObjPos: qpos.fset.Position(obj.Pos()).String(),
				Desc:   fmt.Sprintf("%s %s", obj.Kind, obj.Name),
			}, nil
		}

		// Qualified identifier?
		if pkg := guru.PackageForQualIdent(qpos.path, id); pkg != "" {
			srcdir := filepath.Dir(qpos.fset.File(qpos.start).Name())
			tok, pos, err := guru.FindPackageMember(q.Build, qpos.fset, srcdir, pkg, id.Name)
			if err != nil {
				return nil, err
			}
			return &serial.Definition{
				ObjPos: qpos.fset.Position(pos).String(),
				Desc:   fmt.Sprintf("%s %s.%s", tok, pkg, id.Name),
			}, nil
		}

		// Fall back on the type checker.
	}

	// Set loader.Config, same allowErrors() function result except CgoEnabled = false
	q.Build.CgoEnabled = true
	// Run the type checker.
	lconf := loader.Config{
		Build:       q.Build,
		AllowErrors: true,
		// AllErrors makes the parser always return an AST instead of
		// bailing out after 10 errors and returning an empty ast.File.
		ParserMode: parser.AllErrors,
		TypeChecker: types.Config{
			Error: func(err error) {},
		},
	}

	if _, err := importQueryPackage(q.Pos, &lconf); err != nil {
		return nil, err
	}

	// Load/parse/type-check the program.
	lprog, err := lconf.Load()
	if err != nil {
		return nil, err
	}

	qpos, err := parseQueryPos(lprog, q.Pos, false)
	if err != nil {
		return nil, err
	}

	id, _ := qpos.path[0].(*ast.Ident)
	if id == nil {
		err := errors.New("no identifier here")
		return nil, err
	}

	obj := qpos.info.ObjectOf(id)
	if obj == nil {
		// Happens for y in "switch y := x.(type)",
		// and the package declaration,
		// but I think that's all.
		err := errors.New("no object for identifier")
		return nil, err
	}

	if !obj.Pos().IsValid() {
		err := errors.Errorf("%s is built in", obj.Name())
		return nil, err
	}

	return &serial.Definition{
		ObjPos: qpos.fset.Position(obj.Pos()).String(),
		Desc:   qpos.ObjectString(obj),
	}, nil
}

// A QueryPos represents the position provided as input to a query:
// a textual extent in the program's source code, the AST node it
// corresponds to, and the package to which it belongs.
// Instances are created by parseQueryPos.
type queryPos struct {
	fset       *token.FileSet
	start, end token.Pos           // source extent of query
	path       []ast.Node          // AST path from query node to root of ast.File
	exact      bool                // 2nd result of PathEnclosingInterval
	info       *loader.PackageInfo // type info for the queried package (nil for fastQueryPos)
}

// TypeString prints type T relative to the query position.
func (qpos *queryPos) typeString(T types.Type) string {
	return types.TypeString(T, types.RelativeTo(qpos.info.Pkg))
}

// ObjectString prints object obj relative to the query position.
func (qpos *queryPos) ObjectString(obj types.Object) string {
	return types.ObjectString(obj, types.RelativeTo(qpos.info.Pkg))
}

// SelectionString prints selection sel relative to the query position.
func (qpos *queryPos) selectionString(sel *types.Selection) string {
	return types.SelectionString(sel, types.RelativeTo(qpos.info.Pkg))
}

// parseOctothorpDecimal returns the numeric value if s matches "#%d",
// otherwise -1.
func parseOctothorpDecimal(s string) int {
	if s != "" && s[0] == '#' {
		if s, err := strconv.ParseInt(s[1:], 10, 32); err == nil {
			return int(s)
		}
	}
	return -1
}

// parsePos parses a string of the form "file:pos" or
// file:start,end" where pos, start, end match #%d and represent byte
// offsets, and returns its components.
//
// (Numbers without a '#' prefix are reserved for future use,
// e.g. to indicate line/column positions.)
//
func parsePos(pos string) (filename string, startOffset, endOffset int, err error) {
	if pos == "" {
		err = fmt.Errorf("no source position specified")
		return
	}

	colon := strings.LastIndex(pos, ":")
	if colon < 0 {
		err = fmt.Errorf("bad position syntax %q", pos)
		return
	}
	filename, offset := pos[:colon], pos[colon+1:]
	startOffset = -1
	endOffset = -1
	if hyphen := strings.Index(offset, ","); hyphen < 0 {
		// e.g. "foo.go:#123"
		startOffset = parseOctothorpDecimal(offset)
		endOffset = startOffset
	} else {
		// e.g. "foo.go:#123,#456"
		startOffset = parseOctothorpDecimal(offset[:hyphen])
		endOffset = parseOctothorpDecimal(offset[hyphen+1:])
	}
	if startOffset < 0 || endOffset < 0 {
		err = fmt.Errorf("invalid offset %q in query position", offset)
		return
	}
	return
}

// fileOffsetToPos translates the specified file-relative byte offsets
// into token.Pos form.  It returns an error if the file was not found
// or the offsets were out of bounds.
//
func fileOffsetToPos(file *token.File, startOffset, endOffset int) (start, end token.Pos, err error) {
	// Range check [start..end], inclusive of both end-points.

	if 0 <= startOffset && startOffset <= file.Size() {
		start = file.Pos(int(startOffset))
	} else {
		err = fmt.Errorf("start position is beyond end of file")
		return
	}

	if 0 <= endOffset && endOffset <= file.Size() {
		end = file.Pos(int(endOffset))
	} else {
		err = fmt.Errorf("end position is beyond end of file")
		return
	}

	return
}

// sameFile returns true if x and y have the same basename and denote
// the same file.
//
func sameFile(x, y string) bool {
	if filepath.Base(x) == filepath.Base(y) { // (optimisation)
		if xi, err := os.Stat(x); err == nil {
			if yi, err := os.Stat(y); err == nil {
				return os.SameFile(xi, yi)
			}
		}
	}
	return false
}

// fastQueryPos parses the position string and returns a queryPos.
// It parses only a single file and does not run the type checker.
func fastQueryPos(ctxt *build.Context, pos string) (*queryPos, error) {
	filename, startOffset, endOffset, err := parsePos(pos)
	if err != nil {
		return nil, err
	}

	// Parse the file, opening it the file via the build.Context
	// so that we observe the effects of the -modified flag.
	fset := token.NewFileSet()
	fset.AddFile(filename, fset.Base(), startOffset)
	cwd, _ := os.Getwd()
	f, err := buildutil.ParseFile(fset, ctxt, nil, filepath.Clean(cwd), filename, parser.Mode(0))
	// ParseFile usually returns a partial file along with an error.
	// Only fail if there is no file.
	if f == nil {
		return nil, err
	}
	if !f.Pos().IsValid() {
		return nil, fmt.Errorf("%s is not a Go source file", filename)
	}

	start, end, err := fileOffsetToPos(fset.File(f.Pos()), startOffset, endOffset)
	if err != nil {
		return nil, err
	}

	path, exact := astutil.PathEnclosingInterval(f, start, end)
	if path == nil {
		return nil, fmt.Errorf("no syntax here")
	}

	return &queryPos{fset, start, end, path, exact, nil}, nil
}

// ParseQueryPos parses the source query position pos and returns the
// AST node of the loaded program lprog that it identifies.
// If needExact, it must identify a single AST subtree;
// this is appropriate for queries that allow fairly arbitrary syntax,
// e.g. "describe".
//
func parseQueryPos(lprog *loader.Program, pos string, needExact bool) (*queryPos, error) {
	filename, startOffset, endOffset, err := parsePos(pos)
	if err != nil {
		return nil, err
	}

	// Find the named file among those in the loaded program.
	var file *token.File
	lprog.Fset.Iterate(func(f *token.File) bool {
		if sameFile(filename, f.Name()) {
			file = f
			return false // done
		}
		return true // continue
	})
	if file == nil {
		return nil, fmt.Errorf("file %s not found in loaded program", filename)
	}

	start, end, err := fileOffsetToPos(file, startOffset, endOffset)
	if err != nil {
		return nil, err
	}
	info, path, exact := lprog.PathEnclosingInterval(start, end)
	if path == nil {
		return nil, fmt.Errorf("no syntax here")
	}
	if needExact && !exact {
		return nil, fmt.Errorf("ambiguous selection within %s", astutil.NodeDescription(path[0]))
	}
	return &queryPos{lprog.Fset, start, end, path, exact, info}, nil
}

// importQueryPackage finds the package P containing the
// query position and tells conf to import it.
// It returns the package's path.
func importQueryPackage(pos string, conf *loader.Config) (string, error) {
	fqpos, err := fastQueryPos(conf.Build, pos)
	if err != nil {
		return "", err // bad query
	}
	filename := fqpos.fset.File(fqpos.start).Name()

	_, importPath, err := guru.GuessImportPath(filename, conf.Build)
	if err != nil {
		return "", err // can't find GOPATH dir
	}

	// Check that it's possible to load the queried package.
	// (e.g. guru tests contain different 'package' decls in same dir.)
	// Keep consistent with logic in loader/util.go!
	cfg2 := *conf.Build
	cfg2.CgoEnabled = false
	bp, err := cfg2.Import(importPath, "", 0)
	if err != nil {
		return "", err // no files for package
	}

	switch pkgContainsFile(bp, filename) {
	case 'T':
		conf.ImportWithTests(importPath)
	case 'X':
		conf.ImportWithTests(importPath)
		importPath += "_test" // for TypeCheckFuncBodies
	case 'G':
		conf.Import(importPath)
	default:
		// This happens for ad-hoc packages like
		// $GOROOT/src/net/http/triv.go.
		return "", fmt.Errorf("package %q doesn't contain file %s",
			importPath, filename)
	}

	conf.TypeCheckFuncBodies = func(p string) bool { return p == importPath }

	return importPath, nil
}

// pkgContainsFile reports whether file was among the packages Go
// files, Test files, eXternal test files, or not found.
func pkgContainsFile(bp *build.Package, filename string) byte {
	for i, files := range [][]string{bp.GoFiles, bp.TestGoFiles, bp.XTestGoFiles} {
		for _, file := range files {
			if sameFile(filepath.Join(bp.Dir, file), filename) {
				return "GTX"[i]
			}
		}
	}
	return 0 // not found
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/kennygrant/sanitize"
)

const (
	ignoreTranslations = true
	downloadPath       = "./downloads/"
)

var (
	preferredMimeTypes    = []string{"video/webm", "video/mp4", "video/ogg", "audio/ogg", "audio/opus", "audio/mpeg", "application/x-subrip"}
	extensionForMimeTypes = make(map[string]string)
)

type Conference struct {
	Acronym        string `json:"acronym"`
	AspectRatio    string `json:"aspect_ratio"`
	ImagesURL      string `json:"images_url"`
	LogoURL        string `json:"logo_url"`
	RecordingsURL  string `json:"recordings_url"`
	ScheduleURL    string `json:"schedule_url"`
	Slug           string `json:"slug"`
	Title          string `json:"title"`
	UpdatedAt      string `json:"updated_at"`
	URL            string `json:"url"`
	WebgenLocation string `json:"webgen_location"`
}

type Conferences struct {
	Conferences []Conference `json:"conferences"`
}

func priorityForMimeType(mime string) int {
	for i, v := range preferredMimeTypes {
		if strings.ToLower(mime) == v {
			return i
		}
	}

	return -1
}

func findConferences(url string) (Conferences, error) {
	ci := Conferences{}

	r, err := http.Get(url)
	if err != nil {
		return ci, err
	}
	defer r.Body.Close()

	err = json.NewDecoder(r.Body).Decode(&ci)
	return ci, err
}

func main() {
	extensionForMimeTypes["video/webm"] = "webm"
	extensionForMimeTypes["video/mp4"] = "mp4"
	extensionForMimeTypes["video/ogg"] = "ogm"
	extensionForMimeTypes["audio/ogg"] = "ogg"
	extensionForMimeTypes["audio/opus"] = "opus"
	extensionForMimeTypes["audio/mpeg"] = "mp3"

	ci, err := findConferences("https://api.media.ccc.de/public/conferences")
	if err != nil {
		panic(err)
	}

	for _, v := range ci.Conferences {
		fmt.Printf("Conference: %s, URL: %s\n", v.Title, v.URL)

		events, err := findEvents(v.URL)
		if err != nil {
			panic(err)
		}

		for _, e := range events.Events {
			desc := strings.Replace(sanitize.HTML(e.Description), "\n", "", -1)
			if len(desc) > 48 {
				desc = desc[:45] + "..."
			}
			if len(desc) > 0 {
				desc = " - " + desc
			}
			fmt.Printf("Event: %s%s\n", e.Title, desc)

			media, err := findMedia(e.URL)
			if err != nil {
				panic(err)
			}

			bestMatch := Recording{}
			highestPriority := -1
			if len(media.Recordings) == 0 {
				panic("No recordings found for this event!")
			}
			for _, m := range media.Recordings {
				if ignoreTranslations && m.Language != e.OriginalLanguage {
					continue
				}

				if m.Width == 0 {
					fmt.Printf("\tFound audio (%s): %d minutes (HD: %t, %dMiB) %s\n", m.MimeType, m.Length/60, m.HighQuality, m.Size, m.URL)
				} else {
					fmt.Printf("\tFound video (%s): %d minutes, %dx%d (HD: %t, %dMiB) %s\n", m.MimeType, m.Length/60, m.Width, m.Height, m.HighQuality, m.Size, m.URL)
				}

				prio := priorityForMimeType(m.MimeType)
				if prio < 0 {
					panic("Unknown mimetype encountered:" + m.MimeType)
				}

				if prio < highestPriority ||
					(prio == highestPriority && m.Width > bestMatch.Width) ||
					highestPriority == -1 {
					highestPriority = prio
					bestMatch = m
				}
			}

			err = download(v, e, bestMatch)
			if err != nil {
				panic(err)
			}

			fmt.Println()
		}
	}

	fmt.Println("Done.")
}
<commit_msg>Added cli params and let users filter by conference acronym<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"strings"

	"github.com/kennygrant/sanitize"
)

const ()

var (
	preferredMimeTypes    = []string{"video/webm", "video/mp4", "video/ogg", "audio/ogg", "audio/opus", "audio/mpeg", "application/x-subrip"}
	extensionForMimeTypes = make(map[string]string)

	downloadPath       string
	acronym            string
	ignoreTranslations bool
)

type Conference struct {
	Acronym        string `json:"acronym"`
	AspectRatio    string `json:"aspect_ratio"`
	ImagesURL      string `json:"images_url"`
	LogoURL        string `json:"logo_url"`
	RecordingsURL  string `json:"recordings_url"`
	ScheduleURL    string `json:"schedule_url"`
	Slug           string `json:"slug"`
	Title          string `json:"title"`
	UpdatedAt      string `json:"updated_at"`
	URL            string `json:"url"`
	WebgenLocation string `json:"webgen_location"`
}

type Conferences struct {
	Conferences []Conference `json:"conferences"`
}

func priorityForMimeType(mime string) int {
	for i, v := range preferredMimeTypes {
		if strings.ToLower(mime) == v {
			return i
		}
	}

	return -1
}

func findConferences(url string) (Conferences, error) {
	ci := Conferences{}

	r, err := http.Get(url)
	if err != nil {
		return ci, err
	}
	defer r.Body.Close()

	err = json.NewDecoder(r.Body).Decode(&ci)
	return ci, err
}

func main() {
	flag.StringVar(&acronym, "acronym", "", "download only media belonging to this conference-acronym (e.g. '33c3')")
	flag.StringVar(&downloadPath, "destination", "./downloads/", "where to store downloaded media")
	flag.BoolVar(&ignoreTranslations, "ignoreTranslations", true, "do not download talk translations")
	flag.Parse()

	extensionForMimeTypes["video/webm"] = "webm"
	extensionForMimeTypes["video/mp4"] = "mp4"
	extensionForMimeTypes["video/ogg"] = "ogm"
	extensionForMimeTypes["audio/ogg"] = "ogg"
	extensionForMimeTypes["audio/opus"] = "opus"
	extensionForMimeTypes["audio/mpeg"] = "mp3"

	ci, err := findConferences("https://api.media.ccc.de/public/conferences")
	if err != nil {
		panic(err)
	}

	for _, v := range ci.Conferences {
		fmt.Printf("Conference: %s, acronym: %s (URL: %s)\n", v.Title, v.Acronym, v.URL)
		if len(acronym) > 0 && strings.ToLower(acronym) != strings.ToLower(v.Acronym) {
			continue
		}

		events, err := findEvents(v.URL)
		if err != nil {
			panic(err)
		}

		for _, e := range events.Events {
			desc := strings.Replace(sanitize.HTML(e.Description), "\n", "", -1)
			if len(desc) > 48 {
				desc = desc[:45] + "..."
			}
			if len(desc) > 0 {
				desc = " - " + desc
			}
			fmt.Printf("Event: %s%s\n", e.Title, desc)

			media, err := findMedia(e.URL)
			if err != nil {
				panic(err)
			}

			bestMatch := Recording{}
			highestPriority := -1
			if len(media.Recordings) == 0 {
				panic("No recordings found for this event!")
			}
			for _, m := range media.Recordings {
				if ignoreTranslations && m.Language != e.OriginalLanguage {
					continue
				}

				if m.Width == 0 {
					fmt.Printf("\tFound audio (%s): %d minutes (HD: %t, %dMiB) %s\n", m.MimeType, m.Length/60, m.HighQuality, m.Size, m.URL)
				} else {
					fmt.Printf("\tFound video (%s): %d minutes, %dx%d (HD: %t, %dMiB) %s\n", m.MimeType, m.Length/60, m.Width, m.Height, m.HighQuality, m.Size, m.URL)
				}

				prio := priorityForMimeType(m.MimeType)
				if prio < 0 {
					panic("Unknown mimetype encountered:" + m.MimeType)
				}

				if prio < highestPriority ||
					(prio == highestPriority && m.Width > bestMatch.Width) ||
					highestPriority == -1 {
					highestPriority = prio
					bestMatch = m
				}
			}

			err = download(v, e, bestMatch)
			if err != nil {
				panic(err)
			}

			fmt.Println()
		}
	}

	fmt.Println("Done.")
}
<|endoftext|>


<commit_before>// Package moka provides a mocking framework for the Go programming language.
// Moka works very well with the Ginkgo testing framework, but can be easily
// used with any other testing framework, including the testing package from
// the standard library.
package moka

// FailHandler is the type required for Moka fail handler functions. It matches
// the type of the Ginkgo `Fail` function.
type FailHandler func(message string, callerSkip ...int)

var globalFailHandler FailHandler

// RegisterDoublesFailHandler registers a function as the global fail handler
// used by newly instantiated Moka doubles.
func RegisterDoublesFailHandler(failHandler FailHandler) {
	globalFailHandler = failHandler
}

// AllowanceTarget wraps a Double to enable the configuration of allowed
// interactions on it.
type AllowanceTarget struct {
	double Double
}

// AllowDouble wraps a Double in an `AllowanceTarget`.
func AllowDouble(double Double) AllowanceTarget {
	return AllowanceTarget{double: double}
}

// To configures the interaction built by the provided `InteractionBuilder` on
// the wrapper `Double`.
func (t AllowanceTarget) To(interactionBuilder InteractionBuilder) {
	t.double.addInteraction(interactionBuilder.build())
}

// ExpectationTarget wraps a Double to enable the configuration of expected
// interactions on it.
type ExpectationTarget struct {
	double Double
}

// ExpectDouble wraps a Double in an `ExpectationTarget`.
func ExpectDouble(double Double) ExpectationTarget {
	return ExpectationTarget{double: double}
}

// To configures the interaction built by the provided `InteractionBuilder` on
// the wrapper `Double`.
func (t ExpectationTarget) To(interactionBuilder InteractionBuilder) {
	t.double.addInteraction(newExpectedInteraction(interactionBuilder.build()))
}

// VerifyCalls verifies that all expected interactions on the wrapper `Double`
// have actually happened.
func VerifyCalls(double Double) {
	double.verifyInteractions()
}

// InteractionBuilder provides a fluid API to build interactions.
type InteractionBuilder struct {
	methodName   string
	args         []interface{}
	returnValues []interface{}
}

// ReceiveCallTo allows to specify the method name of the interaction.
func ReceiveCallTo(methodName string) InteractionBuilder {
	return InteractionBuilder{methodName: methodName}
}

// With allows to specify the expected arguments of the interaction.
func (b InteractionBuilder) With(args ...interface{}) InteractionBuilder {
	return InteractionBuilder{methodName: b.methodName, args: args, returnValues: b.returnValues}
}

// AndReturn allows to specify the return value of the interaction.
func (b InteractionBuilder) AndReturn(returnValues ...interface{}) InteractionBuilder {
	return InteractionBuilder{methodName: b.methodName, args: b.args, returnValues: returnValues}
}

func (b InteractionBuilder) build() interaction {
	return newAllowedInteraction(b.methodName, b.args, b.returnValues)
}
<commit_msg>Fix typos in docstrings<commit_after>// Package moka provides a mocking framework for the Go programming language.
// Moka works very well with the Ginkgo testing framework, but can be easily
// used with any other testing framework, including the testing package from
// the standard library.
package moka

// FailHandler is the type required for Moka fail handler functions. It matches
// the type of the Ginkgo `Fail` function.
type FailHandler func(message string, callerSkip ...int)

var globalFailHandler FailHandler

// RegisterDoublesFailHandler registers a function as the global fail handler
// used by newly instantiated Moka doubles.
func RegisterDoublesFailHandler(failHandler FailHandler) {
	globalFailHandler = failHandler
}

// AllowanceTarget wraps a Double to enable the configuration of allowed
// interactions on it.
type AllowanceTarget struct {
	double Double
}

// AllowDouble wraps a Double in an `AllowanceTarget`.
func AllowDouble(double Double) AllowanceTarget {
	return AllowanceTarget{double: double}
}

// To configures the interaction built by the provided `InteractionBuilder` on
// the wrapped `Double`.
func (t AllowanceTarget) To(interactionBuilder InteractionBuilder) {
	t.double.addInteraction(interactionBuilder.build())
}

// ExpectationTarget wraps a Double to enable the configuration of expected
// interactions on it.
type ExpectationTarget struct {
	double Double
}

// ExpectDouble wraps a Double in an `ExpectationTarget`.
func ExpectDouble(double Double) ExpectationTarget {
	return ExpectationTarget{double: double}
}

// To configures the interaction built by the provided `InteractionBuilder` on
// the wrapped `Double`.
func (t ExpectationTarget) To(interactionBuilder InteractionBuilder) {
	t.double.addInteraction(newExpectedInteraction(interactionBuilder.build()))
}

// VerifyCalls verifies that all expected interactions on the wrapper `Double`
// have actually happened.
func VerifyCalls(double Double) {
	double.verifyInteractions()
}

// InteractionBuilder provides a fluid API to build interactions.
type InteractionBuilder struct {
	methodName   string
	args         []interface{}
	returnValues []interface{}
}

// ReceiveCallTo allows to specify the method name of the interaction.
func ReceiveCallTo(methodName string) InteractionBuilder {
	return InteractionBuilder{methodName: methodName}
}

// With allows to specify the expected arguments of the interaction.
func (b InteractionBuilder) With(args ...interface{}) InteractionBuilder {
	return InteractionBuilder{methodName: b.methodName, args: args, returnValues: b.returnValues}
}

// AndReturn allows to specify the return value of the interaction.
func (b InteractionBuilder) AndReturn(returnValues ...interface{}) InteractionBuilder {
	return InteractionBuilder{methodName: b.methodName, args: b.args, returnValues: returnValues}
}

func (b InteractionBuilder) build() interaction {
	return newAllowedInteraction(b.methodName, b.args, b.returnValues)
}
<|endoftext|>


<commit_before>package unit

import (
	"fmt"
	"log"
	"strings"

	"github.com/innovate-technologies/Dispatch/dispatchd/etcdcache"

	"strconv"

	"github.com/spf13/afero"

	"github.com/innovate-technologies/Dispatch/dispatchd/config"
	"github.com/innovate-technologies/Dispatch/dispatchd/etcdclient"
	"github.com/innovate-technologies/Dispatch/dispatchd/unit/state"
	"github.com/innovate-technologies/Dispatch/interfaces"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"github.com/coreos/go-systemd/dbus"
	"golang.org/x/net/context"
)

const unitPath = "/var/run/dispatch/"

var (
	// Config is a pointer need to be set to the main configuration
	Config *config.ConfigurationInfo
	ctx    = context.Background()
	//EtcdAPI is the etcd keys api
	EtcdAPI interfaces.EtcdAPI = etcdclient.GetEtcdv3()
	// DBusConnection is the connection to the system's D-Bus
	DBusConnection DBusConnectionInterface
	// FS is the file system to be used
	FS = afero.NewOsFs()
)

// UnitInterface is the interface to a Unit
type UnitInterface interface {
	Start()
	Stop()
	Restart()
	Create()
	PutOnQueue()
	SaveOnEtcd()
	Destroy()
	LoadAndWatch()
	Watch()
	SetState(state.State)
	SetDesiredState(s state.State)
}

// Unit is a struct containing all info of a specific unit
type Unit struct {
	Name         string `json:"name"`
	Machine      string `json:"machine"`
	Template     string `json:"template,omitempty"` // is set with template name if from Template
	Global       string `json:"global,omitempty"`   // is set with global name if from global
	State        state.State
	DesiredState state.State
	Ports        []int64 `json:"ports"`
	Constraints  map[string]string
	UnitContent  string `json:"unitContent"`
	onEtcd       bool
	onDisk       bool
	etcdName     string
	etcdCache    *etcdcache.EtcdCache
	disableCache bool
	runContext   context.Context
	runCancel    context.CancelFunc
}

// GetAll returns all units in our zone
func GetAll() ([]Unit, error) {
	cache, err := etcdcache.NewForPrefix(fmt.Sprintf("/dispatch/%s/units", Config.Zone))
	if err != nil {
		return nil, err
	}

	units := []Unit{}
	hadUnitNames := map[string]bool{}

	for _, kv := range cache.GetAll() {
		pathParts := strings.Split(string(kv.Key), "/")
		if _, ok := hadUnitNames[pathParts[4]]; !ok {
			units = append(units, NewFromEtcd(pathParts[4]))
			hadUnitNames[pathParts[4]] = true
		}

	}

	return units, nil
}

// New returns a new Unit
func New() Unit {
	setUpDBus()
	return Unit{}
}

// NewFromEtcd creates a new unit with info from etcd
func NewFromEtcd(name string) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}

	unit := New()
	fillFromEtcd(&unit, name)

	return unit
}

// NewFromEtcdWithCache creates a new unit with info from etcd using a specified cache
func NewFromEtcdWithCache(name string, cache *etcdcache.EtcdCache) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}
	unit := New()
	unit.etcdCache = cache
	fillFromEtcd(&unit, name)

	return unit
}

func fillFromEtcd(unit *Unit, name string) {
	unit.etcdName = name
	unit.onEtcd = true
	unit.Name = unit.getKeyFromEtcd("name")
	unit.Machine = unit.getKeyFromEtcd("machine")
	unit.Template = unit.getKeyFromEtcd("template")
	unit.Global = unit.getKeyFromEtcd("global")
	unit.UnitContent = unit.getKeyFromEtcd("unit")
	unit.DesiredState = state.ForString(unit.getKeyFromEtcd("desiredState"))
	unit.State = state.ForString(unit.getKeyFromEtcd("state"))

	unit.Ports = []int64{}
	portsStringArray := strings.Split(unit.getKeyFromEtcd("ports"), ",")
	for _, portString := range portsStringArray {
		port, err := strconv.ParseInt(portString, 10, 64)
		if err == nil {
			unit.Ports = append(unit.Ports, port)
		}
	}
}

// Start starts the specific unit
func (unit *Unit) Start() {
	log.Println("Starting unit", unit.Name)
	unit.SetState(state.Starting)
	c := make(chan string)
	_, err := DBusConnection.StartUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error starting unit", unit.Name, err)
		return
	}
	<-c
	unit.SetState(state.Active)
}

// Stop stops the unit
func (unit *Unit) Stop() {
	log.Println("Stopping unit", unit.Name)
	c := make(chan string)
	_, err := DBusConnection.StopUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error stopping unit", unit.Name, err)
		log.Println("Killing unit", unit.Name)
		DBusConnection.KillUnit(unit.Name, 9) // the big guns!
		unit.SetState(state.Dead)
		return
	}
	result := <-c
	if result == "done" {
		log.Println("Stopped unit", unit.Name)
		unit.SetState(state.Dead)
	}
}

// Restart restarts the unit
func (unit *Unit) Restart() {
	unit.Stop()
	unit.Start()
}

// Create writes the unit to the disk
func (unit *Unit) Create() {
	if unit.Name == "" {
		log.Println("Error starting unit with no name")
		if unit.etcdName != "" {
			// Faulty unit
			unit.Name = unit.etcdName
			unit.Destroy()
		}
		return // can't create file without a name
	}
	unit.runContext, unit.runCancel = context.WithCancel(context.Background())

	thisUnitPath := unitPath + unit.Name

	fileContent := []byte(unit.UnitContent)

	FS.MkdirAll(unitPath, 0755)

	FS.Remove(thisUnitPath) // make sure the old unit is gone
	file, err := FS.Create(thisUnitPath)
	if err != nil {
		panic(err)
	}
	_, err = file.Write(fileContent)
	if err != nil {
		panic(err)
	}

	c := make(chan string)
	_, stopErr := DBusConnection.StopUnit(unit.Name, "fail", c) // stop unit to make sure new one is loaded
	if stopErr == nil {
		<-c // wait on completion
	}

	unit.onDisk = true
	_, dberr := DBusConnection.LinkUnitFiles([]string{thisUnitPath}, true, true)
	fmt.Println(dberr)
	DBusConnection.Reload()
}

// PutOnQueue places a specific unit on the queue
func (unit *Unit) PutOnQueue() {
	if unit.Global != "" {
		return
	}
	log.Println("Placing", unit.Name, "on queue")
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/queue/%s", Config.Zone, unit.Name), unit.Name)
}

func (unit *Unit) normalizeName() {
	if !strings.HasSuffix(unit.Name, ".service") {
		unit.Name += ".service"
	}
}

// SaveOnEtcd saves the unit to etcd
func (unit *Unit) SaveOnEtcd() {
	unit.normalizeName()

	log.Println("Saving", unit.Name, "to etcd")

	unit.setKeyOnEtcd("name", unit.Name)
	unit.setKeyOnEtcd("unit", unit.UnitContent)
	unit.setKeyOnEtcd("template", unit.Template)
	unit.setKeyOnEtcd("desiredState", unit.DesiredState.String())

	portStrings := []string{}
	for _, port := range unit.Ports {
		portStrings = append(portStrings, strconv.FormatInt(port, 10))
	}
	unit.setKeyOnEtcd("ports", strings.Join(portStrings, ","))

	if unit.Global != "" {
		unit.setKeyOnEtcd("global", unit.Global)
		EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/globals/%s", Config.Zone, unit.Name), unit.Name)
	}

	unit.onEtcd = true
}

// Destroy destroys the given unit
func (unit *Unit) Destroy() {
	log.Println("Destroying unit", unit.Name)
	if unit.runCancel != nil {
		unit.runCancel()
	}

	unit.Stop() // just making sure

	FS.Remove(unitPath + unit.Name)
	unit.onDisk = false
	DBusConnection.Reload()

	if unit.Name == "" {
		// oopsie
		unit.onEtcd = false //probably not
		return
	}

	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/units/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/queue/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/machines/%s/units/%s", Config.Zone, unit.Machine, unit.Name), etcd.WithPrefix())
	if unit.Global != "" {
		EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/globals/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	}
	unit.etcdCache = etcdcache.New() // clear out all old cache!
	unit.onEtcd = false
}

// LoadAndWatch loads the unit to the system and follows the desired state
func (unit *Unit) LoadAndWatch() {
	unit.Create()
	unit.becomeDesiredState()
	go unit.Watch()
}

func (unit *Unit) becomeDesiredState() {
	fmt.Println("desiredState:", unit.DesiredState)
	if unit.DesiredState == state.Active {
		unit.Start()
	} else if unit.DesiredState == state.Dead {
		unit.Stop()
	} else if unit.DesiredState == state.Destroy {
		unit.Destroy()
	}
}

// Watch creates and etcd watcher for the desired state of a specific unit
func (unit *Unit) Watch() {
	chans := EtcdAPI.Watch(unit.runContext, fmt.Sprintf("/dispatch/%s/units/%s/desiredState", Config.Zone, unit.Name), etcd.WithPrefix())
	for resp := range chans {
		for _, ev := range resp.Events {
			if ev.IsModify() || ev.IsCreate() {
				unit.DesiredState = state.ForString(string(ev.Kv.Value))
				unit.becomeDesiredState()
			}
			if ev.Type == mvccpb.DELETE {
				break
			}
		}
	}
}

// WaitOnDestroy waits for the unit to enter a destroyed state
func (unit *Unit) WaitOnDestroy() {
	response, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	if err != nil || response.Count == 0 { // Destroyed already
		return
	}

	chans := EtcdAPI.Watch(context.Background(), fmt.Sprintf("/dispatch/%s/units/%s/name", Config.Zone, unit.Name), etcd.WithPrefix())
	for resp := range chans {
		for _, ev := range resp.Events {
			if ev.Type == mvccpb.DELETE {
				break
			}
		}
	}
}

// WaitOnState waits for the unit to enter a specific state
func (unit *Unit) WaitOnState(s state.State) {
	stateInfo, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name))
	if err != nil || stateInfo.Count == 0 { // hmmm
		return
	}
	if string(stateInfo.Kvs[0].Value) == s.String() {
		return
	}

	chans := EtcdAPI.Watch(context.Background(), fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name), etcd.WithPrefix())
	for resp := range chans {
		for _, ev := range resp.Events {
			if string(ev.Kv.Value) == s.String() {
				break
			}
		}
	}
}

// SetState sets the state of the unit and updates etcd
func (unit *Unit) SetState(s state.State) {
	if unit.Global != "" || !unit.isHealthy() {
		return
	}
	unit.State = s
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name), s.String())
}

// SetDesiredState sets the desiredstate of the unit and updates etcd
func (unit *Unit) SetDesiredState(s state.State) {
	unit.DesiredState = s
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/desiredState", Config.Zone, unit.Name), s.String())
}

func (unit *Unit) isHealthy() bool {
	unit.disableCache = true
	name := unit.getKeyFromEtcd("name")
	if name == "" {
		unit.Destroy()
		return false
	}
	return true
}

func (unit *Unit) getKeyFromEtcd(key string) string {
	if unit.etcdCache != nil {
		if kv, err := unit.etcdCache.Get(key); err == nil {
			return string(kv.Value)
		}
	}
	if unit.etcdName == "" && unit.Name != "" {
		unit.etcdName = unit.Name
	}
	response, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s/%s", Config.Zone, unit.etcdName, key))
	if err != nil || response.Count == 0 {
		return ""
	}
	return string(response.Kvs[0].Value)
}

func (unit *Unit) setKeyOnEtcd(key, content string) {
	if unit.etcdCache != nil {
		unit.etcdCache.Invalidate(key)
	}
	if unit.etcdName == "" && unit.Name != "" {
		unit.etcdName = unit.Name
	}
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/%s", Config.Zone, unit.etcdName, key), content)
}

func setUpDBus() {
	if DBusConnection == nil {
		var err error
		DBusConnection, err = dbus.NewSystemdConnection()
		if err != nil {
			panic(err)
		}
	}
}

// KillAllOldUnits makes sure all old Dispatch spawned unit files on the system are deleted
func KillAllOldUnits() {
	setUpDBus()

	FS.MkdirAll(unitPath, 0755) // maybe we have a first run
	files, err := afero.Afero{Fs: FS}.ReadDir(unitPath)
	if err != nil {
		panic(err)
	}
	for _, file := range files {
		log.Println("Stopping unit", file.Name())
		DBusConnection.KillUnit(file.Name(), 9) // do we care at this point?
		FS.Remove(unitPath + file.Name())
	}

	DBusConnection.Reload()
}
<commit_msg>Make WaitOnDestroy great again (spoiler: but crappy coded)<commit_after>package unit

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/innovate-technologies/Dispatch/dispatchd/etcdcache"

	"strconv"

	"github.com/spf13/afero"

	"github.com/innovate-technologies/Dispatch/dispatchd/config"
	"github.com/innovate-technologies/Dispatch/dispatchd/etcdclient"
	"github.com/innovate-technologies/Dispatch/dispatchd/unit/state"
	"github.com/innovate-technologies/Dispatch/interfaces"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"github.com/coreos/go-systemd/dbus"
	"golang.org/x/net/context"
)

const unitPath = "/var/run/dispatch/"

var (
	// Config is a pointer need to be set to the main configuration
	Config *config.ConfigurationInfo
	ctx    = context.Background()
	//EtcdAPI is the etcd keys api
	EtcdAPI interfaces.EtcdAPI = etcdclient.GetEtcdv3()
	// DBusConnection is the connection to the system's D-Bus
	DBusConnection DBusConnectionInterface
	// FS is the file system to be used
	FS = afero.NewOsFs()
)

// UnitInterface is the interface to a Unit
type UnitInterface interface {
	Start()
	Stop()
	Restart()
	Create()
	PutOnQueue()
	SaveOnEtcd()
	Destroy()
	LoadAndWatch()
	Watch()
	SetState(state.State)
	SetDesiredState(s state.State)
}

// Unit is a struct containing all info of a specific unit
type Unit struct {
	Name         string `json:"name"`
	Machine      string `json:"machine"`
	Template     string `json:"template,omitempty"` // is set with template name if from Template
	Global       string `json:"global,omitempty"`   // is set with global name if from global
	State        state.State
	DesiredState state.State
	Ports        []int64 `json:"ports"`
	Constraints  map[string]string
	UnitContent  string `json:"unitContent"`
	onEtcd       bool
	onDisk       bool
	etcdName     string
	etcdCache    *etcdcache.EtcdCache
	disableCache bool
	runContext   context.Context
	runCancel    context.CancelFunc
}

// GetAll returns all units in our zone
func GetAll() ([]Unit, error) {
	cache, err := etcdcache.NewForPrefix(fmt.Sprintf("/dispatch/%s/units", Config.Zone))
	if err != nil {
		return nil, err
	}

	units := []Unit{}
	hadUnitNames := map[string]bool{}

	for _, kv := range cache.GetAll() {
		pathParts := strings.Split(string(kv.Key), "/")
		if _, ok := hadUnitNames[pathParts[4]]; !ok {
			units = append(units, NewFromEtcd(pathParts[4]))
			hadUnitNames[pathParts[4]] = true
		}

	}

	return units, nil
}

// New returns a new Unit
func New() Unit {
	setUpDBus()
	return Unit{}
}

// NewFromEtcd creates a new unit with info from etcd
func NewFromEtcd(name string) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}

	unit := New()
	fillFromEtcd(&unit, name)

	return unit
}

// NewFromEtcdWithCache creates a new unit with info from etcd using a specified cache
func NewFromEtcdWithCache(name string, cache *etcdcache.EtcdCache) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}
	unit := New()
	unit.etcdCache = cache
	fillFromEtcd(&unit, name)

	return unit
}

func fillFromEtcd(unit *Unit, name string) {
	unit.etcdName = name
	unit.onEtcd = true
	unit.Name = unit.getKeyFromEtcd("name")
	unit.Machine = unit.getKeyFromEtcd("machine")
	unit.Template = unit.getKeyFromEtcd("template")
	unit.Global = unit.getKeyFromEtcd("global")
	unit.UnitContent = unit.getKeyFromEtcd("unit")
	unit.DesiredState = state.ForString(unit.getKeyFromEtcd("desiredState"))
	unit.State = state.ForString(unit.getKeyFromEtcd("state"))

	unit.Ports = []int64{}
	portsStringArray := strings.Split(unit.getKeyFromEtcd("ports"), ",")
	for _, portString := range portsStringArray {
		port, err := strconv.ParseInt(portString, 10, 64)
		if err == nil {
			unit.Ports = append(unit.Ports, port)
		}
	}
}

// Start starts the specific unit
func (unit *Unit) Start() {
	log.Println("Starting unit", unit.Name)
	unit.SetState(state.Starting)
	c := make(chan string)
	_, err := DBusConnection.StartUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error starting unit", unit.Name, err)
		return
	}
	<-c
	unit.SetState(state.Active)
}

// Stop stops the unit
func (unit *Unit) Stop() {
	log.Println("Stopping unit", unit.Name)
	c := make(chan string)
	_, err := DBusConnection.StopUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error stopping unit", unit.Name, err)
		log.Println("Killing unit", unit.Name)
		DBusConnection.KillUnit(unit.Name, 9) // the big guns!
		unit.SetState(state.Dead)
		return
	}
	result := <-c
	if result == "done" {
		log.Println("Stopped unit", unit.Name)
		unit.SetState(state.Dead)
	}
}

// Restart restarts the unit
func (unit *Unit) Restart() {
	unit.Stop()
	unit.Start()
}

// Create writes the unit to the disk
func (unit *Unit) Create() {
	if unit.Name == "" {
		log.Println("Error starting unit with no name")
		if unit.etcdName != "" {
			// Faulty unit
			unit.Name = unit.etcdName
			unit.Destroy()
		}
		return // can't create file without a name
	}
	unit.runContext, unit.runCancel = context.WithCancel(context.Background())

	thisUnitPath := unitPath + unit.Name

	fileContent := []byte(unit.UnitContent)

	FS.MkdirAll(unitPath, 0755)

	FS.Remove(thisUnitPath) // make sure the old unit is gone
	file, err := FS.Create(thisUnitPath)
	if err != nil {
		panic(err)
	}
	_, err = file.Write(fileContent)
	if err != nil {
		panic(err)
	}

	c := make(chan string)
	_, stopErr := DBusConnection.StopUnit(unit.Name, "fail", c) // stop unit to make sure new one is loaded
	if stopErr == nil {
		<-c // wait on completion
	}

	unit.onDisk = true
	_, dberr := DBusConnection.LinkUnitFiles([]string{thisUnitPath}, true, true)
	fmt.Println(dberr)
	DBusConnection.Reload()
}

// PutOnQueue places a specific unit on the queue
func (unit *Unit) PutOnQueue() {
	if unit.Global != "" {
		return
	}
	log.Println("Placing", unit.Name, "on queue")
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/queue/%s", Config.Zone, unit.Name), unit.Name)
}

func (unit *Unit) normalizeName() {
	if !strings.HasSuffix(unit.Name, ".service") {
		unit.Name += ".service"
	}
}

// SaveOnEtcd saves the unit to etcd
func (unit *Unit) SaveOnEtcd() {
	unit.normalizeName()

	log.Println("Saving", unit.Name, "to etcd")

	unit.setKeyOnEtcd("name", unit.Name)
	unit.setKeyOnEtcd("unit", unit.UnitContent)
	unit.setKeyOnEtcd("template", unit.Template)
	unit.setKeyOnEtcd("desiredState", unit.DesiredState.String())

	portStrings := []string{}
	for _, port := range unit.Ports {
		portStrings = append(portStrings, strconv.FormatInt(port, 10))
	}
	unit.setKeyOnEtcd("ports", strings.Join(portStrings, ","))

	if unit.Global != "" {
		unit.setKeyOnEtcd("global", unit.Global)
		EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/globals/%s", Config.Zone, unit.Name), unit.Name)
	}

	unit.onEtcd = true
}

// Destroy destroys the given unit
func (unit *Unit) Destroy() {
	log.Println("Destroying unit", unit.Name)
	if unit.runCancel != nil {
		unit.runCancel()
	}

	unit.Stop() // just making sure

	FS.Remove(unitPath + unit.Name)
	unit.onDisk = false
	DBusConnection.Reload()

	if unit.Name == "" {
		// oopsie
		unit.onEtcd = false //probably not
		return
	}

	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/units/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/queue/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/machines/%s/units/%s", Config.Zone, unit.Machine, unit.Name), etcd.WithPrefix())
	if unit.Global != "" {
		EtcdAPI.Delete(ctx, fmt.Sprintf("/dispatch/%s/globals/%s", Config.Zone, unit.Name), etcd.WithPrefix())
	}
	unit.etcdCache = etcdcache.New() // clear out all old cache!
	unit.onEtcd = false
}

// LoadAndWatch loads the unit to the system and follows the desired state
func (unit *Unit) LoadAndWatch() {
	unit.Create()
	unit.becomeDesiredState()
	go unit.Watch()
}

func (unit *Unit) becomeDesiredState() {
	fmt.Println("desiredState:", unit.DesiredState)
	if unit.DesiredState == state.Active {
		unit.Start()
	} else if unit.DesiredState == state.Dead {
		unit.Stop()
	} else if unit.DesiredState == state.Destroy {
		unit.Destroy()
	}
}

// Watch creates and etcd watcher for the desired state of a specific unit
func (unit *Unit) Watch() {
	chans := EtcdAPI.Watch(unit.runContext, fmt.Sprintf("/dispatch/%s/units/%s/desiredState", Config.Zone, unit.Name), etcd.WithPrefix())
	for resp := range chans {
		for _, ev := range resp.Events {
			if ev.IsModify() || ev.IsCreate() {
				unit.DesiredState = state.ForString(string(ev.Kv.Value))
				unit.becomeDesiredState()
			}
			if ev.Type == mvccpb.DELETE {
				break
			}
		}
	}
}

// WaitOnDestroy waits for the unit to enter a destroyed state
func (unit *Unit) WaitOnDestroy() {
	unit.disableCache = true
	for unit.getKeyFromEtcd("name") != "" {
		time.Sleep(100 * time.Millisecond)
	}
}

// WaitOnState waits for the unit to enter a specific state
func (unit *Unit) WaitOnState(s state.State) {
	stateInfo, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name))
	if err != nil || stateInfo.Count == 0 { // hmmm
		return
	}
	if string(stateInfo.Kvs[0].Value) == s.String() {
		return
	}

	chans := EtcdAPI.Watch(context.Background(), fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name), etcd.WithPrefix())
	for resp := range chans {
		for _, ev := range resp.Events {
			if string(ev.Kv.Value) == s.String() {
				break
			}
		}
	}
}

// SetState sets the state of the unit and updates etcd
func (unit *Unit) SetState(s state.State) {
	if unit.Global != "" || !unit.isHealthy() {
		return
	}
	unit.State = s
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/state", Config.Zone, unit.Name), s.String())
}

// SetDesiredState sets the desiredstate of the unit and updates etcd
func (unit *Unit) SetDesiredState(s state.State) {
	unit.DesiredState = s
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/desiredState", Config.Zone, unit.Name), s.String())
}

func (unit *Unit) isHealthy() bool {
	unit.disableCache = true
	name := unit.getKeyFromEtcd("name")
	if name == "" {
		unit.Destroy()
		return false
	}
	return true
}

func (unit *Unit) getKeyFromEtcd(key string) string {
	if unit.etcdCache != nil {
		if kv, err := unit.etcdCache.Get(key); err == nil {
			return string(kv.Value)
		}
	}
	if unit.etcdName == "" && unit.Name != "" {
		unit.etcdName = unit.Name
	}
	response, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/%s/units/%s/%s", Config.Zone, unit.etcdName, key))
	if err != nil || response.Count == 0 {
		return ""
	}
	return string(response.Kvs[0].Value)
}

func (unit *Unit) setKeyOnEtcd(key, content string) {
	if unit.etcdCache != nil {
		unit.etcdCache.Invalidate(key)
	}
	if unit.etcdName == "" && unit.Name != "" {
		unit.etcdName = unit.Name
	}
	EtcdAPI.Put(ctx, fmt.Sprintf("/dispatch/%s/units/%s/%s", Config.Zone, unit.etcdName, key), content)
}

func setUpDBus() {
	if DBusConnection == nil {
		var err error
		DBusConnection, err = dbus.NewSystemdConnection()
		if err != nil {
			panic(err)
		}
	}
}

// KillAllOldUnits makes sure all old Dispatch spawned unit files on the system are deleted
func KillAllOldUnits() {
	setUpDBus()

	FS.MkdirAll(unitPath, 0755) // maybe we have a first run
	files, err := afero.Afero{Fs: FS}.ReadDir(unitPath)
	if err != nil {
		panic(err)
	}
	for _, file := range files {
		log.Println("Stopping unit", file.Name())
		DBusConnection.KillUnit(file.Name(), 9) // do we care at this point?
		FS.Remove(unitPath + file.Name())
	}

	DBusConnection.Reload()
}
<|endoftext|>