<commit_before>package main

import (
	"bytes"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"

	"code.google.com/p/go.crypto/ssh"
	"github.com/flynn/go-shlex"
)

const PrereceiveHookTmpl = `#!/bin/bash
set -eo pipefail; while read oldrev newrev refname; do
[[ $refname = "refs/heads/master" ]] && git archive $newrev | {{RECEIVER}} "$RECEIVE_REPO" "$newrev" | sed -$([[ $(uname) == "Darwin" ]] && echo l || echo u) "s/^/"$'\e[1G'"/"
done
`

var prereceiveHook []byte

var port = flag.String("p", "22", "port to listen on")
var repoPath = flag.String("r", "/tmp/repos", "path to repo cache")
var noAuth = flag.Bool("n", false, "disable client authentication")
var keys = flag.String("k", "", "pem file containing private keys (read from SSH_PRIVATE_KEYS by default)")

var authChecker []string

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %v [options] <authchecker> <receiver>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if flag.NArg() < 2 {
		flag.Usage()
		os.Exit(64)
	}
	var err error
	authChecker, err = shlex.Split(flag.Arg(0))
	if err != nil {
		log.Fatalln("Unable to parse authchecker command:", err)
	}

	receiver, err := shlex.Split(flag.Arg(1))
	if err != nil {
		log.Fatalln("Unable to parse receiver command:", err)
	}
	receiver[0], err = filepath.Abs(receiver[0])
	if err != nil {
		log.Fatalln("Invalid receiver path:", err)
	}
	prereceiveHook = []byte(strings.Replace(PrereceiveHookTmpl, "{{RECEIVER}}", strings.Join(receiver, " "), 1))

	var config *ssh.ServerConfig
	if *noAuth {
		config = &ssh.ServerConfig{NoClientAuth: true}
	} else {
		config = &ssh.ServerConfig{PublicKeyCallback: checkAuth}
	}

	if keyEnv := os.Getenv("SSH_PRIVATE_KEYS"); keyEnv != "" {
		if err := parseKeys(config, []byte(keyEnv)); err != nil {
			log.Fatalln("Failed to load private keys:", err)
		}
	} else {
		pemBytes, err := ioutil.ReadFile(*keys)
		if err != nil {
			log.Fatalln("Failed to load private keys:", err)
		}
		if err := parseKeys(config, pemBytes); err != nil {
			log.Fatalln("Failed to parse private keys:", err)
		}
	}

	if p := os.Getenv("PORT"); p != "" && *port == "22" {
		*port = p
	}
	listener, err := net.Listen("tcp", ":"+*port)
	if err != nil {
		log.Fatalln("Failed to listen for connections:", err)
	}
	for {
		// SSH connections just house multiplexed connections
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Failed to accept incoming connection:", err)
			continue
		}
		go handleConn(conn, config)
	}
}

func parseKeys(conf *ssh.ServerConfig, pemData []byte) error {
	var found bool
	for {
		var block *pem.Block
		block, pemData = pem.Decode(pemData)
		if block == nil {
			if !found {
				return errors.New("no private keys found")
			}
			return nil
		}
		if err := addKey(conf, block); err != nil {
			return err
		}
		found = true
	}
}

func addKey(conf *ssh.ServerConfig, block *pem.Block) (err error) {
	var key interface{}
	switch block.Type {
	case "RSA PRIVATE KEY":
		key, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	case "EC PRIVATE KEY":
		key, err = x509.ParseECPrivateKey(block.Bytes)
	case "DSA PRIVATE KEY":
		key, err = ssh.ParseDSAPrivateKey(block.Bytes)
	default:
		return fmt.Errorf("unsupported key type %q", block.Type)
	}
	if err != nil {
		return err
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return err
	}
	conf.AddHostKey(signer)
	return nil
}

func handleConn(conn net.Conn, conf *ssh.ServerConfig) {
	defer conn.Close()
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, conf)
	if err != nil {
		log.Println("Failed to handshake:", err)
		return
	}

	go ssh.DiscardRequests(reqs)

	for ch := range chans {
		if ch.ChannelType() != "session" {
			ch.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		go handleChannel(sshConn, ch)
	}
}

func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel) {
	ch, reqs, err := newChan.Accept()
	if err != nil {
		log.Println("newChan.Accept failed:", err)
		return
	}
	defer ch.Close()
	for req := range reqs {
		switch req.Type {
		case "exec":
			fail := func(at string, err error) {
				log.Printf("%s failed: %s", at, err)
				ch.Stderr().Write([]byte("Internal error.\n"))
			}
			if req.WantReply {
				req.Reply(true, nil)
			}
			cmdline := string(req.Payload[4:])
			cmdargs, err := shlex.Split(cmdline)
			if err != nil || len(cmdargs) != 2 {
				ch.Stderr().Write([]byte("Invalid arguments.\n"))
				return
			}
			if cmdargs[0] != "git-receive-pack" {
				ch.Stderr().Write([]byte("Only `git push` is supported.\n"))
				return
			}
			cmdargs[1] = strings.TrimSuffix(strings.TrimPrefix(cmdargs[1], "/"), ".git")
			if strings.Contains(cmdargs[1], "..") {
				ch.Stderr().Write([]byte("Invalid repo.\n"))
				return
			}

			if err := ensureCacheRepo(cmdargs[1]); err != nil {
				fail("ensureCacheRepo", err)
				return
			}
			cmd := exec.Command("git-shell", "-c", cmdargs[0]+" '"+cmdargs[1]+"'")
			cmd.Dir = *repoPath
			cmd.Env = append(os.Environ(),
				"RECEIVE_USER="+conn.User(),
				"RECEIVE_REPO="+cmdargs[1],
			)
			done, err := attachCmd(cmd, ch, ch.Stderr(), ch)
			if err != nil {
				fail("attachCmd", err)
				return
			}
			if err := cmd.Start(); err != nil {
				fail("cmd.Start", err)
				return
			}
			done.Wait()
			status, err := exitStatus(cmd.Wait())
			if err != nil {
				fail("exitStatus", err)
				return
			}
			if _, err := ch.SendRequest("exit-status", false, ssh.Marshal(&status)); err != nil {
				fail("sendExit", err)
			}
			return
		case "env":
			if req.WantReply {
				req.Reply(true, nil)
			}
		}
	}
}

var ErrUnauthorized = errors.New("gitreceive: user is unauthorized")

func checkAuth(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	status, err := exitStatus(exec.Command(authChecker[0],
		append(authChecker[1:], conn.User(), string(bytes.TrimSpace(ssh.MarshalAuthorizedKey(key))))...).Run())
	if err != nil {
		return nil, err
	}
	if status.Status == 0 {
		return nil, nil
	}
	return nil, ErrUnauthorized
}

func attachCmd(cmd *exec.Cmd, stdout, stderr io.Writer, stdin io.Reader) (*sync.WaitGroup, error) {
	var wg sync.WaitGroup
	wg.Add(2)

	stdinIn, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdoutOut, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderrOut, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}

	go func() {
		io.Copy(stdinIn, stdin)
		stdinIn.Close()
	}()
	go func() {
		io.Copy(stdout, stdoutOut)
		wg.Done()
	}()
	go func() {
		io.Copy(stderr, stderrOut)
		wg.Done()
	}()

	return &wg, nil
}

type exitStatusMsg struct {
	Status uint32
}

func exitStatus(err error) (exitStatusMsg, error) {
	if err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			// There is no platform independent way to retrieve
			// the exit code, but the following will work on Unix
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				return exitStatusMsg{uint32(status.ExitStatus())}, nil
			}
		}
		return exitStatusMsg{0}, err
	}
	return exitStatusMsg{0}, nil
}

var cacheMtx sync.Mutex

func ensureCacheRepo(path string) error {
	cacheMtx.Lock()
	defer cacheMtx.Unlock()

	cachePath := *repoPath + "/" + path
	if _, err := os.Stat(cachePath); os.IsNotExist(err) {
		os.MkdirAll(cachePath, 0755)
		cmd := exec.Command("git", "init", "--bare")
		cmd.Dir = cachePath
		err = cmd.Run()
		if err != nil {
			return err
		}
	}
	return ioutil.WriteFile(cachePath+"/hooks/pre-receive", prereceiveHook, 0755)
}
<commit_msg>Clear "remote:" even where an output line is shorter than it<commit_after>package main

import (
	"bytes"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"

	"code.google.com/p/go.crypto/ssh"
	"github.com/flynn/go-shlex"
)

const PrereceiveHookTmpl = `#!/bin/bash
set -eo pipefail; while read oldrev newrev refname; do
[[ $refname = "refs/heads/master" ]] && git archive $newrev | {{RECEIVER}} "$RECEIVE_REPO" "$newrev" | sed -$([[ $(uname) == "Darwin" ]] && echo l || echo u) "s/^/"$'\e[1G\e[K'"/"
done
`

var prereceiveHook []byte

var port = flag.String("p", "22", "port to listen on")
var repoPath = flag.String("r", "/tmp/repos", "path to repo cache")
var noAuth = flag.Bool("n", false, "disable client authentication")
var keys = flag.String("k", "", "pem file containing private keys (read from SSH_PRIVATE_KEYS by default)")

var authChecker []string

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %v [options] <authchecker> <receiver>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if flag.NArg() < 2 {
		flag.Usage()
		os.Exit(64)
	}
	var err error
	authChecker, err = shlex.Split(flag.Arg(0))
	if err != nil {
		log.Fatalln("Unable to parse authchecker command:", err)
	}

	receiver, err := shlex.Split(flag.Arg(1))
	if err != nil {
		log.Fatalln("Unable to parse receiver command:", err)
	}
	receiver[0], err = filepath.Abs(receiver[0])
	if err != nil {
		log.Fatalln("Invalid receiver path:", err)
	}
	prereceiveHook = []byte(strings.Replace(PrereceiveHookTmpl, "{{RECEIVER}}", strings.Join(receiver, " "), 1))

	var config *ssh.ServerConfig
	if *noAuth {
		config = &ssh.ServerConfig{NoClientAuth: true}
	} else {
		config = &ssh.ServerConfig{PublicKeyCallback: checkAuth}
	}

	if keyEnv := os.Getenv("SSH_PRIVATE_KEYS"); keyEnv != "" {
		if err := parseKeys(config, []byte(keyEnv)); err != nil {
			log.Fatalln("Failed to load private keys:", err)
		}
	} else {
		pemBytes, err := ioutil.ReadFile(*keys)
		if err != nil {
			log.Fatalln("Failed to load private keys:", err)
		}
		if err := parseKeys(config, pemBytes); err != nil {
			log.Fatalln("Failed to parse private keys:", err)
		}
	}

	if p := os.Getenv("PORT"); p != "" && *port == "22" {
		*port = p
	}
	listener, err := net.Listen("tcp", ":"+*port)
	if err != nil {
		log.Fatalln("Failed to listen for connections:", err)
	}
	for {
		// SSH connections just house multiplexed connections
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Failed to accept incoming connection:", err)
			continue
		}
		go handleConn(conn, config)
	}
}

func parseKeys(conf *ssh.ServerConfig, pemData []byte) error {
	var found bool
	for {
		var block *pem.Block
		block, pemData = pem.Decode(pemData)
		if block == nil {
			if !found {
				return errors.New("no private keys found")
			}
			return nil
		}
		if err := addKey(conf, block); err != nil {
			return err
		}
		found = true
	}
}

func addKey(conf *ssh.ServerConfig, block *pem.Block) (err error) {
	var key interface{}
	switch block.Type {
	case "RSA PRIVATE KEY":
		key, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	case "EC PRIVATE KEY":
		key, err = x509.ParseECPrivateKey(block.Bytes)
	case "DSA PRIVATE KEY":
		key, err = ssh.ParseDSAPrivateKey(block.Bytes)
	default:
		return fmt.Errorf("unsupported key type %q", block.Type)
	}
	if err != nil {
		return err
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return err
	}
	conf.AddHostKey(signer)
	return nil
}

func handleConn(conn net.Conn, conf *ssh.ServerConfig) {
	defer conn.Close()
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, conf)
	if err != nil {
		log.Println("Failed to handshake:", err)
		return
	}

	go ssh.DiscardRequests(reqs)

	for ch := range chans {
		if ch.ChannelType() != "session" {
			ch.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		go handleChannel(sshConn, ch)
	}
}

func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel) {
	ch, reqs, err := newChan.Accept()
	if err != nil {
		log.Println("newChan.Accept failed:", err)
		return
	}
	defer ch.Close()
	for req := range reqs {
		switch req.Type {
		case "exec":
			fail := func(at string, err error) {
				log.Printf("%s failed: %s", at, err)
				ch.Stderr().Write([]byte("Internal error.\n"))
			}
			if req.WantReply {
				req.Reply(true, nil)
			}
			cmdline := string(req.Payload[4:])
			cmdargs, err := shlex.Split(cmdline)
			if err != nil || len(cmdargs) != 2 {
				ch.Stderr().Write([]byte("Invalid arguments.\n"))
				return
			}
			if cmdargs[0] != "git-receive-pack" {
				ch.Stderr().Write([]byte("Only `git push` is supported.\n"))
				return
			}
			cmdargs[1] = strings.TrimSuffix(strings.TrimPrefix(cmdargs[1], "/"), ".git")
			if strings.Contains(cmdargs[1], "..") {
				ch.Stderr().Write([]byte("Invalid repo.\n"))
				return
			}

			if err := ensureCacheRepo(cmdargs[1]); err != nil {
				fail("ensureCacheRepo", err)
				return
			}
			cmd := exec.Command("git-shell", "-c", cmdargs[0]+" '"+cmdargs[1]+"'")
			cmd.Dir = *repoPath
			cmd.Env = append(os.Environ(),
				"RECEIVE_USER="+conn.User(),
				"RECEIVE_REPO="+cmdargs[1],
			)
			done, err := attachCmd(cmd, ch, ch.Stderr(), ch)
			if err != nil {
				fail("attachCmd", err)
				return
			}
			if err := cmd.Start(); err != nil {
				fail("cmd.Start", err)
				return
			}
			done.Wait()
			status, err := exitStatus(cmd.Wait())
			if err != nil {
				fail("exitStatus", err)
				return
			}
			if _, err := ch.SendRequest("exit-status", false, ssh.Marshal(&status)); err != nil {
				fail("sendExit", err)
			}
			return
		case "env":
			if req.WantReply {
				req.Reply(true, nil)
			}
		}
	}
}

var ErrUnauthorized = errors.New("gitreceive: user is unauthorized")

func checkAuth(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	status, err := exitStatus(exec.Command(authChecker[0],
		append(authChecker[1:], conn.User(), string(bytes.TrimSpace(ssh.MarshalAuthorizedKey(key))))...).Run())
	if err != nil {
		return nil, err
	}
	if status.Status == 0 {
		return nil, nil
	}
	return nil, ErrUnauthorized
}

func attachCmd(cmd *exec.Cmd, stdout, stderr io.Writer, stdin io.Reader) (*sync.WaitGroup, error) {
	var wg sync.WaitGroup
	wg.Add(2)

	stdinIn, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdoutOut, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderrOut, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}

	go func() {
		io.Copy(stdinIn, stdin)
		stdinIn.Close()
	}()
	go func() {
		io.Copy(stdout, stdoutOut)
		wg.Done()
	}()
	go func() {
		io.Copy(stderr, stderrOut)
		wg.Done()
	}()

	return &wg, nil
}

type exitStatusMsg struct {
	Status uint32
}

func exitStatus(err error) (exitStatusMsg, error) {
	if err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			// There is no platform independent way to retrieve
			// the exit code, but the following will work on Unix
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				return exitStatusMsg{uint32(status.ExitStatus())}, nil
			}
		}
		return exitStatusMsg{0}, err
	}
	return exitStatusMsg{0}, nil
}

var cacheMtx sync.Mutex

func ensureCacheRepo(path string) error {
	cacheMtx.Lock()
	defer cacheMtx.Unlock()

	cachePath := *repoPath + "/" + path
	if _, err := os.Stat(cachePath); os.IsNotExist(err) {
		os.MkdirAll(cachePath, 0755)
		cmd := exec.Command("git", "init", "--bare")
		cmd.Dir = cachePath
		err = cmd.Run()
		if err != nil {
			return err
		}
	}
	return ioutil.WriteFile(cachePath+"/hooks/pre-receive", prereceiveHook, 0755)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
	"os"
	"path/filepath"

	"./arch"
	"./loader"
	"./models"
	"./syscalls"
)

type Usercorn struct {
	*Unicorn
	loader       models.Loader
	interpLoader models.Loader

	base       uint64
	interpBase uint64
	entry      uint64
	binEntry   uint64

	StackBase   uint64
	DataSegment models.Segment
	Verbose     bool
	TraceSys    bool
	TraceMem    bool
	TraceExec   bool
	TraceReg    bool
	LoadPrefix  string
	status      models.StatusDiff
	stacktrace  models.Stacktrace

	// deadlock detection
	lastBlock uint64
	lastCode  uint64
	deadlock  int
}

func NewUsercorn(exe string, prefix string) (*Usercorn, error) {
	l, err := loader.LoadFile(exe)
	if err != nil {
		return nil, err
	}
	a, os, err := arch.GetArch(l.Arch(), l.OS())
	if err != nil {
		return nil, err
	}
	unicorn, err := NewUnicorn(a, os, l.ByteOrder())
	if err != nil {
		return nil, err
	}
	ds, de := l.DataSegment()
	u := &Usercorn{
		Unicorn:     unicorn,
		loader:      l,
		LoadPrefix:  prefix,
		DataSegment: models.Segment{ds, de},
	}
	u.status = models.StatusDiff{U: u, Color: true}
	u.interpBase, u.entry, u.base, u.binEntry, err = u.mapBinary(u.loader)
	if err != nil {
		return nil, err
	}
	return u, nil
}

func (u *Usercorn) Run(args []string, env []string) error {
	if err := u.addHooks(); err != nil {
		return err
	}
	if err := u.setupStack(); err != nil {
		return err
	}
	if u.OS.Init != nil {
		if err := u.OS.Init(u, args, env); err != nil {
			return err
		}
	}
	if u.Verbose {
		fmt.Fprintf(os.Stderr, "[entry @ 0x%x]\n", u.entry)
		dis, err := u.Disas(u.entry, 64)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
		} else {
			fmt.Fprintln(os.Stderr, dis)
		}
		sp, err := u.RegRead(u.arch.SP)
		if err != nil {
			return err
		}
		buf := make([]byte, u.StackBase+STACK_SIZE-sp)
		if err := u.MemReadInto(buf, sp); err != nil {
			return err
		}
		fmt.Fprintf(os.Stderr, "[stack @ 0x%x]\n", sp)
		fmt.Fprintf(os.Stderr, "%s\n", HexDump(sp, buf[:], u.arch))
	}
	if u.Verbose || u.TraceReg {
		u.status.Changes().Print(true, false)
	}
	if u.Verbose {
		fmt.Fprintln(os.Stderr, "=====================================")
		fmt.Fprintln(os.Stderr, "==== Program output begins here. ====")
		fmt.Fprintln(os.Stderr, "=====================================")
	}
	if u.TraceReg || u.TraceExec {
		sp, _ := u.RegRead(u.arch.SP)
		u.stacktrace.Update(u.entry, sp)
	}
	err := u.Unicorn.Start(u.entry, 0xffffffffffffffff)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Registers:")
		u.status.Changes().Print(true, false)
		fmt.Fprintln(os.Stderr, "Stacktrace:")
		u.stacktrace.Print(u)
	}
	return err
}

func (u *Usercorn) Loader() models.Loader {
	return u.loader
}

func (u *Usercorn) InterpBase() uint64 {
	// points to interpreter base or 0
	return u.interpBase
}

func (u *Usercorn) Base() uint64 {
	// points to program base
	return u.base
}

func (u *Usercorn) Entry() uint64 {
	// points to effective program entry: either an interpreter or the binary
	return u.entry
}

func (u *Usercorn) BinEntry() uint64 {
	// points to binary entry, even if an interpreter is used
	return u.binEntry
}

func (u *Usercorn) PosixInit(args, env []string, auxv []byte) error {
	// end marker
	if err := u.Push(0); err != nil {
		return err
	}
	// auxv
	if err := u.PushBytes(auxv); err != nil {
		return err
	}
	// envp
	envp, err := u.pushStrings(env...)
	if err != nil {
		return err
	}
	if err := u.pushAddrs(envp); err != nil {
		return err
	}
	// argv
	argv, err := u.pushStrings(args...)
	if err != nil {
		return err
	}
	if err := u.pushAddrs(argv); err != nil {
		return err
	}
	// argc
	return u.Push(uint64(len(args)))
}

func (u *Usercorn) PrefixPath(path string, force bool) string {
	if filepath.IsAbs(path) && u.LoadPrefix != "" {
		target := filepath.Join(u.LoadPrefix, path)
		_, err := os.Stat(target)
		exists := !os.IsNotExist(err)
		if force || exists {
			return target
		}
	}
	return path
}

func (u *Usercorn) Symbolicate(addr uint64) (string, error) {
	var symbolicate = func(addr uint64, symbols []models.Symbol) (result models.Symbol, distance uint64) {
		if len(symbols) == 0 {
			return
		}
		nearest := make(map[uint64][]models.Symbol)
		var min int64 = -1
		for _, sym := range symbols {
			dist := int64(addr - sym.Start)
			if dist > 0 && (sym.Start+uint64(dist) <= sym.End || sym.End == 0) && sym.Name != "" {
				if dist < min || min == -1 {
					min = dist
				}
				nearest[uint64(dist)] = append(nearest[uint64(dist)], sym)
			}
		}
		if len(nearest) > 0 {
			sym := nearest[uint64(min)][0]
			return sym, uint64(min)
		}
		return
	}
	symbols, err := u.loader.Symbols()
	if err != nil {
		return "", err
	}
	var interpSym []models.Symbol
	if u.interpLoader != nil {
		interpSym, _ = u.interpLoader.Symbols()
	}
	sym, sdist := symbolicate(addr-u.base, symbols)
	isym, idist := symbolicate(addr-u.interpBase, interpSym)
	if idist < sdist && isym.Name != "" || sym.Name == "" {
		sym = isym
		sdist = idist
	}
	if sym.Name != "" {
		return fmt.Sprintf("%s+0x%x", sym.Name, sdist), nil
	}
	return "", nil
}

func (u *Usercorn) Brk(addr uint64) (uint64, error) {
	// TODO: this is linux specific
	s := u.DataSegment
	if addr > 0 {
		u.MemMap(s.End, addr)
		s.End = addr
	}
	return s.End, nil
}

func (u *Usercorn) addHooks() error {
	if u.TraceExec || u.TraceReg {
		u.HookAdd(uc.HOOK_BLOCK, func(_ uc.Unicorn, addr uint64, size uint32) {
			sym, _ := u.Symbolicate(addr)
			if sym != "" {
				sym = " (" + sym + ")"
			}
			blockLine := fmt.Sprintf("| block%s @0x%x", sym, addr)
			if !u.TraceExec && u.TraceReg && u.deadlock == 0 {
				changes := u.status.Changes()
				if changes.Count() > 0 {
					fmt.Fprintln(os.Stderr, blockLine)
					changes.Print(true, true)
				}
			} else {
				fmt.Fprintln(os.Stderr, blockLine)
			}
			if sp, err := u.RegRead(u.arch.SP); err == nil {
				u.stacktrace.Update(addr, sp)
			}
			u.lastBlock = addr
		})
	}
	if u.TraceExec {
		u.HookAdd(uc.HOOK_CODE, func(_ uc.Unicorn, addr uint64, size uint32) {
			var changes *models.Changes
			if addr == u.lastCode || u.TraceReg && u.TraceExec {
				changes = u.status.Changes()
			}
			if u.TraceExec {
				if u.TraceReg {
					changes.Print(true, true)
				}
				dis, _ := u.Disas(addr, uint64(size))
				fmt.Fprintln(os.Stderr, dis)
			}
			if addr == u.lastCode {
				u.deadlock++
				if changes.Count() > 0 {
					if u.TraceReg {
						changes.Print(true, true)
					}
					u.deadlock = 0
				}
				if u.deadlock > 2 {
					sym, _ := u.Symbolicate(addr)
					if sym != "" {
						sym = " (" + sym + ")"
					}
					fmt.Fprintf(os.Stderr, "FATAL: deadlock detected at 0x%x%s\n", addr, sym)
					changes.Print(true, false)
					u.Stop()
				}
			} else {
				u.deadlock = 0
			}
			u.lastCode = addr
		})
	}
	if u.TraceMem {
		u.HookAdd(uc.HOOK_MEM_READ_WRITE, func(_ uc.Unicorn, access int, addr uint64, size int, value int64) {
			if access == uc.MEM_WRITE {
				fmt.Fprintf(os.Stderr, "MEM_WRITE")
			} else {
				fmt.Fprintf(os.Stderr, "MEM_READ")
			}
			fmt.Fprintf(os.Stderr, " 0x%x %d 0x%x\n", addr, size, value)
		})
	}
	u.HookAdd(uc.HOOK_MEM_INVALID, func(_ uc.Unicorn, access int, addr uint64, size int, value int64) bool {
		if access == uc.MEM_WRITE {
			fmt.Fprintf(os.Stderr, "invalid write")
		} else {
			fmt.Fprintf(os.Stderr, "invalid read")
		}
		fmt.Fprintf(os.Stderr, ": @0x%x, 0x%x = 0x%x\n", addr, size, value)
		return false
	})
	u.HookAdd(uc.HOOK_INTR, func(_ uc.Unicorn, intno uint32) {
		u.OS.Interrupt(u, intno)
	})
	return nil
}

func (u *Usercorn) mapBinary(l models.Loader) (interpBase, entry, base, realEntry uint64, err error) {
	var dynamic bool
	switch l.Type() {
	case loader.EXEC:
		dynamic = false
	case loader.DYN:
		dynamic = true
	default:
		err = errors.New("Unsupported file load type.")
		return
	}
	segments, err := l.Segments()
	if err != nil {
		return
	}
	// merge overlapping segments
	merged := make([]*models.Segment, 0, len(segments))
outer:
	for _, seg := range segments {
		addr, size := align(seg.Addr, seg.Size, true)
		s := &models.Segment{addr, addr + size}
		for _, s2 := range merged {
			if s2.Overlaps(s) {
				s2.Merge(s)
				continue outer
			}
		}
		merged = append(merged, s)
	}
	// map merged segments
	var loadBias uint64
	for _, seg := range merged {
		size := seg.End - seg.Start
		if dynamic && seg.Start == 0 && loadBias == 0 {
			loadBias, err = u.Mmap(0x1000000, size)
		} else {
			err = u.MemMap(loadBias+seg.Start, seg.End-seg.Start)
		}
		if err != nil {
			return
		}
	}
	// write segment memory
	for _, seg := range segments {
		if err = u.MemWrite(loadBias+seg.Addr, seg.Data); err != nil {
			return
		}
	}
	entry = loadBias + l.Entry()
	// load interpreter if present
	interp := l.Interp()
	if interp != "" {
		var bin models.Loader
		bin, err = loader.LoadFile(u.PrefixPath(interp, true))
		if err != nil {
			return
		}
		u.interpLoader = bin
		_, _, interpBias, interpEntry, err := u.mapBinary(bin)
		return interpBias, interpEntry, loadBias, entry, err
	} else {
		return 0, entry, loadBias, entry, nil
	}
}

func (u *Usercorn) setupStack() error {
	stack, err := u.Mmap(STACK_BASE, STACK_SIZE)
	if err != nil {
		return err
	}
	u.StackBase = stack
	if err := u.RegWrite(u.arch.SP, stack+STACK_SIZE); err != nil {
		return err
	}
	return nil
}

func (u *Usercorn) pushStrings(args ...string) ([]uint64, error) {
	// TODO: does anything case if these are actually on the stack?
	argvSize := 0
	for _, v := range args {
		argvSize += len(v) + 1
	}
	argvAddr, err := u.Mmap(0, uint64(argvSize))
	if err != nil {
		return nil, err
	}
	buf := make([]byte, argvSize)
	addrs := make([]uint64, 0, len(args)+1)
	var pos uint64
	for i := len(args) - 1; i >= 0; i-- {
		copy(buf[pos:], []byte(args[i]))
		addrs = append(addrs, argvAddr+pos)
		pos += uint64(len(args[i]) + 1)
	}
	u.MemWrite(argvAddr, buf)
	return addrs, nil
}

func (u *Usercorn) pushAddrs(addrs []uint64) error {
	if err := u.Push(0); err != nil {
		return err
	}
	for _, v := range addrs {
		if err := u.Push(v); err != nil {
			return err
		}
	}
	return nil
}

func (u *Usercorn) Syscall(num int, name string, getArgs func(n int) ([]uint64, error)) (uint64, error) {
	if name == "" {
		panic(fmt.Sprintf("Syscall missing: %d", num))
	}
	return syscalls.Call(u, num, name, getArgs, u.TraceSys)
}
<commit_msg>fix symbols in glibc test (#42)<commit_after>package main

import (
	"errors"
	"fmt"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
	"os"
	"path/filepath"

	"./arch"
	"./loader"
	"./models"
	"./syscalls"
)

type Usercorn struct {
	*Unicorn
	loader       models.Loader
	interpLoader models.Loader

	base       uint64
	interpBase uint64
	entry      uint64
	binEntry   uint64

	StackBase   uint64
	DataSegment models.Segment
	Verbose     bool
	TraceSys    bool
	TraceMem    bool
	TraceExec   bool
	TraceReg    bool
	LoadPrefix  string
	status      models.StatusDiff
	stacktrace  models.Stacktrace

	// deadlock detection
	lastBlock uint64
	lastCode  uint64
	deadlock  int
}

func NewUsercorn(exe string, prefix string) (*Usercorn, error) {
	l, err := loader.LoadFile(exe)
	if err != nil {
		return nil, err
	}
	a, os, err := arch.GetArch(l.Arch(), l.OS())
	if err != nil {
		return nil, err
	}
	unicorn, err := NewUnicorn(a, os, l.ByteOrder())
	if err != nil {
		return nil, err
	}
	ds, de := l.DataSegment()
	u := &Usercorn{
		Unicorn:     unicorn,
		loader:      l,
		LoadPrefix:  prefix,
		DataSegment: models.Segment{ds, de},
	}
	u.status = models.StatusDiff{U: u, Color: true}
	u.interpBase, u.entry, u.base, u.binEntry, err = u.mapBinary(u.loader)
	if err != nil {
		return nil, err
	}
	return u, nil
}

func (u *Usercorn) Run(args []string, env []string) error {
	if err := u.addHooks(); err != nil {
		return err
	}
	if err := u.setupStack(); err != nil {
		return err
	}
	if u.OS.Init != nil {
		if err := u.OS.Init(u, args, env); err != nil {
			return err
		}
	}
	if u.Verbose {
		fmt.Fprintf(os.Stderr, "[entry @ 0x%x]\n", u.entry)
		dis, err := u.Disas(u.entry, 64)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
		} else {
			fmt.Fprintln(os.Stderr, dis)
		}
		sp, err := u.RegRead(u.arch.SP)
		if err != nil {
			return err
		}
		buf := make([]byte, u.StackBase+STACK_SIZE-sp)
		if err := u.MemReadInto(buf, sp); err != nil {
			return err
		}
		fmt.Fprintf(os.Stderr, "[stack @ 0x%x]\n", sp)
		fmt.Fprintf(os.Stderr, "%s\n", HexDump(sp, buf[:], u.arch))
	}
	if u.Verbose || u.TraceReg {
		u.status.Changes().Print(true, false)
	}
	if u.Verbose {
		fmt.Fprintln(os.Stderr, "=====================================")
		fmt.Fprintln(os.Stderr, "==== Program output begins here. ====")
		fmt.Fprintln(os.Stderr, "=====================================")
	}
	if u.TraceReg || u.TraceExec {
		sp, _ := u.RegRead(u.arch.SP)
		u.stacktrace.Update(u.entry, sp)
	}
	err := u.Unicorn.Start(u.entry, 0xffffffffffffffff)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Registers:")
		u.status.Changes().Print(true, false)
		fmt.Fprintln(os.Stderr, "Stacktrace:")
		u.stacktrace.Print(u)
	}
	return err
}

func (u *Usercorn) Loader() models.Loader {
	return u.loader
}

func (u *Usercorn) InterpBase() uint64 {
	// points to interpreter base or 0
	return u.interpBase
}

func (u *Usercorn) Base() uint64 {
	// points to program base
	return u.base
}

func (u *Usercorn) Entry() uint64 {
	// points to effective program entry: either an interpreter or the binary
	return u.entry
}

func (u *Usercorn) BinEntry() uint64 {
	// points to binary entry, even if an interpreter is used
	return u.binEntry
}

func (u *Usercorn) PosixInit(args, env []string, auxv []byte) error {
	// end marker
	if err := u.Push(0); err != nil {
		return err
	}
	// auxv
	if err := u.PushBytes(auxv); err != nil {
		return err
	}
	// envp
	envp, err := u.pushStrings(env...)
	if err != nil {
		return err
	}
	if err := u.pushAddrs(envp); err != nil {
		return err
	}
	// argv
	argv, err := u.pushStrings(args...)
	if err != nil {
		return err
	}
	if err := u.pushAddrs(argv); err != nil {
		return err
	}
	// argc
	return u.Push(uint64(len(args)))
}

func (u *Usercorn) PrefixPath(path string, force bool) string {
	if filepath.IsAbs(path) && u.LoadPrefix != "" {
		target := filepath.Join(u.LoadPrefix, path)
		_, err := os.Stat(target)
		exists := !os.IsNotExist(err)
		if force || exists {
			return target
		}
	}
	return path
}

func (u *Usercorn) Symbolicate(addr uint64) (string, error) {
	var symbolicate = func(addr uint64, symbols []models.Symbol) (result models.Symbol, distance uint64) {
		if len(symbols) == 0 {
			return
		}
		nearest := make(map[uint64][]models.Symbol)
		var min int64 = -1
		for _, sym := range symbols {
			dist := int64(addr - sym.Start)
			if dist > 0 && (sym.Start+uint64(dist) <= sym.End || sym.End == 0) && sym.Name != "" {
				if dist < min || min == -1 {
					min = dist
				}
				nearest[uint64(dist)] = append(nearest[uint64(dist)], sym)
			}
		}
		if len(nearest) > 0 {
			sym := nearest[uint64(min)][0]
			return sym, uint64(min)
		}
		return
	}
	symbols, _ := u.loader.Symbols()
	var interpSym []models.Symbol
	if u.interpLoader != nil {
		interpSym, _ = u.interpLoader.Symbols()
	}
	sym, sdist := symbolicate(addr-u.base, symbols)
	isym, idist := symbolicate(addr-u.interpBase, interpSym)
	if idist < sdist && isym.Name != "" || sym.Name == "" {
		sym = isym
		sdist = idist
	}
	if sym.Name != "" {
		return fmt.Sprintf("%s+0x%x", sym.Name, sdist), nil
	}
	return "", nil
}

func (u *Usercorn) Brk(addr uint64) (uint64, error) {
	// TODO: this is linux specific
	s := u.DataSegment
	if addr > 0 {
		u.MemMap(s.End, addr)
		s.End = addr
	}
	return s.End, nil
}

func (u *Usercorn) addHooks() error {
	if u.TraceExec || u.TraceReg {
		u.HookAdd(uc.HOOK_BLOCK, func(_ uc.Unicorn, addr uint64, size uint32) {
			sym, _ := u.Symbolicate(addr)
			if sym != "" {
				sym = " (" + sym + ")"
			}
			blockLine := fmt.Sprintf("| block%s @0x%x", sym, addr)
			if !u.TraceExec && u.TraceReg && u.deadlock == 0 {
				changes := u.status.Changes()
				if changes.Count() > 0 {
					fmt.Fprintln(os.Stderr, blockLine)
					changes.Print(true, true)
				}
			} else {
				fmt.Fprintln(os.Stderr, blockLine)
			}
			if sp, err := u.RegRead(u.arch.SP); err == nil {
				u.stacktrace.Update(addr, sp)
			}
			u.lastBlock = addr
		})
	}
	if u.TraceExec {
		u.HookAdd(uc.HOOK_CODE, func(_ uc.Unicorn, addr uint64, size uint32) {
			var changes *models.Changes
			if addr == u.lastCode || u.TraceReg && u.TraceExec {
				changes = u.status.Changes()
			}
			if u.TraceExec {
				if u.TraceReg {
					changes.Print(true, true)
				}
				dis, _ := u.Disas(addr, uint64(size))
				fmt.Fprintln(os.Stderr, dis)
			}
			if addr == u.lastCode {
				u.deadlock++
				if changes.Count() > 0 {
					if u.TraceReg {
						changes.Print(true, true)
					}
					u.deadlock = 0
				}
				if u.deadlock > 2 {
					sym, _ := u.Symbolicate(addr)
					if sym != "" {
						sym = " (" + sym + ")"
					}
					fmt.Fprintf(os.Stderr, "FATAL: deadlock detected at 0x%x%s\n", addr, sym)
					changes.Print(true, false)
					u.Stop()
				}
			} else {
				u.deadlock = 0
			}
			u.lastCode = addr
		})
	}
	if u.TraceMem {
		u.HookAdd(uc.HOOK_MEM_READ_WRITE, func(_ uc.Unicorn, access int, addr uint64, size int, value int64) {
			if access == uc.MEM_WRITE {
				fmt.Fprintf(os.Stderr, "MEM_WRITE")
			} else {
				fmt.Fprintf(os.Stderr, "MEM_READ")
			}
			fmt.Fprintf(os.Stderr, " 0x%x %d 0x%x\n", addr, size, value)
		})
	}
	u.HookAdd(uc.HOOK_MEM_INVALID, func(_ uc.Unicorn, access int, addr uint64, size int, value int64) bool {
		if access == uc.MEM_WRITE {
			fmt.Fprintf(os.Stderr, "invalid write")
		} else {
			fmt.Fprintf(os.Stderr, "invalid read")
		}
		fmt.Fprintf(os.Stderr, ": @0x%x, 0x%x = 0x%x\n", addr, size, value)
		return false
	})
	u.HookAdd(uc.HOOK_INTR, func(_ uc.Unicorn, intno uint32) {
		u.OS.Interrupt(u, intno)
	})
	return nil
}

func (u *Usercorn) mapBinary(l models.Loader) (interpBase, entry, base, realEntry uint64, err error) {
	var dynamic bool
	switch l.Type() {
	case loader.EXEC:
		dynamic = false
	case loader.DYN:
		dynamic = true
	default:
		err = errors.New("Unsupported file load type.")
		return
	}
	segments, err := l.Segments()
	if err != nil {
		return
	}
	// merge overlapping segments
	merged := make([]*models.Segment, 0, len(segments))
outer:
	for _, seg := range segments {
		addr, size := align(seg.Addr, seg.Size, true)
		s := &models.Segment{addr, addr + size}
		for _, s2 := range merged {
			if s2.Overlaps(s) {
				s2.Merge(s)
				continue outer
			}
		}
		merged = append(merged, s)
	}
	// map merged segments
	var loadBias uint64
	for _, seg := range merged {
		size := seg.End - seg.Start
		if dynamic && seg.Start == 0 && loadBias == 0 {
			loadBias, err = u.Mmap(0x1000000, size)
		} else {
			err = u.MemMap(loadBias+seg.Start, seg.End-seg.Start)
		}
		if err != nil {
			return
		}
	}
	// write segment memory
	for _, seg := range segments {
		if err = u.MemWrite(loadBias+seg.Addr, seg.Data); err != nil {
			return
		}
	}
	entry = loadBias + l.Entry()
	// load interpreter if present
	interp := l.Interp()
	if interp != "" {
		var bin models.Loader
		bin, err = loader.LoadFile(u.PrefixPath(interp, true))
		if err != nil {
			return
		}
		u.interpLoader = bin
		_, _, interpBias, interpEntry, err := u.mapBinary(bin)
		return interpBias, interpEntry, loadBias, entry, err
	} else {
		return 0, entry, loadBias, entry, nil
	}
}

func (u *Usercorn) setupStack() error {
	stack, err := u.Mmap(STACK_BASE, STACK_SIZE)
	if err != nil {
		return err
	}
	u.StackBase = stack
	if err := u.RegWrite(u.arch.SP, stack+STACK_SIZE); err != nil {
		return err
	}
	return nil
}

func (u *Usercorn) pushStrings(args ...string) ([]uint64, error) {
	// TODO: does anything case if these are actually on the stack?
	argvSize := 0
	for _, v := range args {
		argvSize += len(v) + 1
	}
	argvAddr, err := u.Mmap(0, uint64(argvSize))
	if err != nil {
		return nil, err
	}
	buf := make([]byte, argvSize)
	addrs := make([]uint64, 0, len(args)+1)
	var pos uint64
	for i := len(args) - 1; i >= 0; i-- {
		copy(buf[pos:], []byte(args[i]))
		addrs = append(addrs, argvAddr+pos)
		pos += uint64(len(args[i]) + 1)
	}
	u.MemWrite(argvAddr, buf)
	return addrs, nil
}

func (u *Usercorn) pushAddrs(addrs []uint64) error {
	if err := u.Push(0); err != nil {
		return err
	}
	for _, v := range addrs {
		if err := u.Push(v); err != nil {
			return err
		}
	}
	return nil
}

func (u *Usercorn) Syscall(num int, name string, getArgs func(n int) ([]uint64, error)) (uint64, error) {
	if name == "" {
		panic(fmt.Sprintf("Syscall missing: %d", num))
	}
	return syscalls.Call(u, num, name, getArgs, u.TraceSys)
}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"log"
	"os"

	"github.com/fatih/color"
)

var verbose bool

func SetVerbose(v bool) {
	verbose = v
}

func LogInfo(info ...interface{}) {
	color.Set(color.FgWhite)
	log.Printf("INFO: %s", fmt.Sprintln(info...))
	color.Unset()
}

func LogInfof(format string, info ...interface{}) {
	if verbose {
		color.Set(color.FgWhite)
		log.Printf("INFO: %s", fmt.Sprintf(format, info...))
		color.Unset()
	}
}

func LogAttention(info ...interface{}) {
	if verbose {
		color.Set(color.FgYellow)
		log.Printf("INFO: %s", fmt.Sprintln(info...))
		color.Unset()
	}
}

func LogAttentionf(format string, info ...interface{}) {
	if verbose {
		color.Set(color.FgYellow)
		log.Printf("INFO: %s", fmt.Sprintf(format, info...))
		color.Unset()
	}
}

func LogWarn(warn ...interface{}) {
	if verbose {
		color.Set(color.FgMagenta)
		log.Printf("WARN: %s", fmt.Sprintln(warn...))
		color.Unset()
	}
}

func LogWarnf(format string, warn ...interface{}) {
	if verbose {
		color.Set(color.FgMagenta)
		log.Printf("WARN: %s", fmt.Sprintf(format, warn...))
		color.Unset()
	}
}

func LogAlert(alert ...interface{}) {
	if verbose {
		color.Set(color.FgCyan)
		log.Printf("ALERT: %s", fmt.Sprintln(alert...))
		color.Unset()
	}
}

func LogAlertf(format string, alert ...interface{}) {
	if verbose {
		color.Set(color.FgCyan)
		log.Printf("ALERT: %s", fmt.Sprintf(format, alert...))
		color.Unset()
	}
}

func LogError(err ...interface{}) {
	color.Set(color.FgHiRed)
	log.Printf("ERROR: %s", fmt.Sprintln(err...))
	color.Unset()
}

func LogErrorf(format string, err ...interface{}) {
	color.Set(color.FgHiRed)
	log.Printf("ERROR: %s", fmt.Sprintf(format, err...))
	color.Unset()
}

func LogFatal(code int, err ...interface{}) {
	LogErrorf("FATAL: %s", fmt.Sprintln(err...))
	os.Exit(code)
}
<commit_msg>Spruced up log utility function with additional options for logging in colour.<commit_after>package util

import (
	"fmt"
	"log"
	"os"

	"github.com/fatih/color"
)

var verbose bool

func SetVerbose(v bool) {
	verbose = v
}

func LogColour(out string, attr color.Attribute) {
	if verbose {
		color.Set(attr)
		log.Printf(out)
		color.Unset()
	}
}

func LogGreen(out string) {
	LogColour(out, color.FgGreen)
}

func LogWhite(out string) {
	LogColour(out, color.FgWhite)
}

func LogMagenta(out string) {
	LogColour(out, color.FgMagenta)
}

func LogYellow(out string) {
	LogColour(out, color.FgYellow)
}

func LogCyan(out string) {
	LogColour(out, color.FgCyan)
}

func LogRed(out string) {
	LogColour(out, color.FgRed)
}

func LogRedBright(out string) {
	LogColour(out, color.FgHiRed)
}

func LogInfo(info ...interface{}) {
	LogWhite("INFO:  " + fmt.Sprintln(info...))
}

func LogInfof(format string, info ...interface{}) {
	LogWhite("INFO:  " + fmt.Sprintf(format, info...))
}

func LogOk(info ...interface{}) {
	LogGreen("OK:    " + fmt.Sprintln(info...))
}

func LogOkf(format string, info ...interface{}) {
	LogGreen("OK:    " + fmt.Sprintf(format, info...))
}

func LogAttention(info ...interface{}) {
	LogYellow("ATTN:  " + fmt.Sprintln(info...))
}

func LogAttentionf(format string, info ...interface{}) {
	LogYellow("ATTN:  " + fmt.Sprintf(format, info...))
}

func LogWarn(warn ...interface{}) {
	LogMagenta("WARN:  " + fmt.Sprintln(warn...))
}

func LogWarnf(format string, warn ...interface{}) {
	LogMagenta("WARN:  " + fmt.Sprintf(format, warn...))
}

func LogAlert(alert ...interface{}) {
	LogCyan("ALERT: " + fmt.Sprintln(alert...))
}

func LogAlertf(format string, alert ...interface{}) {
	LogCyan("ALERT: " + fmt.Sprintf(format, alert...))
}

func LogError(err ...interface{}) {
	orig := verbose
	verbose = true
	LogRed("ERROR: " + fmt.Sprintln(err...))
	verbose = orig
}

func LogErrorf(format string, err ...interface{}) {
	orig := verbose
	verbose = true
	LogRed("ERROR: " + fmt.Sprintf(format, err...))
	verbose = orig
}

func LogFatal(code int, err ...interface{}) {
	LogErrorf("FATAL: %s", fmt.Sprintln(err...))
	os.Exit(code)
}
<|endoftext|>


<commit_before>package cgolmnl

/*
#include <libmnl/libmnl.h>
*/
import "C"
import (
	"unsafe"
	"errors"
)

/*
 * attr.go
 */
func (attr *Nlattr) GetType() uint16 { return AttrGetType(attr) }
func (attr *Nlattr) GetLen() uint16 { return AttrGetLen(attr) }
func (attr *Nlattr) PayloadLen() uint16 { return AttrGetPayloadLen(attr) }
func (attr *Nlattr) Payload() unsafe.Pointer { return AttrGetPayload(attr) }
func (attr *Nlattr) PayloadBytes() []byte { return AttrGetPayloadBytes(attr) } // added
func (attr *Nlattr) Ok(size int) bool { return AttrOk(attr, size) }
func (attr *Nlattr) Next() *Nlattr { return AttrNext(attr) }
func (attr *Nlattr) TypeValid(max uint16) (int, error) { return AttrTypeValid(attr, max) }
func (attr *Nlattr) Validate(data_type AttrDataType) (int, error) { return AttrValidate(attr, data_type) }
func (attr *Nlattr) Validate2(data_type AttrDataType, exp_len Size_t) (int, error) { return AttrValidate2(attr, data_type, exp_len) }
func (nlh *Nlmsghdr) Parse(offset Size_t, cb MnlAttrCb, data interface{}) (int, error) { return AttrParse(nlh, offset, cb, data) }
func (attr *Nlattr) ParseNested(cb MnlAttrCb, data interface{}) (int, error) { return AttrParseNested(attr, cb, data) }
func (attr *Nlattr) U8() uint8 { return AttrGetU8(attr) }
func (attr *Nlattr) U16() uint16 { return AttrGetU16(attr) }
func (attr *Nlattr) U32() uint32 { return AttrGetU32(attr) }
func (attr *Nlattr) U64() uint64 { return AttrGetU64(attr) }
func (attr *Nlattr) Str() string { return AttrGetStr(attr) }
func (nlh *Nlmsghdr) Put(attr_type uint16, data []byte) { AttrPut(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutU8(attr_type uint16, data uint8) { AttrPutU8(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutU16(attr_type uint16, data uint16) { AttrPutU16(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutU32(attr_type uint16, data uint32) { AttrPutU32(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutU64(attr_type uint16, data uint64) { AttrPutU64(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutStr(attr_type uint16, data string) { AttrPutStr(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutStrz(attr_type uint16, data string) { AttrPutStrz(nlh, attr_type, data) }
func (nlh *Nlmsghdr) NestStart(attr_type uint16) *Nlattr { return AttrNestStart(nlh, attr_type) }
func (nlh *Nlmsghdr) PutCheck(buflen Size_t, attr_type uint16, data []byte) bool { return AttrPutCheck(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutU8Check(buflen Size_t, attr_type uint16, data uint8) bool { return AttrPutU8Check(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutU16Check(buflen Size_t, attr_type uint16, data uint16) bool { return AttrPutU16Check(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutU32Check(buflen Size_t, attr_type uint16, data uint32) bool { return AttrPutU32Check(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutU64Check(buflen Size_t, attr_type uint16, data uint64) bool { return AttrPutU64Check(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutStrCheck(buflen Size_t, attr_type uint16, data string) bool { return AttrPutStrCheck(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutStrzCheck(buflen Size_t, attr_type uint16, data string) bool { return AttrPutStrzCheck(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) NestEnd(start *Nlattr) { AttrNestEnd(nlh, start) }
func (nlh *Nlmsghdr) NestCancel(start *Nlattr) { AttrNestCancel(nlh, start) }

/*
 * nlmsg.go
 */
func (nlh *Nlmsghdr) PayloadLen() Size_t { return NlmsgGetPayloadLen(nlh) }
func (nlh *Nlmsghdr) PutExtraHeader(size Size_t) unsafe.Pointer { return NlmsgPutExtraHeader(nlh, size) }
func (nlh *Nlmsghdr) Payload() unsafe.Pointer { return NlmsgGetPayload(nlh) }
func (nlh *Nlmsghdr) PayloadOffset(offset Size_t) unsafe.Pointer { return NlmsgGetPayloadOffset(nlh, offset) }
func (nlh *Nlmsghdr) Ok(size int) bool { return NlmsgOk(nlh, size) }
func (nlh *Nlmsghdr) Next(size int) (*Nlmsghdr, int) { return NlmsgNext(nlh, size) }
func (nlh *Nlmsghdr) PayloadTail() unsafe.Pointer { return NlmsgGetPayloadTail(nlh) }
func (nlh *Nlmsghdr) SeqOk(seq uint32) bool { return NlmsgSeqOk(nlh, seq) }
func (nlh *Nlmsghdr) PortidOk(portid uint32) bool { return NlmsgPortidOk(nlh, portid) }
func (b *NlmsgBatchDescriptor) Stop() { NlmsgBatchStop(b) }
func (b *NlmsgBatchDescriptor) Next() bool { return NlmsgBatchNext(b) }
func (b *NlmsgBatchDescriptor) Reset() { NlmsgBatchReset(b) }
func (b *NlmsgBatchDescriptor) Size() Size_t { return NlmsgBatchSize(b) }
func (b *NlmsgBatchDescriptor) Head() unsafe.Pointer { return NlmsgBatchHead(b) }
func (b *NlmsgBatchDescriptor) HeadBytes() []byte { return NlmsgBatchHeadBytes(b) }
func (b *NlmsgBatchDescriptor) Current() unsafe.Pointer { return NlmsgBatchCurrent(b) }
func (b *NlmsgBatchDescriptor) IsEmpty() bool { return NlmsgBatchIsEmpty(b) }

// helper function
func NewNlmsghdr(size int) (*Nlmsghdr, error) {
	if size < int(MNL_NLMSG_HDRLEN) {
		return nil, errors.New("too short size")
	}
	b := make([]byte, size)
	return (*Nlmsghdr)(unsafe.Pointer(&b[0])), nil
}

func NlmsghdrPointer(b []byte) *Nlmsghdr {
	return (*Nlmsghdr)(unsafe.Pointer(&b[0]))
}

func (nlh *Nlmsghdr) PutHeader() {
	C.mnl_nlmsg_put_header(unsafe.Pointer(nlh))
}

func PutNewNlmsghdr(size int) (*Nlmsghdr, error) {
	nlh, err := NewNlmsghdr(size)
	if err != nil {
		return nil, err
	}
	nlh.PutHeader()
	return nlh, nil
}

/*
 * socket.go
 */
func (nl *SocketDescriptor) Fd() int { return SocketGetFd(nl) }
func (nl *SocketDescriptor) Portid() uint32 { return SocketGetPortid(nl) }
func (nl *SocketDescriptor) Bind(groups uint, pid Pid_t) error { return SocketBind(nl, groups, pid) }
func (nl *SocketDescriptor) Sendto(buf []byte) (Ssize_t, error) { return SocketSendto(nl, buf) }
func (nl *SocketDescriptor) SendNlmsg(nlh *Nlmsghdr) (Ssize_t, error) { return SocketSendNlmsg(nl, nlh) }
func (nl *SocketDescriptor) Recvfrom(buf []byte) (Ssize_t, error) { return SocketRecvfrom(nl, buf) }
func (nl *SocketDescriptor) Close() error { return SocketClose(nl) }
func (nl *SocketDescriptor) Setsockopt(optype int, buf []byte) error { return SocketSetsockopt(nl, optype, buf) }
func (nl *SocketDescriptor) Sockopt(optype int, size Socklen_t) ([]byte, error) { return SocketGetsockopt(nl, optype, size) }

<commit_msg>attr: add constructor<commit_after>package cgolmnl

/*
#include <libmnl/libmnl.h>
*/
import "C"
import (
	"unsafe"
	"errors"
)

/*
 * attr.go
 */
func (attr *Nlattr) GetType() uint16 { return AttrGetType(attr) }
func (attr *Nlattr) GetLen() uint16 { return AttrGetLen(attr) }
func (attr *Nlattr) PayloadLen() uint16 { return AttrGetPayloadLen(attr) }
func (attr *Nlattr) Payload() unsafe.Pointer { return AttrGetPayload(attr) }
func (attr *Nlattr) PayloadBytes() []byte { return AttrGetPayloadBytes(attr) } // added
func (attr *Nlattr) Ok(size int) bool { return AttrOk(attr, size) }
func (attr *Nlattr) Next() *Nlattr { return AttrNext(attr) }
func (attr *Nlattr) TypeValid(max uint16) (int, error) { return AttrTypeValid(attr, max) }
func (attr *Nlattr) Validate(data_type AttrDataType) (int, error) { return AttrValidate(attr, data_type) }
func (attr *Nlattr) Validate2(data_type AttrDataType, exp_len Size_t) (int, error) { return AttrValidate2(attr, data_type, exp_len) }
func (nlh *Nlmsghdr) Parse(offset Size_t, cb MnlAttrCb, data interface{}) (int, error) { return AttrParse(nlh, offset, cb, data) }
func (attr *Nlattr) ParseNested(cb MnlAttrCb, data interface{}) (int, error) { return AttrParseNested(attr, cb, data) }
func (attr *Nlattr) U8() uint8 { return AttrGetU8(attr) }
func (attr *Nlattr) U16() uint16 { return AttrGetU16(attr) }
func (attr *Nlattr) U32() uint32 { return AttrGetU32(attr) }
func (attr *Nlattr) U64() uint64 { return AttrGetU64(attr) }
func (attr *Nlattr) Str() string { return AttrGetStr(attr) }
func (nlh *Nlmsghdr) Put(attr_type uint16, data []byte) { AttrPut(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutU8(attr_type uint16, data uint8) { AttrPutU8(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutU16(attr_type uint16, data uint16) { AttrPutU16(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutU32(attr_type uint16, data uint32) { AttrPutU32(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutU64(attr_type uint16, data uint64) { AttrPutU64(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutStr(attr_type uint16, data string) { AttrPutStr(nlh, attr_type, data) }
func (nlh *Nlmsghdr) PutStrz(attr_type uint16, data string) { AttrPutStrz(nlh, attr_type, data) }
func (nlh *Nlmsghdr) NestStart(attr_type uint16) *Nlattr { return AttrNestStart(nlh, attr_type) }
func (nlh *Nlmsghdr) PutCheck(buflen Size_t, attr_type uint16, data []byte) bool { return AttrPutCheck(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutU8Check(buflen Size_t, attr_type uint16, data uint8) bool { return AttrPutU8Check(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutU16Check(buflen Size_t, attr_type uint16, data uint16) bool { return AttrPutU16Check(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutU32Check(buflen Size_t, attr_type uint16, data uint32) bool { return AttrPutU32Check(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutU64Check(buflen Size_t, attr_type uint16, data uint64) bool { return AttrPutU64Check(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutStrCheck(buflen Size_t, attr_type uint16, data string) bool { return AttrPutStrCheck(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) PutStrzCheck(buflen Size_t, attr_type uint16, data string) bool { return AttrPutStrzCheck(nlh, buflen, attr_type, data) }
func (nlh *Nlmsghdr) NestEnd(start *Nlattr) { AttrNestEnd(nlh, start) }
func (nlh *Nlmsghdr) NestCancel(start *Nlattr) { AttrNestCancel(nlh, start) }

// helper function
func NewNlattr(size int) (*Nlattr, error) {
	if size < SizeofNlattr {
		return nil, errors.New("too short size")
	}
	b := make([]byte, size)
	return (*Nlattr)(unsafe.Pointer(&b[0])), nil
}

func NlattrPointer(b []byte) *Nlattr {
	// ???: check buf len
	//      len(b) >= SizeofNlattr
	//      nla.len <= len(b)
	return (*Nlattr)(unsafe.Pointer(&b[0]))
}

/*
 * nlmsg.go
 */
func (nlh *Nlmsghdr) PayloadLen() Size_t { return NlmsgGetPayloadLen(nlh) }
func (nlh *Nlmsghdr) PutExtraHeader(size Size_t) unsafe.Pointer { return NlmsgPutExtraHeader(nlh, size) }
func (nlh *Nlmsghdr) Payload() unsafe.Pointer { return NlmsgGetPayload(nlh) }
func (nlh *Nlmsghdr) PayloadOffset(offset Size_t) unsafe.Pointer { return NlmsgGetPayloadOffset(nlh, offset) }
func (nlh *Nlmsghdr) Ok(size int) bool { return NlmsgOk(nlh, size) }
func (nlh *Nlmsghdr) Next(size int) (*Nlmsghdr, int) { return NlmsgNext(nlh, size) }
func (nlh *Nlmsghdr) PayloadTail() unsafe.Pointer { return NlmsgGetPayloadTail(nlh) }
func (nlh *Nlmsghdr) SeqOk(seq uint32) bool { return NlmsgSeqOk(nlh, seq) }
func (nlh *Nlmsghdr) PortidOk(portid uint32) bool { return NlmsgPortidOk(nlh, portid) }
func (b *NlmsgBatchDescriptor) Stop() { NlmsgBatchStop(b) }
func (b *NlmsgBatchDescriptor) Next() bool { return NlmsgBatchNext(b) }
func (b *NlmsgBatchDescriptor) Reset() { NlmsgBatchReset(b) }
func (b *NlmsgBatchDescriptor) Size() Size_t { return NlmsgBatchSize(b) }
func (b *NlmsgBatchDescriptor) Head() unsafe.Pointer { return NlmsgBatchHead(b) }
func (b *NlmsgBatchDescriptor) HeadBytes() []byte { return NlmsgBatchHeadBytes(b) }
func (b *NlmsgBatchDescriptor) Current() unsafe.Pointer { return NlmsgBatchCurrent(b) }
func (b *NlmsgBatchDescriptor) IsEmpty() bool { return NlmsgBatchIsEmpty(b) }

// helper function
func NewNlmsghdr(size int) (*Nlmsghdr, error) {
	if size < int(MNL_NLMSG_HDRLEN) {
		return nil, errors.New("too short size")
	}
	b := make([]byte, size)
	return (*Nlmsghdr)(unsafe.Pointer(&b[0])), nil
}

func NlmsghdrPointer(b []byte) *Nlmsghdr {
	return (*Nlmsghdr)(unsafe.Pointer(&b[0]))
}

func (nlh *Nlmsghdr) PutHeader() {
	// ???: check buf len
	//      len(b) >= SizeofNlmsghdr
	//      nlh.len <= len(buf)
	C.mnl_nlmsg_put_header(unsafe.Pointer(nlh))
}

func PutNewNlmsghdr(size int) (*Nlmsghdr, error) {
	nlh, err := NewNlmsghdr(size)
	if err != nil {
		return nil, err
	}
	nlh.PutHeader()
	return nlh, nil
}

/*
 * socket.go
 */
func (nl *SocketDescriptor) Fd() int { return SocketGetFd(nl) }
func (nl *SocketDescriptor) Portid() uint32 { return SocketGetPortid(nl) }
func (nl *SocketDescriptor) Bind(groups uint, pid Pid_t) error { return SocketBind(nl, groups, pid) }
func (nl *SocketDescriptor) Sendto(buf []byte) (Ssize_t, error) { return SocketSendto(nl, buf) }
func (nl *SocketDescriptor) SendNlmsg(nlh *Nlmsghdr) (Ssize_t, error) { return SocketSendNlmsg(nl, nlh) }
func (nl *SocketDescriptor) Recvfrom(buf []byte) (Ssize_t, error) { return SocketRecvfrom(nl, buf) }
func (nl *SocketDescriptor) Close() error { return SocketClose(nl) }
func (nl *SocketDescriptor) Setsockopt(optype int, buf []byte) error { return SocketSetsockopt(nl, optype, buf) }
func (nl *SocketDescriptor) Sockopt(optype int, size Socklen_t) ([]byte, error) { return SocketGetsockopt(nl, optype, size) }

<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/goadesign/goa/goagen/codegen"
	"github.com/goadesign/goa/goagen/meta"
	"github.com/goadesign/goa/goagen/utils"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	// These are packages required by the generated code but not by goagen.
	// We list them here so that `go get` picks them up.
	_ "gopkg.in/yaml.v2"
)

func main() {
	var (
		files            []string
		err              error
		terminatedByUser bool
	)

	// rootCmd is the base command used when goagen is called with no argument.
	rootCmd := &cobra.Command{
		Use:   "goagen",
		Short: "goa code generation tool",
		Long: `The goagen tool generates artifacts from a goa service design package.

Each command supported by the tool produces a specific type of artifacts. For example
the "app" command generates the code that supports the service controllers.

The "bootstrap" command runs the "app", "main", "client" and "swagger" commands generating the
controllers supporting code and main skeleton code (if not already present) as well as a client
package and tool and the Swagger specification for the API.
`}
	var (
		cwd, designPkg string
		debug          bool
	)
	cwd, err = os.Getwd()
	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		os.Exit(1)
	}
	rootCmd.PersistentFlags().StringVarP(&cwd, "out", "o", cwd, "output directory")
	rootCmd.PersistentFlags().StringVarP(&designPkg, "design", "d", "", "design package import path")
	rootCmd.PersistentFlags().BoolVar(&debug, "debug", false, "enable debug mode, does not cleanup temporary files.")

	// appCmd implements the "app" command.
	var (
		pkg    string
		notest bool
	)
	appCmd := &cobra.Command{
		Use:   "app",
		Short: "Generate application code",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genapp", c) },
	}
	appCmd.Flags().StringVar(&pkg, "pkg", "app", "Name of generated Go package containing controllers supporting code (contexts, media types, user types etc.)")
	appCmd.Flags().BoolVar(&notest, "notest", false, "Prevent generation of test helpers")
	rootCmd.AddCommand(appCmd)

	// mainCmd implements the "main" command.
	var (
		force bool
	)
	mainCmd := &cobra.Command{
		Use:   "main",
		Short: "Generate application scaffolding",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genmain", c) },
	}
	mainCmd.Flags().BoolVar(&force, "force", false, "overwrite existing files")
	rootCmd.AddCommand(mainCmd)

	// clientCmd implements the "client" command.
	var (
		toolDir, tool string
		notool        bool
	)
	clientCmd := &cobra.Command{
		Use:   "client",
		Short: "Generate client package and tool",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genclient", c) },
	}
	clientCmd.Flags().StringVar(&pkg, "pkg", "client", "Name of generated client Go package")
	clientCmd.Flags().StringVar(&toolDir, "tooldir", "tool", "Name of generated tool directory")
	clientCmd.Flags().StringVar(&tool, "tool", "[API-name]-cli", "Name of generated tool")
	clientCmd.Flags().BoolVar(&notool, "notool", false, "Prevent generation of cli tool")
	rootCmd.AddCommand(clientCmd)

	// swaggerCmd implements the "swagger" command.
	swaggerCmd := &cobra.Command{
		Use:   "swagger",
		Short: "Generate Swagger",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genswagger", c) },
	}
	rootCmd.AddCommand(swaggerCmd)

	// jsCmd implements the "js" command.
	var (
		timeout      = time.Duration(20) * time.Second
		scheme, host string
		noexample    bool
	)
	jsCmd := &cobra.Command{
		Use:   "js",
		Short: "Generate JavaScript client",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genjs", c) },
	}
	jsCmd.Flags().DurationVar(&timeout, "timeout", timeout, `the duration before the request times out.`)
	jsCmd.Flags().StringVar(&scheme, "scheme", "", `the URL scheme used to make requests to the API, defaults to the scheme defined in the API design if any.`)
	jsCmd.Flags().StringVar(&host, "host", "", `the API hostname, defaults to the hostname defined in the API design if any`)
	jsCmd.Flags().BoolVar(&noexample, "noexample", false, `Skip generation of example HTML and controller`)
	rootCmd.AddCommand(jsCmd)

	// schemaCmd implements the "schema" command.
	schemaCmd := &cobra.Command{
		Use:   "schema",
		Short: "Generate JSON Schema",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genschema", c) },
	}
	rootCmd.AddCommand(schemaCmd)

	// genCmd implements the "gen" command.
	var (
		pkgPath string
	)
	genCmd := &cobra.Command{
		Use:   "gen",
		Short: "Run third-party generator",
		Run:   func(c *cobra.Command, _ []string) { files, err = runGen(c) },
	}
	genCmd.Flags().StringVar(&pkgPath, "pkg-path", "", "Package import path of generator. The package must implement the Generate global function.")
	rootCmd.AddCommand(genCmd)

	// boostrapCmd implements the "bootstrap" command.
	bootCmd := &cobra.Command{
		Use:   "bootstrap",
		Short: `Equivalent to running the "app", "main", "client" and "swagger" commands.`,
		Run: func(c *cobra.Command, a []string) {
			appCmd.Run(c, a)
			if err != nil {
				return
			}
			prev := files

			mainCmd.Run(c, a)
			if err != nil {
				return
			}
			prev = append(prev, files...)

			clientCmd.Run(c, a)
			if err != nil {
				return
			}
			prev = append(prev, files...)

			swaggerCmd.Run(c, a)
			files = append(prev, files...)
		},
	}
	rootCmd.AddCommand(bootCmd)

	// Now proceed with code generation
	cleanup := func() {
		for _, f := range files {
			os.RemoveAll(f)
		}
	}

	go utils.Catch(nil, func() {
		terminatedByUser = true
	})

	rootCmd.Execute()

	if terminatedByUser {
		cleanup()
		return
	}

	if err != nil {
		cleanup()
		fmt.Fprintln(os.Stderr, err.Error())
		os.Exit(1)
	}

	rels := make([]string, len(files))
	cd, _ := os.Getwd()
	for i, f := range files {
		r, err := filepath.Rel(cd, f)
		if err == nil {
			rels[i] = r
		} else {
			rels[i] = f
		}
	}
	fmt.Println(strings.Join(rels, "\n"))
}

func run(pkg string, c *cobra.Command) ([]string, error) {
	pkgPath := fmt.Sprintf("github.com/goadesign/goa/goagen/gen_%s", pkg[3:])
	pkgSrcPath, err := codegen.PackageSourcePath(pkgPath)
	if err != nil {
		return nil, fmt.Errorf("invalid plugin package import path: %s", err)
	}
	pkgName, err := codegen.PackageName(pkgSrcPath)
	if err != nil {
		return nil, fmt.Errorf("invalid package import path: %s", err)
	}
	return generate(pkgName, pkgPath, c)
}

func runGen(c *cobra.Command) ([]string, error) {
	pkgPath := c.Flag("pkg-path").Value.String()
	pkgSrcPath, err := codegen.PackageSourcePath(pkgPath)
	if err != nil {
		return nil, fmt.Errorf("invalid plugin package import path: %s", err)
	}
	pkgName, err := codegen.PackageName(pkgSrcPath)
	if err != nil {
		return nil, fmt.Errorf("invalid plugin package import path: %s", err)
	}
	return generate(pkgName, pkgPath, c)
}

func generate(pkgName, pkgPath string, c *cobra.Command) ([]string, error) {
	m := make(map[string]string)
	c.Flags().Visit(func(f *pflag.Flag) {
		if f.Name != "pkg-path" {
			m[f.Name] = f.Value.String()
		}
	})
	if _, ok := m["out"]; !ok {
		m["out"] = c.Flag("out").DefValue
	}
	gen, err := meta.NewGenerator(
		pkgName+".Generate",
		[]*codegen.ImportSpec{codegen.SimpleImport(pkgPath)},
		m,
	)
	if err != nil {
		return nil, err
	}
	return gen.Generate()
}
<commit_msg>Add `goagen commands` command.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/goadesign/goa/goagen/codegen"
	"github.com/goadesign/goa/goagen/meta"
	"github.com/goadesign/goa/goagen/utils"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	// These are packages required by the generated code but not by goagen.
	// We list them here so that `go get` picks them up.
	_ "gopkg.in/yaml.v2"
)

func main() {
	var (
		files            []string
		err              error
		terminatedByUser bool
	)

	// rootCmd is the base command used when goagen is called with no argument.
	rootCmd := &cobra.Command{
		Use:   "goagen",
		Short: "goa code generation tool",
		Long: `The goagen tool generates artifacts from a goa service design package.

Each command supported by the tool produces a specific type of artifacts. For example
the "app" command generates the code that supports the service controllers.

The "bootstrap" command runs the "app", "main", "client" and "swagger" commands generating the
controllers supporting code and main skeleton code (if not already present) as well as a client
package and tool and the Swagger specification for the API.
`}
	var (
		cwd, designPkg string
		debug          bool
	)
	cwd, err = os.Getwd()
	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		os.Exit(1)
	}
	rootCmd.PersistentFlags().StringVarP(&cwd, "out", "o", cwd, "output directory")
	rootCmd.PersistentFlags().StringVarP(&designPkg, "design", "d", "", "design package import path")
	rootCmd.PersistentFlags().BoolVar(&debug, "debug", false, "enable debug mode, does not cleanup temporary files.")

	// appCmd implements the "app" command.
	var (
		pkg    string
		notest bool
	)
	appCmd := &cobra.Command{
		Use:   "app",
		Short: "Generate application code",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genapp", c) },
	}
	appCmd.Flags().StringVar(&pkg, "pkg", "app", "Name of generated Go package containing controllers supporting code (contexts, media types, user types etc.)")
	appCmd.Flags().BoolVar(&notest, "notest", false, "Prevent generation of test helpers")
	rootCmd.AddCommand(appCmd)

	// mainCmd implements the "main" command.
	var (
		force bool
	)
	mainCmd := &cobra.Command{
		Use:   "main",
		Short: "Generate application scaffolding",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genmain", c) },
	}
	mainCmd.Flags().BoolVar(&force, "force", false, "overwrite existing files")
	rootCmd.AddCommand(mainCmd)

	// clientCmd implements the "client" command.
	var (
		toolDir, tool string
		notool        bool
	)
	clientCmd := &cobra.Command{
		Use:   "client",
		Short: "Generate client package and tool",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genclient", c) },
	}
	clientCmd.Flags().StringVar(&pkg, "pkg", "client", "Name of generated client Go package")
	clientCmd.Flags().StringVar(&toolDir, "tooldir", "tool", "Name of generated tool directory")
	clientCmd.Flags().StringVar(&tool, "tool", "[API-name]-cli", "Name of generated tool")
	clientCmd.Flags().BoolVar(&notool, "notool", false, "Prevent generation of cli tool")
	rootCmd.AddCommand(clientCmd)

	// swaggerCmd implements the "swagger" command.
	swaggerCmd := &cobra.Command{
		Use:   "swagger",
		Short: "Generate Swagger",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genswagger", c) },
	}
	rootCmd.AddCommand(swaggerCmd)

	// jsCmd implements the "js" command.
	var (
		timeout      = time.Duration(20) * time.Second
		scheme, host string
		noexample    bool
	)
	jsCmd := &cobra.Command{
		Use:   "js",
		Short: "Generate JavaScript client",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genjs", c) },
	}
	jsCmd.Flags().DurationVar(&timeout, "timeout", timeout, `the duration before the request times out.`)
	jsCmd.Flags().StringVar(&scheme, "scheme", "", `the URL scheme used to make requests to the API, defaults to the scheme defined in the API design if any.`)
	jsCmd.Flags().StringVar(&host, "host", "", `the API hostname, defaults to the hostname defined in the API design if any`)
	jsCmd.Flags().BoolVar(&noexample, "noexample", false, `Skip generation of example HTML and controller`)
	rootCmd.AddCommand(jsCmd)

	// schemaCmd implements the "schema" command.
	schemaCmd := &cobra.Command{
		Use:   "schema",
		Short: "Generate JSON Schema",
		Run:   func(c *cobra.Command, _ []string) { files, err = run("genschema", c) },
	}
	rootCmd.AddCommand(schemaCmd)

	// genCmd implements the "gen" command.
	var (
		pkgPath string
	)
	genCmd := &cobra.Command{
		Use:   "gen",
		Short: "Run third-party generator",
		Run:   func(c *cobra.Command, _ []string) { files, err = runGen(c) },
	}
	genCmd.Flags().StringVar(&pkgPath, "pkg-path", "", "Package import path of generator. The package must implement the Generate global function.")
	rootCmd.AddCommand(genCmd)

	// boostrapCmd implements the "bootstrap" command.
	bootCmd := &cobra.Command{
		Use:   "bootstrap",
		Short: `Equivalent to running the "app", "main", "client" and "swagger" commands.`,
		Run: func(c *cobra.Command, a []string) {
			appCmd.Run(c, a)
			if err != nil {
				return
			}
			prev := files

			mainCmd.Run(c, a)
			if err != nil {
				return
			}
			prev = append(prev, files...)

			clientCmd.Run(c, a)
			if err != nil {
				return
			}
			prev = append(prev, files...)

			swaggerCmd.Run(c, a)
			files = append(prev, files...)
		},
	}
	bootCmd.Flags().AddFlagSet(appCmd.Flags())
	bootCmd.Flags().AddFlagSet(mainCmd.Flags())
	bootCmd.Flags().AddFlagSet(clientCmd.Flags())
	bootCmd.Flags().AddFlagSet(swaggerCmd.Flags())
	rootCmd.AddCommand(bootCmd)

	// cmdsCmd implements the commands command
	// It lists all the commands and flags in JSON to enable shell integrations.
	cmdsCmd := &cobra.Command{
		Use:   "commands",
		Short: "Lists all commands and flags in JSON",
		Run:   func(c *cobra.Command, _ []string) { runCommands(rootCmd) },
	}
	rootCmd.AddCommand(cmdsCmd)

	// Now proceed with code generation
	cleanup := func() {
		for _, f := range files {
			os.RemoveAll(f)
		}
	}

	go utils.Catch(nil, func() {
		terminatedByUser = true
	})

	rootCmd.Execute()

	if terminatedByUser {
		cleanup()
		return
	}

	if err != nil {
		cleanup()
		fmt.Fprintln(os.Stderr, err.Error())
		os.Exit(1)
	}

	rels := make([]string, len(files))
	cd, _ := os.Getwd()
	for i, f := range files {
		r, err := filepath.Rel(cd, f)
		if err == nil {
			rels[i] = r
		} else {
			rels[i] = f
		}
	}
	fmt.Println(strings.Join(rels, "\n"))
}

func run(pkg string, c *cobra.Command) ([]string, error) {
	pkgPath := fmt.Sprintf("github.com/goadesign/goa/goagen/gen_%s", pkg[3:])
	pkgSrcPath, err := codegen.PackageSourcePath(pkgPath)
	if err != nil {
		return nil, fmt.Errorf("invalid plugin package import path: %s", err)
	}
	pkgName, err := codegen.PackageName(pkgSrcPath)
	if err != nil {
		return nil, fmt.Errorf("invalid package import path: %s", err)
	}
	return generate(pkgName, pkgPath, c)
}

func runGen(c *cobra.Command) ([]string, error) {
	pkgPath := c.Flag("pkg-path").Value.String()
	pkgSrcPath, err := codegen.PackageSourcePath(pkgPath)
	if err != nil {
		return nil, fmt.Errorf("invalid plugin package import path: %s", err)
	}
	pkgName, err := codegen.PackageName(pkgSrcPath)
	if err != nil {
		return nil, fmt.Errorf("invalid plugin package import path: %s", err)
	}
	return generate(pkgName, pkgPath, c)
}

func generate(pkgName, pkgPath string, c *cobra.Command) ([]string, error) {
	m := make(map[string]string)
	c.Flags().Visit(func(f *pflag.Flag) {
		if f.Name != "pkg-path" {
			m[f.Name] = f.Value.String()
		}
	})
	if _, ok := m["out"]; !ok {
		m["out"] = c.Flag("out").DefValue
	}
	gen, err := meta.NewGenerator(
		pkgName+".Generate",
		[]*codegen.ImportSpec{codegen.SimpleImport(pkgPath)},
		m,
	)
	if err != nil {
		return nil, err
	}
	return gen.Generate()
}

type (
	rootCommand struct {
		Name     string     `json:"name"`
		Commands []*command `json:"commands"`
		Flags    []*flag    `json:"flags"`
	}

	flag struct {
		Long        string `json:"long,omitempty"`
		Short       string `json:"short,omitempty"`
		Description string `json:"description,omitempty"`
		Argument    string `json:"argument,omitempty"`
		Required    bool   `json:"required"`
	}

	command struct {
		Name  string  `json:"name"`
		Flags []*flag `json:"flags,omitempty"`
	}
)

func runCommands(root *cobra.Command) {
	var (
		gblFlags []*flag
		cmds     []*command
	)
	root.Flags().VisitAll(func(fl *pflag.Flag) {
		gblFlags = append(gblFlags, flagJSON(fl))
	})
	cmds = make([]*command, len(root.Commands())-2)
	j := 0
	for _, cm := range root.Commands() {
		if cm.Name() == "help" || cm.Name() == "commands" {
			continue
		}
		cmds[j] = cmdJSON(cm, gblFlags)
		j++
	}
	rc := rootCommand{os.Args[0], cmds, gblFlags}
	b, _ := json.MarshalIndent(rc, "", "    ")
	fmt.Println(string(b))
}

// Lots of assumptions in here, it's OK for what we are doing
// Remember to update as goagen commands and flags evolve
//
// The flag argument values use variable names that cary semantic:
// $DIR for file system directories, $DESIGN_PKG for import path to Go goa design Go packages, $PKG
// for import path to any Go package.
func flagJSON(fl *pflag.Flag) *flag {
	f := &flag{Long: fl.Name, Short: fl.Shorthand, Description: fl.Usage}
	f.Required = fl.Name == "pkg-path" || fl.Name == "design"
	switch fl.Name {
	case "out":
		f.Argument = "$DIR"
	case "design":
		f.Argument = "$DESIGN_PKG"
	case "pkg-path":
		f.Argument = "$PKG"
	}
	return f
}

func cmdJSON(cm *cobra.Command, flags []*flag) *command {
	res := make([]*flag, len(flags))
	for i, fl := range flags {
		f := *fl
		res[i] = &f
	}
	cm.Flags().VisitAll(func(fl *pflag.Flag) {
		res = append(res, flagJSON(fl))
	})
	return &command{cm.Name(), res}
}
<|endoftext|>


<commit_before>package gqt_test

import (
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/guardian/gqt/runner"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
)

var _ = Describe("graph flags", func() {
	var (
		client               *runner.RunningGarden
		layersPath           string
		diffPath             string
		mntPath              string
		nonDefaultRootfsPath string
		persistentImages     []string
	)

	numLayersInGraph := func() int {
		layerFiles, err := ioutil.ReadDir(layersPath)
		Expect(err).ToNot(HaveOccurred())
		diffFiles, err := ioutil.ReadDir(diffPath)
		Expect(err).ToNot(HaveOccurred())
		mntFiles, err := ioutil.ReadDir(mntPath)
		Expect(err).ToNot(HaveOccurred())

		numLayerFiles := len(layerFiles)
		Expect(numLayerFiles).To(Equal(len(diffFiles)))
		Expect(numLayerFiles).To(Equal(len(mntFiles)))
		return numLayerFiles
	}

	expectLayerCountAfterGraphCleanupToBe := func(layerCount int) {
		nonPersistantRootfsContainer, err := client.Create(garden.ContainerSpec{
			RootFSPath: nonDefaultRootfsPath,
		})
		Expect(err).ToNot(HaveOccurred())
		Expect(client.Destroy(nonPersistantRootfsContainer.Handle())).To(Succeed())
		Expect(numLayersInGraph()).To(Equal(layerCount + 2)) // +2 for the layers created for the nondefaultrootfs container
	}

	peaDoesNotLeaveDebrisInGraphOnExit := func(imageURI string) {
		var (
			pea                             garden.Process
			numLayersAfterContainerCreation int
		)

		JustBeforeEach(func() {
			imageURI := imageURI
			container, err := client.Create(garden.ContainerSpec{
				RootFSPath: imageURI,
			})
			Expect(err).ToNot(HaveOccurred())

			numLayersAfterContainerCreation = numLayersInGraph()

			pea, err = container.Run(garden.ProcessSpec{
				Path:  "/bin/sh",
				Args:  []string{"-c", "exit 0"},
				Image: garden.ImageRef{URI: imageURI},
			}, garden.ProcessIO{})
			Expect(err).NotTo(HaveOccurred())
		})

		It("does not leave debris in the graph on exit", func() {
			Expect(pea.Wait()).To(Equal(0))

			numLayersAfterPeaExits := numLayersInGraph()
			Expect(numLayersAfterPeaExits).To(Equal(numLayersAfterContainerCreation))
		})
	}

	BeforeEach(func() {
		var err error
		nonDefaultRootfsPath, err = ioutil.TempDir("", "tmpRootfs")
		Expect(err).ToNot(HaveOccurred())
		// temporary workaround as runc expects a /tmp dir to exist in the container rootfs
		err = os.Mkdir(filepath.Join(nonDefaultRootfsPath, "tmp"), 0700)
		Expect(err).ToNot(HaveOccurred())
	})

	JustBeforeEach(func() {
		config.PersistentImages = persistentImages
		config = resetImagePluginConfig()
		client = runner.Start(config)

		layersPath = path.Join(*client.GraphDir, "aufs", "layers")
		diffPath = path.Join(*client.GraphDir, "aufs", "diff")
		mntPath = path.Join(*client.GraphDir, "aufs", "mnt")
	})

	AfterEach(func() {
		Expect(os.RemoveAll(nonDefaultRootfsPath)).To(Succeed())
		Expect(client.DestroyAndStop()).To(Succeed())
	})

	Describe("--graph-cleanup-threshold-in-megabytes", func() {
		JustBeforeEach(func() {
			Expect(numLayersInGraph()).To(Equal(0))
			container, err := client.Create(garden.ContainerSpec{
				RootFSPath: "docker:///busybox",
			})
			Expect(err).ToNot(HaveOccurred())
			Expect(client.Destroy(container.Handle())).To(Succeed())
		})

		Context("when the graph cleanup threshold is set to -1", func() {
			BeforeEach(func() {
				config.GraphCleanupThresholdMB = intptr(-1)
			})

			It("does NOT clean up the graph directory on create", func() {
				initialNumberOfLayers := numLayersInGraph()
				anotherContainer, err := client.Create(garden.ContainerSpec{})
				Expect(err).ToNot(HaveOccurred())

				Expect(numLayersInGraph()).To(BeNumerically(">", initialNumberOfLayers), "after creation, should NOT have deleted anything")
				Expect(client.Destroy(anotherContainer.Handle())).To(Succeed())
			})

			Context("and we run a pea with the same image as the container", func() {
				peaDoesNotLeaveDebrisInGraphOnExit("docker:///busybox")
			})

			Context("and we run a pea with another image", func() {
				peaDoesNotLeaveDebrisInGraphOnExit("docker:///alpine")
			})
		})

		Context("when the graph cleanup threshold is set to 0", func() {
			BeforeEach(func() {
				config.GraphCleanupThresholdMB = intptr(0)
			})

			Context("and we run a pea with the same image as the container", func() {
				peaDoesNotLeaveDebrisInGraphOnExit("docker:///busybox")
			})

			Context("and we run a pea with another image", func() {
				peaDoesNotLeaveDebrisInGraphOnExit("docker:///alpine")
			})
		})

		Context("when the graph cleanup threshold is exceeded", func() {
			BeforeEach(func() {
				config.GraphCleanupThresholdMB = intptr(0)
				persistentImages = []string{}
			})

			Context("when there are other rootfs layers in the graph dir", func() {
				It("cleans up the graph directory on container creation (and not on destruction)", func() {
					Expect(numLayersInGraph()).To(BeNumerically(">", 0))

					anotherContainer, err := client.Create(garden.ContainerSpec{})
					Expect(err).ToNot(HaveOccurred())

					Expect(numLayersInGraph()).To(Equal(3), "after creation, should have deleted everything other than the default rootfs, uid translation layer and container layer")
					Expect(client.Destroy(anotherContainer.Handle())).To(Succeed())
					Expect(numLayersInGraph()).To(Equal(2), "should not garbage collect parent layers on destroy")
				})
			})
		})

		Context("when the graph cleanup threshold is not exceeded", func() {
			BeforeEach(func() {
				config.GraphCleanupThresholdMB = intptr(1024)
			})

			It("does not cleanup", func() {
				// threshold is not yet exceeded
				Expect(numLayersInGraph()).To(Equal(3))

				anotherContainer, err := client.Create(garden.ContainerSpec{})
				Expect(err).ToNot(HaveOccurred())

				Expect(numLayersInGraph()).To(Equal(6))
				Expect(client.Destroy(anotherContainer.Handle())).To(Succeed())
			})
		})
	})

	Describe("--persistentImage", func() {
		BeforeEach(func() {
			config.GraphCleanupThresholdMB = intptr(0)
		})

		Context("when set", func() {
			JustBeforeEach(func() {
				Eventually(client, "60s").Should(gbytes.Say("retain.retained"))
			})

			Context("and local images are used", func() {
				BeforeEach(func() {
					persistentImages = []string{defaultTestRootFS}
				})

				Describe("graph cleanup for a rootfs on the whitelist", func() {
					It("keeps the rootfs in the graph", func() {
						container, err := client.Create(garden.ContainerSpec{
							RootFSPath: persistentImages[0],
						})
						Expect(err).ToNot(HaveOccurred())
						Expect(client.Destroy(container.Handle())).To(Succeed())

						expectLayerCountAfterGraphCleanupToBe(2)
					})

					Context("which is a symlink", func() {
						BeforeEach(func() {
							Expect(os.MkdirAll("/var/vcap/packages", 0755)).To(Succeed())
							err := exec.Command("ln", "-s", defaultTestRootFS, "/var/vcap/packages/busybox").Run()
							Expect(err).ToNot(HaveOccurred())

							persistentImages = []string{"/var/vcap/packages/busybox"}
						})

						AfterEach(func() {
							Expect(os.RemoveAll("/var/vcap/packages")).To(Succeed())
						})

						It("keeps the rootfs in the graph", func() {
							container, err := client.Create(garden.ContainerSpec{
								RootFSPath: persistentImages[0],
							})
							Expect(err).ToNot(HaveOccurred())
							Expect(client.Destroy(container.Handle())).To(Succeed())

							expectLayerCountAfterGraphCleanupToBe(2)
						})
					})
				})

				Describe("graph cleanup for a rootfs not on the whitelist", func() {
					It("cleans up all unused images from the graph", func() {
						container, err := client.Create(garden.ContainerSpec{
							RootFSPath: nonDefaultRootfsPath,
						})
						Expect(err).ToNot(HaveOccurred())
						Expect(client.Destroy(container.Handle())).To(Succeed())

						expectLayerCountAfterGraphCleanupToBe(0)
					})
				})
			})

			Context("and docker images are used", func() {
				BeforeEach(func() {
					persistentImages = []string{
						"docker:///busybox",
						"docker:///ubuntu",
						"docker://banana/bananatest",
					}
				})

				Describe("graph cleanup for a rootfs on the whitelist", func() {
					It("keeps the rootfs in the graph", func() {
						numLayersBeforeDockerPull := numLayersInGraph()
						container, err := client.Create(garden.ContainerSpec{
							RootFSPath: persistentImages[0],
						})
						Expect(err).ToNot(HaveOccurred())
						Expect(client.Destroy(container.Handle())).To(Succeed())
						numLayersInImage := numLayersInGraph() - numLayersBeforeDockerPull

						expectLayerCountAfterGraphCleanupToBe(numLayersInImage)
					})
				})

				Describe("graph cleanup for a rootfs not on the whitelist", func() {
					It("cleans up all unused images from the graph", func() {
						container, err := client.Create(garden.ContainerSpec{
							RootFSPath: "docker:///cfgarden/garden-busybox",
						})
						Expect(err).ToNot(HaveOccurred())
						Expect(client.Destroy(container.Handle())).To(Succeed())

						expectLayerCountAfterGraphCleanupToBe(0)
					})
				})
			})
		})

		Context("when it is not set", func() {
			BeforeEach(func() {
				persistentImages = []string{}
			})

			It("cleans up all unused images from the graph", func() {
				defaultRootfsContainer, err := client.Create(garden.ContainerSpec{})
				Expect(err).ToNot(HaveOccurred())

				nonDefaultRootfsContainer, err := client.Create(garden.ContainerSpec{
					RootFSPath: nonDefaultRootfsPath,
				})
				Expect(err).ToNot(HaveOccurred())

				Expect(client.Destroy(defaultRootfsContainer.Handle())).To(Succeed())
				Expect(client.Destroy(nonDefaultRootfsContainer.Handle())).To(Succeed())

				expectLayerCountAfterGraphCleanupToBe(0)
			})
		})
	})
})
<commit_msg>Don't run shed GQTs on xenial<commit_after>package gqt_test

import (
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/guardian/gqt/runner"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
)

var _ = Describe("deprecated graph tests (aka garden-shed)", func() {
	var (
		client               *runner.RunningGarden
		layersPath           string
		diffPath             string
		mntPath              string
		nonDefaultRootfsPath string
		persistentImages     []string
	)

	numLayersInGraph := func() int {
		layerFiles, err := ioutil.ReadDir(layersPath)
		Expect(err).ToNot(HaveOccurred())
		diffFiles, err := ioutil.ReadDir(diffPath)
		Expect(err).ToNot(HaveOccurred())
		mntFiles, err := ioutil.ReadDir(mntPath)
		Expect(err).ToNot(HaveOccurred())

		numLayerFiles := len(layerFiles)
		Expect(numLayerFiles).To(Equal(len(diffFiles)))
		Expect(numLayerFiles).To(Equal(len(mntFiles)))
		return numLayerFiles
	}

	expectLayerCountAfterGraphCleanupToBe := func(layerCount int) {
		nonPersistantRootfsContainer, err := client.Create(garden.ContainerSpec{
			RootFSPath: nonDefaultRootfsPath,
		})
		Expect(err).ToNot(HaveOccurred())
		Expect(client.Destroy(nonPersistantRootfsContainer.Handle())).To(Succeed())
		Expect(numLayersInGraph()).To(Equal(layerCount + 2)) // +2 for the layers created for the nondefaultrootfs container
	}

	peaDoesNotLeaveDebrisInGraphOnExit := func(imageURI string) {
		var (
			pea                             garden.Process
			numLayersAfterContainerCreation int
		)

		JustBeforeEach(func() {
			imageURI := imageURI
			container, err := client.Create(garden.ContainerSpec{
				RootFSPath: imageURI,
			})
			Expect(err).ToNot(HaveOccurred())

			numLayersAfterContainerCreation = numLayersInGraph()

			pea, err = container.Run(garden.ProcessSpec{
				Path:  "/bin/sh",
				Args:  []string{"-c", "exit 0"},
				Image: garden.ImageRef{URI: imageURI},
			}, garden.ProcessIO{})
			Expect(err).NotTo(HaveOccurred())
		})

		It("does not leave debris in the graph on exit", func() {
			Expect(pea.Wait()).To(Equal(0))

			numLayersAfterPeaExits := numLayersInGraph()
			Expect(numLayersAfterPeaExits).To(Equal(numLayersAfterContainerCreation))
		})
	}

	BeforeEach(func() {
		if runningOnXenial() {
			Skip("garden-shed not supported on xenial")
		}

		var err error
		nonDefaultRootfsPath, err = ioutil.TempDir("", "tmpRootfs")
		Expect(err).ToNot(HaveOccurred())
		// temporary workaround as runc expects a /tmp dir to exist in the container rootfs
		err = os.Mkdir(filepath.Join(nonDefaultRootfsPath, "tmp"), 0700)
		Expect(err).ToNot(HaveOccurred())
	})

	JustBeforeEach(func() {
		config.PersistentImages = persistentImages
		config = resetImagePluginConfig()
		client = runner.Start(config)

		layersPath = path.Join(*client.GraphDir, "aufs", "layers")
		diffPath = path.Join(*client.GraphDir, "aufs", "diff")
		mntPath = path.Join(*client.GraphDir, "aufs", "mnt")
	})

	AfterEach(func() {
		Expect(os.RemoveAll(nonDefaultRootfsPath)).To(Succeed())
		Expect(client.DestroyAndStop()).To(Succeed())
	})

	Describe("--graph-cleanup-threshold-in-megabytes", func() {
		JustBeforeEach(func() {
			Expect(numLayersInGraph()).To(Equal(0))
			container, err := client.Create(garden.ContainerSpec{
				RootFSPath: "docker:///busybox",
			})
			Expect(err).ToNot(HaveOccurred())
			Expect(client.Destroy(container.Handle())).To(Succeed())
		})

		Context("when the graph cleanup threshold is set to -1", func() {
			BeforeEach(func() {
				config.GraphCleanupThresholdMB = intptr(-1)
			})

			It("does NOT clean up the graph directory on create", func() {
				initialNumberOfLayers := numLayersInGraph()
				anotherContainer, err := client.Create(garden.ContainerSpec{})
				Expect(err).ToNot(HaveOccurred())

				Expect(numLayersInGraph()).To(BeNumerically(">", initialNumberOfLayers), "after creation, should NOT have deleted anything")
				Expect(client.Destroy(anotherContainer.Handle())).To(Succeed())
			})

			Context("and we run a pea with the same image as the container", func() {
				peaDoesNotLeaveDebrisInGraphOnExit("docker:///busybox")
			})

			Context("and we run a pea with another image", func() {
				peaDoesNotLeaveDebrisInGraphOnExit("docker:///alpine")
			})
		})

		Context("when the graph cleanup threshold is set to 0", func() {
			BeforeEach(func() {
				config.GraphCleanupThresholdMB = intptr(0)
			})

			Context("and we run a pea with the same image as the container", func() {
				peaDoesNotLeaveDebrisInGraphOnExit("docker:///busybox")
			})

			Context("and we run a pea with another image", func() {
				peaDoesNotLeaveDebrisInGraphOnExit("docker:///alpine")
			})
		})

		Context("when the graph cleanup threshold is exceeded", func() {
			BeforeEach(func() {
				config.GraphCleanupThresholdMB = intptr(0)
				persistentImages = []string{}
			})

			Context("when there are other rootfs layers in the graph dir", func() {
				It("cleans up the graph directory on container creation (and not on destruction)", func() {
					Expect(numLayersInGraph()).To(BeNumerically(">", 0))

					anotherContainer, err := client.Create(garden.ContainerSpec{})
					Expect(err).ToNot(HaveOccurred())

					Expect(numLayersInGraph()).To(Equal(3), "after creation, should have deleted everything other than the default rootfs, uid translation layer and container layer")
					Expect(client.Destroy(anotherContainer.Handle())).To(Succeed())
					Expect(numLayersInGraph()).To(Equal(2), "should not garbage collect parent layers on destroy")
				})
			})
		})

		Context("when the graph cleanup threshold is not exceeded", func() {
			BeforeEach(func() {
				config.GraphCleanupThresholdMB = intptr(1024)
			})

			It("does not cleanup", func() {
				// threshold is not yet exceeded
				Expect(numLayersInGraph()).To(Equal(3))

				anotherContainer, err := client.Create(garden.ContainerSpec{})
				Expect(err).ToNot(HaveOccurred())

				Expect(numLayersInGraph()).To(Equal(6))
				Expect(client.Destroy(anotherContainer.Handle())).To(Succeed())
			})
		})
	})

	Describe("--persistentImage", func() {
		BeforeEach(func() {
			config.GraphCleanupThresholdMB = intptr(0)
		})

		Context("when set", func() {
			JustBeforeEach(func() {
				Eventually(client, "60s").Should(gbytes.Say("retain.retained"))
			})

			Context("and local images are used", func() {
				BeforeEach(func() {
					persistentImages = []string{defaultTestRootFS}
				})

				Describe("graph cleanup for a rootfs on the whitelist", func() {
					It("keeps the rootfs in the graph", func() {
						container, err := client.Create(garden.ContainerSpec{
							RootFSPath: persistentImages[0],
						})
						Expect(err).ToNot(HaveOccurred())
						Expect(client.Destroy(container.Handle())).To(Succeed())

						expectLayerCountAfterGraphCleanupToBe(2)
					})

					Context("which is a symlink", func() {
						BeforeEach(func() {
							Expect(os.MkdirAll("/var/vcap/packages", 0755)).To(Succeed())
							err := exec.Command("ln", "-s", defaultTestRootFS, "/var/vcap/packages/busybox").Run()
							Expect(err).ToNot(HaveOccurred())

							persistentImages = []string{"/var/vcap/packages/busybox"}
						})

						AfterEach(func() {
							Expect(os.RemoveAll("/var/vcap/packages")).To(Succeed())
						})

						It("keeps the rootfs in the graph", func() {
							container, err := client.Create(garden.ContainerSpec{
								RootFSPath: persistentImages[0],
							})
							Expect(err).ToNot(HaveOccurred())
							Expect(client.Destroy(container.Handle())).To(Succeed())

							expectLayerCountAfterGraphCleanupToBe(2)
						})
					})
				})

				Describe("graph cleanup for a rootfs not on the whitelist", func() {
					It("cleans up all unused images from the graph", func() {
						container, err := client.Create(garden.ContainerSpec{
							RootFSPath: nonDefaultRootfsPath,
						})
						Expect(err).ToNot(HaveOccurred())
						Expect(client.Destroy(container.Handle())).To(Succeed())

						expectLayerCountAfterGraphCleanupToBe(0)
					})
				})
			})

			Context("and docker images are used", func() {
				BeforeEach(func() {
					persistentImages = []string{
						"docker:///busybox",
						"docker:///ubuntu",
						"docker://banana/bananatest",
					}
				})

				Describe("graph cleanup for a rootfs on the whitelist", func() {
					It("keeps the rootfs in the graph", func() {
						numLayersBeforeDockerPull := numLayersInGraph()
						container, err := client.Create(garden.ContainerSpec{
							RootFSPath: persistentImages[0],
						})
						Expect(err).ToNot(HaveOccurred())
						Expect(client.Destroy(container.Handle())).To(Succeed())
						numLayersInImage := numLayersInGraph() - numLayersBeforeDockerPull

						expectLayerCountAfterGraphCleanupToBe(numLayersInImage)
					})
				})

				Describe("graph cleanup for a rootfs not on the whitelist", func() {
					It("cleans up all unused images from the graph", func() {
						container, err := client.Create(garden.ContainerSpec{
							RootFSPath: "docker:///cfgarden/garden-busybox",
						})
						Expect(err).ToNot(HaveOccurred())
						Expect(client.Destroy(container.Handle())).To(Succeed())

						expectLayerCountAfterGraphCleanupToBe(0)
					})
				})
			})
		})

		Context("when it is not set", func() {
			BeforeEach(func() {
				persistentImages = []string{}
			})

			It("cleans up all unused images from the graph", func() {
				defaultRootfsContainer, err := client.Create(garden.ContainerSpec{})
				Expect(err).ToNot(HaveOccurred())

				nonDefaultRootfsContainer, err := client.Create(garden.ContainerSpec{
					RootFSPath: nonDefaultRootfsPath,
				})
				Expect(err).ToNot(HaveOccurred())

				Expect(client.Destroy(defaultRootfsContainer.Handle())).To(Succeed())
				Expect(client.Destroy(nonDefaultRootfsContainer.Handle())).To(Succeed())

				expectLayerCountAfterGraphCleanupToBe(0)
			})
		})
	})
})

// runningOnXenial() checks the /etc/os-release file to determine the current
// OS. It will only return true if the /etc/os-release file exists and contains
// the string "Xenial".
func runningOnXenial() bool {
	checkCmd := exec.Command("grep", "Xenial", filepath.Join("/", "etc", "os-release"))
	return checkCmd.Run() == nil
}
<|endoftext|>