<commit_before>package physical

import (
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/armon/go-metrics"
	"github.com/samuel/go-zookeeper/zk"
)

// ZookeeperBackend is a physical backend that stores data at specific
// prefix within Zookeeper. It is used in production situations as
// it allows Vault to run on multiple machines in a highly-available manner.
type ZookeeperBackend struct {
	path   string
	client *zk.Conn
}

// newZookeeperBackend constructs a Zookeeper backend using the given API client
// and the prefix in the KV store.
func newZookeeperBackend(conf map[string]string) (Backend, error) {
	// Get the path in Zookeeper
	path, ok := conf["path"]
	if !ok {
		path = "vault/"
	}

	// Ensure path is suffixed and prefixed (zk requires prefix /)
	if !strings.HasSuffix(path, "/") {
		path += "/"
	}
	if !strings.HasPrefix(path, "/") {
		path = "/" + path
	}

	// Configure the client
	machines := strings.Split(conf["address"], ",")

	client, _, err := zk.Connect(machines, time.Second)
	if err != nil {
		return nil, fmt.Errorf("client setup failed: %v", err)
	}

	// Setup the backend
	c := &ZookeeperBackend{
		path:   path,
		client: client,
	}
	return c, nil
}

// zookeeper requires nodes to be there before set and get
func (c *ZookeeperBackend) ensurePath(path string) {

	nodes := strings.Split(path, "/")

	acl := zk.WorldACL(zk.PermAll)

	fullPath := ""
	for _, node := range nodes {
		if strings.TrimSpace(node) != "" {
			fullPath += "/" + node

			if exists, _, _ := c.client.Exists(fullPath); !exists {
				c.client.Create(fullPath, nil, int32(0), acl)
			}
		}
	}
}

// Put is used to insert or update an entry
func (c *ZookeeperBackend) Put(entry *Entry) error {
	defer metrics.MeasureSince([]string{"zookeeper", "put"}, time.Now())

	fullPath := c.path + entry.Key

	c.ensurePath(fullPath)

	_, err := c.client.Set(fullPath, entry.Value, 0)

	return err
}

// Get is used to fetch an entry
func (c *ZookeeperBackend) Get(key string) (*Entry, error) {
	defer metrics.MeasureSince([]string{"zookeeper", "get"}, time.Now())

	fullPath := c.path + key

	c.ensurePath(fullPath)

	value, _, err := c.client.Get(fullPath)

	if err != nil {
		return nil, err
	}
	if value == nil {
		return nil, nil
	}
	ent := &Entry{
		Key:   key,
		Value: value,
	}
	return ent, nil
}

// Delete is used to permanently delete an entry
func (c *ZookeeperBackend) Delete(key string) error {
	defer metrics.MeasureSince([]string{"zookeeper", "delete"}, time.Now())

	fullPath := c.path + key

	exists, _, err := c.client.Exists(fullPath)

	if err != nil {
		return err
	}

	if exists {
		return c.client.Delete(fullPath, -1)
	} else {
		return nil
	}
}

// List is used ot list all the keys under a given
// prefix, up to the next prefix.
func (c *ZookeeperBackend) List(prefix string) ([]string, error) {
	defer metrics.MeasureSince([]string{"zookeeper", "list"}, time.Now())

	fullPath := strings.TrimSuffix(c.path+prefix, "/")

	c.ensurePath(fullPath)

	result, _, _ := c.client.Children(fullPath)

	children := []string{}

	for _, key := range result {
		children = append(children, key)

		nodeChildren, _, _ := c.client.Children(fullPath + "/" + key)
		if nodeChildren != nil && len(nodeChildren) > 0 {
			children = append(children, key+"/")
		}
	}

	sort.Strings(children)

	return children, nil
}

// // Lock is used for mutual exclusion based on the given key.
// func (c *ZookeeperBackend) LockWith(key, value string) (Lock, error) {
// 	// Create the lock
// 	opts := &api.LockOptions{
// 		Key:         key,
// 		Value:       []byte(value),
// 		SessionName: "Vault Lock",
// 	}
// 	lock, err := c.client.LockOpts(opts)
// 	if err != nil {
// 		return nil, fmt.Errorf("failed to create lock: %v", err)
// 	}
// 	cl := &ZookeeperLock{
// 		client: c.client,
// 		key:    key,
// 		lock:   lock,
// 	}
// 	return cl, nil
// }

// // DetectHostAddr is used to detect the host address by asking the Zookeeper agent
// func (c *ZookeeperBackend) DetectHostAddr() (string, error) {
// 	agent := c.client.Agent()
// 	self, err := agent.Self()
// 	if err != nil {
// 		return "", err
// 	}
// 	addr := self["Member"]["Addr"].(string)
// 	return addr, nil
// }

// ZookeeperLock is used to provide the Lock interface backed by Zookeeper
// type ZookeeperLock struct {
// 	client *zk.Conn
// 	key    string
// 	lock   *api.Lock
// }

// func (c *ZookeeperLock) Lock(stopCh <-chan struct{}) (<-chan struct{}, error) {
// 	return c.lock.Lock(stopCh)
// }

// func (c *ZookeeperLock) Unlock() error {
// 	return c.lock.Unlock()
// }

// func (c *ZookeeperLock) Value() (bool, string, error) {
// 	kv := c.client.KV()
// 	pair, _, err := kv.Get(c.key, nil)
// 	if err != nil {
// 		return false, "", err
// 	}
// 	if pair == nil {
// 		return false, "", nil
// 	}
// 	held := pair.Session != ""
// 	value := string(pair.Value)
// 	return held, value, nil
// }
<commit_msg>limit round trips on zk delete<commit_after>package physical

import (
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/armon/go-metrics"
	"github.com/samuel/go-zookeeper/zk"
)

// ZookeeperBackend is a physical backend that stores data at specific
// prefix within Zookeeper. It is used in production situations as
// it allows Vault to run on multiple machines in a highly-available manner.
type ZookeeperBackend struct {
	path   string
	client *zk.Conn
}

// newZookeeperBackend constructs a Zookeeper backend using the given API client
// and the prefix in the KV store.
func newZookeeperBackend(conf map[string]string) (Backend, error) {
	// Get the path in Zookeeper
	path, ok := conf["path"]
	if !ok {
		path = "vault/"
	}

	// Ensure path is suffixed and prefixed (zk requires prefix /)
	if !strings.HasSuffix(path, "/") {
		path += "/"
	}
	if !strings.HasPrefix(path, "/") {
		path = "/" + path
	}

	// Configure the client
	machines := strings.Split(conf["address"], ",")

	client, _, err := zk.Connect(machines, time.Second)
	if err != nil {
		return nil, fmt.Errorf("client setup failed: %v", err)
	}

	// Setup the backend
	c := &ZookeeperBackend{
		path:   path,
		client: client,
	}
	return c, nil
}

// zookeeper requires nodes to be there before set and get
func (c *ZookeeperBackend) ensurePath(path string) {

	nodes := strings.Split(path, "/")

	acl := zk.WorldACL(zk.PermAll)

	fullPath := ""
	for _, node := range nodes {
		if strings.TrimSpace(node) != "" {
			fullPath += "/" + node

			if exists, _, _ := c.client.Exists(fullPath); !exists {
				c.client.Create(fullPath, nil, int32(0), acl)
			}
		}
	}
}

// Put is used to insert or update an entry
func (c *ZookeeperBackend) Put(entry *Entry) error {
	defer metrics.MeasureSince([]string{"zookeeper", "put"}, time.Now())

	fullPath := c.path + entry.Key

	c.ensurePath(fullPath)

	_, err := c.client.Set(fullPath, entry.Value, 0)

	return err
}

// Get is used to fetch an entry
func (c *ZookeeperBackend) Get(key string) (*Entry, error) {
	defer metrics.MeasureSince([]string{"zookeeper", "get"}, time.Now())

	fullPath := c.path + key

	c.ensurePath(fullPath)

	value, _, err := c.client.Get(fullPath)

	if err != nil {
		return nil, err
	}
	if value == nil {
		return nil, nil
	}
	ent := &Entry{
		Key:   key,
		Value: value,
	}
	return ent, nil
}

// Delete is used to permanently delete an entry
func (c *ZookeeperBackend) Delete(key string) error {
	defer metrics.MeasureSince([]string{"zookeeper", "delete"}, time.Now())

	fullPath := c.path + key

	err := c.client.Delete(fullPath, -1)

	if err == zk.ErrNoNode {
		return nil
	} else {
		return err
	}
}

// List is used ot list all the keys under a given
// prefix, up to the next prefix.
func (c *ZookeeperBackend) List(prefix string) ([]string, error) {
	defer metrics.MeasureSince([]string{"zookeeper", "list"}, time.Now())

	fullPath := strings.TrimSuffix(c.path+prefix, "/")

	c.ensurePath(fullPath)

	result, _, _ := c.client.Children(fullPath)

	children := []string{}

	for _, key := range result {
		children = append(children, key)

		nodeChildren, _, _ := c.client.Children(fullPath + "/" + key)
		if nodeChildren != nil && len(nodeChildren) > 0 {
			children = append(children, key+"/")
		}
	}

	sort.Strings(children)

	return children, nil
}

// // Lock is used for mutual exclusion based on the given key.
// func (c *ZookeeperBackend) LockWith(key, value string) (Lock, error) {
// 	// Create the lock
// 	opts := &api.LockOptions{
// 		Key:         key,
// 		Value:       []byte(value),
// 		SessionName: "Vault Lock",
// 	}
// 	lock, err := c.client.LockOpts(opts)
// 	if err != nil {
// 		return nil, fmt.Errorf("failed to create lock: %v", err)
// 	}
// 	cl := &ZookeeperLock{
// 		client: c.client,
// 		key:    key,
// 		lock:   lock,
// 	}
// 	return cl, nil
// }

// // DetectHostAddr is used to detect the host address by asking the Zookeeper agent
// func (c *ZookeeperBackend) DetectHostAddr() (string, error) {
// 	agent := c.client.Agent()
// 	self, err := agent.Self()
// 	if err != nil {
// 		return "", err
// 	}
// 	addr := self["Member"]["Addr"].(string)
// 	return addr, nil
// }

// ZookeeperLock is used to provide the Lock interface backed by Zookeeper
// type ZookeeperLock struct {
// 	client *zk.Conn
// 	key    string
// 	lock   *api.Lock
// }

// func (c *ZookeeperLock) Lock(stopCh <-chan struct{}) (<-chan struct{}, error) {
// 	return c.lock.Lock(stopCh)
// }

// func (c *ZookeeperLock) Unlock() error {
// 	return c.lock.Unlock()
// }

// func (c *ZookeeperLock) Value() (bool, string, error) {
// 	kv := c.client.KV()
// 	pair, _, err := kv.Get(c.key, nil)
// 	if err != nil {
// 		return false, "", err
// 	}
// 	if pair == nil {
// 		return false, "", nil
// 	}
// 	held := pair.Session != ""
// 	value := string(pair.Value)
// 	return held, value, nil
// }
<|endoftext|>


<commit_before>package compact

import (
	"github.com/bcicen/ctop/config"
	"github.com/bcicen/ctop/logging"
	"github.com/bcicen/ctop/models"
	ui "github.com/gizak/termui"
)

var log = logging.Init()

type CompactCol interface {
	ui.GridBufferer
	Reset()
	Highlight()
	UnHighlight()
	SetMeta(models.Meta)
	SetMetrics(models.Metrics)
}

type Compact struct {
	Bg     *RowBg
	Cols   []CompactCol
	X, Y   int
	Width  int
	Height int
}

func NewCompact(id string) *Compact {
	// truncate container id
	if len(id) > 12 {
		id = id[:12]
	}
	row := &Compact{
		Bg: NewRowBg(),
		Cols: []CompactCol{
			NewStatus(),
			&NameCol{NewTextCol("-")},
			&CIDCol{NewTextCol(id)},
			&CPUCol{NewGaugeCol()},
			&MemCol{NewGaugeCol()},
			&NetCol{NewTextCol("-")},
			&IOCol{NewTextCol("-")},
			&PIDCol{NewTextCol("-")},
		},
		X:      1,
		Height: 1,
	}
	return row
}

func (row *Compact) SetMeta(m models.Meta) {
	for _, w := range row.Cols {
		w.SetMeta(m)
	}
}

func (row *Compact) SetMetrics(m models.Metrics) {
	for _, w := range row.Cols {
		w.SetMetrics(m)
	}
}

// Set gauges, counters, etc. to default unread values
func (row *Compact) Reset() {
	for _, w := range row.Cols {
		w.Reset()
	}
}

func (row *Compact) GetHeight() int { return row.Height }
func (row *Compact) SetX(x int)     { row.X = x }

func (row *Compact) SetY(y int) {
	if y == row.Y {
		return
	}

	row.Bg.Y = y
	for _, w := range row.Cols {
		w.SetY(y)
	}
	row.Y = y
}

func (row *Compact) SetWidth(width int) {
	if width == row.Width {
		return
	}
	x := row.X

	row.Bg.SetX(x + colWidths[0] + 1)
	row.Bg.SetWidth(width)

	autoWidth := calcWidth(width)
	for n, w := range row.Cols {
		// set static width, if provided
		if colWidths[n] != 0 {
			w.SetX(x)
			w.SetWidth(colWidths[n])
			x += colWidths[n]
			continue
		}
		// else use auto width
		w.SetX(x)
		w.SetWidth(autoWidth)
		x += autoWidth + colSpacing
	}
	row.Width = width
}

func (row *Compact) Buffer() ui.Buffer {
	buf := ui.NewBuffer()
	buf.Merge(row.Bg.Buffer())
	for _, w := range row.Cols {
		buf.Merge(w.Buffer())
	}
	return buf
}

func (row *Compact) Highlight() {
	row.Cols[1].Highlight()
	if config.GetSwitchVal("fullRowCursor") {
		for _, w := range row.Cols {
			w.Highlight()
		}
	}
}

func (row *Compact) UnHighlight() {
	row.Cols[1].UnHighlight()
	if config.GetSwitchVal("fullRowCursor") {
		for _, w := range row.Cols {
			w.UnHighlight()
		}
	}
}

type RowBg struct {
	*ui.Par
}

func NewRowBg() *RowBg {
	bg := ui.NewPar("")
	bg.Height = 1
	bg.Border = false
	bg.Bg = ui.ThemeAttr("par.text.bg")
	return &RowBg{bg}
}

func (w *RowBg) Highlight()   { w.Bg = ui.ThemeAttr("par.text.fg") }
func (w *RowBg) UnHighlight() { w.Bg = ui.ThemeAttr("par.text.bg") }
<commit_msg>continuing compact widget refactor<commit_after><|endoftext|>


<commit_before>package apps

import (
	"io"
	"net/url"
	"regexp"
	"time"

	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/sirupsen/logrus"
)

var slugReg = regexp.MustCompile(`^[a-z0-9\-]+$`)

// Operation is the type of operation the installer is created for.
type Operation int

const (
	// Install operation for installing an application
	Install Operation = iota + 1
	// Update operation for updating an application
	Update
	// Delete operation for deleting an application
	Delete
)

// Installer is used to install or update applications.
type Installer struct {
	fetcher  Fetcher
	op       Operation
	fs       Copier
	db       couchdb.Database
	endState State

	man  Manifest
	src  *url.URL
	slug string

	err  error
	errc chan error
	manc chan Manifest
	log  *logrus.Entry
}

// InstallerOptions provides the slug name of the application along with the
// source URL.
type InstallerOptions struct {
	Type        AppType
	Operation   Operation
	Slug        string
	SourceURL   string
	Deactivated bool
}

// Fetcher interface should be implemented by the underlying transport
// used to fetch the application data.
type Fetcher interface {
	// FetchManifest should returns an io.ReadCloser to read the
	// manifest data
	FetchManifest(src *url.URL) (io.ReadCloser, error)
	// Fetch should download the application and install it in the given
	// directory.
	Fetch(src *url.URL, fs Copier, man Manifest) error
}

// NewInstaller creates a new Installer
func NewInstaller(db couchdb.Database, fs Copier, opts *InstallerOptions) (*Installer, error) {
	if opts.Operation == 0 {
		panic("Missing installer operation")
	}
	if opts.Type != Webapp && opts.Type != Konnector {
		panic("Bad or missing installer type")
	}

	slug := opts.Slug
	if slug == "" || !slugReg.MatchString(slug) {
		return nil, ErrInvalidSlugName
	}

	// For konnectors applications, we actually create a tar archive in which the
	// sources are stored before copying the archive into the application
	// storage.
	if opts.Type == Konnector {
		fs = newTarCopier(fs, KonnectorArchiveName)
	}

	man, err := GetBySlug(db, slug, opts.Type)
	if opts.Operation == Install {
		if err == nil {
			return nil, ErrAlreadyExists
		}
		if err != ErrNotFound {
			return nil, err
		}
		err = nil
		switch opts.Type {
		case Webapp:
			man = &WebappManifest{}
		case Konnector:
			man = &KonnManifest{}
		}
	} else if err != nil {
		return nil, err
	}

	var src *url.URL
	switch opts.Operation {
	case Install:
		if opts.SourceURL == "" {
			return nil, ErrMissingSource
		}
		src, err = url.Parse(opts.SourceURL)
	case Update, Delete:
		var srcString string
		if opts.SourceURL == "" {
			srcString = man.Source()
		} else {
			srcString = opts.SourceURL
		}
		src, err = url.Parse(srcString)
	}
	if err != nil {
		return nil, err
	}

	var endState State
	if opts.Deactivated || man.State() == Installed {
		endState = Installed
	} else {
		endState = Ready
	}

	log := db.Logger()

	var fetcher Fetcher
	switch src.Scheme {
	case "git":
		fetcher = newGitFetcher(opts.Type, log)
	default:
		return nil, ErrNotSupportedSource
	}

	return &Installer{
		fetcher:  fetcher,
		op:       opts.Operation,
		db:       db,
		fs:       fs,
		endState: endState,

		man:  man,
		src:  src,
		slug: slug,

		errc: make(chan error, 1),
		manc: make(chan Manifest, 2),
		log:  log,
	}, nil
}

// Run will install, update or delete the application linked to the installer,
// depending on specified operation. It will report its progress or error (see
// Poll method) and should be run asynchronously.
func (i *Installer) Run() {
	defer i.endOfProc()
	switch i.op {
	case Install:
		i.man, i.err = i.install()
	case Update:
		i.man, i.err = i.update()
	case Delete:
		i.man, i.err = i.delete()
	}
	return
}

// RunSync does the same work as Run but can be used synchronously.
func (i *Installer) RunSync() (Manifest, error) {
	go i.Run()
	for {
		man, done, err := i.Poll()
		if err != nil {
			return nil, err
		}
		if done {
			return man, nil
		}
	}
}

func (i *Installer) endOfProc() {
	man, err := i.man, i.err
	if man == nil || err == ErrBadState {
		i.errc <- err
		return
	}
	if err != nil {
		man.SetState(Errored)
		man.SetError(err)
		man.Update(i.db)
		i.errc <- err
		return
	}
	man.SetState(i.endState)
	man.Update(i.db)
	i.manc <- i.man
}

// install will perform the installation of an application. It returns the
// freshly fetched manifest from the source along with a possible error in case
// the installation went wrong.
//
// Note that the fetched manifest is returned even if an error occurred while
// upgrading.
func (i *Installer) install() (Manifest, error) {
	i.log.Infof("[apps] Start install: %s %s", i.slug, i.src.String())
	man := i.man
	if err := i.ReadManifest(Installing, man); err != nil {
		return nil, err
	}
	if err := man.Create(i.db); err != nil {
		return man, err
	}
	i.manc <- man
	return man, i.fetcher.Fetch(i.src, i.fs, man)
}

// update will perform the update of an already installed application. It
// returns the freshly fetched manifest from the source along with a possible
// error in case the update went wrong.
//
// Note that the fetched manifest is returned even if an error occurred while
// upgrading.
func (i *Installer) update() (Manifest, error) {
	i.log.Infof("[apps] Start update: %s %s", i.slug, i.src.String())
	man := i.man
	if err := i.checkState(man); err != nil {
		return nil, err
	}
	if err := i.ReadManifest(Upgrading, man); err != nil {
		return man, err
	}
	if err := man.Update(i.db); err != nil {
		return man, err
	}
	i.manc <- man
	return man, i.fetcher.Fetch(i.src, i.fs, man)
}

func (i *Installer) delete() (Manifest, error) {
	i.log.Infof("[apps] Start delete: %s %s", i.slug, i.src.String())
	man := i.man
	if err := i.checkState(man); err != nil {
		return nil, err
	}
	return man, i.man.Delete(i.db)
}

// checkState returns whether or not the manifest is in the right state to
// perform an update or deletion.
func (i *Installer) checkState(man Manifest) error {
	state := man.State()
	if state == Ready ||
		state == Installed ||
		state == Errored {
		return nil
	}
	timeAgo := time.Now().Add(-2 * time.Minute)
	if man.LastUpdate().Before(timeAgo) {
		return nil
	}
	return ErrBadState
}

// ReadManifest will fetch the manifest and read its JSON content into the
// passed manifest pointer.
//
// The State field of the manifest will be set to the specified state.
func (i *Installer) ReadManifest(state State, man Manifest) error {
	r, err := i.fetcher.FetchManifest(i.src)
	if err != nil {
		return err
	}
	defer r.Close()
	man.SetState(state)
	return man.ReadManifest(io.LimitReader(r, ManifestMaxSize), i.slug, i.src.String())
}

// Poll should be used to monitor the progress of the Installer.
func (i *Installer) Poll() (Manifest, bool, error) {
	select {
	case man := <-i.manc:
		state := man.State()
		done := (state == Ready || state == Installed)
		return man, done, nil
	case err := <-i.errc:
		return nil, false, err
	}
}
<commit_msg>Increase last update duration<commit_after>package apps

import (
	"io"
	"net/url"
	"regexp"
	"time"

	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/sirupsen/logrus"
)

var slugReg = regexp.MustCompile(`^[a-z0-9\-]+$`)

// Operation is the type of operation the installer is created for.
type Operation int

const (
	// Install operation for installing an application
	Install Operation = iota + 1
	// Update operation for updating an application
	Update
	// Delete operation for deleting an application
	Delete
)

// Installer is used to install or update applications.
type Installer struct {
	fetcher  Fetcher
	op       Operation
	fs       Copier
	db       couchdb.Database
	endState State

	man  Manifest
	src  *url.URL
	slug string

	err  error
	errc chan error
	manc chan Manifest
	log  *logrus.Entry
}

// InstallerOptions provides the slug name of the application along with the
// source URL.
type InstallerOptions struct {
	Type        AppType
	Operation   Operation
	Slug        string
	SourceURL   string
	Deactivated bool
}

// Fetcher interface should be implemented by the underlying transport
// used to fetch the application data.
type Fetcher interface {
	// FetchManifest should returns an io.ReadCloser to read the
	// manifest data
	FetchManifest(src *url.URL) (io.ReadCloser, error)
	// Fetch should download the application and install it in the given
	// directory.
	Fetch(src *url.URL, fs Copier, man Manifest) error
}

// NewInstaller creates a new Installer
func NewInstaller(db couchdb.Database, fs Copier, opts *InstallerOptions) (*Installer, error) {
	if opts.Operation == 0 {
		panic("Missing installer operation")
	}
	if opts.Type != Webapp && opts.Type != Konnector {
		panic("Bad or missing installer type")
	}

	slug := opts.Slug
	if slug == "" || !slugReg.MatchString(slug) {
		return nil, ErrInvalidSlugName
	}

	// For konnectors applications, we actually create a tar archive in which the
	// sources are stored before copying the archive into the application
	// storage.
	if opts.Type == Konnector {
		fs = newTarCopier(fs, KonnectorArchiveName)
	}

	man, err := GetBySlug(db, slug, opts.Type)
	if opts.Operation == Install {
		if err == nil {
			return nil, ErrAlreadyExists
		}
		if err != ErrNotFound {
			return nil, err
		}
		err = nil
		switch opts.Type {
		case Webapp:
			man = &WebappManifest{}
		case Konnector:
			man = &KonnManifest{}
		}
	} else if err != nil {
		return nil, err
	}

	var src *url.URL
	switch opts.Operation {
	case Install:
		if opts.SourceURL == "" {
			return nil, ErrMissingSource
		}
		src, err = url.Parse(opts.SourceURL)
	case Update, Delete:
		var srcString string
		if opts.SourceURL == "" {
			srcString = man.Source()
		} else {
			srcString = opts.SourceURL
		}
		src, err = url.Parse(srcString)
	}
	if err != nil {
		return nil, err
	}

	var endState State
	if opts.Deactivated || man.State() == Installed {
		endState = Installed
	} else {
		endState = Ready
	}

	log := db.Logger()

	var fetcher Fetcher
	switch src.Scheme {
	case "git":
		fetcher = newGitFetcher(opts.Type, log)
	default:
		return nil, ErrNotSupportedSource
	}

	return &Installer{
		fetcher:  fetcher,
		op:       opts.Operation,
		db:       db,
		fs:       fs,
		endState: endState,

		man:  man,
		src:  src,
		slug: slug,

		errc: make(chan error, 1),
		manc: make(chan Manifest, 2),
		log:  log,
	}, nil
}

// Run will install, update or delete the application linked to the installer,
// depending on specified operation. It will report its progress or error (see
// Poll method) and should be run asynchronously.
func (i *Installer) Run() {
	defer i.endOfProc()
	switch i.op {
	case Install:
		i.man, i.err = i.install()
	case Update:
		i.man, i.err = i.update()
	case Delete:
		i.man, i.err = i.delete()
	}
	return
}

// RunSync does the same work as Run but can be used synchronously.
func (i *Installer) RunSync() (Manifest, error) {
	go i.Run()
	for {
		man, done, err := i.Poll()
		if err != nil {
			return nil, err
		}
		if done {
			return man, nil
		}
	}
}

func (i *Installer) endOfProc() {
	man, err := i.man, i.err
	if man == nil || err == ErrBadState {
		i.errc <- err
		return
	}
	if err != nil {
		man.SetState(Errored)
		man.SetError(err)
		man.Update(i.db)
		i.errc <- err
		return
	}
	man.SetState(i.endState)
	man.Update(i.db)
	i.manc <- i.man
}

// install will perform the installation of an application. It returns the
// freshly fetched manifest from the source along with a possible error in case
// the installation went wrong.
//
// Note that the fetched manifest is returned even if an error occurred while
// upgrading.
func (i *Installer) install() (Manifest, error) {
	i.log.Infof("[apps] Start install: %s %s", i.slug, i.src.String())
	man := i.man
	if err := i.ReadManifest(Installing, man); err != nil {
		return nil, err
	}
	if err := man.Create(i.db); err != nil {
		return man, err
	}
	i.manc <- man
	return man, i.fetcher.Fetch(i.src, i.fs, man)
}

// update will perform the update of an already installed application. It
// returns the freshly fetched manifest from the source along with a possible
// error in case the update went wrong.
//
// Note that the fetched manifest is returned even if an error occurred while
// upgrading.
func (i *Installer) update() (Manifest, error) {
	i.log.Infof("[apps] Start update: %s %s", i.slug, i.src.String())
	man := i.man
	if err := i.checkState(man); err != nil {
		return nil, err
	}
	if err := i.ReadManifest(Upgrading, man); err != nil {
		return man, err
	}
	if err := man.Update(i.db); err != nil {
		return man, err
	}
	i.manc <- man
	return man, i.fetcher.Fetch(i.src, i.fs, man)
}

func (i *Installer) delete() (Manifest, error) {
	i.log.Infof("[apps] Start delete: %s %s", i.slug, i.src.String())
	man := i.man
	if err := i.checkState(man); err != nil {
		return nil, err
	}
	return man, i.man.Delete(i.db)
}

// checkState returns whether or not the manifest is in the right state to
// perform an update or deletion.
func (i *Installer) checkState(man Manifest) error {
	state := man.State()
	if state == Ready ||
		state == Installed ||
		state == Errored {
		return nil
	}
	if time.Since(man.LastUpdate()) > 15*time.Minute {
		return nil
	}
	return ErrBadState
}

// ReadManifest will fetch the manifest and read its JSON content into the
// passed manifest pointer.
//
// The State field of the manifest will be set to the specified state.
func (i *Installer) ReadManifest(state State, man Manifest) error {
	r, err := i.fetcher.FetchManifest(i.src)
	if err != nil {
		return err
	}
	defer r.Close()
	man.SetState(state)
	return man.ReadManifest(io.LimitReader(r, ManifestMaxSize), i.slug, i.src.String())
}

// Poll should be used to monitor the progress of the Installer.
func (i *Installer) Poll() (Manifest, bool, error) {
	select {
	case man := <-i.manc:
		state := man.State()
		done := (state == Ready || state == Installed)
		return man, done, nil
	case err := <-i.errc:
		return nil, false, err
	}
}
<|endoftext|>


<commit_before>package influxdb_output

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"io/ioutil"
	"log"
	"math"
	"strings"
	"time"

	influxdb2 "github.com/influxdata/influxdb-client-go/v2"
	"github.com/karimra/gnmic/formatters"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/proto"
)

const (
	defaultURL               = "http://localhost:8086"
	defaultBatchSize         = 1000
	defaultFlushTimer        = 10 * time.Second
	defaultHealthCheckPeriod = 30 * time.Second

	numWorkers    = 1
	loggingPrefix = "[influxdb_output] "
)

func init() {
	outputs.Register("influxdb", func() outputs.Output {
		return &InfluxDBOutput{
			Cfg:       &Config{},
			eventChan: make(chan *formatters.EventMsg),
			reset:     make(chan struct{}),
			startSig:  make(chan struct{}),
			logger:    log.New(ioutil.Discard, loggingPrefix, log.LstdFlags|log.Lmicroseconds),
		}
	})
}

type InfluxDBOutput struct {
	Cfg       *Config
	client    influxdb2.Client
	logger    *log.Logger
	cancelFn  context.CancelFunc
	eventChan chan *formatters.EventMsg
	reset     chan struct{}
	startSig  chan struct{}
	wasUP     bool
	evps      []formatters.EventProcessor
	dbVersion string
}
type Config struct {
	URL                string        `mapstructure:"url,omitempty"`
	Org                string        `mapstructure:"org,omitempty"`
	Bucket             string        `mapstructure:"bucket,omitempty"`
	Token              string        `mapstructure:"token,omitempty"`
	BatchSize          uint          `mapstructure:"batch-size,omitempty"`
	FlushTimer         time.Duration `mapstructure:"flush-timer,omitempty"`
	UseGzip            bool          `mapstructure:"use-gzip,omitempty"`
	EnableTLS          bool          `mapstructure:"enable-tls,omitempty"`
	HealthCheckPeriod  time.Duration `mapstructure:"health-check-period,omitempty"`
	Debug              bool          `mapstructure:"debug,omitempty"`
	EventProcessors    []string      `mapstructure:"event-processors,omitempty"`
	EnableMetrics      bool          `mapstructure:"enable-metrics,omitempty"`
	OverrideTimestamps bool          `mapstructure:"override-timestamps,omitempty"`
}

func (k *InfluxDBOutput) String() string {
	b, err := json.Marshal(k)
	if err != nil {
		return ""
	}
	return string(b)
}

func (i *InfluxDBOutput) SetLogger(logger *log.Logger) {
	if logger != nil && i.logger != nil {
		i.logger.SetOutput(logger.Writer())
		i.logger.SetFlags(logger.Flags())
	}
}

func (i *InfluxDBOutput) SetEventProcessors(ps map[string]map[string]interface{}, logger *log.Logger, tcs map[string]interface{}) {
	for _, epName := range i.Cfg.EventProcessors {
		if epCfg, ok := ps[epName]; ok {
			epType := ""
			for k := range epCfg {
				epType = k
				break
			}
			if in, ok := formatters.EventProcessors[epType]; ok {
				ep := in()
				err := ep.Init(epCfg[epType], formatters.WithLogger(logger), formatters.WithTargets(tcs))
				if err != nil {
					i.logger.Printf("failed initializing event processor '%s' of type='%s': %v", epName, epType, err)
					continue
				}
				i.evps = append(i.evps, ep)
				i.logger.Printf("added event processor '%s' of type=%s to influxdb output", epName, epType)
				continue
			}
			i.logger.Printf("%q event processor has an unknown type=%q", epName, epType)
			continue
		}
		i.logger.Printf("%q event processor not found!", epName)
	}
}

func (i *InfluxDBOutput) Init(ctx context.Context, name string, cfg map[string]interface{}, opts ...outputs.Option) error {
	err := outputs.DecodeConfig(cfg, i.Cfg)
	if err != nil {
		return err
	}
	for _, opt := range opts {
		opt(i)
	}
	if i.Cfg.URL == "" {
		i.Cfg.URL = defaultURL
	}
	if i.Cfg.BatchSize == 0 {
		i.Cfg.BatchSize = defaultBatchSize
	}
	if i.Cfg.FlushTimer == 0 {
		i.Cfg.FlushTimer = defaultFlushTimer
	}
	if i.Cfg.HealthCheckPeriod == 0 {
		i.Cfg.HealthCheckPeriod = defaultHealthCheckPeriod
	}

	iopts := influxdb2.DefaultOptions().
		SetUseGZip(i.Cfg.UseGzip).
		SetBatchSize(i.Cfg.BatchSize).
		SetFlushInterval(uint(i.Cfg.FlushTimer.Milliseconds()))
	if i.Cfg.EnableTLS {
		iopts.SetTLSConfig(&tls.Config{
			InsecureSkipVerify: true,
		})
	}
	if i.Cfg.Debug {
		iopts.SetLogLevel(3)
	}
	ctx, i.cancelFn = context.WithCancel(ctx)
CRCLIENT:
	i.client = influxdb2.NewClientWithOptions(i.Cfg.URL, i.Cfg.Token, iopts)
	// start influx health check
	err = i.health(ctx)
	if err != nil {
		i.logger.Printf("failed to check influxdb health: %v", err)
		time.Sleep(10 * time.Second)
		goto CRCLIENT
	}
	i.wasUP = true
	go i.healthCheck(ctx)
	i.logger.Printf("initialized influxdb client: %s", i.String())

	for k := 0; k < numWorkers; k++ {
		go i.worker(ctx, k)
	}
	go func() {
		<-ctx.Done()
		i.Close()
	}()
	return nil
}

func (i *InfluxDBOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := formatters.ResponseToEventMsgs(measName, rsp, meta, i.evps...)
		if err != nil {
			i.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			select {
			case <-ctx.Done():
				return
			case <-i.reset:
				return
			case i.eventChan <- ev:
			}
		}
	}
}

func (i *InfluxDBOutput) WriteEvent(ctx context.Context, ev *formatters.EventMsg) {}

func (i *InfluxDBOutput) Close() error {
	i.logger.Printf("closing client...")
	i.cancelFn()
	i.logger.Printf("closed.")
	return nil
}
func (i *InfluxDBOutput) RegisterMetrics(reg *prometheus.Registry) {}

func (i *InfluxDBOutput) healthCheck(ctx context.Context) {
	ticker := time.NewTicker(i.Cfg.HealthCheckPeriod)
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			i.health(ctx)
		}
	}
}

func (i *InfluxDBOutput) health(ctx context.Context) error {
	res, err := i.client.Health(ctx)
	if err != nil {
		i.logger.Printf("failed health check: %v", err)
		if i.wasUP {
			close(i.reset)
			i.reset = make(chan struct{})
		}
		return err
	}
	if res != nil {
		if res.Version != nil {
			i.dbVersion = *res.Version
		}
		b, err := json.Marshal(res)
		if err != nil {
			i.logger.Printf("failed to marshal health check result: %v", err)
			i.logger.Printf("health check result: %+v", res)
			if i.wasUP {
				close(i.reset)
				i.reset = make(chan struct{})
			}
			return err
		}
		i.wasUP = true
		close(i.startSig)
		i.startSig = make(chan struct{})
		i.logger.Printf("health check result: %s", string(b))
		return nil
	}
	i.wasUP = true
	close(i.startSig)
	i.startSig = make(chan struct{})
	i.logger.Print("health check result is nil")
	return nil
}

func (i *InfluxDBOutput) worker(ctx context.Context, idx int) {
	firstStart := true
START:
	if !firstStart {
		i.logger.Printf("worker-%d waiting for client recovery", idx)
		<-i.startSig
	}
	i.logger.Printf("starting worker-%d", idx)
	writer := i.client.WriteAPI(i.Cfg.Org, i.Cfg.Bucket)
	//defer writer.Flush()
	for {
		select {
		case <-ctx.Done():
			if ctx.Err() != nil {
				i.logger.Printf("worker-%d err=%v", idx, ctx.Err())
			}
			i.logger.Printf("worker-%d terminating...", idx)
			return
		case ev := <-i.eventChan:
			for n, v := range ev.Values {
				switch v := v.(type) {
				case *gnmi.Decimal64:
					ev.Values[n] = float64(v.Digits) / math.Pow10(int(v.Precision))
				}
			}
			if ev.Timestamp == 0 || i.Cfg.OverrideTimestamps {
				ev.Timestamp = time.Now().UnixNano()
			}
			i.convertUints(ev)
			writer.WritePoint(influxdb2.NewPoint(ev.Name, ev.Tags, ev.Values, time.Unix(0, ev.Timestamp)))
		case <-i.reset:
			firstStart = false
			i.logger.Printf("resetting worker-%d...", idx)
			goto START
		case err := <-writer.Errors():
			i.logger.Printf("worker-%d write error: %v", idx, err)
		}
	}
}

func (i *InfluxDBOutput) SetName(name string)        {}
func (i *InfluxDBOutput) SetClusterName(name string) {}

func (i *InfluxDBOutput) convertUints(ev *formatters.EventMsg) {
	if !strings.HasPrefix(i.dbVersion, "1.8") {
		return
	}
	for k, v := range ev.Values {
		switch v := v.(type) {
		case uint:
			ev.Values[k] = int(v)
		case uint8:
			ev.Values[k] = int(v)
		case uint16:
			ev.Values[k] = int(v)
		case uint32:
			ev.Values[k] = int(v)
		case uint64:
			ev.Values[k] = int(v)
		}
	}
}
<commit_msg> influxdb output<commit_after>package influxdb_output

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"io/ioutil"
	"log"
	"math"
	"strings"
	"text/template"
	"time"

	influxdb2 "github.com/influxdata/influxdb-client-go/v2"
	"github.com/karimra/gnmic/formatters"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/proto"
)

const (
	defaultURL               = "http://localhost:8086"
	defaultBatchSize         = 1000
	defaultFlushTimer        = 10 * time.Second
	defaultHealthCheckPeriod = 30 * time.Second

	numWorkers    = 1
	loggingPrefix = "[influxdb_output] "
)

func init() {
	outputs.Register("influxdb", func() outputs.Output {
		return &InfluxDBOutput{
			Cfg:       &Config{},
			eventChan: make(chan *formatters.EventMsg),
			reset:     make(chan struct{}),
			startSig:  make(chan struct{}),
			logger:    log.New(ioutil.Discard, loggingPrefix, log.LstdFlags|log.Lmicroseconds),
		}
	})
}

type InfluxDBOutput struct {
	Cfg       *Config
	client    influxdb2.Client
	logger    *log.Logger
	cancelFn  context.CancelFunc
	eventChan chan *formatters.EventMsg
	reset     chan struct{}
	startSig  chan struct{}
	wasUP     bool
	evps      []formatters.EventProcessor
	dbVersion string

	targetTpl *template.Template
}
type Config struct {
	URL                string        `mapstructure:"url,omitempty"`
	Org                string        `mapstructure:"org,omitempty"`
	Bucket             string        `mapstructure:"bucket,omitempty"`
	Token              string        `mapstructure:"token,omitempty"`
	BatchSize          uint          `mapstructure:"batch-size,omitempty"`
	FlushTimer         time.Duration `mapstructure:"flush-timer,omitempty"`
	UseGzip            bool          `mapstructure:"use-gzip,omitempty"`
	EnableTLS          bool          `mapstructure:"enable-tls,omitempty"`
	HealthCheckPeriod  time.Duration `mapstructure:"health-check-period,omitempty"`
	Debug              bool          `mapstructure:"debug,omitempty"`
	AddTarget          string        `mapstructure:"add-target,omitempty"`
	TargetTemplate     string        `mapstructure:"target-template,omitempty"`
	EventProcessors    []string      `mapstructure:"event-processors,omitempty"`
	EnableMetrics      bool          `mapstructure:"enable-metrics,omitempty"`
	OverrideTimestamps bool          `mapstructure:"override-timestamps,omitempty"`
}

func (k *InfluxDBOutput) String() string {
	b, err := json.Marshal(k)
	if err != nil {
		return ""
	}
	return string(b)
}

func (i *InfluxDBOutput) SetLogger(logger *log.Logger) {
	if logger != nil && i.logger != nil {
		i.logger.SetOutput(logger.Writer())
		i.logger.SetFlags(logger.Flags())
	}
}

func (i *InfluxDBOutput) SetEventProcessors(ps map[string]map[string]interface{}, logger *log.Logger, tcs map[string]interface{}) {
	for _, epName := range i.Cfg.EventProcessors {
		if epCfg, ok := ps[epName]; ok {
			epType := ""
			for k := range epCfg {
				epType = k
				break
			}
			if in, ok := formatters.EventProcessors[epType]; ok {
				ep := in()
				err := ep.Init(epCfg[epType], formatters.WithLogger(logger), formatters.WithTargets(tcs))
				if err != nil {
					i.logger.Printf("failed initializing event processor '%s' of type='%s': %v", epName, epType, err)
					continue
				}
				i.evps = append(i.evps, ep)
				i.logger.Printf("added event processor '%s' of type=%s to influxdb output", epName, epType)
				continue
			}
			i.logger.Printf("%q event processor has an unknown type=%q", epName, epType)
			continue
		}
		i.logger.Printf("%q event processor not found!", epName)
	}
}

func (i *InfluxDBOutput) Init(ctx context.Context, name string, cfg map[string]interface{}, opts ...outputs.Option) error {
	err := outputs.DecodeConfig(cfg, i.Cfg)
	if err != nil {
		return err
	}
	for _, opt := range opts {
		opt(i)
	}
	if i.Cfg.URL == "" {
		i.Cfg.URL = defaultURL
	}
	if i.Cfg.BatchSize == 0 {
		i.Cfg.BatchSize = defaultBatchSize
	}
	if i.Cfg.FlushTimer == 0 {
		i.Cfg.FlushTimer = defaultFlushTimer
	}
	if i.Cfg.HealthCheckPeriod == 0 {
		i.Cfg.HealthCheckPeriod = defaultHealthCheckPeriod
	}

	iopts := influxdb2.DefaultOptions().
		SetUseGZip(i.Cfg.UseGzip).
		SetBatchSize(i.Cfg.BatchSize).
		SetFlushInterval(uint(i.Cfg.FlushTimer.Milliseconds()))
	if i.Cfg.EnableTLS {
		iopts.SetTLSConfig(&tls.Config{
			InsecureSkipVerify: true,
		})
	}
	if i.Cfg.TargetTemplate == "" {
		i.targetTpl = outputs.DefaultTargetTemplate
	} else if i.Cfg.AddTarget != "" {
		i.targetTpl, err = template.New("target-template").
			Funcs(outputs.TemplateFuncs).
			Parse(i.Cfg.TargetTemplate)
		if err != nil {
			return err
		}
	}
	if i.Cfg.Debug {
		iopts.SetLogLevel(3)
	}
	ctx, i.cancelFn = context.WithCancel(ctx)
CRCLIENT:
	i.client = influxdb2.NewClientWithOptions(i.Cfg.URL, i.Cfg.Token, iopts)
	// start influx health check
	err = i.health(ctx)
	if err != nil {
		i.logger.Printf("failed to check influxdb health: %v", err)
		time.Sleep(10 * time.Second)
		goto CRCLIENT
	}
	i.wasUP = true
	go i.healthCheck(ctx)
	i.logger.Printf("initialized influxdb client: %s", i.String())

	for k := 0; k < numWorkers; k++ {
		go i.worker(ctx, k)
	}
	go func() {
		<-ctx.Done()
		i.Close()
	}()
	return nil
}

func (i *InfluxDBOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	err := outputs.AddSubscriptionTarget(rsp, meta, i.Cfg.AddTarget, i.targetTpl)
	if err != nil {
		i.logger.Printf("failed to add target to the response: %v", err)
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := formatters.ResponseToEventMsgs(measName, rsp, meta, i.evps...)
		if err != nil {
			i.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			select {
			case <-ctx.Done():
				return
			case <-i.reset:
				return
			case i.eventChan <- ev:
			}
		}
	}
}

func (i *InfluxDBOutput) WriteEvent(ctx context.Context, ev *formatters.EventMsg) {}

func (i *InfluxDBOutput) Close() error {
	i.logger.Printf("closing client...")
	i.cancelFn()
	i.logger.Printf("closed.")
	return nil
}
func (i *InfluxDBOutput) RegisterMetrics(reg *prometheus.Registry) {}

func (i *InfluxDBOutput) healthCheck(ctx context.Context) {
	ticker := time.NewTicker(i.Cfg.HealthCheckPeriod)
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			i.health(ctx)
		}
	}
}

func (i *InfluxDBOutput) health(ctx context.Context) error {
	res, err := i.client.Health(ctx)
	if err != nil {
		i.logger.Printf("failed health check: %v", err)
		if i.wasUP {
			close(i.reset)
			i.reset = make(chan struct{})
		}
		return err
	}
	if res != nil {
		if res.Version != nil {
			i.dbVersion = *res.Version
		}
		b, err := json.Marshal(res)
		if err != nil {
			i.logger.Printf("failed to marshal health check result: %v", err)
			i.logger.Printf("health check result: %+v", res)
			if i.wasUP {
				close(i.reset)
				i.reset = make(chan struct{})
			}
			return err
		}
		i.wasUP = true
		close(i.startSig)
		i.startSig = make(chan struct{})
		i.logger.Printf("health check result: %s", string(b))
		return nil
	}
	i.wasUP = true
	close(i.startSig)
	i.startSig = make(chan struct{})
	i.logger.Print("health check result is nil")
	return nil
}

func (i *InfluxDBOutput) worker(ctx context.Context, idx int) {
	firstStart := true
START:
	if !firstStart {
		i.logger.Printf("worker-%d waiting for client recovery", idx)
		<-i.startSig
	}
	i.logger.Printf("starting worker-%d", idx)
	writer := i.client.WriteAPI(i.Cfg.Org, i.Cfg.Bucket)
	//defer writer.Flush()
	for {
		select {
		case <-ctx.Done():
			if ctx.Err() != nil {
				i.logger.Printf("worker-%d err=%v", idx, ctx.Err())
			}
			i.logger.Printf("worker-%d terminating...", idx)
			return
		case ev := <-i.eventChan:
			for n, v := range ev.Values {
				switch v := v.(type) {
				case *gnmi.Decimal64:
					ev.Values[n] = float64(v.Digits) / math.Pow10(int(v.Precision))
				}
			}
			if ev.Timestamp == 0 || i.Cfg.OverrideTimestamps {
				ev.Timestamp = time.Now().UnixNano()
			}
			i.convertUints(ev)
			writer.WritePoint(influxdb2.NewPoint(ev.Name, ev.Tags, ev.Values, time.Unix(0, ev.Timestamp)))
		case <-i.reset:
			firstStart = false
			i.logger.Printf("resetting worker-%d...", idx)
			goto START
		case err := <-writer.Errors():
			i.logger.Printf("worker-%d write error: %v", idx, err)
		}
	}
}

func (i *InfluxDBOutput) SetName(name string)        {}
func (i *InfluxDBOutput) SetClusterName(name string) {}

func (i *InfluxDBOutput) convertUints(ev *formatters.EventMsg) {
	if !strings.HasPrefix(i.dbVersion, "1.8") {
		return
	}
	for k, v := range ev.Values {
		switch v := v.(type) {
		case uint:
			ev.Values[k] = int(v)
		case uint8:
			ev.Values[k] = int(v)
		case uint16:
			ev.Values[k] = int(v)
		case uint32:
			ev.Values[k] = int(v)
		case uint64:
			ev.Values[k] = int(v)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Mender Software AS
//
//    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 images

import (
	"net/http"
	"strconv"
	"time"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/asaskevich/govalidator"
	"github.com/mendersoftware/artifacts/mvc"
	"github.com/pkg/errors"
)

// API input validation constants
const (
	DefaultDownloadLinkExpire = 60
	DefaultUploadLinkExpire   = 60

	// AWS limitation is 1 week
	MaxLinkExpire = 60 * 7 * 24
)

var (
	ErrIDNotUUIDv4        = errors.New("ID is not UUIDv4")
	ErrInvalidExpireParam = errors.New("Invalid expire parameter")
)

type ImagesModeler interface {
	ListImages(filters map[string]string) ([]*SoftwareImage, error)
	UploadLink(imageID string, expire time.Duration) (*Link, error)
	DownloadLink(imageID string, expire time.Duration) (*Link, error)
	GetImage(id string) (*SoftwareImage, error)
	DeleteImage(imageID string) error
	CreateImage(constructorData *SoftwareImageConstructor) (string, error)
	EditImage(id string, constructorData *SoftwareImageConstructor) (bool, error)
}

type SoftwareImagesController struct {
	views mvc.RESTViewDefaults
	model ImagesModeler
}

func NewSoftwareImagesController(model ImagesModeler, views mvc.RESTViewDefaults) *SoftwareImagesController {
	return &SoftwareImagesController{
		model: model,
		views: views,
	}
}

func (s *SoftwareImagesController) GetImage(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	image, err := s.model.GetImage(id)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	if image == nil {
		s.views.RenderErrorNotFound(w)
	}

	s.views.RenderSuccessGet(w, image)
}

func (s *SoftwareImagesController) ListImages(w rest.ResponseWriter, r *rest.Request) {

	list, err := s.model.ListImages(r.PathParams)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	s.views.RenderSuccessGet(w, list)
}

func (s *SoftwareImagesController) UploadLink(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	expire := DefaultDownloadLinkExpire
	expireStr := r.URL.Query().Get("expire")

	// Validate input
	if !govalidator.IsNull(expireStr) {
		if !s.validExpire(expireStr) {
			s.views.RenderError(w, ErrInvalidExpireParam, http.StatusBadRequest)
			return
		}

		var err error
		expire, err = strconv.Atoi(expireStr)
		if err != nil {
			s.views.RenderError(w, err, http.StatusInternalServerError)
			return
		}
	}

	link, err := s.model.UploadLink(id, time.Duration(expire)*time.Minute)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	if link == nil {
		s.views.RenderErrorNotFound(w)
		return
	}

	s.views.RenderSuccessGet(w, link)
}

func (s *SoftwareImagesController) DownloadLink(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	expire := DefaultUploadLinkExpire
	expireStr := r.URL.Query().Get("expire")

	// Validate input
	if !govalidator.IsNull(expireStr) {
		if !s.validExpire(expireStr) {
			s.views.RenderError(w, ErrInvalidExpireParam, http.StatusBadRequest)
			return
		}

		var err error
		expire, err = strconv.Atoi(expireStr)
		if err != nil {
			s.views.RenderError(w, err, http.StatusInternalServerError)
			return
		}
	}

	link, err := s.model.DownloadLink(id, time.Duration(expire)*time.Minute)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	if link == nil {
		s.views.RenderErrorNotFound(w)
		return
	}

	s.views.RenderSuccessGet(w, link)
}

func (s *SoftwareImagesController) validExpire(expire string) bool {

	if govalidator.IsNull(expire) {
		return false
	}

	if !govalidator.IsInt(expire) {
		return false
	}

	number, err := strconv.ParseFloat(expire, 64)
	if err != nil {
		return false
	}

	if !govalidator.InRange(number, 0, MaxLinkExpire) {
		return false
	}

	return true
}

func (s *SoftwareImagesController) DeleteImage(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	if err := s.model.DeleteImage(id); err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	s.views.RenderSuccessDelete(w)
}

func (s *SoftwareImagesController) EditImage(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	var constructor *SoftwareImageConstructor

	if err := r.DecodeJsonPayload(&constructor); err != nil {
		s.views.RenderError(w, err, http.StatusBadRequest)
		return
	}

	if err := constructor.Validate(); err != nil {
		s.views.RenderError(w, err, http.StatusBadRequest)
		return
	}

	found, err := s.model.EditImage(id, constructor)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	if !found {
		s.views.RenderErrorNotFound(w)
		return
	}

	s.views.RenderSuccessPut(w)
}

func (s *SoftwareImagesController) NewImage(w rest.ResponseWriter, r *rest.Request) {

	var constructor *SoftwareImageConstructor

	if err := r.DecodeJsonPayload(&constructor); err != nil {
		s.views.RenderError(w, err, http.StatusBadRequest)
		return
	}

	if err := constructor.Validate(); err != nil {
		s.views.RenderError(w, err, http.StatusBadRequest)
		return
	}

	id, err := s.model.CreateImage(constructor)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	s.views.RenderSuccessPost(w, r, id)
}
<commit_msg>Fix missing return in images controller GetImage()<commit_after>// Copyright 2016 Mender Software AS
//
//    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 images

import (
	"net/http"
	"strconv"
	"time"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/asaskevich/govalidator"
	"github.com/mendersoftware/artifacts/mvc"
	"github.com/pkg/errors"
)

// API input validation constants
const (
	DefaultDownloadLinkExpire = 60
	DefaultUploadLinkExpire   = 60

	// AWS limitation is 1 week
	MaxLinkExpire = 60 * 7 * 24
)

var (
	ErrIDNotUUIDv4        = errors.New("ID is not UUIDv4")
	ErrInvalidExpireParam = errors.New("Invalid expire parameter")
)

type ImagesModeler interface {
	ListImages(filters map[string]string) ([]*SoftwareImage, error)
	UploadLink(imageID string, expire time.Duration) (*Link, error)
	DownloadLink(imageID string, expire time.Duration) (*Link, error)
	GetImage(id string) (*SoftwareImage, error)
	DeleteImage(imageID string) error
	CreateImage(constructorData *SoftwareImageConstructor) (string, error)
	EditImage(id string, constructorData *SoftwareImageConstructor) (bool, error)
}

type SoftwareImagesController struct {
	views mvc.RESTViewDefaults
	model ImagesModeler
}

func NewSoftwareImagesController(model ImagesModeler, views mvc.RESTViewDefaults) *SoftwareImagesController {
	return &SoftwareImagesController{
		model: model,
		views: views,
	}
}

func (s *SoftwareImagesController) GetImage(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	image, err := s.model.GetImage(id)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	if image == nil {
		s.views.RenderErrorNotFound(w)
		return
	}

	s.views.RenderSuccessGet(w, image)
}

func (s *SoftwareImagesController) ListImages(w rest.ResponseWriter, r *rest.Request) {

	list, err := s.model.ListImages(r.PathParams)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	s.views.RenderSuccessGet(w, list)
}

func (s *SoftwareImagesController) UploadLink(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	expire := DefaultDownloadLinkExpire
	expireStr := r.URL.Query().Get("expire")

	// Validate input
	if !govalidator.IsNull(expireStr) {
		if !s.validExpire(expireStr) {
			s.views.RenderError(w, ErrInvalidExpireParam, http.StatusBadRequest)
			return
		}

		var err error
		expire, err = strconv.Atoi(expireStr)
		if err != nil {
			s.views.RenderError(w, err, http.StatusInternalServerError)
			return
		}
	}

	link, err := s.model.UploadLink(id, time.Duration(expire)*time.Minute)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	if link == nil {
		s.views.RenderErrorNotFound(w)
		return
	}

	s.views.RenderSuccessGet(w, link)
}

func (s *SoftwareImagesController) DownloadLink(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	expire := DefaultUploadLinkExpire
	expireStr := r.URL.Query().Get("expire")

	// Validate input
	if !govalidator.IsNull(expireStr) {
		if !s.validExpire(expireStr) {
			s.views.RenderError(w, ErrInvalidExpireParam, http.StatusBadRequest)
			return
		}

		var err error
		expire, err = strconv.Atoi(expireStr)
		if err != nil {
			s.views.RenderError(w, err, http.StatusInternalServerError)
			return
		}
	}

	link, err := s.model.DownloadLink(id, time.Duration(expire)*time.Minute)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	if link == nil {
		s.views.RenderErrorNotFound(w)
		return
	}

	s.views.RenderSuccessGet(w, link)
}

func (s *SoftwareImagesController) validExpire(expire string) bool {

	if govalidator.IsNull(expire) {
		return false
	}

	if !govalidator.IsInt(expire) {
		return false
	}

	number, err := strconv.ParseFloat(expire, 64)
	if err != nil {
		return false
	}

	if !govalidator.InRange(number, 0, MaxLinkExpire) {
		return false
	}

	return true
}

func (s *SoftwareImagesController) DeleteImage(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	if err := s.model.DeleteImage(id); err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	s.views.RenderSuccessDelete(w)
}

func (s *SoftwareImagesController) EditImage(w rest.ResponseWriter, r *rest.Request) {

	id := r.PathParam("id")

	if !govalidator.IsUUIDv4(id) {
		s.views.RenderError(w, ErrIDNotUUIDv4, http.StatusBadRequest)
		return
	}

	var constructor *SoftwareImageConstructor

	if err := r.DecodeJsonPayload(&constructor); err != nil {
		s.views.RenderError(w, err, http.StatusBadRequest)
		return
	}

	if err := constructor.Validate(); err != nil {
		s.views.RenderError(w, err, http.StatusBadRequest)
		return
	}

	found, err := s.model.EditImage(id, constructor)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	if !found {
		s.views.RenderErrorNotFound(w)
		return
	}

	s.views.RenderSuccessPut(w)
}

func (s *SoftwareImagesController) NewImage(w rest.ResponseWriter, r *rest.Request) {

	var constructor *SoftwareImageConstructor

	if err := r.DecodeJsonPayload(&constructor); err != nil {
		s.views.RenderError(w, err, http.StatusBadRequest)
		return
	}

	if err := constructor.Validate(); err != nil {
		s.views.RenderError(w, err, http.StatusBadRequest)
		return
	}

	id, err := s.model.CreateImage(constructor)
	if err != nil {
		s.views.RenderError(w, err, http.StatusInternalServerError)
		return
	}

	s.views.RenderSuccessPost(w, r, id)
}
<|endoftext|>


<commit_before>package whale

import (
	"strings"
	"time"
	"strconv"
)

import (
	"butler/null"
	"gokogiri/xpath"
	"golog"
	"rubex"
	tp "tritium/proto"
	"steno"
	"go-cache"
	"go-cache/arc"
)

type Whale struct {
	Log *golog.Logger
	Debugger steno.Debugger
	RegexpCache       cache.Cache
	XPathCache        cache.Cache
}

type EngineContext struct {
	Functions                []*Function
	Types                    []string
	Exports                  [][]string
	Logs                     []string
	ImportedFiles            []string
	Env                      map[string]string
	MatchStack               []string
	MatchShouldContinueStack []bool
	Yields                   []*YieldBlock
	*Whale
	*tp.Transform
	Rrules []*tp.RewriteRule

	InnerReplacer     *rubex.Regexp
	HeaderContentType *rubex.Regexp

	// Debug info
	Filename          string
	HadError          bool
	Deadline time.Time
	Mobjects []MemoryObject
	MessagePath string
}

const OutputBufferSize = 500 * 1024 //500KB
const defaultMobjects = 4
const TimeoutError = "EngineTimeout"

func NewEngine(logger *golog.Logger, debugger steno.Debugger) *Whale {
	e := &Whale{
		Log: logger,
		Debugger: debugger,
		RegexpCache:              arc.NewARCache(1000),
		XPathCache:               arc.NewARCache(1000),
	}
	e.RegexpCache.SetCleanFunc(CleanRegexpObject)
	e.XPathCache.SetCleanFunc(CleanXpathExpObject)
	return e
}

func NewEngineCtx(eng *Whale, vars map[string]string, transform *tp.Transform, rrules []*tp.RewriteRule, deadline time.Time, messagePath string) (ctx *EngineContext) {
	ctx = &EngineContext{
		Whale:                    eng,
		Exports:                  make([][]string, 0),
		Logs:                     make([]string, 0),
		Env:                      vars,
		Transform:                transform,
		Rrules:                   rrules,
		MatchStack:               make([]string, 0),
		MatchShouldContinueStack: make([]bool, 0),
		Yields:                   make([]*YieldBlock, 0),
		HadError:                 false,

		Deadline:                 deadline,
		Mobjects:                 make([]MemoryObject, 0, defaultMobjects),
		MessagePath:              messagePath,
	}
	return
}

func (eng *Whale) Free() {
	eng.RegexpCache.Reset()
	eng.XPathCache.Reset()
}

func (eng *Whale) Run(transform *tp.Transform, rrules []*tp.RewriteRule, input interface{}, vars map[string]string, deadline time.Time, customer, project, messagePath string) (output string, exports [][]string, logs []string) {
	ctx := NewEngineCtx(eng, vars, transform, rrules, deadline, messagePath)
	defer ctx.Free()
	ctx.Yields = append(ctx.Yields, &YieldBlock{Vars: make(map[string]interface{})})
	ctx.UsePackage(transform.Pkg)
	scope := &Scope{Value: input.(string)}
	obj := transform.Objects[0]
	ctx.Filename = null.GetString(obj.Name)
	ctx.RunInstruction(scope, obj.Root)
	output = scope.Value.(string)
	exports = ctx.Exports
	logs = ctx.Logs
	return
}

func (eng *Whale) GetCacheStats() (int, int, int, int) {
	return eng.RegexpCache.GetHitRate(), eng.RegexpCache.GetUsageRate(), eng.XPathCache.GetHitRate(), eng.XPathCache.GetUsageRate()
}

func (ctx *EngineContext) Free() {
    for _, o := range ctx.Mobjects {
		if o != nil {
			o.Free()
		}
	}
}

func (ctx *EngineContext) RunInstruction(scope *Scope, ins *tp.Instruction) (returnValue interface{}) {
	defer func() {
		//TODO Stack traces seem to get truncated on syslog...
		if x := recover(); x != nil {
			err, ok := x.(error)
			errString := ""
			if ok {
				errString = err.Error()
			} else {
				errString = x.(string)
			}
			if errString != TimeoutError {
				if ctx.HadError == false {
					ctx.HadError = true
					errString = errString + "\n" + ins.Type.String() + " " + null.GetString(ins.Value) + "\n\n\nTritium Stack\n=========\n\n"
				}
				errString = errString + ctx.FileAndLine(ins) + "\n"
			}
			panic(errString)
		}
	}()

	ctx.Whale.Debugger.TrapInstruction(ctx.Filename, ins, scope.Value, scope.Index)

	if time.Now().After(ctx.Deadline) {
		panic(TimeoutError)
	}

	// If our object is invalid, then skip it
	if null.GetBool(ins.IsValid) == false {
		panic("Invalid instruction. Should have stopped before linking!")
	}
	indent := ""
	for i := 0; i < len(ctx.Yields); i++ {
		indent += "\t"
	}

	returnValue = ""
	switch *ins.Type {
	case tp.Instruction_BLOCK:
		for _, child := range ins.Children {
			returnValue = ctx.RunInstruction(scope, child)
		}
	case tp.Instruction_TEXT:
		returnValue = null.GetString(ins.Value)
	case tp.Instruction_LOCAL_VAR:
		name := null.GetString(ins.Value)
		vars := ctx.Vars()
		if len(ins.Arguments) > 0 {
			vars[name] = ctx.RunInstruction(scope, ins.Arguments[0])
		}
		if len(ins.Children) > 0 {
			ts := &Scope{Value: ctx.Vars()[name]}
			for _, child := range ins.Children {
				ctx.RunInstruction(ts, child)
			}
			vars[name] = ts.Value
		}
		returnValue = vars[name]
	case tp.Instruction_IMPORT:
		obj := ctx.Objects[int(null.GetInt32(ins.ObjectId))]
		curFile := ctx.Filename
		ctx.Filename = null.GetString(obj.Name)
		start := time.Now()
		index := ctx.AddLog("__IMPORT__FILE__:" + ctx.Filename)
		for _, child := range obj.Root.Children {
			ctx.RunInstruction(scope, child)
		}
		timeSpent := time.Since(start).Seconds()
		ctx.UpdateLog(index, "__IMPORT__FILE__:" + strconv.Itoa(int(timeSpent*1000)) + ":" + ctx.Filename)
		ctx.Filename = curFile
	case tp.Instruction_FUNCTION_CALL:
		fun := ctx.Functions[int(null.GetInt32(ins.FunctionId))]
		args := make([]interface{}, len(ins.Arguments))
		for i, argIns := range ins.Arguments {
			args[i] = ctx.RunInstruction(scope, argIns)
		}

		if null.GetBool(fun.BuiltIn) {
			if f := builtInFunctions[fun.Name]; f != nil {
				returnValue = f(ctx, scope, ins, args)
				if returnValue == nil {
					returnValue = ""
				}
			} else {
				panic("missing function: " + fun.Name)
			}
		} else {
			// We are using a user-defined function
			//println("Resetting localvar")
			// Setup the new local var
			vars := make(map[string]interface{}, len(args))
			for i, arg := range fun.Args {
				vars[null.GetString(arg.Name)] = args[i]
			}
			yieldBlock := &YieldBlock{
				Ins:  ins,
				Vars: vars,
			}
			// PUSH!
			ctx.PushYieldBlock(yieldBlock)
			for _, child := range fun.Instruction.Children {
				returnValue = ctx.RunInstruction(scope, child)
			}
			// POP!
			ctx.PopYieldBlock()
		}
	}

	return
}

func (ctx *EngineContext) ShouldContinue() (result bool) {
	if len(ctx.MatchShouldContinueStack) > 0 {
		result = ctx.MatchShouldContinueStack[len(ctx.MatchShouldContinueStack)-1]
	} else {
		result = false
	}
	return
}

func (ctx *EngineContext) MatchTarget() string {
	return ctx.MatchStack[len(ctx.MatchStack)-1]
}

func (ctx *EngineContext) PushYieldBlock(b *YieldBlock) {
	ctx.Yields = append(ctx.Yields, b)
}

func (ctx *EngineContext) PopYieldBlock() (b *YieldBlock) {
	num := len(ctx.Yields)
	if num > 0 {
		b = ctx.Yields[num-1]
		ctx.Yields = ctx.Yields[:num-1]
	}
	return
}

func (ctx *EngineContext) HasYieldBlock() bool {
	return len(ctx.Yields) > 0
}

func (ctx *EngineContext) TopYieldBlock() (b *YieldBlock) {
	num := len(ctx.Yields)
	if num > 0 {
		b = ctx.Yields[num-1]
	}
	return
}

func (ctx *EngineContext) Vars() map[string]interface{} {
	b := ctx.TopYieldBlock()
	if b != nil {
		return b.Vars
	}
	return nil
}

func (ctx *EngineContext) FileAndLine(ins *tp.Instruction) string {
	lineNum := strconv.Itoa(int(null.GetInt32(ins.LineNumber)))
	return (ctx.Filename + ":" + lineNum)
}

func (ctx *EngineContext) UsePackage(pkg *tp.Package) {
	ctx.Types = make([]string, len(pkg.Types))
	for i, t := range pkg.Types {
		ctx.Types[i] = null.GetString(t.Name)
	}

	ctx.Functions = make([]*Function, len(pkg.Functions))
	for i, f := range pkg.Functions {
		name := null.GetString(f.Name)
		for _, a := range f.Args {
			typeString := ctx.Types[int(null.GetInt32(a.TypeId))]
			name = name + "." + typeString
		}
		fun := &Function{
			Name:     name,
			Function: f,
		}
		ctx.Functions[i] = fun
	}
}

func (ctx *EngineContext) GetRegexp(pattern, options string) (r *rubex.Regexp) {
	sig := pattern + "/" + options
	object, err := ctx.RegexpCache.Get(sig)
	if err != nil {
		mode := rubex.ONIG_OPTION_DEFAULT
		if strings.Index(options, "i") >= 0 {
			mode = rubex.ONIG_OPTION_IGNORECASE
		}
		if strings.Index(options, "m") >= 0 {
			mode = rubex.ONIG_OPTION_MULTILINE
		}
		var err error
		r, err = rubex.NewRegexp(pattern, mode)
		if err == nil {
			//ctx.AddMemoryObject(r)
			ctx.RegexpCache.Set(sig, &RegexpObject{Regexp:r})
		}
		return r
	}
	return object.(*RegexpObject).Regexp
}

func (ctx *EngineContext) GetXpathExpr(p string) (e *xpath.Expression) {
	object, err := ctx.XPathCache.Get(p)
	if err != nil {
		e = xpath.Compile(p)
		if e != nil {
			//ctx.AddMemoryObject(e)
			ctx.XPathCache.Set(p, &XpathExpObject{Expression: e})
		} else {
			ctx.AddLog("Invalid XPath used: " + p)
		}
		return e
	}
	return object.(*XpathExpObject).Expression
}

func (ctx *EngineContext) AddExport(exports []string) {
	ctx.Exports = append(ctx.Exports, exports)
}

func (ctx *EngineContext) AddLog(log string) int {
	//ctx.Log.Info("TRITIUM: " + log)
	index := len(ctx.Logs)
	ctx.Logs = append(ctx.Logs, log)
	return index
}

func (ctx *EngineContext) UpdateLog(index int, log string) {
	//ctx.Log.Info("TRITIUM: " + log)
	if index >= 0 && index < len(ctx.Logs) {
		ctx.Logs[index] = log
	}
}

func (ctx *EngineContext) SetEnv(key, val string) {
	ctx.Env[key] = val
}

func (ctx *EngineContext) GetEnv(key string) (val string) {
	val = ctx.Env[key]
	return
}

func (ctx *EngineContext) SetVar(key string, val interface{}) {
	b := ctx.TopYieldBlock()
	if b != nil {
		b.Vars[key] = val
	}
}

func (ctx *EngineContext) GetVar(key string) (val interface{}) {
	b := ctx.TopYieldBlock()
	if b != nil {
		val = b.Vars[key]
	}
	return
}

func (ctx *EngineContext) GetInnerReplacer() (r *rubex.Regexp) {
	return ctx.GetRegexp(`[\\$](\d)`, "i")
}

func (ctx *EngineContext) GetHeaderContentTypeRegex() (r *rubex.Regexp) {
	return ctx.GetRegexp(`<meta\s+http-equiv="content-type"\s+content="(.*?)"`, "i")
}

func (ctx *EngineContext) GetOutputBuffer() (b []byte) {
	//b = ctx.OutputBuffer
	return
}

func (ctx *EngineContext) Logger() (logger *golog.Logger) {
	logger = ctx.Log
	return
}

func (ctx *EngineContext) PushMatchStack(match string) {
	ctx.MatchStack = append(ctx.MatchStack, match)
}

func (ctx *EngineContext) PopMatchStack() (match string) {
	if num := len(ctx.MatchStack); num > 0 {
		match = ctx.MatchStack[num-1]
		ctx.MatchStack = ctx.MatchStack[:num-1]
	}
	return
}

func (ctx *EngineContext) PushShouldContinueStack(cont bool) {
	ctx.MatchShouldContinueStack = append(ctx.MatchShouldContinueStack, cont)
}

func (ctx *EngineContext) PopShouldContinueStack() (cont bool) {
	if num := len(ctx.MatchShouldContinueStack); num > 0 {
		cont = ctx.MatchShouldContinueStack[num-1]
		ctx.MatchShouldContinueStack = ctx.MatchShouldContinueStack[:num-1]
	}
	return
}

func (ctx *EngineContext) SetShouldContinue(cont bool) {
	if num := len(ctx.MatchShouldContinueStack); num > 0 {
		ctx.MatchShouldContinueStack[num-1] = cont
	}
}

func (ctx *EngineContext) AddMemoryObject(o MemoryObject) {
	ctx.Mobjects = append(ctx.Mobjects, o)
}
<commit_msg>pass env to trap<commit_after>package whale

import (
	"strings"
	"time"
	"strconv"
)

import (
	"butler/null"
	"gokogiri/xpath"
	"golog"
	"rubex"
	tp "tritium/proto"
	"steno"
	"go-cache"
	"go-cache/arc"
)

type Whale struct {
	Log *golog.Logger
	Debugger steno.Debugger
	RegexpCache       cache.Cache
	XPathCache        cache.Cache
}

type EngineContext struct {
	Functions                []*Function
	Types                    []string
	Exports                  [][]string
	Logs                     []string
	ImportedFiles            []string
	Env                      map[string]string
	MatchStack               []string
	MatchShouldContinueStack []bool
	Yields                   []*YieldBlock
	*Whale
	*tp.Transform
	Rrules []*tp.RewriteRule

	InnerReplacer     *rubex.Regexp
	HeaderContentType *rubex.Regexp

	// Debug info
	Filename          string
	HadError          bool
	Deadline time.Time
	Mobjects []MemoryObject
	MessagePath string
}

const OutputBufferSize = 500 * 1024 //500KB
const defaultMobjects = 4
const TimeoutError = "EngineTimeout"

func NewEngine(logger *golog.Logger, debugger steno.Debugger) *Whale {
	e := &Whale{
		Log: logger,
		Debugger: debugger,
		RegexpCache:              arc.NewARCache(1000),
		XPathCache:               arc.NewARCache(1000),
	}
	e.RegexpCache.SetCleanFunc(CleanRegexpObject)
	e.XPathCache.SetCleanFunc(CleanXpathExpObject)
	return e
}

func NewEngineCtx(eng *Whale, vars map[string]string, transform *tp.Transform, rrules []*tp.RewriteRule, deadline time.Time, messagePath string) (ctx *EngineContext) {
	ctx = &EngineContext{
		Whale:                    eng,
		Exports:                  make([][]string, 0),
		Logs:                     make([]string, 0),
		Env:                      vars,
		Transform:                transform,
		Rrules:                   rrules,
		MatchStack:               make([]string, 0),
		MatchShouldContinueStack: make([]bool, 0),
		Yields:                   make([]*YieldBlock, 0),
		HadError:                 false,

		Deadline:                 deadline,
		Mobjects:                 make([]MemoryObject, 0, defaultMobjects),
		MessagePath:              messagePath,
	}
	return
}

func (eng *Whale) Free() {
	eng.RegexpCache.Reset()
	eng.XPathCache.Reset()
}

func (eng *Whale) Run(transform *tp.Transform, rrules []*tp.RewriteRule, input interface{}, vars map[string]string, deadline time.Time, customer, project, messagePath string) (output string, exports [][]string, logs []string) {
	ctx := NewEngineCtx(eng, vars, transform, rrules, deadline, messagePath)
	defer ctx.Free()
	ctx.Yields = append(ctx.Yields, &YieldBlock{Vars: make(map[string]interface{})})
	ctx.UsePackage(transform.Pkg)
	scope := &Scope{Value: input.(string)}
	obj := transform.Objects[0]
	ctx.Filename = null.GetString(obj.Name)
	ctx.RunInstruction(scope, obj.Root)
	output = scope.Value.(string)
	exports = ctx.Exports
	logs = ctx.Logs
	return
}

func (eng *Whale) GetCacheStats() (int, int, int, int) {
	return eng.RegexpCache.GetHitRate(), eng.RegexpCache.GetUsageRate(), eng.XPathCache.GetHitRate(), eng.XPathCache.GetUsageRate()
}

func (ctx *EngineContext) Free() {
    for _, o := range ctx.Mobjects {
		if o != nil {
			o.Free()
		}
	}
}

func (ctx *EngineContext) RunInstruction(scope *Scope, ins *tp.Instruction) (returnValue interface{}) {
	defer func() {
		//TODO Stack traces seem to get truncated on syslog...
		if x := recover(); x != nil {
			err, ok := x.(error)
			errString := ""
			if ok {
				errString = err.Error()
			} else {
				errString = x.(string)
			}
			if errString != TimeoutError {
				if ctx.HadError == false {
					ctx.HadError = true
					errString = errString + "\n" + ins.Type.String() + " " + null.GetString(ins.Value) + "\n\n\nTritium Stack\n=========\n\n"
				}
				errString = errString + ctx.FileAndLine(ins) + "\n"
			}
			panic(errString)
		}
	}()

	ctx.Whale.Debugger.TrapInstruction(ctx.Filename, ctx.Env, ins, scope.Value, scope.Index)

	if time.Now().After(ctx.Deadline) {
		panic(TimeoutError)
	}

	// If our object is invalid, then skip it
	if null.GetBool(ins.IsValid) == false {
		panic("Invalid instruction. Should have stopped before linking!")
	}
	indent := ""
	for i := 0; i < len(ctx.Yields); i++ {
		indent += "\t"
	}

	returnValue = ""
	switch *ins.Type {
	case tp.Instruction_BLOCK:
		for _, child := range ins.Children {
			returnValue = ctx.RunInstruction(scope, child)
		}
	case tp.Instruction_TEXT:
		returnValue = null.GetString(ins.Value)
	case tp.Instruction_LOCAL_VAR:
		name := null.GetString(ins.Value)
		vars := ctx.Vars()
		if len(ins.Arguments) > 0 {
			vars[name] = ctx.RunInstruction(scope, ins.Arguments[0])
		}
		if len(ins.Children) > 0 {
			ts := &Scope{Value: ctx.Vars()[name]}
			for _, child := range ins.Children {
				ctx.RunInstruction(ts, child)
			}
			vars[name] = ts.Value
		}
		returnValue = vars[name]
	case tp.Instruction_IMPORT:
		obj := ctx.Objects[int(null.GetInt32(ins.ObjectId))]
		curFile := ctx.Filename
		ctx.Filename = null.GetString(obj.Name)
		start := time.Now()
		index := ctx.AddLog("__IMPORT__FILE__:" + ctx.Filename)
		for _, child := range obj.Root.Children {
			ctx.RunInstruction(scope, child)
		}
		timeSpent := time.Since(start).Seconds()
		ctx.UpdateLog(index, "__IMPORT__FILE__:" + strconv.Itoa(int(timeSpent*1000)) + ":" + ctx.Filename)
		ctx.Filename = curFile
	case tp.Instruction_FUNCTION_CALL:
		fun := ctx.Functions[int(null.GetInt32(ins.FunctionId))]
		args := make([]interface{}, len(ins.Arguments))
		for i, argIns := range ins.Arguments {
			args[i] = ctx.RunInstruction(scope, argIns)
		}

		if null.GetBool(fun.BuiltIn) {
			if f := builtInFunctions[fun.Name]; f != nil {
				returnValue = f(ctx, scope, ins, args)
				if returnValue == nil {
					returnValue = ""
				}
			} else {
				panic("missing function: " + fun.Name)
			}
		} else {
			// We are using a user-defined function
			//println("Resetting localvar")
			// Setup the new local var
			vars := make(map[string]interface{}, len(args))
			for i, arg := range fun.Args {
				vars[null.GetString(arg.Name)] = args[i]
			}
			yieldBlock := &YieldBlock{
				Ins:  ins,
				Vars: vars,
			}
			// PUSH!
			ctx.PushYieldBlock(yieldBlock)
			for _, child := range fun.Instruction.Children {
				returnValue = ctx.RunInstruction(scope, child)
			}
			// POP!
			ctx.PopYieldBlock()
		}
	}

	return
}

func (ctx *EngineContext) ShouldContinue() (result bool) {
	if len(ctx.MatchShouldContinueStack) > 0 {
		result = ctx.MatchShouldContinueStack[len(ctx.MatchShouldContinueStack)-1]
	} else {
		result = false
	}
	return
}

func (ctx *EngineContext) MatchTarget() string {
	return ctx.MatchStack[len(ctx.MatchStack)-1]
}

func (ctx *EngineContext) PushYieldBlock(b *YieldBlock) {
	ctx.Yields = append(ctx.Yields, b)
}

func (ctx *EngineContext) PopYieldBlock() (b *YieldBlock) {
	num := len(ctx.Yields)
	if num > 0 {
		b = ctx.Yields[num-1]
		ctx.Yields = ctx.Yields[:num-1]
	}
	return
}

func (ctx *EngineContext) HasYieldBlock() bool {
	return len(ctx.Yields) > 0
}

func (ctx *EngineContext) TopYieldBlock() (b *YieldBlock) {
	num := len(ctx.Yields)
	if num > 0 {
		b = ctx.Yields[num-1]
	}
	return
}

func (ctx *EngineContext) Vars() map[string]interface{} {
	b := ctx.TopYieldBlock()
	if b != nil {
		return b.Vars
	}
	return nil
}

func (ctx *EngineContext) FileAndLine(ins *tp.Instruction) string {
	lineNum := strconv.Itoa(int(null.GetInt32(ins.LineNumber)))
	return (ctx.Filename + ":" + lineNum)
}

func (ctx *EngineContext) UsePackage(pkg *tp.Package) {
	ctx.Types = make([]string, len(pkg.Types))
	for i, t := range pkg.Types {
		ctx.Types[i] = null.GetString(t.Name)
	}

	ctx.Functions = make([]*Function, len(pkg.Functions))
	for i, f := range pkg.Functions {
		name := null.GetString(f.Name)
		for _, a := range f.Args {
			typeString := ctx.Types[int(null.GetInt32(a.TypeId))]
			name = name + "." + typeString
		}
		fun := &Function{
			Name:     name,
			Function: f,
		}
		ctx.Functions[i] = fun
	}
}

func (ctx *EngineContext) GetRegexp(pattern, options string) (r *rubex.Regexp) {
	sig := pattern + "/" + options
	object, err := ctx.RegexpCache.Get(sig)
	if err != nil {
		mode := rubex.ONIG_OPTION_DEFAULT
		if strings.Index(options, "i") >= 0 {
			mode = rubex.ONIG_OPTION_IGNORECASE
		}
		if strings.Index(options, "m") >= 0 {
			mode = rubex.ONIG_OPTION_MULTILINE
		}
		var err error
		r, err = rubex.NewRegexp(pattern, mode)
		if err == nil {
			//ctx.AddMemoryObject(r)
			ctx.RegexpCache.Set(sig, &RegexpObject{Regexp:r})
		}
		return r
	}
	return object.(*RegexpObject).Regexp
}

func (ctx *EngineContext) GetXpathExpr(p string) (e *xpath.Expression) {
	object, err := ctx.XPathCache.Get(p)
	if err != nil {
		e = xpath.Compile(p)
		if e != nil {
			//ctx.AddMemoryObject(e)
			ctx.XPathCache.Set(p, &XpathExpObject{Expression: e})
		} else {
			ctx.AddLog("Invalid XPath used: " + p)
		}
		return e
	}
	return object.(*XpathExpObject).Expression
}

func (ctx *EngineContext) AddExport(exports []string) {
	ctx.Exports = append(ctx.Exports, exports)
}

func (ctx *EngineContext) AddLog(log string) int {
	//ctx.Log.Info("TRITIUM: " + log)
	index := len(ctx.Logs)
	ctx.Logs = append(ctx.Logs, log)
	return index
}

func (ctx *EngineContext) UpdateLog(index int, log string) {
	//ctx.Log.Info("TRITIUM: " + log)
	if index >= 0 && index < len(ctx.Logs) {
		ctx.Logs[index] = log
	}
}

func (ctx *EngineContext) SetEnv(key, val string) {
	ctx.Env[key] = val
}

func (ctx *EngineContext) GetEnv(key string) (val string) {
	val = ctx.Env[key]
	return
}

func (ctx *EngineContext) SetVar(key string, val interface{}) {
	b := ctx.TopYieldBlock()
	if b != nil {
		b.Vars[key] = val
	}
}

func (ctx *EngineContext) GetVar(key string) (val interface{}) {
	b := ctx.TopYieldBlock()
	if b != nil {
		val = b.Vars[key]
	}
	return
}

func (ctx *EngineContext) GetInnerReplacer() (r *rubex.Regexp) {
	return ctx.GetRegexp(`[\\$](\d)`, "i")
}

func (ctx *EngineContext) GetHeaderContentTypeRegex() (r *rubex.Regexp) {
	return ctx.GetRegexp(`<meta\s+http-equiv="content-type"\s+content="(.*?)"`, "i")
}

func (ctx *EngineContext) GetOutputBuffer() (b []byte) {
	//b = ctx.OutputBuffer
	return
}

func (ctx *EngineContext) Logger() (logger *golog.Logger) {
	logger = ctx.Log
	return
}

func (ctx *EngineContext) PushMatchStack(match string) {
	ctx.MatchStack = append(ctx.MatchStack, match)
}

func (ctx *EngineContext) PopMatchStack() (match string) {
	if num := len(ctx.MatchStack); num > 0 {
		match = ctx.MatchStack[num-1]
		ctx.MatchStack = ctx.MatchStack[:num-1]
	}
	return
}

func (ctx *EngineContext) PushShouldContinueStack(cont bool) {
	ctx.MatchShouldContinueStack = append(ctx.MatchShouldContinueStack, cont)
}

func (ctx *EngineContext) PopShouldContinueStack() (cont bool) {
	if num := len(ctx.MatchShouldContinueStack); num > 0 {
		cont = ctx.MatchShouldContinueStack[num-1]
		ctx.MatchShouldContinueStack = ctx.MatchShouldContinueStack[:num-1]
	}
	return
}

func (ctx *EngineContext) SetShouldContinue(cont bool) {
	if num := len(ctx.MatchShouldContinueStack); num > 0 {
		ctx.MatchShouldContinueStack[num-1] = cont
	}
}

func (ctx *EngineContext) AddMemoryObject(o MemoryObject) {
	ctx.Mobjects = append(ctx.Mobjects, o)
}
<|endoftext|>


<commit_before>/*
Package httprs provides a ReadSeeker for http.Response.Body.

Usage :

	resp, err := http.Get(url)
	rs := httprs.NewHttpReadSeeker(resp)
	defer rs.Close()
	io.ReadFull(rs, buf) // reads the first bytes from the response body
	rs.Seek(1024, 0) // moves the position, but does no range request
	io.ReadFull(rs, buf) // does a range request and reads from the response body

If you want use a specific http.Client for additional range requests :
	rs := httprs.NewHttpReadSeeker(resp, client)
*/
package httprs

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

	"github.com/mitchellh/copystructure"
)

// A HttpReadSeeker reads from a http.Response.Body. It can Seek
// by doing range requests.
type HttpReadSeeker struct {
	c       *http.Client
	req     *http.Request
	res     *http.Response
	r       io.ReadCloser
	pos     int64
	canSeek bool
}

var _ io.ReadCloser = (*HttpReadSeeker)(nil)
var _ io.Seeker = (*HttpReadSeeker)(nil)

var (
	// ErrNoContentLength is returned by Seek when the initial http response did not include a Content-Length header
	ErrNoContentLength = errors.New("Content-Length was not set")
	// ErrRangeRequestsNotSupported is returned by Seek and Read
	// when the remote server does not allow range requests (Accept-Ranges was not set)
	ErrRangeRequestsNotSupported = errors.New("Range requests are not supported by the remote server")
	// ErrInvalidRange is returned by Read when trying to read past the end of the file
	ErrInvalidRange = errors.New("Invalid range")
	// ErrContentHasChanged is returned by Read when the content has changed since the first request
	ErrContentHasChanged = errors.New("Content has changed since first request")
)

// NewHttpReadSeeker returns a HttpReadSeeker, using the http.Response and, optionaly, the http.Client
// that needs to be used for future range requests. If no http.Client is given, http.DefaultClient will
// be used.
//
// res.Request will be reused for range requests, headers may be added/removed
func NewHttpReadSeeker(res *http.Response, client ...*http.Client) *HttpReadSeeker {
	r := &HttpReadSeeker{
		req:     res.Request,
		res:     res,
		r:       res.Body,
		canSeek: (res.Header.Get("Accept-Ranges") == "bytes"),
	}
	if len(client) > 0 {
		r.c = client[0]
	} else {
		r.c = http.DefaultClient
	}
	return r
}

// Clone clones the reader to enable parallel downloads of ranges
func (r *HttpReadSeeker) Clone() (*HttpReadSeeker, error) {
	req, err := copystructure.Copy(r.req)
	if err != nil {
		return nil, err
	}
	return &HttpReadSeeker{
		req:     req.(*http.Request),
		res:     r.res,
		r:       nil,
		canSeek: r.canSeek,
		c:       r.c,
	}, nil
}

// Read reads from the response body. It does a range request if Seek was called before.
//
// May return ErrRangeRequestsNotSupported, ErrInvalidRange or ErrContentHasChanged
func (r *HttpReadSeeker) Read(p []byte) (n int, err error) {
	if r.r == nil {
		err = r.rangeRequest()
	}
	if r.r != nil {
		n, err = r.r.Read(p)
		r.pos += int64(n)
	}
	return
}

// ReadAt reads from the response body starting at offset off.
//
// May return ErrRangeRequestsNotSupported, ErrInvalidRange or ErrContentHasChanged
func (r *HttpReadSeeker) ReadAt(p []byte, off int64) (n int, err error) {
	r.Seek(off, 0)
	return r.Read(p)
}

// Close closes the response body
func (r *HttpReadSeeker) Close() error {
	if r.r != nil {
		return r.r.Close()
	}
	return nil
}

// Seek moves the reader position to a new offset.
//
// It does not send http requests, allowing for multiple seeks without overhead.
// The http request will be sent by the next Read call.
//
// May return ErrNoContentLength or ErrRangeRequestsNotSupported
func (r *HttpReadSeeker) Seek(offset int64, whence int) (int64, error) {
	if !r.canSeek {
		return 0, ErrRangeRequestsNotSupported
	}
	var err error
	switch whence {
	case 0:
	case 1:
		offset += r.pos
	case 2:
		if r.res.ContentLength <= 0 {
			return 0, ErrNoContentLength
		}
		offset = r.res.ContentLength - offset
	}
	if r.r != nil && r.pos != offset {
		err = r.r.Close()
		r.r = nil
	}
	r.pos = offset
	return r.pos, err
}

func (r *HttpReadSeeker) rangeRequest() error {
	r.req.Header.Set("Range", fmt.Sprintf("bytes=%d-", r.pos))
	etag, last := r.res.Header.Get("ETag"), r.res.Header.Get("Last-Modified")
	switch {
	case last != "":
		r.req.Header.Set("If-Range", last)
	case etag != "":
		r.req.Header.Set("If-Range", etag)
	}

	res, err := r.c.Do(r.req)
	if err != nil {
		return err
	}
	switch res.StatusCode {
	case http.StatusRequestedRangeNotSatisfiable:
		return ErrInvalidRange
	case http.StatusOK:
		return ErrContentHasChanged
	case http.StatusPartialContent:
		r.r = res.Body
		return nil
	}
	return ErrRangeRequestsNotSupported
}
<commit_msg># This is a combination of 3 commits. # The first commit's message is: Read, instead of seeking<commit_after>/*
Package httprs provides a ReadSeeker for http.Response.Body.

Usage :

	resp, err := http.Get(url)
	rs := httprs.NewHttpReadSeeker(resp)
	defer rs.Close()
	io.ReadFull(rs, buf) // reads the first bytes from the response body
	rs.Seek(1024, 0) // moves the position, but does no range request
	io.ReadFull(rs, buf) // does a range request and reads from the response body

If you want use a specific http.Client for additional range requests :
	rs := httprs.NewHttpReadSeeker(resp, client)
*/
package httprs

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

	"github.com/mitchellh/copystructure"
)

const seekReadInsteadOfRequest = 1024

// A HttpReadSeeker reads from a http.Response.Body. It can Seek
// by doing range requests.
type HttpReadSeeker struct {
	c       *http.Client
	req     *http.Request
	res     *http.Response
	r       io.ReadCloser
	pos     int64
	canSeek bool
}

var _ io.ReadCloser = (*HttpReadSeeker)(nil)
var _ io.Seeker = (*HttpReadSeeker)(nil)

var (
	// ErrNoContentLength is returned by Seek when the initial http response did not include a Content-Length header
	ErrNoContentLength = errors.New("Content-Length was not set")
	// ErrRangeRequestsNotSupported is returned by Seek and Read
	// when the remote server does not allow range requests (Accept-Ranges was not set)
	ErrRangeRequestsNotSupported = errors.New("Range requests are not supported by the remote server")
	// ErrInvalidRange is returned by Read when trying to read past the end of the file
	ErrInvalidRange = errors.New("Invalid range")
	// ErrContentHasChanged is returned by Read when the content has changed since the first request
	ErrContentHasChanged = errors.New("Content has changed since first request")
)

// NewHttpReadSeeker returns a HttpReadSeeker, using the http.Response and, optionaly, the http.Client
// that needs to be used for future range requests. If no http.Client is given, http.DefaultClient will
// be used.
//
// res.Request will be reused for range requests, headers may be added/removed
func NewHttpReadSeeker(res *http.Response, client ...*http.Client) *HttpReadSeeker {
	r := &HttpReadSeeker{
		req:     res.Request,
		res:     res,
		r:       res.Body,
		canSeek: (res.Header.Get("Accept-Ranges") == "bytes"),
	}
	if len(client) > 0 {
		r.c = client[0]
	} else {
		r.c = http.DefaultClient
	}
	return r
}

// Clone clones the reader to enable parallel downloads of ranges
func (r *HttpReadSeeker) Clone() (*HttpReadSeeker, error) {
	req, err := copystructure.Copy(r.req)
	if err != nil {
		return nil, err
	}
	return &HttpReadSeeker{
		req:     req.(*http.Request),
		res:     r.res,
		r:       nil,
		canSeek: r.canSeek,
		c:       r.c,
	}, nil
}

// Read reads from the response body. It does a range request if Seek was called before.
//
// May return ErrRangeRequestsNotSupported, ErrInvalidRange or ErrContentHasChanged
func (r *HttpReadSeeker) Read(p []byte) (n int, err error) {
	if r.r == nil {
		err = r.rangeRequest()
	}
	if r.r != nil {
		n, err = r.r.Read(p)
		r.pos += int64(n)
	}
	return
}

// ReadAt reads from the response body starting at offset off.
//
// May return ErrRangeRequestsNotSupported, ErrInvalidRange or ErrContentHasChanged
func (r *HttpReadSeeker) ReadAt(p []byte, off int64) (n int, err error) {
	r.Seek(off, 0)
	return r.Read(p)
}

// Close closes the response body
func (r *HttpReadSeeker) Close() error {
	if r.r != nil {
		return r.r.Close()
	}
	return nil
}

// Seek moves the reader position to a new offset.
//
// It does not send http requests, allowing for multiple seeks without overhead.
// The http request will be sent by the next Read call.
//
// May return ErrNoContentLength or ErrRangeRequestsNotSupported
func (r *HttpReadSeeker) Seek(offset int64, whence int) (int64, error) {
	if !r.canSeek {
		return 0, ErrRangeRequestsNotSupported
	}
	var err error
	switch whence {
	case 0:
	case 1:
		offset += r.pos
	case 2:
		if r.res.ContentLength <= 0 {
			return 0, ErrNoContentLength
		}
		offset = r.res.ContentLength - offset
	}
	if r.r != nil {
		// Try to read, which is cheaper than doing a request
		if r.pos < offset && offset-r.pos <= seekReadInsteadOfRequest {
			_, err := io.CopyN(ioutil.Discard, r, offset-r.pos)
			if err != nil {
				return 0, err
			}
		}

		if r.pos != offset {
			err = r.r.Close()
			r.r = nil
		}
	}
	r.pos = offset
	return r.pos, err
}

func (r *HttpReadSeeker) rangeRequest() error {
	r.req.Header.Set("Range", fmt.Sprintf("bytes=%d-", r.pos))
	etag, last := r.res.Header.Get("ETag"), r.res.Header.Get("Last-Modified")
	switch {
	case last != "":
		r.req.Header.Set("If-Range", last)
	case etag != "":
		r.req.Header.Set("If-Range", etag)
	}

	res, err := r.c.Do(r.req)
	if err != nil {
		return err
	}
	switch res.StatusCode {
	case http.StatusRequestedRangeNotSatisfiable:
		return ErrInvalidRange
	case http.StatusOK:
		return ErrContentHasChanged
	case http.StatusPartialContent:
		r.r = res.Body
		return nil
	}
	return ErrRangeRequestsNotSupported
}
<|endoftext|>