<commit_before>package ytdl

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
)

func getDownloadURL(format Format, htmlPlayerFile string) (*url.URL, error) {
	var sig string
	if s, ok := format.meta["s"]; ok && len(s.(string)) > 0 {
		tokens, err := getSigTokens(htmlPlayerFile)
		if err != nil {
			return nil, err
		}
		sig = decipherTokens(tokens, s.(string))
	} else {
		if s, ok := format.meta["sig"]; ok {
			sig = s.(string)
		}
	}
	var urlString string
	if s, ok := format.meta["url"]; ok {
		urlString = s.(string)
	} else if s, ok := format.meta["stream"]; ok {
		if c, ok := format.meta["conn"]; ok {
			urlString = c.(string)
			if urlString[len(urlString)-1] != '/' {
				urlString += "/"
			}
		}
		urlString += s.(string)
	} else {
		return nil, fmt.Errorf("Couldn't extract url from format")
	}
	urlString, err := url.QueryUnescape(urlString)
	if err != nil {
		return nil, err
	}
	u, err := url.Parse(urlString)
	if err != nil {
		return nil, err
	}
	query := u.Query()
	query.Set("ratebypass", "yes")
	if len(sig) > 0 {
		query.Set("sig", sig)
	}
	u.RawQuery = query.Encode()
	return u, nil
}

func decipherTokens(tokens []string, sig string) string {
	var pos int
	sigSplit := strings.Split(sig, "")
	for i, l := 0, len(tokens); i < l; i++ {
		tok := tokens[i]
		if len(tok) > 1 {
			pos, _ = strconv.Atoi(string(tok[1:]))
			pos = ^^pos
		}
		switch string(tok[0]) {
		case "r":
			reverseStringSlice(sigSplit)
		case "w":
			s := sigSplit[0]
			sigSplit[0] = sigSplit[pos]
			sigSplit[pos] = s
		case "s":
			sigSplit = sigSplit[pos:]
		case "p":
			sigSplit = sigSplit[pos:]
		}
	}
	return strings.Join(sigSplit, "")
}

const (
	jsvarStr   = "[a-zA-Z_\\$][a-zA-Z_0-9]*"
	reverseStr = ":function\\(a\\)\\{" +
		"(?:return )?a\\.reverse\\(\\)" +
		"\\}"
	sliceStr = ":function\\(a,b\\)\\{" +
		"return a\\.slice\\(b\\)" +
		"\\}"
	spliceStr = ":function\\(a,b\\)\\{" +
		"a\\.splice\\(0,b\\)" +
		"\\}"
	swapStr = ":function\\(a,b\\)\\{" +
		"var c=a\\[0\\];a\\[0\\]=a\\[b(?:%a\\.length)?\\];a\\[b(?:%a\\.length)?\\]=c(?:;return a)?" +
		"\\}"
)

var actionsObjRegexp = regexp.MustCompile(fmt.Sprintf(
	"var (%s)=\\{((?:(?:%s%s|%s%s|%s%s|%s%s),?\\n?)+)\\};", jsvarStr, jsvarStr, reverseStr, jsvarStr, sliceStr, jsvarStr, spliceStr, jsvarStr, swapStr))

var actionsFuncRegexp = regexp.MustCompile(fmt.Sprintf(
	"function(?: %s)?\\(a\\)\\{"+
		"a=a\\.split\\(\"\"\\);\\s*"+
		"((?:(?:a=)?%s\\.%s\\(a,\\d+\\);)+)"+
		"return a\\.join\\(\"\"\\)"+
		"\\}", jsvarStr, jsvarStr, jsvarStr))

var reverseRegexp = regexp.MustCompile(fmt.Sprintf(
	"(?m)(?:^|,)(%s)%s", jsvarStr, reverseStr))
var sliceRegexp = regexp.MustCompile(fmt.Sprintf(
	"(?m)(?:^|,)(%s)%s", jsvarStr, sliceStr))
var spliceRegexp = regexp.MustCompile(fmt.Sprintf(
	"(?m)(?:^|,)(%s)%s", jsvarStr, spliceStr))
var swapRegexp = regexp.MustCompile(fmt.Sprintf(
	"(?m)(?:^|,)(%s)%s", jsvarStr, swapStr))

func getSigTokens(htmlPlayerFile string) ([]string, error) {
	u, _ := url.Parse(youtubeBaseURL)
	p, err := url.Parse(htmlPlayerFile)
	if err != nil {
		return nil, err
	}
	resp, err := http.Get(u.ResolveReference(p).String())
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Error fetching signature tokens, status code %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	bodyString := string(body)
	if err != nil {
		return nil, err
	}

	objResult := actionsObjRegexp.FindStringSubmatch(bodyString)
	funcResult := actionsFuncRegexp.FindStringSubmatch(bodyString)

	if len(objResult) < 3 || len(funcResult) < 2 {
		return nil, fmt.Errorf("Error parsing signature tokens")
	}
	obj := strings.Replace(objResult[1], "$", "\\$", -1)
	objBody := strings.Replace(objResult[2], "$", "\\$", -1)
	funcBody := strings.Replace(funcResult[1], "$", "\\$", -1)

	var reverseKey, sliceKey, spliceKey, swapKey string
	var result []string

	if result = reverseRegexp.FindStringSubmatch(objBody); len(result) > 1 {
		reverseKey = strings.Replace(result[1], "$", "\\$", -1)
	}
	if result = sliceRegexp.FindStringSubmatch(objBody); len(result) > 1 {
		sliceKey = strings.Replace(result[1], "$", "\\$", -1)
	}
	if result = spliceRegexp.FindStringSubmatch(objBody); len(result) > 1 {
		spliceKey = strings.Replace(result[1], "$", "\\$", -1)
	}
	if result = swapRegexp.FindStringSubmatch(objBody); len(result) > 1 {
		swapKey = strings.Replace(result[1], "$", "\\$", -1)
	}

	keys := []string{reverseKey, sliceKey, spliceKey, swapKey}
	regex, err := regexp.Compile(fmt.Sprintf("(?:a=)?%s\\.(%s)\\(a,(\\d+)\\)", obj, strings.Join(keys, "|")))
	if err != nil {
		return nil, err
	}
	results := regex.FindAllStringSubmatch(funcBody, -1)
	var tokens []string
	for _, s := range results {
		switch s[1] {
		case swapKey:
			tokens = append(tokens, "w"+s[2])
		case reverseKey:
			tokens = append(tokens, "r")
		case sliceKey:
			tokens = append(tokens, "s"+s[2])
		case spliceKey:
			tokens = append(tokens, "p"+s[2])
		}
	}
	return tokens, nil
}
<commit_msg>Copy fix from fent/node-ytdl-core<commit_after>package ytdl

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
)

func getDownloadURL(format Format, htmlPlayerFile string) (*url.URL, error) {
	var sig string
	if s, ok := format.meta["s"]; ok && len(s.(string)) > 0 {
		tokens, err := getSigTokens(htmlPlayerFile)
		if err != nil {
			return nil, err
		}
		sig = decipherTokens(tokens, s.(string))
	} else {
		if s, ok := format.meta["sig"]; ok {
			sig = s.(string)
		}
	}
	var urlString string
	if s, ok := format.meta["url"]; ok {
		urlString = s.(string)
	} else if s, ok := format.meta["stream"]; ok {
		if c, ok := format.meta["conn"]; ok {
			urlString = c.(string)
			if urlString[len(urlString)-1] != '/' {
				urlString += "/"
			}
		}
		urlString += s.(string)
	} else {
		return nil, fmt.Errorf("Couldn't extract url from format")
	}
	urlString, err := url.QueryUnescape(urlString)
	if err != nil {
		return nil, err
	}
	u, err := url.Parse(urlString)
	if err != nil {
		return nil, err
	}
	query := u.Query()
	query.Set("ratebypass", "yes")
	if len(sig) > 0 {
		sigParam := "signature"
		if v, ok := format.meta["sp"].(string); ok && v != "" {
			sigParam = v
		}

		query.Set(sigParam, sig)
	}
	u.RawQuery = query.Encode()
	return u, nil
}

func decipherTokens(tokens []string, sig string) string {
	var pos int
	sigSplit := strings.Split(sig, "")
	for i, l := 0, len(tokens); i < l; i++ {
		tok := tokens[i]
		if len(tok) > 1 {
			pos, _ = strconv.Atoi(string(tok[1:]))
			pos = ^^pos
		}
		switch string(tok[0]) {
		case "r":
			reverseStringSlice(sigSplit)
		case "w":
			s := sigSplit[0]
			sigSplit[0] = sigSplit[pos]
			sigSplit[pos] = s
		case "s":
			sigSplit = sigSplit[pos:]
		case "p":
			sigSplit = sigSplit[pos:]
		}
	}
	return strings.Join(sigSplit, "")
}

const (
	jsvarStr   = "[a-zA-Z_\\$][a-zA-Z_0-9]*"
	reverseStr = ":function\\(a\\)\\{" +
		"(?:return )?a\\.reverse\\(\\)" +
		"\\}"
	sliceStr = ":function\\(a,b\\)\\{" +
		"return a\\.slice\\(b\\)" +
		"\\}"
	spliceStr = ":function\\(a,b\\)\\{" +
		"a\\.splice\\(0,b\\)" +
		"\\}"
	swapStr = ":function\\(a,b\\)\\{" +
		"var c=a\\[0\\];a\\[0\\]=a\\[b(?:%a\\.length)?\\];a\\[b(?:%a\\.length)?\\]=c(?:;return a)?" +
		"\\}"
)

var actionsObjRegexp = regexp.MustCompile(fmt.Sprintf(
	"var (%s)=\\{((?:(?:%s%s|%s%s|%s%s|%s%s),?\\n?)+)\\};", jsvarStr, jsvarStr, reverseStr, jsvarStr, sliceStr, jsvarStr, spliceStr, jsvarStr, swapStr))

var actionsFuncRegexp = regexp.MustCompile(fmt.Sprintf(
	"function(?: %s)?\\(a\\)\\{"+
		"a=a\\.split\\(\"\"\\);\\s*"+
		"((?:(?:a=)?%s\\.%s\\(a,\\d+\\);)+)"+
		"return a\\.join\\(\"\"\\)"+
		"\\}", jsvarStr, jsvarStr, jsvarStr))

var reverseRegexp = regexp.MustCompile(fmt.Sprintf(
	"(?m)(?:^|,)(%s)%s", jsvarStr, reverseStr))
var sliceRegexp = regexp.MustCompile(fmt.Sprintf(
	"(?m)(?:^|,)(%s)%s", jsvarStr, sliceStr))
var spliceRegexp = regexp.MustCompile(fmt.Sprintf(
	"(?m)(?:^|,)(%s)%s", jsvarStr, spliceStr))
var swapRegexp = regexp.MustCompile(fmt.Sprintf(
	"(?m)(?:^|,)(%s)%s", jsvarStr, swapStr))

func getSigTokens(htmlPlayerFile string) ([]string, error) {
	u, _ := url.Parse(youtubeBaseURL)
	p, err := url.Parse(htmlPlayerFile)
	if err != nil {
		return nil, err
	}
	resp, err := http.Get(u.ResolveReference(p).String())
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Error fetching signature tokens, status code %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	bodyString := string(body)
	if err != nil {
		return nil, err
	}

	objResult := actionsObjRegexp.FindStringSubmatch(bodyString)
	funcResult := actionsFuncRegexp.FindStringSubmatch(bodyString)

	if len(objResult) < 3 || len(funcResult) < 2 {
		return nil, fmt.Errorf("Error parsing signature tokens")
	}
	obj := strings.Replace(objResult[1], "$", "\\$", -1)
	objBody := strings.Replace(objResult[2], "$", "\\$", -1)
	funcBody := strings.Replace(funcResult[1], "$", "\\$", -1)

	var reverseKey, sliceKey, spliceKey, swapKey string
	var result []string

	if result = reverseRegexp.FindStringSubmatch(objBody); len(result) > 1 {
		reverseKey = strings.Replace(result[1], "$", "\\$", -1)
	}
	if result = sliceRegexp.FindStringSubmatch(objBody); len(result) > 1 {
		sliceKey = strings.Replace(result[1], "$", "\\$", -1)
	}
	if result = spliceRegexp.FindStringSubmatch(objBody); len(result) > 1 {
		spliceKey = strings.Replace(result[1], "$", "\\$", -1)
	}
	if result = swapRegexp.FindStringSubmatch(objBody); len(result) > 1 {
		swapKey = strings.Replace(result[1], "$", "\\$", -1)
	}

	keys := []string{reverseKey, sliceKey, spliceKey, swapKey}
	regex, err := regexp.Compile(fmt.Sprintf("(?:a=)?%s\\.(%s)\\(a,(\\d+)\\)", obj, strings.Join(keys, "|")))
	if err != nil {
		return nil, err
	}
	results := regex.FindAllStringSubmatch(funcBody, -1)
	var tokens []string
	for _, s := range results {
		switch s[1] {
		case swapKey:
			tokens = append(tokens, "w"+s[2])
		case reverseKey:
			tokens = append(tokens, "r")
		case sliceKey:
			tokens = append(tokens, "s"+s[2])
		case spliceKey:
			tokens = append(tokens, "p"+s[2])
		}
	}
	return tokens, nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/elazarl/goproxy"

	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"
)

const (
	appendLog int = iota
	reopenLog int = iota
)

// Digest auth. operation type
const (
	validateUser int = iota
	getNonce     int = iota
)

// Digest auth. resp status
const (
	authOk     int = iota
	authFailed int = iota
	nonceOk    int = iota
)

const proxyForwardedForHeader = "X-Forwarded-For"

type Meta struct {
	action int
	req    *http.Request
	resp   *http.Response
	err    error
	time   time.Time
}

type BasicAuthRequest struct {
	data        *BasicAuthData
	respChannel chan *BasicAuthResponse
}

type BasicAuthResponse struct {
	status bool
}

type DigestAuthRequest struct {
	data        *DigestAuthData
	op          int
	respChannel chan *DigestAuthResponse
}

type DigestAuthResponse struct {
	data   string
	status int
}

func makeBasicAuthValidator(auth *BasicAuth) BasicAuthFunc {
	channel := make(chan *BasicAuthRequest)
	validator := func() {
		for e := range channel {
			status := auth.Validate(e.data)
			e.respChannel <- &BasicAuthResponse{status}
		}
	}

	go validator()

	return func(authData *BasicAuthData) *BasicAuthResponse {
		request := &BasicAuthRequest{authData, make(chan *BasicAuthResponse)}
		channel <- request
		return <-request.respChannel
	}
}

func makeDigestAuthValidator(auth *DigestAuth) DigestAuthFunc {
	channel := make(chan *DigestAuthRequest)
	processor := func() {
		for e := range channel {
			var response *DigestAuthResponse
			switch e.op {
			case validateUser:
				status := auth.Validate(e.data)
				if status {
					response = &DigestAuthResponse{status: authOk}
				} else {
					response = &DigestAuthResponse{status: authFailed}
				}
			case getNonce:
				nonce := auth.NewNonce()
				response = &DigestAuthResponse{status: nonceOk, data: nonce}
			default:
				panic("unexpected operation type")
			}
			e.respChannel <- response
		}
	}

	go processor()

	f := func(authData *DigestAuthData, op int) *DigestAuthResponse {
		request := &DigestAuthRequest{data: authData, op: op, respChannel: make(chan *DigestAuthResponse)}
		channel <- request
		return <-request.respChannel
	}

	return f
}

func fprintf(nr *int64, err *error, w io.Writer, pat string, a ...interface{}) {
	if *err != nil {
		return
	}
	var n int
	n, *err = fmt.Fprintf(w, pat, a...)
	*nr += int64(n)
}

func write(nr *int64, err *error, w io.Writer, b []byte) {
	if *err != nil {
		return
	}
	var n int
	n, *err = w.Write(b)
	*nr += int64(n)
}

func (m *Meta) WriteTo(w io.Writer) (nr int64, err error) {
	if m.resp != nil {
		if m.resp.Request != nil {
			fprintf(&nr, &err, w,
				"%v %v %v %v %v %v\n",
				m.time.Format(time.RFC3339),
				m.resp.Request.RemoteAddr,
				m.resp.Request.Method,
				m.resp.Request.URL,
				m.resp.StatusCode,
				m.resp.ContentLength)
		} else {
			fprintf(&nr, &err, w,
				"%v %v %v %v %v %v\n",
				m.time.Format(time.RFC3339),
				"-",
				"-",
				"-",
				m.resp.StatusCode,
				m.resp.ContentLength)
		}
	} else if m.req != nil {
		fprintf(&nr, &err, w,
			"%v %v %v %v %v %v\n",
			m.time.Format(time.RFC3339),
			m.req.RemoteAddr,
			m.req.Method,
			m.req.URL,
			"-",
			"-")
	}

	return
}

type HttpLogger struct {
	path        string
	logChannel  chan *Meta
	errorChanel chan error
}

func NewLogger(conf *Configuration) *HttpLogger {
	var fh *os.File

	if conf.AccessLog != "" {
		var err error
		fh, err = os.OpenFile(conf.AccessLog, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			log.Fatalf("Couldn't open log file: %v", err)
		}
	}

	logger := &HttpLogger{conf.AccessLog, make(chan *Meta), make(chan error)}

	go func() {
		for m := range logger.logChannel {
			if fh != nil {
				switch m.action {
				case appendLog:
					if _, err := m.WriteTo(fh); err != nil {
						log.Println("Can't write meta", err)
					}
				case reopenLog:
					fh.Close()
					var err error
					fh, err = os.OpenFile(conf.AccessLog, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
					if err != nil {
						log.Fatalf("Couldn't reopen log file: %v", err)
					}
				}
			}
		}
		logger.errorChanel <- fh.Close()
	}()

	return logger
}

var emptyResp = &http.Response{}
var emptyReq = &http.Request{}

func (logger *HttpLogger) LogReq(req *http.Request, ctx *goproxy.ProxyCtx) {
	if req == nil {
		req = emptyReq
	}

	logger.LogMeta(&Meta{
		action: appendLog,
		req:    req,
		err:    ctx.Error,
		time:   time.Now()})
}

func (logger *HttpLogger) LogResp(resp *http.Response, ctx *goproxy.ProxyCtx) {
	if resp == nil {
		resp = emptyResp
	}

	logger.LogMeta(&Meta{
		action: appendLog,
		resp:   resp,
		err:    ctx.Error,
		time:   time.Now()})
}

func (logger *HttpLogger) LogMeta(m *Meta) {
	logger.logChannel <- m
}

func (logger *HttpLogger) Close() error {
	close(logger.logChannel)
	return <-logger.errorChanel
}

func (logger *HttpLogger) Reopen() {
	logger.LogMeta(&Meta{action: reopenLog})
}

func setAllowedNetworksHandler(conf *Configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.AllowedNetworks != nil && len(conf.AllowedNetworks) > 0 {
		proxy.OnRequest(goproxy.Not(sourceIpMatches(conf.AllowedNetworks))).HandleConnect(goproxy.AlwaysReject)
		proxy.OnRequest(goproxy.Not(sourceIpMatches(conf.AllowedNetworks))).DoFunc(
			func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
				return req, goproxy.NewResponse(req, goproxy.ContentTypeHtml, http.StatusForbidden, "Access denied")
			})
	}

	if conf.DisallowedNetworks != nil && len(conf.DisallowedNetworks) > 0 {
		proxy.OnRequest(sourceIpMatches(conf.DisallowedNetworks)).HandleConnect(goproxy.AlwaysReject)
		proxy.OnRequest(sourceIpMatches(conf.DisallowedNetworks)).DoFunc(
			func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
				return req, goproxy.NewResponse(req, goproxy.ContentTypeHtml, http.StatusForbidden, "Access denied")
			})
	}
}

func sourceIpMatches(networks []string) goproxy.ReqConditionFunc {
	cidrs := make([](*net.IPNet), len(networks))
	for idx, network := range networks {
		_, cidrnet, _ := net.ParseCIDR(network)
		cidrs[idx] = cidrnet
	}

	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		ip, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			ctx.Warnf("coudn't parse remote address %v: %v", req.RemoteAddr, err)
			return false
		}
		addr := net.ParseIP(ip)
		for _, network := range cidrs {
			if network.Contains(addr) {
				return true
			}
		}
		return false
	}
}

func setAllowedConnectPortsHandler(conf *Configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.AllowedConnectPorts != nil && len(conf.AllowedConnectPorts) > 0 {
		ports := make([]string, len(conf.AllowedConnectPorts))
		for i, v := range conf.AllowedConnectPorts {
			ports[i] = ":" + v
		}
		rx := "(" + strings.Join(ports, "|") + ")$"
		proxy.OnRequest(goproxy.Not(goproxy.ReqHostMatches(regexp.MustCompile(rx)))).HandleConnect(goproxy.AlwaysReject)
	}
}

func setForwardedForHeaderHandler(conf *Configuration, proxy *goproxy.ProxyHttpServer) {
	f := func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		ip, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			ctx.Warnf("coudn't parse remote address %v: %v", req.RemoteAddr, err)
			return req, nil
		}

		switch conf.ForwardedFor {
		case "on":
			header := req.Header.Get(proxyForwardedForHeader)
			if header == "" {
				req.Header.Add(proxyForwardedForHeader, ip)
			} else {
				header = header + ", " + ip
				req.Header.Del(proxyForwardedForHeader)
				req.Header.Add(proxyForwardedForHeader, header)
			}
		case "delete":
			req.Header.Del(proxyForwardedForHeader)
		case "truncate":
			req.Header.Del(proxyForwardedForHeader)
			req.Header.Add(proxyForwardedForHeader, ip)
		}

		return req, nil
	}

	proxy.OnRequest().DoFunc(f)
}

func makeCustomDial(localAddr *net.TCPAddr) func(string, string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		remoteAddr, err := net.ResolveTCPAddr(network, addr)
		if err != nil {
			return nil, err
		}
		c, err := net.DialTCP(network, localAddr, remoteAddr)
		return c, err
	}
}

func createProxy(conf *Configuration) *goproxy.ProxyHttpServer {
	proxy := goproxy.NewProxyHttpServer()
	setActivityLog(conf, proxy)

	if conf.BindIP != "" {
		var laddr string
		addressOk := true
		if ip := net.ParseIP(conf.BindIP); ip != nil {
			if ip4 := ip.To4(); ip4 != nil {
				laddr = conf.BindIP + ":0"
			} else if ip16 := ip.To16(); ip16 != nil {
				laddr = "[" + conf.BindIP + "]:0"
			} else {
				proxy.Logger.Printf("[WARN] couldn't use \"%v\" as outgoing request address.\n", conf.BindIP)
				addressOk = false
			}
		}
		if addressOk {
			if addr, err := net.ResolveTCPAddr("tcp", laddr); err == nil {
				proxy.Tr.Dial = makeCustomDial(addr)
			} else {
				proxy.Logger.Printf("[WARN] couldn't use \"%v\" as outgoing request address. %v\n", conf.BindIP, err)
			}
		}
	}

	return proxy
}

func setActivityLog(conf *Configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.ActivityLog != "" {
		fh, err := os.OpenFile(conf.ActivityLog, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			log.Fatalf("couldn't open activity log file %v: %v", conf.ActivityLog, err)
		}
		proxy.Logger = log.New(fh, "", log.LstdFlags)
	}
}

func main() {
	config := flag.String("config", "microproxy.json", "proxy configuration file")
	verbose := flag.Bool("v", false, "enable verbose debug mode")

	flag.Parse()

	conf := NewConfiguration(*config)

	proxy := createProxy(conf)
	proxy.Verbose = *verbose

	logger := NewLogger(conf)

	if conf.AuthFile != "" {
		if conf.AuthType == "basic" {
			auth, err := NewBasicAuthFromFile(conf.AuthFile)
			if err != nil {
				proxy.Logger.Fatalf("couldn't create basic auth structure: %v\n", err)
			}
			ProxyBasic(proxy, conf.AuthRealm, makeBasicAuthValidator(auth))
		} else {
			auth, err := NewDigestAuthFromFile(conf.AuthFile)
			if err != nil {
				proxy.Logger.Fatalf("couldn't create digest auth structure: %v\n", err)
			}
			ProxyDigest(proxy, conf.AuthRealm, makeDigestAuthValidator(auth))
		}
	}

	proxy.OnRequest().HandleConnectFunc(
		func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
			if ctx.Req == nil {
				ctx.Req = emptyReq
			}

			logger.LogMeta(&Meta{
				action: appendLog,
				req:    ctx.Req,
				err:    ctx.Error,
				time:   time.Now()})

			return goproxy.OkConnect, host
		})

	setAllowedConnectPortsHandler(conf, proxy)
	setAllowedNetworksHandler(conf, proxy)
	setForwardedForHeaderHandler(conf, proxy)

	proxy.OnResponse().DoFunc(
		func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
			logger.LogResp(resp, ctx)
			return resp
		})

	signalChannel := make(chan os.Signal)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGUSR1)

	go func() {
		for {
			sig := <-signalChannel
			switch sig {
			case os.Interrupt:
				proxy.Logger.Println("got interrupt signal, exiting")
				logger.Close()
				os.Exit(0)
			case syscall.SIGUSR1:
				proxy.Logger.Println("got USR1 signal, reopening logs")
				// Reopen access log
				logger.Reopen()
				// Reopen activity log
				setActivityLog(conf, proxy)
			}
		}
	}()

	proxy.Logger.Println("starting proxy")

	log.Fatal(http.ListenAndServe(conf.Listen, proxy))
}
<commit_msg>go fmt<commit_after>package main

import (
	"github.com/elazarl/goproxy"

	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"
)

const (
	appendLog int = iota
	reopenLog int = iota
)

// Digest auth. operation type
const (
	validateUser int = iota
	getNonce     int = iota
)

// Digest auth. resp status
const (
	authOk     int = iota
	authFailed int = iota
	nonceOk    int = iota
)

const proxyForwardedForHeader = "X-Forwarded-For"

type Meta struct {
	action int
	req    *http.Request
	resp   *http.Response
	err    error
	time   time.Time
}

type BasicAuthRequest struct {
	data        *BasicAuthData
	respChannel chan *BasicAuthResponse
}

type BasicAuthResponse struct {
	status bool
}

type DigestAuthRequest struct {
	data        *DigestAuthData
	op          int
	respChannel chan *DigestAuthResponse
}

type DigestAuthResponse struct {
	data   string
	status int
}

func makeBasicAuthValidator(auth *BasicAuth) BasicAuthFunc {
	channel := make(chan *BasicAuthRequest)
	validator := func() {
		for e := range channel {
			status := auth.Validate(e.data)
			e.respChannel <- &BasicAuthResponse{status}
		}
	}

	go validator()

	return func(authData *BasicAuthData) *BasicAuthResponse {
		request := &BasicAuthRequest{authData, make(chan *BasicAuthResponse)}
		channel <- request
		return <-request.respChannel
	}
}

func makeDigestAuthValidator(auth *DigestAuth) DigestAuthFunc {
	channel := make(chan *DigestAuthRequest)
	processor := func() {
		for e := range channel {
			var response *DigestAuthResponse
			switch e.op {
			case validateUser:
				status := auth.Validate(e.data)
				if status {
					response = &DigestAuthResponse{status: authOk}
				} else {
					response = &DigestAuthResponse{status: authFailed}
				}
			case getNonce:
				nonce := auth.NewNonce()
				response = &DigestAuthResponse{status: nonceOk, data: nonce}
			default:
				panic("unexpected operation type")
			}
			e.respChannel <- response
		}
	}

	go processor()

	f := func(authData *DigestAuthData, op int) *DigestAuthResponse {
		request := &DigestAuthRequest{data: authData, op: op, respChannel: make(chan *DigestAuthResponse)}
		channel <- request
		return <-request.respChannel
	}

	return f
}

func fprintf(nr *int64, err *error, w io.Writer, pat string, a ...interface{}) {
	if *err != nil {
		return
	}
	var n int
	n, *err = fmt.Fprintf(w, pat, a...)
	*nr += int64(n)
}

func write(nr *int64, err *error, w io.Writer, b []byte) {
	if *err != nil {
		return
	}
	var n int
	n, *err = w.Write(b)
	*nr += int64(n)
}

func (m *Meta) WriteTo(w io.Writer) (nr int64, err error) {
	if m.resp != nil {
		if m.resp.Request != nil {
			fprintf(&nr, &err, w,
				"%v %v %v %v %v %v\n",
				m.time.Format(time.RFC3339),
				m.resp.Request.RemoteAddr,
				m.resp.Request.Method,
				m.resp.Request.URL,
				m.resp.StatusCode,
				m.resp.ContentLength)
		} else {
			fprintf(&nr, &err, w,
				"%v %v %v %v %v %v\n",
				m.time.Format(time.RFC3339),
				"-",
				"-",
				"-",
				m.resp.StatusCode,
				m.resp.ContentLength)
		}
	} else if m.req != nil {
		fprintf(&nr, &err, w,
			"%v %v %v %v %v %v\n",
			m.time.Format(time.RFC3339),
			m.req.RemoteAddr,
			m.req.Method,
			m.req.URL,
			"-",
			"-")
	}

	return
}

type HttpLogger struct {
	path        string
	logChannel  chan *Meta
	errorChanel chan error
}

func NewLogger(conf *Configuration) *HttpLogger {
	var fh *os.File

	if conf.AccessLog != "" {
		var err error
		fh, err = os.OpenFile(conf.AccessLog, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			log.Fatalf("Couldn't open log file: %v", err)
		}
	}

	logger := &HttpLogger{conf.AccessLog, make(chan *Meta), make(chan error)}

	go func() {
		for m := range logger.logChannel {
			if fh != nil {
				switch m.action {
				case appendLog:
					if _, err := m.WriteTo(fh); err != nil {
						log.Println("Can't write meta", err)
					}
				case reopenLog:
					fh.Close()
					var err error
					fh, err = os.OpenFile(conf.AccessLog, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
					if err != nil {
						log.Fatalf("Couldn't reopen log file: %v", err)
					}
				}
			}
		}
		logger.errorChanel <- fh.Close()
	}()

	return logger
}

var emptyResp = &http.Response{}
var emptyReq = &http.Request{}

func (logger *HttpLogger) LogReq(req *http.Request, ctx *goproxy.ProxyCtx) {
	if req == nil {
		req = emptyReq
	}

	logger.LogMeta(&Meta{
		action: appendLog,
		req:    req,
		err:    ctx.Error,
		time:   time.Now()})
}

func (logger *HttpLogger) LogResp(resp *http.Response, ctx *goproxy.ProxyCtx) {
	if resp == nil {
		resp = emptyResp
	}

	logger.LogMeta(&Meta{
		action: appendLog,
		resp:   resp,
		err:    ctx.Error,
		time:   time.Now()})
}

func (logger *HttpLogger) LogMeta(m *Meta) {
	logger.logChannel <- m
}

func (logger *HttpLogger) Close() error {
	close(logger.logChannel)
	return <-logger.errorChanel
}

func (logger *HttpLogger) Reopen() {
	logger.LogMeta(&Meta{action: reopenLog})
}

func setAllowedNetworksHandler(conf *Configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.AllowedNetworks != nil && len(conf.AllowedNetworks) > 0 {
		proxy.OnRequest(goproxy.Not(sourceIpMatches(conf.AllowedNetworks))).HandleConnect(goproxy.AlwaysReject)
		proxy.OnRequest(goproxy.Not(sourceIpMatches(conf.AllowedNetworks))).DoFunc(
			func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
				return req, goproxy.NewResponse(req, goproxy.ContentTypeHtml, http.StatusForbidden, "Access denied")
			})
	}

	if conf.DisallowedNetworks != nil && len(conf.DisallowedNetworks) > 0 {
		proxy.OnRequest(sourceIpMatches(conf.DisallowedNetworks)).HandleConnect(goproxy.AlwaysReject)
		proxy.OnRequest(sourceIpMatches(conf.DisallowedNetworks)).DoFunc(
			func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
				return req, goproxy.NewResponse(req, goproxy.ContentTypeHtml, http.StatusForbidden, "Access denied")
			})
	}
}

func sourceIpMatches(networks []string) goproxy.ReqConditionFunc {
	cidrs := make([](*net.IPNet), len(networks))
	for idx, network := range networks {
		_, cidrnet, _ := net.ParseCIDR(network)
		cidrs[idx] = cidrnet
	}

	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		ip, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			ctx.Warnf("coudn't parse remote address %v: %v", req.RemoteAddr, err)
			return false
		}
		addr := net.ParseIP(ip)
		for _, network := range cidrs {
			if network.Contains(addr) {
				return true
			}
		}
		return false
	}
}

func setAllowedConnectPortsHandler(conf *Configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.AllowedConnectPorts != nil && len(conf.AllowedConnectPorts) > 0 {
		ports := make([]string, len(conf.AllowedConnectPorts))
		for i, v := range conf.AllowedConnectPorts {
			ports[i] = ":" + v
		}
		rx := "(" + strings.Join(ports, "|") + ")$"
		proxy.OnRequest(goproxy.Not(goproxy.ReqHostMatches(regexp.MustCompile(rx)))).HandleConnect(goproxy.AlwaysReject)
	}
}

func setForwardedForHeaderHandler(conf *Configuration, proxy *goproxy.ProxyHttpServer) {
	f := func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		ip, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			ctx.Warnf("coudn't parse remote address %v: %v", req.RemoteAddr, err)
			return req, nil
		}

		switch conf.ForwardedFor {
		case "on":
			header := req.Header.Get(proxyForwardedForHeader)
			if header == "" {
				req.Header.Add(proxyForwardedForHeader, ip)
			} else {
				header = header + ", " + ip
				req.Header.Del(proxyForwardedForHeader)
				req.Header.Add(proxyForwardedForHeader, header)
			}
		case "delete":
			req.Header.Del(proxyForwardedForHeader)
		case "truncate":
			req.Header.Del(proxyForwardedForHeader)
			req.Header.Add(proxyForwardedForHeader, ip)
		}

		return req, nil
	}

	proxy.OnRequest().DoFunc(f)
}

func makeCustomDial(localAddr *net.TCPAddr) func(string, string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		remoteAddr, err := net.ResolveTCPAddr(network, addr)
		if err != nil {
			return nil, err
		}
		c, err := net.DialTCP(network, localAddr, remoteAddr)
		return c, err
	}
}

func createProxy(conf *Configuration) *goproxy.ProxyHttpServer {
	proxy := goproxy.NewProxyHttpServer()
	setActivityLog(conf, proxy)

	if conf.BindIP != "" {
		var laddr string
		addressOk := true
		if ip := net.ParseIP(conf.BindIP); ip != nil {
			if ip4 := ip.To4(); ip4 != nil {
				laddr = conf.BindIP + ":0"
			} else if ip16 := ip.To16(); ip16 != nil {
				laddr = "[" + conf.BindIP + "]:0"
			} else {
				proxy.Logger.Printf("[WARN] couldn't use \"%v\" as outgoing request address.\n", conf.BindIP)
				addressOk = false
			}
		}
		if addressOk {
			if addr, err := net.ResolveTCPAddr("tcp", laddr); err == nil {
				proxy.Tr.Dial = makeCustomDial(addr)
			} else {
				proxy.Logger.Printf("[WARN] couldn't use \"%v\" as outgoing request address. %v\n", conf.BindIP, err)
			}
		}
	}

	return proxy
}

func setActivityLog(conf *Configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.ActivityLog != "" {
		fh, err := os.OpenFile(conf.ActivityLog, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			log.Fatalf("couldn't open activity log file %v: %v", conf.ActivityLog, err)
		}
		proxy.Logger = log.New(fh, "", log.LstdFlags)
	}
}

func main() {
	config := flag.String("config", "microproxy.json", "proxy configuration file")
	verbose := flag.Bool("v", false, "enable verbose debug mode")

	flag.Parse()

	conf := NewConfiguration(*config)

	proxy := createProxy(conf)
	proxy.Verbose = *verbose

	logger := NewLogger(conf)

	if conf.AuthFile != "" {
		if conf.AuthType == "basic" {
			auth, err := NewBasicAuthFromFile(conf.AuthFile)
			if err != nil {
				proxy.Logger.Fatalf("couldn't create basic auth structure: %v\n", err)
			}
			ProxyBasic(proxy, conf.AuthRealm, makeBasicAuthValidator(auth))
		} else {
			auth, err := NewDigestAuthFromFile(conf.AuthFile)
			if err != nil {
				proxy.Logger.Fatalf("couldn't create digest auth structure: %v\n", err)
			}
			ProxyDigest(proxy, conf.AuthRealm, makeDigestAuthValidator(auth))
		}
	}

	proxy.OnRequest().HandleConnectFunc(
		func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
			if ctx.Req == nil {
				ctx.Req = emptyReq
			}

			logger.LogMeta(&Meta{
				action: appendLog,
				req:    ctx.Req,
				err:    ctx.Error,
				time:   time.Now()})

			return goproxy.OkConnect, host
		})

	setAllowedConnectPortsHandler(conf, proxy)
	setAllowedNetworksHandler(conf, proxy)
	setForwardedForHeaderHandler(conf, proxy)

	proxy.OnResponse().DoFunc(
		func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
			logger.LogResp(resp, ctx)
			return resp
		})

	signalChannel := make(chan os.Signal)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGUSR1)

	go func() {
		for {
			sig := <-signalChannel
			switch sig {
			case os.Interrupt:
				proxy.Logger.Println("got interrupt signal, exiting")
				logger.Close()
				os.Exit(0)
			case syscall.SIGUSR1:
				proxy.Logger.Println("got USR1 signal, reopening logs")
				// Reopen access log
				logger.Reopen()
				// Reopen activity log
				setActivityLog(conf, proxy)
			}
		}
	}()

	proxy.Logger.Println("starting proxy")

	log.Fatal(http.ListenAndServe(conf.Listen, proxy))
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"
	"strings"

	"github.com/DimShadoWWW/power-pg/common"
	"github.com/op/go-logging"
)

var (
	connid = uint64(0)
)

// Pkg PostgreSQL package structure
type Pkg struct {
	Type    byte
	Content []byte
}

// Start function
func Start(localHost, remoteHost *string, remotePort *string, msgBytes chan []byte, msgCh chan Pkg, recreate bool, log *logging.Logger) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost, remotePort)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
			log:    log,
		}
		go p.start(msgBytes, msgCh, recreate)
	}
}

func getResolvedAddresses(localHost, remoteHost, remotePort *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", *remoteHost, *remotePort))
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
	result        *[]string
	log           *logging.Logger
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(msgBytes chan []byte, msgCh chan Pkg, recreate bool) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, msgBytes, msgCh, recreate, p.log)
	go p.pipe(p.rconn, p.lconn, nil, nil, recreate, p.log)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, msgBytes chan []byte, msgCh chan Pkg, recreate bool, log *logging.Logger) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(ReadBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	// spaces := regexp.MustCompile("[\n\t ]+")
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r ReadBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			if msgBytes != nil {
				log.Debug("Readed bytes: %d\n", n)
			}
			b := buff[:n]
			log.Info("Readed: %v\n", b)
			msgBytes <- b
			//write out result
			if !recreate {
				n, err = dst.Write(b)
				if err != nil {
					p.err("Write failed '%s'\n", err)
					return
				}
			}
			r = buff[:n]
			log.Debug("PostgreSQL full message: %s\n", string(r))
			// if msgCh != nil {
			// 						msgCh <- 	fmt.Sprintf("%#v", string(buff[:n]))}
			if msgBytes != nil {
				log.Debug("Remaining bytes: %d\n", remainingBytes)
			}
			if msgBytes != nil {
				log.Debug("newPacket : %v\n", newPacket)
			}
			if msgBytes != nil {
				log.Debug("len(r) : %v\n", len(r))
			}
			fmt.Println("3")
			// NewP:
			fmt.Println("4")
			if newPacket || (len(msg) > 4 && len(r) > 4 && remainingBytes == 0) {
				fmt.Println("5")
				// remainingBytes = 0
				newPacket = false
				if msgBytes != nil && msg != "" {
					log.Debug("2 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
				}
				var msg []byte
				log.Debug("1 n: %d\n", n)
				t := r.Byte()
				n = n - 1
				log.Debug("2 n: %d\n", n)
				fmt.Println("t: ", string(t))
				switch t {
				case 'Q', 'B', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'P', 'p', 'S', 'X':
					// case 'B', 'P':
					// c.rxReadyForQuery(r)
					log.Debug("PostgreSQL pkg type: %s\n", string(t))
					remainingBytes = r.Int32()
					remainingBytes = remainingBytes - 4
					r = r[:remainingBytes]
					n = remainingBytes
					if msgCh != nil {
						log.Debug("1 Remaining bytes: %d \tn: %d\n", remainingBytes, n)
					}
					if remainingBytes < 4 {
						fmt.Println("ERROR: remainingBytes can't be less than 4 bytes if int32")
					} else {
						log.Debug("1 r: %v\n", msg)
						log.Debug("1 string(r): %s\n", string(msg))
						s := strings.Index(string(msg), string([]byte{0}))
						n = remainingBytes - s
						msg = msg[:]
						log.Debug("2 r: %v\n", msg)
						log.Debug("2 string(r): %s\n", string(msg))

						if msgCh != nil {
							log.Debug("2 Remaining bytes: %d \tn: %d\n", remainingBytes, n)
						}
						if remainingBytes > 0 {
							log.Debug("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))

							newPacket = true
							msg = append(msg, r.Next(remainingBytes)[:]...)
							// msg = spaces.ReplaceAll(msg, []byte{' '})
							remainingBytes = n - remainingBytes
							if msgBytes != nil {
								log.Debug("3.1 Remaining bytes: %d \tmsg: %v\n", remainingBytes, msg)
							}

							if msgCh != nil {
								msgCh <- Pkg{
									Type:    t,
									Content: msg,
								}
							}
							remainingBytes = 0
							// goto NewP
							// } else {
							// 	newPacket = false
							// 	msg = append(msg, r.Next(remainingBytes)[:]...)
							// 	// msg = bytes.Replace(msg, []byte("\n\t"), []byte(" "), -1)
							// 	msg = spaces.ReplaceAll(msg, []byte{' '})
							// 	// msg = []byte(stripchars(string(msg),
							// 	// 	"\n\t"))
							// 	remainingBytes = remainingBytes - n
							// 	if msgBytes != nil {
							// 		log.Debug("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							// 	}
						}
					}
					// case :
					// 	fmt.Println("TODO")
					// 	// c.rxReadyForQuery(r)
					// 	remainingBytes = r.int32()
					// 	remainingBytes = remainingBytes - 4
					// 	if remainingBytes > 0 {
					// 		if remainingBytes <= n {
					// 			newPacket = true
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = n - remainingBytes
					// 			if msgCh != nil {
					// msgCh <- 	fmt.Sprintf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))}
					// 			// fmt.Println(msg)
					// 			goto NewP
					// 		} else {
					// 			newPacket = false
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = remainingBytes - n
					// 			if msgCh != nil {
					// msgCh <- 	fmt.Sprintf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))}
					// 		}
					// 	}
					// case rowDescription:
					// case dataRow:
					// case bindComplete:
					// case commandComplete:
					// 	commandTag = CommandTag(r.readCString())
				// case 'Q', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'p', 'S', 'X':
				default:
					fmt.Println("6")
					remainingBytes = 0
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
				remainingBytes = 0
			} else {
				fmt.Println("7")
				remainingBytes = 0
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// if msgCh != nil {
			// msgCh <- 	fmt.Sprintf("t: %#v\n", t)}

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}

func stripchars(str, chr string) string {
	return strings.Map(func(r rune) rune {
		if strings.IndexRune(chr, r) < 0 {
			return r
		}
		return -1
	}, str)
}
<commit_msg>Update<commit_after>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"
	"strings"

	"github.com/DimShadoWWW/power-pg/common"
	"github.com/op/go-logging"
)

var (
	connid = uint64(0)
)

// Pkg PostgreSQL package structure
type Pkg struct {
	Type    byte
	Content []byte
}

// Start function
func Start(localHost, remoteHost *string, remotePort *string, msgBytes chan []byte, msgCh chan Pkg, recreate bool, log *logging.Logger) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost, remotePort)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
			log:    log,
		}
		go p.start(msgBytes, msgCh, recreate)
	}
}

func getResolvedAddresses(localHost, remoteHost, remotePort *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", *remoteHost, *remotePort))
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
	result        *[]string
	log           *logging.Logger
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(msgBytes chan []byte, msgCh chan Pkg, recreate bool) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, msgBytes, msgCh, recreate, p.log)
	go p.pipe(p.rconn, p.lconn, nil, nil, recreate, p.log)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, msgBytes chan []byte, msgCh chan Pkg, recreate bool, log *logging.Logger) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(ReadBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	// spaces := regexp.MustCompile("[\n\t ]+")
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r ReadBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			if msgBytes != nil {
				log.Debug("Readed bytes: %d\n", n)
			}
			b := buff[:n]
			log.Info("Readed: %v\n", b)
			msgBytes <- b
			//write out result
			if !recreate {
				n, err = dst.Write(b)
				if err != nil {
					p.err("Write failed '%s'\n", err)
					return
				}
			}
			r = buff[:n]
			log.Debug("PostgreSQL full message: %s\n", string(r))
			// if msgCh != nil {
			// 						msgCh <- 	fmt.Sprintf("%#v", string(buff[:n]))}
			if msgBytes != nil {
				log.Debug("Remaining bytes: %d\n", remainingBytes)
			}
			if msgBytes != nil {
				log.Debug("newPacket : %v\n", newPacket)
			}
			if msgBytes != nil {
				log.Debug("len(r) : %v\n", len(r))
			}
			fmt.Println("3")
			// NewP:
			fmt.Println("4")
			if newPacket || (len(msg) > 4 && len(r) > 4 && remainingBytes == 0) {
				fmt.Println("5")
				// remainingBytes = 0
				newPacket = false
				if msgBytes != nil && msg != "" {
					log.Debug("2 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
				}
				var msg []byte
				log.Debug("1 n: %d\n", n)
				t := r.Byte()
				n = n - 1
				log.Debug("2 n: %d\n", n)
				fmt.Println("t: ", string(t))
				switch t {
				case 'Q', 'B', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'P', 'p', 'S', 'X':
					// case 'B', 'P':
					// c.rxReadyForQuery(r)
					log.Debug("PostgreSQL pkg type: %s\n", string(t))
					remainingBytes = r.Int32()
					remainingBytes = remainingBytes - 4
					r = r[:remainingBytes]
					n = remainingBytes
					if msgCh != nil {
						log.Debug("1 Remaining bytes: %d \tn: %d\n", remainingBytes, n)
					}
					if remainingBytes < 4 {
						fmt.Println("ERROR: remainingBytes can't be less than 4 bytes if int32")
					} else {
						log.Debug("1 r: %v\n", r)
						log.Debug("1 string(r): %s\n", string(r))
						s := strings.Index(string(r), string([]byte{0}))
						n = remainingBytes - s
						r = r[s:]
						log.Debug("2 r: %v\n", r)
						log.Debug("2 string(r): %s\n", string(r))

						if msgCh != nil {
							log.Debug("2 Remaining bytes: %d \tn: %d\n", remainingBytes, n)
						}
						if remainingBytes > 0 {
							log.Debug("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))

							newPacket = true
							msg = append(msg, r.Next(remainingBytes)[:]...)
							// msg = spaces.ReplaceAll(msg, []byte{' '})
							remainingBytes = n - remainingBytes
							if msgBytes != nil {
								log.Debug("3.1 Remaining bytes: %d \tmsg: %v\n", remainingBytes, msg)
							}

							if msgCh != nil {
								msgCh <- Pkg{
									Type:    t,
									Content: msg,
								}
							}
							remainingBytes = 0
							// goto NewP
							// } else {
							// 	newPacket = false
							// 	msg = append(msg, r.Next(remainingBytes)[:]...)
							// 	// msg = bytes.Replace(msg, []byte("\n\t"), []byte(" "), -1)
							// 	msg = spaces.ReplaceAll(msg, []byte{' '})
							// 	// msg = []byte(stripchars(string(msg),
							// 	// 	"\n\t"))
							// 	remainingBytes = remainingBytes - n
							// 	if msgBytes != nil {
							// 		log.Debug("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							// 	}
						}
					}
					// case :
					// 	fmt.Println("TODO")
					// 	// c.rxReadyForQuery(r)
					// 	remainingBytes = r.int32()
					// 	remainingBytes = remainingBytes - 4
					// 	if remainingBytes > 0 {
					// 		if remainingBytes <= n {
					// 			newPacket = true
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = n - remainingBytes
					// 			if msgCh != nil {
					// msgCh <- 	fmt.Sprintf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))}
					// 			// fmt.Println(msg)
					// 			goto NewP
					// 		} else {
					// 			newPacket = false
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = remainingBytes - n
					// 			if msgCh != nil {
					// msgCh <- 	fmt.Sprintf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))}
					// 		}
					// 	}
					// case rowDescription:
					// case dataRow:
					// case bindComplete:
					// case commandComplete:
					// 	commandTag = CommandTag(r.readCString())
				// case 'Q', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'p', 'S', 'X':
				default:
					fmt.Println("6")
					remainingBytes = 0
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
				remainingBytes = 0
			} else {
				fmt.Println("7")
				remainingBytes = 0
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// if msgCh != nil {
			// msgCh <- 	fmt.Sprintf("t: %#v\n", t)}

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}

func stripchars(str, chr string) string {
	return strings.Map(func(r rune) rune {
		if strings.IndexRune(chr, r) < 0 {
			return r
		}
		return -1
	}, str)
}
<|endoftext|>


<commit_before>package proxy

import (
	"strings"

	log "github.com/Sirupsen/logrus"
)

// Proxy defines proxy rules for a route
type Proxy struct {
	PreserveHostHeader          bool     `bson:"preserve_host_header" json:"preserve_host_header"`
	ListenPath                  string   `bson:"listen_path" json:"listen_path" valid:"required"`
	TargetURL                   string   `bson:"target_url" json:"target_url" valid:"url,required"`
	StripListenPath             bool     `bson:"strip_listen_path" json:"strip_listen_path"`
	AppendListenPath            bool     `bson:"append_listen_path" json:"append_listen_path"`
	EnableLoadBalancing         bool     `bson:"enable_load_balancing" json:"enable_load_balancing"`
	TargetList                  []string `bson:"target_list" json:"target_list"`
	CheckHostAgainstUptimeTests bool     `bson:"check_host_against_uptime_tests" json:"check_host_against_uptime_tests"`
	Methods                     []string `bson:"methods" json:"methods"`
}

//Validate validates proxy data
func Validate(proxy Proxy) bool {
	if proxy.ListenPath == "" {
		log.Warning("Listen path is empty")
		return false
	}

	if strings.Contains(proxy.ListenPath, " ") {
		log.Warning("Listen path contains spaces, is invalid")
		return false
	}

	return true
}
<commit_msg>Added new Route struct<commit_after>package proxy

import (
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/hellofresh/janus/router"
)

// Route is the container for a proxy and it's handlers
type Route struct {
	proxy    Proxy
	handlers []router.Constructor
}

// NewRoute creates an instance of Route
func NewRoute(proxy Proxy, handlers ...router.Constructor) *Route {
	return &Route{proxy, handlers}
}

// Proxy defines proxy rules for a route
type Proxy struct {
	PreserveHostHeader          bool     `bson:"preserve_host_header" json:"preserve_host_header"`
	ListenPath                  string   `bson:"listen_path" json:"listen_path" valid:"required"`
	TargetURL                   string   `bson:"target_url" json:"target_url" valid:"url,required"`
	StripListenPath             bool     `bson:"strip_listen_path" json:"strip_listen_path"`
	AppendListenPath            bool     `bson:"append_listen_path" json:"append_listen_path"`
	EnableLoadBalancing         bool     `bson:"enable_load_balancing" json:"enable_load_balancing"`
	TargetList                  []string `bson:"target_list" json:"target_list"`
	CheckHostAgainstUptimeTests bool     `bson:"check_host_against_uptime_tests" json:"check_host_against_uptime_tests"`
	Methods                     []string `bson:"methods" json:"methods"`
}

// Validate validates proxy data
func Validate(proxy Proxy) bool {
	if proxy.ListenPath == "" {
		log.Warning("Listen path is empty")
		return false
	}

	if strings.Contains(proxy.ListenPath, " ") {
		log.Warning("Listen path contains spaces, is invalid")
		return false
	}

	return true
}
<|endoftext|>


<commit_before>package proxy

import (
	"fmt"
	"net/url"
	"sort"

	"github.com/mailgun/vulcand/Godeps/_workspace/src/github.com/mailgun/oxy/memmetrics"
	"github.com/mailgun/vulcand/engine"
)

func (mx *mux) frontendStats(key engine.FrontendKey) (*engine.RoundTripStats, error) {
	f, ok := mx.frontends[key]
	if !ok {
		return nil, fmt.Errorf("%v not found", key)
	}
	return f.watcher.rtStats()
}

func (mx *mux) backendStats(key engine.BackendKey) (*engine.RoundTripStats, error) {
	m, err := memmetrics.NewRTMetrics()
	if err != nil {
		return nil, err
	}
	for _, f := range mx.frontends {
		if f.backend.backend.Id != key.Id {
			continue
		}
		if err := f.watcher.collectMetrics(m); err != nil {
			return nil, err
		}
	}
	return engine.NewRoundTripStats(m)
}

func (mx *mux) serverStats(key engine.ServerKey) (*engine.RoundTripStats, error) {
	b, ok := mx.backends[key.BackendKey]
	if !ok {
		return nil, fmt.Errorf("%v not found", key.BackendKey)
	}
	srv, ok := b.findServer(key)
	if !ok {
		return nil, fmt.Errorf("%v not found", key)
	}

	u, err := url.Parse(srv.URL)
	if err != nil {
		return nil, err
	}

	m, err := memmetrics.NewRTMetrics()
	if err != nil {
		return nil, err
	}
	for _, f := range mx.frontends {
		if f.backend.backend.Id != key.BackendKey.Id {
			continue
		}
		if err := f.watcher.collectServerMetrics(m, u); err != nil {
			return nil, err
		}
	}
	return engine.NewRoundTripStats(m)
}

func (mx *mux) topFrontends(key *engine.BackendKey) ([]engine.Frontend, error) {
	frontends := []engine.Frontend{}
	for _, m := range mx.frontends {
		if key != nil && key.Id != m.backend.backend.Id {
			continue
		}
		f := m.frontend
		stats, err := m.watcher.rtStats()
		if err != nil {
			return nil, err
		}
		f.Stats = stats
		frontends = append(frontends, f)
	}
	sort.Stable(&frontendSorter{frontends: frontends})
	return frontends, nil
}

func (mx *mux) topServers(key *engine.BackendKey) ([]engine.Server, error) {
	servers := []engine.Server{}
	for _, f := range mx.frontends {
		if key != nil && key.Id != f.backend.backend.Id {
			continue
		}
		for _, s := range f.backend.servers {
			u, err := url.Parse(s.URL)
			if err != nil {
				return nil, err
			}
			stats, err := f.watcher.rtServerStats(u)
			if err != nil {
				return nil, err
			}
			s.Stats = stats
			servers = append(servers, s)
		}
	}
	sort.Stable(&serverSorter{es: servers})
	return servers, nil
}

type frontendSorter struct {
	frontends []engine.Frontend
}

func (s *frontendSorter) Len() int {
	return len(s.frontends)
}

func (s *frontendSorter) Swap(i, j int) {
	s.frontends[i], s.frontends[j] = s.frontends[j], s.frontends[i]
}

func (s *frontendSorter) Less(i, j int) bool {
	return cmpStats(s.frontends[i].Stats, s.frontends[j].Stats)
}

type serverSorter struct {
	es []engine.Server
}

func (s *serverSorter) Len() int {
	return len(s.es)
}

func (s *serverSorter) Swap(i, j int) {
	s.es[i], s.es[j] = s.es[j], s.es[i]
}

func (s *serverSorter) Less(i, j int) bool {
	return cmpStats(s.es[i].Stats, s.es[j].Stats)
}

func cmpStats(s1, s2 *engine.RoundTripStats) bool {
	// Items that have network errors go first
	if s1.NetErrorRatio() != 0 || s2.NetErrorRatio() != 0 {
		return s1.NetErrorRatio() > s2.NetErrorRatio()
	}

	// Items that have application level errors go next
	if s1.AppErrorRatio() != 0 || s2.AppErrorRatio() != 0 {
		return s1.AppErrorRatio() > s2.AppErrorRatio()
	}

	// More highly loaded items go next
	return s1.Counters.Total > s2.Counters.Total
}
<commit_msg>Fix reporting<commit_after>package proxy

import (
	"fmt"
	"net/url"
	"sort"

	"github.com/mailgun/vulcand/Godeps/_workspace/src/github.com/mailgun/oxy/memmetrics"
	"github.com/mailgun/vulcand/engine"
)

func (mx *mux) frontendStats(key engine.FrontendKey) (*engine.RoundTripStats, error) {
	f, ok := mx.frontends[key]
	if !ok {
		return nil, fmt.Errorf("%v not found", key)
	}
	return f.watcher.rtStats()
}

func (mx *mux) backendStats(key engine.BackendKey) (*engine.RoundTripStats, error) {
	m, err := memmetrics.NewRTMetrics()
	if err != nil {
		return nil, err
	}
	for _, f := range mx.frontends {
		if f.backend.backend.Id != key.Id {
			continue
		}
		if err := f.watcher.collectMetrics(m); err != nil {
			return nil, err
		}
	}
	return engine.NewRoundTripStats(m)
}

func (mx *mux) serverStats(key engine.ServerKey) (*engine.RoundTripStats, error) {
	b, ok := mx.backends[key.BackendKey]
	if !ok {
		return nil, fmt.Errorf("%v not found", key.BackendKey)
	}
	srv, ok := b.findServer(key)
	if !ok {
		return nil, fmt.Errorf("%v not found", key)
	}

	u, err := url.Parse(srv.URL)
	if err != nil {
		return nil, err
	}

	m, err := memmetrics.NewRTMetrics()
	if err != nil {
		return nil, err
	}
	for _, f := range mx.frontends {
		if f.backend.backend.Id != key.BackendKey.Id {
			continue
		}
		if err := f.watcher.collectServerMetrics(m, u); err != nil {
			return nil, err
		}
	}
	return engine.NewRoundTripStats(m)
}

func (mx *mux) topFrontends(key *engine.BackendKey) ([]engine.Frontend, error) {
	frontends := []engine.Frontend{}
	for _, m := range mx.frontends {
		if key != nil && key.Id != m.backend.backend.Id {
			continue
		}
		f := m.frontend
		stats, err := m.watcher.rtStats()
		if err != nil {
			return nil, err
		}
		f.Stats = stats
		frontends = append(frontends, f)
	}
	sort.Stable(&frontendSorter{frontends: frontends})
	return frontends, nil
}

func (mx *mux) topServers(key *engine.BackendKey) ([]engine.Server, error) {
	metrics := map[string]*sval{}
	for _, f := range mx.frontends {
		if key != nil && key.Id != f.backend.backend.Id {
			continue
		}
		for _, s := range f.backend.servers {
			val, ok := metrics[s.URL]
			if !ok {
				sval, err := newSval(s)
				if err != nil {
					return nil, err
				}
				metrics[s.URL] = sval
				val = sval
			}
			if err := f.watcher.collectServerMetrics(val.m, val.u); err != nil {
				return nil, err
			}
		}
	}
	servers := make([]engine.Server, 0, len(metrics))
	for _, v := range metrics {
		stats, err := engine.NewRoundTripStats(v.m)
		if err != nil {
			return nil, err
		}
		v.srv.Stats = stats
		servers = append(servers, *v.srv)
	}
	sort.Stable(&serverSorter{es: servers})
	return servers, nil
}

type sval struct {
	u   *url.URL
	srv *engine.Server
	m   *memmetrics.RTMetrics
}

func newSval(s engine.Server) (*sval, error) {
	m, err := memmetrics.NewRTMetrics()
	if err != nil {
		return nil, err
	}
	u, err := url.Parse(s.URL)
	if err != nil {
		return nil, err
	}
	return &sval{srv: &s, m: m, u: u}, nil
}

type frontendSorter struct {
	frontends []engine.Frontend
}

func (s *frontendSorter) Len() int {
	return len(s.frontends)
}

func (s *frontendSorter) Swap(i, j int) {
	s.frontends[i], s.frontends[j] = s.frontends[j], s.frontends[i]
}

func (s *frontendSorter) Less(i, j int) bool {
	return cmpStats(s.frontends[i].Stats, s.frontends[j].Stats)
}

type serverSorter struct {
	es []engine.Server
}

func (s *serverSorter) Len() int {
	return len(s.es)
}

func (s *serverSorter) Swap(i, j int) {
	s.es[i], s.es[j] = s.es[j], s.es[i]
}

func (s *serverSorter) Less(i, j int) bool {
	return cmpStats(s.es[i].Stats, s.es[j].Stats)
}

func cmpStats(s1, s2 *engine.RoundTripStats) bool {
	// Items that have network errors go first
	if s1.NetErrorRatio() != 0 || s2.NetErrorRatio() != 0 {
		return s1.NetErrorRatio() > s2.NetErrorRatio()
	}

	// Items that have application level errors go next
	if s1.AppErrorRatio() != 0 || s2.AppErrorRatio() != 0 {
		return s1.AppErrorRatio() > s2.AppErrorRatio()
	}

	// More highly loaded items go next
	return s1.Counters.Total > s2.Counters.Total
}
<|endoftext|>


<commit_before>package jsonry

import (
	"bytes"
	"errors"
	"fmt"

	"code.cloudfoundry.org/cli/types"

	"encoding/json"
	"reflect"
	"strings"
)

var (
	jsonNumberType             = reflect.TypeOf(json.Number("0"))
	mapOfNullStringType        = reflect.TypeOf(map[string]types.NullString{})
	mapOfNullStringPointerType = reflect.PtrTo(mapOfNullStringType)
)

func Unmarshal(data []byte, store interface{}) error {
	storeValue, err := relectOnAndCheck(store)
	if err != nil {
		return err
	}

	var tree interface{}
	if err := unmarshalJSON(data, &tree); err != nil {
		return err
	}

	return unmarshal(storeValue, tree)
}

func relectOnAndCheck(store interface{}) (reflect.Value, error) {
	p := reflect.ValueOf(store)
	if kind := p.Kind(); kind != reflect.Ptr {
		return reflect.Value{}, errors.New("the storage object must be a pointer")
	}

	v := p.Elem()
	if kind := v.Kind(); kind != reflect.Struct {
		return reflect.Value{}, errors.New("the storage object pointer must point to a struct")
	}

	return v, nil
}

func unmarshalJSON(data []byte, store interface{}) error {
	d := json.NewDecoder(bytes.NewBuffer(data))
	d.UseNumber()
	return d.Decode(store)
}

func unmarshal(storeValue reflect.Value, tree interface{}) error {
	storeValue = depointerify(storeValue)
	storeType := storeValue.Type()

	for i := 0; i < storeType.NumField(); i++ {
		field := storeType.Field(i)
		path := computePath(field)
		if value, ok := navigateAndFetch(path, tree); ok {
			if err := set(field, storeValue.Field(i), value); err != nil {
				return err
			}
		}
	}

	return nil
}

func depointerify(v reflect.Value) reflect.Value {
	if v.Kind() != reflect.Ptr {
		return v
	}

	n := reflect.New(v.Type().Elem())
	v.Set(n)
	return n.Elem()
}

func actualKind(t reflect.Type) reflect.Kind {
	if k := t.Kind(); k != reflect.Ptr {
		return k
	}

	return t.Elem().Kind()
}

func computePath(field reflect.StructField) []string {
	if tag := field.Tag.Get("jsonry"); tag != "" {
		return strings.Split(tag, ".")
	}

	if tag := field.Tag.Get("json"); tag != "" {
		parts := strings.Split(tag, ",")
		if parts[0] != "" {
			return []string{parts[0]}
		}
	}

	return []string{strings.ToLower(field.Name)}
}

func navigateAndFetch(path []string, tree interface{}) (interface{}, bool) {
	branch := path[0]

	node, ok := tree.(map[string]interface{})
	if !ok {
		return nil, false
	}

	val, ok := node[branch]
	if !ok {
		return nil, false
	}

	if len(path) > 1 {
		if vals, ok := val.([]interface{}); ok {
			return iterateAndFetch(path, vals), true
		}

		return navigateAndFetch(path[1:], val)
	}

	return val, true
}

func iterateAndFetch(path []string, vals []interface{}) interface{} {
	var results []interface{}

	for _, v := range vals {
		r, _ := navigateAndFetch(path[1:], v)
		results = append(results, r)
	}

	return results
}

func set(field reflect.StructField, fieldValue reflect.Value, value interface{}) error {
	switch actualKind(field.Type) {
	case reflect.Struct:
		return unmarshal(fieldValue, value)
	case reflect.Slice:
		return setSlice(field.Name, fieldValue, value)
	default:
		return setValue(field.Name, fieldValue, value)
	}
}

func setValue(fieldName string, store reflect.Value, value interface{}) error {
	vv := valueOfWithDenumberification(value)

	// Special case: map[string]types.NullString
	store = depointerify(store)
	if store.Type().AssignableTo(mapOfNullStringType) {
		if mv, ok := tryToConvertMapOfNullStrings(value); ok {
			vv = mv
		}
	}

	if vv.Type().AssignableTo(store.Type()) {
		store.Set(vv)
		return nil
	}
	if vv.Type().ConvertibleTo(store.Type()) {
		store.Set(vv.Convert(store.Type()))
		return nil
	}

	return fmt.Errorf(
		"could not convert value '%v' type '%s' to '%s' for field '%s'",
		value,
		reflect.TypeOf(value),
		store.Type(),
		fieldName,
	)
}

func setSlice(fieldName string, store reflect.Value, value interface{}) error {
	vs, ok := value.([]interface{})
	if !ok {
		return fmt.Errorf(
			"could not convert value '%v' type '%s' to '%s' for field '%s' because it is not a list type",
			value,
			reflect.TypeOf(value),
			store.Type(),
			fieldName,
		)
	}

	store = depointerify(store)
	elemType := store.Type().Elem()
	arr := reflect.MakeSlice(reflect.SliceOf(elemType), len(vs), len(vs))
	for i, v := range vs {
		vv := valueOfWithDenumberification(v)

		if vv.Type().AssignableTo(elemType) {
			arr.Index(i).Set(vv)
			continue
		}
		if vv.Type().ConvertibleTo(elemType) {
			arr.Index(i).Set(vv.Convert(elemType))
			continue
		}

		return fmt.Errorf(
			"could not convert value '%v' type '%s' to '%s' for field '%s' index %d",
			v,
			vv.Type(),
			elemType,
			fieldName,
			i,
		)
	}

	store.Set(arr.Slice(0, arr.Len()))
	return nil
}

func tryToConvertMapOfNullStrings(value interface{}) (reflect.Value, bool) {
	if source, ok := value.(map[string]interface{}); ok {
		destination := reflect.MakeMap(mapOfNullStringType)

		for k, v := range source {
			ns := types.NewNullString()
			if s, ok := v.(string); ok {
				ns = types.NewNullString(s)
			}
			destination.SetMapIndex(reflect.ValueOf(k), reflect.ValueOf(ns))
		}

		return destination, true
	}

	return reflect.Value{}, false
}

func valueOfWithDenumberification(v interface{}) reflect.Value {
	vv := reflect.ValueOf(v)

	if vv.Type() == jsonNumberType {
		num := json.Number(vv.String())
		// Extend to support other number types as needed
		if i64, err := num.Int64(); err == nil {
			return reflect.ValueOf(int(i64))
		}
	}

	return vv
}
<commit_msg>Remove unused variable<commit_after>package jsonry

import (
	"bytes"
	"errors"
	"fmt"

	"code.cloudfoundry.org/cli/types"

	"encoding/json"
	"reflect"
	"strings"
)

var (
	jsonNumberType      = reflect.TypeOf(json.Number("0"))
	mapOfNullStringType = reflect.TypeOf(map[string]types.NullString{})
)

func Unmarshal(data []byte, store interface{}) error {
	storeValue, err := relectOnAndCheck(store)
	if err != nil {
		return err
	}

	var tree interface{}
	if err := unmarshalJSON(data, &tree); err != nil {
		return err
	}

	return unmarshal(storeValue, tree)
}

func relectOnAndCheck(store interface{}) (reflect.Value, error) {
	p := reflect.ValueOf(store)
	if kind := p.Kind(); kind != reflect.Ptr {
		return reflect.Value{}, errors.New("the storage object must be a pointer")
	}

	v := p.Elem()
	if kind := v.Kind(); kind != reflect.Struct {
		return reflect.Value{}, errors.New("the storage object pointer must point to a struct")
	}

	return v, nil
}

func unmarshalJSON(data []byte, store interface{}) error {
	d := json.NewDecoder(bytes.NewBuffer(data))
	d.UseNumber()
	return d.Decode(store)
}

func unmarshal(storeValue reflect.Value, tree interface{}) error {
	storeValue = depointerify(storeValue)
	storeType := storeValue.Type()

	for i := 0; i < storeType.NumField(); i++ {
		field := storeType.Field(i)
		path := computePath(field)
		if value, ok := navigateAndFetch(path, tree); ok {
			if err := set(field, storeValue.Field(i), value); err != nil {
				return err
			}
		}
	}

	return nil
}

func depointerify(v reflect.Value) reflect.Value {
	if v.Kind() != reflect.Ptr {
		return v
	}

	n := reflect.New(v.Type().Elem())
	v.Set(n)
	return n.Elem()
}

func actualKind(t reflect.Type) reflect.Kind {
	if k := t.Kind(); k != reflect.Ptr {
		return k
	}

	return t.Elem().Kind()
}

func computePath(field reflect.StructField) []string {
	if tag := field.Tag.Get("jsonry"); tag != "" {
		return strings.Split(tag, ".")
	}

	if tag := field.Tag.Get("json"); tag != "" {
		parts := strings.Split(tag, ",")
		if parts[0] != "" {
			return []string{parts[0]}
		}
	}

	return []string{strings.ToLower(field.Name)}
}

func navigateAndFetch(path []string, tree interface{}) (interface{}, bool) {
	branch := path[0]

	node, ok := tree.(map[string]interface{})
	if !ok {
		return nil, false
	}

	val, ok := node[branch]
	if !ok {
		return nil, false
	}

	if len(path) > 1 {
		if vals, ok := val.([]interface{}); ok {
			return iterateAndFetch(path, vals), true
		}

		return navigateAndFetch(path[1:], val)
	}

	return val, true
}

func iterateAndFetch(path []string, vals []interface{}) interface{} {
	var results []interface{}

	for _, v := range vals {
		r, _ := navigateAndFetch(path[1:], v)
		results = append(results, r)
	}

	return results
}

func set(field reflect.StructField, fieldValue reflect.Value, value interface{}) error {
	switch actualKind(field.Type) {
	case reflect.Struct:
		return unmarshal(fieldValue, value)
	case reflect.Slice:
		return setSlice(field.Name, fieldValue, value)
	default:
		return setValue(field.Name, fieldValue, value)
	}
}

func setValue(fieldName string, store reflect.Value, value interface{}) error {
	vv := valueOfWithDenumberification(value)

	// Special case: map[string]types.NullString
	store = depointerify(store)
	if store.Type().AssignableTo(mapOfNullStringType) {
		if mv, ok := tryToConvertMapOfNullStrings(value); ok {
			vv = mv
		}
	}

	if vv.Type().AssignableTo(store.Type()) {
		store.Set(vv)
		return nil
	}
	if vv.Type().ConvertibleTo(store.Type()) {
		store.Set(vv.Convert(store.Type()))
		return nil
	}

	return fmt.Errorf(
		"could not convert value '%v' type '%s' to '%s' for field '%s'",
		value,
		reflect.TypeOf(value),
		store.Type(),
		fieldName,
	)
}

func setSlice(fieldName string, store reflect.Value, value interface{}) error {
	vs, ok := value.([]interface{})
	if !ok {
		return fmt.Errorf(
			"could not convert value '%v' type '%s' to '%s' for field '%s' because it is not a list type",
			value,
			reflect.TypeOf(value),
			store.Type(),
			fieldName,
		)
	}

	store = depointerify(store)
	elemType := store.Type().Elem()
	arr := reflect.MakeSlice(reflect.SliceOf(elemType), len(vs), len(vs))
	for i, v := range vs {
		vv := valueOfWithDenumberification(v)

		if vv.Type().AssignableTo(elemType) {
			arr.Index(i).Set(vv)
			continue
		}
		if vv.Type().ConvertibleTo(elemType) {
			arr.Index(i).Set(vv.Convert(elemType))
			continue
		}

		return fmt.Errorf(
			"could not convert value '%v' type '%s' to '%s' for field '%s' index %d",
			v,
			vv.Type(),
			elemType,
			fieldName,
			i,
		)
	}

	store.Set(arr.Slice(0, arr.Len()))
	return nil
}

func tryToConvertMapOfNullStrings(value interface{}) (reflect.Value, bool) {
	if source, ok := value.(map[string]interface{}); ok {
		destination := reflect.MakeMap(mapOfNullStringType)

		for k, v := range source {
			ns := types.NewNullString()
			if s, ok := v.(string); ok {
				ns = types.NewNullString(s)
			}
			destination.SetMapIndex(reflect.ValueOf(k), reflect.ValueOf(ns))
		}

		return destination, true
	}

	return reflect.Value{}, false
}

func valueOfWithDenumberification(v interface{}) reflect.Value {
	vv := reflect.ValueOf(v)

	if vv.Type() == jsonNumberType {
		num := json.Number(vv.String())
		// Extend to support other number types as needed
		if i64, err := num.Int64(); err == nil {
			return reflect.ValueOf(int(i64))
		}
	}

	return vv
}
<|endoftext|>


<commit_before>package web

import (
	"crypto/tls"
	"log"
	"net"
	"net/http"
	"os"
	"sync"
)

var (
	_app  *Application
	_once sync.Once
)

// Data Data
type Data interface{}

// Callback function
type Callback func(ctx *Context) (Data, error)

// PanicCallback function
type PanicCallback func(http.ResponseWriter, *http.Request, interface{})

// Application is type of a web.Application
type Application struct {
	trees      map[string]*node
	logger     *log.Logger
	panic      PanicCallback
	paramsPool sync.Pool
	maxParams  uint16

	NotFound http.Handler
}

// CreateApplication return a singleton web.Application
func CreateApplication() *Application {
	_once.Do(func() {
		_app = &Application{}
	})
	return _app
}

// SetLogger set Logger
func (app *Application) SetLogger(logger *log.Logger) {
	app.logger = logger
}

// SetPanic set Logger
func (app *Application) SetPanic(panic PanicCallback) {
	app.panic = panic
}

// Use Add the given callback function to this application.middlewares.
func (app *Application) Use(path string, callback Callback) {

}

// On add event
func (app *Application) On(name string, cb Callback) {

}

// Get method
func (app *Application) Get(path string, cb Callback) {
	app.addRoute(http.MethodGet, path, cb)
}

// Head method
func (app *Application) Head(path string, cb Callback) {
	app.addRoute(http.MethodHead, path, cb)
}

// Post method
func (app *Application) Post(path string, cb Callback) {
	app.addRoute(http.MethodPost, path, cb)
}

// Put method
func (app *Application) Put(path string, cb Callback) {
	app.addRoute(http.MethodPut, path, cb)
}

// Patch method
func (app *Application) Patch(path string, cb Callback) {
	app.addRoute(http.MethodPatch, path, cb)
}

// Delete method
func (app *Application) Delete(path string, cb Callback) {
	app.addRoute(http.MethodDelete, path, cb)
}

// Options method
func (app *Application) Options(path string, cb Callback) {
	app.addRoute(http.MethodOptions, path, cb)
}

func (app *Application) addRoute(method, path string, cb Callback) {

	if method == "" {
		panic("method must not be empty")
	}

	if len(path) < 1 || path[0] != '/' {
		panic("path must begin with '/' in path '" + path + "'")
	}

	if cb == nil {
		panic("callback must not be nil")
	}

	if app.trees == nil {
		app.trees = make(map[string]*node)
	}

	root := app.trees[method]

	if root == nil {
		root = new(node)
		app.trees[method] = root
	}

	root.addRoute(path, cb)

	if pc := countParams(path); pc > app.maxParams {
		app.maxParams = pc
	}
}

// ServeFiles ("/src/*filepath", http.Dir("/var/www"))
func (app *Application) ServeFiles(path string, root http.FileSystem) {
	if len(path) < 10 || path[len(path)-10:] != "/*filepath" {
		panic("path must end with /*filepath in path '" + path + "'")
	}

	fileServer := http.FileServer(root)

	app.Get(path, func(ctx *Context) (Data, error) {
		ctx.r.URL.Path = ctx.Param("filepath")
		fileServer.ServeHTTP(ctx.w, ctx.r)
		return nil, nil
	})
}

// ServeHTTP w, r
func (app *Application) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer app.recv(w, r)
	path := r.URL.Path

	if root := app.trees[r.Method]; root != nil {

		if callback, params, _ := root.getValue(path, app.getParams); callback != nil {

			ctx := createContext(w, r, params)
			app.putParams(params)

			val, err := callback(ctx)

			if err != nil {

				switch err {
				case ErrUnauthorized:
					ctx.SetStatus(http.StatusUnauthorized)
				case ErrForbidden:
					ctx.SetStatus(http.StatusForbidden)
				default:
					ctx.SetStatus(http.StatusBadRequest)
				}

				app.logf("%s %s %d %s %s %d %v", r.RemoteAddr, r.Host, ctx.UserID(), r.Method, path, ctx.Status(), err)

				if err := ctx.write(err.Error()); err != nil {
					app.logf("ctx.write err: %v", err)
				}

				return
			}

			if val != nil {
				if err := ctx.write(val); err != nil {
					app.logf("ctx.write err: %v", err)
				}
			}

			return
		}
	}

	if app.NotFound != nil {
		app.NotFound.ServeHTTP(w, r)
	} else {
		http.NotFound(w, r)
	}
}

// ListenAndServe Serve with options on addr
func (app *Application) ListenAndServe(addr string, fns ...func(*http.Server)) error {

	l, err := net.Listen("tcp", addr)

	if err != nil {
		return err
	}

	defer l.Close()

	return app.serve(addr, l, fns...)
}

// ListenAndServeTLS Serve with tls and options on addr
func (app *Application) ListenAndServeTLS(addr string, tlsConfig *tls.Config, fns ...func(*http.Server)) error {

	l, err := tls.Listen("tcp", addr, tlsConfig)

	if err != nil {
		return err
	}

	defer l.Close()

	return app.serve(addr, l, fns...)
}

func (app *Application) serve(addr string, listener net.Listener, fns ...func(*http.Server)) error {

	mux := http.NewServeMux()

	mux.Handle("/", app)

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

	for _, fn := range fns {
		fn(srv)
	}

	if app.paramsPool.New == nil && app.maxParams > 0 {
		app.paramsPool.New = func() interface{} {
			ps := make(Params, 0, app.maxParams)
			return &ps
		}
	}

	if err := srv.Serve(listener); err != nil {
		return err
	}

	if err := srv.Close(); err != nil {
		return err
	}

	return nil
}

// Inspect method
func (app *Application) Inspect() string {
	return ""
}

// logf write log
func (app *Application) logf(format string, v ...interface{}) {
	if app.logger == nil {
		app.logger = log.New(os.Stdout, "", log.Ldate|log.Ltime)
	}

	app.logger.Printf(format, v...)
}

func (app *Application) getParams() *Params {
	ps := app.paramsPool.Get().(*Params)
	*ps = (*ps)[0:0] // reset slice
	return ps
}

func (app *Application) putParams(ps *Params) {
	if ps != nil {
		app.paramsPool.Put(ps)
	}
}

func (app *Application) recv(w http.ResponseWriter, r *http.Request) {
	if rcv := recover(); rcv != nil {
		if app.panic != nil {
			app.panic(w, r, rcv)
		} else {
			app.logf("%s %s %s %s rcv: %v", r.RemoteAddr, r.Host, r.Method, r.URL.Path, rcv)
		}
	}
}
<commit_msg>clean<commit_after>package web

import (
	"crypto/tls"
	"log"
	"net"
	"net/http"
	"os"
	"sync"
)

var (
	_app  *Application
	_once sync.Once
)

// Data Data
type Data interface{}

// Callback function
type Callback func(ctx *Context) (Data, error)

// PanicCallback function
type PanicCallback func(http.ResponseWriter, *http.Request, interface{})

// Application is type of a web.Application
type Application struct {
	trees      map[string]*node
	logger     *log.Logger
	panic      PanicCallback
	paramsPool sync.Pool
	maxParams  uint16

	NotFound http.Handler
}

// CreateApplication return a singleton web.Application
func CreateApplication() *Application {
	_once.Do(func() {
		_app = &Application{}
	})
	return _app
}

// SetLogger set Logger
func (app *Application) SetLogger(logger *log.Logger) {
	app.logger = logger
}

// SetPanic set Logger
func (app *Application) SetPanic(panic PanicCallback) {
	app.panic = panic
}

// Use Add the given callback function to this application.middlewares.
func (app *Application) Use(path string, callback Callback) {

}

// On add event
func (app *Application) On(name string, cb Callback) {

}

// Get method
func (app *Application) Get(path string, cb Callback) {
	app.addRoute(http.MethodGet, path, cb)
}

// Head method
func (app *Application) Head(path string, cb Callback) {
	app.addRoute(http.MethodHead, path, cb)
}

// Post method
func (app *Application) Post(path string, cb Callback) {
	app.addRoute(http.MethodPost, path, cb)
}

// Put method
func (app *Application) Put(path string, cb Callback) {
	app.addRoute(http.MethodPut, path, cb)
}

// Patch method
func (app *Application) Patch(path string, cb Callback) {
	app.addRoute(http.MethodPatch, path, cb)
}

// Delete method
func (app *Application) Delete(path string, cb Callback) {
	app.addRoute(http.MethodDelete, path, cb)
}

// Options method
func (app *Application) Options(path string, cb Callback) {
	app.addRoute(http.MethodOptions, path, cb)
}

func (app *Application) addRoute(method, path string, cb Callback) {

	if method == "" {
		panic("method must not be empty")
	}

	if len(path) < 1 || path[0] != '/' {
		panic("path must begin with '/' in path '" + path + "'")
	}

	if cb == nil {
		panic("callback must not be nil")
	}

	if app.trees == nil {
		app.trees = make(map[string]*node)
	}

	root := app.trees[method]

	if root == nil {
		root = new(node)
		app.trees[method] = root
	}

	root.addRoute(path, cb)

	if pc := countParams(path); pc > app.maxParams {
		app.maxParams = pc
	}

	if app.paramsPool.New == nil && app.maxParams > 0 {
		app.paramsPool.New = func() interface{} {
			ps := make(Params, 0, app.maxParams)
			return &ps
		}
	}
}

// ServeFiles ("/src/*filepath", http.Dir("/var/www"))
func (app *Application) ServeFiles(path string, root http.FileSystem) {
	if len(path) < 10 || path[len(path)-10:] != "/*filepath" {
		panic("path must end with /*filepath in path '" + path + "'")
	}

	fileServer := http.FileServer(root)

	app.Get(path, func(ctx *Context) (Data, error) {
		ctx.r.URL.Path = ctx.Param("filepath")
		fileServer.ServeHTTP(ctx.w, ctx.r)
		return nil, nil
	})
}

// ServeHTTP w, r
func (app *Application) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer app.recv(w, r)
	path := r.URL.Path

	if root := app.trees[r.Method]; root != nil {

		if callback, params, _ := root.getValue(path, app.getParams); callback != nil {

			ctx := createContext(w, r, params)
			app.putParams(params)

			val, err := callback(ctx)

			if err != nil {

				switch err {
				case ErrUnauthorized:
					ctx.SetStatus(http.StatusUnauthorized)
				case ErrForbidden:
					ctx.SetStatus(http.StatusForbidden)
				default:
					ctx.SetStatus(http.StatusBadRequest)
				}

				app.logf("%s %s %d %s %s %d %v", r.RemoteAddr, r.Host, ctx.UserID(), r.Method, path, ctx.Status(), err)

				if err := ctx.write(err.Error()); err != nil {
					app.logf("ctx.write err: %v", err)
				}

				return
			}

			if val != nil {
				if err := ctx.write(val); err != nil {
					app.logf("ctx.write err: %v", err)
				}
			}

			return
		}
	}

	if app.NotFound != nil {
		app.NotFound.ServeHTTP(w, r)
	} else {
		http.NotFound(w, r)
	}
}

// ListenAndServe Serve with options on addr
func (app *Application) ListenAndServe(addr string, fns ...func(*http.Server)) error {

	l, err := net.Listen("tcp", addr)

	if err != nil {
		return err
	}

	defer l.Close()

	return app.serve(addr, l, fns...)
}

// ListenAndServeTLS Serve with tls and options on addr
func (app *Application) ListenAndServeTLS(addr string, tlsConfig *tls.Config, fns ...func(*http.Server)) error {

	l, err := tls.Listen("tcp", addr, tlsConfig)

	if err != nil {
		return err
	}

	defer l.Close()

	return app.serve(addr, l, fns...)
}

func (app *Application) serve(addr string, listener net.Listener, fns ...func(*http.Server)) error {

	mux := http.NewServeMux()

	mux.Handle("/", app)

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

	for _, fn := range fns {
		fn(srv)
	}

	if err := srv.Serve(listener); err != nil {
		return err
	}

	if err := srv.Close(); err != nil {
		return err
	}

	return nil
}

// Inspect method
func (app *Application) Inspect() string {
	return ""
}

// logf write log
func (app *Application) logf(format string, v ...interface{}) {
	if app.logger == nil {
		app.logger = log.New(os.Stdout, "", log.Ldate|log.Ltime)
	}

	app.logger.Printf(format, v...)
}

func (app *Application) getParams() *Params {
	ps := app.paramsPool.Get().(*Params)
	*ps = (*ps)[0:0] // reset slice
	return ps
}

func (app *Application) putParams(ps *Params) {
	if ps != nil {
		app.paramsPool.Put(ps)
	}
}

func (app *Application) recv(w http.ResponseWriter, r *http.Request) {
	if rcv := recover(); rcv != nil {
		if app.panic != nil {
			app.panic(w, r, rcv)
		} else {
			app.logf("%s %s %s %s rcv: %v", r.RemoteAddr, r.Host, r.Method, r.URL.Path, rcv)
		}
	}
}
<|endoftext|>


<commit_before>package mirror

import (
	"os"
	"os/exec"
	"testing"
)

func TestFlock(t *testing.T) {
	t.Parallel()

	// container-based image cannot run this test.
	if os.Getenv("TRAVIS") == "true" {
		t.Skip()
	}

	cmd := exec.Command("flock", "t/mirror.toml", "sleep", "0.2")
	err := cmd.Start()
	if err != nil {
		t.Skip()
		return
	}

	f, err := os.Open("t/mirror.toml")
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()

	fl := Flock{f}
	if err = fl.Lock(); err == nil {
		t.Error(`err = fl.Lock(); err == nil`)
	} else {
		t.Log(err)
	}

	cmd.Wait()
	if err = fl.Lock(); err != nil {
		t.Fatal(err)
	}
	if err = fl.Unlock(); err != nil {
		t.Error(err)
	}
}
<commit_msg>[mirror] fix unstable flock test.<commit_after>package mirror

import (
	"os"
	"os/exec"
	"testing"
	"time"
)

func TestFlock(t *testing.T) {
	t.Parallel()

	cmd := exec.Command("flock", "t/mirror.toml", "sleep", "0.2")
	err := cmd.Start()
	if err != nil {
		t.Skip()
		return
	}
	time.Sleep(100 * time.Millisecond)

	f, err := os.Open("t/mirror.toml")
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()

	fl := Flock{f}
	if err = fl.Lock(); err == nil {
		t.Error(`err = fl.Lock(); err == nil`)
	} else {
		t.Log(err)
	}

	cmd.Wait()
	if err = fl.Lock(); err != nil {
		t.Fatal(err)
	}
	if err = fl.Unlock(); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>package gopisysfs

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"syscall"
	"time"
)

const (
	sys_i2c   = "sys/class/i2c-dev"
	i2c_SLAVE = 0x703
)

func I2CListDevices() ([]string, error) {
	devdir := file(sys_i2c)
	files, err := ioutil.ReadDir(devdir)
	if err != nil {
		return nil, err
	}

	names := []string{}
	for _, f := range files {
		if f.IsDir() {
			continue
		}
		name := f.Name()
		dev := filepath.Join("/dev", name)
		info("I2C Checking %v\n", dev)
		if _, err := os.Stat(dev); err != nil {
			continue
		}
		names = append(names, dev)
	}
	return names, nil
}

type I2CRecording struct {
	Timestamp time.Time
	Data      []byte
}

func copyBytes(buffer []byte, count int) []byte {
	ret := make([]byte, count)
	copy(ret, buffer)
	return ret
}

// I2CPoll establishes a connection to a slave I2C device and periodically reads a fixed number of bytes from that device.
// The dev, address, and bytes parameters indicates which device to read and how much to read each time.
// The bufferdepth determines how deep the returned channel's buffer is.
// All samples taken after the buffer is filled will be discarded until space is available.
// An unbuffered return is supported, and guarantees that a receive on that channel gets the most recent sample.
// The interval indicates the period to sample at.
// The returned channel will be closed if there's an error reading the device or the poller is closed using the returned termination function.
// Call the termination function returned when you no longer need to receive polling data.
func I2CPoll(dev string, address int, bytes int, bufferdepth int, interval time.Duration) (<-chan I2CRecording, func(), error) {

	ctrl, err := os.OpenFile(dev, os.O_RDWR, 0)
	if err != nil {
		return nil, nil, err
	}

	killer := make(chan bool, 1)

	termfn := func() {
		killer <- true
	}

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(ctrl.Fd()), i2c_SLAVE, uintptr(address))
	if errno != 0 {
		return nil, nil, errno
	}

	buffer := make([]byte, bytes)
	n, err := ctrl.Read(buffer)
	if err != nil {
		ctrl.Close()
		return nil, nil, err
	}

	// unbuffered channel - reader only gets data when asking to receive it, and they get the most recently available value.
	data := make(chan Recording, 0)
	record := I2CRecording{time.Now(), copyBytes(buffer, n)}

	go func() {
		defer close(data)
		defer ctrl.Close()

		tick := time.NewTicker(interval)
		defer tick.Stop()

		var stamp time.Time

		// we do some nil channel tricks to manipulate the select statement. dest is part of that.
		dest := data

		for {
			select {
			case <-killer:
				return
			case dest <- record:
				// disable dest until there's a new record.
				dest = nil
			case stamp = <-tick.C:
				n, err := ctrl.Read(buffer)
				if err != nil {
					info("I2C Unexpected error reading %v: %v\n", dev, err)
					return
				}
				record = I2CRecording{stamp, copyBytes(buffer, n)}
				// indicate there's data to send and reenable dest.
				dest = data
			}
		}
	}()

	return data, termfn, nil

}
<commit_msg>Correct incomplete fix<commit_after>package gopisysfs

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"syscall"
	"time"
)

const (
	sys_i2c   = "sys/class/i2c-dev"
	i2c_SLAVE = 0x703
)

func I2CListDevices() ([]string, error) {
	devdir := file(sys_i2c)
	files, err := ioutil.ReadDir(devdir)
	if err != nil {
		return nil, err
	}

	names := []string{}
	for _, f := range files {
		if f.IsDir() {
			continue
		}
		name := f.Name()
		dev := filepath.Join("/dev", name)
		info("I2C Checking %v\n", dev)
		if _, err := os.Stat(dev); err != nil {
			continue
		}
		names = append(names, dev)
	}
	return names, nil
}

type I2CRecording struct {
	Timestamp time.Time
	Data      []byte
}

func copyBytes(buffer []byte, count int) []byte {
	ret := make([]byte, count)
	copy(ret, buffer)
	return ret
}

// I2CPoll establishes a connection to a slave I2C device and periodically reads a fixed number of bytes from that device.
// The dev, address, and bytes parameters indicates which device to read and how much to read each time.
// The bufferdepth determines how deep the returned channel's buffer is.
// All samples taken after the buffer is filled will be discarded until space is available.
// An unbuffered return is supported, and guarantees that a receive on that channel gets the most recent sample.
// The interval indicates the period to sample at.
// The returned channel will be closed if there's an error reading the device or the poller is closed using the returned termination function.
// Call the termination function returned when you no longer need to receive polling data.
func I2CPoll(dev string, address int, bytes int, bufferdepth int, interval time.Duration) (<-chan I2CRecording, func(), error) {

	ctrl, err := os.OpenFile(dev, os.O_RDWR, 0)
	if err != nil {
		return nil, nil, err
	}

	killer := make(chan bool, 1)

	termfn := func() {
		killer <- true
	}

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(ctrl.Fd()), i2c_SLAVE, uintptr(address))
	if errno != 0 {
		return nil, nil, errno
	}

	buffer := make([]byte, bytes)
	n, err := ctrl.Read(buffer)
	if err != nil {
		ctrl.Close()
		return nil, nil, err
	}

	// unbuffered channel - reader only gets data when asking to receive it, and they get the most recently available value.
	data := make(chan I2CRecording, 0)
	record := I2CRecording{time.Now(), copyBytes(buffer, n)}

	go func() {
		defer close(data)
		defer ctrl.Close()

		tick := time.NewTicker(interval)
		defer tick.Stop()

		var stamp time.Time

		// we do some nil channel tricks to manipulate the select statement. dest is part of that.
		dest := data

		for {
			select {
			case <-killer:
				return
			case dest <- record:
				// disable dest until there's a new record.
				dest = nil
			case stamp = <-tick.C:
				n, err := ctrl.Read(buffer)
				if err != nil {
					info("I2C Unexpected error reading %v: %v\n", dev, err)
					return
				}
				record = I2CRecording{stamp, copyBytes(buffer, n)}
				// indicate there's data to send and reenable dest.
				dest = data
			}
		}
	}()

	return data, termfn, nil

}
<|endoftext|>