<commit_before>// Copyright 2013 Martini Authors
// Copyright 2014 The Macaron 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 csrf is a middleware that generates and validates CSRF tokens for Macaron.
package csrf

import (
	"net/http"

	"github.com/Unknwon/com"
	"github.com/go-macaron/session"
	"gopkg.in/macaron.v1"
)

const _VERSION = "0.0.5"

func Version() string {
	return _VERSION
}

// CSRF represents a CSRF service and is used to get the current token and validate a suspect token.
type CSRF interface {
	// Return HTTP header to search for token.
	GetHeaderName() string
	// Return form value to search for token.
	GetFormName() string
	// Return cookie name to search for token.
	GetCookieName() string
	// Return cookie path
	GetCookiePath() string
	// Return the token.
	GetToken() string
	// Validate by token.
	ValidToken(t string) bool
	// Error replies to the request with a custom function when ValidToken fails.
	Error(w http.ResponseWriter)
}

type csrf struct {
	// Header name value for setting and getting csrf token.
	Header string
	// Form name value for setting and getting csrf token.
	Form string
	// Cookie name value for setting and getting csrf token.
	Cookie string
	//Cookie path
	CookiePath string
	// Token generated to pass via header, cookie, or hidden form value.
	Token string
	// This value must be unique per user.
	ID string
	// Secret used along with the unique id above to generate the Token.
	Secret string
	// ErrorFunc is the custom function that replies to the request when ValidToken fails.
	ErrorFunc func(w http.ResponseWriter)
}

// GetHeaderName returns the name of the HTTP header for csrf token.
func (c *csrf) GetHeaderName() string {
	return c.Header
}

// GetFormName returns the name of the form value for csrf token.
func (c *csrf) GetFormName() string {
	return c.Form
}

// GetCookieName returns the name of the cookie for csrf token.
func (c *csrf) GetCookieName() string {
	return c.Cookie
}

// GetCookiePath returns the path of the cookie for csrf token.
func (c *csrf) GetCookiePath() string {
	return c.CookiePath
}

// GetToken returns the current token. This is typically used
// to populate a hidden form in an HTML template.
func (c *csrf) GetToken() string {
	return c.Token
}

// ValidToken validates the passed token against the existing Secret and ID.
func (c *csrf) ValidToken(t string) bool {
	return ValidToken(t, c.Secret, c.ID, "POST")
}

// Error replies to the request when ValidToken fails.
func (c *csrf) Error(w http.ResponseWriter) {
	c.ErrorFunc(w)
}

// Options maintains options to manage behavior of Generate.
type Options struct {
	// The global secret value used to generate Tokens.
	Secret string
	// HTTP header used to set and get token.
	Header string
	// Form value used to set and get token.
	Form string
	// Cookie value used to set and get token.
	Cookie string
	// Cookie path.
	CookiePath string
	// Key used for getting the unique ID per user.
	SessionKey string
	// oldSeesionKey saves old value corresponding to SessionKey.
	oldSeesionKey string
	// If true, send token via X-CSRFToken header.
	SetHeader bool
	// If true, send token via _csrf cookie.
	SetCookie bool
	// Set the Secure flag to true on the cookie.
	Secure bool
	// Disallow Origin appear in request header.
	Origin bool
	// The function called when Validate fails.
	ErrorFunc func(w http.ResponseWriter)
}

func prepareOptions(options []Options) Options {
	var opt Options
	if len(options) > 0 {
		opt = options[0]
	}

	// Defaults.
	if len(opt.Secret) == 0 {
		opt.Secret = string(com.RandomCreateBytes(10))
	}
	if len(opt.Header) == 0 {
		opt.Header = "X-CSRFToken"
	}
	if len(opt.Form) == 0 {
		opt.Form = "_csrf"
	}
	if len(opt.Cookie) == 0 {
		opt.Cookie = "_csrf"
	}
	if len(opt.CookiePath) == 0 {
		opt.CookiePath = "/"
	}
	if len(opt.SessionKey) == 0 {
		opt.SessionKey = "uid"
	}
	opt.oldSeesionKey = "_old_" + opt.SessionKey
	if opt.ErrorFunc == nil {
		opt.ErrorFunc = func(w http.ResponseWriter) {
			http.Error(w, "Invalid csrf token.", http.StatusBadRequest)
		}
	}

	return opt
}

// Generate maps CSRF to each request. If this request is a Get request, it will generate a new token.
// Additionally, depending on options set, generated tokens will be sent via Header and/or Cookie.
func Generate(options ...Options) macaron.Handler {
	opt := prepareOptions(options)
	return func(ctx *macaron.Context, sess session.Store) {
		x := &csrf{
			Secret:     opt.Secret,
			Header:     opt.Header,
			Form:       opt.Form,
			Cookie:     opt.Cookie,
			CookiePath: opt.CookiePath,
			ErrorFunc:  opt.ErrorFunc,
		}
		ctx.MapTo(x, (*CSRF)(nil))

		if opt.Origin && len(ctx.Req.Header.Get("Origin")) > 0 {
			return
		}

		x.ID = "0"
		uid := sess.Get(opt.SessionKey)
		if uid != nil {
			x.ID = com.ToStr(uid)
		}

		needsNew := false
		oldUid := sess.Get(opt.oldSeesionKey)
		if oldUid == nil || oldUid.(string) != x.ID {
			needsNew = true
			sess.Set(opt.oldSeesionKey, x.ID)
		} else {
			// If cookie present, map existing token, else generate a new one.
			if val := ctx.GetCookie(opt.Cookie); len(val) > 0 {
				// FIXME: test coverage.
				x.Token = val
			} else {
				needsNew = true
			}
		}

		if needsNew {
			// FIXME: actionId.
			x.Token = GenerateToken(x.Secret, x.ID, "POST")
		}

		if opt.SetCookie {
			ctx.SetCookie(opt.Cookie, x.Token, 86400, opt.CookiePath)
		}
		if opt.SetHeader {
			ctx.Resp.Header().Add(opt.Header, x.Token)
		}
	}
}

// Csrfer maps CSRF to each request. If this request is a Get request, it will generate a new token.
// Additionally, depending on options set, generated tokens will be sent via Header and/or Cookie.
func Csrfer(options ...Options) macaron.Handler {
	return Generate(options...)
}

// Validate should be used as a per route middleware. It attempts to get a token from a "X-CSRFToken"
// HTTP header and then a "_csrf" form value. If one of these is found, the token will be validated
// using ValidToken. If this validation fails, custom Error is sent in the reply.
// If neither a header or form value is found, http.StatusBadRequest is sent.
func Validate(ctx *macaron.Context, x CSRF) {
	if token := ctx.Req.Header.Get(x.GetHeaderName()); len(token) > 0 {
		if !x.ValidToken(token) {
			ctx.SetCookie(x.GetCookieName(), "", -1, x.GetCookiePath())
			x.Error(ctx.Resp)
		}
		return
	}
	if token := ctx.Req.FormValue(x.GetFormName()); len(token) > 0 {
		if !x.ValidToken(token) {
			ctx.SetCookie(x.GetCookieName(), "", -1, x.GetCookiePath())
			x.Error(ctx.Resp)
		}
		return
	}

	http.Error(ctx.Resp, "Bad Request: no CSRF token present", http.StatusBadRequest)
}
<commit_msg>Fix set CSRF every request which would never expire<commit_after>// Copyright 2013 Martini Authors
// Copyright 2014 The Macaron 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 csrf is a middleware that generates and validates CSRF tokens for Macaron.
package csrf

import (
	"net/http"
	"time"

	"github.com/Unknwon/com"
	"github.com/go-macaron/session"
	"gopkg.in/macaron.v1"
)

const _VERSION = "0.1.0"

func Version() string {
	return _VERSION
}

// CSRF represents a CSRF service and is used to get the current token and validate a suspect token.
type CSRF interface {
	// Return HTTP header to search for token.
	GetHeaderName() string
	// Return form value to search for token.
	GetFormName() string
	// Return cookie name to search for token.
	GetCookieName() string
	// Return cookie path
	GetCookiePath() string
	// Return the token.
	GetToken() string
	// Validate by token.
	ValidToken(t string) bool
	// Error replies to the request with a custom function when ValidToken fails.
	Error(w http.ResponseWriter)
}

type csrf struct {
	// Header name value for setting and getting csrf token.
	Header string
	// Form name value for setting and getting csrf token.
	Form string
	// Cookie name value for setting and getting csrf token.
	Cookie string
	//Cookie path
	CookiePath string
	// Token generated to pass via header, cookie, or hidden form value.
	Token string
	// This value must be unique per user.
	ID string
	// Secret used along with the unique id above to generate the Token.
	Secret string
	// ErrorFunc is the custom function that replies to the request when ValidToken fails.
	ErrorFunc func(w http.ResponseWriter)
}

// GetHeaderName returns the name of the HTTP header for csrf token.
func (c *csrf) GetHeaderName() string {
	return c.Header
}

// GetFormName returns the name of the form value for csrf token.
func (c *csrf) GetFormName() string {
	return c.Form
}

// GetCookieName returns the name of the cookie for csrf token.
func (c *csrf) GetCookieName() string {
	return c.Cookie
}

// GetCookiePath returns the path of the cookie for csrf token.
func (c *csrf) GetCookiePath() string {
	return c.CookiePath
}

// GetToken returns the current token. This is typically used
// to populate a hidden form in an HTML template.
func (c *csrf) GetToken() string {
	return c.Token
}

// ValidToken validates the passed token against the existing Secret and ID.
func (c *csrf) ValidToken(t string) bool {
	return ValidToken(t, c.Secret, c.ID, "POST")
}

// Error replies to the request when ValidToken fails.
func (c *csrf) Error(w http.ResponseWriter) {
	c.ErrorFunc(w)
}

// Options maintains options to manage behavior of Generate.
type Options struct {
	// The global secret value used to generate Tokens.
	Secret string
	// HTTP header used to set and get token.
	Header string
	// Form value used to set and get token.
	Form string
	// Cookie value used to set and get token.
	Cookie string
	// Cookie path.
	CookiePath string
	// Key used for getting the unique ID per user.
	SessionKey string
	// oldSeesionKey saves old value corresponding to SessionKey.
	oldSeesionKey string
	// If true, send token via X-CSRFToken header.
	SetHeader bool
	// If true, send token via _csrf cookie.
	SetCookie bool
	// Set the Secure flag to true on the cookie.
	Secure bool
	// Disallow Origin appear in request header.
	Origin bool
	// The function called when Validate fails.
	ErrorFunc func(w http.ResponseWriter)
}

func prepareOptions(options []Options) Options {
	var opt Options
	if len(options) > 0 {
		opt = options[0]
	}

	// Defaults.
	if len(opt.Secret) == 0 {
		opt.Secret = string(com.RandomCreateBytes(10))
	}
	if len(opt.Header) == 0 {
		opt.Header = "X-CSRFToken"
	}
	if len(opt.Form) == 0 {
		opt.Form = "_csrf"
	}
	if len(opt.Cookie) == 0 {
		opt.Cookie = "_csrf"
	}
	if len(opt.CookiePath) == 0 {
		opt.CookiePath = "/"
	}
	if len(opt.SessionKey) == 0 {
		opt.SessionKey = "uid"
	}
	opt.oldSeesionKey = "_old_" + opt.SessionKey
	if opt.ErrorFunc == nil {
		opt.ErrorFunc = func(w http.ResponseWriter) {
			http.Error(w, "Invalid csrf token.", http.StatusBadRequest)
		}
	}

	return opt
}

// Generate maps CSRF to each request. If this request is a Get request, it will generate a new token.
// Additionally, depending on options set, generated tokens will be sent via Header and/or Cookie.
func Generate(options ...Options) macaron.Handler {
	opt := prepareOptions(options)
	return func(ctx *macaron.Context, sess session.Store) {
		x := &csrf{
			Secret:     opt.Secret,
			Header:     opt.Header,
			Form:       opt.Form,
			Cookie:     opt.Cookie,
			CookiePath: opt.CookiePath,
			ErrorFunc:  opt.ErrorFunc,
		}
		ctx.MapTo(x, (*CSRF)(nil))

		if opt.Origin && len(ctx.Req.Header.Get("Origin")) > 0 {
			return
		}

		x.ID = "0"
		uid := sess.Get(opt.SessionKey)
		if uid != nil {
			x.ID = com.ToStr(uid)
		}

		needsNew := false
		oldUid := sess.Get(opt.oldSeesionKey)
		if oldUid == nil || oldUid.(string) != x.ID {
			needsNew = true
			sess.Set(opt.oldSeesionKey, x.ID)
		} else {
			// If cookie present, map existing token, else generate a new one.
			if val := ctx.GetCookie(opt.Cookie); len(val) > 0 {
				// FIXME: test coverage.
				x.Token = val
			} else {
				needsNew = true
			}
		}

		if needsNew {
			// FIXME: actionId.
			x.Token = GenerateToken(x.Secret, x.ID, "POST")
			if opt.SetCookie {
				ctx.SetCookie(opt.Cookie, x.Token, 0, opt.CookiePath, "", false, true, time.Now().AddDate(0, 0, 1))
			}
		}

		if opt.SetHeader {
			ctx.Resp.Header().Add(opt.Header, x.Token)
		}
	}
}

// Csrfer maps CSRF to each request. If this request is a Get request, it will generate a new token.
// Additionally, depending on options set, generated tokens will be sent via Header and/or Cookie.
func Csrfer(options ...Options) macaron.Handler {
	return Generate(options...)
}

// Validate should be used as a per route middleware. It attempts to get a token from a "X-CSRFToken"
// HTTP header and then a "_csrf" form value. If one of these is found, the token will be validated
// using ValidToken. If this validation fails, custom Error is sent in the reply.
// If neither a header or form value is found, http.StatusBadRequest is sent.
func Validate(ctx *macaron.Context, x CSRF) {
	if token := ctx.Req.Header.Get(x.GetHeaderName()); len(token) > 0 {
		if !x.ValidToken(token) {
			ctx.SetCookie(x.GetCookieName(), "", -1, x.GetCookiePath())
			x.Error(ctx.Resp)
		}
		return
	}
	if token := ctx.Req.FormValue(x.GetFormName()); len(token) > 0 {
		if !x.ValidToken(token) {
			ctx.SetCookie(x.GetCookieName(), "", -1, x.GetCookiePath())
			x.Error(ctx.Resp)
		}
		return
	}

	http.Error(ctx.Resp, "Bad Request: no CSRF token present", http.StatusBadRequest)
}
<|endoftext|>


<commit_before>package csrf

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

	"github.com/gorilla/context"
	"github.com/gorilla/securecookie"
)

// CSRF token length in bytes.
const tokenLength = 32

// Context/session keys & prefixes
const (
	tokenKey    string = "gorilla.csrf.Token"
	errorKey    string = "gorilla.csrf.Error"
	cookieName  string = "_gorilla_csrf"
	errorPrefix string = "gorilla/csrf: "
)

var (
	// The name value used in form fields.
	fieldName = tokenKey
	// defaultAge sets the default MaxAge for cookies.
	defaultAge = 3600 * 12
	// The default HTTP request header to inspect
	headerName = "X-CSRF-Token"
	// Idempotent (safe) methods as defined by RFC7231 section 4.2.2.
	safeMethods = []string{"GET", "HEAD", "OPTIONS", "TRACE"}
)

// TemplateTag provides a default template tag - e.g. {{ .csrfField }} - for use
// with the TemplateField function.
var TemplateTag = "csrfField"

var (
	// ErrNoReferer is returned when a HTTPS request provides an empty Referer
	// header.
	ErrNoReferer = errors.New("referer not supplied")
	// ErrBadReferer is returned when the scheme & host in the URL do not match
	// the supplied Referer header.
	ErrBadReferer = errors.New("referer invalid")
	// ErrNoToken is returned if no CSRF token is supplied in the request.
	ErrNoToken = errors.New("CSRF token not found in request")
	// ErrBadToken is returned if the CSRF token in the request does not match
	// the token in the session, or is otherwise malformed.
	ErrBadToken = errors.New("CSRF token invalid")
)

type csrf struct {
	h    http.Handler
	sc   *securecookie.SecureCookie
	st   store
	opts options
}

// options contains the optional settings for the CSRF middleware.
type options struct {
	MaxAge int
	Domain string
	Path   string
	// Note that the function and field names match the case of the associated
	// http.Cookie field instead of the "correct" HTTPOnly name that golint suggests.
	HttpOnly      bool
	Secure        bool
	RequestHeader string
	FieldName     string
	ErrorHandler  http.Handler
	CookieName    string
}

// Protect is HTTP middleware that provides Cross-Site Request Forgery
// protection.
//
// It securely generates a masked (unique-per-request) token that
// can be embedded in the HTTP response (e.g. form field or HTTP header).
// The original (unmasked) token is stored in the session, which is inaccessible
// by an attacker (provided you are using HTTPS). Subsequent requests are
// expected to include this token, which is compared against the session token.
// Requests that do not provide a matching token are served with a HTTP 403
// 'Forbidden' error response.
//
// Example:
//	package main
//
//	import (
//		"github.com/elithrar/protect"
//		"github.com/gorilla/mux"
//	)
//
//	func main() {
//	  r := mux.NewRouter()
//
//	  mux.HandlerFunc("/signup", GetSignupForm)
//	  // POST requests without a valid token will return a HTTP 403 Forbidden.
//	  mux.HandlerFunc("/signup/post", PostSignupForm)
//
//	  // Add the middleware to your router.
//	  http.ListenAndServe(":8000",
//            // Note that the authentication key provided should be 32 bytes
//            // long and persist across application restarts.
//			  csrf.Protect([]byte("32-byte-long-auth-key"))(r))
//	}
//
//	func GetSignupForm(w http.ResponseWriter, r *http.Request) {
//		// signup_form.tmpl just needs a {{ .csrfField }} template tag for
//		// csrf.TemplateField to inject the CSRF token into. Easy!
//		t.ExecuteTemplate(w, "signup_form.tmpl", map[string]interface{
//			csrf.TemplateTag: csrf.TemplateField(r),
//		})
//		// We could also retrieve the token directly from csrf.Token(r) and
//		// set it in the request header - w.Header.Set("X-CSRF-Token", token)
//		// This is useful if your sending JSON to clients or a front-end JavaScript
//		// framework.
//	}
//
func Protect(authKey []byte, opts ...func(*csrf)) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		cs := parseOptions(h, opts...)

		// Set the defaults if no options have been specified
		if cs.opts.ErrorHandler == nil {
			cs.opts.ErrorHandler = http.HandlerFunc(unauthorizedHandler)
		}

		if cs.opts.MaxAge < 1 {
			// Default of 12 hours
			cs.opts.MaxAge = defaultAge
		}

		if cs.opts.FieldName == "" {
			cs.opts.FieldName = fieldName
		}

		if cs.opts.CookieName == "" {
			cs.opts.CookieName = cookieName
		}

		if cs.opts.RequestHeader == "" {
			cs.opts.RequestHeader = headerName
		}

		// Create an authenticated securecookie instance.
		if cs.sc == nil {
			cs.sc = securecookie.New(authKey, nil)
			// Set the MaxAge of the underlying securecookie.
			cs.sc.MaxAge(cs.opts.MaxAge)
		}

		if cs.st == nil {
			// Default to the cookieStore
			cs.st = &cookieStore{
				name:     cs.opts.CookieName,
				maxAge:   cs.opts.MaxAge,
				secure:   cs.opts.Secure,
				httpOnly: cs.opts.HttpOnly,
				sc:       cs.sc,
			}
		}

		return cs
	}
}

// Implements http.Handler for the csrf type.
func (cs *csrf) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Retrieve the token from the session.
	// An error represents either a cookie that failed HMAC validation
	// or that doesn't exist.
	realToken, err := cs.st.Get(r)
	if err != nil || len(realToken) != tokenLength {
		// If there was an error retrieving the token, the token doesn't exist
		// yet, or it's the wrong length, generate a new token.
		// Note that the new token will (correctly) fail validation downstream
		// as it will no longer match the request token.
		realToken, err = generateRandomBytes(tokenLength)
		if err != nil {
			envError(r, err)
			cs.opts.ErrorHandler.ServeHTTP(w, r)
			return
		}

		// Save the new (real) token in the session store.
		err = cs.st.Save(realToken, w)
		if err != nil {
			envError(r, err)
			cs.opts.ErrorHandler.ServeHTTP(w, r)
			return
		}
	}

	// Save the masked token to the request context
	context.Set(r, tokenKey, mask(realToken, r))

	// HTTP methods not defined as idempotent ("safe") under RFC7231 require
	// inspection.
	if !contains(safeMethods, r.Method) {
		// Enforce an origin check for HTTPS connections. As per the Django CSRF
		// implementation (https://goo.gl/vKA7GE) the Referer header is almost
		// always present for same-domain HTTP requests.
		if r.URL.Scheme == "https" {
			// Fetch the Referer value. Call the error handler if it's empty or
			// otherwise fails to parse.
			referer, err := url.Parse(r.Referer())
			if err != nil || referer.String() == "" {
				envError(r, ErrNoReferer)
				cs.opts.ErrorHandler.ServeHTTP(w, r)
				return
			}

			if sameOrigin(r.URL, referer) == false {
				envError(r, ErrBadReferer)
				cs.opts.ErrorHandler.ServeHTTP(w, r)
				return
			}
		}

		// If the token returned from the session store is nil for non-idempotent
		// ("unsafe") methods, call the error handler.
		if realToken == nil {
			envError(r, ErrNoToken)
			cs.opts.ErrorHandler.ServeHTTP(w, r)
			return
		}

		// Retrieve the combined token (pad + masked) token and unmask it.
		requestToken := unmask(cs.requestToken(r))

		// Compare the request token against the real token
		if !compareTokens(requestToken, realToken) {
			envError(r, ErrBadToken)
			cs.opts.ErrorHandler.ServeHTTP(w, r)
			return
		}

	}

	// Set the Vary: Cookie header to protect clients from caching the response.
	w.Header().Add("Vary", "Cookie")

	// Call the wrapped handler/router on success.
	cs.h.ServeHTTP(w, r)
	// Clear the request context after the handler has completed.
	context.Clear(r)
}

// unauthorizedhandler sets a HTTP 403 Forbidden status and writes the
// CSRF failure reason to the response.
func unauthorizedHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, fmt.Sprintf("%s - %s",
		http.StatusText(http.StatusForbidden), FailureReason(r)),
		http.StatusForbidden)
	return
}
<commit_msg>[feature]: Use securecookie.JSONEncoder{}<commit_after>package csrf

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

	"github.com/gorilla/context"
	"github.com/gorilla/securecookie"
)

// CSRF token length in bytes.
const tokenLength = 32

// Context/session keys & prefixes
const (
	tokenKey    string = "gorilla.csrf.Token"
	errorKey    string = "gorilla.csrf.Error"
	cookieName  string = "_gorilla_csrf"
	errorPrefix string = "gorilla/csrf: "
)

var (
	// The name value used in form fields.
	fieldName = tokenKey
	// defaultAge sets the default MaxAge for cookies.
	defaultAge = 3600 * 12
	// The default HTTP request header to inspect
	headerName = "X-CSRF-Token"
	// Idempotent (safe) methods as defined by RFC7231 section 4.2.2.
	safeMethods = []string{"GET", "HEAD", "OPTIONS", "TRACE"}
)

// TemplateTag provides a default template tag - e.g. {{ .csrfField }} - for use
// with the TemplateField function.
var TemplateTag = "csrfField"

var (
	// ErrNoReferer is returned when a HTTPS request provides an empty Referer
	// header.
	ErrNoReferer = errors.New("referer not supplied")
	// ErrBadReferer is returned when the scheme & host in the URL do not match
	// the supplied Referer header.
	ErrBadReferer = errors.New("referer invalid")
	// ErrNoToken is returned if no CSRF token is supplied in the request.
	ErrNoToken = errors.New("CSRF token not found in request")
	// ErrBadToken is returned if the CSRF token in the request does not match
	// the token in the session, or is otherwise malformed.
	ErrBadToken = errors.New("CSRF token invalid")
)

type csrf struct {
	h    http.Handler
	sc   *securecookie.SecureCookie
	st   store
	opts options
}

// options contains the optional settings for the CSRF middleware.
type options struct {
	MaxAge int
	Domain string
	Path   string
	// Note that the function and field names match the case of the associated
	// http.Cookie field instead of the "correct" HTTPOnly name that golint suggests.
	HttpOnly      bool
	Secure        bool
	RequestHeader string
	FieldName     string
	ErrorHandler  http.Handler
	CookieName    string
}

// Protect is HTTP middleware that provides Cross-Site Request Forgery
// protection.
//
// It securely generates a masked (unique-per-request) token that
// can be embedded in the HTTP response (e.g. form field or HTTP header).
// The original (unmasked) token is stored in the session, which is inaccessible
// by an attacker (provided you are using HTTPS). Subsequent requests are
// expected to include this token, which is compared against the session token.
// Requests that do not provide a matching token are served with a HTTP 403
// 'Forbidden' error response.
//
// Example:
//	package main
//
//	import (
//		"github.com/elithrar/protect"
//		"github.com/gorilla/mux"
//	)
//
//	func main() {
//	  r := mux.NewRouter()
//
//	  mux.HandlerFunc("/signup", GetSignupForm)
//	  // POST requests without a valid token will return a HTTP 403 Forbidden.
//	  mux.HandlerFunc("/signup/post", PostSignupForm)
//
//	  // Add the middleware to your router.
//	  http.ListenAndServe(":8000",
//            // Note that the authentication key provided should be 32 bytes
//            // long and persist across application restarts.
//			  csrf.Protect([]byte("32-byte-long-auth-key"))(r))
//	}
//
//	func GetSignupForm(w http.ResponseWriter, r *http.Request) {
//		// signup_form.tmpl just needs a {{ .csrfField }} template tag for
//		// csrf.TemplateField to inject the CSRF token into. Easy!
//		t.ExecuteTemplate(w, "signup_form.tmpl", map[string]interface{
//			csrf.TemplateTag: csrf.TemplateField(r),
//		})
//		// We could also retrieve the token directly from csrf.Token(r) and
//		// set it in the request header - w.Header.Set("X-CSRF-Token", token)
//		// This is useful if your sending JSON to clients or a front-end JavaScript
//		// framework.
//	}
//
func Protect(authKey []byte, opts ...func(*csrf)) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		cs := parseOptions(h, opts...)

		// Set the defaults if no options have been specified
		if cs.opts.ErrorHandler == nil {
			cs.opts.ErrorHandler = http.HandlerFunc(unauthorizedHandler)
		}

		if cs.opts.MaxAge < 1 {
			// Default of 12 hours
			cs.opts.MaxAge = defaultAge
		}

		if cs.opts.FieldName == "" {
			cs.opts.FieldName = fieldName
		}

		if cs.opts.CookieName == "" {
			cs.opts.CookieName = cookieName
		}

		if cs.opts.RequestHeader == "" {
			cs.opts.RequestHeader = headerName
		}

		// Create an authenticated securecookie instance.
		if cs.sc == nil {
			cs.sc = securecookie.New(authKey, nil)
			// Use JSON serialization (faster than one-off gob encoding)
			cs.sc.SetSerializer(securecookie.JSONEncoder{})
			// Set the MaxAge of the underlying securecookie.
			cs.sc.MaxAge(cs.opts.MaxAge)
		}

		if cs.st == nil {
			// Default to the cookieStore
			cs.st = &cookieStore{
				name:     cs.opts.CookieName,
				maxAge:   cs.opts.MaxAge,
				secure:   cs.opts.Secure,
				httpOnly: cs.opts.HttpOnly,
				sc:       cs.sc,
			}
		}

		return cs
	}
}

// Implements http.Handler for the csrf type.
func (cs *csrf) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Retrieve the token from the session.
	// An error represents either a cookie that failed HMAC validation
	// or that doesn't exist.
	realToken, err := cs.st.Get(r)
	if err != nil || len(realToken) != tokenLength {
		// If there was an error retrieving the token, the token doesn't exist
		// yet, or it's the wrong length, generate a new token.
		// Note that the new token will (correctly) fail validation downstream
		// as it will no longer match the request token.
		realToken, err = generateRandomBytes(tokenLength)
		if err != nil {
			envError(r, err)
			cs.opts.ErrorHandler.ServeHTTP(w, r)
			return
		}

		// Save the new (real) token in the session store.
		err = cs.st.Save(realToken, w)
		if err != nil {
			envError(r, err)
			cs.opts.ErrorHandler.ServeHTTP(w, r)
			return
		}
	}

	// Save the masked token to the request context
	context.Set(r, tokenKey, mask(realToken, r))

	// HTTP methods not defined as idempotent ("safe") under RFC7231 require
	// inspection.
	if !contains(safeMethods, r.Method) {
		// Enforce an origin check for HTTPS connections. As per the Django CSRF
		// implementation (https://goo.gl/vKA7GE) the Referer header is almost
		// always present for same-domain HTTP requests.
		if r.URL.Scheme == "https" {
			// Fetch the Referer value. Call the error handler if it's empty or
			// otherwise fails to parse.
			referer, err := url.Parse(r.Referer())
			if err != nil || referer.String() == "" {
				envError(r, ErrNoReferer)
				cs.opts.ErrorHandler.ServeHTTP(w, r)
				return
			}

			if sameOrigin(r.URL, referer) == false {
				envError(r, ErrBadReferer)
				cs.opts.ErrorHandler.ServeHTTP(w, r)
				return
			}
		}

		// If the token returned from the session store is nil for non-idempotent
		// ("unsafe") methods, call the error handler.
		if realToken == nil {
			envError(r, ErrNoToken)
			cs.opts.ErrorHandler.ServeHTTP(w, r)
			return
		}

		// Retrieve the combined token (pad + masked) token and unmask it.
		requestToken := unmask(cs.requestToken(r))

		// Compare the request token against the real token
		if !compareTokens(requestToken, realToken) {
			envError(r, ErrBadToken)
			cs.opts.ErrorHandler.ServeHTTP(w, r)
			return
		}

	}

	// Set the Vary: Cookie header to protect clients from caching the response.
	w.Header().Add("Vary", "Cookie")

	// Call the wrapped handler/router on success.
	cs.h.ServeHTTP(w, r)
	// Clear the request context after the handler has completed.
	context.Clear(r)
}

// unauthorizedhandler sets a HTTP 403 Forbidden status and writes the
// CSRF failure reason to the response.
func unauthorizedHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, fmt.Sprintf("%s - %s",
		http.StatusText(http.StatusForbidden), FailureReason(r)),
		http.StatusForbidden)
	return
}
<|endoftext|>


<commit_before>package interpose

import (
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
)

func TestServeHTTP(t *testing.T) {
	result := ""
	response := httptest.NewRecorder()

	middle := New()

	middle.Use(func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
			result += "0"
			next.ServeHTTP(rw, req)
			result += "0"
		})
	})

	middle.Use(func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
			result += "1"
			next.ServeHTTP(rw, req)
			result += "1"
		})
	})

	middle.UseHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		result += "2"
	}))

	middle.ServeHTTP(response, (*http.Request)(nil))
	expect(t, result, "01210")
}

func TestEmptyMiddleware(t *testing.T) {
	result := ""
	response := httptest.NewRecorder()

	middle := New()

	middle.ServeHTTP(response, (*http.Request)(nil))
	expect(t, result, "")
}

func expect(t *testing.T, a interface{}, b interface{}) {
	if a != b {
		t.Errorf("Expected %v (type %v) - Got %v (type %v)", b, reflect.TypeOf(b), a, reflect.TypeOf(a))
	}
}
<commit_msg>Updates the tests to add a benchmark demonstrating the difference between the "compiled" Interpose (i.e., calling .Handler() ) and the "raw" Interpose (i.e., not calling .Handler() ). NB, on my Macbook Air, the compiled Interpose is ~25% faster, requring 500 nanoseconds less time per request than the raw Interpose.<commit_after>package interpose

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"

	"github.com/interpose/middleware"
)

func BasicMiddleware() *Middleware {
	middle := New()

	middle.Use(middleware.Json())
	middle.Use(middleware.Buffer())

	middle.Use(func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
			fmt.Fprint(rw, "0")
			next.ServeHTTP(rw, req)
			fmt.Fprint(rw, "0")
		})
	})

	middle.Use(func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
			fmt.Fprint(rw, "1")
			next.ServeHTTP(rw, req)
			fmt.Fprint(rw, "1")
		})
	})

	middle.UseHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		fmt.Fprint(rw, "2")
	}))

	return middle
}

func TestCompiledMiddleware(t *testing.T) {
	response := httptest.NewRecorder()

	middle := BasicMiddleware().Handler()

	middle.ServeHTTP(response, (*http.Request)(nil))
	out, err := ioutil.ReadAll(response.Body)
	if err != nil {
		t.Error(err)
	}
	expect(t, string(out), "01210")
}

func TestServeHTTP(t *testing.T) {
	response := httptest.NewRecorder()

	middle := BasicMiddleware()

	middle.ServeHTTP(response, (*http.Request)(nil))
	out, err := ioutil.ReadAll(response.Body)
	if err != nil {
		t.Error(err)
	}
	expect(t, string(out), "01210")
}

func TestEmptyMiddleware(t *testing.T) {
	response := httptest.NewRecorder()

	middle := New()

	middle.ServeHTTP(response, (*http.Request)(nil))
	out, err := ioutil.ReadAll(response.Body)
	if err != nil {
		t.Error(err)
	}
	expect(t, string(out), "")
}

func BenchmarkCompiled(b *testing.B) {
	response := httptest.NewRecorder()

	middle := BasicMiddleware().Handler()

	for i := 0; i < b.N; i++ {
		middle.ServeHTTP(response, (*http.Request)(nil))
	}
}

func BenchmarkUncompiled(b *testing.B) {
	response := httptest.NewRecorder()

	middle := BasicMiddleware()

	for i := 0; i < b.N; i++ {
		middle.ServeHTTP(response, (*http.Request)(nil))
	}
}

func BenchmarkEmpty(b *testing.B) {
	response := httptest.NewRecorder()

	middle := New()
	middle.UseHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		return
	}))

	for i := 0; i < b.N; i++ {
		middle.ServeHTTP(response, (*http.Request)(nil))
	}
}

func expect(t *testing.T, a interface{}, b interface{}) {
	if a != b {
		t.Errorf("Expected %v (type %v) - Got %v (type %v)", b, reflect.TypeOf(b), a, reflect.TypeOf(a))
	}
}
<|endoftext|>


<commit_before>package zfs_test

import (
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"
	"time"

	zfs "github.com/mistifyio/go-zfs/v3"
)

func sleep(delay int) {
	time.Sleep(time.Duration(delay) * time.Second)
}

func pow2(x int) int64 {
	return int64(math.Pow(2, float64(x)))
}

// https://github.com/benbjohnson/testing
// assert fails the test if the condition is false.
func assert(t *testing.T, condition bool, msg string, v ...interface{}) {
	t.Helper()

	if !condition {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: "+msg+"\033[39m\n\n", append([]interface{}{filepath.Base(file), line}, v...)...)
		t.FailNow()
	}
}

// ok fails the test if an err is not nil.
func ok(t *testing.T, err error) {
	t.Helper()

	if err != nil {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: unexpected error: %s\033[39m\n\n", filepath.Base(file), line, err.Error())
		t.FailNow()
	}
}

// nok fails the test if an err is nil.
func nok(t *testing.T, err error) {
	t.Helper()

	if err == nil {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: expected error: %s\033[39m\n\n", filepath.Base(file), line, err.Error())
		t.FailNow()
	}
}

// equals fails the test if exp is not equal to act.
func equals(t *testing.T, exp, act interface{}) {
	t.Helper()

	if !reflect.DeepEqual(exp, act) {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d:\n\n\texp: %#v\n\n\tgot: %#v\033[39m\n\n", filepath.Base(file), line, exp, act)
		t.FailNow()
	}
}

func zpoolTest(t *testing.T, fn func()) {
	t.Helper()

	tempfiles := make([]string, 3)
	for i := range tempfiles {
		f, _ := ioutil.TempFile("/tmp/", "zfs-")
		err := f.Truncate(pow2(30))
		f.Close()
		ok(t, err)
		tempfiles[i] = f.Name()
		defer os.Remove(f.Name()) // nolint:revive // its ok to defer to end of func
	}

	pool, err := zfs.CreateZpool("test", nil, tempfiles...)
	ok(t, err)
	defer pool.Destroy()
	ok(t, err)
	fn()
}

func TestDatasets(t *testing.T) {
	t.Helper()

	zpoolTest(t, func() {
		_, err := zfs.Datasets("")
		ok(t, err)

		ds, err := zfs.GetDataset("test")
		ok(t, err)
		equals(t, zfs.DatasetFilesystem, ds.Type)
		equals(t, "", ds.Origin)
		if runtime.GOOS != "solaris" {
			assert(t, ds.Logicalused != 0, "Logicalused is not greater than 0")
		}
	})
}

func TestDatasetGetProperty(t *testing.T) {
	zpoolTest(t, func() {
		ds, err := zfs.GetDataset("test")
		ok(t, err)

		prop, err := ds.GetProperty("foobarbaz")
		nok(t, err)
		equals(t, "", prop)

		prop, err = ds.GetProperty("compression")
		ok(t, err)
		equals(t, "off", prop)
	})
}

func TestSnapshots(t *testing.T) {
	zpoolTest(t, func() {
		snapshots, err := zfs.Snapshots("")
		ok(t, err)

		for _, snapshot := range snapshots {
			equals(t, zfs.DatasetSnapshot, snapshot.Type)
		}
	})
}

func TestFilesystems(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/filesystem-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestCreateFilesystemWithProperties(t *testing.T) {
	zpoolTest(t, func() {
		props := map[string]string{
			"compression": "lz4",
		}

		f, err := zfs.CreateFilesystem("test/filesystem-test", props)
		ok(t, err)

		equals(t, "lz4", f.Compression)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestVolumes(t *testing.T) {
	zpoolTest(t, func() {
		v, err := zfs.CreateVolume("test/volume-test", uint64(pow2(23)), nil)
		ok(t, err)

		// volumes are sometimes "busy" if you try to manipulate them right away
		sleep(1)

		equals(t, zfs.DatasetVolume, v.Type)
		volumes, err := zfs.Volumes("")
		ok(t, err)

		for _, volume := range volumes {
			equals(t, zfs.DatasetVolume, volume.Type)
		}

		ok(t, v.Destroy(zfs.DestroyDefault))
	})
}

func TestSnapshot(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		s, err := f.Snapshot("test", false)
		ok(t, err)

		equals(t, zfs.DatasetSnapshot, s.Type)

		equals(t, "test/snapshot-test@test", s.Name)

		ok(t, s.Destroy(zfs.DestroyDefault))

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestClone(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		s, err := f.Snapshot("test", false)
		ok(t, err)

		equals(t, zfs.DatasetSnapshot, s.Type)
		equals(t, "test/snapshot-test@test", s.Name)

		c, err := s.Clone("test/clone-test", nil)
		ok(t, err)

		equals(t, zfs.DatasetFilesystem, c.Type)

		ok(t, c.Destroy(zfs.DestroyDefault))

		ok(t, s.Destroy(zfs.DestroyDefault))

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestSendSnapshot(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		s, err := f.Snapshot("test", false)
		ok(t, err)

		file, _ := ioutil.TempFile("/tmp/", "zfs-")
		defer file.Close()
		err = file.Truncate(pow2(30))
		ok(t, err)
		defer os.Remove(file.Name())

		err = s.SendSnapshot(file)
		ok(t, err)

		ok(t, s.Destroy(zfs.DestroyDefault))

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestChildren(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		s, err := f.Snapshot("test", false)
		ok(t, err)

		equals(t, zfs.DatasetSnapshot, s.Type)
		equals(t, "test/snapshot-test@test", s.Name)

		children, err := f.Children(0)
		ok(t, err)

		equals(t, 1, len(children))
		equals(t, "test/snapshot-test@test", children[0].Name)

		ok(t, s.Destroy(zfs.DestroyDefault))
		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestListZpool(t *testing.T) {
	zpoolTest(t, func() {
		pools, err := zfs.ListZpools()
		ok(t, err)
		for _, pool := range pools {
			if pool.Name == "test" {
				equals(t, "test", pool.Name)
				return
			}
		}
		t.Fatal("Failed to find test pool")
	})
}

func TestRollback(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		s1, err := f.Snapshot("test", false)
		ok(t, err)

		_, err = f.Snapshot("test2", false)
		ok(t, err)

		s3, err := f.Snapshot("test3", false)
		ok(t, err)

		err = s3.Rollback(false)
		ok(t, err)

		err = s1.Rollback(false)
		assert(t, err != nil, "should error when rolling back beyond most recent without destroyMoreRecent = true")

		err = s1.Rollback(true)
		ok(t, err)

		ok(t, s1.Destroy(zfs.DestroyDefault))

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestDiff(t *testing.T) {
	zpoolTest(t, func() {
		fs, err := zfs.CreateFilesystem("test/origin", nil)
		ok(t, err)

		linkedFile, err := os.Create(filepath.Join(fs.Mountpoint, "linked"))
		ok(t, err)

		movedFile, err := os.Create(filepath.Join(fs.Mountpoint, "file"))
		ok(t, err)

		snapshot, err := fs.Snapshot("snapshot", false)
		ok(t, err)

		unicodeFile, err := os.Create(filepath.Join(fs.Mountpoint, "i ❤ unicode"))
		ok(t, err)

		err = os.Rename(movedFile.Name(), movedFile.Name()+"-new")
		ok(t, err)

		err = os.Link(linkedFile.Name(), linkedFile.Name()+"_hard")
		ok(t, err)

		inodeChanges, err := fs.Diff(snapshot.Name)
		ok(t, err)
		equals(t, 4, len(inodeChanges))

		unicodePath := "/test/origin/i\x040\x1c2\x135\x144\x040unicode"
		wants := map[string]*zfs.InodeChange{
			"/test/origin/linked": {
				Type:                 zfs.File,
				Change:               zfs.Modified,
				ReferenceCountChange: 1,
			},
			"/test/origin/file": {
				Type:    zfs.File,
				Change:  zfs.Renamed,
				NewPath: "/test/origin/file-new",
			},
			"/test/origin/i ❤ unicode": {
				Path:   "❤❤ unicode ❤❤",
				Type:   zfs.File,
				Change: zfs.Created,
			},
			unicodePath: {
				Path:   "❤❤ unicode ❤❤",
				Type:   zfs.File,
				Change: zfs.Created,
			},
			"/test/origin/": {
				Type:   zfs.Directory,
				Change: zfs.Modified,
			},
		}
		for _, change := range inodeChanges {
			want := wants[change.Path]
			want.Path = change.Path
			delete(wants, change.Path)

			equals(t, want, change)
		}

		equals(t, 1, len(wants))
		for _, want := range wants {
			equals(t, "❤❤ unicode ❤❤", want.Path)
		}

		ok(t, movedFile.Close())
		ok(t, unicodeFile.Close())
		ok(t, linkedFile.Close())
		ok(t, snapshot.Destroy(zfs.DestroyForceUmount))
		ok(t, fs.Destroy(zfs.DestroyForceUmount))
	})
}
<commit_msg>test: Avoid reptitive/duplicate error logging and quitting<commit_after>package zfs_test

import (
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"
	"time"

	zfs "github.com/mistifyio/go-zfs/v3"
)

func sleep(delay int) {
	time.Sleep(time.Duration(delay) * time.Second)
}

func pow2(x int) int64 {
	return int64(math.Pow(2, float64(x)))
}

// https://github.com/benbjohnson/testing
// assert fails the test if the condition is false.
func _assert(t *testing.T, condition bool, msg string, v ...interface{}) {
	t.Helper()

	if !condition {
		_, file, line, _ := runtime.Caller(2)
		fmt.Printf("\033[31m%s:%d: "+msg+"\033[39m\n\n", append([]interface{}{filepath.Base(file), line}, v...)...)
		t.FailNow()
	}
}

func assert(t *testing.T, condition bool, msg string, v ...interface{}) {
	t.Helper()
	_assert(t, condition, msg, v...)
}

// ok fails the test if an err is not nil.
func ok(t *testing.T, err error) {
	t.Helper()
	_assert(t, err == nil, "unexpected error: %v", err)
}

// nok fails the test if an err is nil.
func nok(t *testing.T, err error) {
	t.Helper()
	_assert(t, err != nil, "expected error, got nil")
}

// equals fails the test if exp is not equal to act.
func equals(t *testing.T, exp, act interface{}) {
	t.Helper()
	_assert(t, reflect.DeepEqual(exp, act), "exp: %#v\n\ngot: %#v", exp, act)
}

func zpoolTest(t *testing.T, fn func()) {
	t.Helper()

	tempfiles := make([]string, 3)
	for i := range tempfiles {
		f, _ := ioutil.TempFile("/tmp/", "zfs-")
		err := f.Truncate(pow2(30))
		f.Close()
		ok(t, err)
		tempfiles[i] = f.Name()
		defer os.Remove(f.Name()) // nolint:revive // its ok to defer to end of func
	}

	pool, err := zfs.CreateZpool("test", nil, tempfiles...)
	ok(t, err)
	defer pool.Destroy()
	ok(t, err)
	fn()
}

func TestDatasets(t *testing.T) {
	t.Helper()

	zpoolTest(t, func() {
		_, err := zfs.Datasets("")
		ok(t, err)

		ds, err := zfs.GetDataset("test")
		ok(t, err)
		equals(t, zfs.DatasetFilesystem, ds.Type)
		equals(t, "", ds.Origin)
		if runtime.GOOS != "solaris" {
			assert(t, ds.Logicalused != 0, "Logicalused is not greater than 0")
		}
	})
}

func TestDatasetGetProperty(t *testing.T) {
	zpoolTest(t, func() {
		ds, err := zfs.GetDataset("test")
		ok(t, err)

		prop, err := ds.GetProperty("foobarbaz")
		nok(t, err)
		equals(t, "", prop)

		prop, err = ds.GetProperty("compression")
		ok(t, err)
		equals(t, "off", prop)
	})
}

func TestSnapshots(t *testing.T) {
	zpoolTest(t, func() {
		snapshots, err := zfs.Snapshots("")
		ok(t, err)

		for _, snapshot := range snapshots {
			equals(t, zfs.DatasetSnapshot, snapshot.Type)
		}
	})
}

func TestFilesystems(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/filesystem-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestCreateFilesystemWithProperties(t *testing.T) {
	zpoolTest(t, func() {
		props := map[string]string{
			"compression": "lz4",
		}

		f, err := zfs.CreateFilesystem("test/filesystem-test", props)
		ok(t, err)

		equals(t, "lz4", f.Compression)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestVolumes(t *testing.T) {
	zpoolTest(t, func() {
		v, err := zfs.CreateVolume("test/volume-test", uint64(pow2(23)), nil)
		ok(t, err)

		// volumes are sometimes "busy" if you try to manipulate them right away
		sleep(1)

		equals(t, zfs.DatasetVolume, v.Type)
		volumes, err := zfs.Volumes("")
		ok(t, err)

		for _, volume := range volumes {
			equals(t, zfs.DatasetVolume, volume.Type)
		}

		ok(t, v.Destroy(zfs.DestroyDefault))
	})
}

func TestSnapshot(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		s, err := f.Snapshot("test", false)
		ok(t, err)

		equals(t, zfs.DatasetSnapshot, s.Type)

		equals(t, "test/snapshot-test@test", s.Name)

		ok(t, s.Destroy(zfs.DestroyDefault))

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestClone(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		s, err := f.Snapshot("test", false)
		ok(t, err)

		equals(t, zfs.DatasetSnapshot, s.Type)
		equals(t, "test/snapshot-test@test", s.Name)

		c, err := s.Clone("test/clone-test", nil)
		ok(t, err)

		equals(t, zfs.DatasetFilesystem, c.Type)

		ok(t, c.Destroy(zfs.DestroyDefault))

		ok(t, s.Destroy(zfs.DestroyDefault))

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestSendSnapshot(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		s, err := f.Snapshot("test", false)
		ok(t, err)

		file, _ := ioutil.TempFile("/tmp/", "zfs-")
		defer file.Close()
		err = file.Truncate(pow2(30))
		ok(t, err)
		defer os.Remove(file.Name())

		err = s.SendSnapshot(file)
		ok(t, err)

		ok(t, s.Destroy(zfs.DestroyDefault))

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestChildren(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		s, err := f.Snapshot("test", false)
		ok(t, err)

		equals(t, zfs.DatasetSnapshot, s.Type)
		equals(t, "test/snapshot-test@test", s.Name)

		children, err := f.Children(0)
		ok(t, err)

		equals(t, 1, len(children))
		equals(t, "test/snapshot-test@test", children[0].Name)

		ok(t, s.Destroy(zfs.DestroyDefault))
		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestListZpool(t *testing.T) {
	zpoolTest(t, func() {
		pools, err := zfs.ListZpools()
		ok(t, err)
		for _, pool := range pools {
			if pool.Name == "test" {
				equals(t, "test", pool.Name)
				return
			}
		}
		t.Fatal("Failed to find test pool")
	})
}

func TestRollback(t *testing.T) {
	zpoolTest(t, func() {
		f, err := zfs.CreateFilesystem("test/snapshot-test", nil)
		ok(t, err)

		filesystems, err := zfs.Filesystems("")
		ok(t, err)

		for _, filesystem := range filesystems {
			equals(t, zfs.DatasetFilesystem, filesystem.Type)
		}

		s1, err := f.Snapshot("test", false)
		ok(t, err)

		_, err = f.Snapshot("test2", false)
		ok(t, err)

		s3, err := f.Snapshot("test3", false)
		ok(t, err)

		err = s3.Rollback(false)
		ok(t, err)

		err = s1.Rollback(false)
		assert(t, err != nil, "should error when rolling back beyond most recent without destroyMoreRecent = true")

		err = s1.Rollback(true)
		ok(t, err)

		ok(t, s1.Destroy(zfs.DestroyDefault))

		ok(t, f.Destroy(zfs.DestroyDefault))
	})
}

func TestDiff(t *testing.T) {
	zpoolTest(t, func() {
		fs, err := zfs.CreateFilesystem("test/origin", nil)
		ok(t, err)

		linkedFile, err := os.Create(filepath.Join(fs.Mountpoint, "linked"))
		ok(t, err)

		movedFile, err := os.Create(filepath.Join(fs.Mountpoint, "file"))
		ok(t, err)

		snapshot, err := fs.Snapshot("snapshot", false)
		ok(t, err)

		unicodeFile, err := os.Create(filepath.Join(fs.Mountpoint, "i ❤ unicode"))
		ok(t, err)

		err = os.Rename(movedFile.Name(), movedFile.Name()+"-new")
		ok(t, err)

		err = os.Link(linkedFile.Name(), linkedFile.Name()+"_hard")
		ok(t, err)

		inodeChanges, err := fs.Diff(snapshot.Name)
		ok(t, err)
		equals(t, 4, len(inodeChanges))

		unicodePath := "/test/origin/i\x040\x1c2\x135\x144\x040unicode"
		wants := map[string]*zfs.InodeChange{
			"/test/origin/linked": {
				Type:                 zfs.File,
				Change:               zfs.Modified,
				ReferenceCountChange: 1,
			},
			"/test/origin/file": {
				Type:    zfs.File,
				Change:  zfs.Renamed,
				NewPath: "/test/origin/file-new",
			},
			"/test/origin/i ❤ unicode": {
				Path:   "❤❤ unicode ❤❤",
				Type:   zfs.File,
				Change: zfs.Created,
			},
			unicodePath: {
				Path:   "❤❤ unicode ❤❤",
				Type:   zfs.File,
				Change: zfs.Created,
			},
			"/test/origin/": {
				Type:   zfs.Directory,
				Change: zfs.Modified,
			},
		}
		for _, change := range inodeChanges {
			want := wants[change.Path]
			want.Path = change.Path
			delete(wants, change.Path)

			equals(t, want, change)
		}

		equals(t, 1, len(wants))
		for _, want := range wants {
			equals(t, "❤❤ unicode ❤❤", want.Path)
		}

		ok(t, movedFile.Close())
		ok(t, unicodeFile.Close())
		ok(t, linkedFile.Close())
		ok(t, snapshot.Destroy(zfs.DestroyForceUmount))
		ok(t, fs.Destroy(zfs.DestroyForceUmount))
	})
}
<|endoftext|>


<commit_before>package soapboxd

import (
	"bytes"
	"encoding/json"
	"time"

	pb "github.com/adhocteam/soapbox/proto"
	gpb "github.com/golang/protobuf/ptypes/timestamp"
	"github.com/pkg/errors"

	"golang.org/x/net/context"
)

func (s *server) ListEnvironments(ctx context.Context, req *pb.ListEnvironmentRequest) (*pb.ListEnvironmentResponse, error) {
	listSQL := "SELECT id, application_id, name, slug, vars, created_at FROM environments WHERE application_id = $1 ORDER BY id"
	rows, err := s.db.Query(listSQL, req.GetApplicationId())
	if err != nil {
		return nil, errors.Wrap(err, "querying db for environments")
	}
	var envs []*pb.Environment
	for rows.Next() {
		env := &pb.Environment{
			CreatedAt: new(gpb.Timestamp),
		}
		var createdAt time.Time
		var vars []byte
		dest := []interface{}{
			&env.Id,
			&env.ApplicationId,
			&env.Name,
			&env.Slug,
			&vars,
			&createdAt,
		}
		if err := rows.Scan(dest...); err != nil {
			return nil, errors.Wrap(err, "scanning db row")
		}
		setPbTimestamp(env.CreatedAt, createdAt)
		if err := json.Unmarshal(vars, &env.Vars); err != nil {
			return nil, errors.Wrap(err, "unmarshalling env vars JSON")
		}
		envs = append(envs, env)
	}
	if err := rows.Err(); err != nil {
		return nil, errors.Wrap(err, "iterating over db result set")
	}
	res := &pb.ListEnvironmentResponse{Environments: envs}
	return res, nil
}

func (s *server) GetEnvironment(ctx context.Context, req *pb.GetEnvironmentRequest) (*pb.Environment, error) {
	getSQL := "SELECT id, application_id, name, slug, vars, created_at FROM environments WHERE id = $1"
	var env pb.Environment
	var vars []byte
	dest := []interface{}{
		&env.Id,
		&env.ApplicationId,
		&env.Name,
		&env.Slug,
		&vars,
		&env.CreatedAt,
	}
	if err := s.db.QueryRow(getSQL, req.GetId()).Scan(dest...); err != nil {
		return nil, errors.Wrap(err, "scanning db row")
	}
	if err := json.Unmarshal(vars, &env.Vars); err != nil {
		return nil, errors.Wrap(err, "unmarshalling env vars JSON")
	}
	return &env, nil
}

func (s *server) CreateEnvironment(ctx context.Context, req *pb.Environment) (*pb.Environment, error) {
	// TODO(paulsmith): can we even do this in XO??
	insertSQL := "INSERT INTO environments (application_id, name, slug, vars) VALUES ($1, $2, $3, $4) RETURNING id, created_at"

	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(req.Vars); err != nil {
		return nil, errors.Wrap(err, "encoding env vars as JSON")
	}

	args := []interface{}{
		req.GetApplicationId(),
		req.GetName(),
		slugify(req.GetName()),
		buf.String(),
	}

	var id int

	if err := s.db.QueryRow(insertSQL, args...).Scan(&id, &req.CreatedAt); err != nil {
		return nil, errors.Wrap(err, "inserting in to db")
	}

	req.Id = int32(id)

	return req, nil
}

func (s *server) DestroyEnvironment(ctx context.Context, req *pb.DestroyEnvironmentRequest) (*pb.Empty, error) {
	deleteSQL := "DELETE FROM environments WHERE id = $1"
	if _, err := s.db.Exec(deleteSQL, req.GetId()); err != nil {
		return nil, errors.Wrap(err, "deleting row from db")
	}
	return &pb.Empty{}, nil
}

func (s *server) CopyEnvironment(context.Context, *pb.CopyEnvironmentRequest) (*pb.Environment, error) {
	return nil, nil
}
<commit_msg>Fix a few missed timestamp conversions<commit_after>package soapboxd

import (
	"bytes"
	"encoding/json"
	"time"

	pb "github.com/adhocteam/soapbox/proto"
	gpb "github.com/golang/protobuf/ptypes/timestamp"
	"github.com/pkg/errors"

	"golang.org/x/net/context"
)

func (s *server) ListEnvironments(ctx context.Context, req *pb.ListEnvironmentRequest) (*pb.ListEnvironmentResponse, error) {
	listSQL := "SELECT id, application_id, name, slug, vars, created_at FROM environments WHERE application_id = $1 ORDER BY id"
	rows, err := s.db.Query(listSQL, req.GetApplicationId())
	if err != nil {
		return nil, errors.Wrap(err, "querying db for environments")
	}
	var envs []*pb.Environment
	for rows.Next() {
		env := &pb.Environment{
			CreatedAt: new(gpb.Timestamp),
		}
		var createdAt time.Time
		var vars []byte
		dest := []interface{}{
			&env.Id,
			&env.ApplicationId,
			&env.Name,
			&env.Slug,
			&vars,
			&createdAt,
		}
		if err := rows.Scan(dest...); err != nil {
			return nil, errors.Wrap(err, "scanning db row")
		}
		setPbTimestamp(env.CreatedAt, createdAt)
		if err := json.Unmarshal(vars, &env.Vars); err != nil {
			return nil, errors.Wrap(err, "unmarshalling env vars JSON")
		}
		envs = append(envs, env)
	}
	if err := rows.Err(); err != nil {
		return nil, errors.Wrap(err, "iterating over db result set")
	}
	res := &pb.ListEnvironmentResponse{Environments: envs}
	return res, nil
}

func (s *server) GetEnvironment(ctx context.Context, req *pb.GetEnvironmentRequest) (*pb.Environment, error) {
	getSQL := "SELECT id, application_id, name, slug, vars, created_at FROM environments WHERE id = $1"
	env := &pb.Environment{
		CreatedAt: new(gpb.Timestamp),
	}
	var createdAt time.Time
	var vars []byte
	dest := []interface{}{
		&env.Id,
		&env.ApplicationId,
		&env.Name,
		&env.Slug,
		&vars,
		&createdAt,
	}
	if err := s.db.QueryRow(getSQL, req.GetId()).Scan(dest...); err != nil {
		return nil, errors.Wrap(err, "scanning db row")
	}
	setPbTimestamp(env.CreatedAt, createdAt)
	if err := json.Unmarshal(vars, &env.Vars); err != nil {
		return nil, errors.Wrap(err, "unmarshalling env vars JSON")
	}
	return env, nil
}

func (s *server) CreateEnvironment(ctx context.Context, req *pb.Environment) (*pb.Environment, error) {
	// TODO(paulsmith): can we even do this in XO??
	insertSQL := "INSERT INTO environments (application_id, name, slug, vars) VALUES ($1, $2, $3, $4) RETURNING id, created_at"

	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(req.Vars); err != nil {
		return nil, errors.Wrap(err, "encoding env vars as JSON")
	}

	args := []interface{}{
		req.GetApplicationId(),
		req.GetName(),
		slugify(req.GetName()),
		buf.String(),
	}

	var id int
	var createdAt time.Time

	if err := s.db.QueryRow(insertSQL, args...).Scan(&id, &createdAt); err != nil {
		return nil, errors.Wrap(err, "inserting in to db")
	}

	req.Id = int32(id)
	setPbTimestamp(req.CreatedAt, createdAt)

	return req, nil
}

func (s *server) DestroyEnvironment(ctx context.Context, req *pb.DestroyEnvironmentRequest) (*pb.Empty, error) {
	deleteSQL := "DELETE FROM environments WHERE id = $1"
	if _, err := s.db.Exec(deleteSQL, req.GetId()); err != nil {
		return nil, errors.Wrap(err, "deleting row from db")
	}
	return &pb.Empty{}, nil
}

func (s *server) CopyEnvironment(context.Context, *pb.CopyEnvironmentRequest) (*pb.Environment, error) {
	return nil, nil
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/cezarsa/form"
	"github.com/ghodss/yaml"
	"github.com/tsuru/gnuflag"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/event"
)

type EventList struct {
	fs     *gnuflag.FlagSet
	filter eventFilter
}

type eventFilter struct {
	KindName    string
	Target      string
	TargetValue string
	OwnerName   string
	Running     bool
}

func (f *eventFilter) queryString(client *cmd.Client) (url.Values, error) {
	values, err := form.EncodeToValues(f)
	if err != nil {
		return nil, err
	}
	for k, v := range values {
		values.Del(k)
		values[strings.ToLower(k)] = v
	}
	if !f.Running {
		values.Del("running")
	}
	return values, nil
}

func (f *eventFilter) flags(fs *gnuflag.FlagSet) {
	name := "Filter events by kind name"
	fs.StringVar(&f.KindName, "kind", "", name)
	fs.StringVar(&f.KindName, "k", "", name)
	name = "Filter events by target name"
	fs.StringVar(&f.Target, "target", "", name)
	fs.StringVar(&f.Target, "t", "", name)
	name = "Filter events by target value"
	fs.StringVar(&f.TargetValue, "target-value", "", name)
	fs.StringVar(&f.TargetValue, "v", "", name)
	name = "Filter events by owner name"
	fs.StringVar(&f.OwnerName, "owner", "", name)
	fs.StringVar(&f.OwnerName, "o", "", name)
	name = "Shows only currently running events"
	fs.BoolVar(&f.Running, "running", false, name)
	fs.BoolVar(&f.Running, "r", false, name)
}

func (c *EventList) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "event-list",
		Usage: "event-list [-k kindName]",
		Desc:  `Lists events possibly filtering them.`,
	}
}

func (c *EventList) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		c.filter.flags(c.fs)
	}
	return c.fs
}

func (c *EventList) Run(context *cmd.Context, client *cmd.Client) error {
	qs, err := c.filter.queryString(client)
	if err != nil {
		return err
	}
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events?%s", qs.Encode()))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusNoContent {
		return nil
	}
	result, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	var evts []event.Event
	err = json.Unmarshal(result, &evts)
	if err != nil {
		return fmt.Errorf("unable to unmarshal %q: %s", string(result), err)
	}
	return c.Show(evts, context)
}

var reEmailShort = regexp.MustCompile(`@.*$`)

func (c *EventList) Show(evts []event.Event, context *cmd.Context) error {
	tbl := cmd.NewTable()
	tbl.Headers = cmd.Row{"ID", "Start (duration)", "Success", "Owner", "Kind", "Target"}
	for i := range evts {
		evt := &evts[i]
		if evt.Target.Type == "container" {
			evt.Target.Value = evt.Target.Value[:12]
		}
		fullTarget := fmt.Sprintf("%s: %s", evt.Target.Type, evt.Target.Value)
		startFmt := evt.StartTime.Format(time.RFC822Z)
		owner := reEmailShort.ReplaceAllString(evt.Owner.Name, "@…")
		var ts, success string
		if evt.Running {
			ts = fmt.Sprintf("%s (…)", startFmt)
			success = "…"
		} else {
			ts = fmt.Sprintf("%s (%v)", startFmt, evt.EndTime.Sub(evt.StartTime))
			success = fmt.Sprintf("%v", evt.Error == "")
			if evt.CancelInfo.Canceled {
				success += " ✗"
			}
		}
		row := cmd.Row{evt.UniqueID.Hex(), ts, success, owner, evt.Kind.Name, fullTarget}
		var color string
		if evt.Running {
			color = "yellow"
		} else if evt.CancelInfo.Canceled {
			color = "magenta"
		} else if evt.Error != "" {
			color = "red"
		}
		if color != "" {
			for i, v := range row {
				if v != "" {
					row[i] = cmd.Colorfy(v, color, "", "")
				}
			}
		}
		tbl.AddRow(row)
	}
	fmt.Fprintf(context.Stdout, "%s", tbl.String())
	return nil
}

type EventInfo struct{}

func (c *EventInfo) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "event-info",
		Usage:   "event-info <event-id>",
		Desc:    `Show detailed information about one single event.`,
		MinArgs: 1,
		MaxArgs: 1,
	}
}

func (c *EventInfo) Run(context *cmd.Context, client *cmd.Client) error {
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events/%s", context.Args[0]))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	result, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	var evt event.Event
	err = json.Unmarshal(result, &evt)
	if err != nil {
		return fmt.Errorf("unable to unmarshal %q: %s", string(result), err)
	}
	return c.Show(evt, context)
}

func (c *EventInfo) Show(evt event.Event, context *cmd.Context) error {
	type item struct {
		label string
		value string
	}
	startFmt := evt.StartTime.Format(time.RFC822Z)
	var endFmt string
	if evt.Running {
		endFmt = fmt.Sprintf("running (%v)", time.Now().Sub(evt.StartTime))
	} else {
		endFmt = fmt.Sprintf("%s (%v)", evt.EndTime.Format(time.RFC822Z), evt.EndTime.Sub(evt.StartTime))
	}
	items := []item{
		{"ID", evt.UniqueID.Hex()},
		{"Start", startFmt},
		{"End", endFmt},
		{"Target", fmt.Sprintf("%s(%s)", evt.Target.Type, evt.Target.Value)},
		{"Kind", fmt.Sprintf("%s(%s)", evt.Kind.Type, evt.Kind.Name)},
		{"Owner", fmt.Sprintf("%s(%s)", evt.Owner.Type, evt.Owner.Name)},
	}
	successful := evt.Error == ""
	successfulStr := strconv.FormatBool(successful)
	if successful {
		if evt.Running {
			successfulStr = "…"
		}
		items = append(items, item{"Success", successfulStr})
	} else {
		redError := cmd.Colorfy(fmt.Sprintf("%q", evt.Error), "red", "", "")
		redSuccess := cmd.Colorfy(successfulStr, "red", "", "")
		items = append(items, []item{
			{"Success", redSuccess},
			{"Error", redError},
		}...)
	}
	items = append(items, []item{
		{"Cancelable", strconv.FormatBool(evt.Cancelable)},
		{"Canceled", strconv.FormatBool(evt.CancelInfo.Canceled)},
	}...)
	if evt.CancelInfo.Canceled {
		items = append(items, []item{
			{"  Reason", evt.CancelInfo.Reason},
			{"  By", evt.CancelInfo.Owner},
			{"  At", evt.CancelInfo.AckTime.Format(time.RFC822Z)},
		}...)
	}
	labels := []string{"Start", "End", "Other"}
	for i, fn := range []func(interface{}) error{evt.StartData, evt.EndData, evt.OtherData} {
		var data interface{}
		err := fn(&data)
		if err == nil && data != nil {
			str, err := yaml.Marshal(data)
			if err == nil {
				padded := padLines(string(str), "    ")
				items = append(items, item{fmt.Sprintf("%s Custom Data", labels[i]), "\n" + padded})
			}
		}
	}
	if evt.Log != "" {
		items = append(items, item{"Log", "\n" + padLines(evt.Log, "    ")})
	}
	var maxSz int
	for _, item := range items {
		sz := len(item.label)
		if len(item.value) > 0 && item.value[0] != '\n' && sz > maxSz {
			maxSz = sz
		}
	}
	for _, item := range items {
		count := (maxSz - len(item.label)) + 1
		var pad string
		if count > 0 && len(item.value) > 0 && item.value[0] != '\n' {
			pad = strings.Repeat(" ", count)
		}
		label := cmd.Colorfy(item.label+":", "cyan", "", "")
		fmt.Fprintf(context.Stdout, "%s%s%s\n", label, pad, item.value)
	}
	return nil
}

var rePadLines = regexp.MustCompile(`(?m)^(.+)`)

func padLines(s string, pad string) string {
	return rePadLines.ReplaceAllString(s, pad+`$1`)
}

type EventCancel struct {
	cmd.ConfirmationCommand
}

func (c *EventCancel) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "event-cancel",
		Usage:   "event-cancel <event-id> <reason> [-y]",
		Desc:    `Cancel a running event.`,
		MinArgs: 2,
	}
}

func (c *EventCancel) Run(context *cmd.Context, client *cmd.Client) error {
	if !c.Confirm(context, "Are you sure you want to cancel this event?") {
		return nil
	}
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events/%s/cancel", context.Args[0]))
	if err != nil {
		return err
	}
	v := url.Values{}
	v.Set("reason", strings.Join(context.Args[1:], " "))
	request, err := http.NewRequest("POST", u, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	fmt.Fprintln(context.Stdout, "Event successfully canceled.")
	return nil
}
<commit_msg>lint: avoid copying values containing mutexes<commit_after>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/cezarsa/form"
	"github.com/ghodss/yaml"
	"github.com/tsuru/gnuflag"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/event"
)

type EventList struct {
	fs     *gnuflag.FlagSet
	filter eventFilter
}

type eventFilter struct {
	KindName    string
	Target      string
	TargetValue string
	OwnerName   string
	Running     bool
}

func (f *eventFilter) queryString(client *cmd.Client) (url.Values, error) {
	values, err := form.EncodeToValues(f)
	if err != nil {
		return nil, err
	}
	for k, v := range values {
		values.Del(k)
		values[strings.ToLower(k)] = v
	}
	if !f.Running {
		values.Del("running")
	}
	return values, nil
}

func (f *eventFilter) flags(fs *gnuflag.FlagSet) {
	name := "Filter events by kind name"
	fs.StringVar(&f.KindName, "kind", "", name)
	fs.StringVar(&f.KindName, "k", "", name)
	name = "Filter events by target name"
	fs.StringVar(&f.Target, "target", "", name)
	fs.StringVar(&f.Target, "t", "", name)
	name = "Filter events by target value"
	fs.StringVar(&f.TargetValue, "target-value", "", name)
	fs.StringVar(&f.TargetValue, "v", "", name)
	name = "Filter events by owner name"
	fs.StringVar(&f.OwnerName, "owner", "", name)
	fs.StringVar(&f.OwnerName, "o", "", name)
	name = "Shows only currently running events"
	fs.BoolVar(&f.Running, "running", false, name)
	fs.BoolVar(&f.Running, "r", false, name)
}

func (c *EventList) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "event-list",
		Usage: "event-list [-k kindName]",
		Desc:  `Lists events possibly filtering them.`,
	}
}

func (c *EventList) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		c.filter.flags(c.fs)
	}
	return c.fs
}

func (c *EventList) Run(context *cmd.Context, client *cmd.Client) error {
	qs, err := c.filter.queryString(client)
	if err != nil {
		return err
	}
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events?%s", qs.Encode()))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusNoContent {
		return nil
	}
	result, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	var evts []event.Event
	err = json.Unmarshal(result, &evts)
	if err != nil {
		return fmt.Errorf("unable to unmarshal %q: %s", string(result), err)
	}
	return c.Show(evts, context)
}

var reEmailShort = regexp.MustCompile(`@.*$`)

func (c *EventList) Show(evts []event.Event, context *cmd.Context) error {
	tbl := cmd.NewTable()
	tbl.Headers = cmd.Row{"ID", "Start (duration)", "Success", "Owner", "Kind", "Target"}
	for i := range evts {
		evt := &evts[i]
		if evt.Target.Type == "container" {
			evt.Target.Value = evt.Target.Value[:12]
		}
		fullTarget := fmt.Sprintf("%s: %s", evt.Target.Type, evt.Target.Value)
		startFmt := evt.StartTime.Format(time.RFC822Z)
		owner := reEmailShort.ReplaceAllString(evt.Owner.Name, "@…")
		var ts, success string
		if evt.Running {
			ts = fmt.Sprintf("%s (…)", startFmt)
			success = "…"
		} else {
			ts = fmt.Sprintf("%s (%v)", startFmt, evt.EndTime.Sub(evt.StartTime))
			success = fmt.Sprintf("%v", evt.Error == "")
			if evt.CancelInfo.Canceled {
				success += " ✗"
			}
		}
		row := cmd.Row{evt.UniqueID.Hex(), ts, success, owner, evt.Kind.Name, fullTarget}
		var color string
		if evt.Running {
			color = "yellow"
		} else if evt.CancelInfo.Canceled {
			color = "magenta"
		} else if evt.Error != "" {
			color = "red"
		}
		if color != "" {
			for i, v := range row {
				if v != "" {
					row[i] = cmd.Colorfy(v, color, "", "")
				}
			}
		}
		tbl.AddRow(row)
	}
	fmt.Fprintf(context.Stdout, "%s", tbl.String())
	return nil
}

type EventInfo struct{}

func (c *EventInfo) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "event-info",
		Usage:   "event-info <event-id>",
		Desc:    `Show detailed information about one single event.`,
		MinArgs: 1,
		MaxArgs: 1,
	}
}

func (c *EventInfo) Run(context *cmd.Context, client *cmd.Client) error {
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events/%s", context.Args[0]))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	result, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	var evt event.Event
	err = json.Unmarshal(result, &evt)
	if err != nil {
		return fmt.Errorf("unable to unmarshal %q: %s", string(result), err)
	}
	return c.Show(&evt, context)
}

func (c *EventInfo) Show(evt *event.Event, context *cmd.Context) error {
	type item struct {
		label string
		value string
	}
	startFmt := evt.StartTime.Format(time.RFC822Z)
	var endFmt string
	if evt.Running {
		endFmt = fmt.Sprintf("running (%v)", time.Now().Sub(evt.StartTime))
	} else {
		endFmt = fmt.Sprintf("%s (%v)", evt.EndTime.Format(time.RFC822Z), evt.EndTime.Sub(evt.StartTime))
	}
	items := []item{
		{"ID", evt.UniqueID.Hex()},
		{"Start", startFmt},
		{"End", endFmt},
		{"Target", fmt.Sprintf("%s(%s)", evt.Target.Type, evt.Target.Value)},
		{"Kind", fmt.Sprintf("%s(%s)", evt.Kind.Type, evt.Kind.Name)},
		{"Owner", fmt.Sprintf("%s(%s)", evt.Owner.Type, evt.Owner.Name)},
	}
	successful := evt.Error == ""
	successfulStr := strconv.FormatBool(successful)
	if successful {
		if evt.Running {
			successfulStr = "…"
		}
		items = append(items, item{"Success", successfulStr})
	} else {
		redError := cmd.Colorfy(fmt.Sprintf("%q", evt.Error), "red", "", "")
		redSuccess := cmd.Colorfy(successfulStr, "red", "", "")
		items = append(items, []item{
			{"Success", redSuccess},
			{"Error", redError},
		}...)
	}
	items = append(items, []item{
		{"Cancelable", strconv.FormatBool(evt.Cancelable)},
		{"Canceled", strconv.FormatBool(evt.CancelInfo.Canceled)},
	}...)
	if evt.CancelInfo.Canceled {
		items = append(items, []item{
			{"  Reason", evt.CancelInfo.Reason},
			{"  By", evt.CancelInfo.Owner},
			{"  At", evt.CancelInfo.AckTime.Format(time.RFC822Z)},
		}...)
	}
	labels := []string{"Start", "End", "Other"}
	for i, fn := range []func(interface{}) error{evt.StartData, evt.EndData, evt.OtherData} {
		var data interface{}
		err := fn(&data)
		if err == nil && data != nil {
			str, err := yaml.Marshal(data)
			if err == nil {
				padded := padLines(string(str), "    ")
				items = append(items, item{fmt.Sprintf("%s Custom Data", labels[i]), "\n" + padded})
			}
		}
	}
	if evt.Log != "" {
		items = append(items, item{"Log", "\n" + padLines(evt.Log, "    ")})
	}
	var maxSz int
	for _, item := range items {
		sz := len(item.label)
		if len(item.value) > 0 && item.value[0] != '\n' && sz > maxSz {
			maxSz = sz
		}
	}
	for _, item := range items {
		count := (maxSz - len(item.label)) + 1
		var pad string
		if count > 0 && len(item.value) > 0 && item.value[0] != '\n' {
			pad = strings.Repeat(" ", count)
		}
		label := cmd.Colorfy(item.label+":", "cyan", "", "")
		fmt.Fprintf(context.Stdout, "%s%s%s\n", label, pad, item.value)
	}
	return nil
}

var rePadLines = regexp.MustCompile(`(?m)^(.+)`)

func padLines(s string, pad string) string {
	return rePadLines.ReplaceAllString(s, pad+`$1`)
}

type EventCancel struct {
	cmd.ConfirmationCommand
}

func (c *EventCancel) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "event-cancel",
		Usage:   "event-cancel <event-id> <reason> [-y]",
		Desc:    `Cancel a running event.`,
		MinArgs: 2,
	}
}

func (c *EventCancel) Run(context *cmd.Context, client *cmd.Client) error {
	if !c.Confirm(context, "Are you sure you want to cancel this event?") {
		return nil
	}
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events/%s/cancel", context.Args[0]))
	if err != nil {
		return err
	}
	v := url.Values{}
	v.Set("reason", strings.Join(context.Args[1:], " "))
	request, err := http.NewRequest("POST", u, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	fmt.Fprintln(context.Stdout, "Event successfully canceled.")
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	microerror "github.com/giantswarm/microkit/error"
)

type VPC struct {
	CidrBlock string
	Name      string
	id        string
	AWSEntity
}

func (v VPC) findExisting() (*ec2.Vpc, error) {
	vpcs, err := v.Clients.EC2.DescribeVpcs(&ec2.DescribeVpcsInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name: aws.String(fmt.Sprintf("tag:%s", tagKeyName)),
				Values: []*string{
					aws.String(v.Name),
				},
			},
		},
	})
	if err != nil {
		return nil, microerror.MaskAny(err)
	}

	if len(vpcs.Vpcs) < 1 {
		return nil, microerror.MaskAny(vpcFindError)
	}

	return vpcs.Vpcs[0], nil
}

func (v *VPC) checkIfExists() (bool, error) {
	_, err := v.findExisting()
	if IsVpcFindError(err) {
		return false, nil
	} else if err != nil {
		return false, microerror.MaskAny(err)
	}

	return true, nil
}

func (v *VPC) CreateIfNotExists() (bool, error) {
	exists, err := v.checkIfExists()
	if err != nil {
		return false, microerror.MaskAny(err)
	}

	if exists {
		return false, nil
	}

	if err := v.CreateOrFail(); err != nil {
		return false, microerror.MaskAny(err)
	}

	return true, nil
}

func (v *VPC) CreateOrFail() error {
	vpc, err := v.Clients.EC2.CreateVpc(&ec2.CreateVpcInput{
		CidrBlock: aws.String(v.CidrBlock),
	})
	if err != nil {
		return microerror.MaskAny(err)
	}
	vpcID := *vpc.Vpc.VpcId

	if err := v.Clients.EC2.WaitUntilVpcAvailable(&ec2.DescribeVpcsInput{
		VpcIds: []*string{
			aws.String(vpcID),
		},
	}); err != nil {
		return microerror.MaskAny(err)
	}

	if _, err := v.Clients.EC2.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{
			aws.String(vpcID),
		},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String(tagKeyName),
				Value: aws.String(v.Name),
			},
		},
	}); err != nil {
		return microerror.MaskAny(err)
	}

	// These attributes are required for a VPC with private Hosted Zones.
	if _, err := v.Clients.EC2.ModifyVpcAttribute(&ec2.ModifyVpcAttributeInput{
		EnableDnsHostnames: &ec2.AttributeBooleanValue{
			Value: aws.Bool(true),
		},
		EnableDnsSupport: &ec2.AttributeBooleanValue{
			Value: aws.Bool(true),
		},
		VpcId: aws.String(vpcID),
	}); err != nil {
		return microerror.MaskAny(err)
	}

	v.id = vpcID

	return nil
}

func (v *VPC) Delete() error {
	vpc, err := v.findExisting()
	if err != nil {
		return microerror.MaskAny(err)
	}

	if _, err := v.Clients.EC2.DeleteVpc(&ec2.DeleteVpcInput{
		VpcId: vpc.VpcId,
	}); err != nil {
		return microerror.MaskAny(err)
	}

	return nil
}

func (v VPC) GetID() (string, error) {
	if v.id != "" {
		return v.id, nil
	}

	vpc, err := v.findExisting()
	if err != nil {
		return "", microerror.MaskAny(err)
	}

	return *vpc.VpcId, nil
}
<commit_msg>vpc: attributes must be modified separately<commit_after>package aws

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	microerror "github.com/giantswarm/microkit/error"
)

type VPC struct {
	CidrBlock string
	Name      string
	id        string
	AWSEntity
}

func (v VPC) findExisting() (*ec2.Vpc, error) {
	vpcs, err := v.Clients.EC2.DescribeVpcs(&ec2.DescribeVpcsInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name: aws.String(fmt.Sprintf("tag:%s", tagKeyName)),
				Values: []*string{
					aws.String(v.Name),
				},
			},
		},
	})
	if err != nil {
		return nil, microerror.MaskAny(err)
	}

	if len(vpcs.Vpcs) < 1 {
		return nil, microerror.MaskAny(vpcFindError)
	}

	return vpcs.Vpcs[0], nil
}

func (v *VPC) checkIfExists() (bool, error) {
	_, err := v.findExisting()
	if IsVpcFindError(err) {
		return false, nil
	} else if err != nil {
		return false, microerror.MaskAny(err)
	}

	return true, nil
}

func (v *VPC) CreateIfNotExists() (bool, error) {
	exists, err := v.checkIfExists()
	if err != nil {
		return false, microerror.MaskAny(err)
	}

	if exists {
		return false, nil
	}

	if err := v.CreateOrFail(); err != nil {
		return false, microerror.MaskAny(err)
	}

	return true, nil
}

func (v *VPC) CreateOrFail() error {
	vpc, err := v.Clients.EC2.CreateVpc(&ec2.CreateVpcInput{
		CidrBlock: aws.String(v.CidrBlock),
	})
	if err != nil {
		return microerror.MaskAny(err)
	}
	vpcID := *vpc.Vpc.VpcId

	if err := v.Clients.EC2.WaitUntilVpcAvailable(&ec2.DescribeVpcsInput{
		VpcIds: []*string{
			aws.String(vpcID),
		},
	}); err != nil {
		return microerror.MaskAny(err)
	}

	if _, err := v.Clients.EC2.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{
			aws.String(vpcID),
		},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String(tagKeyName),
				Value: aws.String(v.Name),
			},
		},
	}); err != nil {
		return microerror.MaskAny(err)
	}

	// These attributes are required for a VPC with private Hosted Zones.
	if _, err := v.Clients.EC2.ModifyVpcAttribute(&ec2.ModifyVpcAttributeInput{
		EnableDnsHostnames: &ec2.AttributeBooleanValue{
			Value: aws.Bool(true),
		},
		VpcId: aws.String(vpcID),
	}); err != nil {
		return microerror.MaskAny(err)
	}

	if _, err := v.Clients.EC2.ModifyVpcAttribute(&ec2.ModifyVpcAttributeInput{
		EnableDnsSupport: &ec2.AttributeBooleanValue{
			Value: aws.Bool(true),
		},
		VpcId: aws.String(vpcID),
	}); err != nil {
		return microerror.MaskAny(err)
	}

	v.id = vpcID

	return nil
}

func (v *VPC) Delete() error {
	vpc, err := v.findExisting()
	if err != nil {
		return microerror.MaskAny(err)
	}

	if _, err := v.Clients.EC2.DeleteVpc(&ec2.DeleteVpcInput{
		VpcId: vpc.VpcId,
	}); err != nil {
		return microerror.MaskAny(err)
	}

	return nil
}

func (v VPC) GetID() (string, error) {
	if v.id != "" {
		return v.id, nil
	}

	vpc, err := v.findExisting()
	if err != nil {
		return "", microerror.MaskAny(err)
	}

	return *vpc.VpcId, nil
}
<|endoftext|>


<commit_before>/*
Current limitations:

	- GSS-API authentication is not supported
	- only SOCKS version 5 is supported
	- TCP bind and UDP not yet supported

Example http client over SOCKS5:

	proxy := &socks.Proxy{"127.0.0.1:1080"}
	tr := &http.Transport{
		Dial: func(net, addr string) (net.Conn, error) {
			return proxy.Dial(net, addr)
		},
	}
	client := &http.Client{Transport: tr}
	resp, err := client.Get("https://example.com")
*/
package socks

import (
	"errors"
	"io"
	"net"
	"strconv"
)

const (
	protocolVersion = 5

	defaultPort = 1080

	authNone             = 0
	authGssApi           = 1
	authUsernamePassword = 2
	authUnavailable      = 0xff

	commandTcpConnect   = 1
	commandTcpBind      = 2
	commandUdpAssociate = 3

	addressTypeIPv4   = 1
	addressTypeDomain = 3
	addressTypeIPv6   = 4

	statusRequestGranted          = 0
	statusGeneralFailure          = 1
	statusConnectionNotAllowed    = 2
	statusNetworkUnreachable      = 3
	statusHostUnreachable         = 4
	statusConnectionRefused       = 5
	statusTtlExpired              = 6
	statusCommandNotSupport       = 7
	statusAddressTypeNotSupported = 8
)

var (
	ErrAuthFailed             = errors.New("authentication failed")
	ErrInvalidProxyResponse   = errors.New("invalid proxy response")
	ErrNoAcceptableAuthMethod = errors.New("no acceptable authentication method")

	statusErrors = map[byte]error{
		statusGeneralFailure:          errors.New("general failure"),
		statusConnectionNotAllowed:    errors.New("connection not allowed by ruleset"),
		statusNetworkUnreachable:      errors.New("network unreachable"),
		statusHostUnreachable:         errors.New("host unreachable"),
		statusConnectionRefused:       errors.New("connection refused by destination host"),
		statusTtlExpired:              errors.New("TTL expired"),
		statusCommandNotSupport:       errors.New("command not supported / protocol error"),
		statusAddressTypeNotSupported: errors.New("address type not supported"),
	}
)

type Proxy struct {
	Addr     string
	Username string
	Password string
}

func (p *Proxy) Dial(network, addr string) (net.Conn, error) {
	host, strPort, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(strPort)
	if err != nil {
		return nil, err
	}

	conn, err := net.Dial("tcp", p.Addr)
	if err != nil {
		return nil, err
	}

	buf := make([]byte, 16)

	buf[0] = protocolVersion
	if p.Username != "" {
		buf = buf[:4]
		buf[1] = 2 // num auth methods
		buf[2] = authNone
		buf[3] = authUsernamePassword
	} else {
		buf = buf[:3]
		buf[1] = 1 // num auth methods
		buf[2] = authNone
	}

	_, err = conn.Write(buf)
	if err != nil {
		conn.Close()
		return nil, err
	}

	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		conn.Close()
		return nil, err
	}
	if buf[0] != protocolVersion {
		conn.Close()
		return nil, ErrInvalidProxyResponse
	}
	err = nil
	switch buf[1] {
	default:
		err = ErrInvalidProxyResponse
	case authUnavailable:
		err = ErrNoAcceptableAuthMethod
	case authGssApi:
		err = ErrNoAcceptableAuthMethod
	case authUsernamePassword:
		buf = buf[:3+len(p.Username)+len(p.Password)]
		buf[0] = 1 // version
		buf[1] = byte(len(p.Username))
		copy(buf[2:], p.Username)
		buf[2+len(p.Username)] = byte(len(p.Password))
		copy(buf[3+len(p.Username):], p.Password)
		if _, err = conn.Write(buf); err != nil {
			conn.Close()
			return nil, err
		}
		if _, err = io.ReadFull(conn, buf[:2]); err != nil {
			conn.Close()
			return nil, err
		}
		if buf[0] != 1 { // version
			err = ErrInvalidProxyResponse
		} else if buf[1] != 0 { // 0 = succes, else auth failed
			err = ErrAuthFailed
		}
	case authNone:
		// Do nothing
	}
	if err != nil {
		conn.Close()
		return nil, err
	}

	buf = buf[:7+len(host)]
	buf[0] = protocolVersion
	buf[1] = commandTcpConnect
	buf[2] = 0 // reserved
	buf[3] = addressTypeDomain
	buf[4] = byte(len(host))
	copy(buf[5:], host)
	buf[5+len(host)] = byte(port >> 8)
	buf[6+len(host)] = byte(port & 0xff)
	if _, err := conn.Write(buf); err != nil {
		conn.Close()
		return nil, err
	}

	if _, err := io.ReadFull(conn, buf[:4]); err != nil {
		conn.Close()
		return nil, err
	}

	if buf[0] != protocolVersion {
		conn.Close()
		return nil, ErrInvalidProxyResponse
	}

	if buf[1] != statusRequestGranted {
		conn.Close()
		err := statusErrors[buf[1]]
		if err == nil {
			err = ErrInvalidProxyResponse
		}
		return nil, err
	}

	paddr := &proxiedAddr{net: network}

	switch buf[3] {
	default:
		conn.Close()
		return nil, ErrInvalidProxyResponse
	case addressTypeIPv4:
		if _, err := io.ReadFull(conn, buf[:4]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = net.IP(buf).String()
	case addressTypeIPv6:
		if _, err := io.ReadFull(conn, buf[:16]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = net.IP(buf).String()
	case addressTypeDomain:
		if _, err := io.ReadFull(conn, buf[:1]); err != nil {
			conn.Close()
			return nil, err
		}
		domainLen := buf[0]
		if _, err := io.ReadFull(conn, buf[:domainLen]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = string(buf[:domainLen])
	}

	// Port
	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		conn.Close()
		return nil, err
	}
	paddr.port = int(buf[0])<<8 | int(buf[1])

	return &proxiedConn{
		conn:       conn,
		boundAddr:  paddr,
		remoteAddr: &proxiedAddr{network, host, port},
	}, nil
}
<commit_msg>Fix an issue with long usernames and/or passwords<commit_after>/*
Current limitations:

	- GSS-API authentication is not supported
	- only SOCKS version 5 is supported
	- TCP bind and UDP not yet supported

Example http client over SOCKS5:

	proxy := &socks.Proxy{"127.0.0.1:1080"}
	tr := &http.Transport{
		Dial: func(net, addr string) (net.Conn, error) {
			return proxy.Dial(net, addr)
		},
	}
	client := &http.Client{Transport: tr}
	resp, err := client.Get("https://example.com")
*/
package socks

import (
	"errors"
	"io"
	"net"
	"strconv"
)

const (
	protocolVersion = 5

	defaultPort = 1080

	authNone             = 0
	authGssApi           = 1
	authUsernamePassword = 2
	authUnavailable      = 0xff

	commandTcpConnect   = 1
	commandTcpBind      = 2
	commandUdpAssociate = 3

	addressTypeIPv4   = 1
	addressTypeDomain = 3
	addressTypeIPv6   = 4

	statusRequestGranted          = 0
	statusGeneralFailure          = 1
	statusConnectionNotAllowed    = 2
	statusNetworkUnreachable      = 3
	statusHostUnreachable         = 4
	statusConnectionRefused       = 5
	statusTtlExpired              = 6
	statusCommandNotSupport       = 7
	statusAddressTypeNotSupported = 8
)

var (
	ErrAuthFailed             = errors.New("authentication failed")
	ErrInvalidProxyResponse   = errors.New("invalid proxy response")
	ErrNoAcceptableAuthMethod = errors.New("no acceptable authentication method")

	statusErrors = map[byte]error{
		statusGeneralFailure:          errors.New("general failure"),
		statusConnectionNotAllowed:    errors.New("connection not allowed by ruleset"),
		statusNetworkUnreachable:      errors.New("network unreachable"),
		statusHostUnreachable:         errors.New("host unreachable"),
		statusConnectionRefused:       errors.New("connection refused by destination host"),
		statusTtlExpired:              errors.New("TTL expired"),
		statusCommandNotSupport:       errors.New("command not supported / protocol error"),
		statusAddressTypeNotSupported: errors.New("address type not supported"),
	}
)

type Proxy struct {
	Addr     string
	Username string
	Password string
}

func (p *Proxy) Dial(network, addr string) (net.Conn, error) {
	host, strPort, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(strPort)
	if err != nil {
		return nil, err
	}

	conn, err := net.Dial("tcp", p.Addr)
	if err != nil {
		return nil, err
	}

	buf := make([]byte, 16+len(p.Username)+len(p.Password))

	// Initial greeting

	buf[0] = protocolVersion
	if p.Username != "" {
		buf = buf[:4]
		buf[1] = 2 // num auth methods
		buf[2] = authNone
		buf[3] = authUsernamePassword
	} else {
		buf = buf[:3]
		buf[1] = 1 // num auth methods
		buf[2] = authNone
	}

	_, err = conn.Write(buf)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Server's auth choice

	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		conn.Close()
		return nil, err
	}
	if buf[0] != protocolVersion {
		conn.Close()
		return nil, ErrInvalidProxyResponse
	}
	err = nil
	switch buf[1] {
	default:
		err = ErrInvalidProxyResponse
	case authUnavailable:
		err = ErrNoAcceptableAuthMethod
	case authGssApi:
		err = ErrNoAcceptableAuthMethod
	case authUsernamePassword:
		buf = buf[:3+len(p.Username)+len(p.Password)]
		buf[0] = 1 // version
		buf[1] = byte(len(p.Username))
		copy(buf[2:], p.Username)
		buf[2+len(p.Username)] = byte(len(p.Password))
		copy(buf[3+len(p.Username):], p.Password)
		if _, err = conn.Write(buf); err != nil {
			conn.Close()
			return nil, err
		}
		if _, err = io.ReadFull(conn, buf[:2]); err != nil {
			conn.Close()
			return nil, err
		}
		if buf[0] != 1 { // version
			err = ErrInvalidProxyResponse
		} else if buf[1] != 0 { // 0 = succes, else auth failed
			err = ErrAuthFailed
		}
	case authNone:
		// Do nothing
	}
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Command / connection request

	buf = buf[:7+len(host)]
	buf[0] = protocolVersion
	buf[1] = commandTcpConnect
	buf[2] = 0 // reserved
	buf[3] = addressTypeDomain
	buf[4] = byte(len(host))
	copy(buf[5:], host)
	buf[5+len(host)] = byte(port >> 8)
	buf[6+len(host)] = byte(port & 0xff)
	if _, err := conn.Write(buf); err != nil {
		conn.Close()
		return nil, err
	}

	// Server response

	if _, err := io.ReadFull(conn, buf[:4]); err != nil {
		conn.Close()
		return nil, err
	}

	if buf[0] != protocolVersion {
		conn.Close()
		return nil, ErrInvalidProxyResponse
	}

	if buf[1] != statusRequestGranted {
		conn.Close()
		err := statusErrors[buf[1]]
		if err == nil {
			err = ErrInvalidProxyResponse
		}
		return nil, err
	}

	paddr := &proxiedAddr{net: network}

	switch buf[3] {
	default:
		conn.Close()
		return nil, ErrInvalidProxyResponse
	case addressTypeIPv4:
		if _, err := io.ReadFull(conn, buf[:4]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = net.IP(buf).String()
	case addressTypeIPv6:
		if _, err := io.ReadFull(conn, buf[:16]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = net.IP(buf).String()
	case addressTypeDomain:
		if _, err := io.ReadFull(conn, buf[:1]); err != nil {
			conn.Close()
			return nil, err
		}
		domainLen := buf[0]
		if _, err := io.ReadFull(conn, buf[:domainLen]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = string(buf[:domainLen])
	}

	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		conn.Close()
		return nil, err
	}
	paddr.port = int(buf[0])<<8 | int(buf[1])

	return &proxiedConn{
		conn:       conn,
		boundAddr:  paddr,
		remoteAddr: &proxiedAddr{network, host, port},
	}, nil
}
<|endoftext|>