<commit_before>// package forwarder implements http handler that forwards requests to remote server
// and serves back the response
// websocket proxying support based on https://github.com/yhat/wsutil
package forward

import (
	"crypto/tls"
	"io"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/websocket"
	"github.com/vulcand/oxy/utils"
)

// ReqRewriter can alter request headers and body
type ReqRewriter interface {
	Rewrite(r *http.Request)
}

type optSetter func(f *Forwarder) error

// PassHostHeader specifies if a client's Host header field should
// be delegated
func PassHostHeader(b bool) optSetter {
	return func(f *Forwarder) error {
		f.passHost = b
		return nil
	}
}

// StreamResponse forces streaming body (flushes response directly to client)
func StreamResponse(b bool) optSetter {
	return func(f *Forwarder) error {
		f.httpForwarder.streamResponse = b
		return nil
	}
}

// RoundTripper sets a new http.RoundTripper
// Forwarder will use http.DefaultTransport as a default round tripper
func RoundTripper(r http.RoundTripper) optSetter {
	return func(f *Forwarder) error {
		f.roundTripper = r
		return nil
	}
}

// Rewriter defines a request rewriter for the HTTP forwarder
func Rewriter(r ReqRewriter) optSetter {
	return func(f *Forwarder) error {
		f.httpForwarder.rewriter = r
		return nil
	}
}

// WebsocketRewriter defines a request rewriter for the websocket forwarder
func WebsocketRewriter(r ReqRewriter) optSetter {
	return func(f *Forwarder) error {
		f.websocketForwarder.rewriter = r
		return nil
	}
}

// ErrorHandler is a functional argument that sets error handler of the server
func ErrorHandler(h utils.ErrorHandler) optSetter {
	return func(f *Forwarder) error {
		f.errHandler = h
		return nil
	}
}

// Logger specifies the logger to use.
// Forwarder will default to oxyutils.NullLogger if no logger has been specified
func Logger(l utils.Logger) optSetter {
	return func(f *Forwarder) error {
		f.log = l
		return nil
	}
}

// Forwarder wraps two traffic forwarding implementations: HTTP and websockets.
// It decides based on the specified request which implementation to use
type Forwarder struct {
	*httpForwarder
	*websocketForwarder
	*handlerContext
}

// handlerContext defines a handler context for error reporting and logging
type handlerContext struct {
	errHandler utils.ErrorHandler
	log        utils.Logger
}

// httpForwarder is a handler that can reverse proxy
// HTTP traffic
type httpForwarder struct {
	roundTripper   http.RoundTripper
	rewriter       ReqRewriter
	passHost       bool
	streamResponse bool
}

// websocketForwarder is a handler that can reverse proxy
// websocket traffic
type websocketForwarder struct {
	rewriter        ReqRewriter
	TLSClientConfig *tls.Config
}

// New creates an instance of Forwarder based on the provided list of configuration options
func New(setters ...optSetter) (*Forwarder, error) {
	f := &Forwarder{
		httpForwarder:      &httpForwarder{},
		websocketForwarder: &websocketForwarder{},
		handlerContext:     &handlerContext{},
	}
	for _, s := range setters {
		if err := s(f); err != nil {
			return nil, err
		}
	}
	if f.httpForwarder.roundTripper == nil {
		f.httpForwarder.roundTripper = http.DefaultTransport
	}
	f.websocketForwarder.TLSClientConfig = f.httpForwarder.roundTripper.(*http.Transport).TLSClientConfig
	if f.httpForwarder.rewriter == nil {
		h, err := os.Hostname()
		if err != nil {
			h = "localhost"
		}
		f.httpForwarder.rewriter = &HeaderRewriter{TrustForwardHeader: true, Hostname: h}
	}
	if f.log == nil {
		f.log = utils.NullLogger
	}
	if f.errHandler == nil {
		f.errHandler = utils.DefaultHandler
	}
	return f, nil
}

// ServeHTTP decides which forwarder to use based on the specified
// request and delegates to the proper implementation
func (f *Forwarder) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if isWebsocketRequest(req) {
		f.websocketForwarder.serveHTTP(w, req, f.handlerContext)
	} else {
		f.httpForwarder.serveHTTP(w, req, f.handlerContext)
	}
}

// serveHTTP forwards HTTP traffic using the configured transport
func (f *httpForwarder) serveHTTP(w http.ResponseWriter, req *http.Request, ctx *handlerContext) {
	start := time.Now().UTC()

	response, err := f.roundTripper.RoundTrip(f.copyRequest(req, req.URL))

	if err != nil {
		ctx.log.Errorf("Error forwarding to %v, err: %v", req.URL, err)
		ctx.errHandler.ServeHTTP(w, req, err)
		return
	}

	utils.CopyHeaders(w.Header(), response.Header)
	// Remove hop-by-hop headers.
	utils.RemoveHeaders(w.Header(), HopHeaders...)

	announcedTrailerKeyCount := len(response.Trailer)
	if announcedTrailerKeyCount > 0 {
		trailerKeys := make([]string, 0, announcedTrailerKeyCount)
		for k := range response.Trailer {
			trailerKeys = append(trailerKeys, k)
		}
		w.Header().Add("Trailer", strings.Join(trailerKeys, ", "))
	}

	w.WriteHeader(response.StatusCode)

	stream := f.streamResponse
	if !stream {
		contentType, err := utils.GetHeaderMediaType(response.Header, ContentType)
		if err == nil {
			stream = contentType == "text/event-stream"
		}
	}

	flush := stream || req.ProtoMajor == 2
	written, err := io.Copy(newResponseFlusher(w, flush), response.Body)
	if err != nil {
		ctx.log.Errorf("Error copying upstream response body: %v", err)
		ctx.errHandler.ServeHTTP(w, req, err)
		return
	}

	defer response.Body.Close()

	forceSetTrailers := len(response.Trailer) != announcedTrailerKeyCount
	shallowCopyTrailers(w.Header(), response.Trailer, forceSetTrailers)

	if written != 0 {
		w.Header().Set(ContentLength, strconv.FormatInt(written, 10))
	}

	if req.TLS != nil {
		ctx.log.Infof("Round trip: %v, code: %v, duration: %v tls:version: %x, tls:resume:%t, tls:csuite:%x, tls:server:%v",
			req.URL, response.StatusCode, time.Now().UTC().Sub(start),
			req.TLS.Version,
			req.TLS.DidResume,
			req.TLS.CipherSuite,
			req.TLS.ServerName)
	} else {
		ctx.log.Infof("Round trip: %v, code: %v, duration: %v",
			req.URL, response.StatusCode, time.Now().UTC().Sub(start))
	}

}

// copyRequest makes a copy of the specified request to be sent using the configured
// transport
func (f *httpForwarder) copyRequest(req *http.Request, u *url.URL) *http.Request {
	outReq := new(http.Request)
	*outReq = *req // includes shallow copies of maps, but we handle this below

	outReq.URL = utils.CopyURL(req.URL)
	outReq.URL.Scheme = u.Scheme
	outReq.URL.Host = u.Host
	outReq.URL.Opaque = req.RequestURI
	// raw query is already included in RequestURI, so ignore it to avoid dupes
	outReq.URL.RawQuery = ""
	// Do not pass client Host header unless optsetter PassHostHeader is set.
	if !f.passHost {
		outReq.Host = u.Host
	}
	outReq.Proto = "HTTP/1.1"
	outReq.ProtoMajor = 1
	outReq.ProtoMinor = 1

	// Overwrite close flag so we can keep persistent connection for the backend servers
	outReq.Close = false

	outReq.Header = make(http.Header)
	utils.CopyHeaders(outReq.Header, req.Header)

	if f.rewriter != nil {
		f.rewriter.Rewrite(outReq)
	}

	if req.ContentLength == 0 {
		// https://github.com/golang/go/issues/16036: nil Body for http.Transport retries
		outReq.Body = nil
	}

	return outReq
}

// serveHTTP forwards websocket traffic
func (f *websocketForwarder) serveHTTP(w http.ResponseWriter, req *http.Request, ctx *handlerContext) {
	outReq := f.copyRequest(req, req.URL)

	dialer := websocket.DefaultDialer
	if outReq.URL.Scheme == "wss" && f.TLSClientConfig != nil {
		dialer.TLSClientConfig = f.TLSClientConfig.Clone()
		dialer.TLSClientConfig.NextProtos = []string{"http/1.1"}
	}
	targetConn, resp, err := dialer.Dial(outReq.URL.String(), outReq.Header)
	if err != nil {
		if resp == nil {
			ctx.errHandler.ServeHTTP(w, req, err)
		} else {
			ctx.log.Errorf("Error dialing %q: %v with resp: %d %s", outReq.Host, err, resp.StatusCode, resp.Status)
			hijacker, ok := w.(http.Hijacker)
			if !ok {
				ctx.log.Errorf("%s can not be hijack", reflect.TypeOf(w))
				ctx.errHandler.ServeHTTP(w, req, err)
				return
			}

			conn, _, err := hijacker.Hijack()
			if err != nil {
				ctx.log.Errorf("Failed to hijack responseWriter")
				ctx.errHandler.ServeHTTP(w, req, err)
				return
			}
			defer conn.Close()

			err = resp.Write(conn)
			if err != nil {
				ctx.log.Errorf("Failed to forward response")
				ctx.errHandler.ServeHTTP(w, req, err)
				return
			}
		}
		return
	}

	//Only the targetConn choose to CheckOrigin or not
	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool {
		return true
	}}

	utils.RemoveHeaders(resp.Header, WebsocketUpgradeHeaders...)
	underlyingConn, err := upgrader.Upgrade(w, req, resp.Header)
	if err != nil {
		ctx.log.Errorf("Error while upgrading connection : %v", err)
		return
	}
	defer underlyingConn.Close()
	defer targetConn.Close()

	errc := make(chan error, 2)
	replicate := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}

	go replicate(targetConn.UnderlyingConn(), underlyingConn.UnderlyingConn())

	// Try to read the first message
	t, msg, err := targetConn.ReadMessage()
	if err != nil {
		ctx.log.Errorf("Couldn't read first message : %v", err)
	} else {
		underlyingConn.WriteMessage(t, msg)
	}

	go replicate(underlyingConn.UnderlyingConn(), targetConn.UnderlyingConn())
	<-errc

}

// copyRequest makes a copy of the specified request.
func (f *websocketForwarder) copyRequest(req *http.Request, u *url.URL) (outReq *http.Request) {
	outReq = new(http.Request)
	*outReq = *req // includes shallow copies of maps, but we handle this below

	outReq.URL = utils.CopyURL(req.URL)
	outReq.URL.Scheme = u.Scheme

	//sometimes backends might be registered as HTTP/HTTPS servers so translate URLs to websocket URLs.
	switch u.Scheme {
	case "https":
		outReq.URL.Scheme = "wss"
	case "http":
		outReq.URL.Scheme = "ws"
	}

	if requestURI, err := url.ParseRequestURI(outReq.RequestURI); err == nil {
		if requestURI.RawPath != "" {
			outReq.URL.Path = requestURI.RawPath
		} else {
			outReq.URL.Path = requestURI.Path
		}
		outReq.URL.RawQuery = requestURI.RawQuery
	}

	outReq.URL.Host = u.Host

	outReq.Header = make(http.Header)
	//gorilla websocket use this header to set the request.Host tested in checkSameOrigin
	outReq.Header.Set("Host", outReq.Host)
	utils.CopyHeaders(outReq.Header, req.Header)
	utils.RemoveHeaders(outReq.Header, WebsocketDialHeaders...)

	if f.rewriter != nil {
		f.rewriter.Rewrite(outReq)
	}
	return outReq
}

// isWebsocketRequest determines if the specified HTTP request is a
// websocket handshake request
func isWebsocketRequest(req *http.Request) bool {
	containsHeader := func(name, value string) bool {
		items := strings.Split(req.Header.Get(name), ",")
		for _, item := range items {
			if value == strings.ToLower(strings.TrimSpace(item)) {
				return true
			}
		}
		return false
	}
	return containsHeader(Connection, "upgrade") && containsHeader(Upgrade, "websocket")
}

func shallowCopyTrailers(dstHeader, srcTrailer http.Header, forceSetTrailers bool) {
	for k, vv := range srcTrailer {
		if forceSetTrailers {
			k = http.TrailerPrefix + k
		}
		dstHeader[k] = vv
	}
}
<commit_msg>Use gorilla readMessage and writeMessage instead of just an io.Copy<commit_after>// package forwarder implements http handler that forwards requests to remote server
// and serves back the response
// websocket proxying support based on https://github.com/yhat/wsutil
package forward

import (
	"crypto/tls"
	"io"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/websocket"
	"github.com/vulcand/oxy/utils"
)

// ReqRewriter can alter request headers and body
type ReqRewriter interface {
	Rewrite(r *http.Request)
}

type optSetter func(f *Forwarder) error

// PassHostHeader specifies if a client's Host header field should
// be delegated
func PassHostHeader(b bool) optSetter {
	return func(f *Forwarder) error {
		f.passHost = b
		return nil
	}
}

// StreamResponse forces streaming body (flushes response directly to client)
func StreamResponse(b bool) optSetter {
	return func(f *Forwarder) error {
		f.httpForwarder.streamResponse = b
		return nil
	}
}

// RoundTripper sets a new http.RoundTripper
// Forwarder will use http.DefaultTransport as a default round tripper
func RoundTripper(r http.RoundTripper) optSetter {
	return func(f *Forwarder) error {
		f.roundTripper = r
		return nil
	}
}

// Rewriter defines a request rewriter for the HTTP forwarder
func Rewriter(r ReqRewriter) optSetter {
	return func(f *Forwarder) error {
		f.httpForwarder.rewriter = r
		return nil
	}
}

// WebsocketRewriter defines a request rewriter for the websocket forwarder
func WebsocketRewriter(r ReqRewriter) optSetter {
	return func(f *Forwarder) error {
		f.websocketForwarder.rewriter = r
		return nil
	}
}

// ErrorHandler is a functional argument that sets error handler of the server
func ErrorHandler(h utils.ErrorHandler) optSetter {
	return func(f *Forwarder) error {
		f.errHandler = h
		return nil
	}
}

// Logger specifies the logger to use.
// Forwarder will default to oxyutils.NullLogger if no logger has been specified
func Logger(l utils.Logger) optSetter {
	return func(f *Forwarder) error {
		f.log = l
		return nil
	}
}

// Forwarder wraps two traffic forwarding implementations: HTTP and websockets.
// It decides based on the specified request which implementation to use
type Forwarder struct {
	*httpForwarder
	*websocketForwarder
	*handlerContext
}

// handlerContext defines a handler context for error reporting and logging
type handlerContext struct {
	errHandler utils.ErrorHandler
	log        utils.Logger
}

// httpForwarder is a handler that can reverse proxy
// HTTP traffic
type httpForwarder struct {
	roundTripper   http.RoundTripper
	rewriter       ReqRewriter
	passHost       bool
	streamResponse bool
}

// websocketForwarder is a handler that can reverse proxy
// websocket traffic
type websocketForwarder struct {
	rewriter        ReqRewriter
	TLSClientConfig *tls.Config
}

// New creates an instance of Forwarder based on the provided list of configuration options
func New(setters ...optSetter) (*Forwarder, error) {
	f := &Forwarder{
		httpForwarder:      &httpForwarder{},
		websocketForwarder: &websocketForwarder{},
		handlerContext:     &handlerContext{},
	}
	for _, s := range setters {
		if err := s(f); err != nil {
			return nil, err
		}
	}
	if f.httpForwarder.roundTripper == nil {
		f.httpForwarder.roundTripper = http.DefaultTransport
	}
	f.websocketForwarder.TLSClientConfig = f.httpForwarder.roundTripper.(*http.Transport).TLSClientConfig
	if f.httpForwarder.rewriter == nil {
		h, err := os.Hostname()
		if err != nil {
			h = "localhost"
		}
		f.httpForwarder.rewriter = &HeaderRewriter{TrustForwardHeader: true, Hostname: h}
	}
	if f.log == nil {
		f.log = utils.NullLogger
	}
	if f.errHandler == nil {
		f.errHandler = utils.DefaultHandler
	}
	return f, nil
}

// ServeHTTP decides which forwarder to use based on the specified
// request and delegates to the proper implementation
func (f *Forwarder) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if isWebsocketRequest(req) {
		f.websocketForwarder.serveHTTP(w, req, f.handlerContext)
	} else {
		f.httpForwarder.serveHTTP(w, req, f.handlerContext)
	}
}

// serveHTTP forwards HTTP traffic using the configured transport
func (f *httpForwarder) serveHTTP(w http.ResponseWriter, req *http.Request, ctx *handlerContext) {
	start := time.Now().UTC()

	response, err := f.roundTripper.RoundTrip(f.copyRequest(req, req.URL))

	if err != nil {
		ctx.log.Errorf("Error forwarding to %v, err: %v", req.URL, err)
		ctx.errHandler.ServeHTTP(w, req, err)
		return
	}

	utils.CopyHeaders(w.Header(), response.Header)
	// Remove hop-by-hop headers.
	utils.RemoveHeaders(w.Header(), HopHeaders...)

	announcedTrailerKeyCount := len(response.Trailer)
	if announcedTrailerKeyCount > 0 {
		trailerKeys := make([]string, 0, announcedTrailerKeyCount)
		for k := range response.Trailer {
			trailerKeys = append(trailerKeys, k)
		}
		w.Header().Add("Trailer", strings.Join(trailerKeys, ", "))
	}

	w.WriteHeader(response.StatusCode)

	stream := f.streamResponse
	if !stream {
		contentType, err := utils.GetHeaderMediaType(response.Header, ContentType)
		if err == nil {
			stream = contentType == "text/event-stream"
		}
	}

	flush := stream || req.ProtoMajor == 2
	written, err := io.Copy(newResponseFlusher(w, flush), response.Body)
	if err != nil {
		ctx.log.Errorf("Error copying upstream response body: %v", err)
		ctx.errHandler.ServeHTTP(w, req, err)
		return
	}

	defer response.Body.Close()

	forceSetTrailers := len(response.Trailer) != announcedTrailerKeyCount
	shallowCopyTrailers(w.Header(), response.Trailer, forceSetTrailers)

	if written != 0 {
		w.Header().Set(ContentLength, strconv.FormatInt(written, 10))
	}

	if req.TLS != nil {
		ctx.log.Infof("Round trip: %v, code: %v, duration: %v tls:version: %x, tls:resume:%t, tls:csuite:%x, tls:server:%v",
			req.URL, response.StatusCode, time.Now().UTC().Sub(start),
			req.TLS.Version,
			req.TLS.DidResume,
			req.TLS.CipherSuite,
			req.TLS.ServerName)
	} else {
		ctx.log.Infof("Round trip: %v, code: %v, duration: %v",
			req.URL, response.StatusCode, time.Now().UTC().Sub(start))
	}

}

// copyRequest makes a copy of the specified request to be sent using the configured
// transport
func (f *httpForwarder) copyRequest(req *http.Request, u *url.URL) *http.Request {
	outReq := new(http.Request)
	*outReq = *req // includes shallow copies of maps, but we handle this below

	outReq.URL = utils.CopyURL(req.URL)
	outReq.URL.Scheme = u.Scheme
	outReq.URL.Host = u.Host
	outReq.URL.Opaque = req.RequestURI
	// raw query is already included in RequestURI, so ignore it to avoid dupes
	outReq.URL.RawQuery = ""
	// Do not pass client Host header unless optsetter PassHostHeader is set.
	if !f.passHost {
		outReq.Host = u.Host
	}
	outReq.Proto = "HTTP/1.1"
	outReq.ProtoMajor = 1
	outReq.ProtoMinor = 1

	// Overwrite close flag so we can keep persistent connection for the backend servers
	outReq.Close = false

	outReq.Header = make(http.Header)
	utils.CopyHeaders(outReq.Header, req.Header)

	if f.rewriter != nil {
		f.rewriter.Rewrite(outReq)
	}

	if req.ContentLength == 0 {
		// https://github.com/golang/go/issues/16036: nil Body for http.Transport retries
		outReq.Body = nil
	}

	return outReq
}

// serveHTTP forwards websocket traffic
func (f *websocketForwarder) serveHTTP(w http.ResponseWriter, req *http.Request, ctx *handlerContext) {
	outReq := f.copyRequest(req, req.URL)

	dialer := websocket.DefaultDialer
	if outReq.URL.Scheme == "wss" && f.TLSClientConfig != nil {
		dialer.TLSClientConfig = f.TLSClientConfig.Clone()
		dialer.TLSClientConfig.NextProtos = []string{"http/1.1"}
	}
	targetConn, resp, err := dialer.Dial(outReq.URL.String(), outReq.Header)
	if err != nil {
		if resp == nil {
			ctx.errHandler.ServeHTTP(w, req, err)
		} else {
			ctx.log.Errorf("Error dialing %q: %v with resp: %d %s", outReq.Host, err, resp.StatusCode, resp.Status)
			hijacker, ok := w.(http.Hijacker)
			if !ok {
				ctx.log.Errorf("%s can not be hijack", reflect.TypeOf(w))
				ctx.errHandler.ServeHTTP(w, req, err)
				return
			}

			conn, _, err := hijacker.Hijack()
			if err != nil {
				ctx.log.Errorf("Failed to hijack responseWriter")
				ctx.errHandler.ServeHTTP(w, req, err)
				return
			}
			defer conn.Close()

			err = resp.Write(conn)
			if err != nil {
				ctx.log.Errorf("Failed to forward response")
				ctx.errHandler.ServeHTTP(w, req, err)
				return
			}
		}
		return
	}

	//Only the targetConn choose to CheckOrigin or not
	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool {
		return true
	}}

	utils.RemoveHeaders(resp.Header, WebsocketUpgradeHeaders...)
	underlyingConn, err := upgrader.Upgrade(w, req, resp.Header)
	if err != nil {
		ctx.log.Errorf("Error while upgrading connection : %v", err)
		return
	}
	defer underlyingConn.Close()
	defer targetConn.Close()

	errc := make(chan error, 2)

	replicateWebsocketConn := func(dst, src *websocket.Conn, dstName, srcName string) {
		var err error
		for {
			msgType, msg, err := src.ReadMessage()
			if err != nil {
				ctx.log.Errorf("vulcand/oxy/forward/websocket: Error when copying from %s to %s using ReadMessage: %v", srcName, dstName, err)
				break
			}
			err = dst.WriteMessage(msgType, msg)
			if err != nil {
				ctx.log.Errorf("vulcand/oxy/forward/websocket: Error when copying from %s to %s using WriteMessage: %v", srcName, dstName, err)
				break
			} else {
				ctx.log.Infof("vulcand/oxy/forward/websocket: Copying from %s to %s completed without error.", srcName, dstName)
			}
		}
		errc <- err
	}

	go replicateWebsocketConn(underlyingConn, targetConn, "client", "backend")
	go replicateWebsocketConn(targetConn, underlyingConn, "backend", "client")

	<-errc

}

// copyRequest makes a copy of the specified request.
func (f *websocketForwarder) copyRequest(req *http.Request, u *url.URL) (outReq *http.Request) {
	outReq = new(http.Request)
	*outReq = *req // includes shallow copies of maps, but we handle this below

	outReq.URL = utils.CopyURL(req.URL)
	outReq.URL.Scheme = u.Scheme

	//sometimes backends might be registered as HTTP/HTTPS servers so translate URLs to websocket URLs.
	switch u.Scheme {
	case "https":
		outReq.URL.Scheme = "wss"
	case "http":
		outReq.URL.Scheme = "ws"
	}

	if requestURI, err := url.ParseRequestURI(outReq.RequestURI); err == nil {
		if requestURI.RawPath != "" {
			outReq.URL.Path = requestURI.RawPath
		} else {
			outReq.URL.Path = requestURI.Path
		}
		outReq.URL.RawQuery = requestURI.RawQuery
	}

	outReq.URL.Host = u.Host

	outReq.Header = make(http.Header)
	//gorilla websocket use this header to set the request.Host tested in checkSameOrigin
	outReq.Header.Set("Host", outReq.Host)
	utils.CopyHeaders(outReq.Header, req.Header)
	utils.RemoveHeaders(outReq.Header, WebsocketDialHeaders...)

	if f.rewriter != nil {
		f.rewriter.Rewrite(outReq)
	}
	return outReq
}

// isWebsocketRequest determines if the specified HTTP request is a
// websocket handshake request
func isWebsocketRequest(req *http.Request) bool {
	containsHeader := func(name, value string) bool {
		items := strings.Split(req.Header.Get(name), ",")
		for _, item := range items {
			if value == strings.ToLower(strings.TrimSpace(item)) {
				return true
			}
		}
		return false
	}
	return containsHeader(Connection, "upgrade") && containsHeader(Upgrade, "websocket")
}

func shallowCopyTrailers(dstHeader, srcTrailer http.Header, forceSetTrailers bool) {
	for k, vv := range srcTrailer {
		if forceSetTrailers {
			k = http.TrailerPrefix + k
		}
		dstHeader[k] = vv
	}
}
<|endoftext|>


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

package core

import (
	"log"
	"sort"
	"sync"

	"github.com/go-ego/riot/types"
	"github.com/go-ego/riot/utils"
)

// Ranker ranker
type Ranker struct {
	lock struct {
		sync.RWMutex
		fields map[uint64]interface{}
		docs   map[uint64]bool
		// new
		content map[uint64]string
		attri   map[uint64]interface{}
	}

	idOnly      bool
	initialized bool
}

// Init init ranker
func (ranker *Ranker) Init(onlyID ...bool) {
	if ranker.initialized == true {
		log.Fatal("The Ranker can not be initialized twice.")
	}
	ranker.initialized = true

	if len(onlyID) > 0 {
		ranker.idOnly = onlyID[0]
	}

	ranker.lock.fields = make(map[uint64]interface{})
	ranker.lock.docs = make(map[uint64]bool)

	if !ranker.idOnly {
		// new
		ranker.lock.content = make(map[uint64]string)
		ranker.lock.attri = make(map[uint64]interface{})
	}
}

// AddDoc add doc
// 给某个文档添加评分字段
func (ranker *Ranker) AddDoc(
	// docId uint64, fields interface{}, content string, attri interface{}) {
	docId uint64, fields interface{}, content ...interface{}) {
	if ranker.initialized == false {
		log.Fatal("The Ranker has not been initialized.")
	}

	ranker.lock.Lock()
	ranker.lock.fields[docId] = fields
	ranker.lock.docs[docId] = true

	if !ranker.idOnly {
		// new
		if len(content) > 0 {
			ranker.lock.content[docId] = content[0].(string)
		}

		if len(content) > 1 {
			ranker.lock.attri[docId] = content[1]
			// ranker.lock.attri[docId] = attri
		}
	}

	ranker.lock.Unlock()
}

// RemoveDoc 删除某个文档的评分字段
func (ranker *Ranker) RemoveDoc(docId uint64) {
	if ranker.initialized == false {
		log.Fatal("The Ranker has not been initialized.")
	}

	ranker.lock.Lock()
	delete(ranker.lock.fields, docId)
	delete(ranker.lock.docs, docId)

	if !ranker.idOnly {
		// new
		delete(ranker.lock.content, docId)
		delete(ranker.lock.attri, docId)
	}

	ranker.lock.Unlock()
}

func maxOutput(options types.RankOpts, docsLen int) (int, int) {
	var start, end int
	if options.MaxOutputs != 0 {
		start = utils.MinInt(options.OutputOffset, docsLen)
		end = utils.MinInt(options.OutputOffset+options.MaxOutputs, docsLen)
		return start, end
	}

	start = utils.MinInt(options.OutputOffset, docsLen)
	end = docsLen
	return start, end
}

func (ranker *Ranker) rankOutIDs(docs []types.IndexedDoc, options types.RankOpts,
	countDocsOnly bool) (outputDocs types.ScoredIDs, numDocs int) {
	for _, d := range docs {
		ranker.lock.RLock()
		// 判断 doc 是否存在
		if _, ok := ranker.lock.docs[d.DocId]; ok {

			fs := ranker.lock.fields[d.DocId]
			ranker.lock.RUnlock()

			// 计算评分并剔除没有分值的文档
			scores := options.ScoringCriteria.Score(d, fs)
			if len(scores) > 0 {
				if !countDocsOnly {
					outputDocs = append(outputDocs,
						types.ScoredID{
							DocId:            d.DocId,
							Scores:           scores,
							TokenSnippetLocs: d.TokenSnippetLocs,
							TokenLocs:        d.TokenLocs,
						})
				}
				numDocs++
			}
		} else {
			ranker.lock.RUnlock()
		}
	}

	return
}

// RankDocID rank docs by types.ScoredIDs
func (ranker *Ranker) RankDocID(docs []types.IndexedDoc,
	options types.RankOpts, countDocsOnly bool) (types.ScoredIDs, int) {

	outputDocs, numDocs := ranker.rankOutIDs(docs, options, countDocsOnly)

	// 排序
	if !countDocsOnly {
		if options.ReverseOrder {
			sort.Sort(sort.Reverse(outputDocs))
		} else {
			sort.Sort(outputDocs)
		}
		// 当用户要求只返回部分结果时返回部分结果
		docsLen := len(outputDocs)
		start, end := maxOutput(options, docsLen)

		return outputDocs[start:end], numDocs
	}

	return outputDocs, numDocs
}

func (ranker *Ranker) rankOutDocs(docs []types.IndexedDoc, options types.RankOpts,
	countDocsOnly bool) (outputDocs types.ScoredDocs, numDocs int) {
	for _, d := range docs {
		ranker.lock.RLock()
		// 判断 doc 是否存在
		if _, ok := ranker.lock.docs[d.DocId]; ok {

			fs := ranker.lock.fields[d.DocId]
			content := ranker.lock.content[d.DocId]
			attri := ranker.lock.attri[d.DocId]
			ranker.lock.RUnlock()

			// 计算评分并剔除没有分值的文档
			scores := options.ScoringCriteria.Score(d, fs)
			if len(scores) > 0 {
				if !countDocsOnly {
					scoredID := types.ScoredID{
						DocId:            d.DocId,
						Scores:           scores,
						TokenSnippetLocs: d.TokenSnippetLocs,
						TokenLocs:        d.TokenLocs,
					}

					outputDocs = append(outputDocs,
						types.ScoredDoc{
							ScoredID: scoredID,
							// new
							Fields:  fs,
							Content: content,
							Attri:   attri,
						})
				}
				numDocs++
			}
		} else {
			ranker.lock.RUnlock()
		}
	}

	return
}

// RankDocs rank docs by types.ScoredDocs
func (ranker *Ranker) RankDocs(docs []types.IndexedDoc,
	options types.RankOpts, countDocsOnly bool) (types.ScoredDocs, int) {

	outputDocs, numDocs := ranker.rankOutDocs(docs, options, countDocsOnly)

	// 排序
	if !countDocsOnly {
		if options.ReverseOrder {
			sort.Sort(sort.Reverse(outputDocs))
		} else {
			sort.Sort(outputDocs)
		}
		// 当用户要求只返回部分结果时返回部分结果
		docsLen := len(outputDocs)
		start, end := maxOutput(options, docsLen)

		return outputDocs[start:end], numDocs
	}

	return outputDocs, numDocs
}

// Rank rank docs
// 给文档评分并排序
func (ranker *Ranker) Rank(docs []types.IndexedDoc,
	options types.RankOpts, countDocsOnly bool) (interface{}, int) {

	if ranker.initialized == false {
		log.Fatal("The Ranker has not been initialized.")
	}

	// 对每个文档评分
	if ranker.idOnly {
		outputDocs, numDocs := ranker.RankDocID(docs, options, countDocsOnly)
		return outputDocs, numDocs
	}

	outputDocs, numDocs := ranker.RankDocs(docs, options, countDocsOnly)
	return outputDocs, numDocs
}
<commit_msg>update ranker code<commit_after>// Copyright 2013 Hui Chen
// Copyright 2016 ego authors
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package core

import (
	"log"
	"sort"
	"sync"

	"github.com/go-ego/riot/types"
	"github.com/go-ego/riot/utils"
)

// Ranker ranker
type Ranker struct {
	lock struct {
		sync.RWMutex

		fields map[string]interface{}
		docs   map[string]bool
		// new
		content map[string]string
		attri   map[string]interface{}
	}

	idOnly      bool
	initialized bool
}

// Init init ranker
func (ranker *Ranker) Init(onlyID ...bool) {
	if ranker.initialized == true {
		log.Fatal("The Ranker can not be initialized twice.")
	}
	ranker.initialized = true

	if len(onlyID) > 0 {
		ranker.idOnly = onlyID[0]
	}

	ranker.lock.fields = make(map[string]interface{})
	ranker.lock.docs = make(map[string]bool)

	if !ranker.idOnly {
		// new
		ranker.lock.content = make(map[string]string)
		ranker.lock.attri = make(map[string]interface{})
	}
}

// AddDoc add doc
// 给某个文档添加评分字段
func (ranker *Ranker) AddDoc(
	// docId uint64, fields interface{}, content string, attri interface{}) {
	docId string, fields interface{}, content ...interface{}) {
	if ranker.initialized == false {
		log.Fatal("The Ranker has not been initialized.")
	}

	ranker.lock.Lock()
	ranker.lock.fields[docId] = fields
	ranker.lock.docs[docId] = true

	if !ranker.idOnly {
		// new
		if len(content) > 0 {
			ranker.lock.content[docId] = content[0].(string)
		}

		if len(content) > 1 {
			ranker.lock.attri[docId] = content[1]
			// ranker.lock.attri[docId] = attri
		}
	}

	ranker.lock.Unlock()
}

// RemoveDoc 删除某个文档的评分字段
func (ranker *Ranker) RemoveDoc(docId string) {
	if ranker.initialized == false {
		log.Fatal("The Ranker has not been initialized.")
	}

	ranker.lock.Lock()
	delete(ranker.lock.fields, docId)
	delete(ranker.lock.docs, docId)

	if !ranker.idOnly {
		// new
		delete(ranker.lock.content, docId)
		delete(ranker.lock.attri, docId)
	}

	ranker.lock.Unlock()
}

func maxOutput(options types.RankOpts, docsLen int) (int, int) {
	var start, end int
	if options.MaxOutputs != 0 {
		start = utils.MinInt(options.OutputOffset, docsLen)
		end = utils.MinInt(options.OutputOffset+options.MaxOutputs, docsLen)
		return start, end
	}

	start = utils.MinInt(options.OutputOffset, docsLen)
	end = docsLen
	return start, end
}

func (ranker *Ranker) rankOutIDs(docs []types.IndexedDoc, options types.RankOpts,
	countDocsOnly bool) (outputDocs types.ScoredIDs, numDocs int) {
	for _, d := range docs {
		ranker.lock.RLock()
		// 判断 doc 是否存在
		if _, ok := ranker.lock.docs[d.DocId]; ok {

			fs := ranker.lock.fields[d.DocId]
			ranker.lock.RUnlock()

			// 计算评分并剔除没有分值的文档
			scores := options.ScoringCriteria.Score(d, fs)
			if len(scores) > 0 {
				if !countDocsOnly {
					outputDocs = append(outputDocs,
						types.ScoredID{
							DocId:            d.DocId,
							Scores:           scores,
							TokenSnippetLocs: d.TokenSnippetLocs,
							TokenLocs:        d.TokenLocs,
						})
				}
				numDocs++
			}
		} else {
			ranker.lock.RUnlock()
		}
	}

	return
}

// RankDocID rank docs by types.ScoredIDs
func (ranker *Ranker) RankDocID(docs []types.IndexedDoc,
	options types.RankOpts, countDocsOnly bool) (types.ScoredIDs, int) {

	outputDocs, numDocs := ranker.rankOutIDs(docs, options, countDocsOnly)

	// 排序
	if !countDocsOnly {
		if options.ReverseOrder {
			sort.Sort(sort.Reverse(outputDocs))
		} else {
			sort.Sort(outputDocs)
		}
		// 当用户要求只返回部分结果时返回部分结果
		docsLen := len(outputDocs)
		start, end := maxOutput(options, docsLen)

		return outputDocs[start:end], numDocs
	}

	return outputDocs, numDocs
}

func (ranker *Ranker) rankOutDocs(docs []types.IndexedDoc, options types.RankOpts,
	countDocsOnly bool) (outputDocs types.ScoredDocs, numDocs int) {
	for _, d := range docs {
		ranker.lock.RLock()
		// 判断 doc 是否存在
		if _, ok := ranker.lock.docs[d.DocId]; ok {

			fs := ranker.lock.fields[d.DocId]
			content := ranker.lock.content[d.DocId]
			attri := ranker.lock.attri[d.DocId]
			ranker.lock.RUnlock()

			// 计算评分并剔除没有分值的文档
			scores := options.ScoringCriteria.Score(d, fs)
			if len(scores) > 0 {
				if !countDocsOnly {
					scoredID := types.ScoredID{
						DocId:            d.DocId,
						Scores:           scores,
						TokenSnippetLocs: d.TokenSnippetLocs,
						TokenLocs:        d.TokenLocs,
					}

					outputDocs = append(outputDocs,
						types.ScoredDoc{
							ScoredID: scoredID,
							// new
							Fields:  fs,
							Content: content,
							Attri:   attri,
						})
				}
				numDocs++
			}
		} else {
			ranker.lock.RUnlock()
		}
	}

	return
}

// RankDocs rank docs by types.ScoredDocs
func (ranker *Ranker) RankDocs(docs []types.IndexedDoc,
	options types.RankOpts, countDocsOnly bool) (types.ScoredDocs, int) {

	outputDocs, numDocs := ranker.rankOutDocs(docs, options, countDocsOnly)

	// 排序
	if !countDocsOnly {
		if options.ReverseOrder {
			sort.Sort(sort.Reverse(outputDocs))
		} else {
			sort.Sort(outputDocs)
		}
		// 当用户要求只返回部分结果时返回部分结果
		docsLen := len(outputDocs)
		start, end := maxOutput(options, docsLen)

		return outputDocs[start:end], numDocs
	}

	return outputDocs, numDocs
}

// Rank rank docs
// 给文档评分并排序
func (ranker *Ranker) Rank(docs []types.IndexedDoc,
	options types.RankOpts, countDocsOnly bool) (interface{}, int) {

	if ranker.initialized == false {
		log.Fatal("The Ranker has not been initialized.")
	}

	// 对每个文档评分
	if ranker.idOnly {
		outputDocs, numDocs := ranker.RankDocID(docs, options, countDocsOnly)
		return outputDocs, numDocs
	}

	outputDocs, numDocs := ranker.RankDocs(docs, options, countDocsOnly)
	return outputDocs, numDocs
}
<|endoftext|>


<commit_before>package gpt

// Type constants for the GUID for type of partition, see https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_entries
type Type string

// List of GUID partition types
const (
	Unused                   Type = "00000000-0000-0000-0000-000000000000"
	MbrBoot                  Type = "024DEE41-33E7-11D3-9D69-0008C781F39F"
	EFISystemPartition       Type = "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"
	BiosBoot                 Type = "21686148-6449-6E6F-744E-656564454649"
	MicrosoftReserved        Type = "E3C9E316-0B5C-4DB8-817D-F92DF00215AE"
	MicrosoftBasicData       Type = "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7"
	MicrosoftLDMMetadata     Type = "5808C8AA-7E8F-42E0-85D2-E1E90434CFB3"
	MicrosoftLDMData         Type = "AF9B60A0-1431-4F62-BC68-3311714A69AD"
	MicrosoftWindowsRecovery Type = "DE94BBA4-06D1-4D40-A16A-BFD50179D6AC"
	LinuxFilesystem          Type = "0FC63DAF-8483-4772-8E79-3D69D8477DE4"
	LinuxRaid                Type = "A19D880F-05FC-4D3B-A006-743F0F84911E"
	LinuxRootX86             Type = "44479540-F297-41B2-9AF7-D131D5F0458A"
	LinuxRootX86_64          Type = "4F68BCE3-E8CD-4DB1-96E7-FBCAF984B709"
	LinuxRootArm32           Type = "69DAD710-2CE4-4E3C-B16C-21A1D49ABED3"
	LinuxRootArm64           Type = "B921B045-1DF0-41C3-AF44-4C6F280D3FAE"
	LinuxSwap                Type = "0657FD6D-A4AB-43C4-84E5-0933C84B4F4F"
	LinuxLVM                 Type = "E6D6D379-F507-44C2-A23C-238F2A3DF928"
	LinuxDMCrypt             Type = "7FFEC5C9-2D00-49B7-8941-3EA10A5586B7"
	LinuxLUKS                Type = "CA7D7CCB-63ED-4C53-861C-1742536059CC"
	VMWareFilesystem         Type = "AA31E02A-400F-11DB-9590-000C2911D1B8"
)
<commit_msg>Add GUID partition types (#147)<commit_after>package gpt

// Type constants for the GUID for type of partition, see https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_entries
type Type string

// List of GUID partition types
const (
	Unused                   Type = "00000000-0000-0000-0000-000000000000"
	EFISystemPartition       Type = "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"
	MBRPartitionScheme       Type = "024DEE41-33E7-11D3-9D69-0008C781F39F"
	IntelFastFlash           Type = "D3BFE2DE-3DAF-11DF-BA40-E3A556D89593"
	BIOSBoot                 Type = "21686148-6449-6E6F-744E-656564454649"
	SonyBootPartition        Type = "F4019732-066E-4E12-8273-346C5641494F"
	LenovoBootPartition      Type = "BFBFAFE7-A34F-448A-9A5B-6213EB736C22"
	PowerPCPRePBoot          Type = "9E1A2D38-C612-4316-AA26-8B49521E5A8B"
	ONIEBoot                 Type = "7412F7D5-A156-4B13-81DC-867174929325"
	ONIEConfig               Type = "D4E6E2CD-4469-46F3-B5CB-1BFF57AFC149"
	MicrosoftReserved        Type = "E3C9E316-0B5C-4DB8-817D-F92DF00215AE"
	MicrosoftBasicData       Type = "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7"
	MicrosoftLDMMetadata     Type = "5808C8AA-7E8F-42E0-85D2-E1E90434CFB3"
	MicrosoftLDMData         Type = "AF9B60A0-1431-4F62-BC68-3311714A69AD"
	MicrosoftWindowsRecovery Type = "DE94BBA4-06D1-4D40-A16A-BFD50179D6AC"
	IBMGeneralParallelFs     Type = "37AFFC90-EF7D-4E96-91C3-2D7AE055B174"
	MicrosoftStorageSpaces   Type = "E75CAF8F-F680-4CEE-AFA3-B001E56EFC2D"
	HPUXData                 Type = "75894C1E-3AEB-11D3-B7C1-7B03A0000000"
	HPUXService              Type = "E2A1E728-32E3-11D6-A682-7B03A0000000"
	LinuxSwap                Type = "0657FD6D-A4AB-43C4-84E5-0933C84B4F4F"
	LinuxFilesystem          Type = "0FC63DAF-8483-4772-8E79-3D69D8477DE4"
	LinuxServerData          Type = "3B8F8425-20E0-4F3B-907F-1A25A76F98E8"
	LinuxRootX86             Type = "44479540-F297-41B2-9AF7-D131D5F0458A"
	LinuxRootArm             Type = "69DAD710-2CE4-4E3C-B16C-21A1D49ABED3"
	LinuxRootX86_64          Type = "4F68BCE3-E8CD-4DB1-96E7-FBCAF984B709"
	LinuxRootArm64           Type = "B921B045-1DF0-41C3-AF44-4C6F280D3FAE"
	LinuxRootIA64            Type = "993D8D3D-F80E-4225-855A-9DAF8ED7EA97"
	LinuxReserved            Type = "8DA63339-0007-60C0-C436-083AC8230908"
	LinuxHome                Type = "933AC7E1-2EB4-4F13-B844-0E14E2AEF915"
	LinuxRAID                Type = "A19D880F-05FC-4D3B-A006-743F0F84911E"
	LinuxExtendedBoot        Type = "BC13C2FF-59E6-4262-A352-B275FD6F7172"
	LinuxLVM                 Type = "E6D6D379-F507-44C2-A23C-238F2A3DF928"
	LinuxDMCrypt             Type = "7FFEC5C9-2D00-49B7-8941-3EA10A5586B7"
	LinuxLUKS                Type = "CA7D7CCB-63ED-4C53-861C-1742536059CC"
	FreeBSDData              Type = "516E7CB4-6ECF-11D6-8FF8-00022D09712B"
	FreeBSDBoot              Type = "83BD6B9D-7F41-11DC-BE0B-001560B84F0F"
	FreeBSDSwap              Type = "516E7CB5-6ECF-11D6-8FF8-00022D09712B"
	FreeBSDNANDFS            Type = "74BA7DD9-A689-11E1-BD04-00E081286ACF"
	FreeBSDUFS               Type = "516E7CB6-6ECF-11D6-8FF8-00022D09712B"
	FreeBSDZFS               Type = "516E7CBA-6ECF-11D6-8FF8-00022D09712B"
	FreeBSDVinum             Type = "516E7CB8-6ECF-11D6-8FF8-00022D09712B"
	AppleHFS                 Type = "48465300-0000-11AA-AA11-00306543ECAC"
	AppleUFS                 Type = "55465300-0000-11AA-AA11-00306543ECAC"
	AppleAPFS                Type = "7C3457EF-0000-11AA-AA11-00306543ECAC"
	AppleRAID                Type = "52414944-0000-11AA-AA11-00306543ECAC"
	AppleRAIDOffline         Type = "52414944-5F4F-11AA-AA11-00306543ECAC"
	AppleBoot                Type = "426F6F74-0000-11AA-AA11-00306543ECAC"
	AppleLabel               Type = "4C616265-6C00-11AA-AA11-00306543ECAC"
	AppleTVRecovery          Type = "5265636F-7665-11AA-AA11-00306543ECAC"
	AppleCoreStorage         Type = "53746F72-6167-11AA-AA11-00306543ECAC"
	DragonflyLabel32         Type = "9D087404-1CA5-11DC-8817-01301BB8A9F5"
	DragonflySwap            Type = "9D58FDBD-1CA5-11DC-8817-01301BB8A9F5"
	DragonflyUFS             Type = "9D94CE7C-1CA5-11DC-8817-01301BB8A9F5"
	DragonflyVINUM           Type = "9DD4478F-1CA5-11DC-8817-01301BB8A9F5"
	DragonflyCCD             Type = "DBD5211B-1CA5-11DC-8817-01301BB8A9F5"
	DragonflyLabel64         Type = "3D48CE54-1D16-11DC-8696-01301BB8A9F5"
	DragonflyLegacy          Type = "BD215AB2-1D16-11DC-8696-01301BB8A9F5"
	DragonflyHAMMER          Type = "61DC63AC-6E38-11DC-8513-01301BB8A9F5"
	DragonflyHAMMER2         Type = "5CBB9AD1-862D-11DC-A94D-01301BB8A9F5"
	SolarisBoot              Type = "6A82CB45-1DD2-11B2-99A6-080020736631"
	SolarisRoot              Type = "6A85CF4D-1DD2-11B2-99A6-080020736631"
	SolarisUsrAndAppleZFS    Type = "6A898CC3-1DD2-11B2-99A6-080020736631"
	SolarisSwap              Type = "6A87C46F-1DD2-11B2-99A6-080020736631"
	SolarisBackup            Type = "6A8B642B-1DD2-11B2-99A6-080020736631"
	SolarisVar               Type = "6A8EF2E9-1DD2-11B2-99A6-080020736631"
	SolarisHome              Type = "6A90BA39-1DD2-11B2-99A6-080020736631"
	SolarisAlternateSector   Type = "6A9283A5-1DD2-11B2-99A6-080020736631"
	SolarisReserved1         Type = "6A945A3B-1DD2-11B2-99A6-080020736631"
	SolarisReserved2         Type = "6A9630D1-1DD2-11B2-99A6-080020736631"
	SolarisReserved3         Type = "6A980767-1DD2-11B2-99A6-080020736631"
	SolarisReserved4         Type = "6A96237F-1DD2-11B2-99A6-080020736631"
	SolarisReserved5         Type = "6A8D2AC7-1DD2-11B2-99A6-080020736631"
	NetBSDSwap               Type = "49F48D32-B10E-11DC-B99B-0019D1879648"
	NetBSDFFS                Type = "49F48D5A-B10E-11DC-B99B-0019D1879648"
	NetBSDLFS                Type = "49F48D82-B10E-11DC-B99B-0019D1879648"
	NetBSDConcatenated       Type = "2DB519C4-B10E-11DC-B99B-0019D1879648"
	NetBSDEncrypted          Type = "2DB519EC-B10E-11DC-B99B-0019D1879648"
	NetBSDRAID               Type = "49F48DAA-B10E-11DC-B99B-0019D1879648"
	ChromeOSFirmware         Type = "CAB6E88E-ABF3-4102-A07A-D4BB9BE3C1D3"
	ChromeOSKernel           Type = "FE3A2A5D-4F32-41A7-B725-ACCC3285A309"
	ChromeOSRootFs           Type = "3CB8E202-3B7E-47DD-8A3C-7FF2A13CFCEC"
	ChromeOSReserved         Type = "2E0A753D-9E48-43B0-8337-B15192CB1B5E"
	MidnightBSDData          Type = "85D5E45A-237C-11E1-B4B3-E89A8F7FC3A7"
	MidnightBSDBoot          Type = "85D5E45E-237C-11E1-B4B3-E89A8F7FC3A7"
	MidnightBSDSwap          Type = "85D5E45B-237C-11E1-B4B3-E89A8F7FC3A7"
	MidnightBSDUFS           Type = "0394EF8B-237E-11E1-B4B3-E89A8F7FC3A7"
	MidnightBSDZFS           Type = "85D5E45D-237C-11E1-B4B3-E89A8F7FC3A7"
	MidnightBSDVinum         Type = "85D5E45C-237C-11E1-B4B3-E89A8F7FC3A7"
	CephJournal              Type = "45B0969E-9B03-4F30-B4C6-B4B80CEFF106"
	CephEncryptedJournal     Type = "45B0969E-9B03-4F30-B4C6-5EC00CEFF106"
	CephOSD                  Type = "4FBD7E29-9D25-41B8-AFD0-062C0CEFF05D"
	CephCryptOSD             Type = "4FBD7E29-9D25-41B8-AFD0-5EC00CEFF05D"
	CephDiskInCreation       Type = "89C57F98-2FE5-4DC0-89C1-F3AD0CEFF2BE"
	CephCryptDiskInCreation  Type = "89C57F98-2FE5-4DC0-89C1-5EC00CEFF2BE"
	VMwareVMFS               Type = "AA31E02A-400F-11DB-9590-000C2911D1B8"
	VMwareDiagnostic         Type = "9D275380-40AD-11DB-BF97-000C2911D1B8"
	VMwareVirtualSAN         Type = "381CFCCC-7288-11E0-92EE-000C2911D0B2"
	VMwareVirsto             Type = "77719A0C-A4A0-11E3-A47E-000C29745A24"
	VMwareReserved           Type = "9198EFFC-31C0-11DB-8F78-000C2911D1B8"
	OpenBSDData              Type = "824CC7A0-36A8-11E3-890A-952519AD3F61"
	QNX6FileSystem           Type = "CEF5A9AD-73BC-4601-89F3-CDEEEEE321A1"
	Plan9Partition           Type = "C91818F9-8025-47AF-89D2-F030D7000C2C"
	HiFiveUnleashedFSBL      Type = "5B193300-FC78-40CD-8002-E86C45580B47"
	HiFiveUnleashedBBL       Type = "2E54B353-1271-4842-806F-E436D6AF6985"
)
<|endoftext|>


<commit_before>// Hellofs implements a simple "hello world" file system.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	_ "bazil.org/fuse/fs/fstestutil"
)

var Usage = func() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s MOUNTPOINT\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	flag.Usage = Usage
	flag.Parse()

	if flag.NArg() != 1 {
		Usage()
		os.Exit(2)
	}
	mountpoint := flag.Arg(0)

	c, err := fuse.Mount(
		mountpoint,
		fuse.FSName("helloworld"),
		fuse.Subtype("hellofs"),
		fuse.LocalVolume(),
		fuse.VolumeName("Hello world!"),
	)
	if err != nil {
		log.Fatal(err)
	}
	defer c.Close()

	err = fs.Serve(c, FS{})
	if err != nil {
		log.Fatal(err)
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; err != nil {
		log.Fatal(err)
	}
}

// FS implements the hello world file system.
type FS struct{}

func (FS) Root() (fs.Node, fuse.Error) {
	return Dir{}, nil
}

// Dir implements both Node and Handle for the root directory.
type Dir struct{}

func (Dir) Attr() fuse.Attr {
	return fuse.Attr{Inode: 1, Mode: os.ModeDir | 0555}
}

func (Dir) Lookup(name string, intr fs.Intr) (fs.Node, fuse.Error) {
	if name == "hello" {
		return File{}, nil
	}
	return nil, fuse.ENOENT
}

var dirDirs = []fuse.Dirent{
	{Inode: 2, Name: "hello", Type: fuse.DT_File},
}

func (Dir) ReadDir(intr fs.Intr) ([]fuse.Dirent, fuse.Error) {
	return dirDirs, nil
}

// File implements both Node and Handle for the hello file.
type File struct{}

const greeting = "hello, world\n"

func (File) Attr() fuse.Attr {
	return fuse.Attr{Inode: 2, Mode: 0444, Size: uint64(len(greeting))}
}

func (File) ReadAll(intr fs.Intr) ([]byte, fuse.Error) {
	return []byte(greeting), nil
}
<commit_msg>Initial comment describing fuse-gdrive<commit_after>// The fuse-gdrive command makes your Google Drive files accessible as a local mount point.
// It implements a user space filesystem, using the Fuse and Google Drive APIs,
// to allow you to access your files in Google Drive just like a regular local
// filesystem.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	_ "bazil.org/fuse/fs/fstestutil"
)

var Usage = func() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s MOUNTPOINT\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	flag.Usage = Usage
	flag.Parse()

	if flag.NArg() != 1 {
		Usage()
		os.Exit(2)
	}
	mountpoint := flag.Arg(0)

	c, err := fuse.Mount(
		mountpoint,
		fuse.FSName("helloworld"),
		fuse.Subtype("hellofs"),
		fuse.LocalVolume(),
		fuse.VolumeName("Hello world!"),
	)
	if err != nil {
		log.Fatal(err)
	}
	defer c.Close()

	err = fs.Serve(c, FS{})
	if err != nil {
		log.Fatal(err)
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; err != nil {
		log.Fatal(err)
	}
}

// FS implements the hello world file system.
type FS struct{}

func (FS) Root() (fs.Node, fuse.Error) {
	return Dir{}, nil
}

// Dir implements both Node and Handle for the root directory.
type Dir struct{}

func (Dir) Attr() fuse.Attr {
	return fuse.Attr{Inode: 1, Mode: os.ModeDir | 0555}
}

func (Dir) Lookup(name string, intr fs.Intr) (fs.Node, fuse.Error) {
	if name == "hello" {
		return File{}, nil
	}
	return nil, fuse.ENOENT
}

var dirDirs = []fuse.Dirent{
	{Inode: 2, Name: "hello", Type: fuse.DT_File},
}

func (Dir) ReadDir(intr fs.Intr) ([]fuse.Dirent, fuse.Error) {
	return dirDirs, nil
}

// File implements both Node and Handle for the hello file.
type File struct{}

const greeting = "hello, world\n"

func (File) Attr() fuse.Attr {
	return fuse.Attr{Inode: 2, Mode: 0444, Size: uint64(len(greeting))}
}

func (File) ReadAll(intr fs.Intr) ([]byte, fuse.Error) {
	return []byte(greeting), nil
}
<|endoftext|>


<commit_before>package game

import (
	"math"
	"time"

	"github.com/Bredgren/game1/game/camera"
	"github.com/Bredgren/game1/game/sprite"
	"github.com/Bredgren/geo"
	"github.com/hajimehoshi/ebiten"
)

const (
	playerMoveSpeed = 500
	playerGravity   = 50
	playerJumpSpeed = 700
	playerJumpTime  = 500 * time.Millisecond
	playerPunchTime = 200 * time.Millisecond
	playerPunchGap  = 100 * time.Millisecond
)

type playerState int

const (
	awaken playerState = iota
	idle
	playerMove
	playerPunch
)

type player struct {
	pos geo.Vec
	vel geo.Vec

	left             bool    // Move left button is down
	right            bool    // Move right button is down
	move             float64 // Gampad axis for movement
	jump             bool    // Jump button is down
	punch            bool    // Punch button is down
	punchAxis        geo.Vec
	punchWithGamepad bool

	canJump   bool
	isJumping bool
	jumpTime  time.Duration
	flipDir   bool

	punchTime time.Duration
	punchGap  time.Duration

	state playerState

	currentSprite *sprite.Sprite
	awakenSprite  sprite.Sprite
	idleSprite    sprite.Sprite
	moveSprite    sprite.Sprite
	punchSprite   sprite.Sprite
}

func newPlayer() *player {
	p := &player{
		canJump:   true,
		isJumping: false,
		jumpTime:  0,

		state:        idle,
		awakenSprite: sprite.Get("awaken"),
		idleSprite:   sprite.Get("idle"),
		moveSprite:   sprite.Get("move"),
		punchSprite:  sprite.Get("punch"),
	}

	p.currentSprite = &p.idleSprite

	return p
}

func (p *player) awaken() {
	p.currentSprite = &p.awakenSprite
	p.currentSprite.Begin(false)
	p.state = awaken
}

func (p *player) awoke() bool {
	return p.awakenSprite.Ended()
}

func (p *player) doIdle() {
	p.state = idle
	p.currentSprite = &p.idleSprite
}

func (p *player) shouldStartPunch() bool {
	p.punchWithGamepad = p.punchAxis.Len2() > 0.25
	return p.punchGap <= 0 && (p.punch || p.punchWithGamepad)
}

func (p *player) doPunch() {
	p.punch = false
	p.state = playerPunch
	p.currentSprite = &p.punchSprite
	p.punchTime = playerPunchTime
}

func (p *player) update(dt time.Duration) {
	p.punchGap -= dt

	switch p.state {
	case awaken:
		if p.awoke() {
			p.doIdle()
		}
	case idle:
		p.updateMove()
		if p.move != 0 {
			p.state = playerMove
			p.currentSprite = &p.moveSprite
		}
		if p.shouldStartPunch() {
			p.doPunch()
		}
	case playerMove:
		p.updateMove()
		if p.move == 0 {
			p.doIdle()
		}
		if p.shouldStartPunch() {
			p.doPunch()
		}
	case playerPunch:
		p.updateMove()
		p.punchTime -= dt
		if p.punchTime <= 0 {
			p.punchGap = playerPunchGap
			if p.move == 0 {
				p.state = playerMove
				p.currentSprite = &p.moveSprite
			} else {
				p.doIdle()
			}
		}
	}

	switch p.state {
	case awaken:
	case idle:
		p.updateMovement(dt)
	case playerMove:
		p.updateMovement(dt)
	case playerPunch:
		p.updateMovement(dt)
	}

	p.currentSprite.Update(dt)

	// Reset Move for next frame, it will be set each frame by user input.
	p.move = 0
}

func (p *player) draw(dst *ebiten.Image, cam *camera.Camera) {
	pos := cam.ScreenCoords(p.pos)
	geom := ebiten.GeoM{}

	size := p.currentSprite.Size()
	bounds := geo.RectWH(size.XY())
	bounds.SetBottomMid(pos.XY())

	switch p.state {
	case awaken:
	case idle:
	case playerMove:
		p.flipDir = p.vel.X < 0
	case playerPunch:
		mPos := geo.VecXYi(ebiten.CursorPosition())
		pos = geo.VecXY(bounds.Mid())

		if p.punchWithGamepad {
			mPos = pos.Plus(p.punchAxis)
		}

		p.flipDir = mPos.X < pos.X
		angle := mPos.Minus(pos).Angle()
		if p.flipDir {
			angle += math.Pi
			angle *= -1
		}
		geom.Translate(-size.X/2, -size.Y/2)
		geom.Rotate(-angle)
		geom.Translate(size.X/2, size.Y/2)
	}

	if p.flipDir {
		geom.Scale(-1, 1)
		geom.Translate(size.X, 0)
	}
	p.currentSprite.Draw(dst, pos, geom)
}

func (p *player) updateMove() {
	if p.move == 0 { // If gamepad axis isn't being used, check left/right buttons.
		if p.left {
			p.move = -1
		}
		if p.right {
			p.move = 1
		}
	}
}

func (p *player) updateMovement(dt time.Duration) {
	p.vel.X = p.move * playerMoveSpeed

	// Check if it's time to jump before handling jump the jump state so that we start
	// jumping as soon as possible
	if !p.isJumping && p.jump && p.canJump {
		p.isJumping = true
		p.jumpTime = playerJumpTime
	}

	if p.isJumping {
		if p.jumpTime <= 0 || !p.jump {
			p.isJumping = false
		} else {
			p.jumpTime -= dt
			p.vel.Y = geo.Lerp(0, -playerJumpSpeed, p.jumpTime.Seconds()/playerJumpTime.Seconds())
		}
	} else {
		p.vel.Y += playerGravity
	}

	p.pos.Add(p.vel.Times(dt.Seconds()))

	// Ground collision
	if p.pos.Y > 0 {
		p.pos.Y = 0
		p.canJump = true
	} else {
		p.canJump = false
	}
}

func (p *player) Pos() geo.Vec {
	return p.pos
}

func (p *player) SetPos(pos geo.Vec) {
	p.pos = pos
}

func (p *player) handleLeft(down bool) bool {
	p.left = down
	return false
}

func (p *player) handleRight(down bool) bool {
	p.right = down
	return false
}

func (p *player) handleMove(val float64) bool {
	p.move = val
	return false
}

func (p *player) handleJump(down bool) bool {
	p.jump = down
	return false
}

func (p *player) handlePunch(down bool) bool {
	p.punch = down
	return false
}

func (p *player) handlePunchH(val float64) bool {
	p.punchAxis.X = val
	// p.handlePunchAxis()
	return false
}

func (p *player) handlePunchV(val float64) bool {
	p.punchAxis.Y = -val
	// p.handlePunchAxis()
	return false
}

// func (p *player) handlePunchAxis() {
// 	if p.punchAxis.Len2() > 0.25 {
// 		p.punch =
// 	}
// }
<commit_msg>Put transition to move state in function<commit_after>package game

import (
	"math"
	"time"

	"github.com/Bredgren/game1/game/camera"
	"github.com/Bredgren/game1/game/sprite"
	"github.com/Bredgren/geo"
	"github.com/hajimehoshi/ebiten"
)

const (
	playerMoveSpeed = 500
	playerGravity   = 50
	playerJumpSpeed = 700
	playerJumpTime  = 500 * time.Millisecond
	playerPunchTime = 200 * time.Millisecond
	playerPunchGap  = 100 * time.Millisecond
)

type playerState int

const (
	awaken playerState = iota
	idle
	playerMove
	playerPunch
)

type player struct {
	pos geo.Vec
	vel geo.Vec

	left             bool    // Move left button is down
	right            bool    // Move right button is down
	move             float64 // Gampad axis for movement
	jump             bool    // Jump button is down
	punch            bool    // Punch button is down
	punchAxis        geo.Vec
	punchWithGamepad bool

	canJump   bool
	isJumping bool
	jumpTime  time.Duration
	flipDir   bool

	punchTime time.Duration
	punchGap  time.Duration

	state playerState

	currentSprite *sprite.Sprite
	awakenSprite  sprite.Sprite
	idleSprite    sprite.Sprite
	moveSprite    sprite.Sprite
	punchSprite   sprite.Sprite
}

func newPlayer() *player {
	p := &player{
		canJump:   true,
		isJumping: false,
		jumpTime:  0,

		state:        idle,
		awakenSprite: sprite.Get("awaken"),
		idleSprite:   sprite.Get("idle"),
		moveSprite:   sprite.Get("move"),
		punchSprite:  sprite.Get("punch"),
	}

	p.currentSprite = &p.idleSprite

	return p
}

func (p *player) awaken() {
	p.currentSprite = &p.awakenSprite
	p.currentSprite.Begin(false)
	p.state = awaken
}

func (p *player) awoke() bool {
	return p.awakenSprite.Ended()
}

func (p *player) doIdle() {
	p.state = idle
	p.currentSprite = &p.idleSprite
}

func (p *player) doMove() {
	p.state = playerMove
	p.currentSprite = &p.moveSprite
}

func (p *player) shouldStartPunch() bool {
	p.punchWithGamepad = p.punchAxis.Len2() > 0.25
	return p.punchGap <= 0 && (p.punch || p.punchWithGamepad)
}

func (p *player) doPunch() {
	p.punch = false
	p.state = playerPunch
	p.currentSprite = &p.punchSprite
	p.punchTime = playerPunchTime
}

func (p *player) update(dt time.Duration) {
	p.punchGap -= dt

	switch p.state {
	case awaken:
		if p.awoke() {
			p.doIdle()
		}
	case idle:
		p.updateMove()
		if p.move != 0 {
			p.doMove()
		}
		if p.shouldStartPunch() {
			p.doPunch()
		}
	case playerMove:
		p.updateMove()
		if p.move == 0 {
			p.doIdle()
		}
		if p.shouldStartPunch() {
			p.doPunch()
		}
	case playerPunch:
		p.updateMove()
		p.punchTime -= dt
		if p.punchTime <= 0 {
			p.punchGap = playerPunchGap
			if p.move == 0 {
				p.doMove()
			} else {
				p.doIdle()
			}
		}
	}

	switch p.state {
	case awaken:
	case idle:
		p.updateMovement(dt)
	case playerMove:
		p.updateMovement(dt)
	case playerPunch:
		p.updateMovement(dt)
	}

	p.currentSprite.Update(dt)

	// Reset Move for next frame, it will be set each frame by user input.
	p.move = 0
}

func (p *player) draw(dst *ebiten.Image, cam *camera.Camera) {
	pos := cam.ScreenCoords(p.pos)
	geom := ebiten.GeoM{}

	size := p.currentSprite.Size()
	bounds := geo.RectWH(size.XY())
	bounds.SetBottomMid(pos.XY())

	switch p.state {
	case awaken:
	case idle:
	case playerMove:
		p.flipDir = p.vel.X < 0
	case playerPunch:
		mPos := geo.VecXYi(ebiten.CursorPosition())
		pos = geo.VecXY(bounds.Mid())

		if p.punchWithGamepad {
			mPos = pos.Plus(p.punchAxis)
		}

		p.flipDir = mPos.X < pos.X
		angle := mPos.Minus(pos).Angle()
		if p.flipDir {
			angle += math.Pi
			angle *= -1
		}
		geom.Translate(-size.X/2, -size.Y/2)
		geom.Rotate(-angle)
		geom.Translate(size.X/2, size.Y/2)
	}

	if p.flipDir {
		geom.Scale(-1, 1)
		geom.Translate(size.X, 0)
	}
	p.currentSprite.Draw(dst, pos, geom)
}

func (p *player) updateMove() {
	if p.move == 0 { // If gamepad axis isn't being used, check left/right buttons.
		if p.left {
			p.move = -1
		}
		if p.right {
			p.move = 1
		}
	}
}

func (p *player) updateMovement(dt time.Duration) {
	p.vel.X = p.move * playerMoveSpeed

	// Check if it's time to jump before handling jump the jump state so that we start
	// jumping as soon as possible
	if !p.isJumping && p.jump && p.canJump {
		p.isJumping = true
		p.jumpTime = playerJumpTime
	}

	if p.isJumping {
		if p.jumpTime <= 0 || !p.jump {
			p.isJumping = false
		} else {
			p.jumpTime -= dt
			p.vel.Y = geo.Lerp(0, -playerJumpSpeed, p.jumpTime.Seconds()/playerJumpTime.Seconds())
		}
	} else {
		p.vel.Y += playerGravity
	}

	p.pos.Add(p.vel.Times(dt.Seconds()))

	// Ground collision
	if p.pos.Y > 0 {
		p.pos.Y = 0
		p.canJump = true
	} else {
		p.canJump = false
	}
}

func (p *player) Pos() geo.Vec {
	return p.pos
}

func (p *player) SetPos(pos geo.Vec) {
	p.pos = pos
}

func (p *player) handleLeft(down bool) bool {
	p.left = down
	return false
}

func (p *player) handleRight(down bool) bool {
	p.right = down
	return false
}

func (p *player) handleMove(val float64) bool {
	p.move = val
	return false
}

func (p *player) handleJump(down bool) bool {
	p.jump = down
	return false
}

func (p *player) handlePunch(down bool) bool {
	p.punch = down
	return false
}

func (p *player) handlePunchH(val float64) bool {
	p.punchAxis.X = val
	// p.handlePunchAxis()
	return false
}

func (p *player) handlePunchV(val float64) bool {
	p.punchAxis.Y = -val
	// p.handlePunchAxis()
	return false
}

// func (p *player) handlePunchAxis() {
// 	if p.punchAxis.Len2() > 0.25 {
// 		p.punch =
// 	}
// }
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package generators has the generators for the set-gen utility.
package generators

import (
	"io"

	"k8s.io/gengo/args"
	"k8s.io/gengo/generator"
	"k8s.io/gengo/namer"
	"k8s.io/gengo/types"

	"k8s.io/klog"
)

// NameSystems returns the name system used by the generators in this package.
func NameSystems() namer.NameSystems {
	return namer.NameSystems{
		"public":  namer.NewPublicNamer(0),
		"private": namer.NewPrivateNamer(0),
		"raw":     namer.NewRawNamer("", nil),
	}
}

// DefaultNameSystem returns the default name system for ordering the types to be
// processed by the generators in this package.
func DefaultNameSystem() string {
	return "public"
}

// Packages makes the sets package definition.
func Packages(_ *generator.Context, arguments *args.GeneratorArgs) generator.Packages {
	boilerplate, err := arguments.LoadGoBoilerplate()
	if err != nil {
		klog.Fatalf("Failed loading boilerplate: %v", err)
	}

	return generator.Packages{&generator.DefaultPackage{
		PackageName: "sets",
		PackagePath: arguments.OutputPackagePath,
		HeaderText:  boilerplate,
		PackageDocumentation: []byte(
			`// Package sets has auto-generated set types.
`),
		// GeneratorFunc returns a list of generators. Each generator makes a
		// single file.
		GeneratorFunc: func(c *generator.Context) (generators []generator.Generator) {
			generators = []generator.Generator{
				// Always generate a "doc.go" file.
				generator.DefaultGen{OptionalName: "doc"},
				// Make a separate file for the Empty type, since it's shared by every type.
				generator.DefaultGen{
					OptionalName: "empty",
					OptionalBody: []byte(emptyTypeDecl),
				},
			}
			// Since we want a file per type that we generate a set for, we
			// have to provide a function for this.
			for _, t := range c.Order {
				generators = append(generators, &genSet{
					DefaultGen: generator.DefaultGen{
						// Use the privatized version of the
						// type name as the file name.
						//
						// TODO: make a namer that converts
						// camelCase to '-' separation for file
						// names?
						OptionalName: c.Namers["private"].Name(t),
					},
					outputPackage: arguments.OutputPackagePath,
					typeToMatch:   t,
					imports:       generator.NewImportTracker(),
				})
			}
			return generators
		},
		FilterFunc: func(c *generator.Context, t *types.Type) bool {
			// It would be reasonable to filter by the type's package here.
			// It might be necessary if your input directory has a big
			// import graph.
			switch t.Kind {
			case types.Map, types.Slice, types.Pointer:
				// These types can't be keys in a map.
				return false
			case types.Builtin:
				return true
			case types.Struct:
				// Only some structs can be keys in a map. This is triggered by the line
				// // +genset
				// or
				// // +genset=true
				return extractBoolTagOrDie("genset", t.CommentLines) == true
			}
			return false
		},
	}}
}

// genSet produces a file with a set for a single type.
type genSet struct {
	generator.DefaultGen
	outputPackage string
	typeToMatch   *types.Type
	imports       namer.ImportTracker
}

// Filter ignores all but one type because we're making a single file per type.
func (g *genSet) Filter(c *generator.Context, t *types.Type) bool { return t == g.typeToMatch }

func (g *genSet) Namers(c *generator.Context) namer.NameSystems {
	return namer.NameSystems{
		"raw": namer.NewRawNamer(g.outputPackage, g.imports),
	}
}

func (g *genSet) Imports(c *generator.Context) (imports []string) {
	return append(g.imports.ImportLines(), "reflect", "sort")
}

// args constructs arguments for templates. Usage:
// g.args(t, "key1", value1, "key2", value2, ...)
//
// 't' is loaded with the key 'type'.
//
// We could use t directly as the argument, but doing it this way makes it easy
// to mix in additional parameters. This feature is not used in this set
// generator, but is present as an example.
func (g *genSet) args(t *types.Type, kv ...interface{}) interface{} {
	m := map[interface{}]interface{}{"type": t}
	for i := 0; i < len(kv)/2; i++ {
		m[kv[i*2]] = kv[i*2+1]
	}
	return m
}

// GenerateType makes the body of a file implementing a set for type t.
func (g *genSet) GenerateType(c *generator.Context, t *types.Type, w io.Writer) error {
	sw := generator.NewSnippetWriter(w, c, "$", "$")
	sw.Do(setCode, g.args(t))
	sw.Do("func less$.type|public$(lhs, rhs $.type|raw$) bool {\n", g.args(t))
	g.lessBody(sw, t)
	sw.Do("}\n", g.args(t))
	return sw.Error()
}

func (g *genSet) lessBody(sw *generator.SnippetWriter, t *types.Type) {
	// TODO: make this recursive, handle pointers and multiple nested structs...
	switch t.Kind {
	case types.Struct:
		for _, m := range types.FlattenMembers(t.Members) {
			sw.Do("if lhs.$.Name$ < rhs.$.Name$ { return true }\n", m)
			sw.Do("if lhs.$.Name$ > rhs.$.Name$ { return false }\n", m)
		}
		sw.Do("return false\n", nil)
	default:
		sw.Do("return lhs < rhs\n", nil)
	}
}

// written to the "empty.go" file.
var emptyTypeDecl = `
// Empty is public since it is used by some internal API objects for conversions between external
// string arrays and internal sets, and conversion logic requires public types today.
type Empty struct{}
`

// Written for every type. If you've never used text/template before:
// $.type$ refers to the source type; |public means to
// call the function giving the public name, |raw the raw type name.
var setCode = `// sets.$.type|public$ is a set of $.type|raw$s, implemented via map[$.type|raw$]struct{} for minimal memory consumption.
type $.type|public$ map[$.type|raw$]Empty

// New$.type|public$ creates a $.type|public$ from a list of values.
func New$.type|public$(items ...$.type|raw$) $.type|public$ {
	ss := $.type|public${}
	ss.Insert(items...)
	return ss
}

// $.type|public$KeySet creates a $.type|public$ from a keys of a map[$.type|raw$](? extends interface{}).
// If the value passed in is not actually a map, this will panic.
func $.type|public$KeySet(theMap interface{}) $.type|public$ {
	v := reflect.ValueOf(theMap)
	ret := $.type|public${}

	for _, keyValue := range v.MapKeys() {
		ret.Insert(keyValue.Interface().($.type|raw$))
	}
	return ret
}

// Insert adds items to the set.
func (s $.type|public$) Insert(items ...$.type|raw$) {
	for _, item := range items {
		s[item] = Empty{}
	}
}

// Delete removes all items from the set.
func (s $.type|public$) Delete(items ...$.type|raw$) {
	for _, item := range items {
		delete(s, item)
	}
}

// Has returns true if and only if item is contained in the set.
func (s $.type|public$) Has(item $.type|raw$) bool {
	_, contained := s[item]
	return contained
}

// HasAll returns true if and only if all items are contained in the set.
func (s $.type|public$) HasAll(items ...$.type|raw$) bool {
	for _, item := range items {
		if !s.Has(item) {
			return false
		}
	}
	return true
}

// HasAny returns true if any items are contained in the set.
func (s $.type|public$) HasAny(items ...$.type|raw$) bool {
	for _, item := range items {
		if s.Has(item) {
			return true
		}
	}
	return false
}

// Difference returns a set of objects that are not in s2
// For example:
// s1 = {a1, a2, a3}
// s2 = {a1, a2, a4, a5}
// s1.Difference(s2) = {a3}
// s2.Difference(s1) = {a4, a5}
func (s $.type|public$) Difference(s2 $.type|public$) $.type|public$ {
	result := New$.type|public$()
	for key := range s {
		if !s2.Has(key) {
			result.Insert(key)
		}
	}
	return result
}

// Union returns a new set which includes items in either s1 or s2.
// For example:
// s1 = {a1, a2}
// s2 = {a3, a4}
// s1.Union(s2) = {a1, a2, a3, a4}
// s2.Union(s1) = {a1, a2, a3, a4}
func (s1 $.type|public$) Union(s2 $.type|public$) $.type|public$ {
	result := New$.type|public$()
	for key := range s1 {
		result.Insert(key)
	}
	for key := range s2 {
		result.Insert(key)
	}
	return result
}

// Intersection returns a new set which includes the item in BOTH s1 and s2
// For example:
// s1 = {a1, a2}
// s2 = {a2, a3}
// s1.Intersection(s2) = {a2}
func (s1 $.type|public$) Intersection(s2 $.type|public$) $.type|public$ {
	var walk, other $.type|public$
	result := New$.type|public$()
	if s1.Len() < s2.Len() {
		walk = s1
		other = s2
	} else {
		walk = s2
		other = s1
	}
	for key := range walk {
		if other.Has(key) {
			result.Insert(key)
		}
	}
	return result
}

// IsSuperset returns true if and only if s1 is a superset of s2.
func (s1 $.type|public$) IsSuperset(s2 $.type|public$) bool {
	for item := range s2 {
		if !s1.Has(item) {
			return false
		}
	}
	return true
}

// Equal returns true if and only if s1 is equal (as a set) to s2.
// Two sets are equal if their membership is identical.
// (In practice, this means same elements, order doesn't matter)
func (s1 $.type|public$) Equal(s2 $.type|public$) bool {
	return len(s1) == len(s2) && s1.IsSuperset(s2)
}

type sortableSliceOf$.type|public$ []$.type|raw$

func (s sortableSliceOf$.type|public$) Len() int { return len(s) }
func (s sortableSliceOf$.type|public$) Less(i, j int) bool { return less$.type|public$(s[i], s[j]) }
func (s sortableSliceOf$.type|public$) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// List returns the contents as a sorted $.type|raw$ slice.
func (s $.type|public$) List() []$.type|raw$ {
	res := make(sortableSliceOf$.type|public$, 0, len(s))
	for key := range s {
		res = append(res, key)
	}
	sort.Sort(res)
	return []$.type|raw$(res)
}

// UnsortedList returns the slice with contents in random order.
func (s $.type|public$) UnsortedList() []$.type|raw$ {
	res :=make([]$.type|raw$, 0, len(s))
	for key := range s {
		res = append(res, key)
	}
	return res
}

// Returns a single element from the set.
func (s $.type|public$) PopAny() ($.type|raw$, bool) {
	for key := range s {
		s.Delete(key)
		return key, true
	}
	var zeroValue $.type|raw$
	return zeroValue, false
}

// Len returns the size of the set.
func (s $.type|public$) Len() int {
	return len(s)
}

`
<commit_msg>Sets: Make Insert and Delete return the Set<commit_after>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package generators has the generators for the set-gen utility.
package generators

import (
	"io"

	"k8s.io/gengo/args"
	"k8s.io/gengo/generator"
	"k8s.io/gengo/namer"
	"k8s.io/gengo/types"

	"k8s.io/klog"
)

// NameSystems returns the name system used by the generators in this package.
func NameSystems() namer.NameSystems {
	return namer.NameSystems{
		"public":  namer.NewPublicNamer(0),
		"private": namer.NewPrivateNamer(0),
		"raw":     namer.NewRawNamer("", nil),
	}
}

// DefaultNameSystem returns the default name system for ordering the types to be
// processed by the generators in this package.
func DefaultNameSystem() string {
	return "public"
}

// Packages makes the sets package definition.
func Packages(_ *generator.Context, arguments *args.GeneratorArgs) generator.Packages {
	boilerplate, err := arguments.LoadGoBoilerplate()
	if err != nil {
		klog.Fatalf("Failed loading boilerplate: %v", err)
	}

	return generator.Packages{&generator.DefaultPackage{
		PackageName: "sets",
		PackagePath: arguments.OutputPackagePath,
		HeaderText:  boilerplate,
		PackageDocumentation: []byte(
			`// Package sets has auto-generated set types.
`),
		// GeneratorFunc returns a list of generators. Each generator makes a
		// single file.
		GeneratorFunc: func(c *generator.Context) (generators []generator.Generator) {
			generators = []generator.Generator{
				// Always generate a "doc.go" file.
				generator.DefaultGen{OptionalName: "doc"},
				// Make a separate file for the Empty type, since it's shared by every type.
				generator.DefaultGen{
					OptionalName: "empty",
					OptionalBody: []byte(emptyTypeDecl),
				},
			}
			// Since we want a file per type that we generate a set for, we
			// have to provide a function for this.
			for _, t := range c.Order {
				generators = append(generators, &genSet{
					DefaultGen: generator.DefaultGen{
						// Use the privatized version of the
						// type name as the file name.
						//
						// TODO: make a namer that converts
						// camelCase to '-' separation for file
						// names?
						OptionalName: c.Namers["private"].Name(t),
					},
					outputPackage: arguments.OutputPackagePath,
					typeToMatch:   t,
					imports:       generator.NewImportTracker(),
				})
			}
			return generators
		},
		FilterFunc: func(c *generator.Context, t *types.Type) bool {
			// It would be reasonable to filter by the type's package here.
			// It might be necessary if your input directory has a big
			// import graph.
			switch t.Kind {
			case types.Map, types.Slice, types.Pointer:
				// These types can't be keys in a map.
				return false
			case types.Builtin:
				return true
			case types.Struct:
				// Only some structs can be keys in a map. This is triggered by the line
				// // +genset
				// or
				// // +genset=true
				return extractBoolTagOrDie("genset", t.CommentLines) == true
			}
			return false
		},
	}}
}

// genSet produces a file with a set for a single type.
type genSet struct {
	generator.DefaultGen
	outputPackage string
	typeToMatch   *types.Type
	imports       namer.ImportTracker
}

// Filter ignores all but one type because we're making a single file per type.
func (g *genSet) Filter(c *generator.Context, t *types.Type) bool { return t == g.typeToMatch }

func (g *genSet) Namers(c *generator.Context) namer.NameSystems {
	return namer.NameSystems{
		"raw": namer.NewRawNamer(g.outputPackage, g.imports),
	}
}

func (g *genSet) Imports(c *generator.Context) (imports []string) {
	return append(g.imports.ImportLines(), "reflect", "sort")
}

// args constructs arguments for templates. Usage:
// g.args(t, "key1", value1, "key2", value2, ...)
//
// 't' is loaded with the key 'type'.
//
// We could use t directly as the argument, but doing it this way makes it easy
// to mix in additional parameters. This feature is not used in this set
// generator, but is present as an example.
func (g *genSet) args(t *types.Type, kv ...interface{}) interface{} {
	m := map[interface{}]interface{}{"type": t}
	for i := 0; i < len(kv)/2; i++ {
		m[kv[i*2]] = kv[i*2+1]
	}
	return m
}

// GenerateType makes the body of a file implementing a set for type t.
func (g *genSet) GenerateType(c *generator.Context, t *types.Type, w io.Writer) error {
	sw := generator.NewSnippetWriter(w, c, "$", "$")
	sw.Do(setCode, g.args(t))
	sw.Do("func less$.type|public$(lhs, rhs $.type|raw$) bool {\n", g.args(t))
	g.lessBody(sw, t)
	sw.Do("}\n", g.args(t))
	return sw.Error()
}

func (g *genSet) lessBody(sw *generator.SnippetWriter, t *types.Type) {
	// TODO: make this recursive, handle pointers and multiple nested structs...
	switch t.Kind {
	case types.Struct:
		for _, m := range types.FlattenMembers(t.Members) {
			sw.Do("if lhs.$.Name$ < rhs.$.Name$ { return true }\n", m)
			sw.Do("if lhs.$.Name$ > rhs.$.Name$ { return false }\n", m)
		}
		sw.Do("return false\n", nil)
	default:
		sw.Do("return lhs < rhs\n", nil)
	}
}

// written to the "empty.go" file.
var emptyTypeDecl = `
// Empty is public since it is used by some internal API objects for conversions between external
// string arrays and internal sets, and conversion logic requires public types today.
type Empty struct{}
`

// Written for every type. If you've never used text/template before:
// $.type$ refers to the source type; |public means to
// call the function giving the public name, |raw the raw type name.
var setCode = `// sets.$.type|public$ is a set of $.type|raw$s, implemented via map[$.type|raw$]struct{} for minimal memory consumption.
type $.type|public$ map[$.type|raw$]Empty

// New$.type|public$ creates a $.type|public$ from a list of values.
func New$.type|public$(items ...$.type|raw$) $.type|public$ {
	ss := $.type|public${}
	ss.Insert(items...)
	return ss
}

// $.type|public$KeySet creates a $.type|public$ from a keys of a map[$.type|raw$](? extends interface{}).
// If the value passed in is not actually a map, this will panic.
func $.type|public$KeySet(theMap interface{}) $.type|public$ {
	v := reflect.ValueOf(theMap)
	ret := $.type|public${}

	for _, keyValue := range v.MapKeys() {
		ret.Insert(keyValue.Interface().($.type|raw$))
	}
	return ret
}

// Insert adds items to the set.
func (s $.type|public$) Insert(items ...$.type|raw$) $.type|public$ {
	for _, item := range items {
		s[item] = Empty{}
	}
	return s
}

// Delete removes all items from the set.
func (s $.type|public$) Delete(items ...$.type|raw$) $.type|public$ {
	for _, item := range items {
		delete(s, item)
	}
	return s
}

// Has returns true if and only if item is contained in the set.
func (s $.type|public$) Has(item $.type|raw$) bool {
	_, contained := s[item]
	return contained
}

// HasAll returns true if and only if all items are contained in the set.
func (s $.type|public$) HasAll(items ...$.type|raw$) bool {
	for _, item := range items {
		if !s.Has(item) {
			return false
		}
	}
	return true
}

// HasAny returns true if any items are contained in the set.
func (s $.type|public$) HasAny(items ...$.type|raw$) bool {
	for _, item := range items {
		if s.Has(item) {
			return true
		}
	}
	return false
}

// Difference returns a set of objects that are not in s2
// For example:
// s1 = {a1, a2, a3}
// s2 = {a1, a2, a4, a5}
// s1.Difference(s2) = {a3}
// s2.Difference(s1) = {a4, a5}
func (s $.type|public$) Difference(s2 $.type|public$) $.type|public$ {
	result := New$.type|public$()
	for key := range s {
		if !s2.Has(key) {
			result.Insert(key)
		}
	}
	return result
}

// Union returns a new set which includes items in either s1 or s2.
// For example:
// s1 = {a1, a2}
// s2 = {a3, a4}
// s1.Union(s2) = {a1, a2, a3, a4}
// s2.Union(s1) = {a1, a2, a3, a4}
func (s1 $.type|public$) Union(s2 $.type|public$) $.type|public$ {
	result := New$.type|public$()
	for key := range s1 {
		result.Insert(key)
	}
	for key := range s2 {
		result.Insert(key)
	}
	return result
}

// Intersection returns a new set which includes the item in BOTH s1 and s2
// For example:
// s1 = {a1, a2}
// s2 = {a2, a3}
// s1.Intersection(s2) = {a2}
func (s1 $.type|public$) Intersection(s2 $.type|public$) $.type|public$ {
	var walk, other $.type|public$
	result := New$.type|public$()
	if s1.Len() < s2.Len() {
		walk = s1
		other = s2
	} else {
		walk = s2
		other = s1
	}
	for key := range walk {
		if other.Has(key) {
			result.Insert(key)
		}
	}
	return result
}

// IsSuperset returns true if and only if s1 is a superset of s2.
func (s1 $.type|public$) IsSuperset(s2 $.type|public$) bool {
	for item := range s2 {
		if !s1.Has(item) {
			return false
		}
	}
	return true
}

// Equal returns true if and only if s1 is equal (as a set) to s2.
// Two sets are equal if their membership is identical.
// (In practice, this means same elements, order doesn't matter)
func (s1 $.type|public$) Equal(s2 $.type|public$) bool {
	return len(s1) == len(s2) && s1.IsSuperset(s2)
}

type sortableSliceOf$.type|public$ []$.type|raw$

func (s sortableSliceOf$.type|public$) Len() int { return len(s) }
func (s sortableSliceOf$.type|public$) Less(i, j int) bool { return less$.type|public$(s[i], s[j]) }
func (s sortableSliceOf$.type|public$) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// List returns the contents as a sorted $.type|raw$ slice.
func (s $.type|public$) List() []$.type|raw$ {
	res := make(sortableSliceOf$.type|public$, 0, len(s))
	for key := range s {
		res = append(res, key)
	}
	sort.Sort(res)
	return []$.type|raw$(res)
}

// UnsortedList returns the slice with contents in random order.
func (s $.type|public$) UnsortedList() []$.type|raw$ {
	res :=make([]$.type|raw$, 0, len(s))
	for key := range s {
		res = append(res, key)
	}
	return res
}

// Returns a single element from the set.
func (s $.type|public$) PopAny() ($.type|raw$, bool) {
	for key := range s {
		s.Delete(key)
		return key, true
	}
	var zeroValue $.type|raw$
	return zeroValue, false
}

// Len returns the size of the set.
func (s $.type|public$) Len() int {
	return len(s)
}

`
<|endoftext|>


<commit_before>// Copyright 2012 SocialCode. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package gelf

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"compress/zlib"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"path"
	"runtime"
	"strings"
	"sync"
	"time"
)

// Writer implements io.Writer and is used to send both discrete
// messages to a graylog2 server, or data from a stream-oriented
// interface (like the functions in log).
type Writer struct {
	mu               sync.Mutex
	conn             net.Conn
	hostname         string
	Facility         string // defaults to current process name
	CompressionLevel int    // one of the consts from compress/flate
	CompressionType  CompressType
}

// What compression type the writer should use when sending messages
// to the graylog2 server
type CompressType int

const (
	CompressGzip CompressType = iota
	CompressZlib
)

// Message represents the contents of the GELF message.  It is gzipped
// before sending.
type Message struct {
	Version  string                 `json:"version"`
	Host     string                 `json:"host"`
	Short    string                 `json:"short_message"`
	Full     string                 `json:"full_message"`
	TimeUnix float64                `json:"timestamp"`
	Level    int32                  `json:"level"`
	Facility string                 `json:"facility"`
	Extra    map[string]interface{} `json:"-"`
}

// Used to control GELF chunking.  Should be less than (MTU - len(UDP
// header)).
//
// TODO: generate dynamically using Path MTU Discovery?
const (
	ChunkSize        = 1420
	chunkedHeaderLen = 12
	chunkedDataLen   = ChunkSize - chunkedHeaderLen
)

var (
	magicChunked = []byte{0x1e, 0x0f}
	magicZlib    = []byte{0x78}
	magicGzip    = []byte{0x1f, 0x8b}
)

// Syslog severity levels
const (
	LOG_EMERG   = int32(0)
	LOG_ALERT   = int32(1)
	LOG_CRIT    = int32(2)
	LOG_ERR     = int32(3)
	LOG_WARNING = int32(4)
	LOG_NOTICE  = int32(5)
	LOG_INFO    = int32(6)
	LOG_DEBUG   = int32(7)
)

// numChunks returns the number of GELF chunks necessary to transmit
// the given compressed buffer.
func numChunks(b []byte) int {
	lenB := len(b)
	if lenB <= ChunkSize {
		return 1
	}
	return len(b)/chunkedDataLen + 1
}

// New returns a new GELF Writer.  This writer can be used to send the
// output of the standard Go log functions to a central GELF server by
// passing it to log.SetOutput()
func NewWriter(addr string) (*Writer, error) {
	var err error
	w := new(Writer)
	w.CompressionLevel = flate.BestSpeed

	if w.conn, err = net.Dial("udp", addr); err != nil {
		return nil, err
	}
	if w.hostname, err = os.Hostname(); err != nil {
		return nil, err
	}

	w.Facility = path.Base(os.Args[0])

	return w, nil
}

// writes the gzip compressed byte array to the connection as a series
// of GELF chunked messages.  The header format is documented at
// https://github.com/Graylog2/graylog2-docs/wiki/GELF as:
//
//     2-byte magic (0x1e 0x0f), 8 byte id, 1 byte sequence id, 1 byte
//     total, chunk-data
func (w *Writer) writeChunked(zBytes []byte) (err error) {
	b := make([]byte, 0, ChunkSize)
	buf := bytes.NewBuffer(b)
	nChunksI := numChunks(zBytes)
	if nChunksI > 255 {
		return fmt.Errorf("msg too large, would need %d chunks", nChunksI)
	}
	nChunks := uint8(nChunksI)
	// use urandom to get a unique message id
	msgId := make([]byte, 8)
	n, err := io.ReadFull(rand.Reader, msgId)
	if err != nil || n != 8 {
		return fmt.Errorf("rand.Reader: %d/%s", n, err)
	}

	bytesLeft := len(zBytes)
	for i := uint8(0); i < nChunks; i++ {
		buf.Reset()
		// manually write header.  Don't care about
		// host/network byte order, because the spec only
		// deals in individual bytes.
		buf.Write(magicChunked) //magic
		buf.Write(msgId)
		buf.WriteByte(i)
		buf.WriteByte(nChunks)
		// slice out our chunk from zBytes
		chunkLen := chunkedDataLen
		if chunkLen > bytesLeft {
			chunkLen = bytesLeft
		}
		off := int(i) * chunkedDataLen
		chunk := zBytes[off : off+chunkLen]
		buf.Write(chunk)

		// write this chunk, and make sure the write was good
		n, err := w.conn.Write(buf.Bytes())
		if err != nil {
			return fmt.Errorf("Write (chunk %d/%d): %s", i,
				nChunks, err)
		}
		if n != len(buf.Bytes()) {
			return fmt.Errorf("Write len: (chunk %d/%d) (%d/%d)",
				i, nChunks, n, len(buf.Bytes()))
		}

		bytesLeft -= chunkLen
	}

	if bytesLeft != 0 {
		return fmt.Errorf("error: %d bytes left after sending", bytesLeft)
	}
	return nil
}

// WriteMessage sends the specified message to the GELF server
// specified in the call to New().  It assumes all the fields are
// filled out appropriately.  In general, clients will want to use
// Write, rather than WriteMessage.
func (w *Writer) WriteMessage(m *Message) (err error) {
	mBuf := bytes.NewBuffer(nil)
	if _, err = m.MarshalJSONBuf(mBuf); err != nil {
		return
	}
	mBytes := mBuf.Bytes()

	var zBuf bytes.Buffer
	var zw io.WriteCloser
	switch w.CompressionType {
	case CompressGzip:
		zw, err = gzip.NewWriterLevel(&zBuf, w.CompressionLevel)
	case CompressZlib:
		zw, err = zlib.NewWriterLevel(&zBuf, w.CompressionLevel)
	default:
		panic(fmt.Sprintf("unknown compression type %d",
			w.CompressionType))
	}
	if err != nil {
		return
	}
	if _, err = zw.Write(mBytes); err != nil {
		return
	}
	zw.Close()

	zBytes := zBuf.Bytes()
	if numChunks(zBytes) > 1 {
		return w.writeChunked(zBytes)
	}

	n, err := w.conn.Write(zBytes)
	if err != nil {
		return
	}
	if n != len(zBytes) {
		return fmt.Errorf("bad write (%d/%d)", n, len(zBytes))
	}

	return nil
}

// Close connection and interrupt blocked Read or Write operations
func (w *Writer) Close() error {
	return w.conn.Close()
}

/*
func (w *Writer) Alert(m string) (err error)
func (w *Writer) Close() error
func (w *Writer) Crit(m string) (err error)
func (w *Writer) Debug(m string) (err error)
func (w *Writer) Emerg(m string) (err error)
func (w *Writer) Err(m string) (err error)
func (w *Writer) Info(m string) (err error)
func (w *Writer) Notice(m string) (err error)
func (w *Writer) Warning(m string) (err error)
*/

// getCaller returns the filename and the line info of a function
// further down in the call stack.  Passing 0 in as callDepth would
// return info on the function calling getCallerIgnoringLog, 1 the
// parent function, and so on.  Any suffixes passed to getCaller are
// path fragments like "/pkg/log/log.go", and functions in the call
// stack from that file are ignored.
func getCaller(callDepth int, suffixesToIgnore ...string) (file string, line int) {
	// bump by 1 to ignore the getCaller (this) stackframe
	callDepth++
outer:
	for {
		var ok bool
		_, file, line, ok = runtime.Caller(callDepth)
		if !ok {
			file = "???"
			line = 0
			break
		}

		for _, s := range suffixesToIgnore {
			if strings.HasSuffix(file, s) {
				callDepth++
				continue outer
			}
		}
		break
	}
	return
}

func getCallerIgnoringLogMulti(callDepth int) (string, int) {
	// the +1 is to ignore this (getCallerIgnoringLogMulti) frame
	return getCaller(callDepth+1, "/pkg/log/log.go", "/pkg/io/multi.go")
}

// Write encodes the given string in a GELF message and sends it to
// the server specified in New().
func (w *Writer) Write(p []byte) (n int, err error) {

	// 1 for the function that called us.
	file, line := getCallerIgnoringLogMulti(1)

	// remove trailing and leading whitespace
	p = bytes.TrimSpace(p)

	// If there are newlines in the message, use the first line
	// for the short message and set the full message to the
	// original input.  If the input has no newlines, stick the
	// whole thing in Short.
	short := p
	full := []byte("")
	if i := bytes.IndexRune(p, '\n'); i > 0 {
		short = p[:i]
		full = p
	}

	m := Message{
		Version:  "1.1",
		Host:     w.hostname,
		Short:    string(short),
		Full:     string(full),
		TimeUnix: float64(time.Now().Unix()),
		Level:    6, // info
		Facility: w.Facility,
		Extra: map[string]interface{}{
			"_file": file,
			"_line": line,
		},
	}

	if err = w.WriteMessage(&m); err != nil {
		return 0, err
	}

	return len(p), nil
}

func (m *Message) MarshalJSONBuf(buf *bytes.Buffer) (int, error) {
	var err error
	var b, eb []byte

	b, err = json.Marshal(m)
	if err != nil {
		return 0, err
	}

	if len(m.Extra) == 0 {
		return buf.Write(b)
	}

	if eb, err = json.Marshal(m.Extra); err != nil {
		return 0, err
	}

	// merge serialized message + serialized extra map
	b[len(b)-1] = ','
	b = append(b, eb[1:len(eb)]...)
	return buf.Write(b)
}

func (m *Message) UnmarshalJSON(data []byte) error {
	i := make(map[string]interface{}, 16)
	if err := json.Unmarshal(data, &i); err != nil {
		return err
	}
	for k, v := range i {
		if k[0] == '_' {
			if m.Extra == nil {
				m.Extra = make(map[string]interface{}, 1)
			}
			m.Extra[k] = v
			continue
		}
		switch k {
		case "version":
			m.Version = v.(string)
		case "host":
			m.Host = v.(string)
		case "short_message":
			m.Short = v.(string)
		case "full_message":
			m.Full = v.(string)
		case "timestamp":
			m.TimeUnix = v.(float64)
		case "level":
			m.Level = int32(v.(float64))
		case "facility":
			m.Facility = v.(string)
		}
	}
	return nil
}
<commit_msg>add bytes pool for json and gzip<commit_after>// Copyright 2012 SocialCode. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package gelf

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"compress/zlib"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"path"
	"runtime"
	"strings"
	"sync"
	"time"
)

// Writer implements io.Writer and is used to send both discrete
// messages to a graylog2 server, or data from a stream-oriented
// interface (like the functions in log).
type Writer struct {
	mu               sync.Mutex
	conn             net.Conn
	hostname         string
	Facility         string // defaults to current process name
	CompressionLevel int    // one of the consts from compress/flate
	CompressionType  CompressType
}

// What compression type the writer should use when sending messages
// to the graylog2 server
type CompressType int

const (
	CompressGzip CompressType = iota
	CompressZlib
)

// Message represents the contents of the GELF message.  It is gzipped
// before sending.
type Message struct {
	Version  string                 `json:"version"`
	Host     string                 `json:"host"`
	Short    string                 `json:"short_message"`
	Full     string                 `json:"full_message"`
	TimeUnix float64                `json:"timestamp"`
	Level    int32                  `json:"level"`
	Facility string                 `json:"facility"`
	Extra    map[string]interface{} `json:"-"`
}

// Used to control GELF chunking.  Should be less than (MTU - len(UDP
// header)).
//
// TODO: generate dynamically using Path MTU Discovery?
const (
	ChunkSize        = 1420
	chunkedHeaderLen = 12
	chunkedDataLen   = ChunkSize - chunkedHeaderLen
)

var (
	magicChunked = []byte{0x1e, 0x0f}
	magicZlib    = []byte{0x78}
	magicGzip    = []byte{0x1f, 0x8b}
)

// Syslog severity levels
const (
	LOG_EMERG   = int32(0)
	LOG_ALERT   = int32(1)
	LOG_CRIT    = int32(2)
	LOG_ERR     = int32(3)
	LOG_WARNING = int32(4)
	LOG_NOTICE  = int32(5)
	LOG_INFO    = int32(6)
	LOG_DEBUG   = int32(7)
)

// numChunks returns the number of GELF chunks necessary to transmit
// the given compressed buffer.
func numChunks(b []byte) int {
	lenB := len(b)
	if lenB <= ChunkSize {
		return 1
	}
	return len(b)/chunkedDataLen + 1
}

// New returns a new GELF Writer.  This writer can be used to send the
// output of the standard Go log functions to a central GELF server by
// passing it to log.SetOutput()
func NewWriter(addr string) (*Writer, error) {
	var err error
	w := new(Writer)
	w.CompressionLevel = flate.BestSpeed

	if w.conn, err = net.Dial("udp", addr); err != nil {
		return nil, err
	}
	if w.hostname, err = os.Hostname(); err != nil {
		return nil, err
	}

	w.Facility = path.Base(os.Args[0])

	return w, nil
}

// writes the gzip compressed byte array to the connection as a series
// of GELF chunked messages.  The header format is documented at
// https://github.com/Graylog2/graylog2-docs/wiki/GELF as:
//
//     2-byte magic (0x1e 0x0f), 8 byte id, 1 byte sequence id, 1 byte
//     total, chunk-data
func (w *Writer) writeChunked(zBytes []byte) (err error) {
	b := make([]byte, 0, ChunkSize)
	buf := bytes.NewBuffer(b)
	nChunksI := numChunks(zBytes)
	if nChunksI > 255 {
		return fmt.Errorf("msg too large, would need %d chunks", nChunksI)
	}
	nChunks := uint8(nChunksI)
	// use urandom to get a unique message id
	msgId := make([]byte, 8)
	n, err := io.ReadFull(rand.Reader, msgId)
	if err != nil || n != 8 {
		return fmt.Errorf("rand.Reader: %d/%s", n, err)
	}

	bytesLeft := len(zBytes)
	for i := uint8(0); i < nChunks; i++ {
		buf.Reset()
		// manually write header.  Don't care about
		// host/network byte order, because the spec only
		// deals in individual bytes.
		buf.Write(magicChunked) //magic
		buf.Write(msgId)
		buf.WriteByte(i)
		buf.WriteByte(nChunks)
		// slice out our chunk from zBytes
		chunkLen := chunkedDataLen
		if chunkLen > bytesLeft {
			chunkLen = bytesLeft
		}
		off := int(i) * chunkedDataLen
		chunk := zBytes[off : off+chunkLen]
		buf.Write(chunk)

		// write this chunk, and make sure the write was good
		n, err := w.conn.Write(buf.Bytes())
		if err != nil {
			return fmt.Errorf("Write (chunk %d/%d): %s", i,
				nChunks, err)
		}
		if n != len(buf.Bytes()) {
			return fmt.Errorf("Write len: (chunk %d/%d) (%d/%d)",
				i, nChunks, n, len(buf.Bytes()))
		}

		bytesLeft -= chunkLen
	}

	if bytesLeft != 0 {
		return fmt.Errorf("error: %d bytes left after sending", bytesLeft)
	}
	return nil
}

// 1k bytes buffer by default
var bufPool = sync.Pool{
	New: func() interface{} {
		return bytes.NewBuffer(make([]byte, 0, 1024))
	},
}

func newBuffer() *bytes.Buffer {
	b := bufPool.Get().(*bytes.Buffer)
	if b != nil {
		b.Reset()
		return b
	}
	return bytes.NewBuffer(nil)
}

// WriteMessage sends the specified message to the GELF server
// specified in the call to New().  It assumes all the fields are
// filled out appropriately.  In general, clients will want to use
// Write, rather than WriteMessage.
func (w *Writer) WriteMessage(m *Message) (err error) {
	mBuf := newBuffer()
	defer bufPool.Put(mBuf)
	if _, err = m.MarshalJSONBuf(mBuf); err != nil {
		return
	}
	mBytes := mBuf.Bytes()

	zBuf := newBuffer()
	defer bufPool.Put(zBuf)

	var zw io.WriteCloser
	switch w.CompressionType {
	case CompressGzip:
		zw, err = gzip.NewWriterLevel(zBuf, w.CompressionLevel)
	case CompressZlib:
		zw, err = zlib.NewWriterLevel(zBuf, w.CompressionLevel)
	default:
		panic(fmt.Sprintf("unknown compression type %d",
			w.CompressionType))
	}
	if err != nil {
		return
	}
	if _, err = zw.Write(mBytes); err != nil {
		return
	}
	zw.Close()

	zBytes := zBuf.Bytes()
	if numChunks(zBytes) > 1 {
		return w.writeChunked(zBytes)
	}

	n, err := w.conn.Write(zBytes)
	if err != nil {
		return
	}
	if n != len(zBytes) {
		return fmt.Errorf("bad write (%d/%d)", n, len(zBytes))
	}

	return nil
}

// Close connection and interrupt blocked Read or Write operations
func (w *Writer) Close() error {
	return w.conn.Close()
}

/*
func (w *Writer) Alert(m string) (err error)
func (w *Writer) Close() error
func (w *Writer) Crit(m string) (err error)
func (w *Writer) Debug(m string) (err error)
func (w *Writer) Emerg(m string) (err error)
func (w *Writer) Err(m string) (err error)
func (w *Writer) Info(m string) (err error)
func (w *Writer) Notice(m string) (err error)
func (w *Writer) Warning(m string) (err error)
*/

// getCaller returns the filename and the line info of a function
// further down in the call stack.  Passing 0 in as callDepth would
// return info on the function calling getCallerIgnoringLog, 1 the
// parent function, and so on.  Any suffixes passed to getCaller are
// path fragments like "/pkg/log/log.go", and functions in the call
// stack from that file are ignored.
func getCaller(callDepth int, suffixesToIgnore ...string) (file string, line int) {
	// bump by 1 to ignore the getCaller (this) stackframe
	callDepth++
outer:
	for {
		var ok bool
		_, file, line, ok = runtime.Caller(callDepth)
		if !ok {
			file = "???"
			line = 0
			break
		}

		for _, s := range suffixesToIgnore {
			if strings.HasSuffix(file, s) {
				callDepth++
				continue outer
			}
		}
		break
	}
	return
}

func getCallerIgnoringLogMulti(callDepth int) (string, int) {
	// the +1 is to ignore this (getCallerIgnoringLogMulti) frame
	return getCaller(callDepth+1, "/pkg/log/log.go", "/pkg/io/multi.go")
}

// Write encodes the given string in a GELF message and sends it to
// the server specified in New().
func (w *Writer) Write(p []byte) (n int, err error) {

	// 1 for the function that called us.
	file, line := getCallerIgnoringLogMulti(1)

	// remove trailing and leading whitespace
	p = bytes.TrimSpace(p)

	// If there are newlines in the message, use the first line
	// for the short message and set the full message to the
	// original input.  If the input has no newlines, stick the
	// whole thing in Short.
	short := p
	full := []byte("")
	if i := bytes.IndexRune(p, '\n'); i > 0 {
		short = p[:i]
		full = p
	}

	m := Message{
		Version:  "1.1",
		Host:     w.hostname,
		Short:    string(short),
		Full:     string(full),
		TimeUnix: float64(time.Now().Unix()),
		Level:    6, // info
		Facility: w.Facility,
		Extra: map[string]interface{}{
			"_file": file,
			"_line": line,
		},
	}

	if err = w.WriteMessage(&m); err != nil {
		return 0, err
	}

	return len(p), nil
}

func (m *Message) MarshalJSONBuf(buf *bytes.Buffer) (int, error) {
	var err error
	var b, eb []byte

	b, err = json.Marshal(m)
	if err != nil {
		return 0, err
	}

	if len(m.Extra) == 0 {
		return buf.Write(b)
	}

	if eb, err = json.Marshal(m.Extra); err != nil {
		return 0, err
	}

	// merge serialized message + serialized extra map
	b[len(b)-1] = ','
	b = append(b, eb[1:len(eb)]...)
	return buf.Write(b)
}

func (m *Message) UnmarshalJSON(data []byte) error {
	i := make(map[string]interface{}, 16)
	if err := json.Unmarshal(data, &i); err != nil {
		return err
	}
	for k, v := range i {
		if k[0] == '_' {
			if m.Extra == nil {
				m.Extra = make(map[string]interface{}, 1)
			}
			m.Extra[k] = v
			continue
		}
		switch k {
		case "version":
			m.Version = v.(string)
		case "host":
			m.Host = v.(string)
		case "short_message":
			m.Short = v.(string)
		case "full_message":
			m.Full = v.(string)
		case "timestamp":
			m.TimeUnix = v.(float64)
		case "level":
			m.Level = int32(v.(float64))
		case "facility":
			m.Facility = v.(string)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/nick11roberts/geozip"
)

func main() {
	lat, lon := 45.321, 164.4533
	bucket := geozip.Encode(lat, lon)
	fmt.Println("")
}
<commit_msg>About to rename the repo...<commit_after>package main

import (
	"fmt"

	"github.com/nick11roberts/geozip"
)

func main() {
	lat, lon := 45.321, 164.4533
	bucket := geozip.Encode(lat, lon)
	fmt.Println(bucket)
}
<|endoftext|>