<commit_before>package main

import (
	"bytes"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/json-iterator/go"
	"github.com/valyala/fasthttp"
)

type domainConfig struct {
	apiReplaces                []replace
	apiOfficialLongpollReplace replace
	apiLongpollReplace         replace
	siteHlsReplace             replace
}

var client = &fasthttp.Client{
	Name: "vk-proxy",
}
var domains = make(map[string]*domainConfig)
var json = jsoniter.ConfigFastest

func getDomainConfig(domain string) *domainConfig {
	cfg, ok := domains[domain]
	if !ok {
		var replacementStart = []byte(`\/\/` + domain + `\/_\/`)
		cfg = &domainConfig{}
		cfg.apiReplaces = []replace{
			newStringReplace(`"https:\/\/vk.com\/video_hls.php`, `"https:\/\/`+domain+`\/@vk.com\/video_hls.php`),
			newRegexFastReplace(`\\/\\/(?:pu\.vk\.com|[-_a-zA-Z0-9]{1,15}\.(?:userapi\.com|vk-cdn\.net|vk\.me|vkuser(?:live|video|audio)\.(?:net|com)))\\/`,
				func(src, dst []byte, start, end int) []byte {
					if start < 7 || !bytes.Equal(src[start-7:start], jsonUrlPrefix) {
						return append(dst, src[start:end]...)
					}
					dst = append(dst, replacementStart...)
					dst = append(dst, src[start+4:end]...)
					return dst
				}),
			newRegexReplace(`"https:\\/\\/vk\.com\\/(images\\/|doc-?[0-9]+_)`, `"https:\/\/`+domain+`\/_\/vk.com\/$1`),
		}
		cfg.apiOfficialLongpollReplace = newStringReplace(`"server":"api.vk.com\/newuim`, `"server":"`+domain+`\/_\/api.vk.com\/newuim`)
		cfg.apiLongpollReplace = newStringReplace(`"server":"`, `"server":"`+domain+`\/_\/`)
		cfg.siteHlsReplace = newRegexReplace(`https:\/\/([-_a-zA-Z0-9]+\.(?:userapi\.com|vk-cdn\.net|vk\.me|vkuser(?:live|video)\.(?:net|com)))\/`, `https://`+domain+`/_/$1/`)
		domains[domain] = cfg
	}
	return cfg
}

func reverseProxyHandler(ctx *fasthttp.RequestCtx) {
	defer func() {
		if r := recover(); r != nil {
			ctx.Logger().Printf("panic when proxying the request: %s", r)
			ctx.Response.Reset()
			ctx.SetStatusCode(500)
			ctx.SetBodyString("500 Internal Server Error")
		}
	}()

	var config *domainConfig
	if Config.domain != "" {
		config = getDomainConfig(Config.domain)
	} else {
		config = getDomainConfig(string(ctx.Host()))
	}

	if !preRequest(ctx) {
		ctx.Response.SetStatusCode(400)
		ctx.Response.SetBodyString("400 Bad Request")
		return
	}

	// In case of redirect
	if ctx.Response.StatusCode() != 200 {
		trackRequest(ctx, 0)
		return
	}

	err := client.DoTimeout(&ctx.Request, &ctx.Response, 30*time.Second)
	if err == nil {
		err = postResponse(config, ctx)
	}

	if err != nil {
		ctx.Logger().Printf("error when proxying the request: %s", err)
		ctx.Response.Reset()
		if strings.Contains(err.Error(), "timed out") || strings.Contains(err.Error(), "timeout") {
			ctx.SetStatusCode(408)
			ctx.SetBodyString("408 Request Timeout")
		} else {
			ctx.SetStatusCode(500)
			ctx.SetBodyString("500 Internal Server Error")
		}
		trackRequest(ctx, 0)
	}
}

func preRequest(ctx *fasthttp.RequestCtx) bool {
	req := &ctx.Request
	path := req.RequestURI()
	if bytes.HasPrefix(path, atPath) {
		slashIndex := bytes.IndexRune(path[1:], '/')
		if slashIndex == -1 {
			return false
		}
		endpoint := []byte(path[4:slashIndex+1])
		if !bytes.Equal(endpoint, siteHost) && !bytes.HasSuffix(endpoint, siteHostRoot) {
			return false
		}
		req.Header.SetHostBytes(endpoint)
		req.SetRequestURIBytes([]byte(path[1+slashIndex:]))
	} else if bytes.HasPrefix(path, awayPath) {
		to := string(req.URI().QueryArgs().Peek("to"))
		if to == "" {
			return false
		}
		to, err := url.QueryUnescape(to)
		if err != nil {
			return false
		}
		ctx.Redirect(to, 301)
		return true
	} else {
		req.SetHostBytes(apiHost)
	}
	// After req.URI() call it is impossible to modify URI
	req.URI().SetSchemeBytes(https)
	if Config.gzipUpstream {
		req.Header.SetBytesKV(acceptEncoding, gzip)
	} else {
		req.Header.DelBytes(acceptEncoding)
	}
	return true
}

func postResponse(config *domainConfig, ctx *fasthttp.RequestCtx) error {
	res := &ctx.Response
	res.Header.DelBytes(setCookie)
	res.Header.SetBytesKV(serverHeader, vkProxyName)
	var body []byte
	if bytes.Contains(res.Header.PeekBytes(contentEncoding), gzip) {
		res.Header.DelBytes(contentEncoding)
		var err error
		body, err = res.BodyGunzip()
		if err != nil {
			return err
		}
	} else {
		body = res.Body()
	}
	uri := ctx.Request.URI()
	path := uri.Path()

	if bytes.Equal(uri.Host(), siteHost) {
		if bytes.Equal(path, videoHlsPath) {
			body = config.siteHlsReplace.apply(body)
		}
	} else {
		for _, replace := range config.apiReplaces {
			body = replace.apply(body)
		}

		// Replace longpoll server
		if bytes.Equal(path, apiLongpollPath) {
			body = config.apiLongpollReplace.apply(body)
		} else

		// Replace longpoll server for official app
		if bytes.Equal(path, apiOfficialLongpollPath) ||
			bytes.Equal(path, apiOfficialLongpollPath2) ||
			bytes.Equal(path, apiOfficialLongpollPath3) {
			body = config.apiOfficialLongpollReplace.apply(body)
		}

		// Clear feed from SPAM
		if Config.removeAdsFromFeed {
			if bytes.Equal(path, apiOfficialNewsfeedPath) ||
				bytes.Equal(path, apiNewsfeedGet) {
				var parsed map[string]interface{}
				if err := json.Unmarshal(body, &parsed); err == nil {
					if parsed["response"] != nil {
						response := parsed["response"].(map[string]interface{})
						if response["items"] != nil {
							newItems, modified := filterFeed(response["items"].([]interface{}))
							if modified {
								response["items"] = newItems
								body, err = json.Marshal(parsed)
							}
						}
					}
				}
			}
		}
	}
	res.SetBody(body)

	if Config.debug {
		log.Println(string(path) + "\n" + string(body))
	}

	trackRequest(ctx, len(body))
	return nil
}

func filterFeed(items []interface{}) ([]interface{}, bool) {
	removed := 0
	for i := len(items) - 1; i >= 0; i-- {
		post := items[i].(map[string]interface{})
		if post["type"] == "ads" || (post["type"] == "post" && post["marked_as_ads"] != nil && post["marked_as_ads"].(float64) == 1) {
			items[i] = items[len(items)-1]
			items[len(items)-1] = nil
			items = items[:len(items)-1]
			removed++
		}
	}
	if removed > 0 {
		newItems := make([]interface{}, len(items))
		copy(newItems, items)
		return newItems, true
	}
	return nil, false
}
<commit_msg>Починены стикеры<commit_after>package main

import (
	"bytes"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/json-iterator/go"
	"github.com/valyala/fasthttp"
)

type domainConfig struct {
	apiReplaces                []replace
	apiOfficialLongpollReplace replace
	apiLongpollReplace         replace
	siteHlsReplace             replace
}

var client = &fasthttp.Client{
	Name: "vk-proxy",
}
var domains = make(map[string]*domainConfig)
var json = jsoniter.ConfigFastest

func getDomainConfig(domain string) *domainConfig {
	cfg, ok := domains[domain]
	if !ok {
		var replacementStart = []byte(`\/\/` + domain + `\/_\/`)
		cfg = &domainConfig{}
		cfg.apiReplaces = []replace{
			newStringReplace(`"https:\/\/vk.com\/video_hls.php`, `"https:\/\/`+domain+`\/@vk.com\/video_hls.php`),
			newRegexFastReplace(`\\/\\/(?:pu\.vk\.com|[-_a-zA-Z0-9]{1,15}\.(?:userapi\.com|vk-cdn\.net|vk\.me|vkuser(?:live|video|audio)\.(?:net|com)))\\/`,
				func(src, dst []byte, start, end int) []byte {
					if start < 7 || !bytes.Equal(src[start-7:start], jsonUrlPrefix) {
						return append(dst, src[start:end]...)
					}
					dst = append(dst, replacementStart...)
					dst = append(dst, src[start+4:end]...)
					return dst
				}),
			newRegexReplace(`"https:\\/\\/vk\.com\\/((?:\\/)?images\\/|doc-?[0-9]+_)`, `"https:\/\/`+domain+`\/_\/vk.com\/$1`),
		}
		cfg.apiOfficialLongpollReplace = newStringReplace(`"server":"api.vk.com\/newuim`, `"server":"`+domain+`\/_\/api.vk.com\/newuim`)
		cfg.apiLongpollReplace = newStringReplace(`"server":"`, `"server":"`+domain+`\/_\/`)
		cfg.siteHlsReplace = newRegexReplace(`https:\/\/([-_a-zA-Z0-9]+\.(?:userapi\.com|vk-cdn\.net|vk\.me|vkuser(?:live|video)\.(?:net|com)))\/`, `https://`+domain+`/_/$1/`)
		domains[domain] = cfg
	}
	return cfg
}

func reverseProxyHandler(ctx *fasthttp.RequestCtx) {
	defer func() {
		if r := recover(); r != nil {
			ctx.Logger().Printf("panic when proxying the request: %s", r)
			ctx.Response.Reset()
			ctx.SetStatusCode(500)
			ctx.SetBodyString("500 Internal Server Error")
		}
	}()

	var config *domainConfig
	if Config.domain != "" {
		config = getDomainConfig(Config.domain)
	} else {
		config = getDomainConfig(string(ctx.Host()))
	}

	if !preRequest(ctx) {
		ctx.Response.SetStatusCode(400)
		ctx.Response.SetBodyString("400 Bad Request")
		return
	}

	// In case of redirect
	if ctx.Response.StatusCode() != 200 {
		trackRequest(ctx, 0)
		return
	}

	err := client.DoTimeout(&ctx.Request, &ctx.Response, 30*time.Second)
	if err == nil {
		err = postResponse(config, ctx)
	}

	if err != nil {
		ctx.Logger().Printf("error when proxying the request: %s", err)
		ctx.Response.Reset()
		if strings.Contains(err.Error(), "timed out") || strings.Contains(err.Error(), "timeout") {
			ctx.SetStatusCode(408)
			ctx.SetBodyString("408 Request Timeout")
		} else {
			ctx.SetStatusCode(500)
			ctx.SetBodyString("500 Internal Server Error")
		}
		trackRequest(ctx, 0)
	}
}

func preRequest(ctx *fasthttp.RequestCtx) bool {
	req := &ctx.Request
	path := req.RequestURI()
	if bytes.HasPrefix(path, atPath) {
		slashIndex := bytes.IndexRune(path[1:], '/')
		if slashIndex == -1 {
			return false
		}
		endpoint := []byte(path[4:slashIndex+1])
		if !bytes.Equal(endpoint, siteHost) && !bytes.HasSuffix(endpoint, siteHostRoot) {
			return false
		}
		req.Header.SetHostBytes(endpoint)
		req.SetRequestURIBytes([]byte(path[1+slashIndex:]))
	} else if bytes.HasPrefix(path, awayPath) {
		to := string(req.URI().QueryArgs().Peek("to"))
		if to == "" {
			return false
		}
		to, err := url.QueryUnescape(to)
		if err != nil {
			return false
		}
		ctx.Redirect(to, 301)
		return true
	} else {
		req.SetHostBytes(apiHost)
	}
	// After req.URI() call it is impossible to modify URI
	req.URI().SetSchemeBytes(https)
	if Config.gzipUpstream {
		req.Header.SetBytesKV(acceptEncoding, gzip)
	} else {
		req.Header.DelBytes(acceptEncoding)
	}
	return true
}

func postResponse(config *domainConfig, ctx *fasthttp.RequestCtx) error {
	res := &ctx.Response
	res.Header.DelBytes(setCookie)
	res.Header.SetBytesKV(serverHeader, vkProxyName)
	var body []byte
	if bytes.Contains(res.Header.PeekBytes(contentEncoding), gzip) {
		res.Header.DelBytes(contentEncoding)
		var err error
		body, err = res.BodyGunzip()
		if err != nil {
			return err
		}
	} else {
		body = res.Body()
	}
	uri := ctx.Request.URI()
	path := uri.Path()

	if bytes.Equal(uri.Host(), siteHost) {
		if bytes.Equal(path, videoHlsPath) {
			body = config.siteHlsReplace.apply(body)
		}
	} else {
		for _, replace := range config.apiReplaces {
			body = replace.apply(body)
		}

		// Replace longpoll server
		if bytes.Equal(path, apiLongpollPath) {
			body = config.apiLongpollReplace.apply(body)
		} else

		// Replace longpoll server for official app
		if bytes.Equal(path, apiOfficialLongpollPath) ||
			bytes.Equal(path, apiOfficialLongpollPath2) ||
			bytes.Equal(path, apiOfficialLongpollPath3) {
			body = config.apiOfficialLongpollReplace.apply(body)
		}

		// Clear feed from SPAM
		if Config.removeAdsFromFeed {
			if bytes.Equal(path, apiOfficialNewsfeedPath) ||
				bytes.Equal(path, apiNewsfeedGet) {
				var parsed map[string]interface{}
				if err := json.Unmarshal(body, &parsed); err == nil {
					if parsed["response"] != nil {
						response := parsed["response"].(map[string]interface{})
						if response["items"] != nil {
							newItems, modified := filterFeed(response["items"].([]interface{}))
							if modified {
								response["items"] = newItems
								body, err = json.Marshal(parsed)
							}
						}
					}
				}
			}
		}
	}
	res.SetBody(body)

	if Config.debug {
		log.Println(string(path) + "\n" + string(body))
	}

	trackRequest(ctx, len(body))
	return nil
}

func filterFeed(items []interface{}) ([]interface{}, bool) {
	removed := 0
	for i := len(items) - 1; i >= 0; i-- {
		post := items[i].(map[string]interface{})
		if post["type"] == "ads" || (post["type"] == "post" && post["marked_as_ads"] != nil && post["marked_as_ads"].(float64) == 1) {
			items[i] = items[len(items)-1]
			items[len(items)-1] = nil
			items = items[:len(items)-1]
			removed++
		}
	}
	if removed > 0 {
		newItems := make([]interface{}, len(items))
		copy(newItems, items)
		return newItems, true
	}
	return nil, false
}
<|endoftext|>


<commit_before>package solidproxy

import (
	"net/http"
	"net/url"

	"github.com/elazarl/goproxy"
)

var (
	proxy = goproxy.NewProxyHttpServer()
)

func InitProxy(conf *ServerConfig) {
	proxy.OnResponse().DoFunc(func(r *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
		// CORS
		r.Header.Set("Access-Control-Allow-Credentials", "true")
		r.Header.Set("Access-Control-Expose-Headers", "User, Triples, Location, Link, Vary, Last-Modified, Content-Length")
		r.Header.Set("Access-Control-Max-Age", "60")

		// Retry with server credentials if authentication is required
		if r.StatusCode == 401 {
			Logger.Println("Resource " + ctx.Req.URL.String() + " requires authentication (HTTP 401).")
			req, err := http.NewRequest("GET", ctx.Req.URL.String(), ctx.Req.Body)
			req.Header.Add("On-Behalf-Of", conf.User)
			resp, err := agentClient.Do(req)
			if err != nil {
				Logger.Fatal("GET:", err)
			}
			Logger.Println("Retrying with server credentials...received data with status HTTP", resp.StatusCode)
			return resp
		}
		Logger.Println("Received data with status HTTP", r.StatusCode)
		return r
	})
}

func ProxyHandler(w http.ResponseWriter, req *http.Request) {
	uri, err := url.Parse(req.FormValue("uri"))
	if err != nil {
		Logger.Println("Error parsing URL:", req.URL, err.Error())
	}
	req.URL = uri
	req.Host = uri.Host
	req.RequestURI = uri.RequestURI()
	Logger.Println("Proxying request for:", req.URL)
	proxy.ServeHTTP(w, req)
	return
}
<commit_msg>adds cookie auth support<commit_after>package solidproxy

import (
	"crypto/tls"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/elazarl/goproxy"
)

var (
	proxy    = goproxy.NewProxyHttpServer()
	cookies  = map[string]map[string][]*http.Cookie{}
	cookiesL = new(sync.RWMutex)
)

func InitProxy(conf *ServerConfig) {
	proxy.OnResponse().DoFunc(func(r *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
		// CORS
		r.Header.Set("Access-Control-Allow-Credentials", "true")
		r.Header.Set("Access-Control-Expose-Headers", "User, Triples, Location, Link, Vary, Last-Modified, Content-Length")
		r.Header.Set("Access-Control-Max-Age", "60")

		// Retry with server credentials if authentication is required
		if r.StatusCode == 401 {
			defer TimeTrack(time.Now(), "Fetching")
			var resp *http.Response
			var client *http.Client
			var solutionMsg string
			req, err := http.NewRequest("GET", ctx.Req.URL.String(), ctx.Req.Body)
			req.Header.Add("On-Behalf-Of", conf.User)
			// Retry the request
			if len(cookies[conf.User]) > 0 { // Use existing cookie
				req.AddCookie(cookies[conf.User][ctx.Req.Host][0])
				// Create the client
				client = &http.Client{
					Transport: &http.Transport{
						TLSClientConfig: &tls.Config{
							InsecureSkipVerify: conf.Insecure,
						},
					},
				}
				solutionMsg = "Retrying with cookies"
			} else { // Using WebIDTLS client
				client = webidTlsClient
				solutionMsg = "Retrying with WebID-TLS"
			}
			resp, err = client.Do(req)
			if err != nil {
				Logger.Fatal("GET:", err)
			}
			// Store cookies per user and request host
			if len(resp.Cookies()) > 0 {
				cookiesL.Lock()
				// Should store cookies based on domain value AND path from cookie
				cookies[conf.User] = map[string][]*http.Cookie{}
				cookies[conf.User][ctx.Req.Host] = resp.Cookies()
				Logger.Printf("Cookies: %+v\n", cookies)
				cookiesL.Unlock()
			}
			Logger.Println("Resource "+ctx.Req.URL.String(),
				"requires authentication (HTTP 401).", solutionMsg,
				"resulted in HTTP", resp.StatusCode)
			return resp
		}
		Logger.Println("Received data with status HTTP", r.StatusCode)
		return r
	})
}

func ProxyHandler(w http.ResponseWriter, req *http.Request) {
	toProxy := req.FormValue("uri")
	if len(toProxy) == 0 {
		w.WriteHeader(500)
		w.Write([]byte("No URI was provided to the proxy!"))
		return
	}

	uri, err := url.ParseRequestURI(toProxy)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("Error parsing URL: " + req.URL.String() + " " + err.Error()))
		Logger.Println("Error parsing URL:", req.URL, err.Error())
		return
	}
	req.URL = uri
	req.Host = uri.Host
	req.RequestURI = uri.RequestURI()
	Logger.Println("Proxying request for:", req.URL)
	proxy.ServeHTTP(w, req)
	return
}

func TimeTrack(start time.Time, name string) {
	elapsed := time.Since(start)
	Logger.Printf("%s finished in %s", name, elapsed)
}
<|endoftext|>


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

package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/kelseyhightower/envconfig"
)

const version = "0.1.1"

var printVersion bool

func init() {
	flag.BoolVar(&printVersion, "v", false, "print version and exit")
	flag.Parse()
}

type config struct {
	BindAddress    string `envconfig:"BIND_ADDRESS" default:":9000"`
	HipacheAddress string `envconfig:"HIPACHE_ADDRESS" required:"true"`
}

func main() {
	if printVersion {
		fmt.Printf("hipache-healthcheck-proxy %s\n", version)
		return
	}
	var c config
	err := envconfig.Process("", &c)
	if err != nil {
		log.Fatal(err)
	}
	if _, err = url.Parse(c.HipacheAddress); err != nil {
		log.Fatalf("failed to parse hipache address: %s", err)
	}
	client := http.Client{
		Transport: &http.Transport{
			DialContext: (&net.Dialer{Timeout: time.Second}).DialContext,
		},
		Timeout: 2 * time.Second,
	}

	log.Printf("starting on %s...", c.BindAddress)
	err = http.ListenAndServe(c.BindAddress, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		req, err := http.NewRequest("GET", c.HipacheAddress, nil)
		if err != nil {
			http.Error(w, "failed to process request", http.StatusInternalServerError)
			return
		}
		req.Host = "__ping__"
		resp, err := client.Do(req)
		if err != nil {
			http.Error(w, err.Error(), http.StatusServiceUnavailable)
			return
		}
		defer resp.Body.Close()
		for h := range resp.Header {
			w.Header().Set(h, resp.Header.Get(h))
		}
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	}))
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Bump to 0.1.2<commit_after>// Copyright 2016 EF CTX. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/kelseyhightower/envconfig"
)

const version = "0.1.2"

var printVersion bool

func init() {
	flag.BoolVar(&printVersion, "v", false, "print version and exit")
	flag.Parse()
}

type config struct {
	BindAddress    string `envconfig:"BIND_ADDRESS" default:":9000"`
	HipacheAddress string `envconfig:"HIPACHE_ADDRESS" required:"true"`
}

func main() {
	if printVersion {
		fmt.Printf("hipache-healthcheck-proxy %s\n", version)
		return
	}
	var c config
	err := envconfig.Process("", &c)
	if err != nil {
		log.Fatal(err)
	}
	if _, err = url.Parse(c.HipacheAddress); err != nil {
		log.Fatalf("failed to parse hipache address: %s", err)
	}
	client := http.Client{
		Transport: &http.Transport{
			DialContext: (&net.Dialer{Timeout: time.Second}).DialContext,
		},
		Timeout: 2 * time.Second,
	}

	log.Printf("starting on %s...", c.BindAddress)
	err = http.ListenAndServe(c.BindAddress, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		req, err := http.NewRequest("GET", c.HipacheAddress, nil)
		if err != nil {
			http.Error(w, "failed to process request", http.StatusInternalServerError)
			return
		}
		req.Host = "__ping__"
		resp, err := client.Do(req)
		if err != nil {
			http.Error(w, err.Error(), http.StatusServiceUnavailable)
			return
		}
		defer resp.Body.Close()
		for h := range resp.Header {
			w.Header().Set(h, resp.Header.Get(h))
		}
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	}))
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package apiproxy

import (
	"github.com/sourcegraph/httpcache"
	"net/http"
	"net/http/httputil"
	"net/url"
)

// NewCachingSingleHostReverseProxy constructs a caching reverse proxy handler for
// target. If cache is nil, a volatile, in-memory cache is used.
func NewCachingSingleHostReverseProxy(target *url.URL, cache httpcache.Cache) *httputil.ReverseProxy {
	proxy := NewSingleHostReverseProxy(target)
	if cache == nil {
		cache = httpcache.NewMemoryCache()
	}
	proxy.Transport = httpcache.NewTransport(cache)
	return proxy
}

// NewSingleHostReverseProxy wraps net/http/httputil.NewSingleHostReverseProxy
// and sets the Host header.
func NewSingleHostReverseProxy(url *url.URL) *httputil.ReverseProxy {
	proxy := httputil.NewSingleHostReverseProxy(url)
	oldDirector := proxy.Director
	proxy.Director = func(r *http.Request) {
		oldDirector(r)
		r.Host = url.Host
	}
	return proxy
}
<commit_msg>Clarify docs<commit_after>package apiproxy

import (
	"github.com/sourcegraph/httpcache"
	"net/http"
	"net/http/httputil"
	"net/url"
)

// NewCachingSingleHostReverseProxy constructs a caching reverse proxy handler for
// target. If cache is nil, a volatile, in-memory cache is used.
func NewCachingSingleHostReverseProxy(target *url.URL, cache httpcache.Cache) *httputil.ReverseProxy {
	proxy := NewSingleHostReverseProxy(target)
	if cache == nil {
		cache = httpcache.NewMemoryCache()
	}
	proxy.Transport = httpcache.NewTransport(cache)
	return proxy
}

// NewSingleHostReverseProxy wraps net/http/httputil.NewSingleHostReverseProxy
// and sets the Host header based on the target URL.
func NewSingleHostReverseProxy(url *url.URL) *httputil.ReverseProxy {
	proxy := httputil.NewSingleHostReverseProxy(url)
	oldDirector := proxy.Director
	proxy.Director = func(r *http.Request) {
		oldDirector(r)
		r.Host = url.Host
	}
	return proxy
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/astaxie/beego"
	//"runtime"
)

type MessageStruct struct {
	Message string
}

type JsonController struct {
	beego.Controller
}

func (this *JsonController) Get() {
	m := MessageStruct{"Hello, world"}
	this.Data["json"] = &m
	this.ServeJson()
}

func main() {
	//don't need this set, beego default set it
	//runtime.GOMAXPROCS(runtime.NumCPU())
	beego.RunMode = "prod"
	beego.Router("/json", &JsonController{})
	beego.Run()
}
<commit_msg>verify Beego json fix<commit_after>package main

import (
	"github.com/astaxie/beego"
	//"runtime"
)

type MessageStruct struct {
	message string
}

type JsonController struct {
	beego.Controller
}

func (this *JsonController) Get() {
	m := MessageStruct{"Hello, World!"}
	this.Data["json"] = &m
	this.ServeJson()
}

func main() {
	//don't need this set, beego default set it
	//runtime.GOMAXPROCS(runtime.NumCPU())
	beego.RunMode = "prod"
	beego.Router("/json", &JsonController{})
	beego.Run()
}
<|endoftext|>


<commit_before>package channels

// BatchingChannel implements the Channel interface, with the change that instead of producing individual elements
// on Out(), it batches together the entire internal buffer each time. Trying to construct an unbuffered batching channel
// will panic, that configuration is not supported (and provides no benefit over an unbuffered NativeChannel).
type BatchingChannel struct {
	input, output chan interface{}
	buffer        []interface{}
	size          BufferCap
}

func NewBatchingChannel(size BufferCap) *BatchingChannel {
	if size == None {
		panic("channels: BatchingChannel does not support unbuffered behaviour")
	}
	if size < 0 && size != Infinity {
		panic("channels: invalid negative size in NewBatchingChannel")
	}
	ch := &BatchingChannel{make(chan interface{}), make(chan interface{}), nil, size}
	go ch.batchingBuffer()
	return ch
}

func (ch *BatchingChannel) In() chan<- interface{} {
	return ch.input
}

func (ch *BatchingChannel) Out() <-chan interface{} {
	return ch.output
}

func (ch *BatchingChannel) Len() int {
	return len(ch.buffer)
}

func (ch *BatchingChannel) Cap() BufferCap {
	return ch.size
}

func (ch *BatchingChannel) Close() {
	close(ch.input)
}

func (ch *BatchingChannel) shutdown() {
	ch.output <- ch.buffer
	close(ch.output)
}

func (ch *BatchingChannel) batchingBuffer() {
	for {
		if len(ch.buffer) == 0 {
			elem, open := <-ch.input
			if open {
				ch.buffer = append(ch.buffer, elem)
			} else {
				ch.shutdown()
				return
			}
		} else if ch.size != Infinity && len(ch.buffer) >= int(ch.size) {
			ch.output <- ch.buffer
			ch.buffer = nil
		} else {
			select {
			case elem, open := <-ch.input:
				if open {
					ch.buffer = append(ch.buffer, elem)
				} else {
					ch.shutdown()
					return
				}
			case ch.output <- ch.buffer:
				ch.buffer = nil
			}
		}
	}
}
<commit_msg>Document BatchingChannel.Out<commit_after>package channels

// BatchingChannel implements the Channel interface, with the change that instead of producing individual elements
// on Out(), it batches together the entire internal buffer each time. Trying to construct an unbuffered batching channel
// will panic, that configuration is not supported (and provides no benefit over an unbuffered NativeChannel).
type BatchingChannel struct {
	input, output chan interface{}
	buffer        []interface{}
	size          BufferCap
}

func NewBatchingChannel(size BufferCap) *BatchingChannel {
	if size == None {
		panic("channels: BatchingChannel does not support unbuffered behaviour")
	}
	if size < 0 && size != Infinity {
		panic("channels: invalid negative size in NewBatchingChannel")
	}
	ch := &BatchingChannel{make(chan interface{}), make(chan interface{}), nil, size}
	go ch.batchingBuffer()
	return ch
}

func (ch *BatchingChannel) In() chan<- interface{} {
	return ch.input
}

// Out returns a <-chan interface{} in order that BatchingChannel conforms to the standard Channel interface provided
// by this package, however each output value is guaranteed to be of type []interface{} - a slice collecting the most
// recent batch of values sent on the In channel. The slice is guaranteed to not be empty or nil. In practice the net
// result is that you need an additional type assertion to access the underlying values.
func (ch *BatchingChannel) Out() <-chan interface{} {
	return ch.output
}

func (ch *BatchingChannel) Len() int {
	return len(ch.buffer)
}

func (ch *BatchingChannel) Cap() BufferCap {
	return ch.size
}

func (ch *BatchingChannel) Close() {
	close(ch.input)
}

func (ch *BatchingChannel) shutdown() {
	ch.output <- ch.buffer
	close(ch.output)
}

func (ch *BatchingChannel) batchingBuffer() {
	for {
		if len(ch.buffer) == 0 {
			elem, open := <-ch.input
			if open {
				ch.buffer = append(ch.buffer, elem)
			} else {
				ch.shutdown()
				return
			}
		} else if ch.size != Infinity && len(ch.buffer) >= int(ch.size) {
			ch.output <- ch.buffer
			ch.buffer = nil
		} else {
			select {
			case elem, open := <-ch.input:
				if open {
					ch.buffer = append(ch.buffer, elem)
				} else {
					ch.shutdown()
					return
				}
			case ch.output <- ch.buffer:
				ch.buffer = nil
			}
		}
	}
}
<|endoftext|>


<commit_before>/*
 *    Copyright (C) 2014 Christian Muehlhaeuser
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors:
 *      Christian Muehlhaeuser <muesli@gmail.com>
 */

package bees

// Modules provide events, which are described in an EventDescriptor.
type EventDescriptor struct {
	Namespace   string
	Name        string
	Description string
	Options     []PlaceholderDescriptor
}

// Modules offer actions, which are described in an ActionDescriptor.
type ActionDescriptor struct {
	Namespace   string
	Name        string
	Description string
	Options     []PlaceholderDescriptor
}

// A PlaceholderDescriptor shows which in & out values a module expects and returns.
type PlaceholderDescriptor struct {
	Name        string
	Description string
	Type        string
}

// A BeeOptionDescriptor shows which config values a module expects.
type BeeOptionDescriptor struct {
	Name        string
	Description string
	Type        string
	Mandatory   bool
}

// Returns the ActionDescriptor matching an action.
func GetActionDescriptor(action *Action) ActionDescriptor {
	factory := (*GetFactory((*GetBee(action.Bee)).Namespace()))
	for _, ac := range factory.Actions() {
		if ac.Name == action.Name {
			return ac
		}
	}

	return ActionDescriptor{}
}

// Returns the EventDescriptor matching an event.
func GetEventDescriptor(event *Event) EventDescriptor {
	factory := (*GetFactory((*GetBee(event.Bee)).Namespace()))
	for _, ev := range factory.Events() {
		if ev.Name == event.Name {
			return ev
		}
	}

	return EventDescriptor{}
}
<commit_msg>Panic if the bee isn't found<commit_after>/*
 *    Copyright (C) 2014 Christian Muehlhaeuser
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors:
 *      Christian Muehlhaeuser <muesli@gmail.com>
 */

package bees

// Modules provide events, which are described in an EventDescriptor.
type EventDescriptor struct {
	Namespace   string
	Name        string
	Description string
	Options     []PlaceholderDescriptor
}

// Modules offer actions, which are described in an ActionDescriptor.
type ActionDescriptor struct {
	Namespace   string
	Name        string
	Description string
	Options     []PlaceholderDescriptor
}

// A PlaceholderDescriptor shows which in & out values a module expects and returns.
type PlaceholderDescriptor struct {
	Name        string
	Description string
	Type        string
}

// A BeeOptionDescriptor shows which config values a module expects.
type BeeOptionDescriptor struct {
	Name        string
	Description string
	Type        string
	Mandatory   bool
}

// Returns the ActionDescriptor matching an action.
func GetActionDescriptor(action *Action) ActionDescriptor {
	bee := GetBee(action.Bee)
	if bee == nil {
		panic("Bee " + action.Bee + " not registered")
	}
	factory := (*GetFactory((*bee).Namespace()))
	for _, ac := range factory.Actions() {
		if ac.Name == action.Name {
			return ac
		}
	}

	return ActionDescriptor{}
}

// Returns the EventDescriptor matching an event.
func GetEventDescriptor(event *Event) EventDescriptor {
	bee := GetBee(event.Bee)
	if bee == nil {
		panic("Bee " + event.Bee + " not registered")
	}
	factory := (*GetFactory((*bee).Namespace()))
	for _, ev := range factory.Events() {
		if ev.Name == event.Name {
			return ev
		}
	}

	return EventDescriptor{}
}
<|endoftext|>


<commit_before>package relay

import (
	"log"
	"sync"
	"time"

	"github.com/armon/relay/broker"
)

// retryBroker implements a resilient AMQP broker which automatically replaces
// the underlying channel if an unexpected error is observed. Due to AMQP's
// persistent connections, this becomes necessary to survive even minor network
// hiccups which may temporarily sever the connection to the server.
//
// The retryBroker does not come without a price, however. Clients are still
// responsible for deduplicating messages or handling them idempotently, as
// duplications are a regular happening in the face of network errors due to
// message redelivery provided by the protocol. It is also important to note
// that acknowledgements - both ack and nack - cannot be retried across
// connection resets. This stems from the fact that a given consumer's channel
// cannot be revived and must be thrown away if the consumer is to be reused.
type retryBroker struct {
	broker   *relayBroker
	attempts int
	min      time.Duration
	max      time.Duration
}

// RetryBroker returns a new retrying broker with the given settings.
func (r *Relay) RetryBroker(attempts int, min, max time.Duration) *retryBroker {
	return &retryBroker{&relayBroker{r}, attempts, min, max}
}

// Publisher returns a new retrying broker.Publisher.
func (rb *retryBroker) Publisher(queue string) (broker.Publisher, error) {
	return &retryPublisher{
		broker:   rb.broker,
		queue:    queue,
		attempts: rb.attempts,
		min:      rb.min,
		max:      rb.max,
	}, nil
}

// Consumer returns a new retrying broker.Consumer.
func (rb *retryBroker) Consumer(queue string) (broker.Consumer, error) {
	return &retryConsumer{
		broker:   rb.broker,
		queue:    queue,
		attempts: rb.attempts,
		min:      rb.min,
		max:      rb.max,
	}, nil
}

// Close closes the broker.
func (rb *retryBroker) Close() error {
	if rb.broker != nil {
		return rb.broker.Close()
	}
	return nil
}

// retryConsumer tracks the relay broker and an associated consumer, allowing
// the consumer to be replaced transparently in the face of failure.
type retryConsumer struct {
	broker   broker.Broker
	cons     broker.Consumer
	queue    string
	attempts int
	min      time.Duration
	max      time.Duration
	l        sync.RWMutex
}

// consumer is used to connect the consumer to the relay queue.
func (rc *retryConsumer) consumer(create bool) (broker.Consumer, error) {
	// Check for an existing consumer
	rc.l.RLock()
	cons := rc.cons
	rc.l.RUnlock()

	if cons != nil || !create {
		return cons, nil
	}

	// Make a new consumer
	rc.l.Lock()
	defer rc.l.Unlock()

	cons, err := rc.broker.Consumer(rc.queue)
	if err != nil {
		return nil, err
	}
	rc.cons = cons
	return cons, nil
}

// discard is used to remove a known-bad consumer.
func (rc *retryConsumer) discard(cons broker.Consumer) {
	if cons == nil {
		return
	}
	cons.Close()

	rc.l.Lock()
	defer rc.l.Unlock()
	if cons == rc.cons {
		rc.cons = nil
	}
}

// Close closes the consumer.
func (rc *retryConsumer) Close() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Close()
}

// Ack marks message(s) as delivered.
func (rc *retryConsumer) Ack() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Ack()
}

// Nack sends message(s) back to the queue.
func (rc *retryConsumer) Nack() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Nack()
}

// ConsumeAck is used to consume with automatic ack.
func (rc *retryConsumer) ConsumeAck(out interface{}) error {
	if err := rc.Consume(out); err != nil {
		return err
	}
	if err := rc.Ack(); err != nil {
		return err
	}
	return nil
}

// Consume consumes a single message from the queue.
func (rc *retryConsumer) Consume(out interface{}) error {
	return rc.ConsumeTimeout(out, -1)
}

// ConsumeTimeout consumes a single message from the queue with an upper bound
// on the time spent waiting. This behaves slightly different in the retry
// broker, as we may encounter errors getting a valid consumer. In this case,
// the timeout may actually be longer than specified to allow reconnection
// attempts to take place.
func (rc *retryConsumer) ConsumeTimeout(out interface{}, timeout time.Duration) error {
	wait := rc.min
	for i := 0; ; i++ {
		cons, err := rc.consumer(true)
		if err != nil {
			goto RETRY
		}

		err = cons.ConsumeTimeout(out, timeout)
		if err == nil {
			break
		}

		// Respect the case where the consume legitimately times out.
		if err == TimedOut {
			return err
		}

	RETRY:
		log.Printf("[ERR] relay: consumer got error: %v", err)
		rc.discard(cons)
		if i == rc.attempts {
			log.Printf("[ERR] relay: consumer giving up after %d attempts", i)
			return err
		}
		log.Printf("[DEBUG] relay: consumer retrying in %s", wait)
		time.Sleep(wait)
		wait *= 2
		if wait > rc.max {
			wait = rc.max
		}
	}
	return nil
}

// retryPublisher wraps a relay broker and an associated publisher, allowing
// the underlying publisher's channel to be automatically swapped out if any
// errors are seen during a publish.
type retryPublisher struct {
	broker   broker.Broker
	pub      broker.Publisher
	queue    string
	attempts int
	min      time.Duration
	max      time.Duration
	l        sync.RWMutex
}

// publisher is used to connect the publisher to the relay queue.
func (rp *retryPublisher) publisher(create bool) (broker.Publisher, error) {
	// Check for an existing publisher
	rp.l.RLock()
	pub := rp.pub
	rp.l.RUnlock()

	if pub != nil || !create {
		return pub, nil
	}

	// Make a new publisher
	rp.l.Lock()
	defer rp.l.Unlock()

	pub, err := rp.broker.Publisher(rp.queue)
	if err != nil {
		return nil, err
	}
	rp.pub = pub
	return pub, nil
}

// discard is used to remove a known-bad publisher.
func (rp *retryPublisher) discard(pub broker.Publisher) {
	if pub == nil {
		return
	}
	pub.Close()

	rp.l.Lock()
	defer rp.l.Unlock()
	if pub == rp.pub {
		rp.pub = nil
	}
}

// Close closes the publisher.
func (rp *retryPublisher) Close() error {
	pub, err := rp.publisher(false)
	if err != nil || pub == nil {
		return err
	}
	return pub.Close()
}

// Publish publishes a single message to the queue. If an error is encountered,
// the broker automatically tries to replace the underlying channel.
func (rp *retryPublisher) Publish(in interface{}) error {
	wait := rp.min
	for i := 0; ; i++ {
		pub, err := rp.publisher(true)
		if err != nil {
			goto RETRY
		}

		err = pub.Publish(in)
		if err == nil {
			break
		}

	RETRY:
		log.Printf("[ERR] relay: publisher got error: %v", err)
		rp.discard(pub)
		if i == rp.attempts {
			log.Printf("[ERR] relay: publisher giving up after %d attempts", i)
			return err
		}
		log.Printf("[DEBUG] relay: publisher retrying in %s", wait)
		time.Sleep(wait)
		wait *= 2
		if wait > rp.max {
			wait = rp.max
		}
	}
	return nil
}
<commit_msg>Simplify retry backoff<commit_after>package relay

import (
	"log"
	"sync"
	"time"

	"github.com/armon/relay/broker"
)

// retryBroker implements a resilient AMQP broker which automatically replaces
// the underlying channel if an unexpected error is observed. Due to AMQP's
// persistent connections, this becomes necessary to survive even minor network
// hiccups which may temporarily sever the connection to the server.
//
// The retryBroker does not come without a price, however. Clients are still
// responsible for deduplicating messages or handling them idempotently, as
// duplications are a regular happening in the face of network errors due to
// message redelivery provided by the protocol. It is also important to note
// that acknowledgements - both ack and nack - cannot be retried across
// connection resets. This stems from the fact that a given consumer's channel
// cannot be revived and must be thrown away if the consumer is to be reused.
type retryBroker struct {
	broker   *relayBroker
	attempts int
	min      time.Duration
	max      time.Duration
}

// RetryBroker returns a new retrying broker with the given settings.
func (r *Relay) RetryBroker(attempts int, min, max time.Duration) *retryBroker {
	return &retryBroker{&relayBroker{r}, attempts, min, max}
}

// Publisher returns a new retrying broker.Publisher.
func (rb *retryBroker) Publisher(queue string) (broker.Publisher, error) {
	return &retryPublisher{
		broker:   rb.broker,
		queue:    queue,
		attempts: rb.attempts,
		min:      rb.min,
		max:      rb.max,
	}, nil
}

// Consumer returns a new retrying broker.Consumer.
func (rb *retryBroker) Consumer(queue string) (broker.Consumer, error) {
	return &retryConsumer{
		broker:   rb.broker,
		queue:    queue,
		attempts: rb.attempts,
		min:      rb.min,
		max:      rb.max,
	}, nil
}

// Close closes the broker.
func (rb *retryBroker) Close() error {
	if rb.broker != nil {
		return rb.broker.Close()
	}
	return nil
}

// retryConsumer tracks the relay broker and an associated consumer, allowing
// the consumer to be replaced transparently in the face of failure.
type retryConsumer struct {
	broker   broker.Broker
	cons     broker.Consumer
	queue    string
	attempts int
	min      time.Duration
	max      time.Duration
	l        sync.RWMutex
}

// consumer is used to connect the consumer to the relay queue.
func (rc *retryConsumer) consumer(create bool) (broker.Consumer, error) {
	// Check for an existing consumer
	rc.l.RLock()
	cons := rc.cons
	rc.l.RUnlock()

	if cons != nil || !create {
		return cons, nil
	}

	// Make a new consumer
	rc.l.Lock()
	defer rc.l.Unlock()

	cons, err := rc.broker.Consumer(rc.queue)
	if err != nil {
		return nil, err
	}
	rc.cons = cons
	return cons, nil
}

// discard is used to remove a known-bad consumer.
func (rc *retryConsumer) discard(cons broker.Consumer) {
	if cons == nil {
		return
	}
	cons.Close()

	rc.l.Lock()
	defer rc.l.Unlock()
	if cons == rc.cons {
		rc.cons = nil
	}
}

// Close closes the consumer.
func (rc *retryConsumer) Close() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Close()
}

// Ack marks message(s) as delivered.
func (rc *retryConsumer) Ack() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Ack()
}

// Nack sends message(s) back to the queue.
func (rc *retryConsumer) Nack() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Nack()
}

// ConsumeAck is used to consume with automatic ack.
func (rc *retryConsumer) ConsumeAck(out interface{}) error {
	if err := rc.Consume(out); err != nil {
		return err
	}
	if err := rc.Ack(); err != nil {
		return err
	}
	return nil
}

// Consume consumes a single message from the queue.
func (rc *retryConsumer) Consume(out interface{}) error {
	return rc.ConsumeTimeout(out, -1)
}

// ConsumeTimeout consumes a single message from the queue with an upper bound
// on the time spent waiting. This behaves slightly different in the retry
// broker, as we may encounter errors getting a valid consumer. In this case,
// the timeout may actually be longer than specified to allow reconnection
// attempts to take place.
func (rc *retryConsumer) ConsumeTimeout(out interface{}, timeout time.Duration) error {
	for i := 0; ; i++ {
		cons, err := rc.consumer(true)
		if err != nil {
			goto RETRY
		}

		err = cons.ConsumeTimeout(out, timeout)
		if err == nil {
			break
		}

		// Respect the case where the consume legitimately times out.
		if err == TimedOut {
			return err
		}

	RETRY:
		log.Printf("[ERR] relay: consumer got error: %v", err)
		rc.discard(cons)
		if i == rc.attempts {
			log.Printf("[ERR] relay: consumer giving up after %d attempts", i)
			return err
		}
		wait := rc.min * (1 << uint(i))
		if wait > rc.max {
			wait = rc.max
		}
		log.Printf("[DEBUG] relay: consumer retrying in %s", wait)
		time.Sleep(wait)
	}
	return nil
}

// retryPublisher wraps a relay broker and an associated publisher, allowing
// the underlying publisher's channel to be automatically swapped out if any
// errors are seen during a publish.
type retryPublisher struct {
	broker   broker.Broker
	pub      broker.Publisher
	queue    string
	attempts int
	min      time.Duration
	max      time.Duration
	l        sync.RWMutex
}

// publisher is used to connect the publisher to the relay queue.
func (rp *retryPublisher) publisher(create bool) (broker.Publisher, error) {
	// Check for an existing publisher
	rp.l.RLock()
	pub := rp.pub
	rp.l.RUnlock()

	if pub != nil || !create {
		return pub, nil
	}

	// Make a new publisher
	rp.l.Lock()
	defer rp.l.Unlock()

	pub, err := rp.broker.Publisher(rp.queue)
	if err != nil {
		return nil, err
	}
	rp.pub = pub
	return pub, nil
}

// discard is used to remove a known-bad publisher.
func (rp *retryPublisher) discard(pub broker.Publisher) {
	if pub == nil {
		return
	}
	pub.Close()

	rp.l.Lock()
	defer rp.l.Unlock()
	if pub == rp.pub {
		rp.pub = nil
	}
}

// Close closes the publisher.
func (rp *retryPublisher) Close() error {
	pub, err := rp.publisher(false)
	if err != nil || pub == nil {
		return err
	}
	return pub.Close()
}

// Publish publishes a single message to the queue. If an error is encountered,
// the broker automatically tries to replace the underlying channel.
func (rp *retryPublisher) Publish(in interface{}) error {
	for i := 0; ; i++ {
		pub, err := rp.publisher(true)
		if err != nil {
			goto RETRY
		}

		err = pub.Publish(in)
		if err == nil {
			break
		}

	RETRY:
		log.Printf("[ERR] relay: publisher got error: %v", err)
		rp.discard(pub)
		if i == rp.attempts {
			log.Printf("[ERR] relay: publisher giving up after %d attempts", i)
			return err
		}
		wait := rp.min * (1 << uint(i))
		if wait > rp.max {
			wait = rp.max
		}
		log.Printf("[DEBUG] relay: publisher retrying in %s", wait)
		time.Sleep(wait)
	}
	return nil
}
<|endoftext|>


<commit_before>package stdbus

import (
	"sync"

	"github.com/CyCoreSystems/ari"
)

// subscriptionEventBufferSize defines the number of events that each
// subscription will queue before accepting more events.
var subscriptionEventBufferSize = 100

// bus is an event bus for ARI events.  It receives and
// redistributes events based on a subscription
// model.
type bus struct {
	subs []*subscription // The list of subscriptions

	mu sync.Mutex

	closed bool
}

// New creates and returns the event bus.
func New() ari.Bus {
	b := &bus{
		subs: []*subscription{},
	}

	return b
}

// Close closes out all subscriptions in the bus.
func (b *bus) Close() {
	if b.closed {
		return
	}
	b.closed = true

	for _, s := range b.subs {
		s.Cancel()
	}
}

// Send sends the message to the bus
func (b *bus) Send(e ari.Event) {
	// Disseminate the message to the subscribers
	for _, s := range b.subs {
		for _, k := range e.Keys() {
			if s.key.Match(k) {
				for _, topic := range s.events {
					if topic == e.GetType() || topic == ari.Events.All {
						select {
						case s.C <- e:
						default: // never block
						}
					}
				}
			}
		}
	}
}

// Subscribe returns a subscription to the given list
// of event types
func (b *bus) Subscribe(key *ari.Key, eTypes ...string) ari.Subscription {
	s := newSubscription(b, key, eTypes...)
	b.add(s)
	return s
}

// add appends a new subscription to the bus
func (b *bus) add(s *subscription) {
	b.mu.Lock()
	b.subs = append(b.subs, s)
	b.mu.Unlock()
}

// remove deletes the given subscription from the bus
func (b *bus) remove(s *subscription) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for i, si := range b.subs {
		if s == si {
			// Subs are pointers, so we have to explicitly remove them
			// to prevent memory leaks
			b.subs[i] = b.subs[len(b.subs)-1] // replace the current with the end
			b.subs[len(b.subs)-1] = nil       // remove the end
			b.subs = b.subs[:len(b.subs)-1]   // lop off the end
			return
		}
	}
}

// A Subscription is a wrapped channel for receiving
// events from the ARI event bus.
type subscription struct {
	key    *ari.Key
	b      *bus     // reference to the event bus
	events []string // list of events to listen for

	closed bool           // channel closure protection flag
	C      chan ari.Event // channel for sending events to the subscriber
}

// newSubscription creates a new, unattached subscription
func newSubscription(b *bus, key *ari.Key, eTypes ...string) *subscription {
	return &subscription{
		key:    key,
		b:      b,
		events: eTypes,
		C:      make(chan ari.Event, subscriptionEventBufferSize),
	}
}

// Events returns the events channel
func (s *subscription) Events() <-chan ari.Event {
	return s.C
}

// Cancel cancels the subscription and removes it from
// the event bus.
func (s *subscription) Cancel() {
	if s == nil {
		return
	}

	// Remove the subscription from the bus
	if s.b != nil {
		s.b.remove(s)
	}

	if s.closed {
		return
	}

	// Close the subscription's deliver channel
	if s.C != nil {
		s.closed = true
		close(s.C)
	}
}
<commit_msg>bus: one event per matched subscription:w<commit_after>package stdbus

import (
	"sync"

	"github.com/CyCoreSystems/ari"
)

// subscriptionEventBufferSize defines the number of events that each
// subscription will queue before accepting more events.
var subscriptionEventBufferSize = 100

// bus is an event bus for ARI events.  It receives and
// redistributes events based on a subscription
// model.
type bus struct {
	subs []*subscription // The list of subscriptions

	mu sync.Mutex

	closed bool
}

// New creates and returns the event bus.
func New() ari.Bus {
	b := &bus{
		subs: []*subscription{},
	}

	return b
}

// Close closes out all subscriptions in the bus.
func (b *bus) Close() {
	if b.closed {
		return
	}
	b.closed = true

	for _, s := range b.subs {
		s.Cancel()
	}
}

// Send sends the message to the bus
func (b *bus) Send(e ari.Event) {
	var matched bool

	// Disseminate the message to the subscribers
	for _, s := range b.subs {
		matched = false
		for _, k := range e.Keys() {
			if matched {
				break
			}
			if s.key.Match(k) {
				matched = true
				for _, topic := range s.events {
					if topic == e.GetType() || topic == ari.Events.All {
						select {
						case s.C <- e:
						default: // never block
						}
					}
				}
			}
		}
	}
}

// Subscribe returns a subscription to the given list
// of event types
func (b *bus) Subscribe(key *ari.Key, eTypes ...string) ari.Subscription {
	s := newSubscription(b, key, eTypes...)
	b.add(s)
	return s
}

// add appends a new subscription to the bus
func (b *bus) add(s *subscription) {
	b.mu.Lock()
	b.subs = append(b.subs, s)
	b.mu.Unlock()
}

// remove deletes the given subscription from the bus
func (b *bus) remove(s *subscription) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for i, si := range b.subs {
		if s == si {
			// Subs are pointers, so we have to explicitly remove them
			// to prevent memory leaks
			b.subs[i] = b.subs[len(b.subs)-1] // replace the current with the end
			b.subs[len(b.subs)-1] = nil       // remove the end
			b.subs = b.subs[:len(b.subs)-1]   // lop off the end
			return
		}
	}
}

// A Subscription is a wrapped channel for receiving
// events from the ARI event bus.
type subscription struct {
	key    *ari.Key
	b      *bus     // reference to the event bus
	events []string // list of events to listen for

	closed bool           // channel closure protection flag
	C      chan ari.Event // channel for sending events to the subscriber
}

// newSubscription creates a new, unattached subscription
func newSubscription(b *bus, key *ari.Key, eTypes ...string) *subscription {
	return &subscription{
		key:    key,
		b:      b,
		events: eTypes,
		C:      make(chan ari.Event, subscriptionEventBufferSize),
	}
}

// Events returns the events channel
func (s *subscription) Events() <-chan ari.Event {
	return s.C
}

// Cancel cancels the subscription and removes it from
// the event bus.
func (s *subscription) Cancel() {
	if s == nil {
		return
	}

	// Remove the subscription from the bus
	if s.b != nil {
		s.b.remove(s)
	}

	if s.closed {
		return
	}

	// Close the subscription's deliver channel
	if s.C != nil {
		s.closed = true
		close(s.C)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"path"
)

// #cgo LDFLAGS: -lavformat -lavutil
// #include "libavutil/dict.h"
// #include "libavformat/avformat.h"
import "C"

func readTags(name string) (tags map[string]string, err error) {
	var fmtCtx *C.AVFormatContext
	if C.avformat_open_input(&fmtCtx, C.CString(name), nil, nil) != 0 {
		return nil, errors.New("C.avformat_open_input")
	}
	emptyCString := C.CString("")
	tags = make(map[string]string)
	for i := C.av_dict_get(fmtCtx.metadata, emptyCString, nil, C.AV_DICT_IGNORE_SUFFIX);
		i != nil;
		i = C.av_dict_get(fmtCtx.metadata, emptyCString, i, C.AV_DICT_IGNORE_SUFFIX) {
		tags[C.GoString(i.key)] = C.GoString(i.value)
	}
	C.avformat_close_input(&fmtCtx)
	return
}

var (
	library = flag.String("library", "", "the path to the library directory")
	port    = flag.Uint("port", 8080, "the port to listen on")
)

func handler(w http.ResponseWriter, r *http.Request) {
	fis, err := ioutil.ReadDir(*library)
	if err != nil {
		fmt.Fprintln(w, "Error reading files")
		return
	}
	for _, fi := range fis {
		name := fi.Name()
		tags, err := readTags(path.Join(*library, name))
		if err != nil {
			fmt.Fprintf(w, "%s [%s]\n", name, err)
		} else {
			fmt.Fprintln(w, tags)
		}
	}
}

func main() {
	flag.Parse()
	if *library == "" {
		log.Fatal("no library path specified")
	}
	C.av_register_all()
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", *port), nil))
}
<commit_msg>Use absolute include path<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"path"
)

// #cgo LDFLAGS: -lavformat -lavutil
// #include <libavutil/dict.h>
// #include <libavformat/avformat.h>
import "C"

func readTags(name string) (tags map[string]string, err error) {
	var fmtCtx *C.AVFormatContext
	if C.avformat_open_input(&fmtCtx, C.CString(name), nil, nil) != 0 {
		return nil, errors.New("C.avformat_open_input")
	}
	emptyCString := C.CString("")
	tags = make(map[string]string)
	for i := C.av_dict_get(fmtCtx.metadata, emptyCString, nil, C.AV_DICT_IGNORE_SUFFIX);
		i != nil;
		i = C.av_dict_get(fmtCtx.metadata, emptyCString, i, C.AV_DICT_IGNORE_SUFFIX) {
		tags[C.GoString(i.key)] = C.GoString(i.value)
	}
	C.avformat_close_input(&fmtCtx)
	return
}

var (
	library = flag.String("library", "", "the path to the library directory")
	port    = flag.Uint("port", 8080, "the port to listen on")
)

func handler(w http.ResponseWriter, r *http.Request) {
	fis, err := ioutil.ReadDir(*library)
	if err != nil {
		fmt.Fprintln(w, "Error reading files")
		return
	}
	for _, fi := range fis {
		name := fi.Name()
		tags, err := readTags(path.Join(*library, name))
		if err != nil {
			fmt.Fprintf(w, "%s [%s]\n", name, err)
		} else {
			fmt.Fprintln(w, tags)
		}
	}
}

func main() {
	flag.Parse()
	if *library == "" {
		log.Fatal("no library path specified")
	}
	C.av_register_all()
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", *port), nil))
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"mime"
	"os"
	"path/filepath"
	"strings"
)

var (
	doBasename bool
	doFullpath bool
	targets    []string
)

func init() {
	chkHelp()
	flags()
}

// Check whether user requested help.
func chkHelp() {
	if len(os.Args) < 2 {
		return
	}

	switch os.Args[1] {
	case "-h", "h", "help", "--help", "-H", "H", "HELP", "--HELP", "-help", "--h", "--H":
		help(0)
	}

}

// Print help and exit with a status code.
func help(status int) {
	defer os.Exit(status)
	fmt.Printf(
		"%s\n\n  %s\n\n  %s\n%s\n\n  %s\n%s\n%s\n%s\n%s\n\n  %s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n",
		"rmext",
		"Usage: rmext file...",
		"Description:",
		"    Print the basename of files WITHOUT their extension.",
		"Options:",
		"   -b         Print the basename of files without their path",
		"   -f         Print the full path of files",
		"   (default)  Print the path of files exactly as they were entered",
		"                (minus any file extensions)",
		"Notes:",
		"    This program relies on the system's list of known mimetypes.",
		"	 On Linux, usually the file '/etc/mime.types' stores all",
		"      MIME types known by the system.",
		"	 On Windows, the registry mess stores them.",
		"    In rare cases, the system may not have a mime entry for an",
		"      uncommon file type, and this program will not remove the",
		"      extension associated with that file type.",
		"    An example is '.go' files; Linux distributions usually do not",
		"      include a mime type for '.go' files by default (yet).",
	)
}

// Parse user arguments and modify global variables accordingly.
func flags() {
	// Program requires at least one user argument.
	// Print help and exit with status 1 if none were received.
	if len(os.Args) < 2 {
		help(1)
	}

	// Parse commandline arguments.
	flag.BoolVar(&doBasename, "b", false, "")
	flag.BoolVar(&doFullpath, "f", false, "")
	flag.Parse()

	// Modify global variables based on commandline arguments.
	targets = os.Args[1:]
	if doBasename && doFullpath {
		help(1)
	}
	if !doBasename && !doFullpath {
		return
	}
	bools := []string{"-b", "-f"}
	targets = filter(targets, bools...)
	if len(targets) == 0 {
		help(1)
	}
	return
}

// Remove elements in a slice (if they exist).
// Only remove EXACT matches.
func filter(slc []string, args ...string) (filtered []string) {
	for _, s := range slc {
		if slcHas(slc, s) {
			continue
		}
		filtered = append(filtered, s)
	}
	return
}

// Check whether a slice contains a string.
// Only return true if an element in the slice EXACTLY matches the string.
// If testing for more than one string,
//   return true if ANY of them match an element in the slice.
func slcHas(slc []string, args ...string) bool {
	for _, s := range slc {
		for _, a := range args {
			if s == a {
				return true
			}
		}
	}
	return false
}

func main() {
	defer os.Exit(0)
	var err error
	for _, t := range targets {
		switch {
		case doBasename:
			t = filepath.Base(t)
		case doFullpath:
			t, err = filepath.Abs(t)
			chkerr(err)
		}
		base, _ := splitExt(t)
		fmt.Println(base)
	}
}

// Exit with status 1 if an error occurs.
func chkerr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

// Split the extension off a filename.
// Return the basename and the extension.
func splitExt(filename string) (base, ext string) {
	base = filepath.Clean(filename)
	for {
		testext := filepath.Ext(base)
		if testext == "" || mime.TypeByExtension(testext) == "" {
			return
		}
		ext = concat(testext, ext)
		base = strings.TrimSuffix(base, testext)
	}
}

// Concatenate strings.
func concat(slc ...string) string {
	b := bytes.NewBuffer(nil)
	defer b.Reset()
	for _, s := range slc {
		b.WriteString(s)
	}
	return b.String()
}
<commit_msg>Cleanup<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"mime"
	"os"
	"path/filepath"
	"strings"
)

var (
	PrintBase bool
	PrintFull bool
	Paths     []string
)

func init() {
	if helpNeeded() {
		help(1)
	} else if helpWanted() {
		help(0)
	}
	flags()
}

func helpNeeded() bool {
	if noArgs := (len(os.Args) == 1); noArgs {
		return true
	}
	return false
}

func helpWanted() bool {
	switch os.Args[1] {
	case "-h", "h", "help", "--help", "-H", "H", "HELP", "--HELP", "-help", "--h", "--H":
		return true
	}
	return false
}

// Print help and exit with a status code.
func help(status int) {
	defer os.Exit(status)
	fmt.Printf(
		"%s\n\n  %s\n\n  %s\n%s\n\n  %s\n%s\n%s\n%s\n%s\n\n  %s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n",
		"rmext",
		"Usage: rmext file...",
		"Description:",
		"    Print the basename of files WITHOUT their extension.",
		"Options:",
		"   -b         Print the basename of files without their path",
		"   -f         Print the full path of files",
		"   (default)  Print the path of files exactly as they were entered",
		"                (minus any file extensions)",
		"Notes:",
		"    This program relies on the system's list of known mimetypes.",
		"	 On Linux, usually the file '/etc/mime.types' stores all",
		"      MIME types known by the system.",
		"	 On Windows, the registry mess stores them.",
		"    In rare cases, the system may not have a mime entry for an",
		"      uncommon file type, and this program will not remove the",
		"      extension associated with that file type.",
		"    An example is '.go' files; Linux distributions usually do not",
		"      include a mime type for '.go' files by default (yet).",
	)
}

func flags() {
	flag.BoolVar(&PrintBase, "b", false, "")
	flag.BoolVar(&PrintFull, "f", false, "")
	flag.Parse()

	// Modify global variables based on commandline arguments.
	Paths = os.Args[1:]
	if PrintBase && PrintFull {
		help(1)
	}
	if !PrintBase && !PrintFull {
		return
	}
	bools := []string{"-b", "-f"}
	Paths = filter(Paths, bools...)
	if len(Paths) == 0 {
		help(1)
	}
	return
}

// Remove elements in a slice (if they exist).
// Only remove EXACT matches.
func filter(slc []string, args ...string) (filtered []string) {
	for _, s := range slc {
		if slcHas(slc, s) {
			continue
		}
		filtered = append(filtered, s)
	}
	return
}

// Check whether a slice contains a string.
// Only return true if an element in the slice EXACTLY matches the string.
// If testing for more than one string,
//   return true if ANY of them match an element in the slice.
func slcHas(slc []string, args ...string) bool {
	for _, s := range slc {
		for _, a := range args {
			if s == a {
				return true
			}
		}
	}
	return false
}

func main() {
	defer os.Exit(0)
	var err error
	for _, t := range Paths {
		switch {
		case PrintBase:
			t = filepath.Base(t)
		case PrintFull:
			t, err = filepath.Abs(t)
			chkerr(err)
		}
		base, _ := splitExt(t)
		fmt.Println(base)
	}
}

// Exit with status 1 if an error occurs.
func chkerr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

// Split the extension off a filename.
// Return the basename and the extension.
func splitExt(filename string) (base, ext string) {
	base = filepath.Clean(filename)
	for {
		testext := filepath.Ext(base)
		if testext == "" || mime.TypeByExtension(testext) == "" {
			return
		}
		ext = concat(testext, ext)
		base = strings.TrimSuffix(base, testext)
	}
}

// Concatenate strings.
func concat(slc ...string) string {
	b := bytes.NewBuffer(nil)
	defer b.Reset()
	for _, s := range slc {
		b.WriteString(s)
	}
	return b.String()
}
<|endoftext|>


<commit_before>package gobot

import (
	"fmt"
	"log"
	"math/rand"
	"time"
)

type JSONRobot struct {
	Name        string            `json:"name"`
	Commands    []string          `json:"commands"`
	Connections []*JSONConnection `json:"connections"`
	Devices     []*JSONDevice     `json:"devices"`
}

type Robot struct {
	Name        string                                              `json:"-"`
	Commands    map[string]func(map[string]interface{}) interface{} `json:"-"`
	Work        func()                                              `json:"-"`
	connections connections                                         `json:"-"`
	devices     devices                                             `json:"-"`
}

type Robots []*Robot

func (r Robots) Start() {
	for _, robot := range r {
		robot.Start()
	}
}

func (r Robots) Each(f func(*Robot)) {
	for _, robot := range r {
		f(robot)
	}
}

func NewRobot(name string, v ...interface{}) *Robot {
	r := &Robot{
		Name:     name,
		Commands: make(map[string]func(map[string]interface{}) interface{}),
	}
	r.initName()
	log.Println("Initializing Robot", r.Name, "...")
	if len(v) > 0 {
		if v[0] == nil {
			v[0] = []Connection{}
		}
		r.initConnections(v[0].([]Connection))
	}
	if len(v) > 1 {
		if v[1] == nil {
			v[1] = []Device{}
		}
		r.initDevices(v[1].([]Device))
	}
	if len(v) > 2 {
		if v[2] == nil {
			v[2] = func() {}
		}
		r.Work = v[2].(func())
	}
	return r
}

func (r *Robot) AddCommand(name string, f func(map[string]interface{}) interface{}) {
	r.Commands[name] = f
}

func (r *Robot) Start() {
	log.Println("Starting Robot", r.Name, "...")
	if err := r.Connections().Start(); err != nil {
		panic("Could not start connections")
	}
	if err := r.Devices().Start(); err != nil {
		panic("Could not start devices")
	}
	if r.Work != nil {
		log.Println("Starting work...")
		r.Work()
	}
}

func (r *Robot) initName() {
	if r.Name == "" {
		rand.Seed(time.Now().UTC().UnixNano())
		i := rand.Int()
		r.Name = fmt.Sprintf("Robot%v", i)
	}
}

func (r *Robot) initConnections(c []Connection) {
	r.connections = make(connections, len(c))
	log.Println("Initializing connections...")
	for i, connection := range c {
		log.Println("Initializing connection", FieldByNamePtr(connection, "Name"), "...")
		r.connections[i] = NewConnection(connection, r)
	}
}

func (r *Robot) initDevices(d []Device) {
	r.devices = make([]*device, len(d))
	log.Println("Initializing devices...")
	for i, device := range d {
		log.Println("Initializing device", FieldByNamePtr(device, "Name"), "...")
		r.devices[i] = NewDevice(device, r)
	}
}

func (r *Robot) Devices() devices {
	return devices(r.devices)
}

func (r *Robot) Device(name string) *device {
	if r == nil {
		return nil
	}
	for _, device := range r.devices {
		if device.Name == name {
			return device
		}
	}
	return nil
}

func (r *Robot) Connections() connections {
	return connections(r.connections)
}

func (r *Robot) Connection(name string) *connection {
	if r == nil {
		return nil
	}
	for _, connection := range r.connections {
		if connection.Name == name {
			return connection
		}
	}
	return nil
}

func (r *Robot) ToJSON() *JSONRobot {
	jsonRobot := &JSONRobot{
		Name:        r.Name,
		Commands:    []string{},
		Connections: []*JSONConnection{},
		Devices:     []*JSONDevice{},
	}
	for command := range r.Commands {
		jsonRobot.Commands = append(jsonRobot.Commands, command)
	}
	for _, device := range r.Devices() {
		jsonDevice := device.ToJSON()
		jsonRobot.Connections = append(jsonRobot.Connections, jsonDevice.Connection)
		jsonRobot.Devices = append(jsonRobot.Devices, jsonDevice)
	}
	return jsonRobot
}
<commit_msg>Add AddDevice and AddConnection functions<commit_after>package gobot

import (
	"fmt"
	"log"
	"math/rand"
	"time"
)

type JSONRobot struct {
	Name        string            `json:"name"`
	Commands    []string          `json:"commands"`
	Connections []*JSONConnection `json:"connections"`
	Devices     []*JSONDevice     `json:"devices"`
}

type Robot struct {
	Name        string                                              `json:"-"`
	Commands    map[string]func(map[string]interface{}) interface{} `json:"-"`
	Work        func()                                              `json:"-"`
	connections connections                                         `json:"-"`
	devices     devices                                             `json:"-"`
}

type Robots []*Robot

func (r Robots) Start() {
	for _, robot := range r {
		robot.Start()
	}
}

func (r Robots) Each(f func(*Robot)) {
	for _, robot := range r {
		f(robot)
	}
}

func NewRobot(name string, v ...interface{}) *Robot {
	r := &Robot{
		Name:        name,
		Commands:    make(map[string]func(map[string]interface{}) interface{}),
		connections: connections{},
		devices:     devices{},
	}
	r.initName()
	log.Println("Initializing Robot", r.Name, "...")
	if len(v) > 0 {
		if v[0] == nil {
			v[0] = []Connection{}
		}
		r.initConnections(v[0].([]Connection))
	}
	if len(v) > 1 {
		if v[1] == nil {
			v[1] = []Device{}
		}
		r.initDevices(v[1].([]Device))
	}
	if len(v) > 2 {
		if v[2] == nil {
			v[2] = func() {}
		}
		r.Work = v[2].(func())
	}
	return r
}

func (r *Robot) AddDevice(d Device) *device {
	device := NewDevice(d, r)
	r.devices = append(r.devices, device)
	return device
}

func (r *Robot) AddConnection(c Connection) *connection {
	connection := NewConnection(c, r)
	r.connections = append(r.connections, connection)
	return connection
}

func (r *Robot) AddCommand(name string, f func(map[string]interface{}) interface{}) {
	r.Commands[name] = f
}

func (r *Robot) Start() {
	log.Println("Starting Robot", r.Name, "...")
	if err := r.Connections().Start(); err != nil {
		panic("Could not start connections")
	}
	if err := r.Devices().Start(); err != nil {
		panic("Could not start devices")
	}
	if r.Work != nil {
		log.Println("Starting work...")
		r.Work()
	}
}

func (r *Robot) initName() {
	if r.Name == "" {
		rand.Seed(time.Now().UTC().UnixNano())
		i := rand.Int()
		r.Name = fmt.Sprintf("Robot%v", i)
	}
}

func (r *Robot) initConnections(c []Connection) {
	log.Println("Initializing connections...")
	for _, connection := range c {
		log.Println("Initializing connection", FieldByNamePtr(connection, "Name"), "...")
		r.AddConnection(connection)
	}
}

func (r *Robot) initDevices(d []Device) {
	log.Println("Initializing devices...")
	for _, device := range d {
		log.Println("Initializing device", FieldByNamePtr(device, "Name"), "...")
		r.AddDevice(device)
	}
}

func (r *Robot) Devices() devices {
	return devices(r.devices)
}

func (r *Robot) Device(name string) *device {
	if r == nil {
		return nil
	}
	for _, device := range r.devices {
		if device.Name == name {
			return device
		}
	}
	return nil
}

func (r *Robot) Connections() connections {
	return connections(r.connections)
}

func (r *Robot) Connection(name string) *connection {
	if r == nil {
		return nil
	}
	for _, connection := range r.connections {
		if connection.Name == name {
			return connection
		}
	}
	return nil
}

func (r *Robot) ToJSON() *JSONRobot {
	jsonRobot := &JSONRobot{
		Name:        r.Name,
		Commands:    []string{},
		Connections: []*JSONConnection{},
		Devices:     []*JSONDevice{},
	}
	for command := range r.Commands {
		jsonRobot.Commands = append(jsonRobot.Commands, command)
	}
	for _, device := range r.Devices() {
		jsonDevice := device.ToJSON()
		jsonRobot.Connections = append(jsonRobot.Connections, jsonDevice.Connection)
		jsonRobot.Devices = append(jsonRobot.Devices, jsonDevice)
	}
	return jsonRobot
}
<|endoftext|>


<commit_before>package main

import (
	twodee "../libs/twodee"
)

const (
	UpLayer twodee.GameEventType = iota
	DownLayer
	UpWaterLevel
	DownWaterLevel
	PlayerMove
	GameIsClosing
	PlayOutdoorMusic
	PlayExploreMusic
	PlayWarningMusic
	PlayDangerMusic
	PauseMusic
	ResumeMusic
	MenuPauseMusic
	MenuResumeMusic
	PlayerTouchedItem
	PlayerUsedItem
	PlayerDestroyedItem
	sentinel
)

const (
	NumGameEventTypes = int(sentinel)
)

type MoveDirection byte

const (
	None  MoveDirection = iota
	North MoveDirection = 1 << (iota - 1)
	East
	South
	West
)

type PlayerMoveEvent struct {
	*twodee.BasicGameEvent
	Dir     MoveDirection
	Inverse bool
}

func NewPlayerMoveEvent(direction MoveDirection) (e *PlayerMoveEvent) {
	e = &PlayerMoveEvent{
		twodee.NewBasicGameEvent(PlayerMove),
		direction,
		false,
	}
	return
}

func NewInversePlayerMoveEvent(direction MoveDirection) (e *PlayerMoveEvent) {
	e = &PlayerMoveEvent{
		twodee.NewBasicGameEvent(PlayerMove),
		direction,
		true,
	}
	return
}

type PlayerTouchedItemEvent struct {
	*twodee.BasicGameEvent
	Item *Item
}

func NewPlayerTouchedItemEvent(i *Item) (e *PlayerTouchedItemEvent) {
	e = &PlayerTouchedItemEvent{
		twodee.NewBasicGameEvent(PlayerTouchedItem),
		i,
	}
	return
}

type PlayerUsedItemEvent struct {
	*twodee.BasicGameEvent
	Item *Item
}

func NewPlayerUsedItemEvent(i *Item) (e *PlayerUsedItemEvent) {
	e = &PlayerUsedItemEvent{
		twodee.NewBasicGameEvent(PlayerUsedItem),
		i,
	}
	return
}

type PlayerDestroyedItemEvent struct {
	*twodee.BasicGameEvent
	Item *Item
}

func NewPlayerDestroyedItemEvent(i *Item) (e *PlayerDestroyedItemEvent) {
	e = &PlayerDestroyedItemEvent{
		twodee.NewBasicGameEvent(PlayerDestroyedItem),
		i,
	}
	return
}
<commit_msg>Move some anonymous embedded pointer fields to anonymous struct fields in the event structs.<commit_after>package main

import (
	twodee "../libs/twodee"
)

const (
	UpLayer twodee.GameEventType = iota
	DownLayer
	UpWaterLevel
	DownWaterLevel
	PlayerMove
	GameIsClosing
	PlayOutdoorMusic
	PlayExploreMusic
	PlayWarningMusic
	PlayDangerMusic
	PauseMusic
	ResumeMusic
	MenuPauseMusic
	MenuResumeMusic
	PlayerTouchedItem
	PlayerUsedItem
	PlayerDestroyedItem
	sentinel
)

const (
	NumGameEventTypes = int(sentinel)
)

type MoveDirection byte

const (
	None  MoveDirection = iota
	North MoveDirection = 1 << (iota - 1)
	East
	South
	West
)

type PlayerMoveEvent struct {
	twodee.BasicGameEvent
	Dir     MoveDirection
	Inverse bool
}

func NewPlayerMoveEvent(direction MoveDirection) (e *PlayerMoveEvent) {
	e = &PlayerMoveEvent{
		*twodee.NewBasicGameEvent(PlayerMove),
		direction,
		false,
	}
	return
}

func NewInversePlayerMoveEvent(direction MoveDirection) (e *PlayerMoveEvent) {
	e = &PlayerMoveEvent{
		*twodee.NewBasicGameEvent(PlayerMove),
		direction,
		true,
	}
	return
}

type PlayerTouchedItemEvent struct {
	twodee.BasicGameEvent
	Item *Item
}

func NewPlayerTouchedItemEvent(i *Item) (e *PlayerTouchedItemEvent) {
	e = &PlayerTouchedItemEvent{
		*twodee.NewBasicGameEvent(PlayerTouchedItem),
		i,
	}
	return
}

type PlayerUsedItemEvent struct {
	twodee.BasicGameEvent
	Item *Item
}

func NewPlayerUsedItemEvent(i *Item) (e *PlayerUsedItemEvent) {
	e = &PlayerUsedItemEvent{
		*twodee.NewBasicGameEvent(PlayerUsedItem),
		i,
	}
	return
}

type PlayerDestroyedItemEvent struct {
	twodee.BasicGameEvent
	Item *Item
}

func NewPlayerDestroyedItemEvent(i *Item) (e *PlayerDestroyedItemEvent) {
	e = &PlayerDestroyedItemEvent{
		*twodee.NewBasicGameEvent(PlayerDestroyedItem),
		i,
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"

	"github.com/flynn/flynn-controller/client"
	"github.com/flynn/strowger/types"
)

var cmdRouteAddHTTP = &Command{
	Run:   runRouteAddHTTP,
	Usage: "route-add-http [-s <service>] [-c <tls-cert>] [-k <tls-key>] <domain>",
	Short: "add a HTTP route",
	Long:  `Add a HTTP route to an app`,
}

var routeHTTPService string
var tlsCertPath string
var tlsKeyPath string
var sticky bool

func init() {
	cmdRouteAddHTTP.Flag.StringVarP(&routeHTTPService, "service", "s", "", "service name to route domain to (defaults to APPNAME-web)")
	cmdRouteAddHTTP.Flag.StringVarP(&tlsCertPath, "tls-cert", "c", "", "path to PEM encoded certificate for TLS, - for stdin")
	cmdRouteAddHTTP.Flag.StringVarP(&tlsKeyPath, "tls-key", "k", "", "path to PEM encoded private key for TLS, - for stdin")
	cmdRouteAddHTTP.Flag.BoolVarP(&sticky, "sticky", "s", false, "enable cookie-based sticky routing")
}

func runRouteAddHTTP(cmd *Command, args []string, client *controller.Client) error {
	var tlsCert []byte
	var tlsKey []byte

	if len(args) != 1 {
		cmd.printUsage(true)
	}

	if routeHTTPService == "" {
		routeHTTPService = mustApp() + "-web"
	}

	if tlsCertPath != "" && tlsKeyPath != "" {
		var stdin []byte
		var err error

		if tlsCertPath == "-" || tlsKeyPath == "-" {
			stdin, err = ioutil.ReadAll(os.Stdin)
			if err != nil {
				return fmt.Errorf("Failed to read from stdin: %s", err)
			}
		}

		tlsCert, err = readPEM("CERTIFICATE", tlsCertPath, stdin)
		if err != nil {
			return errors.New("Failed to read TLS Cert")
		}
		tlsKey, err = readPEM("PRIVATE KEY", tlsKeyPath, stdin)
		if err != nil {
			return errors.New("Failed to read TLS Key")
		}
	} else if tlsCertPath != "" || tlsKeyPath != "" {
		return errors.New("Both the TLS certificate AND private key need to be specified")
	}

	hr := &strowger.HTTPRoute{
		Service: routeHTTPService,
		Domain:  args[0],
		TLSCert: string(tlsCert),
		TLSKey:  string(tlsKey),
		Sticky:  sticky,
	}
	route := hr.ToRoute()
	if err := client.CreateRoute(mustApp(), route); err != nil {
		return err
	}
	fmt.Println(route.ID)
	return nil
}

func readPEM(typ string, path string, stdin []byte) ([]byte, error) {
	if path == "-" {
		var buf bytes.Buffer
		var block *pem.Block
		for {
			block, stdin = pem.Decode(stdin)
			if block == nil {
				break
			}
			if block.Type == typ {
				pem.Encode(&buf, block)
			}
		}
		if buf.Len() > 0 {
			return buf.Bytes(), nil
		}
		return nil, errors.New("No PEM blocks found in stdin")
	}
	return ioutil.ReadFile(path)
}

var cmdRoutes = &Command{
	Run:   runRoutes,
	Usage: "routes",
	Short: "list routes",
	Long:  `list routes for application"`,
}

func runRoutes(cmd *Command, args []string, client *controller.Client) error {
	if len(args) != 0 {
		cmd.printUsage(true)
	}
	routes, err := client.RouteList(mustApp())
	if err != nil {
		return err
	}

	w := tabWriter()
	defer w.Flush()

	var route, protocol, service string
	listRec(w, "ROUTE", "SERVICE", "ID")
	for _, k := range routes {
		switch k.Type {
		case "tcp":
			protocol = "tcp"
			route = strconv.Itoa(k.TCPRoute().Port)
			service = k.TCPRoute().Service
		case "http":
			route = k.HTTPRoute().Domain
			service = k.TCPRoute().Service
			if k.HTTPRoute().TLSCert == "" {
				protocol = "http"
			} else {
				protocol = "https"
			}
		}
		listRec(w, protocol+":"+route, service, k.ID)
	}
	return nil
}

var cmdRouteRemove = &Command{
	Run:   runRouteRemove,
	Usage: "route-remove <id>",
	Short: "remove a route",
	Long:  "Command route-remove removes a route from the Flynn controller.",
}

func runRouteRemove(cmd *Command, args []string, client *controller.Client) error {
	if len(args) != 1 {
		cmd.printUsage(true)
	}
	routeID := args[0]

	if err := client.DeleteRoute(mustApp(), routeID); err != nil {
		return err
	}
	fmt.Printf("Route %s removed.\n", routeID)
	return nil
}
<commit_msg>Fix typo<commit_after>package main

import (
	"bytes"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"

	"github.com/flynn/flynn-controller/client"
	"github.com/flynn/strowger/types"
)

var cmdRouteAddHTTP = &Command{
	Run:   runRouteAddHTTP,
	Usage: "route-add-http [-s <service>] [-c <tls-cert>] [-k <tls-key>] <domain>",
	Short: "add a HTTP route",
	Long:  `Add a HTTP route to an app`,
}

var routeHTTPService string
var tlsCertPath string
var tlsKeyPath string
var sticky bool

func init() {
	cmdRouteAddHTTP.Flag.StringVarP(&routeHTTPService, "service", "s", "", "service name to route domain to (defaults to APPNAME-web)")
	cmdRouteAddHTTP.Flag.StringVarP(&tlsCertPath, "tls-cert", "c", "", "path to PEM encoded certificate for TLS, - for stdin")
	cmdRouteAddHTTP.Flag.StringVarP(&tlsKeyPath, "tls-key", "k", "", "path to PEM encoded private key for TLS, - for stdin")
	cmdRouteAddHTTP.Flag.BoolVarP(&sticky, "sticky", "t", false, "enable cookie-based sticky routing")
}

func runRouteAddHTTP(cmd *Command, args []string, client *controller.Client) error {
	var tlsCert []byte
	var tlsKey []byte

	if len(args) != 1 {
		cmd.printUsage(true)
	}

	if routeHTTPService == "" {
		routeHTTPService = mustApp() + "-web"
	}

	if tlsCertPath != "" && tlsKeyPath != "" {
		var stdin []byte
		var err error

		if tlsCertPath == "-" || tlsKeyPath == "-" {
			stdin, err = ioutil.ReadAll(os.Stdin)
			if err != nil {
				return fmt.Errorf("Failed to read from stdin: %s", err)
			}
		}

		tlsCert, err = readPEM("CERTIFICATE", tlsCertPath, stdin)
		if err != nil {
			return errors.New("Failed to read TLS Cert")
		}
		tlsKey, err = readPEM("PRIVATE KEY", tlsKeyPath, stdin)
		if err != nil {
			return errors.New("Failed to read TLS Key")
		}
	} else if tlsCertPath != "" || tlsKeyPath != "" {
		return errors.New("Both the TLS certificate AND private key need to be specified")
	}

	hr := &strowger.HTTPRoute{
		Service: routeHTTPService,
		Domain:  args[0],
		TLSCert: string(tlsCert),
		TLSKey:  string(tlsKey),
		Sticky:  sticky,
	}
	route := hr.ToRoute()
	if err := client.CreateRoute(mustApp(), route); err != nil {
		return err
	}
	fmt.Println(route.ID)
	return nil
}

func readPEM(typ string, path string, stdin []byte) ([]byte, error) {
	if path == "-" {
		var buf bytes.Buffer
		var block *pem.Block
		for {
			block, stdin = pem.Decode(stdin)
			if block == nil {
				break
			}
			if block.Type == typ {
				pem.Encode(&buf, block)
			}
		}
		if buf.Len() > 0 {
			return buf.Bytes(), nil
		}
		return nil, errors.New("No PEM blocks found in stdin")
	}
	return ioutil.ReadFile(path)
}

var cmdRoutes = &Command{
	Run:   runRoutes,
	Usage: "routes",
	Short: "list routes",
	Long:  `list routes for application"`,
}

func runRoutes(cmd *Command, args []string, client *controller.Client) error {
	if len(args) != 0 {
		cmd.printUsage(true)
	}
	routes, err := client.RouteList(mustApp())
	if err != nil {
		return err
	}

	w := tabWriter()
	defer w.Flush()

	var route, protocol, service string
	listRec(w, "ROUTE", "SERVICE", "ID")
	for _, k := range routes {
		switch k.Type {
		case "tcp":
			protocol = "tcp"
			route = strconv.Itoa(k.TCPRoute().Port)
			service = k.TCPRoute().Service
		case "http":
			route = k.HTTPRoute().Domain
			service = k.TCPRoute().Service
			if k.HTTPRoute().TLSCert == "" {
				protocol = "http"
			} else {
				protocol = "https"
			}
		}
		listRec(w, protocol+":"+route, service, k.ID)
	}
	return nil
}

var cmdRouteRemove = &Command{
	Run:   runRouteRemove,
	Usage: "route-remove <id>",
	Short: "remove a route",
	Long:  "Command route-remove removes a route from the Flynn controller.",
}

func runRouteRemove(cmd *Command, args []string, client *controller.Client) error {
	if len(args) != 1 {
		cmd.printUsage(true)
	}
	routeID := args[0]

	if err := client.DeleteRoute(mustApp(), routeID); err != nil {
		return err
	}
	fmt.Printf("Route %s removed.\n", routeID)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"log"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/dchest/siphash"
	"github.com/hashicorp/golang-lru"
	"github.com/miekg/dns"
)

const (
	// MinTTL Minimum TTL
	MinTTL = 60
	// MaxTTL Maximum TTL
	MaxTTL = 604800
	// VacuumPeriod Vacuum period in seconds
	VacuumPeriod = 30
)

// SipHashKey SipHash secret key
type SipHashKey struct {
	k1 uint64
	k2 uint64
}

// UpstreamServer Upstream server
type UpstreamServer struct {
	addr    string
	offline bool
}

// UpstreamServers List of upstream servers
type UpstreamServers struct {
	lock    sync.RWMutex
	servers []UpstreamServer
	live    []string
}

var (
	address            = flag.String("listen", ":53", "Address to listen to (TCP and UDP)")
	upstreamServersStr = flag.String("upstream", "8.8.8.8:53", "Comma-delimited list of upstream servers")
	upstreamServers    *UpstreamServers
	cacheSize          = flag.Int("cachesize", 10000000, "Cache size (default=10000000)")
	memSize            = flag.Uint64("memsize", 1*1024, "Memory size in MB (default=1GB)")
	minLabelsCount     = flag.Int("minlabels", 2, "Minimum number of labels (default=2)")
	cache              *lru.ARCCache
	sipHashKey         = SipHashKey{k1: 0, k2: 0}
)

func parseUpstreamServers(str string) (*UpstreamServers, error) {
	servers := []UpstreamServer{}
	live := []string{}
	for _, addr := range strings.Split(str, ",") {
		server := UpstreamServer{addr: addr}
		servers = append(servers, server)
		live = append(live, addr)
	}
	res := UpstreamServers{servers: servers, live: live}
	return &res, nil
}

func randUint64() uint64 {
	buf := make([]byte, 8)
	length, err := rand.Read(buf)
	if err != nil || length != len(buf) {
		log.Fatal("RNG failure")
	}
	return binary.LittleEndian.Uint64(buf)
}

func main() {
	flag.Parse()
	*memSize *= 1024 * 1024
	if *cacheSize < 2 {
		log.Fatal("Cache size too small")
	}
	cache, _ = lru.NewARC(*cacheSize)
	upstreamServers, _ = parseUpstreamServers(*upstreamServersStr)
	sipHashKey = SipHashKey{k1: randUint64(), k2: randUint64()}
	fmt.Println("RPDNS")
	dns.HandleFunc(".", route)
	udpServer := &dns.Server{Addr: *address, Net: "udp"}
	tcpServer := &dns.Server{Addr: *address, Net: "tcp"}
	go func() {
		log.Fatal(udpServer.ListenAndServe())
	}()
	go func() {
		log.Fatal(tcpServer.ListenAndServe())
	}()
	vacuumThread()
}

// CacheKey Key for a cache entry
type CacheKey struct {
	Name   string `dns:"cdomain-name"`
	Qtype  uint16
	DNSSEC bool
}

// CacheVal Value for a cache entry
type CacheVal struct {
	ValidUntil time.Time
	Response   *dns.Msg
}

func getKey(req *dns.Msg) (*CacheKey, error) {
	questions := req.Question
	if len(questions) != 1 {
		return nil, errors.New("Invalid number of questions")
	}
	question := questions[0]
	if question.Qclass != dns.ClassINET {
		return nil, errors.New("Unsupported question class")
	}
	if dns.CountLabel(question.Name) < *minLabelsCount {
		return nil, errors.New("Not enough labels")
	}
	dnssec := false
	for _, extra := range req.Extra {
		if extra.Header().Rrtype == dns.TypeOPT {
			dnssec = extra.(*dns.OPT).Do()
		}
	}
	CacheKey := CacheKey{Name: strings.ToLower(question.Name), Qtype: question.Qtype, DNSSEC: dnssec}
	return &CacheKey, nil
}

func getMaxPayloadSize(req *dns.Msg) uint16 {
	opt := req.IsEdns0()
	if opt == nil {
		return dns.MinMsgSize
	}
	maxPayloadSize := opt.UDPSize()
	if maxPayloadSize < dns.MinMsgSize {
		maxPayloadSize = dns.MinMsgSize
	}
	return maxPayloadSize
}

func pickUpstream(req *dns.Msg) (*string, error) {
	name := strings.ToLower(req.Question[0].Name)
	h := siphash.Hash(sipHashKey.k1, sipHashKey.k2, []byte(name))
	upstreamServers.lock.RLock()
	liveCount := len(upstreamServers.live)
	if liveCount <= 0 {
		upstreamServers.lock.RUnlock()
		return nil, errors.New("All upstream servers are down")
	}
	i := h % uint64(len(upstreamServers.live))
	res := upstreamServers.live[i]
	upstreamServers.lock.RUnlock()
	return &res, nil
}

func markDead(addr string) {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	servers := upstreamServers.servers
	live := []string{}
	for i, server := range upstreamServers.servers {
		if server.addr == addr {
			servers[i].offline = true
		} else if server.offline == false {
			live = append(live, server.addr)
		}
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
}

func resetUpstreamServers() {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	servers := upstreamServers.servers
	if len(servers) == len(upstreamServers.live) {
		return
	}
	live := []string{}
	for i, server := range upstreamServers.servers {
		servers[i].offline = false
		live = append(live, server.addr)
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
}

func resolve(req *dns.Msg, dnssec bool) (*dns.Msg, error) {
	extra2 := []dns.RR{}
	for _, extra := range req.Extra {
		if extra.Header().Rrtype != dns.TypeOPT {
			extra2 = append(extra2, extra)
		}
	}
	req.Extra = extra2
	req.SetEdns0(dns.MaxMsgSize, dnssec)
	addr, err := pickUpstream(req)
	if err != nil {
		return nil, err
	}
	client := &dns.Client{Net: "udp"}
	client.SingleInflight = true
	resolved, _, err := client.Exchange(req, *addr)
	if err != nil {
		markDead(*addr)
		return nil, err
	}
	if resolved.Truncated {
		client = &dns.Client{Net: "tcp"}
		client.SingleInflight = true
		resolved, _, err = client.Exchange(req, *addr)
	}
	if err != nil {
		return nil, err
	}
	resolved.Compress = true
	return resolved, nil
}

func getMinTTL(resp *dns.Msg) time.Duration {
	ttl := uint32(MaxTTL)
	for _, rr := range resp.Answer {
		if rr.Header().Ttl < ttl {
			ttl = rr.Header().Ttl
		}
	}
	if ttl < MinTTL {
		ttl = MaxTTL
	}
	return time.Duration(ttl) * time.Second
}

func sendTruncated(w dns.ResponseWriter, msgHdr dns.MsgHdr) {
	emptyResp := new(dns.Msg)
	emptyResp.MsgHdr = msgHdr
	emptyResp.Truncated = true
	w.WriteMsg(emptyResp)
}

func vacuumThread() {
	for {
		time.Sleep(VacuumPeriod * time.Second)
		resetUpstreamServers()
		memStats := new(runtime.MemStats)
		runtime.ReadMemStats(memStats)
		if memStats.Alloc > (*memSize)*1024*1024 {
			cache.Purge()
		}
	}
}

func failWithRcode(w dns.ResponseWriter, r *dns.Msg, rCode int) {
	m := new(dns.Msg)
	m.SetRcode(r, rCode)
	w.WriteMsg(m)
}

func route(w dns.ResponseWriter, req *dns.Msg) {
	keyP, err := getKey(req)
	if err != nil {
		failWithRcode(w, req, dns.RcodeRefused)
		return
	}
	maxPayloadSize := getMaxPayloadSize(req)
	var resp *dns.Msg
	cacheValP, _ := cache.Get(*keyP)
	if cacheValP != nil {
		cacheVal := cacheValP.(CacheVal)
		remaining := -time.Since(cacheVal.ValidUntil)
		if remaining > 0 {
			resp = cacheVal.Response.Copy()
			resp.Id = req.Id
			resp.Question = req.Question
		}
	}
	if resp == nil {
		resp, err = resolve(req, keyP.DNSSEC)
		if err != nil {
			dns.HandleFailed(w, req)
			return
		}
		validUntil := time.Now().Add(getMinTTL(resp))
		cache.Add(*keyP, CacheVal{ValidUntil: validUntil, Response: resp})
	}
	packed, _ := resp.Pack()
	packedLen := len(packed)
	if uint16(packedLen) > maxPayloadSize {
		sendTruncated(w, resp.MsgHdr)
	} else {
		w.WriteMsg(resp)
	}
}
<commit_msg>Slightly change how upstream servers are chosen<commit_after>package main

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/dchest/siphash"
	"github.com/hashicorp/golang-lru"
	"github.com/miekg/dns"
)

const (
	// MinTTL Minimum TTL
	MinTTL = 60
	// MaxTTL Maximum TTL
	MaxTTL = 604800
	// VacuumPeriod Vacuum period in seconds
	VacuumPeriod = 30
)

// SipHashKey SipHash secret key
type SipHashKey struct {
	k1 uint64
	k2 uint64
}

// UpstreamServer Upstream server
type UpstreamServer struct {
	addr    string
	offline bool
}

// UpstreamServers List of upstream servers
type UpstreamServers struct {
	lock    sync.RWMutex
	servers []UpstreamServer
	live    []string
}

var (
	address            = flag.String("listen", ":53", "Address to listen to (TCP and UDP)")
	upstreamServersStr = flag.String("upstream", "8.8.8.8:53", "Comma-delimited list of upstream servers")
	upstreamServers    *UpstreamServers
	cacheSize          = flag.Int("cachesize", 10000000, "Cache size (default=10000000)")
	memSize            = flag.Uint64("memsize", 1*1024, "Memory size in MB (default=1GB)")
	minLabelsCount     = flag.Int("minlabels", 2, "Minimum number of labels (default=2)")
	cache              *lru.ARCCache
	sipHashKey         = SipHashKey{k1: 0, k2: 0}
)

func parseUpstreamServers(str string) (*UpstreamServers, error) {
	servers := []UpstreamServer{}
	live := []string{}
	for _, addr := range strings.Split(str, ",") {
		server := UpstreamServer{addr: addr}
		servers = append(servers, server)
		live = append(live, addr)
	}
	res := UpstreamServers{servers: servers, live: live}
	return &res, nil
}

func randUint64() uint64 {
	buf := make([]byte, 8)
	length, err := rand.Read(buf)
	if err != nil || length != len(buf) {
		log.Fatal("RNG failure")
	}
	return binary.LittleEndian.Uint64(buf)
}

func main() {
	flag.Parse()
	*memSize *= 1024 * 1024
	if *cacheSize < 2 {
		log.Fatal("Cache size too small")
	}
	cache, _ = lru.NewARC(*cacheSize)
	upstreamServers, _ = parseUpstreamServers(*upstreamServersStr)
	sipHashKey = SipHashKey{k1: randUint64(), k2: randUint64()}
	fmt.Println("RPDNS")
	dns.HandleFunc(".", route)
	udpServer := &dns.Server{Addr: *address, Net: "udp"}
	tcpServer := &dns.Server{Addr: *address, Net: "tcp"}
	go func() {
		log.Fatal(udpServer.ListenAndServe())
	}()
	go func() {
		log.Fatal(tcpServer.ListenAndServe())
	}()
	vacuumThread()
}

// CacheKey Key for a cache entry
type CacheKey struct {
	Name   string `dns:"cdomain-name"`
	Qtype  uint16
	DNSSEC bool
}

// CacheVal Value for a cache entry
type CacheVal struct {
	ValidUntil time.Time
	Response   *dns.Msg
}

func getKey(req *dns.Msg) (*CacheKey, error) {
	questions := req.Question
	if len(questions) != 1 {
		return nil, errors.New("Invalid number of questions")
	}
	question := questions[0]
	if question.Qclass != dns.ClassINET {
		return nil, errors.New("Unsupported question class")
	}
	if dns.CountLabel(question.Name) < *minLabelsCount {
		return nil, errors.New("Not enough labels")
	}
	dnssec := false
	for _, extra := range req.Extra {
		if extra.Header().Rrtype == dns.TypeOPT {
			dnssec = extra.(*dns.OPT).Do()
		}
	}
	CacheKey := CacheKey{Name: strings.ToLower(question.Name), Qtype: question.Qtype, DNSSEC: dnssec}
	return &CacheKey, nil
}

func getMaxPayloadSize(req *dns.Msg) uint16 {
	opt := req.IsEdns0()
	if opt == nil {
		return dns.MinMsgSize
	}
	maxPayloadSize := opt.UDPSize()
	if maxPayloadSize < dns.MinMsgSize {
		maxPayloadSize = dns.MinMsgSize
	}
	return maxPayloadSize
}

func pickUpstream(req *dns.Msg) (*string, error) {
	name := strings.ToLower(req.Question[0].Name)
	h := siphash.Hash(sipHashKey.k1, sipHashKey.k2, []byte(name))
	upstreamServers.lock.RLock()
	liveCount := uint64(len(upstreamServers.live))
	if liveCount <= 0 {
		upstreamServers.lock.RUnlock()
		return nil, errors.New("All upstream servers are down")
	}
	i := h / (math.MaxUint64 / liveCount)
	if i >= liveCount {
		i = liveCount - 1
	}
	res := upstreamServers.live[i]
	upstreamServers.lock.RUnlock()
	return &res, nil
}

func markDead(addr string) {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	servers := upstreamServers.servers
	live := []string{}
	for i, server := range upstreamServers.servers {
		if server.addr == addr {
			servers[i].offline = true
		} else if server.offline == false {
			live = append(live, server.addr)
		}
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
}

func resetUpstreamServers() {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	servers := upstreamServers.servers
	if len(servers) == len(upstreamServers.live) {
		return
	}
	live := []string{}
	for i, server := range upstreamServers.servers {
		servers[i].offline = false
		live = append(live, server.addr)
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
}

func resolve(req *dns.Msg, dnssec bool) (*dns.Msg, error) {
	extra2 := []dns.RR{}
	for _, extra := range req.Extra {
		if extra.Header().Rrtype != dns.TypeOPT {
			extra2 = append(extra2, extra)
		}
	}
	req.Extra = extra2
	req.SetEdns0(dns.MaxMsgSize, dnssec)
	addr, err := pickUpstream(req)
	if err != nil {
		return nil, err
	}
	client := &dns.Client{Net: "udp"}
	client.SingleInflight = true
	resolved, _, err := client.Exchange(req, *addr)
	if err != nil {
		markDead(*addr)
		return nil, err
	}
	if resolved.Truncated {
		client = &dns.Client{Net: "tcp"}
		client.SingleInflight = true
		resolved, _, err = client.Exchange(req, *addr)
	}
	if err != nil {
		return nil, err
	}
	resolved.Compress = true
	return resolved, nil
}

func getMinTTL(resp *dns.Msg) time.Duration {
	ttl := uint32(MaxTTL)
	for _, rr := range resp.Answer {
		if rr.Header().Ttl < ttl {
			ttl = rr.Header().Ttl
		}
	}
	if ttl < MinTTL {
		ttl = MaxTTL
	}
	return time.Duration(ttl) * time.Second
}

func sendTruncated(w dns.ResponseWriter, msgHdr dns.MsgHdr) {
	emptyResp := new(dns.Msg)
	emptyResp.MsgHdr = msgHdr
	emptyResp.Truncated = true
	w.WriteMsg(emptyResp)
}

func vacuumThread() {
	for {
		time.Sleep(VacuumPeriod * time.Second)
		resetUpstreamServers()
		memStats := new(runtime.MemStats)
		runtime.ReadMemStats(memStats)
		if memStats.Alloc > (*memSize)*1024*1024 {
			cache.Purge()
		}
	}
}

func failWithRcode(w dns.ResponseWriter, r *dns.Msg, rCode int) {
	m := new(dns.Msg)
	m.SetRcode(r, rCode)
	w.WriteMsg(m)
}

func route(w dns.ResponseWriter, req *dns.Msg) {
	keyP, err := getKey(req)
	if err != nil {
		failWithRcode(w, req, dns.RcodeRefused)
		return
	}
	maxPayloadSize := getMaxPayloadSize(req)
	var resp *dns.Msg
	cacheValP, _ := cache.Get(*keyP)
	if cacheValP != nil {
		cacheVal := cacheValP.(CacheVal)
		remaining := -time.Since(cacheVal.ValidUntil)
		if remaining > 0 {
			resp = cacheVal.Response.Copy()
			resp.Id = req.Id
			resp.Question = req.Question
		}
	}
	if resp == nil {
		resp, err = resolve(req, keyP.DNSSEC)
		if err != nil {
			dns.HandleFailed(w, req)
			return
		}
		validUntil := time.Now().Add(getMinTTL(resp))
		cache.Add(*keyP, CacheVal{ValidUntil: validUntil, Response: resp})
	}
	packed, _ := resp.Pack()
	packedLen := len(packed)
	if uint16(packedLen) > maxPayloadSize {
		sendTruncated(w, resp.MsgHdr)
	} else {
		w.WriteMsg(resp)
	}
}
<|endoftext|>