<commit_before>/*
 Copyright 2021 The CloudEvents Authors
 SPDX-License-Identifier: Apache-2.0
*/

package http

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudevents/sdk-go/v2/binding"
	cecontext "github.com/cloudevents/sdk-go/v2/context"
	"github.com/cloudevents/sdk-go/v2/event"
	"github.com/cloudevents/sdk-go/v2/protocol"
)

const (
	// DefaultShutdownTimeout defines the default timeout given to the http.Server when calling Shutdown.
	DefaultShutdownTimeout = time.Minute * 1
)

type msgErr struct {
	msg    *Message
	respFn protocol.ResponseFn
	err    error
}

// Default error codes that we retry on - string isn't used, it's just there so
// people know what each error code's title is.
// To modify this use Option
var defaultRetriableErrors = map[int]string{
	404: "Not Found",
	413: "Payload Too Large",
	425: "Too Early",
	429: "Too Many Requests",
	502: "Bad Gateway",
	503: "Service Unavailable",
	504: "Gateway Timeout",
}

// Protocol acts as both a http client and a http handler.
type Protocol struct {
	Target          *url.URL
	RequestTemplate *http.Request
	Client          *http.Client
	incoming        chan msgErr

	// OptionsHandlerFn handles the OPTIONS method requests and is intended to
	// implement the abuse protection spec:
	// https://github.com/cloudevents/spec/blob/v1.0/http-webhook.md#4-abuse-protection
	OptionsHandlerFn http.HandlerFunc
	WebhookConfig    *WebhookConfig

	GetHandlerFn    http.HandlerFunc
	DeleteHandlerFn http.HandlerFunc

	// To support Opener:

	// ShutdownTimeout defines the timeout given to the http.Server when calling Shutdown.
	// If 0, DefaultShutdownTimeout is used.
	ShutdownTimeout time.Duration

	// Port is the port configured to bind the receiver to. Defaults to 8080.
	// If you want to know the effective port you're listening to, use GetListeningPort()
	Port int
	// Path is the path to bind the receiver to. Defaults to "/".
	Path string

	// Receive Mutex
	reMu sync.Mutex
	// Handler is the handler the http Server will use. Use this to reuse the
	// http server. If nil, the Protocol will create a one.
	Handler *http.ServeMux

	listener          atomic.Value
	roundTripper      http.RoundTripper
	server            *http.Server
	handlerRegistered bool
	middleware        []Middleware

	isRetriableFunc IsRetriable
}

func New(opts ...Option) (*Protocol, error) {
	p := &Protocol{
		incoming: make(chan msgErr),
		Port:     -1,
	}
	if err := p.applyOptions(opts...); err != nil {
		return nil, err
	}

	if p.Client == nil {
		p.Client = http.DefaultClient
	}

	if p.roundTripper != nil {
		p.Client.Transport = p.roundTripper
	}

	if p.ShutdownTimeout == 0 {
		p.ShutdownTimeout = DefaultShutdownTimeout
	}

	if p.isRetriableFunc == nil {
		p.isRetriableFunc = defaultIsRetriableFunc
	}

	return p, nil
}

// NewObserved creates an HTTP protocol with trace propagating middleware.
// Deprecated: now this behaves like New and it will be removed in future releases,
// setup the http observed protocol using the opencensus separate module NewObservedHttp
var NewObserved = New

func (p *Protocol) applyOptions(opts ...Option) error {
	for _, fn := range opts {
		if err := fn(p); err != nil {
			return err
		}
	}
	return nil
}

// Send implements binding.Sender
func (p *Protocol) Send(ctx context.Context, m binding.Message, transformers ...binding.Transformer) error {
	if ctx == nil {
		return fmt.Errorf("nil Context")
	} else if m == nil {
		return fmt.Errorf("nil Message")
	}

	msg, err := p.Request(ctx, m, transformers...)
	if err != nil && !protocol.IsACK(err) {
		var res *Result
		if protocol.ResultAs(err, &res) {
			if message, ok := msg.(*Message); ok {
				buf := new(bytes.Buffer)
				buf.ReadFrom(message.BodyReader)
				errorStr := buf.String()
				err = NewResult(res.StatusCode, "%s", errorStr)
			}
		}
	}
	return err
}

// Request implements binding.Requester
func (p *Protocol) Request(ctx context.Context, m binding.Message, transformers ...binding.Transformer) (binding.Message, error) {
	if ctx == nil {
		return nil, fmt.Errorf("nil Context")
	} else if m == nil {
		return nil, fmt.Errorf("nil Message")
	}

	var err error
	defer func() { _ = m.Finish(err) }()

	req := p.makeRequest(ctx)

	if p.Client == nil || req == nil || req.URL == nil {
		return nil, fmt.Errorf("not initialized: %#v", p)
	}

	if err = WriteRequest(ctx, m, req, transformers...); err != nil {
		return nil, err
	}

	return p.do(ctx, req)
}

func (p *Protocol) makeRequest(ctx context.Context) *http.Request {
	req := &http.Request{
		Method: http.MethodPost,
		Header: HeaderFrom(ctx),
	}

	if p.RequestTemplate != nil {
		req.Method = p.RequestTemplate.Method
		req.URL = p.RequestTemplate.URL
		req.Close = p.RequestTemplate.Close
		req.Host = p.RequestTemplate.Host
		copyHeadersEnsure(p.RequestTemplate.Header, &req.Header)
	}

	if p.Target != nil {
		req.URL = p.Target
	}

	// Override the default request with target from context.
	if target := cecontext.TargetFrom(ctx); target != nil {
		req.URL = target
	}
	return req.WithContext(ctx)
}

// Ensure to is a non-nil map before copying
func copyHeadersEnsure(from http.Header, to *http.Header) {
	if len(from) > 0 {
		if *to == nil {
			*to = http.Header{}
		}
		copyHeaders(from, *to)
	}
}

func copyHeaders(from, to http.Header) {
	if from == nil || to == nil {
		return
	}
	for header, values := range from {
		for _, value := range values {
			to.Add(header, value)
		}
	}
}

// Receive the next incoming HTTP request as a CloudEvent.
// Returns non-nil error if the incoming HTTP request fails to parse as a CloudEvent
// Returns io.EOF if the receiver is closed.
func (p *Protocol) Receive(ctx context.Context) (binding.Message, error) {
	if ctx == nil {
		return nil, fmt.Errorf("nil Context")
	}

	msg, fn, err := p.Respond(ctx)
	// No-op the response when finish is invoked.
	if msg != nil {
		return binding.WithFinish(msg, func(err error) {
			if fn != nil {
				_ = fn(ctx, nil, nil)
			}
		}), err
	} else {
		return nil, err
	}
}

// Respond receives the next incoming HTTP request as a CloudEvent and waits
// for the response callback to invoked before continuing.
// Returns non-nil error if the incoming HTTP request fails to parse as a CloudEvent
// Returns io.EOF if the receiver is closed.
func (p *Protocol) Respond(ctx context.Context) (binding.Message, protocol.ResponseFn, error) {
	if ctx == nil {
		return nil, nil, fmt.Errorf("nil Context")
	}

	select {
	case in, ok := <-p.incoming:
		if !ok {
			return nil, nil, io.EOF
		}

		if in.msg == nil {
			return nil, in.respFn, in.err
		}
		return in.msg, in.respFn, in.err

	case <-ctx.Done():
		return nil, nil, io.EOF
	}
}

// ServeHTTP implements http.Handler.
// Blocks until ResponseFn is invoked.
func (p *Protocol) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// Filter the GET style methods:
	switch req.Method {
	case http.MethodOptions:
		if p.OptionsHandlerFn == nil {
			rw.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		p.OptionsHandlerFn(rw, req)
		return

	case http.MethodGet:
		if p.GetHandlerFn == nil {
			rw.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		p.GetHandlerFn(rw, req)
		return

	case http.MethodDelete:
		if p.DeleteHandlerFn == nil {
			rw.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		p.DeleteHandlerFn(rw, req)
		return
	}

	m := NewMessageFromHttpRequest(req)
	if m == nil {
		// Should never get here unless ServeHTTP is called directly.
		p.incoming <- msgErr{msg: nil, err: binding.ErrUnknownEncoding}
		rw.WriteHeader(http.StatusBadRequest)
		return // if there was no message, return.
	}

	var finishErr error
	m.OnFinish = func(err error) error {
		finishErr = err
		return nil
	}

	wg := sync.WaitGroup{}
	wg.Add(1)
	var fn protocol.ResponseFn = func(ctx context.Context, respMsg binding.Message, res protocol.Result, transformers ...binding.Transformer) error {
		// Unblock the ServeHTTP after the reply is written
		defer func() {
			wg.Done()
		}()

		if finishErr != nil {
			http.Error(rw, fmt.Sprintf("Cannot forward CloudEvent: %s", finishErr), http.StatusInternalServerError)
			return finishErr
		}

		status := http.StatusOK
		if res != nil {
			var result *Result
			switch {
			case protocol.ResultAs(res, &result):
				if result.StatusCode > 100 && result.StatusCode < 600 {
					status = result.StatusCode
				}

			case !protocol.IsACK(res):
				// Map client errors to http status code
				validationError := event.ValidationError{}
				if errors.As(res, &validationError) {
					status = http.StatusBadRequest
					rw.Header().Set("content-type", "text/plain")
					rw.WriteHeader(status)
					_, _ = rw.Write([]byte(validationError.Error()))
					return validationError
				} else if errors.Is(res, binding.ErrUnknownEncoding) {
					status = http.StatusUnsupportedMediaType
				} else {
					status = http.StatusInternalServerError
				}
			}
		}

		if respMsg != nil {
			err := WriteResponseWriter(ctx, respMsg, status, rw, transformers...)
			return respMsg.Finish(err)
		}

		rw.WriteHeader(status)
		return nil
	}

	p.incoming <- msgErr{msg: m, respFn: fn} // Send to Request
	// Block until ResponseFn is invoked
	wg.Wait()
}

func defaultIsRetriableFunc(sc int) bool {
	_, ok := defaultRetriableErrors[sc]
	return ok
}
<commit_msg>Finish() no longer used binding-message in Send() (#689)<commit_after>/*
 Copyright 2021 The CloudEvents Authors
 SPDX-License-Identifier: Apache-2.0
*/

package http

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudevents/sdk-go/v2/binding"
	cecontext "github.com/cloudevents/sdk-go/v2/context"
	"github.com/cloudevents/sdk-go/v2/event"
	"github.com/cloudevents/sdk-go/v2/protocol"
)

const (
	// DefaultShutdownTimeout defines the default timeout given to the http.Server when calling Shutdown.
	DefaultShutdownTimeout = time.Minute * 1
)

type msgErr struct {
	msg    *Message
	respFn protocol.ResponseFn
	err    error
}

// Default error codes that we retry on - string isn't used, it's just there so
// people know what each error code's title is.
// To modify this use Option
var defaultRetriableErrors = map[int]string{
	404: "Not Found",
	413: "Payload Too Large",
	425: "Too Early",
	429: "Too Many Requests",
	502: "Bad Gateway",
	503: "Service Unavailable",
	504: "Gateway Timeout",
}

// Protocol acts as both a http client and a http handler.
type Protocol struct {
	Target          *url.URL
	RequestTemplate *http.Request
	Client          *http.Client
	incoming        chan msgErr

	// OptionsHandlerFn handles the OPTIONS method requests and is intended to
	// implement the abuse protection spec:
	// https://github.com/cloudevents/spec/blob/v1.0/http-webhook.md#4-abuse-protection
	OptionsHandlerFn http.HandlerFunc
	WebhookConfig    *WebhookConfig

	GetHandlerFn    http.HandlerFunc
	DeleteHandlerFn http.HandlerFunc

	// To support Opener:

	// ShutdownTimeout defines the timeout given to the http.Server when calling Shutdown.
	// If 0, DefaultShutdownTimeout is used.
	ShutdownTimeout time.Duration

	// Port is the port configured to bind the receiver to. Defaults to 8080.
	// If you want to know the effective port you're listening to, use GetListeningPort()
	Port int
	// Path is the path to bind the receiver to. Defaults to "/".
	Path string

	// Receive Mutex
	reMu sync.Mutex
	// Handler is the handler the http Server will use. Use this to reuse the
	// http server. If nil, the Protocol will create a one.
	Handler *http.ServeMux

	listener          atomic.Value
	roundTripper      http.RoundTripper
	server            *http.Server
	handlerRegistered bool
	middleware        []Middleware

	isRetriableFunc IsRetriable
}

func New(opts ...Option) (*Protocol, error) {
	p := &Protocol{
		incoming: make(chan msgErr),
		Port:     -1,
	}
	if err := p.applyOptions(opts...); err != nil {
		return nil, err
	}

	if p.Client == nil {
		p.Client = http.DefaultClient
	}

	if p.roundTripper != nil {
		p.Client.Transport = p.roundTripper
	}

	if p.ShutdownTimeout == 0 {
		p.ShutdownTimeout = DefaultShutdownTimeout
	}

	if p.isRetriableFunc == nil {
		p.isRetriableFunc = defaultIsRetriableFunc
	}

	return p, nil
}

// NewObserved creates an HTTP protocol with trace propagating middleware.
// Deprecated: now this behaves like New and it will be removed in future releases,
// setup the http observed protocol using the opencensus separate module NewObservedHttp
var NewObserved = New

func (p *Protocol) applyOptions(opts ...Option) error {
	for _, fn := range opts {
		if err := fn(p); err != nil {
			return err
		}
	}
	return nil
}

// Send implements binding.Sender
func (p *Protocol) Send(ctx context.Context, m binding.Message, transformers ...binding.Transformer) error {
	if ctx == nil {
		return fmt.Errorf("nil Context")
	} else if m == nil {
		return fmt.Errorf("nil Message")
	}

	msg, err := p.Request(ctx, m, transformers...)
	if msg != nil {
		defer func() { _ = msg.Finish(err) }()
	}
	if err != nil && !protocol.IsACK(err) {
		var res *Result
		if protocol.ResultAs(err, &res) {
			if message, ok := msg.(*Message); ok {
				buf := new(bytes.Buffer)
				buf.ReadFrom(message.BodyReader)
				errorStr := buf.String()
				err = NewResult(res.StatusCode, "%s", errorStr)
			}
		}
	}
	return err
}

// Request implements binding.Requester
func (p *Protocol) Request(ctx context.Context, m binding.Message, transformers ...binding.Transformer) (binding.Message, error) {
	if ctx == nil {
		return nil, fmt.Errorf("nil Context")
	} else if m == nil {
		return nil, fmt.Errorf("nil Message")
	}

	var err error
	defer func() { _ = m.Finish(err) }()

	req := p.makeRequest(ctx)

	if p.Client == nil || req == nil || req.URL == nil {
		return nil, fmt.Errorf("not initialized: %#v", p)
	}

	if err = WriteRequest(ctx, m, req, transformers...); err != nil {
		return nil, err
	}

	return p.do(ctx, req)
}

func (p *Protocol) makeRequest(ctx context.Context) *http.Request {
	req := &http.Request{
		Method: http.MethodPost,
		Header: HeaderFrom(ctx),
	}

	if p.RequestTemplate != nil {
		req.Method = p.RequestTemplate.Method
		req.URL = p.RequestTemplate.URL
		req.Close = p.RequestTemplate.Close
		req.Host = p.RequestTemplate.Host
		copyHeadersEnsure(p.RequestTemplate.Header, &req.Header)
	}

	if p.Target != nil {
		req.URL = p.Target
	}

	// Override the default request with target from context.
	if target := cecontext.TargetFrom(ctx); target != nil {
		req.URL = target
	}
	return req.WithContext(ctx)
}

// Ensure to is a non-nil map before copying
func copyHeadersEnsure(from http.Header, to *http.Header) {
	if len(from) > 0 {
		if *to == nil {
			*to = http.Header{}
		}
		copyHeaders(from, *to)
	}
}

func copyHeaders(from, to http.Header) {
	if from == nil || to == nil {
		return
	}
	for header, values := range from {
		for _, value := range values {
			to.Add(header, value)
		}
	}
}

// Receive the next incoming HTTP request as a CloudEvent.
// Returns non-nil error if the incoming HTTP request fails to parse as a CloudEvent
// Returns io.EOF if the receiver is closed.
func (p *Protocol) Receive(ctx context.Context) (binding.Message, error) {
	if ctx == nil {
		return nil, fmt.Errorf("nil Context")
	}

	msg, fn, err := p.Respond(ctx)
	// No-op the response when finish is invoked.
	if msg != nil {
		return binding.WithFinish(msg, func(err error) {
			if fn != nil {
				_ = fn(ctx, nil, nil)
			}
		}), err
	} else {
		return nil, err
	}
}

// Respond receives the next incoming HTTP request as a CloudEvent and waits
// for the response callback to invoked before continuing.
// Returns non-nil error if the incoming HTTP request fails to parse as a CloudEvent
// Returns io.EOF if the receiver is closed.
func (p *Protocol) Respond(ctx context.Context) (binding.Message, protocol.ResponseFn, error) {
	if ctx == nil {
		return nil, nil, fmt.Errorf("nil Context")
	}

	select {
	case in, ok := <-p.incoming:
		if !ok {
			return nil, nil, io.EOF
		}

		if in.msg == nil {
			return nil, in.respFn, in.err
		}
		return in.msg, in.respFn, in.err

	case <-ctx.Done():
		return nil, nil, io.EOF
	}
}

// ServeHTTP implements http.Handler.
// Blocks until ResponseFn is invoked.
func (p *Protocol) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// Filter the GET style methods:
	switch req.Method {
	case http.MethodOptions:
		if p.OptionsHandlerFn == nil {
			rw.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		p.OptionsHandlerFn(rw, req)
		return

	case http.MethodGet:
		if p.GetHandlerFn == nil {
			rw.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		p.GetHandlerFn(rw, req)
		return

	case http.MethodDelete:
		if p.DeleteHandlerFn == nil {
			rw.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		p.DeleteHandlerFn(rw, req)
		return
	}

	m := NewMessageFromHttpRequest(req)
	if m == nil {
		// Should never get here unless ServeHTTP is called directly.
		p.incoming <- msgErr{msg: nil, err: binding.ErrUnknownEncoding}
		rw.WriteHeader(http.StatusBadRequest)
		return // if there was no message, return.
	}

	var finishErr error
	m.OnFinish = func(err error) error {
		finishErr = err
		return nil
	}

	wg := sync.WaitGroup{}
	wg.Add(1)
	var fn protocol.ResponseFn = func(ctx context.Context, respMsg binding.Message, res protocol.Result, transformers ...binding.Transformer) error {
		// Unblock the ServeHTTP after the reply is written
		defer func() {
			wg.Done()
		}()

		if finishErr != nil {
			http.Error(rw, fmt.Sprintf("Cannot forward CloudEvent: %s", finishErr), http.StatusInternalServerError)
			return finishErr
		}

		status := http.StatusOK
		if res != nil {
			var result *Result
			switch {
			case protocol.ResultAs(res, &result):
				if result.StatusCode > 100 && result.StatusCode < 600 {
					status = result.StatusCode
				}

			case !protocol.IsACK(res):
				// Map client errors to http status code
				validationError := event.ValidationError{}
				if errors.As(res, &validationError) {
					status = http.StatusBadRequest
					rw.Header().Set("content-type", "text/plain")
					rw.WriteHeader(status)
					_, _ = rw.Write([]byte(validationError.Error()))
					return validationError
				} else if errors.Is(res, binding.ErrUnknownEncoding) {
					status = http.StatusUnsupportedMediaType
				} else {
					status = http.StatusInternalServerError
				}
			}
		}

		if respMsg != nil {
			err := WriteResponseWriter(ctx, respMsg, status, rw, transformers...)
			return respMsg.Finish(err)
		}

		rw.WriteHeader(status)
		return nil
	}

	p.incoming <- msgErr{msg: m, respFn: fn} // Send to Request
	// Block until ResponseFn is invoked
	wg.Wait()
}

func defaultIsRetriableFunc(sc int) bool {
	_, ok := defaultRetriableErrors[sc]
	return ok
}
<|endoftext|>


<commit_before>package libproxy

import (
	"io"
	"net"
	"syscall"

	"github.com/Sirupsen/logrus"
)

// TCPProxy is a proxy for TCP connections. It implements the Proxy interface to
// handle TCP traffic forwarding between the frontend and backend addresses.
type TCPProxy struct {
	listener     *net.TCPListener
	frontendAddr *net.TCPAddr
	backendAddr  *net.TCPAddr
}

// NewTCPProxy creates a new TCPProxy.
func NewTCPProxy(frontendAddr, backendAddr *net.TCPAddr) (*TCPProxy, error) {
	listener, err := net.ListenTCP("tcp", frontendAddr)
	if err != nil {
		return nil, err
	}
	// If the port in frontendAddr was 0 then ListenTCP will have a picked
	// a port to listen on, hence the call to Addr to get that actual port:
	return &TCPProxy{
		listener:     listener,
		frontendAddr: listener.Addr().(*net.TCPAddr),
		backendAddr:  backendAddr,
	}, nil
}

func (proxy *TCPProxy) clientLoop(client *net.TCPConn, quit chan bool) {
	backend, err := net.DialTCP("tcp", nil, proxy.backendAddr)
	if err != nil {
		logrus.Printf("Can't forward traffic to backend tcp/%v: %s\n", proxy.backendAddr, err)
		client.Close()
		return
	}

	event := make(chan int64)
	var broker = func(to, from *net.TCPConn) {
		written, err := io.Copy(to, from)
		if err != nil {
			// If the socket we are writing to is shutdown with
			// SHUT_WR, forward it to the other end of the pipe:
			if err, ok := err.(*net.OpError); ok && err.Err == syscall.EPIPE {
				from.CloseWrite()
			}
		}
		to.CloseRead()
		event <- written
	}

	go broker(client, backend)
	go broker(backend, client)

	var transferred int64
	for i := 0; i < 2; i++ {
		select {
		case written := <-event:
			transferred += written
		case <-quit:
			// Interrupt the two brokers and "join" them.
			client.Close()
			backend.Close()
			for ; i < 2; i++ {
				transferred += <-event
			}
			return
		}
	}
	client.Close()
	backend.Close()
}

// Run starts forwarding the traffic using TCP.
func (proxy *TCPProxy) Run() {
	quit := make(chan bool)
	defer close(quit)
	for {
		client, err := proxy.listener.Accept()
		if err != nil {
			logrus.Printf("Stopping proxy on tcp/%v for tcp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err)
			return
		}
		go proxy.clientLoop(client.(*net.TCPConn), quit)
	}
}

// Close stops forwarding the traffic.
func (proxy *TCPProxy) Close() { proxy.listener.Close() }

// FrontendAddr returns the TCP address on which the proxy is listening.
func (proxy *TCPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr }

// BackendAddr returns the TCP proxied address.
func (proxy *TCPProxy) BackendAddr() net.Addr { return proxy.backendAddr }
<commit_msg>proxy: TCPProxy connects to TCP but can read from other things<commit_after>package libproxy

import (
	"io"
	"net"
	"syscall"

	"github.com/Sirupsen/logrus"
)

type Conn interface {
	io.Reader
	io.Writer
	io.Closer
	CloseRead() error
	CloseWrite() error
}

// TCPProxy is a proxy for TCP connections. It implements the Proxy interface to
// handle TCP traffic forwarding between the frontend and backend addresses.
type TCPProxy struct {
	listener     *net.TCPListener
	frontendAddr *net.TCPAddr
	backendAddr  *net.TCPAddr
}

// NewTCPProxy creates a new TCPProxy.
func NewTCPProxy(frontendAddr, backendAddr *net.TCPAddr) (*TCPProxy, error) {
	listener, err := net.ListenTCP("tcp", frontendAddr)
	if err != nil {
		return nil, err
	}
	// If the port in frontendAddr was 0 then ListenTCP will have a picked
	// a port to listen on, hence the call to Addr to get that actual port:
	return &TCPProxy{
		listener:     listener,
		frontendAddr: listener.Addr().(*net.TCPAddr),
		backendAddr:  backendAddr,
	}, nil
}

func (proxy *TCPProxy) clientLoop(client Conn, quit chan bool) {
	backend, err := net.DialTCP("tcp", nil, proxy.backendAddr)
	if err != nil {
		logrus.Printf("Can't forward traffic to backend tcp/%v: %s\n", proxy.backendAddr, err)
		client.Close()
		return
	}

	event := make(chan int64)
	var broker = func(to, from Conn) {
		written, err := io.Copy(to, from)
		if err != nil {
			// If the socket we are writing to is shutdown with
			// SHUT_WR, forward it to the other end of the pipe:
			if err, ok := err.(*net.OpError); ok && err.Err == syscall.EPIPE {
				from.CloseWrite()
			}
		}
		to.CloseRead()
		event <- written
	}

	go broker(client, backend)
	go broker(backend, client)

	var transferred int64
	for i := 0; i < 2; i++ {
		select {
		case written := <-event:
			transferred += written
		case <-quit:
			// Interrupt the two brokers and "join" them.
			client.Close()
			backend.Close()
			for ; i < 2; i++ {
				transferred += <-event
			}
			return
		}
	}
	client.Close()
	backend.Close()
}

// Run starts forwarding the traffic using TCP.
func (proxy *TCPProxy) Run() {
	quit := make(chan bool)
	defer close(quit)
	for {
		client, err := proxy.listener.Accept()
		if err != nil {
			logrus.Printf("Stopping proxy on tcp/%v for tcp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err)
			return
		}
		go proxy.clientLoop(client.(Conn), quit)
	}
}

// Close stops forwarding the traffic.
func (proxy *TCPProxy) Close() { proxy.listener.Close() }

// FrontendAddr returns the TCP address on which the proxy is listening.
func (proxy *TCPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr }

// BackendAddr returns the TCP proxied address.
func (proxy *TCPProxy) BackendAddr() net.Addr { return proxy.backendAddr }
<|endoftext|>


<commit_before>package goat

import (
	"bytes"
	"net/http/httptest"
	"testing"
)

func TestWriteError(t *testing.T) {
	// In
	code := 500
	err := "foo"

	// Expected
	json := `{
  "error": "` + err + `"
}
`
	buf := bytes.NewBufferString(json)

	w := httptest.NewRecorder()
	WriteError(w, code, err)

	// Test code
	if w.Code != code {
		t.Errorf("WriteError should set Code to %i, but did set it to %i", code, w.Code)
	}

	// Test body
	if w.Body == nil {
		t.Errorf("WriteError should set Body to %s, but didn't", json)
	} else if !bytes.Equal(w.Body.Bytes(), buf.Bytes()) {
		t.Errorf("WriteError should set Body to %v, but did set it to %v", buf, w.Body)
	}
}

func TestWriteJSON(t *testing.T) {
	in := map[string]string{
		"foo":   "bar",
		"knock": "knock",
	}
	json := `{
  "foo": "bar",
  "knock": "knock"
}
`
	buf := bytes.NewBufferString(json)

	w := httptest.NewRecorder()
	WriteJSON(w, in)

	if w.Body == nil {
		t.Errorf("WriteJSON should set the Body to %s, but didn't", json)
	} else if !bytes.Equal(w.Body.Bytes(), buf.Bytes()) {
		t.Errorf("WriteJSON set the Body to %v, but should set it to %v", buf, w.Body)
	}

	w = httptest.NewRecorder()
	if err := WriteJSON(w, nil); err != nil {
		t.Errorf("WriteJSON should return an error, but didn't")
	}
}
<commit_msg>Switch to strings for comparison<commit_after>package goat

import (
	"bytes"
	"net/http/httptest"
	"testing"
)

func TestWriteError(t *testing.T) {
	// In
	code := 500
	err := "foo"

	// Expected
	json := `{
  "error": "` + err + `"
}
`
	buf := bytes.NewBufferString(json)

	w := httptest.NewRecorder()
	WriteError(w, code, err)

	// Test code
	if w.Code != code {
		t.Errorf("WriteError should set Code to %i, but did set it to %i", code, w.Code)
	}

	// Test body
	if w.Body == nil {
		t.Errorf("WriteError should set Body to %s, but didn't", json)
	} else if string(w.Body.Bytes()) == string(buf.Bytes()) {
		t.Errorf("WriteError should set Body to %v, but did set it to %v", buf, w.Body)
	}
}

func TestWriteJSON(t *testing.T) {
	in := map[string]string{
		"foo":   "bar",
		"knock": "knock",
	}
	json := `{
  "foo": "bar",
  "knock": "knock"
}
`
	buf := bytes.NewBufferString(json)

	w := httptest.NewRecorder()
	WriteJSON(w, in)

	if w.Body == nil {
		t.Errorf("WriteJSON should set the Body to %s, but didn't", json)
	} else if string(w.Body.Bytes()) == string(buf.Bytes()) {
		t.Errorf("WriteJSON set the Body to %v, but should set it to %v", buf, w.Body)
	}

	// Test error
	w = httptest.NewRecorder()
	if err := WriteJSON(w, nil); err != nil {
		t.Errorf("WriteJSON should return an error, but didn't")
	}
}
<|endoftext|>


<commit_before>package goka

import (
	"fmt"
	"sync"

	"github.com/facebookgo/ensure"
	"github.com/golang/mock/gomock"
	"github.com/golang/protobuf/proto"
	metrics "github.com/rcrowley/go-metrics"

	"github.com/lovoo/goka/codec"
	"github.com/lovoo/goka/kafka"
	"github.com/lovoo/goka/storage"
)

// EmitHandler abstracts a function that allows to overwrite kafkamock's Emit function to
// simulate producer errors
type EmitHandler func(topic string, key string, value []byte) *kafka.Promise

type gomockPanicker struct {
	reporter gomock.TestReporter
}

func (gm *gomockPanicker) Errorf(format string, args ...interface{}) {
	gm.reporter.Errorf(format, args...)
}
func (gm *gomockPanicker) Fatalf(format string, args ...interface{}) {
	defer panic(fmt.Sprintf(format, args...))
	gm.reporter.Fatalf(format, args...)
}

// NewMockController returns a *gomock.Controller using a wrapped testing.T (or whatever)
// which panics on a Fatalf. This is necessary when using a mock in kafkamock. Otherwise it will
// freeze on an unexpected call.
func NewMockController(t gomock.TestReporter) *gomock.Controller {
	return gomock.NewController(&gomockPanicker{reporter: t})
}

// KafkaMock is allows interacting with a test processor
type KafkaMock struct {
	t       Tester
	storage storage.Storage

	offset         int64
	tableOffset    int64
	incomingEvents chan kafka.Event
	consumerEvents chan kafka.Event
	// Stores a map of all topics that are handled by the processor.
	// Every time an emit is called, those messages for handled topics are relayed
	// after the consume-function has finished.
	// All other messages are stored in the emitted-slice for further inspection
	handledTopics map[string]bool
	groupTopic    string
	emitted       []*kafka.Message

	groupTableCreator func() (string, []byte)
	callQueue         []func()
	wg                sync.WaitGroup

	consumerMock *consumerMock

	producerMock *producerMock
	emitHandler  EmitHandler
	topicMgrMock *topicMgrMock
}

// Tester abstracts the interface we assume from the test case.
// Will most likely be *testing.T
type Tester interface {
	Errorf(format string, args ...interface{})
	Fatalf(format string, args ...interface{})
	Fatal(a ...interface{})
}

// NewKafkaMock returns a new testprocessor mocking every external service
func NewKafkaMock(t Tester, groupName Group) *KafkaMock {
	kafkaMock := &KafkaMock{
		storage:        storage.NewMock(new(codec.Bytes)),
		t:              t,
		incomingEvents: make(chan kafka.Event),
		consumerEvents: make(chan kafka.Event),
		handledTopics:  make(map[string]bool),
		groupTopic:     tableName(groupName),
	}
	kafkaMock.consumerMock = newConsumerMock(kafkaMock)
	kafkaMock.producerMock = newProducerMock(kafkaMock.handleEmit)
	kafkaMock.topicMgrMock = newTopicMgrMock(kafkaMock)

	return kafkaMock
}

func (km *KafkaMock) SetGroupTableCreator(creator func() (string, []byte)) {
	km.groupTableCreator = creator
}

// ProcessorOptions returns the options that must be passed to NewProcessor
// to use the Mock. It essentially replaces the consumer/producer/topicmanager with a mock.
// For convenience, the storage is also mocked.
// For example, a normal call to NewProcessor like this
//     NewProcessor(brokers, group, subscriptions,
//                       option_a,
//                       option_b,
//                       option_c,
//     )
// would become in the unit test:
// kafkaMock := NewKafkaMock(t)
// NewProcessor(brokers, group, subscriptions,
//                       append(kafkaMock.ProcessorOptions(),
//                       option_a,
//                       option_b,
//                       option_c,
//                       )...,
//     )
func (km *KafkaMock) ProcessorOptions() []ProcessorOption {
	return []ProcessorOption{
		WithStorageBuilder(func(topic string, partition int32, c Codec, reg metrics.Registry) (storage.Storage, error) {
			return km.storage, nil
		}),
		WithConsumer(km.consumerMock),
		WithProducer(km.producerMock),
		WithTopicManager(km.topicMgrMock),
		WithPartitionChannelSize(0),
	}
}

// initProtocol initiates the protocol with the client basically making the KafkaMock
// usable.
func (km *KafkaMock) initProtocol() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("recovered from panic")
		}
	}()
	km.consumerEvents <- &kafka.Assignment{
		0: -1,
	}

	for km.groupTableCreator != nil {
		key, value := km.groupTableCreator()
		if key == "" || value == nil {
			break
		}
		km.consumerEvents <- &kafka.Message{
			Topic:     km.groupTopic,
			Partition: 0,
			Offset:    km.tableOffset,
			Key:       key,
			Value:     value,
		}
	}

	km.consumerEvents <- &kafka.EOF{Partition: 0}
	km.consumerEvents <- &kafka.NOP{Partition: -1}

	for ev := range km.incomingEvents {
		km.consumerEvents <- ev
	}
}

// ConsumeProto simulates a message on kafka in a topic with a key.
func (km *KafkaMock) ConsumeProto(topic string, key string, msg proto.Message) {
	data, err := proto.Marshal(msg)
	if err != nil {
		if km.t != nil {
			km.t.Errorf("Error marshaling message for consume: %v", err)
		}
	}

	km.consumeData(topic, key, data)
	km.makeCalls()
}

func (km *KafkaMock) ConsumeString(topic string, key string, msg string) {
	km.consumeData(topic, key, []byte(msg))
	km.makeCalls()
}

// Helper function consuming marshalled data. This function is used by ConsumeProto by the test case
// as well as any emit calls of the processor being tested.
func (km *KafkaMock) consumeData(topic string, key string, data []byte) {
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("recovered from panic: %+v\n", r)
		}
	}()
	km.offset++
	kafkaMsg := &kafka.Message{
		Topic:     topic,
		Partition: 0,
		Offset:    km.offset,

		Key:   key,
		Value: data,
	}
	// send message to processing goroutine
	km.incomingEvents <- kafkaMsg
	// wait until partition processing goroutine processes message by requiring it to read
	// the following NOP.
	km.incomingEvents <- &kafka.NOP{Partition: 0}
	km.incomingEvents <- &kafka.NOP{Partition: 0}

	// wait util processor goroutine is ready
	km.incomingEvents <- &kafka.NOP{Partition: -1}
	km.incomingEvents <- &kafka.NOP{Partition: -1}
}

func (km *KafkaMock) consumeError(err error) {
	km.incomingEvents <- &kafka.Error{Err: err}
	// no need to send NOP (actuallly we can't, otherwise we might panic
	// as the channels are already closed due to the error first).
}

// ValueForKey attempts to get a value from KafkaMock's storage.
func (km *KafkaMock) ValueForKey(key string) interface{} {
	item, err := km.storage.Get(key)
	ensure.Nil(km.t, err)
	return item
}

// SetValue sets a value in the storage.
func (km *KafkaMock) SetValue(key string, value interface{}) {
	err := km.storage.Set(key, value)
	ensure.Nil(km.t, err)
}

func (km *KafkaMock) ReplaceEmitHandler(emitter EmitHandler) {
	km.producerMock.emitter = emitter
}

// ExpectEmit ensures a message exists in passed topic and key. The message may be
// inspected/unmarshalled by a passed expecter function.
func (km *KafkaMock) ExpectEmit(topic string, key string, expecter func(value []byte)) {
	for i := len(km.emitted) - 1; i >= 0; i-- {
		msg := km.emitted[i]
		if msg.Topic != topic || msg.Key != key {
			continue
		}
		if expecter != nil {
			expecter(msg.Value)
		}
		// remove element from slice
		// https://github.com/golang/go/wiki/SliceTricks
		km.emitted = append(km.emitted[:i], km.emitted[i+1:]...)
		return
	}

	km.t.Errorf("Expected emit for key %s in topic %s was not present.", key, topic)
}

// Finish marks the kafkamock that there is no emit to be expected.
// Set @param fail to true, if kafkamock is supposed to fail the test case in case
// of remaining emits.
// Clears the list of emits either case.
// This should always be called at the end of a test case to make sure
// no emits of prior test cases are stuck in the list and mess with the test results.
func (km *KafkaMock) Finish(fail bool) {
	if len(km.emitted) > 0 {
		if fail {
			km.t.Errorf("The following emits are still in the list, although it's supposed to be empty:")
			for _, emitted := range km.emitted {
				km.t.Errorf("    topic: %s key: %s", emitted.Topic, emitted.Key)
			}
		}
	}
	km.emitted = make([]*kafka.Message, 0)
}

// handleEmit handles an Emit-call on the producerMock.
// This takes care of queueing calls
// to handled topics or putting the emitted messages in the emitted-messages-list
func (km *KafkaMock) handleEmit(topic string, key string, value []byte) *kafka.Promise {
	promise := kafka.NewPromise()
	if topic == km.groupTopic {
		return promise.Finish(nil)
	}
	if _, hasTopic := km.handledTopics[topic]; hasTopic {
		km.newCall(func() {
			km.consumeData(topic, key, value)
		})
	} else {
		km.offset++
		km.emitted = append(km.emitted, &kafka.Message{
			Topic:  topic,
			Key:    key,
			Value:  value,
			Offset: km.offset,
		})
	}
	return promise.Finish(nil)
}

// creates a new call being executed after the consume function has run.
func (km *KafkaMock) newCall(call func()) {
	km.wg.Add(1)
	km.callQueue = append(km.callQueue, call)
}

// executes all calls on the call queue.
// Executing calls may put new calls on the queue (if they emit something),
// so this function executes until no further calls are being made.
func (km *KafkaMock) makeCalls() {
	go func() {
		for len(km.callQueue) > 0 {
			call := km.callQueue[0]
			call()
			km.callQueue = km.callQueue[1:]
			km.wg.Done()
		}
	}()
	km.wg.Wait()
}

type consumerMock struct {
	kafkaMock *KafkaMock
}

func newConsumerMock(kafkaMock *KafkaMock) *consumerMock {
	return &consumerMock{
		kafkaMock: kafkaMock,
	}
}

// Events returns the event channel of the consumer mock
func (km *consumerMock) Events() <-chan kafka.Event {
	return km.kafkaMock.consumerEvents
}

// Subscribe marks the consumer to subscribe to passed topics.
// The consumerMock simply marks the topics as handled to make sure to
// pass emitted messages back to the processor.
func (km *consumerMock) Subscribe(topics map[string]int64) error {
	for topic := range topics {
		km.kafkaMock.handledTopics[topic] = true
	}
	go km.kafkaMock.initProtocol()
	return nil
}

// AddGroupPartition adds a partition for group consumption.
// No action required in the mock.
func (km *consumerMock) AddGroupPartition(partition int32) {
}

// Commit commits an offest.
// No action required in the mock.
func (km *consumerMock) Commit(topic string, partition int32, offset int64) error {
	return nil
}

// AddPartition marks the topic as a state topic.
// The mock has to know the state topic to ignore emit calls (which would never be consumed)
func (km *consumerMock) AddPartition(topic string, partition int32, initialOffset int64) {
}

// RemovePartition removes a partition from a topic.
// No action required in the mock.
func (km *consumerMock) RemovePartition(topic string, partition int32) {
}

// Close closes the consumer.
// No action required in the mock.
func (km *consumerMock) Close() error {
	close(km.kafkaMock.incomingEvents)
	close(km.kafkaMock.consumerEvents)
	fmt.Println("closed consumer mock")
	return nil
}

type topicMgrMock struct {
	kafkaMock *KafkaMock
}

// EnsureTableExists checks that a table (log-compacted topic) exists, or create one if possible
func (tm *topicMgrMock) EnsureTableExists(topic string, npar int) error {
	return nil
}

// EnsureStreamExists checks that a stream topic exists, or create one if possible
func (tm *topicMgrMock) EnsureStreamExists(topic string, npar int) error {
	return nil
}

// Partitions returns the number of partitions of a topic, that are assigned to the running
// instance, i.e. it doesn't represent all partitions of a topic.
func (tm *topicMgrMock) Partitions(topic string) ([]int32, error) {
	tm.kafkaMock.handledTopics[topic] = true
	return []int32{0}, nil
}

// Close closes the topic manager.
// No action required in the mock.
func (tm *topicMgrMock) Close() error {
	fmt.Println("closing topic manager")
	return nil
}

func newTopicMgrMock(kafkaMock *KafkaMock) *topicMgrMock {
	return &topicMgrMock{
		kafkaMock: kafkaMock,
	}
}

type producerMock struct {
	emitter EmitHandler
}

func newProducerMock(emitter EmitHandler) *producerMock {
	return &producerMock{
		emitter: emitter,
	}
}

// Emit emits messages to arbitrary topics.
// The mock simply forwards the emit to the KafkaMock which takes care of queueing calls
// to handled topics or putting the emitted messages in the emitted-messages-list
func (p *producerMock) Emit(topic string, key string, value []byte) *kafka.Promise {
	return p.emitter(topic, key, value)
}

// Close closes the producer mock
// No action required in the mock.
func (p *producerMock) Close() error {
	fmt.Println("Closing producer mock")
	return nil
}
<commit_msg>expect emit in correct order. Added expectall-function<commit_after>package goka

import (
	"fmt"
	"sync"

	"github.com/facebookgo/ensure"
	"github.com/golang/mock/gomock"
	"github.com/golang/protobuf/proto"
	metrics "github.com/rcrowley/go-metrics"

	"github.com/lovoo/goka/codec"
	"github.com/lovoo/goka/kafka"
	"github.com/lovoo/goka/storage"
)

// EmitHandler abstracts a function that allows to overwrite kafkamock's Emit function to
// simulate producer errors
type EmitHandler func(topic string, key string, value []byte) *kafka.Promise

type gomockPanicker struct {
	reporter gomock.TestReporter
}

func (gm *gomockPanicker) Errorf(format string, args ...interface{}) {
	gm.reporter.Errorf(format, args...)
}
func (gm *gomockPanicker) Fatalf(format string, args ...interface{}) {
	defer panic(fmt.Sprintf(format, args...))
	gm.reporter.Fatalf(format, args...)
}

// NewMockController returns a *gomock.Controller using a wrapped testing.T (or whatever)
// which panics on a Fatalf. This is necessary when using a mock in kafkamock. Otherwise it will
// freeze on an unexpected call.
func NewMockController(t gomock.TestReporter) *gomock.Controller {
	return gomock.NewController(&gomockPanicker{reporter: t})
}

// KafkaMock is allows interacting with a test processor
type KafkaMock struct {
	t       Tester
	storage storage.Storage

	offset         int64
	tableOffset    int64
	incomingEvents chan kafka.Event
	consumerEvents chan kafka.Event
	// Stores a map of all topics that are handled by the processor.
	// Every time an emit is called, those messages for handled topics are relayed
	// after the consume-function has finished.
	// All other messages are stored in the emitted-slice for further inspection
	handledTopics map[string]bool
	groupTopic    string
	emitted       []*kafka.Message

	groupTableCreator func() (string, []byte)
	callQueue         []func()
	wg                sync.WaitGroup

	consumerMock *consumerMock

	producerMock *producerMock
	emitHandler  EmitHandler
	topicMgrMock *topicMgrMock
}

// Tester abstracts the interface we assume from the test case.
// Will most likely be *testing.T
type Tester interface {
	Errorf(format string, args ...interface{})
	Fatalf(format string, args ...interface{})
	Fatal(a ...interface{})
}

// NewKafkaMock returns a new testprocessor mocking every external service
func NewKafkaMock(t Tester, groupName Group) *KafkaMock {
	kafkaMock := &KafkaMock{
		storage:        storage.NewMock(new(codec.Bytes)),
		t:              t,
		incomingEvents: make(chan kafka.Event),
		consumerEvents: make(chan kafka.Event),
		handledTopics:  make(map[string]bool),
		groupTopic:     tableName(groupName),
	}
	kafkaMock.consumerMock = newConsumerMock(kafkaMock)
	kafkaMock.producerMock = newProducerMock(kafkaMock.handleEmit)
	kafkaMock.topicMgrMock = newTopicMgrMock(kafkaMock)

	return kafkaMock
}

func (km *KafkaMock) SetGroupTableCreator(creator func() (string, []byte)) {
	km.groupTableCreator = creator
}

// ProcessorOptions returns the options that must be passed to NewProcessor
// to use the Mock. It essentially replaces the consumer/producer/topicmanager with a mock.
// For convenience, the storage is also mocked.
// For example, a normal call to NewProcessor like this
//     NewProcessor(brokers, group, subscriptions,
//                       option_a,
//                       option_b,
//                       option_c,
//     )
// would become in the unit test:
// kafkaMock := NewKafkaMock(t)
// NewProcessor(brokers, group, subscriptions,
//                       append(kafkaMock.ProcessorOptions(),
//                       option_a,
//                       option_b,
//                       option_c,
//                       )...,
//     )
func (km *KafkaMock) ProcessorOptions() []ProcessorOption {
	return []ProcessorOption{
		WithStorageBuilder(func(topic string, partition int32, c Codec, reg metrics.Registry) (storage.Storage, error) {
			return km.storage, nil
		}),
		WithConsumer(km.consumerMock),
		WithProducer(km.producerMock),
		WithTopicManager(km.topicMgrMock),
		WithPartitionChannelSize(0),
	}
}

// initProtocol initiates the protocol with the client basically making the KafkaMock
// usable.
func (km *KafkaMock) initProtocol() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("recovered from panic")
		}
	}()
	km.consumerEvents <- &kafka.Assignment{
		0: -1,
	}

	for km.groupTableCreator != nil {
		key, value := km.groupTableCreator()
		if key == "" || value == nil {
			break
		}
		km.consumerEvents <- &kafka.Message{
			Topic:     km.groupTopic,
			Partition: 0,
			Offset:    km.tableOffset,
			Key:       key,
			Value:     value,
		}
	}

	km.consumerEvents <- &kafka.EOF{Partition: 0}
	km.consumerEvents <- &kafka.NOP{Partition: -1}

	for ev := range km.incomingEvents {
		km.consumerEvents <- ev
	}
}

// ConsumeProto simulates a message on kafka in a topic with a key.
func (km *KafkaMock) ConsumeProto(topic string, key string, msg proto.Message) {
	data, err := proto.Marshal(msg)
	if err != nil {
		if km.t != nil {
			km.t.Errorf("Error marshaling message for consume: %v", err)
		}
	}

	km.consumeData(topic, key, data)
	km.makeCalls()
}

func (km *KafkaMock) ConsumeString(topic string, key string, msg string) {
	km.consumeData(topic, key, []byte(msg))
	km.makeCalls()
}

// Helper function consuming marshalled data. This function is used by ConsumeProto by the test case
// as well as any emit calls of the processor being tested.
func (km *KafkaMock) consumeData(topic string, key string, data []byte) {
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("recovered from panic: %+v\n", r)
		}
	}()
	km.offset++
	kafkaMsg := &kafka.Message{
		Topic:     topic,
		Partition: 0,
		Offset:    km.offset,

		Key:   key,
		Value: data,
	}
	// send message to processing goroutine
	km.incomingEvents <- kafkaMsg
	// wait until partition processing goroutine processes message by requiring it to read
	// the following NOP.
	km.incomingEvents <- &kafka.NOP{Partition: 0}
	km.incomingEvents <- &kafka.NOP{Partition: 0}

	// wait util processor goroutine is ready
	km.incomingEvents <- &kafka.NOP{Partition: -1}
	km.incomingEvents <- &kafka.NOP{Partition: -1}
}

func (km *KafkaMock) consumeError(err error) {
	km.incomingEvents <- &kafka.Error{Err: err}
	// no need to send NOP (actuallly we can't, otherwise we might panic
	// as the channels are already closed due to the error first).
}

// ValueForKey attempts to get a value from KafkaMock's storage.
func (km *KafkaMock) ValueForKey(key string) interface{} {
	item, err := km.storage.Get(key)
	ensure.Nil(km.t, err)
	return item
}

// SetValue sets a value in the storage.
func (km *KafkaMock) SetValue(key string, value interface{}) {
	err := km.storage.Set(key, value)
	ensure.Nil(km.t, err)
}

func (km *KafkaMock) ReplaceEmitHandler(emitter EmitHandler) {
	km.producerMock.emitter = emitter
}

// ExpectEmit ensures a message exists in passed topic and key. The message may be
// inspected/unmarshalled by a passed expecter function.
func (km *KafkaMock) ExpectEmit(topic string, key string, expecter func(value []byte)) {
	for i := 0; i < len(km.emitted); i++ {
		msg := km.emitted[i]
		if msg.Topic != topic || msg.Key != key {
			continue
		}
		if expecter != nil {
			expecter(msg.Value)
		}
		// remove element from slice
		// https://github.com/golang/go/wiki/SliceTricks
		km.emitted = append(km.emitted[:i], km.emitted[i+1:]...)
		return
	}

	km.t.Errorf("Expected emit for key %s in topic %s was not present.", key, topic)
}

// ExpectAllEmitted calls passed expected-emit-handler function for all emitted values and clears the
// emitted values
func (km *KafkaMock) ExpectAllEmitted(handler func(topic string, key string, value []byte)) {
	for _, emitted := range km.emitted {
		handler(emitted.Topic, emitted.Key, emitted.Value)
	}
	km.emitted = make([]*kafka.Message, 0)
}

// Finish marks the kafkamock that there is no emit to be expected.
// Set @param fail to true, if kafkamock is supposed to fail the test case in case
// of remaining emits.
// Clears the list of emits either case.
// This should always be called at the end of a test case to make sure
// no emits of prior test cases are stuck in the list and mess with the test results.
func (km *KafkaMock) Finish(fail bool) {
	if len(km.emitted) > 0 {
		if fail {
			km.t.Errorf("The following emits are still in the list, although it's supposed to be empty:")
			for _, emitted := range km.emitted {
				km.t.Errorf("    topic: %s key: %s", emitted.Topic, emitted.Key)
			}
		}
	}
	km.emitted = make([]*kafka.Message, 0)
}

// handleEmit handles an Emit-call on the producerMock.
// This takes care of queueing calls
// to handled topics or putting the emitted messages in the emitted-messages-list
func (km *KafkaMock) handleEmit(topic string, key string, value []byte) *kafka.Promise {
	promise := kafka.NewPromise()
	if topic == km.groupTopic {
		return promise.Finish(nil)
	}
	if _, hasTopic := km.handledTopics[topic]; hasTopic {
		km.newCall(func() {
			km.consumeData(topic, key, value)
		})
	} else {
		km.offset++
		km.emitted = append(km.emitted, &kafka.Message{
			Topic:  topic,
			Key:    key,
			Value:  value,
			Offset: km.offset,
		})
	}
	return promise.Finish(nil)
}

// creates a new call being executed after the consume function has run.
func (km *KafkaMock) newCall(call func()) {
	km.wg.Add(1)
	km.callQueue = append(km.callQueue, call)
}

// executes all calls on the call queue.
// Executing calls may put new calls on the queue (if they emit something),
// so this function executes until no further calls are being made.
func (km *KafkaMock) makeCalls() {
	go func() {
		for len(km.callQueue) > 0 {
			call := km.callQueue[0]
			call()
			km.callQueue = km.callQueue[1:]
			km.wg.Done()
		}
	}()
	km.wg.Wait()
}

type consumerMock struct {
	kafkaMock *KafkaMock
}

func newConsumerMock(kafkaMock *KafkaMock) *consumerMock {
	return &consumerMock{
		kafkaMock: kafkaMock,
	}
}

// Events returns the event channel of the consumer mock
func (km *consumerMock) Events() <-chan kafka.Event {
	return km.kafkaMock.consumerEvents
}

// Subscribe marks the consumer to subscribe to passed topics.
// The consumerMock simply marks the topics as handled to make sure to
// pass emitted messages back to the processor.
func (km *consumerMock) Subscribe(topics map[string]int64) error {
	for topic := range topics {
		km.kafkaMock.handledTopics[topic] = true
	}
	go km.kafkaMock.initProtocol()
	return nil
}

// AddGroupPartition adds a partition for group consumption.
// No action required in the mock.
func (km *consumerMock) AddGroupPartition(partition int32) {
}

// Commit commits an offest.
// No action required in the mock.
func (km *consumerMock) Commit(topic string, partition int32, offset int64) error {
	return nil
}

// AddPartition marks the topic as a state topic.
// The mock has to know the state topic to ignore emit calls (which would never be consumed)
func (km *consumerMock) AddPartition(topic string, partition int32, initialOffset int64) {
}

// RemovePartition removes a partition from a topic.
// No action required in the mock.
func (km *consumerMock) RemovePartition(topic string, partition int32) {
}

// Close closes the consumer.
// No action required in the mock.
func (km *consumerMock) Close() error {
	close(km.kafkaMock.incomingEvents)
	close(km.kafkaMock.consumerEvents)
	fmt.Println("closed consumer mock")
	return nil
}

type topicMgrMock struct {
	kafkaMock *KafkaMock
}

// EnsureTableExists checks that a table (log-compacted topic) exists, or create one if possible
func (tm *topicMgrMock) EnsureTableExists(topic string, npar int) error {
	return nil
}

// EnsureStreamExists checks that a stream topic exists, or create one if possible
func (tm *topicMgrMock) EnsureStreamExists(topic string, npar int) error {
	return nil
}

// Partitions returns the number of partitions of a topic, that are assigned to the running
// instance, i.e. it doesn't represent all partitions of a topic.
func (tm *topicMgrMock) Partitions(topic string) ([]int32, error) {
	tm.kafkaMock.handledTopics[topic] = true
	return []int32{0}, nil
}

// Close closes the topic manager.
// No action required in the mock.
func (tm *topicMgrMock) Close() error {
	fmt.Println("closing topic manager")
	return nil
}

func newTopicMgrMock(kafkaMock *KafkaMock) *topicMgrMock {
	return &topicMgrMock{
		kafkaMock: kafkaMock,
	}
}

type producerMock struct {
	emitter EmitHandler
}

func newProducerMock(emitter EmitHandler) *producerMock {
	return &producerMock{
		emitter: emitter,
	}
}

// Emit emits messages to arbitrary topics.
// The mock simply forwards the emit to the KafkaMock which takes care of queueing calls
// to handled topics or putting the emitted messages in the emitted-messages-list
func (p *producerMock) Emit(topic string, key string, value []byte) *kafka.Promise {
	return p.emitter(topic, key, value)
}

// Close closes the producer mock
// No action required in the mock.
func (p *producerMock) Close() error {
	fmt.Println("Closing producer mock")
	return nil
}
<|endoftext|>


<commit_before>package sortedmap

import (
	"testing"
	"time"
)

func TestKeys(t *testing.T) {
	sm, _, err := newSortedMapFromRandRecords(300)
	if err != nil {
		t.Fatal(err)
	}
	i := 0
	keys := sm.Keys()
	for _, key := range keys {
		if key == nil {
			t.Fatal("Key's value is nil.")
		}
		i++
	}
	if i == 0 {
		t.Fatal("The returned slice was empty.")
	}
}

func TestBoundedKeys(t *testing.T) {
	sm, _, err := newSortedMapFromRandRecords(300)
	if err != nil {
		t.Fatal(err)
	}
	i := 0
	keys, ok := sm.BoundedKeys(time.Time{}, time.Now())
	if !ok {
		t.Fatal("No values fall between or are equal to the given bounds.")
	}
	for _, key := range keys {
		if key == nil {
			t.Fatal("Key's value is nil.")
		}
		i++
	}
	if i == 0 {
		t.Fatal("The returned slice was empty.")
	}
}<commit_msg>Added test for missing condition<commit_after>package sortedmap

import (
	"testing"
	"time"
)

func TestKeys(t *testing.T) {
	sm, _, err := newSortedMapFromRandRecords(300)
	if err != nil {
		t.Fatal(err)
	}
	i := 0
	keys := sm.Keys()
	for _, key := range keys {
		if key == nil {
			t.Fatal("Key's value is nil.")
		}
		i++
	}
	if i == 0 {
		t.Fatal("The returned slice was empty.")
	}
}

func TestBoundedKeys(t *testing.T) {
	sm, _, err := newSortedMapFromRandRecords(300)
	if err != nil {
		t.Fatal(err)
	}
	i := 0
	keys, ok := sm.BoundedKeys(time.Time{}, time.Now())
	if !ok {
		t.Fatal("No values fall between or are equal to the given bounds.")
	}
	for _, key := range keys {
		if key == nil {
			t.Fatal("Key's value is nil.")
		}
		i++
	}
	if i == 0 {
		t.Fatal("The returned slice was empty.")
	}
}

func TestBoundedKeysWithNoBoundsReturned(t *testing.T) {
	sm, _, err := newSortedMapFromRandRecords(300)
	if err != nil {
		t.Fatal(err)
	}
	if _, ok := sm.BoundedKeys(time.Now().Add(-1 * time.Second), time.Now()); ok {
		t.Fatal("Values fall between or are equal to the given bounds when it should not have returned bounds.")
	}
}<|endoftext|>


<commit_before>package kinesumer

import (
	"errors"
	"math/rand"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/remind101/kinesumer/checkpointers/empty"
	k "github.com/remind101/kinesumer/interface"
	"github.com/remind101/kinesumer/provisioners/empty"
)

type Kinesumer struct {
	Kinesis      k.Kinesis
	Checkpointer k.Checkpointer
	Provisioner  k.Provisioner
	Stream       string
	Options      *KinesumerOptions
	records      chan k.Record
	stop         chan Unit
	stopped      chan Unit
	nRunning     int
	rand         *rand.Rand
}

type KinesumerOptions struct {
	ListStreamsLimit    int64
	DescribeStreamLimit int64
	GetRecordsLimit     int64
	PollTime            int
	MaxShardWorkers     int
	Handlers            k.Handlers
}

var DefaultKinesumerOptions = KinesumerOptions{
	// These values are the hard limits set by Amazon

	ListStreamsLimit:    1000,
	DescribeStreamLimit: 10000,
	GetRecordsLimit:     10000,

	PollTime:        2000,
	MaxShardWorkers: 50,
	Handlers:        k.DefaultHandlers{},
}

func NewDefaultKinesumer(awsAccessKey, awsSecretKey, awsRegion, stream string) (*Kinesumer, error) {
	return NewKinesumer(
		kinesis.New(
			&aws.Config{
				Credentials: credentials.NewStaticCredentials(awsAccessKey, awsSecretKey, ""),
				Region:      &awsRegion,
			},
		),
		nil,
		nil,
		nil,
		stream,
		nil,
	)
}

func NewKinesumer(kinesis k.Kinesis, checkpointer k.Checkpointer, provisioner k.Provisioner,
	randSource rand.Source, stream string, opt *KinesumerOptions) (*Kinesumer, error) {

	if kinesis == nil {
		return nil, NewError(ECrit, "Kinesis object must not be nil", nil)
	}

	if checkpointer == nil {
		checkpointer = emptycheckpointer.Checkpointer{}
	}

	if provisioner == nil {
		provisioner = emptyprovisioner.Provisioner{}
	}

	if randSource == nil {
		randSource = rand.NewSource(time.Now().UnixNano())
	}

	if len(stream) == 0 {
		return nil, NewError(ECrit, "Stream name can't be empty", nil)
	}

	if opt == nil {
		tmp := DefaultKinesumerOptions
		opt = &tmp
	}

	return &Kinesumer{
		Kinesis:      kinesis,
		Checkpointer: checkpointer,
		Provisioner:  provisioner,
		Stream:       stream,
		Options:      opt,
		records:      make(chan k.Record, opt.GetRecordsLimit*2+10),
		rand:         rand.New(randSource),
	}, nil
}

func (kin *Kinesumer) GetStreams() (streams []string, err error) {
	streams = make([]string, 0)
	err = kin.Kinesis.ListStreamsPages(&kinesis.ListStreamsInput{
		Limit: &kin.Options.ListStreamsLimit,
	}, func(sts *kinesis.ListStreamsOutput, _ bool) bool {
		streams = append(streams, aws.StringValueSlice(sts.StreamNames)...)
		return true
	})
	return
}

func (kin *Kinesumer) StreamExists() (found bool, err error) {
	streams, err := kin.GetStreams()
	if err != nil {
		return
	}
	for _, stream := range streams {
		if stream == kin.Stream {
			return true, nil
		}
	}
	return
}

func (kin *Kinesumer) GetShards() (shards []*kinesis.Shard, err error) {
	shards = make([]*kinesis.Shard, 0)
	err = kin.Kinesis.DescribeStreamPages(&kinesis.DescribeStreamInput{
		Limit:      &kin.Options.DescribeStreamLimit,
		StreamName: &kin.Stream,
	}, func(desc *kinesis.DescribeStreamOutput, _ bool) bool {
		if desc == nil || desc.StreamDescription == nil {
			err = errors.New("Stream could not be described")
			return false
		}
		if aws.StringValue(desc.StreamDescription.StreamStatus) == "DELETING" {
			err = errors.New("Stream is being deleted")
			return false
		}
		shards = append(shards, desc.StreamDescription.Shards...)
		return true
	})
	return
}

func (kin *Kinesumer) LaunchShardWorker(shards []*kinesis.Shard) (int, error) {
	perm := kin.rand.Perm(len(shards))
	for _, j := range perm {
		err := kin.Provisioner.TryAcquire(aws.StringValue(shards[j].ShardID))
		if err == nil {
			worker := &ShardWorker{
				kinesis:         kin.Kinesis,
				shard:           shards[j],
				checkpointer:    kin.Checkpointer,
				stream:          kin.Stream,
				pollTime:        kin.Options.PollTime,
				stop:            kin.stop,
				stopped:         kin.stopped,
				c:               kin.records,
				provisioner:     kin.Provisioner,
				handlers:        kin.Options.Handlers,
				GetRecordsLimit: kin.Options.GetRecordsLimit,
			}
			kin.Options.Handlers.Go(func() {
				worker.RunWorker()
			})
			kin.nRunning++
			return j, nil
		}
	}
	return 0, errors.New("Could not launch worker")
}

func (kin *Kinesumer) Begin() (err error) {
	shards, err := kin.GetShards()
	if err != nil {
		return
	}

	err = kin.Checkpointer.Begin(kin.Options.Handlers)
	if err != nil {
		return
	}

	n := kin.Options.MaxShardWorkers
	if n <= 0 || len(shards) < n {
		n = len(shards)
	}

	start := time.Now()
	tryTime := 2 * kin.Provisioner.TTL()

	kin.stop = make(chan Unit, n)
	kin.stopped = make(chan Unit, n)
	for kin.nRunning < n && len(shards) > 0 && time.Now().Sub(start) < tryTime {
		for i := 0; i < n; i++ {
			j, err := kin.LaunchShardWorker(shards)
			if err != nil {
				kin.Options.Handlers.Err(NewError(EWarn, "Could not start shard worker", err))
			} else {
				shards = append(shards[:j], shards[j+1:]...)
			}
		}
		time.Sleep(500 * time.Millisecond)
	}

	return
}

func (kin *Kinesumer) End() {
	for kin.nRunning > 0 {
		select {
		case <-kin.stopped:
			kin.nRunning--
		case kin.stop <- Unit{}:
		}
	}
	kin.Checkpointer.End()
}

func (kin *Kinesumer) Records() <-chan k.Record {
	return kin.records
}
<commit_msg>Fixed counting bug<commit_after>package kinesumer

import (
	"errors"
	"fmt"
	"math/rand"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/remind101/kinesumer/checkpointers/empty"
	k "github.com/remind101/kinesumer/interface"
	"github.com/remind101/kinesumer/provisioners/empty"
)

type Kinesumer struct {
	Kinesis      k.Kinesis
	Checkpointer k.Checkpointer
	Provisioner  k.Provisioner
	Stream       string
	Options      *KinesumerOptions
	records      chan k.Record
	stop         chan Unit
	stopped      chan Unit
	nRunning     int
	rand         *rand.Rand
}

type KinesumerOptions struct {
	ListStreamsLimit    int64
	DescribeStreamLimit int64
	GetRecordsLimit     int64
	PollTime            int
	MaxShardWorkers     int
	Handlers            k.Handlers
}

var DefaultKinesumerOptions = KinesumerOptions{
	// These values are the hard limits set by Amazon

	ListStreamsLimit:    1000,
	DescribeStreamLimit: 10000,
	GetRecordsLimit:     10000,

	PollTime:        2000,
	MaxShardWorkers: 50,
	Handlers:        k.DefaultHandlers{},
}

func NewDefaultKinesumer(awsAccessKey, awsSecretKey, awsRegion, stream string) (*Kinesumer, error) {
	return NewKinesumer(
		kinesis.New(
			&aws.Config{
				Credentials: credentials.NewStaticCredentials(awsAccessKey, awsSecretKey, ""),
				Region:      &awsRegion,
			},
		),
		nil,
		nil,
		nil,
		stream,
		nil,
	)
}

func NewKinesumer(kinesis k.Kinesis, checkpointer k.Checkpointer, provisioner k.Provisioner,
	randSource rand.Source, stream string, opt *KinesumerOptions) (*Kinesumer, error) {

	if kinesis == nil {
		return nil, NewError(ECrit, "Kinesis object must not be nil", nil)
	}

	if checkpointer == nil {
		checkpointer = emptycheckpointer.Checkpointer{}
	}

	if provisioner == nil {
		provisioner = emptyprovisioner.Provisioner{}
	}

	if randSource == nil {
		randSource = rand.NewSource(time.Now().UnixNano())
	}

	if len(stream) == 0 {
		return nil, NewError(ECrit, "Stream name can't be empty", nil)
	}

	if opt == nil {
		tmp := DefaultKinesumerOptions
		opt = &tmp
	}

	return &Kinesumer{
		Kinesis:      kinesis,
		Checkpointer: checkpointer,
		Provisioner:  provisioner,
		Stream:       stream,
		Options:      opt,
		records:      make(chan k.Record, opt.GetRecordsLimit*2+10),
		rand:         rand.New(randSource),
	}, nil
}

func (kin *Kinesumer) GetStreams() (streams []string, err error) {
	streams = make([]string, 0)
	err = kin.Kinesis.ListStreamsPages(&kinesis.ListStreamsInput{
		Limit: &kin.Options.ListStreamsLimit,
	}, func(sts *kinesis.ListStreamsOutput, _ bool) bool {
		streams = append(streams, aws.StringValueSlice(sts.StreamNames)...)
		return true
	})
	return
}

func (kin *Kinesumer) StreamExists() (found bool, err error) {
	streams, err := kin.GetStreams()
	if err != nil {
		return
	}
	for _, stream := range streams {
		if stream == kin.Stream {
			return true, nil
		}
	}
	return
}

func (kin *Kinesumer) GetShards() (shards []*kinesis.Shard, err error) {
	shards = make([]*kinesis.Shard, 0)
	err = kin.Kinesis.DescribeStreamPages(&kinesis.DescribeStreamInput{
		Limit:      &kin.Options.DescribeStreamLimit,
		StreamName: &kin.Stream,
	}, func(desc *kinesis.DescribeStreamOutput, _ bool) bool {
		if desc == nil || desc.StreamDescription == nil {
			err = errors.New("Stream could not be described")
			return false
		}
		if aws.StringValue(desc.StreamDescription.StreamStatus) == "DELETING" {
			err = errors.New("Stream is being deleted")
			return false
		}
		shards = append(shards, desc.StreamDescription.Shards...)
		return true
	})
	return
}

func (kin *Kinesumer) LaunchShardWorker(shards []*kinesis.Shard) (int, error) {
	perm := kin.rand.Perm(len(shards))
	for _, j := range perm {
		err := kin.Provisioner.TryAcquire(aws.StringValue(shards[j].ShardID))
		if err == nil {
			worker := &ShardWorker{
				kinesis:         kin.Kinesis,
				shard:           shards[j],
				checkpointer:    kin.Checkpointer,
				stream:          kin.Stream,
				pollTime:        kin.Options.PollTime,
				stop:            kin.stop,
				stopped:         kin.stopped,
				c:               kin.records,
				provisioner:     kin.Provisioner,
				handlers:        kin.Options.Handlers,
				GetRecordsLimit: kin.Options.GetRecordsLimit,
			}
			kin.Options.Handlers.Go(func() {
				worker.RunWorker()
			})
			kin.nRunning++
			return j, nil
		}
	}
	return 0, errors.New("Could not launch worker")
}

func (kin *Kinesumer) Begin() (err error) {
	shards, err := kin.GetShards()
	if err != nil {
		return
	}

	err = kin.Checkpointer.Begin(kin.Options.Handlers)
	if err != nil {
		return
	}

	n := kin.Options.MaxShardWorkers
	if n <= 0 || len(shards) < n {
		n = len(shards)
	}

	start := time.Now()
	tryTime := 2 * kin.Provisioner.TTL()

	kin.stop = make(chan Unit, n)
	kin.stopped = make(chan Unit, n)
	for kin.nRunning < n && len(shards) > 0 && time.Now().Sub(start) < tryTime {
		for i := kin.nRunning; i < n; i++ {
			j, err := kin.LaunchShardWorker(shards)
			if err != nil {
				kin.Options.Handlers.Err(NewError(EWarn, "Could not start shard worker", err))
			} else {
				shards = append(shards[:j], shards[j+1:]...)
			}
		}
		time.Sleep(500 * time.Millisecond)
	}

	kin.Options.Handlers.Err(NewError(EWarn, fmt.Sprintf(EInfo, "%v/%v workers started", kin.nRunning, n), nil))

	return
}

func (kin *Kinesumer) End() {
	for kin.nRunning > 0 {
		select {
		case <-kin.stopped:
			kin.nRunning--
		case kin.stop <- Unit{}:
		}
	}
	kin.Checkpointer.End()
}

func (kin *Kinesumer) Records() <-chan k.Record {
	return kin.records
}
<|endoftext|>


<commit_before>package mp

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
)

type MediaPlayer struct {
	player       Backend
	stateChange  chan StateChange
	volumeChange chan int

	// Two channels to handle passing of the play state without causing race
	// conditions.
	playstateChan     chan PlayState
	playstateChanPing chan bool

	streams map[string]string // map of YouTube ID to stream gotten from youtube-dl
}

func New(stateChange chan StateChange, volumeChange chan int) *MediaPlayer {
	p := MediaPlayer{}

	p.playstateChan = make(chan PlayState)
	p.playstateChanPing = make(chan bool)
	p.stateChange = stateChange
	p.volumeChange = volumeChange
	p.streams = make(map[string]string)

	p.player = &MPV{}
	playerEventChan := p.player.initialize()

	go p.run(playerEventChan)

	return &p
}

func (p *MediaPlayer) Quit() {
	go func() {
		p.player.quit()
	}()

	// TODO: fix race conditions
	close(p.stateChange)
	close(p.volumeChange)
}

// GetPlaystate returns the play state synchronously (may block)
func (p *MediaPlayer) GetPlaystate() *PlayState {
	p.playstateChanPing <- false
	ps := <-p.playstateChan
	return &ps
}

// changePlaystate changes the play state inside a callback
// The *PlayState argument in the callback is the PlayState that can be
// changed, but it can only be accessed inside the callback (outside of
// it, race conditions can occur).
func (p *MediaPlayer) changePlaystate(callback func(*PlayState)) {
	p.playstateChanPing <- true
	ps := <-p.playstateChan
	callback(&ps)
	p.playstateChan <- ps
}

// SetPlaystate changes the play state to the specified arguments
// This function doesn't block, but changes may not be immediately applied.
func (p *MediaPlayer) SetPlaystate(playlist []string, index int, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Playlist = playlist
		ps.Index = index
		ps.Position = position

		if len(ps.Playlist) > 0 {
			p.startPlaying(ps)
		} else {
			p.Stop()
		}
	})
}

func (p *MediaPlayer) startPlaying(ps *PlayState) {
	p.setPlayState(ps, STATE_BUFFERING)

	videoId := ps.Playlist[ps.Index]
	position := ps.Position

	// do not use the playstate inside this goroutine to prevent race conditions
	go func() {
		streamUrl := p.getYouTubeStream(videoId)
		p.player.play(streamUrl, position)
	}()
}

func (p *MediaPlayer) getYouTubeStream(videoId string) string {
	if stream, ok := p.streams[videoId]; ok {
		return stream
	}

	youtubeUrl := "http://www.youtube.com/watch?v=" + videoId

	fmt.Println("Fetching YouTube stream...", youtubeUrl)
	cmd := exec.Command("youtube-dl", "-g", "-f", "bestaudio", youtubeUrl)
	cmd.Stderr = os.Stderr
	output, err := cmd.Output()
	fmt.Println("Got stream.")
	if err != nil {
		panic(err)
	}
	stream := strings.TrimSpace(string(output))
	p.streams[videoId] = stream
	return stream
}

func (p *MediaPlayer) setPlayState(ps *PlayState, state State) {
	if state == ps.State {
		fmt.Printf("WARNING: state %d did not change\n", state)
	}
	ps.State = state

	go func() {
		p.stateChange <- StateChange{ps.State, ps.Position}
	}()
}

func (p *MediaPlayer) UpdatePlaylist(playlist []string) {
	go p.changePlaystate(func(ps *PlayState) {

		if len(ps.Playlist) == 0 {

			if ps.State == STATE_PLAYING {
				// just to be sure
				panic("empty playlist while playing")
			}
			ps.Playlist = playlist

			if ps.Index >= len(playlist) {
				// this appears to be the normal behavior of YouTube
				ps.Index = len(playlist) - 1
			}

			if ps.State == STATE_STOPPED {
				ps.Position = 0
				p.startPlaying(ps)
			}

		} else {
			videoId := ps.Playlist[ps.Index]
			ps.Playlist = playlist
			p.setPlaylistIndex(ps, videoId)
		}
	})
}

func (p *MediaPlayer) SetVideo(videoId string, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		p.setPlaylistIndex(ps, videoId)
		ps.Position = position
		p.startPlaying(ps)
	})
}

func (p *MediaPlayer) setPlaylistIndex(ps *PlayState, videoId string) {
	newIndex := -1
	for i, v := range ps.Playlist {
		if v == videoId {
			if newIndex >= 0 {
				fmt.Fprintln(os.Stderr, "WARNING: videoId exists twice in playlist")
				break
			}
			newIndex = i
			// no 'break' so duplicate video entries can be checked
		}
	}

	if newIndex == -1 {
		// don't know how to proceed
		panic("current video does not exist in new playlist")
	}

	ps.Index = newIndex
}

// Pause pauses the currently playing video
func (p *MediaPlayer) Pause() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PLAYING {
			fmt.Printf("Warning: pause while in state %d\n", ps.State)
		}

		p.player.pause()
	})
}

// Resume resumes playback when it was paused
func (p *MediaPlayer) Resume() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			fmt.Printf("Warning: resume while in state %d\n", ps.State)
		}

		p.player.resume()
	})
}

// Seek jumps to the specified position
func (p *MediaPlayer) Seek(position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			fmt.Printf("Warning: state is not paused while seeking (state: %d)\n", ps.State)
		}

		p.player.setPosition(position)
		p.updatePosition(ps)
	})
}

// updatePosition gets the current position from the player
func (p *MediaPlayer) updatePosition(ps *PlayState) {
	// Also give 0 position on buffering because we can't ask the player where
	// we are while we are fetching the YouTube stream.
	if ps.State == STATE_STOPPED || ps.State == STATE_BUFFERING {
		ps.Position = 0
		return
	}

	ps.Position = p.player.getPosition()
}

func (p *MediaPlayer) SetVolume(volume int) {
	p.changePlaystate(func(ps *PlayState) {
		ps.Volume = volume
		p.player.setVolume(volume)
	})
}

func (p *MediaPlayer) Stop() {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Playlist = []string{}
		// Do not set ps.Index to 0, it may be needed for UpdatePlaylist:
		// Stop is called before UpdatePlaylist when removing the currently
		// playing video from the playlist.
		ps.Position = 0
		ps.State = STATE_STOPPED
		p.player.stop()
	})
}

func (p *MediaPlayer) run(playerEventChan chan State) {
	ps := PlayState{}

	ps.Volume = -1

	for {
		select {
		case replace := <-p.playstateChanPing:
			p.updatePosition(&ps)
			p.playstateChan <- ps
			if replace {
				ps = <-p.playstateChan
			}

		case event, ok := <-playerEventChan:
			if !ok {
				// player has quit, and closed channel
				return
			}

			if event == ps.State {
				// status hasn't changed
				// especially libvlc may send multiple events, ignore those
				continue
			}

			switch event {
			case STATE_PLAYING:
				p.setPlayState(&ps, STATE_PLAYING)

				if ps.Volume == -1 {
					ps.Volume = 100
					p.player.setVolume(ps.Volume)
					p.volumeChange <- ps.Volume
				}

			case STATE_PAUSED:
				p.setPlayState(&ps, STATE_PAUSED)

			case STATE_STOPPED:
				if ps.State == STATE_BUFFERING {
					// Especially VLC may keep sending 'stopped' events
					// while the next track is already buffering.
					// Ignore those events.
					continue
				}
				if ps.Index+1 < len(ps.Playlist) {
					// there are more videos, play the next
					ps.Index++
					ps.Position = 0
					p.startPlaying(&ps)
				} else {
					// signal that the video has stopped playing
					// this resets the position but keeps the playlist
					ps.Position = 0
					p.setPlayState(&ps, STATE_STOPPED)
				}
			}
		}
	}
}
<commit_msg>Use HTTPS everywhere possible<commit_after>package mp

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
)

type MediaPlayer struct {
	player       Backend
	stateChange  chan StateChange
	volumeChange chan int

	// Two channels to handle passing of the play state without causing race
	// conditions.
	playstateChan     chan PlayState
	playstateChanPing chan bool

	streams map[string]string // map of YouTube ID to stream gotten from youtube-dl
}

func New(stateChange chan StateChange, volumeChange chan int) *MediaPlayer {
	p := MediaPlayer{}

	p.playstateChan = make(chan PlayState)
	p.playstateChanPing = make(chan bool)
	p.stateChange = stateChange
	p.volumeChange = volumeChange
	p.streams = make(map[string]string)

	p.player = &MPV{}
	playerEventChan := p.player.initialize()

	go p.run(playerEventChan)

	return &p
}

func (p *MediaPlayer) Quit() {
	go func() {
		p.player.quit()
	}()

	// TODO: fix race conditions
	close(p.stateChange)
	close(p.volumeChange)
}

// GetPlaystate returns the play state synchronously (may block)
func (p *MediaPlayer) GetPlaystate() *PlayState {
	p.playstateChanPing <- false
	ps := <-p.playstateChan
	return &ps
}

// changePlaystate changes the play state inside a callback
// The *PlayState argument in the callback is the PlayState that can be
// changed, but it can only be accessed inside the callback (outside of
// it, race conditions can occur).
func (p *MediaPlayer) changePlaystate(callback func(*PlayState)) {
	p.playstateChanPing <- true
	ps := <-p.playstateChan
	callback(&ps)
	p.playstateChan <- ps
}

// SetPlaystate changes the play state to the specified arguments
// This function doesn't block, but changes may not be immediately applied.
func (p *MediaPlayer) SetPlaystate(playlist []string, index int, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Playlist = playlist
		ps.Index = index
		ps.Position = position

		if len(ps.Playlist) > 0 {
			p.startPlaying(ps)
		} else {
			p.Stop()
		}
	})
}

func (p *MediaPlayer) startPlaying(ps *PlayState) {
	p.setPlayState(ps, STATE_BUFFERING)

	videoId := ps.Playlist[ps.Index]
	position := ps.Position

	// do not use the playstate inside this goroutine to prevent race conditions
	go func() {
		streamUrl := p.getYouTubeStream(videoId)
		p.player.play(streamUrl, position)
	}()
}

func (p *MediaPlayer) getYouTubeStream(videoId string) string {
	if stream, ok := p.streams[videoId]; ok {
		return stream
	}

	youtubeUrl := "https://www.youtube.com/watch?v=" + videoId

	fmt.Println("Fetching YouTube stream...", youtubeUrl)
	cmd := exec.Command("youtube-dl", "-g", "-f", "bestaudio", youtubeUrl)
	cmd.Stderr = os.Stderr
	output, err := cmd.Output()
	fmt.Println("Got stream.")
	if err != nil {
		panic(err)
	}
	stream := strings.TrimSpace(string(output))
	p.streams[videoId] = stream
	return stream
}

func (p *MediaPlayer) setPlayState(ps *PlayState, state State) {
	if state == ps.State {
		fmt.Printf("WARNING: state %d did not change\n", state)
	}
	ps.State = state

	go func() {
		p.stateChange <- StateChange{ps.State, ps.Position}
	}()
}

func (p *MediaPlayer) UpdatePlaylist(playlist []string) {
	go p.changePlaystate(func(ps *PlayState) {

		if len(ps.Playlist) == 0 {

			if ps.State == STATE_PLAYING {
				// just to be sure
				panic("empty playlist while playing")
			}
			ps.Playlist = playlist

			if ps.Index >= len(playlist) {
				// this appears to be the normal behavior of YouTube
				ps.Index = len(playlist) - 1
			}

			if ps.State == STATE_STOPPED {
				ps.Position = 0
				p.startPlaying(ps)
			}

		} else {
			videoId := ps.Playlist[ps.Index]
			ps.Playlist = playlist
			p.setPlaylistIndex(ps, videoId)
		}
	})
}

func (p *MediaPlayer) SetVideo(videoId string, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		p.setPlaylistIndex(ps, videoId)
		ps.Position = position
		p.startPlaying(ps)
	})
}

func (p *MediaPlayer) setPlaylistIndex(ps *PlayState, videoId string) {
	newIndex := -1
	for i, v := range ps.Playlist {
		if v == videoId {
			if newIndex >= 0 {
				fmt.Fprintln(os.Stderr, "WARNING: videoId exists twice in playlist")
				break
			}
			newIndex = i
			// no 'break' so duplicate video entries can be checked
		}
	}

	if newIndex == -1 {
		// don't know how to proceed
		panic("current video does not exist in new playlist")
	}

	ps.Index = newIndex
}

// Pause pauses the currently playing video
func (p *MediaPlayer) Pause() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PLAYING {
			fmt.Printf("Warning: pause while in state %d\n", ps.State)
		}

		p.player.pause()
	})
}

// Resume resumes playback when it was paused
func (p *MediaPlayer) Resume() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			fmt.Printf("Warning: resume while in state %d\n", ps.State)
		}

		p.player.resume()
	})
}

// Seek jumps to the specified position
func (p *MediaPlayer) Seek(position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			fmt.Printf("Warning: state is not paused while seeking (state: %d)\n", ps.State)
		}

		p.player.setPosition(position)
		p.updatePosition(ps)
	})
}

// updatePosition gets the current position from the player
func (p *MediaPlayer) updatePosition(ps *PlayState) {
	// Also give 0 position on buffering because we can't ask the player where
	// we are while we are fetching the YouTube stream.
	if ps.State == STATE_STOPPED || ps.State == STATE_BUFFERING {
		ps.Position = 0
		return
	}

	ps.Position = p.player.getPosition()
}

func (p *MediaPlayer) SetVolume(volume int) {
	p.changePlaystate(func(ps *PlayState) {
		ps.Volume = volume
		p.player.setVolume(volume)
	})
}

func (p *MediaPlayer) Stop() {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Playlist = []string{}
		// Do not set ps.Index to 0, it may be needed for UpdatePlaylist:
		// Stop is called before UpdatePlaylist when removing the currently
		// playing video from the playlist.
		ps.Position = 0
		ps.State = STATE_STOPPED
		p.player.stop()
	})
}

func (p *MediaPlayer) run(playerEventChan chan State) {
	ps := PlayState{}

	ps.Volume = -1

	for {
		select {
		case replace := <-p.playstateChanPing:
			p.updatePosition(&ps)
			p.playstateChan <- ps
			if replace {
				ps = <-p.playstateChan
			}

		case event, ok := <-playerEventChan:
			if !ok {
				// player has quit, and closed channel
				return
			}

			if event == ps.State {
				// status hasn't changed
				// especially libvlc may send multiple events, ignore those
				continue
			}

			switch event {
			case STATE_PLAYING:
				p.setPlayState(&ps, STATE_PLAYING)

				if ps.Volume == -1 {
					ps.Volume = 100
					p.player.setVolume(ps.Volume)
					p.volumeChange <- ps.Volume
				}

			case STATE_PAUSED:
				p.setPlayState(&ps, STATE_PAUSED)

			case STATE_STOPPED:
				if ps.State == STATE_BUFFERING {
					// Especially VLC may keep sending 'stopped' events
					// while the next track is already buffering.
					// Ignore those events.
					continue
				}
				if ps.Index+1 < len(ps.Playlist) {
					// there are more videos, play the next
					ps.Index++
					ps.Position = 0
					p.startPlaying(&ps)
				} else {
					// signal that the video has stopped playing
					// this resets the position but keeps the playlist
					ps.Position = 0
					p.setPlayState(&ps, STATE_STOPPED)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>//go:build !providerless
// +build !providerless

/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package aws

import (
	"encoding/json"
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

// IPPermissionSet maps IP strings of strings to EC2 IpPermissions
type IPPermissionSet map[string]*ec2.IpPermission

// IPPermissionPredicate is an predicate to test whether IPPermission matches some condition.
type IPPermissionPredicate interface {
	// Test checks whether specified IPPermission matches condition.
	Test(perm *ec2.IpPermission) bool
}

// NewIPPermissionSet creates a new IPPermissionSet
func NewIPPermissionSet(items ...*ec2.IpPermission) IPPermissionSet {
	s := make(IPPermissionSet)
	s.Insert(items...)
	return s
}

// Ungroup splits permissions out into individual permissions
// EC2 will combine permissions with the same port but different SourceRanges together, for example
// We ungroup them so we can process them
func (s IPPermissionSet) Ungroup() IPPermissionSet {
	l := []*ec2.IpPermission{}
	for _, p := range s.List() {
		if len(p.IpRanges) <= 1 {
			l = append(l, p)
			continue
		}
		for _, ipRange := range p.IpRanges {
			c := &ec2.IpPermission{}
			*c = *p
			c.IpRanges = []*ec2.IpRange{ipRange}
			l = append(l, c)
		}
	}

	l2 := []*ec2.IpPermission{}
	for _, p := range l {
		if len(p.UserIdGroupPairs) <= 1 {
			l2 = append(l2, p)
			continue
		}
		for _, u := range p.UserIdGroupPairs {
			c := &ec2.IpPermission{}
			*c = *p
			c.UserIdGroupPairs = []*ec2.UserIdGroupPair{u}
			l2 = append(l, c)
		}
	}

	l3 := []*ec2.IpPermission{}
	for _, p := range l2 {
		if len(p.PrefixListIds) <= 1 {
			l3 = append(l3, p)
			continue
		}
		for _, v := range p.PrefixListIds {
			c := &ec2.IpPermission{}
			*c = *p
			c.PrefixListIds = []*ec2.PrefixListId{v}
			l3 = append(l3, c)
		}
	}

	return NewIPPermissionSet(l3...)
}

// Insert adds items to the set.
func (s IPPermissionSet) Insert(items ...*ec2.IpPermission) {
	for _, p := range items {
		k := keyForIPPermission(p)
		s[k] = p
	}
}

// Delete delete permission from the set.
func (s IPPermissionSet) Delete(items ...*ec2.IpPermission) {
	for _, p := range items {
		k := keyForIPPermission(p)
		delete(s, k)
	}
}

// DeleteIf delete permission from the set if permission matches predicate.
func (s IPPermissionSet) DeleteIf(predicate IPPermissionPredicate) {
	for k, p := range s {
		if predicate.Test(p) {
			delete(s, k)
		}
	}
}

// List returns the contents as a slice.  Order is not defined.
func (s IPPermissionSet) List() []*ec2.IpPermission {
	res := make([]*ec2.IpPermission, 0, len(s))
	for _, v := range s {
		res = append(res, v)
	}
	return res
}

// IsSuperset returns true if and only if s is a superset of s2.
func (s IPPermissionSet) IsSuperset(s2 IPPermissionSet) bool {
	for k := range s2 {
		_, found := s[k]
		if !found {
			return false
		}
	}
	return true
}

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

// Difference returns a set of objects that are not in s2
// For example:
// s1 = {a1, a2, a3}
// s2 = {a1, a2, a4, a5}
// s1.Difference(s2) = {a3}
// s2.Difference(s1) = {a4, a5}
func (s IPPermissionSet) Difference(s2 IPPermissionSet) IPPermissionSet {
	result := NewIPPermissionSet()
	for k, v := range s {
		_, found := s2[k]
		if !found {
			result[k] = v
		}
	}
	return result
}

// Len returns the size of the set.
func (s IPPermissionSet) Len() int {
	return len(s)
}

func keyForIPPermission(p *ec2.IpPermission) string {
	v, err := json.Marshal(p)
	if err != nil {
		panic(fmt.Sprintf("error building JSON representation of ec2.IpPermission: %v", err))
	}
	return string(v)
}

var _ IPPermissionPredicate = IPPermissionMatchDesc{}

// IPPermissionMatchDesc checks whether specific IPPermission contains description.
type IPPermissionMatchDesc struct {
	Description string
}

// Test whether specific IPPermission contains description.
func (p IPPermissionMatchDesc) Test(perm *ec2.IpPermission) bool {
	for _, v4Range := range perm.IpRanges {
		if aws.StringValue(v4Range.Description) == p.Description {
			return true
		}
	}
	for _, v6Range := range perm.Ipv6Ranges {
		if aws.StringValue(v6Range.Description) == p.Description {
			return true
		}
	}
	for _, prefixListID := range perm.PrefixListIds {
		if aws.StringValue(prefixListID.Description) == p.Description {
			return true
		}
	}
	for _, group := range perm.UserIdGroupPairs {
		if aws.StringValue(group.Description) == p.Description {
			return true
		}
	}
	return false
}

var _ IPPermissionPredicate = IPPermissionNotMatch{}

// IPPermissionNotMatch is the *not* operator for Predicate
type IPPermissionNotMatch struct {
	Predicate IPPermissionPredicate
}

// Test whether specific IPPermission not match the embed predicate.
func (p IPPermissionNotMatch) Test(perm *ec2.IpPermission) bool {
	return !p.Predicate.Test(perm)
}
<commit_msg>make map cap<commit_after>//go:build !providerless
// +build !providerless

/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package aws

import (
	"encoding/json"
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

// IPPermissionSet maps IP strings of strings to EC2 IpPermissions
type IPPermissionSet map[string]*ec2.IpPermission

// IPPermissionPredicate is an predicate to test whether IPPermission matches some condition.
type IPPermissionPredicate interface {
	// Test checks whether specified IPPermission matches condition.
	Test(perm *ec2.IpPermission) bool
}

// NewIPPermissionSet creates a new IPPermissionSet
func NewIPPermissionSet(items ...*ec2.IpPermission) IPPermissionSet {
	s := make(IPPermissionSet, len(items))
	s.Insert(items...)
	return s
}

// Ungroup splits permissions out into individual permissions
// EC2 will combine permissions with the same port but different SourceRanges together, for example
// We ungroup them so we can process them
func (s IPPermissionSet) Ungroup() IPPermissionSet {
	l := []*ec2.IpPermission{}
	for _, p := range s.List() {
		if len(p.IpRanges) <= 1 {
			l = append(l, p)
			continue
		}
		for _, ipRange := range p.IpRanges {
			c := &ec2.IpPermission{}
			*c = *p
			c.IpRanges = []*ec2.IpRange{ipRange}
			l = append(l, c)
		}
	}

	l2 := []*ec2.IpPermission{}
	for _, p := range l {
		if len(p.UserIdGroupPairs) <= 1 {
			l2 = append(l2, p)
			continue
		}
		for _, u := range p.UserIdGroupPairs {
			c := &ec2.IpPermission{}
			*c = *p
			c.UserIdGroupPairs = []*ec2.UserIdGroupPair{u}
			l2 = append(l, c)
		}
	}

	l3 := []*ec2.IpPermission{}
	for _, p := range l2 {
		if len(p.PrefixListIds) <= 1 {
			l3 = append(l3, p)
			continue
		}
		for _, v := range p.PrefixListIds {
			c := &ec2.IpPermission{}
			*c = *p
			c.PrefixListIds = []*ec2.PrefixListId{v}
			l3 = append(l3, c)
		}
	}

	return NewIPPermissionSet(l3...)
}

// Insert adds items to the set.
func (s IPPermissionSet) Insert(items ...*ec2.IpPermission) {
	for _, p := range items {
		k := keyForIPPermission(p)
		s[k] = p
	}
}

// Delete delete permission from the set.
func (s IPPermissionSet) Delete(items ...*ec2.IpPermission) {
	for _, p := range items {
		k := keyForIPPermission(p)
		delete(s, k)
	}
}

// DeleteIf delete permission from the set if permission matches predicate.
func (s IPPermissionSet) DeleteIf(predicate IPPermissionPredicate) {
	for k, p := range s {
		if predicate.Test(p) {
			delete(s, k)
		}
	}
}

// List returns the contents as a slice.  Order is not defined.
func (s IPPermissionSet) List() []*ec2.IpPermission {
	res := make([]*ec2.IpPermission, 0, len(s))
	for _, v := range s {
		res = append(res, v)
	}
	return res
}

// IsSuperset returns true if and only if s is a superset of s2.
func (s IPPermissionSet) IsSuperset(s2 IPPermissionSet) bool {
	for k := range s2 {
		_, found := s[k]
		if !found {
			return false
		}
	}
	return true
}

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

// Difference returns a set of objects that are not in s2
// For example:
// s1 = {a1, a2, a3}
// s2 = {a1, a2, a4, a5}
// s1.Difference(s2) = {a3}
// s2.Difference(s1) = {a4, a5}
func (s IPPermissionSet) Difference(s2 IPPermissionSet) IPPermissionSet {
	result := NewIPPermissionSet()
	for k, v := range s {
		_, found := s2[k]
		if !found {
			result[k] = v
		}
	}
	return result
}

// Len returns the size of the set.
func (s IPPermissionSet) Len() int {
	return len(s)
}

func keyForIPPermission(p *ec2.IpPermission) string {
	v, err := json.Marshal(p)
	if err != nil {
		panic(fmt.Sprintf("error building JSON representation of ec2.IpPermission: %v", err))
	}
	return string(v)
}

var _ IPPermissionPredicate = IPPermissionMatchDesc{}

// IPPermissionMatchDesc checks whether specific IPPermission contains description.
type IPPermissionMatchDesc struct {
	Description string
}

// Test whether specific IPPermission contains description.
func (p IPPermissionMatchDesc) Test(perm *ec2.IpPermission) bool {
	for _, v4Range := range perm.IpRanges {
		if aws.StringValue(v4Range.Description) == p.Description {
			return true
		}
	}
	for _, v6Range := range perm.Ipv6Ranges {
		if aws.StringValue(v6Range.Description) == p.Description {
			return true
		}
	}
	for _, prefixListID := range perm.PrefixListIds {
		if aws.StringValue(prefixListID.Description) == p.Description {
			return true
		}
	}
	for _, group := range perm.UserIdGroupPairs {
		if aws.StringValue(group.Description) == p.Description {
			return true
		}
	}
	return false
}

var _ IPPermissionPredicate = IPPermissionNotMatch{}

// IPPermissionNotMatch is the *not* operator for Predicate
type IPPermissionNotMatch struct {
	Predicate IPPermissionPredicate
}

// Test whether specific IPPermission not match the embed predicate.
func (p IPPermissionNotMatch) Test(perm *ec2.IpPermission) bool {
	return !p.Predicate.Test(perm)
}
<|endoftext|>


<commit_before>package fate

import (
	"math/rand"
	"sort"
)

// tokset maintains a set of tokens as a sorted slice of integers.
//
// 1-byte tokens (<= 0xFF) are in buf[0:t2]
// 2-byte tokens (<= 0xFFFF) are in buf[t2:t3]
// 3-byte tokens (<= 0xFFFFFF) are in buf[t3:]
//
// They're stored little-endian. Adds are O(log N). Choosing a random
// token in the set is O(1).
//
// tokens greater than 0xFFFFFF are not currently supported. This is
// enough token space to handle the Web 1T corpus.
type tokset struct {
	buf []byte

	// count of 2-byte tokens, count of 1-byte tokens
	c2 uint16
	c1 uint8
}

func (t *tokset) Add(tok token) bool {
	size := len(t.buf)

	switch {
	case tok <= 0xFF:
		if size == 0 {
			t.buf = append(t.buf, byte(tok))
			t.c1++
			return false
		}

		return t.add1(tok)
	case tok <= 0xFFFF:
		if size == 0 {
			t.buf = append(t.buf, byte(tok), byte(tok>>8))
			t.c2++
			return false
		}

		return t.add2(tok)
	case tok <= 0xFFFFFF:
		if size == 0 {
			t.buf = append(t.buf, byte(tok), byte(tok>>8), byte(tok>>16))
			return false
		}

		return t.add3(tok)
	}

	panic("oops")
}

func (t *tokset) span1() []byte {
	return t.buf[0:t.c1]
}

func (t *tokset) span2() []byte {
	return t.buf[int(t.c1) : int(t.c1)+2*int(t.c2)]
}

func (t *tokset) span3() []byte {
	return t.buf[int(t.c1)+2*int(t.c2):]
}

func (t *tokset) add1(tok token) bool {
	span := t.span1()
	loc := sort.Search(len(span), func(i int) bool {
		return token(span[i]) >= tok
	})

	if loc < len(span) && token(span[loc]) == tok {
		return true
	}

	t.buf = append(t.buf, 0)
	copy(t.buf[loc+1:], t.buf[loc:])
	t.buf[loc] = byte(tok)

	t.c1++

	return false
}

func (t *tokset) add2(tok token) bool {
	span := t.span2()
	idx := sort.Search(len(span)/2, func(i int) bool {
		return unpack2(span[2*i:]) >= tok
	})

	if idx < len(span)/2 && unpack2(span[2*idx:]) == tok {
		return true
	}

	t.buf = append(t.buf, 0, 0)

	loc := int(t.c1) + 2*idx
	copy(t.buf[loc+2:], t.buf[loc:])
	put2(t.buf[loc:], tok)

	t.c2++

	return false
}
func (t *tokset) add3(tok token) bool {
	span := t.span3()
	idx := sort.Search(len(span)/3, func(i int) bool {
		return unpack3(span[3*i:]) >= tok
	})

	if idx < len(span)/3 && unpack3(span[3*idx:]) == tok {
		return true
	}

	t.buf = append(t.buf, 0, 0, 0)

	loc := int(t.c1) + 2*int(t.c2) + 3*idx
	copy(t.buf[loc+3:], t.buf[loc:])
	put3(t.buf[loc:], tok)

	return false
}

func (t *tokset) Len() int {
	if t == nil {
		return 0
	}

	return int(t.c1) + int(t.c2) + len(t.span3())/3
}

func (t *tokset) Tokens() []token {
	if t == nil {
		return nil
	}

	var tokens []token
	for _, val := range t.span1() {
		tokens = append(tokens, token(val))
	}

	span2 := t.span2()
	for i := 0; i < len(span2); i += 2 {
		tokens = append(tokens, unpack2(span2[i:]))
	}

	span3 := t.span3()
	for i := 0; i < len(span3); i += 3 {
		tokens = append(tokens, unpack3(span3[i:]))
	}

	return tokens
}

func put2(buf []byte, tok token) {
	buf[0] = byte(tok)
	buf[1] = byte(tok >> 8)
}

func put3(buf []byte, tok token) {
	buf[0] = byte(tok)
	buf[1] = byte(tok >> 8)
	buf[2] = byte(tok >> 16)
}

func unpack2(buf []byte) token {
	return token(buf[0]) | token(buf[1])<<8
}

func unpack3(buf []byte) token {
	return token(buf[0]) | token(buf[1])<<8 | token(buf[2])<<16
}

func (t tokset) Choice(r *rand.Rand) token {
	index := r.Intn(t.Len())

	switch {
	case index < int(t.c1):
		return token(t.buf[index])
	case index < int(t.c1)+int(t.c2):
		span := t.span2()
		return unpack2(span[2*(index-int(t.c1)):])
	case index < t.Len():
		span := t.span3()
		return unpack3(span[3*(index-(int(t.c2)+int(t.c1))):])
	}

	panic("oops")
}
<commit_msg>Fix the tokset comment for c1 & c2<commit_after>package fate

import (
	"math/rand"
	"sort"
)

// tokset maintains a set of tokens as a sorted slice of integers.
//
// 1-byte tokens (<= 0xFF) are in buf[0:c1]
// 2-byte tokens (<= 0xFFFF) are in buf[c1:c1+2*c2]
// 3-byte tokens (<= 0xFFFFFF) are in buf[c1+2*c2:]
//
// They're stored little-endian. Adds are O(log N). Choosing a random
// token in the set is O(1).
//
// tokens greater than 0xFFFFFF are not currently supported. This is
// enough token space to handle the Web 1T corpus.
type tokset struct {
	buf []byte

	// count of 2-byte tokens, count of 1-byte tokens
	c2 uint16
	c1 uint8
}

func (t *tokset) Add(tok token) bool {
	size := len(t.buf)

	switch {
	case tok <= 0xFF:
		if size == 0 {
			t.buf = append(t.buf, byte(tok))
			t.c1++
			return false
		}

		return t.add1(tok)
	case tok <= 0xFFFF:
		if size == 0 {
			t.buf = append(t.buf, byte(tok), byte(tok>>8))
			t.c2++
			return false
		}

		return t.add2(tok)
	case tok <= 0xFFFFFF:
		if size == 0 {
			t.buf = append(t.buf, byte(tok), byte(tok>>8), byte(tok>>16))
			return false
		}

		return t.add3(tok)
	}

	panic("oops")
}

func (t *tokset) span1() []byte {
	return t.buf[0:t.c1]
}

func (t *tokset) span2() []byte {
	return t.buf[int(t.c1) : int(t.c1)+2*int(t.c2)]
}

func (t *tokset) span3() []byte {
	return t.buf[int(t.c1)+2*int(t.c2):]
}

func (t *tokset) add1(tok token) bool {
	span := t.span1()
	loc := sort.Search(len(span), func(i int) bool {
		return token(span[i]) >= tok
	})

	if loc < len(span) && token(span[loc]) == tok {
		return true
	}

	t.buf = append(t.buf, 0)
	copy(t.buf[loc+1:], t.buf[loc:])
	t.buf[loc] = byte(tok)

	t.c1++

	return false
}

func (t *tokset) add2(tok token) bool {
	span := t.span2()
	idx := sort.Search(len(span)/2, func(i int) bool {
		return unpack2(span[2*i:]) >= tok
	})

	if idx < len(span)/2 && unpack2(span[2*idx:]) == tok {
		return true
	}

	t.buf = append(t.buf, 0, 0)

	loc := int(t.c1) + 2*idx
	copy(t.buf[loc+2:], t.buf[loc:])
	put2(t.buf[loc:], tok)

	t.c2++

	return false
}
func (t *tokset) add3(tok token) bool {
	span := t.span3()
	idx := sort.Search(len(span)/3, func(i int) bool {
		return unpack3(span[3*i:]) >= tok
	})

	if idx < len(span)/3 && unpack3(span[3*idx:]) == tok {
		return true
	}

	t.buf = append(t.buf, 0, 0, 0)

	loc := int(t.c1) + 2*int(t.c2) + 3*idx
	copy(t.buf[loc+3:], t.buf[loc:])
	put3(t.buf[loc:], tok)

	return false
}

func (t *tokset) Len() int {
	if t == nil {
		return 0
	}

	return int(t.c1) + int(t.c2) + len(t.span3())/3
}

func (t *tokset) Tokens() []token {
	if t == nil {
		return nil
	}

	var tokens []token
	for _, val := range t.span1() {
		tokens = append(tokens, token(val))
	}

	span2 := t.span2()
	for i := 0; i < len(span2); i += 2 {
		tokens = append(tokens, unpack2(span2[i:]))
	}

	span3 := t.span3()
	for i := 0; i < len(span3); i += 3 {
		tokens = append(tokens, unpack3(span3[i:]))
	}

	return tokens
}

func put2(buf []byte, tok token) {
	buf[0] = byte(tok)
	buf[1] = byte(tok >> 8)
}

func put3(buf []byte, tok token) {
	buf[0] = byte(tok)
	buf[1] = byte(tok >> 8)
	buf[2] = byte(tok >> 16)
}

func unpack2(buf []byte) token {
	return token(buf[0]) | token(buf[1])<<8
}

func unpack3(buf []byte) token {
	return token(buf[0]) | token(buf[1])<<8 | token(buf[2])<<16
}

func (t tokset) Choice(r *rand.Rand) token {
	index := r.Intn(t.Len())

	switch {
	case index < int(t.c1):
		return token(t.buf[index])
	case index < int(t.c1)+int(t.c2):
		span := t.span2()
		return unpack2(span[2*(index-int(t.c1)):])
	case index < t.Len():
		span := t.span3()
		return unpack3(span[3*(index-(int(t.c2)+int(t.c1))):])
	}

	panic("oops")
}
<|endoftext|>