<commit_before>package archive

import (
	"fmt"
	"io"
	"path"
	"time"

	"crypto"
	"crypto/sha512"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/clearsign"
	"golang.org/x/crypto/openpgp/packet"

	"pault.ag/go/blobstore"
	"pault.ag/go/debian/control"
	"pault.ag/go/debian/dependency"
	"pault.ag/go/debian/transput"
)

// Archive {{{

// Core Archive abstrcation. This contains helpers to write out package files,
// as well as handles creating underlying abstractions (such as Suites).
type Archive struct {
	store      blobstore.Store
	signingKey *openpgp.Entity
}

// Create a new Archive at the given `root` on the filesystem, with the
// openpgp.Entity `signer` (an Entity which contains an OpenPGP Private
// Key).
//
// This interface is intended to *write* Archives, not *read* them. Extra
// steps must be taken to load an Archive over the network, and attention
// must be paid when handling the Cryptographic chain of trust.
func New(path string, signer *openpgp.Entity) (*Archive, error) {
	store, err := blobstore.Load(path)
	if err != nil {
		return nil, err
	}

	return &Archive{
		store:      *store,
		signingKey: signer,
	}, nil
}

// Use the default backend to remove any unlinked files from the Blob store.
//
// If files you care about are not linked onto the stage, they will be removed
// by the garbage collector. Decruft only when you're sure the stage has been
// set.
func (a Archive) Decruft() error {
	return a.store.GC(blobstore.DumbGarbageCollector{})
}

// Get the current state of files in the Archive, which is an underlying
// call to the blobstore Paths function.
func (a Archive) Paths() (ArchiveState, error) {
	return a.store.Paths()
}

func (a Archive) Open(o blobstore.Object) (io.ReadCloser, error) {
	return a.store.Open(o)
}

// Given a list of objects, link them to the keyed paths.
func (a Archive) Link(blobs ArchiveState) error {
	for path, obj := range blobs {
		if err := a.store.Link(obj, path); err != nil {
			return err
		}
	}
	return nil
}

// Create a new Release object from a Suite, passing off the Name, Description
// and constructing the rest of the goodies.
//
// This will be an entirely empty object, without anything read off disk.
func newRelease(suite Suite) (*Release, error) {
	when := time.Now()

	var validUntil string = ""
	if suite.features.Duration != "" {
		duration, err := time.ParseDuration(suite.features.Duration)
		if err != nil {
			return nil, err
		}
		validUntil = when.Add(duration).Format(time.RFC1123Z)
	}

	release := Release{
		Suite:       suite.Name,
		Description: suite.Description,
		ValidUntil:  validUntil,
		Origin:      suite.Origin,
		Label:       suite.Label,
		Version:     suite.Version,
	}
	release.Date = when.Format(time.RFC1123Z)
	release.Architectures = []dependency.Arch{}
	release.Components = []string{}
	release.SHA256 = []control.SHA256FileHash{}
	release.SHA1 = []control.SHA1FileHash{}
	release.SHA512 = []control.SHA512FileHash{}
	release.MD5Sum = []control.MD5FileHash{}
	return &release, nil
}

// This is a set of file changes ready to be passed to `Link` to link in.
// Basically, this maps file paths to blobstore objects, which will be
// swapped in all at once. This allows errors to avoid mutating state in
// the archive.
type ArchiveState map[string]blobstore.Object

// Engross a Suite for signing and final commit into the blobstore. This
// will return handle(s) to the signed and ready Objects, fit for passage
// to Link.
//
// This will contain all the related Packages and Release files.
func (a Archive) Engross(suite Suite) (ArchiveState, error) {
	release, err := newRelease(suite)
	if err != nil {
		return nil, err
	}

	files := ArchiveState{}
	arches := map[dependency.Arch]bool{}

	for name, component := range suite.components {
		release.Components = append(release.Components, name)
		for arch, writer := range component.packageWriters {
			arches[arch] = true

			// For each Binary entry, do the same as above (todo: someone
			// DRY this out a bit. I'm too lazy.

			suitePath := path.Join(name, fmt.Sprintf("binary-%s", arch),
				"Packages")

			obj, err := a.store.Commit(*writer.handle)
			if err != nil {
				return nil, err
			}

			for _, hasher := range writer.hashers {
				fileHash := control.FileHashFromHasher(suitePath, *hasher)
				release.AddHash(fileHash)
			}

			filePath := path.Join("dists", suite.Name, suitePath)
			files[filePath] = *obj
		}
	}

	for arch, _ := range arches {
		release.Architectures = append(release.Architectures, arch)
	}

	/* Now, let's do some magic */

	// Now, let's write out the Release file (and sign it normally)
	obj, sig, err := suite.archive.encodeSigned(release)
	if err != nil {
		return nil, err
	}

	filePath := path.Join("dists", suite.Name, "Release")
	files[filePath] = *obj
	files[fmt.Sprintf("%s.gpg", filePath)] = *sig

	// Ditto with the clearsigned version (Should we merge the two above?)
	obj, err = suite.archive.encodeClearsigned(release)
	if err != nil {
		return nil, err
	}

	files[path.Join("dists", suite.Name, "InRelease")] = *obj

	return files, nil
}

// Given a control.Marshal'able object, encode it to the blobstore, while
// also clearsigning the data.
func (a Archive) encodeClearsigned(data interface{}) (*blobstore.Object, error) {

	if a.signingKey == nil {
		return nil, fmt.Errorf("No signing key loaded")
	}

	fd, err := a.store.Create()
	if err != nil {
		return nil, err
	}

	defer fd.Close()
	wc, err := clearsign.Encode(fd, a.signingKey.PrivateKey, nil)
	if err != nil {
		return nil, err
	}

	encoder, err := control.NewEncoder(wc)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	if err := wc.Close(); err != nil {
		return nil, err
	}

	return a.store.Commit(*fd)
}

// Given a control.Marshal'able object, encode it to the blobstore, while
// also doing a detached OpenPGP signature. The objects returned (in order)
// are data, commited to the blobstore, the signature for that object, commited
// to the blobstore, and any error(s), finally.
func (a Archive) encodeSigned(data interface{}) (*blobstore.Object, *blobstore.Object, error) {
	/* Right, so, the trick here is that we secretly call out to encode,
	 * but tap it with a pipe into the signing code */

	if a.signingKey == nil {
		return nil, nil, fmt.Errorf("No signing key loaded")
	}

	signature, err := a.store.Create()
	if err != nil {
		return nil, nil, err
	}
	defer signature.Close()

	hash := sha512.New()

	obj, err := a.encode(data, hash)
	if err != nil {
		return nil, nil, err
	}

	sig := new(packet.Signature)
	sig.SigType = packet.SigTypeBinary
	sig.PubKeyAlgo = a.signingKey.PrivateKey.PubKeyAlgo

	sig.Hash = crypto.SHA512

	sig.CreationTime = new(packet.Config).Now()
	sig.IssuerKeyId = &(a.signingKey.PrivateKey.KeyId)

	err = sig.Sign(hash, a.signingKey.PrivateKey, nil)
	if err != nil {
		return nil, nil, err
	}

	if err := sig.Serialize(signature); err != nil {
		return nil, nil, err
	}

	sigObj, err := a.store.Commit(*signature)
	if err != nil {
		return nil, nil, err
	}

	return obj, sigObj, nil

}

// Encode a given control.Marshal'able object into the Blobstore, and return
// a handle to its object.
//
// The optinal argument `tap` will be written to as the object gets sent into
// the blobstore. This may be useful if you wish to have a copy of the data
// going into the store.
func (a Archive) encode(data interface{}, tap io.Writer) (*blobstore.Object, error) {
	fd, err := a.store.Create()
	if err != nil {
		return nil, err
	}

	var writer io.Writer = fd
	if tap != nil {
		writer = io.MultiWriter(fd, tap)
	}

	encoder, err := control.NewEncoder(writer)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	return a.store.Commit(*fd)
}

// }}}

// Suite {{{

// Abstraction to handle writing data into a Suite. This is a write-only
// target, and is not intended to read a Release file.
//
// This contains no state read off disk, and is purely for writing to.
type Suite struct {
	control.Paragraph

	archive *Archive

	Name        string `control:"Suite"`
	Description string
	Origin      string
	Label       string
	Version     string

	Pool       Pool                  `control:"-"`
	components map[string]*Component `control:"-"`

	features struct {
		Hashes   []string
		Duration string
	} `control:"-"`
}

// Get a handle to write a given Suite from an Archive.
// The suite will be entirely blank, and attributes will not be
// read from the existing files, if any.
func (a Archive) Suite(name string) (*Suite, error) {
	suite := Suite{
		Name:       name,
		archive:    &a,
		components: map[string]*Component{},
	}

	suite.Pool = Pool{store: a.store, suite: &suite}
	suite.features.Hashes = []string{"sha256", "sha1", "sha512"}
	suite.features.Duration = "240h"

	return &suite, nil
}

// Get or create a Component for a given Suite. If no such Component
// has been created so far, this will create a new object, otherwise
// it will return the existing entry.
//
// This contains no state read off disk, and is purely for writing to.
func (s Suite) Component(name string) (*Component, error) {
	if _, ok := s.components[name]; !ok {
		comp, err := newComponent(&s)
		if err != nil {
			return nil, err
		}
		s.components[name] = comp
		return comp, nil
	}
	el := s.components[name]
	return el, nil
}

// }}}

// Component {{{

// Small wrapper to represent a Component of a Suite, which, at its core
// is simply a set of Indexes to be written to.
//
// This contains no state read off disk, and is purely for writing to.
type Component struct {
	suite          *Suite
	packageWriters map[dependency.Arch]*IndexWriter
}

// Create a new Component, configured for use.
func newComponent(suite *Suite) (*Component, error) {
	return &Component{
		suite:          suite,
		packageWriters: map[dependency.Arch]*IndexWriter{},
	}, nil
}

// Get a given IndexWriter for an arch, or create one if none exists.
func (c *Component) getWriter(arch dependency.Arch) (*IndexWriter, error) {
	if _, ok := c.packageWriters[arch]; !ok {
		writer, err := newIndexWriter(c.suite)
		if err != nil {
			return nil, err
		}
		c.packageWriters[arch] = writer
	}
	return c.packageWriters[arch], nil
}

// Add a given Package to a Package List. Under the hood, this will
// get or create a IndexWriter, and invoke the .Add method on the
// Package Writer.
func (c *Component) AddPackage(pkg Package) error {
	writer, err := c.getWriter(pkg.Architecture)
	if err != nil {
		return err
	}
	return writer.Add(pkg)
}

// }}}

// IndexWriter {{{

// This writer represents a Package list - which is to say, a list of
// binary .deb files, for a particular Architecture, in a particular Component
// in a particular Suite, in a particular Archive.
//
// This is not an encapsulation to store the entire Index in memory, rather,
// it's a wrapper to help write Package entries into the Index.
type IndexWriter struct {
	archive *Archive

	handle  *blobstore.Writer
	closer  func() error
	encoder *control.Encoder

	hashers []*transput.Hasher
}

func getHashers(suite *Suite) (io.Writer, []*transput.Hasher, error) {
	ret := []*transput.Hasher{}
	writers := []io.Writer{}

	for _, algo := range suite.features.Hashes {
		hasher, err := transput.NewHasher(algo)
		if err != nil {
			return nil, nil, err
		}
		writers = append(writers, hasher)
		ret = append(ret, hasher)
	}

	return io.MultiWriter(writers...), ret, nil
}

// given a Suite, create a new Package Writer, configured with
// the appropriate Hashing, and targeting a new file blob in the
// underlying blobstore.
func newIndexWriter(suite *Suite) (*IndexWriter, error) {
	handle, err := suite.archive.store.Create()
	if err != nil {
		return nil, err
	}

	writer, hashers, err := getHashers(suite)
	if err != nil {
		return nil, err
	}

	encoder, err := control.NewEncoder(io.MultiWriter(writer, handle))
	if err != nil {
		handle.Close()
		return nil, err
	}

	return &IndexWriter{
		archive: suite.archive,
		closer:  handle.Close,
		encoder: encoder,
		handle:  handle,
		hashers: hashers,
	}, nil
}

// Write a Package entry into the Packages index.
func (p IndexWriter) Add(data interface{}) error {
	return p.encoder.Encode(data)
}

// }}}

// vim: foldmethod=marker
<commit_msg>add path<commit_after>package archive

import (
	"fmt"
	"io"
	"path"
	"time"

	"crypto"
	"crypto/sha512"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/clearsign"
	"golang.org/x/crypto/openpgp/packet"

	"pault.ag/go/blobstore"
	"pault.ag/go/debian/control"
	"pault.ag/go/debian/dependency"
	"pault.ag/go/debian/transput"
)

// Archive {{{

// Core Archive abstrcation. This contains helpers to write out package files,
// as well as handles creating underlying abstractions (such as Suites).
type Archive struct {
	store      blobstore.Store
	signingKey *openpgp.Entity
	path       string
}

// Create a new Archive at the given `root` on the filesystem, with the
// openpgp.Entity `signer` (an Entity which contains an OpenPGP Private
// Key).
//
// This interface is intended to *write* Archives, not *read* them. Extra
// steps must be taken to load an Archive over the network, and attention
// must be paid when handling the Cryptographic chain of trust.
func New(path string, signer *openpgp.Entity) (*Archive, error) {
	store, err := blobstore.Load(path)
	if err != nil {
		return nil, err
	}

	return &Archive{
		store:      *store,
		signingKey: signer,
		path:       path,
	}, nil
}

func (a Archive) Path() string {
	return a.path
}

// Use the default backend to remove any unlinked files from the Blob store.
//
// If files you care about are not linked onto the stage, they will be removed
// by the garbage collector. Decruft only when you're sure the stage has been
// set.
func (a Archive) Decruft() error {
	return a.store.GC(blobstore.DumbGarbageCollector{})
}

// Get the current state of files in the Archive, which is an underlying
// call to the blobstore Paths function.
func (a Archive) Paths() (ArchiveState, error) {
	return a.store.Paths()
}

// Get an io.ReadCloser from the underlying Blobstore. This is a passthrough
// to the underlying blobstore function.
func (a Archive) Open(o blobstore.Object) (io.ReadCloser, error) {
	return a.store.Open(o)
}

// Given a list of objects, link them to the keyed paths.
func (a Archive) Link(blobs ArchiveState) error {
	for path, obj := range blobs {
		if err := a.store.Link(obj, path); err != nil {
			return err
		}
	}
	return nil
}

// Create a new Release object from a Suite, passing off the Name, Description
// and constructing the rest of the goodies.
//
// This will be an entirely empty object, without anything read off disk.
func newRelease(suite Suite) (*Release, error) {
	when := time.Now()

	var validUntil string = ""
	if suite.features.Duration != "" {
		duration, err := time.ParseDuration(suite.features.Duration)
		if err != nil {
			return nil, err
		}
		validUntil = when.Add(duration).Format(time.RFC1123Z)
	}

	release := Release{
		Suite:       suite.Name,
		Description: suite.Description,
		ValidUntil:  validUntil,
		Origin:      suite.Origin,
		Label:       suite.Label,
		Version:     suite.Version,
	}
	release.Date = when.Format(time.RFC1123Z)
	release.Architectures = []dependency.Arch{}
	release.Components = []string{}
	release.SHA256 = []control.SHA256FileHash{}
	release.SHA1 = []control.SHA1FileHash{}
	release.SHA512 = []control.SHA512FileHash{}
	release.MD5Sum = []control.MD5FileHash{}
	return &release, nil
}

// This is a set of file changes ready to be passed to `Link` to link in.
// Basically, this maps file paths to blobstore objects, which will be
// swapped in all at once. This allows errors to avoid mutating state in
// the archive.
type ArchiveState map[string]blobstore.Object

// Engross a Suite for signing and final commit into the blobstore. This
// will return handle(s) to the signed and ready Objects, fit for passage
// to Link.
//
// This will contain all the related Packages and Release files.
func (a Archive) Engross(suite Suite) (ArchiveState, error) {
	release, err := newRelease(suite)
	if err != nil {
		return nil, err
	}

	files := ArchiveState{}
	arches := map[dependency.Arch]bool{}

	for name, component := range suite.components {
		release.Components = append(release.Components, name)
		for arch, writer := range component.packageWriters {
			arches[arch] = true

			// For each Binary entry, do the same as above (todo: someone
			// DRY this out a bit. I'm too lazy.

			suitePath := path.Join(name, fmt.Sprintf("binary-%s", arch),
				"Packages")

			obj, err := a.store.Commit(*writer.handle)
			if err != nil {
				return nil, err
			}

			for _, hasher := range writer.hashers {
				fileHash := control.FileHashFromHasher(suitePath, *hasher)
				release.AddHash(fileHash)
			}

			filePath := path.Join("dists", suite.Name, suitePath)
			files[filePath] = *obj
		}
	}

	for arch, _ := range arches {
		release.Architectures = append(release.Architectures, arch)
	}

	/* Now, let's do some magic */

	// Now, let's write out the Release file (and sign it normally)
	obj, sig, err := suite.archive.encodeSigned(release)
	if err != nil {
		return nil, err
	}

	filePath := path.Join("dists", suite.Name, "Release")
	files[filePath] = *obj
	files[fmt.Sprintf("%s.gpg", filePath)] = *sig

	// Ditto with the clearsigned version (Should we merge the two above?)
	obj, err = suite.archive.encodeClearsigned(release)
	if err != nil {
		return nil, err
	}

	files[path.Join("dists", suite.Name, "InRelease")] = *obj

	return files, nil
}

// Given a control.Marshal'able object, encode it to the blobstore, while
// also clearsigning the data.
func (a Archive) encodeClearsigned(data interface{}) (*blobstore.Object, error) {

	if a.signingKey == nil {
		return nil, fmt.Errorf("No signing key loaded")
	}

	fd, err := a.store.Create()
	if err != nil {
		return nil, err
	}

	defer fd.Close()
	wc, err := clearsign.Encode(fd, a.signingKey.PrivateKey, nil)
	if err != nil {
		return nil, err
	}

	encoder, err := control.NewEncoder(wc)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	if err := wc.Close(); err != nil {
		return nil, err
	}

	return a.store.Commit(*fd)
}

// Given a control.Marshal'able object, encode it to the blobstore, while
// also doing a detached OpenPGP signature. The objects returned (in order)
// are data, commited to the blobstore, the signature for that object, commited
// to the blobstore, and any error(s), finally.
func (a Archive) encodeSigned(data interface{}) (*blobstore.Object, *blobstore.Object, error) {
	/* Right, so, the trick here is that we secretly call out to encode,
	 * but tap it with a pipe into the signing code */

	if a.signingKey == nil {
		return nil, nil, fmt.Errorf("No signing key loaded")
	}

	signature, err := a.store.Create()
	if err != nil {
		return nil, nil, err
	}
	defer signature.Close()

	hash := sha512.New()

	obj, err := a.encode(data, hash)
	if err != nil {
		return nil, nil, err
	}

	sig := new(packet.Signature)
	sig.SigType = packet.SigTypeBinary
	sig.PubKeyAlgo = a.signingKey.PrivateKey.PubKeyAlgo

	sig.Hash = crypto.SHA512

	sig.CreationTime = new(packet.Config).Now()
	sig.IssuerKeyId = &(a.signingKey.PrivateKey.KeyId)

	err = sig.Sign(hash, a.signingKey.PrivateKey, nil)
	if err != nil {
		return nil, nil, err
	}

	if err := sig.Serialize(signature); err != nil {
		return nil, nil, err
	}

	sigObj, err := a.store.Commit(*signature)
	if err != nil {
		return nil, nil, err
	}

	return obj, sigObj, nil

}

// Encode a given control.Marshal'able object into the Blobstore, and return
// a handle to its object.
//
// The optinal argument `tap` will be written to as the object gets sent into
// the blobstore. This may be useful if you wish to have a copy of the data
// going into the store.
func (a Archive) encode(data interface{}, tap io.Writer) (*blobstore.Object, error) {
	fd, err := a.store.Create()
	if err != nil {
		return nil, err
	}

	var writer io.Writer = fd
	if tap != nil {
		writer = io.MultiWriter(fd, tap)
	}

	encoder, err := control.NewEncoder(writer)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	return a.store.Commit(*fd)
}

// }}}

// Suite {{{

// Abstraction to handle writing data into a Suite. This is a write-only
// target, and is not intended to read a Release file.
//
// This contains no state read off disk, and is purely for writing to.
type Suite struct {
	control.Paragraph

	archive *Archive

	Name        string `control:"Suite"`
	Description string
	Origin      string
	Label       string
	Version     string

	Pool       Pool                  `control:"-"`
	components map[string]*Component `control:"-"`

	features struct {
		Hashes   []string
		Duration string
	} `control:"-"`
}

// Get a handle to write a given Suite from an Archive.
// The suite will be entirely blank, and attributes will not be
// read from the existing files, if any.
func (a Archive) Suite(name string) (*Suite, error) {
	suite := Suite{
		Name:       name,
		archive:    &a,
		components: map[string]*Component{},
	}

	suite.Pool = Pool{store: a.store, suite: &suite}
	suite.features.Hashes = []string{"sha256", "sha1", "sha512"}
	suite.features.Duration = "240h"

	return &suite, nil
}

// Get or create a Component for a given Suite. If no such Component
// has been created so far, this will create a new object, otherwise
// it will return the existing entry.
//
// This contains no state read off disk, and is purely for writing to.
func (s Suite) Component(name string) (*Component, error) {
	if _, ok := s.components[name]; !ok {
		comp, err := newComponent(&s)
		if err != nil {
			return nil, err
		}
		s.components[name] = comp
		return comp, nil
	}
	el := s.components[name]
	return el, nil
}

// }}}

// Component {{{

// Small wrapper to represent a Component of a Suite, which, at its core
// is simply a set of Indexes to be written to.
//
// This contains no state read off disk, and is purely for writing to.
type Component struct {
	suite          *Suite
	packageWriters map[dependency.Arch]*IndexWriter
}

// Create a new Component, configured for use.
func newComponent(suite *Suite) (*Component, error) {
	return &Component{
		suite:          suite,
		packageWriters: map[dependency.Arch]*IndexWriter{},
	}, nil
}

// Get a given IndexWriter for an arch, or create one if none exists.
func (c *Component) getWriter(arch dependency.Arch) (*IndexWriter, error) {
	if _, ok := c.packageWriters[arch]; !ok {
		writer, err := newIndexWriter(c.suite)
		if err != nil {
			return nil, err
		}
		c.packageWriters[arch] = writer
	}
	return c.packageWriters[arch], nil
}

// Add a given Package to a Package List. Under the hood, this will
// get or create a IndexWriter, and invoke the .Add method on the
// Package Writer.
func (c *Component) AddPackage(pkg Package) error {
	writer, err := c.getWriter(pkg.Architecture)
	if err != nil {
		return err
	}
	return writer.Add(pkg)
}

// }}}

// IndexWriter {{{

// This writer represents a Package list - which is to say, a list of
// binary .deb files, for a particular Architecture, in a particular Component
// in a particular Suite, in a particular Archive.
//
// This is not an encapsulation to store the entire Index in memory, rather,
// it's a wrapper to help write Package entries into the Index.
type IndexWriter struct {
	archive *Archive

	handle  *blobstore.Writer
	closer  func() error
	encoder *control.Encoder

	hashers []*transput.Hasher
}

func getHashers(suite *Suite) (io.Writer, []*transput.Hasher, error) {
	ret := []*transput.Hasher{}
	writers := []io.Writer{}

	for _, algo := range suite.features.Hashes {
		hasher, err := transput.NewHasher(algo)
		if err != nil {
			return nil, nil, err
		}
		writers = append(writers, hasher)
		ret = append(ret, hasher)
	}

	return io.MultiWriter(writers...), ret, nil
}

// given a Suite, create a new Package Writer, configured with
// the appropriate Hashing, and targeting a new file blob in the
// underlying blobstore.
func newIndexWriter(suite *Suite) (*IndexWriter, error) {
	handle, err := suite.archive.store.Create()
	if err != nil {
		return nil, err
	}

	writer, hashers, err := getHashers(suite)
	if err != nil {
		return nil, err
	}

	encoder, err := control.NewEncoder(io.MultiWriter(writer, handle))
	if err != nil {
		handle.Close()
		return nil, err
	}

	return &IndexWriter{
		archive: suite.archive,
		closer:  handle.Close,
		encoder: encoder,
		handle:  handle,
		hashers: hashers,
	}, nil
}

// Write a Package entry into the Packages index.
func (p IndexWriter) Add(data interface{}) error {
	return p.encoder.Encode(data)
}

// }}}

// vim: foldmethod=marker
<|endoftext|>


<commit_before>package esatompub

import (
	"crypto/rand"
	"database/sql"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/xtracdev/goes"
	"github.com/xtracdev/orapub"
	"os"
	"strconv"
)

const defaultFeedThreshold = 100

var FeedThreshold = defaultFeedThreshold

func ReadFeedThresholdFromEnv() {
	thresholdOverride := os.Getenv("FEED_THRESHOLD")
	if thresholdOverride != "" {
		threshold, err := strconv.Atoi(thresholdOverride)
		if err != nil {
			log.Warnf("Attempted to override threshold with non integer: %s", thresholdOverride)
			return
		}

		log.Infof("Overriding default feed threshold with %d", threshold)
		FeedThreshold = threshold
	}
}

func readPreviousFeedId(tx *sql.Tx) (sql.NullString, error) {
	//TODO: Need to lock feeds or concurrent processes might overwrite each other's feed update
	log.Info("Select last feed id")
	var feedid sql.NullString
	rows, err := tx.Query("select feedid from feeds where event_time = (select max(event_time) from feeds)")
	if err != nil {
		log.Warn(err.Error())
		return feedid, err
	}

	defer rows.Close()
	for rows.Next() {
		//Only one row can be returned at mpst
		if err := rows.Scan(&feedid); err != nil {
			return feedid, err
		}
	}

	return feedid, nil
}

func writeEventToRecentTable(tx *sql.Tx, event *goes.Event) error {
	log.Info("insert event into recent")
	_, err := tx.Exec("insert into recent (aggregate_id, version,typecode, payload) values(:1,:2,:3,:4)",
		event.Source, event.Version, event.TypeCode, event.Payload)
	return err
}

func getRecentFeedCount(tx *sql.Tx) (int, error) {
	log.Info("get current count")
	var count int
	err := tx.QueryRow("select count(*) from recent where feedid is null").Scan(&count)
	return count, err
}

func createNewFeed(tx *sql.Tx, currentFeedId sql.NullString) error {
	log.Infof("Feed threshold of %d met", FeedThreshold)
	var prevFeedId sql.NullString
	uuidStr, err := uuid()
	if err != nil {
		return nil
	}

	if currentFeedId.Valid {
		prevFeedId = currentFeedId

	}
	currentFeedId = sql.NullString{String: uuidStr, Valid: true}

	log.Info("Update feed ids")
	_, err = tx.Exec("update recent set feedid = :1 where feedid is null", currentFeedId)
	if err != nil {
		return err
	}

	log.Info("Insert into feeds %v, %v", currentFeedId, prevFeedId)
	_, err = tx.Exec("insert into feeds (feedid, previous) values (:1, :2)",
		currentFeedId, prevFeedId)
	return err
}

func NewESAtomPubProcessor() orapub.EventProcessor {
	return orapub.EventProcessor{
		Initialize: func(db *sql.DB) error {
			return nil
		},
		Processor: func(db *sql.DB, event *goes.Event) error {
			log.Info("Processor invoked")

			//Need a transaction to group the work in this method
			log.Info("create transaction")
			tx, err := db.Begin()
			if err != nil {
				return err
			}

			defer tx.Rollback()

			//Get the current feed id
			feedid, err := readPreviousFeedId(tx)
			log.Infof("previous feed id is %s", feedid.String)

			//Insert current row
			err = writeEventToRecentTable(tx, event)
			if err != nil {
				return err
			}

			//Get current count of records in the current feed
			count, err := getRecentFeedCount(tx)
			if err != nil {
				return err
			}
			log.Infof("current count is %d", count)

			//Threshold met
			if count == FeedThreshold {
				err := createNewFeed(tx, feedid)
				if err != nil {
					return err
				}
			}

			log.Info("commit txn")
			tx.Commit()

			return nil
		},
	}
}

func uuid() (string, error) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%X-%X-%X-%X-%X", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]), nil

}
<commit_msg>Extract inline event processor code into separate function<commit_after>package esatompub

import (
	"crypto/rand"
	"database/sql"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/xtracdev/goes"
	"github.com/xtracdev/orapub"
	"os"
	"strconv"
)

const defaultFeedThreshold = 100

var FeedThreshold = defaultFeedThreshold

func ReadFeedThresholdFromEnv() {
	thresholdOverride := os.Getenv("FEED_THRESHOLD")
	if thresholdOverride != "" {
		threshold, err := strconv.Atoi(thresholdOverride)
		if err != nil {
			log.Warnf("Attempted to override threshold with non integer: %s", thresholdOverride)
			return
		}

		log.Infof("Overriding default feed threshold with %d", threshold)
		FeedThreshold = threshold
	}
}

func readPreviousFeedId(tx *sql.Tx) (sql.NullString, error) {
	//TODO: Need to lock feeds or concurrent processes might overwrite each other's feed update
	log.Info("Select last feed id")
	var feedid sql.NullString
	rows, err := tx.Query("select feedid from feeds where event_time = (select max(event_time) from feeds)")
	if err != nil {
		log.Warn(err.Error())
		return feedid, err
	}

	defer rows.Close()
	for rows.Next() {
		//Only one row can be returned at mpst
		if err := rows.Scan(&feedid); err != nil {
			return feedid, err
		}
	}

	return feedid, nil
}

func writeEventToRecentTable(tx *sql.Tx, event *goes.Event) error {
	log.Info("insert event into recent")
	_, err := tx.Exec("insert into recent (aggregate_id, version,typecode, payload) values(:1,:2,:3,:4)",
		event.Source, event.Version, event.TypeCode, event.Payload)
	return err
}

func getRecentFeedCount(tx *sql.Tx) (int, error) {
	log.Info("get current count")
	var count int
	err := tx.QueryRow("select count(*) from recent where feedid is null").Scan(&count)
	return count, err
}

func createNewFeed(tx *sql.Tx, currentFeedId sql.NullString) error {
	log.Infof("Feed threshold of %d met", FeedThreshold)
	var prevFeedId sql.NullString
	uuidStr, err := uuid()
	if err != nil {
		return nil
	}

	if currentFeedId.Valid {
		prevFeedId = currentFeedId

	}
	currentFeedId = sql.NullString{String: uuidStr, Valid: true}

	log.Info("Update feed ids")
	_, err = tx.Exec("update recent set feedid = :1 where feedid is null", currentFeedId)
	if err != nil {
		return err
	}

	log.Info("Insert into feeds %v, %v", currentFeedId, prevFeedId)
	_, err = tx.Exec("insert into feeds (feedid, previous) values (:1, :2)",
		currentFeedId, prevFeedId)
	return err
}

func processEvent(db *sql.DB, event *goes.Event) error {
	log.Info("Processor invoked")

	//Need a transaction to group the work in this method
	log.Info("create transaction")
	tx, err := db.Begin()
	if err != nil {
		return err
	}

	defer tx.Rollback()

	//Get the current feed id
	feedid, err := readPreviousFeedId(tx)
	log.Infof("previous feed id is %s", feedid.String)

	//Insert current row
	err = writeEventToRecentTable(tx, event)
	if err != nil {
		return err
	}

	//Get current count of records in the current feed
	count, err := getRecentFeedCount(tx)
	if err != nil {
		return err
	}
	log.Infof("current count is %d", count)

	//Threshold met
	if count == FeedThreshold {
		err := createNewFeed(tx, feedid)
		if err != nil {
			return err
		}
	}

	log.Info("commit txn")
	tx.Commit()

	return nil
}

func NewESAtomPubProcessor() orapub.EventProcessor {
	return orapub.EventProcessor{
		Initialize: func(db *sql.DB) error {
			return nil
		},
		Processor: func(db *sql.DB, event *goes.Event) error {
			return processEvent(db, event)
		},
	}
}

func uuid() (string, error) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%X-%X-%X-%X-%X", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]), nil

}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative 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 leaderelection

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/kelseyhightower/envconfig"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/client-go/tools/leaderelection/resourcelock"

	cm "knative.dev/pkg/configmap"
)

const (
	configMapNameEnv = "CONFIG_LEADERELECTION_NAME"
	// knativeResourceLock is the only supported lock mechanism for Knative.
	knativeResourceLock = resourcelock.LeasesResourceLock
)

// MaxBuckets is the maximum number of buckets to allow users to define.
// This is a variable so that it may be customized in the binary entrypoint.
var MaxBuckets uint32 = 10

// NewConfigFromMap returns a Config for the given map, or an error.
func NewConfigFromMap(data map[string]string) (*Config, error) {
	config := defaultConfig()

	if err := cm.Parse(data,
		cm.AsDuration("leaseDuration", &config.LeaseDuration),
		cm.AsDuration("renewDeadline", &config.RenewDeadline),
		cm.AsDuration("retryPeriod", &config.RetryPeriod),

		cm.AsUint32("buckets", &config.Buckets),
	); err != nil {
		return nil, err
	}

	if config.Buckets < 1 || config.Buckets > MaxBuckets {
		return nil, fmt.Errorf("buckets: value must be between %d <= %d <= %d", 1, config.Buckets, MaxBuckets)
	}
	return config, nil
}

// NewConfigFromConfigMap returns a new Config from the given ConfigMap.
func NewConfigFromConfigMap(configMap *corev1.ConfigMap) (*Config, error) {
	if configMap == nil {
		config := defaultConfig()
		return config, nil
	}
	return NewConfigFromMap(configMap.Data)
}

// Config represents the leader election config for a set of components
// contained within a single namespace. Typically these will correspond to a
// single source repository, viz: serving or eventing.
type Config struct {
	Buckets       uint32
	LeaseDuration time.Duration
	RenewDeadline time.Duration
	RetryPeriod   time.Duration

	// This field is deprecated and will be removed once downstream
	// repositories have removed their validation of it.
	// TODO(https://github.com/knative/pkg/issues/1478): Remove this field.
	EnabledComponents sets.String
}

func (c *Config) GetComponentConfig(name string) ComponentConfig {
	return ComponentConfig{
		Component:     name,
		Buckets:       c.Buckets,
		LeaseDuration: c.LeaseDuration,
		RenewDeadline: c.RenewDeadline,
		RetryPeriod:   c.RetryPeriod,
	}
}

func defaultConfig() *Config {
	return &Config{
		Buckets:       1,
		LeaseDuration: 15 * time.Second,
		RenewDeadline: 10 * time.Second,
		RetryPeriod:   2 * time.Second,
	}
}

// ComponentConfig represents the leader election config for a single component.
type ComponentConfig struct {
	Component     string
	Buckets       uint32
	LeaseDuration time.Duration
	RenewDeadline time.Duration
	RetryPeriod   time.Duration
	// LeaseName is a function to customize the lease name given the index i.
	// If not present, a name in format {Component}.{queue-name}.{i}-of-{Buckets}
	// will be use.
	// Autoscaler need to know the Lease names to filter out Leases which are not
	// used for Autoscaler. Instead of exposing the names from leadelection package,
	// we let Autoscaler to pass them in.
	LeaseName func(i uint32) string `json:"-"`
	// Identity is the unique string identifying a resource lock holder across
	// all participants in an election. If not present, a new unique string will
	// be generated to be used as identity for each BuildElector call.
	// Autoscaler uses the pod IP as identity.
	Identity string
}

// statefulSetID is a envconfig Decodable controller ordinal and name.
type statefulSetID struct {
	ssName  string
	ordinal int
}

func (ssID *statefulSetID) Decode(v string) error {
	if i := strings.LastIndex(v, "-"); i != -1 {
		ui, err := strconv.Atoi(v[i+1:])
		ssID.ordinal = ui
		ssID.ssName = v[:i]
		return err
	}
	return fmt.Errorf("%q is not a valid stateful set controller ordinal", v)
}

var _ envconfig.Decoder = (*statefulSetID)(nil)

// statefulSetConfig represents the required information for a StatefulSet service.
type statefulSetConfig struct {
	StatefulSetID statefulSetID `envconfig:"STATEFUL_CONTROLLER_ORDINAL" required:"true"`
	ServiceName   string        `envconfig:"STATEFUL_SERVICE_NAME" required:"true"`
	Port          string        `envconfig:"STATEFUL_SERVICE_PORT" default:"80"`
	Protocol      string        `envconfig:"STATEFUL_SERVICE_PROTOCOL" default:"http"`
}

// newStatefulSetConfig builds a stateful set LE config.
func newStatefulSetConfig() (*statefulSetConfig, error) {
	ssc := &statefulSetConfig{}
	if err := envconfig.Process("", ssc); err != nil {
		return nil, err
	}
	return ssc, nil
}

// ConfigMapName returns the name of the configmap to read for leader election
// settings.
func ConfigMapName() string {
	cm := os.Getenv(configMapNameEnv)
	if cm == "" {
		return "config-leader-election"
	}
	return cm
}

// UniqueID returns a unique ID for use with a leader elector that prevents from
// pods running on the same host from colliding with one another.
func UniqueID() (string, error) {
	id, err := os.Hostname()
	if err != nil {
		return "", err
	}

	return id + "_" + string(uuid.NewUUID()), nil
}
<commit_msg>drop deprecated enabledComponents field (#2162)<commit_after>/*
Copyright 2020 The Knative 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 leaderelection

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/kelseyhightower/envconfig"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/client-go/tools/leaderelection/resourcelock"

	cm "knative.dev/pkg/configmap"
)

const (
	configMapNameEnv = "CONFIG_LEADERELECTION_NAME"
	// knativeResourceLock is the only supported lock mechanism for Knative.
	knativeResourceLock = resourcelock.LeasesResourceLock
)

// MaxBuckets is the maximum number of buckets to allow users to define.
// This is a variable so that it may be customized in the binary entrypoint.
var MaxBuckets uint32 = 10

// NewConfigFromMap returns a Config for the given map, or an error.
func NewConfigFromMap(data map[string]string) (*Config, error) {
	config := defaultConfig()

	if err := cm.Parse(data,
		cm.AsDuration("leaseDuration", &config.LeaseDuration),
		cm.AsDuration("renewDeadline", &config.RenewDeadline),
		cm.AsDuration("retryPeriod", &config.RetryPeriod),

		cm.AsUint32("buckets", &config.Buckets),
	); err != nil {
		return nil, err
	}

	if config.Buckets < 1 || config.Buckets > MaxBuckets {
		return nil, fmt.Errorf("buckets: value must be between %d <= %d <= %d", 1, config.Buckets, MaxBuckets)
	}
	return config, nil
}

// NewConfigFromConfigMap returns a new Config from the given ConfigMap.
func NewConfigFromConfigMap(configMap *corev1.ConfigMap) (*Config, error) {
	if configMap == nil {
		config := defaultConfig()
		return config, nil
	}
	return NewConfigFromMap(configMap.Data)
}

// Config represents the leader election config for a set of components
// contained within a single namespace. Typically these will correspond to a
// single source repository, viz: serving or eventing.
type Config struct {
	Buckets       uint32
	LeaseDuration time.Duration
	RenewDeadline time.Duration
	RetryPeriod   time.Duration
}

func (c *Config) GetComponentConfig(name string) ComponentConfig {
	return ComponentConfig{
		Component:     name,
		Buckets:       c.Buckets,
		LeaseDuration: c.LeaseDuration,
		RenewDeadline: c.RenewDeadline,
		RetryPeriod:   c.RetryPeriod,
	}
}

func defaultConfig() *Config {
	return &Config{
		Buckets:       1,
		LeaseDuration: 15 * time.Second,
		RenewDeadline: 10 * time.Second,
		RetryPeriod:   2 * time.Second,
	}
}

// ComponentConfig represents the leader election config for a single component.
type ComponentConfig struct {
	Component     string
	Buckets       uint32
	LeaseDuration time.Duration
	RenewDeadline time.Duration
	RetryPeriod   time.Duration
	// LeaseName is a function to customize the lease name given the index i.
	// If not present, a name in format {Component}.{queue-name}.{i}-of-{Buckets}
	// will be use.
	// Autoscaler need to know the Lease names to filter out Leases which are not
	// used for Autoscaler. Instead of exposing the names from leadelection package,
	// we let Autoscaler to pass them in.
	LeaseName func(i uint32) string `json:"-"`
	// Identity is the unique string identifying a resource lock holder across
	// all participants in an election. If not present, a new unique string will
	// be generated to be used as identity for each BuildElector call.
	// Autoscaler uses the pod IP as identity.
	Identity string
}

// statefulSetID is a envconfig Decodable controller ordinal and name.
type statefulSetID struct {
	ssName  string
	ordinal int
}

func (ssID *statefulSetID) Decode(v string) error {
	if i := strings.LastIndex(v, "-"); i != -1 {
		ui, err := strconv.Atoi(v[i+1:])
		ssID.ordinal = ui
		ssID.ssName = v[:i]
		return err
	}
	return fmt.Errorf("%q is not a valid stateful set controller ordinal", v)
}

var _ envconfig.Decoder = (*statefulSetID)(nil)

// statefulSetConfig represents the required information for a StatefulSet service.
type statefulSetConfig struct {
	StatefulSetID statefulSetID `envconfig:"STATEFUL_CONTROLLER_ORDINAL" required:"true"`
	ServiceName   string        `envconfig:"STATEFUL_SERVICE_NAME" required:"true"`
	Port          string        `envconfig:"STATEFUL_SERVICE_PORT" default:"80"`
	Protocol      string        `envconfig:"STATEFUL_SERVICE_PROTOCOL" default:"http"`
}

// newStatefulSetConfig builds a stateful set LE config.
func newStatefulSetConfig() (*statefulSetConfig, error) {
	ssc := &statefulSetConfig{}
	if err := envconfig.Process("", ssc); err != nil {
		return nil, err
	}
	return ssc, nil
}

// ConfigMapName returns the name of the configmap to read for leader election
// settings.
func ConfigMapName() string {
	cm := os.Getenv(configMapNameEnv)
	if cm == "" {
		return "config-leader-election"
	}
	return cm
}

// UniqueID returns a unique ID for use with a leader elector that prevents from
// pods running on the same host from colliding with one another.
func UniqueID() (string, error) {
	id, err := os.Hostname()
	if err != nil {
		return "", err
	}

	return id + "_" + string(uuid.NewUUID()), nil
}
<|endoftext|>


<commit_before>// Copyright 2016 clair 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 migrations

import "github.com/remind101/migrate"

func init() {
	RegisterMigration(migrate.Migration{
		ID: 6,
		Up: migrate.Queries([]string{
			`ALTER TABLE Namespace ADD COLUMN version_format varchar(128);`,
		}),
		Down: migrate.Queries([]string{
			`ALTER TABLE Namespace DROP COLUMN version_format;`,
		}),
	})
}
<commit_msg>backfill version_format column<commit_after>// Copyright 2016 clair 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 migrations

import (
	"database/sql"
	"strings"

	"github.com/remind101/migrate"
)

func init() {
	RegisterMigration(migrate.Migration{
		ID: 6,
		Up: func(tx *sql.Tx) error {
			_, err := tx.Exec(`ALTER TABLE Namespace ADD COLUMN version_format varchar(128);`)
			if err != nil {
				return err
			}

			rows, err := tx.Query(`SELECT name FROM Namespace FOR UPDATE;`)
			if err != nil {
				return err
			}

			for rows.Next() {
				var nsName string
				err = rows.Scan(&nsName)
				if err != nil {
					return err
				}

				// Backfill rpm namespaces.
				if strings.HasPrefix(nsName, "rhel") ||
					strings.HasPrefix(nsName, "centos") ||
					strings.HasPrefix(nsName, "fedora") ||
					strings.HasPrefix(nsName, "amzn") ||
					strings.HasPrefix(nsName, "scientific") ||
					strings.HasPrefix(nsName, "ol") ||
					strings.HasPrefix(nsName, "oracle") {
					_, err := tx.Exec(`UPDATE Namespace SET version_format = 'rpm' WHERE name = ?;`, nsName)
					if err != nil {
						return err
					}
				} else {
					// Fallback to dpkg.
					_, err := tx.Exec(`UPDATE Namespace SET version_format = 'dpkg' WHERE name = ?;`, nsName)
					if err != nil {
						return err
					}
				}
			}

			return nil
		},
		Down: migrate.Queries([]string{
			`ALTER TABLE Namespace DROP COLUMN version_format;`,
		}),
	})
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strconv"

	awssdk "github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/elbv2"

	"github.com/gliderlabs/registrator/bridge"
	eureka "github.com/hudl/fargo"
)

// LBInfo represents a ELBv2 endpoint
type LBInfo struct {
	DNSName string
	Port    int64
}

var lbCache = make(map[string]*LBInfo)

// Helper function to retrieve all target groups
func getAllTargetGroups(svc *elbv2.ELBV2) ([]*elbv2.DescribeTargetGroupsOutput, error) {
	var tgs []*elbv2.DescribeTargetGroupsOutput
	var e error
	var mark *string

	// Get first page of groups
	tg, e := getTargetGroupsPage(svc, mark)

	if e != nil {
		return nil, e
	}
	tgs = append(tgs, tg)
	mark = tg.NextMarker

	// Page through all remaining target groups generating a slice of DescribeTargetGroupOutputs
	for mark != nil {
		tg, e = getTargetGroupsPage(svc, mark)
		tgs = append(tgs, tg)
		mark = tg.NextMarker
		if e != nil {
			return nil, e
		}
	}
	return tgs, e
}

// Helper function to get a page of target groups
func getTargetGroupsPage(svc *elbv2.ELBV2, marker *string) (*elbv2.DescribeTargetGroupsOutput, error) {
	params := &elbv2.DescribeTargetGroupsInput{
		PageSize: awssdk.Int64(400),
		Marker:   marker,
	}
	tg, e := svc.DescribeTargetGroups(params)

	if e != nil {
		log.Printf("An error occurred using DescribeTargetGroups: %s \n", e.Error())
		return nil, e
	}
	return tg, nil
}

// getELBV2ForContainer returns an LBInfo struct with the load balancer DNS name and listener port for a given instanceId and port
// if an error occurs, or the target is not found, an empty LBInfo is returned.
// Return the DNS:port pair as an identifier to put in the container's registration metadata
// Pass it the instanceID for the docker host, and the the host port to lookup the associated ELB.
// useCache parameter, if true, will retrieve ELBv2 details from memory, rather than calling AWS.
// this is only really safe to use for heartbeat calls, as details can change dynamically
func getELBV2ForContainer(instanceID string, port int64, useCache bool) (lbinfo *LBInfo, err error) {

	// Retrieve from basic cache (for heartbeats)
	cacheKey := instanceID + "_" + strconv.FormatInt(port, 10)
	if val, ok := lbCache[cacheKey]; ok && useCache {
		log.Println("Retrieving value from cache.")
		return val, nil
	}

	var lbArns []*string
	var lbPort *int64
	var tgArn string
	info := &LBInfo{}

	sess, err := session.NewSession()
	if err != nil {
		message := fmt.Errorf("Failed to create session connecting to AWS: %s", err)
		return nil, message
	}

	// Need to set the region here - we'll get it from instance metadata
	awsMetadata := GetMetadata()
	svc := elbv2.New(sess, awssdk.NewConfig().WithRegion(awsMetadata.Region))

	// TODO Note: There could be thousands of these, and we need to check them all.  Seems to be no
	// other way to retrieve a TG via instance/port with current API
	tgslice, err := getAllTargetGroups(svc)
	if err != nil || tgslice == nil {
		message := fmt.Errorf("Failed to retrieve Target Groups: %s", err)
		return nil, message
	}

	// Check each target group's target list for a matching port and instanceID
	// TODO Assumption: that that there is only one LB for the target group (though the data structure allows more)
	for _, tgs := range tgslice {
		for _, tg := range tgs.TargetGroups {

			thParams := &elbv2.DescribeTargetHealthInput{
				TargetGroupArn: awssdk.String(*tg.TargetGroupArn),
			}

			tarH, err := svc.DescribeTargetHealth(thParams)
			if err != nil {
				log.Printf("An error occurred using DescribeTargetHealth: %s \n", err.Error())
				return nil, err
			}

			for _, thd := range tarH.TargetHealthDescriptions {
				if *thd.Target.Port == port && *thd.Target.Id == instanceID {
					lbArns = tg.LoadBalancerArns
					tgArn = *tg.TargetGroupArn
					break
				}
			}
		}
		if lbArns != nil && tgArn != "" {
			break
		}
	}

	if err != nil || lbArns == nil {
		message := fmt.Errorf("failed to retrieve load balancer ARN")
		return nil, message
	}

	// Loop through the load balancer listeners to get the listener port for the target group
	lsnrParams := &elbv2.DescribeListenersInput{
		LoadBalancerArn: lbArns[0],
	}
	lnrData, err := svc.DescribeListeners(lsnrParams)
	if err != nil {
		log.Printf("An error occurred using DescribeListeners: %s \n", err.Error())
		return nil, err
	}
	for _, listener := range lnrData.Listeners {
		for _, act := range listener.DefaultActions {
			if *act.TargetGroupArn == tgArn {
				log.Printf("Found matching listener: %v", *listener.ListenerArn)
				lbPort = listener.Port
				break
			}
		}
	}
	if lbPort == nil {
		message := fmt.Errorf("error: Unable to identify listener port for ELBv2")
		return nil, message
	}

	// Get more information on the load balancer to retrieve the DNSName
	lbParams := &elbv2.DescribeLoadBalancersInput{
		LoadBalancerArns: lbArns,
	}
	lbData, err := svc.DescribeLoadBalancers(lbParams)
	if err != nil {
		log.Printf("An error occurred using DescribeLoadBalancers: %s \n", err.Error())
		return nil, err
	}
	log.Printf("LB Endpoint for Instance:%v Port:%v, Target Group:%v, is: %s:%s\n", instanceID, port, tgArn, *lbData.LoadBalancers[0].DNSName, strconv.FormatInt(*lbPort, 10))

	info.DNSName = *lbData.LoadBalancers[0].DNSName
	info.Port = *lbPort

	// Add to a basic cache for heartbeats
	lbCache[cacheKey] = info

	return info, nil
}

// CheckELBFlags - Helper function to check if the correct config flags are set to use ELBs
func CheckELBFlags(service *bridge.Service) bool {
	if service.Attrs["eureka_use_elbv2_endpoint"] != "" && service.Attrs["eureka_datacenterinfo_name"] != eureka.MyOwn {
		v, err := strconv.ParseBool(service.Attrs["eureka_use_elbv2_endpoint"])
		if err != nil {
			log.Printf("eureka: eureka_use_elbv2_endpoint must be valid boolean, was %v : %s", v, err)
			return false
		}
		return v
	}
	return false
}

// Helper function to create a registration struct, and change container registration
// useCache parameter is passed to getELBV2ForContainer
func setRegInfo(service *bridge.Service, registration *eureka.Instance, useCache bool) *eureka.Instance {

	awsMetadata := GetMetadata()

	elbMetadata, err := getELBV2ForContainer(awsMetadata.InstanceID, int64(registration.Port), useCache)

	if err != nil {
		log.Printf("Unable to find associated ELBv2 for: %s, Error: %s\n", registration.HostName, err)
		return nil
	}

	elbStrPort := strconv.FormatInt(elbMetadata.Port, 10)
	elbEndpoint := elbMetadata.DNSName + "_" + elbStrPort

	registration.SetMetadataString("has-elbv2", "true")
	registration.SetMetadataString("elbv2-endpoint", elbEndpoint)

	elbReg := new(eureka.Instance)

	// Put a little metadata in here as required - setting InstanceID to the ELB endpoint prevents double registration
	elbReg.DataCenterInfo.Metadata = eureka.AmazonMetadataType{
		PublicHostname: elbMetadata.DNSName,
		HostName:       elbMetadata.DNSName,
		InstanceID:     elbEndpoint,
	}

	elbReg.Port = int(elbMetadata.Port)
	elbReg.IPAddr = elbMetadata.DNSName
	elbReg.App = service.Name
	elbReg.VipAddress = elbReg.IPAddr
	elbReg.HostName = elbMetadata.DNSName
	elbReg.DataCenterInfo.Name = eureka.Amazon
	elbReg.SetMetadataString("is-elbv2", "true")
	elbReg.Status = eureka.UP
	return elbReg
}

// RegisterELBv2 - If specified, also register an ELBv2 (application load balancer, ALB) endpoint in eureka, and alter service name
// for container registrations.  This will mean traffic is directed to the ALB rather than directly to containers
// though they are still registered in eureka for information purposes
func RegisterELBv2(service *bridge.Service, registration *eureka.Instance, client eureka.EurekaConnection) {
	if CheckELBFlags(service) {
		log.Printf("Found ELBv2 flags, will attempt to register LB for: %s\n", registration.HostName)
		elbReg := setRegInfo(service, registration, false)
		if elbReg != nil {
			client.RegisterInstance(elbReg)
		}
	}
}

// DeregisterELBv2 - If specified, and all containers are gone, also deregister the ELBv2 (application load balancer, ALB) endpoint in eureka.
//
func DeregisterELBv2(service *bridge.Service, albEndpoint string, client eureka.EurekaConnection) {
	if CheckELBFlags(service) {
		// Check if there are any containers around with this ALB still attached
		log.Printf("Found ELBv2 flags, will check if it needs to be deregistered too, for: %v\n", albEndpoint)
		appName := "CONTAINER_" + service.Name

		app, err := client.GetApp(appName)
		if err != nil {
			log.Printf("Unable to retrieve app metadata for %s: %s\n", appName, err)
			return
		}

		if app != nil {
			for _, instance := range app.Instances {
				val, err := instance.Metadata.GetString("elbv2_endpoint")
				if err == nil && val == albEndpoint {
					log.Printf("Eureka entry still present for one or more ALB linked containers: %s\n", val)
					return
				}
			}
		}

		if app == nil {
			log.Printf("Removing eureka entry for ELBv2: %s\n", albEndpoint)
			elbReg := new(eureka.Instance)
			elbReg.App = service.Name
			elbReg.HostName = albEndpoint // This uses the full endpoint identifier so eureka can find it to remove
			client.DeregisterInstance(elbReg)
		}
	}
}

// HeartbeatELBv2 - Send a heartbeat to eureka for this ELBv2 registration.  Every host running registrator will send heartbeats, meaning they will
// be received more frequently than the --ttl-refresh interval if there are multiple hosts running registrator.
//
func HeartbeatELBv2(service *bridge.Service, registration *eureka.Instance, client eureka.EurekaConnection) {
	if CheckELBFlags(service) {
		log.Printf("Heartbeating ELBv2 for container: %s)\n", registration.HostName)

		elbReg := setRegInfo(service, registration, true) // Can safely use cache when heartbeating
		if elbReg != nil {
			client.HeartBeatInstance(elbReg)
		}
	}
}
<commit_msg>Add caching fix for ELB<commit_after>package aws

import (
	"fmt"
	"log"
	"strconv"

	awssdk "github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/elbv2"

	"github.com/gliderlabs/registrator/bridge"
	eureka "github.com/hudl/fargo"
)

// LBInfo represents a ELBv2 endpoint
type LBInfo struct {
	DNSName string
	Port    int64
}

var lbCache = make(map[string]*LBInfo)

// Helper function to retrieve all target groups
func getAllTargetGroups(svc *elbv2.ELBV2) ([]*elbv2.DescribeTargetGroupsOutput, error) {
	var tgs []*elbv2.DescribeTargetGroupsOutput
	var e error
	var mark *string

	// Get first page of groups
	tg, e := getTargetGroupsPage(svc, mark)

	if e != nil {
		return nil, e
	}
	tgs = append(tgs, tg)
	mark = tg.NextMarker

	// Page through all remaining target groups generating a slice of DescribeTargetGroupOutputs
	for mark != nil {
		tg, e = getTargetGroupsPage(svc, mark)
		tgs = append(tgs, tg)
		mark = tg.NextMarker
		if e != nil {
			return nil, e
		}
	}
	return tgs, e
}

// Helper function to get a page of target groups
func getTargetGroupsPage(svc *elbv2.ELBV2, marker *string) (*elbv2.DescribeTargetGroupsOutput, error) {
	params := &elbv2.DescribeTargetGroupsInput{
		PageSize: awssdk.Int64(400),
		Marker:   marker,
	}
	tg, e := svc.DescribeTargetGroups(params)

	if e != nil {
		log.Printf("An error occurred using DescribeTargetGroups: %s \n", e.Error())
		return nil, e
	}
	return tg, nil
}

// getELBV2ForContainer returns an LBInfo struct with the load balancer DNS name and listener port for a given instanceId and port
// if an error occurs, or the target is not found, an empty LBInfo is returned.
// Return the DNS:port pair as an identifier to put in the container's registration metadata
// Pass it the instanceID for the docker host, and the the host port to lookup the associated ELB.
// useCache parameter, if true, will retrieve ELBv2 details from memory, rather than calling AWS.
// this is only really safe to use for heartbeat calls, as details can change dynamically
func getELBV2ForContainer(containerID string, instanceID string, port int64, useCache bool) (lbinfo *LBInfo, err error) {

	// Retrieve from basic cache (for heartbeats)
	cacheKey := instanceID + "_" + containerID
	if val, ok := lbCache[cacheKey]; ok && useCache {
		log.Println("Retrieving value from cache.")
		return val, nil
	}

	var lbArns []*string
	var lbPort *int64
	var tgArn string
	info := &LBInfo{}

	sess, err := session.NewSession()
	if err != nil {
		message := fmt.Errorf("Failed to create session connecting to AWS: %s", err)
		return nil, message
	}

	// Need to set the region here - we'll get it from instance metadata
	awsMetadata := GetMetadata()
	svc := elbv2.New(sess, awssdk.NewConfig().WithRegion(awsMetadata.Region))

	// TODO Note: There could be thousands of these, and we need to check them all.  Seems to be no
	// other way to retrieve a TG via instance/port with current API
	tgslice, err := getAllTargetGroups(svc)
	if err != nil || tgslice == nil {
		message := fmt.Errorf("Failed to retrieve Target Groups: %s", err)
		return nil, message
	}

	// Check each target group's target list for a matching port and instanceID
	// TODO Assumption: that that there is only one LB for the target group (though the data structure allows more)
	for _, tgs := range tgslice {
		for _, tg := range tgs.TargetGroups {

			thParams := &elbv2.DescribeTargetHealthInput{
				TargetGroupArn: awssdk.String(*tg.TargetGroupArn),
			}

			tarH, err := svc.DescribeTargetHealth(thParams)
			if err != nil {
				log.Printf("An error occurred using DescribeTargetHealth: %s \n", err.Error())
				return nil, err
			}

			for _, thd := range tarH.TargetHealthDescriptions {
				if *thd.Target.Port == port && *thd.Target.Id == instanceID {
					lbArns = tg.LoadBalancerArns
					tgArn = *tg.TargetGroupArn
					break
				}
			}
		}
		if lbArns != nil && tgArn != "" {
			break
		}
	}

	if err != nil || lbArns == nil {
		message := fmt.Errorf("failed to retrieve load balancer ARN")
		return nil, message
	}

	// Loop through the load balancer listeners to get the listener port for the target group
	lsnrParams := &elbv2.DescribeListenersInput{
		LoadBalancerArn: lbArns[0],
	}
	lnrData, err := svc.DescribeListeners(lsnrParams)
	if err != nil {
		log.Printf("An error occurred using DescribeListeners: %s \n", err.Error())
		return nil, err
	}
	for _, listener := range lnrData.Listeners {
		for _, act := range listener.DefaultActions {
			if *act.TargetGroupArn == tgArn {
				log.Printf("Found matching listener: %v", *listener.ListenerArn)
				lbPort = listener.Port
				break
			}
		}
	}
	if lbPort == nil {
		message := fmt.Errorf("error: Unable to identify listener port for ELBv2")
		return nil, message
	}

	// Get more information on the load balancer to retrieve the DNSName
	lbParams := &elbv2.DescribeLoadBalancersInput{
		LoadBalancerArns: lbArns,
	}
	lbData, err := svc.DescribeLoadBalancers(lbParams)
	if err != nil {
		log.Printf("An error occurred using DescribeLoadBalancers: %s \n", err.Error())
		return nil, err
	}
	log.Printf("LB Endpoint for Instance:%v Port:%v, Target Group:%v, is: %s:%s\n", instanceID, port, tgArn, *lbData.LoadBalancers[0].DNSName, strconv.FormatInt(*lbPort, 10))

	info.DNSName = *lbData.LoadBalancers[0].DNSName
	info.Port = *lbPort

	// Add to a basic cache for heartbeats
	lbCache[cacheKey] = info

	return info, nil
}

// CheckELBFlags - Helper function to check if the correct config flags are set to use ELBs
func CheckELBFlags(service *bridge.Service) bool {
	if service.Attrs["eureka_use_elbv2_endpoint"] != "" && service.Attrs["eureka_datacenterinfo_name"] != eureka.MyOwn {
		v, err := strconv.ParseBool(service.Attrs["eureka_use_elbv2_endpoint"])
		if err != nil {
			log.Printf("eureka: eureka_use_elbv2_endpoint must be valid boolean, was %v : %s", v, err)
			return false
		}
		return v
	}
	return false
}

// Helper function to create a registration struct, and change container registration
// useCache parameter is passed to getELBV2ForContainer
func setRegInfo(service *bridge.Service, registration *eureka.Instance, useCache bool) *eureka.Instance {

	awsMetadata := GetMetadata()

	elbMetadata, err := getELBV2ForContainer(service.Origin.ContainerID, awsMetadata.InstanceID, int64(registration.Port), useCache)

	if err != nil {
		log.Printf("Unable to find associated ELBv2 for: %s, Error: %s\n", registration.HostName, err)
		return nil
	}

	elbStrPort := strconv.FormatInt(elbMetadata.Port, 10)
	elbEndpoint := elbMetadata.DNSName + "_" + elbStrPort

	registration.SetMetadataString("has-elbv2", "true")
	registration.SetMetadataString("elbv2-endpoint", elbEndpoint)

	elbReg := new(eureka.Instance)

	// Put a little metadata in here as required - setting InstanceID to the ELB endpoint prevents double registration
	elbReg.DataCenterInfo.Metadata = eureka.AmazonMetadataType{
		PublicHostname: elbMetadata.DNSName,
		HostName:       elbMetadata.DNSName,
		InstanceID:     elbEndpoint,
	}

	elbReg.Port = int(elbMetadata.Port)
	elbReg.IPAddr = elbMetadata.DNSName
	elbReg.App = service.Name
	elbReg.VipAddress = elbReg.IPAddr
	elbReg.HostName = elbMetadata.DNSName
	elbReg.DataCenterInfo.Name = eureka.Amazon
	elbReg.SetMetadataString("is-elbv2", "true")
	elbReg.Status = eureka.UP
	return elbReg
}

// RegisterELBv2 - If specified, also register an ELBv2 (application load balancer, ALB) endpoint in eureka, and alter service name
// for container registrations.  This will mean traffic is directed to the ALB rather than directly to containers
// though they are still registered in eureka for information purposes
func RegisterELBv2(service *bridge.Service, registration *eureka.Instance, client eureka.EurekaConnection) {
	if CheckELBFlags(service) {
		log.Printf("Found ELBv2 flags, will attempt to register LB for: %s\n", registration.HostName)
		elbReg := setRegInfo(service, registration, false)
		if elbReg != nil {
			client.RegisterInstance(elbReg)
		}
	}
}

// DeregisterELBv2 - If specified, and all containers are gone, also deregister the ELBv2 (application load balancer, ALB) endpoint in eureka.
//
func DeregisterELBv2(service *bridge.Service, albEndpoint string, client eureka.EurekaConnection) {
	if CheckELBFlags(service) {
		// Check if there are any containers around with this ALB still attached
		log.Printf("Found ELBv2 flags, will check if it needs to be deregistered too, for: %v\n", albEndpoint)
		appName := "CONTAINER_" + service.Name

		app, err := client.GetApp(appName)
		if err != nil {
			log.Printf("Unable to retrieve app metadata for %s: %s\n", appName, err)
			return
		}

		if app != nil {
			for _, instance := range app.Instances {
				val, err := instance.Metadata.GetString("elbv2_endpoint")
				if err == nil && val == albEndpoint {
					log.Printf("Eureka entry still present for one or more ALB linked containers: %s\n", val)
					return
				}
			}
		}

		if app == nil {
			log.Printf("Removing eureka entry for ELBv2: %s\n", albEndpoint)
			elbReg := new(eureka.Instance)
			elbReg.App = service.Name
			elbReg.HostName = albEndpoint // This uses the full endpoint identifier so eureka can find it to remove
			client.DeregisterInstance(elbReg)
		}
	}
}

// HeartbeatELBv2 - Send a heartbeat to eureka for this ELBv2 registration.  Every host running registrator will send heartbeats, meaning they will
// be received more frequently than the --ttl-refresh interval if there are multiple hosts running registrator.
//
func HeartbeatELBv2(service *bridge.Service, registration *eureka.Instance, client eureka.EurekaConnection) {
	if CheckELBFlags(service) {
		log.Printf("Heartbeating ELBv2 for container: %s)\n", registration.HostName)

		elbReg := setRegInfo(service, registration, true) // Can safely use cache when heartbeating
		if elbReg != nil {
			client.HeartBeatInstance(elbReg)
		}
	}
}
<|endoftext|>


<commit_before>// Package imgur implements a Service which adds !commands for imgur image search
package imgur

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/matrix-org/go-neb/types"
	"github.com/matrix-org/gomatrix"
)

// ServiceType of the Imgur service
const ServiceType = "imgur"

var httpClient = &http.Client{}

type imgurGalleryImage struct {
	ID          string `json:"id"`          // id	string	The ID for the image
	Title       string `json:"title"`       // title	string	The title of the image.
	Description string `json:"description"` // description	string	Description of the image.
	DateTime    int64  `json:"datetime"`    // datetime	integer	Time inserted into the gallery, epoch time
	Type        string `json:"type"`        // type	string	Image MIME type.
	Animated    bool   `json:"animated"`    // animated	boolean	is the image animated
	Width       int    `json:"width"`       // width	integer	The width of the image in pixels
	Height      int    `json:"height"`      // height	integer	The height of the image in pixels
	Size        int64  `json:"size"`        // size	integer	The size of the image in bytes
	Views       int64  `json:"views"`       // views	integer	The number of image views
	Link        string `json:"link"`        // link	string	The direct link to the the image. (Note: if fetching an animated GIF that was over 20MB in original size, a .gif thumbnail will be returned)
	Gifv        string `json:"gifv"`        // gifv	string	OPTIONAL, The .gifv link. Only available if the image is animated and type is 'image/gif'.
	MP4         string `json:"mp4"`         // mp4	string	OPTIONAL, The direct link to the .mp4. Only available if the image is animated and type is 'image/gif'.
	MP4Size     int64  `json:"mp4_size"`    // mp4_size	integer	OPTIONAL, The Content-Length of the .mp4. Only available if the image is animated and type is 'image/gif'. Note that a zero value (0) is possible if the video has not yet been generated
	Looping     bool   `json:"looping"`     // looping	boolean	OPTIONAL, Whether the image has a looping animation. Only available if the image is animated and type is 'image/gif'.
	NSFW        bool   `json:"nsfw"`        // nsfw	boolean	Indicates if the image has been marked as nsfw or not. Defaults to null if information is not available.
	Topic       string `json:"topic"`       // topic	string	Topic of the gallery image.
	Section     string `json:"section"`     // section	string	If the image has been categorized by our backend then this will contain the section the image belongs in. (funny, cats, adviceanimals, wtf, etc)
	IsAlbum     bool   `json:"is_album"`    // is_album	boolean	If it's an album or not
	// ** Uninplemented fields **
	// bandwidth	integer	Bandwidth consumed by the image in bytes
	// deletehash	string	OPTIONAL, the deletehash, if you're logged in as the image owner
	// comment_count	int	Number of comments on the gallery image.
	// topic_id	integer	Topic ID of the gallery image.
	// vote	string	The current user's vote on the album. null if not signed in or if the user hasn't voted on it.
	// favorite	boolean	Indicates if the current user favorited the image. Defaults to false if not signed in.
	// account_url	string	The username of the account that uploaded it, or null.
	// account_id	integer	The account ID of the account that uploaded it, or null.
	// ups	integer	Upvotes for the image
	// downs	integer	Number of downvotes for the image
	// points	integer	Upvotes minus downvotes
	// score	integer	Imgur popularity score
}

type imgurGalleryAlbum struct {
	ID          string              `json:"id"`           // id	string	The ID for the album
	Title       string              `json:"title"`        // title	string	The title of the album.
	Description string              `json:"description"`  // description	string	Description of the album.
	DateTime    int64               `json:"datetime"`     // datetime	integer	Time inserted into the gallery, epoch time
	Views       int64               `json:"views"`        // views	integer	The number of album views
	Link        string              `json:"link"`         // link	string	The URL link to the album
	NSFW        bool                `json:"nsfw"`         // nsfw	boolean	Indicates if the album has been marked as nsfw or not. Defaults to null if information is not available.
	Topic       string              `json:"topic"`        // topic	string	Topic of the gallery album.
	IsAlbum     bool                `json:"is_album"`     // is_album	boolean	If it's an album or not
	Cover       string              `json:"cover"`        // cover	string	The ID of the album cover image
	CoverWidth  int                 `json:"cover_width"`  // cover_width	integer	The width, in pixels, of the album cover image
	CoverHeight int                 `json:"cover_height"` // cover_height	integer	The height, in pixels, of the album cover image
	ImagesCount int                 `json:"images_count"` // images_count	integer	The total number of images in the album
	Images      []imgurGalleryImage `json:"images"`       // images	Array of Images	An array of all the images in the album (only available when requesting the direct album)

	// ** Unimplemented fields **
	// account_url	string	The account username or null if it's anonymous.
	// account_id	integer	The account ID of the account that uploaded it, or null.
	// privacy	string	The privacy level of the album, you can only view public if not logged in as album owner
	// layout	string	The view layout of the album.
	// views	integer	The number of image views
	// ups	integer	Upvotes for the image
	// downs	integer	Number of downvotes for the image
	// points	integer	Upvotes minus downvotes
	// score	integer	Imgur popularity score
	// vote	string	The current user's vote on the album. null if not signed in or if the user hasn't voted on it.
	// favorite	boolean	Indicates if the current user favorited the album. Defaults to false if not signed in.
	// comment_count	int	Number of comments on the gallery album.
	// topic_id	integer	Topic ID of the gallery album.
}

type imgurSearchResponse struct {
	Data    []json.RawMessage `json:"data"`
	Success bool              `json:"success"` // Request completed successfully
	Status  int               `json:"status"`  // HTTP response code
}

// Service contains the Config fields for the Imgur service.
//
// Example request:
//   {
//			"client_id": "AIzaSyA4FD39..."
//			"client_secret": "ASdsaijwdfASD..."
//   }
type Service struct {
	types.DefaultService
	// The Imgur client ID
	ClientID string `json:"client_id"`
	// The API key to use when making HTTP requests to Imgur.
	ClientSecret string `json:"client_secret"`
}

// Commands supported:
//    !imgur some_search_query_without_quotes
// Responds with a suitable image into the same room as the command.
func (s *Service) Commands(client *gomatrix.Client) []types.Command {
	return []types.Command{
		types.Command{
			Path: []string{"imgur", "help"},
			Command: func(roomID, userID string, args []string) (interface{}, error) {
				return usageMessage(), nil
			},
		},
		types.Command{
			Path: []string{"imgur"},
			Command: func(roomID, userID string, args []string) (interface{}, error) {
				return s.cmdImgurImgSearch(client, roomID, userID, args)
			},
		},
	}
}

// usageMessage returns a matrix TextMessage representation of the service usage
func usageMessage() *gomatrix.TextMessage {
	return &gomatrix.TextMessage{"m.notice",
		`Usage: !imgur image_search_text`}
}

func (s *Service) cmdImgurImgSearch(client *gomatrix.Client, roomID, userID string, args []string) (interface{}, error) {

	if len(args) < 1 {
		return usageMessage(), nil
	}

	// Get the query text to search for.
	querySentence := strings.Join(args, " ")

	searchResultImage, searchResultAlbum, err := s.text2imgImgur(querySentence)

	if err != nil {
		return nil, err
	}

	// Image returned
	if searchResultImage != nil {
		var imgURL = searchResultImage.Link
		if imgURL == "" {
			return gomatrix.TextMessage{
				MsgType: "m.notice",
				Body:    "No image found!",
			}, nil
		}

		// FIXME -- Sometimes upload fails with a cryptic error - "msg=Upload request failed code=400"
		// log.Printf("Uploading image at: %s", imgURL)
		resUpload, err := client.UploadLink(imgURL)
		if err != nil {
			return nil, fmt.Errorf("Failed to upload Imgur image (%s) to matrix: %s", imgURL, err.Error())
		}

		return gomatrix.ImageMessage{
			MsgType: "m.image",
			Body:    querySentence,
			URL:     resUpload.ContentURI,
			Info: gomatrix.ImageInfo{
				Height:   uint(searchResultImage.Height),
				Width:    uint(searchResultImage.Width),
				Mimetype: searchResultImage.Type,
			},
		}, nil
	} else if searchResultAlbum != nil {
		return gomatrix.TextMessage{
			MsgType: "m.notice",
			Body:    "Search returned an album - Not currently supported",
		}, nil
	} else {
		return gomatrix.TextMessage{
			MsgType: "m.notice",
			Body:    "No image found!",
		}, nil
	}
}

// text2imgImgur returns info about an image or an album
func (s *Service) text2imgImgur(query string) (*imgurGalleryImage, *imgurGalleryAlbum, error) {
	log.Info("Searching Imgur for an image of a ", query)

	query = url.QueryEscape(query)
	var base = "https://api.imgur.com/3/gallery/search"
	var sort = "time"  // time | viral | top
	var window = "all" // day | week | month | year | all
	var page = 1
	var urlString = fmt.Sprintf("%s/%s/%s/%d?q=%s", base, sort, window, page, query)
	// var urlString = fmt.Sprintf("%s?q=%s", base, query)

	u, err := url.Parse(urlString)
	if err != nil {
		return nil, nil, err
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return nil, nil, err
	}

	req.Header.Add("Authorization", "Client-ID "+s.ClientID)
	res, err := httpClient.Do(req)
	if res != nil {
		defer res.Body.Close()
	}
	if err != nil {
		return nil, nil, err
	}
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		return nil, nil, fmt.Errorf("Request error: %d, %s", res.StatusCode, response2String(res))
	}

	var searchResults imgurSearchResponse
	if err := json.NewDecoder(res.Body).Decode(&searchResults); err != nil {
		return nil, nil, fmt.Errorf("No images found - %s", err.Error())
	} else if len(searchResults.Data) < 1 {
		return nil, nil, fmt.Errorf("No images found")
	}

	log.Printf("%d results were returned from Imgur", len(searchResults.Data))
	// Return a random image result
	var images []imgurGalleryImage
	for i := 0; i < len(searchResults.Data); i++ {
		var image imgurGalleryImage
		if err := json.Unmarshal(searchResults.Data[i], &image); err == nil && !image.IsAlbum {
			images = append(images, image)
		}
	}
	if len(images) > 0 {
		var r = 0
		if len(images) > 1 {
			r = rand.Intn(len(images) - 1)
		}
		return &images[r], nil, nil
	}

	return nil, nil, fmt.Errorf("No images found")
}

// response2String returns a string representation of an HTTP response body
func response2String(res *http.Response) string {
	bs, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "Failed to decode response body"
	}
	str := string(bs)
	return str
}

// Initialise the service
func init() {
	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
		return &Service{
			DefaultService: types.NewDefaultService(serviceID, serviceUserID, ServiceType),
		}
	})
}
<commit_msg>Inline redundant variable<commit_after>// Package imgur implements a Service which adds !commands for imgur image search
package imgur

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/matrix-org/go-neb/types"
	"github.com/matrix-org/gomatrix"
)

// ServiceType of the Imgur service
const ServiceType = "imgur"

var httpClient = &http.Client{}

type imgurGalleryImage struct {
	ID          string `json:"id"`          // id	string	The ID for the image
	Title       string `json:"title"`       // title	string	The title of the image.
	Description string `json:"description"` // description	string	Description of the image.
	DateTime    int64  `json:"datetime"`    // datetime	integer	Time inserted into the gallery, epoch time
	Type        string `json:"type"`        // type	string	Image MIME type.
	Animated    bool   `json:"animated"`    // animated	boolean	is the image animated
	Width       int    `json:"width"`       // width	integer	The width of the image in pixels
	Height      int    `json:"height"`      // height	integer	The height of the image in pixels
	Size        int64  `json:"size"`        // size	integer	The size of the image in bytes
	Views       int64  `json:"views"`       // views	integer	The number of image views
	Link        string `json:"link"`        // link	string	The direct link to the the image. (Note: if fetching an animated GIF that was over 20MB in original size, a .gif thumbnail will be returned)
	Gifv        string `json:"gifv"`        // gifv	string	OPTIONAL, The .gifv link. Only available if the image is animated and type is 'image/gif'.
	MP4         string `json:"mp4"`         // mp4	string	OPTIONAL, The direct link to the .mp4. Only available if the image is animated and type is 'image/gif'.
	MP4Size     int64  `json:"mp4_size"`    // mp4_size	integer	OPTIONAL, The Content-Length of the .mp4. Only available if the image is animated and type is 'image/gif'. Note that a zero value (0) is possible if the video has not yet been generated
	Looping     bool   `json:"looping"`     // looping	boolean	OPTIONAL, Whether the image has a looping animation. Only available if the image is animated and type is 'image/gif'.
	NSFW        bool   `json:"nsfw"`        // nsfw	boolean	Indicates if the image has been marked as nsfw or not. Defaults to null if information is not available.
	Topic       string `json:"topic"`       // topic	string	Topic of the gallery image.
	Section     string `json:"section"`     // section	string	If the image has been categorized by our backend then this will contain the section the image belongs in. (funny, cats, adviceanimals, wtf, etc)
	IsAlbum     bool   `json:"is_album"`    // is_album	boolean	If it's an album or not
	// ** Uninplemented fields **
	// bandwidth	integer	Bandwidth consumed by the image in bytes
	// deletehash	string	OPTIONAL, the deletehash, if you're logged in as the image owner
	// comment_count	int	Number of comments on the gallery image.
	// topic_id	integer	Topic ID of the gallery image.
	// vote	string	The current user's vote on the album. null if not signed in or if the user hasn't voted on it.
	// favorite	boolean	Indicates if the current user favorited the image. Defaults to false if not signed in.
	// account_url	string	The username of the account that uploaded it, or null.
	// account_id	integer	The account ID of the account that uploaded it, or null.
	// ups	integer	Upvotes for the image
	// downs	integer	Number of downvotes for the image
	// points	integer	Upvotes minus downvotes
	// score	integer	Imgur popularity score
}

type imgurGalleryAlbum struct {
	ID          string              `json:"id"`           // id	string	The ID for the album
	Title       string              `json:"title"`        // title	string	The title of the album.
	Description string              `json:"description"`  // description	string	Description of the album.
	DateTime    int64               `json:"datetime"`     // datetime	integer	Time inserted into the gallery, epoch time
	Views       int64               `json:"views"`        // views	integer	The number of album views
	Link        string              `json:"link"`         // link	string	The URL link to the album
	NSFW        bool                `json:"nsfw"`         // nsfw	boolean	Indicates if the album has been marked as nsfw or not. Defaults to null if information is not available.
	Topic       string              `json:"topic"`        // topic	string	Topic of the gallery album.
	IsAlbum     bool                `json:"is_album"`     // is_album	boolean	If it's an album or not
	Cover       string              `json:"cover"`        // cover	string	The ID of the album cover image
	CoverWidth  int                 `json:"cover_width"`  // cover_width	integer	The width, in pixels, of the album cover image
	CoverHeight int                 `json:"cover_height"` // cover_height	integer	The height, in pixels, of the album cover image
	ImagesCount int                 `json:"images_count"` // images_count	integer	The total number of images in the album
	Images      []imgurGalleryImage `json:"images"`       // images	Array of Images	An array of all the images in the album (only available when requesting the direct album)

	// ** Unimplemented fields **
	// account_url	string	The account username or null if it's anonymous.
	// account_id	integer	The account ID of the account that uploaded it, or null.
	// privacy	string	The privacy level of the album, you can only view public if not logged in as album owner
	// layout	string	The view layout of the album.
	// views	integer	The number of image views
	// ups	integer	Upvotes for the image
	// downs	integer	Number of downvotes for the image
	// points	integer	Upvotes minus downvotes
	// score	integer	Imgur popularity score
	// vote	string	The current user's vote on the album. null if not signed in or if the user hasn't voted on it.
	// favorite	boolean	Indicates if the current user favorited the album. Defaults to false if not signed in.
	// comment_count	int	Number of comments on the gallery album.
	// topic_id	integer	Topic ID of the gallery album.
}

type imgurSearchResponse struct {
	Data    []json.RawMessage `json:"data"`
	Success bool              `json:"success"` // Request completed successfully
	Status  int               `json:"status"`  // HTTP response code
}

// Service contains the Config fields for the Imgur service.
//
// Example request:
//   {
//			"client_id": "AIzaSyA4FD39..."
//			"client_secret": "ASdsaijwdfASD..."
//   }
type Service struct {
	types.DefaultService
	// The Imgur client ID
	ClientID string `json:"client_id"`
	// The API key to use when making HTTP requests to Imgur.
	ClientSecret string `json:"client_secret"`
}

// Commands supported:
//    !imgur some_search_query_without_quotes
// Responds with a suitable image into the same room as the command.
func (s *Service) Commands(client *gomatrix.Client) []types.Command {
	return []types.Command{
		types.Command{
			Path: []string{"imgur", "help"},
			Command: func(roomID, userID string, args []string) (interface{}, error) {
				return usageMessage(), nil
			},
		},
		types.Command{
			Path: []string{"imgur"},
			Command: func(roomID, userID string, args []string) (interface{}, error) {
				return s.cmdImgurImgSearch(client, roomID, userID, args)
			},
		},
	}
}

// usageMessage returns a matrix TextMessage representation of the service usage
func usageMessage() *gomatrix.TextMessage {
	return &gomatrix.TextMessage{"m.notice",
		`Usage: !imgur image_search_text`}
}

func (s *Service) cmdImgurImgSearch(client *gomatrix.Client, roomID, userID string, args []string) (interface{}, error) {

	if len(args) < 1 {
		return usageMessage(), nil
	}

	// Get the query text to search for.
	querySentence := strings.Join(args, " ")

	searchResultImage, searchResultAlbum, err := s.text2imgImgur(querySentence)

	if err != nil {
		return nil, err
	}

	// Image returned
	if searchResultImage != nil {
		var imgURL = searchResultImage.Link
		if imgURL == "" {
			return gomatrix.TextMessage{
				MsgType: "m.notice",
				Body:    "No image found!",
			}, nil
		}

		// FIXME -- Sometimes upload fails with a cryptic error - "msg=Upload request failed code=400"
		// log.Printf("Uploading image at: %s", imgURL)
		resUpload, err := client.UploadLink(imgURL)
		if err != nil {
			return nil, fmt.Errorf("Failed to upload Imgur image (%s) to matrix: %s", imgURL, err.Error())
		}

		return gomatrix.ImageMessage{
			MsgType: "m.image",
			Body:    querySentence,
			URL:     resUpload.ContentURI,
			Info: gomatrix.ImageInfo{
				Height:   uint(searchResultImage.Height),
				Width:    uint(searchResultImage.Width),
				Mimetype: searchResultImage.Type,
			},
		}, nil
	} else if searchResultAlbum != nil {
		return gomatrix.TextMessage{
			MsgType: "m.notice",
			Body:    "Search returned an album - Not currently supported",
		}, nil
	} else {
		return gomatrix.TextMessage{
			MsgType: "m.notice",
			Body:    "No image found!",
		}, nil
	}
}

// text2imgImgur returns info about an image or an album
func (s *Service) text2imgImgur(query string) (*imgurGalleryImage, *imgurGalleryAlbum, error) {
	log.Info("Searching Imgur for an image of a ", query)

	query = url.QueryEscape(query)

	var sort = "time"  // time | viral | top
	var window = "all" // day | week | month | year | all
	var page = 1
	var urlString = fmt.Sprintf("%s/%s/%s/%d?q=%s", "https://api.imgur.com/3/gallery/search", sort, window, page, query)
	// var urlString = fmt.Sprintf("%s?q=%s", base, query)

	u, err := url.Parse(urlString)
	if err != nil {
		return nil, nil, err
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return nil, nil, err
	}

	req.Header.Add("Authorization", "Client-ID "+s.ClientID)
	res, err := httpClient.Do(req)
	if res != nil {
		defer res.Body.Close()
	}
	if err != nil {
		return nil, nil, err
	}
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		return nil, nil, fmt.Errorf("Request error: %d, %s", res.StatusCode, response2String(res))
	}

	var searchResults imgurSearchResponse
	if err := json.NewDecoder(res.Body).Decode(&searchResults); err != nil {
		return nil, nil, fmt.Errorf("No images found - %s", err.Error())
	} else if len(searchResults.Data) < 1 {
		return nil, nil, fmt.Errorf("No images found")
	}

	log.Printf("%d results were returned from Imgur", len(searchResults.Data))
	// Return a random image result
	var images []imgurGalleryImage
	for i := 0; i < len(searchResults.Data); i++ {
		var image imgurGalleryImage
		if err := json.Unmarshal(searchResults.Data[i], &image); err == nil && !image.IsAlbum {
			images = append(images, image)
		}
	}
	if len(images) > 0 {
		var r = 0
		if len(images) > 1 {
			r = rand.Intn(len(images) - 1)
		}
		return &images[r], nil, nil
	}

	return nil, nil, fmt.Errorf("No images found")
}

// response2String returns a string representation of an HTTP response body
func response2String(res *http.Response) string {
	bs, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "Failed to decode response body"
	}
	str := string(bs)
	return str
}

// Initialise the service
func init() {
	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
		return &Service{
			DefaultService: types.NewDefaultService(serviceID, serviceUserID, ServiceType),
		}
	})
}
<|endoftext|>