<commit_before>// Copyright 2022 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package execute

import (
	"context"
	"fmt"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/gae/service/datastore"

	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/run"
	"go.chromium.org/luci/cv/internal/tryjob"
)

type state struct {
	_kind string         `gae:"$kind,TryjobExecutionState"`
	_id   int64          `gae:"$id,1"`
	Run   *datastore.Key `gae:"$parent"`
	// EVersion is the version of this state. Start with 1.
	EVersion int64
	State    *tryjob.ExecutionState
}

func loadExecutionState(ctx context.Context, rid common.RunID) (*tryjob.ExecutionState, int64, error) {
	s := &state{Run: datastore.MakeKey(ctx, run.RunKind, string(rid))}
	switch err := datastore.Get(ctx, s); {
	case err == datastore.ErrNoSuchEntity:
		return nil, 0, nil
	case err != nil:
		return nil, 0, errors.Annotate(err, "failed to load tryjob execution state of run %q", rid).Err()
	default:
		return s.State, s.EVersion, nil
	}
}

// Do executes the tryjob requirement for a run.
//
// This function is idempotent so it is safe to retry.
func Do(ctx context.Context, r *run.Run, t *tryjob.ExecuteTryjobsPayload, shouldStop func() bool) error {
	execState, stateVer, err := loadExecutionState(ctx, r.ID)
	switch {
	case err != nil:
		return err
	case execState == nil:
		execState = &tryjob.ExecutionState{}
	}

	plan, err := prepExecutionPlan(ctx, execState, r, t.GetTryjobsUpdated(), t.GetRequirementChanged())
	if err != nil {
		return err
	}
	if err := plan.execute(ctx, r, execState); err != nil {
		return err
	}

	var innerErr error
	err = datastore.RunInTransaction(ctx, func(c context.Context) (err error) {
		defer func() { innerErr = err }()
		switch _, latestStateVer, err := loadExecutionState(ctx, r.ID); {
		case err != nil:
			return err
		case latestStateVer != stateVer:
			return errors.Reason("execution state has changed. before: %d, current: %d", stateVer, latestStateVer).Tag(transient.Tag).Err()
		default:
			s := &state{
				Run:      datastore.MakeKey(ctx, run.RunKind, string(r.ID)),
				EVersion: latestStateVer + 1,
				State:    execState,
			}
			return errors.Annotate(datastore.Put(ctx, s), "failed to save execution state").Tag(transient.Tag).Err()
		}
	}, nil)

	switch {
	case innerErr != nil:
		return innerErr
	case err != nil:
		return errors.Annotate(err, "failed to commit transaction").Tag(transient.Tag).Err()
	default:
		return nil
	}
}

type plan struct {
	// new tryjob to trigger due to new builder in requirement.
	trigger []*tryjob.Definition

	// retry the existing tryjob due to transient failure.
	retry map[*tryjob.Definition]*tryjob.ExecutionState_Execution

	// discard the tryjob due to the builder is removed in the Project Config.
	discarded map[*tryjob.Definition]*tryjob.ExecutionState_Execution
}

func prepExecutionPlan(ctx context.Context, execState *tryjob.ExecutionState, r *run.Run, tryjobsUpdated []int64, reqmtChanged bool) (*plan, error) {
	switch {
	case len(tryjobsUpdated) > 0 && reqmtChanged:
		panic(fmt.Errorf("the executor can't handle requirement update and tryjobs update at the same time"))
	case len(tryjobsUpdated) > 0:
		tryjobs := make([]*tryjob.Tryjob, len(tryjobsUpdated))
		if err := datastore.Get(ctx, tryjobs); err != nil {
			return nil, errors.Annotate(err, "failed to load tryjobs").Tag(transient.Tag).Err()
		}
		// TODO(robertocn): update the state and compute the plan
		return &plan{}, nil
	case reqmtChanged:
		// TODO(yiwzhang): compute the plan
		return &plan{}, nil
	default:
		panic(fmt.Errorf("the executor is called without any update on either tryjobs or requirement"))
	}
}

// execute executes the plan and mutate the state.
func (p plan) execute(ctx context.Context, r *run.Run, execState *tryjob.ExecutionState) error {
	panic("implement")
}
<commit_msg>[cv] [tj] Load tryjobs correctly in prepExecutionPlan.<commit_after>// Copyright 2022 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package execute

import (
	"context"
	"fmt"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/gae/service/datastore"

	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/run"
	"go.chromium.org/luci/cv/internal/tryjob"
)

type state struct {
	_kind string         `gae:"$kind,TryjobExecutionState"`
	_id   int64          `gae:"$id,1"`
	Run   *datastore.Key `gae:"$parent"`
	// EVersion is the version of this state. Start with 1.
	EVersion int64
	State    *tryjob.ExecutionState
}

func loadExecutionState(ctx context.Context, rid common.RunID) (*tryjob.ExecutionState, int64, error) {
	s := &state{Run: datastore.MakeKey(ctx, run.RunKind, string(rid))}
	switch err := datastore.Get(ctx, s); {
	case err == datastore.ErrNoSuchEntity:
		return nil, 0, nil
	case err != nil:
		return nil, 0, errors.Annotate(err, "failed to load tryjob execution state of run %q", rid).Err()
	default:
		return s.State, s.EVersion, nil
	}
}

// Do executes the tryjob requirement for a run.
//
// This function is idempotent so it is safe to retry.
func Do(ctx context.Context, r *run.Run, t *tryjob.ExecuteTryjobsPayload, shouldStop func() bool) error {
	execState, stateVer, err := loadExecutionState(ctx, r.ID)
	switch {
	case err != nil:
		return err
	case execState == nil:
		execState = &tryjob.ExecutionState{}
	}

	plan, err := prepExecutionPlan(ctx, execState, r, t.GetTryjobsUpdated(), t.GetRequirementChanged())
	if err != nil {
		return err
	}
	if err := plan.execute(ctx, r, execState); err != nil {
		return err
	}

	var innerErr error
	err = datastore.RunInTransaction(ctx, func(c context.Context) (err error) {
		defer func() { innerErr = err }()
		switch _, latestStateVer, err := loadExecutionState(ctx, r.ID); {
		case err != nil:
			return err
		case latestStateVer != stateVer:
			return errors.Reason("execution state has changed. before: %d, current: %d", stateVer, latestStateVer).Tag(transient.Tag).Err()
		default:
			s := &state{
				Run:      datastore.MakeKey(ctx, run.RunKind, string(r.ID)),
				EVersion: latestStateVer + 1,
				State:    execState,
			}
			return errors.Annotate(datastore.Put(ctx, s), "failed to save execution state").Tag(transient.Tag).Err()
		}
	}, nil)

	switch {
	case innerErr != nil:
		return innerErr
	case err != nil:
		return errors.Annotate(err, "failed to commit transaction").Tag(transient.Tag).Err()
	default:
		return nil
	}
}

type plan struct {
	// new tryjob to trigger due to new builder in requirement.
	trigger []*tryjob.Definition

	// retry the existing tryjob due to transient failure.
	retry map[*tryjob.Definition]*tryjob.ExecutionState_Execution

	// discard the tryjob due to the builder is removed in the Project Config.
	discarded map[*tryjob.Definition]*tryjob.ExecutionState_Execution
}

func prepExecutionPlan(ctx context.Context, execState *tryjob.ExecutionState, r *run.Run, tryjobsUpdated []int64, reqmtChanged bool) (*plan, error) {
	switch {
	case len(tryjobsUpdated) > 0 && reqmtChanged:
		panic(fmt.Errorf("the executor can't handle requirement update and tryjobs update at the same time"))
	case len(tryjobsUpdated) > 0:
		tryjobs := make([]*tryjob.Tryjob, len(tryjobsUpdated))
		for i, tjid := range tryjobsUpdated {
			tryjobs[i] = &tryjob.Tryjob{ID: common.TryjobID(tjid)}
		}
		if err := datastore.Get(ctx, tryjobs); err != nil {
			return nil, errors.Annotate(err, "failed to load tryjobs").Tag(transient.Tag).Err()
		}
		// TODO(robertocn): update the state and compute the plan
		return &plan{}, nil
	case reqmtChanged:
		// TODO(yiwzhang): compute the plan
		return &plan{}, nil
	default:
		panic(fmt.Errorf("the executor is called without any update on either tryjobs or requirement"))
	}
}

// execute executes the plan and mutate the state.
func (p plan) execute(ctx context.Context, r *run.Run, execState *tryjob.ExecutionState) error {
	panic("implement")
}
<|endoftext|>


<commit_before>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package gvisor provides support for gVisor, user-space kernel, testing.
// See https://github.com/google/gvisor
package gvisor

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/google/syzkaller/pkg/config"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/vm/vmimpl"
)

func init() {
	vmimpl.Register("gvisor", ctor)
}

type Config struct {
	Count     int    `json:"count"` // number of VMs to use
	RunscArgs string `json:"runsc_args"`
}

type Pool struct {
	env *vmimpl.Env
	cfg *Config
}

type instance struct {
	cfg      *Config
	image    string
	debug    bool
	rootDir  string
	imageDir string
	name     string
	port     int
	cmd      *exec.Cmd
	merger   *vmimpl.OutputMerger
}

func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
	cfg := &Config{
		Count: 1,
	}
	if err := config.LoadData(env.Config, cfg); err != nil {
		return nil, fmt.Errorf("failed to parse vm config: %v", err)
	}
	if cfg.Count < 1 || cfg.Count > 1000 {
		return nil, fmt.Errorf("invalid config param count: %v, want [1, 1000]", cfg.Count)
	}
	if env.Debug {
		cfg.Count = 1
	}
	if !osutil.IsExist(env.Image) {
		return nil, fmt.Errorf("image file %q does not exist", env.Image)
	}
	pool := &Pool{
		cfg: cfg,
		env: env,
	}
	return pool, nil
}

func (pool *Pool) Count() int {
	return pool.cfg.Count
}

func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
	rootDir := filepath.Clean(filepath.Join(workdir, "..", "gvisor_root"))
	imageDir := filepath.Join(workdir, "image")
	bundleDir := filepath.Join(workdir, "bundle")
	osutil.MkdirAll(rootDir)
	osutil.MkdirAll(bundleDir)
	osutil.MkdirAll(imageDir)

	caps := ""
	for _, c := range sandboxCaps {
		if caps != "" {
			caps += ", "
		}
		caps += "\"" + c + "\""
	}
	vmConfig := fmt.Sprintf(configTempl, imageDir, caps)
	if err := osutil.WriteFile(filepath.Join(bundleDir, "config.json"), []byte(vmConfig)); err != nil {
		return nil, err
	}
	if err := osutil.CopyFile(os.Args[0], filepath.Join(imageDir, "init")); err != nil {
		return nil, err
	}

	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, err
	}
	var tee io.Writer
	if pool.env.Debug {
		tee = os.Stdout
	}
	merger := vmimpl.NewOutputMerger(tee)
	merger.Add("gvisor", rpipe)

	inst := &instance{
		cfg:      pool.cfg,
		image:    pool.env.Image,
		debug:    pool.env.Debug,
		rootDir:  rootDir,
		imageDir: imageDir,
		name:     fmt.Sprintf("%v-%v", pool.env.Name, index),
		merger:   merger,
	}

	// Kill the previous instance in case it's still running.
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	time.Sleep(3 * time.Second)

	cmd := inst.runscCmd("run", "-bundle", bundleDir, inst.name)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		merger.Wait()
		return nil, err
	}
	inst.cmd = cmd
	wpipe.Close()

	if err := inst.waitBoot(); err != nil {
		inst.Close()
		return nil, err
	}
	return inst, nil
}

func (inst *instance) waitBoot() error {
	errorMsg := []byte("FATAL ERROR:")
	bootedMsg := []byte(initStartMsg)
	timeout := time.NewTimer(time.Minute)
	defer timeout.Stop()
	var output []byte
	for {
		select {
		case out := <-inst.merger.Output:
			output = append(output, out...)
			if pos := bytes.Index(output, errorMsg); pos != -1 {
				end := bytes.IndexByte(output[pos:], '\n')
				if end == -1 {
					end = len(output)
				} else {
					end += pos
				}
				return vmimpl.BootError{
					Title:  string(output[pos:end]),
					Output: output,
				}
			}
			if bytes.Contains(output, bootedMsg) {
				return nil
			}
		case err := <-inst.merger.Err:
			return vmimpl.BootError{
				Title:  fmt.Sprintf("runsc failed: %v", err),
				Output: output,
			}
		case <-timeout.C:
			return vmimpl.BootError{
				Title:  "init process did not start",
				Output: output,
			}
		}
	}
}

func (inst *instance) runscCmd(add ...string) *exec.Cmd {
	args := []string{
		"-root", inst.rootDir,
		"-network=none",
	}
	args = append(args, strings.Split(inst.cfg.RunscArgs, " ")...)
	args = append(args, add...)
	cmd := osutil.Command(inst.image, args...)
	cmd.Env = []string{
		"GOTRACEBACK=all",
		"GORACE=halt_on_error=1",
	}
	return cmd
}

func (inst *instance) Close() {
	time.Sleep(3 * time.Second)
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	inst.cmd.Process.Kill()
	inst.merger.Wait()
	inst.cmd.Wait()
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	time.Sleep(3 * time.Second)
}

func (inst *instance) Forward(port int) (string, error) {
	if inst.port != 0 {
		return "", fmt.Errorf("forward port is already setup")
	}
	inst.port = port
	return "stdin", nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	fname := filepath.Base(hostSrc)
	if err := osutil.CopyFile(hostSrc, filepath.Join(inst.imageDir, fname)); err != nil {
		return "", err
	}
	if err := os.Chmod(inst.imageDir, 0777); err != nil {
		return "", err
	}
	return filepath.Join("/", fname), nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
	<-chan []byte, <-chan error, error) {
	args := []string{"exec", "-user=0:0"}
	for _, c := range sandboxCaps {
		args = append(args, "-cap", c)
	}
	args = append(args, inst.name)
	args = append(args, strings.Split(command, " ")...)
	cmd := inst.runscCmd(args...)

	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	defer wpipe.Close()
	inst.merger.Add("cmd", rpipe)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe

	guestSock, err := inst.guestProxy()
	if err != nil {
		return nil, nil, err
	}
	if guestSock != nil {
		defer guestSock.Close()
		cmd.Stdin = guestSock
	}

	if err := cmd.Start(); err != nil {
		return nil, nil, err
	}
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	go func() {
		select {
		case <-time.After(timeout):
			signal(vmimpl.ErrTimeout)
		case <-stop:
			signal(vmimpl.ErrTimeout)
		case err := <-inst.merger.Err:
			cmd.Process.Kill()
			if cmdErr := cmd.Wait(); cmdErr == nil {
				// If the command exited successfully, we got EOF error from merger.
				// But in this case no error has happened and the EOF is expected.
				err = nil
			}
			signal(err)
			return
		}
		cmd.Process.Kill()
		cmd.Wait()
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) guestProxy() (*os.File, error) {
	if inst.port == 0 {
		return nil, nil
	}
	// One does not simply let gvisor guest connect to host tcp port.
	// We create a unix socket, pass it to guest in stdin.
	// Guest will use it instead of dialing manager directly.
	// On host we connect to manager tcp port and proxy between the tcp and unix connections.
	socks, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
	if err != nil {
		return nil, err
	}
	hostSock := os.NewFile(uintptr(socks[0]), "host unix proxy")
	guestSock := os.NewFile(uintptr(socks[1]), "guest unix proxy")
	conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%v", inst.port))
	if err != nil {
		conn.Close()
		hostSock.Close()
		guestSock.Close()
		return nil, err
	}
	go func() {
		io.Copy(hostSock, conn)
		hostSock.Close()
	}()
	go func() {
		io.Copy(conn, hostSock)
		conn.Close()
	}()
	return guestSock, nil
}

func (inst *instance) Diagnose() bool {
	osutil.Run(time.Minute, inst.runscCmd("debug", "-stacks", inst.name))
	return true
}

func init() {
	if os.Getenv("SYZ_GVISOR_PROXY") != "" {
		fmt.Fprintf(os.Stderr, initStartMsg)
		select {}
	}
}

const initStartMsg = "SYZKALLER INIT STARTED\n"

const configTempl = `
{
	"root": {
		"path": "%[1]v",
		"readonly": true
	},
	"process":{
                "args": ["/init"],
                "cwd": "/tmp",
                "env": ["SYZ_GVISOR_PROXY=1"],
                "capabilities": {
                	"bounding": [%[2]v],
                	"effective": [%[2]v],
                	"inheritable": [%[2]v],
                	"permitted": [%[2]v],
                	"ambient": [%[2]v]
                }
	}
}
`

var sandboxCaps = []string{
	"CAP_CHOWN", "CAP_DAC_OVERRIDE", "CAP_DAC_READ_SEARCH", "CAP_FOWNER", "CAP_FSETID",
	"CAP_KILL", "CAP_SETGID", "CAP_SETUID", "CAP_SETPCAP", "CAP_LINUX_IMMUTABLE",
	"CAP_NET_BIND_SERVICE", "CAP_NET_BROADCAST", "CAP_NET_ADMIN", "CAP_NET_RAW",
	"CAP_IPC_LOCK", "CAP_IPC_OWNER", "CAP_SYS_MODULE", "CAP_SYS_RAWIO", "CAP_SYS_CHROOT",
	"CAP_SYS_PTRACE", "CAP_SYS_PACCT", "CAP_SYS_ADMIN", "CAP_SYS_BOOT", "CAP_SYS_NICE",
	"CAP_SYS_RESOURCE", "CAP_SYS_TIME", "CAP_SYS_TTY_CONFIG", "CAP_MKNOD", "CAP_LEASE",
	"CAP_AUDIT_WRITE", "CAP_AUDIT_CONTROL", "CAP_SETFCAP", "CAP_MAC_OVERRIDE", "CAP_MAC_ADMIN",
	"CAP_SYSLOG", "CAP_WAKE_ALARM", "CAP_BLOCK_SUSPEND", "CAP_AUDIT_READ",
}
<commit_msg>vm/gvisor: fix empty runsc_args<commit_after>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package gvisor provides support for gVisor, user-space kernel, testing.
// See https://github.com/google/gvisor
package gvisor

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/google/syzkaller/pkg/config"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/vm/vmimpl"
)

func init() {
	vmimpl.Register("gvisor", ctor)
}

type Config struct {
	Count     int    `json:"count"` // number of VMs to use
	RunscArgs string `json:"runsc_args"`
}

type Pool struct {
	env *vmimpl.Env
	cfg *Config
}

type instance struct {
	cfg      *Config
	image    string
	debug    bool
	rootDir  string
	imageDir string
	name     string
	port     int
	cmd      *exec.Cmd
	merger   *vmimpl.OutputMerger
}

func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
	cfg := &Config{
		Count: 1,
	}
	if err := config.LoadData(env.Config, cfg); err != nil {
		return nil, fmt.Errorf("failed to parse vm config: %v", err)
	}
	if cfg.Count < 1 || cfg.Count > 1000 {
		return nil, fmt.Errorf("invalid config param count: %v, want [1, 1000]", cfg.Count)
	}
	if env.Debug {
		cfg.Count = 1
	}
	if !osutil.IsExist(env.Image) {
		return nil, fmt.Errorf("image file %q does not exist", env.Image)
	}
	pool := &Pool{
		cfg: cfg,
		env: env,
	}
	return pool, nil
}

func (pool *Pool) Count() int {
	return pool.cfg.Count
}

func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
	rootDir := filepath.Clean(filepath.Join(workdir, "..", "gvisor_root"))
	imageDir := filepath.Join(workdir, "image")
	bundleDir := filepath.Join(workdir, "bundle")
	osutil.MkdirAll(rootDir)
	osutil.MkdirAll(bundleDir)
	osutil.MkdirAll(imageDir)

	caps := ""
	for _, c := range sandboxCaps {
		if caps != "" {
			caps += ", "
		}
		caps += "\"" + c + "\""
	}
	vmConfig := fmt.Sprintf(configTempl, imageDir, caps)
	if err := osutil.WriteFile(filepath.Join(bundleDir, "config.json"), []byte(vmConfig)); err != nil {
		return nil, err
	}
	if err := osutil.CopyFile(os.Args[0], filepath.Join(imageDir, "init")); err != nil {
		return nil, err
	}

	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, err
	}
	var tee io.Writer
	if pool.env.Debug {
		tee = os.Stdout
	}
	merger := vmimpl.NewOutputMerger(tee)
	merger.Add("gvisor", rpipe)

	inst := &instance{
		cfg:      pool.cfg,
		image:    pool.env.Image,
		debug:    pool.env.Debug,
		rootDir:  rootDir,
		imageDir: imageDir,
		name:     fmt.Sprintf("%v-%v", pool.env.Name, index),
		merger:   merger,
	}

	// Kill the previous instance in case it's still running.
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	time.Sleep(3 * time.Second)

	cmd := inst.runscCmd("run", "-bundle", bundleDir, inst.name)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		merger.Wait()
		return nil, err
	}
	inst.cmd = cmd
	wpipe.Close()

	if err := inst.waitBoot(); err != nil {
		inst.Close()
		return nil, err
	}
	return inst, nil
}

func (inst *instance) waitBoot() error {
	errorMsg := []byte("FATAL ERROR:")
	bootedMsg := []byte(initStartMsg)
	timeout := time.NewTimer(time.Minute)
	defer timeout.Stop()
	var output []byte
	for {
		select {
		case out := <-inst.merger.Output:
			output = append(output, out...)
			if pos := bytes.Index(output, errorMsg); pos != -1 {
				end := bytes.IndexByte(output[pos:], '\n')
				if end == -1 {
					end = len(output)
				} else {
					end += pos
				}
				return vmimpl.BootError{
					Title:  string(output[pos:end]),
					Output: output,
				}
			}
			if bytes.Contains(output, bootedMsg) {
				return nil
			}
		case err := <-inst.merger.Err:
			return vmimpl.BootError{
				Title:  fmt.Sprintf("runsc failed: %v", err),
				Output: output,
			}
		case <-timeout.C:
			return vmimpl.BootError{
				Title:  "init process did not start",
				Output: output,
			}
		}
	}
}

func (inst *instance) runscCmd(add ...string) *exec.Cmd {
	args := []string{
		"-root", inst.rootDir,
		"-network=none",
	}
	if inst.cfg.RunscArgs != "" {
		args = append(args, strings.Split(inst.cfg.RunscArgs, " ")...)
	}
	args = append(args, add...)
	cmd := osutil.Command(inst.image, args...)
	cmd.Env = []string{
		"GOTRACEBACK=all",
		"GORACE=halt_on_error=1",
	}
	return cmd
}

func (inst *instance) Close() {
	time.Sleep(3 * time.Second)
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	inst.cmd.Process.Kill()
	inst.merger.Wait()
	inst.cmd.Wait()
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	time.Sleep(3 * time.Second)
}

func (inst *instance) Forward(port int) (string, error) {
	if inst.port != 0 {
		return "", fmt.Errorf("forward port is already setup")
	}
	inst.port = port
	return "stdin", nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	fname := filepath.Base(hostSrc)
	if err := osutil.CopyFile(hostSrc, filepath.Join(inst.imageDir, fname)); err != nil {
		return "", err
	}
	if err := os.Chmod(inst.imageDir, 0777); err != nil {
		return "", err
	}
	return filepath.Join("/", fname), nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
	<-chan []byte, <-chan error, error) {
	args := []string{"exec", "-user=0:0"}
	for _, c := range sandboxCaps {
		args = append(args, "-cap", c)
	}
	args = append(args, inst.name)
	args = append(args, strings.Split(command, " ")...)
	cmd := inst.runscCmd(args...)

	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	defer wpipe.Close()
	inst.merger.Add("cmd", rpipe)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe

	guestSock, err := inst.guestProxy()
	if err != nil {
		return nil, nil, err
	}
	if guestSock != nil {
		defer guestSock.Close()
		cmd.Stdin = guestSock
	}

	if err := cmd.Start(); err != nil {
		return nil, nil, err
	}
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	go func() {
		select {
		case <-time.After(timeout):
			signal(vmimpl.ErrTimeout)
		case <-stop:
			signal(vmimpl.ErrTimeout)
		case err := <-inst.merger.Err:
			cmd.Process.Kill()
			if cmdErr := cmd.Wait(); cmdErr == nil {
				// If the command exited successfully, we got EOF error from merger.
				// But in this case no error has happened and the EOF is expected.
				err = nil
			}
			signal(err)
			return
		}
		cmd.Process.Kill()
		cmd.Wait()
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) guestProxy() (*os.File, error) {
	if inst.port == 0 {
		return nil, nil
	}
	// One does not simply let gvisor guest connect to host tcp port.
	// We create a unix socket, pass it to guest in stdin.
	// Guest will use it instead of dialing manager directly.
	// On host we connect to manager tcp port and proxy between the tcp and unix connections.
	socks, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
	if err != nil {
		return nil, err
	}
	hostSock := os.NewFile(uintptr(socks[0]), "host unix proxy")
	guestSock := os.NewFile(uintptr(socks[1]), "guest unix proxy")
	conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%v", inst.port))
	if err != nil {
		conn.Close()
		hostSock.Close()
		guestSock.Close()
		return nil, err
	}
	go func() {
		io.Copy(hostSock, conn)
		hostSock.Close()
	}()
	go func() {
		io.Copy(conn, hostSock)
		conn.Close()
	}()
	return guestSock, nil
}

func (inst *instance) Diagnose() bool {
	osutil.Run(time.Minute, inst.runscCmd("debug", "-stacks", inst.name))
	return true
}

func init() {
	if os.Getenv("SYZ_GVISOR_PROXY") != "" {
		fmt.Fprintf(os.Stderr, initStartMsg)
		select {}
	}
}

const initStartMsg = "SYZKALLER INIT STARTED\n"

const configTempl = `
{
	"root": {
		"path": "%[1]v",
		"readonly": true
	},
	"process":{
                "args": ["/init"],
                "cwd": "/tmp",
                "env": ["SYZ_GVISOR_PROXY=1"],
                "capabilities": {
                	"bounding": [%[2]v],
                	"effective": [%[2]v],
                	"inheritable": [%[2]v],
                	"permitted": [%[2]v],
                	"ambient": [%[2]v]
                }
	}
}
`

var sandboxCaps = []string{
	"CAP_CHOWN", "CAP_DAC_OVERRIDE", "CAP_DAC_READ_SEARCH", "CAP_FOWNER", "CAP_FSETID",
	"CAP_KILL", "CAP_SETGID", "CAP_SETUID", "CAP_SETPCAP", "CAP_LINUX_IMMUTABLE",
	"CAP_NET_BIND_SERVICE", "CAP_NET_BROADCAST", "CAP_NET_ADMIN", "CAP_NET_RAW",
	"CAP_IPC_LOCK", "CAP_IPC_OWNER", "CAP_SYS_MODULE", "CAP_SYS_RAWIO", "CAP_SYS_CHROOT",
	"CAP_SYS_PTRACE", "CAP_SYS_PACCT", "CAP_SYS_ADMIN", "CAP_SYS_BOOT", "CAP_SYS_NICE",
	"CAP_SYS_RESOURCE", "CAP_SYS_TIME", "CAP_SYS_TTY_CONFIG", "CAP_MKNOD", "CAP_LEASE",
	"CAP_AUDIT_WRITE", "CAP_AUDIT_CONTROL", "CAP_SETFCAP", "CAP_MAC_OVERRIDE", "CAP_MAC_ADMIN",
	"CAP_SYSLOG", "CAP_WAKE_ALARM", "CAP_BLOCK_SUSPEND", "CAP_AUDIT_READ",
}
<|endoftext|>


<commit_before>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package gvisor provides support for gVisor, user-space kernel, testing.
// See https://github.com/google/gvisor
package gvisor

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/google/syzkaller/pkg/config"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/vm/vmimpl"
)

func init() {
	vmimpl.Register("gvisor", ctor)
}

type Config struct {
	Count int `json:"count"` // number of VMs to use
}

type Pool struct {
	env *vmimpl.Env
	cfg *Config
}

type instance struct {
	cfg      *Config
	image    string
	debug    bool
	rootDir  string
	imageDir string
	name     string
	port     int
	cmd      *exec.Cmd
	merger   *vmimpl.OutputMerger
}

func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
	cfg := &Config{
		Count: 1,
	}
	if err := config.LoadData(env.Config, cfg); err != nil {
		return nil, fmt.Errorf("failed to parse vm config: %v", err)
	}
	if cfg.Count < 1 || cfg.Count > 1000 {
		return nil, fmt.Errorf("invalid config param count: %v, want [1, 1000]", cfg.Count)
	}
	if env.Debug {
		cfg.Count = 1
	}
	if !osutil.IsExist(env.Image) {
		return nil, fmt.Errorf("image file %q does not exist", env.Image)
	}
	pool := &Pool{
		cfg: cfg,
		env: env,
	}
	return pool, nil
}

func (pool *Pool) Count() int {
	return pool.cfg.Count
}

func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
	rootDir := filepath.Clean(filepath.Join(workdir, "..", "gvisor_root"))
	imageDir := filepath.Join(workdir, "image")
	bundleDir := filepath.Join(workdir, "bundle")
	osutil.MkdirAll(rootDir)
	osutil.MkdirAll(bundleDir)
	osutil.MkdirAll(imageDir)
	vmConfig := fmt.Sprintf(configTempl, imageDir)
	if err := osutil.WriteFile(filepath.Join(bundleDir, "config.json"), []byte(vmConfig)); err != nil {
		return nil, err
	}
	if err := osutil.CopyFile(os.Args[0], filepath.Join(imageDir, "init")); err != nil {
		return nil, err
	}

	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, err
	}
	var tee io.Writer
	if pool.env.Debug {
		tee = os.Stdout
	}
	merger := vmimpl.NewOutputMerger(tee)
	merger.Add("gvisor", rpipe)

	inst := &instance{
		cfg:      pool.cfg,
		image:    pool.env.Image,
		debug:    pool.env.Debug,
		rootDir:  rootDir,
		imageDir: imageDir,
		name:     fmt.Sprintf("%v-%v", pool.env.Name, index),
		merger:   merger,
	}

	// Kill the previous instance in case it's still running.
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	time.Sleep(3 * time.Second)

	cmd := inst.runscCmd("run", "-bundle", bundleDir, inst.name)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		merger.Wait()
		return nil, err
	}
	inst.cmd = cmd
	wpipe.Close()

	if err := inst.waitBoot(); err != nil {
		inst.Close()
		return nil, err
	}
	return inst, nil
}

func (inst *instance) waitBoot() error {
	errorMsg := []byte("FATAL ERROR:")
	bootedMsg := []byte(initStartMsg)
	timeout := time.NewTimer(time.Minute)
	defer timeout.Stop()
	var output []byte
	for {
		select {
		case out := <-inst.merger.Output:
			output = append(output, out...)
			if pos := bytes.Index(output, errorMsg); pos != -1 {
				end := bytes.IndexByte(output[pos:], '\n')
				if end == -1 {
					end = len(output)
				} else {
					end += pos
				}
				return vmimpl.BootError{
					Title:  string(output[pos:end]),
					Output: output,
				}
			}
			if bytes.Contains(output, bootedMsg) {
				return nil
			}
		case err := <-inst.merger.Err:
			return vmimpl.BootError{
				Title:  fmt.Sprintf("runsc failed: %v", err),
				Output: output,
			}
		case <-timeout.C:
			return vmimpl.BootError{
				Title:  "init process did not start",
				Output: output,
			}
		}
	}
}

func (inst *instance) runscCmd(add ...string) *exec.Cmd {
	args := []string{
		"-root", inst.rootDir,
		"-network=none",
	}
	args = append(args, add...)
	cmd := osutil.Command(inst.image, args...)
	cmd.Env = []string{
		"GOTRACEBACK=all",
		"GORACE=halt_on_error=1",
	}
	return cmd
}

func (inst *instance) Close() {
	time.Sleep(3 * time.Second)
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	inst.cmd.Process.Kill()
	inst.merger.Wait()
	inst.cmd.Wait()
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	time.Sleep(3 * time.Second)
}

func (inst *instance) Forward(port int) (string, error) {
	if inst.port != 0 {
		return "", fmt.Errorf("forward port is already setup")
	}
	inst.port = port
	return "stdin", nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	fname := filepath.Base(hostSrc)
	if err := osutil.CopyFile(hostSrc, filepath.Join(inst.imageDir, fname)); err != nil {
		return "", err
	}
	if err := os.Chmod(inst.imageDir, 0777); err != nil {
		return "", err
	}
	return filepath.Join("/", fname), nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
	<-chan []byte, <-chan error, error) {
	args := []string{"exec", "-user=0:0"}
	for _, c := range sandboxCaps {
		args = append(args, "-cap", c)
	}
	args = append(args, inst.name)
	args = append(args, strings.Split(command, " ")...)
	cmd := inst.runscCmd(args...)

	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	defer wpipe.Close()
	inst.merger.Add("cmd", rpipe)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe

	guestSock, err := inst.guestProxy()
	if err != nil {
		return nil, nil, err
	}
	if guestSock != nil {
		defer guestSock.Close()
		cmd.Stdin = guestSock
	}

	if err := cmd.Start(); err != nil {
		return nil, nil, err
	}
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	go func() {
		select {
		case <-time.After(timeout):
			signal(vmimpl.ErrTimeout)
		case <-stop:
			signal(vmimpl.ErrTimeout)
		case err := <-inst.merger.Err:
			cmd.Process.Kill()
			if cmdErr := cmd.Wait(); cmdErr == nil {
				// If the command exited successfully, we got EOF error from merger.
				// But in this case no error has happened and the EOF is expected.
				err = nil
			}
			signal(err)
			return
		}
		cmd.Process.Kill()
		cmd.Wait()
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) guestProxy() (*os.File, error) {
	if inst.port == 0 {
		return nil, nil
	}
	// One does not simply let gvisor guest connect to host tcp port.
	// We create a unix socket, pass it to guest in stdin.
	// Guest will use it instead of dialing manager directly.
	// On host we connect to manager tcp port and proxy between the tcp and unix connections.
	socks, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
	if err != nil {
		return nil, err
	}
	hostSock := os.NewFile(uintptr(socks[0]), "host unix proxy")
	guestSock := os.NewFile(uintptr(socks[1]), "guest unix proxy")
	conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%v", inst.port))
	if err != nil {
		conn.Close()
		hostSock.Close()
		guestSock.Close()
		return nil, err
	}
	go func() {
		io.Copy(hostSock, conn)
		hostSock.Close()
	}()
	go func() {
		io.Copy(conn, hostSock)
		conn.Close()
	}()
	return guestSock, nil
}

func (inst *instance) Diagnose() bool {
	osutil.Run(time.Minute, inst.runscCmd("debug", "-stacks", inst.name))
	return true
}

func init() {
	if os.Getenv("SYZ_GVISOR_PROXY") != "" {
		fmt.Fprintf(os.Stderr, initStartMsg)
		select {}
	}
}

const initStartMsg = "SYZKALLER INIT STARTED\n"

const configTempl = `
{
	"root": {
		"path": "%v",
		"readonly": true
	},
	"process":{
                "args": ["/init"],
                "cwd": "/tmp",
                "env": ["SYZ_GVISOR_PROXY=1"]
	}
}
`

var sandboxCaps = []string{
	"CAP_CHOWN", "CAP_DAC_OVERRIDE", "CAP_DAC_READ_SEARCH", "CAP_FOWNER", "CAP_FSETID",
	"CAP_KILL", "CAP_SETGID", "CAP_SETUID", "CAP_SETPCAP", "CAP_LINUX_IMMUTABLE",
	"CAP_NET_BIND_SERVICE", "CAP_NET_BROADCAST", "CAP_NET_ADMIN", "CAP_NET_RAW",
	"CAP_IPC_LOCK", "CAP_IPC_OWNER", "CAP_SYS_MODULE", "CAP_SYS_RAWIO", "CAP_SYS_CHROOT",
	"CAP_SYS_PTRACE", "CAP_SYS_PACCT", "CAP_SYS_ADMIN", "CAP_SYS_BOOT", "CAP_SYS_NICE",
	"CAP_SYS_RESOURCE", "CAP_SYS_TIME", "CAP_SYS_TTY_CONFIG", "CAP_MKNOD", "CAP_LEASE",
	"CAP_AUDIT_WRITE", "CAP_AUDIT_CONTROL", "CAP_SETFCAP", "CAP_MAC_OVERRIDE", "CAP_MAC_ADMIN",
	"CAP_SYSLOG", "CAP_WAKE_ALARM", "CAP_BLOCK_SUSPEND", "CAP_AUDIT_READ",
}
<commit_msg>vm/gvisor: allow specifying arbitrary runsc flags<commit_after>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package gvisor provides support for gVisor, user-space kernel, testing.
// See https://github.com/google/gvisor
package gvisor

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/google/syzkaller/pkg/config"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/vm/vmimpl"
)

func init() {
	vmimpl.Register("gvisor", ctor)
}

type Config struct {
	Count     int    `json:"count"` // number of VMs to use
	RunscArgs string `json:"runsc_args"`
}

type Pool struct {
	env *vmimpl.Env
	cfg *Config
}

type instance struct {
	cfg      *Config
	image    string
	debug    bool
	rootDir  string
	imageDir string
	name     string
	port     int
	cmd      *exec.Cmd
	merger   *vmimpl.OutputMerger
}

func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
	cfg := &Config{
		Count: 1,
	}
	if err := config.LoadData(env.Config, cfg); err != nil {
		return nil, fmt.Errorf("failed to parse vm config: %v", err)
	}
	if cfg.Count < 1 || cfg.Count > 1000 {
		return nil, fmt.Errorf("invalid config param count: %v, want [1, 1000]", cfg.Count)
	}
	if env.Debug {
		cfg.Count = 1
	}
	if !osutil.IsExist(env.Image) {
		return nil, fmt.Errorf("image file %q does not exist", env.Image)
	}
	pool := &Pool{
		cfg: cfg,
		env: env,
	}
	return pool, nil
}

func (pool *Pool) Count() int {
	return pool.cfg.Count
}

func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
	rootDir := filepath.Clean(filepath.Join(workdir, "..", "gvisor_root"))
	imageDir := filepath.Join(workdir, "image")
	bundleDir := filepath.Join(workdir, "bundle")
	osutil.MkdirAll(rootDir)
	osutil.MkdirAll(bundleDir)
	osutil.MkdirAll(imageDir)
	vmConfig := fmt.Sprintf(configTempl, imageDir)
	if err := osutil.WriteFile(filepath.Join(bundleDir, "config.json"), []byte(vmConfig)); err != nil {
		return nil, err
	}
	if err := osutil.CopyFile(os.Args[0], filepath.Join(imageDir, "init")); err != nil {
		return nil, err
	}

	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, err
	}
	var tee io.Writer
	if pool.env.Debug {
		tee = os.Stdout
	}
	merger := vmimpl.NewOutputMerger(tee)
	merger.Add("gvisor", rpipe)

	inst := &instance{
		cfg:      pool.cfg,
		image:    pool.env.Image,
		debug:    pool.env.Debug,
		rootDir:  rootDir,
		imageDir: imageDir,
		name:     fmt.Sprintf("%v-%v", pool.env.Name, index),
		merger:   merger,
	}

	// Kill the previous instance in case it's still running.
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	time.Sleep(3 * time.Second)

	cmd := inst.runscCmd("run", "-bundle", bundleDir, inst.name)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		merger.Wait()
		return nil, err
	}
	inst.cmd = cmd
	wpipe.Close()

	if err := inst.waitBoot(); err != nil {
		inst.Close()
		return nil, err
	}
	return inst, nil
}

func (inst *instance) waitBoot() error {
	errorMsg := []byte("FATAL ERROR:")
	bootedMsg := []byte(initStartMsg)
	timeout := time.NewTimer(time.Minute)
	defer timeout.Stop()
	var output []byte
	for {
		select {
		case out := <-inst.merger.Output:
			output = append(output, out...)
			if pos := bytes.Index(output, errorMsg); pos != -1 {
				end := bytes.IndexByte(output[pos:], '\n')
				if end == -1 {
					end = len(output)
				} else {
					end += pos
				}
				return vmimpl.BootError{
					Title:  string(output[pos:end]),
					Output: output,
				}
			}
			if bytes.Contains(output, bootedMsg) {
				return nil
			}
		case err := <-inst.merger.Err:
			return vmimpl.BootError{
				Title:  fmt.Sprintf("runsc failed: %v", err),
				Output: output,
			}
		case <-timeout.C:
			return vmimpl.BootError{
				Title:  "init process did not start",
				Output: output,
			}
		}
	}
}

func (inst *instance) runscCmd(add ...string) *exec.Cmd {
	args := []string{
		"-root", inst.rootDir,
		"-network=none",
	}
	args = append(args, strings.Split(inst.cfg.RunscArgs, " ")...)
	args = append(args, add...)
	cmd := osutil.Command(inst.image, args...)
	cmd.Env = []string{
		"GOTRACEBACK=all",
		"GORACE=halt_on_error=1",
	}
	return cmd
}

func (inst *instance) Close() {
	time.Sleep(3 * time.Second)
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	inst.cmd.Process.Kill()
	inst.merger.Wait()
	inst.cmd.Wait()
	osutil.Run(time.Minute, inst.runscCmd("delete", "-force", inst.name))
	time.Sleep(3 * time.Second)
}

func (inst *instance) Forward(port int) (string, error) {
	if inst.port != 0 {
		return "", fmt.Errorf("forward port is already setup")
	}
	inst.port = port
	return "stdin", nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	fname := filepath.Base(hostSrc)
	if err := osutil.CopyFile(hostSrc, filepath.Join(inst.imageDir, fname)); err != nil {
		return "", err
	}
	if err := os.Chmod(inst.imageDir, 0777); err != nil {
		return "", err
	}
	return filepath.Join("/", fname), nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
	<-chan []byte, <-chan error, error) {
	args := []string{"exec", "-user=0:0"}
	for _, c := range sandboxCaps {
		args = append(args, "-cap", c)
	}
	args = append(args, inst.name)
	args = append(args, strings.Split(command, " ")...)
	cmd := inst.runscCmd(args...)

	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	defer wpipe.Close()
	inst.merger.Add("cmd", rpipe)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe

	guestSock, err := inst.guestProxy()
	if err != nil {
		return nil, nil, err
	}
	if guestSock != nil {
		defer guestSock.Close()
		cmd.Stdin = guestSock
	}

	if err := cmd.Start(); err != nil {
		return nil, nil, err
	}
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	go func() {
		select {
		case <-time.After(timeout):
			signal(vmimpl.ErrTimeout)
		case <-stop:
			signal(vmimpl.ErrTimeout)
		case err := <-inst.merger.Err:
			cmd.Process.Kill()
			if cmdErr := cmd.Wait(); cmdErr == nil {
				// If the command exited successfully, we got EOF error from merger.
				// But in this case no error has happened and the EOF is expected.
				err = nil
			}
			signal(err)
			return
		}
		cmd.Process.Kill()
		cmd.Wait()
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) guestProxy() (*os.File, error) {
	if inst.port == 0 {
		return nil, nil
	}
	// One does not simply let gvisor guest connect to host tcp port.
	// We create a unix socket, pass it to guest in stdin.
	// Guest will use it instead of dialing manager directly.
	// On host we connect to manager tcp port and proxy between the tcp and unix connections.
	socks, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
	if err != nil {
		return nil, err
	}
	hostSock := os.NewFile(uintptr(socks[0]), "host unix proxy")
	guestSock := os.NewFile(uintptr(socks[1]), "guest unix proxy")
	conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%v", inst.port))
	if err != nil {
		conn.Close()
		hostSock.Close()
		guestSock.Close()
		return nil, err
	}
	go func() {
		io.Copy(hostSock, conn)
		hostSock.Close()
	}()
	go func() {
		io.Copy(conn, hostSock)
		conn.Close()
	}()
	return guestSock, nil
}

func (inst *instance) Diagnose() bool {
	osutil.Run(time.Minute, inst.runscCmd("debug", "-stacks", inst.name))
	return true
}

func init() {
	if os.Getenv("SYZ_GVISOR_PROXY") != "" {
		fmt.Fprintf(os.Stderr, initStartMsg)
		select {}
	}
}

const initStartMsg = "SYZKALLER INIT STARTED\n"

const configTempl = `
{
	"root": {
		"path": "%v",
		"readonly": true
	},
	"process":{
                "args": ["/init"],
                "cwd": "/tmp",
                "env": ["SYZ_GVISOR_PROXY=1"]
	}
}
`

var sandboxCaps = []string{
	"CAP_CHOWN", "CAP_DAC_OVERRIDE", "CAP_DAC_READ_SEARCH", "CAP_FOWNER", "CAP_FSETID",
	"CAP_KILL", "CAP_SETGID", "CAP_SETUID", "CAP_SETPCAP", "CAP_LINUX_IMMUTABLE",
	"CAP_NET_BIND_SERVICE", "CAP_NET_BROADCAST", "CAP_NET_ADMIN", "CAP_NET_RAW",
	"CAP_IPC_LOCK", "CAP_IPC_OWNER", "CAP_SYS_MODULE", "CAP_SYS_RAWIO", "CAP_SYS_CHROOT",
	"CAP_SYS_PTRACE", "CAP_SYS_PACCT", "CAP_SYS_ADMIN", "CAP_SYS_BOOT", "CAP_SYS_NICE",
	"CAP_SYS_RESOURCE", "CAP_SYS_TIME", "CAP_SYS_TTY_CONFIG", "CAP_MKNOD", "CAP_LEASE",
	"CAP_AUDIT_WRITE", "CAP_AUDIT_CONTROL", "CAP_SETFCAP", "CAP_MAC_OVERRIDE", "CAP_MAC_ADMIN",
	"CAP_SYSLOG", "CAP_WAKE_ALARM", "CAP_BLOCK_SUSPEND", "CAP_AUDIT_READ",
}
<|endoftext|>


<commit_before>package heroku

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"

	"code.google.com/p/go-uuid/uuid"
)

var DefaultTransport = &Transport{}

var DefaultClient = &http.Client{
	Transport: DefaultTransport,
}

type Transport struct {
	// Username is the HTTP basic auth username for API calls made by this Client.
	Username string

	// Password is the HTTP basic auth password for API calls made by this Client.
	Password string

	// UserAgent to be provided in API requests. Set to DefaultUserAgent if not
	// specified.
	UserAgent string

	// Debug mode can be used to dump the full request and response to stdout.
	Debug bool

	// AdditionalHeaders are extra headers to add to each HTTP request sent by
	// this Client.
	AdditionalHeaders http.Header

	// Transport is the HTTP transport to use when making requests.
	// It will default to http.DefaultTransport if nil.
	Transport http.RoundTripper
}

func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	if t.Transport == nil {
		t.Transport = http.DefaultTransport
	}

	// Making a copy of the Request so that
	// we don't modify the Request we were given.
	req = cloneRequest(req)

	if t.UserAgent != "" {
		req.Header.Set("User-Agent", t.UserAgent)
	}

	req.Header.Set("Accept", "application/vnd.heroku+json; version=3")
	req.Header.Set("Request-Id", uuid.New())
	req.SetBasicAuth(t.Username, t.Password)
	for k, v := range t.AdditionalHeaders {
		req.Header[k] = v
	}

	if t.Debug {
		dump, err := httputil.DumpRequestOut(req, true)
		if err != nil {
			log.Println(err)
		} else {
			os.Stderr.Write(dump)
			os.Stderr.Write([]byte{'\n', '\n'})
		}
	}

	resp, err := t.Transport.RoundTrip(req)
	if err != nil {
		return nil, err
	}

	if t.Debug {
		dump, err := httputil.DumpResponse(resp, true)
		if err != nil {
			log.Println(err)
		} else {
			os.Stderr.Write(dump)
			os.Stderr.Write([]byte{'\n'})
		}
	}

	if err = checkResponse(resp); err != nil {
		return nil, err
	}

	return resp, nil
}

type Error struct {
	error
	ID  string
	URL string
}

func checkResponse(resp *http.Response) error {
	if resp.StatusCode/100 != 2 { // 200, 201, 202, etc
		var e struct {
			Message string
			ID      string
			URL     string `json:"url"`
		}
		err := json.NewDecoder(resp.Body).Decode(&e)
		if err != nil {
			return fmt.Errorf("encountered an error : %s", resp.Status)
		}
		return Error{error: errors.New(e.Message), ID: e.ID, URL: e.URL}
	}
	if msg := resp.Header.Get("X-Heroku-Warning"); msg != "" {
		log.Println(strings.TrimSpace(msg))
	}
	return nil
}

// cloneRequest returns a clone of the provided *http.Request.
func cloneRequest(req *http.Request) *http.Request {
	// shallow copy of the struct
	clone := new(http.Request)
	*clone = *req
	// deep copy of the Header
	clone.Header = make(http.Header)
	for k, s := range req.Header {
		clone.Header[k] = s
	}
	return clone
}
<commit_msg>Close http.Response.Body on err in Transport<commit_after>package heroku

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"

	"code.google.com/p/go-uuid/uuid"
)

var DefaultTransport = &Transport{}

var DefaultClient = &http.Client{
	Transport: DefaultTransport,
}

type Transport struct {
	// Username is the HTTP basic auth username for API calls made by this Client.
	Username string

	// Password is the HTTP basic auth password for API calls made by this Client.
	Password string

	// UserAgent to be provided in API requests. Set to DefaultUserAgent if not
	// specified.
	UserAgent string

	// Debug mode can be used to dump the full request and response to stdout.
	Debug bool

	// AdditionalHeaders are extra headers to add to each HTTP request sent by
	// this Client.
	AdditionalHeaders http.Header

	// Transport is the HTTP transport to use when making requests.
	// It will default to http.DefaultTransport if nil.
	Transport http.RoundTripper
}

func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	if t.Transport == nil {
		t.Transport = http.DefaultTransport
	}

	// Making a copy of the Request so that
	// we don't modify the Request we were given.
	req = cloneRequest(req)

	if t.UserAgent != "" {
		req.Header.Set("User-Agent", t.UserAgent)
	}

	req.Header.Set("Accept", "application/vnd.heroku+json; version=3")
	req.Header.Set("Request-Id", uuid.New())
	req.SetBasicAuth(t.Username, t.Password)
	for k, v := range t.AdditionalHeaders {
		req.Header[k] = v
	}

	if t.Debug {
		dump, err := httputil.DumpRequestOut(req, true)
		if err != nil {
			log.Println(err)
		} else {
			os.Stderr.Write(dump)
			os.Stderr.Write([]byte{'\n', '\n'})
		}
	}

	resp, err := t.Transport.RoundTrip(req)
	if err != nil {
		if resp != nil {
			resp.Body.Close()
		}
		return nil, err
	}

	if t.Debug {
		dump, err := httputil.DumpResponse(resp, true)
		if err != nil {
			log.Println(err)
		} else {
			os.Stderr.Write(dump)
			os.Stderr.Write([]byte{'\n'})
		}
	}

	if err = checkResponse(resp); err != nil {
		if resp != nil {
			resp.Body.Close()
		}
		return nil, err
	}

	return resp, nil
}

type Error struct {
	error
	ID  string
	URL string
}

func checkResponse(resp *http.Response) error {
	if resp.StatusCode/100 != 2 { // 200, 201, 202, etc
		var e struct {
			Message string
			ID      string
			URL     string `json:"url"`
		}
		err := json.NewDecoder(resp.Body).Decode(&e)
		if err != nil {
			return fmt.Errorf("encountered an error : %s", resp.Status)
		}
		return Error{error: errors.New(e.Message), ID: e.ID, URL: e.URL}
	}
	if msg := resp.Header.Get("X-Heroku-Warning"); msg != "" {
		log.Println(strings.TrimSpace(msg))
	}
	return nil
}

// cloneRequest returns a clone of the provided *http.Request.
func cloneRequest(req *http.Request) *http.Request {
	// shallow copy of the struct
	clone := new(http.Request)
	*clone = *req
	// deep copy of the Header
	clone.Header = make(http.Header)
	for k, s := range req.Header {
		clone.Header[k] = s
	}
	return clone
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2021 The GoPlus Authors (goplus.org). All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package format

import (
	"bytes"
	"go/types"
	"path"
	"strconv"

	"github.com/goplus/gop/ast"
	"github.com/goplus/gop/format"
	"github.com/goplus/gop/parser"
	"github.com/goplus/gop/token"
)

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

func GopstyleSource(src []byte, filename ...string) (ret []byte, err error) {
	var fname string
	if filename != nil {
		fname = filename[0]
	}
	fset := token.NewFileSet()
	var f *ast.File
	if f, err = parser.ParseFile(fset, fname, src, parser.ParseComments); err == nil {
		Gopstyle(f)
		var buf bytes.Buffer
		if err = format.Node(&buf, fset, f); err == nil {
			ret = buf.Bytes()
		}
	}
	return
}

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

func Gopstyle(file *ast.File) {
	if identEqual(file.Name, "main") {
		file.NoPkgDecl = true
	}
	if idx := findFuncDecl(file.Decls, "main"); idx >= 0 {
		last := len(file.Decls) - 1
		if idx == last {
			file.NoEntrypoint = true
			// TODO: idx != last: swap main func to last
			// TODO: should also swap file.Comments
			/*
				fn := file.Decls[idx]
				copy(file.Decls[idx:], file.Decls[idx+1:])
				file.Decls[last] = fn
			*/
		}
	}
	formatFile(file)
}

func findFuncDecl(decls []ast.Decl, name string) int {
	for i, decl := range decls {
		if fn, ok := decl.(*ast.FuncDecl); ok {
			if identEqual(fn.Name, name) {
				return i
			}
		}
	}
	return -1
}

func findDecl(decls []ast.Decl, v ast.Decl) int {
	for i, decl := range decls {
		if decl == v {
			return i
		}
	}
	return -1
}

func findSpec(specs []ast.Spec, v ast.Spec) int {
	for i, spec := range specs {
		if spec == v {
			return i
		}
	}
	return -1
}

func deleteDecl(decls []ast.Decl, v ast.Decl) []ast.Decl {
	if idx := findDecl(decls, v); idx >= 0 {
		decls = append(decls[:idx], decls[idx+1:]...)
	}
	return decls
}

func deleteSpec(specs []ast.Spec, v ast.Spec) []ast.Spec {
	if idx := findSpec(specs, v); idx >= 0 {
		specs = append(specs[:idx], specs[idx+1:]...)
	}
	return specs
}

func startWithLowerCase(v *ast.Ident) {
	if c := v.Name[0]; c >= 'A' && c <= 'Z' {
		v.Name = string(c+('a'-'A')) + v.Name[1:]
	}
}

func identEqual(v *ast.Ident, name string) bool {
	return v != nil && v.Name == name
}

func toString(l *ast.BasicLit) string {
	if l.Kind == token.STRING {
		s, err := strconv.Unquote(l.Value)
		if err == nil {
			return s
		}
	}
	panic("TODO: toString - convert ast.BasicLit to string failed")
}

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

type importCtx struct {
	pkgPath string
	decl    *ast.GenDecl
	spec    *ast.ImportSpec
	isUsed  bool
}

type formatCtx struct {
	imports map[string]*importCtx
	scope   *types.Scope
}

func (ctx *formatCtx) enterBlock() *types.Scope {
	old := ctx.scope
	ctx.scope = types.NewScope(old, token.NoPos, token.NoPos, "")
	return old
}

func (ctx *formatCtx) leaveBlock(old *types.Scope) {
	ctx.scope = old
}

func formatFile(file *ast.File) {
	var funcs []*ast.FuncDecl
	ctx := &formatCtx{
		imports: make(map[string]*importCtx),
		scope:   types.NewScope(nil, token.NoPos, token.NoPos, ""),
	}
	for _, decl := range file.Decls {
		switch v := decl.(type) {
		case *ast.FuncDecl:
			// delay the process, because package level vars need to be processed first.
			funcs = append(funcs, v)
		case *ast.GenDecl:
			switch v.Tok {
			case token.IMPORT:
				for _, item := range v.Specs {
					var spec = item.(*ast.ImportSpec)
					var pkgPath = toString(spec.Path)
					var name string
					if spec.Name == nil {
						name = path.Base(pkgPath) // TODO: open pkgPath to get pkgName
					} else {
						name = spec.Name.Name
						if name == "." || name == "_" {
							continue
						}
					}
					ctx.imports[name] = &importCtx{pkgPath: pkgPath, decl: v, spec: spec}
				}
			default:
				formatGenDecl(ctx, v)
			}
		}
	}

	for _, fn := range funcs {
		formatFuncDecl(ctx, fn)
	}
	for _, imp := range ctx.imports {
		if imp.pkgPath == "fmt" && !imp.isUsed {
			if len(imp.decl.Specs) == 1 {
				file.Decls = deleteDecl(file.Decls, imp.decl)
			} else {
				imp.decl.Specs = deleteSpec(imp.decl.Specs, imp.spec)
			}
		}
	}
}

func formatGenDecl(ctx *formatCtx, v *ast.GenDecl) {
	switch v.Tok {
	case token.VAR, token.CONST:
		for _, item := range v.Specs {
			spec := item.(*ast.ValueSpec)
			formatType(ctx, spec.Type, &spec.Type)
			formatExprs(ctx, spec.Values)
			for _, name := range spec.Names {
				o := types.NewParam(token.NoPos, nil, name.Name, types.Typ[types.UntypedNil])
				ctx.scope.Insert(o)
			}
		}
	case token.TYPE:
		for _, item := range v.Specs {
			spec := item.(*ast.TypeSpec)
			formatType(ctx, spec.Type, &spec.Type)
		}
	}
}

func formatFuncDecl(ctx *formatCtx, v *ast.FuncDecl) {
	formatFuncType(ctx, v.Type)
	formatBlockStmt(ctx, v.Body)
}

/*
func fillVarCtx(ctx *formatCtx, spec *ast.ValueSpec) {
	for _, name := range spec.Names {
		ctx.scp.addVar(name.Name)
	}
}

type scope struct {
	vars []map[string]bool
}

func newScope() *scope {
	return &scope{
		vars: []map[string]bool{make(map[string]bool)},
	}
}

func (s *scope) addVar(name string) {
	s.vars[len(s.vars)-1][name] = true
}

func (s *scope) containsVar(name string) bool {
	for _, m := range s.vars {
		if m[name] {
			return true
		}
	}

	return false
}

func (s *scope) enterScope() {
	s.vars = append(s.vars, make(map[string]bool))
}

func (s *scope) exitScope() {
	s.vars = s.vars[:len(s.vars)-1]
}
*/
// -----------------------------------------------------------------------------
<commit_msg>ctx.insert<commit_after>/*
 * Copyright (c) 2021 The GoPlus Authors (goplus.org). All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package format

import (
	"bytes"
	"go/types"
	"path"
	"strconv"

	"github.com/goplus/gop/ast"
	"github.com/goplus/gop/format"
	"github.com/goplus/gop/parser"
	"github.com/goplus/gop/token"
)

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

func GopstyleSource(src []byte, filename ...string) (ret []byte, err error) {
	var fname string
	if filename != nil {
		fname = filename[0]
	}
	fset := token.NewFileSet()
	var f *ast.File
	if f, err = parser.ParseFile(fset, fname, src, parser.ParseComments); err == nil {
		Gopstyle(f)
		var buf bytes.Buffer
		if err = format.Node(&buf, fset, f); err == nil {
			ret = buf.Bytes()
		}
	}
	return
}

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

func Gopstyle(file *ast.File) {
	if identEqual(file.Name, "main") {
		file.NoPkgDecl = true
	}
	if idx := findFuncDecl(file.Decls, "main"); idx >= 0 {
		last := len(file.Decls) - 1
		if idx == last {
			file.NoEntrypoint = true
			// TODO: idx != last: swap main func to last
			// TODO: should also swap file.Comments
			/*
				fn := file.Decls[idx]
				copy(file.Decls[idx:], file.Decls[idx+1:])
				file.Decls[last] = fn
			*/
		}
	}
	formatFile(file)
}

func findFuncDecl(decls []ast.Decl, name string) int {
	for i, decl := range decls {
		if fn, ok := decl.(*ast.FuncDecl); ok {
			if identEqual(fn.Name, name) {
				return i
			}
		}
	}
	return -1
}

func findDecl(decls []ast.Decl, v ast.Decl) int {
	for i, decl := range decls {
		if decl == v {
			return i
		}
	}
	return -1
}

func findSpec(specs []ast.Spec, v ast.Spec) int {
	for i, spec := range specs {
		if spec == v {
			return i
		}
	}
	return -1
}

func deleteDecl(decls []ast.Decl, v ast.Decl) []ast.Decl {
	if idx := findDecl(decls, v); idx >= 0 {
		decls = append(decls[:idx], decls[idx+1:]...)
	}
	return decls
}

func deleteSpec(specs []ast.Spec, v ast.Spec) []ast.Spec {
	if idx := findSpec(specs, v); idx >= 0 {
		specs = append(specs[:idx], specs[idx+1:]...)
	}
	return specs
}

func startWithLowerCase(v *ast.Ident) {
	if c := v.Name[0]; c >= 'A' && c <= 'Z' {
		v.Name = string(c+('a'-'A')) + v.Name[1:]
	}
}

func identEqual(v *ast.Ident, name string) bool {
	return v != nil && v.Name == name
}

func toString(l *ast.BasicLit) string {
	if l.Kind == token.STRING {
		s, err := strconv.Unquote(l.Value)
		if err == nil {
			return s
		}
	}
	panic("TODO: toString - convert ast.BasicLit to string failed")
}

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

type importCtx struct {
	pkgPath string
	decl    *ast.GenDecl
	spec    *ast.ImportSpec
	isUsed  bool
}

type formatCtx struct {
	imports map[string]*importCtx
	scope   *types.Scope
}

func (ctx *formatCtx) insert(name string) {
	o := types.NewParam(token.NoPos, nil, name, types.Typ[types.UntypedNil])
	ctx.scope.Insert(o)
}

func (ctx *formatCtx) enterBlock() *types.Scope {
	old := ctx.scope
	ctx.scope = types.NewScope(old, token.NoPos, token.NoPos, "")
	return old
}

func (ctx *formatCtx) leaveBlock(old *types.Scope) {
	ctx.scope = old
}

func formatFile(file *ast.File) {
	var funcs []*ast.FuncDecl
	ctx := &formatCtx{
		imports: make(map[string]*importCtx),
		scope:   types.NewScope(nil, token.NoPos, token.NoPos, ""),
	}
	for _, decl := range file.Decls {
		switch v := decl.(type) {
		case *ast.FuncDecl:
			// delay the process, because package level vars need to be processed first.
			funcs = append(funcs, v)
		case *ast.GenDecl:
			switch v.Tok {
			case token.IMPORT:
				for _, item := range v.Specs {
					var spec = item.(*ast.ImportSpec)
					var pkgPath = toString(spec.Path)
					var name string
					if spec.Name == nil {
						name = path.Base(pkgPath) // TODO: open pkgPath to get pkgName
					} else {
						name = spec.Name.Name
						if name == "." || name == "_" {
							continue
						}
					}
					ctx.imports[name] = &importCtx{pkgPath: pkgPath, decl: v, spec: spec}
				}
			default:
				formatGenDecl(ctx, v)
			}
		}
	}

	for _, fn := range funcs {
		formatFuncDecl(ctx, fn)
	}
	for _, imp := range ctx.imports {
		if imp.pkgPath == "fmt" && !imp.isUsed {
			if len(imp.decl.Specs) == 1 {
				file.Decls = deleteDecl(file.Decls, imp.decl)
			} else {
				imp.decl.Specs = deleteSpec(imp.decl.Specs, imp.spec)
			}
		}
	}
}

func formatGenDecl(ctx *formatCtx, v *ast.GenDecl) {
	switch v.Tok {
	case token.VAR, token.CONST:
		for _, item := range v.Specs {
			spec := item.(*ast.ValueSpec)
			formatType(ctx, spec.Type, &spec.Type)
			formatExprs(ctx, spec.Values)
			for _, name := range spec.Names {
				ctx.insert(name.Name)
			}
		}
	case token.TYPE:
		for _, item := range v.Specs {
			spec := item.(*ast.TypeSpec)
			formatType(ctx, spec.Type, &spec.Type)
		}
	}
}

func formatFuncDecl(ctx *formatCtx, v *ast.FuncDecl) {
	formatFuncType(ctx, v.Type)
	formatBlockStmt(ctx, v.Body)
}

/*
func fillVarCtx(ctx *formatCtx, spec *ast.ValueSpec) {
	for _, name := range spec.Names {
		ctx.scp.addVar(name.Name)
	}
}

type scope struct {
	vars []map[string]bool
}

func newScope() *scope {
	return &scope{
		vars: []map[string]bool{make(map[string]bool)},
	}
}

func (s *scope) addVar(name string) {
	s.vars[len(s.vars)-1][name] = true
}

func (s *scope) containsVar(name string) bool {
	for _, m := range s.vars {
		if m[name] {
			return true
		}
	}

	return false
}

func (s *scope) enterScope() {
	s.vars = append(s.vars, make(map[string]bool))
}

func (s *scope) exitScope() {
	s.vars = s.vars[:len(s.vars)-1]
}
*/
// -----------------------------------------------------------------------------
<|endoftext|>


<commit_before>package composers

import (
	"math/rand"

	"errors"

	m "github.com/borodiychuk/patchwork/models"
)

// Shuffle just picks random sample and turns it by random angle. Shuffle guarantees that similar elements have no common sides
type Shuffle struct{}

// Compose composes a picture out of provided set of samples
func (c *Shuffle) Compose(canvas *m.Canvas, samples []m.Sample) error {
	// Elements to exclude from list to avoid similar elements staying together
	topElement := ""
	leftElement := ""

	for i := 0; i < canvas.Length*canvas.Width; i++ {
		col, row, err := canvas.XYforIndex(i)
		if err != nil {
			return err
		}
		if row > 0 {
			topElementIndex, err := canvas.IndexForXY(col, row-1)
			if err != nil {
				return err
			}
			topElement = canvas.Elements[topElementIndex].Sample.ID()
		}
		// This piece of code with cycle is not really optimal and could utilize some caching for better performance. But from other side,
		// that is not a tool for real-time data processing, so nothing bad happens if it stays so: not so fast
		// but semantically correct.
		allowedSamples := []m.Sample{}
		for _, s := range samples {
			if s.ID() != topElement && s.ID() != leftElement {
				allowedSamples = append(allowedSamples, s)
			}
		}
		if len(allowedSamples) == 0 {
			return errors.New("Too few samples to compose the pattern")
		}
		e := &m.Element{}
		e.Sample = allowedSamples[rand.Intn(len(allowedSamples))]
		e.Rotation = rand.Intn(4)
		canvas.Elements = append(canvas.Elements, e)
		leftElement = e.Sample.ID()
	}
	return nil
}

// Seed seeds the ranom number generator
func (c *Shuffle) Seed(seed int64) {
	rand.Seed(seed)
}
<commit_msg>Optimize shuffle pattern<commit_after>package composers

import (
	"math/rand"

	"errors"

	m "github.com/borodiychuk/patchwork/models"
)

// Shuffle just picks random sample and turns it by random angle. Shuffle guarantees that similar elements have no common sides
type Shuffle struct{}

// Compose composes a picture out of provided set of samples
func (c *Shuffle) Compose(canvas *m.Canvas, samples []m.Sample) error {
	// Elements to exclude from list to avoid similar elements staying together
	forbiddenSamples := []string{}

	for i := 0; i < canvas.Length*canvas.Width; i++ {
		col, row, err := canvas.XYforIndex(i)
		if err != nil {
			return err
		}
		// This piece of code with cycle is not really optimal and could utilize some caching for better performance. But from other side,
		// that is not a tool for real-time data processing, so nothing bad happens if it stays so: not so fast
		// but semantically correct.
		allowedSamples := []m.Sample{}
		for _, s := range samples {
			allowed := true
			for _, fs := range forbiddenSamples {
				if s.ID() == fs {
					allowed = false
					break
				}
			}
			if allowed {
				allowedSamples = append(allowedSamples, s)
			}
		}
		if len(allowedSamples) == 0 {
			return errors.New("Too few samples to compose the pattern")
		}
		e := &m.Element{}
		e.Sample = allowedSamples[rand.Intn(len(allowedSamples))]
		e.Rotation = rand.Intn(4)
		canvas.Elements = append(canvas.Elements, e)

		forbiddenSamples = []string{e.Sample.ID()}
		if topElementIndex, err := canvas.IndexForXY(col, row-1); err == nil {
			forbiddenSamples = append(forbiddenSamples, canvas.Elements[topElementIndex].Sample.ID())
		}
		if topLeftElementIndex, err := canvas.IndexForXY(col-1, row-1); err == nil {
			forbiddenSamples = append(forbiddenSamples, canvas.Elements[topLeftElementIndex].Sample.ID())
		}
		if topRightElementIndex, err := canvas.IndexForXY(col+1, row-1); err == nil {
			forbiddenSamples = append(forbiddenSamples, canvas.Elements[topRightElementIndex].Sample.ID())
		}
	}
	return nil
}

// Seed seeds the ranom number generator
func (c *Shuffle) Seed(seed int64) {
	rand.Seed(seed)
}
<|endoftext|>


<commit_before>package mesoslog

// Package sandboxfinder helps to look up the sandbox page on the mesos-master
// where the stdout and stderr streams for a running task can be viewed

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	// PageLength is the amount of data we want to consume during log tailing
	PageLength int    = 5000
	TailURIFmt string = "http://%s:5051/files/read.json?path=%s&offset=%v&length=%v"
)

// MesosClient holds state about the current Master node.  Allows method receivers to obtain these values
type MesosClient struct {
	Host      string
	Port      int
	MasterURL string
	State     *masterState
	Options   *MesosClientOptions
}

type MesosClientOptions struct {
	SearchCompletedTasks bool
	ShowLatestOnly       bool
}

// NewMesosClient - Creates a new MesosClient
// {host} - the host/ip of the mesos master node
// {port} - the port # of the mesos master node
func NewMesosClient(host string, port int) (*MesosClient, error) {
	return NewMesosClientWithOptions(host, port, nil)
}

// NewMesosClient - Creates a new MesosClient
// {host} - the host/ip of the mesos master node
// {port} - the port # of the mesos master node
// {options} - client options - optional
func NewMesosClientWithOptions(host string, port int, options *MesosClientOptions) (*MesosClient, error) {
	masterURL, err := getMasterRedirect(host, port)
	if err != nil {
		return nil, err
	}

	if options == nil {
		options = &MesosClientOptions{}
	}

	state, err := getMasterState(masterURL)
	if err != nil {
		return nil, err
	}
	return &MesosClient{
		Host:      host,
		Port:      port,
		MasterURL: masterURL,
		State:     state,
		Options:   options,
	}, nil
}

// GetAppNames - List all unique app names aka task names running in the Mesos cluster
func (c *MesosClient) GetAppNames() (map[string]int, error) {
	apps := findApps(c.State)
	if apps == nil || len(apps) == 0 {
		return nil, fmt.Errorf("Applications could not be retrieved")
	}
	return apps, nil
}

// GetLog - Gets/Downloads logs for a [appID]
// {appID} - the task name / app identifier
// {logtype} - the desired log type STDOUT | STDERR
// {dir} - optional output dir which is used to download vs stdout
func (c *MesosClient) GetLog(appID string, logtype LogType, dir string) ([]*LogOut, error) {
	var result []*LogOut

	taskInfo := findTask(c.State, appID)
	tasks := taskInfo.Tasks
	if c.Options.SearchCompletedTasks {
		tasks = taskInfo.CompletedTasks
	}

	sort.Sort(SortTasksByLatestTimestamp(tasks))

	if tasks == nil || len(tasks) == 0 {
		return nil, fmt.Errorf("application could not be found")
	}

	if c.Options.ShowLatestOnly {
		tasks = tasks[:1]
	}

	for _, task := range tasks {
		slaveInfo, err := c.getSlaveInfo(task)
		if err != nil {
			return nil, err
		}

		url := fmt.Sprintf("http://%s:5051/files/download.json?path=%s/", slaveInfo.Slave.Hostname, slaveInfo.Directory)

		var filename string
		if dir != "" {
			filename = filepath.Join(dir, fmt.Sprintf("%s_%s.txt", task.ID, logtype.String()))
		}
		data, err := download(url+logtype.String(), filename)
		if err != nil {
			return nil, err
		}
		result = append(result, &LogOut{TaskID: task.ID, AppID: appID, Log: data})
	}
	return result, nil
}

// TailLog - Tails the logs for a [appID]
// {appID} - the task name / app identifier
// {logtype} - the desired log type STDOUT | STDERR
// {duration} - poll frequency in seconds
func (c *MesosClient) TailLog(appID string, logtype LogType, duration int) error {
	tasks := findTask(c.State, appID).Tasks

	if tasks == nil || len(tasks) == 0 {
		return fmt.Errorf("application could not be found")
	}

	var chans []<-chan string
	for _, task := range tasks {

		slaveInfo, err := c.getSlaveInfo(task)
		if err != nil {
			return err
		}
		o := c.asyncTail(task, slaveInfo, logtype, duration)
		chans = append(chans, o)

	}
	output := merge(chans...)
	for {
		fmt.Print(<-output)
	}
	return nil
}

// GetAppNameForTaskID - Attempts to find the Mesos Application name for the given TaskID
// {taskID} - the task identifier
func (c *MesosClient) GetAppNameForTaskID(taskID string) (string, error) {
	ti := findTask(c.State, taskID)
	if ti != nil {
		if c.Options.SearchCompletedTasks == false {
			for _, task := range ti.Tasks {
				return task.Name, nil
			}
		} else {
			for _, task := range ti.CompletedTasks {
				return task.Name, nil
			}
		}
	}
	return "", fmt.Errorf("application could not be found")

}

func (c *MesosClient) GetAppNameForPath(path string) string {
	pathArr := strings.Split(path, "/")
	if strings.HasPrefix(path, "/") {
		pathArr = pathArr[1:]
	}

	rpath := reversePath(pathArr)
	return strings.Join(rpath, ".")
}

func reversePath(s []string) []string {
	for i := 0; i < len(s)/2; i++ {
		j := len(s) - i - 1
		s[i], s[j] = s[j], s[i]
	}
	return s
}

func (c *MesosClient) asyncTail(task *mstateTask, s *slaveInfo, lt LogType, duration int) <-chan string {
	ch := make(chan string)
	path := fmt.Sprintf("%s/%s", s.Directory, lt.String())
	go func() {
		offset := 0
		for {
			url := fmt.Sprintf(TailURIFmt, s.Slave.Hostname, path, offset, PageLength)

			resp, err := download(url, "")
			if err != nil {
				fmt.Printf("Error: %s", err.Error())
				continue
			}
			var rd readData
			json.Unmarshal([]byte(resp), &rd)

			if len(rd.Data) < 5 {
				time.Sleep(time.Duration(duration) * time.Second)
				continue
			}
			offset += len(rd.Data)
			ch <- decorateLog(task.ID, rd.Data)
		}
	}()
	return ch
}

func merge(cs ...<-chan string) <-chan string {
	var wg sync.WaitGroup
	out := make(chan string)

	// Start an output goroutine for each input channel in cs.  output
	// copies values from c to out until c is closed, then calls wg.Done.
	output := func(c <-chan string) {
		for n := range c {
			out <- n
		}
		wg.Done()
	}
	wg.Add(len(cs))
	for _, c := range cs {
		go output(c)
	}

	// Start a goroutine to close out once all the output goroutines are
	// done.  This must start after the wg.Add call.
	go func() {
		wg.Wait()
		close(out)
	}()
	return out
}

func decorateLog(name, data string) string {
	lines := strings.Split(data, "\n")
	buf := new(bytes.Buffer)
	taskIdx := strings.Index(name, ".")
	dec := name[0:taskIdx]
	taskId := name[taskIdx+1:]
	dec = fmt.Sprintf("%s.%s", dec, taskId[0:strings.Index(taskId, "-")])
	for _, l := range lines {
		if len(l) > 0 {
			buf.WriteString(fmt.Sprintf("[%s] %s\n", dec, l))
		}
	}
	return buf.String()
}

func (c *MesosClient) getSlaveInfo(task *mstateTask) (*slaveInfo, error) {

	slave := findSlave(c.State, task.SlaveID)
	if slave == nil {
		return nil, fmt.Errorf("invalid state.json; referenced slave not present")
	}

	slaveURL, err := constructSlaveURL(slave)
	if err != nil {
		return nil, err
	}

	slaveState, err := getSlaveState(slaveURL)
	if err != nil {
		return nil, err
	}

	directory := findDirectory(slaveState, task, c.Options.SearchCompletedTasks)

	if directory == "" {
		return nil, fmt.Errorf("couldn't locate directory on slave")
	}
	return &slaveInfo{Slave: slave, State: slaveState, Directory: directory}, nil
}

func getMasterRedirect(host string, port int) (string, error) {
	url := fmt.Sprintf("http://%s:%d/master/redirect", host, port)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return "", err
	}
	resp, err := http.DefaultTransport.RoundTrip(req)
	if err != nil {
		return "", err
	}
	resp.Body.Close()
	loc, err := resp.Location()
	if err != nil {
		return "", err
	}
	return loc.String(), nil
}

func getMasterState(masterURL string) (*masterState, error) {
	url := masterURL + "/state.json"
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var mstate masterState
	err = json.Unmarshal(data, &mstate)
	if err != nil {
		return nil, err
	}
	return &mstate, nil
}

func findTask(state *masterState, appID string) *taskInfo {
	taskInfo := &taskInfo{
		Tasks:          []*mstateTask{},
		CompletedTasks: []*mstateTask{},
	}
	for _, framework := range state.Frameworks {
		for _, task := range framework.Tasks {
			if task.Name == appID || task.ID == appID {
				task.UpdateLastState(findTaskLastState(task))
				taskInfo.Tasks = append(taskInfo.Tasks, task)
			}
		}

		if framework.CompletedTasks != nil {
			for _, task := range framework.CompletedTasks {
				if task.Name == appID || task.ID == appID {
					task.UpdateLastState(findTaskLastState(task))
					taskInfo.CompletedTasks = append(taskInfo.CompletedTasks, task)
				}
			}
		}
	}
	return taskInfo
}

func findTaskLastState(task *mstateTask) *mstateTaskStatus {
	statuses := task.Statuses

	if statuses != nil && len(statuses) > 0 {
		return statuses[len(statuses)-1]
	}
	return nil
}

func findApps(state *masterState) map[string]int {
	m := make(map[string]int)
	for _, framework := range state.Frameworks {
		for _, task := range framework.Tasks {
			m[task.Name]++
		}
	}
	return m
}

func findSlave(state *masterState, slaveID string) *mstateSlave {
	for _, slave := range state.Slaves {
		if slave.ID == slaveID {
			return slave
		}
	}
	return nil
}

func constructSlaveURL(slave *mstateSlave) (*url.URL, error) {
	parts := strings.SplitN(slave.Pid, "@", 2)
	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid slave pid %s", slave.Pid)
	}
	slaveName := parts[0]

	hostAndPort := strings.Split(parts[1], ":")
	port := "80"
	if len(hostAndPort) > 1 {
		port = hostAndPort[1]
	}

	host := fmt.Sprintf("%s:%s", slave.Hostname, port)
	path := fmt.Sprintf("%s/state.json", slaveName)

	return &url.URL{
		Scheme: "http",
		Host:   host,
		Opaque: fmt.Sprintf("//%s/%s", host, path),
	}, nil
}

func getSlaveState(slaveURL *url.URL) (*slaveState, error) {
	req, err := http.NewRequest("GET", slaveURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.URL = slaveURL
	client := http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var sstate slaveState
	err = json.Unmarshal(data, &sstate)
	if err != nil {
		return nil, err
	}
	return &sstate, nil
}

func findDirectory(sstate *slaveState, task *mstateTask, completedTasks bool) string {

	for _, framework := range sstate.Frameworks {
		if framework.ID != task.FrameworkID {
			continue
		}

		executors := framework.Executors
		if completedTasks {
			executors = framework.CompletedExecutors
		}

		for _, executor := range executors {
			if executor.ID == task.ExecutorID || executor.ID == task.ID {
				return executor.Directory
			}
		}
	}
	return ""
}

func download(slaveURL string, filename string) (string, error) {
	resp, err := http.Get(slaveURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if filename != "" {
		if e := writeFile(filename, resp.Body); e != nil {
			return "", e
		}
		return filename, nil
	}
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	return string(data), nil

}

func writeFile(filename string, r io.Reader) error {
	if err := os.MkdirAll(filepath.Dir(filename), 0700); err != nil {
		return err
	}
	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, r)
	if err != nil {
		return err
	}
	return nil

}
<commit_msg>Expose tail log func for consumption by other apps by adding ability to pass in channel to receive streaming logs<commit_after>package mesoslog

// Package sandboxfinder helps to look up the sandbox page on the mesos-master
// where the stdout and stderr streams for a running task can be viewed

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	// PageLength is the amount of data we want to consume during log tailing
	PageLength int    = 5000
	TailURIFmt string = "http://%s:5051/files/read.json?path=%s&offset=%v&length=%v"
)

// MesosClient holds state about the current Master node.  Allows method receivers to obtain these values
type MesosClient struct {
	Host      string
	Port      int
	MasterURL string
	State     *masterState
	Options   *MesosClientOptions
}

type MesosClientOptions struct {
	SearchCompletedTasks bool
	ShowLatestOnly       bool
}

// NewMesosClient - Creates a new MesosClient
// {host} - the host/ip of the mesos master node
// {port} - the port # of the mesos master node
func NewMesosClient(host string, port int) (*MesosClient, error) {
	return NewMesosClientWithOptions(host, port, nil)
}

// NewMesosClient - Creates a new MesosClient
// {host} - the host/ip of the mesos master node
// {port} - the port # of the mesos master node
// {options} - client options - optional
func NewMesosClientWithOptions(host string, port int, options *MesosClientOptions) (*MesosClient, error) {
	masterURL, err := getMasterRedirect(host, port)
	if err != nil {
		return nil, err
	}

	if options == nil {
		options = &MesosClientOptions{}
	}

	state, err := getMasterState(masterURL)
	if err != nil {
		return nil, err
	}
	return &MesosClient{
		Host:      host,
		Port:      port,
		MasterURL: masterURL,
		State:     state,
		Options:   options,
	}, nil
}

// GetAppNames - List all unique app names aka task names running in the Mesos cluster
func (c *MesosClient) GetAppNames() (map[string]int, error) {
	apps := findApps(c.State)
	if apps == nil || len(apps) == 0 {
		return nil, fmt.Errorf("Applications could not be retrieved")
	}
	return apps, nil
}

// GetLog - Gets/Downloads logs for a [appID]
// {appID} - the task name / app identifier
// {logtype} - the desired log type STDOUT | STDERR
// {dir} - optional output dir which is used to download vs stdout
func (c *MesosClient) GetLog(appID string, logtype LogType, dir string) ([]*LogOut, error) {
	var result []*LogOut

	taskInfo := findTask(c.State, appID)
	tasks := taskInfo.Tasks
	if c.Options.SearchCompletedTasks {
		tasks = taskInfo.CompletedTasks
	}

	sort.Sort(SortTasksByLatestTimestamp(tasks))

	if tasks == nil || len(tasks) == 0 {
		return nil, fmt.Errorf("application could not be found")
	}

	if c.Options.ShowLatestOnly {
		tasks = tasks[:1]
	}

	for _, task := range tasks {
		slaveInfo, err := c.getSlaveInfo(task)
		if err != nil {
			return nil, err
		}

		url := fmt.Sprintf("http://%s:5051/files/download.json?path=%s/", slaveInfo.Slave.Hostname, slaveInfo.Directory)

		var filename string
		if dir != "" {
			filename = filepath.Join(dir, fmt.Sprintf("%s_%s.txt", task.ID, logtype.String()))
		}
		data, err := download(url+logtype.String(), filename)
		if err != nil {
			return nil, err
		}
		result = append(result, &LogOut{TaskID: task.ID, AppID: appID, Log: data})
	}
	return result, nil
}

func (c *MesosClient) TailLogToChannel(appID string, logtype LogType, duration int, target chan<- string) error {
	tasks := findTask(c.State, appID).Tasks

	if tasks == nil || len(tasks) == 0 {
		return fmt.Errorf("application could not be found")
	}

	var chans []<-chan string
	for _, task := range tasks {

		slaveInfo, err := c.getSlaveInfo(task)
		if err != nil {
			return err
		}
		o := c.asyncTail(task, slaveInfo, logtype, duration)
		chans = append(chans, o)

	}
	output := merge(chans...)
	for msg := range output {
		target <- msg
	}
	return nil
}

// TailLog - Tails the logs for a [appID]
// {appID} - the task name / app identifier
// {logtype} - the desired log type STDOUT | STDERR
// {duration} - poll frequency in seconds
func (c *MesosClient) TailLog(appID string, logtype LogType, duration int) error {

	target := make(chan string)

	go c.TailLogToChannel(appID, logtype, duration, target)

	for msg := range target {
		fmt.Println(msg)
	}

	return nil
}

// GetAppNameForTaskID - Attempts to find the Mesos Application name for the given TaskID
// {taskID} - the task identifier
func (c *MesosClient) GetAppNameForTaskID(taskID string) (string, error) {
	ti := findTask(c.State, taskID)
	if ti != nil {
		if c.Options.SearchCompletedTasks == false {
			for _, task := range ti.Tasks {
				return task.Name, nil
			}
		} else {
			for _, task := range ti.CompletedTasks {
				return task.Name, nil
			}
		}
	}
	return "", fmt.Errorf("application could not be found")

}

func (c *MesosClient) GetAppNameForPath(path string) string {
	pathArr := strings.Split(path, "/")
	if strings.HasPrefix(path, "/") {
		pathArr = pathArr[1:]
	}

	rpath := reversePath(pathArr)
	return strings.Join(rpath, ".")
}

func reversePath(s []string) []string {
	for i := 0; i < len(s)/2; i++ {
		j := len(s) - i - 1
		s[i], s[j] = s[j], s[i]
	}
	return s
}

func (c *MesosClient) asyncTail(task *mstateTask, s *slaveInfo, lt LogType, duration int) <-chan string {
	ch := make(chan string)
	path := fmt.Sprintf("%s/%s", s.Directory, lt.String())
	go func() {
		offset := 0
		for {
			url := fmt.Sprintf(TailURIFmt, s.Slave.Hostname, path, offset, PageLength)

			resp, err := download(url, "")
			if err != nil {
				fmt.Printf("Error: %s", err.Error())
				continue
			}
			var rd readData
			json.Unmarshal([]byte(resp), &rd)

			if len(rd.Data) < 5 {
				time.Sleep(time.Duration(duration) * time.Second)
				continue
			}
			offset += len(rd.Data)
			ch <- decorateLog(task.ID, rd.Data)
		}
	}()
	return ch
}

func merge(cs ...<-chan string) <-chan string {
	var wg sync.WaitGroup
	out := make(chan string)

	// Start an output goroutine for each input channel in cs.  output
	// copies values from c to out until c is closed, then calls wg.Done.
	output := func(c <-chan string) {
		for n := range c {
			out <- n
		}
		wg.Done()
	}
	wg.Add(len(cs))
	for _, c := range cs {
		go output(c)
	}

	// Start a goroutine to close out once all the output goroutines are
	// done.  This must start after the wg.Add call.
	go func() {
		wg.Wait()
		close(out)
	}()
	return out
}

func decorateLog(name, data string) string {
	lines := strings.Split(data, "\n")
	buf := new(bytes.Buffer)
	taskIdx := strings.Index(name, ".")
	dec := name[0:taskIdx]
	taskId := name[taskIdx+1:]
	dec = fmt.Sprintf("%s.%s", dec, taskId[0:strings.Index(taskId, "-")])
	for _, l := range lines {
		if len(l) > 0 {
			buf.WriteString(fmt.Sprintf("[%s] %s\n", dec, l))
		}
	}
	return buf.String()
}

func (c *MesosClient) getSlaveInfo(task *mstateTask) (*slaveInfo, error) {

	slave := findSlave(c.State, task.SlaveID)
	if slave == nil {
		return nil, fmt.Errorf("invalid state.json; referenced slave not present")
	}

	slaveURL, err := constructSlaveURL(slave)
	if err != nil {
		return nil, err
	}

	slaveState, err := getSlaveState(slaveURL)
	if err != nil {
		return nil, err
	}

	directory := findDirectory(slaveState, task, c.Options.SearchCompletedTasks)

	if directory == "" {
		return nil, fmt.Errorf("couldn't locate directory on slave")
	}
	return &slaveInfo{Slave: slave, State: slaveState, Directory: directory}, nil
}

func getMasterRedirect(host string, port int) (string, error) {
	url := fmt.Sprintf("http://%s:%d/master/redirect", host, port)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return "", err
	}
	resp, err := http.DefaultTransport.RoundTrip(req)
	if err != nil {
		return "", err
	}
	resp.Body.Close()
	loc, err := resp.Location()
	if err != nil {
		return "", err
	}
	return loc.String(), nil
}

func getMasterState(masterURL string) (*masterState, error) {
	url := masterURL + "/state.json"
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var mstate masterState
	err = json.Unmarshal(data, &mstate)
	if err != nil {
		return nil, err
	}
	return &mstate, nil
}

func findTask(state *masterState, appID string) *taskInfo {
	taskInfo := &taskInfo{
		Tasks:          []*mstateTask{},
		CompletedTasks: []*mstateTask{},
	}
	for _, framework := range state.Frameworks {
		for _, task := range framework.Tasks {
			if task.Name == appID || task.ID == appID {
				task.UpdateLastState(findTaskLastState(task))
				taskInfo.Tasks = append(taskInfo.Tasks, task)
			}
		}

		if framework.CompletedTasks != nil {
			for _, task := range framework.CompletedTasks {
				if task.Name == appID || task.ID == appID {
					task.UpdateLastState(findTaskLastState(task))
					taskInfo.CompletedTasks = append(taskInfo.CompletedTasks, task)
				}
			}
		}
	}
	return taskInfo
}

func findTaskLastState(task *mstateTask) *mstateTaskStatus {
	statuses := task.Statuses

	if statuses != nil && len(statuses) > 0 {
		return statuses[len(statuses)-1]
	}
	return nil
}

func findApps(state *masterState) map[string]int {
	m := make(map[string]int)
	for _, framework := range state.Frameworks {
		for _, task := range framework.Tasks {
			m[task.Name]++
		}
	}
	return m
}

func findSlave(state *masterState, slaveID string) *mstateSlave {
	for _, slave := range state.Slaves {
		if slave.ID == slaveID {
			return slave
		}
	}
	return nil
}

func constructSlaveURL(slave *mstateSlave) (*url.URL, error) {
	parts := strings.SplitN(slave.Pid, "@", 2)
	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid slave pid %s", slave.Pid)
	}
	slaveName := parts[0]

	hostAndPort := strings.Split(parts[1], ":")
	port := "80"
	if len(hostAndPort) > 1 {
		port = hostAndPort[1]
	}

	host := fmt.Sprintf("%s:%s", slave.Hostname, port)
	path := fmt.Sprintf("%s/state.json", slaveName)

	return &url.URL{
		Scheme: "http",
		Host:   host,
		Opaque: fmt.Sprintf("//%s/%s", host, path),
	}, nil
}

func getSlaveState(slaveURL *url.URL) (*slaveState, error) {
	req, err := http.NewRequest("GET", slaveURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.URL = slaveURL
	client := http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var sstate slaveState
	err = json.Unmarshal(data, &sstate)
	if err != nil {
		return nil, err
	}
	return &sstate, nil
}

func findDirectory(sstate *slaveState, task *mstateTask, completedTasks bool) string {

	for _, framework := range sstate.Frameworks {
		if framework.ID != task.FrameworkID {
			continue
		}

		executors := framework.Executors
		if completedTasks {
			executors = framework.CompletedExecutors
		}

		for _, executor := range executors {
			if executor.ID == task.ExecutorID || executor.ID == task.ID {
				return executor.Directory
			}
		}
	}
	return ""
}

func download(slaveURL string, filename string) (string, error) {
	resp, err := http.Get(slaveURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if filename != "" {
		if e := writeFile(filename, resp.Body); e != nil {
			return "", e
		}
		return filename, nil
	}
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	return string(data), nil

}

func writeFile(filename string, r io.Reader) error {
	if err := os.MkdirAll(filepath.Dir(filename), 0700); err != nil {
		return err
	}
	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, r)
	if err != nil {
		return err
	}
	return nil

}
<|endoftext|>