<commit_before>package archive

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strings"

	"compress/bzip2"
	"compress/gzip"
	"xi2.org/x/xz"

	"golang.org/x/crypto/openpgp"
	"pault.ag/go/debian/control"
	"pault.ag/go/debian/dependency"
)

// Common Types {{{

type Closer func() error
type PathReader func(path string) (io.Reader, Closer, error)

// known compression types {{{

type compressionReader func(io.Reader) (io.Reader, error)

func gzipNewReader(r io.Reader) (io.Reader, error) {
	return gzip.NewReader(r)
}

func xzNewReader(r io.Reader) (io.Reader, error) {
	return xz.NewReader(r, 0)
}

func bzipNewReader(r io.Reader) (io.Reader, error) {
	return bzip2.NewReader(r), nil
}

var knownCompressionAlgorithms = map[string]compressionReader{
	".gz":  gzipNewReader,
	".bz2": bzipNewReader,
	".xz":  xzNewReader,
}

// }}}

// }}}

// Archive {{{

type Archive struct {
	root       string
	pathReader PathReader
	keyring    *openpgp.EntityList
}

// Get an archive Suite for a given Archive {{{

func (a Archive) Suite(name string) (*Suite, error) {
	inReleasePath := path.Join("dists", name, "InRelease")
	reader, closer, err := a.getFile(inReleasePath)
	if err != nil {
		return nil, err
	}
	defer closer()

	release, err := LoadInRelease(reader, a.keyring)
	if err != nil {
		return nil, err
	}
	return &Suite{
		Release: *release,
		archive: a,
	}, nil
}

// }}}

// Suite {{{

// Suite struct {{{

type Suite struct {
	Release Release
	archive Archive
}

// }}}

// HasArch {{{

func (s Suite) HasArch(arch dependency.Arch) bool {
	for _, el := range s.Release.Architectures {
		if el.Is(&arch) {
			return true
		}
	}
	return false
}

// }}}

// HasComponent {{{

func (s Suite) HasComponent(component string) bool {
	for _, el := range s.Release.Components {
		if el == component {
			return true
		}
	}
	return false
}

// }}}

// getHashedFileReader - get a Hashed File {{{

func (s Suite) getHashedFileReader(suitePath string) (io.Reader, Closer, control.FileHashValidators, error) {
	packagesPath := path.Join("dists", s.Release.Suite, suitePath)
	reader, closer, err := s.archive.getFile(packagesPath)

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

	hashes := s.Release.Indices()[suitePath]
	if len(hashes) == 0 {
		closer()
		return nil, nil, nil, fmt.Errorf("Undeclared file in InRelease")
	}

	validators, err := hashes.Validators()
	if err != nil {
		closer()
		return nil, nil, nil, err
	}

	validationReader := io.TeeReader(reader, validators.Writer())
	return validationReader, closer, validators, err
}

// }}}

// Packages index for a Suite {{{

func (s Suite) Packages(component string, arch dependency.Arch) ([]Package, error) {
	ret := []Package{}

	if !s.HasComponent(component) {
		return ret, fmt.Errorf("No such component: '%s'", component)
	}
	if !s.HasArch(arch) {
		return ret, fmt.Errorf("No such arch: '%s'", arch.String())
	}
	suitePath := path.Join(
		component,
		fmt.Sprintf("binary-%s", arch.String()),
		"Packages",
	)

	validationReader, closer, validators, err := s.getHashedFileReader(suitePath)
	if err != nil {
		return []Package{}, err
	}
	defer closer()

	packages, err := LoadPackages(validationReader)
	if err != nil {
		return []Package{}, err
	}

	for {
		next, err := packages.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return []Package{}, err
		}
		ret = append(ret, *next)
	}

	if !validators.Validate() {
		return nil, fmt.Errorf("Index hashes don't match!")
	}

	return ret, nil
}

// }}}

// Sources index for a Suite {{{

func (s Suite) Sources(component string) ([]Source, error) {
	ret := []Source{}

	if !s.HasComponent(component) {
		return ret, fmt.Errorf("No such component: '%s'", component)
	}
	suitePath := path.Join(component, "source", "Sources.gz")

	validationReader, closer, validators, err := s.getHashedFileReader(suitePath)
	if err != nil {
		return []Source{}, err
	}
	defer closer()

	packages, err := LoadSources(validationReader)
	if err != nil {
		return []Source{}, err
	}

	for {
		next, err := packages.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return []Source{}, err
		}
		ret = append(ret, *next)
	}

	if !validators.Validate() {
		return nil, fmt.Errorf("Index hashes don't match!")
	}

	return ret, nil
}

// }}}

// }}}

// getFile wrapper {{{

func (a Archive) getFile(requestPath string) (io.Reader, Closer, error) {
	archivePath := a.root + "/" + requestPath
	fmt.Printf("%s\n", archivePath)
	reader, closer, err := a.pathReader(archivePath)
	if err != nil {
		return nil, nil, err
	}

	for suffix, decompressor := range knownCompressionAlgorithms {
		if strings.HasSuffix(requestPath, suffix) {
			newReader, err := decompressor(reader)
			if err != nil {
				closer()
				return nil, nil, err
			}
			return newReader, closer, nil
		}
	}

	return reader, closer, err
}

// }}}

// Constructors {{{

func NewArchive(
	root string,
	pathReader PathReader,
	keyring *openpgp.EntityList,
) Archive {
	return Archive{
		root:       root,
		pathReader: pathReader,
		keyring:    keyring,
	}
}

// Filesystem Archive {{{

func NewFilesystemArchive(root string, keyring *openpgp.EntityList) Archive {
	return NewArchive(root, filesystemPathReader, keyring)
}

func NewHttpArchive(root string, keyring *openpgp.EntityList) Archive {
	return NewArchive(root, httpPathReader, keyring)
}

// }}}

// }}}

// Readers {{{

func filesystemPathReader(path string) (io.Reader, Closer, error) {
	fd, err := os.Open(path)
	return fd, fd.Close, err
}

func httpPathReader(uri string) (io.Reader, Closer, error) {
	resp, err := http.Get(uri)
	if err != nil {
		return nil, nil, err
	}
	return resp.Body, resp.Body.Close, nil
}

// }}}

// }}}

// vim: foldmethod=marker
<commit_msg>Hash before the decompressor<commit_after>package archive

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strings"

	"compress/bzip2"
	"compress/gzip"
	"xi2.org/x/xz"

	"golang.org/x/crypto/openpgp"
	"pault.ag/go/debian/control"
	"pault.ag/go/debian/dependency"
)

// Common Types {{{

type Closer func() error
type PathReader func(path string) (io.Reader, Closer, error)

// known compression types {{{

type compressionReader func(io.Reader) (io.Reader, error)

func gzipNewReader(r io.Reader) (io.Reader, error) {
	return gzip.NewReader(r)
}

func xzNewReader(r io.Reader) (io.Reader, error) {
	return xz.NewReader(r, 0)
}

func bzipNewReader(r io.Reader) (io.Reader, error) {
	return bzip2.NewReader(r), nil
}

var knownCompressionAlgorithms = map[string]compressionReader{
	".gz":  gzipNewReader,
	".bz2": bzipNewReader,
	".xz":  xzNewReader,
}

// }}}

// }}}

// Archive {{{

type Archive struct {
	root       string
	pathReader PathReader
	keyring    *openpgp.EntityList
}

// Get an archive Suite for a given Archive {{{

func (a Archive) Suite(name string) (*Suite, error) {
	inReleasePath := path.Join("dists", name, "InRelease")
	reader, closer, err := a.getFile(inReleasePath, nil)
	if err != nil {
		return nil, err
	}
	defer closer()

	release, err := LoadInRelease(reader, a.keyring)
	if err != nil {
		return nil, err
	}
	return &Suite{
		Release: *release,
		archive: a,
	}, nil
}

// }}}

// Suite {{{

// Suite struct {{{

type Suite struct {
	Release Release
	archive Archive
}

// }}}

// HasArch {{{

func (s Suite) HasArch(arch dependency.Arch) bool {
	for _, el := range s.Release.Architectures {
		if el.Is(&arch) {
			return true
		}
	}
	return false
}

// }}}

// HasComponent {{{

func (s Suite) HasComponent(component string) bool {
	for _, el := range s.Release.Components {
		if el == component {
			return true
		}
	}
	return false
}

// }}}

// getHashedFileReader - get a Hashed File {{{

func (s Suite) getHashedFileReader(suitePath string) (io.Reader, Closer, control.FileHashValidators, error) {
	hashes := s.Release.Indices()[suitePath]
	if len(hashes) == 0 {
		return nil, nil, nil, fmt.Errorf("Undeclared file in InRelease")
	}

	validators, err := hashes.Validators()
	if err != nil {
		return nil, nil, nil, err
	}

	packagesPath := path.Join("dists", s.Release.Suite, suitePath)
	reader, closer, err := s.archive.getFile(packagesPath, validators.Writer())

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

	return reader, closer, validators, err
}

// }}}

// Packages index for a Suite {{{

func (s Suite) Packages(component string, arch dependency.Arch) ([]Package, error) {
	ret := []Package{}

	if !s.HasComponent(component) {
		return ret, fmt.Errorf("No such component: '%s'", component)
	}
	if !s.HasArch(arch) {
		return ret, fmt.Errorf("No such arch: '%s'", arch.String())
	}
	suitePath := path.Join(
		component,
		fmt.Sprintf("binary-%s", arch.String()),
		"Packages",
	)

	validationReader, closer, validators, err := s.getHashedFileReader(suitePath)
	if err != nil {
		return []Package{}, err
	}
	defer closer()

	packages, err := LoadPackages(validationReader)
	if err != nil {
		return []Package{}, err
	}

	for {
		next, err := packages.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return []Package{}, err
		}
		ret = append(ret, *next)
	}

	if !validators.Validate() {
		return nil, fmt.Errorf("Index hashes don't match!")
	}

	return ret, nil
}

// }}}

// Sources index for a Suite {{{

func (s Suite) Sources(component string) ([]Source, error) {
	ret := []Source{}

	if !s.HasComponent(component) {
		return ret, fmt.Errorf("No such component: '%s'", component)
	}
	suitePath := path.Join(component, "source", "Sources.gz")

	validationReader, closer, validators, err := s.getHashedFileReader(suitePath)
	if err != nil {
		return []Source{}, err
	}
	defer closer()

	packages, err := LoadSources(validationReader)
	if err != nil {
		return []Source{}, err
	}

	for {
		next, err := packages.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return []Source{}, err
		}
		ret = append(ret, *next)
	}

	if !validators.Validate() {
		return nil, fmt.Errorf("Index hashes don't match!")
	}

	return ret, nil
}

// }}}

// }}}

// getFile wrapper {{{

func (a Archive) getFile(requestPath string, tee io.Writer) (io.Reader, Closer, error) {
	archivePath := a.root + "/" + requestPath
	reader, closer, err := a.pathReader(archivePath)
	if err != nil {
		return nil, nil, err
	}

	if tee != nil {
		reader = io.TeeReader(reader, tee)
	}

	for suffix, decompressor := range knownCompressionAlgorithms {
		if strings.HasSuffix(requestPath, suffix) {
			newReader, err := decompressor(reader)
			if err != nil {
				closer()
				return nil, nil, err
			}
			return newReader, closer, nil
		}
	}

	return reader, closer, err
}

// }}}

// Constructors {{{

func NewArchive(
	root string,
	pathReader PathReader,
	keyring *openpgp.EntityList,
) Archive {
	return Archive{
		root:       root,
		pathReader: pathReader,
		keyring:    keyring,
	}
}

// Filesystem Archive {{{

func NewFilesystemArchive(root string, keyring *openpgp.EntityList) Archive {
	return NewArchive(root, filesystemPathReader, keyring)
}

func NewHttpArchive(root string, keyring *openpgp.EntityList) Archive {
	return NewArchive(root, httpPathReader, keyring)
}

// }}}

// }}}

// Readers {{{

func filesystemPathReader(path string) (io.Reader, Closer, error) {
	fd, err := os.Open(path)
	return fd, fd.Close, err
}

func httpPathReader(uri string) (io.Reader, Closer, error) {
	resp, err := http.Get(uri)
	if err != nil {
		return nil, nil, err
	}
	return resp.Body, resp.Body.Close, nil
}

// }}}

// }}}

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


<commit_before>package archive

import (
	"golang.org/x/crypto/openpgp"

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

// 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
}

// 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
}

// Get a Suite for a given Archive. Information will be loaded from the
// InRelease file (if it exists) into the Suite object.
//
// The Suite object contains neat abstractions such as Components, and a
// number of helpers to collection information across all components, such
// as the `Arches()` helper.
func (a Archive) Suite(name string) (*Suite, error) {
	/* Get the Release / InRelease */
	suite := Suite{
		archive:            &a,
		packageCollections: map[string]PackageCollections{},
	}

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

	return &suite, nil
}

// Use the default backend to remove any unlinked files from the Blob store.
func (a Archive) Decruft() error {
	return a.store.GC(blobstore.DumbGarbageCollector{})
}

type Suite struct {
	control.Paragraph

	archive *Archive

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

	/* Componenet        ~= PackageCollections
	 * Compoenent + Arch ~= PackageCollection */
	packageCollections map[string]PackageCollections `control:"-"`
	Pool               Pool                          `control:"-"`

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

func (s Suite) PackageCollections(name string) PackageCollections {
	if _, ok := s.packageCollections[name]; !ok {
		s.packageCollections[name] = PackageCollections{}
	}
	return s.packageCollections[name]
}

type PackageCollections map[dependency.Arch]PackageCollection

func (p PackageCollections) get(arch dependency.Arch) PackageCollection {
	if _, ok := p[arch]; !ok {
		p[arch] = PackageCollection{}
	}
	return p[arch]
}

func (p PackageCollections) Add(pkg Package) error {
	collection := p.get(pkg.Architecture)
	return collection.Add(pkg)
}

// Component is a section of the Archive, which is a set of Binary packages
// that are provided to the end user. Debian has three major ones, `main`,
// `contrib` and `non-free`.
type PackageCollection struct {
	encoder *control.Encoder
	writer  *blobstore.Writer
	archive *Archive
	/* XXX: Add a flag to ensure alpha sorted entries */
}

func (p PackageCollection) Add(pkg Package) error {
	return p.encoder.Encode(pkg)
}

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

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

	"golang.org/x/crypto/openpgp"

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

// 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
}

// 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
}

func (a Archive) Suite(name string) (*Suite, error) {
	suite := Suite{
		archive:    &a,
		components: map[string]Component{},
	}

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

	return &suite, nil
}

// Use the default backend to remove any unlinked files from the Blob store.
func (a Archive) Decruft() error {
	return a.store.GC(blobstore.DumbGarbageCollector{})
}

func (a Archive) Link(blobs map[string]blobstore.Object) error {
	for path, obj := range blobs {
		if err := a.store.Link(obj, path); err != nil {
			return err
		}
	}
	return nil
}

func (suite Suite) newRelease() Release {
	when := time.Now()
	release := Release{
		Suite:       suite.Name,
		Description: suite.Description,
		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
}

//
func (a Archive) Engross(suite Suite) (map[string]blobstore.Object, error) {
	// release := suite.newRelease()

	files := map[string]blobstore.Object{}

	for name, component := range suite.components {
		for arch, writer := range component.packageWriters {

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

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

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

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

	return files, nil
}

//
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:"-"`
}

/////////////////////////////////////////////////////////

func (s Suite) Component(name string) (*Component, error) {
	if _, ok := s.components[name]; !ok {
		comp, err := newComponent(s.archive)
		if err != nil {
			return nil, err
		}
		s.components[name] = *comp
		return comp, nil
	}
	el := s.components[name]
	return &el, nil
}

func newComponent(archive *Archive) (*Component, error) {
	return &Component{
		archive:        archive,
		packageWriters: map[dependency.Arch]*PackageWriter{},
	}, nil
}

type Component struct {
	archive        *Archive
	packageWriters map[dependency.Arch]*PackageWriter
	// sourceWriter *SourceWriter
}

func (c *Component) getWriter(arch dependency.Arch) (*PackageWriter, error) {
	if _, ok := c.packageWriters[arch]; !ok {
		writer, err := newPackageWriter(c.archive)
		if err != nil {
			return nil, err
		}
		c.packageWriters[arch] = writer
	}
	return c.packageWriters[arch], nil
}

func (c *Component) AddPackage(pkg Package) error {
	writer, err := c.getWriter(pkg.Architecture)
	if err != nil {
		return err
	}
	return writer.Add(pkg)
}

type packageWriters map[dependency.Arch]*PackageWriter

func newPackageWriter(archive *Archive) (*PackageWriter, error) {

	/* So, a PackageWriter is the thing we use to create a Packages entry
	 * for a given suite/component/binary-arch */

	handle, err := archive.store.Create()
	if err != nil {
		return nil, err
	}

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

	return &PackageWriter{
		archive: archive,
		writer:  handle,
		closer:  handle.Close,
		encoder: encoder,
		handle:  handle,
	}, nil
}

type PackageWriter struct {
	archive *Archive

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

func (p PackageWriter) Add(pkg Package) error {
	return p.encoder.Encode(pkg)
}

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


<commit_before>package torrent

import (
	"time"

	"github.com/cenkalti/rain/internal/metainfo"
)

var (
	publicPeerIDPrefix                    = "-RN" + Version + "-"
	publicExtensionHandshakeClientVersion = "Rain " + Version
	trackerHTTPPublicUserAgent            = "Rain/" + Version
)

func init() {
	metainfo.Creator = publicExtensionHandshakeClientVersion
}

// Config for Session.
type Config struct {
	// Database file to save resume data.
	Database string
	// DataDir is where files are downloaded.
	DataDir string
	// If true, torrent files are saved into <data_dir>/<torrent_id>/<torrent_name>.
	// Useful if downloading the same torrent from multiple sources.
	DataDirIncludesTorrentID bool
	// New torrents will be listened at selected port in this range.
	PortBegin, PortEnd uint16
	// At start, client will set max open files limit to this number. (like "ulimit -n" command)
	MaxOpenFiles uint64
	// Enable peer exchange protocol.
	PEXEnabled bool
	// Resume data (bitfield & stats) are saved to disk at interval to keep IO lower.
	ResumeWriteInterval time.Duration
	// Peer id is prefixed with this string. See BEP 20. Remaining bytes of peer id will be randomized.
	// Only applies to private torrents.
	PrivatePeerIDPrefix string
	// Client version that is sent in BEP 10 handshake message.
	// Only applies to private torrents.
	PrivateExtensionHandshakeClientVersion string
	// URL to the blocklist file in CIDR format.
	BlocklistURL string
	// When to refresh blocklist
	BlocklistUpdateInterval time.Duration
	// HTTP timeout for downloading blocklist
	BlocklistUpdateTimeout time.Duration
	// Do not contact tracker if it's IP is blocked
	BlocklistEnabledForTrackers bool
	// Do not connect to peer if it's IP is blocked
	BlocklistEnabledForOutgoingConnections bool
	// Do not accept connections from peer if it's IP is blocked
	BlocklistEnabledForIncomingConnections bool
	// Do not accept response larger than this size
	BlocklistMaxResponseSize int64
	// Time to wait when adding torrent with AddURI().
	TorrentAddHTTPTimeout time.Duration
	// Maximum allowed size to be received by metadata extension.
	MaxMetadataSize uint
	// Maximum allowed size to be read when adding torrent.
	MaxTorrentSize uint
	// Maximum allowed number of pieces in a torrent.
	MaxPieces uint32
	// Time to wait when resolving host names for trackers and peers.
	DNSResolveTimeout time.Duration
	// Global download speed limit in KB/s.
	SpeedLimitDownload int64
	// Global upload speed limit in KB/s.
	SpeedLimitUpload int64
	// Start torrent automatically if it was running when previous session was closed.
	ResumeOnStartup bool

	// Enable RPC server
	RPCEnabled bool
	// Host to listen for RPC server
	RPCHost string
	// Listen port for RPC server
	RPCPort int
	// Time to wait for ongoing requests before shutting down RPC HTTP server.
	RPCShutdownTimeout time.Duration

	// Enable DHT node.
	DHTEnabled bool
	// DHT node will listen on this IP.
	DHTHost string
	// DHT node will listen on this UDP port.
	DHTPort uint16
	// DHT announce interval
	DHTAnnounceInterval time.Duration
	// Minimum announce interval when announcing to DHT.
	DHTMinAnnounceInterval time.Duration
	// Known routers to bootstrap local DHT node.
	DHTBootstrapNodes []string

	// Number of peer addresses to request in announce request.
	TrackerNumWant int
	// Time to wait for announcing stopped event.
	// Stopped event is sent to the tracker when torrent is stopped.
	TrackerStopTimeout time.Duration
	// When the client needs new peer addresses to connect, it ask to the tracker.
	// To prevent spamming the tracker an interval is set to wait before the next announce.
	TrackerMinAnnounceInterval time.Duration
	// Total time to wait for response to be read.
	// This includes ConnectTimeout and TLSHandshakeTimeout.
	TrackerHTTPTimeout time.Duration
	// User agent sent when communicating with HTTP trackers.
	// Only applies to private torrents.
	TrackerHTTPPrivateUserAgent string
	// Max number of bytes in a tracker response.
	TrackerHTTPMaxResponseSize uint
	// Check and validate TLS ceritificates.
	TrackerHTTPVerifyTLS bool

	// Number of unchoked peers.
	UnchokedPeers int
	// Number of optimistic unchoked peers.
	OptimisticUnchokedPeers int
	// Max number of blocks allowed to be queued without dropping any.
	MaxRequestsIn int
	// Max number of blocks requested from a peer but not received yet.
	// `rreq` value from extended handshake cannot exceed this limit.
	MaxRequestsOut int
	// Number of bloks requested from peer if it does not send `rreq` value in extended handshake.
	DefaultRequestsOut int
	// Time to wait for a requested block to be received before marking peer as snubbed
	RequestTimeout time.Duration
	// Max number of running downloads on piece in endgame mode, snubbed and choed peers don't count
	EndgameMaxDuplicateDownloads int
	// Max number of outgoing connections to dial
	MaxPeerDial int
	// Max number of incoming connections to accept
	MaxPeerAccept int
	// Running metadata downloads, snubbed peers don't count
	ParallelMetadataDownloads int
	// Time to wait for TCP connection to open.
	PeerConnectTimeout time.Duration
	// Time to wait for BitTorrent handshake to complete.
	PeerHandshakeTimeout time.Duration
	// When peer has started to send piece block, if it does not send any bytes in PieceReadTimeout, the connection is closed.
	PieceReadTimeout time.Duration
	// Max number of peer addresses to keep in connect queue.
	MaxPeerAddresses int
	// Number of allowed-fast messages to send after handshake.
	AllowedFastSet int

	// Number of bytes to read when a piece is requested by a peer.
	ReadCacheBlockSize int64
	// Number of cached bytes for piece read requests.
	ReadCacheSize int64
	// Read bytes for a piece part expires after duration.
	ReadCacheTTL time.Duration
	// Number of read operations to do in parallel.
	ParallelReads uint
	// Number of write operations to do in parallel.
	ParallelWrites uint
	// Number of bytes allocated in memory for downloading piece data.
	WriteCacheSize int64

	// When the client want to connect a peer, first it tries to do encrypted handshake.
	// If it does not work, it connects to same peer again and does unencrypted handshake.
	// This behavior can be changed via this variable.
	DisableOutgoingEncryption bool
	// Dial only encrypted connections.
	ForceOutgoingEncryption bool
	// Do not accept unencrypted connections.
	ForceIncomingEncryption bool

	// TCP connect timeout for WebSeed sources
	WebseedDialTimeout time.Duration
	// TLS handshake timeout for WebSeed sources
	WebseedTLSHandshakeTimeout time.Duration
	// HTTP header timeout for WebSeed sources
	WebseedResponseHeaderTimeout time.Duration
	// HTTP body read timeout for Webseed sources
	WebseedResponseBodyReadTimeout time.Duration
	// Retry interval for restarting failed downloads
	WebseedRetryInterval time.Duration
	// Verify TLS certificate for WebSeed URLs
	WebseedVerifyTLS bool
	// Limit the number of WebSeed sources in torrent.
	WebseedMaxSources int
	// Number of maximum simulateous downloads from WebSeed sources.
	WebseedMaxDownloads int
}

// DefaultConfig for Session. Do not pass zero value Config to NewSession. Copy this struct and modify instead.
var DefaultConfig = Config{
	// Session
	Database:                               "~/rain/session.db",
	DataDir:                                "~/rain/data",
	DataDirIncludesTorrentID:               true,
	PortBegin:                              50000,
	PortEnd:                                60000,
	MaxOpenFiles:                           10240,
	PEXEnabled:                             true,
	ResumeWriteInterval:                    30 * time.Second,
	PrivatePeerIDPrefix:                    "-RN" + Version + "-",
	PrivateExtensionHandshakeClientVersion: "Rain " + Version,
	BlocklistUpdateInterval:                24 * time.Hour,
	BlocklistUpdateTimeout:                 10 * time.Minute,
	BlocklistEnabledForTrackers:            true,
	BlocklistEnabledForOutgoingConnections: true,
	BlocklistEnabledForIncomingConnections: true,
	BlocklistMaxResponseSize:               100 << 20,
	TorrentAddHTTPTimeout:                  30 * time.Second,
	MaxMetadataSize:                        10 << 20,
	MaxTorrentSize:                         10 << 20,
	MaxPieces:                              64 << 10,
	DNSResolveTimeout:                      5 * time.Second,
	ResumeOnStartup:                        true,

	// RPC Server
	RPCEnabled:         true,
	RPCHost:            "127.0.0.1",
	RPCPort:            7246,
	RPCShutdownTimeout: 5 * time.Second,

	// Tracker
	TrackerNumWant:              200,
	TrackerStopTimeout:          5 * time.Second,
	TrackerMinAnnounceInterval:  time.Minute,
	TrackerHTTPTimeout:          10 * time.Second,
	TrackerHTTPPrivateUserAgent: "Rain/" + Version,
	TrackerHTTPMaxResponseSize:  2 << 20,
	TrackerHTTPVerifyTLS:        true,

	// DHT node
	DHTEnabled:             true,
	DHTHost:                "0.0.0.0",
	DHTPort:                7246,
	DHTAnnounceInterval:    30 * time.Minute,
	DHTMinAnnounceInterval: time.Minute,
	DHTBootstrapNodes: []string{
		"router.bittorrent.com:6881",
		"dht.transmissionbt.com:6881",
		"router.utorrent.com:6881",
		"dht.libtorrent.org:25401",
		"dht.aelitis.com:6881",
	},

	// Peer
	UnchokedPeers:                3,
	OptimisticUnchokedPeers:      1,
	MaxRequestsIn:                250,
	MaxRequestsOut:               250,
	DefaultRequestsOut:           50,
	RequestTimeout:               20 * time.Second,
	EndgameMaxDuplicateDownloads: 20,
	MaxPeerDial:                  80,
	MaxPeerAccept:                20,
	ParallelMetadataDownloads:    2,
	PeerConnectTimeout:           5 * time.Second,
	PeerHandshakeTimeout:         10 * time.Second,
	PieceReadTimeout:             30 * time.Second,
	MaxPeerAddresses:             2000,
	AllowedFastSet:               10,

	// IO
	ReadCacheBlockSize: 128 << 10,
	ReadCacheSize:      256 << 20,
	ReadCacheTTL:       1 * time.Minute,
	ParallelReads:      1,
	ParallelWrites:     1,
	WriteCacheSize:     1 << 30,

	// Webseed settings
	WebseedDialTimeout:             10 * time.Second,
	WebseedTLSHandshakeTimeout:     10 * time.Second,
	WebseedResponseHeaderTimeout:   10 * time.Second,
	WebseedResponseBodyReadTimeout: 10 * time.Second,
	WebseedRetryInterval:           time.Minute,
	WebseedVerifyTLS:               true,
	WebseedMaxSources:              10,
	WebseedMaxDownloads:            4,
}
<commit_msg>increase default max metadata size<commit_after>package torrent

import (
	"time"

	"github.com/cenkalti/rain/internal/metainfo"
)

var (
	publicPeerIDPrefix                    = "-RN" + Version + "-"
	publicExtensionHandshakeClientVersion = "Rain " + Version
	trackerHTTPPublicUserAgent            = "Rain/" + Version
)

func init() {
	metainfo.Creator = publicExtensionHandshakeClientVersion
}

// Config for Session.
type Config struct {
	// Database file to save resume data.
	Database string
	// DataDir is where files are downloaded.
	DataDir string
	// If true, torrent files are saved into <data_dir>/<torrent_id>/<torrent_name>.
	// Useful if downloading the same torrent from multiple sources.
	DataDirIncludesTorrentID bool
	// New torrents will be listened at selected port in this range.
	PortBegin, PortEnd uint16
	// At start, client will set max open files limit to this number. (like "ulimit -n" command)
	MaxOpenFiles uint64
	// Enable peer exchange protocol.
	PEXEnabled bool
	// Resume data (bitfield & stats) are saved to disk at interval to keep IO lower.
	ResumeWriteInterval time.Duration
	// Peer id is prefixed with this string. See BEP 20. Remaining bytes of peer id will be randomized.
	// Only applies to private torrents.
	PrivatePeerIDPrefix string
	// Client version that is sent in BEP 10 handshake message.
	// Only applies to private torrents.
	PrivateExtensionHandshakeClientVersion string
	// URL to the blocklist file in CIDR format.
	BlocklistURL string
	// When to refresh blocklist
	BlocklistUpdateInterval time.Duration
	// HTTP timeout for downloading blocklist
	BlocklistUpdateTimeout time.Duration
	// Do not contact tracker if it's IP is blocked
	BlocklistEnabledForTrackers bool
	// Do not connect to peer if it's IP is blocked
	BlocklistEnabledForOutgoingConnections bool
	// Do not accept connections from peer if it's IP is blocked
	BlocklistEnabledForIncomingConnections bool
	// Do not accept response larger than this size
	BlocklistMaxResponseSize int64
	// Time to wait when adding torrent with AddURI().
	TorrentAddHTTPTimeout time.Duration
	// Maximum allowed size to be received by metadata extension.
	MaxMetadataSize uint
	// Maximum allowed size to be read when adding torrent.
	MaxTorrentSize uint
	// Maximum allowed number of pieces in a torrent.
	MaxPieces uint32
	// Time to wait when resolving host names for trackers and peers.
	DNSResolveTimeout time.Duration
	// Global download speed limit in KB/s.
	SpeedLimitDownload int64
	// Global upload speed limit in KB/s.
	SpeedLimitUpload int64
	// Start torrent automatically if it was running when previous session was closed.
	ResumeOnStartup bool

	// Enable RPC server
	RPCEnabled bool
	// Host to listen for RPC server
	RPCHost string
	// Listen port for RPC server
	RPCPort int
	// Time to wait for ongoing requests before shutting down RPC HTTP server.
	RPCShutdownTimeout time.Duration

	// Enable DHT node.
	DHTEnabled bool
	// DHT node will listen on this IP.
	DHTHost string
	// DHT node will listen on this UDP port.
	DHTPort uint16
	// DHT announce interval
	DHTAnnounceInterval time.Duration
	// Minimum announce interval when announcing to DHT.
	DHTMinAnnounceInterval time.Duration
	// Known routers to bootstrap local DHT node.
	DHTBootstrapNodes []string

	// Number of peer addresses to request in announce request.
	TrackerNumWant int
	// Time to wait for announcing stopped event.
	// Stopped event is sent to the tracker when torrent is stopped.
	TrackerStopTimeout time.Duration
	// When the client needs new peer addresses to connect, it ask to the tracker.
	// To prevent spamming the tracker an interval is set to wait before the next announce.
	TrackerMinAnnounceInterval time.Duration
	// Total time to wait for response to be read.
	// This includes ConnectTimeout and TLSHandshakeTimeout.
	TrackerHTTPTimeout time.Duration
	// User agent sent when communicating with HTTP trackers.
	// Only applies to private torrents.
	TrackerHTTPPrivateUserAgent string
	// Max number of bytes in a tracker response.
	TrackerHTTPMaxResponseSize uint
	// Check and validate TLS ceritificates.
	TrackerHTTPVerifyTLS bool

	// Number of unchoked peers.
	UnchokedPeers int
	// Number of optimistic unchoked peers.
	OptimisticUnchokedPeers int
	// Max number of blocks allowed to be queued without dropping any.
	MaxRequestsIn int
	// Max number of blocks requested from a peer but not received yet.
	// `rreq` value from extended handshake cannot exceed this limit.
	MaxRequestsOut int
	// Number of bloks requested from peer if it does not send `rreq` value in extended handshake.
	DefaultRequestsOut int
	// Time to wait for a requested block to be received before marking peer as snubbed
	RequestTimeout time.Duration
	// Max number of running downloads on piece in endgame mode, snubbed and choed peers don't count
	EndgameMaxDuplicateDownloads int
	// Max number of outgoing connections to dial
	MaxPeerDial int
	// Max number of incoming connections to accept
	MaxPeerAccept int
	// Running metadata downloads, snubbed peers don't count
	ParallelMetadataDownloads int
	// Time to wait for TCP connection to open.
	PeerConnectTimeout time.Duration
	// Time to wait for BitTorrent handshake to complete.
	PeerHandshakeTimeout time.Duration
	// When peer has started to send piece block, if it does not send any bytes in PieceReadTimeout, the connection is closed.
	PieceReadTimeout time.Duration
	// Max number of peer addresses to keep in connect queue.
	MaxPeerAddresses int
	// Number of allowed-fast messages to send after handshake.
	AllowedFastSet int

	// Number of bytes to read when a piece is requested by a peer.
	ReadCacheBlockSize int64
	// Number of cached bytes for piece read requests.
	ReadCacheSize int64
	// Read bytes for a piece part expires after duration.
	ReadCacheTTL time.Duration
	// Number of read operations to do in parallel.
	ParallelReads uint
	// Number of write operations to do in parallel.
	ParallelWrites uint
	// Number of bytes allocated in memory for downloading piece data.
	WriteCacheSize int64

	// When the client want to connect a peer, first it tries to do encrypted handshake.
	// If it does not work, it connects to same peer again and does unencrypted handshake.
	// This behavior can be changed via this variable.
	DisableOutgoingEncryption bool
	// Dial only encrypted connections.
	ForceOutgoingEncryption bool
	// Do not accept unencrypted connections.
	ForceIncomingEncryption bool

	// TCP connect timeout for WebSeed sources
	WebseedDialTimeout time.Duration
	// TLS handshake timeout for WebSeed sources
	WebseedTLSHandshakeTimeout time.Duration
	// HTTP header timeout for WebSeed sources
	WebseedResponseHeaderTimeout time.Duration
	// HTTP body read timeout for Webseed sources
	WebseedResponseBodyReadTimeout time.Duration
	// Retry interval for restarting failed downloads
	WebseedRetryInterval time.Duration
	// Verify TLS certificate for WebSeed URLs
	WebseedVerifyTLS bool
	// Limit the number of WebSeed sources in torrent.
	WebseedMaxSources int
	// Number of maximum simulateous downloads from WebSeed sources.
	WebseedMaxDownloads int
}

// DefaultConfig for Session. Do not pass zero value Config to NewSession. Copy this struct and modify instead.
var DefaultConfig = Config{
	// Session
	Database:                               "~/rain/session.db",
	DataDir:                                "~/rain/data",
	DataDirIncludesTorrentID:               true,
	PortBegin:                              50000,
	PortEnd:                                60000,
	MaxOpenFiles:                           10240,
	PEXEnabled:                             true,
	ResumeWriteInterval:                    30 * time.Second,
	PrivatePeerIDPrefix:                    "-RN" + Version + "-",
	PrivateExtensionHandshakeClientVersion: "Rain " + Version,
	BlocklistUpdateInterval:                24 * time.Hour,
	BlocklistUpdateTimeout:                 10 * time.Minute,
	BlocklistEnabledForTrackers:            true,
	BlocklistEnabledForOutgoingConnections: true,
	BlocklistEnabledForIncomingConnections: true,
	BlocklistMaxResponseSize:               100 << 20,
	TorrentAddHTTPTimeout:                  30 * time.Second,
	MaxMetadataSize:                        30 << 20,
	MaxTorrentSize:                         10 << 20,
	MaxPieces:                              64 << 10,
	DNSResolveTimeout:                      5 * time.Second,
	ResumeOnStartup:                        true,

	// RPC Server
	RPCEnabled:         true,
	RPCHost:            "127.0.0.1",
	RPCPort:            7246,
	RPCShutdownTimeout: 5 * time.Second,

	// Tracker
	TrackerNumWant:              200,
	TrackerStopTimeout:          5 * time.Second,
	TrackerMinAnnounceInterval:  time.Minute,
	TrackerHTTPTimeout:          10 * time.Second,
	TrackerHTTPPrivateUserAgent: "Rain/" + Version,
	TrackerHTTPMaxResponseSize:  2 << 20,
	TrackerHTTPVerifyTLS:        true,

	// DHT node
	DHTEnabled:             true,
	DHTHost:                "0.0.0.0",
	DHTPort:                7246,
	DHTAnnounceInterval:    30 * time.Minute,
	DHTMinAnnounceInterval: time.Minute,
	DHTBootstrapNodes: []string{
		"router.bittorrent.com:6881",
		"dht.transmissionbt.com:6881",
		"router.utorrent.com:6881",
		"dht.libtorrent.org:25401",
		"dht.aelitis.com:6881",
	},

	// Peer
	UnchokedPeers:                3,
	OptimisticUnchokedPeers:      1,
	MaxRequestsIn:                250,
	MaxRequestsOut:               250,
	DefaultRequestsOut:           50,
	RequestTimeout:               20 * time.Second,
	EndgameMaxDuplicateDownloads: 20,
	MaxPeerDial:                  80,
	MaxPeerAccept:                20,
	ParallelMetadataDownloads:    2,
	PeerConnectTimeout:           5 * time.Second,
	PeerHandshakeTimeout:         10 * time.Second,
	PieceReadTimeout:             30 * time.Second,
	MaxPeerAddresses:             2000,
	AllowedFastSet:               10,

	// IO
	ReadCacheBlockSize: 128 << 10,
	ReadCacheSize:      256 << 20,
	ReadCacheTTL:       1 * time.Minute,
	ParallelReads:      1,
	ParallelWrites:     1,
	WriteCacheSize:     1 << 30,

	// Webseed settings
	WebseedDialTimeout:             10 * time.Second,
	WebseedTLSHandshakeTimeout:     10 * time.Second,
	WebseedResponseHeaderTimeout:   10 * time.Second,
	WebseedResponseBodyReadTimeout: 10 * time.Second,
	WebseedRetryInterval:           time.Minute,
	WebseedVerifyTLS:               true,
	WebseedMaxSources:              10,
	WebseedMaxDownloads:            4,
}
<|endoftext|>


<commit_before>package extensions

import (
	"bytes"
	cryptoRand "crypto/rand"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode"
)

type Version struct {
	Major          int
	Minor          int
	Revision       int
	MajorString    string
	MinorString    string
	RevisionString string
	Value          string
}

type FilePath struct {
	Name string `json:"Name"`
	Path string `json:"Path"`
	Type string `json:"Type"`
}

func (obj *FilePath) ToString() (str string) {
	return obj.Path + string(os.PathSeparator) + obj.Name + " | " + obj.Type
}

func GetDecimalAndStringFromHex(twoDigitHexCode string) (decimalValue int, asciiString string, err error) {
	decoded, err := hex.DecodeString(twoDigitHexCode)
	if err != nil {
		return
	}
  	decimalValue = int(decoded[0])
	asciiString = string(string(decoded))
	return
}

/*
* leftPad and rightPad just repoeat the padStr the indicated
* number of times
*
 */
func LeftPad(s string, padStr string, pLen int) string {
	return strings.Repeat(padStr, pLen) + s
}
func RightPad(s string, padStr string, pLen int) string {
	return s + strings.Repeat(padStr, pLen)
}

/* the Pad2Len functions are generally assumed to be padded with short sequences of strings
* in many cases with a single character sequence
*
* so we assume we can build the string out as if the char seq is 1 char and then
* just substr the string if it is longer than needed
*
* this means we are wasting some cpu and memory work
* but this always get us to want we want it to be
*
* in short not optimized to for massive string work
*
* If the overallLen is shorter than the original string length
* the string will be shortened to this length (substr)
*
 */
func RightPad2Len(s string, padStr string, overallLen int) string {
	var padCountInt int
	padCountInt = 1 + ((overallLen - len(padStr)) / len(padStr))
	var retStr = s + strings.Repeat(padStr, padCountInt)
	return retStr[:overallLen]
}

func LeftPad2Len(s string, padStr string, overallLen int) string {
	var padCountInt int
	padCountInt = 1 + ((overallLen - len(padStr)) / len(padStr))
	var retStr = strings.Repeat(padStr, padCountInt) + s
	return retStr[(len(retStr) - overallLen):]
}

func TrimSuffix(s, suffix string) string {
	if strings.HasSuffix(s, suffix) {
		s = s[:len(s)-len(suffix)]
	}
	return s
}

func PrintKiloBytes(bytes int64) string {

	var kilobytes float64
	kilobytes = float64(bytes / 1024)

	return fmt.Sprint(FloatToString(kilobytes, 2), " kB")
}

func PrintMegaBytes(bytes int64) string {

	var kilobytes float64
	kilobytes = float64(bytes / 1024)

	var megabytes float64
	megabytes = kilobytes / 1024 // cast to type float64

	return fmt.Sprint(FloatToString(megabytes, 2), " MB")
}

func IsPrintable(s string) bool {
	for _, c := range s {
		if (c < 32 || c > 126) && c != 10 && c != 13 && c != 9 {
			return false
		}
	}
	return true
}

func PrintZettaBytes(bytes int64) string {

	var kilobytes float64
	kilobytes = float64(bytes / 1024)

	var megabytes float64
	megabytes = (kilobytes / 1024) // cast to type float64

	var gigabytes float64
	gigabytes = (megabytes / 1024)

	var terabytes float64
	terabytes = (gigabytes / 1024)

	var petabytes float64
	petabytes = (terabytes / 1024)

	var exabytes float64
	exabytes = (petabytes / 1024)

	var zettabytes float64
	zettabytes = (exabytes / 1024)

	return fmt.Sprint(FloatToString(zettabytes, 2), " ZB")
}

func FloatToString(input_num float64, decimals int) string {
	// to convert a float number to a string
	return strconv.FormatFloat(input_num, 'f', decimals, 64)
}

func Round(x, unit float64) float64 {
	return float64(int64(x/unit+0.5)) * unit
}

func StringToInt(val string) int {

	r, err := strconv.Atoi(val)
	if err != nil {
		return 0
	}
	return r
}

func StringToUInt64(val string) uint64 {
	i, err := strconv.ParseUint(val, 10, 64)
	if err != nil {
		return 0
	}
	return i
}

func IntToString(val int) string {
	return strconv.Itoa(val)
}

func IntToBool(val int) bool {
	return val != 0
}

func Int32ToString(val int32) string {
	return strconv.Itoa(Int32ToInt(val))
}

func Int64ToString(val int64) string {
	return strconv.FormatInt(val, 10)
}

func Int64ToInt32(val int64) (ret int) {
	tempLong := ((val >> 32) << 32) //shift it right then left 32 bits, which zeroes the lower half of the long
	ret = (int)(val - tempLong)
	return ret
}

func Int32ToInt(val int32) (ret int) {
	tempLong := ((val >> 32) << 32) //shift it right then left 32 bits, which zeroes the lower half of the long
	ret = (int)(val - tempLong)
	return ret
}

// Int32ToDWORD converts an int32 to DWORD byte array.
func Int32ToDWORD(val int32) (value []byte) {
	buff := new(bytes.Buffer)
	err := binary.Write(buff, binary.BigEndian, val)
	if err != nil {
		return
	}
	value = buff.Bytes()
	return
}

// DWordToInt converts an 4 byte to a int32
func DWordToInt(data []byte) (num int32) {
	buff := bytes.NewReader(data)
	err := binary.Read(buff, binary.BigEndian, &num)
	if err != nil {
		return
	}

	return
}

func BoolToString(val bool) string {
	return strconv.FormatBool(val)
}

func StringToBool(val string) bool {
	r, err := strconv.ParseBool(val)
	if err != nil {
		return false
	}
	return r
}

func (obj *Version) Init(value string) {
	versionInfo := strings.Split(value, ".")

	obj.MajorString = versionInfo[0]
	obj.MinorString = versionInfo[1]
	obj.RevisionString = versionInfo[2]
	obj.Value = value

	if val, err := strconv.Atoi(versionInfo[0]); err == nil {
		obj.Major = val
	}

	if val, err := strconv.Atoi(versionInfo[1]); err == nil {
		obj.Minor = val
	}

	if val, err := strconv.Atoi(versionInfo[2]); err == nil {
		obj.Revision = val
	}
}

func GenPackageImport(name string, imports []string) string {

	val := "package " + name + "\n\n"
	val += "import(\n"
	for _, imp := range imports {
		if imp == "" {
			continue
		}
		val += "\t\"" + imp + "\"\n"
	}
	val += ")\n\n"

	return val
}

func MakeFirstLowerCase(s string) string {

	if len(s) < 2 {
		return strings.ToLower(s)
	}

	bts := []byte(s)

	lc := bytes.ToLower([]byte{bts[0]})
	rest := bts[1:]

	return string(bytes.Join([][]byte{lc, rest}, nil))
}

func ExtractArgsWithinBrackets(str string) (res []string) {

	brackets := &unicode.RangeTable{
		R16: []unicode.Range16{
			// {0x0028, 0x0029, 1}, // ( )
			// {0x005b, 0x005d, 1}, // [ ]
			{0x007b, 0x007d, 1}, // { }
		},
	}

	isBracket := func(r rune) bool {
		if unicode.In(r, brackets) {
			return true
		}
		return false
	}

	res = strings.FieldsFunc(str, isBracket)
	return
}

func Random(min, max int) int {
	rand.Seed(time.Now().UTC().UnixNano())
	return rand.Intn(max-min) + min
}

//Between will return substring between two strings
func Between(value string, a string, b string) string {
	// Get substring between two strings.
	posFirst := strings.Index(value, a)
	if posFirst == -1 {
		return ""
	}
	posLast := strings.Index(value, b)
	if posLast == -1 {
		return ""
	}
	posFirstAdjusted := posFirst + len(a)
	if posFirstAdjusted >= posLast {
		return ""
	}
	return value[posFirstAdjusted:posLast]
}

//SyncMapAny will return true if there are any items in the sync.Map
func SyncMapAny(x *sync.Map) (ok bool) {
	x.Range(func(key interface{}, value interface{}) bool {
		ok = true
		return false
	})
	return
}

//SyncMapLength will return true the length of items in the sync.Map
func SyncMapLength(x *sync.Map) (length int) {
	x.Range(func(key interface{}, value interface{}) bool {
		length = length + 1
		return true
	})
	return
}

// NewUUID generates a random UUID according to RFC 4122
func NewUUID() (string, error) {
	uuid := make([]byte, 16)
	n, err := io.ReadFull(cryptoRand.Reader, uuid)
	if n != len(uuid) || err != nil {
		return "", err
	}
	// variant bits; see section 4.1.1
	uuid[8] = uuid[8]&^0xc0 | 0x80
	// version 4 (pseudo-random); see section 4.1.3
	uuid[6] = uuid[6]&^0xf0 | 0x40
	return fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
}

//RandomString returns a random string of length
func RandomString(strlen int) string {
	rand.Seed(time.Now().UTC().UnixNano())
	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
	result := make([]byte, strlen)
	for i := 0; i < strlen; i++ {
		result[i] = chars[rand.Intn(len(chars))]
	}
	return string(result)
}

//BigEndianUInt16 will convert a uint16 to 2 bytes
func BigEndianUInt16(value uint16) (value1 byte, value2 byte) {
	inject := make([]byte, 2)
	binary.BigEndian.PutUint16(inject, value)
	value1 = inject[0]
	value2 = inject[1]
	return
}

//BigEndianUInt32 will convert a uint32 to 4 bytes
func BigEndianUInt32(value uint32) (value1 byte, value2 byte, value3 byte, value4 byte) {
	inject := make([]byte, 4)
	binary.BigEndian.PutUint32(inject, value)
	value1 = inject[0]
	value2 = inject[1]
	value3 = inject[2]
	value4 = inject[3]
	return
}

//IpAddressToUint32 will convert a string ip address to UInt32
func IpAddressToUint32(ip string) (value uint32) {
	binary.Read(bytes.NewBuffer(net.ParseIP(ip).To4()), binary.BigEndian, &value)
	return
}
<commit_msg>Float func<commit_after>package extensions

import (
	"bytes"
	cryptoRand "crypto/rand"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode"
)

type Version struct {
	Major          int
	Minor          int
	Revision       int
	MajorString    string
	MinorString    string
	RevisionString string
	Value          string
}

type FilePath struct {
	Name string `json:"Name"`
	Path string `json:"Path"`
	Type string `json:"Type"`
}

func (obj *FilePath) ToString() (str string) {
	return obj.Path + string(os.PathSeparator) + obj.Name + " | " + obj.Type
}

func GetDecimalAndStringFromHex(twoDigitHexCode string) (decimalValue int, asciiString string, err error) {
	decoded, err := hex.DecodeString(twoDigitHexCode)
	if err != nil {
		return
	}
	decimalValue = int(decoded[0])
	asciiString = string(string(decoded))
	return
}

/*
* leftPad and rightPad just repoeat the padStr the indicated
* number of times
*
 */
func LeftPad(s string, padStr string, pLen int) string {
	return strings.Repeat(padStr, pLen) + s
}
func RightPad(s string, padStr string, pLen int) string {
	return s + strings.Repeat(padStr, pLen)
}

/* the Pad2Len functions are generally assumed to be padded with short sequences of strings
* in many cases with a single character sequence
*
* so we assume we can build the string out as if the char seq is 1 char and then
* just substr the string if it is longer than needed
*
* this means we are wasting some cpu and memory work
* but this always get us to want we want it to be
*
* in short not optimized to for massive string work
*
* If the overallLen is shorter than the original string length
* the string will be shortened to this length (substr)
*
 */
func RightPad2Len(s string, padStr string, overallLen int) string {
	var padCountInt int
	padCountInt = 1 + ((overallLen - len(padStr)) / len(padStr))
	var retStr = s + strings.Repeat(padStr, padCountInt)
	return retStr[:overallLen]
}

func LeftPad2Len(s string, padStr string, overallLen int) string {
	var padCountInt int
	padCountInt = 1 + ((overallLen - len(padStr)) / len(padStr))
	var retStr = strings.Repeat(padStr, padCountInt) + s
	return retStr[(len(retStr) - overallLen):]
}

func TrimSuffix(s, suffix string) string {
	if strings.HasSuffix(s, suffix) {
		s = s[:len(s)-len(suffix)]
	}
	return s
}

func PrintKiloBytes(bytes int64) string {

	var kilobytes float64
	kilobytes = float64(bytes / 1024)

	return fmt.Sprint(FloatToString(kilobytes, 2), " kB")
}

func PrintMegaBytes(bytes int64) string {

	var kilobytes float64
	kilobytes = float64(bytes / 1024)

	var megabytes float64
	megabytes = kilobytes / 1024 // cast to type float64

	return fmt.Sprint(FloatToString(megabytes, 2), " MB")
}

func IsPrintable(s string) bool {
	for _, c := range s {
		if (c < 32 || c > 126) && c != 10 && c != 13 && c != 9 {
			return false
		}
	}
	return true
}

func PrintZettaBytes(bytes int64) string {

	var kilobytes float64
	kilobytes = float64(bytes / 1024)

	var megabytes float64
	megabytes = (kilobytes / 1024) // cast to type float64

	var gigabytes float64
	gigabytes = (megabytes / 1024)

	var terabytes float64
	terabytes = (gigabytes / 1024)

	var petabytes float64
	petabytes = (terabytes / 1024)

	var exabytes float64
	exabytes = (petabytes / 1024)

	var zettabytes float64
	zettabytes = (exabytes / 1024)

	return fmt.Sprint(FloatToString(zettabytes, 2), " ZB")
}

func FloatToString(input_num float64, decimals int) string {
	// to convert a float number to a string
	return strconv.FormatFloat(input_num, 'f', decimals, 64)
}

func Round(x, unit float64) float64 {
	return float64(int64(x/unit+0.5)) * unit
}

func StringToInt(val string) int {

	r, err := strconv.Atoi(val)
	if err != nil {
		return 0
	}
	return r
}

func StringToFloat(val string, bitSize int) (r float64) {
	r, err := strconv.ParseFloat(val, bitSize)
	if err != nil {
		return 0
	}
	return r
}

func StringToUInt64(val string) uint64 {
	i, err := strconv.ParseUint(val, 10, 64)
	if err != nil {
		return 0
	}
	return i
}

func IntToString(val int) string {
	return strconv.Itoa(val)
}

func IntToBool(val int) bool {
	return val != 0
}

func Int32ToString(val int32) string {
	return strconv.Itoa(Int32ToInt(val))
}

func Int64ToString(val int64) string {
	return strconv.FormatInt(val, 10)
}

func Int64ToInt32(val int64) (ret int) {
	tempLong := ((val >> 32) << 32) //shift it right then left 32 bits, which zeroes the lower half of the long
	ret = (int)(val - tempLong)
	return ret
}

func Int32ToInt(val int32) (ret int) {
	tempLong := ((val >> 32) << 32) //shift it right then left 32 bits, which zeroes the lower half of the long
	ret = (int)(val - tempLong)
	return ret
}

// Int32ToDWORD converts an int32 to DWORD byte array.
func Int32ToDWORD(val int32) (value []byte) {
	buff := new(bytes.Buffer)
	err := binary.Write(buff, binary.BigEndian, val)
	if err != nil {
		return
	}
	value = buff.Bytes()
	return
}

// DWordToInt converts an 4 byte to a int32
func DWordToInt(data []byte) (num int32) {
	buff := bytes.NewReader(data)
	err := binary.Read(buff, binary.BigEndian, &num)
	if err != nil {
		return
	}

	return
}

func BoolToString(val bool) string {
	return strconv.FormatBool(val)
}

func StringToBool(val string) bool {
	r, err := strconv.ParseBool(val)
	if err != nil {
		return false
	}
	return r
}

func (obj *Version) Init(value string) {
	versionInfo := strings.Split(value, ".")

	obj.MajorString = versionInfo[0]
	obj.MinorString = versionInfo[1]
	obj.RevisionString = versionInfo[2]
	obj.Value = value

	if val, err := strconv.Atoi(versionInfo[0]); err == nil {
		obj.Major = val
	}

	if val, err := strconv.Atoi(versionInfo[1]); err == nil {
		obj.Minor = val
	}

	if val, err := strconv.Atoi(versionInfo[2]); err == nil {
		obj.Revision = val
	}
}

func GenPackageImport(name string, imports []string) string {

	val := "package " + name + "\n\n"
	val += "import(\n"
	for _, imp := range imports {
		if imp == "" {
			continue
		}
		val += "\t\"" + imp + "\"\n"
	}
	val += ")\n\n"

	return val
}

func MakeFirstLowerCase(s string) string {

	if len(s) < 2 {
		return strings.ToLower(s)
	}

	bts := []byte(s)

	lc := bytes.ToLower([]byte{bts[0]})
	rest := bts[1:]

	return string(bytes.Join([][]byte{lc, rest}, nil))
}

func ExtractArgsWithinBrackets(str string) (res []string) {

	brackets := &unicode.RangeTable{
		R16: []unicode.Range16{
			// {0x0028, 0x0029, 1}, // ( )
			// {0x005b, 0x005d, 1}, // [ ]
			{0x007b, 0x007d, 1}, // { }
		},
	}

	isBracket := func(r rune) bool {
		if unicode.In(r, brackets) {
			return true
		}
		return false
	}

	res = strings.FieldsFunc(str, isBracket)
	return
}

func Random(min, max int) int {
	rand.Seed(time.Now().UTC().UnixNano())
	return rand.Intn(max-min) + min
}

//Between will return substring between two strings
func Between(value string, a string, b string) string {
	// Get substring between two strings.
	posFirst := strings.Index(value, a)
	if posFirst == -1 {
		return ""
	}
	posLast := strings.Index(value, b)
	if posLast == -1 {
		return ""
	}
	posFirstAdjusted := posFirst + len(a)
	if posFirstAdjusted >= posLast {
		return ""
	}
	return value[posFirstAdjusted:posLast]
}

//SyncMapAny will return true if there are any items in the sync.Map
func SyncMapAny(x *sync.Map) (ok bool) {
	x.Range(func(key interface{}, value interface{}) bool {
		ok = true
		return false
	})
	return
}

//SyncMapLength will return true the length of items in the sync.Map
func SyncMapLength(x *sync.Map) (length int) {
	x.Range(func(key interface{}, value interface{}) bool {
		length = length + 1
		return true
	})
	return
}

// NewUUID generates a random UUID according to RFC 4122
func NewUUID() (string, error) {
	uuid := make([]byte, 16)
	n, err := io.ReadFull(cryptoRand.Reader, uuid)
	if n != len(uuid) || err != nil {
		return "", err
	}
	// variant bits; see section 4.1.1
	uuid[8] = uuid[8]&^0xc0 | 0x80
	// version 4 (pseudo-random); see section 4.1.3
	uuid[6] = uuid[6]&^0xf0 | 0x40
	return fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
}

//RandomString returns a random string of length
func RandomString(strlen int) string {
	rand.Seed(time.Now().UTC().UnixNano())
	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
	result := make([]byte, strlen)
	for i := 0; i < strlen; i++ {
		result[i] = chars[rand.Intn(len(chars))]
	}
	return string(result)
}

//BigEndianUInt16 will convert a uint16 to 2 bytes
func BigEndianUInt16(value uint16) (value1 byte, value2 byte) {
	inject := make([]byte, 2)
	binary.BigEndian.PutUint16(inject, value)
	value1 = inject[0]
	value2 = inject[1]
	return
}

//BigEndianUInt32 will convert a uint32 to 4 bytes
func BigEndianUInt32(value uint32) (value1 byte, value2 byte, value3 byte, value4 byte) {
	inject := make([]byte, 4)
	binary.BigEndian.PutUint32(inject, value)
	value1 = inject[0]
	value2 = inject[1]
	value3 = inject[2]
	value4 = inject[3]
	return
}

//IpAddressToUint32 will convert a string ip address to UInt32
func IpAddressToUint32(ip string) (value uint32) {
	binary.Read(bytes.NewBuffer(net.ParseIP(ip).To4()), binary.BigEndian, &value)
	return
}
<|endoftext|>


<commit_before><commit_msg>ran gofmt<commit_after><|endoftext|>


<commit_before>package main

import (
	"io"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"path"
	"stathat.com/c/consistent"
	"strings"
)

// Sattolo shuffle
// https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#Sattolo.27s_algorithm
func shuffle(list []string) {
	for i := range list {
		j := rand.Intn(i + 1)
		list[i], list[j] = list[j], list[i]
	}
}

type CacheProxy struct {
	backends *consistent.Consistent
}

func New() *CacheProxy {
	c := new(CacheProxy)
	c.backends = consistent.New()
	return c
}

func (c *CacheProxy) Manager() *http.ServeMux {
	manager := http.NewServeMux()

	// RESTful interface for cache members
	// GET /members/
	// 	  Returns a list of current cache members

	// PUT /members/member[:port]
	//    Adds "member" as a cache member

	// DELETE /members/member[:port]
	//    Delets "member" from cache ring

	manager.HandleFunc("/members/", func(w http.ResponseWriter, req *http.Request) {
		switch {
		case req.Method == "PUT":
			c.backends.Add(path.Base(req.URL.Path))
		case req.Method == "DELETE":
			c.backends.Remove(path.Base(req.URL.Path))
		}
		io.WriteString(w, strings.Join(c.backends.Members(), ","))
	})

	return manager
}

func (c *CacheProxy) Select(path string) []string {
	selection, _ := c.backends.GetN(path, 3)
	shuffle(selection)
	return selection
}

// Reverse proxy that selects the backend by nearest match to the request URL
// on the consistent hash ring.
func (c *CacheProxy) Proxy() *httputil.ReverseProxy {
	return &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			req.URL.Scheme = "http"
			backends := c.Select(req.URL.Path)
			req.URL.Host = backends[0]
			req.Header.Set("X-Backends", strings.Join(backends, ","))
		},
		Transport: &http.Transport{},
	}
}

func main() {
	proxy := New()

	// Replace with runtime arguments for initialization
	proxy.backends.Add("localhost:9091")
	proxy.backends.Add("localhost:9092")
	proxy.backends.Add("localhost:9093")
	proxy.backends.Add("localhost:9094")
	proxy.backends.Add("localhost:9095")

	// Initialize and run manager server in background via goroutine
	go http.ListenAndServe(":9190", proxy.Manager())

	http.ListenAndServe(":9090", proxy.Proxy())
}
<commit_msg>Removing unnecessary CacheProxy methods for HTTP servers<commit_after>package main

import (
	"io"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"path"
	"stathat.com/c/consistent"
	"strings"
)

// Sattolo shuffle
// https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#Sattolo.27s_algorithm
func shuffle(list []string) {
	for i := range list {
		j := rand.Intn(i + 1)
		list[i], list[j] = list[j], list[i]
	}
}

type CacheProxy struct {
	backends *consistent.Consistent
	manager  *http.ServeMux
	proxy    *httputil.ReverseProxy
}

func New() *CacheProxy {
	c := new(CacheProxy)
	c.backends = consistent.New()

	// RESTful interface for cache members
	c.manager = http.NewServeMux()
	c.manager.HandleFunc("/members/", func(w http.ResponseWriter, req *http.Request) {
		switch {
		case req.Method == "PUT":
			// Adds "member" as a cache member
			c.backends.Add(path.Base(req.URL.Path))
		case req.Method == "DELETE":
			// Deletes "member" from cache ring
			c.backends.Remove(path.Base(req.URL.Path))
		}
		// Always returns a list of current cache members
		io.WriteString(w, strings.Join(c.backends.Members(), ","))
	})

	// Reverse proxy that selects the backend by nearest match to the request URL
	// on the consistent hash ring.
	c.proxy = &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			req.URL.Scheme = "http"
			backends := c.Select(req.URL.Path)
			req.URL.Host = backends[0]
			req.Header.Set("X-Backends", strings.Join(backends, ","))
		},
		Transport: &http.Transport{},
	}

	return c
}

func (c *CacheProxy) Select(path string) []string {
	selection, _ := c.backends.GetN(path, 3)
	shuffle(selection)
	return selection
}

func main() {
	proxy := New()

	// Replace with runtime arguments for initialization
	proxy.backends.Add("localhost:9091")
	proxy.backends.Add("localhost:9092")
	proxy.backends.Add("localhost:9093")
	proxy.backends.Add("localhost:9094")
	proxy.backends.Add("localhost:9095")

	// Initialize and run manager server in background via goroutine
	go http.ListenAndServe(":9190", proxy.manager)

	http.ListenAndServe(":9090", proxy.proxy)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"sync"

	"github.com/Sirupsen/logrus"
	"gopkg.in/tomb.v1"

	"net"
)

// Proxy represents the proxy in its entirity with all its links. The main
// responsibility of Proxy is to accept new client and create Links between the
// client and upstream.
//
// Client <-> toxiproxy <-> Upstream
//
type Proxy struct {
	sync.Mutex

	Name     string `json:"name"`
	Listen   string `json:"listen"`
	Upstream string `json:"upstream"`
	Enabled  bool   `json:"enabled"`

	started chan error

	tomb        tomb.Tomb
	connections map[string]net.Conn
	upToxics    *ToxicCollection
	downToxics  *ToxicCollection
}

var ErrProxyAlreadyStarted = errors.New("Proxy already started")

func NewProxy() *Proxy {
	proxy := &Proxy{
		started:     make(chan error),
		connections: make(map[string]net.Conn),
	}
	proxy.upToxics = NewToxicCollection(proxy)
	proxy.downToxics = NewToxicCollection(proxy)
	return proxy
}

func (proxy *Proxy) Start() error {
	proxy.Lock()
	defer proxy.Unlock()

	if proxy.Enabled {
		return ErrProxyAlreadyStarted
	}
	proxy.Enabled = true

	go proxy.server()
	err := <-proxy.started
	// Disable the proxy again if it failed to start
	proxy.Enabled = err == nil
	return err
}

// server runs the Proxy server, accepting new clients and creating Links to
// connect them to upstreams.
func (proxy *Proxy) server() {
	ln, err := net.Listen("tcp", proxy.Listen)
	if err != nil {
		proxy.started <- err
		return
	}

	proxy.Listen = ln.Addr().String()
	proxy.started <- nil

	logrus.WithFields(logrus.Fields{
		"name":     proxy.Name,
		"proxy":    proxy.Listen,
		"upstream": proxy.Upstream,
	}).Info("Started proxy")

	acceptTomb := tomb.Tomb{}
	defer acceptTomb.Done()

	// This channel is to kill the blocking Accept() call below by closing the
	// net.Listener.
	go func() {
		<-proxy.tomb.Dying()

		// Notify ln.Accept() that the shutdown was safe
		acceptTomb.Killf("Shutting down from stop()")
		// Unblock ln.Accept()
		err := ln.Close()
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"proxy":  proxy.Name,
				"listen": proxy.Listen,
				"err":    err,
			}).Warn("Attempted to close an already closed proxy server")
		}

		// Wait for the accept loop to finish processing
		acceptTomb.Wait()
		proxy.tomb.Done()
	}()

	for {
		client, err := ln.Accept()
		if err != nil {
			// This is to confirm we're being shut down in a legit way. Unfortunately,
			// Go doesn't export the error when it's closed from Close() so we have to
			// sync up with a channel here.
			//
			// See http://zhen.org/blog/graceful-shutdown-of-go-net-dot-listeners/
			select {
			case <-acceptTomb.Dying():
			default:
				logrus.WithFields(logrus.Fields{
					"proxy":  proxy.Name,
					"listen": proxy.Listen,
					"err":    err,
				}).Warn("Error while accepting client")
			}
			return
		}

		logrus.WithFields(logrus.Fields{
			"name":     proxy.Name,
			"client":   client.RemoteAddr(),
			"proxy":    proxy.Listen,
			"upstream": proxy.Upstream,
		}).Info("Accepted client")

		upstream, err := net.Dial("tcp", proxy.Upstream)
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"name":     proxy.Name,
				"client":   client.RemoteAddr(),
				"proxy":    proxy.Listen,
				"upstream": proxy.Upstream,
			}).Error("Unable to open connection to upstream")
			client.Close()
			continue
		}

		name := client.RemoteAddr().String()
		proxy.Lock()
		proxy.connections[name+"client"] = client
		proxy.connections[name+"upstream"] = upstream
		proxy.Unlock()
		proxy.upToxics.StartLink(name+"client", client, upstream)
		proxy.downToxics.StartLink(name+"upstream", upstream, client)
	}
}

func (proxy *Proxy) RemoveConnection(name string) {
	proxy.Lock()
	defer proxy.Unlock()
	delete(proxy.connections, name)
}

func (proxy *Proxy) Stop() {
	proxy.Lock()
	defer proxy.Unlock()
	if !proxy.Enabled {
		return
	}
	proxy.Enabled = false
	proxy.Unlock()

	proxy.tomb.Killf("Shutting down from stop()")
	proxy.tomb.Wait() // Wait until we stop accepting new connections

	proxy.Lock()
	for _, conn := range proxy.connections {
		conn.Close()
	}

	logrus.WithFields(logrus.Fields{
		"name":     proxy.Name,
		"proxy":    proxy.Listen,
		"upstream": proxy.Upstream,
	}).Info("Terminated proxy")
}
<commit_msg>Change Lock/Unlock in proxy.Stop()<commit_after>package main

import (
	"errors"
	"sync"

	"github.com/Sirupsen/logrus"
	"gopkg.in/tomb.v1"

	"net"
)

// Proxy represents the proxy in its entirity with all its links. The main
// responsibility of Proxy is to accept new client and create Links between the
// client and upstream.
//
// Client <-> toxiproxy <-> Upstream
//
type Proxy struct {
	sync.Mutex

	Name     string `json:"name"`
	Listen   string `json:"listen"`
	Upstream string `json:"upstream"`
	Enabled  bool   `json:"enabled"`

	started chan error

	tomb        tomb.Tomb
	connections map[string]net.Conn
	upToxics    *ToxicCollection
	downToxics  *ToxicCollection
}

var ErrProxyAlreadyStarted = errors.New("Proxy already started")

func NewProxy() *Proxy {
	proxy := &Proxy{
		started:     make(chan error),
		connections: make(map[string]net.Conn),
	}
	proxy.upToxics = NewToxicCollection(proxy)
	proxy.downToxics = NewToxicCollection(proxy)
	return proxy
}

func (proxy *Proxy) Start() error {
	proxy.Lock()
	defer proxy.Unlock()

	if proxy.Enabled {
		return ErrProxyAlreadyStarted
	}
	proxy.Enabled = true

	go proxy.server()
	err := <-proxy.started
	// Disable the proxy again if it failed to start
	proxy.Enabled = err == nil
	return err
}

// server runs the Proxy server, accepting new clients and creating Links to
// connect them to upstreams.
func (proxy *Proxy) server() {
	ln, err := net.Listen("tcp", proxy.Listen)
	if err != nil {
		proxy.started <- err
		return
	}

	proxy.Listen = ln.Addr().String()
	proxy.started <- nil

	logrus.WithFields(logrus.Fields{
		"name":     proxy.Name,
		"proxy":    proxy.Listen,
		"upstream": proxy.Upstream,
	}).Info("Started proxy")

	acceptTomb := tomb.Tomb{}
	defer acceptTomb.Done()

	// This channel is to kill the blocking Accept() call below by closing the
	// net.Listener.
	go func() {
		<-proxy.tomb.Dying()

		// Notify ln.Accept() that the shutdown was safe
		acceptTomb.Killf("Shutting down from stop()")
		// Unblock ln.Accept()
		err := ln.Close()
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"proxy":  proxy.Name,
				"listen": proxy.Listen,
				"err":    err,
			}).Warn("Attempted to close an already closed proxy server")
		}

		// Wait for the accept loop to finish processing
		acceptTomb.Wait()
		proxy.tomb.Done()
	}()

	for {
		client, err := ln.Accept()
		if err != nil {
			// This is to confirm we're being shut down in a legit way. Unfortunately,
			// Go doesn't export the error when it's closed from Close() so we have to
			// sync up with a channel here.
			//
			// See http://zhen.org/blog/graceful-shutdown-of-go-net-dot-listeners/
			select {
			case <-acceptTomb.Dying():
			default:
				logrus.WithFields(logrus.Fields{
					"proxy":  proxy.Name,
					"listen": proxy.Listen,
					"err":    err,
				}).Warn("Error while accepting client")
			}
			return
		}

		logrus.WithFields(logrus.Fields{
			"name":     proxy.Name,
			"client":   client.RemoteAddr(),
			"proxy":    proxy.Listen,
			"upstream": proxy.Upstream,
		}).Info("Accepted client")

		upstream, err := net.Dial("tcp", proxy.Upstream)
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"name":     proxy.Name,
				"client":   client.RemoteAddr(),
				"proxy":    proxy.Listen,
				"upstream": proxy.Upstream,
			}).Error("Unable to open connection to upstream")
			client.Close()
			continue
		}

		name := client.RemoteAddr().String()
		proxy.Lock()
		proxy.connections[name+"client"] = client
		proxy.connections[name+"upstream"] = upstream
		proxy.Unlock()
		proxy.upToxics.StartLink(name+"client", client, upstream)
		proxy.downToxics.StartLink(name+"upstream", upstream, client)
	}
}

func (proxy *Proxy) RemoveConnection(name string) {
	proxy.Lock()
	defer proxy.Unlock()
	delete(proxy.connections, name)
}

func (proxy *Proxy) Stop() {
	proxy.Lock()
	if !proxy.Enabled {
		proxy.Unlock()
		return
	}
	proxy.Enabled = false
	proxy.Unlock()

	proxy.tomb.Killf("Shutting down from stop()")
	proxy.tomb.Wait() // Wait until we stop accepting new connections

	proxy.Lock()
	defer proxy.Unlock()
	for _, conn := range proxy.connections {
		conn.Close()
	}

	logrus.WithFields(logrus.Fields{
		"name":     proxy.Name,
		"proxy":    proxy.Listen,
		"upstream": proxy.Upstream,
	}).Info("Terminated proxy")
}
<|endoftext|>


<commit_before>// Copyright 2015 The rkt 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 main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/appc/spec/schema/types"
	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/coreos/gexpect"
	"github.com/coreos/rkt/tests/testutils"
)

const (
	rmImageReferenced = `rkt: image %q is referenced by some containers, cannot remove.`
	rmImageOk         = "rkt: successfully removed aci for image:"

	unreferencedACI = "rkt-unreferencedACI.aci"
	unreferencedApp = "coreos.com/rkt-unreferenced"
	referencedApp   = "coreos.com/rkt-inspect"

	stage1App = "coreos.com/rkt/stage1"
)

func TestImageRunRmName(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmd = fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), referencedACI)
	spawnAndWaitOrFail(t, cmd, true)

	t.Logf("Retrieving stage1 image name")
	stage1ImageName := getImageName(t, ctx, stage1App)

	t.Logf("Removing stage1 image (should work)")
	if err := removeImage(ctx, stage1ImageName, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, referencedApp, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, unreferencedApp, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestImageRunRmID(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	t.Logf("Fetching %s", imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmd = fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), referencedACI)
	t.Logf("Running %s", referencedACI)
	spawnAndWaitOrFail(t, cmd, true)

	t.Logf("Retrieving stage1 image ID")
	stage1ImageID, err := getImageId(ctx, stage1App)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", referencedApp)
	referencedImageID, err := getImageId(ctx, referencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", unreferencedApp)
	unreferencedImageID, err := getImageId(ctx, unreferencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Removing stage1 image (should work)")
	if err := removeImage(ctx, stage1ImageID, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, referencedImageID, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, unreferencedImageID, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestImagePrepareRmNameRun(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	t.Logf("Fetching %s", imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmds := strings.Fields(ctx.Cmd())
	prepareCmd := exec.Command(cmds[0], cmds[1:]...)
	prepareCmd.Args = append(prepareCmd.Args, "--insecure-options=image", "prepare", referencedACI)
	output, err := prepareCmd.Output()
	if err != nil {
		t.Fatalf("Cannot read the output: %v", err)
	}

	podIDStr := strings.TrimSpace(string(output))
	podID, err := types.NewUUID(podIDStr)
	if err != nil {
		t.Fatalf("%q is not a valid UUID: %v", podIDStr, err)
	}

	t.Logf("Retrieving stage1 image name")
	stage1ImageName := getImageName(t, ctx, stage1App)

	t.Logf("Removing stage1 image (should work)")
	if err := removeImage(ctx, stage1ImageName, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, referencedApp, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, unreferencedApp, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	cmd = fmt.Sprintf("%s run-prepared --mds-register=false %s", ctx.Cmd(), podID.String())
	t.Logf("Running %s", referencedACI)
	spawnAndWaitOrFail(t, cmd, true)
}

func TestImagePrepareRmIDRun(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	t.Logf("Fetching %s", imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmds := strings.Fields(ctx.Cmd())
	prepareCmd := exec.Command(cmds[0], cmds[1:]...)
	prepareCmd.Args = append(prepareCmd.Args, "--insecure-options=image", "prepare", referencedACI)
	output, err := prepareCmd.Output()
	if err != nil {
		t.Fatalf("Cannot read the output: %v", err)
	}

	podIDStr := strings.TrimSpace(string(output))
	podID, err := types.NewUUID(podIDStr)
	if err != nil {
		t.Fatalf("%q is not a valid UUID: %v", podIDStr, err)
	}

	t.Logf("Retrieving stage1 imageID")
	stage1ImageID, err := getImageId(ctx, stage1App)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", referencedApp)
	referencedImageID, err := getImageId(ctx, referencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", unreferencedApp)
	unreferencedImageID, err := getImageId(ctx, unreferencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Removing stage1 image (should work)")
	if err := removeImage(ctx, stage1ImageID, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, referencedImageID, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, unreferencedImageID, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	cmd = fmt.Sprintf("%s run-prepared --mds-register=false %s", ctx.Cmd(), podID.String())
	t.Logf("Running %s", referencedACI)
	spawnAndWaitOrFail(t, cmd, true)
}

func getImageName(t *testing.T, ctx *testutils.RktRunCtx, name string) string {
	cmd := fmt.Sprintf(`/bin/sh -c "%s image list --fields=name --no-legend | grep %s | cut -d: -f1"`, ctx.Cmd(), name)
	child := spawnOrFail(t, cmd)
	imageName, err := child.ReadLine()
	imageName = strings.TrimSpace(imageName)
	imageName = string(bytes.Trim([]byte(imageName), "\x00"))
	if err != nil {
		t.Fatalf("Cannot exec: %v", err)
	}
	waitOrFail(t, child, true)
	return imageName
}

func getImageId(ctx *testutils.RktRunCtx, name string) (string, error) {
	cmd := fmt.Sprintf(`/bin/sh -c "%s image list --fields=id,name --no-legend | grep %s | awk '{print $1}'"`, ctx.Cmd(), name)
	child, err := gexpect.Spawn(cmd)
	if err != nil {
		return "", fmt.Errorf("Cannot exec rkt: %v", err)
	}
	imageID, err := child.ReadLine()
	imageID = strings.TrimSpace(imageID)
	imageID = string(bytes.Trim([]byte(imageID), "\x00"))
	if err != nil {
		return "", fmt.Errorf("Cannot exec: %v", err)
	}
	if err := child.Wait(); err != nil {
		return "", fmt.Errorf("rkt didn't terminate correctly: %v", err)
	}
	return imageID, nil
}

func removeImage(ctx *testutils.RktRunCtx, image string, shouldWork bool) error {
	expect := fmt.Sprintf(rmImageReferenced, image)
	if shouldWork {
		expect = rmImageOk
	}
	cmd := fmt.Sprintf("%s image rm %s", ctx.Cmd(), image)
	child, err := gexpect.Spawn(cmd)
	if err != nil {
		return fmt.Errorf("Cannot exec: %v", err)
	}
	if err := expectWithOutput(child, expect); err != nil {
		return fmt.Errorf("Expected %q but not found: %v", expect, err)
	}
	if err := child.Wait(); err != nil {
		return fmt.Errorf("rkt didn't terminate correctly: %v", err)
	}
	return nil
}
<commit_msg>tests: add test for removing duplicate images<commit_after>// Copyright 2015 The rkt 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 main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/appc/spec/schema/types"
	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/coreos/gexpect"
	"github.com/coreos/rkt/tests/testutils"
)

const (
	rmImageReferenced = `rkt: image %q is referenced by some containers, cannot remove.`
	rmImageOk         = "rkt: successfully removed aci for image:"

	unreferencedACI = "rkt-unreferencedACI.aci"
	unreferencedApp = "coreos.com/rkt-unreferenced"
	referencedApp   = "coreos.com/rkt-inspect"

	stage1App = "coreos.com/rkt/stage1"
)

func TestImageRunRmName(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmd = fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), referencedACI)
	spawnAndWaitOrFail(t, cmd, true)

	t.Logf("Retrieving stage1 image name")
	stage1ImageName := getImageName(t, ctx, stage1App)

	t.Logf("Removing stage1 image (should work)")
	if err := removeImage(ctx, []string{stage1ImageName}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, []string{referencedApp}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, []string{unreferencedApp}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestImageRunRmID(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	t.Logf("Fetching %s", imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmd = fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), referencedACI)
	t.Logf("Running %s", referencedACI)
	spawnAndWaitOrFail(t, cmd, true)

	t.Logf("Retrieving stage1 image ID")
	stage1ImageID, err := getImageId(ctx, stage1App)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", referencedApp)
	referencedImageID, err := getImageId(ctx, referencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", unreferencedApp)
	unreferencedImageID, err := getImageId(ctx, unreferencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Removing stage1 image (should work)")
	if err := removeImage(ctx, []string{stage1ImageID}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, []string{referencedImageID}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, []string{unreferencedImageID}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestImageRunRmDuplicate(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	t.Logf("Fetching %s", imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmd = fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), referencedACI)
	t.Logf("Running %s", referencedACI)
	spawnAndWaitOrFail(t, cmd, true)

	t.Logf("Retrieving %s image ID", referencedApp)
	referencedImageID, err := getImageId(ctx, referencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", unreferencedApp)
	unreferencedImageID, err := getImageId(ctx, unreferencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, []string{referencedApp, referencedImageID}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, []string{unreferencedImageID, unreferencedApp}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestImagePrepareRmNameRun(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	t.Logf("Fetching %s", imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmds := strings.Fields(ctx.Cmd())
	prepareCmd := exec.Command(cmds[0], cmds[1:]...)
	prepareCmd.Args = append(prepareCmd.Args, "--insecure-options=image", "prepare", referencedACI)
	output, err := prepareCmd.Output()
	if err != nil {
		t.Fatalf("Cannot read the output: %v", err)
	}

	podIDStr := strings.TrimSpace(string(output))
	podID, err := types.NewUUID(podIDStr)
	if err != nil {
		t.Fatalf("%q is not a valid UUID: %v", podIDStr, err)
	}

	t.Logf("Retrieving stage1 image name")
	stage1ImageName := getImageName(t, ctx, stage1App)

	t.Logf("Removing stage1 image (should work)")
	if err := removeImage(ctx, []string{stage1ImageName}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, []string{referencedApp}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, []string{unreferencedApp}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	cmd = fmt.Sprintf("%s run-prepared --mds-register=false %s", ctx.Cmd(), podID.String())
	t.Logf("Running %s", referencedACI)
	spawnAndWaitOrFail(t, cmd, true)
}

func TestImagePrepareRmIDRun(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	t.Logf("Fetching %s", imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmds := strings.Fields(ctx.Cmd())
	prepareCmd := exec.Command(cmds[0], cmds[1:]...)
	prepareCmd.Args = append(prepareCmd.Args, "--insecure-options=image", "prepare", referencedACI)
	output, err := prepareCmd.Output()
	if err != nil {
		t.Fatalf("Cannot read the output: %v", err)
	}

	podIDStr := strings.TrimSpace(string(output))
	podID, err := types.NewUUID(podIDStr)
	if err != nil {
		t.Fatalf("%q is not a valid UUID: %v", podIDStr, err)
	}

	t.Logf("Retrieving stage1 imageID")
	stage1ImageID, err := getImageId(ctx, stage1App)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", referencedApp)
	referencedImageID, err := getImageId(ctx, referencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", unreferencedApp)
	unreferencedImageID, err := getImageId(ctx, unreferencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Removing stage1 image (should work)")
	if err := removeImage(ctx, []string{stage1ImageID}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, []string{referencedImageID}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, []string{unreferencedImageID}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	cmd = fmt.Sprintf("%s run-prepared --mds-register=false %s", ctx.Cmd(), podID.String())
	t.Logf("Running %s", referencedACI)
	spawnAndWaitOrFail(t, cmd, true)
}

func TestImagePrepareRmDuplicate(t *testing.T) {
	imageFile := patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
	defer os.Remove(imageFile)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	cmd := fmt.Sprintf("%s --insecure-options=image fetch %s", ctx.Cmd(), imageFile)
	t.Logf("Fetching %s", imageFile)
	spawnAndWaitOrFail(t, cmd, true)

	// at this point we know that RKT_INSPECT_IMAGE env var is not empty
	referencedACI := os.Getenv("RKT_INSPECT_IMAGE")
	cmds := strings.Fields(ctx.Cmd())
	prepareCmd := exec.Command(cmds[0], cmds[1:]...)
	prepareCmd.Args = append(prepareCmd.Args, "--insecure-options=image", "prepare", referencedACI)
	output, err := prepareCmd.Output()
	if err != nil {
		t.Fatalf("Cannot read the output: %v", err)
	}

	podIDStr := strings.TrimSpace(string(output))
	podID, err := types.NewUUID(podIDStr)
	if err != nil {
		t.Fatalf("%q is not a valid UUID: %v", podIDStr, err)
	}

	t.Logf("Retrieving %s image ID", referencedApp)
	referencedImageID, err := getImageId(ctx, referencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Retrieving %s image ID", unreferencedApp)
	unreferencedImageID, err := getImageId(ctx, unreferencedApp)
	if err != nil {
		t.Fatalf("rkt didn't terminate correctly: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", referencedApp)
	if err := removeImage(ctx, []string{referencedApp, referencedImageID}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	t.Logf("Removing image for app %s (should work)", unreferencedApp)
	if err := removeImage(ctx, []string{unreferencedImageID, unreferencedApp}, true); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	cmd = fmt.Sprintf("%s run-prepared --mds-register=false %s", ctx.Cmd(), podID.String())
	t.Logf("Running %s", referencedACI)
	spawnAndWaitOrFail(t, cmd, true)
}

func getImageName(t *testing.T, ctx *testutils.RktRunCtx, name string) string {
	cmd := fmt.Sprintf(`/bin/sh -c "%s image list --fields=name --no-legend | grep %s | cut -d: -f1"`, ctx.Cmd(), name)
	child := spawnOrFail(t, cmd)
	imageName, err := child.ReadLine()
	imageName = strings.TrimSpace(imageName)
	imageName = string(bytes.Trim([]byte(imageName), "\x00"))
	if err != nil {
		t.Fatalf("Cannot exec: %v", err)
	}
	waitOrFail(t, child, true)
	return imageName
}

func getImageId(ctx *testutils.RktRunCtx, name string) (string, error) {
	cmd := fmt.Sprintf(`/bin/sh -c "%s image list --fields=id,name --no-legend | grep %s | awk '{print $1}'"`, ctx.Cmd(), name)
	child, err := gexpect.Spawn(cmd)
	if err != nil {
		return "", fmt.Errorf("Cannot exec rkt: %v", err)
	}
	imageID, err := child.ReadLine()
	imageID = strings.TrimSpace(imageID)
	imageID = string(bytes.Trim([]byte(imageID), "\x00"))
	if err != nil {
		return "", fmt.Errorf("Cannot exec: %v", err)
	}
	if err := child.Wait(); err != nil {
		return "", fmt.Errorf("rkt didn't terminate correctly: %v", err)
	}
	return imageID, nil
}

func removeImage(ctx *testutils.RktRunCtx, images []string, shouldWork bool) error {
	cmd := fmt.Sprintf("%s image rm", ctx.Cmd())
	for _, image := range images {
		cmd += fmt.Sprintf(" %s", image)
	}
	expect := fmt.Sprintf(rmImageReferenced, images[0])
	if shouldWork {
		expect = rmImageOk
	}
	child, err := gexpect.Spawn(cmd)
	if err != nil {
		return fmt.Errorf("Cannot exec: %v", err)
	}
	if err := expectWithOutput(child, expect); err != nil {
		return fmt.Errorf("Expected %q but not found: %v", expect, err)
	}
	if err := child.Wait(); err != nil {
		return fmt.Errorf("rkt didn't terminate correctly: %v", err)
	}
	return nil
}
<|endoftext|>