<commit_before>// Package grpcutil implements a gRPC HTTP handler and OpenAPI proxy. Useful to serve gRPC requests from an existing HTTP servers.
package grpcutil

import (
	"crypto/tls"
	"crypto/x509"
	"log"
	"net/http"
	"strings"

	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// ServiceBinding are the gRPC server, and gRPC HTTP gateway to which the service will be bound to.
type ServiceBinding struct {
	GRPCServer        *grpc.Server
	GRPCGatewayClient *grpc.ClientConn
	GRPCGatewayMuxer  *runtime.ServeMux
}

// ServiceRegisterFn defines a function type for registering gRPC services.
type ServiceRegisterFn func(ServiceBinding) error

// Option implements http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type Option func(*options)

type options struct {
	serverOpts []grpc.ServerOption
	services   []ServiceRegisterFn
	cert       *tls.Certificate
	port       string
}

// WithServerOpts sets gRPC server options. Optional.
func WithServerOpts(opts []grpc.ServerOption) Option {
	return func(o *options) {
		o.serverOpts = opts
	}
}

// WithServices sets the list of services to register with the gRPC and OpenAPI servers. Required.
func WithServices(services []ServiceRegisterFn) Option {
	return func(o *options) {
		o.services = services
	}
}

// WithTLSCert sets the TLS certificate to use by the gRPC server and OpenAPI gRPC client. Required.
func WithTLSCert(cert *tls.Certificate) Option {
	return func(o *options) {
		o.cert = cert
	}
}

// WithPort sets the port for the OpenAPI gRPC client to use when connecting to the gRPC service.
func WithPort(port string) Option {
	return func(o *options) {
		o.port = port
	}
}

// Handler serves gRPC and OpenAPI requests or hands over to the next handler if they are not gRPC
// or application/json requests.
func Handler(h http.Handler, opts ...Option) http.Handler {
	options := new(options)
	for _, opt := range opts {
		opt(options)
	}

	if options.cert == nil {
		log.Fatal("grpcutil: TLS certificate required")
	}

	if options.port == "" {
		log.Fatal("grpcutil: port required")
	}

	if len(options.services) == 0 {
		log.Fatal("grpcutil: the list of gRPC services is required")
	}

	options.serverOpts = append(options.serverOpts, grpc.Creds(credentials.NewServerTLSFromCert(options.cert)))
	server := grpc.NewServer(options.serverOpts...)

	certPool := x509.NewCertPool()
	x509Cert, err := x509.ParseCertificate(options.cert.Certificate[0])
	if err != nil {
		log.Fatalf("grpcutil: failed parsing x509 certificate: %+v, ", err)
	}
	certPool.AddCert(x509Cert)
	clientCreds := credentials.NewClientTLSFromCert(certPool, "")

	clientOpts := []grpc.DialOption{
		grpc.WithTransportCredentials(clientCreds),
		grpc.WithUserAgent("grpc-gw"),
	}

	address := "localhost:" + options.port
	clientConn, err := grpc.Dial(address, clientOpts...)
	if err != nil {
		log.Fatalf("failed to connect to local gRPC server: %v", err)
	}

	gwMuxer := runtime.NewServeMux()

	serviceBinding := ServiceBinding{
		GRPCServer:        server,
		GRPCGatewayClient: clientConn,
		GRPCGatewayMuxer:  gwMuxer,
	}

	log.Println("registering GRPC services...")
	for _, register := range options.services {
		if err := register(serviceBinding); err != nil {
			log.Fatalf("failed to register service: %+v", err)
		}
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		contentType := r.Header.Get("Content-Type")
		if r.ProtoMajor == 2 && strings.Contains(contentType, "application/grpc") {
			server.ServeHTTP(w, r)
			return
		}

		accept := r.Header.Get("Accept")
		if strings.Contains(contentType, "application/json") ||
			strings.Contains(accept, "application/json") {
			gwMuxer.ServeHTTP(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}
<commit_msg>grpcutil: add option to skip specific request paths<commit_after>// Package grpcutil implements a gRPC HTTP handler and OpenAPI proxy. Useful to serve gRPC requests from an existing HTTP servers.
package grpcutil

import (
	"crypto/tls"
	"crypto/x509"
	"log"
	"net/http"
	"strings"

	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// ServiceBinding are the gRPC server, and gRPC HTTP gateway to which the service will be bound to.
type ServiceBinding struct {
	GRPCServer        *grpc.Server
	GRPCGatewayClient *grpc.ClientConn
	GRPCGatewayMuxer  *runtime.ServeMux
}

// ServiceRegisterFn defines a function type for registering gRPC services.
type ServiceRegisterFn func(ServiceBinding) error

// Option implements http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type Option func(*options)

type options struct {
	serverOpts []grpc.ServerOption
	services   []ServiceRegisterFn
	cert       *tls.Certificate
	port       string
	skipPaths  []string
}

// WithServerOpts sets gRPC server options. Optional.
func WithServerOpts(opts []grpc.ServerOption) Option {
	return func(o *options) {
		o.serverOpts = opts
	}
}

// WithServices sets the list of services to register with the gRPC and OpenAPI servers. Required.
func WithServices(services []ServiceRegisterFn) Option {
	return func(o *options) {
		o.services = services
	}
}

// WithTLSCert sets the TLS certificate to use by the gRPC server and OpenAPI gRPC client. Required.
func WithTLSCert(cert *tls.Certificate) Option {
	return func(o *options) {
		o.cert = cert
	}
}

// WithPort sets the port for the OpenAPI gRPC client to use when connecting to the gRPC service.
func WithPort(port string) Option {
	return func(o *options) {
		o.port = port
	}
}

// WithSkipPath allows other handlers to serve static JSON files by instructing the GRPC Gateway Muxer to
// skip serving the given prefixed paths.
func WithSkipPath(path ...string) Option {
	return func(o *options) {
		o.skipPaths = path
	}
}

// Handler serves gRPC and OpenAPI requests or hands over to the next handler if they are not gRPC
// or application/json requests.
func Handler(h http.Handler, opts ...Option) http.Handler {
	options := new(options)
	for _, opt := range opts {
		opt(options)
	}

	if options.cert == nil {
		log.Fatal("grpcutil: TLS certificate required")
	}

	if options.port == "" {
		log.Fatal("grpcutil: port required")
	}

	if len(options.services) == 0 {
		log.Fatal("grpcutil: the list of gRPC services is required")
	}

	options.serverOpts = append(options.serverOpts, grpc.Creds(credentials.NewServerTLSFromCert(options.cert)))
	server := grpc.NewServer(options.serverOpts...)

	certPool := x509.NewCertPool()
	x509Cert, err := x509.ParseCertificate(options.cert.Certificate[0])
	if err != nil {
		log.Fatalf("grpcutil: failed parsing x509 certificate: %+v, ", err)
	}
	certPool.AddCert(x509Cert)
	clientCreds := credentials.NewClientTLSFromCert(certPool, "")

	clientOpts := []grpc.DialOption{
		grpc.WithTransportCredentials(clientCreds),
		grpc.WithUserAgent("grpc-gw"),
	}

	address := "localhost:" + options.port
	clientConn, err := grpc.Dial(address, clientOpts...)
	if err != nil {
		log.Fatalf("failed to connect to local gRPC server: %v", err)
	}

	gwMuxer := runtime.NewServeMux()

	serviceBinding := ServiceBinding{
		GRPCServer:        server,
		GRPCGatewayClient: clientConn,
		GRPCGatewayMuxer:  gwMuxer,
	}

	log.Println("registering GRPC services...")
	for _, register := range options.services {
		if err := register(serviceBinding); err != nil {
			log.Fatalf("failed to register service: %+v", err)
		}
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		contentType := r.Header.Get("Content-Type")
		if r.ProtoMajor == 2 && strings.Contains(contentType, "application/grpc") {
			server.ServeHTTP(w, r)
			return
		}

		for _, p := range options.skipPaths {
			if strings.HasPrefix(r.URL.Path, p) {
				h.ServeHTTP(w, r)
				return
			}
		}

		accept := r.Header.Get("Accept")
		if strings.Contains(contentType, "application/json") ||
			strings.Contains(accept, "application/json") {
			gwMuxer.ServeHTTP(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}
<|endoftext|>


<commit_before>// Package grpcutil implements a gRPC HTTP handler and OpenAPI proxy. Useful to serve gRPC requests from an existing HTTP servers.
package grpcutil

import (
	"crypto/tls"
	"crypto/x509"
	"log"
	"net"
	"net/http"
	"strings"

	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// ServiceBinding are the gRPC server, and gRPC HTTP gateway to which the service will be bound to.
type ServiceBinding struct {
	GRPCServer        *grpc.Server
	GRPCGatewayClient *grpc.ClientConn
	GRPCGatewayMuxer  *runtime.ServeMux
}

// ServiceRegisterFn defines a function type for registering gRPC services.
type ServiceRegisterFn func(ServiceBinding) error

// Option implements http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type Option func(*options)

type options struct {
	serverOpts   []grpc.ServerOption
	services     []ServiceRegisterFn
	cert         *tls.Certificate
	compressor   grpc.Compressor
	decompressor grpc.Decompressor
	errorHandler runtime.ProtoErrorHandlerFunc
	port         string
	skipPaths    []string
}

// WithServerOpts sets gRPC server options. Optional.
func WithServerOpts(opts []grpc.ServerOption) Option {
	return func(o *options) {
		o.serverOpts = opts
	}
}

// WithServices sets the list of services to register with the gRPC and OpenAPI servers. Required.
func WithServices(services []ServiceRegisterFn) Option {
	return func(o *options) {
		o.services = services
	}
}

// WithTLSCert sets the TLS certificate to use by the gRPC server and OpenAPI gRPC client. Required.
func WithTLSCert(cert *tls.Certificate) Option {
	return func(o *options) {
		o.cert = cert
	}
}

// WithPort sets the port for the OpenAPI gRPC client to use when connecting to the gRPC service.
func WithPort(port string) Option {
	return func(o *options) {
		o.port = port
	}
}

// WithCompressor sets gRPC server and gateway client compressor.
// Use this function instead of passing compressor through grcputil.WithServerOpts()
func WithCompressor(c grpc.Compressor) Option {
	return func(o *options) {
		o.compressor = c
	}
}

// WithDecompressor sets gRPC server and gateway client decompressor.
// Use this function instead of passing a decompressor through grcputil.WithServerOpts()
func WithDecompressor(d grpc.Decompressor) Option {
	return func(o *options) {
		o.decompressor = d
	}
}

// WithSkipPath allows other handlers to serve static JSON files by instructing the GRPC Gateway Muxer to
// skip serving the given prefixed paths.
func WithSkipPath(path ...string) Option {
	return func(o *options) {
		o.skipPaths = path
	}
}

// WithErrorHandler sets a custom gRPC gateway error handler.
func WithErrorHandler(fn runtime.ProtoErrorHandlerFunc) Option {
	return func(o *options) {
		o.errorHandler = fn
	}
}

// Handler serves gRPC and OpenAPI requests or hands over to the next handler if they are not gRPC
// or application/json requests.
func Handler(h http.Handler, opts ...Option) http.Handler {
	options := new(options)
	for _, opt := range opts {
		opt(options)
	}

	if options.cert == nil {
		log.Fatal("grpcutil: TLS certificate required")
	}

	if options.port == "" {
		log.Fatal("grpcutil: port required")
	}

	if len(options.services) == 0 {
		log.Fatal("grpcutil: the list of gRPC services is required")
	}

	if options.cert == nil {
		log.Fatal("grpcutil: a TLS key pair is required")
	}

	options.serverOpts = append(options.serverOpts,
		grpc.Creds(credentials.NewServerTLSFromCert(options.cert)))

	// In order to support self-signed certificates we make sure to add
	// the provided certificate to gRPC gateway cert pool.
	certPool := x509.NewCertPool()
	x509Cert, err := x509.ParseCertificate(options.cert.Certificate[0])
	if err != nil {
		log.Fatalf("grpcutil: failed parsing x509 certificate: %+v, ", err)
	}
	certPool.AddCert(x509Cert)
	clientCreds := credentials.NewClientTLSFromCert(certPool, "")

	clientOpts := []grpc.DialOption{
		grpc.WithTransportCredentials(clientCreds),
		grpc.WithUserAgent("grpc-gw"),
	}

	if options.compressor != nil {
		options.serverOpts = append(options.serverOpts,
			grpc.RPCCompressor(options.compressor))

		if options.decompressor == nil {
			log.Fatal("grpcutil: gRPC gateway client needs the decompressor")
		}
	}

	if options.decompressor != nil {
		options.serverOpts = append(options.serverOpts, grpc.RPCDecompressor(options.decompressor))
		clientOpts = append(clientOpts, grpc.WithDecompressor(options.decompressor))
	}

	server := grpc.NewServer(options.serverOpts...)
	address := net.JoinHostPort("localhost", options.port)
	clientConn, err := grpc.Dial(address, clientOpts...)
	if err != nil {
		log.Fatalf("failed to connect to local gRPC server: %v", err)
	}

	var muxOpts []runtime.ServeMuxOption
	if options.errorHandler != nil {
		muxOpts = append(muxOpts, runtime.WithProtoErrorHandler(options.errorHandler))
	}

	gwMuxer := runtime.NewServeMux(muxOpts...)
	serviceBinding := ServiceBinding{
		GRPCServer:        server,
		GRPCGatewayClient: clientConn,
		GRPCGatewayMuxer:  gwMuxer,
	}

	log.Println("registering GRPC services...")
	for _, register := range options.services {
		if err := register(serviceBinding); err != nil {
			log.Fatalf("failed to register service: %+v", err)
		}
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		contentType := r.Header.Get("Content-Type")
		if r.ProtoMajor == 2 && strings.Contains(contentType, "application/grpc") {
			server.ServeHTTP(w, r)
			return
		}

		for _, p := range options.skipPaths {
			if strings.HasPrefix(r.URL.Path, p) {
				h.ServeHTTP(w, r)
				return
			}
		}

		accept := r.Header.Get("Accept")
		if strings.Contains(contentType, "application/json") ||
			strings.Contains(accept, "application/json") {
			gwMuxer.ServeHTTP(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}
<commit_msg>grpcutil: Enable passing grpc-gateway mux options<commit_after>// Package grpcutil implements a gRPC HTTP handler and OpenAPI proxy. Useful to serve gRPC requests from an existing HTTP servers.
package grpcutil

import (
	"crypto/tls"
	"crypto/x509"
	"log"
	"net"
	"net/http"
	"strings"

	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// ServiceBinding are the gRPC server, and gRPC HTTP gateway to which the service will be bound to.
type ServiceBinding struct {
	GRPCServer        *grpc.Server
	GRPCGatewayClient *grpc.ClientConn
	GRPCGatewayMuxer  *runtime.ServeMux
}

// ServiceRegisterFn defines a function type for registering gRPC services.
type ServiceRegisterFn func(ServiceBinding) error

// Option implements http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type Option func(*options)

type options struct {
	serverOpts   []grpc.ServerOption
	gwServerOpts []runtime.ServeMuxOption
	services     []ServiceRegisterFn
	cert         *tls.Certificate
	compressor   grpc.Compressor
	decompressor grpc.Decompressor
	port         string
	skipPaths    []string
}

// WithServerOpts sets gRPC server options. Optional.
func WithServerOpts(opts []grpc.ServerOption) Option {
	return func(o *options) {
		o.serverOpts = opts
	}
}

// WithGWServerOpts sets GRPC Gateway server options. Optional.
func WithGWServerOpts(opts []runtime.ServeMuxOption) Option {
	return func(o *options) {
		o.gwServerOpts = opts
	}
}

// WithServices sets the list of services to register with the gRPC and OpenAPI servers. Required.
func WithServices(services []ServiceRegisterFn) Option {
	return func(o *options) {
		o.services = services
	}
}

// WithTLSCert sets the TLS certificate to use by the gRPC server and OpenAPI gRPC client. Required.
func WithTLSCert(cert *tls.Certificate) Option {
	return func(o *options) {
		o.cert = cert
	}
}

// WithPort sets the port for the OpenAPI gRPC client to use when connecting to the gRPC service.
func WithPort(port string) Option {
	return func(o *options) {
		o.port = port
	}
}

// WithCompressor sets gRPC server and gateway client compressor.
// Use this function instead of passing compressor through grcputil.WithServerOpts()
func WithCompressor(c grpc.Compressor) Option {
	return func(o *options) {
		o.compressor = c
	}
}

// WithDecompressor sets gRPC server and gateway client decompressor.
// Use this function instead of passing a decompressor through grcputil.WithServerOpts()
func WithDecompressor(d grpc.Decompressor) Option {
	return func(o *options) {
		o.decompressor = d
	}
}

// WithSkipPath allows other handlers to serve static JSON files by instructing the GRPC Gateway Muxer to
// skip serving the given prefixed paths.
func WithSkipPath(path ...string) Option {
	return func(o *options) {
		o.skipPaths = path
	}
}

// Handler serves gRPC and OpenAPI requests or hands over to the next handler if they are not gRPC
// or application/json requests.
func Handler(h http.Handler, opts ...Option) http.Handler {
	options := new(options)
	for _, opt := range opts {
		opt(options)
	}

	if options.cert == nil {
		log.Fatal("grpcutil: TLS certificate required")
	}

	if options.port == "" {
		log.Fatal("grpcutil: port required")
	}

	if len(options.services) == 0 {
		log.Fatal("grpcutil: the list of gRPC services is required")
	}

	if options.cert == nil {
		log.Fatal("grpcutil: a TLS key pair is required")
	}

	options.serverOpts = append(options.serverOpts,
		grpc.Creds(credentials.NewServerTLSFromCert(options.cert)))

	// In order to support self-signed certificates we make sure to add
	// the provided certificate to gRPC gateway cert pool.
	certPool := x509.NewCertPool()
	x509Cert, err := x509.ParseCertificate(options.cert.Certificate[0])
	if err != nil {
		log.Fatalf("grpcutil: failed parsing x509 certificate: %+v, ", err)
	}
	certPool.AddCert(x509Cert)
	clientCreds := credentials.NewClientTLSFromCert(certPool, "")

	clientOpts := []grpc.DialOption{
		grpc.WithTransportCredentials(clientCreds),
		grpc.WithUserAgent("grpc-gw"),
	}

	if options.compressor != nil {
		options.serverOpts = append(options.serverOpts,
			grpc.RPCCompressor(options.compressor))

		if options.decompressor == nil {
			log.Fatal("grpcutil: gRPC gateway client needs the decompressor")
		}
	}

	if options.decompressor != nil {
		options.serverOpts = append(options.serverOpts, grpc.RPCDecompressor(options.decompressor))
		clientOpts = append(clientOpts, grpc.WithDecompressor(options.decompressor))
	}

	server := grpc.NewServer(options.serverOpts...)
	address := net.JoinHostPort("localhost", options.port)
	clientConn, err := grpc.Dial(address, clientOpts...)
	if err != nil {
		log.Fatalf("failed to connect to local gRPC server: %v", err)
	}

	gwMuxer := runtime.NewServeMux(options.gwServerOpts...)
	serviceBinding := ServiceBinding{
		GRPCServer:        server,
		GRPCGatewayClient: clientConn,
		GRPCGatewayMuxer:  gwMuxer,
	}

	log.Println("registering GRPC services...")
	for _, register := range options.services {
		if err := register(serviceBinding); err != nil {
			log.Fatalf("failed to register service: %+v", err)
		}
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		contentType := r.Header.Get("Content-Type")
		if r.ProtoMajor == 2 && strings.Contains(contentType, "application/grpc") {
			server.ServeHTTP(w, r)
			return
		}

		for _, p := range options.skipPaths {
			if strings.HasPrefix(r.URL.Path, p) {
				h.ServeHTTP(w, r)
				return
			}
		}

		accept := r.Header.Get("Accept")
		if strings.Contains(contentType, "application/json") ||
			strings.Contains(accept, "application/json") {
			gwMuxer.ServeHTTP(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}
<|endoftext|>


<commit_before>package core

import (
	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

type GeneralMessageUpdateEventBuilder struct {
	model.ClockConsumer
	model.UUIDConsumer

	partner *Partner
}

func NewGeneralMessageUpdateEventBuilder(partner *Partner) GeneralMessageUpdateEventBuilder {
	return GeneralMessageUpdateEventBuilder{partner: partner}
}

func (builder *GeneralMessageUpdateEventBuilder) SetGeneralMessageDeliveryUpdateEvents(event *[]*model.SituationUpdateEvent, xmlResponse *siri.XMLGeneralMessageDelivery) {
	xmlGeneralMessageEvents := xmlResponse.XMLGeneralMessages()
	if len(xmlGeneralMessageEvents) == 0 {
		return
	}

	for _, xmlGeneralMessageEvents := range xmlGeneralMessageEvents {
		builder.buildGeneralMessageUpdateEvent(event, xmlGeneralMessageEvents, xmlResponse.ProducerRef())
	}
}

func (builder *GeneralMessageUpdateEventBuilder) SetGeneralMessageResponseUpdateEvents(event *[]*model.SituationUpdateEvent, xmlResponse *siri.XMLGeneralMessageResponse) {
	xmlGeneralMessageEvents := xmlResponse.XMLGeneralMessages()
	if len(xmlGeneralMessageEvents) == 0 {
		return
	}

	for _, xmlGeneralMessageEvents := range xmlGeneralMessageEvents {
		builder.buildGeneralMessageUpdateEvent(event, xmlGeneralMessageEvents, xmlResponse.ProducerRef())
	}
}

func (builder *GeneralMessageUpdateEventBuilder) buildGeneralMessageUpdateEvent(event *[]*model.SituationUpdateEvent, xmlGeneralMessageEvent *siri.XMLGeneralMessage, producerRef string) {
	if xmlGeneralMessageEvent.Content() == nil {
		return
	}

	situationEvent := &model.SituationUpdateEvent{
		Origin:            string(builder.partner.Slug()),
		CreatedAt:         builder.Clock().Now(),
		RecordedAt:        xmlGeneralMessageEvent.RecordedAtTime(),
		SituationObjectID: model.NewObjectID(builder.partner.Setting("remote_objectid_kind"), xmlGeneralMessageEvent.InfoMessageIdentifier()),
		Version:           xmlGeneralMessageEvent.InfoMessageVersion(),
		ProducerRef:       producerRef,
	}
	situationEvent.SetId(model.SituationUpdateRequestId(builder.NewUUID()))

	situationEvent.SituationAttributes.Format = xmlGeneralMessageEvent.FormatRef()
	situationEvent.SituationAttributes.Channel = xmlGeneralMessageEvent.InfoChannelRef()
	situationEvent.SituationAttributes.ValidUntil = xmlGeneralMessageEvent.ValidUntilTime()

	content := xmlGeneralMessageEvent.Content().(siri.IDFGeneralMessageStructure)
	for _, xmlMessage := range content.Messages() {
		message := &model.Message{
			Content:             xmlMessage.MessageText(),
			Type:                xmlMessage.MessageType(),
			NumberOfLines:       xmlMessage.NumberOfLines(),
			NumberOfCharPerLine: xmlMessage.NumberOfCharPerLine(),
		}
		situationEvent.SituationAttributes.Messages = append(situationEvent.SituationAttributes.Messages, message)
	}

	builder.setReferences(situationEvent, &content)

	*event = append(*event, situationEvent)
}

func (builder *GeneralMessageUpdateEventBuilder) setReferences(event *model.SituationUpdateEvent, content *siri.IDFGeneralMessageStructure) {
	tx := builder.partner.Referential().NewTransaction()
	defer tx.Close()

	remoteObjectidKind := builder.partner.Setting("remote_objectid_kind")

	if lineRefs := content.LineRef(); len(lineRefs) != 0 {
		for i := range lineRefs {
			ref := model.NewReference(model.NewObjectID(remoteObjectidKind, lineRefs[i]))
			ref.Type = "LineRef"
			event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
		}
	}
	if stopPointRefs := content.StopPointRef(); len(stopPointRefs) != 0 {
		for i := range stopPointRefs {
			ref := model.NewReference(model.NewObjectID(remoteObjectidKind, stopPointRefs[i]))
			ref.Type = "StopPointRef"
			event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
		}
	}
	if journeyPatternRefs := content.JourneyPatternRef(); len(journeyPatternRefs) != 0 {
		for i := range journeyPatternRefs {
			ref := model.NewReference(model.NewObjectID(remoteObjectidKind, journeyPatternRefs[i]))
			ref.Type = "JourneyPatternRef"
			event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
		}
	}
	if destinationRefs := content.DestinationRef(); len(destinationRefs) != 0 {
		for i := range destinationRefs {
			ref := model.NewReference(model.NewObjectID(remoteObjectidKind, destinationRefs[i]))
			ref.Type = "DestinationRef"
			event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
		}
	}
	if routeRefs := content.RouteRef(); len(routeRefs) != 0 {
		for i := range routeRefs {
			ref := model.NewReference(model.NewObjectID(remoteObjectidKind, routeRefs[i]))
			ref.Type = "RouteRef"
			event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
		}
	}
	if groupOfLinesRefs := content.GroupOfLinesRef(); len(groupOfLinesRefs) != 0 {
		for i := range groupOfLinesRefs {
			ref := model.NewReference(model.NewObjectID(remoteObjectidKind, groupOfLinesRefs[i]))
			ref.Type = "GroupOfLinesRef"
			event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
		}
	}
	if lineSections := content.LineSections(); len(lineSections) != 0 {
		for _, lineSection := range lineSections {
			builder.handleLineSection(tx, remoteObjectidKind, lineSection, event)
		}
	}
}

func (builder *GeneralMessageUpdateEventBuilder) handleLineSection(tx *model.Transaction, remoteObjectidKind string, lineSection *siri.IDFLineSectionStructure, event *model.SituationUpdateEvent) {
	references := model.NewReferences()

	lineRef := model.NewReference(model.NewObjectID(remoteObjectidKind, lineSection.LineRef()))
	references.Set("LineRef", *lineRef)

	firstStopRef := model.NewReference(model.NewObjectID(remoteObjectidKind, lineSection.FirstStop()))
	references.Set("FirstStop", *firstStopRef)

	lastStopRef := model.NewReference(model.NewObjectID(remoteObjectidKind, lineSection.LastStop()))
	references.Set("LastStop", *lastStopRef)

	event.SituationAttributes.LineSections = append(event.SituationAttributes.LineSections, &references)
}
<commit_msg>Cleaner GeneralMessageUpdateEventBuilder#setReferences<commit_after>package core

import (
	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

type GeneralMessageUpdateEventBuilder struct {
	model.ClockConsumer
	model.UUIDConsumer

	partner *Partner
}

func NewGeneralMessageUpdateEventBuilder(partner *Partner) GeneralMessageUpdateEventBuilder {
	return GeneralMessageUpdateEventBuilder{partner: partner}
}

func (builder *GeneralMessageUpdateEventBuilder) SetGeneralMessageDeliveryUpdateEvents(event *[]*model.SituationUpdateEvent, xmlResponse *siri.XMLGeneralMessageDelivery) {
	xmlGeneralMessageEvents := xmlResponse.XMLGeneralMessages()
	if len(xmlGeneralMessageEvents) == 0 {
		return
	}

	for _, xmlGeneralMessageEvents := range xmlGeneralMessageEvents {
		builder.buildGeneralMessageUpdateEvent(event, xmlGeneralMessageEvents, xmlResponse.ProducerRef())
	}
}

func (builder *GeneralMessageUpdateEventBuilder) SetGeneralMessageResponseUpdateEvents(event *[]*model.SituationUpdateEvent, xmlResponse *siri.XMLGeneralMessageResponse) {
	xmlGeneralMessageEvents := xmlResponse.XMLGeneralMessages()
	if len(xmlGeneralMessageEvents) == 0 {
		return
	}

	for _, xmlGeneralMessageEvents := range xmlGeneralMessageEvents {
		builder.buildGeneralMessageUpdateEvent(event, xmlGeneralMessageEvents, xmlResponse.ProducerRef())
	}
}

func (builder *GeneralMessageUpdateEventBuilder) buildGeneralMessageUpdateEvent(event *[]*model.SituationUpdateEvent, xmlGeneralMessageEvent *siri.XMLGeneralMessage, producerRef string) {
	if xmlGeneralMessageEvent.Content() == nil {
		return
	}

	situationEvent := &model.SituationUpdateEvent{
		Origin:            string(builder.partner.Slug()),
		CreatedAt:         builder.Clock().Now(),
		RecordedAt:        xmlGeneralMessageEvent.RecordedAtTime(),
		SituationObjectID: model.NewObjectID(builder.partner.Setting("remote_objectid_kind"), xmlGeneralMessageEvent.InfoMessageIdentifier()),
		Version:           xmlGeneralMessageEvent.InfoMessageVersion(),
		ProducerRef:       producerRef,
	}
	situationEvent.SetId(model.SituationUpdateRequestId(builder.NewUUID()))

	situationEvent.SituationAttributes.Format = xmlGeneralMessageEvent.FormatRef()
	situationEvent.SituationAttributes.Channel = xmlGeneralMessageEvent.InfoChannelRef()
	situationEvent.SituationAttributes.ValidUntil = xmlGeneralMessageEvent.ValidUntilTime()

	content := xmlGeneralMessageEvent.Content().(siri.IDFGeneralMessageStructure)
	for _, xmlMessage := range content.Messages() {
		message := &model.Message{
			Content:             xmlMessage.MessageText(),
			Type:                xmlMessage.MessageType(),
			NumberOfLines:       xmlMessage.NumberOfLines(),
			NumberOfCharPerLine: xmlMessage.NumberOfCharPerLine(),
		}
		situationEvent.SituationAttributes.Messages = append(situationEvent.SituationAttributes.Messages, message)
	}

	builder.setReferences(situationEvent, &content)

	*event = append(*event, situationEvent)
}

func (builder *GeneralMessageUpdateEventBuilder) setReferences(event *model.SituationUpdateEvent, content *siri.IDFGeneralMessageStructure) {
	tx := builder.partner.Referential().NewTransaction()
	defer tx.Close()

	remoteObjectidKind := builder.partner.Setting("remote_objectid_kind")

	for _, lineref := range content.LineRef() {
		ref := model.NewReference(model.NewObjectID(remoteObjectidKind, lineref))
		ref.Type = "LineRef"
		event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
	}
	for _, stoppointref := range content.StopPointRef() {
		ref := model.NewReference(model.NewObjectID(remoteObjectidKind, stoppointref))
		ref.Type = "StopPointRef"
		event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
	}
	for _, journeypatternref := range content.JourneyPatternRef() {
		ref := model.NewReference(model.NewObjectID(remoteObjectidKind, journeypatternref))
		ref.Type = "JourneyPatternRef"
		event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
	}
	for _, destinationref := range content.DestinationRef() {
		ref := model.NewReference(model.NewObjectID(remoteObjectidKind, destinationref))
		ref.Type = "DestinationRef"
		event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
	}
	for _, routeref := range content.RouteRef() {
		ref := model.NewReference(model.NewObjectID(remoteObjectidKind, routeref))
		ref.Type = "RouteRef"
		event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
	}
	for _, groupoflinesref := range content.GroupOfLinesRef() {
		ref := model.NewReference(model.NewObjectID(remoteObjectidKind, groupoflinesref))
		ref.Type = "GroupOfLinesRef"
		event.SituationAttributes.References = append(event.SituationAttributes.References, ref)
	}
	for _, lineSection := range content.LineSections() {
		builder.handleLineSection(tx, remoteObjectidKind, lineSection, event)
	}
}

func (builder *GeneralMessageUpdateEventBuilder) handleLineSection(tx *model.Transaction, remoteObjectidKind string, lineSection *siri.IDFLineSectionStructure, event *model.SituationUpdateEvent) {
	references := model.NewReferences()

	lineRef := model.NewReference(model.NewObjectID(remoteObjectidKind, lineSection.LineRef()))
	references.Set("LineRef", *lineRef)

	firstStopRef := model.NewReference(model.NewObjectID(remoteObjectidKind, lineSection.FirstStop()))
	references.Set("FirstStop", *firstStopRef)

	lastStopRef := model.NewReference(model.NewObjectID(remoteObjectidKind, lineSection.LastStop()))
	references.Set("LastStop", *lastStopRef)

	event.SituationAttributes.LineSections = append(event.SituationAttributes.LineSections, &references)
}
<|endoftext|>


<commit_before>package raft

import (
	"testing"
)

// #5.2-p3s1: A candidate wins an election if it receives votes from a
// majority of the servers in the full cluster for the same term.
// #RFS-L1a: Upon election: send initial empty AppendEntries RPCs (heartbeat)
// to each server;
func TestCM_RpcRVR_Candidate_CandidateWinsElectionIfItReceivesMajorityOfVotes(t *testing.T) {
	mcm, mrs := testSetupMCM_Candidate_Figure7LeaderLine(t)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()

	// s2 grants vote - should stay as candidate
	mcm.pcm.rpcReply("s2", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}

	// s3 denies vote - should stay as candidate
	mcm.pcm.rpcReply("s3", &RpcRequestVoteReply{serverTerm, false})
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}

	// s4 grants vote - should become leader
	mcm.pcm.rpcReply("s4", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != LEADER {
		t.Fatal()
	}

	// leader setup
	lastLogIndex, lastLogTerm := getIndexAndTermOfLastEntry(mcm.pcm.log)
	expectedRpc := &RpcAppendEntries{
		serverTerm,
		lastLogIndex,
		lastLogTerm,
		[]LogEntry{},
		0, // TODO: tests for this?!
	}
	expectedRpcs := []mockSentRpc{
		{"s2", expectedRpc},
		{"s3", expectedRpc},
		{"s4", expectedRpc},
		{"s5", expectedRpc},
	}
	mrs.checkSentRpcs(t, expectedRpcs)

	// s5 grants vote - should stay leader
	mcm.pcm.rpcReply("s5", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != LEADER {
		t.Fatal()
	}
}

// #5.2-p5s1: The third possible outcome is that a candidate neither
// wins nor loses the election; ... votes could be split so that no
// candidate obtains a majority.
// #5.2-p5s2: When this happens, each candidate will time out and
// start a new election by incrementing its term and initiating
// another round of RequestVote RPCs.
func TestCM_RpcRVR_Candidate_StartNewElectionOnElectionTimeout(t *testing.T) {
	mcm, mrs := testSetupMCM_Candidate_Figure7LeaderLine(t)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()

	// s2 grants vote - should stay as candidate
	mcm.pcm.rpcReply("s2", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}

	// s3 denies vote - should stay as candidate
	mcm.pcm.rpcReply("s3", &RpcRequestVoteReply{serverTerm, false})
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}

	// no more votes - election timeout causes a new election
	testCM_FollowerOrCandidate_StartsElectionOnElectionTimeout_Part2(t, mcm, mrs, testCurrentTerm+2)
}

// Extra: follower ignores vote
func TestCM_RpcRVR_Follower_Ignores(t *testing.T) {
	mcm, mrs := testSetupMCM_Follower_Figure7LeaderLine(t)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()

	// s2 grants vote - ignore
	mcm.pcm.rpcReply("s2", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}
	mrs.checkSentRpcs(t, []mockSentRpc{})

	// s3 denies vote - ignore
	mcm.pcm.rpcReply("s3", &RpcRequestVoteReply{serverTerm, false})
	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}
	mrs.checkSentRpcs(t, []mockSentRpc{})
}
<commit_msg>Leader ignores RequestVoteReply<commit_after>package raft

import (
	"testing"
)

// #5.2-p3s1: A candidate wins an election if it receives votes from a
// majority of the servers in the full cluster for the same term.
// #RFS-L1a: Upon election: send initial empty AppendEntries RPCs (heartbeat)
// to each server;
func TestCM_RpcRVR_Candidate_CandidateWinsElectionIfItReceivesMajorityOfVotes(t *testing.T) {
	mcm, mrs := testSetupMCM_Candidate_Figure7LeaderLine(t)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()

	// s2 grants vote - should stay as candidate
	mcm.pcm.rpcReply("s2", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}

	// s3 denies vote - should stay as candidate
	mcm.pcm.rpcReply("s3", &RpcRequestVoteReply{serverTerm, false})
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}

	// s4 grants vote - should become leader
	mcm.pcm.rpcReply("s4", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != LEADER {
		t.Fatal()
	}

	// leader setup
	lastLogIndex, lastLogTerm := getIndexAndTermOfLastEntry(mcm.pcm.log)
	expectedRpc := &RpcAppendEntries{
		serverTerm,
		lastLogIndex,
		lastLogTerm,
		[]LogEntry{},
		0, // TODO: tests for this?!
	}
	expectedRpcs := []mockSentRpc{
		{"s2", expectedRpc},
		{"s3", expectedRpc},
		{"s4", expectedRpc},
		{"s5", expectedRpc},
	}
	mrs.checkSentRpcs(t, expectedRpcs)

	// s5 grants vote - should stay leader
	mcm.pcm.rpcReply("s5", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != LEADER {
		t.Fatal()
	}
}

// #5.2-p5s1: The third possible outcome is that a candidate neither
// wins nor loses the election; ... votes could be split so that no
// candidate obtains a majority.
// #5.2-p5s2: When this happens, each candidate will time out and
// start a new election by incrementing its term and initiating
// another round of RequestVote RPCs.
func TestCM_RpcRVR_Candidate_StartNewElectionOnElectionTimeout(t *testing.T) {
	mcm, mrs := testSetupMCM_Candidate_Figure7LeaderLine(t)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()

	// s2 grants vote - should stay as candidate
	mcm.pcm.rpcReply("s2", &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}

	// s3 denies vote - should stay as candidate
	mcm.pcm.rpcReply("s3", &RpcRequestVoteReply{serverTerm, false})
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}

	// no more votes - election timeout causes a new election
	testCM_FollowerOrCandidate_StartsElectionOnElectionTimeout_Part2(t, mcm, mrs, testCurrentTerm+2)
}

// Extra: follower or leader ignores vote
func TestCM_RpcRVR_FollowerOrLeader_Ignores(t *testing.T) {
	f := func(setup func(t *testing.T) (mcm *managedConsensusModule, mrs *mockRpcSender)) {
		mcm, mrs := setup(t)
		serverTerm := mcm.pcm.persistentState.GetCurrentTerm()
		beforeState := mcm.pcm.getServerState()

		// s2 grants vote - ignore
		mcm.pcm.rpcReply("s2", &RpcRequestVoteReply{serverTerm, true})
		if mcm.pcm.getServerState() != beforeState {
			t.Fatal()
		}
		mrs.checkSentRpcs(t, []mockSentRpc{})

		// s3 denies vote - ignore
		mcm.pcm.rpcReply("s3", &RpcRequestVoteReply{serverTerm, false})
		if mcm.pcm.getServerState() != beforeState {
			t.Fatal()
		}
		mrs.checkSentRpcs(t, []mockSentRpc{})
	}

	f(testSetupMCM_Leader_Figure7LeaderLine)
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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.

// This file implements a cadvisor datasource, that collects metrics from an instance
// of cadvisor running on a specific host.

package kubelet

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/golang/glog"
	cadvisor "github.com/google/cadvisor/info/v1"
	kubelet_client "k8s.io/heapster/metrics/sources/kubelet/util"
	"k8s.io/kubernetes/pkg/kubelet/api/v1alpha1/stats"
)

type Host struct {
	IP       net.IP
	Port     int
	Resource string
}

func (h Host) String() string {
	return net.JoinHostPort(h.IP.String(), strconv.Itoa(h.Port))
}

type KubeletClient struct {
	config *kubelet_client.KubeletClientConfig
	client *http.Client
}

type ErrNotFound struct {
	endpoint string
}

func (err *ErrNotFound) Error() string {
	return fmt.Sprintf("%q not found", err.endpoint)
}

func IsNotFoundError(err error) bool {
	_, isNotFound := err.(*ErrNotFound)
	return isNotFound
}

func sampleContainerStats(stats []*cadvisor.ContainerStats) []*cadvisor.ContainerStats {
	if len(stats) == 0 {
		return []*cadvisor.ContainerStats{}
	}
	return []*cadvisor.ContainerStats{stats[len(stats)-1]}
}

func (self *KubeletClient) postRequestAndGetValue(client *http.Client, req *http.Request, value interface{}) error {
	response, err := client.Do(req)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return fmt.Errorf("failed to read response body - %v", err)
	}
	if response.StatusCode == http.StatusNotFound {
		return &ErrNotFound{req.URL.String()}
	} else if response.StatusCode != http.StatusOK {
		return fmt.Errorf("request failed - %q, response: %q", response.Status, string(body))
	}

	kubeletAddr := "[unknown]"
	if req.URL != nil {
		kubeletAddr = req.URL.Host
	}
	glog.V(10).Infof("Raw response from Kubelet at %s: %s", kubeletAddr, string(body))

	err = json.Unmarshal(body, value)
	if err != nil {
		return fmt.Errorf("failed to parse output. Response: %q. Error: %v", string(body), err)
	}
	return nil
}

func (self *KubeletClient) parseStat(containerInfo *cadvisor.ContainerInfo) *cadvisor.ContainerInfo {
	containerInfo.Stats = sampleContainerStats(containerInfo.Stats)
	if len(containerInfo.Aliases) > 0 {
		containerInfo.Name = containerInfo.Aliases[0]
	}
	return containerInfo
}

// TODO(vmarmol): Use Kubernetes' if we export it as an API.
type statsRequest struct {
	// The name of the container for which to request stats.
	// Default: /
	ContainerName string `json:"containerName,omitempty"`

	// Max number of stats to return.
	// If start and end time are specified this limit is ignored.
	// Default: 60
	NumStats int `json:"num_stats,omitempty"`

	// Start time for which to query information.
	// If omitted, the beginning of time is assumed.
	Start time.Time `json:"start,omitempty"`

	// End time for which to query information.
	// If omitted, current time is assumed.
	End time.Time `json:"end,omitempty"`

	// Whether to also include information from subcontainers.
	// Default: false.
	Subcontainers bool `json:"subcontainers,omitempty"`
}

// Get stats for all non-Kubernetes containers.
func (self *KubeletClient) GetAllRawContainers(host Host, start, end time.Time) ([]cadvisor.ContainerInfo, error) {
	scheme := "http"
	if self.config != nil && self.config.EnableHttps {
		scheme = "https"
	}

	url := url.URL{
		Scheme: scheme,
		Host:   host.String(),
		Path:   "/stats/container/",
	}

	return self.getAllContainers(url.String(), start, end)
}

func (self *KubeletClient) GetSummary(host Host) (*stats.Summary, error) {
	url := url.URL{
		Scheme: "http",
		Host:   host.String(),
		Path:   "/stats/summary/",
	}
	if self.config != nil && self.config.EnableHttps {
		url.Scheme = "https"
	}

	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		return nil, err
	}
	summary := &stats.Summary{}
	client := self.client
	if client == nil {
		client = http.DefaultClient
	}
	err = self.postRequestAndGetValue(client, req, summary)
	return summary, err
}

func (self *KubeletClient) GetPort() int {
	return int(self.config.Port)
}

func (self *KubeletClient) getAllContainers(url string, start, end time.Time) ([]cadvisor.ContainerInfo, error) {
	// Request data from all subcontainers.
	request := statsRequest{
		ContainerName: "/",
		NumStats:      1,
		Start:         start,
		End:           end,
		Subcontainers: true,
	}
	body, err := json.Marshal(request)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	var containers map[string]cadvisor.ContainerInfo
	client := self.client
	if client == nil {
		client = http.DefaultClient
	}
	err = self.postRequestAndGetValue(client, req, &containers)
	if err != nil {
		return nil, fmt.Errorf("failed to get all container stats from Kubelet URL %q: %v", url, err)
	}

	result := make([]cadvisor.ContainerInfo, 0, len(containers))
	for _, containerInfo := range containers {
		cont := self.parseStat(&containerInfo)
		if cont != nil {
			result = append(result, *cont)
		}
	}
	return result, nil
}

func NewKubeletClient(kubeletConfig *kubelet_client.KubeletClientConfig) (*KubeletClient, error) {
	transport, err := kubelet_client.MakeTransport(kubeletConfig)
	if err != nil {
		return nil, err
	}
	c := &http.Client{
		Transport: transport,
		Timeout:   kubeletConfig.HTTPTimeout,
	}
	return &KubeletClient{
		config: kubeletConfig,
		client: c,
	}, nil
}
<commit_msg>Code refactoring and simplification<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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.

// This file implements a cadvisor datasource, that collects metrics from an instance
// of cadvisor running on a specific host.

package kubelet

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/golang/glog"
	cadvisor "github.com/google/cadvisor/info/v1"
	kubelet_client "k8s.io/heapster/metrics/sources/kubelet/util"
	"k8s.io/kubernetes/pkg/kubelet/api/v1alpha1/stats"
)

type Host struct {
	IP       net.IP
	Port     int
	Resource string
}

func (h Host) String() string {
	return net.JoinHostPort(h.IP.String(), strconv.Itoa(h.Port))
}

type KubeletClient struct {
	config *kubelet_client.KubeletClientConfig
	client *http.Client
}

type ErrNotFound struct {
	endpoint string
}

func (err *ErrNotFound) Error() string {
	return fmt.Sprintf("%q not found", err.endpoint)
}

func IsNotFoundError(err error) bool {
	_, isNotFound := err.(*ErrNotFound)
	return isNotFound
}

func sampleContainerStats(stats []*cadvisor.ContainerStats) []*cadvisor.ContainerStats {
	if len(stats) == 0 {
		return []*cadvisor.ContainerStats{}
	}
	return []*cadvisor.ContainerStats{stats[len(stats)-1]}
}

func (self *KubeletClient) postRequestAndGetValue(client *http.Client, req *http.Request, value interface{}) error {
	response, err := client.Do(req)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return fmt.Errorf("failed to read response body - %v", err)
	}
	if response.StatusCode == http.StatusNotFound {
		return &ErrNotFound{req.URL.String()}
	} else if response.StatusCode != http.StatusOK {
		return fmt.Errorf("request failed - %q, response: %q", response.Status, string(body))
	}

	kubeletAddr := "[unknown]"
	if req.URL != nil {
		kubeletAddr = req.URL.Host
	}
	glog.V(10).Infof("Raw response from Kubelet at %s: %s", kubeletAddr, string(body))

	err = json.Unmarshal(body, value)
	if err != nil {
		return fmt.Errorf("failed to parse output. Response: %q. Error: %v", string(body), err)
	}
	return nil
}

func (self *KubeletClient) parseStat(containerInfo *cadvisor.ContainerInfo) *cadvisor.ContainerInfo {
	containerInfo.Stats = sampleContainerStats(containerInfo.Stats)
	if len(containerInfo.Aliases) > 0 {
		containerInfo.Name = containerInfo.Aliases[0]
	}
	return containerInfo
}

// TODO(vmarmol): Use Kubernetes' if we export it as an API.
type statsRequest struct {
	// The name of the container for which to request stats.
	// Default: /
	ContainerName string `json:"containerName,omitempty"`

	// Max number of stats to return.
	// If start and end time are specified this limit is ignored.
	// Default: 60
	NumStats int `json:"num_stats,omitempty"`

	// Start time for which to query information.
	// If omitted, the beginning of time is assumed.
	Start time.Time `json:"start,omitempty"`

	// End time for which to query information.
	// If omitted, current time is assumed.
	End time.Time `json:"end,omitempty"`

	// Whether to also include information from subcontainers.
	// Default: false.
	Subcontainers bool `json:"subcontainers,omitempty"`
}

func (self *KubeletClient) getScheme() string {
	if self.config != nil && self.config.EnableHttps {
		return "https"
	} else {
		return "http"
	}
}

func (self *KubeletClient) getUrl(host Host, path string) string {
	url := url.URL{
		Scheme: self.getScheme(),
		Host:   host.String(),
		Path:   path,
	}

	return url.String()
}

// Get stats for all non-Kubernetes containers.
func (self *KubeletClient) GetAllRawContainers(host Host, start, end time.Time) ([]cadvisor.ContainerInfo, error) {
	url := self.getUrl(host, "/stats/container/")

	return self.getAllContainers(url, start, end)
}

func (self *KubeletClient) GetSummary(host Host) (*stats.Summary, error) {
	url := self.getUrl(host, "/stats/summary/")

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	summary := &stats.Summary{}
	client := self.client
	if client == nil {
		client = http.DefaultClient
	}
	err = self.postRequestAndGetValue(client, req, summary)
	return summary, err
}

func (self *KubeletClient) GetPort() int {
	return int(self.config.Port)
}

func (self *KubeletClient) getAllContainers(url string, start, end time.Time) ([]cadvisor.ContainerInfo, error) {
	// Request data from all subcontainers.
	request := statsRequest{
		ContainerName: "/",
		NumStats:      1,
		Start:         start,
		End:           end,
		Subcontainers: true,
	}
	body, err := json.Marshal(request)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	var containers map[string]cadvisor.ContainerInfo
	client := self.client
	if client == nil {
		client = http.DefaultClient
	}
	err = self.postRequestAndGetValue(client, req, &containers)
	if err != nil {
		return nil, fmt.Errorf("failed to get all container stats from Kubelet URL %q: %v", url, err)
	}

	result := make([]cadvisor.ContainerInfo, 0, len(containers))
	for _, containerInfo := range containers {
		cont := self.parseStat(&containerInfo)
		if cont != nil {
			result = append(result, *cont)
		}
	}
	return result, nil
}

func NewKubeletClient(kubeletConfig *kubelet_client.KubeletClientConfig) (*KubeletClient, error) {
	transport, err := kubelet_client.MakeTransport(kubeletConfig)
	if err != nil {
		return nil, err
	}
	c := &http.Client{
		Transport: transport,
		Timeout:   kubeletConfig.HTTPTimeout,
	}
	return &KubeletClient{
		config: kubeletConfig,
		client: c,
	}, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc.
//
// 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 cpuload

import (
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/stretchrcom/testify/assert"

	"github.com/soumya92/barista/bar"
	"github.com/soumya92/barista/outputs"
	"github.com/soumya92/barista/scheduler"
	testBar "github.com/soumya92/barista/testing/bar"
)

var syncMutex sync.Mutex
var simulatedLoads LoadAvg
var simulatedCount int
var simulatedErr error

func shouldError(err error) {
	syncMutex.Lock()
	defer syncMutex.Unlock()
	simulatedErr = err
	simulatedCount = 0
	simulatedLoads = LoadAvg{}
}

func shouldReturn(loads ...float64) {
	syncMutex.Lock()
	defer syncMutex.Unlock()
	simulatedErr = nil
	simulatedCount = len(loads)
	for i, l := range loads {
		if i < len(simulatedLoads) {
			simulatedLoads[i] = l
		}
	}
}

var mockloadavg = func(out *LoadAvg, count int) (int, error) {
	syncMutex.Lock()
	defer syncMutex.Unlock()
	for i := 0; i < count; i++ {
		if i < len(simulatedLoads) {
			out[i] = simulatedLoads[i]
		}
	}
	return simulatedCount, simulatedErr
}

func TestCpuload(t *testing.T) {
	assert := assert.New(t)
	getloadavg = mockloadavg
	testBar.New(t)

	shouldReturn(0, 0, 0)

	load := New()
	testBar.Run(load)

	testBar.LatestOutput().AssertText(
		[]string{"0.00"}, "on start")

	shouldReturn(1, 2, 3)
	testBar.AssertNoOutput("until refresh")
	testBar.Tick()
	testBar.NextOutput().AssertText(
		[]string{"1.00"}, "on next tick")

	load.OutputTemplate(outputs.TextTemplate(`{{.Min5 | printf "%.2f"}}`))
	testBar.NextOutput().AssertText(
		[]string{"2.00"}, "on output format change")

	load.UrgentWhen(func(l LoadAvg) bool {
		return l.Min15() > 2
	})
	testBar.NextOutput().AssertEqual(
		outputs.Text("2.00").Urgent(true),
		"on urgent function change")

	load.OutputColor(func(l LoadAvg) bar.Color {
		return bar.Color("red")
	})
	testBar.NextOutput().AssertEqual(
		outputs.Text("2.00").Urgent(true).Color(bar.Color("red")),
		"on color function change")

	shouldReturn(0, 0, 0)
	testBar.Tick()
	testBar.NextOutput().AssertEqual(
		outputs.Text("0.00").Urgent(false).Color(bar.Color("red")),
		"on next tick")

	load.RefreshInterval(time.Minute)
	testBar.AssertNoOutput("on refresh interval change")

	beforeTick := scheduler.Now()
	afterTick := scheduler.NextTick()
	testBar.NextOutput().Expect("on next tick")
	assert.Equal(time.Minute, afterTick.Sub(beforeTick))

	testBar.AssertNoOutput("until next tick")
}

func TestErrors(t *testing.T) {
	assert := assert.New(t)
	getloadavg = mockloadavg
	testBar.New(t)

	load := New()
	testBar.Run(load)

	shouldReturn(1)
	testBar.Tick()
	errs := testBar.LatestOutput().AssertError("on next tick with error")
	assert.Equal("getloadavg: 1", errs[0], "error string contains getloadavg code")
	testBar.Click(0) // to restart.

	shouldReturn(1, 2, 3, 4, 5)
	errs = testBar.LatestOutput().AssertError("on next tick with error")
	assert.Equal("getloadavg: 5", errs[0], "error string contains getloadavg code")
	testBar.Click(0)

	shouldError(errors.New("test"))
	errs = testBar.LatestOutput().AssertError("on next tick with error")
	assert.Equal("test", errs[0], "error string is passed through")
}
<commit_msg>Fix race in cpuload test<commit_after>// Copyright 2017 Google Inc.
//
// 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 cpuload

import (
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/stretchrcom/testify/assert"

	"github.com/soumya92/barista/bar"
	"github.com/soumya92/barista/outputs"
	"github.com/soumya92/barista/scheduler"
	testBar "github.com/soumya92/barista/testing/bar"
)

var syncMutex sync.Mutex
var simulatedLoads LoadAvg
var simulatedCount int
var simulatedErr error

func shouldError(err error) {
	syncMutex.Lock()
	defer syncMutex.Unlock()
	simulatedErr = err
	simulatedCount = 0
	simulatedLoads = LoadAvg{}
}

func shouldReturn(loads ...float64) {
	syncMutex.Lock()
	defer syncMutex.Unlock()
	simulatedErr = nil
	simulatedCount = len(loads)
	for i, l := range loads {
		if i < len(simulatedLoads) {
			simulatedLoads[i] = l
		}
	}
}

var mockloadavg = func(out *LoadAvg, count int) (int, error) {
	syncMutex.Lock()
	defer syncMutex.Unlock()
	for i := 0; i < count; i++ {
		if i < len(simulatedLoads) {
			out[i] = simulatedLoads[i]
		}
	}
	return simulatedCount, simulatedErr
}

func TestCpuload(t *testing.T) {
	assert := assert.New(t)
	getloadavg = mockloadavg
	testBar.New(t)

	shouldReturn(0, 0, 0)

	load := New()
	testBar.Run(load)

	testBar.LatestOutput().AssertText(
		[]string{"0.00"}, "on start")

	shouldReturn(1, 2, 3)
	testBar.AssertNoOutput("until refresh")
	testBar.Tick()
	testBar.NextOutput().AssertText(
		[]string{"1.00"}, "on next tick")

	load.OutputTemplate(outputs.TextTemplate(`{{.Min5 | printf "%.2f"}}`))
	testBar.NextOutput().AssertText(
		[]string{"2.00"}, "on output format change")

	load.UrgentWhen(func(l LoadAvg) bool {
		return l.Min15() > 2
	})
	testBar.NextOutput().AssertEqual(
		outputs.Text("2.00").Urgent(true),
		"on urgent function change")

	load.OutputColor(func(l LoadAvg) bar.Color {
		return bar.Color("red")
	})
	testBar.NextOutput().AssertEqual(
		outputs.Text("2.00").Urgent(true).Color(bar.Color("red")),
		"on color function change")

	shouldReturn(0, 0, 0)
	testBar.Tick()
	testBar.NextOutput().AssertEqual(
		outputs.Text("0.00").Urgent(false).Color(bar.Color("red")),
		"on next tick")

	load.RefreshInterval(time.Minute)
	testBar.AssertNoOutput("on refresh interval change")

	beforeTick := scheduler.Now()
	afterTick := scheduler.NextTick()
	testBar.NextOutput().Expect("on next tick")
	assert.Equal(time.Minute, afterTick.Sub(beforeTick))

	testBar.AssertNoOutput("until next tick")
}

func TestErrors(t *testing.T) {
	assert := assert.New(t)
	getloadavg = mockloadavg
	testBar.New(t)

	load := New()
	testBar.Run(load)

	shouldReturn(1)
	testBar.Tick()
	errs := testBar.LatestOutput().AssertError("on next tick with error")
	assert.Equal("getloadavg: 1", errs[0], "error string contains getloadavg code")

	shouldReturn(1, 2, 3, 4, 5)
	testBar.Click(0) // to restart.
	errs = testBar.LatestOutput().AssertError("on next tick with error")
	assert.Equal("getloadavg: 5", errs[0], "error string contains getloadavg code")

	shouldError(errors.New("test"))
	testBar.Click(0)
	errs = testBar.LatestOutput().AssertError("on next tick with error")
	assert.Equal("test", errs[0], "error string is passed through")
}
<|endoftext|>


<commit_before>package contractor

import (
	"bytes"
	"os"
	"strconv"
	"testing"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/modules/renter/proto"
	"github.com/NebulousLabs/Sia/types"
)

// memPersist implements the persister interface in-memory.
type memPersist contractorPersist

func (m *memPersist) save(data contractorPersist) error { *m = memPersist(data); return nil }
func (m *memPersist) update(...journalUpdate) error     { return nil }
func (m memPersist) load(data *contractorPersist) error { *data = contractorPersist(m); return nil }
func (m memPersist) Close() error                       { return nil }

// TestSaveLoad tests that the contractor can save and load itself.
func TestSaveLoad(t *testing.T) {
	// create contractor with mocked persist dependency
	c := &Contractor{
		persist: new(memPersist),
	}

	// add some fake contracts
	c.contracts = proto.NewContractSet([]modules.RenterContract{
		{ID: types.FileContractID{0}, HostPublicKey: types.SiaPublicKey{Key: []byte("foo")}},
		{ID: types.FileContractID{1}, HostPublicKey: types.SiaPublicKey{Key: []byte("bar")}},
		{ID: types.FileContractID{2}, HostPublicKey: types.SiaPublicKey{Key: []byte("baz")}},
	})
	c.renewedIDs = map[types.FileContractID]types.FileContractID{
		{0}: {1},
		{1}: {2},
		{2}: {3},
	}
	c.cachedRevisions = map[types.FileContractID]cachedRevision{
		{0}: {Revision: types.FileContractRevision{ParentID: types.FileContractID{0}}},
		{1}: {Revision: types.FileContractRevision{ParentID: types.FileContractID{1}}},
		{2}: {Revision: types.FileContractRevision{ParentID: types.FileContractID{2}}},
	}
	c.oldContracts = map[types.FileContractID]modules.RenterContract{
		{0}: {ID: types.FileContractID{0}, HostPublicKey: types.SiaPublicKey{Key: []byte("foo")}},
		{1}: {ID: types.FileContractID{1}, HostPublicKey: types.SiaPublicKey{Key: []byte("bar")}},
		{2}: {ID: types.FileContractID{2}, HostPublicKey: types.SiaPublicKey{Key: []byte("baz")}},
	}

	// save, clear, and reload
	err := c.save()
	if err != nil {
		t.Fatal(err)
	}
	c.hdb = stubHostDB{}
	c.contracts = proto.NewContractSet(nil)
	c.renewedIDs = make(map[types.FileContractID]types.FileContractID)
	c.cachedRevisions = make(map[types.FileContractID]cachedRevision)
	c.oldContracts = make(map[types.FileContractID]modules.RenterContract)
	err = c.load()
	if err != nil {
		t.Fatal(err)
	}
	// check that all fields were restored
	_, ok0 := c.contracts.View(types.FileContractID{0})
	_, ok1 := c.contracts.View(types.FileContractID{1})
	_, ok2 := c.contracts.View(types.FileContractID{2})
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("contracts were not restored properly:", c.contracts)
	}
	_, ok0 = c.renewedIDs[types.FileContractID{0}]
	_, ok1 = c.renewedIDs[types.FileContractID{1}]
	_, ok2 = c.renewedIDs[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("renewed IDs were not restored properly:", c.renewedIDs)
	}
	_, ok0 = c.cachedRevisions[types.FileContractID{0}]
	_, ok1 = c.cachedRevisions[types.FileContractID{1}]
	_, ok2 = c.cachedRevisions[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("cached revisions were not restored properly:", c.cachedRevisions)
	}
	_, ok0 = c.oldContracts[types.FileContractID{0}]
	_, ok1 = c.oldContracts[types.FileContractID{1}]
	_, ok2 = c.oldContracts[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("oldContracts were not restored properly:", c.oldContracts)
	}

	// use stdPersist instead of mock
	c.persist = newPersist(build.TempDir("contractor", t.Name()))
	os.MkdirAll(build.TempDir("contractor", t.Name()), 0700)

	// save, clear, and reload
	err = c.save()
	if err != nil {
		t.Fatal(err)
	}
	c.contracts = proto.NewContractSet(nil)
	c.renewedIDs = make(map[types.FileContractID]types.FileContractID)
	c.cachedRevisions = make(map[types.FileContractID]cachedRevision)
	c.oldContracts = make(map[types.FileContractID]modules.RenterContract)
	err = c.load()
	if err != nil {
		t.Fatal(err)
	}
	// check that all fields were restored
	_, ok0 = c.contracts.View(types.FileContractID{0})
	_, ok1 = c.contracts.View(types.FileContractID{1})
	_, ok2 = c.contracts.View(types.FileContractID{2})
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("contracts were not restored properly:", c.contracts)
	}
	_, ok0 = c.renewedIDs[types.FileContractID{0}]
	_, ok1 = c.renewedIDs[types.FileContractID{1}]
	_, ok2 = c.renewedIDs[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("renewed IDs were not restored properly:", c.renewedIDs)
	}
	_, ok0 = c.cachedRevisions[types.FileContractID{0}]
	_, ok1 = c.cachedRevisions[types.FileContractID{1}]
	_, ok2 = c.cachedRevisions[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("cached revisions were not restored properly:", c.cachedRevisions)
	}
	_, ok0 = c.oldContracts[types.FileContractID{0}]
	_, ok1 = c.oldContracts[types.FileContractID{1}]
	_, ok2 = c.oldContracts[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("oldContracts were not restored properly:", c.oldContracts)
	}
}

// blockCS is a consensusSet that calls ProcessConsensusChange on its blocks.
type blockCS struct {
	blocks []types.Block
}

func (cs blockCS) ConsensusSetSubscribe(s modules.ConsensusSetSubscriber, _ modules.ConsensusChangeID, _ <-chan struct{}) error {
	s.ProcessConsensusChange(modules.ConsensusChange{
		AppliedBlocks: cs.blocks,
	})
	return nil
}

func (blockCS) Synced() bool { return true }

func (blockCS) Unsubscribe(modules.ConsensusSetSubscriber) { return }

// TestPubKeyScanner tests that the pubkeyScanner type correctly identifies
// public keys in the blockchain.
func TestPubKeyScanner(t *testing.T) {
	// create pubkeys, announcements, and contracts
	var contracts []modules.RenterContract
	var blocks []types.Block
	var pubkeys []types.SiaPublicKey
	for i := 0; i < 3; i++ {
		// generate a keypair
		sk, pk := crypto.GenerateKeyPair()
		spk := types.SiaPublicKey{
			Algorithm: types.SignatureEd25519,
			Key:       pk[:],
		}
		pubkeys = append(pubkeys, spk)

		// create an announcement and add it to cs
		addr := modules.NetAddress("foo.bar:999" + strconv.Itoa(i))
		ann, err := modules.CreateAnnouncement(addr, spk, sk)
		if err != nil {
			t.Fatal(err)
		}
		blocks = append(blocks, types.Block{
			Transactions: []types.Transaction{{
				ArbitraryData: [][]byte{ann},
			}},
		})

		id := types.FileContractID{byte(i)}
		contracts = append(contracts, modules.RenterContract{ID: id, NetAddress: addr})
	}
	// overwrite the first pubkey with a new one, using the same netaddress.
	// The contractor should use the newer pubkey.
	sk, pk := crypto.GenerateKeyPair()
	spk := types.SiaPublicKey{
		Algorithm: types.SignatureEd25519,
		Key:       pk[:],
	}
	pubkeys[0] = spk
	ann, err := modules.CreateAnnouncement("foo.bar:9990", spk, sk)
	if err != nil {
		t.Fatal(err)
	}
	blocks = append(blocks, types.Block{
		Transactions: []types.Transaction{{
			ArbitraryData: [][]byte{ann},
		}},
	})

	// create contractor with mocked persist and cs dependencies
	c := &Contractor{
		persist:   new(memPersist),
		cs:        blockCS{blocks},
		contracts: proto.NewContractSet(contracts),
	}

	// save, clear, and reload
	err = c.save()
	if err != nil {
		t.Fatal(err)
	}
	c.contracts = proto.NewContractSet(nil)
	err = c.load()
	if err != nil {
		t.Fatal(err)
	}
	// check that contracts were loaded and have their pubkeys filled in
	for i, pk := range pubkeys {
		id := types.FileContractID{byte(i)}
		contract, ok := c.contracts.View(id)
		if !ok {
			t.Fatal("contracts were not restored properly:", c.contracts)
		}
		// check that pubkey was filled in
		if !bytes.Equal(contract.HostPublicKey.Key, pk.Key) {
			t.Errorf("contract has wrong pubkey: expected %q, got %q", pk.String(), contract.HostPublicKey.String())
		}
	}
}
<commit_msg>appease go vet by resolving t.Fatal mutex copy<commit_after>package contractor

import (
	"bytes"
	"os"
	"strconv"
	"testing"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/modules/renter/proto"
	"github.com/NebulousLabs/Sia/types"
)

// memPersist implements the persister interface in-memory.
type memPersist contractorPersist

func (m *memPersist) save(data contractorPersist) error { *m = memPersist(data); return nil }
func (m *memPersist) update(...journalUpdate) error     { return nil }
func (m memPersist) load(data *contractorPersist) error { *data = contractorPersist(m); return nil }
func (m memPersist) Close() error                       { return nil }

// TestSaveLoad tests that the contractor can save and load itself.
func TestSaveLoad(t *testing.T) {
	// create contractor with mocked persist dependency
	c := &Contractor{
		persist: new(memPersist),
	}

	// add some fake contracts
	c.contracts = proto.NewContractSet([]modules.RenterContract{
		{ID: types.FileContractID{0}, HostPublicKey: types.SiaPublicKey{Key: []byte("foo")}},
		{ID: types.FileContractID{1}, HostPublicKey: types.SiaPublicKey{Key: []byte("bar")}},
		{ID: types.FileContractID{2}, HostPublicKey: types.SiaPublicKey{Key: []byte("baz")}},
	})
	c.renewedIDs = map[types.FileContractID]types.FileContractID{
		{0}: {1},
		{1}: {2},
		{2}: {3},
	}
	c.cachedRevisions = map[types.FileContractID]cachedRevision{
		{0}: {Revision: types.FileContractRevision{ParentID: types.FileContractID{0}}},
		{1}: {Revision: types.FileContractRevision{ParentID: types.FileContractID{1}}},
		{2}: {Revision: types.FileContractRevision{ParentID: types.FileContractID{2}}},
	}
	c.oldContracts = map[types.FileContractID]modules.RenterContract{
		{0}: {ID: types.FileContractID{0}, HostPublicKey: types.SiaPublicKey{Key: []byte("foo")}},
		{1}: {ID: types.FileContractID{1}, HostPublicKey: types.SiaPublicKey{Key: []byte("bar")}},
		{2}: {ID: types.FileContractID{2}, HostPublicKey: types.SiaPublicKey{Key: []byte("baz")}},
	}

	// save, clear, and reload
	err := c.save()
	if err != nil {
		t.Fatal(err)
	}
	c.hdb = stubHostDB{}
	c.contracts = proto.NewContractSet(nil)
	c.renewedIDs = make(map[types.FileContractID]types.FileContractID)
	c.cachedRevisions = make(map[types.FileContractID]cachedRevision)
	c.oldContracts = make(map[types.FileContractID]modules.RenterContract)
	err = c.load()
	if err != nil {
		t.Fatal(err)
	}
	// check that all fields were restored
	_, ok0 := c.contracts.View(types.FileContractID{0})
	_, ok1 := c.contracts.View(types.FileContractID{1})
	_, ok2 := c.contracts.View(types.FileContractID{2})
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("contracts were not restored properly:", c.contracts.Len())
	}
	_, ok0 = c.renewedIDs[types.FileContractID{0}]
	_, ok1 = c.renewedIDs[types.FileContractID{1}]
	_, ok2 = c.renewedIDs[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("renewed IDs were not restored properly:", c.renewedIDs)
	}
	_, ok0 = c.cachedRevisions[types.FileContractID{0}]
	_, ok1 = c.cachedRevisions[types.FileContractID{1}]
	_, ok2 = c.cachedRevisions[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("cached revisions were not restored properly:", c.cachedRevisions)
	}
	_, ok0 = c.oldContracts[types.FileContractID{0}]
	_, ok1 = c.oldContracts[types.FileContractID{1}]
	_, ok2 = c.oldContracts[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("oldContracts were not restored properly:", c.oldContracts)
	}

	// use stdPersist instead of mock
	c.persist = newPersist(build.TempDir("contractor", t.Name()))
	os.MkdirAll(build.TempDir("contractor", t.Name()), 0700)

	// save, clear, and reload
	err = c.save()
	if err != nil {
		t.Fatal(err)
	}
	c.contracts = proto.NewContractSet(nil)
	c.renewedIDs = make(map[types.FileContractID]types.FileContractID)
	c.cachedRevisions = make(map[types.FileContractID]cachedRevision)
	c.oldContracts = make(map[types.FileContractID]modules.RenterContract)
	err = c.load()
	if err != nil {
		t.Fatal(err)
	}
	// check that all fields were restored
	_, ok0 = c.contracts.View(types.FileContractID{0})
	_, ok1 = c.contracts.View(types.FileContractID{1})
	_, ok2 = c.contracts.View(types.FileContractID{2})
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("contracts were not restored properly:", c.contracts.Len())
	}
	_, ok0 = c.renewedIDs[types.FileContractID{0}]
	_, ok1 = c.renewedIDs[types.FileContractID{1}]
	_, ok2 = c.renewedIDs[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("renewed IDs were not restored properly:", c.renewedIDs)
	}
	_, ok0 = c.cachedRevisions[types.FileContractID{0}]
	_, ok1 = c.cachedRevisions[types.FileContractID{1}]
	_, ok2 = c.cachedRevisions[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("cached revisions were not restored properly:", c.cachedRevisions)
	}
	_, ok0 = c.oldContracts[types.FileContractID{0}]
	_, ok1 = c.oldContracts[types.FileContractID{1}]
	_, ok2 = c.oldContracts[types.FileContractID{2}]
	if !ok0 || !ok1 || !ok2 {
		t.Fatal("oldContracts were not restored properly:", c.oldContracts)
	}
}

// blockCS is a consensusSet that calls ProcessConsensusChange on its blocks.
type blockCS struct {
	blocks []types.Block
}

func (cs blockCS) ConsensusSetSubscribe(s modules.ConsensusSetSubscriber, _ modules.ConsensusChangeID, _ <-chan struct{}) error {
	s.ProcessConsensusChange(modules.ConsensusChange{
		AppliedBlocks: cs.blocks,
	})
	return nil
}

func (blockCS) Synced() bool { return true }

func (blockCS) Unsubscribe(modules.ConsensusSetSubscriber) { return }

// TestPubKeyScanner tests that the pubkeyScanner type correctly identifies
// public keys in the blockchain.
func TestPubKeyScanner(t *testing.T) {
	// create pubkeys, announcements, and contracts
	var contracts []modules.RenterContract
	var blocks []types.Block
	var pubkeys []types.SiaPublicKey
	for i := 0; i < 3; i++ {
		// generate a keypair
		sk, pk := crypto.GenerateKeyPair()
		spk := types.SiaPublicKey{
			Algorithm: types.SignatureEd25519,
			Key:       pk[:],
		}
		pubkeys = append(pubkeys, spk)

		// create an announcement and add it to cs
		addr := modules.NetAddress("foo.bar:999" + strconv.Itoa(i))
		ann, err := modules.CreateAnnouncement(addr, spk, sk)
		if err != nil {
			t.Fatal(err)
		}
		blocks = append(blocks, types.Block{
			Transactions: []types.Transaction{{
				ArbitraryData: [][]byte{ann},
			}},
		})

		id := types.FileContractID{byte(i)}
		contracts = append(contracts, modules.RenterContract{ID: id, NetAddress: addr})
	}
	// overwrite the first pubkey with a new one, using the same netaddress.
	// The contractor should use the newer pubkey.
	sk, pk := crypto.GenerateKeyPair()
	spk := types.SiaPublicKey{
		Algorithm: types.SignatureEd25519,
		Key:       pk[:],
	}
	pubkeys[0] = spk
	ann, err := modules.CreateAnnouncement("foo.bar:9990", spk, sk)
	if err != nil {
		t.Fatal(err)
	}
	blocks = append(blocks, types.Block{
		Transactions: []types.Transaction{{
			ArbitraryData: [][]byte{ann},
		}},
	})

	// create contractor with mocked persist and cs dependencies
	c := &Contractor{
		persist:   new(memPersist),
		cs:        blockCS{blocks},
		contracts: proto.NewContractSet(contracts),
	}

	// save, clear, and reload
	err = c.save()
	if err != nil {
		t.Fatal(err)
	}
	c.contracts = proto.NewContractSet(nil)
	err = c.load()
	if err != nil {
		t.Fatal(err)
	}
	// check that contracts were loaded and have their pubkeys filled in
	for i, pk := range pubkeys {
		id := types.FileContractID{byte(i)}
		contract, ok := c.contracts.View(id)
		if !ok {
			t.Fatal("contracts were not restored properly:", c.contracts.Len())
		}
		// check that pubkey was filled in
		if !bytes.Equal(contract.HostPublicKey.Key, pk.Key) {
			t.Errorf("contract has wrong pubkey: expected %q, got %q", pk.String(), contract.HostPublicKey.String())
		}
	}
}
<|endoftext|>


<commit_before>// Package s3 signs HTTP requests for Amazon S3 and compatible services.
package s3

// See
// http://docs.amazonwebservices.com/AmazonS3/2006-03-01/dev/RESTAuthentication.html.

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"io"
	"net/http"
	"sort"
	"strings"
)

var signParams = map[string]bool{
	"acl":                          true,
	"delete":                       true,
	"lifecycle":                    true,
	"location":                     true,
	"logging":                      true,
	"notification":                 true,
	"partNumber":                   true,
	"policy":                       true,
	"requestPayment":               true,
	"response-cache-control":       true,
	"response-content-disposition": true,
	"response-content-encoding":    true,
	"response-content-language":    true,
	"response-content-type":        true,
	"response-expires":             true,
	"torrent":                      true,
	"uploadId":                     true,
	"uploads":                      true,
	"versionId":                    true,
	"versioning":                   true,
	"versions":                     true,
	"website":                      true,
}

// Keys holds a set of Amazon Security Credentials.
type Keys struct {
	AccessKey string
	SecretKey string

	// SecurityToken is used for temporary security credentials.
	// If set, it will be added to header field X-Amz-Security-Token
	// before signing a request.
	SecurityToken string
	// See http://docs.aws.amazon.com/AmazonS3/latest/dev/MakingRequests.html#TypesofSecurityCredentials
}

// IdentityBucket returns subdomain.
// It is designed to be used with S3-compatible services that
// treat the entire subdomain as the bucket name, for example
// storage.io.
func IdentityBucket(subdomain string) string {
	return subdomain
}

// AmazonBucket returns everything up to the last '.' in subdomain.
// It is designed to be used with the Amazon service.
//   "johnsmith.s3"           becomes "johnsmith"
//   "johnsmith.s3-eu-west-1" becomes "johnsmith"
func AmazonBucket(subdomain string) string {
	s := strings.Split(subdomain, ".")
	return strings.Join(s[:len(s)-1], ".")
}

// DefaultService is the default Service used by Sign.
var DefaultService = &Service{Domain: "amazonaws.com"}

// Sign signs an HTTP request with the given S3 keys.
//
// This function is a wrapper around DefaultService.Sign.
func Sign(r *http.Request, k Keys) {
	DefaultService.Sign(r, k)
}

// Service represents an S3-compatible service.
type Service struct {
	// Domain is the service's root domain. It is used to extract
	// the subdomain from an http.Request before passing the
	// subdomain to Bucket.
	Domain string

	// Bucket derives the bucket name from a subdomain.
	// If nil, AmazonBucket is used.
	Bucket func(subdomain string) string
}

// Sign signs an HTTP request with the given S3 keys for use on service s.
func (s *Service) Sign(r *http.Request, k Keys) {
	if k.SecurityToken != "" {
		r.Header.Set("X-Amz-Security-Token", k.SecurityToken)
	}
	h := hmac.New(sha1.New, []byte(k.SecretKey))
	s.writeSigData(h, r)
	sig := make([]byte, base64.StdEncoding.EncodedLen(h.Size()))
	base64.StdEncoding.Encode(sig, h.Sum(nil))
	r.Header.Set("Authorization", "AWS "+k.AccessKey+":"+string(sig))
}

func (s *Service) writeSigData(w io.Writer, r *http.Request) {
	w.Write([]byte(r.Method))
	w.Write([]byte{'\n'})
	w.Write([]byte(r.Header.Get("content-md5")))
	w.Write([]byte{'\n'})
	w.Write([]byte(r.Header.Get("content-type")))
	w.Write([]byte{'\n'})
	if _, ok := r.Header["X-Amz-Date"]; !ok {
		w.Write([]byte(r.Header.Get("date")))
	}
	w.Write([]byte{'\n'})
	writeAmzHeaders(w, r)
	s.writeResource(w, r)
}

func (s *Service) writeResource(w io.Writer, r *http.Request) {
	s.writeVhostBucket(w, strings.ToLower(r.Host))
	path := r.URL.RequestURI()
	if r.URL.RawQuery != "" {
		path = path[:len(path)-len(r.URL.RawQuery)-1]
	}
	w.Write([]byte(path))
	s.writeSubResource(w, r)
}

func (s *Service) writeVhostBucket(w io.Writer, host string) {
	if i := strings.Index(host, ":"); i != -1 {
		host = host[:i]
	}

	if host == s.Domain {
		// no vhost - do nothing
	} else if strings.HasSuffix(host, "."+s.Domain) {
		// vhost - bucket may be in prefix
		b := s.Bucket
		if b == nil {
			b = AmazonBucket
		}
		bucket := b(host[:len(host)-len(s.Domain)-1])

		if bucket != "" {
			w.Write([]byte{'/'})
			w.Write([]byte(bucket))
		}
	} else {
		// cname - bucket is host
		w.Write([]byte{'/'})
		w.Write([]byte(host))
	}
}

func (s *Service) writeSubResource(w io.Writer, r *http.Request) {
	var a []string
	for k, vs := range r.URL.Query() {
		if signParams[k] {
			for _, v := range vs {
				if v == "" {
					a = append(a, k)
				} else {
					a = append(a, k+"="+v)
				}
			}
		}
	}
	sort.Strings(a)
	var p byte = '?'
	for _, s := range a {
		w.Write([]byte{p})
		w.Write([]byte(s))
		p = '&'
	}
}

func writeAmzHeaders(w io.Writer, r *http.Request) {
	var keys []string
	for k, _ := range r.Header {
		if strings.HasPrefix(strings.ToLower(k), "x-amz-") {
			keys = append(keys, k)
		}
	}

	sort.Strings(keys)
	var a []string
	for _, k := range keys {
		v := r.Header[k]
		a = append(a, strings.ToLower(k)+":"+strings.Join(v, ","))
	}
	for _, h := range a {
		w.Write([]byte(h))
		w.Write([]byte{'\n'})
	}
}
<commit_msg>one more AmazonBucket example<commit_after>// Package s3 signs HTTP requests for Amazon S3 and compatible services.
package s3

// See
// http://docs.amazonwebservices.com/AmazonS3/2006-03-01/dev/RESTAuthentication.html.

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"io"
	"net/http"
	"sort"
	"strings"
)

var signParams = map[string]bool{
	"acl":                          true,
	"delete":                       true,
	"lifecycle":                    true,
	"location":                     true,
	"logging":                      true,
	"notification":                 true,
	"partNumber":                   true,
	"policy":                       true,
	"requestPayment":               true,
	"response-cache-control":       true,
	"response-content-disposition": true,
	"response-content-encoding":    true,
	"response-content-language":    true,
	"response-content-type":        true,
	"response-expires":             true,
	"torrent":                      true,
	"uploadId":                     true,
	"uploads":                      true,
	"versionId":                    true,
	"versioning":                   true,
	"versions":                     true,
	"website":                      true,
}

// Keys holds a set of Amazon Security Credentials.
type Keys struct {
	AccessKey string
	SecretKey string

	// SecurityToken is used for temporary security credentials.
	// If set, it will be added to header field X-Amz-Security-Token
	// before signing a request.
	SecurityToken string
	// See http://docs.aws.amazon.com/AmazonS3/latest/dev/MakingRequests.html#TypesofSecurityCredentials
}

// IdentityBucket returns subdomain.
// It is designed to be used with S3-compatible services that
// treat the entire subdomain as the bucket name, for example
// storage.io.
func IdentityBucket(subdomain string) string {
	return subdomain
}

// AmazonBucket returns everything up to the last '.' in subdomain.
// It is designed to be used with the Amazon service.
//   "johnsmith.s3"           becomes "johnsmith"
//   "johnsmith.s3-eu-west-1" becomes "johnsmith"
//   "www.example.com.s3"     becomes "www.example.com"
func AmazonBucket(subdomain string) string {
	s := strings.Split(subdomain, ".")
	return strings.Join(s[:len(s)-1], ".")
}

// DefaultService is the default Service used by Sign.
var DefaultService = &Service{Domain: "amazonaws.com"}

// Sign signs an HTTP request with the given S3 keys.
//
// This function is a wrapper around DefaultService.Sign.
func Sign(r *http.Request, k Keys) {
	DefaultService.Sign(r, k)
}

// Service represents an S3-compatible service.
type Service struct {
	// Domain is the service's root domain. It is used to extract
	// the subdomain from an http.Request before passing the
	// subdomain to Bucket.
	Domain string

	// Bucket derives the bucket name from a subdomain.
	// If nil, AmazonBucket is used.
	Bucket func(subdomain string) string
}

// Sign signs an HTTP request with the given S3 keys for use on service s.
func (s *Service) Sign(r *http.Request, k Keys) {
	if k.SecurityToken != "" {
		r.Header.Set("X-Amz-Security-Token", k.SecurityToken)
	}
	h := hmac.New(sha1.New, []byte(k.SecretKey))
	s.writeSigData(h, r)
	sig := make([]byte, base64.StdEncoding.EncodedLen(h.Size()))
	base64.StdEncoding.Encode(sig, h.Sum(nil))
	r.Header.Set("Authorization", "AWS "+k.AccessKey+":"+string(sig))
}

func (s *Service) writeSigData(w io.Writer, r *http.Request) {
	w.Write([]byte(r.Method))
	w.Write([]byte{'\n'})
	w.Write([]byte(r.Header.Get("content-md5")))
	w.Write([]byte{'\n'})
	w.Write([]byte(r.Header.Get("content-type")))
	w.Write([]byte{'\n'})
	if _, ok := r.Header["X-Amz-Date"]; !ok {
		w.Write([]byte(r.Header.Get("date")))
	}
	w.Write([]byte{'\n'})
	writeAmzHeaders(w, r)
	s.writeResource(w, r)
}

func (s *Service) writeResource(w io.Writer, r *http.Request) {
	s.writeVhostBucket(w, strings.ToLower(r.Host))
	path := r.URL.RequestURI()
	if r.URL.RawQuery != "" {
		path = path[:len(path)-len(r.URL.RawQuery)-1]
	}
	w.Write([]byte(path))
	s.writeSubResource(w, r)
}

func (s *Service) writeVhostBucket(w io.Writer, host string) {
	if i := strings.Index(host, ":"); i != -1 {
		host = host[:i]
	}

	if host == s.Domain {
		// no vhost - do nothing
	} else if strings.HasSuffix(host, "."+s.Domain) {
		// vhost - bucket may be in prefix
		b := s.Bucket
		if b == nil {
			b = AmazonBucket
		}
		bucket := b(host[:len(host)-len(s.Domain)-1])

		if bucket != "" {
			w.Write([]byte{'/'})
			w.Write([]byte(bucket))
		}
	} else {
		// cname - bucket is host
		w.Write([]byte{'/'})
		w.Write([]byte(host))
	}
}

func (s *Service) writeSubResource(w io.Writer, r *http.Request) {
	var a []string
	for k, vs := range r.URL.Query() {
		if signParams[k] {
			for _, v := range vs {
				if v == "" {
					a = append(a, k)
				} else {
					a = append(a, k+"="+v)
				}
			}
		}
	}
	sort.Strings(a)
	var p byte = '?'
	for _, s := range a {
		w.Write([]byte{p})
		w.Write([]byte(s))
		p = '&'
	}
}

func writeAmzHeaders(w io.Writer, r *http.Request) {
	var keys []string
	for k, _ := range r.Header {
		if strings.HasPrefix(strings.ToLower(k), "x-amz-") {
			keys = append(keys, k)
		}
	}

	sort.Strings(keys)
	var a []string
	for _, k := range keys {
		v := r.Header[k]
		a = append(a, strings.ToLower(k)+":"+strings.Join(v, ","))
	}
	for _, h := range a {
		w.Write([]byte(h))
		w.Write([]byte{'\n'})
	}
}
<|endoftext|>


<commit_before>package gorpc

import (
	"bytes"
	"sync"
	"time"

	"github.com/gsdocker/gserrors"
	"github.com/gsdocker/gslogger"
)

// Sink rpc sink
type Sink interface {
	Handler
	MessageChannel
	Unregister(dispatcher Dispatcher)
	Register(dispatcher Dispatcher)
}

type _Promise struct {
	promise chan *Response
	err     error
	timer   *time.Timer
}

// Promise .
type Promise interface {
	Wait() (callReturn *Response, err error)
	Notify(callReturn *Response, err error)
	Timeout()
	Cancel()
}

// NewPromise .
func NewPromise(timeout time.Duration, f func()) Promise {
	promise := &_Promise{
		promise: make(chan *Response, 1),
	}

	promise.timer = time.AfterFunc(timeout, func() {
		f()
		promise.Timeout()

	})

	return promise
}

func (promise *_Promise) Wait() (callReturn *Response, err error) {
	return <-promise.promise, promise.err
}

func (promise *_Promise) Notify(callReturn *Response, err error) {

	promise.timer.Stop()

	promise.err = err

	promise.promise <- callReturn
}

func (promise *_Promise) Timeout() {
	promise.err = ErrTimeout
	promise.promise <- nil
}

func (promise *_Promise) Cancel() {
	promise.err = ErrCanceled
	close(promise.promise)
}

// _Sink .
type _Sink struct {
	gslogger.Log                       // Mixin Log APIs
	sync.RWMutex                       // mutex
	name         string                // name
	timeout      time.Duration         // timeout
	seqID        uint16                // sequence id
	dispatchers  map[uint16]Dispatcher // register dispatchers
	promises     map[uint16]Promise    // rpc promise
	cached       chan *Message         // cached message
	cachedTask   chan func()           // cached task
	processors   int                   // task processors
}

// NewSink .
func NewSink(name string, timeout time.Duration, cached int, processors int) Sink {
	return &_Sink{
		Log:         gslogger.Get("sink"),
		name:        name,
		timeout:     timeout,
		dispatchers: make(map[uint16]Dispatcher),
		promises:    make(map[uint16]Promise),
		cached:      make(chan *Message, cached),
		cachedTask:  make(chan func(), cached),
		processors:  processors,
	}
}

func (sink *_Sink) String() string {
	return sink.name
}

// Promise .
func (sink *_Sink) Promise() (Promise, uint16) {

	sink.Lock()
	defer sink.Unlock()

	for {
		seqID := sink.seqID

		sink.seqID++

		_, ok := sink.promises[seqID]

		if ok {
			continue
		}

		promise := NewPromise(sink.timeout, func() {
			sink.Lock()
			defer sink.Unlock()

			delete(sink.promises, seqID)
		})

		sink.promises[seqID] = promise

		return promise, seqID
	}
}

// Send .
func (sink *_Sink) Send(call *Request) (Future, error) {

	promise, id := sink.Promise()

	call.ID = id

	var buff bytes.Buffer

	err := WriteRequest(&buff, call)

	if err != nil {
		promise.Cancel()
		return nil, err
	}

	message := NewMessage()

	message.Code = CodeRequest

	message.Content = buff.Bytes()

	err = sink.SendMessage(message)

	if err != nil {
		promise.Cancel()
		return nil, err
	}

	return promise, nil
}

func (sink *_Sink) SendMessage(message *Message) (err error) {

	defer func() {
		if e := recover(); e != nil {
			err = ErrClosed
		}
	}()

	select {
	case sink.cached <- message:
	default:
		err = gserrors.Newf(ErrOverflow, "sink(%s) cached overflow", sink)
	}

	return
}

// Register .
func (sink *_Sink) Register(dispatcher Dispatcher) {
	sink.Lock()
	defer sink.Unlock()

	sink.dispatchers[dispatcher.ID()] = dispatcher
}

// Unreigster .
func (sink *_Sink) Unregister(dispatcher Dispatcher) {
	sink.Lock()
	defer sink.Unlock()

	if dispatcher, ok := sink.dispatchers[dispatcher.ID()]; ok && dispatcher == dispatcher {
		delete(sink.dispatchers, dispatcher.ID())
	}
}

func (sink *_Sink) dispatchResponse(response *Response) {
	sink.Lock()
	defer sink.Unlock()

	if promise, ok := sink.promises[response.ID]; ok {
		go promise.Notify(response, nil)
		delete(sink.promises, response.ID)
		return
	}

	sink.W("unhandle response(%d)(%d)", response.ID, response.Service)
}

func (sink *_Sink) dispatch(id uint16) (dispatcher Dispatcher, ok bool) {
	sink.RLock()
	defer sink.RUnlock()

	dispatcher, ok = sink.dispatchers[id]

	return
}

func (sink *_Sink) dispatchRequest(message *Message) error {

	request, err := ReadRequest(bytes.NewBuffer(message.Content))

	if err != nil {
		sink.E("[%s] unmarshal request error\n%s", sink.name, err)
		return err
	}

	if dispatcher, ok := sink.dispatch(request.Service); ok {

		response, err := dispatcher.Dispatch(request)

		if err != nil {
			sink.E("dispatch request(%d)(%d:%d) error\n%s", request.ID, request.Service, request.Method, err)
			return nil
		}

		var buff bytes.Buffer

		err = WriteResponse(&buff, response)

		if err != nil {
			sink.E("marshal request(%d)(%d:%d)'s response error\n%s", request.ID, request.Service, request.Method, err)
			return nil
		}

		message.Code = CodeResponse

		message.Content = buff.Bytes()

		go sink.SendMessage(message)

		return nil
	}

	sink.W("[%s] unhandle request(%d)(%d:%d)", sink.name, request.ID, request.Service, request.Method)

	return nil
}

func (sink *_Sink) OpenHandler(context Context) error {

	go func() {
		for message := range sink.cached {
			context.WriteReadPipline(message)
		}
	}()

	for i := 0; i < sink.processors; i++ {
		go func() {
			for f := range sink.cachedTask {
				sink.protectedCall(f)
			}
		}()
	}

	return nil
}

func (sink *_Sink) protectedCall(f func()) {

	defer func() {
		if e := recover(); e != nil {

			if _, ok := e.(error); ok {
				sink.E("%s", gserrors.Newf(e.(error), "catched exception"))
			} else {
				sink.E("%s", gserrors.Newf(ErrUnknown, "catched exception :%s", e))
			}
		}
	}()

	f()
}

func (sink *_Sink) CloseHandler(context Context) {
	close(sink.cached)
	close(sink.cachedTask)
}
func (sink *_Sink) HandleError(context Context, err error) error {
	return err
}

func (sink *_Sink) HandleRead(context Context, message *Message) (*Message, error) {
	return message, nil
}

func (sink *_Sink) pushTask(f func()) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = ErrClosed
		}
	}()

	select {
	case sink.cachedTask <- f:
	default:
		err = gserrors.Newf(ErrOverflow, "sink task(%s) task cached overflow", sink)
	}

	return
}

func (sink *_Sink) HandleWrite(context Context, message *Message) (*Message, error) {

	switch message.Code {
	case CodeRequest:

		err := sink.pushTask(func() { sink.dispatchRequest(message) })

		if err != nil {
			return nil, err
		}

		return nil, nil

	case CodeResponse:

		response, err := ReadResponse(bytes.NewBuffer(message.Content))

		if err != nil {
			sink.E("[%s] unmarshal response error\n%s", sink.name, err)
			return nil, err
		}

		err = sink.pushTask(func() { sink.dispatchResponse(response) })

		if err != nil {
			return nil, err
		}

		return nil, nil

	default:

		sink.W("[%s] unsupport message(%s)", sink.name, message.Code)

		return message, nil
	}

}
<commit_msg>catch recluse chan exception<commit_after>package gorpc

import (
	"bytes"
	"sync"
	"time"

	"github.com/gsdocker/gserrors"
	"github.com/gsdocker/gslogger"
)

// Sink rpc sink
type Sink interface {
	Handler
	MessageChannel
	Unregister(dispatcher Dispatcher)
	Register(dispatcher Dispatcher)
}

type _Promise struct {
	promise chan *Response
	err     error
	timer   *time.Timer
}

// Promise .
type Promise interface {
	Wait() (callReturn *Response, err error)
	Notify(callReturn *Response, err error)
	Timeout()
	Cancel()
}

// NewPromise .
func NewPromise(timeout time.Duration, f func()) Promise {
	promise := &_Promise{
		promise: make(chan *Response, 1),
	}

	promise.timer = time.AfterFunc(timeout, func() {
		f()
		promise.Timeout()

	})

	return promise
}

func (promise *_Promise) Wait() (callReturn *Response, err error) {
	return <-promise.promise, promise.err
}

func (promise *_Promise) Notify(callReturn *Response, err error) {

	promise.timer.Stop()

	promise.err = err

	promise.promise <- callReturn
}

func (promise *_Promise) Timeout() {
	promise.err = ErrTimeout
	promise.promise <- nil
}

func (promise *_Promise) Cancel() {
	promise.err = ErrCanceled
	close(promise.promise)
}

// _Sink .
type _Sink struct {
	gslogger.Log                       // Mixin Log APIs
	sync.RWMutex                       // mutex
	name         string                // name
	timeout      time.Duration         // timeout
	seqID        uint16                // sequence id
	dispatchers  map[uint16]Dispatcher // register dispatchers
	promises     map[uint16]Promise    // rpc promise
	cached       chan *Message         // cached message
	cachedTask   chan func()           // cached task
	processors   int                   // task processors
}

// NewSink .
func NewSink(name string, timeout time.Duration, cached int, processors int) Sink {
	return &_Sink{
		Log:         gslogger.Get("sink"),
		name:        name,
		timeout:     timeout,
		dispatchers: make(map[uint16]Dispatcher),
		promises:    make(map[uint16]Promise),
		cached:      make(chan *Message, cached),
		cachedTask:  make(chan func(), cached),
		processors:  processors,
	}
}

func (sink *_Sink) String() string {
	return sink.name
}

// Promise .
func (sink *_Sink) Promise() (Promise, uint16) {

	sink.Lock()
	defer sink.Unlock()

	for {
		seqID := sink.seqID

		sink.seqID++

		_, ok := sink.promises[seqID]

		if ok {
			continue
		}

		promise := NewPromise(sink.timeout, func() {
			sink.Lock()
			defer sink.Unlock()

			delete(sink.promises, seqID)
		})

		sink.promises[seqID] = promise

		return promise, seqID
	}
}

// Send .
func (sink *_Sink) Send(call *Request) (Future, error) {

	promise, id := sink.Promise()

	call.ID = id

	var buff bytes.Buffer

	err := WriteRequest(&buff, call)

	if err != nil {
		promise.Cancel()
		return nil, err
	}

	message := NewMessage()

	message.Code = CodeRequest

	message.Content = buff.Bytes()

	err = sink.SendMessage(message)

	if err != nil {
		promise.Cancel()
		return nil, err
	}

	return promise, nil
}

func (sink *_Sink) SendMessage(message *Message) (err error) {

	defer func() {
		if e := recover(); e != nil {
			err = ErrClosed
		}
	}()

	select {
	case sink.cached <- message:
	default:
		err = gserrors.Newf(ErrOverflow, "sink(%s) cached overflow", sink)
	}

	return
}

// Register .
func (sink *_Sink) Register(dispatcher Dispatcher) {
	sink.Lock()
	defer sink.Unlock()

	sink.dispatchers[dispatcher.ID()] = dispatcher
}

// Unreigster .
func (sink *_Sink) Unregister(dispatcher Dispatcher) {
	sink.Lock()
	defer sink.Unlock()

	if dispatcher, ok := sink.dispatchers[dispatcher.ID()]; ok && dispatcher == dispatcher {
		delete(sink.dispatchers, dispatcher.ID())
	}
}

func (sink *_Sink) dispatchResponse(response *Response) {
	sink.Lock()
	defer sink.Unlock()

	if promise, ok := sink.promises[response.ID]; ok {
		go promise.Notify(response, nil)
		delete(sink.promises, response.ID)
		return
	}

	sink.W("unhandle response(%d)(%d)", response.ID, response.Service)
}

func (sink *_Sink) dispatch(id uint16) (dispatcher Dispatcher, ok bool) {
	sink.RLock()
	defer sink.RUnlock()

	dispatcher, ok = sink.dispatchers[id]

	return
}

func (sink *_Sink) dispatchRequest(message *Message) error {

	request, err := ReadRequest(bytes.NewBuffer(message.Content))

	if err != nil {
		sink.E("[%s] unmarshal request error\n%s", sink.name, err)
		return err
	}

	if dispatcher, ok := sink.dispatch(request.Service); ok {

		response, err := dispatcher.Dispatch(request)

		if err != nil {
			sink.E("dispatch request(%d)(%d:%d) error\n%s", request.ID, request.Service, request.Method, err)
			return nil
		}

		var buff bytes.Buffer

		err = WriteResponse(&buff, response)

		if err != nil {
			sink.E("marshal request(%d)(%d:%d)'s response error\n%s", request.ID, request.Service, request.Method, err)
			return nil
		}

		message.Code = CodeResponse

		message.Content = buff.Bytes()

		go sink.SendMessage(message)

		return nil
	}

	sink.W("[%s] unhandle request(%d)(%d:%d)", sink.name, request.ID, request.Service, request.Method)

	return nil
}

func (sink *_Sink) OpenHandler(context Context) error {

	go func() {
		for message := range sink.cached {
			context.WriteReadPipline(message)
		}
	}()

	for i := 0; i < sink.processors; i++ {
		go func() {
			for f := range sink.cachedTask {
				sink.protectedCall(f)
			}
		}()
	}

	return nil
}

func (sink *_Sink) protectedCall(f func()) {

	defer func() {
		if e := recover(); e != nil {

			if _, ok := e.(error); ok {
				sink.E("%s", gserrors.Newf(e.(error), "catched exception"))
			} else {
				sink.E("%s", gserrors.Newf(ErrUnknown, "catched exception :%s", e))
			}
		}
	}()

	f()
}

func (sink *_Sink) CloseHandler(context Context) {

	defer recover()

	close(sink.cached)
	close(sink.cachedTask)
}
func (sink *_Sink) HandleError(context Context, err error) error {
	return err
}

func (sink *_Sink) HandleRead(context Context, message *Message) (*Message, error) {
	return message, nil
}

func (sink *_Sink) pushTask(f func()) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = ErrClosed
		}
	}()

	select {
	case sink.cachedTask <- f:
	default:
		err = gserrors.Newf(ErrOverflow, "sink task(%s) task cached overflow", sink)
	}

	return
}

func (sink *_Sink) HandleWrite(context Context, message *Message) (*Message, error) {

	switch message.Code {
	case CodeRequest:

		err := sink.pushTask(func() { sink.dispatchRequest(message) })

		if err != nil {
			return nil, err
		}

		return nil, nil

	case CodeResponse:

		response, err := ReadResponse(bytes.NewBuffer(message.Content))

		if err != nil {
			sink.E("[%s] unmarshal response error\n%s", sink.name, err)
			return nil, err
		}

		err = sink.pushTask(func() { sink.dispatchResponse(response) })

		if err != nil {
			return nil, err
		}

		return nil, nil

	default:

		sink.W("[%s] unsupport message(%s)", sink.name, message.Code)

		return message, nil
	}

}
<|endoftext|>