<commit_before>// Copyright 2014 Unknwon
//
// 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 seago

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"path"
	"path/filepath"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/seago/seago/inject"
)

// Locale reprents a localization interface.
type Locale interface {
	Language() string
	Tr(string, ...interface{}) string
}

// Body is the request's body.
type RequestBody struct {
	reader io.ReadCloser
}

func (rb *RequestBody) Bytes() ([]byte, error) {
	return ioutil.ReadAll(rb.reader)
}

func (rb *RequestBody) String() (string, error) {
	data, err := rb.Bytes()
	return string(data), err
}

func (rb *RequestBody) ReadCloser() io.ReadCloser {
	return rb.reader
}

// A Request represents an HTTP request received by a server or to be sent by a client.
type Request struct {
	*http.Request
}

func (r *Request) Body() *RequestBody {
	return &RequestBody{r.Request.Body}
}

// Context represents the runtime context of current request of Seago instance.
// It is the integration of most frequently used middlewares and helper methods.
type Context struct {
	inject.Injector
	handlers []Handler
	action   Handler
	index    int

	*Router
	Request  Request
	Response ResponseWriter
	params   Params

	File   *multipart.File
	Render // Not nil only if you use macaran.Render middleware.
	Locale
	Data    map[string]interface{}
	statics map[string]*http.Dir // Registered static paths.
}

func (c *Context) handler() Handler {
	if c.index < len(c.handlers) {
		return c.handlers[c.index]
	}
	if c.index == len(c.handlers) {
		return c.action
	}
	panic("invalid index for context handler")
}

func (c *Context) Next() {
	c.index += 1
	c.run()
}

func (c *Context) Written() bool {
	return c.Response.Written()
}

func (c *Context) run() {
	for c.index <= len(c.handlers) {
		vals, err := c.Invoke(c.handler())
		if err != nil {
			panic(err)
		}
		c.index += 1

		// if the handler returned something, write it to the http response
		if len(vals) > 0 {
			ev := c.GetVal(reflect.TypeOf(ReturnHandler(nil)))
			handleReturn := ev.Interface().(ReturnHandler)
			handleReturn(c, vals)
		}

		if c.Written() {
			return
		}
	}
}

// RemoteAddr returns more real IP address.
func (ctx *Context) RemoteAddr() string {
	addr := ctx.Request.Header.Get("X-Real-IP")
	if len(addr) == 0 {
		addr = ctx.Request.Header.Get("X-Forwarded-For")
		if addr == "" {
			addr = ctx.Request.RemoteAddr
			if i := strings.LastIndex(addr, ":"); i > -1 {
				addr = addr[:i]
			}
		}
	}
	return addr
}

func (ctx *Context) renderHTML(status int, setName, tplName string, data ...interface{}) {
	if ctx.Render == nil {
		panic("renderer middleware hasn't been registered")
	}
	if len(data) == 0 {
		ctx.Render.HTMLSet(status, setName, tplName, ctx.Data)
	} else {
		ctx.Render.HTMLSet(status, setName, tplName, data[0])
		if len(data) > 1 {
			ctx.Render.HTMLSet(status, setName, tplName, data[0], data[1].(HTMLOptions))
		}
	}
}

// HTML calls Render.HTML but allows less arguments.
func (ctx *Context) HTML(status int, name string, data ...interface{}) {
	ctx.renderHTML(status, _DEFAULT_TPL_SET_NAME, name, data...)
}

// HTML calls Render.HTMLSet but allows less arguments.
func (ctx *Context) HTMLSet(status int, setName, tplName string, data ...interface{}) {
	ctx.renderHTML(status, setName, tplName, data...)
}

f//URL跳转
func (ctx *Context) Redirect(location string, status ...int) {
	code := http.StatusFound
	if len(status) == 1 {
		code = status[0]
	}

	http.Redirect(ctx.Response, ctx.Request.Request, location, code)
}

//获取http 请求中 GET POST PUT DELETE等方法的参数以及路由规则中的参数
func (ctx *Context) GetParam(name string) Value {
	return Value(ctx.params[name])
}

// GetFile returns information about user upload file by given form field name.
func (ctx *Context) GetFile(name string) (multipart.File, *multipart.FileHeader, error) {
	return ctx.Request.FormFile(name)
}

// SetCookie sets given cookie value to response header.
func (ctx *Context) SetCookie(name string, value string, others ...interface{}) {
	cookie := http.Cookie{}
	cookie.Name = name
	cookie.Value = value

	if len(others) > 0 {
		switch v := others[0].(type) {
		case int:
			cookie.MaxAge = v
		case int64:
			cookie.MaxAge = int(v)
		case int32:
			cookie.MaxAge = int(v)
		}
	}

	// default "/"
	if len(others) > 1 {
		if v, ok := others[1].(string); ok && len(v) > 0 {
			cookie.Path = v
		}
	} else {
		cookie.Path = "/"
	}

	// default empty
	if len(others) > 2 {
		if v, ok := others[2].(string); ok && len(v) > 0 {
			cookie.Domain = v
		}
	}

	// default empty
	if len(others) > 3 {
		switch v := others[3].(type) {
		case bool:
			cookie.Secure = v
		default:
			if others[3] != nil {
				cookie.Secure = true
			}
		}
	}

	// default false. for session cookie default true
	if len(others) > 4 {
		if v, ok := others[4].(bool); ok && v {
			cookie.HttpOnly = true
		}
	}

	ctx.Response.Header().Add("Set-Cookie", cookie.String())
}

// GetCookie returns given cookie value from request header.
func (ctx *Context) GetCookie(name string) string {
	cookie, err := ctx.Request.Cookie(name)
	if err != nil {
		return ""
	}
	return cookie.Value
}

var defaultCookieSecret string

// SetDefaultCookieSecret sets global default secure cookie secret.
func (m *Seago) SetDefaultCookieSecret(secret string) {
	defaultCookieSecret = secret
}

// SetSecureCookie sets given cookie value to response header with default secret string.
func (ctx *Context) SetSecureCookie(name, value string, others ...interface{}) {
	ctx.SetSuperSecureCookie(defaultCookieSecret, name, value, others...)
}

// GetSecureCookie returns given cookie value from request header with default secret string.
func (ctx *Context) GetSecureCookie(key string) (string, bool) {
	return ctx.GetSuperSecureCookie(defaultCookieSecret, key)
}

// SetSuperSecureCookie sets given cookie value to response header with secret string.
func (ctx *Context) SetSuperSecureCookie(Secret, name, value string, others ...interface{}) {
	vs := base64.URLEncoding.EncodeToString([]byte(value))
	timestamp := strconv.FormatInt(time.Now().UnixNano(), 10)
	h := hmac.New(sha1.New, []byte(Secret))
	fmt.Fprintf(h, "%s%s", vs, timestamp)
	sig := fmt.Sprintf("%02x", h.Sum(nil))
	cookie := strings.Join([]string{vs, timestamp, sig}, "|")
	ctx.SetCookie(name, cookie, others...)
}

// GetSuperSecureCookie returns given cookie value from request header with secret string.
func (ctx *Context) GetSuperSecureCookie(Secret, key string) (string, bool) {
	val := ctx.GetCookie(key)
	if val == "" {
		return "", false
	}

	parts := strings.SplitN(val, "|", 3)

	if len(parts) != 3 {
		return "", false
	}

	vs := parts[0]
	timestamp := parts[1]
	sig := parts[2]

	h := hmac.New(sha1.New, []byte(Secret))
	fmt.Fprintf(h, "%s%s", vs, timestamp)

	if fmt.Sprintf("%02x", h.Sum(nil)) != sig {
		return "", false
	}
	res, _ := base64.URLEncoding.DecodeString(vs)
	return string(res), true
}

// ServeFile serves given file to response.
func (ctx *Context) ServeFile(file string, names ...string) {
	var name string
	if len(names) > 0 {
		name = names[0]
	} else {
		name = path.Base(file)
	}
	ctx.Response.Header().Set("Content-Description", "File Transfer")
	ctx.Response.Header().Set("Content-Type", "application/octet-stream")
	ctx.Response.Header().Set("Content-Disposition", "attachment; filename="+name)
	ctx.Response.Header().Set("Content-Transfer-Encoding", "binary")
	ctx.Response.Header().Set("Expires", "0")
	ctx.Response.Header().Set("Cache-Control", "must-revalidate")
	ctx.Response.Header().Set("Pragma", "public")
	http.ServeFile(ctx.Response, ctx.Request.Request, file)
}

// ChangeStaticPath changes static path from old to new one.
func (ctx *Context) ChangeStaticPath(oldPath, newPath string) {
	if !filepath.IsAbs(oldPath) {
		oldPath = filepath.Join(Root, oldPath)
	}
	if _, ok := ctx.statics[oldPath]; ok {
		dir := ctx.statics[oldPath]
		delete(ctx.statics, oldPath)

		if !filepath.IsAbs(newPath) {
			newPath = filepath.Join(Root, newPath)
		}
		*dir = http.Dir(newPath)
		ctx.statics[newPath] = dir
	}
}
<commit_msg>消除多余字符<commit_after>// Copyright 2014 Unknwon
//
// 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 seago

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"path"
	"path/filepath"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/seago/seago/inject"
)

// Locale reprents a localization interface.
type Locale interface {
	Language() string
	Tr(string, ...interface{}) string
}

// Body is the request's body.
type RequestBody struct {
	reader io.ReadCloser
}

func (rb *RequestBody) Bytes() ([]byte, error) {
	return ioutil.ReadAll(rb.reader)
}

func (rb *RequestBody) String() (string, error) {
	data, err := rb.Bytes()
	return string(data), err
}

func (rb *RequestBody) ReadCloser() io.ReadCloser {
	return rb.reader
}

// A Request represents an HTTP request received by a server or to be sent by a client.
type Request struct {
	*http.Request
}

func (r *Request) Body() *RequestBody {
	return &RequestBody{r.Request.Body}
}

// Context represents the runtime context of current request of Seago instance.
// It is the integration of most frequently used middlewares and helper methods.
type Context struct {
	inject.Injector
	handlers []Handler
	action   Handler
	index    int

	*Router
	Request  Request
	Response ResponseWriter
	params   Params

	File   *multipart.File
	Render // Not nil only if you use macaran.Render middleware.
	Locale
	Data    map[string]interface{}
	statics map[string]*http.Dir // Registered static paths.
}

func (c *Context) handler() Handler {
	if c.index < len(c.handlers) {
		return c.handlers[c.index]
	}
	if c.index == len(c.handlers) {
		return c.action
	}
	panic("invalid index for context handler")
}

func (c *Context) Next() {
	c.index += 1
	c.run()
}

func (c *Context) Written() bool {
	return c.Response.Written()
}

func (c *Context) run() {
	for c.index <= len(c.handlers) {
		vals, err := c.Invoke(c.handler())
		if err != nil {
			panic(err)
		}
		c.index += 1

		// if the handler returned something, write it to the http response
		if len(vals) > 0 {
			ev := c.GetVal(reflect.TypeOf(ReturnHandler(nil)))
			handleReturn := ev.Interface().(ReturnHandler)
			handleReturn(c, vals)
		}

		if c.Written() {
			return
		}
	}
}

// RemoteAddr returns more real IP address.
func (ctx *Context) RemoteAddr() string {
	addr := ctx.Request.Header.Get("X-Real-IP")
	if len(addr) == 0 {
		addr = ctx.Request.Header.Get("X-Forwarded-For")
		if addr == "" {
			addr = ctx.Request.RemoteAddr
			if i := strings.LastIndex(addr, ":"); i > -1 {
				addr = addr[:i]
			}
		}
	}
	return addr
}

func (ctx *Context) renderHTML(status int, setName, tplName string, data ...interface{}) {
	if ctx.Render == nil {
		panic("renderer middleware hasn't been registered")
	}
	if len(data) == 0 {
		ctx.Render.HTMLSet(status, setName, tplName, ctx.Data)
	} else {
		ctx.Render.HTMLSet(status, setName, tplName, data[0])
		if len(data) > 1 {
			ctx.Render.HTMLSet(status, setName, tplName, data[0], data[1].(HTMLOptions))
		}
	}
}

// HTML calls Render.HTML but allows less arguments.
func (ctx *Context) HTML(status int, name string, data ...interface{}) {
	ctx.renderHTML(status, _DEFAULT_TPL_SET_NAME, name, data...)
}

// HTML calls Render.HTMLSet but allows less arguments.
func (ctx *Context) HTMLSet(status int, setName, tplName string, data ...interface{}) {
	ctx.renderHTML(status, setName, tplName, data...)
}

//URL跳转
func (ctx *Context) Redirect(location string, status ...int) {
	code := http.StatusFound
	if len(status) == 1 {
		code = status[0]
	}

	http.Redirect(ctx.Response, ctx.Request.Request, location, code)
}

//获取http 请求中 GET POST PUT DELETE等方法的参数以及路由规则中的参数
func (ctx *Context) GetParam(name string) Value {
	return Value(ctx.params[name])
}

// GetFile returns information about user upload file by given form field name.
func (ctx *Context) GetFile(name string) (multipart.File, *multipart.FileHeader, error) {
	return ctx.Request.FormFile(name)
}

// SetCookie sets given cookie value to response header.
func (ctx *Context) SetCookie(name string, value string, others ...interface{}) {
	cookie := http.Cookie{}
	cookie.Name = name
	cookie.Value = value

	if len(others) > 0 {
		switch v := others[0].(type) {
		case int:
			cookie.MaxAge = v
		case int64:
			cookie.MaxAge = int(v)
		case int32:
			cookie.MaxAge = int(v)
		}
	}

	// default "/"
	if len(others) > 1 {
		if v, ok := others[1].(string); ok && len(v) > 0 {
			cookie.Path = v
		}
	} else {
		cookie.Path = "/"
	}

	// default empty
	if len(others) > 2 {
		if v, ok := others[2].(string); ok && len(v) > 0 {
			cookie.Domain = v
		}
	}

	// default empty
	if len(others) > 3 {
		switch v := others[3].(type) {
		case bool:
			cookie.Secure = v
		default:
			if others[3] != nil {
				cookie.Secure = true
			}
		}
	}

	// default false. for session cookie default true
	if len(others) > 4 {
		if v, ok := others[4].(bool); ok && v {
			cookie.HttpOnly = true
		}
	}

	ctx.Response.Header().Add("Set-Cookie", cookie.String())
}

// GetCookie returns given cookie value from request header.
func (ctx *Context) GetCookie(name string) string {
	cookie, err := ctx.Request.Cookie(name)
	if err != nil {
		return ""
	}
	return cookie.Value
}

var defaultCookieSecret string

// SetDefaultCookieSecret sets global default secure cookie secret.
func (m *Seago) SetDefaultCookieSecret(secret string) {
	defaultCookieSecret = secret
}

// SetSecureCookie sets given cookie value to response header with default secret string.
func (ctx *Context) SetSecureCookie(name, value string, others ...interface{}) {
	ctx.SetSuperSecureCookie(defaultCookieSecret, name, value, others...)
}

// GetSecureCookie returns given cookie value from request header with default secret string.
func (ctx *Context) GetSecureCookie(key string) (string, bool) {
	return ctx.GetSuperSecureCookie(defaultCookieSecret, key)
}

// SetSuperSecureCookie sets given cookie value to response header with secret string.
func (ctx *Context) SetSuperSecureCookie(Secret, name, value string, others ...interface{}) {
	vs := base64.URLEncoding.EncodeToString([]byte(value))
	timestamp := strconv.FormatInt(time.Now().UnixNano(), 10)
	h := hmac.New(sha1.New, []byte(Secret))
	fmt.Fprintf(h, "%s%s", vs, timestamp)
	sig := fmt.Sprintf("%02x", h.Sum(nil))
	cookie := strings.Join([]string{vs, timestamp, sig}, "|")
	ctx.SetCookie(name, cookie, others...)
}

// GetSuperSecureCookie returns given cookie value from request header with secret string.
func (ctx *Context) GetSuperSecureCookie(Secret, key string) (string, bool) {
	val := ctx.GetCookie(key)
	if val == "" {
		return "", false
	}

	parts := strings.SplitN(val, "|", 3)

	if len(parts) != 3 {
		return "", false
	}

	vs := parts[0]
	timestamp := parts[1]
	sig := parts[2]

	h := hmac.New(sha1.New, []byte(Secret))
	fmt.Fprintf(h, "%s%s", vs, timestamp)

	if fmt.Sprintf("%02x", h.Sum(nil)) != sig {
		return "", false
	}
	res, _ := base64.URLEncoding.DecodeString(vs)
	return string(res), true
}

// ServeFile serves given file to response.
func (ctx *Context) ServeFile(file string, names ...string) {
	var name string
	if len(names) > 0 {
		name = names[0]
	} else {
		name = path.Base(file)
	}
	ctx.Response.Header().Set("Content-Description", "File Transfer")
	ctx.Response.Header().Set("Content-Type", "application/octet-stream")
	ctx.Response.Header().Set("Content-Disposition", "attachment; filename="+name)
	ctx.Response.Header().Set("Content-Transfer-Encoding", "binary")
	ctx.Response.Header().Set("Expires", "0")
	ctx.Response.Header().Set("Cache-Control", "must-revalidate")
	ctx.Response.Header().Set("Pragma", "public")
	http.ServeFile(ctx.Response, ctx.Request.Request, file)
}

// ChangeStaticPath changes static path from old to new one.
func (ctx *Context) ChangeStaticPath(oldPath, newPath string) {
	if !filepath.IsAbs(oldPath) {
		oldPath = filepath.Join(Root, oldPath)
	}
	if _, ok := ctx.statics[oldPath]; ok {
		dir := ctx.statics[oldPath]
		delete(ctx.statics, oldPath)

		if !filepath.IsAbs(newPath) {
			newPath = filepath.Join(Root, newPath)
		}
		*dir = http.Dir(newPath)
		ctx.statics[newPath] = dir
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"strings"
)

// Context is a type that is passed through to
// each Handler action in a cli application. Context
// can be used to retrieve context-specific args and
// parsed command-line options.
type Context struct {
	context.Context
	App           *App
	Command       *Command
	shellComplete bool
	setFlags      map[string]bool
	flagSet       *flag.FlagSet
	parentContext *Context
}

// NewContext creates a new context. For use in when invoking an App or Command action.
func NewContext(app *App, set *flag.FlagSet, parentCtx *Context) *Context {
	c := &Context{
		App:           app,
		flagSet:       set,
		parentContext: parentCtx,
	}

	if parentCtx != nil {
		c.Context = parentCtx.Context
		c.shellComplete = parentCtx.shellComplete

		if parentCtx.flagSet == nil {
			parentCtx.flagSet = &flag.FlagSet{}
		}
	}

	c.Command = &Command{}

	if c.Context == nil {
		c.Context = context.Background()
	}

	return c
}

// NumFlags returns the number of flags set
func (c *Context) NumFlags() int {
	return c.flagSet.NFlag()
}

// Set sets a context flag to a value.
func (c *Context) Set(name, value string) error {
	return c.flagSet.Set(name, value)
}

// IsSet determines if the flag was actually set
func (c *Context) IsSet(name string) bool {
	if fs := lookupFlagSet(name, c); fs != nil {
		if fs := lookupFlagSet(name, c); fs != nil {
			isSet := false
			fs.Visit(func(f *flag.Flag) {
				if f.Name == name {
					isSet = true
				}
			})
			if isSet {
				return true
			}
		}

		f := lookupFlag(name, c)
		if f == nil {
			return false
		}

		return f.IsSet()
	}

	return false
}

// LocalFlagNames returns a slice of flag names used in this context.
func (c *Context) LocalFlagNames() []string {
	var names []string
	c.flagSet.Visit(makeFlagNameVisitor(&names))
	return names
}

// FlagNames returns a slice of flag names used by the this context and all of
// its parent contexts.
func (c *Context) FlagNames() []string {
	var names []string
	for _, ctx := range c.Lineage() {
		ctx.flagSet.Visit(makeFlagNameVisitor(&names))
	}
	return names
}

// Lineage returns *this* context and all of its ancestor contexts in order from
// child to parent
func (c *Context) Lineage() []*Context {
	var lineage []*Context

	for cur := c; cur != nil; cur = cur.parentContext {
		lineage = append(lineage, cur)
	}

	return lineage
}

// Value returns the value of the flag corresponding to `name`
func (c *Context) Value(name string) interface{} {
	return c.flagSet.Lookup(name).Value.(flag.Getter).Get()
}

// Args returns the command line arguments associated with the context.
func (c *Context) Args() Args {
	ret := args(c.flagSet.Args())
	return &ret
}

// NArg returns the number of the command line arguments.
func (c *Context) NArg() int {
	return c.Args().Len()
}

func lookupFlag(name string, ctx *Context) Flag {
	for _, c := range ctx.Lineage() {
		if c.Command == nil {
			continue
		}

		for _, f := range c.Command.Flags {
			for _, n := range f.Names() {
				if n == name {
					return f
				}
			}
		}
	}

	if ctx.App != nil {
		for _, f := range ctx.App.Flags {
			for _, n := range f.Names() {
				if n == name {
					return f
				}
			}
		}
	}

	return nil
}

func lookupFlagSet(name string, ctx *Context) *flag.FlagSet {
	for _, c := range ctx.Lineage() {
		if f := c.flagSet.Lookup(name); f != nil {
			return c.flagSet
		}
	}

	return nil
}

func copyFlag(name string, ff *flag.Flag, set *flag.FlagSet) {
	switch ff.Value.(type) {
	case Serializer:
		_ = set.Set(name, ff.Value.(Serializer).Serialize())
	default:
		_ = set.Set(name, ff.Value.String())
	}
}

func normalizeFlags(flags []Flag, set *flag.FlagSet) error {
	visited := make(map[string]bool)
	set.Visit(func(f *flag.Flag) {
		visited[f.Name] = true
	})
	for _, f := range flags {
		parts := f.Names()
		if len(parts) == 1 {
			continue
		}
		var ff *flag.Flag
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if visited[name] {
				if ff != nil {
					return errors.New("Cannot use two forms of the same flag: " + name + " " + ff.Name)
				}
				ff = set.Lookup(name)
			}
		}
		if ff == nil {
			continue
		}
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if !visited[name] {
				copyFlag(name, ff, set)
			}
		}
	}
	return nil
}

func makeFlagNameVisitor(names *[]string) func(*flag.Flag) {
	return func(f *flag.Flag) {
		nameParts := strings.Split(f.Name, ",")
		name := strings.TrimSpace(nameParts[0])

		for _, part := range nameParts {
			part = strings.TrimSpace(part)
			if len(part) > len(name) {
				name = part
			}
		}

		if name != "" {
			*names = append(*names, name)
		}
	}
}

type requiredFlagsErr interface {
	error
	getMissingFlags() []string
}

type errRequiredFlags struct {
	missingFlags []string
}

func (e *errRequiredFlags) Error() string {
	numberOfMissingFlags := len(e.missingFlags)
	if numberOfMissingFlags == 1 {
		return fmt.Sprintf("Required flag %q not set", e.missingFlags[0])
	}
	joinedMissingFlags := strings.Join(e.missingFlags, ", ")
	return fmt.Sprintf("Required flags %q not set", joinedMissingFlags)
}

func (e *errRequiredFlags) getMissingFlags() []string {
	return e.missingFlags
}

func checkRequiredFlags(flags []Flag, context *Context) requiredFlagsErr {
	var missingFlags []string
	for _, f := range flags {
		if rf, ok := f.(RequiredFlag); ok && rf.IsRequired() {
			var flagPresent bool
			var flagName string

			for _, key := range f.Names() {
				if len(key) > 1 {
					flagName = key
				}

				if context.IsSet(strings.TrimSpace(key)) {
					flagPresent = true
				}
			}

			if !flagPresent && flagName != "" {
				missingFlags = append(missingFlags, flagName)
			}
		}
	}

	if len(missingFlags) != 0 {
		return &errRequiredFlags{missingFlags: missingFlags}
	}

	return nil
}
<commit_msg>reduce diff<commit_after>package cli

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"strings"
)

// Context is a type that is passed through to
// each Handler action in a cli application. Context
// can be used to retrieve context-specific args and
// parsed command-line options.
type Context struct {
	context.Context
	App           *App
	Command       *Command
	shellComplete bool
	setFlags      map[string]bool
	flagSet       *flag.FlagSet
	parentContext *Context
}

// NewContext creates a new context. For use in when invoking an App or Command action.
func NewContext(app *App, set *flag.FlagSet, parentCtx *Context) *Context {
	c := &Context{App: app, flagSet: set, parentContext: parentCtx}

	if parentCtx != nil {
		c.Context = parentCtx.Context
		c.shellComplete = parentCtx.shellComplete

		if parentCtx.flagSet == nil {
			parentCtx.flagSet = &flag.FlagSet{}
		}
	}

	c.Command = &Command{}

	if c.Context == nil {
		c.Context = context.Background()
	}

	return c
}

// NumFlags returns the number of flags set
func (c *Context) NumFlags() int {
	return c.flagSet.NFlag()
}

// Set sets a context flag to a value.
func (c *Context) Set(name, value string) error {
	return c.flagSet.Set(name, value)
}

// IsSet determines if the flag was actually set
func (c *Context) IsSet(name string) bool {
	if fs := lookupFlagSet(name, c); fs != nil {
		if fs := lookupFlagSet(name, c); fs != nil {
			isSet := false
			fs.Visit(func(f *flag.Flag) {
				if f.Name == name {
					isSet = true
				}
			})
			if isSet {
				return true
			}
		}

		f := lookupFlag(name, c)
		if f == nil {
			return false
		}

		return f.IsSet()
	}

	return false
}

// LocalFlagNames returns a slice of flag names used in this context.
func (c *Context) LocalFlagNames() []string {
	var names []string
	c.flagSet.Visit(makeFlagNameVisitor(&names))
	return names
}

// FlagNames returns a slice of flag names used by the this context and all of
// its parent contexts.
func (c *Context) FlagNames() []string {
	var names []string
	for _, ctx := range c.Lineage() {
		ctx.flagSet.Visit(makeFlagNameVisitor(&names))
	}
	return names
}

// Lineage returns *this* context and all of its ancestor contexts in order from
// child to parent
func (c *Context) Lineage() []*Context {
	var lineage []*Context

	for cur := c; cur != nil; cur = cur.parentContext {
		lineage = append(lineage, cur)
	}

	return lineage
}

// Value returns the value of the flag corresponding to `name`
func (c *Context) Value(name string) interface{} {
	return c.flagSet.Lookup(name).Value.(flag.Getter).Get()
}

// Args returns the command line arguments associated with the context.
func (c *Context) Args() Args {
	ret := args(c.flagSet.Args())
	return &ret
}

// NArg returns the number of the command line arguments.
func (c *Context) NArg() int {
	return c.Args().Len()
}

func lookupFlag(name string, ctx *Context) Flag {
	for _, c := range ctx.Lineage() {
		if c.Command == nil {
			continue
		}

		for _, f := range c.Command.Flags {
			for _, n := range f.Names() {
				if n == name {
					return f
				}
			}
		}
	}

	if ctx.App != nil {
		for _, f := range ctx.App.Flags {
			for _, n := range f.Names() {
				if n == name {
					return f
				}
			}
		}
	}

	return nil
}

func lookupFlagSet(name string, ctx *Context) *flag.FlagSet {
	for _, c := range ctx.Lineage() {
		if f := c.flagSet.Lookup(name); f != nil {
			return c.flagSet
		}
	}

	return nil
}

func copyFlag(name string, ff *flag.Flag, set *flag.FlagSet) {
	switch ff.Value.(type) {
	case Serializer:
		_ = set.Set(name, ff.Value.(Serializer).Serialize())
	default:
		_ = set.Set(name, ff.Value.String())
	}
}

func normalizeFlags(flags []Flag, set *flag.FlagSet) error {
	visited := make(map[string]bool)
	set.Visit(func(f *flag.Flag) {
		visited[f.Name] = true
	})
	for _, f := range flags {
		parts := f.Names()
		if len(parts) == 1 {
			continue
		}
		var ff *flag.Flag
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if visited[name] {
				if ff != nil {
					return errors.New("Cannot use two forms of the same flag: " + name + " " + ff.Name)
				}
				ff = set.Lookup(name)
			}
		}
		if ff == nil {
			continue
		}
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if !visited[name] {
				copyFlag(name, ff, set)
			}
		}
	}
	return nil
}

func makeFlagNameVisitor(names *[]string) func(*flag.Flag) {
	return func(f *flag.Flag) {
		nameParts := strings.Split(f.Name, ",")
		name := strings.TrimSpace(nameParts[0])

		for _, part := range nameParts {
			part = strings.TrimSpace(part)
			if len(part) > len(name) {
				name = part
			}
		}

		if name != "" {
			*names = append(*names, name)
		}
	}
}

type requiredFlagsErr interface {
	error
	getMissingFlags() []string
}

type errRequiredFlags struct {
	missingFlags []string
}

func (e *errRequiredFlags) Error() string {
	numberOfMissingFlags := len(e.missingFlags)
	if numberOfMissingFlags == 1 {
		return fmt.Sprintf("Required flag %q not set", e.missingFlags[0])
	}
	joinedMissingFlags := strings.Join(e.missingFlags, ", ")
	return fmt.Sprintf("Required flags %q not set", joinedMissingFlags)
}

func (e *errRequiredFlags) getMissingFlags() []string {
	return e.missingFlags
}

func checkRequiredFlags(flags []Flag, context *Context) requiredFlagsErr {
	var missingFlags []string
	for _, f := range flags {
		if rf, ok := f.(RequiredFlag); ok && rf.IsRequired() {
			var flagPresent bool
			var flagName string

			for _, key := range f.Names() {
				if len(key) > 1 {
					flagName = key
				}

				if context.IsSet(strings.TrimSpace(key)) {
					flagPresent = true
				}
			}

			if !flagPresent && flagName != "" {
				missingFlags = append(missingFlags, flagName)
			}
		}
	}

	if len(missingFlags) != 0 {
		return &errRequiredFlags{missingFlags: missingFlags}
	}

	return nil
}
<|endoftext|>


<commit_before>package echo

import (
	"context"
	"io"
	"mime/multipart"
	"net/http"
	"time"

	"github.com/admpub/events"

	"github.com/webx-top/echo/engine"
	"github.com/webx-top/echo/logger"
	"github.com/webx-top/echo/param"
	pkgCode "github.com/webx-top/echo/code"
)

// Context represents context for the current request. It holds request and
// response objects, path parameters, data and registered handler.
type Context interface {
	context.Context
	events.Emitter
	SetEmitter(events.Emitter)

	//Transaction
	SetTransaction(t Transaction)
	Transaction() Transaction
	Begin() error
	Rollback() error
	Commit() error
	End(succeed bool) error

	//Standard Context
	StdContext() context.Context
	SetStdContext(context.Context)

	Validator
	SetValidator(Validator)
	Translator
	SetTranslator(Translator)
	Request() engine.Request
	Response() engine.Response
	Handle(Context) error
	Logger() logger.Logger
	Object() *xContext
	Echo() *Echo
	Route() *Route
	Reset(engine.Request, engine.Response)

	//----------------
	// Param
	//----------------

	Path() string
	P(int, ...string) string
	Param(string, ...string) string
	// ParamNames returns path parameter names.
	ParamNames() []string
	ParamValues() []string
	SetParamValues(values ...string)
	// Host
	HostNames() []string
	HostValues() []string
	HostParam(string, ...string) string
	HostP(int, ...string) string
	setHostParamValues([]string, []string)

	// Queries returns the query parameters as map. It is an alias for `engine.URL#Query()`.
	Queries() map[string][]string
	QueryValues(string) []string
	QueryxValues(string) param.StringSlice
	Query(string, ...string) string

	//----------------
	// Form data
	//----------------

	Form(string, ...string) string
	FormValues(string) []string
	FormxValues(string) param.StringSlice
	// Forms returns the form parameters as map. It is an alias for `engine.Request#Form().All()`.
	Forms() map[string][]string

	// Param+
	Px(int, ...string) param.String
	Paramx(string, ...string) param.String
	Queryx(string, ...string) param.String
	Formx(string, ...string) param.String
	// string to param.String
	Atop(string) param.String
	ToParamString(string) param.String
	ToStringSlice([]string) param.StringSlice

	//----------------
	// Context data
	//----------------

	Set(string, interface{})
	Get(string, ...interface{}) interface{}
	Delete(...string)
	Stored() Store
	Internal() *param.SafeMap

	//----------------
	// Bind
	//----------------

	Bind(interface{}, ...FormDataFilter) error
	MustBind(interface{}, ...FormDataFilter) error

	//----------------
	// Response data
	//----------------

	Render(string, interface{}, ...int) error
	HTML(string, ...int) error
	String(string, ...int) error
	Blob([]byte, ...int) error
	JSON(interface{}, ...int) error
	JSONBlob([]byte, ...int) error
	JSONP(string, interface{}, ...int) error
	XML(interface{}, ...int) error
	XMLBlob([]byte, ...int) error
	Stream(func(io.Writer) bool)
	SSEvent(string, chan interface{}) error
	File(string, ...http.FileSystem) error
	Attachment(io.Reader, string, ...bool) error
	NoContent(...int) error
	Redirect(string, ...int) error
	Error(err error)
	NewError(code pkgCode.Code, msg string, args ...interface{}) *Error
	SetCode(int)
	Code() int
	SetData(Data)
	Data() Data

	// ServeContent sends static content from `io.Reader` and handles caching
	// via `If-Modified-Since` request header. It automatically sets `Content-Type`
	// and `Last-Modified` response headers.
	ServeContent(io.Reader, string, time.Time) error
	ServeCallbackContent(func(Context) (io.Reader, error), string, time.Time) error

	//----------------
	// FuncMap
	//----------------

	SetFunc(string, interface{})
	GetFunc(string) interface{}
	ResetFuncs(map[string]interface{})
	Funcs() map[string]interface{}
	PrintFuncs()

	//----------------
	// Render
	//----------------
	SetAuto(on bool) Context
	Fetch(string, interface{}) ([]byte, error)
	SetRenderer(Renderer)

	//----------------
	// Cookie
	//----------------

	SetCookieOptions(*CookieOptions)
	CookieOptions() *CookieOptions
	NewCookie(string, string) *Cookie
	Cookie() Cookier
	GetCookie(string) string
	SetCookie(string, string, ...interface{})

	//----------------
	// Session
	//----------------

	SetSessionOptions(*SessionOptions)
	SessionOptions() *SessionOptions
	SetSessioner(Sessioner)
	Session() Sessioner
	Flash(...string) interface{}

	//----------------
	// Request data
	//----------------

	Header(string) string
	IsAjax() bool
	IsPjax() bool
	PjaxContainer() string
	Method() string
	Format() string
	SetFormat(string)
	IsPost() bool
	IsGet() bool
	IsPut() bool
	IsDel() bool
	IsHead() bool
	IsPatch() bool
	IsOptions() bool
	IsSecure() bool
	IsWebsocket() bool
	IsUpload() bool
	ResolveContentType() string
	WithFormatExtension(bool)
	ResolveFormat() string
	Accept() *Accepts
	Protocol() string
	Site() string
	Scheme() string
	Domain() string
	Host() string
	Proxy() []string
	Referer() string
	Port() int
	RealIP() string
	HasAnyRequest() bool

	MapForm(i interface{}, names ...string) error
	MapData(i interface{}, data map[string][]string, names ...string) error
	SaveUploadedFile(fieldName string, saveAbsPath string, saveFileName ...string) (*multipart.FileHeader, error)
	SaveUploadedFileToWriter(string, io.Writer) (*multipart.FileHeader, error)
	//Multiple file upload
	SaveUploadedFiles(fieldName string, savePath func(*multipart.FileHeader) (string, error)) error
	SaveUploadedFilesToWriter(fieldName string, writer func(*multipart.FileHeader) (io.Writer, error)) error

	//----------------
	// Hook
	//----------------

	AddPreResponseHook(func() error) Context
	SetPreResponseHook(...func() error) Context
}
<commit_msg>update<commit_after>package echo

import (
	"context"
	"io"
	"mime/multipart"
	"net/http"
	"time"

	"github.com/admpub/events"

	"github.com/webx-top/echo/engine"
	"github.com/webx-top/echo/logger"
	"github.com/webx-top/echo/param"
	pkgCode "github.com/webx-top/echo/code"
)

// Context represents context for the current request. It holds request and
// response objects, path parameters, data and registered handler.
type Context interface {
	context.Context
	events.Emitter
	SetEmitter(events.Emitter)

	//Transaction
	SetTransaction(t Transaction)
	Transaction() Transaction
	Begin() error
	Rollback() error
	Commit() error
	End(succeed bool) error

	//Standard Context
	StdContext() context.Context
	SetStdContext(context.Context)

	Validator
	SetValidator(Validator)
	Translator
	SetTranslator(Translator)
	Request() engine.Request
	Response() engine.Response
	Handle(Context) error
	Logger() logger.Logger
	Object() *xContext
	Echo() *Echo
	Route() *Route
	Reset(engine.Request, engine.Response)

	//----------------
	// Param
	//----------------

	Path() string
	P(int, ...string) string
	Param(string, ...string) string
	// ParamNames returns path parameter names.
	ParamNames() []string
	ParamValues() []string
	SetParamValues(values ...string)
	// Host
	HostNames() []string
	HostValues() []string
	HostParam(string, ...string) string
	HostP(int, ...string) string
	setHostParamValues([]string, []string)

	// Queries returns the query parameters as map. It is an alias for `engine.URL#Query()`.
	Queries() map[string][]string
	QueryValues(string) []string
	QueryxValues(string) param.StringSlice
	Query(string, ...string) string

	//----------------
	// Form data
	//----------------

	Form(string, ...string) string
	FormValues(string) []string
	FormxValues(string) param.StringSlice
	// Forms returns the form parameters as map. It is an alias for `engine.Request#Form().All()`.
	Forms() map[string][]string

	// Param+
	Px(int, ...string) param.String
	Paramx(string, ...string) param.String
	Queryx(string, ...string) param.String
	Formx(string, ...string) param.String
	// string to param.String
	Atop(string) param.String
	ToParamString(string) param.String
	ToStringSlice([]string) param.StringSlice

	//----------------
	// Context data
	//----------------

	Set(string, interface{})
	Get(string, ...interface{}) interface{}
	Delete(...string)
	Stored() Store
	Internal() *param.SafeMap

	//----------------
	// Bind
	//----------------

	Bind(interface{}, ...FormDataFilter) error
	MustBind(interface{}, ...FormDataFilter) error

	//----------------
	// Response data
	//----------------

	Render(string, interface{}, ...int) error
	HTML(string, ...int) error
	String(string, ...int) error
	Blob([]byte, ...int) error
	JSON(interface{}, ...int) error
	JSONBlob([]byte, ...int) error
	JSONP(string, interface{}, ...int) error
	XML(interface{}, ...int) error
	XMLBlob([]byte, ...int) error
	Stream(func(io.Writer) bool)
	SSEvent(string, chan interface{}) error
	File(string, ...http.FileSystem) error
	Attachment(io.Reader, string, ...bool) error
	NoContent(...int) error
	Redirect(string, ...int) error
	Error(err error)
	NewError(code pkgCode.Code, msg string, args ...interface{}) *Error
	SetCode(int)
	Code() int
	SetData(Data)
	Data() Data

	// ServeContent sends static content from `io.Reader` and handles caching
	// via `If-Modified-Since` request header. It automatically sets `Content-Type`
	// and `Last-Modified` response headers.
	ServeContent(io.Reader, string, time.Time) error
	ServeCallbackContent(func(Context) (io.Reader, error), string, time.Time) error

	//----------------
	// FuncMap
	//----------------

	SetFunc(string, interface{})
	GetFunc(string) interface{}
	ResetFuncs(map[string]interface{})
	Funcs() map[string]interface{}
	PrintFuncs()

	//----------------
	// Render
	//----------------
	SetAuto(on bool) Context
	Fetch(string, interface{}) ([]byte, error)
	SetRenderer(Renderer)

	//----------------
	// Cookie
	//----------------

	SetCookieOptions(*CookieOptions)
	CookieOptions() *CookieOptions
	NewCookie(string, string) *Cookie
	Cookie() Cookier
	GetCookie(string) string
	// SetCookie @param:key,value,maxAge(seconds),path(/),domain,secure,httpOnly,sameSite(lax/strict/default)
	SetCookie(string, string, ...interface{})

	//----------------
	// Session
	//----------------

	SetSessionOptions(*SessionOptions)
	SessionOptions() *SessionOptions
	SetSessioner(Sessioner)
	Session() Sessioner
	Flash(...string) interface{}

	//----------------
	// Request data
	//----------------

	Header(string) string
	IsAjax() bool
	IsPjax() bool
	PjaxContainer() string
	Method() string
	Format() string
	SetFormat(string)
	IsPost() bool
	IsGet() bool
	IsPut() bool
	IsDel() bool
	IsHead() bool
	IsPatch() bool
	IsOptions() bool
	IsSecure() bool
	IsWebsocket() bool
	IsUpload() bool
	ResolveContentType() string
	WithFormatExtension(bool)
	ResolveFormat() string
	Accept() *Accepts
	Protocol() string
	Site() string
	Scheme() string
	Domain() string
	Host() string
	Proxy() []string
	Referer() string
	Port() int
	RealIP() string
	HasAnyRequest() bool

	MapForm(i interface{}, names ...string) error
	MapData(i interface{}, data map[string][]string, names ...string) error
	SaveUploadedFile(fieldName string, saveAbsPath string, saveFileName ...string) (*multipart.FileHeader, error)
	SaveUploadedFileToWriter(string, io.Writer) (*multipart.FileHeader, error)
	//Multiple file upload
	SaveUploadedFiles(fieldName string, savePath func(*multipart.FileHeader) (string, error)) error
	SaveUploadedFilesToWriter(fieldName string, writer func(*multipart.FileHeader) (io.Writer, error)) error

	//----------------
	// Hook
	//----------------

	AddPreResponseHook(func() error) Context
	SetPreResponseHook(...func() error) Context
}
<|endoftext|>


<commit_before>package chuper

import (
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/PuerkitoBio/fetchbot"
	"github.com/PuerkitoBio/goquery"
)

const (
	DefaultCrawlDelay      = 5 * time.Second
	DefaultCrawlPoliteness = false
	DefaultUserAgent       = fetchbot.DefaultUserAgent
)

var (
	DefaultHTTPClient = http.DefaultClient

	DefaultCache = NewMemoryCache()

	DefaultErrorHandler = fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		fmt.Printf("chuper - %s - error: %s %s - %s\n", time.Now().Format(time.RFC3339), ctx.Cmd.Method(), ctx.Cmd.URL(), err)
	})

	DefaultLogHandlerFunc = func(ctx *fetchbot.Context, res *http.Response, err error) {
		if err == nil {
			fmt.Printf("chuper - %s - info: [%d] %s %s - %s\n", time.Now().Format(time.RFC3339), res.StatusCode, ctx.Cmd.Method(), ctx.Cmd.URL(), res.Header.Get("Content-Type"))
		}
	}
)

type Crawler struct {
	CrawlDelay      time.Duration
	CrawlDuration   time.Duration
	CrawlPoliteness bool
	HTTPClient      fetchbot.Doer
	Cache           Cache
	ErrorHandler    fetchbot.Handler
	UserAgent       string
	BasicAuthUser   string
	BasicAuthPass   string
	LogHandlerFunc  func(ctx *fetchbot.Context, res *http.Response, err error)

	mux *fetchbot.Mux
	f   *fetchbot.Fetcher
	q   *fetchbot.Queue
}

// New returns an initialized Crawler.
func New() *Crawler {
	return &Crawler{
		CrawlDelay:      DefaultCrawlDelay,
		CrawlPoliteness: DefaultCrawlPoliteness,
		HTTPClient:      DefaultHTTPClient,
		Cache:           DefaultCache,
		ErrorHandler:    DefaultErrorHandler,
		UserAgent:       DefaultUserAgent,
		LogHandlerFunc:  DefaultLogHandlerFunc,
		mux:             fetchbot.NewMux(),
	}
}

func (c *Crawler) Start() Enqueuer {
	c.mux.HandleErrors(c.ErrorHandler)
	l := newLogHandler(c.mux, c.LogHandlerFunc)

	f := fetchbot.New(l)
	f.CrawlDelay = c.CrawlDelay
	f.DisablePoliteness = !c.CrawlPoliteness
	f.HttpClient = c.HTTPClient
	f.UserAgent = c.UserAgent

	c.f = f
	c.q = c.f.Start()

	if c.CrawlDuration > 0 {
		go func() {
			t := time.After(c.CrawlDuration)
			<-t
			c.q.Close()
		}()
	}

	return &Queue{c.q}
}

func (c *Crawler) Block() {
	c.q.Block()
}

func (c *Crawler) Finish() {
	c.q.Close()
}

func (c *Crawler) Enqueue(method string, rawURL ...string) error {
	for _, u := range rawURL {
		url, err := url.Parse(u)
		if err != nil {
			return err
		}

		ok := true
		if c.mustCache() {
			ok, _ = c.Cache.SetNX(u, true)
		}

		if ok {
			var cmd fetchbot.Command
			if c.BasicAuthUser != "" && c.BasicAuthPass != "" {
				cmd = &CmdBasicAuth{&fetchbot.Cmd{U: url, M: method}, url, c.BasicAuthUser, c.BasicAuthPass}
			} else {
				cmd = &fetchbot.Cmd{U: url, M: method}
			}
			if err := c.q.Send(cmd); err != nil {
				return err
			}

		}
	}
	return nil
}

func (c *Crawler) EnqueueWithSource(method string, URL string, sourceURL string) (bool, error) {
	ok := true
	if c.mustCache() {
		ok, _ = c.Cache.SetNX(URL, true)
	}
	if ok {
		u, err := url.Parse(URL)
		if err != nil {
			return ok, err
		}
		s, err := url.Parse(sourceURL)
		if err != nil {
			return ok, err
		}

		var cmd fetchbot.Command
		if c.BasicAuthUser != "" && c.BasicAuthPass != "" {
			cmd = &CmdBasicAuth{&fetchbot.Cmd{U: u, M: method}, s, c.BasicAuthUser, c.BasicAuthPass}
		} else {
			cmd = &Cmd{&fetchbot.Cmd{U: u, M: method}, s}
		}
		if err := c.q.Send(cmd); err != nil {
			return ok, err
		}
	}

	return ok, nil
}

type ResponseCriteria struct {
	Method      string
	ContentType string
	Status      int
	MinStatus   int
	MaxStatus   int
	Path        string
	Host        string
}

func (c *Crawler) Match(r *ResponseCriteria) *fetchbot.ResponseMatcher {
	m := c.mux.Response()

	if r.Method != "" {
		m.Method(r.Method)
	}

	if r.ContentType != "" {
		m.ContentType(r.ContentType)
	}

	if r.Status != 0 {
		m.Status(r.Status)
	} else {
		if r.MinStatus != 0 && r.MaxStatus != 0 {
			m.StatusRange(r.MinStatus, r.MaxStatus)
		} else {
			if r.MinStatus != 0 {
				m.Status(r.MinStatus)
			}
			if r.MaxStatus != 0 {
				m.Status(r.MaxStatus)
			}
		}
	}

	if r.Path != "" {
		m.Path(r.Path)
	}

	if r.Host != "" {
		m.Host(r.Host)
	}

	return m
}

func (c *Crawler) Register(rc *ResponseCriteria, procs ...Processor) {
	m := c.Match(rc)
	h := newDocHandler(c.Cache, procs...)
	m.Handler(h)
}

func (c *Crawler) mustCache() bool {
	if c.Cache == nil {
		return false
	}
	return true
}

func newLogHandler(wrapped fetchbot.Handler, f func(ctx *fetchbot.Context, res *http.Response, err error)) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		f(ctx, res, err)
		wrapped.Handle(ctx, res, err)
	})
}

func newDocHandler(cache Cache, procs ...Processor) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		context := &Ctx{ctx, cache}
		doc, err := goquery.NewDocumentFromResponse(res)
		if err != nil {
			fmt.Printf("chuper - %s - error: %s %s - %s\n", time.Now().Format(time.RFC3339), ctx.Cmd.Method(), ctx.Cmd.URL(), err)
			return
		}
		for _, p := range procs {
			ok := p.Process(context, doc)
			if !ok {
				return
			}
		}
	})
}
<commit_msg>Remove BasicAuth fields from Crawler<commit_after>package chuper

import (
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/PuerkitoBio/fetchbot"
	"github.com/PuerkitoBio/goquery"
)

const (
	DefaultCrawlDelay      = 5 * time.Second
	DefaultCrawlPoliteness = false
	DefaultUserAgent       = fetchbot.DefaultUserAgent
)

var (
	DefaultHTTPClient = http.DefaultClient

	DefaultCache = NewMemoryCache()

	DefaultErrorHandler = fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		fmt.Printf("chuper - %s - error: %s %s - %s\n", time.Now().Format(time.RFC3339), ctx.Cmd.Method(), ctx.Cmd.URL(), err)
	})

	DefaultLogHandlerFunc = func(ctx *fetchbot.Context, res *http.Response, err error) {
		if err == nil {
			fmt.Printf("chuper - %s - info: [%d] %s %s - %s\n", time.Now().Format(time.RFC3339), res.StatusCode, ctx.Cmd.Method(), ctx.Cmd.URL(), res.Header.Get("Content-Type"))
		}
	}
)

type Crawler struct {
	CrawlDelay      time.Duration
	CrawlDuration   time.Duration
	CrawlPoliteness bool
	HTTPClient      fetchbot.Doer
	Cache           Cache
	ErrorHandler    fetchbot.Handler
	UserAgent       string
	LogHandlerFunc  func(ctx *fetchbot.Context, res *http.Response, err error)

	mux *fetchbot.Mux
	f   *fetchbot.Fetcher
	q   *fetchbot.Queue
}

// New returns an initialized Crawler.
func New() *Crawler {
	return &Crawler{
		CrawlDelay:      DefaultCrawlDelay,
		CrawlPoliteness: DefaultCrawlPoliteness,
		HTTPClient:      DefaultHTTPClient,
		Cache:           DefaultCache,
		ErrorHandler:    DefaultErrorHandler,
		UserAgent:       DefaultUserAgent,
		LogHandlerFunc:  DefaultLogHandlerFunc,
		mux:             fetchbot.NewMux(),
	}
}

func (c *Crawler) Start() Enqueuer {
	c.mux.HandleErrors(c.ErrorHandler)
	l := newLogHandler(c.mux, c.LogHandlerFunc)

	f := fetchbot.New(l)
	f.CrawlDelay = c.CrawlDelay
	f.DisablePoliteness = !c.CrawlPoliteness
	f.HttpClient = c.HTTPClient
	f.UserAgent = c.UserAgent

	c.f = f
	c.q = c.f.Start()

	if c.CrawlDuration > 0 {
		go func() {
			t := time.After(c.CrawlDuration)
			<-t
			c.q.Close()
		}()
	}

	return &Queue{c.q}
}

func (c *Crawler) Block() {
	c.q.Block()
}

func (c *Crawler) Finish() {
	c.q.Close()
}

func (c *Crawler) Enqueue(method string, rawURL ...string) error {
	for _, u := range rawURL {
		url, err := url.Parse(u)
		if err != nil {
			return err
		}

		ok := true
		if c.mustCache() {
			ok, _ = c.Cache.SetNX(u, true)
		}

		if ok {
			var cmd fetchbot.Command
			if c.BasicAuthUser != "" && c.BasicAuthPass != "" {
				cmd = &CmdBasicAuth{&fetchbot.Cmd{U: url, M: method}, url, c.BasicAuthUser, c.BasicAuthPass}
			} else {
				cmd = &fetchbot.Cmd{U: url, M: method}
			}
			if err := c.q.Send(cmd); err != nil {
				return err
			}

		}
	}
	return nil
}

func (c *Crawler) EnqueueWithSource(method string, URL string, sourceURL string) (bool, error) {
	ok := true
	if c.mustCache() {
		ok, _ = c.Cache.SetNX(URL, true)
	}
	if ok {
		u, err := url.Parse(URL)
		if err != nil {
			return ok, err
		}
		s, err := url.Parse(sourceURL)
		if err != nil {
			return ok, err
		}

		var cmd fetchbot.Command
		if c.BasicAuthUser != "" && c.BasicAuthPass != "" {
			cmd = &CmdBasicAuth{&fetchbot.Cmd{U: u, M: method}, s, c.BasicAuthUser, c.BasicAuthPass}
		} else {
			cmd = &Cmd{&fetchbot.Cmd{U: u, M: method}, s}
		}
		if err := c.q.Send(cmd); err != nil {
			return ok, err
		}
	}

	return ok, nil
}

type ResponseCriteria struct {
	Method      string
	ContentType string
	Status      int
	MinStatus   int
	MaxStatus   int
	Path        string
	Host        string
}

func (c *Crawler) Match(r *ResponseCriteria) *fetchbot.ResponseMatcher {
	m := c.mux.Response()

	if r.Method != "" {
		m.Method(r.Method)
	}

	if r.ContentType != "" {
		m.ContentType(r.ContentType)
	}

	if r.Status != 0 {
		m.Status(r.Status)
	} else {
		if r.MinStatus != 0 && r.MaxStatus != 0 {
			m.StatusRange(r.MinStatus, r.MaxStatus)
		} else {
			if r.MinStatus != 0 {
				m.Status(r.MinStatus)
			}
			if r.MaxStatus != 0 {
				m.Status(r.MaxStatus)
			}
		}
	}

	if r.Path != "" {
		m.Path(r.Path)
	}

	if r.Host != "" {
		m.Host(r.Host)
	}

	return m
}

func (c *Crawler) Register(rc *ResponseCriteria, procs ...Processor) {
	m := c.Match(rc)
	h := newDocHandler(c.Cache, procs...)
	m.Handler(h)
}

func (c *Crawler) mustCache() bool {
	if c.Cache == nil {
		return false
	}
	return true
}

func newLogHandler(wrapped fetchbot.Handler, f func(ctx *fetchbot.Context, res *http.Response, err error)) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		f(ctx, res, err)
		wrapped.Handle(ctx, res, err)
	})
}

func newDocHandler(cache Cache, procs ...Processor) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		context := &Ctx{ctx, cache}
		doc, err := goquery.NewDocumentFromResponse(res)
		if err != nil {
			fmt.Printf("chuper - %s - error: %s %s - %s\n", time.Now().Format(time.RFC3339), ctx.Cmd.Method(), ctx.Cmd.URL(), err)
			return
		}
		for _, p := range procs {
			ok := p.Process(context, doc)
			if !ok {
				return
			}
		}
	})
}
<|endoftext|>


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

package webseclab

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"strings"
)

func DoLabTestStandard(w http.ResponseWriter, r *http.Request) *LabResp {
	// split r.URL.RawQuery into a param map
	rawParams := make(map[string][]string)
	ParseRawQuery(rawParams, r.URL.RawQuery)
	input := &InData{}
	inputRaw, ok := rawParams["in"]
	if ok && len(inputRaw) > 0 {
		input.InRaw = inputRaw[0]
		unesc, err := url.QueryUnescape(input.InRaw)
		if err != nil {
			fmt.Printf("ERROR in url.QueryUnescape on %s\n", input.InRaw)
		}
		input.In = unesc
	}
	err := DoTemplate(w, r.URL.Path, input)
	if err != nil {
		httpcode := http.StatusInternalServerError
		log.Printf("Error returned from DoTemplate: %s\n", err)
		if _, ok := err.(ErrNotFound); ok {
			httpcode = http.StatusNotFound
		}
		return &LabResp{Err: err, Code: httpcode}
	}
	return &LabResp{Err: nil, Code: http.StatusOK}
}

type ErrNotFound struct {
	text string
}

func NewErrNotFound(text string) (e ErrNotFound) {
	e = ErrNotFound{text: text}
	return
}

func (e ErrNotFound) Error() string {
	return e.text
}

// DoTemplate opens the template file and processes the template with the passed input
// if the URL Path ends with ".ok", it uses an HTML context-escaped template
// (html/template - safe version), otherwise - text/template (exploitable)
func DoTemplate(w io.Writer, path string, input *InData) (err error) {
	if input == nil {
		return errors.New("ERROR Internal - Nil passed to DoTemplate as InData")
	}
	if path == "" {
		return errors.New("ERROR Internal - empty path passed to DoTemplate")
	}
	if len(path) == 1 {
		return errors.New("ERROR Internal - too short (len = 1) path passed to DoTemplate: " + path)
	}
	hasOk := strings.HasSuffix(path, ".ok")
	var fpath string // filepath
	if hasOk {
		fpath = path[1 : len(path)-3]
	} else {
		fpath = path[1:]
	}
	if hasOk {
		// try text template with the literal ".ok" first for special cases (templates with broken HTML)
		_, ok := LookupTextTemplate(fpath + ".ok")
		// text/template - no context-sensitive escaping
		if ok {
			return doTextTemplate(w, fpath+".ok", input)
		}
		return doHtmlTemplate(w, fpath, input)
	}
	return doTextTemplate(w, fpath, input)
}

func doHtmlTemplate(w io.Writer, fpath string, input *InData) (err error) {
	// html/template - context-sensitive escaping
	tmpl, ok := LookupHtmlTemplate(fpath)
	if ok == false {
		return errors.New("Error in DoTemplate - html template " + fpath + " not found.")
	}
	parts := strings.Split(fpath, "/")
	err = tmpl.ExecuteTemplate(w, parts[len(parts)-1], *input)
	if err != nil {
		log.Printf("Error in DoTemplate (html) - tmpl.Execute: %s\n", err)
		return err
	}
	return nil
}

func doTextTemplate(w io.Writer, fpath string, input *InData) (err error) {
	tmpl, ok := LookupTextTemplate(fpath)
	if ok == false {
		err := NewErrNotFound("Error in DoTemplate - text template " + fpath + " not found.")
		return err
	}
	parts := strings.Split(fpath, "/")
	err = tmpl.ExecuteTemplate(w, parts[len(parts)-1], *input)
	if err != nil {
		log.Printf("Error in DoTemplate (text) - tmpl.Execute's err: %s\n", err)
		return err
	}
	return nil
}

// HandleFilterBased factors out the common handling for the "typical" -
// filter-based tests (listed in the map in custom.go, filters are in filters.go)
func HandleFilterBased(w http.ResponseWriter, r *http.Request, filters []filter) *LabResp {
	var input InData
	rawParams := make(map[string][]string)
	ParseRawQuery(rawParams, r.URL.RawQuery)
	inputRaw, ok := rawParams["in"]
	if ok && len(inputRaw) > 0 {
		input.InRaw = inputRaw[0]
		unesc, err := url.QueryUnescape(input.InRaw)
		if err != nil {
			log.Printf("Error in %s: %s\n", r.URL.Path, err)
			return &LabResp{Err: nil, Code: http.StatusInternalServerError}
		}
		tr := NewTransformer(filters...)
		input.In = tr.Transform(unesc)
	}
	err := DoTemplate(w, r.URL.Path, &input)
	if err != nil {
		log.Printf("Error in DoTemplate: %s\n", err)
		return &LabResp{Err: nil, Code: http.StatusInternalServerError}
	}
	return &LabResp{Err: nil, Code: http.StatusOK}
}
<commit_msg>set explicitly Content-type to text/html to prevent content sniffing<commit_after>// Copyright 2015, Yahoo Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package webseclab

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"
)

func DoLabTestStandard(w http.ResponseWriter, r *http.Request) *LabResp {
	// split r.URL.RawQuery into a param map
	rawParams := make(map[string][]string)
	ParseRawQuery(rawParams, r.URL.RawQuery)
	input := &InData{}
	inputRaw, ok := rawParams["in"]
	if ok && len(inputRaw) > 0 {
		input.InRaw = inputRaw[0]
		unesc, err := url.QueryUnescape(input.InRaw)
		if err != nil {
			fmt.Printf("ERROR in url.QueryUnescape on %s\n", input.InRaw)
		}
		input.In = unesc
	}
	err := DoTemplate(w, r.URL.Path, input)
	if err != nil {
		httpcode := http.StatusInternalServerError
		log.Printf("Error returned from DoTemplate: %s\n", err)
		if _, ok := err.(ErrNotFound); ok {
			httpcode = http.StatusNotFound
		}
		return &LabResp{Err: err, Code: httpcode}
	}
	return &LabResp{Err: nil, Code: http.StatusOK}
}

type ErrNotFound struct {
	text string
}

func NewErrNotFound(text string) (e ErrNotFound) {
	e = ErrNotFound{text: text}
	return
}

func (e ErrNotFound) Error() string {
	return e.text
}

// DoTemplate opens the template file and processes the template with the passed input
// if the URL Path ends with ".ok", it uses an HTML context-escaped template
// (html/template - safe version), otherwise - text/template (exploitable)
func DoTemplate(w http.ResponseWriter, path string, input *InData) (err error) {
	if input == nil {
		return errors.New("ERROR Internal - Nil passed to DoTemplate as InData")
	}
	if path == "" {
		return errors.New("ERROR Internal - empty path passed to DoTemplate")
	}
	if len(path) == 1 {
		return errors.New("ERROR Internal - too short (len = 1) path passed to DoTemplate: " + path)
	}
	hasOk := strings.HasSuffix(path, ".ok")
	var fpath string // filepath
	if hasOk {
		fpath = path[1 : len(path)-3]
	} else {
		fpath = path[1:]
	}
	if hasOk {
		// try text template with the literal ".ok" first for special cases (templates with broken HTML)
		_, ok := LookupTextTemplate(fpath + ".ok")
		// text/template - no context-sensitive escaping
		if ok {
			return doTextTemplate(w, fpath+".ok", input)
		}
		return doHtmlTemplate(w, fpath, input)
	}
	return doTextTemplate(w, fpath, input)
}

func doHtmlTemplate(w http.ResponseWriter, fpath string, input *InData) (err error) {
	// html/template - context-sensitive escaping
	tmpl, ok := LookupHtmlTemplate(fpath)
	if ok == false {
		return errors.New("Error in DoTemplate - html template " + fpath + " not found.")
	}
	w.Header().Set("Content-type", "text/html; charset=utf-8")
	parts := strings.Split(fpath, "/")
	err = tmpl.ExecuteTemplate(w, parts[len(parts)-1], *input)
	if err != nil {
		log.Printf("Error in DoTemplate (html) - tmpl.Execute: %s\n", err)
		return err
	}
	return nil
}

func doTextTemplate(w http.ResponseWriter, fpath string, input *InData) (err error) {
	tmpl, ok := LookupTextTemplate(fpath)
	if ok == false {
		err := NewErrNotFound("Error in DoTemplate - text template " + fpath + " not found.")
		return err
	}
	w.Header().Set("Content-type", "text/html; charset=utf-8")
	parts := strings.Split(fpath, "/")
	err = tmpl.ExecuteTemplate(w, parts[len(parts)-1], *input)
	if err != nil {
		log.Printf("Error in DoTemplate (text) - tmpl.Execute's err: %s\n", err)
		return err
	}
	return nil
}

// HandleFilterBased factors out the common handling for the "typical" -
// filter-based tests (listed in the map in custom.go, filters are in filters.go)
func HandleFilterBased(w http.ResponseWriter, r *http.Request, filters []filter) *LabResp {
	var input InData
	rawParams := make(map[string][]string)
	ParseRawQuery(rawParams, r.URL.RawQuery)
	inputRaw, ok := rawParams["in"]
	if ok && len(inputRaw) > 0 {
		input.InRaw = inputRaw[0]
		unesc, err := url.QueryUnescape(input.InRaw)
		if err != nil {
			log.Printf("Error in %s: %s\n", r.URL.Path, err)
			return &LabResp{Err: nil, Code: http.StatusInternalServerError}
		}
		tr := NewTransformer(filters...)
		input.In = tr.Transform(unesc)
	}
	err := DoTemplate(w, r.URL.Path, &input)
	if err != nil {
		log.Printf("Error in DoTemplate: %s\n", err)
		return &LabResp{Err: nil, Code: http.StatusInternalServerError}
	}
	return &LabResp{Err: nil, Code: http.StatusOK}
}
<|endoftext|>


<commit_before>package buffalo

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/gobuffalo/buffalo/render"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
)

// assert that DefaultContext is implementing Context
var _ Context = &DefaultContext{}
var _ context.Context = &DefaultContext{}

// DefaultContext is, as its name implies, a default
// implementation of the Context interface.
type DefaultContext struct {
	context.Context
	response    http.ResponseWriter
	request     *http.Request
	params      url.Values
	logger      Logger
	session     *Session
	contentType string
	data        map[string]interface{}
	flash       *Flash
}

// Response returns the original Response for the request.
func (d *DefaultContext) Response() http.ResponseWriter {
	return d.response
}

// Request returns the original Request.
func (d *DefaultContext) Request() *http.Request {
	return d.request
}

// Params returns all of the parameters for the request,
// including both named params and query string parameters.
// These parameters are automatically available in templates
// as "{{.params}}".
func (d *DefaultContext) Params() ParamValues {
	return d.params
}

// Logger returns the Logger for this context.
func (d *DefaultContext) Logger() Logger {
	return d.logger
}

// Param returns a param, either named or query string,
// based on the key.
func (d *DefaultContext) Param(key string) string {
	return d.Params().Get(key)
}

// ParamInt tries to convert the requested parameter to
// an int. It will return an error if there is a problem.
func (d *DefaultContext) ParamInt(key string) (int, error) {
	k := d.Params().Get(key)
	i, err := strconv.Atoi(k)
	return i, errors.WithMessage(err, fmt.Sprintf("could not convert %s to an int", k))
}

// Set a value onto the Context. Any value set onto the Context
// will be automatically available in templates.
func (d *DefaultContext) Set(key string, value interface{}) {
	d.data[key] = value
}

// Get is deprecated. Please use Value instead.
func (d *DefaultContext) Get(key string) interface{} {
	warningMsg := "Context#Get is deprecated. Please use Context#Value instead."

	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("Context#Get is deprecated. Please use Context#Value instead. Called from %s:%d", file, no)
	}

	d.Logger().Warn(warningMsg)
	return d.Value(key)
}

// Value that has previously stored on the context.
func (d *DefaultContext) Value(key interface{}) interface{} {
	if k, ok := key.(string); ok {
		if v, ok := d.data[k]; ok {
			return v
		}
	}
	return d.Context.Value(key)
}

// Session for the associated Request.
func (d *DefaultContext) Session() *Session {
	return d.session
}

// Flash messages for the associated Request.
func (d *DefaultContext) Flash() *Flash {
	return d.flash
}

// Render a status code and render.Renderer to the associated Response.
// The request parameters will be made available to the render.Renderer
// "{{.params}}". Any values set onto the Context will also automatically
// be made available to the render.Renderer. To render "no content" pass
// in a nil render.Renderer.
func (d *DefaultContext) Render(status int, rr render.Renderer) error {
	now := time.Now()
	defer func() {
		d.LogField("render", time.Now().Sub(now))
	}()
	if rr != nil {
		data := d.data
		pp := map[string]string{}
		for k, v := range d.params {
			pp[k] = v[0]
		}
		data["params"] = pp
		data["flash"] = d.Flash().data
		bb := &bytes.Buffer{}

		err := rr.Render(bb, data)
		if err != nil {
			return HTTPError{Status: 500, Cause: errors.WithStack(err)}
		}

		if d.Session() != nil {
			d.Flash().Clear()
			d.Flash().persist(d.Session())
		}

		d.Response().Header().Set("Content-Type", rr.ContentType())
		d.Response().WriteHeader(status)
		_, err = io.Copy(d.Response(), bb)
		if err != nil {
			return HTTPError{Status: 500, Cause: errors.WithStack(err)}
		}

		return nil
	}
	d.Response().WriteHeader(status)
	return nil
}

// Bind the interface to the request.Body. The type of binding
// is dependent on the "Content-Type" for the request. If the type
// is "application/json" it will use "json.NewDecoder". If the type
// is "application/xml" it will use "xml.NewDecoder". The default
// binder is "http://www.gorillatoolkit.org/pkg/schema".
func (d *DefaultContext) Bind(value interface{}) error {
	ct := strings.ToLower(d.Request().Header.Get("Content-Type"))
	if ct != "" {
		cts := strings.Split(ct, ";")
		c := cts[0]
		if b, ok := binders[strings.TrimSpace(c)]; ok {
			return b(d.Request(), value)
		}
		return errors.Errorf("could not find a binder for %s", c)
	}
	return errors.New("blank content type")
}

// LogField adds the key/value pair onto the Logger to be printed out
// as part of the request logging. This allows you to easily add things
// like metrics (think DB times) to your request.
func (d *DefaultContext) LogField(key string, value interface{}) {
	d.logger = d.logger.WithField(key, value)
}

// LogFields adds the key/value pairs onto the Logger to be printed out
// as part of the request logging. This allows you to easily add things
// like metrics (think DB times) to your request.
func (d *DefaultContext) LogFields(values map[string]interface{}) {
	d.logger = d.logger.WithFields(values)
}

func (d *DefaultContext) Error(status int, err error) error {
	return HTTPError{Status: status, Cause: errors.WithStack(err)}
}

// Websocket returns an upgraded github.com/gorilla/websocket.Conn
// that can then be used to work with websockets easily.
func (d *DefaultContext) Websocket() (*websocket.Conn, error) {
	return defaultUpgrader.Upgrade(d.Response(), d.Request(), nil)
}

// Redirect a request with the given status to the given URL.
func (d *DefaultContext) Redirect(status int, url string, args ...interface{}) error {
	d.Flash().persist(d.Session())

	http.Redirect(d.Response(), d.Request(), fmt.Sprintf(url, args...), status)
	return nil
}

// Data contains all the values set through Get/Set.
func (d *DefaultContext) Data() map[string]interface{} {
	return d.data
}

var defaultUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}
<commit_msg>Update docs to reflect current state of template context. closes #312<commit_after>package buffalo

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/gobuffalo/buffalo/render"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
)

// assert that DefaultContext is implementing Context
var _ Context = &DefaultContext{}
var _ context.Context = &DefaultContext{}

// DefaultContext is, as its name implies, a default
// implementation of the Context interface.
type DefaultContext struct {
	context.Context
	response    http.ResponseWriter
	request     *http.Request
	params      url.Values
	logger      Logger
	session     *Session
	contentType string
	data        map[string]interface{}
	flash       *Flash
}

// Response returns the original Response for the request.
func (d *DefaultContext) Response() http.ResponseWriter {
	return d.response
}

// Request returns the original Request.
func (d *DefaultContext) Request() *http.Request {
	return d.request
}

// Params returns all of the parameters for the request,
// including both named params and query string parameters.
func (d *DefaultContext) Params() ParamValues {
	return d.params
}

// Logger returns the Logger for this context.
func (d *DefaultContext) Logger() Logger {
	return d.logger
}

// Param returns a param, either named or query string,
// based on the key.
func (d *DefaultContext) Param(key string) string {
	return d.Params().Get(key)
}

// ParamInt tries to convert the requested parameter to
// an int. It will return an error if there is a problem.
func (d *DefaultContext) ParamInt(key string) (int, error) {
	k := d.Params().Get(key)
	i, err := strconv.Atoi(k)
	return i, errors.WithMessage(err, fmt.Sprintf("could not convert %s to an int", k))
}

// Set a value onto the Context. Any value set onto the Context
// will be automatically available in templates.
func (d *DefaultContext) Set(key string, value interface{}) {
	d.data[key] = value
}

// Get is deprecated. Please use Value instead.
func (d *DefaultContext) Get(key string) interface{} {
	warningMsg := "Context#Get is deprecated. Please use Context#Value instead."

	_, file, no, ok := runtime.Caller(1)
	if ok {
		warningMsg = fmt.Sprintf("Context#Get is deprecated. Please use Context#Value instead. Called from %s:%d", file, no)
	}

	d.Logger().Warn(warningMsg)
	return d.Value(key)
}

// Value that has previously stored on the context.
func (d *DefaultContext) Value(key interface{}) interface{} {
	if k, ok := key.(string); ok {
		if v, ok := d.data[k]; ok {
			return v
		}
	}
	return d.Context.Value(key)
}

// Session for the associated Request.
func (d *DefaultContext) Session() *Session {
	return d.session
}

// Flash messages for the associated Request.
func (d *DefaultContext) Flash() *Flash {
	return d.flash
}

// Render a status code and render.Renderer to the associated Response.
// The request parameters will be made available to the render.Renderer
// "{{.params}}". Any values set onto the Context will also automatically
// be made available to the render.Renderer. To render "no content" pass
// in a nil render.Renderer.
func (d *DefaultContext) Render(status int, rr render.Renderer) error {
	now := time.Now()
	defer func() {
		d.LogField("render", time.Now().Sub(now))
	}()
	if rr != nil {
		data := d.data
		pp := map[string]string{}
		for k, v := range d.params {
			pp[k] = v[0]
		}
		data["params"] = pp
		data["flash"] = d.Flash().data
		bb := &bytes.Buffer{}

		err := rr.Render(bb, data)
		if err != nil {
			return HTTPError{Status: 500, Cause: errors.WithStack(err)}
		}

		if d.Session() != nil {
			d.Flash().Clear()
			d.Flash().persist(d.Session())
		}

		d.Response().Header().Set("Content-Type", rr.ContentType())
		d.Response().WriteHeader(status)
		_, err = io.Copy(d.Response(), bb)
		if err != nil {
			return HTTPError{Status: 500, Cause: errors.WithStack(err)}
		}

		return nil
	}
	d.Response().WriteHeader(status)
	return nil
}

// Bind the interface to the request.Body. The type of binding
// is dependent on the "Content-Type" for the request. If the type
// is "application/json" it will use "json.NewDecoder". If the type
// is "application/xml" it will use "xml.NewDecoder". The default
// binder is "http://www.gorillatoolkit.org/pkg/schema".
func (d *DefaultContext) Bind(value interface{}) error {
	ct := strings.ToLower(d.Request().Header.Get("Content-Type"))
	if ct != "" {
		cts := strings.Split(ct, ";")
		c := cts[0]
		if b, ok := binders[strings.TrimSpace(c)]; ok {
			return b(d.Request(), value)
		}
		return errors.Errorf("could not find a binder for %s", c)
	}
	return errors.New("blank content type")
}

// LogField adds the key/value pair onto the Logger to be printed out
// as part of the request logging. This allows you to easily add things
// like metrics (think DB times) to your request.
func (d *DefaultContext) LogField(key string, value interface{}) {
	d.logger = d.logger.WithField(key, value)
}

// LogFields adds the key/value pairs onto the Logger to be printed out
// as part of the request logging. This allows you to easily add things
// like metrics (think DB times) to your request.
func (d *DefaultContext) LogFields(values map[string]interface{}) {
	d.logger = d.logger.WithFields(values)
}

func (d *DefaultContext) Error(status int, err error) error {
	return HTTPError{Status: status, Cause: errors.WithStack(err)}
}

// Websocket returns an upgraded github.com/gorilla/websocket.Conn
// that can then be used to work with websockets easily.
func (d *DefaultContext) Websocket() (*websocket.Conn, error) {
	return defaultUpgrader.Upgrade(d.Response(), d.Request(), nil)
}

// Redirect a request with the given status to the given URL.
func (d *DefaultContext) Redirect(status int, url string, args ...interface{}) error {
	d.Flash().persist(d.Session())

	http.Redirect(d.Response(), d.Request(), fmt.Sprintf(url, args...), status)
	return nil
}

// Data contains all the values set through Get/Set.
func (d *DefaultContext) Data() map[string]interface{} {
	return d.data
}

var defaultUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}
<|endoftext|>


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

package main

import (
	"context"
	"flag"
	"math/rand"
	"path/filepath"
	"sync"
	"time"

	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/app/crash"
	"github.com/google/gapid/core/event/task"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/gapis/service"
)

type stresstestVerb struct{ StressTestFlags }

func init() {
	app.AddVerb(&app.Verb{
		Name:      "stress-test",
		ShortHelp: "Performs evil things on GAPIS to try to break it",
		Action:    &stresstestVerb{},
	})
}

func (verb *stresstestVerb) Run(ctx context.Context, flags flag.FlagSet) error {
	if flags.NArg() != 1 {
		app.Usage(ctx, "Exactly one gfx trace file expected, got %d", flags.NArg())
		return nil
	}

	client, err := getGapis(ctx, verb.Gapis, verb.Gapir)
	if err != nil {
		return log.Err(ctx, err, "Failed to connect to the GAPIS server")
	}
	defer client.Close()

	filepath, err := filepath.Abs(flags.Arg(0))
	ctx = log.V{"filepath": filepath}.Bind(ctx)
	if err != nil {
		return log.Err(ctx, err, "Could not find capture file")
	}

	c, err := client.LoadCapture(ctx, filepath)
	if err != nil {
		return log.Err(ctx, err, "Failed to load the capture file")
	}

	boxedCapture, err := client.Get(ctx, c.Path())
	if err != nil {
		return log.Err(ctx, err, "Failed to load the capture")
	}
	count := int(boxedCapture.(*service.Capture).NumCommands)

	wg := sync.WaitGroup{}

	for i := 0; i < 1000; i++ {
		at := uint64(rand.Intn(count - 1))
		duration := time.Second + time.Duration(rand.Intn(int(time.Second*10)))
		wg.Add(1)

		crash.Go(func() {
			defer wg.Done()
			ctx, _ := task.WithTimeout(ctx, duration)
			boxedTree, err := client.Get(ctx, c.Command(at).StateAfter().Tree().Path())
			if err != nil {
				return
			}

			tree := boxedTree.(*service.StateTree)

			if _, err := client.Get(ctx, tree.Root.Path()); err != nil {
				return
			}
		})
	}

	wg.Wait()

	return nil
}
<commit_msg>cmd/gapit: Make stresstest trigger #1123<commit_after>// Copyright (C) 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"flag"
	"math/rand"
	"path/filepath"
	"sync"
	"time"

	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/app/crash"
	"github.com/google/gapid/core/event/task"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/gapis/api"
	"github.com/google/gapid/gapis/service"
)

type stresstestVerb struct{ StressTestFlags }

func init() {
	app.AddVerb(&app.Verb{
		Name:      "stress-test",
		ShortHelp: "Performs evil things on GAPIS to try to break it",
		Action:    &stresstestVerb{},
	})
}

func (verb *stresstestVerb) Run(ctx context.Context, flags flag.FlagSet) error {
	if flags.NArg() != 1 {
		app.Usage(ctx, "Exactly one gfx trace file expected, got %d", flags.NArg())
		return nil
	}

	client, err := getGapis(ctx, verb.Gapis, verb.Gapir)
	if err != nil {
		return log.Err(ctx, err, "Failed to connect to the GAPIS server")
	}
	defer client.Close()

	filepath, err := filepath.Abs(flags.Arg(0))
	ctx = log.V{"filepath": filepath}.Bind(ctx)
	if err != nil {
		return log.Err(ctx, err, "Could not find capture file")
	}

	c, err := client.LoadCapture(ctx, filepath)
	if err != nil {
		return log.Err(ctx, err, "Failed to load the capture file")
	}

	boxedCapture, err := client.Get(ctx, c.Path())
	if err != nil {
		return log.Err(ctx, err, "Failed to load the capture")
	}
	count := int(boxedCapture.(*service.Capture).NumCommands)

	wg := sync.WaitGroup{}

	for l := 0; l < 10; l++ {
		for i := 0; i < 10000; i++ {
			at := uint64(rand.Intn(count - 1))
			duration := time.Second + time.Duration(rand.Intn(int(time.Second*10)))
			wg.Add(1)

			const (
				getStateAfter = iota
				getMesh
				getCount
			)
			method := rand.Intn(getCount)

			crash.Go(func() {
				defer wg.Done()
				ctx, _ := task.WithTimeout(ctx, duration)

				switch method {
				case getStateAfter:
					boxedTree, err := client.Get(ctx, c.Command(at).StateAfter().Tree().Path())
					if err == nil {
						tree := boxedTree.(*service.StateTree)
						client.Get(ctx, tree.Root.Path())
					}

				case getMesh:
					boxedMesh, err := client.Get(ctx, c.Command(at).Mesh(true).Path())
					if err == nil {
						mesh := boxedMesh.(*api.Mesh)
						_ = mesh
					}
				}
			})
		}
		wg.Wait()
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2021 Timo Savola. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import m "gate.computer/gate/cmd/gate-daemon/main"

func main() { m.Main() }
<commit_msg>cmd/gate-daemon: import sqlite driver<commit_after>// Copyright (c) 2021 Timo Savola. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	m "gate.computer/gate/cmd/gate-daemon/main"
	_ "modernc.org/sqlite"
)

func main() { m.Main() }
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"strings"

	"github.com/funkygao/gocli"
	"github.com/tylertreat/comcast/throttler"
)

type Chaos struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Chaos) Run(args []string) (exitCode int) {
	var (
		ipv6, dryrun        bool
		device, mode, ports string
		loss                float64
		latency             int
	)
	cmdFlags := flag.NewFlagSet("chaos", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&device, "i", "", "")
	cmdFlags.StringVar(&mode, "mode", "start", "")
	cmdFlags.IntVar(&latency, "latency", -1, "")
	cmdFlags.StringVar(&ports, "ports", "", "")
	cmdFlags.Float64Var(&loss, "loss", 0, "")
	cmdFlags.BoolVar(&dryrun, "dryrun", false, "")
	cmdFlags.BoolVar(&ipv6, "ipv6", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	log.SetOutput(os.Stdout)
	var stop = false
	switch mode {
	case "start":
		stop = false
	case "stop":
		stop = true
	default:
		this.Ui.Output(this.Help())
		return 2
	}
	throttler.Run(&throttler.Config{
		Device:      device,
		Stop:        stop,
		Latency:     latency,
		DryRun:      dryrun,
		TargetPorts: strings.Split(ports, ","),
		PacketLoss:  loss,
	})

	return
}

func (*Chaos) Synopsis() string {
	return "Chaos monkey which simulates common network problems"
}

func (this *Chaos) Help() string {
	var underlying string
	switch runtime.GOOS {
	case "darwin", "freebsd":
		underlying = "ipfw"

	case "linux":
		underlying = "tc"
	}

	help := fmt.Sprintf(`
Usage: %s chaos [options]

    %s

    A wrapper of %s

Options:

    -mode <start|stop>

    -i eth0

    -latency n in ms

    -loss 0.01
     1%% packet loss.

    -ports comma seperated port list

    -dryrun

`, this.Cmd, this.Synopsis(), underlying)
	return strings.TrimSpace(help)
}
<commit_msg>mode is required<commit_after>package command

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"strings"

	"github.com/funkygao/gocli"
	"github.com/tylertreat/comcast/throttler"
)

type Chaos struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Chaos) Run(args []string) (exitCode int) {
	var (
		ipv6, dryrun        bool
		device, mode, ports string
		loss                float64
		latency             int
	)
	cmdFlags := flag.NewFlagSet("chaos", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&device, "i", "", "")
	cmdFlags.StringVar(&mode, "mode", "", "")
	cmdFlags.IntVar(&latency, "latency", -1, "")
	cmdFlags.StringVar(&ports, "ports", "", "")
	cmdFlags.Float64Var(&loss, "loss", 0, "")
	cmdFlags.BoolVar(&dryrun, "dryrun", false, "")
	cmdFlags.BoolVar(&ipv6, "ipv6", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	log.SetOutput(os.Stdout)
	var stop = false
	switch mode {
	case "start":
		stop = false
	case "stop":
		stop = true
	default:
		this.Ui.Output(this.Help())
		return 2
	}
	throttler.Run(&throttler.Config{
		Device:      device,
		Stop:        stop,
		Latency:     latency,
		DryRun:      dryrun,
		TargetPorts: strings.Split(ports, ","),
		PacketLoss:  loss,
	})

	return
}

func (*Chaos) Synopsis() string {
	return "Chaos monkey which simulates common network problems"
}

func (this *Chaos) Help() string {
	var underlying string
	switch runtime.GOOS {
	case "darwin", "freebsd":
		underlying = "ipfw"

	case "linux":
		underlying = "tc"
	}

	help := fmt.Sprintf(`
Usage: %s chaos [options]

    %s

    A wrapper of %s

Options:

    -mode <start|stop>

    -i eth0

    -latency n in ms

    -loss 0.01
     1%% packet loss.

    -ports comma seperated port list

    -dryrun

`, this.Cmd, this.Synopsis(), underlying)
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"fmt"
	"io"
	"os"

	"github.com/spf13/cobra"

	"k8s.io/helm/cmd/helm/helmpath"
	"k8s.io/helm/pkg/repo"
)

type repoRemoveCmd struct {
	out  io.Writer
	name string
	home helmpath.Home
}

func newRepoRemoveCmd(out io.Writer) *cobra.Command {
	remove := &repoRemoveCmd{
		out: out,
	}

	cmd := &cobra.Command{
		Use:     "remove [flags] [NAME]",
		Aliases: []string{"rm"},
		Short:   "remove a chart repository",
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := checkArgsLength(len(args), "name of chart repository"); err != nil {
				return err
			}
			remove.name = args[0]
			remove.home = helmpath.Home(homePath())

			return remove.run()
		},
	}

	return cmd
}

func (r *repoRemoveCmd) run() error {
	return removeRepoLine(r.out, r.name, r.home)
}

func removeRepoLine(out io.Writer, name string, home helmpath.Home) error {
	repoFile := home.RepositoryFile()
	r, err := repo.LoadRepositoriesFile(repoFile)
	if err != nil {
		return err
	}

	if !r.Remove(name) {
		return fmt.Errorf("no repo named %q found", name)
	}
	if err := r.WriteFile(repoFile, 0644); err != nil {
		return err
	}

	if err := removeRepoCache(name, home); err != nil {
		return err
	}

	fmt.Fprintf(out, "%q has been removed from your repositories", name)

	return nil
}

func removeRepoCache(name string, home helmpath.Home) error {
	if _, err := os.Stat(home.CacheIndex(name)); err == nil {
		err = os.Remove(home.CacheIndex(name))
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>ref(helm): add new line after repo remove msg<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"fmt"
	"io"
	"os"

	"github.com/spf13/cobra"

	"k8s.io/helm/cmd/helm/helmpath"
	"k8s.io/helm/pkg/repo"
)

type repoRemoveCmd struct {
	out  io.Writer
	name string
	home helmpath.Home
}

func newRepoRemoveCmd(out io.Writer) *cobra.Command {
	remove := &repoRemoveCmd{
		out: out,
	}

	cmd := &cobra.Command{
		Use:     "remove [flags] [NAME]",
		Aliases: []string{"rm"},
		Short:   "remove a chart repository",
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := checkArgsLength(len(args), "name of chart repository"); err != nil {
				return err
			}
			remove.name = args[0]
			remove.home = helmpath.Home(homePath())

			return remove.run()
		},
	}

	return cmd
}

func (r *repoRemoveCmd) run() error {
	return removeRepoLine(r.out, r.name, r.home)
}

func removeRepoLine(out io.Writer, name string, home helmpath.Home) error {
	repoFile := home.RepositoryFile()
	r, err := repo.LoadRepositoriesFile(repoFile)
	if err != nil {
		return err
	}

	if !r.Remove(name) {
		return fmt.Errorf("no repo named %q found", name)
	}
	if err := r.WriteFile(repoFile, 0644); err != nil {
		return err
	}

	if err := removeRepoCache(name, home); err != nil {
		return err
	}

	fmt.Fprintf(out, "%q has been removed from your repositories\n", name)

	return nil
}

func removeRepoCache(name string, home helmpath.Home) error {
	if _, err := os.Stat(home.CacheIndex(name)); err == nil {
		err = os.Remove(home.CacheIndex(name))
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
   Copyright (c) 2016 VMware, Inc. All Rights Reserved.
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package replication

import (
	"net"
)

func retry(err error) bool {
	if err == nil {
		return false
	}
	return isTemporary(err)
}

func isTemporary(err error) bool {
	if netErr, ok := err.(net.Error); ok {
		return netErr.Temporary()
	}
	return false
}
<commit_msg>retry the job which failed because the network error<commit_after>/*
   Copyright (c) 2016 VMware, Inc. All Rights Reserved.
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package replication

import (
	"net"
	"syscall"
)

const (
	connectionRefusedErr syscall.Errno = 0x6f
)

func retry(err error) bool {
	if err == nil {
		return false
	}
	return isNetworkErr(err)
}

func isTemporary(err error) bool {
	if netErr, ok := err.(net.Error); ok {
		return netErr.Temporary()
	}
	return false
}

func isNetworkErr(err error) bool {
	_, ok := err.(net.Error)
	return ok
}
<|endoftext|>


<commit_before>package executor

import (
	"log"
	"os"
	"strings"
	"testing"

	"github.com/hashicorp/nomad/client/testutil"
)

func TestExecScriptCheckWithIsolation(t *testing.T) {
	testutil.ExecCompatible(t)

	execCmd := ExecCommand{Cmd: "/bin/echo", Args: []string{"hello world"}}
	ctx, allocDir := testExecutorContextWithChroot(t)
	defer allocDir.Destroy()

	execCmd.FSIsolation = true
	execCmd.ResourceLimits = true
	execCmd.User = dstructs.DefaultUnpriviledgedUser

	executor := NewExecutor(log.New(os.Stdout, "", log.LstdFlags))

	if err := executor.SetContext(ctx); err != nil {
		t.Fatalf("Unexpected error")
	}

	_, err := executor.LaunchCmd(&execCmd)
	if err != nil {
		t.Fatalf("error in launching command: %v", err)
	}

	check := &ExecScriptCheck{
		id:          "foo",
		cmd:         "/bin/echo",
		args:        []string{"hello", "world"},
		taskDir:     ctx.TaskDir,
		FSIsolation: true,
	}

	res := check.Run()
	expectedOutput := "hello world"
	expectedExitCode := 0
	if res.Err != nil {
		t.Fatalf("err: %v", res.Err)
	}
	if strings.TrimSpace(res.Output) != expectedOutput {
		t.Fatalf("output expected: %v, actual: %v", expectedOutput, res.Output)
	}

	if res.ExitCode != expectedExitCode {
		t.Fatalf("exitcode expected: %v, actual: %v", expectedExitCode, res.ExitCode)
	}
}
<commit_msg>Fix import<commit_after>package executor

import (
	"log"
	"os"
	"strings"
	"testing"

	dstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/testutil"
)

func TestExecScriptCheckWithIsolation(t *testing.T) {
	testutil.ExecCompatible(t)

	execCmd := ExecCommand{Cmd: "/bin/echo", Args: []string{"hello world"}}
	ctx, allocDir := testExecutorContextWithChroot(t)
	defer allocDir.Destroy()

	execCmd.FSIsolation = true
	execCmd.ResourceLimits = true
	execCmd.User = dstructs.DefaultUnpriviledgedUser

	executor := NewExecutor(log.New(os.Stdout, "", log.LstdFlags))

	if err := executor.SetContext(ctx); err != nil {
		t.Fatalf("Unexpected error")
	}

	_, err := executor.LaunchCmd(&execCmd)
	if err != nil {
		t.Fatalf("error in launching command: %v", err)
	}

	check := &ExecScriptCheck{
		id:          "foo",
		cmd:         "/bin/echo",
		args:        []string{"hello", "world"},
		taskDir:     ctx.TaskDir,
		FSIsolation: true,
	}

	res := check.Run()
	expectedOutput := "hello world"
	expectedExitCode := 0
	if res.Err != nil {
		t.Fatalf("err: %v", res.Err)
	}
	if strings.TrimSpace(res.Output) != expectedOutput {
		t.Fatalf("output expected: %v, actual: %v", expectedOutput, res.Output)
	}

	if res.ExitCode != expectedExitCode {
		t.Fatalf("exitcode expected: %v, actual: %v", expectedExitCode, res.ExitCode)
	}
}
<|endoftext|>