<commit_before>package main

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

func TestKillContainer(t *testing.T) {
	runCmd := exec.Command(dockerBinary, "run", "-d", "busybox", "sh", "-c", "sleep 10")
	out, _, err := runCommandWithOutput(runCmd)
	errorOut(err, t, fmt.Sprintf("run failed with errors: %v", err))

	cleanedContainerID := stripTrailingCharacters(out)

	inspectCmd := exec.Command(dockerBinary, "inspect", cleanedContainerID)
	inspectOut, _, err := runCommandWithOutput(inspectCmd)
	errorOut(err, t, fmt.Sprintf("out should've been a container id: %v %v", inspectOut, err))

	killCmd := exec.Command(dockerBinary, "kill", cleanedContainerID)
	out, _, err = runCommandWithOutput(killCmd)
	errorOut(err, t, fmt.Sprintf("failed to kill container: %v %v", out, err))

	listRunningContainersCmd := exec.Command(dockerBinary, "ps", "-q")
	out, _, err = runCommandWithOutput(listRunningContainersCmd)
	errorOut(err, t, fmt.Sprintf("failed to list running containers: %v", err))

	if strings.Contains(out, cleanedContainerID) {
		t.Fatal("killed container is still running")
	}

	deleteContainer(cleanedContainerID)

	logDone("kill - kill container running sleep 10")
}

func TestKillDifferentUserContainer(t *testing.T) {
	runCmd := exec.Command(dockerBinary, "run", "-u", "daemon", "-d", "busybox", "sh", "-c", "sleep 10")
	out, _, err := runCommandWithOutput(runCmd)
	errorOut(err, t, fmt.Sprintf("run failed with errors: %v", err))

	cleanedContainerID := stripTrailingCharacters(out)

	inspectCmd := exec.Command(dockerBinary, "inspect", cleanedContainerID)
	inspectOut, _, err := runCommandWithOutput(inspectCmd)
	errorOut(err, t, fmt.Sprintf("out should've been a container id: %v %v", inspectOut, err))

	killCmd := exec.Command(dockerBinary, "kill", cleanedContainerID)
	out, _, err = runCommandWithOutput(killCmd)
	errorOut(err, t, fmt.Sprintf("failed to kill container: %v %v", out, err))

	listRunningContainersCmd := exec.Command(dockerBinary, "ps", "-q")
	out, _, err = runCommandWithOutput(listRunningContainersCmd)
	errorOut(err, t, fmt.Sprintf("failed to list running containers: %v", err))

	if strings.Contains(out, cleanedContainerID) {
		t.Fatal("killed container is still running")
	}

	deleteContainer(cleanedContainerID)

	logDone("kill - kill container running sleep 10 from a different user")
}
<commit_msg>Cleanup errorOut resp in kill test<commit_after>package main

import (
	"os/exec"
	"strings"
	"testing"
)

func TestKillContainer(t *testing.T) {
	runCmd := exec.Command(dockerBinary, "run", "-d", "busybox", "sh", "-c", "sleep 10")
	out, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	cleanedContainerID := stripTrailingCharacters(out)

	inspectCmd := exec.Command(dockerBinary, "inspect", cleanedContainerID)
	if out, _, err = runCommandWithOutput(inspectCmd); err != nil {
		t.Fatalf("out should've been a container id: %s, %v", out, err)
	}

	killCmd := exec.Command(dockerBinary, "kill", cleanedContainerID)
	if out, _, err = runCommandWithOutput(killCmd); err != nil {
		t.Fatalf("failed to kill container: %s, %v", out, err)
	}

	listRunningContainersCmd := exec.Command(dockerBinary, "ps", "-q")
	out, _, err = runCommandWithOutput(listRunningContainersCmd)
	if err != nil {
		t.Fatalf("failed to list running containers: %s, %v", out, err)
	}

	if strings.Contains(out, cleanedContainerID) {
		t.Fatal("killed container is still running")
	}

	deleteContainer(cleanedContainerID)

	logDone("kill - kill container running sleep 10")
}

func TestKillDifferentUserContainer(t *testing.T) {
	runCmd := exec.Command(dockerBinary, "run", "-u", "daemon", "-d", "busybox", "sh", "-c", "sleep 10")
	out, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	cleanedContainerID := stripTrailingCharacters(out)

	inspectCmd := exec.Command(dockerBinary, "inspect", cleanedContainerID)
	if out, _, err = runCommandWithOutput(inspectCmd); err != nil {
		t.Fatalf("out should've been a container id: %s, %v", out, err)
	}

	killCmd := exec.Command(dockerBinary, "kill", cleanedContainerID)
	if out, _, err = runCommandWithOutput(killCmd); err != nil {
		t.Fatalf("failed to kill container: %s, %v", out, err)
	}

	listRunningContainersCmd := exec.Command(dockerBinary, "ps", "-q")
	out, _, err = runCommandWithOutput(listRunningContainersCmd)
	if err != nil {
		t.Fatalf("failed to list running containers: %s, %v", out, err)
	}

	if strings.Contains(out, cleanedContainerID) {
		t.Fatal("killed container is still running")
	}

	deleteContainer(cleanedContainerID)

	logDone("kill - kill container running sleep 10 from a different user")
}
<|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.

// Package federationout handles requests from other federation servers for data.
package federationout

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/google/exposure-notifications-server/internal/federationin/model"
	"github.com/google/exposure-notifications-server/internal/logging"
	"github.com/google/exposure-notifications-server/internal/pb"

	coredb "github.com/google/exposure-notifications-server/internal/database"
	"github.com/google/exposure-notifications-server/internal/federationout/database"
	publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"

	"github.com/google/exposure-notifications-server/internal/serverenv"
	"google.golang.org/api/idtoken"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

const (
	authHeader = "authorization"
	bearer     = "Bearer"
)

// Compile time assert that this server implements the required grpc interface.
var _ pb.FederationServer = (*Server)(nil)

type iterateExposuresFunc func(context.Context, publishdb.IterateExposuresCriteria, func(*publishmodel.Exposure) error) (string, error)

// NewServer builds a new FederationServer.
func NewServer(env *serverenv.ServerEnv, config *Config) pb.FederationServer {
	return &Server{
		env:       env,
		db:        database.New(env.Database()),
		publishdb: publishdb.New(env.Database()),
		config:    config,
	}
}

type Server struct {
	env       *serverenv.ServerEnv
	db        *database.FederationOutDB
	publishdb *publishdb.PublishDB
	config    *Config
}

type authKey struct{}

// Fetch implements the FederationServer Fetch endpoint.
func (s Server) Fetch(ctx context.Context, req *pb.FederationFetchRequest) (*pb.FederationFetchResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, s.config.Timeout)
	defer cancel()
	logger := logging.FromContext(ctx)
	response, err := s.fetch(ctx, req, s.publishdb.IterateExposures, publishmodel.TruncateWindow(time.Now(), s.config.TruncateWindow)) // Don't fetch the current window, which isn't complete yet. TODO(squee1945): should I double this for safety?
	if err != nil {
		s.env.MetricsExporter(ctx).WriteInt("federation-fetch-failed", true, 1)
		logger.Errorf("Fetch error: %v", err)
		return nil, errors.New("internal error")
	}
	return response, nil
}

func (s Server) fetch(ctx context.Context, req *pb.FederationFetchRequest, itFunc iterateExposuresFunc, fetchUntil time.Time) (*pb.FederationFetchResponse, error) {
	logger := logging.FromContext(ctx)
	metrics := s.env.MetricsExporter(ctx)

	for i := range req.RegionIdentifiers {
		req.RegionIdentifiers[i] = strings.ToUpper(req.RegionIdentifiers[i])
	}
	for i := range req.ExcludeRegionIdentifiers {
		req.ExcludeRegionIdentifiers[i] = strings.ToUpper(req.ExcludeRegionIdentifiers[i])
	}
	metrics.WriteInt("federation-fetch-regions-requested", false, len(req.RegionIdentifiers))
	metrics.WriteInt("federation-fetch-regions-excluded", false, len(req.ExcludeRegionIdentifiers))

	logger.Infof("Processing client request %#v", req)

	// If there is a FederationAuthorization on the context, set the query to operate within its limits.
	if auth, ok := ctx.Value(authKey{}).(*model.FederationOutAuthorization); ok {
		// For included regions, we INTERSECT the requested included regions with the configured included regions.
		req.RegionIdentifiers = intersect(req.RegionIdentifiers, auth.IncludeRegions)
		// For excluded regions, we UNION the the requested excluded regions with the configured excluded regions.
		req.ExcludeRegionIdentifiers = union(req.ExcludeRegionIdentifiers, auth.ExcludeRegions)
	}

	criteria := publishdb.IterateExposuresCriteria{
		IncludeRegions:      req.RegionIdentifiers,
		ExcludeRegions:      req.ExcludeRegionIdentifiers,
		SinceTimestamp:      time.Unix(req.LastFetchResponseKeyTimestamp, 0),
		UntilTimestamp:      fetchUntil,
		LastCursor:          req.NextFetchToken,
		OnlyLocalProvenance: true, // Do not return results that came from other federation partners.
	}

	logger.Infof("Query criteria: %#v", criteria)

	// Filter included countries in memory.
	includedRegions := map[string]struct{}{}
	for _, region := range req.RegionIdentifiers {
		includedRegions[region] = struct{}{}
	}

	// Filter excluded countries in memory, using a map for efficiency.
	excludedRegions := map[string]struct{}{}
	for _, region := range req.ExcludeRegionIdentifiers {
		excludedRegions[region] = struct{}{}
	}

	ctrMap := map[string]*pb.ContactTracingResponse{} // local index into the response being assembled; keyed on unique set of regions.
	ctiMap := map[string]*pb.ContactTracingInfo{}     // local index into the response being assembled; keys on unique set of (ctrMap key, transmissionRisk, verificationAuthorityName)
	response := &pb.FederationFetchResponse{}
	count := 0
	cursor, err := itFunc(ctx, criteria, func(inf *publishmodel.Exposure) error {
		// If the diagnosis key is empty, it's malformed, so skip it.
		if len(inf.ExposureKey) == 0 {
			logger.Debugf("Exposure %s missing ExposureKey, skipping.", inf.ExposureKey)
			return nil
		}

		// If there are no regions on the exposure, it's malformed, so skip it.
		if len(inf.Regions) == 0 {
			logger.Debugf("Exposure %s missing Regions, skipping.", inf.ExposureKey)
			return nil
		}

		// Filter out non-LocalProvenance results; we should not re-federate.
		// This may already be handled by the database query and is included here for completeness.
		if !inf.LocalProvenance {
			logger.Debugf("Exposure %s not LocalProvenance, skipping.", inf.ExposureKey)
			return nil
		}

		// If all the regions on the record are excluded, skip it.
		skip := true
		for _, region := range inf.Regions {
			if _, excluded := excludedRegions[region]; !excluded {
				// At least one region for the exposure is NOT excluded, so we don't skip this record.
				skip = false
				break
			}
		}
		if skip {
			logger.Debugf("Exposure %s contains only excluded regions, skipping.", inf.ExposureKey)
			return nil
		}

		// If filtering on a region (len(includedRegions) > 0) and none of the regions on the record are included, skip it.
		if len(includedRegions) > 0 {
			skip = true
			for _, region := range inf.Regions {
				if _, included := includedRegions[region]; included {
					skip = false
					break
				}
			}
			if skip {
				logger.Debugf("Exposure %s does not contain requested regions, skipping.", inf.ExposureKey)
				return nil
			}
		}

		// Find, or create, the ContactTracingResponse based on the unique set of regions.
		sort.Strings(inf.Regions)
		ctrKey := strings.Join(inf.Regions, "::")
		ctr := ctrMap[ctrKey]
		if ctr == nil {
			ctr = &pb.ContactTracingResponse{RegionIdentifiers: inf.Regions}
			ctrMap[ctrKey] = ctr
			response.Response = append(response.Response, ctr)
		}

		// Find, or create, the ContactTracingInfo for (ctrKey, transmissionRisk).
		ctiKey := fmt.Sprintf("%s::%d", ctrKey, inf.TransmissionRisk)
		cti := ctiMap[ctiKey]
		if cti == nil {
			cti = &pb.ContactTracingInfo{TransmissionRisk: int32(inf.TransmissionRisk)}
			ctiMap[ctiKey] = cti
			ctr.ContactTracingInfo = append(ctr.ContactTracingInfo, cti)
		}

		// Add the key to the ContactTracingInfo.
		cti.ExposureKeys = append(cti.ExposureKeys, &pb.ExposureKey{
			ExposureKey:    inf.ExposureKey,
			IntervalNumber: inf.IntervalNumber,
			IntervalCount:  inf.IntervalCount,
		})

		created := inf.CreatedAt.Unix()
		if created > response.FetchResponseKeyTimestamp {
			response.FetchResponseKeyTimestamp = created
		}

		count++
		return nil
	})
	if err != nil {
		metrics.WriteInt("federation-fetch-error", true, 1)
		if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
			logger.Infof("Fetch request reached time out, returning partial response.")
			response.PartialResponse = true
			response.NextFetchToken = cursor
		} else {
			return nil, err
		}
	}
	metrics.WriteInt("federation-fetch-count", false, count)
	logger.Infof("Sent %d keys", count)
	return response, nil
}

// AuthInterceptor validates incoming OIDC bearer token and adds corresponding FederationAuthorization record to the context.
func (s Server) AuthInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	logger := logging.FromContext(ctx)
	metrics := s.env.MetricsExporter(ctx)

	raw, err := rawToken(ctx)
	if err != nil {
		logger.Infof("Invalid headers: %v", err)
		return nil, err
	}

	token, err := idtoken.Validate(ctx, raw, "")
	if err != nil {
		logger.Infof("Invalid token: %v", err)
		metrics.WriteInt("federation-fetch-invalid-auth-token", true, 1)
		return nil, status.Errorf(codes.Unauthenticated, "Invalid token")
	}

	auth, err := s.db.GetFederationOutAuthorization(ctx, token.Issuer, token.Subject)
	if err != nil {
		if errors.Is(err, coredb.ErrNotFound) {
			metrics.WriteInt("federation-fetch-unauthorized", true, 1)
			logger.Infof("Authorization not found (issuer %q, subject %s)", token.Issuer, token.Subject)
			return nil, status.Errorf(codes.Unauthenticated, "Invalid issuer/subject")
		}
		logger.Errorf("Failed to fetch authorization (issuer %q, subject %s): %v", token.Issuer, token.Subject, err)
		metrics.WriteInt("federation-fetch-internal-error", true, 1)
		return nil, status.Errorf(codes.Internal, "Internal error")
	}

	if auth.Audience != "" && auth.Audience != token.Audience {
		metrics.WriteInt("federation-fetch-invalid-audience", true, 1)
		logger.Infof("Invalid audience, got %q, want %q", token.Audience, auth.Audience)
		return nil, status.Errorf(codes.Unauthenticated, "Invalid audience")
	}

	// Store the FederationAuthorization on the context.
	logger.Infof("Caller: issuer %q subject %q", auth.Issuer, auth.Subject)
	ctx = context.WithValue(ctx, authKey{}, auth)
	return handler(ctx, req)
}

func rawToken(ctx context.Context) (string, error) {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return "", status.Errorf(codes.Unauthenticated, "Missing metadata")
	}
	if _, ok := md[authHeader]; !ok {
		return "", status.Errorf(codes.Unauthenticated, "Missing authorization header [1]")
	}
	if len(md[authHeader]) == 0 {
		return "", status.Errorf(codes.Unauthenticated, "Missing authorization header [2]")
	}
	if len(md[authHeader]) > 1 {
		return "", status.Errorf(codes.Unauthenticated, "Multiple authorization headers")
	}

	authHeader := md[authHeader][0]
	if !strings.HasPrefix(authHeader, bearer) {
		return "", status.Errorf(codes.Unauthenticated, "Invalid authorization header")
	}
	rawToken := strings.TrimSpace(strings.TrimPrefix(authHeader, bearer))
	return rawToken, nil
}

func intersect(aa, bb []string) []string {
	if len(aa) == 0 || len(bb) == 0 {
		return nil
	}
	var result []string
	for _, a := range aa {
		for _, b := range bb {
			if a == b {
				result = append(result, a)
				break
			}
		}
	}
	return result
}

func union(aa, bb []string) []string {
	if len(aa) == 0 {
		return bb
	}
	if len(bb) == 0 {
		return aa
	}
	m := map[string]struct{}{}
	for _, a := range aa {
		m[a] = struct{}{}
	}
	for _, b := range bb {
		m[b] = struct{}{}
	}
	var result []string
	for k := range m {
		result = append(result, k)
	}
	sort.Strings(result)
	return result
}
<commit_msg>fixed map allocation (#418)<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.

// Package federationout handles requests from other federation servers for data.
package federationout

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/google/exposure-notifications-server/internal/federationin/model"
	"github.com/google/exposure-notifications-server/internal/logging"
	"github.com/google/exposure-notifications-server/internal/pb"

	coredb "github.com/google/exposure-notifications-server/internal/database"
	"github.com/google/exposure-notifications-server/internal/federationout/database"
	publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"

	"github.com/google/exposure-notifications-server/internal/serverenv"
	"google.golang.org/api/idtoken"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

const (
	authHeader = "authorization"
	bearer     = "Bearer"
)

// Compile time assert that this server implements the required grpc interface.
var _ pb.FederationServer = (*Server)(nil)

type iterateExposuresFunc func(context.Context, publishdb.IterateExposuresCriteria, func(*publishmodel.Exposure) error) (string, error)

// NewServer builds a new FederationServer.
func NewServer(env *serverenv.ServerEnv, config *Config) pb.FederationServer {
	return &Server{
		env:       env,
		db:        database.New(env.Database()),
		publishdb: publishdb.New(env.Database()),
		config:    config,
	}
}

type Server struct {
	env       *serverenv.ServerEnv
	db        *database.FederationOutDB
	publishdb *publishdb.PublishDB
	config    *Config
}

type authKey struct{}

// Fetch implements the FederationServer Fetch endpoint.
func (s Server) Fetch(ctx context.Context, req *pb.FederationFetchRequest) (*pb.FederationFetchResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, s.config.Timeout)
	defer cancel()
	logger := logging.FromContext(ctx)
	response, err := s.fetch(ctx, req, s.publishdb.IterateExposures, publishmodel.TruncateWindow(time.Now(), s.config.TruncateWindow)) // Don't fetch the current window, which isn't complete yet. TODO(squee1945): should I double this for safety?
	if err != nil {
		s.env.MetricsExporter(ctx).WriteInt("federation-fetch-failed", true, 1)
		logger.Errorf("Fetch error: %v", err)
		return nil, errors.New("internal error")
	}
	return response, nil
}

func (s Server) fetch(ctx context.Context, req *pb.FederationFetchRequest, itFunc iterateExposuresFunc, fetchUntil time.Time) (*pb.FederationFetchResponse, error) {
	logger := logging.FromContext(ctx)
	metrics := s.env.MetricsExporter(ctx)

	for i := range req.RegionIdentifiers {
		req.RegionIdentifiers[i] = strings.ToUpper(req.RegionIdentifiers[i])
	}
	for i := range req.ExcludeRegionIdentifiers {
		req.ExcludeRegionIdentifiers[i] = strings.ToUpper(req.ExcludeRegionIdentifiers[i])
	}
	metrics.WriteInt("federation-fetch-regions-requested", false, len(req.RegionIdentifiers))
	metrics.WriteInt("federation-fetch-regions-excluded", false, len(req.ExcludeRegionIdentifiers))

	logger.Infof("Processing client request %#v", req)

	// If there is a FederationAuthorization on the context, set the query to operate within its limits.
	if auth, ok := ctx.Value(authKey{}).(*model.FederationOutAuthorization); ok {
		// For included regions, we INTERSECT the requested included regions with the configured included regions.
		req.RegionIdentifiers = intersect(req.RegionIdentifiers, auth.IncludeRegions)
		// For excluded regions, we UNION the the requested excluded regions with the configured excluded regions.
		req.ExcludeRegionIdentifiers = union(req.ExcludeRegionIdentifiers, auth.ExcludeRegions)
	}

	criteria := publishdb.IterateExposuresCriteria{
		IncludeRegions:      req.RegionIdentifiers,
		ExcludeRegions:      req.ExcludeRegionIdentifiers,
		SinceTimestamp:      time.Unix(req.LastFetchResponseKeyTimestamp, 0),
		UntilTimestamp:      fetchUntil,
		LastCursor:          req.NextFetchToken,
		OnlyLocalProvenance: true, // Do not return results that came from other federation partners.
	}

	logger.Infof("Query criteria: %#v", criteria)

	// Filter included countries in memory.
	includedRegions := make(map[string]struct{}, len(req.RegionIdentifiers))
	for _, region := range req.RegionIdentifiers {
		includedRegions[region] = struct{}{}
	}

	// Filter excluded countries in memory, using a map for efficiency.
	excludedRegions := make(map[string]struct{}, len(req.ExcludeRegionIdentifiers))
	for _, region := range req.ExcludeRegionIdentifiers {
		excludedRegions[region] = struct{}{}
	}

	ctrMap := map[string]*pb.ContactTracingResponse{} // local index into the response being assembled; keyed on unique set of regions.
	ctiMap := map[string]*pb.ContactTracingInfo{}     // local index into the response being assembled; keys on unique set of (ctrMap key, transmissionRisk, verificationAuthorityName)
	response := &pb.FederationFetchResponse{}
	count := 0
	cursor, err := itFunc(ctx, criteria, func(inf *publishmodel.Exposure) error {
		// If the diagnosis key is empty, it's malformed, so skip it.
		if len(inf.ExposureKey) == 0 {
			logger.Debugf("Exposure %s missing ExposureKey, skipping.", inf.ExposureKey)
			return nil
		}

		// If there are no regions on the exposure, it's malformed, so skip it.
		if len(inf.Regions) == 0 {
			logger.Debugf("Exposure %s missing Regions, skipping.", inf.ExposureKey)
			return nil
		}

		// Filter out non-LocalProvenance results; we should not re-federate.
		// This may already be handled by the database query and is included here for completeness.
		if !inf.LocalProvenance {
			logger.Debugf("Exposure %s not LocalProvenance, skipping.", inf.ExposureKey)
			return nil
		}

		// If all the regions on the record are excluded, skip it.
		skip := true
		for _, region := range inf.Regions {
			if _, excluded := excludedRegions[region]; !excluded {
				// At least one region for the exposure is NOT excluded, so we don't skip this record.
				skip = false
				break
			}
		}
		if skip {
			logger.Debugf("Exposure %s contains only excluded regions, skipping.", inf.ExposureKey)
			return nil
		}

		// If filtering on a region (len(includedRegions) > 0) and none of the regions on the record are included, skip it.
		if len(includedRegions) > 0 {
			skip = true
			for _, region := range inf.Regions {
				if _, included := includedRegions[region]; included {
					skip = false
					break
				}
			}
			if skip {
				logger.Debugf("Exposure %s does not contain requested regions, skipping.", inf.ExposureKey)
				return nil
			}
		}

		// Find, or create, the ContactTracingResponse based on the unique set of regions.
		sort.Strings(inf.Regions)
		ctrKey := strings.Join(inf.Regions, "::")
		ctr := ctrMap[ctrKey]
		if ctr == nil {
			ctr = &pb.ContactTracingResponse{RegionIdentifiers: inf.Regions}
			ctrMap[ctrKey] = ctr
			response.Response = append(response.Response, ctr)
		}

		// Find, or create, the ContactTracingInfo for (ctrKey, transmissionRisk).
		ctiKey := fmt.Sprintf("%s::%d", ctrKey, inf.TransmissionRisk)
		cti := ctiMap[ctiKey]
		if cti == nil {
			cti = &pb.ContactTracingInfo{TransmissionRisk: int32(inf.TransmissionRisk)}
			ctiMap[ctiKey] = cti
			ctr.ContactTracingInfo = append(ctr.ContactTracingInfo, cti)
		}

		// Add the key to the ContactTracingInfo.
		cti.ExposureKeys = append(cti.ExposureKeys, &pb.ExposureKey{
			ExposureKey:    inf.ExposureKey,
			IntervalNumber: inf.IntervalNumber,
			IntervalCount:  inf.IntervalCount,
		})

		created := inf.CreatedAt.Unix()
		if created > response.FetchResponseKeyTimestamp {
			response.FetchResponseKeyTimestamp = created
		}

		count++
		return nil
	})
	if err != nil {
		metrics.WriteInt("federation-fetch-error", true, 1)
		if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
			logger.Infof("Fetch request reached time out, returning partial response.")
			response.PartialResponse = true
			response.NextFetchToken = cursor
		} else {
			return nil, err
		}
	}
	metrics.WriteInt("federation-fetch-count", false, count)
	logger.Infof("Sent %d keys", count)
	return response, nil
}

// AuthInterceptor validates incoming OIDC bearer token and adds corresponding FederationAuthorization record to the context.
func (s Server) AuthInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	logger := logging.FromContext(ctx)
	metrics := s.env.MetricsExporter(ctx)

	raw, err := rawToken(ctx)
	if err != nil {
		logger.Infof("Invalid headers: %v", err)
		return nil, err
	}

	token, err := idtoken.Validate(ctx, raw, "")
	if err != nil {
		logger.Infof("Invalid token: %v", err)
		metrics.WriteInt("federation-fetch-invalid-auth-token", true, 1)
		return nil, status.Errorf(codes.Unauthenticated, "Invalid token")
	}

	auth, err := s.db.GetFederationOutAuthorization(ctx, token.Issuer, token.Subject)
	if err != nil {
		if errors.Is(err, coredb.ErrNotFound) {
			metrics.WriteInt("federation-fetch-unauthorized", true, 1)
			logger.Infof("Authorization not found (issuer %q, subject %s)", token.Issuer, token.Subject)
			return nil, status.Errorf(codes.Unauthenticated, "Invalid issuer/subject")
		}
		logger.Errorf("Failed to fetch authorization (issuer %q, subject %s): %v", token.Issuer, token.Subject, err)
		metrics.WriteInt("federation-fetch-internal-error", true, 1)
		return nil, status.Errorf(codes.Internal, "Internal error")
	}

	if auth.Audience != "" && auth.Audience != token.Audience {
		metrics.WriteInt("federation-fetch-invalid-audience", true, 1)
		logger.Infof("Invalid audience, got %q, want %q", token.Audience, auth.Audience)
		return nil, status.Errorf(codes.Unauthenticated, "Invalid audience")
	}

	// Store the FederationAuthorization on the context.
	logger.Infof("Caller: issuer %q subject %q", auth.Issuer, auth.Subject)
	ctx = context.WithValue(ctx, authKey{}, auth)
	return handler(ctx, req)
}

func rawToken(ctx context.Context) (string, error) {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return "", status.Errorf(codes.Unauthenticated, "Missing metadata")
	}
	if _, ok := md[authHeader]; !ok {
		return "", status.Errorf(codes.Unauthenticated, "Missing authorization header [1]")
	}
	if len(md[authHeader]) == 0 {
		return "", status.Errorf(codes.Unauthenticated, "Missing authorization header [2]")
	}
	if len(md[authHeader]) > 1 {
		return "", status.Errorf(codes.Unauthenticated, "Multiple authorization headers")
	}

	authHeader := md[authHeader][0]
	if !strings.HasPrefix(authHeader, bearer) {
		return "", status.Errorf(codes.Unauthenticated, "Invalid authorization header")
	}
	rawToken := strings.TrimSpace(strings.TrimPrefix(authHeader, bearer))
	return rawToken, nil
}

func intersect(aa, bb []string) []string {
	if len(aa) == 0 || len(bb) == 0 {
		return nil
	}
	var result []string
	for _, a := range aa {
		for _, b := range bb {
			if a == b {
				result = append(result, a)
				break
			}
		}
	}
	return result
}

func union(aa, bb []string) []string {
	if len(aa) == 0 {
		return bb
	}
	if len(bb) == 0 {
		return aa
	}
	m := map[string]struct{}{}
	for _, a := range aa {
		m[a] = struct{}{}
	}
	for _, b := range bb {
		m[b] = struct{}{}
	}
	result := make([]string, 0, len(m))
	for k := range m {
		result = append(result, k)
	}
	sort.Strings(result)
	return result
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"

	"github.com/aws/aws-sdk-go/internal/model/api"
)

func main() {
	dir, _ := os.Open("apis")
	names, _ := dir.Readdirnames(0)
	for _, name := range names {
		m, _ := filepath.Glob(filepath.Join("apis", name, "*.normal.json"))

		sort.Strings(m)
		f := m[len(m)-1]
		a := api.API{}
		a.Attach(f)
		fmt.Printf("%s\t%s\n", a.Metadata.ServiceFullName, a.Metadata.APIVersion)
	}
}
<commit_msg>Fix empty APIs in api_info task<commit_after>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"

	"github.com/aws/aws-sdk-go/internal/model/api"
)

func main() {
	dir, _ := os.Open("apis")
	names, _ := dir.Readdirnames(0)
	for _, name := range names {
		m, _ := filepath.Glob(filepath.Join("apis", name, "*.normal.json"))
		if len(m) == 0 {
			continue
		}

		sort.Strings(m)
		f := m[len(m)-1]
		a := api.API{}
		a.Attach(f)
		fmt.Printf("%s\t%s\n", a.Metadata.ServiceFullName, a.Metadata.APIVersion)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/peterbourgon/gokit/addsvc/pb"
	"github.com/peterbourgon/gokit/metrics"
	"github.com/peterbourgon/gokit/server"
	"golang.org/x/net/context"
)

// A binding wraps an Endpoint so that it's usable by a transport. grpcBinding
// makes an Endpoint usable over gRPC.
type grpcBinding struct{ server.Endpoint }

// Add implements the proto3 AddServer by forwarding to the wrapped Endpoint.
func (b grpcBinding) Add(ctx context.Context, req *pb.AddRequest) (*pb.AddReply, error) {
	addReq := request{req.A, req.B}
	r, err := b.Endpoint(ctx, addReq)
	if err != nil {
		return nil, err
	}

	resp, ok := r.(*response)
	if !ok {
		return nil, server.ErrBadCast
	}

	return &pb.AddReply{
		V: resp.V,
	}, nil
}

func grpcInstrument(requests metrics.Counter, duration metrics.Histogram) func(pb.AddServer) pb.AddServer {
	return func(next pb.AddServer) pb.AddServer {
		return grpcInstrumented{requests, duration, next}
	}
}

type grpcInstrumented struct {
	requests metrics.Counter
	duration metrics.Histogram
	next     pb.AddServer
}

func (i grpcInstrumented) Add(ctx context.Context, req *pb.AddRequest) (*pb.AddReply, error) {
	i.requests.Add(1)
	defer func(begin time.Time) { i.duration.Observe(time.Since(begin).Nanoseconds()) }(time.Now())
	return i.next.Add(ctx, req)
}
<commit_msg>addsvc: fix import order<commit_after>package main

import (
	"time"

	"golang.org/x/net/context"

	"github.com/peterbourgon/gokit/addsvc/pb"
	"github.com/peterbourgon/gokit/metrics"
	"github.com/peterbourgon/gokit/server"
)

// A binding wraps an Endpoint so that it's usable by a transport. grpcBinding
// makes an Endpoint usable over gRPC.
type grpcBinding struct{ server.Endpoint }

// Add implements the proto3 AddServer by forwarding to the wrapped Endpoint.
func (b grpcBinding) Add(ctx context.Context, req *pb.AddRequest) (*pb.AddReply, error) {
	addReq := request{req.A, req.B}
	r, err := b.Endpoint(ctx, addReq)
	if err != nil {
		return nil, err
	}

	resp, ok := r.(*response)
	if !ok {
		return nil, server.ErrBadCast
	}

	return &pb.AddReply{
		V: resp.V,
	}, nil
}

func grpcInstrument(requests metrics.Counter, duration metrics.Histogram) func(pb.AddServer) pb.AddServer {
	return func(next pb.AddServer) pb.AddServer {
		return grpcInstrumented{requests, duration, next}
	}
}

type grpcInstrumented struct {
	requests metrics.Counter
	duration metrics.Histogram
	next     pb.AddServer
}

func (i grpcInstrumented) Add(ctx context.Context, req *pb.AddRequest) (*pb.AddReply, error) {
	i.requests.Add(1)
	defer func(begin time.Time) { i.duration.Observe(time.Since(begin).Nanoseconds()) }(time.Now())
	return i.next.Add(ctx, req)
}
<|endoftext|>


<commit_before>package totp

import (
	"bytes"
	"encoding/base32"
	"encoding/base64"
	"fmt"
	"image/png"

	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/logical/framework"
	otplib "github.com/pquerna/otp"
	totplib "github.com/pquerna/otp/totp"
)

func pathListKeys(b *backend) *framework.Path {
	return &framework.Path{
		Pattern: "keys/?$",

		Callbacks: map[logical.Operation]framework.OperationFunc{
			logical.ListOperation: b.pathKeyList,
		},

		HelpSynopsis:    pathKeyHelpSyn,
		HelpDescription: pathKeyHelpDesc,
	}
}

func pathKeys(b *backend) *framework.Path {
	return &framework.Path{
		Pattern: "keys/" + framework.GenericNameRegex("name"),
		Fields: map[string]*framework.FieldSchema{
			"name": {
				Type:        framework.TypeString,
				Description: "Name of the key.",
			},

			"generate": {
				Type:        framework.TypeBool,
				Default:     false,
				Description: "Determines if a key is generated by vault or another service.",
			},

			"key_size": {
				Type:        framework.TypeInt,
				Default:     20,
				Description: "Determines the size in bytes of the generated key.",
			},

			"key": {
				Type:        framework.TypeString,
				Description: "The shared master key used to generate a TOTP token.",
			},

			"issuer": {
				Type:        framework.TypeString,
				Description: `The name of the key's issuing organization.`,
			},

			"account_name": {
				Type:        framework.TypeString,
				Description: `The name of the account associated with the key.`,
			},

			"period": {
				Type:        framework.TypeDurationSecond,
				Default:     30,
				Description: `The length of time used to generate a counter for the TOTP token calculation.`,
			},

			"algorithm": {
				Type:        framework.TypeString,
				Default:     "SHA1",
				Description: `The hashing algorithm used to generate the TOTP token. Options include SHA1, SHA256 and SHA512`,
			},

			"digits": {
				Type:        framework.TypeInt,
				Default:     6,
				Description: `The number of digits in the generated TOTP token. This value can either be 6 or 8`,
			},

			"skew": {
				Type:        framework.TypeInt,
				Default:     0,
				Description: `The number of delay periods that are allowed when validating a TOTP token. This value can either be 0 or 1`,
			},

			"qr_size": {
				Type:        framework.TypeInt,
				Default:     200,
				Description: `The pixel size of the generated square QR code.`,
			},
		},

		Callbacks: map[logical.Operation]framework.OperationFunc{
			logical.ReadOperation:   b.pathKeyRead,
			logical.UpdateOperation: b.pathKeyCreate,
			logical.DeleteOperation: b.pathKeyDelete,
		},

		HelpSynopsis:    pathKeyHelpSyn,
		HelpDescription: pathKeyHelpDesc,
	}
}

func (b *backend) Key(s logical.Storage, n string) (*keyEntry, error) {
	entry, err := s.Get("key/" + n)
	if err != nil {
		return nil, err
	}
	if entry == nil {
		return nil, nil
	}

	var result keyEntry
	if err := entry.DecodeJSON(&result); err != nil {
		return nil, err
	}

	return &result, nil
}

func (b *backend) pathKeyDelete(
	req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	err := req.Storage.Delete("key/" + data.Get("name").(string))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (b *backend) pathKeyRead(
	req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	key, err := b.Key(req.Storage, data.Get("name").(string))
	if err != nil {
		return nil, err
	}
	if key == nil {
		return nil, nil
	}

	// Translate algorithm back to string
	algorithm := key.Algorithm.String()

	// Return values of key
	return &logical.Response{
		Data: map[string]interface{}{
			"issuer":       key.Issuer,
			"account_name": key.AccountName,
			"period":       key.Period,
			"algorithm":    algorithm,
			"digits":       key.Digits,
		},
	}, nil
}

func (b *backend) pathKeyList(
	req *logical.Request, d *framework.FieldData) (*logical.Response, error) {
	entries, err := req.Storage.List("key/")
	if err != nil {
		return nil, err
	}

	return logical.ListResponse(entries), nil
}

func (b *backend) pathKeyCreate(
	req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	name := data.Get("name").(string)
	generate := data.Get("generate").(bool)
	keyString := data.Get("key").(string)
	issuer := data.Get("issuer").(string)
	accountName := data.Get("account_name").(string)
	period := data.Get("period").(int)
	algorithm := data.Get("algorithm").(string)
	digits := data.Get("digits").(int)
	skew := data.Get("skew").(int)
	qrSize := data.Get("qr_size").(int)
	keySize := data.Get("key_size").(int)

	// Translate digits and algorithm to a format the totp library understands
	var keyDigits otplib.Digits
	switch digits {
	case 6:
		keyDigits = otplib.DigitsSix
	case 8:
		keyDigits = otplib.DigitsEight
	default:
		return logical.ErrorResponse("the digit value can only be 6 or 8"), nil
	}

	var keyAlgorithm otplib.Algorithm
	switch algorithm {
	case "SHA1":
		keyAlgorithm = otplib.AlgorithmSHA1
	case "SHA256":
		keyAlgorithm = otplib.AlgorithmSHA256
	case "SHA512":
		keyAlgorithm = otplib.AlgorithmSHA512
	default:
		return logical.ErrorResponse("the algorithm value is not valid"), nil
	}

	// Enforce input value requirements
	if period <= 0 {
		return logical.ErrorResponse("the period value must be greater than zero"), nil
	}

	switch skew {
	case 0:
	case 1:
	default:
		return logical.ErrorResponse("the skew value must be 0 or 1"), nil
	}

	if qrSize <= 0 {
		return logical.ErrorResponse("the qr_size value must be greater than zero"), nil
	}

	if keySize <= 0 {
		return logical.ErrorResponse("the key_size value must be greater than zero"), nil
	}

	// Period, Skew and Key Size need to be unsigned ints
	uintPeriod := uint(period)
	uintSkew := uint(skew)
	uintKeySize := uint(keySize)

	url := ""

	var response logical.Response

	switch generate {
	case true:
		// If the key is generated, Account Name and Issuer are required.
		if accountName == "" {
			return logical.ErrorResponse("the account_name value is required for generated keys"), nil
		}

		if issuer == "" {
			return logical.ErrorResponse("the issuer value is required for generated keys"), nil
		}

		// Generate a new key
		keyObject, err := totplib.Generate(totplib.GenerateOpts{
			Issuer:      issuer,
			AccountName: accountName,
			Period:      uintPeriod,
			Digits:      keyDigits,
			Algorithm:   keyAlgorithm,
			SecretSize:  uintKeySize,
		})

		if err != nil {
			return logical.ErrorResponse("an error occured while generating a key"), nil
		}

		// Get key string value
		keyString = keyObject.Secret()

		// Prepare the url and barcode
		url = keyObject.String()
		barcode, err := keyObject.Image(qrSize, qrSize)

		if err != nil {
			return logical.ErrorResponse("an error occured while generating a QR code image"), nil
		}

		var buff bytes.Buffer
		png.Encode(&buff, barcode)
		b64Barcode := base64.StdEncoding.EncodeToString(buff.Bytes())
		response = logical.Response{
			Data: map[string]interface{}{
				"url":     url,
				"barcode": b64Barcode,
			},
		}
	default:
		if keyString == "" {
			return logical.ErrorResponse("the key value is required"), nil
		}

		_, err := base32.StdEncoding.DecodeString(keyString)

		if err != nil {
			return logical.ErrorResponse(fmt.Sprintf(
				"invalid key value: %s", err)), nil
		}
	}

	// Store it
	entry, err := logical.StorageEntryJSON("key/"+name, &keyEntry{
		Key:         keyString,
		Issuer:      issuer,
		AccountName: accountName,
		Period:      uintPeriod,
		Algorithm:   keyAlgorithm,
		Digits:      keyDigits,
		Skew:        uintSkew,
		QRSize:      qrSize,
		URL:         url,
		Generate:    generate,
	})
	if err != nil {
		return nil, err
	}
	if err := req.Storage.Put(entry); err != nil {
		return nil, err
	}

	switch generate {
	case true:
		return &response, nil
	default:
		return nil, nil
	}
}

type keyEntry struct {
	Key         string           `json:"key" mapstructure:"key" structs:"key"`
	Issuer      string           `json:"issuer" mapstructure:"issuer" structs:"issuer"`
	AccountName string           `json:"account_name" mapstructure:"account_name" structs:"account_name"`
	Period      uint             `json:"period" mapstructure:"period" structs:"period"`
	Algorithm   otplib.Algorithm `json:"algorithm" mapstructure:"algorithm" structs:"algorithm"`
	Digits      otplib.Digits    `json:"digits" mapstructure:"digits" structs:"digits"`
	Skew        uint             `json:"skew" mapstructure:"skew" structs:"skew"`
	QRSize      int              `json:"qr_size" mapstructure:"qr_size" structs:"qr_size"`
	URL         string           `json:"url" mapstructure:"url" structs:"url"`
	Generate    bool             `json:"generate" mapstructure:"generate" structs:"generate"`
}

const pathKeyHelpSyn = `
Manage the keys that can be created with this backend.
`

const pathKeyHelpDesc = `
This path lets you manage the keys that can be created with this backend.

`
<commit_msg>Added ability to pass in TOTP urls<commit_after>package totp

import (
	"bytes"
	"encoding/base32"
	"encoding/base64"
	"fmt"
	"image/png"
	"net/url"
	"strconv"
	"strings"

	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/logical/framework"
	otplib "github.com/pquerna/otp"
	totplib "github.com/pquerna/otp/totp"
)

func pathListKeys(b *backend) *framework.Path {
	return &framework.Path{
		Pattern: "keys/?$",

		Callbacks: map[logical.Operation]framework.OperationFunc{
			logical.ListOperation: b.pathKeyList,
		},

		HelpSynopsis:    pathKeyHelpSyn,
		HelpDescription: pathKeyHelpDesc,
	}
}

func pathKeys(b *backend) *framework.Path {
	return &framework.Path{
		Pattern: "keys/" + framework.GenericNameRegex("name"),
		Fields: map[string]*framework.FieldSchema{
			"name": {
				Type:        framework.TypeString,
				Description: "Name of the key.",
			},

			"generate": {
				Type:        framework.TypeBool,
				Default:     false,
				Description: "Determines if a key is generated by vault or another service.",
			},

			"key_size": {
				Type:        framework.TypeInt,
				Default:     20,
				Description: "Determines the size in bytes of the generated key.",
			},

			"key": {
				Type:        framework.TypeString,
				Description: "The shared master key used to generate a TOTP token.",
			},

			"issuer": {
				Type:        framework.TypeString,
				Description: `The name of the key's issuing organization.`,
			},

			"account_name": {
				Type:        framework.TypeString,
				Description: `The name of the account associated with the key.`,
			},

			"period": {
				Type:        framework.TypeDurationSecond,
				Default:     30,
				Description: `The length of time used to generate a counter for the TOTP token calculation.`,
			},

			"algorithm": {
				Type:        framework.TypeString,
				Default:     "SHA1",
				Description: `The hashing algorithm used to generate the TOTP token. Options include SHA1, SHA256 and SHA512`,
			},

			"digits": {
				Type:        framework.TypeInt,
				Default:     6,
				Description: `The number of digits in the generated TOTP token. This value can either be 6 or 8`,
			},

			"skew": {
				Type:        framework.TypeInt,
				Default:     0,
				Description: `The number of delay periods that are allowed when validating a TOTP token. This value can either be 0 or 1`,
			},

			"qr_size": {
				Type:        framework.TypeInt,
				Default:     200,
				Description: `The pixel size of the generated square QR code.`,
			},

			"url": {
				Type:        framework.TypeString,
				Description: `A TOTP url string containing all of the parameters for key setup.`,
			},
		},

		Callbacks: map[logical.Operation]framework.OperationFunc{
			logical.ReadOperation:   b.pathKeyRead,
			logical.UpdateOperation: b.pathKeyCreate,
			logical.DeleteOperation: b.pathKeyDelete,
		},

		HelpSynopsis:    pathKeyHelpSyn,
		HelpDescription: pathKeyHelpDesc,
	}
}

func (b *backend) Key(s logical.Storage, n string) (*keyEntry, error) {
	entry, err := s.Get("key/" + n)
	if err != nil {
		return nil, err
	}
	if entry == nil {
		return nil, nil
	}

	var result keyEntry
	if err := entry.DecodeJSON(&result); err != nil {
		return nil, err
	}

	return &result, nil
}

func (b *backend) pathKeyDelete(
	req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	err := req.Storage.Delete("key/" + data.Get("name").(string))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (b *backend) pathKeyRead(
	req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	key, err := b.Key(req.Storage, data.Get("name").(string))
	if err != nil {
		return nil, err
	}
	if key == nil {
		return nil, nil
	}

	// Translate algorithm back to string
	algorithm := key.Algorithm.String()

	// Return values of key
	return &logical.Response{
		Data: map[string]interface{}{
			"issuer":       key.Issuer,
			"account_name": key.AccountName,
			"period":       key.Period,
			"algorithm":    algorithm,
			"digits":       key.Digits,
		},
	}, nil
}

func (b *backend) pathKeyList(
	req *logical.Request, d *framework.FieldData) (*logical.Response, error) {
	entries, err := req.Storage.List("key/")
	if err != nil {
		return nil, err
	}

	return logical.ListResponse(entries), nil
}

func (b *backend) pathKeyCreate(
	req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	name := data.Get("name").(string)
	generate := data.Get("generate").(bool)
	keyString := data.Get("key").(string)
	issuer := data.Get("issuer").(string)
	accountName := data.Get("account_name").(string)
	period := data.Get("period").(int)
	algorithm := data.Get("algorithm").(string)
	digits := data.Get("digits").(int)
	skew := data.Get("skew").(int)
	qrSize := data.Get("qr_size").(int)
	keySize := data.Get("key_size").(int)
	inputURL := data.Get("url").(string)

	// Read parameters from url if given
	if inputURL != "" {
		//Parse url
		urlObject, err := url.Parse(inputURL)

		if err != nil {
			return logical.ErrorResponse("an error occured while parsing url string"), nil
		}

		//Set up query object
		urlQuery := urlObject.Query()
		path := strings.TrimPrefix(urlObject.Path, "/")
		index := strings.Index(path, ":")

		//Read issuer
		urlIssuer := urlQuery.Get("issuer")
		if urlIssuer != "" {
			issuer = urlIssuer
		} else {
			if index != -1 {
				issuer = path[:index]
			}
		}

		//Read account name
		if index == -1 {
			accountName = path
		} else {
			accountName = path[index+1:]
		}

		//Read key string
		keyString = urlQuery.Get("secret")

		//Read period
		periodQuery, err := strconv.Atoi(urlQuery.Get("period"))

		if err == nil {
			period = periodQuery
		}

		//Read digits
		digitsQuery, err := strconv.Atoi(urlQuery.Get("digits"))

		if err == nil {
			digits = digitsQuery
		}

		//Read algorithm
		algorithmQuery := urlQuery.Get("algorithm")

		if algorithmQuery != "" {
			algorithm = algorithmQuery
		}
	}

	// Translate digits and algorithm to a format the totp library understands
	var keyDigits otplib.Digits
	switch digits {
	case 6:
		keyDigits = otplib.DigitsSix
	case 8:
		keyDigits = otplib.DigitsEight
	default:
		return logical.ErrorResponse("the digit value can only be 6 or 8"), nil
	}

	var keyAlgorithm otplib.Algorithm
	switch algorithm {
	case "SHA1":
		keyAlgorithm = otplib.AlgorithmSHA1
	case "SHA256":
		keyAlgorithm = otplib.AlgorithmSHA256
	case "SHA512":
		keyAlgorithm = otplib.AlgorithmSHA512
	default:
		return logical.ErrorResponse("the algorithm value is not valid"), nil
	}

	// Enforce input value requirements
	if period <= 0 {
		return logical.ErrorResponse("the period value must be greater than zero"), nil
	}

	switch skew {
	case 0:
	case 1:
	default:
		return logical.ErrorResponse("the skew value must be 0 or 1"), nil
	}

	if qrSize <= 0 {
		return logical.ErrorResponse("the qr_size value must be greater than zero"), nil
	}

	if keySize <= 0 {
		return logical.ErrorResponse("the key_size value must be greater than zero"), nil
	}

	// Period, Skew and Key Size need to be unsigned ints
	uintPeriod := uint(period)
	uintSkew := uint(skew)
	uintKeySize := uint(keySize)

	urlString := ""

	var response logical.Response

	switch generate {
	case true:
		// If the key is generated, Account Name and Issuer are required.
		if accountName == "" {
			return logical.ErrorResponse("the account_name value is required for generated keys"), nil
		}

		if issuer == "" {
			return logical.ErrorResponse("the issuer value is required for generated keys"), nil
		}

		// Generate a new key
		keyObject, err := totplib.Generate(totplib.GenerateOpts{
			Issuer:      issuer,
			AccountName: accountName,
			Period:      uintPeriod,
			Digits:      keyDigits,
			Algorithm:   keyAlgorithm,
			SecretSize:  uintKeySize,
		})

		if err != nil {
			return logical.ErrorResponse("an error occured while generating a key"), nil
		}

		// Get key string value
		keyString = keyObject.Secret()

		// Prepare the url and barcode
		urlString = keyObject.String()
		barcode, err := keyObject.Image(qrSize, qrSize)

		if err != nil {
			return logical.ErrorResponse("an error occured while generating a QR code image"), nil
		}

		var buff bytes.Buffer
		png.Encode(&buff, barcode)
		b64Barcode := base64.StdEncoding.EncodeToString(buff.Bytes())
		response = logical.Response{
			Data: map[string]interface{}{
				"url":     urlString,
				"barcode": b64Barcode,
			},
		}
	default:
		if keyString == "" {
			return logical.ErrorResponse("the key value is required"), nil
		}

		_, err := base32.StdEncoding.DecodeString(keyString)

		if err != nil {
			return logical.ErrorResponse(fmt.Sprintf(
				"invalid key value: %s", err)), nil
		}
	}

	// Store it
	entry, err := logical.StorageEntryJSON("key/"+name, &keyEntry{
		Key:         keyString,
		Issuer:      issuer,
		AccountName: accountName,
		Period:      uintPeriod,
		Algorithm:   keyAlgorithm,
		Digits:      keyDigits,
		Skew:        uintSkew,
		QRSize:      qrSize,
		URL:         urlString,
		Generate:    generate,
	})
	if err != nil {
		return nil, err
	}
	if err := req.Storage.Put(entry); err != nil {
		return nil, err
	}

	switch generate {
	case true:
		return &response, nil
	default:
		return nil, nil
	}
}

type keyEntry struct {
	Key         string           `json:"key" mapstructure:"key" structs:"key"`
	Issuer      string           `json:"issuer" mapstructure:"issuer" structs:"issuer"`
	AccountName string           `json:"account_name" mapstructure:"account_name" structs:"account_name"`
	Period      uint             `json:"period" mapstructure:"period" structs:"period"`
	Algorithm   otplib.Algorithm `json:"algorithm" mapstructure:"algorithm" structs:"algorithm"`
	Digits      otplib.Digits    `json:"digits" mapstructure:"digits" structs:"digits"`
	Skew        uint             `json:"skew" mapstructure:"skew" structs:"skew"`
	QRSize      int              `json:"qr_size" mapstructure:"qr_size" structs:"qr_size"`
	URL         string           `json:"url" mapstructure:"url" structs:"url"`
	Generate    bool             `json:"generate" mapstructure:"generate" structs:"generate"`
}

const pathKeyHelpSyn = `
Manage the keys that can be created with this backend.
`

const pathKeyHelpDesc = `
This path lets you manage the keys that can be created with this backend.

`
<|endoftext|>


<commit_before>package sshego

import (
	"context"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"time"

	ssh "github.com/glycerine/sshego/xendor/github.com/glycerine/xcryptossh"
)

var ErrShutdown = fmt.Errorf("shutting down")

// Tricorder records (holds) three key objects:
//   an *ssh.Client, the underlyign net.Conn, and a
//   set of ssh.Channel
//
// Tricorder supports auto reconnect when disconnected.
//
type Tricorder struct {

	// optional, parent can provide us
	// a Halter, and we will AddDownstream.
	ParentHalt *ssh.Halter

	// should only reflect close of the internal sshChannel, not cli nor nc.
	ChannelHalt *ssh.Halter

	cfg *SshegoConfig

	cli         *ssh.Client
	nc          io.Closer
	uhp         *UHP
	sshChannels map[net.Conn]context.CancelFunc

	lastSshChannel ssh.Channel

	getChannelCh      chan *getChannelTicket
	getCliCh          chan *ssh.Client
	getNcCh           chan io.Closer
	reconnectNeededCh chan *UHP

	dc *DialConfig
}

/*
NewTricorder has got to wait to allocate
ssh.Channel until requested. Otherwise we
make too many, and get them mixed up.
*/
func (cfg *SshegoConfig) NewTricorder(halt *ssh.Halter, dc *DialConfig, sshClient *ssh.Client, sshChan net.Conn) (tri *Tricorder) {

	tri = &Tricorder{
		dc:          dc,
		cfg:         cfg,
		ParentHalt:  halt,
		ChannelHalt: ssh.NewHalter(),

		sshChannels: make(map[net.Conn]context.CancelFunc),

		reconnectNeededCh: make(chan *UHP, 1),
		getChannelCh:      make(chan *getChannelTicket),
		getCliCh:          make(chan *ssh.Client),
		getNcCh:           make(chan io.Closer),
	}
	if tri.ParentHalt != nil {
		tri.ParentHalt.AddDownstream(tri.ChannelHalt)
	}
	// first call to Subscribe is here.
	cfg.ClientReconnectNeededTower.Subscribe(tri.reconnectNeededCh)
	if sshClient != nil {
		tri.cli = sshClient
		tri.nc = sshClient.NcCloser()
	}
	if sshChan != nil {
		tri.sshChannels[sshChan] = nil
	}

	tri.startReconnectLoop()
	return tri
}

// CustomInprocStreamChanName is how sshego/reptile specific
// channels are named.
//const CustomInprocStreamChanName = "custom-inproc-stream"
const CustomInprocStreamChanName = "direct-tcpip"

func (t *Tricorder) closeChannels() {
	if len(t.sshChannels) > 0 {
		for ch, cancel := range t.sshChannels {
			ch.Close()
			if cancel != nil {
				cancel()
			}
		}
	}
	t.sshChannels = make(map[net.Conn]context.CancelFunc)
}

func (t *Tricorder) startReconnectLoop() {
	go func() {
		defer func() {
			t.ChannelHalt.RequestStop()
			t.ChannelHalt.MarkDone()
			if t.ParentHalt != nil {
				t.ParentHalt.RemoveDownstream(t.ChannelHalt)
			}
			t.closeChannels()
		}()
		for {
			select {
			case <-t.ChannelHalt.ReqStopChan():
				return
			case uhp := <-t.reconnectNeededCh:
				pp("Tricorder sees reconnectNeeded!!")
				t.uhp = uhp
				t.closeChannels()

				t.ChannelHalt.RequestStop()
				t.ChannelHalt.MarkDone()

				if t.ParentHalt != nil {
					t.ParentHalt.RemoveDownstream(t.ChannelHalt)
				}
				t.ChannelHalt = ssh.NewHalter()
				if t.ParentHalt != nil {
					t.ParentHalt.AddDownstream(t.ChannelHalt)
				}

				t.cli = nil
				t.nc = nil
				// need to reconnect!
				t.helperNewClientConnect()

				// provide current state
			case t.getCliCh <- t.cli:
			case t.getNcCh <- t.nc:

				// bring up a new channel
			case tk := <-t.getChannelCh:
				t.helperGetChannel(tk)
			}
		}
	}()
}

// only reconnect, don't open any new channels!
func (t *Tricorder) helperNewClientConnect() {

	pp("Tricorder.helperNewClientConnect starting!")
	//	time.Sleep(time.Second * 5)
	//	pp("Tricorder.helperNewClientConnect is done pausing!!")

	destHost, port, err := splitHostPort(t.uhp.HostPort)
	_, _ = destHost, port
	panicOn(err)

	ctx := context.Background()
	pw := ""
	_ = pw
	toptUrl := ""
	_ = toptUrl
	//t.cfg.AddIfNotKnown = false
	var sshcli *ssh.Client
	tries := 3
	pause := 1000 * time.Millisecond
	if t.cfg.KnownHosts == nil {
		panic("problem! t.cfg.KnownHosts is nil")
	}
	if t.cfg.PrivateKeyPath == "" {
		panic("problem! t.cfg.PrivateKeyPath is empty")
	}

	const skipDownstreamChannelCreation = true
	for i := 0; i < tries; i++ {
		pp("Tricorder.helperNewClientConnect() calilng t.dc.Dial(), i=%v", i)
		_, sshcli, _, err = t.dc.Dial(ctx, skipDownstreamChannelCreation)
		//		sshcli, sshChan, err = t.cfg.SSHConnect(ctx, t.cfg.KnownHosts, t.uhp.User, t.cfg.PrivateKeyPath, destHost, int64(port), pw, toptUrl, t.ChannelHalt)
		if err == nil {
			break

		} else {
			if strings.Contains(err.Error(), "connection refused") {
				pp("Tricorder.helperNewClientConnect: ignoring 'connection refused' and retrying.")
				time.Sleep(pause)
				continue
			}
			pp("err = '%v'. retrying", err)
			time.Sleep(pause)
			continue
		}
	}
	panicOn(err)
	pp("good: Tricorder.helperNewClientConnect succeeded.")
	t.cli = sshcli
	t.nc = t.cli.NcCloser()
}

func (t *Tricorder) helperGetChannel(tk *getChannelTicket) {

	pp("Tricorder.helperGetChannel starting!")

	var ch net.Conn
	var err error
	if t.cli == nil {
		pp("Tricorder.helperGetChannel: saw nil cli, so making new client")
		t.helperNewClientConnect()
	}

	pp("Tricorder.helperGetChannel: had cli already, so calling t.cli.Dial()")

	// for now assume we are doing a "direct-tcpip" forward
	hp := strings.Trim(t.dc.DownstreamHostPort, "\n\r\t ")
	pp("Tricorder.helperGetChannel dialing hp='%v'", hp)
	ch, err = t.cli.Dial("tcp", hp)
	if ch != nil {
		t.sshChannels[ch] = nil
	}
	/*

		bkg := context.Background()
		ctx, cancelOpenChannelCtx := context.WithDeadline(bkg, time.Now().Add(5*time.Second))

		defer cancelOpenChannelCtx() // TODO: is this right??

		ch, in, err := t.cli.OpenChannel(ctx, CustomInprocStreamChanName, nil)
		if err == nil {
			t.lastSshChannel = ch
			discardCtx, discardCtxCancel := context.WithCancel(bkg)
			go DiscardRequestsExceptKeepalives(discardCtx, in, t.ChannelHalt.ReqStopChan())
			t.sshChannels[ch] = discardCtxCancel

			if ch != nil && t.cfg.IdleTimeoutDur > 0 {
				ch.SetIdleTimeout(t.cfg.IdleTimeoutDur)
			}
		}
	*/

	tk.sshChannel = ch
	tk.err = err

	close(tk.done)
}

type getChannelTicket struct {
	done         chan struct{}
	sshChannel   net.Conn
	destHostPort string
	err          error
}

func newGetChannelTicket() *getChannelTicket {
	return &getChannelTicket{
		done: make(chan struct{}),
	}
}

func (t *Tricorder) SSHChannel() (net.Conn, error) {
	tk := newGetChannelTicket()
	t.getChannelCh <- tk
	<-tk.done
	return tk.sshChannel, tk.err
}

func (t *Tricorder) Cli() (cli *ssh.Client, err error) {
	select {
	case cli = <-t.getCliCh:
	case <-t.ChannelHalt.ReqStopChan():
		err = ErrShutdown
	}
	return
}

func (t *Tricorder) Nc() (nc io.Closer, err error) {
	select {
	case nc = <-t.getNcCh:
	case <-t.ChannelHalt.ReqStopChan():
		err = ErrShutdown
	}
	return
}

func splitHostPort(hostport string) (host string, port int, err error) {
	sPort := ""
	host, sPort, err = net.SplitHostPort(hostport)
	if err != nil {
		err = fmt.Errorf("bad addr '%s': net.SplitHostPort() gave: %s", hostport, err)
		return
	}
	if host == "" {
		host = "127.0.0.1"
	}
	if len(sPort) == 0 {
		err = fmt.Errorf("no port found in '%s'", hostport)
		return
	}
	var prt uint64
	prt, err = strconv.ParseUint(sPort, 10, 16)
	if err != nil {
		return
	}
	port = int(prt)
	return
}

/*
func (t *Tricorder) Write(p []byte) (n int, err error) {
	return t.lastSshChannel.Write(p)
}

func (t *Tricorder) Read(p []byte) (n int, err error) {
	return t.lastSshChannel.Read(p)
}
*/
<commit_msg>cleanup<commit_after>package sshego

import (
	"context"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"time"

	ssh "github.com/glycerine/sshego/xendor/github.com/glycerine/xcryptossh"
)

var ErrShutdown = fmt.Errorf("shutting down")

// Tricorder records (holds) three key objects:
//   an *ssh.Client, the underlyign net.Conn, and a
//   set of ssh.Channel
//
// Tricorder supports auto reconnect when disconnected.
//
type Tricorder struct {

	// optional, parent can provide us
	// a Halter, and we will AddDownstream.
	ParentHalt *ssh.Halter

	// should only reflect close of the internal sshChannel, not cli nor nc.
	ChannelHalt *ssh.Halter

	cfg *SshegoConfig

	cli         *ssh.Client
	nc          io.Closer
	uhp         *UHP
	sshChannels map[net.Conn]context.CancelFunc

	lastSshChannel ssh.Channel

	getChannelCh      chan *getChannelTicket
	getCliCh          chan *ssh.Client
	getNcCh           chan io.Closer
	reconnectNeededCh chan *UHP

	dc *DialConfig
}

/*
NewTricorder has got to wait to allocate
ssh.Channel until requested. Otherwise we
make too many, and get them mixed up.
*/
func (cfg *SshegoConfig) NewTricorder(halt *ssh.Halter, dc *DialConfig, sshClient *ssh.Client, sshChan net.Conn) (tri *Tricorder) {

	tri = &Tricorder{
		dc:          dc,
		cfg:         cfg,
		ParentHalt:  halt,
		ChannelHalt: ssh.NewHalter(),

		sshChannels: make(map[net.Conn]context.CancelFunc),

		reconnectNeededCh: make(chan *UHP, 1),
		getChannelCh:      make(chan *getChannelTicket),
		getCliCh:          make(chan *ssh.Client),
		getNcCh:           make(chan io.Closer),
	}
	if tri.ParentHalt != nil {
		tri.ParentHalt.AddDownstream(tri.ChannelHalt)
	}
	// first call to Subscribe is here.
	cfg.ClientReconnectNeededTower.Subscribe(tri.reconnectNeededCh)
	if sshClient != nil {
		tri.cli = sshClient
		tri.nc = sshClient.NcCloser()
	}
	if sshChan != nil {
		tri.sshChannels[sshChan] = nil
	}

	tri.startReconnectLoop()
	return tri
}

// CustomInprocStreamChanName is how sshego/reptile specific
// channels are named.
//const CustomInprocStreamChanName = "custom-inproc-stream"
const CustomInprocStreamChanName = "direct-tcpip"

func (t *Tricorder) closeChannels() {
	if len(t.sshChannels) > 0 {
		for ch, cancel := range t.sshChannels {
			ch.Close()
			if cancel != nil {
				cancel()
			}
		}
	}
	t.sshChannels = make(map[net.Conn]context.CancelFunc)
}

func (t *Tricorder) startReconnectLoop() {
	go func() {
		defer func() {
			t.ChannelHalt.RequestStop()
			t.ChannelHalt.MarkDone()
			if t.ParentHalt != nil {
				t.ParentHalt.RemoveDownstream(t.ChannelHalt)
			}
			t.closeChannels()
		}()
		for {
			select {
			case <-t.ChannelHalt.ReqStopChan():
				return
			case uhp := <-t.reconnectNeededCh:
				pp("Tricorder sees reconnectNeeded!!")
				t.uhp = uhp
				t.closeChannels()

				t.ChannelHalt.RequestStop()
				t.ChannelHalt.MarkDone()

				if t.ParentHalt != nil {
					t.ParentHalt.RemoveDownstream(t.ChannelHalt)
				}
				t.ChannelHalt = ssh.NewHalter()
				if t.ParentHalt != nil {
					t.ParentHalt.AddDownstream(t.ChannelHalt)
				}

				t.cli = nil
				t.nc = nil
				// need to reconnect!
				t.helperNewClientConnect()

				// provide current state
			case t.getCliCh <- t.cli:
			case t.getNcCh <- t.nc:

				// bring up a new channel
			case tk := <-t.getChannelCh:
				t.helperGetChannel(tk)
			}
		}
	}()
}

// only reconnect, don't open any new channels!
func (t *Tricorder) helperNewClientConnect() {

	pp("Tricorder.helperNewClientConnect starting!")

	destHost, port, err := splitHostPort(t.uhp.HostPort)
	_, _ = destHost, port
	panicOn(err)

	ctx := context.Background()
	pw := ""
	_ = pw
	toptUrl := ""
	_ = toptUrl
	//t.cfg.AddIfNotKnown = false
	var sshcli *ssh.Client
	tries := 3
	pause := 1000 * time.Millisecond
	if t.cfg.KnownHosts == nil {
		panic("problem! t.cfg.KnownHosts is nil")
	}
	if t.cfg.PrivateKeyPath == "" {
		panic("problem! t.cfg.PrivateKeyPath is empty")
	}

	const skipDownstreamChannelCreation = true
	for i := 0; i < tries; i++ {
		pp("Tricorder.helperNewClientConnect() calilng t.dc.Dial(), i=%v", i)
		_, sshcli, _, err = t.dc.Dial(ctx, skipDownstreamChannelCreation)
		//		sshcli, sshChan, err = t.cfg.SSHConnect(ctx, t.cfg.KnownHosts, t.uhp.User, t.cfg.PrivateKeyPath, destHost, int64(port), pw, toptUrl, t.ChannelHalt)

		if err == nil {
			break

		} else {
			if strings.Contains(err.Error(), "connection refused") {
				pp("Tricorder.helperNewClientConnect: ignoring 'connection refused' and retrying.")
				time.Sleep(pause)
				continue
			}
			pp("err = '%v'. retrying", err)
			time.Sleep(pause)
			continue
		}
	}
	panicOn(err)
	pp("good: Tricorder.helperNewClientConnect succeeded.")
	t.cli = sshcli
	t.nc = t.cli.NcCloser()
}

func (t *Tricorder) helperGetChannel(tk *getChannelTicket) {

	pp("Tricorder.helperGetChannel starting!")

	var ch net.Conn
	var err error
	if t.cli == nil {
		pp("Tricorder.helperGetChannel: saw nil cli, so making new client")
		t.helperNewClientConnect()
	}

	pp("Tricorder.helperGetChannel: had cli already, so calling t.cli.Dial()")

	// for now assume we are doing a "direct-tcpip" forward
	hp := strings.Trim(t.dc.DownstreamHostPort, "\n\r\t ")
	pp("Tricorder.helperGetChannel dialing hp='%v'", hp)
	ch, err = t.cli.Dial("tcp", hp)
	if ch != nil {
		t.sshChannels[ch] = nil
	}
	/*

		bkg := context.Background()
		ctx, cancelOpenChannelCtx := context.WithDeadline(bkg, time.Now().Add(5*time.Second))

		defer cancelOpenChannelCtx() // TODO: is this right??

		ch, in, err := t.cli.OpenChannel(ctx, CustomInprocStreamChanName, nil)
		if err == nil {
			t.lastSshChannel = ch
			discardCtx, discardCtxCancel := context.WithCancel(bkg)
			go DiscardRequestsExceptKeepalives(discardCtx, in, t.ChannelHalt.ReqStopChan())
			t.sshChannels[ch] = discardCtxCancel

			if ch != nil && t.cfg.IdleTimeoutDur > 0 {
				ch.SetIdleTimeout(t.cfg.IdleTimeoutDur)
			}
		}
	*/

	tk.sshChannel = ch
	tk.err = err

	close(tk.done)
}

type getChannelTicket struct {
	done         chan struct{}
	sshChannel   net.Conn
	destHostPort string
	err          error
}

func newGetChannelTicket() *getChannelTicket {
	return &getChannelTicket{
		done: make(chan struct{}),
	}
}

func (t *Tricorder) SSHChannel() (net.Conn, error) {
	tk := newGetChannelTicket()
	t.getChannelCh <- tk
	<-tk.done
	return tk.sshChannel, tk.err
}

func (t *Tricorder) Cli() (cli *ssh.Client, err error) {
	select {
	case cli = <-t.getCliCh:
	case <-t.ChannelHalt.ReqStopChan():
		err = ErrShutdown
	}
	return
}

func (t *Tricorder) Nc() (nc io.Closer, err error) {
	select {
	case nc = <-t.getNcCh:
	case <-t.ChannelHalt.ReqStopChan():
		err = ErrShutdown
	}
	return
}

func splitHostPort(hostport string) (host string, port int, err error) {
	sPort := ""
	host, sPort, err = net.SplitHostPort(hostport)
	if err != nil {
		err = fmt.Errorf("bad addr '%s': net.SplitHostPort() gave: %s", hostport, err)
		return
	}
	if host == "" {
		host = "127.0.0.1"
	}
	if len(sPort) == 0 {
		err = fmt.Errorf("no port found in '%s'", hostport)
		return
	}
	var prt uint64
	prt, err = strconv.ParseUint(sPort, 10, 16)
	if err != nil {
		return
	}
	port = int(prt)
	return
}

/*
func (t *Tricorder) Write(p []byte) (n int, err error) {
	return t.lastSshChannel.Write(p)
}

func (t *Tricorder) Read(p []byte) (n int, err error) {
	return t.lastSshChannel.Read(p)
}
*/
<|endoftext|>


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

// Package gopathwalk is like filepath.Walk but specialized for finding Go
// packages, particularly in $GOPATH and $GOROOT.
package gopathwalk

import (
	"bufio"
	"bytes"
	"fmt"
	"go/build"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/tools/internal/fastwalk"
)

// Options controls the behavior of a Walk call.
type Options struct {
	Debug          bool // Enable debug logging
	ModulesEnabled bool // Search module caches. Also disables legacy goimports ignore rules.
}

// RootType indicates the type of a Root.
type RootType int

const (
	RootUnknown RootType = iota
	RootGOROOT
	RootGOPATH
	RootCurrentModule
	RootModuleCache
	RootOther
)

// A Root is a starting point for a Walk.
type Root struct {
	Path string
	Type RootType
}

// SrcDirsRoots returns the roots from build.Default.SrcDirs(). Not modules-compatible.
func SrcDirsRoots(ctx *build.Context) []Root {
	var roots []Root
	roots = append(roots, Root{filepath.Join(ctx.GOROOT, "src"), RootGOROOT})
	for _, p := range filepath.SplitList(ctx.GOPATH) {
		roots = append(roots, Root{filepath.Join(p, "src"), RootGOPATH})
	}
	return roots
}

// Walk walks Go source directories ($GOROOT, $GOPATH, etc) to find packages.
// For each package found, add will be called (concurrently) with the absolute
// paths of the containing source directory and the package directory.
// add will be called concurrently.
func Walk(roots []Root, add func(root Root, dir string), opts Options) {
	WalkSkip(roots, add, func(Root, string) bool { return false }, opts)
}

// WalkSkip walks Go source directories ($GOROOT, $GOPATH, etc) to find packages.
// For each package found, add will be called (concurrently) with the absolute
// paths of the containing source directory and the package directory.
// For each directory that will be scanned, skip will be called (concurrently)
// with the absolute paths of the containing source directory and the directory.
// If skip returns false on a directory it will be processed.
// add will be called concurrently.
// skip will be called concurrently.
func WalkSkip(roots []Root, add func(root Root, dir string), skip func(root Root, dir string) bool, opts Options) {
	for _, root := range roots {
		walkDir(root, add, skip, opts)
	}
}

func walkDir(root Root, add func(Root, string), skip func(root Root, dir string) bool, opts Options) {
	if _, err := os.Stat(root.Path); os.IsNotExist(err) {
		if opts.Debug {
			log.Printf("skipping nonexistent directory: %v", root.Path)
		}
		return
	}
	start := time.Now()
	if opts.Debug {
		log.Printf("gopathwalk: scanning %s", root.Path)
	}
	w := &walker{
		root: root,
		add:  add,
		skip: skip,
		opts: opts,
	}
	w.init()
	if err := fastwalk.Walk(root.Path, w.walk); err != nil {
		log.Printf("gopathwalk: scanning directory %v: %v", root.Path, err)
	}

	if opts.Debug {
		log.Printf("gopathwalk: scanned %s in %v", root.Path, time.Since(start))
	}
}

// walker is the callback for fastwalk.Walk.
type walker struct {
	root Root                    // The source directory to scan.
	add  func(Root, string)      // The callback that will be invoked for every possible Go package dir.
	skip func(Root, string) bool // The callback that will be invoked for every dir. dir is skipped if it returns true.
	opts Options                 // Options passed to Walk by the user.

	ignoredDirs []os.FileInfo // The ignored directories, loaded from .goimportsignore files.
}

// init initializes the walker based on its Options.
func (w *walker) init() {
	var ignoredPaths []string
	if w.root.Type == RootModuleCache {
		ignoredPaths = []string{"cache"}
	}
	if !w.opts.ModulesEnabled && w.root.Type == RootGOPATH {
		ignoredPaths = w.getIgnoredDirs(w.root.Path)
		ignoredPaths = append(ignoredPaths, "v", "mod")
	}

	for _, p := range ignoredPaths {
		full := filepath.Join(w.root.Path, p)
		if fi, err := os.Stat(full); err == nil {
			w.ignoredDirs = append(w.ignoredDirs, fi)
			if w.opts.Debug {
				log.Printf("Directory added to ignore list: %s", full)
			}
		} else if w.opts.Debug {
			log.Printf("Error statting ignored directory: %v", err)
		}
	}
}

// getIgnoredDirs reads an optional config file at <path>/.goimportsignore
// of relative directories to ignore when scanning for go files.
// The provided path is one of the $GOPATH entries with "src" appended.
func (w *walker) getIgnoredDirs(path string) []string {
	file := filepath.Join(path, ".goimportsignore")
	slurp, err := ioutil.ReadFile(file)
	if w.opts.Debug {
		if err != nil {
			log.Print(err)
		} else {
			log.Printf("Read %s", file)
		}
	}
	if err != nil {
		return nil
	}

	var ignoredDirs []string
	bs := bufio.NewScanner(bytes.NewReader(slurp))
	for bs.Scan() {
		line := strings.TrimSpace(bs.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		ignoredDirs = append(ignoredDirs, line)
	}
	return ignoredDirs
}

func (w *walker) shouldSkipDir(fi os.FileInfo, dir string) bool {
	for _, ignoredDir := range w.ignoredDirs {
		if os.SameFile(fi, ignoredDir) {
			return true
		}
	}
	if w.skip != nil {
		// Check with the user specified callback.
		return w.skip(w.root, dir)
	}
	return false
}

func (w *walker) walk(path string, typ os.FileMode) error {
	dir := filepath.Dir(path)
	if typ.IsRegular() {
		if dir == w.root.Path && (w.root.Type == RootGOROOT || w.root.Type == RootGOPATH) {
			// Doesn't make sense to have regular files
			// directly in your $GOPATH/src or $GOROOT/src.
			return fastwalk.SkipFiles
		}
		if !strings.HasSuffix(path, ".go") {
			return nil
		}

		w.add(w.root, dir)
		return fastwalk.SkipFiles
	}
	if typ == os.ModeDir {
		base := filepath.Base(path)
		if base == "" || base[0] == '.' || base[0] == '_' ||
			base == "testdata" ||
			(w.root.Type == RootGOROOT && w.opts.ModulesEnabled && base == "vendor") ||
			(!w.opts.ModulesEnabled && base == "node_modules") {
			return filepath.SkipDir
		}
		fi, err := os.Lstat(path)
		if err == nil && w.shouldSkipDir(fi, path) {
			return filepath.SkipDir
		}
		return nil
	}
	if typ == os.ModeSymlink {
		base := filepath.Base(path)
		if strings.HasPrefix(base, ".#") {
			// Emacs noise.
			return nil
		}
		fi, err := os.Lstat(path)
		if err != nil {
			// Just ignore it.
			return nil
		}
		if w.shouldTraverse(dir, fi) {
			return fastwalk.TraverseLink
		}
	}
	return nil
}

// shouldTraverse reports whether the symlink fi, found in dir,
// should be followed.  It makes sure symlinks were never visited
// before to avoid symlink loops.
func (w *walker) shouldTraverse(dir string, fi os.FileInfo) bool {
	path := filepath.Join(dir, fi.Name())
	target, err := filepath.EvalSymlinks(path)
	if err != nil {
		return false
	}
	ts, err := os.Stat(target)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return false
	}
	if !ts.IsDir() {
		return false
	}
	if w.shouldSkipDir(ts, dir) {
		return false
	}
	// Check for symlink loops by statting each directory component
	// and seeing if any are the same file as ts.
	for {
		parent := filepath.Dir(path)
		if parent == path {
			// Made it to the root without seeing a cycle.
			// Use this symlink.
			return true
		}
		parentInfo, err := os.Stat(parent)
		if err != nil {
			return false
		}
		if os.SameFile(ts, parentInfo) {
			// Cycle. Don't traverse.
			return false
		}
		path = parent
	}

}
<commit_msg>internal/gopathwalk/walk: add missing function descriptions and renaming<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package gopathwalk is like filepath.Walk but specialized for finding Go
// packages, particularly in $GOPATH and $GOROOT.
package gopathwalk

import (
	"bufio"
	"bytes"
	"fmt"
	"go/build"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/tools/internal/fastwalk"
)

// Options controls the behavior of a Walk call.
type Options struct {
	Debug          bool // Enable debug logging
	ModulesEnabled bool // Search module caches. Also disables legacy goimports ignore rules.
}

// RootType indicates the type of a Root.
type RootType int

const (
	RootUnknown RootType = iota
	RootGOROOT
	RootGOPATH
	RootCurrentModule
	RootModuleCache
	RootOther
)

// A Root is a starting point for a Walk.
type Root struct {
	Path string
	Type RootType
}

// SrcDirsRoots returns the roots from build.Default.SrcDirs(). Not modules-compatible.
func SrcDirsRoots(ctx *build.Context) []Root {
	var roots []Root
	roots = append(roots, Root{filepath.Join(ctx.GOROOT, "src"), RootGOROOT})
	for _, p := range filepath.SplitList(ctx.GOPATH) {
		roots = append(roots, Root{filepath.Join(p, "src"), RootGOPATH})
	}
	return roots
}

// Walk walks Go source directories ($GOROOT, $GOPATH, etc) to find packages.
// For each package found, add will be called (concurrently) with the absolute
// paths of the containing source directory and the package directory.
// add will be called concurrently.
func Walk(roots []Root, add func(root Root, dir string), opts Options) {
	WalkSkip(roots, add, func(Root, string) bool { return false }, opts)
}

// WalkSkip walks Go source directories ($GOROOT, $GOPATH, etc) to find packages.
// For each package found, add will be called (concurrently) with the absolute
// paths of the containing source directory and the package directory.
// For each directory that will be scanned, skip will be called (concurrently)
// with the absolute paths of the containing source directory and the directory.
// If skip returns false on a directory it will be processed.
// add will be called concurrently.
// skip will be called concurrently.
func WalkSkip(roots []Root, add func(root Root, dir string), skip func(root Root, dir string) bool, opts Options) {
	for _, root := range roots {
		walkDir(root, add, skip, opts)
	}
}

// walkDir creates a walker and starts fastwalk with this walker.
func walkDir(root Root, add func(Root, string), skip func(root Root, dir string) bool, opts Options) {
	if _, err := os.Stat(root.Path); os.IsNotExist(err) {
		if opts.Debug {
			log.Printf("skipping nonexistent directory: %v", root.Path)
		}
		return
	}
	start := time.Now()
	if opts.Debug {
		log.Printf("gopathwalk: scanning %s", root.Path)
	}
	w := &walker{
		root: root,
		add:  add,
		skip: skip,
		opts: opts,
	}
	w.init()
	if err := fastwalk.Walk(root.Path, w.walk); err != nil {
		log.Printf("gopathwalk: scanning directory %v: %v", root.Path, err)
	}

	if opts.Debug {
		log.Printf("gopathwalk: scanned %s in %v", root.Path, time.Since(start))
	}
}

// walker is the callback for fastwalk.Walk.
type walker struct {
	root Root                    // The source directory to scan.
	add  func(Root, string)      // The callback that will be invoked for every possible Go package dir.
	skip func(Root, string) bool // The callback that will be invoked for every dir. dir is skipped if it returns true.
	opts Options                 // Options passed to Walk by the user.

	ignoredDirs []os.FileInfo // The ignored directories, loaded from .goimportsignore files.
}

// init initializes the walker based on its Options
func (w *walker) init() {
	var ignoredPaths []string
	if w.root.Type == RootModuleCache {
		ignoredPaths = []string{"cache"}
	}
	if !w.opts.ModulesEnabled && w.root.Type == RootGOPATH {
		ignoredPaths = w.getIgnoredDirs(w.root.Path)
		ignoredPaths = append(ignoredPaths, "v", "mod")
	}

	for _, p := range ignoredPaths {
		full := filepath.Join(w.root.Path, p)
		if fi, err := os.Stat(full); err == nil {
			w.ignoredDirs = append(w.ignoredDirs, fi)
			if w.opts.Debug {
				log.Printf("Directory added to ignore list: %s", full)
			}
		} else if w.opts.Debug {
			log.Printf("Error statting ignored directory: %v", err)
		}
	}
}

// getIgnoredDirs reads an optional config file at <path>/.goimportsignore
// of relative directories to ignore when scanning for go files.
// The provided path is one of the $GOPATH entries with "src" appended.
func (w *walker) getIgnoredDirs(path string) []string {
	file := filepath.Join(path, ".goimportsignore")
	slurp, err := ioutil.ReadFile(file)
	if w.opts.Debug {
		if err != nil {
			log.Print(err)
		} else {
			log.Printf("Read %s", file)
		}
	}
	if err != nil {
		return nil
	}

	var ignoredDirs []string
	bs := bufio.NewScanner(bytes.NewReader(slurp))
	for bs.Scan() {
		line := strings.TrimSpace(bs.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		ignoredDirs = append(ignoredDirs, line)
	}
	return ignoredDirs
}

// shouldSkipDir reports whether the file should be skipped or not.
func (w *walker) shouldSkipDir(fi os.FileInfo, dir string) bool {
	for _, ignoredDir := range w.ignoredDirs {
		if os.SameFile(fi, ignoredDir) {
			return true
		}
	}
	if w.skip != nil {
		// Check with the user specified callback.
		return w.skip(w.root, dir)
	}
	return false
}

// walk walks through the given path.
func (w *walker) walk(path string, typ os.FileMode) error {
	dir := filepath.Dir(path)
	if typ.IsRegular() {
		if dir == w.root.Path && (w.root.Type == RootGOROOT || w.root.Type == RootGOPATH) {
			// Doesn't make sense to have regular files
			// directly in your $GOPATH/src or $GOROOT/src.
			return fastwalk.SkipFiles
		}
		if !strings.HasSuffix(path, ".go") {
			return nil
		}

		w.add(w.root, dir)
		return fastwalk.SkipFiles
	}
	if typ == os.ModeDir {
		base := filepath.Base(path)
		if base == "" || base[0] == '.' || base[0] == '_' ||
			base == "testdata" ||
			(w.root.Type == RootGOROOT && w.opts.ModulesEnabled && base == "vendor") ||
			(!w.opts.ModulesEnabled && base == "node_modules") {
			return filepath.SkipDir
		}
		fi, err := os.Lstat(path)
		if err == nil && w.shouldSkipDir(fi, path) {
			return filepath.SkipDir
		}
		return nil
	}
	if typ == os.ModeSymlink {
		base := filepath.Base(path)
		if strings.HasPrefix(base, ".#") {
			// Emacs noise.
			return nil
		}
		fi, err := os.Lstat(path)
		if err != nil {
			// Just ignore it.
			return nil
		}
		if w.shouldTraverse(dir, fi) {
			return fastwalk.TraverseLink
		}
	}
	return nil
}

// shouldTraverse reports whether the symlink fi, found in dir,
// should be followed.  It makes sure symlinks were never visited
// before to avoid symlink loops.
func (w *walker) shouldTraverse(dir string, fi os.FileInfo) bool {
	path := filepath.Join(dir, fi.Name())
	target, err := filepath.EvalSymlinks(path)
	if err != nil {
		return false
	}
	ts, err := os.Stat(target)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return false
	}
	if !ts.IsDir() {
		return false
	}
	if w.shouldSkipDir(ts, dir) {
		return false
	}
	// Check for symlink loops by statting each directory component
	// and seeing if any are the same file as ts.
	for {
		parent := filepath.Dir(path)
		if parent == path {
			// Made it to the root without seeing a cycle.
			// Use this symlink.
			return true
		}
		parentInfo, err := os.Stat(parent)
		if err != nil {
			return false
		}
		if os.SameFile(ts, parentInfo) {
			// Cycle. Don't traverse.
			return false
		}
		path = parent
	}

}
<|endoftext|>