<commit_before>package nvic

// IRQ represents Cortex-M external interrupt.
type IRQ byte

// Enabled returns true if handler for irq is enabled.
func (irq IRQ) Enabled() bool {
	return r.ISER.Bit(irq)
}

// Enable enables handler for irq.
func (irq IRQ) Enable() {
	r.ISER.SetBit(irq)
}

// Disable disables handler for irq. To disable some class of exceptions in
// atomic way see PRIMASK, FAULTMASK, BASEPRI registers in cortexm package..
func (irq IRQ) Disable() {
	r.ICER.SetBit(irq)
}

// Pending returns true if irq is pending.
func (irq IRQ) Pending() bool {
	return r.ISPR.Bit(irq)
}

// SetPending generates irq.
func (irq IRQ) SetPending() {
	r.ISPR.SetBit(irq)
}

// ClearPending clears pending flag for irq.
func (irq IRQ) ClearPending() {
	r.ICPR.SetBit(irq)
}

// Active returns true if CPU is handling irq.
func (irq IRQ) Active() bool {
	return r.IABR.Bit(irq)
}

// Prio returns priority level for irq.
func (irq IRQ) Prio() int {
	return int(r.IPR.Byte(irq))
}

// SetPrio sets priority level for irq.
func (irq IRQ) SetPrio(p int) {
	r.IPR.SetByte(irq, byte(p))
}

// Trig generates irq. This function can be call from user level if
// USERSETMPEND bit is set in scb.CCR.
func (irq IRQ) Trig() {
	sti.Store(uint32(irq))
}
<commit_msg>cortexm/nvic: a bug in IRQ type definition<commit_after>package nvic

// IRQ represents Cortex-M external interrupt.
type IRQ byte // BUG: ARMv7-M supports up to 496 external interrupts

// Enabled returns true if handler for irq is enabled.
func (irq IRQ) Enabled() bool {
	return r.ISER.Bit(irq)
}

// Enable enables handler for irq.
func (irq IRQ) Enable() {
	r.ISER.SetBit(irq)
}

// Disable disables handler for irq. To disable some class of exceptions in
// atomic way see PRIMASK, FAULTMASK, BASEPRI registers in cortexm package..
func (irq IRQ) Disable() {
	r.ICER.SetBit(irq)
}

// Pending returns true if irq is pending.
func (irq IRQ) Pending() bool {
	return r.ISPR.Bit(irq)
}

// SetPending generates irq.
func (irq IRQ) SetPending() {
	r.ISPR.SetBit(irq)
}

// ClearPending clears pending flag for irq.
func (irq IRQ) ClearPending() {
	r.ICPR.SetBit(irq)
}

// Active returns true if CPU is handling irq.
func (irq IRQ) Active() bool {
	return r.IABR.Bit(irq)
}

// Prio returns priority level for irq.
func (irq IRQ) Prio() int {
	return int(r.IPR.Byte(irq))
}

// SetPrio sets priority level for irq.
func (irq IRQ) SetPrio(p int) {
	r.IPR.SetByte(irq, byte(p))
}

// Trig generates irq. This function can be call from user level if
// USERSETMPEND bit is set in scb.CCR.
func (irq IRQ) Trig() {
	sti.Store(uint32(irq))
}
<|endoftext|>


<commit_before>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"github.com/globocom/tsuru/provision"
	"io"
	"strconv"
	"time"
)

func init() {
	provision.Register("fake", &FakeProvisioner{})
}

// Fake implementation for provision.Unit.
type FakeUnit struct {
	name    string
	machine int
	actions []string
}

func (u *FakeUnit) GetName() string {
	u.actions = append(u.actions, "getname")
	return u.name
}

func (u *FakeUnit) GetMachine() int {
	u.actions = append(u.actions, "getmachine")
	return u.machine
}

// Fake implementation for provision.App.
type FakeApp struct {
	name      string
	framework string
	units     []provision.AppUnit
	logs      []string
	actions   []string
}

func NewFakeApp(name, framework string, units int) *FakeApp {
	app := FakeApp{
		name:      name,
		framework: framework,
		units:     make([]provision.AppUnit, units),
	}
	for i := 0; i < units; i++ {
		app.units[i] = &FakeUnit{name: name, machine: i + 1}
	}
	return &app
}

func (a *FakeApp) Log(message, source string) error {
	a.logs = append(a.logs, source+message)
	a.actions = append(a.actions, "log "+source+" - "+message)
	return nil
}

func (a *FakeApp) GetName() string {
	a.actions = append(a.actions, "getname")
	return a.name
}

func (a *FakeApp) GetFramework() string {
	a.actions = append(a.actions, "getframework")
	return a.framework
}

func (a *FakeApp) ProvisionUnits() []provision.AppUnit {
	a.actions = append(a.actions, "getunits")
	return a.units
}

type Cmd struct {
	Cmd  string
	Args []string
	App  provision.App
}

type failure struct {
	method string
	err    error
}

// Fake implementation for provision.Provisioner.
type FakeProvisioner struct {
	apps     []provision.App
	Cmds     []Cmd
	outputs  chan []byte
	failures chan failure
}

func NewFakeProvisioner() *FakeProvisioner {
	p := FakeProvisioner{}
	p.outputs = make(chan []byte, 8)
	p.failures = make(chan failure, 8)
	return &p
}

func (p *FakeProvisioner) getError(method string) error {
	select {
	case fail := <-p.failures:
		if fail.method == method {
			return fail.err
		}
		p.failures <- fail
	case <-time.After(1e6):
	}
	return nil
}

func (p *FakeProvisioner) GetCmds(cmd string, app provision.App) []Cmd {
	var cmds []Cmd
	for _, c := range p.Cmds {
		if c.Cmd == cmd && app.GetName() == c.App.GetName() {
			cmds = append(cmds, c)
		}
	}
	return cmds
}

func (p *FakeProvisioner) FindApp(app provision.App) int {
	for i, a := range p.apps {
		if a.GetName() == app.GetName() {
			return i
		}
	}
	return -1
}

func (p *FakeProvisioner) PrepareOutput(b []byte) {
	p.outputs <- b
}

func (p *FakeProvisioner) PrepareFailure(method string, err error) {
	p.failures <- failure{method, err}
}

func (p *FakeProvisioner) Reset() {
	close(p.outputs)
	close(p.failures)
	p.outputs = make(chan []byte, 8)
	p.failures = make(chan failure, 8)
}

func (p *FakeProvisioner) Provision(app provision.App) error {
	if err := p.getError("Provision"); err != nil {
		return err
	}
	index := p.FindApp(app)
	if index > -1 {
		return &provision.Error{Reason: "App already provisioned."}
	}
	p.apps = append(p.apps, app)
	return nil
}

func (p *FakeProvisioner) Destroy(app provision.App) error {
	if err := p.getError("Destroy"); err != nil {
		return err
	}
	index := p.FindApp(app)
	if index == -1 {
		return &provision.Error{Reason: "App is not provisioned."}
	}
	copy(p.apps[index:], p.apps[index+1:])
	p.apps = p.apps[:len(p.apps)-1]
	return nil
}

func (p *FakeProvisioner) ExecuteCommand(stdout, stderr io.Writer, app provision.App, cmd string, args ...string) error {
	var (
		output []byte
		err    error
	)
	command := Cmd{
		Cmd:  cmd,
		Args: args,
		App:  app,
	}
	p.Cmds = append(p.Cmds, command)
	select {
	case output = <-p.outputs:
		stdout.Write(output)
	case fail := <-p.failures:
		select {
		case output = <-p.outputs:
			stderr.Write(output)
		case <-time.After(1e6):
		}
		if fail.method == "ExecuteCommand" {
			err = fail.err
		} else {
			p.failures <- fail
		}
	}
	return err
}

func (p *FakeProvisioner) CollectStatus() ([]provision.Unit, error) {
	if err := p.getError("CollectStatus"); err != nil {
		return nil, err
	}
	units := make([]provision.Unit, len(p.apps))
	for i, app := range p.apps {
		unit := provision.Unit{
			Name:    "somename",
			AppName: app.GetName(),
			Type:    app.GetFramework(),
			Status:  "started",
			Ip:      "10.10.10." + strconv.Itoa(i+1),
			Machine: i + 1,
		}
		units = append(units, unit)
	}
	return units, nil
}
<commit_msg>testing: improve GetCmds method<commit_after>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"github.com/globocom/tsuru/provision"
	"io"
	"strconv"
	"time"
)

func init() {
	provision.Register("fake", &FakeProvisioner{})
}

// Fake implementation for provision.Unit.
type FakeUnit struct {
	name    string
	machine int
	actions []string
}

func (u *FakeUnit) GetName() string {
	u.actions = append(u.actions, "getname")
	return u.name
}

func (u *FakeUnit) GetMachine() int {
	u.actions = append(u.actions, "getmachine")
	return u.machine
}

// Fake implementation for provision.App.
type FakeApp struct {
	name      string
	framework string
	units     []provision.AppUnit
	logs      []string
	actions   []string
}

func NewFakeApp(name, framework string, units int) *FakeApp {
	app := FakeApp{
		name:      name,
		framework: framework,
		units:     make([]provision.AppUnit, units),
	}
	for i := 0; i < units; i++ {
		app.units[i] = &FakeUnit{name: name, machine: i + 1}
	}
	return &app
}

func (a *FakeApp) Log(message, source string) error {
	a.logs = append(a.logs, source+message)
	a.actions = append(a.actions, "log "+source+" - "+message)
	return nil
}

func (a *FakeApp) GetName() string {
	a.actions = append(a.actions, "getname")
	return a.name
}

func (a *FakeApp) GetFramework() string {
	a.actions = append(a.actions, "getframework")
	return a.framework
}

func (a *FakeApp) ProvisionUnits() []provision.AppUnit {
	a.actions = append(a.actions, "getunits")
	return a.units
}

type Cmd struct {
	Cmd  string
	Args []string
	App  provision.App
}

type failure struct {
	method string
	err    error
}

// Fake implementation for provision.Provisioner.
type FakeProvisioner struct {
	apps     []provision.App
	Cmds     []Cmd
	outputs  chan []byte
	failures chan failure
}

func NewFakeProvisioner() *FakeProvisioner {
	p := FakeProvisioner{}
	p.outputs = make(chan []byte, 8)
	p.failures = make(chan failure, 8)
	return &p
}

func (p *FakeProvisioner) getError(method string) error {
	select {
	case fail := <-p.failures:
		if fail.method == method {
			return fail.err
		}
		p.failures <- fail
	case <-time.After(1e6):
	}
	return nil
}

// GetCmds returns a list of commands executed in an app. If you don't specify
// the command (""), it will return all commands executed in the given app.
func (p *FakeProvisioner) GetCmds(cmd string, app provision.App) []Cmd {
	var cmds []Cmd
	for _, c := range p.Cmds {
		if (cmd == "" || c.Cmd == cmd) && app.GetName() == c.App.GetName() {
			cmds = append(cmds, c)
		}
	}
	return cmds
}

func (p *FakeProvisioner) FindApp(app provision.App) int {
	for i, a := range p.apps {
		if a.GetName() == app.GetName() {
			return i
		}
	}
	return -1
}

func (p *FakeProvisioner) PrepareOutput(b []byte) {
	p.outputs <- b
}

func (p *FakeProvisioner) PrepareFailure(method string, err error) {
	p.failures <- failure{method, err}
}

func (p *FakeProvisioner) Reset() {
	close(p.outputs)
	close(p.failures)
	p.outputs = make(chan []byte, 8)
	p.failures = make(chan failure, 8)
	p.Cmds = nil
}

func (p *FakeProvisioner) Provision(app provision.App) error {
	if err := p.getError("Provision"); err != nil {
		return err
	}
	index := p.FindApp(app)
	if index > -1 {
		return &provision.Error{Reason: "App already provisioned."}
	}
	p.apps = append(p.apps, app)
	return nil
}

func (p *FakeProvisioner) Destroy(app provision.App) error {
	if err := p.getError("Destroy"); err != nil {
		return err
	}
	index := p.FindApp(app)
	if index == -1 {
		return &provision.Error{Reason: "App is not provisioned."}
	}
	copy(p.apps[index:], p.apps[index+1:])
	p.apps = p.apps[:len(p.apps)-1]
	return nil
}

func (p *FakeProvisioner) ExecuteCommand(stdout, stderr io.Writer, app provision.App, cmd string, args ...string) error {
	var (
		output []byte
		err    error
	)
	command := Cmd{
		Cmd:  cmd,
		Args: args,
		App:  app,
	}
	p.Cmds = append(p.Cmds, command)
	select {
	case output = <-p.outputs:
		stdout.Write(output)
	case fail := <-p.failures:
		select {
		case output = <-p.outputs:
			stderr.Write(output)
		case <-time.After(1e6):
		}
		if fail.method == "ExecuteCommand" {
			err = fail.err
		} else {
			p.failures <- fail
		}
	}
	return err
}

func (p *FakeProvisioner) CollectStatus() ([]provision.Unit, error) {
	if err := p.getError("CollectStatus"); err != nil {
		return nil, err
	}
	units := make([]provision.Unit, len(p.apps))
	for i, app := range p.apps {
		unit := provision.Unit{
			Name:    "somename",
			AppName: app.GetName(),
			Type:    app.GetFramework(),
			Status:  "started",
			Ip:      "10.10.10." + strconv.Itoa(i+1),
			Machine: i + 1,
		}
		units = append(units, unit)
	}
	return units, nil
}
<|endoftext|>


<commit_before>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"github.com/globocom/tsuru/provision"
	"io"
	"strconv"
)

func init() {
	provision.Register("fake", &FakeProvisioner{})
}

// Fake implementation for provision.Unit.
type FakeUnit struct {
	name    string
	machine int
	actions []string
}

func (u *FakeUnit) GetName() string {
	u.actions = append(u.actions, "getname")
	return u.name
}

func (u *FakeUnit) GetMachine() int {
	u.actions = append(u.actions, "getmachine")
	return u.machine
}

// Fake implementation for provision.App.
type FakeApp struct {
	name      string
	framework string
	units     []provision.AppUnit
	logs      []string
	actions   []string
}

func NewFakeApp(name, framework string, units int) *FakeApp {
	app := FakeApp{
		name:      name,
		framework: framework,
		units:     make([]provision.AppUnit, units),
	}
	for i := 0; i < units; i++ {
		app.units[i] = &FakeUnit{name: name, machine: i + 1}
	}
	return &app
}

func (a *FakeApp) Log(message, source string) error {
	a.logs = append(a.logs, source+message)
	a.actions = append(a.actions, "log "+source+" - "+message)
	return nil
}

func (a *FakeApp) GetName() string {
	a.actions = append(a.actions, "getname")
	return a.name
}

func (a *FakeApp) GetFramework() string {
	a.actions = append(a.actions, "getframework")
	return a.framework
}

func (a *FakeApp) GetUnits() []provision.AppUnit {
	a.actions = append(a.actions, "getunits")
	return a.units
}

type Cmd struct {
	Cmd  string
	Args []string
	App  provision.App
}

// Fake implementation for provision.Provisioner.
type FakeProvisioner struct {
	apps []provision.App
	Cmds []Cmd
}

func (p *FakeProvisioner) FindApp(app provision.App) int {
	for i, a := range p.apps {
		if a.GetName() == app.GetName() {
			return i
		}
	}
	return -1
}

func (p *FakeProvisioner) Provision(app provision.App) *provision.Error {
	index := p.FindApp(app)
	if index > -1 {
		return &provision.Error{Reason: "App already provisioned."}
	}
	p.apps = append(p.apps, app)
	return nil
}

func (p *FakeProvisioner) Destroy(app provision.App) *provision.Error {
	index := p.FindApp(app)
	if index == -1 {
		return &provision.Error{Reason: "App is not provisioned."}
	}
	copy(p.apps[index:], p.apps[index+1:])
	p.apps = p.apps[:len(p.apps)-1]
	return nil
}

func (p *FakeProvisioner) ExecuteCommand(w io.Writer, app provision.App, cmd string, args ...string) error {
	command := Cmd{
		Cmd:  cmd,
		Args: args,
		App:  app,
	}
	p.Cmds = append(p.Cmds, command)
	return nil
}

func (p *FakeProvisioner) CollectStatus() ([]provision.Unit, *provision.Error) {
	units := make([]provision.Unit, len(p.apps))
	for i, app := range p.apps {
		unit := provision.Unit{
			Name:    "somename",
			AppName: app.GetName(),
			Type:    app.GetFramework(),
			Status:  "started",
			Ip:      "10.10.10." + strconv.Itoa(i+1),
			Machine: i + 1,
		}
		units = append(units, unit)
	}
	return units, nil
}
<commit_msg>testing: update FakeApp to match new provision.App interface<commit_after>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"github.com/globocom/tsuru/provision"
	"io"
	"strconv"
	"time"
)

func init() {
	provision.Register("fake", &FakeProvisioner{})
}

// Fake implementation for provision.Unit.
type FakeUnit struct {
	name    string
	machine int
	actions []string
}

func (u *FakeUnit) GetName() string {
	u.actions = append(u.actions, "getname")
	return u.name
}

func (u *FakeUnit) GetMachine() int {
	u.actions = append(u.actions, "getmachine")
	return u.machine
}

// Fake implementation for provision.App.
type FakeApp struct {
	name      string
	framework string
	units     []provision.AppUnit
	logs      []string
	actions   []string
}

func NewFakeApp(name, framework string, units int) *FakeApp {
	app := FakeApp{
		name:      name,
		framework: framework,
		units:     make([]provision.AppUnit, units),
	}
	for i := 0; i < units; i++ {
		app.units[i] = &FakeUnit{name: name, machine: i + 1}
	}
	return &app
}

func (a *FakeApp) Log(message, source string) error {
	a.logs = append(a.logs, source+message)
	a.actions = append(a.actions, "log "+source+" - "+message)
	return nil
}

func (a *FakeApp) GetName() string {
	a.actions = append(a.actions, "getname")
	return a.name
}

func (a *FakeApp) GetFramework() string {
	a.actions = append(a.actions, "getframework")
	return a.framework
}

func (a *FakeApp) GetProvisionUnits() []provision.AppUnit {
	a.actions = append(a.actions, "getunits")
	return a.units
}

type Cmd struct {
	Cmd  string
	Args []string
	App  provision.App
}

// Fake implementation for provision.Provisioner.
type FakeProvisioner struct {
	apps     []provision.App
	Cmds     []Cmd
	outputs  chan []byte
	failures chan error
}

func NewFakeProvisioner() *FakeProvisioner {
	p := FakeProvisioner{}
	p.outputs = make(chan []byte, 8)
	p.failures = make(chan error, 8)
	return &p
}

func (p *FakeProvisioner) FindApp(app provision.App) int {
	for i, a := range p.apps {
		if a.GetName() == app.GetName() {
			return i
		}
	}
	return -1
}

func (p *FakeProvisioner) PrepareOutput(b []byte) {
	p.outputs <- b
}

func (p *FakeProvisioner) PrepareCommandFailure(err error) {
	p.failures <- err
}

func (p *FakeProvisioner) Reset() {
	close(p.outputs)
	close(p.failures)
	p.outputs = make(chan []byte, 8)
	p.failures = make(chan error, 8)
}

func (p *FakeProvisioner) Provision(app provision.App) *provision.Error {
	index := p.FindApp(app)
	if index > -1 {
		return &provision.Error{Reason: "App already provisioned."}
	}
	p.apps = append(p.apps, app)
	return nil
}

func (p *FakeProvisioner) Destroy(app provision.App) *provision.Error {
	index := p.FindApp(app)
	if index == -1 {
		return &provision.Error{Reason: "App is not provisioned."}
	}
	copy(p.apps[index:], p.apps[index+1:])
	p.apps = p.apps[:len(p.apps)-1]
	return nil
}

func (p *FakeProvisioner) ExecuteCommand(w io.Writer, app provision.App, cmd string, args ...string) error {
	command := Cmd{
		Cmd:  cmd,
		Args: args,
		App:  app,
	}
	p.Cmds = append(p.Cmds, command)
	var output []byte
	select {
	case output = <-p.outputs:
		w.Write(output)
	case <-time.After(1e9):
	}
	return nil
}

func (p *FakeProvisioner) CollectStatus() ([]provision.Unit, *provision.Error) {
	units := make([]provision.Unit, len(p.apps))
	for i, app := range p.apps {
		unit := provision.Unit{
			Name:    "somename",
			AppName: app.GetName(),
			Type:    app.GetFramework(),
			Status:  "started",
			Ip:      "10.10.10." + strconv.Itoa(i+1),
			Machine: i + 1,
		}
		units = append(units, unit)
	}
	return units, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 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 cas provides remote-apis-sdks client with luci integration.
package cas

import (
	"context"
	"strings"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/hardcoded/chromeinfra"
)

// NewClient returns luci auth configured Client for RBE-CAS.
func NewClient(ctx context.Context, instance string, opts auth.Options, readOnly bool) (*client.Client, error) {
	project := strings.Split(instance, "/")[1]
	var role string
	if readOnly {
		role = "cas-read-only"
	} else {
		role = "cas-read-write"
	}

	// Construct auth.Options.
	opts.ActAsServiceAccount = role + "@" + project + ".iam.gserviceaccount.com"
	opts.ActViaLUCIRealm = "@internal:" + project + "/" + role
	opts.Scopes = []string{"https://www.googleapis.com/auth/cloud-platform"}

	if strings.HasSuffix(project, "-dev") || strings.HasSuffix(project, "-staging") {
		// use dev token server for dev/staging projects.
		opts.TokenServerHost = chromeinfra.TokenServerDevHost
	}

	a := auth.NewAuthenticator(ctx, auth.SilentLogin, opts)
	creds, err := a.PerRPCCredentials()
	if err != nil {
		return nil, errors.Annotate(err, "failed to get PerRPCCredentials").Err()
	}

	client, err := client.NewClient(ctx, instance,
		client.DialParams{
			Service:            "remotebuildexecution.googleapis.com:443",
			TransportCredsOnly: true,
		}, &client.PerRPCCreds{Creds: creds})
	if err != nil {
		return nil, errors.Annotate(err, "failed to create client").Err()
	}

	// Set restricted permission for written files.
	client.DirMode = 0700
	client.ExecutableMode = 0700
	client.RegularMode = 0600
	client.UtilizeLocality = true
	return client, nil
}
<commit_msg>[cas] set lower cas concurrecy than default<commit_after>// Copyright 2020 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 cas provides remote-apis-sdks client with luci integration.
package cas

import (
	"context"
	"runtime"
	"strings"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/hardcoded/chromeinfra"
)

// NewClient returns luci auth configured Client for RBE-CAS.
func NewClient(ctx context.Context, instance string, opts auth.Options, readOnly bool) (*client.Client, error) {
	project := strings.Split(instance, "/")[1]
	var role string
	if readOnly {
		role = "cas-read-only"
	} else {
		role = "cas-read-write"
	}

	// Construct auth.Options.
	opts.ActAsServiceAccount = role + "@" + project + ".iam.gserviceaccount.com"
	opts.ActViaLUCIRealm = "@internal:" + project + "/" + role
	opts.Scopes = []string{"https://www.googleapis.com/auth/cloud-platform"}

	if strings.HasSuffix(project, "-dev") || strings.HasSuffix(project, "-staging") {
		// use dev token server for dev/staging projects.
		opts.TokenServerHost = chromeinfra.TokenServerDevHost
	}

	a := auth.NewAuthenticator(ctx, auth.SilentLogin, opts)
	creds, err := a.PerRPCCredentials()
	if err != nil {
		return nil, errors.Annotate(err, "failed to get PerRPCCredentials").Err()
	}

	client, err := client.NewClient(ctx, instance,
		client.DialParams{
			Service:            "remotebuildexecution.googleapis.com:443",
			TransportCredsOnly: true,
		}, &client.PerRPCCreds{Creds: creds},
		// This is for better file write performance on Windows (http://b/171672371#comment6).
		client.CASConcurrency(runtime.NumCPU()))
	if err != nil {
		return nil, errors.Annotate(err, "failed to create client").Err()
	}

	// Set restricted permission for written files.
	client.DirMode = 0700
	client.ExecutableMode = 0700
	client.RegularMode = 0600
	client.UtilizeLocality = true
	return client, nil
}
<|endoftext|>


<commit_before>package sqsadaptor

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
)

func init() {
}

func TestAdaptorConstruction(t *testing.T) {
	testsqs := NewSQSAdaptor("testqueue")
	assert.Equal(t, testsqs.QueueURL, "testqueue")
}

func TestJSON(t *testing.T) {
	result := Result{
		"2016-01-01 10:00:00",
		"example.com",
		"Fetch",
		"928429348",
		200,
		238947,
		2398,
		"Finished",
	}
	str, jsonerr := jsonFromResult(result)
	if jsonerr != nil {
		fmt.Println("JSON error")
		return
	}
	assert.Equal(t, str, "{\"time\":\"2016-01-01 10:00:00\",\"host\":\"example.com\",\"type\":\"Fetch\",\"requestID\":\"928429348\",\"status\":200,\"elapsed\":238947,\"bytes\":2398,\"state\":\"Finished\"}")
}
<commit_msg>Update sqsadaptor test to match implementation<commit_after>package sqsadaptor

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
)

func init() {
}

func TestAdaptorConstruction(t *testing.T) {
	testsqs := NewSQSAdaptor("testqueue")
	assert.Equal(t, testsqs.QueueURL, "testqueue")
}

func TestJSON(t *testing.T) {
	result := Result{
		"2016-01-01 10:00:00",
		"example.com",
		"Fetch",
		200,
		int64(12345), // elapsed, first byte
		int64(6789),  // elapsed, last byte
		int64(4567),  // elapsed, total
		4711,         //bytes
		"Success",
		"aws-lambda-4711", // AWS lamba instance

	}
	str, jsonerr := jsonFromResult(result)
	if jsonerr != nil {
		fmt.Println("JSON error")
		return
	}
	assert.Equal(t, str, "{\"time\":\"2016-01-01 10:00:00\",\"host\":\"example.com\",\"type\":\"Fetch\",\"status\":200,\"elapsed-first-byte\":12345,\"elapsed-last-byte\":6789,\"elapsed\":4567,\"bytes\":4711,\"state\":\"Success\",\"instance\":\"aws-lambda-4711\"}")
}
<|endoftext|>


<commit_before>package driver

import (
        "bytes"
        "encoding/json"
        "fmt"
        "log"
        "os"
        "os/exec"
        "regexp"
        "runtime"
        "strings"
        "syscall"
        "time"

        "github.com/hashicorp/nomad/client/config"
        "github.com/hashicorp/nomad/nomad/structs"
)

var (
        reRktVersion  = regexp.MustCompile("rkt version ([\\d\\.]+).+")
        reAppcVersion = regexp.MustCompile("appc version ([\\d\\.]+).+")
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
        DriverContext
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
        proc   *os.Process
        name   string
        logger *log.Logger
        waitCh chan error
        doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
        Pid  int
        Name string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
        return &RktDriver{*ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
        // Only enable if we are root when running on non-windows systems.
        if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
                d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
                return false, nil
        }

        outBytes, err := exec.Command("rkt", "version").Output()
        if err != nil {
                return false, nil
        }
        out := strings.TrimSpace(string(outBytes))

        rktMatches := reRktVersion.FindStringSubmatch(out)
        appcMatches := reRktVersion.FindStringSubmatch(out)
        if len(rktMatches) != 2 || len(appcMatches) != 2 {
                return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
        }

        node.Attributes["driver.rkt"] = "true"
        node.Attributes["driver.rkt.version"] = rktMatches[0]
        node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

        return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
        trust_prefix, ok := task.Config["trust_prefix"]
        if !ok || trust_prefix == "" {
                return nil, fmt.Errorf("Missing trust prefix for rkt")
        }

        // Add the given trust prefix
        var outBuf, errBuf bytes.Buffer
        cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trust_prefix))
        cmd.Stdout = &outBuf
        cmd.Stderr = &errBuf
        d.logger.Printf("[DEBUG] driver.rkt: starting rkt command: %q", cmd.Args)
        if err := cmd.Run(); err != nil {
                return nil, fmt.Errorf(
                        "Error running rkt: %s\n\nOutput: %s\n\nError: %s",
                        err, outBuf.String(), errBuf.String())
        }
        d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trust_prefix)

        name, ok := task.Config["name"]
        if !ok || name == "" {
                return nil, fmt.Errorf("Missing ACI name for rkt")
        }

        // Run the ACI
        var aoutBuf, aerrBuf bytes.Buffer
        acmd := exec.Command("rkt", "run", "--mds-register=false", name)
        acmd.Stdout = &aoutBuf
        acmd.Stderr = &aerrBuf
        d.logger.Printf("[DEBUG] driver:rkt: starting rkt command: %q", acmd.Args)
        if err := acmd.Run(); err != nil {
                return nil, fmt.Errorf(
                        "Error running rkt: %s\n\nOutput: %s\n\nError: %s",
                        err, aoutBuf.String(), aerrBuf.String())
        }
        d.logger.Printf("[DEBUG] driver.rkt: started ACI: %q", name)
        h := &rktHandle{
                proc:   acmd.Process,
                name:   name,
                logger: d.logger,
                doneCh: make(chan struct{}),
                waitCh: make(chan error, 1),
        }
        go h.run()
        return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
        // Parse the handle
        pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
        qpid := &rktPID{}
        if err := json.Unmarshal(pidBytes, qpid); err != nil {
                return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
        }

        // Find the process
        proc, err := os.FindProcess(qpid.Pid)
        if proc == nil || err != nil {
                return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
        }

        // Return a driver handle
        h := &rktHandle{
                proc:   proc,
                name:   qpid.Name,
                logger: d.logger,
                doneCh: make(chan struct{}),
                waitCh: make(chan error, 1),
        }

        go h.run()
        return h, nil
}

func (h *rktHandle) ID() string {
        // Return a handle to the PID
        pid := &rktPID{
                Pid:  h.proc.Pid,
                Name: h.name,
        }
        data, err := json.Marshal(pid)
        if err != nil {
                h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
        }
        return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan error {
        return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
        // Update is not possible
        return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
        h.proc.Signal(os.Interrupt)
        select {
        case <-h.doneCh:
                return nil
        case <-time.After(5 * time.Second):
                return h.proc.Kill()
        }
}

func (h *rktHandle) run() {
        ps, err := h.proc.Wait()
        close(h.doneCh)
        if err != nil {
                h.waitCh <- err
        } else if !ps.Success() {
                h.waitCh <- fmt.Errorf("task exited with error")
        }
        close(h.waitCh)
}
<commit_msg>Use Start to run commands<commit_after>package driver

import (
        "bytes"
        "encoding/json"
        "fmt"
        "log"
        "os"
        "os/exec"
        "regexp"
        "runtime"
        "strings"
        "syscall"
        "time"

        "github.com/hashicorp/nomad/client/config"
        "github.com/hashicorp/nomad/nomad/structs"
)

var (
        reRktVersion  = regexp.MustCompile("rkt version ([\\d\\.]+).+")
        reAppcVersion = regexp.MustCompile("appc version ([\\d\\.]+).+")
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
        DriverContext
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
        proc   *os.Process
        name   string
        logger *log.Logger
        waitCh chan error
        doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
        Pid  int
        Name string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
        return &RktDriver{*ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
        // Only enable if we are root when running on non-windows systems.
        if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
                d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
                return false, nil
        }

        outBytes, err := exec.Command("rkt", "version").Output()
        if err != nil {
                return false, nil
        }
        out := strings.TrimSpace(string(outBytes))

        rktMatches := reRktVersion.FindStringSubmatch(out)
        appcMatches := reRktVersion.FindStringSubmatch(out)
        if len(rktMatches) != 2 || len(appcMatches) != 2 {
                return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
        }

        node.Attributes["driver.rkt"] = "true"
        node.Attributes["driver.rkt.version"] = rktMatches[0]
        node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

        return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
        trust_prefix, ok := task.Config["trust_prefix"]
        if !ok || trust_prefix == "" {
                return nil, fmt.Errorf("Missing trust prefix for rkt")
        }

        // Add the given trust prefix
        var outBuf, errBuf bytes.Buffer
        cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trust_prefix))
        cmd.Stdout = &outBuf
        cmd.Stderr = &errBuf
        d.logger.Printf("[DEBUG] driver.rkt: starting rkt command: %q", cmd.Args)
        if err := cmd.Start(); err != nil {
                return nil, fmt.Errorf(
                        "Error running rkt: %s\n\nOutput: %s\n\nError: %s",
                        err, outBuf.String(), errBuf.String())
        }
        d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trust_prefix)

        name, ok := task.Config["name"]
        if !ok || name == "" {
                return nil, fmt.Errorf("Missing ACI name for rkt")
        }

        // Run the ACI
        var aoutBuf, aerrBuf bytes.Buffer
        acmd := exec.Command("rkt", "run", "--mds-register=false", name)
        acmd.Stdout = &aoutBuf
        acmd.Stderr = &aerrBuf
        d.logger.Printf("[DEBUG] driver:rkt: starting rkt command: %q", acmd.Args)
        if err := acmd.Start(); err != nil {
                return nil, fmt.Errorf(
                        "Error running rkt: %s\n\nOutput: %s\n\nError: %s",
                        err, aoutBuf.String(), aerrBuf.String())
        }
        d.logger.Printf("[DEBUG] driver.rkt: started ACI: %q", name)
        h := &rktHandle{
                proc:   acmd.Process,
                name:   name,
                logger: d.logger,
                doneCh: make(chan struct{}),
                waitCh: make(chan error, 1),
        }
        go h.run()
        return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
        // Parse the handle
        pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
        qpid := &rktPID{}
        if err := json.Unmarshal(pidBytes, qpid); err != nil {
                return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
        }

        // Find the process
        proc, err := os.FindProcess(qpid.Pid)
        if proc == nil || err != nil {
                return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
        }

        // Return a driver handle
        h := &rktHandle{
                proc:   proc,
                name:   qpid.Name,
                logger: d.logger,
                doneCh: make(chan struct{}),
                waitCh: make(chan error, 1),
        }

        go h.run()
        return h, nil
}

func (h *rktHandle) ID() string {
        // Return a handle to the PID
        pid := &rktPID{
                Pid:  h.proc.Pid,
                Name: h.name,
        }
        data, err := json.Marshal(pid)
        if err != nil {
                h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
        }
        return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan error {
        return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
        // Update is not possible
        return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
        h.proc.Signal(os.Interrupt)
        select {
        case <-h.doneCh:
                return nil
        case <-time.After(5 * time.Second):
                return h.proc.Kill()
        }
}

func (h *rktHandle) run() {
        ps, err := h.proc.Wait()
        close(h.doneCh)
        if err != nil {
                h.waitCh <- err
        } else if !ps.Success() {
                h.waitCh <- fmt.Errorf("task exited with error")
        }
        close(h.waitCh)
}
<|endoftext|>


<commit_before>package driver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/go-plugin"
	"github.com/hashicorp/go-version"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/executor"
	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/helper/discover"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/mitchellh/mapstructure"
)

var (
	reRktVersion  = regexp.MustCompile(`rkt [vV]ersion[:]? (\d[.\d]+)`)
	reAppcVersion = regexp.MustCompile(`appc [vV]ersion[:]? (\d[.\d]+)`)
)

const (
	// minRktVersion is the earliest supported version of rkt. rkt added support
	// for CPU and memory isolators in 0.14.0. We cannot support an earlier
	// version to maintain an uniform interface across all drivers
	minRktVersion = "0.14.0"

	// bytesToMB is the conversion from bytes to megabytes.
	bytesToMB = 1024 * 1024
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
	DriverContext
	fingerprint.StaticFingerprinter
}

type RktDriverConfig struct {
	ImageName        string   `mapstructure:"image"`
	Args             []string `mapstructure:"args"`
	DNSServers       []string `mapstructure:"dns_servers"`        // DNS Server for containers
	DNSSearchDomains []string `mapstructure:"dns_search_domains"` // DNS Search domains for containers
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
	pluginClient   *plugin.Client
	executorPid    int
	executor       executor.Executor
	allocDir       *allocdir.AllocDir
	logger         *log.Logger
	killTimeout    time.Duration
	maxKillTimeout time.Duration
	waitCh         chan *cstructs.WaitResult
	doneCh         chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
	PluginConfig   *PluginReattachConfig
	AllocDir       *allocdir.AllocDir
	ExecutorPid    int
	KillTimeout    time.Duration
	MaxKillTimeout time.Duration
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
	return &RktDriver{DriverContext: *ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	// Only enable if we are root when running on non-windows systems.
	if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
		d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
		return false, nil
	}

	outBytes, err := exec.Command("rkt", "version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	rktMatches := reRktVersion.FindStringSubmatch(out)
	appcMatches := reAppcVersion.FindStringSubmatch(out)
	if len(rktMatches) != 2 || len(appcMatches) != 2 {
		return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
	}

	node.Attributes["driver.rkt"] = "1"
	node.Attributes["driver.rkt.version"] = rktMatches[1]
	node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

	minVersion, _ := version.NewVersion(minRktVersion)
	currentVersion, _ := version.NewVersion(node.Attributes["driver.rkt.version"])
	if currentVersion.LessThan(minVersion) {
		// Do not allow rkt < 0.14.0
		d.logger.Printf("[WARN] driver.rkt: please upgrade rkt to a version >= %s", minVersion)
		node.Attributes["driver.rkt"] = "0"
	}
	return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	var driverConfig RktDriverConfig
	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
		return nil, err
	}
	// Validate that the config is valid.
	img := driverConfig.ImageName
	if img == "" {
		return nil, fmt.Errorf("Missing ACI image for rkt")
	}

	// Get the tasks local directory.
	taskName := d.DriverContext.taskName
	taskDir, ok := ctx.AllocDir.TaskDirs[taskName]
	if !ok {
		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
	}

	// Build the command.
	var cmdArgs []string

	// Add the given trust prefix
	trustPrefix, trustCmd := task.Config["trust_prefix"]
	insecure := false
	if trustCmd {
		var outBuf, errBuf bytes.Buffer
		cmd := exec.Command("rkt", "trust", "--skip-fingerprint-review=true", fmt.Sprintf("--prefix=%s", trustPrefix))
		cmd.Stdout = &outBuf
		cmd.Stderr = &errBuf
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf("Error running rkt trust: %s\n\nOutput: %s\n\nError: %s",
				err, outBuf.String(), errBuf.String())
		}
		d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trustPrefix)
	} else {
		// Disble signature verification if the trust command was not run.
		insecure = true
	}

	cmdArgs = append(cmdArgs, "run", "--interactive")
	cmdArgs = append(cmdArgs, fmt.Sprintf("--volume=%s,kind=host,source=%s", task.Name, ctx.AllocDir.SharedDir))
	cmdArgs = append(cmdArgs, fmt.Sprintf("--mount=volume=%s,target=%s", task.Name, ctx.AllocDir.SharedDir))
	cmdArgs = append(cmdArgs, img)
	if insecure == true {
		cmdArgs = append(cmdArgs, "--insecure-options=all")
	}

	// Inject enviornment variables
	for k, v := range d.taskEnv.EnvMap() {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--set-env=%v=%v", k, v))
	}

	// Check if the user has overriden the exec command.
	if execCmd, ok := task.Config["command"]; ok {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--exec=%v", execCmd))
	}

	if task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Memory limit cannot be zero")
	}
	if task.Resources.CPU == 0 {
		return nil, fmt.Errorf("CPU limit cannot be zero")
	}

	// Add memory isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--memory=%vM", int64(task.Resources.MemoryMB)*bytesToMB))

	// Add CPU isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--cpu=%vm", int64(task.Resources.CPU)))

	// Add DNS servers
	for _, ip := range driverConfig.DNSServers {
		if err := net.ParseIP(ip); err == nil {
			msg := fmt.Errorf("invalid ip address for container dns server %q", ip)
			d.logger.Printf("[DEBUG] driver.rkt: %v", msg)
			return nil, msg
		} else {
			cmdArgs = append(cmdArgs, fmt.Sprintf("--dns=%s", ip))
		}
	}

	// set DNS search domains
	for _, domain := range driverConfig.DNSSearchDomains {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--dns-search=%s", domain))
	}

	// Add user passed arguments.
	if len(driverConfig.Args) != 0 {
		parsed := d.taskEnv.ParseAndReplace(driverConfig.Args)

		// Need to start arguments with "--"
		if len(parsed) > 0 {
			cmdArgs = append(cmdArgs, "--")
		}

		for _, arg := range parsed {
			cmdArgs = append(cmdArgs, fmt.Sprintf("%v", arg))
		}
	}

	bin, err := discover.NomadExecutable()
	if err != nil {
		return nil, fmt.Errorf("unable to find the nomad binary: %v", err)
	}

	pluginLogFile := filepath.Join(taskDir, fmt.Sprintf("%s-executor.out", task.Name))
	pluginConfig := &plugin.ClientConfig{
		Cmd: exec.Command(bin, "executor", pluginLogFile),
	}

	execIntf, pluginClient, err := createExecutor(pluginConfig, d.config.LogOutput, d.config)
	if err != nil {
		return nil, err
	}
	executorCtx := &executor.ExecutorContext{
		TaskEnv:          d.taskEnv,
		AllocDir:         ctx.AllocDir,
		TaskName:         task.Name,
		TaskResources:    task.Resources,
		UnprivilegedUser: false,
		LogConfig:        task.LogConfig,
	}

	absPath, err := GetAbsolutePath("rkt")
	if err != nil {
		return nil, err
	}

	ps, err := execIntf.LaunchCmd(&executor.ExecCommand{Cmd: absPath, Args: cmdArgs}, executorCtx)
	if err != nil {
		pluginClient.Kill()
		return nil, fmt.Errorf("error starting process via the plugin: %v", err)
	}

	d.logger.Printf("[DEBUG] driver.rkt: started ACI %q with: %v", img, cmdArgs)
	maxKill := d.DriverContext.config.MaxKillTimeout
	h := &rktHandle{
		pluginClient:   pluginClient,
		executor:       execIntf,
		executorPid:    ps.Pid,
		allocDir:       ctx.AllocDir,
		logger:         d.logger,
		killTimeout:    GetKillTimeout(task.KillTimeout, maxKill),
		maxKillTimeout: maxKill,
		doneCh:         make(chan struct{}),
		waitCh:         make(chan *cstructs.WaitResult, 1),
	}
	go h.run()
	return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
	id := &rktPID{}
	if err := json.Unmarshal(pidBytes, id); err != nil {
		return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
	}

	pluginConfig := &plugin.ClientConfig{
		Reattach: id.PluginConfig.PluginConfig(),
	}
	executor, pluginClient, err := createExecutor(pluginConfig, d.config.LogOutput, d.config)
	if err != nil {
		d.logger.Println("[ERROR] driver.rkt: error connecting to plugin so destroying plugin pid and user pid")
		if e := destroyPlugin(id.PluginConfig.Pid, id.ExecutorPid); e != nil {
			d.logger.Printf("[ERROR] driver.rkt: error destroying plugin and executor pid: %v", e)
		}
		return nil, fmt.Errorf("error connecting to plugin: %v", err)
	}

	// Return a driver handle
	h := &rktHandle{
		pluginClient:   pluginClient,
		executorPid:    id.ExecutorPid,
		allocDir:       id.AllocDir,
		executor:       executor,
		logger:         d.logger,
		killTimeout:    id.KillTimeout,
		maxKillTimeout: id.MaxKillTimeout,
		doneCh:         make(chan struct{}),
		waitCh:         make(chan *cstructs.WaitResult, 1),
	}

	go h.run()
	return h, nil
}

func (h *rktHandle) ID() string {
	// Return a handle to the PID
	pid := &rktPID{
		PluginConfig:   NewPluginReattachConfig(h.pluginClient.ReattachConfig()),
		KillTimeout:    h.killTimeout,
		MaxKillTimeout: h.maxKillTimeout,
		ExecutorPid:    h.executorPid,
		AllocDir:       h.allocDir,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
	}
	return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan *cstructs.WaitResult {
	return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
	// Store the updated kill timeout.
	h.killTimeout = GetKillTimeout(task.KillTimeout, h.maxKillTimeout)
	h.executor.UpdateLogConfig(task.LogConfig)

	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
	h.executor.ShutDown()
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(h.killTimeout):
		return h.executor.Exit()
	}
}

func (h *rktHandle) run() {
	ps, err := h.executor.Wait()
	close(h.doneCh)
	if ps.ExitCode == 0 && err != nil {
		if e := killProcess(h.executorPid); e != nil {
			h.logger.Printf("[ERROR] driver.rkt: error killing user process: %v", e)
		}
		if e := h.allocDir.UnmountAll(); e != nil {
			h.logger.Printf("[ERROR] driver.rkt: unmounting dev,proc and alloc dirs failed: %v", e)
		}
	}
	h.waitCh <- cstructs.NewWaitResult(ps.ExitCode, 0, err)
	close(h.waitCh)
	h.pluginClient.Kill()
}
<commit_msg>Remove unnecessary --interactive flag<commit_after>package driver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/go-plugin"
	"github.com/hashicorp/go-version"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/executor"
	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/helper/discover"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/mitchellh/mapstructure"
)

var (
	reRktVersion  = regexp.MustCompile(`rkt [vV]ersion[:]? (\d[.\d]+)`)
	reAppcVersion = regexp.MustCompile(`appc [vV]ersion[:]? (\d[.\d]+)`)
)

const (
	// minRktVersion is the earliest supported version of rkt. rkt added support
	// for CPU and memory isolators in 0.14.0. We cannot support an earlier
	// version to maintain an uniform interface across all drivers
	minRktVersion = "0.14.0"

	// bytesToMB is the conversion from bytes to megabytes.
	bytesToMB = 1024 * 1024
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
	DriverContext
	fingerprint.StaticFingerprinter
}

type RktDriverConfig struct {
	ImageName        string   `mapstructure:"image"`
	Args             []string `mapstructure:"args"`
	DNSServers       []string `mapstructure:"dns_servers"`        // DNS Server for containers
	DNSSearchDomains []string `mapstructure:"dns_search_domains"` // DNS Search domains for containers
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
	pluginClient   *plugin.Client
	executorPid    int
	executor       executor.Executor
	allocDir       *allocdir.AllocDir
	logger         *log.Logger
	killTimeout    time.Duration
	maxKillTimeout time.Duration
	waitCh         chan *cstructs.WaitResult
	doneCh         chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
	PluginConfig   *PluginReattachConfig
	AllocDir       *allocdir.AllocDir
	ExecutorPid    int
	KillTimeout    time.Duration
	MaxKillTimeout time.Duration
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
	return &RktDriver{DriverContext: *ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	// Only enable if we are root when running on non-windows systems.
	if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
		d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
		return false, nil
	}

	outBytes, err := exec.Command("rkt", "version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	rktMatches := reRktVersion.FindStringSubmatch(out)
	appcMatches := reAppcVersion.FindStringSubmatch(out)
	if len(rktMatches) != 2 || len(appcMatches) != 2 {
		return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
	}

	node.Attributes["driver.rkt"] = "1"
	node.Attributes["driver.rkt.version"] = rktMatches[1]
	node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

	minVersion, _ := version.NewVersion(minRktVersion)
	currentVersion, _ := version.NewVersion(node.Attributes["driver.rkt.version"])
	if currentVersion.LessThan(minVersion) {
		// Do not allow rkt < 0.14.0
		d.logger.Printf("[WARN] driver.rkt: please upgrade rkt to a version >= %s", minVersion)
		node.Attributes["driver.rkt"] = "0"
	}
	return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	var driverConfig RktDriverConfig
	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
		return nil, err
	}
	// Validate that the config is valid.
	img := driverConfig.ImageName
	if img == "" {
		return nil, fmt.Errorf("Missing ACI image for rkt")
	}

	// Get the tasks local directory.
	taskName := d.DriverContext.taskName
	taskDir, ok := ctx.AllocDir.TaskDirs[taskName]
	if !ok {
		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
	}

	// Build the command.
	var cmdArgs []string

	// Add the given trust prefix
	trustPrefix, trustCmd := task.Config["trust_prefix"]
	insecure := false
	if trustCmd {
		var outBuf, errBuf bytes.Buffer
		cmd := exec.Command("rkt", "trust", "--skip-fingerprint-review=true", fmt.Sprintf("--prefix=%s", trustPrefix))
		cmd.Stdout = &outBuf
		cmd.Stderr = &errBuf
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf("Error running rkt trust: %s\n\nOutput: %s\n\nError: %s",
				err, outBuf.String(), errBuf.String())
		}
		d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trustPrefix)
	} else {
		// Disble signature verification if the trust command was not run.
		insecure = true
	}

	cmdArgs = append(cmdArgs, "run")
	cmdArgs = append(cmdArgs, fmt.Sprintf("--volume=%s,kind=host,source=%s", task.Name, ctx.AllocDir.SharedDir))
	cmdArgs = append(cmdArgs, fmt.Sprintf("--mount=volume=%s,target=%s", task.Name, ctx.AllocDir.SharedDir))
	cmdArgs = append(cmdArgs, img)
	if insecure == true {
		cmdArgs = append(cmdArgs, "--insecure-options=all")
	}

	// Inject enviornment variables
	for k, v := range d.taskEnv.EnvMap() {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--set-env=%v=%v", k, v))
	}

	// Check if the user has overriden the exec command.
	if execCmd, ok := task.Config["command"]; ok {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--exec=%v", execCmd))
	}

	if task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Memory limit cannot be zero")
	}
	if task.Resources.CPU == 0 {
		return nil, fmt.Errorf("CPU limit cannot be zero")
	}

	// Add memory isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--memory=%vM", int64(task.Resources.MemoryMB)*bytesToMB))

	// Add CPU isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--cpu=%vm", int64(task.Resources.CPU)))

	// Add DNS servers
	for _, ip := range driverConfig.DNSServers {
		if err := net.ParseIP(ip); err == nil {
			msg := fmt.Errorf("invalid ip address for container dns server %q", ip)
			d.logger.Printf("[DEBUG] driver.rkt: %v", msg)
			return nil, msg
		} else {
			cmdArgs = append(cmdArgs, fmt.Sprintf("--dns=%s", ip))
		}
	}

	// set DNS search domains
	for _, domain := range driverConfig.DNSSearchDomains {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--dns-search=%s", domain))
	}

	// Add user passed arguments.
	if len(driverConfig.Args) != 0 {
		parsed := d.taskEnv.ParseAndReplace(driverConfig.Args)

		// Need to start arguments with "--"
		if len(parsed) > 0 {
			cmdArgs = append(cmdArgs, "--")
		}

		for _, arg := range parsed {
			cmdArgs = append(cmdArgs, fmt.Sprintf("%v", arg))
		}
	}

	bin, err := discover.NomadExecutable()
	if err != nil {
		return nil, fmt.Errorf("unable to find the nomad binary: %v", err)
	}

	pluginLogFile := filepath.Join(taskDir, fmt.Sprintf("%s-executor.out", task.Name))
	pluginConfig := &plugin.ClientConfig{
		Cmd: exec.Command(bin, "executor", pluginLogFile),
	}

	execIntf, pluginClient, err := createExecutor(pluginConfig, d.config.LogOutput, d.config)
	if err != nil {
		return nil, err
	}
	executorCtx := &executor.ExecutorContext{
		TaskEnv:          d.taskEnv,
		AllocDir:         ctx.AllocDir,
		TaskName:         task.Name,
		TaskResources:    task.Resources,
		UnprivilegedUser: false,
		LogConfig:        task.LogConfig,
	}

	absPath, err := GetAbsolutePath("rkt")
	if err != nil {
		return nil, err
	}

	ps, err := execIntf.LaunchCmd(&executor.ExecCommand{Cmd: absPath, Args: cmdArgs}, executorCtx)
	if err != nil {
		pluginClient.Kill()
		return nil, fmt.Errorf("error starting process via the plugin: %v", err)
	}

	d.logger.Printf("[DEBUG] driver.rkt: started ACI %q with: %v", img, cmdArgs)
	maxKill := d.DriverContext.config.MaxKillTimeout
	h := &rktHandle{
		pluginClient:   pluginClient,
		executor:       execIntf,
		executorPid:    ps.Pid,
		allocDir:       ctx.AllocDir,
		logger:         d.logger,
		killTimeout:    GetKillTimeout(task.KillTimeout, maxKill),
		maxKillTimeout: maxKill,
		doneCh:         make(chan struct{}),
		waitCh:         make(chan *cstructs.WaitResult, 1),
	}
	go h.run()
	return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
	id := &rktPID{}
	if err := json.Unmarshal(pidBytes, id); err != nil {
		return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
	}

	pluginConfig := &plugin.ClientConfig{
		Reattach: id.PluginConfig.PluginConfig(),
	}
	executor, pluginClient, err := createExecutor(pluginConfig, d.config.LogOutput, d.config)
	if err != nil {
		d.logger.Println("[ERROR] driver.rkt: error connecting to plugin so destroying plugin pid and user pid")
		if e := destroyPlugin(id.PluginConfig.Pid, id.ExecutorPid); e != nil {
			d.logger.Printf("[ERROR] driver.rkt: error destroying plugin and executor pid: %v", e)
		}
		return nil, fmt.Errorf("error connecting to plugin: %v", err)
	}

	// Return a driver handle
	h := &rktHandle{
		pluginClient:   pluginClient,
		executorPid:    id.ExecutorPid,
		allocDir:       id.AllocDir,
		executor:       executor,
		logger:         d.logger,
		killTimeout:    id.KillTimeout,
		maxKillTimeout: id.MaxKillTimeout,
		doneCh:         make(chan struct{}),
		waitCh:         make(chan *cstructs.WaitResult, 1),
	}

	go h.run()
	return h, nil
}

func (h *rktHandle) ID() string {
	// Return a handle to the PID
	pid := &rktPID{
		PluginConfig:   NewPluginReattachConfig(h.pluginClient.ReattachConfig()),
		KillTimeout:    h.killTimeout,
		MaxKillTimeout: h.maxKillTimeout,
		ExecutorPid:    h.executorPid,
		AllocDir:       h.allocDir,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
	}
	return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan *cstructs.WaitResult {
	return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
	// Store the updated kill timeout.
	h.killTimeout = GetKillTimeout(task.KillTimeout, h.maxKillTimeout)
	h.executor.UpdateLogConfig(task.LogConfig)

	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
	h.executor.ShutDown()
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(h.killTimeout):
		return h.executor.Exit()
	}
}

func (h *rktHandle) run() {
	ps, err := h.executor.Wait()
	close(h.doneCh)
	if ps.ExitCode == 0 && err != nil {
		if e := killProcess(h.executorPid); e != nil {
			h.logger.Printf("[ERROR] driver.rkt: error killing user process: %v", e)
		}
		if e := h.allocDir.UnmountAll(); e != nil {
			h.logger.Printf("[ERROR] driver.rkt: unmounting dev,proc and alloc dirs failed: %v", e)
		}
	}
	h.waitCh <- cstructs.NewWaitResult(ps.ExitCode, 0, err)
	close(h.waitCh)
	h.pluginClient.Kill()
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	goprompt "github.com/c-bata/go-prompt"
	"github.com/c-bata/go-prompt/completer"
	homedir "github.com/mitchellh/go-homedir"
	"github.com/nsf/termbox-go"
	"github.com/openconfig/goyang/pkg/yang"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

var promptMode bool
var promptHistory []string
var schemaTree *yang.Entry
var promptModeCmd = &cobra.Command{
	Use:   "prompt",
	Short: "enter the interactive gnmic prompt mode",
	RunE: func(cmd *cobra.Command, args []string) error {
		err := pathCmdRun(promptDirs, promptFiles, promptExcluded, true)
		if err == nil {
			promptMode = true
		}
		// load history
		promptHistory = make([]string, 0, 256)
		home, err := homedir.Dir()
		if err != nil {
			return err
		}
		content, err := ioutil.ReadFile(home + "/.gnmic.history")
		if err != nil {
			return err
		}
		history := strings.Split(string(content), "\n")
		for i := range history {
			if history[i] != "" {
				promptHistory = append(promptHistory, history[i])
			}
		}
		return err
	},
	PostRun: func(cmd *cobra.Command, args []string) {
		cmd.ResetFlags()
		initPromptFlags(cmd)
	},
	SilenceUsage: true,
}

var promptQuitCmd = &cobra.Command{
	Use:   "quit",
	Short: "quit the gnmic-prompt",
	Run: func(cmd *cobra.Command, args []string) {
		// save history
		home, err := homedir.Dir()
		if err != nil {
			os.Exit(0)
		}
		f, err := os.Create(home + "/.gnmic.history")
		if err != nil {
			os.Exit(0)
		}
		l := len(promptHistory)
		if l > 128 {
			promptHistory = promptHistory[l-128:]
		}
		for i := range promptHistory {
			f.WriteString(promptHistory[i] + "\n")
		}
		f.Close()
		os.Exit(0)
	},
}

func init() {
	rootCmd.AddCommand(promptModeCmd)
	initPromptFlags(promptModeCmd)
}

var promptFiles []string
var promptExcluded []string
var promptDirs []string

// used to init or reset pathCmd flags for gnmic-prompt mode
func initPromptFlags(cmd *cobra.Command) {
	cmd.Flags().StringArrayVarP(&promptFiles, "file", "", []string{}, "yang files to get the paths")
	cmd.Flags().StringArrayVarP(&promptExcluded, "exclude", "", []string{}, "yang modules to be excluded from path generation")
	cmd.Flags().StringArrayVarP(&promptDirs, "dir", "", []string{}, "directories to search yang includes and imports")
	viper.BindPFlag("prompt-file", cmd.LocalFlags().Lookup("file"))
	viper.BindPFlag("prompt-exclude", cmd.LocalFlags().Lookup("exclude"))
	viper.BindPFlag("prompt-dir", cmd.LocalFlags().Lookup("dir"))
}

func findMatchedXPATH(entry *yang.Entry, word string, cursor int) []goprompt.Suggest {
	suggestions := make([]goprompt.Suggest, 0, 4)
	cword := word[cursor:]
	for name, child := range entry.Dir {
		pathelem := "/" + name
		if strings.HasPrefix(pathelem, cword) {
			node := ""
			if os.PathSeparator != '/' {
				node = fmt.Sprintf("%s%s", word[:cursor], pathelem)
				suggestions = append(suggestions, goprompt.Suggest{Text: node, Description: child.Description})
			} else {
				if len(cword) >= 1 && cword[0] == '/' {
					node = name
				} else {
					node = pathelem
				}
				suggestions = append(suggestions, goprompt.Suggest{Text: node, Description: child.Description})
			}
			if child.Key != "" { // list
				keylist := strings.Split(child.Key, " ")
				for _, key := range keylist {
					node = fmt.Sprintf("%s[%s=*]", node, key)
				}
				suggestions = append(suggestions, goprompt.Suggest{Text: node, Description: child.Description})
			}
		} else if strings.HasPrefix(cword, pathelem) {
			var prevC rune
			var bracketCount int
			var endIndex int = -1
			var stop bool
			for i, c := range cword {
				switch c {
				case '[':
					bracketCount++
				case ']':
					if prevC != '\\' {
						bracketCount--
						endIndex = i
					}
				case '/':
					if i != 0 && bracketCount == 0 {
						endIndex = i
						stop = true
					}
				}
				if stop {
					break
				}
				prevC = c
			}
			if bracketCount == 0 {
				if endIndex >= 0 {
					suggestions = append(suggestions, findMatchedXPATH(child, word, cursor+endIndex)...)
				} else {
					suggestions = append(suggestions, findMatchedXPATH(child, word, cursor+len(pathelem))...)
				}
			}
		}
	}
	return suggestions
}

var filePathCompleter = completer.FilePathCompleter{
	IgnoreCase: true,
	Filter: func(fi os.FileInfo) bool {
		fmt.Println(fi.Name())
		return fi.IsDir() || !strings.HasPrefix(fi.Name(), ".")
	},
}

var yangPathCompleter = completer.FilePathCompleter{
	IgnoreCase: true,
	Filter: func(fi os.FileInfo) bool {
		fmt.Println(fi.Name())
		return fi.IsDir() || strings.HasSuffix(fi.Name(), ".yang")
	},
}

var dirPathCompleter = completer.FilePathCompleter{
	IgnoreCase: true,
	Filter: func(fi os.FileInfo) bool {
		return fi.IsDir()
	},
}

func findDynamicSuggestions(annotation string, doc goprompt.Document) []goprompt.Suggest {
	switch annotation {
	case "XPATH":
		word := doc.GetWordBeforeCursor()
		suggestions := make([]goprompt.Suggest, 0, 16)
		for _, entry := range schemaTree.Dir {
			suggestions = append(suggestions, findMatchedXPATH(entry, word, 0)...)
		}
		return suggestions
	case "FILE":
		return filePathCompleter.Complete(doc)
	case "YANG":
		return yangPathCompleter.Complete(doc)
	case "MODEL":
		suggestions := make([]goprompt.Suggest, 0, len(schemaTree.Dir))
		for name := range schemaTree.Dir {
			suggestions = append(suggestions, goprompt.Suggest{Text: name})
		}
		return goprompt.FilterHasPrefix(suggestions, doc.GetWordBeforeCursor(), true)
	case "DIR":
		return dirPathCompleter.Complete(doc)
	case "ENCODING":
		suggestions := make([]goprompt.Suggest, 0, len(encodings))
		for _, name := range encodings {
			suggestions = append(suggestions, goprompt.Suggest{Text: name})
		}
		return goprompt.FilterHasPrefix(suggestions, doc.GetWordBeforeCursor(), true)
	}
	return []goprompt.Suggest{}
}

func showCommandArguments(b *goprompt.Buffer) {
	doc := b.Document()
	showLocalFlags := false
	command := rootCmd
	args := strings.Fields(doc.CurrentLine())
	if found, _, err := command.Find(args); err == nil {
		if command != found {
			showLocalFlags = true
		}
		command = found
	}
	maxNameLen := 0
	suggestions := make([]goprompt.Suggest, 0, 32)
	if command.HasAvailableSubCommands() {
		for _, c := range command.Commands() {
			if c.Hidden {
				continue
			}
			length := len(c.Name())
			if maxNameLen < length {
				maxNameLen = length
			}
			suggestions = append(suggestions, goprompt.Suggest{Text: c.Name(), Description: c.Short})
		}
	}
	if showLocalFlags {
		addFlags := func(flag *pflag.Flag) {
			if flag.Hidden {
				return
			}
			length := len(flag.Name)
			if maxNameLen < length+2 {
				maxNameLen = length + 2
			}
			suggestions = append(suggestions, goprompt.Suggest{Text: "--" + flag.Name, Description: flag.Usage})
		}
		command.LocalFlags().VisitAll(addFlags)
	}
	suggestions = goprompt.FilterHasPrefix(suggestions, doc.GetWordBeforeCursor(), true)
	if len(suggestions) == 0 {
		return
	}
	if err := termbox.Init(); err != nil {
		logger.Fatalf("%v", err)
	}
	w, _ := termbox.Size()
	termbox.Close()
	fmt.Printf("\n")
	maxDescLen := w - maxNameLen - 6
	format := fmt.Sprintf("  %%-%ds : %%-%ds\n", maxNameLen, maxDescLen)
	for i := range suggestions {
		length := len(suggestions[i].Description)
		if length > maxDescLen {
			fmt.Printf(format, suggestions[i].Text, suggestions[i].Description[:maxDescLen])
		} else {
			fmt.Printf(format, suggestions[i].Text, suggestions[i].Description)
		}
	}
	fmt.Printf("\n")
}

// ExecutePrompt load and run gnmic-prompt mode.
func ExecutePrompt() {
	var err error
	schemaTree, err = loadSchemaZip()
	if err != nil {
		schemaTree = buildRootEntry()
	}
	rootCmd.AddCommand(promptQuitCmd)
	rootCmd.RemoveCommand(promptModeCmd)
	shell := &cmdPrompt{
		RootCmd: rootCmd,
		GoPromptOptions: []goprompt.Option{
			goprompt.OptionTitle("gnmic-prompt"),
			goprompt.OptionPrefix("gnmic> "),
			goprompt.OptionHistory(promptHistory),
			goprompt.OptionMaxSuggestion(5),
			goprompt.OptionPrefixTextColor(goprompt.Yellow),
			// goprompt.OptionPreviewSuggestionTextColor(goprompt.Yellow),
			goprompt.OptionPreviewSuggestionBGColor(goprompt.Black),
			goprompt.OptionSuggestionTextColor(goprompt.White),
			goprompt.OptionSuggestionBGColor(goprompt.Black),
			goprompt.OptionSelectedSuggestionTextColor(goprompt.Black),
			goprompt.OptionSelectedSuggestionBGColor(goprompt.White),
			// goprompt.OptionDescriptionTextColor(goprompt.White),
			goprompt.OptionDescriptionBGColor(goprompt.Yellow),
			goprompt.OptionSelectedDescriptionTextColor(goprompt.Black),
			goprompt.OptionSelectedDescriptionBGColor(goprompt.White),
			goprompt.OptionScrollbarBGColor(goprompt.White),
			goprompt.OptionAddASCIICodeBind(goprompt.ASCIICodeBind{
				ASCIICode: []byte{0x3f}, Fn: showCommandArguments}),
			goprompt.OptionCompletionWordSeparator(completer.FilePathCompletionSeparator),
		},
	}
	shell.Run()
}

// Reference: https://github.com/stromland/cobra-prompt
// cmdPrompt requires RootCmd to run
type cmdPrompt struct {
	// RootCmd is the start point, all its sub commands and flags will be available as suggestions
	RootCmd *cobra.Command

	// GoPromptOptions is for customize go-prompt
	// see https://github.com/c-bata/go-prompt/blob/master/option.go
	GoPromptOptions []goprompt.Option
}

// Run will automatically generate suggestions for all cobra commands
// and flags defined by RootCmd and execute the selected commands.
func (co cmdPrompt) Run() {
	p := goprompt.New(
		func(in string) {
			promptArgs := strings.Fields(in)
			os.Args = append([]string{os.Args[0]}, promptArgs...)
			if len(promptArgs) > 0 {
				err := co.RootCmd.Execute()
				if err == nil && in != "" {
					promptHistory = append(promptHistory, in)
				}
			}
		},
		func(d goprompt.Document) []goprompt.Suggest {
			return findSuggestions(co, d)
		},
		co.GoPromptOptions...,
	)
	p.Run()
}

func findSuggestions(co cmdPrompt, doc goprompt.Document) []goprompt.Suggest {
	showLocalFlags := false
	command := co.RootCmd
	args := strings.Fields(doc.CurrentLine())
	if found, _, err := command.Find(args); err == nil {
		if command != found {
			showLocalFlags = true
		}
		command = found
	}

	suggestions := make([]goprompt.Suggest, 0, 32)
	if command.HasAvailableSubCommands() {
		for _, c := range command.Commands() {
			if !c.Hidden {
				suggestions = append(suggestions, goprompt.Suggest{Text: c.Name(), Description: c.Short})
			}
		}
	}

	// check flag annotation for the dynamic suggestion
	annotation := ""
	argnum := len(args)
	wordBefore := doc.GetWordBeforeCursor()
	if wordBefore == "" {
		if argnum >= 1 {
			annotation = command.Annotations[args[argnum-1]]
		}
	} else {
		if argnum >= 2 {
			annotation = command.Annotations[args[argnum-2]]
		}
	}
	if annotation != "" {
		return append(suggestions, findDynamicSuggestions(annotation, doc)...)
	}

	if showLocalFlags {
		// load local flags of the command
		addFlags := func(flag *pflag.Flag) {
			if flag.Hidden {
				return
			}
			suggestions = append(suggestions, goprompt.Suggest{Text: "--" + flag.Name, Description: flag.Usage})
		}
		command.LocalFlags().VisitAll(addFlags)
		// command.InheritedFlags().VisitAll(addFlags)
	} else {

		// persistent flags are shown if run.
		addFlags := func(flag *pflag.Flag) {
			if flag.Hidden {
				return
			}
			// if strings.HasPrefix(doc.GetWordBeforeCursor(), "--") {
			// 	suggestions = append(suggestions, goprompt.Suggest{Text: "--" + flag.Name, Description: flag.Usage})
			// } else if strings.HasPrefix(doc.GetWordBeforeCursor(), "-") && flag.Shorthand != "" {
			// 	suggestions = append(suggestions, goprompt.Suggest{Text: "-" + flag.Shorthand, Description: flag.Usage})
			// }
			if strings.HasPrefix(doc.GetWordBeforeCursor(), "-") {
				suggestions = append(suggestions, goprompt.Suggest{Text: "--" + flag.Name, Description: flag.Usage})
			}
		}
		command.LocalFlags().VisitAll(addFlags)
	}
	return goprompt.FilterHasPrefix(suggestions, doc.GetWordBeforeCursor(), true)
}
<commit_msg>add suggestions from flags with annotations before subCommands<commit_after>package cmd

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	goprompt "github.com/c-bata/go-prompt"
	"github.com/c-bata/go-prompt/completer"
	homedir "github.com/mitchellh/go-homedir"
	"github.com/nsf/termbox-go"
	"github.com/openconfig/goyang/pkg/yang"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

var promptMode bool
var promptHistory []string
var schemaTree *yang.Entry
var promptModeCmd = &cobra.Command{
	Use:   "prompt",
	Short: "enter the interactive gnmic prompt mode",
	RunE: func(cmd *cobra.Command, args []string) error {
		err := pathCmdRun(promptDirs, promptFiles, promptExcluded, true)
		if err == nil {
			promptMode = true
		}
		// load history
		promptHistory = make([]string, 0, 256)
		home, err := homedir.Dir()
		if err != nil {
			return err
		}
		content, err := ioutil.ReadFile(home + "/.gnmic.history")
		if err != nil {
			return err
		}
		history := strings.Split(string(content), "\n")
		for i := range history {
			if history[i] != "" {
				promptHistory = append(promptHistory, history[i])
			}
		}
		return err
	},
	PostRun: func(cmd *cobra.Command, args []string) {
		cmd.ResetFlags()
		initPromptFlags(cmd)
	},
	SilenceUsage: true,
}

var promptQuitCmd = &cobra.Command{
	Use:   "quit",
	Short: "quit the gnmic-prompt",
	Run: func(cmd *cobra.Command, args []string) {
		// save history
		home, err := homedir.Dir()
		if err != nil {
			os.Exit(0)
		}
		f, err := os.Create(home + "/.gnmic.history")
		if err != nil {
			os.Exit(0)
		}
		l := len(promptHistory)
		if l > 128 {
			promptHistory = promptHistory[l-128:]
		}
		for i := range promptHistory {
			f.WriteString(promptHistory[i] + "\n")
		}
		f.Close()
		os.Exit(0)
	},
}

func init() {
	rootCmd.AddCommand(promptModeCmd)
	initPromptFlags(promptModeCmd)
}

var promptFiles []string
var promptExcluded []string
var promptDirs []string

// used to init or reset pathCmd flags for gnmic-prompt mode
func initPromptFlags(cmd *cobra.Command) {
	cmd.Flags().StringArrayVarP(&promptFiles, "file", "", []string{}, "yang files to get the paths")
	cmd.Flags().StringArrayVarP(&promptExcluded, "exclude", "", []string{}, "yang modules to be excluded from path generation")
	cmd.Flags().StringArrayVarP(&promptDirs, "dir", "", []string{}, "directories to search yang includes and imports")
	viper.BindPFlag("prompt-file", cmd.LocalFlags().Lookup("file"))
	viper.BindPFlag("prompt-exclude", cmd.LocalFlags().Lookup("exclude"))
	viper.BindPFlag("prompt-dir", cmd.LocalFlags().Lookup("dir"))
}

func findMatchedXPATH(entry *yang.Entry, word string, cursor int) []goprompt.Suggest {
	suggestions := make([]goprompt.Suggest, 0, 4)
	cword := word[cursor:]
	for name, child := range entry.Dir {
		pathelem := "/" + name
		if strings.HasPrefix(pathelem, cword) {
			node := ""
			if os.PathSeparator != '/' {
				node = fmt.Sprintf("%s%s", word[:cursor], pathelem)
				suggestions = append(suggestions, goprompt.Suggest{Text: node, Description: child.Description})
			} else {
				if len(cword) >= 1 && cword[0] == '/' {
					node = name
				} else {
					node = pathelem
				}
				suggestions = append(suggestions, goprompt.Suggest{Text: node, Description: child.Description})
			}
			if child.Key != "" { // list
				keylist := strings.Split(child.Key, " ")
				for _, key := range keylist {
					node = fmt.Sprintf("%s[%s=*]", node, key)
				}
				suggestions = append(suggestions, goprompt.Suggest{Text: node, Description: child.Description})
			}
		} else if strings.HasPrefix(cword, pathelem) {
			var prevC rune
			var bracketCount int
			var endIndex int = -1
			var stop bool
			for i, c := range cword {
				switch c {
				case '[':
					bracketCount++
				case ']':
					if prevC != '\\' {
						bracketCount--
						endIndex = i
					}
				case '/':
					if i != 0 && bracketCount == 0 {
						endIndex = i
						stop = true
					}
				}
				if stop {
					break
				}
				prevC = c
			}
			if bracketCount == 0 {
				if endIndex >= 0 {
					suggestions = append(suggestions, findMatchedXPATH(child, word, cursor+endIndex)...)
				} else {
					suggestions = append(suggestions, findMatchedXPATH(child, word, cursor+len(pathelem))...)
				}
			}
		}
	}
	return suggestions
}

var filePathCompleter = completer.FilePathCompleter{
	IgnoreCase: true,
	Filter: func(fi os.FileInfo) bool {
		fmt.Println(fi.Name())
		return fi.IsDir() || !strings.HasPrefix(fi.Name(), ".")
	},
}

var yangPathCompleter = completer.FilePathCompleter{
	IgnoreCase: true,
	Filter: func(fi os.FileInfo) bool {
		fmt.Println(fi.Name())
		return fi.IsDir() || strings.HasSuffix(fi.Name(), ".yang")
	},
}

var dirPathCompleter = completer.FilePathCompleter{
	IgnoreCase: true,
	Filter: func(fi os.FileInfo) bool {
		return fi.IsDir()
	},
}

func findDynamicSuggestions(annotation string, doc goprompt.Document) []goprompt.Suggest {
	switch annotation {
	case "XPATH":
		word := doc.GetWordBeforeCursor()
		suggestions := make([]goprompt.Suggest, 0, 16)
		for _, entry := range schemaTree.Dir {
			suggestions = append(suggestions, findMatchedXPATH(entry, word, 0)...)
		}
		return suggestions
	case "FILE":
		return filePathCompleter.Complete(doc)
	case "YANG":
		return yangPathCompleter.Complete(doc)
	case "MODEL":
		suggestions := make([]goprompt.Suggest, 0, len(schemaTree.Dir))
		for name := range schemaTree.Dir {
			suggestions = append(suggestions, goprompt.Suggest{Text: name})
		}
		return goprompt.FilterHasPrefix(suggestions, doc.GetWordBeforeCursor(), true)
	case "DIR":
		return dirPathCompleter.Complete(doc)
	case "ENCODING":
		suggestions := make([]goprompt.Suggest, 0, len(encodings))
		for _, name := range encodings {
			suggestions = append(suggestions, goprompt.Suggest{Text: name})
		}
		return goprompt.FilterHasPrefix(suggestions, doc.GetWordBeforeCursor(), true)
	}
	return []goprompt.Suggest{}
}

func showCommandArguments(b *goprompt.Buffer) {
	doc := b.Document()
	showLocalFlags := false
	command := rootCmd
	args := strings.Fields(doc.CurrentLine())
	if found, _, err := command.Find(args); err == nil {
		if command != found {
			showLocalFlags = true
		}
		command = found
	}
	maxNameLen := 0
	suggestions := make([]goprompt.Suggest, 0, 32)
	if command.HasAvailableSubCommands() {
		for _, c := range command.Commands() {
			if c.Hidden {
				continue
			}
			length := len(c.Name())
			if maxNameLen < length {
				maxNameLen = length
			}
			suggestions = append(suggestions, goprompt.Suggest{Text: c.Name(), Description: c.Short})
		}
	}
	if showLocalFlags {
		addFlags := func(flag *pflag.Flag) {
			if flag.Hidden {
				return
			}
			length := len(flag.Name)
			if maxNameLen < length+2 {
				maxNameLen = length + 2
			}
			suggestions = append(suggestions, goprompt.Suggest{Text: "--" + flag.Name, Description: flag.Usage})
		}
		command.LocalFlags().VisitAll(addFlags)
	}
	suggestions = goprompt.FilterHasPrefix(suggestions, doc.GetWordBeforeCursor(), true)
	if len(suggestions) == 0 {
		return
	}
	if err := termbox.Init(); err != nil {
		logger.Fatalf("%v", err)
	}
	w, _ := termbox.Size()
	termbox.Close()
	fmt.Printf("\n")
	maxDescLen := w - maxNameLen - 6
	format := fmt.Sprintf("  %%-%ds : %%-%ds\n", maxNameLen, maxDescLen)
	for i := range suggestions {
		length := len(suggestions[i].Description)
		if length > maxDescLen {
			fmt.Printf(format, suggestions[i].Text, suggestions[i].Description[:maxDescLen])
		} else {
			fmt.Printf(format, suggestions[i].Text, suggestions[i].Description)
		}
	}
	fmt.Printf("\n")
}

// ExecutePrompt load and run gnmic-prompt mode.
func ExecutePrompt() {
	var err error
	schemaTree, err = loadSchemaZip()
	if err != nil {
		schemaTree = buildRootEntry()
	}
	rootCmd.AddCommand(promptQuitCmd)
	rootCmd.RemoveCommand(promptModeCmd)
	shell := &cmdPrompt{
		RootCmd: rootCmd,
		GoPromptOptions: []goprompt.Option{
			goprompt.OptionTitle("gnmic-prompt"),
			goprompt.OptionPrefix("gnmic> "),
			goprompt.OptionHistory(promptHistory),
			goprompt.OptionMaxSuggestion(5),
			goprompt.OptionPrefixTextColor(goprompt.Yellow),
			// goprompt.OptionPreviewSuggestionTextColor(goprompt.Yellow),
			goprompt.OptionPreviewSuggestionBGColor(goprompt.Black),
			goprompt.OptionSuggestionTextColor(goprompt.White),
			goprompt.OptionSuggestionBGColor(goprompt.Black),
			goprompt.OptionSelectedSuggestionTextColor(goprompt.Black),
			goprompt.OptionSelectedSuggestionBGColor(goprompt.White),
			// goprompt.OptionDescriptionTextColor(goprompt.White),
			goprompt.OptionDescriptionBGColor(goprompt.Yellow),
			goprompt.OptionSelectedDescriptionTextColor(goprompt.Black),
			goprompt.OptionSelectedDescriptionBGColor(goprompt.White),
			goprompt.OptionScrollbarBGColor(goprompt.White),
			goprompt.OptionAddASCIICodeBind(goprompt.ASCIICodeBind{
				ASCIICode: []byte{0x3f}, Fn: showCommandArguments}),
			goprompt.OptionCompletionWordSeparator(completer.FilePathCompletionSeparator),
		},
	}
	shell.Run()
}

// Reference: https://github.com/stromland/cobra-prompt
// cmdPrompt requires RootCmd to run
type cmdPrompt struct {
	// RootCmd is the start point, all its sub commands and flags will be available as suggestions
	RootCmd *cobra.Command

	// GoPromptOptions is for customize go-prompt
	// see https://github.com/c-bata/go-prompt/blob/master/option.go
	GoPromptOptions []goprompt.Option
}

// Run will automatically generate suggestions for all cobra commands
// and flags defined by RootCmd and execute the selected commands.
func (co cmdPrompt) Run() {
	p := goprompt.New(
		func(in string) {
			promptArgs := strings.Fields(in)
			os.Args = append([]string{os.Args[0]}, promptArgs...)
			if len(promptArgs) > 0 {
				err := co.RootCmd.Execute()
				if err == nil && in != "" {
					promptHistory = append(promptHistory, in)
				}
			}
		},
		func(d goprompt.Document) []goprompt.Suggest {
			return findSuggestions(co, d)
		},
		co.GoPromptOptions...,
	)
	p.Run()
}

func findSuggestions(co cmdPrompt, doc goprompt.Document) []goprompt.Suggest {
	showLocalFlags := false
	command := co.RootCmd
	args := strings.Fields(doc.CurrentLine())
	if found, _, err := command.Find(args); err == nil {
		if command != found {
			showLocalFlags = true
		}
		command = found
	}

	suggestions := make([]goprompt.Suggest, 0, 32)

	// check flag annotation for the dynamic suggestion
	annotation := ""
	argnum := len(args)
	wordBefore := doc.GetWordBeforeCursor()
	if wordBefore == "" {
		if argnum >= 1 {
			annotation = command.Annotations[args[argnum-1]]
		}
	} else {
		if argnum >= 2 {
			annotation = command.Annotations[args[argnum-2]]
		}
	}
	if annotation != "" {
		return append(suggestions, findDynamicSuggestions(annotation, doc)...)
	}
	// add sub commands suggestions if they exist
	if command.HasAvailableSubCommands() {
		for _, c := range command.Commands() {
			if !c.Hidden {
				suggestions = append(suggestions, goprompt.Suggest{Text: c.Name(), Description: c.Short})
			}
		}
	}
	if showLocalFlags {
		// load local flags of the command
		addFlags := func(flag *pflag.Flag) {
			if flag.Hidden {
				return
			}
			suggestions = append(suggestions, goprompt.Suggest{Text: "--" + flag.Name, Description: flag.Usage})
		}
		command.LocalFlags().VisitAll(addFlags)
		// command.InheritedFlags().VisitAll(addFlags)
	} else {

		// persistent flags are shown if run.
		addFlags := func(flag *pflag.Flag) {
			if flag.Hidden {
				return
			}
			// if strings.HasPrefix(doc.GetWordBeforeCursor(), "--") {
			// 	suggestions = append(suggestions, goprompt.Suggest{Text: "--" + flag.Name, Description: flag.Usage})
			// } else if strings.HasPrefix(doc.GetWordBeforeCursor(), "-") && flag.Shorthand != "" {
			// 	suggestions = append(suggestions, goprompt.Suggest{Text: "-" + flag.Shorthand, Description: flag.Usage})
			// }
			if strings.HasPrefix(doc.GetWordBeforeCursor(), "-") {
				suggestions = append(suggestions, goprompt.Suggest{Text: "--" + flag.Name, Description: flag.Usage})
			}
		}
		command.LocalFlags().VisitAll(addFlags)
	}
	return goprompt.FilterHasPrefix(suggestions, doc.GetWordBeforeCursor(), true)
}
<|endoftext|>


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

package cmd

import (
	"crypto/tls"
	"encoding/json"
	"io"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"time"

	"go.uber.org/zap"

	"github.com/lvzhihao/wechat/core"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// serverCmd represents the server command
var serverCmd = &cobra.Command{
	Use:   "server",
	Short: "Wechat Proxy Server",
	Long: `Auto Refresh Access Token. For example:

wechat server --app_id=xxxx --app_secret=xxxx`,
	Run: func(cmd *cobra.Command, args []string) {
		//logger
		logger, _ := zap.NewProduction()
		defer logger.Sync() // flushes buffer, if any
		//wechat client
		client, err := core.New(&core.ClientConfig{
			AppId:          viper.GetString("app_id"),
			AppSecret:      viper.GetString("app_secret"),
			DefaultTimeout: 10 * time.Second,
		})
		//wechat config error, panic
		if err != nil {
			logger.Panic("wechat config error", zap.Any("error", err))
		}
		logger.Info("Wechat Connecting...", zap.String("token", client.FetchToken()))
		hijack(client, logger)
		logger.Info("Proxy Running...", zap.String("addr", viper.GetString("proxy_addr")))
		logger.Fatal("Proxy Stop...", zap.Any("info", http.ListenAndServeTLS(
			viper.GetString("proxy_addr"),
			viper.GetString("tls_cert"),
			viper.GetString("tls_key"),
			nil,
		)))
	},
}

func errResult(code int, msg string) string {
	b, _ := json.Marshal(struct {
		errcode int    `json:"errcode"`
		errmsg  string `json:"errmsg"`
	}{
		code,
		msg,
	})
	return string(b)
}

func randStringRunes(n int) string {
	var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

func hijack(c *core.Client, l *zap.Logger) {
	sugar := l.Sugar()
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		requestId := randStringRunes(40)
		w.Header().Set("X-REQUEST-ID", requestId)

		hj, ok := w.(http.Hijacker)
		if !ok {
			http.Error(w, errResult(-2, "not a hijacker"), http.StatusInternalServerError)
			return
		}

		in, _, err := hj.Hijack()
		if err != nil {
			l.Error("Hijack error", zap.Any("url", r.URL), zap.Error(err), zap.String("request_id", requestId))
			http.Error(w, errResult(-2, "hijack error"), http.StatusInternalServerError)
			return
		}
		defer in.Close()

		r.URL.Scheme = "https"
		r.URL.Host = "api.weixin.qq.com:443"
		v := r.URL.Query()
		v.Set("access_token", c.FetchToken())
		r.URL.RawQuery = v.Encode()

		dial, err := net.Dial("tcp", r.URL.Host)
		tls_conn := tls.Client(dial, &tls.Config{
			InsecureSkipVerify: false,
			ServerName:         "api.weixin.qq.com", //must config, see tls.config
		})
		out, _ := httputil.NewClientConn(tls_conn, nil).Hijack()
		err = r.Write(out)
		if err != nil {
			l.Error("Error copying request", zap.Any("url", r.URL), zap.Error(err), zap.String("request_id", requestId))
			http.Error(w, errResult(-2, "error copying request"), http.StatusInternalServerError)
			return
		}
		defer out.Close()

		errc := make(chan error, 2)
		cp := func(dst io.Writer, src io.Reader) {
			_, err := io.Copy(dst, src)
			errc <- err
		}

		sugar.Debugf("Request: %v", r)

		go cp(out, in)
		go cp(in, out)
		err = <-errc
		if err != nil && err != io.EOF {
			l.Error("proxy error", zap.Any("url", r.URL), zap.Error(err), zap.String("request_id", requestId))
		}
	})
}

func init() {
	RootCmd.AddCommand(serverCmd)

	rand.Seed(time.Now().UnixNano())

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// serverCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// serverCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
	viper.Set("app_id", serverCmd.Flags().String("app_id", "", "第三方用户唯一凭证"))
	viper.Set("app_secret", serverCmd.Flags().String("app_secret", "", "第三方用户唯一凭证密钥"))
	viper.Set("proxy_addr", serverCmd.Flags().String("proxy_addr", ":9099", "代理监听地址"))
	viper.Set("tls_cert", serverCmd.Flags().String("tls_cert", "./pem/server.cert", "ssl证书"))
	viper.Set("tls_key", serverCmd.Flags().String("tls_key", "./pem/server.key", "ssl证书"))
}
<commit_msg>proxy<commit_after>// Copyright © 2017 edwin <edwin.lzh@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"crypto/tls"
	"encoding/json"
	"io"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"time"

	"go.uber.org/zap"

	"github.com/lvzhihao/wechat/core"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// serverCmd represents the server command
var serverCmd = &cobra.Command{
	Use:   "server",
	Short: "Wechat Proxy Server",
	Long: `Auto Refresh Access Token. For example:

wechat server --app_id=xxxx --app_secret=xxxx`,
	Run: func(cmd *cobra.Command, args []string) {
		//logger
		logger, _ := zap.NewProduction()
		defer logger.Sync() // flushes buffer, if any
		//wechat client
		client, err := core.New(&core.ClientConfig{
			AppId:          viper.GetString("app_id"),
			AppSecret:      viper.GetString("app_secret"),
			DefaultTimeout: 10 * time.Second,
		})
		//wechat config error, panic
		if err != nil {
			logger.Panic("wechat config error", zap.Any("error", err))
		}
		logger.Info("Wechat Connecting...", zap.String("token", client.FetchToken()))
		proxy(client, logger)
		logger.Info("Proxy Running...", zap.String("addr", viper.GetString("proxy_addr")))
		logger.Fatal("Proxy Stop...", zap.Any("info", http.ListenAndServeTLS(
			viper.GetString("proxy_addr"),
			viper.GetString("tls_cert"),
			viper.GetString("tls_key"),
			nil,
		)))
	},
}

func errResult(code int, msg string) string {
	b, _ := json.Marshal(struct {
		errcode int    `json:"errcode"`
		errmsg  string `json:"errmsg"`
	}{
		code,
		msg,
	})
	return string(b)
}

func randStringRunes(n int) string {
	var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

func proxy(c *core.Client, l *zap.Logger) {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		requestId := randStringRunes(40)
		w.Header().Set("X-REQUEST-ID", requestId)

		r.URL.Scheme = "https"
		r.URL.Host = "api.weixin.qq.com:443"
		v := r.URL.Query()
		v.Set("access_token", c.FetchToken())
		r.URL.RawQuery = v.Encode()

		dial, err := net.Dial("tcp", r.URL.Host)
		tls_conn := tls.Client(dial, &tls.Config{
			InsecureSkipVerify: false,
			ServerName:         "api.weixin.qq.com", //must config, see tls.config
		})
		conn := httputil.NewClientConn(tls_conn, nil)
		rsp, err := conn.Do(r)
		l.Sugar().Infof("Request: %v", r)
		if err != nil {
			l.Error("Error copying request", zap.Any("url", r.URL), zap.Error(err), zap.String("request_id", requestId))
			http.Error(w, errResult(-2, "error copying request"), http.StatusInternalServerError)
		} else {
			for k, v := range rsp.Header {
				for _, h := range v {
					w.Header().Set(k, h)
				}
			}
			w.WriteHeader(rsp.StatusCode)
			io.Copy(w, rsp.Body)
		}
	})
}

func init() {
	RootCmd.AddCommand(serverCmd)

	rand.Seed(time.Now().UnixNano())

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// serverCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// serverCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
	viper.Set("app_id", serverCmd.Flags().String("app_id", "", "第三方用户唯一凭证"))
	viper.Set("app_secret", serverCmd.Flags().String("app_secret", "", "第三方用户唯一凭证密钥"))
	viper.Set("proxy_addr", serverCmd.Flags().String("proxy_addr", ":9099", "代理监听地址"))
	viper.Set("tls_cert", serverCmd.Flags().String("tls_cert", "./pem/server.cert", "ssl证书"))
	viper.Set("tls_key", serverCmd.Flags().String("tls_key", "./pem/server.key", "ssl证书"))
}
<|endoftext|>