<commit_before>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package ghttp

import (
	"context"
	"fmt"
	"github.com/gogf/gf/internal/intlog"
	"github.com/gogf/gf/os/gres"
	"github.com/gogf/gf/os/gsession"
	"github.com/gogf/gf/os/gview"
	"github.com/gogf/gf/util/guid"
	"net/http"
	"strings"
	"time"

	"github.com/gogf/gf/os/gtime"
	"github.com/gogf/gf/text/gregex"
)

// Request is the context object for a request.
type Request struct {
	*http.Request
	Server          *Server                // Server.
	Cookie          *Cookie                // Cookie.
	Session         *gsession.Session      // Session.
	Response        *Response              // Corresponding Response of this request.
	Router          *Router                // Matched Router for this request. Note that it's not available in HOOK handler.
	EnterTime       int64                  // Request starting time in microseconds.
	LeaveTime       int64                  // Request ending time in microseconds.
	Middleware      *Middleware            // Middleware manager.
	StaticFile      *StaticFile            // Static file object for static file serving.
	context         context.Context        // Custom context for internal usage purpose.
	handlers        []*handlerParsedItem   // All matched handlers containing handler, hook and middleware for this request.
	hasHookHandler  bool                   // A bool marking whether there's hook handler in the handlers for performance purpose.
	hasServeHandler bool                   // A bool marking whether there's serving handler in the handlers for performance purpose.
	parsedQuery     bool                   // A bool marking whether the GET parameters parsed.
	parsedBody      bool                   // A bool marking whether the request body parsed.
	parsedForm      bool                   // A bool marking whether request Form parsed for HTTP method PUT, POST, PATCH.
	paramsMap       map[string]interface{} // Custom parameters map.
	routerMap       map[string]string      // Router parameters map, which might be nil if there're no router parameters.
	queryMap        map[string]interface{} // Query parameters map, which is nil if there's no query string.
	formMap         map[string]interface{} // Form parameters map, which is nil if there's no form data from client.
	bodyMap         map[string]interface{} // Body parameters map, which might be nil if there're no body content.
	error           error                  // Current executing error of the request.
	exit            bool                   // A bool marking whether current request is exited.
	parsedHost      string                 // The parsed host name for current host used by GetHost function.
	clientIp        string                 // The parsed client ip for current host used by GetClientIp function.
	bodyContent     []byte                 // Request body content.
	isFileRequest   bool                   // A bool marking whether current request is file serving.
	viewObject      *gview.View            // Custom template view engine object for this response.
	viewParams      gview.Params           // Custom template view variables for this response.
}

// StaticFile is the file struct for static file service.
type StaticFile struct {
	File  *gres.File // Resource file object.
	Path  string     // File path.
	IsDir bool       // Is directory.
}

// newRequest creates and returns a new request object.
func newRequest(s *Server, r *http.Request, w http.ResponseWriter) *Request {
	request := &Request{
		Server:    s,
		Request:   r,
		Response:  newResponse(s, w),
		EnterTime: gtime.TimestampMilli(),
	}
	request.Cookie = GetCookie(request)
	request.Session = s.sessionManager.New(request.GetSessionId())
	request.Response.Request = request
	request.Middleware = &Middleware{
		request: request,
	}
	// Custom session id creating function.
	err := request.Session.SetIdFunc(func(ttl time.Duration) string {
		var (
			address = request.RemoteAddr
			header  = fmt.Sprintf("%v", request.Header)
		)
		intlog.Print(address, header)
		return guid.S([]byte(address), []byte(header))
	})
	if err != nil {
		panic(err)
	}
	return request
}

// WebSocket upgrades current request as a websocket request.
// It returns a new WebSocket object if success, or the error if failure.
// Note that the request should be a websocket request, or it will surely fail upgrading.
func (r *Request) WebSocket() (*WebSocket, error) {
	if conn, err := wsUpgrader.Upgrade(r.Response.Writer, r.Request, nil); err == nil {
		return &WebSocket{
			conn,
		}, nil
	} else {
		return nil, err
	}
}

// Exit exits executing of current HTTP handler.
func (r *Request) Exit() {
	panic(gEXCEPTION_EXIT)
}

// ExitAll exits executing of current and following HTTP handlers.
func (r *Request) ExitAll() {
	r.exit = true
	panic(gEXCEPTION_EXIT_ALL)
}

// ExitHook exits executing of current and following HTTP HOOK handlers.
func (r *Request) ExitHook() {
	panic(gEXCEPTION_EXIT_HOOK)
}

// IsExited checks and returns whether current request is exited.
func (r *Request) IsExited() bool {
	return r.exit
}

// GetHost returns current request host name, which might be a domain or an IP without port.
func (r *Request) GetHost() string {
	if len(r.parsedHost) == 0 {
		array, _ := gregex.MatchString(`(.+):(\d+)`, r.Host)
		if len(array) > 1 {
			r.parsedHost = array[1]
		} else {
			r.parsedHost = r.Host
		}
	}
	return r.parsedHost
}

// IsFileRequest checks and returns whether current request is serving file.
func (r *Request) IsFileRequest() bool {
	return r.isFileRequest
}

// IsAjaxRequest checks and returns whether current request is an AJAX request.
func (r *Request) IsAjaxRequest() bool {
	return strings.EqualFold(r.Header.Get("X-Requested-With"), "XMLHttpRequest")
}

// GetClientIp returns the client ip of this request without port.
func (r *Request) GetClientIp() string {
	if len(r.clientIp) == 0 {
		if r.clientIp = r.Header.Get("X-Real-IP"); r.clientIp == "" {
			array, _ := gregex.MatchString(`(.+):(\d+)`, r.RemoteAddr)
			if len(array) > 1 {
				r.clientIp = array[1]
			} else {
				r.clientIp = r.RemoteAddr
			}
		}
	}
	return r.clientIp
}

// GetUrl returns current URL of this request.
func (r *Request) GetUrl() string {
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	return fmt.Sprintf(`%s://%s%s`, scheme, r.Host, r.URL.String())
}

// GetSessionId retrieves and returns session id from cookie or header.
func (r *Request) GetSessionId() string {
	id := r.Cookie.GetSessionId()
	if id == "" {
		id = r.Header.Get(r.Server.GetSessionIdName())
	}
	return id
}

// GetReferer returns referer of this request.
func (r *Request) GetReferer() string {
	return r.Header.Get("Referer")
}

// GetError returns the error occurs in the procedure of the request.
// It returns nil if there's no error.
func (r *Request) GetError() error {
	return r.error
}
<commit_msg>add get client realIp<commit_after>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package ghttp

import (
	"context"
	"fmt"
	"github.com/gogf/gf/internal/intlog"
	"github.com/gogf/gf/os/gres"
	"github.com/gogf/gf/os/gsession"
	"github.com/gogf/gf/os/gview"
	"github.com/gogf/gf/text/gstr"
	"github.com/gogf/gf/util/guid"
	"net/http"
	"strings"
	"time"

	"github.com/gogf/gf/os/gtime"
	"github.com/gogf/gf/text/gregex"
)

// Request is the context object for a request.
type Request struct {
	*http.Request
	Server          *Server                // Server.
	Cookie          *Cookie                // Cookie.
	Session         *gsession.Session      // Session.
	Response        *Response              // Corresponding Response of this request.
	Router          *Router                // Matched Router for this request. Note that it's not available in HOOK handler.
	EnterTime       int64                  // Request starting time in microseconds.
	LeaveTime       int64                  // Request ending time in microseconds.
	Middleware      *Middleware            // Middleware manager.
	StaticFile      *StaticFile            // Static file object for static file serving.
	context         context.Context        // Custom context for internal usage purpose.
	handlers        []*handlerParsedItem   // All matched handlers containing handler, hook and middleware for this request.
	hasHookHandler  bool                   // A bool marking whether there's hook handler in the handlers for performance purpose.
	hasServeHandler bool                   // A bool marking whether there's serving handler in the handlers for performance purpose.
	parsedQuery     bool                   // A bool marking whether the GET parameters parsed.
	parsedBody      bool                   // A bool marking whether the request body parsed.
	parsedForm      bool                   // A bool marking whether request Form parsed for HTTP method PUT, POST, PATCH.
	paramsMap       map[string]interface{} // Custom parameters map.
	routerMap       map[string]string      // Router parameters map, which might be nil if there're no router parameters.
	queryMap        map[string]interface{} // Query parameters map, which is nil if there's no query string.
	formMap         map[string]interface{} // Form parameters map, which is nil if there's no form data from client.
	bodyMap         map[string]interface{} // Body parameters map, which might be nil if there're no body content.
	error           error                  // Current executing error of the request.
	exit            bool                   // A bool marking whether current request is exited.
	parsedHost      string                 // The parsed host name for current host used by GetHost function.
	clientIp        string                 // The parsed client ip for current host used by GetClientIp function.
	clientRealIp    string                 // The parsed client real ip for current host used by GetClientIp function.
	bodyContent     []byte                 // Request body content.
	isFileRequest   bool                   // A bool marking whether current request is file serving.
	viewObject      *gview.View            // Custom template view engine object for this response.
	viewParams      gview.Params           // Custom template view variables for this response.
}

// StaticFile is the file struct for static file service.
type StaticFile struct {
	File  *gres.File // Resource file object.
	Path  string     // File path.
	IsDir bool       // Is directory.
}

// newRequest creates and returns a new request object.
func newRequest(s *Server, r *http.Request, w http.ResponseWriter) *Request {
	request := &Request{
		Server:    s,
		Request:   r,
		Response:  newResponse(s, w),
		EnterTime: gtime.TimestampMilli(),
	}
	request.Cookie = GetCookie(request)
	request.Session = s.sessionManager.New(request.GetSessionId())
	request.Response.Request = request
	request.Middleware = &Middleware{
		request: request,
	}
	// Custom session id creating function.
	err := request.Session.SetIdFunc(func(ttl time.Duration) string {
		var (
			address = request.RemoteAddr
			header  = fmt.Sprintf("%v", request.Header)
		)
		intlog.Print(address, header)
		return guid.S([]byte(address), []byte(header))
	})
	if err != nil {
		panic(err)
	}
	return request
}

// WebSocket upgrades current request as a websocket request.
// It returns a new WebSocket object if success, or the error if failure.
// Note that the request should be a websocket request, or it will surely fail upgrading.
func (r *Request) WebSocket() (*WebSocket, error) {
	if conn, err := wsUpgrader.Upgrade(r.Response.Writer, r.Request, nil); err == nil {
		return &WebSocket{
			conn,
		}, nil
	} else {
		return nil, err
	}
}

// Exit exits executing of current HTTP handler.
func (r *Request) Exit() {
	panic(gEXCEPTION_EXIT)
}

// ExitAll exits executing of current and following HTTP handlers.
func (r *Request) ExitAll() {
	r.exit = true
	panic(gEXCEPTION_EXIT_ALL)
}

// ExitHook exits executing of current and following HTTP HOOK handlers.
func (r *Request) ExitHook() {
	panic(gEXCEPTION_EXIT_HOOK)
}

// IsExited checks and returns whether current request is exited.
func (r *Request) IsExited() bool {
	return r.exit
}

// GetHost returns current request host name, which might be a domain or an IP without port.
func (r *Request) GetHost() string {
	if len(r.parsedHost) == 0 {
		array, _ := gregex.MatchString(`(.+):(\d+)`, r.Host)
		if len(array) > 1 {
			r.parsedHost = array[1]
		} else {
			r.parsedHost = r.Host
		}
	}
	return r.parsedHost
}

// IsFileRequest checks and returns whether current request is serving file.
func (r *Request) IsFileRequest() bool {
	return r.isFileRequest
}

// IsAjaxRequest checks and returns whether current request is an AJAX request.
func (r *Request) IsAjaxRequest() bool {
	return strings.EqualFold(r.Header.Get("X-Requested-With"), "XMLHttpRequest")
}

// GetClientIp returns the client ip of this request without port.
func (r *Request) GetClientIp() string {
	if len(r.clientIp) == 0 {
		if r.clientIp = r.Header.Get("X-Real-IP"); r.clientIp == "" {
			array, _ := gregex.MatchString(`(.+):(\d+)`, r.RemoteAddr)
			if len(array) > 1 {
				r.clientIp = array[1]
			} else {
				r.clientIp = r.RemoteAddr
			}
		}
	}
	return r.clientIp
}

func (r *Request) GetClientRealIp() string {
	if len(r.clientRealIp) == 0 {
		readlIps := r.Header.Get("X-Forwarded-For")
		if readlIps != "" && len(readlIps) != 0 && !strings.EqualFold("unknown", readlIps) {
			ipArray := gstr.Split(readlIps, ",")
			//	g.Log().Info("ipArray",ipArray)
			r.clientRealIp = ipArray[0]
		}
		if r.clientRealIp == "" || len(readlIps) == 0 || strings.EqualFold("unknown", readlIps) {
			r.clientRealIp = r.Header.Get("Proxy-Client-IP")
		}
		if r.clientRealIp == "" || len(readlIps) == 0 ||  strings.EqualFold("unknown", readlIps) {
			r.clientRealIp = r.Header.Get("WL-Proxy-Client-IP")
		}
		if r.clientRealIp == "" || len(readlIps) == 0 ||  strings.EqualFold("unknown", readlIps) {
			r.clientRealIp = r.Header.Get("HTTP_CLIENT_IP")
		}
		if r.clientRealIp == "" || len(readlIps) == 0 ||  strings.EqualFold("unknown", readlIps) {
			r.clientRealIp = r.Header.Get("HTTP_X_FORWARDED_FOR")
		}
		if r.clientRealIp == "" || len(readlIps) == 0 ||  strings.EqualFold("unknown", readlIps) {
			r.clientRealIp = r.Header.Get("X-Real-IP")
		}
		if r.clientRealIp == "" || len(readlIps) == 0 ||  strings.EqualFold("unknown", readlIps) {
			r.clientRealIp = r.RemoteAddr
		}
	}
	return r.clientRealIp
}

// GetUrl returns current URL of this request.
func (r *Request) GetUrl() string {
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	return fmt.Sprintf(`%s://%s%s`, scheme, r.Host, r.URL.String())
}

// GetSessionId retrieves and returns session id from cookie or header.
func (r *Request) GetSessionId() string {
	id := r.Cookie.GetSessionId()
	if id == "" {
		id = r.Header.Get(r.Server.GetSessionIdName())
	}
	return id
}

// GetReferer returns referer of this request.
func (r *Request) GetReferer() string {
	return r.Header.Get("Referer")
}

// GetError returns the error occurs in the procedure of the request.
// It returns nil if there's no error.
func (r *Request) GetError() error {
	return r.error
}
<|endoftext|>


<commit_before>// Package mucp transparently forwards the incoming request using a go-micro client.
package mucp

import (
	"context"
	"io"
	"strings"
	"sync"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/codec/bytes"
	"github.com/micro/go-micro/config/options"
	"github.com/micro/go-micro/network/proxy"
	"github.com/micro/go-micro/network/router"
	"github.com/micro/go-micro/server"

	"github.com/micro/go-micro/network/router/table"
)

// Proxy will transparently proxy requests to an endpoint.
// If no endpoint is specified it will call a service using the client.
type Proxy struct {
	// embed options
	options.Options

	// Endpoint specifies the fixed service endpoint to call.
	Endpoint string

	// The client to use for outbound requests
	Client client.Client

	// The router for routes
	Router router.Router

	// A fib of routes service:address
	sync.RWMutex
	Routes map[string]map[uint64]table.Route
}

// read client request and write to server
func readLoop(r server.Request, s client.Stream) error {
	// request to backend server
	req := s.Request()

	for {
		// get data from client
		//  no need to decode it
		body, err := r.Read()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		// get the header from client
		hdr := r.Header()
		msg := &codec.Message{
			Type:   codec.Request,
			Header: hdr,
			Body:   body,
		}

		// write the raw request
		err = req.Codec().Write(msg, nil)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

// toNodes returns a list of node addresses from given routes
func toNodes(routes map[uint64]table.Route) []string {
	var nodes []string
	for _, node := range routes {
		address := node.Address
		if len(node.Gateway) > 0 {
			address = node.Gateway
		}
		nodes = append(nodes, address)
	}
	return nodes
}

func (p *Proxy) getRoute(service string) ([]string, error) {
	// lookup the route cache first
	p.Lock()
	routes, ok := p.Routes[service]
	if ok {
		p.Unlock()
		return toNodes(routes), nil
	}
	p.Routes[service] = make(map[uint64]table.Route)
	p.Unlock()

	// if the router is broken return error
	if status := p.Router.Status(); status.Code == router.Error {
		return nil, status.Error
	}

	// lookup the routes in the router
	results, err := p.Router.Lookup(table.NewQuery(table.QueryService(service)))
	if err != nil {
		return nil, err
	}

	// update the proxy cache
	p.Lock()
	for _, route := range results {
		p.Routes[service][route.Hash()] = route
	}
	routes = p.Routes[service]
	p.Unlock()

	return toNodes(routes), nil
}

// ServeRequest honours the server.Router interface
func (p *Proxy) ServeRequest(ctx context.Context, req server.Request, rsp server.Response) error {
	// service name
	service := req.Service()
	endpoint := req.Endpoint()
	var addresses []string

	// call a specific backend endpoint either by name or address
	if len(p.Endpoint) > 0 {
		// address:port
		if parts := strings.Split(p.Endpoint, ":"); len(parts) > 1 {
			addresses = []string{p.Endpoint}
		} else {
			// get route for endpoint from router
			addr, err := p.getRoute(p.Endpoint)
			if err != nil {
				return err
			}
			// set the address
			addresses = addr
			// set the name
			service = p.Endpoint
		}
	} else {
		// no endpoint was specified just lookup the route
		// get route for endpoint from router
		addr, err := p.getRoute(service)
		if err != nil {
			return err
		}
		addresses = addr
	}

	var opts []client.CallOption

	// set address if available
	if len(addresses) > 0 {
		opts = append(opts, client.WithAddress(addresses...))
	}

	// read initial request
	body, err := req.Read()
	if err != nil {
		return err
	}

	// create new request with raw bytes body
	creq := p.Client.NewRequest(service, endpoint, &bytes.Frame{body}, client.WithContentType(req.ContentType()))

	// create new stream
	stream, err := p.Client.Stream(ctx, creq, opts...)
	if err != nil {
		return err
	}
	defer stream.Close()

	// create client request read loop
	go readLoop(req, stream)

	// get raw response
	resp := stream.Response()

	// create server response write loop
	for {
		// read backend response body
		body, err := resp.Read()
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}

		// read backend response header
		hdr := resp.Header()

		// write raw response header to client
		rsp.WriteHeader(hdr)

		// write raw response body to client
		err = rsp.Write(body)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}

	return nil
}

// NewSingleHostProxy returns a proxy which sends requests to a single backend
func NewSingleHostProxy(endpoint string) *Proxy {
	return &Proxy{
		Options:  options.NewOptions(),
		Endpoint: endpoint,
	}
}

// NewProxy returns a new proxy which will route based on mucp headers
func NewProxy(opts ...options.Option) proxy.Proxy {
	p := new(Proxy)
	p.Options = options.NewOptions(opts...)
	p.Options.Init(options.WithString("mucp"))

	// get endpoint
	ep, ok := p.Options.Values().Get("proxy.endpoint")
	if ok {
		p.Endpoint = ep.(string)
	}

	// get client
	c, ok := p.Options.Values().Get("proxy.client")
	if ok {
		p.Client = c.(client.Client)
	}

	// set the default client
	if p.Client == nil {
		p.Client = client.DefaultClient
	}

	// get router
	r, ok := p.Options.Values().Get("proxy.router")
	if ok {
		p.Router = r.(router.Router)
	}

	// create default router and start it
	if p.Router == nil {
		p.Router = router.DefaultRouter
	}

	// routes cache
	p.Routes = make(map[string]map[uint64]table.Route)

	return p
}
<commit_msg>Adds route watcher to mucp.Proxy<commit_after>// Package mucp transparently forwards the incoming request using a go-micro client.
package mucp

import (
	"context"
	"fmt"
	"io"
	"strings"
	"sync"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/codec/bytes"
	"github.com/micro/go-micro/config/options"
	"github.com/micro/go-micro/network/proxy"
	"github.com/micro/go-micro/network/router"
	"github.com/micro/go-micro/server"

	"github.com/micro/go-micro/network/router/table"
)

// Proxy will transparently proxy requests to an endpoint.
// If no endpoint is specified it will call a service using the client.
type Proxy struct {
	// embed options
	options.Options

	// Endpoint specifies the fixed service endpoint to call.
	Endpoint string

	// The client to use for outbound requests
	Client client.Client

	// The router for routes
	Router router.Router

	// A fib of routes service:address
	sync.RWMutex
	Routes map[string]map[uint64]table.Route

	// The channel to monitor watcher errors
	errChan chan error
}

// read client request and write to server
func readLoop(r server.Request, s client.Stream) error {
	// request to backend server
	req := s.Request()

	for {
		// get data from client
		//  no need to decode it
		body, err := r.Read()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		// get the header from client
		hdr := r.Header()
		msg := &codec.Message{
			Type:   codec.Request,
			Header: hdr,
			Body:   body,
		}

		// write the raw request
		err = req.Codec().Write(msg, nil)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

// toNodes returns a list of node addresses from given routes
func toNodes(routes map[uint64]table.Route) []string {
	var nodes []string
	for _, node := range routes {
		address := node.Address
		if len(node.Gateway) > 0 {
			address = node.Gateway
		}
		nodes = append(nodes, address)
	}
	return nodes
}

func (p *Proxy) getRoute(service string) ([]string, error) {
	// lookup the route cache first
	p.Lock()
	routes, ok := p.Routes[service]
	if ok {
		p.Unlock()
		return toNodes(routes), nil
	}
	p.Routes[service] = make(map[uint64]table.Route)
	p.Unlock()

	// if the router is broken return error
	if status := p.Router.Status(); status.Code == router.Error {
		return nil, status.Error
	}

	// lookup the routes in the router
	results, err := p.Router.Lookup(table.NewQuery(table.QueryService(service)))
	if err != nil {
		return nil, err
	}

	// update the proxy cache
	p.Lock()
	for _, route := range results {
		p.Routes[service][route.Hash()] = route
	}
	routes = p.Routes[service]
	p.Unlock()

	return toNodes(routes), nil
}

// manageRouteCache applies action on a given route to Proxy route cache
func (p *Proxy) manageRouteCache(route table.Route, action string) error {
	switch action {
	case "create", "update":
		if _, ok := p.Routes[route.Service]; !ok {
			p.Routes[route.Service] = make(map[uint64]table.Route)
		}
		p.Routes[route.Service][route.Hash()] = route
	case "delete":
		if _, ok := p.Routes[route.Service]; !ok {
			return fmt.Errorf("route not found")
		}
		delete(p.Routes[route.Service], route.Hash())
	default:
		return fmt.Errorf("unknown action: %s", action)
	}

	return nil
}

// watchRoutes watches service routes and updates proxy cache
func (p *Proxy) watchRoutes() {
	// this is safe to do as the only way watchRoutes returns is
	// when some error is written into error channel - we want to bail then
	defer close(p.errChan)

	// route watcher
	w, err := p.Router.Watch()
	if err != nil {
		p.errChan <- err
		return
	}

	for {
		event, err := w.Next()
		if err != nil {
			p.errChan <- err
			return
		}

		p.Lock()
		if err := p.manageRouteCache(event.Route, fmt.Sprintf("%s", event.Type)); err != nil {
			// TODO: should we bail here?
			continue
		}
		p.Unlock()
	}
}

// ServeRequest honours the server.Router interface
func (p *Proxy) ServeRequest(ctx context.Context, req server.Request, rsp server.Response) error {
	// service name
	service := req.Service()
	endpoint := req.Endpoint()
	var addresses []string

	// call a specific backend endpoint either by name or address
	if len(p.Endpoint) > 0 {
		// address:port
		if parts := strings.Split(p.Endpoint, ":"); len(parts) > 1 {
			addresses = []string{p.Endpoint}
		} else {
			// get route for endpoint from router
			addr, err := p.getRoute(p.Endpoint)
			if err != nil {
				return err
			}
			// set the address
			addresses = addr
			// set the name
			service = p.Endpoint
		}
	} else {
		// no endpoint was specified just lookup the route
		// get route for endpoint from router
		addr, err := p.getRoute(service)
		if err != nil {
			return err
		}
		addresses = addr
	}

	var opts []client.CallOption

	// set address if available
	if len(addresses) > 0 {
		opts = append(opts, client.WithAddress(addresses...))
	}

	// read initial request
	body, err := req.Read()
	if err != nil {
		return err
	}

	// create new request with raw bytes body
	creq := p.Client.NewRequest(service, endpoint, &bytes.Frame{body}, client.WithContentType(req.ContentType()))

	// create new stream
	stream, err := p.Client.Stream(ctx, creq, opts...)
	if err != nil {
		return err
	}
	defer stream.Close()

	// create client request read loop
	go readLoop(req, stream)

	// get raw response
	resp := stream.Response()

	// route watcher error
	var watchErr error

	// create server response write loop
	for {
		select {
		case err := <-p.errChan:
			if err != nil {
				watchErr = err
			}
			return watchErr
		default:
			// read backend response body
			body, err := resp.Read()
			if err == io.EOF {
				return nil
			} else if err != nil {
				return err
			}

			// read backend response header
			hdr := resp.Header()

			// write raw response header to client
			rsp.WriteHeader(hdr)

			// write raw response body to client
			err = rsp.Write(body)
			if err == io.EOF {
				return nil
			} else if err != nil {
				return err
			}
		}
	}

	return nil
}

// NewSingleHostProxy returns a proxy which sends requests to a single backend
func NewSingleHostProxy(endpoint string) *Proxy {
	return &Proxy{
		Options:  options.NewOptions(),
		Endpoint: endpoint,
	}
}

// NewProxy returns a new proxy which will route based on mucp headers
func NewProxy(opts ...options.Option) proxy.Proxy {
	p := new(Proxy)
	p.Options = options.NewOptions(opts...)
	p.Options.Init(options.WithString("mucp"))

	// get endpoint
	ep, ok := p.Options.Values().Get("proxy.endpoint")
	if ok {
		p.Endpoint = ep.(string)
	}

	// get client
	c, ok := p.Options.Values().Get("proxy.client")
	if ok {
		p.Client = c.(client.Client)
	}

	// set the default client
	if p.Client == nil {
		p.Client = client.DefaultClient
	}

	// get router
	r, ok := p.Options.Values().Get("proxy.router")
	if ok {
		p.Router = r.(router.Router)
	}

	// create default router and start it
	if p.Router == nil {
		p.Router = router.DefaultRouter
	}

	// routes cache
	p.Routes = make(map[string]map[uint64]table.Route)

	// watch router service routes
	p.errChan = make(chan error, 1)
	go p.watchRoutes()

	return p
}
<|endoftext|>


<commit_before>package prelude

const jsmapping = `
var $jsObjectPtr, $jsErrorPtr;

var $needsExternalization = function(t) {
  switch (t.kind) {
    case $kindBool:
    case $kindInt:
    case $kindInt8:
    case $kindInt16:
    case $kindInt32:
    case $kindUint:
    case $kindUint8:
    case $kindUint16:
    case $kindUint32:
    case $kindUintptr:
    case $kindFloat32:
    case $kindFloat64:
      return false;
    default:
      return t !== $jsObjectPtr;
  }
};

var $externalize = function(v, t) {
  if (t === $jsObjectPtr) {
    return v;
  }
  switch (t.kind) {
  case $kindBool:
  case $kindInt:
  case $kindInt8:
  case $kindInt16:
  case $kindInt32:
  case $kindUint:
  case $kindUint8:
  case $kindUint16:
  case $kindUint32:
  case $kindUintptr:
  case $kindFloat32:
  case $kindFloat64:
    return v;
  case $kindInt64:
  case $kindUint64:
    return $flatten64(v);
  case $kindArray:
    if ($needsExternalization(t.elem)) {
      return $mapArray(v, function(e) { return $externalize(e, t.elem); });
    }
    return v;
  case $kindFunc:
    return $externalizeFunction(v, t, false);
  case $kindInterface:
    if (v === $ifaceNil) {
      return null;
    }
    if (v.constructor === $jsObjectPtr) {
      return v.$val.object;
    }
    return $externalize(v.$val, v.constructor);
  case $kindMap:
    var m = {};
    var keys = $keys(v);
    for (var i = 0; i < keys.length; i++) {
      var entry = v[keys[i]];
      m[$externalize(entry.k, t.key)] = $externalize(entry.v, t.elem);
    }
    return m;
  case $kindPtr:
    if (v === t.nil) {
      return null;
    }
    return $externalize(v.$get(), t.elem);
  case $kindSlice:
    if ($needsExternalization(t.elem)) {
      return $mapArray($sliceToArray(v), function(e) { return $externalize(e, t.elem); });
    }
    return $sliceToArray(v);
  case $kindString:
    if (v.search(/^[\x00-\x7F]*$/) !== -1) {
      return v;
    }
    var s = "", r;
    for (var i = 0; i < v.length; i += r[1]) {
      r = $decodeRune(v, i);
      s += String.fromCharCode(r[0]);
    }
    return s;
  case $kindStruct:
    var timePkg = $packages["time"];
    if (timePkg && v.constructor === timePkg.Time.ptr) {
      var milli = $div64(v.UnixNano(), new $Int64(0, 1000000));
      return new Date($flatten64(milli));
    }

    var noJsObject = {};
    var searchJsObject = function(v, t) {
      if (t === $jsObjectPtr) {
        return v;
      }
      switch (t.kind) {
      case $kindPtr:
        if (v === t.nil) {
          return noJsObject;
        }
        return searchJsObject(v.$get(), t.elem);
      case $kindStruct:
        var f = t.fields[0];
        return searchJsObject(v[f.prop], f.typ);
      case $kindInterface:
        return searchJsObject(v.$val, v.constructor);
      default:
        return noJsObject;
      }
    };
    var o = searchJsObject(v, t);
    if (o !== noJsObject) {
      return o;
    }

    o = {};
    for (var i = 0; i < t.fields.length; i++) {
      var f = t.fields[i];
      if (f.pkg !== "") { /* not exported */
        continue;
      }
      o[f.name] = $externalize(v[f.prop], f.typ);
    }
    return o;
  }
  $panic(new $String("cannot externalize " + t.string));
};

var $externalizeFunction = function(v, t, passThis) {
  if (v === $throwNilPointerError) {
    return null;
  }
  if (v.$externalizeWrapper === undefined) {
    $checkForDeadlock = false;
    v.$externalizeWrapper = function() {
      var args = [];
      for (var i = 0; i < t.params.length; i++) {
        if (t.variadic && i === t.params.length - 1) {
          var vt = t.params[i].elem, varargs = [];
          for (var j = i; j < arguments.length; j++) {
            varargs.push($internalize(arguments[j], vt));
          }
          args.push(new (t.params[i])(varargs));
          break;
        }
        args.push($internalize(arguments[i], t.params[i]));
      }
      var canBlock = $curGoroutine.canBlock;
      $curGoroutine.canBlock = false;
      try {
        var result = v.apply(passThis ? this : undefined, args);
      } finally {
        $curGoroutine.canBlock = canBlock;
      }
      switch (t.results.length) {
      case 0:
        return;
      case 1:
        return $externalize(result, t.results[0]);
      default:
        for (var i = 0; i < t.results.length; i++) {
          result[i] = $externalize(result[i], t.results[i]);
        }
        return result;
      }
    };
  }
  return v.$externalizeWrapper;
};

var $internalize = function(v, t, recv) {
  if (t === $jsObjectPtr) {
    return v;
  }
  if (t === $jsObjectPtr.elem) {
    $panic(new $String("cannot internalize js.Object, use *js.Object instead"));
  }
  switch (t.kind) {
  case $kindBool:
    return !!v;
  case $kindInt:
    return parseInt(v);
  case $kindInt8:
    return parseInt(v) << 24 >> 24;
  case $kindInt16:
    return parseInt(v) << 16 >> 16;
  case $kindInt32:
    return parseInt(v) >> 0;
  case $kindUint:
    return parseInt(v);
  case $kindUint8:
    return parseInt(v) << 24 >>> 24;
  case $kindUint16:
    return parseInt(v) << 16 >>> 16;
  case $kindUint32:
  case $kindUintptr:
    return parseInt(v) >>> 0;
  case $kindInt64:
  case $kindUint64:
    return new t(0, v);
  case $kindFloat32:
  case $kindFloat64:
    return parseFloat(v);
  case $kindArray:
    if (v.length !== t.len) {
      $throwRuntimeError("got array with wrong size from JavaScript native");
    }
    return $mapArray(v, function(e) { return $internalize(e, t.elem); });
  case $kindFunc:
    return function() {
      var args = [];
      for (var i = 0; i < t.params.length; i++) {
        if (t.variadic && i === t.params.length - 1) {
          var vt = t.params[i].elem, varargs = arguments[i];
          for (var j = 0; j < varargs.$length; j++) {
            args.push($externalize(varargs.$array[varargs.$offset + j], vt));
          }
          break;
        }
        args.push($externalize(arguments[i], t.params[i]));
      }
      var result = v.apply(recv, args);
      switch (t.results.length) {
      case 0:
        return;
      case 1:
        return $internalize(result, t.results[0]);
      default:
        for (var i = 0; i < t.results.length; i++) {
          result[i] = $internalize(result[i], t.results[i]);
        }
        return result;
      }
    };
  case $kindInterface:
    if (t.methods.length !== 0) {
      $panic(new $String("cannot internalize " + t.string));
    }
    if (v === null) {
      return $ifaceNil;
    }
    if (v === undefined) {
      return new $jsObjectPtr(undefined);
    }
    switch (v.constructor) {
    case Int8Array:
      return new ($sliceType($Int8))(v);
    case Int16Array:
      return new ($sliceType($Int16))(v);
    case Int32Array:
      return new ($sliceType($Int))(v);
    case Uint8Array:
      return new ($sliceType($Uint8))(v);
    case Uint16Array:
      return new ($sliceType($Uint16))(v);
    case Uint32Array:
      return new ($sliceType($Uint))(v);
    case Float32Array:
      return new ($sliceType($Float32))(v);
    case Float64Array:
      return new ($sliceType($Float64))(v);
    case Array:
      return $internalize(v, $sliceType($emptyInterface));
    case Boolean:
      return new $Bool(!!v);
    case Date:
      var timePkg = $packages["time"];
      if (timePkg) {
        return new timePkg.Time(timePkg.Unix(new $Int64(0, 0), new $Int64(0, v.getTime() * 1000000)));
      } else {
        /* time package is not present, internalize as &js.Object{Date} so it can be externalized into original Date. */
        return new $jsObjectPtr(v);
      }
    case Function:
      var funcType = $funcType([$sliceType($emptyInterface)], [$jsObjectPtr], true);
      return new funcType($internalize(v, funcType));
    case Number:
      return new $Float64(parseFloat(v));
    case String:
      return new $String($internalize(v, $String));
    default:
      if ($global.Node && v instanceof $global.Node) {
        return new $jsObjectPtr(v);
      }
      var mapType = $mapType($String, $emptyInterface);
      return new mapType($internalize(v, mapType));
    }
  case $kindMap:
    var m = new $Map();
    var keys = $keys(v);
    for (var i = 0; i < keys.length; i++) {
      var k = $internalize(keys[i], t.key);
      m[t.key.keyFor(k)] = { k: k, v: $internalize(v[keys[i]], t.elem) };
    }
    return m;
  case $kindPtr:
    if (t.elem.kind === $kindStruct) {
      return $internalize(v, t.elem);
    }
  case $kindSlice:
    return new t($mapArray(v, function(e) { return $internalize(e, t.elem); }));
  case $kindString:
    v = String(v);
    if (v.search(/^[\x00-\x7F]*$/) !== -1) {
      return v;
    }
    var s = "";
    for (var i = 0; i < v.length; i++) {
      s += $encodeRune(v.charCodeAt(i));
    }
    return s;
  case $kindStruct:
    var noJsObject = {};
    var searchJsObject = function(t) {
      if (t === $jsObjectPtr) {
        return v;
      }
      if (t === $jsObjectPtr.elem) {
        $panic(new $String("cannot internalize js.Object, use *js.Object instead"));
      }
      switch (t.kind) {
      case $kindPtr:
        return searchJsObject(t.elem);
      case $kindStruct:
        var f = t.fields[0];
        var o = searchJsObject(f.typ);
        if (o !== noJsObject) {
          var n = new t.ptr();
          n[f.prop] = o;
          return n;
        }
        return noJsObject;
      default:
        return noJsObject;
      }
    };
    var o = searchJsObject(t);
    if (o !== noJsObject) {
      return o;
    }
  }
  $panic(new $String("cannot internalize " + t.string));
};
`
<commit_msg>Code style changes.<commit_after>package prelude

const jsmapping = `
var $jsObjectPtr, $jsErrorPtr;

var $needsExternalization = function(t) {
  switch (t.kind) {
    case $kindBool:
    case $kindInt:
    case $kindInt8:
    case $kindInt16:
    case $kindInt32:
    case $kindUint:
    case $kindUint8:
    case $kindUint16:
    case $kindUint32:
    case $kindUintptr:
    case $kindFloat32:
    case $kindFloat64:
      return false;
    default:
      return t !== $jsObjectPtr;
  }
};

var $externalize = function(v, t) {
  if (t === $jsObjectPtr) {
    return v;
  }
  switch (t.kind) {
  case $kindBool:
  case $kindInt:
  case $kindInt8:
  case $kindInt16:
  case $kindInt32:
  case $kindUint:
  case $kindUint8:
  case $kindUint16:
  case $kindUint32:
  case $kindUintptr:
  case $kindFloat32:
  case $kindFloat64:
    return v;
  case $kindInt64:
  case $kindUint64:
    return $flatten64(v);
  case $kindArray:
    if ($needsExternalization(t.elem)) {
      return $mapArray(v, function(e) { return $externalize(e, t.elem); });
    }
    return v;
  case $kindFunc:
    return $externalizeFunction(v, t, false);
  case $kindInterface:
    if (v === $ifaceNil) {
      return null;
    }
    if (v.constructor === $jsObjectPtr) {
      return v.$val.object;
    }
    return $externalize(v.$val, v.constructor);
  case $kindMap:
    var m = {};
    var keys = $keys(v);
    for (var i = 0; i < keys.length; i++) {
      var entry = v[keys[i]];
      m[$externalize(entry.k, t.key)] = $externalize(entry.v, t.elem);
    }
    return m;
  case $kindPtr:
    if (v === t.nil) {
      return null;
    }
    return $externalize(v.$get(), t.elem);
  case $kindSlice:
    if ($needsExternalization(t.elem)) {
      return $mapArray($sliceToArray(v), function(e) { return $externalize(e, t.elem); });
    }
    return $sliceToArray(v);
  case $kindString:
    if (v.search(/^[\x00-\x7F]*$/) !== -1) {
      return v;
    }
    var s = "", r;
    for (var i = 0; i < v.length; i += r[1]) {
      r = $decodeRune(v, i);
      s += String.fromCharCode(r[0]);
    }
    return s;
  case $kindStruct:
    var timePkg = $packages["time"];
    if (timePkg !== undefined && v.constructor === timePkg.Time.ptr) {
      var milli = $div64(v.UnixNano(), new $Int64(0, 1000000));
      return new Date($flatten64(milli));
    }

    var noJsObject = {};
    var searchJsObject = function(v, t) {
      if (t === $jsObjectPtr) {
        return v;
      }
      switch (t.kind) {
      case $kindPtr:
        if (v === t.nil) {
          return noJsObject;
        }
        return searchJsObject(v.$get(), t.elem);
      case $kindStruct:
        var f = t.fields[0];
        return searchJsObject(v[f.prop], f.typ);
      case $kindInterface:
        return searchJsObject(v.$val, v.constructor);
      default:
        return noJsObject;
      }
    };
    var o = searchJsObject(v, t);
    if (o !== noJsObject) {
      return o;
    }

    o = {};
    for (var i = 0; i < t.fields.length; i++) {
      var f = t.fields[i];
      if (f.pkg !== "") { /* not exported */
        continue;
      }
      o[f.name] = $externalize(v[f.prop], f.typ);
    }
    return o;
  }
  $panic(new $String("cannot externalize " + t.string));
};

var $externalizeFunction = function(v, t, passThis) {
  if (v === $throwNilPointerError) {
    return null;
  }
  if (v.$externalizeWrapper === undefined) {
    $checkForDeadlock = false;
    v.$externalizeWrapper = function() {
      var args = [];
      for (var i = 0; i < t.params.length; i++) {
        if (t.variadic && i === t.params.length - 1) {
          var vt = t.params[i].elem, varargs = [];
          for (var j = i; j < arguments.length; j++) {
            varargs.push($internalize(arguments[j], vt));
          }
          args.push(new (t.params[i])(varargs));
          break;
        }
        args.push($internalize(arguments[i], t.params[i]));
      }
      var canBlock = $curGoroutine.canBlock;
      $curGoroutine.canBlock = false;
      try {
        var result = v.apply(passThis ? this : undefined, args);
      } finally {
        $curGoroutine.canBlock = canBlock;
      }
      switch (t.results.length) {
      case 0:
        return;
      case 1:
        return $externalize(result, t.results[0]);
      default:
        for (var i = 0; i < t.results.length; i++) {
          result[i] = $externalize(result[i], t.results[i]);
        }
        return result;
      }
    };
  }
  return v.$externalizeWrapper;
};

var $internalize = function(v, t, recv) {
  if (t === $jsObjectPtr) {
    return v;
  }
  if (t === $jsObjectPtr.elem) {
    $panic(new $String("cannot internalize js.Object, use *js.Object instead"));
  }
  switch (t.kind) {
  case $kindBool:
    return !!v;
  case $kindInt:
    return parseInt(v);
  case $kindInt8:
    return parseInt(v) << 24 >> 24;
  case $kindInt16:
    return parseInt(v) << 16 >> 16;
  case $kindInt32:
    return parseInt(v) >> 0;
  case $kindUint:
    return parseInt(v);
  case $kindUint8:
    return parseInt(v) << 24 >>> 24;
  case $kindUint16:
    return parseInt(v) << 16 >>> 16;
  case $kindUint32:
  case $kindUintptr:
    return parseInt(v) >>> 0;
  case $kindInt64:
  case $kindUint64:
    return new t(0, v);
  case $kindFloat32:
  case $kindFloat64:
    return parseFloat(v);
  case $kindArray:
    if (v.length !== t.len) {
      $throwRuntimeError("got array with wrong size from JavaScript native");
    }
    return $mapArray(v, function(e) { return $internalize(e, t.elem); });
  case $kindFunc:
    return function() {
      var args = [];
      for (var i = 0; i < t.params.length; i++) {
        if (t.variadic && i === t.params.length - 1) {
          var vt = t.params[i].elem, varargs = arguments[i];
          for (var j = 0; j < varargs.$length; j++) {
            args.push($externalize(varargs.$array[varargs.$offset + j], vt));
          }
          break;
        }
        args.push($externalize(arguments[i], t.params[i]));
      }
      var result = v.apply(recv, args);
      switch (t.results.length) {
      case 0:
        return;
      case 1:
        return $internalize(result, t.results[0]);
      default:
        for (var i = 0; i < t.results.length; i++) {
          result[i] = $internalize(result[i], t.results[i]);
        }
        return result;
      }
    };
  case $kindInterface:
    if (t.methods.length !== 0) {
      $panic(new $String("cannot internalize " + t.string));
    }
    if (v === null) {
      return $ifaceNil;
    }
    if (v === undefined) {
      return new $jsObjectPtr(undefined);
    }
    switch (v.constructor) {
    case Int8Array:
      return new ($sliceType($Int8))(v);
    case Int16Array:
      return new ($sliceType($Int16))(v);
    case Int32Array:
      return new ($sliceType($Int))(v);
    case Uint8Array:
      return new ($sliceType($Uint8))(v);
    case Uint16Array:
      return new ($sliceType($Uint16))(v);
    case Uint32Array:
      return new ($sliceType($Uint))(v);
    case Float32Array:
      return new ($sliceType($Float32))(v);
    case Float64Array:
      return new ($sliceType($Float64))(v);
    case Array:
      return $internalize(v, $sliceType($emptyInterface));
    case Boolean:
      return new $Bool(!!v);
    case Date:
      var timePkg = $packages["time"];
      if (timePkg === undefined) {
        /* time package is not present, internalize as &js.Object{Date} so it can be externalized into original Date. */
        return new $jsObjectPtr(v);
      }
      return new timePkg.Time(timePkg.Unix(new $Int64(0, 0), new $Int64(0, v.getTime() * 1000000)));
    case Function:
      var funcType = $funcType([$sliceType($emptyInterface)], [$jsObjectPtr], true);
      return new funcType($internalize(v, funcType));
    case Number:
      return new $Float64(parseFloat(v));
    case String:
      return new $String($internalize(v, $String));
    default:
      if ($global.Node && v instanceof $global.Node) {
        return new $jsObjectPtr(v);
      }
      var mapType = $mapType($String, $emptyInterface);
      return new mapType($internalize(v, mapType));
    }
  case $kindMap:
    var m = new $Map();
    var keys = $keys(v);
    for (var i = 0; i < keys.length; i++) {
      var k = $internalize(keys[i], t.key);
      m[t.key.keyFor(k)] = { k: k, v: $internalize(v[keys[i]], t.elem) };
    }
    return m;
  case $kindPtr:
    if (t.elem.kind === $kindStruct) {
      return $internalize(v, t.elem);
    }
  case $kindSlice:
    return new t($mapArray(v, function(e) { return $internalize(e, t.elem); }));
  case $kindString:
    v = String(v);
    if (v.search(/^[\x00-\x7F]*$/) !== -1) {
      return v;
    }
    var s = "";
    for (var i = 0; i < v.length; i++) {
      s += $encodeRune(v.charCodeAt(i));
    }
    return s;
  case $kindStruct:
    var noJsObject = {};
    var searchJsObject = function(t) {
      if (t === $jsObjectPtr) {
        return v;
      }
      if (t === $jsObjectPtr.elem) {
        $panic(new $String("cannot internalize js.Object, use *js.Object instead"));
      }
      switch (t.kind) {
      case $kindPtr:
        return searchJsObject(t.elem);
      case $kindStruct:
        var f = t.fields[0];
        var o = searchJsObject(f.typ);
        if (o !== noJsObject) {
          var n = new t.ptr();
          n[f.prop] = o;
          return n;
        }
        return noJsObject;
      default:
        return noJsObject;
      }
    };
    var o = searchJsObject(t);
    if (o !== noJsObject) {
      return o;
    }
  }
  $panic(new $String("cannot internalize " + t.string));
};
`
<|endoftext|>


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

package subtitle

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"strings"
	"time"
)

func ReadSrt(data []byte) Book {
	var book Book
	var script Script

	b := bytes.NewBuffer(data)

	const (
		STATE_IDX = iota
		STATE_TS
		STATE_SCRIPT
	)

	state := STATE_IDX
	for {
		line, err := b.ReadString('\n')
		if err != nil {
			break
		}
		line = strings.TrimRight(line, "\r\n")
		/* log.Printf("line = '%s'", line) */

		switch state {
		case STATE_IDX:
			/* log.Println("STATE_IDX") */
			_, err := fmt.Sscanln(line, &script.Idx)
			if err != nil {
				log.Fatalf("failed to parse index! in \"%s\" : %s",
				line,  err)
			}
			state = STATE_TS

		case STATE_TS:
			/* log.Println("STATE_TS") */
			var sH, sM, sS, sMs int
			var eH, eM, eS, eMs int
			_, err := fmt.Sscanf(line,
				"%d:%d:%d,%d --> %d:%d:%d,%d",
				&sH, &sM, &sS, &sMs,
				&eH, &eM, &eS, &eMs)
			if err != nil {
				log.Fatalln("failed to parse timestamp!")
			}

			startMs := sMs + sS*1000 + sM*60*1000 + sH*60*60*1000
			script.Start = time.Duration(startMs) * time.Millisecond

			endMs := eMs + eS*1000 + eM*60*1000 + eH*60*60*1000
			script.End = time.Duration(endMs) * time.Millisecond

			script.Text = ""
			/* log.Println("script = ", script) */
			state = STATE_SCRIPT

		case STATE_SCRIPT:
			/* log.Println("STATE_SCRIPT") */
			if line == "" {
				/* log.Println("script = ", script) */
				book = append(book, script)
				state = STATE_IDX
			} else {
				if script.Text != "" {
					script.Text += "\n"
				}
				script.Text += line
			}
		}

	}
	/* log.Println("book = ", book) */
	return book
}

func ReadSrtFile(filename string) Book {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalln("faile to read file, ", filename)
	}
	return ReadSrt(data)
}
<commit_msg>subtitle: Skip UTF-8 BOM<commit_after>// Copyright 2013, Homin Lee. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package subtitle

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"strings"
	"time"
)

func ReadSrt(data []byte) Book {
	var book Book
	var script Script

	b := bytes.NewBuffer(data)

	const (
		STATE_IDX = iota
		STATE_TS
		STATE_SCRIPT
	)

	state := STATE_IDX
	for {
		line, err := b.ReadString('\n')
		if err != nil {
			break
		}
		line = strings.TrimRight(line, "\r\n")
		/* log.Printf("line = '%s'", line) */

		switch state {
		case STATE_IDX:
			/* log.Println("STATE_IDX") */
			_, err := fmt.Sscanln(line, &script.Idx)
			if err != nil {
				log.Fatalf("failed to parse index! in \"%s\" : %s",
				line,  err)
			}
			state = STATE_TS

		case STATE_TS:
			/* log.Println("STATE_TS") */
			var sH, sM, sS, sMs int
			var eH, eM, eS, eMs int
			_, err := fmt.Sscanf(line,
				"%d:%d:%d,%d --> %d:%d:%d,%d",
				&sH, &sM, &sS, &sMs,
				&eH, &eM, &eS, &eMs)
			if err != nil {
				log.Fatalln("failed to parse timestamp!")
			}

			startMs := sMs + sS*1000 + sM*60*1000 + sH*60*60*1000
			script.Start = time.Duration(startMs) * time.Millisecond

			endMs := eMs + eS*1000 + eM*60*1000 + eH*60*60*1000
			script.End = time.Duration(endMs) * time.Millisecond

			script.Text = ""
			/* log.Println("script = ", script) */
			state = STATE_SCRIPT

		case STATE_SCRIPT:
			/* log.Println("STATE_SCRIPT") */
			if line == "" {
				/* log.Println("script = ", script) */
				book = append(book, script)
				state = STATE_IDX
			} else {
				if script.Text != "" {
					script.Text += "\n"
				}
				script.Text += line
			}
		}

	}
	/* log.Println("book = ", book) */
	return book
}

func ReadSrtFile(filename string) Book {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalln("faile to read file, ", filename)
	}

	// skip UTF-8 BOM if exists
	if bytes.Equal(data[:3], []byte{0xEF, 0xBB, 0xBF}) {
		data = data[3:]
	}

	return ReadSrt(data)
}
<|endoftext|>


<commit_before>package state

import (
	"fmt"
	"launchpad.net/gozk/zookeeper"
	"launchpad.net/juju-core/state/presence"
	"strconv"
	"strings"
)

// RelationRole defines the role of a relation endpoint.
type RelationRole string

const (
	RoleProvider RelationRole = "provider"
	RoleRequirer RelationRole = "requirer"
	RolePeer     RelationRole = "peer"
)

// counterpartRole returns the RelationRole that this RelationRole
// can relate to.
// This should remain an internal method because the relation
// model does not guarantee that for every role there will
// necessarily exist a single counterpart role that is sensible
// for basing algorithms upon.
func (r RelationRole) counterpartRole() RelationRole {
	switch r {
	case RoleProvider:
		return RoleRequirer
	case RoleRequirer:
		return RoleProvider
	case RolePeer:
		return RolePeer
	}
	panic(fmt.Errorf("unknown RelationRole: %q", r))
}

// RelationScope describes the scope of a relation endpoint.
type RelationScope string

const (
	ScopeGlobal    RelationScope = "global"
	ScopeContainer RelationScope = "container"
)

// RelationEndpoint represents one endpoint of a relation.
type RelationEndpoint struct {
	ServiceName   string
	Interface     string
	RelationName  string
	RelationRole  RelationRole
	RelationScope RelationScope
}

// CanRelateTo returns whether a relation may be established between e and other.
func (e *RelationEndpoint) CanRelateTo(other *RelationEndpoint) bool {
	if e.Interface != other.Interface {
		return false
	}
	if e.RelationRole == RolePeer {
		// Peer relations do not currently work with multiple endpoints.
		return false
	}
	return e.RelationRole.counterpartRole() == other.RelationRole
}

// String returns the unique identifier of the relation endpoint.
func (e RelationEndpoint) String() string {
	return e.ServiceName + ":" + e.RelationName
}

// describeEndpoints returns a string describing the relation defined by
// endpoints, for use in various contexts (including error messages).
func describeEndpoints(endpoints []RelationEndpoint) string {
	names := []string{}
	for _, ep := range endpoints {
		names = append(names, ep.String())
	}
	return strings.Join(names, " ")
}

// Relation represents a relation between one or two service endpoints.
type Relation struct {
	st        *State
	key       string
	endpoints []RelationEndpoint
}

func (r *Relation) String() string {
	return fmt.Sprintf("relation %q", describeEndpoints(r.endpoints))
}

// Id returns the integer part of the internal relation key. This is
// exposed because the unit agent needs to expose a value derived from
// this (as JUJU_RELATION_ID) to allow relation hooks to differentiate
// between relations with different services.
func (r *Relation) Id() int {
	keyParts := strings.Split(r.key, "-")
	id, err := strconv.Atoi(keyParts[1])
	if err != nil {
		panic(fmt.Errorf("relation key %q in unknown format", r.key))
	}
	return id
}

// Endpoint returns the endpoint of the relation for the named service.
// If the service is not part of the relation, an error will be returned.
func (r *Relation) Endpoint(serviceName string) (RelationEndpoint, error) {
	for _, ep := range r.endpoints {
		if ep.ServiceName == serviceName {
			return ep, nil
		}
	}
	return RelationEndpoint{}, fmt.Errorf("service %q is not a member of %q", serviceName, r)
}

// RelatedEndpoints returns the endpoints of the relation r with which
// units of the named service will establish relations. If the service
// is not part of the relation r, an error will be returned.
func (r *Relation) RelatedEndpoints(serviceName string) ([]RelationEndpoint, error) {
	local, err := r.Endpoint(serviceName)
	if err != nil {
		return nil, err
	}
	role := local.RelationRole.counterpartRole()
	var eps []RelationEndpoint
	for _, ep := range r.endpoints {
		if ep.RelationRole == role {
			eps = append(eps, ep)
		}
	}
	if eps == nil {
		return nil, fmt.Errorf("no endpoints of %q relate to service %q", r, serviceName)
	}
	return eps, nil
}

// Unit returns a RelationUnit for the supplied unit.
func (r *Relation) Unit(u *Unit) (*RelationUnit, error) {
	ep, err := r.Endpoint(u.serviceName)
	if err != nil {
		return nil, err
	}
	path := []string{"/relations", r.key}
	if ep.RelationScope == ScopeContainer {
		container := u.principalKey
		if container == "" {
			container = u.key
		}
		path = append(path, container)
	}
	return &RelationUnit{
		st:       r.st,
		relation: r,
		unit:     u,
		role:     ep.RelationRole,
		scope:    unitScopePath(strings.Join(path, "/")),
	}, nil
}

// RelationUnit holds information about a single unit in a relation, and
// allows clients to conveniently access unit-specific functionality.
type RelationUnit struct {
	st       *State
	relation *Relation
	unit     *Unit
	role     RelationRole
	scope    unitScopePath
}

// Join joins the unit to the relation, such that other units watching the
// relation will observe its presence and changes to its settings.
func (ru *RelationUnit) Join() (*presence.Pinger, error) {
	if err := ru.scope.prepareJoin(ru.st.zk, ru.role); err != nil {
		return nil, err
	}
	// Private address should be set at agent startup.
	address, err := ru.unit.PrivateAddress()
	if err != nil {
		return nil, err
	}
	settings, err := readConfigNode(ru.st.zk, ru.scope.settingsPath(ru.unit.key))
	if err != nil {
		return nil, err
	}
	settings.Set("private-address", address)
	if _, err := settings.Write(); err != nil {
		return nil, err
	}
	presencePath := ru.scope.presencePath(ru.role, ru.unit.key)
	return presence.StartPinger(ru.st.zk, presencePath, agentPingerPeriod)
}

// Watch returns a RelationUnitsWatcher which notifies of relevant changes to
// other units in the relation.
func (ru *RelationUnit) Watch() *RelationUnitsWatcher {
	return newRelationUnitsWatcher(ru.scope, ru.role.counterpartRole(), ru.unit)
}

// Settings returns a ConfigNode which allows access to the unit's settings
// within the relation.
func (ru *RelationUnit) Settings() (*ConfigNode, error) {
	return readConfigNode(ru.st.zk, ru.scope.settingsPath(ru.unit.key))
}

// unitScopePath represents a common zookeeper path used by the set of units
// that can (transitively) affect one another within the context of a
// particular relation. For a globally-scoped relation, every unit is part of
// the same scope; but for a container-scoped relation, each unit is is a
// scope shared only with the units that share a container.
// Thus, unitScopePaths will take one of the following forms:
//
//   /relations/<relation-id>
//   /relations/<relation-id>/<container-id>
type unitScopePath string

// settingsPath returns the path to the relation unit settings node for the
// unit identified by unitKey, or to the relation scope settings node if
// unitKey is empty.
func (s unitScopePath) settingsPath(unitKey string) string {
	return s.subpath("settings", unitKey)
}

// presencePath returns the path to the relation unit presence node for a
// unit (identified by unitKey) of a service acting as role; or to the relation
// scope role node if unitKey is empty.
func (s unitScopePath) presencePath(role RelationRole, unitKey string) string {
	return s.subpath(string(role), unitKey)
}

// prepareJoin ensures that ZooKeeper nodes exist such that a unit of a
// service with the supplied role will be able to join the relation.
func (s unitScopePath) prepareJoin(zk *zookeeper.Conn, role RelationRole) error {
	paths := []string{string(s), s.settingsPath(""), s.presencePath(role, "")}
	for _, path := range paths {
		if _, err := zk.Create(path, "", 0, zkPermAll); err != nil {
			if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
				continue
			}
			return err
		}
	}
	return nil
}

// subpath returns an absolute ZooKeeper path to the node whose path relative
// to the scope node is composed of parts. Empty parts will be stripped.
func (s unitScopePath) subpath(parts ...string) string {
	path := string(s)
	for _, part := range parts {
		if part != "" {
			path = path + "/" + part
		}
	}
	return path
}
<commit_msg>address review points<commit_after>package state

import (
	"fmt"
	"launchpad.net/gozk/zookeeper"
	"launchpad.net/juju-core/state/presence"
	"strconv"
	"strings"
)

// RelationRole defines the role of a relation endpoint.
type RelationRole string

const (
	RoleProvider RelationRole = "provider"
	RoleRequirer RelationRole = "requirer"
	RolePeer     RelationRole = "peer"
)

// counterpartRole returns the RelationRole that this RelationRole
// can relate to.
// This should remain an internal method because the relation
// model does not guarantee that for every role there will
// necessarily exist a single counterpart role that is sensible
// for basing algorithms upon.
func (r RelationRole) counterpartRole() RelationRole {
	switch r {
	case RoleProvider:
		return RoleRequirer
	case RoleRequirer:
		return RoleProvider
	case RolePeer:
		return RolePeer
	}
	panic(fmt.Errorf("unknown RelationRole: %q", r))
}

// RelationScope describes the scope of a relation endpoint.
type RelationScope string

const (
	ScopeGlobal    RelationScope = "global"
	ScopeContainer RelationScope = "container"
)

// RelationEndpoint represents one endpoint of a relation.
type RelationEndpoint struct {
	ServiceName   string
	Interface     string
	RelationName  string
	RelationRole  RelationRole
	RelationScope RelationScope
}

// CanRelateTo returns whether a relation may be established between e and other.
func (e *RelationEndpoint) CanRelateTo(other *RelationEndpoint) bool {
	if e.Interface != other.Interface {
		return false
	}
	if e.RelationRole == RolePeer {
		// Peer relations do not currently work with multiple endpoints.
		return false
	}
	return e.RelationRole.counterpartRole() == other.RelationRole
}

// String returns the unique identifier of the relation endpoint.
func (e RelationEndpoint) String() string {
	return e.ServiceName + ":" + e.RelationName
}

// describeEndpoints returns a string describing the relation defined by
// endpoints, for use in various contexts (including error messages).
func describeEndpoints(endpoints []RelationEndpoint) string {
	names := []string{}
	for _, ep := range endpoints {
		names = append(names, ep.String())
	}
	return strings.Join(names, " ")
}

// Relation represents a relation between one or two service endpoints.
type Relation struct {
	st        *State
	key       string
	endpoints []RelationEndpoint
}

func (r *Relation) String() string {
	return fmt.Sprintf("relation %q", describeEndpoints(r.endpoints))
}

// Id returns the integer part of the internal relation key. This is
// exposed because the unit agent needs to expose a value derived from
// this (as JUJU_RELATION_ID) to allow relation hooks to differentiate
// between relations with different services.
func (r *Relation) Id() int {
	keyParts := strings.Split(r.key, "-")
	id, err := strconv.Atoi(keyParts[1])
	if err != nil {
		panic(fmt.Errorf("relation key %q in unknown format", r.key))
	}
	return id
}

// Endpoint returns the endpoint of the relation for the named service.
// If the service is not part of the relation, an error will be returned.
func (r *Relation) Endpoint(serviceName string) (RelationEndpoint, error) {
	for _, ep := range r.endpoints {
		if ep.ServiceName == serviceName {
			return ep, nil
		}
	}
	return RelationEndpoint{}, fmt.Errorf("service %q is not a member of %q", serviceName, r)
}

// RelatedEndpoints returns the endpoints of the relation r with which
// units of the named service will establish relations. If the service
// is not part of the relation r, an error will be returned.
func (r *Relation) RelatedEndpoints(serviceName string) ([]RelationEndpoint, error) {
	local, err := r.Endpoint(serviceName)
	if err != nil {
		return nil, err
	}
	role := local.RelationRole.counterpartRole()
	var eps []RelationEndpoint
	for _, ep := range r.endpoints {
		if ep.RelationRole == role {
			eps = append(eps, ep)
		}
	}
	if eps == nil {
		return nil, fmt.Errorf("no endpoints of %q relate to service %q", r, serviceName)
	}
	return eps, nil
}

// Unit returns a RelationUnit for the supplied unit.
func (r *Relation) Unit(u *Unit) (*RelationUnit, error) {
	ep, err := r.Endpoint(u.serviceName)
	if err != nil {
		return nil, err
	}
	path := []string{"/relations", r.key}
	if ep.RelationScope == ScopeContainer {
		container := u.principalKey
		if container == "" {
			container = u.key
		}
		path = append(path, container)
	}
	return &RelationUnit{
		st:       r.st,
		relation: r,
		unit:     u,
		role:     ep.RelationRole,
		scope:    unitScopePath(strings.Join(path, "/")),
	}, nil
}

// RelationUnit holds information about a single unit in a relation, and
// allows clients to conveniently access unit-specific functionality.
type RelationUnit struct {
	st       *State
	relation *Relation
	unit     *Unit
	role     RelationRole
	scope    unitScopePath
}

// Join joins the unit to the relation, such that other units watching the
// relation will observe its presence and changes to its settings.
func (ru *RelationUnit) Join() (p *presence.Pinger, err error) {
	errorContextf(&err, "can't join unit %q to relation %q", ru.unit, ru.relation)
	if err = ru.scope.prepareJoin(ru.st.zk, ru.role); err != nil {
		return nil, err
	}
	// Private address should be set at agent startup.
	address, err := ru.unit.PrivateAddress()
	if err != nil {
		return nil, err
	}
	settings, err := readConfigNode(ru.st.zk, ru.scope.settingsPath(ru.unit.key))
	if err != nil {
		return nil, err
	}
	settings.Set("private-address", address)
	if _, err = settings.Write(); err != nil {
		return nil, err
	}
	presencePath := ru.scope.presencePath(ru.role, ru.unit.key)
	return presence.StartPinger(ru.st.zk, presencePath, agentPingerPeriod)
}

// Watch returns a watcher that notifies when any other unit in
// the relation joins, departs, or has its settings changed.
func (ru *RelationUnit) Watch() *RelationUnitsWatcher {
	return newRelationUnitsWatcher(ru.scope, ru.role.counterpartRole(), ru.unit)
}

// Settings returns a ConfigNode which allows access to the unit's settings
// within the relation.
func (ru *RelationUnit) Settings() (*ConfigNode, error) {
	return readConfigNode(ru.st.zk, ru.scope.settingsPath(ru.unit.key))
}

// unitScopePath represents a common zookeeper path used by the set of units
// that can (transitively) affect one another within the context of a
// particular relation. For a globally-scoped relation, every unit is part of
// the same scope; but for a container-scoped relation, each unit is is a
// scope shared only with the units that share a container.
// Thus, unitScopePaths will take one of the following forms:
//
//   /relations/<relation-id>
//   /relations/<relation-id>/<container-id>
type unitScopePath string

// settingsPath returns the path to the relation unit settings node for the
// unit identified by unitKey, or to the relation scope settings node if
// unitKey is empty.
func (s unitScopePath) settingsPath(unitKey string) string {
	return s.subpath("settings", unitKey)
}

// presencePath returns the path to the relation unit presence node for a
// unit (identified by unitKey) of a service acting as role; or to the relation
// scope role node if unitKey is empty.
func (s unitScopePath) presencePath(role RelationRole, unitKey string) string {
	return s.subpath(string(role), unitKey)
}

// prepareJoin ensures that ZooKeeper nodes exist such that a unit of a
// service with the supplied role will be able to join the relation.
func (s unitScopePath) prepareJoin(zk *zookeeper.Conn, role RelationRole) error {
	paths := []string{string(s), s.settingsPath(""), s.presencePath(role, "")}
	for _, path := range paths {
		if _, err := zk.Create(path, "", 0, zkPermAll); err != nil {
			if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
				continue
			}
			return err
		}
	}
	return nil
}

// subpath returns an absolute ZooKeeper path to the node whose path relative
// to the scope node is composed of parts. Empty parts will be stripped.
func (s unitScopePath) subpath(parts ...string) string {
	path := string(s)
	for _, part := range parts {
		if part != "" {
			path = path + "/" + part
		}
	}
	return path
}
<|endoftext|>


<commit_before>package tdb

import (
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"github.com/Knetic/govaluate"
	"github.com/davecgh/go-spew/spew"
	"github.com/getlantern/bytemap"
	. "github.com/getlantern/tdb/expr"
	"github.com/getlantern/tdb/sql"

	"github.com/stretchr/testify/assert"
	"testing"
)

func TestRoundTime(t *testing.T) {
	ts := time.Date(2015, 5, 6, 7, 8, 9, 10, time.UTC)
	rounded := roundTime(ts, time.Second)
	expected := time.Date(2015, 5, 6, 7, 8, 9, 0, time.UTC)
	assert.Equal(t, expected, rounded)
}

func TestIntegration(t *testing.T) {
	epoch := time.Date(2015, time.January, 1, 0, 0, 0, 0, time.UTC)

	tmpDir, err := ioutil.TempDir("", "tdbtest")
	if !assert.NoError(t, err, "Unable to create temp directory") {
		return
	}
	defer os.RemoveAll(tmpDir)

	tmpFile, err := ioutil.TempFile("", "tdbschema")
	if !assert.NoError(t, err, "Unable to create temp file") {
		return
	}
	tmpFile.Close()

	resolution := time.Millisecond
	hotPeriod := 2 * resolution

	schemaA := `
Test_a:
  hotperiod: 2ms
  retentionperiod: 200ms
  sql: >
    SELECT
      SUM(i) AS i,
      SUM(ii) AS ii,
      AVG(i * ii) AS iii
    FROM inbound
    WHERE r = 'A'
    GROUP BY period(1ms)
`
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaA), 0644)
	if !assert.NoError(t, err, "Unable to write schemaA") {
		return
	}

	db, err := NewDB(&DBOpts{
		Dir:        tmpDir,
		BatchSize:  1,
		SchemaFile: tmpFile.Name(),
	})
	if !assert.NoError(t, err, "Unable to create DB") {
		return
	}

	schemaB := schemaA + `
view_a:
  hotperiod: 2ms
  retentionperiod: 200ms
  sql: >
    SELECT
      SUM(i) AS i,
      SUM(ii) AS ii,
      AVG(i * ii) AS iii
    FROM inbound
    WHERE r = 'A'
    GROUP BY u, period(1ms)`
	log.Debug("Writing schemaB")
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaB), 0644)
	if !assert.NoError(t, err, "Unable to write schemaB") {
		return
	}

	viewCreated := false
	for i := 0; i < 50; i++ {
		time.Sleep(100 * time.Millisecond)
		if db.getTable("view_a") != nil {
			viewCreated = true
		}
	}
	if !assert.True(t, viewCreated, "View failed to create within 5 seconds") {
		return
	}

	now := epoch
	advance := func(d time.Duration) {
		time.Sleep(250 * time.Millisecond)
		now = now.Add(d)
		db.getTable("test_a").clock.Advance(now)
		db.getTable("view_a").clock.Advance(now)
		time.Sleep(250 * time.Millisecond)
		for _, table := range []string{"test_a", "view_a"} {
			stats := db.TableStats(table)
			log.Debugf("%v (%v)\tFiltered Out Points: %d\tInserted Points: %d\tDropped Points: %d\tHot Keys: %d\tArchived Buckets: %d", table, db.Now(table).In(time.UTC), stats.FilteredPoints, stats.InsertedPoints, stats.DroppedPoints, stats.HotKeys, stats.ArchivedBuckets)
		}
	}

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  1,
			"ii": 2,
		},
	})

	// This should get excluded by the filter
	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "B",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  1,
			"ii": 2,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  10,
			"ii": 20,
		},
	})

	advance(resolution)

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  111,
			"ii": 222,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 2,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  31,
			"ii": 42,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 2,
			"b": true,
		},
		Vals: map[string]float64{
			"i":  30000,
			"ii": 40000,
		},
	})

	advance(hotPeriod * 10)

	query := func(table string, from time.Time, to time.Time, dim string, field string) (map[int][]float64, error) {
		filter, queryErr := govaluate.NewEvaluableExpression("!b")
		if queryErr != nil {
			return nil, queryErr
		}
		fromOffset := from.Sub(now)
		toOffset := to.Sub(now)
		result := make(map[int][]float64, 0)
		_, err = db.runQuery(&query{
			table:       table,
			fields:      []string{field},
			asOfOffset:  fromOffset,
			untilOffset: toOffset,
			filter:      filter,
			onValues: func(keybytes bytemap.ByteMap, resultField string, vals []float64) {
				key := keybytes.AsMap()
				log.Debugf("%v : %v : %v", key, field, vals)
				if field == resultField {
					result[key[dim].(int)] = vals
				}
			},
		})
		return result, err
	}

	result, err := query("Test_A", epoch.Add(-1*resolution), epoch, "u", "i")
	if assert.NoError(t, err, "Unable to run query") {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{11}, result[1])
		}
	}

	result, err = query("Test_A", epoch.Add(-1*resolution), epoch, "u", "iii")
	if assert.NoError(t, err, "Unable to run query") {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{101}, result[1])
		}
	}

	result, err = query("Test_A", epoch.Add(-2*resolution), epoch.Add(resolution*2), "u", "ii")
	if assert.NoError(t, err, "Unable to run query") {
		if assert.Len(t, result, 2) {
			assert.Equal(t, []float64{222, 22, 0, 0}, result[1])
			assert.Equal(t, []float64{42, 0, 0, 0}, result[2])
		}
	}

	result, err = query("Test_A", epoch.Add(-2*resolution), epoch.Add(resolution*2), "u", "ii")
	log.Debug(result)
	if assert.NoError(t, err, "Unable to run query") {
		if assert.Len(t, result, 2) {
			assert.Equal(t, []float64{222, 22, 0, 0}, result[1])
			assert.Equal(t, []float64{42, 0, 0, 0}, result[2])
		}
	}

	testAggregateQuery(t, db, now, epoch, resolution)
}

func testAggregateQuery(t *testing.T, db *DB, now time.Time, epoch time.Time, resolution time.Duration) {
	scalingFactor := 5

	aq, err := db.SQLQuery(fmt.Sprintf(`
SELECT
	SUM(ii) AS sum_ii,
	COUNT(ii) AS count_ii,
	AVG(ii) * 2 AS avg_ii,
	MIN(ii) AS min_ii
FROM test_a
ASOF '%v' UNTIL '%v'
WHERE b != true
GROUP BY r, period('%v')
HAVING SUM(sum_ii) = 286
ORDER BY AVG(avg_ii) DESC
`, epoch.Add(-1*resolution).Sub(now), epoch.Add(3*resolution).Sub(now), resolution*time.Duration(scalingFactor)))
	if !assert.NoError(t, err, "Unable to create SQL query") {
		return
	}

	result, err := aq.Run()
	if !assert.NoError(t, err, "Unable to run query") {
		return
	}

	log.Debugf("%v -> %v", result.AsOf, result.Until)
	if !assert.Len(t, result.Entries, 1, "Wrong number of entries, perhaps HAVING isn't working") {
		return
	}
	entry := result.Entries[0]
	if !assert.EqualValues(t, "A", entry.Dims["r"], "Wrong dim, result may be sorted incorrectly") {
		return
	}
	if !assert.Len(t, entry.Fields["avg_ii"], 1, "Wrong number of periods, bucketing may not be working correctly") {
		return
	}
	avg := float64(286) / float64(scalingFactor)
	assert.EqualValues(t, 286, entry.Fields["sum_ii"][0].Get(), "Wrong derived value, bucketing may not be working correctly")
	if !assert.EqualValues(t, avg, entry.Fields["avg_ii"][0].Get(), "Wrong derived value, bucketing may not be working correctly") {
		t.Log(spew.Sprint(entry.Fields["avg_ii"][0]))
	}
	assert.EqualValues(t, 0, entry.Fields["min_ii"][0].Get(), "Wrong derived value, bucketing may not be working correctly")
	fields := make([]string, 0, len(result.Fields))
	for _, field := range result.Fields {
		fields = append(fields, field.Name)
	}
	assert.Equal(t, []string{"sum_ii", "count_ii", "avg_ii", "min_ii"}, fields)

	// Test defaults
	aq = db.Query(&sql.Query{
		From:       "test_a",
		Fields:     []sql.Field{sql.Field{Expr: SUM("ii"), Name: "sum_ii"}},
		AsOfOffset: epoch.Add(-1 * resolution).Sub(now),
	})

	result, err = aq.Run()
	if assert.NoError(t, err, "Unable to run query with defaults") {
		assert.Equal(t, []string{"b", "r", "u"}, result.GroupBy)
		assert.NotNil(t, result.Until)
	}
}
<commit_msg>Added sleep to make test pass repetably<commit_after>package tdb

import (
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"github.com/Knetic/govaluate"
	"github.com/davecgh/go-spew/spew"
	"github.com/getlantern/bytemap"
	. "github.com/getlantern/tdb/expr"
	"github.com/getlantern/tdb/sql"

	"github.com/stretchr/testify/assert"
	"testing"
)

func TestRoundTime(t *testing.T) {
	ts := time.Date(2015, 5, 6, 7, 8, 9, 10, time.UTC)
	rounded := roundTime(ts, time.Second)
	expected := time.Date(2015, 5, 6, 7, 8, 9, 0, time.UTC)
	assert.Equal(t, expected, rounded)
}

func TestIntegration(t *testing.T) {
	epoch := time.Date(2015, time.January, 1, 0, 0, 0, 0, time.UTC)

	tmpDir, err := ioutil.TempDir("", "tdbtest")
	if !assert.NoError(t, err, "Unable to create temp directory") {
		return
	}
	defer os.RemoveAll(tmpDir)

	tmpFile, err := ioutil.TempFile("", "tdbschema")
	if !assert.NoError(t, err, "Unable to create temp file") {
		return
	}
	tmpFile.Close()

	resolution := time.Millisecond
	hotPeriod := 2 * resolution

	schemaA := `
Test_a:
  hotperiod: 2ms
  retentionperiod: 200ms
  sql: >
    SELECT
      SUM(i) AS i,
      SUM(ii) AS ii,
      AVG(i * ii) AS iii
    FROM inbound
    WHERE r = 'A'
    GROUP BY period(1ms)
`
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaA), 0644)
	if !assert.NoError(t, err, "Unable to write schemaA") {
		return
	}

	db, err := NewDB(&DBOpts{
		Dir:        tmpDir,
		BatchSize:  1,
		SchemaFile: tmpFile.Name(),
	})
	if !assert.NoError(t, err, "Unable to create DB") {
		return
	}

	schemaB := schemaA + `
view_a:
  hotperiod: 2ms
  retentionperiod: 200ms
  sql: >
    SELECT
      SUM(i) AS i,
      SUM(ii) AS ii,
      AVG(i * ii) AS iii
    FROM inbound
    WHERE r = 'A'
    GROUP BY u, period(1ms)`
	log.Debug("Writing schemaB")
	err = ioutil.WriteFile(tmpFile.Name(), []byte(schemaB), 0644)
	if !assert.NoError(t, err, "Unable to write schemaB") {
		return
	}

	viewCreated := false
	for i := 0; i < 50; i++ {
		time.Sleep(100 * time.Millisecond)
		if db.getTable("view_a") != nil {
			viewCreated = true
		}
	}
	if !assert.True(t, viewCreated, "View failed to create within 5 seconds") {
		return
	}

	now := epoch
	advance := func(d time.Duration) {
		time.Sleep(250 * time.Millisecond)
		now = now.Add(d)
		db.getTable("test_a").clock.Advance(now)
		db.getTable("view_a").clock.Advance(now)
		time.Sleep(250 * time.Millisecond)
		for _, table := range []string{"test_a", "view_a"} {
			stats := db.TableStats(table)
			log.Debugf("%v (%v)\tFiltered Out Points: %d\tInserted Points: %d\tDropped Points: %d\tHot Keys: %d\tArchived Buckets: %d", table, db.Now(table).In(time.UTC), stats.FilteredPoints, stats.InsertedPoints, stats.DroppedPoints, stats.HotKeys, stats.ArchivedBuckets)
		}
	}

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  1,
			"ii": 2,
		},
	})

	// This should get excluded by the filter
	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "B",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  1,
			"ii": 2,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  10,
			"ii": 20,
		},
	})

	advance(resolution)

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 1,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  111,
			"ii": 222,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 2,
			"b": false,
		},
		Vals: map[string]float64{
			"i":  31,
			"ii": 42,
		},
	})

	db.Insert("inbound", &Point{
		Ts: now,
		Dims: map[string]interface{}{
			"r": "A",
			"u": 2,
			"b": true,
		},
		Vals: map[string]float64{
			"i":  30000,
			"ii": 40000,
		},
	})

	advance(hotPeriod * 10)

	query := func(table string, from time.Time, to time.Time, dim string, field string) (map[int][]float64, error) {
		filter, queryErr := govaluate.NewEvaluableExpression("!b")
		if queryErr != nil {
			return nil, queryErr
		}
		fromOffset := from.Sub(now)
		toOffset := to.Sub(now)
		result := make(map[int][]float64, 0)
		_, err = db.runQuery(&query{
			table:       table,
			fields:      []string{field},
			asOfOffset:  fromOffset,
			untilOffset: toOffset,
			filter:      filter,
			onValues: func(keybytes bytemap.ByteMap, resultField string, vals []float64) {
				key := keybytes.AsMap()
				log.Debugf("%v : %v : %v", key, field, vals)
				if field == resultField {
					result[key[dim].(int)] = vals
				}
			},
		})
		return result, err
	}

	// Give archiver time to catch up
	time.Sleep(1 * time.Second)

	result, err := query("Test_A", epoch.Add(-1*resolution), epoch, "u", "i")
	if assert.NoError(t, err, "Unable to run query") {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{11}, result[1])
		}
	}

	result, err = query("Test_A", epoch.Add(-1*resolution), epoch, "u", "iii")
	if assert.NoError(t, err, "Unable to run query") {
		if assert.Len(t, result, 1) {
			assert.Equal(t, []float64{101}, result[1])
		}
	}

	result, err = query("Test_A", epoch.Add(-2*resolution), epoch.Add(resolution*2), "u", "ii")
	if assert.NoError(t, err, "Unable to run query") {
		if assert.Len(t, result, 2) {
			assert.Equal(t, []float64{222, 22, 0, 0}, result[1])
			assert.Equal(t, []float64{42, 0, 0, 0}, result[2])
		}
	}

	result, err = query("Test_A", epoch.Add(-2*resolution), epoch.Add(resolution*2), "u", "ii")
	log.Debug(result)
	if assert.NoError(t, err, "Unable to run query") {
		if assert.Len(t, result, 2) {
			assert.Equal(t, []float64{222, 22, 0, 0}, result[1])
			assert.Equal(t, []float64{42, 0, 0, 0}, result[2])
		}
	}

	testAggregateQuery(t, db, now, epoch, resolution)
}

func testAggregateQuery(t *testing.T, db *DB, now time.Time, epoch time.Time, resolution time.Duration) {
	scalingFactor := 5

	aq, err := db.SQLQuery(fmt.Sprintf(`
SELECT
	SUM(ii) AS sum_ii,
	COUNT(ii) AS count_ii,
	AVG(ii) * 2 AS avg_ii,
	MIN(ii) AS min_ii
FROM test_a
ASOF '%v' UNTIL '%v'
WHERE b != true
GROUP BY r, period('%v')
HAVING SUM(sum_ii) = 286
ORDER BY AVG(avg_ii) DESC
`, epoch.Add(-1*resolution).Sub(now), epoch.Add(3*resolution).Sub(now), resolution*time.Duration(scalingFactor)))
	if !assert.NoError(t, err, "Unable to create SQL query") {
		return
	}

	result, err := aq.Run()
	if !assert.NoError(t, err, "Unable to run query") {
		return
	}

	log.Debugf("%v -> %v", result.AsOf, result.Until)
	if !assert.Len(t, result.Entries, 1, "Wrong number of entries, perhaps HAVING isn't working") {
		return
	}
	entry := result.Entries[0]
	if !assert.EqualValues(t, "A", entry.Dims["r"], "Wrong dim, result may be sorted incorrectly") {
		return
	}
	if !assert.Len(t, entry.Fields["avg_ii"], 1, "Wrong number of periods, bucketing may not be working correctly") {
		return
	}
	avg := float64(286) / float64(scalingFactor)
	assert.EqualValues(t, 286, entry.Fields["sum_ii"][0].Get(), "Wrong derived value, bucketing may not be working correctly")
	if !assert.EqualValues(t, avg, entry.Fields["avg_ii"][0].Get(), "Wrong derived value, bucketing may not be working correctly") {
		t.Log(spew.Sprint(entry.Fields["avg_ii"][0]))
	}
	assert.EqualValues(t, 0, entry.Fields["min_ii"][0].Get(), "Wrong derived value, bucketing may not be working correctly")
	fields := make([]string, 0, len(result.Fields))
	for _, field := range result.Fields {
		fields = append(fields, field.Name)
	}
	assert.Equal(t, []string{"sum_ii", "count_ii", "avg_ii", "min_ii"}, fields)

	// Test defaults
	aq = db.Query(&sql.Query{
		From:       "test_a",
		Fields:     []sql.Field{sql.Field{Expr: SUM("ii"), Name: "sum_ii"}},
		AsOfOffset: epoch.Add(-1 * resolution).Sub(now),
	})

	result, err = aq.Run()
	if assert.NoError(t, err, "Unable to run query with defaults") {
		assert.Equal(t, []string{"b", "r", "u"}, result.GroupBy)
		assert.NotNil(t, result.Until)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"
)

// Console flags
var (
	listen                = flag.String("l", ":8888", "port to accept requests")
	targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")
	altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")
	debug                 = flag.Bool("debug", false, "more logging, showing ignored output")
	productionTimeout     = flag.Int("a.timeout", 2500, "timeout in milliseconds for production traffic")
	alternateTimeout      = flag.Int("b.timeout", 1000, "timeout in milliseconds for alternate site traffic")
	productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic")
	alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")
	percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")
	tlsPrivateKey         = flag.String("key.file", "", "path to the TLS private key file")
	tlsCertificate        = flag.String("cert.file", "", "path to the TLS certificate file")
)

// Sets the request URL.
//
// This turns a inbound request (a request without URL) into an outbound request.
func setRequestTarget(request *http.Request, target *string) {
	URL, err := url.Parse("http://" + *target + request.URL.String())
	if err != nil {
		log.Println(err)
	}
	request.URL = URL
}

// Sends a request and returns the response.
func handleRequest(request *http.Request, timeout time.Duration) (*http.Response) {
	transport := &http.Transport{
		// NOTE(girone): DialTLS is not needed here, because the teeproxy works
		// as an SSL terminator.
		Dial: (&net.Dialer{  // go1.8 deprecated: Use DialContext instead
			Timeout: timeout,
			KeepAlive: 10 * timeout,
		}).Dial,
		// Always close connections to the alternative and production servers.
		DisableKeepAlives: true,
		//IdleConnTimeout: timeout,  // go1.8
		TLSHandshakeTimeout: timeout,
		ResponseHeaderTimeout: timeout,
		ExpectContinueTimeout: timeout,
	}
	// Do not use http.Client here, because it's higher level and processes
	// redirects internally, which is not what we want.
	//client := &http.Client{
	//	Timeout: timeout,
	//	Transport: transport,
	//}
	//response, err := client.Do(request)
	response, err := transport.RoundTrip(request)
	if err != nil {
		log.Println("Request failed:", err)
	}
	return response
}

// handler contains the address of the main Target and the one for the Alternative target
type handler struct {
	Target      string
	Alternative string
	Randomizer  rand.Rand
}

// ServeHTTP duplicates the incoming request (req) and does the request to the
// Target and the Alternate target discading the Alternate response
func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var productionRequest, alternativeRequest *http.Request
	updateForwardedHeaders(req)
	if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {
		alternativeRequest, productionRequest = DuplicateRequest(req)
		go func() {
			defer func() {
				if r := recover(); r != nil && *debug {
					log.Println("Recovered in ServeHTTP(alternate request) from:", r)
				}
			}()

			setRequestTarget(alternativeRequest, altTarget)

			if *alternateHostRewrite {
				alternativeRequest.Host = h.Alternative
			}

			timeout := time.Duration(*alternateTimeout) * time.Millisecond
			// This keeps responses from the alternative target away from the outside world.
			_ = handleRequest(alternativeRequest, timeout)
		}()
	} else {
		productionRequest = req
	}
	defer func() {
		if r := recover(); r != nil && *debug {
			log.Println("Recovered in ServeHTTP(production request) from:", r)
		}
	}()

	setRequestTarget(productionRequest, targetProduction)

	if *productionHostRewrite {
		productionRequest.Host = h.Target
	}

	timeout := time.Duration(*productionTimeout) * time.Millisecond
	resp := handleRequest(productionRequest, timeout)

	if resp != nil {
		defer resp.Body.Close()

		// Forward response headers.
		for k, v := range resp.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(resp.StatusCode)

		// Forward response body.
		body, _ := ioutil.ReadAll(resp.Body)
		w.Write(body)
	}
}

func main() {
	flag.Parse()

	log.Printf("Starting teeproxy at %s sending to A: %s and B: %s",
	           *listen, *targetProduction, *altTarget)

	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error

	var listener net.Listener

	if len(*tlsPrivateKey) > 0 {
		cer, err := tls.LoadX509KeyPair(*tlsCertificate, *tlsPrivateKey)
		if err != nil {
			log.Fatalf("Failed to load certficate: %s and private key: %s", *tlsCertificate, *tlsPrivateKey)
		}

		config := &tls.Config{Certificates: []tls.Certificate{cer}}
		listener, err = tls.Listen("tcp", *listen, config)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	} else {
		listener, err = net.Listen("tcp", *listen)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	}

	h := handler{
		Target:      *targetProduction,
		Alternative: *altTarget,
		Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),
	}

	server := &http.Server{
		Handler: h,
	}
	server.Serve(listener)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {
	b1 := new(bytes.Buffer)
	b2 := new(bytes.Buffer)
	w := io.MultiWriter(b1, b2)
	io.Copy(w, request.Body)
	defer request.Body.Close()
	request1 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b1},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	request2 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b2},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	return
}

func updateForwardedHeaders(request *http.Request) {
	positionOfColon := strings.LastIndex(request.RemoteAddr, ":")
	var remoteIP string
	if positionOfColon != -1 {
		remoteIP = request.RemoteAddr[:positionOfColon]
	} else {
		Logger.Printf("The default format of request.RemoteAddr should be IP:Port but was %s\n", remoteIP)
		remoteIP = request.RemoteAddr
	}
	insertOrExtendForwardedHeader(request, remoteIP)
	insertOrExtendXFFHeader(request, remoteIP)
}

const XFF_HEADER = "X-Forwarded-For"

func insertOrExtendXFFHeader(request *http.Request, remoteIP string) {
	header := request.Header.Get(XFF_HEADER)
	if header != "" {
		// extend
		request.Header.Set(XFF_HEADER, header + ", " + remoteIP)
	} else {
		// insert
		request.Header.Set(XFF_HEADER, remoteIP)
	}
}

const FORWARDED_HEADER = "Forwarded"

// Implementation according to rfc7239
func insertOrExtendForwardedHeader(request *http.Request, remoteIP string) {
	extension := "for=" + remoteIP
	header := request.Header.Get(FORWARDED_HEADER)
	if header != "" {
		// extend
		request.Header.Set(FORWARDED_HEADER, header + ", " + extension)
	} else {
		// insert
		request.Header.Set(FORWARDED_HEADER, extension)
	}
}
<commit_msg>Make forwarding of client IP optional<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"
)

// Console flags
var (
	listen                = flag.String("l", ":8888", "port to accept requests")
	targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")
	altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")
	debug                 = flag.Bool("debug", false, "more logging, showing ignored output")
	productionTimeout     = flag.Int("a.timeout", 2500, "timeout in milliseconds for production traffic")
	alternateTimeout      = flag.Int("b.timeout", 1000, "timeout in milliseconds for alternate site traffic")
	productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic")
	alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")
	percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")
	tlsPrivateKey         = flag.String("key.file", "", "path to the TLS private key file")
	tlsCertificate        = flag.String("cert.file", "", "path to the TLS certificate file")
	forwardClientIP       = flag.Bool("forward-client-ip", false, "enable forwarding of the client IP to the backend using the 'X-Forwarded-For' and 'Forwarded' headers")
)

// Sets the request URL.
//
// This turns a inbound request (a request without URL) into an outbound request.
func setRequestTarget(request *http.Request, target *string) {
	URL, err := url.Parse("http://" + *target + request.URL.String())
	if err != nil {
		log.Println(err)
	}
	request.URL = URL
}

// Sends a request and returns the response.
func handleRequest(request *http.Request, timeout time.Duration) (*http.Response) {
	transport := &http.Transport{
		// NOTE(girone): DialTLS is not needed here, because the teeproxy works
		// as an SSL terminator.
		Dial: (&net.Dialer{  // go1.8 deprecated: Use DialContext instead
			Timeout: timeout,
			KeepAlive: 10 * timeout,
		}).Dial,
		// Always close connections to the alternative and production servers.
		DisableKeepAlives: true,
		//IdleConnTimeout: timeout,  // go1.8
		TLSHandshakeTimeout: timeout,
		ResponseHeaderTimeout: timeout,
		ExpectContinueTimeout: timeout,
	}
	// Do not use http.Client here, because it's higher level and processes
	// redirects internally, which is not what we want.
	//client := &http.Client{
	//	Timeout: timeout,
	//	Transport: transport,
	//}
	//response, err := client.Do(request)
	response, err := transport.RoundTrip(request)
	if err != nil {
		log.Println("Request failed:", err)
	}
	return response
}

// handler contains the address of the main Target and the one for the Alternative target
type handler struct {
	Target      string
	Alternative string
	Randomizer  rand.Rand
}

// ServeHTTP duplicates the incoming request (req) and does the request to the
// Target and the Alternate target discading the Alternate response
func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var productionRequest, alternativeRequest *http.Request
	if *forwardClientIP {
		updateForwardedHeaders(req)
	}
	if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {
		alternativeRequest, productionRequest = DuplicateRequest(req)
		go func() {
			defer func() {
				if r := recover(); r != nil && *debug {
					log.Println("Recovered in ServeHTTP(alternate request) from:", r)
				}
			}()

			setRequestTarget(alternativeRequest, altTarget)

			if *alternateHostRewrite {
				alternativeRequest.Host = h.Alternative
			}

			timeout := time.Duration(*alternateTimeout) * time.Millisecond
			// This keeps responses from the alternative target away from the outside world.
			_ = handleRequest(alternativeRequest, timeout)
		}()
	} else {
		productionRequest = req
	}
	defer func() {
		if r := recover(); r != nil && *debug {
			log.Println("Recovered in ServeHTTP(production request) from:", r)
		}
	}()

	setRequestTarget(productionRequest, targetProduction)

	if *productionHostRewrite {
		productionRequest.Host = h.Target
	}

	timeout := time.Duration(*productionTimeout) * time.Millisecond
	resp := handleRequest(productionRequest, timeout)

	if resp != nil {
		defer resp.Body.Close()

		// Forward response headers.
		for k, v := range resp.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(resp.StatusCode)

		// Forward response body.
		body, _ := ioutil.ReadAll(resp.Body)
		w.Write(body)
	}
}

func main() {
	flag.Parse()

	log.Printf("Starting teeproxy at %s sending to A: %s and B: %s",
	           *listen, *targetProduction, *altTarget)

	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error

	var listener net.Listener

	if len(*tlsPrivateKey) > 0 {
		cer, err := tls.LoadX509KeyPair(*tlsCertificate, *tlsPrivateKey)
		if err != nil {
			log.Fatalf("Failed to load certficate: %s and private key: %s", *tlsCertificate, *tlsPrivateKey)
		}

		config := &tls.Config{Certificates: []tls.Certificate{cer}}
		listener, err = tls.Listen("tcp", *listen, config)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	} else {
		listener, err = net.Listen("tcp", *listen)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	}

	h := handler{
		Target:      *targetProduction,
		Alternative: *altTarget,
		Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),
	}

	server := &http.Server{
		Handler: h,
	}
	server.Serve(listener)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {
	b1 := new(bytes.Buffer)
	b2 := new(bytes.Buffer)
	w := io.MultiWriter(b1, b2)
	io.Copy(w, request.Body)
	defer request.Body.Close()
	request1 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b1},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	request2 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b2},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	return
}

func updateForwardedHeaders(request *http.Request) {
	positionOfColon := strings.LastIndex(request.RemoteAddr, ":")
	var remoteIP string
	if positionOfColon != -1 {
		remoteIP = request.RemoteAddr[:positionOfColon]
	} else {
		Logger.Printf("The default format of request.RemoteAddr should be IP:Port but was %s\n", remoteIP)
		remoteIP = request.RemoteAddr
	}
	insertOrExtendForwardedHeader(request, remoteIP)
	insertOrExtendXFFHeader(request, remoteIP)
}

const XFF_HEADER = "X-Forwarded-For"

func insertOrExtendXFFHeader(request *http.Request, remoteIP string) {
	header := request.Header.Get(XFF_HEADER)
	if header != "" {
		// extend
		request.Header.Set(XFF_HEADER, header + ", " + remoteIP)
	} else {
		// insert
		request.Header.Set(XFF_HEADER, remoteIP)
	}
}

const FORWARDED_HEADER = "Forwarded"

// Implementation according to rfc7239
func insertOrExtendForwardedHeader(request *http.Request, remoteIP string) {
	extension := "for=" + remoteIP
	header := request.Header.Get(FORWARDED_HEADER)
	if header != "" {
		// extend
		request.Header.Set(FORWARDED_HEADER, header + ", " + extension)
	} else {
		// insert
		request.Header.Set(FORWARDED_HEADER, extension)
	}
}
<|endoftext|>