<commit_before>package main

import (
	"errors"
	"fmt"
	"sort"

	"github.com/Urethramancer/lbapi"
	"github.com/ryanuber/columnize"
)

// DNSGetCmd arguments.
type DNSGetCmd struct {
	A     DNSGetACmd     `command:"a" description:"Get A records for a domain." alias:"A"`
	AAAA  DNSGetAAAACmd  `command:"aaaa" description:"Get AAAA records for a domain." alias:"AAAA"`
	CNAME DNSGetCNAMECmd `command:"cname" description:"Get CNAME records for a domain." alias:"CNAME"`
	MX    DNSGetMXCmd    `command:"mx" description:"Get MX records for a domain." alias:"MX"`
	NS    DNSGetNSCmd    `command:"ns" description:"Get NS records for a domain." alias:"NS"`
	TXT   DNSGetTXTCmd   `command:"txt" description:"Get TXT records for a domain." alias:"TXT"`
	SRV   DNSGetSRVCmd   `command:"srv" description:"Get SRV records for a domain." alias:"SRV"`
	All   DNSGetAllCmd   `command:"all" description:"Get all records for a domain."`
}

// DNSGetACmd arguments.
type DNSGetACmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute A record fetch.
func (cmd *DNSGetACmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "A")
}

// DNSGetAAAACmd arguments.
type DNSGetAAAACmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute AAAA record fetch.
func (cmd *DNSGetAAAACmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "AAAA")
}

// DNSGetCNAMECmd arguments.
type DNSGetCNAMECmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute CNAME record fetch.
func (cmd *DNSGetCNAMECmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "CNAME")
}

// DNSGetMXCmd arguments.
type DNSGetMXCmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute MX record fetch.
func (cmd *DNSGetMXCmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "MX")
}

// DNSGetNSCmd arguments.
type DNSGetNSCmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute NS record fetch.
func (cmd *DNSGetNSCmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "NS")
}

// DNSGetTXTCmd arguments.
type DNSGetTXTCmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute TXT record fetch.
func (cmd *DNSGetTXTCmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "TXT")
}

// Print the requested records of a type.
func printRecord(domain, value, host, t string) error {
	var err error

	dns, err := getRecords(domain, value, host, t)
	if err != nil {
		return err
	}

	sort.Sort(dns)
	cc := columnize.DefaultConfig()
	cc.Delim = "\t"
	cc.Glue = "  "
	var s []string
	var pri bool
	switch t {
	case "mx", "MX", "srv", "SRV":
		pri = true
	}
	if pri {
		s = []string{fmt.Sprintln("Host\tAddress\tType\tTTL\tStatus\tPriority")}
	} else {
		s = []string{fmt.Sprintln("Host\tAddress\tType\tTTL\tStatus")}
	}
	for _, r := range dns {
		if pri {
			s = append(s, fmt.Sprintf("%s\t%s\t%s\t%d\t%s\t%d", r.Host, r.Value, r.Type, r.TTL, r.Status, r.Priority))
		} else {
			s = append(s, fmt.Sprintf("%s\t%s\t%s\t%d\t%s", r.Host, r.Value, r.Type, r.TTL, r.Status))
		}
	}
	res := columnize.Format(s, cc)
	pr(res)
	return nil
}

func getRecords(domain, value, host, t string) (lbapi.DNSRecords, error) {
	var err error
	var records lbapi.DNSRecords
	page := 1
	var count int64
	for {
		var recs *lbapi.DNSRecordList
		recs, err = client.GetDNSRecords(domain, host, value, t, page)
		if err != nil {
			return nil, err
		}

		if recs.Count == 0 {
			return nil, errors.New("no records matching arguments")
		}

		records = append(records, recs.Records...)
		page++
		count += recs.Count
		if count >= recs.MaxRecords {
			break
		}
	}
	return records, nil
}

// DNSGetAllCmd arguments.
type DNSGetAllCmd struct {
	Compact bool `short:"c" long:"compact" description:"Print records all in one table."`
	Args    struct {
		Domain string `required:"true" positional-arg-name:"DOMAIN" description:"Domain name."`
	} `positional-args:"true"`
}

// Execute a complete record fetch for the domain.
func (cmd *DNSGetAllCmd) Execute(args []string) error {
	var everything, a, aaaa, cname, mx, ns, txt, srv lbapi.DNSRecords
	var err error
	prn("Fetching A records…")
	a, err = getRecords(cmd.Args.Domain, "", "", "A")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(a)
	}

	prn("Fetching AAAA records…")
	aaaa, err = getRecords(cmd.Args.Domain, "", "", "AAAA")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(aaaa)
	}

	prn("Fetching CNAME records…")
	cname, err = getRecords(cmd.Args.Domain, "", "", "CNAME")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(cname)
	}

	prn("Fetching MX records…")
	mx, err = getRecords(cmd.Args.Domain, "", "", "MX")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(mx)
	}

	prn("Fetching NS records…")
	ns, err = getRecords(cmd.Args.Domain, "", "", "NS")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(ns)
	}

	prn("Fetching TXT records…")
	txt, err = getRecords(cmd.Args.Domain, "", "", "TXT")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(txt)
	}

	prn("Fetching SRV records…")
	srv, err = getRecords(cmd.Args.Domain, "", "", "SRV")
	if err != nil {
		pr("done.")
		pr("none found.")
	} else {
		sort.Sort(srv)
	}

	pr("Done: %v", everything)
	if cmd.Compact {
		everything = append(everything, a...)
		everything = append(everything, aaaa...)
		everything = append(everything, cname...)
		everything = append(everything, mx...)
		everything = append(everything, ns...)
		everything = append(everything, txt...)
		everything = append(everything, srv...)
		pr("%v", everything)
	} else {
		cc := columnize.DefaultConfig()
		cc.Delim = "\t"
		cc.Glue = "  "
		var s []string
		var res string
		base1 := []string{fmt.Sprintln("Host\tAddress\tType\tTTL\tStatus")}
		base2 := []string{fmt.Sprintln("Host\tAddress\tType\tTTL\tStatus\tPriority")}

		if len(a) > 0 {
			for _, r := range a {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("A records:\n%s\n", res)
		}

		if len(aaaa) > 0 {
			for _, r := range aaaa {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("AAAA records:\n%s\n", res)
		}

		if len(cname) > 0 {
			for _, r := range cname {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("CNAME records:\n%s\n", res)
		}

		if len(mx) > 0 {
			for _, r := range mx {
				s = append(base2, fmt.Sprintf("%s\t%s\t%d\t%s\t%d", r.Host, r.Value, r.TTL, r.Status, r.Priority))
			}
			res = columnize.Format(s, cc)
			pr("MX records:\n%s\n", res)
		}

		if len(ns) > 0 {
			for _, r := range ns {
				s = append(base2, fmt.Sprintf("%s\t%s\t%d\t%s\t%d", r.Host, r.Value, r.TTL, r.Status, r.Priority))
			}
			res = columnize.Format(s, cc)
			pr("NS records:\n%s\n", res)
		}

		if len(txt) > 0 {
			for _, r := range txt {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("TXT records:\n%s\n", res)
		}

		if len(txt) > 0 {
			for _, r := range txt {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("TXT records:\n%s\n", res)
		}
	}
	return nil
}

// DNSGetSRVCmd arguments.
type DNSGetSRVCmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute SRV record fetch.
func (cmd *DNSGetSRVCmd) Execute(args []string) error {
	srv, err := getRecords(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "SRV")
	if err != nil {
		return err
	}

	if len(srv) == 0 {
		return errors.New("no records found")
	}

	s := []string{fmt.Sprintln("Service name\tHost name\tPort\tWeight\tStatus")}
	for _, r := range srv {
		s = append(s, fmt.Sprintf("%s\t%s\t%d\t%d\t%s", r.Host, r.Value, r.Port, r.Weight, r.Status))
	}

	cc := columnize.DefaultConfig()
	cc.Delim = "\t"
	cc.Glue = "  "
	res := columnize.Format(s, cc)
	pr(res)
	return nil
}
<commit_msg>Type wasn't necessary in the output.<commit_after>package main

import (
	"errors"
	"fmt"
	"sort"

	"github.com/Urethramancer/lbapi"
	"github.com/ryanuber/columnize"
)

// DNSGetCmd arguments.
type DNSGetCmd struct {
	A     DNSGetACmd     `command:"a" description:"Get A records for a domain." alias:"A"`
	AAAA  DNSGetAAAACmd  `command:"aaaa" description:"Get AAAA records for a domain." alias:"AAAA"`
	CNAME DNSGetCNAMECmd `command:"cname" description:"Get CNAME records for a domain." alias:"CNAME"`
	MX    DNSGetMXCmd    `command:"mx" description:"Get MX records for a domain." alias:"MX"`
	NS    DNSGetNSCmd    `command:"ns" description:"Get NS records for a domain." alias:"NS"`
	TXT   DNSGetTXTCmd   `command:"txt" description:"Get TXT records for a domain." alias:"TXT"`
	SRV   DNSGetSRVCmd   `command:"srv" description:"Get SRV records for a domain." alias:"SRV"`
	All   DNSGetAllCmd   `command:"all" description:"Get all records for a domain."`
}

// DNSGetACmd arguments.
type DNSGetACmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute A record fetch.
func (cmd *DNSGetACmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "A")
}

// DNSGetAAAACmd arguments.
type DNSGetAAAACmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute AAAA record fetch.
func (cmd *DNSGetAAAACmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "AAAA")
}

// DNSGetCNAMECmd arguments.
type DNSGetCNAMECmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute CNAME record fetch.
func (cmd *DNSGetCNAMECmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "CNAME")
}

// DNSGetMXCmd arguments.
type DNSGetMXCmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute MX record fetch.
func (cmd *DNSGetMXCmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "MX")
}

// DNSGetNSCmd arguments.
type DNSGetNSCmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute NS record fetch.
func (cmd *DNSGetNSCmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "NS")
}

// DNSGetTXTCmd arguments.
type DNSGetTXTCmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute TXT record fetch.
func (cmd *DNSGetTXTCmd) Execute(args []string) error {
	return printRecord(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "TXT")
}

// Print the requested records of a type.
func printRecord(domain, value, host, t string) error {
	var err error

	dns, err := getRecords(domain, value, host, t)
	if err != nil {
		return err
	}

	sort.Sort(dns)
	cc := columnize.DefaultConfig()
	cc.Delim = "\t"
	cc.Glue = "  "
	var s []string
	var pri bool
	switch t {
	case "mx", "MX", "srv", "SRV":
		pri = true
	}
	if pri {
		s = []string{fmt.Sprintln("Host\tAddress\tTTL\tStatus\tPriority")}
	} else {
		s = []string{fmt.Sprintln("Host\tAddress\tTTL\tStatus")}
	}
	for _, r := range dns {
		if pri {
			s = append(s, fmt.Sprintf("%s\t%s\t%d\t%s\t%d", r.Host, r.Value, r.TTL, r.Status, r.Priority))
		} else {
			s = append(s, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
		}
	}
	res := columnize.Format(s, cc)
	pr(res)
	return nil
}

func getRecords(domain, value, host, t string) (lbapi.DNSRecords, error) {
	var err error
	var records lbapi.DNSRecords
	page := 1
	var count int64
	for {
		var recs *lbapi.DNSRecordList
		recs, err = client.GetDNSRecords(domain, host, value, t, page)
		if err != nil {
			return nil, err
		}

		if recs.Count == 0 {
			return nil, errors.New("no records matching arguments")
		}

		records = append(records, recs.Records...)
		page++
		count += recs.Count
		if count >= recs.MaxRecords {
			break
		}
	}
	return records, nil
}

// DNSGetAllCmd arguments.
type DNSGetAllCmd struct {
	Compact bool `short:"c" long:"compact" description:"Print records all in one table."`
	Args    struct {
		Domain string `required:"true" positional-arg-name:"DOMAIN" description:"Domain name."`
	} `positional-args:"true"`
}

// Execute a complete record fetch for the domain.
func (cmd *DNSGetAllCmd) Execute(args []string) error {
	var everything, a, aaaa, cname, mx, ns, txt, srv lbapi.DNSRecords
	var err error
	prn("Fetching A records…")
	a, err = getRecords(cmd.Args.Domain, "", "", "A")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(a)
	}

	prn("Fetching AAAA records…")
	aaaa, err = getRecords(cmd.Args.Domain, "", "", "AAAA")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(aaaa)
	}

	prn("Fetching CNAME records…")
	cname, err = getRecords(cmd.Args.Domain, "", "", "CNAME")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(cname)
	}

	prn("Fetching MX records…")
	mx, err = getRecords(cmd.Args.Domain, "", "", "MX")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(mx)
	}

	prn("Fetching NS records…")
	ns, err = getRecords(cmd.Args.Domain, "", "", "NS")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(ns)
	}

	prn("Fetching TXT records…")
	txt, err = getRecords(cmd.Args.Domain, "", "", "TXT")
	if err != nil {
		pr("none found.")
	} else {
		pr("done.")
		sort.Sort(txt)
	}

	prn("Fetching SRV records…")
	srv, err = getRecords(cmd.Args.Domain, "", "", "SRV")
	if err != nil {
		pr("done.")
		pr("none found.")
	} else {
		sort.Sort(srv)
	}

	pr("Done: %v", everything)
	if cmd.Compact {
		everything = append(everything, a...)
		everything = append(everything, aaaa...)
		everything = append(everything, cname...)
		everything = append(everything, mx...)
		everything = append(everything, ns...)
		everything = append(everything, txt...)
		everything = append(everything, srv...)
		pr("%v", everything)
	} else {
		cc := columnize.DefaultConfig()
		cc.Delim = "\t"
		cc.Glue = "  "
		var s []string
		var res string
		base1 := []string{fmt.Sprintln("Host\tAddress\tType\tTTL\tStatus")}
		base2 := []string{fmt.Sprintln("Host\tAddress\tType\tTTL\tStatus\tPriority")}

		if len(a) > 0 {
			for _, r := range a {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("A records:\n%s\n", res)
		}

		if len(aaaa) > 0 {
			for _, r := range aaaa {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("AAAA records:\n%s\n", res)
		}

		if len(cname) > 0 {
			for _, r := range cname {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("CNAME records:\n%s\n", res)
		}

		if len(mx) > 0 {
			for _, r := range mx {
				s = append(base2, fmt.Sprintf("%s\t%s\t%d\t%s\t%d", r.Host, r.Value, r.TTL, r.Status, r.Priority))
			}
			res = columnize.Format(s, cc)
			pr("MX records:\n%s\n", res)
		}

		if len(ns) > 0 {
			for _, r := range ns {
				s = append(base2, fmt.Sprintf("%s\t%s\t%d\t%s\t%d", r.Host, r.Value, r.TTL, r.Status, r.Priority))
			}
			res = columnize.Format(s, cc)
			pr("NS records:\n%s\n", res)
		}

		if len(txt) > 0 {
			for _, r := range txt {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("TXT records:\n%s\n", res)
		}

		if len(txt) > 0 {
			for _, r := range txt {
				s = append(base1, fmt.Sprintf("%s\t%s\t%d\t%s", r.Host, r.Value, r.TTL, r.Status))
			}
			res = columnize.Format(s, cc)
			pr("TXT records:\n%s\n", res)
		}
	}
	return nil
}

// DNSGetSRVCmd arguments.
type DNSGetSRVCmd struct {
	Args DNSGetArgs `positional-args:"true"`
}

// Execute SRV record fetch.
func (cmd *DNSGetSRVCmd) Execute(args []string) error {
	srv, err := getRecords(cmd.Args.Domain, cmd.Args.Value, cmd.Args.Host, "SRV")
	if err != nil {
		return err
	}

	if len(srv) == 0 {
		return errors.New("no records found")
	}

	s := []string{fmt.Sprintln("Service name\tHost name\tPort\tWeight\tStatus")}
	for _, r := range srv {
		s = append(s, fmt.Sprintf("%s\t%s\t%d\t%d\t%s", r.Host, r.Value, r.Port, r.Weight, r.Status))
	}

	cc := columnize.DefaultConfig()
	cc.Delim = "\t"
	cc.Glue = "  "
	res := columnize.Format(s, cc)
	pr(res)
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package pprof

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"testing"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
)

type suite struct {
}

var _ = gc.Suite(&suite{})

func TestSuite(t *testing.T) {
	if runtime.GOOS != "linux" {
		t.Skipf("skipping pprof tests, %q not supported", runtime.GOOS)
	}
	gc.TestingT(t)
}

func (s *suite) TestFilename(c *gc.C) {
	got := Filename
	want := fmt.Sprintf("pprof.pprof.test.%d", os.Getpid())
	c.Assert(got, gc.Equals, want)
}

func (s *suite) TestPprofStartReturnsNonNilShutdownFn(c *gc.C) {
	stop := Start(filepath.Join(c.MkDir(), Filename))
	c.Assert(stop, gc.NotNil)
	defer stop()
}

func (s *suite) TestPprofStart(c *gc.C) {
	path := filepath.Join(c.MkDir(), Filename)
	_, err := os.Stat(path)
	c.Assert(os.IsNotExist(err), jc.IsTrue)

	stop := Start(path)
	_, err = os.Stat(path)
	c.Assert(err, gc.IsNil)

	err = stop()
	c.Assert(err, gc.IsNil)
	_, err = os.Stat(path)
	c.Assert(os.IsNotExist(err), jc.IsTrue)
}

func (s *suite) TestPprofStartWithExistingSocketFile(c *gc.C) {
	path := filepath.Join(c.MkDir(), Filename)
	w, err := os.Create(path)
	c.Assert(err, gc.IsNil)

	w.Write([]byte("not a socket"))
	err = w.Close() // can ignore error from w.Write
	c.Assert(err, gc.IsNil)

	stop := Start(path)
	defer stop()
	fi, err := os.Stat(path)
	c.Assert(err, gc.IsNil)
	c.Assert(fi.Mode()&os.ModeSocket != 0, jc.IsTrue)
}

type pprofSuite struct {
	stop func() error
	path string
}

var _ = gc.Suite(&pprofSuite{})

func (s *pprofSuite) SetUpSuite(c *gc.C) {
	s.path = filepath.Join(c.MkDir(), Filename)
	s.stop = Start(s.path)
}

func (s *pprofSuite) TearDownSuite(c *gc.C) {
	s.stop()
}

func (s *pprofSuite) call(c *gc.C, url string) []byte {
	conn, err := net.Dial("unix", s.path)
	c.Assert(err, gc.IsNil)
	defer conn.Close()

	_, err = fmt.Fprintf(conn, "GET %s HTTP/1.0\r\n\r\n", url)
	c.Assert(err, gc.IsNil)

	buf, err := ioutil.ReadAll(conn)
	c.Assert(err, gc.IsNil)
	return buf
}

func (s *pprofSuite) TestCmdLine(c *gc.C) {
	buf := s.call(c, "/debug/pprof/cmdline")
	c.Assert(buf, gc.NotNil)
	matches(c, buf, ".*github.com/juju/juju/cmd/pprof/_test/pprof.test")
}

func (s *pprofSuite) TestGoroutineProfile(c *gc.C) {
	buf := s.call(c, "/debug/pprof/goroutine")
	c.Assert(buf, gc.NotNil)
	matches(c, buf, `^goroutine profile: total \d+`)
}

// matches fails if regex is not found in the contents of b.
// b is expected to be the response from the pprof http server, and will
// contain some HTTP preamble that should be ignored.
func matches(c *gc.C, b []byte, regex string) {
	re, err := regexp.Compile(regex)
	c.Assert(err, gc.IsNil)
	r := bytes.NewReader(b)
	sc := bufio.NewScanner(r)
	for sc.Scan() {
		if re.MatchString(sc.Text()) {
			return
		}
	}
	c.Fatalf("%q did not match regex %q", string(b), regex)
}
<commit_msg>Check all stop() errors<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package pprof

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"testing"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
)

type suite struct {
}

var _ = gc.Suite(&suite{})

func TestSuite(t *testing.T) {
	if runtime.GOOS != "linux" {
		t.Skipf("skipping pprof tests, %q not supported", runtime.GOOS)
	}
	gc.TestingT(t)
}

func (s *suite) TestFilename(c *gc.C) {
	got := Filename
	want := fmt.Sprintf("pprof.pprof.test.%d", os.Getpid())
	c.Assert(got, gc.Equals, want)
}

func (s *suite) TestPprofStartReturnsNonNilShutdownFn(c *gc.C) {
	stop := Start(filepath.Join(c.MkDir(), Filename))
	c.Assert(stop, gc.NotNil)
	c.Assert(stop(), jc.ErrorIsNil)
}

func (s *suite) TestPprofStart(c *gc.C) {
	path := filepath.Join(c.MkDir(), Filename)
	_, err := os.Stat(path)
	c.Assert(os.IsNotExist(err), jc.IsTrue)

	stop := Start(path)
	_, err = os.Stat(path)
	c.Assert(err, jc.ErrorIsNil)

	err = stop()
	c.Assert(err, jc.ErrorIsNil)
	_, err = os.Stat(path)
	c.Assert(os.IsNotExist(err), jc.IsTrue)
}

func (s *suite) TestPprofStartWithExistingSocketFile(c *gc.C) {
	path := filepath.Join(c.MkDir(), Filename)
	w, err := os.Create(path)
	c.Assert(err, jc.ErrorIsNil)

	w.Write([]byte("not a socket"))
	err = w.Close() // can ignore error from w.Write
	c.Assert(err, jc.ErrorIsNil)

	stop := Start(path)
	defer stop()
	fi, err := os.Stat(path)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(fi.Mode()&os.ModeSocket != 0, jc.IsTrue)
}

type pprofSuite struct {
	stop func() error
	path string
}

var _ = gc.Suite(&pprofSuite{})

func (s *pprofSuite) SetUpSuite(c *gc.C) {
	s.path = filepath.Join(c.MkDir(), Filename)
	s.stop = Start(s.path)
}

func (s *pprofSuite) TearDownSuite(c *gc.C) {
	err := s.stop()
	c.Assert(err, jc.ErrorIsNil)
}

func (s *pprofSuite) call(c *gc.C, url string) []byte {
	conn, err := net.Dial("unix", s.path)
	c.Assert(err, jc.ErrorIsNil)
	defer conn.Close()

	_, err = fmt.Fprintf(conn, "GET %s HTTP/1.0\r\n\r\n", url)
	c.Assert(err, jc.ErrorIsNil)

	buf, err := ioutil.ReadAll(conn)
	c.Assert(err, jc.ErrorIsNil)
	return buf
}

func (s *pprofSuite) TestCmdLine(c *gc.C) {
	buf := s.call(c, "/debug/pprof/cmdline")
	c.Assert(buf, gc.NotNil)
	matches(c, buf, ".*github.com/juju/juju/cmd/pprof/_test/pprof.test")
}

func (s *pprofSuite) TestGoroutineProfile(c *gc.C) {
	buf := s.call(c, "/debug/pprof/goroutine")
	c.Assert(buf, gc.NotNil)
	matches(c, buf, `^goroutine profile: total \d+`)
}

// matches fails if regex is not found in the contents of b.
// b is expected to be the response from the pprof http server, and will
// contain some HTTP preamble that should be ignored.
func matches(c *gc.C, b []byte, regex string) {
	re, err := regexp.Compile(regex)
	c.Assert(err, jc.ErrorIsNil)
	r := bytes.NewReader(b)
	sc := bufio.NewScanner(r)
	for sc.Scan() {
		if re.MatchString(sc.Text()) {
			return
		}
	}
	c.Fatalf("%q did not match regex %q", string(b), regex)
}
<|endoftext|>


<commit_before>package backend

import (
	"context"
	"fmt"
	"path/filepath"

	"github.com/containerd/containerd"
	"github.com/containerd/go-cni"
	"github.com/opencontainers/runtime-spec/specs-go"
)

//go:generate go run github.com/maxbrunsfeld/counterfeiter/v6 github.com/containerd/go-cni.CNI

// CNINetworkConfig provides configuration for CNINetwork to override the
// defaults.
//
type CNINetworkConfig struct {
	// BridgeName is the name that the bridge set up in the current network
	// namespace to connect the veth's to.
	//
	BridgeName string

	// NetworkName is the virtual name used to identify the managed network.
	//
	NetworkName string

	// Subnet is the subnet (in CIDR notation) which the veths should be
	// added to.
	//
	Subnet string
}

const (
	// fileStoreWorkDir is a default directory used for storing
	// container-related files
	//
	fileStoreWorkDir = "/tmp"

	// binariesDir corresponds to the directory where CNI plugins have their
	// binaries in.
	//
	binariesDir = "/usr/local/concourse/bin"
)

var (
	// defaultNameServers is the default set of nameservers used.
	//
	defaultNameServers = []string{"8.8.8.8"}

	// defaultCNINetworkConfig is the default configuration for the CNI network
	// created to put concourse containers into.
	//
	defaultCNINetworkConfig = CNINetworkConfig{
		BridgeName:  "concourse0",
		NetworkName: "concourse",
		Subnet:      "10.80.0.0/16",
	}
)

func (c CNINetworkConfig) ToJSON() string {
	const networksConfListFormat = `{
  "cniVersion": "0.4.0",
  "name": "%s",
  "plugins": [
    {
      "type": "bridge",
      "bridge": "%s",
      "isGateway": true,
      "ipMasq": true,
      "ipam": {
        "type": "host-local",
        "subnet": "%s",
        "routes": [
          {
            "dst": "0.0.0.0/0"
          }
        ]
      }
    },
    {
      "type": "firewall"
    }
  ]
}`

	return fmt.Sprintf(networksConfListFormat,
		c.NetworkName, c.BridgeName, c.Subnet,
	)
}

// CNINetworkOpt defines a functional option that when applied, modifies the
// configuration of a CNINetwork.
//
type CNINetworkOpt func(n *cniNetwork)

// WithCNIBinariesDir is the directory where the binaries necessary for setting
// up the network live.
//
func WithCNIBinariesDir(dir string) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.binariesDir = dir
	}
}

// WithNameServers sets the set of nameservers to be configured for the
// /etc/resolv.conf inside the containers.
//
func WithNameServers(nameservers []string) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.nameServers = nameservers
	}
}

// WithCNIClient is an implementor of the CNI interface for reaching out to CNI
// plugins.
//
func WithCNIClient(c cni.CNI) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.client = c
	}
}

// WithCNINetworkConfig provides a custom CNINetworkConfig to be used by the CNI
// client at startup time.
//
func WithCNINetworkConfig(c CNINetworkConfig) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.config = c
	}
}

// WithCNIFileStore changes the default FileStore used to store files that
// belong to network configurations for containers.
//
func WithCNIFileStore(f FileStore) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.store = f
	}
}

type cniNetwork struct {
	client      cni.CNI
	store       FileStore
	config      CNINetworkConfig
	nameServers []string
	binariesDir string
}

var _ Network = (*cniNetwork)(nil)

func NewCNINetwork(opts ...CNINetworkOpt) (*cniNetwork, error) {
	var err error

	n := &cniNetwork{
		binariesDir: binariesDir,
		config:      defaultCNINetworkConfig,
		nameServers: defaultNameServers,
	}

	for _, opt := range opts {
		opt(n)
	}

	if n.store == nil {
		n.store = NewFileStore(fileStoreWorkDir)
	}

	if n.client == nil {
		n.client, err = cni.New(cni.WithPluginDir([]string{n.binariesDir}))
		if err != nil {
			return nil, fmt.Errorf("cni init: %w", err)
		}

		err = n.client.Load(
			cni.WithConfListBytes([]byte(n.config.ToJSON())),
			cni.WithLoNetwork,
		)
		if err != nil {
			return nil, fmt.Errorf("cni configuration loading: %w", err)
		}
	}

	return n, nil
}

func (n cniNetwork) SetupMounts(handle string) ([]specs.Mount, error) {
	if handle == "" {
		return nil, ErrInvalidInput("empty handle")
	}

	etcHosts, err := n.store.Create(
		filepath.Join(handle, "/hosts"),
		[]byte("127.0.0.1 localhost"),
	)
	if err != nil {
		return nil, fmt.Errorf("creating /etc/hosts: %w", err)
	}

	resolvConf, err := n.store.Create(
		filepath.Join(handle, "/resolv.conf"),
		n.generateResolvConfContents(),
	)
	if err != nil {
		return nil, fmt.Errorf("creating /etc/resolv.conf: %w", err)
	}

	return []specs.Mount{
		{
			Destination: "/etc/hosts",
			Type:        "bind",
			Source:      etcHosts,
			Options:     []string{"bind", "rw"},
		}, {
			Destination: "/etc/resolv.conf",
			Type:        "bind",
			Source:      resolvConf,
			Options:     []string{"bind", "rw"},
		},
	}, nil
}

func (n cniNetwork) generateResolvConfContents() []byte {
	contents := ""
	for _, n := range n.nameServers {
		contents = contents + fmt.Sprintf("nameserver %s\n", n)
	}

	return []byte(contents)
}

func (n cniNetwork) Add(ctx context.Context, task containerd.Task) error {
	if task == nil {
		return ErrInvalidInput("nil task")
	}

	id, netns := netId(task), netNsPath(task)

	_, err := n.client.Setup(ctx, id, netns)
	if err != nil {
		return fmt.Errorf("cni net setup: %w", err)
	}

	return nil
}

func (n cniNetwork) Remove(ctx context.Context, task containerd.Task) error {
	if task == nil {
		return ErrInvalidInput("nil task")
	}

	id, netns := netId(task), netNsPath(task)

	err := n.client.Remove(ctx, id, netns)
	if err != nil {
		return fmt.Errorf("cni net teardown: %w", err)
	}

	return nil
}

func netId(task containerd.Task) string {
	return task.ID()
}

func netNsPath(task containerd.Task) string {
	return fmt.Sprintf("/proc/%d/ns/net", task.Pid())
}
<commit_msg>worker: remove unecessary formatting<commit_after>package backend

import (
	"context"
	"fmt"
	"path/filepath"

	"github.com/containerd/containerd"
	"github.com/containerd/go-cni"
	"github.com/opencontainers/runtime-spec/specs-go"
)

//go:generate go run github.com/maxbrunsfeld/counterfeiter/v6 github.com/containerd/go-cni.CNI

// CNINetworkConfig provides configuration for CNINetwork to override the
// defaults.
//
type CNINetworkConfig struct {
	// BridgeName is the name that the bridge set up in the current network
	// namespace to connect the veth's to.
	//
	BridgeName string

	// NetworkName is the virtual name used to identify the managed network.
	//
	NetworkName string

	// Subnet is the subnet (in CIDR notation) which the veths should be
	// added to.
	//
	Subnet string
}

const (
	// fileStoreWorkDir is a default directory used for storing
	// container-related files
	//
	fileStoreWorkDir = "/tmp"

	// binariesDir corresponds to the directory where CNI plugins have their
	// binaries in.
	//
	binariesDir = "/usr/local/concourse/bin"
)

var (
	// defaultNameServers is the default set of nameservers used.
	//
	defaultNameServers = []string{"8.8.8.8"}

	// defaultCNINetworkConfig is the default configuration for the CNI network
	// created to put concourse containers into.
	//
	defaultCNINetworkConfig = CNINetworkConfig{
		BridgeName:  "concourse0",
		NetworkName: "concourse",
		Subnet:      "10.80.0.0/16",
	}
)

func (c CNINetworkConfig) ToJSON() string {
	const networksConfListFormat = `{
  "cniVersion": "0.4.0",
  "name": "%s",
  "plugins": [
    {
      "type": "bridge",
      "bridge": "%s",
      "isGateway": true,
      "ipMasq": true,
      "ipam": {
        "type": "host-local",
        "subnet": "%s",
        "routes": [
          {
            "dst": "0.0.0.0/0"
          }
        ]
      }
    },
    {
      "type": "firewall"
    }
  ]
}`

	return fmt.Sprintf(networksConfListFormat,
		c.NetworkName, c.BridgeName, c.Subnet,
	)
}

// CNINetworkOpt defines a functional option that when applied, modifies the
// configuration of a CNINetwork.
//
type CNINetworkOpt func(n *cniNetwork)

// WithCNIBinariesDir is the directory where the binaries necessary for setting
// up the network live.
//
func WithCNIBinariesDir(dir string) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.binariesDir = dir
	}
}

// WithNameServers sets the set of nameservers to be configured for the
// /etc/resolv.conf inside the containers.
//
func WithNameServers(nameservers []string) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.nameServers = nameservers
	}
}

// WithCNIClient is an implementor of the CNI interface for reaching out to CNI
// plugins.
//
func WithCNIClient(c cni.CNI) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.client = c
	}
}

// WithCNINetworkConfig provides a custom CNINetworkConfig to be used by the CNI
// client at startup time.
//
func WithCNINetworkConfig(c CNINetworkConfig) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.config = c
	}
}

// WithCNIFileStore changes the default FileStore used to store files that
// belong to network configurations for containers.
//
func WithCNIFileStore(f FileStore) CNINetworkOpt {
	return func(n *cniNetwork) {
		n.store = f
	}
}

type cniNetwork struct {
	client      cni.CNI
	store       FileStore
	config      CNINetworkConfig
	nameServers []string
	binariesDir string
}

var _ Network = (*cniNetwork)(nil)

func NewCNINetwork(opts ...CNINetworkOpt) (*cniNetwork, error) {
	var err error

	n := &cniNetwork{
		binariesDir: binariesDir,
		config:      defaultCNINetworkConfig,
		nameServers: defaultNameServers,
	}

	for _, opt := range opts {
		opt(n)
	}

	if n.store == nil {
		n.store = NewFileStore(fileStoreWorkDir)
	}

	if n.client == nil {
		n.client, err = cni.New(cni.WithPluginDir([]string{n.binariesDir}))
		if err != nil {
			return nil, fmt.Errorf("cni init: %w", err)
		}

		err = n.client.Load(
			cni.WithConfListBytes([]byte(n.config.ToJSON())),
			cni.WithLoNetwork,
		)
		if err != nil {
			return nil, fmt.Errorf("cni configuration loading: %w", err)
		}
	}

	return n, nil
}

func (n cniNetwork) SetupMounts(handle string) ([]specs.Mount, error) {
	if handle == "" {
		return nil, ErrInvalidInput("empty handle")
	}

	etcHosts, err := n.store.Create(
		filepath.Join(handle, "/hosts"),
		[]byte("127.0.0.1 localhost"),
	)
	if err != nil {
		return nil, fmt.Errorf("creating /etc/hosts: %w", err)
	}

	resolvConf, err := n.store.Create(
		filepath.Join(handle, "/resolv.conf"),
		n.generateResolvConfContents(),
	)
	if err != nil {
		return nil, fmt.Errorf("creating /etc/resolv.conf: %w", err)
	}

	return []specs.Mount{
		{
			Destination: "/etc/hosts",
			Type:        "bind",
			Source:      etcHosts,
			Options:     []string{"bind", "rw"},
		}, {
			Destination: "/etc/resolv.conf",
			Type:        "bind",
			Source:      resolvConf,
			Options:     []string{"bind", "rw"},
		},
	}, nil
}

func (n cniNetwork) generateResolvConfContents() []byte {
	contents := ""
	for _, n := range n.nameServers {
		contents = contents + "nameserver " + n + "\n"
	}

	return []byte(contents)
}

func (n cniNetwork) Add(ctx context.Context, task containerd.Task) error {
	if task == nil {
		return ErrInvalidInput("nil task")
	}

	id, netns := netId(task), netNsPath(task)

	_, err := n.client.Setup(ctx, id, netns)
	if err != nil {
		return fmt.Errorf("cni net setup: %w", err)
	}

	return nil
}

func (n cniNetwork) Remove(ctx context.Context, task containerd.Task) error {
	if task == nil {
		return ErrInvalidInput("nil task")
	}

	id, netns := netId(task), netNsPath(task)

	err := n.client.Remove(ctx, id, netns)
	if err != nil {
		return fmt.Errorf("cni net teardown: %w", err)
	}

	return nil
}

func netId(task containerd.Task) string {
	return task.ID()
}

func netNsPath(task containerd.Task) string {
	return fmt.Sprintf("/proc/%d/ns/net", task.Pid())
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"reflect"
	"strings"

	. "github.com/tendermint/go-common"
	"github.com/tendermint/go-wire"
	"github.com/tendermint/tmsp/types"

	"github.com/codegangsta/cli"
)

// connection is a global variable so it can be reused by the console
var conn net.Conn

func main() {
	app := cli.NewApp()
	app.Name = "tmsp-cli"
	app.Usage = "tmsp-cli [command] [args...]"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "address",
			Value: "tcp://127.0.0.1:46658",
			Usage: "address of application socket",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "batch",
			Usage: "Run a batch of tmsp commands against an application",
			Action: func(c *cli.Context) {
				cmdBatch(app, c)
			},
		},
		{
			Name:  "console",
			Usage: "Start an interactive tmsp console for multiple commands",
			Action: func(c *cli.Context) {
				cmdConsole(app, c)
			},
		},
		{
			Name:  "echo",
			Usage: "Have the application echo a message",
			Action: func(c *cli.Context) {
				cmdEcho(c)
			},
		},
		{
			Name:  "info",
			Usage: "Get some info about the application",
			Action: func(c *cli.Context) {
				cmdInfo(c)
			},
		},
		{
			Name:  "set_option",
			Usage: "Set an option on the application",
			Action: func(c *cli.Context) {
				cmdSetOption(c)
			},
		},
		{
			Name:  "append_tx",
			Usage: "Append a new tx to application",
			Action: func(c *cli.Context) {
				cmdAppendTx(c)
			},
		},
		{
			Name:  "check_tx",
			Usage: "Validate a tx",
			Action: func(c *cli.Context) {
				cmdCheckTx(c)
			},
		},
		{
			Name:  "get_hash",
			Usage: "Get application Merkle root hash",
			Action: func(c *cli.Context) {
				cmdGetHash(c)
			},
		},
		{
			Name:  "query",
			Usage: "Query application state",
			Action: func(c *cli.Context) {
				cmdQuery(c)
			},
		},
	}
	app.Before = before
	app.Run(os.Args)

}

func before(c *cli.Context) error {
	if conn == nil {
		var err error
		conn, err = Connect(c.GlobalString("address"))
		if err != nil {
			Exit(err.Error())
		}
	}
	return nil
}

//--------------------------------------------------------------------------------

func cmdBatch(app *cli.App, c *cli.Context) {
	bufReader := bufio.NewReader(os.Stdin)
	for {
		line, more, err := bufReader.ReadLine()
		if more {
			Exit("input line is too long")
		} else if err == io.EOF {
			break
		} else if len(line) == 0 {
			continue
		} else if err != nil {
			Exit(err.Error())
		}
		args := []string{"tmsp"}
		args = append(args, strings.Split(string(line), " ")...)
		app.Run(args)
	}
}

func cmdConsole(app *cli.App, c *cli.Context) {
	for {
		fmt.Printf("\n> ")
		bufReader := bufio.NewReader(os.Stdin)
		line, more, err := bufReader.ReadLine()
		if more {
			Exit("input is too long")
		} else if err != nil {
			Exit(err.Error())
		}

		args := []string{"tmsp"}
		args = append(args, strings.Split(string(line), " ")...)
		app.Run(args)
	}
}

// Have the application echo a message
func cmdEcho(c *cli.Context) {
	args := c.Args()
	if len(args) != 1 {
		Exit("echo takes 1 argument")
	}
	res, err := makeRequest(conn, types.RequestEcho{args[0]})
	if err != nil {
		Exit(err.Error())
	}
	fmt.Println("->", res)
}

// Get some info from the application
func cmdInfo(c *cli.Context) {
	res, err := makeRequest(conn, types.RequestInfo{})
	if err != nil {
		Exit(err.Error())
	}
	fmt.Println("->", res)
}

// Set an option on the application
func cmdSetOption(c *cli.Context) {
	args := c.Args()
	if len(args) != 2 {
		Exit("set_option takes 2 arguments (key, value)")
	}
	_, err := makeRequest(conn, types.RequestSetOption{args[0], args[1]})
	if err != nil {
		Exit(err.Error())
	}
	fmt.Println("->", Fmt("%s=%s", args[0], args[1]))
}

// Append a new tx to application
func cmdAppendTx(c *cli.Context) {
	args := c.Args()
	if len(args) != 1 {
		Exit("append_tx takes 1 argument")
	}
	txString := args[0]
	tx := []byte(txString)
	if len(txString) > 2 && strings.HasPrefix(txString, "0x") {
		var err error
		tx, err = hex.DecodeString(txString[2:])
		if err != nil {
			Exit(err.Error())
		}
	}

	res, err := makeRequest(conn, types.RequestAppendTx{tx})
	if err != nil {
		Exit(err.Error())
	}
	fmt.Println("->", res)
}

// Validate a tx
func cmdCheckTx(c *cli.Context) {
	args := c.Args()
	if len(args) != 1 {
		Exit("append_tx takes 1 argument")
	}
	txString := args[0]
	tx := []byte(txString)
	if len(txString) > 2 && strings.HasPrefix(txString, "0x") {
		var err error
		tx, err = hex.DecodeString(txString[2:])
		if err != nil {
			Exit(err.Error())
		}
	}

	res, err := makeRequest(conn, types.RequestCheckTx{tx})
	if err != nil {
		Exit(err.Error())
	}
	fmt.Println("->", res)
}

// Get application Merkle root hash
func cmdGetHash(c *cli.Context) {
	res, err := makeRequest(conn, types.RequestGetHash{})
	if err != nil {
		Exit(err.Error())
	}
	fmt.Printf("%X\n", res.(types.ResponseGetHash).Hash)
}

// Query application state
func cmdQuery(c *cli.Context) {
	args := c.Args()
	if len(args) != 1 {
		Exit("append_tx takes 1 argument")
	}
	queryString := args[0]
	query := []byte(queryString)
	if len(queryString) > 2 && strings.HasPrefix(queryString, "0x") {
		var err error
		query, err = hex.DecodeString(queryString[2:])
		if err != nil {
			Exit(err.Error())
		}
	}

	res, err := makeRequest(conn, types.RequestQuery{query})
	if err != nil {
		Exit(err.Error())
	}
	fmt.Println("->", res)
}

//--------------------------------------------------------------------------------

func makeRequest(conn net.Conn, req types.Request) (types.Response, error) {
	var n int
	var err error

	// Write desired request
	wire.WriteBinaryLengthPrefixed(struct{ types.Request }{req}, conn, &n, &err)
	if err != nil {
		return nil, err
	}

	// Write flush request
	wire.WriteBinaryLengthPrefixed(struct{ types.Request }{types.RequestFlush{}}, conn, &n, &err)
	if err != nil {
		return nil, err
	}

	// Read desired response
	var res types.Response
	wire.ReadBinaryPtrLengthPrefixed(&res, conn, 0, &n, &err)
	if err != nil {
		return nil, err
	}

	// Read flush response
	var resFlush types.Response
	wire.ReadBinaryPtrLengthPrefixed(&resFlush, conn, 0, &n, &err)
	if err != nil {
		return nil, err
	}
	if _, ok := resFlush.(types.ResponseFlush); !ok {
		return nil, errors.New(Fmt("Expected types.ResponseFlush but got %v instead", reflect.TypeOf(resFlush)))
	}

	return res, nil
}
<commit_msg>Do not exit upon error, for tmsp-cli console<commit_after>package main

import (
	"bufio"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"reflect"
	"strings"

	. "github.com/tendermint/go-common"
	"github.com/tendermint/go-wire"
	"github.com/tendermint/tmsp/types"

	"github.com/codegangsta/cli"
)

// connection is a global variable so it can be reused by the console
var conn net.Conn

func main() {
	app := cli.NewApp()
	app.Name = "tmsp-cli"
	app.Usage = "tmsp-cli [command] [args...]"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "address",
			Value: "tcp://127.0.0.1:46658",
			Usage: "address of application socket",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "batch",
			Usage: "Run a batch of tmsp commands against an application",
			Action: func(c *cli.Context) {
				cmdBatch(app, c)
			},
		},
		{
			Name:  "console",
			Usage: "Start an interactive tmsp console for multiple commands",
			Action: func(c *cli.Context) {
				cmdConsole(app, c)
			},
		},
		{
			Name:  "echo",
			Usage: "Have the application echo a message",
			Action: func(c *cli.Context) {
				cmdEcho(c)
			},
		},
		{
			Name:  "info",
			Usage: "Get some info about the application",
			Action: func(c *cli.Context) {
				cmdInfo(c)
			},
		},
		{
			Name:  "set_option",
			Usage: "Set an option on the application",
			Action: func(c *cli.Context) {
				cmdSetOption(c)
			},
		},
		{
			Name:  "append_tx",
			Usage: "Append a new tx to application",
			Action: func(c *cli.Context) {
				cmdAppendTx(c)
			},
		},
		{
			Name:  "check_tx",
			Usage: "Validate a tx",
			Action: func(c *cli.Context) {
				cmdCheckTx(c)
			},
		},
		{
			Name:  "get_hash",
			Usage: "Get application Merkle root hash",
			Action: func(c *cli.Context) {
				cmdGetHash(c)
			},
		},
		{
			Name:  "query",
			Usage: "Query application state",
			Action: func(c *cli.Context) {
				cmdQuery(c)
			},
		},
	}
	app.Before = before
	app.Run(os.Args)

}

func before(c *cli.Context) error {
	if conn == nil {
		var err error
		conn, err = Connect(c.GlobalString("address"))
		if err != nil {
			Exit(err.Error())
		}
	}
	return nil
}

//--------------------------------------------------------------------------------

func cmdBatch(app *cli.App, c *cli.Context) {
	bufReader := bufio.NewReader(os.Stdin)
	for {
		line, more, err := bufReader.ReadLine()
		if more {
			fmt.Println("input line is too long")
			return
		} else if err == io.EOF {
			break
		} else if len(line) == 0 {
			continue
		} else if err != nil {
			fmt.Println(err.Error())
			return
		}
		args := []string{"tmsp"}
		args = append(args, strings.Split(string(line), " ")...)
		app.Run(args)
	}
}

func cmdConsole(app *cli.App, c *cli.Context) {
	for {
		fmt.Printf("\n> ")
		bufReader := bufio.NewReader(os.Stdin)
		line, more, err := bufReader.ReadLine()
		if more {
			fmt.Println("input is too long")
			return
		} else if err != nil {
			fmt.Println(err.Error())
			return
		}

		args := []string{"tmsp"}
		args = append(args, strings.Split(string(line), " ")...)
		app.Run(args)
	}
}

// Have the application echo a message
func cmdEcho(c *cli.Context) {
	args := c.Args()
	if len(args) != 1 {
		fmt.Println("echo takes 1 argument")
		return
	}
	res, err := makeRequest(conn, types.RequestEcho{args[0]})
	if err != nil {
		fmt.Println(err.Error())
		return
	}
	fmt.Println("->", res)
}

// Get some info from the application
func cmdInfo(c *cli.Context) {
	res, err := makeRequest(conn, types.RequestInfo{})
	if err != nil {
		fmt.Println(err.Error())
		return
	}
	fmt.Println("->", res)
}

// Set an option on the application
func cmdSetOption(c *cli.Context) {
	args := c.Args()
	if len(args) != 2 {
		fmt.Println("set_option takes 2 arguments (key, value)")
		return
	}
	_, err := makeRequest(conn, types.RequestSetOption{args[0], args[1]})
	if err != nil {
		fmt.Println(err.Error())
		return
	}
	fmt.Println("->", Fmt("%s=%s", args[0], args[1]))
}

// Append a new tx to application
func cmdAppendTx(c *cli.Context) {
	args := c.Args()
	if len(args) != 1 {
		fmt.Println("append_tx takes 1 argument")
		return
	}
	txString := args[0]
	tx := []byte(txString)
	if len(txString) > 2 && strings.HasPrefix(txString, "0x") {
		var err error
		tx, err = hex.DecodeString(txString[2:])
		if err != nil {
			fmt.Println(err.Error())
			return
		}
	}

	res, err := makeRequest(conn, types.RequestAppendTx{tx})
	if err != nil {
		fmt.Println(err.Error())
		return
	}
	fmt.Println("->", res)
}

// Validate a tx
func cmdCheckTx(c *cli.Context) {
	args := c.Args()
	if len(args) != 1 {
		fmt.Println("append_tx takes 1 argument")
		return
	}
	txString := args[0]
	tx := []byte(txString)
	if len(txString) > 2 && strings.HasPrefix(txString, "0x") {
		var err error
		tx, err = hex.DecodeString(txString[2:])
		if err != nil {
			fmt.Println(err.Error())
			return
		}
	}

	res, err := makeRequest(conn, types.RequestCheckTx{tx})
	if err != nil {
		fmt.Println(err.Error())
		return
	}
	fmt.Println("->", res)
}

// Get application Merkle root hash
func cmdGetHash(c *cli.Context) {
	res, err := makeRequest(conn, types.RequestGetHash{})
	if err != nil {
		fmt.Println(err.Error())
		return
	}
	fmt.Printf("%X\n", res.(types.ResponseGetHash).Hash)
}

// Query application state
func cmdQuery(c *cli.Context) {
	args := c.Args()
	if len(args) != 1 {
		fmt.Println("append_tx takes 1 argument")
		return
	}
	queryString := args[0]
	query := []byte(queryString)
	if len(queryString) > 2 && strings.HasPrefix(queryString, "0x") {
		var err error
		query, err = hex.DecodeString(queryString[2:])
		if err != nil {
			fmt.Println(err.Error())
			return
		}
	}

	res, err := makeRequest(conn, types.RequestQuery{query})
	if err != nil {
		fmt.Println(err.Error())
		return
	}
	fmt.Println("->", res)
}

//--------------------------------------------------------------------------------

func makeRequest(conn net.Conn, req types.Request) (types.Response, error) {
	var n int
	var err error

	// Write desired request
	wire.WriteBinaryLengthPrefixed(struct{ types.Request }{req}, conn, &n, &err)
	if err != nil {
		return nil, err
	}

	// Write flush request
	wire.WriteBinaryLengthPrefixed(struct{ types.Request }{types.RequestFlush{}}, conn, &n, &err)
	if err != nil {
		return nil, err
	}

	// Read desired response
	var res types.Response
	wire.ReadBinaryPtrLengthPrefixed(&res, conn, 0, &n, &err)
	if err != nil {
		return nil, err
	}

	// Read flush response
	var resFlush types.Response
	wire.ReadBinaryPtrLengthPrefixed(&resFlush, conn, 0, &n, &err)
	if err != nil {
		return nil, err
	}
	if _, ok := resFlush.(types.ResponseFlush); !ok {
		return nil, errors.New(Fmt("Expected types.ResponseFlush but got %v instead", reflect.TypeOf(resFlush)))
	}

	return res, nil
}
<|endoftext|>


<commit_before>// Package versions provides an implementation of an action on "versions".
package versions

import (
	"context"
	"fmt"
	"os"
	"regexp"
	"text/tabwriter"

	"github.com/actions-on-google/gactions/api/sdk"
	"github.com/actions-on-google/gactions/project"
	"github.com/actions-on-google/gactions/project/studio"
	"github.com/spf13/cobra"
)

var versionIDRegExp = regexp.MustCompile(`^projects/[^/]+/versions/(?P<versionID>[^/]+)$`)
var modifiedOnRegExp = regexp.MustCompile(`(?P<date>\d{4}-\d{2}-\d{2})+T+(?P<time>\d{2}:\d{2}:\d{2})(\.\d{6})+Z`)

// AddCommand adds the release-channels list sub-command to the passed in root command.
func AddCommand(ctx context.Context, root *cobra.Command, project project.Project) {
	versions := &cobra.Command{
		Use:   "versions",
		Short: "This command performs versions specified actions.",
		Long:  "This command performs versions specified actions",
		Args:  cobra.MinimumNArgs(1),
	}
	list := &cobra.Command{
		Use:   "list",
		Short: "This command lists all versions and their metadata.",
		Long:  "This command lists all versions and their metadata.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			studioProj, ok := project.(studio.Studio)
			if !ok {
				return fmt.Errorf("can not convert %T to %T", project, studio.Studio{})
			}
			pid, err := cmd.Flags().GetString("project-id")
			if err != nil {
				return err
			}
			if err := (&studioProj).SetProjectID(pid); err != nil {
				return err
			}
			res, err := sdk.ListVersionsJSON(ctx, studioProj)
			if err != nil {
				return err
			}
			return printVersions(res)
		},
	}
	list.Flags().String("project-id", "", "List versions of the project specified by the ID. The value provided in this flag will overwrite the value from settings file, if present.")
	versions.AddCommand(list)
	root.AddCommand(versions)
}

func printVersions(versions []project.Version) error {
	w := new(tabwriter.Writer)
	// Format in tab-separated columns with a tab stop of 8.
	w.Init(os.Stdout, 20, 8, 1, '\t', 0)
	fmt.Fprintln(w, "Version\tStatus\tLast Modified By\tModified On\t")
	for _, version := range versions {
		fmt.Fprintf(w, "%v\t%v\t%v\t%v\t\n", versionID(version.ID), version.State.Message, version.LastModifiedBy, formatModifiedOn(version.ModifiedOn))
	}
	fmt.Fprintf(w, "To learn more about release channels, visit https://developers.google.com/assistant/actionssdk/reference/rest/Shared.Types/ReleaseChannel.")
	fmt.Fprintln(w)
	return w.Flush()
}

func versionID(version string) string {
	versionIDMatch := versionIDRegExp.FindStringSubmatch(version)
	if versionIDMatch == nil {
		return "N/A"
	}
	return versionIDMatch[versionIDRegExp.SubexpIndex("versionID")]
}

func formatModifiedOn(modifiedOn string) string {
	modifiedOnMatch := modifiedOnRegExp.FindStringSubmatch(modifiedOn)
	if modifiedOnMatch == nil {
		return "N/A"
	}

	return modifiedOnMatch[modifiedOnRegExp.SubexpIndex("date")] + " " + modifiedOnMatch[modifiedOnRegExp.SubexpIndex("time")]
}
<commit_msg>Update versions command description<commit_after>// Package versions provides an implementation of an action on "versions".
package versions

import (
	"context"
	"fmt"
	"os"
	"regexp"
	"text/tabwriter"

	"github.com/actions-on-google/gactions/api/sdk"
	"github.com/actions-on-google/gactions/project"
	"github.com/actions-on-google/gactions/project/studio"
	"github.com/spf13/cobra"
)

var versionIDRegExp = regexp.MustCompile(`^projects/[^/]+/versions/(?P<versionID>[^/]+)$`)
var modifiedOnRegExp = regexp.MustCompile(`(?P<date>\d{4}-\d{2}-\d{2})+T+(?P<time>\d{2}:\d{2}:\d{2})(\.\d{6})+Z`)

// AddCommand adds the release-channels list sub-command to the passed in root command.
func AddCommand(ctx context.Context, root *cobra.Command, project project.Project) {
	versions := &cobra.Command{
		Use:   "versions",
		Short: "This command performs versions specific actions.",
		Long:  "This command performs versions specific actions",
		Args:  cobra.MinimumNArgs(1),
	}
	list := &cobra.Command{
		Use:   "list",
		Short: "This command lists all versions and their metadata.",
		Long:  "This command lists all versions and their metadata.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			studioProj, ok := project.(studio.Studio)
			if !ok {
				return fmt.Errorf("can not convert %T to %T", project, studio.Studio{})
			}
			pid, err := cmd.Flags().GetString("project-id")
			if err != nil {
				return err
			}
			if err := (&studioProj).SetProjectID(pid); err != nil {
				return err
			}
			res, err := sdk.ListVersionsJSON(ctx, studioProj)
			if err != nil {
				return err
			}
			return printVersions(res)
		},
	}
	list.Flags().String("project-id", "", "List versions of the project specified by the ID. The value provided in this flag will overwrite the value from settings file, if present.")
	versions.AddCommand(list)
	root.AddCommand(versions)
}

func printVersions(versions []project.Version) error {
	w := new(tabwriter.Writer)
	// Format in tab-separated columns with a tab stop of 8.
	w.Init(os.Stdout, 20, 8, 1, '\t', 0)
	fmt.Fprintln(w, "Version\tStatus\tLast Modified By\tModified On\t")
	for _, version := range versions {
		fmt.Fprintf(w, "%v\t%v\t%v\t%v\t\n", versionID(version.ID), version.State.Message, version.LastModifiedBy, formatModifiedOn(version.ModifiedOn))
	}
	fmt.Fprintf(w, "To learn more about release channels, visit https://developers.google.com/assistant/actionssdk/reference/rest/Shared.Types/ReleaseChannel.")
	fmt.Fprintln(w)
	return w.Flush()
}

func versionID(version string) string {
	versionIDMatch := versionIDRegExp.FindStringSubmatch(version)
	if versionIDMatch == nil {
		return "N/A"
	}
	return versionIDMatch[versionIDRegExp.SubexpIndex("versionID")]
}

func formatModifiedOn(modifiedOn string) string {
	modifiedOnMatch := modifiedOnRegExp.FindStringSubmatch(modifiedOn)
	if modifiedOnMatch == nil {
		return "N/A"
	}

	return modifiedOnMatch[modifiedOnRegExp.SubexpIndex("date")] + " " + modifiedOnMatch[modifiedOnRegExp.SubexpIndex("time")]
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv4/client4"
	"github.com/insomniacslk/dhcp/netboot"
	"github.com/u-root/u-root/pkg/boot/stboot"
	"github.com/vishvananda/netlink"
)

var debug = func(string, ...interface{}) {}

var (
	dryRun  = flag.Bool("dryrun", false, "Do everything except booting the loaded kernel")
	doDebug = flag.Bool("d", false, "Print debug output")
)

const (
	eth                = "eth0"
	rootCACertPath     = "/root/LetsEncrypt_Authority_X3.pem"
	entropyAvail       = "/proc/sys/kernel/random/entropy_avail"
	interfaceUpTimeout = 10 * time.Second
)

var banner = `
  _____ _______   _____   ____   ____________
 / ____|__   __|  |  _ \ / __ \ / __ \__   __|
| (___    | |     | |_) | |  | | |  | | | |
 \___ \   | |     |  _ <| |  | | |  | | | |
 ____) |  | |     | |_) | |__| | |__| | | |
|_____/   |_|     |____/ \____/ \____/  |_|

`

var check = `
           //\\
OS is     //  \\
valid    //   //
        //   //
 //\\  //   //
//  \\//   //
\\        //
 \\      //
  \\    //
   \\__//
`

func main() {
	log.SetPrefix("stboot: ")

	flag.Parse()
	if *doDebug {
		debug = log.Printf
	}
	log.Print(banner)

	vars, err := stboot.FindHostVarsInInitramfs()
	if err != nil {
		log.Fatalf("Cant find Netvars at all: %v", err)
	}

	if *doDebug {
		str, _ := json.MarshalIndent(vars, "", "  ")
		log.Printf("Host variables: %s", str)
	}

	debug("Configuring network interfaces")

	if vars.HostIP != "" {
		err = configureStaticNetwork(vars, *doDebug)
	} else {
		err = configureDHCPNetwork()
	}

	if err != nil {
		log.Fatalf("Can not set up IO: %v", err)
	}

	ballPath := path.Join("root/", stboot.BallName)
	url, err := url.Parse(vars.BootstrapURL)
	if err != nil {
		log.Fatalf("Invalid bootstrap URL: %v", err)
	}
	url.Path = path.Join(url.Path, stboot.BallName)
	err = downloadFromHTTPS(url.String(), ballPath)
	if err != nil {
		log.Fatalf("Downloading bootball from %s failed: %v", url, err)
	}

	ball, err := stboot.BootBallFromArchie(ballPath)
	if err != nil {
		log.Fatal("Cannot open bootball")
	}

	// Just choose the first Bootconfig for now
	log.Printf("Pick the first boot configuration")
	var index = 0
	bc, err := ball.GetBootConfigByIndex(index)
	if err != nil {
		log.Fatalf("Cannot get boot configuration %d: %v", index, err)
	}

	if *doDebug {
		str, _ := json.MarshalIndent(*bc, "", "  ")
		log.Printf("Bootconfig: %s", str)
	}

	n, err := ball.VerifyBootconfigByName(bc.Name)
	if err != nil {
		log.Fatalf("Bootconfig %d seems to be not trustworthy: %v", index, err)
	}
	if n < vars.MinimalSignaturesMatch {
		log.Fatalf("Did not found enough valid signatures. %d valid, %d required", n, vars.MinimalSignaturesMatch)
	}
	log.Printf("Bootconfig '%s' passed verification", bc.Name)
	log.Print(check)

	if *dryRun {
		debug("Dryrun mode: will not boot")
		return
	}

	log.Println("Starting up new kernel.")

	if err := bc.Boot(); err != nil {
		log.Printf("Failed to boot kernel %s: %v", bc.Kernel, err)
	}
	// if we reach this point, no boot configuration succeeded
	log.Print("No boot configuration succeeded")

	return
}

func configureStaticNetwork(vars stboot.HostVars, doDebug bool) error {
	log.Printf("Setup network configuration with IP: " + vars.HostIP)
	addr, err := netlink.ParseAddr(vars.HostIP)
	if err != nil {
		return fmt.Errorf("Error parsing HostIP string to CIDR format address: %v", err)
	}

	iface, err := netlink.LinkByName(eth)
	if err = netlink.AddrAdd(iface, addr); err != nil {
		return fmt.Errorf("Error retrieving interface by name: %v", err)
	}

	if err = netlink.LinkSetUp(iface); err != nil {
		return fmt.Errorf("Error bringing up interface:%v with error: %v", eth, err)
	}

	gateway, err := netlink.ParseAddr(vars.DefaultGateway)
	if err != nil {
		return fmt.Errorf("Error parsing GatewayIP string to CIDR format address: %v", err)
	}

	r := &netlink.Route{LinkIndex: iface.Attrs().Index, Gw: gateway.IPNet.IP}
	if err = netlink.RouteAdd(r); err != nil {
		return fmt.Errorf("Error setting default gateway: %v", err)
	}

	return nil
}

func configureDHCPNetwork() error {

	log.Printf("Trying to configure network configuration dynamically...")
	attempts := 10
	var conversation []*dhcpv4.DHCPv4

	_, err := netboot.IfUp(eth, interfaceUpTimeout)
	if err != nil {
		return fmt.Errorf("Ifup with %s failed: %v", eth, err)
	}
	if attempts < 1 {
		attempts = 1
	}

	client := client4.NewClient()
	for attempt := 0; attempt < attempts; attempt++ {
		log.Printf("Attempt to get DHCP lease %d of %d for interface %s", attempt+1, attempts, eth)
		conversation, err = client.Exchange(eth)

		if err != nil && attempt < attempts {
			log.Printf("Error: %v", err)
			continue
		}
		break
	}

	if conversation[3] == nil {
		return fmt.Errorf("Gateway is null")
	}
	netbootConfig, err := netboot.GetNetConfFromPacketv4(conversation[3])

	if err != nil {
		return err
	}

	err = netboot.ConfigureInterface(eth, netbootConfig)

	if err != nil {
		return err
	}

	iface, err := netlink.LinkByName(eth)
	if err != nil {
		return fmt.Errorf("Error getting Link by Name: %v", err)
	}
	gateway, err := netlink.ParseAddr(netbootConfig.Routers[0].String() + "/24")
	if err != nil {
		return fmt.Errorf("Error parsing GatewayIP string to CIDR format address: %v", err)
	}
	r := &netlink.Route{LinkIndex: iface.Attrs().Index, Gw: gateway.IPNet.IP}
	if err = netlink.RouteAdd(r); err != nil {
		return fmt.Errorf("Error setting default gateway: %v", err)
	}

	return nil
}

func downloadFromHTTPS(url string, destination string) error {
	roots := x509.NewCertPool()
	if err := loadHTTPSCertificate(roots); err != nil {
		return fmt.Errorf("Failed to load root certificate: %v", err)
	}

	// setup https client
	client := http.Client{
		Transport: (&http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
				DualStack: true,
			}).DialContext,
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
			TLSClientConfig: (&tls.Config{
				RootCAs: roots,
			}),
		}),
	}

	// check available kernel entropy
	e, err := ioutil.ReadFile(entropyAvail)
	es := strings.TrimSpace(string(e))
	entr, err := strconv.Atoi(es)
	if err != nil {
		return fmt.Errorf("Cannot evaluate entropy, %v", err)
	}
	log.Printf("Available kernel entropy: %d", entr)
	if entr < 128 {
		log.Print("WARNING: low entropy!")
		log.Printf("%s : %d", entropyAvail, entr)
	}
	// get remote boot bundle
	log.Print("Downloading bootball ...")
	resp, err := client.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("non-200 HTTP status: %d", resp.StatusCode)
	}
	f, err := os.Create(destination)
	if err != nil {
		return fmt.Errorf("failed create boot config file: %v", err)
	}
	defer f.Close()

	_, err = io.Copy(f, resp.Body)
	if err != nil {
		return fmt.Errorf("failed to save bootball: %v", err)
	}

	return nil
}

// loadHTTPSCertificate loads the certificate needed
// for HTTPS and verifies it.
func loadHTTPSCertificate(roots *x509.CertPool) error {
	log.Printf("Load %s as CA certificate", rootCACertPath)
	rootCertBytes, err := ioutil.ReadFile(rootCACertPath)
	if err != nil {
		return err
	}
	rootCertPem, _ := pem.Decode(rootCertBytes)
	if rootCertPem.Type != "CERTIFICATE" {
		return fmt.Errorf("Failed decoding certificate: Certificate is of the wrong type. PEM Type is: %s", rootCertPem.Type)
	}
	ok := roots.AppendCertsFromPEM([]byte(rootCertBytes))
	if !ok {
		return fmt.Errorf("Error parsing CA root certificate")
	}
	return nil
}
<commit_msg>cmds/boot/stboot: fix dhcp usage<commit_after>package main

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv4/client4"
	"github.com/insomniacslk/dhcp/netboot"
	"github.com/u-root/u-root/pkg/boot/stboot"
	"github.com/vishvananda/netlink"
)

var debug = func(string, ...interface{}) {}

var (
	dryRun  = flag.Bool("dryrun", false, "Do everything except booting the loaded kernel")
	doDebug = flag.Bool("d", false, "Print debug output")
)

const (
	eth                = "eth0"
	rootCACertPath     = "/root/LetsEncrypt_Authority_X3.pem"
	entropyAvail       = "/proc/sys/kernel/random/entropy_avail"
	interfaceUpTimeout = 10 * time.Second
)

var banner = `
  _____ _______   _____   ____   ____________
 / ____|__   __|  |  _ \ / __ \ / __ \__   __|
| (___    | |     | |_) | |  | | |  | | | |
 \___ \   | |     |  _ <| |  | | |  | | | |
 ____) |  | |     | |_) | |__| | |__| | | |
|_____/   |_|     |____/ \____/ \____/  |_|

`

var check = `
           //\\
OS is     //  \\
valid    //   //
        //   //
 //\\  //   //
//  \\//   //
\\        //
 \\      //
  \\    //
   \\__//
`

func main() {
	log.SetPrefix("stboot: ")

	flag.Parse()
	if *doDebug {
		debug = log.Printf
	}
	log.Print(banner)

	vars, err := stboot.FindHostVarsInInitramfs()
	if err != nil {
		log.Fatalf("Cant find Netvars at all: %v", err)
	}

	if *doDebug {
		str, _ := json.MarshalIndent(vars, "", "  ")
		log.Printf("Host variables: %s", str)
	}

	debug("Configuring network interfaces")

	if vars.HostIP != "" {
		err = configureStaticNetwork(vars, *doDebug)
	} else {
		err = configureDHCPNetwork()
	}

	if err != nil {
		log.Fatalf("Can not set up IO: %v", err)
	}

	ballPath := path.Join("root/", stboot.BallName)
	url, err := url.Parse(vars.BootstrapURL)
	if err != nil {
		log.Fatalf("Invalid bootstrap URL: %v", err)
	}
	url.Path = path.Join(url.Path, stboot.BallName)
	err = downloadFromHTTPS(url.String(), ballPath)
	if err != nil {
		log.Fatalf("Downloading bootball from %s failed: %v", url, err)
	}

	ball, err := stboot.BootBallFromArchie(ballPath)
	if err != nil {
		log.Fatal("Cannot open bootball")
	}

	// Just choose the first Bootconfig for now
	log.Printf("Pick the first boot configuration")
	var index = 0
	bc, err := ball.GetBootConfigByIndex(index)
	if err != nil {
		log.Fatalf("Cannot get boot configuration %d: %v", index, err)
	}

	if *doDebug {
		str, _ := json.MarshalIndent(*bc, "", "  ")
		log.Printf("Bootconfig: %s", str)
	}

	n, err := ball.VerifyBootconfigByName(bc.Name)
	if err != nil {
		log.Fatalf("Bootconfig %d seems to be not trustworthy: %v", index, err)
	}
	if n < vars.MinimalSignaturesMatch {
		log.Fatalf("Did not found enough valid signatures. %d valid, %d required", n, vars.MinimalSignaturesMatch)
	}
	log.Printf("Bootconfig '%s' passed verification", bc.Name)
	log.Print(check)

	if *dryRun {
		debug("Dryrun mode: will not boot")
		return
	}

	log.Println("Starting up new kernel.")

	if err := bc.Boot(); err != nil {
		log.Printf("Failed to boot kernel %s: %v", bc.Kernel, err)
	}
	// if we reach this point, no boot configuration succeeded
	log.Print("No boot configuration succeeded")

	return
}

func configureStaticNetwork(vars stboot.HostVars, doDebug bool) error {
	log.Printf("Setup network configuration with IP: " + vars.HostIP)
	addr, err := netlink.ParseAddr(vars.HostIP)
	if err != nil {
		return fmt.Errorf("Error parsing HostIP string to CIDR format address: %v", err)
	}

	iface, err := netlink.LinkByName(eth)
	if err != nil {
		return fmt.Errorf("Error retrieving interface %s: %v", eth, err)
	}

	if err = netlink.AddrAdd(iface, addr); err != nil {
		return fmt.Errorf("Error adding address: %v", err)
	}

	if err = netlink.LinkSetUp(iface); err != nil {
		return fmt.Errorf("Error bringing up interface %s: %v", eth, err)
	}

	gateway, err := netlink.ParseAddr(vars.DefaultGateway)
	if err != nil {
		return fmt.Errorf("Error parsing GatewayIP string to CIDR format address: %v", err)
	}

	r := &netlink.Route{LinkIndex: iface.Attrs().Index, Gw: gateway.IPNet.IP}
	if err = netlink.RouteAdd(r); err != nil {
		return fmt.Errorf("Error setting default gateway: %v", err)
	}

	return nil
}

func configureDHCPNetwork() error {

	log.Printf("Trying to configure network configuration dynamically...")
	attempts := 10
	var conversation []*dhcpv4.DHCPv4

	_, err := netboot.IfUp(eth, interfaceUpTimeout)
	if err != nil {
		return fmt.Errorf("Error bringing up interface %s: %v", eth, err)
	}
	if attempts < 1 {
		attempts = 1
	}

	client := client4.NewClient()
	for attempt := 0; attempt < attempts; attempt++ {
		log.Printf("Attempt to get DHCP lease %d of %d for interface %s", attempt+1, attempts, eth)
		conversation, err = client.Exchange(eth)

		if err != nil && attempt < attempts {
			log.Printf("Error: %v", err)
			continue
		}
		break
	}

	if conversation[3] == nil {
		return fmt.Errorf("Gateway is null")
	}

	netbootConfig, err := netboot.GetNetConfFromPacketv4(conversation[3])
	if err != nil {
		return err
	}

	if *doDebug {
		str, _ := json.MarshalIndent(netbootConfig, "", "  ")
		log.Printf("Network configuration: %s", str)
	}

	return netboot.ConfigureInterface(eth, netbootConfig)
}

func downloadFromHTTPS(url string, destination string) error {
	roots := x509.NewCertPool()
	if err := loadHTTPSCertificate(roots); err != nil {
		return fmt.Errorf("Failed to load root certificate: %v", err)
	}

	// setup https client
	client := http.Client{
		Transport: (&http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
				DualStack: true,
			}).DialContext,
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
			TLSClientConfig: (&tls.Config{
				RootCAs: roots,
			}),
		}),
	}

	// check available kernel entropy
	e, err := ioutil.ReadFile(entropyAvail)
	es := strings.TrimSpace(string(e))
	entr, err := strconv.Atoi(es)
	if err != nil {
		return fmt.Errorf("Cannot evaluate entropy, %v", err)
	}
	log.Printf("Available kernel entropy: %d", entr)
	if entr < 128 {
		log.Print("WARNING: low entropy!")
		log.Printf("%s : %d", entropyAvail, entr)
	}
	// get remote boot bundle
	log.Print("Downloading bootball ...")
	resp, err := client.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("non-200 HTTP status: %d", resp.StatusCode)
	}
	f, err := os.Create(destination)
	if err != nil {
		return fmt.Errorf("failed create boot config file: %v", err)
	}
	defer f.Close()

	_, err = io.Copy(f, resp.Body)
	if err != nil {
		return fmt.Errorf("failed to save bootball: %v", err)
	}

	return nil
}

// loadHTTPSCertificate loads the certificate needed
// for HTTPS and verifies it.
func loadHTTPSCertificate(roots *x509.CertPool) error {
	log.Printf("Load %s as CA certificate", rootCACertPath)
	rootCertBytes, err := ioutil.ReadFile(rootCACertPath)
	if err != nil {
		return err
	}
	rootCertPem, _ := pem.Decode(rootCertBytes)
	if rootCertPem.Type != "CERTIFICATE" {
		return fmt.Errorf("Failed decoding certificate: Certificate is of the wrong type. PEM Type is: %s", rootCertPem.Type)
	}
	ok := roots.AppendCertsFromPEM([]byte(rootCertBytes))
	if !ok {
		return fmt.Errorf("Error parsing CA root certificate")
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
	"github.com/uber/zanzibar/codegen"
	"github.com/uber/zanzibar/runtime"
)

var configFile = flag.String("config", "", "the config file path")

const templateDir = "./codegen/templates/*.tmpl"

type stackTracer interface {
	StackTrace() errors.StackTrace
}

func checkError(err error, message string) {
	if err != nil {
		fmt.Printf("%s:\n%s\n", message, err)

		causeErr := errors.Cause(err)
		if causeErr, ok := causeErr.(stackTracer); ok {
			fmt.Printf("%+v \n", causeErr.StackTrace())
		}

		os.Exit(1)
	}
}

func main() {
	flag.Parse()
	if *configFile == "" {
		flag.Usage()
		os.Exit(1)
		return
	}

	configRoot := filepath.Dir(*configFile)
	config := zanzibar.NewStaticConfigOrDie([]string{
		*configFile,
	}, nil)

	configRoot, err := filepath.Abs(configRoot)
	checkError(
		err, fmt.Sprintf("can not get abs path of config dir %s", configRoot),
	)

	copyright, err := ioutil.ReadFile(filepath.Join(
		configRoot,
		config.MustGetString("copyrightHeader"),
	))
	if err != nil {
		// Default to an empty copyright for now
		copyright = []byte("")
	}

	packageHelper, err := codegen.NewPackageHelper(
		config.MustGetString("packageRoot"),
		configRoot,
		config.MustGetString("middlewareConfig"),
		config.MustGetString("thriftRootDir"),
		config.MustGetString("genCodePackage"),
		config.MustGetString("targetGenDir"),
		string(copyright),
	)
	checkError(
		err, fmt.Sprintf("Can't build package helper %s", configRoot),
	)

	moduleSystem, err := codegen.NewDefaultModuleSystem(packageHelper)
	checkError(
		err, fmt.Sprintf("Error creating module system %s", configRoot),
	)

	// TODO: All components should be generated by the module system
	fmt.Printf("Generating module system components:\n")
	_, err = moduleSystem.GenerateBuild(
		packageHelper.PackageRoot(),
		configRoot,
		packageHelper.CodeGenTargetPath(),
	)
	checkError(err, "Failed to generate module system components")
}
<commit_msg>adapt codegen runner to new config interface<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
	"github.com/uber/zanzibar/codegen"
	"github.com/uber/zanzibar/runtime"
)

var configFile = flag.String("config", "", "the config file path")

const templateDir = "./codegen/templates/*.tmpl"

type stackTracer interface {
	StackTrace() errors.StackTrace
}

func checkError(err error, message string) {
	if err != nil {
		fmt.Printf("%s:\n%s\n", message, err)

		causeErr := errors.Cause(err)
		if causeErr, ok := causeErr.(stackTracer); ok {
			fmt.Printf("%+v \n", causeErr.StackTrace())
		}

		os.Exit(1)
	}
}

func main() {
	flag.Parse()
	if *configFile == "" {
		flag.Usage()
		os.Exit(1)
		return
	}

	configRoot := filepath.Dir(*configFile)
	config := zanzibar.NewStaticConfigOrDie([]*zanzibar.ConfigOption{
		zanzibar.ConfigFilePath(*configFile),
	}, nil)

	configRoot, err := filepath.Abs(configRoot)
	checkError(
		err, fmt.Sprintf("can not get abs path of config dir %s", configRoot),
	)

	copyright, err := ioutil.ReadFile(filepath.Join(
		configRoot,
		config.MustGetString("copyrightHeader"),
	))
	if err != nil {
		// Default to an empty copyright for now
		copyright = []byte("")
	}

	packageHelper, err := codegen.NewPackageHelper(
		config.MustGetString("packageRoot"),
		configRoot,
		config.MustGetString("middlewareConfig"),
		config.MustGetString("thriftRootDir"),
		config.MustGetString("genCodePackage"),
		config.MustGetString("targetGenDir"),
		string(copyright),
	)
	checkError(
		err, fmt.Sprintf("Can't build package helper %s", configRoot),
	)

	moduleSystem, err := codegen.NewDefaultModuleSystem(packageHelper)
	checkError(
		err, fmt.Sprintf("Error creating module system %s", configRoot),
	)

	// TODO: All components should be generated by the module system
	fmt.Printf("Generating module system components:\n")
	_, err = moduleSystem.GenerateBuild(
		packageHelper.PackageRoot(),
		configRoot,
		packageHelper.CodeGenTargetPath(),
	)
	checkError(err, "Failed to generate module system components")
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"fmt"
	"sync"
)

type nodeCacheEntry struct {
	core     *nodeCore
	refCount int
}

// nodeCacheStandard implements the NodeCache interface by tracking
// the reference counts of nodeStandard Nodes, and using their member
// fields to construct paths.
type nodeCacheStandard struct {
	folderBranch FolderBranch

	lock         sync.RWMutex
	nodes        map[BlockRef]*nodeCacheEntry
	rootWrappers []func(Node) Node
}

var _ NodeCache = (*nodeCacheStandard)(nil)

func newNodeCacheStandard(fb FolderBranch) *nodeCacheStandard {
	return &nodeCacheStandard{
		folderBranch: fb,
		nodes:        make(map[BlockRef]*nodeCacheEntry),
	}
}

// lock must be locked for writing by the caller
func (ncs *nodeCacheStandard) forgetLocked(core *nodeCore) {
	ref := core.pathNode.Ref()

	entry, ok := ncs.nodes[ref]
	if !ok {
		return
	}
	if entry.core != core {
		return
	}

	entry.refCount--
	if entry.refCount <= 0 {
		delete(ncs.nodes, ref)
	}
}

// should be called only by nodeStandardFinalizer().
func (ncs *nodeCacheStandard) forget(core *nodeCore) {
	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	ncs.forgetLocked(core)
}

// lock must be held for writing by the caller
func (ncs *nodeCacheStandard) newChildForParentLocked(parent Node) (*nodeStandard, error) {
	nodeStandard, ok := parent.Unwrap().(*nodeStandard)
	if !ok {
		return nil, ParentNodeNotFoundError{BlockRef{}}
	}

	ref := nodeStandard.core.pathNode.Ref()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil, ParentNodeNotFoundError{ref}
	}
	if nodeStandard.core != entry.core {
		return nil, ParentNodeNotFoundError{ref}
	}
	return nodeStandard, nil
}
func (ncs *nodeCacheStandard) wrapNodeStandard(
	n Node, rootWrappers []func(Node) Node) Node {
	ns, ok := n.(*nodeStandard)
	if !ok {
		return n
	}
	if ns.core.parent != nil {
		return ns.core.parent.WrapChild(n)
	}
	for _, f := range rootWrappers {
		n = f(n)
	}
	return n
}

func (ncs *nodeCacheStandard) makeNodeStandardForEntryLocked(
	entry *nodeCacheEntry) *nodeStandard {
	entry.refCount++
	return makeNodeStandard(entry.core)
}

// GetOrCreate implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) GetOrCreate(
	ptr BlockPointer, name string, parent Node, et EntryType) (
	n Node, err error) {
	var rootWrappers []func(Node) Node
	defer func() {
		if n != nil {
			n = ncs.wrapNodeStandard(n, rootWrappers)
		}
	}()

	if !ptr.IsValid() {
		// Temporary code to track down bad block
		// pointers. Remove when not needed anymore.
		panic(InvalidBlockRefError{ptr.Ref()})
	}

	if name == "" {
		return nil, EmptyNameError{ptr.Ref()}
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	rootWrappers = ncs.rootWrappers
	entry, ok := ncs.nodes[ptr.Ref()]
	if ok {
		// If the entry happens to be unlinked, we may be in a
		// situation where a node got unlinked and then recreated, but
		// someone held onto a node the whole time and so it never got
		// removed from the cache.  In that case, forcibly remove it
		// from the cache to make room for the new node.
		if parent != nil && entry.core.parent == nil {
			delete(ncs.nodes, ptr.Ref())
		} else {
			return ncs.makeNodeStandardForEntryLocked(entry), nil
		}
	}

	if parent != nil {
		// Make sure a child can be made for this parent.
		_, err := ncs.newChildForParentLocked(parent)
		if err != nil {
			return nil, err
		}
	}

	entry = &nodeCacheEntry{
		core: newNodeCore(ptr, name, parent, ncs, et),
	}
	ncs.nodes[ptr.Ref()] = entry
	return ncs.makeNodeStandardForEntryLocked(entry), nil
}

// Get implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Get(ref BlockRef) (n Node) {
	if ref == (BlockRef{}) {
		return nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	var rootWrappers []func(Node) Node
	defer func() {
		if n != nil {
			n = ncs.wrapNodeStandard(n, rootWrappers)
		}
	}()

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	rootWrappers = ncs.rootWrappers
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil
	}
	return ncs.makeNodeStandardForEntryLocked(entry)
}

// UpdatePointer implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) UpdatePointer(
	oldRef BlockRef, newPtr BlockPointer) (updatedNode NodeID) {
	if oldRef == (BlockRef{}) && newPtr == (BlockPointer{}) {
		return nil
	}

	if !oldRef.IsValid() {
		panic(fmt.Sprintf("invalid oldRef %s with newPtr %s", oldRef, newPtr))
	}

	if !newPtr.IsValid() {
		panic(fmt.Sprintf("invalid newPtr %s with oldRef %s", newPtr, oldRef))
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[oldRef]
	if !ok {
		return nil
	}

	// Cannot update the pointer for an unlinked node.
	if entry.core.cachedPath.isValid() {
		return nil
	}

	entry.core.pathNode.BlockPointer = newPtr
	delete(ncs.nodes, oldRef)
	ncs.nodes[newPtr.Ref()] = entry
	return entry.core
}

// Move implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Move(
	ref BlockRef, newParent Node, newName string) (undoFn func(), err error) {
	if ref == (BlockRef{}) {
		return nil, nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	if newName == "" {
		return nil, EmptyNameError{ref}
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil, nil
	}

	newParentNS, err := ncs.newChildForParentLocked(newParent)
	if err != nil {
		return nil, err
	}

	oldParent := entry.core.parent
	oldName := entry.core.pathNode.Name

	entry.core.parent = newParentNS
	entry.core.pathNode.Name = newName

	return func() {
		entry.core.parent = oldParent
		entry.core.pathNode.Name = oldName
	}, nil
}

// Unlink implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Unlink(
	ref BlockRef, oldPath path, oldDe DirEntry) (undoFn func()) {
	if ref == (BlockRef{}) {
		return nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil
	}

	if entry.core.cachedPath.isValid() {
		// Already unlinked!
		return nil
	}

	oldParent := entry.core.parent
	oldName := entry.core.pathNode.Name

	entry.core.cachedPath = oldPath
	entry.core.cachedDe = oldDe
	entry.core.parent = nil
	entry.core.pathNode.Name = ""

	return func() {
		entry.core.cachedPath = path{}
		entry.core.cachedDe = DirEntry{}
		entry.core.parent = oldParent
		entry.core.pathNode.Name = oldName
	}
}

// IsUnlinked implements the NodeCache interface for
// nodeCacheStandard.
func (ncs *nodeCacheStandard) IsUnlinked(node Node) bool {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()

	ns, ok := node.Unwrap().(*nodeStandard)
	if !ok {
		return false
	}

	return ns.core.cachedPath.isValid()
}

// UnlinkedDirEntry implements the NodeCache interface for
// nodeCacheStandard.
func (ncs *nodeCacheStandard) UnlinkedDirEntry(node Node) DirEntry {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()

	ns, ok := node.Unwrap().(*nodeStandard)
	if !ok {
		return DirEntry{}
	}

	return ns.core.cachedDe
}

// UpdateUnlinkedDirEntry implements the NodeCache interface for
// nodeCacheStandard.
func (ncs *nodeCacheStandard) UpdateUnlinkedDirEntry(
	node Node, newDe DirEntry) {
	ncs.lock.Lock()
	defer ncs.lock.Unlock()

	ns, ok := node.Unwrap().(*nodeStandard)
	if !ok {
		return
	}

	ns.core.cachedDe = newDe
}

// PathFromNode implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) PathFromNode(node Node) (p path) {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()

	ns, ok := node.Unwrap().(*nodeStandard)
	if !ok {
		p.path = nil
		return
	}

	for ns != nil {
		core := ns.core
		if core.parent == nil && len(core.cachedPath.path) > 0 {
			// The node was unlinked, but is still in use, so use its
			// cached path.  The path is already reversed, so append
			// it backwards one-by-one to the existing path.  If this
			// is the first node, we can just optimize by returning
			// the complete cached path.
			if len(p.path) == 0 {
				return core.cachedPath
			}
			for i := len(core.cachedPath.path) - 1; i >= 0; i-- {
				p.path = append(p.path, core.cachedPath.path[i])
			}
			break
		}

		p.path = append(p.path, *core.pathNode)
		if core.parent != nil {
			ns = core.parent.Unwrap().(*nodeStandard)
		} else {
			break
		}
	}

	// need to reverse the path nodes
	for i := len(p.path)/2 - 1; i >= 0; i-- {
		opp := len(p.path) - 1 - i
		p.path[i], p.path[opp] = p.path[opp], p.path[i]
	}

	// TODO: would it make any sense to cache the constructed path?
	p.FolderBranch = ncs.folderBranch
	return
}

// AllNodes implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) AllNodes() (nodes []Node) {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()
	nodes = make([]Node, 0, len(ncs.nodes))
	for _, entry := range ncs.nodes {
		nodes = append(nodes, ncs.makeNodeStandardForEntryLocked(entry))
	}
	return nodes
}

// AllNodeChildren implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) AllNodeChildren(n Node) (nodes []Node) {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()
	nodes = make([]Node, 0, len(ncs.nodes))
	entryIDs := make(map[NodeID]bool)
	for _, entry := range ncs.nodes {
		var pathIDs []NodeID
		parent := entry.core.parent
		for parent != nil {
			// If the node's parent is what we're looking for (or on
			// the path to what we're looking for), include it in the
			// list.
			parentID := parent.GetID()
			if parentID == n.GetID() || entryIDs[parentID] {
				nodes = append(nodes, ncs.makeNodeStandardForEntryLocked(entry))
				for _, id := range pathIDs {
					entryIDs[id] = true
				}
				entryIDs[entry.core] = true
				break
			}

			// Otherwise, remember this parent and continue back
			// toward the root.
			pathIDs = append(pathIDs, parentID)
			ns, ok := parent.Unwrap().(*nodeStandard)
			if !ok {
				break
			}
			parent = ns.core.parent
		}
	}
	return nodes
}

func (ncs *nodeCacheStandard) AddRootWrapper(f func(Node) Node) {
	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	ncs.rootWrappers = append(ncs.rootWrappers, f)
}
<commit_msg>node_cache: don't access a node's parent without holding lock<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"fmt"
	"sync"
)

type nodeCacheEntry struct {
	core     *nodeCore
	refCount int
}

// nodeCacheStandard implements the NodeCache interface by tracking
// the reference counts of nodeStandard Nodes, and using their member
// fields to construct paths.
type nodeCacheStandard struct {
	folderBranch FolderBranch

	lock         sync.RWMutex
	nodes        map[BlockRef]*nodeCacheEntry
	rootWrappers []func(Node) Node
}

var _ NodeCache = (*nodeCacheStandard)(nil)

func newNodeCacheStandard(fb FolderBranch) *nodeCacheStandard {
	return &nodeCacheStandard{
		folderBranch: fb,
		nodes:        make(map[BlockRef]*nodeCacheEntry),
	}
}

// lock must be locked for writing by the caller
func (ncs *nodeCacheStandard) forgetLocked(core *nodeCore) {
	ref := core.pathNode.Ref()

	entry, ok := ncs.nodes[ref]
	if !ok {
		return
	}
	if entry.core != core {
		return
	}

	entry.refCount--
	if entry.refCount <= 0 {
		delete(ncs.nodes, ref)
	}
}

// should be called only by nodeStandardFinalizer().
func (ncs *nodeCacheStandard) forget(core *nodeCore) {
	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	ncs.forgetLocked(core)
}

// lock must be held for writing by the caller
func (ncs *nodeCacheStandard) newChildForParentLocked(parent Node) (*nodeStandard, error) {
	nodeStandard, ok := parent.Unwrap().(*nodeStandard)
	if !ok {
		return nil, ParentNodeNotFoundError{BlockRef{}}
	}

	ref := nodeStandard.core.pathNode.Ref()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil, ParentNodeNotFoundError{ref}
	}
	if nodeStandard.core != entry.core {
		return nil, ParentNodeNotFoundError{ref}
	}
	return nodeStandard, nil
}
func (ncs *nodeCacheStandard) wrapNodeStandard(
	n Node, rootWrappers []func(Node) Node, parent Node) Node {
	if parent != nil {
		return parent.WrapChild(n)
	}
	for _, f := range rootWrappers {
		n = f(n)
	}
	return n
}

func (ncs *nodeCacheStandard) makeNodeStandardForEntryLocked(
	entry *nodeCacheEntry) *nodeStandard {
	entry.refCount++
	return makeNodeStandard(entry.core)
}

// GetOrCreate implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) GetOrCreate(
	ptr BlockPointer, name string, parent Node, et EntryType) (
	n Node, err error) {
	var rootWrappers []func(Node) Node
	defer func() {
		if n != nil {
			n = ncs.wrapNodeStandard(n, rootWrappers, parent)
		}
	}()

	if !ptr.IsValid() {
		// Temporary code to track down bad block
		// pointers. Remove when not needed anymore.
		panic(InvalidBlockRefError{ptr.Ref()})
	}

	if name == "" {
		return nil, EmptyNameError{ptr.Ref()}
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	rootWrappers = ncs.rootWrappers
	entry, ok := ncs.nodes[ptr.Ref()]
	if ok {
		// If the entry happens to be unlinked, we may be in a
		// situation where a node got unlinked and then recreated, but
		// someone held onto a node the whole time and so it never got
		// removed from the cache.  In that case, forcibly remove it
		// from the cache to make room for the new node.
		if parent != nil && entry.core.parent == nil {
			delete(ncs.nodes, ptr.Ref())
		} else {
			return ncs.makeNodeStandardForEntryLocked(entry), nil
		}
	}

	if parent != nil {
		// Make sure a child can be made for this parent.
		_, err := ncs.newChildForParentLocked(parent)
		if err != nil {
			return nil, err
		}
	}

	entry = &nodeCacheEntry{
		core: newNodeCore(ptr, name, parent, ncs, et),
	}
	ncs.nodes[ptr.Ref()] = entry
	return ncs.makeNodeStandardForEntryLocked(entry), nil
}

// Get implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Get(ref BlockRef) (n Node) {
	if ref == (BlockRef{}) {
		return nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	var rootWrappers []func(Node) Node
	var parent Node
	defer func() {
		if n != nil {
			n = ncs.wrapNodeStandard(n, rootWrappers, parent)
		}
	}()

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	rootWrappers = ncs.rootWrappers
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil
	}
	ns := ncs.makeNodeStandardForEntryLocked(entry)
	parent = ns.core.parent // get while under lock
	return ns
}

// UpdatePointer implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) UpdatePointer(
	oldRef BlockRef, newPtr BlockPointer) (updatedNode NodeID) {
	if oldRef == (BlockRef{}) && newPtr == (BlockPointer{}) {
		return nil
	}

	if !oldRef.IsValid() {
		panic(fmt.Sprintf("invalid oldRef %s with newPtr %s", oldRef, newPtr))
	}

	if !newPtr.IsValid() {
		panic(fmt.Sprintf("invalid newPtr %s with oldRef %s", newPtr, oldRef))
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[oldRef]
	if !ok {
		return nil
	}

	// Cannot update the pointer for an unlinked node.
	if entry.core.cachedPath.isValid() {
		return nil
	}

	entry.core.pathNode.BlockPointer = newPtr
	delete(ncs.nodes, oldRef)
	ncs.nodes[newPtr.Ref()] = entry
	return entry.core
}

// Move implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Move(
	ref BlockRef, newParent Node, newName string) (undoFn func(), err error) {
	if ref == (BlockRef{}) {
		return nil, nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	if newName == "" {
		return nil, EmptyNameError{ref}
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil, nil
	}

	newParentNS, err := ncs.newChildForParentLocked(newParent)
	if err != nil {
		return nil, err
	}

	oldParent := entry.core.parent
	oldName := entry.core.pathNode.Name

	entry.core.parent = newParentNS
	entry.core.pathNode.Name = newName

	return func() {
		entry.core.parent = oldParent
		entry.core.pathNode.Name = oldName
	}, nil
}

// Unlink implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Unlink(
	ref BlockRef, oldPath path, oldDe DirEntry) (undoFn func()) {
	if ref == (BlockRef{}) {
		return nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil
	}

	if entry.core.cachedPath.isValid() {
		// Already unlinked!
		return nil
	}

	oldParent := entry.core.parent
	oldName := entry.core.pathNode.Name

	entry.core.cachedPath = oldPath
	entry.core.cachedDe = oldDe
	entry.core.parent = nil
	entry.core.pathNode.Name = ""

	return func() {
		entry.core.cachedPath = path{}
		entry.core.cachedDe = DirEntry{}
		entry.core.parent = oldParent
		entry.core.pathNode.Name = oldName
	}
}

// IsUnlinked implements the NodeCache interface for
// nodeCacheStandard.
func (ncs *nodeCacheStandard) IsUnlinked(node Node) bool {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()

	ns, ok := node.Unwrap().(*nodeStandard)
	if !ok {
		return false
	}

	return ns.core.cachedPath.isValid()
}

// UnlinkedDirEntry implements the NodeCache interface for
// nodeCacheStandard.
func (ncs *nodeCacheStandard) UnlinkedDirEntry(node Node) DirEntry {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()

	ns, ok := node.Unwrap().(*nodeStandard)
	if !ok {
		return DirEntry{}
	}

	return ns.core.cachedDe
}

// UpdateUnlinkedDirEntry implements the NodeCache interface for
// nodeCacheStandard.
func (ncs *nodeCacheStandard) UpdateUnlinkedDirEntry(
	node Node, newDe DirEntry) {
	ncs.lock.Lock()
	defer ncs.lock.Unlock()

	ns, ok := node.Unwrap().(*nodeStandard)
	if !ok {
		return
	}

	ns.core.cachedDe = newDe
}

// PathFromNode implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) PathFromNode(node Node) (p path) {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()

	ns, ok := node.Unwrap().(*nodeStandard)
	if !ok {
		p.path = nil
		return
	}

	for ns != nil {
		core := ns.core
		if core.parent == nil && len(core.cachedPath.path) > 0 {
			// The node was unlinked, but is still in use, so use its
			// cached path.  The path is already reversed, so append
			// it backwards one-by-one to the existing path.  If this
			// is the first node, we can just optimize by returning
			// the complete cached path.
			if len(p.path) == 0 {
				return core.cachedPath
			}
			for i := len(core.cachedPath.path) - 1; i >= 0; i-- {
				p.path = append(p.path, core.cachedPath.path[i])
			}
			break
		}

		p.path = append(p.path, *core.pathNode)
		if core.parent != nil {
			ns = core.parent.Unwrap().(*nodeStandard)
		} else {
			break
		}
	}

	// need to reverse the path nodes
	for i := len(p.path)/2 - 1; i >= 0; i-- {
		opp := len(p.path) - 1 - i
		p.path[i], p.path[opp] = p.path[opp], p.path[i]
	}

	// TODO: would it make any sense to cache the constructed path?
	p.FolderBranch = ncs.folderBranch
	return
}

// AllNodes implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) AllNodes() (nodes []Node) {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()
	nodes = make([]Node, 0, len(ncs.nodes))
	for _, entry := range ncs.nodes {
		nodes = append(nodes, ncs.makeNodeStandardForEntryLocked(entry))
	}
	return nodes
}

// AllNodeChildren implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) AllNodeChildren(n Node) (nodes []Node) {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()
	nodes = make([]Node, 0, len(ncs.nodes))
	entryIDs := make(map[NodeID]bool)
	for _, entry := range ncs.nodes {
		var pathIDs []NodeID
		parent := entry.core.parent
		for parent != nil {
			// If the node's parent is what we're looking for (or on
			// the path to what we're looking for), include it in the
			// list.
			parentID := parent.GetID()
			if parentID == n.GetID() || entryIDs[parentID] {
				nodes = append(nodes, ncs.makeNodeStandardForEntryLocked(entry))
				for _, id := range pathIDs {
					entryIDs[id] = true
				}
				entryIDs[entry.core] = true
				break
			}

			// Otherwise, remember this parent and continue back
			// toward the root.
			pathIDs = append(pathIDs, parentID)
			ns, ok := parent.Unwrap().(*nodeStandard)
			if !ok {
				break
			}
			parent = ns.core.parent
		}
	}
	return nodes
}

func (ncs *nodeCacheStandard) AddRootWrapper(f func(Node) Node) {
	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	ncs.rootWrappers = append(ncs.rootWrappers, f)
}
<|endoftext|>