<commit_before>package web

import (
    "bytes"
    "container/vector"
    "crypto/hmac"
    "encoding/base64"
    "fmt"
    "http"
    "io/ioutil"
    "log"
    "os"
    "path"
    "reflect"
    "regexp"
    "strconv"
    "strings"
    "time"
)

//secret key used to store cookies
var secret = ""

type conn interface {
    StartResponse(status int)
    SetHeader(hdr string, val string, unique bool)
    Write(data []byte) (n int, err os.Error)
    Close()
}

type Context struct {
    *Request
    *conn
    responseStarted bool
}

func (ctx *Context) StartResponse(status int) {
    ctx.conn.StartResponse(status)
    ctx.responseStarted = true
}

func (ctx *Context) Write(data []byte) (n int, err os.Error) {
    if !ctx.responseStarted {
        ctx.StartResponse(200)
    }
    return ctx.conn.Write(data)
}
func (ctx *Context) WriteString(content string) {
    ctx.Write(strings.Bytes(content))
}

func (ctx *Context) Abort(status int, body string) {
    ctx.StartResponse(status)
    ctx.WriteString(body)
}

func (ctx *Context) Redirect(status int, url string) {
    ctx.SetHeader("Location", url, true)
    ctx.StartResponse(status)
    ctx.WriteString("Redirecting to: " + url)
}

func (ctx *Context) NotFound(message string) {
    ctx.StartResponse(404)
    ctx.WriteString(message)
}

//Sets a cookie -- duration is the amount of time in seconds. 0 = forever
func (ctx *Context) SetCookie(name string, value string, age int64) {
    if age == 0 {
        //do some really long time
    }

    utctime := time.UTC()
    utc1 := time.SecondsToUTC(utctime.Seconds() + 60*30)
    expires := utc1.RFC1123()
    expires = expires[0:len(expires)-3] + "GMT"
    cookie := fmt.Sprintf("%s=%s; expires=%s", name, value, expires)
    ctx.SetHeader("Set-Cookie", cookie, false)
}

func SetCookieSecret(key string) { secret = key }

func getCookieSig(val []byte, timestamp string) string {
    hm := hmac.NewSHA1(strings.Bytes(secret))

    hm.Write(val)
    hm.Write(strings.Bytes(timestamp))

    hex := fmt.Sprintf("%02x", hm.Sum())
    return hex
}

func (ctx *Context) SetSecureCookie(name string, val string, age int64) {
    //base64 encode the val
    if len(secret) == 0 {
        log.Stderrf("Secret Key for secure cookies has not been set. Please call web.SetCookieSecret\n")
        return
    }
    var buf bytes.Buffer
    encoder := base64.NewEncoder(base64.StdEncoding, &buf)
    encoder.Write(strings.Bytes(val))
    encoder.Close()
    vs := buf.String()
    vb := buf.Bytes()

    timestamp := strconv.Itoa64(time.Seconds())

    sig := getCookieSig(vb, timestamp)

    cookie := strings.Join([]string{vs, timestamp, sig}, "|")

    ctx.SetCookie(name, cookie, age)
}

func (ctx *Context) GetSecureCookie(name string) (string, bool) {

    cookie, ok := ctx.Request.Cookies[name]

    if !ok {
        return "", false
    }

    parts := strings.Split(cookie, "|", 3)

    val := parts[0]
    timestamp := parts[1]
    sig := parts[2]

    if getCookieSig(strings.Bytes(val), timestamp) != sig {
        return "", false
    }

    ts, _ := strconv.Atoi64(timestamp)

    if time.Seconds()-31*86400 > ts {
        return "", false
    }

    buf := bytes.NewBufferString(val)
    encoder := base64.NewDecoder(base64.StdEncoding, buf)

    res, _ := ioutil.ReadAll(encoder)
    return string(res), true
}

var contextType reflect.Type
var staticDir string

func init() {
    contextType = reflect.Typeof(Context{})
    SetStaticDir("static")
}

type route struct {
    r       string
    cr      *regexp.Regexp
    method  string
    handler *reflect.FuncValue
}

var routes = make(map[*regexp.Regexp]route)

func addRoute(r string, method string, handler interface{}) {
    cr, err := regexp.Compile(r)
    if err != nil {
        log.Stderrf("Error in route regex %q\n", r)
        return
    }
    fv := reflect.NewValue(handler).(*reflect.FuncValue)
    routes[cr] = route{r, cr, method, fv}
}

type httpConn struct {
    conn *http.Conn
}

func (c *httpConn) StartResponse(status int) { c.conn.WriteHeader(status) }

func (c *httpConn) SetHeader(hdr string, val string, unique bool) {
    //right now unique can't be implemented through the http package.
    //see issue 488
    c.conn.SetHeader(hdr, val)
}

func (c *httpConn) WriteString(content string) {
    buf := bytes.NewBufferString(content)
    c.conn.Write(buf.Bytes())
}

func (c *httpConn) Write(content []byte) (n int, err os.Error) {
    return c.conn.Write(content)
}

func (c *httpConn) Close() {
    rwc, buf, _ := c.conn.Hijack()
    if buf != nil {
        buf.Flush()
    }

    if rwc != nil {
        rwc.Close()
    }
}

func httpHandler(c *http.Conn, req *http.Request) {
    conn := httpConn{c}

    wreq := newRequest(req)

    routeHandler(wreq, &conn)
}

func routeHandler(req *Request, c conn) {
    requestPath := req.URL.Path

    //log the request
    if len(req.URL.RawQuery) == 0 {
        log.Stdout(requestPath)
    } else {
        log.Stdout(requestPath + "?" + req.URL.RawQuery)
    }

    //parse the form data (if it exists)
    perr := req.parseParams()
    if perr != nil {
        log.Stderrf("Failed to parse form data %q", perr.String())
    }

    //parse the cookies
    perr = req.parseCookies()
    if perr != nil {
        log.Stderrf("Failed to parse cookies %q", perr.String())
    }

    ctx := Context{req, &c, false}

    //try to serve a static file
    staticFile := path.Join(staticDir, requestPath)
    if fileExists(staticFile) {
        serveFile(&ctx, staticFile)
        return
    }

    //set default encoding
    ctx.SetHeader("Content-Type", "text/html; charset=utf-8", true)
    ctx.SetHeader("Server", "web.go", true)

    for cr, route := range routes {
        if req.Method != route.method {
            continue
        }

        if !cr.MatchString(requestPath) {
            continue
        }
        match := cr.MatchStrings(requestPath)

        if len(match[0]) != len(requestPath) {
            continue
        }

        var args vector.Vector

        handlerType := route.handler.Type().(*reflect.FuncType)

        //check if the first arg in the handler is a context type
        if handlerType.NumIn() > 0 {
            if a0, ok := handlerType.In(0).(*reflect.PtrType); ok {
                typ := a0.Elem()
                if typ == contextType {
                    args.Push(reflect.NewValue(&ctx))
                }
            }
        }

        for _, arg := range match[1:] {
            args.Push(reflect.NewValue(arg))
        }

        if args.Len() != handlerType.NumIn() {
            log.Stderrf("Incorrect number of arguments for %s\n", requestPath)
            ctx.Abort(500, "Server Error")
            return
        }

        valArgs := make([]reflect.Value, args.Len())
        for i := 0; i < args.Len(); i++ {
            valArgs[i] = args.At(i).(reflect.Value)
        }

        ret := route.handler.Call(valArgs)

        if len(ret) == 0 {
            return
        }

        sval, ok := ret[0].(*reflect.StringValue)

        if ok && !ctx.responseStarted {
            ctx.StartResponse(200)
            ctx.WriteString(sval.Get())
        }

        return
    }

    ctx.Abort(404, "Page not found")
}

//runs the web application and serves http requests
func Run(addr string) {
    http.Handle("/", http.HandlerFunc(httpHandler))

    log.Stdoutf("web.go serving %s", addr)
    err := http.ListenAndServe(addr, nil)
    if err != nil {
        log.Exit("ListenAndServe:", err)
    }
}

//runs the web application and serves scgi requests
func RunScgi(addr string) {
    log.Stdoutf("web.go serving scgi %s", addr)
    listenAndServeScgi(addr)
}

//runs the web application by serving fastcgi requests
func RunFcgi(addr string) {
    log.Stdoutf("web.go serving fcgi %s", addr)
    listenAndServeFcgi(addr)
}

//Adds a handler for the 'GET' http method.
func Get(route string, handler interface{}) { addRoute(route, "GET", handler) }

//Adds a handler for the 'POST' http method.
func Post(route string, handler interface{}) { addRoute(route, "POST", handler) }

//Adds a handler for the 'PUT' http method.
func Put(route string, handler interface{}) { addRoute(route, "PUT", handler) }

//Adds a handler for the 'DELETE' http method.
func Delete(route string, handler interface{}) {
    addRoute(route, "DELETE", handler)
}

func dirExists(dir string) bool {
    d, e := os.Stat(dir)
    switch {
    case e != nil:
        return false
    case !d.IsDirectory():
        return false
    }

    return true
}

func fileExists(dir string) bool {
    d, e := os.Stat(dir)
    switch {
    case e != nil:
        return false
    case !d.IsRegular():
        return false
    }

    return true
}

type dirError string

func (path dirError) String() string { return "Failed to set directory " + string(path) }

func getCwd() string { return os.Getenv("PWD") }

//changes the location of the static directory. by default, it's under the 'static' folder
//of the directory containing the web application
func SetStaticDir(dir string) os.Error {
    cwd := getCwd()
    sd := path.Join(cwd, dir)
    if !dirExists(sd) {
        return dirError(sd)

    }
    staticDir = sd

    return nil
}

func Urlencode(data map[string]string) string {
    var buf bytes.Buffer
    for k, v := range (data) {
        buf.WriteString(http.URLEscape(k))
        buf.WriteByte('=')
        buf.WriteString(http.URLEscape(v))
        buf.WriteByte('&')
    }
    s := buf.String()
    return s[0 : len(s)-1]
}

//copied from go's http package, because it's not public
var statusText = map[int]string{
    http.StatusContinue: "Continue",
    http.StatusSwitchingProtocols: "Switching Protocols",

    http.StatusOK: "OK",
    http.StatusCreated: "Created",
    http.StatusAccepted: "Accepted",
    http.StatusNonAuthoritativeInfo: "Non-Authoritative Information",
    http.StatusNoContent: "No Content",
    http.StatusResetContent: "Reset Content",
    http.StatusPartialContent: "Partial Content",

    http.StatusMultipleChoices: "Multiple Choices",
    http.StatusMovedPermanently: "Moved Permanently",
    http.StatusFound: "Found",
    http.StatusSeeOther: "See Other",
    http.StatusNotModified: "Not Modified",
    http.StatusUseProxy: "Use Proxy",
    http.StatusTemporaryRedirect: "Temporary Redirect",

    http.StatusBadRequest: "Bad Request",
    http.StatusUnauthorized: "Unauthorized",
    http.StatusPaymentRequired: "Payment Required",
    http.StatusForbidden: "Forbidden",
    http.StatusNotFound: "Not Found",
    http.StatusMethodNotAllowed: "Method Not Allowed",
    http.StatusNotAcceptable: "Not Acceptable",
    http.StatusProxyAuthRequired: "Proxy Authentication Required",
    http.StatusRequestTimeout: "Request Timeout",
    http.StatusConflict: "Conflict",
    http.StatusGone: "Gone",
    http.StatusLengthRequired: "Length Required",
    http.StatusPreconditionFailed: "Precondition Failed",
    http.StatusRequestEntityTooLarge: "Request Entity Too Large",
    http.StatusRequestURITooLong: "Request URI Too Long",
    http.StatusUnsupportedMediaType: "Unsupported Media Type",
    http.StatusRequestedRangeNotSatisfiable: "Requested Range Not Satisfiable",
    http.StatusExpectationFailed: "Expectation Failed",

    http.StatusInternalServerError: "Internal Server Error",
    http.StatusNotImplemented: "Not Implemented",
    http.StatusBadGateway: "Bad Gateway",
    http.StatusServiceUnavailable: "Service Unavailable",
    http.StatusGatewayTimeout: "Gateway Timeout",
    http.StatusHTTPVersionNotSupported: "HTTP Version Not Supported",
}
<commit_msg>Update web.go to compile on -release<commit_after>package web

import (
    "bytes"
    "container/vector"
    "crypto/hmac"
    "encoding/base64"
    "fmt"
    "http"
    "io/ioutil"
    "log"
    "os"
    "path"
    "reflect"
    "regexp"
    "strconv"
    "strings"
    "time"
)

//secret key used to store cookies
var secret = ""

type conn interface {
    StartResponse(status int)
    SetHeader(hdr string, val string, unique bool)
    Write(data []byte) (n int, err os.Error)
    Close()
}

type Context struct {
    *Request
    *conn
    responseStarted bool
}

func (ctx *Context) StartResponse(status int) {
    ctx.conn.StartResponse(status)
    ctx.responseStarted = true
}

func (ctx *Context) Write(data []byte) (n int, err os.Error) {
    if !ctx.responseStarted {
        ctx.StartResponse(200)
    }
    return ctx.conn.Write(data)
}
func (ctx *Context) WriteString(content string) {
    ctx.Write(strings.Bytes(content))
}

func (ctx *Context) Abort(status int, body string) {
    ctx.StartResponse(status)
    ctx.WriteString(body)
}

func (ctx *Context) Redirect(status int, url string) {
    ctx.SetHeader("Location", url, true)
    ctx.StartResponse(status)
    ctx.WriteString("Redirecting to: " + url)
}

func (ctx *Context) NotFound(message string) {
    ctx.StartResponse(404)
    ctx.WriteString(message)
}

//Sets a cookie -- duration is the amount of time in seconds. 0 = forever
func (ctx *Context) SetCookie(name string, value string, age int64) {
    if age == 0 {
        //do some really long time
    }

    utctime := time.UTC()
    utc1 := time.SecondsToUTC(utctime.Seconds() + 60*30)
    expires := utc1.Format(time.RFC1123)
    expires = expires[0:len(expires)-3] + "GMT"
    cookie := fmt.Sprintf("%s=%s; expires=%s", name, value, expires)
    ctx.SetHeader("Set-Cookie", cookie, false)
}

func SetCookieSecret(key string) { secret = key }

func getCookieSig(val []byte, timestamp string) string {
    hm := hmac.NewSHA1(strings.Bytes(secret))

    hm.Write(val)
    hm.Write(strings.Bytes(timestamp))

    hex := fmt.Sprintf("%02x", hm.Sum())
    return hex
}

func (ctx *Context) SetSecureCookie(name string, val string, age int64) {
    //base64 encode the val
    if len(secret) == 0 {
        log.Stderrf("Secret Key for secure cookies has not been set. Please call web.SetCookieSecret\n")
        return
    }
    var buf bytes.Buffer
    encoder := base64.NewEncoder(base64.StdEncoding, &buf)
    encoder.Write(strings.Bytes(val))
    encoder.Close()
    vs := buf.String()
    vb := buf.Bytes()

    timestamp := strconv.Itoa64(time.Seconds())

    sig := getCookieSig(vb, timestamp)

    cookie := strings.Join([]string{vs, timestamp, sig}, "|")

    ctx.SetCookie(name, cookie, age)
}

func (ctx *Context) GetSecureCookie(name string) (string, bool) {

    cookie, ok := ctx.Request.Cookies[name]

    if !ok {
        return "", false
    }

    parts := strings.Split(cookie, "|", 3)

    val := parts[0]
    timestamp := parts[1]
    sig := parts[2]

    if getCookieSig(strings.Bytes(val), timestamp) != sig {
        return "", false
    }

    ts, _ := strconv.Atoi64(timestamp)

    if time.Seconds()-31*86400 > ts {
        return "", false
    }

    buf := bytes.NewBufferString(val)
    encoder := base64.NewDecoder(base64.StdEncoding, buf)

    res, _ := ioutil.ReadAll(encoder)
    return string(res), true
}

var contextType reflect.Type
var staticDir string

func init() {
    contextType = reflect.Typeof(Context{})
    SetStaticDir("static")
}

type route struct {
    r       string
    cr      *regexp.Regexp
    method  string
    handler *reflect.FuncValue
}

var routes = make(map[*regexp.Regexp]route)

func addRoute(r string, method string, handler interface{}) {
    cr, err := regexp.Compile(r)
    if err != nil {
        log.Stderrf("Error in route regex %q\n", r)
        return
    }
    fv := reflect.NewValue(handler).(*reflect.FuncValue)
    routes[cr] = route{r, cr, method, fv}
}

type httpConn struct {
    conn *http.Conn
}

func (c *httpConn) StartResponse(status int) { c.conn.WriteHeader(status) }

func (c *httpConn) SetHeader(hdr string, val string, unique bool) {
    //right now unique can't be implemented through the http package.
    //see issue 488
    c.conn.SetHeader(hdr, val)
}

func (c *httpConn) WriteString(content string) {
    buf := bytes.NewBufferString(content)
    c.conn.Write(buf.Bytes())
}

func (c *httpConn) Write(content []byte) (n int, err os.Error) {
    return c.conn.Write(content)
}

func (c *httpConn) Close() {
    rwc, buf, _ := c.conn.Hijack()
    if buf != nil {
        buf.Flush()
    }

    if rwc != nil {
        rwc.Close()
    }
}

func httpHandler(c *http.Conn, req *http.Request) {
    conn := httpConn{c}

    wreq := newRequest(req)

    routeHandler(wreq, &conn)
}

func routeHandler(req *Request, c conn) {
    requestPath := req.URL.Path

    //log the request
    if len(req.URL.RawQuery) == 0 {
        log.Stdout(requestPath)
    } else {
        log.Stdout(requestPath + "?" + req.URL.RawQuery)
    }

    //parse the form data (if it exists)
    perr := req.parseParams()
    if perr != nil {
        log.Stderrf("Failed to parse form data %q", perr.String())
    }

    //parse the cookies
    perr = req.parseCookies()
    if perr != nil {
        log.Stderrf("Failed to parse cookies %q", perr.String())
    }

    ctx := Context{req, &c, false}

    //try to serve a static file
    staticFile := path.Join(staticDir, requestPath)
    if fileExists(staticFile) {
        serveFile(&ctx, staticFile)
        return
    }

    //set default encoding
    ctx.SetHeader("Content-Type", "text/html; charset=utf-8", true)
    ctx.SetHeader("Server", "web.go", true)

    for cr, route := range routes {
        if req.Method != route.method {
            continue
        }

        if !cr.MatchString(requestPath) {
            continue
        }
        match := cr.MatchStrings(requestPath)

        if len(match[0]) != len(requestPath) {
            continue
        }

        var args vector.Vector

        handlerType := route.handler.Type().(*reflect.FuncType)

        //check if the first arg in the handler is a context type
        if handlerType.NumIn() > 0 {
            if a0, ok := handlerType.In(0).(*reflect.PtrType); ok {
                typ := a0.Elem()
                if typ == contextType {
                    args.Push(reflect.NewValue(&ctx))
                }
            }
        }

        for _, arg := range match[1:] {
            args.Push(reflect.NewValue(arg))
        }

        if args.Len() != handlerType.NumIn() {
            log.Stderrf("Incorrect number of arguments for %s\n", requestPath)
            ctx.Abort(500, "Server Error")
            return
        }

        valArgs := make([]reflect.Value, args.Len())
        for i := 0; i < args.Len(); i++ {
            valArgs[i] = args.At(i).(reflect.Value)
        }

        ret := route.handler.Call(valArgs)

        if len(ret) == 0 {
            return
        }

        sval, ok := ret[0].(*reflect.StringValue)

        if ok && !ctx.responseStarted {
            ctx.StartResponse(200)
            ctx.WriteString(sval.Get())
        }

        return
    }

    ctx.Abort(404, "Page not found")
}

//runs the web application and serves http requests
func Run(addr string) {
    http.Handle("/", http.HandlerFunc(httpHandler))

    log.Stdoutf("web.go serving %s", addr)
    err := http.ListenAndServe(addr, nil)
    if err != nil {
        log.Exit("ListenAndServe:", err)
    }
}

//runs the web application and serves scgi requests
func RunScgi(addr string) {
    log.Stdoutf("web.go serving scgi %s", addr)
    listenAndServeScgi(addr)
}

//runs the web application by serving fastcgi requests
func RunFcgi(addr string) {
    log.Stdoutf("web.go serving fcgi %s", addr)
    listenAndServeFcgi(addr)
}

//Adds a handler for the 'GET' http method.
func Get(route string, handler interface{}) { addRoute(route, "GET", handler) }

//Adds a handler for the 'POST' http method.
func Post(route string, handler interface{}) { addRoute(route, "POST", handler) }

//Adds a handler for the 'PUT' http method.
func Put(route string, handler interface{}) { addRoute(route, "PUT", handler) }

//Adds a handler for the 'DELETE' http method.
func Delete(route string, handler interface{}) {
    addRoute(route, "DELETE", handler)
}

func dirExists(dir string) bool {
    d, e := os.Stat(dir)
    switch {
    case e != nil:
        return false
    case !d.IsDirectory():
        return false
    }

    return true
}

func fileExists(dir string) bool {
    d, e := os.Stat(dir)
    switch {
    case e != nil:
        return false
    case !d.IsRegular():
        return false
    }

    return true
}

type dirError string

func (path dirError) String() string { return "Failed to set directory " + string(path) }

func getCwd() string { return os.Getenv("PWD") }

//changes the location of the static directory. by default, it's under the 'static' folder
//of the directory containing the web application
func SetStaticDir(dir string) os.Error {
    cwd := getCwd()
    sd := path.Join(cwd, dir)
    if !dirExists(sd) {
        return dirError(sd)

    }
    staticDir = sd

    return nil
}

func Urlencode(data map[string]string) string {
    var buf bytes.Buffer
    for k, v := range (data) {
        buf.WriteString(http.URLEscape(k))
        buf.WriteByte('=')
        buf.WriteString(http.URLEscape(v))
        buf.WriteByte('&')
    }
    s := buf.String()
    return s[0 : len(s)-1]
}

//copied from go's http package, because it's not public
var statusText = map[int]string{
    http.StatusContinue: "Continue",
    http.StatusSwitchingProtocols: "Switching Protocols",

    http.StatusOK: "OK",
    http.StatusCreated: "Created",
    http.StatusAccepted: "Accepted",
    http.StatusNonAuthoritativeInfo: "Non-Authoritative Information",
    http.StatusNoContent: "No Content",
    http.StatusResetContent: "Reset Content",
    http.StatusPartialContent: "Partial Content",

    http.StatusMultipleChoices: "Multiple Choices",
    http.StatusMovedPermanently: "Moved Permanently",
    http.StatusFound: "Found",
    http.StatusSeeOther: "See Other",
    http.StatusNotModified: "Not Modified",
    http.StatusUseProxy: "Use Proxy",
    http.StatusTemporaryRedirect: "Temporary Redirect",

    http.StatusBadRequest: "Bad Request",
    http.StatusUnauthorized: "Unauthorized",
    http.StatusPaymentRequired: "Payment Required",
    http.StatusForbidden: "Forbidden",
    http.StatusNotFound: "Not Found",
    http.StatusMethodNotAllowed: "Method Not Allowed",
    http.StatusNotAcceptable: "Not Acceptable",
    http.StatusProxyAuthRequired: "Proxy Authentication Required",
    http.StatusRequestTimeout: "Request Timeout",
    http.StatusConflict: "Conflict",
    http.StatusGone: "Gone",
    http.StatusLengthRequired: "Length Required",
    http.StatusPreconditionFailed: "Precondition Failed",
    http.StatusRequestEntityTooLarge: "Request Entity Too Large",
    http.StatusRequestURITooLong: "Request URI Too Long",
    http.StatusUnsupportedMediaType: "Unsupported Media Type",
    http.StatusRequestedRangeNotSatisfiable: "Requested Range Not Satisfiable",
    http.StatusExpectationFailed: "Expectation Failed",

    http.StatusInternalServerError: "Internal Server Error",
    http.StatusNotImplemented: "Not Implemented",
    http.StatusBadGateway: "Bad Gateway",
    http.StatusServiceUnavailable: "Service Unavailable",
    http.StatusGatewayTimeout: "Gateway Timeout",
    http.StatusHTTPVersionNotSupported: "HTTP Version Not Supported",
}
<|endoftext|>


<commit_before>// Package wmi provides a WQL interface for WMI on Windows.
package wmi

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	ole "github.com/mjibson/go-ole"
	"github.com/mjibson/go-ole/oleutil"
)

var l = log.New(os.Stdout, "", log.LstdFlags)

var (
	ErrInvalidEntityType = errors.New("wmi: invalid entity type")
)

func init() {
	go func() {
		runtime.LockOSThread()
		ole.CoInitializeEx(0, 2)
		for f := range mainfunc {
			f()
		}
	}()
}

// queue of work to run in main thread.
var mainfunc = make(chan func())

// do runs f on the main thread.
func do(f func() error) (err error) {
	done := make(chan bool, 1)
	mainfunc <- func() {
		err = f()
		done <- true
	}
	<-done
	return
}

// QueryNamespace invokes Query with the given namespace on the local machine.
func QueryNamespace(query string, dst interface{}, namespace string) error {
	return Query(query, dst, nil, namespace)
}

var lock = sync.Mutex{}

// Query runs the WQL query and appends the values to dst.
//
// dst must have type *[]S or *[]*S, for some struct type S. Fields selected in
// the query must have the same name in dst. Supported types are all signed and
// unsigned integers, time.Time, string, bool. Array types are not supported.
// See wmi_test.go for some examples.
//
// By default, the local machine and default namespace are used. These can be
// changed using connectServerArgs. See
// http://msdn.microsoft.com/en-us/library/aa393720.aspx for details.
func Query(query string, dst interface{}, connectServerArgs ...interface{}) (queryErr error) {
	f := func() error {
		defer func() {
			if e := recover(); e != nil {
				queryErr = e.(error)
			}
		}()
		dv := reflect.ValueOf(dst)
		if dv.Kind() != reflect.Ptr || dv.IsNil() {
			return ErrInvalidEntityType
		}
		dv = dv.Elem()
		mat, elemType := checkMultiArg(dv)
		if mat == multiArgTypeInvalid {
			return ErrInvalidEntityType
		}

		unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
		if err != nil {
			return err
		}
		defer unknown.Release()

		wmi, err := unknown.QueryInterface(ole.IID_IDispatch)
		if err != nil {
			return err
		}
		defer wmi.Release()

		// service is a SWbemServices
		serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer", connectServerArgs...)
		if err != nil {
			return err
		}
		service := serviceRaw.ToIDispatch()
		defer service.Release()

		// result is a SWBemObjectSet
		resultRaw, err := oleutil.CallMethod(service, "ExecQuery", query)
		if err != nil {
			return err
		}
		result := resultRaw.ToIDispatch()
		defer result.Release()

		count, err := oleInt64(result, "Count")
		if err != nil {
			return err
		}

		var errFieldMismatch error
		for i := int64(0); i < count; i++ {
			err := func() error {
				// item is a SWbemObject, but really a Win32_Process
				itemRaw, err := oleutil.CallMethod(result, "ItemIndex", i)
				if err != nil {
					return err
				}
				item := itemRaw.ToIDispatch()
				defer item.Release()

				ev := reflect.New(elemType)
				if err = loadEntity(ev.Interface(), item); err != nil {
					if _, ok := err.(*ErrFieldMismatch); ok {
						// We continue loading entities even in the face of field mismatch errors.
						// If we encounter any other error, that other error is returned. Otherwise,
						// an ErrFieldMismatch is returned.
						errFieldMismatch = err
					} else {
						return err
					}
				}
				if mat != multiArgTypeStructPtr {
					ev = ev.Elem()
				}
				dv.Set(reflect.Append(dv, ev))
				return nil
			}()
			if err != nil {
				return err
			}
		}
		return errFieldMismatch
	}
	r := do(f)
	if r != nil && queryErr == nil {
		queryErr = r
	}
	return
}

// ErrFieldMismatch is returned when a field is to be loaded into a different
// type than the one it was stored from, or when a field is missing or
// unexported in the destination struct.
// StructType is the type of the struct pointed to by the destination argument.
type ErrFieldMismatch struct {
	StructType reflect.Type
	FieldName  string
	Reason     string
}

func (e *ErrFieldMismatch) Error() string {
	return fmt.Sprintf("wmi: cannot load field %q into a %q: %s",
		e.FieldName, e.StructType, e.Reason)
}

var timeType = reflect.TypeOf(time.Time{})

// loadEntity loads a SWbemObject into a struct pointer.
func loadEntity(dst interface{}, src *ole.IDispatch) (errFieldMismatch error) {
	v := reflect.ValueOf(dst).Elem()
	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		n := v.Type().Field(i).Name
		if !f.CanSet() {
			return &ErrFieldMismatch{
				StructType: f.Type(),
				FieldName:  n,
				Reason:     "CanSet() is false",
			}
		}
		prop, err := oleutil.GetProperty(src, n)
		if err != nil {
			errFieldMismatch = &ErrFieldMismatch{
				StructType: f.Type(),
				FieldName:  n,
				Reason:     "no such struct field",
			}
			continue
		}
		switch val := prop.Value(); reflect.ValueOf(val).Kind() {
		case reflect.Int64:
			iv := val.(int64)
			switch f.Kind() {
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				f.SetInt(iv)
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				f.SetUint(uint64(iv))
			default:
				return &ErrFieldMismatch{
					StructType: f.Type(),
					FieldName:  n,
					Reason:     "not an integer class",
				}
			}
		case reflect.String:
			sv := val.(string)
			iv, err := strconv.ParseInt(sv, 10, 64)
			switch f.Kind() {
			case reflect.String:
				f.SetString(sv)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				if err != nil {
					return err
				}
				f.SetInt(iv)
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				if err != nil {
					return err
				}
				f.SetUint(uint64(iv))
			case reflect.Struct:
				switch f.Type() {
				case timeType:
					if len(sv) == 25 {
						sv = sv[:22] + "0" + sv[22:]
					}
					t, err := time.Parse("20060102150405.000000-0700", sv)
					if err != nil {
						return err
					}
					f.Set(reflect.ValueOf(t))
				}
			}
		case reflect.Bool:
			bv := val.(bool)
			switch f.Kind() {
			case reflect.Bool:
				f.SetBool(bv)
			default:
				return &ErrFieldMismatch{
					StructType: f.Type(),
					FieldName:  n,
					Reason:     "not a bool",
				}
			}
		}
	}
	return errFieldMismatch
}

type multiArgType int

const (
	multiArgTypeInvalid multiArgType = iota
	multiArgTypeStruct
	multiArgTypeStructPtr
)

// checkMultiArg checks that v has type []S, []*S for some struct type S.
//
// It returns what category the slice's elements are, and the reflect.Type
// that represents S.
func checkMultiArg(v reflect.Value) (m multiArgType, elemType reflect.Type) {
	if v.Kind() != reflect.Slice {
		return multiArgTypeInvalid, nil
	}
	elemType = v.Type().Elem()
	switch elemType.Kind() {
	case reflect.Struct:
		return multiArgTypeStruct, elemType
	case reflect.Ptr:
		elemType = elemType.Elem()
		if elemType.Kind() == reflect.Struct {
			return multiArgTypeStructPtr, elemType
		}
	}
	return multiArgTypeInvalid, nil
}

func oleInt64(item *ole.IDispatch, prop string) (int64, error) {
	v, err := oleutil.GetProperty(item, prop)
	if err != nil {
		return 0, err
	}
	i := int64(v.Val)
	return i, nil
}

func CreateQuery(src interface{}, where string) string {
	var b bytes.Buffer
	b.WriteString("SELECT ")
	s := reflect.Indirect(reflect.ValueOf(src))
	t := s.Type()
	if s.Kind() == reflect.Slice {
		t = t.Elem()
	}
	if t.Kind() != reflect.Struct {
		return ""
	}
	var fields []string
	for i := 0; i < t.NumField(); i++ {
		fields = append(fields, t.Field(i).Name)
	}
	b.WriteString(strings.Join(fields, ", "))
	b.WriteString(" FROM ")
	b.WriteString(t.Name())
	b.WriteString(" " + where)
	return b.String()
}
<commit_msg>This appears to work better<commit_after>// Package wmi provides a WQL interface for WMI on Windows.
package wmi

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	ole "github.com/mjibson/go-ole"
	"github.com/mjibson/go-ole/oleutil"
)

var l = log.New(os.Stdout, "", log.LstdFlags)

var (
	ErrInvalidEntityType = errors.New("wmi: invalid entity type")
)

func init() {
	go func() {
		runtime.LockOSThread()
		ole.CoInitializeEx(0, 0)
		for f := range mainfunc {
			f()
		}
	}()
}

// queue of work to run in main thread.
var mainfunc = make(chan func())

// do runs f on the main thread.
func do(f func() error) (err error) {
	done := make(chan bool, 1)
	mainfunc <- func() {
		err = f()
		done <- true
	}
	<-done
	return
}

// QueryNamespace invokes Query with the given namespace on the local machine.
func QueryNamespace(query string, dst interface{}, namespace string) error {
	return Query(query, dst, nil, namespace)
}

var lock = sync.Mutex{}

// Query runs the WQL query and appends the values to dst.
//
// dst must have type *[]S or *[]*S, for some struct type S. Fields selected in
// the query must have the same name in dst. Supported types are all signed and
// unsigned integers, time.Time, string, bool. Array types are not supported.
// See wmi_test.go for some examples.
//
// By default, the local machine and default namespace are used. These can be
// changed using connectServerArgs. See
// http://msdn.microsoft.com/en-us/library/aa393720.aspx for details.
func Query(query string, dst interface{}, connectServerArgs ...interface{}) (queryErr error) {
	f := func() error {
		defer func() {
			if e := recover(); e != nil {
				queryErr = e.(error)
			}
		}()
		dv := reflect.ValueOf(dst)
		if dv.Kind() != reflect.Ptr || dv.IsNil() {
			return ErrInvalidEntityType
		}
		dv = dv.Elem()
		mat, elemType := checkMultiArg(dv)
		if mat == multiArgTypeInvalid {
			return ErrInvalidEntityType
		}

		unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
		if err != nil {
			return err
		}
		defer unknown.Release()

		wmi, err := unknown.QueryInterface(ole.IID_IDispatch)
		if err != nil {
			return err
		}
		defer wmi.Release()

		// service is a SWbemServices
		serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer", connectServerArgs...)
		if err != nil {
			return err
		}
		service := serviceRaw.ToIDispatch()
		defer service.Release()

		// result is a SWBemObjectSet
		resultRaw, err := oleutil.CallMethod(service, "ExecQuery", query)
		if err != nil {
			return err
		}
		result := resultRaw.ToIDispatch()
		defer result.Release()

		count, err := oleInt64(result, "Count")
		if err != nil {
			return err
		}

		var errFieldMismatch error
		for i := int64(0); i < count; i++ {
			err := func() error {
				// item is a SWbemObject, but really a Win32_Process
				itemRaw, err := oleutil.CallMethod(result, "ItemIndex", i)
				if err != nil {
					return err
				}
				item := itemRaw.ToIDispatch()
				defer item.Release()

				ev := reflect.New(elemType)
				if err = loadEntity(ev.Interface(), item); err != nil {
					if _, ok := err.(*ErrFieldMismatch); ok {
						// We continue loading entities even in the face of field mismatch errors.
						// If we encounter any other error, that other error is returned. Otherwise,
						// an ErrFieldMismatch is returned.
						errFieldMismatch = err
					} else {
						return err
					}
				}
				if mat != multiArgTypeStructPtr {
					ev = ev.Elem()
				}
				dv.Set(reflect.Append(dv, ev))
				return nil
			}()
			if err != nil {
				return err
			}
		}
		return errFieldMismatch
	}
	r := do(f)
	if r != nil && queryErr == nil {
		queryErr = r
	}
	return
}

// ErrFieldMismatch is returned when a field is to be loaded into a different
// type than the one it was stored from, or when a field is missing or
// unexported in the destination struct.
// StructType is the type of the struct pointed to by the destination argument.
type ErrFieldMismatch struct {
	StructType reflect.Type
	FieldName  string
	Reason     string
}

func (e *ErrFieldMismatch) Error() string {
	return fmt.Sprintf("wmi: cannot load field %q into a %q: %s",
		e.FieldName, e.StructType, e.Reason)
}

var timeType = reflect.TypeOf(time.Time{})

// loadEntity loads a SWbemObject into a struct pointer.
func loadEntity(dst interface{}, src *ole.IDispatch) (errFieldMismatch error) {
	v := reflect.ValueOf(dst).Elem()
	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		n := v.Type().Field(i).Name
		if !f.CanSet() {
			return &ErrFieldMismatch{
				StructType: f.Type(),
				FieldName:  n,
				Reason:     "CanSet() is false",
			}
		}
		prop, err := oleutil.GetProperty(src, n)
		if err != nil {
			errFieldMismatch = &ErrFieldMismatch{
				StructType: f.Type(),
				FieldName:  n,
				Reason:     "no such struct field",
			}
			continue
		}
		switch val := prop.Value(); reflect.ValueOf(val).Kind() {
		case reflect.Int64:
			iv := val.(int64)
			switch f.Kind() {
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				f.SetInt(iv)
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				f.SetUint(uint64(iv))
			default:
				return &ErrFieldMismatch{
					StructType: f.Type(),
					FieldName:  n,
					Reason:     "not an integer class",
				}
			}
		case reflect.String:
			sv := val.(string)
			iv, err := strconv.ParseInt(sv, 10, 64)
			switch f.Kind() {
			case reflect.String:
				f.SetString(sv)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				if err != nil {
					return err
				}
				f.SetInt(iv)
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				if err != nil {
					return err
				}
				f.SetUint(uint64(iv))
			case reflect.Struct:
				switch f.Type() {
				case timeType:
					if len(sv) == 25 {
						sv = sv[:22] + "0" + sv[22:]
					}
					t, err := time.Parse("20060102150405.000000-0700", sv)
					if err != nil {
						return err
					}
					f.Set(reflect.ValueOf(t))
				}
			}
		case reflect.Bool:
			bv := val.(bool)
			switch f.Kind() {
			case reflect.Bool:
				f.SetBool(bv)
			default:
				return &ErrFieldMismatch{
					StructType: f.Type(),
					FieldName:  n,
					Reason:     "not a bool",
				}
			}
		}
	}
	return errFieldMismatch
}

type multiArgType int

const (
	multiArgTypeInvalid multiArgType = iota
	multiArgTypeStruct
	multiArgTypeStructPtr
)

// checkMultiArg checks that v has type []S, []*S for some struct type S.
//
// It returns what category the slice's elements are, and the reflect.Type
// that represents S.
func checkMultiArg(v reflect.Value) (m multiArgType, elemType reflect.Type) {
	if v.Kind() != reflect.Slice {
		return multiArgTypeInvalid, nil
	}
	elemType = v.Type().Elem()
	switch elemType.Kind() {
	case reflect.Struct:
		return multiArgTypeStruct, elemType
	case reflect.Ptr:
		elemType = elemType.Elem()
		if elemType.Kind() == reflect.Struct {
			return multiArgTypeStructPtr, elemType
		}
	}
	return multiArgTypeInvalid, nil
}

func oleInt64(item *ole.IDispatch, prop string) (int64, error) {
	v, err := oleutil.GetProperty(item, prop)
	if err != nil {
		return 0, err
	}
	i := int64(v.Val)
	return i, nil
}

func CreateQuery(src interface{}, where string) string {
	var b bytes.Buffer
	b.WriteString("SELECT ")
	s := reflect.Indirect(reflect.ValueOf(src))
	t := s.Type()
	if s.Kind() == reflect.Slice {
		t = t.Elem()
	}
	if t.Kind() != reflect.Struct {
		return ""
	}
	var fields []string
	for i := 0; i < t.NumField(); i++ {
		fields = append(fields, t.Field(i).Name)
	}
	b.WriteString(strings.Join(fields, ", "))
	b.WriteString(" FROM ")
	b.WriteString(t.Name())
	b.WriteString(" " + where)
	return b.String()
}
<|endoftext|>


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

package dns

import (
	"net"
	"time"
)

// Envelope is used when doing a transfer with a remote server.
type Envelope struct {
	RR    []RR  // The set of RRs in the answer section of the AXFR reply message.
	Error error // If something went wrong, this contains the error.
}

type Transfer struct {
	Conn
	DialTimeout    time.Duration // net.DialTimeout (ns), defaults to 2 * 1e9
	ReadTimeout    time.Duration // net.Conn.SetReadTimeout value for connections (ns), defaults to 2 * 1e9
	WriteTimeout   time.Duration // net.Conn.SetWriteTimeout value for connections (ns), defaults to 2 * 1e9
	tsigTimersOnly bool
}

// In performs a [AI]XFR request (depends on the message's Qtype). It returns
// a channel of *Envelope on which the replies from the server are sent.
// At the end of the transfer the channel is closed.
// The messages are TSIG checked if needed, no other post-processing is performed.
// The caller must dissect the returned messages.
//
// Basic use pattern for receiving an AXFR:
//
//	// m contains the AXFR request
//	t := new(dns.Transfer)
//	c, e := t.In(m, "127.0.0.1:53")
//	for env := range c
//		// ... deal with env.RR or env.Error
//	}

func (t *Transfer) In(q *Msg, a string, env chan *Envelope) (err error) {
	co := new(Conn)
	timeout := dnsTimeout
	if t.DialTimeout != 0 {
		timeout = t.DialTimeout
	}
	co.Conn, err = net.DialTimeout("tcp", a, timeout)
	if err != nil {
		return err
	}
	if q.Question[0].Qtype == TypeAXFR {
		go t.inAxfr(q.Id, env)
		return nil
	}
	if q.Question[0].Qtype == TypeIXFR {
		go t.inIxfr(q.Id, env)
		return nil
	}
	return nil // TODO(miek): some error
}

func (t *Transfer) inAxfr(id uint16, c chan *Envelope) {
	first := true
	defer t.Close()
	defer close(c)
	timeout := dnsTimeout
	if t.ReadTimeout != 0 {
		timeout = t.ReadTimeout
	}
	for {
		t.SetReadDeadline(time.Now().Add(timeout))
		in, err := t.ReadMsg()
		if err != nil {
			c <- &Envelope{nil, err}
			return
		}
		if id != in.Id {
			c <- &Envelope{in.Answer, ErrId}
			return
		}
		if first {
			if !isSOAFirst(in) {
				c <- &Envelope{in.Answer, ErrSoa}
				return
			}
			first = !first
			// only one answer that is SOA, receive more
			if len(in.Answer) == 1 {
				t.tsigTimersOnly = true
				c <- &Envelope{in.Answer, nil}
				continue
			}
		}

		if !first {
			t.tsigTimersOnly = true // Subsequent envelopes use this.
			if isSOALast(in) {
				c <- &Envelope{in.Answer, nil}
				return
			}
			c <- &Envelope{in.Answer, nil}
		}
	}
	panic("dns: not reached")
}

func (t *Transfer) inIxfr(id uint16, c chan *Envelope) {
	serial := uint32(0) // The first serial seen is the current server serial
	first := true
	defer t.Close()
	defer close(c)
	timeout := dnsTimeout
	if t.ReadTimeout != 0 {
		timeout = t.ReadTimeout
	}
	for {
		// re-read 'n stuff must be pushed down
		t.SetReadDeadline(time.Now().Add(timeout))
		in, err := t.ReadMsg()
		if err != nil {
			c <- &Envelope{in.Answer, err}
			return
		}
		if id != in.Id {
			c <- &Envelope{in.Answer, ErrId}
			return
		}
		if first {
			// A single SOA RR signals "no changes"
			if len(in.Answer) == 1 && isSOAFirst(in) {
				c <- &Envelope{in.Answer, nil}
				return
			}

			// Check if the returned answer is ok
			if !isSOAFirst(in) {
				c <- &Envelope{in.Answer, ErrSoa}
				return
			}
			// This serial is important
			serial = in.Answer[0].(*SOA).Serial
			first = !first
			// continue // TODO(miek)
		}

		// Now we need to check each message for SOA records, to see what we need to do
		if !first {
			t.tsigTimersOnly = true
			// If the last record in the IXFR contains the servers' SOA,  we should quit
			if v, ok := in.Answer[len(in.Answer)-1].(*SOA); ok {
				if v.Serial == serial {
					c <- &Envelope{in.Answer, nil}
					return
				}
			}
			c <- &Envelope{in.Answer, nil}
		}
	}
}

func (t *Transfer) Out(w ResponseWriter, q *Msg, a string) (chan *Envelope, error) {
	ch := make(chan *Envelope)
	r := new(Msg)
	r.SetReply(q)
	r.Authoritative = true
	go func() {
		for x := range ch {
			// assume it fits TODO(miek): fix
			r.Answer = append(r.Answer, x.RR...)
			if err := w.WriteMsg(r); err != nil {
				return
			}
		}
		//		w.TsigTimersOnly(true)
		//		rep.Answer = nil
	}()
	return ch, nil
}

// ReadMsg reads a message from the transfer connection t.
func (t *Transfer) ReadMsg() (*Msg, error) {
	m := new(Msg)
	p := make([]byte, MaxMsgSize)
	n, err := t.Read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	if ts := m.IsTsig(); t != nil {
		if _, ok := t.TsigSecret[ts.Hdr.Name]; !ok {
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		err = TsigVerify(p, t.TsigSecret[ts.Hdr.Name], t.tsigRequestMAC, t.tsigTimersOnly)
	}
	return m, err
}

// WriteMsg write a message throught the transfer connection t.
func (t *Transfer) WriteMsg(m *Msg) (err error) {
	var out []byte
	if ts := m.IsTsig(); t != nil {
		if _, ok := t.TsigSecret[ts.Hdr.Name]; !ok {
			return ErrSecret
		}
		out, t.tsigRequestMAC, err = TsigGenerate(m, t.TsigSecret[ts.Hdr.Name], t.tsigRequestMAC, t.tsigTimersOnly)
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	if _, err = t.Write(out); err != nil {
		return err
	}
	return nil
}

/*

*/

func isSOAFirst(in *Msg) bool {
	if len(in.Answer) > 0 {
		return in.Answer[0].Header().Rrtype == TypeSOA
	}
	return false
}

func isSOALast(in *Msg) bool {
	if len(in.Answer) > 0 {
		return in.Answer[len(in.Answer)-1].Header().Rrtype == TypeSOA
	}
	return false
}
<commit_msg>Fix xfr some more<commit_after>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package dns

import (
	"net"
	"time"
)

// Envelope is used when doing a transfer with a remote server.
type Envelope struct {
	RR    []RR  // The set of RRs in the answer section of the AXFR reply message.
	Error error // If something went wrong, this contains the error.
}

type Transfer struct {
	Conn
	DialTimeout    time.Duration // net.DialTimeout (ns), defaults to 2 * 1e9
	ReadTimeout    time.Duration // net.Conn.SetReadTimeout value for connections (ns), defaults to 2 * 1e9
	WriteTimeout   time.Duration // net.Conn.SetWriteTimeout value for connections (ns), defaults to 2 * 1e9
	tsigTimersOnly bool
}

// In performs an incoming transfer with the server in a.
func (t *Transfer) In(q *Msg, a string) (env chan *Envelope, err error) {
	co := new(Conn)
	timeout := dnsTimeout
	if t.DialTimeout != 0 {
		timeout = t.DialTimeout
	}
	co.Conn, err = net.DialTimeout("tcp", a, timeout)
	if err != nil {
		return nil, err
	}
	env = make(chan *Envelope)
	go func() {
		if q.Question[0].Qtype == TypeAXFR {
			go t.inAxfr(q.Id, env)
			return
		}
		if q.Question[0].Qtype == TypeIXFR {
			go t.inIxfr(q.Id, env)
			return
		}
	}()
	return env, nil
}

func (t *Transfer) inAxfr(id uint16, c chan *Envelope) {
	first := true
	defer t.Close()
	defer close(c)
	timeout := dnsTimeout
	if t.ReadTimeout != 0 {
		timeout = t.ReadTimeout
	}
	for {
		t.SetReadDeadline(time.Now().Add(timeout))
		in, err := t.ReadMsg()
		if err != nil {
			c <- &Envelope{nil, err}
			return
		}
		if id != in.Id {
			c <- &Envelope{in.Answer, ErrId}
			return
		}
		if first {
			if !isSOAFirst(in) {
				c <- &Envelope{in.Answer, ErrSoa}
				return
			}
			first = !first
			// only one answer that is SOA, receive more
			if len(in.Answer) == 1 {
				t.tsigTimersOnly = true
				c <- &Envelope{in.Answer, nil}
				continue
			}
		}

		if !first {
			t.tsigTimersOnly = true // Subsequent envelopes use this.
			if isSOALast(in) {
				c <- &Envelope{in.Answer, nil}
				return
			}
			c <- &Envelope{in.Answer, nil}
		}
	}
	panic("dns: not reached")
}

func (t *Transfer) inIxfr(id uint16, c chan *Envelope) {
	serial := uint32(0) // The first serial seen is the current server serial
	first := true
	defer t.Close()
	defer close(c)
	timeout := dnsTimeout
	if t.ReadTimeout != 0 {
		timeout = t.ReadTimeout
	}
	for {
		t.SetReadDeadline(time.Now().Add(timeout))
		in, err := t.ReadMsg()
		if err != nil {
			c <- &Envelope{in.Answer, err}
			return
		}
		if id != in.Id {
			c <- &Envelope{in.Answer, ErrId}
			return
		}
		if first {
			// A single SOA RR signals "no changes"
			if len(in.Answer) == 1 && isSOAFirst(in) {
				c <- &Envelope{in.Answer, nil}
				return
			}

			// Check if the returned answer is ok
			if !isSOAFirst(in) {
				c <- &Envelope{in.Answer, ErrSoa}
				return
			}
			// This serial is important
			serial = in.Answer[0].(*SOA).Serial
			first = !first
			// continue // TODO(miek)
		}

		// Now we need to check each message for SOA records, to see what we need to do
		if !first {
			t.tsigTimersOnly = true
			// If the last record in the IXFR contains the servers' SOA,  we should quit
			if v, ok := in.Answer[len(in.Answer)-1].(*SOA); ok {
				if v.Serial == serial {
					c <- &Envelope{in.Answer, nil}
					return
				}
			}
			c <- &Envelope{in.Answer, nil}
		}
	}
}

func (t *Transfer) Out(w ResponseWriter, q *Msg, a string) (chan *Envelope, error) {
	ch := make(chan *Envelope)
	r := new(Msg)
	r.SetReply(q)
	r.Authoritative = true
	go func() {
		for x := range ch {
			// assume it fits TODO(miek): fix
			r.Answer = append(r.Answer, x.RR...)
			if err := w.WriteMsg(r); err != nil {
				return
			}
		}
		w.TsigTimersOnly(true)
		r.Answer = nil
	}()
	return ch, nil
}

// ReadMsg reads a message from the transfer connection t.
func (t *Transfer) ReadMsg() (*Msg, error) {
	m := new(Msg)
	p := make([]byte, MaxMsgSize)
	n, err := t.Read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	if ts := m.IsTsig(); t != nil {
		if _, ok := t.TsigSecret[ts.Hdr.Name]; !ok {
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		err = TsigVerify(p, t.TsigSecret[ts.Hdr.Name], t.tsigRequestMAC, t.tsigTimersOnly)
	}
	return m, err
}

// WriteMsg write a message throught the transfer connection t.
func (t *Transfer) WriteMsg(m *Msg) (err error) {
	var out []byte
	if ts := m.IsTsig(); t != nil {
		if _, ok := t.TsigSecret[ts.Hdr.Name]; !ok {
			return ErrSecret
		}
		out, t.tsigRequestMAC, err = TsigGenerate(m, t.TsigSecret[ts.Hdr.Name], t.tsigRequestMAC, t.tsigTimersOnly)
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	if _, err = t.Write(out); err != nil {
		return err
	}
	return nil
}

/*

*/

func isSOAFirst(in *Msg) bool {
	if len(in.Answer) > 0 {
		return in.Answer[0].Header().Rrtype == TypeSOA
	}
	return false
}

func isSOALast(in *Msg) bool {
	if len(in.Answer) > 0 {
		return in.Answer[len(in.Answer)-1].Header().Rrtype == TypeSOA
	}
	return false
}
<|endoftext|>


<commit_before>package jsonfilelog

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/fsnotify/fsnotify"
	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/daemon/logger"
	"github.com/docker/docker/pkg/filenotify"
	"github.com/docker/docker/pkg/ioutils"
	"github.com/docker/docker/pkg/jsonlog"
	"github.com/docker/docker/pkg/tailfile"
)

const maxJSONDecodeRetry = 20000

func decodeLogLine(dec *json.Decoder, l *jsonlog.JSONLog) (*logger.Message, error) {
	l.Reset()
	if err := dec.Decode(l); err != nil {
		return nil, err
	}
	msg := &logger.Message{
		Source:    l.Stream,
		Timestamp: l.Created,
		Line:      []byte(l.Log),
		Attrs:     l.Attrs,
	}
	return msg, nil
}

// ReadLogs implements the logger's LogReader interface for the logs
// created by this driver.
func (l *JSONFileLogger) ReadLogs(config logger.ReadConfig) *logger.LogWatcher {
	logWatcher := logger.NewLogWatcher()

	go l.readLogs(logWatcher, config)
	return logWatcher
}

func (l *JSONFileLogger) readLogs(logWatcher *logger.LogWatcher, config logger.ReadConfig) {
	defer close(logWatcher.Msg)

	// lock so the read stream doesn't get corrupted due to rotations or other log data written while we read
	// This will block writes!!!
	l.mu.Lock()

	pth := l.writer.LogPath()
	var files []io.ReadSeeker
	for i := l.writer.MaxFiles(); i > 1; i-- {
		f, err := os.Open(fmt.Sprintf("%s.%d", pth, i-1))
		if err != nil {
			if !os.IsNotExist(err) {
				logWatcher.Err <- err
				break
			}
			continue
		}
		defer f.Close()

		files = append(files, f)
	}

	latestFile, err := os.Open(pth)
	if err != nil {
		logWatcher.Err <- err
		l.mu.Unlock()
		return
	}
	defer latestFile.Close()

	if config.Tail != 0 {
		tailer := ioutils.MultiReadSeeker(append(files, latestFile)...)
		tailFile(tailer, logWatcher, config.Tail, config.Since)
	}

	// close all the rotated files
	for _, f := range files {
		if err := f.(io.Closer).Close(); err != nil {
			logrus.WithField("logger", "json-file").Warnf("error closing tailed log file: %v", err)
		}
	}

	if !config.Follow {
		if err := latestFile.Close(); err != nil {
			logrus.Errorf("Error closing file: %v", err)
		}
		l.mu.Unlock()
		return
	}

	if config.Tail >= 0 {
		latestFile.Seek(0, os.SEEK_END)
	}

	l.readers[logWatcher] = struct{}{}
	l.mu.Unlock()

	notifyRotate := l.writer.NotifyRotate()
	followLogs(latestFile, logWatcher, notifyRotate, config.Since)

	l.mu.Lock()
	delete(l.readers, logWatcher)
	l.mu.Unlock()

	l.writer.NotifyRotateEvict(notifyRotate)
}

func tailFile(f io.ReadSeeker, logWatcher *logger.LogWatcher, tail int, since time.Time) {
	var rdr io.Reader
	rdr = f
	if tail > 0 {
		ls, err := tailfile.TailFile(f, tail)
		if err != nil {
			logWatcher.Err <- err
			return
		}
		rdr = bytes.NewBuffer(bytes.Join(ls, []byte("\n")))
	}
	dec := json.NewDecoder(rdr)
	l := &jsonlog.JSONLog{}
	for {
		msg, err := decodeLogLine(dec, l)
		if err != nil {
			if err != io.EOF {
				logWatcher.Err <- err
			}
			return
		}
		if !since.IsZero() && msg.Timestamp.Before(since) {
			continue
		}
		select {
		case <-logWatcher.WatchClose():
			return
		case logWatcher.Msg <- msg:
		}
	}
}

func watchFile(name string) (filenotify.FileWatcher, error) {
	fileWatcher, err := filenotify.New()
	if err != nil {
		return nil, err
	}

	if err := fileWatcher.Add(name); err != nil {
		logrus.WithField("logger", "json-file").Warnf("falling back to file poller due to error: %v", err)
		fileWatcher.Close()
		fileWatcher = filenotify.NewPollingWatcher()

		if err := fileWatcher.Add(name); err != nil {
			fileWatcher.Close()
			logrus.Debugf("error watching log file for modifications: %v", err)
			return nil, err
		}
	}
	return fileWatcher, nil
}

func followLogs(f *os.File, logWatcher *logger.LogWatcher, notifyRotate chan interface{}, since time.Time) {
	dec := json.NewDecoder(f)
	l := &jsonlog.JSONLog{}

	name := f.Name()
	fileWatcher, err := watchFile(name)
	if err != nil {
		logWatcher.Err <- err
		return
	}
	defer func() {
		f.Close()
		fileWatcher.Remove(name)
		fileWatcher.Close()
	}()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		select {
		case <-logWatcher.WatchClose():
			fileWatcher.Remove(name)
			cancel()
		case <-ctx.Done():
			return
		}
	}()

	var retries int
	handleRotate := func() error {
		f.Close()
		fileWatcher.Remove(name)

		// retry when the file doesn't exist
		for retries := 0; retries <= 5; retries++ {
			f, err = os.Open(name)
			if err == nil || !os.IsNotExist(err) {
				break
			}
		}
		if err != nil {
			return err
		}
		if err := fileWatcher.Add(name); err != nil {
			return err
		}
		dec = json.NewDecoder(f)
		return nil
	}

	errRetry := errors.New("retry")
	errDone := errors.New("done")
	waitRead := func() error {
		select {
		case e := <-fileWatcher.Events():
			switch e.Op {
			case fsnotify.Write:
				dec = json.NewDecoder(f)
				return nil
			case fsnotify.Rename, fsnotify.Remove:
				select {
				case <-notifyRotate:
				case <-ctx.Done():
					return errDone
				}
				if err := handleRotate(); err != nil {
					return err
				}
				return nil
			}
			return errRetry
		case err := <-fileWatcher.Errors():
			logrus.Debug("logger got error watching file: %v", err)
			// Something happened, let's try and stay alive and create a new watcher
			if retries <= 5 {
				fileWatcher.Close()
				fileWatcher, err = watchFile(name)
				if err != nil {
					return err
				}
				retries++
				return errRetry
			}
			return err
		case <-ctx.Done():
			return errDone
		}
	}

	handleDecodeErr := func(err error) error {
		if err == io.EOF {
			for err := waitRead(); err != nil; {
				if err == errRetry {
					// retry the waitRead
					continue
				}
				return err
			}
			return nil
		}
		// try again because this shouldn't happen
		if _, ok := err.(*json.SyntaxError); ok && retries <= maxJSONDecodeRetry {
			dec = json.NewDecoder(f)
			retries++
			return nil
		}
		// io.ErrUnexpectedEOF is returned from json.Decoder when there is
		// remaining data in the parser's buffer while an io.EOF occurs.
		// If the json logger writes a partial json log entry to the disk
		// while at the same time the decoder tries to decode it, the race condition happens.
		if err == io.ErrUnexpectedEOF && retries <= maxJSONDecodeRetry {
			reader := io.MultiReader(dec.Buffered(), f)
			dec = json.NewDecoder(reader)
			retries++
			return nil
		}
		return err
	}

	// main loop
	for {
		msg, err := decodeLogLine(dec, l)
		if err != nil {
			if err := handleDecodeErr(err); err != nil {
				if err == errDone {
					return
				}
				// we got an unrecoverable error, so return
				logWatcher.Err <- err
				return
			}
			// ready to try again
			continue
		}

		retries = 0 // reset retries since we've succeeded
		if !since.IsZero() && msg.Timestamp.Before(since) {
			continue
		}
		select {
		case logWatcher.Msg <- msg:
		case <-ctx.Done():
			logWatcher.Msg <- msg
			for {
				msg, err := decodeLogLine(dec, l)
				if err != nil {
					return
				}
				if !since.IsZero() && msg.Timestamp.Before(since) {
					continue
				}
				logWatcher.Msg <- msg
			}
		}
	}
}
<commit_msg>Fix cpu spin waiting for log write events<commit_after>package jsonfilelog

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/fsnotify/fsnotify"
	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/daemon/logger"
	"github.com/docker/docker/pkg/filenotify"
	"github.com/docker/docker/pkg/ioutils"
	"github.com/docker/docker/pkg/jsonlog"
	"github.com/docker/docker/pkg/tailfile"
)

const maxJSONDecodeRetry = 20000

func decodeLogLine(dec *json.Decoder, l *jsonlog.JSONLog) (*logger.Message, error) {
	l.Reset()
	if err := dec.Decode(l); err != nil {
		return nil, err
	}
	msg := &logger.Message{
		Source:    l.Stream,
		Timestamp: l.Created,
		Line:      []byte(l.Log),
		Attrs:     l.Attrs,
	}
	return msg, nil
}

// ReadLogs implements the logger's LogReader interface for the logs
// created by this driver.
func (l *JSONFileLogger) ReadLogs(config logger.ReadConfig) *logger.LogWatcher {
	logWatcher := logger.NewLogWatcher()

	go l.readLogs(logWatcher, config)
	return logWatcher
}

func (l *JSONFileLogger) readLogs(logWatcher *logger.LogWatcher, config logger.ReadConfig) {
	defer close(logWatcher.Msg)

	// lock so the read stream doesn't get corrupted due to rotations or other log data written while we read
	// This will block writes!!!
	l.mu.Lock()

	pth := l.writer.LogPath()
	var files []io.ReadSeeker
	for i := l.writer.MaxFiles(); i > 1; i-- {
		f, err := os.Open(fmt.Sprintf("%s.%d", pth, i-1))
		if err != nil {
			if !os.IsNotExist(err) {
				logWatcher.Err <- err
				break
			}
			continue
		}
		defer f.Close()

		files = append(files, f)
	}

	latestFile, err := os.Open(pth)
	if err != nil {
		logWatcher.Err <- err
		l.mu.Unlock()
		return
	}
	defer latestFile.Close()

	if config.Tail != 0 {
		tailer := ioutils.MultiReadSeeker(append(files, latestFile)...)
		tailFile(tailer, logWatcher, config.Tail, config.Since)
	}

	// close all the rotated files
	for _, f := range files {
		if err := f.(io.Closer).Close(); err != nil {
			logrus.WithField("logger", "json-file").Warnf("error closing tailed log file: %v", err)
		}
	}

	if !config.Follow {
		if err := latestFile.Close(); err != nil {
			logrus.Errorf("Error closing file: %v", err)
		}
		l.mu.Unlock()
		return
	}

	if config.Tail >= 0 {
		latestFile.Seek(0, os.SEEK_END)
	}

	l.readers[logWatcher] = struct{}{}
	l.mu.Unlock()

	notifyRotate := l.writer.NotifyRotate()
	followLogs(latestFile, logWatcher, notifyRotate, config.Since)

	l.mu.Lock()
	delete(l.readers, logWatcher)
	l.mu.Unlock()

	l.writer.NotifyRotateEvict(notifyRotate)
}

func tailFile(f io.ReadSeeker, logWatcher *logger.LogWatcher, tail int, since time.Time) {
	var rdr io.Reader
	rdr = f
	if tail > 0 {
		ls, err := tailfile.TailFile(f, tail)
		if err != nil {
			logWatcher.Err <- err
			return
		}
		rdr = bytes.NewBuffer(bytes.Join(ls, []byte("\n")))
	}
	dec := json.NewDecoder(rdr)
	l := &jsonlog.JSONLog{}
	for {
		msg, err := decodeLogLine(dec, l)
		if err != nil {
			if err != io.EOF {
				logWatcher.Err <- err
			}
			return
		}
		if !since.IsZero() && msg.Timestamp.Before(since) {
			continue
		}
		select {
		case <-logWatcher.WatchClose():
			return
		case logWatcher.Msg <- msg:
		}
	}
}

func watchFile(name string) (filenotify.FileWatcher, error) {
	fileWatcher, err := filenotify.New()
	if err != nil {
		return nil, err
	}

	if err := fileWatcher.Add(name); err != nil {
		logrus.WithField("logger", "json-file").Warnf("falling back to file poller due to error: %v", err)
		fileWatcher.Close()
		fileWatcher = filenotify.NewPollingWatcher()

		if err := fileWatcher.Add(name); err != nil {
			fileWatcher.Close()
			logrus.Debugf("error watching log file for modifications: %v", err)
			return nil, err
		}
	}
	return fileWatcher, nil
}

func followLogs(f *os.File, logWatcher *logger.LogWatcher, notifyRotate chan interface{}, since time.Time) {
	dec := json.NewDecoder(f)
	l := &jsonlog.JSONLog{}

	name := f.Name()
	fileWatcher, err := watchFile(name)
	if err != nil {
		logWatcher.Err <- err
		return
	}
	defer func() {
		f.Close()
		fileWatcher.Remove(name)
		fileWatcher.Close()
	}()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		select {
		case <-logWatcher.WatchClose():
			fileWatcher.Remove(name)
			cancel()
		case <-ctx.Done():
			return
		}
	}()

	var retries int
	handleRotate := func() error {
		f.Close()
		fileWatcher.Remove(name)

		// retry when the file doesn't exist
		for retries := 0; retries <= 5; retries++ {
			f, err = os.Open(name)
			if err == nil || !os.IsNotExist(err) {
				break
			}
		}
		if err != nil {
			return err
		}
		if err := fileWatcher.Add(name); err != nil {
			return err
		}
		dec = json.NewDecoder(f)
		return nil
	}

	errRetry := errors.New("retry")
	errDone := errors.New("done")
	waitRead := func() error {
		select {
		case e := <-fileWatcher.Events():
			switch e.Op {
			case fsnotify.Write:
				dec = json.NewDecoder(f)
				return nil
			case fsnotify.Rename, fsnotify.Remove:
				select {
				case <-notifyRotate:
				case <-ctx.Done():
					return errDone
				}
				if err := handleRotate(); err != nil {
					return err
				}
				return nil
			}
			return errRetry
		case err := <-fileWatcher.Errors():
			logrus.Debug("logger got error watching file: %v", err)
			// Something happened, let's try and stay alive and create a new watcher
			if retries <= 5 {
				fileWatcher.Close()
				fileWatcher, err = watchFile(name)
				if err != nil {
					return err
				}
				retries++
				return errRetry
			}
			return err
		case <-ctx.Done():
			return errDone
		}
	}

	handleDecodeErr := func(err error) error {
		if err == io.EOF {
			for {
				err := waitRead()
				if err == nil {
					break
				}
				if err == errRetry {
					continue
				}
				return err
			}
			return nil
		}
		// try again because this shouldn't happen
		if _, ok := err.(*json.SyntaxError); ok && retries <= maxJSONDecodeRetry {
			dec = json.NewDecoder(f)
			retries++
			return nil
		}
		// io.ErrUnexpectedEOF is returned from json.Decoder when there is
		// remaining data in the parser's buffer while an io.EOF occurs.
		// If the json logger writes a partial json log entry to the disk
		// while at the same time the decoder tries to decode it, the race condition happens.
		if err == io.ErrUnexpectedEOF && retries <= maxJSONDecodeRetry {
			reader := io.MultiReader(dec.Buffered(), f)
			dec = json.NewDecoder(reader)
			retries++
			return nil
		}
		return err
	}

	// main loop
	for {
		msg, err := decodeLogLine(dec, l)
		if err != nil {
			if err := handleDecodeErr(err); err != nil {
				if err == errDone {
					return
				}
				// we got an unrecoverable error, so return
				logWatcher.Err <- err
				return
			}
			// ready to try again
			continue
		}

		retries = 0 // reset retries since we've succeeded
		if !since.IsZero() && msg.Timestamp.Before(since) {
			continue
		}
		select {
		case logWatcher.Msg <- msg:
		case <-ctx.Done():
			logWatcher.Msg <- msg
			for {
				msg, err := decodeLogLine(dec, l)
				if err != nil {
					return
				}
				if !since.IsZero() && msg.Timestamp.Before(since) {
					continue
				}
				logWatcher.Msg <- msg
			}
		}
	}
}
<|endoftext|>


<commit_before>package gorm

import (
	"github.com/gkarlik/quark-go/data/access/rdbms"
	"github.com/gkarlik/quark-go/logger"
	"github.com/jinzhu/gorm"
)

// DbContext represents relational database access mechanism using GORM library.
type DbContext struct {
	DB              *gorm.DB // raw gorm database object
	isInTransaction bool     // indicates if database context has already started the database transaction
}

// NewDbContext creates DbContext instance with specified SQL dialect and connection string.
// It should be create once per goroutine and reused. It opens database connection which should be closed in defer function using Dispose method.
func NewDbContext(dialect string, connStr string) (*DbContext, error) {
	logger.Log().DebugWithFields(logger.LogFields{
		"dialect":           dialect,
		"connection_string": connStr,
		"component":         componentName,
	}, "Creating new database context")

	db, err := gorm.Open(dialect, connStr)
	if err != nil {
		logger.Log().ErrorWithFields(logger.LogFields{
			"error":     err,
			"component": componentName,
		}, "Cannot create database context")

		return nil, err
	}

	return &DbContext{
		DB:              db,
		isInTransaction: false,
	}, nil
}

// IsInTransaction indicates if DbContext has started DbTransaction.
func (c *DbContext) IsInTransaction() bool {
	return c.isInTransaction
}

// BeginTransaction starts new database transaction. Panics if DbContext is already in transaction.
func (c *DbContext) BeginTransaction() rdbms.DbTransaction {
	if c.IsInTransaction() {
		logger.Log().PanicWithFields(logger.LogFields{
			"component": componentName,
		}, "DbContext is already in transaction. Nested transactions are not supported!")
	}

	db := c.DB.Begin()

	return &DbTransaction{
		context: &DbContext{
			DB:              db,
			isInTransaction: true,
		},
	}
}

// Dispose closes database context and cleans up DbContext instance.
func (c *DbContext) Dispose() {
	logger.Log().InfoWithFields(logger.LogFields{"component": componentName}, "Closing database context")

	if c.DB != nil {
		err := c.DB.Close()
		if err != nil {
			logger.Log().ErrorWithFields(logger.LogFields{
				"error":     err,
				"component": componentName,
			}, "Cannot close database context")
		}
		c.DB = nil
	}
}
<commit_msg>Fixing log message.<commit_after>package gorm

import (
	"github.com/gkarlik/quark-go/data/access/rdbms"
	"github.com/gkarlik/quark-go/logger"
	"github.com/jinzhu/gorm"
)

// DbContext represents relational database access mechanism using GORM library.
type DbContext struct {
	DB              *gorm.DB // raw gorm database object
	isInTransaction bool     // indicates if database context has already started the database transaction
}

// NewDbContext creates DbContext instance with specified SQL dialect and connection string.
// It should be create once per goroutine and reused. It opens database connection which should be closed in defer function using Dispose method.
func NewDbContext(dialect string, connStr string) (*DbContext, error) {
	logger.Log().DebugWithFields(logger.LogFields{
		"dialect":           dialect,
		"connection_string": connStr,
		"component":         componentName,
	}, "Creating new database context")

	db, err := gorm.Open(dialect, connStr)
	if err != nil {
		logger.Log().ErrorWithFields(logger.LogFields{
			"error":     err,
			"component": componentName,
		}, "Cannot create database context")

		return nil, err
	}

	return &DbContext{
		DB:              db,
		isInTransaction: false,
	}, nil
}

// IsInTransaction indicates if DbContext has started DbTransaction.
func (c *DbContext) IsInTransaction() bool {
	return c.isInTransaction
}

// BeginTransaction starts new database transaction. Panics if DbContext is already in transaction.
func (c *DbContext) BeginTransaction() rdbms.DbTransaction {
	if c.IsInTransaction() {
		logger.Log().PanicWithFields(logger.LogFields{
			"component": componentName,
		}, "DbContext is already in transaction. Nested transactions are not supported!")
	}

	db := c.DB.Begin()

	return &DbTransaction{
		context: &DbContext{
			DB:              db,
			isInTransaction: true,
		},
	}
}

// Dispose closes database context and cleans up DbContext instance.
func (c *DbContext) Dispose() {
	logger.Log().InfoWithFields(logger.LogFields{"component": componentName}, "Disposing database context")

	if c.DB != nil {
		err := c.DB.Close()
		if err != nil {
			logger.Log().ErrorWithFields(logger.LogFields{
				"error":     err,
				"component": componentName,
			}, "Cannot close database context")
		}
		c.DB = nil
	}
}
<|endoftext|>


<commit_before>// Yam (Yet Another Mux)

package main

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

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

func DefaultOptionsHandler(route *Route) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		methods := []string{}
		for key, _ := range route.handlers {
			methods = append(methods, key)
		}
		w.Header().Add("Allow", strings.Join(methods, ", "))
	})
}

func DefaultTraceHandler(route *Route) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		dump, _ := httputil.DumpRequest(r, false)
		w.Write(dump)
	})
}

type Config struct {
	Options        bool
	OptionsHandler func(*Route) http.Handler
	Trace          bool
	TraceHandler   func(*Route) http.Handler
}

func NewConfig() *Config {
	return &Config{
		Options:        true,
		OptionsHandler: DefaultOptionsHandler,
		Trace:          false,
		TraceHandler:   DefaultTraceHandler,
	}
}

type Yam struct {
	Root   *Route
	Config *Config
}

func New() *Yam {
	y := &Yam{}
	y.Root = &Route{yam: y}
	y.Config = NewConfig()

	return y
}

func (y *Yam) Route(path string) *Route {
	route := y.Root.Route(path)
	route.yam = y

	return route
}

func (y *Yam) SetConfig(c *Config) {
	y.Config = c
}

func route(path string, router *Route, y *Yam) *Route {
	parts := strings.Split(path, "/")[1:]
	routes := router.Routes

	fmt.Println("Start Router:", router.path)
	fmt.Println("Stat Path:", path)
	fullPath := router.path + path

	for i, part := range parts {
		fmt.Println("Part:", part)
		if i == len(parts)-1 {

			for _, route := range routes {
				if route.leaf == part {
					fmt.Println("Route Exists")
					fmt.Println("--------------")
					return route
				}
			}

			route := &Route{leaf: part, path: fullPath, yam: y}
			fmt.Println("Add:", route.path)
			fmt.Println("Router:", router.path)
			router.Routes = append(router.Routes, route)

			fmt.Println("--------------")

			return route

		} else {
			for _, route := range routes {
				if route.leaf == part {
					fmt.Println("Leaf:", route.leaf)
					router = route
				} else {
					fmt.Println("Router:", router.path)
					route := &Route{leaf: part, path: router.path + path, yam: y}
					router.Routes = append(router.Routes, route)
					router = route
				}
			}

		}
	}

	return nil
}

func (y *Yam) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(r.URL.Path, "/")[1:]
	fmt.Println(parts)
	routes := y.Root.Routes

	for i, part := range parts {
		fmt.Println(part)
		for _, route := range routes {
			fmt.Println("Leaf:", route.leaf)
			match := false
			// Pattern Match
			if strings.HasPrefix(route.leaf, ":") {
				fmt.Println("Pattern Match")
				match = true
				values := url.Values{}
				values.Add(route.leaf, part)
				r.URL.RawQuery = url.Values(values).Encode() + "&" + r.URL.RawQuery
			} else { // Exact match
				fmt.Println("Exact Match")
				if route.leaf == part {
					match = true
				}
			}

			if match {
				fmt.Println("Leaf ==", part)
				if i < len(parts)-1 {
					routes = route.Routes
					break
				} else {
					fmt.Println("Found: ", route.path)

					handler := route.handlers[r.Method]
					if handler != nil {
						handler.ServeHTTP(w, r)
						return
					}

					fmt.Println("No handler for method")
					w.WriteHeader(http.StatusMethodNotAllowed)
					return

				}
			}
		}
	}

	// If we get here then we have not found a route
	fmt.Println("Not Found")
	w.WriteHeader(http.StatusNotFound)
}

type Route struct {
	leaf   string   // a part of a URL path, /foo/bar - a leaf would be foo and bar
	path   string   // full url path
	Routes []*Route // Routes that live under this route

	yam *Yam // Reference to Yam and global configuraion

	handlers map[string]http.Handler
}

func (r *Route) Route(path string) *Route {
	r = route(path, r, r.yam)

	if r.handlers == nil {
		r.handlers = make(map[string]http.Handler)
	}

	if r.yam.Config.Options {
		if r.handlers["OPTIONS"] == nil {
			r.handlers["OPTIONS"] = r.yam.Config.OptionsHandler(r)
		}
	}

	if r.yam.Config.Trace {
		if r.handlers["TRACE"] == nil {
			r.handlers["TRACE"] = r.yam.Config.TraceHandler(r)
		}
	}

	return r
}

func (r *Route) Add(method string, handler http.Handler) *Route {
	r.handlers[method] = handler

	return r
}

func (r *Route) Head(h handler) *Route {
	r.Add("HEAD", http.HandlerFunc(h))

	return r
}

func (r *Route) Get(h handler) *Route {
	r.Add("GET", http.HandlerFunc(h))

	// Implement the HEAD handler by default for all GET requests - HEAD
	// should not return a body so we wrap it in a middleware
	head := func(n http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Serve the handler
			n.ServeHTTP(w, r)
			// Flush the body so we don't write to the client
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		})
	}

	// Apply the head middleware to the head handler
	r.Add("HEAD", head(http.HandlerFunc(h)))

	return r
}

func (r *Route) Post(h handler) *Route {
	r.Add("POST", http.HandlerFunc(h))

	return r
}

func (r *Route) Put(h handler) *Route {
	r.Add("PUT", http.HandlerFunc(h))

	return r
}

func (r *Route) Delete(h handler) *Route {
	r.Add("DELETE", http.HandlerFunc(h))

	return r
}

func (r *Route) Patch(h handler) *Route {
	r.Add("PATCH", http.HandlerFunc(h))

	return r
}

func GetRootHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get Root Handler"))
}

func GetFooHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get Foo Handler"))
}

func GetAHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get A Handler"))
}

func GetBHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get B Handler"))
}

func GetCHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get C Handler"))
}

func GetDHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get D Handler"))
}

func main() {
	y := New()
	y.Config.Trace = true

	y.Route("/").Get(GetRootHandler)
	y.Route("/get").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("GET"))
	})
	y.Route("/post").Post(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("POST"))
	})
	y.Route("/put").Put(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PUT"))
	})
	y.Route("/patch").Patch(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PATCH"))
	})
	y.Route("/delete").Delete(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("DELETE"))
	})

	a := y.Route("/a").Get(GetAHandler)
	a.Route("/b").Get(GetBHandler)
	a.Route("/b").Put(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PUT B Handler"))
	})
	c := a.Route("/b/c").Get(GetCHandler)
	c.Route("/d").Get(GetDHandler)
	e := c.Route("/d/e").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("E Handler"))
	})
	e.Route("/f").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("F Handler"))
	})

	// Pattern Matching
	a.Route("/:foo").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("A :foo Handler\n"))
		w.Write([]byte(r.URL.Query().Get(":foo")))
	})

	bar := a.Route("/:foo/:bar").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("/a/:foo/:bar Handler\n"))
		w.Write([]byte(r.URL.Query().Get(":foo")))
		w.Write([]byte("\n"))
		w.Write([]byte(r.URL.Query().Get(":bar")))
	})

	bar.Route("/baz").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Foo", "Bar")
		w.Write([]byte("baz\n"))
		w.Write([]byte(r.URL.Query().Get(":foo")))
		w.Write([]byte("\n"))
		w.Write([]byte(r.URL.Query().Get(":bar")))
	})

	fmt.Printf("%+v\n", y)

	http.ListenAndServe(":5000", y)
}
<commit_msg>Code cleanup and doc strigs<commit_after>// Yam (Yet Another Mux)

package main

import (
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

// An easy type to reference standard http handler functions
type handler func(http.ResponseWriter, *http.Request)

// Configuration type that allows configuration of YAM
type Config struct {
	Options        bool
	OptionsHandler func(*Route) http.Handler
	Trace          bool
	TraceHandler   func(*Route) http.Handler
}

// Constructs a new Config instance with default values
func NewConfig() *Config {
	return &Config{
		Options:        true,
		OptionsHandler: DefaultOptionsHandler,
		Trace:          false,
		TraceHandler:   DefaultTraceHandler,
	}
}

// Base YAM type contain the Root Routes and Configuration
type Yam struct {
	Root   *Route
	Config *Config
}

// Constructs a new YAM instance with default configuration
func New() *Yam {
	y := &Yam{}
	y.Root = &Route{yam: y}
	y.Config = NewConfig()

	return y
}

// Creates a new base Route - Effectively a constructor for Route
func (y *Yam) Route(path string) *Route {
	route := y.Root.Route(path)
	route.yam = y

	return route
}

// Registers a new Route on the Path, building a Tree structure of Routes
func route(path string, router *Route, y *Yam) *Route {
	parts := strings.Split(path, "/")[1:]
	routes := router.Routes
	fullPath := router.path + path

	for i, part := range parts {
		if i == len(parts)-1 {
			for _, route := range routes {
				if route.leaf == part {
					return route
				}
			}
			route := &Route{leaf: part, path: fullPath, yam: y}
			router.Routes = append(router.Routes, route)
			return route
		} else {
			for _, route := range routes {
				if route.leaf == part {
					router = route
				} else {
					route := &Route{leaf: part, path: router.path + path, yam: y}
					router.Routes = append(router.Routes, route)
					router = route
				}
			}
		}
	}

	return nil
}

// Implements the http.Handler Interface.  Finds the correct handler for
// a path based on the path and http verb of the request.
func (y *Yam) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(r.URL.Path, "/")[1:]
	routes := y.Root.Routes
	for i, part := range parts {
		for _, route := range routes {
			match := false
			// Pattern Match
			if strings.HasPrefix(route.leaf, ":") {
				match = true
				values := url.Values{}
				values.Add(route.leaf, part)
				r.URL.RawQuery = url.Values(values).Encode() + "&" + r.URL.RawQuery
			} else { // Exact match
				if route.leaf == part {
					match = true
				}
			}
			// Did we get a match
			if match {
				// If we are not at the end of the path, then we go round again
				if i < len(parts)-1 {
					routes = route.Routes
					break
				} else { // We are at the end of the path - we have a route
					handler := route.handlers[r.Method]
					// Do we have a handler for this Verb
					if handler != nil {
						// Yes, serve and return
						handler.ServeHTTP(w, r)
						return
					}
					// We do not - Serve a 405 Method Not Allowed
					w.WriteHeader(http.StatusMethodNotAllowed)
					return
				}
			}
		}
	}

	// If we get here then we have not found a route
	w.WriteHeader(http.StatusNotFound)
}

// This type contains all the handlers for each path, each Route can also hold
// a list of child routes
type Route struct {
	leaf   string   // a part of a URL path, /foo/bar - a leaf would be foo and bar
	path   string   // full url path
	Routes []*Route // Routes that live under this route

	yam *Yam // Reference to Yam and global configuration

	// Verb handlers
	handlers map[string]http.Handler
}

// Adds a new route to the tree, and depending on configuration implements
// default handler implementation for OPTIONS and TRACE requests
func (r *Route) Route(path string) *Route {
	r = route(path, r, r.yam)

	if r.handlers == nil {
		r.handlers = make(map[string]http.Handler)
	}

	if r.yam.Config.Options {
		if r.handlers["OPTIONS"] == nil {
			r.handlers["OPTIONS"] = r.yam.Config.OptionsHandler(r)
		}
	}

	if r.yam.Config.Trace {
		if r.handlers["TRACE"] == nil {
			r.handlers["TRACE"] = r.yam.Config.TraceHandler(r)
		}
	}

	return r
}

// Adds a new handler to the route based on http Verb
func (r *Route) Add(method string, handler http.Handler) *Route {
	r.handlers[method] = handler

	return r
}

// Set a HEAD request handler for the route
func (r *Route) Head(h handler) *Route {
	r.Add("HEAD", http.HandlerFunc(h))

	return r
}

// Set a OPTIONS request for the route, overrides default implementation
func (r *Route) Options(h handler) *Route {
	r.Add("OPTIONS", http.HandlerFunc(h))

	return r
}

// Set a TRACE request for the route, overrides default implementation
func (r *Route) Trace(h handler) *Route {
	r.Add("TRACE", http.HandlerFunc(h))

	return r
}

// Set a GET request handler for the Route. A default HEAD request implementation
// will also be implemented since HEAD requests should perform the same as a GET
// request but simply not return the response body.
func (r *Route) Get(h handler) *Route {
	r.Add("GET", http.HandlerFunc(h))

	// Implement the HEAD handler by default for all GET requests - HEAD
	// should not return a body so we wrap it in a middleware
	head := func(n http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Serve the handler
			n.ServeHTTP(w, r)
			// Flush the body so we don't write to the client
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		})
	}

	// Apply the head middleware to the head handler
	r.Add("HEAD", head(http.HandlerFunc(h)))

	return r
}

// Set a POST request handler for the route.
func (r *Route) Post(h handler) *Route {
	r.Add("POST", http.HandlerFunc(h))

	return r
}

// Set a PUT request handler for the route.
func (r *Route) Put(h handler) *Route {
	r.Add("PUT", http.HandlerFunc(h))

	return r
}

// Set a DELETE request handler for the route.
func (r *Route) Delete(h handler) *Route {
	r.Add("DELETE", http.HandlerFunc(h))

	return r
}

// Set a PATCH request handler for the route.
func (r *Route) Patch(h handler) *Route {
	r.Add("PATCH", http.HandlerFunc(h))

	return r
}

// Default HTTP Handler function for OPTIONS requests. Adds the Allow header
// with the http verbs the route supports
func DefaultOptionsHandler(route *Route) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		methods := []string{}
		for key, _ := range route.handlers {
			methods = append(methods, key)
		}
		w.Header().Add("Allow", strings.Join(methods, ", "))
	})
}

// Default HTTP handler function for TRACE requests. Dumps the request as the Response.
func DefaultTraceHandler(route *Route) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		dump, _ := httputil.DumpRequest(r, false)
		w.Write(dump)
	})
}
<|endoftext|>


<commit_before>package promhandler

import (
	"bufio"
	"fmt"
	"net/http"

	"github.com/jacksontj/dataman/metrics"
)

func Handler(c metrics.Collectable) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		ch := make(chan metrics.MetricPoint)
		go func() {
			defer close(ch)
			c.Collect(ctx, ch)
		}()

		writer := bufio.NewWriter(w)

		for metricPoint := range ch {
			_, err := writer.WriteString(metricPoint.String())
			// TODO: do something with the error
			fmt.Println("err", err)
			writer.WriteByte('\n')
		}
		writer.Flush()
	}
}
<commit_msg>Handle error<commit_after>package promhandler

import (
	"bufio"
	"net/http"

	"github.com/jacksontj/dataman/metrics"
)

func Handler(c metrics.Collectable) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		ch := make(chan metrics.MetricPoint)
		go func() {
			defer close(ch)
			c.Collect(ctx, ch)
		}()

		writer := bufio.NewWriter(w)

		for metricPoint := range ch {
			if _, err := writer.WriteString(metricPoint.String()); err != nil {
				return
			}
			writer.WriteByte('\n')
		}
		writer.Flush()
	}
}
<|endoftext|>


<commit_before>package circonus

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strconv"
	"testing"

	"github.com/circonus-labs/circonus-gometrics"

	"github.com/go-kit/kit/metrics2/generic"
	"github.com/go-kit/kit/metrics2/teststat"
)

func TestCounter(t *testing.T) {
	// The only way to extract values from Circonus is to pose as a Circonus
	// server and receive real HTTP writes.
	const name = "abc"
	var val int64
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var res map[string]struct {
			Value int64 `json:"_value"` // reverse-engineered :\
		}
		json.NewDecoder(r.Body).Decode(&res)
		val = res[name].Value
	}))
	defer s.Close()

	// Set up a Circonus object, submitting to our HTTP server.
	m := circonusgometrics.NewCirconusMetrics()
	m.SubmissionUrl = s.URL
	counter := New(m).NewCounter(name)
	value := func() float64 { m.Flush(); return float64(val) }

	// Engage.
	if err := teststat.TestCounter(counter, value); err != nil {
		t.Fatal(err)
	}
}

func TestGauge(t *testing.T) {
	const name = "def"
	var val int64
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var res map[string]struct {
			Value int64 `json:"_value"`
		}
		json.NewDecoder(r.Body).Decode(&res)
		val = res[name].Value
	}))
	defer s.Close()

	m := circonusgometrics.NewCirconusMetrics()
	m.SubmissionUrl = s.URL
	gauge := New(m).NewGauge(name)
	value := func() float64 { m.Flush(); return float64(val) }

	if err := teststat.TestGauge(gauge, value); err != nil {
		t.Fatal(err)
	}
}

func TestHistogram(t *testing.T) {
	const name = "ghi"

	// Circonus just emits bucketed counts. We'll dump them into a generic
	// histogram (losing some precision) and take statistics from there. Note
	// this does assume that the generic histogram computes statistics properly,
	// but we have another test for that :)
	re := regexp.MustCompile(`^H\[([0-9\.e\+]+)\]=([0-9]+)$`) // H[1.2e+03]=456

	var p50, p90, p95, p99 float64
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var res map[string]struct {
			Values []string `json:"_value"` // reverse-engineered :\
		}
		json.NewDecoder(r.Body).Decode(&res)

		h := generic.NewHistogram(len(res[name].Values)) // match tbe bucket counts
		for _, v := range res[name].Values {
			match := re.FindStringSubmatch(v)
			f, _ := strconv.ParseFloat(match[1], 64)
			n, _ := strconv.ParseInt(match[2], 10, 64)
			for i := int64(0); i < n; i++ {
				h.Observe(f)
			}
		}

		p50 = h.Quantile(0.50)
		p90 = h.Quantile(0.90)
		p95 = h.Quantile(0.95)
		p99 = h.Quantile(0.99)
	}))
	defer s.Close()

	m := circonusgometrics.NewCirconusMetrics()
	m.SubmissionUrl = s.URL
	histogram := New(m).NewHistogram(name)
	quantiles := func() (float64, float64, float64, float64) { m.Flush(); return p50, p90, p95, p99 }

	// Circonus metrics, because they do their own bucketing, are less precise
	// than other systems. So, we bump the tolerance to 5 percent.
	if err := teststat.TestHistogram(histogram, quantiles, 0.05); err != nil {
		t.Fatal(err)
	}
}
<commit_msg>A new Circonus API appears<commit_after>package circonus

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strconv"
	"testing"

	"github.com/circonus-labs/circonus-gometrics"
	"github.com/circonus-labs/circonus-gometrics/checkmgr"

	"github.com/go-kit/kit/metrics2/generic"
	"github.com/go-kit/kit/metrics2/teststat"
)

func TestCounter(t *testing.T) {
	// The only way to extract values from Circonus is to pose as a Circonus
	// server and receive real HTTP writes.
	const name = "abc"
	var val int64
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var res map[string]struct {
			Value int64 `json:"_value"` // reverse-engineered :\
		}
		json.NewDecoder(r.Body).Decode(&res)
		val = res[name].Value
	}))
	defer s.Close()

	// Set up a Circonus object, submitting to our HTTP server.
	m := newCirconusMetrics(s.URL)
	counter := New(m).NewCounter(name)
	value := func() float64 { m.Flush(); return float64(val) }

	// Engage.
	if err := teststat.TestCounter(counter, value); err != nil {
		t.Fatal(err)
	}
}

func TestGauge(t *testing.T) {
	const name = "def"
	var val int64
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var res map[string]struct {
			Value int64 `json:"_value"`
		}
		json.NewDecoder(r.Body).Decode(&res)
		val = res[name].Value
	}))
	defer s.Close()

	m := newCirconusMetrics(s.URL)
	gauge := New(m).NewGauge(name)
	value := func() float64 { m.Flush(); return float64(val) }

	if err := teststat.TestGauge(gauge, value); err != nil {
		t.Fatal(err)
	}
}

func TestHistogram(t *testing.T) {
	const name = "ghi"

	// Circonus just emits bucketed counts. We'll dump them into a generic
	// histogram (losing some precision) and take statistics from there. Note
	// this does assume that the generic histogram computes statistics properly,
	// but we have another test for that :)
	re := regexp.MustCompile(`^H\[([0-9\.e\+]+)\]=([0-9]+)$`) // H[1.2e+03]=456

	var p50, p90, p95, p99 float64
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var res map[string]struct {
			Values []string `json:"_value"` // reverse-engineered :\
		}
		json.NewDecoder(r.Body).Decode(&res)

		h := generic.NewHistogram(len(res[name].Values)) // match tbe bucket counts
		for _, v := range res[name].Values {
			match := re.FindStringSubmatch(v)
			f, _ := strconv.ParseFloat(match[1], 64)
			n, _ := strconv.ParseInt(match[2], 10, 64)
			for i := int64(0); i < n; i++ {
				h.Observe(f)
			}
		}

		p50 = h.Quantile(0.50)
		p90 = h.Quantile(0.90)
		p95 = h.Quantile(0.95)
		p99 = h.Quantile(0.99)
	}))
	defer s.Close()

	m := newCirconusMetrics(s.URL)
	histogram := New(m).NewHistogram(name)
	quantiles := func() (float64, float64, float64, float64) { m.Flush(); return p50, p90, p95, p99 }

	// Circonus metrics, because they do their own bucketing, are less precise
	// than other systems. So, we bump the tolerance to 5 percent.
	if err := teststat.TestHistogram(histogram, quantiles, 0.05); err != nil {
		t.Fatal(err)
	}
}

func newCirconusMetrics(url string) *circonusgometrics.CirconusMetrics {
	m, err := circonusgometrics.NewCirconusMetrics(&circonusgometrics.Config{
		CheckManager: checkmgr.Config{
			Check: checkmgr.CheckConfig{
				SubmissionUrl: url,
			},
		},
	})
	if err != nil {
		panic(fmt.Sprintln("URL=", url, "--", err))
	}
	return m
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"fmt"
)

func CheckIsAPIOwner(handler func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		tykAuthKey := r.Header.Get("X-Tyk-Authorisation")
		if tykAuthKey != config.Secret {
			// Error
			log.Warning("Attempted administrative access with invalid or missing key!")

			responseMessage := createError("Method not supported")
			w.WriteHeader(403)
			fmt.Fprintf(w, string(responseMessage))

			return
		}

		handler(w, r)

	}
}
<commit_msg>documenting<commit_after>package main

import (
	"net/http"
	"fmt"
)

// CheckIsAPIOwner will ensure that the accessor of the tyk API has the correct security credentials - this is a
// shared secret between the client and the owner and is set in the tyk.conf file. This should never be made public!
func CheckIsAPIOwner(handler func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		tykAuthKey := r.Header.Get("X-Tyk-Authorisation")
		if tykAuthKey != config.Secret {
			// Error
			log.Warning("Attempted administrative access with invalid or missing key!")

			responseMessage := createError("Method not supported")
			w.WriteHeader(403)
			fmt.Fprintf(w, string(responseMessage))

			return
		}

		handler(w, r)

	}
}
<|endoftext|>