<commit_before>package libvirt

import (
	"encoding/xml"
	"fmt"
	libvirt "github.com/dmacvicar/libvirt-go"
	"github.com/hashicorp/terraform/helper/schema"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
)

// network transparent image
type image interface {
	Size() (int64, error)
	WriteToStream(*libvirt.VirStream) error
	String() string
}

type localImage struct {
	path string
}

func (i *localImage) String() string {
	return i.path
}

func (i *localImage) Size() (int64, error) {
	file, err := os.Open(i.path)
	if err != nil {
		return 0, err
	}

	fi, err := file.Stat()
	if err != nil {
		return 0, err
	}
	return fi.Size(), nil
}

func (i *localImage) WriteToStream(stream *libvirt.VirStream) error {
	file, err := os.Open(i.path)
	defer file.Close()
	if err != nil {
		return fmt.Errorf("Error while opening %s: %s", i.path, err)
	}

	n, err := io.Copy(stream, file)
	if err != nil {
		return fmt.Errorf("Error while reading %s: %s", i.path, err)
	}
	log.Printf("%d bytes uploaded\n", n)
	return nil
}

type httpImage struct {
	url *url.URL
}

func (i *httpImage) String() string {
	return i.url.String()
}

func (i *httpImage) Size() (int64, error) {
	response, err := http.Head(i.url.String())
	if err != nil {
		return 0, err
	}
	length, err := strconv.Atoi(response.Header.Get("Content-Length"))
	if err != nil {
		return 0, err
	}
	return int64(length), nil
}

func (i *httpImage) WriteToStream(stream *libvirt.VirStream) error {
	response, err := http.Get(i.url.String())
	defer response.Body.Close()
	if err != nil {
		return fmt.Errorf("Error while downloading %s: %s", i.url.String(), err)
	}

	n, err := io.Copy(stream, response.Body)
	if err != nil {
		return fmt.Errorf("Error while downloading %s: %s", i.url.String(), err)
	}
	log.Printf("%d bytes uploaded\n", n)
	return nil
}

func newImage(source string) (image, error) {
	url, err := url.Parse(source)
	if err != nil {
		return nil, fmt.Errorf("Can't parse source '%s' as url: %s", source, err)
	}

	if strings.HasPrefix(url.Scheme, "http") {
		return &httpImage{url: url}, nil
	} else if url.Scheme == "file" || url.Scheme == "" {
		return &localImage{path: url.Path}, nil
	} else {
		return nil, fmt.Errorf("Don't know how to read from '%s': %s", url.String(), err)
	}
}

func volumeCommonSchema() map[string]*schema.Schema {
	return map[string]*schema.Schema{
		"name": &schema.Schema{
			Type:     schema.TypeString,
			Required: true,
			ForceNew: true,
		},
		"pool": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			Default:  "default",
			ForceNew: true,
		},
		"source": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"size": &schema.Schema{
			Type:     schema.TypeInt,
			Optional: true,
			Computed: true,
			ForceNew: true,
		},
		"base_volume_id": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"base_volume_pool": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"base_volume_name": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
	}
}

func resourceLibvirtVolume() *schema.Resource {
	return &schema.Resource{
		Create: resourceLibvirtVolumeCreate,
		Read:   resourceLibvirtVolumeRead,
		Delete: resourceLibvirtVolumeDelete,
		Schema: volumeCommonSchema(),
	}
}

func remoteImageSize(url string) (int, error) {
	response, err := http.Head(url)
	if err != nil {
		return 0, err
	}
	length, err := strconv.Atoi(response.Header.Get("Content-Length"))
	if err != nil {
		return 0, err
	}
	return length, nil
}

func resourceLibvirtVolumeCreate(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	poolName := "default"
	if _, ok := d.GetOk("pool"); ok {
		poolName = d.Get("pool").(string)
	}

	pool, err := virConn.LookupStoragePoolByName(poolName)
	if err != nil {
		return fmt.Errorf("can't find storage pool '%s'", poolName)
	}
	defer pool.Free()

	// Refresh the pool of the volume so that libvirt knows it is
	// not longer in use.
	WaitForSuccess("Error refreshing pool for volume", func() error {
		return pool.Refresh(0)
	})

	volumeDef := newDefVolume()

	if name, ok := d.GetOk("name"); ok {
		volumeDef.Name = name.(string)
	}

	// an existing image was given, this mean we can't choose size
	if url, ok := d.GetOk("source"); ok {
		// source and size conflict
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'source' is given (the size will be set to the size of the source image.")
		}
		if _, ok := d.GetOk("base_volume_id"); ok {
			return fmt.Errorf("'base_volume_id' can't be specified when also 'source' is given.")
		}

		if _, ok := d.GetOk("base_volume_name"); ok {
			return fmt.Errorf("'base_volume_name' can't be specified when also 'source' is given.")
		}

		img, err := newImage(url.(string))
		if err != nil {
			return err
		}

		size, err := img.Size()
		if err != nil {
			return err
		}
		log.Printf("Image %s image is: %d bytes", img, size)

		volumeDef.Capacity.Unit = "B"
		volumeDef.Capacity.Amount = size
	} else {
		_, noSize := d.GetOk("size")
		_, noBaseVol := d.GetOk("base_volume_id")

		if noSize && noBaseVol {
			return fmt.Errorf("'size' needs to be specified if no 'source' or 'base_volume_id' is given.")
		}
		volumeDef.Capacity.Amount = int64(d.Get("size").(int))
	}

	if baseVolumeId, ok := d.GetOk("base_volume_id"); ok {
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'base_volume_id' is given (the size will be set to the size of the backing image.")
		}

		if _, ok := d.GetOk("base_volume_name"); ok {
			return fmt.Errorf("'base_volume_name' can't be specified when also 'base_volume_id' is given.")
		}

		volumeDef.BackingStore = new(defBackingStore)
		volumeDef.BackingStore.Format.Type = "qcow2"
		baseVolume, err := virConn.LookupStorageVolByKey(baseVolumeId.(string))
		if err != nil {
			return fmt.Errorf("Can't retrieve volume %s", baseVolumeId.(string))
		}
		baseVolPath, err := baseVolume.GetPath()
		if err != nil {
			return fmt.Errorf("can't get name for base image '%s'", baseVolumeId)
		}
		volumeDef.BackingStore.Path = baseVolPath
	}

	if baseVolumeName, ok := d.GetOk("base_volume_name"); ok {
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'base_volume_name' is given (the size will be set to the size of the backing image.")
		}

		baseVolumePool := pool
		if _, ok := d.GetOk("base_volume_pool"); ok {
			baseVolumePoolName := d.Get("base_volume_pool").(string)
			baseVolumePool, err = virConn.LookupStoragePoolByName(baseVolumePoolName)
			if err != nil {
				return fmt.Errorf("can't find storage pool '%s'", baseVolumePoolName)
			}
			defer baseVolumePool.Free()
		}

		volumeDef.BackingStore = new(defBackingStore)
		volumeDef.BackingStore.Format.Type = "qcow2"
		baseVolume, err := baseVolumePool.LookupStorageVolByName(baseVolumeName.(string))
		if err != nil {
			return fmt.Errorf("Can't retrieve volume %s", baseVolumeName.(string))
		}
		baseVolPath, err := baseVolume.GetPath()
		if err != nil {
			return fmt.Errorf("can't get name for base image '%s'", baseVolumeName)
		}
		volumeDef.BackingStore.Path = baseVolPath
	}

	volumeDefXml, err := xml.Marshal(volumeDef)
	if err != nil {
		return fmt.Errorf("Error serializing libvirt volume: %s", err)
	}

	// create the volume
	volume, err := pool.StorageVolCreateXML(string(volumeDefXml), 0)
	if err != nil {
		return fmt.Errorf("Error creating libvirt volume: %s", err)
	}
	defer volume.Free()

	// we use the key as the id
	key, err := volume.GetKey()
	if err != nil {
		return fmt.Errorf("Error retrieving volume key: %s", err)
	}
	d.SetId(key)
	log.Printf("[INFO] Volume ID: %s", d.Id())

	// upload source if present
	if source, ok := d.GetOk("source"); ok {
		stream, err := libvirt.NewVirStream(virConn, 0)
		if err != nil {
			return err
		}
		defer stream.Close()

		img, err := newImage(source.(string))
		if err != nil {
			return err
		}

		volume.Upload(stream, 0, uint64(volumeDef.Capacity.Amount), 0)
		err = img.WriteToStream(stream)
		if err != nil {
			return err
		}
	}

	return resourceLibvirtVolumeRead(d, meta)
}

func resourceLibvirtVolumeRead(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	volume, err := virConn.LookupStorageVolByKey(d.Id())
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s", d.Id())
	}
	defer volume.Free()

	volName, err := volume.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving volume name: %s", err)
	}

	volPool, err := volume.LookupPoolByVolume()
	if err != nil {
		return fmt.Errorf("Error retrieving pool for volume: %s", err)
	}
	defer volPool.Free()

	volPoolName, err := volPool.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving pool name: %s", err)
	}

	d.Set("pool", volPoolName)
	d.Set("name", volName)

	info, err := volume.GetInfo()
	if err != nil {
		return fmt.Errorf("Error retrieving volume name: %s", err)
	}
	d.Set("size", info.GetCapacityInBytes())

	return nil
}

func resourceLibvirtVolumeDelete(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	return RemoveVolume(virConn, d.Id())
}
<commit_msg>Start storage pool if volume was not found<commit_after>package libvirt

import (
	"encoding/xml"
	"fmt"
	libvirt "github.com/dmacvicar/libvirt-go"
	"github.com/hashicorp/terraform/helper/schema"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
)

// network transparent image
type image interface {
	Size() (int64, error)
	WriteToStream(*libvirt.VirStream) error
	String() string
}

type localImage struct {
	path string
}

func (i *localImage) String() string {
	return i.path
}

func (i *localImage) Size() (int64, error) {
	file, err := os.Open(i.path)
	if err != nil {
		return 0, err
	}

	fi, err := file.Stat()
	if err != nil {
		return 0, err
	}
	return fi.Size(), nil
}

func (i *localImage) WriteToStream(stream *libvirt.VirStream) error {
	file, err := os.Open(i.path)
	defer file.Close()
	if err != nil {
		return fmt.Errorf("Error while opening %s: %s", i.path, err)
	}

	n, err := io.Copy(stream, file)
	if err != nil {
		return fmt.Errorf("Error while reading %s: %s", i.path, err)
	}
	log.Printf("%d bytes uploaded\n", n)
	return nil
}

type httpImage struct {
	url *url.URL
}

func (i *httpImage) String() string {
	return i.url.String()
}

func (i *httpImage) Size() (int64, error) {
	response, err := http.Head(i.url.String())
	if err != nil {
		return 0, err
	}
	length, err := strconv.Atoi(response.Header.Get("Content-Length"))
	if err != nil {
		return 0, err
	}
	return int64(length), nil
}

func (i *httpImage) WriteToStream(stream *libvirt.VirStream) error {
	response, err := http.Get(i.url.String())
	defer response.Body.Close()
	if err != nil {
		return fmt.Errorf("Error while downloading %s: %s", i.url.String(), err)
	}

	n, err := io.Copy(stream, response.Body)
	if err != nil {
		return fmt.Errorf("Error while downloading %s: %s", i.url.String(), err)
	}
	log.Printf("%d bytes uploaded\n", n)
	return nil
}

func newImage(source string) (image, error) {
	url, err := url.Parse(source)
	if err != nil {
		return nil, fmt.Errorf("Can't parse source '%s' as url: %s", source, err)
	}

	if strings.HasPrefix(url.Scheme, "http") {
		return &httpImage{url: url}, nil
	} else if url.Scheme == "file" || url.Scheme == "" {
		return &localImage{path: url.Path}, nil
	} else {
		return nil, fmt.Errorf("Don't know how to read from '%s': %s", url.String(), err)
	}
}

func volumeCommonSchema() map[string]*schema.Schema {
	return map[string]*schema.Schema{
		"name": &schema.Schema{
			Type:     schema.TypeString,
			Required: true,
			ForceNew: true,
		},
		"pool": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			Default:  "default",
			ForceNew: true,
		},
		"source": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"size": &schema.Schema{
			Type:     schema.TypeInt,
			Optional: true,
			Computed: true,
			ForceNew: true,
		},
		"base_volume_id": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"base_volume_pool": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"base_volume_name": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
	}
}

func resourceLibvirtVolume() *schema.Resource {
	return &schema.Resource{
		Create: resourceLibvirtVolumeCreate,
		Read:   resourceLibvirtVolumeRead,
		Delete: resourceLibvirtVolumeDelete,
		Schema: volumeCommonSchema(),
	}
}

func remoteImageSize(url string) (int, error) {
	response, err := http.Head(url)
	if err != nil {
		return 0, err
	}
	length, err := strconv.Atoi(response.Header.Get("Content-Length"))
	if err != nil {
		return 0, err
	}
	return length, nil
}

func resourceLibvirtVolumeCreate(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	poolName := "default"
	if _, ok := d.GetOk("pool"); ok {
		poolName = d.Get("pool").(string)
	}

	pool, err := virConn.LookupStoragePoolByName(poolName)
	if err != nil {
		return fmt.Errorf("can't find storage pool '%s'", poolName)
	}
	defer pool.Free()

	// Refresh the pool of the volume so that libvirt knows it is
	// not longer in use.
	WaitForSuccess("Error refreshing pool for volume", func() error {
		return pool.Refresh(0)
	})

	volumeDef := newDefVolume()

	if name, ok := d.GetOk("name"); ok {
		volumeDef.Name = name.(string)
	}

	// an existing image was given, this mean we can't choose size
	if url, ok := d.GetOk("source"); ok {
		// source and size conflict
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'source' is given (the size will be set to the size of the source image.")
		}
		if _, ok := d.GetOk("base_volume_id"); ok {
			return fmt.Errorf("'base_volume_id' can't be specified when also 'source' is given.")
		}

		if _, ok := d.GetOk("base_volume_name"); ok {
			return fmt.Errorf("'base_volume_name' can't be specified when also 'source' is given.")
		}

		img, err := newImage(url.(string))
		if err != nil {
			return err
		}

		size, err := img.Size()
		if err != nil {
			return err
		}
		log.Printf("Image %s image is: %d bytes", img, size)

		volumeDef.Capacity.Unit = "B"
		volumeDef.Capacity.Amount = size
	} else {
		_, noSize := d.GetOk("size")
		_, noBaseVol := d.GetOk("base_volume_id")

		if noSize && noBaseVol {
			return fmt.Errorf("'size' needs to be specified if no 'source' or 'base_volume_id' is given.")
		}
		volumeDef.Capacity.Amount = int64(d.Get("size").(int))
	}

	if baseVolumeId, ok := d.GetOk("base_volume_id"); ok {
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'base_volume_id' is given (the size will be set to the size of the backing image.")
		}

		if _, ok := d.GetOk("base_volume_name"); ok {
			return fmt.Errorf("'base_volume_name' can't be specified when also 'base_volume_id' is given.")
		}

		volumeDef.BackingStore = new(defBackingStore)
		volumeDef.BackingStore.Format.Type = "qcow2"
		baseVolume, err := virConn.LookupStorageVolByKey(baseVolumeId.(string))
		if err != nil {
			return fmt.Errorf("Can't retrieve volume %s", baseVolumeId.(string))
		}
		baseVolPath, err := baseVolume.GetPath()
		if err != nil {
			return fmt.Errorf("can't get name for base image '%s'", baseVolumeId)
		}
		volumeDef.BackingStore.Path = baseVolPath
	}

	if baseVolumeName, ok := d.GetOk("base_volume_name"); ok {
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'base_volume_name' is given (the size will be set to the size of the backing image.")
		}

		baseVolumePool := pool
		if _, ok := d.GetOk("base_volume_pool"); ok {
			baseVolumePoolName := d.Get("base_volume_pool").(string)
			baseVolumePool, err = virConn.LookupStoragePoolByName(baseVolumePoolName)
			if err != nil {
				return fmt.Errorf("can't find storage pool '%s'", baseVolumePoolName)
			}
			defer baseVolumePool.Free()
		}

		volumeDef.BackingStore = new(defBackingStore)
		volumeDef.BackingStore.Format.Type = "qcow2"
		baseVolume, err := baseVolumePool.LookupStorageVolByName(baseVolumeName.(string))
		if err != nil {
			return fmt.Errorf("Can't retrieve volume %s", baseVolumeName.(string))
		}
		baseVolPath, err := baseVolume.GetPath()
		if err != nil {
			return fmt.Errorf("can't get name for base image '%s'", baseVolumeName)
		}
		volumeDef.BackingStore.Path = baseVolPath
	}

	volumeDefXml, err := xml.Marshal(volumeDef)
	if err != nil {
		return fmt.Errorf("Error serializing libvirt volume: %s", err)
	}

	// create the volume
	volume, err := pool.StorageVolCreateXML(string(volumeDefXml), 0)
	if err != nil {
		return fmt.Errorf("Error creating libvirt volume: %s", err)
	}
	defer volume.Free()

	// we use the key as the id
	key, err := volume.GetKey()
	if err != nil {
		return fmt.Errorf("Error retrieving volume key: %s", err)
	}
	d.SetId(key)
	log.Printf("[INFO] Volume ID: %s", d.Id())

	// upload source if present
	if source, ok := d.GetOk("source"); ok {
		stream, err := libvirt.NewVirStream(virConn, 0)
		if err != nil {
			return err
		}
		defer stream.Close()

		img, err := newImage(source.(string))
		if err != nil {
			return err
		}

		volume.Upload(stream, 0, uint64(volumeDef.Capacity.Amount), 0)
		err = img.WriteToStream(stream)
		if err != nil {
			return err
		}
	}

	return resourceLibvirtVolumeRead(d, meta)
}

func resourceLibvirtVolumeRead(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	volume, err := virConn.LookupStorageVolByKey(d.Id())
	if err != nil {
		virErr := err.(libvirt.VirError)
		if virErr.Code == libvirt.VIR_ERR_NO_STORAGE_VOL {
			volId := d.Id()
			volPoolName := d.Get("pool").(string)
			log.Printf("[INFO] Volume %s not found, attempting to start pool %s", volId, volPoolName)

			volPool, err := virConn.LookupStoragePoolByName(volPoolName)
			if err != nil {
				return fmt.Errorf("Error retrieving pool %s for volume %s: %s", volPoolName, volId, err)
			}
			defer volPool.Free()

			err = volPool.Create(0)
			if err != nil {
				return fmt.Errorf("Error starting pool %s: %s", volPoolName, err)
			}

			// attempt a new lookup
			volume, err = virConn.LookupStorageVolByKey(d.Id())
			if err != nil {
				return fmt.Errorf("Can't retrieve volume %s", d.Id())
			}
		} else {
			return fmt.Errorf("Can't retrieve volume %s", d.Id())
		}
	}
	defer volume.Free()

	volName, err := volume.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving volume name: %s", err)
	}

	volPool, err := volume.LookupPoolByVolume()
	if err != nil {
		return fmt.Errorf("Error retrieving pool for volume: %s", err)
	}
	defer volPool.Free()

	volPoolName, err := volPool.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving pool name: %s", err)
	}

	d.Set("pool", volPoolName)
	d.Set("name", volName)

	info, err := volume.GetInfo()
	if err != nil {
		return fmt.Errorf("Error retrieving volume name: %s", err)
	}
	d.Set("size", info.GetCapacityInBytes())

	return nil
}

func resourceLibvirtVolumeDelete(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	return RemoveVolume(virConn, d.Id())
}
<|endoftext|>


<commit_before>package formats

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/martinlindhe/formats/parse"
	"github.com/martinlindhe/formats/parse/archive"
	"github.com/martinlindhe/formats/parse/av"
	"github.com/martinlindhe/formats/parse/bin"
	"github.com/martinlindhe/formats/parse/doc"
	"github.com/martinlindhe/formats/parse/exe"
	"github.com/martinlindhe/formats/parse/font"
	"github.com/martinlindhe/formats/parse/image"
	"github.com/martinlindhe/formats/parse/macos"
	"github.com/martinlindhe/formats/parse/windows"
)

type Parser func(*parse.ParseChecker) (*parse.ParsedLayout, error)

var (
	parsers = map[string]Parser{
		"7z":             archive.SevenZIP,
		"arj":            archive.ARJ,
		"bzip2":          archive.BZIP2,
		"cab":            archive.CAB,
		"deb":            archive.DEB,
		"gzip":           archive.GZIP,
		"iso":            archive.ISO,
		"lzma":           archive.LZMA,
		"rar":            archive.RAR,
		"td2":            archive.TD2,
		"vdi":            archive.VDI,
		"wim":            archive.WIM,
		"winimg":         archive.WINIMG,
		"xz":             archive.XZ,
		"zip":            archive.ZIP,
		"aiff":           av.AIFF,
		"asf":            av.ASF,
		"caf":            av.CAF,
		"flac":           av.FLAC,
		"flv":            av.FLV,
		"midi":           av.MIDI,
		"mkv":            av.MKV,
		"mp3":            av.MP3,
		"mp4":            av.MP4,
		"ogg":            av.OGG,
		"riff":           av.RIFF,
		"dbm":            bin.GnuDBM,
		"gpg":            bin.GnuPG,
		"mo":             bin.GnuMO,
		"mapledb":        bin.MapleDB,
		"sqlite3":        bin.SQLITE3,
		"pdb":            bin.PDB,
		"rdb":            bin.RDB,
		"terminfo":       bin.Terminfo,
		"timezone":       bin.Timezone,
		"xkm":            bin.XKM,
		"chm":            doc.CHM,
		"hlp":            doc.HLP,
		"pdf":            doc.PDF,
		"rtf":            doc.RTF,
		"wri":            doc.WRI,
		"gb-rom":         exe.GameboyROM,
		"gba-rom":        exe.GBAROM,
		"n64-rom":        exe.N64ROM,
		"dex":            exe.DEX,
		"elf":            exe.ELF,
		"java":           exe.JavaClass,
		"lua":            exe.LUA,
		"macho":          exe.MachO,
		"mz":             exe.MZ,
		"python":         exe.PythonBytecode,
		"swf":            exe.SWF,
		"vbe":            exe.VBE,
		"eot":            font.EOT,
		"otf":            font.OTF,
		"pfb":            font.PFB,
		"ps1":            font.PSType1,
		"ttc":            font.TTC,
		"ttf":            font.TTF,
		"woff":           font.WOFF,
		"woff2":          font.WOFF2,
		"x11-snf":        font.X11FontSNF,
		"bmp":            image.BMP,
		"gif":            image.GIF,
		"icns":           image.ICNS,
		"ico":            image.ICO,
		"jpeg":           image.JPEG,
		"pcx":            image.PCX,
		"png":            image.PNG,
		"tga":            image.TGA, // XXX has too loose detection
		"tiff":           image.TIFF,
		"xcursor":        image.XCursor,
		"bom_store":      macos.BOMStore,
		"bplist":         macos.BPLIST,
		"code_directory": macos.CodeDirectory,
		"ds_store":       macos.DSStore,
		"keychain":       macos.Keychain,
		"ari8":           windows.ARI8,
		"dxbc":           windows.DXBC,
		"grp":            windows.GRP,
		"hwrs":           windows.HWRS,
		"lnk":            windows.LNK,
		"ole-cf":         windows.OLECF,
		"p7x":            windows.P7X,
		"pif":            windows.PIF,
		"pri":            windows.PRI,
		"rbs":            windows.RBS,
		"rsrc":           windows.RSRC,
		"uce":            windows.UCE,
		"xbf":            windows.XBF,
	}
)

type MatchingParsers []parse.ParsedLayout

func (mp *MatchingParsers) First() *parse.ParsedLayout {

	for _, parser := range *mp {
		return &parser
	}
	return nil
}

func (mp *MatchingParsers) ChoseOne(file *os.File) (*parse.ParsedLayout, error) {

	i := 1
	fmt.Println("multiple parsers matched input file, please choose one:\n")
	for name, _ := range parsers {
		fmt.Printf("%d: %s\n", i, name)
		i++
	}

	reader := bufio.NewReader(os.Stdin)
	text, _ := reader.ReadString('\n')
	text = strings.TrimSpace(text)

	u, err := strconv.ParseUint(text, 10, 64)
	if err != nil {
		return nil, fmt.Errorf("invalid input")
	}

	fileSize, _ := fileSize(file)
	layout := parse.ParsedLayout{
		FileName: fileGetName(file),
		FileSize: fileSize}
	checker := parse.ParseChecker{
		File:         file,
		ParsedLayout: layout}

	i = 1
	for name, parser := range parsers {
		if i == int(u) {

			pl, err2 := parser(&checker)
			if err2 != nil {
				fmt.Println("XXX parser", name, "failed")
				return nil, err2
			}
			return pl, nil
		}
		i++
	}

	return nil, fmt.Errorf("selection not in list")
}

// returns all matching parsers
func MatchAll(file *os.File) (MatchingParsers, error) {

	fileSize, _ := fileSize(file)
	layout := parse.ParsedLayout{
		FileName: fileGetName(file),
		FileSize: fileSize}
	checker := parse.ParseChecker{
		File:         file,
		ParsedLayout: layout}
	var m MatchingParsers

	var err error
	checker.Header, err = readHeaderChunk(file)
	if err != nil {
		fmt.Println("warning: MatchAll failed reading header chunk")
		return nil, err
	}

	for name, parser := range parsers {

		pl, err2 := parser(&checker)
		if err2 != nil {
			fmt.Println("warning: parser", name, "failed with", err2)
			continue
		}
		if pl == nil {
			continue
		}
		if pl.FormatName == "" {
			pl.FormatName = name
		}
		m = append(m, *pl)
	}

	if len(m) == 0 {
		// try text detector
		text, err := parse.Text(&checker)
		if err == nil {
			m = append(m, *text)
		} else {
			// fall back to raw
			raw, _ := parse.RAW(&checker)
			m = append(m, *raw)
		}
	}

	return m, nil
}

// ParseLayout returns a ParsedLayout for the file
func ParseLayout(file *os.File) (*parse.ParsedLayout, error) {

	all, err := MatchAll(file)
	return all.First(), err
}

// slice to expand, new length in bytes
func expandByteSlice(b []byte, newLen int64) []byte {

	i := int64(len(b))
	newLen = newLen - i
	return append(b[:i], append(make([]byte, newLen), b[i:]...)...)
}

func readHeaderChunk(file *os.File) ([]byte, error) {

	fileSize, _ := fileSize(file)
	maxHeaderLen := int64(0xffff)
	len := fileSize
	if len > maxHeaderLen {
		len = maxHeaderLen
	}

	b := make([]byte, len)

	file.Seek(0, os.SEEK_SET)
	if err := binary.Read(file, binary.LittleEndian, &b); err != nil {
		fmt.Println("error: failed to read header!", err)
		return nil, err
	}

	// resize to maxHeaderLen
	return expandByteSlice(b, maxHeaderLen), nil
}

func fileGetName(file *os.File) string {
	stat, _ := file.Stat()
	return stat.Name()
}
<commit_msg>cmd/formats: selector if file matches multiple parsers, finishing touches<commit_after>package formats

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/martinlindhe/formats/parse"
	"github.com/martinlindhe/formats/parse/archive"
	"github.com/martinlindhe/formats/parse/av"
	"github.com/martinlindhe/formats/parse/bin"
	"github.com/martinlindhe/formats/parse/doc"
	"github.com/martinlindhe/formats/parse/exe"
	"github.com/martinlindhe/formats/parse/font"
	"github.com/martinlindhe/formats/parse/image"
	"github.com/martinlindhe/formats/parse/macos"
	"github.com/martinlindhe/formats/parse/windows"
)

type Parser func(*parse.ParseChecker) (*parse.ParsedLayout, error)

var (
	parsers = map[string]Parser{
		"7z":             archive.SevenZIP,
		"arj":            archive.ARJ,
		"bzip2":          archive.BZIP2,
		"cab":            archive.CAB,
		"deb":            archive.DEB,
		"gzip":           archive.GZIP,
		"iso":            archive.ISO,
		"lzma":           archive.LZMA,
		"rar":            archive.RAR,
		"td2":            archive.TD2,
		"vdi":            archive.VDI,
		"wim":            archive.WIM,
		"winimg":         archive.WINIMG,
		"xz":             archive.XZ,
		"zip":            archive.ZIP,
		"aiff":           av.AIFF,
		"asf":            av.ASF,
		"caf":            av.CAF,
		"flac":           av.FLAC,
		"flv":            av.FLV,
		"midi":           av.MIDI,
		"mkv":            av.MKV,
		"mp3":            av.MP3,
		"mp4":            av.MP4,
		"ogg":            av.OGG,
		"riff":           av.RIFF,
		"dbm":            bin.GnuDBM,
		"gpg":            bin.GnuPG,
		"mo":             bin.GnuMO,
		"mapledb":        bin.MapleDB,
		"sqlite3":        bin.SQLITE3,
		"pdb":            bin.PDB,
		"rdb":            bin.RDB,
		"terminfo":       bin.Terminfo,
		"timezone":       bin.Timezone,
		"xkm":            bin.XKM,
		"chm":            doc.CHM,
		"hlp":            doc.HLP,
		"pdf":            doc.PDF,
		"rtf":            doc.RTF,
		"wri":            doc.WRI,
		"gb-rom":         exe.GameboyROM,
		"gba-rom":        exe.GBAROM,
		"n64-rom":        exe.N64ROM,
		"dex":            exe.DEX,
		"elf":            exe.ELF,
		"java":           exe.JavaClass,
		"lua":            exe.LUA,
		"macho":          exe.MachO,
		"mz":             exe.MZ,
		"python":         exe.PythonBytecode,
		"swf":            exe.SWF,
		"vbe":            exe.VBE,
		"eot":            font.EOT,
		"otf":            font.OTF,
		"pfb":            font.PFB,
		"ps1":            font.PSType1,
		"ttc":            font.TTC,
		"ttf":            font.TTF,
		"woff":           font.WOFF,
		"woff2":          font.WOFF2,
		"x11-snf":        font.X11FontSNF,
		"bmp":            image.BMP,
		"gif":            image.GIF,
		"icns":           image.ICNS,
		"ico":            image.ICO,
		"jpeg":           image.JPEG,
		"pcx":            image.PCX,
		"png":            image.PNG,
		"tga":            image.TGA, // XXX has too loose detection
		"tiff":           image.TIFF,
		"xcursor":        image.XCursor,
		"bom_store":      macos.BOMStore,
		"bplist":         macos.BPLIST,
		"code_directory": macos.CodeDirectory,
		"ds_store":       macos.DSStore,
		"keychain":       macos.Keychain,
		"ari8":           windows.ARI8,
		"dxbc":           windows.DXBC,
		"grp":            windows.GRP,
		"hwrs":           windows.HWRS,
		"lnk":            windows.LNK,
		"ole-cf":         windows.OLECF,
		"p7x":            windows.P7X,
		"pif":            windows.PIF,
		"pri":            windows.PRI,
		"rbs":            windows.RBS,
		"rsrc":           windows.RSRC,
		"uce":            windows.UCE,
		"xbf":            windows.XBF,
	}
)

type MatchingParsers []parse.ParsedLayout

func (mp *MatchingParsers) First() *parse.ParsedLayout {

	for _, parser := range *mp {
		return &parser
	}
	return nil
}

func (mp *MatchingParsers) ChoseOne(file *os.File) (*parse.ParsedLayout, error) {

	i := 1
	fmt.Println("multiple parsers matched input file, please choose one:\n")
	for _, pl := range *mp {
		fmt.Printf("%d: %s\n", i, pl.FormatName)
		i++
	}

	reader := bufio.NewReader(os.Stdin)
	text, _ := reader.ReadString('\n')
	text = strings.TrimSpace(text)

	u, err := strconv.ParseUint(text, 10, 64)
	if err != nil {
		return nil, fmt.Errorf("invalid input")
	}

	i = 1
	for _, pl := range *mp {
		if i == int(u) {
			return &pl, nil
		}
		i++
	}

	return nil, fmt.Errorf("selection not in list")
}

// returns all matching parsers
func MatchAll(file *os.File) (MatchingParsers, error) {

	fileSize, _ := fileSize(file)
	layout := parse.ParsedLayout{
		FileName: fileGetName(file),
		FileSize: fileSize}
	checker := parse.ParseChecker{
		File:         file,
		ParsedLayout: layout}
	var m MatchingParsers

	var err error
	checker.Header, err = readHeaderChunk(file)
	if err != nil {
		fmt.Println("warning: MatchAll failed reading header chunk")
		return nil, err
	}

	for name, parser := range parsers {

		pl, err2 := parser(&checker)
		if err2 != nil {
			fmt.Println("warning: parser", name, "failed with", err2)
			continue
		}
		if pl == nil {
			continue
		}
		if pl.FormatName == "" {
			pl.FormatName = name
		}
		m = append(m, *pl)
	}

	if len(m) == 0 {
		// try text detector
		text, err := parse.Text(&checker)
		if err == nil {
			m = append(m, *text)
		} else {
			// fall back to raw
			raw, _ := parse.RAW(&checker)
			m = append(m, *raw)
		}
	}

	return m, nil
}

// ParseLayout returns a ParsedLayout for the file
func ParseLayout(file *os.File) (*parse.ParsedLayout, error) {

	all, err := MatchAll(file)
	return all.First(), err
}

// slice to expand, new length in bytes
func expandByteSlice(b []byte, newLen int64) []byte {

	i := int64(len(b))
	newLen = newLen - i
	return append(b[:i], append(make([]byte, newLen), b[i:]...)...)
}

func readHeaderChunk(file *os.File) ([]byte, error) {

	fileSize, _ := fileSize(file)
	maxHeaderLen := int64(0xffff)
	len := fileSize
	if len > maxHeaderLen {
		len = maxHeaderLen
	}

	b := make([]byte, len)

	file.Seek(0, os.SEEK_SET)
	if err := binary.Read(file, binary.LittleEndian, &b); err != nil {
		fmt.Println("error: failed to read header!", err)
		return nil, err
	}

	// resize to maxHeaderLen
	return expandByteSlice(b, maxHeaderLen), nil
}

func fileGetName(file *os.File) string {
	stat, _ := file.Stat()
	return stat.Name()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sync"

	"github.com/nsf/termbox-go"

	"github.com/prattmic/lesser/lineio"
)

type size struct {
	x int
	y int
}

type Event int

const (
	// EventQuit requests an application exit.
	EventQuit Event = iota

	// EventRefresh requests a display refresh.
	EventRefresh
)

type Mode int

const (
	// ModeNormal is the standard mode, allowing file navigation.
	ModeNormal Mode = iota

	// ModeSearchEntry is search entry mode. Key presses are added
	// to the search string.
	ModeSearchEntry
)

type Lesser struct {
	// src is the source file being displayed.
	src lineio.LineReader

	// tabStop is the number of spaces per tab.
	tabStop int

	// events is used to notify the main goroutine of events.
	events chan Event

	// mu locks the fields below.
	mu sync.Mutex

	// size is the size of the file display.
	// There is a statusbar beneath the display.
	size size

	// line is the line number of the first line of the display.
	line int64

	// mode is the viewer mode.
	mode Mode

	// regexp is the search regexp specified by the user.
	// Must only be modified by the event goroutine.
	regexp string
}

// lastLine returns the last line on the display.  It may be beyond the end
// of the file, if the file is short enough.
// mu must be held on call.
func (l *Lesser) lastLine() int64 {
	return l.line + int64(l.size.y) - 1
}

// scrollUp moves the display up (i.e., decrements the first line number).
// You cannot scroll beyond the beginning of the file.
// refreshScreen must be called for the display to actually update.
func (l *Lesser) scrollUp() {
	l.mu.Lock()
	defer l.mu.Unlock()

	if l.line > 1 {
		l.line--
	}
}

// scrollDown moves the display down (i.e., increments the first line number).
// refreshScreen must be called for the display to actually update.
func (l *Lesser) scrollDown() {
	l.mu.Lock()
	defer l.mu.Unlock()

	// You can only scroll down if the next line exists.
	if l.src.LineExists(l.lastLine() + 1) {
		l.line++
	}
}

// scrollTop moves to the first line.
// refreshScreen must be called for the display to actually update.
func (l *Lesser) scrollTop() {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.line = 1
}

// scrollBottom moves to the last line.
// refreshScreen must be called for the display to actually update.
func (l *Lesser) scrollBottom() {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.line = 1

	// TODO(prattmic): binary search
	for l.src.LineExists(l.lastLine() + 1) {
		l.line++
	}
}

func (l *Lesser) handleEvent(e termbox.Event) {
	l.mu.Lock()
	mode := l.mode
	l.mu.Unlock()

	if e.Type != termbox.EventKey {
		return
	}

	c := e.Ch

	switch mode {
	case ModeNormal:
		switch c {
		case 'q':
			l.events <- EventQuit
		case 'j':
			l.scrollDown()
			l.events <- EventRefresh
		case 'k':
			l.scrollUp()
			l.events <- EventRefresh
		case 'g':
			l.scrollTop()
			l.events <- EventRefresh
		case 'G':
			l.scrollBottom()
			l.events <- EventRefresh
		case '/':
			l.mu.Lock()
			l.mode = ModeSearchEntry
			l.mu.Unlock()
			l.events <- EventRefresh
		}
	case ModeSearchEntry:
		switch c {
		case 0:
			switch e.Key {
			case termbox.KeyEnter:
				l.search(l.regexp)
				l.mu.Lock()
				l.mode = ModeNormal
				l.regexp = ""
				l.mu.Unlock()
				l.events <- EventRefresh
			}
		default:
			l.mu.Lock()
			l.regexp += string(c)
			l.mu.Unlock()
			l.events <- EventRefresh
		}
	}
}

func (l *Lesser) listenEvents() {
	for {
		e := termbox.PollEvent()
		l.handleEvent(e)
	}
}

type searchResult struct {
	line    int64
	matches [][]int
	err     error
}

func (l *Lesser) search(s string) {
	reg, err := regexp.Compile(s)
	if err != nil {
		// TODO(prattmic): display a better error
		log.Printf("regexp failed to compile: %v", err)
		return
	}

	results := make(chan searchResult, 100)

	searchLine := func(line int64) {
		r, err := l.src.SearchLine(reg, line)
		if err != nil {
			r = nil
		}

		results <- searchResult{
			line:    line,
			matches: r,
			err:     err,
		}
	}

	nextLine := int64(1)
	// Spawn initial search goroutines
	for ; nextLine <= 5; nextLine++ {
		go searchLine(nextLine)
	}

	all := make([]searchResult, 0)

	// Collect results, start searching next lines until we start
	// hitting EOF.
	for {
		ret := <-results
		all = append(all, ret)

		// We started hitting errors on a previous line,
		// there is no reason to search later lines.
		if ret.err != nil {
			break
		}

		go searchLine(nextLine)
		nextLine++
	}

	// Collect the remaing results.
	for int64(len(all)) < nextLine-1 {
		all = append(all, <-results)
	}

	log.Printf("Results: %+v\n", all)
}

// statusBar renders the status bar.
// mu must be held on call.
func (l *Lesser) statusBar() {
	// The statusbar is just below the display.

	// Clear the statusbar
	for i := 0; i < l.size.x; i++ {
		termbox.SetCell(i, l.size.y, ' ', 0, 0)
	}

	switch l.mode {
	case ModeNormal:
		// Just a colon and a cursor
		termbox.SetCell(0, l.size.y, ':', 0, 0)
		termbox.SetCursor(1, l.size.y)
	case ModeSearchEntry:
		// / and search string
		termbox.SetCell(0, l.size.y, '/', 0, 0)
		for i, c := range l.regexp {
			termbox.SetCell(1+i, l.size.y, c, 0, 0)
		}
		termbox.SetCursor(1+len(l.regexp), l.size.y)
	}
}

func (l *Lesser) refreshScreen() error {
	l.mu.Lock()
	defer l.mu.Unlock()

	for y := 0; y < l.size.y; y++ {
		buf := make([]byte, l.size.x)

		_, err := l.src.ReadLine(buf, l.line+int64(y))
		// EOF just means the line was shorter than the display.
		if err != nil && err != io.EOF {
			return err
		}

		var tabOffset int
		for i, c := range buf {
			// If there are tabs, we may get to the end of the
			// display before we run out of characters.
			if i >= l.size.x {
				break
			}

			if c == '\t' {
				// Clear the tab spaces
				for j := 0; j < l.tabStop; j++ {
					termbox.SetCell(tabOffset+j, y, ' ', 0, 0)
				}
				tabOffset += l.tabStop - 1
			} else {
				termbox.SetCell(tabOffset+i, y, rune(c), 0, 0)
			}
		}
	}

	l.statusBar()

	termbox.Flush()

	return nil
}

func (l *Lesser) Run() {
	go l.listenEvents()

	err := l.refreshScreen()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to refresh screen: %v\n", err)
		return
	}

	for {
		e := <-l.events

		switch e {
		case EventQuit:
			return
		case EventRefresh:
			err = l.refreshScreen()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Failed to refresh screen: %v\n", err)
				return
			}
		}
	}
}

func NewLesser(f *os.File, ts int) Lesser {
	x, y := termbox.Size()

	return Lesser{
		src:     lineio.NewLineReader(f),
		tabStop: ts,
		// Save one line for statusbar.
		size:   size{x: x, y: y - 1},
		line:   1,
		events: make(chan Event, 1),
		mode:   ModeNormal,
	}
}
<commit_msg>Highlight search results<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sync"

	"github.com/nsf/termbox-go"

	"github.com/prattmic/lesser/lineio"
)

type size struct {
	x int
	y int
}

type Event int

const (
	// EventQuit requests an application exit.
	EventQuit Event = iota

	// EventRefresh requests a display refresh.
	EventRefresh
)

type Mode int

const (
	// ModeNormal is the standard mode, allowing file navigation.
	ModeNormal Mode = iota

	// ModeSearchEntry is search entry mode. Key presses are added
	// to the search string.
	ModeSearchEntry
)

type Lesser struct {
	// src is the source file being displayed.
	src lineio.LineReader

	// tabStop is the number of spaces per tab.
	tabStop int

	// events is used to notify the main goroutine of events.
	events chan Event

	// mu locks the fields below.
	mu sync.Mutex

	// size is the size of the file display.
	// There is a statusbar beneath the display.
	size size

	// line is the line number of the first line of the display.
	line int64

	// mode is the viewer mode.
	mode Mode

	// regexp is the search regexp specified by the user.
	// Must only be modified by the event goroutine.
	regexp string

	// searchResults are the results for the current search.
	// They should be highlighted.
	searchResults searchResults
}

// lastLine returns the last line on the display.  It may be beyond the end
// of the file, if the file is short enough.
// mu must be held on call.
func (l *Lesser) lastLine() int64 {
	return l.line + int64(l.size.y) - 1
}

// scrollUp moves the display up (i.e., decrements the first line number).
// You cannot scroll beyond the beginning of the file.
// refreshScreen must be called for the display to actually update.
func (l *Lesser) scrollUp() {
	l.mu.Lock()
	defer l.mu.Unlock()

	if l.line > 1 {
		l.line--
	}
}

// scrollDown moves the display down (i.e., increments the first line number).
// refreshScreen must be called for the display to actually update.
func (l *Lesser) scrollDown() {
	l.mu.Lock()
	defer l.mu.Unlock()

	// You can only scroll down if the next line exists.
	if l.src.LineExists(l.lastLine() + 1) {
		l.line++
	}
}

// scrollTop moves to the first line.
// refreshScreen must be called for the display to actually update.
func (l *Lesser) scrollTop() {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.line = 1
}

// scrollBottom moves to the last line.
// refreshScreen must be called for the display to actually update.
func (l *Lesser) scrollBottom() {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.line = 1

	// TODO(prattmic): binary search
	for l.src.LineExists(l.lastLine() + 1) {
		l.line++
	}
}

func (l *Lesser) handleEvent(e termbox.Event) {
	l.mu.Lock()
	mode := l.mode
	l.mu.Unlock()

	if e.Type != termbox.EventKey {
		return
	}

	c := e.Ch

	switch mode {
	case ModeNormal:
		switch c {
		case 'q':
			l.events <- EventQuit
		case 'j':
			l.scrollDown()
			l.events <- EventRefresh
		case 'k':
			l.scrollUp()
			l.events <- EventRefresh
		case 'g':
			l.scrollTop()
			l.events <- EventRefresh
		case 'G':
			l.scrollBottom()
			l.events <- EventRefresh
		case '/':
			l.mu.Lock()
			l.mode = ModeSearchEntry
			l.mu.Unlock()
			l.events <- EventRefresh
		}
	case ModeSearchEntry:
		switch c {
		case 0:
			switch e.Key {
			case termbox.KeyEnter:
				r := l.search(l.regexp)
				l.mu.Lock()
				l.mode = ModeNormal
				l.regexp = ""
				l.searchResults = r
				l.mu.Unlock()
				l.events <- EventRefresh
			}
		default:
			l.mu.Lock()
			l.regexp += string(c)
			l.mu.Unlock()
			l.events <- EventRefresh
		}
	}
}

func (l *Lesser) listenEvents() {
	for {
		e := termbox.PollEvent()
		l.handleEvent(e)
	}
}

// searchResult describes search matches on a single line.
type searchResult struct {
	line    int64
	matches [][]int
	err     error
}

// matchesChar returns true if the search result contains a match for
// character index c.
func (s searchResult) matchesChar(c int) bool {
	for _, match := range s.matches {
		if len(match) < 2 {
			continue
		}

		if c >= match[0] && c < match[1] {
			return true
		}
	}
	return false
}

type searchResults []searchResult

// findLine finds the result for a specific line, if any.
// TODO(prattmic): Make a much more efficient data structure.
func (s searchResults) findLine(line int64) searchResult {
	for _, r := range s {
		if r.line == line {
			return r
		}
	}

	return searchResult{}
}

func (l *Lesser) search(s string) []searchResult {
	reg, err := regexp.Compile(s)
	if err != nil {
		// TODO(prattmic): display a better error
		log.Printf("regexp failed to compile: %v", err)
		return nil
	}

	results := make(chan searchResult, 100)

	searchLine := func(line int64) {
		r, err := l.src.SearchLine(reg, line)
		if err != nil {
			r = nil
		}

		results <- searchResult{
			line:    line,
			matches: r,
			err:     err,
		}
	}

	nextLine := int64(1)
	// Spawn initial search goroutines
	for ; nextLine <= 5; nextLine++ {
		go searchLine(nextLine)
	}

	all := make([]searchResult, 0)

	// Collect results, start searching next lines until we start
	// hitting EOF.
	for {
		ret := <-results
		all = append(all, ret)

		// We started hitting errors on a previous line,
		// there is no reason to search later lines.
		if ret.err != nil {
			break
		}

		go searchLine(nextLine)
		nextLine++
	}

	// Collect the remaing results.
	for int64(len(all)) < nextLine-1 {
		all = append(all, <-results)
	}

	return all
}

// statusBar renders the status bar.
// mu must be held on call.
func (l *Lesser) statusBar() {
	// The statusbar is just below the display.

	// Clear the statusbar
	for i := 0; i < l.size.x; i++ {
		termbox.SetCell(i, l.size.y, ' ', 0, 0)
	}

	switch l.mode {
	case ModeNormal:
		// Just a colon and a cursor
		termbox.SetCell(0, l.size.y, ':', 0, 0)
		termbox.SetCursor(1, l.size.y)
	case ModeSearchEntry:
		// / and search string
		termbox.SetCell(0, l.size.y, '/', 0, 0)
		for i, c := range l.regexp {
			termbox.SetCell(1+i, l.size.y, c, 0, 0)
		}
		termbox.SetCursor(1+len(l.regexp), l.size.y)
	}
}

func (l *Lesser) refreshScreen() error {
	l.mu.Lock()
	defer l.mu.Unlock()

	for y := 0; y < l.size.y; y++ {
		buf := make([]byte, l.size.x)
		line := l.line + int64(y)

		_, err := l.src.ReadLine(buf, line)
		// EOF just means the line was shorter than the display.
		if err != nil && err != io.EOF {
			return err
		}

		highlight := l.searchResults.findLine(line)

		var tabOffset int
		for i, c := range buf {
			// If there are tabs, we may get to the end of the
			// display before we run out of characters.
			if i >= l.size.x {
				break
			}

			fg := termbox.ColorDefault
			bg := termbox.ColorDefault

			// Highlight matches
			if highlight.matchesChar(i) {
				fg = termbox.ColorBlack
				bg = termbox.ColorWhite
			}

			if c == '\t' {
				// Clear the tab spaces
				for j := 0; j < l.tabStop; j++ {
					termbox.SetCell(tabOffset+j, y, ' ', 0, 0)
				}
				tabOffset += l.tabStop - 1
			} else {
				termbox.SetCell(tabOffset+i, y, rune(c), fg, bg)
			}
		}
	}

	l.statusBar()

	termbox.Flush()

	return nil
}

func (l *Lesser) Run() {
	go l.listenEvents()

	err := l.refreshScreen()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to refresh screen: %v\n", err)
		return
	}

	for {
		e := <-l.events

		switch e {
		case EventQuit:
			return
		case EventRefresh:
			err = l.refreshScreen()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Failed to refresh screen: %v\n", err)
				return
			}
		}
	}
}

func NewLesser(f *os.File, ts int) Lesser {
	x, y := termbox.Size()

	return Lesser{
		src:     lineio.NewLineReader(f),
		tabStop: ts,
		// Save one line for statusbar.
		size:   size{x: x, y: y - 1},
		line:   1,
		events: make(chan Event, 1),
		mode:   ModeNormal,
	}
}
<|endoftext|>


<commit_before>package bagman_test

import (
	"github.com/APTrust/bagman/bagman"
	"path/filepath"
	"strings"
	"testing"
)

/* ------ NOTE -------
Many vars used here, such as gopath, testDataPath, sampleGood, sampleNoBagit, etc.
are defined in bag_test.
*/

// Make sure we can untar properly formatted tar files.
func TestValidTarFile(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == false {
		t.Errorf("Bag should be valid, but validator says it isn't: %s",
			validator.ErrorMessage)
	}
}

// Make sure Untar doesn't blow up when it gets an invalid
// or corrupt tar file. It should return a TarResult with an
// Error property.
func TestInvalidTarFile(t *testing.T) {
	validator, err := bagman.NewValidator(invalidTarFile)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag should NOT be valid, but validator says it is")
	}
}

// Make sure each of the bad bags produces an error in the BagReadResult.
// The underlying bagins library prints some warnings to stderr, so we
// include a note that those are expected.
func TestBadBagsReturnError(t *testing.T) {
	for _, tarFile := range badFiles {
		validator, err := bagman.NewValidator(tarFile)
		if err != nil {
			t.Errorf("Error creating validator: %s", err)
			return
		}
		if validator.IsValid() == true {
			t.Errorf("Bag '%s' should NOT be valid, but validator says it is", tarFile)
		}
		if validator.ErrorMessage == "" {
			t.Errorf("Invalid bag '%s' should have a specific error message", tarFile)
		}
	}
}

func TestBadFolderName(t *testing.T) {
	validator, err := bagman.NewValidator(sampleWrongFolderName)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleWrongFolderName)
	}
	if !strings.Contains(validator.ErrorMessage, "should untar to a folder named") {
		t.Errorf("Untarring file '%s' should have generated an 'incorrect file name' error.",
			sampleWrongFolderName)
	}
}

func TestBadAccessValue(t *testing.T) {
	validator, err := bagman.NewValidator(sampleBadAccess)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleWrongFolderName)
	}
	if !strings.Contains(validator.ErrorMessage, "access (rights) value") {
		t.Errorf("File '%s' should have generated an 'invalid access value' error.",
			sampleBadAccess)
	}
}

func TestMissingTitle(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoTitle)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoTitle)
	}
	if !strings.Contains(validator.ErrorMessage, "Title is missing") {
		t.Errorf("File '%s' should have generated a missing title error.",
			sampleNoTitle)
	}
}

func TestBadChecksums(t *testing.T) {
	validator, err := bagman.NewValidator(sampleBadChecksums)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoTitle)
	}
	if !strings.Contains(validator.ErrorMessage, "The following checksums could not be verified:") {
		t.Errorf("File '%s' should have generated a missing title error.",
			sampleBadChecksums)
	}
}

func TestMissingDataFile(t *testing.T) {
	validator, err := bagman.NewValidator(sampleMissingDataFile)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleMissingDataFile)
	}
	if !strings.Contains(validator.ErrorMessage, "The following checksums could not be verified:") {
		t.Errorf("File '%s' should have generated an error describing a missing file.",
			sampleMissingDataFile)
	}
}

func TestNoBagInfo(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoBagInfo)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoBagInfo)
	}
	if !strings.Contains(validator.ErrorMessage, "Unable to find tagfile bag-info.txt") {
		t.Errorf("File '%s' should have generated an error describing a missing bag-info.txt file.",
			sampleNoBagInfo)
	}
}

func TestNoBagit(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoBagit)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoBagit)
	}
	if !strings.Contains(validator.ErrorMessage, "Bag is missing bagit.txt file") {
		t.Errorf("File '%s' should have generated an error describing a missing bagit.txt file.",
			sampleNoBagit)
	}
}

func TestNoMd5Manifest(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoMd5Manifest)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoMd5Manifest)
	}
	if validator.ErrorMessage != "Required checksum file manifest-md5.txt is missing." {
		t.Errorf("File '%s' should have generated an error describing a missing manifest-md5.txt file.",
			sampleNoMd5Manifest)
		t.Errorf(validator.ErrorMessage)
	}
}

func TestNoAPTrustInfo(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoAPTrustInfo)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoAPTrustInfo)
	}
	if !strings.Contains(validator.ErrorMessage, "Unable to find tagfile aptrust-info.txt") {
		t.Errorf("File '%s' should have generated an error describing a missing aptrust-info.txt file.",
			sampleNoAPTrustInfo)
		t.Errorf(validator.ErrorMessage)
	}
}

func TestNoDataDir(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoDataDir)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoDataDir)
	}
	if validator.ErrorMessage != "Bag is missing the data directory, which should contain the payload files." {
		t.Errorf("File '%s' should have generated an error describing a missing data directory.",
			sampleNoDataDir)
		t.Errorf(validator.ErrorMessage)
	}
}

// Make sure we can validate a bag that is not yet tarred.
func TestValidateFromDirectory(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGoodUntarred)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == false {
		t.Errorf("%s should be valid, but got error: %s",
			sampleGoodUntarred, validator.ErrorMessage)
	}
}

func TestUntarredDir(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if strings.HasSuffix(validator.UntarredDir(), "bagman/testdata/example.edu.sample_good") == false {
		t.Errorf("UntarredDir() should be 'bagman/testdata/example.edu.sample_good', but got %s", validator.UntarredDir())
	}
}

func TestInstitutionDomain(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	domain, err := validator.InstitutionDomain()
	if err != nil {
		t.Error(err)
	}
	if domain != "example.edu" {
		t.Errorf("InstitutionDomain() should have returned 'example.edu', but got %s",
			domain)
	}
}

func TestTarFileName(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.TarFileName() != "example.edu.sample_good.tar" {
		t.Errorf("Tar file name should be example.edu.sample_good.tar")
	}

	validator, err = bagman.NewValidator(sampleGoodUntarred)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.TarFileName() != "example.edu.sample_good.tar" {
		t.Errorf("Tar file name should be example.edu.sample_good.tar")
	}

}

func TestFileType(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	fileType, err := validator.FileType()
	if err != nil {
		t.Error(err)
	}
	if fileType != bagman.VAL_TYPE_TAR {
		t.Errorf("File type for %s should be tar", sampleGood)
	}

	validator, err = bagman.NewValidator(sampleGoodUntarred)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	fileType, err = validator.FileType()
	if err != nil {
		t.Error(err)
	}
	if fileType != bagman.VAL_TYPE_DIR {
		t.Errorf("File type for %s should be directory", sampleGoodUntarred)
	}

	pathToJsonFile := filepath.Join(testDataPath, "intel_obj.json")
	validator, err = bagman.NewValidator(pathToJsonFile)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	fileType, err = validator.FileType()
	if err == nil {
		t.Error("Invalid file type should have returned an error.")
	}
	if fileType != bagman.VAL_TYPE_ERR {
		t.Errorf("File type for %s should be invalid", pathToJsonFile)
	}

}
<commit_msg>Test that missing bag is reported correctly<commit_after>package bagman_test

import (
	"github.com/APTrust/bagman/bagman"
	"path/filepath"
	"strings"
	"testing"
)

/* ------ NOTE -------
Many vars used here, such as gopath, testDataPath, sampleGood, sampleNoBagit, etc.
are defined in bag_test.
*/

// Make sure we can untar properly formatted tar files.
func TestValidTarFile(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == false {
		t.Errorf("Bag should be valid, but validator says it isn't: %s",
			validator.ErrorMessage)
	}
}

// Make sure Untar doesn't blow up when it gets an invalid
// or corrupt tar file. It should return a TarResult with an
// Error property.
func TestInvalidTarFile(t *testing.T) {
	validator, err := bagman.NewValidator(invalidTarFile)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag should NOT be valid, but validator says it is")
	}
}

// Make sure each of the bad bags produces an error in the BagReadResult.
// The underlying bagins library prints some warnings to stderr, so we
// include a note that those are expected.
func TestBadBagsReturnError(t *testing.T) {
	for _, tarFile := range badFiles {
		validator, err := bagman.NewValidator(tarFile)
		if err != nil {
			t.Errorf("Error creating validator: %s", err)
			return
		}
		if validator.IsValid() == true {
			t.Errorf("Bag '%s' should NOT be valid, but validator says it is", tarFile)
		}
		if validator.ErrorMessage == "" {
			t.Errorf("Invalid bag '%s' should have a specific error message", tarFile)
		}
	}
}

func TestBadFolderName(t *testing.T) {
	validator, err := bagman.NewValidator(sampleWrongFolderName)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleWrongFolderName)
	}
	if !strings.Contains(validator.ErrorMessage, "should untar to a folder named") {
		t.Errorf("Untarring file '%s' should have generated an 'incorrect file name' error.",
			sampleWrongFolderName)
	}
}

func TestBadAccessValue(t *testing.T) {
	validator, err := bagman.NewValidator(sampleBadAccess)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleWrongFolderName)
	}
	if !strings.Contains(validator.ErrorMessage, "access (rights) value") {
		t.Errorf("File '%s' should have generated an 'invalid access value' error.",
			sampleBadAccess)
	}
}

func TestMissingTitle(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoTitle)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoTitle)
	}
	if !strings.Contains(validator.ErrorMessage, "Title is missing") {
		t.Errorf("File '%s' should have generated a missing title error.",
			sampleNoTitle)
	}
}

func TestBadChecksums(t *testing.T) {
	validator, err := bagman.NewValidator(sampleBadChecksums)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoTitle)
	}
	if !strings.Contains(validator.ErrorMessage, "The following checksums could not be verified:") {
		t.Errorf("File '%s' should have generated a missing title error.",
			sampleBadChecksums)
	}
}

func TestMissingDataFile(t *testing.T) {
	validator, err := bagman.NewValidator(sampleMissingDataFile)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleMissingDataFile)
	}
	if !strings.Contains(validator.ErrorMessage, "The following checksums could not be verified:") {
		t.Errorf("File '%s' should have generated an error describing a missing file.",
			sampleMissingDataFile)
	}
}

func TestNoBagInfo(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoBagInfo)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoBagInfo)
	}
	if !strings.Contains(validator.ErrorMessage, "Unable to find tagfile bag-info.txt") {
		t.Errorf("File '%s' should have generated an error describing a missing bag-info.txt file.",
			sampleNoBagInfo)
	}
}

func TestNoBagit(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoBagit)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoBagit)
	}
	if !strings.Contains(validator.ErrorMessage, "Bag is missing bagit.txt file") {
		t.Errorf("File '%s' should have generated an error describing a missing bagit.txt file.",
			sampleNoBagit)
	}
}

func TestNoMd5Manifest(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoMd5Manifest)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoMd5Manifest)
	}
	if validator.ErrorMessage != "Required checksum file manifest-md5.txt is missing." {
		t.Errorf("File '%s' should have generated an error describing a missing manifest-md5.txt file.",
			sampleNoMd5Manifest)
		t.Errorf(validator.ErrorMessage)
	}
}

func TestNoAPTrustInfo(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoAPTrustInfo)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoAPTrustInfo)
	}
	if !strings.Contains(validator.ErrorMessage, "Unable to find tagfile aptrust-info.txt") {
		t.Errorf("File '%s' should have generated an error describing a missing aptrust-info.txt file.",
			sampleNoAPTrustInfo)
		t.Errorf(validator.ErrorMessage)
	}
}

func TestNoDataDir(t *testing.T) {
	validator, err := bagman.NewValidator(sampleNoDataDir)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == true {
		t.Errorf("Bag '%s' should NOT be valid, but validator says it is", sampleNoDataDir)
	}
	if validator.ErrorMessage != "Bag is missing the data directory, which should contain the payload files." {
		t.Errorf("File '%s' should have generated an error describing a missing data directory.",
			sampleNoDataDir)
		t.Errorf(validator.ErrorMessage)
	}
}

// Make sure we can validate a bag that is not yet tarred.
func TestValidateFromDirectory(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGoodUntarred)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() == false {
		t.Errorf("%s should be valid, but got error: %s",
			sampleGoodUntarred, validator.ErrorMessage)
	}
}

func TestUntarredDir(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if strings.HasSuffix(validator.UntarredDir(), "bagman/testdata/example.edu.sample_good") == false {
		t.Errorf("UntarredDir() should be 'bagman/testdata/example.edu.sample_good', but got %s", validator.UntarredDir())
	}
}

func TestNonExistentFile(t *testing.T) {
	validator, err := bagman.NewValidator("Example.Edu.SomeBogusFile.tar")
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.IsValid() {
		t.Errorf("Non-existent file should not validate!")
	}
	if !strings.Contains(validator.ErrorMessage, "no such file or directory") {
		t.Errorf("Error message for missing bag file should say the file is missing!")
	}
}

func TestInstitutionDomain(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	domain, err := validator.InstitutionDomain()
	if err != nil {
		t.Error(err)
	}
	if domain != "example.edu" {
		t.Errorf("InstitutionDomain() should have returned 'example.edu', but got %s",
			domain)
	}
}

func TestTarFileName(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.TarFileName() != "example.edu.sample_good.tar" {
		t.Errorf("Tar file name should be example.edu.sample_good.tar")
	}

	validator, err = bagman.NewValidator(sampleGoodUntarred)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	if validator.TarFileName() != "example.edu.sample_good.tar" {
		t.Errorf("Tar file name should be example.edu.sample_good.tar")
	}

}

func TestFileType(t *testing.T) {
	validator, err := bagman.NewValidator(sampleGood)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	fileType, err := validator.FileType()
	if err != nil {
		t.Error(err)
	}
	if fileType != bagman.VAL_TYPE_TAR {
		t.Errorf("File type for %s should be tar", sampleGood)
	}

	validator, err = bagman.NewValidator(sampleGoodUntarred)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	fileType, err = validator.FileType()
	if err != nil {
		t.Error(err)
	}
	if fileType != bagman.VAL_TYPE_DIR {
		t.Errorf("File type for %s should be directory", sampleGoodUntarred)
	}

	pathToJsonFile := filepath.Join(testDataPath, "intel_obj.json")
	validator, err = bagman.NewValidator(pathToJsonFile)
	if err != nil {
		t.Errorf("Error creating validator: %s", err)
		return
	}
	fileType, err = validator.FileType()
	if err == nil {
		t.Error("Invalid file type should have returned an error.")
	}
	if fileType != bagman.VAL_TYPE_ERR {
		t.Errorf("File type for %s should be invalid", pathToJsonFile)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"github.com/btcsuite/btcd/btcec"
	"golang.org/x/crypto/ripemd160"
)

// TODO(roasbeef): Might need to change? due to the PRG* requirements? But
// yeh, once I add the LN specific stuff FSLength will change. Also doesn't
// 20 hops seem a bit excessive? 7?
const NumMaxHops = 20
const FSLength = 48

// Hmm appears that they use k = 128 throughout the paper?

// HMAC -> SHA-256
//  * or could use Poly1035: https://godoc.org/golang.org/x/crypto/poly1305
//  * but, the paper specs: {0, 1}^k x {0, 1}* -> {0, 1}^k
//  * Poly1035 is actually: {0, 1}^k x {0, 1}* -> {0, 1}^(2/k)
//  * Also with Poly, I guess the key is treated as a nonce, tagging two messages
//    with the same key allows an attacker to forge message or something like that

// Size of a forwarding segment is 32 bytes, the MAC is 16 bytes, so c = 48 bytes
//  * NOTE: this doesn't include adding R to the forwarding segment, and w/e esle

// Hmmm since each uses diff key, just use AES-CTR with blank nonce, given key,
// encrypt plaintext of all zeroes, this'll give us our len(plaintext) rand bytes.
// PRG0 -> {0, 1}^k -> {0, 1}^r(c+k) or {0, 1}^1280  (assuming 20 hops, like rusty, but, is that too large? maybe, idk)
// PRG1 -> {0, 1}^k -> {0, 1}^r(c+k) or {0, 1}^1280  (assuming 20 hops)
// PRG2 -> {0, 1}^k -> {0, 1}^rc or {0, 1}^960 (assuming 20 hops, c=48)
//  * NOTE: in second version of paper (accepted to CCS'15), all the PRG*'s are like PRG2
//  * so makes it simpler

// PRP -> AES? or
//  * {0, 1}^k x {0, 1}^a -> {0, 1}^a

// Do we need AEAD for the below? Or are is the per-hop MAC okay?
// ENC: AES-CTR or CHACHA20?

// DEC: AES-CTR or CHACHA20?

// h_op: G^* -> {0, 1}^k
//  * op (elem of) {MAC, PRGO, PRG!, PRP, ENC, DEC}
//  * key gen for the above essentially

// RoutingSegment...
// NOTE: Length of routing segment in the paper is 8 bytes (enough for their
// imaginary network, I guess). But, looking like they'll be (20 + 33 bytes)
// 53 bytes. Or 52 if we use curve25519
type RoutingSegment struct {
	nextHop *btcec.PublicKey // NOTE: or, is this a LN addr? w/e that is?
	// nextHop [32]byte
	rCommitment [ripemd160.Size]byte

	// stuff perhaps?
}

// ForwardingSegment....
type ForwardingSegment struct {
	// Here's hash(R), attempt to make an HTLC with the next hop. If
	// successful, then pass along the onion so we can finish getting the
	// payment circuit set up.
	// TODO(roasbeef): Do we create HTLC's with the minimum amount
	// possible? 1 satoshi or is it 1 mili-satoshi?
	rs RoutingSegment

	// To defend against replay attacks. Intermediate nodes will drop the
	// FS if it deems it's expired.
	expiration uint64

	// Key shared by intermediate node with the source, used to peel a layer
	// off the onion for the next hop.
	sharedSymmetricKey [32]byte // TODO(roasbeef): or, 16?
}

// AnonymousHeader...
type AnonymousHeader struct {
	// Forwarding info for the current hop. When serialized, it'll be
	// encrypted with SV, the secret key for this node known to no-one but
	// the node. It also contains a secret key shared with this node and the
	// source, so it can peel off a layer of the onion for the next hop.
	FS ForwardingSegment

	mac [32]byte // TODO(roasbeef): or, 16?
}

// CommonHeader...
type CommonHeader struct {
	// TODO(roasbeef): maybe can use this to extend HORNET with additiona control signals
	// for LN nodes?
	controlType uint8
	hops        uint8
	nonce       [8]byte // either interpreted as EXP or nonce, little-endian? idk
}

// DataPacket...
type DataPacket struct {
	Chdr  CommonHeader
	Ahdr  AnonymousHeader             // TODO(roasbeef): MAC in ahdr includes the chdr?
	Onion [FSLength * NumMaxHops]byte // TODO(roasbeef): or, is it NumMaxHops - 1?
}

type SphinxHeader struct {
}

// SessionSetupPacket...
type SessionSetupPacket struct {
	Chdr      CommonHeader
	Shdr      SphinxHeader
	Sp        SphinxPayload
	FsPayload [FSLength * NumMaxHops]byte // ? r*c
	// TODO(roabeef): hmm does this implcitly mean messages are a max of 48 bytes?
}

// will add other LN specific control packets below, I guess
<commit_msg>add sphinx payload struct back so hornet.go can compile<commit_after>package main

import (
	"github.com/btcsuite/btcd/btcec"
	"golang.org/x/crypto/ripemd160"
)

// TODO(roasbeef): Might need to change? due to the PRG* requirements? But
// yeh, once I add the LN specific stuff FSLength will change. Also doesn't
// 20 hops seem a bit excessive? 7?
const NumMaxHops = 20
const FSLength = 48

// Hmm appears that they use k = 128 throughout the paper?

// HMAC -> SHA-256
//  * or could use Poly1035: https://godoc.org/golang.org/x/crypto/poly1305
//  * but, the paper specs: {0, 1}^k x {0, 1}* -> {0, 1}^k
//  * Poly1035 is actually: {0, 1}^k x {0, 1}* -> {0, 1}^(2/k)
//  * Also with Poly, I guess the key is treated as a nonce, tagging two messages
//    with the same key allows an attacker to forge message or something like that

// Size of a forwarding segment is 32 bytes, the MAC is 16 bytes, so c = 48 bytes
//  * NOTE: this doesn't include adding R to the forwarding segment, and w/e esle

// Hmmm since each uses diff key, just use AES-CTR with blank nonce, given key,
// encrypt plaintext of all zeroes, this'll give us our len(plaintext) rand bytes.
// PRG0 -> {0, 1}^k -> {0, 1}^r(c+k) or {0, 1}^1280  (assuming 20 hops, like rusty, but, is that too large? maybe, idk)
// PRG1 -> {0, 1}^k -> {0, 1}^r(c+k) or {0, 1}^1280  (assuming 20 hops)
// PRG2 -> {0, 1}^k -> {0, 1}^rc or {0, 1}^960 (assuming 20 hops, c=48)
//  * NOTE: in second version of paper (accepted to CCS'15), all the PRG*'s are like PRG2
//  * so makes it simpler

// PRP -> AES? or
//  * {0, 1}^k x {0, 1}^a -> {0, 1}^a

// Do we need AEAD for the below? Or are is the per-hop MAC okay?
// ENC: AES-CTR or CHACHA20?

// DEC: AES-CTR or CHACHA20?

// h_op: G^* -> {0, 1}^k
//  * op (elem of) {MAC, PRGO, PRG!, PRP, ENC, DEC}
//  * key gen for the above essentially

// RoutingSegment...
// NOTE: Length of routing segment in the paper is 8 bytes (enough for their
// imaginary network, I guess). But, looking like they'll be (20 + 33 bytes)
// 53 bytes. Or 52 if we use curve25519
type RoutingSegment struct {
	nextHop *btcec.PublicKey // NOTE: or, is this a LN addr? w/e that is?
	// nextHop [32]byte
	rCommitment [ripemd160.Size]byte

	// stuff perhaps?
}

// SphinxPayload...
type SphinxPayload struct {
}

// ForwardingSegment....
type ForwardingSegment struct {
	// Here's hash(R), attempt to make an HTLC with the next hop. If
	// successful, then pass along the onion so we can finish getting the
	// payment circuit set up.
	// TODO(roasbeef): Do we create HTLC's with the minimum amount
	// possible? 1 satoshi or is it 1 mili-satoshi?
	rs RoutingSegment

	// To defend against replay attacks. Intermediate nodes will drop the
	// FS if it deems it's expired.
	expiration uint64

	// Key shared by intermediate node with the source, used to peel a layer
	// off the onion for the next hop.
	sharedSymmetricKey [32]byte // TODO(roasbeef): or, 16?
}

// AnonymousHeader...
type AnonymousHeader struct {
	// Forwarding info for the current hop. When serialized, it'll be
	// encrypted with SV, the secret key for this node known to no-one but
	// the node. It also contains a secret key shared with this node and the
	// source, so it can peel off a layer of the onion for the next hop.
	FS ForwardingSegment

	mac [32]byte // TODO(roasbeef): or, 16?
}

// CommonHeader...
type CommonHeader struct {
	// TODO(roasbeef): maybe can use this to extend HORNET with additiona control signals
	// for LN nodes?
	controlType uint8
	hops        uint8
	nonce       [8]byte // either interpreted as EXP or nonce, little-endian? idk
}

// DataPacket...
type DataPacket struct {
	Chdr  CommonHeader
	Ahdr  AnonymousHeader             // TODO(roasbeef): MAC in ahdr includes the chdr?
	Onion [FSLength * NumMaxHops]byte // TODO(roasbeef): or, is it NumMaxHops - 1?
}

type SphinxHeader struct {
}

// SessionSetupPacket...
type SessionSetupPacket struct {
	Chdr      CommonHeader
	Shdr      SphinxHeader
	Sp        SphinxPayload
	FsPayload [FSLength * NumMaxHops]byte // ? r*c
	// TODO(roabeef): hmm does this implcitly mean messages are a max of 48 bytes?
}

// will add other LN specific control packets below, I guess
<|endoftext|>


<commit_before>package shell_test

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"testing"
	"time"

	"github.com/buildkite/agent/bootstrap/shell"
	"github.com/buildkite/bintest"
)

func TestRunAndCaptureWithTTY(t *testing.T) {
	sshKeygen, err := bintest.CompileProxy("ssh-keygen")
	if err != nil {
		t.Fatal(err)
	}
	defer sshKeygen.Close()

	sh := newShellForTest(t)
	sh.PTY = true

	go func() {
		call := <-sshKeygen.Ch
		fmt.Fprintln(call.Stdout, "Llama party! 🎉")
		call.Exit(0)
	}()

	actual, err := sh.RunAndCapture(sshKeygen.Path, "-f", "my_hosts", "-F", "llamas.com")
	if err != nil {
		t.Error(err)
	}

	if expected := "Llama party! 🎉"; string(actual) != expected {
		t.Errorf("Expected %q, got %q", expected, actual)
	}
}

func TestRunAndCaptureWithExitCode(t *testing.T) {
	sshKeygen, err := bintest.CompileProxy("ssh-keygen")
	if err != nil {
		t.Fatal(err)
	}
	defer sshKeygen.Close()

	sh := newShellForTest(t)

	go func() {
		call := <-sshKeygen.Ch
		fmt.Fprintln(call.Stdout, "Llama drama! 🚨")
		call.Exit(24)
	}()

	_, err = sh.RunAndCapture(sshKeygen.Path)
	if err == nil {
		t.Error("Expected an error, got nil")
	}

	if exitCode := shell.GetExitCode(err); exitCode != 24 {
		t.Fatalf("Expected %d, got %d", 24, exitCode)
	}
}

func TestRun(t *testing.T) {
	sshKeygen, err := bintest.CompileProxy("ssh-keygen")
	if err != nil {
		t.Fatal(err)
	}
	defer sshKeygen.Close()

	out := &bytes.Buffer{}

	sh := newShellForTest(t)
	sh.PTY = false
	sh.Writer = out
	sh.Logger = &shell.WriterLogger{Writer: out, Ansi: false}

	go func() {
		call := <-sshKeygen.Ch
		fmt.Fprintln(call.Stdout, "Llama party! 🎉")
		call.Exit(0)
	}()

	if err = sh.Run(sshKeygen.Path, "-f", "my_hosts", "-F", "llamas.com"); err != nil {
		t.Fatal(err)
	}

	actual := out.String()

	promptPrefix := "$"
	if runtime.GOOS == "windows" {
		promptPrefix = ">"
	}

	if expected := promptPrefix + " " + sshKeygen.Path + " -f my_hosts -F llamas.com\nLlama party! 🎉\n"; actual != expected {
		t.Fatalf("Expected %q, got %q", expected, actual)
	}
}

func TestContextCancelTerminates(t *testing.T) {
	sleepCmd, err := bintest.CompileProxy("sleep")
	if err != nil {
		t.Fatal(err)
	}
	defer sleepCmd.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sh, err := shell.NewWithContext(ctx)
	if err != nil {
		t.Fatal(err)
	}

	sh.Logger = shell.DiscardLogger

	go func() {
		call := <-sleepCmd.Ch
		time.Sleep(time.Second * 60)
		call.Exit(0)
	}()

	cancel()

	err = sh.Run(sleepCmd.Path)
	if !shell.IsExitSignaled(err) {
		t.Fatalf("Expected signal exit, got %#v", err)
	}
}

func TestInterrupt(t *testing.T) {
	sleepCmd, err := bintest.CompileProxy("sleep")
	if err != nil {
		t.Fatal(err)
	}
	defer sleepCmd.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sh, err := shell.NewWithContext(ctx)
	if err != nil {
		t.Fatal(err)
	}

	sh.Logger = shell.DiscardLogger

	go func() {
		call := <-sleepCmd.Ch
		time.Sleep(time.Second * 10)
		call.Exit(0)
	}()

	// interrupt the process after 50ms
	go func() {
		<-time.After(time.Millisecond * 50)
		sh.Interrupt()
	}()

	err = sh.Run(sleepCmd.Path)
	if err == nil {
		t.Error("Expected an error")
	}
}

func TestDefaultWorkingDirFromSystem(t *testing.T) {
	sh, err := shell.New()
	if err != nil {
		t.Fatal(err)
	}

	currentWd, _ := os.Getwd()
	if actual := sh.Getwd(); actual != currentWd {
		t.Fatalf("Expected working dir %q, got %q", currentWd, actual)
	}
}

func TestWorkingDir(t *testing.T) {
	tempDir, err := ioutil.TempDir("", "shelltest")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tempDir)

	// macos has a symlinked temp dir
	if runtime.GOOS == "darwin" {
		tempDir, _ = filepath.EvalSymlinks(tempDir)
	}

	dirs := []string{tempDir, "my", "test", "dirs"}

	if err := os.MkdirAll(filepath.Join(dirs...), 0700); err != nil {
		t.Fatal(err)
	}

	currentWd, _ := os.Getwd()

	sh, err := shell.New()
	sh.Logger = shell.DiscardLogger

	if err != nil {
		t.Fatal(err)
	}

	for idx := range dirs {
		dir := filepath.Join(dirs[0 : idx+1]...)

		if err := sh.Chdir(dir); err != nil {
			t.Fatal(err)
		}

		if actual := sh.Getwd(); actual != dir {
			t.Fatalf("Expected working dir %q, got %q", dir, actual)
		}

		var out string

		// there is no pwd for windows, and getting it requires using a shell builtin
		if runtime.GOOS == "windows" {
			out, err = sh.RunAndCapture("cmd", "/c", "echo", "%cd%")
			if err != nil {
				t.Fatal(err)
			}
		} else {
			out, err = sh.RunAndCapture("pwd")
			if err != nil {
				t.Fatal(err)
			}
		}

		if actual := out; actual != dir {
			t.Fatalf("Expected working dir (from pwd command) %q, got %q", dir, actual)
		}
	}

	afterWd, _ := os.Getwd()
	if afterWd != currentWd {
		t.Fatalf("Expected working dir to be the same as before shell commands ran")
	}
}

func TestLockFileRetriesAndTimesOut(t *testing.T) {
	dir, err := ioutil.TempDir("", "shelltest")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	sh := newShellForTest(t)
	sh.Logger = shell.DiscardLogger

	lockPath := filepath.Join(dir, "my.lock")

	// acquire a lock in another process
	cmd, err := acquireLockInOtherProcess(lockPath)
	if err != nil {
		t.Fatal(err)
	}

	defer cmd.Process.Kill()

	// acquire lock
	_, err = sh.LockFile(lockPath, time.Second*2)
	if err != context.DeadlineExceeded {
		t.Fatalf("Expected DeadlineExceeded error, got %v", err)
	}
}

func acquireLockInOtherProcess(lockfile string) (*exec.Cmd, error) {
	cmd := exec.Command(os.Args[0], "-test.run=TestAcquiringLockHelperProcess", "--", lockfile)
	cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1"}

	err := cmd.Start()
	if err != nil {
		return cmd, err
	}

	// wait for the above process to get a lock
	for {
		if _, err = os.Stat(lockfile); os.IsNotExist(err) {
			time.Sleep(time.Millisecond * 10)
			continue
		}
		break
	}

	return cmd, nil
}

// TestAcquiringLockHelperProcess isn't a real test. It's used as a helper process
func TestAcquiringLockHelperProcess(t *testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}

	fileName := os.Args[len(os.Args)-1]
	sh := newShellForTest(t)

	log.Printf("Locking %s", fileName)
	if _, err := sh.LockFile(fileName, time.Second*10); err != nil {
		os.Exit(1)
	}

	log.Printf("Acquired lock %s", fileName)
	c := make(chan struct{})
	<-c
}

func newShellForTest(t *testing.T) *shell.Shell {
	sh, err := shell.New()
	if err != nil {
		t.Fatal(err)
	}
	sh.Logger = shell.DiscardLogger
	return sh
}
<commit_msg>Disable some more tests that fail under windows<commit_after>package shell_test

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"testing"
	"time"

	"github.com/buildkite/agent/bootstrap/shell"
	"github.com/buildkite/bintest"
)

func TestRunAndCaptureWithTTY(t *testing.T) {
	sshKeygen, err := bintest.CompileProxy("ssh-keygen")
	if err != nil {
		t.Fatal(err)
	}
	defer sshKeygen.Close()

	sh := newShellForTest(t)
	sh.PTY = true

	go func() {
		call := <-sshKeygen.Ch
		fmt.Fprintln(call.Stdout, "Llama party! 🎉")
		call.Exit(0)
	}()

	actual, err := sh.RunAndCapture(sshKeygen.Path, "-f", "my_hosts", "-F", "llamas.com")
	if err != nil {
		t.Error(err)
	}

	if expected := "Llama party! 🎉"; string(actual) != expected {
		t.Errorf("Expected %q, got %q", expected, actual)
	}
}

func TestRunAndCaptureWithExitCode(t *testing.T) {
	sshKeygen, err := bintest.CompileProxy("ssh-keygen")
	if err != nil {
		t.Fatal(err)
	}
	defer sshKeygen.Close()

	sh := newShellForTest(t)

	go func() {
		call := <-sshKeygen.Ch
		fmt.Fprintln(call.Stdout, "Llama drama! 🚨")
		call.Exit(24)
	}()

	_, err = sh.RunAndCapture(sshKeygen.Path)
	if err == nil {
		t.Error("Expected an error, got nil")
	}

	if exitCode := shell.GetExitCode(err); exitCode != 24 {
		t.Fatalf("Expected %d, got %d", 24, exitCode)
	}
}

func TestRun(t *testing.T) {
	sshKeygen, err := bintest.CompileProxy("ssh-keygen")
	if err != nil {
		t.Fatal(err)
	}
	defer sshKeygen.Close()

	out := &bytes.Buffer{}

	sh := newShellForTest(t)
	sh.PTY = false
	sh.Writer = out
	sh.Logger = &shell.WriterLogger{Writer: out, Ansi: false}

	go func() {
		call := <-sshKeygen.Ch
		fmt.Fprintln(call.Stdout, "Llama party! 🎉")
		call.Exit(0)
	}()

	if err = sh.Run(sshKeygen.Path, "-f", "my_hosts", "-F", "llamas.com"); err != nil {
		t.Fatal(err)
	}

	actual := out.String()

	promptPrefix := "$"
	if runtime.GOOS == "windows" {
		promptPrefix = ">"
	}

	if expected := promptPrefix + " " + sshKeygen.Path + " -f my_hosts -F llamas.com\nLlama party! 🎉\n"; actual != expected {
		t.Fatalf("Expected %q, got %q", expected, actual)
	}
}

func TestContextCancelTerminates(t *testing.T) {
	if runtime.GOOS == `windows` {
		t.Skip("Not supported in windows")
	}

	sleepCmd, err := bintest.CompileProxy("sleep")
	if err != nil {
		t.Fatal(err)
	}
	defer sleepCmd.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sh, err := shell.NewWithContext(ctx)
	if err != nil {
		t.Fatal(err)
	}

	sh.Logger = shell.DiscardLogger

	go func() {
		call := <-sleepCmd.Ch
		time.Sleep(time.Second * 60)
		call.Exit(0)
	}()

	cancel()

	err = sh.Run(sleepCmd.Path)
	if !shell.IsExitSignaled(err) {
		t.Fatalf("Expected signal exit, got %#v", err)
	}
}

func TestInterrupt(t *testing.T) {
	if runtime.GOOS == `windows` {
		t.Skip("Not supported in windows")
	}

	sleepCmd, err := bintest.CompileProxy("sleep")
	if err != nil {
		t.Fatal(err)
	}
	defer sleepCmd.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sh, err := shell.NewWithContext(ctx)
	if err != nil {
		t.Fatal(err)
	}

	sh.Logger = shell.DiscardLogger

	go func() {
		call := <-sleepCmd.Ch
		time.Sleep(time.Second * 10)
		call.Exit(0)
	}()

	// interrupt the process after 50ms
	go func() {
		<-time.After(time.Millisecond * 50)
		sh.Interrupt()
	}()

	err = sh.Run(sleepCmd.Path)
	if err == nil {
		t.Error("Expected an error")
	}
}

func TestDefaultWorkingDirFromSystem(t *testing.T) {
	sh, err := shell.New()
	if err != nil {
		t.Fatal(err)
	}

	currentWd, _ := os.Getwd()
	if actual := sh.Getwd(); actual != currentWd {
		t.Fatalf("Expected working dir %q, got %q", currentWd, actual)
	}
}

func TestWorkingDir(t *testing.T) {
	tempDir, err := ioutil.TempDir("", "shelltest")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tempDir)

	// macos has a symlinked temp dir
	if runtime.GOOS == "darwin" {
		tempDir, _ = filepath.EvalSymlinks(tempDir)
	}

	dirs := []string{tempDir, "my", "test", "dirs"}

	if err := os.MkdirAll(filepath.Join(dirs...), 0700); err != nil {
		t.Fatal(err)
	}

	currentWd, _ := os.Getwd()

	sh, err := shell.New()
	sh.Logger = shell.DiscardLogger

	if err != nil {
		t.Fatal(err)
	}

	for idx := range dirs {
		dir := filepath.Join(dirs[0 : idx+1]...)

		if err := sh.Chdir(dir); err != nil {
			t.Fatal(err)
		}

		if actual := sh.Getwd(); actual != dir {
			t.Fatalf("Expected working dir %q, got %q", dir, actual)
		}

		var out string

		// there is no pwd for windows, and getting it requires using a shell builtin
		if runtime.GOOS == "windows" {
			out, err = sh.RunAndCapture("cmd", "/c", "echo", "%cd%")
			if err != nil {
				t.Fatal(err)
			}
		} else {
			out, err = sh.RunAndCapture("pwd")
			if err != nil {
				t.Fatal(err)
			}
		}

		if actual := out; actual != dir {
			t.Fatalf("Expected working dir (from pwd command) %q, got %q", dir, actual)
		}
	}

	afterWd, _ := os.Getwd()
	if afterWd != currentWd {
		t.Fatalf("Expected working dir to be the same as before shell commands ran")
	}
}

func TestLockFileRetriesAndTimesOut(t *testing.T) {
	dir, err := ioutil.TempDir("", "shelltest")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	sh := newShellForTest(t)
	sh.Logger = shell.DiscardLogger

	lockPath := filepath.Join(dir, "my.lock")

	// acquire a lock in another process
	cmd, err := acquireLockInOtherProcess(lockPath)
	if err != nil {
		t.Fatal(err)
	}

	defer cmd.Process.Kill()

	// acquire lock
	_, err = sh.LockFile(lockPath, time.Second*2)
	if err != context.DeadlineExceeded {
		t.Fatalf("Expected DeadlineExceeded error, got %v", err)
	}
}

func acquireLockInOtherProcess(lockfile string) (*exec.Cmd, error) {
	cmd := exec.Command(os.Args[0], "-test.run=TestAcquiringLockHelperProcess", "--", lockfile)
	cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1"}

	err := cmd.Start()
	if err != nil {
		return cmd, err
	}

	// wait for the above process to get a lock
	for {
		if _, err = os.Stat(lockfile); os.IsNotExist(err) {
			time.Sleep(time.Millisecond * 10)
			continue
		}
		break
	}

	return cmd, nil
}

// TestAcquiringLockHelperProcess isn't a real test. It's used as a helper process
func TestAcquiringLockHelperProcess(t *testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}

	fileName := os.Args[len(os.Args)-1]
	sh := newShellForTest(t)

	log.Printf("Locking %s", fileName)
	if _, err := sh.LockFile(fileName, time.Second*10); err != nil {
		os.Exit(1)
	}

	log.Printf("Acquired lock %s", fileName)
	c := make(chan struct{})
	<-c
}

func newShellForTest(t *testing.T) *shell.Shell {
	sh, err := shell.New()
	if err != nil {
		t.Fatal(err)
	}
	sh.Logger = shell.DiscardLogger
	return sh
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 e2e

import (
	"github.com/onsi/ginkgo"

	"github.com/jetstack/cert-manager/test/e2e/framework"
	"github.com/jetstack/cert-manager/test/e2e/framework/addon"
)

var (
	cfg = framework.DefaultConfig
)

var _ = ginkgo.SynchronizedBeforeSuite(func() []byte {
	addon.InitGlobals(cfg)

	ginkgo.By("Provisioning shared cluster addons")

	err := addon.ProvisionGlobals(cfg)
	if err != nil {
		framework.Failf("Error provisioning global addons: %v", err)
	}

	return nil
}, func([]byte) {
	addon.InitGlobals(cfg)

	ginkgo.By("Configuring details for shared cluster addons")

	err := addon.SetupGlobals(cfg)
	if err != nil {
		framework.Failf("Error configuring global addons: %v", err)
	}
})

var globalLogs string

var _ = ginkgo.SynchronizedAfterSuite(func() {},
	func() {
		ginkgo.By("Retrieving logs for global addons")
		var err error
		globalLogs, err = addon.GlobalLogs()
		if err != nil {
			ginkgo.GinkgoWriter.Write([]byte("Failed to retrieve global addon logs: " + err.Error()))
		}

		ginkgo.By("Cleaning up the provisioned globals")
		err = addon.DeprovisionGlobals(cfg)
		if err != nil {
			framework.Failf("Error deprovisioning global addons: %v", err)
		}
	})
<commit_msg>Add missing call to GinkgoWriter.Write for global addon logs<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 e2e

import (
	"github.com/onsi/ginkgo"

	"github.com/jetstack/cert-manager/test/e2e/framework"
	"github.com/jetstack/cert-manager/test/e2e/framework/addon"
)

var (
	cfg = framework.DefaultConfig
)

var _ = ginkgo.SynchronizedBeforeSuite(func() []byte {
	addon.InitGlobals(cfg)

	ginkgo.By("Provisioning shared cluster addons")

	err := addon.ProvisionGlobals(cfg)
	if err != nil {
		framework.Failf("Error provisioning global addons: %v", err)
	}

	return nil
}, func([]byte) {
	addon.InitGlobals(cfg)

	ginkgo.By("Configuring details for shared cluster addons")

	err := addon.SetupGlobals(cfg)
	if err != nil {
		framework.Failf("Error configuring global addons: %v", err)
	}
})

var globalLogs string

var _ = ginkgo.SynchronizedAfterSuite(func() {},
	func() {
		ginkgo.By("Retrieving logs for global addons")
		var err error
		globalLogs, err = addon.GlobalLogs()
		if err != nil {
			ginkgo.GinkgoWriter.Write([]byte("Failed to retrieve global addon logs: " + err.Error()))
		}

		ginkgo.GinkgoWriter.Write([]byte(globalLogs))

		ginkgo.By("Cleaning up the provisioned globals")
		err = addon.DeprovisionGlobals(cfg)
		if err != nil {
			framework.Failf("Error deprovisioning global addons: %v", err)
		}
	})
<|endoftext|>