<commit_before>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/netx"
	"github.com/getlantern/preconn"
	"github.com/getlantern/proxy/filters"
)

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			safeClose(downstream)
			err = errors.New("Recovered from panic handling connection: %v", p)
		}
	}()

	err = proxy.handle(ctx, downstreamIn, downstream, nil)
	return
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection: %v", p)
		}
	}()

	conn.Close()
}

func (proxy *proxy) logInitialReadError(downstream net.Conn, err error) error {
	rem := downstream.RemoteAddr()
	r := ""
	if rem != nil {
		r = rem.String()
	}
	txt := err.Error()
	// Ignore our generated error that should have already been reported.
	if strings.HasPrefix(txt, "Client Hello has no cipher suites") {
		log.Debugf("No cipher suites in common -- old Lantern client")
		return err
	}
	// These errors should all typically be internal go errors, typically with TLS. Break them up
	// for stackdriver grouping.
	if strings.Contains(txt, "oversized") {
		return log.Errorf("Oversized record on initial read: %v from %v", err, r)
	}
	if strings.Contains(txt, "first record does not") {
		return log.Errorf("Not a TLS client connection: %v from %v", err, r)
	}
	return log.Errorf("Initial ReadRequest: %v from %v", err, r)
}

func (proxy *proxy) handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn, upstream net.Conn) error {
	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
	}()

	downstreamBuffered := bufio.NewReader(downstreamIn)
	fctx := filters.WrapContext(withAwareConn(ctx), downstream)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = remoteAddr.String()
		}
		if origURLScheme(ctx) == "" {
			fctx = fctx.
				WithValue(ctxKeyOrigURLScheme, req.URL.Scheme).
				WithValue(ctxKeyOrigURLHost, req.URL.Host).
				WithValue(ctxKeyOrigHost, req.Host)
		}
	}

	if err != nil {
		if isUnexpected(err) {
			errResp := proxy.OnError(fctx, req, true, err)
			if errResp != nil {
				proxy.writeResponse(downstream, req, errResp)
			}

			return proxy.logInitialReadError(downstream, err)
		}
		return nil
	}

	var next filters.Next
	if req.Method == http.MethodConnect {
		next = proxy.nextCONNECT(downstream)
	} else {
		var tr idleClosingTransport
		if upstream != nil {
			setUpstreamForAwareConn(fctx, upstream)
			tr = &addressLoggingTransport{
				Transport: &http.Transport{
					DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
						// always use the supplied upstream connection, but don't allow it to
						// be closed by the transport
						return &noCloseConn{upstream}, nil
					},
					// this transport is only used once, don't keep any idle connections,
					// however still allow the transport to close the connection after using
					// it
					MaxIdleConnsPerHost: -1,
				},
				upstream: upstream,
			}
		} else {
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					conn, err := proxy.Dial(ctx, false, net, addr)
					if err == nil {
						// On first dialing conn, handle RequestAware
						setUpstreamForAwareConn(ctx, conn)
						handleRequestAware(ctx)
					}
					return conn, err
				},
				IdleConnTimeout: proxy.IdleTimeout,
				// since we have one transport per downstream connection, we don't need
				// more than this
				MaxIdleConnsPerHost: 1,
			}
		}

		defer tr.CloseIdleConnections()
		next = func(ctx filters.Context, modifiedReq *http.Request) (*http.Response, filters.Context, error) {
			modifiedReq = modifiedReq.WithContext(ctx)
			setRequestForAwareConn(ctx, modifiedReq)
			handleRequestAware(ctx)
			resp, err := tr.RoundTrip(prepareRequest(modifiedReq))
			handleResponseAware(ctx, modifiedReq, resp, err)
			if err != nil {
				err = errors.New("Unable to round-trip http request to upstream: %v", err)
			}
			return resp, ctx, err
		}
	}

	return proxy.processRequests(fctx, req.RemoteAddr, req, downstream, downstreamBuffered, next)
}

func (proxy *proxy) processRequests(ctx filters.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, next filters.Next) error {
	var readErr error
	var resp *http.Response
	var err error

	for {
		if req.URL.Scheme == "" {
			req.URL.Scheme = origURLScheme(ctx)
		}
		if req.URL.Host == "" {
			req.URL.Host = origURLHost(ctx)
		}
		if req.Host == "" {
			req.Host = origHost(ctx)
		}
		resp, ctx, err = proxy.Filter.Apply(ctx, req, next)
		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, false, err)
			if resp != nil {
				log.Debugf("Closing client connection on error: %v", err)
				// On error, we will always close the connection
				resp.Close = true
			}
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return log.Errorf("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if err != nil {
			// We encountered an error on round-tripping, stop now
			return err
		}

		upstream := upstreamConn(ctx)
		upstreamAddr := upstreamAddr(ctx)
		isConnect := upstream != nil || upstreamAddr != ""

		buffered := downstreamBuffered.Buffered()
		if buffered > 0 {
			b, _ := downstreamBuffered.Peek(buffered)
			downstream = preconn.Wrap(downstream, b)
		}

		if isConnect {
			return proxy.proceedWithConnect(ctx, req, upstreamAddr, upstream, downstream)
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				errResp := proxy.OnError(ctx, req, true, readErr)
				if errResp != nil {
					proxy.writeResponse(downstream, req, errResp)
				}
				return log.Errorf("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		ctx = ctx.IncrementRequestNumber()
		req.RemoteAddr = remoteAddr
		req = req.WithContext(ctx)
	}
}

func handleRequestAware(ctx context.Context) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case RequestAware:
			req := requestForAwareConn(ctx)
			t.OnRequest(req)
		}
		return true
	})
}

func handleResponseAware(ctx context.Context, req *http.Request, resp *http.Response, err error) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case ResponseAware:
			t.OnResponse(req, resp, err)
		}
		return true
	})
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	if resp.ProtoMajor == 0 {
		resp.ProtoMajor = 1
		resp.ProtoMinor = 1
	}
	belowHTTP11 := !resp.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
		proxy.addIdleKeepAlive(resp.Header)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	req.Close = false

	// Request Header
	newHeader := make(http.Header)
	copyHeadersForForwarding(newHeader, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	newHeader.Set("Host", req.Host)
	req.Header = newHeader

	// Request URL
	req.URL = cloneURL(req.URL)
	// If req.URL.Scheme was blank, it's http. Otherwise, it's https and we leave
	// it alone.
	if req.URL.Scheme == "" {
		req.URL.Scheme = "http"
	}
	// We need to make sure the host is defined in the URL (not the actual URI)
	req.URL.Host = req.Host

	userAgent := req.UserAgent()
	if userAgent == "" {
		req.Header.Del("User-Agent")
	} else {
		req.Header.Set("User-Agent", userAgent)
	}

	return req
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		case "Keep-Alive":
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	if err == nil {
		return false
	}
	if err == io.EOF {
		return false
	}
	// This is okay per the HTTP spec.
	// See https://www.w3.org/Protocols/rfc2616/rfc2616-sec8.html#sec8.1.4
	if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
		return false
	}

	text := err.Error()
	return !strings.HasSuffix(text, "EOF") &&
		!strings.Contains(text, "i/o timeout") &&
		!strings.Contains(text, "Use of idled network connection") &&
		!strings.Contains(text, "use of closed network connection") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "broken pipe") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "connection reset by peer")
}

func defaultFilter(ctx filters.Context, req *http.Request, next filters.Next) (*http.Response, filters.Context, error) {
	return next(ctx, req)
}

func defaultOnError(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
	return nil
}

type idleClosingTransport interface {
	RoundTrip(req *http.Request) (*http.Response, error)
	CloseIdleConnections()
}

type addressLoggingTransport struct {
	*http.Transport
	upstream net.Conn
}

func (alt *addressLoggingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := alt.Transport.RoundTrip(req)
	if err != nil {
		err = errors.New("Error round-tripping to %v: %v", alt.upstream.RemoteAddr(), err)
	}
	return resp, err
}
<commit_msg>do not forward Proxy-Connection as well<commit_after>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/netx"
	"github.com/getlantern/preconn"
	"github.com/getlantern/proxy/filters"
)

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			safeClose(downstream)
			err = errors.New("Recovered from panic handling connection: %v", p)
		}
	}()

	err = proxy.handle(ctx, downstreamIn, downstream, nil)
	return
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection: %v", p)
		}
	}()

	conn.Close()
}

func (proxy *proxy) logInitialReadError(downstream net.Conn, err error) error {
	rem := downstream.RemoteAddr()
	r := ""
	if rem != nil {
		r = rem.String()
	}
	txt := err.Error()
	// Ignore our generated error that should have already been reported.
	if strings.HasPrefix(txt, "Client Hello has no cipher suites") {
		log.Debugf("No cipher suites in common -- old Lantern client")
		return err
	}
	// These errors should all typically be internal go errors, typically with TLS. Break them up
	// for stackdriver grouping.
	if strings.Contains(txt, "oversized") {
		return log.Errorf("Oversized record on initial read: %v from %v", err, r)
	}
	if strings.Contains(txt, "first record does not") {
		return log.Errorf("Not a TLS client connection: %v from %v", err, r)
	}
	return log.Errorf("Initial ReadRequest: %v from %v", err, r)
}

func (proxy *proxy) handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn, upstream net.Conn) error {
	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
	}()

	downstreamBuffered := bufio.NewReader(downstreamIn)
	fctx := filters.WrapContext(withAwareConn(ctx), downstream)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = remoteAddr.String()
		}
		if origURLScheme(ctx) == "" {
			fctx = fctx.
				WithValue(ctxKeyOrigURLScheme, req.URL.Scheme).
				WithValue(ctxKeyOrigURLHost, req.URL.Host).
				WithValue(ctxKeyOrigHost, req.Host)
		}
	}

	if err != nil {
		if isUnexpected(err) {
			errResp := proxy.OnError(fctx, req, true, err)
			if errResp != nil {
				proxy.writeResponse(downstream, req, errResp)
			}

			return proxy.logInitialReadError(downstream, err)
		}
		return nil
	}

	var next filters.Next
	if req.Method == http.MethodConnect {
		next = proxy.nextCONNECT(downstream)
	} else {
		var tr idleClosingTransport
		if upstream != nil {
			setUpstreamForAwareConn(fctx, upstream)
			tr = &addressLoggingTransport{
				Transport: &http.Transport{
					DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
						// always use the supplied upstream connection, but don't allow it to
						// be closed by the transport
						return &noCloseConn{upstream}, nil
					},
					// this transport is only used once, don't keep any idle connections,
					// however still allow the transport to close the connection after using
					// it
					MaxIdleConnsPerHost: -1,
				},
				upstream: upstream,
			}
		} else {
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					conn, err := proxy.Dial(ctx, false, net, addr)
					if err == nil {
						// On first dialing conn, handle RequestAware
						setUpstreamForAwareConn(ctx, conn)
						handleRequestAware(ctx)
					}
					return conn, err
				},
				IdleConnTimeout: proxy.IdleTimeout,
				// since we have one transport per downstream connection, we don't need
				// more than this
				MaxIdleConnsPerHost: 1,
			}
		}

		defer tr.CloseIdleConnections()
		next = func(ctx filters.Context, modifiedReq *http.Request) (*http.Response, filters.Context, error) {
			modifiedReq = modifiedReq.WithContext(ctx)
			setRequestForAwareConn(ctx, modifiedReq)
			handleRequestAware(ctx)
			resp, err := tr.RoundTrip(prepareRequest(modifiedReq))
			handleResponseAware(ctx, modifiedReq, resp, err)
			if err != nil {
				err = errors.New("Unable to round-trip http request to upstream: %v", err)
			}
			return resp, ctx, err
		}
	}

	return proxy.processRequests(fctx, req.RemoteAddr, req, downstream, downstreamBuffered, next)
}

func (proxy *proxy) processRequests(ctx filters.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, next filters.Next) error {
	var readErr error
	var resp *http.Response
	var err error

	for {
		if req.URL.Scheme == "" {
			req.URL.Scheme = origURLScheme(ctx)
		}
		if req.URL.Host == "" {
			req.URL.Host = origURLHost(ctx)
		}
		if req.Host == "" {
			req.Host = origHost(ctx)
		}
		resp, ctx, err = proxy.Filter.Apply(ctx, req, next)
		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, false, err)
			if resp != nil {
				log.Debugf("Closing client connection on error: %v", err)
				// On error, we will always close the connection
				resp.Close = true
			}
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return log.Errorf("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if err != nil {
			// We encountered an error on round-tripping, stop now
			return err
		}

		upstream := upstreamConn(ctx)
		upstreamAddr := upstreamAddr(ctx)
		isConnect := upstream != nil || upstreamAddr != ""

		buffered := downstreamBuffered.Buffered()
		if buffered > 0 {
			b, _ := downstreamBuffered.Peek(buffered)
			downstream = preconn.Wrap(downstream, b)
		}

		if isConnect {
			return proxy.proceedWithConnect(ctx, req, upstreamAddr, upstream, downstream)
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				errResp := proxy.OnError(ctx, req, true, readErr)
				if errResp != nil {
					proxy.writeResponse(downstream, req, errResp)
				}
				return log.Errorf("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		ctx = ctx.IncrementRequestNumber()
		req.RemoteAddr = remoteAddr
		req = req.WithContext(ctx)
	}
}

func handleRequestAware(ctx context.Context) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case RequestAware:
			req := requestForAwareConn(ctx)
			t.OnRequest(req)
		}
		return true
	})
}

func handleResponseAware(ctx context.Context, req *http.Request, resp *http.Response, err error) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case ResponseAware:
			t.OnResponse(req, resp, err)
		}
		return true
	})
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	if resp.ProtoMajor == 0 {
		resp.ProtoMajor = 1
		resp.ProtoMinor = 1
	}
	belowHTTP11 := !resp.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
		proxy.addIdleKeepAlive(resp.Header)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	req.Close = false

	// Request Header
	newHeader := make(http.Header)
	copyHeadersForForwarding(newHeader, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	newHeader.Set("Host", req.Host)
	req.Header = newHeader

	// Request URL
	req.URL = cloneURL(req.URL)
	// If req.URL.Scheme was blank, it's http. Otherwise, it's https and we leave
	// it alone.
	if req.URL.Scheme == "" {
		req.URL.Scheme = "http"
	}
	// We need to make sure the host is defined in the URL (not the actual URI)
	req.URL.Host = req.Host

	userAgent := req.UserAgent()
	if userAgent == "" {
		req.Header.Del("User-Agent")
	} else {
		req.Header.Set("User-Agent", userAgent)
	}

	return req
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		case "Keep-Alive":
		case "Proxy-Connection": // It's non-standard but every browser adds it when sending request to proxy.
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	if err == nil {
		return false
	}
	if err == io.EOF {
		return false
	}
	// This is okay per the HTTP spec.
	// See https://www.w3.org/Protocols/rfc2616/rfc2616-sec8.html#sec8.1.4
	if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
		return false
	}

	text := err.Error()
	return !strings.HasSuffix(text, "EOF") &&
		!strings.Contains(text, "i/o timeout") &&
		!strings.Contains(text, "Use of idled network connection") &&
		!strings.Contains(text, "use of closed network connection") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "broken pipe") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "connection reset by peer")
}

func defaultFilter(ctx filters.Context, req *http.Request, next filters.Next) (*http.Response, filters.Context, error) {
	return next(ctx, req)
}

func defaultOnError(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
	return nil
}

type idleClosingTransport interface {
	RoundTrip(req *http.Request) (*http.Response, error)
	CloseIdleConnections()
}

type addressLoggingTransport struct {
	*http.Transport
	upstream net.Conn
}

func (alt *addressLoggingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := alt.Transport.RoundTrip(req)
	if err != nil {
		err = errors.New("Error round-tripping to %v: %v", alt.upstream.RemoteAddr(), err)
	}
	return resp, err
}
<|endoftext|>


<commit_before>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/proxy/filters"
)

type contextKey string

const (
	ctxKeyDownstream         = contextKey("downstream")
	ctxKeyDownstreamBuffered = contextKey("downstreamBuffered")
)

// DownstreamConn retrieves the downstream connection from the given Context.
func DownstreamConn(ctx context.Context) net.Conn {
	return ctx.Value(ctxKeyDownstream).(net.Conn)
}

// DownstreamBuffered retrieves the downstream buffered reader from the given
// Context.
func DownstreamBuffered(ctx context.Context) *bufio.Reader {
	return ctx.Value(ctxKeyDownstreamBuffered).(*bufio.Reader)
}

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
	if opts.IdleTimeout > 0 {
		opts.Filter = filters.Join(filters.FilterFunc(func(ctx context.Context, req *http.Request, next filters.Next) (*http.Response, error) {
			resp, err := next(ctx, req)
			if resp != nil {
				opts.addIdleKeepAlive(resp.Header)
			}
			return resp, err
		}), opts.Filter)
	}
}

type httpInterceptor struct {
	idleTimeout time.Duration
	onRequest   func(req *http.Request) *http.Request
	onResponse  func(resp *http.Response) *http.Response
	onError     func(req *http.Request, err error) *http.Response
	dial        DialFunc
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstream net.Conn) error {
	downstreamBuffered := bufio.NewReader(downstream)
	ctx = context.WithValue(ctx, ctxKeyDownstream, downstream)
	ctx = context.WithValue(ctx, ctxKeyDownstreamBuffered, downstreamBuffered)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = downstream.RemoteAddr().String()
		}
	}
	if err != nil {
		errResp := proxy.OnError(ctx, req, err)
		if errResp != nil {
			proxy.writeResponse(downstream, req, errResp)
		}
		return err
	}

	if req.Method == http.MethodConnect {
		return proxy.handleCONNECT(ctx, downstream, req)
	}
	return proxy.handleHTTP(ctx, downstream, downstreamBuffered, req)
}

func (proxy *proxy) handleHTTP(ctx context.Context, downstream net.Conn, downstreamBuffered *bufio.Reader, req *http.Request) error {
	tr := &http.Transport{
		DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
			return proxy.Dial(false, net, addr)
		},
		IdleConnTimeout: proxy.IdleTimeout,
		// since we have one transport per downstream connection, we don't need
		// more than this
		MaxIdleConnsPerHost: 1,
	}

	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
		tr.CloseIdleConnections()
	}()

	return proxy.processRequests(ctx, req.RemoteAddr, req, downstream, downstreamBuffered, tr)
}

func (proxy *proxy) processRequests(ctx context.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, tr *http.Transport) error {
	var readErr error

	for {
		resp, err := proxy.Filter.Apply(ctx, req, func(ctx context.Context, modifiedReq *http.Request) (*http.Response, error) {
			return tr.RoundTrip(prepareRequest(modifiedReq))
		})

		if err != nil {
			resp = proxy.OnError(ctx, req, err)
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return errors.New("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				return errors.New("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		req.RemoteAddr = remoteAddr
	}
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	belowHTTP11 := !resp.Request.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	outReq := new(http.Request)
	// Beware, this will make a shallow copy. We have to copy all maps
	*outReq = *req

	outReq.Proto = "HTTP/1.1"
	outReq.ProtoMajor = 1
	outReq.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	outReq.Close = false

	// Request Header
	outReq.Header = make(http.Header)
	copyHeadersForForwarding(outReq.Header, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	outReq.Header.Set("Host", req.Host)

	// Request URL
	outReq.URL = cloneURL(req.URL)
	// We know that is going to be HTTP always because HTTPS isn't forwarded.
	// We need to hardcode it here because req.URL.Scheme can be undefined, since
	// client request don't need to use absolute URIs
	outReq.URL.Scheme = "http"
	// We need to make sure the host is defined in the URL (not the actual URI)
	outReq.URL.Host = req.Host
	outReq.URL.RawQuery = req.URL.RawQuery
	outReq.Body = req.Body

	userAgent := req.UserAgent()
	if userAgent == "" {
		outReq.Header.Del("User-Agent")
	} else {
		outReq.Header.Set("User-Agent", userAgent)
	}

	return outReq
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		// case "Keep-Alive":
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	text := err.Error()
	return !strings.HasSuffix(text, "EOF") && !strings.Contains(text, "use of closed network connection") && !strings.Contains(text, "Use of idled network connection") && !strings.Contains(text, "broken pipe")
}

func defaultFilter(ctx context.Context, req *http.Request, next filters.Next) (*http.Response, error) {
	return next(ctx, req)
}

func defaultOnError(ctx context.Context, req *http.Request, err error) *http.Response {
	return nil
}
<commit_msg>Only generating error response if original response was nil<commit_after>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/proxy/filters"
)

type contextKey string

const (
	ctxKeyDownstream         = contextKey("downstream")
	ctxKeyDownstreamBuffered = contextKey("downstreamBuffered")
)

// DownstreamConn retrieves the downstream connection from the given Context.
func DownstreamConn(ctx context.Context) net.Conn {
	return ctx.Value(ctxKeyDownstream).(net.Conn)
}

// DownstreamBuffered retrieves the downstream buffered reader from the given
// Context.
func DownstreamBuffered(ctx context.Context) *bufio.Reader {
	return ctx.Value(ctxKeyDownstreamBuffered).(*bufio.Reader)
}

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
	if opts.IdleTimeout > 0 {
		opts.Filter = filters.Join(filters.FilterFunc(func(ctx context.Context, req *http.Request, next filters.Next) (*http.Response, error) {
			resp, err := next(ctx, req)
			if resp != nil {
				opts.addIdleKeepAlive(resp.Header)
			}
			return resp, err
		}), opts.Filter)
	}
}

type httpInterceptor struct {
	idleTimeout time.Duration
	onRequest   func(req *http.Request) *http.Request
	onResponse  func(resp *http.Response) *http.Response
	onError     func(req *http.Request, err error) *http.Response
	dial        DialFunc
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstream net.Conn) error {
	downstreamBuffered := bufio.NewReader(downstream)
	ctx = context.WithValue(ctx, ctxKeyDownstream, downstream)
	ctx = context.WithValue(ctx, ctxKeyDownstreamBuffered, downstreamBuffered)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = downstream.RemoteAddr().String()
		}
	}
	if err != nil {
		errResp := proxy.OnError(ctx, req, err)
		if errResp != nil {
			proxy.writeResponse(downstream, req, errResp)
		}
		return err
	}

	if req.Method == http.MethodConnect {
		return proxy.handleCONNECT(ctx, downstream, req)
	}
	return proxy.handleHTTP(ctx, downstream, downstreamBuffered, req)
}

func (proxy *proxy) handleHTTP(ctx context.Context, downstream net.Conn, downstreamBuffered *bufio.Reader, req *http.Request) error {
	tr := &http.Transport{
		DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
			return proxy.Dial(false, net, addr)
		},
		IdleConnTimeout: proxy.IdleTimeout,
		// since we have one transport per downstream connection, we don't need
		// more than this
		MaxIdleConnsPerHost: 1,
	}

	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
		tr.CloseIdleConnections()
	}()

	return proxy.processRequests(ctx, req.RemoteAddr, req, downstream, downstreamBuffered, tr)
}

func (proxy *proxy) processRequests(ctx context.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, tr *http.Transport) error {
	var readErr error

	for {
		resp, err := proxy.Filter.Apply(ctx, req, func(ctx context.Context, modifiedReq *http.Request) (*http.Response, error) {
			return tr.RoundTrip(prepareRequest(modifiedReq))
		})

		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, err)
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return errors.New("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				return errors.New("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		req.RemoteAddr = remoteAddr
	}
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	belowHTTP11 := !resp.Request.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	outReq := new(http.Request)
	// Beware, this will make a shallow copy. We have to copy all maps
	*outReq = *req

	outReq.Proto = "HTTP/1.1"
	outReq.ProtoMajor = 1
	outReq.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	outReq.Close = false

	// Request Header
	outReq.Header = make(http.Header)
	copyHeadersForForwarding(outReq.Header, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	outReq.Header.Set("Host", req.Host)

	// Request URL
	outReq.URL = cloneURL(req.URL)
	// We know that is going to be HTTP always because HTTPS isn't forwarded.
	// We need to hardcode it here because req.URL.Scheme can be undefined, since
	// client request don't need to use absolute URIs
	outReq.URL.Scheme = "http"
	// We need to make sure the host is defined in the URL (not the actual URI)
	outReq.URL.Host = req.Host
	outReq.URL.RawQuery = req.URL.RawQuery
	outReq.Body = req.Body

	userAgent := req.UserAgent()
	if userAgent == "" {
		outReq.Header.Del("User-Agent")
	} else {
		outReq.Header.Set("User-Agent", userAgent)
	}

	return outReq
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		// case "Keep-Alive":
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	text := err.Error()
	return !strings.HasSuffix(text, "EOF") && !strings.Contains(text, "use of closed network connection") && !strings.Contains(text, "Use of idled network connection") && !strings.Contains(text, "broken pipe")
}

func defaultFilter(ctx context.Context, req *http.Request, next filters.Next) (*http.Response, error) {
	return next(ctx, req)
}

func defaultOnError(ctx context.Context, req *http.Request, err error) *http.Response {
	return nil
}
<|endoftext|>


<commit_before>package proxy

import (
	"bytes"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"

	"github.com/asciimoo/filtron/rule"
	"github.com/asciimoo/filtron/types"
)

var transport *http.Transport = &http.Transport{
	DisableKeepAlives: false,
}

var client *http.Client = &http.Client{Transport: transport}

type Proxy struct {
	NumberOfRequests uint
	target           string
	rules            *[]*rule.Rule
}

func Listen(address, target string, rules *[]*rule.Rule) *Proxy {
	log.Println("Proxy listens on", address)
	p := &Proxy{0, target, rules}
	s := http.NewServeMux()
	s.HandleFunc("/", p.Handler)
	go func(address string, s *http.ServeMux) {
		http.ListenAndServe(address, s)
	}(address, s)
	return p
}

func (p *Proxy) Handler(w http.ResponseWriter, r *http.Request) {

	respState := types.UNTOUCHED
	for _, rule := range *p.rules {
		s := rule.Validate(r, w, respState)
		if s > respState {
			respState = s
		}
	}
	if respState == types.SERVED {
		//w.WriteHeader(429)
		//w.Write([]byte("Rate limit exceeded"))
		//log.Println("Blocked:", uri.String())
		return
	}

	uri, err := url.Parse(p.target)
	fatal(err)
	err = r.ParseForm()
	fatal(err)
	uri.Path = path.Join(uri.Path, r.URL.Path)
	uri.RawQuery = r.URL.RawQuery

	var appRequest *http.Request
	if r.Method == "POST" || r.Method == "PUT" {
		appRequest, err = http.NewRequest(r.Method, uri.String(), bytes.NewBufferString(r.PostForm.Encode()))
	} else {
		appRequest, err = http.NewRequest(r.Method, uri.String(), nil)
	}
	fatal(err)
	copyHeaders(&r.Header, &appRequest.Header)

	resp, err := transport.RoundTrip(appRequest)
	if err != nil {
		log.Println("Response error:", err, resp)
		w.WriteHeader(429)
		return
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	fatal(err)

	dH := w.Header()
	copyHeaders(&resp.Header, &dH)
	w.WriteHeader(resp.StatusCode)

	w.Write(body)
}

func (p *Proxy) ReloadRules(filename string) error {
	rules, err := rule.ParseJSON(filename)
	if err != nil {
		return err
	}
	p.rules = &rules
	return nil
}

func fatal(err error) {
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
}

func copyHeaders(source *http.Header, dest *http.Header) {
	for n, v := range *source {
		if n == "Connection" || n == "Accept-Encoding" {
			continue
		}
		for _, vv := range v {
			dest.Add(n, vv)
		}
	}
}
<commit_msg>[fix] parse request data before validation<commit_after>package proxy

import (
	"bytes"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"

	"github.com/asciimoo/filtron/rule"
	"github.com/asciimoo/filtron/types"
)

var transport *http.Transport = &http.Transport{
	DisableKeepAlives: false,
}

var client *http.Client = &http.Client{Transport: transport}

type Proxy struct {
	NumberOfRequests uint
	target           string
	rules            *[]*rule.Rule
}

func Listen(address, target string, rules *[]*rule.Rule) *Proxy {
	log.Println("Proxy listens on", address)
	p := &Proxy{0, target, rules}
	s := http.NewServeMux()
	s.HandleFunc("/", p.Handler)
	go func(address string, s *http.ServeMux) {
		http.ListenAndServe(address, s)
	}(address, s)
	return p
}

func (p *Proxy) Handler(w http.ResponseWriter, r *http.Request) {

	err := r.ParseForm()
	fatal(err)

	respState := types.UNTOUCHED
	for _, rule := range *p.rules {
		s := rule.Validate(r, w, respState)
		if s > respState {
			respState = s
		}
	}
	if respState == types.SERVED {
		return
	}

	uri, err := url.Parse(p.target)
	fatal(err)
	uri.Path = path.Join(uri.Path, r.URL.Path)
	uri.RawQuery = r.URL.RawQuery

	var appRequest *http.Request
	if r.Method == "POST" || r.Method == "PUT" {
		appRequest, err = http.NewRequest(r.Method, uri.String(), bytes.NewBufferString(r.PostForm.Encode()))
	} else {
		appRequest, err = http.NewRequest(r.Method, uri.String(), nil)
	}
	fatal(err)
	copyHeaders(&r.Header, &appRequest.Header)

	resp, err := transport.RoundTrip(appRequest)
	if err != nil {
		log.Println("Response error:", err, resp)
		w.WriteHeader(429)
		return
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	fatal(err)

	dH := w.Header()
	copyHeaders(&resp.Header, &dH)
	w.WriteHeader(resp.StatusCode)

	w.Write(body)
}

func (p *Proxy) ReloadRules(filename string) error {
	rules, err := rule.ParseJSON(filename)
	if err != nil {
		return err
	}
	p.rules = &rules
	return nil
}

func fatal(err error) {
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
}

func copyHeaders(source *http.Header, dest *http.Header) {
	for n, v := range *source {
		if n == "Connection" || n == "Accept-Encoding" {
			continue
		}
		for _, vv := range v {
			dest.Add(n, vv)
		}
	}
}
<|endoftext|>


<commit_before>package proxy

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

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

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

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

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

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

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

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

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

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			// fmt.Printf("%#v", string(buff[:n]))
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			fmt.Printf("newPacket : %v\n", newPacket)
			fmt.Printf("len(r) : %v\n", len(r))
			// if remainingBytes > 0 {
			// 	if remainingBytes <= n {
			// 		fmt.Println("1")
			// 		newPacket = true
			// 		msg = msg + string(r.next(remainingBytes))
			// 		fmt.Printf("2 Remaining bytes: %d\n", remainingBytes)
			// 		remainingBytes = n - remainingBytes
			// 		fmt.Printf("3 Remaining bytes: %d\n", remainingBytes)
			// 		// fmt.Println("msg: ", string(msg))
			// 	} else {
			// 		fmt.Println("2")
			// 		newPacket = false
			// 		msg = msg + string(r.next(remainingBytes))
			// 		remainingBytes = remainingBytes - n
			// 	}
			//
			// 	fmt.Printf("1 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
			// }
			fmt.Println("3")
		NewP:
			fmt.Println("4")
			if newPacket && len(r) > 4 && remainingBytes == 0 {
				fmt.Println("5")
				// remainingBytes = 0
				newPacket = false
				fmt.Printf("2 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
				var msg []byte
				t := r.byte()
				n = n - 1
				fmt.Println("t: ", string(t))
				switch t {
				case 'Q', 'B', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'P', 'p', 'S', 'X':
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					if remainingBytes < 4 {
						fmt.Errorf("ERROR: remainingBytes can't be less than 4 bytes if int32")
					} else {
						remainingBytes = remainingBytes - 4
						if remainingBytes > 0 {
							if remainingBytes <= n {
								newPacket = true
								msg = append(msg, r.next(remainingBytes)[:]...)
								msg = []byte(stripchars(string(msg),
									"\n\t"))
								// msg = []byte(strings.Replace(string(msg), `\n`, "", -1))
								// msg = []byte(strings.Replace(string(msg), `\t`, "", -1))
								remainingBytes = n - remainingBytes
								// fmt.Printf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
								// fmt.Printf("3 Remaining bytes: %d \tmsg: %v\n", remainingBytes, msg)
								fmt.Printf("3 Remaining bytes: %d \tmsg: ", remainingBytes)
								for _, v := range msg {
									fmt.Printf("'%v': '%s'  ", v, string(v))
								}
								fmt.Println(msg)
								fmt.Println(string(msg))
								goto NewP
							} else {
								newPacket = false
								msg = append(msg, r.next(remainingBytes)[:]...)
								// msg = bytes.Replace(msg, []byte("\n\t"), []byte(" "), -1)
								msg = []byte(stripchars(string(msg),
									"\n\t"))
								remainingBytes = remainingBytes - n
								fmt.Printf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							}
						}
					}
					// case :
					// 	fmt.Println("TODO")
					// 	// c.rxReadyForQuery(r)
					// 	remainingBytes = r.int32()
					// 	remainingBytes = remainingBytes - 4
					// 	if remainingBytes > 0 {
					// 		if remainingBytes <= n {
					// 			newPacket = true
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = n - remainingBytes
					// 			fmt.Printf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
					// 			// fmt.Println(msg)
					// 			goto NewP
					// 		} else {
					// 			newPacket = false
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = remainingBytes - n
					// 			fmt.Printf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
					// 		}
					// 	}
					// case rowDescription:
					// case dataRow:
					// case bindComplete:
					// case commandComplete:
					// 	commandTag = CommandTag(r.readCString())
				default:
					fmt.Println("6")
					remainingBytes = 0
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			} else {
				fmt.Println("7")
				remainingBytes = 0
			}
			// r = append(r, buff[:]...)

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

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

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

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

}

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

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

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

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

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

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

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

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

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

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

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

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

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	spaces := regexp.MustCompile("[\n\t ]+")
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			// fmt.Printf("%#v", string(buff[:n]))
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			fmt.Printf("newPacket : %v\n", newPacket)
			fmt.Printf("len(r) : %v\n", len(r))
			// if remainingBytes > 0 {
			// 	if remainingBytes <= n {
			// 		fmt.Println("1")
			// 		newPacket = true
			// 		msg = msg + string(r.next(remainingBytes))
			// 		fmt.Printf("2 Remaining bytes: %d\n", remainingBytes)
			// 		remainingBytes = n - remainingBytes
			// 		fmt.Printf("3 Remaining bytes: %d\n", remainingBytes)
			// 		// fmt.Println("msg: ", string(msg))
			// 	} else {
			// 		fmt.Println("2")
			// 		newPacket = false
			// 		msg = msg + string(r.next(remainingBytes))
			// 		remainingBytes = remainingBytes - n
			// 	}
			//
			// 	fmt.Printf("1 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
			// }
			fmt.Println("3")
		NewP:
			fmt.Println("4")
			if newPacket && len(r) > 4 && remainingBytes == 0 {
				fmt.Println("5")
				// remainingBytes = 0
				newPacket = false
				fmt.Printf("2 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
				var msg []byte
				t := r.byte()
				n = n - 1
				fmt.Println("t: ", string(t))
				switch t {
				case 'Q', 'B', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'P', 'p', 'S', 'X':
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					if remainingBytes < 4 {
						fmt.Errorf("ERROR: remainingBytes can't be less than 4 bytes if int32")
					} else {
						remainingBytes = remainingBytes - 4
						if remainingBytes > 0 {
							if remainingBytes <= n {
								newPacket = true
								msg = append(msg, r.next(remainingBytes)[:]...)
								msg = spaces.ReplaceAll(msg, []byte{' '})
								// msg = []byte(stripchars(string(msg),
								// 	"\n\t"))
								// msg = []byte(strings.Replace(string(msg), `\n`, "", -1))
								// msg = []byte(strings.Replace(string(msg), `\t`, "", -1))
								remainingBytes = n - remainingBytes
								// fmt.Printf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
								// fmt.Printf("3 Remaining bytes: %d \tmsg: %v\n", remainingBytes, msg)
								fmt.Printf("3 Remaining bytes: %d \tmsg: ", remainingBytes)
								for _, v := range msg {
									fmt.Printf("'%v': '%s'  ", v, string(v))
								}
								fmt.Println(msg)
								fmt.Println(string(msg))
								goto NewP
							} else {
								newPacket = false
								msg = append(msg, r.next(remainingBytes)[:]...)
								// msg = bytes.Replace(msg, []byte("\n\t"), []byte(" "), -1)
								msg = spaces.ReplaceAll(msg, []byte{' '})
								// msg = []byte(stripchars(string(msg),
								// 	"\n\t"))
								remainingBytes = remainingBytes - n
								fmt.Printf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							}
						}
					}
					// case :
					// 	fmt.Println("TODO")
					// 	// c.rxReadyForQuery(r)
					// 	remainingBytes = r.int32()
					// 	remainingBytes = remainingBytes - 4
					// 	if remainingBytes > 0 {
					// 		if remainingBytes <= n {
					// 			newPacket = true
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = n - remainingBytes
					// 			fmt.Printf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
					// 			// fmt.Println(msg)
					// 			goto NewP
					// 		} else {
					// 			newPacket = false
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = remainingBytes - n
					// 			fmt.Printf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
					// 		}
					// 	}
					// case rowDescription:
					// case dataRow:
					// case bindComplete:
					// case commandComplete:
					// 	commandTag = CommandTag(r.readCString())
				default:
					fmt.Println("6")
					remainingBytes = 0
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			} else {
				fmt.Println("7")
				remainingBytes = 0
			}
			// r = append(r, buff[:]...)

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

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

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

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

}

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

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

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


<commit_before>package proxy

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

	sbAudio "github.com/dh1tw/remoteAudio/sb_audio"
	"github.com/gogo/protobuf/proto"
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/client"
)

type AudioServer struct {
	sync.RWMutex
	name           string
	serviceName    string
	client         client.Client
	rpc            sbAudio.ServerClient
	stateSub       broker.Subscriber
	rxAddress      string
	txAddress      string
	stateAddress   string
	rxOn           bool
	txUser         string
	latency        int
	notifyChangeCb func()
	closePing      chan struct{}
}

func NewAudioServer(name string, client client.Client, opts ...Option) (*AudioServer, error) {

	serviceName := fmt.Sprintf("shackbus.radio.%s.audio", name)

	as := &AudioServer{
		name:         name,
		serviceName:  serviceName,
		client:       client,
		rxAddress:    fmt.Sprintf("%s.rx", serviceName),
		txAddress:    fmt.Sprintf("%s.tx", serviceName),
		stateAddress: fmt.Sprintf("%s.state", serviceName),
		txUser:       "",
		closePing:    make(chan struct{}),
	}

	as.rpc = sbAudio.NewServerClient(as.serviceName, as.client)

	if err := as.getCapabilities(); err != nil {
		return nil, err
	}
	if err := as.getState(); err != nil {
		return nil, err
	}

	sub, err := as.client.Options().Broker.Subscribe(as.stateAddress, as.stateUpdateCb)
	if err != nil {
		return nil, err
	}

	as.stateSub = sub

	// start a go routine to ping our service every second
	go func() {
		for {
			select {
			case <-time.After(time.Second):

				ping, err := as.ping()

				if err != nil {
					log.Println("ping:", err)
					break
				}
				as.Lock()
				as.latency = ping
				as.Unlock()
				go as.notifyChangeCb()
			case <-as.closePing:
				return
			}
		}
	}()

	return as, nil
}

func (as *AudioServer) Close() {
	as.Lock()
	defer as.Unlock()

	as.stateSub.Unsubscribe()
	close(as.closePing)
	as.rpc = nil
	as.client = nil
	as.notifyChangeCb = nil
}

func (as *AudioServer) ping() (int, error) {

	ping := time.Now().UnixNano() / int64(time.Millisecond)

	pingMsg := sbAudio.PingPong{
		Ping: ping,
	}

	pong, err := as.rpc.Ping(context.Background(), &pingMsg)
	if err != nil {
		return 0, err
	}

	now := time.Now().UnixNano() / int64(time.Millisecond)

	res := int((now - pong.Ping) / 2)
	return res, nil
}

func (as *AudioServer) SetNotifyCb(f func()) {
	as.Lock()
	defer as.Unlock()
	as.notifyChangeCb = f
}

func (as *AudioServer) Name() string {
	as.RLock()
	defer as.RUnlock()
	return as.name
}

func (as *AudioServer) ServiceName() string {
	as.RLock()
	defer as.RUnlock()
	return as.serviceName
}

func (as *AudioServer) RxOn() bool {
	as.RLock()
	defer as.RUnlock()
	return as.rxOn
}

func (as *AudioServer) RxAddress() string {
	as.RLock()
	defer as.RUnlock()
	return as.rxAddress
}

func (as *AudioServer) TxAddress() string {
	as.RLock()
	defer as.RUnlock()
	return as.txAddress
}

func (as *AudioServer) StartRxStream() error {
	_, err := as.rpc.StartStream(context.Background(), &sbAudio.None{})
	if err != nil {
		return err
	}
	as.Lock()
	as.rxOn = true
	as.Unlock()
	return nil
}

func (as *AudioServer) StopRxStream() error {
	_, err := as.rpc.StopStream(context.Background(), &sbAudio.None{})
	if err != nil {
		return err
	}
	as.Lock()
	as.rxOn = false
	as.Unlock()
	return nil
}

func (as *AudioServer) TxUser() string {
	as.RLock()
	defer as.RUnlock()
	return as.txUser
}

func (as *AudioServer) Latency() int {
	as.RLock()
	defer as.RUnlock()
	return as.latency
}

func (as *AudioServer) stateUpdateCb(msg broker.Publication) error {

	newState := sbAudio.State{}

	if err := proto.Unmarshal(msg.Message().Body, &newState); err != nil {
		return err
	}
	as.Lock()
	defer as.Unlock()

	as.rxOn = newState.GetRxOn()
	as.txUser = newState.GetTxUser()

	if as.notifyChangeCb != nil {
		go as.notifyChangeCb()
	}

	// TBD check if something has changed and notify subcribers

	return nil
}

func (as *AudioServer) getState() error {
	state, err := as.rpc.GetState(context.Background(), &sbAudio.None{})
	if err != nil {
		return fmt.Errorf("getState: %v", err)
	}
	as.Lock()
	defer as.Unlock()
	as.rxOn = state.RxOn
	as.txUser = state.TxUser

	return nil
}

func (as *AudioServer) getCapabilities() error {
	caps, err := as.rpc.GetCapabilities(context.Background(), &sbAudio.None{})
	if err != nil {
		return fmt.Errorf("getCapabilities: %v", err)
	}
	as.Lock()
	defer as.Unlock()
	as.rxAddress = caps.GetRxStreamAddress()
	if len(as.rxAddress) == 0 {
		return fmt.Errorf("getCapabilities: RxStreamAddress empty")
	}
	as.txAddress = caps.GetTxStreamAddress()
	if len(as.txAddress) == 0 {
		return fmt.Errorf("getCapabilities: TxStreamAddress empty")
	}
	as.stateAddress = caps.GetStateUpdatesAddress()
	if len(as.stateAddress) == 0 {
		return fmt.Errorf("getCapabilities: StateUpdatesAddress empty")
	}
	return nil
}
<commit_msg>increased ping time to 3 sec<commit_after>package proxy

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

	sbAudio "github.com/dh1tw/remoteAudio/sb_audio"
	"github.com/gogo/protobuf/proto"
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/client"
)

type AudioServer struct {
	sync.RWMutex
	name           string
	serviceName    string
	client         client.Client
	rpc            sbAudio.ServerClient
	stateSub       broker.Subscriber
	rxAddress      string
	txAddress      string
	stateAddress   string
	rxOn           bool
	txUser         string
	latency        int
	notifyChangeCb func()
	closePing      chan struct{}
}

func NewAudioServer(name string, client client.Client, opts ...Option) (*AudioServer, error) {

	serviceName := fmt.Sprintf("shackbus.radio.%s.audio", name)

	as := &AudioServer{
		name:         name,
		serviceName:  serviceName,
		client:       client,
		rxAddress:    fmt.Sprintf("%s.rx", serviceName),
		txAddress:    fmt.Sprintf("%s.tx", serviceName),
		stateAddress: fmt.Sprintf("%s.state", serviceName),
		txUser:       "",
		closePing:    make(chan struct{}),
	}

	as.rpc = sbAudio.NewServerClient(as.serviceName, as.client)

	if err := as.getCapabilities(); err != nil {
		return nil, err
	}
	if err := as.getState(); err != nil {
		return nil, err
	}

	sub, err := as.client.Options().Broker.Subscribe(as.stateAddress, as.stateUpdateCb)
	if err != nil {
		return nil, err
	}

	as.stateSub = sub

	// start a go routine to ping our service every second
	go func() {
		for {
			select {
			case <-time.After(time.Second * 3):

				ping, err := as.ping()

				if err != nil {
					log.Println("ping:", err) // should be fatal?
					break
				}
				as.Lock()
				as.latency = ping
				as.Unlock()
				go as.notifyChangeCb()
			case <-as.closePing:
				return
			}
		}
	}()

	return as, nil
}

func (as *AudioServer) Close() {
	as.Lock()
	defer as.Unlock()

	as.stateSub.Unsubscribe()
	close(as.closePing)
	as.rpc = nil
	as.client = nil
	as.notifyChangeCb = nil
}

func (as *AudioServer) ping() (int, error) {

	ping := time.Now().UnixNano() / int64(time.Millisecond)

	pingMsg := sbAudio.PingPong{
		Ping: ping,
	}

	pong, err := as.rpc.Ping(context.Background(), &pingMsg)
	if err != nil {
		return 0, err
	}

	now := time.Now().UnixNano() / int64(time.Millisecond)

	res := int((now - pong.Ping) / 2)
	return res, nil
}

func (as *AudioServer) SetNotifyCb(f func()) {
	as.Lock()
	defer as.Unlock()
	as.notifyChangeCb = f
}

func (as *AudioServer) Name() string {
	as.RLock()
	defer as.RUnlock()
	return as.name
}

func (as *AudioServer) ServiceName() string {
	as.RLock()
	defer as.RUnlock()
	return as.serviceName
}

func (as *AudioServer) RxOn() bool {
	as.RLock()
	defer as.RUnlock()
	return as.rxOn
}

func (as *AudioServer) RxAddress() string {
	as.RLock()
	defer as.RUnlock()
	return as.rxAddress
}

func (as *AudioServer) TxAddress() string {
	as.RLock()
	defer as.RUnlock()
	return as.txAddress
}

func (as *AudioServer) StartRxStream() error {
	_, err := as.rpc.StartStream(context.Background(), &sbAudio.None{})
	if err != nil {
		return err
	}
	as.Lock()
	as.rxOn = true
	as.Unlock()
	return nil
}

func (as *AudioServer) StopRxStream() error {
	_, err := as.rpc.StopStream(context.Background(), &sbAudio.None{})
	if err != nil {
		return err
	}
	as.Lock()
	as.rxOn = false
	as.Unlock()
	return nil
}

func (as *AudioServer) TxUser() string {
	as.RLock()
	defer as.RUnlock()
	return as.txUser
}

func (as *AudioServer) Latency() int {
	as.RLock()
	defer as.RUnlock()
	return as.latency
}

func (as *AudioServer) stateUpdateCb(msg broker.Publication) error {

	newState := sbAudio.State{}

	if err := proto.Unmarshal(msg.Message().Body, &newState); err != nil {
		return err
	}
	as.Lock()
	defer as.Unlock()

	as.rxOn = newState.GetRxOn()
	as.txUser = newState.GetTxUser()

	if as.notifyChangeCb != nil {
		go as.notifyChangeCb()
	}

	// TBD check if something has changed and notify subcribers

	return nil
}

func (as *AudioServer) getState() error {
	state, err := as.rpc.GetState(context.Background(), &sbAudio.None{})
	if err != nil {
		return fmt.Errorf("getState: %v", err)
	}
	as.Lock()
	defer as.Unlock()
	as.rxOn = state.RxOn
	as.txUser = state.TxUser

	return nil
}

func (as *AudioServer) getCapabilities() error {
	caps, err := as.rpc.GetCapabilities(context.Background(), &sbAudio.None{})
	if err != nil {
		return fmt.Errorf("getCapabilities: %v", err)
	}
	as.Lock()
	defer as.Unlock()
	as.rxAddress = caps.GetRxStreamAddress()
	if len(as.rxAddress) == 0 {
		return fmt.Errorf("getCapabilities: RxStreamAddress empty")
	}
	as.txAddress = caps.GetTxStreamAddress()
	if len(as.txAddress) == 0 {
		return fmt.Errorf("getCapabilities: TxStreamAddress empty")
	}
	as.stateAddress = caps.GetStateUpdatesAddress()
	if len(as.stateAddress) == 0 {
		return fmt.Errorf("getCapabilities: StateUpdatesAddress empty")
	}
	return nil
}
<|endoftext|>


<commit_before>package mpb

import (
	"io"
	"io/ioutil"
	"time"
)

type proxyReader struct {
	io.ReadCloser
	bar *Bar
}

func (x *proxyReader) Read(p []byte) (int, error) {
	n, err := x.ReadCloser.Read(p)
	x.bar.IncrBy(n)
	if err == io.EOF {
		go x.bar.SetTotal(-1, true)
	}
	return n, err
}

type proxyWriterTo struct {
	io.ReadCloser // *proxyReader
	wt            io.WriterTo
	bar           *Bar
}

func (x *proxyWriterTo) WriteTo(w io.Writer) (int64, error) {
	n, err := x.wt.WriteTo(w)
	x.bar.IncrInt64(n)
	if err == io.EOF {
		go x.bar.SetTotal(-1, true)
	}
	return n, err
}

type ewmaProxyReader struct {
	io.ReadCloser // *proxyReader
	bar           *Bar
}

func (x *ewmaProxyReader) Read(p []byte) (int, error) {
	start := time.Now()
	n, err := x.ReadCloser.Read(p)
	if n > 0 {
		x.bar.DecoratorEwmaUpdate(time.Since(start))
	}
	return n, err
}

type ewmaProxyWriterTo struct {
	io.ReadCloser             // *ewmaProxyReader
	wt            io.WriterTo // *proxyWriterTo
	bar           *Bar
}

func (x *ewmaProxyWriterTo) WriteTo(w io.Writer) (int64, error) {
	start := time.Now()
	n, err := x.wt.WriteTo(w)
	if n > 0 {
		x.bar.DecoratorEwmaUpdate(time.Since(start))
	}
	return n, err
}

func (bar *Bar) newProxyReader(r io.Reader) io.ReadCloser {
	rc := toReadCloser(r)
	wt, isWriterTo := r.(io.WriterTo)
	if bar.hasEwmaDecorators {
		if isWriterTo {
			rc = &ewmaProxyWriterTo{rc, wt, bar}
		} else {
			rc = &ewmaProxyReader{rc, bar}
		}
	} else if isWriterTo {
		rc = &proxyWriterTo{rc, wt, bar}
	} else {
		rc = &proxyReader{rc, bar}
	}
	return rc
}

func toReadCloser(r io.Reader) io.ReadCloser {
	if rc, ok := r.(io.ReadCloser); ok {
		return rc
	}
	return ioutil.NopCloser(r)
}
<commit_msg>fixup! refactor newProxyReader<commit_after>package mpb

import (
	"io"
	"io/ioutil"
	"time"
)

type proxyReader struct {
	io.ReadCloser
	bar *Bar
}

func (x proxyReader) Read(p []byte) (int, error) {
	n, err := x.ReadCloser.Read(p)
	x.bar.IncrBy(n)
	if err == io.EOF {
		go x.bar.SetTotal(-1, true)
	}
	return n, err
}

type proxyWriterTo struct {
	proxyReader
	wt io.WriterTo
}

func (x proxyWriterTo) WriteTo(w io.Writer) (int64, error) {
	n, err := x.wt.WriteTo(w)
	x.bar.IncrInt64(n)
	if err == io.EOF {
		go x.bar.SetTotal(-1, true)
	}
	return n, err
}

type ewmaProxyReader struct {
	proxyReader
}

func (x ewmaProxyReader) Read(p []byte) (int, error) {
	start := time.Now()
	n, err := x.proxyReader.Read(p)
	if n > 0 {
		x.bar.DecoratorEwmaUpdate(time.Since(start))
	}
	return n, err
}

type ewmaProxyWriterTo struct {
	ewmaProxyReader
	wt proxyWriterTo
}

func (x ewmaProxyWriterTo) WriteTo(w io.Writer) (int64, error) {
	start := time.Now()
	n, err := x.wt.WriteTo(w)
	if n > 0 {
		x.bar.DecoratorEwmaUpdate(time.Since(start))
	}
	return n, err
}

func (b *Bar) newProxyReader(r io.Reader) (rc io.ReadCloser) {
	pr := proxyReader{toReadCloser(r), b}
	if wt, ok := r.(io.WriterTo); ok {
		pw := proxyWriterTo{pr, wt}
		if b.hasEwmaDecorators {
			rc = ewmaProxyWriterTo{ewmaProxyReader{pr}, pw}
		} else {
			rc = pw
		}
	} else if b.hasEwmaDecorators {
		rc = ewmaProxyReader{pr}
	} else {
		rc = pr
	}
	return rc
}

func toReadCloser(r io.Reader) io.ReadCloser {
	if rc, ok := r.(io.ReadCloser); ok {
		return rc
	}
	return ioutil.NopCloser(r)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/urfave/cli"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
)

var packages []string

func main() {
	app := cli.NewApp()

	app.Name = "builder"
	app.Usage = "Generates a new urfave/cli build!"

	app.Commands = cli.Commands{
		cli.Command{
			Name:   "vet",
			Action: VetActionFunc,
		},
		cli.Command{
			Name:   "test",
			Action: TestActionFunc,
		},
		cli.Command{
			Name:   "gfmrun",
			Action: GfmrunActionFunc,
		},
		cli.Command{
			Name:   "toc",
			Action: TocActionFunc,
		},
		cli.Command{
			Name:   "generate",
			Action: GenActionFunc,
		},
	}

	packages = []string{"cli", "altsrc"}

	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func VetActionFunc(_ *cli.Context) error {
	return exec.Command("go", "vet").Run()
}

func TestActionFunc(c *cli.Context) error {
	for _, pkg := range packages {
		var packageName string

		if pkg == "cli" {
			packageName = fmt.Sprintf("github.com/urfave/cli")
		} else {
			packageName = fmt.Sprintf("github.com/urfave/cli/%s", pkg)
		}

		coverProfile := fmt.Sprintf("--coverprofile=%s.coverprofile", pkg)

		err := exec.Command(
			"go", "test", "-v", coverProfile, packageName,
		).Run()

		if err != nil {
			return err
		}
	}

	return testCleanup()
}

func testCleanup() error {
	var out bytes.Buffer

	for _, pkg := range packages {
		file, err := os.Open(fmt.Sprintf("%s.coverprofile", pkg))
		if err != nil {
			return err
		}

		b, err := ioutil.ReadAll(file)
		if err != nil {
			return err
		}

		out.Write(b)
		err = file.Close()
		if err != nil {
			return err
		}

		err = os.Remove(fmt.Sprintf("%s.coverprofile", pkg))
		if err != nil {
			return err
		}
	}

	outFile, err := os.Create("coverage.txt")
	if err != nil {
		return err
	}

	_, err = out.WriteTo(outFile)
	if err != nil {
		return err
	}

	err = outFile.Close()
	if err != nil {
		return err
	}

	return nil
}

func GfmrunActionFunc(_ *cli.Context) error {
	file, err := os.Open("README.md")
	if err != nil {
		return err
	}

	var counter int
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "package main") {
			counter++
		}
	}

	err = scanner.Err()
	if err != nil {
		return err
	}

	return exec.Command("gfmrun", "-c", fmt.Sprint(counter), "-s", "README.md").Run()
}

func TocActionFunc(_ *cli.Context) error {
	err := exec.Command("node_modules/.bin/markdown-toc", "-i", "README.md").Run()
	if err != nil {
		return err
	}

	err = exec.Command("git", "diff", "--exit-code").Run()
	if err != nil {
		return err
	}

	return nil
}

func GenActionFunc(_ *cli.Context) error {
	err := exec.Command("go", "generate", "flag-gen/main.go").Run()
	if err != nil {
		return err
	}

	err = exec.Command("go", "generate", "cli.go").Run()
	if err != nil {
		return err
	}

	err = exec.Command("git", "diff", "--exit-code").Run()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>move packages slice to global scope<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/urfave/cli"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
)

var packages = []string{"cli", "altsrc"}

func main() {
	app := cli.NewApp()

	app.Name = "builder"
	app.Usage = "Generates a new urfave/cli build!"

	app.Commands = cli.Commands{
		cli.Command{
			Name:   "vet",
			Action: VetActionFunc,
		},
		cli.Command{
			Name:   "test",
			Action: TestActionFunc,
		},
		cli.Command{
			Name:   "gfmrun",
			Action: GfmrunActionFunc,
		},
		cli.Command{
			Name:   "toc",
			Action: TocActionFunc,
		},
		cli.Command{
			Name:   "generate",
			Action: GenActionFunc,
		},
	}

	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func VetActionFunc(_ *cli.Context) error {
	return exec.Command("go", "vet").Run()
}

func TestActionFunc(c *cli.Context) error {
	for _, pkg := range packages {
		var packageName string

		if pkg == "cli" {
			packageName = fmt.Sprintf("github.com/urfave/cli")
		} else {
			packageName = fmt.Sprintf("github.com/urfave/cli/%s", pkg)
		}

		coverProfile := fmt.Sprintf("--coverprofile=%s.coverprofile", pkg)

		err := exec.Command(
			"go", "test", "-v", coverProfile, packageName,
		).Run()

		if err != nil {
			return err
		}
	}

	return testCleanup()
}

func testCleanup() error {
	var out bytes.Buffer

	for _, pkg := range packages {
		file, err := os.Open(fmt.Sprintf("%s.coverprofile", pkg))
		if err != nil {
			return err
		}

		b, err := ioutil.ReadAll(file)
		if err != nil {
			return err
		}

		out.Write(b)
		err = file.Close()
		if err != nil {
			return err
		}

		err = os.Remove(fmt.Sprintf("%s.coverprofile", pkg))
		if err != nil {
			return err
		}
	}

	outFile, err := os.Create("coverage.txt")
	if err != nil {
		return err
	}

	_, err = out.WriteTo(outFile)
	if err != nil {
		return err
	}

	err = outFile.Close()
	if err != nil {
		return err
	}

	return nil
}

func GfmrunActionFunc(_ *cli.Context) error {
	file, err := os.Open("README.md")
	if err != nil {
		return err
	}

	var counter int
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "package main") {
			counter++
		}
	}

	err = scanner.Err()
	if err != nil {
		return err
	}

	return exec.Command("gfmrun", "-c", fmt.Sprint(counter), "-s", "README.md").Run()
}

func TocActionFunc(_ *cli.Context) error {
	err := exec.Command("node_modules/.bin/markdown-toc", "-i", "README.md").Run()
	if err != nil {
		return err
	}

	err = exec.Command("git", "diff", "--exit-code").Run()
	if err != nil {
		return err
	}

	return nil
}

func GenActionFunc(_ *cli.Context) error {
	err := exec.Command("go", "generate", "flag-gen/main.go").Run()
	if err != nil {
		return err
	}

	err = exec.Command("go", "generate", "cli.go").Run()
	if err != nil {
		return err
	}

	err = exec.Command("git", "diff", "--exit-code").Run()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package build

import (
	"fmt"
	"path/filepath"

	"github.com/sourcegraph/makex"
	"sourcegraph.com/sourcegraph/srcgraph/buildstore"
	"sourcegraph.com/sourcegraph/srcgraph/config"
	"sourcegraph.com/sourcegraph/srcgraph/grapher2"
	_ "sourcegraph.com/sourcegraph/srcgraph/toolchain/all_toolchains"
	"sourcegraph.com/sourcegraph/srcgraph/unit"
)

func init() {
	RegisterRuleMaker("graph", makeGraphRules)
	buildstore.RegisterDataType("graph.v0", &grapher2.Output{})
}

func makeGraphRules(c *config.Repository, dataDir string, existing []makex.Rule) ([]makex.Rule, error) {
	var rules []makex.Rule
	for _, u := range c.SourceUnits {
		rules = append(rules, &GraphSourceUnitRule{dataDir, u})
	}
	return rules, nil
}

type GraphSourceUnitRule struct {
	dataDir string
	Unit    unit.SourceUnit
}

func (r *GraphSourceUnitRule) Target() string {
	return filepath.Join(r.dataDir, SourceUnitDataFilename(&grapher2.Output{}, r.Unit))
}

func (r *GraphSourceUnitRule) Prereqs() []string { return r.Unit.Paths() }

func (r *GraphSourceUnitRule) Recipes() []string {
	return []string{
		"mkdir -p `dirname $@`",
		fmt.Sprintf("srcgraph -v graph -json %q 1> $@", unit.MakeID(r.Unit)),
	}
}
<commit_msg>rename<commit_after>package build

import (
	"fmt"
	"path/filepath"

	"github.com/sourcegraph/makex"
	"sourcegraph.com/sourcegraph/srcgraph/buildstore"
	"sourcegraph.com/sourcegraph/srcgraph/config"
	"sourcegraph.com/sourcegraph/srcgraph/grapher2"
	_ "sourcegraph.com/sourcegraph/srcgraph/toolchain/all_toolchains"
	"sourcegraph.com/sourcegraph/srcgraph/unit"
)

func init() {
	RegisterRuleMaker("graph", makeGraphRules)
	buildstore.RegisterDataType("graph.v0", &grapher2.Output{})
}

func makeGraphRules(c *config.Repository, dataDir string, existing []makex.Rule) ([]makex.Rule, error) {
	var rules []makex.Rule
	for _, u := range c.SourceUnits {
		rules = append(rules, &GraphUnitRule{dataDir, u})
	}
	return rules, nil
}

type GraphUnitRule struct {
	dataDir string
	Unit    unit.SourceUnit
}

func (r *GraphUnitRule) Target() string {
	return filepath.Join(r.dataDir, SourceUnitDataFilename(&grapher2.Output{}, r.Unit))
}

func (r *GraphUnitRule) Prereqs() []string { return r.Unit.Paths() }

func (r *GraphUnitRule) Recipes() []string {
	return []string{
		"mkdir -p `dirname $@`",
		fmt.Sprintf("srcgraph -v graph -json %q 1> $@", unit.MakeID(r.Unit)),
	}
}
<|endoftext|>


<commit_before>package router

import "github.com/AsynkronIT/protoactor-go/actor"

// process serves as a proxy to the router implementation and forwards messages directly to the routee. This
// optimization avoids serializing router messages through an actor
type process struct {
	router *actor.PID
	state  Interface
}

func (ref *process) SendUserMessage(pid *actor.PID, message interface{}, sender *actor.PID) {
	if _, ok := message.(ManagementMessage); ok {
		r, _ := actor.ProcessRegistry.Get(ref.router)
		r.SendUserMessage(pid, message, sender)
	} else {
		ref.state.RouteMessage(message, sender)
	}
}

func (ref *process) SendSystemMessage(pid *actor.PID, message interface{}) {
	r, _ := actor.ProcessRegistry.Get(ref.router)
	r.SendSystemMessage(pid, message)
}

func (ref *process) Stop(pid *actor.PID) {
	actor.ProcessRegistry.Remove(pid)
	ref.SendSystemMessage(pid, &actor.Stop{})
}
<commit_msg>fix: short-term fix to ensure router properly cleans up<commit_after>package router

import (
	"sync"
	"sync/atomic"

	"github.com/AsynkronIT/protoactor-go/actor"
)

// process serves as a proxy to the router implementation and forwards messages directly to the routee. This
// optimization avoids serializing router messages through an actor
type process struct {
	router   *actor.PID
	state    Interface
	mu       sync.Mutex
	watchers actor.PIDSet
	stopping int32
}

func (ref *process) SendUserMessage(pid *actor.PID, message interface{}, sender *actor.PID) {
	if _, ok := message.(ManagementMessage); ok {
		r, _ := actor.ProcessRegistry.Get(ref.router)
		r.SendUserMessage(pid, message, sender)
	} else {
		ref.state.RouteMessage(message, sender)
	}
}

func (ref *process) SendSystemMessage(pid *actor.PID, message interface{}) {
	switch msg := message.(type) {
	case *actor.Watch:
		if atomic.LoadInt32(&ref.stopping) == 1 {
			if r, ok := actor.ProcessRegistry.Get(msg.Watcher); ok {
				r.SendSystemMessage(msg.Watcher, &actor.Terminated{Who: pid})
			}
			return
		}
		ref.mu.Lock()
		ref.watchers.Add(msg.Watcher)
		ref.mu.Unlock()

	case *actor.Unwatch:
		ref.mu.Lock()
		ref.watchers.Remove(msg.Watcher)
		ref.mu.Unlock()

	case *actor.Stop:
		term := &actor.Terminated{Who: pid}
		ref.mu.Lock()
		ref.watchers.ForEach(func(_ int, other actor.PID) {
			if r, ok := actor.ProcessRegistry.Get(&other); ok {
				r.SendSystemMessage(&other, term)
			}
		})
		ref.mu.Unlock()

	default:
		r, _ := actor.ProcessRegistry.Get(ref.router)
		r.SendSystemMessage(pid, message)

	}
}

func (ref *process) Stop(pid *actor.PID) {
	if atomic.SwapInt32(&ref.stopping, 1) == 1 {
		return
	}

	ref.router.StopFuture().Wait()
	actor.ProcessRegistry.Remove(pid)
	ref.SendSystemMessage(pid, &actor.Stop{})
}
<|endoftext|>


<commit_before>package kloud

import (
	"fmt"
	"koding/kites/kloud/eventer"
	"koding/kites/kloud/protocol"

	"github.com/koding/kite"
)

type domainArgs struct {
	DomainName string
}

type domainFunc func(*protocol.Machine, *domainArgs) (interface{}, error)

func (k *Kloud) domainHandler(r *kite.Request, fn domainFunc) (resp interface{}, err error) {
	args := &domainArgs{}
	if err := r.Args.One().Unmarshal(args); err != nil {
		return nil, err
	}

	if args.DomainName == "" {
		return nil, fmt.Errorf("domain name argument is empty")
	}

	m, err := k.PrepareMachine(r)
	if err != nil {
		return nil, err
	}

	// PreparMachine is locking for us, so unlock after we are done
	defer k.Locker.Unlock(m.Id)

	if m.IpAddress == "" {
		return nil, fmt.Errorf("ip address is not defined")
	}

	// fake eventer to avoid panics if someone tries to use the eventer
	m.Eventer = &eventer.Events{}

	if err := k.Domainer.Validate(args.DomainName, r.Username); err != nil {
		return nil, err
	}

	return fn(m, args)
}

func (k *Kloud) DomainAdd(r *kite.Request) (resp interface{}, reqErr error) {
	addFunc := func(m *protocol.Machine, args *domainArgs) (interface{}, error) {
		// a non nil means the domain exists
		if _, err := k.Domainer.Get(args.DomainName); err == nil {
			return nil, fmt.Errorf("domain record does exists")
		}

		// now assign the machine ip to the given domain name
		if err := k.Domainer.Create(args.DomainName, m.IpAddress); err != nil {
			return nil, err
		}

		domain := &protocol.Domain{
			Username:  m.Username,
			MachineId: m.Id,
			Name:      args.DomainName,
		}

		if err := k.DomainStorage.Add(domain); err != nil {
			return nil, err
		}

		return true, nil
	}

	return k.domainHandler(r, addFunc)
}

func (k *Kloud) DomainRemove(r *kite.Request) (resp interface{}, reqErr error) {
	removeFunc := func(m *protocol.Machine, args *domainArgs) (interface{}, error) {
		if err := k.Domainer.Delete(args.DomainName, m.IpAddress); err != nil {
			return nil, err
		}

		if err := k.DomainStorage.Delete(args.DomainName); err != nil {
			return nil, err
		}

		return true, nil
	}

	return k.domainHandler(r, removeFunc)
}

func (k *Kloud) DomainUnset(r *kite.Request) (resp interface{}, reqErr error) {
	unsetFunc := func(m *protocol.Machine, args *domainArgs) (interface{}, error) {
		// be sure the domain does exist in storage before we delete the domain
		if _, err := k.DomainStorage.Get(args.DomainName); err != nil {
			return nil, fmt.Errorf("domain does not exists in DB")
		}

		if err := k.Domainer.Delete(args.DomainName, m.IpAddress); err != nil {
			return nil, err
		}

		// remove the machineID for the given domain. The document is still
		// there, but it'sn associated anymore with this domain.
		if err := k.DomainStorage.UpdateMachine(args.DomainName, ""); err != nil {
			return nil, err
		}

		return true, nil
	}

	return k.domainHandler(r, unsetFunc)
}

func (k *Kloud) DomainSet(r *kite.Request) (resp interface{}, reqErr error) {
	setFunc := func(m *protocol.Machine, args *domainArgs) (interface{}, error) {
		// be sure the domain does exist in storage before we create the domain
		if _, err := k.DomainStorage.Get(args.DomainName); err != nil {
			return nil, fmt.Errorf("domain does not exists in DB")
		}

		if err := k.Domainer.Create(args.DomainName, m.IpAddress); err != nil {
			return nil, err
		}

		// addign the machineID for the given domain.
		if err := k.DomainStorage.UpdateMachine(args.DomainName, m.Id); err != nil {
			return nil, err
		}

		return true, nil
	}

	return k.domainHandler(r, setFunc)
}
<commit_msg>kloud/domain: fix another edge case<commit_after>package kloud

import (
	"fmt"
	"koding/kites/kloud/eventer"
	"koding/kites/kloud/protocol"

	"github.com/koding/kite"
)

type domainArgs struct {
	DomainName string
}

type domainFunc func(*protocol.Machine, *domainArgs) (interface{}, error)

func (k *Kloud) domainHandler(r *kite.Request, fn domainFunc) (resp interface{}, err error) {
	args := &domainArgs{}
	if err := r.Args.One().Unmarshal(args); err != nil {
		return nil, err
	}

	if args.DomainName == "" {
		return nil, fmt.Errorf("domain name argument is empty")
	}

	m, err := k.PrepareMachine(r)
	if err != nil {
		return nil, err
	}

	// PreparMachine is locking for us, so unlock after we are done
	defer k.Locker.Unlock(m.Id)

	if m.IpAddress == "" {
		return nil, fmt.Errorf("ip address is not defined")
	}

	// fake eventer to avoid panics if someone tries to use the eventer
	m.Eventer = &eventer.Events{}

	if err := k.Domainer.Validate(args.DomainName, r.Username); err != nil {
		return nil, err
	}

	return fn(m, args)
}

func (k *Kloud) DomainAdd(r *kite.Request) (resp interface{}, reqErr error) {
	addFunc := func(m *protocol.Machine, args *domainArgs) (interface{}, error) {
		// a non nil means the domain exists
		if _, err := k.Domainer.Get(args.DomainName); err == nil {
			return nil, fmt.Errorf("domain record does exists")
		}

		// now assign the machine ip to the given domain name
		if err := k.Domainer.Create(args.DomainName, m.IpAddress); err != nil {
			return nil, err
		}

		domain := &protocol.Domain{
			Username:  m.Username,
			MachineId: m.Id,
			Name:      args.DomainName,
		}

		if err := k.DomainStorage.Add(domain); err != nil {
			return nil, err
		}

		return true, nil
	}

	return k.domainHandler(r, addFunc)
}

func (k *Kloud) DomainRemove(r *kite.Request) (resp interface{}, reqErr error) {
	removeFunc := func(m *protocol.Machine, args *domainArgs) (interface{}, error) {
		// do not return on error because it might be already delete via unset
		k.Domainer.Delete(args.DomainName, m.IpAddress)

		if err := k.DomainStorage.Delete(args.DomainName); err != nil {
			return nil, err
		}

		return true, nil
	}

	return k.domainHandler(r, removeFunc)
}

func (k *Kloud) DomainUnset(r *kite.Request) (resp interface{}, reqErr error) {
	unsetFunc := func(m *protocol.Machine, args *domainArgs) (interface{}, error) {
		// be sure the domain does exist in storage before we delete the domain
		if _, err := k.DomainStorage.Get(args.DomainName); err != nil {
			return nil, fmt.Errorf("domain does not exists in DB")
		}

		if err := k.Domainer.Delete(args.DomainName, m.IpAddress); err != nil {
			return nil, err
		}

		// remove the machineID for the given domain. The document is still
		// there, but it'sn associated anymore with this domain.
		if err := k.DomainStorage.UpdateMachine(args.DomainName, ""); err != nil {
			return nil, err
		}

		return true, nil
	}

	return k.domainHandler(r, unsetFunc)
}

func (k *Kloud) DomainSet(r *kite.Request) (resp interface{}, reqErr error) {
	setFunc := func(m *protocol.Machine, args *domainArgs) (interface{}, error) {
		// be sure the domain does exist in storage before we create the domain
		if _, err := k.DomainStorage.Get(args.DomainName); err != nil {
			return nil, fmt.Errorf("domain does not exists in DB")
		}

		if err := k.Domainer.Create(args.DomainName, m.IpAddress); err != nil {
			return nil, err
		}

		// addign the machineID for the given domain.
		if err := k.DomainStorage.UpdateMachine(args.DomainName, m.Id); err != nil {
			return nil, err
		}

		return true, nil
	}

	return k.domainHandler(r, setFunc)
}
<|endoftext|>