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

// Package syntax parses regular expressions into syntax trees.
// WORK IN PROGRESS.
package syntax

// Note to implementers:
// In this package, re is always a *Regexp and r is always a rune.

import (
	"bytes"
	"strconv"
	"strings"
	"unicode"
)

// A Regexp is a node in a regular expression syntax tree.
type Regexp struct {
	Op       Op // operator
	Flags    Flags
	Sub      []*Regexp  // subexpressions, if any
	Sub0     [1]*Regexp // storage for short Sub
	Rune     []rune     // matched runes, for OpLiteral, OpCharClass
	Rune0    [2]rune    // storage for short Rune
	Min, Max int        // min, max for OpRepeat
	Cap      int        // capturing index, for OpCapture
	Name     string     // capturing name, for OpCapture
}

// An Op is a single regular expression operator.
type Op uint8

// Operators are listed in precedence order, tightest binding to weakest.
// Character class operators are listed simplest to most complex
// (OpLiteral, OpCharClass, OpAnyCharNotNL, OpAnyChar).

const (
	OpNoMatch        Op = 1 + iota // matches no strings
	OpEmptyMatch                   // matches empty string
	OpLiteral                      // matches Runes sequence
	OpCharClass                    // matches Runes interpreted as range pair list
	OpAnyCharNotNL                 // matches any character
	OpAnyChar                      // matches any character
	OpBeginLine                    // matches empty string at beginning of line
	OpEndLine                      // matches empty string at end of line
	OpBeginText                    // matches empty string at beginning of text
	OpEndText                      // matches empty string at end of text
	OpWordBoundary                 // matches word boundary `\b`
	OpNoWordBoundary               // matches word non-boundary `\B`
	OpCapture                      // capturing subexpression with index Cap, optional name Name
	OpStar                         // matches Sub[0] zero or more times
	OpPlus                         // matches Sub[0] one or more times
	OpQuest                        // matches Sub[0] zero or one times
	OpRepeat                       // matches Sub[0] at least Min times, at most Max (Max == -1 is no limit)
	OpConcat                       // matches concatenation of Subs
	OpAlternate                    // matches alternation of Subs
)

const opPseudo Op = 128 // where pseudo-ops start

// Equal returns true if x and y have identical structure.
func (x *Regexp) Equal(y *Regexp) bool {
	if x == nil || y == nil {
		return x == y
	}
	if x.Op != y.Op {
		return false
	}
	switch x.Op {
	case OpEndText:
		// The parse flags remember whether this is \z or \Z.
		if x.Flags&WasDollar != y.Flags&WasDollar {
			return false
		}

	case OpLiteral, OpCharClass:
		if len(x.Rune) != len(y.Rune) {
			return false
		}
		for i, r := range x.Rune {
			if r != y.Rune[i] {
				return false
			}
		}

	case OpAlternate, OpConcat:
		if len(x.Sub) != len(y.Sub) {
			return false
		}
		for i, sub := range x.Sub {
			if !sub.Equal(y.Sub[i]) {
				return false
			}
		}

	case OpStar, OpPlus, OpQuest:
		if x.Flags&NonGreedy != y.Flags&NonGreedy || !x.Sub[0].Equal(y.Sub[0]) {
			return false
		}

	case OpRepeat:
		if x.Flags&NonGreedy != y.Flags&NonGreedy || x.Min != y.Min || x.Max != y.Max || !x.Sub[0].Equal(y.Sub[0]) {
			return false
		}

	case OpCapture:
		if x.Cap != y.Cap || x.Name != y.Name || !x.Sub[0].Equal(y.Sub[0]) {
			return false
		}
	}
	return true
}

// writeRegexp writes the Perl syntax for the regular expression re to b.
func writeRegexp(b *bytes.Buffer, re *Regexp) {
	switch re.Op {
	default:
		b.WriteString("<invalid op" + strconv.Itoa(int(re.Op)) + ">")
	case OpNoMatch:
		b.WriteString(`[^\x00-\x{10FFFF}]`)
	case OpEmptyMatch:
		b.WriteString(`(?:)`)
	case OpLiteral:
		if re.Flags&FoldCase != 0 {
			b.WriteString(`(?i:`)
		}
		for _, r := range re.Rune {
			escape(b, r, false)
		}
		if re.Flags&FoldCase != 0 {
			b.WriteString(`)`)
		}
	case OpCharClass:
		if len(re.Rune)%2 != 0 {
			b.WriteString(`[invalid char class]`)
			break
		}
		b.WriteRune('[')
		if len(re.Rune) == 0 {
			b.WriteString(`^\x00-\x{10FFFF}`)
		} else if re.Rune[0] == 0 && re.Rune[len(re.Rune)-1] == unicode.MaxRune {
			// Contains 0 and MaxRune.  Probably a negated class.
			// Print the gaps.
			b.WriteRune('^')
			for i := 1; i < len(re.Rune)-1; i += 2 {
				lo, hi := re.Rune[i]+1, re.Rune[i+1]-1
				escape(b, lo, lo == '-')
				if lo != hi {
					b.WriteRune('-')
					escape(b, hi, hi == '-')
				}
			}
		} else {
			for i := 0; i < len(re.Rune); i += 2 {
				lo, hi := re.Rune[i], re.Rune[i+1]
				escape(b, lo, lo == '-')
				if lo != hi {
					b.WriteRune('-')
					escape(b, hi, hi == '-')
				}
			}
		}
		b.WriteRune(']')
	case OpAnyCharNotNL:
		b.WriteString(`(?-s:.)`)
	case OpAnyChar:
		b.WriteString(`(?s:.)`)
	case OpBeginLine:
		b.WriteRune('^')
	case OpEndLine:
		b.WriteRune('$')
	case OpBeginText:
		b.WriteString(`\A`)
	case OpEndText:
		if re.Flags&WasDollar != 0 {
			b.WriteString(`(?-m:$)`)
		} else {
			b.WriteString(`\z`)
		}
	case OpWordBoundary:
		b.WriteString(`\b`)
	case OpNoWordBoundary:
		b.WriteString(`\B`)
	case OpCapture:
		if re.Name != "" {
			b.WriteString(`(?P<`)
			b.WriteString(re.Name)
			b.WriteRune('>')
		} else {
			b.WriteRune('(')
		}
		if re.Sub[0].Op != OpEmptyMatch {
			writeRegexp(b, re.Sub[0])
		}
		b.WriteRune(')')
	case OpStar, OpPlus, OpQuest, OpRepeat:
		if sub := re.Sub[0]; sub.Op > OpCapture || sub.Op == OpLiteral && len(sub.Rune) > 1 {
			b.WriteString(`(?:`)
			writeRegexp(b, sub)
			b.WriteString(`)`)
		} else {
			writeRegexp(b, sub)
		}
		switch re.Op {
		case OpStar:
			b.WriteRune('*')
		case OpPlus:
			b.WriteRune('+')
		case OpQuest:
			b.WriteRune('?')
		case OpRepeat:
			b.WriteRune('{')
			b.WriteString(strconv.Itoa(re.Min))
			if re.Max != re.Min {
				b.WriteRune(',')
				if re.Max >= 0 {
					b.WriteString(strconv.Itoa(re.Max))
				}
			}
			b.WriteRune('}')
		}
		if re.Flags&NonGreedy != 0 {
			b.WriteRune('?')
		}
	case OpConcat:
		for _, sub := range re.Sub {
			if sub.Op == OpAlternate {
				b.WriteString(`(?:`)
				writeRegexp(b, sub)
				b.WriteString(`)`)
			} else {
				writeRegexp(b, sub)
			}
		}
	case OpAlternate:
		for i, sub := range re.Sub {
			if i > 0 {
				b.WriteRune('|')
			}
			writeRegexp(b, sub)
		}
	}
}

func (re *Regexp) String() string {
	var b bytes.Buffer
	writeRegexp(&b, re)
	return b.String()
}

const meta = `\.+*?()|[]{}^$`

func escape(b *bytes.Buffer, r rune, force bool) {
	if unicode.IsPrint(r) {
		if strings.IndexRune(meta, r) >= 0 || force {
			b.WriteRune('\\')
		}
		b.WriteRune(r)
		return
	}

	switch r {
	case '\a':
		b.WriteString(`\a`)
	case '\f':
		b.WriteString(`\f`)
	case '\n':
		b.WriteString(`\n`)
	case '\r':
		b.WriteString(`\r`)
	case '\t':
		b.WriteString(`\t`)
	case '\v':
		b.WriteString(`\v`)
	default:
		if r < 0x100 {
			b.WriteString(`\x`)
			s := strconv.FormatInt(int64(r), 16)
			if len(s) == 1 {
				b.WriteRune('0')
			}
			b.WriteString(s)
			break
		}
		b.WriteString(`\x{`)
		b.WriteString(strconv.FormatInt(int64(r), 16))
		b.WriteString(`}`)
	}
}

// MaxCap walks the regexp to find the maximum capture index.
func (re *Regexp) MaxCap() int {
	m := 0
	if re.Op == OpCapture {
		m = re.Cap
	}
	for _, sub := range re.Sub {
		if n := sub.MaxCap(); m < n {
			m = n
		}
	}
	return m
}

// CapNames walks the regexp to find the names of capturing groups.
func (re *Regexp) CapNames() []string {
	names := make([]string, re.MaxCap()+1)
	re.capNames(names)
	return names
}

func (re *Regexp) capNames(names []string) {
	if re.Op == OpCapture {
		names[re.Cap] = re.Name
	}
	for _, sub := range re.Sub {
		sub.capNames(names)
	}
}
<commit_msg>regexp/syntax: delete old package comment<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package syntax

// Note to implementers:
// In this package, re is always a *Regexp and r is always a rune.

import (
	"bytes"
	"strconv"
	"strings"
	"unicode"
)

// A Regexp is a node in a regular expression syntax tree.
type Regexp struct {
	Op       Op // operator
	Flags    Flags
	Sub      []*Regexp  // subexpressions, if any
	Sub0     [1]*Regexp // storage for short Sub
	Rune     []rune     // matched runes, for OpLiteral, OpCharClass
	Rune0    [2]rune    // storage for short Rune
	Min, Max int        // min, max for OpRepeat
	Cap      int        // capturing index, for OpCapture
	Name     string     // capturing name, for OpCapture
}

// An Op is a single regular expression operator.
type Op uint8

// Operators are listed in precedence order, tightest binding to weakest.
// Character class operators are listed simplest to most complex
// (OpLiteral, OpCharClass, OpAnyCharNotNL, OpAnyChar).

const (
	OpNoMatch        Op = 1 + iota // matches no strings
	OpEmptyMatch                   // matches empty string
	OpLiteral                      // matches Runes sequence
	OpCharClass                    // matches Runes interpreted as range pair list
	OpAnyCharNotNL                 // matches any character
	OpAnyChar                      // matches any character
	OpBeginLine                    // matches empty string at beginning of line
	OpEndLine                      // matches empty string at end of line
	OpBeginText                    // matches empty string at beginning of text
	OpEndText                      // matches empty string at end of text
	OpWordBoundary                 // matches word boundary `\b`
	OpNoWordBoundary               // matches word non-boundary `\B`
	OpCapture                      // capturing subexpression with index Cap, optional name Name
	OpStar                         // matches Sub[0] zero or more times
	OpPlus                         // matches Sub[0] one or more times
	OpQuest                        // matches Sub[0] zero or one times
	OpRepeat                       // matches Sub[0] at least Min times, at most Max (Max == -1 is no limit)
	OpConcat                       // matches concatenation of Subs
	OpAlternate                    // matches alternation of Subs
)

const opPseudo Op = 128 // where pseudo-ops start

// Equal returns true if x and y have identical structure.
func (x *Regexp) Equal(y *Regexp) bool {
	if x == nil || y == nil {
		return x == y
	}
	if x.Op != y.Op {
		return false
	}
	switch x.Op {
	case OpEndText:
		// The parse flags remember whether this is \z or \Z.
		if x.Flags&WasDollar != y.Flags&WasDollar {
			return false
		}

	case OpLiteral, OpCharClass:
		if len(x.Rune) != len(y.Rune) {
			return false
		}
		for i, r := range x.Rune {
			if r != y.Rune[i] {
				return false
			}
		}

	case OpAlternate, OpConcat:
		if len(x.Sub) != len(y.Sub) {
			return false
		}
		for i, sub := range x.Sub {
			if !sub.Equal(y.Sub[i]) {
				return false
			}
		}

	case OpStar, OpPlus, OpQuest:
		if x.Flags&NonGreedy != y.Flags&NonGreedy || !x.Sub[0].Equal(y.Sub[0]) {
			return false
		}

	case OpRepeat:
		if x.Flags&NonGreedy != y.Flags&NonGreedy || x.Min != y.Min || x.Max != y.Max || !x.Sub[0].Equal(y.Sub[0]) {
			return false
		}

	case OpCapture:
		if x.Cap != y.Cap || x.Name != y.Name || !x.Sub[0].Equal(y.Sub[0]) {
			return false
		}
	}
	return true
}

// writeRegexp writes the Perl syntax for the regular expression re to b.
func writeRegexp(b *bytes.Buffer, re *Regexp) {
	switch re.Op {
	default:
		b.WriteString("<invalid op" + strconv.Itoa(int(re.Op)) + ">")
	case OpNoMatch:
		b.WriteString(`[^\x00-\x{10FFFF}]`)
	case OpEmptyMatch:
		b.WriteString(`(?:)`)
	case OpLiteral:
		if re.Flags&FoldCase != 0 {
			b.WriteString(`(?i:`)
		}
		for _, r := range re.Rune {
			escape(b, r, false)
		}
		if re.Flags&FoldCase != 0 {
			b.WriteString(`)`)
		}
	case OpCharClass:
		if len(re.Rune)%2 != 0 {
			b.WriteString(`[invalid char class]`)
			break
		}
		b.WriteRune('[')
		if len(re.Rune) == 0 {
			b.WriteString(`^\x00-\x{10FFFF}`)
		} else if re.Rune[0] == 0 && re.Rune[len(re.Rune)-1] == unicode.MaxRune {
			// Contains 0 and MaxRune.  Probably a negated class.
			// Print the gaps.
			b.WriteRune('^')
			for i := 1; i < len(re.Rune)-1; i += 2 {
				lo, hi := re.Rune[i]+1, re.Rune[i+1]-1
				escape(b, lo, lo == '-')
				if lo != hi {
					b.WriteRune('-')
					escape(b, hi, hi == '-')
				}
			}
		} else {
			for i := 0; i < len(re.Rune); i += 2 {
				lo, hi := re.Rune[i], re.Rune[i+1]
				escape(b, lo, lo == '-')
				if lo != hi {
					b.WriteRune('-')
					escape(b, hi, hi == '-')
				}
			}
		}
		b.WriteRune(']')
	case OpAnyCharNotNL:
		b.WriteString(`(?-s:.)`)
	case OpAnyChar:
		b.WriteString(`(?s:.)`)
	case OpBeginLine:
		b.WriteRune('^')
	case OpEndLine:
		b.WriteRune('$')
	case OpBeginText:
		b.WriteString(`\A`)
	case OpEndText:
		if re.Flags&WasDollar != 0 {
			b.WriteString(`(?-m:$)`)
		} else {
			b.WriteString(`\z`)
		}
	case OpWordBoundary:
		b.WriteString(`\b`)
	case OpNoWordBoundary:
		b.WriteString(`\B`)
	case OpCapture:
		if re.Name != "" {
			b.WriteString(`(?P<`)
			b.WriteString(re.Name)
			b.WriteRune('>')
		} else {
			b.WriteRune('(')
		}
		if re.Sub[0].Op != OpEmptyMatch {
			writeRegexp(b, re.Sub[0])
		}
		b.WriteRune(')')
	case OpStar, OpPlus, OpQuest, OpRepeat:
		if sub := re.Sub[0]; sub.Op > OpCapture || sub.Op == OpLiteral && len(sub.Rune) > 1 {
			b.WriteString(`(?:`)
			writeRegexp(b, sub)
			b.WriteString(`)`)
		} else {
			writeRegexp(b, sub)
		}
		switch re.Op {
		case OpStar:
			b.WriteRune('*')
		case OpPlus:
			b.WriteRune('+')
		case OpQuest:
			b.WriteRune('?')
		case OpRepeat:
			b.WriteRune('{')
			b.WriteString(strconv.Itoa(re.Min))
			if re.Max != re.Min {
				b.WriteRune(',')
				if re.Max >= 0 {
					b.WriteString(strconv.Itoa(re.Max))
				}
			}
			b.WriteRune('}')
		}
		if re.Flags&NonGreedy != 0 {
			b.WriteRune('?')
		}
	case OpConcat:
		for _, sub := range re.Sub {
			if sub.Op == OpAlternate {
				b.WriteString(`(?:`)
				writeRegexp(b, sub)
				b.WriteString(`)`)
			} else {
				writeRegexp(b, sub)
			}
		}
	case OpAlternate:
		for i, sub := range re.Sub {
			if i > 0 {
				b.WriteRune('|')
			}
			writeRegexp(b, sub)
		}
	}
}

func (re *Regexp) String() string {
	var b bytes.Buffer
	writeRegexp(&b, re)
	return b.String()
}

const meta = `\.+*?()|[]{}^$`

func escape(b *bytes.Buffer, r rune, force bool) {
	if unicode.IsPrint(r) {
		if strings.IndexRune(meta, r) >= 0 || force {
			b.WriteRune('\\')
		}
		b.WriteRune(r)
		return
	}

	switch r {
	case '\a':
		b.WriteString(`\a`)
	case '\f':
		b.WriteString(`\f`)
	case '\n':
		b.WriteString(`\n`)
	case '\r':
		b.WriteString(`\r`)
	case '\t':
		b.WriteString(`\t`)
	case '\v':
		b.WriteString(`\v`)
	default:
		if r < 0x100 {
			b.WriteString(`\x`)
			s := strconv.FormatInt(int64(r), 16)
			if len(s) == 1 {
				b.WriteRune('0')
			}
			b.WriteString(s)
			break
		}
		b.WriteString(`\x{`)
		b.WriteString(strconv.FormatInt(int64(r), 16))
		b.WriteString(`}`)
	}
}

// MaxCap walks the regexp to find the maximum capture index.
func (re *Regexp) MaxCap() int {
	m := 0
	if re.Op == OpCapture {
		m = re.Cap
	}
	for _, sub := range re.Sub {
		if n := sub.MaxCap(); m < n {
			m = n
		}
	}
	return m
}

// CapNames walks the regexp to find the names of capturing groups.
func (re *Regexp) CapNames() []string {
	names := make([]string, re.MaxCap()+1)
	re.capNames(names)
	return names
}

func (re *Regexp) capNames(names []string) {
	if re.Op == OpCapture {
		names[re.Cap] = re.Name
	}
	for _, sub := range re.Sub {
		sub.capNames(names)
	}
}
<|endoftext|>


<commit_before>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/netx"
	"github.com/getlantern/preconn"
	"github.com/getlantern/proxy/filters"
)

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			safeClose(downstream)
			err = errors.New("Recovered from panic handling connection: %v", p)
		}
	}()

	err = proxy.handle(ctx, downstreamIn, downstream, nil)
	return
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection: %v", p)
		}
	}()

	conn.Close()
}

func (proxy *proxy) logInitialReadError(downstream net.Conn, err error) error {
	loc := downstream.LocalAddr()
	l := ""
	if loc != nil {
		l = loc.String()
	}
	rem := downstream.RemoteAddr()
	r := ""
	if rem != nil {
		r = rem.String()
	}
	return log.Errorf("Error in initial ReadRequest from %v to %v: %v", r, l, err)
}

func (proxy *proxy) handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn, upstream net.Conn) error {
	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
	}()

	downstreamBuffered := bufio.NewReader(downstreamIn)
	fctx := filters.WrapContext(withAwareConn(ctx), downstream)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = downstream.RemoteAddr().String()
		}
		if origURLScheme(ctx) == "" {
			fctx = fctx.
				WithValue(ctxKeyOrigURLScheme, req.URL.Scheme).
				WithValue(ctxKeyOrigURLHost, req.URL.Host).
				WithValue(ctxKeyOrigHost, req.Host)
		}
	}

	if err != nil {
		if isUnexpected(err) {
			errResp := proxy.OnError(fctx, req, true, err)
			if errResp != nil {
				proxy.writeResponse(downstream, req, errResp)
			}

			return proxy.logInitialReadError(downstream, err)
		}
		return nil
	}

	var next filters.Next
	if req.Method == http.MethodConnect {
		next = proxy.nextCONNECT(downstream)
	} else {
		var tr *http.Transport
		if upstream != nil {
			setUpstreamForAwareConn(fctx, upstream)
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					// always use the supplied upstream connection, but don't allow it to
					// be closed by the transport
					return &noCloseConn{upstream}, nil
				},
				// this transport is only used once, don't keep any idle connections,
				// however still allow the transport to close the connection after using
				// it
				MaxIdleConnsPerHost: -1,
			}
		} else {
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					conn, err := proxy.Dial(ctx, false, net, addr)
					if err == nil {
						// On first dialing conn, handle RequestAware
						setUpstreamForAwareConn(ctx, conn)
						handleRequestAware(ctx)
					}
					return conn, err
				},
				IdleConnTimeout: proxy.IdleTimeout,
				// since we have one transport per downstream connection, we don't need
				// more than this
				MaxIdleConnsPerHost: 1,
			}
		}

		defer tr.CloseIdleConnections()
		next = func(ctx filters.Context, modifiedReq *http.Request) (*http.Response, filters.Context, error) {
			modifiedReq = modifiedReq.WithContext(ctx)
			setRequestForAwareConn(ctx, modifiedReq)
			handleRequestAware(ctx)
			resp, err := tr.RoundTrip(prepareRequest(modifiedReq))
			handleResponseAware(ctx, modifiedReq, resp, err)
			if err != nil {
				err = errors.New("Unable to round-trip http request to upstream: %v", err)
			}
			return resp, ctx, err
		}
	}

	return proxy.processRequests(fctx, req.RemoteAddr, req, downstream, downstreamBuffered, next)
}

func (proxy *proxy) processRequests(ctx filters.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, next filters.Next) error {
	var readErr error
	var resp *http.Response
	var err error

	for {
		if req.URL.Scheme == "" {
			req.URL.Scheme = origURLScheme(ctx)
		}
		if req.URL.Host == "" {
			req.URL.Host = origURLHost(ctx)
		}
		if req.Host == "" {
			req.Host = origHost(ctx)
		}
		resp, ctx, err = proxy.Filter.Apply(ctx, req, next)
		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, false, err)
			if resp != nil {
				// On error, we will always close the connection
				resp.Close = true
			}
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return log.Errorf("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if err != nil {
			// We encountered an error on round-tripping, stop now
			return err
		}

		upstream := upstreamConn(ctx)
		upstreamAddr := upstreamAddr(ctx)
		isConnect := upstream != nil || upstreamAddr != ""

		buffered := downstreamBuffered.Buffered()
		if buffered > 0 {
			b, _ := downstreamBuffered.Peek(buffered)
			downstream = preconn.Wrap(downstream, b)
		}

		if isConnect {
			return proxy.proceedWithConnect(ctx, req, upstreamAddr, upstream, downstream)
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				errResp := proxy.OnError(ctx, req, true, readErr)
				if errResp != nil {
					proxy.writeResponse(downstream, req, errResp)
				}
				return log.Errorf("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		ctx = ctx.IncrementRequestNumber()
		req.RemoteAddr = remoteAddr
		req = req.WithContext(ctx)
	}
}

func handleRequestAware(ctx context.Context) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case RequestAware:
			req := requestForAwareConn(ctx)
			t.OnRequest(req)
		}
		return true
	})
}

func handleResponseAware(ctx context.Context, req *http.Request, resp *http.Response, err error) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case ResponseAware:
			t.OnResponse(req, resp, err)
		}
		return true
	})
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	if resp.ProtoMajor == 0 {
		resp.ProtoMajor = 1
		resp.ProtoMinor = 1
	}
	belowHTTP11 := !resp.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
		proxy.addIdleKeepAlive(resp.Header)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	req.Close = false

	// Request Header
	newHeader := make(http.Header)
	copyHeadersForForwarding(newHeader, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	newHeader.Set("Host", req.Host)
	req.Header = newHeader

	// Request URL
	req.URL = cloneURL(req.URL)
	// If req.URL.Scheme was blank, it's http. Otherwise, it's https and we leave
	// it alone.
	if req.URL.Scheme == "" {
		req.URL.Scheme = "http"
	}
	// We need to make sure the host is defined in the URL (not the actual URI)
	req.URL.Host = req.Host

	userAgent := req.UserAgent()
	if userAgent == "" {
		req.Header.Del("User-Agent")
	} else {
		req.Header.Set("User-Agent", userAgent)
	}

	return req
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		case "Keep-Alive":
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	if err == nil {
		return false
	}
	if err == io.EOF {
		return false
	}
	// This is okay per the HTTP spec.
	// See https://www.w3.org/Protocols/rfc2616/rfc2616-sec8.html#sec8.1.4
	if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
		return false
	}

	text := err.Error()
	return !strings.HasSuffix(text, "EOF") &&
		!strings.Contains(text, "i/o timeout") &&
		!strings.Contains(text, "Use of idled network connection") &&
		!strings.Contains(text, "use of closed network connection") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "broken pipe") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "connection reset by peer")
}

func defaultFilter(ctx filters.Context, req *http.Request, next filters.Next) (*http.Response, filters.Context, error) {
	return next(ctx, req)
}

func defaultOnError(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
	return nil
}
<commit_msg>fix log ordering for error grouping<commit_after>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/netx"
	"github.com/getlantern/preconn"
	"github.com/getlantern/proxy/filters"
)

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			safeClose(downstream)
			err = errors.New("Recovered from panic handling connection: %v", p)
		}
	}()

	err = proxy.handle(ctx, downstreamIn, downstream, nil)
	return
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection: %v", p)
		}
	}()

	conn.Close()
}

func (proxy *proxy) logInitialReadError(downstream net.Conn, err error) error {
	loc := downstream.LocalAddr()
	l := ""
	if loc != nil {
		l = loc.String()
	}
	rem := downstream.RemoteAddr()
	r := ""
	if rem != nil {
		r = rem.String()
	}
	return log.Errorf("Initial ReadRequest: %v from %v to %v", err, r, l)
}

func (proxy *proxy) handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn, upstream net.Conn) error {
	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
	}()

	downstreamBuffered := bufio.NewReader(downstreamIn)
	fctx := filters.WrapContext(withAwareConn(ctx), downstream)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = downstream.RemoteAddr().String()
		}
		if origURLScheme(ctx) == "" {
			fctx = fctx.
				WithValue(ctxKeyOrigURLScheme, req.URL.Scheme).
				WithValue(ctxKeyOrigURLHost, req.URL.Host).
				WithValue(ctxKeyOrigHost, req.Host)
		}
	}

	if err != nil {
		if isUnexpected(err) {
			errResp := proxy.OnError(fctx, req, true, err)
			if errResp != nil {
				proxy.writeResponse(downstream, req, errResp)
			}

			return proxy.logInitialReadError(downstream, err)
		}
		return nil
	}

	var next filters.Next
	if req.Method == http.MethodConnect {
		next = proxy.nextCONNECT(downstream)
	} else {
		var tr *http.Transport
		if upstream != nil {
			setUpstreamForAwareConn(fctx, upstream)
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					// always use the supplied upstream connection, but don't allow it to
					// be closed by the transport
					return &noCloseConn{upstream}, nil
				},
				// this transport is only used once, don't keep any idle connections,
				// however still allow the transport to close the connection after using
				// it
				MaxIdleConnsPerHost: -1,
			}
		} else {
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					conn, err := proxy.Dial(ctx, false, net, addr)
					if err == nil {
						// On first dialing conn, handle RequestAware
						setUpstreamForAwareConn(ctx, conn)
						handleRequestAware(ctx)
					}
					return conn, err
				},
				IdleConnTimeout: proxy.IdleTimeout,
				// since we have one transport per downstream connection, we don't need
				// more than this
				MaxIdleConnsPerHost: 1,
			}
		}

		defer tr.CloseIdleConnections()
		next = func(ctx filters.Context, modifiedReq *http.Request) (*http.Response, filters.Context, error) {
			modifiedReq = modifiedReq.WithContext(ctx)
			setRequestForAwareConn(ctx, modifiedReq)
			handleRequestAware(ctx)
			resp, err := tr.RoundTrip(prepareRequest(modifiedReq))
			handleResponseAware(ctx, modifiedReq, resp, err)
			if err != nil {
				err = errors.New("Unable to round-trip http request to upstream: %v", err)
			}
			return resp, ctx, err
		}
	}

	return proxy.processRequests(fctx, req.RemoteAddr, req, downstream, downstreamBuffered, next)
}

func (proxy *proxy) processRequests(ctx filters.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, next filters.Next) error {
	var readErr error
	var resp *http.Response
	var err error

	for {
		if req.URL.Scheme == "" {
			req.URL.Scheme = origURLScheme(ctx)
		}
		if req.URL.Host == "" {
			req.URL.Host = origURLHost(ctx)
		}
		if req.Host == "" {
			req.Host = origHost(ctx)
		}
		resp, ctx, err = proxy.Filter.Apply(ctx, req, next)
		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, false, err)
			if resp != nil {
				// On error, we will always close the connection
				resp.Close = true
			}
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return log.Errorf("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if err != nil {
			// We encountered an error on round-tripping, stop now
			return err
		}

		upstream := upstreamConn(ctx)
		upstreamAddr := upstreamAddr(ctx)
		isConnect := upstream != nil || upstreamAddr != ""

		buffered := downstreamBuffered.Buffered()
		if buffered > 0 {
			b, _ := downstreamBuffered.Peek(buffered)
			downstream = preconn.Wrap(downstream, b)
		}

		if isConnect {
			return proxy.proceedWithConnect(ctx, req, upstreamAddr, upstream, downstream)
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				errResp := proxy.OnError(ctx, req, true, readErr)
				if errResp != nil {
					proxy.writeResponse(downstream, req, errResp)
				}
				return log.Errorf("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		ctx = ctx.IncrementRequestNumber()
		req.RemoteAddr = remoteAddr
		req = req.WithContext(ctx)
	}
}

func handleRequestAware(ctx context.Context) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case RequestAware:
			req := requestForAwareConn(ctx)
			t.OnRequest(req)
		}
		return true
	})
}

func handleResponseAware(ctx context.Context, req *http.Request, resp *http.Response, err error) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case ResponseAware:
			t.OnResponse(req, resp, err)
		}
		return true
	})
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	if resp.ProtoMajor == 0 {
		resp.ProtoMajor = 1
		resp.ProtoMinor = 1
	}
	belowHTTP11 := !resp.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
		proxy.addIdleKeepAlive(resp.Header)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	req.Close = false

	// Request Header
	newHeader := make(http.Header)
	copyHeadersForForwarding(newHeader, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	newHeader.Set("Host", req.Host)
	req.Header = newHeader

	// Request URL
	req.URL = cloneURL(req.URL)
	// If req.URL.Scheme was blank, it's http. Otherwise, it's https and we leave
	// it alone.
	if req.URL.Scheme == "" {
		req.URL.Scheme = "http"
	}
	// We need to make sure the host is defined in the URL (not the actual URI)
	req.URL.Host = req.Host

	userAgent := req.UserAgent()
	if userAgent == "" {
		req.Header.Del("User-Agent")
	} else {
		req.Header.Set("User-Agent", userAgent)
	}

	return req
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		case "Keep-Alive":
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	if err == nil {
		return false
	}
	if err == io.EOF {
		return false
	}
	// This is okay per the HTTP spec.
	// See https://www.w3.org/Protocols/rfc2616/rfc2616-sec8.html#sec8.1.4
	if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
		return false
	}

	text := err.Error()
	return !strings.HasSuffix(text, "EOF") &&
		!strings.Contains(text, "i/o timeout") &&
		!strings.Contains(text, "Use of idled network connection") &&
		!strings.Contains(text, "use of closed network connection") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "broken pipe") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "connection reset by peer")
}

func defaultFilter(ctx filters.Context, req *http.Request, next filters.Next) (*http.Response, filters.Context, error) {
	return next(ctx, req)
}

func defaultOnError(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
	return nil
}
<|endoftext|>


<commit_before>package vulcan

import (
	. "github.com/mailgun/vulcan/limit"
	. "github.com/mailgun/vulcan/loadbalance"
	. "github.com/mailgun/vulcan/loadbalance/roundrobin"
	"github.com/mailgun/vulcan/netutils"
	. "github.com/mailgun/vulcan/route"
	. "github.com/mailgun/vulcan/upstream"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"time"
)

type ProxySuite struct {
	authHeaders http.Header
}

var _ = Suite(&ProxySuite{
	authHeaders: http.Header{
		"Authorization": []string{"Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="}}})

func (s *ProxySuite) SetUpTest(c *C) {
}

func (s *ProxySuite) Get(c *C, requestUrl string, header http.Header, body string) (*http.Response, []byte) {
	request, _ := http.NewRequest("GET", requestUrl, strings.NewReader(body))
	netutils.CopyHeaders(request.Header, header)
	request.Close = true
	// the HTTP lib treats Host as a special header.  it only respects the value on req.Host, and ignores
	// values in req.Headers
	if header.Get("Host") != "" {
		request.Host = header.Get("Host")
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		c.Fatalf("Get: %v", err)
	}

	bodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		c.Fatalf("Get body failed: %v", err)
	}
	return response, bodyBytes
}

func (s *ProxySuite) Post(c *C, requestUrl string, header http.Header, body url.Values) (*http.Response, []byte) {
	request, _ := http.NewRequest("POST", requestUrl, strings.NewReader(body.Encode()))
	netutils.CopyHeaders(request.Header, header)
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	request.Close = true
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		c.Fatalf("Post: %v", err)
	}

	bodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		c.Fatalf("Post body failed: %v", err)
	}
	return response, bodyBytes
}

type WebHandler func(http.ResponseWriter, *http.Request)

func (s *ProxySuite) newServer(handler WebHandler) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(handler))
}

func (s *ProxySuite) newProxyWithParams(
	l LoadBalancer,
	r Limiter,
	readTimeout time.Duration,
	dialTimeout time.Duration) *httptest.Server {

	proxySettings := ProxySettings{
		Router: &MatchAll{
			Location: &BaseLocation{LoadBalancer: l, Limiter: r},
		},
		HttpReadTimeout: readTimeout,
		HttpDialTimeout: dialTimeout,
	}

	proxy, err := NewReverseProxy(proxySettings)
	if err != nil {
		panic(err)
	}
	return httptest.NewServer(proxy)
}

func (s *ProxySuite) newProxy(l LoadBalancer, r Limiter) *httptest.Server {
	return s.newProxyWithParams(l, r, time.Duration(0), time.Duration(0))
}

func (s *ProxySuite) newUpstream(url string) Upstream {
	u, err := NewUpstreamFromString(url)
	if err != nil {
		panic(err)
	}
	return u
}

// Success, make sure we've successfully proxied the response
func (s *ProxySuite) TestSuccess(c *C) {
	upstream := s.newServer(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Hi, I'm upstream"))
	})
	defer upstream.Close()

	proxy := s.newProxy(NewRoundRobin(s.newUpstream(upstream.URL)), nil)
	defer proxy.Close()

	response, bodyBytes := s.Get(c, proxy.URL, s.authHeaders, "hello!")
	c.Assert(response.StatusCode, Equals, http.StatusOK)
	c.Assert(string(bodyBytes), Equals, "Hi, I'm upstream")
}
<commit_msg>Fix proxy test<commit_after>package vulcan

import (
	. "github.com/mailgun/vulcan/limit"
	. "github.com/mailgun/vulcan/loadbalance"
	. "github.com/mailgun/vulcan/loadbalance/roundrobin"
	"github.com/mailgun/vulcan/netutils"
	. "github.com/mailgun/vulcan/route"
	. "github.com/mailgun/vulcan/upstream"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"time"
)

type ProxySuite struct {
	authHeaders http.Header
}

var _ = Suite(&ProxySuite{
	authHeaders: http.Header{
		"Authorization": []string{"Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="}}})

func (s *ProxySuite) SetUpTest(c *C) {
}

func (s *ProxySuite) Get(c *C, requestUrl string, header http.Header, body string) (*http.Response, []byte) {
	request, _ := http.NewRequest("GET", requestUrl, strings.NewReader(body))
	netutils.CopyHeaders(request.Header, header)
	request.Close = true
	// the HTTP lib treats Host as a special header.  it only respects the value on req.Host, and ignores
	// values in req.Headers
	if header.Get("Host") != "" {
		request.Host = header.Get("Host")
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		c.Fatalf("Get: %v", err)
	}

	bodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		c.Fatalf("Get body failed: %v", err)
	}
	return response, bodyBytes
}

func (s *ProxySuite) Post(c *C, requestUrl string, header http.Header, body url.Values) (*http.Response, []byte) {
	request, _ := http.NewRequest("POST", requestUrl, strings.NewReader(body.Encode()))
	netutils.CopyHeaders(request.Header, header)
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	request.Close = true
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		c.Fatalf("Post: %v", err)
	}

	bodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		c.Fatalf("Post body failed: %v", err)
	}
	return response, bodyBytes
}

type WebHandler func(http.ResponseWriter, *http.Request)

func (s *ProxySuite) newServer(handler WebHandler) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(handler))
}

func (s *ProxySuite) newProxyWithParams(
	l LoadBalancer,
	r Limiter,
	readTimeout time.Duration,
	dialTimeout time.Duration) *httptest.Server {

	proxySettings := ProxySettings{
		Router: &MatchAll{
			Location: &BaseLocation{LoadBalancer: l, Limiter: r},
		},
		HttpReadTimeout: readTimeout,
		HttpDialTimeout: dialTimeout,
	}

	proxy, err := NewReverseProxy(proxySettings)
	if err != nil {
		panic(err)
	}
	return httptest.NewServer(proxy)
}

func (s *ProxySuite) newProxy(l LoadBalancer, r Limiter) *httptest.Server {
	return s.newProxyWithParams(l, r, time.Duration(0), time.Duration(0))
}

// Success, make sure we've successfully proxied the response
func (s *ProxySuite) TestSuccess(c *C) {
	upstream := s.newServer(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Hi, I'm upstream"))
	})
	defer upstream.Close()

	proxy := s.newProxy(NewRoundRobin(MustParseUpstream(upstream.URL)), nil)
	defer proxy.Close()

	response, bodyBytes := s.Get(c, proxy.URL, s.authHeaders, "hello!")
	c.Assert(response.StatusCode, Equals, http.StatusOK)
	c.Assert(string(bodyBytes), Equals, "Hi, I'm upstream")
}
<|endoftext|>


<commit_before>package slashquery

import (
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"

	"github.com/nbari/violetear"
	"github.com/slashquery/resolver"
)

func TestProxy(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("slashquery /?"))
	}))
	defer ts.Close()

	r, err := resolver.New("4.2.2.2")
	if err != nil {
		t.Fatal(err)
	}

	routes := make(map[string]*Route)
	routes["test"] = &Route{
		URL: ts.URL,
	}
	sq := &Slashquery{
		Config:    make(map[string]string),
		Resolver:  r,
		Routes:    routes,
		Servers:   make(map[string]*Servers),
		Upstreams: make(map[string]*Upstream),
	}
	if err := sq.Setup(); err != nil {
		t.Error(err)
	}
	sq.ResolveUpstreams()
	testRoute := sq.Routes["test"]
	expect(t, testRoute.URL, ts.URL)
	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Error(err)
	}
	testUpstream, ok := sq.Upstreams[u.Host]
	if !ok {
		t.Errorf("Expecting upstream: %s", testUpstream)
	}
	testServers, ok := sq.Servers[u.Host]
	if !ok {
		t.Errorf("Expecting servers: %s", testServers)
	}
	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		t.Error(err)
	}
	expect(t, testServers.Addresses[0], host)

	// test router
	router := violetear.New()
	router.Verbose = false
	router.LogRequests = false
	router.Handle("/*", sq.Proxy("test"), "GET,HEAD")

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", "/", nil)
	router.ServeHTTP(w, req)

	expect(t, w.Body.String(), "slashquery /?")
}
<commit_msg>test rawQuery<commit_after>package slashquery

import (
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"

	"github.com/nbari/violetear"
	"github.com/slashquery/resolver"
)

func TestProxy(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("slashquery /?"))
	}))
	defer ts.Close()

	r, err := resolver.New("4.2.2.2")
	if err != nil {
		t.Fatal(err)
	}

	routes := make(map[string]*Route)
	routes["test"] = &Route{
		URL: ts.URL,
	}
	sq := &Slashquery{
		Config:    make(map[string]string),
		Resolver:  r,
		Routes:    routes,
		Servers:   make(map[string]*Servers),
		Upstreams: make(map[string]*Upstream),
	}
	if err := sq.Setup(); err != nil {
		t.Error(err)
	}
	sq.ResolveUpstreams()
	testRoute := sq.Routes["test"]
	expect(t, testRoute.URL, ts.URL)
	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Error(err)
	}
	testUpstream, ok := sq.Upstreams[u.Host]
	if !ok {
		t.Errorf("Expecting upstream: %s", testUpstream)
	}
	testServers, ok := sq.Servers[u.Host]
	if !ok {
		t.Errorf("Expecting servers: %s", testServers)
	}
	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		t.Error(err)
	}
	expect(t, testServers.Addresses[0], host)

	// test router
	router := violetear.New()
	router.Verbose = false
	router.LogRequests = false
	router.Handle("/*", sq.Proxy("test"), "GET,HEAD")

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", "/", nil)
	router.ServeHTTP(w, req)

	expect(t, w.Body.String(), "slashquery /?")
}

func TestProxyQuery(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(fmt.Sprintf("%s", r.URL)))
	}))
	defer ts.Close()

	r, err := resolver.New("4.2.2.2")
	if err != nil {
		t.Fatal(err)
	}

	routes := make(map[string]*Route)
	routes["test"] = &Route{
		URL: fmt.Sprintf("%s/test?slash=query", ts.URL),
	}
	sq := &Slashquery{
		Config:    make(map[string]string),
		Resolver:  r,
		Routes:    routes,
		Servers:   make(map[string]*Servers),
		Upstreams: make(map[string]*Upstream),
	}
	if err := sq.Setup(); err != nil {
		t.Error(err)
	}
	sq.ResolveUpstreams()
	testRoute := sq.Routes["test"]
	expect(t, testRoute.rawQuery, "slash=query")
	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Error(err)
	}
	expect(t, fmt.Sprintf("http://%s", u.Host), ts.URL)
	testUpstream, ok := sq.Upstreams[u.Host]
	if !ok {
		t.Errorf("Expecting upstream: %s", testUpstream)
	}
	testServers, ok := sq.Servers[u.Host]
	if !ok {
		t.Errorf("Expecting servers: %s", testServers)
	}
	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		t.Error(err)
	}
	expect(t, testServers.Addresses[0], host)

	// test router
	router := violetear.New()
	router.Verbose = false
	router.LogRequests = false
	router.Handle("/*", sq.Proxy("test"), "GET,HEAD")

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", "/?foo=bar", nil)
	router.ServeHTTP(w, req)

	expect(t, w.Body.String(), "/test?slash=query&foo=bar")
}
<|endoftext|>


<commit_before>// Copyright © 2016 Zlatko Čalušić
//
// Use of this source code is governed by an MIT-style license that can be found in the LICENSE file.

// Package sysinfo is a pure Go library providing Linux OS / kernel / hardware system information.
package sysinfo

// SysInfo struct encapsulates all other information structs.
type SysInfo struct {
	Meta    Meta            `json:"sysinfo"`
	Node    Node            `json:"node"`
	OS      OS              `json:"os"`
	Kernel  Kernel          `json:"kernel"`
	Product Product         `json:"product"`
	Board   Board           `json:"board"`
	Chassis Chassis         `json:"chassis"`
	BIOS    BIOS            `json:"bios"`
	CPU     CPU             `json:"cpu"`
	Memory  Memory          `json:"memory"`
	Storage []StorageDevice `json:"storage"`
	Network []NetworkDevice `json:"network"`
}

// GetSysInfo gathers all available system information.
func (si *SysInfo) GetSysInfo() {
	// Meta info
	si.getMetaInfo()

	// Software info
	si.getNodeInfo()
	si.getOSInfo()
	si.getKernelInfo()

	// Hardware info
	si.getProductInfo()
	si.getBoardInfo()
	si.getChassisInfo()
	si.getBIOSInfo()
	si.getCPUInfo()
	si.getMemoryInfo()
	si.getStorageInfo()
	si.getNetworkInfo()
}
<commit_msg>Use omitempty tag for storage & network arrays<commit_after>// Copyright © 2016 Zlatko Čalušić
//
// Use of this source code is governed by an MIT-style license that can be found in the LICENSE file.

// Package sysinfo is a pure Go library providing Linux OS / kernel / hardware system information.
package sysinfo

// SysInfo struct encapsulates all other information structs.
type SysInfo struct {
	Meta    Meta            `json:"sysinfo"`
	Node    Node            `json:"node"`
	OS      OS              `json:"os"`
	Kernel  Kernel          `json:"kernel"`
	Product Product         `json:"product"`
	Board   Board           `json:"board"`
	Chassis Chassis         `json:"chassis"`
	BIOS    BIOS            `json:"bios"`
	CPU     CPU             `json:"cpu"`
	Memory  Memory          `json:"memory"`
	Storage []StorageDevice `json:"storage,omitempty"`
	Network []NetworkDevice `json:"network,omitempty"`
}

// GetSysInfo gathers all available system information.
func (si *SysInfo) GetSysInfo() {
	// Meta info
	si.getMetaInfo()

	// Software info
	si.getNodeInfo()
	si.getOSInfo()
	si.getKernelInfo()

	// Hardware info
	si.getProductInfo()
	si.getBoardInfo()
	si.getChassisInfo()
	si.getBIOSInfo()
	si.getCPUInfo()
	si.getMemoryInfo()
	si.getStorageInfo()
	si.getNetworkInfo()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"testing"

	. "github.com/flynn/go-check"
	"github.com/gorilla/sessions"
)

const (
	testLoginToken    = "test-login-token"
	testControllerKey = "test-controller-key"
)

// Hook gocheck up to the "go test" runner
func Test(t *testing.T) { TestingT(t) }

type S struct {
	srv        *httptest.Server
	cookiePath string
}

var _ = Suite(&S{})

func (s *S) SetUpSuite(c *C) {
	cookiePath := "/"
	s.srv = httptest.NewServer(APIHandler(&Config{
		SessionStore:  sessions.NewCookieStore([]byte("session-secret")),
		LoginToken:    testLoginToken,
		ControllerKey: testControllerKey,
		CookiePath:    cookiePath,
	}))
}

func (s *S) testAuthenticated(c *C, client *http.Client) {
	res, err := client.Get(s.srv.URL + "/config")
	c.Assert(err, IsNil)
	var conf *UserConfig
	c.Assert(json.NewDecoder(res.Body).Decode(&conf), IsNil)
	c.Assert(conf.User, Not(IsNil))
	c.Assert(conf.User.ControllerKey, Equals, testControllerKey)
}

func (s *S) TestUserSessionJSON(c *C) {
	loginInfo := LoginInfo{Token: testLoginToken}
	data, err := json.Marshal(&loginInfo)
	c.Assert(err, IsNil)
	jar, err := cookiejar.New(&cookiejar.Options{})
	c.Assert(err, IsNil)
	client := &http.Client{
		Jar: jar,
	}
	res, err := client.Post(s.srv.URL+"/user/sessions", "application/json", bytes.NewReader(data))
	c.Assert(err, IsNil)
	c.Assert(res.StatusCode, Equals, 200)
	s.testAuthenticated(c, client)
}

func (s *S) TestUserSessionForm(c *C) {
	data := url.Values{}
	data.Set("token", testLoginToken)
	jar, err := cookiejar.New(&cookiejar.Options{})
	c.Assert(err, IsNil)
	client := &http.Client{
		Jar: jar,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	res, err := client.PostForm(s.srv.URL+"/user/sessions", data)
	c.Assert(err, IsNil)
	c.Assert(res.StatusCode, Equals, 302)
	c.Assert(res.Header.Get("Location"), Equals, "/")
	s.testAuthenticated(c, client)
}
<commit_msg>dashboard: Fix typo<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"testing"

	. "github.com/flynn/go-check"
	"github.com/gorilla/sessions"
)

const (
	testLoginToken    = "test-login-token"
	testControllerKey = "test-controller-key"
)

// Hook gocheck up to the "go test" runner
func Test(t *testing.T) { TestingT(t) }

type S struct {
	srv        *httptest.Server
	cookiePath string
}

var _ = Suite(&S{})

func (s *S) SetUpSuite(c *C) {
	s.cookiePath = "/"
	s.srv = httptest.NewServer(APIHandler(&Config{
		SessionStore:  sessions.NewCookieStore([]byte("session-secret")),
		LoginToken:    testLoginToken,
		ControllerKey: testControllerKey,
		CookiePath:    s.cookiePath,
	}))
}

func (s *S) testAuthenticated(c *C, client *http.Client) {
	res, err := client.Get(s.srv.URL + "/config")
	c.Assert(err, IsNil)
	var conf *UserConfig
	c.Assert(json.NewDecoder(res.Body).Decode(&conf), IsNil)
	c.Assert(conf.User, Not(IsNil))
	c.Assert(conf.User.ControllerKey, Equals, testControllerKey)
}

func (s *S) TestUserSessionJSON(c *C) {
	loginInfo := LoginInfo{Token: testLoginToken}
	data, err := json.Marshal(&loginInfo)
	c.Assert(err, IsNil)
	jar, err := cookiejar.New(&cookiejar.Options{})
	c.Assert(err, IsNil)
	client := &http.Client{
		Jar: jar,
	}
	res, err := client.Post(s.srv.URL+"/user/sessions", "application/json", bytes.NewReader(data))
	c.Assert(err, IsNil)
	c.Assert(res.StatusCode, Equals, 200)
	s.testAuthenticated(c, client)
}

func (s *S) TestUserSessionForm(c *C) {
	data := url.Values{}
	data.Set("token", testLoginToken)
	jar, err := cookiejar.New(&cookiejar.Options{})
	c.Assert(err, IsNil)
	client := &http.Client{
		Jar: jar,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	res, err := client.PostForm(s.srv.URL+"/user/sessions", data)
	c.Assert(err, IsNil)
	c.Assert(res.StatusCode, Equals, 302)
	c.Assert(res.Header.Get("Location"), Equals, s.cookiePath)
	s.testAuthenticated(c, client)
}
<|endoftext|>


<commit_before>// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See the LICENSE file in the project root for license information.

package arm

import (
	"fmt"
	"testing"

	"github.com/Azure/azure-sdk-for-go/arm/compute"
	"github.com/Azure/packer-azure/packer/builder/azure/common/constants"
	"github.com/mitchellh/multistep"
)

func TestStepCaptureImageShouldFailIfStepFails(t *testing.T) {

	var testSubject = &StepCaptureImage{
		capture: func(string, string, *compute.VirtualMachineCaptureParameters) error {
			return fmt.Errorf("!! Unit Test FAIL !!")
		},
		say:   func(message string) {},
		error: func(e error) {},
	}

	stateBag := createTestStateBagStepCaptureImage()

	var result = testSubject.Run(stateBag)
	if result != multistep.ActionHalt {
		t.Fatalf("Expected the step to return 'ActionHalt', but got '%d'.", result)
	}

	if _, ok := stateBag.GetOk(constants.Error); ok == false {
		t.Fatalf("Expected the step to set stateBag['%s'], but it was not.", constants.Error)
	}

	if _, ok := stateBag.GetOk(constants.ArmIsValid); ok == true {
		t.Fatalf("Expected the step to not set stateBag['%s'], but it was.", constants.ArmIsValid)
	}
}

func TestStepCaptureImageShouldPassIfStepPasses(t *testing.T) {
	var testSubject = &StepCaptureImage{
		capture: func(string, string, *compute.VirtualMachineCaptureParameters) error { return nil },
		say:     func(message string) {},
		error:   func(e error) {},
	}

	stateBag := createTestStateBagStepCaptureImage()

	var result = testSubject.Run(stateBag)
	if result != multistep.ActionContinue {
		t.Fatalf("Expected the step to return 'ActionContinue', but got '%d'.", result)
	}

	if _, ok := stateBag.GetOk(constants.Error); ok == true {
		t.Fatalf("Expected the step to not set stateBag['%s'], but it was.", constants.Error)
	}

	if _, ok := stateBag.GetOk(constants.ArmIsValid); ok == false {
		t.Fatalf("Expected the step to set stateBag['%s'], but it was not.", constants.ArmIsValid)
	}
}

func TestStepCaptureImageShouldTakeStepArgumentsFromStateBag(t *testing.T) {
	var actualResourceGroupName string
	var actualComputeName string
	var actualVirtualMachineCaptureParameters *compute.VirtualMachineCaptureParameters

	var testSubject = &StepCaptureImage{
		capture: func(resourceGroupName string, computeName string, parameters *compute.VirtualMachineCaptureParameters) error {
			actualResourceGroupName = resourceGroupName
			actualComputeName = computeName
			actualVirtualMachineCaptureParameters = parameters

			return nil
		},
		say:   func(message string) {},
		error: func(e error) {},
	}

	stateBag := createTestStateBagStepCaptureImage()
	var result = testSubject.Run(stateBag)

	if result != multistep.ActionContinue {
		t.Fatalf("Expected the step to return 'ActionContinue', but got '%d'.", result)
	}

	var expectedComputeName = stateBag.Get(constants.ArmComputeName).(string)
	var expectedResourceGroupName = stateBag.Get(constants.ArmResourceGroupName).(string)
	var expectedVirtualMachineCaptureParameters = stateBag.Get(constants.ArmVirtualMachineCaptureParameters).(*compute.VirtualMachineCaptureParameters)

	if actualComputeName != expectedComputeName {
		t.Fatalf("Expected StepCaptureImage to source 'constants.ArmComputeName' from the state bag, but it did not.")
	}

	if actualResourceGroupName != expectedResourceGroupName {
		t.Fatalf("Expected StepCaptureImage to source 'constants.ArmResourceGroupName' from the state bag, but it did not.")
	}

	if actualVirtualMachineCaptureParameters != expectedVirtualMachineCaptureParameters {
		t.Fatalf("Expected StepCaptureImage to source 'constants.ArmVirtualMachineCaptureParameters' from the state bag, but it did not.")
	}
}

func createTestStateBagStepCaptureImage() multistep.StateBag {
	stateBag := new(multistep.BasicStateBag)

	stateBag.Put(constants.ArmComputeName, "Unit Test: ComputeName")
	stateBag.Put(constants.ArmResourceGroupName, "Unit Test: ResourceGroupName")
	stateBag.Put(constants.ArmVirtualMachineCaptureParameters, &compute.VirtualMachineCaptureParameters{})

	return stateBag
}
<commit_msg>Fix test case.<commit_after>// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See the LICENSE file in the project root for license information.

package arm

import (
	"fmt"
	"testing"

	"github.com/Azure/azure-sdk-for-go/arm/compute"
	"github.com/Azure/packer-azure/packer/builder/azure/common/constants"
	"github.com/mitchellh/multistep"
)

func TestStepCaptureImageShouldFailIfStepFails(t *testing.T) {

	var testSubject = &StepCaptureImage{
		capture: func(string, string, *compute.VirtualMachineCaptureParameters) error {
			return fmt.Errorf("!! Unit Test FAIL !!")
		},
		say:   func(message string) {},
		error: func(e error) {},
	}

	stateBag := createTestStateBagStepCaptureImage()

	var result = testSubject.Run(stateBag)
	if result != multistep.ActionHalt {
		t.Fatalf("Expected the step to return 'ActionHalt', but got '%d'.", result)
	}

	if _, ok := stateBag.GetOk(constants.Error); ok == false {
		t.Fatalf("Expected the step to set stateBag['%s'], but it was not.", constants.Error)
	}
}

func TestStepCaptureImageShouldPassIfStepPasses(t *testing.T) {
	var testSubject = &StepCaptureImage{
		capture: func(string, string, *compute.VirtualMachineCaptureParameters) error { return nil },
		say:     func(message string) {},
		error:   func(e error) {},
	}

	stateBag := createTestStateBagStepCaptureImage()

	var result = testSubject.Run(stateBag)
	if result != multistep.ActionContinue {
		t.Fatalf("Expected the step to return 'ActionContinue', but got '%d'.", result)
	}

	if _, ok := stateBag.GetOk(constants.Error); ok == true {
		t.Fatalf("Expected the step to not set stateBag['%s'], but it was.", constants.Error)
	}
}

func TestStepCaptureImageShouldTakeStepArgumentsFromStateBag(t *testing.T) {
	var actualResourceGroupName string
	var actualComputeName string
	var actualVirtualMachineCaptureParameters *compute.VirtualMachineCaptureParameters

	var testSubject = &StepCaptureImage{
		capture: func(resourceGroupName string, computeName string, parameters *compute.VirtualMachineCaptureParameters) error {
			actualResourceGroupName = resourceGroupName
			actualComputeName = computeName
			actualVirtualMachineCaptureParameters = parameters

			return nil
		},
		say:   func(message string) {},
		error: func(e error) {},
	}

	stateBag := createTestStateBagStepCaptureImage()
	var result = testSubject.Run(stateBag)

	if result != multistep.ActionContinue {
		t.Fatalf("Expected the step to return 'ActionContinue', but got '%d'.", result)
	}

	var expectedComputeName = stateBag.Get(constants.ArmComputeName).(string)
	var expectedResourceGroupName = stateBag.Get(constants.ArmResourceGroupName).(string)
	var expectedVirtualMachineCaptureParameters = stateBag.Get(constants.ArmVirtualMachineCaptureParameters).(*compute.VirtualMachineCaptureParameters)

	if actualComputeName != expectedComputeName {
		t.Fatalf("Expected StepCaptureImage to source 'constants.ArmComputeName' from the state bag, but it did not.")
	}

	if actualResourceGroupName != expectedResourceGroupName {
		t.Fatalf("Expected StepCaptureImage to source 'constants.ArmResourceGroupName' from the state bag, but it did not.")
	}

	if actualVirtualMachineCaptureParameters != expectedVirtualMachineCaptureParameters {
		t.Fatalf("Expected StepCaptureImage to source 'constants.ArmVirtualMachineCaptureParameters' from the state bag, but it did not.")
	}
}

func createTestStateBagStepCaptureImage() multistep.StateBag {
	stateBag := new(multistep.BasicStateBag)

	stateBag.Put(constants.ArmComputeName, "Unit Test: ComputeName")
	stateBag.Put(constants.ArmResourceGroupName, "Unit Test: ResourceGroupName")
	stateBag.Put(constants.ArmVirtualMachineCaptureParameters, &compute.VirtualMachineCaptureParameters{})

	return stateBag
}
<|endoftext|>


<commit_before>package robo

import (
	"testing"
)

var segmentTests = []struct {
	s  segment
	in string
	n  int
}{
	{segment{0, "/foo", nil}, "/foo", 4},
	{segment{0, "/foo", nil}, "/bar", 0},
	{segment{0, "/foo", nil}, "/foo/bar", 4},
	{segment{0, "/foo/", nil}, "/foo/bar", 5},

	{segment{1, "", nil}, "foo", 3},
	{segment{1, "", nil}, "foo/bar", 3},
	{segment{1, "", nil}, "/bar", 0},

	{segment{2, "", []rune{'a', 'z'}}, "foo", 3},
	{segment{2, "", []rune{'a', 'z'}}, "foo/", 3},
	{segment{2, "", []rune{'a', 'z'}}, "/foo", 0},
	{segment{2, "", []rune{'a', 'z', '0', '9'}}, "1a23foo", 7},

	{segment{3, "", nil}, "foo", 3},
	{segment{3, "", nil}, "foo/bar", 7},
	{segment{3, "", nil}, "/a/b/c//", 8},
}

func TestSegmentMatch(t *testing.T) {
	for _, test := range segmentTests {
		if n := test.s.Match(test.in); n != test.n {
			t.Errorf("%+v.Match(%q)\n", test.s, test.in)
			t.Errorf("  got  %d\n", n)
			t.Errorf("  want %d\n", test.n)
		}
	}
}

var patternTests = []struct {
	p  pattern
	in string
	ok bool
	m  map[string]string
}{
	{[]segment{segment{0, "/", nil}}, "/", true, nil},
	{[]segment{segment{0, "/foo/", nil}, segment{3, "", nil}}, "/foo/bar", true, nil},
	{[]segment{segment{0, "/foo/", nil}, segment{3, "", nil}}, "/foo-bar", false, nil},
	{[]segment{segment{0, "/", nil}, segment{1, "a", nil}}, "/foo", true, map[string]string{"a": "foo"}},
	{[]segment{segment{0, "/", nil}, segment{1, "b", nil}}, "/foo-bar", true, map[string]string{"b": "foo-bar"}},
	{[]segment{segment{0, "/", nil}, segment{1, "b", nil}}, "/foo/bar", false, nil},
	{[]segment{segment{0, "/", nil}, segment{1, "a", nil}, segment{0, "/", nil}, segment{2, "b", []rune{'a', 'z'}}}, "/foo/bar", true, map[string]string{"a": "foo", "b": "bar"}},
	{[]segment{segment{0, "/", nil}, segment{1, "a", nil}, segment{0, "/", nil}, segment{2, "b", []rune{'a', 'z'}}}, "/foo/123", false, nil},
}

func TestPatternMatch(t *testing.T) {
	for _, test := range patternTests {
		ok, m := test.p.Match(test.in)
		if ok != test.ok || len(m) != len(test.m) {
			goto fail
		}

		for key := range m {
			if m[key] != test.m[key] {
				goto fail
			}
		}

		continue

	fail:
		t.Errorf("%+v.Match(%q):", test.p, test.in)
		t.Errorf("  want %v, %+v", test.ok, test.m)
		t.Errorf("  want %v, %+v", ok, m)
	}
}
<commit_msg>Tidy pattern_test.go<commit_after>package robo

import (
	"testing"
)

var segmentTests = []struct {
	segment segment
	input   string
	n       int
}{
	{segment{0, "/foo", nil}, "/foo", 4},
	{segment{0, "/foo", nil}, "/bar", 0},
	{segment{0, "/foo", nil}, "/foo/bar", 4},
	{segment{0, "/foo/", nil}, "/foo/bar", 5},

	{segment{1, "", nil}, "foo", 3},
	{segment{1, "", nil}, "foo/bar", 3},
	{segment{1, "", nil}, "/bar", 0},

	{segment{2, "", []rune{'a', 'z'}}, "foo", 3},
	{segment{2, "", []rune{'a', 'z'}}, "foo/", 3},
	{segment{2, "", []rune{'a', 'z'}}, "/foo", 0},
	{segment{2, "", []rune{'a', 'z', '0', '9'}}, "1a23foo", 7},

	{segment{3, "", nil}, "foo", 3},
	{segment{3, "", nil}, "foo/bar", 7},
	{segment{3, "", nil}, "/a/b/c//", 8},
}

func TestSegmentMatch(t *testing.T) {
	for _, test := range segmentTests {
		if n := test.segment.Match(test.input); n != test.n {
			t.Errorf("%+v.Match(%q)", test.segment, test.input)
			t.Errorf("  got  %d", n)
			t.Errorf("  want %d", test.n)
		}
	}
}

var patternTests = []struct {
	pattern pattern
	input   string
	ok      bool
	params  map[string]string
}{
	{[]segment{segment{0, "/", nil}}, "/", true, nil},
	{[]segment{segment{0, "/foo/", nil}, segment{3, "", nil}}, "/foo/bar", true, nil},
	{[]segment{segment{0, "/foo/", nil}, segment{3, "", nil}}, "/foo-bar", false, nil},
	{[]segment{segment{0, "/", nil}, segment{1, "a", nil}}, "/foo", true, map[string]string{"a": "foo"}},
	{[]segment{segment{0, "/", nil}, segment{1, "b", nil}}, "/foo-bar", true, map[string]string{"b": "foo-bar"}},
	{[]segment{segment{0, "/", nil}, segment{1, "b", nil}}, "/foo/bar", false, nil},
	{[]segment{segment{0, "/", nil}, segment{1, "a", nil}, segment{0, "/", nil}, segment{2, "b", []rune{'a', 'z'}}}, "/foo/bar", true, map[string]string{"a": "foo", "b": "bar"}},
	{[]segment{segment{0, "/", nil}, segment{1, "a", nil}, segment{0, "/", nil}, segment{2, "b", []rune{'a', 'z'}}}, "/foo/123", false, nil},
}

func TestPatternMatch(t *testing.T) {
	for _, test := range patternTests {
		ok, params := test.pattern.Match(test.input)
		if ok != test.ok || len(params) != len(test.params) {
			goto fail
		}

		for key := range params {
			if params[key] != test.params[key] {
				goto fail
			}
		}

		continue

	fail:
		t.Errorf("%+v.Match(%q):", test.pattern, test.input)
		t.Errorf("  want %v, %+v", test.ok, test.params)
		t.Errorf("  want %v, %+v", ok, params)
	}
}
<|endoftext|>


<commit_before>package scribo

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"

	"github.com/gorilla/mux"
)

// Version specifies what version of Scribo we're on.
const Version = "0.1"

// App defines the complete structure of a web application including a
// router for multiplexing, storages, assets, templates, cookies, and more.
// This should be the primary interface for working with Scribo.
type App struct {
	StaticDir   string
	TemplateDir string
	Templates   *template.Template
	Router      *mux.Router
	DB          *sql.DB
}

// CreateApp allows you to easily instantiate an App instance.
func CreateApp() *App {
	// Instantiate the app
	app := new(App)

	// Connect to the database
	app.DB = ConnectDB()

	// Set the static and template directories
	root, _ := os.Getwd()
	app.StaticDir = path.Join(root, "assets")
	app.TemplateDir = path.Join(root, "templates")

	// Load the templates from the template directory
	tmplGlob := path.Join(app.TemplateDir, "*")
	app.Templates = template.Must(template.ParseGlob(tmplGlob))

	app.Router = mux.NewRouter().StrictSlash(true)

	// Add all routes
	for _, route := range routes {
		app.AddRoute(route)
	}

	// Add a static file server pointing at the assets directory
	app.AddStatic(app.StaticDir)

	return app
}

// Run the web application via the associated router.
func (app *App) Run(port int) {
	addr := fmt.Sprintf(":%d", port)

	name, err := os.Hostname()
	if err != nil {
		name = "localhost"
	}

	log.Printf("Starting server at http://%s:%d (use CTRL+C to quit)", name, port)
	log.Fatal(http.ListenAndServe(addr, app.Router))
}

// AddRoute allows you to add a handler for a specific route to the router.
func (app *App) AddRoute(route Route) {
	var handler http.Handler
	handler = route.Handler(app)

	if route.Authorize {
		handler = Authenticate(app, handler)
	}

	handler = Logger(app, handler)
	// handler = Debugger(app, handler)

	app.Router.
		Methods(route.Methods...).
		Path(route.Pattern).
		Name(route.Name).
		Handler(handler)
}

// AddStatic creates a handler to serve static files.
func (app *App) AddStatic(staticDir string) {
	static := http.StripPrefix("/assets/", http.FileServer(http.Dir(staticDir)))
	app.Router.PathPrefix("/assets/").Handler(Logger(app, static))
}

// Abort is a handler to terminate the request with no error message
func (app *App) Abort(w http.ResponseWriter, statusCode int) {
	w.WriteHeader(statusCode)
}

// Error is a handler to terminate the request with an error message.
func (app *App) Error(w http.ResponseWriter, err error, statusCode int) {
	http.Error(w, err.Error(), statusCode)
}

// JSONAbort is a handler to terminate the request with a JSON response
func (app *App) JSONAbort(w http.ResponseWriter, statusCode int) {
	w.Header().Set(CTKEY, CTJSON)
	w.WriteHeader(statusCode)

	response := make(map[string]string)
	response["code"] = strconv.Itoa(statusCode)
	response["reason"] = http.StatusText(statusCode)

	if err := json.NewEncoder(w).Encode(response); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// JSONError is a handler to terminate the request with a JSON response
func (app *App) JSONError(w http.ResponseWriter, err error, statusCode int) {
	w.Header().Set(CTKEY, CTJSON)
	w.WriteHeader(statusCode)

	response := make(map[string]string)
	response["code"] = strconv.Itoa(statusCode)
	response["error"] = err.Error()

	if err := json.NewEncoder(w).Encode(response); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<commit_msg>version bump<commit_after>package scribo

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"

	"github.com/gorilla/mux"
)

// Version specifies what version of Scribo we're on.
const Version = "1.0"

// App defines the complete structure of a web application including a
// router for multiplexing, storages, assets, templates, cookies, and more.
// This should be the primary interface for working with Scribo.
type App struct {
	StaticDir   string
	TemplateDir string
	Templates   *template.Template
	Router      *mux.Router
	DB          *sql.DB
}

// CreateApp allows you to easily instantiate an App instance.
func CreateApp() *App {
	// Instantiate the app
	app := new(App)

	// Connect to the database
	app.DB = ConnectDB()

	// Set the static and template directories
	root, _ := os.Getwd()
	app.StaticDir = path.Join(root, "assets")
	app.TemplateDir = path.Join(root, "templates")

	// Load the templates from the template directory
	tmplGlob := path.Join(app.TemplateDir, "*")
	app.Templates = template.Must(template.ParseGlob(tmplGlob))

	app.Router = mux.NewRouter().StrictSlash(true)

	// Add all routes
	for _, route := range routes {
		app.AddRoute(route)
	}

	// Add a static file server pointing at the assets directory
	app.AddStatic(app.StaticDir)

	return app
}

// Run the web application via the associated router.
func (app *App) Run(port int) {
	addr := fmt.Sprintf(":%d", port)

	name, err := os.Hostname()
	if err != nil {
		name = "localhost"
	}

	log.Printf("Starting server at http://%s:%d (use CTRL+C to quit)", name, port)
	log.Fatal(http.ListenAndServe(addr, app.Router))
}

// AddRoute allows you to add a handler for a specific route to the router.
func (app *App) AddRoute(route Route) {
	var handler http.Handler
	handler = route.Handler(app)

	if route.Authorize {
		handler = Authenticate(app, handler)
	}

	handler = Logger(app, handler)
	// handler = Debugger(app, handler)

	app.Router.
		Methods(route.Methods...).
		Path(route.Pattern).
		Name(route.Name).
		Handler(handler)
}

// AddStatic creates a handler to serve static files.
func (app *App) AddStatic(staticDir string) {
	static := http.StripPrefix("/assets/", http.FileServer(http.Dir(staticDir)))
	app.Router.PathPrefix("/assets/").Handler(Logger(app, static))
}

// Abort is a handler to terminate the request with no error message
func (app *App) Abort(w http.ResponseWriter, statusCode int) {
	w.WriteHeader(statusCode)
}

// Error is a handler to terminate the request with an error message.
func (app *App) Error(w http.ResponseWriter, err error, statusCode int) {
	http.Error(w, err.Error(), statusCode)
}

// JSONAbort is a handler to terminate the request with a JSON response
func (app *App) JSONAbort(w http.ResponseWriter, statusCode int) {
	w.Header().Set(CTKEY, CTJSON)
	w.WriteHeader(statusCode)

	response := make(map[string]string)
	response["code"] = strconv.Itoa(statusCode)
	response["reason"] = http.StatusText(statusCode)

	if err := json.NewEncoder(w).Encode(response); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// JSONError is a handler to terminate the request with a JSON response
func (app *App) JSONError(w http.ResponseWriter, err error, statusCode int) {
	w.Header().Set(CTKEY, CTJSON)
	w.WriteHeader(statusCode)

	response := make(map[string]string)
	response["code"] = strconv.Itoa(statusCode)
	response["error"] = err.Error()

	if err := json.NewEncoder(w).Encode(response); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"log"
	"os"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
)

// SQSHandler hello world
type SQSHandler struct {
	client        sqsiface.SQSAPI
	messageID     string
	messageBody   string
	receiptHandle string
	queueURL      string
	awsRegion     string
}

// SQSClient hello world
type SQSClient struct {
	queueURL  string
	awsRegion string
	sqsClient sqsiface.SQSAPI
}

func (handler *SQSHandler) id() *string {
	return &handler.messageID
}

func (handler *SQSHandler) body() *string {
	return &handler.messageBody
}

func (handler *SQSHandler) initialize() {
	handler.newClient(sqs.New(session.New()))
	handler.queueURL = os.Getenv("TASK_QUEUE_URL")
}

func (handler *SQSHandler) newClient(client sqsiface.SQSAPI) {
	handler.client = client
}

func (handler *SQSHandler) receive() bool {
	receiveMessageParams := &sqs.ReceiveMessageInput{
		QueueUrl:            aws.String(handler.queueURL),
		MaxNumberOfMessages: aws.Int64(1),
		WaitTimeSeconds:     aws.Int64(20),
	}
	receiveMessageResponse, receiveMessageError := handler.client.ReceiveMessage(receiveMessageParams)

	if receiveMessageError != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		log.Println("E: ", receiveMessageError.Error())
		return false
	}
	if len(receiveMessageResponse.Messages) == 0 {
		log.Println("I: ", "No messages retrieved from queue")
		return false
	}

	handler.messageBody = *receiveMessageResponse.Messages[0].Body
	handler.messageID = *receiveMessageResponse.Messages[0].MessageId
	handler.receiptHandle = *receiveMessageResponse.Messages[0].ReceiptHandle

	writeFileError := ioutil.WriteFile("payload.json", []byte(handler.messageBody), 0644)
	if writeFileError != nil {
		panic(writeFileError)
	}
	return true
}

func (handler *SQSHandler) success() {
	deleteMessageParams := &sqs.DeleteMessageInput{
		QueueUrl:      aws.String(handler.queueURL),
		ReceiptHandle: aws.String(handler.receiptHandle),
	}
	_, deleteMessageError := handler.client.DeleteMessage(deleteMessageParams)

	if deleteMessageError != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		log.Println(deleteMessageError.Error())
		return
	}
}
<commit_msg>SQS should retry 30 times<commit_after>package main

import (
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
)

// SQSHandler hello world
type SQSHandler struct {
	client        sqsiface.SQSAPI
	messageID     string
	messageBody   string
	receiptHandle string
	queueURL      string
	awsRegion     string
}

// SQSClient hello world
type SQSClient struct {
	queueURL  string
	awsRegion string
	sqsClient sqsiface.SQSAPI
}

func (handler *SQSHandler) id() *string {
	return &handler.messageID
}

func (handler *SQSHandler) body() *string {
	return &handler.messageBody
}

func (handler *SQSHandler) initialize() {
	handler.newClient(sqs.New(session.New(), &aws.Config{
		MaxRetries: aws.Int(30),
		HTTPClient: &http.Client{
			Timeout: 30 * time.Second,
		},
	}))
	handler.queueURL = os.Getenv("TASK_QUEUE_URL")
}

func (handler *SQSHandler) newClient(client sqsiface.SQSAPI) {
	handler.client = client
}

func (handler *SQSHandler) receive() bool {
	receiveMessageParams := &sqs.ReceiveMessageInput{
		QueueUrl:            aws.String(handler.queueURL),
		MaxNumberOfMessages: aws.Int64(1),
		WaitTimeSeconds:     aws.Int64(20),
	}
	receiveMessageResponse, receiveMessageError := handler.client.ReceiveMessage(receiveMessageParams)

	if receiveMessageError != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		log.Println("E: ", receiveMessageError.Error())
		return false
	}
	if len(receiveMessageResponse.Messages) == 0 {
		log.Println("I: ", "No messages retrieved from queue")
		return false
	}

	handler.messageBody = *receiveMessageResponse.Messages[0].Body
	handler.messageID = *receiveMessageResponse.Messages[0].MessageId
	handler.receiptHandle = *receiveMessageResponse.Messages[0].ReceiptHandle

	writeFileError := ioutil.WriteFile("payload.json", []byte(handler.messageBody), 0644)
	if writeFileError != nil {
		panic(writeFileError)
	}
	return true
}

func (handler *SQSHandler) success() {
	deleteMessageParams := &sqs.DeleteMessageInput{
		QueueUrl:      aws.String(handler.queueURL),
		ReceiptHandle: aws.String(handler.receiptHandle),
	}
	_, deleteMessageError := handler.client.DeleteMessage(deleteMessageParams)

	if deleteMessageError != nil {
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				log.Fatalf("AWS SDK Error: %s %s", awsErr.Code(), awsErr.Message())
				// The message cannot be deleted from the SQS queue!
			}
		}
		return
	}
}
<|endoftext|>