<commit_before>package formatter

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"

	"github.com/docker/distribution/reference"
	"github.com/docker/docker/api/types"
	units "github.com/docker/go-units"
)

const (
	defaultDiskUsageImageTableFormat     = "table {{.Repository}}\t{{.Tag}}\t{{.ID}}\t{{.CreatedSince}} ago\t{{.VirtualSize}}\t{{.SharedSize}}\t{{.UniqueSize}}\t{{.Containers}}"
	defaultDiskUsageContainerTableFormat = "table {{.ID}}\t{{.Image}}\t{{.Command}}\t{{.LocalVolumes}}\t{{.Size}}\t{{.RunningFor}} ago\t{{.Status}}\t{{.Names}}"
	defaultDiskUsageVolumeTableFormat    = "table {{.Name}}\t{{.Links}}\t{{.Size}}"
	defaultDiskUsageTableFormat          = "table {{.Type}}\t{{.TotalCount}}\t{{.Active}}\t{{.Size}}\t{{.Reclaimable}}"

	typeHeader        = "TYPE"
	totalHeader       = "TOTAL"
	activeHeader      = "ACTIVE"
	reclaimableHeader = "RECLAIMABLE"
	containersHeader  = "CONTAINERS"
	sharedSizeHeader  = "SHARED SIZE"
	uniqueSizeHeader  = "UNIQUE SiZE"
)

// DiskUsageContext contains disk usage specific information required by the formatter, encapsulate a Context struct.
type DiskUsageContext struct {
	Context
	Verbose    bool
	LayersSize int64
	Images     []*types.ImageSummary
	Containers []*types.Container
	Volumes    []*types.Volume
}

func (ctx *DiskUsageContext) startSubsection(format string) (*template.Template, error) {
	ctx.buffer = bytes.NewBufferString("")
	ctx.header = ""
	ctx.Format = Format(format)
	ctx.preFormat()

	return ctx.parseFormat()
}

func (ctx *DiskUsageContext) Write() {
	if ctx.Verbose == false {
		ctx.buffer = bytes.NewBufferString("")
		ctx.Format = defaultDiskUsageTableFormat
		ctx.preFormat()

		tmpl, err := ctx.parseFormat()
		if err != nil {
			return
		}

		err = ctx.contextFormat(tmpl, &diskUsageImagesContext{
			totalSize: ctx.LayersSize,
			images:    ctx.Images,
		})
		if err != nil {
			return
		}
		err = ctx.contextFormat(tmpl, &diskUsageContainersContext{
			containers: ctx.Containers,
		})
		if err != nil {
			return
		}

		err = ctx.contextFormat(tmpl, &diskUsageVolumesContext{
			volumes: ctx.Volumes,
		})
		if err != nil {
			return
		}

		ctx.postFormat(tmpl, &diskUsageContainersContext{containers: []*types.Container{}})

		return
	}

	// First images
	tmpl, err := ctx.startSubsection(defaultDiskUsageImageTableFormat)
	if err != nil {
		return
	}

	ctx.Output.Write([]byte("Images space usage:\n\n"))
	for _, i := range ctx.Images {
		repo := "<none>"
		tag := "<none>"
		if len(i.RepoTags) > 0 && !isDangling(*i) {
			// Only show the first tag
			ref, err := reference.ParseNamed(i.RepoTags[0])
			if err != nil {
				continue
			}
			if nt, ok := ref.(reference.NamedTagged); ok {
				repo = ref.Name()
				tag = nt.Tag()
			}
		}

		err = ctx.contextFormat(tmpl, &imageContext{
			repo:  repo,
			tag:   tag,
			trunc: true,
			i:     *i,
		})
		if err != nil {
			return
		}
	}
	ctx.postFormat(tmpl, &imageContext{})

	// Now containers
	ctx.Output.Write([]byte("\nContainers space usage:\n\n"))
	tmpl, err = ctx.startSubsection(defaultDiskUsageContainerTableFormat)
	if err != nil {
		return
	}
	for _, c := range ctx.Containers {
		// Don't display the virtual size
		c.SizeRootFs = 0
		err = ctx.contextFormat(tmpl, &containerContext{
			trunc: true,
			c:     *c,
		})
		if err != nil {
			return
		}
	}
	ctx.postFormat(tmpl, &containerContext{})

	// And volumes
	ctx.Output.Write([]byte("\nLocal Volumes space usage:\n\n"))
	tmpl, err = ctx.startSubsection(defaultDiskUsageVolumeTableFormat)
	if err != nil {
		return
	}
	for _, v := range ctx.Volumes {
		err = ctx.contextFormat(tmpl, &volumeContext{
			v: *v,
		})
		if err != nil {
			return
		}
	}
	ctx.postFormat(tmpl, &volumeContext{v: types.Volume{}})
}

type diskUsageImagesContext struct {
	HeaderContext
	totalSize int64
	images    []*types.ImageSummary
}

func (c *diskUsageImagesContext) MarshalJSON() ([]byte, error) {
	return marshalJSON(c)
}

func (c *diskUsageImagesContext) Type() string {
	c.AddHeader(typeHeader)
	return "Images"
}

func (c *diskUsageImagesContext) TotalCount() string {
	c.AddHeader(totalHeader)
	return fmt.Sprintf("%d", len(c.images))
}

func (c *diskUsageImagesContext) Active() string {
	c.AddHeader(activeHeader)
	used := 0
	for _, i := range c.images {
		if i.Containers > 0 {
			used++
		}
	}

	return fmt.Sprintf("%d", used)
}

func (c *diskUsageImagesContext) Size() string {
	c.AddHeader(sizeHeader)
	return units.HumanSize(float64(c.totalSize))

}

func (c *diskUsageImagesContext) Reclaimable() string {
	var used int64

	c.AddHeader(reclaimableHeader)
	for _, i := range c.images {
		if i.Containers != 0 {
			used += i.Size
		}
	}

	reclaimable := c.totalSize - used
	if c.totalSize > 0 {
		return fmt.Sprintf("%s (%v%%)", units.HumanSize(float64(reclaimable)), (reclaimable*100)/c.totalSize)
	}
	return fmt.Sprintf("%s", units.HumanSize(float64(reclaimable)))
}

type diskUsageContainersContext struct {
	HeaderContext
	verbose    bool
	containers []*types.Container
}

func (c *diskUsageContainersContext) MarshalJSON() ([]byte, error) {
	return marshalJSON(c)
}

func (c *diskUsageContainersContext) Type() string {
	c.AddHeader(typeHeader)
	return "Containers"
}

func (c *diskUsageContainersContext) TotalCount() string {
	c.AddHeader(totalHeader)
	return fmt.Sprintf("%d", len(c.containers))
}

func (c *diskUsageContainersContext) isActive(container types.Container) bool {
	return strings.Contains(container.State, "running") ||
		strings.Contains(container.State, "paused") ||
		strings.Contains(container.State, "restarting")
}

func (c *diskUsageContainersContext) Active() string {
	c.AddHeader(activeHeader)
	used := 0
	for _, container := range c.containers {
		if c.isActive(*container) {
			used++
		}
	}

	return fmt.Sprintf("%d", used)
}

func (c *diskUsageContainersContext) Size() string {
	var size int64

	c.AddHeader(sizeHeader)
	for _, container := range c.containers {
		size += container.SizeRw
	}

	return units.HumanSize(float64(size))
}

func (c *diskUsageContainersContext) Reclaimable() string {
	var reclaimable int64
	var totalSize int64

	c.AddHeader(reclaimableHeader)
	for _, container := range c.containers {
		if !c.isActive(*container) {
			reclaimable += container.SizeRw
		}
		totalSize += container.SizeRw
	}

	if totalSize > 0 {
		return fmt.Sprintf("%s (%v%%)", units.HumanSize(float64(reclaimable)), (reclaimable*100)/totalSize)
	}

	return fmt.Sprintf("%s", units.HumanSize(float64(reclaimable)))
}

type diskUsageVolumesContext struct {
	HeaderContext
	verbose bool
	volumes []*types.Volume
}

func (c *diskUsageVolumesContext) MarshalJSON() ([]byte, error) {
	return marshalJSON(c)
}

func (c *diskUsageVolumesContext) Type() string {
	c.AddHeader(typeHeader)
	return "Local Volumes"
}

func (c *diskUsageVolumesContext) TotalCount() string {
	c.AddHeader(totalHeader)
	return fmt.Sprintf("%d", len(c.volumes))
}

func (c *diskUsageVolumesContext) Active() string {
	c.AddHeader(activeHeader)

	used := 0
	for _, v := range c.volumes {
		if v.UsageData.RefCount > 0 {
			used++
		}
	}

	return fmt.Sprintf("%d", used)
}

func (c *diskUsageVolumesContext) Size() string {
	var size int64

	c.AddHeader(sizeHeader)
	for _, v := range c.volumes {
		if v.UsageData.Size != -1 {
			size += v.UsageData.Size
		}
	}

	return units.HumanSize(float64(size))
}

func (c *diskUsageVolumesContext) Reclaimable() string {
	var reclaimable int64
	var totalSize int64

	c.AddHeader(reclaimableHeader)
	for _, v := range c.volumes {
		if v.UsageData.Size != -1 {
			if v.UsageData.RefCount == 0 {
				reclaimable += v.UsageData.Size
			}
			totalSize += v.UsageData.Size
		}
	}

	if totalSize > 0 {
		return fmt.Sprintf("%s (%v%%)", units.HumanSize(float64(reclaimable)), (reclaimable*100)/totalSize)
	}

	return fmt.Sprintf("%s", units.HumanSize(float64(reclaimable)))
}
<commit_msg>Ensure proper value is used when computing reclaimable space<commit_after>package formatter

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"

	"github.com/docker/distribution/reference"
	"github.com/docker/docker/api/types"
	units "github.com/docker/go-units"
)

const (
	defaultDiskUsageImageTableFormat     = "table {{.Repository}}\t{{.Tag}}\t{{.ID}}\t{{.CreatedSince}} ago\t{{.VirtualSize}}\t{{.SharedSize}}\t{{.UniqueSize}}\t{{.Containers}}"
	defaultDiskUsageContainerTableFormat = "table {{.ID}}\t{{.Image}}\t{{.Command}}\t{{.LocalVolumes}}\t{{.Size}}\t{{.RunningFor}} ago\t{{.Status}}\t{{.Names}}"
	defaultDiskUsageVolumeTableFormat    = "table {{.Name}}\t{{.Links}}\t{{.Size}}"
	defaultDiskUsageTableFormat          = "table {{.Type}}\t{{.TotalCount}}\t{{.Active}}\t{{.Size}}\t{{.Reclaimable}}"

	typeHeader        = "TYPE"
	totalHeader       = "TOTAL"
	activeHeader      = "ACTIVE"
	reclaimableHeader = "RECLAIMABLE"
	containersHeader  = "CONTAINERS"
	sharedSizeHeader  = "SHARED SIZE"
	uniqueSizeHeader  = "UNIQUE SiZE"
)

// DiskUsageContext contains disk usage specific information required by the formatter, encapsulate a Context struct.
type DiskUsageContext struct {
	Context
	Verbose    bool
	LayersSize int64
	Images     []*types.ImageSummary
	Containers []*types.Container
	Volumes    []*types.Volume
}

func (ctx *DiskUsageContext) startSubsection(format string) (*template.Template, error) {
	ctx.buffer = bytes.NewBufferString("")
	ctx.header = ""
	ctx.Format = Format(format)
	ctx.preFormat()

	return ctx.parseFormat()
}

func (ctx *DiskUsageContext) Write() {
	if ctx.Verbose == false {
		ctx.buffer = bytes.NewBufferString("")
		ctx.Format = defaultDiskUsageTableFormat
		ctx.preFormat()

		tmpl, err := ctx.parseFormat()
		if err != nil {
			return
		}

		err = ctx.contextFormat(tmpl, &diskUsageImagesContext{
			totalSize: ctx.LayersSize,
			images:    ctx.Images,
		})
		if err != nil {
			return
		}
		err = ctx.contextFormat(tmpl, &diskUsageContainersContext{
			containers: ctx.Containers,
		})
		if err != nil {
			return
		}

		err = ctx.contextFormat(tmpl, &diskUsageVolumesContext{
			volumes: ctx.Volumes,
		})
		if err != nil {
			return
		}

		ctx.postFormat(tmpl, &diskUsageContainersContext{containers: []*types.Container{}})

		return
	}

	// First images
	tmpl, err := ctx.startSubsection(defaultDiskUsageImageTableFormat)
	if err != nil {
		return
	}

	ctx.Output.Write([]byte("Images space usage:\n\n"))
	for _, i := range ctx.Images {
		repo := "<none>"
		tag := "<none>"
		if len(i.RepoTags) > 0 && !isDangling(*i) {
			// Only show the first tag
			ref, err := reference.ParseNamed(i.RepoTags[0])
			if err != nil {
				continue
			}
			if nt, ok := ref.(reference.NamedTagged); ok {
				repo = ref.Name()
				tag = nt.Tag()
			}
		}

		err = ctx.contextFormat(tmpl, &imageContext{
			repo:  repo,
			tag:   tag,
			trunc: true,
			i:     *i,
		})
		if err != nil {
			return
		}
	}
	ctx.postFormat(tmpl, &imageContext{})

	// Now containers
	ctx.Output.Write([]byte("\nContainers space usage:\n\n"))
	tmpl, err = ctx.startSubsection(defaultDiskUsageContainerTableFormat)
	if err != nil {
		return
	}
	for _, c := range ctx.Containers {
		// Don't display the virtual size
		c.SizeRootFs = 0
		err = ctx.contextFormat(tmpl, &containerContext{
			trunc: true,
			c:     *c,
		})
		if err != nil {
			return
		}
	}
	ctx.postFormat(tmpl, &containerContext{})

	// And volumes
	ctx.Output.Write([]byte("\nLocal Volumes space usage:\n\n"))
	tmpl, err = ctx.startSubsection(defaultDiskUsageVolumeTableFormat)
	if err != nil {
		return
	}
	for _, v := range ctx.Volumes {
		err = ctx.contextFormat(tmpl, &volumeContext{
			v: *v,
		})
		if err != nil {
			return
		}
	}
	ctx.postFormat(tmpl, &volumeContext{v: types.Volume{}})
}

type diskUsageImagesContext struct {
	HeaderContext
	totalSize int64
	images    []*types.ImageSummary
}

func (c *diskUsageImagesContext) MarshalJSON() ([]byte, error) {
	return marshalJSON(c)
}

func (c *diskUsageImagesContext) Type() string {
	c.AddHeader(typeHeader)
	return "Images"
}

func (c *diskUsageImagesContext) TotalCount() string {
	c.AddHeader(totalHeader)
	return fmt.Sprintf("%d", len(c.images))
}

func (c *diskUsageImagesContext) Active() string {
	c.AddHeader(activeHeader)
	used := 0
	for _, i := range c.images {
		if i.Containers > 0 {
			used++
		}
	}

	return fmt.Sprintf("%d", used)
}

func (c *diskUsageImagesContext) Size() string {
	c.AddHeader(sizeHeader)
	return units.HumanSize(float64(c.totalSize))

}

func (c *diskUsageImagesContext) Reclaimable() string {
	var used int64

	c.AddHeader(reclaimableHeader)
	for _, i := range c.images {
		if i.Containers != 0 {
			if i.VirtualSize == -1 || i.SharedSize == -1 {
				continue
			}
			used += i.VirtualSize - i.SharedSize
		}
	}

	reclaimable := c.totalSize - used
	if c.totalSize > 0 {
		return fmt.Sprintf("%s (%v%%)", units.HumanSize(float64(reclaimable)), (reclaimable*100)/c.totalSize)
	}
	return fmt.Sprintf("%s", units.HumanSize(float64(reclaimable)))
}

type diskUsageContainersContext struct {
	HeaderContext
	verbose    bool
	containers []*types.Container
}

func (c *diskUsageContainersContext) MarshalJSON() ([]byte, error) {
	return marshalJSON(c)
}

func (c *diskUsageContainersContext) Type() string {
	c.AddHeader(typeHeader)
	return "Containers"
}

func (c *diskUsageContainersContext) TotalCount() string {
	c.AddHeader(totalHeader)
	return fmt.Sprintf("%d", len(c.containers))
}

func (c *diskUsageContainersContext) isActive(container types.Container) bool {
	return strings.Contains(container.State, "running") ||
		strings.Contains(container.State, "paused") ||
		strings.Contains(container.State, "restarting")
}

func (c *diskUsageContainersContext) Active() string {
	c.AddHeader(activeHeader)
	used := 0
	for _, container := range c.containers {
		if c.isActive(*container) {
			used++
		}
	}

	return fmt.Sprintf("%d", used)
}

func (c *diskUsageContainersContext) Size() string {
	var size int64

	c.AddHeader(sizeHeader)
	for _, container := range c.containers {
		size += container.SizeRw
	}

	return units.HumanSize(float64(size))
}

func (c *diskUsageContainersContext) Reclaimable() string {
	var reclaimable int64
	var totalSize int64

	c.AddHeader(reclaimableHeader)
	for _, container := range c.containers {
		if !c.isActive(*container) {
			reclaimable += container.SizeRw
		}
		totalSize += container.SizeRw
	}

	if totalSize > 0 {
		return fmt.Sprintf("%s (%v%%)", units.HumanSize(float64(reclaimable)), (reclaimable*100)/totalSize)
	}

	return fmt.Sprintf("%s", units.HumanSize(float64(reclaimable)))
}

type diskUsageVolumesContext struct {
	HeaderContext
	verbose bool
	volumes []*types.Volume
}

func (c *diskUsageVolumesContext) MarshalJSON() ([]byte, error) {
	return marshalJSON(c)
}

func (c *diskUsageVolumesContext) Type() string {
	c.AddHeader(typeHeader)
	return "Local Volumes"
}

func (c *diskUsageVolumesContext) TotalCount() string {
	c.AddHeader(totalHeader)
	return fmt.Sprintf("%d", len(c.volumes))
}

func (c *diskUsageVolumesContext) Active() string {
	c.AddHeader(activeHeader)

	used := 0
	for _, v := range c.volumes {
		if v.UsageData.RefCount > 0 {
			used++
		}
	}

	return fmt.Sprintf("%d", used)
}

func (c *diskUsageVolumesContext) Size() string {
	var size int64

	c.AddHeader(sizeHeader)
	for _, v := range c.volumes {
		if v.UsageData.Size != -1 {
			size += v.UsageData.Size
		}
	}

	return units.HumanSize(float64(size))
}

func (c *diskUsageVolumesContext) Reclaimable() string {
	var reclaimable int64
	var totalSize int64

	c.AddHeader(reclaimableHeader)
	for _, v := range c.volumes {
		if v.UsageData.Size != -1 {
			if v.UsageData.RefCount == 0 {
				reclaimable += v.UsageData.Size
			}
			totalSize += v.UsageData.Size
		}
	}

	if totalSize > 0 {
		return fmt.Sprintf("%s (%v%%)", units.HumanSize(float64(reclaimable)), (reclaimable*100)/totalSize)
	}

	return fmt.Sprintf("%s", units.HumanSize(float64(reclaimable)))
}
<|endoftext|>


<commit_before>// Responsible for parsing lshw output

package utils

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	ssh "github.com/dorzheh/infra/comm/common"
	"github.com/dorzheh/infra/utils/lshw"
	"github.com/dorzheh/mxj"
)

// CpuInfo contains CPU description and properties
type CpuInfo struct {
	Cpus int
	Desc map[string]interface{}
	Cap  map[string]interface{}
}

// supported NIC types
type NicType string

const (
	NicTypePhys   NicType = "physical"
	NicTypeOVS    NicType = "openvswitch"
	NicTypeBridge NicType = "bridge"
)

// NIC information
type NicInfo struct {
	// port name (eth0,br0...)
	Name string

	// NIC driver(bridge,openvswitch...)
	Driver string

	// Description
	Desc string

	// PCI Address
	PCIAddr string

	// Port type
	Type NicType
}

type HwInfoParser struct {
	run       func(string) (string, error)
	cacheFile string
	cmd       string
}

// NewHwInfoParser constructs new lshw parser
// The output will be represented in JSON format
func NewHwInfoParser(cacheFile, lshwpath string, sshconf *ssh.Config) (*HwInfoParser, error) {
	i := new(HwInfoParser)
	i.run = RunFunc(sshconf)
	if lshwpath == "" {
		out, err := i.run("which lshw")
		if err != nil {
			return nil, fmt.Errorf("%s [%v]", out, err)
		}
		lshwpath = out
	} else {
		if sshconf != nil {
			dir, err := UploadBinaries(sshconf, lshwpath)
			if err != nil {
				return nil, err
			}
			lshwpath = filepath.Join(dir, filepath.Base(lshwpath))

		}
	}

	lshwconf := &lshw.Config{[]lshw.Class{lshw.All}, lshw.FormatJSON}
	l, err := lshw.New(lshwpath, lshwconf)
	if err != nil {
		return nil, err
	}
	i.cmd = l.Cmd()
	i.cacheFile = cacheFile
	return i, nil
}

// Parse parses lshw output
func (i *HwInfoParser) Parse() error {
	out, err := i.run(i.cmd)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(i.cacheFile, []byte(out), 0)
}

// CpuInfo gathers information related to installed CPUs
func (i *HwInfoParser) CpuInfo() (*CpuInfo, error) {
	if _, err := os.Stat(i.cacheFile); err != nil {
		if err = i.Parse(); err != nil {
			return nil, err
		}
	}

	out, err := mxj.ReadMapsFromJsonFile(i.cacheFile)
	if err != nil {
		return nil, err
	}

	c := new(CpuInfo)
	cpustr, err := i.run(`grep -c processor /proc/cpuinfo`)
	if err != nil {
		return nil, err
	}

	c.Cpus, err = strconv.Atoi(strings.Trim(cpustr, "\n"))
	if err != nil {
		return nil, err
	}

	c.Desc = make(map[string]interface{})
	c.Cap = make(map[string]interface{})
	for _, s := range out {
		r, _ := s.ValuesForPath("children.children")
		for _, n := range r {
			ch := n.(map[string]interface{})
			if ch["id"] == "cpu:0" {
				for k, v := range ch {
					if k != "capabilities" {
						c.Desc[k] = v
					}
				}
				for k, v := range ch["capabilities"].(map[string]interface{}) {
					c.Cap[k] = v
				}
			}
		}
	}
	return c, nil
}

// NicInfo gathers information related to installed NICs
func (i *HwInfoParser) NicsInfo(supNicVendors []string) ([]*NicInfo, error) {
	if _, err := os.Stat(i.cacheFile); err != nil {
		if err = i.Parse(); err != nil {
			return nil, err
		}
	}
	out, err := mxj.ReadMapsFromJsonFile(i.cacheFile)
	if err != nil {
		return nil, err
	}

	nics := make([]*NicInfo, 0)
	deep := []string{"children.children.children.children", "children"}
	for _, m := range out {
		for _, d := range deep {
			r, _ := m.ValuesForPath(d)
			for _, n := range r {
				ch := n.(map[string]interface{})
				if ch["description"] == "Ethernet interface" {
					name := ch["logicalname"].(string)
					if name == "ovs-system" {
						continue
					}
					nic := new(NicInfo)
					nic.Name = name
					driver := ch["configuration"].(map[string]interface{})["driver"].(string)
					switch driver {
					case "tun":
						continue
					case "openvswitch":
						nic.Desc = "Open vSwitch interface"
						nic.Type = NicTypeOVS
					default:
						prod, ok := ch["product"].(string)
						if ok {
							vendor, _ := ch["vendor"].(string)
							if len(supNicVendors) > 0 {
								found := false
								for _, v := range supNicVendors {
									if v == vendor {
										found = true
										break
									}
								}
								if !found {
									continue
								}

							}
							if _, err := i.run(fmt.Sprintf("[ -d /sys/class/net/%s/master ]", ch["logicalname"].(string))); err == nil {
								continue
							}
							nic.PCIAddr = ch["businfo"].(string)
							nic.Desc = vendor + " " + prod
							nic.Type = NicTypePhys
						}
					}
					nic.Driver = driver
					nics = append(nics, nic)
				}
			}
		}
	}

	// lshw is unable to find linux bridges so let's do it manually
	res, err := i.run(`out="";for n in /sys/class/net/*;do [ -d $n/bridge ] && out="$out ${n##/sys/class/net/}";done;echo $out`)
	if err != nil {
		return nil, err
	}
	if res != "" {
		for _, n := range strings.Split(res, " ") {
			br := &NicInfo{
				Name:   n,
				Driver: "bridge",
				Desc:   "Bridge interface",
				Type:   NicTypeBridge,
			}
			nics = append(nics, br)
		}
	}
	return nics, nil
}

// RAMSize gathers information related to the installed amount of RAM
func (i *HwInfoParser) RAMSize() (uint, error) {
	out, err := i.run("grep MemTotal /proc/meminfo")
	if err != nil {
		return 0, err
	}
	var col3 string
	var ramsize uint
	fmt.Sscanf(out, "MemTotal: %d %s", &ramsize, &col3)
	return ramsize, nil
}
<commit_msg>Update hwinfo.go<commit_after>// Responsible for parsing lshw output

package utils

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	ssh "github.com/dorzheh/infra/comm/common"
	"github.com/dorzheh/infra/utils/lshw"
	"github.com/dorzheh/mxj"
)

// CpuInfo contains CPU description and properties
type CpuInfo struct {
	Cpus int
	Desc map[string]interface{}
	Cap  map[string]interface{}
}

// supported NIC types
type NicType string

const (
	NicTypePhys   NicType = "physical"
	NicTypeOVS    NicType = "openvswitch"
	NicTypeBridge NicType = "bridge"
)

// NIC information
type NicInfo struct {
	// port name (eth0,br0...)
	Name string

	// NIC driver(bridge,openvswitch...)
	Driver string

	// Description
	Desc string

	// PCI Address
	PCIAddr string

	// Port type
	Type NicType
}

type HwInfoParser struct {
	run       func(string) (string, error)
	cacheFile string
	cmd       string
}

// NewHwInfoParser constructs new lshw parser
// The output will be represented in JSON format
func NewHwInfoParser(cacheFile, lshwpath string, sshconf *ssh.Config) (*HwInfoParser, error) {
	i := new(HwInfoParser)
	i.run = RunFunc(sshconf)
	if lshwpath == "" {
		out, err := i.run("which lshw")
		if err != nil {
			return nil, fmt.Errorf("%s [%v]", out, err)
		}
		lshwpath = out
	} else {
		if sshconf != nil {
			dir, err := UploadBinaries(sshconf, lshwpath)
			if err != nil {
				return nil, err
			}
			lshwpath = filepath.Join(dir, filepath.Base(lshwpath))

		}
	}

	lshwconf := &lshw.Config{[]lshw.Class{lshw.All}, lshw.FormatJSON}
	l, err := lshw.New(lshwpath, lshwconf)
	if err != nil {
		return nil, err
	}
	i.cmd = l.Cmd()
	i.cacheFile = cacheFile
	return i, nil
}

// Parse parses lshw output
func (i *HwInfoParser) Parse() error {
	out, err := i.run(i.cmd)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(i.cacheFile, []byte(out), 0)
}

// CpuInfo gathers information related to installed CPUs
func (i *HwInfoParser) CpuInfo() (*CpuInfo, error) {
	if _, err := os.Stat(i.cacheFile); err != nil {
		if err = i.Parse(); err != nil {
			return nil, err
		}
	}

	out, err := mxj.ReadMapsFromJsonFile(i.cacheFile)
	if err != nil {
		return nil, err
	}

	c := new(CpuInfo)
	cpustr, err := i.run(`grep -c ^processor /proc/cpuinfo`)
	if err != nil {
		return nil, err
	}

	c.Cpus, err = strconv.Atoi(strings.Trim(cpustr, "\n"))
	if err != nil {
		return nil, err
	}

	c.Desc = make(map[string]interface{})
	c.Cap = make(map[string]interface{})
	for _, s := range out {
		r, _ := s.ValuesForPath("children.children")
		for _, n := range r {
			ch := n.(map[string]interface{})
			if ch["id"] == "cpu:0" {
				for k, v := range ch {
					if k != "capabilities" {
						c.Desc[k] = v
					}
				}
				for k, v := range ch["capabilities"].(map[string]interface{}) {
					c.Cap[k] = v
				}
			}
		}
	}
	return c, nil
}

// NicInfo gathers information related to installed NICs
func (i *HwInfoParser) NicsInfo(supNicVendors []string) ([]*NicInfo, error) {
	if _, err := os.Stat(i.cacheFile); err != nil {
		if err = i.Parse(); err != nil {
			return nil, err
		}
	}
	out, err := mxj.ReadMapsFromJsonFile(i.cacheFile)
	if err != nil {
		return nil, err
	}

	nics := make([]*NicInfo, 0)
	deep := []string{"children.children.children.children", "children"}
	for _, m := range out {
		for _, d := range deep {
			r, _ := m.ValuesForPath(d)
			for _, n := range r {
				ch := n.(map[string]interface{})
				if ch["description"] == "Ethernet interface" {
					name := ch["logicalname"].(string)
					if name == "ovs-system" {
						continue
					}
					nic := new(NicInfo)
					nic.Name = name
					driver := ch["configuration"].(map[string]interface{})["driver"].(string)
					switch driver {
					case "tun":
						continue
					case "openvswitch":
						nic.Desc = "Open vSwitch interface"
						nic.Type = NicTypeOVS
					default:
						prod, ok := ch["product"].(string)
						if ok {
							vendor, _ := ch["vendor"].(string)
							if len(supNicVendors) > 0 {
								found := false
								for _, v := range supNicVendors {
									if v == vendor {
										found = true
										break
									}
								}
								if !found {
									continue
								}

							}
							if _, err := i.run(fmt.Sprintf("[ -d /sys/class/net/%s/master ]", ch["logicalname"].(string))); err == nil {
								continue
							}
							nic.PCIAddr = ch["businfo"].(string)
							nic.Desc = vendor + " " + prod
							nic.Type = NicTypePhys
						}
					}
					nic.Driver = driver
					nics = append(nics, nic)
				}
			}
		}
	}

	// lshw is unable to find linux bridges so let's do it manually
	res, err := i.run(`out="";for n in /sys/class/net/*;do [ -d $n/bridge ] && out="$out ${n##/sys/class/net/}";done;echo $out`)
	if err != nil {
		return nil, err
	}
	if res != "" {
		for _, n := range strings.Split(res, " ") {
			br := &NicInfo{
				Name:   n,
				Driver: "bridge",
				Desc:   "Bridge interface",
				Type:   NicTypeBridge,
			}
			nics = append(nics, br)
		}
	}
	return nics, nil
}

// RAMSize gathers information related to the installed amount of RAM
func (i *HwInfoParser) RAMSize() (uint, error) {
	out, err := i.run("grep MemTotal /proc/meminfo")
	if err != nil {
		return 0, err
	}
	var col3 string
	var ramsize uint
	fmt.Sscanf(out, "MemTotal: %d %s", &ramsize, &col3)
	return ramsize, nil
}
<|endoftext|>


<commit_before>// Copyright 2014 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package gossip

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/gossip/resolver"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/rpc"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/stop"
)

// TestGossipInfoStore verifies operation of gossip instance infostore.
func TestGossipInfoStore(t *testing.T) {
	defer leaktest.AfterTest(t)()
	stopper := stop.NewStopper()
	defer stopper.Stop()
	rpcContext := rpc.NewContext(nil, nil, stopper)
	g := New(rpcContext, nil, stopper)
	// Have to call g.SetNodeID before call g.AddInfo
	g.SetNodeID(roachpb.NodeID(1))
	slice := []byte("b")
	if err := g.AddInfo("s", slice, time.Hour); err != nil {
		t.Fatal(err)
	}
	if val, err := g.GetInfo("s"); !bytes.Equal(val, slice) || err != nil {
		t.Errorf("error fetching string: %v", err)
	}
	if _, err := g.GetInfo("s2"); err == nil {
		t.Errorf("expected error fetching nonexistent key \"s2\"")
	}
}

func TestGossipGetNextBootstrapAddress(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer resolver.SetLookupTimeout(time.Minute)()

	// Set up an http server for testing the http load balancer.
	i := 0
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		i++
		fmt.Fprintf(w, `{"address": {"network": "tcp", "address": "10.10.0.%d:12345"}}`, i)
	})
	s := httptest.NewServer(handler)
	defer s.Close()

	resolverSpecs := []string{
		"127.0.0.1:9000",
		"tcp=127.0.0.1:9001",
		"unix=/tmp/unix-socket12345",
		fmt.Sprintf("http-lb=%s", s.Listener.Addr()),
		"foo=127.0.0.1:9003", // error should not resolve.
		"http-lb=",           // error should not resolve.
		"localhost:9004",
	}

	resolvers := []resolver.Resolver{}
	for _, rs := range resolverSpecs {
		resolver, err := resolver.NewResolver(&base.Context{Insecure: true}, rs)
		if err == nil {
			resolvers = append(resolvers, resolver)
		}
	}
	if len(resolvers) != 5 {
		t.Errorf("expected 5 resolvers; got %d", len(resolvers))
	}
	g := New(nil, resolvers, nil)

	// Using specified resolvers, fetch bootstrap addresses 10 times
	// and verify the results match expected addresses.
	expAddresses := []string{
		"127.0.0.1:9000",
		"127.0.0.1:9001",
		"/tmp/unix-socket12345",
		"10.10.0.1:12345",
		"localhost:9004",
		"10.10.0.2:12345",
		"10.10.0.3:12345",
		"10.10.0.4:12345",
		"10.10.0.5:12345",
		"10.10.0.6:12345",
	}
	for i := 0; i < len(expAddresses); i++ {
		if addr := g.getNextBootstrapAddress(); addr == nil {
			t.Errorf("%d: unexpected nil addr when expecting %s", i, expAddresses[i])
		} else if addrStr := addr.String(); addrStr != expAddresses[i] {
			t.Errorf("%d: expected addr %s; got %s", i, expAddresses[i], addrStr)
		}
	}
}

// TestGossipCullNetwork verifies that a client will be culled from
// the network periodically (at cullInterval duration intervals).
func TestGossipCullNetwork(t *testing.T) {
	defer leaktest.AfterTest(t)()

	// Create the local gossip and minPeers peers.
	stopper := stop.NewStopper()
	defer stopper.Stop()
	local := startGossip(1, stopper, t)
	local.SetCullInterval(5 * time.Millisecond)
	peers := []*Gossip{}
	for i := 0; i < minPeers; i++ {
		peers = append(peers, startGossip(roachpb.NodeID(i+2), stopper, t))
	}

	// Start clients to all peers and start the local gossip's manage routine.
	local.mu.Lock()
	for _, p := range peers {
		pAddr := p.is.NodeAddr
		local.startClient(&pAddr, stopper)
	}
	local.mu.Unlock()
	local.manage()

	util.SucceedsSoon(t, func() error {
		// Verify that a client is closed within the cull interval.
		if len(local.Outgoing()) == minPeers-1 {
			return nil
		}
		return errors.New("no network culling occurred")
	})
}
<commit_msg>gossip: remove useless slice<commit_after>// Copyright 2014 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package gossip

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/gossip/resolver"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/rpc"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/stop"
)

// TestGossipInfoStore verifies operation of gossip instance infostore.
func TestGossipInfoStore(t *testing.T) {
	defer leaktest.AfterTest(t)()
	stopper := stop.NewStopper()
	defer stopper.Stop()
	rpcContext := rpc.NewContext(nil, nil, stopper)
	g := New(rpcContext, nil, stopper)
	// Have to call g.SetNodeID before call g.AddInfo
	g.SetNodeID(roachpb.NodeID(1))
	slice := []byte("b")
	if err := g.AddInfo("s", slice, time.Hour); err != nil {
		t.Fatal(err)
	}
	if val, err := g.GetInfo("s"); !bytes.Equal(val, slice) || err != nil {
		t.Errorf("error fetching string: %v", err)
	}
	if _, err := g.GetInfo("s2"); err == nil {
		t.Errorf("expected error fetching nonexistent key \"s2\"")
	}
}

func TestGossipGetNextBootstrapAddress(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer resolver.SetLookupTimeout(time.Minute)()

	// Set up an http server for testing the http load balancer.
	i := 0
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		i++
		fmt.Fprintf(w, `{"address": {"network": "tcp", "address": "10.10.0.%d:12345"}}`, i)
	})
	s := httptest.NewServer(handler)
	defer s.Close()

	resolverSpecs := []string{
		"127.0.0.1:9000",
		"tcp=127.0.0.1:9001",
		"unix=/tmp/unix-socket12345",
		fmt.Sprintf("http-lb=%s", s.Listener.Addr()),
		"foo=127.0.0.1:9003", // error should not resolve.
		"http-lb=",           // error should not resolve.
		"localhost:9004",
	}

	resolvers := []resolver.Resolver{}
	for _, rs := range resolverSpecs {
		resolver, err := resolver.NewResolver(&base.Context{Insecure: true}, rs)
		if err == nil {
			resolvers = append(resolvers, resolver)
		}
	}
	if len(resolvers) != 5 {
		t.Errorf("expected 5 resolvers; got %d", len(resolvers))
	}
	g := New(nil, resolvers, nil)

	// Using specified resolvers, fetch bootstrap addresses 10 times
	// and verify the results match expected addresses.
	expAddresses := []string{
		"127.0.0.1:9000",
		"127.0.0.1:9001",
		"/tmp/unix-socket12345",
		"10.10.0.1:12345",
		"localhost:9004",
		"10.10.0.2:12345",
		"10.10.0.3:12345",
		"10.10.0.4:12345",
		"10.10.0.5:12345",
		"10.10.0.6:12345",
	}
	for i := 0; i < len(expAddresses); i++ {
		if addr := g.getNextBootstrapAddress(); addr == nil {
			t.Errorf("%d: unexpected nil addr when expecting %s", i, expAddresses[i])
		} else if addrStr := addr.String(); addrStr != expAddresses[i] {
			t.Errorf("%d: expected addr %s; got %s", i, expAddresses[i], addrStr)
		}
	}
}

// TestGossipCullNetwork verifies that a client will be culled from
// the network periodically (at cullInterval duration intervals).
func TestGossipCullNetwork(t *testing.T) {
	defer leaktest.AfterTest(t)()

	stopper := stop.NewStopper()
	defer stopper.Stop()
	local := startGossip(1, stopper, t)
	local.SetCullInterval(5 * time.Millisecond)

	local.mu.Lock()
	for i := 0; i < minPeers; i++ {
		peer := startGossip(roachpb.NodeID(i+2), stopper, t)
		local.startClient(&peer.is.NodeAddr, stopper)
	}
	local.mu.Unlock()
	local.manage()

	util.SucceedsSoon(t, func() error {
		// Verify that a client is closed within the cull interval.
		if len(local.Outgoing()) == minPeers-1 {
			return nil
		}
		return errors.New("no network culling occurred")
	})
}
<|endoftext|>


<commit_before>package crypto_test

import (
	"crypto/cipher"
	"testing"

	. "github.com/v2ray/v2ray-core/common/crypto"
)

const benchSize = 1024 * 1024

func benchmarkStream(b *testing.B, c cipher.Stream) {
	b.SetBytes(benchSize)
	input := make([]byte, benchSize)
	output := make([]byte, benchSize)
	for i := 0; i < b.N; i++ {
		c.XORKeyStream(output, input)
	}
}

func BenchmarkChaCha20(b *testing.B) {
	key := make([]byte, 32)
	nonce := make([]byte, 8)
	c := NewChaCha20Stream(key, nonce)
	benchmarkStream(b, c)
}

func BenchmarkChaCha20IETF(b *testing.B) {
	key := make([]byte, 32)
	nonce := make([]byte, 12)
	c := NewChaCha20Stream(key, nonce)
	benchmarkStream(b, c)
}

func BenchmarkAESEncryption(b *testing.B) {
	key := make([]byte, 32)
	iv := make([]byte, 16)
	c := NewAesEncryptionStream(key, iv)

	benchmarkStream(b, c)
}

func BenchmarkAESDecryption(b *testing.B) {
	key := make([]byte, 32)
	iv := make([]byte, 16)
	c := NewAesDecryptionStream(key, iv)

	benchmarkStream(b, c)
}
<commit_msg>reset timer before benchmark<commit_after>package crypto_test

import (
	"crypto/cipher"
	"testing"

	. "github.com/v2ray/v2ray-core/common/crypto"
)

const benchSize = 1024 * 1024

func benchmarkStream(b *testing.B, c cipher.Stream) {
	b.SetBytes(benchSize)
	input := make([]byte, benchSize)
	output := make([]byte, benchSize)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		c.XORKeyStream(output, input)
	}
}

func BenchmarkChaCha20(b *testing.B) {
	key := make([]byte, 32)
	nonce := make([]byte, 8)
	c := NewChaCha20Stream(key, nonce)
	benchmarkStream(b, c)
}

func BenchmarkChaCha20IETF(b *testing.B) {
	key := make([]byte, 32)
	nonce := make([]byte, 12)
	c := NewChaCha20Stream(key, nonce)
	benchmarkStream(b, c)
}

func BenchmarkAESEncryption(b *testing.B) {
	key := make([]byte, 32)
	iv := make([]byte, 16)
	c := NewAesEncryptionStream(key, iv)

	benchmarkStream(b, c)
}

func BenchmarkAESDecryption(b *testing.B) {
	key := make([]byte, 32)
	iv := make([]byte, 16)
	c := NewAesDecryptionStream(key, iv)

	benchmarkStream(b, c)
}
<|endoftext|>


<commit_before>package lib

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/cheggaaa/pb"
	"github.com/nightlyone/lockfile"
	"github.com/pivotal-golang/archiver/extractor"

	"github.com/subutai-io/base/agent/config"
	"github.com/subutai-io/base/agent/lib/container"
	"github.com/subutai-io/base/agent/lib/gpg"
	"github.com/subutai-io/base/agent/lib/template"
	"github.com/subutai-io/base/agent/log"
)

var (
	lock lockfile.Lockfile
)

type templ struct {
	name      string
	file      string
	version   string
	branch    string
	id        string
	owner     []string
	signature map[string]string
}

type metainfo struct {
	ID    string            `json:"id"`
	Name  string            `json:"name"`
	Owner []string          `json:"owner"`
	File  string            `json:"filename"`
	Signs map[string]string `json:"signature"`
}

func templId(t *templ, kurjun *http.Client, token string) {
	var meta metainfo

	url := config.Cdn.Kurjun + "/template/info?name=" + t.name + "&token=" + token
	if t.name == "management" && len(t.branch) != 0 {
		url = config.Cdn.Kurjun + "/template/info?name=" + t.name + "&version=" + t.version + "-" + t.branch + "&token=" + token
	} else if t.name == "management" {
		url = config.Cdn.Kurjun + "/template/info?name=" + t.name + "&version=" + t.version + "&token=" + token
	}

	response, err := kurjun.Get(url)
	defer response.Body.Close()
	log.Debug("Retrieving id, get: " + url)

	if err == nil && response.StatusCode == 204 && t.name == "management" {
		log.Warn("Cannot get management with specified version, trying without version")
		response, err = kurjun.Get(config.Cdn.Kurjun + "/template/info?name=" + t.name + "&token=" + token)
	}
	if log.Check(log.WarnLevel, "Getting kurjun response", err) || response.StatusCode != 200 {
		return
	}

	defer response.Body.Close()
	body, err := ioutil.ReadAll(response.Body)

	if log.Check(log.WarnLevel, "Parsing response body", json.Unmarshal(body, &meta)) {
		return
	}

	if t.name != meta.Name {
		log.Info("Found: " + t.name + " -> " + meta.Name)
		t.name = meta.Name
	}
	if len(t.owner) == 0 {
		t.owner = meta.Owner
	}
	t.id = meta.ID
	t.file = meta.File
	t.signature = meta.Signs
}

func md5sum(filePath string) string {
	file, err := os.Open(filePath)
	if err != nil {
		return ""
	}
	defer file.Close()

	hash := md5.New()
	if _, err := io.Copy(hash, file); err != nil {
		return ""
	}
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func checkLocal(t templ) bool {
	var response string
	files, _ := ioutil.ReadDir(config.Agent.LxcPrefix + "tmpdir")
	for _, f := range files {
		if t.file == f.Name() {
			if len(t.id) == 0 {
				fmt.Print("Cannot verify local template. Trust anyway? (y/n)")
				_, err := fmt.Scanln(&response)
				log.Check(log.FatalLevel, "Reading input", err)
				if response == "y" {
					return true
				}
				return false
			}
			if id := strings.Split(t.id, "."); len(id) > 0 && id[len(id)-1] == md5sum(config.Agent.LxcPrefix+"tmpdir/"+f.Name()) {
				return true
			}
		}
	}
	return false
}

func download(t templ, kurjun *http.Client, token string) bool {
	if len(t.id) == 0 {
		return false
	}
	out, err := os.Create(config.Agent.LxcPrefix + "tmpdir/" + t.file)
	log.Check(log.FatalLevel, "Creating file "+t.file, err)
	defer out.Close()
	log.Info("Downloading " + t.name)

	url := config.Cdn.Kurjun + "/template/download?id=" + t.id
	if len(t.owner) > 0 {
		url = config.Cdn.Kurjun + "/template/" + t.owner[0] + "/" + t.file
	}
	response, err := kurjun.Get(url)
	log.Check(log.FatalLevel, "Getting "+url, err)

	defer response.Body.Close()
	bar := pb.New(int(response.ContentLength)).SetUnits(pb.U_BYTES)
	bar.Start()
	rd := bar.NewProxyReader(response.Body)

	_, err = io.Copy(out, rd)
	for c := 0; err != nil && c < 5; _, err = io.Copy(out, rd) {
		log.Info("Download interrupted, retrying")
		time.Sleep(3 * time.Second)
		c++

		//Repeating GET request to CDN, while need to continue interrupted download
		out, err = os.Create(config.Agent.LxcPrefix + "tmpdir/" + t.file)
		log.Check(log.FatalLevel, "Creating file "+t.file, err)
		defer out.Close()
		response, err = kurjun.Get(config.Cdn.Kurjun + "/template/download?id=" + t.id + "&token=" + token)
		log.Check(log.FatalLevel, "Getting "+config.Cdn.Kurjun+"/template/download?id="+t.id+"&token="+token, err)
		defer response.Body.Close()
		bar = pb.New(int(response.ContentLength)).SetUnits(pb.U_BYTES)
		bar.Start()
		rd = bar.NewProxyReader(response.Body)
	}

	log.Check(log.FatalLevel, "Writing response body to file", err)

	if id := strings.Split(t.id, "."); len(id) > 0 && id[len(id)-1] == md5sum(config.Agent.LxcPrefix+"tmpdir/"+t.file) {
		return true
	}
	log.Error("Failed to check MD5 after download. Please check your connection and try again.")
	return false
}

func idToName(id string, kurjun *http.Client, token string) string {
	var meta metainfo

	//Since only kurjun knows template's ID, we cannot define if we have template already installed in system by ID as we do it by name, so unreachable kurjun in this case is a deadend for us
	//To omit this issue we should add ID into template config and use this ID as a "primary key" to any request
	response, err := kurjun.Get(config.Cdn.Kurjun + "/template/info?id=" + id + "&token=" + token)
	log.Check(log.ErrorLevel, "Getting kurjun response", err)
	defer response.Body.Close()

	body, err := ioutil.ReadAll(response.Body)

	if string(body) == "Not found" {
		log.Error("Template with id \"" + id + "\" not found")
	}
	log.Check(log.ErrorLevel, "Parsing response body", json.Unmarshal(body, &meta))

	return meta.Name
}

func lockSubutai(file string) bool {
	lock, err := lockfile.New("/var/run/lock/subutai." + file)
	if log.Check(log.DebugLevel, "Init lock "+file, err) {
		return false
	}

	err = lock.TryLock()
	if log.Check(log.DebugLevel, "Locking file "+file, err) {
		return false
	}

	return true
}

func unlockSubutai() {
	lock.Unlock()
}

func LxcImport(name, version, token string) {
	var kurjun *http.Client

	if container.IsContainer(name) && name == "management" && len(token) > 1 {
		gpg.ExchageAndEncrypt("management", token)
		return
	}

	if id := strings.Split(name, "id:"); len(id) > 1 {
		kurjun, _ = config.CheckKurjun()
		name = idToName(id[1], kurjun, token)
	}

	var t templ

	t.name = name
	if line := strings.Split(t.name, "/"); len(line) > 1 {
		t.name = line[1]
		t.owner = append(t.owner, line[0])
	}

	log.Info("Importing " + name)
	for !lockSubutai(t.name + ".import") {
		time.Sleep(time.Second * 1)
	}
	defer unlockSubutai()

	if container.IsContainer(t.name) {
		log.Info(t.name + " instance exist")
		return
	}

	t.version = config.Template.Version
	t.branch = config.Template.Branch
	if len(version) != 0 {
		if strings.Contains(version, "-") {
			verstr := strings.Split(version, "-")
			if len(verstr) == 2 {
				t.version = verstr[0]
				t.branch = verstr[1]
			} else {
				log.Error("Invalid version")
			}
		} else {
			if strings.Contains(version, ".") {
				t.version = version
				t.branch = config.Template.Branch
			} else {
				t.version = config.Template.Version
				t.branch = version
			}
		}
	}

	log.Info("Version: " + t.version + ", branch: " + t.branch)

	if t.branch == "" || t.name != "management" {
		t.file = t.name + "-subutai-template_" + t.version + "_" + config.Template.Arch + ".tar.gz"
	} else {
		t.file = t.name + "-subutai-template_" + t.version + "-" + t.branch + "_" + config.Template.Arch + ".tar.gz"
	}

	if kurjun == nil {
		kurjun, _ = config.CheckKurjun()
	}
	if kurjun != nil {
		templId(&t, kurjun, token)
	} else {
		log.Info("Trying to import from local storage")
	}

	if len(t.id) != 0 && len(t.signature) == 0 {
		log.Warn("Template is not signed")
	}

	for owner, signature := range t.signature {
		// if v.Author == "public" || v.Author == "subutai" || v.Author == "jenkins" {
		signedhash := gpg.VerifySignature(gpg.KurjunUserPK(owner), signature)
		if t.id != signedhash {
			log.Error("Signature does not match with template hash")
		}
		log.Info("Digital signature and file integrity verified")
		log.Debug("Signature belongs to " + owner)
		break
		// }
	}

	if !checkLocal(t) && !download(t, kurjun, token) {
		log.Error(t.name + " template not found")
	}

	log.Info("Unpacking template " + t.name)
	tgz := extractor.NewTgz()
	templdir := config.Agent.LxcPrefix + "tmpdir/" + t.name
	tgz.Extract(config.Agent.LxcPrefix+"tmpdir/"+t.file, templdir)
	parent := container.GetConfigItem(templdir+"/config", "subutai.parent")
	if parent != "" && parent != t.name && !container.IsTemplate(parent) {
		log.Info("Parent template required: " + parent)
		LxcImport(parent, "", token)
	}

	log.Info("Installing template " + t.name)
	template.Install(parent, t.name)
	// TODO following lines kept for back compatibility with old templates, should be deleted when all templates will be replaced.
	os.Rename(config.Agent.LxcPrefix+t.name+"/"+t.name+"-home", config.Agent.LxcPrefix+t.name+"/home")
	os.Rename(config.Agent.LxcPrefix+t.name+"/"+t.name+"-var", config.Agent.LxcPrefix+t.name+"/var")
	os.Rename(config.Agent.LxcPrefix+t.name+"/"+t.name+"-opt", config.Agent.LxcPrefix+t.name+"/opt")
	log.Check(log.FatalLevel, "Removing temp dir "+templdir, os.RemoveAll(templdir))

	if t.name == "management" {
		template.MngInit()
		return
	}

	container.SetContainerConf(t.name, [][]string{
		{"lxc.rootfs", config.Agent.LxcPrefix + t.name + "/rootfs"},
		{"lxc.mount", config.Agent.LxcPrefix + t.name + "/fstab"},
		{"lxc.hook.pre-start", ""},
		{"lxc.include", config.Agent.AppPrefix + "share/lxc/config/ubuntu.common.conf"},
		{"lxc.include", config.Agent.AppPrefix + "share/lxc/config/ubuntu.userns.conf"},
		{"subutai.config.path", config.Agent.AppPrefix + "etc"},
		{"lxc.network.script.up", config.Agent.AppPrefix + "bin/create_ovs_interface"},
		{"lxc.mount.entry", config.Agent.LxcPrefix + t.name + "/home home none bind,rw 0 0"},
		{"lxc.mount.entry", config.Agent.LxcPrefix + t.name + "/opt opt none bind,rw 0 0"},
		{"lxc.mount.entry", config.Agent.LxcPrefix + t.name + "/var var none bind,rw 0 0"},
	})
}
<commit_msg>Added priority for template owners.<commit_after>package lib

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/cheggaaa/pb"
	"github.com/nightlyone/lockfile"
	"github.com/pivotal-golang/archiver/extractor"

	"github.com/subutai-io/base/agent/config"
	"github.com/subutai-io/base/agent/lib/container"
	"github.com/subutai-io/base/agent/lib/gpg"
	"github.com/subutai-io/base/agent/lib/template"
	"github.com/subutai-io/base/agent/log"
)

var (
	lock   lockfile.Lockfile
	owners = []string{"subutai", "public", "jenkins", "docker", ""}
)

type templ struct {
	name      string
	file      string
	version   string
	branch    string
	id        string
	owner     []string
	signature map[string]string
}

type metainfo struct {
	ID    string            `json:"id"`
	Name  string            `json:"name"`
	Owner []string          `json:"owner"`
	File  string            `json:"filename"`
	Signs map[string]string `json:"signature"`
}

func templId(t *templ, kurjun *http.Client, token string) {
	var meta metainfo

	url := config.Cdn.Kurjun + "/template/info?name=" + t.name + "&token=" + token
	if t.name == "management" && len(t.branch) != 0 {
		url = config.Cdn.Kurjun + "/template/info?name=" + t.name + "&version=" + t.version + "-" + t.branch + "&token=" + token
	} else if t.name == "management" {
		url = config.Cdn.Kurjun + "/template/info?name=" + t.name + "&version=" + t.version + "&token=" + token
	}

	response, err := kurjun.Get(url)
	defer response.Body.Close()
	log.Debug("Retrieving id, get: " + url)

	if err == nil && response.StatusCode == 204 && t.name == "management" {
		log.Warn("Cannot get management with specified version, trying without version")
		response, err = kurjun.Get(config.Cdn.Kurjun + "/template/info?name=" + t.name + "&token=" + token)
	}
	if log.Check(log.WarnLevel, "Getting kurjun response", err) || response.StatusCode != 200 {
		return
	}

	defer response.Body.Close()
	body, err := ioutil.ReadAll(response.Body)

	if log.Check(log.WarnLevel, "Parsing response body", json.Unmarshal(body, &meta)) {
		return
	}

	if t.name != meta.Name {
		log.Info("Found: " + t.name + " -> " + meta.Name)
		t.name = meta.Name
	}
	t.id = meta.ID
	t.file = meta.File
	t.signature = meta.Signs
}

func md5sum(filePath string) string {
	file, err := os.Open(filePath)
	if err != nil {
		return ""
	}
	defer file.Close()

	hash := md5.New()
	if _, err := io.Copy(hash, file); err != nil {
		return ""
	}
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func checkLocal(t templ) bool {
	var response string
	files, _ := ioutil.ReadDir(config.Agent.LxcPrefix + "tmpdir")
	for _, f := range files {
		if t.file == f.Name() {
			if len(t.id) == 0 {
				fmt.Print("Cannot verify local template. Trust anyway? (y/n)")
				_, err := fmt.Scanln(&response)
				log.Check(log.FatalLevel, "Reading input", err)
				if response == "y" {
					return true
				}
				return false
			}
			if id := strings.Split(t.id, "."); len(id) > 0 && id[len(id)-1] == md5sum(config.Agent.LxcPrefix+"tmpdir/"+f.Name()) {
				return true
			}
		}
	}
	return false
}

func download(t templ, kurjun *http.Client, token string) bool {
	if len(t.id) == 0 {
		return false
	}
	out, err := os.Create(config.Agent.LxcPrefix + "tmpdir/" + t.file)
	log.Check(log.FatalLevel, "Creating file "+t.file, err)
	defer out.Close()

	url := config.Cdn.Kurjun + "/template/download?id=" + t.id
	if len(t.owner) > 0 {
		url = config.Cdn.Kurjun + "/template/" + t.owner[0] + "/" + t.file
	}
	response, err := kurjun.Get(url)
	log.Check(log.FatalLevel, "Getting "+url, err)

	defer response.Body.Close()
	bar := pb.New(int(response.ContentLength)).SetUnits(pb.U_BYTES)
	bar.Start()
	rd := bar.NewProxyReader(response.Body)

	_, err = io.Copy(out, rd)
	for c := 0; err != nil && c < 5; _, err = io.Copy(out, rd) {
		log.Info("Download interrupted, retrying")
		time.Sleep(3 * time.Second)
		c++

		//Repeating GET request to CDN, while need to continue interrupted download
		out, err = os.Create(config.Agent.LxcPrefix + "tmpdir/" + t.file)
		log.Check(log.FatalLevel, "Creating file "+t.file, err)
		defer out.Close()
		response, err = kurjun.Get(url)
		log.Check(log.FatalLevel, "Getting "+url, err)
		defer response.Body.Close()
		bar = pb.New(int(response.ContentLength)).SetUnits(pb.U_BYTES)
		bar.Start()
		rd = bar.NewProxyReader(response.Body)
	}

	log.Check(log.FatalLevel, "Writing response body to file", err)

	if id := strings.Split(t.id, "."); len(id) > 0 && id[len(id)-1] == md5sum(config.Agent.LxcPrefix+"tmpdir/"+t.file) {
		return true
	}
	return false
}

func idToName(id string, kurjun *http.Client, token string) string {
	var meta metainfo

	//Since only kurjun knows template's ID, we cannot define if we have template already installed in system by ID as we do it by name, so unreachable kurjun in this case is a deadend for us
	//To omit this issue we should add ID into template config and use this ID as a "primary key" to any request
	response, err := kurjun.Get(config.Cdn.Kurjun + "/template/info?id=" + id + "&token=" + token)
	log.Check(log.ErrorLevel, "Getting kurjun response", err)
	defer response.Body.Close()

	body, err := ioutil.ReadAll(response.Body)

	if string(body) == "Not found" {
		log.Error("Template with id \"" + id + "\" not found")
	}
	log.Check(log.ErrorLevel, "Parsing response body", json.Unmarshal(body, &meta))

	return meta.Name
}

func lockSubutai(file string) bool {
	lock, err := lockfile.New("/var/run/lock/subutai." + file)
	if log.Check(log.DebugLevel, "Init lock "+file, err) {
		return false
	}

	err = lock.TryLock()
	if log.Check(log.DebugLevel, "Locking file "+file, err) {
		return false
	}

	return true
}

func unlockSubutai() {
	lock.Unlock()
}

func LxcImport(name, version, token string) {
	var kurjun *http.Client

	if container.IsContainer(name) && name == "management" && len(token) > 1 {
		gpg.ExchageAndEncrypt("management", token)
		return
	}

	if id := strings.Split(name, "id:"); len(id) > 1 {
		kurjun, _ = config.CheckKurjun()
		name = idToName(id[1], kurjun, token)
	}

	var t templ

	t.name = name
	if line := strings.Split(t.name, "/"); len(line) > 1 {
		t.name = line[1]
		t.owner = append(t.owner, line[0])
	}

	log.Info("Importing " + name)
	for !lockSubutai(t.name + ".import") {
		time.Sleep(time.Second * 1)
	}
	defer unlockSubutai()

	if container.IsContainer(t.name) {
		log.Info(t.name + " instance exist")
		return
	}

	t.version = config.Template.Version
	t.branch = config.Template.Branch
	if len(version) != 0 {
		if strings.Contains(version, "-") {
			verstr := strings.Split(version, "-")
			if len(verstr) == 2 {
				t.version = verstr[0]
				t.branch = verstr[1]
			} else {
				log.Error("Invalid version")
			}
		} else {
			if strings.Contains(version, ".") {
				t.version = version
				t.branch = config.Template.Branch
			} else {
				t.version = config.Template.Version
				t.branch = version
			}
		}
	}

	log.Info("Version: " + t.version + ", branch: " + t.branch)

	if t.branch == "" || t.name != "management" {
		t.file = t.name + "-subutai-template_" + t.version + "_" + config.Template.Arch + ".tar.gz"
	} else {
		t.file = t.name + "-subutai-template_" + t.version + "-" + t.branch + "_" + config.Template.Arch + ".tar.gz"
	}

	if kurjun == nil {
		kurjun, _ = config.CheckKurjun()
	}
	if kurjun != nil {
		templId(&t, kurjun, token)
	} else {
		log.Info("Trying to import from local storage")
	}

	if len(t.id) != 0 && len(t.signature) == 0 {
		log.Warn("Template is not signed")
	}

	for owner, signature := range t.signature {
		// if v.Author == "public" || v.Author == "subutai" || v.Author == "jenkins" {
		signedhash := gpg.VerifySignature(gpg.KurjunUserPK(owner), signature)
		if t.id != signedhash {
			log.Error("Signature does not match with template hash")
		}
		log.Info("Digital signature and file integrity verified")
		log.Debug("Signature belongs to " + owner)
		break
		// }
	}

	if !checkLocal(t) {
		log.Info("Downloading " + t.name)
		downloaded := false
		if len(t.owner) == 0 {
			for _, owner := range owners {
				if t.owner = []string{owner}; len(owner) == 0 {
					t.owner = []string{}
				}
				if download(t, kurjun, token) {
					downloaded = true
					break
				}
			}
		}
		if !downloaded && !download(t, kurjun, token) {
			log.Error("Failed to download template " + t.name)
		}
	}

	time.Sleep(time.Millisecond * 200) // Added sleep to prevent output collision with progress bar.

	log.Info("Unpacking template " + t.name)
	tgz := extractor.NewTgz()
	templdir := config.Agent.LxcPrefix + "tmpdir/" + t.name
	tgz.Extract(config.Agent.LxcPrefix+"tmpdir/"+t.file, templdir)
	parent := container.GetConfigItem(templdir+"/config", "subutai.parent")
	if parent != "" && parent != t.name && !container.IsTemplate(parent) {
		log.Info("Parent template required: " + parent)
		LxcImport(parent, "", token)
	}

	log.Info("Installing template " + t.name)
	template.Install(parent, t.name)
	// TODO following lines kept for back compatibility with old templates, should be deleted when all templates will be replaced.
	os.Rename(config.Agent.LxcPrefix+t.name+"/"+t.name+"-home", config.Agent.LxcPrefix+t.name+"/home")
	os.Rename(config.Agent.LxcPrefix+t.name+"/"+t.name+"-var", config.Agent.LxcPrefix+t.name+"/var")
	os.Rename(config.Agent.LxcPrefix+t.name+"/"+t.name+"-opt", config.Agent.LxcPrefix+t.name+"/opt")
	log.Check(log.FatalLevel, "Removing temp dir "+templdir, os.RemoveAll(templdir))

	if t.name == "management" {
		template.MngInit()
		return
	}

	container.SetContainerConf(t.name, [][]string{
		{"lxc.rootfs", config.Agent.LxcPrefix + t.name + "/rootfs"},
		{"lxc.mount", config.Agent.LxcPrefix + t.name + "/fstab"},
		{"lxc.hook.pre-start", ""},
		{"lxc.include", config.Agent.AppPrefix + "share/lxc/config/ubuntu.common.conf"},
		{"lxc.include", config.Agent.AppPrefix + "share/lxc/config/ubuntu.userns.conf"},
		{"subutai.config.path", config.Agent.AppPrefix + "etc"},
		{"lxc.network.script.up", config.Agent.AppPrefix + "bin/create_ovs_interface"},
		{"lxc.mount.entry", config.Agent.LxcPrefix + t.name + "/home home none bind,rw 0 0"},
		{"lxc.mount.entry", config.Agent.LxcPrefix + t.name + "/opt opt none bind,rw 0 0"},
		{"lxc.mount.entry", config.Agent.LxcPrefix + t.name + "/var var none bind,rw 0 0"},
	})
}
<|endoftext|>


<commit_before>// Copyright ©2012 The bíogo Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package pals

import (
	"code.google.com/p/biogo.interval"
	"code.google.com/p/biogo/feat"

	"fmt"
	"unsafe"
)

var duplicatePair = fmt.Errorf("pals: attempt to add duplicate feature pair to pile")

// Note Location must be comparable according to http://golang.org/ref/spec#Comparison_operators.
type PileInterval struct {
	Start, End int
	Location   feat.Feature
	Pairs      []*Pair
	overlap    int
}

func (i *PileInterval) Overlap(b interval.IntRange) bool {
	return i.End-i.overlap >= b.Start && i.Start <= b.End-i.overlap
}
func (i *PileInterval) ID() uintptr              { return uintptr(unsafe.Pointer(i)) }
func (i *PileInterval) Range() interval.IntRange { return interval.IntRange{i.Start, i.End} }

type containQuery struct {
	start, end int
	slop       int
	location   feat.Feature
}

func (q containQuery) Overlap(b interval.IntRange) bool {
	return b.Start <= q.start+q.slop && b.End >= q.end-q.slop
}
func (q containQuery) ID() uintptr              { return 0 }
func (q containQuery) Range() interval.IntRange { return interval.IntRange{q.start, q.end} }

// A Piler performs the aggregation of feature pairs according to the description in section 2.3
// of Edgar and Myers (2005) using an interval tree, giving O(nlogn) time but better space complexity
// and flexibility with feature overlap.
type Piler struct {
	intervals map[feat.Feature]*interval.IntTree
	seen      map[sp]struct{}
	overlap   int
}

type (
	sf struct {
		loc  feat.Feature
		s, e int
	}

	sp struct {
		a, b sf
	}
)

// NewPiler creates a Piler object ready for piling feature pairs.
func NewPiler(overlap int) *Piler {
	return &Piler{
		intervals: make(map[feat.Feature]*interval.IntTree),
		seen:      make(map[sp]struct{}),
		overlap:   overlap,
	}
}

// Add adds a feature pair to the piler incorporating the features into piles where appropriate.
func (p *Piler) Add(fp *Pair) error {
	a := sf{fp.A.Location(), fp.A.Start(), fp.A.End()}
	b := sf{fp.B.Location(), fp.B.Start(), fp.B.End()}
	ab, ba := sp{a, b}, sp{b, a}

	if _, ok := p.seen[ab]; ok {
		return duplicatePair
	}
	if _, ok := p.seen[ba]; ok {
		return duplicatePair
	}

	p.merge(&PileInterval{fp.A.Start(), fp.A.End(), fp.A.Location(), []*Pair{fp}, p.overlap})
	p.merge(&PileInterval{fp.B.Start(), fp.B.End(), fp.B.Location(), nil, p.overlap})
	p.seen[ab] = struct{}{}

	return nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// merge merges an interval into the tree moving location meta data from the replaced intervals
// into the new interval.
func (p *Piler) merge(pi *PileInterval) {
	var (
		f  = true
		r  []interval.IntInterface
		qi = &PileInterval{Start: pi.Start, End: pi.End}
	)
	t, ok := p.intervals[pi.Location]
	if !ok {
		t = &interval.IntTree{}
		p.intervals[pi.Location] = t
	}
	t.DoMatching(
		func(e interval.IntInterface) (done bool) {
			iv := e.(*PileInterval)
			r = append(r, e)
			pi.Pairs = append(pi.Pairs, iv.Pairs...)
			if f {
				pi.Start = min(iv.Start, pi.Start)
				f = false
			}
			pi.End = max(iv.End, pi.End)
			return
		},
		qi,
	)
	for _, d := range r {
		t.Delete(d, false)
	}
	t.Insert(pi, false)
}

// A PileFilter is used to determine whether a Pair is included in a Pile
type PileFilter func(a, b feat.Feature, pa, pb *PileInterval) bool

// Piles returns a slice of piles determined by application of the filter function f to
// the feature pairs that have been added to the piler.
func (p *Piler) Piles(f PileFilter) ([]*Pile, error) {
	var (
		pm  = make(map[*PileInterval]*Pile)
		err error
	)
	for _, t := range p.intervals {
		t.Do(
			func(e interval.IntInterface) (done bool) {
				var (
					pa = e.(*PileInterval)
					pb *PileInterval
				)
				for _, pp := range pa.Pairs {
					pb, err = p.pile(pp.B)
					if err != nil {
						return true // Terminate Do() and allow Piles() to return err.
					}

					if f != nil && !f(pp.A, pp.B, pa, pb) {
						continue
					}
					if wp, ok := pm[pa]; !ok {
						tp := &Pile{
							Loc: pa.Location, From: pa.Start, To: pa.End,
							Images: []*Pair{pp},
						}
						pp.A.(*Feature).Loc = tp
						pm[pa] = tp
					} else {
						pp.A.(*Feature).Loc = wp
						wp.Images = append(wp.Images, pp)
					}
					if wp, ok := pm[pb]; !ok {
						tp := &Pile{
							Loc: pb.Location, From: pb.Start, To: pb.End,
							Images: []*Pair{pp.Invert()},
						}
						pp.B.(*Feature).Loc = tp
						pm[pb] = tp
					} else {
						pp.B.(*Feature).Loc = wp
						wp.Images = append(wp.Images, pp.Invert())
					}
				}
				return
			},
		)
	}
	if err != nil {
		return nil, err
	}

	piles := make([]*Pile, 0, len(pm))
	for _, pile := range pm {
		piles = append(piles, pile)
	}

	return piles, nil
}

// Pile returns a Pile representation of the pile containing i.
// An error is returned if more than one pile would be returned.
func (p *Piler) Pile(q feat.Feature) (*Pile, error) {
	pi, err := p.pile(q)
	if err != nil {
		return nil, err
	}
	return &Pile{
		From:   pi.Start,
		To:     pi.End,
		Loc:    pi.Location,
		Images: pi.Pairs,
	}, nil
}

func (p *Piler) pile(q feat.Feature) (*PileInterval, error) {
	var (
		qi = containQuery{
			start:    q.Start(),
			end:      q.End(),
			location: q.Location(),
			slop:     p.overlap,
		}
		t  = p.intervals[qi.location]
		c  = 0
		pt interval.IntInterface
	)

	t.DoMatching(
		func(e interval.IntInterface) (done bool) {
			c++
			pt = e
			return
		},
		qi,
	)

	// Sanity check: no pile should overlap any other pile within overlap constraints
	// TODO: Should this be a panic?
	if c > 1 {
		return nil, fmt.Errorf("pals: internal inconsistency - too many results: %d", c)
	}

	return pt.(*PileInterval), nil
}
<commit_msg>Make types safer to change<commit_after>// Copyright ©2012 The bíogo Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package pals

import (
	"code.google.com/p/biogo.interval"
	"code.google.com/p/biogo/feat"

	"fmt"
	"unsafe"
)

var duplicatePair = fmt.Errorf("pals: attempt to add duplicate feature pair to pile")

// Note Location must be comparable according to http://golang.org/ref/spec#Comparison_operators.
type PileInterval struct {
	Start, End int
	Location   feat.Feature
	Pairs      []*Pair
	overlap    int
}

func (i *PileInterval) Overlap(b interval.IntRange) bool {
	return i.End-i.overlap >= b.Start && i.Start <= b.End-i.overlap
}
func (i *PileInterval) ID() uintptr              { return uintptr(unsafe.Pointer(i)) }
func (i *PileInterval) Range() interval.IntRange { return interval.IntRange{Start: i.Start, End: i.End} }

type containQuery struct {
	start, end int
	slop       int
	location   feat.Feature
}

func (q containQuery) Overlap(b interval.IntRange) bool {
	return b.Start <= q.start+q.slop && b.End >= q.end-q.slop
}
func (q containQuery) ID() uintptr              { return 0 }
func (q containQuery) Range() interval.IntRange { return interval.IntRange{Start: q.start, End: q.end} }

// A Piler performs the aggregation of feature pairs according to the description in section 2.3
// of Edgar and Myers (2005) using an interval tree, giving O(nlogn) time but better space complexity
// and flexibility with feature overlap.
type Piler struct {
	intervals map[feat.Feature]*interval.IntTree
	seen      map[[2]sf]struct{}
	overlap   int
}

type sf struct {
	loc  feat.Feature
	s, e int
}

// NewPiler creates a Piler object ready for piling feature pairs.
func NewPiler(overlap int) *Piler {
	return &Piler{
		intervals: make(map[feat.Feature]*interval.IntTree),
		seen:      make(map[[2]sf]struct{}),
		overlap:   overlap,
	}
}

// Add adds a feature pair to the piler incorporating the features into piles where appropriate.
func (p *Piler) Add(fp *Pair) error {
	a := sf{loc: fp.A.Location(), s: fp.A.Start(), e: fp.A.End()}
	b := sf{loc: fp.B.Location(), s: fp.B.Start(), e: fp.B.End()}
	ab, ba := [2]sf{a, b}, [2]sf{b, a}

	if _, ok := p.seen[ab]; ok {
		return duplicatePair
	}
	if _, ok := p.seen[ba]; ok {
		return duplicatePair
	}

	p.merge(&PileInterval{Start: fp.A.Start(), End: fp.A.End(), Location: fp.A.Location(), Pairs: []*Pair{fp}, overlap: p.overlap})
	p.merge(&PileInterval{Start: fp.B.Start(), End: fp.B.End(), Location: fp.B.Location(), Pairs: nil, overlap: p.overlap})
	p.seen[ab] = struct{}{}

	return nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// merge merges an interval into the tree moving location meta data from the replaced intervals
// into the new interval.
func (p *Piler) merge(pi *PileInterval) {
	var (
		f  = true
		r  []interval.IntInterface
		qi = &PileInterval{Start: pi.Start, End: pi.End}
	)
	t, ok := p.intervals[pi.Location]
	if !ok {
		t = &interval.IntTree{}
		p.intervals[pi.Location] = t
	}
	t.DoMatching(
		func(e interval.IntInterface) (done bool) {
			iv := e.(*PileInterval)
			r = append(r, e)
			pi.Pairs = append(pi.Pairs, iv.Pairs...)
			if f {
				pi.Start = min(iv.Start, pi.Start)
				f = false
			}
			pi.End = max(iv.End, pi.End)
			return
		},
		qi,
	)
	for _, d := range r {
		t.Delete(d, false)
	}
	t.Insert(pi, false)
}

// A PileFilter is used to determine whether a Pair is included in a Pile
type PileFilter func(a, b feat.Feature, pa, pb *PileInterval) bool

// Piles returns a slice of piles determined by application of the filter function f to
// the feature pairs that have been added to the piler.
func (p *Piler) Piles(f PileFilter) ([]*Pile, error) {
	var (
		pm  = make(map[*PileInterval]*Pile)
		err error
	)
	for _, t := range p.intervals {
		t.Do(
			func(e interval.IntInterface) (done bool) {
				var (
					pa = e.(*PileInterval)
					pb *PileInterval
				)
				for _, pp := range pa.Pairs {
					pb, err = p.pile(pp.B)
					if err != nil {
						return true // Terminate Do() and allow Piles() to return err.
					}

					if f != nil && !f(pp.A, pp.B, pa, pb) {
						continue
					}
					if wp, ok := pm[pa]; !ok {
						tp := &Pile{
							Loc: pa.Location, From: pa.Start, To: pa.End,
							Images: []*Pair{pp},
						}
						pp.A.(*Feature).Loc = tp
						pm[pa] = tp
					} else {
						pp.A.(*Feature).Loc = wp
						wp.Images = append(wp.Images, pp)
					}
					if wp, ok := pm[pb]; !ok {
						tp := &Pile{
							Loc: pb.Location, From: pb.Start, To: pb.End,
							Images: []*Pair{pp.Invert()},
						}
						pp.B.(*Feature).Loc = tp
						pm[pb] = tp
					} else {
						pp.B.(*Feature).Loc = wp
						wp.Images = append(wp.Images, pp.Invert())
					}
				}
				return
			},
		)
	}
	if err != nil {
		return nil, err
	}

	piles := make([]*Pile, 0, len(pm))
	for _, pile := range pm {
		piles = append(piles, pile)
	}

	return piles, nil
}

// Pile returns a Pile representation of the pile containing i.
// An error is returned if more than one pile would be returned.
func (p *Piler) Pile(q feat.Feature) (*Pile, error) {
	pi, err := p.pile(q)
	if err != nil {
		return nil, err
	}
	return &Pile{
		From:   pi.Start,
		To:     pi.End,
		Loc:    pi.Location,
		Images: pi.Pairs,
	}, nil
}

func (p *Piler) pile(q feat.Feature) (*PileInterval, error) {
	var (
		qi = containQuery{
			start:    q.Start(),
			end:      q.End(),
			location: q.Location(),
			slop:     p.overlap,
		}
		t  = p.intervals[qi.location]
		c  = 0
		pt interval.IntInterface
	)

	t.DoMatching(
		func(e interval.IntInterface) (done bool) {
			c++
			pt = e
			return
		},
		qi,
	)

	// Sanity check: no pile should overlap any other pile within overlap constraints
	// TODO: Should this be a panic?
	if c > 1 {
		return nil, fmt.Errorf("pals: internal inconsistency - too many results: %d", c)
	}

	return pt.(*PileInterval), nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// 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 (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/GoogleCloudPlatform/cos-customizer/src/pkg/config"
	"github.com/GoogleCloudPlatform/cos-customizer/src/pkg/fs"
	"github.com/GoogleCloudPlatform/cos-customizer/src/pkg/provisioner"
	"github.com/GoogleCloudPlatform/cos-customizer/src/pkg/utils"

	"cloud.google.com/go/storage"
	"github.com/google/subcommands"
	"google.golang.org/api/iterator"
)

const (
	gpuScript          = "install_gpu.sh"
	installerContainer = "gcr.io/cos-cloud/cos-gpu-installer:v20210204"
)

// TODO(b/121332360): Move most GPU functionality to cos-gpu-installer
var (
	validGPUs = []string{"nvidia-tesla-k80", "nvidia-tesla-p100", "nvidia-tesla-v100"}
)

// InstallGPU implements subcommands.Command for the "install-gpu" command.
// This command configures the current image build process to customize the result image
// with GPU drivers.
type InstallGPU struct {
	NvidiaDriverVersion  string
	NvidiaDriverMd5sum   string
	NvidiaInstallDirHost string
	gpuType              string
	getValidDrivers      bool
	gpuDataDir           string
}

// Name implements subcommands.Command.Name.
func (*InstallGPU) Name() string {
	return "install-gpu"
}

// Synopsis implements subcommands.Command.Synopsis.
func (*InstallGPU) Synopsis() string {
	return "Configure the image build with GPU drivers."
}

// Usage implements subcommands.Command.Usage.
func (*InstallGPU) Usage() string {
	return `install-gpu [flags]
`
}

// SetFlags implements subcommands.Command.SetFlags.
func (i *InstallGPU) SetFlags(f *flag.FlagSet) {
	f.StringVar(&i.NvidiaDriverVersion, "version", "", "Driver version to install. Can also be the name of an nvidia installer present in the "+
		"directory specified by '-deps-dir'; e.g., NVIDIA-Linux-x86_64-450.51.06.run.")
	f.StringVar(&i.NvidiaDriverMd5sum, "md5sum", "", "Md5sum of the driver to install.")
	f.StringVar(&i.NvidiaInstallDirHost, "install-dir", "/var/lib/nvidia",
		"Location to install drivers on the image.")
	f.StringVar(
		&i.gpuType, "gpu-type", "nvidia-tesla-p100",
		fmt.Sprintf("The type of GPU to verify drivers for. Must be one of: %v", validGPUs))
	f.BoolVar(
		&i.getValidDrivers, "get-valid-drivers", false,
		"Print the list of supported GPU driver versions. If this flag is given, no other actions will be taken.")
	f.StringVar(&i.gpuDataDir, "deps-dir", "", "If provided, the local directory to search for cos-gpu-installer data dependencies. "+
		"The exact data dependencies that must be present in this directory depends on the version of cos-gpu-installer "+
		"used by cos-customizer. Do not expect this flag to be stable; it exists for compatibility with pre-release COS images.")
}

func validDriverVersions(ctx context.Context, gcsClient *storage.Client) (map[string]bool, error) {
	// We gather the set of valid drivers from the set of drivers provided by Nvidia in their GCS bucket.
	// Nominally, paths we care about in this bucket look like 'tesla/<version>/<binaries>'. Version 390.46 has
	// a deprecated path structure, and since it's supported by cos-gpu-installer, we special case that here.
	validDrivers := map[string]bool{"390.46": true}
	query := &storage.Query{Prefix: "tesla/"}
	it := gcsClient.Bucket("nvidia-drivers-us-public").Objects(ctx, query)
	for {
		objAttrs, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}
		// Example object: tesla/396.26/NVIDIA-Linux-x86_64-396.26-diagnostic.run
		if splitPath := strings.SplitN(objAttrs.Name, "/", 3); len(splitPath) > 1 {
			validDrivers[splitPath[1]] = true
		}
	}
	return validDrivers, nil
}

func (i *InstallGPU) validate(ctx context.Context, gcsClient *storage.Client, files *fs.Files, provConfig *provisioner.Config) error {
	isValidGPU := false
	for _, g := range validGPUs {
		if i.gpuType == g {
			isValidGPU = true
			break
		}
	}
	if !isValidGPU {
		return fmt.Errorf("%q is an invalid GPU type. Must be one of: %v", i.gpuType, validGPUs)
	}
	if i.NvidiaDriverVersion == "" {
		return fmt.Errorf("version must be set")
	}
	var gpuAlreadyConf bool
	for _, s := range provConfig.Steps {
		if s.Type == "InstallGPU" {
			gpuAlreadyConf = true
			break
		}
	}
	if gpuAlreadyConf {
		return fmt.Errorf("install-gpu can only be invoked once in an image build process. Only one driver version can be installed on the image")
	}
	if strings.HasSuffix(i.NvidiaDriverVersion, ".run") {
		log.Printf("driver version is set to %q, which looks like an nvidia installer file", i.NvidiaDriverVersion)
		if i.gpuDataDir == "" {
			return errors.New(`"-deps-dir" must be set when the version is specified as an nvidia installer file`)
		}
		fileName := filepath.Join(i.gpuDataDir, i.NvidiaDriverVersion)
		info, err := os.Stat(fileName)
		if os.IsNotExist(err) {
			return fmt.Errorf("nvidia installer file at %q does not exist", fileName)
		}
		if info.IsDir() {
			return fmt.Errorf("nvidia installer file at %q is a directory", fileName)
		}
	} else {
		validDrivers, err := validDriverVersions(ctx, gcsClient)
		if err != nil {
			return err
		}
		if !validDrivers[i.NvidiaDriverVersion] {
			var drivers []string
			for d := range validDrivers {
				drivers = append(drivers, d)
			}
			return fmt.Errorf("driver version %s is not valid; valid driver versions are: %v", i.NvidiaDriverVersion, drivers)
		}
	}
	return nil
}

func (i *InstallGPU) templateScript(scriptPath string) error {
	setCOSDownloadGCS := ""
	if i.gpuDataDir != "" {
		setCOSDownloadGCS = "true"
	}
	data := struct {
		NvidiaDriverVersion  string
		NvidiaDriverMd5sum   string
		NvidiaInstallDirHost string
		SetCOSDownloadGCS    string
	}{
		NvidiaDriverVersion:  utils.QuoteForShell(i.NvidiaDriverVersion),
		NvidiaDriverMd5sum:   utils.QuoteForShell(i.NvidiaDriverMd5sum),
		NvidiaInstallDirHost: utils.QuoteForShell(i.NvidiaInstallDirHost),
		SetCOSDownloadGCS:    utils.QuoteForShell(setCOSDownloadGCS),
	}
	tmpl, err := template.New(filepath.Base(scriptPath)).ParseFiles(scriptPath)
	if err != nil {
		return err
	}
	w, err := os.Create(scriptPath)
	if err != nil {
		return err
	}
	defer w.Close()
	return tmpl.Execute(w, data)
}

func (i *InstallGPU) updateBuildConfig(configPath string) error {
	buildConfig := &config.Build{}
	configFile, err := os.OpenFile(configPath, os.O_RDWR, 0666)
	if err != nil {
		return err
	}
	defer configFile.Close()
	if err := config.Load(configFile, buildConfig); err != nil {
		return err
	}
	buildConfig.GPUType = i.gpuType
	if i.gpuDataDir != "" {
		files, err := ioutil.ReadDir(i.gpuDataDir)
		if err != nil {
			return fmt.Errorf("error reading dir %q: %v", i.gpuDataDir, err)
		}
		for _, f := range files {
			if f.Mode().IsRegular() {
				buildConfig.GCSFiles = append(buildConfig.GCSFiles, filepath.Join(i.gpuDataDir, f.Name()))
			}
		}
	}
	if _, err := configFile.Seek(0, 0); err != nil {
		return err
	}
	return config.SaveConfigToFile(configFile, buildConfig)
}

func (i *InstallGPU) updateProvConfig(provConfig *provisioner.Config) error {
	buf, err := json.Marshal(&provisioner.InstallGPUStep{
		NvidiaDriverVersion:      i.NvidiaDriverVersion,
		NvidiaDriverMD5Sum:       i.NvidiaDriverMd5sum,
		NvidiaInstallDirHost:     i.NvidiaInstallDirHost,
		NvidiaInstallerContainer: installerContainer,
		// GCSDepsPrefix will be converted into a gs:// address by the preloader
		// package.
		GCSDepsPrefix: i.gpuDataDir,
	})
	if err != nil {
		return err
	}
	provConfig.Steps = append(provConfig.Steps, provisioner.StepConfig{
		Type: "InstallGPU",
		Args: json.RawMessage(buf),
	})
	return nil
}

// Execute implements subcommands.Command.Execute. It configures the current image build process to
// customize the result image with GPU drivers.
func (i *InstallGPU) Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
	if f.NArg() != 0 {
		f.Usage()
		return subcommands.ExitUsageError
	}
	if len(args) < 2 {
		log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
	}
	files, ok := args[0].(*fs.Files)
	if !ok {
		log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
	}
	serviceClients, ok := args[1].(ServiceClients)
	if !ok {
		log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
	}
	_, gcsClient, err := serviceClients(ctx, true)
	if err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	defer gcsClient.Close()
	if i.getValidDrivers {
		validDrivers, err := validDriverVersions(ctx, gcsClient)
		if err != nil {
			log.Println(err)
			return subcommands.ExitFailure
		}
		var drivers []string
		for d := range validDrivers {
			drivers = append(drivers, d)
		}
		log.Printf("Valid driver versions are: %v\n", drivers)
		return subcommands.ExitSuccess
	}
	var provConfig provisioner.Config
	if err := config.LoadFromFile(files.ProvConfig, &provConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	if err := i.validate(ctx, gcsClient, files, &provConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	if err := i.updateBuildConfig(files.BuildConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	if err := i.updateProvConfig(&provConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	if err := config.SaveConfigToPath(files.ProvConfig, &provConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	return subcommands.ExitSuccess
}
<commit_msg>Update cos-gpu-installer to v20210319<commit_after>// Copyright 2018 Google LLC
//
// 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 (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/GoogleCloudPlatform/cos-customizer/src/pkg/config"
	"github.com/GoogleCloudPlatform/cos-customizer/src/pkg/fs"
	"github.com/GoogleCloudPlatform/cos-customizer/src/pkg/provisioner"
	"github.com/GoogleCloudPlatform/cos-customizer/src/pkg/utils"

	"cloud.google.com/go/storage"
	"github.com/google/subcommands"
	"google.golang.org/api/iterator"
)

const (
	gpuScript          = "install_gpu.sh"
	installerContainer = "gcr.io/cos-cloud/cos-gpu-installer:v20210319"
)

// TODO(b/121332360): Move most GPU functionality to cos-gpu-installer
var (
	validGPUs = []string{"nvidia-tesla-k80", "nvidia-tesla-p100", "nvidia-tesla-v100"}
)

// InstallGPU implements subcommands.Command for the "install-gpu" command.
// This command configures the current image build process to customize the result image
// with GPU drivers.
type InstallGPU struct {
	NvidiaDriverVersion  string
	NvidiaDriverMd5sum   string
	NvidiaInstallDirHost string
	gpuType              string
	getValidDrivers      bool
	gpuDataDir           string
}

// Name implements subcommands.Command.Name.
func (*InstallGPU) Name() string {
	return "install-gpu"
}

// Synopsis implements subcommands.Command.Synopsis.
func (*InstallGPU) Synopsis() string {
	return "Configure the image build with GPU drivers."
}

// Usage implements subcommands.Command.Usage.
func (*InstallGPU) Usage() string {
	return `install-gpu [flags]
`
}

// SetFlags implements subcommands.Command.SetFlags.
func (i *InstallGPU) SetFlags(f *flag.FlagSet) {
	f.StringVar(&i.NvidiaDriverVersion, "version", "", "Driver version to install. Can also be the name of an nvidia installer present in the "+
		"directory specified by '-deps-dir'; e.g., NVIDIA-Linux-x86_64-450.51.06.run.")
	f.StringVar(&i.NvidiaDriverMd5sum, "md5sum", "", "Md5sum of the driver to install.")
	f.StringVar(&i.NvidiaInstallDirHost, "install-dir", "/var/lib/nvidia",
		"Location to install drivers on the image.")
	f.StringVar(
		&i.gpuType, "gpu-type", "nvidia-tesla-p100",
		fmt.Sprintf("The type of GPU to verify drivers for. Must be one of: %v", validGPUs))
	f.BoolVar(
		&i.getValidDrivers, "get-valid-drivers", false,
		"Print the list of supported GPU driver versions. If this flag is given, no other actions will be taken.")
	f.StringVar(&i.gpuDataDir, "deps-dir", "", "If provided, the local directory to search for cos-gpu-installer data dependencies. "+
		"The exact data dependencies that must be present in this directory depends on the version of cos-gpu-installer "+
		"used by cos-customizer. Do not expect this flag to be stable; it exists for compatibility with pre-release COS images.")
}

func validDriverVersions(ctx context.Context, gcsClient *storage.Client) (map[string]bool, error) {
	// We gather the set of valid drivers from the set of drivers provided by Nvidia in their GCS bucket.
	// Nominally, paths we care about in this bucket look like 'tesla/<version>/<binaries>'. Version 390.46 has
	// a deprecated path structure, and since it's supported by cos-gpu-installer, we special case that here.
	validDrivers := map[string]bool{"390.46": true}
	query := &storage.Query{Prefix: "tesla/"}
	it := gcsClient.Bucket("nvidia-drivers-us-public").Objects(ctx, query)
	for {
		objAttrs, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}
		// Example object: tesla/396.26/NVIDIA-Linux-x86_64-396.26-diagnostic.run
		if splitPath := strings.SplitN(objAttrs.Name, "/", 3); len(splitPath) > 1 {
			validDrivers[splitPath[1]] = true
		}
	}
	return validDrivers, nil
}

func (i *InstallGPU) validate(ctx context.Context, gcsClient *storage.Client, files *fs.Files, provConfig *provisioner.Config) error {
	isValidGPU := false
	for _, g := range validGPUs {
		if i.gpuType == g {
			isValidGPU = true
			break
		}
	}
	if !isValidGPU {
		return fmt.Errorf("%q is an invalid GPU type. Must be one of: %v", i.gpuType, validGPUs)
	}
	if i.NvidiaDriverVersion == "" {
		return fmt.Errorf("version must be set")
	}
	var gpuAlreadyConf bool
	for _, s := range provConfig.Steps {
		if s.Type == "InstallGPU" {
			gpuAlreadyConf = true
			break
		}
	}
	if gpuAlreadyConf {
		return fmt.Errorf("install-gpu can only be invoked once in an image build process. Only one driver version can be installed on the image")
	}
	if strings.HasSuffix(i.NvidiaDriverVersion, ".run") {
		log.Printf("driver version is set to %q, which looks like an nvidia installer file", i.NvidiaDriverVersion)
		if i.gpuDataDir == "" {
			return errors.New(`"-deps-dir" must be set when the version is specified as an nvidia installer file`)
		}
		fileName := filepath.Join(i.gpuDataDir, i.NvidiaDriverVersion)
		info, err := os.Stat(fileName)
		if os.IsNotExist(err) {
			return fmt.Errorf("nvidia installer file at %q does not exist", fileName)
		}
		if info.IsDir() {
			return fmt.Errorf("nvidia installer file at %q is a directory", fileName)
		}
	} else {
		validDrivers, err := validDriverVersions(ctx, gcsClient)
		if err != nil {
			return err
		}
		if !validDrivers[i.NvidiaDriverVersion] {
			var drivers []string
			for d := range validDrivers {
				drivers = append(drivers, d)
			}
			return fmt.Errorf("driver version %s is not valid; valid driver versions are: %v", i.NvidiaDriverVersion, drivers)
		}
	}
	return nil
}

func (i *InstallGPU) templateScript(scriptPath string) error {
	setCOSDownloadGCS := ""
	if i.gpuDataDir != "" {
		setCOSDownloadGCS = "true"
	}
	data := struct {
		NvidiaDriverVersion  string
		NvidiaDriverMd5sum   string
		NvidiaInstallDirHost string
		SetCOSDownloadGCS    string
	}{
		NvidiaDriverVersion:  utils.QuoteForShell(i.NvidiaDriverVersion),
		NvidiaDriverMd5sum:   utils.QuoteForShell(i.NvidiaDriverMd5sum),
		NvidiaInstallDirHost: utils.QuoteForShell(i.NvidiaInstallDirHost),
		SetCOSDownloadGCS:    utils.QuoteForShell(setCOSDownloadGCS),
	}
	tmpl, err := template.New(filepath.Base(scriptPath)).ParseFiles(scriptPath)
	if err != nil {
		return err
	}
	w, err := os.Create(scriptPath)
	if err != nil {
		return err
	}
	defer w.Close()
	return tmpl.Execute(w, data)
}

func (i *InstallGPU) updateBuildConfig(configPath string) error {
	buildConfig := &config.Build{}
	configFile, err := os.OpenFile(configPath, os.O_RDWR, 0666)
	if err != nil {
		return err
	}
	defer configFile.Close()
	if err := config.Load(configFile, buildConfig); err != nil {
		return err
	}
	buildConfig.GPUType = i.gpuType
	if i.gpuDataDir != "" {
		files, err := ioutil.ReadDir(i.gpuDataDir)
		if err != nil {
			return fmt.Errorf("error reading dir %q: %v", i.gpuDataDir, err)
		}
		for _, f := range files {
			if f.Mode().IsRegular() {
				buildConfig.GCSFiles = append(buildConfig.GCSFiles, filepath.Join(i.gpuDataDir, f.Name()))
			}
		}
	}
	if _, err := configFile.Seek(0, 0); err != nil {
		return err
	}
	return config.SaveConfigToFile(configFile, buildConfig)
}

func (i *InstallGPU) updateProvConfig(provConfig *provisioner.Config) error {
	buf, err := json.Marshal(&provisioner.InstallGPUStep{
		NvidiaDriverVersion:      i.NvidiaDriverVersion,
		NvidiaDriverMD5Sum:       i.NvidiaDriverMd5sum,
		NvidiaInstallDirHost:     i.NvidiaInstallDirHost,
		NvidiaInstallerContainer: installerContainer,
		// GCSDepsPrefix will be converted into a gs:// address by the preloader
		// package.
		GCSDepsPrefix: i.gpuDataDir,
	})
	if err != nil {
		return err
	}
	provConfig.Steps = append(provConfig.Steps, provisioner.StepConfig{
		Type: "InstallGPU",
		Args: json.RawMessage(buf),
	})
	return nil
}

// Execute implements subcommands.Command.Execute. It configures the current image build process to
// customize the result image with GPU drivers.
func (i *InstallGPU) Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
	if f.NArg() != 0 {
		f.Usage()
		return subcommands.ExitUsageError
	}
	if len(args) < 2 {
		log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
	}
	files, ok := args[0].(*fs.Files)
	if !ok {
		log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
	}
	serviceClients, ok := args[1].(ServiceClients)
	if !ok {
		log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
	}
	_, gcsClient, err := serviceClients(ctx, true)
	if err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	defer gcsClient.Close()
	if i.getValidDrivers {
		validDrivers, err := validDriverVersions(ctx, gcsClient)
		if err != nil {
			log.Println(err)
			return subcommands.ExitFailure
		}
		var drivers []string
		for d := range validDrivers {
			drivers = append(drivers, d)
		}
		log.Printf("Valid driver versions are: %v\n", drivers)
		return subcommands.ExitSuccess
	}
	var provConfig provisioner.Config
	if err := config.LoadFromFile(files.ProvConfig, &provConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	if err := i.validate(ctx, gcsClient, files, &provConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	if err := i.updateBuildConfig(files.BuildConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	if err := i.updateProvConfig(&provConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	if err := config.SaveConfigToPath(files.ProvConfig, &provConfig); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	return subcommands.ExitSuccess
}
<|endoftext|>


<commit_before>package x86_test

import (
	"bufio"
	"bytes"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"testing"
)

const testdata = `
MOVQ AX, AX -> MOVQ AX, AX

LEAQ name(SB), AX -> MOVQ name@GOT(SB), AX
LEAQ name+10(SB), AX -> MOVQ name@GOT(SB), AX; ADDQ $10, AX
MOVQ $name(SB), AX -> MOVQ name@GOT(SB), AX
MOVQ $name+10(SB), AX -> MOVQ name@GOT(SB), AX; ADDQ $10, AX

MOVQ name(SB), AX -> MOVQ name@GOT(SB), R15; MOVQ (R15), AX
MOVQ name+10(SB), AX -> MOVQ name@GOT(SB), R15; MOVQ 10(R15), AX

CMPQ name(SB), $0 -> MOVQ name@GOT(SB), R15; CMPQ (R15), $0

MOVQ $1, name(SB) -> MOVQ name@GOT(SB), R15; MOVQ $1, (R15)
MOVQ $1, name+10(SB) -> MOVQ name@GOT(SB), R15; MOVQ $1, 10(R15)
`

type ParsedTestData struct {
	input              string
	marks              []int
	marker_to_input    map[int][]string
	marker_to_expected map[int][]string
	marker_to_output   map[int][]string
}

const marker_start = 1234

func parseTestData(t *testing.T) *ParsedTestData {
	r := &ParsedTestData{}
	scanner := bufio.NewScanner(strings.NewReader(testdata))
	r.marker_to_input = make(map[int][]string)
	r.marker_to_expected = make(map[int][]string)
	marker := marker_start
	input_insns := []string{}
	for scanner.Scan() {
		line := scanner.Text()
		if len(strings.TrimSpace(line)) == 0 {
			continue
		}
		parts := strings.Split(line, "->")
		if len(parts) != 2 {
			t.Fatalf("malformed line %v", line)
		}
		r.marks = append(r.marks, marker)
		marker_insn := fmt.Sprintf("MOVQ $%d, AX", marker)
		input_insns = append(input_insns, marker_insn)
		for _, input_insn := range strings.Split(parts[0], ";") {
			input_insns = append(input_insns, input_insn)
			r.marker_to_input[marker] = append(r.marker_to_input[marker], normalize(input_insn))
		}
		for _, expected_insn := range strings.Split(parts[1], ";") {
			r.marker_to_expected[marker] = append(r.marker_to_expected[marker], normalize(expected_insn))
		}
		marker++
	}
	r.input = "TEXT ·foo(SB),$0\n" + strings.Join(input_insns, "\n") + "\n"
	return r
}

var spaces_re *regexp.Regexp = regexp.MustCompile("\\s+")
var marker_re *regexp.Regexp = regexp.MustCompile("MOVQ \\$([0-9]+), AX")

func normalize(s string) string {
	return spaces_re.ReplaceAllLiteralString(strings.TrimSpace(s), " ")
}

func asmOutput(t *testing.T, s string) []byte {
	tmpdir, err := ioutil.TempDir("", "progedittest")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tmpdir)
	tmpfile, err := os.Create(filepath.Join(tmpdir, "input.s"))
	if err != nil {
		t.Fatal(err)
	}
	defer tmpfile.Close()
	_, err = tmpfile.WriteString(s)
	if err != nil {
		t.Fatal(err)
	}
	gofolder := filepath.Join(build.Default.GOROOT, "bin")
	if gobin := os.Getenv("GOBIN"); len(gobin) != 0 {
		gofolder = gobin
	}

	cmd := exec.Command(
		filepath.Join(gofolder, "go"), "tool", "asm", "-S", "-dynlink",
		"-o", filepath.Join(tmpdir, "output.6"), tmpfile.Name())

	var env []string
	for _, v := range os.Environ() {
		if !strings.HasPrefix(v, "GOARCH=") {
			env = append(env, v)
		}
	}
	cmd.Env = append(env, "GOARCH=amd64")
	asmout, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("error %s output %s", err, asmout)
	}
	return asmout
}

func parseOutput(t *testing.T, td *ParsedTestData, asmout []byte) {
	scanner := bufio.NewScanner(bytes.NewReader(asmout))
	marker := regexp.MustCompile("MOVQ \\$([0-9]+), AX")
	mark := -1
	td.marker_to_output = make(map[int][]string)
	for scanner.Scan() {
		line := scanner.Text()
		if line[0] != '\t' {
			continue
		}
		parts := strings.SplitN(line, "\t", 3)
		if len(parts) != 3 {
			continue
		}
		n := normalize(parts[2])
		mark_matches := marker.FindStringSubmatch(n)
		if mark_matches != nil {
			mark, _ = strconv.Atoi(mark_matches[1])
			if _, ok := td.marker_to_input[mark]; !ok {
				t.Fatalf("unexpected marker %d", mark)
			}
		} else if mark != -1 {
			td.marker_to_output[mark] = append(td.marker_to_output[mark], n)
		}
	}
}

func TestDynlink(t *testing.T) {
	if runtime.GOOS == "nacl" || runtime.GOOS == "android" || (runtime.GOOS == "darwin" && runtime.GOARCH == "arm") {
		// iOS and nacl cannot fork
		t.Skipf("skipping on %s/%s", runtime.GOOS, runtime.GOARCH)
	}
	testdata := parseTestData(t)
	asmout := asmOutput(t, testdata.input)
	parseOutput(t, testdata, asmout)
	for _, m := range testdata.marks {
		i := strings.Join(testdata.marker_to_input[m], "; ")
		o := strings.Join(testdata.marker_to_output[m], "; ")
		e := strings.Join(testdata.marker_to_expected[m], "; ")
		if o != e {
			if o == i {
				t.Errorf("%s was unchanged; should have become %s", i, e)
			} else {
				t.Errorf("%s became %s; should have become %s", i, o, e)
			}
		} else if i != e {
			t.Logf("%s correctly became %s", i, o)
		}
	}
}
<commit_msg>cmd/internal/obj/x86: skip test on darwin/arm64<commit_after>package x86_test

import (
	"bufio"
	"bytes"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"testing"
)

const testdata = `
MOVQ AX, AX -> MOVQ AX, AX

LEAQ name(SB), AX -> MOVQ name@GOT(SB), AX
LEAQ name+10(SB), AX -> MOVQ name@GOT(SB), AX; ADDQ $10, AX
MOVQ $name(SB), AX -> MOVQ name@GOT(SB), AX
MOVQ $name+10(SB), AX -> MOVQ name@GOT(SB), AX; ADDQ $10, AX

MOVQ name(SB), AX -> MOVQ name@GOT(SB), R15; MOVQ (R15), AX
MOVQ name+10(SB), AX -> MOVQ name@GOT(SB), R15; MOVQ 10(R15), AX

CMPQ name(SB), $0 -> MOVQ name@GOT(SB), R15; CMPQ (R15), $0

MOVQ $1, name(SB) -> MOVQ name@GOT(SB), R15; MOVQ $1, (R15)
MOVQ $1, name+10(SB) -> MOVQ name@GOT(SB), R15; MOVQ $1, 10(R15)
`

type ParsedTestData struct {
	input              string
	marks              []int
	marker_to_input    map[int][]string
	marker_to_expected map[int][]string
	marker_to_output   map[int][]string
}

const marker_start = 1234

func parseTestData(t *testing.T) *ParsedTestData {
	r := &ParsedTestData{}
	scanner := bufio.NewScanner(strings.NewReader(testdata))
	r.marker_to_input = make(map[int][]string)
	r.marker_to_expected = make(map[int][]string)
	marker := marker_start
	input_insns := []string{}
	for scanner.Scan() {
		line := scanner.Text()
		if len(strings.TrimSpace(line)) == 0 {
			continue
		}
		parts := strings.Split(line, "->")
		if len(parts) != 2 {
			t.Fatalf("malformed line %v", line)
		}
		r.marks = append(r.marks, marker)
		marker_insn := fmt.Sprintf("MOVQ $%d, AX", marker)
		input_insns = append(input_insns, marker_insn)
		for _, input_insn := range strings.Split(parts[0], ";") {
			input_insns = append(input_insns, input_insn)
			r.marker_to_input[marker] = append(r.marker_to_input[marker], normalize(input_insn))
		}
		for _, expected_insn := range strings.Split(parts[1], ";") {
			r.marker_to_expected[marker] = append(r.marker_to_expected[marker], normalize(expected_insn))
		}
		marker++
	}
	r.input = "TEXT ·foo(SB),$0\n" + strings.Join(input_insns, "\n") + "\n"
	return r
}

var spaces_re *regexp.Regexp = regexp.MustCompile("\\s+")
var marker_re *regexp.Regexp = regexp.MustCompile("MOVQ \\$([0-9]+), AX")

func normalize(s string) string {
	return spaces_re.ReplaceAllLiteralString(strings.TrimSpace(s), " ")
}

func asmOutput(t *testing.T, s string) []byte {
	tmpdir, err := ioutil.TempDir("", "progedittest")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tmpdir)
	tmpfile, err := os.Create(filepath.Join(tmpdir, "input.s"))
	if err != nil {
		t.Fatal(err)
	}
	defer tmpfile.Close()
	_, err = tmpfile.WriteString(s)
	if err != nil {
		t.Fatal(err)
	}
	gofolder := filepath.Join(build.Default.GOROOT, "bin")
	if gobin := os.Getenv("GOBIN"); len(gobin) != 0 {
		gofolder = gobin
	}

	cmd := exec.Command(
		filepath.Join(gofolder, "go"), "tool", "asm", "-S", "-dynlink",
		"-o", filepath.Join(tmpdir, "output.6"), tmpfile.Name())

	var env []string
	for _, v := range os.Environ() {
		if !strings.HasPrefix(v, "GOARCH=") {
			env = append(env, v)
		}
	}
	cmd.Env = append(env, "GOARCH=amd64")
	asmout, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("error %s output %s", err, asmout)
	}
	return asmout
}

func parseOutput(t *testing.T, td *ParsedTestData, asmout []byte) {
	scanner := bufio.NewScanner(bytes.NewReader(asmout))
	marker := regexp.MustCompile("MOVQ \\$([0-9]+), AX")
	mark := -1
	td.marker_to_output = make(map[int][]string)
	for scanner.Scan() {
		line := scanner.Text()
		if line[0] != '\t' {
			continue
		}
		parts := strings.SplitN(line, "\t", 3)
		if len(parts) != 3 {
			continue
		}
		n := normalize(parts[2])
		mark_matches := marker.FindStringSubmatch(n)
		if mark_matches != nil {
			mark, _ = strconv.Atoi(mark_matches[1])
			if _, ok := td.marker_to_input[mark]; !ok {
				t.Fatalf("unexpected marker %d", mark)
			}
		} else if mark != -1 {
			td.marker_to_output[mark] = append(td.marker_to_output[mark], n)
		}
	}
}

func TestDynlink(t *testing.T) {
	iOS := runtime.GOOS == "darwin" && (runtime.GOARCH == "arm" || runtime.GOARCH == "arm64")
	if runtime.GOOS == "nacl" || runtime.GOOS == "android" || iOS {
		t.Skipf("skipping on %s/%s, cannot fork", runtime.GOOS, runtime.GOARCH)
	}
	testdata := parseTestData(t)
	asmout := asmOutput(t, testdata.input)
	parseOutput(t, testdata, asmout)
	for _, m := range testdata.marks {
		i := strings.Join(testdata.marker_to_input[m], "; ")
		o := strings.Join(testdata.marker_to_output[m], "; ")
		e := strings.Join(testdata.marker_to_expected[m], "; ")
		if o != e {
			if o == i {
				t.Errorf("%s was unchanged; should have become %s", i, e)
			} else {
				t.Errorf("%s became %s; should have become %s", i, o, e)
			}
		} else if i != e {
			t.Logf("%s correctly became %s", i, o)
		}
	}
}
<|endoftext|>