<commit_before>package graph

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"github.com/docker/distribution/digest"
	"github.com/docker/docker/daemon/events"
	"github.com/docker/docker/graph/tags"
	"github.com/docker/docker/image"
	"github.com/docker/docker/pkg/broadcaster"
	"github.com/docker/docker/pkg/parsers"
	"github.com/docker/docker/pkg/stringid"
	"github.com/docker/docker/registry"
	"github.com/docker/docker/utils"
	"github.com/docker/libtrust"
)

// ErrNameIsNotExist returned when there is no image with requested name.
var ErrNameIsNotExist = errors.New("image with specified name does not exist")

// TagStore manages repositories. It encompasses the Graph used for versioned
// storage, as well as various services involved in pushing and pulling
// repositories.
type TagStore struct {
	path  string
	graph *Graph
	// Repositories is a map of repositories, indexed by name.
	Repositories map[string]Repository
	trustKey     libtrust.PrivateKey
	sync.Mutex
	// FIXME: move push/pull-related fields
	// to a helper type
	pullingPool     map[string]*broadcaster.Buffered
	pushingPool     map[string]*broadcaster.Buffered
	registryService *registry.Service
	eventsService   *events.Events
}

// Repository maps tags to image IDs.
type Repository map[string]string

// Update updates repository mapping with content of repository 'u'.
func (r Repository) Update(u Repository) {
	for k, v := range u {
		r[k] = v
	}
}

// Contains returns true if the contents of Repository u are wholly contained
// in Repository r.
func (r Repository) Contains(u Repository) bool {
	for k, v := range u {
		// if u's key is not present in r OR u's key is present, but not the same value
		if rv, ok := r[k]; !ok || (ok && rv != v) {
			return false
		}
	}
	return true
}

// TagStoreConfig provides parameters for a new TagStore.
type TagStoreConfig struct {
	// Graph is the versioned image store
	Graph *Graph
	// Key is the private key to use for signing manifests.
	Key libtrust.PrivateKey
	// Registry is the registry service to use for TLS configuration and
	// endpoint lookup.
	Registry *registry.Service
	// Events is the events service to use for logging.
	Events *events.Events
}

// NewTagStore creates a new TagStore at specified path, using the parameters
// and services provided in cfg.
func NewTagStore(path string, cfg *TagStoreConfig) (*TagStore, error) {
	abspath, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}

	store := &TagStore{
		path:            abspath,
		graph:           cfg.Graph,
		trustKey:        cfg.Key,
		Repositories:    make(map[string]Repository),
		pullingPool:     make(map[string]*broadcaster.Buffered),
		pushingPool:     make(map[string]*broadcaster.Buffered),
		registryService: cfg.Registry,
		eventsService:   cfg.Events,
	}
	// Load the json file if it exists, otherwise create it.
	if err := store.reload(); os.IsNotExist(err) {
		if err := store.save(); err != nil {
			return nil, err
		}
	} else if err != nil {
		return nil, err
	}
	return store, nil
}

func (store *TagStore) save() error {
	// Store the json ball
	jsonData, err := json.Marshal(store)
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(store.path, jsonData, 0600); err != nil {
		return err
	}
	return nil
}

func (store *TagStore) reload() error {
	f, err := os.Open(store.path)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := json.NewDecoder(f).Decode(&store); err != nil {
		return err
	}
	return nil
}

// LookupImage returns pointer to an Image struct corresponding to the given
// name. The name can include an optional tag; otherwise the default tag will
// be used.
func (store *TagStore) LookupImage(name string) (*image.Image, error) {
	// FIXME: standardize on returning nil when the image doesn't exist, and err for everything else
	// (so we can pass all errors here)
	repoName, ref := parsers.ParseRepositoryTag(name)
	if ref == "" {
		ref = tags.DefaultTag
	}
	var (
		err error
		img *image.Image
	)

	img, err = store.GetImage(repoName, ref)
	if err != nil {
		return nil, err
	}

	if img != nil {
		return img, nil
	}

	// name must be an image ID.
	store.Lock()
	defer store.Unlock()
	if img, err = store.graph.Get(name); err != nil {
		return nil, err
	}

	return img, nil
}

// GetID returns ID for image name.
func (store *TagStore) GetID(name string) (string, error) {
	repoName, ref := parsers.ParseRepositoryTag(name)
	if ref == "" {
		ref = tags.DefaultTag
	}
	store.Lock()
	defer store.Unlock()
	repoName = registry.NormalizeLocalName(repoName)
	repo, ok := store.Repositories[repoName]
	if !ok {
		return "", ErrNameIsNotExist
	}
	id, ok := repo[ref]
	if !ok {
		return "", ErrNameIsNotExist
	}
	return id, nil
}

// ByID returns a reverse-lookup table of all the names which refer to each
// image - e.g. {"43b5f19b10584": {"base:latest", "base:v1"}}
func (store *TagStore) ByID() map[string][]string {
	store.Lock()
	defer store.Unlock()
	byID := make(map[string][]string)
	for repoName, repository := range store.Repositories {
		for tag, id := range repository {
			name := utils.ImageReference(repoName, tag)
			if _, exists := byID[id]; !exists {
				byID[id] = []string{name}
			} else {
				byID[id] = append(byID[id], name)
				sort.Strings(byID[id])
			}
		}
	}
	return byID
}

// HasReferences returns whether or not the given image is referenced in one or
// more repositories.
func (store *TagStore) HasReferences(img *image.Image) bool {
	return len(store.ByID()[img.ID]) > 0
}

// ImageName returns name of an image, given the image's ID.
func (store *TagStore) ImageName(id string) string {
	if names, exists := store.ByID()[id]; exists && len(names) > 0 {
		return names[0]
	}
	return stringid.TruncateID(id)
}

// DeleteAll removes images identified by a specific ID from the store.
func (store *TagStore) DeleteAll(id string) error {
	names, exists := store.ByID()[id]
	if !exists || len(names) == 0 {
		return nil
	}
	for _, name := range names {
		if strings.Contains(name, ":") {
			nameParts := strings.Split(name, ":")
			if _, err := store.Delete(nameParts[0], nameParts[1]); err != nil {
				return err
			}
		} else {
			if _, err := store.Delete(name, ""); err != nil {
				return err
			}
		}
	}
	return nil
}

// Delete deletes a repository or a specific tag. If ref is empty, the entire
// repository named repoName will be deleted; otherwise only the tag named by
// ref will be deleted.
func (store *TagStore) Delete(repoName, ref string) (bool, error) {
	store.Lock()
	defer store.Unlock()
	deleted := false
	if err := store.reload(); err != nil {
		return false, err
	}

	repoName = registry.NormalizeLocalName(repoName)

	if ref == "" {
		// Delete the whole repository.
		delete(store.Repositories, repoName)
		return true, store.save()
	}

	repoRefs, exists := store.Repositories[repoName]
	if !exists {
		return false, fmt.Errorf("No such repository: %s", repoName)
	}

	if _, exists := repoRefs[ref]; exists {
		delete(repoRefs, ref)
		if len(repoRefs) == 0 {
			delete(store.Repositories, repoName)
		}
		deleted = true
	}

	return deleted, store.save()
}

// Tag creates a tag in the repository reponame, pointing to the image named
// imageName. If force is true, an existing tag with the same name may be
// overwritten.
func (store *TagStore) Tag(repoName, tag, imageName string, force bool) error {
	return store.setLoad(repoName, tag, imageName, force, nil)
}

// setLoad stores the image to the store.
// If the imageName is already in the repo then a '-f' flag should be used to replace existing image.
func (store *TagStore) setLoad(repoName, tag, imageName string, force bool, out io.Writer) error {
	img, err := store.LookupImage(imageName)
	store.Lock()
	defer store.Unlock()
	if err != nil {
		return err
	}
	if tag == "" {
		tag = tags.DefaultTag
	}
	if err := validateRepoName(repoName); err != nil {
		return err
	}
	if err := tags.ValidateTagName(tag); err != nil {
		return err
	}
	if err := store.reload(); err != nil {
		return err
	}
	var repo Repository
	repoName = registry.NormalizeLocalName(repoName)
	if r, exists := store.Repositories[repoName]; exists {
		repo = r
		if old, exists := store.Repositories[repoName][tag]; exists {

			if !force {
				return fmt.Errorf("Conflict: Tag %s:%s is already set to image %s, if you want to replace it, please use -f option", repoName, tag, old[:12])
			}

			if old != img.ID && out != nil {

				fmt.Fprintf(out, "The image %s:%s already exists, renaming the old one with ID %s to empty string\n", repoName, tag, old[:12])

			}
		}
	} else {
		repo = make(map[string]string)
		store.Repositories[repoName] = repo
	}
	repo[tag] = img.ID
	return store.save()
}

// SetDigest creates a digest reference to an image ID.
func (store *TagStore) SetDigest(repoName, digest, imageName string) error {
	img, err := store.LookupImage(imageName)
	if err != nil {
		return err
	}

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

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

	store.Lock()
	defer store.Unlock()
	if err := store.reload(); err != nil {
		return err
	}

	repoName = registry.NormalizeLocalName(repoName)
	repoRefs, exists := store.Repositories[repoName]
	if !exists {
		repoRefs = Repository{}
		store.Repositories[repoName] = repoRefs
	} else if oldID, exists := repoRefs[digest]; exists && oldID != img.ID {
		return fmt.Errorf("Conflict: Digest %s is already set to image %s", digest, oldID)
	}

	repoRefs[digest] = img.ID
	return store.save()
}

// Get returns the Repository tag/image map for a given repository.
func (store *TagStore) Get(repoName string) (Repository, error) {
	store.Lock()
	defer store.Unlock()
	if err := store.reload(); err != nil {
		return nil, err
	}
	repoName = registry.NormalizeLocalName(repoName)
	if r, exists := store.Repositories[repoName]; exists {
		return r, nil
	}
	return nil, nil
}

// GetImage returns a pointer to an Image structure describing the image
// referred to by refOrID inside repository repoName.
func (store *TagStore) GetImage(repoName, refOrID string) (*image.Image, error) {
	repo, err := store.Get(repoName)

	if err != nil {
		return nil, err
	}
	if repo == nil {
		return nil, nil
	}

	store.Lock()
	defer store.Unlock()
	if imgID, exists := repo[refOrID]; exists {
		return store.graph.Get(imgID)
	}

	// If no matching tag is found, search through images for a matching image id
	// iff it looks like a short ID or would look like a short ID
	if stringid.IsShortID(stringid.TruncateID(refOrID)) {
		for _, revision := range repo {
			if strings.HasPrefix(revision, refOrID) {
				return store.graph.Get(revision)
			}
		}
	}

	return nil, nil
}

// GetRepoRefs returns a map with image IDs as keys, and slices listing
// repo/tag references as the values. It covers all repositories.
func (store *TagStore) GetRepoRefs() map[string][]string {
	store.Lock()
	reporefs := make(map[string][]string)

	for name, repository := range store.Repositories {
		for tag, id := range repository {
			shortID := stringid.TruncateID(id)
			reporefs[shortID] = append(reporefs[shortID], utils.ImageReference(name, tag))
		}
	}
	store.Unlock()
	return reporefs
}

// validateRepoName validates the name of a repository.
func validateRepoName(name string) error {
	if name == "" {
		return fmt.Errorf("Repository name can't be empty")
	}
	if name == "scratch" {
		return fmt.Errorf("'scratch' is a reserved name")
	}
	return nil
}

func validateDigest(dgst string) error {
	if dgst == "" {
		return errors.New("digest can't be empty")
	}
	if _, err := digest.ParseDigest(dgst); err != nil {
		return err
	}
	return nil
}

// poolAdd checks if a push or pull is already running, and returns
// (broadcaster, true) if a running operation is found. Otherwise, it creates a
// new one and returns (broadcaster, false).
func (store *TagStore) poolAdd(kind, key string) (*broadcaster.Buffered, bool) {
	store.Lock()
	defer store.Unlock()

	if p, exists := store.pullingPool[key]; exists {
		return p, true
	}
	if p, exists := store.pushingPool[key]; exists {
		return p, true
	}

	broadcaster := broadcaster.NewBuffered()

	switch kind {
	case "pull":
		store.pullingPool[key] = broadcaster
	case "push":
		store.pushingPool[key] = broadcaster
	default:
		panic("Unknown pool type")
	}

	return broadcaster, false
}

func (store *TagStore) poolRemoveWithError(kind, key string, broadcasterResult error) error {
	store.Lock()
	defer store.Unlock()
	switch kind {
	case "pull":
		if broadcaster, exists := store.pullingPool[key]; exists {
			broadcaster.CloseWithError(broadcasterResult)
			delete(store.pullingPool, key)
		}
	case "push":
		if broadcaster, exists := store.pushingPool[key]; exists {
			broadcaster.CloseWithError(broadcasterResult)
			delete(store.pushingPool, key)
		}
	default:
		return fmt.Errorf("Unknown pool type")
	}
	return nil
}

func (store *TagStore) poolRemove(kind, key string) error {
	return store.poolRemoveWithError(kind, key, nil)
}
<commit_msg>Remove dead graph code<commit_after>package graph

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"github.com/docker/distribution/digest"
	"github.com/docker/docker/daemon/events"
	"github.com/docker/docker/graph/tags"
	"github.com/docker/docker/image"
	"github.com/docker/docker/pkg/broadcaster"
	"github.com/docker/docker/pkg/parsers"
	"github.com/docker/docker/pkg/stringid"
	"github.com/docker/docker/registry"
	"github.com/docker/docker/utils"
	"github.com/docker/libtrust"
)

// ErrNameIsNotExist returned when there is no image with requested name.
var ErrNameIsNotExist = errors.New("image with specified name does not exist")

// TagStore manages repositories. It encompasses the Graph used for versioned
// storage, as well as various services involved in pushing and pulling
// repositories.
type TagStore struct {
	path  string
	graph *Graph
	// Repositories is a map of repositories, indexed by name.
	Repositories map[string]Repository
	trustKey     libtrust.PrivateKey
	sync.Mutex
	// FIXME: move push/pull-related fields
	// to a helper type
	pullingPool     map[string]*broadcaster.Buffered
	pushingPool     map[string]*broadcaster.Buffered
	registryService *registry.Service
	eventsService   *events.Events
}

// Repository maps tags to image IDs.
type Repository map[string]string

// TagStoreConfig provides parameters for a new TagStore.
type TagStoreConfig struct {
	// Graph is the versioned image store
	Graph *Graph
	// Key is the private key to use for signing manifests.
	Key libtrust.PrivateKey
	// Registry is the registry service to use for TLS configuration and
	// endpoint lookup.
	Registry *registry.Service
	// Events is the events service to use for logging.
	Events *events.Events
}

// NewTagStore creates a new TagStore at specified path, using the parameters
// and services provided in cfg.
func NewTagStore(path string, cfg *TagStoreConfig) (*TagStore, error) {
	abspath, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}

	store := &TagStore{
		path:            abspath,
		graph:           cfg.Graph,
		trustKey:        cfg.Key,
		Repositories:    make(map[string]Repository),
		pullingPool:     make(map[string]*broadcaster.Buffered),
		pushingPool:     make(map[string]*broadcaster.Buffered),
		registryService: cfg.Registry,
		eventsService:   cfg.Events,
	}
	// Load the json file if it exists, otherwise create it.
	if err := store.reload(); os.IsNotExist(err) {
		if err := store.save(); err != nil {
			return nil, err
		}
	} else if err != nil {
		return nil, err
	}
	return store, nil
}

func (store *TagStore) save() error {
	// Store the json ball
	jsonData, err := json.Marshal(store)
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(store.path, jsonData, 0600); err != nil {
		return err
	}
	return nil
}

func (store *TagStore) reload() error {
	f, err := os.Open(store.path)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := json.NewDecoder(f).Decode(&store); err != nil {
		return err
	}
	return nil
}

// LookupImage returns pointer to an Image struct corresponding to the given
// name. The name can include an optional tag; otherwise the default tag will
// be used.
func (store *TagStore) LookupImage(name string) (*image.Image, error) {
	// FIXME: standardize on returning nil when the image doesn't exist, and err for everything else
	// (so we can pass all errors here)
	repoName, ref := parsers.ParseRepositoryTag(name)
	if ref == "" {
		ref = tags.DefaultTag
	}
	var (
		err error
		img *image.Image
	)

	img, err = store.GetImage(repoName, ref)
	if err != nil {
		return nil, err
	}

	if img != nil {
		return img, nil
	}

	// name must be an image ID.
	store.Lock()
	defer store.Unlock()
	if img, err = store.graph.Get(name); err != nil {
		return nil, err
	}

	return img, nil
}

// GetID returns ID for image name.
func (store *TagStore) GetID(name string) (string, error) {
	repoName, ref := parsers.ParseRepositoryTag(name)
	if ref == "" {
		ref = tags.DefaultTag
	}
	store.Lock()
	defer store.Unlock()
	repoName = registry.NormalizeLocalName(repoName)
	repo, ok := store.Repositories[repoName]
	if !ok {
		return "", ErrNameIsNotExist
	}
	id, ok := repo[ref]
	if !ok {
		return "", ErrNameIsNotExist
	}
	return id, nil
}

// ByID returns a reverse-lookup table of all the names which refer to each
// image - e.g. {"43b5f19b10584": {"base:latest", "base:v1"}}
func (store *TagStore) ByID() map[string][]string {
	store.Lock()
	defer store.Unlock()
	byID := make(map[string][]string)
	for repoName, repository := range store.Repositories {
		for tag, id := range repository {
			name := utils.ImageReference(repoName, tag)
			if _, exists := byID[id]; !exists {
				byID[id] = []string{name}
			} else {
				byID[id] = append(byID[id], name)
				sort.Strings(byID[id])
			}
		}
	}
	return byID
}

// HasReferences returns whether or not the given image is referenced in one or
// more repositories.
func (store *TagStore) HasReferences(img *image.Image) bool {
	return len(store.ByID()[img.ID]) > 0
}

// Delete deletes a repository or a specific tag. If ref is empty, the entire
// repository named repoName will be deleted; otherwise only the tag named by
// ref will be deleted.
func (store *TagStore) Delete(repoName, ref string) (bool, error) {
	store.Lock()
	defer store.Unlock()
	deleted := false
	if err := store.reload(); err != nil {
		return false, err
	}

	repoName = registry.NormalizeLocalName(repoName)

	if ref == "" {
		// Delete the whole repository.
		delete(store.Repositories, repoName)
		return true, store.save()
	}

	repoRefs, exists := store.Repositories[repoName]
	if !exists {
		return false, fmt.Errorf("No such repository: %s", repoName)
	}

	if _, exists := repoRefs[ref]; exists {
		delete(repoRefs, ref)
		if len(repoRefs) == 0 {
			delete(store.Repositories, repoName)
		}
		deleted = true
	}

	return deleted, store.save()
}

// Tag creates a tag in the repository reponame, pointing to the image named
// imageName. If force is true, an existing tag with the same name may be
// overwritten.
func (store *TagStore) Tag(repoName, tag, imageName string, force bool) error {
	return store.setLoad(repoName, tag, imageName, force, nil)
}

// setLoad stores the image to the store.
// If the imageName is already in the repo then a '-f' flag should be used to replace existing image.
func (store *TagStore) setLoad(repoName, tag, imageName string, force bool, out io.Writer) error {
	img, err := store.LookupImage(imageName)
	store.Lock()
	defer store.Unlock()
	if err != nil {
		return err
	}
	if tag == "" {
		tag = tags.DefaultTag
	}
	if err := validateRepoName(repoName); err != nil {
		return err
	}
	if err := tags.ValidateTagName(tag); err != nil {
		return err
	}
	if err := store.reload(); err != nil {
		return err
	}
	var repo Repository
	repoName = registry.NormalizeLocalName(repoName)
	if r, exists := store.Repositories[repoName]; exists {
		repo = r
		if old, exists := store.Repositories[repoName][tag]; exists {

			if !force {
				return fmt.Errorf("Conflict: Tag %s:%s is already set to image %s, if you want to replace it, please use -f option", repoName, tag, old[:12])
			}

			if old != img.ID && out != nil {

				fmt.Fprintf(out, "The image %s:%s already exists, renaming the old one with ID %s to empty string\n", repoName, tag, old[:12])

			}
		}
	} else {
		repo = make(map[string]string)
		store.Repositories[repoName] = repo
	}
	repo[tag] = img.ID
	return store.save()
}

// SetDigest creates a digest reference to an image ID.
func (store *TagStore) SetDigest(repoName, digest, imageName string) error {
	img, err := store.LookupImage(imageName)
	if err != nil {
		return err
	}

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

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

	store.Lock()
	defer store.Unlock()
	if err := store.reload(); err != nil {
		return err
	}

	repoName = registry.NormalizeLocalName(repoName)
	repoRefs, exists := store.Repositories[repoName]
	if !exists {
		repoRefs = Repository{}
		store.Repositories[repoName] = repoRefs
	} else if oldID, exists := repoRefs[digest]; exists && oldID != img.ID {
		return fmt.Errorf("Conflict: Digest %s is already set to image %s", digest, oldID)
	}

	repoRefs[digest] = img.ID
	return store.save()
}

// Get returns the Repository tag/image map for a given repository.
func (store *TagStore) Get(repoName string) (Repository, error) {
	store.Lock()
	defer store.Unlock()
	if err := store.reload(); err != nil {
		return nil, err
	}
	repoName = registry.NormalizeLocalName(repoName)
	if r, exists := store.Repositories[repoName]; exists {
		return r, nil
	}
	return nil, nil
}

// GetImage returns a pointer to an Image structure describing the image
// referred to by refOrID inside repository repoName.
func (store *TagStore) GetImage(repoName, refOrID string) (*image.Image, error) {
	repo, err := store.Get(repoName)

	if err != nil {
		return nil, err
	}
	if repo == nil {
		return nil, nil
	}

	store.Lock()
	defer store.Unlock()
	if imgID, exists := repo[refOrID]; exists {
		return store.graph.Get(imgID)
	}

	// If no matching tag is found, search through images for a matching image id
	// iff it looks like a short ID or would look like a short ID
	if stringid.IsShortID(stringid.TruncateID(refOrID)) {
		for _, revision := range repo {
			if strings.HasPrefix(revision, refOrID) {
				return store.graph.Get(revision)
			}
		}
	}

	return nil, nil
}

// validateRepoName validates the name of a repository.
func validateRepoName(name string) error {
	if name == "" {
		return fmt.Errorf("Repository name can't be empty")
	}
	if name == "scratch" {
		return fmt.Errorf("'scratch' is a reserved name")
	}
	return nil
}

func validateDigest(dgst string) error {
	if dgst == "" {
		return errors.New("digest can't be empty")
	}
	if _, err := digest.ParseDigest(dgst); err != nil {
		return err
	}
	return nil
}

// poolAdd checks if a push or pull is already running, and returns
// (broadcaster, true) if a running operation is found. Otherwise, it creates a
// new one and returns (broadcaster, false).
func (store *TagStore) poolAdd(kind, key string) (*broadcaster.Buffered, bool) {
	store.Lock()
	defer store.Unlock()

	if p, exists := store.pullingPool[key]; exists {
		return p, true
	}
	if p, exists := store.pushingPool[key]; exists {
		return p, true
	}

	broadcaster := broadcaster.NewBuffered()

	switch kind {
	case "pull":
		store.pullingPool[key] = broadcaster
	case "push":
		store.pushingPool[key] = broadcaster
	default:
		panic("Unknown pool type")
	}

	return broadcaster, false
}

func (store *TagStore) poolRemoveWithError(kind, key string, broadcasterResult error) error {
	store.Lock()
	defer store.Unlock()
	switch kind {
	case "pull":
		if broadcaster, exists := store.pullingPool[key]; exists {
			broadcaster.CloseWithError(broadcasterResult)
			delete(store.pullingPool, key)
		}
	case "push":
		if broadcaster, exists := store.pushingPool[key]; exists {
			broadcaster.CloseWithError(broadcasterResult)
			delete(store.pushingPool, key)
		}
	default:
		return fmt.Errorf("Unknown pool type")
	}
	return nil
}

func (store *TagStore) poolRemove(kind, key string) error {
	return store.poolRemoveWithError(kind, key, nil)
}
<|endoftext|>


<commit_before>package wavs

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

	"github.com/gordonklaus/portaudio"
	log "github.com/sirupsen/logrus"
	"github.com/youpy/go-wav"
)

const (
	WAVS = "./wavs"
	BUF  = 128
)

type Wavs struct {
	wavs    map[string]*wavFile
	stream  *portaudio.Stream
	stopAll chan struct{}
}

type wavFile struct {
	name      string
	buf       []float32
	length    int
	active    chan struct{}
	remaining int
}

func InitWavs() *Wavs {
	err := portaudio.Initialize()
	if err != nil {
		panic(err)
	}

	files, _ := ioutil.ReadDir(WAVS)

	wavs := &Wavs{
		wavs:    map[string]*wavFile{},
		stopAll: make(chan struct{}),
	}

	for _, f := range files {
		if strings.HasSuffix(f.Name(), ".wav") {
			wavs.wavs[f.Name()] = initWav(f)
		}
	}

	wavs.stream, err = portaudio.OpenDefaultStream(0, 1, 44100, BUF, wavs.cb)
	if err != nil {
		panic(err)
	}
	err = wavs.stream.Start()
	if err != nil {
		panic(err)
	}

	return wavs
}

func initWav(f os.FileInfo) *wavFile {
	file, _ := os.Open(WAVS + "/" + f.Name())
	reader := wav.NewReader(file)
	defer file.Close()

	buf := make([]float32, f.Size())
	loc := 0
	for {
		samples, err := reader.ReadSamples()
		if err != nil {
			if err == io.EOF {
				break
			} else {
				panic(fmt.Sprintf("ReadSamples failed on file %s with error: %s", f.Name(), err))
			}
		}

		for _, sample := range samples {
			buf[loc] = float32(reader.FloatValue(sample, 0))
			loc += 1
		}
	}

	w := wavFile{
		name:   f.Name(),
		buf:    make([]float32, loc),
		length: loc,
		active: make(chan struct{}),
	}
	copy(w.buf, buf[:])

	return &w
}

func (w *Wavs) cb(output [][]float32) {
	for _, wv := range w.wavs {
		select {
		case <-wv.active:
			wv.remaining = wv.length
		default:
		}
	}

	select {
	case <-w.stopAll:
		for _, wv := range w.wavs {
			wv.remaining = 0
		}
	default:
	}

	out := make([]float32, BUF)
	for _, wv := range w.wavs {
		for i := 0; i < BUF; i++ {
			if wv.remaining > i {
				// TODO: do this correctly
				// add dynamic range compressor
				out[i] += wv.buf[wv.length-wv.remaining+i]
			}
		}
		wv.remaining -= BUF
	}
	copy(output[0], out)
}

func (w *Wavs) Play(name string) {
	wav, ok := w.wavs[name]
	if !ok {
		log.Debugf("Unknown wav file: %s", name)
		return
	}

	wav.active <- struct{}{}
}

func (w *Wavs) StopAll() {
	w.stopAll <- struct{}{}
}

func (w *Wavs) Durations() map[string]time.Duration {
	durations := map[string]time.Duration{}

	for name, wav := range w.wavs {
		durations[name] = time.Duration(wav.length*1000/44100) * time.Millisecond
	}

	return durations
}

func (w *Wavs) Close() {
	w.stream.Stop()
	w.stream.Close()
	portaudio.Terminate()

	log.Debugf("Wavs Closed")
}
<commit_msg>fix 32-bit overflow when calculating wave file duration<commit_after>package wavs

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

	"github.com/gordonklaus/portaudio"
	log "github.com/sirupsen/logrus"
	"github.com/youpy/go-wav"
)

const (
	WAVS = "./wavs"
	BUF  = 128
)

type Wavs struct {
	wavs    map[string]*wavFile
	stream  *portaudio.Stream
	stopAll chan struct{}
}

type wavFile struct {
	name      string
	buf       []float32
	length    int
	active    chan struct{}
	remaining int
}

func InitWavs() *Wavs {
	err := portaudio.Initialize()
	if err != nil {
		panic(err)
	}

	files, _ := ioutil.ReadDir(WAVS)

	wavs := &Wavs{
		wavs:    map[string]*wavFile{},
		stopAll: make(chan struct{}),
	}

	for _, f := range files {
		if strings.HasSuffix(f.Name(), ".wav") {
			wavs.wavs[f.Name()] = initWav(f)
		}
	}

	wavs.stream, err = portaudio.OpenDefaultStream(0, 1, 44100, BUF, wavs.cb)
	if err != nil {
		panic(err)
	}
	err = wavs.stream.Start()
	if err != nil {
		panic(err)
	}

	return wavs
}

func initWav(f os.FileInfo) *wavFile {
	file, _ := os.Open(WAVS + "/" + f.Name())
	reader := wav.NewReader(file)
	defer file.Close()

	buf := make([]float32, f.Size())
	loc := 0
	for {
		samples, err := reader.ReadSamples()
		if err != nil {
			if err == io.EOF {
				break
			} else {
				panic(fmt.Sprintf("ReadSamples failed on file %s with error: %s", f.Name(), err))
			}
		}

		for _, sample := range samples {
			buf[loc] = float32(reader.FloatValue(sample, 0))
			loc += 1
		}
	}

	w := wavFile{
		name:   f.Name(),
		buf:    make([]float32, loc),
		length: loc,
		active: make(chan struct{}),
	}
	copy(w.buf, buf[:])

	return &w
}

func (w *Wavs) cb(output [][]float32) {
	for _, wv := range w.wavs {
		select {
		case <-wv.active:
			wv.remaining = wv.length
		default:
		}
	}

	select {
	case <-w.stopAll:
		for _, wv := range w.wavs {
			wv.remaining = 0
		}
	default:
	}

	out := make([]float32, BUF)
	for _, wv := range w.wavs {
		for i := 0; i < BUF; i++ {
			if wv.remaining > i {
				// TODO: do this correctly
				// add dynamic range compressor
				out[i] += wv.buf[wv.length-wv.remaining+i]
			}
		}
		wv.remaining -= BUF
	}
	copy(output[0], out)
}

func (w *Wavs) Play(name string) {
	wav, ok := w.wavs[name]
	if !ok {
		log.Debugf("Unknown wav file: %s", name)
		return
	}

	wav.active <- struct{}{}
}

func (w *Wavs) StopAll() {
	w.stopAll <- struct{}{}
}

func (w *Wavs) Durations() map[string]time.Duration {
	durations := map[string]time.Duration{}

	for name, wav := range w.wavs {
		durations[name] = time.Duration(float64(wav.length)/44.1) * time.Millisecond
	}

	return durations
}

func (w *Wavs) Close() {
	w.stream.Stop()
	w.stream.Close()
	portaudio.Terminate()

	log.Debugf("Wavs Closed")
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Kubernetes Authors.

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

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

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

package apimachinery

import (
	"context"
	"strings"

	"github.com/onsi/ginkgo"
	"k8s.io/apimachinery/pkg/runtime/schema"

	v1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
	"k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apiextensions-apiserver/test/integration/fixtures"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/util/json"
	"k8s.io/apiserver/pkg/storage/names"
	"k8s.io/client-go/dynamic"
	"k8s.io/kubernetes/test/e2e/framework"
)

var _ = SIGDescribe("CustomResourceValidationRules [Privileged:ClusterAdmin][Alpha][Feature:CustomResourceValidationExpressions]", func() {
	f := framework.NewDefaultFramework("crd-validation-expressions")

	var apiExtensionClient *clientset.Clientset
	ginkgo.BeforeEach(func() {
		var err error
		apiExtensionClient, err = clientset.NewForConfig(f.ClientConfig())
		framework.ExpectNoError(err, "initializing apiExtensionClient")
	})

	customResourceClient := func(crd *v1.CustomResourceDefinition) (dynamic.NamespaceableResourceInterface, schema.GroupVersionResource) {
		gvrs := fixtures.GetGroupVersionResourcesOfCustomResource(crd)
		if len(gvrs) != 1 {
			ginkgo.Fail("Expected one version in custom resource definition")
		}
		gvr := gvrs[0]
		return f.DynamicClient.Resource(gvr), gvr
	}
	unmarshallSchema := func(schemaJson []byte) *v1.JSONSchemaProps {
		var c v1.JSONSchemaProps
		err := json.Unmarshal(schemaJson, &c)
		framework.ExpectNoError(err, "unmarshalling OpenAPIv3 schema")
		return &c
	}

	var schemaWithValidationExpression = unmarshallSchema([]byte(`{
	   "type":"object",
	   "properties":{
		  "spec":{
			 "type":"object",
			 "x-kubernetes-validations":[
		       { "rule":"self.x + self.y > 0" }
	         ],
			 "properties":{
				"x":{ "type":"integer" },
				"y":{ "type":"integer" }
			 }
		  },
		  "status":{
			 "type":"object",
			 "x-kubernetes-validations":[
				{ "rule":"self.health == 'ok' || self.health == 'unhealthy'" }
			 ],
			 "properties":{
				"health":{ "type":"string" }
			 }
		  }
	   }
	}`))
	ginkgo.It("MUST NOT fail validation for create of a custom resource that satisfies the x-kubernetes-validator rules", func() {
		ginkgo.By("Creating a custom resource definition with validation rules")
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithValidationExpression, false)
		crd, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectNoError(err, "creating CustomResourceDefinition")
		defer func() {
			err = fixtures.DeleteV1CustomResourceDefinition(crd, apiExtensionClient)
			framework.ExpectNoError(err, "deleting CustomResourceDefinition")
		}()

		ginkgo.By("Creating a custom resource with values that are allowed by the validation rules set on the custom resource definition")
		crClient, gvr := customResourceClient(crd)
		name1 := names.SimpleNameGenerator.GenerateName("cr-1")
		_, err = crClient.Namespace(f.Namespace.Name).Create(context.TODO(), &unstructured.Unstructured{Object: map[string]interface{}{
			"apiVersion": gvr.Group + "/" + gvr.Version,
			"kind":       crd.Spec.Names.Kind,
			"metadata": map[string]interface{}{
				"name":      name1,
				"namespace": f.Namespace.Name,
			},
			"spec": map[string]interface{}{
				"x": int64(1),
				"y": int64(0),
			},
		}}, metav1.CreateOptions{})
		framework.ExpectNoError(err, "validation rules satisfied")
	})
	ginkgo.It("MUST fail validation for create of a custom resource that does not satisfy the x-kubernetes-validator rules", func() {
		ginkgo.By("Creating a custom resource definition with validation rules")
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithValidationExpression, false)
		crd, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectNoError(err, "creating CustomResourceDefinition")
		defer func() {
			err = fixtures.DeleteV1CustomResourceDefinition(crd, apiExtensionClient)
			framework.ExpectNoError(err, "deleting CustomResourceDefinition")
		}()

		ginkgo.By("Creating a custom resource with values that fail the validation rules set on the custom resource definition")
		crClient, gvr := customResourceClient(crd)
		name1 := names.SimpleNameGenerator.GenerateName("cr-1")
		_, err = crClient.Namespace(f.Namespace.Name).Create(context.TODO(), &unstructured.Unstructured{Object: map[string]interface{}{
			"apiVersion": gvr.Group + "/" + gvr.Version,
			"kind":       crd.Spec.Names.Kind,
			"metadata": map[string]interface{}{
				"name":      name1,
				"namespace": f.Namespace.Name,
			},
			"spec": map[string]interface{}{
				"x": int64(0),
				"y": int64(0),
			},
		}}, metav1.CreateOptions{})
		framework.ExpectError(err, "validation rules not satisfied")
		expectedErrMsg := "failed rule"
		if !strings.Contains(err.Error(), expectedErrMsg) {
			framework.Failf("expect error contains %q, got %q", expectedErrMsg, err.Error())
		}
	})

	ginkgo.It("MUST fail create of a custom resource definition that contains a x-kubernetes-validator rule that refers to a property that do not exist", func() {
		ginkgo.By("Defining a custom resource definition with a validation rule that refers to a property that do not exist")
		var schemaWithInvalidValidationRule = unmarshallSchema([]byte(`{
		   "type":"object",
		   "properties":{
			  "spec":{
				 "type":"object",
				 "x-kubernetes-validations":[
				   { "rule":"self.z == 100" }
				 ],
				 "properties":{
					"x":{ "type":"integer" }
				 }
			  }
		   }
		}`))
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithInvalidValidationRule, false)
		_, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectError(err, "creating CustomResourceDefinition with a validation rule that refers to a property that do not exist")
		expectedErrMsg := "undefined field 'z'"
		if !strings.Contains(err.Error(), expectedErrMsg) {
			framework.Failf("expect error contains %q, got %q", expectedErrMsg, err.Error())
		}
	})

	ginkgo.It("MUST fail create of a custom resource definition that contains an x-kubernetes-validations rule that contains a syntax error", func() {
		ginkgo.By("Defining a custom resource definition that contains a validation rule with a syntax error")
		var schemaWithSyntaxErrorRule = unmarshallSchema([]byte(`{
			"type":"object",
			"properties":{
			   "spec":{
				  "type":"object",
				  "x-kubernetes-validations":[
					{ "rule":"self = 42" }
				  ]
			   }
			}
		}`))
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithSyntaxErrorRule, false)
		_, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectError(err, "creating a CustomResourceDefinition with a validation rule that contains a syntax error")
		expectedErrMsg := "Syntax error"
		if !strings.Contains(err.Error(), expectedErrMsg) {
			framework.Failf("expected error message to contain %q, got %q", expectedErrMsg, err.Error())
		}
	})

	ginkgo.It("MUST fail create of a custom resource definition that contains an x-kubernetes-validations rule that exceeds the estimated cost limit", func() {
		ginkgo.By("Defining a custom resource definition that contains a validation rule that exceeds the cost limit")
		var schemaWithExpensiveRule = unmarshallSchema([]byte(`{
			"type": "object",
			"properties": {
				"spec": {
					"type": "object",
					"properties": {
						"x": {
							"type": "array",
							"items": {
								"type": "array",
								"items": {
									"type": "string"
								},
								"x-kubernetes-validations": [{
									"rule": "self.all(s, s == \"string constant\")"
								}]
							}
						}
					}
				}
			}
		}`))
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithExpensiveRule, false)
		_, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectError(err, "creating a CustomResourceDefinition with a validation rule that exceeds the cost limit")
		expectedErrMsg := "exceeded budget"
		if !strings.Contains(err.Error(), expectedErrMsg) {
			framework.Failf("expected error message to contain %q, got %q", expectedErrMsg, err.Error())
		}
	})
})
<commit_msg>Fix incorrect wording in expectedErrMsg.<commit_after>/*
Copyright 2021 The Kubernetes Authors.

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

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

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

package apimachinery

import (
	"context"
	"strings"

	"github.com/onsi/ginkgo"
	"k8s.io/apimachinery/pkg/runtime/schema"

	v1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
	"k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apiextensions-apiserver/test/integration/fixtures"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/util/json"
	"k8s.io/apiserver/pkg/storage/names"
	"k8s.io/client-go/dynamic"
	"k8s.io/kubernetes/test/e2e/framework"
)

var _ = SIGDescribe("CustomResourceValidationRules [Privileged:ClusterAdmin][Alpha][Feature:CustomResourceValidationExpressions]", func() {
	f := framework.NewDefaultFramework("crd-validation-expressions")

	var apiExtensionClient *clientset.Clientset
	ginkgo.BeforeEach(func() {
		var err error
		apiExtensionClient, err = clientset.NewForConfig(f.ClientConfig())
		framework.ExpectNoError(err, "initializing apiExtensionClient")
	})

	customResourceClient := func(crd *v1.CustomResourceDefinition) (dynamic.NamespaceableResourceInterface, schema.GroupVersionResource) {
		gvrs := fixtures.GetGroupVersionResourcesOfCustomResource(crd)
		if len(gvrs) != 1 {
			ginkgo.Fail("Expected one version in custom resource definition")
		}
		gvr := gvrs[0]
		return f.DynamicClient.Resource(gvr), gvr
	}
	unmarshallSchema := func(schemaJson []byte) *v1.JSONSchemaProps {
		var c v1.JSONSchemaProps
		err := json.Unmarshal(schemaJson, &c)
		framework.ExpectNoError(err, "unmarshalling OpenAPIv3 schema")
		return &c
	}

	var schemaWithValidationExpression = unmarshallSchema([]byte(`{
	   "type":"object",
	   "properties":{
		  "spec":{
			 "type":"object",
			 "x-kubernetes-validations":[
		       { "rule":"self.x + self.y > 0" }
	         ],
			 "properties":{
				"x":{ "type":"integer" },
				"y":{ "type":"integer" }
			 }
		  },
		  "status":{
			 "type":"object",
			 "x-kubernetes-validations":[
				{ "rule":"self.health == 'ok' || self.health == 'unhealthy'" }
			 ],
			 "properties":{
				"health":{ "type":"string" }
			 }
		  }
	   }
	}`))
	ginkgo.It("MUST NOT fail validation for create of a custom resource that satisfies the x-kubernetes-validator rules", func() {
		ginkgo.By("Creating a custom resource definition with validation rules")
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithValidationExpression, false)
		crd, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectNoError(err, "creating CustomResourceDefinition")
		defer func() {
			err = fixtures.DeleteV1CustomResourceDefinition(crd, apiExtensionClient)
			framework.ExpectNoError(err, "deleting CustomResourceDefinition")
		}()

		ginkgo.By("Creating a custom resource with values that are allowed by the validation rules set on the custom resource definition")
		crClient, gvr := customResourceClient(crd)
		name1 := names.SimpleNameGenerator.GenerateName("cr-1")
		_, err = crClient.Namespace(f.Namespace.Name).Create(context.TODO(), &unstructured.Unstructured{Object: map[string]interface{}{
			"apiVersion": gvr.Group + "/" + gvr.Version,
			"kind":       crd.Spec.Names.Kind,
			"metadata": map[string]interface{}{
				"name":      name1,
				"namespace": f.Namespace.Name,
			},
			"spec": map[string]interface{}{
				"x": int64(1),
				"y": int64(0),
			},
		}}, metav1.CreateOptions{})
		framework.ExpectNoError(err, "validation rules satisfied")
	})
	ginkgo.It("MUST fail validation for create of a custom resource that does not satisfy the x-kubernetes-validator rules", func() {
		ginkgo.By("Creating a custom resource definition with validation rules")
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithValidationExpression, false)
		crd, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectNoError(err, "creating CustomResourceDefinition")
		defer func() {
			err = fixtures.DeleteV1CustomResourceDefinition(crd, apiExtensionClient)
			framework.ExpectNoError(err, "deleting CustomResourceDefinition")
		}()

		ginkgo.By("Creating a custom resource with values that fail the validation rules set on the custom resource definition")
		crClient, gvr := customResourceClient(crd)
		name1 := names.SimpleNameGenerator.GenerateName("cr-1")
		_, err = crClient.Namespace(f.Namespace.Name).Create(context.TODO(), &unstructured.Unstructured{Object: map[string]interface{}{
			"apiVersion": gvr.Group + "/" + gvr.Version,
			"kind":       crd.Spec.Names.Kind,
			"metadata": map[string]interface{}{
				"name":      name1,
				"namespace": f.Namespace.Name,
			},
			"spec": map[string]interface{}{
				"x": int64(0),
				"y": int64(0),
			},
		}}, metav1.CreateOptions{})
		framework.ExpectError(err, "validation rules not satisfied")
		expectedErrMsg := "failed rule"
		if !strings.Contains(err.Error(), expectedErrMsg) {
			framework.Failf("expect error contains %q, got %q", expectedErrMsg, err.Error())
		}
	})

	ginkgo.It("MUST fail create of a custom resource definition that contains a x-kubernetes-validator rule that refers to a property that do not exist", func() {
		ginkgo.By("Defining a custom resource definition with a validation rule that refers to a property that do not exist")
		var schemaWithInvalidValidationRule = unmarshallSchema([]byte(`{
		   "type":"object",
		   "properties":{
			  "spec":{
				 "type":"object",
				 "x-kubernetes-validations":[
				   { "rule":"self.z == 100" }
				 ],
				 "properties":{
					"x":{ "type":"integer" }
				 }
			  }
		   }
		}`))
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithInvalidValidationRule, false)
		_, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectError(err, "creating CustomResourceDefinition with a validation rule that refers to a property that do not exist")
		expectedErrMsg := "undefined field 'z'"
		if !strings.Contains(err.Error(), expectedErrMsg) {
			framework.Failf("expect error contains %q, got %q", expectedErrMsg, err.Error())
		}
	})

	ginkgo.It("MUST fail create of a custom resource definition that contains an x-kubernetes-validations rule that contains a syntax error", func() {
		ginkgo.By("Defining a custom resource definition that contains a validation rule with a syntax error")
		var schemaWithSyntaxErrorRule = unmarshallSchema([]byte(`{
			"type":"object",
			"properties":{
			   "spec":{
				  "type":"object",
				  "x-kubernetes-validations":[
					{ "rule":"self = 42" }
				  ]
			   }
			}
		}`))
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithSyntaxErrorRule, false)
		_, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectError(err, "creating a CustomResourceDefinition with a validation rule that contains a syntax error")
		expectedErrMsg := "Syntax error"
		if !strings.Contains(err.Error(), expectedErrMsg) {
			framework.Failf("expected error message to contain %q, got %q", expectedErrMsg, err.Error())
		}
	})

	ginkgo.It("MUST fail create of a custom resource definition that contains an x-kubernetes-validations rule that exceeds the estimated cost limit", func() {
		ginkgo.By("Defining a custom resource definition that contains a validation rule that exceeds the cost limit")
		var schemaWithExpensiveRule = unmarshallSchema([]byte(`{
			"type": "object",
			"properties": {
				"spec": {
					"type": "object",
					"properties": {
						"x": {
							"type": "array",
							"items": {
								"type": "array",
								"items": {
									"type": "string"
								},
								"x-kubernetes-validations": [{
									"rule": "self.all(s, s == \"string constant\")"
								}]
							}
						}
					}
				}
			}
		}`))
		crd := fixtures.NewRandomNameV1CustomResourceDefinitionWithSchema(v1.NamespaceScoped, schemaWithExpensiveRule, false)
		_, err := fixtures.CreateNewV1CustomResourceDefinitionWatchUnsafe(crd, apiExtensionClient)
		framework.ExpectError(err, "creating a CustomResourceDefinition with a validation rule that exceeds the cost limit")
		expectedErrMsg := "exceeds budget"
		if !strings.Contains(err.Error(), expectedErrMsg) {
			framework.Failf("expected error message to contain %q, got %q", expectedErrMsg, err.Error())
		}
	})
})
<|endoftext|>


<commit_before>package main

import (
	"compress/zlib"
	"encoding/gob"
	. "github.com/ogier/libtcod-go/tcod"
	"io/ioutil"
	"os"
	"path"
	"strings"
)

type Transition uint8

type MoveTransition uint8

// Transitions
const (
	INSTANT = iota
	FADE_IN
	MOVE_IN
	RANDOMIZE
	NB_TRANSITIONS
)

// Move transitions
const (
	MOVE_UP = iota
	MOVE_DOWN
	MOVE_LEFT
	MOVE_RIGHT
	NB_MOVE_TRANSITIONS
)

const TRANSITION_TIME = 700 // time in milliseconds

type IAsciiArt interface {
	Render(console IConsole, transition Transition, first bool)
}

type AsciiArtGallery struct {
	secondary      IConsole
	firstTime      int // time in milliseconds when first render occurred
	moveTransition MoveTransition
	arts           []IAsciiArt
	artfiles       []string
	last           int
	loadedArtsChan chan bool
	loadedArts     bool
}

// TODO switch to go version  from May 2010
// TODO decode ascii arts ,
// TODO store as binary file
// CODE LenName Name  WIDTH, HEIGHT  OFFSETX OFFSETY  and stream of colors

type AsciiArt struct {
	name             string
	chars            []string
	colors           [][]Color
	width, height    int
	offsetX, offsetY int
}

var asciiArts []IAsciiArt

var aas *AsciiArtGallery

func NewAsciiArtGallery() *AsciiArtGallery {
	result := &AsciiArtGallery{}
	result.loadedArts = false
	result.loadedArtsChan = make(chan bool)
	result.goLoadAsciiArts()
	return result
}

func (self *AsciiArtGallery) isArtsLoaded() bool {
	// if arts not loaded then wait for loading
	if !self.loadedArts {
		if _, ok := <-self.loadedArtsChan; ok {
			self.loadedArts = true
		}
	}
	return self.loadedArts
}

// loads ascii arts from data directory
func (self *AsciiArtGallery) goLoadAsciiArts() {
	self.loadedArts = false
	go func() {
		files := []string{}
		entries, err := ioutil.ReadDir("data/ascii")
		if err != nil {
			panic("Cannot read files from data/ascii")
		}
		for _, f := range entries {
			if strings.HasSuffix(f.Name(), ".dat") {
				files = append(files, path.Join("data/ascii", f.Name()))
			}
		}

		self.arts = make([]IAsciiArt, len(files))
		self.artfiles = files
		self.loadedArtsChan <- true
	}()
}

func (self *AsciiArtGallery) getAsciiArt(id int) IAsciiArt {
	var result IAsciiArt
	if self.arts[id] == nil {
		art, err := NewAsciiArtFromFile(self.artfiles[id])
		if err != nil {
			panic("Cannot load ascii file " + self.artfiles[id])
		}
		self.arts[id] = art
	}
	result = self.arts[id]
	return result
}

// select random ASCII art but not the one before
func (self *AsciiArtGallery) randomAsciiArt() IAsciiArt {
	if self.isArtsLoaded() {
		var id int

		// find id to return
		if len(self.arts) == 0 {
			panic("No arts found!")
		} else if len(self.arts) == 1 {
			id = 0
		} else {
			for {
				id = random.GetInt(0, len(self.arts)-1)
				if self.last != id {
					break
				}
			}
			self.last = id
		}
		return self.getAsciiArt(id)

	}
	return nil
}

func NewAsciiArtFromFile(fname string) (art IAsciiArt, err error) {

	fin, err := os.Open(fname)
	if err != nil {
		return nil, err
	}
	i, err := zlib.NewReader(fin)
	defer i.Close()
	if err != nil {
		return nil, err
	}
	d := gob.NewDecoder(i)
	a := AsciiArt{}
	d.Decode(&a)
	return &a, nil

}

func (self *AsciiArt) PutChar(console IConsole, x, y int) {
	color := self.colors[y+self.offsetY][x+self.offsetX]
	// char := int(self.chars[y+self.offsetY][x+self.offsetX])
	console.PutChar(x, y, 'x', BKGND_SET)
	console.SetCharForeground(x, y, color)
	console.SetCharBackground(x, y, COLOR_BLACK, BKGND_SET)

}

func (self *AsciiArt) Draw(console IConsole) {
	console.SetDefaultBackground(COLOR_BLACK)
	console.Clear()
	for x := 0; x < min(console.GetWidth(), self.width); x++ {
		for y := 0; y < min(console.GetHeight(), self.height); y++ {
			self.PutChar(console, x, y)
		}
	}
}

func (self *AsciiArt) Render(console IConsole, transition Transition, first bool) {
	if first {
		aas.firstTime = int(SysElapsedMilliseconds())
	}
	elapsed := int(SysElapsedMilliseconds()) - aas.firstTime

	switch transition {

	case INSTANT:
		{
			self.Draw(console)

		}
	case RANDOMIZE:
		{
			if first {
				console.SetDefaultBackground(COLOR_BLACK)
			}
			for i := 0; i < 500; i++ {
				x := random.GetInt(0, console.GetWidth()-1)
				y := random.GetInt(0, console.GetHeight()-1)

				// if in our picture then draw it
				if x < self.width && y < self.height {
					self.PutChar(console, x, y)
				} else {
					// else draw black background
					console.PutChar(x, y, ' ', BKGND_SET)
					console.SetCharForeground(x, y, COLOR_BLACK)
					console.SetCharBackground(x, y, COLOR_BLACK, BKGND_SET)
				}
			}

		}
	case FADE_IN:
		{
			if first {
				aas.secondary = NewConsole(console.GetWidth(), console.GetHeight())
				console.Blit(0, 0, console.GetWidth(), console.GetHeight(), aas.secondary, 0, 0, 1.0, 1.0)
			}
			console.SetDefaultBackground(COLOR_BLACK)
			console.Clear()
			// fade out old
			if elapsed <= TRANSITION_TIME {
				alpha := 1.0 - float32(elapsed)/TRANSITION_TIME
				aas.secondary.Blit(0, 0, console.GetWidth(), console.GetHeight(), console, 0, 0, alpha, alpha)
				// fade in new
			} else {
				self.Draw(aas.secondary)
				alpha := minf((float32(elapsed)-TRANSITION_TIME)/TRANSITION_TIME, 1)
				aas.secondary.Blit(0, 0, console.GetWidth(), console.GetHeight(), console, 0, 0, alpha, alpha)
			}

		}
	case MOVE_IN:
		{
			if first {
				aas.secondary = NewConsole(console.GetWidth(), console.GetHeight())
				aas.moveTransition = MoveTransition(random.GetInt(0, NB_MOVE_TRANSITIONS-1))
				self.Draw(aas.secondary)
			}
			widthPart := max(1, min(console.GetWidth(), int(float32(elapsed)/TRANSITION_TIME*float32(console.GetWidth()))))
			heightPart := max(1, min(console.GetHeight(), int(float32(elapsed)/TRANSITION_TIME*float32(console.GetHeight()))))

			switch aas.moveTransition {
			case MOVE_LEFT:
				aas.secondary.Blit(
					0, 0, widthPart, console.GetHeight(),
					console, 0, 0, 1.0, 1.0)
			case MOVE_RIGHT:
				aas.secondary.Blit(
					console.GetWidth()-widthPart, 0, widthPart, console.GetHeight(),
					console, console.GetWidth()-widthPart, 0, 1.0, 1.0)
			case MOVE_DOWN:
				aas.secondary.Blit(0, 0, console.GetWidth(), heightPart,
					console, 0, 0, 1.0, 1.0)
			case MOVE_UP:
				aas.secondary.Blit(
					0, console.GetHeight()-heightPart, console.GetWidth(), heightPart,
					console, 0, console.GetHeight()-heightPart, 1.0, 1.0)

			}

		}
	default:
		{
			// do nothing
			self.Render(console, FADE_IN, first)
			// TODO
			// rain/driple -> transitions -> similar to MOVE_IN transitions but for each streak sepreately
			// with different speed
			// fade in
			// have 2 consoles -> save current as screenshot
			// // and fill then new one and blit it
			// move in -> <-
			//  rain driple like matrix effect
		}
	}
}
<commit_msg>Update asciiart.go<commit_after>package main

import (
	"compress/zlib"
	"encoding/gob"
	. "github.com/afolmert/libtcod-go/tcod"
	"io/ioutil"
	"os"
	"path"
	"strings"
)

type Transition uint8

type MoveTransition uint8

// Transitions
const (
	INSTANT = iota
	FADE_IN
	MOVE_IN
	RANDOMIZE
	NB_TRANSITIONS
)

// Move transitions
const (
	MOVE_UP = iota
	MOVE_DOWN
	MOVE_LEFT
	MOVE_RIGHT
	NB_MOVE_TRANSITIONS
)

const TRANSITION_TIME = 700 // time in milliseconds

type IAsciiArt interface {
	Render(console IConsole, transition Transition, first bool)
}

type AsciiArtGallery struct {
	secondary      IConsole
	firstTime      int // time in milliseconds when first render occurred
	moveTransition MoveTransition
	arts           []IAsciiArt
	artfiles       []string
	last           int
	loadedArtsChan chan bool
	loadedArts     bool
}

// TODO switch to go version  from May 2010
// TODO decode ascii arts ,
// TODO store as binary file
// CODE LenName Name  WIDTH, HEIGHT  OFFSETX OFFSETY  and stream of colors

type AsciiArt struct {
	name             string
	chars            []string
	colors           [][]Color
	width, height    int
	offsetX, offsetY int
}

var asciiArts []IAsciiArt

var aas *AsciiArtGallery

func NewAsciiArtGallery() *AsciiArtGallery {
	result := &AsciiArtGallery{}
	result.loadedArts = false
	result.loadedArtsChan = make(chan bool)
	result.goLoadAsciiArts()
	return result
}

func (self *AsciiArtGallery) isArtsLoaded() bool {
	// if arts not loaded then wait for loading
	if !self.loadedArts {
		if _, ok := <-self.loadedArtsChan; ok {
			self.loadedArts = true
		}
	}
	return self.loadedArts
}

// loads ascii arts from data directory
func (self *AsciiArtGallery) goLoadAsciiArts() {
	self.loadedArts = false
	go func() {
		files := []string{}
		entries, err := ioutil.ReadDir("data/ascii")
		if err != nil {
			panic("Cannot read files from data/ascii")
		}
		for _, f := range entries {
			if strings.HasSuffix(f.Name(), ".dat") {
				files = append(files, path.Join("data/ascii", f.Name()))
			}
		}

		self.arts = make([]IAsciiArt, len(files))
		self.artfiles = files
		self.loadedArtsChan <- true
	}()
}

func (self *AsciiArtGallery) getAsciiArt(id int) IAsciiArt {
	var result IAsciiArt
	if self.arts[id] == nil {
		art, err := NewAsciiArtFromFile(self.artfiles[id])
		if err != nil {
			panic("Cannot load ascii file " + self.artfiles[id])
		}
		self.arts[id] = art
	}
	result = self.arts[id]
	return result
}

// select random ASCII art but not the one before
func (self *AsciiArtGallery) randomAsciiArt() IAsciiArt {
	if self.isArtsLoaded() {
		var id int

		// find id to return
		if len(self.arts) == 0 {
			panic("No arts found!")
		} else if len(self.arts) == 1 {
			id = 0
		} else {
			for {
				id = random.GetInt(0, len(self.arts)-1)
				if self.last != id {
					break
				}
			}
			self.last = id
		}
		return self.getAsciiArt(id)

	}
	return nil
}

func NewAsciiArtFromFile(fname string) (art IAsciiArt, err error) {

	fin, err := os.Open(fname)
	if err != nil {
		return nil, err
	}
	i, err := zlib.NewReader(fin)
	defer i.Close()
	if err != nil {
		return nil, err
	}
	d := gob.NewDecoder(i)
	a := AsciiArt{}
	d.Decode(&a)
	return &a, nil

}

func (self *AsciiArt) PutChar(console IConsole, x, y int) {
	color := self.colors[y+self.offsetY][x+self.offsetX]
	// char := int(self.chars[y+self.offsetY][x+self.offsetX])
	console.PutChar(x, y, 'x', BKGND_SET)
	console.SetCharForeground(x, y, color)
	console.SetCharBackground(x, y, COLOR_BLACK, BKGND_SET)

}

func (self *AsciiArt) Draw(console IConsole) {
	console.SetDefaultBackground(COLOR_BLACK)
	console.Clear()
	for x := 0; x < min(console.GetWidth(), self.width); x++ {
		for y := 0; y < min(console.GetHeight(), self.height); y++ {
			self.PutChar(console, x, y)
		}
	}
}

func (self *AsciiArt) Render(console IConsole, transition Transition, first bool) {
	if first {
		aas.firstTime = int(SysElapsedMilliseconds())
	}
	elapsed := int(SysElapsedMilliseconds()) - aas.firstTime

	switch transition {

	case INSTANT:
		{
			self.Draw(console)

		}
	case RANDOMIZE:
		{
			if first {
				console.SetDefaultBackground(COLOR_BLACK)
			}
			for i := 0; i < 500; i++ {
				x := random.GetInt(0, console.GetWidth()-1)
				y := random.GetInt(0, console.GetHeight()-1)

				// if in our picture then draw it
				if x < self.width && y < self.height {
					self.PutChar(console, x, y)
				} else {
					// else draw black background
					console.PutChar(x, y, ' ', BKGND_SET)
					console.SetCharForeground(x, y, COLOR_BLACK)
					console.SetCharBackground(x, y, COLOR_BLACK, BKGND_SET)
				}
			}

		}
	case FADE_IN:
		{
			if first {
				aas.secondary = NewConsole(console.GetWidth(), console.GetHeight())
				console.Blit(0, 0, console.GetWidth(), console.GetHeight(), aas.secondary, 0, 0, 1.0, 1.0)
			}
			console.SetDefaultBackground(COLOR_BLACK)
			console.Clear()
			// fade out old
			if elapsed <= TRANSITION_TIME {
				alpha := 1.0 - float32(elapsed)/TRANSITION_TIME
				aas.secondary.Blit(0, 0, console.GetWidth(), console.GetHeight(), console, 0, 0, alpha, alpha)
				// fade in new
			} else {
				self.Draw(aas.secondary)
				alpha := minf((float32(elapsed)-TRANSITION_TIME)/TRANSITION_TIME, 1)
				aas.secondary.Blit(0, 0, console.GetWidth(), console.GetHeight(), console, 0, 0, alpha, alpha)
			}

		}
	case MOVE_IN:
		{
			if first {
				aas.secondary = NewConsole(console.GetWidth(), console.GetHeight())
				aas.moveTransition = MoveTransition(random.GetInt(0, NB_MOVE_TRANSITIONS-1))
				self.Draw(aas.secondary)
			}
			widthPart := max(1, min(console.GetWidth(), int(float32(elapsed)/TRANSITION_TIME*float32(console.GetWidth()))))
			heightPart := max(1, min(console.GetHeight(), int(float32(elapsed)/TRANSITION_TIME*float32(console.GetHeight()))))

			switch aas.moveTransition {
			case MOVE_LEFT:
				aas.secondary.Blit(
					0, 0, widthPart, console.GetHeight(),
					console, 0, 0, 1.0, 1.0)
			case MOVE_RIGHT:
				aas.secondary.Blit(
					console.GetWidth()-widthPart, 0, widthPart, console.GetHeight(),
					console, console.GetWidth()-widthPart, 0, 1.0, 1.0)
			case MOVE_DOWN:
				aas.secondary.Blit(0, 0, console.GetWidth(), heightPart,
					console, 0, 0, 1.0, 1.0)
			case MOVE_UP:
				aas.secondary.Blit(
					0, console.GetHeight()-heightPart, console.GetWidth(), heightPart,
					console, 0, console.GetHeight()-heightPart, 1.0, 1.0)

			}

		}
	default:
		{
			// do nothing
			self.Render(console, FADE_IN, first)
			// TODO
			// rain/driple -> transitions -> similar to MOVE_IN transitions but for each streak sepreately
			// with different speed
			// fade in
			// have 2 consoles -> save current as screenshot
			// // and fill then new one and blit it
			// move in -> <-
			//  rain driple like matrix effect
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/belogik/goes"
	flags "github.com/jessevdk/go-flags"
	"github.com/miku/dupsquash"
)

type Work struct {
	Indices    []string
	Connection dupsquash.SearchConnection

	Fields        []string
	LikeText      string
	MinTermFreq   int
	MaxQueryTerms int
	Size          int

	Values []string
}

func Worker(in chan *Work, out chan [][]string, wg *sync.WaitGroup) {
	defer wg.Done()
	for work := range in {
		var query = map[string]interface{}{
			"query": map[string]interface{}{
				"more_like_this": map[string]interface{}{
					"fields":          work.Fields,
					"like_text":       work.LikeText,
					"min_term_freq":   work.MinTermFreq,
					"max_query_terms": work.MaxQueryTerms,
				},
			},
			"size": work.Size,
		}
		queryResults := Query(work.Connection, &work.Indices, &query)
		var results [][]string
		for _, result := range queryResults {
			parts := append(work.Values, result...)
			results = append(results, parts)
		}
		out <- results
	}
}

// FanInWriter writes the channel content to the writer
func FanInWriter(writer io.Writer, in chan [][]string, done chan bool) {
	for results := range in {
		for _, result := range results {
			writer.Write([]byte(strings.Join(result, "\t")))
			writer.Write([]byte("\n"))
		}
	}
	done <- true
}

// Query runs `query` over connection `conn` on `indices` and returns a slice of string slices
func Query(conn dupsquash.SearchConnection, indices *[]string, query *map[string]interface{}) [][]string {
	extraArgs := make(url.Values, 1)
	searchResults, err := conn.Search(*query, *indices, []string{""}, extraArgs)
	if err != nil {
		log.Fatalln(err)
	}
	results := make([][]string, len(searchResults.Hits.Hits))
	for i, hit := range searchResults.Hits.Hits {
		results[i] = []string{hit.Index, hit.Type, hit.Id, strconv.FormatFloat(hit.Score, 'f', 3, 64)}
	}
	return results
}

func main() {
	var opts struct {
		ElasticSearchHost string `long:"host" default:"localhost" description:"elasticsearch host" value-name:"HOST"`
		ElasticSearchPort string `long:"port" default:"9200" description:"elasticsearch port" value-name:"PORT"`

		Like string `short:"l" long:"like" description:"string to compare" value-name:"STRING"`

		LikeFile      string `short:"i" long:"file" description:"input file (TSV) with strings to compare" value-name:"FILENAME"`
		FileColumn    string `short:"f" long:"column" default:"1" description:"which column(s) to pick for the comparison" value-name:"COLUMN[S]"`
		FileDelimiter string `long:"delimiter" default:"\t" description:"column delimiter of the file" value-name:"DELIM"`
		FileNullValue string `long:"null-value" default:"<NULL>" description:"value that indicates empty value in input file" value-name:"STRING"`

		Index         string `long:"index" description:"index or indices (space separated)" value-name:"NAME[S]"`
		IndexFields   string `short:"x" default:"content.245.a content.245.b" long:"index-fields"description:"which index fields to use for comparison" value-name:"NAME[S]"`
		MinTermFreq   int    `long:"min-term-freq" description:"passed on lucene option" default:"1" value-name:"N"`
		MaxQueryTerms int    `long:"max-query-terms" description:"passed on lucene option" default:"25" value-name:"N"`
		Size          int    `short:"s" long:"size" description:"number of results per query" default:"5" value-name:"N"`

		NumWorkers  int    `short:"w" long:"workers" default:"0" description:"number of workers, 0 means number of available cpus" value-name:"N"`
		ShowVersion bool   `short:"V" default:"false" long:"version" description:"show version and exit"`
		ShowHelp    bool   `short:"h" default:"false" long:"help" description:"show this help message"`
		CpuProfile  string `long:"cpuprofile" description:"write pprof file" value-name:"FILE"`
	}

	argparser := flags.NewParser(&opts, flags.PrintErrors|flags.PassDoubleDash)
	_, err := argparser.Parse()
	if err != nil {
		log.Fatal(err)
	}

	if opts.ShowVersion {
		fmt.Printf("%s\n", dupsquash.AppVersion)
		return
	}

	argparser.Usage = fmt.Sprintf("[OPTIONS]")
	if opts.ShowHelp {
		argparser.WriteHelp(os.Stdout)
		return
	}

	if opts.NumWorkers == 0 {
		opts.NumWorkers = runtime.NumCPU()
		runtime.GOMAXPROCS(opts.NumWorkers)
	}

	if opts.NumWorkers < 0 {
		log.Fatal("value for --workers must be non-negative")
	}

	conn := goes.NewConnection(opts.ElasticSearchHost, opts.ElasticSearchPort)
	fields := strings.Fields(opts.IndexFields)
	indices := strings.Fields(opts.Index)

	if opts.LikeFile != "" {
		if _, err := os.Stat(opts.LikeFile); os.IsNotExist(err) {
			log.Fatalf("no such file or directory: %s\n", opts.LikeFile)
		}

		file, err := os.Open(opts.LikeFile)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		projector, err := dupsquash.ParseIndices(opts.FileColumn)
		if err != nil {
			log.Fatalf("could not parse column indices: %s\n", opts.FileColumn)
		}

		queue := make(chan *Work)
		results := make(chan [][]string)
		done := make(chan bool)

		writer := bufio.NewWriter(os.Stdout)
		defer writer.Flush()
		go FanInWriter(writer, results, done)

		var wg sync.WaitGroup
		for i := 0; i < opts.NumWorkers; i++ {
			wg.Add(1)
			go Worker(queue, results, &wg)
		}

		for scanner.Scan() {
			values := strings.Split(scanner.Text(), opts.FileDelimiter)
			likeText, err := dupsquash.ConcatenateValuesNull(values, projector, opts.FileNullValue)
			if err != nil {
				log.Fatal(err)
			}

			work := Work{
				Indices:       indices,
				Connection:    conn,
				Fields:        fields,
				LikeText:      likeText,
				MinTermFreq:   opts.MinTermFreq,
				MaxQueryTerms: opts.MaxQueryTerms,
				Size:          opts.Size,
				Values:        values,
			}
			queue <- &work
		}

		if err := scanner.Err(); err != nil {
			log.Fatal(err)
		}

		close(queue)
		wg.Wait()
		close(results)
		select {
		case <-time.After(1e9):
			break
		case <-done:
			break
		}
		return
	}

	if opts.Like != "" {
		var query = map[string]interface{}{
			"query": map[string]interface{}{
				"more_like_this": map[string]interface{}{
					"fields":          fields,
					"like_text":       opts.Like,
					"min_term_freq":   opts.MinTermFreq,
					"max_query_terms": opts.MaxQueryTerms,
				},
			},
			"size": opts.Size,
		}

		results := Query(conn, &indices, &query)
		for _, result := range results {
			fmt.Println(strings.Join(result, "\t"))
		}
		return
	}
}
<commit_msg>added cpu profiling<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/belogik/goes"
	flags "github.com/jessevdk/go-flags"
	"github.com/miku/dupsquash"
)

type Work struct {
	Indices    []string
	Connection dupsquash.SearchConnection

	Fields        []string
	LikeText      string
	MinTermFreq   int
	MaxQueryTerms int
	Size          int

	Values []string
}

func Worker(in chan *Work, out chan [][]string, wg *sync.WaitGroup) {
	defer wg.Done()
	for work := range in {
		var query = map[string]interface{}{
			"query": map[string]interface{}{
				"more_like_this": map[string]interface{}{
					"fields":          work.Fields,
					"like_text":       work.LikeText,
					"min_term_freq":   work.MinTermFreq,
					"max_query_terms": work.MaxQueryTerms,
				},
			},
			"size": work.Size,
		}
		queryResults := Query(work.Connection, &work.Indices, &query)
		var results [][]string
		for _, result := range queryResults {
			parts := append(work.Values, result...)
			results = append(results, parts)
		}
		out <- results
	}
}

// FanInWriter writes the channel content to the writer
func FanInWriter(writer io.Writer, in chan [][]string, done chan bool) {
	for results := range in {
		for _, result := range results {
			writer.Write([]byte(strings.Join(result, "\t")))
			writer.Write([]byte("\n"))
		}
	}
	done <- true
}

// Query runs `query` over connection `conn` on `indices` and returns a slice of string slices
func Query(conn dupsquash.SearchConnection, indices *[]string, query *map[string]interface{}) [][]string {
	extraArgs := make(url.Values, 1)
	searchResults, err := conn.Search(*query, *indices, []string{""}, extraArgs)
	if err != nil {
		log.Fatalln(err)
	}
	results := make([][]string, len(searchResults.Hits.Hits))
	for i, hit := range searchResults.Hits.Hits {
		results[i] = []string{hit.Index, hit.Type, hit.Id, strconv.FormatFloat(hit.Score, 'f', 3, 64)}
	}
	return results
}

func main() {
	var opts struct {
		ElasticSearchHost string `long:"host" default:"localhost" description:"elasticsearch host" value-name:"HOST"`
		ElasticSearchPort string `long:"port" default:"9200" description:"elasticsearch port" value-name:"PORT"`

		Like string `short:"l" long:"like" description:"string to compare" value-name:"STRING"`

		LikeFile      string `short:"i" long:"file" description:"input file (TSV) with strings to compare" value-name:"FILENAME"`
		FileColumn    string `short:"f" long:"column" default:"1" description:"which column(s) to pick for the comparison" value-name:"COLUMN[S]"`
		FileDelimiter string `long:"delimiter" default:"\t" description:"column delimiter of the file" value-name:"DELIM"`
		FileNullValue string `long:"null-value" default:"<NULL>" description:"value that indicates empty value in input file" value-name:"STRING"`

		Index         string `long:"index" description:"index or indices (space separated)" value-name:"NAME[S]"`
		IndexFields   string `short:"x" default:"content.245.a content.245.b" long:"index-fields"description:"which index fields to use for comparison" value-name:"NAME[S]"`
		MinTermFreq   int    `long:"min-term-freq" description:"passed on lucene option" default:"1" value-name:"N"`
		MaxQueryTerms int    `long:"max-query-terms" description:"passed on lucene option" default:"25" value-name:"N"`
		Size          int    `short:"s" long:"size" description:"number of results per query" default:"5" value-name:"N"`

		NumWorkers  int    `short:"w" long:"workers" default:"0" description:"number of workers, 0 means number of available cpus" value-name:"N"`
		ShowVersion bool   `short:"V" default:"false" long:"version" description:"show version and exit"`
		ShowHelp    bool   `short:"h" default:"false" long:"help" description:"show this help message"`
		CpuProfile  string `long:"cpuprofile" description:"write pprof file" value-name:"FILE"`
	}

	argparser := flags.NewParser(&opts, flags.PrintErrors|flags.PassDoubleDash)
	_, err := argparser.Parse()
	if err != nil {
		log.Fatal(err)
	}

	if opts.ShowVersion {
		fmt.Printf("%s\n", dupsquash.AppVersion)
		return
	}

	argparser.Usage = fmt.Sprintf("[OPTIONS]")
	if opts.ShowHelp {
		argparser.WriteHelp(os.Stdout)
		return
	}

	if opts.CpuProfile != "" {
		f, err := os.Create(opts.CpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if opts.NumWorkers == 0 {
		opts.NumWorkers = runtime.NumCPU()
		runtime.GOMAXPROCS(opts.NumWorkers)
	}

	if opts.NumWorkers < 0 {
		log.Fatal("value for --workers must be non-negative")
	}

	conn := goes.NewConnection(opts.ElasticSearchHost, opts.ElasticSearchPort)
	fields := strings.Fields(opts.IndexFields)
	indices := strings.Fields(opts.Index)

	if opts.LikeFile != "" {
		if _, err := os.Stat(opts.LikeFile); os.IsNotExist(err) {
			log.Fatalf("no such file or directory: %s\n", opts.LikeFile)
		}

		file, err := os.Open(opts.LikeFile)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		projector, err := dupsquash.ParseIndices(opts.FileColumn)
		if err != nil {
			log.Fatalf("could not parse column indices: %s\n", opts.FileColumn)
		}

		queue := make(chan *Work)
		results := make(chan [][]string)
		done := make(chan bool)

		writer := bufio.NewWriter(os.Stdout)
		defer writer.Flush()
		go FanInWriter(writer, results, done)

		var wg sync.WaitGroup
		for i := 0; i < opts.NumWorkers; i++ {
			wg.Add(1)
			go Worker(queue, results, &wg)
		}

		for scanner.Scan() {
			values := strings.Split(scanner.Text(), opts.FileDelimiter)
			likeText, err := dupsquash.ConcatenateValuesNull(values, projector, opts.FileNullValue)
			if err != nil {
				log.Fatal(err)
			}

			work := Work{
				Indices:       indices,
				Connection:    conn,
				Fields:        fields,
				LikeText:      likeText,
				MinTermFreq:   opts.MinTermFreq,
				MaxQueryTerms: opts.MaxQueryTerms,
				Size:          opts.Size,
				Values:        values,
			}
			queue <- &work
		}

		if err := scanner.Err(); err != nil {
			log.Fatal(err)
		}

		close(queue)
		wg.Wait()
		close(results)
		select {
		case <-time.After(1e9):
			break
		case <-done:
			break
		}
		return
	}

	if opts.Like != "" {
		var query = map[string]interface{}{
			"query": map[string]interface{}{
				"more_like_this": map[string]interface{}{
					"fields":          fields,
					"like_text":       opts.Like,
					"min_term_freq":   opts.MinTermFreq,
					"max_query_terms": opts.MaxQueryTerms,
				},
			},
			"size": opts.Size,
		}

		results := Query(conn, &indices, &query)
		for _, result := range results {
			fmt.Println(strings.Join(result, "\t"))
		}
		return
	}
}
<|endoftext|>


<commit_before>package ip2region

import (
	"fmt"
	"regexp"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/webx-top/com"
)

func TestIPWAN(t *testing.T) {
	wan, err := GetWANIP(0)
	if err != nil {
		t.Fatal(err)
	}
	fmt.Println(`IPv4:`, wan.IPv4)
	fmt.Println(`IPv6:`, wan.IPv6)
	r := regexp.MustCompile(IPv6Rule)
	matches := r.FindAllStringSubmatch(` 	2001:3CA1:010F:001A:121B:0000:0000:0010`, 1)
	com.Dump(matches)
	expected := [][]string{
		[]string{
			"2001:3CA1:010F:001A:121B:0000:0000:0010",
			"2001:3CA1:010F:001A:121B:0000:0000:0010",
			"", "",
		},
	}
	assert.Equal(t, expected, matches)
	//panic(``)
}
<commit_msg>improved<commit_after>package ip2region

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/webx-top/com"
)

func TestIPWAN(t *testing.T) {
	wan, err := GetWANIP(0)
	if err != nil {
		t.Fatal(err)
	}
	fmt.Println(`IPv4:`, wan.IPv4)
	fmt.Println(`IPv6:`, wan.IPv6)
	matches := ipv6Regexp.FindAllStringSubmatch(` 	2001:3CA1:010F:001A:121B:0000:0000:0010`, 1)
	com.Dump(matches)
	expected := [][]string{
		[]string{
			"2001:3CA1:010F:001A:121B:0000:0000:0010",
			"2001:3CA1:010F:001A:121B:0000:0000:0010",
			"", "",
		},
	}
	assert.Equal(t, expected, matches)
	//panic(``)
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package blockfile provides methods for reading packets from blockfiles
// generated by stenotype.
package blockfile

import (
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
	"unsafe"

	"github.com/google/gopacket"
	"github.com/google/stenographer/base"
	"github.com/google/stenographer/filecache"
	"github.com/google/stenographer/indexfile"
	"github.com/google/stenographer/query"
	"github.com/google/stenographer/stats"
	"golang.org/x/net/context"
)

// #include <linux/if_packet.h>
import "C"

var (
	v                = base.V // Verbose logging
	packetReadNanos  = stats.S.Get("packet_read_nanos")
	packetScanNanos  = stats.S.Get("packet_scan_nanos")
	packetsRead      = stats.S.Get("packets_read")
	packetsScanned   = stats.S.Get("packets_scanned")
	packetBlocksRead = stats.S.Get("packets_blocks_read")
)

// BlockFile provides an interface to a single stenotype file on disk and its
// associated index.
type BlockFile struct {
	name string
	f    *filecache.CachedFile
	i    *indexfile.IndexFile
	mu   sync.RWMutex // Stops Close() from invalidating a file before a current query is done with it.
	done chan struct{}
}

// NewBlockFile opens up a named block file (and its index), returning a handle
// which can be used to look up packets.
func NewBlockFile(filename string, fc *filecache.Cache) (*BlockFile, error) {
	v(1, "Blockfile opening: %q", filename)
	i, err := indexfile.NewIndexFile(indexfile.IndexPathFromBlockfilePath(filename), fc)
	if err != nil {
		return nil, fmt.Errorf("could not open index for %q: %v", filename, err)
	}
	return &BlockFile{
		f:    fc.Open(filename),
		i:    i,
		name: filename,
		done: make(chan struct{}),
	}, nil
}

// Name returns the name of the file underlying this blockfile.
func (b *BlockFile) Name() string {
	return b.name
}

// readPacket reads a single packet from the file at the given position.
// It updates the passed in CaptureInfo with information on the packet.
func (b *BlockFile) readPacket(pos int64, ci *gopacket.CaptureInfo) ([]byte, error) {
	// 28 bytes actually isn't the entire packet header, but it's all the fields
	// that we care about.
	packetsRead.Increment()
	defer packetReadNanos.NanoTimer()()
	var dataBuf [28]byte
	_, err := b.f.ReadAt(dataBuf[:], pos)
	if err != nil {
		return nil, err
	}
	pkt := (*C.struct_tpacket3_hdr)(unsafe.Pointer(&dataBuf[0]))
	*ci = gopacket.CaptureInfo{
		Timestamp:     time.Unix(int64(pkt.tp_sec), int64(pkt.tp_nsec)),
		Length:        int(pkt.tp_len),
		CaptureLength: int(pkt.tp_snaplen),
	}
	out := make([]byte, ci.CaptureLength)
	pos += int64(pkt.tp_mac)
	_, err = b.f.ReadAt(out, pos)
	return out, err
}

// Close cleans up this blockfile.
func (b *BlockFile) Close() (err error) {
	v(2, "Blockfile closing: %q", b.name)
	close(b.done)
	b.mu.Lock()
	defer b.mu.Unlock()
	v(3, "Blockfile closing file descriptors: %q", b.name)
	if e := b.i.Close(); e != nil {
		err = e
	}
	if e := b.f.Close(); e != nil {
		err = e
	}
	b.i, b.f = nil, nil
	return
}

// allPacketsIter implements Iter.
type allPacketsIter struct {
	*BlockFile
	blockData        []byte
	block            *C.struct_tpacket_hdr_v1
	pkt              *C.struct_tpacket3_hdr
	blockPacketsRead int
	blockOffset      int64
	packetOffset     int // offset of packet in block
	err              error
	done             bool
}

func (a *allPacketsIter) Next() bool {
	defer packetScanNanos.NanoTimer()()
	if a.err != nil || a.done {
		return false
	}
	for a.block == nil || a.blockPacketsRead == int(a.block.num_pkts) {
		packetBlocksRead.Increment()
		a.blockData = make([]byte, 1<<20)
		_, err := a.f.ReadAt(a.blockData[:], a.blockOffset)
		if err == io.EOF {
			a.done = true
			return false
		} else if err != nil {
			a.err = fmt.Errorf("could not read block at %v: %v", a.blockOffset, err)
			return false
		}
		baseHdr := (*C.struct_tpacket_block_desc)(unsafe.Pointer(&a.blockData[0]))
		a.block = (*C.struct_tpacket_hdr_v1)(unsafe.Pointer(&baseHdr.hdr[0]))
		a.blockOffset += 1 << 20
		a.blockPacketsRead = 0
		a.pkt = nil
	}
	a.blockPacketsRead++
	if a.pkt == nil {
		a.packetOffset = int(a.block.offset_to_first_pkt)
	} else if a.pkt.tp_next_offset != 0 {
		a.packetOffset += int(a.pkt.tp_next_offset)
	} else {
		a.err = errors.New("block format currently not supported")
		return false
	}
	a.pkt = (*C.struct_tpacket3_hdr)(unsafe.Pointer(&a.blockData[a.packetOffset]))
	packetsScanned.Increment()
	return true
}

func (a *allPacketsIter) Packet() *base.Packet {
	start := a.packetOffset + int(a.pkt.tp_mac)
	buf := a.blockData[start : start+int(a.pkt.tp_snaplen)]
	p := &base.Packet{Data: buf}
	p.CaptureInfo.Timestamp = time.Unix(int64(a.pkt.tp_sec), int64(a.pkt.tp_nsec))
	p.CaptureInfo.Length = int(a.pkt.tp_len)
	p.CaptureInfo.CaptureLength = int(a.pkt.tp_snaplen)
	return p
}

func (a *allPacketsIter) Err() error {
	return a.err
}

// AllPackets returns a packet channel to which all packets in the blockfile are
// sent.
func (b *BlockFile) AllPackets() *base.PacketChan {
	b.mu.RLock()
	c := base.NewPacketChan(100)
	go func() {
		defer b.mu.RUnlock()
		pkts := &allPacketsIter{BlockFile: b}
		for pkts.Next() {
			c.Send(pkts.Packet())
		}
		c.Close(pkts.Err())
	}()
	return c
}

// Positions returns the positions in the blockfile of all packets matched by
// the passed-in query.
func (b *BlockFile) Positions(ctx context.Context, q query.Query) (base.Positions, error) {
	b.mu.RLock()
	defer b.mu.RUnlock()
	return b.positionsLocked(ctx, q)
}

// positionsLocked returns the positions in the blockfile of all packets matched by
// the passed-in query.  b.mu must be locked.
func (b *BlockFile) positionsLocked(ctx context.Context, q query.Query) (base.Positions, error) {
	if b.i == nil || b.f == nil {
		// If we're closed, just return nothing.
		return nil, nil
	}
	return q.LookupIn(ctx, b.i)
}

// Lookup returns all packets in the blockfile matched by the passed-in query.
func (b *BlockFile) Lookup(ctx context.Context, q query.Query, out *base.PacketChan) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	var ci gopacket.CaptureInfo
	v(2, "Blockfile %q looking up query %q", b.name, q.String())
	start := time.Now()
	positions, err := b.positionsLocked(ctx, q)
	if err != nil {
		out.Close(fmt.Errorf("index lookup failure: %v", err))
		return
	}
	if positions.IsAllPositions() {
		v(2, "Blockfile %q reading all packets", b.name)
		iter := &allPacketsIter{BlockFile: b}
	all_packets_loop:
		for iter.Next() {
			select {
			case <-ctx.Done():
				v(2, "Blockfile %q canceling packet read", b.name)
				break all_packets_loop
			case <-b.done:
				v(2, "Blockfile %q closing, breaking out of query", b.name)
				break all_packets_loop
			case out.C <- iter.Packet():
			}
		}
		if iter.Err() != nil {
			out.Close(fmt.Errorf("error reading all packets from %q: %v", b.name, iter.Err()))
			return
		}
	} else {
		v(2, "Blockfile %q reading %v packets", b.name, len(positions))
	query_packets_loop:
		for _, pos := range positions {
			buffer, err := b.readPacket(pos, &ci)
			if err != nil {
				v(2, "Blockfile %q error reading packet: %v", b.name, err)
				out.Close(fmt.Errorf("error reading packets from %q @ %v: %v", b.name, pos, err))
				return
			}
			select {
			case <-ctx.Done():
				v(2, "Blockfile %q canceling packet read", b.name)
				break query_packets_loop
			case <-b.done:
				v(2, "Blockfile %q closing, breaking out of query", b.name)
				break query_packets_loop
			case out.C <- &base.Packet{Data: buffer, CaptureInfo: ci}:
			}
		}
	}
	v(2, "Blockfile %q finished reading all packets in %v", b.name, time.Since(start))
	out.Close(ctx.Err())
}

// DumpIndex dumps out a "human-readable" debug version of the blockfile's index
// to the given writer.
func (b *BlockFile) DumpIndex(out io.Writer, start, finish []byte) {
	b.mu.RLock()
	defer b.mu.RUnlock()
	b.i.Dump(out, start, finish)
}
<commit_msg>Add BlockFile Size.<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package blockfile provides methods for reading packets from blockfiles
// generated by stenotype.
package blockfile

import (
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
	"unsafe"

	"github.com/google/gopacket"
	"github.com/google/stenographer/base"
	"github.com/google/stenographer/filecache"
	"github.com/google/stenographer/indexfile"
	"github.com/google/stenographer/query"
	"github.com/google/stenographer/stats"
	"golang.org/x/net/context"
)

// #include <linux/if_packet.h>
import "C"

var (
	v                = base.V // Verbose logging
	packetReadNanos  = stats.S.Get("packet_read_nanos")
	packetScanNanos  = stats.S.Get("packet_scan_nanos")
	packetsRead      = stats.S.Get("packets_read")
	packetsScanned   = stats.S.Get("packets_scanned")
	packetBlocksRead = stats.S.Get("packets_blocks_read")
)

// BlockFile provides an interface to a single stenotype file on disk and its
// associated index.
type BlockFile struct {
	name string
	f    *filecache.CachedFile
	i    *indexfile.IndexFile
	mu   sync.RWMutex // Stops Close() from invalidating a file before a current query is done with it.
	done chan struct{}
	size int64
}

// NewBlockFile opens up a named block file (and its index), returning a handle
// which can be used to look up packets.
func NewBlockFile(filename string, fc *filecache.Cache) (*BlockFile, error) {
	v(1, "Blockfile opening: %q", filename)
	i, err := indexfile.NewIndexFile(indexfile.IndexPathFromBlockfilePath(filename), fc)
	if err != nil {
		return nil, fmt.Errorf("could not open index for %q: %v", filename, err)
	}
	f := fc.Open(filename)
	s, err := f.Stat()
	if err != nil {
		return nil, fmt.Errorf("could not stat file %q: %v", filename, err)
	}
	return &BlockFile{
		f:    f,
		i:    i,
		name: filename,
		done: make(chan struct{}),
		size: s.Size(),
	}, nil
}

// Name returns the name of the file underlying this blockfile.
func (b *BlockFile) Name() string {
	return b.name
}

// Size returns the size of the blockfile in bytes.
func (b *BlockFile) Size() int64 {
	return b.size
}

// readPacket reads a single packet from the file at the given position.
// It updates the passed in CaptureInfo with information on the packet.
func (b *BlockFile) readPacket(pos int64, ci *gopacket.CaptureInfo) ([]byte, error) {
	// 28 bytes actually isn't the entire packet header, but it's all the fields
	// that we care about.
	packetsRead.Increment()
	defer packetReadNanos.NanoTimer()()
	var dataBuf [28]byte
	_, err := b.f.ReadAt(dataBuf[:], pos)
	if err != nil {
		return nil, err
	}
	pkt := (*C.struct_tpacket3_hdr)(unsafe.Pointer(&dataBuf[0]))
	*ci = gopacket.CaptureInfo{
		Timestamp:     time.Unix(int64(pkt.tp_sec), int64(pkt.tp_nsec)),
		Length:        int(pkt.tp_len),
		CaptureLength: int(pkt.tp_snaplen),
	}
	out := make([]byte, ci.CaptureLength)
	pos += int64(pkt.tp_mac)
	_, err = b.f.ReadAt(out, pos)
	return out, err
}

// Close cleans up this blockfile.
func (b *BlockFile) Close() (err error) {
	v(2, "Blockfile closing: %q", b.name)
	close(b.done)
	b.mu.Lock()
	defer b.mu.Unlock()
	v(3, "Blockfile closing file descriptors: %q", b.name)
	if e := b.i.Close(); e != nil {
		err = e
	}
	if e := b.f.Close(); e != nil {
		err = e
	}
	b.i, b.f = nil, nil
	return
}

// allPacketsIter implements Iter.
type allPacketsIter struct {
	*BlockFile
	blockData        []byte
	block            *C.struct_tpacket_hdr_v1
	pkt              *C.struct_tpacket3_hdr
	blockPacketsRead int
	blockOffset      int64
	packetOffset     int // offset of packet in block
	err              error
	done             bool
}

func (a *allPacketsIter) Next() bool {
	defer packetScanNanos.NanoTimer()()
	if a.err != nil || a.done {
		return false
	}
	for a.block == nil || a.blockPacketsRead == int(a.block.num_pkts) {
		packetBlocksRead.Increment()
		a.blockData = make([]byte, 1<<20)
		_, err := a.f.ReadAt(a.blockData[:], a.blockOffset)
		if err == io.EOF {
			a.done = true
			return false
		} else if err != nil {
			a.err = fmt.Errorf("could not read block at %v: %v", a.blockOffset, err)
			return false
		}
		baseHdr := (*C.struct_tpacket_block_desc)(unsafe.Pointer(&a.blockData[0]))
		a.block = (*C.struct_tpacket_hdr_v1)(unsafe.Pointer(&baseHdr.hdr[0]))
		a.blockOffset += 1 << 20
		a.blockPacketsRead = 0
		a.pkt = nil
	}
	a.blockPacketsRead++
	if a.pkt == nil {
		a.packetOffset = int(a.block.offset_to_first_pkt)
	} else if a.pkt.tp_next_offset != 0 {
		a.packetOffset += int(a.pkt.tp_next_offset)
	} else {
		a.err = errors.New("block format currently not supported")
		return false
	}
	a.pkt = (*C.struct_tpacket3_hdr)(unsafe.Pointer(&a.blockData[a.packetOffset]))
	packetsScanned.Increment()
	return true
}

func (a *allPacketsIter) Packet() *base.Packet {
	start := a.packetOffset + int(a.pkt.tp_mac)
	buf := a.blockData[start : start+int(a.pkt.tp_snaplen)]
	p := &base.Packet{Data: buf}
	p.CaptureInfo.Timestamp = time.Unix(int64(a.pkt.tp_sec), int64(a.pkt.tp_nsec))
	p.CaptureInfo.Length = int(a.pkt.tp_len)
	p.CaptureInfo.CaptureLength = int(a.pkt.tp_snaplen)
	return p
}

func (a *allPacketsIter) Err() error {
	return a.err
}

// AllPackets returns a packet channel to which all packets in the blockfile are
// sent.
func (b *BlockFile) AllPackets() *base.PacketChan {
	b.mu.RLock()
	c := base.NewPacketChan(100)
	go func() {
		defer b.mu.RUnlock()
		pkts := &allPacketsIter{BlockFile: b}
		for pkts.Next() {
			c.Send(pkts.Packet())
		}
		c.Close(pkts.Err())
	}()
	return c
}

// Positions returns the positions in the blockfile of all packets matched by
// the passed-in query.
func (b *BlockFile) Positions(ctx context.Context, q query.Query) (base.Positions, error) {
	b.mu.RLock()
	defer b.mu.RUnlock()
	return b.positionsLocked(ctx, q)
}

// positionsLocked returns the positions in the blockfile of all packets matched by
// the passed-in query.  b.mu must be locked.
func (b *BlockFile) positionsLocked(ctx context.Context, q query.Query) (base.Positions, error) {
	if b.i == nil || b.f == nil {
		// If we're closed, just return nothing.
		return nil, nil
	}
	return q.LookupIn(ctx, b.i)
}

// Lookup returns all packets in the blockfile matched by the passed-in query.
func (b *BlockFile) Lookup(ctx context.Context, q query.Query, out *base.PacketChan) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	var ci gopacket.CaptureInfo
	v(2, "Blockfile %q looking up query %q", b.name, q.String())
	start := time.Now()
	positions, err := b.positionsLocked(ctx, q)
	if err != nil {
		out.Close(fmt.Errorf("index lookup failure: %v", err))
		return
	}
	if positions.IsAllPositions() {
		v(2, "Blockfile %q reading all packets", b.name)
		iter := &allPacketsIter{BlockFile: b}
	all_packets_loop:
		for iter.Next() {
			select {
			case <-ctx.Done():
				v(2, "Blockfile %q canceling packet read", b.name)
				break all_packets_loop
			case <-b.done:
				v(2, "Blockfile %q closing, breaking out of query", b.name)
				break all_packets_loop
			case out.C <- iter.Packet():
			}
		}
		if iter.Err() != nil {
			out.Close(fmt.Errorf("error reading all packets from %q: %v", b.name, iter.Err()))
			return
		}
	} else {
		v(2, "Blockfile %q reading %v packets", b.name, len(positions))
	query_packets_loop:
		for _, pos := range positions {
			buffer, err := b.readPacket(pos, &ci)
			if err != nil {
				v(2, "Blockfile %q error reading packet: %v", b.name, err)
				out.Close(fmt.Errorf("error reading packets from %q @ %v: %v", b.name, pos, err))
				return
			}
			select {
			case <-ctx.Done():
				v(2, "Blockfile %q canceling packet read", b.name)
				break query_packets_loop
			case <-b.done:
				v(2, "Blockfile %q closing, breaking out of query", b.name)
				break query_packets_loop
			case out.C <- &base.Packet{Data: buffer, CaptureInfo: ci}:
			}
		}
	}
	v(2, "Blockfile %q finished reading all packets in %v", b.name, time.Since(start))
	out.Close(ctx.Err())
}

// DumpIndex dumps out a "human-readable" debug version of the blockfile's index
// to the given writer.
func (b *BlockFile) DumpIndex(out io.Writer, start, finish []byte) {
	b.mu.RLock()
	defer b.mu.RUnlock()
	b.i.Dump(out, start, finish)
}
<|endoftext|>