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

// Package redirect provides hooks to register HTTP handlers that redirect old
// godoc paths to their new equivalents and assist in accessing the issue
// tracker, wiki, code review system, etc.
package redirect

import (
	"net/http"
	"regexp"
)

// Register registers HTTP handlers that redirect old godoc paths to their new
// equivalents and assist in accessing the issue tracker, wiki, code review
// system, etc. If mux is nil it uses http.DefaultServeMux.
func Register(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}
	handlePathRedirects(mux, pkgRedirects, "/pkg/")
	handlePathRedirects(mux, cmdRedirects, "/cmd/")
	for prefix, redirect := range prefixHelpers {
		p := "/" + prefix + "/"
		mux.Handle(p, PrefixHandler(p, redirect))
	}
	for path, redirect := range redirects {
		mux.Handle(path, Handler(redirect))
	}
}

func handlePathRedirects(mux *http.ServeMux, redirects map[string]string, prefix string) {
	for source, target := range redirects {
		h := Handler(prefix + target + "/")
		p := prefix + source
		mux.Handle(p, h)
		mux.Handle(p+"/", h)
	}
}

// Packages that were renamed between r60 and go1.
var pkgRedirects = map[string]string{
	"asn1":              "encoding/asn1",
	"big":               "math/big",
	"cmath":             "math/cmplx",
	"csv":               "encoding/csv",
	"exec":              "os/exec",
	"exp/template/html": "html/template",
	"gob":               "encoding/gob",
	"http":              "net/http",
	"http/cgi":          "net/http/cgi",
	"http/fcgi":         "net/http/fcgi",
	"http/httptest":     "net/http/httptest",
	"http/pprof":        "net/http/pprof",
	"json":              "encoding/json",
	"mail":              "net/mail",
	"rand":              "math/rand",
	"rpc":               "net/rpc",
	"rpc/jsonrpc":       "net/rpc/jsonrpc",
	"scanner":           "text/scanner",
	"smtp":              "net/smtp",
	"tabwriter":         "text/tabwriter",
	"template":          "text/template",
	"template/parse":    "text/template/parse",
	"url":               "net/url",
	"utf16":             "unicode/utf16",
	"utf8":              "unicode/utf8",
	"xml":               "encoding/xml",
}

// Commands that were renamed between r60 and go1.
var cmdRedirects = map[string]string{
	"gofix":     "fix",
	"goinstall": "go",
	"gopack":    "pack",
	"gotest":    "go",
	"govet":     "vet",
	"goyacc":    "yacc",
}

var redirects = map[string]string{
	"/blog":       "/blog/",
	"/build":      "http://build.golang.org",
	"/change":     "https://code.google.com/p/go/source/list",
	"/cl":         "https://gocodereview.appspot.com/",
	"/cmd/godoc/": "http://godoc.org/code.google.com/p/go.tools/cmd/godoc/",
	"/cmd/vet/":   "http://godoc.org/code.google.com/p/go.tools/cmd/vet/",
	"/issue":      "https://code.google.com/p/go/issues",
	"/issue/new":  "https://code.google.com/p/go/issues/entry",
	"/issues":     "https://code.google.com/p/go/issues",
	"/play":       "http://play.golang.org",
	"/ref":        "/doc/#references",
	"/ref/":       "/doc/#references",
	"/ref/mem":    "/doc/mem",
	"/ref/spec":   "/doc/spec",
	"/talks":      "http://talks.golang.org",
	"/tour":       "http://tour.golang.org",
	"/wiki":       "https://code.google.com/p/go-wiki/w/list",

	"/doc/articles/c_go_cgo.html":                    "/blog/c-go-cgo",
	"/doc/articles/concurrency_patterns.html":        "/blog/go-concurrency-patterns-timing-out-and",
	"/doc/articles/defer_panic_recover.html":         "/blog/defer-panic-and-recover",
	"/doc/articles/error_handling.html":              "/blog/error-handling-and-go",
	"/doc/articles/gobs_of_data.html":                "/blog/gobs-of-data",
	"/doc/articles/godoc_documenting_go_code.html":   "/blog/godoc-documenting-go-code",
	"/doc/articles/gos_declaration_syntax.html":      "/blog/gos-declaration-syntax",
	"/doc/articles/image_draw.html":                  "/blog/go-imagedraw-package",
	"/doc/articles/image_package.html":               "/blog/go-image-package",
	"/doc/articles/json_and_go.html":                 "/blog/json-and-go",
	"/doc/articles/json_rpc_tale_of_interfaces.html": "/blog/json-rpc-tale-of-interfaces",
	"/doc/articles/laws_of_reflection.html":          "/blog/laws-of-reflection",
	"/doc/articles/race_detector.html":               "/blog/race-detector",
	"/doc/articles/slices_usage_and_internals.html":  "/blog/go-slices-usage-and-internals",
	"/doc/go_for_cpp_programmers.html":               "https://code.google.com/p/go-wiki/wiki/GoForCPPProgrammers",
	"/doc/go_tutorial.html":                          "http://tour.golang.org/",
}

var prefixHelpers = map[string]string{
	"change": "https://code.google.com/p/go/source/detail?r=",
	"cl":     "https://codereview.appspot.com/",
	"issue":  "https://code.google.com/p/go/issues/detail?id=",
	"play":   "http://play.golang.org/",
	"talks":  "http://talks.golang.org/",
	"wiki":   "https://code.google.com/p/go-wiki/wiki/",
}

func Handler(target string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, target, http.StatusMovedPermanently)
	})
}

var validId = regexp.MustCompile(`^[A-Za-z0-9-]*$`)

func PrefixHandler(prefix, baseURL string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if p := r.URL.Path; p == prefix {
			// redirect /prefix/ to /prefix
			http.Redirect(w, r, p[:len(p)-1], http.StatusFound)
			return
		}
		id := r.URL.Path[len(prefix):]
		if !validId.MatchString(id) {
			http.Error(w, "Not found", http.StatusNotFound)
			return
		}
		target := baseURL + id
		http.Redirect(w, r, target, http.StatusFound)
	})
}
<commit_msg>go.tools/godoc: remove Go 1.2 redirects<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package redirect provides hooks to register HTTP handlers that redirect old
// godoc paths to their new equivalents and assist in accessing the issue
// tracker, wiki, code review system, etc.
package redirect

import (
	"net/http"
	"regexp"
)

// Register registers HTTP handlers that redirect old godoc paths to their new
// equivalents and assist in accessing the issue tracker, wiki, code review
// system, etc. If mux is nil it uses http.DefaultServeMux.
func Register(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}
	handlePathRedirects(mux, pkgRedirects, "/pkg/")
	handlePathRedirects(mux, cmdRedirects, "/cmd/")
	for prefix, redirect := range prefixHelpers {
		p := "/" + prefix + "/"
		mux.Handle(p, PrefixHandler(p, redirect))
	}
	for path, redirect := range redirects {
		mux.Handle(path, Handler(redirect))
	}
}

func handlePathRedirects(mux *http.ServeMux, redirects map[string]string, prefix string) {
	for source, target := range redirects {
		h := Handler(prefix + target + "/")
		p := prefix + source
		mux.Handle(p, h)
		mux.Handle(p+"/", h)
	}
}

// Packages that were renamed between r60 and go1.
var pkgRedirects = map[string]string{
	"asn1":              "encoding/asn1",
	"big":               "math/big",
	"cmath":             "math/cmplx",
	"csv":               "encoding/csv",
	"exec":              "os/exec",
	"exp/template/html": "html/template",
	"gob":               "encoding/gob",
	"http":              "net/http",
	"http/cgi":          "net/http/cgi",
	"http/fcgi":         "net/http/fcgi",
	"http/httptest":     "net/http/httptest",
	"http/pprof":        "net/http/pprof",
	"json":              "encoding/json",
	"mail":              "net/mail",
	"rand":              "math/rand",
	"rpc":               "net/rpc",
	"rpc/jsonrpc":       "net/rpc/jsonrpc",
	"scanner":           "text/scanner",
	"smtp":              "net/smtp",
	"tabwriter":         "text/tabwriter",
	"template":          "text/template",
	"template/parse":    "text/template/parse",
	"url":               "net/url",
	"utf16":             "unicode/utf16",
	"utf8":              "unicode/utf8",
	"xml":               "encoding/xml",
}

// Commands that were renamed between r60 and go1.
var cmdRedirects = map[string]string{
	"gofix":     "fix",
	"goinstall": "go",
	"gopack":    "pack",
	"gotest":    "go",
	"govet":     "vet",
	"goyacc":    "yacc",
}

var redirects = map[string]string{
	"/blog":       "/blog/",
	"/build":      "http://build.golang.org",
	"/change":     "https://code.google.com/p/go/source/list",
	"/cl":         "https://gocodereview.appspot.com/",
	"/cmd/godoc/": "http://godoc.org/code.google.com/p/go.tools/cmd/godoc/",
	"/cmd/vet/":   "http://godoc.org/code.google.com/p/go.tools/cmd/vet/",
	"/issue":      "https://code.google.com/p/go/issues",
	"/issue/new":  "https://code.google.com/p/go/issues/entry",
	"/issues":     "https://code.google.com/p/go/issues",
	"/play":       "http://play.golang.org",
	// THESE DO NOT WORK FOR GO 1.1 SITES SUCH AS golang.org.
	// DO NOT RE-ENABLE.
	// "/ref":        "/doc/#references",
	// "/ref/":       "/doc/#references",
	// "/ref/mem":    "/doc/mem",
	// "/ref/spec":   "/doc/spec",
	"/talks": "http://talks.golang.org",
	"/tour":  "http://tour.golang.org",
	"/wiki":  "https://code.google.com/p/go-wiki/w/list",

	"/doc/articles/c_go_cgo.html":                    "/blog/c-go-cgo",
	"/doc/articles/concurrency_patterns.html":        "/blog/go-concurrency-patterns-timing-out-and",
	"/doc/articles/defer_panic_recover.html":         "/blog/defer-panic-and-recover",
	"/doc/articles/error_handling.html":              "/blog/error-handling-and-go",
	"/doc/articles/gobs_of_data.html":                "/blog/gobs-of-data",
	"/doc/articles/godoc_documenting_go_code.html":   "/blog/godoc-documenting-go-code",
	"/doc/articles/gos_declaration_syntax.html":      "/blog/gos-declaration-syntax",
	"/doc/articles/image_draw.html":                  "/blog/go-imagedraw-package",
	"/doc/articles/image_package.html":               "/blog/go-image-package",
	"/doc/articles/json_and_go.html":                 "/blog/json-and-go",
	"/doc/articles/json_rpc_tale_of_interfaces.html": "/blog/json-rpc-tale-of-interfaces",
	"/doc/articles/laws_of_reflection.html":          "/blog/laws-of-reflection",
	"/doc/articles/race_detector.html":               "/blog/race-detector",
	"/doc/articles/slices_usage_and_internals.html":  "/blog/go-slices-usage-and-internals",
	"/doc/go_for_cpp_programmers.html":               "https://code.google.com/p/go-wiki/wiki/GoForCPPProgrammers",
	"/doc/go_tutorial.html":                          "http://tour.golang.org/",
}

var prefixHelpers = map[string]string{
	"change": "https://code.google.com/p/go/source/detail?r=",
	"cl":     "https://codereview.appspot.com/",
	"issue":  "https://code.google.com/p/go/issues/detail?id=",
	"play":   "http://play.golang.org/",
	"talks":  "http://talks.golang.org/",
	"wiki":   "https://code.google.com/p/go-wiki/wiki/",
}

func Handler(target string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, target, http.StatusMovedPermanently)
	})
}

var validId = regexp.MustCompile(`^[A-Za-z0-9-]*$`)

func PrefixHandler(prefix, baseURL string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if p := r.URL.Path; p == prefix {
			// redirect /prefix/ to /prefix
			http.Redirect(w, r, p[:len(p)-1], http.StatusFound)
			return
		}
		id := r.URL.Path[len(prefix):]
		if !validId.MatchString(id) {
			http.Error(w, "Not found", http.StatusNotFound)
			return
		}
		target := baseURL + id
		http.Redirect(w, r, target, http.StatusFound)
	})
}
<|endoftext|>


<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()
}

// 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
}
<commit_msg>update the description<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()
}

// OnHTMLDetach deregister a function. Function will not be execute after detached
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 main

// DNS based service discovery
// ---------------------------------------------------------------------------
//
// Marathon: "/path/to/application"
// DNS-query: application.to.path.$basedomain (A, AAAA, TXT, SRV)
// DNS-reply:
//   A 		=> list of IPv4 addresses
//   AAAA => list of IPv6 addresses
//   SRV 	=> ip:port array per task
//   TXT 	=> app labels

// TODO: DNS forwarding
// TODO: DNS proxy cache (for speeding up)

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/christianparpart/go-marathon/marathon"
	"github.com/miekg/dns"
)

type DnsManager struct {
	Verbose     bool
	ServiceAddr net.IP
	ServicePort uint
	BaseName    string
	DnsTTL      time.Duration
	PushSRV     bool
	server      *dns.Server
	db          map[string]*appEntry
	dbMutex     sync.Mutex
}

type appEntry struct {
	ipAddresses []net.IP
	app         *marathon.App
}

func (manager *DnsManager) Setup() error {
	dns.HandleFunc(manager.BaseName, manager.dnsHandler)

	manager.server = &dns.Server{
		Addr:       fmt.Sprintf("%v:%v", manager.ServiceAddr, manager.ServicePort),
		Net:        "udp",
		TsigSecret: nil,
	}

	go func() {
		err := manager.server.ListenAndServe()
		if err != nil {
			log.Fatal(err)
		}
	}()
	return nil
}

func (manager *DnsManager) Shutdown() {
	manager.server.Shutdown()
}

func (manager *DnsManager) Log(msg string) {
	if manager.Verbose {
		log.Printf("[dns]: %v\n", msg)
	}
}

func (manager *DnsManager) Apply(apps []*marathon.App, force bool) error {
	manager.dbMutex.Lock()
	manager.db = make(map[string]*appEntry)
	manager.dbMutex.Unlock()

	for _, app := range apps {
		if err := manager.updateApp(app); err != nil {
			return err
		}
	}

	return nil
}

func (manager *DnsManager) Update(app *marathon.App, taskID string) error {
	return manager.updateApp(app)
}

func (manager *DnsManager) updateApp(app *marathon.App) error {
	var ipAddresses []net.IP

	for _, task := range app.Tasks {
		ip, err := net.ResolveIPAddr("ip", task.Host)
		if err != nil {
			return err
		}
		ipAddresses = append(ipAddresses, ip.IP)
	}

	var reversed = manager.makeDnsNameFromAppId(app.Id)
	var entry = &appEntry{
		ipAddresses: ipAddresses,
		app:         app,
	}

	manager.dbMutex.Lock()
	manager.db[reversed] = entry
	manager.dbMutex.Unlock()

	return nil
}

func (manager *DnsManager) makeDnsNameFromAppId(appID string) string {
	var parts = strings.Split(appID, "/")[1:]
	var reversedParts []string
	for i := range parts {
		reversedParts = append(reversedParts, parts[len(parts)-i-1])
	}
	var reversed = strings.Join(reversedParts, ".")

	return reversed
}

func (manager *DnsManager) Remove(appID string, taskID string, app *marathon.App) error {
	if app == nil {
		manager.dbMutex.Lock()
		delete(manager.db, manager.makeDnsNameFromAppId(appID))
		manager.dbMutex.Unlock()
		return nil
	}
	return manager.updateApp(app)
}

func (manager *DnsManager) dnsHandler(w dns.ResponseWriter, req *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(req)

	name := req.Question[0].Name
	name = strings.TrimSuffix(name, "."+manager.BaseName)

	manager.dbMutex.Lock()
	entry, ok := manager.db[name]
	manager.dbMutex.Unlock()

	if ok {
		switch req.Question[0].Qtype {
		case dns.TypeSRV:
			m.Answer = manager.makeAllSRV(entry)
		case dns.TypeA:
			m.Answer = manager.makeAllA(entry)
			if manager.PushSRV {
				m.Extra = manager.makeAllSRV(entry)
			}
		}
	}

	w.WriteMsg(m)
}

func (manager *DnsManager) makeAllA(entry *appEntry) []dns.RR {
	var result []dns.RR

	for _, ip := range entry.ipAddresses {
		rr := &dns.A{
			Hdr: dns.RR_Header{
				Ttl:    uint32(manager.DnsTTL.Seconds()),
				Name:   manager.BaseName,
				Class:  dns.ClassINET,
				Rrtype: dns.TypeA,
			},
			A: ip.To4(),
		}
		result = append(result, rr)
	}

	return result
}

func (manager *DnsManager) makeAllSRV(entry *appEntry) []dns.RR {
	var result []dns.RR

	for _, task := range entry.app.Tasks {
		for _, port := range task.Ports {
			rr := &dns.SRV{
				Hdr: dns.RR_Header{
					Ttl:    uint32(manager.DnsTTL.Seconds()),
					Name:   manager.BaseName,
					Class:  dns.ClassINET,
					Rrtype: dns.TypeSRV,
				},
				Port:     uint16(port),
				Target:   task.Host + ".",
				Weight:   1,
				Priority: 1,
			}
			result = append(result, rr)
		}
	}

	return result
}
<commit_msg>[dns] support TCP along-side to UDP<commit_after>package main

// DNS based service discovery
// ---------------------------------------------------------------------------
//
// Marathon: "/path/to/application"
// DNS-query: application.to.path.$basedomain (A, AAAA, TXT, SRV)
// DNS-reply:
//   A 		=> list of IPv4 addresses
//   AAAA => list of IPv6 addresses
//   SRV 	=> ip:port array per task
//   TXT 	=> app labels

// TODO: DNS forwarding
// TODO: DNS proxy cache (for speeding up)

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/christianparpart/go-marathon/marathon"
	"github.com/miekg/dns"
)

type DnsManager struct {
	Verbose     bool
	ServiceAddr net.IP
	ServicePort uint
	BaseName    string
	DnsTTL      time.Duration
	PushSRV     bool
	udpServer   *dns.Server
	tcpServer   *dns.Server
	db          map[string]*appEntry
	dbMutex     sync.Mutex
}

type appEntry struct {
	ipAddresses []net.IP
	app         *marathon.App
}

func (manager *DnsManager) Setup() error {
	dns.HandleFunc(manager.BaseName, manager.dnsHandler)

	go func() {
		manager.udpServer = &dns.Server{
			Addr:       fmt.Sprintf("%v:%v", manager.ServiceAddr, manager.ServicePort),
			Net:        "udp",
			TsigSecret: nil,
		}
		err := manager.udpServer.ListenAndServe()
		if err != nil {
			log.Fatal(err)
		}
	}()

	go func() {
		manager.tcpServer = &dns.Server{
			Addr:       fmt.Sprintf("%v:%v", manager.ServiceAddr, manager.ServicePort),
			Net:        "tcp",
			TsigSecret: nil,
		}
		err := manager.tcpServer.ListenAndServe()
		if err != nil {
			log.Fatal(err)
		}
	}()

	return nil
}

func (manager *DnsManager) Shutdown() {
	manager.udpServer.Shutdown()
	manager.tcpServer.Shutdown()
}

func (manager *DnsManager) Log(msg string) {
	if manager.Verbose {
		log.Printf("[dns]: %v\n", msg)
	}
}

func (manager *DnsManager) Apply(apps []*marathon.App, force bool) error {
	manager.dbMutex.Lock()
	manager.db = make(map[string]*appEntry)
	manager.dbMutex.Unlock()

	for _, app := range apps {
		if err := manager.updateApp(app); err != nil {
			return err
		}
	}

	return nil
}

func (manager *DnsManager) Update(app *marathon.App, taskID string) error {
	return manager.updateApp(app)
}

func (manager *DnsManager) updateApp(app *marathon.App) error {
	var ipAddresses []net.IP

	for _, task := range app.Tasks {
		ip, err := net.ResolveIPAddr("ip", task.Host)
		if err != nil {
			return err
		}
		ipAddresses = append(ipAddresses, ip.IP)
	}

	var reversed = manager.makeDnsNameFromAppId(app.Id)
	var entry = &appEntry{
		ipAddresses: ipAddresses,
		app:         app,
	}

	manager.dbMutex.Lock()
	manager.db[reversed] = entry
	manager.dbMutex.Unlock()

	return nil
}

func (manager *DnsManager) makeDnsNameFromAppId(appID string) string {
	var parts = strings.Split(appID, "/")[1:]
	var reversedParts []string
	for i := range parts {
		reversedParts = append(reversedParts, parts[len(parts)-i-1])
	}
	var reversed = strings.Join(reversedParts, ".")

	return reversed
}

func (manager *DnsManager) Remove(appID string, taskID string, app *marathon.App) error {
	if app == nil {
		manager.dbMutex.Lock()
		delete(manager.db, manager.makeDnsNameFromAppId(appID))
		manager.dbMutex.Unlock()
		return nil
	}
	return manager.updateApp(app)
}

func (manager *DnsManager) dnsHandler(w dns.ResponseWriter, req *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(req)

	name := req.Question[0].Name
	name = strings.TrimSuffix(name, "."+manager.BaseName)

	manager.dbMutex.Lock()
	entry, ok := manager.db[name]
	manager.dbMutex.Unlock()

	if ok {
		switch req.Question[0].Qtype {
		case dns.TypeSRV:
			m.Answer = manager.makeAllSRV(entry)
		case dns.TypeA:
			m.Answer = manager.makeAllA(entry)
			if manager.PushSRV {
				m.Extra = manager.makeAllSRV(entry)
			}
		}
	}

	w.WriteMsg(m)
}

func (manager *DnsManager) makeAllA(entry *appEntry) []dns.RR {
	var result []dns.RR

	for _, ip := range entry.ipAddresses {
		rr := &dns.A{
			Hdr: dns.RR_Header{
				Ttl:    uint32(manager.DnsTTL.Seconds()),
				Name:   manager.BaseName,
				Class:  dns.ClassINET,
				Rrtype: dns.TypeA,
			},
			A: ip.To4(),
		}
		result = append(result, rr)
	}

	return result
}

func (manager *DnsManager) makeAllSRV(entry *appEntry) []dns.RR {
	var result []dns.RR

	for _, task := range entry.app.Tasks {
		for _, port := range task.Ports {
			rr := &dns.SRV{
				Hdr: dns.RR_Header{
					Ttl:    uint32(manager.DnsTTL.Seconds()),
					Name:   manager.BaseName,
					Class:  dns.ClassINET,
					Rrtype: dns.TypeSRV,
				},
				Port:     uint16(port),
				Target:   task.Host + ".",
				Weight:   1,
				Priority: 1,
			}
			result = append(result, rr)
		}
	}

	return result
}
<|endoftext|>


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

package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/dnstap/golang-dnstap"
)

var (
	flagReadTcp    = flag.String("l", "", "read dnstap payloads from tcp/ip")
	flagWriteTcp   = flag.String("T", "", "write dnstap payloads to tcp/ip address")
	flagTimeout    = flag.Duration("t", 0, "I/O timeout for tcp/ip and unix domain sockets")
	flagReadFile   = flag.String("r", "", "read dnstap payloads from file")
	flagReadSock   = flag.String("u", "", "read dnstap payloads from unix socket")
	flagWriteUnix  = flag.String("U", "", "write dnstap payloads to unix socket")
	flagWriteFile  = flag.String("w", "-", "write output to file")
	flagAppendFile = flag.Bool("a", false, "append to the given file, do not overwrite. valid only when outputting a text or YAML file.")
	flagQuietText  = flag.Bool("q", false, "use quiet text output")
	flagYamlText   = flag.Bool("y", false, "use verbose YAML output")
	flagJsonText   = flag.Bool("j", false, "use verbose JSON output")
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTION]...\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, `
Quiet text output format mnemonics:
    AQ: AUTH_QUERY
    AR: AUTH_RESPONSE
    RQ: RESOLVER_QUERY
    RR: RESOLVER_RESPONSE
    CQ: CLIENT_QUERY
    CR: CLIENT_RESPONSE
    FQ: FORWARDER_QUERY
    FR: FORWARDER_RESPONSE
    SQ: STUB_QUERY
    SR: STUB_RESPONSE
    TQ: TOOL_QUERY
    TR: TOOL_RESPONSE
`)
}

func outputOpener(fname string, text, yaml, json, doAppend bool) func() dnstap.Output {
	return func() dnstap.Output {
		var o dnstap.Output
		var err error
		if text {
			o, err = dnstap.NewTextOutputFromFilename(fname, dnstap.TextFormat, doAppend)
		} else if yaml {
			o, err = dnstap.NewTextOutputFromFilename(fname, dnstap.YamlFormat, doAppend)
		} else if json {
			o, err = dnstap.NewTextOutputFromFilename(fname, dnstap.JsonFormat, doAppend)
		} else {
			o, err = dnstap.NewFrameStreamOutputFromFilename(fname)
		}

		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Failed to open output file: %s\n", err)
			os.Exit(1)
		}

		go o.RunOutputLoop()
		return o
	}
}

func outputLoop(opener func() dnstap.Output, data <-chan []byte, done chan<- struct{}) {
	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, os.Interrupt, syscall.SIGHUP)
	o := opener()
	defer func() {
		o.Close()
		close(done)
		os.Exit(0)
	}()
	for {
		select {
		case b, ok := <-data:
			if !ok {
				return
			}
			o.GetOutputChannel() <- b
		case sig := <-sigch:
			if sig == syscall.SIGHUP {
				o.Close()
				o = opener()
				continue
			}
			return
		}
	}
}

func main() {
	var err error
	var i dnstap.Input

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.SetFlags(0)
	flag.Usage = usage

	// Handle command-line arguments.
	flag.Parse()

	if *flagReadFile == "" && *flagReadSock == "" && *flagReadTcp == "" {
		fmt.Fprintf(os.Stderr, "dnstap: Error: no inputs specified.\n")
		os.Exit(1)
	}

	if *flagWriteFile == "-" {
		if *flagQuietText == false && *flagYamlText == false && *flagJsonText == false {
			*flagQuietText = true
		}
	}

	if *flagAppendFile == true {
		if *flagWriteFile == "-" || *flagWriteFile == "" {
			fmt.Fprintf(os.Stderr, "dnstap: Error: -a must specify the file output path.\n")
			os.Exit(1)
		}
	}
	if *flagYamlText == true && *flagJsonText == true {
		fmt.Fprintf(os.Stderr, "dnstap: Error: specify exactly one of -y or -j.\n")
		os.Exit(1)
	}

	if *flagReadFile != "" && *flagReadSock != "" && *flagReadTcp != "" {
		fmt.Fprintf(os.Stderr, "dnstap: Error: specify exactly one of -r, -u or -l.\n")
		os.Exit(1)
	}

	var output chan []byte
	outDone := make(chan struct{})

	// Start the output loop.
	if *flagWriteTcp == "" && *flagWriteUnix == "" {
		output = make(chan []byte, 1)
		opener := outputOpener(*flagWriteFile, *flagQuietText, *flagYamlText,
			*flagJsonText, *flagAppendFile)
		go outputLoop(opener, output, outDone)
	} else {
		var addr net.Addr
		var sockOutput *dnstap.FrameStreamSockOutput
		if *flagWriteTcp != "" {
			addr, err = net.ResolveTCPAddr("tcp", *flagWriteTcp)
		} else {
			addr, err = net.ResolveUnixAddr("unix", *flagWriteUnix)
		}
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Error: invalid address: %v", err)
			os.Exit(1)
		}
		sockOutput, err = dnstap.NewFrameStreamSockOutput(addr,
			&dnstap.SockOutputConfig{Timeout: *flagTimeout})
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Error: failed to open socket output: %v", err)
			os.Exit(1)
		}
		output = sockOutput.GetOutputChannel()
		go sockOutput.RunOutputLoop()
		close(outDone)
	}

	// Open the input and start the input loop.
	if *flagReadFile != "" {
		i, err = dnstap.NewFrameStreamInputFromFilename(*flagReadFile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Failed to open input file: %s\n", err)
			os.Exit(1)
		}
		fmt.Fprintf(os.Stderr, "dnstap: opened input file %s\n", *flagReadFile)
	} else if *flagReadSock != "" {
		var si *dnstap.FrameStreamSockInput
		si, err = dnstap.NewFrameStreamSockInputFromPath(*flagReadSock)
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Failed to open input socket: %s\n", err)
			os.Exit(1)
		}
		si.SetTimeout(*flagTimeout)
		i = si
		fmt.Fprintf(os.Stderr, "dnstap: opened input socket %s\n", *flagReadSock)
	} else if *flagReadTcp != "" {
		l, err := net.Listen("tcp", *flagReadTcp)
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Failed to listen: %s\n", err)
			os.Exit(1)
		}
		si := dnstap.NewFrameStreamSockInput(l)
		si.SetTimeout(*flagTimeout)
		i = si
	}
	i.ReadInto(output)

	// Wait for input loop to finish.
	i.Wait()
	close(output)

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

package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/dnstap/golang-dnstap"
)

var (
	flagReadTcp    = flag.String("l", "", "read dnstap payloads from tcp/ip")
	flagWriteTcp   = flag.String("T", "", "write dnstap payloads to tcp/ip address")
	flagTimeout    = flag.Duration("t", 0, "I/O timeout for tcp/ip and unix domain sockets")
	flagReadFile   = flag.String("r", "", "read dnstap payloads from file")
	flagReadSock   = flag.String("u", "", "read dnstap payloads from unix socket")
	flagWriteUnix  = flag.String("U", "", "write dnstap payloads to unix socket")
	flagWriteFile  = flag.String("w", "-", "write output to file")
	flagAppendFile = flag.Bool("a", false, "append to the given file, do not overwrite. valid only when outputting a text or YAML file.")
	flagQuietText  = flag.Bool("q", false, "use quiet text output")
	flagYamlText   = flag.Bool("y", false, "use verbose YAML output")
	flagJsonText   = flag.Bool("j", false, "use verbose JSON output")
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTION]...\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, `
Quiet text output format mnemonics:
    AQ: AUTH_QUERY
    AR: AUTH_RESPONSE
    RQ: RESOLVER_QUERY
    RR: RESOLVER_RESPONSE
    CQ: CLIENT_QUERY
    CR: CLIENT_RESPONSE
    FQ: FORWARDER_QUERY
    FR: FORWARDER_RESPONSE
    SQ: STUB_QUERY
    SR: STUB_RESPONSE
    TQ: TOOL_QUERY
    TR: TOOL_RESPONSE
`)
}

func outputOpener(fname string, text, yaml, json, doAppend bool) func() dnstap.Output {
	return func() dnstap.Output {
		var o dnstap.Output
		var err error
		if text {
			o, err = dnstap.NewTextOutputFromFilename(fname, dnstap.TextFormat, doAppend)
		} else if yaml {
			o, err = dnstap.NewTextOutputFromFilename(fname, dnstap.YamlFormat, doAppend)
		} else if json {
			o, err = dnstap.NewTextOutputFromFilename(fname, dnstap.JsonFormat, doAppend)
		} else {
			o, err = dnstap.NewFrameStreamOutputFromFilename(fname)
		}

		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Failed to open output file: %s\n", err)
			os.Exit(1)
		}

		go o.RunOutputLoop()
		return o
	}
}

func outputLoop(opener func() dnstap.Output, data <-chan []byte, done chan<- struct{}) {
	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, os.Interrupt, syscall.SIGHUP)
	o := opener()
	defer func() {
		o.Close()
		close(done)
		os.Exit(0)
	}()
	for {
		select {
		case b, ok := <-data:
			if !ok {
				return
			}
			o.GetOutputChannel() <- b
		case sig := <-sigch:
			if sig == syscall.SIGHUP {
				o.Close()
				o = opener()
				continue
			}
			return
		}
	}
}

func main() {
	var err error
	var i dnstap.Input

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.SetFlags(0)
	flag.Usage = usage

	// Handle command-line arguments.
	flag.Parse()

	if *flagReadFile == "" && *flagReadSock == "" && *flagReadTcp == "" {
		fmt.Fprintf(os.Stderr, "dnstap: Error: no inputs specified.\n")
		os.Exit(1)
	}

	if *flagWriteFile == "-" || *flagWriteFile == "" {
		if *flagQuietText == false && *flagYamlText == false && *flagJsonText == false {
			*flagQuietText = true
		}
	}

	if *flagAppendFile == true {
		if *flagWriteFile == "-" || *flagWriteFile == "" {
			fmt.Fprintf(os.Stderr, "dnstap: Error: -a must specify the file output path.\n")
			os.Exit(1)
		}
	}
	if *flagYamlText == true && *flagJsonText == true {
		fmt.Fprintf(os.Stderr, "dnstap: Error: specify exactly one of -y or -j.\n")
		os.Exit(1)
	}

	if *flagReadFile != "" && *flagReadSock != "" && *flagReadTcp != "" {
		fmt.Fprintf(os.Stderr, "dnstap: Error: specify exactly one of -r, -u or -l.\n")
		os.Exit(1)
	}

	var output chan []byte
	outDone := make(chan struct{})

	// Start the output loop.
	if *flagWriteTcp == "" && *flagWriteUnix == "" {
		output = make(chan []byte, 1)
		opener := outputOpener(*flagWriteFile, *flagQuietText, *flagYamlText,
			*flagJsonText, *flagAppendFile)
		go outputLoop(opener, output, outDone)
	} else {
		var addr net.Addr
		var sockOutput *dnstap.FrameStreamSockOutput
		if *flagWriteTcp != "" {
			addr, err = net.ResolveTCPAddr("tcp", *flagWriteTcp)
		} else {
			addr, err = net.ResolveUnixAddr("unix", *flagWriteUnix)
		}
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Error: invalid address: %v", err)
			os.Exit(1)
		}
		sockOutput, err = dnstap.NewFrameStreamSockOutput(addr,
			&dnstap.SockOutputConfig{Timeout: *flagTimeout})
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Error: failed to open socket output: %v", err)
			os.Exit(1)
		}
		output = sockOutput.GetOutputChannel()
		go sockOutput.RunOutputLoop()
		close(outDone)
	}

	// Open the input and start the input loop.
	if *flagReadFile != "" {
		i, err = dnstap.NewFrameStreamInputFromFilename(*flagReadFile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Failed to open input file: %s\n", err)
			os.Exit(1)
		}
		fmt.Fprintf(os.Stderr, "dnstap: opened input file %s\n", *flagReadFile)
	} else if *flagReadSock != "" {
		var si *dnstap.FrameStreamSockInput
		si, err = dnstap.NewFrameStreamSockInputFromPath(*flagReadSock)
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Failed to open input socket: %s\n", err)
			os.Exit(1)
		}
		si.SetTimeout(*flagTimeout)
		i = si
		fmt.Fprintf(os.Stderr, "dnstap: opened input socket %s\n", *flagReadSock)
	} else if *flagReadTcp != "" {
		l, err := net.Listen("tcp", *flagReadTcp)
		if err != nil {
			fmt.Fprintf(os.Stderr, "dnstap: Failed to listen: %s\n", err)
			os.Exit(1)
		}
		si := dnstap.NewFrameStreamSockInput(l)
		si.SetTimeout(*flagTimeout)
		i = si
	}
	i.ReadInto(output)

	// Wait for input loop to finish.
	i.Wait()
	close(output)

	<-outDone
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"
)

func main() {

	// subscribe to SIGINT signals
	stopChan := make(chan os.Signal)
	signal.Notify(stopChan, os.Interrupt)

	mux := http.NewServeMux()

	mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(5 * time.Second)
		w.Write([]byte("Finished!\n"))
	}))

	srv := &http.Server{Addr: ":8081", Handler: mux}

	go func() {
		// service connections
		if err := srv.ListenAndServe(); err != nil {
			log.Printf("listen: %s\n", err)
		}
	}()

	<-stopChan // wait for SIGINT
	log.Println("Shutting down server...")

	srv.Shutdown(context.Background())

	log.Println("gracefully stopped stopped")

}
<commit_msg>Add a timeout to the graceful shutdown example, since that's probably best practice<commit_after>package main

import (
	"context"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"
)

func main() {

	// subscribe to SIGINT signals
	stopChan := make(chan os.Signal)
	signal.Notify(stopChan, os.Interrupt)

	mux := http.NewServeMux()

	mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(5 * time.Second)
		w.Write([]byte("Finished!\n"))
	}))

	srv := &http.Server{Addr: ":8081", Handler: mux}

	go func() {
		// service connections
		if err := srv.ListenAndServe(); err != nil {
			log.Printf("listen: %s\n", err)
		}
	}()

	<-stopChan // wait for SIGINT
	log.Println("Shutting down server...")

	// shut down gracefully, but wait no longer than 5 seconds before halting
	ctx, _ := context.WithTimeout(context.Background(), 5*time.Second)
	srv.Shutdown(ctx)

	log.Println("gracefully stopped stopped")

}
<|endoftext|>


<commit_before>package editor

import (
	"strings"
)

// basic represents a basic editor.
// Valid positions are in range [0, len(e.buf)].
type basic struct {
	pos int
	buf []rune
}

// move moves the position to the given position.
// Given a invalid position, move sets the position at the end of the buffer.
func (e *basic) move(to int) {
	switch {
	case to >= len(e.buf):
		e.pos = len(e.buf)
	case to <= 0:
		e.pos = 0
	default:
		e.pos = to
	}
}

// insert inserts s into the buffer at the given position.
// Given a invalid position, insert considers the position to be at the end of the buffer.
func (e *basic) insert(s []rune, at int) {
	switch {
	case at < 0:
		at = 0
	case at > len(e.buf):
		at = len(e.buf)
	}
	switch at {
	case 0:
		e.buf = append(s, e.buf...)
	case len(e.buf):
		e.buf = append(e.buf, s...)
	default:
		x := append(e.buf[:at:at], s...)
		e.buf = append(x, e.buf[at:]...)
	}
	if at <= e.pos {
		e.pos += len(s)
	}
}

func max(m, n int) int {
	if m > n {
		return m
	}
	return n
}

func min(m, n int) int {
	if m < n {
		return m
	}
	return n
}

func constrain(n, low, high int) int {
	if n < low {
		return low
	}
	if n > high {
		return high
	}
	return n
}

// delete deletes runes from the buffer [from, to].
// Given a invalid position, delete considers the position to be at the end of the buffer.
func (e *basic) delete(from, to int) {
	left := constrain(min(from, to), 0, len(e.buf))
	right := constrain(max(from, to), 0, len(e.buf))
	switch {
	case left == 0:
		e.buf = e.buf[right:]
	case right == len(e.buf):
		e.buf = e.buf[:left]
	default:
		e.buf = append(e.buf[:left], e.buf[right:]...)
	}
	switch {
	case e.pos < left:
	case right < e.pos:
		e.pos = e.pos - (right - left)
	default:
		e.pos = left
	}
}

// slice slices the buffer [from, to].
// Given a invalid position, slice considers the position to be at the end of the buffer.
func (e *basic) slice(from, to int) []rune {
	left := constrain(min(from, to), 0, len(e.buf))
	right := constrain(max(from, to), 0, len(e.buf))
	s := make([]rune, right-left)
	copy(s, e.buf[left:right])
	return s
}

func (e *basic) index(ch rune, start int) int {
	start = constrain(start, 0, len(e.buf))
	for i := start; i < len(e.buf); i++ {
		if e.buf[i] == ch {
			return i
		}
	}
	return -1
}

func (e *basic) lastIndex(ch rune, last int) int {
	last = constrain(last, 0, len(e.buf))
	for i := last - 1; i >= 0; i-- {
		if e.buf[i] == ch {
			return i
		}
	}
	return -1
}

// indexFunc(f, start, true) == indexFunc(func(r) bool { return !f(r) }, start, false)
func (e *basic) indexFunc(f func(rune) bool, start int, truth bool) int {
	start = constrain(start, 0, len(e.buf))
	for i := start; i < len(e.buf); i++ {
		if f(e.buf[i]) == truth {
			return i
		}
	}
	return -1
}

// lastIndexFunc(f, start, true) == lastIndexFunc(func(r) bool { return !f(r) }, start, false)
func (e *basic) lastIndexFunc(f func(rune) bool, last int, truth bool) int {
	last = constrain(last, 0, len(e.buf))
	for i := last - 1; i >= 0; i-- {
		if f(e.buf[i]) == truth {
			return i
		}
	}
	return -1
}

func (e *basic) replace(s []rune, at int) {
	if s == nil {
		return
	}
	// Q. What should we do when `at` < 0 or len(e.buf) < `at`?
	switch {
	case len(e.buf) <= at:
		e.buf = append(e.buf, []rune(strings.Repeat(" ", at-len(e.buf)))...)
		e.buf = append(e.buf, s...)
	case at+len(s) <= 0:
		// no-op
	case at < 0:
		for i := 0; i < at+len(s); i++ {
			e.buf[i] = s[i-at]
		}
	case len(e.buf) < at+len(s):
		v := make([]rune, at+len(s))
		copy(v, e.buf)
		for i := at; i < at+len(s); i++ {
			v[i] = s[i-at]
		}
		e.buf = v
	default:
		for i := at; i < at+len(s); i++ {
			e.buf[i] = s[i-at]
		}
	}
}

func (e *basic) siege(from int, to int, r rune) {
	switch r {
	case '\'', '"', '`', '@', '*', '+', '_', '|', '$':
		e.insert([]rune{r}, to)
		e.insert([]rune{r}, from)
	case '(', ')', 'b':
		e.insert([]rune{')'}, to)
		e.insert([]rune{'('}, from)
	case '{', '}', 'B':
		e.insert([]rune{'}'}, to)
		e.insert([]rune{'{'}, from)
	case '[', ']', 'r':
		e.insert([]rune{']'}, to)
		e.insert([]rune{'['}, from)
	case '<', '>', 'a':
		e.insert([]rune{'>'}, to)
		e.insert([]rune{'<'}, from)
	}
}

func (e *basic) wordForward() {
	switch n := len(e.buf) - e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = len(e.buf)
		return
	}
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.indexFunc(isWhitespace, e.pos+1, false); i > 0 {
			e.pos = i
			return
		}
	case isKeyword(r):
		if i := e.indexFunc(isKeyword, e.pos+1, false); i > 0 {
			if !isWhitespace(e.buf[i]) {
				e.pos = i
				return
			}
			if i := e.indexFunc(isWhitespace, i+1, false); i > 0 {
				e.pos = i
				return
			}
		}
	default:
		if i := e.indexFunc(isSymbol, e.pos+1, false); i > 0 {
			if isKeyword(e.buf[i]) {
				e.pos = i
				return
			}
			if i := e.indexFunc(isWhitespace, i+1, false); i > 0 {
				e.pos = i
				return
			}
		}
	}
	e.pos = len(e.buf)
}

func (e *basic) wordBackward() {
	switch e.pos {
	case 0:
		return
	case 1:
		e.pos = 0
		return
	}

	n := e.pos - 1
	switch r := e.buf[n]; {
	case isWhitespace(r):
		n = e.lastIndexFunc(isWhitespace, n, false)
		if n < 0 {
			e.pos = 0
			return
		}
	}

	switch r := e.buf[n]; {
	case isKeyword(r):
		if i := e.lastIndexFunc(isKeyword, n, false); i >= 0 {
			e.pos = i + 1
			return
		}
	default:
		if i := e.lastIndexFunc(isSymbol, n, false); i >= 0 {
			e.pos = i + 1
			return
		}
	}
	e.pos = 0
}

func (e *basic) wordForwardNonBlank() {
	i := e.indexFunc(isWhitespace, e.pos, true)
	if i < 0 {
		e.pos = len(e.buf)
		return
	}
	i = e.indexFunc(isWhitespace, i+1, false)
	if i < 0 {
		e.pos = len(e.buf)
		return
	}
	e.pos = i
}

func (e *basic) wordBackwardNonBlank() {
	i := e.lastIndexFunc(isWhitespace, e.pos, false)
	if i < 0 {
		e.pos = 0
		return
	}
	i = e.lastIndexFunc(isWhitespace, i, true)
	if i < 0 {
		e.pos = 0
		return
	}
	e.pos = i + 1
}

func (e *basic) wordEnd() {
	switch n := len(e.buf) - e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = len(e.buf)
		return
	}
	e.pos++
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.indexFunc(isWhitespace, e.pos+1, false); i > 0 {
			switch r := e.buf[i]; {
			case isKeyword(r):
				if i := e.indexFunc(isKeyword, i+1, false); i > 0 {
					e.pos = i - 1
					return
				}
			default:
				if i := e.indexFunc(isSymbol, i+1, false); i > 0 {
					e.pos = i - 1
					return
				}
			}
		}
	case isKeyword(r):
		if i := e.indexFunc(isKeyword, e.pos+1, false); i > 0 {
			e.pos = i - 1
			return
		}
	default:
		if i := e.indexFunc(isSymbol, e.pos+1, false); i > 0 {
			e.pos = i - 1
			return
		}
	}
	e.pos = len(e.buf) - 1
}

func (e *basic) wordEndNonBlank() {
	switch n := len(e.buf) - e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = len(e.buf)
		return
	}
	e.pos++
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.indexFunc(isWhitespace, e.pos+1, false); i > 0 {
			if i := e.indexFunc(isWhitespace, i+1, true); i > 0 {
				e.pos = i - 1
				return
			}
		}
	default:
		if i := e.indexFunc(isWhitespace, e.pos+1, true); i > 0 {
			e.pos = i - 1
			return
		}
	}
	e.pos = len(e.buf) - 1
}

func (e *basic) wordEndBackward() {
	switch n := e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = 0
		return
	}
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.lastIndexFunc(isWhitespace, e.pos, false); i > 0 {
			e.pos = i
			return
		}
	case isKeyword(r):
		if i := e.lastIndexFunc(isKeyword, e.pos, false); i > 0 {
			switch {
			case isWhitespace(e.buf[i]):
				if i := e.lastIndexFunc(isWhitespace, i, false); i > 0 {
					e.pos = i
					return
				}
			default:
				e.pos = i
				return
			}
		}
	default:
		if i := e.lastIndexFunc(isSymbol, e.pos, false); i > 0 {
			switch {
			case isWhitespace(e.buf[i]):
				if i := e.lastIndexFunc(isWhitespace, i, false); i > 0 {
					e.pos = i
					return
				}
			default:
				e.pos = i
				return
			}
		}
	}
	e.pos = 0
}

func (e *basic) wordEndBackwardNonBlank() {
	switch n := e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = 0
		return
	}
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.lastIndexFunc(isWhitespace, e.pos, false); i > 0 {
			e.pos = i
			return
		}
	default:
		if i := e.lastIndexFunc(isWhitespace, e.pos, true); i > 0 {
			if i := e.lastIndexFunc(isWhitespace, i, false); i > 0 {
				e.pos = i
				return
			}
		}
	}
	e.pos = 0
}
<commit_msg>Fix typos<commit_after>package editor

import (
	"strings"
)

// basic represents a basic editor.
// Valid positions are in range [0, len(e.buf)].
type basic struct {
	pos int
	buf []rune
}

// move moves the position to the given position.
// Given an invalid position, move sets the position at the end of the buffer.
func (e *basic) move(to int) {
	switch {
	case to >= len(e.buf):
		e.pos = len(e.buf)
	case to <= 0:
		e.pos = 0
	default:
		e.pos = to
	}
}

// insert inserts s into the buffer at the given position.
// Given an invalid position, insert considers the position to be at the end of the buffer.
func (e *basic) insert(s []rune, at int) {
	switch {
	case at < 0:
		at = 0
	case at > len(e.buf):
		at = len(e.buf)
	}
	switch at {
	case 0:
		e.buf = append(s, e.buf...)
	case len(e.buf):
		e.buf = append(e.buf, s...)
	default:
		x := append(e.buf[:at:at], s...)
		e.buf = append(x, e.buf[at:]...)
	}
	if at <= e.pos {
		e.pos += len(s)
	}
}

func max(m, n int) int {
	if m > n {
		return m
	}
	return n
}

func min(m, n int) int {
	if m < n {
		return m
	}
	return n
}

func constrain(n, low, high int) int {
	if n < low {
		return low
	}
	if n > high {
		return high
	}
	return n
}

// delete deletes runes from the buffer [from, to].
// Given an invalid position, delete considers the position to be at the end of the buffer.
func (e *basic) delete(from, to int) {
	left := constrain(min(from, to), 0, len(e.buf))
	right := constrain(max(from, to), 0, len(e.buf))
	switch {
	case left == 0:
		e.buf = e.buf[right:]
	case right == len(e.buf):
		e.buf = e.buf[:left]
	default:
		e.buf = append(e.buf[:left], e.buf[right:]...)
	}
	switch {
	case e.pos < left:
	case right < e.pos:
		e.pos = e.pos - (right - left)
	default:
		e.pos = left
	}
}

// slice slices the buffer [from, to].
// Given an invalid position, slice considers the position to be at the end of the buffer.
func (e *basic) slice(from, to int) []rune {
	left := constrain(min(from, to), 0, len(e.buf))
	right := constrain(max(from, to), 0, len(e.buf))
	s := make([]rune, right-left)
	copy(s, e.buf[left:right])
	return s
}

func (e *basic) index(ch rune, start int) int {
	start = constrain(start, 0, len(e.buf))
	for i := start; i < len(e.buf); i++ {
		if e.buf[i] == ch {
			return i
		}
	}
	return -1
}

func (e *basic) lastIndex(ch rune, last int) int {
	last = constrain(last, 0, len(e.buf))
	for i := last - 1; i >= 0; i-- {
		if e.buf[i] == ch {
			return i
		}
	}
	return -1
}

// indexFunc(f, start, true) == indexFunc(func(r) bool { return !f(r) }, start, false)
func (e *basic) indexFunc(f func(rune) bool, start int, truth bool) int {
	start = constrain(start, 0, len(e.buf))
	for i := start; i < len(e.buf); i++ {
		if f(e.buf[i]) == truth {
			return i
		}
	}
	return -1
}

// lastIndexFunc(f, start, true) == lastIndexFunc(func(r) bool { return !f(r) }, start, false)
func (e *basic) lastIndexFunc(f func(rune) bool, last int, truth bool) int {
	last = constrain(last, 0, len(e.buf))
	for i := last - 1; i >= 0; i-- {
		if f(e.buf[i]) == truth {
			return i
		}
	}
	return -1
}

func (e *basic) replace(s []rune, at int) {
	if s == nil {
		return
	}
	// Q. What should we do when `at` < 0 or len(e.buf) < `at`?
	switch {
	case len(e.buf) <= at:
		e.buf = append(e.buf, []rune(strings.Repeat(" ", at-len(e.buf)))...)
		e.buf = append(e.buf, s...)
	case at+len(s) <= 0:
		// no-op
	case at < 0:
		for i := 0; i < at+len(s); i++ {
			e.buf[i] = s[i-at]
		}
	case len(e.buf) < at+len(s):
		v := make([]rune, at+len(s))
		copy(v, e.buf)
		for i := at; i < at+len(s); i++ {
			v[i] = s[i-at]
		}
		e.buf = v
	default:
		for i := at; i < at+len(s); i++ {
			e.buf[i] = s[i-at]
		}
	}
}

func (e *basic) siege(from int, to int, r rune) {
	switch r {
	case '\'', '"', '`', '@', '*', '+', '_', '|', '$':
		e.insert([]rune{r}, to)
		e.insert([]rune{r}, from)
	case '(', ')', 'b':
		e.insert([]rune{')'}, to)
		e.insert([]rune{'('}, from)
	case '{', '}', 'B':
		e.insert([]rune{'}'}, to)
		e.insert([]rune{'{'}, from)
	case '[', ']', 'r':
		e.insert([]rune{']'}, to)
		e.insert([]rune{'['}, from)
	case '<', '>', 'a':
		e.insert([]rune{'>'}, to)
		e.insert([]rune{'<'}, from)
	}
}

func (e *basic) wordForward() {
	switch n := len(e.buf) - e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = len(e.buf)
		return
	}
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.indexFunc(isWhitespace, e.pos+1, false); i > 0 {
			e.pos = i
			return
		}
	case isKeyword(r):
		if i := e.indexFunc(isKeyword, e.pos+1, false); i > 0 {
			if !isWhitespace(e.buf[i]) {
				e.pos = i
				return
			}
			if i := e.indexFunc(isWhitespace, i+1, false); i > 0 {
				e.pos = i
				return
			}
		}
	default:
		if i := e.indexFunc(isSymbol, e.pos+1, false); i > 0 {
			if isKeyword(e.buf[i]) {
				e.pos = i
				return
			}
			if i := e.indexFunc(isWhitespace, i+1, false); i > 0 {
				e.pos = i
				return
			}
		}
	}
	e.pos = len(e.buf)
}

func (e *basic) wordBackward() {
	switch e.pos {
	case 0:
		return
	case 1:
		e.pos = 0
		return
	}

	n := e.pos - 1
	switch r := e.buf[n]; {
	case isWhitespace(r):
		n = e.lastIndexFunc(isWhitespace, n, false)
		if n < 0 {
			e.pos = 0
			return
		}
	}

	switch r := e.buf[n]; {
	case isKeyword(r):
		if i := e.lastIndexFunc(isKeyword, n, false); i >= 0 {
			e.pos = i + 1
			return
		}
	default:
		if i := e.lastIndexFunc(isSymbol, n, false); i >= 0 {
			e.pos = i + 1
			return
		}
	}
	e.pos = 0
}

func (e *basic) wordForwardNonBlank() {
	i := e.indexFunc(isWhitespace, e.pos, true)
	if i < 0 {
		e.pos = len(e.buf)
		return
	}
	i = e.indexFunc(isWhitespace, i+1, false)
	if i < 0 {
		e.pos = len(e.buf)
		return
	}
	e.pos = i
}

func (e *basic) wordBackwardNonBlank() {
	i := e.lastIndexFunc(isWhitespace, e.pos, false)
	if i < 0 {
		e.pos = 0
		return
	}
	i = e.lastIndexFunc(isWhitespace, i, true)
	if i < 0 {
		e.pos = 0
		return
	}
	e.pos = i + 1
}

func (e *basic) wordEnd() {
	switch n := len(e.buf) - e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = len(e.buf)
		return
	}
	e.pos++
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.indexFunc(isWhitespace, e.pos+1, false); i > 0 {
			switch r := e.buf[i]; {
			case isKeyword(r):
				if i := e.indexFunc(isKeyword, i+1, false); i > 0 {
					e.pos = i - 1
					return
				}
			default:
				if i := e.indexFunc(isSymbol, i+1, false); i > 0 {
					e.pos = i - 1
					return
				}
			}
		}
	case isKeyword(r):
		if i := e.indexFunc(isKeyword, e.pos+1, false); i > 0 {
			e.pos = i - 1
			return
		}
	default:
		if i := e.indexFunc(isSymbol, e.pos+1, false); i > 0 {
			e.pos = i - 1
			return
		}
	}
	e.pos = len(e.buf) - 1
}

func (e *basic) wordEndNonBlank() {
	switch n := len(e.buf) - e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = len(e.buf)
		return
	}
	e.pos++
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.indexFunc(isWhitespace, e.pos+1, false); i > 0 {
			if i := e.indexFunc(isWhitespace, i+1, true); i > 0 {
				e.pos = i - 1
				return
			}
		}
	default:
		if i := e.indexFunc(isWhitespace, e.pos+1, true); i > 0 {
			e.pos = i - 1
			return
		}
	}
	e.pos = len(e.buf) - 1
}

func (e *basic) wordEndBackward() {
	switch n := e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = 0
		return
	}
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.lastIndexFunc(isWhitespace, e.pos, false); i > 0 {
			e.pos = i
			return
		}
	case isKeyword(r):
		if i := e.lastIndexFunc(isKeyword, e.pos, false); i > 0 {
			switch {
			case isWhitespace(e.buf[i]):
				if i := e.lastIndexFunc(isWhitespace, i, false); i > 0 {
					e.pos = i
					return
				}
			default:
				e.pos = i
				return
			}
		}
	default:
		if i := e.lastIndexFunc(isSymbol, e.pos, false); i > 0 {
			switch {
			case isWhitespace(e.buf[i]):
				if i := e.lastIndexFunc(isWhitespace, i, false); i > 0 {
					e.pos = i
					return
				}
			default:
				e.pos = i
				return
			}
		}
	}
	e.pos = 0
}

func (e *basic) wordEndBackwardNonBlank() {
	switch n := e.pos; {
	case n < 1:
		return
	case n == 1:
		e.pos = 0
		return
	}
	switch r := e.buf[e.pos]; {
	case isWhitespace(r):
		if i := e.lastIndexFunc(isWhitespace, e.pos, false); i > 0 {
			e.pos = i
			return
		}
	default:
		if i := e.lastIndexFunc(isWhitespace, e.pos, true); i > 0 {
			if i := e.lastIndexFunc(isWhitespace, i, false); i > 0 {
				e.pos = i
				return
			}
		}
	}
	e.pos = 0
}
<|endoftext|>


<commit_before>package aufs

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
)

func exists(pth string) bool {
	if _, err := os.Stat(pth); err != nil {
		return false
	}
	return true
}

func (a *AufsDriver) Migrate(pth string) error {
	fis, err := ioutil.ReadDir(pth)
	if err != nil {
		return err
	}
	for _, fi := range fis {
		if fi.IsDir() && exists(path.Join(pth, fi.Name(), "layer")) && !a.Exists(fi.Name()) {
			if err := tryRelocate(path.Join(pth, fi.Name(), "layer"), path.Join(a.rootPath(), "diff", fi.Name())); err != nil {
				return err
			}
			if err := a.Create(fi.Name(), ""); err != nil {
				return err
			}
		}
	}
	return nil
}

// tryRelocate will try to rename the old path to the new pack and if
// the operation fails, it will fallback to a symlink
func tryRelocate(oldPath, newPath string) error {
	if err := os.Rename(oldPath, newPath); err != nil {
		if sErr := os.Symlink(oldPath, newPath); sErr != nil {
			return fmt.Errorf("Unable to relocate %s to %s: Rename err %s Symlink err %s", oldPath, newPath, err, sErr)
		}
	}
	return nil
}
<commit_msg>Migrate images with parent relationship<commit_after>package aufs

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"time"
)

type imageMetadata struct {
	ID            string    `json:"id"`
	ParentID      string    `json:"parent,omitempty"`
	Created       time.Time `json:"created"`
	DockerVersion string    `json:"docker_version,omitempty"`
	Architecture  string    `json:"architecture,omitempty"`

	parent *imageMetadata
}

func pathExists(pth string) bool {
	if _, err := os.Stat(pth); err != nil {
		return false
	}
	return true
}

// Migrate existing images and containers from docker < 0.7.x
func (a *AufsDriver) Migrate(pth string) error {
	fis, err := ioutil.ReadDir(pth)
	if err != nil {
		return err
	}
	var (
		metadata = make(map[string]*imageMetadata)
		current  *imageMetadata
		exists   bool
	)

	// Load metadata
	for _, fi := range fis {
		if id := fi.Name(); fi.IsDir() && pathExists(path.Join(pth, id, "layer")) && !a.Exists(id) {
			if current, exists = metadata[id]; !exists {
				current, err = loadMetadata(pth, id)
				if err != nil {
					return err
				}
				metadata[id] = current
			}
		}
	}

	// Recreate tree
	for _, v := range metadata {
		v.parent = metadata[v.ParentID]
	}

	// Perform image migration
	for _, v := range metadata {
		if err := migrateImage(v, a, pth); err != nil {
			return err
		}
	}
	return nil
}

func migrateImage(m *imageMetadata, a *AufsDriver, pth string) error {
	if !pathExists(path.Join(a.rootPath(), "diff", m.ID)) {
		if m.parent != nil {
			migrateImage(m.parent, a, pth)
		}
		if err := tryRelocate(path.Join(pth, m.ID, "layer"), path.Join(a.rootPath(), "diff", m.ID)); err != nil {
			return err
		}

		if err := a.Create(m.ID, m.ParentID); err != nil {
			return err
		}
	}
	return nil
}

// tryRelocate will try to rename the old path to the new pack and if
// the operation fails, it will fallback to a symlink
func tryRelocate(oldPath, newPath string) error {
	if err := os.Rename(oldPath, newPath); err != nil {
		if sErr := os.Symlink(oldPath, newPath); sErr != nil {
			return fmt.Errorf("Unable to relocate %s to %s: Rename err %s Symlink err %s", oldPath, newPath, err, sErr)
		}
	}
	return nil
}

func loadMetadata(pth, id string) (*imageMetadata, error) {
	f, err := os.Open(path.Join(pth, id, "json"))
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var (
		out = &imageMetadata{}
		dec = json.NewDecoder(f)
	)

	if err := dec.Decode(out); err != nil {
		return nil, err
	}
	return out, nil
}
<|endoftext|>


<commit_before>package hoverfly

import (
	"bytes"
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/boltdb/bolt"
)

// Cache - cache interface used to store and retrieve request/response payloads
type Cache interface {
	Set(key, value []byte) error
	Get(key []byte) ([]byte, error)
	GetAllRequests() ([]Payload, error)
	RecordsCount() (int, error)
	DeleteData() error
	GetAllKeys() (map[string]bool, error)
	CloseDB()
}

// NewBoltDBCache - returns new BoltCache instance
func NewBoltDBCache(db *bolt.DB, bucket []byte) BoltCache {
	cache := BoltCache{
		DS:             db,
		RequestsBucket: []byte(bucket),
	}
	return cache
}

// RequestsBucketName - default name for BoltDB bucket
const RequestsBucketName = "rqbucket"

// Cache - provides access to BoltDB and holds current bucket name
type Cache struct {
	DS             *bolt.DB
	RequestsBucket []byte
}

// GetDB - returns open BoltDB database with read/write permissions or goes down in flames if
// something bad happends
func GetDB(name string) *bolt.DB {
	log.WithFields(log.Fields{
		"databaseName": name,
	}).Info("Initiating database")
	db, err := bolt.Open(name, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}

	return db
}

// Set - saves given key and value pair to cache
func (c *Cache) Set(key, value []byte) error {
	err := c.DS.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists(c.RequestsBucket)
		if err != nil {
			return err
		}
		err = bucket.Put(key, value)
		if err != nil {
			return err
		}
		return nil
	})

	return err
}

// Get - searches for given key in the cache and returns value if found
func (c *Cache) Get(key []byte) (value []byte, err error) {

	err = c.DS.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(c.RequestsBucket)
		if bucket == nil {
			return fmt.Errorf("Bucket %q not found!", c.RequestsBucket)
		}
		// "Byte slices returned from Bolt are only valid during a transaction."
		var buffer bytes.Buffer
		val := bucket.Get(key)

		// If it doesn't exist then it will return nil
		if val == nil {
			return fmt.Errorf("key %q not found \n", key)
		}

		buffer.Write(val)
		value = buffer.Bytes()
		return nil
	})

	return
}

// GetAllRequests - returns all captured requests/responses
func (c *Cache) GetAllRequests() (payloads []Payload, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)
		if b == nil {
			// bucket doesn't exist
			return nil
		}
		c := b.Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			pl, err := decodePayload(v)
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
					"json":  v,
				}).Warning("Failed to deserialize bytes to payload.")
			} else {
				payloads = append(payloads, *pl)
			}
		}
		return nil
	})
	return
}

// RecordsCount - returns records count
func (c *Cache) RecordsCount() (count int, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)
		if b == nil {
			// bucket doesn't exist
			return nil
		}

		count = b.Stats().KeyN

		return nil
	})
	return
}

// DeleteBucket - deletes bucket with all saved data
func (c *Cache) DeleteBucket(name []byte) (err error) {
	err = c.DS.Update(func(tx *bolt.Tx) error {
		err = tx.DeleteBucket(name)
		if err != nil {
			log.WithFields(log.Fields{
				"error": err.Error(),
				"name":  string(name),
			}).Warning("Failed to delete bucket")

		}
		return err
	})
	return
}

// GetAllKeys - gets all current keys
func (c *Cache) GetAllKeys() (keys map[string]bool, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)

		keys = make(map[string]bool)

		if b == nil {
			// bucket doesn't exist
			return nil
		}
		c := b.Cursor()

		for k, _ := c.First(); k != nil; k, _ = c.Next() {
			keys[string(k)] = true
		}
		return nil
	})
	return
}
<commit_msg>Cache renamed to BoltCache<commit_after>package hoverfly

import (
	"bytes"
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/boltdb/bolt"
)

// Cache - cache interface used to store and retrieve request/response payloads
type Cache interface {
	Set(key, value []byte) error
	Get(key []byte) ([]byte, error)
	GetAllRequests() ([]Payload, error)
	RecordsCount() (int, error)
	DeleteData() error
	GetAllKeys() (map[string]bool, error)
	CloseDB()
}

// NewBoltDBCache - returns new BoltCache instance
func NewBoltDBCache(db *bolt.DB, bucket []byte) BoltCache {
	cache := BoltCache{
		DS:             db,
		RequestsBucket: []byte(bucket),
	}
	return cache
}

// RequestsBucketName - default name for BoltDB bucket
const RequestsBucketName = "rqbucket"

// Cache - provides access to BoltDB and holds current bucket name
type BoltCache struct {
	DS             *bolt.DB
	RequestsBucket []byte
}

// GetDB - returns open BoltDB database with read/write permissions or goes down in flames if
// something bad happends
func GetDB(name string) *bolt.DB {
	log.WithFields(log.Fields{
		"databaseName": name,
	}).Info("Initiating database")
	db, err := bolt.Open(name, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}

	return db
}

// CloseDB - closes database
func (c *BoltCache) CloseDB() {
	c.DS.Close()
}

// Set - saves given key and value pair to cache
func (c *BoltCache) Set(key, value []byte) error {
	err := c.DS.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists(c.RequestsBucket)
		if err != nil {
			return err
		}
		err = bucket.Put(key, value)
		if err != nil {
			return err
		}
		return nil
	})

	return err
}

// Get - searches for given key in the cache and returns value if found
func (c *BoltCache) Get(key []byte) (value []byte, err error) {

	err = c.DS.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(c.RequestsBucket)
		if bucket == nil {
			return fmt.Errorf("Bucket %q not found!", c.RequestsBucket)
		}
		// "Byte slices returned from Bolt are only valid during a transaction."
		var buffer bytes.Buffer
		val := bucket.Get(key)

		// If it doesn't exist then it will return nil
		if val == nil {
			return fmt.Errorf("key %q not found \n", key)
		}

		buffer.Write(val)
		value = buffer.Bytes()
		return nil
	})

	return
}

// GetAllRequests - returns all captured requests/responses
func (c *BoltCache) GetAllRequests() (payloads []Payload, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)
		if b == nil {
			// bucket doesn't exist
			return nil
		}
		c := b.Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			pl, err := decodePayload(v)
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
					"json":  v,
				}).Warning("Failed to deserialize bytes to payload.")
			} else {
				payloads = append(payloads, *pl)
			}
		}
		return nil
	})
	return
}

// RecordsCount - returns records count
func (c *BoltCache) RecordsCount() (count int, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)
		if b == nil {
			// bucket doesn't exist
			return nil
		}

		count = b.Stats().KeyN

		return nil
	})
	return
}

// DeleteData - deletes bucket with all saved data
func (c *BoltCache) DeleteData() error {
	err := c.DeleteBucket(c.RequestsBucket)
	return err
}

// DeleteBucket - deletes bucket with all saved data
func (c *BoltCache) DeleteBucket(name []byte) (err error) {
	err = c.DS.Update(func(tx *bolt.Tx) error {
		err = tx.DeleteBucket(name)
		if err != nil {
			log.WithFields(log.Fields{
				"error": err.Error(),
				"name":  string(name),
			}).Warning("Failed to delete bucket")

		}
		return err
	})
	return
}

// GetAllKeys - gets all current keys
func (c *BoltCache) GetAllKeys() (keys map[string]bool, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)

		keys = make(map[string]bool)

		if b == nil {
			// bucket doesn't exist
			return nil
		}
		c := b.Cursor()

		for k, _ := c.First(); k != nil; k, _ = c.Next() {
			keys[string(k)] = true
		}
		return nil
	})
	return
}
<|endoftext|>


<commit_before>// The reason that so many of these functions seem like duplication is simply for the better error
// messages so each driver doesn't need to deal with validation on these simple field types

package drivers

import (
	"os"
	"strconv"
	"strings"

	"github.com/pkg/errors"
)

// Config is a map with helper functions
type Config map[string]interface{}

// MustString retrieves a string that must exist and must be a string, it must also not be the empty string
func (c Config) MustString(key string) string {
	s, ok := c[key]
	if !ok {
		panic(errors.Errorf("failed to find key %s in config", key))
	}

	str, ok := s.(string)
	if !ok {
		panic(errors.Errorf("found key %s in config, but it was not a string (%T)", key, s))
	}

	if len(str) == 0 {
		panic(errors.Errorf("found key %s in config, but it was an empty string", key))
	}

	return str
}

// MustInt retrieves a string that must exist and must be an int, and it must not be 0
func (c Config) MustInt(key string) int {
	i, ok := c[key]
	if !ok {
		panic(errors.Errorf("failed to find key %s in config", key))
	}

	var integer int
	switch t := i.(type) {
	case int:
		integer = t
	case float64:
		integer = int(t)
	case string:
		integer, _ = strconv.Atoi(t)
	default:
		panic(errors.Errorf("found key %s in config, but it was not an int (%T)", key, i))
	}

	if integer == 0 {
		panic(errors.Errorf("found key %s in config, but its value was 0", key))
	}

	return integer
}

// String retrieves a non-empty string, the bool says if it exists, is of appropriate type,
// and has a non-zero length or not.
func (c Config) String(key string) (string, bool) {
	s, ok := c[key]
	if !ok {
		return "", false
	}

	str, ok := s.(string)
	if !ok {
		return "", false
	}

	if len(str) == 0 {
		return "", false
	}

	return str, true
}

// DefaultString retrieves a non-empty string or the default value provided.
func (c Config) DefaultString(key, def string) string {
	str, ok := c.String(key)
	if !ok {
		return def
	}

	return str
}

// Int retrieves an int, the bool says if it exists, is of the appropriate type,
// and is non-zero. Coerces float64 to int because JSON and Javascript kinda suck.
func (c Config) Int(key string) (int, bool) {
	i, ok := c[key]
	if !ok {
		return 0, false
	}

	var integer int
	switch t := i.(type) {
	case int:
		integer = t
	case float64:
		integer = int(t)
	case string:
		integer, _ = strconv.Atoi(t)
	default:
		return 0, false
	}

	if integer == 0 {
		return 0, false
	}

	return integer, true
}

// DefaultInt retrieves a non-zero int or the default value provided.
func (c Config) DefaultInt(key string, def int) int {
	i, ok := c.Int(key)
	if !ok {
		return def
	}

	return i
}

// StringSlice retrieves an string slice, the bool says if it exists, is of the appropriate type,
// is non-nil and non-zero length
func (c Config) StringSlice(key string) ([]string, bool) {
	ss, ok := c[key]
	if !ok {
		return nil, false
	}

	var slice []string
	if intfSlice, ok := ss.([]interface{}); ok {
		for _, i := range intfSlice {
			slice = append(slice, i.(string))
		}
	} else if stringSlice, ok := ss.([]string); ok {
		slice = stringSlice
	} else {
		return nil, false
	}

	// Also detects nil
	if len(slice) == 0 {
		return nil, false
	}
	return slice, true
}

// DefaultEnv grabs a value from the environment or a default.
// This is shared by drivers to get config for testing.
func DefaultEnv(key, def string) string {
	val := os.Getenv(key)
	if len(val) == 0 {
		val = def
	}
	return val
}

// TablesFromList takes a whitelist or blacklist and returns
// the table names.
func TablesFromList(list []string) []string {
	if len(list) == 0 {
		return nil
	}

	var tables []string
	for _, i := range list {
		splits := strings.Split(i, ".")

		if len(splits) == 1 {
			tables = append(tables, splits[0])
		}
	}

	return tables
}

// ColumnsFromList takes a whitelist or blacklist and returns
// the columns for a given table.
func ColumnsFromList(list []string, tablename string) []string {
	if len(list) == 0 {
		return nil
	}

	var columns []string
	for _, i := range list {
		splits := strings.Split(i, ".")

		if len(splits) != 2 {
			continue
		}

		if splits[0] == tablename {
			columns = append(columns, splits[1])
		}
	}

	return columns
}
<commit_msg>Don't ignore parse Int error<commit_after>// The reason that so many of these functions seem like duplication is simply for the better error
// messages so each driver doesn't need to deal with validation on these simple field types

package drivers

import (
	"os"
	"strconv"
	"strings"

	"github.com/pkg/errors"
)

// Config is a map with helper functions
type Config map[string]interface{}

// MustString retrieves a string that must exist and must be a string, it must also not be the empty string
func (c Config) MustString(key string) string {
	s, ok := c[key]
	if !ok {
		panic(errors.Errorf("failed to find key %s in config", key))
	}

	str, ok := s.(string)
	if !ok {
		panic(errors.Errorf("found key %s in config, but it was not a string (%T)", key, s))
	}

	if len(str) == 0 {
		panic(errors.Errorf("found key %s in config, but it was an empty string", key))
	}

	return str
}

// MustInt retrieves a string that must exist and must be an int, and it must not be 0
func (c Config) MustInt(key string) int {
	i, ok := c[key]
	if !ok {
		panic(errors.Errorf("failed to find key %s in config", key))
	}

	var integer int
	switch t := i.(type) {
	case int:
		integer = t
	case float64:
		integer = int(t)
	case string:
		var err error
		integer, err = strconv.Atoi(t)
		if err != nil {
			panic(errors.Errorf("fail to parse %v to int: %v", t, err))
		}
	default:
		panic(errors.Errorf("found key %s in config, but it was not an int (%T)", key, i))
	}

	if integer == 0 {
		panic(errors.Errorf("found key %s in config, but its value was 0", key))
	}

	return integer
}

// String retrieves a non-empty string, the bool says if it exists, is of appropriate type,
// and has a non-zero length or not.
func (c Config) String(key string) (string, bool) {
	s, ok := c[key]
	if !ok {
		return "", false
	}

	str, ok := s.(string)
	if !ok {
		return "", false
	}

	if len(str) == 0 {
		return "", false
	}

	return str, true
}

// DefaultString retrieves a non-empty string or the default value provided.
func (c Config) DefaultString(key, def string) string {
	str, ok := c.String(key)
	if !ok {
		return def
	}

	return str
}

// Int retrieves an int, the bool says if it exists, is of the appropriate type,
// and is non-zero. Coerces float64 to int because JSON and Javascript kinda suck.
func (c Config) Int(key string) (int, bool) {
	i, ok := c[key]
	if !ok {
		return 0, false
	}

	var integer int
	switch t := i.(type) {
	case int:
		integer = t
	case float64:
		integer = int(t)
	case string:
		var err error
		integer, err = strconv.Atoi(t)
		if err != nil {
			return 0, false
		}
	default:
		return 0, false
	}

	if integer == 0 {
		return 0, false
	}

	return integer, true
}

// DefaultInt retrieves a non-zero int or the default value provided.
func (c Config) DefaultInt(key string, def int) int {
	i, ok := c.Int(key)
	if !ok {
		return def
	}

	return i
}

// StringSlice retrieves an string slice, the bool says if it exists, is of the appropriate type,
// is non-nil and non-zero length
func (c Config) StringSlice(key string) ([]string, bool) {
	ss, ok := c[key]
	if !ok {
		return nil, false
	}

	var slice []string
	if intfSlice, ok := ss.([]interface{}); ok {
		for _, i := range intfSlice {
			slice = append(slice, i.(string))
		}
	} else if stringSlice, ok := ss.([]string); ok {
		slice = stringSlice
	} else {
		return nil, false
	}

	// Also detects nil
	if len(slice) == 0 {
		return nil, false
	}
	return slice, true
}

// DefaultEnv grabs a value from the environment or a default.
// This is shared by drivers to get config for testing.
func DefaultEnv(key, def string) string {
	val := os.Getenv(key)
	if len(val) == 0 {
		val = def
	}
	return val
}

// TablesFromList takes a whitelist or blacklist and returns
// the table names.
func TablesFromList(list []string) []string {
	if len(list) == 0 {
		return nil
	}

	var tables []string
	for _, i := range list {
		splits := strings.Split(i, ".")

		if len(splits) == 1 {
			tables = append(tables, splits[0])
		}
	}

	return tables
}

// ColumnsFromList takes a whitelist or blacklist and returns
// the columns for a given table.
func ColumnsFromList(list []string, tablename string) []string {
	if len(list) == 0 {
		return nil
	}

	var columns []string
	for _, i := range list {
		splits := strings.Split(i, ".")

		if len(splits) != 2 {
			continue
		}

		if splits[0] == tablename {
			columns = append(columns, splits[1])
		}
	}

	return columns
}
<|endoftext|>