<commit_before>package tunnel

import (
	"os"
	"sync"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/config/options"
	"github.com/micro/go-micro/proxy/mucp"
	"github.com/micro/go-micro/router"
	"github.com/micro/go-micro/server"
	trn "github.com/micro/go-micro/transport"
	tun "github.com/micro/go-micro/tunnel"
	"github.com/micro/go-micro/tunnel/transport"
	"github.com/micro/go-micro/util/log"
)

var (
	// Name of the router microservice
	Name = "go.micro.tunnel"
	// Address is the tunnel microservice bind address
	Address = ":8084"
	// Tunnel is the tunnel bind address
	Tunnel = ":9096"
	// Channel is the name of the tunnel session
	Channel = "tunnel"
	// Router is the router gossip bind address
	Router = ":9093"
	// Network is the network id
	Network = "local"
)

func accept(l tun.Listener, exit chan struct{}) error {
	for {
		// accept a connection
		c, err := l.Accept()
		if err != nil {
			return err
		}

		select {
		case <-exit:
			return nil
		default:
			m := new(trn.Message)
			if err := c.Recv(m); err != nil {
				return err
			}
			// TODO: do something with the message
		}
	}
}

// run runs the micro server
func run(ctx *cli.Context, srvOpts ...micro.Option) {
	// Init plugins
	for _, p := range Plugins() {
		p.Init(ctx)
	}

	if len(ctx.GlobalString("server_name")) > 0 {
		Name = ctx.GlobalString("server_name")
	}
	if len(ctx.String("address")) > 0 {
		Address = ctx.String("address")
	}
	if len(ctx.String("network_address")) > 0 {
		Network = ctx.String("network")
	}
	if len(ctx.String("tunnel_address")) > 0 {
		Tunnel = ctx.String("tunnel")
	}
	// default gateway address
	var gateway string
	if len(ctx.String("gateway_address")) > 0 {
		gateway = ctx.String("gateway")
	}

	// create a tunnel
	t := tun.NewTunnel(
		tun.Address(Tunnel),
	)

	// connect the tunnel
	if err := t.Connect(); err != nil {
		log.Logf("[tunnel] failed to connect: %s", err)
		os.Exit(1)
	}

	// listen on tunnel
	l, err := t.Listen(Channel)
	if err != nil {
		log.Logf("[tunnel] failed to listen: %s", err)
		os.Exit(1)
	}

	// TODO: go accept tunnel connections

	// Initialise service
	service := micro.NewService(
		micro.Name(Name),
		micro.Address(Address),
		micro.RegisterTTL(time.Duration(ctx.GlobalInt("register_ttl"))*time.Second),
		micro.RegisterInterval(time.Duration(ctx.GlobalInt("register_interval"))*time.Second),
	)

	// local tunnel router
	r := router.NewRouter(
		router.Id(service.Server().Options().Id),
		router.Registry(service.Client().Options().Registry),
		router.Address(Router),
		router.Network(Network),
		router.Gateway(gateway),
	)

	// create tunnel client with tunnel transport
	tr := transport.NewTransport(
		transport.WithTunnel(t),
	)
	c := client.NewClient(
		client.Transport(tr),
	)

	// local proxy
	localProxy := mucp.NewProxy(
		options.WithValue("proxy.router", r),
		options.WithValue("proxy.client", c),
	)

	// local server
	localSrv := server.NewServer(
		server.WithRouter(localProxy),
	)

	// init server
	service.Init(
		micro.Server(localSrv),
	)

	var wg sync.WaitGroup

	// error channel to collect errors and bail
	errChan := make(chan error, 2)

	// exit channelt o exit tunnel listener
	exit := make(chan struct{})

	// accept new tunnel connections
	wg.Add(1)
	go func() {
		defer wg.Done()
		errChan <- accept(l, exit)
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		errChan <- service.Run()
	}()

	// we block here until either service or server fails
	if err := <-errChan; err != nil {
		log.Logf("[tunnel] error running the tunnel: %v", err)
	}

	log.Log("[tunnel] attempting to stop the tunnel")

	// close tunnel listener
	close(exit)

	if err := l.Close(); err != nil {
		log.Logf("[tunnel] error closing the tunnel: %v", err)
	}

	// stop the router
	if err := r.Stop(); err != nil {
		log.Logf("[tunnel] error stopping tunnel router:%v", err)
	}

	wg.Wait()

	log.Logf("[tunnel] successfully stopped")
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "tunnel",
		Usage: "Run the micro network tunnel",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "tunnel_address",
				Usage:  "Set the micro tunnel address :9096",
				EnvVar: "MICRO_TUNNEL_ADDRESS",
			},
			cli.StringFlag{
				Name:   "network_address",
				Usage:  "Set the micro network address: local",
				EnvVar: "MICRO_NETWORK_ADDRESS",
			},
			cli.StringFlag{
				Name:   "gateway_address",
				Usage:  "Set the micro default gateway address :9094",
				EnvVar: "MICRO_GATEWAY_ADDRESS",
			},
		},
		Action: func(ctx *cli.Context) {
			run(ctx, options...)
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<commit_msg>Add tunnel server listener<commit_after>package tunnel

import (
	"os"
	"sync"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/config/options"
	"github.com/micro/go-micro/proxy/mucp"
	"github.com/micro/go-micro/router"
	"github.com/micro/go-micro/server"
	tun "github.com/micro/go-micro/tunnel"
	"github.com/micro/go-micro/tunnel/transport"
	"github.com/micro/go-micro/util/log"
)

var (
	// Name of the router microservice
	Name = "go.micro.tunnel"
	// Address is the tunnel microservice bind address
	Address = ":8084"
	// Tunnel is the tunnel bind address
	Tunnel = ":9096"
	// Channel is the name of the tunnel session
	Channel = "tunnel"
	// Router is the router gossip bind address
	Router = ":9093"
	// Network is the network id
	Network = "local"
)

// run runs the micro server
func run(ctx *cli.Context, srvOpts ...micro.Option) {
	// Init plugins
	for _, p := range Plugins() {
		p.Init(ctx)
	}

	if len(ctx.GlobalString("server_name")) > 0 {
		Name = ctx.GlobalString("server_name")
	}
	if len(ctx.String("address")) > 0 {
		Address = ctx.String("address")
	}
	if len(ctx.String("network_address")) > 0 {
		Network = ctx.String("network")
	}
	if len(ctx.String("tunnel_address")) > 0 {
		Tunnel = ctx.String("tunnel")
	}
	// default gateway address
	var gateway string
	if len(ctx.String("gateway_address")) > 0 {
		gateway = ctx.String("gateway")
	}

	// create a tunnel
	t := tun.NewTunnel(
		tun.Address(Tunnel),
	)

	// connect the tunnel
	if err := t.Connect(); err != nil {
		log.Logf("[tunnel] failed to connect: %s", err)
		os.Exit(1)
	}

	// Initialise service
	service := micro.NewService(
		micro.Name(Name),
		micro.Address(Address),
		micro.RegisterTTL(time.Duration(ctx.GlobalInt("register_ttl"))*time.Second),
		micro.RegisterInterval(time.Duration(ctx.GlobalInt("register_interval"))*time.Second),
	)

	// local tunnel router
	r := router.NewRouter(
		router.Id(service.Server().Options().Id),
		router.Registry(service.Client().Options().Registry),
		router.Address(Router),
		router.Network(Network),
		router.Gateway(gateway),
	)

	// create tunnel client with tunnel transport
	tunTransport := transport.NewTransport(
		transport.WithTunnel(t),
	)

	// local server client talks to tunnel
	localSrvClient := client.NewClient(
		client.Transport(tunTransport),
	)

	// local proxy
	localProxy := mucp.NewProxy(
		options.WithValue("proxy.router", r),
		options.WithValue("proxy.client", localSrvClient),
	)

	// local server
	localSrv := server.NewServer(
		server.WithRouter(localProxy),
	)

	// init server
	service.Init(
		micro.Server(localSrv),
	)

	// local transport client
	tunSrvClient := client.NewClient(
		client.Transport(service.Options().Transport),
	)

	// local proxy
	tunProxy := mucp.NewProxy(
		options.WithValue("proxy.client", tunSrvClient),
	)

	// local server
	tunSrv := server.NewServer(
		server.Transport(tunTransport),
		server.WithRouter(tunProxy),
	)

	var wg sync.WaitGroup

	// error channel to collect errors and bail
	errChan := make(chan error, 2)

	wg.Add(1)
	go func() {
		defer wg.Done()
		errChan <- service.Run()
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		errChan <- tunSrv.Start()
	}()

	// we block here until either service or server fails
	if err := <-errChan; err != nil {
		log.Logf("[tunnel] error running the tunnel: %v", err)
	}

	log.Log("[tunnel] attempting to stop the tunnel")

	// stop the router
	if err := r.Stop(); err != nil {
		log.Logf("[tunnel] error stopping tunnel router:%v", err)
	}

	wg.Wait()

	log.Logf("[tunnel] successfully stopped")
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "tunnel",
		Usage: "Run the micro network tunnel",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "tunnel_address",
				Usage:  "Set the micro tunnel address :9096",
				EnvVar: "MICRO_TUNNEL_ADDRESS",
			},
			cli.StringFlag{
				Name:   "network_address",
				Usage:  "Set the micro network address: local",
				EnvVar: "MICRO_NETWORK_ADDRESS",
			},
			cli.StringFlag{
				Name:   "gateway_address",
				Usage:  "Set the micro default gateway address :9094",
				EnvVar: "MICRO_GATEWAY_ADDRESS",
			},
		},
		Action: func(ctx *cli.Context) {
			run(ctx, options...)
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC.
//
// 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.

// Binary itest is an integration test for the API with the Stackdriver.
// To run the test:
//   go run lib/auditlogsapi/itest/main.go --alsologtostderr --project=ghasemloo-hcls-fa8 --user="subject"
package main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"strings"
	"time"

	glog "github.com/golang/glog" /* copybara-comment */
	"cloud.google.com/go/logging" /* copybara-comment: logging */
	"github.com/google/go-cmp/cmp" /* copybara-comment */
	"google.golang.org/api/option" /* copybara-comment: option */
	"google.golang.org/grpc/credentials" /* copybara-comment: credentials */
	"google.golang.org/grpc/credentials/oauth" /* copybara-comment: oauth */
	"google.golang.org/grpc" /* copybara-comment */
	"google.golang.org/protobuf/testing/protocmp" /* copybara-comment */
	"github.com/pborman/uuid" /* copybara-comment */
	"github.com/GoogleCloudPlatform/healthcare-federated-access-services/lib/auditlog" /* copybara-comment: auditlog */
	"github.com/GoogleCloudPlatform/healthcare-federated-access-services/lib/auditlogsapi" /* copybara-comment: auditlogsapi */
	"github.com/GoogleCloudPlatform/healthcare-federated-access-services/lib/httputils" /* copybara-comment: httputils */

	lgrpcpb "google.golang.org/genproto/googleapis/logging/v2" /* copybara-comment: logging_go_grpc */
	lpb "google.golang.org/genproto/googleapis/logging/v2" /* copybara-comment: logging_go_proto */
	dpb "github.com/golang/protobuf/ptypes/duration" /* copybara-comment */
	apb "github.com/GoogleCloudPlatform/healthcare-federated-access-services/proto/auditlogs/v0" /* copybara-comment: auditlogs_go_proto */
)

var (
	projectID = flag.String("project", "", "GCP project ID")
	userID    = flag.String("user", "", "user id (cirrently the subject of tokens)")
	sdlAddr   = flag.String("sdl_addr", "logging.googleapis.com:443", "The address for Stackdriver Logging API")
)

func main() {
	ctx := context.Background()
	flag.Parse()
	conn := NewGRPCClient(ctx, *sdlAddr)
	defer conn.Close()

	projectName := "projects/" + *projectID

	sdlc := lgrpcpb.NewLoggingServiceV2Client(conn)
	logger := NewLogger(ctx, conn, *projectID)

	s := auditlogsapi.NewAuditLogs(sdlc, *projectID)

	TestListLogEntriesRequest(ctx, sdlc, projectName, *userID)
	TestListAuditLogFromProject(ctx, s, projectName, *userID)
	TestAuditLog(ctx, s, sdlc, logger, projectName, *userID)

	glog.Exit()
}

// TestListLogEntriesRequest reads logs audit logs from Stackdriver Logging and displays them.
// Checks that the way we call the Stackdriver Logging actually works.
func TestListLogEntriesRequest(ctx context.Context, c lgrpcpb.LoggingServiceV2Client, projectName string, userID string) {
	glog.Infof("## TestListLogEntriesRequest started. ##")
	defer glog.Infof("## TestListLogEntriesRequest finished. ##\n\n")

	since := time.Date(2020, time.January, 1, 0, 0, 0, 0, time.UTC)
	filters := []string{
		fmt.Sprintf("timestamp>=%q", since.Format(time.RFC3339)),
		`logName="` + projectName + `/logs/federated-access-audit"`,
		`labels.token_subject="` + userID + `"`,
	}
	req := &lpb.ListLogEntriesRequest{
		ResourceNames: []string{projectName},
		PageSize:      1000,
		OrderBy:       "timestamp desc",
		Filter:        strings.Join(filters, " AND "),
	}
	glog.Infof("Request: %v\n\n", req)
	resp, err := c.ListLogEntries(ctx, req)
	if err != nil {
		glog.Errorf("ListLogEntries() failed: %v", err)
		return
	}
	glog.Infof("Response: %v\n\n", resp)
}

// TestListAuditLogFromProject checks the intergration of both auditlogsapi packages with Stackdriver logging.
// Gets auditlogs for the specified user with default options for listing and logs them.
func TestListAuditLogFromProject(ctx context.Context, s *auditlogsapi.AuditLogs, projectName string, userID string) {
	glog.Infof("## TestListAuditLogFromProject started. ##")
	defer glog.Infof("## TestListAuditLogFromProject finished. ##\n\n")
	resp, err := s.ListAuditLogs(ctx, &apb.ListAuditLogsRequest{Parent: "users/" + userID})
	if err != nil {
		glog.Errorf("ListAuditLogs() failed: %v", err)
		return
	}
	for _, l := range resp.GetAuditLogs() {
		glog.Infof("AuditLog: %v\n\n", l)
	}
}

// TestAuditLog checks the intergration of both auditlog and auditlogsapi packages with Stackdriver logging.
// It creates audit logs for a user (a generated uuid) and lists and logs them.
func TestAuditLog(ctx context.Context, s *auditlogsapi.AuditLogs, c lgrpcpb.LoggingServiceV2Client, logger *logging.Client, projectName string, userID string) {
	glog.Infof("## TestAuditLog started. ##")
	defer glog.Infof("## TestAuditLog finished. ##\n\n")

	// Generate random user id if no user is specified.
	randomUser := false
	if userID == "" {
		userID = "fake-user-id-" + uuid.New()
		randomUser = true
	}

	// Write an access and a policy log.
	al := &auditlog.AccessLog{
		TokenID:         "fake-token-id",
		TokenSubject:    userID,
		TokenIssuer:     "fake-issuer-id",
		TracingID:       "fake-tracing-id",
		RequestMethod:   "fake-method",
		RequestEndpoint: "fake-endpoint",
		RequestIP:       "fake-requester-ip",
		ErrorType:       "fake-error-type",
		ResponseCode:    1234,
		Request:         httputils.MustNewReq(http.MethodGet, "http://fake.org/fake-path", nil),
		PassAuthCheck:   true,
		Payload:         "fake-reason",
	}
	auditlog.WriteAccessLog(ctx, logger, al)

	pl := &auditlog.PolicyDecisionLog{
		TokenID:       "fake-token-id",
		TokenSubject:  userID,
		TokenIssuer:   "fake-issuer-id",
		Resource:      "fake-resource",
		TTL:           time.Hour.String(),
		PassAuthCheck: true,
		ErrorType:     "fake-error-type",
		Message:       "fake-reason",
	}
	auditlog.WritePolicyDecisionLog(logger, pl)

	// It takes a while before written logs are visible on Stackdriver.
	var got *apb.ListAuditLogsResponse
	end := time.Now().Add(time.Minute)
	for len(got.GetAuditLogs()) < 2 && time.Now().Before(end) {
		var err error
		got, err = s.ListAuditLogs(ctx, &apb.ListAuditLogsRequest{Parent: "users/" + userID})
		if err != nil {
			glog.Errorf("ListAuditLogs() failed: %v", err)
			return
		}
		time.Sleep(time.Second)
	}

	for _, l := range got.GetAuditLogs() {
		glog.Infof("AuditLog: %v\n\n", l)
	}

	if !randomUser {
		return
	}
	// If it was a randomly generated user, we can check if the logs are correct.

	if got.AuditLogs[0].AccessLog == nil {
		got.AuditLogs[0], got.AuditLogs[1] = got.AuditLogs[1], got.AuditLogs[0]
	}
	got.AuditLogs[0].Name = ""
	got.AuditLogs[0].AccessLog.Time = nil
	got.AuditLogs[1].Name = ""
	got.AuditLogs[1].PolicyLog.Time = nil

	want := &apb.ListAuditLogsResponse{
		AuditLogs: []*apb.AuditLog{
			{
				Name: "",
				AccessLog: &apb.AccessLog{
					ServiceName:      "unset-serviceinfo-Name",
					ServiceType:      "unset-serviceinfo-Type",
					TokenId:          "fake-token-id",
					TokenSubject:     userID,
					TokenIssuer:      "fake-issuer-id",
					Decision:         apb.Decision_PASS,
					ErrorType:        "fake-error-type",
					Reason:           "fake-reason",
					MethodName:       http.MethodGet,
					ResourceName:     "fake-endpoint",
					TracingId:        "fake-tracing-id",
					CallerIp:         "fake-requester-ip",
					HttpResponseCode: 1234,
					HttpRequest:      nil,
				},
			},
			{
				Name: "",
				PolicyLog: &apb.PolicyLog{
					ServiceName:  "unset-serviceinfo-Name",
					ServiceType:  "unset-serviceinfo-Type",
					TokenId:      "fake-token-id",
					TokenSubject: userID,
					TokenIssuer:  "fake-issuer-id",
					Decision:     apb.Decision_PASS,
					ErrorType:    "fake-error-type",
					Reason:       "fake-reason",
					ResourceName: "fake-resource",
					Ttl:          &dpb.Duration{Seconds: int64(time.Hour / time.Second)},
				},
			},
		},
	}
	if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" {
		glog.Errorf("ListAuditLogs() returned diff (-want +got):\n%s", diff)
	}
}

// NewGRPCClient creates a new GRPC client connect to the provided address.
func NewGRPCClient(ctx context.Context, addr string, opts ...grpc.DialOption) *grpc.ClientConn {
	opts = append(opts, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")))
	creds, err := oauth.NewApplicationDefault(ctx)
	if err != nil {
		glog.Exitf("oauth.NewApplicationDefault() failed: %v", err)
	}
	opts = append(opts, grpc.WithPerRPCCredentials(creds))
	conn, err := grpc.Dial(addr, opts...)
	if err != nil {
		glog.Exitf("Failed to connect to %q: %v", addr, err)
	}
	return conn
}

// NewLogger creates a new logger.
func NewLogger(ctx context.Context, conn *grpc.ClientConn, projectName string) *logging.Client {
	auditlog.LogSync = true
	logger, err := logging.NewClient(ctx, projectName, option.WithGRPCConn(conn))
	if err != nil {
		glog.Exitf("logging.NewClient(ctx,%v) failed: %v", projectName, err)
	}
	return logger
}
<commit_msg>explicitly specify the scope to be "https://www.googleapis.com/auth/cloud-platform"<commit_after>// Copyright 2020 Google LLC.
//
// 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.

// Binary itest is an integration test for the API with the Stackdriver.
// To run the test:
//   go run lib/auditlogsapi/itest/main.go --alsologtostderr --project=ghasemloo-hcls-fa8 --user="subject"
package main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"strings"
	"time"

	glog "github.com/golang/glog" /* copybara-comment */
	"cloud.google.com/go/logging" /* copybara-comment: logging */
	"github.com/google/go-cmp/cmp" /* copybara-comment */
	"google.golang.org/api/option" /* copybara-comment: option */
	"google.golang.org/grpc/credentials" /* copybara-comment: credentials */
	"google.golang.org/grpc/credentials/oauth" /* copybara-comment: oauth */
	"google.golang.org/grpc" /* copybara-comment */
	"google.golang.org/protobuf/testing/protocmp" /* copybara-comment */
	"github.com/pborman/uuid" /* copybara-comment */
	"github.com/GoogleCloudPlatform/healthcare-federated-access-services/lib/auditlog" /* copybara-comment: auditlog */
	"github.com/GoogleCloudPlatform/healthcare-federated-access-services/lib/auditlogsapi" /* copybara-comment: auditlogsapi */
	"github.com/GoogleCloudPlatform/healthcare-federated-access-services/lib/httputils" /* copybara-comment: httputils */

	lgrpcpb "google.golang.org/genproto/googleapis/logging/v2" /* copybara-comment: logging_go_grpc */
	lpb "google.golang.org/genproto/googleapis/logging/v2" /* copybara-comment: logging_go_proto */
	dpb "github.com/golang/protobuf/ptypes/duration" /* copybara-comment */
	apb "github.com/GoogleCloudPlatform/healthcare-federated-access-services/proto/auditlogs/v0" /* copybara-comment: auditlogs_go_proto */
)

var (
	projectID = flag.String("project", "", "GCP project ID")
	userID    = flag.String("user", "", "user id (cirrently the subject of tokens)")
	sdlAddr   = flag.String("sdl_addr", "logging.googleapis.com:443", "The address for Stackdriver Logging API")
)

func main() {
	ctx := context.Background()
	flag.Parse()
	conn := NewGRPCClient(ctx, *sdlAddr)
	defer conn.Close()

	projectName := "projects/" + *projectID

	sdlc := lgrpcpb.NewLoggingServiceV2Client(conn)
	logger := NewLogger(ctx, conn, *projectID)

	s := auditlogsapi.NewAuditLogs(sdlc, *projectID)

	TestListLogEntriesRequest(ctx, sdlc, projectName, *userID)
	TestListAuditLogFromProject(ctx, s, projectName, *userID)
	TestAuditLog(ctx, s, sdlc, logger, projectName, *userID)

	glog.Exit()
}

// TestListLogEntriesRequest reads logs audit logs from Stackdriver Logging and displays them.
// Checks that the way we call the Stackdriver Logging actually works.
func TestListLogEntriesRequest(ctx context.Context, c lgrpcpb.LoggingServiceV2Client, projectName string, userID string) {
	glog.Infof("## TestListLogEntriesRequest started. ##")
	defer glog.Infof("## TestListLogEntriesRequest finished. ##\n\n")

	since := time.Date(2020, time.January, 1, 0, 0, 0, 0, time.UTC)
	filters := []string{
		fmt.Sprintf("timestamp>=%q", since.Format(time.RFC3339)),
		`logName="` + projectName + `/logs/federated-access-audit"`,
		`labels.token_subject="` + userID + `"`,
	}
	req := &lpb.ListLogEntriesRequest{
		ResourceNames: []string{projectName},
		PageSize:      1000,
		OrderBy:       "timestamp desc",
		Filter:        strings.Join(filters, " AND "),
	}
	glog.Infof("Request: %v\n\n", req)
	resp, err := c.ListLogEntries(ctx, req)
	if err != nil {
		glog.Errorf("ListLogEntries() failed: %v", err)
		return
	}
	glog.Infof("Response: %v\n\n", resp)
}

// TestListAuditLogFromProject checks the intergration of both auditlogsapi packages with Stackdriver logging.
// Gets auditlogs for the specified user with default options for listing and logs them.
func TestListAuditLogFromProject(ctx context.Context, s *auditlogsapi.AuditLogs, projectName string, userID string) {
	glog.Infof("## TestListAuditLogFromProject started. ##")
	defer glog.Infof("## TestListAuditLogFromProject finished. ##\n\n")
	resp, err := s.ListAuditLogs(ctx, &apb.ListAuditLogsRequest{Parent: "users/" + userID})
	if err != nil {
		glog.Errorf("ListAuditLogs() failed: %v", err)
		return
	}
	for _, l := range resp.GetAuditLogs() {
		glog.Infof("AuditLog: %v\n\n", l)
	}
}

// TestAuditLog checks the intergration of both auditlog and auditlogsapi packages with Stackdriver logging.
// It creates audit logs for a user (a generated uuid) and lists and logs them.
func TestAuditLog(ctx context.Context, s *auditlogsapi.AuditLogs, c lgrpcpb.LoggingServiceV2Client, logger *logging.Client, projectName string, userID string) {
	glog.Infof("## TestAuditLog started. ##")
	defer glog.Infof("## TestAuditLog finished. ##\n\n")

	// Generate random user id if no user is specified.
	randomUser := false
	if userID == "" {
		userID = "fake-user-id-" + uuid.New()
		randomUser = true
	}

	// Write an access and a policy log.
	al := &auditlog.AccessLog{
		TokenID:         "fake-token-id",
		TokenSubject:    userID,
		TokenIssuer:     "fake-issuer-id",
		TracingID:       "fake-tracing-id",
		RequestMethod:   "fake-method",
		RequestEndpoint: "fake-endpoint",
		RequestIP:       "fake-requester-ip",
		ErrorType:       "fake-error-type",
		ResponseCode:    1234,
		Request:         httputils.MustNewReq(http.MethodGet, "http://fake.org/fake-path", nil),
		PassAuthCheck:   true,
		Payload:         "fake-reason",
	}
	auditlog.WriteAccessLog(ctx, logger, al)

	pl := &auditlog.PolicyDecisionLog{
		TokenID:       "fake-token-id",
		TokenSubject:  userID,
		TokenIssuer:   "fake-issuer-id",
		Resource:      "fake-resource",
		TTL:           time.Hour.String(),
		PassAuthCheck: true,
		ErrorType:     "fake-error-type",
		Message:       "fake-reason",
	}
	auditlog.WritePolicyDecisionLog(logger, pl)

	// It takes a while before written logs are visible on Stackdriver.
	var got *apb.ListAuditLogsResponse
	end := time.Now().Add(time.Minute)
	for len(got.GetAuditLogs()) < 2 && time.Now().Before(end) {
		var err error
		got, err = s.ListAuditLogs(ctx, &apb.ListAuditLogsRequest{Parent: "users/" + userID})
		if err != nil {
			glog.Errorf("ListAuditLogs() failed: %v", err)
			return
		}
		time.Sleep(time.Second)
	}

	for _, l := range got.GetAuditLogs() {
		glog.Infof("AuditLog: %v\n\n", l)
	}

	if !randomUser {
		return
	}
	// If it was a randomly generated user, we can check if the logs are correct.

	if got.AuditLogs[0].AccessLog == nil {
		got.AuditLogs[0], got.AuditLogs[1] = got.AuditLogs[1], got.AuditLogs[0]
	}
	got.AuditLogs[0].Name = ""
	got.AuditLogs[0].AccessLog.Time = nil
	got.AuditLogs[1].Name = ""
	got.AuditLogs[1].PolicyLog.Time = nil

	want := &apb.ListAuditLogsResponse{
		AuditLogs: []*apb.AuditLog{
			{
				Name: "",
				AccessLog: &apb.AccessLog{
					ServiceName:      "unset-serviceinfo-Name",
					ServiceType:      "unset-serviceinfo-Type",
					TokenId:          "fake-token-id",
					TokenSubject:     userID,
					TokenIssuer:      "fake-issuer-id",
					Decision:         apb.Decision_PASS,
					ErrorType:        "fake-error-type",
					Reason:           "fake-reason",
					MethodName:       http.MethodGet,
					ResourceName:     "fake-endpoint",
					TracingId:        "fake-tracing-id",
					CallerIp:         "fake-requester-ip",
					HttpResponseCode: 1234,
					HttpRequest:      nil,
				},
			},
			{
				Name: "",
				PolicyLog: &apb.PolicyLog{
					ServiceName:  "unset-serviceinfo-Name",
					ServiceType:  "unset-serviceinfo-Type",
					TokenId:      "fake-token-id",
					TokenSubject: userID,
					TokenIssuer:  "fake-issuer-id",
					Decision:     apb.Decision_PASS,
					ErrorType:    "fake-error-type",
					Reason:       "fake-reason",
					ResourceName: "fake-resource",
					Ttl:          &dpb.Duration{Seconds: int64(time.Hour / time.Second)},
				},
			},
		},
	}
	if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" {
		glog.Errorf("ListAuditLogs() returned diff (-want +got):\n%s", diff)
	}
}

// NewGRPCClient creates a new GRPC client connect to the provided address.
func NewGRPCClient(ctx context.Context, addr string, opts ...grpc.DialOption) *grpc.ClientConn {
	opts = append(opts, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")))
	creds, err := oauth.NewApplicationDefault(ctx, "https://www.googleapis.com/auth/cloud-platform")
	if err != nil {
		glog.Exitf("oauth.NewApplicationDefault() failed: %v", err)
	}
	opts = append(opts, grpc.WithPerRPCCredentials(creds))
	conn, err := grpc.Dial(addr, opts...)
	if err != nil {
		glog.Exitf("Failed to connect to %q: %v", addr, err)
	}
	return conn
}

// NewLogger creates a new logger.
func NewLogger(ctx context.Context, conn *grpc.ClientConn, projectName string) *logging.Client {
	auditlog.LogSync = true
	logger, err := logging.NewClient(ctx, projectName, option.WithGRPCConn(conn))
	if err != nil {
		glog.Exitf("logging.NewClient(ctx,%v) failed: %v", projectName, err)
	}
	return logger
}
<|endoftext|>


<commit_before>package api

// StatusCode represents a valid LXD operation and container status.
type StatusCode int

// LXD status codes.
const (
	OperationCreated StatusCode = 100
	Started          StatusCode = 101
	Stopped          StatusCode = 102
	Running          StatusCode = 103
	Cancelling       StatusCode = 104
	Pending          StatusCode = 105
	Starting         StatusCode = 106
	Stopping         StatusCode = 107
	Aborting         StatusCode = 108
	Freezing         StatusCode = 109
	Frozen           StatusCode = 110
	Thawed           StatusCode = 111
	Error            StatusCode = 112

	Success StatusCode = 200

	Failure   StatusCode = 400
	Cancelled StatusCode = 401
)

// String returns a suitable string representation for the status code.
func (o StatusCode) String() string {
	return map[StatusCode]string{
		OperationCreated: "Operation created",
		Started:          "Started",
		Stopped:          "Stopped",
		Running:          "Running",
		Cancelling:       "Cancelling",
		Pending:          "Pending",
		Success:          "Success",
		Failure:          "Failure",
		Cancelled:        "Cancelled",
		Starting:         "Starting",
		Stopping:         "Stopping",
		Aborting:         "Aborting",
		Freezing:         "Freezing",
		Frozen:           "Frozen",
		Thawed:           "Thawed",
		Error:            "Error",
	}[o]
}

// IsFinal will return true if the status code indicates an end state.
func (o StatusCode) IsFinal() bool {
	return int(o) >= 200
}
<commit_msg>shared/api: Add Ready state<commit_after>package api

// StatusCode represents a valid LXD operation and container status.
type StatusCode int

// LXD status codes.
const (
	OperationCreated StatusCode = 100
	Started          StatusCode = 101
	Stopped          StatusCode = 102
	Running          StatusCode = 103
	Cancelling       StatusCode = 104
	Pending          StatusCode = 105
	Starting         StatusCode = 106
	Stopping         StatusCode = 107
	Aborting         StatusCode = 108
	Freezing         StatusCode = 109
	Frozen           StatusCode = 110
	Thawed           StatusCode = 111
	Error            StatusCode = 112
	Ready            StatusCode = 113

	Success StatusCode = 200

	Failure   StatusCode = 400
	Cancelled StatusCode = 401
)

// String returns a suitable string representation for the status code.
func (o StatusCode) String() string {
	return map[StatusCode]string{
		OperationCreated: "Operation created",
		Started:          "Started",
		Stopped:          "Stopped",
		Running:          "Running",
		Cancelling:       "Cancelling",
		Pending:          "Pending",
		Success:          "Success",
		Failure:          "Failure",
		Cancelled:        "Cancelled",
		Starting:         "Starting",
		Stopping:         "Stopping",
		Aborting:         "Aborting",
		Freezing:         "Freezing",
		Frozen:           "Frozen",
		Thawed:           "Thawed",
		Error:            "Error",
		Ready:            "Ready",
	}[o]
}

// IsFinal will return true if the status code indicates an end state.
func (o StatusCode) IsFinal() bool {
	return int(o) >= 200
}
<|endoftext|>


<commit_before>package node

import (
	"log"
	"net/http"
	"os"
	"time"

	"github.com/hashicorp/raft"

	"github.com/abrander/gansoi/database"
)

type (
	// Node represents a single gansoi node.
	Node struct {
		raft   *raft.Raft
		stream *HTTPStream
	}
)

// NewNode will initialize a new node.
func NewNode(db *database.Database, peerStore *PeerStore) (*Node, error) {
	var err error
	n := &Node{}

	// Raft config.
	conf := raft.DefaultConfig()
	conf.HeartbeatTimeout = 1000 * time.Millisecond
	conf.ElectionTimeout = 1000 * time.Millisecond
	conf.LeaderLeaseTimeout = 500 * time.Millisecond
	conf.CommitTimeout = 200 * time.Millisecond
	conf.Logger = log.New(os.Stdout, "      RAFT ", log.Lmicroseconds|log.Lshortfile)

	// Set up nice HTTP based transport.
	n.stream, err = NewHTTPStream(peerStore.Self())
	if err != nil {
		panic(err)
	}

	logger := log.New(os.Stdout, " TRANSPORT ", log.Ldate|log.Lmicroseconds|log.Lshortfile)
	transport := raft.NewNetworkTransportWithLogger(n.stream, 1, 0, logger)

	n.raft, err = raft.NewRaft(
		conf,
		db,
		raft.NewInmemStore(),
		raft.NewInmemStore(),
		raft.NewDiscardSnapshotStore(),
		peerStore,
		transport)
	if err != nil {
		return nil, err
	}

	return n, nil
}

// ServeHTTP implements the http.Handler interface.
func (n *Node) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	n.stream.ServeHTTP(w, r)
}
<commit_msg>Implement Get()/Set() for Node.<commit_after>package node

import (
	"log"
	"net/http"
	"os"
	"time"

	"github.com/hashicorp/raft"

	"github.com/abrander/gansoi/database"
)

type (
	// Node represents a single gansoi node.
	Node struct {
		db     *database.Database
		raft   *raft.Raft
		leader bool
		stream *HTTPStream
	}
)

// NewNode will initialize a new node.
func NewNode(db *database.Database, peerStore *PeerStore) (*Node, error) {
	var err error
	n := &Node{
		db: db,
	}

	// Raft config.
	conf := raft.DefaultConfig()
	conf.HeartbeatTimeout = 1000 * time.Millisecond
	conf.ElectionTimeout = 1000 * time.Millisecond
	conf.LeaderLeaseTimeout = 500 * time.Millisecond
	conf.CommitTimeout = 200 * time.Millisecond
	conf.Logger = log.New(os.Stdout, "      RAFT ", log.Lmicroseconds|log.Lshortfile)

	// Set up nice HTTP based transport.
	n.stream, err = NewHTTPStream(peerStore.Self())
	if err != nil {
		panic(err)
	}

	logger := log.New(os.Stdout, " TRANSPORT ", log.Ldate|log.Lmicroseconds|log.Lshortfile)
	transport := raft.NewNetworkTransportWithLogger(n.stream, 1, 0, logger)

	n.raft, err = raft.NewRaft(
		conf,
		db,
		raft.NewInmemStore(),
		raft.NewInmemStore(),
		raft.NewDiscardSnapshotStore(),
		peerStore,
		transport)
	if err != nil {
		return nil, err
	}

	lch := n.raft.LeaderCh()
	go func() {
		select {
		case leader := <-lch:
			n.leader = leader
		}
	}()

	return n, nil
}

// ServeHTTP implements the http.Handler interface.
func (n *Node) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	n.stream.ServeHTTP(w, r)
}

// Set will set a key in the generic Raft-backed key/value store.
func (n *Node) Set(key string, value []byte) error {
	if !n.leader {
		/* FIXME: Relay to leader somehow */
		return raft.ErrNotLeader
	}

	entry := database.NewLogEntry(database.CommandSet, key, value)
	n.raft.Apply(entry.Byte(), time.Minute)

	return nil
}

// Get will retrieve a key from the generic Raft-backed K/V store.
func (n *Node) Get(key string) ([]byte, error) {
	return n.db.Get(key)
}
<|endoftext|>


<commit_before>package paxos

import (
    "fmt"
    "strings"
    "strconv"

    "borg/assert"
    "testing"
    "container/vector"
)

const (
    iFrom = iota
    iTo
    iCmd
    iRnd
    iNumParts
)

func accept(me uint64, ins, outs chan string) {
    var rnd, vrnd uint64
    var vval string

    ch, sent := make(chan int), 0
    for in := range ins {
        parts := strings.Split(in, ":", iNumParts)
        if len(parts) != iNumParts {
            continue
        }

        inTo, _ := strconv.Btoui64(parts[iTo], 10)
        if inTo != me && parts[iTo] != "*" {
            continue
        }

        switch parts[iCmd] {
        case "INVITE":
            i, err := strconv.Btoui64(parts[iRnd], 10)
            if err != nil { continue }

            inFrom, err := strconv.Btoui64(parts[iFrom], 10)
            if err != nil { continue }

            switch {
                case i <= rnd:
                case i > rnd:
                    rnd = i

                    sent++
                    outTo := inFrom // reply to the sender
                    msg := fmt.Sprintf(
                        "%d:%d:ACCEPT:%d:%d:%s",
                        me,
                        outTo,
                        i,
                        vrnd,
                        vval,
                    )
                    go func(msg string) { outs <- msg ; ch <- 1 }(msg)
            }
        }
    }

    for x := 0; x < sent; x++ {
        <-ch
    }

    close(outs)
}

// TESTING

func slurp(ch chan string) (got []string) {
    for x := range ch { (*vector.StringVector)(&got).Push(x) }
    return
}

func TestAcceptsInvite(t *testing.T) {
    ins := make(chan string)
    outs := make(chan string)

    exp := []string{"2:1:ACCEPT:1:0:"}

    go accept(2, ins, outs)
    // Send a message with no senderId
    ins <- "1:*:INVITE:1"
    close(ins)

    // outs was closed; therefore all messages have been processed
    assert.Equal(t, exp, slurp(outs), "")
}

func TestIgnoresStaleInvites(t *testing.T) {
    ins := make(chan string)
    outs := make(chan string)

    exp := []string{"2:1:ACCEPT:2:0:"}

    go accept(2, ins, outs)
    // Send a message with no senderId
    ins <- "1:*:INVITE:2"
    ins <- "1:*:INVITE:1"
    close(ins)

    // outs was closed; therefore all messages have been processed
    assert.Equal(t, exp, slurp(outs), "")
}

func TestIgnoresMalformedMessages(t *testing.T) {
    totest := []string{
        "x", // too few separators
        "x:x", // too few separators
        "x:x:x", // too few separators
        "x:x:x:x:x", // too many separators
        "1:*:INVITE:x", // invalid round number
        "1:*:x:1", // unknown command
        "1:x:INVITE:1", // invalid to address
        "1:7:INVITE:1", // valid but incorrect to address
        "X:*:INVITE:1", // invalid from address
    }

    for _, msg := range(totest) {
        ins := make(chan string)
        outs := make(chan string)

        exp := []string{}

        go accept(2, ins, outs)
        // Send a message with no senderId
        ins <- msg
        close(ins)

        // outs was closed; therefore all messages have been processed
        assert.Equal(t, exp, slurp(outs), "")
    }
}
<commit_msg>keith style<commit_after>package paxos

import (
    "fmt"
    "strings"
    "strconv"

    "borg/assert"
    "testing"
    "container/vector"
)

const (
    iFrom = iota
    iTo
    iCmd
    iRnd
    iNumParts
)

func accept(me uint64, ins, outs chan string) {
    var rnd, vrnd uint64
    var vval string

    ch, sent := make(chan int), 0
    for in := range ins {
        parts := strings.Split(in, ":", iNumParts)
        if len(parts) != iNumParts {
            continue
        }

        inTo, _ := strconv.Btoui64(parts[iTo], 10)
        if inTo != me && parts[iTo] != "*" {
            continue
        }

        switch parts[iCmd] {
        case "INVITE":
            i, err := strconv.Btoui64(parts[iRnd], 10)
            if err != nil { continue }

            inFrom, err := strconv.Btoui64(parts[iFrom], 10)
            if err != nil { continue }

            switch {
                case i <= rnd:
                case i > rnd:
                    rnd = i

                    sent++
                    outTo := inFrom // reply to the sender
                    msg := fmt.Sprintf("%d:%d:ACCEPT:%d:%d:%s",
                        me,
                        outTo,
                        i,
                        vrnd,
                        vval,
                    )
                    go func(msg string) { outs <- msg ; ch <- 1 }(msg)
            }
        }
    }

    for x := 0; x < sent; x++ {
        <-ch
    }

    close(outs)
}

// TESTING

func slurp(ch chan string) (got []string) {
    for x := range ch { (*vector.StringVector)(&got).Push(x) }
    return
}

func TestAcceptsInvite(t *testing.T) {
    ins := make(chan string)
    outs := make(chan string)

    exp := []string{"2:1:ACCEPT:1:0:"}

    go accept(2, ins, outs)
    // Send a message with no senderId
    ins <- "1:*:INVITE:1"
    close(ins)

    // outs was closed; therefore all messages have been processed
    assert.Equal(t, exp, slurp(outs), "")
}

func TestIgnoresStaleInvites(t *testing.T) {
    ins := make(chan string)
    outs := make(chan string)

    exp := []string{"2:1:ACCEPT:2:0:"}

    go accept(2, ins, outs)
    // Send a message with no senderId
    ins <- "1:*:INVITE:2"
    ins <- "1:*:INVITE:1"
    close(ins)

    // outs was closed; therefore all messages have been processed
    assert.Equal(t, exp, slurp(outs), "")
}

func TestIgnoresMalformedMessages(t *testing.T) {
    totest := []string{
        "x", // too few separators
        "x:x", // too few separators
        "x:x:x", // too few separators
        "x:x:x:x:x", // too many separators
        "1:*:INVITE:x", // invalid round number
        "1:*:x:1", // unknown command
        "1:x:INVITE:1", // invalid to address
        "1:7:INVITE:1", // valid but incorrect to address
        "X:*:INVITE:1", // invalid from address
    }

    for _, msg := range(totest) {
        ins := make(chan string)
        outs := make(chan string)

        exp := []string{}

        go accept(2, ins, outs)
        // Send a message with no senderId
        ins <- msg
        close(ins)

        // outs was closed; therefore all messages have been processed
        assert.Equal(t, exp, slurp(outs), "")
    }
}
<|endoftext|>


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

package registry

import (
	"bytes"
	"fmt"
	"hash"
	"hash/fnv"
	"sort"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"

	"github.com/prometheus/statsd_exporter/pkg/clock"
	"github.com/prometheus/statsd_exporter/pkg/mapper"
	"github.com/prometheus/statsd_exporter/pkg/metrics"
)

// uncheckedCollector wraps a Collector but its Describe method yields no Desc.
// This allows incoming metrics to have inconsistent label sets
type uncheckedCollector struct {
	c prometheus.Collector
}

func (u uncheckedCollector) Describe(_ chan<- *prometheus.Desc) {}
func (u uncheckedCollector) Collect(c chan<- prometheus.Metric) {
	u.c.Collect(c)
}

type Registry struct {
	Registerer prometheus.Registerer
	Metrics    map[string]metrics.Metric
	Mapper     *mapper.MetricMapper
	// The below value and label variables are allocated in the registry struct
	// so that we don't have to allocate them every time have to compute a label
	// hash.
	ValueBuf, NameBuf bytes.Buffer
	Hasher            hash.Hash64
}

func NewRegistry(reg prometheus.Registerer, mapper *mapper.MetricMapper) *Registry {
	return &Registry{
		Registerer: reg,
		Metrics:    make(map[string]metrics.Metric),
		Mapper:     mapper,
		Hasher:     fnv.New64a(),
	}
}

func (r *Registry) MetricConflicts(metricName string, metricType metrics.MetricType) bool {
	vector, hasMetrics := r.Metrics[metricName]
	if !hasMetrics {
		// No metrics.Metric with this name exists
		return false
	}

	if vector.MetricType == metricType {
		// We've found a copy of this metrics.Metric with this type, but different
		// labels, so it's safe to create a new one.
		return false
	}

	// The metrics.Metric exists, but it's of a different type than we're trying to
	// create.
	return true
}

func (r *Registry) StoreCounter(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vec *prometheus.CounterVec, c prometheus.Counter, ttl time.Duration) {
	r.Store(metricName, hash, labels, vec, c, metrics.CounterMetricType, ttl)
}

func (r *Registry) StoreGauge(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vec *prometheus.GaugeVec, g prometheus.Gauge, ttl time.Duration) {
	r.Store(metricName, hash, labels, vec, g, metrics.GaugeMetricType, ttl)
}

func (r *Registry) StoreHistogram(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vec *prometheus.HistogramVec, o prometheus.Observer, ttl time.Duration) {
	r.Store(metricName, hash, labels, vec, o, metrics.HistogramMetricType, ttl)
}

func (r *Registry) StoreSummary(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vec *prometheus.SummaryVec, o prometheus.Observer, ttl time.Duration) {
	r.Store(metricName, hash, labels, vec, o, metrics.SummaryMetricType, ttl)
}

func (r *Registry) Store(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vh metrics.VectorHolder, mh metrics.MetricHolder, metricType metrics.MetricType, ttl time.Duration) {
	metric, hasMetrics := r.Metrics[metricName]
	if !hasMetrics {
		metric.MetricType = metricType
		metric.Vectors = make(map[metrics.NameHash]*metrics.Vector)
		metric.Metrics = make(map[metrics.ValueHash]*metrics.RegisteredMetric)

		r.Metrics[metricName] = metric
	}

	v, ok := metric.Vectors[hash.Names]
	if !ok {
		v = &metrics.Vector{Holder: vh}
		metric.Vectors[hash.Names] = v
	}

	now := clock.Now()
	rm, ok := metric.Metrics[hash.Values]
	if !ok {
		rm = &metrics.RegisteredMetric{
			LastRegisteredAt: now,
			Labels:           labels,
			TTL:              ttl,
			Metric:           mh,
			VecKey:           hash.Names,
		}
		metric.Metrics[hash.Values] = rm
		v.RefCount++
		return
	}
	rm.LastRegisteredAt = now
	// Update ttl from mapping
	rm.TTL = ttl
}

func (r *Registry) Get(metricName string, hash metrics.LabelHash, metricType metrics.MetricType) (metrics.VectorHolder, metrics.MetricHolder) {
	metric, hasMetric := r.Metrics[metricName]

	if !hasMetric {
		return nil, nil
	}
	if metric.MetricType != metricType {
		return nil, nil
	}

	rm, ok := metric.Metrics[hash.Values]
	if ok {
		now := clock.Now()
		rm.LastRegisteredAt = now
		return metric.Vectors[hash.Names].Holder, rm.Metric
	}

	vector, ok := metric.Vectors[hash.Names]
	if ok {
		return vector.Holder, nil
	}

	return nil, nil
}

func (r *Registry) GetCounter(metricName string, labels prometheus.Labels, help string, mapping *mapper.MetricMapping, metricsCount *prometheus.GaugeVec) (prometheus.Counter, error) {
	hash, labelNames := r.HashLabels(labels)
	vh, mh := r.Get(metricName, hash, metrics.CounterMetricType)
	if mh != nil {
		return mh.(prometheus.Counter), nil
	}

	if r.MetricConflicts(metricName, metrics.CounterMetricType) {
		return nil, fmt.Errorf("metric with name %s is already registered", metricName)
	}

	err := r.checkHistogramNameCollision(metricName)
	if err != nil {
		return nil, err
	}

	var counterVec *prometheus.CounterVec
	if vh == nil {
		metricsCount.WithLabelValues("counter").Inc()
		counterVec = prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: metricName,
			Help: help,
		}, labelNames)

		if err := r.Registerer.Register(uncheckedCollector{counterVec}); err != nil {
			return nil, err
		}
	} else {
		counterVec = vh.(*prometheus.CounterVec)
	}

	var counter prometheus.Counter
	if counter, err = counterVec.GetMetricWith(labels); err != nil {
		return nil, err
	}
	r.StoreCounter(metricName, hash, labels, counterVec, counter, mapping.Ttl)

	return counter, nil
}

func (r *Registry) checkHistogramNameCollision(metricName string) error {
	histogramSuffixes := []string{"_bucket", "_count", "_sum"}
	for _, suffix := range histogramSuffixes {
		if strings.HasSuffix(metricName, suffix) {
			if r.MetricConflicts(strings.TrimSuffix(metricName, suffix), metrics.CounterMetricType) {
				return fmt.Errorf("metric with name %s is already registered", metricName)
			}
		}
	}
	return nil
}

func (r *Registry) GetGauge(metricName string, labels prometheus.Labels, help string, mapping *mapper.MetricMapping, metricsCount *prometheus.GaugeVec) (prometheus.Gauge, error) {
	hash, labelNames := r.HashLabels(labels)
	vh, mh := r.Get(metricName, hash, metrics.GaugeMetricType)
	if mh != nil {
		return mh.(prometheus.Gauge), nil
	}

	if r.MetricConflicts(metricName, metrics.GaugeMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}

	err := r.checkHistogramNameCollision(metricName)
	if err != nil {
		return nil, fmt.Errorf("metrics.Metric with name %s conflicts with previously registered histogram", metricName)
	}

	var gaugeVec *prometheus.GaugeVec
	if vh == nil {
		metricsCount.WithLabelValues("gauge").Inc()
		gaugeVec = prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Name: metricName,
			Help: help,
		}, labelNames)

		if err := r.Registerer.Register(uncheckedCollector{gaugeVec}); err != nil {
			return nil, err
		}
	} else {
		gaugeVec = vh.(*prometheus.GaugeVec)
	}

	var gauge prometheus.Gauge
	if gauge, err = gaugeVec.GetMetricWith(labels); err != nil {
		return nil, err
	}
	r.StoreGauge(metricName, hash, labels, gaugeVec, gauge, mapping.Ttl)

	return gauge, nil
}

func (r *Registry) GetHistogram(metricName string, labels prometheus.Labels, help string, mapping *mapper.MetricMapping, metricsCount *prometheus.GaugeVec) (prometheus.Observer, error) {
	hash, labelNames := r.HashLabels(labels)
	vh, mh := r.Get(metricName, hash, metrics.HistogramMetricType)
	if mh != nil {
		return mh.(prometheus.Observer), nil
	}

	if r.MetricConflicts(metricName, metrics.HistogramMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_sum", metrics.HistogramMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_count", metrics.HistogramMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_bucket", metrics.HistogramMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}

	var histogramVec *prometheus.HistogramVec
	if vh == nil {
		metricsCount.WithLabelValues("histogram").Inc()
		buckets := r.Mapper.Defaults.HistogramOptions.Buckets
		if mapping.HistogramOptions != nil && len(mapping.HistogramOptions.Buckets) > 0 {
			buckets = mapping.HistogramOptions.Buckets
		}
		histogramVec = prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    metricName,
			Help:    help,
			Buckets: buckets,
		}, labelNames)

		if err := r.Registerer.Register(uncheckedCollector{histogramVec}); err != nil {
			return nil, err
		}
	} else {
		histogramVec = vh.(*prometheus.HistogramVec)
	}

	var observer prometheus.Observer
	var err error
	if observer, err = histogramVec.GetMetricWith(labels); err != nil {
		return nil, err
	}
	r.StoreHistogram(metricName, hash, labels, histogramVec, observer, mapping.Ttl)

	return observer, nil
}

func (r *Registry) GetSummary(metricName string, labels prometheus.Labels, help string, mapping *mapper.MetricMapping, metricsCount *prometheus.GaugeVec) (prometheus.Observer, error) {
	hash, labelNames := r.HashLabels(labels)
	vh, mh := r.Get(metricName, hash, metrics.SummaryMetricType)
	if mh != nil {
		return mh.(prometheus.Observer), nil
	}

	if r.MetricConflicts(metricName, metrics.SummaryMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_sum", metrics.SummaryMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_count", metrics.SummaryMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}

	var summaryVec *prometheus.SummaryVec
	if vh == nil {
		metricsCount.WithLabelValues("summary").Inc()
		quantiles := r.Mapper.Defaults.SummaryOptions.Quantiles
		if mapping != nil && mapping.SummaryOptions != nil && len(mapping.SummaryOptions.Quantiles) > 0 {
			quantiles = mapping.SummaryOptions.Quantiles
		}

		summaryOptions := mapper.SummaryOptions{
			MaxAge:     r.Mapper.Defaults.SummaryOptions.MaxAge,
			AgeBuckets: r.Mapper.Defaults.SummaryOptions.AgeBuckets,
			BufCap:     r.Mapper.Defaults.SummaryOptions.BufCap,
		}

		if mapping != nil && mapping.SummaryOptions != nil {
			summaryOptions = *mapping.SummaryOptions
		}

		objectives := make(map[float64]float64)
		for _, q := range quantiles {
			objectives[q.Quantile] = q.Error
		}
		// In the case of no mapping file, explicitly define the default quantiles
		if len(objectives) == 0 {
			objectives = map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001}
		}
		summaryVec = prometheus.NewSummaryVec(prometheus.SummaryOpts{
			Name:       metricName,
			Help:       help,
			Objectives: objectives,
			MaxAge:     summaryOptions.MaxAge,
			AgeBuckets: summaryOptions.AgeBuckets,
			BufCap:     summaryOptions.BufCap,
		}, labelNames)

		if err := r.Registerer.Register(uncheckedCollector{summaryVec}); err != nil {
			return nil, err
		}
	} else {
		summaryVec = vh.(*prometheus.SummaryVec)
	}

	var observer prometheus.Observer
	var err error
	if observer, err = summaryVec.GetMetricWith(labels); err != nil {
		return nil, err
	}
	r.StoreSummary(metricName, hash, labels, summaryVec, observer, mapping.Ttl)

	return observer, nil
}

func (r *Registry) RemoveStaleMetrics() {
	now := clock.Now()
	// delete timeseries with expired ttl
	for _, metric := range r.Metrics {
		for hash, rm := range metric.Metrics {
			if rm.TTL == 0 {
				continue
			}
			if rm.LastRegisteredAt.Add(rm.TTL).Before(now) {
				metric.Vectors[rm.VecKey].Holder.Delete(rm.Labels)
				metric.Vectors[rm.VecKey].RefCount--
				delete(metric.Metrics, hash)
			}
		}
	}
}

// Calculates a hash of both the label names and values.
func (r *Registry) HashLabels(labels prometheus.Labels) (metrics.LabelHash, []string) {
	r.Hasher.Reset()
	r.NameBuf.Reset()
	r.ValueBuf.Reset()
	labelNames := make([]string, 0, len(labels))

	for labelName := range labels {
		labelNames = append(labelNames, labelName)
	}
	sort.Strings(labelNames)

	r.ValueBuf.WriteByte(model.SeparatorByte)
	for _, labelName := range labelNames {
		r.ValueBuf.WriteString(labels[labelName])
		r.ValueBuf.WriteByte(model.SeparatorByte)

		r.NameBuf.WriteString(labelName)
		r.NameBuf.WriteByte(model.SeparatorByte)
	}

	lh := metrics.LabelHash{}
	r.Hasher.Write(r.NameBuf.Bytes())
	lh.Names = metrics.NameHash(r.Hasher.Sum64())

	// Now add the values to the names we've already hashed.
	r.Hasher.Write(r.ValueBuf.Bytes())
	lh.Values = metrics.ValueHash(r.Hasher.Sum64())

	return lh, labelNames
}
<commit_msg>Use same error message<commit_after>// Copyright 2013 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package registry

import (
	"bytes"
	"fmt"
	"hash"
	"hash/fnv"
	"sort"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"

	"github.com/prometheus/statsd_exporter/pkg/clock"
	"github.com/prometheus/statsd_exporter/pkg/mapper"
	"github.com/prometheus/statsd_exporter/pkg/metrics"
)

// uncheckedCollector wraps a Collector but its Describe method yields no Desc.
// This allows incoming metrics to have inconsistent label sets
type uncheckedCollector struct {
	c prometheus.Collector
}

func (u uncheckedCollector) Describe(_ chan<- *prometheus.Desc) {}
func (u uncheckedCollector) Collect(c chan<- prometheus.Metric) {
	u.c.Collect(c)
}

type Registry struct {
	Registerer prometheus.Registerer
	Metrics    map[string]metrics.Metric
	Mapper     *mapper.MetricMapper
	// The below value and label variables are allocated in the registry struct
	// so that we don't have to allocate them every time have to compute a label
	// hash.
	ValueBuf, NameBuf bytes.Buffer
	Hasher            hash.Hash64
}

func NewRegistry(reg prometheus.Registerer, mapper *mapper.MetricMapper) *Registry {
	return &Registry{
		Registerer: reg,
		Metrics:    make(map[string]metrics.Metric),
		Mapper:     mapper,
		Hasher:     fnv.New64a(),
	}
}

func (r *Registry) MetricConflicts(metricName string, metricType metrics.MetricType) bool {
	vector, hasMetrics := r.Metrics[metricName]
	if !hasMetrics {
		// No metrics.Metric with this name exists
		return false
	}

	if vector.MetricType == metricType {
		// We've found a copy of this metrics.Metric with this type, but different
		// labels, so it's safe to create a new one.
		return false
	}

	// The metrics.Metric exists, but it's of a different type than we're trying to
	// create.
	return true
}

func (r *Registry) StoreCounter(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vec *prometheus.CounterVec, c prometheus.Counter, ttl time.Duration) {
	r.Store(metricName, hash, labels, vec, c, metrics.CounterMetricType, ttl)
}

func (r *Registry) StoreGauge(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vec *prometheus.GaugeVec, g prometheus.Gauge, ttl time.Duration) {
	r.Store(metricName, hash, labels, vec, g, metrics.GaugeMetricType, ttl)
}

func (r *Registry) StoreHistogram(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vec *prometheus.HistogramVec, o prometheus.Observer, ttl time.Duration) {
	r.Store(metricName, hash, labels, vec, o, metrics.HistogramMetricType, ttl)
}

func (r *Registry) StoreSummary(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vec *prometheus.SummaryVec, o prometheus.Observer, ttl time.Duration) {
	r.Store(metricName, hash, labels, vec, o, metrics.SummaryMetricType, ttl)
}

func (r *Registry) Store(metricName string, hash metrics.LabelHash, labels prometheus.Labels, vh metrics.VectorHolder, mh metrics.MetricHolder, metricType metrics.MetricType, ttl time.Duration) {
	metric, hasMetrics := r.Metrics[metricName]
	if !hasMetrics {
		metric.MetricType = metricType
		metric.Vectors = make(map[metrics.NameHash]*metrics.Vector)
		metric.Metrics = make(map[metrics.ValueHash]*metrics.RegisteredMetric)

		r.Metrics[metricName] = metric
	}

	v, ok := metric.Vectors[hash.Names]
	if !ok {
		v = &metrics.Vector{Holder: vh}
		metric.Vectors[hash.Names] = v
	}

	now := clock.Now()
	rm, ok := metric.Metrics[hash.Values]
	if !ok {
		rm = &metrics.RegisteredMetric{
			LastRegisteredAt: now,
			Labels:           labels,
			TTL:              ttl,
			Metric:           mh,
			VecKey:           hash.Names,
		}
		metric.Metrics[hash.Values] = rm
		v.RefCount++
		return
	}
	rm.LastRegisteredAt = now
	// Update ttl from mapping
	rm.TTL = ttl
}

func (r *Registry) Get(metricName string, hash metrics.LabelHash, metricType metrics.MetricType) (metrics.VectorHolder, metrics.MetricHolder) {
	metric, hasMetric := r.Metrics[metricName]

	if !hasMetric {
		return nil, nil
	}
	if metric.MetricType != metricType {
		return nil, nil
	}

	rm, ok := metric.Metrics[hash.Values]
	if ok {
		now := clock.Now()
		rm.LastRegisteredAt = now
		return metric.Vectors[hash.Names].Holder, rm.Metric
	}

	vector, ok := metric.Vectors[hash.Names]
	if ok {
		return vector.Holder, nil
	}

	return nil, nil
}

func (r *Registry) GetCounter(metricName string, labels prometheus.Labels, help string, mapping *mapper.MetricMapping, metricsCount *prometheus.GaugeVec) (prometheus.Counter, error) {
	hash, labelNames := r.HashLabels(labels)
	vh, mh := r.Get(metricName, hash, metrics.CounterMetricType)
	if mh != nil {
		return mh.(prometheus.Counter), nil
	}

	if r.MetricConflicts(metricName, metrics.CounterMetricType) {
		return nil, fmt.Errorf("metric with name %s is already registered", metricName)
	}

	err := r.checkHistogramNameCollision(metricName)
	if err != nil {
		return nil, err
	}

	var counterVec *prometheus.CounterVec
	if vh == nil {
		metricsCount.WithLabelValues("counter").Inc()
		counterVec = prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: metricName,
			Help: help,
		}, labelNames)

		if err := r.Registerer.Register(uncheckedCollector{counterVec}); err != nil {
			return nil, err
		}
	} else {
		counterVec = vh.(*prometheus.CounterVec)
	}

	var counter prometheus.Counter
	if counter, err = counterVec.GetMetricWith(labels); err != nil {
		return nil, err
	}
	r.StoreCounter(metricName, hash, labels, counterVec, counter, mapping.Ttl)

	return counter, nil
}

func (r *Registry) checkHistogramNameCollision(metricName string) error {
	histogramSuffixes := []string{"_bucket", "_count", "_sum"}
	for _, suffix := range histogramSuffixes {
		if strings.HasSuffix(metricName, suffix) {
			if r.MetricConflicts(strings.TrimSuffix(metricName, suffix), metrics.CounterMetricType) {
				return fmt.Errorf("metric with name %s is already registered", metricName)
			}
		}
	}
	return nil
}

func (r *Registry) GetGauge(metricName string, labels prometheus.Labels, help string, mapping *mapper.MetricMapping, metricsCount *prometheus.GaugeVec) (prometheus.Gauge, error) {
	hash, labelNames := r.HashLabels(labels)
	vh, mh := r.Get(metricName, hash, metrics.GaugeMetricType)
	if mh != nil {
		return mh.(prometheus.Gauge), nil
	}

	if r.MetricConflicts(metricName, metrics.GaugeMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}

	err := r.checkHistogramNameCollision(metricName)
	if err != nil {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}

	var gaugeVec *prometheus.GaugeVec
	if vh == nil {
		metricsCount.WithLabelValues("gauge").Inc()
		gaugeVec = prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Name: metricName,
			Help: help,
		}, labelNames)

		if err := r.Registerer.Register(uncheckedCollector{gaugeVec}); err != nil {
			return nil, err
		}
	} else {
		gaugeVec = vh.(*prometheus.GaugeVec)
	}

	var gauge prometheus.Gauge
	if gauge, err = gaugeVec.GetMetricWith(labels); err != nil {
		return nil, err
	}
	r.StoreGauge(metricName, hash, labels, gaugeVec, gauge, mapping.Ttl)

	return gauge, nil
}

func (r *Registry) GetHistogram(metricName string, labels prometheus.Labels, help string, mapping *mapper.MetricMapping, metricsCount *prometheus.GaugeVec) (prometheus.Observer, error) {
	hash, labelNames := r.HashLabels(labels)
	vh, mh := r.Get(metricName, hash, metrics.HistogramMetricType)
	if mh != nil {
		return mh.(prometheus.Observer), nil
	}

	if r.MetricConflicts(metricName, metrics.HistogramMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_sum", metrics.HistogramMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_count", metrics.HistogramMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_bucket", metrics.HistogramMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}

	var histogramVec *prometheus.HistogramVec
	if vh == nil {
		metricsCount.WithLabelValues("histogram").Inc()
		buckets := r.Mapper.Defaults.HistogramOptions.Buckets
		if mapping.HistogramOptions != nil && len(mapping.HistogramOptions.Buckets) > 0 {
			buckets = mapping.HistogramOptions.Buckets
		}
		histogramVec = prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    metricName,
			Help:    help,
			Buckets: buckets,
		}, labelNames)

		if err := r.Registerer.Register(uncheckedCollector{histogramVec}); err != nil {
			return nil, err
		}
	} else {
		histogramVec = vh.(*prometheus.HistogramVec)
	}

	var observer prometheus.Observer
	var err error
	if observer, err = histogramVec.GetMetricWith(labels); err != nil {
		return nil, err
	}
	r.StoreHistogram(metricName, hash, labels, histogramVec, observer, mapping.Ttl)

	return observer, nil
}

func (r *Registry) GetSummary(metricName string, labels prometheus.Labels, help string, mapping *mapper.MetricMapping, metricsCount *prometheus.GaugeVec) (prometheus.Observer, error) {
	hash, labelNames := r.HashLabels(labels)
	vh, mh := r.Get(metricName, hash, metrics.SummaryMetricType)
	if mh != nil {
		return mh.(prometheus.Observer), nil
	}

	if r.MetricConflicts(metricName, metrics.SummaryMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_sum", metrics.SummaryMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}
	if r.MetricConflicts(metricName+"_count", metrics.SummaryMetricType) {
		return nil, fmt.Errorf("metrics.Metric with name %s is already registered", metricName)
	}

	var summaryVec *prometheus.SummaryVec
	if vh == nil {
		metricsCount.WithLabelValues("summary").Inc()
		quantiles := r.Mapper.Defaults.SummaryOptions.Quantiles
		if mapping != nil && mapping.SummaryOptions != nil && len(mapping.SummaryOptions.Quantiles) > 0 {
			quantiles = mapping.SummaryOptions.Quantiles
		}

		summaryOptions := mapper.SummaryOptions{
			MaxAge:     r.Mapper.Defaults.SummaryOptions.MaxAge,
			AgeBuckets: r.Mapper.Defaults.SummaryOptions.AgeBuckets,
			BufCap:     r.Mapper.Defaults.SummaryOptions.BufCap,
		}

		if mapping != nil && mapping.SummaryOptions != nil {
			summaryOptions = *mapping.SummaryOptions
		}

		objectives := make(map[float64]float64)
		for _, q := range quantiles {
			objectives[q.Quantile] = q.Error
		}
		// In the case of no mapping file, explicitly define the default quantiles
		if len(objectives) == 0 {
			objectives = map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001}
		}
		summaryVec = prometheus.NewSummaryVec(prometheus.SummaryOpts{
			Name:       metricName,
			Help:       help,
			Objectives: objectives,
			MaxAge:     summaryOptions.MaxAge,
			AgeBuckets: summaryOptions.AgeBuckets,
			BufCap:     summaryOptions.BufCap,
		}, labelNames)

		if err := r.Registerer.Register(uncheckedCollector{summaryVec}); err != nil {
			return nil, err
		}
	} else {
		summaryVec = vh.(*prometheus.SummaryVec)
	}

	var observer prometheus.Observer
	var err error
	if observer, err = summaryVec.GetMetricWith(labels); err != nil {
		return nil, err
	}
	r.StoreSummary(metricName, hash, labels, summaryVec, observer, mapping.Ttl)

	return observer, nil
}

func (r *Registry) RemoveStaleMetrics() {
	now := clock.Now()
	// delete timeseries with expired ttl
	for _, metric := range r.Metrics {
		for hash, rm := range metric.Metrics {
			if rm.TTL == 0 {
				continue
			}
			if rm.LastRegisteredAt.Add(rm.TTL).Before(now) {
				metric.Vectors[rm.VecKey].Holder.Delete(rm.Labels)
				metric.Vectors[rm.VecKey].RefCount--
				delete(metric.Metrics, hash)
			}
		}
	}
}

// Calculates a hash of both the label names and values.
func (r *Registry) HashLabels(labels prometheus.Labels) (metrics.LabelHash, []string) {
	r.Hasher.Reset()
	r.NameBuf.Reset()
	r.ValueBuf.Reset()
	labelNames := make([]string, 0, len(labels))

	for labelName := range labels {
		labelNames = append(labelNames, labelName)
	}
	sort.Strings(labelNames)

	r.ValueBuf.WriteByte(model.SeparatorByte)
	for _, labelName := range labelNames {
		r.ValueBuf.WriteString(labels[labelName])
		r.ValueBuf.WriteByte(model.SeparatorByte)

		r.NameBuf.WriteString(labelName)
		r.NameBuf.WriteByte(model.SeparatorByte)
	}

	lh := metrics.LabelHash{}
	r.Hasher.Write(r.NameBuf.Bytes())
	lh.Names = metrics.NameHash(r.Hasher.Sum64())

	// Now add the values to the names we've already hashed.
	r.Hasher.Write(r.ValueBuf.Bytes())
	lh.Values = metrics.ValueHash(r.Hasher.Sum64())

	return lh, labelNames
}
<|endoftext|>


<commit_before>package remotedialer

import (
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
)

var (
	Token = "X-API-Tunnel-Token"
	ID    = "X-API-Tunnel-ID"
)

func (s *Server) AddPeer(url, id, token string) {
	if s.PeerID == "" || s.PeerToken == "" {
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	peer := peer{
		url:    url,
		id:     id,
		token:  token,
		cancel: cancel,
	}

	logrus.Infof("Adding peer %s, %s", url, id)

	s.peerLock.Lock()
	defer s.peerLock.Unlock()

	if p, ok := s.peers[id]; ok {
		if p.equals(peer) {
			return
		}
		p.cancel()
	}

	s.peers[id] = peer
	go peer.start(ctx, s)
}

func (s *Server) RemovePeer(id string) {
	s.peerLock.Lock()
	defer s.peerLock.Unlock()

	if p, ok := s.peers[id]; ok {
		logrus.Infof("Removing peer %s", id)
		p.cancel()
	}
	delete(s.peers, id)
}

type peer struct {
	url, id, token string
	cancel         func()
}

func (p peer) equals(other peer) bool {
	return p.url == other.url &&
		p.id == other.id &&
		p.token == other.token
}

func (p *peer) start(ctx context.Context, s *Server) {
	headers := http.Header{
		ID:    {s.PeerID},
		Token: {s.PeerToken},
	}

	dialer := &websocket.Dialer{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}

	for {
		select {
		case <-ctx.Done():
			break
		default:
		}

		ws, _, err := dialer.Dial(p.url, headers)
		if err != nil {
			logrus.Errorf("Failed to connect to peer %s %#v: %v", p.url, headers, err)
			time.Sleep(5 * time.Second)
			continue
		}

		session := newClientSession(func(string, string) bool { return true }, ws)
		session.dialer = func(network, address string) (net.Conn, error) {
			parts := strings.SplitN(network, "::", 2)
			if len(parts) != 2 {
				return nil, fmt.Errorf("invalid clientKey/proto: %s", network)
			}
			return s.Dial(parts[0], 15*time.Second, parts[1], address)
		}

		s.sessions.addListener(session)
		_, err = session.serve()
		s.sessions.removeListener(session)
		session.Close()

		if err != nil {
			logrus.Errorf("Failed to serve peer connection %s: %v", p.id, err)
		}

		ws.Close()
		time.Sleep(5 * time.Second)
	}
}
<commit_msg>Don't log token on connect failure<commit_after>package remotedialer

import (
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
)

var (
	Token = "X-API-Tunnel-Token"
	ID    = "X-API-Tunnel-ID"
)

func (s *Server) AddPeer(url, id, token string) {
	if s.PeerID == "" || s.PeerToken == "" {
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	peer := peer{
		url:    url,
		id:     id,
		token:  token,
		cancel: cancel,
	}

	logrus.Infof("Adding peer %s, %s", url, id)

	s.peerLock.Lock()
	defer s.peerLock.Unlock()

	if p, ok := s.peers[id]; ok {
		if p.equals(peer) {
			return
		}
		p.cancel()
	}

	s.peers[id] = peer
	go peer.start(ctx, s)
}

func (s *Server) RemovePeer(id string) {
	s.peerLock.Lock()
	defer s.peerLock.Unlock()

	if p, ok := s.peers[id]; ok {
		logrus.Infof("Removing peer %s", id)
		p.cancel()
	}
	delete(s.peers, id)
}

type peer struct {
	url, id, token string
	cancel         func()
}

func (p peer) equals(other peer) bool {
	return p.url == other.url &&
		p.id == other.id &&
		p.token == other.token
}

func (p *peer) start(ctx context.Context, s *Server) {
	headers := http.Header{
		ID:    {s.PeerID},
		Token: {s.PeerToken},
	}

	dialer := &websocket.Dialer{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}

	for {
		select {
		case <-ctx.Done():
			break
		default:
		}

		ws, _, err := dialer.Dial(p.url, headers)
		if err != nil {
			logrus.Errorf("Failed to connect to peer %s [local ID=%s]: %v", p.url, s.PeerID, err)
			time.Sleep(5 * time.Second)
			continue
		}

		session := newClientSession(func(string, string) bool { return true }, ws)
		session.dialer = func(network, address string) (net.Conn, error) {
			parts := strings.SplitN(network, "::", 2)
			if len(parts) != 2 {
				return nil, fmt.Errorf("invalid clientKey/proto: %s", network)
			}
			return s.Dial(parts[0], 15*time.Second, parts[1], address)
		}

		s.sessions.addListener(session)
		_, err = session.serve()
		s.sessions.removeListener(session)
		session.Close()

		if err != nil {
			logrus.Errorf("Failed to serve peer connection %s: %v", p.id, err)
		}

		ws.Close()
		time.Sleep(5 * time.Second)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The go-hep Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// root-ls dumps the content of a ROOT file
package main // import "go-hep.org/x/hep/rootio/cmd/root-ls"

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"text/tabwriter"

	"go-hep.org/x/hep/rootio"
)

var (
	doProf = flag.String("profile", "", "filename of cpuprofile")
	doSI   = flag.Bool("sinfos", false, "print StreamerInfos")
	doTree = flag.Bool("t", false, "print Tree(s) (recursively)")
)

func main() {
	flag.Parse()

	if *doProf != "" {
		f, err := os.Create(*doProf)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if flag.NArg() <= 0 {
		fmt.Fprintf(os.Stderr, "**error** you need to give a ROOT file\n")
		flag.Usage()
		os.Exit(1)
	}

	for ii, fname := range flag.Args() {

		if ii > 0 {
			fmt.Printf("\n")
		}

		fmt.Printf("=== [%s] ===\n", fname)
		f, err := rootio.Open(fname)
		if err != nil {
			fmt.Fprintf(os.Stderr, "rootio: failed to open [%s]: %v\n", fname, err)
			os.Exit(1)
		}
		fmt.Printf("version: %v\n", f.Version())
		if *doSI {
			fmt.Printf("streamer-infos:\n")
			sinfos := f.StreamerInfo()
			for _, v := range sinfos {
				name := v.Name()
				fmt.Printf(" StreamerInfo for %q version=%d title=%q\n", name, v.ClassVersion(), v.Title())
				w := tabwriter.NewWriter(os.Stdout, 8, 4, 1, ' ', 0)
				for _, elm := range v.Elements() {
					fmt.Fprintf(w, "  %s\t%s\toffset=%3d\ttype=%3d\tsize=%3d\t %s\n", elm.TypeName(), elm.Name(), elm.Offset(), elm.Type(), elm.Size(), elm.Title())
				}
				w.Flush()
			}
		}

		w := tabwriter.NewWriter(os.Stdout, 8, 4, 1, ' ', 0)
		for _, k := range f.Keys() {
			walk(w, k)
		}
		w.Flush()
	}
}

func walk(w io.Writer, k rootio.Key) {
	obj := k.Value()
	if *doTree {
		tree, ok := obj.(rootio.Tree)
		if ok {
			w := tabwriter.NewWriter(w, 8, 4, 1, ' ', 0)
			fmt.Fprintf(w, "%s\t%s\t%s\t(entries=%d)\n", k.Class(), k.Name(), k.Title(), tree.Entries())
			displayBranches(w, tree, 2)
			w.Flush()
			return
		}
	}
	fmt.Fprintf(w, "%s\t%s\t%s\t(cycle=%d)\n", k.Class(), k.Name(), k.Title(), k.Cycle())
	if dir, ok := obj.(rootio.Directory); ok {
		w := newWindent(2, w)
		for _, k := range dir.Keys() {
			walk(w, k)
		}
		w.Flush()
	}
}

type windent struct {
	hdr []byte
	w   io.Writer
}

func newWindent(n int, w io.Writer) *windent {
	return &windent{
		hdr: bytes.Repeat([]byte(" "), n),
		w:   w,
	}
}

func (w *windent) Write(data []byte) (int, error) {
	return w.w.Write(append(w.hdr, data...))
}

func (w *windent) Flush() error {
	ww, ok := w.w.(flusher)
	if !ok {
		return nil
	}
	return ww.Flush()
}

type flusher interface {
	Flush() error
}

type brancher interface {
	Branches() []rootio.Branch
}

func displayBranches(w io.Writer, bres brancher, indent int) {
	branches := bres.Branches()
	if len(branches) <= 0 {
		return
	}
	ww := newWindent(indent, w)
	for _, b := range branches {
		fmt.Fprintf(ww, "%s\t%q\t%v\n", b.Name(), b.Title(), b.Class())
		displayBranches(ww, b, 2)
	}
	ww.Flush()
}
<commit_msg>rootio/cmd/root-ls: improve documentation<commit_after>// Copyright 2017 The go-hep Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// root-ls dumps the content of a ROOT file.
//
// Usage: root-ls [options] file1.root [file2.root [...]]
//
// ex:
//
//  $> root-ls -t ./testdata/graphs.root ./testdata/small-flat-tree.root
//  === [./testdata/graphs.root] ===
//  version: 60806
//  TGraph            tg      graph without errors         (cycle=1)
//  TGraphErrors      tge     graph with errors            (cycle=1)
//  TGraphAsymmErrors tgae    graph with asymmetric errors (cycle=1)
//
//  === [./testdata/small-flat-tree.root] ===
//  version: 60804
//  TTree          tree                 my tree title (entries=100)
//    Int32        "Int32/I"            TBranch
//    Int64        "Int64/L"            TBranch
//    UInt32       "UInt32/i"           TBranch
//    UInt64       "UInt64/l"           TBranch
//    Float32      "Float32/F"          TBranch
//    Float64      "Float64/D"          TBranch
//    ArrayInt32   "ArrayInt32[10]/I"   TBranch
//    ArrayInt64   "ArrayInt64[10]/L"   TBranch
//    ArrayUInt32  "ArrayInt32[10]/i"   TBranch
//    ArrayUInt64  "ArrayInt64[10]/l"   TBranch
//    ArrayFloat32 "ArrayFloat32[10]/F" TBranch
//    ArrayFloat64 "ArrayFloat64[10]/D" TBranch
//    N            "N/I"                TBranch
//    SliceInt32   "SliceInt32[N]/I"    TBranch
//    SliceInt64   "SliceInt64[N]/L"    TBranch
//    SliceUInt32  "SliceInt32[N]/i"    TBranch
//    SliceUInt64  "SliceInt64[N]/l"    TBranch
//    SliceFloat32 "SliceFloat32[N]/F"  TBranch
//    SliceFloat64 "SliceFloat64[N]/D"  TBranch
//
package main // import "go-hep.org/x/hep/rootio/cmd/root-ls"

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"text/tabwriter"

	"go-hep.org/x/hep/rootio"
)

var (
	doProf = flag.String("profile", "", "filename of cpuprofile")
	doSI   = flag.Bool("sinfos", false, "print StreamerInfos")
	doTree = flag.Bool("t", false, "print Tree(s) (recursively)")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(
			os.Stderr,
			`Usage: root-ls [options] file1.root [file2.root [...]]

ex:
 $> root-ls ./testdata/graphs.root

options:
`,
		)
		flag.PrintDefaults()
	}

	flag.Parse()

	if *doProf != "" {
		f, err := os.Create(*doProf)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if flag.NArg() <= 0 {
		fmt.Fprintf(os.Stderr, "error: you need to give a ROOT file\n\n")
		flag.Usage()
		os.Exit(1)
	}

	for ii, fname := range flag.Args() {

		if ii > 0 {
			fmt.Printf("\n")
		}

		fmt.Printf("=== [%s] ===\n", fname)
		f, err := rootio.Open(fname)
		if err != nil {
			fmt.Fprintf(os.Stderr, "rootio: failed to open [%s]: %v\n", fname, err)
			os.Exit(1)
		}
		fmt.Printf("version: %v\n", f.Version())
		if *doSI {
			fmt.Printf("streamer-infos:\n")
			sinfos := f.StreamerInfo()
			for _, v := range sinfos {
				name := v.Name()
				fmt.Printf(" StreamerInfo for %q version=%d title=%q\n", name, v.ClassVersion(), v.Title())
				w := tabwriter.NewWriter(os.Stdout, 8, 4, 1, ' ', 0)
				for _, elm := range v.Elements() {
					fmt.Fprintf(w, "  %s\t%s\toffset=%3d\ttype=%3d\tsize=%3d\t %s\n", elm.TypeName(), elm.Name(), elm.Offset(), elm.Type(), elm.Size(), elm.Title())
				}
				w.Flush()
			}
		}

		w := tabwriter.NewWriter(os.Stdout, 8, 4, 1, ' ', 0)
		for _, k := range f.Keys() {
			walk(w, k)
		}
		w.Flush()
	}
}

func walk(w io.Writer, k rootio.Key) {
	obj := k.Value()
	if *doTree {
		tree, ok := obj.(rootio.Tree)
		if ok {
			w := tabwriter.NewWriter(w, 8, 4, 1, ' ', 0)
			fmt.Fprintf(w, "%s\t%s\t%s\t(entries=%d)\n", k.Class(), k.Name(), k.Title(), tree.Entries())
			displayBranches(w, tree, 2)
			w.Flush()
			return
		}
	}
	fmt.Fprintf(w, "%s\t%s\t%s\t(cycle=%d)\n", k.Class(), k.Name(), k.Title(), k.Cycle())
	if dir, ok := obj.(rootio.Directory); ok {
		w := newWindent(2, w)
		for _, k := range dir.Keys() {
			walk(w, k)
		}
		w.Flush()
	}
}

type windent struct {
	hdr []byte
	w   io.Writer
}

func newWindent(n int, w io.Writer) *windent {
	return &windent{
		hdr: bytes.Repeat([]byte(" "), n),
		w:   w,
	}
}

func (w *windent) Write(data []byte) (int, error) {
	return w.w.Write(append(w.hdr, data...))
}

func (w *windent) Flush() error {
	ww, ok := w.w.(flusher)
	if !ok {
		return nil
	}
	return ww.Flush()
}

type flusher interface {
	Flush() error
}

type brancher interface {
	Branches() []rootio.Branch
}

func displayBranches(w io.Writer, bres brancher, indent int) {
	branches := bres.Branches()
	if len(branches) <= 0 {
		return
	}
	ww := newWindent(indent, w)
	for _, b := range branches {
		fmt.Fprintf(ww, "%s\t%q\t%v\n", b.Name(), b.Title(), b.Class())
		displayBranches(ww, b, 2)
	}
	ww.Flush()
}
<|endoftext|>


<commit_before>package lane

import (
	"fmt"
	"sync"
)

// PQType represents a priority queue ordering kind (see MAXPQ and MINPQ)
type PQType int

const (
	MAXPQ PQType = iota
	MINPQ
)

type item struct {
	value    interface{}
	priority int
}

// PQueue is a heap priority queue data structure implementation.
// It can be whether max or min ordered and it is synchronized
// and is safe for concurrent operations.
type PQueue struct {
	sync.RWMutex
	items      []*item
	elemsCount int
	comparator func(int, int) bool
}

func newItem(value interface{}, priority int) *item {
	return &item{
		value:    value,
		priority: priority,
	}
}

func (i *item) String() string {
	return fmt.Sprintf("<item value:%s priority:%d>", i.value, i.priority)
}

// NewPQueue creates a new priority queue with the provided pqtype
// ordering type
func NewPQueue(pqType PQType) *PQueue {
	var cmp func(int, int) bool

	if pqType == MAXPQ {
		cmp = max
	} else {
		cmp = min
	}

	items := make([]*item, 1)
	items[0] = nil // Heap queue first element should always be nil

	return &PQueue{
		items:      items,
		elemsCount: 0,
		comparator: cmp,
	}
}

// Push the value item into the priority queue with provided priority.
func (pq *PQueue) Push(value interface{}, priority int) {
	item := newItem(value, priority)

	pq.Lock()
	pq.items = append(pq.items, item)
	pq.elemsCount += 1
	pq.swim(pq.size())
	pq.Unlock()
}

// Pop and returns the highest/lowest priority item (depending on whether
// you're using a MINPQ or MAXPQ) from the priority queue
func (pq *PQueue) Pop() (interface{}, int) {
	pq.Lock()
	defer pq.Unlock()

	if pq.size() < 1 {
		return nil, 0
	}

	var max *item = pq.items[1]

	pq.exch(1, pq.size())
	pq.items = pq.items[0:pq.size()]
	pq.elemsCount -= 1
	pq.sink(1)

	return max.value, max.priority
}

// Head returns the highest/lowest priority item (depending on whether
// you're using a MINPQ or MAXPQ) from the priority queue
func (pq *PQueue) Head() (interface{}, int) {
	pq.RLock()
	defer pq.RUnlock()

	if pq.size() < 1 {
		return nil, 0
	}

	headValue := pq.items[1].value
	headPriority := pq.items[1].priority

	return headValue, headPriority
}

// Size returns the elements present in the priority queue count
func (pq *PQueue) Size() int {
	pq.RLock()
	defer pq.RUnlock()
	return pq.size()
}

func (pq *PQueue) size() int{
	return pq.elemsCount
}

func max(i, j int) bool {
	return i < j
}

func min(i, j int) bool {
	return i > j
}

func (pq *PQueue) less(i, j int) bool {
	return pq.comparator(pq.items[i].priority, pq.items[j].priority)
}

func (pq *PQueue) exch(i, j int) {
	var tmpItem *item = pq.items[i]

	pq.items[i] = pq.items[j]
	pq.items[j] = tmpItem
}

func (pq *PQueue) swim(k int) {
	for k > 1 && pq.less(k/2, k) {
		pq.exch(k/2, k)
	}

	k = k / 2
}

func (pq *PQueue) sink(k int) {
	for 2*k <= pq.size() {
		var j int = 2 * k

		if j < pq.size() && pq.less(j, j+1) {
			j++
		}

		if !pq.less(k, j) {
			break
		}

		pq.exch(k, j)
		k = j
	}
}
<commit_msg>Update #7 #8 : fix pqueue swim method iterator<commit_after>package lane

import (
	"fmt"
	"sync"
)

// PQType represents a priority queue ordering kind (see MAXPQ and MINPQ)
type PQType int

const (
	MAXPQ PQType = iota
	MINPQ
)

type item struct {
	value    interface{}
	priority int
}

// PQueue is a heap priority queue data structure implementation.
// It can be whether max or min ordered and it is synchronized
// and is safe for concurrent operations.
type PQueue struct {
	sync.RWMutex
	items      []*item
	elemsCount int
	comparator func(int, int) bool
}

func newItem(value interface{}, priority int) *item {
	return &item{
		value:    value,
		priority: priority,
	}
}

func (i *item) String() string {
	return fmt.Sprintf("<item value:%s priority:%d>", i.value, i.priority)
}

// NewPQueue creates a new priority queue with the provided pqtype
// ordering type
func NewPQueue(pqType PQType) *PQueue {
	var cmp func(int, int) bool

	if pqType == MAXPQ {
		cmp = max
	} else {
		cmp = min
	}

	items := make([]*item, 1)
	items[0] = nil // Heap queue first element should always be nil

	return &PQueue{
		items:      items,
		elemsCount: 0,
		comparator: cmp,
	}
}

// Push the value item into the priority queue with provided priority.
func (pq *PQueue) Push(value interface{}, priority int) {
	item := newItem(value, priority)

	pq.Lock()
	pq.items = append(pq.items, item)
	pq.elemsCount += 1
	pq.swim(pq.size())
	pq.Unlock()
}

// Pop and returns the highest/lowest priority item (depending on whether
// you're using a MINPQ or MAXPQ) from the priority queue
func (pq *PQueue) Pop() (interface{}, int) {
	pq.Lock()
	defer pq.Unlock()

	if pq.size() < 1 {
		return nil, 0
	}

	var max *item = pq.items[1]

	pq.exch(1, pq.size())
	pq.items = pq.items[0:pq.size()]
	pq.elemsCount -= 1
	pq.sink(1)

	return max.value, max.priority
}

// Head returns the highest/lowest priority item (depending on whether
// you're using a MINPQ or MAXPQ) from the priority queue
func (pq *PQueue) Head() (interface{}, int) {
	pq.RLock()
	defer pq.RUnlock()

	if pq.size() < 1 {
		return nil, 0
	}

	headValue := pq.items[1].value
	headPriority := pq.items[1].priority

	return headValue, headPriority
}

// Size returns the elements present in the priority queue count
func (pq *PQueue) Size() int {
	pq.RLock()
	defer pq.RUnlock()
	return pq.size()
}

func (pq *PQueue) size() int {
	return pq.elemsCount
}

func max(i, j int) bool {
	return i < j
}

func min(i, j int) bool {
	return i > j
}

func (pq *PQueue) less(i, j int) bool {
	return pq.comparator(pq.items[i].priority, pq.items[j].priority)
}

func (pq *PQueue) exch(i, j int) {
	var tmpItem *item = pq.items[i]

	pq.items[i] = pq.items[j]
	pq.items[j] = tmpItem
}

func (pq *PQueue) swim(k int) {
	for k > 1 && pq.less(k/2, k) {
		pq.exch(k/2, k)
		k = k / 2
	}

}

func (pq *PQueue) sink(k int) {
	for 2*k <= pq.size() {
		var j int = 2 * k

		if j < pq.size() && pq.less(j, j+1) {
			j++
		}

		if !pq.less(k, j) {
			break
		}

		pq.exch(k, j)
		k = j
	}
}
<|endoftext|>


<commit_before>// Package safecookie provides secure encoding and decoding for cookies which
// provides both confidentiality and authenticity against both active and
// passive attackers.
//
// It does so by sealing a cookie's value with an AEAD using the cookie's name
// as the authenticated data. If the cookie's name or value change at all, the
// opening process will fail.
//
// This provides some important guarantees:
//
//     - No one who does not have the secret key can read the cookie's plaintext
//       value.
//
//     - No one who does not have the secret key can create a cookie which will
//       be considered valid.
//
//     - Any cookie which has had its name or value changed will be considered
//       invalid.
//
// Cookie values are gobs, so be sure to register the types via gob.Register.
package safecookie

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"errors"
	"net/http"
)

var (
	// ErrInvalidCookie is returned if the cookie is invalid.
	ErrInvalidCookie = errors.New("invalid cookie")
)

// SafeCookie seals cookies and opens them.
type SafeCookie struct {
	// AEAD is the Authenticated Encryption And Data algorithm to use for
	// encrypting and decrypting cookie values.
	AEAD cipher.AEAD
}

// NewGCM returns a new AES-GCM-based SafeCookie instance given a 128-, 192-, or
// 256-bit key.
func NewGCM(key []byte) (*SafeCookie, error) {
	b, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(b)
	if err != nil {
		return nil, err
	}

	return &SafeCookie{AEAD: gcm}, nil
}

// Seal marshals the given value into a gob, encrypts that using the cookie's
// name as authenticated data, and sets the cookie's value to the Base64-encoded
// ciphertext.
func (sc *SafeCookie) Seal(e interface{}, c *http.Cookie) error {
	w := bytes.NewBuffer(nil)
	if err := gob.NewEncoder(w).Encode(e); err != nil {
		return err
	}

	nonce := make([]byte, sc.AEAD.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return err
	}

	ciphertext := sc.AEAD.Seal(nonce, nonce, w.Bytes(), []byte(c.Name))
	c.Value = base64.URLEncoding.EncodeToString(ciphertext)

	return nil
}

// Open decodes the cookie's value as Base64, decrypts it (authenticating the
// cookie name), and unmarshals the resulting gob into the given pointer. If the
// cookie is invalid, it returns ErrInvalidCookie.
func (sc *SafeCookie) Open(c *http.Cookie, e interface{}) error {
	b, err := base64.URLEncoding.DecodeString(c.Value)
	if err != nil || len(b) <= sc.AEAD.NonceSize() {
		return ErrInvalidCookie
	}

	nonce := b[:sc.AEAD.NonceSize()]
	ciphertext := b[sc.AEAD.NonceSize():]

	b, err = sc.AEAD.Open(nil, nonce, ciphertext, []byte(c.Name))
	if err != nil {
		return ErrInvalidCookie
	}

	return gob.NewDecoder(bytes.NewReader(b)).Decode(e)
}
<commit_msg>Docfix.<commit_after>// Package safecookie provides secure encoding and decoding for cookies which
// provides both confidentiality and authenticity against both active and
// passive attackers.
//
// It does so by sealing a cookie's value with an AEAD using the cookie's name
// as the authenticated data. If the cookie's name or value change at all, the
// opening process will fail.
//
// This provides some important guarantees:
//
//     - No one who does not have the secret key can read the cookie's plaintext
//       value.
//     - No one who does not have the secret key can create a cookie which will
//       be considered valid.
//     - Any cookie which has had its name or value changed will be considered
//       invalid.
//
// Cookie values are gobs, so be sure to register the types via gob.Register.
package safecookie

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"errors"
	"net/http"
)

var (
	// ErrInvalidCookie is returned if the cookie is invalid.
	ErrInvalidCookie = errors.New("invalid cookie")
)

// SafeCookie seals cookies and opens them.
type SafeCookie struct {
	// AEAD is the Authenticated Encryption And Data algorithm to use for
	// encrypting and decrypting cookie values.
	AEAD cipher.AEAD
}

// NewGCM returns a new AES-GCM-based SafeCookie instance given a 128-, 192-, or
// 256-bit key.
func NewGCM(key []byte) (*SafeCookie, error) {
	b, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(b)
	if err != nil {
		return nil, err
	}

	return &SafeCookie{AEAD: gcm}, nil
}

// Seal marshals the given value into a gob, encrypts that using the cookie's
// name as authenticated data, and sets the cookie's value to the Base64-encoded
// ciphertext.
func (sc *SafeCookie) Seal(e interface{}, c *http.Cookie) error {
	w := bytes.NewBuffer(nil)
	if err := gob.NewEncoder(w).Encode(e); err != nil {
		return err
	}

	nonce := make([]byte, sc.AEAD.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return err
	}

	ciphertext := sc.AEAD.Seal(nonce, nonce, w.Bytes(), []byte(c.Name))
	c.Value = base64.URLEncoding.EncodeToString(ciphertext)

	return nil
}

// Open decodes the cookie's value as Base64, decrypts it (authenticating the
// cookie name), and unmarshals the resulting gob into the given pointer. If the
// cookie is invalid, it returns ErrInvalidCookie.
func (sc *SafeCookie) Open(c *http.Cookie, e interface{}) error {
	b, err := base64.URLEncoding.DecodeString(c.Value)
	if err != nil || len(b) <= sc.AEAD.NonceSize() {
		return ErrInvalidCookie
	}

	nonce := b[:sc.AEAD.NonceSize()]
	ciphertext := b[sc.AEAD.NonceSize():]

	b, err = sc.AEAD.Open(nil, nonce, ciphertext, []byte(c.Name))
	if err != nil {
		return ErrInvalidCookie
	}

	return gob.NewDecoder(bytes.NewReader(b)).Decode(e)
}
<|endoftext|>