<commit_before>// Package colly implements a HTTP scraping framework
package colly

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/html"
	"golang.org/x/net/html/charset"

	"github.com/PuerkitoBio/goquery"
)

// Collector provides the scraper instance for a scraping job
type Collector struct {
	// UserAgent is the User-Agent string used by HTTP requests
	UserAgent string
	// MaxDepth limits the recursion depth of visited URLs.
	// Set it to 0 for infinite recursion (default).
	MaxDepth int
	// AllowedDomains is a domain whitelist.
	// Leave it blank to allow any domains to be visited
	AllowedDomains []string
	// AllowURLRevisit allows multiple downloads of the same URL
	AllowURLRevisit bool
	// MaxBodySize is the limit of the retrieved response body in bytes.
	// `0` means unlimited.
	// The default value for MaxBodySize is 10MB (10 * 1024 * 1024 bytes).
	MaxBodySize int
	// CacheDir specifies a location where GET requests are cached as files.
	// When it's not defined, caching is disabled.
	CacheDir          string
	visitedURLs       []string
	htmlCallbacks     map[string]HTMLCallback
	requestCallbacks  []RequestCallback
	responseCallbacks []ResponseCallback
	backend           *httpBackend
	wg                *sync.WaitGroup
	lock              *sync.Mutex
}

// Request is the representation of a HTTP request made by a Collector
type Request struct {
	// URL is the parsed URL of the HTTP request
	URL *url.URL
	// Headers contains the Request's HTTP headers
	Headers *http.Header
	// Ctx is a context between a Request and a Response
	Ctx *Context
	// Depth is the number of the parents of this request
	Depth     int
	collector *Collector
}

// Response is the representation of a HTTP response made by a Collector
type Response struct {
	// StatusCode is the status code of the Response
	StatusCode int
	// Body is the content of the Response
	Body []byte
	// Ctx is a context between a Request and a Response
	Ctx *Context
	// Request is the Request object of the response
	Request *Request
	// Headers contains the Response's HTTP headers
	Headers *http.Header
}

// HTMLElement is the representation of a HTML tag.
type HTMLElement struct {
	// Name is the name of the tag
	Name       string
	Text       string
	attributes []html.Attribute
	// Request is the request object of the element's HTML document
	Request *Request
	// Response is the Response object of the element's HTML document
	Response *Response
	// DOM is the goquery parsed DOM object of the page. DOM is relative
	// to the current HTMLElement
	DOM *goquery.Selection
}

// Context provides a tiny layer for passing data between callbacks
type Context struct {
	contextMap map[string]string
	lock       *sync.RWMutex
}

// RequestCallback is a type alias for OnRequest callback functions
type RequestCallback func(*Request)

// ResponseCallback is a type alias for OnResponse callback functions
type ResponseCallback func(*Response)

// HTMLCallback is a type alias for OnHTML callback functions
type HTMLCallback func(*HTMLElement)

// NewCollector creates a new Collector instance with default configuration
func NewCollector() *Collector {
	c := &Collector{}
	c.Init()
	return c
}

// NewContext initializes a new Context instance
func NewContext() *Context {
	return &Context{
		contextMap: make(map[string]string),
		lock:       &sync.RWMutex{},
	}
}

// Init initializes the Collector's private variables and sets default
// configuration for the Collector
func (c *Collector) Init() {
	c.UserAgent = "colly - https://github.com/asciimoo/colly"
	c.MaxDepth = 0
	c.visitedURLs = make([]string, 0, 8)
	c.htmlCallbacks = make(map[string]HTMLCallback, 0)
	c.requestCallbacks = make([]RequestCallback, 0, 8)
	c.responseCallbacks = make([]ResponseCallback, 0, 8)
	c.MaxBodySize = 10 * 1024 * 1024
	c.backend = &httpBackend{}
	c.backend.Init()
	c.wg = &sync.WaitGroup{}
	c.lock = &sync.Mutex{}
}

// Visit starts Collector's collecting job by creating a
// request to the URL specified in parameter.
// Visit also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (c *Collector) Visit(URL string) error {
	return c.scrape(URL, "GET", 1, nil, nil)
}

// Post starts a collector job by creating a POST request.
// Post also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (c *Collector) Post(URL string, requestData map[string]string) error {
	return c.scrape(URL, "POST", 1, createFormReader(requestData), nil)
}

// PostRaw starts a collector job by creating a POST request with raw binary data.
// Post also calls the previously provided callbacks
func (c *Collector) PostRaw(URL string, requestData []byte) error {
	return c.scrape(URL, "POST", 1, bytes.NewReader(requestData), nil)
}

func (c *Collector) scrape(u, method string, depth int, requestData io.Reader, ctx *Context) error {
	c.wg.Add(1)
	defer c.wg.Done()
	if err := c.requestCheck(u, depth); err != nil {
		return err
	}
	parsedURL, err := url.Parse(u)
	if err != nil {
		return err
	}
	if parsedURL.Scheme == "" {
		parsedURL.Scheme = "http"
	}
	if !c.isDomainAllowed(parsedURL.Host) {
		return errors.New("Forbidden domain")
	}
	req, err := http.NewRequest(method, parsedURL.String(), requestData)
	if err != nil {
		return err
	}
	req.Header.Set("User-Agent", c.UserAgent)
	if ctx == nil {
		ctx = NewContext()
	}
	request := &Request{
		URL:       parsedURL,
		Headers:   &req.Header,
		Ctx:       ctx,
		Depth:     depth,
		collector: c,
	}

	c.handleOnRequest(request)

	if method == "POST" && req.Header.Get("Content-Type") == "" {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	}
	response, err := c.backend.Cache(req, c.MaxBodySize, c.CacheDir)
	// TODO add OnError callback to handle these cases
	if err != nil {
		return err
	}
	response.Ctx = ctx
	response.Request = request
	response.fixCharset()

	c.handleOnResponse(response)

	c.handleOnHTML(request, response)

	return nil
}

func (c *Collector) requestCheck(u string, depth int) error {
	if u == "" {
		return errors.New("Missing URL")
	}
	if c.MaxDepth > 0 && c.MaxDepth < depth {
		return errors.New("Max depth limit reached")
	}
	if !c.AllowURLRevisit {
		for _, u2 := range c.visitedURLs {
			if u2 == u {
				return errors.New("URL already visited")
			}
		}
		c.lock.Lock()
		c.visitedURLs = append(c.visitedURLs, u)
		c.lock.Unlock()
	}
	return nil
}

func (c *Collector) isDomainAllowed(domain string) bool {
	if c.AllowedDomains == nil || len(c.AllowedDomains) == 0 {
		return true
	}
	for _, d2 := range c.AllowedDomains {
		if d2 == domain {
			return true
		}
	}
	return false
}

// Wait returns when the collector jobs are finished
func (c *Collector) Wait() {
	c.wg.Wait()
}

// OnRequest registers a function. Function will be executed on every
// request made by the Collector
func (c *Collector) OnRequest(f RequestCallback) {
	c.lock.Lock()
	c.requestCallbacks = append(c.requestCallbacks, f)
	c.lock.Unlock()
}

// OnResponse registers a function. Function will be executed on every response
func (c *Collector) OnResponse(f ResponseCallback) {
	c.lock.Lock()
	c.responseCallbacks = append(c.responseCallbacks, f)
	c.lock.Unlock()
}

// OnHTML registers a function. Function will be executed on every HTML
// element matched by the `goquerySelector` parameter.
// `goquerySelector` is a selector used by https://github.com/PuerkitoBio/goquery
func (c *Collector) OnHTML(goquerySelector string, f HTMLCallback) {
	c.lock.Lock()
	c.htmlCallbacks[goquerySelector] = f
	c.lock.Unlock()
}

// WithTransport allows you to set a custom http.Transport for this collector.
func (c *Collector) WithTransport(transport *http.Transport) {
	c.backend.Client.Transport = transport
}

// DisableCookies turns off cookie handling for this collector
func (c *Collector) DisableCookies() {
	c.backend.Client.Jar = nil
}

// SetRequestTimeout overrides the default timeout (10 seconds) for this collector
func (c *Collector) SetRequestTimeout(timeout time.Duration) {
	c.backend.Client.Timeout = timeout
}

func (c *Collector) handleOnRequest(r *Request) {
	for _, f := range c.requestCallbacks {
		f(r)
	}
}

func (c *Collector) handleOnResponse(r *Response) {
	for _, f := range c.responseCallbacks {
		f(r)
	}
}

func (c *Collector) handleOnHTML(req *Request, resp *Response) {
	if strings.Index(strings.ToLower(resp.Headers.Get("Content-Type")), "html") == -1 {
		return
	}
	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(resp.Body))
	if err != nil {
		return
	}
	for expr, f := range c.htmlCallbacks {
		doc.Find(expr).Each(func(i int, s *goquery.Selection) {
			for _, n := range s.Nodes {
				e := &HTMLElement{
					Name:       n.Data,
					Request:    req,
					Response:   resp,
					Text:       goquery.NewDocumentFromNode(n).Text(),
					DOM:        s,
					attributes: n.Attr,
				}
				f(e)
			}
		})
	}
}

// Limit adds a new `LimitRule` to the collector
func (c *Collector) Limit(rule *LimitRule) error {
	return c.backend.Limit(rule)
}

// Limits adds new `LimitRule`s to the collector
func (c *Collector) Limits(rules []*LimitRule) error {
	return c.backend.Limits(rules)
}

// SetCookies handles the receipt of the cookies in a reply for the given URL
func (c *Collector) SetCookies(URL string, cookies []*http.Cookie) error {
	if c.backend.Client.Jar == nil {
		return errors.New("Cookie jar is not available")
	}
	u, err := url.Parse(URL)
	if err != nil {
		return err
	}
	c.backend.Client.Jar.SetCookies(u, cookies)
	return nil
}

// Cookies returns the cookies to send in a request for the given URL.
func (c *Collector) Cookies(URL string) []*http.Cookie {
	if c.backend.Client.Jar == nil {
		return nil
	}
	u, err := url.Parse(URL)
	if err != nil {
		return nil
	}
	return c.backend.Client.Jar.Cookies(u)
}

// Attr returns the selected attribute of a HTMLElement or empty string
// if no attribute found
func (h *HTMLElement) Attr(k string) string {
	for _, a := range h.attributes {
		if a.Key == k {
			return a.Val
		}
	}
	return ""
}

// AbsoluteURL returns with the resolved absolute URL of an URL chunk.
// AbsoluteURL returns empty string if the URL chunk is a fragment or
// could not be parsed
func (r *Request) AbsoluteURL(u string) string {
	if strings.HasPrefix(u, "#") {
		return ""
	}
	absURL, err := r.URL.Parse(u)
	if err != nil {
		return ""
	}
	absURL.Fragment = ""
	if absURL.Scheme == "//" {
		absURL.Scheme = r.URL.Scheme
	}
	return absURL.String()
}

// Visit continues Collector's collecting job by creating a
// request and preserves the Context of the previous request.
// Visit also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (r *Request) Visit(URL string) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "GET", r.Depth+1, nil, r.Ctx)
}

// Post continues a collector job by creating a POST request and preserves the Context
// of the previous request.
// Post also calls the previously provided OnRequest, OnResponse, OnHTML callbacks
func (r *Request) Post(URL string, requestData map[string]string) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "POST", r.Depth+1, createFormReader(requestData), r.Ctx)
}

// PostRaw starts a collector job by creating a POST request with raw binary data.
// PostRaw preserves the Context of the previous request
// and calls the previously provided callbacks
func (r *Request) PostRaw(URL string, requestData []byte) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "POST", r.Depth+1, bytes.NewReader(requestData), r.Ctx)
}

// UnmarshalBinary decodes Context value to nil
// This function is used by request caching
func (c *Context) UnmarshalBinary(_ []byte) error {
	return nil
}

// MarshalBinary encodes Context value
// This function is used by request caching
func (c *Context) MarshalBinary() (_ []byte, _ error) {
	return nil, nil
}

// Put stores a value in Context
func (c *Context) Put(key, value string) {
	c.lock.Lock()
	c.contextMap[key] = value
	c.lock.Unlock()
}

// Get retrieves a value from Context. If no value found for `k`
// Get returns an empty string if key not found
func (c *Context) Get(key string) string {
	c.lock.RLock()
	if v, ok := c.contextMap[key]; ok {
		c.lock.RUnlock()
		return v
	}
	c.lock.RUnlock()
	return ""
}

func createFormReader(data map[string]string) io.Reader {
	form := url.Values{}
	for k, v := range data {
		form.Add(k, v)
	}
	return strings.NewReader(form.Encode())
}

func (r *Response) fixCharset() {
	contentType := strings.ToLower(r.Headers.Get("Content-Type"))
	if strings.Index(contentType, "charset") == -1 {
		return
	}
	if strings.Index(contentType, "utf-8") != -1 || strings.Index(contentType, "utf8") != -1 {
		return
	}
	encodedBodyReader, err := charset.NewReader(bytes.NewReader(r.Body), contentType)
	if err != nil {
		return
	}
	tmpBody, err := ioutil.ReadAll(encodedBodyReader)
	if err != nil {
		return
	}
	r.Body = tmpBody
}
<commit_msg>remove html handler from collector<commit_after>// Package colly implements a HTTP scraping framework
package colly

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/html"
	"golang.org/x/net/html/charset"

	"github.com/PuerkitoBio/goquery"
)

// Collector provides the scraper instance for a scraping job
type Collector struct {
	// UserAgent is the User-Agent string used by HTTP requests
	UserAgent string
	// MaxDepth limits the recursion depth of visited URLs.
	// Set it to 0 for infinite recursion (default).
	MaxDepth int
	// AllowedDomains is a domain whitelist.
	// Leave it blank to allow any domains to be visited
	AllowedDomains []string
	// AllowURLRevisit allows multiple downloads of the same URL
	AllowURLRevisit bool
	// MaxBodySize is the limit of the retrieved response body in bytes.
	// `0` means unlimited.
	// The default value for MaxBodySize is 10MB (10 * 1024 * 1024 bytes).
	MaxBodySize int
	// CacheDir specifies a location where GET requests are cached as files.
	// When it's not defined, caching is disabled.
	CacheDir          string
	visitedURLs       []string
	htmlCallbacks     map[string]HTMLCallback
	requestCallbacks  []RequestCallback
	responseCallbacks []ResponseCallback
	backend           *httpBackend
	wg                *sync.WaitGroup
	lock              *sync.Mutex
}

// Request is the representation of a HTTP request made by a Collector
type Request struct {
	// URL is the parsed URL of the HTTP request
	URL *url.URL
	// Headers contains the Request's HTTP headers
	Headers *http.Header
	// Ctx is a context between a Request and a Response
	Ctx *Context
	// Depth is the number of the parents of this request
	Depth     int
	collector *Collector
}

// Response is the representation of a HTTP response made by a Collector
type Response struct {
	// StatusCode is the status code of the Response
	StatusCode int
	// Body is the content of the Response
	Body []byte
	// Ctx is a context between a Request and a Response
	Ctx *Context
	// Request is the Request object of the response
	Request *Request
	// Headers contains the Response's HTTP headers
	Headers *http.Header
}

// HTMLElement is the representation of a HTML tag.
type HTMLElement struct {
	// Name is the name of the tag
	Name       string
	Text       string
	attributes []html.Attribute
	// Request is the request object of the element's HTML document
	Request *Request
	// Response is the Response object of the element's HTML document
	Response *Response
	// DOM is the goquery parsed DOM object of the page. DOM is relative
	// to the current HTMLElement
	DOM *goquery.Selection
}

// Context provides a tiny layer for passing data between callbacks
type Context struct {
	contextMap map[string]string
	lock       *sync.RWMutex
}

// RequestCallback is a type alias for OnRequest callback functions
type RequestCallback func(*Request)

// ResponseCallback is a type alias for OnResponse callback functions
type ResponseCallback func(*Response)

// HTMLCallback is a type alias for OnHTML callback functions
type HTMLCallback func(*HTMLElement)

// NewCollector creates a new Collector instance with default configuration
func NewCollector() *Collector {
	c := &Collector{}
	c.Init()
	return c
}

// NewContext initializes a new Context instance
func NewContext() *Context {
	return &Context{
		contextMap: make(map[string]string),
		lock:       &sync.RWMutex{},
	}
}

// Init initializes the Collector's private variables and sets default
// configuration for the Collector
func (c *Collector) Init() {
	c.UserAgent = "colly - https://github.com/asciimoo/colly"
	c.MaxDepth = 0
	c.visitedURLs = make([]string, 0, 8)
	c.htmlCallbacks = make(map[string]HTMLCallback, 0)
	c.requestCallbacks = make([]RequestCallback, 0, 8)
	c.responseCallbacks = make([]ResponseCallback, 0, 8)
	c.MaxBodySize = 10 * 1024 * 1024
	c.backend = &httpBackend{}
	c.backend.Init()
	c.wg = &sync.WaitGroup{}
	c.lock = &sync.Mutex{}
}

// Visit starts Collector's collecting job by creating a
// request to the URL specified in parameter.
// Visit also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (c *Collector) Visit(URL string) error {
	return c.scrape(URL, "GET", 1, nil, nil)
}

// Post starts a collector job by creating a POST request.
// Post also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (c *Collector) Post(URL string, requestData map[string]string) error {
	return c.scrape(URL, "POST", 1, createFormReader(requestData), nil)
}

// PostRaw starts a collector job by creating a POST request with raw binary data.
// Post also calls the previously provided callbacks
func (c *Collector) PostRaw(URL string, requestData []byte) error {
	return c.scrape(URL, "POST", 1, bytes.NewReader(requestData), nil)
}

func (c *Collector) scrape(u, method string, depth int, requestData io.Reader, ctx *Context) error {
	c.wg.Add(1)
	defer c.wg.Done()
	if err := c.requestCheck(u, depth); err != nil {
		return err
	}
	parsedURL, err := url.Parse(u)
	if err != nil {
		return err
	}
	if parsedURL.Scheme == "" {
		parsedURL.Scheme = "http"
	}
	if !c.isDomainAllowed(parsedURL.Host) {
		return errors.New("Forbidden domain")
	}
	req, err := http.NewRequest(method, parsedURL.String(), requestData)
	if err != nil {
		return err
	}
	req.Header.Set("User-Agent", c.UserAgent)
	if ctx == nil {
		ctx = NewContext()
	}
	request := &Request{
		URL:       parsedURL,
		Headers:   &req.Header,
		Ctx:       ctx,
		Depth:     depth,
		collector: c,
	}

	c.handleOnRequest(request)

	if method == "POST" && req.Header.Get("Content-Type") == "" {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	}
	response, err := c.backend.Cache(req, c.MaxBodySize, c.CacheDir)
	// TODO add OnError callback to handle these cases
	if err != nil {
		return err
	}
	response.Ctx = ctx
	response.Request = request
	response.fixCharset()

	c.handleOnResponse(response)

	c.handleOnHTML(request, response)

	return nil
}

func (c *Collector) requestCheck(u string, depth int) error {
	if u == "" {
		return errors.New("Missing URL")
	}
	if c.MaxDepth > 0 && c.MaxDepth < depth {
		return errors.New("Max depth limit reached")
	}
	if !c.AllowURLRevisit {
		for _, u2 := range c.visitedURLs {
			if u2 == u {
				return errors.New("URL already visited")
			}
		}
		c.lock.Lock()
		c.visitedURLs = append(c.visitedURLs, u)
		c.lock.Unlock()
	}
	return nil
}

func (c *Collector) isDomainAllowed(domain string) bool {
	if c.AllowedDomains == nil || len(c.AllowedDomains) == 0 {
		return true
	}
	for _, d2 := range c.AllowedDomains {
		if d2 == domain {
			return true
		}
	}
	return false
}

// Wait returns when the collector jobs are finished
func (c *Collector) Wait() {
	c.wg.Wait()
}

// OnRequest registers a function. Function will be executed on every
// request made by the Collector
func (c *Collector) OnRequest(f RequestCallback) {
	c.lock.Lock()
	c.requestCallbacks = append(c.requestCallbacks, f)
	c.lock.Unlock()
}

// OnResponse registers a function. Function will be executed on every response
func (c *Collector) OnResponse(f ResponseCallback) {
	c.lock.Lock()
	c.responseCallbacks = append(c.responseCallbacks, f)
	c.lock.Unlock()
}

// OnHTML registers a function. Function will be executed on every HTML
// element matched by the `goquerySelector` parameter.
// `goquerySelector` is a selector used by https://github.com/PuerkitoBio/goquery
func (c *Collector) OnHTML(goquerySelector string, f HTMLCallback) {
	c.lock.Lock()
	c.htmlCallbacks[goquerySelector] = f
	c.lock.Unlock()
}

// OnHTML deregister a function. Function will not be execute after register
func (c *Collector) OnHTMLDetach(goquerySelector string) {
	c.lock.Lock()
	delete(c.htmlCallbacks, goquerySelector)
	c.lock.Unlock()
}

// WithTransport allows you to set a custom http.Transport for this collector.
func (c *Collector) WithTransport(transport *http.Transport) {
	c.backend.Client.Transport = transport
}

// DisableCookies turns off cookie handling for this collector
func (c *Collector) DisableCookies() {
	c.backend.Client.Jar = nil
}

// SetRequestTimeout overrides the default timeout (10 seconds) for this collector
func (c *Collector) SetRequestTimeout(timeout time.Duration) {
	c.backend.Client.Timeout = timeout
}

func (c *Collector) handleOnRequest(r *Request) {
	for _, f := range c.requestCallbacks {
		f(r)
	}
}

func (c *Collector) handleOnResponse(r *Response) {
	for _, f := range c.responseCallbacks {
		f(r)
	}
}

func (c *Collector) handleOnHTML(req *Request, resp *Response) {
	if strings.Index(strings.ToLower(resp.Headers.Get("Content-Type")), "html") == -1 {
		return
	}
	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(resp.Body))
	if err != nil {
		return
	}
	for expr, f := range c.htmlCallbacks {
		doc.Find(expr).Each(func(i int, s *goquery.Selection) {
			for _, n := range s.Nodes {
				e := &HTMLElement{
					Name:       n.Data,
					Request:    req,
					Response:   resp,
					Text:       goquery.NewDocumentFromNode(n).Text(),
					DOM:        s,
					attributes: n.Attr,
				}
				f(e)
			}
		})
	}
}

// Limit adds a new `LimitRule` to the collector
func (c *Collector) Limit(rule *LimitRule) error {
	return c.backend.Limit(rule)
}

// Limits adds new `LimitRule`s to the collector
func (c *Collector) Limits(rules []*LimitRule) error {
	return c.backend.Limits(rules)
}

// SetCookies handles the receipt of the cookies in a reply for the given URL
func (c *Collector) SetCookies(URL string, cookies []*http.Cookie) error {
	if c.backend.Client.Jar == nil {
		return errors.New("Cookie jar is not available")
	}
	u, err := url.Parse(URL)
	if err != nil {
		return err
	}
	c.backend.Client.Jar.SetCookies(u, cookies)
	return nil
}

// Cookies returns the cookies to send in a request for the given URL.
func (c *Collector) Cookies(URL string) []*http.Cookie {
	if c.backend.Client.Jar == nil {
		return nil
	}
	u, err := url.Parse(URL)
	if err != nil {
		return nil
	}
	return c.backend.Client.Jar.Cookies(u)
}

// Attr returns the selected attribute of a HTMLElement or empty string
// if no attribute found
func (h *HTMLElement) Attr(k string) string {
	for _, a := range h.attributes {
		if a.Key == k {
			return a.Val
		}
	}
	return ""
}

// AbsoluteURL returns with the resolved absolute URL of an URL chunk.
// AbsoluteURL returns empty string if the URL chunk is a fragment or
// could not be parsed
func (r *Request) AbsoluteURL(u string) string {
	if strings.HasPrefix(u, "#") {
		return ""
	}
	absURL, err := r.URL.Parse(u)
	if err != nil {
		return ""
	}
	absURL.Fragment = ""
	if absURL.Scheme == "//" {
		absURL.Scheme = r.URL.Scheme
	}
	return absURL.String()
}

// Visit continues Collector's collecting job by creating a
// request and preserves the Context of the previous request.
// Visit also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (r *Request) Visit(URL string) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "GET", r.Depth+1, nil, r.Ctx)
}

// Post continues a collector job by creating a POST request and preserves the Context
// of the previous request.
// Post also calls the previously provided OnRequest, OnResponse, OnHTML callbacks
func (r *Request) Post(URL string, requestData map[string]string) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "POST", r.Depth+1, createFormReader(requestData), r.Ctx)
}

// PostRaw starts a collector job by creating a POST request with raw binary data.
// PostRaw preserves the Context of the previous request
// and calls the previously provided callbacks
func (r *Request) PostRaw(URL string, requestData []byte) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "POST", r.Depth+1, bytes.NewReader(requestData), r.Ctx)
}

// UnmarshalBinary decodes Context value to nil
// This function is used by request caching
func (c *Context) UnmarshalBinary(_ []byte) error {
	return nil
}

// MarshalBinary encodes Context value
// This function is used by request caching
func (c *Context) MarshalBinary() (_ []byte, _ error) {
	return nil, nil
}

// Put stores a value in Context
func (c *Context) Put(key, value string) {
	c.lock.Lock()
	c.contextMap[key] = value
	c.lock.Unlock()
}

// Get retrieves a value from Context. If no value found for `k`
// Get returns an empty string if key not found
func (c *Context) Get(key string) string {
	c.lock.RLock()
	if v, ok := c.contextMap[key]; ok {
		c.lock.RUnlock()
		return v
	}
	c.lock.RUnlock()
	return ""
}

func createFormReader(data map[string]string) io.Reader {
	form := url.Values{}
	for k, v := range data {
		form.Add(k, v)
	}
	return strings.NewReader(form.Encode())
}

func (r *Response) fixCharset() {
	contentType := strings.ToLower(r.Headers.Get("Content-Type"))
	if strings.Index(contentType, "charset") == -1 {
		return
	}
	if strings.Index(contentType, "utf-8") != -1 || strings.Index(contentType, "utf8") != -1 {
		return
	}
	encodedBodyReader, err := charset.NewReader(bytes.NewReader(r.Body), contentType)
	if err != nil {
		return
	}
	tmpBody, err := ioutil.ReadAll(encodedBodyReader)
	if err != nil {
		return
	}
	r.Body = tmpBody
}
<|endoftext|>


<commit_before>package lib

import (
	"github.com/mitsuse/pushbullet-go"
	"github.com/mitsuse/pushbullet-go/requests"
	"github.com/op/go-logging"
	"gopkg.in/gomail.v2"
	"time"
)

// Sends a Pushbullet-Push containing the given data to the saved API-key.
func sendPush(apiKey string, name string, url string, newStatus string, oldStatus string) {
	logging.MustGetLogger("").Debug("Sending Push about \"" + url + "\"...")

	pb := pushbullet.New(apiKey)

	push := requests.NewNote()
	push.Title = GetConfiguration().Dynamic.Title + " - Status Change"
	push.Body = name + " (" + url + ") went from \"" + oldStatus + "\" to \"" + newStatus + "\"."

	_, err := pb.PostPushesNote(push)
	if err != nil {
		logging.MustGetLogger("").Error("Unable to send Push: ", err)
	}
}

// Sends an email containing the given data to the saved e-mail-address.
// Needs a configured SMTP-server (in config-file).
func sendMail(recipient string, name string, url string, newStatus string, oldStatus string) {
	if GetConfiguration().Mailer.Host == "" || GetConfiguration().Mailer.Host == "smtp.mymail.com" {
		logging.MustGetLogger("").Warning("Not sending email because of missing configuration.")
		return
	}

	logging.MustGetLogger("").Debug("Sending email about \"" + url + "\"...")

	mConf := GetConfiguration().Mailer

	m := gomail.NewMessage()
	m.SetAddressHeader("From", mConf.From, GetConfiguration().Dynamic.Title)
	m.SetHeader("To", recipient)
	m.SetHeader("Subject", "Status Change: "+name)
	m.SetBody("text/html", "Hello,

"+
		""+name+""+" ("+url+") went on "+time.Now().Format("02.01.2006")+" at "+time.Now().Format("15:04:05")+" from \""+oldStatus+"\" to \""+newStatus+"\".

"+
		"Sincerely,
"+GetConfiguration().Dynamic.Title+"

"+
		"This email was sent automatically, please do not respond to it.")

	d := gomail.NewPlainDialer(mConf.Host, mConf.Port, mConf.User, mConf.Password)

	if err := d.DialAndSend(m); err != nil {
		logging.MustGetLogger("").Error("Unable to send email: ", err)
	}
}
<commit_msg>Send Link-Push instead of Note-Push<commit_after>package lib

import (
	"github.com/mitsuse/pushbullet-go"
	"github.com/mitsuse/pushbullet-go/requests"
	"github.com/op/go-logging"
	"gopkg.in/gomail.v2"
	"time"
)

// Sends a Pushbullet-Push containing the given data to the saved API-key.
func sendPush(apiKey string, name string, url string, newStatus string, oldStatus string) {
	logging.MustGetLogger("").Debug("Sending Push about \"" + url + "\"...")

	pb := pushbullet.New(apiKey)

	push := requests.NewLink()
	push.Title = GetConfiguration().Dynamic.Title + " - Status Change"
	push.Body = name + " went from \"" + oldStatus + "\" to \"" + newStatus + "\"."
	push.Url = url

	_, err := pb.PostPushesLink(push)
	if err != nil {
		logging.MustGetLogger("").Error("Unable to send Push: ", err)
	}
}

// Sends an email containing the given data to the saved e-mail-address.
// Needs a configured SMTP-server (in config-file).
func sendMail(recipient string, name string, url string, newStatus string, oldStatus string) {
	if GetConfiguration().Mailer.Host == "" || GetConfiguration().Mailer.Host == "smtp.mymail.com" {
		logging.MustGetLogger("").Warning("Not sending email because of missing configuration.")
		return
	}

	logging.MustGetLogger("").Debug("Sending email about \"" + url + "\"...")

	mConf := GetConfiguration().Mailer

	m := gomail.NewMessage()
	m.SetAddressHeader("From", mConf.From, GetConfiguration().Dynamic.Title)
	m.SetHeader("To", recipient)
	m.SetHeader("Subject", "Status Change: "+name)
	m.SetBody("text/html", "Hello,

"+
		""+name+""+" ("+url+") went on "+time.Now().Format("02.01.2006")+" at "+time.Now().Format("15:04:05")+" from \""+oldStatus+"\" to \""+newStatus+"\".

"+
		"Sincerely,
"+GetConfiguration().Dynamic.Title+"

"+
		"This email was sent automatically, please do not respond to it.")

	d := gomail.NewPlainDialer(mConf.Host, mConf.Port, mConf.User, mConf.Password)

	if err := d.DialAndSend(m); err != nil {
		logging.MustGetLogger("").Error("Unable to send email: ", err)
	}
}
<|endoftext|>


<commit_before>package output

import (
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Jeffail/benthos/v3/internal/component/output"
	"github.com/Jeffail/benthos/v3/internal/docs"
	"github.com/Jeffail/benthos/v3/lib/log"
	"github.com/Jeffail/benthos/v3/lib/metrics"
	"github.com/Jeffail/benthos/v3/lib/response"
	"github.com/Jeffail/benthos/v3/lib/types"
	"github.com/Jeffail/benthos/v3/lib/util/retries"
	"github.com/cenkalti/backoff/v4"
)

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

func init() {
	Constructors[TypeRetry] = TypeSpec{
		constructor: fromSimpleConstructor(NewRetry),
		Summary: `
Attempts to write messages to a child output and if the write fails for any
reason the message is retried either until success or, if the retries or max
elapsed time fields are non-zero, either is reached.`,
		Description: `
All messages in Benthos are always retried on an output error, but this would
usually involve propagating the error back to the source of the message, whereby
it would be reprocessed before reaching the output layer once again.

This output type is useful whenever we wish to avoid reprocessing a message on
the event of a failed send. We might, for example, have a dedupe processor that
we want to avoid reapplying to the same message more than once in the pipeline.

Rather than retrying the same output you may wish to retry the send using a
different output target (a dead letter queue). In which case you should instead
use the ` + "[`try`](/docs/components/outputs/try)" + ` output type.`,
		FieldSpecs: retries.FieldSpecs().Add(
			docs.FieldCommon("output", "A child output.").HasType(docs.FieldOutput),
		),
		Categories: []Category{
			CategoryUtility,
		},
	}
}

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

// RetryConfig contains configuration values for the Retry output type.
type RetryConfig struct {
	Output         *Config `json:"output" yaml:"output"`
	retries.Config `json:",inline" yaml:",inline"`
}

// NewRetryConfig creates a new RetryConfig with default values.
func NewRetryConfig() RetryConfig {
	rConf := retries.NewConfig()
	rConf.MaxRetries = 0
	rConf.Backoff.InitialInterval = "100ms"
	rConf.Backoff.MaxInterval = "1s"
	rConf.Backoff.MaxElapsedTime = "0s"
	return RetryConfig{
		Output: nil,
		Config: retries.NewConfig(),
	}
}

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

type dummyRetryConfig struct {
	Output         interface{} `json:"output" yaml:"output"`
	retries.Config `json:",inline" yaml:",inline"`
}

// MarshalJSON prints an empty object instead of nil.
func (r RetryConfig) MarshalJSON() ([]byte, error) {
	dummy := dummyRetryConfig{
		Output: r.Output,
		Config: r.Config,
	}
	if r.Output == nil {
		dummy.Output = struct{}{}
	}
	return json.Marshal(dummy)
}

// MarshalYAML prints an empty object instead of nil.
func (r RetryConfig) MarshalYAML() (interface{}, error) {
	dummy := dummyRetryConfig{
		Output: r.Output,
		Config: r.Config,
	}
	if r.Output == nil {
		dummy.Output = struct{}{}
	}
	return dummy, nil
}

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

// Retry is an output type that continuously writes a message to a child output
// until the send is successful.
type Retry struct {
	running int32
	conf    RetryConfig

	wrapped     Type
	backoffCtor func() backoff.BackOff

	stats metrics.Type
	log   log.Modular

	transactionsIn  <-chan types.Transaction
	transactionsOut chan types.Transaction

	closeChan  chan struct{}
	closedChan chan struct{}
}

// NewRetry creates a new Retry input type.
func NewRetry(
	conf Config,
	mgr types.Manager,
	log log.Modular,
	stats metrics.Type,
) (Type, error) {
	if conf.Retry.Output == nil {
		return nil, errors.New("cannot create retry output without a child")
	}

	wrapped, err := New(*conf.Retry.Output, mgr, log, stats)
	if err != nil {
		return nil, fmt.Errorf("failed to create output '%v': %v", conf.Retry.Output.Type, err)
	}

	var boffCtor func() backoff.BackOff
	if boffCtor, err = conf.Retry.GetCtor(); err != nil {
		return nil, err
	}

	return &Retry{
		running: 1,
		conf:    conf.Retry,

		log:             log,
		stats:           stats,
		wrapped:         wrapped,
		backoffCtor:     boffCtor,
		transactionsOut: make(chan types.Transaction),

		closeChan:  make(chan struct{}),
		closedChan: make(chan struct{}),
	}, nil
}

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

func (r *Retry) loop() {
	// Metrics paths
	var (
		mRunning      = r.stats.GetGauge("retry.running")
		mCount        = r.stats.GetCounter("retry.count")
		mSuccess      = r.stats.GetCounter("retry.send.success")
		mPartsSuccess = r.stats.GetCounter("retry.parts.send.success")
		mError        = r.stats.GetCounter("retry.send.error")
		mEndOfRetries = r.stats.GetCounter("retry.end_of_retries")
	)

	wg := sync.WaitGroup{}

	defer func() {
		wg.Wait()
		close(r.transactionsOut)
		r.wrapped.CloseAsync()
		err := r.wrapped.WaitForClose(time.Second)
		for ; err != nil; err = r.wrapped.WaitForClose(time.Second) {
		}
		mRunning.Decr(1)
		close(r.closedChan)
	}()
	mRunning.Incr(1)

	errInterruptChan := make(chan struct{})
	var errLooped int64

	for atomic.LoadInt32(&r.running) == 1 {
		// Do not consume another message while pending messages are being
		// reattempted.
		for atomic.LoadInt64(&errLooped) > 0 {
			select {
			case <-errInterruptChan:
			case <-time.After(time.Millisecond * 100):
				// Just incase an interrupt doesn't arrive.
			case <-r.closeChan:
				return
			}
		}

		var tran types.Transaction
		var open bool
		select {
		case tran, open = <-r.transactionsIn:
			if !open {
				return
			}
			mCount.Incr(1)
		case <-r.closeChan:
			return
		}

		rChan := make(chan types.Response)
		select {
		case r.transactionsOut <- types.NewTransaction(tran.Payload, rChan):
		case <-r.closeChan:
			return
		}

		wg.Add(1)
		go func(ts types.Transaction, resChan chan types.Response) {
			var backOff backoff.BackOff
			var resOut types.Response
			var inErrLoop bool

			defer func() {
				wg.Done()
				if inErrLoop {
					atomic.AddInt64(&errLooped, -1)

					// We're exiting our error loop, so (attempt to) interrupt the
					// consumer.
					select {
					case errInterruptChan <- struct{}{}:
					default:
					}
				}
			}()

			for atomic.LoadInt32(&r.running) == 1 {
				var res types.Response
				select {
				case res = <-resChan:
				case <-r.closeChan:
					return
				}

				if res.Error() != nil {
					if !inErrLoop {
						inErrLoop = true
						atomic.AddInt64(&errLooped, 1)
					}

					mError.Incr(1)
					r.log.Errorf("Failed to send message: %v\n", res.Error())
					if backOff == nil {
						backOff = r.backoffCtor()
					}

					nextBackoff := backOff.NextBackOff()
					if nextBackoff == backoff.Stop {
						mEndOfRetries.Incr(1)
						resOut = response.NewNoack()
						break
					}
					select {
					case <-time.After(nextBackoff):
					case <-r.closeChan:
						return
					}

					select {
					case r.transactionsOut <- types.NewTransaction(ts.Payload, resChan):
					case <-r.closeChan:
						return
					}
				} else {
					mSuccess.Incr(1)
					mPartsSuccess.Incr(int64(ts.Payload.Len()))
					resOut = response.NewAck()
					break
				}
			}

			select {
			case ts.ResponseChan <- resOut:
			case <-r.closeChan:
				return
			}
		}(tran, rChan)
	}
}

// Consume assigns a messages channel for the output to read.
func (r *Retry) Consume(ts <-chan types.Transaction) error {
	if r.transactionsIn != nil {
		return types.ErrAlreadyStarted
	}
	if err := r.wrapped.Consume(r.transactionsOut); err != nil {
		return err
	}
	r.transactionsIn = ts
	go r.loop()
	return nil
}

// Connected returns a boolean indicating whether this output is currently
// connected to its target.
func (r *Retry) Connected() bool {
	return r.wrapped.Connected()
}

// MaxInFlight returns the maximum number of in flight messages permitted by the
// output. This value can be used to determine a sensible value for parent
// outputs, but should not be relied upon as part of dispatcher logic.
func (r *Retry) MaxInFlight() (int, bool) {
	return output.GetMaxInFlight(r.wrapped)
}

// CloseAsync shuts down the Retry input and stops processing requests.
func (r *Retry) CloseAsync() {
	if atomic.CompareAndSwapInt32(&r.running, 1, 0) {
		close(r.closeChan)
	}
}

// WaitForClose blocks until the Retry input has closed down.
func (r *Retry) WaitForClose(timeout time.Duration) error {
	select {
	case <-r.closedChan:
	case <-time.After(timeout):
		return types.ErrTimeout
	}
	return nil
}

//------------------------------------------------------------------------------
<commit_msg>Update log level for non-final retries in Retry output (#789)<commit_after>package output

import (
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Jeffail/benthos/v3/internal/component/output"
	"github.com/Jeffail/benthos/v3/internal/docs"
	"github.com/Jeffail/benthos/v3/lib/log"
	"github.com/Jeffail/benthos/v3/lib/metrics"
	"github.com/Jeffail/benthos/v3/lib/response"
	"github.com/Jeffail/benthos/v3/lib/types"
	"github.com/Jeffail/benthos/v3/lib/util/retries"
	"github.com/cenkalti/backoff/v4"
)

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

func init() {
	Constructors[TypeRetry] = TypeSpec{
		constructor: fromSimpleConstructor(NewRetry),
		Summary: `
Attempts to write messages to a child output and if the write fails for any
reason the message is retried either until success or, if the retries or max
elapsed time fields are non-zero, either is reached.`,
		Description: `
All messages in Benthos are always retried on an output error, but this would
usually involve propagating the error back to the source of the message, whereby
it would be reprocessed before reaching the output layer once again.

This output type is useful whenever we wish to avoid reprocessing a message on
the event of a failed send. We might, for example, have a dedupe processor that
we want to avoid reapplying to the same message more than once in the pipeline.

Rather than retrying the same output you may wish to retry the send using a
different output target (a dead letter queue). In which case you should instead
use the ` + "[`try`](/docs/components/outputs/try)" + ` output type.`,
		FieldSpecs: retries.FieldSpecs().Add(
			docs.FieldCommon("output", "A child output.").HasType(docs.FieldOutput),
		),
		Categories: []Category{
			CategoryUtility,
		},
	}
}

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

// RetryConfig contains configuration values for the Retry output type.
type RetryConfig struct {
	Output         *Config `json:"output" yaml:"output"`
	retries.Config `json:",inline" yaml:",inline"`
}

// NewRetryConfig creates a new RetryConfig with default values.
func NewRetryConfig() RetryConfig {
	rConf := retries.NewConfig()
	rConf.MaxRetries = 0
	rConf.Backoff.InitialInterval = "100ms"
	rConf.Backoff.MaxInterval = "1s"
	rConf.Backoff.MaxElapsedTime = "0s"
	return RetryConfig{
		Output: nil,
		Config: retries.NewConfig(),
	}
}

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

type dummyRetryConfig struct {
	Output         interface{} `json:"output" yaml:"output"`
	retries.Config `json:",inline" yaml:",inline"`
}

// MarshalJSON prints an empty object instead of nil.
func (r RetryConfig) MarshalJSON() ([]byte, error) {
	dummy := dummyRetryConfig{
		Output: r.Output,
		Config: r.Config,
	}
	if r.Output == nil {
		dummy.Output = struct{}{}
	}
	return json.Marshal(dummy)
}

// MarshalYAML prints an empty object instead of nil.
func (r RetryConfig) MarshalYAML() (interface{}, error) {
	dummy := dummyRetryConfig{
		Output: r.Output,
		Config: r.Config,
	}
	if r.Output == nil {
		dummy.Output = struct{}{}
	}
	return dummy, nil
}

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

// Retry is an output type that continuously writes a message to a child output
// until the send is successful.
type Retry struct {
	running int32
	conf    RetryConfig

	wrapped     Type
	backoffCtor func() backoff.BackOff

	stats metrics.Type
	log   log.Modular

	transactionsIn  <-chan types.Transaction
	transactionsOut chan types.Transaction

	closeChan  chan struct{}
	closedChan chan struct{}
}

// NewRetry creates a new Retry input type.
func NewRetry(
	conf Config,
	mgr types.Manager,
	log log.Modular,
	stats metrics.Type,
) (Type, error) {
	if conf.Retry.Output == nil {
		return nil, errors.New("cannot create retry output without a child")
	}

	wrapped, err := New(*conf.Retry.Output, mgr, log, stats)
	if err != nil {
		return nil, fmt.Errorf("failed to create output '%v': %v", conf.Retry.Output.Type, err)
	}

	var boffCtor func() backoff.BackOff
	if boffCtor, err = conf.Retry.GetCtor(); err != nil {
		return nil, err
	}

	return &Retry{
		running: 1,
		conf:    conf.Retry,

		log:             log,
		stats:           stats,
		wrapped:         wrapped,
		backoffCtor:     boffCtor,
		transactionsOut: make(chan types.Transaction),

		closeChan:  make(chan struct{}),
		closedChan: make(chan struct{}),
	}, nil
}

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

func (r *Retry) loop() {
	// Metrics paths
	var (
		mRunning      = r.stats.GetGauge("retry.running")
		mCount        = r.stats.GetCounter("retry.count")
		mSuccess      = r.stats.GetCounter("retry.send.success")
		mPartsSuccess = r.stats.GetCounter("retry.parts.send.success")
		mError        = r.stats.GetCounter("retry.send.error")
		mEndOfRetries = r.stats.GetCounter("retry.end_of_retries")
	)

	wg := sync.WaitGroup{}

	defer func() {
		wg.Wait()
		close(r.transactionsOut)
		r.wrapped.CloseAsync()
		err := r.wrapped.WaitForClose(time.Second)
		for ; err != nil; err = r.wrapped.WaitForClose(time.Second) {
		}
		mRunning.Decr(1)
		close(r.closedChan)
	}()
	mRunning.Incr(1)

	errInterruptChan := make(chan struct{})
	var errLooped int64

	for atomic.LoadInt32(&r.running) == 1 {
		// Do not consume another message while pending messages are being
		// reattempted.
		for atomic.LoadInt64(&errLooped) > 0 {
			select {
			case <-errInterruptChan:
			case <-time.After(time.Millisecond * 100):
				// Just incase an interrupt doesn't arrive.
			case <-r.closeChan:
				return
			}
		}

		var tran types.Transaction
		var open bool
		select {
		case tran, open = <-r.transactionsIn:
			if !open {
				return
			}
			mCount.Incr(1)
		case <-r.closeChan:
			return
		}

		rChan := make(chan types.Response)
		select {
		case r.transactionsOut <- types.NewTransaction(tran.Payload, rChan):
		case <-r.closeChan:
			return
		}

		wg.Add(1)
		go func(ts types.Transaction, resChan chan types.Response) {
			var backOff backoff.BackOff
			var resOut types.Response
			var inErrLoop bool

			defer func() {
				wg.Done()
				if inErrLoop {
					atomic.AddInt64(&errLooped, -1)

					// We're exiting our error loop, so (attempt to) interrupt the
					// consumer.
					select {
					case errInterruptChan <- struct{}{}:
					default:
					}
				}
			}()

			for atomic.LoadInt32(&r.running) == 1 {
				var res types.Response
				select {
				case res = <-resChan:
				case <-r.closeChan:
					return
				}

				if res.Error() != nil {
					if !inErrLoop {
						inErrLoop = true
						atomic.AddInt64(&errLooped, 1)
					}

					mError.Incr(1)

					if backOff == nil {
						backOff = r.backoffCtor()
					}

					nextBackoff := backOff.NextBackOff()
					if nextBackoff == backoff.Stop {
						mEndOfRetries.Incr(1)
						r.log.Errorf("Failed to send message: %v\n", res.Error())
						resOut = response.NewNoack()
						break
					} else {
						r.log.Warnf("Failed to send message: %v\n", res.Error())
					}
					select {
					case <-time.After(nextBackoff):
					case <-r.closeChan:
						return
					}

					select {
					case r.transactionsOut <- types.NewTransaction(ts.Payload, resChan):
					case <-r.closeChan:
						return
					}
				} else {
					mSuccess.Incr(1)
					mPartsSuccess.Incr(int64(ts.Payload.Len()))
					resOut = response.NewAck()
					break
				}
			}

			select {
			case ts.ResponseChan <- resOut:
			case <-r.closeChan:
				return
			}
		}(tran, rChan)
	}
}

// Consume assigns a messages channel for the output to read.
func (r *Retry) Consume(ts <-chan types.Transaction) error {
	if r.transactionsIn != nil {
		return types.ErrAlreadyStarted
	}
	if err := r.wrapped.Consume(r.transactionsOut); err != nil {
		return err
	}
	r.transactionsIn = ts
	go r.loop()
	return nil
}

// Connected returns a boolean indicating whether this output is currently
// connected to its target.
func (r *Retry) Connected() bool {
	return r.wrapped.Connected()
}

// MaxInFlight returns the maximum number of in flight messages permitted by the
// output. This value can be used to determine a sensible value for parent
// outputs, but should not be relied upon as part of dispatcher logic.
func (r *Retry) MaxInFlight() (int, bool) {
	return output.GetMaxInFlight(r.wrapped)
}

// CloseAsync shuts down the Retry input and stops processing requests.
func (r *Retry) CloseAsync() {
	if atomic.CompareAndSwapInt32(&r.running, 1, 0) {
		close(r.closeChan)
	}
}

// WaitForClose blocks until the Retry input has closed down.
func (r *Retry) WaitForClose(timeout time.Duration) error {
	select {
	case <-r.closedChan:
	case <-time.After(timeout):
		return types.ErrTimeout
	}
	return nil
}

//------------------------------------------------------------------------------
<|endoftext|>


<commit_before>package govaluate

import (
	"fmt"
	"errors"
	"bytes"
	"math"
	"time"
)

/*
	EvaluableExpression represents a set of ExpressionTokens which, taken together,
	represent an arbitrary expression that can be evaluated down into a single value.
*/
type EvaluableExpression struct {

	/*
		Represents the query format used to output dates. Typically only used when creating SQL or Mongo queries from an expression.
		Defaults to the complete ISO8601 format, including nanoseconds.
	*/
	QueryDateFormat string;

	tokens []ExpressionToken
	inputExpression string
}

/*
	Creates a new EvaluableExpression from the given [expression] string.
	Returns an error if the given expression has invalid syntax.
*/
func NewEvaluableExpression(expression string) (*EvaluableExpression, error) {

	var ret *EvaluableExpression;
	var err error

	ret = new(EvaluableExpression)
	ret.QueryDateFormat = "2006-01-02T15:04:05.999999999Z0700";
	ret.inputExpression = expression;
	ret.tokens, err = parseTokens(expression)

	if(err != nil) {
		return nil, err
	}
	return ret, nil
}

/*
	Evaluate runs the entire expression using the given [parameters].
	Each parameter is mapped from a string to a value, such as "foo" = 1.0.
	If the expression contains a reference to the variable "foo", it will be taken from parameters["foo"].

	This function returns errors if the combination of expression and parameters cannot be run,
	such as if a string parameter is given in an expression that expects it to be a boolean.
	e.g., "foo == true", where foo is any string.
	These errors are almost exclusively returned for parameters not being present, or being of the wrong type.
	Structural problems with the expression (unexpected tokens, unexpected end of expression, etc) are discovered
	during parsing of the expression in NewEvaluableExpression.

	In all non-error circumstances, this returns the single value result of the expression and parameters given.
	e.g., if the expression is "1 + 1", Evaluate will return 2.0.
	e.g., if the expression is "foo + 1" and parameters contains "foo" = 2, Evaluate will return 3.0
*/
func (this EvaluableExpression) Evaluate(parameters map[string]interface{}) (interface{}, error) {

	var stream *tokenStream;

	stream = newTokenStream(this.tokens);
	return evaluateTokens(stream, parameters);
}

func evaluateTokens(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	if(stream.hasNext()) {
		return evaluateLogical(stream, parameters);
	}
	return nil, nil;
}

func evaluateLogical(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateComparator(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = LOGICAL_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break;
		}

		switch(symbol) {

			case OR		:	if(value != nil) {
							return evaluateLogical(stream, parameters);
						} else {
							value, err = evaluateComparator(stream, parameters);
						}
			case AND	:	if(value == nil) {
							return evaluateLogical(stream, parameters);
						} else {
							value, err = evaluateComparator(stream, parameters);
						}
		}

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

	stream.rewind();
	return value, nil;
}

func evaluateComparator(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value, rightValue interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateAdditiveModifier(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = COMPARATOR_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break
		}

		rightValue, err = evaluateAdditiveModifier(stream, parameters);
		if(err != nil) {
			return nil, err;
		}

		switch(symbol) {

			case LT		:	return (value.(float64) < rightValue.(float64)), nil;
			case LTE	:	return (value.(float64) <= rightValue.(float64)), nil;
			case GT		:	return (value.(float64) > rightValue.(float64)), nil;
			case GTE	:	return (value.(float64) >= rightValue.(float64)), nil;
			case EQ		:	return (value == rightValue), nil;
			case NEQ	:	return (value != rightValue), nil;
		}
	}

	stream.rewind();
	return value, nil;
}

func evaluateAdditiveModifier(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value, rightValue interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateMultiplicativeModifier(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = MODIFIER_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break;
		}

		switch(symbol) {

			case PLUS	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						value = value.(float64) + rightValue.(float64);

			case MINUS	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}

						return value.(float64) - rightValue.(float64), nil;

			default		:	stream.rewind();
						return value, nil;
		}
	}

	stream.rewind();
	return value, nil;
}

func evaluateMultiplicativeModifier(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value, rightValue interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateExponentialModifier(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = MODIFIER_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break;
		}

		switch(symbol) {

			case MULTIPLY	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						return value.(float64) * rightValue.(float64), nil;

			case DIVIDE	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						return value.(float64) / rightValue.(float64), nil;

			case MODULUS	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						return math.Mod(value.(float64), rightValue.(float64)), nil;

			default		:	stream.rewind();
						return value, nil;
		}
	}

	stream.rewind();
	return value, nil;
}

func evaluateExponentialModifier(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value, rightValue interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateValue(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = MODIFIER_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break;
		}

		switch(symbol) {

			case EXPONENT	:	rightValue, err = evaluateExponentialModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						return math.Pow(value.(float64), rightValue.(float64)), nil;

			default		:	stream.rewind();
						return value, nil;
		}
	}

	stream.rewind();
	return value, nil;
}

func evaluateValue(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value interface{};
	var errorMessage, variableName string;
	var err error;

	token = stream.next();

	switch(token.Kind) {

		case CLAUSE	:	value, err = evaluateTokens(stream, parameters);
					if(err != nil) {
						return nil, err;
					}

					token = stream.next();
					if(token.Kind != CLAUSE_CLOSE) {

						return nil, errors.New("Unbalanced parenthesis");
					}

					return value, nil;

		case VARIABLE	:	variableName = token.Value.(string);
					value = parameters[variableName];

					if(value == nil) {
						errorMessage = "No parameter '"+ variableName +"' found."
						return nil, errors.New(errorMessage);
					}

					return value, nil;

		case NUMERIC	:	fallthrough
		case STRING	:	fallthrough
		case BOOLEAN	:	return token.Value, nil;
		case TIME	:	return float64(token.Value.(time.Time).Unix()), nil;
		default		:	break;
	}

	stream.rewind();
	return nil, errors.New("Unable to evaluate token kind: " + GetTokenKindString(token.Kind));
}

/*
	Returns a string representing this expression as if it were written in SQL.
	This function assumes that all parameters exist within the same table, and that the table essentially represents
	a serialized object of some sort (e.g., hibernate).
	If your data model is more normalized, you may need to consider iterating through each actual token given by `Tokens()`
	to create your query.

	Boolean values are considered to be "1" for true, "0" for false.

	Times are formatted according to this.QueryDateFormat.
*/
func (this EvaluableExpression) ToSQLQuery() (string, error) {

	var stream *tokenStream;
	var token ExpressionToken;
	var retBuffer bytes.Buffer;
	var toWrite, ret string;

	stream = newTokenStream(this.tokens);

	for(stream.hasNext()) {

		token = stream.next();

		switch(token.Kind) {

			case STRING		:	toWrite = fmt.Sprintf("'%v' ", token.Value);
			case TIME		:	toWrite = fmt.Sprintf("'%s' ", token.Value.(time.Time).Format(this.QueryDateFormat));

			case LOGICALOP		:	switch(LOGICAL_SYMBOLS[token.Value.(string)]) {

								case AND	:	toWrite = "AND ";
								case OR		:	toWrite = "OR ";
							}

			case BOOLEAN		:	if(token.Value.(bool)) {
								toWrite = "1 ";
							} else {
								toWrite = "0 ";
							}

			case VARIABLE		:	toWrite = fmt.Sprintf("[%s] ", token.Value.(string));

			case NUMERIC 		:	toWrite = fmt.Sprintf("%g ", token.Value.(float64));

			case COMPARATOR		:	switch(COMPARATOR_SYMBOLS[token.Value.(string)]) {

								case EQ		:	toWrite = "= ";
								case NEQ	:	toWrite = "<> ";
								default		:	toWrite = fmt.Sprintf("%s ", token.Value.(string));
							}

			case MODIFIER		:	toWrite = fmt.Sprintf("%s ", token.Value.(string));
			case CLAUSE		:	toWrite = "( "
			case CLAUSE_CLOSE	:	toWrite = ") "

			default			:	toWrite = fmt.Sprintf("Unrecognized query token '%s' of kind '%s'", token.Value, token.Kind);
							return "", errors.New(toWrite);
		}

		retBuffer.WriteString(toWrite);
	}

	// trim last space.
	ret = retBuffer.String();
	ret = ret[:len(ret)-1];

	return ret, nil;
}

/*
	Returns an array representing the ExpressionTokens that make up this expression.
*/
func (this EvaluableExpression) Tokens() []ExpressionToken {

	return this.tokens;
}

/*
	Returns the original expression used to create this EvaluableExpression.
*/
func (this EvaluableExpression) String() string {

	return this.inputExpression;
}

func isString(value interface{}) bool {

	switch value.(type) {
		case string	:	return true;
		default		:	break;
	}
	return false;
}
<commit_msg>In-progress mongo queries<commit_after>package govaluate

import (
	"fmt"
	"errors"
	"bytes"
	"math"
	"time"
)

const isoDateFormat string = "2006-01-02T15:04:05.999999999Z0700";

/*
	EvaluableExpression represents a set of ExpressionTokens which, taken together,
	represent an arbitrary expression that can be evaluated down into a single value.
*/
type EvaluableExpression struct {

	/*
		Represents the query format used to output dates. Typically only used when creating SQL or Mongo queries from an expression.
		Defaults to the complete ISO8601 format, including nanoseconds.
	*/
	QueryDateFormat string;

	tokens []ExpressionToken
	inputExpression string
}

/*
	Creates a new EvaluableExpression from the given [expression] string.
	Returns an error if the given expression has invalid syntax.
*/
func NewEvaluableExpression(expression string) (*EvaluableExpression, error) {

	var ret *EvaluableExpression;
	var err error

	ret = new(EvaluableExpression)
	ret.QueryDateFormat = isoDateFormat;
	ret.inputExpression = expression;
	ret.tokens, err = parseTokens(expression)

	if(err != nil) {
		return nil, err
	}
	return ret, nil
}

/*
	Evaluate runs the entire expression using the given [parameters].
	Each parameter is mapped from a string to a value, such as "foo" = 1.0.
	If the expression contains a reference to the variable "foo", it will be taken from parameters["foo"].

	This function returns errors if the combination of expression and parameters cannot be run,
	such as if a string parameter is given in an expression that expects it to be a boolean.
	e.g., "foo == true", where foo is any string.
	These errors are almost exclusively returned for parameters not being present, or being of the wrong type.
	Structural problems with the expression (unexpected tokens, unexpected end of expression, etc) are discovered
	during parsing of the expression in NewEvaluableExpression.

	In all non-error circumstances, this returns the single value result of the expression and parameters given.
	e.g., if the expression is "1 + 1", Evaluate will return 2.0.
	e.g., if the expression is "foo + 1" and parameters contains "foo" = 2, Evaluate will return 3.0
*/
func (this EvaluableExpression) Evaluate(parameters map[string]interface{}) (interface{}, error) {

	var stream *tokenStream;

	stream = newTokenStream(this.tokens);
	return evaluateTokens(stream, parameters);
}

func evaluateTokens(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	if(stream.hasNext()) {
		return evaluateLogical(stream, parameters);
	}
	return nil, nil;
}

func evaluateLogical(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateComparator(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = LOGICAL_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break;
		}

		switch(symbol) {

			case OR		:	if(value != nil) {
							return evaluateLogical(stream, parameters);
						} else {
							value, err = evaluateComparator(stream, parameters);
						}
			case AND	:	if(value == nil) {
							return evaluateLogical(stream, parameters);
						} else {
							value, err = evaluateComparator(stream, parameters);
						}
		}

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

	stream.rewind();
	return value, nil;
}

func evaluateComparator(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value, rightValue interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateAdditiveModifier(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = COMPARATOR_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break
		}

		rightValue, err = evaluateAdditiveModifier(stream, parameters);
		if(err != nil) {
			return nil, err;
		}

		switch(symbol) {

			case LT		:	return (value.(float64) < rightValue.(float64)), nil;
			case LTE	:	return (value.(float64) <= rightValue.(float64)), nil;
			case GT		:	return (value.(float64) > rightValue.(float64)), nil;
			case GTE	:	return (value.(float64) >= rightValue.(float64)), nil;
			case EQ		:	return (value == rightValue), nil;
			case NEQ	:	return (value != rightValue), nil;
		}
	}

	stream.rewind();
	return value, nil;
}

func evaluateAdditiveModifier(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value, rightValue interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateMultiplicativeModifier(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = MODIFIER_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break;
		}

		switch(symbol) {

			case PLUS	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						value = value.(float64) + rightValue.(float64);

			case MINUS	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}

						return value.(float64) - rightValue.(float64), nil;

			default		:	stream.rewind();
						return value, nil;
		}
	}

	stream.rewind();
	return value, nil;
}

func evaluateMultiplicativeModifier(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value, rightValue interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateExponentialModifier(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = MODIFIER_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break;
		}

		switch(symbol) {

			case MULTIPLY	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						return value.(float64) * rightValue.(float64), nil;

			case DIVIDE	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						return value.(float64) / rightValue.(float64), nil;

			case MODULUS	:	rightValue, err = evaluateMultiplicativeModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						return math.Mod(value.(float64), rightValue.(float64)), nil;

			default		:	stream.rewind();
						return value, nil;
		}
	}

	stream.rewind();
	return value, nil;
}

func evaluateExponentialModifier(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value, rightValue interface{};
	var symbol OperatorSymbol;
	var err error;
	var keyFound bool;

	value, err = evaluateValue(stream, parameters);

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

	for stream.hasNext() {

		token = stream.next();

		if(!isString(token.Value)) {
			break;
		}

		symbol, keyFound = MODIFIER_SYMBOLS[token.Value.(string)];
		if(!keyFound) {
			break;
		}

		switch(symbol) {

			case EXPONENT	:	rightValue, err = evaluateExponentialModifier(stream, parameters);
						if(err != nil) {
							return nil, err;
						}
						return math.Pow(value.(float64), rightValue.(float64)), nil;

			default		:	stream.rewind();
						return value, nil;
		}
	}

	stream.rewind();
	return value, nil;
}

func evaluateValue(stream *tokenStream, parameters map[string]interface{}) (interface{}, error) {

	var token ExpressionToken;
	var value interface{};
	var errorMessage, variableName string;
	var err error;

	token = stream.next();

	switch(token.Kind) {

		case CLAUSE	:	value, err = evaluateTokens(stream, parameters);
					if(err != nil) {
						return nil, err;
					}

					token = stream.next();
					if(token.Kind != CLAUSE_CLOSE) {

						return nil, errors.New("Unbalanced parenthesis");
					}

					return value, nil;

		case VARIABLE	:	variableName = token.Value.(string);
					value = parameters[variableName];

					if(value == nil) {
						errorMessage = "No parameter '"+ variableName +"' found."
						return nil, errors.New(errorMessage);
					}

					return value, nil;

		case NUMERIC	:	fallthrough
		case STRING	:	fallthrough
		case BOOLEAN	:	return token.Value, nil;
		case TIME	:	return float64(token.Value.(time.Time).Unix()), nil;
		default		:	break;
	}

	stream.rewind();
	return nil, errors.New("Unable to evaluate token kind: " + GetTokenKindString(token.Kind));
}

/*
	Returns a string representing this expression as if it were written in SQL.
	This function assumes that all parameters exist within the same table, and that the table essentially represents
	a serialized object of some sort (e.g., hibernate).
	If your data model is more normalized, you may need to consider iterating through each actual token given by `Tokens()`
	to create your query.

	Boolean values are considered to be "1" for true, "0" for false.

	Times are formatted according to this.QueryDateFormat.
*/
func (this EvaluableExpression) ToSQLQuery() (string, error) {

	var stream *tokenStream;
	var token ExpressionToken;
	var retBuffer bytes.Buffer;
	var toWrite, ret string;

	stream = newTokenStream(this.tokens);

	for(stream.hasNext()) {

		token = stream.next();

		switch(token.Kind) {

			case STRING		:	toWrite = fmt.Sprintf("'%v' ", token.Value);
			case TIME		:	toWrite = fmt.Sprintf("'%s' ", token.Value.(time.Time).Format(this.QueryDateFormat));

			case LOGICALOP		:	switch(LOGICAL_SYMBOLS[token.Value.(string)]) {

								case AND	:	toWrite = "AND ";
								case OR		:	toWrite = "OR ";
							}

			case BOOLEAN		:	if(token.Value.(bool)) {
								toWrite = "1 ";
							} else {
								toWrite = "0 ";
							}

			case VARIABLE		:	toWrite = fmt.Sprintf("[%s] ", token.Value.(string));

			case NUMERIC 		:	toWrite = fmt.Sprintf("%g ", token.Value.(float64));

			case COMPARATOR		:	switch(COMPARATOR_SYMBOLS[token.Value.(string)]) {

								case EQ		:	toWrite = "= ";
								case NEQ	:	toWrite = "<> ";
								default		:	toWrite = fmt.Sprintf("%s ", token.Value.(string));
							}

			case MODIFIER		:	toWrite = fmt.Sprintf("%s ", token.Value.(string));
			case CLAUSE		:	toWrite = "( "
			case CLAUSE_CLOSE	:	toWrite = ") "

			default			:	toWrite = fmt.Sprintf("Unrecognized query token '%s' of kind '%s'", token.Value, token.Kind);
							return "", errors.New(toWrite);
		}

		retBuffer.WriteString(toWrite);
	}

	// trim last space.
	ret = retBuffer.String();
	ret = ret[:len(ret)-1];

	return ret, nil;
}

/*
	Returns a string representing this expression as if it were written as a Mongo query.
*/
func (this EvaluableExpression) ToMongoQuery() (string, error) {

	var stream *tokenStream;
	var token ExpressionToken;
	var retBuffer bytes.Buffer;
	var toWrite, ret string;

	stream = newTokenStream(this.tokens);

	for(stream.hasNext()) {

		token = stream.next();

		switch(token.Kind) {

			case STRING		:	toWrite = fmt.Sprintf("\"%s\" ", token.Value.(string));
			case TIME		:	toWrite = fmt.Sprintf("ISODate(\"%s\") ", token.Value.(time.Time).Format(isoDateFormat));
			case LOGICALOP		:
			case BOOLEAN		:	if(token.Value.(bool)) {
								toWrite = "true ";
							} else {
								toWrite = "false ";
							}
			case VARIABLE		:	toWrite = fmt.Sprintf("%s ", token.Value.(string));
			case NUMERIC 		:	toWrite = fmt.Sprintf("%g ", token.Value.(float64));
			case COMPARATOR		:
			case CLAUSE		:	fallthrough;
			case CLAUSE_CLOSE	:	continue;

			case MODIFIER		:	toWrite = fmt.Sprintf("Unable to use modifiers in Mongo queries (found '%s')", token.Kind);
							return "", errors.New(toWrite);

			default			:	toWrite = fmt.Sprintf("Unrecognized query token '%s' of kind '%s'", token.Value, token.Kind);
							return "", errors.New(toWrite);
		}

		retBuffer.WriteString(toWrite);
	}

	// trim last space.
	ret = retBuffer.String();
	ret = ret[:len(ret)-1];

	return ret, nil;
}

/*
	Returns an array representing the ExpressionTokens that make up this expression.
*/
func (this EvaluableExpression) Tokens() []ExpressionToken {

	return this.tokens;
}

/*
	Returns the original expression used to create this EvaluableExpression.
*/
func (this EvaluableExpression) String() string {

	return this.inputExpression;
}

func isString(value interface{}) bool {

	switch value.(type) {
		case string	:	return true;
		default		:	break;
	}
	return false;
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

func TestParseFrom(t *testing.T) {
	cases := []struct {
		in   []byte
		want *From
	}{
		{[]byte("ubuntu"), &From{Image: "ubuntu"}},
		{[]byte(" ubuntu"), &From{Image: "ubuntu"}},
		{[]byte("ubuntu "), &From{Image: "ubuntu"}},
		{[]byte("ubuntu:14.04"), &From{Image: "ubuntu", Tag: "14.04"}},
		{[]byte(" ubuntu:14.04"), &From{Image: "ubuntu", Tag: "14.04"}},
		{[]byte("ubuntu:14.04 "), &From{Image: "ubuntu", Tag: "14.04"}},
		{[]byte("ubuntu@12345"), &From{Image: "ubuntu", Digest: "12345"}},
		{[]byte(" ubuntu@12345"), &From{Image: "ubuntu", Digest: "12345"}},
		{[]byte("ubuntu@12345 "), &From{Image: "ubuntu", Digest: "12345"}},
	}
	for _, c := range cases {
		got, _ := ParseFrom(c.in)
		if got.Image != c.want.Image {
			t.Errorf(
				"ParseFrom(%q).Image == %q, want %q",
				c.in,
				got.Image,
				c.want.Image,
			)
		}
		if got.Tag != c.want.Tag {
			t.Errorf(
				"ParseFrom(%q).Tag == %q, want %q",
				c.in,
				got.Tag,
				c.want.Tag,
			)
		}
		if got.Digest != c.want.Digest {
			t.Errorf(
				"ParseFrom(%q).Digest == %q, want %q",
				c.in,
				got.Digest,
				c.want.Digest,
			)
		}
	}
}

func TestParseMaintainer(t *testing.T) {
	cases := []struct {
		in   []byte
		want *Maintainer
	}{
		{[]byte("JohnDoe"), &Maintainer{Name: "JohnDoe"}},
		{[]byte("John Doe"), &Maintainer{Name: "John Doe"}},
		{[]byte(" John Doe"), &Maintainer{Name: "John Doe"}},
		{[]byte("John Doe "), &Maintainer{Name: "John Doe"}},
	}
	for _, c := range cases {
		got, _ := ParseMaintainer(c.in)
		if got.Name != c.want.Name {
			t.Errorf(
				"ParseMaintainer(%q).Name == %q, want %q",
				c.in,
				got.Name,
				c.want.Name,
			)
		}
	}
}

func TestParseExpose(t *testing.T) {
	cases := []struct {
		in   []byte
		want *Expose
	}{
		{[]byte("22"), &Expose{Ports: []int{22}}},
		{[]byte(" 22"), &Expose{Ports: []int{22}}},
		{[]byte("22 "), &Expose{Ports: []int{22}}},
		{[]byte("22 80"), &Expose{Ports: []int{22, 80}}},
		{[]byte(" 22 80"), &Expose{Ports: []int{22, 80}}},
		{[]byte("22 80 "), &Expose{Ports: []int{22, 80}}},
	}
	for _, c := range cases {
		got, _ := ParseExpose(c.in)
		for i, p := range got.Ports {
			if p != c.want.Ports[i] {
				t.Errorf(
					"ParseExpose(%q).Ports == %v, want %v",
					c.in,
					got.Ports,
					c.want.Ports,
				)
			}
		}
	}
}

func TestParseUser(t *testing.T) {
	cases := []struct {
		in   []byte
		want *User
	}{
		{[]byte("root"), &User{Name: "root"}},
		{[]byte(" root"), &User{Name: "root"}},
		{[]byte("root "), &User{Name: "root"}},
	}
	for _, c := range cases {
		got, _ := ParseUser(c.in)
		if got.Name != c.want.Name {
			t.Errorf(
				"ParseUser(%q).Name == %q, want %q",
				c.in,
				got.Name,
				c.want.Name,
			)
		}
	}
}

func TestParseWorkdir(t *testing.T) {
	cases := []struct {
		in   []byte
		want *Workdir
	}{
		{[]byte("/home/foo"), &Workdir{Path: "/home/foo"}},
		{[]byte(" /home/foo"), &Workdir{Path: "/home/foo"}},
		{[]byte("/home/foo "), &Workdir{Path: "/home/foo"}},
	}
	for _, c := range cases {
		got, _ := ParseWorkdir(c.in)
		if got.Path != c.want.Path {
			t.Errorf(
				"ParseWorkdir(%q).Path == %q, want %q",
				c.in,
				got.Path,
				c.want.Path,
			)
		}
	}
}
<commit_msg>Refactor tests to compare structures directly<commit_after>package main

import (
	"testing"
)

func TestParseFrom(t *testing.T) {
	cases := []struct {
		in   []byte
		want *From
	}{
		{[]byte("ubuntu"), &From{Image: "ubuntu"}},
		{[]byte(" ubuntu"), &From{Image: "ubuntu"}},
		{[]byte("ubuntu "), &From{Image: "ubuntu"}},
		{[]byte("ubuntu:14.04"), &From{Image: "ubuntu", Tag: "14.04"}},
		{[]byte(" ubuntu:14.04"), &From{Image: "ubuntu", Tag: "14.04"}},
		{[]byte("ubuntu:14.04 "), &From{Image: "ubuntu", Tag: "14.04"}},
		{[]byte("ubuntu@12345"), &From{Image: "ubuntu", Digest: "12345"}},
		{[]byte(" ubuntu@12345"), &From{Image: "ubuntu", Digest: "12345"}},
		{[]byte("ubuntu@12345 "), &From{Image: "ubuntu", Digest: "12345"}},
	}
	for _, c := range cases {
		got, _ := ParseFrom(c.in)
		if *got != *c.want {
			t.Errorf("ParseFrom(%q) == %v, want %v", c.in, got, c.want)
		}
	}
}

func TestParseMaintainer(t *testing.T) {
	cases := []struct {
		in   []byte
		want *Maintainer
	}{
		{[]byte("JohnDoe"), &Maintainer{Name: "JohnDoe"}},
		{[]byte("John Doe"), &Maintainer{Name: "John Doe"}},
		{[]byte(" John Doe"), &Maintainer{Name: "John Doe"}},
		{[]byte("John Doe "), &Maintainer{Name: "John Doe"}},
	}
	for _, c := range cases {
		got, _ := ParseMaintainer(c.in)
		if *got != *c.want {
			t.Errorf("ParseMaintainer(%q) == %v, want %v", c.in, got, c.want)
		}
	}
}

func TestParseExpose(t *testing.T) {
	cases := []struct {
		in   []byte
		want *Expose
	}{
		{[]byte("22"), &Expose{Ports: []int{22}}},
		{[]byte(" 22"), &Expose{Ports: []int{22}}},
		{[]byte("22 "), &Expose{Ports: []int{22}}},
		{[]byte("22 80"), &Expose{Ports: []int{22, 80}}},
		{[]byte(" 22 80"), &Expose{Ports: []int{22, 80}}},
		{[]byte("22 80 "), &Expose{Ports: []int{22, 80}}},
	}
	for _, c := range cases {
		got, _ := ParseExpose(c.in)
		for i, p := range got.Ports {
			if p != c.want.Ports[i] {
				t.Errorf("ParseExpose(%q) == %v, want %v", c.in, got, c.want)
			}
		}
	}
}

func TestParseUser(t *testing.T) {
	cases := []struct {
		in   []byte
		want *User
	}{
		{[]byte("root"), &User{Name: "root"}},
		{[]byte(" root"), &User{Name: "root"}},
		{[]byte("root "), &User{Name: "root"}},
	}
	for _, c := range cases {
		got, _ := ParseUser(c.in)
		if *got != *c.want {
			t.Errorf("ParseUser(%q) == %v, want %v", c.in, got, c.want)
		}
	}
}

func TestParseWorkdir(t *testing.T) {
	cases := []struct {
		in   []byte
		want *Workdir
	}{
		{[]byte("/home/foo"), &Workdir{Path: "/home/foo"}},
		{[]byte(" /home/foo"), &Workdir{Path: "/home/foo"}},
		{[]byte("/home/foo "), &Workdir{Path: "/home/foo"}},
	}
	for _, c := range cases {
		got, _ := ParseWorkdir(c.in)
		if *got != *c.want {
			t.Errorf("ParseWorkdir(%q) == %v, want %v", c.in, got, c.want)
		}
	}
}
<|endoftext|>


<commit_before>// +build windows, functional

// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 util

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/ec2"
	"github.com/aws/amazon-ecs-agent/agent/ecs_client/model/ecs"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	docker "github.com/fsouza/go-dockerclient"
	"golang.org/x/sys/windows/registry"
)

var ECS *ecs.ECS
var Cluster string

func init() {
	var ecsconfig aws.Config
	if region := os.Getenv("AWS_REGION"); region != "" {
		ecsconfig.Region = &region
	}
	if region := os.Getenv("AWS_DEFAULT_REGION"); region != "" {
		ecsconfig.Region = &region
	}
	if ecsconfig.Region == nil {
		if iid, err := ec2.GetInstanceIdentityDocument(); err == nil {
			ecsconfig.Region = &iid.Region
		}
	}
	if envEndpoint := os.Getenv("ECS_BACKEND_HOST"); envEndpoint != "" {
		ecsconfig.Endpoint = &envEndpoint
	}

	ECS = ecs.New(session.New(&ecsconfig))
	Cluster = "ecs-functional-tests"
	if envCluster := os.Getenv("ECS_CLUSTER"); envCluster != "" {
		Cluster = envCluster
	}
	ECS.CreateCluster(&ecs.CreateClusterInput{
		ClusterName: aws.String(Cluster),
	})
}

// RunAgent launches the agent and returns an object which may be used to reference it.

func RunAgent(t *testing.T, options *AgentOptions) *TestAgent {
	agent := &TestAgent{t: t}

	dockerClient, err := docker.NewClientFromEnv()
	if err != nil {
		t.Fatal(err)
	}
	agent.DockerClient = dockerClient

	tmpdirOverride := os.Getenv("ECS_FTEST_TMP")

	agentTempdir, err := ioutil.TempDir(tmpdirOverride, "ecs_integ_testdata")
	if err != nil {
		t.Fatal("Could not create temp dir for test")
	}
	logdir := filepath.Join(agentTempdir, "log")
	datadir := filepath.Join(agentTempdir, "data")
	os.Setenv("ECS_LOGFILE", logdir+"/log.log")
	os.Setenv("ECS_DATADIR", datadir)
	valueName := fmt.Sprintf("test_path_%d", time.Now().UnixNano())
	fmt.Println("Registry location", valueName)
	os.Setenv("ZZZ_I_KNOW_SETTING_TEST_VALUES_IN_PRODUCTION_IS_NOT_SUPPORTED", valueName)
	os.Setenv("ECS_CLUSTER", Cluster)
	os.Setenv("ECS_ENABLE_TASK_IAM_ROLE", "true")
	os.Setenv("DOCKER_HOST", "npipe:////./pipe/docker_engine")
	os.Setenv("ECS_DISABLE_METRICS", "true")
	os.Setenv("ECS_AUDIT_LOGFILE", logdir+"/audit.log")
	os.Setenv("ECS_LOGLEVEL", "debug")
	os.Setenv("ECS_AVAILABLE_LOGGING_DRIVERS", `["json-file","awslogs"]`)

	t.Log("datadir", datadir)
	os.Mkdir(logdir, 0755)
	os.Mkdir(datadir, 0755)
	agent.TestDir = agentTempdir
	agent.Options = options
	if options == nil {
		agent.Options = &AgentOptions{}
	}
	t.Logf("Created directory %s to store test data in", agentTempdir)

	err = agent.StartAgent()
	if err != nil {
		t.Fatal(err)
	}
	return agent
}

func (agent *TestAgent) StopAgent() error {
	return agent.Process.Kill()
}

func (agent *TestAgent) StartAgent() error {
	agentInvoke := exec.Command(".\\agent.exe")
	if TestDirectory := os.Getenv("ECS_WINDOWS_TEST_DIR"); TestDirectory != "" {
		agentInvoke.Dir = TestDirectory
	}
	agentInvoke.Start()
	agent.Process = agentInvoke.Process
	agent.IntrospectionURL = "http://localhost:51678"
	err := agent.platformIndependentStartAgent()
	return err
}

func (agent *TestAgent) Cleanup() {
	key, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Amazon\ECS Agent\State File`, registry.ALL_ACCESS)
	if err != nil {
		return
	}
	key.DeleteValue(os.Getenv("ZZZ_I_KNOW_SETTING_TEST_VALUES_IN_PRODUCTION_IS_NOT_SUPPORTED"))
	os.Unsetenv("ZZZ_I_KNOW_SETTING_TEST_VALUES_IN_PRODUCTION_IS_NOT_SUPPORTED")
	agent.platformIndependentCleanup()
}
<commit_msg>functional tests: Use t.Log instead of fmt.Println<commit_after>// +build windows, functional

// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 util

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/ec2"
	"github.com/aws/amazon-ecs-agent/agent/ecs_client/model/ecs"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	docker "github.com/fsouza/go-dockerclient"
	"golang.org/x/sys/windows/registry"
)

var ECS *ecs.ECS
var Cluster string

func init() {
	var ecsconfig aws.Config
	if region := os.Getenv("AWS_REGION"); region != "" {
		ecsconfig.Region = &region
	}
	if region := os.Getenv("AWS_DEFAULT_REGION"); region != "" {
		ecsconfig.Region = &region
	}
	if ecsconfig.Region == nil {
		if iid, err := ec2.GetInstanceIdentityDocument(); err == nil {
			ecsconfig.Region = &iid.Region
		}
	}
	if envEndpoint := os.Getenv("ECS_BACKEND_HOST"); envEndpoint != "" {
		ecsconfig.Endpoint = &envEndpoint
	}

	ECS = ecs.New(session.New(&ecsconfig))
	Cluster = "ecs-functional-tests"
	if envCluster := os.Getenv("ECS_CLUSTER"); envCluster != "" {
		Cluster = envCluster
	}
	ECS.CreateCluster(&ecs.CreateClusterInput{
		ClusterName: aws.String(Cluster),
	})
}

// RunAgent launches the agent and returns an object which may be used to reference it.

func RunAgent(t *testing.T, options *AgentOptions) *TestAgent {
	agent := &TestAgent{t: t}

	dockerClient, err := docker.NewClientFromEnv()
	if err != nil {
		t.Fatal(err)
	}
	agent.DockerClient = dockerClient

	tmpdirOverride := os.Getenv("ECS_FTEST_TMP")

	agentTempdir, err := ioutil.TempDir(tmpdirOverride, "ecs_integ_testdata")
	if err != nil {
		t.Fatal("Could not create temp dir for test")
	}
	logdir := filepath.Join(agentTempdir, "log")
	datadir := filepath.Join(agentTempdir, "data")
	os.Setenv("ECS_LOGFILE", logdir+"/log.log")
	os.Setenv("ECS_DATADIR", datadir)
	valueName := fmt.Sprintf("test_path_%d", time.Now().UnixNano())
	t.Log("Registry location", valueName)
	os.Setenv("ZZZ_I_KNOW_SETTING_TEST_VALUES_IN_PRODUCTION_IS_NOT_SUPPORTED", valueName)
	os.Setenv("ECS_CLUSTER", Cluster)
	os.Setenv("ECS_ENABLE_TASK_IAM_ROLE", "true")
	os.Setenv("DOCKER_HOST", "npipe:////./pipe/docker_engine")
	os.Setenv("ECS_DISABLE_METRICS", "true")
	os.Setenv("ECS_AUDIT_LOGFILE", logdir+"/audit.log")
	os.Setenv("ECS_LOGLEVEL", "debug")
	os.Setenv("ECS_AVAILABLE_LOGGING_DRIVERS", `["json-file","awslogs"]`)

	t.Log("datadir", datadir)
	os.Mkdir(logdir, 0755)
	os.Mkdir(datadir, 0755)
	agent.TestDir = agentTempdir
	agent.Options = options
	if options == nil {
		agent.Options = &AgentOptions{}
	}
	t.Logf("Created directory %s to store test data in", agentTempdir)

	err = agent.StartAgent()
	if err != nil {
		t.Fatal(err)
	}
	return agent
}

func (agent *TestAgent) StopAgent() error {
	return agent.Process.Kill()
}

func (agent *TestAgent) StartAgent() error {
	agentInvoke := exec.Command(".\\agent.exe")
	if TestDirectory := os.Getenv("ECS_WINDOWS_TEST_DIR"); TestDirectory != "" {
		agentInvoke.Dir = TestDirectory
	}
	agentInvoke.Start()
	agent.Process = agentInvoke.Process
	agent.IntrospectionURL = "http://localhost:51678"
	err := agent.platformIndependentStartAgent()
	return err
}

func (agent *TestAgent) Cleanup() {
	key, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Amazon\ECS Agent\State File`, registry.ALL_ACCESS)
	if err != nil {
		return
	}
	key.DeleteValue(os.Getenv("ZZZ_I_KNOW_SETTING_TEST_VALUES_IN_PRODUCTION_IS_NOT_SUPPORTED"))
	os.Unsetenv("ZZZ_I_KNOW_SETTING_TEST_VALUES_IN_PRODUCTION_IS_NOT_SUPPORTED")
	agent.platformIndependentCleanup()
}
<|endoftext|>


<commit_before>package URLHandler

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

type ETag string

// A map of extra things to pass to every request handler call
var extras map[string]interface{}

// URLHandler is an interface to describe a request to a URL
//
// After being registered to handle a URL with a RegisterHandler
// call, the URLHandler will handle any requests to that URL by
// delegating to the method for the appropriate HTTP Method being
// called.
//
// All methods receive the http.Request object, and a map of extra
// parameters that have been registered with RegisterExtraParameter
type URLHandler interface {
	// Get will handle an HTTP GET request to this URL and return the
	// content that should be sent to the client
	Get(r *http.Request, params map[string]interface{}) (string, error)

	// Post will handle an HTTP POST request to this URL.
	// Post returns 2 strings: the content to return, an a redirectURL
	// If the redirectURL is not the empty string, the registered
	// URLandler will automatically respond with a 303 return code
	// instead of a 200 return code, and set an appropriate Location:
	// response header
	Post(r *http.Request, params map[string]interface{}) (content, redirectURL string, err error)

	// Put will handle an HTTP PUT request to this URL and return the
	// content that should be sent to the client
	Put(r *http.Request, params map[string]interface{}) (string, error)

	// Delete will handle an HTTP PUT request to this URL and return the
	// content that should be sent to the client
	Delete(r *http.Request, params map[string]interface{}) (string, error)

	// Calculate an ETag to represent the resource being served by
	// this handler, so that a registered handler can return a 304
	// code if the resource hasn't changed.
	ETag(*url.URL, map[string]interface{}) (ETag, error)
}

// handleClientError takes an error from a URLHandler and returns
// an appropriate response if it knows how. Returns true if it's been
// handled, false otherwise
func handleClientError(w http.ResponseWriter, response string, err error) bool {
	switch err.(type) {
	case BadRequestError:
		w.WriteHeader(400)
		fmt.Fprintf(w, response)
		return true
	case ForbiddenError:
		w.WriteHeader(403)
		fmt.Fprintf(w, response)
		return true
	case NotFoundError:
		w.WriteHeader(404)
		fmt.Fprintf(w, response)
		return true
	case InvalidMethodError:
		w.WriteHeader(405)
		fmt.Fprintf(w, response)
		return true
	}
	return false
}

// handleError if a helper function to handle errors from URLHandlers.
// Mostly, it calls handleClientError and then panics if it didn't get
// handled.
func handleError(w http.ResponseWriter, response string, err error) {
	handled := handleClientError(w, response, err)
	if handled {
		return
	}
	panic("Something happened")
}

func enforceIfMatch(success func(*http.Request, map[string]interface{}) (string, error), h URLHandler, w http.ResponseWriter, r *http.Request) {

	etag, _ := h.ETag(r.URL, extras)
	if etag != "" && r.Header.Get("If-Match") == "" {
		w.WriteHeader(428)
		fmt.Fprintf(w, "Must include ETag in If-Match header to ensure resource has not been modified")
		return
	}
	if etag == "" || string(etag) == r.Header.Get("If-Match") {
		response, err := success(r, extras)
		if err != nil {
			handleError(w, response, err)
			return
		}
		fmt.Fprintf(w, response)
	} else {
		w.WriteHeader(412)
		fmt.Fprintf(w, "If-Match header does not match document ETag")
	}
}

// RegisterHandler takes a URLHandler and a url string and registers
// that URLHandler to handle that URL. It automatically registers an
// http.HandleFunc which delegates to the appropriate URLHandler method
func RegisterHandler(h URLHandler, url string) {
	var handler = func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if r := recover(); r != nil {
				w.WriteHeader(500)
				fmt.Fprintf(w, "Unknown server error")
			}
		}()

		switch r.Method {
		case "HEAD":
			etag, err := h.ETag(r.URL, extras)
			switch err.(type) {
			case NotFoundError:
				w.WriteHeader(404)
				return
			}
			if etag != "" {
				w.Header().Add("ETag", string(etag))
			}
		case "GET":
			if etag, _ := h.ETag(r.URL, extras); etag != "" {
				w.Header().Add("ETag", string(etag))
				if string(etag) == r.Header.Get("If-None-Match") {
					w.WriteHeader(304)
					return
				}
			}
			response, err := h.Get(r, extras)
			if err != nil {
				handleError(w, response, err)
				return
			}
			fmt.Fprintf(w, response)
		case "POST":
			response, redirectURL, err := h.Post(r, extras)
			if err != nil {
				handleError(w, response, err)
				return
			}
			if redirectURL != "" {
				w.Header().Add("Location", redirectURL)
				w.WriteHeader(303)
			}
			fmt.Fprintf(w, response)
		case "DELETE":
			enforceIfMatch(h.Delete, h, w, r)
		case "PUT":
			enforceIfMatch(h.Put, h, w, r)
		default:
			w.WriteHeader(501)

		}
	}
	http.HandleFunc(url, handler)
}

// RegisterStaticHandler registers directory to be served by the web
// server on the filesystem without going through the handler function.
func RegisterStaticHandler(url, directoryRoot string) {
	http.Handle(url, http.StripPrefix(url, http.FileServer(http.Dir(directoryRoot))))
}

// RegisterExtraParameter allows you to add arbitrary data to get
// passed to the params parameter of URLHandler handler functions which
// you can retrieve from params[key] (and will need to manually cast to
// the appropriate type.
//
// This is useful for passing, for instance, a pointer to an sql.DB,
// or any configuration data you want to use throughout your web app
func RegisterExtraParameter(key string, obj interface{}) {
	if extras == nil {
		extras = make(map[string]interface{})
	}
	extras[key] = obj
}
<commit_msg>Print error to the console when there's a 500 error<commit_after>package URLHandler

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

type ETag string

// A map of extra things to pass to every request handler call
var extras map[string]interface{}

// URLHandler is an interface to describe a request to a URL
//
// After being registered to handle a URL with a RegisterHandler
// call, the URLHandler will handle any requests to that URL by
// delegating to the method for the appropriate HTTP Method being
// called.
//
// All methods receive the http.Request object, and a map of extra
// parameters that have been registered with RegisterExtraParameter
type URLHandler interface {
	// Get will handle an HTTP GET request to this URL and return the
	// content that should be sent to the client
	Get(r *http.Request, params map[string]interface{}) (string, error)

	// Post will handle an HTTP POST request to this URL.
	// Post returns 2 strings: the content to return, an a redirectURL
	// If the redirectURL is not the empty string, the registered
	// URLandler will automatically respond with a 303 return code
	// instead of a 200 return code, and set an appropriate Location:
	// response header
	Post(r *http.Request, params map[string]interface{}) (content, redirectURL string, err error)

	// Put will handle an HTTP PUT request to this URL and return the
	// content that should be sent to the client
	Put(r *http.Request, params map[string]interface{}) (string, error)

	// Delete will handle an HTTP PUT request to this URL and return the
	// content that should be sent to the client
	Delete(r *http.Request, params map[string]interface{}) (string, error)

	// Calculate an ETag to represent the resource being served by
	// this handler, so that a registered handler can return a 304
	// code if the resource hasn't changed.
	ETag(*url.URL, map[string]interface{}) (ETag, error)
}

// handleClientError takes an error from a URLHandler and returns
// an appropriate response if it knows how. Returns true if it's been
// handled, false otherwise
func handleClientError(w http.ResponseWriter, response string, err error) bool {
	switch err.(type) {
	case BadRequestError:
		w.WriteHeader(400)
		fmt.Fprintf(w, response)
		return true
	case ForbiddenError:
		w.WriteHeader(403)
		fmt.Fprintf(w, response)
		return true
	case NotFoundError:
		w.WriteHeader(404)
		fmt.Fprintf(w, response)
		return true
	case InvalidMethodError:
		w.WriteHeader(405)
		fmt.Fprintf(w, response)
		return true
	}
	return false
}

// handleError if a helper function to handle errors from URLHandlers.
// Mostly, it calls handleClientError and then panics if it didn't get
// handled.
func handleError(w http.ResponseWriter, response string, err error) {
	handled := handleClientError(w, response, err)
	if handled {
		return
	}
	panic("Something happened")
}

func enforceIfMatch(success func(*http.Request, map[string]interface{}) (string, error), h URLHandler, w http.ResponseWriter, r *http.Request) {

	etag, _ := h.ETag(r.URL, extras)
	if etag != "" && r.Header.Get("If-Match") == "" {
		w.WriteHeader(428)
		fmt.Fprintf(w, "Must include ETag in If-Match header to ensure resource has not been modified")
		return
	}
	if etag == "" || string(etag) == r.Header.Get("If-Match") {
		response, err := success(r, extras)
		if err != nil {
			handleError(w, response, err)
			return
		}
		fmt.Fprintf(w, response)
	} else {
		w.WriteHeader(412)
		fmt.Fprintf(w, "If-Match header does not match document ETag")
	}
}

// RegisterHandler takes a URLHandler and a url string and registers
// that URLHandler to handle that URL. It automatically registers an
// http.HandleFunc which delegates to the appropriate URLHandler method
func RegisterHandler(h URLHandler, url string) {
	var handler = func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if r := recover(); r != nil {
				w.WriteHeader(500)
				log.Print(r)
				fmt.Fprintf(w, "Unknown server error")
			}
		}()

		switch r.Method {
		case "HEAD":
			etag, err := h.ETag(r.URL, extras)
			switch err.(type) {
			case NotFoundError:
				w.WriteHeader(404)
				return
			}
			if etag != "" {
				w.Header().Add("ETag", string(etag))
			}
		case "GET":
			if etag, _ := h.ETag(r.URL, extras); etag != "" {
				w.Header().Add("ETag", string(etag))
				if string(etag) == r.Header.Get("If-None-Match") {
					w.WriteHeader(304)
					return
				}
			}
			response, err := h.Get(r, extras)
			if err != nil {
				handleError(w, response, err)
				return
			}
			fmt.Fprintf(w, response)
		case "POST":
			response, redirectURL, err := h.Post(r, extras)
			if err != nil {
				handleError(w, response, err)
				return
			}
			if redirectURL != "" {
				w.Header().Add("Location", redirectURL)
				w.WriteHeader(303)
			}
			fmt.Fprintf(w, response)
		case "DELETE":
			enforceIfMatch(h.Delete, h, w, r)
		case "PUT":
			enforceIfMatch(h.Put, h, w, r)
		default:
			w.WriteHeader(501)

		}
	}
	http.HandleFunc(url, handler)
}

// RegisterStaticHandler registers directory to be served by the web
// server on the filesystem without going through the handler function.
func RegisterStaticHandler(url, directoryRoot string) {
	http.Handle(url, http.StripPrefix(url, http.FileServer(http.Dir(directoryRoot))))
}

// RegisterExtraParameter allows you to add arbitrary data to get
// passed to the params parameter of URLHandler handler functions which
// you can retrieve from params[key] (and will need to manually cast to
// the appropriate type.
//
// This is useful for passing, for instance, a pointer to an sql.DB,
// or any configuration data you want to use throughout your web app
func RegisterExtraParameter(key string, obj interface{}) {
	if extras == nil {
		extras = make(map[string]interface{})
	}
	extras[key] = obj
}
<|endoftext|>


<commit_before>// +build !windows
package duktape

import (
	"os"
	"os/signal"
	"syscall"
	"testing"
	"time"
)

func testCallError() {
	ctx := NewContext()
	defer ctx.DestroyHeap()
	ctx.Error(ErrType, "Nope: ", 500)
}

func testCallErrorf() {
	ctx := NewContext()
	defer ctx.DestroyHeap()
	ctx.Error(ErrType, "Nope: %d", 0xdeadbeef)
}

func waitForAbort(t *testing.T, sig <-chan os.Signal, expectedSig os.Signal) {
	const wait = 500 * time.Millisecond
	select {
	case s := <-sig:
		if s != expectedSig {
			t.Error("Failed to receive %v, but did receive %v", expectedSig, s)
			return
		}
		t.Logf("Received expected signal: %v == %v", expectedSig, s)
	case <-time.After(wait):
		t.Errorf("Failed to receive %v after %v", expectedSig, wait)
	}
}

// TestErrorAbort tests for whether SIGABRT is received when Context.Error[f]
// is called.
func TestErrorAbort(t *testing.T) {
	abrt := os.Signal(syscall.SIGABRT)
	sig := make(chan os.Signal, 1)

	signal.Notify(sig, abrt)
	defer signal.Stop(sig)

	go testCallError()
	waitForAbort(t, sig, abrt)

	go testCallErrorf()
	waitForAbort(t, sig, abrt)
}
<commit_msg>error_test.go omitted<commit_after><|endoftext|>